[BUG] Unexpected CTA-wide __syncthreads() inserted between split T.ws producer/consumer scopes
- Dominant language
- Python
- Stars
- 7.4k
- Forks
- 742
- Avg merge
- 1d 2h
- Merged PRs (30d)
- 108
Description
### Required prerequisites
- [x] I have read the documentation .
- [x] I have searched the [Issue Tracker](https://github.com/tile-ai/tilelang/issues) that this hasn't already been reported. (comment there if it has.)
### What version of TileLang are you using?
0.1.9
### System information
3.12.3 (main, Feb 4 2025, 14:48:35) [GCC 13.3.0] linux
0.1.9
2.8.0a0+5228986c39.nv25.05
PyTorch version: 2.8.0a0+5228986c39.nv25.05
CUDA used to build PyTorch: 12.9
OS: Ubuntu 24.04 x86_64
Python version: 3.12.3
Is CUDA available: True
GPU: NVIDIA H200
Nvidia driver version: 570.86.15
CUDA runtime version: 12.9.41
cuDNN version: 9.x
CPU: x86_64, Intel Xeon, 2 sockets, 192 logical CPUs
Relevant packages:
numpy==1.26.4
torch==2.8.0a0+5228986c39.nv25.5
torchvision==0.22.0a0
triton==3.1.0
pytorch-triton==3.3.0+
torch_tensorrt==2.8.0a0
### Problem description
`T.ws` producer/consumer scopes appear to be serialized by an unexpected CTA-wide `__syncthreads()`.
This can be reproduced directly with the existing TileLang example:
`python3 tilelang/examples/warp_specialize/example_warp_specialize_gemm_barrierpipe_stage2.py
`
The example contains a producer scope and a consumer scope inside the same loop:
```python
for ko in range(T.ceildiv(K, block_K)):
with T.ws(1):
T.mbarrier_wait_parity(...)
T.tma_copy(...)
T.tma_copy(...)
T.mbarrier_arrive(...)
with T.ws(0):
T.mbarrier_wait_parity(...)
T.gemm(...)
T.mbarrier_arrive(...)
```
However, the generated CUDA source contains CTA-wide synchronization between these scopes:
```cpp
for (int ko = 0; ko < ...; ++ko) {
__syncthreads();
if (128 <= ((int)threadIdx.x)) {
// producer: wait empty, issue TMA loads, arrive ready
}
__syncthreads();
if (((int)threadIdx.x) < 128) {
// consumer: wait ready, run WGMMA, arrive empty
}
}
```
This serializes the producer and consumer phases within each pipeline iteration. As a result, the generated code does not appear to implement an overlapped producer/consumer pipeline, even though the source uses separate `T.ws` scopes and explicit mbarrier synchronization.
### Reproducible example code
The Python snippets:
```python
python3 tilelang/examples/warp_specialize/example_warp_specialize_gemm_barrierpipe_stage2.py
```
### Traceback
```pytb
```
### Expected behavior
The producer and consumer warp-specialized roles should be able to run concurrently, synchronized only by the explicit mbarriers in the program.
In other words, the generated code should be closer to a persistent role split such as:
```cpp
if (producer_threads) {
for (...) {
// producer pipeline
}
} else {
for (...) {
// consumer pipeline
}
}
```
without CTA-wide `__syncthreads()` inserted between producer and consumer work in each iteration.
### Additional context
If the producer/consumer split is written manually using `threadIdx.x` / `T.get_thread_binding()` instead of separate `T.ws` scopes, the generated CUDA source does not contain the extra per-iteration `__syncthreads()`.
For example, writing the roles as a single thread-id branch:
```python
tx = T.get_thread_binding()
if tx >= 128:
for ko in range(T.ceildiv(K, block_K)):
T.mbarrier_wait_parity(...)
T.tma_copy(...)
T.tma_copy(...)
T.mbarrier_arrive(...)
else:
for ko in range(T.ceildiv(K, block_K)):
T.mbarrier_wait_parity(...)
T.gemm(...)
T.mbarrier_arrive(...)
```
generates a persistent producer/consumer role split without the extra CTA-wide synchronization between the roles.
Contributor guide
Assessment
This issue has not been assessed yet.