Add wrap interface for calling into rocblas shared library

ROCm · Jun 26, 2024 · e2b1753 · e2b1753
1 parent edc830c
commit e2b1753
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 4 deletions.
diff --git a/third_party/xla/xla/stream_executor/rocm/rocblas_wrapper.h b/third_party/xla/xla/stream_executor/rocm/rocblas_wrapper.h
@@ -270,7 +270,9 @@ using stream_executor::internal::CachedDsoLoader::GetRocblasDsoHandle;
   __macro(rocblas_destroy_handle)               \
   __macro(rocblas_get_stream)                   \
   __macro(rocblas_set_stream)                   \
-  __macro(rocblas_set_atomics_mode)
+  __macro(rocblas_set_atomics_mode)             \
+  __macro(rocblas_get_version_string)           \
+  __macro(rocblas_get_version_string_size)
 
 // clang-format on
 

diff --git a/third_party/xla/xla/stream_executor/rocm/rocm_blas.cc b/third_party/xla/xla/stream_executor/rocm/rocm_blas.cc
@@ -1351,16 +1351,16 @@ absl::Status ROCMBlas::DoBlasGemmStridedBatched(
 }
 
 absl::Status ROCMBlas::GetVersion(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)));