[Lang] Matrix operations for new local matrix representation

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/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 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,27 @@ def prop_setter(instance, value):
     return cls
 
 
+def make_matrix(arr, dt=None):
+    if len(arr) == 0:
+        return impl.expr_init(
+            impl.expr_init_local_tensor([0], ti_python_core.DataType_unknown,
+                                        impl.make_expr_group([])))
+    is_matrix = isinstance(arr[0], Iterable)
+    if dt is None:
+        dt = _make_entries_initializer(is_matrix).infer_dt(arr)
+    if not is_matrix:
+        return impl.expr_init(
+            impl.expr_init_local_tensor([len(arr)], dt,
+                                        impl.make_expr_group(
+                                            [expr.Expr(elt) for elt in arr])))
+    return impl.expr_init(
+        impl.expr_init_local_tensor([len(arr), len(arr[0])], dt,
+                                    impl.make_expr_group([
+                                        expr.Expr(elt) for row in arr
+                                        for elt in row
+                                    ])))
+
+
 class _MatrixBaseImpl:
     def __init__(self, m, n, entries):
         self.m = m