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 aebdc51a6f [FIX][TIRx] Remap typed buffer expressions during
specialization (#20090)
aebdc51a6f is described below
commit aebdc51a6f2ebb6da2ce2783cefb198fd380f12a
Author: Hongyi Jin <[email protected]>
AuthorDate: Wed Aug 5 19:37:39 2026 -0400
[FIX][TIRx] Remap typed buffer expressions during specialization (#20090)
## Motivation
TIRx buffers are represented by typed variables. When
`PrimFunc.specialize` remaps a buffer, every expression that refers to
its `BufferVar` must be remapped as well.
`PrimFuncSpecializer::VisitExpr_(VarNode*)` returned variables that were
not in the scalar specialization map unchanged. This bypassed the base
`StmtExprMutator`, which is responsible for recognizing and remapping
typed buffer variables. As a result, an alias buffer could retain the
original data variable after specialization instead of referring to the
specialized buffer.
## Example
Consider a buffer alias that shares `A.data`:
```python
@T.prim_func(private=True, s_tir=True)
def before(A_handle: T.handle, n: T.int32):
A = T.match_buffer(A_handle, (n,), "int32")
A_alias = T.decl_buffer((n,), "int32", data=A.data)
A_alias[n - 1] = 42
after = before.specialize({before.params[1]: 8})
```
The specialized function should still preserve the alias:
```python
A = T.match_buffer(A_handle, (8,), "int32")
A_alias = T.decl_buffer((8,), "int32", data=A.data)
A_alias[7] = 42
```
Before this PR, `A_alias.data` could retain the pre-specialization
`BufferVar`. After this PR, it is remapped to the specialized `A.data`.
## What changed
- Delegate unmatched `VarNode` expressions to the base
`StmtExprMutator`.
- Add a regression test that specializes a symbolic extent and checks
that the alias relationship is preserved.
## Tests
- Fresh default CMake/Ninja build
- `python -m pytest tests/python/tirx-base/test_tir_specialize.py -q`
(11 passed)
- `pre-commit run --files src/tirx/ir/specialize.cc
tests/python/tirx-base/test_tir_specialize.py`
Ported from
[mlc-ai/tvm@14bf144](https://github.com/mlc-ai/tvm/commit/14bf144057ce4ce95c655f20278bc70a0efe81be).
---
src/tirx/ir/specialize.cc | 7 ++++---
tests/python/tirx-base/test_tir_specialize.py | 18 ++++++++++++++++++
2 files changed, 22 insertions(+), 3 deletions(-)
diff --git a/src/tirx/ir/specialize.cc b/src/tirx/ir/specialize.cc
index 3890829dc1..d1a687af5c 100644
--- a/src/tirx/ir/specialize.cc
+++ b/src/tirx/ir/specialize.cc
@@ -174,12 +174,13 @@ class PrimFuncSpecializer : public StmtExprMutator {
BufferVar VisitBufferUse(const BufferVar& buffer) final { return
GetNewBuffer(buffer); }
Expr VisitExpr_(const VarNode* op) final {
+ Var var = ffi::GetRef<Var>(op);
if (constrained_buffer_params_.count(op)) {
- return ffi::GetRef<Var>(op);
+ return var;
}
- auto it = var_map_.find(ffi::GetRef<Var>(op));
+ auto it = var_map_.find(var);
if (it == var_map_.end()) {
- return ffi::GetRef<Var>(op);
+ return StmtExprMutator::VisitExpr_(op);
} else {
return it->second;
}
diff --git a/tests/python/tirx-base/test_tir_specialize.py
b/tests/python/tirx-base/test_tir_specialize.py
index f47a6dc591..4dffc8dc11 100644
--- a/tests/python/tirx-base/test_tir_specialize.py
+++ b/tests/python/tirx-base/test_tir_specialize.py
@@ -266,6 +266,24 @@ def test_specialize_decl_buffer():
tvm.ir.assert_structural_equal(expected, after)
+def test_specialize_preserves_decl_buffer_alias():
+ @T.prim_func(private=True, s_tir=True)
+ def before(A_handle: T.handle, n: T.int32):
+ A = T.match_buffer(A_handle, (n,), "int32")
+ A_flat = T.decl_buffer((n,), "int32", data=A.data)
+ A_flat[n - 1] = 42
+
+ @T.prim_func(private=True, s_tir=True)
+ def expected(A_handle: T.handle):
+ A = T.match_buffer(A_handle, (8,), "int32")
+ A_flat = T.decl_buffer((8,), "int32", data=A.data)
+ A_flat[7] = 42
+
+ after = before.specialize({before.params[1]: 8})
+
+ tvm.ir.assert_structural_equal(expected, after)
+
+
def test_specialize_buffer_var_to_var():
"""A buffer var may be remapped by specialization
