From e371291c75c8d8caf85fed118b2586fdf842ea8d Mon Sep 17 00:00:00 2001 From: FantasyVR Date: Thu, 10 Nov 2022 23:53:17 +0800 Subject: [PATCH 1/2] spmv return ndarray --- misc/test_sm.py | 36 +++++++++++++++++++++++++++ python/taichi/linalg/sparse_matrix.py | 11 ++++---- 2 files changed, 42 insertions(+), 5 deletions(-) create mode 100644 misc/test_sm.py diff --git a/misc/test_sm.py b/misc/test_sm.py new file mode 100644 index 0000000000000..626845db7a80f --- /dev/null +++ b/misc/test_sm.py @@ -0,0 +1,36 @@ +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 eac2896644af6..bfb0d92ffaadb 100644 --- a/python/taichi/linalg/sparse_matrix.py +++ b/python/taichi/linalg/sparse_matrix.py @@ -1,10 +1,10 @@ from functools import reduce import numpy as np +from taichi.lang._ndarray import Ndarray, ScalarNdarray from taichi.lang.exception import TaichiRuntimeError from taichi.lang.field import Field from taichi.lang.impl import get_runtime -from taichi.lang.matrix import Ndarray from taichi.lang.util import warning from taichi.types import annotations, f32, i32 @@ -222,7 +222,7 @@ 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, y): + def spmv(self, x): """Sparse matrix-vector multiplication using cuSparse. Args: @@ -236,7 +236,7 @@ def spmv(self, x, y): >>> A.build_from_ndarray_cusparse(row_csr, col_csr, value_csr) >>> A.spmv(x, y) """ - if not isinstance(x, Ndarray) or not isinstance(y, Ndarray): + if not isinstance(x, Ndarray): raise TaichiRuntimeError( 'Sparse matrix only supports building from [ti.ndarray, ti.Vector.ndarray, ti.Matrix.ndarray]' ) @@ -244,8 +244,9 @@ def spmv(self, x, y): raise TaichiRuntimeError( f"Dimension mismatch between sparse matrix ({self.n}, {self.m}) and vector ({x.shape})" ) - - self.matrix.spmv(get_runtime().prog, x.arr, y.arr) + res = ScalarNdarray(dtype=x.dtype, arr_shape=(self.n, )) + self.matrix.spmv(get_runtime().prog, x.arr, res.arr) + return res class SparseMatrixBuilder: From 8bdb24b9e5512f49f39b2cd60ae0b7d6dd33cd14 Mon Sep 17 00:00:00 2001 From: FantasyVR Date: Fri, 11 Nov 2022 00:07:47 +0800 Subject: [PATCH 2/2] 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]