Add locks around cuda driver api and sparse calls

The driver API in CUDA is not thread safe(or we are using it incorrectly). We
need to protect these calls with mutexes. The cuSparse library also needs
to be protected with these locks.

Author: umar456

src/api/c/device.cpp

@@ -74,7 +74,7 @@ af_err af_get_active_backend(af_backend *result)
 af_err af_init()
 {
     try {
-        static std::once_flag flag;
+        thread_local std::once_flag flag;
         std::call_once(flag, []() {
                 getDeviceInfo();
             });

src/backend/cuda/jit.cpp

@@ -514,6 +514,7 @@ static kc_entry_t compileKernel(const char *ker_name, string jit_ker)
         reinterpret_cast<void*>(1)
     };
 
+    std::lock_guard<std::mutex> lock(getDriverApiMutex(getActiveDeviceId()));
     CU_CHECK(cuLinkCreate(5, linkOptions, linkOptionValues, &linkState));
     CU_CHECK(cuLinkAddData(linkState, CU_JIT_INPUT_PTX, (void*)ptx.get(),
                            ptx_size, ker_name, 0, NULL, NULL));
@@ -689,6 +690,7 @@ void evalNodes(vector<Param<T> >&outputs, vector<Node *> output_nodes)
         args.push_back((void *)&num_odims);
     }
 
+    std::lock_guard<std::mutex> lock(getDriverApiMutex(getActiveDeviceId()));
     CU_CHECK(cuLaunchKernel(ker,
                             blocks_x,
                             blocks_y,
@@ -739,7 +741,4 @@ template void evalNodes<intl   >(vector<Param<intl   > > &out, vector<Node *> no
 template void evalNodes<uintl  >(vector<Param<uintl  > > &out, vector<Node *> node);
 template void evalNodes<short  >(vector<Param<short  > > &out, vector<Node *> node);
 template void evalNodes<ushort >(vector<Param<ushort > > &out, vector<Node *> node);
-
-
-
 }

src/backend/cuda/platform.cpp

@@ -283,6 +283,10 @@ unsigned getMaxJitSize()
     return length;
 }
 
+std::mutex& getDriverApiMutex(int device) {
+  return DeviceManager::getInstance().driver_api_mutex[device];
+}
+
 int& tlocalActiveDeviceId()
 {
     thread_local int activeDeviceId = 0;

src/backend/cuda/platform.hpp

@@ -11,9 +11,6 @@
 
 #include <cuda.h>
 #include <cuda_runtime.h>
-#include <memory>
-#include <vector>
-#include <string>
 #include <memory.hpp>
 #include <GraphicsResourceManager.hpp>
 #include <cufft.hpp>
@@ -22,6 +19,11 @@
 #include <cusparse.hpp>
 #include <common/types.hpp>
 
+#include <memory>
+#include <mutex>
+#include <string>
+#include <vector>
+
 namespace cuda
 {
 int getBackend();
@@ -41,6 +43,8 @@ void devprop(char* d_name, char* d_platform, char *d_toolkit, char* d_compute);
 
 unsigned getMaxJitSize();
 
+std::mutex& getDriverApiMutex(int device);
+
 int getDeviceCount();
 
 int getActiveDeviceId();
@@ -111,6 +115,8 @@ class DeviceManager
         friend GraphicsResourceManager& interopManager();
 #endif
 
+        friend std::mutex& getDriverApiMutex(int device);
+
         friend std::string getDeviceInfo(int device);
 
         friend std::string getPlatformInfo();
@@ -159,6 +165,7 @@ class DeviceManager
 
         std::unique_ptr<MemoryManagerPinned> pinnedMemManager;
 
+        std::mutex driver_api_mutex[MAX_DEVICES];
 #if defined(WITH_GRAPHICS)
         std::unique_ptr<GraphicsResourceManager> gfxManagers[MAX_DEVICES];
 #endif

src/backend/cuda/sparse_blas.cpp

@@ -142,6 +142,11 @@ Array<T> matmul(const common::SparseArray<T> lhs, const Array<T> rhs,
 
     dim4 rStrides = rhs.strides();
 
+    // NOTE: The cuSparse library seems to be using the driver API in the
+    // implementation. This is causing issues with our JIT kernel generation.
+    // This may be a bug in the cuSparse library.
+    std::lock_guard<std::mutex> lock(getDriverApiMutex(getActiveDeviceId()));
+
     // Create Sparse Matrix Descriptor
     cusparseMatDescr_t descr = 0;
     CUSPARSE_CHECK(cusparseCreateMatDescr(&descr));