unify spmv api

misc/test_build_cusm_from_coo.py

@@ -26,7 +26,7 @@
 A = ti.linalg.SparseMatrix(n=4, m=4, dtype=ti.float32)
 A.build_coo(d_coo_row, d_coo_col, d_coo_val)
 
-A.spmv(x, y)
+y = A @ x
 
 # Check if the results are correct
 equal = True

python/taichi/linalg/sparse_matrix.py

@@ -136,6 +136,14 @@ def __matmul__(self, other):
             assert self.m == other.shape[
                 0], f"Dimension mismatch between sparse matrix ({self.n}, {self.m}) and vector ({other.shape})"
             return self.matrix.mat_vec_mul(other)
+        if isinstance(other, Ndarray):
+            if self.m != other.shape[0]:
+                raise TaichiRuntimeError(
+                    f"Dimension mismatch between sparse matrix ({self.n}, {self.m}) and vector ({other.shape})"
+                )
+            res = ScalarNdarray(dtype=other.dtype, arr_shape=(self.n, ))
+            self.matrix.spmv(get_runtime().prog, other.arr, res.arr)
+            return res
         raise TaichiRuntimeError(
             f"Sparse matrix-matrix/vector multiplication does not support {type(other)} for now. Supported types are SparseMatrix, ti.field, and numpy ndarray."
         )
@@ -222,32 +230,6 @@ def build_coo(self, row_coo, col_coo, value_coo):
             get_runtime().prog.make_sparse_matrix_from_ndarray_cusparse(
                 self.matrix, row_coo.arr, col_coo.arr, value_coo.arr)
 
-    def spmv(self, x):
-        """Sparse matrix-vector multiplication using cuSparse.
-
-        Args:
-            x (ti.ndarray): the vector to be multiplied.
-            y (ti.ndarray): the result of matrix-vector multiplication.
-
-        Example::
-            >>> x = ti.ndarray(shape=4, dtype=val_dt)
-            >>> y = ti.ndarray(shape=4, dtype=val_dt)
-            >>> A = ti.linalg.SparseMatrix(n=4, m=4, dtype=ti.f32)
-            >>> A.build_from_ndarray_cusparse(row_csr, col_csr, value_csr)
-            >>> A.spmv(x, y)
-        """
-        if not isinstance(x, Ndarray):
-            raise TaichiRuntimeError(
-                'Sparse matrix only supports building from [ti.ndarray, ti.Vector.ndarray, ti.Matrix.ndarray]'
-            )
-        if self.m != x.shape[0]:
-            raise TaichiRuntimeError(
-                f"Dimension mismatch between sparse matrix ({self.n}, {self.m}) and vector ({x.shape})"
-            )
-        res = ScalarNdarray(dtype=x.dtype, arr_shape=(self.n, ))
-        self.matrix.spmv(get_runtime().prog, x.arr, res.arr)
-        return res
-
 
 class SparseMatrixBuilder:
     """A python wrap around sparse matrix builder.

tests/python/test_sparse_matrix.py

@@ -411,7 +411,7 @@ def test_gpu_sparse_matrix():
     A.build_coo(d_coo_row, d_coo_col, d_coo_val)
 
     # Compute Y = A @ X
-    A.spmv(X, Y)
+    Y = A @ X
     for i in range(4):
         assert Y[i] == h_Y[i]