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 71f0271942 [REFACTOR][IR] Unify expression subscription realization
(#20246)
71f0271942 is described below
commit 71f0271942532a8011814092e498f4a61dbb97ea
Author: Tianqi Chen <[email protected]>
AuthorDate: Mon Aug 31 19:06:39 2026 -0400
[REFACTOR][IR] Unify expression subscription realization (#20246)
## Summary
- add a lazy `tvm.ir.SubscriptProxy` with type-registered realization
- centralize Python expression operator boundaries through `ExprOperand`
- simplify buffer subscription to full-rank point loads or regions
- preserve tuple, parser, builder, and Relax subscription behavior
---
python/tvm/backend/cuda/ptx/engine.py | 2 +
python/tvm/ir/__init__.py | 2 +
python/tvm/ir/_tensor_expr_overload.py | 4 -
python/tvm/ir/expr.py | 166 +++++++++++++++++++--
python/tvm/relax/expr.py | 29 ----
python/tvm/relax/op/__init__.py | 9 --
python/tvm/script/parser/core/evaluator.py | 3 +
python/tvm/script/parser/core/parser.py | 9 +-
python/tvm/tirx/_buffer_view.py | 4 +-
python/tvm/tirx/buffer.py | 58 +------
python/tvm/tirx/function.py | 5 +
python/tvm/tirx/op.py | 2 +
python/tvm/tirx/script/builder/ir.py | 8 +-
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 +
src/ir/subscript_proxy.cc | 71 +++++++++
src/relax/ir/dependent_type.cc | 17 +++
src/tirx/ir/buffer.cc | 69 ++++++++-
src/tirx/script/printer/buffer.cc | 12 +-
.../test_s_tir_analysis_identify_memcpy.py | 6 +-
tests/python/tirx/test_parser_printer.py | 25 +++-
.../python/tvmscript/test_tvmscript_printer_tir.py | 4 +-
tests/python/tvmscript/test_tvmscript_roundtrip.py | 4 +-
25 files changed, 398 insertions(+), 125 deletions(-)
diff --git a/python/tvm/backend/cuda/ptx/engine.py
b/python/tvm/backend/cuda/ptx/engine.py
index 3897b22991..2529065127 100644
--- a/python/tvm/backend/cuda/ptx/engine.py
+++ b/python/tvm/backend/cuda/ptx/engine.py
@@ -41,6 +41,7 @@ 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
@@ -481,6 +482,7 @@ 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 36534b5635..8a2d241cd1 100644
--- a/python/tvm/ir/__init__.py
+++ b/python/tvm/ir/__init__.py
@@ -37,9 +37,11 @@ from .type import FuncType, OpaqueType, PointerType,
PrimType, TupleType, Type
from .expr import (
Call,
Expr,
+ ExprOperand,
GlobalVar,
OpaqueExpr,
Range,
+ SubscriptProxy,
TensorLoad,
Tuple,
TupleGetItem,
diff --git a/python/tvm/ir/_tensor_expr_overload.py
b/python/tvm/ir/_tensor_expr_overload.py
index d3fcf25ca0..1c42a9ebb6 100644
--- a/python/tvm/ir/_tensor_expr_overload.py
+++ b/python/tvm/ir/_tensor_expr_overload.py
@@ -105,9 +105,5 @@ def __call__(_value, *_args, **_kwargs):
return NotImplemented
-def __getitem__(_value, _index):
- return NotImplemented
-
-
def astype(_value, _dtype, _span=None):
return NotImplemented
diff --git a/python/tvm/ir/expr.py b/python/tvm/ir/expr.py
index 263c2f5821..14d467263b 100644
--- a/python/tvm/ir/expr.py
+++ b/python/tvm/ir/expr.py
@@ -22,7 +22,7 @@ import tvm_ffi
import tvm
-from ..runtime import Object, Scriptable
+from ..runtime import Object, ObjectConvertible, Scriptable
from . import _ffi_api, _overload_prim_expr, _tensor_expr_overload
from .base import Node, Span
@@ -34,6 +34,29 @@ class Expr(Node):
span: Span | None
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)])
+ 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)
+
@tvm_ffi.register_object("ir.OpaqueExpr")
class OpaqueExpr(Expr):
@@ -83,6 +106,8 @@ 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)
@@ -100,98 +125,124 @@ class GlobalVar(Expr):
raise RuntimeError(f"Do not know how to handle GlobalVar.__call__ for
types {arg_types}")
-class _ExprWithOp(Expr, Scriptable):
- """Common type-directed operator behavior for core expressions."""
+def _realize_operand(value):
+ return value._operand() if isinstance(value, ExprOperand) else value
- __hash__ = Expr.__hash__
+
+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)
return _tensor_expr_overload.__add__(self, other)
def __radd__(self, other):
+ self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__radd__(self, other)
return _tensor_expr_overload.__radd__(self, other)
def __sub__(self, other):
+ self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__sub__(self, other)
return _tensor_expr_overload.__sub__(self, other)
def __rsub__(self, other):
+ self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__rsub__(self, other)
return _tensor_expr_overload.__rsub__(self, other)
def __mul__(self, other):
+ self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__mul__(self, other)
return _tensor_expr_overload.__mul__(self, other)
def __rmul__(self, other):
+ self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__rmul__(self, other)
return _tensor_expr_overload.__rmul__(self, other)
def __div__(self, other):
+ self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__div__(self, other)
return _tensor_expr_overload.__div__(self, other)
def __rdiv__(self, other):
+ self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__rdiv__(self, other)
return _tensor_expr_overload.__rdiv__(self, other)
def __truediv__(self, other):
+ self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__truediv__(self, other)
return _tensor_expr_overload.__truediv__(self, other)
def __rtruediv__(self, other):
+ self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__rtruediv__(self, other)
return _tensor_expr_overload.__rtruediv__(self, other)
def __floordiv__(self, other):
+ self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__floordiv__(self, other)
return _tensor_expr_overload.__floordiv__(self, other)
def __rfloordiv__(self, other):
+ self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__rfloordiv__(self, other)
return _tensor_expr_overload.__rfloordiv__(self, other)
def __mod__(self, other):
+ self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__mod__(self, other)
return _tensor_expr_overload.__mod__(self, other)
def __rmod__(self, other):
+ self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__rmod__(self, other)
return _tensor_expr_overload.__rmod__(self, other)
def __pow__(self, other):
+ self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return NotImplemented
return _tensor_expr_overload.__pow__(self, other)
def __rpow__(self, other):
+ self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return NotImplemented
return _tensor_expr_overload.__rpow__(self, other)
def __neg__(self):
+ self = _realize_operand(self)
if is_prim_expr(self):
result = _overload_prim_expr.__neg__(self)
if result is NotImplemented:
@@ -203,56 +254,67 @@ class _ExprWithOp(Expr, Scriptable):
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:
@@ -261,31 +323,37 @@ class _ExprWithOp(Expr, Scriptable):
return NotImplemented
def __lt__(self, other):
+ self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__lt__(self, other)
return _tensor_expr_overload.__lt__(self, other)
def __le__(self, other):
+ self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__le__(self, other)
return _tensor_expr_overload.__le__(self, other)
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)
return _tensor_expr_overload.__gt__(self, other)
def __ge__(self, other):
+ self, other = _realize_operand(self), _realize_operand(other)
if is_prim_expr(self):
return _overload_prim_expr.__ge__(self, other)
return _tensor_expr_overload.__ge__(self, other)
@@ -300,12 +368,14 @@ class _ExprWithOp(Expr, Scriptable):
return self.__nonzero__()
def equal(self, other, span=None):
+ self, other = _realize_operand(self), _realize_operand(other)
result = _overload_prim_expr.equal(self, other, span)
if result is NotImplemented:
raise TypeError("Primitive expression overload equal is not
registered")
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:
@@ -317,6 +387,7 @@ class _ExprWithOp(Expr, Scriptable):
return result
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)
@@ -324,13 +395,88 @@ class _ExprWithOp(Expr, Scriptable):
raise TypeError("Tensor expression overload __call__ is not
registered")
return result
+
+class _ExprWithOp(ExprOperand, Expr, Scriptable):
+ """Common type-directed operator behavior for core expressions."""
+
+ __hash__ = Expr.__hash__
+
+ def _operand(self) -> Expr:
+ return self
+
+
+class SubscriptProxy(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")
+ __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)
+ else:
+ _ffi_api.SubscriptExprCheck(source)
+ self._source = source
+ self._slice = self._flatten(index)
+ self._result = None
+
+ @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):
- if is_prim_expr(self):
- raise TypeError("A primitive-valued expression cannot be indexed")
- result = _tensor_expr_overload.__getitem__(self, index)
- if result is NotImplemented:
- raise TypeError("Tensor expression overload __getitem__ is not
registered")
- return result
+ 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]
+ )
+ 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")
diff --git a/python/tvm/relax/expr.py b/python/tvm/relax/expr.py
index 3c983da029..51c63e21d1 100644
--- a/python/tvm/relax/expr.py
+++ b/python/tvm/relax/expr.py
@@ -206,35 +206,6 @@ class ExprWithOp(Expr, Scriptable):
"""
return tvm.ir.Call(self, args, attrs=attrs)
- def __getitem__(self, index: int) -> "ExprWithOp":
- """Get the i-th element of the tuple or Expr with TupleType.
-
- Parameters
- ----------
- index: int
- The index of the element to be retrieved.
-
- Note
- ----
- This function will be overridden by Tuple and ShapeExpr
-
- Returns
- -------
- result: ExprWithOp
- The result expression.
- """
- try:
- return TupleGetItem(self, index)
- except RuntimeError as err:
- # For Python objects with __getitem__, but without
- # __len__, tuple unpacking is done by iterating over
- # sequential indices until IndexError is raised.
- # Therefore, convert from RuntimeError to IndexError for
- # compatibility.
- if "Index out of bounds" in err.args[0]:
- raise IndexError from err
- raise
-
@tvm_ffi.register_object("relax.expr.If")
class If(ExprWithOp):
diff --git a/python/tvm/relax/op/__init__.py b/python/tvm/relax/op/__init__.py
index 24fa24b02a..21a25e16b4 100644
--- a/python/tvm/relax/op/__init__.py
+++ b/python/tvm/relax/op/__init__.py
@@ -185,19 +185,10 @@ def _register_op_make():
def _rhs(_lhs, rhs):
return expr._binary_rhs_helper(rhs)
- def _getitem(value, index):
- try:
- return expr.TupleGetItem(value, index)
- except RuntimeError as err:
- if "Index out of bounds" in err.args[0]:
- raise IndexError from err
- raise
-
_tensor_expr_overload.astype = lambda lhs, dtype, _span=None:
_ffi_api.astype(lhs, dtype)
_tensor_expr_overload.__call__ = lambda func, *args, attrs=None:
expr.tvm.ir.Call(
func, args, attrs=attrs
)
- _tensor_expr_overload.__getitem__ = _getitem
_tensor_expr_overload.__neg__ = lambda lhs: _ffi_api.negative(lhs)
_tensor_expr_overload.__lt__ = lambda lhs, rhs:
expr._binary_op_helper(lhs, rhs, _ffi_api.less)
_tensor_expr_overload.__le__ = lambda lhs, rhs: expr._binary_op_helper(
diff --git a/python/tvm/script/parser/core/evaluator.py
b/python/tvm/script/parser/core/evaluator.py
index 71a1b3bc8f..06f6dc9b09 100644
--- a/python/tvm/script/parser/core/evaluator.py
+++ b/python/tvm/script/parser/core/evaluator.py
@@ -413,6 +413,9 @@ 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 771be8f3d1..7d0f2887df 100644
--- a/python/tvm/script/parser/core/parser.py
+++ b/python/tvm/script/parser/core/parser.py
@@ -627,7 +627,14 @@ class Parser(doc.NodeVisitor):
for k, v in extra_vars.items():
var_values[k] = v
var_values[ScriptMacro.parser_object_name] = self
- return eval_expr(self, node, var_values)
+ 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 _realize_operand(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 90c4a13a61..3977cf2f3a 100644
--- a/python/tvm/tirx/_buffer_view.py
+++ b/python/tvm/tirx/_buffer_view.py
@@ -572,7 +572,9 @@ class ChunkIndexer:
else:
size = buf.shape[dim] // int(count)
translated.append(slice(pick * size, (pick + 1) * size))
- return buf[tuple(translated)]
+ # ``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()
class TileIndexer:
diff --git a/python/tvm/tirx/buffer.py b/python/tvm/tirx/buffer.py
index d832aa7602..dc62b35f4c 100644
--- a/python/tvm/tirx/buffer.py
+++ b/python/tvm/tirx/buffer.py
@@ -22,7 +22,7 @@ from numbers import Integral
import tvm_ffi
import tvm
-from tvm.ir import PointerType, PrimType, Range, Type
+from tvm.ir import PointerType, PrimType, Type
from tvm.runtime import convert
from . import _buffer_view, _ffi_api
@@ -458,59 +458,6 @@ class _BufferMethods:
"""
return _buffer_view.chunk(self, spec)
- def __getitem__(self, indices):
- if not is_buffer_var(self):
- return _ORIGINAL_VAR_GETITEM(self, indices)
-
- from ..arith import Analyzer # pylint: disable=import-outside-toplevel
- from .expr import BufferLoad, Ramp # pylint:
disable=import-outside-toplevel
- from .stmt import BufferRegion # pylint:
disable=import-outside-toplevel
-
- if not isinstance(indices, tuple | list):
- indices = [indices]
- has_slice = any(isinstance(i, slice) for i in indices)
- has_step = any(
- isinstance(i, slice) and (i.step is not None and i.step != 1) for
i in indices
- )
- has_implicit_slice = len(indices) < len(self.ty.shape)
- analyzer = Analyzer()
- if (has_slice and not has_step) or has_implicit_slice:
- region = []
- for i, index in enumerate(indices):
- if isinstance(index, slice):
- start = 0 if index.start is None else index.start
- stop = self.ty.shape[i] if index.stop is None else
index.stop
- region.append(Range.from_min_extent(start,
analyzer.simplify(stop - start)))
- else:
- region.append(
- Range.from_min_extent(
- index,
- tvm.tirx.expr.IntImm(index.ty, 1) if
tvm.ir.is_prim_expr(index) else 1,
- )
- )
- if has_implicit_slice:
- for i in range(len(indices), len(self.ty.shape)):
- region.append(Range.from_min_extent(0, self.ty.shape[i]))
- return BufferRegion(self, region)
- else:
- expr_indices = []
- for i, index in enumerate(indices):
- if isinstance(index, slice):
- start = 0 if index.start is None else index.start
- stop = self.ty.shape[i] if index.stop is None else
index.stop
- step = 1 if index.step is None else index.step
- # We should ensure the dtype of start is the same with
that of step.
- if tvm.ir.is_prim_expr(start) and isinstance(step, int):
- step = tvm.tirx.expr.IntImm(start.ty, step)
- lanes = analyzer.simplify((stop - start + step - 1) //
step)
- if lanes == 1:
- expr_indices.append(start)
- else:
- expr_indices.append(Ramp(start, step, int(lanes)))
- else:
- expr_indices.append(index)
- return BufferLoad(self, expr_indices)
-
def decl_buffer(
shape,
@@ -581,9 +528,8 @@ def buffer_data_pointer_type(buffer):
# ``tvm.tirx`` therefore augments ``tvm.ir.Var`` process-wide with the legacy
# buffer operation and metadata surface. Non-buffer Vars reject the metadata
# properties with AttributeError, preserving correct ``hasattr`` behavior.
-_ORIGINAL_VAR_GETITEM = tvm.ir.Var.__getitem__
for _name, _value in _BufferMethods.__dict__.items():
- if _name.startswith("__") and _name != "__getitem__":
+ if _name.startswith("__"):
continue
if callable(_value) or isinstance(_value, property):
setattr(tvm.ir.Var, _name, _value)
diff --git a/python/tvm/tirx/function.py b/python/tvm/tirx/function.py
index 38c701075f..7ddc7af40d 100644
--- a/python/tvm/tirx/function.py
+++ b/python/tvm/tirx/function.py
@@ -324,6 +324,10 @@ 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 = []
@@ -339,6 +343,7 @@ 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 9d0eb50932..5621e93195 100644
--- a/python/tvm/tirx/op.py
+++ b/python/tvm/tirx/op.py
@@ -26,6 +26,7 @@ import tvm
from tvm import tirx
from tvm.ir import Call, Expr, 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
@@ -685,6 +686,7 @@ 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 a686506d13..2b136e10ae 100644
--- a/python/tvm/tirx/script/builder/ir.py
+++ b/python/tvm/tirx/script/builder/ir.py
@@ -37,6 +37,7 @@ 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
@@ -328,6 +329,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)
if strides is not None:
strides = [Var(s, "int32") if isinstance(s, str) else s for s in
strides]
else:
@@ -527,6 +529,7 @@ 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
@@ -1845,6 +1848,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)
if strides is not None:
strides = [Var(s, "int32") if isinstance(s, str) else s for s in
strides]
else:
@@ -2069,7 +2073,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 = buf[0]
+ scalar = _realize_operand(buf[0])
if _current_meta_construction_scope() is not None:
return scalar
return scalar_wrapper(scalar)
@@ -2089,7 +2093,7 @@ def decl_scalar(dtype, data, scope, elem_offset=None,
byte_offset=None) -> Tenso
layout=TileLayout(S[1]),
)
assert is_buffer_var(buf)
- scalar = buf[0]
+ scalar = _realize_operand(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 8ad43a96d9..4313ae1e20 100644
--- a/python/tvm/tirx/script/builder/tirx.py
+++ b/python/tvm/tirx/script/builder/tirx.py
@@ -22,6 +22,7 @@ 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
@@ -123,12 +124,14 @@ 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 buffer[[slice(None, None, None) for _ in
range(len(buffer.ty.shape))]]
+ return _realize_operand(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 fd67d6f125..f5f578c3a0 100644
--- a/python/tvm/tirx/script/parser/operation.py
+++ b/python/tvm/tirx/script/parser/operation.py
@@ -19,6 +19,7 @@
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
@@ -37,6 +38,7 @@ 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):
@@ -47,6 +49,7 @@ 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):
@@ -64,6 +67,7 @@ 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)
@@ -165,3 +169,4 @@ 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 607988a4d1..7398f0bc0a 100644
--- a/python/tvm/tirx/script/parser/parser.py
+++ b/python/tvm/tirx/script/parser/parser.py
@@ -24,6 +24,7 @@ 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
@@ -221,6 +222,7 @@ 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 4aaf79113b..cf5fc9e262 100644
--- a/python/tvm/tirx/script/tile.py
+++ b/python/tvm/tirx/script/tile.py
@@ -18,6 +18,7 @@
import functools
+from tvm.ir.expr import _realize_operand
from tvm.tirx import BufferRegion, is_buffer_var
from .builder import tirx as _builder
@@ -30,6 +31,7 @@ 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/src/ir/subscript_proxy.cc b/src/ir/subscript_proxy.cc
new file mode 100644
index 0000000000..a92c961c61
--- /dev/null
+++ b/src/ir/subscript_proxy.cc
@@ -0,0 +1,71 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file subscript_proxy.cc
+ * \brief Type-directed realization for the Python frontend's SubscriptProxy.
+ */
+#include <tvm/ffi/function.h>
+#include <tvm/ffi/reflection/registry.h>
+#include <tvm/ir/expr.h>
+#include <tvm/ir/type.h>
+
+namespace tvm {
+
+using SubscriptSlice = ffi::Array<ffi::Variant<
+ ffi::Tuple<ffi::Optional<PrimExpr>, ffi::Optional<PrimExpr>,
ffi::Optional<PrimExpr>>,
+ PrimExpr>>;
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ namespace refl = tvm::ffi::reflection;
+ refl::EnsureTypeAttrColumn("__subscript_expr_realize__");
+ refl::TypeAttrDef<TupleTypeNode>().def(
+ "__subscript_expr_realize__", [](Expr value, SubscriptSlice slice) ->
ffi::ObjectRef {
+ TVM_FFI_CHECK_EQ(slice.size(), 1, IndexError)
+ << "A tuple expression requires exactly one index";
+ auto index = slice[0].as<PrimExpr>();
+ TVM_FFI_CHECK(index.has_value(), TypeError) << "A tuple expression
requires a point index";
+ const auto* imm = index.value().as<IntImmNode>();
+ TVM_FFI_CHECK(imm != nullptr, TypeError)
+ << "A tuple expression requires a constant integer index";
+ return TupleGetItem(value, static_cast<int>(imm->value));
+ });
+ 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) ->
ffi::ObjectRef {
+ TVM_FFI_CHECK(value.defined(), TypeError) << "Cannot subscript an
undefined expression";
+ static refl::TypeAttrColumn
realize_column("__subscript_expr_realize__");
+ ffi::AnyView packed_realize = realize_column[value->ty->type_index()];
+ TVM_FFI_CHECK(packed_realize != nullptr, TypeError)
+ << "Type " << value->ty->GetTypeKey() << " does not support
subscript";
+ ffi::ObjectRef result =
+ packed_realize.cast<ffi::Function>()(value,
slice).cast<ffi::ObjectRef>();
+ TVM_FFI_CHECK(result.defined(), TypeError)
+ << "__subscript_expr_realize__ for type " <<
value->ty->GetTypeKey()
+ << " returned an undefined object";
+ return result;
+ });
+}
+
+} // namespace tvm
diff --git a/src/relax/ir/dependent_type.cc b/src/relax/ir/dependent_type.cc
index 95bcc3ee04..2ac4898551 100644
--- a/src/relax/ir/dependent_type.cc
+++ b/src/relax/ir/dependent_type.cc
@@ -31,10 +31,27 @@ namespace tvm {
namespace relax {
TVM_FFI_STATIC_INIT_BLOCK() {
+ namespace refl = tvm::ffi::reflection;
AnyTypeNode::RegisterReflection();
ShapeTypeNode::RegisterReflection();
TensorTypeNode::RegisterReflection();
FuncTypeNode::RegisterReflection();
+ refl::TypeAttrDef<TensorTypeNode>().def(
+ "__subscript_expr_realize__",
+ [](Expr value,
+ ffi::Array<ffi::Variant<
+ ffi::Tuple<ffi::Optional<PrimExpr>, ffi::Optional<PrimExpr>,
ffi::Optional<PrimExpr>>,
+ PrimExpr>>
+ slice) -> ffi::ObjectRef {
+ TVM_FFI_CHECK_EQ(slice.size(), 1, IndexError)
+ << "A Relax expression requires exactly one index";
+ auto index = slice[0].as<PrimExpr>();
+ TVM_FFI_CHECK(index.has_value(), TypeError) << "A Relax expression
requires a point index";
+ const auto* imm = index.value().as<IntImmNode>();
+ TVM_FFI_CHECK(imm != nullptr, TypeError)
+ << "A Relax expression requires a constant integer index";
+ return TupleGetItem(value, static_cast<int>(imm->value));
+ });
}
AnyType::AnyType(Span span) : Type(ffi::UnsafeInit{}) {
diff --git a/src/tirx/ir/buffer.cc b/src/tirx/ir/buffer.cc
index 14e0358001..c715aedf0c 100644
--- a/src/tirx/ir/buffer.cc
+++ b/src/tirx/ir/buffer.cc
@@ -29,6 +29,7 @@
#include <tvm/tirx/builtin.h>
#include <tvm/tirx/expr.h>
#include <tvm/tirx/op.h>
+#include <tvm/tirx/stmt.h>
#include <iterator>
#include <list>
@@ -39,7 +40,73 @@
namespace tvm {
namespace tirx {
-TVM_FFI_STATIC_INIT_BLOCK() { BufferTypeNode::RegisterReflection(); }
+namespace {
+
+ffi::ObjectRef RealizeBufferSubscript(
+ Expr value,
+ ffi::Array<ffi::Variant<
+ ffi::Tuple<ffi::Optional<PrimExpr>, ffi::Optional<PrimExpr>,
ffi::Optional<PrimExpr>>,
+ PrimExpr>>
+ slice) {
+ BufferVar buffer = value.as_or_throw<BufferVar>();
+ BufferType buffer_ty = buffer.type();
+ TVM_FFI_CHECK_LE(slice.size(), buffer_ty->shape.size(), IndexError)
+ << "Too many indices for a " << buffer_ty->shape.size() << "-dimensional
buffer";
+
+ bool all_points = slice.size() == buffer_ty->shape.size();
+ for (const auto& item : slice) {
+ if (auto descriptor = item.as<ffi::Tuple<ffi::Optional<PrimExpr>,
ffi::Optional<PrimExpr>,
+ ffi::Optional<PrimExpr>>>()) {
+ all_points = false;
+ ffi::Optional<PrimExpr> step = descriptor.value().get<2>();
+ TVM_FFI_CHECK(!step.has_value() || is_one(step.value()), ValueError)
+ << "Buffer slices with a non-unit step are not supported";
+ }
+ }
+
+ if (all_points) {
+ ffi::Array<PrimExpr> indices;
+ indices.reserve(slice.size());
+ for (const auto& item : slice) {
+ indices.push_back(item.as<PrimExpr>().value());
+ }
+ return BufferLoad(buffer, indices);
+ }
+
+ // Any slice or omitted trailing dimension denotes a region. Rejecting
+ // steps makes the old behavior, where a stride could be silently dropped,
+ // unrepresentable rather than giving it dimension-dependent semantics.
+ arith::Analyzer analyzer;
+ ffi::Array<Range> region;
+ region.reserve(buffer_ty->shape.size());
+ for (size_t i = 0; i < slice.size(); ++i) {
+ if (auto point = slice[i].as<PrimExpr>()) {
+ region.push_back(Range::FromMinExtent(point.value(),
IntImm(point.value().ty(), 1)));
+ } else {
+ auto descriptor = slice[i]
+ .as<ffi::Tuple<ffi::Optional<PrimExpr>,
ffi::Optional<PrimExpr>,
+ ffi::Optional<PrimExpr>>>()
+ .value();
+ PrimExpr start =
descriptor.get<0>().value_or(IntImm(buffer_ty->shape[i].ty(), 0));
+ PrimExpr stop = descriptor.get<1>().value_or(buffer_ty->shape[i]);
+ // Preserve the sole simplification performed by the former Python path.
+ region.push_back(Range::FromMinExtent(start, analyzer->Simplify(stop -
start)));
+ }
+ }
+ for (size_t i = slice.size(); i < buffer_ty->shape.size(); ++i) {
+ region.push_back(
+ Range::FromMinExtent(IntImm(buffer_ty->shape[i].ty(), 0),
buffer_ty->shape[i]));
+ }
+ return BufferRegion(buffer, region);
+}
+
+} // namespace
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ namespace refl = tvm::ffi::reflection;
+ BufferTypeNode::RegisterReflection();
+ refl::TypeAttrDef<BufferTypeNode>().def("__subscript_expr_realize__",
RealizeBufferSubscript);
+}
using IndexMod = tirx::FloorModNode;
using IndexDiv = tirx::FloorDivNode;
diff --git a/src/tirx/script/printer/buffer.cc
b/src/tirx/script/printer/buffer.cc
index a3e2c462e4..c0d58a90c9 100644
--- a/src/tirx/script/printer/buffer.cc
+++ b/src/tirx/script/printer/buffer.cc
@@ -381,6 +381,16 @@ ffi::Array<Doc> BufferIndices(const ffi::Array<PrimExpr>&
indices, const AccessP
return indices_doc;
}
+ffi::Array<Doc> BufferLoadIndices(const ffi::Array<PrimExpr>& indices, const
AccessPath& p,
+ const IRDocsifier& d) {
+ ffi::Array<Doc> indices_doc;
+ indices_doc.reserve(indices.size());
+ for (size_t i = 0; i < indices.size(); ++i) {
+ indices_doc.push_back(d->AsDoc<ExprDoc>(indices[i],
p->Attr("indices")->ArrayItem(i)));
+ }
+ return indices_doc;
+}
+
ffi::Array<Doc> BufferSlices(const ffi::Array<Range>& region, const
AccessPath& p,
const IRDocsifier& d) {
int n = region.size();
@@ -445,7 +455,7 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
return doc.value();
}
- return buffer[BufferIndices(load->indices, p->Attr("indices"), d)];
+ return buffer[BufferLoadIndices(load->indices, p->Attr("indices"),
d)];
});
TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
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 acd33f9814..e4cc34f879 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], B[0:1024])
+ expected = (A[0:1024].to_expr(), B[0:1024].to_expr())
_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], B[0:32, 0:32])
+ expected = (A[0:1024].to_expr(), B[0:32, 0:32].to_expr())
_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], B[0:1024])
+ expected = (A[0:32, 0:32].to_expr(), B[0:1024].to_expr())
_check_memcpy_results(func, expected)
diff --git a/tests/python/tirx/test_parser_printer.py
b/tests/python/tirx/test_parser_printer.py
index 487594ca44..5023290d70 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])
+ assert_structural_equal(reg, A[:, :, 8:16].to_expr())
# a None dim passes an int pick straight through (int → extent-1 region),
# while the chunked dim still narrows to its picked chunk.
@@ -2257,20 +2257,39 @@ 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]
+ br = buf[32:64, 0:32].to_expr()
assert isinstance(br, BufferRegion)
assert br.buffer.same_as(buf)
assert int(br.region[0].extent) == 32
assert int(br.region[1].extent) == 32
+def test_global_call_realizes_buffer_elements():
+ @I.ir_module(s_tir=True)
+ class Module:
+ @T.prim_func(private=True, s_tir=True)
+ def add(a: T.float32, b: T.float32) -> T.float32:
+ return a + b
+
+ @T.prim_func(s_tir=True)
+ def main(
+ A: T.Buffer((16,), "float32"),
+ B: T.Buffer((16,), "float32"),
+ C: T.Buffer((16,), "float32"),
+ ):
+ for i in range(16):
+ C[i] = Module.add(A[i], B[i])
+
+ 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]
+ br1 = buf[32:64, 0:32].to_expr()
assert isinstance(br1, BufferRegion)
# BufferRegion chained slice
diff --git a/tests/python/tvmscript/test_tvmscript_printer_tir.py
b/tests/python/tvmscript/test_tvmscript_printer_tir.py
index 0ecadaba67..8be6466b2b 100644
--- a/tests/python/tvmscript/test_tvmscript_printer_tir.py
+++ b/tests/python/tvmscript/test_tvmscript_printer_tir.py
@@ -1129,7 +1129,7 @@ def test_vload_with_explicit_scalable_data_type():
@T.prim_func(s_tir=True)
def main(A: T.Buffer((128,), "float32"), B: T.Buffer((128,), "float32")):
- B[0:T.vscale() * 4] = A[0:T.vscale() * 4]
+ B[0:T.vscale() * 4] = A[T.Ramp(0, 1, T.vscale() * 4)]
"""
_assert_print(main, expected_output)
@@ -1149,7 +1149,7 @@ def test_vectorize_llvm_pure_intrin():
@T.prim_func(s_tir=True)
def main(A: T.Buffer((4,), "float32"), B: T.Buffer((4,), "float32")):
- A[0:4] = T.call_llvm_pure_intrin("float32x4", "llvm.sqrt", B[0:4])
+ A[0:4] = T.call_llvm_pure_intrin("float32x4", "llvm.sqrt", B[T.Ramp(0, 1,
4)])
"""
_assert_print(main, expected_output)
diff --git a/tests/python/tvmscript/test_tvmscript_roundtrip.py
b/tests/python/tvmscript/test_tvmscript_roundtrip.py
index ddc2c09603..b55d5279f9 100644
--- a/tests/python/tvmscript/test_tvmscript_roundtrip.py
+++ b/tests/python/tvmscript/test_tvmscript_roundtrip.py
@@ -2425,9 +2425,9 @@ def buffer_ramp_access_as_slice_index():
for i in range(128):
A[i : i + 1 : 1] = i
for i in range(4):
- B[i * 32 : i * 32 + 32] = A[i * 32 : i * 32 + 32 : 1] +
T.broadcast(1.0, 32)
+ B[i * 32 : i * 32 + 32] = A[T.Ramp(i * 32, 1, 32)] +
T.broadcast(1.0, 32)
for i in range(4):
- C[i : i + 128 : 4] = B[i : i + 128 : 4] + T.broadcast(1.0, 32)
+ C[i : i + 128 : 4] = B[T.Ramp(i, 4, 32)] + T.broadcast(1.0, 32)
return buffer_ramp_access
