[Lang] Use ti.linalg.sparse_matrix_builder as kernel parameters

examples/simulation/implicit_mass_spring.py

@@ -87,7 +87,7 @@ def init_edges(self):
             rest_len[idx] = (pos[idx1] - pos[idx2]).norm()
 
     @ti.kernel
-    def init_mass_sp(self, M: ti.sparse_matrix_builder()):
+    def init_mass_sp(self, M: ti.linalg.sparse_matrix_builder()):
         for i in range(self.NV):
             if self.invMass[i] != 0.0:
                 mass = 1.0 / self.invMass[i]
@@ -141,7 +141,7 @@ def compute_Jacobians(self):
         self.Jf[1] = ti.Matrix([[self.kf, 0], [0, self.kf]])
 
     @ti.kernel
-    def assemble_K(self, K: ti.sparse_matrix_builder()):
+    def assemble_K(self, K: ti.linalg.sparse_matrix_builder()):
         for i in self.spring:
             idx1, idx2 = self.spring[i][0], self.spring[i][1]
             for m, n in ti.static(ti.ndrange(2, 2)):
@@ -154,7 +154,7 @@ def assemble_K(self, K: ti.sparse_matrix_builder()):
             K[2 * (self.NV - 1) + m, 2 * (self.NV - 1) + n] += self.Jf[1][m, n]
 
     @ti.kernel
-    def assemble_D(self, D: ti.sparse_matrix_builder()):
+    def assemble_D(self, D: ti.linalg.sparse_matrix_builder()):
         for i in self.spring:
             idx1, idx2 = self.spring[i][0], self.spring[i][1]
             for m, n in ti.static(ti.ndrange(2, 2)):
@@ -188,7 +188,7 @@ def update(self, h):
         force = self.force.to_numpy().reshape(2 * self.NV)
         b = (force + h * K @ vel) * h
         # Sparse solver
-        solver = ti.SparseSolver(solver_type="LDLT")
+        solver = ti.linalg.SparseSolver(solver_type="LDLT")
         solver.analyze_pattern(A)
         solver.factorize(A)
         # Solve the linear system

examples/simulation/stable_fluid.py

@@ -69,7 +69,7 @@ def swap(self):
 if use_sparse_matrix:
     # use a sparse matrix to solve Poisson's pressure equation.
     @ti.kernel
-    def fill_laplacian_matrix(A: ti.sparse_matrix_builder()):
+    def fill_laplacian_matrix(A: ti.linalg.sparse_matrix_builder()):
         for i, j in ti.ndrange(res, res):
             row = i * res + j
             center = 0.0

misc/spMv_linear_solve.py

@@ -9,7 +9,8 @@
 
 
 @ti.kernel
-def fill(A: ti.sparse_matrix_builder(), b: ti.template(), interval: ti.i32):
+def fill(A: ti.linalg.sparse_matrix_builder(), b: ti.template(),
+         interval: ti.i32):
     for i in range(n):
         A[i, i] += 2.0
 

misc/sparse_matrix.py

@@ -9,8 +9,8 @@
 
 
 @ti.kernel
-def fill(A: ti.sparse_matrix_builder(), b: ti.sparse_matrix_builder(),
-         interval: ti.i32):
+def fill(A: ti.linalg.sparse_matrix_builder(),
+         b: ti.linalg.sparse_matrix_builder(), interval: ti.i32):
     for i in range(n):
         if i > 0:
             A[i - 1, i] += -1.0

python/taichi/lang/__init__.py

@@ -16,7 +16,6 @@
 from taichi.lang.enums import Layout
 from taichi.lang.exception import InvalidOperationError
 from taichi.lang.impl import *
-from taichi.lang.kernel_arguments import sparse_matrix_builder
 from taichi.lang.kernel_impl import (KernelArgError, KernelDefError,
                                      data_oriented, func, kernel, pyfunc)
 from taichi.lang.matrix import Matrix, Vector

python/taichi/lang/kernel_arguments.py

@@ -35,11 +35,6 @@ def subscript(self, i, j):
         return SparseMatrixEntry(self.ptr, i, j)
 
 
-sparse_matrix_builder = SparseMatrixBuilder
-"""Alias for :class:`~taichi.lang.sparse_matrix.SparseMatrixBuilder`.
-"""
-
-
 def decl_scalar_arg(dtype):
     dtype = cook_dtype(dtype)
     arg_id = _ti_core.decl_arg(dtype, False)

python/taichi/lang/kernel_impl.py

@@ -12,9 +12,9 @@
 from taichi.lang.ast.transformer import ASTTransformerTotal
 from taichi.lang.enums import Layout
 from taichi.lang.exception import TaichiSyntaxError
-from taichi.lang.kernel_arguments import sparse_matrix_builder
 from taichi.lang.shell import _shell_pop_print, oinspect
 from taichi.lang.util import to_taichi_type
