[Lang] Experimental sparse matrix support on CPUs (taichi-dev#2792)

Co-authored-by: Yuanming Hu <[email protected]>
Co-authored-by: Ye Kuang <[email protected]>

CMakeLists.txt

@@ -60,6 +60,8 @@ if (TI_BUILD_TESTS)
   include(cmake/TaichiTests.cmake)
 endif()
 
+include_directories(${PROJECT_SOURCE_DIR}/external/eigen)
+
 message("C++ Flags: ${CMAKE_CXX_FLAGS}")
 message("Build type: ${CMAKE_BUILD_TYPE}")
 

misc/sparse_matrix.py

@@ -0,0 +1,64 @@
+import taichi as ti
+
+ti.init(arch=ti.x64)
+
+n = 8
+
+K = ti.SparseMatrixBuilder(n, n, max_num_triplets=100)
+f = ti.SparseMatrixBuilder(n, 1, max_num_triplets=100)
+
+
+@ti.kernel
+def fill(A: ti.sparse_matrix_builder(), b: ti.sparse_matrix_builder(),
+         interval: ti.i32):
+    for i in range(n):
+        if i > 0:
+            A[i - 1, i] += -1.0
+            A[i, i] += 1
+        if i < n - 1:
+            A[i + 1, i] += -1.0
+            A[i, i] += 1.0
+
+        if i % interval == 0:
+            b[i, 0] += 1.0
+
+
+fill(K, f, 3)
+
+print(">>>> K.print_triplets()")
+K.print_triplets()
+
+A = K.build()
+
+print(">>>> A = K.build()")
+print(A)
+
+print(">>>> Summation: C = A + A")
+C = A + A
+print(C)
+
+print(">>>> Subtraction: D = A - A")
+D = A - A
+print(D)
+
+print(">>>> Multiplication with a scalar on the right: E = A * 3.0")
+E = A * 3.0
+print(E)
+
+print(">>>> Multiplication with a scalar on the left: E = 3.0 * A")
+E = 3.0 * A
+print(E)
+
+print(">>>> Transpose: F = A.transpose()")
+F = A.transpose()
+print(F)
+
+print(">>>> Matrix multiplication: G = E @ A")
+G = E @ A
+print(G)
+
+print(">>>> Element-wise multiplication: H = E * A")
+H = E * A
+print(H)
+
+print(f">>>> Element Access: A[0,0] = {A[0,0]}")

python/taichi/lang/__init__.py

@@ -8,14 +8,16 @@
 from taichi.lang.enums import Layout
 from taichi.lang.exception import InvalidOperationError
 from taichi.lang.impl import *
-from taichi.lang.kernel_arguments import any_arr, ext_arr, template
+from taichi.lang.kernel_arguments import (any_arr, ext_arr,
+                                          sparse_matrix_builder, template)
 from taichi.lang.kernel_impl import (KernelArgError, KernelDefError,
                                      data_oriented, func, kernel, pyfunc)
 from taichi.lang.matrix import Matrix, Vector
 from taichi.lang.ndrange import GroupedNDRange, ndrange
 from taichi.lang.ops import *
 from taichi.lang.quant_impl import quant
 from taichi.lang.runtime_ops import async_flush, sync
+from taichi.lang.sparse_matrix import SparseMatrix, SparseMatrixBuilder
 from taichi.lang.transformer import TaichiSyntaxError
 from taichi.lang.type_factory_impl import type_factory
 from taichi.lang.util import (has_pytorch, is_taichi_class, python_scope,

python/taichi/lang/kernel_arguments.py

@@ -1,8 +1,10 @@
+from taichi.core.primitive_types import u64
 from taichi.core.util import ti_core as _ti_core
 from taichi.lang.enums import Layout
 from taichi.lang.expr import Expr
 from taichi.lang.ext_array import AnyArray, ExtArray
 from taichi.lang.snode import SNode
+from taichi.lang.sparse_matrix import SparseMatrixBuilder
 from taichi.lang.util import cook_dtype, to_taichi_type
 
 
@@ -117,6 +119,39 @@ def extract(self, x):
 """
 
 
+class SparseMatrixEntry:
+    def __init__(self, ptr, i, j):
+        self.ptr = ptr
+        self.i = i
+        self.j = j
+
+    def augassign(self, value, op):
+        from taichi.lang.impl import call_internal, ti_float
+        if op == 'Add':
+            call_internal("insert_triplet", self.ptr, self.i, self.j,
+                          ti_float(value))
+        elif op == 'Sub':
+            call_internal("insert_triplet", self.ptr, self.i, self.j,
+                          -ti_float(value))
+        else:
+            assert False, f"Only operations '+=' and '-=' are supported on sparse matrices."
+
+
+class SparseMatrixProxy:
+    is_taichi_class = True
+
+    def __init__(self, ptr):
+        self.ptr = ptr
+
+    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)
@@ -129,6 +164,13 @@ def decl_ext_arr_arg(dtype, dim):
     return ExtArray(_ti_core.make_external_tensor_expr(dtype, dim, arg_id))
 
 
+def decl_sparse_matrix():
+    ptr_type = cook_dtype(u64)
+    # Treat the sparse matrix argument as a scalar since we only need to pass in the base pointer
+    arg_id = _ti_core.decl_arg(ptr_type, False)
+    return SparseMatrixProxy(_ti_core.make_arg_load_expr(arg_id, ptr_type))
+
+
 def decl_any_arr_arg(dtype, dim, element_shape, layout):
     dtype = cook_dtype(dtype)
     arg_id = _ti_core.decl_arg(dtype, True)

python/taichi/lang/kernel_impl.py

@@ -10,7 +10,8 @@
 from taichi.lang import impl, util
 from taichi.lang.ast_checker import KernelSimplicityASTChecker
 from taichi.lang.exception import TaichiSyntaxError
-from taichi.lang.kernel_arguments import any_arr, ext_arr, template
+from taichi.lang.kernel_arguments import (any_arr, ext_arr,
+                                          sparse_matrix_builder, template)
 from taichi.lang.ndarray import ScalarNdarray
 from taichi.lang.shell import _shell_pop_print, oinspect
 from taichi.lang.transformer import ASTTransformerTotal
@@ -368,6 +369,8 @@ def extract_arguments(self):
                     pass
                 elif id(annotation) in primitive_types.type_ids:
                     pass
+                elif isinstance(annotation, sparse_matrix_builder):
+                    pass
                 else:
                     _taichi_skip_traceback = 1
                     raise KernelDefError(
@@ -483,6 +486,9 @@ def func__(*args):
                     if not isinstance(v, int):
                         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
+                    launch_ctx.set_arg_int(actual_argument_slot, v.get_addr())
                 elif (isinstance(needed, (any_arr, ext_arr)) and self.match_ext_arr(v)) or \
                      (isinstance(needed, any_arr) and isinstance(v, ScalarNdarray)):
                     if isinstance(v, ScalarNdarray):

python/taichi/lang/sparse_matrix.py

@@ -0,0 +1,73 @@
+class SparseMatrix:
+    def __init__(self, n=None, m=None, sm=None):
+        if sm is None:
+            self.n = n
+            self.m = m if m else n
+            from taichi.core.util import ti_core as _ti_core
+            self.matrix = _ti_core.create_sparse_matrix(n, m)
+        else:
+            self.n = sm.num_rows()
+            self.m = sm.num_cols()
+            self.matrix = sm
+
+    def __add__(self, other):
+        assert self.n == other.n and self.m == other.m, f"Dimension mismatch between sparse matrices ({self.n}, {self.m}) and ({other.n}, {other.m})"
+        sm = self.matrix + other.matrix
+        return SparseMatrix(sm=sm)
+
+    def __sub__(self, other):
+        assert self.n == other.n and self.m == other.m, f"Dimension mismatch between sparse matrices ({self.n}, {self.m}) and ({other.n}, {other.m})"
+        sm = self.matrix - other.matrix
+        return SparseMatrix(sm=sm)
+
+    def __mul__(self, other):
+        if isinstance(other, float):
+            sm = self.matrix * other
+            return SparseMatrix(sm=sm)
+        elif isinstance(other, SparseMatrix):
+            assert self.n == other.n and self.m == other.m, f"Dimension mismatch between sparse matrices ({self.n}, {self.m}) and ({other.n}, {other.m})"
+            sm = self.matrix * other.matrix
+            return SparseMatrix(sm=sm)
+
+    def __rmul__(self, other):
+        if isinstance(other, float):
+            sm = other * self.matrix
+            return SparseMatrix(sm=sm)
+
+    def transpose(self):
+        sm = self.matrix.transpose()
+        return SparseMatrix(sm=sm)
+
+    def __matmul__(self, other):
+        assert self.m == other.n, f"Dimension mismatch between sparse matrices ({self.n}, {self.m}) and ({other.n}, {other.m})"
+        sm = self.matrix.matmul(other.matrix)
+        return SparseMatrix(sm=sm)
+
+    def __getitem__(self, indices):
+        return self.matrix.get_element(indices[0], indices[1])
+
+    def __str__(self):
+        return self.matrix.to_string()
+
+    def __repr__(self):
+        return self.matrix.to_string()
+
+
+class SparseMatrixBuilder:
+    def __init__(self, num_rows=None, num_cols=None, max_num_triplets=0):
+        self.num_rows = num_rows
+        self.num_cols = num_cols if num_cols else num_rows
+        if num_rows is not None:
+            from taichi.core.util import ti_core as _ti_core
+            self.ptr = _ti_core.create_sparse_matrix_builder(
+                num_rows, num_cols, max_num_triplets)
+
+    def get_addr(self):
+        return self.ptr.get_addr()
+
+    def print_triplets(self):
+        self.ptr.print_triplets()
+
+    def build(self):
+        sm = self.ptr.build()
+        return SparseMatrix(sm=sm)

python/taichi/lang/stmt_builder.py

@@ -556,6 +556,13 @@ def transform_as_kernel():
                     arg_init.value.args[0] = dt
                     arg_init.value.args[1] = parse_expr("{}".format(array_dim))
                     arg_decls.append(arg_init)
+                elif isinstance(ctx.func.argument_annotations[i],
+                                ti.sparse_matrix_builder):
+                    arg_init = parse_stmt(
+                        'x = ti.lang.kernel_arguments.decl_sparse_matrix()')
+                    arg_init.targets[0].id = arg.arg
+                    ctx.create_variable(arg.arg)
+                    arg_decls.append(arg_init)
                 elif isinstance(ctx.func.argument_annotations[i], ti.any_arr):
                     arg_init = parse_stmt(
                         'x = ti.lang.kernel_arguments.decl_any_arr_arg(0, 0, 0, 0)'

taichi/program/kernel.cpp

@@ -162,7 +162,7 @@ Kernel::LaunchContextBuilder::LaunchContextBuilder(Kernel *kernel)
 
 void Kernel::LaunchContextBuilder::set_arg_float(int arg_id, float64 d) {
   TI_ASSERT_INFO(!kernel_->args[arg_id].is_external_array,
-                 "Assigning scalar value to external(numpy) array argument is "
+                 "Assigning scalar value to external (numpy) array argument is "
                  "not allowed.");
 
   ActionRecorder::get_instance().record(
@@ -198,7 +198,7 @@ void Kernel::LaunchContextBuilder::set_arg_float(int arg_id, float64 d) {
 
 void Kernel::LaunchContextBuilder::set_arg_int(int arg_id, int64 d) {
   TI_ASSERT_INFO(!kernel_->args[arg_id].is_external_array,
-                 "Assigning scalar value to external(numpy) array argument is "
+                 "Assigning scalar value to external (numpy) array argument is "
                  "not allowed.");
 
   ActionRecorder::get_instance().record(
@@ -242,7 +242,7 @@ void Kernel::LaunchContextBuilder::set_arg_external_array(int arg_id,
                                                           uint64 size) {
   TI_ASSERT_INFO(
       kernel_->args[arg_id].is_external_array,
-      "Assigning external(numpy) array to scalar argument is not allowed.");
+      "Assigning external (numpy) array to scalar argument is not allowed.");
 
   ActionRecorder::get_instance().record(
       "set_kernel_arg_ext_ptr",
@@ -256,7 +256,7 @@ void Kernel::LaunchContextBuilder::set_arg_external_array(int arg_id,
 
 void Kernel::LaunchContextBuilder::set_arg_raw(int arg_id, uint64 d) {
   TI_ASSERT_INFO(!kernel_->args[arg_id].is_external_array,
-                 "Assigning scalar value to external(numpy) array argument is "
+                 "Assigning scalar value to external (numpy) array argument is "
                  "not allowed.");
 
   if (!kernel_->is_evaluator) {

taichi/program/program.cpp

@@ -24,6 +24,7 @@
 #include "taichi/program/snode_expr_utils.h"
 #include "taichi/util/statistics.h"
 #include "taichi/math/arithmetic.h"
+
 #if defined(TI_WITH_CC)
 #include "taichi/backends/cc/struct_cc.h"
 #include "taichi/backends/cc/cc_layout.h"

taichi/program/program.h

@@ -31,6 +31,7 @@
 #include "taichi/system/memory_pool.h"
 #include "taichi/system/threading.h"
 #include "taichi/system/unified_allocator.h"
+#include "taichi/program/sparse_matrix.h"
 
 namespace taichi {
 namespace lang {