[Lang] Enable matrix representation in Frontend/CHI IR

python/taichi/lang/ast/ast_transformer.py

@@ -14,7 +14,7 @@
 from taichi.lang.ast.symbol_resolver import ASTResolver
 from taichi.lang.exception import TaichiSyntaxError, TaichiTypeError
 from taichi.lang.field import Field
-from taichi.lang.matrix import (Matrix, MatrixType, _PyScopeMatrixImpl,
+from taichi.lang.matrix import (Matrix, MatrixType, Vector, _PyScopeMatrixImpl,
                                 _TiScopeMatrixImpl)
 from taichi.lang.snode import append
 from taichi.lang.util import in_taichi_scope, is_taichi_class, to_taichi_type
@@ -488,6 +488,11 @@ def build_Call(ctx, node):
             node.ptr = impl.ti_format(*args, **keywords)
             return node.ptr
 
+        if isinstance(node.func,
+                      ast.Attribute) and func == Matrix or func == Vector:
+            node.ptr = matrix.make_matrix(*args, **keywords)
+            return node.ptr
+
         if ASTTransformer.build_call_if_is_builtin(ctx, node, args, keywords):
             return node.ptr
 

python/taichi/lang/expr.py

@@ -38,6 +38,15 @@ def __init__(self, *args, tb=None, dtype=None):
             self.ptr.set_tb(self.tb)
         self.ptr.type_check(impl.get_runtime().prog.config)
 
+    def __getitem__(self, *indices):
+        if not isinstance(indices, (list, tuple)):
+            indices = (indices, )
+
+        indices = make_expr_group(*indices)
+        return Expr(
+            impl.get_runtime().prog.current_ast_builder().expr_indexed_matrix(
+                self.ptr, indices))
+
     def __hash__(self):
         return self.ptr.get_raw_address()
 

python/taichi/lang/impl.py

@@ -35,6 +35,12 @@ def expr_init_local_tensor(shape, element_type, elements):
         shape, element_type, elements)
 
 
+@taichi_scope
+def expr_init_matrix(shape, element_type, elements):
+    return get_runtime().prog.current_ast_builder().expr_alloca_matrix(
+        shape, element_type, elements)
+
+
 @taichi_scope
 def expr_init_shared_array(shape, element_type):
     return get_runtime().prog.current_ast_builder().expr_alloca_shared_array(
@@ -48,6 +54,8 @@ def expr_init(rhs):
     if isinstance(rhs, Matrix) and (hasattr(rhs, "_DIM")):
         return type(rhs)(*rhs.to_list())
     if isinstance(rhs, Matrix):
+        if current_cfg().real_matrix:
+            return rhs
         return Matrix(rhs.to_list())
     if isinstance(rhs, SharedArray):
         return rhs

python/taichi/lang/matrix.py

@@ -97,6 +97,22 @@ def prop_setter(instance, value):
     return cls
 
 
+def make_matrix(arr, dt=None, suppress_warning=False, is_ref=False, **kwargs):
+    if not impl.current_cfg().real_matrix or in_python_scope():
+        return Matrix(arr, dt, suppress_warning, is_ref, **kwargs)
+    cast = (lambda x: ops_mod.cast(x, dt)) if dt else (
+        lambda x: x if isinstance(x, expr.Expr) else expr.Expr(x))
+    if len(arr) == 0:
+        return impl.expr_init(impl.expr_init_matrix([0], dt, []))
+    if not isinstance(arr[0], Iterable):
+        return impl.expr_init(
+            impl.expr_init_matrix([len(arr)], dt,
+                                  [cast(elt).ptr for elt in arr]))
+    return impl.expr_init(
+        impl.expr_init_matrix([len(arr), len(arr[0])], dt,
+                              [cast(elt).ptr for row in arr for elt in row]))
+
+
 class _MatrixBaseImpl:
     def __init__(self, m, n, entries):
         self.m = m

taichi/analysis/data_source_analysis.cpp

@@ -37,6 +37,10 @@ std::vector<Stmt *> get_load_pointers(Stmt *load_stmt) {
     return external_func->arg_stmts;
   } else if (auto ref = load_stmt->cast<ReferenceStmt>()) {
     return {ref->var};
+  } else if (auto matrix_init = load_stmt->cast<MatrixInitStmt>()) {
+    return matrix_init->values;
+  } else if (auto ptr_offset = load_stmt->cast<PtrOffsetStmt>()) {
+    return {ptr_offset->origin};
   } else {
     return std::vector<Stmt *>();
   }
@@ -59,7 +63,7 @@ Stmt *get_store_data(Stmt *store_stmt) {
 
 std::vector<Stmt *> get_store_destination(Stmt *store_stmt) {
   // If store_stmt provides some data sources, return the pointers of the data.
-  if (store_stmt->is<AllocaStmt>() && !store_stmt->ret_type->is<TensorType>()) {
+  if (store_stmt->is<AllocaStmt>()) {
     // The statement itself provides a data source (const [0]).
     return std::vector<Stmt *>(1, store_stmt);
   } else if (auto local_store = store_stmt->cast<LocalStoreStmt>()) {

taichi/analysis/gen_offline_cache_key.cpp

@@ -159,6 +159,14 @@ class ASTSerializer : public IRVisitor, public ExpressionVisitor {
     emit(expr->dual);
   }
 
+  void visit(MatrixExpression *expr) override {
+    emit(ExprOpCode::MatrixExpression);
+    emit(expr->dt);
+    for (auto elt : expr->elements) {
+      emit(elt);
+    }
+  }
+
   void visit(IndexExpression *expr) override {
     emit(ExprOpCode::IndexExpression);
     emit(expr->var);

taichi/analysis/offline_cache_util.cpp

@@ -64,6 +64,7 @@ static std::vector<std::uint8_t> get_offline_cache_key_of_compile_config(
   serializer(config->demote_no_access_mesh_fors);
   serializer(config->experimental_auto_mesh_local);
   serializer(config->auto_mesh_local_default_occupacy);
+  serializer(config->real_matrix);
   serializer.finalize();
 
   return serializer.data;

taichi/analysis/same_statements.cpp

@@ -196,6 +196,24 @@ class IRNodeComparator : public IRVisitor {
     basic_check(stmt);
   }
 
+  void visit(MatrixInitStmt *stmt) override {
+    basic_check(stmt);
+    if (!same)
+      return;
+    auto o = other_node_->as<MatrixInitStmt>();
+    if (stmt->values.size() != o->values.size()) {
+      same = false;
+      return;
+    }
+    for (int i = 0; i < stmt->values.size(); ++i) {
+      other_node_ = o->values[i];
+      stmt->values[i]->accept(this);
+      other_node_ = o;
+      if (!same)
+        return;
+    }
+  }
+
   void visit(IfStmt *stmt) override {
     basic_check(stmt);
     if (!same)