[cublas] Fix race condition in cublas handle deletion (#136)

* [cublas] Fix race condition in cublas handle deletion

This was causing segmentation faults at the end of the program when
running tests.

Cublas handles need to be created per thread and per context, therefore
they're currently being stored in a `static thread_local` map, mapping
contexts to cublas handles.

These handles are deleted in two places, the map destructor called when
the thread is deleted, or the context destructor, where a callback is
registered on the context to delete the cublas handle.

Since there is two places where a handle can be deleted and that these
two will likely run in separate threads, the handle was placed in an
C++ atomic to avoid race conditions, and allocated on the heap so that
it can stay alive when one of either the thread or the context is
deleted.

But there was two issues with this, the main one is that the callback
registered on the context was not given the pointer to the atomic
handle, rather the pointer to the slot in the `thread_local` map where
the atomic pointer was stored. Which meant that whenever the thread
would be deleted before the context, the pointer available to the
callback would be invalid as the map was deleted. The solution for that
is to simply pass directly the pointer to the atomic handle to the
callback.

The second issue is that both of these places were deleting the atomic
in all cases, which means one of the deletion was invalid. The solution
for this is to only do the deletion in the last one being called, that
is to say the one which will get a `nullptr` when atomically accessing
the handle pointer.

With these two patches the tests are running fine.

* [cublas] Fix formatting of cublas handle patch

* Fix half types for DPC++

`half` has been removed from the global scope in DPC++, see
intel/llvm#4818

Author: npmiller

include/oneapi/mkl/types.hpp

@@ -108,10 +108,8 @@ enum class order : char {
 } //namespace mkl
 } //namespace oneapi
 
-// Workaround for supporting ::half for hipSYCL
+// Workaround for supporting ::half for hipSYCL and DPC++
 // TODO: This should be removed after the interface is SYCL2020 conformant
-#ifdef __HIPSYCL__
 using ::cl::sycl::half;
-#endif
 
 #endif //_ONEMKL_TYPES_HPP_

src/blas/backends/cublas/cublas_handle.hpp

@@ -26,28 +26,33 @@ namespace mkl {
 namespace blas {
 namespace cublas {
 
-template<typename T>
+template <typename T>
 struct cublas_handle {
     using handle_container_t = std::unordered_map<T, std::atomic<cublasHandle_t> *>;
     handle_container_t cublas_handle_mapper_{};
-    ~cublas_handle() noexcept(false){
-    for (auto &handle_pair : cublas_handle_mapper_) {
-        cublasStatus_t err;
-        if (handle_pair.second != nullptr) {
-            auto handle = handle_pair.second->exchange(nullptr);
-            if (handle != nullptr) {
-                CUBLAS_ERROR_FUNC(cublasDestroy, err, handle);
-                handle = nullptr;
+    ~cublas_handle() noexcept(false) {
+        for (auto &handle_pair : cublas_handle_mapper_) {
+            cublasStatus_t err;
+            if (handle_pair.second != nullptr) {
+                auto handle = handle_pair.second->exchange(nullptr);
+                if (handle != nullptr) {
+                    CUBLAS_ERROR_FUNC(cublasDestroy, err, handle);
+                    handle = nullptr;
+                }
+                else {
+                    // if the handle is nullptr it means the handle was already
+                    // destroyed by the ContextCallback and we're free to delete the
+                    // atomic object.
+                    delete handle_pair.second;
+                }
+
+                handle_pair.second = nullptr;
             }
-            delete handle_pair.second;
-            handle_pair.second = nullptr;
         }
+        cublas_handle_mapper_.clear();
     }
-    cublas_handle_mapper_.clear();
-}
 };
 
-
 } // namespace cublas
 } // namespace blas
 } // namespace mkl

src/blas/backends/cublas/cublas_scope_handle.cpp

@@ -63,19 +63,21 @@ CublasScopedContextHandler::~CublasScopedContextHandler() noexcept(false) {
 }
 
 void ContextCallback(void *userData) {
-    auto *ptr = static_cast<std::atomic<cublasHandle_t> **>(userData);
+    auto *ptr = static_cast<std::atomic<cublasHandle_t> *>(userData);
     if (!ptr) {
         return;
     }
-    if (*ptr != nullptr) {
-        auto handle = (*ptr)->exchange(nullptr);
-        if (handle != nullptr) {
-            cublasStatus_t err1;
-            CUBLAS_ERROR_FUNC(cublasDestroy, err1, handle);
-            handle = nullptr;
-        }
-        delete *ptr;
-        *ptr = nullptr;
+    auto handle = ptr->exchange(nullptr);
+    if (handle != nullptr) {
+        cublasStatus_t err1;
+        CUBLAS_ERROR_FUNC(cublasDestroy, err1, handle);
+        handle = nullptr;
+    }
+    else {
+        // if the handle is nullptr it means the handle was already destroyed by
+        // the cublas_handle destructor and we're free to delete the atomic
+        // object.
+        delete ptr;
     }
 }
 
@@ -113,9 +115,8 @@ cublasHandle_t CublasScopedContextHandler::get_handle(const cl::sycl::queue &que
     auto insert_iter = handle_helper.cublas_handle_mapper_.insert(
         std::make_pair(piPlacedContext_, new std::atomic<cublasHandle_t>(handle)));
 
-    auto ptr = &(insert_iter.first->second);
-
-    sycl::detail::pi::contextSetExtendedDeleter(placedContext_, ContextCallback, ptr);
+    sycl::detail::pi::contextSetExtendedDeleter(placedContext_, ContextCallback,
+                                                insert_iter.first->second);
 
     return handle;
 }