From 8bdb24b9e5512f49f39b2cd60ae0b7d6dd33cd14 Mon Sep 17 00:00:00 2001 From: FantasyVR Date: Fri, 11 Nov 2022 00:07:47 +0800 Subject: [PATCH] unify spmv api --- misc/test_build_cusm_from_coo.py | 2 +- misc/test_sm.py | 36 --------------------------- python/taichi/linalg/sparse_matrix.py | 34 ++++++------------------- tests/python/test_sparse_matrix.py | 2 +- 4 files changed, 10 insertions(+), 64 deletions(-) delete mode 100644 misc/test_sm.py diff --git a/misc/test_build_cusm_from_coo.py b/misc/test_build_cusm_from_coo.py index 17b46a4384f93..620b95232f7d2 100644 --- a/misc/test_build_cusm_from_coo.py +++ b/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 diff --git a/misc/test_sm.py b/misc/test_sm.py deleted file mode 100644 index 626845db7a80f..0000000000000 --- a/misc/test_sm.py +++ /dev/null @@ -1,36 +0,0 @@ -import numpy as np - -import taichi as ti - -ti.init(arch=ti.cuda, debug=True, offline_cache=False) - -h_coo_row = np.asarray([1, 0, 0, 0, 2, 2, 2, 3, 3], dtype=np.int32) -h_coo_col = np.asarray([1, 0, 2, 3, 0, 2, 3, 1, 3], dtype=np.int32) -h_coo_val = np.asarray([4.0, 1.0, 2.0, 3.0, 5.0, 6.0, 7.0, 8.0, 9.0], - dtype=np.float32) -h_X = np.asarray([1.0, 2.0, 3.0, 4.0], dtype=np.float32) -h_Y = np.asarray([19.0, 8.0, 51.0, 52.0], dtype=np.float32) -# Data structure for building the CSR matrix A using Taichi Sparse Matrix -idx_dt = ti.int32 -val_dt = ti.f32 -d_coo_row = ti.ndarray(shape=9, dtype=idx_dt) -d_coo_col = ti.ndarray(shape=9, dtype=idx_dt) -d_coo_val = ti.ndarray(shape=9, dtype=val_dt) -# Dense vector x -X = ti.ndarray(shape=4, dtype=val_dt) -# Results for A @ x -# Y = ti.ndarray(shape=4, dtype=val_dt) -# Initialize the CSR matrix and vectors with numpy array -d_coo_row.from_numpy(h_coo_row) -d_coo_col.from_numpy(h_coo_col) -d_coo_val.from_numpy(h_coo_val) -X.from_numpy(h_X) -# Y.fill(0.0) -# Define the CSR matrix A -A = ti.linalg.SparseMatrix(n=4, m=4, dtype=ti.f32) -# Build the CSR matrix A with Taichi ndarray -A.build_coo(d_coo_row, d_coo_col, d_coo_val) -# Compute Y = A @ X -Y = A.spmv(X) -for i in range(4): - assert Y[i] == h_Y[i] diff --git a/python/taichi/linalg/sparse_matrix.py b/python/taichi/linalg/sparse_matrix.py index bfb0d92ffaadb..034d3805954a1 100644 --- a/python/taichi/linalg/sparse_matrix.py +++ b/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. diff --git a/tests/python/test_sparse_matrix.py b/tests/python/test_sparse_matrix.py index 57f5e1a44a0ff..e09fd850c7375 100644 --- a/tests/python/test_sparse_matrix.py +++ b/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]