From 0bba1b18d8f4460a595d061e2ae9c2ebe1a4c54b Mon Sep 17 00:00:00 2001 From: Akaash Parthasarathy Date: Sat, 18 Jul 2026 16:48:51 -0400 Subject: [PATCH 1/2] [FIX][RELAX] Track lowered reshape storage aliases --- .../transform/static_plan_block_memory.cc | 6 ++- ...test_transform_static_plan_block_memory.py | 37 +++++++++++++++++++ 2 files changed, 42 insertions(+), 1 deletion(-) diff --git a/src/relax/transform/static_plan_block_memory.cc b/src/relax/transform/static_plan_block_memory.cc index 7d20152ab2ee..b4782c7ca450 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(); + 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 fc5ded659409..52bcd0ff6d32 100644 --- a/tests/python/relax/test_transform_static_plan_block_memory.py +++ b/tests/python/relax/test_transform_static_plan_block_memory.py @@ -1732,6 +1732,43 @@ def main() -> R.Tensor((128,), dtype="float32"): 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: From f0ecefd4d71ea1e3b502a4979c5a35c2375efcd0 Mon Sep 17 00:00:00 2001 From: Akaash Parthasarathy Date: Mon, 17 Aug 2026 02:22:09 -0400 Subject: [PATCH 2/2] [TESTS][RELAX] Update memory planning expectation --- .../python/relax/test_transform_static_plan_block_memory.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) 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 52bcd0ff6d32..2bcf1adf4988 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 main(probs: R.Tensor(("batch_size", "vocab_size"), dtype="float32")) -> R.Te 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")