PR #14854: [ROCm] changes for spack

Imported from GitHub PR openxla/xla#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: -- 588be97790756155b507ac840ef3c15ddf61c683 by Chao Chen <[email protected]>: changes for spack by afzpatel Merging this change closes #14854 PiperOrigin-RevId: 659932164
google · Aug 6, 2024 · 4a72d05 · 4a72d05
1 parent 4de1c44
commit 4a72d05
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/tsl/platform/default/rocm_rocdl_path.cc b/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