[lang] Migrate TensorType expansion for subscription indices from Pyt…

…hon to Frontend IR (#6942)

Issue: #5819

### Brief Summary
For indices of TensorType, instead of scalarizing them at Python level,
it is up to the Frontend IR's consumer to decide whether TensorType'd
indices are acceptable and if we should have it scalarized.

This PR removes `expand_expr` in Expression subscription and migrate the
scalarization logics to the following constructors:

1. MeshIndexConversionExpression::MeshIndexConversionExpression
2. IndexExpression::IndexExpression

python/taichi/lang/_texture.py

@@ -13,7 +13,8 @@ def _get_entries(mat):
     if isinstance(mat, Matrix):
         return mat.entries
     assert isinstance(mat, Expr) and mat.is_tensor()
-    return impl.get_runtime().prog.current_ast_builder().expand_expr([mat.ptr])
+    return impl.get_runtime().prog.current_ast_builder().expand_exprs(
+        [mat.ptr])
 
 
 class TextureSampler:

python/taichi/lang/any_array.py

@@ -78,15 +78,18 @@ def __init__(self, arr, indices_first):
 
     @taichi_scope
     def subscript(self, i, j):
+        ast_builder = impl.get_runtime().prog.current_ast_builder()
+
         indices_second = (i, ) if len(self.arr.element_shape()) == 1 else (i,
                                                                            j)
         if self.arr.layout() == Layout.SOA:
             indices = indices_second + self.indices_first
         else:
             indices = self.indices_first + indices_second
         return Expr(
-            _ti_core.subscript(self.arr.ptr, make_expr_group(*indices),
-                               impl.get_runtime().get_current_src_info()))
+            ast_builder.expr_subscript(
+                self.arr.ptr, make_expr_group(*indices),
+                impl.get_runtime().get_current_src_info()))
 
 
 __all__ = []

python/taichi/lang/ast/ast_transformer.py

@@ -15,7 +15,7 @@
                                                    ReturnStatus)
 from taichi.lang.ast.symbol_resolver import ASTResolver
 from taichi.lang.exception import (TaichiIndexError, TaichiSyntaxError,
-                                   TaichiTypeError)
+                                   TaichiTypeError, handle_exception_from_cpp)
 from taichi.lang.expr import Expr, make_expr_group
 from taichi.lang.field import Field
 from taichi.lang.matrix import Matrix, MatrixType, Vector, is_vector
@@ -156,7 +156,7 @@ def build_assign_unpack(ctx, node_target, values, is_static_assign):
                 raise ValueError(
                     'Matrices with more than one columns cannot be unpacked')
 
-            values = ctx.ast_builder.expand_expr([values.ptr])
+            values = ctx.ast_builder.expand_exprs([values.ptr])
             if len(values) == 1:
                 values = values[0]
 
@@ -302,7 +302,7 @@ def process_generators(ctx, node, now_comp, func, result):
         if isinstance(_iter, impl.Expr) and _iter.ptr.is_tensor():
             shape = _iter.ptr.get_shape()
             flattened = [
-                Expr(x) for x in ctx.ast_builder.expand_expr([_iter.ptr])
+                Expr(x) for x in ctx.ast_builder.expand_exprs([_iter.ptr])
             ]
             _iter = reshape_list(flattened, shape)
 
@@ -514,7 +514,7 @@ def build_Call(ctx, node):
                     # Expand Expr with Matrix-type return into list of Exprs
                     arg_list = [
                         Expr(x)
-                        for x in ctx.ast_builder.expand_expr([arg_list.ptr])
+                        for x in ctx.ast_builder.expand_exprs([arg_list.ptr])
                     ]
 
                 for i in arg_list:
@@ -730,7 +730,7 @@ def build_Return(ctx, node):
             elif isinstance(ctx.func.return_type, MatrixType):
                 values = node.value.ptr
                 if isinstance(values, Expr) and values.ptr.is_tensor():
-                    values = ctx.ast_builder.expand_expr([values.ptr])
+                    values = ctx.ast_builder.expand_exprs([values.ptr])
                 else:
                     assert isinstance(values, Matrix)
                     values = itertools.chain.from_iterable(values.to_list()) if\
@@ -819,12 +819,15 @@ def build_Attribute(ctx, node):
         # we continue to process it as a normal attribute node.
         try:
             build_stmt(ctx, node.value)
