[TIR] Handle DeclBuffer in StorageFlatten's input

src/tir/transforms/storage_flatten.cc

@@ -1431,6 +1431,15 @@ class StorageFlattener : public StmtExprMutator {
     return body;
   }
 
+  Stmt VisitStmt_(const DeclBufferNode* op) final {
+    auto node = Downcast<DeclBuffer>(StmtExprMutator::VisitStmt_(op));
+    const BufferEntry& entry = GetBufferEntry(node->buffer);
+    if (!entry.flattened_buffer.same_as(node->buffer)) {
+      node.CopyOnWrite()->buffer = entry.flattened_buffer;
+    }
+    return std::move(node);
+  }
+
   // AllocateNodes may be present from tvm.tir.ir_builder.  This can
   // be simplified in the future by having AllocateNode hold a buffer,
   // rather than a buffer_var.

tests/python/unittest/test_tir_transform_storage_flatten.py

@@ -165,5 +165,21 @@ def test_flatten_tir():
     )  # StorageFlatten should do nothing to TIR functions
 
 
+class TestPreserveDeclBuffer(tvm.testing.CompareBeforeAfter):
+    transform = tvm.tir.transform.StorageFlatten(64)
+
+    def before():
+        T.func_attr({"from_legacy_te_schedule": True})
+        A = T.decl_buffer([16, 16], "float32")
+        for i, j in T.grid(16, 16):
+            A[i, j] = 0.0
+
+    def expected():
+        T.func_attr({"from_legacy_te_schedule": True})
+        A = T.decl_buffer([256], "float32")
+        for i, j in T.grid(16, 16):
+            A[i * 16 + j] = 0.0
+
+
 if __name__ == "__main__":
     tvm.testing.main()