+from taichi.linalg.sparse_matrix import sparse_matrix_builder
 from taichi.misc.util import obsolete
 from taichi.type import any_arr, primitive_types, template
 
@@ -527,7 +527,7 @@ def func__(*args):
                         raise KernelArgError(i, needed.to_string(), provided)
                     launch_ctx.set_arg_int(actual_argument_slot, int(v))
                 elif isinstance(needed, sparse_matrix_builder):
-                    # Pass only the base pointer of the ti.sparse_matrix_builder() argument
+                    # Pass only the base pointer of the ti.linalg.sparse_matrix_builder() argument
                     launch_ctx.set_arg_int(actual_argument_slot, v.get_addr())
                 elif isinstance(needed, any_arr) and (
                         self.match_ext_arr(v)

python/taichi/lang/stmt_builder.py

@@ -551,7 +551,7 @@ def transform_as_kernel():
                 if isinstance(ctx.func.argument_annotations[i], ti.template):
                     continue
                 if isinstance(ctx.func.argument_annotations[i],
-                              ti.sparse_matrix_builder):
+                              ti.linalg.sparse_matrix_builder):
                     arg_init = parse_stmt(
                         'x = ti.lang.kernel_arguments.decl_sparse_matrix()')
                     arg_init.targets[0].id = arg.arg

python/taichi/linalg/__init__.py

@@ -1,2 +1,3 @@
-from taichi.linalg.sparse_matrix import SparseMatrix, SparseMatrixBuilder
+from taichi.linalg.sparse_matrix import (SparseMatrix, SparseMatrixBuilder,
+                                         sparse_matrix_builder)
 from taichi.linalg.sparse_solver import SparseSolver

python/taichi/linalg/sparse_matrix.py

@@ -151,3 +151,7 @@ def build(self, dtype=f32, format='CSR'):
         """Create a sparse matrix using the triplets"""
         sm = self.ptr.build()
         return SparseMatrix(sm=sm)
+
+
+sparse_matrix_builder = SparseMatrixBuilder
+# Alias for :class:`SparseMatrixBuilder`

tests/python/test_sparse_linear_solver.py

@@ -37,8 +37,8 @@ def test_sparse_LLT_solver(solver_type):
     b = ti.field(ti.f32, shape=n)
 
     @ti.kernel
-    def fill(Abuilder: ti.sparse_matrix_builder(), InputArray: ti.ext_arr(),
-             b: ti.template()):
+    def fill(Abuilder: ti.linalg.sparse_matrix_builder(),
+             InputArray: ti.ext_arr(), b: ti.template()):
         for i, j in ti.ndrange(n, n):
             Abuilder[i, j] += InputArray[i, j]
         for i in range(n):

tests/python/test_sparse_matrix.py

@@ -7,7 +7,7 @@ def test_sparse_matrix_builder():
     Abuilder = ti.linalg.SparseMatrixBuilder(n, n, max_num_triplets=100)
 
     @ti.kernel
-    def fill(Abuilder: ti.sparse_matrix_builder()):
+    def fill(Abuilder: ti.linalg.sparse_matrix_builder()):
         for i, j in ti.ndrange(n, n):
             Abuilder[i, j] += i + j
 
@@ -24,7 +24,7 @@ def test_sparse_matrix_element_access():
     Abuilder = ti.linalg.SparseMatrixBuilder(n, n, max_num_triplets=100)
 
     @ti.kernel
-    def fill(Abuilder: ti.sparse_matrix_builder()):
+    def fill(Abuilder: ti.linalg.sparse_matrix_builder()):
         for i in range(n):
             Abuilder[i, i] += i
 
@@ -40,7 +40,7 @@ def test_sparse_matrix_element_modify():
     Abuilder = ti.linalg.SparseMatrixBuilder(n, n, max_num_triplets=100)
 
     @ti.kernel
-    def fill(Abuilder: ti.sparse_matrix_builder()):
+    def fill(Abuilder: ti.linalg.sparse_matrix_builder()):
         for i in range(n):
             Abuilder[i, i] += i
 
@@ -57,8 +57,8 @@ def test_sparse_matrix_addition():
     Bbuilder = ti.linalg.SparseMatrixBuilder(n, n, max_num_triplets=100)
 
     @ti.kernel
-    def fill(Abuilder: ti.sparse_matrix_builder(),
-             Bbuilder: ti.sparse_matrix_builder()):
+    def fill(Abuilder: ti.linalg.sparse_matrix_builder(),
+             Bbuilder: ti.linalg.sparse_matrix_builder()):
         for i, j in ti.ndrange(n, n):
             Abuilder[i, j] += i + j
             Bbuilder[i, j] += i - j