-        except TaichiIndexError as e:
-            node.value.ptr = None
-            if ASTTransformer.build_attribute_if_is_dynamic_snode_method(
-                    ctx, node):
-                return node.ptr
+        except Exception as e:
+            e = handle_exception_from_cpp(e)
+            if isinstance(e, TaichiIndexError):
+                node.value.ptr = None
+                if ASTTransformer.build_attribute_if_is_dynamic_snode_method(
+                        ctx, node):
+                    return node.ptr
             raise e
+
         if ASTTransformer.build_attribute_if_is_dynamic_snode_method(
                 ctx, node):
             return node.ptr
@@ -837,11 +840,11 @@ def build_Attribute(ctx, node):
                                                       node.attr)
                 attr_len = len(node.attr)
                 if attr_len == 1:
-                    node.ptr = Expr(
-                        _ti_core.subscript(
-                            node.value.ptr.ptr,
-                            make_expr_group(keygroup.index(node.attr)),
-                            impl.get_runtime().get_current_src_info()))
+                    node.ptr = Expr(impl.get_runtime(
+                    ).prog.current_ast_builder().expr_subscript(
+                        node.value.ptr.ptr,
+                        make_expr_group(keygroup.index(node.attr)),
+                        impl.get_runtime().get_current_src_info()))
                 else:
                     node.ptr = Expr(
                         _ti_core.subscript_with_multiple_indices(

python/taichi/lang/exception.py

@@ -56,6 +56,8 @@ def handle_exception_from_cpp(exc):
         return TaichiTypeError(str(exc))
     if isinstance(exc, core.TaichiSyntaxError):
         return TaichiSyntaxError(str(exc))
+    if isinstance(exc, core.TaichiIndexError):
+        return TaichiIndexError(str(exc))
     if isinstance(exc, core.TaichiAssertionError):
         return TaichiAssertionError(str(exc))
     return exc

python/taichi/lang/expr.py

@@ -170,7 +170,7 @@ def _get_flattened_ptrs(val):
             ptrs.extend(_get_flattened_ptrs(item))
         return ptrs
     if isinstance(val, Expr) and val.ptr.is_tensor():
-        return impl.get_runtime().prog.current_ast_builder().expand_expr(
+        return impl.get_runtime().prog.current_ast_builder().expand_exprs(
             [val.ptr])
     return [Expr(val).ptr]
 

python/taichi/lang/impl.py

@@ -9,9 +9,8 @@
 from taichi.lang._texture import RWTextureAccessor
 from taichi.lang.any_array import AnyArray
 from taichi.lang.enums import SNodeGradType
-from taichi.lang.exception import (TaichiCompilationError, TaichiIndexError,
-                                   TaichiRuntimeError, TaichiSyntaxError,
-                                   TaichiTypeError)
+from taichi.lang.exception import (TaichiCompilationError, TaichiRuntimeError,
+                                   TaichiSyntaxError, TaichiTypeError)
 from taichi.lang.expr import Expr, make_expr_group
 from taichi.lang.field import Field, ScalarField
 from taichi.lang.kernel_arguments import SparseMatrixProxy
@@ -132,6 +131,7 @@ def check_validity(x):
 
 @taichi_scope
 def subscript(ast_builder, value, *_indices, skip_reordered=False):
+    ast_builder = get_runtime().prog.current_ast_builder()
     # Directly evaluate in Python for non-Taichi types
     if not isinstance(
             value,
@@ -150,9 +150,6 @@ def subscript(ast_builder, value, *_indices, skip_reordered=False):
         elif isinstance(_index, slice):
             ind = [_index]
             has_slice = True
-        elif isinstance(_index, Expr) and _index.is_tensor():
-            # Expand Expr with TensorType return
-            ind = [Expr(e) for e in ast_builder.expand_expr([_index.ptr])]
         else:
             ind = [_index]
         flattened_indices += ind
@@ -167,7 +164,6 @@ def subscript(ast_builder, value, *_indices, skip_reordered=False):
                 f"The type {type(value)} do not support index of slice type")
     else:
         indices_expr_group = make_expr_group(*indices)
-        index_dim = indices_expr_group.size()
 
     if isinstance(value, SharedArray):
         return value.subscript(*indices)
@@ -178,13 +174,13 @@ def subscript(ast_builder, value, *_indices, skip_reordered=False):
     if isinstance(value,
                   (MeshReorderedScalarFieldProxy,
                    MeshReorderedMatrixFieldProxy)) and not skip_reordered:
-        assert index_dim == 1
+
         reordered_index = tuple([
             Expr(
-                _ti_core.get_index_conversion(value.mesh_ptr,
-                                              value.element_type,
-                                              Expr(indices[0]).ptr,
-                                              ConvType.g2r))
+                ast_builder.mesh_index_conversion(value.mesh_ptr,
+                                                  value.element_type,
+                                                  Expr(indices[0]).ptr,
+                                                  ConvType.g2r))
         ])
         return subscript(ast_builder,
                          value,
@@ -203,29 +199,26 @@ def subscript(ast_builder, value, *_indices, skip_reordered=False):
                 raise RuntimeError(
                     f"Gradient {_var.get_expr_name()} has not been placed, check whether `needs_grad=True`"
                 )
-        field_dim = snode.num_active_indices()
-        if field_dim != index_dim:
-            raise TaichiIndexError(
-                f'Field with dim {field_dim} accessed with indices of dim {index_dim}'
-            )
+
         if isinstance(value, MatrixField):
-            return make_index_expr(value.ptr, indices_expr_group)
+            return Expr(
+                ast_builder.expr_subscript(
+                    value.ptr, indices_expr_group,
+                    get_runtime().get_current_src_info()))
         if isinstance(value, StructField):
             entries = {
                 k: subscript(ast_builder, v, *indices)
                 for k, v in value._items
             }
             entries['__struct_methods'] = value.struct_methods
             return _IntermediateStruct(entries)
-        return make_index_expr(_var, indices_expr_group)
+        return Expr(
+            ast_builder.expr_subscript(_var, indices_expr_group,
+                                       get_runtime().get_current_src_info()))
     if isinstance(value, AnyArray):
-        dim = _ti_core.get_external_tensor_dim(value.ptr)
-        element_dim = len(value.element_shape())
-        if dim != index_dim + element_dim:
-            raise IndexError(
-                f'Field with dim {dim - element_dim} accessed with indices of dim {index_dim}'
-            )
-        return make_index_expr(value.ptr, indices_expr_group)
+        return Expr(
+            ast_builder.expr_subscript(value.ptr, indices_expr_group,
+                                       get_runtime().get_current_src_info()))
     assert isinstance(value, Expr)
     # Index into TensorType
     # value: IndexExpression with ret_type = TensorType
@@ -249,18 +242,14 @@ def subscript(ast_builder, value, *_indices, skip_reordered=False):
                 make_expr_group(i, j) for i in indices[0] for j in indices[1]
             ]
             return_shape = (len(indices[0]), len(indices[1]))
+
         return Expr(
             _ti_core.subscript_with_multiple_indices(
                 value.ptr, multiple_indices, return_shape,
                 get_runtime().get_current_src_info()))
-    return make_index_expr(value.ptr, indices_expr_group)
-
-
-@taichi_scope
-def make_index_expr(_var, indices_expr_group):
     return Expr(
-        _ti_core.subscript(_var, indices_expr_group,
-                           get_runtime().get_current_src_info()))
+        ast_builder.expr_subscript(value.ptr, indices_expr_group,
+                                   get_runtime().get_current_src_info()))
 
 
 class SrcInfoGuard:

python/taichi/lang/kernel_impl.py

@@ -262,7 +262,7 @@ def func_call_rvalue(self, key, args):
                                 impl.Expr) and args[i].ptr.is_tensor():
                     non_template_args.extend([
                         Expr(x) for x in impl.get_runtime().prog.
-                        current_ast_builder().expand_expr([args[i].ptr])
+                        current_ast_builder().expand_exprs([args[i].ptr])
                     ])
                 else:
                     non_template_args.append(args[i])

