feat(collectives): ring allreduce TPUT push model on HOST builtin + InCore composite

[georgebisbas]

Summary

Replaces the pull-model engine of the ring allreduce on both rails with a TPUT push model (remote write), enabling O(1) NeighborBarrier on the HOST builtin and eliminating the pull-model NPU memory-ordering gap.

• HOST builtin (builtin.tensor.allreduce_ring): reduce-scatter + allgather converted from TLOAD/TSTORE pull to pto::comm::TPUT push — TPUT<AtomicAdd> remote-accumulate for RS, non-atomic TPUT for AG. Ordering is pipe_barrier(PIPE_ALL) around every transfer + dsb(DSB_DDR) before TNOTIFY (mirrors the in-tree allgather/all_to_all host builtins; not a GM fence). The O(P²) RoundBarrier is replaced by the O(1) NeighborBarrier (notify/wait the two ring neighbours only), which is NPU-safe because the TPUT write pipeline orders the data ahead of the signal — the pull model could not provide that.
• InCore composite (LowerTensorRingAllReduceRule): replaces pld.tile.remote_load pulls with pld.tile.put pushes (non-atomic TPUT + local reduce, preserving Sum/Max/Min/Prod). Race-free per-subchunk protocol: own-value read → ready barrier → push to right neighbour → push-done barrier → local read+reduce+store; barrier credits stay 2 per subchunk (signal shape unchanged). Ragged/arbitrary lengths and FP16 are preserved via balanced segments + valid shapes, with the shared VEC staging tile narrowed per transfer via tile.set_validshape.

Requires PTOAS >= v0.55 (pypto pins v0.57)

