CSR: ~3x faster creation from scipy CSR. added tocsr_scipy() export method

Author: grlee77

scikits/cuda/cusparse.py

@@ -937,17 +937,30 @@ def __init__(self, handle, descr, csrVal, csrRowPtr, csrColInd, shape):
     # alternative constructor from dense ndarray, gpuarray or cuSPARSE matrix
     @classmethod
     def to_CSR(cls, A, handle):
-        """ convert dense any scipy.sparse matrix formats to cuSPARSE CSR.
+        """ convert dense numpy or gpuarray matrices as well as any
+        scipy.sparse matrix formats to cuSPARSE CSR.
         """
         if has_scipy and isinstance(A, scipy.sparse.spmatrix):
             """Convert scipy.sparse CSR, COO, BSR, etc to cuSPARSE CSR"""
             # converting BSR, COO, etc to CSR
+            if A.dtype.char not in ['f', 'd', 'F', 'D']:
+                raise ValueError("unsupported numpy dtype {}".format(A.dtype))
+
             if not isinstance(A, scipy.sparse.csr_matrix):
                 A = A.tocsr()
 
-            csrRowPtr = gpuarray.to_gpu(A.indptr.astype(np.int32))
-            csrColInd = gpuarray.to_gpu(A.indices.astype(np.int32))
-            csrVal = gpuarray.to_gpu(A.data.astype(A.dtype))
+            # avoid .astype() calls if possible for speed
+            if A.indptr.dtype != np.int32:
+                csrRowPtr = gpuarray.to_gpu(A.indptr.astype(np.int32))
+            else:
+                csrRowPtr = gpuarray.to_gpu(A.indptr)
+
+            if A.indices.dtype != np.int32:
+                csrColInd = gpuarray.to_gpu(A.indices.astype(np.int32))
+            else:
+                csrColInd = gpuarray.to_gpu(A.indices)
+
+            csrVal = gpuarray.to_gpu(A.data)
             descr = cusparseCreateMatDescr()
             cusparseSetMatType(descr, CUSPARSE_MATRIX_TYPE_GENERAL)
             cusparseSetMatIndexBase(descr, CUSPARSE_INDEX_BASE_ZERO)
@@ -1031,6 +1044,14 @@ def getnnz(self):
         """ return number of non-zeros"""
         return self.nnz
 
+    def tocsr_scipy(self):
+        """ return as scipy csr_matrix in host memory """
+        from scipy.sparse import csr_matrix
+        return csr_matrix((self.data.get(),
+                           self.indices.get(),
+                           self.indptr.get()),
+                          shape=self.shape)
+
     def todense(self, lda=None, to_cpu=False, handle=None, stream=None,
                 autosync=True):
         """ return dense gpuarray if to_cpu=False, numpy ndarray if to_cpu=True

tests/test_cusparse.py

@@ -872,6 +872,29 @@ def test_CSR_todense():
         cusparseDestroy(h)
 
 
+def test_CSR_tocsr_scipy():
+    n = 64
+    h = cusparseCreate()
+    try:
+        dtype = np.float32
+        A = 2*np.eye(n)
+
+        # generate a CSR matrix from a dense numpy array
+        A_CSR = CSR.to_CSR(A.astype(dtype), h)
+
+        # convert cusparseCSR back to a dense matrix
+        csr_scipy = A_CSR.tocsr_scipy()
+        # assert_equal(A_CSR.data.get(), csr_scipy.data)
+        # assert_equal(A_CSR.indices.get(), csr_scipy.indices)
+        # assert_equal(A_CSR.indptr.get(), csr_scipy.indptr)
+
+        assert_equal(csr_scipy.indptr, scipy.sparse.csr_matrix(A).indptr)
+        assert_equal(csr_scipy.indices, scipy.sparse.csr_matrix(A).indices)
+        assert_equal(csr_scipy.data, scipy.sparse.csr_matrix(A).data)
+    finally:
+        cusparseDestroy(h)
+
+
 def test_CSR_mv():
     n = 64
     h = cusparseCreate()