[Bugfix] Collect shared buffer on prelude

src/cuda/transform/producer_consumer_ws.cc

@@ -401,17 +401,26 @@ class LocalAccessCollector : public StmtExprVisitor {
 public:
   static LocalAccessSummary Collect(const Stmt &stmt,
                                     const BufferDataToBufferMap &buffer_map) {
-    LocalAccessCollector collector(buffer_map);
+    LocalAccessCollector collector(buffer_map, /*include_shared=*/false);
+    collector.VisitStmt(stmt);
+    return std::move(collector.summary_);
+  }
+
+  static LocalAccessSummary
+  CollectPrelude(const Stmt &stmt, const BufferDataToBufferMap &buffer_map) {
+    LocalAccessCollector collector(buffer_map, /*include_shared=*/true);
     collector.VisitStmt(stmt);
     return std::move(collector.summary_);
   }
 
 private:
-  explicit LocalAccessCollector(const BufferDataToBufferMap &buffer_map)
-      : buffer_data_to_buffer_(buffer_map) {}
+  explicit LocalAccessCollector(const BufferDataToBufferMap &buffer_map,
+                                bool include_shared)
+      : buffer_data_to_buffer_(buffer_map), include_shared_(include_shared) {}
 
-  static bool IsBranchPrivateBuffer(const Buffer &buffer) {
-    return IsFragmentBuffer(buffer) || IsLocalBuffer(buffer, true);
+  bool IsTrackedBuffer(const Buffer &buffer) const {
+    return IsFragmentBuffer(buffer) || IsLocalBuffer(buffer, true) ||
+           (include_shared_ && IsSharedBuffer(buffer));
   }
 
   void VisitStmt_(const BindNode *op) final {
@@ -429,14 +438,14 @@ class LocalAccessCollector : public StmtExprVisitor {
   }
 
   void VisitExpr_(const BufferLoadNode *op) final {
-    if (IsBranchPrivateBuffer(op->buffer)) {
+    if (IsTrackedBuffer(op->buffer)) {
       summary_.read_buffers.insert(op->buffer);
     }
     StmtExprVisitor::VisitExpr_(op);
   }
 
   void VisitStmt_(const BufferStoreNode *op) final {
-    if (IsBranchPrivateBuffer(op->buffer)) {
+    if (IsTrackedBuffer(op->buffer)) {
       summary_.write_buffers.insert(op->buffer);
     }
     StmtExprVisitor::VisitStmt_(op);
@@ -453,10 +462,10 @@ class LocalAccessCollector : public StmtExprVisitor {
   void VisitExpr_(const CallNode *op) final {
     if (auto tile_op = ParseOperator(GetRef<Call>(op)); tile_op.defined()) {
       if (const auto *copy = tile_op.as<CopyNode>()) {
-        if (IsBranchPrivateBuffer(copy->src)) {
+        if (IsTrackedBuffer(copy->src)) {
           summary_.read_buffers.insert(copy->src);
         }
-        if (IsBranchPrivateBuffer(copy->dst)) {
+        if (IsTrackedBuffer(copy->dst)) {
           summary_.write_buffers.insert(copy->dst);
         }
         for (const auto &range : copy->src_range) {
@@ -470,7 +479,7 @@ class LocalAccessCollector : public StmtExprVisitor {
         return;
       }
       if (const auto *fill = tile_op.as<FillNode>()) {
-        if (IsBranchPrivateBuffer(fill->dst)) {
+        if (IsTrackedBuffer(fill->dst)) {
           summary_.write_buffers.insert(fill->dst);
         }
         VisitExpr(fill->value);
@@ -486,7 +495,7 @@ class LocalAccessCollector : public StmtExprVisitor {
       ICHECK_EQ(op->args.size(), 3);
       const auto *base_load = op->args[0].as<BufferLoadNode>();
       ICHECK(base_load);
-      if (IsBranchPrivateBuffer(base_load->buffer)) {
+      if (IsTrackedBuffer(base_load->buffer)) {
         int rw_mask = GetConstAccessMask(op->args[2]);
         if (rw_mask & 1) {
           summary_.read_buffers.insert(base_load->buffer);
@@ -507,8 +516,7 @@ class LocalAccessCollector : public StmtExprVisitor {
       const auto *var = op->args[1].as<VarNode>();
       ICHECK(var);
       auto it = buffer_data_to_buffer_.find(GetRef<Var>(var));
-      if (it != buffer_data_to_buffer_.end() &&
-          IsBranchPrivateBuffer(it->second)) {
+      if (it != buffer_data_to_buffer_.end() && IsTrackedBuffer(it->second)) {
         int rw_mask = GetConstAccessMask(op->args[4]);
         if (rw_mask & 1) {
           summary_.read_buffers.insert(it->second);
@@ -533,6 +541,7 @@ class LocalAccessCollector : public StmtExprVisitor {
   }
 
   const BufferDataToBufferMap &buffer_data_to_buffer_;
+  bool include_shared_{false};
   LocalAccessSummary summary_;
   VarSet bound_vars_;
 };
@@ -549,7 +558,8 @@ ClassifyPreludeStmt(const Stmt &stmt, const BufferDataToBufferMap &buffer_map,
                     const LocalLiveSet &shared_live_seed,
                     const LocalLiveSet &producer_live_seed,
                     const LocalLiveSet &consumer_live_seed) {
-  LocalAccessSummary summary = LocalAccessCollector::Collect(stmt, buffer_map);
+  LocalAccessSummary summary =
+      LocalAccessCollector::CollectPrelude(stmt, buffer_map);
   if (!summary.HasTrackedDefs()) {
     return PreludeStmtPlacement::kKeepSharedPrelude;
   }
@@ -1915,9 +1925,9 @@ class ProducerConsumerWSRewriter : public StmtExprMutator {
     shared_prelude_live_seed_ = {};
     producer_prelude_live_seed_ = {};
     consumer_prelude_live_seed_ = {};
-    producer_prelude_live_seed_.AddUses(LocalAccessCollector::Collect(
+    producer_prelude_live_seed_.AddUses(LocalAccessCollector::CollectPrelude(
         rewritten_producer, buffer_data_to_buffer_));
-    consumer_prelude_live_seed_.AddUses(LocalAccessCollector::Collect(
+    consumer_prelude_live_seed_.AddUses(LocalAccessCollector::CollectPrelude(
         rewritten_consumer, buffer_data_to_buffer_));
 
     // Move pre-loop branch-private initialization next to the branch that
@@ -2448,15 +2458,15 @@ class ProducerConsumerWSRewriter : public StmtExprMutator {
         // the enclosing prelude so existing shared-prelude index math stays
         // common, but treat branch-private buffer uses as consumer-only so
         // their local/fragment initialization does not leak into producer.
-        LocalAccessSummary summary = LocalAccessCollector::Collect(
+        LocalAccessSummary summary = LocalAccessCollector::CollectPrelude(
             op->seq[i], rewriter_->buffer_data_to_buffer_);
         shared_live.vars.insert(summary.read_vars.begin(),
                                 summary.read_vars.end());
         consumer_live.buffers.insert(summary.read_buffers.begin(),
                                      summary.read_buffers.end());
       }
       for (int i = loop_idx - 1; i >= 0; --i) {
-        LocalAccessSummary summary = LocalAccessCollector::Collect(
+        LocalAccessSummary summary = LocalAccessCollector::CollectPrelude(
             op->seq[i], rewriter_->buffer_data_to_buffer_);
         if (!summary.HasTrackedDefs()) {
           shared_live.AddUses(summary);

testing/python/issue/test_tilelang_issue_2443.py

@@ -0,0 +1,65 @@
+import torch
+import tilelang
+import tilelang.testing
+import tilelang.language as T
+
+M, N, K = 1024, 1024, 1024
+block_M, block_N, block_K = 128, 128, 32
+
+
+def make_kernel(disable_ws: bool):
+    pass_configs = {}
+    if disable_ws:
+        pass_configs[tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED] = True
+
+    @tilelang.jit(out_idx=[-1], pass_configs=pass_configs)
+    def gemm_bias_relu(M, N, K, block_M, block_N, block_K, dtype="float16", accum_dtype="float32"):
+        @T.prim_func
+        def main(
+            A: T.Tensor((M, K), dtype),
+            B: T.Tensor((K, N), dtype),
+            bias: T.Tensor((N,), dtype),
+            C: T.Tensor((M, N), dtype),
+        ):
+            with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by):
+                A_shared = T.alloc_shared((block_M, block_K), dtype)
+                B_shared = T.alloc_shared((block_K, block_N), dtype)
+                bias_shared = T.alloc_shared((block_N,), dtype)
+                C_local = T.alloc_fragment((block_M, block_N), accum_dtype)
+
+                # PROLOGUE 1-D copy (correct slice index, correct extent).
+                T.copy(bias[bx * block_N : (bx + 1) * block_N], bias_shared)
+
+                T.clear(C_local)
+                for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=3):
+                    T.copy(A[by * block_M, k * block_K], A_shared)
+                    T.copy(B[k * block_K, bx * block_N], B_shared)
+                    T.gemm(A_shared, B_shared, C_local)
+
+                for i, j in T.Parallel(block_M, block_N):
+                    C_local[i, j] = T.max(C_local[i, j] + bias_shared[j], 0)
+
+                T.copy(C_local, C[by * block_M, bx * block_N])
+
+        return main
+
+    return gemm_bias_relu(M, N, K, block_M, block_N, block_K)
+
+
+@tilelang.testing.requires_cuda_compute_version(9, 0)
+def test_gemm_bias_relu_prologue_shared_copy_with_warp_specialization():
+    torch.manual_seed(0)
+    a = torch.randn(M, K, dtype=torch.float16, device="cuda")
+    b = torch.randn(K, N, dtype=torch.float16, device="cuda")
+    bias = torch.randn(N, dtype=torch.float16, device="cuda")
+    ref = torch.relu(a.float() @ b.float() + bias.float()).half()
+
+    for disable_ws in (False, True):
+        c = make_kernel(disable_ws)(a, b, bias)
+        torch.testing.assert_close(c, ref, rtol=1e-2, atol=1e-2)
+
+    torch.cuda.synchronize()
+
+
+if __name__ == "__main__":
+    tilelang.testing.main()