This is an automated email from the ASF dual-hosted git repository.
tlopex 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 7eaa8eb899 [Fix][Relax] Preserve match-cast storage liveness (#20220)
7eaa8eb899 is described below
commit 7eaa8eb89901d510ffebbefb613294bfd56be66d
Author: Zupeng Wang <[email protected]>
AuthorDate: Sun Aug 30 11:03:08 2026 +0800
[Fix][Relax] Preserve match-cast storage liveness (#20220)
Fixes #20190.
Fixes #20207.
`StaticPlanBlockMemory` propagates storage tokens through ordinary
variable
bindings and lowered reshape/view operations, but not through
`MatchCast`.
A `MatchCast` refines type information without allocating or copying
storage,
so losing the token can make the planner release and reuse the source
allocation
while the alias is still live. This causes silent wrong-code in the
default
Relax pipeline.
This change propagates storage tokens through `MatchCast` bindings and
adds a
regression test showing that the aliased allocation remains live until
its
downstream use.
Tests:
- `python -m pytest -q
tests/python/relax/test_transform_static_plan_block_memory.py`
(`29 passed`)
- `pre-commit run --files
src/relax/transform/static_plan_block_memory.cc
tests/python/relax/test_transform_static_plan_block_memory.py`
- End-to-end compiled-VM reproducers for #20190 and #20207, comparing
NumPy
references before and after the patch
Signed-off-by: Zupeng Wang <[email protected]>
---
src/relax/transform/static_plan_block_memory.cc | 6 ++++
.../test_transform_static_plan_block_memory.py | 40 ++++++++++++++++++++++
2 files changed, 46 insertions(+)
diff --git a/src/relax/transform/static_plan_block_memory.cc
b/src/relax/transform/static_plan_block_memory.cc
index b4782c7ca4..ffb0851f83 100644
--- a/src/relax/transform/static_plan_block_memory.cc
+++ b/src/relax/transform/static_plan_block_memory.cc
@@ -361,6 +361,12 @@ class StorageAllocatorBaseVisitor : public ExprVisitor {
SetTokens(binding->var.get(), token_map_[binding->value.get()]);
}
+ void VisitBinding_(const MatchCastNode* binding) override {
+ ExprVisitor::VisitBinding_(binding);
+ // MatchCast refines the type without changing the underlying storage.
+ SetTokens(binding->var.get(), token_map_[binding->value.get()]);
+ }
+
void VisitBindingBlock_(const DataflowBlockNode* block) override {
// We maintain a block stack for token allocation-site and use-site check.
block_stack_.push_back(block);
diff --git a/tests/python/relax/test_transform_static_plan_block_memory.py
b/tests/python/relax/test_transform_static_plan_block_memory.py
index 2bcf1adf49..a7583eb5bd 100644
--- a/tests/python/relax/test_transform_static_plan_block_memory.py
+++ b/tests/python/relax/test_transform_static_plan_block_memory.py
@@ -1735,6 +1735,46 @@ def test_view():
tvm.ir.assert_structural_equal(after, Expected)
+def test_match_cast_preserves_storage_liveness():
+ @I.ir_module
+ class Before:
+ @T.prim_func(s_tir=True)
+ def copy(A: T.Buffer((16,), "float32"), B: T.Buffer((16,), "float32")):
+ T.evaluate(0)
+
+ @T.prim_func(s_tir=True)
+ def add(
+ A: T.Buffer((16,), "float32"),
+ B: T.Buffer((16,), "float32"),
+ C: T.Buffer((16,), "float32"),
+ ):
+ T.evaluate(0)
+
+ @R.function
+ def main(x: R.Tensor((16,), "float32")) -> R.Tensor((16,), "float32"):
+ R.func_attr({"relax.force_pure": True})
+ cls = Before
+ alloc = R.builtin.alloc_tensor(R.shape([16]), "float32", 0)
+ cls.copy(x, alloc)
+ checked = R.match_cast(alloc, R.Tensor((16,), "float32"))
+ alloc1 = R.builtin.alloc_tensor(R.shape([16]), "float32", 0)
+ cls.copy(x, alloc1)
+ alloc2 = R.builtin.alloc_tensor(R.shape([16]), "float32", 0)
+ cls.add(checked, alloc1, alloc2)
+ return alloc2
+
+ after = relax.transform.StaticPlanBlockMemory()(Before)
+ alloc_storage_op = tvm.ir.Op.get("relax.memory.alloc_storage")
+ storage_allocations = []
+
+ def collect_storage_allocations(expr):
+ if isinstance(expr, relax.Call) and expr.op.same_as(alloc_storage_op):
+ storage_allocations.append(expr)
+
+ relax.analysis.post_order_visit(after["main"], collect_storage_allocations)
+ assert len(storage_allocations) == 2
+
+
def test_builtin_reshape_preserves_storage_liveness():
@I.ir_module
class Before:
