Fixed destroying cublas handle upon thread destruction

oneapi-src · Nov 5, 2024 · c8420f6 · c8420f6
1 parent 4938e48
commit c8420f6
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 6 deletions.
diff --git a/src/blas/backends/cublas/cublas_scope_handle.cpp b/src/blas/backends/cublas/cublas_scope_handle.cpp
@@ -23,22 +23,32 @@ namespace mkl {
 namespace blas {
 namespace cublas {
 
+thread_local std::shared_ptr<cublasHandle_t> CublasScopedContextHandler::cublasHandle{ nullptr };
+
 CublasScopedContextHandler::CublasScopedContextHandler(sycl::interop_handle& ih)
         : ih(ih),
           nativeDevice(ih.get_native_device<sycl::backend::ext_oneapi_cuda>()) {
-    cublasStatus_t err;
-    CUBLAS_ERROR_FUNC(cublasCreate, err, &cublasHandle);
+    if (!cublasHandle) {
+        cublasHandle =
+            std::shared_ptr<cublasHandle_t>(new cublasHandle_t, [](cublasHandle_t* handle) {
+                cublasStatus_t err;
+                CUBLAS_ERROR_FUNC(cublasDestroy, err, *handle);
+                delete handle;
+            });
+        cublasStatus_t err;
+        CUBLAS_ERROR_FUNC(cublasCreate, err, cublasHandle.get());
+    }
 }
 
 cublasHandle_t CublasScopedContextHandler::get_handle(const sycl::queue& queue) {
     CUstream streamId = get_stream(queue);
     cudaStream_t currentStreamId;
     cublasStatus_t err;
-    CUBLAS_ERROR_FUNC(cublasGetStream, err, cublasHandle, &currentStreamId);
+    CUBLAS_ERROR_FUNC(cublasGetStream, err, *cublasHandle, &currentStreamId);
     if (currentStreamId != streamId) {
-        CUBLAS_ERROR_FUNC(cublasSetStream, err, cublasHandle, streamId);
+        CUBLAS_ERROR_FUNC(cublasSetStream, err, *cublasHandle, streamId);
     }
-    return cublasHandle;
+    return *cublasHandle;
 }
 
 CUstream CublasScopedContextHandler::get_stream(const sycl::queue& queue) {

diff --git a/src/blas/backends/cublas/cublas_scope_handle.hpp b/src/blas/backends/cublas/cublas_scope_handle.hpp
@@ -51,7 +51,7 @@ According to NVIDIA:
 class CublasScopedContextHandler {
     sycl::interop_handle& ih;
     CUdevice nativeDevice;
-    cublasHandle_t cublasHandle;
+    static thread_local std::shared_ptr<cublasHandle_t> cublasHandle;
     CUstream get_stream(const sycl::queue& queue);
     sycl::context get_context(const sycl::queue& queue);