[Language] Support tilelang alloc_var(dtype, init=x) (#1092)

* - carry existing local-var initializer map into OpaqueBlockLower, reattach it to
    generated Allocates and the PrimFunc attrs
  - thread the map through FlattenBuffer and StorageRewrite so flattened/merged
    allocations keep their tl.local_var_init annotations
  - teach annotation handling to accept scalar initializers, resolve buffers, and merge
    with existing stat

* lint fix

* enhance

* lint fix

* lint fix

Author: LeiWang1999

src/op/builtin.h

@@ -27,6 +27,7 @@ static constexpr const char *kWarpSpecializationScope =
     "kWarpSpecializationScope";
 static constexpr const char *kCustomWarpSpecialization =
     "kCustomWarpSpecialization";
+static constexpr const char *kLocalVarInit = "tl.local_var_init";
 } // namespace attr
 
 static constexpr const char *kDebugMergeSharedMemoryAllocations =

src/target/codegen_cuda.cc

@@ -2201,8 +2201,16 @@ void CodeGenTileLangCUDA::VisitStmt_(const AllocateNode *op) {
     } else if (scope == "local") {
       stream << ' ' << vid << '[' << constant_size << "];\n";
     } else if (scope == "local.var") {
-      stream << ' ' << vid << " = " << PrintExpr(tir::make_const(op->dtype, 0))
-             << ";\n";
+      PrimExpr init = tir::make_const(op->dtype, 0);
+      auto init_it = op->annotations.find(tl::attr::kLocalVarInit);
+      if (init_it != op->annotations.end()) {
+        PrimExpr user_init = Downcast<PrimExpr>((*init_it).second);
+        if (!user_init.dtype().is_void() && user_init.dtype() != op->dtype) {
+          user_init = tir::Cast(op->dtype, user_init);
+        }
+        init = user_init;
+      }
+      stream << ' ' << vid << " = " << PrintExpr(init) << ";\n";
     } else if (scope != "local.descriptor") {
       ICHECK(false) << "Unsupported scope: " << scope;
     }

src/transform/flatten_buffer.cc

@@ -25,13 +25,16 @@
 #include "tir/transforms/ir_utils.h"
 #include <tvm/arith/iter_affine_map.h>
 #include <tvm/ffi/reflection/registry.h>
+#include <tvm/ir/attrs.h>
 #include <tvm/tir/analysis.h>
 #include <tvm/tir/data_type_rewriter.h>
 #include <tvm/tir/stmt_functor.h>
 #include <tvm/tir/transform.h>
 
 #include <utility>
 
+#include "../op/builtin.h"
+
 namespace tvm {
 namespace tl {
 
@@ -46,6 +49,10 @@ class BufferFlattener : public arith::IRMutatorWithAnalyzer {
   static PrimFunc Flatten(PrimFunc func) {
     arith::Analyzer ana;
     auto pass = BufferFlattener(&ana);
+    if (auto init_map =
+            func->attrs.GetAttr<Map<Var, PrimExpr>>(tl::attr::kLocalVarInit)) {
+      pass.local_var_init_map_ = init_map.value();
+    }
     auto writer = func.CopyOnWrite();
     pass.MarkBufferMapShapes(func);
     writer->body = pass.VisitStmt(func->body);
@@ -198,6 +205,13 @@ class BufferFlattener : public arith::IRMutatorWithAnalyzer {
     if (!new_extents.same_as(alloc->extents)) {
       alloc.CopyOnWrite()->extents = new_extents;
     }
+    if (!local_var_init_map_.empty()) {
+      auto init_it = local_var_init_map_.find(alloc->buffer_var);
+      if (init_it != local_var_init_map_.end()) {
+        const PrimExpr &init = (*init_it).second;
+        alloc.CopyOnWrite()->annotations.Set(tl::attr::kLocalVarInit, init);
+      }
+    }
 
     return std::move(alloc);
   }
@@ -354,6 +368,9 @@ class BufferFlattener : public arith::IRMutatorWithAnalyzer {
 
   /*! \brief The updated external buffer map. */
   Map<Var, Buffer> updated_extern_buffer_map_;
+
+  /*! \brief Local var initializers preserved from block annotations. */
+  Map<Var, PrimExpr> local_var_init_map_;
 };
 
 PrimFunc FlattenBufferRewriter(PrimFunc f) {

src/transform/lower_opaque_block.cc

@@ -22,11 +22,14 @@
  */
 
 #include <tvm/ffi/reflection/registry.h>
+#include <tvm/ir/attrs.h>
 #include <tvm/tir/stmt_functor.h>
 #include <tvm/tir/transform.h>
 
+#include <string>
 #include <utility>
 
+#include "../op/builtin.h"
 #include "tir/transforms/ir_utils.h"
 
 namespace tvm {
@@ -39,10 +42,20 @@ using namespace tir::attr;
  */
 class OpaqueBlockLower : public StmtExprMutator {
 public:
-  static Stmt Rewrite(Stmt body) {
+  static PrimFunc Rewrite(PrimFunc f) {
+    auto fptr = f.CopyOnWrite();
     OpaqueBlockLower lower;
-    lower.storage_align_ = CollectStorageAlignAnnotation(body);
-    return lower(std::move(body));
+    if (auto existing =
+            fptr->attrs.GetAttr<Map<Var, PrimExpr>>(tl::attr::kLocalVarInit)) {
+      lower.local_var_init_map_ = existing.value();
+    }
+    lower.storage_align_ = CollectStorageAlignAnnotation(fptr->body);
+    fptr->body = lower(std::move(fptr->body));
+    if (!lower.local_var_init_map_.empty()) {
+      f = WithAttr(std::move(f), tl::attr::kLocalVarInit,
+                   lower.local_var_init_map_);
+    }
+    return f;
   }
 
 private:
@@ -59,7 +72,13 @@ class OpaqueBlockLower : public StmtExprMutator {
     if (!is_one(predicate)) {
       body = IfThenElse(predicate, std::move(body));
     }
-    // Step 3. Handle allocations in reverse order
+    // Step 3. Handle annotations, block annotations are not preserved by
+    // default.
+    std::vector<std::pair<std::string, PrimExpr>> pragma_attrs;
+    HandleAnnotations(new_block->annotations, &pragma_attrs, /*is_block=*/true,
+                      new_block->alloc_buffers);
+
+    // Step 4. Handle allocations in reverse order
     for (size_t i = new_block->alloc_buffers.size(); i > 0; --i) {
       const Buffer &buffer = new_block->alloc_buffers[i - 1];
       Array<PrimExpr> allocation_shape = GetBufferAllocationShape(buffer);
@@ -74,14 +93,15 @@ class OpaqueBlockLower : public StmtExprMutator {
         }
         allocate_annotations.Set(tir::attr::buffer_dim_align, allocate_aligns);
       }
-
+      auto init_it = local_var_init_map_.find(buffer->data);
+      if (init_it != local_var_init_map_.end()) {
+        const PrimExpr &init = (*init_it).second;
+        allocate_annotations.Set(tl::attr::kLocalVarInit, init);
+      }
       body = Allocate(buffer->data, buffer->dtype, allocation_shape,
                       const_true(), std::move(body), allocate_annotations);
     }
-    // Step 4. Handle annotations, block annotations are not preserved by
-    // default.
-    std::vector<std::pair<std::string, PrimExpr>> pragma_attrs;
-    HandleAnnotations(new_block->annotations, &pragma_attrs, /*is_block=*/true);
+    // Step 5. Insert attribute statements converted from pragmas
     for (auto it = pragma_attrs.rbegin(); it != pragma_attrs.rend(); ++it) {
       body = AttrStmt(Integer(0), it->first, it->second, std::move(body));
     }
@@ -188,13 +208,34 @@ class OpaqueBlockLower : public StmtExprMutator {
   Map<String, ffi::Any>
   HandleAnnotations(const Map<String, ffi::Any> &annotations,
                     std::vector<std::pair<std::string, PrimExpr>> *pragma_attrs,
-                    bool is_block) {
+                    bool is_block,
+                    const Array<Buffer> &alloc_buffers = Array<Buffer>()) {
     Map<String, ffi::Any> preserved_annotations;
     pragma_attrs->clear();
     for (const auto &kv : annotations) {
       const String &key = kv.first;
       if (tir::attr::IsPragmaKey(key)) {
         pragma_attrs->emplace_back(key, ConvertAttrValue(key, kv.second));
+      } else if (key == tl::attr::kLocalVarInit) {
+        if (auto local_init_map = kv.second.try_cast<Map<Var, PrimExpr>>()) {
+          for (const auto &pair : local_init_map.value()) {
+            local_var_init_map_.Set(pair.first, pair.second);
+          }
+        } else if (auto init_expr = kv.second.try_cast<PrimExpr>()) {
+          ICHECK(is_block) << "`" << tl::attr::kLocalVarInit
+                           << "` on non-block annotations is not supported";
+          Buffer target = ResolveLocalVarBuffer(alloc_buffers);
+          if (!target.defined()) {
+            LOG(WARNING) << "Failed to resolve buffer for `"
+                         << tl::attr::kLocalVarInit << "` annotation";
+            continue;
+          }
+          local_var_init_map_.Set(target->data, init_expr.value());
+        } else {
+          LOG(FATAL) << "Expected `" << tl::attr::kLocalVarInit
+                     << "` to be a PrimExpr or Map<Var, PrimExpr>, but got "
+                     << kv.second.GetTypeKey();
+        }
       } else if (!is_block) {
         // the loop annotation is preserved
         preserved_annotations.Set(key, kv.second);
@@ -206,6 +247,19 @@ class OpaqueBlockLower : public StmtExprMutator {
     return preserved_annotations;
   }
 
+  Buffer ResolveLocalVarBuffer(const Array<Buffer> &alloc_buffers) const {
+    for (const Buffer &buffer : alloc_buffers) {
+      std::string scope = buffer.scope();
+      if (scope.find("local.var") != std::string::npos) {
+        return buffer;
+      }
+    }
+    if (!alloc_buffers.empty()) {
+      return alloc_buffers.back();
+    }
+    return Buffer();
+  }
+
   /*! \brief Record the loop_var and loop start value of unit loops, whose
    * extent is one. */
   std::unordered_map<Var, PrimExpr> unit_loop_vars_;
@@ -215,12 +269,13 @@ class OpaqueBlockLower : public StmtExprMutator {
 
   /*! \brief The map from buffer var to its storage alignment information. */
   std::unordered_map<Var, StorageAlignAnnotation> storage_align_;
+
+  /*! \brief Local var initializers collected from block annotations. */
+  Map<Var, PrimExpr> local_var_init_map_;
 };
 
 PrimFunc TLLowerOpaqueBlock(PrimFunc f) {
-  auto fptr = f.CopyOnWrite();
-  fptr->body = OpaqueBlockLower::Rewrite(std::move(fptr->body));
-  return f;
+  return OpaqueBlockLower::Rewrite(std::move(f));
 }
 
 tir::transform::Pass LowerOpaqueBlock() {

src/transform/storage_rewrite.cc

@@ -25,6 +25,7 @@
 #include <tvm/arith/analyzer.h>
 #include <tvm/ffi/function.h>
 #include <tvm/ffi/reflection/registry.h>
+#include <tvm/ir/attrs.h>
 #include <tvm/ir/type.h>
 #include <tvm/target/target_info.h>
 #include <tvm/tir/analysis.h>
@@ -468,8 +469,10 @@ class StoragePlanRewriter : public StmtExprMutator {
   using AllocEntry = LinearAccessPatternFinder::AllocEntry;
 
   Stmt Rewrite(Stmt stmt, bool detect_inplace, bool enable_reuse,
-               bool reuse_require_exact_matched_dtype) {
+               bool reuse_require_exact_matched_dtype,
+               Map<Var, PrimExpr> local_var_init_map = {}) {
     detect_inplace_ = detect_inplace;
+    local_var_init_map_ = std::move(local_var_init_map);
     // plan the rewrite
     LinearAccessPatternFinder finder;
     finder(stmt);
@@ -694,6 +697,17 @@ class StoragePlanRewriter : public StmtExprMutator {
     }
     return body;
   }
+  Map<String, ffi::Any> MakeAllocateAnnotations(const Var &buffer_var) const {
+    Map<String, ffi::Any> annotations;
+    if (local_var_init_map_.defined()) {
+      auto it = local_var_init_map_.find(buffer_var);
+      if (it != local_var_init_map_.end()) {
+        const PrimExpr &init = (*it).second;
+        annotations.Set(tl::attr::kLocalVarInit, init);
+      }
+    }
+    return annotations;
+  }
   // Remap the index
   PrimExpr RemapIndex(DataType dtype, PrimExpr index, StorageEntry *e) {
     if (e->bits_offset == 0)
@@ -766,9 +780,11 @@ class StoragePlanRewriter : public StmtExprMutator {
 
         if (all_allocs_identical) {
           // simply use the original allocation.
-          e->alloc_nest.push_back(
-              Allocate(e->alloc_var, alloc_type, e->allocs[0]->extents,
-                       e->allocs[0]->condition, Evaluate(0)));
+          Map<String, ffi::Any> annotations =
+              MakeAllocateAnnotations(e->alloc_var);
+          e->alloc_nest.push_back(Allocate(
+              e->alloc_var, alloc_type, e->allocs[0]->extents,
+              e->allocs[0]->condition, Evaluate(0), std::move(annotations)));
           if (auto ptr = e->allocs[0]->body.as<DeclBufferNode>()) {
             e->alloc_nest.push_back(DeclBuffer(
                 RemapBuffer(ptr->buffer, e->alloc_var), Evaluate(0)));
@@ -824,9 +840,11 @@ class StoragePlanRewriter : public StmtExprMutator {
             combo_size = combo_size + make_const(DataType::Int(32), 1);
           }
           combo_size = analyzer_.Simplify(combo_size);
-          e->alloc_nest.push_back(Allocate(e->alloc_var, alloc_type,
-                                           {combo_size}, const_true(),
-                                           Evaluate(0)));
+          Map<String, ffi::Any> annotations =
+              MakeAllocateAnnotations(e->alloc_var);
+          e->alloc_nest.push_back(
+              Allocate(e->alloc_var, alloc_type, {combo_size}, const_true(),
+                       Evaluate(0), std::move(annotations)));
           if (IsSpecialTaggedMemory(e->scope)) {
             MemoryInfo info = GetMemoryInfo(e->scope.to_string());
             if (info.defined()) {
@@ -875,8 +893,10 @@ class StoragePlanRewriter : public StmtExprMutator {
     uint64_t type_bits = e->elem_type.bits() * e->elem_type.lanes();
     PrimExpr alloc_size = make_const(e->allocs[0]->extents[0].dtype(),
                                      (total_bits + type_bits - 1) / type_bits);
+    Map<String, ffi::Any> annotations = MakeAllocateAnnotations(e->alloc_var);
     e->alloc_nest.push_back(Allocate(e->alloc_var, e->elem_type, {alloc_size},
-                                     const_true(), Evaluate(0)));
+                                     const_true(), Evaluate(0),
+                                     std::move(annotations)));
     if (info.defined()) {
       ICHECK_LE(total_bits, info->max_num_bits)
           << "Allocation exceed bound of memory tag " << e->scope.to_string();
@@ -1178,6 +1198,8 @@ class StoragePlanRewriter : public StmtExprMutator {
   // Any buffers that is accessed at some point.  DeclBuffer instances
   // that do not appear in this list may be removed.
   std::unordered_set<const BufferNode *> all_buffers_accessed_;
+  // Initial values for local variable buffers.
+  Map<Var, PrimExpr> local_var_init_map_;
   // analyzer
   arith::Analyzer analyzer_;
 };
@@ -1795,7 +1817,7 @@ class VectorTypeRewriter : public StmtExprMutator {
     DLOG(INFO) << "Allocate with " << new_buffer_var << " and "
                << info.new_element_dtype << " extents: " << extents;
     return Allocate(new_buffer_var, info.new_element_dtype, extents,
-                    op->condition, op->body);
+                    op->condition, op->body, op->annotations);
   }
 
   Stmt VisitStmt_(const AllocateConstNode *op) final {
@@ -1941,10 +1963,16 @@ Pass StorageRewrite() {
       // Require exactly same-dtype matching in smem reuse for Vulkan and WebGPU
       reuse_require_exact_matched_dtype = true;
     }
+    Map<Var, PrimExpr> local_var_init_map;
+    if (auto init_map =
+            f->attrs.GetAttr<Map<Var, PrimExpr>>(tl::attr::kLocalVarInit)) {
+      local_var_init_map = init_map.value();
+    }
     auto *n = f.CopyOnWrite();
-    n->body = StoragePlanRewriter().Rewrite(std::move(n->body), detect_inplace,
-                                            enable_reuse,
-                                            reuse_require_exact_matched_dtype);
+    StoragePlanRewriter plan_rewriter;
+    n->body = plan_rewriter.Rewrite(
+        std::move(n->body), detect_inplace, enable_reuse,
+        reuse_require_exact_matched_dtype, std::move(local_var_init_map));
     // Parameters may not be rewritten, but internal allocations may.
     // Vectorization of AllocateConst is currently disabled, as it has
     // indexing issues for types that include padding (e.g. int8x3