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

- This removes `if_gpu_is_configured` guards from targets that are only supposed to be built for GPU. (Also tags them as `gpu` so that they get automatically excluded from the CPU build)
- This removes a bunch of `local_defines = if_cuda_is_configured([GOOGLE_CUDA])` parameters from targets that don't use the `GOOGLE_CUDA` preprocessor definition (anymore).
- Same for `TENSORFLOW_USE_ROCM`
- It tags many gpu-only targets with `gpu` so that they won't be built in a CPU-only build.

PiperOrigin-RevId: 664760472

xla/lit.bzl

@@ -246,6 +246,8 @@ def lit_test(
         )
         test_file = output_file
 
+    env["LD_LIBRARY_PATH"] = lib_dir
+
     native_test(
         name = name,
         src = lit_name,

xla/service/gpu/buffer_comparator_test.cc

@@ -21,9 +21,10 @@ limitations under the License.
 #include <limits>
 #include <vector>
 
+#include "absl/strings/ascii.h"
 #include "xla/primitive_util.h"
 #include "xla/service/gpu/stream_executor_util.h"
-#include "xla/service/hlo_module_config.h"
+#include "xla/service/platform_util.h"
 #include "xla/shape_util.h"
 #include "xla/stream_executor/device_memory.h"
 #include "xla/stream_executor/device_memory_handle.h"
@@ -44,13 +45,11 @@ constexpr double kDefaultTolerance = 0.1;
 class BufferComparatorTest : public testing::Test {
  protected:
   BufferComparatorTest()
-#if GOOGLE_CUDA
-      : platform_(se::PlatformManager::PlatformWithName("CUDA").value()),
-#elif TENSORFLOW_USE_ROCM
-      : platform_(se::PlatformManager::PlatformWithName("ROCM").value()),
-#endif
-        stream_exec_(platform_->ExecutorForDevice(0).value()) {
-  }
+      : platform_(stream_executor::PlatformManager::PlatformWithName(
+                      absl::AsciiStrToUpper(
+                          PlatformUtil::CanonicalPlatformName("gpu").value()))
+                      .value()),
+        stream_exec_(platform_->ExecutorForDevice(0).value()) {}
 
   // Take floats only for convenience. Still uses ElementType internally.
   template <typename ElementType>

xla/service/gpu/fusions/BUILD

@@ -12,6 +12,7 @@ cc_library(
     name = "in_place_dynamic_update_slice_mlir",
     srcs = ["in_place_dynamic_update_slice_mlir.cc"],
     hdrs = ["in_place_dynamic_update_slice_mlir.h"],
+    tags = ["gpu"],
     deps = [
         "//xla:shape_util",
         "//xla:xla_data_proto_cc",
@@ -39,6 +40,7 @@ cc_library(
     name = "copy",
     srcs = ["copy.cc"],
     hdrs = ["copy.h"],
+    tags = ["gpu"],
     deps = [
         ":fusion_emitter",
         "//xla:shape_util",
@@ -61,6 +63,7 @@ cc_library(
     srcs = ["custom.cc"],
     hdrs = ["custom.h"],
     local_defines = if_cuda_is_configured(["GOOGLE_CUDA=1"]),
+    tags = ["gpu"],
     deps = [
         ":fusion_emitter",
         "//xla:literal",
@@ -168,6 +171,7 @@ cc_library(
     name = "fusion_emitter",
     srcs = ["fusion_emitter.cc"],
     hdrs = ["fusion_emitter.h"],
+    tags = ["gpu"],
     visibility = ["//xla/service/gpu:__subpackages__"],
     deps = [
         "//xla:shape_util",
@@ -204,6 +208,7 @@ cc_library(
     name = "fusions",
     srcs = ["fusions.cc"],
     hdrs = ["fusions.h"],
+    tags = ["gpu"],
     visibility = ["//xla/service/gpu:__subpackages__"],
     deps = [
         ":concatenate_mlir",
@@ -245,6 +250,7 @@ cc_library(
     name = "loop_mlir",
     srcs = ["loop_mlir.cc"],
     hdrs = ["loop_mlir.h"],
+    tags = ["gpu"],
     deps = [
         "//xla:shape_util",
         "//xla:status_macros",
@@ -274,6 +280,7 @@ cc_library(
     name = "scatter_mlir",
     srcs = ["scatter_mlir.cc"],
     hdrs = ["scatter_mlir.h"],
+    tags = ["gpu"],
     deps = [
         "//xla:shape_util",
         "//xla:xla_data_proto_cc",
@@ -306,6 +313,7 @@ cc_library(
     name = "transpose_mlir",
     srcs = ["transpose_mlir.cc"],
     hdrs = ["transpose_mlir.h"],
+    tags = ["gpu"],
     deps = [
         ":fusion_emitter",
         "//xla:permutation_util",
@@ -340,6 +348,7 @@ cc_library(
     name = "triton",
     srcs = ["triton.cc"],
     hdrs = ["triton.h"],
+    tags = ["gpu"],
     visibility = ["//xla/service/gpu:__subpackages__"],
     deps = [
         ":fusion_emitter",
@@ -379,7 +388,7 @@ cc_library(
 xla_cc_test(
     name = "triton_test",
     srcs = ["triton_test.cc"],
-    local_defines = if_cuda_is_configured(["GOOGLE_CUDA=1"]),
+    tags = ["gpu"],
     deps = [
         ":fusion_emitter",
         ":fusions",
@@ -403,6 +412,7 @@ cc_library(
     srcs = ["cudnn.cc"],
     hdrs = ["cudnn.h"],
     local_defines = if_cuda_is_configured(["GOOGLE_CUDA=1"]),
+    tags = ["gpu"],
     deps = [
         ":fusion_emitter",
         "//xla/hlo/ir:hlo",
@@ -463,6 +473,7 @@ cc_library(
     name = "thunk_util",
     srcs = ["thunk_util.cc"],
     hdrs = ["thunk_util.h"],
+    tags = ["gpu"],
     visibility = ["//xla/service/gpu:__subpackages__"],
     deps = [
         "//xla:literal",
@@ -482,6 +493,7 @@ cc_library(
     name = "reduction_base",
     srcs = ["reduction_base.cc"],
     hdrs = ["reduction_base.h"],
+    tags = ["gpu"],
     visibility = ["//xla/service/gpu/fusions:__subpackages__"],
     deps = [
         ":fusion_emitter",
@@ -517,6 +529,7 @@ cc_library(
     name = "reduction_mlir",
     srcs = ["reduction_mlir.cc"],
     hdrs = ["reduction_mlir.h"],
+    tags = ["gpu"],
     deps = [
         ":fusion_emitter",
         ":reduction_base",
@@ -558,6 +571,7 @@ cc_library(
 xla_cc_test(
     name = "reduction_base_test",
     srcs = ["reduction_base_test.cc"],
+    tags = ["gpu"],
     deps = [
         ":reduction_base",
         "//xla/service/gpu:gpu_device_info_for_tests",
@@ -572,6 +586,7 @@ cc_library(
     name = "concatenate_mlir",
     srcs = ["concatenate_mlir.cc"],
     hdrs = ["concatenate_mlir.h"],
+    tags = ["gpu"],
     deps = [
         "//xla:shape_util",
         "//xla/hlo/ir:hlo",
@@ -598,6 +613,7 @@ cc_library(
     name = "input_slices_mlir",
     srcs = ["input_slices_mlir.cc"],
     hdrs = ["input_slices_mlir.h"],
+    tags = ["gpu"],
     deps = [
         "//xla:util",
         "//xla:xla_data_proto_cc",

xla/service/gpu/fusions/ir/tests/BUILD

@@ -9,6 +9,7 @@ lit_test_suite(
     name = "tests",
     srcs = glob(["*.mlir"]),
     cfg = "//xla:lit.cfg.py",
+    tags = ["gpu"],
     tools = [
         "//xla/service/gpu/fusions/tools:mlir_fusions_opt",
         "@llvm-project//llvm:FileCheck",

xla/service/gpu/fusions/ir/tests/inlining.mlir

@@ -1,4 +1,4 @@
-// RUN: mlir_fusions_opt %s -split-input-file -xla-erase-dead-functions -inline | FileCheck %s
+// RUN: env LD_DEBUG=files,libs mlir_fusions_opt %s -split-input-file -xla-erase-dead-functions -inline | FileCheck %s
 
 module {
   func.func private @mul(%a: f32, %b: f32) -> f32 {

xla/service/gpu/fusions/legacy/BUILD

@@ -10,6 +10,7 @@ cc_library(
     name = "in_place_dynamic_update_slice",
     srcs = ["in_place_dynamic_update_slice.cc"],
     hdrs = ["in_place_dynamic_update_slice.h"],
+    tags = ["gpu"],
     deps = [
         "//xla/hlo/ir:hlo",
         "//xla/service/gpu:hlo_fusion_analysis",
@@ -33,6 +34,7 @@ cc_library(
 xla_cc_test(
     name = "in_place_dynamic_update_slice_test",
     srcs = ["in_place_dynamic_update_slice_test.cc"],
+    tags = ["gpu"],
     deps = [
         ":in_place_dynamic_update_slice",
         "//xla/service/gpu:gpu_device_info_for_tests",
@@ -53,6 +55,7 @@ cc_library(
     name = "loop",
     srcs = ["loop.cc"],
     hdrs = ["loop.h"],
+    tags = ["gpu"],
     deps = [
         "//xla:shape_util",
         "//xla:util",
@@ -83,6 +86,7 @@ cc_library(
 xla_cc_test(
     name = "loop_test",
     srcs = ["loop_test.cc"],
+    tags = ["gpu"],
     deps = [
         "//xla:status_macros",
         "//xla/service/gpu:gpu_device_info_for_tests",
@@ -105,6 +109,7 @@ cc_library(
     name = "scatter",
     srcs = ["scatter.cc"],
     hdrs = ["scatter.h"],
+    tags = ["gpu"],
     deps = [
         ":loop",
         "//xla:shape_util",
@@ -135,6 +140,7 @@ cc_library(
 xla_cc_test(
     name = "scatter_test",
     srcs = ["scatter_test.cc"],
+    tags = ["gpu"],
     deps = [
         ":scatter",
         "//xla/service/gpu:gpu_device_info_for_tests",
@@ -155,6 +161,7 @@ cc_library(
     name = "tiling_util",
     srcs = ["tiling_util.cc"],
     hdrs = ["tiling_util.h"],
+    tags = ["gpu"],
     visibility = ["//xla/service/gpu:__subpackages__"],
     deps = [
         "//xla:shape_util",
@@ -183,6 +190,7 @@ cc_library(
     name = "reduction",
     srcs = ["reduction.cc"],
     hdrs = ["reduction.h"],
+    tags = ["gpu"],
     deps = [
         ":tiling_util",
         "//xla:shape_util",
@@ -236,6 +244,7 @@ cc_library(
 xla_cc_test(
     name = "reduction_test",
     srcs = ["reduction_test.cc"],
+    tags = ["gpu"],
     deps = [
         ":reduction",
         "//xla/service/gpu:gpu_device_info_for_tests",
@@ -255,6 +264,7 @@ cc_library(
     name = "concatenate",
     srcs = ["concatenate.cc"],
     hdrs = ["concatenate.h"],
+    tags = ["gpu"],
     deps = [
         "//xla:shape_util",
         "//xla/hlo/ir:hlo",
@@ -282,6 +292,7 @@ cc_library(
 xla_cc_test(
     name = "concatenate_test",
     srcs = ["concatenate_test.cc"],
+    tags = ["gpu"],
     deps = [
         ":concatenate",
         "//xla/service/gpu:gpu_device_info_for_tests",
@@ -301,6 +312,7 @@ cc_library(
     name = "transpose",
     srcs = ["transpose.cc"],
     hdrs = ["transpose.h"],
+    tags = ["gpu"],
     deps = [
         ":tiling_util",
         "//xla:permutation_util",
@@ -335,6 +347,7 @@ cc_library(
 xla_cc_test(
     name = "transpose_test",
     srcs = ["transpose_test.cc"],
+    tags = ["gpu"],
     deps = [
         ":transpose",
         "//xla:status_macros",
@@ -357,6 +370,7 @@ cc_library(
     name = "input_slices",
     srcs = ["input_slices.cc"],
     hdrs = ["input_slices.h"],
+    tags = ["gpu"],
     deps = [
         "//xla:shape_util",
         "//xla:util",
@@ -388,6 +402,7 @@ cc_library(
 xla_cc_test(
     name = "input_slices_test",
     srcs = ["input_slices_test.cc"],
+    tags = ["gpu"],
     deps = [
         ":input_slices",
         "//xla/service/gpu:gpu_device_info_for_tests",

xla/service/gpu/fusions/mlir/BUILD

@@ -17,6 +17,7 @@ cc_library(
     name = "computation_partitioner",
     srcs = ["computation_partitioner.cc"],
     hdrs = ["computation_partitioner.h"],
+    tags = ["gpu"],
     deps = [
         ":type_util",
         "//xla:shape_util",
@@ -47,6 +48,7 @@ cc_library(
 xla_cc_test(
     name = "computation_partitioner_test",
     srcs = ["computation_partitioner_test.cc"],
+    tags = ["gpu"],
     deps = [
         ":computation_partitioner",
         "//xla/hlo/ir:hlo",
@@ -64,6 +66,7 @@ cc_library(
     name = "elemental_hlo_to_mlir",
     srcs = ["elemental_hlo_to_mlir.cc"],
     hdrs = ["elemental_hlo_to_mlir.h"],
+    tags = ["gpu"],
     deps = [
         ":computation_partitioner",
         ":type_util",
@@ -113,6 +116,7 @@ cc_library(
 xla_cc_test(
     name = "elemental_hlo_to_mlir_test",
     srcs = ["elemental_hlo_to_mlir_test.cc"],
+    tags = ["gpu"],
     deps = [
         ":computation_partitioner",
         ":elemental_hlo_to_mlir",
@@ -153,6 +157,7 @@ cc_library(
     name = "mlir_fusion_emitter",
     srcs = ["mlir_fusion_emitter.cc"],
     hdrs = ["mlir_fusion_emitter.h"],
+    tags = ["gpu"],
     deps = [
         ":computation_partitioner",
         ":elemental_hlo_to_mlir",
@@ -231,6 +236,7 @@ cc_library(
 xla_cc_test(
     name = "mlir_fusion_emitter_test",
     srcs = ["mlir_fusion_emitter_test.cc"],
+    tags = ["gpu"],
     deps = [
         ":computation_partitioner",
         ":mlir_fusion_emitter",

xla/service/gpu/fusions/tests/BUILD

@@ -10,6 +10,7 @@ lit_test_suite(
     srcs = glob(["**/*.hlo"]),
     cfg = "//xla:lit.cfg.py",
     default_tags = ["requires-gpu-sm80-only"],
+    tags = ["gpu"],
     tools = [
         "//xla/service/gpu/fusions/tools:fusion_to_mlir",
         "//xla/service/gpu/fusions/tools:mlir_fusions_opt",

xla/service/gpu/fusions/tools/BUILD

@@ -8,6 +8,7 @@ package(
 xla_cc_binary(
     name = "mlir_fusions_opt",
     srcs = ["mlir_fusions_opt.cc"],
+    tags = ["gpu"],
     visibility = ["//xla/service/gpu/fusions:__subpackages__"],
     deps = [
         "//xla/mlir_hlo",
@@ -41,6 +42,7 @@ cc_library(
     testonly = 1,
     srcs = ["test_lib.cc"],
     hdrs = ["test_lib.h"],
+    tags = ["gpu"],
     deps = [
         "//xla:status_macros",
         "//xla/hlo/ir:hlo",
@@ -75,6 +77,7 @@ xla_cc_binary(
     name = "fusion_to_mlir",
     testonly = 1,
     srcs = ["fusion_to_mlir.cc"],
+    tags = ["gpu"],
     visibility = ["//xla/service/gpu/fusions:__subpackages__"],
     deps = [
         ":test_lib",
@@ -90,6 +93,7 @@ xla_cc_binary(
     name = "test_correctness",
     testonly = 1,
     srcs = ["test_correctness.cc"],
+    tags = ["gpu"],
     visibility = ["//xla/service/gpu/fusions:__subpackages__"],
     deps = [
         ":test_lib",