[Lang] [ir] Refactor indexing expressions in AST & enforce integer indices

Author: re-xyr

python/taichi/lang/impl.py

@@ -197,10 +197,16 @@ def subscript(value, *_indices, skip_reordered=False):
 
 
 @taichi_scope
-def make_tensor_element_expr(_var, _indices, shape, stride):
+def make_stride_expr(_var, _indices, shape, stride):
     return Expr(
-        _ti_core.make_tensor_element_expr(_var, make_expr_group(*_indices),
-                                          shape, stride))
+        _ti_core.make_stride_expr(_var, make_expr_group(*_indices),
+                                  shape, stride))
+
+
+@taichi_scope
+def make_index_expr(_var, _indices):
+    return Expr(
+        _ti_core.make_index_expr(_var, make_expr_group(*_indices)))
 
 
 class SrcInfoGuard:
@@ -433,6 +439,7 @@ def deactivate_all_snodes():
 
 class _Root:
     """Wrapper around the default root FieldsBuilder instance."""
+
     @staticmethod
     def parent(n=1):
         """Same as :func:`taichi.SNode.parent`"""
@@ -539,10 +546,10 @@ def field(dtype, shape=None, name="", offset=None, needs_grad=False):
     """
 
     if isinstance(shape, numbers.Number):
-        shape = (shape, )
+        shape = (shape,)
 
     if isinstance(offset, numbers.Number):
-        offset = (offset, )
+        offset = (offset,)
 
     if shape is not None and offset is not None:
         assert len(shape) == len(
@@ -583,7 +590,7 @@ def ndarray(dtype, shape, layout=Layout.NULL):
             >>> z = ti.ndarray(matrix_ty, shape=(4, 5), layout=ti.Layout.SOA)  # ndarray of shape (4, 5), each element is a matrix of (3, 4) ti.float scalars.
     """
     if isinstance(shape, numbers.Number):
-        shape = (shape, )
+        shape = (shape,)
     if dtype in types:
         assert layout == Layout.NULL
         return ScalarNdarray(dtype, shape)

python/taichi/lang/matrix.py

@@ -41,7 +41,6 @@ def check(instance, pattern):
 
     for key_group in KEYGROUP_SET:
         for index, attr in enumerate(key_group):
-
             def gen_property(attr, attr_idx, key_group):
                 checker = make_valid_attribs_checker(key_group)
 
@@ -122,7 +121,7 @@ def _linearize_entry_id(self, *args):
         if len(args) == 1 and isinstance(args[0], (list, tuple)):
             args = args[0]
         if len(args) == 1:
-            args = args + (0, )
+            args = args + (0,)
         # TODO(#1004): See if it's possible to support indexing at runtime
         for i, a in enumerate(args):
             if not isinstance(a, int):
@@ -244,19 +243,16 @@ def _subscript(self, is_global_mat, *indices):
         if self.any_array_access:
             return self.any_array_access.subscript(i, j)
         if self.local_tensor_proxy is not None:
-            assert self.dynamic_index_stride is not None
             if len(indices) == 1:
-                return impl.make_tensor_element_expr(self.local_tensor_proxy,
-                                                     (i, ), (self.n, ),
-                                                     self.dynamic_index_stride)
-            return impl.make_tensor_element_expr(self.local_tensor_proxy,
-                                                 (i, j), (self.n, self.m),
-                                                 self.dynamic_index_stride)
+                return impl.make_index_expr(self.local_tensor_proxy,
+                                            (i,))
+            return impl.make_index_expr(self.local_tensor_proxy,
+                                        (i, j))
         if impl.current_cfg(
         ).dynamic_index and is_global_mat and self.dynamic_index_stride:
-            return impl.make_tensor_element_expr(self.entries[0].ptr, (i, j),
-                                                 (self.n, self.m),
-                                                 self.dynamic_index_stride)
+            return impl.make_stride_expr(self.entries[0].ptr, (i, j),
+                                         (self.n, self.m),
+                                         self.dynamic_index_stride)
         return self._get_entry(i, j)
 
     def _calc_slice(self, index, dim):
@@ -318,17 +314,15 @@ def with_dynamic_index(self, arr, dt):
             local_tensor_proxy = impl.expr_init_local_tensor(
                 [len(arr)], dt,
                 expr.make_expr_group([expr.Expr(x) for x in arr]))
-            dynamic_index_stride = 1
             mat = []
             for i in range(len(arr)):
                 mat.append(
                     list([
-                        impl.make_tensor_element_expr(
+                        impl.make_index_expr(
                             local_tensor_proxy,
-                            (expr.Expr(i, dtype=primitive_types.i32), ),
-                            (len(arr), ), dynamic_index_stride)
+                            (expr.Expr(i, dtype=primitive_types.i32),))
                     ]))
