[dev-upstream] Update llvm path + Add wrap interface (#2577)

* [dev-upstream] Update llvm path after compiler change

Added the logic to add new path from ROCm >=6.3.0

To fix, https://ontrack-internal.amd.com/browse/SWDEV-470815

Pushing #2575 to all branches as need.

* Update rocm_configure.bzl

* Update rocblas_wrapper.h

Merging back https://github.com/ROCm/tensorflow-upstream/pull/2572/files

* Update rocm_blas.cc

third_party/gpus/rocm_configure.bzl

@@ -210,6 +210,10 @@ def _rocm_include_path(repository_ctx, rocm_config, bash_bin):
     inc_dirs.append(rocm_config.llvm_path + "/lib/clang/17/include")
     inc_dirs.append(rocm_config.llvm_path + "/lib/clang/18/include")
     inc_dirs.append(rocm_config.llvm_path + "/lib/clang/19/include")
+    if int(rocm_config.rocm_version_number) >= 60200:
+        inc_dirs.append(rocm_toolkit_path + "/lib/llvm/lib/clang/17/include")
+        inc_dirs.append(rocm_toolkit_path + "/lib/llvm/lib/clang/18/include")
+        inc_dirs.append(rocm_toolkit_path + "/lib/llvm/lib/clang/19/include")
 
     # Support hcc based off clang 10.0.0 (for ROCm 3.3)
     rocm_toolkit_path = realpath(repository_ctx, rocm_config.rocm_toolkit_path, bash_bin)

third_party/xla/xla/stream_executor/rocm/rocblas_wrapper.h

@@ -262,18 +262,22 @@ using stream_executor::internal::CachedDsoLoader::GetRocblasDsoHandle;
   __macro(rocblas_gemm_batched_ex_get_solutions)         \
   __macro(rocblas_gemm_batched_ex_get_solutions_by_type) \
   __macro(rocblas_gemm_strided_batched_ex_get_solutions) \
-  __macro(rocblas_is_managing_device_memory)             \
-  __macro(rocblas_is_user_managing_device_memory)        \
-  __macro(rocblas_set_workspace)                         \
-  __macro(rocblas_strsm_batched)                         \
-  __macro(rocblas_dtrsm_batched)                         \
-  __macro(rocblas_ctrsm_batched)                         \
-  __macro(rocblas_ztrsm_batched)                         \
-  __macro(rocblas_create_handle)                         \
-  __macro(rocblas_destroy_handle)                        \
-  __macro(rocblas_get_stream)                            \
-  __macro(rocblas_set_stream)                            \
-  __macro(rocblas_set_atomics_mode)
+  __macro(rocblas_strsm_batched)                \
+  __macro(rocblas_dtrsm_batched)                \
+  __macro(rocblas_ctrsm_batched)                \
+  __macro(rocblas_ztrsm_batched)                \
+  __macro(rocblas_create_handle)                \
+  __macro(rocblas_destroy_handle)               \
+  __macro(rocblas_get_stream)                   \
+  __macro(rocblas_set_stream)                   \
+  __macro(rocblas_set_atomics_mode)             \
+  __macro(rocblas_get_version_string)           \
+  __macro(rocblas_get_version_string_size)      \
+  __macro(rocblas_is_managing_device_memory)    \
+  __macro(rocblas_is_user_managing_device_memory) \
+  __macro(rocblas_set_workspace)                 \
+  __macro(rocblas_create_handle)                         
+
 
 // clang-format on
 

third_party/xla/xla/stream_executor/rocm/rocm_blas.cc

@@ -1259,16 +1259,16 @@ IMPL_DoBlasGemmBatched(float, wrap::rocblas_sgemm_strided_batched)
 }
 
 absl::Status ROCMBlas::GetVersion(std::string *version) {
-#if TF_ROCM_VERSION >= 60300  // Not yet available in ROCM-6.1
+#if TF_ROCM_VERSION > 60100  // Not available in ROCM-6.1
   absl::MutexLock lock{&mu_};
   size_t len = 0;
-  if (auto res = rocblas_get_version_string_size(&len);
+  if (auto res = wrap::rocblas_get_version_string_size(&len);
       res != rocblas_status_success) {
     return absl::InternalError(
         absl::StrCat("GetVersion failed with: ", ToString(res)));
   }
   std::vector<char> buf(len + 1);
-  if (auto res = rocblas_get_version_string(buf.data(), len);
+  if (auto res = wrap::rocblas_get_version_string(buf.data(), len);
       res != rocblas_status_success) {
     return absl::InternalError(
         absl::StrCat("GetVersion failed with: ", ToString(res)));