Remove more GPU/CUDA/ROCm attribute guards from xla/service/gpu

PiperOrigin-RevId: 664760472
openxla · Aug 22, 2024 · 8629ab6 · 8629ab6
1 parent 6fabf59
commit 8629ab6
Show file tree

Hide file tree

Showing 25 changed files with 301 additions and 295 deletions.
diff --git a/build_tools/build.py b/build_tools/build.py
@@ -193,8 +193,9 @@ def nvidia_gpu_build_with_compute_capability(
       image_url=_CUDNN_9_IMAGE,
       target_patterns=_XLA_DEFAULT_TARGET_PATTERNS,
       configs=configs,
-      test_tag_filters=("-no_oss", "requires-gpu-nvidia") + extra_gpu_tags,
-      build_tag_filters=("-no_oss", "requires-gpu-nvidia"),
+      test_tag_filters=("-no_oss", "requires-gpu-nvidia", "gpu")
+      + extra_gpu_tags,
+      build_tag_filters=("-no_oss", "requires-gpu-nvidia", "gpu"),
       options=dict(
           run_under="//tools/ci_build/gpu_build:parallel_gpu_execute",
           repo_env=f"TF_CUDA_COMPUTE_CAPABILITIES={compute_capability/10}",
@@ -376,5 +377,6 @@ def main():
   for cmd in build.commands():
     sh(cmd)
 
+
 if __name__ == "__main__":
   main()
diff --git a/xla/lit.bzl b/xla/lit.bzl
@@ -89,6 +89,7 @@ def lit_test_suite(
     tools = tools or []
     default_tags = default_tags or []
     tags_override = tags_override or {}
+    tags = kwargs.pop("tags", [])
 
     tests = []
     for test_file in srcs:
@@ -107,14 +108,15 @@ def lit_test_suite(
             visibility = visibility,
             env = env,
             timeout = timeout,
-            tags = default_tags + tags_override.get(test_file, []),
+            tags = tags + default_tags + tags_override.get(test_file, []),
             hermetic_cuda_data_dir = hermetic_cuda_data_dir,
             **kwargs
         )
 
     native.test_suite(
         name = name,
         tests = tests,
+        tags = tags,
         **kwargs
     )