replace at::cuda::getCurrentCUDASparseHandle with custom func (fix NVIDIA#308) (NVIDIA#315)

* force initialize cusparse handle

* replace all at::cuda::getCurrentCUDASparseHandle with custom func

* change log

Author: chrischoy

CHANGELOG.md

@@ -5,6 +5,7 @@
 - spmm average cuda function
 - SparseTensor list operators (cat, mean, sum, var)
 - MinkowskiStack containers
+- Replace all at::cuda::getCurrentCUDASparseHandle with custom getCurrentCUDASparseHandle (issue #308)
 
 ## [0.5.1]
 

src/broadcast_gpu.cu

@@ -88,7 +88,7 @@ at::Tensor BroadcastForwardGPU(
       torch::empty({in_feat.size(0), in_feat.size(1)}, in_feat.options());
 
   auto stream = at::cuda::getCurrentCUDAStream();
-  cusparseHandle_t handle = at::cuda::getCurrentCUDASparseHandle();
+  cusparseHandle_t handle = getCurrentCUDASparseHandle();
   cusparseSetStream(handle, stream);
 
   AT_DISPATCH_FLOATING_TYPES(
@@ -158,7 +158,7 @@ std::pair<at::Tensor, at::Tensor> BroadcastBackwardGPU(
   const auto &in_outs = p_map_manager->origin_map(p_in_map_key);
 
   auto stream = at::cuda::getCurrentCUDAStream();
-  cusparseHandle_t handle = at::cuda::getCurrentCUDASparseHandle();
+  cusparseHandle_t handle = getCurrentCUDASparseHandle();
   cusparseSetStream(handle, stream);
 
   AT_DISPATCH_FLOATING_TYPES(

src/global_pooling_gpu.cu

@@ -135,7 +135,7 @@ std::tuple<at::Tensor, at::Tensor> GlobalPoolingForwardGPU(
       case PoolingMode::GLOBAL_AVG_POOLING_KERNEL: {
         const auto &in_outs = p_map_manager->origin_map(p_in_map_key);
         cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();
-        cusparseHandle_t handle = at::cuda::getCurrentCUDASparseHandle();
+        cusparseHandle_t handle = getCurrentCUDASparseHandle();
         cusparseSetStream(handle, stream);
 
         TemplatedAllocator<char> byte_allocator;

src/gpu.cu

@@ -115,6 +115,12 @@ const char *cusparseGetErrorString(cusparseStatus_t error) {
   return "<unknown>";
 }
 
+cusparseHandle_t getCurrentCUDASparseHandle() {
+  cusparseHandle_t handle;
+  CUSPARSE_CHECK(cusparseCreate(&handle));
+  return handle;
+}
+
 static std::string format_size(uint64_t size) {
   std::ostringstream os;
   os.precision(2);

src/gpu.cuh

@@ -163,6 +163,8 @@ const char *cublasGetErrorString(cublasStatus_t error);
 // CUSparse error reporting.
 const char *cusparseGetErrorString(cusparseStatus_t error);
 
+cusparseHandle_t getCurrentCUDASparseHandle();
+
 constexpr uint32_t CUDA_NUM_THREADS = 128;
 
 constexpr uint32_t SHARED_BLOCK_SIZE = 32;

src/local_pooling_gpu.cu

@@ -117,7 +117,7 @@ std::pair<at::Tensor, at::Tensor> LocalPoolingForwardGPU(
       num_nonzero.resize_({out_nrows});
       num_nonzero.zero_();
     }
-    cusparseHandle_t handle = at::cuda::getCurrentCUDASparseHandle();
+    cusparseHandle_t handle = getCurrentCUDASparseHandle();
     cusparseSetStream(handle, stream);
 
     AT_DISPATCH_FLOATING_TYPES(

src/local_pooling_transpose_gpu.cu

@@ -111,7 +111,7 @@ std::pair<at::Tensor, at::Tensor> LocalPoolingTransposeForwardGPU(
   at::Tensor num_nonzero =
       torch::empty({0}, in_feat.options().requires_grad(false));
 
-  cusparseHandle_t handle = at::cuda::getCurrentCUDASparseHandle();
+  cusparseHandle_t handle = getCurrentCUDASparseHandle();
   cusparseSetStream(handle, stream);
 
   AT_DISPATCH_FLOATING_TYPES(

src/spmm.cu

@@ -186,7 +186,11 @@ torch::Tensor coo_spmm(torch::Tensor const &rows, torch::Tensor const &cols,
 
   // Dense matrices have to be contiguous for cusparseSpMM to work
   torch::Tensor const mat2_contig = mat2.contiguous();
-  auto cusparse_handle = at::cuda::getCurrentCUDASparseHandle();
+  // Issue 308
+  // auto cusparse_handle = at::cuda::getCurrentCUDASparseHandle();
+  auto stream = at::cuda::getCurrentCUDAStream();
+  cusparseHandle_t cusparse_handle = getCurrentCUDASparseHandle();
+  cusparseSetStream(cusparse_handle, stream);
 
   torch::Scalar beta = 0;
   torch::Scalar alpha = 1;
@@ -442,7 +446,11 @@ coo_spmm_average(torch::Tensor const &rows, torch::Tensor const &cols,
 
   // Dense matrices have to be contiguous for cusparseSpMM to work
   torch::Tensor const mat2_contig = mat2.contiguous();
-  auto cusparse_handle = at::cuda::getCurrentCUDASparseHandle();
+  // Issue 308
+  // auto cusparse_handle = at::cuda::getCurrentCUDASparseHandle();
+  auto stream = at::cuda::getCurrentCUDAStream();
+  cusparseHandle_t cusparse_handle = getCurrentCUDASparseHandle();
+  cusparseSetStream(cusparse_handle, stream);
 
   torch::Scalar beta = 0;
   torch::Scalar alpha = 1;