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 dbb1157bc8 [FIX][IR] Complete lazy subscript realization (#20251)
dbb1157bc8 is described below
commit dbb1157bc8802a0069b3a99dc405ab17a01bfa91
Author: Hongyi Jin <[email protected]>
AuthorDate: Tue Sep 1 07:44:52 2026 -0400
[FIX][IR] Complete lazy subscript realization (#20251)
## Motivation
Since #20246, expression subscription returns a lazy `SubscriptProxy`
until a consumer realizes the expression. Most parser expression
boundaries already honor that contract, but PTX predicate operand
coercion bypasses the generic realization path. As a result, code such
as `setp(..., buffer[0], ...)` dispatches the proxy itself and rejects
it as an invalid operand.
The delayed realization also means source metadata must travel with the
proxy. Without that, realizing `A[0]` creates a `TensorLoad` without the
exact subscription span, so downstream diagnostics such as TIRx
racecheck and GDN cannot attribute the access to its source location.
## Summary
- realize lazy subscript operands used directly by PTX predicate inputs
and destinations
- preserve parser source spans through `SubscriptProxy` realization
- propagate the span through tuple, Relax, and TIRx subscription
callbacks
This completes the lazy-subscript contract introduced by #20246 and
fixes downstream regressions exposed by the TIRx migration in #20247.
## Testing
- `cmake --build build -j 16`
- `python -m pytest -q
tests/python/tvmscript/test_tvmscript_parser_source.py::test_parser_attaches_span_to_nested_tensor_load
tests/python/tvmscript/test_tvmscript_parser_source.py::test_parser_attaches_span_to_direct_call
tests/python/tirx/codegen/test_ptx_dialect.py::test_ptx_comparison_selection_dispatch`
- downstream wiki-kernel integration: 5 passed
- downstream full `tirx_tools` gate: 2933 passed
---
python/tvm/backend/cuda/ptx/engine.py | 4 ++--
python/tvm/ir/expr.py | 21 ++++++++++++++++++---
python/tvm/script/parser/core/evaluator.py | 17 ++++++++++-------
python/tvm/script/parser/core/parser.py | 8 ++++++--
src/ir/subscript_proxy.cc | 9 +++++----
src/relax/ir/dependent_type.cc | 5 +++--
src/tirx/ir/buffer.cc | 5 +++--
tests/python/tirx/codegen/test_ptx_dialect.py | 5 +++--
.../tvmscript/test_tvmscript_parser_source.py | 22 +++++++++++++++++++++-
9 files changed, 71 insertions(+), 25 deletions(-)
diff --git a/python/tvm/backend/cuda/ptx/engine.py
b/python/tvm/backend/cuda/ptx/engine.py
index 2529065127..5ffdecc785 100644
--- a/python/tvm/backend/cuda/ptx/engine.py
+++ b/python/tvm/backend/cuda/ptx/engine.py
@@ -444,7 +444,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,) = values
+ (value,) = [_realize_operand(value) for value in 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 +457,7 @@ def _coerce_pred_operand(entry, slot, values):
)
return values
if isinstance(value, PredArg):
- value = getattr(value.value, "scalar", value.value)
+ value = _realize_operand(getattr(value.value, "scalar", value.value))
if isinstance(value, bool | int):
return [const(int(value), "uint32")]
ty = getattr(value, "ty", None)
diff --git a/python/tvm/ir/expr.py b/python/tvm/ir/expr.py
index 3cbd19c356..5005737e26 100644
--- a/python/tvm/ir/expr.py
+++ b/python/tvm/ir/expr.py
@@ -39,7 +39,9 @@ class Expr(Node):
# 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)])
+ 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
@@ -444,7 +446,7 @@ class SubscriptProxy(_ExprCallable, ExprOperand,
ObjectConvertible):
containing a slice must be realized before applying a region subscript.
"""
- __slots__ = ("_result", "_slice", "_source")
+ __slots__ = ("_result", "_slice", "_source", "_span")
__hash__ = object.__hash__
def __init__(self, source: Expr, index):
@@ -453,12 +455,23 @@ class SubscriptProxy(_ExprCallable, ExprOperand,
ObjectConvertible):
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,)
@@ -487,7 +500,9 @@ class SubscriptProxy(_ExprCallable, ExprOperand,
ObjectConvertible):
"""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._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")
diff --git a/python/tvm/script/parser/core/evaluator.py
b/python/tvm/script/parser/core/evaluator.py
index 06f6dc9b09..e22abf4b9a 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) -> doc.Name:
+ def _add_intermediate_result(self, value: Any, node: doc.AST) -> doc.Name:
"""Add intermediate result during evaluation into value table.
Parameters
@@ -139,13 +139,16 @@ 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)
+ value = self.parser.annotate_current_source_span(value, node)
name = f"__tvm_tmp_value_{self.new_value_count}"
self.new_value_count += 1
self.value_table[name] = value
@@ -193,13 +196,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)
+ return self._add_intermediate_result(value, node)
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)
+ return self._add_intermediate_result(value, node)
args = []
if (
@@ -262,7 +265,7 @@ class ExprEvaluator:
if isinstance(node, doc.ListComp | doc.SetComp | doc.DictComp):
value = self._eval_expr(node)
- return self._add_intermediate_result(value)
+ return self._add_intermediate_result(value, node)
fields = {}
for field in node.__class__._FIELDS: # pylint:
disable=protected-access
@@ -284,7 +287,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)
+ return self._add_intermediate_result(value, node)
def _eval_lambda(self, node: doc.Lambda) -> Any:
"""The doc AST lambda node evaluating method.
@@ -303,7 +306,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)
+ return self._add_intermediate_result(value, node)
def _eval_bool_op(self, node: doc.BoolOp) -> Any:
"""The doc AST boolean operator node evaluating method.
diff --git a/python/tvm/script/parser/core/parser.py
b/python/tvm/script/parser/core/parser.py
index 7d0f2887df..94c74f24a1 100644
--- a/python/tvm/script/parser/core/parser.py
+++ b/python/tvm/script/parser/core/parser.py
@@ -596,8 +596,12 @@ 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) -> Any:
+ def annotate_current_source_span(self, value: Any, node: doc.AST | None =
None) -> 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
@@ -634,7 +638,7 @@ class Parser(doc.NodeVisitor):
# 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)
+ return self.annotate_current_source_span(_realize_operand(value), node)
def _duplicate_lhs_check(self, target: doc.expr) -> bool | set[str]:
"""Check whether duplicate lhs exists in assignment.
diff --git a/src/ir/subscript_proxy.cc b/src/ir/subscript_proxy.cc
index a92c961c61..c6e8e75109 100644
--- a/src/ir/subscript_proxy.cc
+++ b/src/ir/subscript_proxy.cc
@@ -36,7 +36,8 @@ 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 {
+ "__subscript_expr_realize__",
+ [](Expr value, SubscriptSlice slice, Span span) -> ffi::ObjectRef {
TVM_FFI_CHECK_EQ(slice.size(), 1, IndexError)
<< "A tuple expression requires exactly one index";
auto index = slice[0].as<PrimExpr>();
@@ -44,7 +45,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
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));
+ 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";
@@ -53,14 +54,14 @@ TVM_FFI_STATIC_INIT_BLOCK() {
<< "Type " << value->ty->GetTypeKey() << " does not support subscript";
});
refl::GlobalDef().def(
- "ir.SubscriptExprRealize", [](Expr value, SubscriptSlice slice) ->
ffi::ObjectRef {
+ "ir.SubscriptExprRealize", [](Expr value, SubscriptSlice slice, Span
span) -> 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>();
+ packed_realize.cast<ffi::Function>()(value, slice,
span).cast<ffi::ObjectRef>();
TVM_FFI_CHECK(result.defined(), TypeError)
<< "__subscript_expr_realize__ for type " <<
value->ty->GetTypeKey()
<< " returned an undefined object";
diff --git a/src/relax/ir/dependent_type.cc b/src/relax/ir/dependent_type.cc
index 2ac4898551..ff0bce1fb6 100644
--- a/src/relax/ir/dependent_type.cc
+++ b/src/relax/ir/dependent_type.cc
@@ -42,7 +42,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
ffi::Array<ffi::Variant<
ffi::Tuple<ffi::Optional<PrimExpr>, ffi::Optional<PrimExpr>,
ffi::Optional<PrimExpr>>,
PrimExpr>>
- slice) -> ffi::ObjectRef {
+ slice,
+ Span span) -> ffi::ObjectRef {
TVM_FFI_CHECK_EQ(slice.size(), 1, IndexError)
<< "A Relax expression requires exactly one index";
auto index = slice[0].as<PrimExpr>();
@@ -50,7 +51,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
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));
+ return TupleGetItem(value, static_cast<int>(imm->value), span);
});
}
diff --git a/src/tirx/ir/buffer.cc b/src/tirx/ir/buffer.cc
index c715aedf0c..8c529e2a05 100644
--- a/src/tirx/ir/buffer.cc
+++ b/src/tirx/ir/buffer.cc
@@ -47,7 +47,8 @@ ffi::ObjectRef RealizeBufferSubscript(
ffi::Array<ffi::Variant<
ffi::Tuple<ffi::Optional<PrimExpr>, ffi::Optional<PrimExpr>,
ffi::Optional<PrimExpr>>,
PrimExpr>>
- slice) {
+ slice,
+ Span span) {
BufferVar buffer = value.as_or_throw<BufferVar>();
BufferType buffer_ty = buffer.type();
TVM_FFI_CHECK_LE(slice.size(), buffer_ty->shape.size(), IndexError)
@@ -70,7 +71,7 @@ ffi::ObjectRef RealizeBufferSubscript(
for (const auto& item : slice) {
indices.push_back(item.as<PrimExpr>().value());
}
- return BufferLoad(buffer, indices);
+ return BufferLoad(buffer, indices, span);
}
// Any slice or omitted trailing dimension denotes a region. Rejecting
diff --git a/tests/python/tirx/codegen/test_ptx_dialect.py
b/tests/python/tirx/codegen/test_ptx_dialect.py
index 010fb9ed92..743f5022e0 100644
--- a/tests/python/tirx/codegen/test_ptx_dialect.py
+++ b/tests/python/tirx/codegen/test_ptx_dialect.py
@@ -1356,12 +1356,13 @@ def test_ptx_comparison_selection_dispatch():
q = T.local_scalar("uint32")
d = T.local_scalar("uint32")
f = T.local_scalar("float32")
- T.ptx.setp.lt.s32(p, A[0], A[1]) # one destination
+ p_buffer = T.alloc_local((1,), "uint32")
+ T.ptx.setp.lt.s32(p_buffer[0], A[0], A[1]) # one destination
T.ptx.setp.lt.s32(p, q, A[0], A[1]) # ... two: `p|q`, chosen by arity
T.ptx.setp.lt.and_.s32(p, A[0], A[1], T.ptx.pred(q)) # ... plus a
BoolOp
T.ptx.setp.gt.or_.s32(p, q, A[2], A[3], T.ptx.pred(d)) # ... and both
T.ptx.set.lt.u32.f32(d, f, f) # writes a value, not a predicate
- T.ptx.selp.b32(d, d, p, T.ptx.pred(q)) # predicate selects
+ T.ptx.selp.b32(d, d, p, T.ptx.pred(p_buffer[0])) # predicate selects
T.ptx.slct.ftz.b32.f32(d, d, p, f) # a sign selects
# slct treats d/a/b independently as bit-size values. This mixes all
# three 32-bit carrier classes while c remains exactly .s32.
diff --git a/tests/python/tvmscript/test_tvmscript_parser_source.py
b/tests/python/tvmscript/test_tvmscript_parser_source.py
index da6fff1a87..31c7581d9c 100644
--- a/tests/python/tvmscript/test_tvmscript_parser_source.py
+++ b/tests/python/tvmscript/test_tvmscript_parser_source.py
@@ -24,7 +24,7 @@ import tvm_ffi
import tvm
import tvm.testing
-from tvm.ir import Call, SequentialSpan, assert_structural_equal
+from tvm.ir import Call, SequentialSpan, TensorLoad, assert_structural_equal
from tvm.script import tirx as T
from tvm.script.parser.core import doc_core as doc
from tvm.script.parser.core.diagnostics import Source
@@ -156,6 +156,26 @@ def test_parser_attaches_span_to_direct_call():
assert _span_range(call.span) == _span_range(source.to_span(call_ast))
+def test_parser_attaches_span_to_nested_tensor_load():
+ @_tirx_source
+ def nested_load():
+ source_buffer = T.alloc_buffer((1,), "int32")
+ output = T.alloc_buffer((1,), "int32")
+ output[0] = source_buffer[0] + 1
+
+ source = Source(nested_load)
+ load_ast = source.as_ast().body[0].body[-1].value.left
+ func = T.prim_func(nested_load)
+ load = _find_ir_node(
+ func,
+ lambda node: (
+ isinstance(node, TensorLoad) and getattr(node.source, "name",
None) == "source_buffer"
+ ),
+ )
+
+ assert _span_range(load.span) == _span_range(source.to_span(load_ast))
+
+
def test_parser_retains_inline_call_site_and_definition_spans():
def wait_impl(barrier):
T.cuda.mbarrier_wait(barrier, 0)