-            return local_tensor_proxy, dynamic_index_stride, mat
+            return local_tensor_proxy, mat
 
         def _get_entry_to_infer(self, arr):
             return arr[0]
@@ -348,18 +342,16 @@ def with_dynamic_index(self, arr, dt):
                 expr.make_expr_group(
                     [expr.Expr(x) for row in arr for x in row]))
 
-            dynamic_index_stride = 1
             mat = []
             for i in range(len(arr)):
                 mat.append([])
                 for j in range(len(arr[0])):
                     mat[i].append(
-                        impl.make_tensor_element_expr(
+                        impl.make_index_expr(
                             local_tensor_proxy,
                             (expr.Expr(i, dtype=primitive_types.i32),
-                             expr.Expr(j, dtype=primitive_types.i32)),
-                            (len(arr), len(arr[0])), dynamic_index_stride))
-            return local_tensor_proxy, dynamic_index_stride, mat
+                             expr.Expr(j, dtype=primitive_types.i32))))
+            return local_tensor_proxy, mat
 
         def _get_entry_to_infer(self, arr):
             return arr[0][0]
@@ -413,7 +405,6 @@ class Matrix(TaichiOperations):
 
     def __init__(self, arr, dt=None, suppress_warning=False, is_ref=False):
         local_tensor_proxy = None
-        dynamic_index_stride = None
 
         if not isinstance(arr, (list, tuple, np.ndarray)):
             raise TaichiTypeError(
@@ -440,7 +431,7 @@ def __init__(self, arr, dt=None, suppress_warning=False, is_ref=False):
                     )
                 if dt is None:
                     dt = initializer.infer_dt(arr)
-                local_tensor_proxy, dynamic_index_stride, mat = initializer.with_dynamic_index(
+                local_tensor_proxy, mat = initializer.with_dynamic_index(
                     arr, dt)
 
         self.n, self.m = len(mat), 1
@@ -464,7 +455,7 @@ def __init__(self, arr, dt=None, suppress_warning=False, is_ref=False):
             self._impl = _PyScopeMatrixImpl(m, n, entries)
         else:
             self._impl = _TiScopeMatrixImpl(m, n, entries, local_tensor_proxy,
-                                            dynamic_index_stride)
+                                            None)
 
     def _element_wise_binary(self, foo, other):
         other = self._broadcast_copy(other)
@@ -680,8 +671,8 @@ def E(x, y):
             for i in range(n):
                 for j in range(n):
                     entries[j][i] = inv_determinant * (
-                        E(i + 1, j + 1) * E(i + 2, j + 2) -
-                        E(i + 2, j + 1) * E(i + 1, j + 2))
+                            E(i + 1, j + 1) * E(i + 2, j + 2) -
+                            E(i + 2, j + 1) * E(i + 1, j + 2))
             return Matrix(entries)
         if self.n == 4:
             n = 4
@@ -693,14 +684,14 @@ def E(x, y):
 
             for i in range(n):
                 for j in range(n):
-                    entries[j][i] = inv_determinant * (-1)**(i + j) * ((
-                        E(i + 1, j + 1) *
-                        (E(i + 2, j + 2) * E(i + 3, j + 3) -
-                         E(i + 3, j + 2) * E(i + 2, j + 3)) - E(i + 2, j + 1) *
-                        (E(i + 1, j + 2) * E(i + 3, j + 3) -
-                         E(i + 3, j + 2) * E(i + 1, j + 3)) + E(i + 3, j + 1) *
-                        (E(i + 1, j + 2) * E(i + 2, j + 3) -
-                         E(i + 2, j + 2) * E(i + 1, j + 3))))
+                    entries[j][i] = inv_determinant * (-1) ** (i + j) * ((
+                            E(i + 1, j + 1) *
+                            (E(i + 2, j + 2) * E(i + 3, j + 3) -
+                             E(i + 3, j + 2) * E(i + 2, j + 3)) - E(i + 2, j + 1) *
+                            (E(i + 1, j + 2) * E(i + 3, j + 3) -
+                             E(i + 3, j + 2) * E(i + 1, j + 3)) + E(i + 3, j + 1) *
+                            (E(i + 1, j + 2) * E(i + 2, j + 3) -
+                             E(i + 2, j + 2) * E(i + 1, j + 3))))
             return Matrix(entries)
         raise Exception(
             "Inversions of matrices with sizes >= 5 are not supported")