python/taichi/lang/matrix.py

@@ -1503,7 +1503,7 @@ def __call__(self, *args):
             elif isinstance(x, impl.Expr) and x.ptr.is_tensor():
                 entries += [
                     impl.Expr(e) for e in impl.get_runtime().prog.
-                    current_ast_builder().expand_expr([x.ptr])
+                    current_ast_builder().expand_exprs([x.ptr])
                 ]
             elif isinstance(x, Matrix):
                 entries += x.entries
@@ -1616,7 +1616,7 @@ def __call__(self, *args):
             elif isinstance(x, impl.Expr) and x.ptr.is_tensor():
                 entries += [
                     impl.Expr(e) for e in impl.get_runtime().prog.
-                    current_ast_builder().expand_expr([x.ptr])
+                    current_ast_builder().expand_exprs([x.ptr])
                 ]
             else:
                 entries.append(x)

python/taichi/lang/mesh.py

@@ -605,14 +605,17 @@ def _TetMesh():
 class MeshElementFieldProxy:
     def __init__(self, mesh: MeshInstance, element_type: MeshElementType,
                  entry_expr: impl.Expr):
+        ast_builder = impl.get_runtime().prog.current_ast_builder()
+
         self.mesh = mesh
         self.element_type = element_type
         self.entry_expr = entry_expr
 
         element_field = self.mesh.fields[self.element_type]
         for key, attr in element_field.field_dict.items():
