Fix llvm_gpu_backend build defines

PiperOrigin-RevId: 626282953
tensorflow · Apr 19, 2024 · 6fda823 · 6fda823
1 parent f8bfe73
commit 6fda823
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/third_party/xla/xla/service/gpu/llvm_gpu_backend/BUILD b/third_party/xla/xla/service/gpu/llvm_gpu_backend/BUILD
@@ -2,6 +2,10 @@ load(
     "@local_config_rocm//rocm:build_defs.bzl",
     "if_rocm_is_configured",
 )
+load(
+    "@local_tsl//tsl/platform/default:cuda_build_defs.bzl",
+    "if_cuda_is_configured",
+)
 load("//xla:xla.bzl", "xla_cc_test")
 load("//xla/tsl:tsl.bzl", "internal_visibility")
 
@@ -28,6 +32,7 @@ cc_library(
         "gpu_backend_lib.h",
         "utils.h",
     ],
+    local_defines = if_cuda_is_configured(["GOOGLE_CUDA=1"]) + if_rocm_is_configured(["TENSORFLOW_USE_ROCM=1"]),
     deps = [
         "//xla:status",
         "//xla:status_macros",