Automatic rollback of Use PTX version corresponding to CUDA version

Reverts 99f2569 PiperOrigin-RevId: 623905632
tensorflow · Apr 11, 2024 · 5b63270 · 5b63270
1 parent d7d4935
commit 5b63270
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 17 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,10 +2,6 @@ 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")
 
@@ -32,7 +28,6 @@ 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",

diff --git a/third_party/xla/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc b/third_party/xla/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc
@@ -321,20 +321,13 @@ std::unique_ptr<llvm::TargetMachine> NVPTXGetTargetMachine(
     const DebugOptions& debug_options) {
   // Figure out the exact name of the processor as known to the NVPTX backend
   // from the gpu_architecture flag.
-#if !defined(GOOGLE_CUDA) || CUDA_VERSION < 11000
-  const char feature_str[] = "+ptx65";
-#elif CUDA_VERSION < 13000
-  const char feature_str[] = {'+',
-                              'p',
-                              't',
-                              'x',
-                              '7' + (CUDA_VERSION / 1000) - 11,
-                              '0' + (CUDA_VERSION / 10) % 10};
-#else
-  const char feature_str[] = "+ptx84";
+#if defined(GOOGLE_CUDA) && CUDA_VERSION >= 12010
+  // use ptx81 for CUDA >= 12.1
+  return GetTargetMachine(target_triple, GetSmName(compute_capability),
+                          debug_options, /*feature_str=*/"+ptx81");
 #endif
   return GetTargetMachine(target_triple, GetSmName(compute_capability),
-                          debug_options, feature_str);
+                          debug_options, /*feature_str=*/"+ptx74");
 }
 
 using TargetModuleLinker =