diff --git a/docs/en/dev/distributed_ops.md b/docs/en/dev/distributed_ops.md index 183d6d9c6c..24771362c0 100644 --- a/docs/en/dev/distributed_ops.md +++ b/docs/en/dev/distributed_ops.md @@ -338,7 +338,7 @@ Variable-size all-to-all (MPI_Alltoallv). Flat 2D layouts: - `input` — Tensor or DistributedTensor `[NR*MAX_RECV, SIZE]` - `target` — DistributedTensor `[NR*MAX_RECV, SIZE]` (window-as-result) -- `signal` — DistributedTensor INT32 `[NR, 1]` (single-use Set(1)/wait≥1 barrier) +- `signal` — DistributedTensor INT32 `[NR, 1]` (self-clearing credit barrier; reusable across calls) - `send_counts` — Tensor-like INT32 `[NR]` or `[NR, 1]` (runtime rows per dest) - `recv_counts` — DistributedTensor INT32 `[NR, 1]` (InOut recvcounts) diff --git a/docs/en/dev/passes/40-synthesize_allreduce_signals.md b/docs/en/dev/passes/40-synthesize_allreduce_signals.md index 345c7e972c..023a15f425 100644 --- a/docs/en/dev/passes/40-synthesize_allreduce_signals.md +++ b/docs/en/dev/passes/40-synthesize_allreduce_signals.md @@ -77,13 +77,13 @@ The pass raises `pypto::ValueError` when: expression statement, or return value, - an allreduce appears inside a `for` / `while` loop. -The loop restriction applies to the HOST rail: the `builtin.tensor.allreduce` -kernel (lowered by `LowerHostTensorCollectives`) is not self-clearing — it adds -ready/per-chunk credits via `AtomicAdd(+1)` and never subtracts them — so a -signal synthesized (or explicitly passed) before a loop would be reused on a -later iteration with stale `>=` thresholds. InCore composites lowered by +The loop restriction applies to the HOST rail: `SynthesizeAllReduceSignals` +inserts the synthesized signal allocation immediately before the allreduce +statement, which cannot be placed inside a dynamic loop (a fresh allocation per +iteration under the same name, and every rank must land on the same symmetric +window). InCore composites lowered by [`LowerCompositeOps`](12-lower_composite_ops.md#barrier-signal-protocol) are -loop-safe because that pass emits the self-clearing epilogue. +loop-safe because that pass emits the self-clearing credit-barrier epilogue. ## Pass Properties diff --git a/docs/en/user/distributed/01-collectives.md b/docs/en/user/distributed/01-collectives.md index 0b5ba1b55c..8daffb7141 100644 --- a/docs/en/user/distributed/01-collectives.md +++ b/docs/en/user/distributed/01-collectives.md @@ -108,8 +108,8 @@ Cross-rank barrier — blocks until all ranks arrive. signal = pld.tensor.barrier(signal) ``` -Uses `Set(1)` + `Ge(1)` on the signal. Single-shot; allocate a fresh buffer -before the next barrier. +Uses a self-clearing credit barrier (`AtomicAdd(+1)` / `Ge(1)` with a reset +epilogue), so one signal buffer is reusable across back-to-back calls. ## Broadcast diff --git a/docs/en/user/distributed/02-primitives.md b/docs/en/user/distributed/02-primitives.md index 908b2af204..eb375545e2 100644 --- a/docs/en/user/distributed/02-primitives.md +++ b/docs/en/user/distributed/02-primitives.md @@ -95,9 +95,11 @@ def handshake_step( > **Buffer re-use safety:** Signal cells are zero-initialised by > `alloc_window_buffer`. After `notify`, the signal cell holds the written -> value; after `wait` returns, the caller has observed the barrier. Do not -> reuse the same signal buffer across back-to-back collectives — the protocol -> uses monotonic counters that do not self-reset. Allocate a fresh buffer. +> value; after `wait` returns, the caller has observed the barrier. These +> tile-level `notify`/`wait` primitives use monotonic counters that do not +> self-reset — allocate a fresh buffer per call. The `pld.tensor.*` +> collectives are the exception: their signal buffers are self-clearing and +> reusable across back-to-back calls. ## Tile-Level RMA (`pld.tile.*`) diff --git a/docs/en/user/distributed/04-debugging.md b/docs/en/user/distributed/04-debugging.md index 1a38daab4a..11675f8561 100644 --- a/docs/en/user/distributed/04-debugging.md +++ b/docs/en/user/distributed/04-debugging.md @@ -12,7 +12,7 @@ one rank while the cause is on another. | **Signal cell never reaches expected value** | Wrong `NotifyOp`: used `Set` instead of `AtomicAdd` for a multi-participant barrier | Use `AtomicAdd` when N ranks contribute to the same slot; use `Set` for 1:1 exchanges. | | **Shape mismatch at compile time** | `NR` (world size) used in type annotations without `pl.dynamic` | Wrap runtime-resolved dims in `pl.dynamic("NR")`. The compiler needs the name to bind the runtime value. | | **`TypeError` raised at dispatch** | IO buffer not `.share_memory_()` before `prepare()` — the child processes cannot see a buffer allocated after the fork | Call `.share_memory_()` on every host tensor passed to the worker, before `prepare()`. | -| **Allreduce rejected inside loop** | Signal protocol can't inject a fresh buffer per iteration | Allocate a fresh signal buffer for each allreduce call outside loops; allreduce inside `for`/`while` is currently rejected. | +| **Allreduce rejected inside loop** | HOST-rail allreduce inside a dynamic `for`/`while` is rejected: signal synthesis cannot allocate a fresh signal per iteration (InCore composites are loop-safe via the self-clearing credit-barrier protocol) | Hoist HOST allreduce calls out of the loop. | ## Fatal Pitfalls diff --git a/docs/zh/dev/distributed_ops.md b/docs/zh/dev/distributed_ops.md index a853ca94e8..85892d2376 100644 --- a/docs/zh/dev/distributed_ops.md +++ b/docs/zh/dev/distributed_ops.md @@ -297,7 +297,7 @@ pld.tensor.all_to_all_v( - `input` — Tensor 或 DistributedTensor `[NR*MAX_RECV, SIZE]` - `target` — DistributedTensor `[NR*MAX_RECV, SIZE]`(窗口即结果) -- `signal` — DistributedTensor INT32 `[NR, 1]`(单次使用的 Set(1)/wait≥1 屏障) +- `signal` — DistributedTensor INT32 `[NR, 1]`(自清理信用屏障;可在多次调用间复用) - `send_counts` — Tensor-like INT32 `[NR]` 或 `[NR, 1]`(运行时每目标行数) - `recv_counts` — DistributedTensor INT32 `[NR, 1]`(InOut recvcounts) diff --git a/docs/zh/dev/passes/40-synthesize_allreduce_signals.md b/docs/zh/dev/passes/40-synthesize_allreduce_signals.md index b708041505..bc1da89f9e 100644 --- a/docs/zh/dev/passes/40-synthesize_allreduce_signals.md +++ b/docs/zh/dev/passes/40-synthesize_allreduce_signals.md @@ -68,7 +68,7 @@ alloc / window / allreduce 链路。 - allreduce 作为嵌套表达式出现,而不是直接赋值、表达式语句或 return value; - allreduce 出现在 `for` / `while` 循环内。 -该循环限制针对 HOST 通道:`builtin.tensor.allreduce` kernel(由 `LowerHostTensorCollectives` lower)不是自清理的 —— 它用 `AtomicAdd(+1)` 增加 ready/per-chunk 信用却从不回减 —— 因此循环前合成(或显式传入)的 signal 会在后续迭代中复用残留的 `>=` 阈值。由 [`LowerCompositeOps`](12-lower_composite_ops.md#屏障-信号协议) lower 的 InCore 组合算子则因该 pass 会发出自清理尾声而具备循环安全性。 +该循环限制针对 HOST 通道:`SynthesizeAllReduceSignals` 把合成的 signal 分配插入到 allreduce 语句之前,无法放入动态循环内部(每次迭代在同一名字下重新分配,且每个 rank 必须落在同一个对称 window 上)。由 [`LowerCompositeOps`](12-lower_composite_ops.md#屏障-信号协议) lower 的 InCore 组合算子则因该 pass 会发出自清理信用屏障尾声而具备循环安全性。 ## Pass 属性 diff --git a/docs/zh/user/distributed/01-collectives.md b/docs/zh/user/distributed/01-collectives.md index ff2e795036..a94fb18ded 100644 --- a/docs/zh/user/distributed/01-collectives.md +++ b/docs/zh/user/distributed/01-collectives.md @@ -100,8 +100,8 @@ mesh 路径均支持。Host 内置的 ring 路径(`builtin.tensor.allreduce_ri signal = pld.tensor.barrier(signal) ``` -在 signal 上使用 `Set(1)` + `Ge(1)`。单次使用;下一次 barrier 前需分配新 -buffer。 +在 signal 上使用自清理信用屏障(`AtomicAdd(+1)` / `Ge(1)` 并带重置尾声), +因此同一个 signal buffer 可在连续调用间复用。 ## Broadcast diff --git a/docs/zh/user/distributed/02-primitives.md b/docs/zh/user/distributed/02-primitives.md index d15431a86a..4db3c2ff7b 100644 --- a/docs/zh/user/distributed/02-primitives.md +++ b/docs/zh/user/distributed/02-primitives.md @@ -80,8 +80,9 @@ def handshake_step( `outputs[0] == 1`。rank 1 写入 tag=1,等待来自 rank 0 的 tag 2: `outputs[1] == 2`。结果:`outputs == [[1], [2]]`。 -> **Buffer 重用安全:** Signal 使用单调计数器且不会自重置。不要在背靠背集合通信中 -> 重用同一 signal buffer。每次调用分配新 buffer。 +> **Buffer 重用安全:** Signal 使用单调计数器且不会自重置。这些 tile 级 +> `notify`/`wait` 原语每次调用需分配新 buffer;`pld.tensor.*` 集合通信除外, +> 其 signal buffer 自清理,可在连续调用间复用。 ## Tile 级 RMA (`pld.tile.*`) diff --git a/docs/zh/user/distributed/04-debugging.md b/docs/zh/user/distributed/04-debugging.md index d6c92580ee..516a19aa21 100644 --- a/docs/zh/user/distributed/04-debugging.md +++ b/docs/zh/user/distributed/04-debugging.md @@ -12,7 +12,7 @@ rank 上。 | **Signal cell 永不达到期望值** | 错误 `NotifyOp` | 多参与者屏障用 `AtomicAdd`;1:1 交换用 `Set`。 | | **编译时形状不匹配** | `NR` 未使用 `pl.dynamic` | 将运行时维度包裹在 `pl.dynamic("NR")` 中。 | | **派发时抛出 `TypeError`** | IO buffer 在 `prepare()` 前未调用 `.share_memory_()`——fork 出的子进程看不到 fork 之后分配的 buffer | 在 `prepare()` 之前对每个传给 worker 的 host tensor 调用 `.share_memory_()`。 | -| **循环内 allreduce 被拒绝** | Signal 协议无法每轮注入新 buffer | 在循环外每次调用分配新 signal buffer。 | +| **循环内 allreduce 被拒绝** | HOST 轨的 allreduce 在动态 `for`/`while` 内被拒绝:signal 合成无法为每次迭代分配新 signal(InCore 复合算子通过自清理信用屏障协议可在循环内使用) | 将 HOST allreduce 调用提到循环外。 | ## 致命陷阱 diff --git a/python/pypto/ir/op/distributed/tensor_ops.py b/python/pypto/ir/op/distributed/tensor_ops.py index 164335dc60..0deb84a435 100644 --- a/python/pypto/ir/op/distributed/tensor_ops.py +++ b/python/pypto/ir/op/distributed/tensor_ops.py @@ -238,8 +238,9 @@ def allreduce( ``target`` holds the reduced value. ``signal``, when provided, is a window-bound INT32 matrix used as the cross-rank barrier. Host-level calls may omit it; SynthesizeAllReduceSignals inserts a private signal before - downstream lowering. Explicit signals are single-shot: callers issuing - multiple allreduces must provide a fresh signal for each call. ``op`` + downstream lowering. The signal is self-clearing: the lowering restores + its cells to zero after each call, so one buffer can be reused across + back-to-back calls (and, on the InCore rail, inside for/while loops). ``op`` (:class:`ir.ReduceOp`) selects the reduction operator, defaults to ``ReduceOp.Sum``, and is packed as an ``int`` attr. ``mode`` selects the lowering algorithm: ``"mesh"`` (direct exchange, O(P) windows) or @@ -427,8 +428,8 @@ def all_to_all_v( ``send_counts`` is read at runtime, so the counts may be data-dependent; each count is clamped to the per-peer capacity ``MAX_RECV = - target.shape[0] // NR``. The barrier signal is single-use and must not be - reused inside a ``for``/``while`` loop. + target.shape[0] // NR``. The barrier signal is self-clearing (restored to + zero after each call) and safe to reuse inside a ``for``/``while`` loop. """ actual_span = _get_span_or_capture(span, frame_offset=1) _args: list[Expr] = [input, target, signal, send_counts, recv_counts] diff --git a/python/pypto/language/distributed/op/tensor_ops.py b/python/pypto/language/distributed/op/tensor_ops.py index 3a2a6346f4..36eade1602 100644 --- a/python/pypto/language/distributed/op/tensor_ops.py +++ b/python/pypto/language/distributed/op/tensor_ops.py @@ -941,8 +941,8 @@ def all_to_all_v( side (published value is the clamped logical count, not the physical transfer size). - The barrier ``signal`` is single-use (same Set(1)/wait≥1 protocol as - allreduce) and must not be reused inside a ``for``/``while`` loop. + The barrier ``signal`` is self-clearing (restored to zero after each call) + and safe to reuse inside a ``for``/``while`` loop. Args: input: Flat 2D Tensor or DistributedTensor [NR*MAX_RECV, SIZE] with diff --git a/python/pypto/runtime/builtins/collectives/all_to_all/templates/kernel.cpp.in b/python/pypto/runtime/builtins/collectives/all_to_all/templates/kernel.cpp.in index 7db75bf897..519e6b593e 100644 --- a/python/pypto/runtime/builtins/collectives/all_to_all/templates/kernel.cpp.in +++ b/python/pypto/runtime/builtins/collectives/all_to_all/templates/kernel.cpp.in @@ -118,7 +118,7 @@ extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ in } // ================================================================== - // Phase 2: barrier — notify peers and wait for all (NotifyOp::Set). + // Phase 2: barrier — notify peers and wait for all (NotifyOp::AtomicAdd). // ================================================================== pipe_barrier(PIPE_ALL); dsb(DSB_DDR); @@ -126,7 +126,7 @@ extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ in if (peer == my_rank) continue; __gm__ int32_t *remote_signal = CommRemotePtr(comm_ctx, signal_base + my_rank, peer); pto::comm::Signal sig(remote_signal); - pto::comm::TNOTIFY(sig, static_cast(1), pto::comm::NotifyOp::Set); + pto::comm::TNOTIFY(sig, static_cast(1), pto::comm::NotifyOp::AtomicAdd); } for (int peer = 0; peer < nranks; ++peer) { if (peer == my_rank) continue; @@ -134,5 +134,14 @@ extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ in pto::comm::TWAIT(sig, static_cast(1), pto::comm::WaitCmp::GE); } + // Self-clearing epilogue: each peer's AtomicAdd(+1) left its cell satisfied; + // a local self-notify (TNOTIFY on a local address is the same st_atomic the + // remote path uses) restores every cell to 0 so the signal is reusable + // across calls. + for (int peer = 0; peer < nranks; ++peer) { + if (peer == my_rank) continue; + pto::comm::Signal self_sig(signal_base + peer); + pto::comm::TNOTIFY(self_sig, static_cast(-1), pto::comm::NotifyOp::AtomicAdd); + } pipe_barrier(PIPE_ALL); } diff --git a/python/pypto/runtime/builtins/collectives/all_to_all_v/templates/kernel.cpp.in b/python/pypto/runtime/builtins/collectives/all_to_all_v/templates/kernel.cpp.in index 5096b0eb45..caefcb2f20 100644 --- a/python/pypto/runtime/builtins/collectives/all_to_all_v/templates/kernel.cpp.in +++ b/python/pypto/runtime/builtins/collectives/all_to_all_v/templates/kernel.cpp.in @@ -162,7 +162,7 @@ extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ in } // ================================================================== - // Phase 2: barrier — notify peers and wait for all (NotifyOp::Set). + // Phase 2: barrier — notify peers and wait for all (NotifyOp::AtomicAdd). // ================================================================== pipe_barrier(PIPE_ALL); dsb(DSB_DDR); @@ -170,7 +170,7 @@ extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ in if (peer == my_rank) continue; __gm__ int32_t *remote_signal = CommRemotePtr(comm_ctx, signal_base + my_rank, peer); pto::comm::Signal sig(remote_signal); - pto::comm::TNOTIFY(sig, static_cast(1), pto::comm::NotifyOp::Set); + pto::comm::TNOTIFY(sig, static_cast(1), pto::comm::NotifyOp::AtomicAdd); } for (int peer = 0; peer < nranks; ++peer) { if (peer == my_rank) continue; @@ -178,5 +178,14 @@ extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ in pto::comm::TWAIT(sig, static_cast(1), pto::comm::WaitCmp::GE); } + // Self-clearing epilogue: each peer's AtomicAdd(+1) left its cell satisfied; + // a local self-notify (TNOTIFY on a local address is the same st_atomic the + // remote path uses) restores every cell to 0 so the signal is reusable + // across calls. + for (int peer = 0; peer < nranks; ++peer) { + if (peer == my_rank) continue; + pto::comm::Signal self_sig(signal_base + peer); + pto::comm::TNOTIFY(self_sig, static_cast(-1), pto::comm::NotifyOp::AtomicAdd); + } pipe_barrier(PIPE_ALL); } diff --git a/python/pypto/runtime/builtins/collectives/allgather/templates/kernel.cpp.in b/python/pypto/runtime/builtins/collectives/allgather/templates/kernel.cpp.in index d32afdef28..f80d933e60 100644 --- a/python/pypto/runtime/builtins/collectives/allgather/templates/kernel.cpp.in +++ b/python/pypto/runtime/builtins/collectives/allgather/templates/kernel.cpp.in @@ -121,7 +121,7 @@ extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ in } // ================================================================== - // Phase 2: barrier — notify peers and wait for all (NotifyOp::Set). + // Phase 2: barrier — notify peers and wait for all (NotifyOp::AtomicAdd). // ================================================================== pipe_barrier(PIPE_ALL); dsb(DSB_DDR); @@ -129,7 +129,7 @@ extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ in if (peer == my_rank) continue; __gm__ int32_t *remote_signal = CommRemotePtr(comm_ctx, signal_base + my_rank, peer); pto::comm::Signal sig(remote_signal); - pto::comm::TNOTIFY(sig, static_cast(1), pto::comm::NotifyOp::Set); + pto::comm::TNOTIFY(sig, static_cast(1), pto::comm::NotifyOp::AtomicAdd); } for (int peer = 0; peer < nranks; ++peer) { if (peer == my_rank) continue; @@ -137,5 +137,14 @@ extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ in pto::comm::TWAIT(sig, static_cast(1), pto::comm::WaitCmp::GE); } + // Self-clearing epilogue: each peer's AtomicAdd(+1) left its cell satisfied; + // a local self-notify (TNOTIFY on a local address is the same st_atomic the + // remote path uses) restores every cell to 0 so the signal is reusable + // across calls. + for (int peer = 0; peer < nranks; ++peer) { + if (peer == my_rank) continue; + pto::comm::Signal self_sig(signal_base + peer); + pto::comm::TNOTIFY(self_sig, static_cast(-1), pto::comm::NotifyOp::AtomicAdd); + } pipe_barrier(PIPE_ALL); } diff --git a/python/pypto/runtime/builtins/collectives/allreduce/templates/kernel.cpp.in b/python/pypto/runtime/builtins/collectives/allreduce/templates/kernel.cpp.in index 307b8d7680..26db55c0d3 100644 --- a/python/pypto/runtime/builtins/collectives/allreduce/templates/kernel.cpp.in +++ b/python/pypto/runtime/builtins/collectives/allreduce/templates/kernel.cpp.in @@ -173,5 +173,20 @@ extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ in wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); } + // Self-clearing epilogue: restore every peer's signal cell in MY block's lane + // (in my memory) to 0 so the signal is reusable across calls (the InCore + // credit-barrier protocol, see lower_composite_ops_pass.cpp). Each peer + // notified +1 into the lane this block waited on once per barrier this call + // issued — 1 ready barrier plus one per UB chunk this block processed. + // read_done_expected ends at num_chunks + 2, so the barrier count is + // read_done_expected - 1; the epilogue is a local hardware atomic (TNOTIFY on + // a local address is the same st_atomic the remote path uses). + const int32_t barrier_count = read_done_expected - 1; + for (int peer = 0; peer < nranks; ++peer) { + if (peer == my_rank) continue; + pto::comm::Signal self_sig(signal_base + peer * signal_stride + block_idx); + pto::comm::TNOTIFY(self_sig, -barrier_count, pto::comm::NotifyOp::AtomicAdd); + } + pipe_barrier(PIPE_ALL); } 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..e78d9e4cb9 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 @@ -268,5 +268,19 @@ extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ in // final TLOAD. RoundBarrier(comm_ctx, signal_base + round * signal_cols, my_rank, nranks); + // Self-clearing epilogue: restore every used barrier row to 0 so the signal is + // reusable across calls (the InCore credit-barrier protocol). Each row hosted + // exactly one RoundBarrier, whose notify gave every peer's cell in my memory a + // single +1; undo it with a local self-notify (TNOTIFY on a local address is the + // same st_atomic the remote path uses) so the next call's Ge(1) waits actually + // synchronize instead of passing on stale credits. + for (int r = 0; r < expected_rounds; ++r) { + for (int peer = 0; peer < nranks; ++peer) { + if (peer == my_rank) continue; + pto::comm::Signal self_sig(signal_base + static_cast(r) * signal_cols + peer); + pto::comm::TNOTIFY(self_sig, static_cast(-1), pto::comm::NotifyOp::AtomicAdd); + } + } + pipe_barrier(PIPE_ALL); } diff --git a/python/pypto/runtime/builtins/collectives/barrier/templates/kernel.cpp.in b/python/pypto/runtime/builtins/collectives/barrier/templates/kernel.cpp.in index af341d078e..bfb4a2779d 100644 --- a/python/pypto/runtime/builtins/collectives/barrier/templates/kernel.cpp.in +++ b/python/pypto/runtime/builtins/collectives/barrier/templates/kernel.cpp.in @@ -55,21 +55,37 @@ extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ in int my_rank = static_cast(comm_ctx->rankId); - // NOTE: This barrier is single-use per signal buffer — NotifyOp::Set writes - // a constant 1 and TWAIT checks >= 1. After the first barrier, signal slots - // stay satisfied. The caller must provide a fresh (zero-initialised) signal - // buffer for each barrier call (pld.alloc_window_buffer satisfies this). + // Self-clearing credit barrier: NotifyOp::AtomicAdd writes a constant 1 into + // each peer's cell in my memory and TWAIT checks >= 1. The epilogue below + // restores every cell to 0, so one signal buffer is reusable across calls + // (the InCore credit-barrier protocol). for (int peer = 0; peer < nranks; ++peer) { if (peer == my_rank) continue; __gm__ int32_t *remote_signal = CommRemotePtr(comm_ctx, signal_base + my_rank, peer); pto::comm::Signal sig(remote_signal); - pto::comm::TNOTIFY(sig, static_cast(1), pto::comm::NotifyOp::Set); + pto::comm::TNOTIFY(sig, static_cast(1), pto::comm::NotifyOp::AtomicAdd); } 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(1), pto::comm::WaitCmp::GE); } + + // Self-clearing epilogue: each peer's AtomicAdd(+1) left its cell satisfied; + // a local self-notify restores every cell to 0 so the signal is reusable + // across calls. + // + // Use NotifyOp::Set(0) instead of AtomicAdd(-1): on non-coherent NPU silicon + // a new AIV task dispatch may read a stale cached value for the signal cell; + // AtomicAdd reads-modifies-writes that stale value, while Set unconditionally + // writes 0. Set is safe here because the TWAIT above guarantees all peers' + // AtomicAdd credits have already landed — there is no in-flight write that + // Set could clobber. + for (int peer = 0; peer < nranks; ++peer) { + if (peer == my_rank) continue; + pto::comm::Signal self_sig(signal_base + peer); + pto::comm::TNOTIFY(self_sig, static_cast(0), pto::comm::NotifyOp::Set); + } pipe_barrier(PIPE_ALL); } diff --git a/python/pypto/runtime/builtins/collectives/broadcast/templates/kernel.cpp.in b/python/pypto/runtime/builtins/collectives/broadcast/templates/kernel.cpp.in index 227746b684..50ac027485 100644 --- a/python/pypto/runtime/builtins/collectives/broadcast/templates/kernel.cpp.in +++ b/python/pypto/runtime/builtins/collectives/broadcast/templates/kernel.cpp.in @@ -76,10 +76,19 @@ extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ in TASSIGN(recv_tile, 0x10000); // Signal scheme: each slot is single-writer (rank j writes into my_rank's - // signal[j]). TNOTIFY adds a constant 1 per tile; the TWAIT threshold grows - // 1, 2, 3, … because TWAIT(GE) is non-consuming — the slot retains its - // accumulated value across tiles. This matches the reduce_scatter signal - // protocol (constant notify + incrementing wait threshold). + // signal[j]). Every rank remote-notifies every peer (AtomicAdd(+1) to + // signal[my_rank] on peer's device) and locally waits for every peer's + // signal cell (TWAIT Ge(N) on signal[peer] in its own memory). This + // symmetric write-remote-read-local pattern is the only one that has been + // verified reliable on non-coherent NPU silicon across all collectives; an + // asymmetric read-complete round (non-root remote AtomicAdd → root local + // TWAIT alone) leaves the root polling a locally-cached zero while the + // remote update bypassed the cache, deadlocking the kernel. + // + // The root self-copies its tile before the remote notify so its target + // buffer is up-to-date when non-roots remote-read it, and the + // pipe_barrier(PIPE_ALL) between tiles ensures the store is globally + // visible before the next tile's notify round. int32_t expected = 1; for (int64_t base = 0; base < numel; base += kTileCount) { int64_t chunk64 = numel - base; @@ -126,5 +135,14 @@ extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ in } } + // Self-clearing credit reset — all local, symmetric (every cell received + // exactly tiles credits this call — one per tile iteration). Local + // AtomicAdd(-N) never mixes Set with the AtomicAdd body. + const int32_t tiles = static_cast(expected - 1); + for (int peer = 0; peer < nranks; ++peer) { + if (peer == my_rank) continue; + pto::comm::Signal self_sig(signal_base + peer); + pto::comm::TNOTIFY(self_sig, -tiles, pto::comm::NotifyOp::AtomicAdd); + } pipe_barrier(PIPE_ALL); } diff --git a/python/pypto/runtime/builtins/collectives/reduce_scatter/templates/kernel.cpp.in b/python/pypto/runtime/builtins/collectives/reduce_scatter/templates/kernel.cpp.in index c61106bcf9..b967b2ea93 100644 --- a/python/pypto/runtime/builtins/collectives/reduce_scatter/templates/kernel.cpp.in +++ b/python/pypto/runtime/builtins/collectives/reduce_scatter/templates/kernel.cpp.in @@ -150,5 +150,25 @@ extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ in wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0); } + // Self-clearing epilogue: restore every peer's signal cell (in my memory) to + // 0 so the signal is reusable across calls (the InCore credit-barrier + // protocol, see lower_composite_ops_pass.cpp). Each peer notified +1 into its + // own cell in my memory once per barrier this call issued — 1 ready barrier + // plus one per UB chunk. read_done_expected ends at num_chunks + 2, so the + // barrier count is read_done_expected - 1. + // + // Use NotifyOp::Set(0) instead of AtomicAdd(-N): on non-coherent NPU silicon + // a new AIV task dispatch may read a stale cached value for the signal cell; + // AtomicAdd reads-modifies-writes that stale value, while Set unconditionally + // writes 0. Set is safe here because the per-tile pipe_barrier(PIPE_ALL) + // above guarantees all peers' AtomicAdd credits have already landed — there + // is no in-flight write that Set could clobber. Verified reliable in the + // single-core allreduce and barrier reuse ST suites whose kernels share the + // same protocol. + for (int peer = 0; peer < nranks; ++peer) { + if (peer == my_rank) continue; + pto::comm::Signal self_sig(signal_base + peer); + pto::comm::TNOTIFY(self_sig, static_cast(0), pto::comm::NotifyOp::Set); + } pipe_barrier(PIPE_ALL); } diff --git a/src/ir/op/distributed/allreduce.cpp b/src/ir/op/distributed/allreduce.cpp index 9f681a7408..fde18f6b5c 100644 --- a/src/ir/op/distributed/allreduce.cpp +++ b/src/ir/op/distributed/allreduce.cpp @@ -21,8 +21,10 @@ * UB-sized reduce/barrier/store chunks in * ``src/ir/transforms/lower_composite_ops_pass.cpp``; host-level allreduce is * lowered later by ``LowerHostTensorCollectives``. - * Explicit signal buffers are single-shot: callers issuing multiple allreduces - * must provide a fresh signal for each call. + * The ``signal`` buffer is self-clearing: each lowering restores the barrier + * cells to zero before the call returns (the InCore credit-barrier protocol; + * the host builtins carry the matching epilogue), so one signal is reusable + * across back-to-back calls and, on the InCore rail, inside for/while loops. * * IR signature: * diff --git a/src/ir/op/distributed/collective.cpp b/src/ir/op/distributed/collective.cpp index 98198f8f61..f467a3dd9c 100644 --- a/src/ir/op/distributed/collective.cpp +++ b/src/ir/op/distributed/collective.cpp @@ -687,8 +687,8 @@ REGISTER_OP("pld.tensor.all_to_all_v") .add_argument("target", "Window-bound DistributedTensor [NR*MAX_RECV, SIZE] — staging area for exchange (InOut)") .add_argument("signal", - "Window-bound INT32 DistributedTensor [NR, 1] used as a single-use cross-rank " - "barrier (InOut); not reusable inside for/while loops") + "Window-bound INT32 DistributedTensor [NR, 1] used as a self-clearing cross-rank " + "barrier (InOut); reusable across calls and inside for/while loops") .add_argument("send_counts", "INT32 Tensor [NR] or [NR, 1] — rows to send to each destination, read at " "runtime and clamped to MAX_RECV (Input)") diff --git a/src/ir/transforms/synthesize_allreduce_signals_pass.cpp b/src/ir/transforms/synthesize_allreduce_signals_pass.cpp index 808282faf4..c997c13129 100644 --- a/src/ir/transforms/synthesize_allreduce_signals_pass.cpp +++ b/src/ir/transforms/synthesize_allreduce_signals_pass.cpp @@ -188,8 +188,10 @@ class AllReduceSignalSynthesizer : public IRMutator { << "pld.tensor.allreduce expects target[, signal], got " << call->args_.size() << " positional arguments"; CHECK_SPAN(repeating_scope_depth_ == 0, call->span_) - << "pld.tensor.allreduce is not supported inside a for/while loop. " - "The signal protocol is single-use and cannot reuse a signal across dynamic invocations."; + << "pld.tensor.allreduce is not supported inside a for/while loop on the HOST rail: " + "the synthesized signal binding cannot be allocated per dynamic iteration. " + "Hoist the call out of the loop (InCore composites are loop-safe via the " + "self-clearing credit-barrier protocol)."; } struct SignalNames { diff --git a/tests/st/distributed/test_l3_host_tensor_all_to_all.py b/tests/st/distributed/test_l3_host_tensor_all_to_all.py index aee9eaeeac..bf0fb1df2e 100644 --- a/tests/st/distributed/test_l3_host_tensor_all_to_all.py +++ b/tests/st/distributed/test_l3_host_tensor_all_to_all.py @@ -48,20 +48,18 @@ def _expected_all_to_all(inputs: torch.Tensor) -> torch.Tensor: - """Golden: output[rank, src, j] = src * 1000 + rank * 100 + j.""" - nranks = inputs.shape[0] - rank_idx = torch.arange(nranks, dtype=torch.float32).view(-1, 1, 1) - src_idx = torch.arange(nranks, dtype=torch.float32).view(1, -1, 1) - j = torch.arange(SIZE, dtype=torch.float32).view(1, 1, -1) - return src_idx * 1000 + rank_idx * 100 + j + """Golden: output[rank, src, j] = inputs[src, rank, j] (rank src's chunk + destined for rank ``rank`` lands in rank's slot ``src``). Input-dependent so + distinct per-round data propagates to the expected output.""" + return inputs.permute(1, 0, 2) -def _make_rank_inputs(n_ranks: int) -> torch.Tensor: - """Each rank r fills input[r, d, j] = r * 1000 + d * 100 + j.""" +def _make_rank_inputs(n_ranks: int, round_offset: float = 0.0) -> torch.Tensor: + """Each rank r fills input[r, d, j] = r * 1000 + d * 100 + j (+ round_offset).""" r = torch.arange(n_ranks, dtype=torch.float32).view(-1, 1, 1) d = torch.arange(n_ranks, dtype=torch.float32).view(1, -1, 1) j = torch.arange(SIZE, dtype=torch.float32).view(1, 1, -1) - return r * 1000 + d * 100 + j + return round_offset + r * 1000 + d * 100 + j @pl.program @@ -130,6 +128,91 @@ def host_orch( return outputs +def _build_host_all_to_all_signal_reuse_program(): + """Host all_to_all reusing ONE signal buffer across 2 back-to-back calls. + + The self-clearing epilogue restores the AtomicAdd(+1) cells to 0 after each + call; without it the second call's Ge(1) wait passes on the stale + satisfied cell. + """ + ROUNDS = 2 + + @pl.program + class HostTensorAllToAllSignalReuse: + @pl.function(type=pl.FunctionType.InCore) + def stage_step( + self, + inp: pl.Tensor[[NR, SIZE], pl.FP32], + stage: pl.Out[pld.DistributedTensor[[NR, SIZE], pl.FP32]], + ): + for dest in pl.range(NR): + chunk = pl.load(inp, [dest, 0], [1, SIZE]) + stage = pl.store(chunk, [dest, 0], stage) + + @pl.function(type=pl.FunctionType.Orchestration) + def stage_orch( + self, + inp: pl.Tensor[[NR, SIZE], pl.FP32], + stage: pl.Out[pld.DistributedTensor[[NR, SIZE], pl.FP32]], + ): + self.stage_step(inp, stage) + + @pl.function(type=pl.FunctionType.InCore) + def consume_step( + self, + data: pld.DistributedTensor[[NR, SIZE], pl.FP32], + out: pl.Out[pl.Tensor[[NR, SIZE], pl.FP32]], + ) -> pl.Tensor[[NR, SIZE], pl.FP32]: + for src in pl.range(NR): + row = pl.load(data, [src, 0], [1, SIZE]) + out = pl.store(row, [src, 0], out) + return out + + @pl.function(type=pl.FunctionType.Orchestration) + def consume_orch( + self, + data: pld.DistributedTensor[[NR, SIZE], pl.FP32], + out: pl.Out[pl.Tensor[[NR, SIZE], pl.FP32]], + ) -> pl.Tensor[[NR, SIZE], pl.FP32]: + return self.consume_step(data, out) + + @pl.function(level=pl.Level.HOST, role=pl.Role.Orchestrator) + def host_orch( + self, + inputs: pl.Tensor[[ROUNDS, NR, NR, SIZE], pl.FP32], + outputs: pl.Out[pl.Tensor[[ROUNDS, NR, NR, SIZE], pl.FP32]], + ) -> pl.Tensor[[ROUNDS, NR, NR, SIZE], pl.FP32]: + stage_buf_1 = pld.alloc_window_buffer(pld.world_size() * SIZE * pl.FP32.get_byte()) + stage_buf_2 = pld.alloc_window_buffer(pld.world_size() * SIZE * pl.FP32.get_byte()) + data_buf_1 = pld.alloc_window_buffer(pld.world_size() * SIZE * pl.FP32.get_byte()) + data_buf_2 = pld.alloc_window_buffer(pld.world_size() * SIZE * pl.FP32.get_byte()) + signal_buf = pld.alloc_window_buffer(pld.world_size() * pl.INT32.get_byte()) + signal = pld.window(signal_buf, [pld.world_size()], dtype=pl.INT32) + + # Round 1 — distinct stage/data windows per round, ONE shared signal. + for r in pl.range(pld.world_size()): + stage = pld.window(stage_buf_1, [pld.world_size(), SIZE], dtype=pl.FP32) + self.stage_orch(inputs[0, r], stage, device=r) + stage = pld.window(stage_buf_1, [pld.world_size(), SIZE], dtype=pl.FP32) + data = pld.window(data_buf_1, [pld.world_size(), SIZE], dtype=pl.FP32) + data = pld.tensor.all_to_all(stage, data, signal) + for r in pl.range(pld.world_size()): + self.consume_orch(data, outputs[0, r], device=r) + + # Round 2 — reuse the same signal. + for r in pl.range(pld.world_size()): + stage = pld.window(stage_buf_2, [pld.world_size(), SIZE], dtype=pl.FP32) + self.stage_orch(inputs[1, r], stage, device=r) + stage = pld.window(stage_buf_2, [pld.world_size(), SIZE], dtype=pl.FP32) + data = pld.window(data_buf_2, [pld.world_size(), SIZE], dtype=pl.FP32) + data = pld.tensor.all_to_all(stage, data, signal) + for r in pl.range(pld.world_size()): + self.consume_orch(data, outputs[1, r], device=r) + return outputs + + return HostTensorAllToAllSignalReuse + + class TestL3HostTensorAllToAll: """L3 distributed runtime: HOST-level all-to-all via builtin dispatch.""" @@ -162,6 +245,37 @@ def test_host_tensor_all_to_all(self, test_config, device_ids, n_ranks): f"host all-to-all P={n_ranks} mismatch: max diff = {(outputs - expected).abs().max().item()}" ) + @pytest.mark.parametrize("n_ranks", [2, 4]) + def test_host_tensor_all_to_all_signal_reuse(self, test_config, device_ids, n_ranks): + """Reuse ONE signal buffer across 2 back-to-back all_to_all calls.""" + if len(device_ids) < n_ranks: + pytest.skip(f"host all_to_all P={n_ranks} needs {n_ranks} devices, got {device_ids}") + + rounds = 2 + compiled = ir.compile( + _build_host_all_to_all_signal_reuse_program(), + platform=test_config.platform, + distributed_config=DistributedConfig( + device_ids=device_ids[:n_ranks], + num_sub_workers=0, + ), + ) + variant_dir = compiled.output_dir / "next_levels" / "builtin.tensor.all_to_all__fp32" + assert variant_dir.is_dir() + + # Each round carries a distinct offset so a stale earlier-round result + # (a missed epilogue reset) cannot match the round's golden. + inputs = torch.stack([_make_rank_inputs(n_ranks, round_offset=rd * 10000.0) for rd in range(rounds)]) + outputs = torch.zeros((rounds, n_ranks, n_ranks, SIZE), dtype=torch.float32) + compiled(inputs, outputs) + + for rd in range(rounds): + expected = _expected_all_to_all(inputs[rd]) + assert torch.allclose(outputs[rd], expected), ( + f"host all_to_all signal-reuse round {rd} P={n_ranks} mismatch: " + f"max diff = {(outputs[rd] - expected).abs().max().item()}" + ) + if __name__ == "__main__": pytest.main([__file__, "-v", *sys.argv[1:]]) diff --git a/tests/st/distributed/test_l3_host_tensor_allgather.py b/tests/st/distributed/test_l3_host_tensor_allgather.py index 712bc5675f..f1770541b1 100644 --- a/tests/st/distributed/test_l3_host_tensor_allgather.py +++ b/tests/st/distributed/test_l3_host_tensor_allgather.py @@ -51,9 +51,12 @@ def _expected_allgather(inputs: torch.Tensor, n_ranks: int) -> torch.Tensor: return torch.stack([gathered] * n_ranks).unsqueeze(1) -def _make_rank_inputs(n_ranks: int) -> torch.Tensor: +def _make_rank_inputs(n_ranks: int, round_offset: float = 0.0) -> torch.Tensor: + """Build distinct per-rank rows; ``round_offset`` distinguishes reuse rounds.""" rows = [ - torch.arange(r * 100.0, r * 100.0 + SIZE, dtype=torch.float32).reshape(1, SIZE) + torch.arange(r * 100.0 + round_offset, r * 100.0 + round_offset + SIZE, dtype=torch.float32).reshape( + 1, SIZE + ) for r in range(n_ranks) ] return torch.stack(rows) @@ -135,6 +138,96 @@ def host_orch( return outputs +def _build_host_allgather_signal_reuse_program(): + """Host allgather reusing ONE signal buffer across 2 back-to-back calls. + + The self-clearing epilogue restores the AtomicAdd(+1) cells to 0 after each + call; without it the second call's Ge(1) wait passes on the stale + satisfied cell. + """ + ROUNDS = 2 + + @pl.program + class HostTensorAllGatherSignalReuse: + @pl.function(type=pl.FunctionType.InCore) + def publish_step( + self, + inp: pl.Tensor[[1, SIZE], pl.FP32], + stage: pl.Out[pld.DistributedTensor[[1, SIZE], pl.FP32]], + my_rank: pl.Scalar[pl.INT32], + nranks: pl.Scalar[pl.INT32], + ): + chunk = pl.load(inp, [0, 0], [1, SIZE]) + stage = pl.store(chunk, [0, 0], stage) + + @pl.function(type=pl.FunctionType.Orchestration) + def publish_orch( + self, + inp: pl.Tensor[[1, SIZE], pl.FP32], + stage: pl.Out[pld.DistributedTensor[[1, SIZE], pl.FP32]], + my_rank: pl.Scalar[pl.INT32], + nranks: pl.Scalar[pl.INT32], + ): + self.publish_step(inp, stage, my_rank, nranks) + + @pl.function(type=pl.FunctionType.InCore) + def consume_step( + self, + data: pld.DistributedTensor[[NR, SIZE], pl.FP32], + out: pl.Out[pl.Tensor[[1, NR, SIZE], pl.FP32]], + nranks: pl.Scalar[pl.INT32], + ) -> pl.Tensor[[1, NR, SIZE], pl.FP32]: + for j in pl.range(nranks): + row = pl.load(data, [j, 0], [1, SIZE]) + out = pl.store(row, [0, j, 0], out) + return out + + @pl.function(type=pl.FunctionType.Orchestration) + def consume_orch( + self, + data: pld.DistributedTensor[[NR, SIZE], pl.FP32], + out: pl.Out[pl.Tensor[[1, NR, SIZE], pl.FP32]], + nranks: pl.Scalar[pl.INT32], + ) -> pl.Tensor[[1, NR, SIZE], pl.FP32]: + return self.consume_step(data, out, nranks) + + @pl.function(level=pl.Level.HOST, role=pl.Role.Orchestrator) + def host_orch( + self, + inputs: pl.Tensor[[ROUNDS, NR, 1, SIZE], pl.FP32], + outputs: pl.Out[pl.Tensor[[ROUNDS, NR, 1, NR, SIZE], pl.FP32]], + ) -> pl.Tensor[[ROUNDS, NR, 1, NR, SIZE], pl.FP32]: + stage_buf_1 = pld.alloc_window_buffer(SIZE * pl.FP32.get_byte()) + stage_buf_2 = pld.alloc_window_buffer(SIZE * pl.FP32.get_byte()) + data_buf_1 = pld.alloc_window_buffer(pld.world_size() * SIZE * pl.FP32.get_byte()) + data_buf_2 = pld.alloc_window_buffer(pld.world_size() * SIZE * pl.FP32.get_byte()) + signal_buf = pld.alloc_window_buffer(pld.world_size() * pl.INT32.get_byte()) + signal = pld.window(signal_buf, [pld.world_size()], dtype=pl.INT32) + + # Round 1 — distinct stage/data windows per round, ONE shared signal. + for r in pl.range(pld.world_size()): + stage = pld.window(stage_buf_1, [1, SIZE], dtype=pl.FP32) + self.publish_orch(inputs[0, r], stage, r, pld.world_size(), device=r) + stage = pld.window(stage_buf_1, [1, SIZE], dtype=pl.FP32) + data = pld.window(data_buf_1, [pld.world_size(), SIZE], dtype=pl.FP32) + data = pld.tensor.allgather(stage, data, signal) + for r in pl.range(pld.world_size()): + self.consume_orch(data, outputs[0, r], pld.world_size(), device=r) + + # Round 2 — reuse the same signal. + for r in pl.range(pld.world_size()): + stage = pld.window(stage_buf_2, [1, SIZE], dtype=pl.FP32) + self.publish_orch(inputs[1, r], stage, r, pld.world_size(), device=r) + stage = pld.window(stage_buf_2, [1, SIZE], dtype=pl.FP32) + data = pld.window(data_buf_2, [pld.world_size(), SIZE], dtype=pl.FP32) + data = pld.tensor.allgather(stage, data, signal) + for r in pl.range(pld.world_size()): + self.consume_orch(data, outputs[1, r], pld.world_size(), device=r) + return outputs + + return HostTensorAllGatherSignalReuse + + class TestL3HostTensorAllGather: @pytest.mark.parametrize("n_ranks", [2, 4]) def test_host_tensor_allgather(self, test_config, device_ids, n_ranks): @@ -166,6 +259,37 @@ def test_host_tensor_allgather(self, test_config, device_ids, n_ranks): f"host allgather P={n_ranks} mismatch: max diff = {(outputs - expected).abs().max().item()}" ) + @pytest.mark.parametrize("n_ranks", [2, 4]) + def test_host_tensor_allgather_signal_reuse(self, test_config, device_ids, n_ranks): + """Reuse ONE signal buffer across 2 back-to-back allgather calls.""" + if len(device_ids) < n_ranks: + pytest.skip(f"host allgather P={n_ranks} needs {n_ranks} devices, got {device_ids}") + + rounds = 2 + compiled = ir.compile( + _build_host_allgather_signal_reuse_program(), + platform=test_config.platform, + distributed_config=DistributedConfig( + device_ids=device_ids[:n_ranks], + num_sub_workers=0, + ), + ) + variant_dir = compiled.output_dir / "next_levels" / "builtin.tensor.allgather__fp32" + assert variant_dir.is_dir() + + # Each round carries a distinct offset so a stale earlier-round result + # (a missed epilogue reset) cannot match the round's golden. + inputs = torch.stack([_make_rank_inputs(n_ranks, round_offset=rd * 10000.0) for rd in range(rounds)]) + outputs = torch.zeros((rounds, n_ranks, 1, n_ranks, SIZE), dtype=torch.float32) + compiled(inputs, outputs) + + for rd in range(rounds): + expected = _expected_allgather(inputs[rd], n_ranks) + assert torch.allclose(outputs[rd], expected), ( + f"host allgather signal-reuse round {rd} P={n_ranks} mismatch: " + f"max diff = {(outputs[rd] - expected).abs().max().item()}" + ) + if __name__ == "__main__": sys.exit(pytest.main([__file__, "-v", "-s"] + sys.argv[1:])) diff --git a/tests/st/distributed/test_l3_host_tensor_allreduce.py b/tests/st/distributed/test_l3_host_tensor_allreduce.py index 4be2e4abf0..32265c2d32 100644 --- a/tests/st/distributed/test_l3_host_tensor_allreduce.py +++ b/tests/st/distributed/test_l3_host_tensor_allreduce.py @@ -42,7 +42,9 @@ def _make_rank_inputs( *, dtype: torch.dtype = torch.float32, op_name: str = "sum", + round_offset: float = 0.0, ) -> torch.Tensor: + """Build per-rank distinct inputs; ``round_offset`` distinguishes reuse rounds.""" if op_name == "prod": rows = [ (1.0 + r * 0.125 + torch.arange(size, dtype=torch.float32).remainder(5) * 0.0625).reshape(1, size) @@ -50,7 +52,9 @@ def _make_rank_inputs( ] else: rows = [ - torch.arange(r * 100.0, r * 100.0 + size, dtype=torch.float32).reshape(1, size) + torch.arange( + r * 100.0 + round_offset, r * 100.0 + round_offset + size, dtype=torch.float32 + ).reshape(1, size) for r in range(n_ranks) ] return torch.stack(rows).to(dtype) @@ -269,6 +273,97 @@ def host_orch( return HostTensorAllReduceArbitraryLength +def _build_host_allreduce_signal_reuse(): + """Host allreduce reusing ONE signal buffer across 3 back-to-back calls. + + The program unrolls exactly three rounds (the HOST rail rejects allreduce + under a dynamic-trip-count loop), each reusing the shared ``signal`` — the + self-clearing credit-barrier epilogue's target case. Before the epilogue a + reused signal carried stale credits and the second call's Ge(1) wait passed + spuriously (NPU-visible; the sim executor is sequentially consistent, so the + real proof is the NPU developer gate). + """ + + ROUNDS = 3 + + @pl.program + class HostTensorAllReduceSignalReuse: + @pl.function(type=pl.FunctionType.InCore) + def publish_step( + self, + inp: pl.Tensor[[1, SIZE], pl.FP32], + data: pl.InOut[pld.DistributedTensor[[1, SIZE], pl.FP32]], + ) -> pld.DistributedTensor[[1, SIZE], pl.FP32]: + local = pl.load(inp, [0, 0], [1, SIZE]) + return pl.store(local, [0, 0], data) + + @pl.function(type=pl.FunctionType.Orchestration) + def publish_orch( + self, + inp: pl.Tensor[[1, SIZE], pl.FP32], + data: pl.InOut[pld.DistributedTensor[[1, SIZE], pl.FP32]], + ) -> pld.DistributedTensor[[1, SIZE], pl.FP32]: + return self.publish_step(inp, data) + + @pl.function(type=pl.FunctionType.InCore) + def consume_step( + self, + data: pld.DistributedTensor[[1, SIZE], pl.FP32], + out: pl.Out[pl.Tensor[[1, SIZE], pl.FP32]], + ) -> pl.Tensor[[1, SIZE], pl.FP32]: + reduced = pl.load(data, [0, 0], [1, SIZE]) + return pl.store(reduced, [0, 0], out) + + @pl.function(type=pl.FunctionType.Orchestration) + def consume_orch( + self, + data: pld.DistributedTensor[[1, SIZE], pl.FP32], + out: pl.Out[pl.Tensor[[1, SIZE], pl.FP32]], + ) -> pl.Tensor[[1, SIZE], pl.FP32]: + return self.consume_step(data, out) + + @pl.function(level=pl.Level.HOST, role=pl.Role.Orchestrator) + def host_orch( + self, + inputs: pl.Tensor[[ROUNDS, NR, 1, SIZE], pl.FP32], + outputs: pl.Out[pl.Tensor[[ROUNDS, NR, 1, SIZE], pl.FP32]], + ) -> pl.Tensor[[ROUNDS, NR, 1, SIZE], pl.FP32]: + data_buf = pld.alloc_window_buffer(SIZE * pl.FP32.get_byte()) + signal_buf = pld.alloc_window_buffer(pld.world_size() * pl.INT32.get_byte()) + signal = pld.window(signal_buf, [pld.world_size()], dtype=pl.INT32) + + # Round 1 — every round below reuses the shared ``signal``. + for r in pl.range(pld.world_size()): + data = pld.window(data_buf, [1, SIZE], dtype=pl.FP32) + self.publish_orch(inputs[0, r], data, device=r) + data = pld.window(data_buf, [1, SIZE], dtype=pl.FP32) + data = pld.tensor.allreduce(data, signal, op=pld.ReduceOp.Sum) + for r in pl.range(pld.world_size()): + self.consume_orch(data, outputs[0, r], device=r) + + # Round 2 — reuse the same signal. + for r in pl.range(pld.world_size()): + data = pld.window(data_buf, [1, SIZE], dtype=pl.FP32) + self.publish_orch(inputs[1, r], data, device=r) + data = pld.window(data_buf, [1, SIZE], dtype=pl.FP32) + data = pld.tensor.allreduce(data, signal, op=pld.ReduceOp.Sum) + for r in pl.range(pld.world_size()): + self.consume_orch(data, outputs[1, r], device=r) + + # Round 3 — reuse the same signal again. + for r in pl.range(pld.world_size()): + data = pld.window(data_buf, [1, SIZE], dtype=pl.FP32) + self.publish_orch(inputs[2, r], data, device=r) + data = pld.window(data_buf, [1, SIZE], dtype=pl.FP32) + data = pld.tensor.allreduce(data, signal, op=pld.ReduceOp.Sum) + for r in pl.range(pld.world_size()): + self.consume_orch(data, outputs[2, r], device=r) + + return outputs + + return HostTensorAllReduceSignalReuse + + class TestL3HostTensorAllReduce: @pytest.mark.parametrize("n_ranks", [2, 4]) def test_host_tensor_allreduce(self, test_config, device_ids, n_ranks): @@ -298,6 +393,42 @@ def test_host_tensor_allreduce(self, test_config, device_ids, n_ranks): f"host allreduce P={n_ranks} mismatch: max diff = {(outputs - expected).abs().max().item()}" ) + @pytest.mark.parametrize("n_ranks", [2, 4]) + def test_host_tensor_allreduce_signal_reuse(self, test_config, device_ids, n_ranks): + """Reuse ONE signal buffer across 3 back-to-back allreduce calls. + + The self-clearing credit-barrier epilogue restores the signal to all-zero + after each call; without it a reused signal carries stale credits and the + next call's Ge(1) wait passes spuriously (NPU-visible; sim is sequential). + """ + if len(device_ids) < n_ranks: + pytest.skip(f"host allreduce P={n_ranks} needs {n_ranks} devices, got {device_ids}") + + rounds = 3 + compiled = ir.compile( + _build_host_allreduce_signal_reuse(), + platform=test_config.platform, + distributed_config=DistributedConfig( + device_ids=device_ids[:n_ranks], + num_sub_workers=0, + ), + ) + variant_dir = compiled.output_dir / "next_levels" / "builtin.tensor.allreduce__sum__fp32" + assert variant_dir.is_dir() + + # Each round carries a distinct offset so a stale round-1 result in a + # later round (a missed epilogue reset) cannot match the round's golden. + inputs = torch.stack([_make_rank_inputs(n_ranks, round_offset=rd * 10000.0) for rd in range(rounds)]) + outputs = torch.zeros_like(inputs) + compiled(inputs, outputs) + + for rd in range(rounds): + expected = _expected_allreduce(inputs[rd]) + assert torch.allclose(outputs[rd], expected), ( + f"host allreduce signal-reuse round {rd} P={n_ranks} mismatch: " + f"max diff = {(outputs[rd] - expected).abs().max().item()}" + ) + def test_host_tensor_allreduce_max(self, test_config, device_ids): """Cover a non-default op in the materialized host builtin.""" n_ranks = 2 diff --git a/tests/st/distributed/test_l3_host_tensor_allreduce_multicore.py b/tests/st/distributed/test_l3_host_tensor_allreduce_multicore.py index d99a10d18d..37c0d032bb 100644 --- a/tests/st/distributed/test_l3_host_tensor_allreduce_multicore.py +++ b/tests/st/distributed/test_l3_host_tensor_allreduce_multicore.py @@ -10,9 +10,15 @@ """L3 ST for HOST AllReduce with multiple synchronized AIV blocks. Each case checks the reduced FP32 output and, on device, waits for every -requested signal lane. The lane waits prove that every requested block -started; output correctness additionally proves that active blocks processed -their block-strided chunks. +requested signal lane to be self-cleared back to zero. Because the builtin +clears each lane it used (ready barrier plus per-chunk credits) before it +returns, a lane at zero proves that the owning block started and completed its +epilogue; output correctness additionally proves that active blocks processed +their block-strided chunks. A signal-reuse variant runs two back-to-back calls +through ONE shared signal and checks output correctness for both rounds — if the +per-lane epilogue fails to self-clear, round 2's ready barrier passes spuriously +on stale credits and the reduction reads peers' data too early, producing wrong +output. """ import sys @@ -25,9 +31,13 @@ from pypto.ir.distributed_compiled_program import DistributedConfig -def _make_rank_inputs(n_ranks: int, size: int) -> torch.Tensor: +def _make_rank_inputs(n_ranks: int, size: int, round_offset: float = 0.0) -> torch.Tensor: rows = [ - torch.arange(r * 100.0, r * 100.0 + size, dtype=torch.float32).reshape(1, size) + torch.arange( + round_offset + r * 100.0, + round_offset + r * 100.0 + size, + dtype=torch.float32, + ).reshape(1, size) for r in range(n_ranks) ] return torch.stack(rows) @@ -91,15 +101,17 @@ def consume_step( for peer in pl.range(nr): if peer != my_rank: for lane in pl.range(cores): - # The builtin's start barrier publishes one signal per - # launched block, including blocks with no data chunk. - # TWAIT performs the cache invalidation required for a + # The builtin self-clears each lane it used (ready + # barrier plus per-chunk credits) back to zero before + # it returns, so a lane at zero proves that peer's + # block started and completed its epilogue. TWAIT + # performs the cache invalidation required for a # reliable device-side observation. pld.system.wait( signal=signal, offsets=[peer, lane], - expected=1, - cmp=pld.WaitCmp.Ge, + expected=0, + cmp=pld.WaitCmp.Eq, ) reduced = pl.load( @@ -166,6 +178,122 @@ def host_orch( return HostTensorAllReduceMulticore +def _build_multicore_allreduce_signal_reuse( + n_ranks: int, + size: int, + core_num: int, + signal_stride: int, +): + """Multicore host allreduce reusing ONE signal across 2 back-to-back calls. + + Both rounds pass the same ``signal`` window to ``pld.tensor.allreduce`` with + ``core_num`` blocks. The per-lane self-clearing epilogue must restore every + lane to zero after round 1, or round 2's ready barrier would pass spuriously + on stale credits and the reduction could read peers' data too early. + + Verified indirectly through output correctness: if the signal is not properly + self-cleared after round 1, round 2's Ge(1) ready barrier passes before + peers have published their data, producing wrong output. No on-device signal + readback — TWAIT(Eq 0) on a non-coherent NPU risks reading stale cached + values from a previous AIV task dispatch, and the sim executor is + sequentially consistent so the definitive proof is the NPU developer gate. + """ + nr = n_ranks + sz = size + cores = core_num + stride = signal_stride + stage_rows = 8 if size == 1 else 1 + stage_cols = 1 if size == 1 else ((size + 7) // 8) * 8 + + @pl.program + class HostTensorAllReduceMulticoreSignalReuse: + @pl.function(type=pl.FunctionType.InCore) + def publish_step( + self, + inp: pl.Tensor[[1, sz], pl.FP32], + data: pl.InOut[pld.DistributedTensor[[1, sz], pl.FP32]], + ) -> pld.DistributedTensor[[1, sz], pl.FP32]: + local = pl.load( + inp, + [0, 0], + [stage_rows, stage_cols], + valid_shape=[1, sz], + ) + return pl.store(local, [0, 0], data) + + @pl.function(type=pl.FunctionType.Orchestration) + def publish_orch( + self, + inp: pl.Tensor[[1, sz], pl.FP32], + data: pl.InOut[pld.DistributedTensor[[1, sz], pl.FP32]], + ) -> pld.DistributedTensor[[1, sz], pl.FP32]: + return self.publish_step(inp, data) + + @pl.function(type=pl.FunctionType.InCore) + def consume_step( + self, + data: pld.DistributedTensor[[1, sz], pl.FP32], + out: pl.Out[pl.Tensor[[1, sz], pl.FP32]], + ) -> pl.Tensor[[1, sz], pl.FP32]: + reduced = pl.load( + data, + [0, 0], + [stage_rows, stage_cols], + valid_shape=[1, sz], + ) + return pl.store(reduced, [0, 0], out) + + @pl.function(type=pl.FunctionType.Orchestration) + def consume_orch( + self, + data: pld.DistributedTensor[[1, sz], pl.FP32], + out: pl.Out[pl.Tensor[[1, sz], pl.FP32]], + ) -> pl.Tensor[[1, sz], pl.FP32]: + return self.consume_step(data, out) + + @pl.function(level=pl.Level.HOST, role=pl.Role.Orchestrator) + def host_orch( + self, + inputs: pl.Tensor[[2, nr, 1, sz], pl.FP32], + outputs: pl.Out[pl.Tensor[[2, nr, 1, sz], pl.FP32]], + ) -> pl.Tensor[[2, nr, 1, sz], pl.FP32]: + data_buf = pld.alloc_window_buffer(sz * pl.FP32.get_byte()) + signal_buf = pld.alloc_window_buffer(pld.world_size() * stride * pl.INT32.get_byte()) + signal = pld.window(signal_buf, [pld.world_size(), stride], dtype=pl.INT32) + + # Round 1. + for rank in pl.range(pld.world_size()): + data = pld.window(data_buf, [1, sz], dtype=pl.FP32) + self.publish_orch(inputs[0, rank], data, device=rank) + data = pld.window(data_buf, [1, sz], dtype=pl.FP32) + data = pld.tensor.allreduce( + data, + signal, + op=pld.ReduceOp.Sum, + core_num=cores, + ) + for rank in pl.range(pld.world_size()): + self.consume_orch(data, outputs[0, rank], device=rank) + + # Round 2 — reuse the same signal; stale credits from round 1 would + # make this call's ready barrier pass spuriously. + for rank in pl.range(pld.world_size()): + data = pld.window(data_buf, [1, sz], dtype=pl.FP32) + self.publish_orch(inputs[1, rank], data, device=rank) + data = pld.window(data_buf, [1, sz], dtype=pl.FP32) + data = pld.tensor.allreduce( + data, + signal, + op=pld.ReduceOp.Sum, + core_num=cores, + ) + for rank in pl.range(pld.world_size()): + self.consume_orch(data, outputs[1, rank], device=rank) + return outputs + + return HostTensorAllReduceMulticoreSignalReuse + + class TestL3HostTensorAllReduceMulticore: @pytest.mark.parametrize( ("n_ranks", "core_num", "size", "signal_stride"), @@ -213,12 +341,69 @@ def test_multicore_output_and_signal_lanes( for peer in range(n_ranks): if peer == rank: continue - assert torch.all(signal_outputs[rank, peer, :core_num] >= 1), ( - f"missing signal lane for P={n_ranks}, C={core_num}, size={size}, " + assert torch.all(signal_outputs[rank, peer, :core_num] == 0), ( + f"signal lane not self-cleared for P={n_ranks}, C={core_num}, size={size}, " f"stride={signal_stride}, receiver={rank}, sender={peer}: " f"got {signal_outputs[rank, peer].tolist()}" ) + @pytest.mark.parametrize( + ("n_ranks", "core_num", "size", "signal_stride"), + [ + pytest.param(2, 2, 1, 2, id="reuse-p2-c2-idle-lane"), + pytest.param(2, 4, 4127, 6, id="reuse-p2-c4-wide-stride"), + ], + ) + def test_multicore_allreduce_signal_reuse( + self, + test_config, + device_ids, + n_ranks, + core_num, + size, + signal_stride, + ): + """Reuse ONE multicore signal across 2 back-to-back allreduce calls. + + Output correctness is the definitive proof of correct signal reuse: if + the per-lane epilogue fails to self-clear after round 1, round 2's + Ge(1) ready barrier passes spuriously on stale credits and the reduction + reads peers' data too early — producing wrong output. No on-device + signal readback is attempted because TWAIT(Eq 0) risks reading stale + cached values from a previous AIV task dispatch on non-coherent NPU + silicon; the sequential sim executor cannot distinguish the two cases. + """ + if len(device_ids) < n_ranks: + pytest.skip(f"multicore host allreduce P={n_ranks} needs {n_ranks} devices, got {device_ids}") + + rounds = 2 + program = _build_multicore_allreduce_signal_reuse(n_ranks, size, core_num, signal_stride) + compiled = ir.compile( + program, + platform=test_config.platform, + distributed_config=DistributedConfig( + device_ids=device_ids[:n_ranks], + num_sub_workers=0, + ), + ) + + # Each round carries a distinct offset so a stale round-1 result in a + # later round (a missed epilogue reset) cannot match the round's golden. + inputs = torch.stack( + [_make_rank_inputs(n_ranks, size, round_offset=rd * 10000.0) for rd in range(rounds)] + ) + outputs = torch.zeros((rounds, n_ranks, 1, size), dtype=torch.float32) + + compiled(inputs, outputs) + + for rd in range(rounds): + expected = _expected_allreduce(inputs[rd]) + assert torch.allclose(outputs[rd], expected, rtol=1e-4, atol=1e-5), ( + f"multicore host allreduce signal-reuse round {rd} " + f"P={n_ranks}, C={core_num}, size={size} mismatch: " + f"max diff = {(outputs[rd] - expected).abs().max().item()}" + ) + if __name__ == "__main__": pytest.main([__file__, "-v", *sys.argv[1:]]) diff --git a/tests/st/distributed/test_l3_host_tensor_allreduce_ring.py b/tests/st/distributed/test_l3_host_tensor_allreduce_ring.py index 9d05105a43..bc346de305 100644 --- a/tests/st/distributed/test_l3_host_tensor_allreduce_ring.py +++ b/tests/st/distributed/test_l3_host_tensor_allreduce_ring.py @@ -26,9 +26,12 @@ def _expected_allreduce(inputs: torch.Tensor) -> torch.Tensor: return torch.stack([reduced] * inputs.shape[0]) -def _make_rank_inputs(n_ranks: int) -> torch.Tensor: +def _make_rank_inputs(n_ranks: int, round_offset: float = 0.0) -> torch.Tensor: + """Build distinct per-rank rows; ``round_offset`` distinguishes reuse rounds.""" rows = [ - torch.arange(r * 100.0, r * 100.0 + SIZE, dtype=torch.float32).reshape(1, SIZE) + torch.arange(r * 100.0 + round_offset, r * 100.0 + round_offset + SIZE, dtype=torch.float32).reshape( + 1, SIZE + ) for r in range(n_ranks) ] return torch.stack(rows) @@ -98,6 +101,94 @@ def host_orch( return HostTensorAllReduceRing +def _build_host_ring_allreduce_signal_reuse_program(n_ranks: int): + """Host ring allreduce reusing ONE signal buffer across 3 back-to-back calls. + + Same self-clearing-credit-barrier intent as the mesh reuse test; the ring + signal is the [2*(NR-1)+1, NR] per-round matrix and every used row is reset by + the epilogue.""" + + total_rounds = 2 * (n_ranks - 1) + 1 + rounds = 3 + + @pl.program + class HostTensorAllReduceRingSignalReuse: + @pl.function(type=pl.FunctionType.InCore) + def publish_step( + self, + inp: pl.Tensor[[1, SIZE], pl.FP32], + data: pl.InOut[pld.DistributedTensor[[1, SIZE], pl.FP32]], + ) -> pld.DistributedTensor[[1, SIZE], pl.FP32]: + local = pl.load(inp, [0, 0], [1, SIZE]) + return pl.store(local, [0, 0], data) + + @pl.function(type=pl.FunctionType.Orchestration) + def publish_orch( + self, + inp: pl.Tensor[[1, SIZE], pl.FP32], + data: pl.InOut[pld.DistributedTensor[[1, SIZE], pl.FP32]], + ) -> pld.DistributedTensor[[1, SIZE], pl.FP32]: + return self.publish_step(inp, data) + + @pl.function(type=pl.FunctionType.InCore) + def consume_step( + self, + data: pld.DistributedTensor[[1, SIZE], pl.FP32], + out: pl.Out[pl.Tensor[[1, SIZE], pl.FP32]], + ) -> pl.Tensor[[1, SIZE], pl.FP32]: + reduced = pl.load(data, [0, 0], [1, SIZE]) + return pl.store(reduced, [0, 0], out) + + @pl.function(type=pl.FunctionType.Orchestration) + def consume_orch( + self, + data: pld.DistributedTensor[[1, SIZE], pl.FP32], + out: pl.Out[pl.Tensor[[1, SIZE], pl.FP32]], + ) -> pl.Tensor[[1, SIZE], pl.FP32]: + return self.consume_step(data, out) + + @pl.function(level=pl.Level.HOST, role=pl.Role.Orchestrator) + def host_orch( + self, + inputs: pl.Tensor[[rounds, n_ranks, 1, SIZE], pl.FP32], + outputs: pl.Out[pl.Tensor[[rounds, n_ranks, 1, SIZE], pl.FP32]], + ) -> pl.Tensor[[rounds, n_ranks, 1, SIZE], pl.FP32]: + data_buf = pld.alloc_window_buffer(SIZE * pl.FP32.get_byte()) + signal_buf = pld.alloc_window_buffer(total_rounds * n_ranks * pl.INT32.get_byte()) + signal = pld.window(signal_buf, [total_rounds, n_ranks], dtype=pl.INT32) + + # Round 1 — every round below reuses the shared ``signal``. + for r in pl.range(pld.world_size()): + data = pld.window(data_buf, [1, SIZE], dtype=pl.FP32) + self.publish_orch(inputs[0, r], data, device=r) + data = pld.window(data_buf, [1, SIZE], dtype=pl.FP32) + data = pld.tensor.allreduce(data, signal, op=pld.ReduceOp.Sum, mode="ring") + for r in pl.range(pld.world_size()): + self.consume_orch(data, outputs[0, r], device=r) + + # Round 2 — reuse the same signal. + for r in pl.range(pld.world_size()): + data = pld.window(data_buf, [1, SIZE], dtype=pl.FP32) + self.publish_orch(inputs[1, r], data, device=r) + data = pld.window(data_buf, [1, SIZE], dtype=pl.FP32) + data = pld.tensor.allreduce(data, signal, op=pld.ReduceOp.Sum, mode="ring") + for r in pl.range(pld.world_size()): + self.consume_orch(data, outputs[1, r], device=r) + + # Round 3 — reuse the same signal again. + for r in pl.range(pld.world_size()): + data = pld.window(data_buf, [1, SIZE], dtype=pl.FP32) + self.publish_orch(inputs[2, r], data, device=r) + data = pld.window(data_buf, [1, SIZE], dtype=pl.FP32) + data = pld.tensor.allreduce(data, signal, op=pld.ReduceOp.Sum, mode="ring") + for r in pl.range(pld.world_size()): + self.consume_orch(data, outputs[2, r], device=r) + + return outputs + + return HostTensorAllReduceRingSignalReuse + + class TestL3HostTensorAllReduceRing: @pytest.mark.parametrize("n_ranks", [2, 4]) def test_host_tensor_allreduce_ring(self, test_config, device_ids, n_ranks): @@ -128,6 +219,37 @@ def test_host_tensor_allreduce_ring(self, test_config, device_ids, n_ranks): f"host ring allreduce P={n_ranks} mismatch: max diff = {(outputs - expected).abs().max().item()}" ) + @pytest.mark.parametrize("n_ranks", [2, 4]) + def test_host_tensor_allreduce_ring_signal_reuse(self, test_config, device_ids, n_ranks): + """Reuse ONE ring signal matrix across 3 back-to-back calls.""" + if len(device_ids) < n_ranks: + pytest.skip(f"host ring allreduce P={n_ranks} needs {n_ranks} devices, got {device_ids}") + + rounds = 3 + compiled = ir.compile( + _build_host_ring_allreduce_signal_reuse_program(n_ranks), + platform=test_config.platform, + distributed_config=DistributedConfig( + device_ids=device_ids[:n_ranks], + num_sub_workers=0, + ), + ) + variant_dir = compiled.output_dir / "next_levels" / "builtin.tensor.allreduce_ring__sum__fp32" + assert variant_dir.is_dir() + + # Each round carries a distinct offset so a stale round-1 result in a + # later round (a missed epilogue reset) cannot match the round's golden. + inputs = torch.stack([_make_rank_inputs(n_ranks, round_offset=rd * 10000.0) for rd in range(rounds)]) + outputs = torch.zeros_like(inputs) + compiled(inputs, outputs) + + for rd in range(rounds): + expected = _expected_allreduce(inputs[rd]) + assert torch.allclose(outputs[rd], expected), ( + f"host ring allreduce signal-reuse round {rd} P={n_ranks} mismatch: " + f"max diff = {(outputs[rd] - expected).abs().max().item()}" + ) + if __name__ == "__main__": sys.exit(pytest.main([__file__, "-v", *sys.argv[1:]])) diff --git a/tests/st/distributed/test_l3_host_tensor_barrier.py b/tests/st/distributed/test_l3_host_tensor_barrier.py index 365686dd2e..73945a2fd2 100644 --- a/tests/st/distributed/test_l3_host_tensor_barrier.py +++ b/tests/st/distributed/test_l3_host_tensor_barrier.py @@ -26,9 +26,12 @@ def _expected_peer_swap(inputs: torch.Tensor) -> torch.Tensor: return torch.stack([inputs[1], inputs[0]]) -def _make_rank_inputs(n_ranks: int) -> torch.Tensor: +def _make_rank_inputs(n_ranks: int, round_offset: float = 0.0) -> torch.Tensor: + """Build distinct per-rank rows; ``round_offset`` distinguishes reuse rounds.""" rows = [ - torch.arange(r * 100.0, r * 100.0 + SIZE, dtype=torch.float32).reshape(1, SIZE) + torch.arange(r * 100.0 + round_offset, r * 100.0 + round_offset + SIZE, dtype=torch.float32).reshape( + 1, SIZE + ) for r in range(n_ranks) ] return torch.stack(rows) @@ -97,6 +100,92 @@ def host_orch( return outputs +def _build_host_barrier_signal_reuse_program(): + """Host barrier reusing ONE signal buffer across 2 back-to-back calls. + + The self-clearing epilogue restores the AtomicAdd(+1) cells to 0 after each + call; without it the second call's Ge(1) wait passes on the stale + satisfied cell. + """ + ROUNDS = 2 + + @pl.program + class HostTensorBarrierSignalReuse: + @pl.function(type=pl.FunctionType.InCore) + def publish_step( + self, + inp: pl.Tensor[[1, SIZE], pl.FP32], + data: pl.InOut[pld.DistributedTensor[[1, SIZE], pl.FP32]], + ) -> pld.DistributedTensor[[1, SIZE], pl.FP32]: + local = pl.load(inp, [0, 0], [1, SIZE]) + return pl.store(local, [0, 0], data) + + @pl.function(type=pl.FunctionType.Orchestration) + def publish_orch( + self, + inp: pl.Tensor[[1, SIZE], pl.FP32], + data: pl.InOut[pld.DistributedTensor[[1, SIZE], pl.FP32]], + sig: pld.DistributedTensor[[NR], pl.INT32], + ) -> pld.DistributedTensor[[1, SIZE], pl.FP32]: + return self.publish_step(inp, data) + + @pl.function(type=pl.FunctionType.InCore) + def consume_step( + self, + data: pld.DistributedTensor[[1, SIZE], pl.FP32], + out: pl.Out[pl.Tensor[[1, SIZE], pl.FP32]], + peer: pl.Scalar[pl.INT32], + ) -> pl.Tensor[[1, SIZE], pl.FP32]: + recv = pld.tile.remote_load(data, peer=peer, offsets=[0, 0], shape=[1, SIZE]) + return pl.store(recv, [0, 0], out) + + @pl.function(type=pl.FunctionType.Orchestration) + def consume_orch( + self, + data: pld.DistributedTensor[[1, SIZE], pl.FP32], + out: pl.Out[pl.Tensor[[1, SIZE], pl.FP32]], + peer: pl.Scalar[pl.INT32], + ) -> pl.Tensor[[1, SIZE], pl.FP32]: + return self.consume_step(data, out, peer) + + @pl.function(level=pl.Level.HOST, role=pl.Role.Orchestrator) + def host_orch( + self, + inputs: pl.Tensor[[ROUNDS, NR, 1, SIZE], pl.FP32], + outputs: pl.Out[pl.Tensor[[ROUNDS, NR, 1, SIZE], pl.FP32]], + ) -> pl.Tensor[[ROUNDS, NR, 1, SIZE], pl.FP32]: + data_buf = pld.alloc_window_buffer(SIZE * pl.FP32.get_byte()) + signal_buf = pld.alloc_window_buffer(pld.world_size() * pl.INT32.get_byte()) + signal = pld.window(signal_buf, [pld.world_size()], dtype=pl.INT32) + + # Round 1 — every round below reuses the shared ``signal``. The + # barrier is a bare call (not ``signal = pld.tensor.barrier(signal)``) + # because the rebind would make the next round's input a + # barrier-result var, which MaterializeCommDomainScopes rejects. + for r in pl.range(pld.world_size()): + data = pld.window(data_buf, [1, SIZE], dtype=pl.FP32) + self.publish_orch(inputs[0, r], data, signal, device=r) + pld.tensor.barrier(signal) + for r in pl.range(pld.world_size()): + data = pld.window(data_buf, [1, SIZE], dtype=pl.FP32) + peer = (r + 1) % pld.world_size() + self.consume_orch(data, outputs[0, r], peer, device=r) + + # Round 2 — reuse the same signal. + for r in pl.range(pld.world_size()): + data = pld.window(data_buf, [1, SIZE], dtype=pl.FP32) + self.publish_orch(inputs[1, r], data, signal, device=r) + pld.tensor.barrier(signal) + for r in pl.range(pld.world_size()): + data = pld.window(data_buf, [1, SIZE], dtype=pl.FP32) + peer = (r + 1) % pld.world_size() + self.consume_orch(data, outputs[1, r], peer, device=r) + + return outputs + + return HostTensorBarrierSignalReuse + + class TestL3HostTensorBarrier: @pytest.mark.parametrize("n_ranks", [2]) def test_host_tensor_barrier(self, test_config, device_ids, n_ranks): @@ -126,6 +215,42 @@ def test_host_tensor_barrier(self, test_config, device_ids, n_ranks): f"host barrier P={n_ranks} mismatch: max diff = {(outputs - expected).abs().max().item()}" ) + @pytest.mark.parametrize("n_ranks", [2]) + def test_host_tensor_barrier_signal_reuse(self, test_config, device_ids, n_ranks): + """Reuse ONE signal buffer across 2 back-to-back barrier calls. + + The self-clearing epilogue restores the signal to all-zero after each + call; without it the second call's Ge(1) wait passes on the stale + satisfied cell (NPU-visible; sim is sequentially consistent). + """ + if len(device_ids) < n_ranks: + pytest.skip(f"host barrier P={n_ranks} needs {n_ranks} devices, got {device_ids}") + + rounds = 2 + compiled = ir.compile( + _build_host_barrier_signal_reuse_program(), + platform=test_config.platform, + distributed_config=DistributedConfig( + device_ids=device_ids[:n_ranks], + num_sub_workers=0, + ), + ) + variant_dir = compiled.output_dir / "next_levels" / "builtin.tensor.barrier__fp32" + assert variant_dir.is_dir() + + # Each round carries a distinct offset so a stale earlier-round result + # (a missed epilogue reset) cannot match the round's golden. + inputs = torch.stack([_make_rank_inputs(n_ranks, round_offset=rd * 10000.0) for rd in range(rounds)]) + outputs = torch.zeros_like(inputs) + compiled(inputs, outputs) + + for rd in range(rounds): + expected = _expected_peer_swap(inputs[rd]) + assert torch.allclose(outputs[rd], expected), ( + f"host barrier signal-reuse round {rd} P={n_ranks} mismatch: " + f"max diff = {(outputs[rd] - expected).abs().max().item()}" + ) + if __name__ == "__main__": pytest.main([__file__, "-v", *sys.argv[1:]]) diff --git a/tests/st/distributed/test_l3_host_tensor_broadcast.py b/tests/st/distributed/test_l3_host_tensor_broadcast.py index 363a3654ae..114651c3f0 100644 --- a/tests/st/distributed/test_l3_host_tensor_broadcast.py +++ b/tests/st/distributed/test_l3_host_tensor_broadcast.py @@ -28,9 +28,12 @@ def _expected_broadcast(inputs: torch.Tensor, root: int = ROOT_RANK) -> torch.Te return torch.stack([root_row] * inputs.shape[0]).unsqueeze(1) -def _make_rank_inputs(n_ranks: int) -> torch.Tensor: +def _make_rank_inputs(n_ranks: int, round_offset: float = 0.0) -> torch.Tensor: + """Build distinct per-rank rows; ``round_offset`` distinguishes reuse rounds.""" rows = [ - torch.arange(r * 100.0, r * 100.0 + SIZE, dtype=torch.float32).reshape(1, SIZE) + torch.arange(r * 100.0 + round_offset, r * 100.0 + round_offset + SIZE, dtype=torch.float32).reshape( + 1, SIZE + ) for r in range(n_ranks) ] return torch.stack(rows) @@ -99,6 +102,77 @@ def host_orch( return outputs +def _build_host_broadcast_signal_reuse_program(): + """Host broadcast reusing ONE signal buffer across 2 back-to-back calls. + + The self-clearing epilogue restores the AtomicAdd cells to 0 after each + call; without it the second call's Ge(1) wait passes on stale credits. + """ + ROUNDS = 2 + + @pl.program + class HostTensorBroadcastSignalReuse: + @pl.function(type=pl.FunctionType.InCore) + def publish_step( + self, + inp: pl.Tensor[[1, SIZE], pl.FP32], + data: pl.InOut[pld.DistributedTensor[[1, SIZE], pl.FP32]], + my_rank: pl.Scalar[pl.INT32], + ) -> pld.DistributedTensor[[1, SIZE], pl.FP32]: + if my_rank == ROOT_RANK: + local = pl.load(inp, [0, 0], [1, SIZE]) + return pl.store(local, [0, 0], data) + return data + + @pl.function(type=pl.FunctionType.Orchestration) + def publish_orch( + self, + inp: pl.Tensor[[1, SIZE], pl.FP32], + data: pl.InOut[pld.DistributedTensor[[1, SIZE], pl.FP32]], + my_rank: pl.Scalar[pl.INT32], + ) -> pld.DistributedTensor[[1, SIZE], pl.FP32]: + return self.publish_step(inp, data, my_rank) + + @pl.function(type=pl.FunctionType.InCore) + def consume_step( + self, + data: pld.DistributedTensor[[1, SIZE], pl.FP32], + out: pl.Out[pl.Tensor[[1, SIZE], pl.FP32]], + ) -> pl.Tensor[[1, SIZE], pl.FP32]: + acc = pl.load(data, [0, 0], [1, SIZE]) + return pl.store(acc, [0, 0], out) + + @pl.function(type=pl.FunctionType.Orchestration) + def consume_orch( + self, + data: pld.DistributedTensor[[1, SIZE], pl.FP32], + out: pl.Out[pl.Tensor[[1, SIZE], pl.FP32]], + ) -> pl.Tensor[[1, SIZE], pl.FP32]: + return self.consume_step(data, out) + + @pl.function(level=pl.Level.HOST, role=pl.Role.Orchestrator) + def host_orch( + self, + inputs: pl.Tensor[[ROUNDS, NR, 1, SIZE], pl.FP32], + outputs: pl.Out[pl.Tensor[[ROUNDS, NR, 1, SIZE], pl.FP32]], + ) -> pl.Tensor[[ROUNDS, NR, 1, SIZE], pl.FP32]: + data_buf = pld.alloc_window_buffer(SIZE * pl.FP32.get_byte()) + signal_buf = pld.alloc_window_buffer(pld.world_size() * pl.INT32.get_byte()) + signal = pld.window(signal_buf, [pld.world_size()], dtype=pl.INT32) + + for rd in pl.range(ROUNDS): + for r in pl.range(pld.world_size()): + data = pld.window(data_buf, [1, SIZE], dtype=pl.FP32) + self.publish_orch(inputs[rd, r], data, r, device=r) + data = pld.window(data_buf, [1, SIZE], dtype=pl.FP32) + data = pld.tensor.broadcast(data, signal, root=ROOT_RANK) + for r in pl.range(pld.world_size()): + self.consume_orch(data, outputs[rd, r], device=r) + return outputs + + return HostTensorBroadcastSignalReuse + + class TestL3HostTensorBroadcast: @pytest.mark.parametrize("n_ranks", [2]) def test_host_tensor_broadcast(self, test_config, device_ids, n_ranks): @@ -129,6 +203,37 @@ def test_host_tensor_broadcast(self, test_config, device_ids, n_ranks): ) assert not torch.allclose(outputs[0], inputs[1]), "non-root input leaked into output" + @pytest.mark.parametrize("n_ranks", [2, 4]) + def test_host_tensor_broadcast_signal_reuse(self, test_config, device_ids, n_ranks): + """Reuse ONE signal buffer across 2 back-to-back broadcast calls.""" + if len(device_ids) < n_ranks: + pytest.skip(f"host broadcast P={n_ranks} needs {n_ranks} devices, got {device_ids}") + + rounds = 2 + compiled = ir.compile( + _build_host_broadcast_signal_reuse_program(), + platform=test_config.platform, + distributed_config=DistributedConfig( + device_ids=device_ids[:n_ranks], + num_sub_workers=0, + ), + ) + variant_dir = compiled.output_dir / "next_levels" / "builtin.tensor.broadcast__root0__fp32" + assert variant_dir.is_dir() + + # Each round carries a distinct offset so a stale earlier-round result + # (a missed epilogue reset) cannot match the round's golden. + inputs = torch.stack([_make_rank_inputs(n_ranks, round_offset=rd * 10000.0) for rd in range(rounds)]) + outputs = torch.zeros_like(inputs) + compiled(inputs, outputs) + + for rd in range(rounds): + expected = _expected_broadcast(inputs[rd]) + assert torch.allclose(outputs[rd], expected), ( + f"host broadcast signal-reuse round {rd} P={n_ranks} mismatch: " + f"max diff = {(outputs[rd] - expected).abs().max().item()}" + ) + if __name__ == "__main__": pytest.main([__file__, "-v", *sys.argv[1:]]) diff --git a/tests/st/distributed/test_l3_host_tensor_reduce_scatter.py b/tests/st/distributed/test_l3_host_tensor_reduce_scatter.py index bda48dd98a..90e154618a 100644 --- a/tests/st/distributed/test_l3_host_tensor_reduce_scatter.py +++ b/tests/st/distributed/test_l3_host_tensor_reduce_scatter.py @@ -27,9 +27,12 @@ def _expected_reduce_scatter(inputs: torch.Tensor, n_ranks: int) -> torch.Tensor return torch.stack(chunks).reshape(n_ranks, 1, SIZE) -def _make_rank_inputs(n_ranks: int) -> torch.Tensor: +def _make_rank_inputs(n_ranks: int, round_offset: float = 0.0) -> torch.Tensor: + """Build distinct per-rank rows; ``round_offset`` distinguishes reuse rounds.""" rows = [ - torch.arange(r * 100.0, r * 100.0 + n_ranks * SIZE, dtype=torch.float32).reshape(n_ranks, SIZE) + torch.arange( + r * 100.0 + round_offset, r * 100.0 + round_offset + n_ranks * SIZE, dtype=torch.float32 + ).reshape(n_ranks, SIZE) for r in range(n_ranks) ] return torch.stack(rows) @@ -98,6 +101,77 @@ def host_orch( return outputs +def _build_host_reduce_scatter_signal_reuse_program(): + """Host reduce_scatter reusing ONE signal buffer across 2 back-to-back calls. + + The self-clearing epilogue restores the AtomicAdd cells to 0 after each + call; without it the second call's Ge(1) wait passes on stale credits. + """ + ROUNDS = 2 + + @pl.program + class HostTensorReduceScatterSignalReuse: + @pl.function(type=pl.FunctionType.InCore) + def publish_step( + self, + inp: pl.Tensor[[NR, SIZE], pl.FP32], + data: pl.InOut[pld.DistributedTensor[[NR, SIZE], pl.FP32]], + ) -> pld.DistributedTensor[[NR, SIZE], pl.FP32]: + for j in pl.range(NR): + chunk = pl.load(inp, [j, 0], [1, SIZE]) + data = pl.store(chunk, [j, 0], data) + return data + + @pl.function(type=pl.FunctionType.Orchestration) + def publish_orch( + self, + inp: pl.Tensor[[NR, SIZE], pl.FP32], + data: pl.InOut[pld.DistributedTensor[[NR, SIZE], pl.FP32]], + ) -> pld.DistributedTensor[[NR, SIZE], pl.FP32]: + return self.publish_step(inp, data) + + @pl.function(type=pl.FunctionType.InCore) + def consume_step( + self, + data: pld.DistributedTensor[[NR, SIZE], pl.FP32], + out: pl.Out[pl.Tensor[[1, SIZE], pl.FP32]], + my_rank: pl.Scalar[pl.INT32], + ) -> pl.Tensor[[1, SIZE], pl.FP32]: + acc = pl.load(data, [my_rank, 0], [1, SIZE]) + return pl.store(acc, [0, 0], out) + + @pl.function(type=pl.FunctionType.Orchestration) + def consume_orch( + self, + data: pld.DistributedTensor[[NR, SIZE], pl.FP32], + out: pl.Out[pl.Tensor[[1, SIZE], pl.FP32]], + my_rank: pl.Scalar[pl.INT32], + ) -> pl.Tensor[[1, SIZE], pl.FP32]: + return self.consume_step(data, out, my_rank) + + @pl.function(level=pl.Level.HOST, role=pl.Role.Orchestrator) + def host_orch( + self, + inputs: pl.Tensor[[ROUNDS, NR, NR, SIZE], pl.FP32], + outputs: pl.Out[pl.Tensor[[ROUNDS, NR, 1, SIZE], pl.FP32]], + ) -> pl.Tensor[[ROUNDS, NR, 1, SIZE], pl.FP32]: + data_buf = pld.alloc_window_buffer(NR * SIZE * pl.FP32.get_byte()) + signal_buf = pld.alloc_window_buffer(pld.world_size() * pl.INT32.get_byte()) + signal = pld.window(signal_buf, [pld.world_size()], dtype=pl.INT32) + + for rd in pl.range(ROUNDS): + for r in pl.range(pld.world_size()): + data = pld.window(data_buf, [NR, SIZE], dtype=pl.FP32) + self.publish_orch(inputs[rd, r], data, device=r) + data = pld.window(data_buf, [NR, SIZE], dtype=pl.FP32) + data = pld.tensor.reduce_scatter(data, signal, op=pld.ReduceOp.Sum) + for r in pl.range(pld.world_size()): + self.consume_orch(data, outputs[rd, r], r, device=r) + return outputs + + return HostTensorReduceScatterSignalReuse + + class TestL3HostTensorReduceScatter: @pytest.mark.parametrize("n_ranks", [2]) def test_host_tensor_reduce_scatter(self, test_config, device_ids, n_ranks): @@ -127,6 +201,37 @@ def test_host_tensor_reduce_scatter(self, test_config, device_ids, n_ranks): f"host reduce_scatter P={n_ranks} mismatch: max diff = {(outputs - expected).abs().max().item()}" ) + @pytest.mark.parametrize("n_ranks", [2, 4]) + def test_host_tensor_reduce_scatter_signal_reuse(self, test_config, device_ids, n_ranks): + """Reuse ONE signal buffer across 2 back-to-back reduce_scatter calls.""" + if len(device_ids) < n_ranks: + pytest.skip(f"host reduce_scatter P={n_ranks} needs {n_ranks} devices, got {device_ids}") + + rounds = 2 + compiled = ir.compile( + _build_host_reduce_scatter_signal_reuse_program(), + platform=test_config.platform, + distributed_config=DistributedConfig( + device_ids=device_ids[:n_ranks], + num_sub_workers=0, + ), + ) + variant_dir = compiled.output_dir / "next_levels" / "builtin.tensor.reduce_scatter__sum__fp32" + assert variant_dir.is_dir() + + # Each round carries a distinct offset so a stale earlier-round result + # (a missed epilogue reset) cannot match the round's golden. + inputs = torch.stack([_make_rank_inputs(n_ranks, round_offset=rd * 10000.0) for rd in range(rounds)]) + outputs = torch.zeros((rounds, n_ranks, 1, SIZE), dtype=torch.float32) + compiled(inputs, outputs) + + for rd in range(rounds): + expected = _expected_reduce_scatter(inputs[rd], n_ranks) + assert torch.allclose(outputs[rd], expected), ( + f"host reduce_scatter signal-reuse round {rd} P={n_ranks} mismatch: " + f"max diff = {(outputs[rd] - expected).abs().max().item()}" + ) + if __name__ == "__main__": pytest.main([__file__, "-v", *sys.argv[1:]]) diff --git a/tests/ut/ir/transforms/test_materialize_comm_domain_scopes.py b/tests/ut/ir/transforms/test_materialize_comm_domain_scopes.py index 987ef984ff..d5146a9482 100644 --- a/tests/ut/ir/transforms/test_materialize_comm_domain_scopes.py +++ b/tests/ut/ir/transforms/test_materialize_comm_domain_scopes.py @@ -955,11 +955,10 @@ def host_orch(self): def test_implicit_allreduce_in_loop_is_rejected(): - """The HOST builtin allreduce is not self-clearing (it adds credits without - subtracting them), so a signal reused across a dynamic trip count would pass - its waits on stale state. The loop restriction therefore stays on the HOST - signal-synthesis path; only InCore composites (LowerCompositeOps) are - loop-safe under the self-clearing protocol. + """The HOST-rail signal synthesis cannot allocate a fresh signal per dynamic + iteration, so a synthesized-signal allreduce inside a loop is rejected. + (InCore composites are loop-safe via the self-clearing credit-barrier + protocol; lifting the HOST restriction is tracked separately.) """ @pl.program