[IR] Support local tensor (#2637)

* start working on local tensor

* fix python

* try to solve local load for local tensor

* fix __matmul__

* fix test where PtrOffsetStmt may be the same

* fix typo

* disable local tensor for loop index

* fix store_to_load_forwarding with TensorType Alloca

* fix data type for TensorType Alloca

* local tensor for ti.Matrix

* fix local matrix index

* fix offload with local tensor

* fix AllocaStmt->AllocaStmt/GlobalTemporaryStmt crossing different offloaded tasks

* fix PtrOffsetStmt crossing offloaded tasks issue

* demote atomics for PtrOffsetStmt(AllocaStmt)

* fix initializing to zero with Alloca TensorType

* disable local tensor for python scope

* fix vector_width() with 1 for TensorType

* apply ti.extension.dynamic_index for local tensor

* try to fix Vector.zero()

* fix allocate global (not limited to AllocaStmt and PtrOffsetStmt) within offloaded pass

* disable make_block_local with local tensor, this is related to loop indices issue

* Auto Format

* disable local tensor for loop index

* add missing semicolons

* Update taichi/ir/statements.h

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

* fix one

* fix redundant Matrix.__call__() condition

* reimplement is_unlowered_global_ptr()

* disable local tensor by default

* throw InvalidOperationError when using block_local() and dynamic_index at the same time

* need dt when using local tensor

* Auto Format

* unify global and local with option ti.init(dynamic_index=True)

* Auto Format

* fix trivial tests with dynamic_index turned on by default

* Auto Format

* fix bls tests with dynamic_index disabled

* Auto Format

* action test

* Auto Format

* stop show off

* more flexible _test_features

* Auto Format

* turn off testing packed by default

Co-authored-by: Taichi Gardener <[email protected]>
Co-authored-by: Ye Kuang <[email protected]>

python/taichi/lang/__init__.py

@@ -4,6 +4,7 @@
 
 from taichi.core.util import ti_core as _ti_core
 from taichi.lang import impl
+from taichi.lang.exception import InvalidOperationError
 from taichi.lang.impl import *
 from taichi.lang.kernel_arguments import ext_arr, template
 from taichi.lang.kernel_impl import (KernelArgError, KernelDefError,
@@ -357,6 +358,9 @@ def no_activate(*args):
 
 
 def block_local(*args):
+    if ti.current_cfg().dynamic_index:
+        raise InvalidOperationError(
+            'dynamic_index is not allowed when block_local is turned on.')
     for a in args:
         for v in a.get_field_members():
             _ti_core.insert_snode_access_flag(

python/taichi/lang/expr.py

@@ -50,7 +50,7 @@ def make_var_vector(size):
     exprs = []
     for _ in range(size):
         exprs.append(_ti_core.make_id_expr(''))
-    return ti.Vector(exprs)
+    return ti.Vector(exprs, disable_local_tensor=True)
 
 
 def make_expr_group(*exprs):

python/taichi/lang/impl.py

@@ -19,12 +19,20 @@
 import taichi as ti
 
 
+@taichi_scope
+def expr_init_local_tensor(shape, element_type, elements):
+    return _ti_core.expr_alloca_local_tensor(shape, element_type, elements)
+
+
 @taichi_scope
 def expr_init(rhs):
     if rhs is None:
         return Expr(_ti_core.expr_alloca())
     if is_taichi_class(rhs):
-        return rhs.variable()
+        if rhs.local_tensor_proxy != None:
+            return rhs
+        else:
+            return rhs.variable()
     else:
         if isinstance(rhs, list):
             return [expr_init(e) for e in rhs]
@@ -167,10 +175,17 @@ def subscript(value, *indices):
 
 
 @taichi_scope
-def subscript_with_offset(var, indices, cols, is_aos):
+def local_subscript_with_offset(var, indices):
+    return Expr(
+        _ti_core.local_subscript_with_offset(var, make_expr_group(*indices)))
+
+
+@taichi_scope
+def global_subscript_with_offset(var, indices, cols, is_aos):
     return Expr(
-        _ti_core.subscript_with_offset(var.ptr, make_expr_group(*indices),
-                                       cols, is_aos))
+        _ti_core.global_subscript_with_offset(var.ptr,
+                                              make_expr_group(*indices), cols,
+                                              is_aos))
 
 
 @taichi_scope

python/taichi/lang/linalg.py

@@ -21,7 +21,7 @@ def polar_decompose2d(A, dt):
     scale = (1.0 / ti.sqrt(x * x + y * y))
     c = x * scale
     s = y * scale
-    r = ti.Matrix([[c, -s], [s, c]])
+    r = ti.Matrix([[c, -s], [s, c]], dt=dt)
     return r, r.transpose() @ A
 
 
@@ -79,11 +79,11 @@ def svd2d(A, dt):
         tmp = s1
         s1 = s2
         s2 = tmp
-        V = [[-s, c], [-c, -s]]
+        V = ti.Matrix([[-s, c], [-c, -s]], dt=dt)
     else:
-        V = [[c, s], [-s, c]]
+        V = ti.Matrix([[c, s], [-s, c]], dt=dt)
     U = R @ V
-    return U, ti.Matrix([[s1, ti.cast(0, dt)], [ti.cast(0, dt), s2]]), V
+    return U, ti.Matrix([[s1, ti.cast(0, dt)], [ti.cast(0, dt), s2]], dt=dt), V
 
 
 def svd3d(A, dt, iters=None):
@@ -128,7 +128,7 @@ def svd3d(A, dt, iters=None):
 
 @ti.func
 def eig2x2(A, dt):
-    """Compute the eigenvalues and right eigenvectors (Av=\lambda v) of a 2x2 real matrix.
+    """Compute the eigenvalues and right eigenvectors (Av=lambda v) of a 2x2 real matrix.
 
     Mathematical concept refers to https://en.wikipedia.org/wiki/Eigendecomposition_of_a_matrix.
 
@@ -148,30 +148,28 @@ def eig2x2(A, dt):
     v1 = ti.Vector.zero(dt, 4)
     v2 = ti.Vector.zero(dt, 4)
     if gap > 0:
-        lambda1 = ti.Vector([tr + ti.sqrt(gap), ti.cast(0.0, dt)]) * 0.5
-        lambda2 = ti.Vector([tr - ti.sqrt(gap), ti.cast(0.0, dt)]) * 0.5
+        lambda1 = ti.Vector([tr + ti.sqrt(gap), 0.0], dt=dt) * 0.5
+        lambda2 = ti.Vector([tr - ti.sqrt(gap), 0.0], dt=dt) * 0.5
         A1 = A - lambda1[0] * ti.Matrix.identity(dt, 2)
         A2 = A - lambda2[0] * ti.Matrix.identity(dt, 2)
         if all(A1 == ti.Matrix.zero(dt, 2, 2)) and all(
                 A1 == ti.Matrix.zero(dt, 2, 2)):
             v1 = ti.Vector([0.0, 0.0, 1.0, 0.0]).cast(dt)
             v2 = ti.Vector([1.0, 0.0, 0.0, 0.0]).cast(dt)
         else:
-            v1 = ti.Vector([A2[0, 0], 0.0, A2[1, 0],
-                            0.0]).cast(dt).normalized()
-            v2 = ti.Vector([A1[0, 0], 0.0, A1[1, 0],
-                            0.0]).cast(dt).normalized()
+            v1 = ti.Vector([A2[0, 0], 0.0, A2[1, 0], 0.0], dt=dt).normalized()
+            v2 = ti.Vector([A1[0, 0], 0.0, A1[1, 0], 0.0], dt=dt).normalized()
     else:
-        lambda1 = ti.Vector([tr, ti.sqrt(-gap)]) * 0.5
-        lambda2 = ti.Vector([tr, -ti.sqrt(-gap)]) * 0.5
+        lambda1 = ti.Vector([tr, ti.sqrt(-gap)], dt=dt) * 0.5
+        lambda2 = ti.Vector([tr, -ti.sqrt(-gap)], dt=dt) * 0.5
         A1r = A - lambda1[0] * ti.Matrix.identity(dt, 2)
         A1i = -lambda1[1] * ti.Matrix.identity(dt, 2)
         A2r = A - lambda2[0] * ti.Matrix.identity(dt, 2)
         A2i = -lambda2[1] * ti.Matrix.identity(dt, 2)
-        v1 = ti.Vector([A2r[0, 0], A2i[0, 0], A2r[1, 0],
-                        A2i[1, 0]]).cast(dt).normalized()
-        v2 = ti.Vector([A1r[0, 0], A1i[0, 0], A1r[1, 0],
-                        A1i[1, 0]]).cast(dt).normalized()
+        v1 = ti.Vector([A2r[0, 0], A2i[0, 0], A2r[1, 0], A2i[1, 0]],
+                       dt=dt).normalized()
+        v2 = ti.Vector([A1r[0, 0], A1i[0, 0], A1r[1, 0], A1i[1, 0]],
+                       dt=dt).normalized()
     eigenvalues = ti.Matrix.rows([lambda1, lambda2])
     eigenvectors = ti.Matrix.cols([v1, v2])
 
@@ -180,7 +178,7 @@ def eig2x2(A, dt):
 
 @ti.func
 def sym_eig2x2(A, dt):
-    """Compute the eigenvalues and right eigenvectors (Av=\lambda v) of a 2x2 real symmetric matrix.
+    """Compute the eigenvalues and right eigenvectors (Av=lambda v) of a 2x2 real symmetric matrix.
 
     Mathematical concept refers to https://en.wikipedia.org/wiki/Eigendecomposition_of_a_matrix.
 
@@ -197,7 +195,7 @@ def sym_eig2x2(A, dt):
     gap = tr**2 - 4 * det
     lambda1 = (tr + ti.sqrt(gap)) * 0.5
     lambda2 = (tr - ti.sqrt(gap)) * 0.5
-    eigenvalues = ti.Vector([lambda1, lambda2]).cast(dt)
+    eigenvalues = ti.Vector([lambda1, lambda2], dt=dt)
 
     A1 = A - lambda1 * ti.Matrix.identity(dt, 2)
     A2 = A - lambda2 * ti.Matrix.identity(dt, 2)
@@ -208,8 +206,8 @@ def sym_eig2x2(A, dt):
         v1 = ti.Vector([0.0, 1.0]).cast(dt)
         v2 = ti.Vector([1.0, 0.0]).cast(dt)
     else:
-        v1 = ti.Vector([A2[0, 0], A2[1, 0]]).cast(dt).normalized()
-        v2 = ti.Vector([A1[0, 0], A1[1, 0]]).cast(dt).normalized()
+        v1 = ti.Vector([A2[0, 0], A2[1, 0]], dt=dt).normalized()
+        v2 = ti.Vector([A1[0, 0], A1[1, 0]], dt=dt).normalized()
     eigenvectors = ti.Matrix.cols([v1, v2])
     return eigenvalues, eigenvectors