@@ -760,7 +751,7 @@ def determinant(a):
         if a.n == 3 and a.m == 3:
             return a(0, 0) * (a(1, 1) * a(2, 2) - a(2, 1) * a(1, 2)) - a(
                 1, 0) * (a(0, 1) * a(2, 2) - a(2, 1) * a(0, 2)) + a(
-                    2, 0) * (a(0, 1) * a(1, 2) - a(1, 1) * a(0, 2))
+                2, 0) * (a(0, 1) * a(1, 2) - a(1, 1) * a(0, 2))
         if a.n == 4 and a.m == 4:
             n = 4
 
@@ -769,14 +760,14 @@ def E(x, y):
 
             det = impl.expr_init(0.0)
             for i in range(4):
-                det = det + (-1.0)**i * (
-                    a(i, 0) *
-                    (E(i + 1, 1) *
-                     (E(i + 2, 2) * E(i + 3, 3) - E(i + 3, 2) * E(i + 2, 3)) -
-                     E(i + 2, 1) *
-                     (E(i + 1, 2) * E(i + 3, 3) - E(i + 3, 2) * E(i + 1, 3)) +
-                     E(i + 3, 1) *
-                     (E(i + 1, 2) * E(i + 2, 3) - E(i + 2, 2) * E(i + 1, 3))))
+                det = det + (-1.0) ** i * (
+                        a(i, 0) *
+                        (E(i + 1, 1) *
+                         (E(i + 2, 2) * E(i + 3, 3) - E(i + 3, 2) * E(i + 2, 3)) -
+                         E(i + 2, 1) *
+                         (E(i + 1, 2) * E(i + 3, 3) - E(i + 3, 2) * E(i + 1, 3)) +
+                         E(i + 3, 1) *
+                         (E(i + 1, 2) * E(i + 2, 3) - E(i + 2, 2) * E(i + 1, 3))))
             return det
         raise Exception(
             "Determinants of matrices with sizes >= 5 are not supported")
@@ -908,6 +899,7 @@ def fill(self, val):
             >>> A
             [-1, -1, -1, -1]
         """
+
         def assign_renamed(x, y):
             return ops_mod.assign(x, y)
 
@@ -933,7 +925,7 @@ def to_numpy(self, keep_dims=False):
             array([0, 1, 2, 3])
         """
         as_vector = self.m == 1 and not keep_dims
-        shape_ext = (self.n, ) if as_vector else (self.n, self.m)
+        shape_ext = (self.n,) if as_vector else (self.n, self.m)
         return np.array(self.to_list()).reshape(shape_ext)
 
     @taichi_scope
