[LowerVTCMAlloc] Move LowerVtcmAlloc to after StorageRewrite (apache#12364)

* [LowerVTCMAlloc] Move LowerVtcmAlloc to after StorageRewrite

Vtcm allocations were being moved inside loops even if they were
originally allocated outside of the loops. Normally
PlanAndUpdateBufferAllocationLocation moves allocations as close to use
as possible and then StorageRewrite moves them back out as far as
possible. However, with Vtcm allocation,
PlanAndUpdateBufferAllocationLocation would move the Vtcm allocation
close to the compute, then LowerVtcm would convert the allocation to a
LetStmt. StorageRewrite would not move this LetStmt as it only handles
allocations. Moving LowerVtcmAlloc to after StorageRewrite ensures that
the vtcm allocations are in their final spot before converting them to a
LetStmt.

* fix issues with tagging and storage rewrite

Author: Tristan Konolige

src/driver/driver_api.cc

@@ -204,7 +204,6 @@ Array<tvm::transform::Pass> CreatePassList(bool disable_loop_partition) {
   pass_list.push_back(tir::transform::InjectSoftwarePipeline());
   pass_list.push_back(tir::transform::LowerOpaqueBlock());
   pass_list.push_back(tir::transform::FlattenBuffer());
-  pass_list.push_back(tir::transform::LowerVtcmAlloc());
   pass_list.push_back(tir::transform::BF16Legalize());
   pass_list.push_back(tir::transform::NarrowDataType(32));
   pass_list.push_back(tir::transform::Simplify());
@@ -223,6 +222,8 @@ Array<tvm::transform::Pass> CreatePassList(bool disable_loop_partition) {
   if (!disable_storage_rewrite) {
     pass_list.push_back(tir::transform::StorageRewrite());
   }
+  // LowerVtcmAlloc must occur after any transformations that modify memory allocation locations
+  pass_list.push_back(tir::transform::LowerVtcmAlloc());
   pass_list.push_back(tir::transform::UnrollLoop());
 
   // Add user-defined phase-2 passes

src/tir/transforms/storage_rewrite.cc

@@ -583,8 +583,10 @@ class StoragePlanRewriter : public StmtExprMutator {
   };
 
   // Checks whether the storage_scope is especially tagged for a specific memory.
+  // Special memory is all combined into a single allocation.
   bool IsSpecialTaggedMemory(const StorageScope& scope) {
-    return scope.tag.length() != 0 && scope.tag != ".dyn" && scope.tag != ".workspace";
+    return scope.tag.length() != 0 && scope.tag != ".dyn" && scope.tag != ".workspace" &&
+           scope.tag != ".vtcm";
   }
 
   // Alllocate entry of node.
@@ -655,8 +657,6 @@ class StoragePlanRewriter : public StmtExprMutator {
 
         if (e->allocs.size() == 1) {
           // simply use the original allocation.
-          PrimExpr sz = foldl([](PrimExpr a, PrimExpr b, Span span) { return mul(a, b, span); },
-                              make_const(DataType::Int(32), 1), e->allocs[0]->extents);
           e->new_alloc = Allocate(e->alloc_var, alloc_type, e->allocs[0]->extents,
                                   e->allocs[0]->condition, Evaluate(0));
           if (IsSpecialTaggedMemory(e->scope)) {