[Lang] Experimental sparse matrix support on CPUs

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}")
 

cmake/TaichiCore.cmake

@@ -220,6 +220,8 @@ llvm_map_components_to_libnames(llvm_libs
         )
 target_link_libraries(${LIBRARY_NAME} ${llvm_libs})
 
+include_directories(${LLVM_INCLUDE_DIRS})
+
 if (APPLE AND "${CMAKE_SYSTEM_PROCESSOR}" STREQUAL "arm64")
     llvm_map_components_to_libnames(llvm_aarch64_libs AArch64)
     target_link_libraries(${LIBRARY_NAME} ${llvm_aarch64_libs})

misc/sparse_matrix.py

@@ -0,0 +1,60 @@
+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.SparseMatrixBuilder, b: ti.SparseMatrixBuilder,
+         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(">>>>>>>> Before build: ")
+K.print_triplets()
+
+A = K.build()
+
+print(">>>>>>>> After build: ")
+print(A)
+
+print(">>>>>>> Summation Test: ")
+C = A + A
+print(C)
+
+print(">>>> Subtraction Test")
+D = A - A
+print(D)
+
+print(">>>> Multiplication with scalar")
+E = A * 3.0
+print(E)
+
+print(">>>> Transpose Test")
+F = A.transpose()
+print(F)
+
+print(">>>> Matrix Multiplication")
+G = E @ A
+print(G)
+
+print(">>> Elment-wise Multiplication")
+H = E * A
+print(H)
+
+print(f">>> Element Access: A[0,0] = {A[0,0]}")

python/taichi/lang/__init__.py

@@ -16,6 +16,7 @@
 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 SparseMatrixProxy
 from taichi.lang.util import cook_dtype, to_taichi_type
 
 
@@ -129,6 +131,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 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

@@ -13,6 +13,7 @@
 from taichi.lang.kernel_arguments import any_arr, ext_arr, template
 from taichi.lang.ndarray import ScalarNdarray
 from taichi.lang.shell import _shell_pop_print, oinspect
+from taichi.lang.sparse_matrix import SparseMatrixBuilder
 from taichi.lang.transformer import ASTTransformerTotal
 from taichi.misc.util import obsolete
 
@@ -368,6 +369,8 @@ def extract_arguments(self):
                     pass
                 elif id(annotation) in primitive_types.type_ids:
                     pass
+                elif id(annotation) == id(SparseMatrixBuilder):
+                    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 id(needed) == id(SparseMatrixBuilder):
+                    # Pass only the base pointer of the ti.SparseMatrixBuilder 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,92 @@
+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}), ({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}), ({other.n}, {other.m})"
+        sm = self.matrix - other.matrix
+        return SparseMatrix(sm=sm)
+
+    def __mul__(self, other):
+        if isinstance(other, float):
+            sm = other * self.matrix
+            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}), ({other.n}, {other.m})"
+            sm = self.matrix * other.matrix
+            return SparseMatrix(sm=sm)
+
+    def transpose(self):
+        sm = self.matrix.transpose()
+        return SparseMatrix(sm=sm)
+
+    def __matmul__(self, other):
+        sm = self.matrix.matmult(other.matrix)
+        return SparseMatrix(sm=sm)
+
+    def __getitem__(self, item):
+        return self.matrix.get_coeff(item[0], item[1])
+
+    def __str__(self):
+        return self.matrix.to_string()
+
+    def __repr__(self):
+        return self.matrix.to_string()
+
+
+class SparseMatrixBuilder:
+    def __init__(self, n, m=None, max_num_triplets=0):
+        self.n = n
+        self.m = m if m else n
+        from taichi.core.util import ti_core as _ti_core
+        self.ptr = _ti_core.create_sparse_matrix_builder(
+            n, m, 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)
+
+
+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
+        if op == 'Add':
+            call_internal("insert_triplet", self.ptr, self.i, self.j, float(value))
+        elif op == 'Sub':
+            call_internal("insert_triplet", self.ptr, self.i, self.j, -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)

python/taichi/lang/stmt_builder.py

@@ -6,6 +6,7 @@
 from taichi.lang.ast_resolver import ASTResolver
 from taichi.lang.exception import TaichiSyntaxError
 from taichi.lang.expr_builder import build_expr, build_exprs
+from taichi.lang.sparse_matrix import SparseMatrixBuilder
 from taichi.lang.util import to_taichi_type
 
 import taichi as ti
@@ -555,6 +556,12 @@ 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 ctx.func.argument_annotations[i] is SparseMatrixBuilder:
+                    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/ir/frontend_ir.h

@@ -78,6 +78,22 @@ class FrontendAssertStmt : public Stmt {
   TI_DEFINE_ACCEPT
 };
 
+class FrontendInternalFuncStmt : public Stmt {
+ public:
+  std::string func_name;
+  std::vector<Expr> args;
+
+  FrontendInternalFuncStmt(const std::string &func_name,
+                           std::vector<Expr> args_)
+      : func_name(func_name) {
+    for (auto &a : args_) {
+      args.push_back(load_if_ptr(a));
+    }
+  }
+
+  TI_DEFINE_ACCEPT
+};
+
 class FrontendAssignStmt : public Stmt {
  public:
   Expr lhs, rhs;

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 {