flagos-ai / flagos-ai/FlagTree
[BUG][ppu] tl.atomic_add(..., sem="acquire") causing cross-block spin-barrier deadlock on ppu
- Dominant language
- Python
- Stars
- 350
- Forks
- 149
- Avg merge
- 2d 4h
- Merged PRs (30d)
- 81
Description
# Cross-Block Spin-Barrier Kernel Hangs on PPU
## Environment:
Triton version: 3.5.0+ppu2.0.0.oe
Device: PPU-ZW810E
Driver version: 1.3.2-d7f5a2
## Reproduction
A cross-block synchronization (spin-barrier) kernel hangs on PPU; the process stalls at the device synchronization point and never returns. The same kernel runs correctly on NVIDIA (H20), Hygon, and Iluvatar platforms.
The triggering pattern is shown in the `_barrier` function in the script below. Run with: `python test.py`
```python
import torch
import triton
import triton.language as tl
@triton.jit
def _barrier(ctr, off, NC):
tl.atomic_add(ctr + off, 1, sem="release")
while tl.atomic_add(ctr + off, 0, sem="acquire") < NC:
pass
@triton.jit
def _kern(ctr, part, out, NC: tl.constexpr, ROUNDS: tl.constexpr):
pid = tl.program_id(0)
for r in range(ROUNDS):
tl.atomic_add(part + pid, 1.0) # atomic traffic, as in the real kernel
_barrier(ctr, r, NC)
tl.store(out + pid, tl.atomic_add(ctr + ROUNDS - 1, 0))
def main():
NC, ROUNDS = 64, 32
ctr = torch.zeros(ROUNDS, dtype=torch.int32, device="cuda")
part = torch.zeros(NC, dtype=torch.float32, device="cuda")
out = torch.zeros(NC, dtype=torch.float32, device="cuda")
_kern[(NC,)](
ctr, part, out, NC=NC, ROUNDS=ROUNDS
)
print("OK out[0]=", out[0].item(), flush=True)
if __name__ == "__main__":
main()
```
## Preliminary analysis
The Triton middle-end optimizes tl.atomic_add(ptr, 0, sem="acquire") into a ld.global.gpu.acquire load (the NVIDIA backend performs the same optimization, and the semantics are valid).
After removing `sem="acquire"` from `_barrier` (thereby reverting to the default `acq_rel` mode), it is no longer optimized into `ld.global.gpu.acquire`; the kernel then runs normally on the PPU without hanging.
Therefore, the issue likely lies in the PPU backend's lowering of `ld.global.gpu.acquire`: the spin-read may fail to see values updated by other blocks via atomic operations, causing the loop condition to remain true indefinitely and resulting in an infinite loop.
---
Contributor guide
Research direction
Start by running the provided reproduction with `python test.py` on the PPU, then inspect the middle-end and PPU backend lowering for `tl.atomic_add(ptr, 0, sem="acquire")` and `ld.global.gpu.acquire`. Compare this path with the default `acq_rel` behavior and the NVIDIA backend. Done means the cross-block kernel completes on PPU without hanging and preserves the expected synchronization behavior.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- compilers
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 48/100