This is an automated email from the ASF dual-hosted git repository.
spectrometerHBH 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 e0bbe2c102 [FIX][TIRx] Traverse pointer expressions in tile calls
(#20089)
e0bbe2c102 is described below
commit e0bbe2c102e485b6e947be06a3d6cc1431824747
Author: Hongyi Jin <[email protected]>
AuthorDate: Wed Aug 5 19:37:32 2026 -0400
[FIX][TIRx] Traverse pointer expressions in tile calls (#20089)
## Motivation
`TilePrimitiveCall` arguments and config values can contain any
`tvm.ir.Expr`, not only `PrimExpr`. In particular, `T.address_of(...)`
produces a pointer expression.
The Python `StmtVisitor` and `StmtMutator` only traversed values
accepted by `tvm.ir.is_prim_expr`. Pointer-valued operands were
therefore skipped entirely. Analyses could miss buffer uses below the
pointer expression, and rewrites could leave stale buffer references
behind.
## Example
A tile call can pass an mbarrier address through its config:
```python
@T.prim_func
def copy_async(
A: T.Buffer((8,), "float16"),
B: T.Buffer((8,), "float16"),
mbar: T.Buffer((1,), "uint64"),
):
Tx.copy_async(
B[:],
A[:],
dispatch="tma_auto",
mbar=T.address_of(mbar[0]),
)
```
Before this PR, a `StmtExprVisitor` did not reach the `BufferLoad` for
`mbar[0]`, and a `StmtExprMutator` could not replace it. After this PR,
both visitor and mutator traverse the pointer expression normally.
## What changed
- Treat every `tvm.ir.Expr` in `TilePrimitiveCall` arguments and config
values as traversable.
- Add a regression test that verifies both visiting and replacing the
buffer load beneath `T.address_of`.
## Tests
- `python -m pytest tests/python/tirx/transform/test_stmt_functor.py -q`
(20 passed)
- `pre-commit run --files python/tvm/tirx/stmt_functor.py
tests/python/tirx/transform/test_stmt_functor.py`
Ported from
[mlc-ai/tvm@98b7dec](https://github.com/mlc-ai/tvm/commit/98b7decf5a741d9d47210f7d204085d8acba81b3).
---
python/tvm/tirx/stmt_functor.py | 8 ++---
tests/python/tirx/transform/test_stmt_functor.py | 43 ++++++++++++++++++++++++
2 files changed, 47 insertions(+), 4 deletions(-)
diff --git a/python/tvm/tirx/stmt_functor.py b/python/tvm/tirx/stmt_functor.py
index db418f18b6..1f9755b22a 100644
--- a/python/tvm/tirx/stmt_functor.py
+++ b/python/tvm/tirx/stmt_functor.py
@@ -363,14 +363,14 @@ class StmtVisitor(StmtFunctor):
def visit_op_call_(self, op):
"""Visitor implementation for TilePrimitiveCall."""
for arg in op.args:
- if tvm.ir.is_prim_expr(arg):
+ if isinstance(arg, tvm.ir.Expr):
self.visit_expr(arg)
elif isinstance(arg, tvm.tirx.Stmt):
self.visit_stmt(arg)
elif isinstance(arg, tvm.tirx.BufferRegion):
self.visit_buffer_region_(arg)
for value in op.config.values():
- if tvm.ir.is_prim_expr(value):
+ if isinstance(value, tvm.ir.Expr):
self.visit_expr(value)
elif isinstance(value, tvm.tirx.Stmt):
self.visit_stmt(value)
@@ -842,7 +842,7 @@ class StmtMutator(StmtFunctor):
args_changed = False
for arg in op.args:
- if tvm.ir.is_prim_expr(arg):
+ if isinstance(arg, tvm.ir.Expr):
new_arg = self.visit_expr(arg)
elif isinstance(arg, tvm.tirx.Stmt):
new_arg = self.visit_stmt(arg)
@@ -859,7 +859,7 @@ class StmtMutator(StmtFunctor):
new_config = {}
config_changed = False
for key, value in op.config.items():
- if tvm.ir.is_prim_expr(value):
+ if isinstance(value, tvm.ir.Expr):
new_value = self.visit_expr(value)
elif isinstance(value, tvm.tirx.Stmt):
new_value = self.visit_stmt(value)
diff --git a/tests/python/tirx/transform/test_stmt_functor.py
b/tests/python/tirx/transform/test_stmt_functor.py
index ae7e1aa8fa..3da352a18c 100644
--- a/tests/python/tirx/transform/test_stmt_functor.py
+++ b/tests/python/tirx/transform/test_stmt_functor.py
@@ -1183,6 +1183,49 @@ def test_op_call_config_mutated():
)
+def test_op_call_pointer_config_visited_and_mutated():
+ """Pointer-valued config expressions participate in Python traversal."""
+
+ @T.prim_func
+ def copy_async(
+ A: T.Buffer((8,), "float16"),
+ B: T.Buffer((8,), "float16"),
+ mbar: T.Buffer((1,), "uint64"),
+ ):
+ Tx.copy_async(B[:], A[:], dispatch="tma_auto",
mbar=T.address_of(mbar[0]))
+
+ op_call = copy_async.body
+ assert isinstance(op_call, tir.TilePrimitiveCall)
+ mbar_buffer = copy_async.buffer_map[copy_async.params[2]]
+
+ class LoadCollector(StmtExprVisitor):
+ def __init__(self):
+ super().__init__()
+ self.buffers = []
+
+ def visit_buffer_load_(self, op):
+ self.buffers.append(op.buffer)
+ return super().visit_buffer_load_(op)
+
+ collector = LoadCollector()
+ collector.visit_stmt(op_call)
+ assert any(buffer.same_as(mbar_buffer) for buffer in collector.buffers)
+
+ replacement = tir.decl_buffer((1,), "uint64", name="replacement")
+
+ class ReplaceMbarLoad(StmtExprMutator):
+ def visit_buffer_load_(self, op):
+ new_op = super().visit_buffer_load_(op)
+ if op.buffer.same_as(mbar_buffer):
+ return tir.BufferLoad(replacement, new_op.indices,
new_op.predicate)
+ return new_op
+
+ updated = ReplaceMbarLoad().visit_stmt(op_call)
+ mbar_load = updated.config["mbar"].args[0]
+ assert isinstance(mbar_load, tir.BufferLoad)
+ assert mbar_load.buffer.same_as(replacement)
+
+
def test_op_call_nested_config_visited_and_substituted():
"""Nested selector arrays participate in the core visitor and mutator."""
from tvm.tirx.stmt_functor import post_order_visit, substitute
