This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new af6ae2afaf [REFACTOR][IR] Make expression subscription eager and
remove SubscriptProxy (#20257)
af6ae2afaf is described below
commit af6ae2afafcd947a5f09324d30932415a9ed87a6
Author: Tianqi Chen <[email protected]>
AuthorDate: Tue Sep 1 13:07:19 2026 -0400
[REFACTOR][IR] Make expression subscription eager and remove SubscriptProxy
(#20257)
Lazy subscription required every expression consumer to remember to
realize a proxy, so missed boundaries failed far from the original
subscript. Realize subscriptions immediately through the existing
type-directed hooks and remove SubscriptProxy plus the consumer-side
realization paths it required.
This deliberately removes chained BufferRegion subscription: use A[i, j]
instead of A[i][j]. Tuple bounds now raise IndexError directly from
TupleGetItem so Python sequence iteration retains its existing behavior.
Validation:
- clean RelWithDebInfo build with LLVM and CUDA enabled
- affected Python suites: 560 passed, 5 skipped
- final parser cleanup suite: 201 passed, 1 skipped
- pre-commit run --all-files
---
python/tvm/backend/cuda/ptx/engine.py | 6 +-
python/tvm/ir/__init__.py | 1 -
python/tvm/ir/expr.py | 178 +++------------------
python/tvm/script/parser/core/evaluator.py | 20 +--
python/tvm/script/parser/core/parser.py | 13 +-
python/tvm/tirx/_buffer_view.py | 4 +-
python/tvm/tirx/function.py | 5 -
python/tvm/tirx/op.py | 2 -
python/tvm/tirx/script/builder/ir.py | 10 +-
python/tvm/tirx/script/builder/tirx.py | 5 +-
python/tvm/tirx/script/parser/operation.py | 5 -
python/tvm/tirx/script/parser/parser.py | 2 -
python/tvm/tirx/script/tile.py | 2 -
python/tvm/tirx/stmt.py | 37 +----
src/ir/expr.cc | 6 +-
src/ir/{subscript_proxy.cc => subscript_expr.cc} | 10 +-
tests/python/relax/test_blockbuilder_core.py | 3 +
.../test_s_tir_analysis_identify_memcpy.py | 6 +-
tests/python/tirx/test_parser_printer.py | 31 ++--
19 files changed, 60 insertions(+), 286 deletions(-)
diff --git a/python/tvm/backend/cuda/ptx/engine.py
b/python/tvm/backend/cuda/ptx/engine.py
index 5ffdecc785..3897b22991 100644
--- a/python/tvm/backend/cuda/ptx/engine.py
+++ b/python/tvm/backend/cuda/ptx/engine.py
@@ -41,7 +41,6 @@ from tvm.backend.cuda.codegen.registry import register_codegen
from tvm.backend.cuda.codegen.utils import parse_str
from tvm.backend.cuda.op import cuda_cvta_generic_to_shared, cuda_func_call
from tvm.ir import Call, TensorLoad
-from tvm.ir.expr import _realize_operand
from tvm.ir.op import register_op_attr
from tvm.ir.type import PointerType, PrimType
from tvm.runtime import const
@@ -444,7 +443,7 @@ def _coerce_pred_operand(entry, slot, values):
like any other, and no syntax line offers a non-predicate alternative at
the same position.
"""
- (value,) = [_realize_operand(value) for value in values]
+ (value,) = values
if slot.rw != "r":
# The 0/1 materialization of a .pred result: a "=r" uint32 the caller
# receives through a reference parameter, so it needs a writable
@@ -457,7 +456,7 @@ def _coerce_pred_operand(entry, slot, values):
)
return values
if isinstance(value, PredArg):
- value = _realize_operand(getattr(value.value, "scalar", value.value))
+ value = getattr(value.value, "scalar", value.value)
if isinstance(value, bool | int):
return [const(int(value), "uint32")]
ty = getattr(value, "ty", None)
@@ -482,7 +481,6 @@ def _coerce_pred_operand(entry, slot, values):
def _coerce_typed(entry, slot, values, mod_map):
"""Coerce a dtype-carrying register operand (``rw`` any)."""
- values = [_realize_operand(value) for value in values]
allowed = operand_dtypes(slot, mod_map)
token = operand_type(slot, mod_map)
if slot.rw in ("w", "rw"):
diff --git a/python/tvm/ir/__init__.py b/python/tvm/ir/__init__.py
index 412bd89061..977ec90dcc 100644
--- a/python/tvm/ir/__init__.py
+++ b/python/tvm/ir/__init__.py
@@ -42,7 +42,6 @@ from .expr import (
GlobalVar,
OpaqueExpr,
Range,
- SubscriptProxy,
TensorLoad,
Tuple,
TupleGetItem,
diff --git a/python/tvm/ir/expr.py b/python/tvm/ir/expr.py
index 5005737e26..7bcddca35e 100644
--- a/python/tvm/ir/expr.py
+++ b/python/tvm/ir/expr.py
@@ -22,11 +22,26 @@ import tvm_ffi
import tvm
-from ..runtime import Object, ObjectConvertible, Scriptable
+from ..runtime import Object, Scriptable
from . import _ffi_api, _overload_prim_expr, _tensor_expr_overload
from .base import Node, Span
+def _convert_subscript_index(index):
+ """Convert Python indexing syntax into an FFI subscript descriptor."""
+
+ def convert(value):
+ if value is None or is_prim_expr(value):
+ return value
+ return tvm.tirx.const(value)
+
+ if isinstance(index, slice):
+ return (convert(index.start), convert(index.stop), convert(index.step))
+ if index is Ellipsis or index is None:
+ raise TypeError("Ellipsis and newaxis are not supported in expression
subscriptions")
+ return convert(index)
+
+
@tvm_ffi.register_object("ir.Expr")
class Expr(Node):
"""Base class of all the expressions."""
@@ -35,29 +50,15 @@ class Expr(Node):
ty: "tvm.ir.Type"
def __getitem__(self, index):
- if isinstance(self.ty, tvm.ir.TupleType):
- # Tuple subscription is eager so Python's legacy sequence protocol
- # observes IndexError and terminates tuple iteration/unpacking.
- try:
- return _ffi_api.SubscriptExprRealize(
- self, [SubscriptProxy._convert_index(index)], None
- )
- except RuntimeError as err:
- if "Index out of bounds" in err.args[0]:
- raise IndexError from err
- raise
if self.ty.is_missing():
# Preserve Relax's pre-normalization tuple access: operator calls
# have a missing result type until the block builder infers it.
return TupleGetItem(self, index)
- if isinstance(self.ty, tvm.relax.TensorType):
- # Relax tensor subscription retains its legacy TupleGetItem
- # semantics. Realize each step eagerly so x[i][j] remains two
- # nested tuple-item accesses rather than one multi-axis slice.
- return SubscriptProxy(self, index).to_expr()
- if is_prim_expr(self):
- raise TypeError("A primitive-valued expression cannot be indexed")
- return SubscriptProxy(self, index)
+
+ indices = tuple(index) if isinstance(index, tuple | list) else (index,)
+ return _ffi_api.SubscriptExprRealize(
+ self, [_convert_subscript_index(item) for item in indices], None
+ )
@tvm_ffi.register_object("ir.OpaqueExpr")
@@ -108,8 +109,6 @@ class GlobalVar(Expr):
"""
from .type import PointerType
- args = tuple(_realize_operand(arg) for arg in args)
-
def is_tir_arg(x):
return (
isinstance(x, Number)
@@ -127,140 +126,115 @@ class GlobalVar(Expr):
raise RuntimeError(f"Do not know how to handle GlobalVar.__call__ for
types {arg_types}")
-def _realize_operand(value):
- return value._operand() if isinstance(value, ExprOperand) else value
-
-
class ExprOperand:
"""Python operator surface for anything that denotes an expression."""
__slots__ = ()
__hash__ = object.__hash__
- def _operand(self) -> Expr:
- raise NotImplementedError
-
def expr_ty(self):
"""Return this expression's primitive result type."""
- self = _realize_operand(self)
if is_prim_expr(self):
return self.ty
raise TypeError(f"Expected a primitive-valued expression, but result
type is {self.ty}")
def __add__(self, other):
- self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__add__(self, other)
result = _tensor_expr_overload.__add__(self, other)
return result
def __radd__(self, other):
- self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__radd__(self, other)
result = _tensor_expr_overload.__radd__(self, other)
return result
def __sub__(self, other):
- self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__sub__(self, other)
result = _tensor_expr_overload.__sub__(self, other)
return result
def __rsub__(self, other):
- self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__rsub__(self, other)
result = _tensor_expr_overload.__rsub__(self, other)
return result
def __mul__(self, other):
- self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__mul__(self, other)
result = _tensor_expr_overload.__mul__(self, other)
return result
def __rmul__(self, other):
- self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__rmul__(self, other)
result = _tensor_expr_overload.__rmul__(self, other)
return result
def __div__(self, other):
- self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__div__(self, other)
result = _tensor_expr_overload.__div__(self, other)
return result
def __rdiv__(self, other):
- self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__rdiv__(self, other)
result = _tensor_expr_overload.__rdiv__(self, other)
return result
def __truediv__(self, other):
- self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__truediv__(self, other)
result = _tensor_expr_overload.__truediv__(self, other)
return result
def __rtruediv__(self, other):
- self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__rtruediv__(self, other)
result = _tensor_expr_overload.__rtruediv__(self, other)
return result
def __floordiv__(self, other):
- self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__floordiv__(self, other)
result = _tensor_expr_overload.__floordiv__(self, other)
return result
def __rfloordiv__(self, other):
- self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__rfloordiv__(self, other)
result = _tensor_expr_overload.__rfloordiv__(self, other)
return result
def __mod__(self, other):
- self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__mod__(self, other)
result = _tensor_expr_overload.__mod__(self, other)
return result
def __rmod__(self, other):
- self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__rmod__(self, other)
result = _tensor_expr_overload.__rmod__(self, other)
return result
def __pow__(self, other):
- self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return NotImplemented
result = _tensor_expr_overload.__pow__(self, other)
return result
def __rpow__(self, other):
- self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return NotImplemented
result = _tensor_expr_overload.__rpow__(self, other)
return result
def __neg__(self):
- self = _realize_operand(self)
if is_prim_expr(self):
result = _overload_prim_expr.__neg__(self)
if result is NotImplemented:
@@ -272,67 +246,56 @@ class ExprOperand:
return result
def __lshift__(self, other):
- self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__lshift__(self, other)
return NotImplemented
def __rlshift__(self, other):
- self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__rlshift__(self, other)
return NotImplemented
def __rshift__(self, other):
- self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__rshift__(self, other)
return NotImplemented
def __rrshift__(self, other):
- self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__rrshift__(self, other)
return NotImplemented
def __and__(self, other):
- self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__and__(self, other)
return NotImplemented
def __rand__(self, other):
- self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__rand__(self, other)
return NotImplemented
def __or__(self, other):
- self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__or__(self, other)
return NotImplemented
def __ror__(self, other):
- self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__ror__(self, other)
return NotImplemented
def __xor__(self, other):
- self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__xor__(self, other)
return NotImplemented
def __rxor__(self, other):
- self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__rxor__(self, other)
return NotImplemented
def __invert__(self):
- self = _realize_operand(self)
if is_prim_expr(self):
result = _overload_prim_expr.__invert__(self)
if result is NotImplemented:
@@ -341,40 +304,34 @@ class ExprOperand:
raise TypeError(f"Operator overloading is not supported for expression
type {self.ty}")
def __lt__(self, other):
- self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__lt__(self, other)
result = _tensor_expr_overload.__lt__(self, other)
return result
def __le__(self, other):
- self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__le__(self, other)
result = _tensor_expr_overload.__le__(self, other)
return result
def __eq__(self, other):
- self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__eq__(self, other)
return Object.__eq__(self, other)
def __ne__(self, other):
- self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__ne__(self, other)
return Object.__ne__(self, other)
def __gt__(self, other):
- self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__gt__(self, other)
result = _tensor_expr_overload.__gt__(self, other)
return result
def __ge__(self, other):
- self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__ge__(self, other)
result = _tensor_expr_overload.__ge__(self, other)
@@ -390,7 +347,6 @@ class ExprOperand:
return self.__nonzero__()
def equal(self, other, span=None):
- self, other = _realize_operand(self), _realize_operand(other)
if not is_prim_expr(self):
raise TypeError(f"Operator overloading is not supported for
expression type {self.ty}")
result = _overload_prim_expr.equal(self, other, span)
@@ -399,7 +355,6 @@ class ExprOperand:
return result
def astype(self, dtype, span=None):
- self = _realize_operand(self)
if is_prim_expr(self):
result = _overload_prim_expr.astype(self, dtype, span)
if result is NotImplemented:
@@ -417,7 +372,6 @@ class _ExprCallable:
__slots__ = ()
def __call__(self, *args, attrs=None):
- self, args = _realize_operand(self), tuple(_realize_operand(arg) for
arg in args)
if is_prim_expr(self):
raise TypeError("A primitive-valued expression cannot be called")
result = _tensor_expr_overload.__call__(self, *args, attrs=attrs)
@@ -431,101 +385,11 @@ class ExprWithOp(ExprOperand, Expr, Scriptable):
__hash__ = Expr.__hash__
- def _operand(self) -> Expr:
- return self
-
class _CallableExprWithOp(_ExprCallable, ExprWithOp):
"""Common operator behavior for expression nodes that support function
calls."""
-class SubscriptProxy(_ExprCallable, ExprOperand, ObjectConvertible):
- """An immutable, lazily-realized subscription of an :class:`Expr`.
-
- Point subscriptions may be chained to accumulate dimensions. A proxy
- containing a slice must be realized before applying a region subscript.
- """
-
- __slots__ = ("_result", "_slice", "_source", "_span")
- __hash__ = object.__hash__
-
- def __init__(self, source: Expr, index):
- if isinstance(source, SubscriptProxy):
- if any(isinstance(item, slice) for item in source._slice):
- raise TypeError("Cannot chain a subscription after a slice")
- self._source = source._source
- self._slice = source._slice + self._flatten(index)
- self._span = source._span
- else:
- _ffi_api.SubscriptExprCheck(source)
- self._source = source
- self._slice = self._flatten(index)
- self._span = None
- self._result = None
-
- def with_span(self, span: Span) -> "SubscriptProxy":
- """Return an unrealized proxy carrying its frontend source span."""
- result = object.__new__(SubscriptProxy)
- result._source = self._source
- result._slice = self._slice
- result._span = span
- result._result = None
- return result
-
- @staticmethod
- def _flatten(index):
- return tuple(index) if isinstance(index, tuple | list) else (index,)
-
- @staticmethod
- def _convert_index(index):
- def convert(value):
- value = _realize_operand(value)
- if value is None or is_prim_expr(value):
- return value
- return tvm.tirx.const(value)
-
- if isinstance(index, slice):
- return (convert(index.start), convert(index.stop),
convert(index.step))
- if index is Ellipsis or index is None:
- raise TypeError("Ellipsis and newaxis are not supported in
expression subscriptions")
- return convert(index)
-
- def _operand(self) -> Object:
- return self.to_expr()
-
- def __getitem__(self, index):
- return SubscriptProxy(self, index)
-
- def to_expr(self) -> Object:
- """Realize and cache the subscribed IR object."""
- if self._result is None:
- result = _ffi_api.SubscriptExprRealize(
- self._source,
- [self._convert_index(index) for index in self._slice],
- self._span,
- )
- if not isinstance(result, Object):
- raise TypeError("__subscript_expr_realize__ must return an
Object")
- self._result = result
- return self._result
-
- def asobject(self) -> Object:
- return self.to_expr()
-
- def __tvm_ffi_object__(self) -> Object:
- return self.to_expr()
-
- @property
- def ty(self):
- return self.to_expr().ty
-
- def same_as(self, other):
- return self.to_expr().same_as(_realize_operand(other))
-
- def __repr__(self):
- return f"SubscriptProxy({self._source!r}, {self._slice!r})"
-
-
@tvm_ffi.register_object("ir.Tuple")
class Tuple(_CallableExprWithOp):
"""Tuple expression that groups several fields together.
diff --git a/python/tvm/script/parser/core/evaluator.py
b/python/tvm/script/parser/core/evaluator.py
index e22abf4b9a..71a1b3bc8f 100644
--- a/python/tvm/script/parser/core/evaluator.py
+++ b/python/tvm/script/parser/core/evaluator.py
@@ -131,7 +131,7 @@ class ExprEvaluator:
return result.value
raise TypeError(f"Unexpected result type: {type(result)}")
- def _add_intermediate_result(self, value: Any, node: doc.AST) -> doc.Name:
+ def _add_intermediate_result(self, value: Any) -> doc.Name:
"""Add intermediate result during evaluation into value table.
Parameters
@@ -139,16 +139,13 @@ class ExprEvaluator:
value : Any
The intermediate result.
- node : doc.AST
- The AST node that produced the intermediate result.
-
Returns
-------
name : doc.Name
The doc AST name node with intermediate name for intermediate
result.
"""
if self.parser is not None:
- value = self.parser.annotate_current_source_span(value, node)
+ value = self.parser.annotate_current_source_span(value)
name = f"__tvm_tmp_value_{self.new_value_count}"
self.new_value_count += 1
self.value_table[name] = value
@@ -196,13 +193,13 @@ class ExprEvaluator:
value = self._eval_bool_op(node)
except Exception as err: # pylint: disable=broad-except
self.parser.report_error(node, err)
- return self._add_intermediate_result(value, node)
+ return self._add_intermediate_result(value)
if isinstance(node, doc.IfExp):
try:
value = self._eval_if_exp(node)
except Exception as err: # pylint: disable=broad-except
self.parser.report_error(node, err)
- return self._add_intermediate_result(value, node)
+ return self._add_intermediate_result(value)
args = []
if (
@@ -265,7 +262,7 @@ class ExprEvaluator:
if isinstance(node, doc.ListComp | doc.SetComp | doc.DictComp):
value = self._eval_expr(node)
- return self._add_intermediate_result(value, node)
+ return self._add_intermediate_result(value)
fields = {}
for field in node.__class__._FIELDS: # pylint:
disable=protected-access
@@ -287,7 +284,7 @@ class ExprEvaluator:
value = self._eval_expr(node.__class__(**fields))
except Exception as err: # pylint: disable=broad-except
self.parser.report_error(node, err)
- return self._add_intermediate_result(value, node)
+ return self._add_intermediate_result(value)
def _eval_lambda(self, node: doc.Lambda) -> Any:
"""The doc AST lambda node evaluating method.
@@ -306,7 +303,7 @@ class ExprEvaluator:
value = self._eval_expr(node)
except Exception as err: # pylint: disable=broad-except
self.parser.report_error(node, err)
- return self._add_intermediate_result(value, node)
+ return self._add_intermediate_result(value)
def _eval_bool_op(self, node: doc.BoolOp) -> Any:
"""The doc AST boolean operator node evaluating method.
@@ -416,9 +413,6 @@ class ExprEvaluator:
The evaluation result.
"""
test = self._eval_expr(self._visit(node.test))
- from tvm.ir.expr import _realize_operand # pylint:
disable=import-outside-toplevel
-
- test = _realize_operand(test)
if isinstance(test, bool):
selected = node.body if test else node.orelse
return self._eval_expr(self._visit(selected))
diff --git a/python/tvm/script/parser/core/parser.py
b/python/tvm/script/parser/core/parser.py
index 94c74f24a1..6319eacb34 100644
--- a/python/tvm/script/parser/core/parser.py
+++ b/python/tvm/script/parser/core/parser.py
@@ -596,12 +596,8 @@ class Parser(doc.NodeVisitor):
with
IRBuilder.current().with_source_span(self.diag.source.to_span(node)):
yield
- def annotate_current_source_span(self, value: Any, node: doc.AST | None =
None) -> Any:
+ def annotate_current_source_span(self, value: Any) -> Any:
"""Attach the active parser span to an expression result, when
applicable."""
- from tvm.ir.expr import SubscriptProxy # pylint:
disable=import-outside-toplevel
-
- if isinstance(value, SubscriptProxy) and node is not None:
- return value.with_span(self.diag.source.to_span(node))
if isinstance(value, Object) and IRBuilder.is_in_scope():
return IRBuilder.current()._set_current_source_span(value) #
pylint: disable=protected-access
return value
@@ -633,12 +629,7 @@ class Parser(doc.NodeVisitor):
var_values[ScriptMacro.parser_object_name] = self
value = eval_expr(self, node, var_values)
- # Subscription proxies are an expression-construction detail. Keep
- # them lazy while evaluating a compound Python expression, then
- # normalize once at the parser boundary before statement dispatch.
- from tvm.ir.expr import _realize_operand # pylint:
disable=import-outside-toplevel
-
- return self.annotate_current_source_span(_realize_operand(value), node)
+ return self.annotate_current_source_span(value)
def _duplicate_lhs_check(self, target: doc.expr) -> bool | set[str]:
"""Check whether duplicate lhs exists in assignment.
diff --git a/python/tvm/tirx/_buffer_view.py b/python/tvm/tirx/_buffer_view.py
index 3977cf2f3a..90c4a13a61 100644
--- a/python/tvm/tirx/_buffer_view.py
+++ b/python/tvm/tirx/_buffer_view.py
@@ -572,9 +572,7 @@ class ChunkIndexer:
else:
size = buf.shape[dim] // int(count)
translated.append(slice(pick * size, (pick + 1) * size))
- # ``chunk`` is a view helper whose documented result is a region, not
- # the lazy subscription proxy used by ordinary expression indexing.
- return buf[tuple(translated)].to_expr()
+ return buf[tuple(translated)]
class TileIndexer:
diff --git a/python/tvm/tirx/function.py b/python/tvm/tirx/function.py
index 7ddc7af40d..38c701075f 100644
--- a/python/tvm/tirx/function.py
+++ b/python/tvm/tirx/function.py
@@ -324,10 +324,6 @@ class IndexMap(Object):
mapping = mapping_function(*args, **kwargs)
- from tvm.ir.expr import _realize_operand # pylint:
disable=import-outside-toplevel
-
- mapping = _realize_operand(mapping)
-
initial_indices = args + list(kwargs.values())
final_indices = []
@@ -343,7 +339,6 @@ class IndexMap(Object):
if is_iterable:
for val in mapping:
- val = _realize_operand(val)
if tvm.ir.is_prim_expr(val):
final_indices.append(val)
else:
diff --git a/python/tvm/tirx/op.py b/python/tvm/tirx/op.py
index 29496bcd35..4d9627292a 100644
--- a/python/tvm/tirx/op.py
+++ b/python/tvm/tirx/op.py
@@ -26,7 +26,6 @@ import tvm
from tvm import tirx
from tvm.ir import Call, Expr, ExprWithOp, Op, PointerType, PrimType,
TensorLoad
from tvm.ir.base import Span
-from tvm.ir.expr import _realize_operand
from tvm.ir.type import TensorMapType
from tvm.runtime import const
@@ -686,7 +685,6 @@ def address_of(obj: Buffer | TensorLoad | Var, span: Span |
None = None) -> Expr
call : Expr
The call expression.
"""
- obj = _realize_operand(obj)
if is_buffer_var(obj):
n_dim = len(obj.ty.shape)
buffer_load = BufferLoad(obj, [0] * n_dim)
diff --git a/python/tvm/tirx/script/builder/ir.py
b/python/tvm/tirx/script/builder/ir.py
index 2b136e10ae..c544d35902 100644
--- a/python/tvm/tirx/script/builder/ir.py
+++ b/python/tvm/tirx/script/builder/ir.py
@@ -37,7 +37,6 @@ from tvm import tirx as tir
from tvm.ir import Call, TensorLoad, Type, is_prim_expr
from tvm.ir import register_op_attr as _register_op_attr
from tvm.ir.base import deprecated
-from tvm.ir.expr import _realize_operand
from tvm.runtime import convert
from tvm.script.ir_builder.base import IRBuilder
from tvm.script.ir_builder.ir import meta_var
@@ -329,7 +328,7 @@ def buffer(
The declared buffer.
"""
shape = (shape,) if is_prim_expr(shape) or isinstance(shape, Integral)
else shape
- shape = tuple(_realize_operand(dim) for dim in shape)
+ shape = tuple(shape)
if strides is not None:
strides = [Var(s, "int32") if isinstance(s, str) else s for s in
strides]
else:
@@ -529,7 +528,6 @@ def match_buffer(
res : Buffer
The matched buffer.
"""
- param = _realize_operand(param)
if shape is None:
if isinstance(param, BufferRegion):
dtype = param.buffer.ty.dtype
@@ -1848,7 +1846,7 @@ def decl_buffer(
The declared buffer.
"""
shape = (shape,) if is_prim_expr(shape) or isinstance(shape, Integral)
else shape
- shape = tuple(_realize_operand(dim) for dim in shape)
+ shape = tuple(shape)
if strides is not None:
strides = [Var(s, "int32") if isinstance(s, str) else s for s in
strides]
else:
@@ -2073,7 +2071,7 @@ def alloc_scalar(dtype: str = "float32", scope: str =
"global") -> TensorLoad:
"""Allocate a zero-dimensional buffer (scalar)."""
buf = alloc_buffer(shape=(1,), dtype=dtype, scope=scope,
layout=TileLayout(S[1]))
assert is_buffer_var(buf)
- scalar = _realize_operand(buf[0])
+ scalar = buf[0]
if _current_meta_construction_scope() is not None:
return scalar
return scalar_wrapper(scalar)
@@ -2093,7 +2091,7 @@ def decl_scalar(dtype, data, scope, elem_offset=None,
byte_offset=None) -> Tenso
layout=TileLayout(S[1]),
)
assert is_buffer_var(buf)
- scalar = _realize_operand(buf[0])
+ scalar = buf[0]
if _current_meta_construction_scope() is not None:
return scalar
return scalar_wrapper(scalar)
diff --git a/python/tvm/tirx/script/builder/tirx.py
b/python/tvm/tirx/script/builder/tirx.py
index 4313ae1e20..76d181621c 100644
--- a/python/tvm/tirx/script/builder/tirx.py
+++ b/python/tvm/tirx/script/builder/tirx.py
@@ -22,7 +22,6 @@ from collections.abc import Callable
import tvm
import tvm.tirx.operator as tirx_op
from tvm.ir import Op
-from tvm.ir.expr import _realize_operand
from tvm.tirx import Buffer, BufferRegion, Expr, LambdaExpr, buffer_data,
is_buffer_var
from tvm.tirx.exec_scope import _SCOPE_KIND_TO_NAME, ExecScope
from tvm.tirx.expr import FloatImm, IntImm
@@ -124,14 +123,12 @@ thread = ScopeNamespace("thread", "thread")
def _is_buffer_or_region(x):
- x = _realize_operand(x)
return is_buffer_var(x) or isinstance(x, BufferRegion)
def _to_region(buffer: BufferRegion | Buffer):
- buffer = _realize_operand(buffer)
if is_buffer_var(buffer):
- return _realize_operand(buffer[tuple(slice(None) for _ in
buffer.ty.shape)])
+ return buffer[tuple(slice(None) for _ in buffer.ty.shape)]
assert isinstance(buffer, BufferRegion)
return buffer
diff --git a/python/tvm/tirx/script/parser/operation.py
b/python/tvm/tirx/script/parser/operation.py
index f5f578c3a0..fd67d6f125 100644
--- a/python/tvm/tirx/script/parser/operation.py
+++ b/python/tvm/tirx/script/parser/operation.py
@@ -19,7 +19,6 @@
import tvm
from tvm import tirx
from tvm.ir import PrimType
-from tvm.ir.expr import _realize_operand
from tvm.runtime import DataTypeCode
from tvm.script.parser._core import OpMethod, doc, register_op
from tvm.tirx import IntImm
@@ -38,7 +37,6 @@ def _register_expr_op(ty: type): # pylint:
disable=invalid-name
return ty
def _and(a, b):
- a, b = _realize_operand(a), _realize_operand(b)
if isinstance(a, bool):
a = IntImm("bool", a)
if isinstance(b, bool):
@@ -49,7 +47,6 @@ def _register_expr_op(ty: type): # pylint:
disable=invalid-name
return tirx.And(a, b)
def _or(a, b):
- a, b = _realize_operand(a), _realize_operand(b)
if isinstance(a, bool):
a = IntImm("bool", a)
if isinstance(b, bool):
@@ -67,7 +64,6 @@ def _register_expr_op(ty: type): # pylint:
disable=invalid-name
return dtype_str[0:index]
def _auto_broadcast(a, b, op):
- a, b = _realize_operand(a), _realize_operand(b)
if isinstance(a, int):
if tvm.ir.is_prim_expr(b) or hasattr(b, "expr_ty"):
b_ty = _expr_ty(b)
@@ -169,4 +165,3 @@ def _register_expr_op(ty: type): # pylint:
disable=invalid-name
_register_expr_op(tirx.Expr)
_register_expr_op(tirx.IterVar)
-_register_expr_op(tvm.ir.SubscriptProxy)
diff --git a/python/tvm/tirx/script/parser/parser.py
b/python/tvm/tirx/script/parser/parser.py
index 7398f0bc0a..607988a4d1 100644
--- a/python/tvm/tirx/script/parser/parser.py
+++ b/python/tvm/tirx/script/parser/parser.py
@@ -24,7 +24,6 @@ from typing import Any, TypeVar
import tvm
from tvm.ir import Expr, GlobalVar, PointerType, PrimType, TensorLoad
-from tvm.ir.expr import _realize_operand
from tvm.script.ir_builder import ir as I
from tvm.script.ir_builder.base import IRBuilder
from tvm.script.ir_builder.base import IRBuilderFrame as Frame
@@ -222,7 +221,6 @@ def bind_assign_value(
res : Any
The bound value.
"""
- value = _realize_operand(value)
if var_name in (prim_var_declarations or set()):
# A quoted Buffer shape may have already created this PrimVar. In that
# case ``n = T.int32()`` is match-like syntax: bind the Python name to
diff --git a/python/tvm/tirx/script/tile.py b/python/tvm/tirx/script/tile.py
index cf5fc9e262..4aaf79113b 100644
--- a/python/tvm/tirx/script/tile.py
+++ b/python/tvm/tirx/script/tile.py
@@ -18,7 +18,6 @@
import functools
-from tvm.ir.expr import _realize_operand
from tvm.tirx import BufferRegion, is_buffer_var
from .builder import tirx as _builder
@@ -31,7 +30,6 @@ def _get_arg(args, kwargs, index, name):
def _require_buffer_arg(op_name, arg_name, value):
- value = _realize_operand(value)
if not (is_buffer_var(value) or isinstance(value, BufferRegion)):
raise TypeError(
f"Tx.{op_name} is tile-only and expects `{arg_name}` to be a
Buffer "
diff --git a/python/tvm/tirx/stmt.py b/python/tvm/tirx/stmt.py
index fdccf5ee14..ff3c27a990 100644
--- a/python/tvm/tirx/stmt.py
+++ b/python/tvm/tirx/stmt.py
@@ -33,9 +33,8 @@ from typing import Any
import tvm_ffi
-from tvm.ir import Expr, Range, Span, is_prim_expr
+from tvm.ir import Expr, Range, Span
from tvm.runtime import Object, Scriptable, const
-from tvm.tirx import IntImm
from . import _ffi_api
from .buffer import Buffer
@@ -626,40 +625,6 @@ class BufferRegion(Object, Scriptable):
def __init__(self, buffer: Buffer, region: list[Range]) -> None:
self.__init_handle_by_constructor__(_ffi_api.BufferRegion, buffer,
region) # type: ignore
- def __getitem__(self, indices):
- from ..arith import Analyzer
-
- if not isinstance(indices, tuple | list):
- indices = [indices]
-
- has_step = any(
- isinstance(i, slice) and (i.step is not None and i.step != 1) for
i in indices
- )
- if has_step:
- raise ValueError("BufferRegion slicing does not support steps")
-
- analyzer = Analyzer()
- new_region = []
- for i, index in enumerate(indices):
- old_range = self.region[i]
- if isinstance(index, slice):
- start = 0 if index.start is None else index.start
- stop = old_range.extent if index.stop is None else index.stop
- new_min = old_range.min + start
- new_extent = analyzer.simplify(stop - start)
- new_region.append(Range.from_min_extent(new_min, new_extent))
- else:
- new_min = old_range.min + index
- new_region.append(
- Range.from_min_extent(
- new_min, IntImm(index.ty, 1) if is_prim_expr(index)
else 1
- )
- )
- # Fill remaining dimensions with their original ranges
- for i in range(len(indices), len(self.region)):
- new_region.append(self.region[i])
- return BufferRegion(self.buffer, new_region)
-
@tvm_ffi.register_object("tirx.MatchBufferRegion")
class MatchBufferRegion(Object, Scriptable):
diff --git a/src/ir/expr.cc b/src/ir/expr.cc
index d9c34033d3..67f0318e2a 100644
--- a/src/ir/expr.cc
+++ b/src/ir/expr.cc
@@ -73,11 +73,11 @@ Tuple::Tuple(ffi::Array<Expr> fields, Span span) {
}
TupleGetItem::TupleGetItem(Expr tuple, int index, Span span) {
- TVM_FFI_ICHECK_GE(index, 0) << "Index out of bounds: Tuple " << tuple
- << " cannot be accessed with negative index " <<
index;
+ TVM_FFI_CHECK_GE(index, 0, IndexError) << "Index out of bounds: Tuple " <<
tuple
+ << " cannot be accessed with negative
index " << index;
ffi::ObjectPtr<TupleGetItemNode> node = ffi::make_object<TupleGetItemNode>();
if (const auto* tuple_type = tuple->ty.as<TupleTypeNode>()) {
- TVM_FFI_ICHECK_LT(index, tuple_type->fields.size())
+ TVM_FFI_CHECK_LT(index, tuple_type->fields.size(), IndexError)
<< "Index out of bounds: Tuple " << tuple << " is of size " <<
tuple_type->fields.size()
<< ", and cannot be accessed with index " << index;
node->ty = tuple_type->fields[index];
diff --git a/src/ir/subscript_proxy.cc b/src/ir/subscript_expr.cc
similarity index 84%
rename from src/ir/subscript_proxy.cc
rename to src/ir/subscript_expr.cc
index c6e8e75109..9225a7ad62 100644
--- a/src/ir/subscript_proxy.cc
+++ b/src/ir/subscript_expr.cc
@@ -18,8 +18,8 @@
*/
/*!
- * \file subscript_proxy.cc
- * \brief Type-directed realization for the Python frontend's SubscriptProxy.
+ * \file subscript_expr.cc
+ * \brief Type-directed realization for Python expression subscription.
*/
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
@@ -47,12 +47,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
<< "A tuple expression requires a constant integer index";
return TupleGetItem(value, static_cast<int>(imm->value), span);
});
- refl::GlobalDef().def("ir.SubscriptExprCheck", [](Expr value) {
- TVM_FFI_CHECK(value.defined(), TypeError) << "Cannot subscript an
undefined expression";
- static refl::TypeAttrColumn realize_column("__subscript_expr_realize__");
- TVM_FFI_CHECK(realize_column[value->ty->type_index()] != nullptr,
TypeError)
- << "Type " << value->ty->GetTypeKey() << " does not support subscript";
- });
refl::GlobalDef().def(
"ir.SubscriptExprRealize", [](Expr value, SubscriptSlice slice, Span
span) -> ffi::ObjectRef {
TVM_FFI_CHECK(value.defined(), TypeError) << "Cannot subscript an
undefined expression";
diff --git a/tests/python/relax/test_blockbuilder_core.py
b/tests/python/relax/test_blockbuilder_core.py
index b76acc1ee9..212b9d4388 100644
--- a/tests/python/relax/test_blockbuilder_core.py
+++ b/tests/python/relax/test_blockbuilder_core.py
@@ -333,6 +333,9 @@ def test_tuple_indexing():
y = relax_tuple[1]
tvm.ir.assert_structural_equal(y.ty, shape_y)
+ with pytest.raises(IndexError, match="Index out of bounds"):
+ relax_tuple[2]
+
# Tuple unpacking produces TupleGetItem structs
x_unpack, y_unpack = relax_tuple
tvm.ir.assert_structural_equal(x, x_unpack)
diff --git a/tests/python/s_tir/analysis/test_s_tir_analysis_identify_memcpy.py
b/tests/python/s_tir/analysis/test_s_tir_analysis_identify_memcpy.py
index e4cc34f879..acd33f9814 100644
--- a/tests/python/s_tir/analysis/test_s_tir_analysis_identify_memcpy.py
+++ b/tests/python/s_tir/analysis/test_s_tir_analysis_identify_memcpy.py
@@ -56,7 +56,7 @@ def test_1d():
B[i] = A[i]
A, B = [param for param in func.params if tvm.tirx.is_buffer_var(param)]
- expected = (A[0:1024].to_expr(), B[0:1024].to_expr())
+ expected = (A[0:1024], B[0:1024])
_check_memcpy_results(func, expected)
@@ -118,7 +118,7 @@ def test_1d_input_2d_output_fused_loop():
B[i // 32, i % 32] = A[i]
A, B = [param for param in func.params if tvm.tirx.is_buffer_var(param)]
- expected = (A[0:1024].to_expr(), B[0:32, 0:32].to_expr())
+ expected = (A[0:1024], B[0:32, 0:32])
_check_memcpy_results(func, expected)
@@ -131,7 +131,7 @@ def test_2d_input_1d_output_fused_loop():
B[i] = A[i // 32, i % 32]
A, B = [param for param in func.params if tvm.tirx.is_buffer_var(param)]
- expected = (A[0:32, 0:32].to_expr(), B[0:1024].to_expr())
+ expected = (A[0:32, 0:32], B[0:1024])
_check_memcpy_results(func, expected)
diff --git a/tests/python/tirx/test_parser_printer.py
b/tests/python/tirx/test_parser_printer.py
index 5023290d70..cf950e63e6 100644
--- a/tests/python/tirx/test_parser_printer.py
+++ b/tests/python/tirx/test_parser_printer.py
@@ -2194,7 +2194,7 @@ def test_buffer_chunk_ir():
assert isinstance(reg, BufferRegion)
assert len(reg.region) == 3 # rank-preserving: no extra extent-1 chunk dim
assert (int(reg.region[2].min), int(reg.region[2].extent)) == (8, 8)
- assert_structural_equal(reg, A[:, :, 8:16].to_expr())
+ assert_structural_equal(reg, A[:, :, 8:16])
# a None dim passes an int pick straight through (int → extent-1 region),
# while the chunked dim still narrows to its picked chunk.
@@ -2257,12 +2257,20 @@ def test_buffer_slice_region():
from tvm.tirx.stmt import BufferRegion
buf = tvm.tirx.decl_buffer((128, 64), "float16")
- br = buf[32:64, 0:32].to_expr()
+ br = buf[32:64, 0:32]
assert isinstance(br, BufferRegion)
assert br.buffer.same_as(buf)
assert int(br.region[0].extent) == 32
assert int(br.region[1].extent) == 32
+ load = buf[1, 2]
+ assert isinstance(load, tvm.ir.TensorLoad)
+
+ partial = buf[1]
+ assert isinstance(partial, BufferRegion)
+ with pytest.raises(TypeError):
+ _ = partial[2]
+
def test_global_call_realizes_buffer_elements():
@I.ir_module(s_tir=True)
@@ -2283,25 +2291,6 @@ def test_global_call_realizes_buffer_elements():
assert isinstance(Module["main"], tvm.tirx.PrimFunc)
-def test_buffer_region_slice():
- """Verify BufferRegion slicing returns BufferRegion."""
- from tvm.tirx.stmt import BufferRegion
-
- buf = tvm.tirx.decl_buffer((128, 64), "float16")
-
- br1 = buf[32:64, 0:32].to_expr()
- assert isinstance(br1, BufferRegion)
-
- # BufferRegion chained slice
- br3 = br1[0:16, 0:16]
- assert isinstance(br3, BufferRegion)
- assert br3.buffer.same_as(buf), "chained region slice must reference root
buffer"
- assert int(br3.region[0].min) == 32
- assert int(br3.region[0].extent) == 16
- assert int(br3.region[1].min) == 0
- assert int(br3.region[1].extent) == 16
-
-
def test_roundtrip_serial_unroll_false():
"""T.serial(N, unroll=False) should round-trip."""
