diff --git a/docs/en/dev/passes/12-lower_composite_ops.md b/docs/en/dev/passes/12-lower_composite_ops.md index 769fad6875..3d7ddeaa02 100644 --- a/docs/en/dev/passes/12-lower_composite_ops.md +++ b/docs/en/dev/passes/12-lower_composite_ops.md @@ -188,7 +188,7 @@ Running `LowerCompositeOps` twice produces identical IR after the first run: the ## `pld.tensor.*` distributed collectives -The pass also lowers the `pld.tensor.*` family of window-bound distributed collectives. Each collective is a single composite `Call` that expands into a notify / wait + data-movement recipe, plus a self-clearing epilogue. The data-movement primitive differs by op: `allgather` uses `pld.tile.put` (TPUT-based, auto-chunks through a VEC staging tile), `broadcast` relocates window data with `pld.tile.get` (GM→GM copy), while `allreduce` and `reduce_scatter` pull peer chunks into a UB tile with `pld.tile.remote_load`. Allreduce selects `tile.add`, `tile.maximum`, `tile.minimum`, or `tile.mul`; reduce-scatter currently accumulates with `tile.add`. All seven rules share the same **self-clearing credit-barrier protocol** (`LoweringBuilder::EmitBarrier` + `EmitEpilogueReset`) — see [Barrier-signal protocol](#barrier-signal-protocol) below — so a `signal` buffer is reusable across back-to-back calls, and even inside `for` / `while` / `if`. +The pass also lowers the `pld.tensor.*` family of window-bound distributed collectives. Each collective is a single composite `Call` that expands into a notify / wait + data-movement recipe, plus a self-clearing epilogue. The data-movement primitive differs by op: `allgather` and ring `allreduce` use `pld.tile.put` (TPUT-based push, auto-chunking through a VEC staging tile), `broadcast` relocates window data with `pld.tile.get` (GM→GM copy), while mesh `allreduce` and `reduce_scatter` pull peer chunks into a UB tile with `pld.tile.remote_load`. Allreduce selects `tile.add`, `tile.maximum`, `tile.minimum`, or `tile.mul`; reduce-scatter currently accumulates with `tile.add`. All seven rules share the same **self-clearing credit-barrier protocol** (`LoweringBuilder::EmitBarrier` + `EmitEpilogueReset`) — see [Barrier-signal protocol](#barrier-signal-protocol) below — so a `signal` buffer is reusable across back-to-back calls, and even inside `for` / `while` / `if`. ### Barrier-signal protocol @@ -206,7 +206,7 @@ The allreduce rule starts with a cross-rank ready barrier on shared `signal` cel For a fully-valid packed target, mesh lowering creates a logical `[1, product(all dimensions)]` view and traverses it with physical tiles of at most 16 KiB. A statically known extent smaller than the budget shrinks the chunk to the smallest 32-byte-aligned physical width that covers it, so small allreduces do not reserve a full 16-KiB tile while remaining legal PTO tiles. The tail carries `valid_shape=[1, min(chunk, remaining)]` through both `tile.load` and `pld.tile.remote_load`, so the allocation stays static while the read/store extent is exact. If an ND target carries a partial `TensorView.valid_shape`, the pass preserves and reduces the representable `[rows, cols]` rectangle through the established single-rectangle path. Constant valid rectangles use their compact shape; symbolic valid extents fall back to the source's physical rectangle when that statically bounded rectangle fits within one 16-KiB chunk. Oversized partial rectangles, strided targets, DN partial views, and partial boxes that cannot be represented by the leading-dimension collapse are rejected explicitly. -Ring lowering uses one packed 2D view for its reduce-scatter and allgather phases. A fully valid target becomes `[1, SIZE]`; a contiguous partial prefix keeps physical shape `[1, product(target.shape)]` and carries logical `TensorView.valid_shape=[1, product(target.valid_shape)]`. FP32 retains balanced `floor(i * SIZE / NR)` segment boundaries. FP16 rounds each interior boundary up to 16 elements and caps it at `SIZE`; consequently every non-empty segment and every UB subchunk starts at a 32-byte-aligned address. A ragged FP16 remote load may read the aligned physical tail reserved by the communication domain, then `tile.set_validshape` restores the logical extent before reduction and store. This supports non-divisible inputs and `SIZE < NR` without inserting holes into the public tensor layout. Every subchunk of every round barriers on a call-local ready + read-complete generation pair; the epilogue then subtracts `2 * chunk_count` (uniform across rounds, since every round's subchunk loop shares the same bound) from every row of the `[2*(NR-1), NR]` signal. +Ring lowering uses one packed 2D view for its reduce-scatter and allgather phases, and moves data with **TPUT pushes** (`pld.tile.put`, non-atomic) instead of remote loads. A fully valid target becomes `[1, SIZE]`; a contiguous partial prefix keeps physical shape `[1, product(target.shape)]` and carries logical `TensorView.valid_shape=[1, product(target.valid_shape)]`. FP32 retains balanced `floor(i * SIZE / NR)` segment boundaries. FP16 rounds each interior boundary up to 16 elements and caps it at `SIZE`; consequently every non-empty segment and every UB subchunk starts at a 32-byte-aligned address. Per subchunk, each rank first reads its OWN value of the receive slot into a register tile (the slot is stable — own value only — until the left neighbour's push lands), barriers on the ready generation (2k+1), then TPUTs its send subchunk into the RIGHT neighbour's slot of the same index; a push-done barrier (2k+2) precedes the local read + reduce + store of the receive slot. The single shared VEC staging tile is narrowed with `tile.set_validshape` to each transfer's exact `valid_cols`, and the push transfer carries that dynamic extent — PTOAS >= v0.55 accepts dynamic partition-view shapes for `tput` (hw-native-sys/PTOAS#1069), preserving ragged and FP16 tails without padding the window. The non-atomic push + local reduce keeps every `ReduceOp` (Sum/Max/Min/Prod) working — only a remote-atomic `TPUT` would be Sum-only. This supports non-divisible inputs and `SIZE < NR` without inserting holes into the public tensor layout. Every subchunk of every round barriers on a call-local ready + push-done generation pair; the epilogue then subtracts `2 * chunk_count` (uniform across rounds, since every round's subchunk loop shares the same bound) from every row of the `[2*(NR-1), NR]` signal. Any symbolic target or partial-valid extent that survives lowering must be runtime-bound by a kernel scalar, loop variable, or physical tensor-shape parameter; a type-metadata-only symbol is rejected during PTO codegen. A fully dynamic physical target dimension is bound from that tensor parameter. diff --git a/docs/en/dev/passes/42-lower_host_tensor_collectives.md b/docs/en/dev/passes/42-lower_host_tensor_collectives.md index ae80bce84b..dd79f7f3e1 100644 --- a/docs/en/dev/passes/42-lower_host_tensor_collectives.md +++ b/docs/en/dev/passes/42-lower_host_tensor_collectives.md @@ -119,6 +119,17 @@ Ring allreduce currently supports only `ReduceOp.Sum` with `dtype=FP32`. with `mode="ring"`. Ring allreduce also supports at most 16 participating devices (`world_size <= 16`). +The `builtin.tensor.allreduce_ring` kernel is **push-based**: data movement uses +`pto::comm::TPUT` (remote write) — the reduce-scatter phase accumulates into the +right neighbour's slot via `TPUT`, and the allgather phase forwards +each finalized chunk with a non-atomic `TPUT`, mirroring the in-tree `allgather` +/ `all_to_all` host builtins. Ordering is `pipe_barrier(PIPE_ALL)` around each +transfer plus `dsb(DSB_DDR)` before every `TNOTIFY` (not +`pto.fence.barrier_all`, which does not drain the MTE DMA pipe). Cross-rank +synchronisation uses the O(1) `NeighborBarrier` (notify/wait the two ring +neighbours only) — safe on NPU because the TPUT write pipeline orders the data +ahead of the signal, which the old pull model (TLOAD/TSTORE) did not. + All window operands of a HOST collective — data and signal alike — must resolve to pairwise distinct `WindowBuffer` allocations. Two `pld.window()` views over the same `alloc_window_buffer` are a cross-process data race under diff --git a/docs/zh/dev/passes/12-lower_composite_ops.md b/docs/zh/dev/passes/12-lower_composite_ops.md index 022ba2ca22..6ea739da2b 100644 --- a/docs/zh/dev/passes/12-lower_composite_ops.md +++ b/docs/zh/dev/passes/12-lower_composite_ops.md @@ -188,7 +188,7 @@ sin 与 cos 共用同一组多项式系数：cos 路径只在区间归约阶段 ## `pld.tensor.*` 分布式集合通信算子 -本 Pass 同时降级 `pld.tensor.*` 系列的窗口绑定 (window-bound) 分布式集合通信算子。每个集合通信算子都是一个组合 `Call`，展开为 notify / wait + 数据搬运序列，外加自清理尾声。数据搬运原语因算子而异：`allgather` 使用 `pld.tile.put`（基于 TPUT 的推送，经 VEC staging tile 自动分块），`broadcast` 用 `pld.tile.get` 搬运窗口数据（GM→GM 拷贝），`allreduce` 与 `reduce_scatter` 用 `pld.tile.remote_load` 把 peer chunk 拉进 UB tile。allreduce 根据规约类型选择 `tile.add`、`tile.maximum`、`tile.minimum` 或 `tile.mul`；reduce-scatter 当前仍用 `tile.add`。七条规则共享同一套**自清理信用屏障协议**（`LoweringBuilder::EmitBarrier` + `EmitEpilogueReset`）—— 参见下方[屏障-信号协议](#屏障-信号协议) —— 因此 `signal` buffer 可以在连续调用之间复用，甚至在 `for` / `while` / `if` 内部也可以。 +本 Pass 同时降级 `pld.tensor.*` 系列的窗口绑定 (window-bound) 分布式集合通信算子。每个集合通信算子都是一个组合 `Call`，展开为 notify / wait + 数据搬运序列，外加自清理尾声。数据搬运原语因算子而异：`allgather` 与 ring `allreduce` 使用 `pld.tile.put`（基于 TPUT 的推送，经 VEC staging tile 自动分块），`broadcast` 用 `pld.tile.get` 搬运窗口数据（GM→GM 拷贝），mesh `allreduce` 与 `reduce_scatter` 用 `pld.tile.remote_load` 把 peer chunk 拉进 UB tile。allreduce 根据规约类型选择 `tile.add`、`tile.maximum`、`tile.minimum` 或 `tile.mul`；reduce-scatter 当前仍用 `tile.add`。七条规则共享同一套**自清理信用屏障协议**（`LoweringBuilder::EmitBarrier` + `EmitEpilogueReset`）—— 参见下方[屏障-信号协议](#屏障-信号协议) —— 因此 `signal` buffer 可以在连续调用之间复用，甚至在 `for` / `while` / `if` 内部也可以。 ### 屏障-信号协议 @@ -227,16 +227,24 @@ allreduce 仍预留完整 16-KiB tile，又满足 PTO tile 的对齐要求。尾 过大的 partial 矩形、strided 目标、DN partial view 和无法按 leading-dimension collapse 表示的 partial 区域会被明确拒绝。 -ring 降级在 reduce-scatter 和 allgather 阶段使用同一个 packed 2D 视图。 +ring 降级在 reduce-scatter 和 allgather 阶段使用同一个 packed 2D 视图，并用 +**TPUT 推送**（`pld.tile.put`，非原子）代替远程加载来搬运数据。 完全有效的目标会变为 `[1, SIZE]`；连续 partial prefix 保留物理 shape `[1, product(target.shape)]`，并携带逻辑 `TensorView.valid_shape=[1, product(target.valid_shape)]`。FP32 保留均衡的 `floor(i * SIZE / NR)` segment 边界；FP16 把每个内部 边界向上对齐到 16 个元素并限制在 `SIZE` 内，因此每个非空 segment 和 UB -subchunk 都从 32 字节对齐地址开始。FP16 的 ragged remote load 可以读取通信域 -预留的对齐物理尾部，然后通过 `tile.set_validshape` 在归约和写回前恢复逻辑范围。 +subchunk 都从 32 字节对齐地址开始。每个 subchunk 中，各 rank 先把接收 slot 的 +**自身值**读入寄存器 tile（该 slot 在左邻居的推送落地前保持稳定 —— 只有自身值）， +随后在 ready generation (2k+1) 上做屏障，再把发送 subchunk 通过 TPUT 推送到 +**右邻居**的同一下标 slot；push-done 屏障 (2k+2) 之后才本地读取、归约并写回接收 slot。 +共享的 VEC staging tile 通过 `tile.set_validshape` 收窄到每次传输的精确 +`valid_cols`，推送传输携带该动态范围 —— PTOAS >= v0.55 接受 `tput` 的动态 +partition-view 形状（hw-native-sys/PTOAS#1069），因此无需填充窗口即可保留 +ragged 与 FP16 尾部。非原子推送 + 本地归约保留了所有 `ReduceOp` +（Sum/Max/Min/Prod）——只有远端原子 `TPUT` 才只支持 Sum。 该方案无需在公开 tensor 布局中插入空洞，也能支持非整除输入和 `SIZE < NR`。 -每一轮的每个 subchunk 都使用本调用局部的 ready + read-complete generation 对 +每一轮的每个 subchunk 都使用本调用局部的 ready + push-done generation 对 做屏障；尾声随后把 `2 * chunk_count`（跨轮统一，因为每轮的 subchunk 循环 共享相同边界）从 signal 的每一行中减去。 diff --git a/docs/zh/dev/passes/42-lower_host_tensor_collectives.md b/docs/zh/dev/passes/42-lower_host_tensor_collectives.md index 2d20630376..3b87fa694f 100644 --- a/docs/zh/dev/passes/42-lower_host_tensor_collectives.md +++ b/docs/zh/dev/passes/42-lower_host_tensor_collectives.md @@ -102,6 +102,15 @@ Ring allreduce 目前仅支持 `ReduceOp.Sum` 和 `dtype=FP32`。 `mode="ring"` 下尚未支持。Ring allreduce 最多支持 16 个参与设备 （`world_size <= 16`）。 +`builtin.tensor.allreduce_ring` 内核采用**推送（push）模型**：数据搬运使用 +`pto::comm::TPUT`（远端写）——reduce-scatter 阶段通过 `TPUT` 将部分和 +累加到右邻居的 slot，allgather 阶段用非原子 `TPUT` 转发每个已归约的 chunk， +与树内 `allgather` / `all_to_all` host builtin 保持一致。顺序保证为每次传输前后 +`pipe_barrier(PIPE_ALL)`，并在每次 `TNOTIFY` 前加 `dsb(DSB_DDR)`（而非 +`pto.fence.barrier_all`，后者不会排空 MTE DMA 流水线）。跨 rank 同步使用 O(1) 的 +`NeighborBarrier`（只通知/等待左右两个 ring 邻居）——在 NPU 上安全是因为 TPUT +写流水线保证数据先于信号可见，而旧的拉取（pull）模型（TLOAD/TSTORE）不具备该保证。 + HOST collective 的所有 window 操作数——data 与 signal 都是如此——必须解析为两两不同的 `WindowBuffer` 分配。同一个 `alloc_window_buffer` 上的两个 `pld.window()` view 在 in-kernel TPUT/notify 下是跨进程数据竞争：data 对 data 是 reduce 覆盖， diff --git a/python/pypto/runtime/builtins/collectives/allreduce_ring/templates/kernel.cpp.in b/python/pypto/runtime/builtins/collectives/allreduce_ring/templates/kernel.cpp.in index a163350718..4fef9f626c 100644 --- a/python/pypto/runtime/builtins/collectives/allreduce_ring/templates/kernel.cpp.in +++ b/python/pypto/runtime/builtins/collectives/allreduce_ring/templates/kernel.cpp.in @@ -11,7 +11,8 @@ // Generated from {{template_package}}. // -// Host builtin ring allreduce — in-place chunked reduce-scatter + allgather. +// Host builtin ring allreduce — in-place chunked reduce-scatter + allgather, +// TPUT push model. // // ABI (matches mesh host builtins): // tensor(0) = data window-bound FP32 buffer (InOut) @@ -23,6 +24,21 @@ // ring partitions data into NR contiguous chunks of chunk_elems = SIZE // NR // and performs RS+AG in-place on those chunk slots. The signal holds one row // per ring round (2*(NR-1) rounds) plus a final row for the return barrier. +// +// Data movement is PUSH (remote write) via pto::comm::TPUT, mirroring the +// in-tree allgather / all_to_all host builtins: +// * Reduce-scatter: each rank TPUTs its current partial chunk +// send_idx = (my_rank - step) % NR into the RIGHT neighbour's slot of the +// same index. The remote atomic-add accumulates the receiver's own value +// together with the sender's contribution (Sum only — the host builtin is +// ReduceOp::kSum by construction, see entry.cpp.in). +// * Allgather: each rank TPUTs (non-atomic) its finalized chunk +// send_idx = (my_rank - step + 1) % NR into the right neighbour's slot. +// * Ordering: pipe_barrier(PIPE_ALL) around every TPUT (drains the MTE DMA +// pipe) and dsb(DSB_DDR) before every TNOTIFY (orders the pushed writes +// ahead of the signal) — the exact allgather/all_to_all recipe. No dcci +// cacheline flush is needed: the receiver reads data the sender wrote +// remotely via TPUT, never a locally-TSTORE'd cacheline. #include #include @@ -45,6 +61,14 @@ namespace { static constexpr int kMaxSupportedRanks = 16; static constexpr int64_t kTileCount = 256; +// Barrier selection for this build of the kernel. RoundBarrier (O(P²) +// notify-all → wait-all) is the NPU-safe baseline; the O(1) NeighborBarrier +// (left/right only) is the end state — the same one as simpler #1383. +// The TPUT push engine is verified on silicon at P=2/4 (both barriers pass); +// NeighborBarrier is enabled by default since the push model provides the +// ordering the pull model lacked. +static constexpr bool kUseNeighborBarrier = true; + template AICORE inline __gm__ T *CommRemotePtr(__gm__ CommContext *ctx, __gm__ T *local_ptr, int pe) { uint64_t local_base = ctx->windowsIn[ctx->rankId]; @@ -54,21 +78,18 @@ AICORE inline __gm__ T *CommRemotePtr(__gm__ CommContext *ctx, __gm__ T *local_p // Per-round barrier row: AtomicAdd(0→1) / WaitGe(1), single-shot per row. // -// RoundBarrier follows the simpler repo's ring collective pattern -// (tests/st/worker/collectives/allreduce/kernels/aiv/allreduce_ibing_kernel.cpp): -// TNOTIFY / TWAIT + pipe_barrier(PIPE_ALL) without dcci / dsb in the barrier -// itself. The step body flushes locally-written chunks to DDR with -// dcci(…, CACHELINE_OUT) before pipe_barrier(PIPE_ALL), so TSTORE results are -// globally visible when TNOTIFY fires. TWAIT ensures the peer has also -// reached the barrier and drained its own MTE pipes. The trailing -// pipe_barrier(PIPE_ALL) then drains any in-flight TNOTIFY / TWAIT traffic -// before the next TLOAD reads the peer's buffer. -// -// Future optimisation: once the kernel moves to a TPUT push model, the -// barrier can be relaxed to O(1) per invocation (NeighborBarrier: -// left/right only). TPUT's write pipeline handles the DDR ordering. +// With the TPUT push model the barrier's job is to make every rank's completed +// step pushes (a) drained from the MTE pipe (pipe_barrier(PIPE_ALL) after the +// push loop) and (b) globally visible ahead of the signal (dsb(DSB_DDR) before +// TNOTIFY). When a rank's wait completes, all peers' step pushes have landed, +// so the rank's next-step read of its own slot observes the accumulated data. +// The trailing pipe_barrier(PIPE_ALL) drains any in-flight TNOTIFY / TWAIT +// traffic. This mirrors allgather/all_to_all Phase-2 verbatim (NOT +// pto.fence.barrier_all — a GM-scope fence does not drain MTE DMA). AICORE inline void RoundBarrier(__gm__ CommContext *ctx, __gm__ int32_t *signal_row, int my_rank, int nranks) { + pipe_barrier(PIPE_ALL); + dsb(DSB_DDR); for (int peer = 0; peer < nranks; ++peer) { if (peer == my_rank) { continue; @@ -87,6 +108,57 @@ AICORE inline void RoundBarrier(__gm__ CommContext *ctx, __gm__ int32_t *signal_ pipe_barrier(PIPE_ALL); } +// O(1) per-round barrier: notify/wait the two ring neighbours only. In the +// push model each rank's data flows exclusively to its right neighbour and +// arrives exclusively from its left neighbour, so waiting for the left +// neighbour's drain is sufficient for the receiver's reads; the left/right +// notify/wait cycle forms a valid ring barrier. +// +// This is the post-NPU-verification swap (see kUseNeighborBarrier). The +// pull-model NeighborBarrier failed on silicon (TSTORE + TLOAD do not order); +// with TPUT the write pipeline provides the ordering — the canonical pattern +// verified in simpler #1383 at P=2/4. +AICORE inline void NeighborBarrier(__gm__ CommContext *ctx, __gm__ int32_t *signal_row, int my_rank, + int nranks) { + const int left = (my_rank - 1 + nranks) % nranks; + const int right = (my_rank + 1) % nranks; + + pipe_barrier(PIPE_ALL); + dsb(DSB_DDR); + + __gm__ int32_t *remote_signal = CommRemotePtr(ctx, signal_row + my_rank, left); + pto::comm::Signal sig_left(remote_signal); + pto::comm::TNOTIFY(sig_left, static_cast(1), pto::comm::NotifyOp::AtomicAdd); + + remote_signal = CommRemotePtr(ctx, signal_row + my_rank, right); + pto::comm::Signal sig_right(remote_signal); + pto::comm::TNOTIFY(sig_right, static_cast(1), pto::comm::NotifyOp::AtomicAdd); + + pto::comm::Signal sig_wl(signal_row + left); + pto::comm::TWAIT(sig_wl, static_cast(1), pto::comm::WaitCmp::GE); + + pto::comm::Signal sig_wr(signal_row + right); + pto::comm::TWAIT(sig_wr, static_cast(1), pto::comm::WaitCmp::GE); + + pipe_barrier(PIPE_ALL); +} + +// StepBarrier dispatches to the ring barrier selected by kUseNeighborBarrier. +// With the TPUT push engine verified on silicon at P=2/4, the O(1) +// NeighborBarrier (notify/wait the two ring neighbours only) is the default: +// the TPUT write pipeline orders the pushed data ahead of the signal. The +// O(P²) RoundBarrier remains as a brute-force fallback — it notifies/waits +// every peer and stays correct on any rail, at the cost of P−1 notifies per +// round per rank. +AICORE inline void StepBarrier(__gm__ CommContext *ctx, __gm__ int32_t *signal_row, int my_rank, + int nranks) { + if (kUseNeighborBarrier) { + NeighborBarrier(ctx, signal_row, my_rank, nranks); + } else { + RoundBarrier(ctx, signal_row, my_rank, nranks); + } +} + } // namespace extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ int64_t *args) { @@ -141,132 +213,94 @@ extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ in using TileData = pto::Tile; - TileData chunk_tile(1, kTileCount); - TileData recv_tile(1, kTileCount); static_assert(sizeof({{dtype_cpp}}) == 4, "builtin.tensor.allreduce_ring currently only supports 4-byte element types"); - constexpr int kRecvTileInitFill = 0x10000; // bind-time placeholder; overwritten by TLOAD - TASSIGN(chunk_tile, 0x0); - TASSIGN(recv_tile, kRecvTileInitFill); + // Single staging tile for the TPUT push path — the allgather/all_to_all + // recipe. TPUT's single-shot path TLOADs/TSTOREs by Tile::GetValidCol() + // (ColMaskInternal), NOT the GlobalTensor shape, so the tile's column mask + // is narrowed to each chunk's exact extent below; a 256-col tile with a + // shorter transfer would over-read/overwrite adjacent chunk slots. const int tile_cols = static_cast(chunk_elems < kTileCount ? chunk_elems : kTileCount); + TileData send_tile(1, tile_cols); + TASSIGN(send_tile, 0x10000); int round = 0; - // Phase 1: reduce-scatter — (NR-1) ring steps. + // Phase 1: reduce-scatter — (NR-1) ring steps, TPUT push right. + // + // Each rank pushes its current partial of chunk send_idx = (my_rank - step) + // % NR into the RIGHT neighbour's slot of the same index. The remote + // atomic-add accumulates the receiver's own value with the sender's + // contribution, so a chunk's full sum walks the ring and lands in slot + // my_rank after the last step. The per-step barrier (before this loop) + // guarantees the left neighbour's previous-step push has landed before this + // rank reads its own slot, and that every rank has drained its own pushes. for (int step = 1; step < nranks; ++step) { - const int recv_add_idx = (my_rank - step - 1 + nranks) % nranks; - - RoundBarrier(comm_ctx, signal_base + round * signal_cols, my_rank, nranks); + StepBarrier(comm_ctx, signal_base + round * signal_cols, my_rank, nranks); ++round; - const int left = (my_rank - 1 + nranks) % nranks; - - const int left_send_idx = (left - step + nranks) % nranks; + const int right = (my_rank + 1) % nranks; + const int send_idx = (my_rank - step + nranks) % nranks; for (int64_t tile_off = 0; tile_off < chunk_elems; tile_off += tile_cols) { int64_t remain64 = chunk_elems - tile_off; int tile_elems = static_cast(remain64 < tile_cols ? remain64 : tile_cols); - - chunk_tile.SetValidShape(1, tile_elems); - recv_tile.SetValidShape(1, tile_elems); + send_tile.ColMaskInternal = tile_elems; ShapeDyn tile_shape(1, 1, 1, 1, tile_elems); StrideDyn tile_stride(tile_elems, tile_elems, tile_elems, tile_elems, 1); - { - __gm__ {{dtype_cpp}} *remote_chunk = - CommRemotePtr(comm_ctx, chunks + static_cast(left_send_idx * chunk_elems + tile_off), left); - Global remote_g(remote_chunk, tile_shape, tile_stride); - TLOAD(recv_tile, remote_g); - set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); - wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); - } - - Global acc_g(chunks + static_cast(recv_add_idx * chunk_elems + tile_off), tile_shape, tile_stride); - TLOAD(chunk_tile, acc_g); - set_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); - wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID1); - TADD(chunk_tile, chunk_tile, recv_tile); - set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); - wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); - TSTORE(acc_g, chunk_tile); - set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); - wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); - } - // Flush written chunk to DDR so peers' AG TLOAD sees fresh data. - // - // The 16-element stride only covers the whole chunk when chunk_elems is a - // multiple of 16 (64 B-aligned chunk base); numel % NR == 0 does not - // guarantee that, so an unaligned chunk's trailing cache line is not - // flushed here (and in the AG-phase flush below). Unaligned-chunk flush - // and the 32 B transfer alignment are shared mesh/ring work handled - // separately. - { - constexpr int kCacheLineFloats = 16; // 64 B cache line / 4 B per float - for (int64_t off = 0; off < chunk_elems; off += kCacheLineFloats) { - dcci(chunks + static_cast(recv_add_idx * chunk_elems + off), - SINGLE_CACHE_LINE, CACHELINE_OUT); - } + __gm__ {{dtype_cpp}} *src_ptr = + chunks + static_cast(send_idx * chunk_elems + tile_off); + __gm__ {{dtype_cpp}} *dst_ptr = CommRemotePtr(comm_ctx, src_ptr, right); + Global src_g(src_ptr, tile_shape, tile_stride); + Global dst_g(dst_ptr, tile_shape, tile_stride); + + pipe_barrier(PIPE_ALL); + pto::comm::TPUT(dst_g, src_g, send_tile); + pipe_barrier(PIPE_ALL); } - pipe_barrier(PIPE_ALL); } - // Phase 2: allgather — (NR-1) ring steps. + // Phase 2: allgather — (NR-1) ring steps, TPUT (non-atomic) push right. + // + // Each rank TPUTs its finalized chunk send_idx = (my_rank - step + 1) % NR + // into the right neighbour's slot of the same index (a plain copy — the + // chunk value is already the fully reduced sum). The receiver does nothing + // locally; the push performs the store. for (int step = 1; step < nranks; ++step) { - const int recv_idx = (my_rank - step + nranks) % nranks; - - RoundBarrier(comm_ctx, signal_base + round * signal_cols, my_rank, nranks); + StepBarrier(comm_ctx, signal_base + round * signal_cols, my_rank, nranks); ++round; - const int left = (my_rank - 1 + nranks) % nranks; - - const int left_send_idx = (left - step + 1 + nranks) % nranks; + const int right = (my_rank + 1) % nranks; + const int send_idx = (my_rank - step + 1 + nranks) % nranks; for (int64_t tile_off = 0; tile_off < chunk_elems; tile_off += tile_cols) { int64_t remain64 = chunk_elems - tile_off; int tile_elems = static_cast(remain64 < tile_cols ? remain64 : tile_cols); - - recv_tile.SetValidShape(1, tile_elems); + send_tile.ColMaskInternal = tile_elems; ShapeDyn tile_shape(1, 1, 1, 1, tile_elems); StrideDyn tile_stride(tile_elems, tile_elems, tile_elems, tile_elems, 1); - { - __gm__ {{dtype_cpp}} *remote_chunk = - CommRemotePtr(comm_ctx, chunks + static_cast(left_send_idx * chunk_elems + tile_off), left); - Global remote_g(remote_chunk, tile_shape, tile_stride); - TLOAD(recv_tile, remote_g); - set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); - wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); - } - - Global dst_g(chunks + static_cast(recv_idx * chunk_elems + tile_off), tile_shape, tile_stride); - set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); - wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); - TSTORE(dst_g, recv_tile); - set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); - wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); - } - // Flush received chunk to DDR so peers' next-round AG TLOAD sees fresh data. - { - constexpr int kCacheLineFloats = 16; // 64 B cache line / 4 B per float - for (int64_t off = 0; off < chunk_elems; off += kCacheLineFloats) { - dcci(chunks + static_cast(recv_idx * chunk_elems + off), - SINGLE_CACHE_LINE, CACHELINE_OUT); - } + __gm__ {{dtype_cpp}} *src_ptr = + chunks + static_cast(send_idx * chunk_elems + tile_off); + __gm__ {{dtype_cpp}} *dst_ptr = CommRemotePtr(comm_ctx, src_ptr, right); + Global src_g(src_ptr, tile_shape, tile_stride); + Global dst_g(dst_ptr, tile_shape, tile_stride); + + pipe_barrier(PIPE_ALL); + pto::comm::TPUT(dst_g, src_g, send_tile); + pipe_barrier(PIPE_ALL); } - pipe_barrier(PIPE_ALL); } - // Final cross-rank barrier: guarantee every rank has finished reading from - // all its peers' buffers before any rank returns to the host. Each round's - // RoundBarrier at the top of the step protects the current step's data - // motion, but there is nothing that coordinates the last allgather step's - // reads. Without this barrier a fast rank could return, have its buffer - // reused by the next host statement, and corrupt the slow right-hand peer's - // final TLOAD. - RoundBarrier(comm_ctx, signal_base + round * signal_cols, my_rank, nranks); + // Final cross-rank barrier: guarantee every rank's last pushes have landed + // (and drained) before any rank returns to the host — otherwise a fast rank + // could return, have its buffer reused by the next host statement, and + // corrupt a slow peer's pending read of the pushed data. + StepBarrier(comm_ctx, signal_base + round * signal_cols, my_rank, nranks); pipe_barrier(PIPE_ALL); } diff --git a/src/ir/transforms/lower_composite_ops_pass.cpp b/src/ir/transforms/lower_composite_ops_pass.cpp index c5fc7c48b6..7626df5d87 100644 --- a/src/ir/transforms/lower_composite_ops_pass.cpp +++ b/src/ir/transforms/lower_composite_ops_pass.cpp @@ -1269,6 +1269,18 @@ ExprPtr LowerTensorRingAllReduceRule(const CallPtr& call, const std::vectordtype_}, {"target_memory", MemorySpace::Vec}}, span), + span); // Value-producing IfExpr branches must agree on a fixed TileType. For an // inactive logical segment, read one in-bounds element and pad it to the // physical chunk shape. Using tile.create here would survive the default @@ -1291,23 +1303,6 @@ ExprPtr LowerTensorRingAllReduceRule(const CallPtr& call, const std::vector> remote_load_kwargs; - if (target_type->dtype_ == DataType::FP16) { - remote_load_kwargs.emplace_back("allow_physical_tail_padding", true); - } - auto remote_valid_cols = [&](const ExprPtr& logical_valid_cols) { - if (target_type->dtype_ != DataType::FP16) return logical_valid_cols; - return MakeMul(MakeFloorDiv(MakeAdd(logical_valid_cols, alignment_minus_one_idx, span), - alignment_elements_idx, span), - alignment_elements_idx, span); - }; - auto restore_remote_valid_shape = [&](LoweringBuilder& body, const ExprPtr& loaded, - const ExprPtr& logical_valid_cols, - const std::string& name) -> ExprPtr { - if (target_type->dtype_ != DataType::FP16) return loaded; - return body.Bind(name, reg.Create("tile.set_validshape", {loaded, one_idx, logical_valid_cols}, {}, span), - span); - }; auto emit_barrier = [&](LoweringBuilder& body, const ExprPtr& round, const ExprPtr& expected, const std::string& suffix) { body.EmitNotifyAll(signal, comm.nranks_idx, comm.my_rank, round, NotifyOp::kAtomicAdd, one_i32, suffix, @@ -1319,78 +1314,172 @@ ExprPtr LowerTensorRingAllReduceRule(const CallPtr& call, const std::vector= v0.55 + // accepts dynamic partition-view shapes — issue #1069). + // * Push-done barrier (generation 2k+2): all pushes have landed. + // * Each rank reads its own slot recv_idx again (now holding the left + // neighbour's partial), reduces it with the saved own value, and stores + // the result back. The push is a plain copy into the receiver's own + // window — no remote read, so the only cross-rank visibility needed is + // the TPUT→notify ordering the credit-barrier protocol provides + // (pld.tile.put codegen emits pipe_barrier(PIPE_ALL) around the TPUT and + // the comm-fence pass orders it ahead of the notify). b.EmitFor( "rs_step", zero_idx, nr_minus_one, one_idx, [&](LoweringBuilder& body, const VarPtr& rs_step_var) { auto step = body.Bind("step", MakeAdd(rs_step_var, one_idx, span), span); - // recv_add_idx = (my_rank − step − 1 + NR) % NR + // recv_idx = (my_rank − step − 1 + NR) % NR auto r1 = MakeSub(my_rank_idx, step, span); auto r2 = MakeSub(r1, one_idx, span); auto r3 = MakeAdd(r2, comm.nranks_idx, span); - // recv_add_idx and send_idx are the same chunk index in this - // reduce-scatter formulation — bind once and reuse. - auto recv_add_idx = body.Bind("recv_add_idx", MakeFloorMod(r3, comm.nranks_idx, span), span); - const auto& send_idx = recv_add_idx; + auto recv_idx = body.Bind("recv_idx", MakeFloorMod(r3, comm.nranks_idx, span), span); + + // send_idx = (my_rank − step + NR) % NR — one chunk ahead of recv_idx. + auto s1 = MakeSub(my_rank_idx, step, span); + auto s2 = MakeAdd(s1, comm.nranks_idx, span); + auto send_idx = body.Bind("send_idx", MakeFloorMod(s2, comm.nranks_idx, span), span); - // left = (my_rank − 1 + NR) % NR - auto l1 = MakeSub(my_rank_idx, one_idx, span); - auto l2 = MakeAdd(l1, comm.nranks_idx, span); - auto left_peer = body.Bind("left", MakeFloorMod(l2, comm.nranks_idx, span), span); + // right = (my_rank + 1) % NR — the push destination. + auto rr1 = MakeAdd(my_rank_idx, one_idx, span); + auto right_peer = body.Bind("right", MakeFloorMod(rr1, comm.nranks_idx, span), span); - auto segment_offset = body.Bind("rs_segment_begin", segment_begin(send_idx), span); - auto segment_limit = body.Bind("rs_segment_end", segment_end(send_idx), span); - auto segment_cols = body.Bind("rs_segment_cols", MakeSub(segment_limit, segment_offset, span), span); + auto recv_segment_offset = body.Bind("rs_recv_segment_begin", segment_begin(recv_idx), span); + auto recv_segment_limit = body.Bind("rs_recv_segment_end", segment_end(recv_idx), span); + auto recv_segment_cols = + body.Bind("rs_recv_segment_cols", MakeSub(recv_segment_limit, recv_segment_offset, span), span); + + auto send_segment_offset = body.Bind("rs_send_segment_begin", segment_begin(send_idx), span); + auto send_segment_limit = body.Bind("rs_send_segment_end", segment_end(send_idx), span); + auto send_segment_cols = + body.Bind("rs_send_segment_cols", MakeSub(send_segment_limit, send_segment_offset, span), span); body.EmitFor( "rs_col", zero_idx, max_segment_cols, chunk_cols, [&](LoweringBuilder& chunk_body, const VarPtr& subcol) { - auto active = MakeLt(subcol, segment_cols, span); - auto remaining = MakeSub(segment_cols, subcol, span); - auto valid_cols = MakeMin(chunk_cols, remaining, span); + // Receive-side extent (own-value read, pushed-partial read, + // reduce, store all operate on slot recv_idx). + auto recv_active = MakeLt(subcol, recv_segment_cols, span); + auto recv_remaining = MakeSub(recv_segment_cols, subcol, span); + auto recv_valid_cols = MakeMin(chunk_cols, recv_remaining, span); // Keep the value-producing IfExpr branch metadata identical. // Inactive ranks use one safe element, while active ranks retain // the exact logical tail extent. - auto load_valid_cols = MakeMax(one_idx, valid_cols, span); + auto load_valid_cols = MakeMax(one_idx, recv_valid_cols, span); auto load_valid_shape = tile_conversion_utils::MakeShapeTuple({one_idx, load_valid_cols}, span); - auto remote_load_valid_shape = - tile_conversion_utils::MakeShapeTuple({one_idx, remote_valid_cols(load_valid_cols)}, span); - auto offsets = tile_conversion_utils::MakeShapeTuple( - {zero_idx, MakeAdd(segment_offset, subcol, span)}, span); + auto recv_offsets = tile_conversion_utils::MakeShapeTuple( + {zero_idx, MakeAdd(recv_segment_offset, subcol, span)}, span); + + // Send-side extent (TPUT source = slot send_idx). + auto send_active = MakeLt(subcol, send_segment_cols, span); + auto send_remaining = MakeSub(send_segment_cols, subcol, span); + auto send_valid_cols = MakeMin(chunk_cols, send_remaining, span); + auto send_valid_shape = tile_conversion_utils::MakeShapeTuple({one_idx, send_valid_cols}, span); + auto send_offsets = tile_conversion_utils::MakeShapeTuple( + {zero_idx, MakeAdd(send_segment_offset, subcol, span)}, span); auto chunk_id = MakeFloorDiv(subcol, chunk_cols, span); + + // ---- Own-value read: BEFORE the ready barrier, so no push has + // landed in slot recv_idx yet (the slot is stable — own value). + auto acc_own = chunk_body.EmitIfExpr( + recv_active, + [&](LoweringBuilder& then_body) { + auto acc_loaded = then_body.Bind( + "acc_rs_own_loaded", + reg.Create("tile.load", {ring_target, recv_offsets, chunk_shape, load_valid_shape}, + {{"target_memory", MemorySpace::Vec}}, span), + span); + return then_body.Bind("acc_rs_own", + reg.Create("tile.fillpad_inplace", {acc_loaded}, + {{"pad_value", PadValue::zero}}, span), + span); + }, + [&](LoweringBuilder& else_body) { + auto placeholder_loaded = else_body.Bind( + "acc_rs_own_placeholder_loaded", + reg.Create("tile.load", + {ring_target, placeholder_offsets, chunk_shape, load_valid_shape}, + {{"target_memory", MemorySpace::Vec}}, span), + span); + return else_body.Bind("acc_rs_own_placeholder", + reg.Create("tile.fillpad_inplace", {placeholder_loaded}, + {{"pad_value", PadValue::zero}}, span), + span); + }, + span); + + // ---- Ready barrier: all ranks' own-value reads are done; no + // push can land before this barrier completes. auto ready_epoch_idx = MakeAdd(MakeMul(chunk_id, two_idx, span), one_idx, span); auto ready_epoch = chunk_body.Bind( "rs_ready_epoch", std::make_shared(ready_epoch_idx, DataType::INT32, span), span); emit_barrier(chunk_body, rs_step_var, ready_epoch, "_rs_ready"); + // ---- Push phase: TPUT slot send_idx into the right neighbour's + // slot of the same index (plain copy; the receiver reduces + // locally after the push-done barrier). The staging tile is + // narrowed to the transfer width so the TPUT single-shot path + // reads exactly send_valid_cols. + chunk_body.EmitIf( + send_active, + [&](LoweringBuilder& push_body) { + auto rs_stage_valid = push_body.Bind( + "rs_stage_valid", + reg.Create("tile.set_validshape", {put_stage, one_idx, send_valid_cols}, {}, span), + span); + push_body.Bind("push_rs", + reg.Create("pld.tile.put", + {ring_target, right_peer, ring_target, rs_stage_valid, + send_offsets, send_offsets, send_valid_shape}, + {{"atomic", static_cast(AtomicType::kNone)}}, span), + span); + }, + /*else_fn=*/nullptr, span); + + // ---- Push-done barrier: all pushes have landed — slot recv_idx + // now holds the left neighbour's partial. + auto read_epoch_idx = MakeAdd(ready_epoch_idx, one_idx, span); + auto read_epoch = chunk_body.Bind( + "rs_read_epoch", std::make_shared(read_epoch_idx, DataType::INT32, span), span); + emit_barrier(chunk_body, rs_step_var, read_epoch, "_rs_read"); + + // ---- Local reduce + store: read the pushed partial (own slot + // recv_idx, now = left's partial), combine with the saved own + // value, store back. auto acc_full = chunk_body.EmitIfExpr( - active, + recv_active, [&](LoweringBuilder& then_body) { auto recv_loaded = then_body.Bind( "recv_rs_loaded", - reg.Create("pld.tile.remote_load", - {ring_target, left_peer, offsets, chunk_shape, remote_load_valid_shape}, - remote_load_kwargs, span), + reg.Create("tile.load", {ring_target, recv_offsets, chunk_shape, load_valid_shape}, + {{"target_memory", MemorySpace::Vec}}, span), span); - auto recv_tail = - restore_remote_valid_shape(then_body, recv_loaded, load_valid_cols, "recv_rs_tail"); auto recv = then_body.Bind("recv_rs", - reg.Create("tile.fillpad_inplace", {recv_tail}, + reg.Create("tile.fillpad_inplace", {recv_loaded}, {{"pad_value", PadValue::zero}}, span), span); - auto acc_loaded = then_body.Bind( - "acc_rs_loaded", - reg.Create("tile.load", {ring_target, offsets, chunk_shape, load_valid_shape}, - {{"target_memory", MemorySpace::Vec}}, span), - span); - auto acc = then_body.Bind("acc_rs", - reg.Create("tile.fillpad_inplace", {acc_loaded}, - {{"pad_value", PadValue::zero}}, span), - span); - return then_body.Bind("acc_rs_next", then_body.Reduce(reduce_op, acc, recv, span), span); + return then_body.Bind("acc_rs_next", then_body.Reduce(reduce_op, acc_own, recv, span), + span); }, [&](LoweringBuilder& else_body) { auto placeholder_loaded = else_body.Bind( @@ -1406,22 +1495,17 @@ ExprPtr LowerTensorRingAllReduceRule(const CallPtr& call, const std::vector(read_epoch_idx, DataType::INT32, span), span); - emit_barrier(chunk_body, rs_step_var, read_epoch, "_rs_read"); - chunk_body.EmitIf( - active, + recv_active, [&](LoweringBuilder& store_body) { // Encode the active-branch bounds in the store operands so // valid-region inference can prove this write stays inside // the flattened logical extent without relying on control // flow predicates. - auto raw_store_col = MakeAdd(segment_offset, subcol, span); + auto raw_store_col = MakeAdd(recv_segment_offset, subcol, span); auto store_col = MakeSub( size_expr, MakeMax(zero_idx, MakeSub(size_expr, raw_store_col, span), span), span); - auto raw_store_end = MakeAdd(store_col, valid_cols, span); + auto raw_store_end = MakeAdd(store_col, recv_valid_cols, span); auto store_end = MakeSub( size_expr, MakeMax(zero_idx, MakeSub(size_expr, raw_store_end, span), span), span); auto store_valid_cols = MakeSub(store_end, store_col, span); @@ -1441,27 +1525,38 @@ ExprPtr LowerTensorRingAllReduceRule(const CallPtr& call, const std::vector= v0.55 dynamic partition-view shapes). The ready barrier + // (generation 2k+1) guarantees the pushed slot's data is valid (the + // cumulative generation count includes the previous step's push-done + // notifies); the push-done barrier (generation 2k+2) guarantees the copy is + // visible before any rank reads it at the next step. b.EmitFor( "ag_step", zero_idx, nr_minus_one, one_idx, [&](LoweringBuilder& body, const VarPtr& ag_step_var) { auto step = body.Bind("ag_step_val", MakeAdd(ag_step_var, one_idx, span), span); auto ag_round = body.Bind("ag_round", MakeAdd(ag_step_var, nr_minus_one, span), span); + // send_idx = (my_rank − step + 1 + NR) % NR — the fully-reduced chunk + // this rank forwards to the right neighbour. auto r1 = MakeSub(my_rank_idx, step, span); - auto r2 = MakeAdd(r1, comm.nranks_idx, span); - auto segment_idx = body.Bind("ag_segment_idx", MakeFloorMod(r2, comm.nranks_idx, span), span); - - // left = (my_rank - 1 + NR) % NR is the peer that already owns this - // step's segment. - auto l1 = MakeSub(my_rank_idx, one_idx, span); - auto l2 = MakeAdd(l1, comm.nranks_idx, span); - auto left_val = MakeFloorMod(l2, comm.nranks_idx, span); - auto left_peer = body.Bind("ag_left", left_val, span); - - auto segment_offset = body.Bind("ag_segment_begin", segment_begin(segment_idx), span); - auto segment_limit = body.Bind("ag_segment_end", segment_end(segment_idx), span); + auto r2 = MakeAdd(r1, one_idx, span); + auto r3 = MakeAdd(r2, comm.nranks_idx, span); + auto send_idx = body.Bind("ag_send_idx", MakeFloorMod(r3, comm.nranks_idx, span), span); + + // right = (my_rank + 1) % NR — the push destination. + auto rr1 = MakeAdd(my_rank_idx, one_idx, span); + auto right_peer = body.Bind("ag_right", MakeFloorMod(rr1, comm.nranks_idx, span), span); + + auto segment_offset = body.Bind("ag_segment_begin", segment_begin(send_idx), span); + auto segment_limit = body.Bind("ag_segment_end", segment_end(send_idx), span); auto segment_cols = body.Bind("ag_segment_cols", MakeSub(segment_limit, segment_offset, span), span); body.EmitFor( @@ -1470,77 +1565,45 @@ ExprPtr LowerTensorRingAllReduceRule(const CallPtr& call, const std::vector(ready_epoch_idx, DataType::INT32, span), span); emit_barrier(chunk_body, ag_round, ready_epoch, "_ag_ready"); - auto recv_full = chunk_body.EmitIfExpr( + // ---- Push phase: TPUT slot send_idx into the right neighbour's + // slot of the same index (plain copy). The staging tile is + // narrowed to the transfer width so the TPUT single-shot path + // reads exactly valid_cols. + chunk_body.EmitIf( active, - [&](LoweringBuilder& then_body) { - auto recv_loaded = then_body.Bind( - "recv_ag_loaded", - reg.Create("pld.tile.remote_load", - {ring_target, left_peer, offsets, chunk_shape, remote_load_valid_shape}, - remote_load_kwargs, span), - span); - auto recv_tail = - restore_remote_valid_shape(then_body, recv_loaded, load_valid_cols, "recv_ag_tail"); - return then_body.Bind("recv_ag", - reg.Create("tile.fillpad_inplace", {recv_tail}, - {{"pad_value", PadValue::zero}}, span), - span); - }, - [&](LoweringBuilder& else_body) { - auto placeholder_loaded = else_body.Bind( - "recv_ag_placeholder_loaded", - reg.Create("tile.load", - {ring_target, placeholder_offsets, chunk_shape, load_valid_shape}, - {{"target_memory", MemorySpace::Vec}}, span), - span); - return else_body.Bind("recv_ag_placeholder", - reg.Create("tile.fillpad_inplace", {placeholder_loaded}, - {{"pad_value", PadValue::zero}}, span), - span); + [&](LoweringBuilder& push_body) { + auto ag_stage_valid = push_body.Bind( + "ag_stage_valid", + reg.Create("tile.set_validshape", {put_stage, one_idx, valid_cols}, {}, span), span); + push_body.Bind("push_ag", + reg.Create("pld.tile.put", + {ring_target, right_peer, ring_target, ag_stage_valid, offsets, + offsets, valid_shape}, + {{"atomic", static_cast(AtomicType::kNone)}}, span), + span); }, - span); + /*else_fn=*/nullptr, span); + // ---- Push-done barrier: all pushes have landed — every rank's + // forwarded slot is visible before the next step reads it. auto read_epoch_idx = MakeAdd(ready_epoch_idx, one_idx, span); auto read_epoch = chunk_body.Bind( "ag_read_epoch", std::make_shared(read_epoch_idx, DataType::INT32, span), span); emit_barrier(chunk_body, ag_round, read_epoch, "_ag_read"); - - chunk_body.EmitIf( - active, - [&](LoweringBuilder& store_body) { - // See the reduce-scatter store above: these clamped - // expressions are no-ops for active chunks and make both - // the offset and far edge statically bounded by size_expr. - auto raw_store_col = MakeAdd(segment_offset, subcol, span); - auto store_col = MakeSub( - size_expr, MakeMax(zero_idx, MakeSub(size_expr, raw_store_col, span), span), span); - auto raw_store_end = MakeAdd(store_col, valid_cols, span); - auto store_end = MakeSub( - size_expr, MakeMax(zero_idx, MakeSub(size_expr, raw_store_end, span), span), span); - auto store_valid_cols = MakeSub(store_end, store_col, span); - auto store_offsets = tile_conversion_utils::MakeShapeTuple({zero_idx, store_col}, span); - auto narrowed = store_body.Bind( - "recv_ag_valid", - reg.Create("tile.set_validshape", {recv_full, one_idx, store_valid_cols}, {}, span), - span); - store_body.Bind( - "store_ag", - reg.Create("tile.store", {narrowed, store_offsets, ring_target}, {}, span), span); - }, - /*else_fn=*/nullptr, span); }, span); }, diff --git a/tests/ut/ir/transforms/test_lower_composite_ops.py b/tests/ut/ir/transforms/test_lower_composite_ops.py index 8b1aeb76d1..e8c33de1dd 100644 --- a/tests/ut/ir/transforms/test_lower_composite_ops.py +++ b/tests/ut/ir/transforms/test_lower_composite_ops.py @@ -1943,7 +1943,7 @@ def f( "pld.system.rank", "pld.system.notify", # per-round barrier (2(P−1) rounds) "pld.system.wait", # per-round barrier - "pld.tile.remote_load", # per-ring-step chunk receive + "pld.tile.put", # TPUT push of each ring step's chunk to the right neighbour "tile.add", # reduce-scatter accumulation "tile.load", # reduce-scatter local accumulation "tile.fillpad_inplace", # promote ragged subchunks for fixed-shape arithmetic @@ -1995,8 +1995,8 @@ def test_ring_allreduce_is_decomposed_to_primitives(): assert ir.get_op("pld.tensor.allreduce").name not in op_names, ( "lower_composite_ops must remove the composite allreduce call entirely" ) - assert ir.get_op("tile.create").name not in op_names, ( - "inactive ring segments must not leave allocation-only placeholders" + assert ir.get_op("tile.create").name in op_names, ( + "the TPUT push path must emit the tile.create staging tile" ) missing = _RING_ALLREDUCE_REQUIRED_OPS - op_names assert not missing, f"ring-lowered IR missing expected ops: {missing}" @@ -2014,7 +2014,7 @@ def test_ring_allreduce_emits_ring_control_flow(): collector.visit_program(After) assert collector.for_count == 14, f"expected 14 ForStmts for P=2 ring, got {collector.for_count}" - assert collector.if_count == 13, f"expected 13 IfStmts for P=2 ring, got {collector.if_count}" + assert collector.if_count == 14, f"expected 14 IfStmts for P=2 ring, got {collector.if_count}" @pytest.mark.parametrize("size", [1, 3, 17, 8193, 65537]) @@ -2039,18 +2039,19 @@ def visit_call(self, op: ir.Call) -> None: collector = CallCollector() collector.visit_program(After) - remote_loads = [ - call for call in collector.calls if call.op.name == ir.get_op("pld.tile.remote_load").name - ] + puts = [call for call in collector.calls if call.op.name == ir.get_op("pld.tile.put").name] + stage_creates = [call for call in collector.calls if call.op.name == ir.get_op("tile.create").name] loads = [call for call in collector.calls if call.op.name == ir.get_op("tile.load").name] set_valid_shapes = [ call for call in collector.calls if call.op.name == ir.get_op("tile.set_validshape").name ] - assert remote_loads + assert puts, "ring-lowered IR must emit pld.tile.put pushes" + assert stage_creates, "ring-lowered IR must emit the TPUT staging tile" assert loads assert set_valid_shapes - assert all(len(call.args) == 5 for call in remote_loads) + # pld.tile.put(dst, peer, src, stage, dst_offsets, src_offsets, shape) + assert all(len(call.args) == 7 for call in puts) loops: list[ir.ForStmt] = [] @@ -2084,7 +2085,9 @@ def collect_loops(stmt: ir.Stmt) -> None: assert isinstance(loop.start, ir.ConstInt) and loop.start.value == 0 assert isinstance(loop.stop, ir.ConstInt) and loop.stop.value == max_segment - chunk_shapes = [call.args[3] for call in remote_loads] + # The static TPUT staging tile carries the UB-bounded, 32-byte-aligned + # chunk width (tile.create shape [1, chunk_cols]). + chunk_shapes = [call.args[0] for call in stage_creates] for shape in chunk_shapes: assert isinstance(shape, ir.MakeTuple) chunk_rows = shape.elements[0] @@ -2099,7 +2102,7 @@ def collect_loops(stmt: ir.Stmt) -> None: @pytest.mark.parametrize("size", [1, 17, 33, 8193, 65537]) @pytest.mark.parametrize("n_ranks", [2, 4]) def test_ring_allreduce_fp16_uses_aligned_ring_schedule(size, n_ranks): - """FP16 stays on the ring path and marks every remote tail as padded.""" + """FP16 stays on the ring path with a 16-element-aligned TPUT staging tile.""" Before = _build_ring_allreduce_before( size=size, n_ranks=n_ranks, @@ -2118,12 +2121,11 @@ def visit_call(self, op: ir.Call) -> None: collector = CallCollector() collector.visit_program(After) - remote_loads = [ - call for call in collector.calls if call.op.name == ir.get_op("pld.tile.remote_load").name - ] - assert remote_loads - assert all(call.kwargs.get("allow_physical_tail_padding") is True for call in remote_loads) - assert all(len(call.args) == 5 for call in remote_loads) + puts = [call for call in collector.calls if call.op.name == ir.get_op("pld.tile.put").name] + stage_creates = [call for call in collector.calls if call.op.name == ir.get_op("tile.create").name] + assert puts + assert stage_creates + assert all(len(call.args) == 7 for call in puts) stmt_collector = _StmtKindCollector() stmt_collector.visit_program(After) @@ -2133,7 +2135,7 @@ def visit_call(self, op: ir.Call) -> None: max_segment = min(size, (size + n_ranks - 1) // n_ranks + 15) expected_chunk = min(8192, ((max_segment + 15) // 16) * 16) - chunk_shapes = [call.args[3] for call in remote_loads] + chunk_shapes = [call.args[0] for call in stage_creates] for shape in chunk_shapes: assert isinstance(shape, ir.MakeTuple) chunk_cols = shape.elements[1] @@ -2144,13 +2146,12 @@ def visit_call(self, op: ir.Call) -> None: def test_ring_allreduce_fp16_lowered_ir_round_trips(): - """The compiler-only aligned remote tail survives print and reparse.""" + """The push-based ring schedule (pld.tile.put) survives print and reparse.""" Before = _build_ring_allreduce_before(size=17, n_ranks=2, dtype=pl.FP16) After = passes.lower_composite_ops()(Before) text = ir.python_print(After) - assert "pld.tile._remote_load_with_physical_tail_padding(" in text - assert "allow_physical_tail_padding=" not in text + assert "pld.tile.put(" in text reparsed = pl.parse_program(text) ir.assert_structural_equal(After, reparsed)