Uses getCurrentCUDASparseHandle() instead of THCState

mfbalin · Oct 16, 2020 · 8b565a9 · 8b565a9
1 parent 86f00ab
commit 8b565a9
Show file tree

Hide file tree

Showing 5 changed files with 3 additions and 16 deletions.
diff --git a/sparse-extension/build.out b/sparse-extension/build.out
diff --git a/...on/build/lib.linux-ppc64le-3.6/sparse_coo_tensor_cpp.cpython-36m-powerpc64le-linux-gnu.so b/...on/build/lib.linux-ppc64le-3.6/sparse_coo_tensor_cpp.cpython-36m-powerpc64le-linux-gnu.so
diff --git a/sparse-extension/build/temp.linux-ppc64le-3.6/sparse_coo_tensor.o b/sparse-extension/build/temp.linux-ppc64le-3.6/sparse_coo_tensor.o
diff --git a/sparse-extension/dist/sparse_coo_tensor_cpp-0.0.0-py3.6-linux-ppc64le.egg b/sparse-extension/dist/sparse_coo_tensor_cpp-0.0.0-py3.6-linux-ppc64le.egg
diff --git a/sparse-extension/sparse_coo_tensor.cpp b/sparse-extension/sparse_coo_tensor.cpp
@@ -1,4 +1,5 @@
 #include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
 #include <ATen/Layout.h>
 #include <ATen/Parallel.h>
 #include <ATen/SparseTensorImpl.h>
@@ -122,7 +123,8 @@ void spmm_gpu(const at::Tensor& A_rowindices,
     // cusparseHandle_t handle;
     // CHECK_CUSPARSE(cusparseCreate(&handle));
     auto state = at::globalContext().lazyInitCUDA();
-    auto handle = THCState_getCurrentSparseHandle(state);
+    // auto handle = THCState_getCurrentSparseHandle(state);
+    auto handle = at::cuda::getCurrentCUDASparseHandle();
 
     // Impl1 -- coo2csr + csrmm2
     int nnz = A_values.size(0);