PR #16337: [ROCM] disable setting workspace on ROCM

Imported from GitHub PR #16337 This is a follow-up PR for #16149 which was closed with unmerged commits. Here I disable setting workspace on ROCM. This feature seems to produce sporadic errors because rocblas_set_workspace uses blocking memory functions (e.g. hipMalloc/hipFree) to manipulate scratch memory (which do not work under stream capture). **Disabled** until we have a clear idea on how to make it working on ROCM. @xla-rotation: could you have a look please ? Copybara import of the project: -- 6f42f69 by Pavel Emeliyanenko <[email protected]>: disable workspace Merging this change closes #16337 COPYBARA_INTEGRATE_REVIEW=#16337 from ROCm:ci_rocblas_disable_workspace 6f42f69 PiperOrigin-RevId: 666524181
openxla · Aug 22, 2024 · 119d74c · 119d74c
1 parent 511141e
commit 119d74c
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/xla/stream_executor/rocm/rocm_blas.cc b/xla/stream_executor/rocm/rocm_blas.cc
@@ -366,15 +366,18 @@ absl::Status ROCMBlas::DoBlasInternalImpl(FuncT rocblas_func, Stream *stream,
                  << ToString(ret);
     }
   }
-#if TF_ROCM_VERSION >= 60000
+#if 0
+// pemeliya: the feature is disabled since rocblas does not perform well under
+// graph capture. rocblas_set_workspace seems to use blocking memory functions
+// like hipFree/hipMalloc which result in HIP_ERROR_StreamCaptureUnsupported
   {
     auto *workspace = GetWorkspace();
     auto *wptr = workspace != nullptr ? workspace->opaque() : nullptr;
     size_t wsize = workspace != nullptr ? workspace->size() : 0;
     ret = wrap::rocblas_set_workspace(blas_, wptr, wsize);
     if (err_on_failure && ret != rocblas_status_success) {
-      LOG(ERROR) << "failed to set workspace before " << FuncT::kName << ": "
-                 << ToString(ret);
+      LOG(ERROR) << "failed to set workspace before " << FuncT::kName
+                 << ": " << ToString(ret);
     }
   }
 #endif