Skip to content

[REFACTOR][IR] Unify expression subscription realization - #20246

Merged
tqchen merged 6 commits into
apache:mainfrom
tqchen:tvm-unified-subscriptproxy
Aug 31, 2026
Merged

[REFACTOR][IR] Unify expression subscription realization#20246
tqchen merged 6 commits into
apache:mainfrom
tqchen:tvm-unified-subscriptproxy

Conversation

@tqchen

@tqchen tqchen commented Aug 31, 2026

Copy link
Copy Markdown
Member

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

@tqchen
tqchen force-pushed the tvm-unified-subscriptproxy branch 2 times, most recently from 9c8b881 to 1ad479f Compare August 31, 2026 20:32
@tqchen
tqchen force-pushed the tvm-unified-subscriptproxy branch from 1ad479f to 46ef80d Compare August 31, 2026 21:19
@tqchen
tqchen merged commit 71f0271 into apache:main Aug 31, 2026
8 checks passed
tqchen pushed a commit that referenced this pull request Sep 1, 2026
## 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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants