Sparse CSR CUDA: Add block torch.addmv when mat is sparse (pytorch#68708)

IvanYashchuk · facebook-github-bot · commit a8232ee1bc8f · 2021-12-07T14:02:59.000-08:00
Summary: Pull Request resolved: pytorch#68708 This PR adds block CSR matrix times dense vector multiplication. cc nikitaved pearu cpuhrsch IvanYashchuk ngimel Test Plan: Imported from OSS Reviewed By: pbelevich Differential Revision: D32647694 Pulled By: cpuhrsch fbshipit-source-id: a1c120691c4350284b156fe4259eda684b734b66
diff --git a/aten/src/ATen/cuda/CUDASparseBlas.cpp b/aten/src/ATen/cuda/CUDASparseBlas.cpp
@@ -311,6 +311,87 @@ void bsrmm<c10::complex<double>>(
       ldc));
 }
 
+template <>
+void bsrmv<float>(CUSPARSE_BSRMV_ARGTYPES(float)) {
+  TORCH_CUDASPARSE_CHECK(cusparseSbsrmv(
+      handle,
+      dirA,
+      transA,
+      mb,
+      nb,
+      nnzb,
+      alpha,
+      descrA,
+      bsrValA,
+      bsrRowPtrA,
+      bsrColIndA,
+      blockDim,
+      x,
+      beta,
+      y));
+}
+
+template <>
+void bsrmv<double>(CUSPARSE_BSRMV_ARGTYPES(double)) {
+  TORCH_CUDASPARSE_CHECK(cusparseDbsrmv(
+      handle,
+      dirA,
+      transA,
+      mb,
+      nb,
+      nnzb,
+      alpha,
+      descrA,
+      bsrValA,
+      bsrRowPtrA,
+      bsrColIndA,
+      blockDim,
+      x,
+      beta,
+      y));
+}
+
+template <>
+void bsrmv<c10::complex<float>>(CUSPARSE_BSRMV_ARGTYPES(c10::complex<float>)) {
+  TORCH_CUDASPARSE_CHECK(cusparseCbsrmv(
+      handle,
+      dirA,
+      transA,
+      mb,
+      nb,
+      nnzb,
+      reinterpret_cast<const cuComplex*>(alpha),
+      descrA,
+      reinterpret_cast<const cuComplex*>(bsrValA),
+      bsrRowPtrA,
+      bsrColIndA,
+      blockDim,
+      reinterpret_cast<const cuComplex*>(x),
+      reinterpret_cast<const cuComplex*>(beta),
+      reinterpret_cast<cuComplex*>(y)));
+}
+
+template <>
+void bsrmv<c10::complex<double>>(
+    CUSPARSE_BSRMV_ARGTYPES(c10::complex<double>)) {
+  TORCH_CUDASPARSE_CHECK(cusparseZbsrmv(
+      handle,
+      dirA,
+      transA,
+      mb,
+      nb,
+      nnzb,
+      reinterpret_cast<const cuDoubleComplex*>(alpha),
+      descrA,
+      reinterpret_cast<const cuDoubleComplex*>(bsrValA),
+      bsrRowPtrA,
+      bsrColIndA,
+      blockDim,
+      reinterpret_cast<const cuDoubleComplex*>(x),
+      reinterpret_cast<const cuDoubleComplex*>(beta),
+      reinterpret_cast<cuDoubleComplex*>(y)));
+}
+
 } // namespace sparse
 } // namespace cuda
 } // namespace at
diff --git a/aten/src/ATen/cuda/CUDASparseBlas.h b/aten/src/ATen/cuda/CUDASparseBlas.h
@@ -130,6 +130,30 @@ void bsrmm<c10::complex<float>>(CUSPARSE_BSRMM_ARGTYPES(c10::complex<float>));
 template <>
 void bsrmm<c10::complex<double>>(CUSPARSE_BSRMM_ARGTYPES(c10::complex<double>));
 
+#define CUSPARSE_BSRMV_ARGTYPES(scalar_t)                                    \
+  cusparseHandle_t handle, cusparseDirection_t dirA,                         \
+      cusparseOperation_t transA, int mb, int nb, int nnzb,                  \
+      const scalar_t *alpha, const cusparseMatDescr_t descrA,                \
+      const scalar_t *bsrValA, const int *bsrRowPtrA, const int *bsrColIndA, \
+      int blockDim, const scalar_t *x, const scalar_t *beta, scalar_t *y
+
+template <typename scalar_t>
+inline void bsrmv(CUSPARSE_BSRMV_ARGTYPES(scalar_t)) {
+  TORCH_INTERNAL_ASSERT(
+      false,
+      "at::cuda::sparse::bsrmv: not implemented for ",
+      typeid(scalar_t).name());
+}
+
+template <>
+void bsrmv<float>(CUSPARSE_BSRMV_ARGTYPES(float));
+template <>
+void bsrmv<double>(CUSPARSE_BSRMV_ARGTYPES(double));
+template <>
+void bsrmv<c10::complex<float>>(CUSPARSE_BSRMV_ARGTYPES(c10::complex<float>));
+template <>
+void bsrmv<c10::complex<double>>(CUSPARSE_BSRMV_ARGTYPES(c10::complex<double>));
+
 } // namespace sparse
 } // namespace cuda
 } // namespace at
