feat(distributed): self-clearing barrier signals in host collective kernels

[georgebisbas]

Motivation

The host collective builtin kernels use single-shot credit barriers: each call issues TNOTIFY(+1) / TWAIT(Ge(1)) (or Set(1)) but never restores the INT32 signal cells, so after a call every peer's cell is left non-zero. Any reuse of the same signal buffer makes the next call's Ge(1) wait pass spuriously on stale credits — the barrier stops synchronizing and results are corrupted. The sim executor is sequentially consistent, so this is invisible in CI sim runs but real on non-coherent NPU silicon.

Change

Add a self-clearing epilogue to every host collective builtin kernel that mirrors the InCore composite credit-barrier protocol already shipped for pld.tensor.* collectives (#2175, lower_composite_ops_pass.cpp EmitEpilogueReset): after the final barrier of the call, each rank locally resets every peer's signal cell, so the signal is provably all-zero again before the kernel returns.

Kernel	Credits per peer cell per call	Epilogue
allreduce (mesh)	ready +1 + one per UB chunk	TNOTIFY(-(read_done_expected - 1), AtomicAdd)
allreduce_ring	+1 per used RoundBarrier row	TNOTIFY(-1, AtomicAdd) per row × peer
reduce_scatter	ready +1 + one per tile	TNOTIFY(0, Set)
broadcast	one per tile + final read-complete round	TNOTIFY(-expected, AtomicAdd)
barrier / allgather / all_to_all	+1 (single barrier)	TNOTIFY(-1, AtomicAdd)

All bodies use AtomicAdd in the main loop (single-writer pattern: each rank writes into its own cell on every peer). In the epilogue, most kernels restore zero with AtomicAdd(-N) on the local address. reduce_scatter uses Set(0) because on non-coherent NPU silicon a new AIV task dispatch may read a stale cached cell value — AtomicAdd would read-modify-write that stale value, while Set unconditionally writes zero. Set is safe here because the per-tile pipe_barrier(PIPE_ALL) guarantees all peers' credits have already landed; there are no in-flight writes to clobber.

Broadcast additionally gains a final read-complete round after its tile loop: each rank notifies peers only after its last tile read finishes, so a fast root cannot return and have its window reused while a slow peer still reads the last tile (mirrors the ring kernel's final RoundBarrier).

Tests

Signal-reuse STs — back-to-back calls through ONE shared signal, with distinct per-round payloads so a stale-credit pass fails loudly — for all seven collectives at P=2/P=4 (barrier and reduce_scatter at P=2):

• mesh/ring allreduce, broadcast, reduce_scatter, allgather, all_to_all, barrier
• Multicore allreduce signal reuse (P=2, C=4) — verifies reuse via output correctness rather than on-device signal readback, which is unreliable on non-coherent NPU across AIV task dispatches

Verification

• pre-commit: all hooks pass (ruff, pyright, headers, docs parity)
• sim ST: 24/24 pass (--forked --platform=a2a3sim --device=0,1,2,3), excluding 4 pre-existing Max/Min/FP16 sim SIGSEGVs (reproduced on unmodified kernels)
• NPU silicon ST (P=2, devices 0,1): all host collective signal reuse tests pass including broadcast, allreduce (mesh/ring/multicore), and reduce_scatter
• clang-tidy: N/A (only .cpp.in templates + Python tests changed)
• The on-silicon visibility (stale-credit reset) is what the NPU developer gate verifies — sim is sequentially consistent and cannot observe the bug this fixes.

Follow-up

The in-loop allreduce restriction lift (shared per-host_orch signal synthesis in SynthesizeAllReduceSignals) is tracked separately as #2310, per the review recommendation.

[coderabbitai]

Important

Review skipped

Auto incremental reviews are disabled on this repository.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 54f628a2-40b8-439c-aaf6-4bc63702dd9c

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

📝 Walkthrough

Walkthrough

Kernel templates for mesh allreduce and ring allreduce now include epilogue logic that clears used signal-barrier credits by atomically decrementing peer signal cells after the final barrier. New integration tests validate that a shared signal buffer works correctly across three consecutive allreduce calls for both mesh and ring variants.

Changes

Allreduce signal buffer reuse

Layer / File(s)	Summary
Mesh allreduce kernel epilogue cleanup python/pypto/runtime/builtins/collectives/allreduce/templates/kernel.cpp.in	Adds epilogue logic that computes issued notification counts and atomically subtracts them from each non-local peer signal cell, resetting signal state before the final pipeline barrier.
Ring allreduce kernel epilogue cleanup python/pypto/runtime/builtins/collectives/allreduce_ring/templates/kernel.cpp.in	Adds epilogue cleanup that issues local atomic decrement notifications for each barrier row and non-local peer cell after the final cross-rank barrier.
Mesh allreduce signal-reuse tests tests/st/distributed/test_l3_host_tensor_allreduce.py	Adds a host builder that executes three sequential FP32 sum allreduce rounds with one shared signal buffer, plus a parameterized test for two- and four-rank configurations validating each round.
Ring allreduce signal-reuse tests tests/st/distributed/test_l3_host_tensor_allreduce_ring.py	Adds a host builder that allocates one shared ring signal matrix for three publish/allreduce/consume sequences, plus a parametrized test that skips on insufficient devices, checks the variant directory, and validates all three rounds.

Estimated code review effort: 3 (Moderate) | ~25 minutes

Possibly related PRs

• hw-native-sys/pypto#2175: Implements and tests the same reusable barrier-signal protocol by clearing signal credits in allreduce kernels.
• hw-native-sys/pypto#2275: Addresses signal-buffer lifecycle management for HOST allreduce collectives, including self-clearing barrier signals.
• hw-native-sys/pypto#2160: Extends the same allreduce kernel templates and signal-barrier protocol with cleanup for signal reuse across calls.

Poem

A rabbit hops through signal cells,
Subtracting counts where credit dwells.
Three rounds run, the buffer clean,
Reused again, as if unseen.
🐰 Barriers cleared, the mesh and ring both gleam! ✨

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 31.58% 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 clearly summarizes the self-clearing barrier signal changes in the host allreduce kernels.
Description check	✅ Passed	The description explains the barrier-credit bug, implementation, tests, and verification related to the allreduce kernel changes.

---

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.}

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 987fe66206

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[coderabbitai]

Actionable comments posted: 3

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@tests/st/distributed/test_l3_host_tensor_allreduce_ring.py`:
- Line 237: Update the inputs construction in the distributed allreduce test to
add a distinct round-specific offset to each tensor generated by
_make_rank_inputs, ensuring every reuse round has a different expected reduction
while preserving the existing rank-specific values.

In `@tests/st/distributed/test_l3_host_tensor_allreduce.py`:
- Around line 272-279: Fix _build_host_allreduce_signal_reuse so its rounds
parameter has a valid contract: either remove the configurable rounds argument
and use the fixed three-stage setup, or generate host_orch’s stages and
corresponding output writes from rounds. Ensure all indices remain valid for
values below three and every allocated output slice is written for values above
three.
- Line 411: Update the inputs construction in the distributed allreduce test to
add a round-dependent offset to each round’s data before stacking, while
preserving distinct rank inputs within every round. Ensure later-round expected
results cannot match stale results from earlier rounds.

🪄 Autofix

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 9dba039a-16f1-43ea-9869-9fae17e15947

📥 Commits

Reviewing files that changed from the base of the PR and between a47d300 and 987fe66.

📒 Files selected for processing (4)

• python/pypto/runtime/builtins/collectives/allreduce/templates/kernel.cpp.in
• python/pypto/runtime/builtins/collectives/allreduce_ring/templates/kernel.cpp.in
• tests/st/distributed/test_l3_host_tensor_allreduce.py
• tests/st/distributed/test_l3_host_tensor_allreduce_ring.py

[YunjiQin]

Reviewed the full diff and cross-checked every kernel's credit accounting against the InCore protocol in lower_composite_ops_pass.cpp. The direction is right and the bug analysis (sim is sequentially consistent, so stale credits are only observable on silicon) is accurate.

The AtomicAdd(-N) epilogues in allreduce (mesh), allreduce_ring, barrier, allgather, all_to_all all check out — I verified the credit counts individually and they are exact and race-free:

Kernel	Credits per peer cell	Epilogue
allreduce (mesh)	1 + num_chunks	-(read_done_expected - 1)	correct
allreduce_ring	1 per row, 2*(NR-1)+1 rows	-1 per row	correct — round ends at exactly expected_rounds - 1, no row missed
barrier / allgather / all_to_all	1	-1	correct

Also confirmed: every early-return path (numel <= 0, invalid nranks, ring signal-shape guard) returns before any notify, so no negative cells are left behind; and a local TNOTIFY(signal_base + peer, ...) is the identity mapping of CommRemotePtr(ctx, ptr, my_rank), matching the InCore self-notify.

Three things need to change before this can land, plus one I'd suggest splitting out.

---

1. Set(0) epilogue in broadcast / reduce_scatter must become AtomicAdd(-credits) — blocking

The protocol this PR mirrors states the invariant explicitly (lower_composite_ops_pass.cpp:100-102):

For the same reason kSet must never be mixed with kAtomicAdd on the same cells — a set could clobber an already advanced counter.

broadcast and reduce_scatter notify with AtomicAdd(+1) in the body but reset with Set(0), which is exactly the forbidden combination. Concrete deadlock (reduce_scatter, P >= 2):

rank B: last wait passes -> Set(0) on its own cells -> returns -> next call starts
        -> AtomicAdd(+1) into rank A's cell[B]
rank A: last wait passes only now -> Set(0) wipes B's freshly written credit
rank A: next call waits cell[B] >= 1 -> hangs forever (B will not notify again)

Broadcast has the same shape: the root can return once it has collected every non-root's read-complete notify, and its next-call +1 into a non-root's cell[root] lands while that non-root is still between its own last notify and its Set(0).

The comment's premise — "an owner never writes again until the next call (host-serialized)" — is what fails: the next call is the owner's next write, and the window is precisely between "my last wait passes" and "I execute Set(0)". Per-rank streams are dispatched independently; there is no implicit cross-rank sync between kernel launches.

Worth noting the InCore rail already ships AtomicAdd(-N) for both of these collectives:

• LowerTensorBroadcastRule — EmitBarrier x1 + EmitEpilogueReset(total=1)
• LowerTensorReduceScatterRule — EmitBarrier x2 (ready + post-reduce) + EmitEpilogueReset(total=2)

The only NotifyOp::kSet in that whole file (:2133) is all_to_all_v publishing recv_counts data, not a barrier signal. So the two rails currently disagree on the protocol for the same op.

On the "AtomicAdd deadlocks on 910B" observation: an off-by-one in the credit count looks like a much more likely root cause than a hardware issue. Broadcast's expected ends at num_tiles + 1, so the -expected written in the PR description over-subtracts by one, leaving the cell at -1; the next call's Ge(1) then needs +2 to release — which is exactly the reported symptom. Correct counts:

// broadcast — every cell receives exactly num_tiles credits this call
const int32_t credits = expected - 1;
// reduce_scatter — ready(1) + one per tile, identical to the mesh allreduce
const int32_t credits = read_done_expected - 1;

Please re-test on silicon with these before concluding that Set is required. If it still deadlocks, that would implicate the other five kernels and the already-merged InCore path equally, and the root cause needs to be found before settling the protocol — rather than shipping two mutually exclusive schemes in one commit.

---

2. Remove the redundant notifies in broadcast

Two layers of redundancy in the restructured loop:

• Non-roots notify every peer, but only the root's cell is ever waited on. A non-root waits only on cell[kRoot], so of its P-1 notifies per tile, P-2 land in another non-root's cell[me] where nobody ever reads them — they just accumulate garbage credits (which is part of why Set(0) looked necessary).
• The root's per-tile read-complete wait is not needed within a call. The root writes tile t only in iteration t, and each iteration has a distinct base, so there is no intra-call WAR. That wait only exists to stop the root from returning and letting its window be reused by the next host statement — which needs to happen once, after the loop. That is the "final read-complete round" the PR description proposes; the code does it per tile, which puts every rank in lockstep on every tile and paces the whole broadcast at the slowest reader T times.

int32_t expected = 1;
for (int64_t base = 0; base < numel; base += kTileCount) {
  // ... shape / stride ...
  if (my_rank == kRoot) {
    TLOAD(stage_tile, target_g); /* ... */ TSTORE(target_g, stage_tile); /* ... */
    pipe_barrier(PIPE_ALL);
    for (int peer = 0; peer < nranks; ++peer) {          // root only: data-ready
      if (peer == my_rank) continue;
      pto::comm::Signal sig(CommRemotePtr(comm_ctx, signal_base + my_rank, peer));
      pto::comm::TNOTIFY(sig, static_cast<int32_t>(1), pto::comm::NotifyOp::AtomicAdd);
    }
  } else {
    pipe_barrier(PIPE_ALL);
    pto::comm::Signal sig(signal_base + kRoot);
    pto::comm::TWAIT(sig, expected, pto::comm::WaitCmp::GE);
    // TLOAD(root's target + base) -> TSTORE(own target + base)
  }
  ++expected;
  pipe_barrier(PIPE_ALL);
}
const int32_t tiles = expected - 1;

// One read-complete round after the loop, then a role-aware credit reset.
if (my_rank == kRoot) {
  for (int peer = 0; peer < nranks; ++peer) {
    if (peer == my_rank) continue;
    pto::comm::Signal sig(signal_base + peer);
    pto::comm::TWAIT(sig, static_cast<int32_t>(1), pto::comm::WaitCmp::GE);
  }
  for (int peer = 0; peer < nranks; ++peer) {            // each non-root sent exactly 1 credit
    if (peer == my_rank) continue;
    pto::comm::Signal self_sig(signal_base + peer);
    pto::comm::TNOTIFY(self_sig, static_cast<int32_t>(-1), pto::comm::NotifyOp::AtomicAdd);
  }
} else {
  pto::comm::Signal sig(CommRemotePtr(comm_ctx, signal_base + my_rank, kRoot));
  pto::comm::TNOTIFY(sig, static_cast<int32_t>(1), pto::comm::NotifyOp::AtomicAdd);
  pto::comm::Signal self_sig(signal_base + kRoot);       // root sent one credit per tile
  pto::comm::TNOTIFY(self_sig, -tiles, pto::comm::NotifyOp::AtomicAdd);
}
pipe_barrier(PIPE_ALL);

Remote atomics per tile drop from P*(P-1) to P-1, plus a one-off 2*(P-1) at the end. For P=8 broadcasting 1M elements (4096 tiles) that is roughly 229k -> 29k, about 8x fewer. Non-roots no longer write each other's cells at all, which also removes the last motivation for Set(0).

No new hazard: the root's tile t+1 self-copy and a non-root's tile t read touch different regions, and the ordering that matters (root stores before it notifies; a non-root waits for that notify before it loads) is unchanged.

---

3. Contradictory error message and docs

• synthesize_allreduce_signals_pass.cpp:190-191 still tells users "The signal protocol is single-use and cannot reuse a signal across dynamic invocations." — which directly contradicts the op descriptions and docs this PR updates. Users would see both "self-clearing, reusable" and "single-use" from the same release. It should state the real reason (the HOST-rail signal synthesis limitation), not a protocol property that no longer holds.
• docs/{en,zh}/user/distributed/04-debugging.md: the new remediation "or bind an explicit signal before the loop" does not work. CheckAllReduceCall(call) runs before the args_.size() == 2 branch in all three entry points (AssignStmt / EvalStmt / ReturnStmt), so an explicit signal is rejected just the same. Please drop that half-sentence.
• If two different epilogue schemes end up surviving (see #1), the docs' blanket claim that pld.tensor.* signals are reusable needs to be scoped to "between calls of the same collective" — a buffer shared across a Set-based and an AtomicAdd-based collective would corrupt in both directions.

---

4. Lifting the in-loop allreduce restriction — suggest a follow-up PR

Worth recording since this PR's premise removes the stated justification for it.

pl.range is not unrolled — UnrollLoops only expands ForKind::Unroll (its own error message says "use pl.range() instead"). So the for rd in pl.range(ROUNDS) in this PR's broadcast / reduce_scatter reuse tests is a live ForStmt that reaches codegen and passes in sim. That is direct evidence the host path (MaterializeCommDomainScopes -> LowerHostTensorCollectives -> codegen) already handles a host collective inside a loop; allreduce is blocked solely by the CheckAllReduceCall guard.

The two paths differ in cost though:

• args_.size() == 2 (explicit signal) — the pass synthesizes nothing here, it just visits the args and returns. With the signal now restarting at generation 1 every call and cross-rank skew bounded to one call, this is the same situation as back-to-back calls. Moving the repeating_scope_depth_ check after the args_.size() == 1 branch is enough, plus a for rd in pl.range(N): allreduce(data, signal) ST.
• args_.size() == 1 (synthesized signal) — MakeSignalBinding inserts pld.system.world_size + pld.tensor.alloc_window_buffer + pld.tensor.window immediately before the allreduce statement, so inside a loop it would allocate a window buffer per iteration. That is wrong regardless of the barrier protocol (repeated allocation under one name, and every rank must land on the same symmetric window). The alloc needs hoisting to the nearest non-loop scope — or straight to the host_orch entry, since one self-clearing signal serves every iteration.

Suggest doing this separately: it touches a different pass and needs its own tests, and folding it in would blur the silicon-validation scope of this PR — which right now should be focused on whether broadcast / reduce_scatter still deadlock once they use AtomicAdd(-credits).

[georgebisbas]

@YunjiQin thanks for the thorough review — all three blocking/actionable items are addressed on the pushed head fb09fe10:

1. Set(0) epilogue → AtomicAdd(-credits) (broadcast / reduce_scatter). Both kernels now reset with a local-address TNOTIFY(..., AtomicAdd) credit decrement and never mix kSet with the AtomicAdd(+1) body:

• reduce_scatter: -(read_done_expected - 1) — identical to the mesh allreduce count (1 ready barrier + one per UB chunk).
• broadcast: after the loop, the non-root resets its own cell[kRoot] by -tiles where tiles = expected - 1 — the correct count you derived (the PR description's -expected was the off-by-one).

2. Broadcast redundant notifies removed. Restructured to your suggested shape: the root is the only per-tile notifier (one +1 per non-root); non-roots wait solely on cell[kRoot] (no per-tile completion notify); the read-complete barrier is a single round after the loop — each non-root sends one +1 to root's cell[me], the root waits Ge(1) on each, then both sides do a role-aware AtomicAdd reset. Remote atomics per tile drop from P*(P-1) to P-1, plus the one-off 2*(P-1) round. This also removes the last motivation for Set(0).

3. Contradictory error message and docs fixed. CheckAllReduceCall now states the real reason — the HOST-rail signal-synthesis limitation (a synthesized signal binding cannot be allocated per dynamic iteration) — instead of claiming the protocol is single-use. The 04-debugging.md row (en/zh) drops the non-working "bind an explicit signal before the loop" remediation and the pass doc (en/zh) now explains the synthesis-based restriction accurately.

4. In-loop allreduce lift is not folded into this PR, as you recommended — tracked separately (issue + plan) so the silicon-validation scope stays on broadcast/reduce_scatter AtomicAdd(-credits).

Verified in sim: gate UT 74 + materialize UT 39 passed; all 7 host-collective ST files 34/34 passed (including broadcast + reduce_scatter signal-reuse at P=2/P=4). NPU verification pending on the developer gate.

[georgebisbas]

Rebased onto origin/main + multicore integration

Rebased the branch onto the latest origin/main, which now includes #2160 (multicore HOST AllReduce with core_num signal lanes). Two adjustments landed with it:

1. Allreduce epilogue now resets the block's own signal lane. The self-clearing epilogue previously reset signal_base + peer (the rank's first cell). Under #2160's layout, signal lanes are addressed peer * signal_stride + block_idx, so the old reset would have over-subtracted block-0 lanes and left the other lanes uncleared whenever core_num > 1. The epilogue now uses signal_base + peer * signal_stride + block_idx.

2. Adapted the #2160 multicore ST test to self-clearing semantics. The multicore test observed signal lanes with pld.system.wait(..., expected=1, cmp=Ge) after the allreduce and asserted >= 1. That worked pre-#2279 (single-use signals, no epilogue) but deadlocks with self-clearing signals, because the epilogue resets every lane back to 0 before the kernel returns. The test now waits for each lane to self-clear (expected=0, cmp=Eq) and asserts == 0 — a lane at zero proves that block started and completed its epilogue, and the allreduce completing still proves the ready barrier (all blocks started). This validates the new contract across the c2/c4/stride matrix, including the idle-lane case.

Sim verification (a2a3sim): UT 121 passed; ST 37/37 passed (base + signal-reuse for all seven collectives P=2/P=4, plus the three multicore cases). NPU (onboard) verification is still pending.

[georgebisbas]

Added a multicore signal-reuse ST (two back-to-back calls through ONE shared signal, distinct per-round payloads, core_num 2 and 4) — it proves the per-lane self-clearing epilogue makes the multicore signal reusable across calls. Sim: all 5 multicore cases pass (3 lane cases + 2 reuse cases).

[georgebisbas]

@YunjiQin — heads-up on two follow-up fixes pushed after the on-silicon NPU gate flagged failures:

Fix 1: Multicore allreduce signal reuse test (676a341)

test_multicore_allreduce_signal_reuse[reuse-p2-c4-wide-stride] was failing on non-coherent NPU. The test's consume_step tried to verify cleared signal values on-device via TWAIT(Eq, 0) + pl.load(signal), but a new AIV task dispatch can read stale cached values from the previous dispatch. The kernel correctly cleared the signals (confirmed by test_multicore_output_and_signal_lanes passing), so the fix removes the on-device signal readback and relies on output correctness as the proxy for signal clearing — matching the approach used by the single-core test_host_tensor_allreduce_signal_reuse.

Fix 2: Reduce scatter signal reuse (aa2172b)

test_host_tensor_reduce_scatter_signal_reuse[2] deadlocked with SCHEDULER_TIMEOUT. The kernel's self-clearing epilogue used AtomicAdd(-N) to reset signal cells, but AtomicAdd reads the current (possibly stale-cached) value before writing back. On non-coherent NPU, that stale read produces incorrect cell values, causing asymmetric synchronization between ranks → deadlock. The fix replaces TNOTIFY(self_sig, -barrier_count, NotifyOp::AtomicAdd) with TNOTIFY(self_sig, 0, NotifyOp::Set). Set writes zero unconditionally without reading the current value. This is safe because the per-tile pipe_barrier(PIPE_ALL) guarantees all peers' AtomicAdd credits have already landed — there are no in-flight writes that Set could clobber.

Both verified locally on NPU devices 0,1. The Set(0) approach matches what the barrier kernel already does and is safe under the same reasoning.