+
             global_entry_expr = impl.Expr(
-                _ti_core.get_index_conversion(
+                ast_builder.mesh_index_conversion(
                     self.mesh.mesh_ptr, element_type, entry_expr,
                     ConvType.l2r if element_field.attr_dict[key].reorder else
                     ConvType.l2g))  # transform index space
@@ -622,7 +625,7 @@ def __init__(self, mesh: MeshInstance, element_type: MeshElementType,
                 setattr(
                     self, key,
                     impl.Expr(
-                        _ti_core.subscript(
+                        ast_builder.expr_subscript(
                             attr.ptr, global_entry_expr_group,
                             impl.get_runtime().get_current_src_info())))
             elif isinstance(attr, StructField):
@@ -633,7 +636,7 @@ def __init__(self, mesh: MeshInstance, element_type: MeshElementType,
                 setattr(
                     self, key,
                     impl.Expr(
-                        _ti_core.subscript(
+                        ast_builder.expr_subscript(
                             var, global_entry_expr_group,
                             impl.get_runtime().get_current_src_info())))
 
@@ -650,10 +653,11 @@ def ptr(self):
 
     @property
     def id(self):  # return the global non-reordered index
+        ast_builder = impl.get_runtime().prog.current_ast_builder()
         l2g_expr = impl.Expr(
-            _ti_core.get_index_conversion(self.mesh.mesh_ptr,
-                                          self.element_type, self.entry_expr,
-                                          ConvType.l2g))
+            ast_builder.mesh_index_conversion(self.mesh.mesh_ptr,
+                                              self.element_type,
+                                              self.entry_expr, ConvType.l2g))
         return l2g_expr
 
 

python/taichi/lang/simt/block.py

@@ -54,5 +54,8 @@ def __init__(self, shape, dtype):
 
     @taichi_scope
     def subscript(self, *indices):
-        return impl.make_index_expr(self.shared_array_proxy,
-                                    make_expr_group(*indices))
+        ast_builder = impl.get_runtime().prog.current_ast_builder()
+        return impl.Expr(
+            ast_builder.expr_subscript(
+                self.shared_array_proxy, make_expr_group(*indices),
+                impl.get_runtime().get_current_src_info()))

taichi/common/exceptions.h

@@ -23,6 +23,10 @@ class TaichiSyntaxError : public TaichiExceptionImpl {
   using TaichiExceptionImpl::TaichiExceptionImpl;
 };
 
+class TaichiIndexError : public TaichiExceptionImpl {
+  using TaichiExceptionImpl::TaichiExceptionImpl;
+};
+
 class TaichiRuntimeError : public TaichiExceptionImpl {
   using TaichiExceptionImpl::TaichiExceptionImpl;
 };

taichi/ir/expr.cpp

@@ -26,12 +26,6 @@ Expr bit_cast(const Expr &input, DataType dt) {
   return Expr::make<UnaryOpExpression>(UnaryOpType::cast_bits, input, dt);
 }
 
-Expr Expr::operator[](const ExprGroup &indices) const {
-  TI_ASSERT(is<FieldExpression>() || is<MatrixFieldExpression>() ||
-            is<ExternalTensorExpression>() || is_tensor(expr->ret_type));
-  return Expr::make<IndexExpression>(*this, indices);
-}
-
 Expr &Expr::operator=(const Expr &o) {
   set(o);
   return *this;

taichi/ir/expr.h

@@ -83,8 +83,6 @@ class Expr {
   // std::variant<Expr, std::string> in FrontendPrintStmt.
   Expr &operator=(const Expr &o);
 
-  Expr operator[](const ExprGroup &indices) const;
-
   template <typename T, typename... Args>
   static Expr make(Args &&...args) {
     return Expr(std::make_shared<T>(std::forward<Args>(args)...));