Ring allreduce: unaligned-data handling (dcci flush tail + 32B transfer alignment)

[georgebisbas]

Two alignment gaps in the ring allreduce kernel template that the mesh allreduce kernel already handles but the ring kernel currently does not.

1. dcci flush loop misses the chunk tail

File: python/pypto/runtime/builtins/collectives/allreduce_ring/templates/kernel.cpp.in

The dcci flush loop in both the Reduce-Scatter and AllGather phases steps by 16 elements (i += 16), which only covers the full chunk when chunk_elems % 16 == 0. Compile-time validation only guarantees numel % NR == 0, so a tail cache-line spanning the chunk end goes unflushed.

Counter-example: numel = 18, NR = 2 → chunk_elems = 9 (36 bytes). Chunk 1 starts at byte 36 and spans two cache lines; the 16-element stride misses the second one.

2. 32-byte transfer alignment is missing

The mesh allreduce kernel explicitly rounds ragged spans up to kTransferAlignmentBytes = 32 bytes (in collectives/allreduce/templates/kernel.cpp.in), aligning tile_elems to avoid partial-vector load/store. The ring kernel feeds tile_elems raw into ShapeDyn / StrideDyn with no such rounding.

3. kRecvTileDebugValue rename

The constant name at the top of the kernel reads like leftover debug scaffolding. The TASSIGN call is legitimate (it binds recv_tile to its buffer as a placeholder, overwritten by the subsequent TLOAD), but the constant should be renamed to something like kRecvTileInitFill with a one-line comment.

Suggested approach

• Extract a shared alignment helper used by both mesh and ring kernel templates.
• For the dcci flush: step by 16 elements but emit one additional flush for the final partial cache-line.
• Add an ST case where chunk_elems % 16 != 0 (e.g. numel=18, NR=2).
• Rename kRecvTileDebugValue → kRecvTileInitFill.

Context

Raised in review of PR #2094 by @YunjiQin.

[georgebisbas]

Status check against main (08d3b231):

• Item 3 (kRecvTileDebugValue → kRecvTileInitFill rename) is already done — it landed with feat(distributed): host-orchestrated ring allreduce with builtin.tensor.allreduce_ring #2094 itself, so no separate action needed there.
• Items 1 (dcci flush tail) and 2 (32 B transfer alignment) are still open — confirmed unaddressed in allreduce_ring/templates/kernel.cpp.in. The file still carries an explicit comment acknowledging both gaps.

One thing worth flagging for whoever picks this up: two other open PRs touch this same kernel template and would change the shape of the fix depending on merge order:

• feat(collectives): ring allreduce TPUT push model on HOST builtin + InCore composite #2280 (ring allreduce TPUT push model) replaces the pull-model dcci flush entirely. If it merges first, the dcci-tail-flush half of this issue becomes moot (the loop it patches is gone), while the 32 B alignment fix would need to be re-targeted at the new TPUT path instead.
• feat(distributed): self-clearing barrier signals in host collective kernels #2279 (self-clearing barrier signals in host collective kernels) also edits the same file for an unrelated concern — worth coordinating so both don't land independently and produce a messy rebase.

Suggest deciding the merge order deliberately (dcci fix vs. #2280) rather than letting it happen by accident.

[georgebisbas]

Status update (2026-08-31) — both remaining items are now resolved on main by the ring TPUT-push rewrite (#2280, merged 57e5c315):

• Item 1 (dcci flush loop misses the chunk tail) — moot. The pull-model dcci flush loop no longer exists in allreduce_ring/templates/kernel.cpp.in; ring reduce-scatter and allgather were converted to TPUT push (remote write), which needs no cacheline flush — the receiver reads data the sender wrote remotely, never a locally-TSTOREd line.
• Item 2 (32-byte transfer alignment) — handled. The TPUT path narrows the staging tile to the exact (possibly ragged) transfer extent via ColMaskInternal / tile.set_validshape, so partial tails transfer exactly and never over-read/overwrite adjacent slots. The InCore composite ring ST now exercises ragged sizes {1, 17, 4097, 65537} (Sum/Max/Min/Prod, FP16) at P=2/P=4.
• Item 3 (kRecvTileDebugValue rename) — already done in feat(distributed): host-orchestrated ring allreduce with builtin.tensor.allreduce_ring #2094, as noted in the earlier status check.

Only leftover: the HOST builtin ring rail still requires numel % NR == 0 — enforced by CheckRingChunkConstraints in lower_host_tensor_collectives_pass.cpp, and the HOST ring ST (test_l3_host_tensor_allreduce_ring.py) only exercises uniform sizes. Ragged / arbitrary-length inputs on the HOST ring builtin remain unsupported (the InCore composite rail already handles them). That is the remaining scope of this issue.