diff --git a/src/cuda/transform/producer_consumer_ws.cc b/src/cuda/transform/producer_consumer_ws.cc index 32854aa05d..b54c08248d 100644 --- a/src/cuda/transform/producer_consumer_ws.cc +++ b/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(op)); tile_op.defined()) { if (const auto *copy = tile_op.as()) { - 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()) { - 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(); 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(); ICHECK(var); auto it = buffer_data_to_buffer_.find(GetRef(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,7 +2458,7 @@ 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()); @@ -2456,7 +2466,7 @@ class ProducerConsumerWSRewriter : public StmtExprMutator { 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); diff --git a/testing/python/issue/test_tilelang_issue_2443.py b/testing/python/issue/test_tilelang_issue_2443.py new file mode 100644 index 0000000000..03e67da429 --- /dev/null +++ b/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()