diff --git a/python/taichi/examples/simulation/implicit_mass_spring.py b/python/taichi/examples/simulation/implicit_mass_spring.py index a9cfc743d1907..293fd283b59b3 100644 --- a/python/taichi/examples/simulation/implicit_mass_spring.py +++ b/python/taichi/examples/simulation/implicit_mass_spring.py @@ -179,9 +179,29 @@ def update(self, h): A = self.M - h * D - h**2 * K - vel = self.vel.to_numpy().reshape(2 * self.NV) - force = self.force.to_numpy().reshape(2 * self.NV) - b = (force + h * K @ vel) * h + vel = ti.ndarray(ti.f32, 2 * self.NV) + force = ti.ndarray(ti.f32, 2 * self.NV) + + @ti.kernel + def copy_to(des: ti.types.ndarray(), source: ti.template()): + for i in range(self.NV): + des[2 * i] = source[i][0] + des[2 * i + 1] = source[i][1] + + copy_to(vel, self.vel) + copy_to(force, self.force) + + @ti.kernel + def compute_b(b: ti.types.ndarray(), f: ti.types.ndarray(), + Kv: ti.types.ndarray(), h: ti.f32): + for i in range(2 * self.NV): + b[i] = (f[i] + Kv[i] * h) * h + + # b = (force + h * K @ vel) * h + b = ti.ndarray(ti.f32, 2 * self.NV) + Kv = K @ vel + compute_b(b, force, Kv, h) + # Sparse solver solver = ti.linalg.SparseSolver(solver_type="LDLT") solver.analyze_pattern(A) diff --git a/taichi/program/sparse_matrix.cpp b/taichi/program/sparse_matrix.cpp index 5ed15f6b79618..718ecbf3b887e 100644 --- a/taichi/program/sparse_matrix.cpp +++ b/taichi/program/sparse_matrix.cpp @@ -25,6 +25,11 @@ } \ } +#define INSTANTIATE_SPMV(type, storage) \ + template void \ + EigenSparseMatrix>::spmv( \ + Program *prog, const Ndarray &x, const Ndarray &y); + namespace { using Pair = std::pair; struct key_hash { @@ -226,6 +231,31 @@ void EigenSparseMatrix::build_triplets(void *triplets_adr) { } } +template +void EigenSparseMatrix::spmv(Program *prog, + const Ndarray &x, + const Ndarray &y) { + size_t dX = prog->get_ndarray_data_ptr_as_int(&x); + size_t dY = prog->get_ndarray_data_ptr_as_int(&y); + std::string sdtype = taichi::lang::data_type_name(dtype_); + if (sdtype == "f32") { + Eigen::Map((float *)dY, cols_) = + matrix_.template cast() * + Eigen::Map((float *)dX, cols_); + } else if (sdtype == "f64") { + Eigen::Map((double *)dY, cols_) = + matrix_.template cast() * + Eigen::Map((double *)dX, cols_); + } else { + TI_ERROR("Unsupported sparse matrix data type {}!", sdtype); + } +} + +INSTANTIATE_SPMV(float32, ColMajor) +INSTANTIATE_SPMV(float32, RowMajor) +INSTANTIATE_SPMV(float64, ColMajor) +INSTANTIATE_SPMV(float64, RowMajor) + std::unique_ptr make_sparse_matrix( int rows, int cols, @@ -652,7 +682,7 @@ std::unique_ptr CuSparseMatrix::transpose() const { #endif } -void CuSparseMatrix::spmv(Program *prog, const Ndarray &x, Ndarray &y) { +void CuSparseMatrix::spmv(Program *prog, const Ndarray &x, const Ndarray &y) { #if defined(TI_WITH_CUDA) size_t dX = prog->get_ndarray_data_ptr_as_int(&x); size_t dY = prog->get_ndarray_data_ptr_as_int(&y); diff --git a/taichi/program/sparse_matrix.h b/taichi/program/sparse_matrix.h index 6031e402441cb..3708840e2ee9e 100644 --- a/taichi/program/sparse_matrix.h +++ b/taichi/program/sparse_matrix.h @@ -87,7 +87,7 @@ class SparseMatrix { return nullptr; } - inline DataType get_data_type() { + inline const DataType get_data_type() const { return dtype_; } @@ -197,6 +197,8 @@ class EigenSparseMatrix : public SparseMatrix { return matrix_ * b; } + void spmv(Program *prog, const Ndarray &x, const Ndarray &y); + private: EigenMatrix matrix_; }; @@ -265,7 +267,7 @@ class CuSparseMatrix : public SparseMatrix { void *coo_values_ptr, int nnz) override; - void spmv(Program *prog, const Ndarray &x, Ndarray &y); + void spmv(Program *prog, const Ndarray &x, const Ndarray &y); const void *get_matrix() const override { return &matrix_; diff --git a/taichi/program/sparse_solver.cpp b/taichi/program/sparse_solver.cpp index 1d3f5ba4ff211..5a55add667983 100644 --- a/taichi/program/sparse_solver.cpp +++ b/taichi/program/sparse_solver.cpp @@ -14,6 +14,14 @@ } \ } +#define INSTANTIATE_SOLVER(dt, type, order) \ + using dt##type##order = \ + Eigen::Simplicial##type, Eigen::Lower, \ + Eigen::order##Ordering>; \ + template void \ + EigenSparseSolver>::solve_rf( \ + Program *prog, const SparseMatrix &sm, const Ndarray &b, Ndarray &x); + using Triplets = std::tuple; namespace { struct key_hash { @@ -34,6 +42,9 @@ namespace taichi::lang { template bool EigenSparseSolver::compute( const SparseMatrix &sm) { + if (!is_initialized_) { + SparseSolver::init_solver(sm.num_rows(), sm.num_cols(), sm.get_data_type()); + } GET_EM(sm); solver_.compute(*mat); if (solver_.info() != Eigen::Success) { @@ -44,6 +55,9 @@ bool EigenSparseSolver::compute( template void EigenSparseSolver::analyze_pattern( const SparseMatrix &sm) { + if (!is_initialized_) { + SparseSolver::init_solver(sm.num_rows(), sm.num_cols(), sm.get_data_type()); + } GET_EM(sm); solver_.analyzePattern(*mat); } @@ -66,6 +80,23 @@ bool EigenSparseSolver::info() { return solver_.info() == Eigen::Success; } +template +void EigenSparseSolver::solve_rf( + Program *prog, + const SparseMatrix &sm, + const Ndarray &b, + Ndarray &x) { + size_t db = prog->get_ndarray_data_ptr_as_int(&b); + size_t dX = prog->get_ndarray_data_ptr_as_int(&x); + Eigen::Map((float *)dX, rows_) = + solver_.solve(Eigen::Map((float *)db, cols_)); +} + +INSTANTIATE_SOLVER(float32, LLT, COLAMD) +INSTANTIATE_SOLVER(float32, LDLT, COLAMD) +INSTANTIATE_SOLVER(float32, LLT, AMD) +INSTANTIATE_SOLVER(float32, LDLT, AMD) + CuSparseSolver::CuSparseSolver() { #if defined(TI_WITH_CUDA) if (!CUSPARSEDriver::get_instance().is_loaded()) { diff --git a/taichi/program/sparse_solver.h b/taichi/program/sparse_solver.h index 917fccd0f2917..876452cf93161 100644 --- a/taichi/program/sparse_solver.h +++ b/taichi/program/sparse_solver.h @@ -9,8 +9,19 @@ namespace taichi::lang { class SparseSolver { + protected: + int rows_{0}; + int cols_{0}; + DataType dtype_{PrimitiveType::f32}; + bool is_initialized_{false}; + public: virtual ~SparseSolver() = default; + void init_solver(const int rows, const int cols, const DataType dtype) { + rows_ = rows; + cols_ = cols; + dtype_ = dtype; + } virtual bool compute(const SparseMatrix &sm) = 0; virtual void analyze_pattern(const SparseMatrix &sm) = 0; virtual void factorize(const SparseMatrix &sm) = 0; @@ -46,9 +57,7 @@ class EigenSparseSolver : public SparseSolver { void solve_rf(Program *prog, const SparseMatrix &sm, const Ndarray &b, - Ndarray &x) override { - TI_NOT_IMPLEMENTED; - }; + Ndarray &x) override; bool info() override; }; diff --git a/taichi/python/export_lang.cpp b/taichi/python/export_lang.cpp index 4e131e2b30837..4c49ee6c0cb79 100644 --- a/taichi/python/export_lang.cpp +++ b/taichi/python/export_lang.cpp @@ -1219,6 +1219,7 @@ void export_lang(py::module &m) { .def(float##TYPE() * py::self) \ .def(py::self *py::self) \ .def("matmul", &EigenSparseMatrix::matmul) \ + .def("spmv", &EigenSparseMatrix::spmv) \ .def("transpose", \ &EigenSparseMatrix::transpose) \ .def("get_element", \ diff --git a/tests/python/test_sparse_linear_solver.py b/tests/python/test_sparse_linear_solver.py index 772cfc2382d75..b0e79b66d3bd1 100644 --- a/tests/python/test_sparse_linear_solver.py +++ b/tests/python/test_sparse_linear_solver.py @@ -38,6 +38,39 @@ def fill(Abuilder: ti.types.sparse_matrix_builder(), assert x[i] == test_utils.approx(res[i], rel=1.0) +@pytest.mark.parametrize("dtype", [ti.f32]) +@pytest.mark.parametrize("solver_type", ["LLT", "LDLT", "LU"]) +@pytest.mark.parametrize("ordering", ["AMD", "COLAMD"]) +@test_utils.test(arch=ti.cpu) +def test_sparse_solver_ndarray_vector(dtype, solver_type, ordering): + n = 10 + A = np.random.rand(n, n) + A_psd = np.dot(A, A.transpose()) + Abuilder = ti.linalg.SparseMatrixBuilder(n, n, max_num_triplets=300) + b = ti.ndarray(ti.f32, shape=n) + + @ti.kernel + def fill(Abuilder: ti.types.sparse_matrix_builder(), + InputArray: ti.types.ndarray(), b: ti.types.ndarray()): + for i, j in ti.ndrange(n, n): + Abuilder[i, j] += InputArray[i, j] + for i in range(n): + b[i] = i + 1 + + fill(Abuilder, A_psd, b) + A = Abuilder.build() + solver = ti.linalg.SparseSolver(dtype=dtype, + solver_type=solver_type, + ordering=ordering) + solver.analyze_pattern(A) + solver.factorize(A) + x = solver.solve(b) + + res = np.linalg.solve(A_psd, b.to_numpy()) + for i in range(n): + assert x[i] == test_utils.approx(res[i], rel=1.0) + + @test_utils.test(arch=ti.cuda) def test_gpu_sparse_solver(): from scipy.sparse import coo_matrix diff --git a/tests/python/test_sparse_matrix.py b/tests/python/test_sparse_matrix.py index ce90d4616e15a..bcf0ed72311fc 100644 --- a/tests/python/test_sparse_matrix.py +++ b/tests/python/test_sparse_matrix.py @@ -376,6 +376,34 @@ def fill(Abuilder: ti.types.sparse_matrix_builder(), assert C[i, j] == GT[i][j] +@pytest.mark.parametrize('dtype, storage_format', [(ti.f32, 'col_major'), + (ti.f32, 'row_major'), + (ti.f64, 'col_major'), + (ti.f64, 'row_major')]) +@test_utils.test(arch=ti.cpu) +def test_sparse_matrix_ndarray_vector_multiplication(dtype, storage_format): + n = 2 + Abuilder = ti.linalg.SparseMatrixBuilder(n, + n, + max_num_triplets=100, + dtype=dtype, + storage_format=storage_format) + x = ti.ndarray(dtype, n) + + @ti.kernel + def fill(Abuilder: ti.types.sparse_matrix_builder()): + for i, j in ti.ndrange(n, n): + Abuilder[i, j] += i + j + + fill(Abuilder) + x.fill(1.0) + A = Abuilder.build() + res = A @ x + res_n = res.to_numpy() + assert res_n[0] == 1.0 + assert res_n[1] == 3.0 + + @test_utils.test(arch=ti.cuda) def test_gpu_sparse_matrix(): import numpy as np