PR #14854: [ROCm] changes for spack

Imported from GitHub PR #14854 Thanks @afzpatel for this work and this changes is for spack using ROCm libraries from ROCm/tensorflow-upstream#2578 @xla-rotation Copybara import of the project: -- 588be97 by Chao Chen <[email protected]>: changes for spack by afzpatel Merging this change closes #14854 FUTURE_COPYBARA_INTEGRATE_REVIEW=#14854 from ROCm:ci_spack_change 588be97 PiperOrigin-RevId: 659517988
openxla · Aug 5, 2024 · 83dc55e · 83dc55e
1 parent 81666d8
commit 83dc55e
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 3 deletions.
diff --git a/third_party/tsl/tsl/platform/default/rocm_rocdl_path.cc b/third_party/tsl/tsl/platform/default/rocm_rocdl_path.cc
@@ -26,7 +26,7 @@ limitations under the License.
 
 namespace tsl {
 
-string RocmRoot() {
+std::string RocmRoot() {
 #if TENSORFLOW_USE_ROCM
   if (const char* rocm_path_env = std::getenv("ROCM_PATH")) {
     VLOG(3) << "ROCM root = " << rocm_path_env;
@@ -40,6 +40,12 @@ string RocmRoot() {
 #endif
 }
 
-string RocdlRoot() { return io::JoinPath(RocmRoot(), "amdgcn/bitcode"); }
+std::string RocdlRoot() {
+  if (const char* device_lib_path_env = std::getenv("HIP_DEVICE_LIB_PATH")) {
+    return device_lib_path_env;
+  } else {
+    return io::JoinPath(RocmRoot(), "amdgcn/bitcode");
+  }
+}
 
 }  // namespace tsl
diff --git a/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc b/xla/service/gpu/llvm_gpu_backend/gpu_backend_lib.cc
@@ -17,6 +17,7 @@ limitations under the License.
 
 #include <algorithm>
 #include <cstdint>
+#include <cstdlib>
 #include <fstream>
 #include <functional>
 #include <ios>
@@ -855,7 +856,12 @@ absl::StatusOr<std::vector<uint8_t>> EmitModuleToHsaco(
     ir_fs->flush();
   }
   // Locate lld.
-  std::string lld_path = tsl::io::JoinPath(tsl::RocmRoot(), "llvm/bin");
+  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",