@@ -79,8 +79,8 @@ def test_sparse_matrix_subtraction():
     Bbuilder = ti.linalg.SparseMatrixBuilder(n, n, max_num_triplets=100)
 
     @ti.kernel
-    def fill(Abuilder: ti.sparse_matrix_builder(),
-             Bbuilder: ti.sparse_matrix_builder()):
+    def fill(Abuilder: ti.linalg.sparse_matrix_builder(),
+             Bbuilder: ti.linalg.sparse_matrix_builder()):
         for i, j in ti.ndrange(n, n):
             Abuilder[i, j] += i + j
             Bbuilder[i, j] += i - j
@@ -100,7 +100,7 @@ def test_sparse_matrix_scalar_multiplication():
     Abuilder = ti.linalg.SparseMatrixBuilder(n, n, max_num_triplets=100)
 
     @ti.kernel
-    def fill(Abuilder: ti.sparse_matrix_builder()):
+    def fill(Abuilder: ti.linalg.sparse_matrix_builder()):
         for i, j in ti.ndrange(n, n):
             Abuilder[i, j] += i + j
 
@@ -118,7 +118,7 @@ def test_sparse_matrix_transpose():
     Abuilder = ti.linalg.SparseMatrixBuilder(n, n, max_num_triplets=100)
 
     @ti.kernel
-    def fill(Abuilder: ti.sparse_matrix_builder()):
+    def fill(Abuilder: ti.linalg.sparse_matrix_builder()):
         for i, j in ti.ndrange(n, n):
             Abuilder[i, j] += i + j
 
@@ -137,8 +137,8 @@ def test_sparse_matrix_elementwise_multiplication():
     Bbuilder = ti.linalg.SparseMatrixBuilder(n, n, max_num_triplets=100)
 
     @ti.kernel
-    def fill(Abuilder: ti.sparse_matrix_builder(),
-             Bbuilder: ti.sparse_matrix_builder()):
+    def fill(Abuilder: ti.linalg.sparse_matrix_builder(),
+             Bbuilder: ti.linalg.sparse_matrix_builder()):
         for i, j in ti.ndrange(n, n):
             Abuilder[i, j] += i + j
             Bbuilder[i, j] += i - j
@@ -159,8 +159,8 @@ def test_sparse_matrix_multiplication():
     Bbuilder = ti.linalg.SparseMatrixBuilder(n, n, max_num_triplets=100)
 
     @ti.kernel
-    def fill(Abuilder: ti.sparse_matrix_builder(),
-             Bbuilder: ti.sparse_matrix_builder()):
+    def fill(Abuilder: ti.linalg.sparse_matrix_builder(),
+             Bbuilder: ti.linalg.sparse_matrix_builder()):
         for i, j in ti.ndrange(n, n):
             Abuilder[i, j] += i + j
             Bbuilder[i, j] += i - j
@@ -182,8 +182,8 @@ def test_sparse_matrix_nonsymmetric_multiplication():
     Bbuilder = ti.linalg.SparseMatrixBuilder(k, m, max_num_triplets=100)
 
     @ti.kernel
-    def fill(Abuilder: ti.sparse_matrix_builder(),
-             Bbuilder: ti.sparse_matrix_builder()):
+    def fill(Abuilder: ti.linalg.sparse_matrix_builder(),
+             Bbuilder: ti.linalg.sparse_matrix_builder()):
         for i, j in ti.ndrange(n, k):
             Abuilder[i, j] += i + j
         for i, j in ti.ndrange(k, m):

tests/python/test_spmv.py

@@ -8,7 +8,7 @@ def test_sparse_matrix_vector_multiplication1():
     b = ti.field(ti.f32, shape=n)
 
     @ti.kernel
-    def fill(Abuilder: ti.sparse_matrix_builder(), b: ti.template()):
+    def fill(Abuilder: ti.linalg.sparse_matrix_builder(), b: ti.template()):
         for i, j in ti.ndrange(n, n):
             Abuilder[i, j] += i
 
@@ -29,7 +29,7 @@ def test_sparse_matrix_vector_multiplication2():
     b = ti.field(ti.f32, shape=n)
 
     @ti.kernel
-    def fill(Abuilder: ti.sparse_matrix_builder(), b: ti.template()):
+    def fill(Abuilder: ti.linalg.sparse_matrix_builder(), b: ti.template()):
         for i, j in ti.ndrange(n, n):
             Abuilder[i, j] += i - j
 
@@ -53,7 +53,7 @@ def test_sparse_matrix_vector_multiplication3():
     b = ti.field(ti.f32, shape=n)
 
     @ti.kernel
-    def fill(Abuilder: ti.sparse_matrix_builder(), b: ti.template()):
+    def fill(Abuilder: ti.linalg.sparse_matrix_builder(), b: ti.template()):
         for i, j in ti.ndrange(n, n):
             Abuilder[i, j] += i + j