[Fix][Relax] Track lowered reshape storage aliases

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. */

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")
@@ -1732,6 +1735,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: