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

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)
 

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
 

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
 

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.*`)
 

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
 

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）
 

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 属性
 

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
 

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.*`)
 

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 调用提到循环外。 |
 
 ## 致命陷阱
 

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]

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

python/pypto/runtime/builtins/collectives/all_to_all/templates/kernel.cpp.in

@@ -118,21 +118,30 @@ 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);
   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<int32_t>(1), pto::comm::NotifyOp::Set);
+    pto::comm::TNOTIFY(sig, static_cast<int32_t>(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<int32_t>(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<int32_t>(-1), pto::comm::NotifyOp::AtomicAdd);
+  }
   pipe_barrier(PIPE_ALL);
 }

python/pypto/runtime/builtins/collectives/all_to_all_v/templates/kernel.cpp.in

@@ -162,21 +162,30 @@ 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);
   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<int32_t>(1), pto::comm::NotifyOp::Set);
+    pto::comm::TNOTIFY(sig, static_cast<int32_t>(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<int32_t>(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<int32_t>(-1), pto::comm::NotifyOp::AtomicAdd);
+  }
   pipe_barrier(PIPE_ALL);
 }

python/pypto/runtime/builtins/collectives/allgather/templates/kernel.cpp.in

@@ -121,21 +121,30 @@ 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);
   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<int32_t>(1), pto::comm::NotifyOp::Set);
+    pto::comm::TNOTIFY(sig, static_cast<int32_t>(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<int32_t>(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<int32_t>(-1), pto::comm::NotifyOp::AtomicAdd);
+  }
   pipe_barrier(PIPE_ALL);
 }

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);
 }

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<size_t>(r) * signal_cols + peer);
+      pto::comm::TNOTIFY(self_sig, static_cast<int32_t>(-1), pto::comm::NotifyOp::AtomicAdd);
+    }
+  }
+
   pipe_barrier(PIPE_ALL);
 }

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<int>(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<int32_t>(1), pto::comm::NotifyOp::Set);
+    pto::comm::TNOTIFY(sig, static_cast<int32_t>(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<int32_t>(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<int32_t>(0), pto::comm::NotifyOp::Set);
+  }
   pipe_barrier(PIPE_ALL);
 }