[Lang] Enable definition of local matrices/vectors

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,12 @@ 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)
+            ) and impl.current_cfg().real_matrix and in_taichi_scope():
+            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/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 Matrix(*rhs.to_list(), ndim=rhs.ndim)
     if isinstance(rhs, Matrix):
+        if current_cfg().real_matrix:
+            return rhs
         return Matrix(rhs.to_list(), ndim=rhs.ndim)
     if isinstance(rhs, SharedArray):
         return rhs

python/taichi/lang/matrix.py

@@ -97,6 +97,20 @@ def prop_setter(instance, value):
     return cls
 
 
+def make_matrix(arr, dt=None):
+    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

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

taichi/analysis/offline_cache_util.cpp

@@ -65,6 +65,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)

taichi/codegen/llvm/codegen_llvm.cpp

@@ -124,7 +124,7 @@ void TaskCodeGenLLVM::visit(Block *stmt_list) {
 void TaskCodeGenLLVM::visit(AllocaStmt *stmt) {
   if (stmt->ret_type->is<TensorType>()) {
     auto tensor_type = stmt->ret_type->cast<TensorType>();
-    auto type = tlctx->get_data_type(tensor_type->get_element_type());
+    auto type = tlctx->get_data_type(tensor_type);
     auto array_size = tlctx->get_constant(tensor_type->get_num_elements());
     // Return type is [array_size x type]*.
     if (stmt->is_shared) {
@@ -688,6 +688,11 @@ llvm::Type *TaskCodeGenLLVM::llvm_type(DataType dt) {
     return llvm::Type::getDoubleTy(*llvm_context);
   } else if (dt->is_primitive(PrimitiveTypeID::f16)) {
     return llvm::Type::getHalfTy(*llvm_context);
+  } else if (dt->is<TensorType>()) {
+    auto tensor_type = dt->cast<TensorType>();
+    auto element_type = llvm_type(tensor_type->get_element_type());
+    return llvm::VectorType::get(element_type, tensor_type->get_num_elements(),
+                                 false);
   } else {
     TI_NOT_IMPLEMENTED;
   }
@@ -800,12 +805,20 @@ void TaskCodeGenLLVM::visit(PrintStmt *stmt) {
     if (std::holds_alternative<Stmt *>(content)) {
       auto arg_stmt = std::get<Stmt *>(content);
       auto value = llvm_val[arg_stmt];
-      if (arg_stmt->ret_type->is_primitive(PrimitiveTypeID::f32) ||
-          arg_stmt->ret_type->is_primitive(PrimitiveTypeID::f16))
-        value = builder->CreateFPExt(value,
-                                     tlctx->get_data_type(PrimitiveType::f64));
-      args.push_back(value);
-      formats += data_type_format(arg_stmt->ret_type);
+      if (arg_stmt->ret_type->is<TensorType>()) {
+        auto dtype = arg_stmt->ret_type->cast<TensorType>();
+        for (int i = 0; i < dtype->get_num_elements(); ++i) {
+          args.push_back(builder->CreateExtractElement(value, i));
+        }
+        formats += data_type_format(arg_stmt->ret_type);
+      } else {
+        if (arg_stmt->ret_type->is_primitive(PrimitiveTypeID::f32) ||
+            arg_stmt->ret_type->is_primitive(PrimitiveTypeID::f16))
+          value = builder->CreateFPExt(
+              value, tlctx->get_data_type(PrimitiveType::f64));
+        args.push_back(value);
+        formats += data_type_format(arg_stmt->ret_type);
+      }
     } else {
       auto arg_str = std::get<std::string>(content);
       auto value = builder->CreateGlobalStringPtr(arg_str, "content_string");
@@ -2515,6 +2528,16 @@ void TaskCodeGenLLVM::visit(MeshPatchIndexStmt *stmt) {
   llvm_val[stmt] = get_arg(2);
 }
 
+void TaskCodeGenLLVM::visit(MatrixInitStmt *stmt) {
+  auto type = tlctx->get_data_type(stmt->ret_type->as<TensorType>());
+  llvm::Value *vec = llvm::UndefValue::get(type);
+  for (int i = 0; i < stmt->values.size(); ++i) {
+    auto *elem = llvm_val[stmt->values[i]];
+    vec = builder->CreateInsertElement(vec, elem, i);
+  }
+  llvm_val[stmt] = vec;
+}
+
 void TaskCodeGenLLVM::eliminate_unused_functions() {
   TaichiLLVMContext::eliminate_unused_functions(
       module.get(), [&](std::string func_name) {

taichi/codegen/llvm/codegen_llvm.h

@@ -369,6 +369,8 @@ class TaskCodeGenLLVM : public IRVisitor, public LLVMModuleBuilder {
 
   void visit(ReferenceStmt *stmt) override;
 
+  void visit(MatrixInitStmt *stmt) override;
+
   llvm::Value *create_xlogue(std::unique_ptr<Block> &block);
 
   llvm::Value *create_mesh_xlogue(std::unique_ptr<Block> &block);

taichi/inc/expressions.inc.h

@@ -7,6 +7,7 @@ PER_EXPRESSION(InternalFuncCallExpression)
 PER_EXPRESSION(ExternalTensorExpression)
 PER_EXPRESSION(GlobalVariableExpression)
 PER_EXPRESSION(IndexExpression)
+PER_EXPRESSION(MatrixExpression)
 PER_EXPRESSION(StrideExpression)
 PER_EXPRESSION(RangeAssumptionExpression)
 PER_EXPRESSION(LoopUniqueExpression)

taichi/inc/statements.inc.h

@@ -38,6 +38,7 @@ PER_STATEMENT(LoopUniqueStmt)
 PER_STATEMENT(AssertStmt)
 PER_STATEMENT(ExternalFuncCallStmt)
 PER_STATEMENT(ExternalTensorShapeAlongAxisStmt)
+PER_STATEMENT(MatrixInitStmt)
 
 // Locals with reverse-mode autodiff
 PER_STATEMENT(AdStackAllocaStmt)

taichi/ir/control_flow_graph.cpp

@@ -501,7 +501,8 @@ bool CFGNode::dead_store_elimination(bool after_lower_access) {
       }
     }
     auto load_ptrs = irpass::analysis::get_load_pointers(stmt);
-    if (load_ptrs.size() == 1 && store_ptrs.empty() && stmt->width() == 1) {
+    if (load_ptrs.size() == 1 && store_ptrs.empty() && stmt->width() == 1 &&
+        !stmt->is<PtrOffsetStmt>()) {
       // Identical load elimination
       auto load_ptr = load_ptrs.front();
       if (!after_lower_access ||

taichi/ir/expression_printer.h

@@ -110,6 +110,13 @@ class ExpressionHumanFriendlyPrinter : public ExpressionPrinter {
     }
   }
 
+  void visit(MatrixExpression *expr) override {
+    emit('[');
+    emit_vector(expr->elements);
+    emit(']');
+    emit(fmt::format(" (dt={})", expr->dt->to_string()));
+  }
+
   void visit(IndexExpression *expr) override {
     expr->var->accept(this);
     emit('[');

taichi/ir/frontend_ir.cpp

@@ -429,6 +429,25 @@ Stmt *make_tensor_access(Expression::FlattenContext *ctx,
   return ctx->push_back<PtrOffsetStmt>(var->stmt, offset_stmt);
 }
 
+void MatrixExpression::type_check(CompileConfig *config) {
+  // TODO: typecheck matrix
+  for (auto &arg : elements) {
+    TI_ASSERT_TYPE_CHECKED(arg);
+  }
+}
+
+void MatrixExpression::flatten(FlattenContext *ctx) {
+  // TODO: implement flatten
+  TI_ASSERT(this->dt->is<TensorType>());
+  std::vector<Stmt *> values;
+  for (auto &elt : elements) {
+    flatten_rvalue(elt, ctx);
+    values.push_back(elt->stmt);
+  }
+  stmt = ctx->push_back<MatrixInitStmt>(values);
+  stmt->ret_type = this->dt;
+}
+
 bool IndexExpression::is_field() const {
   return var.is<GlobalVariableExpression>();
 }
@@ -960,6 +979,13 @@ Expr ASTBuilder::expr_alloca() {
   return var;
 }
 
+Expr ASTBuilder::expr_alloca_local_matrix(const std::vector<int> &shape,
+                                          const std::optional<DataType> &dt,
+                                          const std::vector<Expr> &elements) {
+  auto dtype = dt.value_or(PrimitiveType::unknown);
+  return Expr(std::make_shared<MatrixExpression>(elements, shape, dtype));
+}
+
 Expr ASTBuilder::expr_alloca_local_tensor(const std::vector<int> &shape,
                                           const DataType &element_type,
                                           const ExprGroup &elements) {

taichi/ir/frontend_ir.h

@@ -504,6 +504,26 @@ class GlobalVariableExpression : public Expression {
   TI_DEFINE_ACCEPT_FOR_EXPRESSION
 };
 
+class MatrixExpression : public Expression {
+ public:
+  std::vector<Expr> elements;
+  DataType dt;
+
+  MatrixExpression(const std::vector<Expr> &elements,
+                   std::vector<int> shape,
+                   DataType element_type)
+      : elements(elements) {
+    this->dt = DataType(TypeFactory::create_tensor_type(shape, element_type));
+    this->ret_type = this->dt;
+  }
+
+  void type_check(CompileConfig *config) override;
+
+  void flatten(FlattenContext *ctx) override;
+
+  TI_DEFINE_ACCEPT_FOR_EXPRESSION
+};
+
 class IndexExpression : public Expression {
  public:
   // `var` is one of GlobalVariableExpression, ExternalTensorExpression,
@@ -876,6 +896,9 @@ class ASTBuilder {
                                  const ExprGroup &args,
                                  const ExprGroup &outputs);
   Expr expr_alloca();
+  Expr expr_alloca_local_matrix(const std::vector<int> &shape,
+                                const std::optional<DataType> &dt,
+                                const std::vector<Expr> &elements);
   Expr expr_alloca_local_tensor(const std::vector<int> &shape,
                                 const DataType &element_type,
                                 const ExprGroup &elements);

taichi/ir/statements.h

@@ -1807,5 +1807,17 @@ class MeshPatchIndexStmt : public Stmt {
   TI_DEFINE_ACCEPT_AND_CLONE
 };
 
+class MatrixInitStmt : public Stmt {
+ public:
+  std::vector<Stmt *> values;
+
+  MatrixInitStmt(const std::vector<Stmt *> &values) : values(values) {
+    TI_STMT_REG_FIELDS;
+  }
+
+  TI_STMT_DEF_FIELDS(ret_type, values);
+  TI_DEFINE_ACCEPT_AND_CLONE
+};
+
 }  // namespace lang
 }  // namespace taichi

taichi/ir/type.cpp

@@ -87,6 +87,14 @@ std::string TensorType::to_string() const {
   return s;
 }
 
+int TensorType::vector_width() const {
+  int vw = 1;
+  for (auto dim : shape_) {
+    vw *= dim;
+  }
+  return vw;
+}
+
 int Type::vector_width() const {
   return 1;  // TODO: CPU vectorization
 }

taichi/ir/type.h

@@ -175,6 +175,8 @@ class TensorType : public Type {
     return shape_;
   }
 
+  int vector_width() const;
+
   Type *get_compute_type() override {
     return this;
   }