flagos-ai / flagos-ai/FlagTree

[BUG][hygon] Scalar atomic_add(..., 0) poll in a while-loop causing cross-block spin-barrier deadlock on hcu

Open
#1,051 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
350
Forks
149
Avg merge
2d 19h
Merged PRs (30d)
97

Description

# Cross-block spin-barrier kernel deadlocks on HCU (any `sem`)

## Environment

- FlagTree: 0.6.2a1+hcu3.6 (git 367dc57, main branch)
- Device: Hygon DCU, DTK 2604

## Phenomenon

A cross-block synchronization (spin-barrier) kernel deadlocks on hcu when the
spin-read is a **scalar atomic whose result is used in the while loop**
(e.g. `tl.atomic_add(ptr, 0)` polling). The process stalls at device
synchronization and never returns. The hang is deterministic with grid >= 2
CTAs and an asymmetric arrival (one delayed CTA). The same kernel runs
correctly on NVIDIA (H800) — and on hcu, the *same barrier logic* encoded with
a vector atomic + `tl.max` completes with the expected values. On metax
(C550) the same script deadlocks with an identical LLIR structure; mthreads
(MTT S5000) and ppu (PPU-ZW810E, relaxed variant) run correctly.

The hang is independent of `sem`: `relaxed`, `acquire` and the default
(`None` -> `acq_rel`) all deadlock. It also cannot be reproduced with a single
block (within one CU the polling reads stay coherent).

## Reproduction

Run with `timeout 60` (the defect variant hangs on the current hcu backend):

```python
import sys
import torch
import triton
import triton.language as tl

N_CTAS = tl.constexpr(2) # grid size: 2 CTAs (minimum trigger)
BLOCK = tl.constexpr(128) # lanes per CTA
WORK_ITERS = tl.constexpr(2) # delayed-work rounds

@triton.jit
def poll_kernel(counter_ptr, out_ptr, scratch_ptr, VECTOR: tl.constexpr):
lane = tl.arange(0, BLOCK)
pid = tl.program_id(0)
zeros = tl.zeros([BLOCK], dtype=tl.uint32)

if pid == 0: # one delayed CTA -> asymmetric arrival
for t in tl.range(0, WORK_ITERS):
offs = t * BLOCK + lane
x = tl.load(scratch_ptr + offs)
tl.store(scratch_ptr + 4096 + lane, x)

tl.atomic_add(counter_ptr + zeros, 1, mask=lane == 0, sem="release", scope="gpu")

if VECTOR:
# reference: vector atomic read + tl.max uniform exit (safe encoding)
cnt = tl.max(
tl.atomic_add(counter_ptr + zeros, 0, sem=None, scope="gpu"), axis=0
)
while cnt < N_CTAS:
cnt = tl.max(
tl.atomic_add(counter_ptr + zeros, 0, sem=None, scope="gpu"),
axis=0,
)
else:
# defect: scalar atomic read, result used in the while loop
cnt = tl.atomic_add(counter_ptr, 0, sem=None, scope="gpu")
while cnt < N_CTAS:
cnt = tl.atomic_add(counter_ptr, 0, sem=None, scope="gpu")

tl.store(out_ptr + pid + zeros, 1, mask=lane == 0)

def run(vector):
counter = torch.zeros(1, dtype=torch.uint32, device="cuda")
out = torch.zeros(N_CTAS, dtype=torch.int32, device="cuda")
scratch = torch.randn(8192, dtype=torch.float32, device="cuda")
poll_kernel[(N_CTAS,)](counter, out, scratch, VECTOR=vector, num_warps=1)
torch.cuda.synchronize()
c = int(counter.item())
s = int(out.sum().item())
ok = (c == int(N_CTAS)) and (s == int(N_CTAS))
print(f"{'reference' if vector else 'defect'}: "
f"{'MATCH-EXPECTED' if ok else 'WRONG'} "
f"counter={c}/{int(N_CTAS)} out_sum={s}/{int(N_CTAS)}")

if __name__ == "__main__":
mode = sys.argv[1] if len(sys.argv) > 1 else "defect"
run(mode == "reference")
```

