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 cec83d1bad [Fix][Relax] Track lowered reshape storage aliases (#20134)
cec83d1bad is described below
commit cec83d1badecf35ba73dc206dc5d206ad4eea12e
Author: Akaash Parthasarathy <[email protected]>
AuthorDate: Tue Aug 18 22:07:59 2026 -0700
[Fix][Relax] Track lowered reshape storage aliases (#20134)
`StaticPlanBlockMemory` tracks `relax.reshape` as an alias of its input
storage, but did not recognize direct calls to `vm.builtin.reshape`.
`DispatchSortScan` emits these packed reshape calls before static memory
planning. Without alias tracking, the planner could reuse the input
storage while the reshaped view was still live, allowing later
allocations to overwrite its contents.
This PR:
- Recognizes `vm.builtin.reshape` as a storage-aliasing operation
- Preserves the input storage token through the reshape result
- Updates the dynamic-output expected IR for the corrected planning path
---
src/relax/transform/static_plan_block_memory.cc | 6 +++-
.../test_transform_static_plan_block_memory.py | 42 +++++++++++++++++++++-
2 files changed, 46 insertions(+), 2 deletions(-)
diff --git a/src/relax/transform/static_plan_block_memory.cc
b/src/relax/transform/static_plan_block_memory.cc
index 7d20152ab2..b4782c7ca4 100644
--- a/src/relax/transform/static_plan_block_memory.cc
+++ b/src/relax/transform/static_plan_block_memory.cc
@@ -334,7 +334,11 @@ bool IsInplaceMemoryOp(const Expr& op) {
static const Op& reshape_op = Op::Get("relax.reshape");
static const Op& view_op = Op::Get("relax.memory.view");
static const Op& ensure_zero_offset_op =
Op::Get("relax.memory.ensure_zero_offset");
- return op.same_as(reshape_op) || op.same_as(view_op) ||
op.same_as(ensure_zero_offset_op);
+ const auto* extern_func = op.as<ExternFuncNode>();
+ bool is_builtin_reshape =
+ extern_func != nullptr && extern_func->global_symbol ==
"vm.builtin.reshape";
+ return op.same_as(reshape_op) || op.same_as(view_op) ||
op.same_as(ensure_zero_offset_op) ||
+ is_builtin_reshape;
}
/*! \brief The base class for the storage allocation visitor. */
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 fc5ded6594..2bcf1adf49 100644
--- a/tests/python/relax/test_transform_static_plan_block_memory.py
+++ b/tests/python/relax/test_transform_static_plan_block_memory.py
@@ -1658,7 +1658,10 @@ def test_add():
R.dtype("uint8"),
)
storage1: R.Any = R.memory.alloc_storage(
- R.shape([128 * vocab_size]), R.prim_value(0), R.str("global"),
R.dtype("float32")
+ R.shape([32 * vocab_size * 4]),
+ R.prim_value(0),
+ R.str("global"),
+ R.dtype("float32"),
)
alloc1: R.Tensor((batch_size, vocab_size), dtype="float32") =
R.memory.alloc_tensor(
storage1, R.prim_value(0), R.shape([batch_size, vocab_size]),
R.dtype("float32")
@@ -1732,6 +1735,43 @@ def test_view():
tvm.ir.assert_structural_equal(after, Expected)
+def test_builtin_reshape_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)
+
+ @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)
+ reshaped = R.call_packed(
+ "vm.builtin.reshape",
+ alloc,
+ R.shape([16]),
+ ty_args=R.Tensor((16,), "float32"),
+ )
+ alloc1 = R.builtin.alloc_tensor(R.shape([16]), "float32", 0)
+ cls.copy(reshaped, alloc1)
+ alloc2 = R.builtin.alloc_tensor(R.shape([16]), "float32", 0)
+ cls.copy(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_with_dataflow():
@I.ir_module
class Before:
