From ce6b4d3a05dd1216668d20aaaf71b47823503757 Mon Sep 17 00:00:00 2001 From: FantasyVR <6712304+FantasyVR@users.noreply.github.com> Date: Fri, 15 Oct 2021 19:17:06 +0800 Subject: [PATCH] [Lang] Reuse sparse matrix builder (#3199) * [Lang] Reuse sparse matrix builder * [Lang] Use ti.linalg.SparseMatrixBuilder in code * remove SparseMatrixBuilder in taichi/lang/__init__.py --- examples/simulation/implicit_mass_spring.py | 29 ++++++++++----------- examples/simulation/stable_fluid.py | 4 +-- misc/spMv_linear_solve.py | 4 +-- misc/sparse_matrix.py | 4 +-- python/taichi/lang/__init__.py | 1 - taichi/program/sparse_matrix.cpp | 6 +++++ taichi/program/sparse_matrix.h | 2 ++ tests/python/test_sparse_linear_solver.py | 2 +- tests/python/test_spmv.py | 6 ++--- 9 files changed, 32 insertions(+), 26 deletions(-) diff --git a/examples/simulation/implicit_mass_spring.py b/examples/simulation/implicit_mass_spring.py index 7bd284afdd0d8..61ceebda982dd 100644 --- a/examples/simulation/implicit_mass_spring.py +++ b/examples/simulation/implicit_mass_spring.py @@ -33,10 +33,14 @@ def __init__(self, N): self.gravity = ti.Vector([0.0, -2.0]) self.init_pos() self.init_edges() - self.MassBuilder = ti.SparseMatrixBuilder(2 * self.NV, - 2 * self.NV, - max_num_triplets=10000) - + self.MassBuilder = ti.linalg.SparseMatrixBuilder( + 2 * self.NV, 2 * self.NV, max_num_triplets=10000) + self.DBuilder = ti.linalg.SparseMatrixBuilder(2 * self.NV, + 2 * self.NV, + max_num_triplets=10000) + self.KBuilder = ti.linalg.SparseMatrixBuilder(2 * self.NV, + 2 * self.NV, + max_num_triplets=10000) self.init_mass_sp(self.MassBuilder) self.M = self.MassBuilder.build() self.fix_vertex = [self.N, self.NV - 1] @@ -171,17 +175,12 @@ def update(self, h): self.compute_Jacobians() # Assemble global system - DBuilder = ti.SparseMatrixBuilder(2 * self.NV, - 2 * self.NV, - max_num_triplets=10000) - self.assemble_D(DBuilder) - D = DBuilder.build() - - KBuilder = ti.SparseMatrixBuilder(2 * self.NV, - 2 * self.NV, - max_num_triplets=10000) - self.assemble_K(KBuilder) - K = KBuilder.build() + + self.assemble_D(self.DBuilder) + D = self.DBuilder.build() + + self.assemble_K(self.KBuilder) + K = self.KBuilder.build() A = self.M - h * D - h**2 * K diff --git a/examples/simulation/stable_fluid.py b/examples/simulation/stable_fluid.py index 7f91af9e1c405..4d50e6f04df70 100644 --- a/examples/simulation/stable_fluid.py +++ b/examples/simulation/stable_fluid.py @@ -88,12 +88,12 @@ def fill_laplacian_matrix(A: ti.sparse_matrix_builder()): A[row, row] += center N = res * res - K = ti.SparseMatrixBuilder(N, N, max_num_triplets=N * 6) + K = ti.linalg.SparseMatrixBuilder(N, N, max_num_triplets=N * 6) b = ti.field(ti.f32, shape=N) fill_laplacian_matrix(K) L = K.build() - solver = ti.SparseSolver(solver_type="LLT") + solver = ti.linalg.SparseSolver(solver_type="LLT") solver.analyze_pattern(L) solver.factorize(L) diff --git a/misc/spMv_linear_solve.py b/misc/spMv_linear_solve.py index 9e678b6bfc77b..05dc29682dcd8 100644 --- a/misc/spMv_linear_solve.py +++ b/misc/spMv_linear_solve.py @@ -4,7 +4,7 @@ n = 8 -K = ti.SparseMatrixBuilder(n, n, max_num_triplets=100) +K = ti.linalg.SparseMatrixBuilder(n, n, max_num_triplets=100) b = ti.field(ti.f32, shape=n) @@ -30,7 +30,7 @@ def fill(A: ti.sparse_matrix_builder(), b: ti.template(), interval: ti.i32): print(x) print("Solving sparse linear systems Ax = b with the solution x:") -solver = ti.SparseSolver(solver_type="LLT") +solver = ti.linalg.SparseSolver(solver_type="LLT") solver.analyze_pattern(A) solver.factorize(A) x = solver.solve(b) diff --git a/misc/sparse_matrix.py b/misc/sparse_matrix.py index 5d800f27dcf37..be038a2338197 100644 --- a/misc/sparse_matrix.py +++ b/misc/sparse_matrix.py @@ -4,8 +4,8 @@ n = 8 -K = ti.SparseMatrixBuilder(n, n, max_num_triplets=100) -f = ti.SparseMatrixBuilder(n, 1, max_num_triplets=100) +K = ti.linalg.SparseMatrixBuilder(n, n, max_num_triplets=100) +f = ti.linalg.SparseMatrixBuilder(n, 1, max_num_triplets=100) @ti.kernel diff --git a/python/taichi/lang/__init__.py b/python/taichi/lang/__init__.py index 5dfda5c45b969..5e247ef9abfea 100644 --- a/python/taichi/lang/__init__.py +++ b/python/taichi/lang/__init__.py @@ -29,7 +29,6 @@ from taichi.lang.util import (has_pytorch, is_taichi_class, python_scope, taichi_scope, to_numpy_type, to_pytorch_type, to_taichi_type) -from taichi.linalg import SparseMatrix, SparseMatrixBuilder, SparseSolver from taichi.misc.util import deprecated from taichi.profiler import KernelProfiler, get_default_kernel_profiler from taichi.profiler.kernelmetrics import (CuptiMetric, default_cupti_metrics, diff --git a/taichi/program/sparse_matrix.cpp b/taichi/program/sparse_matrix.cpp index 1bc356ac316c4..a42ab8d205f71 100644 --- a/taichi/program/sparse_matrix.cpp +++ b/taichi/program/sparse_matrix.cpp @@ -41,9 +41,15 @@ SparseMatrix SparseMatrixBuilder::build() { } SparseMatrix sm(rows_, cols_); sm.get_matrix().setFromTriplets(triplets.begin(), triplets.end()); + clear(); return sm; } +void SparseMatrixBuilder::clear() { + built_ = false; + num_triplets_ = 0; +} + SparseMatrix::SparseMatrix(Eigen::SparseMatrix &matrix) { this->matrix_ = matrix; } diff --git a/taichi/program/sparse_matrix.h b/taichi/program/sparse_matrix.h index d22f146c306ca..2de867d1ce249 100644 --- a/taichi/program/sparse_matrix.h +++ b/taichi/program/sparse_matrix.h @@ -19,6 +19,8 @@ class SparseMatrixBuilder { SparseMatrix build(); + void clear(); + private: uint64 num_triplets_{0}; void *data_base_ptr_{nullptr}; diff --git a/tests/python/test_sparse_linear_solver.py b/tests/python/test_sparse_linear_solver.py index 7b1d2b2aae6a0..b2f2d6cb837c9 100644 --- a/tests/python/test_sparse_linear_solver.py +++ b/tests/python/test_sparse_linear_solver.py @@ -46,7 +46,7 @@ def fill(Abuilder: ti.sparse_matrix_builder(), InputArray: ti.ext_arr(), fill(Abuilder, Aarray, b) A = Abuilder.build() - solver = ti.SparseSolver(solver_type=solver_type) + solver = ti.linalg.SparseSolver(solver_type=solver_type) solver.analyze_pattern(A) solver.factorize(A) x = solver.solve(b) diff --git a/tests/python/test_spmv.py b/tests/python/test_spmv.py index 90d28b38a1d83..5847eee37592d 100644 --- a/tests/python/test_spmv.py +++ b/tests/python/test_spmv.py @@ -4,7 +4,7 @@ @ti.test(arch=ti.cpu) def test_sparse_matrix_vector_multiplication1(): n = 8 - Abuilder = ti.SparseMatrixBuilder(n, n, max_num_triplets=100) + Abuilder = ti.linalg.SparseMatrixBuilder(n, n, max_num_triplets=100) b = ti.field(ti.f32, shape=n) @ti.kernel @@ -25,7 +25,7 @@ def fill(Abuilder: ti.sparse_matrix_builder(), b: ti.template()): @ti.test(arch=ti.cpu) def test_sparse_matrix_vector_multiplication2(): n = 8 - Abuilder = ti.SparseMatrixBuilder(n, n, max_num_triplets=100) + Abuilder = ti.linalg.SparseMatrixBuilder(n, n, max_num_triplets=100) b = ti.field(ti.f32, shape=n) @ti.kernel @@ -49,7 +49,7 @@ def fill(Abuilder: ti.sparse_matrix_builder(), b: ti.template()): @ti.test(arch=ti.cpu) def test_sparse_matrix_vector_multiplication3(): n = 8 - Abuilder = ti.SparseMatrixBuilder(n, n, max_num_triplets=100) + Abuilder = ti.linalg.SparseMatrixBuilder(n, n, max_num_triplets=100) b = ti.field(ti.f32, shape=n) @ti.kernel