diff --git a/aten/src/ATen/native/sparse/cuda/SparseBlasImpl.cpp b/aten/src/ATen/native/sparse/cuda/SparseBlasImpl.cpp
@@ -109,6 +109,70 @@ void inline col_indices_and_values_resize_(const Tensor& input, int64_t nnz) {
       input.sizes());
 }
 
+void block_sparse_mv(
+    const at::sparse_csr::SparseCsrTensor& mat,
+    const Tensor& vec,
+    const Scalar& beta,
+    const Scalar& alpha,
+    const Tensor& result) {
+  TORCH_INTERNAL_ASSERT_DEBUG_ONLY(mat.is_sparse_csr());
+  // values is expected to be a blocks of sparse matrix
+  TORCH_INTERNAL_ASSERT_DEBUG_ONLY(mat.values().dim() == 3);
+  // blocks are expected to be square
+  TORCH_INTERNAL_ASSERT(mat.values().size(2) == mat.values().size(1));
+  // only block of size > 1 is supported in cuSPARSE
+  TORCH_INTERNAL_ASSERT(mat.values().size(-1) > 1);
+  // blocks are expected to be in row- or column-major order
+  TORCH_INTERNAL_ASSERT(
+      mat.values().is_contiguous() ||
+      mat.values().transpose(-2, -1).is_contiguous());
+
+  const cusparseDirection_t block_layout = mat.values().is_contiguous()
+      ? CUSPARSE_DIRECTION_ROW
+      : CUSPARSE_DIRECTION_COLUMN;
+
+  c10::MaybeOwned<Tensor> result_ = prepare_dense_vector_for_cusparse(result);
+  c10::MaybeOwned<Tensor> vec_ = prepare_dense_vector_for_cusparse(vec);
+
+  auto block_size = cuda_int_cast(mat.values().size(2), "block_size");
+  auto nnzb = cuda_int_cast(mat._nnz(), "nnzb");
+  auto mb = cuda_int_cast(mat.size(0), "mb") / block_size;
+  auto nb = cuda_int_cast(mat.size(1), "nb") / block_size;
+
+  AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES(
+      result.scalar_type(), "block_sparse_mv", [&] {
+        auto beta_ = beta.to<scalar_t>();
+        auto alpha_ = alpha.to<scalar_t>();
+        auto handle = at::cuda::getCurrentCUDASparseHandle();
+        auto desc = at::cuda::sparse::CuSparseMatDescriptor();
+        auto values = mat.values();
+        auto values_data_ptr = values.data_ptr<scalar_t>();
+        auto crow_indices = mat.crow_indices().to(kInt);
+        auto crow_indices_data_ptr = crow_indices.data_ptr<int>();
+        auto col_indices = mat.col_indices().to(kInt);
+        auto col_indices_data_ptr = col_indices.data_ptr<int>();
+        at::cuda::sparse::bsrmv(
+            handle,
+            block_layout,
+            CUSPARSE_OPERATION_NON_TRANSPOSE,
+            mb,
+            nb,
+            nnzb,
+            &alpha_,
+            desc.descriptor(),
+            values_data_ptr,
+            crow_indices_data_ptr,
+            col_indices_data_ptr,
+            block_size,
+            vec_->data_ptr<scalar_t>(),
+            &beta_,
+            result_->data_ptr<scalar_t>());
+      });
+  if (!result.is_same(*result_)) {
+    result.copy_(*result_);
+  }
+}
+
 void block_sparse_mm(
     const at::sparse_csr::SparseCsrTensor& mat1,
     const Tensor& mat2,
@@ -500,6 +564,9 @@ void addmv_out_sparse_csr(
     const Scalar& beta,
     const Scalar& alpha,
     const Tensor& result) {
+  if (mat.values().dim() == 3 && mat.values().size(-1) > 1) {
+    return block_sparse_mv(mat, vec, beta, alpha, result);
+  }
 #if !AT_USE_CUSPARSE_GENERIC_API()
   TORCH_CHECK(
       false,
diff --git a/test/test_sparse_csr.py b/test/test_sparse_csr.py
@@ -585,48 +585,61 @@ def test_csr_matvec(self, device, dtype):
             with self.assertRaisesRegex(RuntimeError, err_msg):
                 csr.matmul(bad_vec)
 
+    def run_test_block_addmm_addmv(self, addmv_addmm, c, a, b, op_b=False, op_out=False, *, dtype=None, device=None):
+        alpha = complex(random.random(), random.random()) if dtype.is_complex else random.random()
+        beta = complex(random.random(), random.random()) if dtype.is_complex else random.random()
+        b = b.mH if (op_b and a.shape == b.shape) else b
+
+        actual = addmv_addmm(c, a, b, alpha=alpha, beta=beta)
+
+        out = torch.empty_like(c.mH if op_out and a.shape == b.shape else c)
+        addmv_addmm(c, a, b, alpha=alpha, beta=beta, out=out)
+
+        a_bsr = sp.bsr_matrix(
+            (
+                a.values().cpu().numpy(),
+                a.col_indices().cpu().numpy(),
+                a.crow_indices().cpu().numpy(),
+            ),
+            shape=a.shape,
+        )
+        expected = alpha * (a_bsr * b.cpu().numpy()) + beta * c.cpu().numpy()
+        self.assertEqual(actual, out)
+        self.assertEqual(actual, expected)
+
     @onlyCUDA
     @unittest.skipIf(not TEST_SCIPY, "SciPy not found")
     @dtypes(torch.float32, torch.float64, torch.complex64, torch.complex128)
     def test_block_addmm(self, device, dtype):
-        def run_test(c, a, b, op_b, op_c, *, alpha=None, beta=None):
-            if dtype.is_complex:
-                alpha = random.random() + 0.3j if alpha is None else alpha
-                beta = random.random() + 0.6j if beta is None else beta
-            else:
-                alpha = random.random() if alpha is None else alpha
-                beta = random.random() if beta is None else beta
-
-            if op_b and a.shape == b.shape:
-                b = b.mH
-
-            actual = torch.addmm(c, a, b, alpha=alpha, beta=beta)
-
-            out = torch.empty_like(c if op_c and a.shape == b.shape else c.mH)
-            torch.addmm(c, a, b, alpha=alpha, beta=beta, out=out)
-
-            a_bsr = sp.bsr_matrix(
-                (
-                    a.values().cpu().numpy(),
-                    a.col_indices().cpu().numpy(),
-                    a.crow_indices().cpu().numpy(),
-                ),
-                shape=a.shape,
-            )
-            expected = alpha * (a_bsr * b.cpu().numpy()) + beta * c.cpu().numpy()
-            self.assertEqual(actual, out)
-            self.assertEqual(actual, expected)
-
         for index_dtype in [torch.int32, torch.int64]:
             for (m, n, k), block_size, noncontiguous in zip(itertools.product([1, 5], repeat=3), [1, 2, 3], [True, False]):
                 nnz = random.randint(0, m * k)
                 a = self.genSparseCSRTensor((m, k), nnz, dtype=dtype, device=device, index_dtype=index_dtype)
                 a_data = make_tensor((nnz, block_size, block_size), dtype=dtype, device=device)
+                a_data = a_data.mT if noncontiguous else a_data   # Test column-major blocks
                 a = torch._sparse_csr_tensor_unsafe(a.crow_indices(), a.col_indices(), a_data, (m * block_size, k * block_size))
                 b = make_tensor((k * block_size, n * block_size), dtype=dtype, device=device, noncontiguous=noncontiguous)
                 c = make_tensor((m * block_size, n * block_size), dtype=dtype, device=device, noncontiguous=noncontiguous)
-                for op_b, op_c in itertools.product([True, False], repeat=2):
-                    run_test(c, a, b, op_b, op_c)
+                for op_b, op_out in itertools.product([True, False], repeat=2):
+                    self.run_test_block_addmm_addmv(torch.addmm, c, a, b, op_b, op_out, dtype=dtype, device=device)
+
+    @onlyCUDA
+    @unittest.skipIf(not TEST_SCIPY, "SciPy not found")
+    @dtypes(torch.float32, torch.float64, torch.complex64, torch.complex128)
+    def test_block_addmv(self, device, dtype):
+        for index_dtype in [torch.int32, torch.int64]:
+            block_sizes = [1, 2, 3]
+            if TEST_WITH_ROCM or not TEST_CUSPARSE_GENERIC:
+                block_sizes = [2, 3]
+            for (m, k), block_size, noncontiguous in zip(itertools.product([1, 5], repeat=2), block_sizes, [True, False]):
+                nnz = random.randint(0, m * k)
+                a = self.genSparseCSRTensor((m, k), nnz, dtype=dtype, device=device, index_dtype=index_dtype)
+                a_data = make_tensor((nnz, block_size, block_size), dtype=dtype, device=device)
+                a_data = a_data.mT if noncontiguous else a_data  # Test column-major blocks
+                a = torch._sparse_csr_tensor_unsafe(a.crow_indices(), a.col_indices(), a_data, (m * block_size, k * block_size))
+                b = make_tensor((k * block_size,), dtype=dtype, device=device, noncontiguous=noncontiguous)
+                c = make_tensor((m * block_size,), dtype=dtype, device=device, noncontiguous=noncontiguous)
+                self.run_test_block_addmm_addmv(torch.addmv, c, a, b, dtype=dtype, device=device)
 
 
     @skipCPUIfNoMklSparse