This PR depends on PTOAS v0.55 (release: hw-native-sys/PTOAS#1069, fixed in PR #1079).

The InCore composite's pld.tile.put transfers carry the exact ragged valid_cols as the partition-view extent. PTOAS ≤ v0.54 rejects dynamic partition-view shapes for pto.comm.tput ('pto.comm.tput' op expects dst to have a positive static shape), so the pure push model cannot compile below v0.55. The HOST builtin does not depend on this (its kernel is hand-written), but the composite rail does.

The requirement is satisfied by the current pin: pypto now pins PTOAS v0.57 (via #2291). The PR is rebased onto current main (2026-08-26, was 76 commits behind; re-rebased twice 2026-08-27 — first onto the #2530 runtime bump adopting the ChipTensor→TaskTensor kernel rename, then onto #2542 adopting the 42-lower_host_tensor_collectives → 43-… docs rename, with the PR's ring-doc edits re-homed) and merges cleanly. The UT tests pin the push structure (pld.tile.put + tile.create staging tile instead of pld.tile.remote_load).

Rebased — merge-order with #2279 (self-clearing signals) resolved

The rebase picked up #2279's self-clearing signal epilogue, which was written for the pull-model RoundBarrier (reset every peer's cell with TNOTIFY(-1) per round). That credit pattern does not match the push model's NeighborBarrier:

• NeighborBarrier credits only the two ring neighbours per round — a single cell when nranks == 2, where both neighbours are the same peer and the cell carries two +1s.
• The feat(distributed): self-clearing barrier signals in host collective kernels #2279 loop would corrupt the unused cells to −1 and, for nranks == 2, leave +1 stale credit in the one used cell — reintroducing the exact stale-credit barrier failure feat(distributed): self-clearing barrier signals in host collective kernels #2279 fixed.

The epilogue now branches on kUseNeighborBarrier: it restores only the two neighbour cells per used row with TNOTIFY(-1, AtomicAdd) (twice on the shared cell when nranks == 2), keeping the all-peer reset for the RoundBarrier fallback. The ring builtin is therefore self-clearing and signal-reuse-safe across back-to-back calls, matching the other host builtins (#2279). The ring signal-reuse ST (test_l3_host_tensor_allreduce_ring.py reuse leg) is the NPU gate for the adapted epilogue.

Issues this PR addresses

• Ring allreduce: unaligned-data handling (dcci flush tail + 32B transfer alignment) #2242 (ring unaligned-data handling): the pull-model dcci-flush tail gap (item 1) is moot — the push model needs no cacheline flush (the receiver reads data the sender wrote remotely via TPUT, never a locally-TSTORE'd line). The 32-byte transfer-alignment concern (item 2) is handled by narrowing the staging tile's column mask (ColMaskInternal / tile.set_validshape) to the exact (possibly ragged) transfer extent, so partial tails transfer exactly and never over-read/overwrite adjacent slots.
• PTOAS: pto.comm.tput rejects dynamic partition-view shapes — blocks variable-size transfer on the wire #2213 (PTOAS dynamic partition-view): closed as superseded by all_to_all_v transfers full MAX_RECV capacity on the wire instead of the runtime row count #2524; the >= v0.55 dependency it describes is satisfied by the v0.57 pin.

Verification (NPU silicon, 910B2, PTOAS v0.55)

All on real NPUs (8x 910B2), P=2 and P=4:

• tests/st/distributed/test_l3_host_tensor_allreduce_ring.py — HOST ring, P=2/4 ✅ (with NeighborBarrier enabled)
• tests/st/distributed/collectives/test_l3_tensor_allreduce_ring_intrinsic.py — InCore ring, P=2/4, sizes {1, 17, 4097, 65537} (ragged + >UB), Sum/Max/Min/Prod, FP16 ✅
• tests/st/distributed/collectives/test_l3_allreduce_ring.py + test_l3_ring_sizing_prewarm.py — no regression ✅
• UTs: test_lower_host_tensor_collectives.py, test_host_orch_distributed.py, test_lower_composite_ops.py (+ numerical) all green ✅

Total: 25/25 ST + 190 UT passed (pre-rebase). The 2026-08-26 rebase + epilogue adaptation re-ran the ring UTs (173/174, the one failure is a pre-existing parser TileView(pad=…) roundtrip gap on main, unrelated to this PR); the 2026-08-27 re-rebases (onto the #2530 runtime bump and the #2542 docs rename) each re-ran the same 173/174. NPU ST should be re-confirmed for the signal-reuse leg.

Trade-off note (ReduceOp)

The HOST builtin is ReduceOp::kSum only by construction, so its TPUT<AtomicAdd> RS is fine. The composite keeps non-atomic push + local reduce to preserve Sum/Max/Min/Prod; only a remote-atomic TPUT<AtomicAdd> variant would be Sum-only (AtomicType has no AtomicMax/Min).

Follow-ups (not in this PR)

• [Feature] Lift the HOST-rail in-loop pld.tensor.allreduce restriction via shared-signal synthesis #2310 — lift the HOST-rail in-loop pld.tensor.allreduce restriction via shared-signal synthesis (the other half of the feat(distributed): self-clearing barrier signals in host collective kernels #2279 review).
• TPUT_ASYNC (pto-isa) for an overlapped / IBing forward phase — optional perf follow-on (simpler ci: bump pto-isa commit to d779cd0 to fix Tpush/Tpop a5sim bug #1383 / plan 50).

Review notes

• Addresses CodeRabbit feedback: the allgather ready-barrier rationale is corrected (counters are per-round; the real guarantee is the previous round's push-done barrier every rank passes before round k), and the nranks == 2 NeighborBarrier behaviour is documented.

[coderabbitai]

📝 Walkthrough

Walkthrough

Changes

Ring allreduce TPUT conversion

Layer / File(s)	Summary
Lowering push schedule src/ir/transforms/lower_composite_ops_pass.cpp, docs/*/dev/passes/12-lower_composite_ops.md	Ring reduce-scatter and allgather now use pld.tile.put with a shared staging tile, dynamic transfer shapes, and ready/push-done barriers.
Kernel push execution and barriers python/pypto/runtime/builtins/collectives/allreduce_ring/templates/kernel.cpp.in, docs/*/dev/passes/40-lower_host_tensor_collectives.md	The kernel uses atomic TPUT for reduce-scatter, non-atomic TPUT for allgather, pipe ordering, dsb, and O(1) neighbor barriers.
Lowering validation updates tests/ut/ir/transforms/test_lower_composite_ops.py	Tests now validate TPUT calls, staging tiles, updated control flow, chunk shapes, FP16 handling, and round-trip output.

Estimated code review effort: 4 (Complex) | ~60 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Rank
  participant NeighborBarrier
  participant RightNeighbor
  Rank->>Rank: Read own receive slot
  Rank->>NeighborBarrier: Wait for ready generation
  Rank->>RightNeighbor: TPUT partial chunk
  Rank->>NeighborBarrier: Wait for push-done generation
  Rank->>Rank: Reduce and store partial
  Rank->>RightNeighbor: TPUT finalized chunk
  Rank->>NeighborBarrier: Wait for allgather completion

Loading

Possibly related PRs

• hw-native-sys/pypto#1967: Refactors collective lowering from remote reads to push-based movement.
• hw-native-sys/pypto#2093: Introduces the related allreduce_ring lowering and kernel changes.
• hw-native-sys/pypto#2161: Modifies ring allreduce TPUT communication, synchronization, and ragged-tail handling.

Poem

A rabbit hops where TPUT flows,
Pushes chunks to neighbors close.
Barriers guard each passing byte,
Reduced by day, gathered by night.
Staging tiles keep tails aligned—
Ring allreduce, neatly bind!

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
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 primary change: replacing ring allreduce with a TPUT push model across the HOST builtin and InCore composite.
Description check	✅ Passed	The description directly explains the TPUT push-model changes, PTOAS dependency, synchronization behavior, testing results, and affected rails.

---

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

[coderabbitai]

Actionable comments posted: 3

🧹 Nitpick comments (2)

src/ir/transforms/lower_composite_ops_pass.cpp (1)

1533-1537: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Correct the allgather ready-barrier rationale.

The comment states the ready barrier is safe because "the cumulative generation count includes the previous step's push-done notifies". Each round uses its own signal row (ag_round), so row k's counters start at zero and carry no credit from row k-1. The actual guarantee comes from the previous round's push-done barrier, which every rank must pass before it enters round k. Restate the rationale so a future change does not remove the previous round's push-done barrier on a false premise. The same wording appears at Lines 1571-1573.

🤖 Prompt for 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.

In `@src/ir/transforms/lower_composite_ops_pass.cpp` around lines 1533 - 1537,
Update the ready-barrier rationale comment around lines 1533-1537 to remove the
incorrect explanation about cumulative generation counts including previous
push-done notifies. Replace it with the correct rationale: the guarantee comes
from the previous round's push-done barrier that every rank must pass before
entering round k. Apply the same correction to the identical wording at lines
1571-1573 to ensure consistency and prevent future maintainers from removing the
previous round's push-done barrier based on the false premise.

python/pypto/runtime/builtins/collectives/allreduce_ring/templates/kernel.cpp.in (1)

121-144: 🩺 Stability & Availability | 🔵 Trivial | 💤 Low value

Add a comment documenting the two-rank NeighborBarrier behavior.

For nranks == 2, both neighbours map to rank 1, so each row cell receives two AtomicAdd writes and is observed twice via the same GE 1 waits. A short comment keeps the two-rank safety invariant explicit; the existing row allocation path already uses a fresh zeroed signal buffer per allreduce call.

🤖 Prompt for 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.

In
`@python/pypto/runtime/builtins/collectives/allreduce_ring/templates/kernel.cpp.in`
around lines 121 - 144, Add a concise comment above NeighborBarrier documenting
the nranks == 2 case: both neighbours resolve to rank 1, each signal cell
receives two AtomicAdd notifications, and the corresponding GE 1 waits observe
the shared cell twice. Keep the existing implementation unchanged and note that
the signal buffer is freshly zeroed per allreduce call.

🤖 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
`@python/pypto/runtime/builtins/collectives/allreduce_ring/templates/kernel.cpp.in`:
- Around line 146-149: Update the stale StepBarrier comment near the barrier
configuration to reflect that kUseNeighborBarrier is already enabled and
NeighborBarrier has been verified on 910B2 for P=2 and P=4. Remove the outdated
instructions to land with RoundBarrier first and defer the swap; keep the
comment focused on the current implementation and its verification status.

In `@src/ir/transforms/lower_composite_ops_pass.cpp`:
- Around line 1439-1454: The tile.set_validshape operation returns a new
narrowed TileType and does not mutate put_stage in-place. Update the
pld.tile.put call to pass the narrowed alias rs_stage_valid instead of
put_stage, so the push uses the correct narrowed staging tile extent. Apply the
same fix to the allgather push path by using ag_stage_valid instead of put_stage
in its corresponding pld.tile.put call.

In `@tests/ut/ir/transforms/test_lower_composite_ops.py`:
- Around line 2102-2109: Add a non-empty assertion for the stage_creates
collection in the FP16 test, matching the sibling test’s assertion and message.
Keep the existing chunk_shapes loop and validation unchanged so the test
requires FP16 staging-tile creation before checking alignment and size
constraints.

---

Nitpick comments:
In
`@python/pypto/runtime/builtins/collectives/allreduce_ring/templates/kernel.cpp.in`:
- Around line 121-144: Add a concise comment above NeighborBarrier documenting
the nranks == 2 case: both neighbours resolve to rank 1, each signal cell
receives two AtomicAdd notifications, and the corresponding GE 1 waits observe
the shared cell twice. Keep the existing implementation unchanged and note that
the signal buffer is freshly zeroed per allreduce call.

In `@src/ir/transforms/lower_composite_ops_pass.cpp`:
- Around line 1533-1537: Update the ready-barrier rationale comment around lines
1533-1537 to remove the incorrect explanation about cumulative generation counts
including previous push-done notifies. Replace it with the correct rationale:
the guarantee comes from the previous round's push-done barrier that every rank
must pass before entering round k. Apply the same correction to the identical
wording at lines 1571-1573 to ensure consistency and prevent future maintainers
from removing the previous round's push-done barrier based on the false premise.

🪄 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: 73efeda0-1ee1-418f-ae6b-d8ff15d8b379

📥 Commits

Reviewing files that changed from the base of the PR and between a47d300 and 6e35e57.

📒 Files selected for processing (7)

• docs/en/dev/passes/12-lower_composite_ops.md
• docs/en/dev/passes/40-lower_host_tensor_collectives.md
• docs/zh/dev/passes/12-lower_composite_ops.md
• docs/zh/dev/passes/40-lower_host_tensor_collectives.md
• python/pypto/runtime/builtins/collectives/allreduce_ring/templates/kernel.cpp.in
• src/ir/transforms/lower_composite_ops_pass.cpp
• tests/ut/ir/transforms/test_lower_composite_ops.py