**Expected values** (analytically derived, compiler-independent):
`counter == 2` (each CTA's lane 0 adds exactly +1; the poll only reads, no
overshoot) and `out_sum == 2` (every CTA writes its completion flag after the
barrier passes). Observed behavior:

| variant | current hcu | NVIDIA H800 (same FlagTree build) |
|---|---|---|
| defect (scalar atomic poll) | HANGS (timeout) | `MATCH-EXPECTED counter=2/2 out_sum=2/2` |
| reference (vector atomic + tl.max) | `MATCH-EXPECTED` | `MATCH-EXPECTED` |

The reference variant always matches the expected values, which proves the
kernel logic is correct — the hang is a compiler lowering issue.

## Preliminary analysis

In `third_party/hcu/lib/TritonHCUGPUToLLVM/LoadStoreOpToLLVM.cpp`, the
`AtomicRMWOpConversion` lowers a **scalar** atomic whose result is used via an
LDS staging path:

```cpp
bool needLdsStaging = !tensorTy && !opResult.use_empty();
```

`AtomicRMWOpsEmitter::emitAtomicRMW` then builds a block-level divergent
branch (`CondBrOp(rmwMask, atomicBlock, endBlock)`): only thread 0 executes
the atomic and writes the result to LDS, and the workgroup barrier
(`b.barrier()`) sits at the join of that divergent branch.

In LLIR the poll loop looks like:

```llvm
br i1 %34, label %.loopexit, label %.preheader.preheader ; %34 = (tid == 0)
.loopexit:
%166 = atomicrmw or ptr addrspace(1) %0, i32 0 syncscope("agent") acq_rel
store i32 %166, ptr addrspace(3) @global_smem
br label %.preheader.preheader
.preheader:
fence syncscope("workgroup") release
tail call void @llvm.amdgcn.s.barrier()
fence syncscope("workgroup") acquire
%167 = load i32, ptr addrspace(3) @global_smem
br i1 %167 < 64, label %169, label %170
169:
br i1 %34, label %.loopexit, label %.preheader ; divergent loop back-edge
```

On GCN, `s.barrier` executed under a partial (divergent) exec mask is
undefined; combined with the cross-CU stale-read window (an early-polling CTA
does not observe the counter increments made by another CU), the loop never
exits and the kernel hangs.

The same source produces an equivalent LDS staging on NVIDIA, but encoded
differently: instruction-level
predication (inline-asm `@$3 atom.global...` / `@$3 st.shared`) with a single
uniform loop back-edge — `bar.sync` always executes with a full mask and the
loop exits uniformly.

| | NVIDIA | hcu |
|---|---|---|
| atomic + smem write | predicated instructions | block-level divergent branch |
| control flow | single uniform back-edge | `br i1 %34` divergent back-edge |
| barrier | always full-mask `bar.sync` | partial-mask `s.barrier` |

## Related issues

- #1029 — `[BUG][ppu] tl.atomic_add(..., sem="acquire")` spin-barrier deadlock
on ppu (same symptom family; different mechanism: acquire->load
optimization vs. the LDS staging structure above)
- #1030 — `[BUG][hygon] tl.load(..., volatile=True)` spin-barrier deadlock on
hygon (volatile flag dropped during lowering; related but distinct)

Contributor guide

Open the contributing guide

Research direction

Start in third_party/hcu/lib/TritonHCUGPUToLLVM/LoadStoreOpToLLVM.cpp, following AtomicRMWOpConversion and AtomicRMWOpsEmitter::emitAtomicRMW. Compare the generated LLIR for the scalar and vector reproductions, then run the provided two-CTA script on HCU. Done means the scalar polling variant completes with counter=2 and out_sum=2 without a divergent barrier deadlock.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp, python
Domain
backend, compilers
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Active
Clarity
Clearly specified
Newbie friendliness
55/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.