[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
taichi-dev · Oct 15, 2021 · ce6b4d3 · ce6b4d3
1 parent fb8e582
commit ce6b4d3
Show file tree

Hide file tree

Showing 9 changed files with 32 additions and 26 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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<float32> &matrix) {
   this->matrix_ = matrix;
 }

diff --git 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
@@ -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
@@ -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