Revert llvm env

ROCm · Nov 18, 2024 · bf81fb2 · bf81fb2
1 parent 167f26c
commit bf81fb2
Show file tree

Hide file tree

Showing 5 changed files with 23 additions and 46 deletions.
diff --git a/third_party/tsl/third_party/gpus/rocm/BUILD.tpl b/third_party/tsl/third_party/gpus/rocm/BUILD.tpl
@@ -53,19 +53,11 @@ cc_library(
     deps = [":rocm_config"],
 )
 
-sh_library(
-    name = "llvm_env",
-    srcs = ["llvm.sh"],
-    data = glob(["%{rocm_root}/llvm/**/*"]),
-    visibility = ["//visibility:public"],
-)
-
 cc_library(
     name = "rocblas",
     srcs = glob(["%{rocm_root}/lib/librocblas*.so*"]),
     hdrs = glob(["%{rocm_root}/include/rocblas/**"]),
     data = glob([
-        "%{rocm_root}/lib/librocblas*.so",
         "%{rocm_root}/lib/rocblas/**",
     ]),
     include_prefix = "rocm",

diff --git a/third_party/tsl/third_party/gpus/rocm/llvm.sh.tpl b/third_party/tsl/third_party/gpus/rocm/llvm.sh.tpl
diff --git a/third_party/tsl/third_party/gpus/rocm_configure.bzl b/third_party/tsl/third_party/gpus/rocm_configure.bzl
@@ -581,7 +581,6 @@ def _create_local_rocm_repository(repository_ctx):
 
     tpl_paths = {labelname: _tpl_path(repository_ctx, labelname) for labelname in [
         "rocm:build_defs.bzl",
-        "rocm:llvm.sh",
         "rocm:BUILD",
         "crosstool:BUILD.rocm",
         "crosstool:hipcc_cc_toolchain_config.bzl",
@@ -613,15 +612,6 @@ def _create_local_rocm_repository(repository_ctx):
 
     have_hipblaslt = "1" if rocm_libs["hipblaslt"] != None else "0"
 
-    repository_ctx.template(
-        "rocm/llvm.sh",
-        tpl_paths["rocm:llvm.sh"],
-        {
-            "%{rocm_root}": str(repository_ctx.path(rocm_config.rocm_toolkit_path)),
-        },
-        executable = True,
-    )
-
     # Set up BUILD file for rocm/
     repository_ctx.template(
         "rocm/build_defs.bzl",

diff --git a/xla/service/gpu/llvm_gpu_backend/BUILD b/xla/service/gpu/llvm_gpu_backend/BUILD
@@ -86,7 +86,6 @@ cc_library(
         "//xla/stream_executor/cuda:cuda_asm_compiler",
     ]) + if_rocm_is_configured([
         "@local_config_rocm//rocm:rocm_headers",
-        "@local_config_rocm//rocm:llvm_env",
         "@llvm-project//llvm:AMDGPUCodeGen",
         "@llvm-project//llvm:AMDGPUAsmParser",
     ]),

diff --git a/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc b/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc
@@ -711,15 +711,11 @@ std::vector<std::string> GetROCDLPaths(std::string gcn_arch_name,
                                        const std::string& rocdl_dir_path) {
   // AMDGPU version-neutral bitcodes.
   static std::vector<std::string>* rocdl_filenames =
-      new std::vector<std::string>{"opencl.bc",
-                                   "ocml.bc",
-                                   "ockl.bc",
-                                   "oclc_finite_only_off.bc",
-                                   "oclc_daz_opt_off.bc",
-                                   "oclc_correctly_rounded_sqrt_on.bc",
-                                   "oclc_unsafe_math_off.bc",
-                                   "oclc_wavefrontsize64_on.bc",
-                                   "oclc_abi_version_500.bc"};
+      new std::vector<std::string>(
+          {"opencl.bc", "ocml.bc", "ockl.bc", "oclc_finite_only_off.bc",
+           "oclc_daz_opt_off.bc", "oclc_correctly_rounded_sqrt_on.bc",
+           "oclc_unsafe_math_off.bc", "oclc_wavefrontsize64_on.bc",
+           "oclc_abi_version_500.bc"});
 
   // Construct full path to ROCDL bitcode libraries.
   std::vector<std::string> result;
@@ -859,21 +855,29 @@ absl::StatusOr<std::vector<uint8_t>> EmitModuleToHsaco(
     module->print(*ir_fs, nullptr);
     ir_fs->flush();
   }
+  // Locate lld.
+  std::string lld_path;
+  if (std::getenv("LLVM_PATH")) {
+    lld_path = tsl::io::JoinPath(std::getenv("LLVM_PATH"), "bin");
+  } else {
+    lld_path = tsl::io::JoinPath(tsl::RocmRoot(), "llvm/bin");
+  }
+  auto lld_program = llvm::sys::findProgramByName("ld.lld", {lld_path});
+  if (!lld_program) {
+    return xla::Internal("unable to find ld.lld in PATH: %s",
+                         lld_program.getError().message());
+  }
   std::vector<llvm::StringRef> lld_args{
-      llvm_ir::AsStringRef("external/local_config_rocm/rocm/llvm.sh"),
-      llvm_ir::AsStringRef("ld.lld"),
-      llvm_ir::AsStringRef("-flavor"),
-      llvm_ir::AsStringRef("gnu"),
-      llvm_ir::AsStringRef("-shared"),
-      llvm_ir::AsStringRef(isabin_path),
-      llvm_ir::AsStringRef("-o"),
+      llvm_ir::AsStringRef("ld.lld"),    llvm_ir::AsStringRef("-flavor"),
+      llvm_ir::AsStringRef("gnu"),       llvm_ir::AsStringRef("-shared"),
+      llvm_ir::AsStringRef(isabin_path), llvm_ir::AsStringRef("-o"),
       llvm_ir::AsStringRef(hsaco_path),
   };
 
   std::string error_message;
-  int lld_result = llvm::sys::ExecuteAndWait(
-      "external/local_config_rocm/rocm/llvm.sh", llvm_ir::AsArrayRef(lld_args),
-      std::nullopt, {}, 0, 0, &error_message);
+  int lld_result =
+      llvm::sys::ExecuteAndWait(*lld_program, llvm_ir::AsArrayRef(lld_args),
+                                std::nullopt, {}, 0, 0, &error_message);
   if (lld_result) {
     return xla::Internal("ld.lld execute fail: %s, error code %d",
                          error_message, lld_result);