[Bugfix] Collect shared buffer on prelude

[ppppqp]

Fixes #2443

Summary

Fixes warp-specialized kernels with a pre-loop shared-memory copy whose destination is consumed by consumer-only post-loop code.
The WS prelude classifier previously tracked local/fragment buffers but ignored shared buffers. As a result, a T.copy on a shared buffer can remain in the common prelude. Due to warp specialization, consumer branch uses a remapped thread index, so it might be lowered to wrong thread extent and produce incorrect result.

For the example listed in the issue, the IR after the tl.PipelinePlanning is:

# ...
            T.copy(T.region(bias[bx * 128], 1, 128), T.region(bias_shared[0], 2, 128))
            T.attr([128, 128], "kWarpSpecializationScope", 0)
            if tx < 128:
               # ...
            else:
               # ...
                for i in T.parallel(128):
                    for j in T.parallel(128):
                        C_local[i, j] = T.max(C_local[i, j] + T.Cast("float32", bias_shared[j]), T.float32(0.0))
                T.copy(T.region(C_local[0, 0], 1, 128, 128), T.region(C[by * 128, bx * 128], 2, 128, 128))

which confirms that T.copy for shared buffer was left out in the prelude, and so the consumer thread gets wrong data

Changes

The fix introduces LocalAccessCollector::CollectPrelude(...) that is specific to prelude collection that include shared buffers. It is specialized to classify prelude statement. Existing LocalAccessCollector::Collect behavior is unchanged. Also renamed isBranchPrivateBuffer to isTrackedBuffer for more accuracy since shared buffer is not necessarily branch private.

After the fix, verified that the numeric result and IR are both correct.

if tx < 128:
               # ...
            else:
                # the copy is now correctly placed in consumer group
                T.copy(T.region(bias[bx * 128], 1, 128), T.region(bias_shared[0], 2, 128))
                # ...
                for i in T.parallel(128):
                    for j in T.parallel(128):
                        C_local[i, j] = T.max(C_local[i, j] + T.Cast("float32", bias_shared[j]), T.float32(0.0))
                T.copy(T.region(C_local[0, 0], 1, 128, 128), T.region(C[by * 128, bx * 128], 2, 128, 128))

Testing

Added a test the mimic the example in the original issue

[github-actions]

👋 Hi! Thank you for contributing to the TileLang project.

Please remember to run pre-commit run --all-files in the root directory of the project to ensure your changes are properly linted and formatted. This will help ensure your contribution passes the format check.

We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀

[coderabbitai]

📝 Walkthrough

Walkthrough

LocalAccessCollector in the warp-specialized producer/consumer rewrite pass gains an include_shared_ flag and a new CollectPrelude entry point that sets it to true. All prelude classification and liveness seed/propagation sites switch from Collect to CollectPrelude, making shared buffer accesses visible during prelude analysis. A new test validates a GEMM+bias+ReLU kernel with a shared-memory bias prologue copy under both warp-specialization modes.

Changes

WS Prelude Shared-Buffer Liveness Fix

Layer / File(s)	Summary
LocalAccessCollector: include_shared flag and CollectPrelude entry point src/cuda/transform/producer_consumer_ws.cc	Adds include_shared_ boolean member; updates IsTrackedBuffer to conditionally include shared buffers; introduces CollectPrelude static factory (include_shared=true) alongside the existing Collect (include_shared=false); updates visitor logic for BufferLoad/BufferStore and tile-op access tracking to use IsTrackedBuffer.
Switch prelude classification and liveness seeds to CollectPrelude src/cuda/transform/producer_consumer_ws.cc	ClassifyPreludeStmt, producer_prelude_live_seed_/consumer_prelude_live_seed_ population, and PropagatePreludeLiveness post-loop sibling collection all switch from Collect to CollectPrelude.
GEMM+bias+ReLU prologue shared-copy WS test testing/python/issue/test_tilelang_issue_2443.py	New test defines a TileLang kernel with shared-memory bias prologue copy, pipelined K-tile GEMM, per-element bias+ReLU, and output copy; runs with and without warp specialization and verifies correctness against PyTorch reference.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Possibly related issues

• [BUG]Warp specialization produces incorrect results for prologue 1-D global→shared copy before pipelined T.gemm #2443: The CollectPrelude changes directly fix the incorrect prologue global→shared copy behavior under warp specialization that this issue reports.

Suggested reviewers

• LeiWang1999
• chengyupku

🐇 A shared buffer hopped into the prelude,
the collector missed it — how rude!
Now include_shared_ is true,
and liveness flows through,
with bias and ReLU correctly queued! 🎉

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title accurately summarizes the main change: collecting shared buffers during prelude analysis.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}