@@ -1128,9 +1120,9 @@ def field(cls,
 
         if shape is not None:
             if isinstance(shape, numbers.Number):
-                shape = (shape, )
+                shape = (shape,)
             if isinstance(offset, numbers.Number):
-                offset = (offset, )
+                offset = (offset,)
 
             if offset is not None:
                 assert len(shape) == len(
@@ -1182,7 +1174,7 @@ def ndarray(cls, n, m, dtype, shape, layout=Layout.AOS):
                 >>> x = ti.Matrix.ndarray(4, 5, ti.f32, shape=(16, 8))
         """
         if isinstance(shape, numbers.Number):
-            shape = (shape, )
+            shape = (shape,)
         return MatrixNdarray(n, m, dtype, shape, layout)
 
     @classmethod
@@ -1202,7 +1194,7 @@ def _Vector_ndarray(cls, n, dtype, shape, layout=Layout.AOS):
                 >>> x = ti.Vector.ndarray(3, ti.f32, shape=(16, 8))
         """
         if isinstance(shape, numbers.Number):
-            shape = (shape, )
+            shape = (shape,)
         return VectorNdarray(n, dtype, shape, layout)
 
     @staticmethod
@@ -1392,6 +1384,7 @@ class _IntermediateMatrix(Matrix):
         m (int): Number of columns of the matrix.
         entries (List[Expr]): All entries of the matrix.
     """
+
     def __init__(self, n, m, entries):
         assert isinstance(entries, list)
         assert n * m == len(entries), "Number of entries doesn't match n * m"
@@ -1411,6 +1404,7 @@ class _MatrixFieldElement(_IntermediateMatrix):
         field (MatrixField): The matrix field.
         indices (taichi_core.ExprGroup): Indices of the element.
     """
+
     def __init__(self, field, indices):
         super().__init__(field.n, field.m, [
             expr.Expr(ti_core.subscript(e.ptr, indices))
@@ -1427,6 +1421,7 @@ class MatrixField(Field):
         n (Int): Number of rows.
         m (Int): Number of columns.
     """
+
     def __init__(self, _vars, n, m):
         assert len(_vars) == n * m
         super().__init__(_vars)
@@ -1472,7 +1467,7 @@ def _calc_dynamic_index_stride(self):
                        i + 1]._offset_bytes_in_parent_cell for path in paths):
                 return
         stride = paths[1][depth_below_lca]._offset_bytes_in_parent_cell - \
-            paths[0][depth_below_lca]._offset_bytes_in_parent_cell
+                 paths[0][depth_below_lca]._offset_bytes_in_parent_cell
         for i in range(2, num_members):
             if stride != paths[i][depth_below_lca]._offset_bytes_in_parent_cell \
                     - paths[i - 1][depth_below_lca]._offset_bytes_in_parent_cell:
@@ -1493,7 +1488,7 @@ def fill(self, val):
         elif isinstance(val,
                         (list, tuple)) and isinstance(val[0], numbers.Number):
             assert self.m == 1
-            val = tuple([(v, ) for v in val])
+            val = tuple([(v,) for v in val])
         elif isinstance(val, Matrix):
             val_tuple = []
             for i in range(val.n):
@@ -1525,7 +1520,7 @@ def to_numpy(self, keep_dims=False, dtype=None):
         if dtype is None:
             dtype = to_numpy_type(self.dtype)
         as_vector = self.m == 1 and not keep_dims
-        shape_ext = (self.n, ) if as_vector else (self.n, self.m)
+        shape_ext = (self.n,) if as_vector else (self.n, self.m)
         arr = np.zeros(self.shape + shape_ext, dtype=dtype)
         from taichi._kernels import matrix_to_ext_arr  # pylint: disable=C0415
         matrix_to_ext_arr(self, arr, as_vector)
@@ -1545,7 +1540,7 @@ def to_torch(self, device=None, keep_dims=False):
         """
         import torch  # pylint: disable=C0415
         as_vector = self.m == 1 and not keep_dims
-        shape_ext = (self.n, ) if as_vector else (self.n, self.m)
+        shape_ext = (self.n,) if as_vector else (self.n, self.m)
         # pylint: disable=E1101
         arr = torch.empty(self.shape + shape_ext,
                           dtype=to_pytorch_type(self.dtype),
@@ -1568,7 +1563,7 @@ def to_paddle(self, place=None, keep_dims=False):
         """
         import paddle  # pylint: disable=C0415
         as_vector = self.m == 1 and not keep_dims
-        shape_ext = (self.n, ) if as_vector else (self.n, self.m)
+        shape_ext = (self.n,) if as_vector else (self.n, self.m)
         # pylint: disable=E1101
         # paddle.empty() doesn't support argument `place``
         arr = paddle.to_tensor(paddle.empty(self.shape + shape_ext,
@@ -1688,6 +1683,7 @@ class MatrixNdarray(Ndarray):
 
         >>> arr = ti.MatrixNdarray(2, 2, ti.f32, shape=(3, 3), layout=Layout.SOA)
     """
+
     def __init__(self, n, m, dtype, shape, layout):
         self.n = n
         self.m = m
@@ -1725,7 +1721,7 @@ def __setitem__(self, key, value):
     @python_scope
     def __getitem__(self, key):
         key = () if key is None else (
-            key, ) if isinstance(key, numbers.Number) else tuple(key)
+            key,) if isinstance(key, numbers.Number) else tuple(key)
         return Matrix(
             [[NdarrayHostAccess(self, key, (i, j)) for j in range(self.m)]
              for i in range(self.n)])
@@ -1786,13 +1782,14 @@ class VectorNdarray(Ndarray):
 
         >>> a = ti.VectorNdarray(3, ti.f32, (3, 3), layout=Layout.SOA)
     """
+
     def __init__(self, n, dtype, shape, layout):
         self.n = n
         super().__init__()
         self.dtype = cook_dtype(dtype)
         self.layout = layout
         self.shape = tuple(shape)
-        self.element_type = TensorType((n, ), self.dtype)
+        self.element_type = TensorType((n,), self.dtype)
         # TODO: pass in element_type, shape, layout directly
         self.arr = impl.get_runtime().prog.create_ndarray(
             self.element_type.dtype, shape, self.element_type.shape, layout)
@@ -1819,9 +1816,9 @@ def __setitem__(self, key, value):
     @python_scope
     def __getitem__(self, key):
         key = () if key is None else (
-            key, ) if isinstance(key, numbers.Number) else tuple(key)
+            key,) if isinstance(key, numbers.Number) else tuple(key)
         return Vector(
-            [NdarrayHostAccess(self, key, (i, )) for i in range(self.n)])
+            [NdarrayHostAccess(self, key, (i,)) for i in range(self.n)])
 
     @python_scope
     def to_numpy(self):