update(distributed): convert all_to_all distributed to push-based lowering

examples/workers/l3/README.md

@@ -66,7 +66,7 @@ Two things to know before reading the example:
 | [`allgather_distributed/`](allgather_distributed/) | One communication domain via `orch.allocate_domain`; each rank stages its slice, synchronizes across ranks, then gathers every rank's window data into a full output. |
 | [`reduce_scatter_distributed/`](reduce_scatter_distributed/) | One communication domain via `orch.allocate_domain`; each rank stages all input chunks, synchronizes, then reduces the per-rank chunk across peers into a rank-local output. |
 | [`broadcast_distributed/`](broadcast_distributed/) | One communication domain via `orch.allocate_domain`; root stages into the window, synchronizes, then every rank reads the root's scratch slot into its output. |
-| [`all_to_all_distributed/`](all_to_all_distributed/) | One communication domain via `orch.allocate_domain`; scratch indexed by destination rank, barrier, then each rank gathers the slice peers sent to it. |
+| [`all_to_all_distributed/`](all_to_all_distributed/) | One communication domain via `orch.allocate_domain`; each rank pushes its per-destination chunks directly to peers' scratch via TPUT, barrier synchronizes, then each rank copies its local result to output. |
 | [`ffn_tp_parallel/`](ffn_tp_parallel/) | Local compute followed by one-domain cross-rank reduction through a domain scratch window. |
 | [`ep_dispatch_combine/`](ep_dispatch_combine/) | MoE-style dispatch/combine over a one-domain communication window. |
 | [`domain_rank_map/`](domain_rank_map/) | Small two-domain example showing domain-local ranks, missing-domain `KeyError`, separate window slices, and real per-domain allreduce. |

examples/workers/l3/all_to_all_distributed/README.md

@@ -2,17 +2,21 @@
 
 This example demonstrates distributed all-to-all (personalized exchange) — each rank sends a different slice to each peer.
 
-## Algorithm (3-Phase Mesh)
+## Algorithm (2-Phase Push)
 
 ```text
-Phase 1: stage-in      for d in 0..P-1: input[d * C] → scratch[d * C] (chunk d is for rank d)
-Phase 2: barrier       mesh barrier (all-to-all notify/wait)
-Phase 3: exchange      for s in 0..P-1: TLOAD(peer s scratch[my_rank * C]) → output[s * C]
+Phase 1: push        for d in 0..P-1: TPUT input[d * C] → peer d's scratch[my_rank * C]
+                     Self-rank handled naturally — CommRemotePtr to self returns local pointer
+Phase 2: barrier     mesh barrier (all-to-all notify/wait); no post-exchange barrier needed
+Phase 3: copy-out    for s in 0..P-1: TLOAD(local scratch[s * C]) → output[s * C]  (purely local)
 ```
 
 **Input**: Each rank owns `nranks * COUNT_PER_RANK` floats (chunk d is payload for rank d).
 **Output**: Each rank receives `nranks * COUNT_PER_RANK` floats (chunk s is what rank s sent to me).
 
+After Phase 1, scratch at offset `src * C` holds the chunk rank `src` pushed here.
+After the barrier, the result is in scratch — Phase 3 is a purely local copy to the output tensor.
+
 ## Usage
 
 ```bash

examples/workers/l3/all_to_all_distributed/kernels/aiv/all_to_all_kernel.cpp

@@ -9,13 +9,23 @@
  * -----------------------------------------------------------------------------------------------------------
  */
 /**
- * AllToAll kernel — symmetric, 3-phase, HCCL-window scratch pattern.
+ * AllToAll kernel — symmetric, 2-phase push-based pattern.
  *
- * Phase 1 (stage-in):   for dest in 0..nranks-1: input[dest*C..) → scratch[dest*C..)
- * Phase 2 (barrier):    signal matrix + TWAIT cross-rank sync
- * Phase 3 (exchange):   for src in 0..nranks-1: TLOAD(peer src scratch[my_rank*C]) → output[src*C..)
+ * Phase 1 (push):     for dest in 0..nranks-1:
+ *                       TPUT input[dest*C..) → peer dest's scratch[my_rank*C..)
+ *                     Self-rank handled naturally: CommRemotePtr to self
+ *                     returns the local pointer, so TPUT degenerates to a
+ *                     local GM write.
+ * Phase 2 (barrier):  signal matrix + TWAIT cross-rank sync.
+ *                     No post-exchange barrier needed — input is read-only
+ *                     and the scratch is write-only from peers, so there is
+ *                     no WAR hazard on the scratch.
+ * Phase 3 (copy-out): for src in 0..nranks-1:
+ *                       TLOAD(local scratch[src*C..)) → TSTORE(output[src*C..))
+ *                     Purely local copy — the scratch already holds the
+ *                     result after the barrier.
  *
- * Scratch is indexed by destination rank: scratch[dest*C] holds the chunk sent to dest.
+ * Scratch at offset src*C holds the chunk that rank src pushed here.
  *
  * args layout:
  *   tensor(0) = input    nranks*COUNT_PER_RANK floats  (INPUT)
@@ -74,35 +84,34 @@ extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ in
         return;
     }
 
-    // signal_base follows the nranks * COUNT_PER_RANK float staging region.
+    // signal_base follows the nranks * COUNT_PER_RANK float data region.
     __gm__ int32_t *signal_base = reinterpret_cast<__gm__ int32_t *>(scratch + nranks * COUNT_PER_RANK);
 
     ShapeDyn shape(1, 1, 1, 1, COUNT_PER_RANK);
     StrideDyn stride(COUNT_PER_RANK, COUNT_PER_RANK, COUNT_PER_RANK, COUNT_PER_RANK, 1);
 
-    TileData stageTile(1, COUNT_PER_RANK);
-    TileData recvTile(1, COUNT_PER_RANK);
-    TASSIGN(stageTile, 0x0);
-    TASSIGN(recvTile, 0x10000);
+    // Single push tile — reused for both Phase 1 push and Phase 3 copy-out.
+    TileData pushTile(1, COUNT_PER_RANK);
+    TASSIGN(pushTile, 0x0);
 
     // ------------------------------------------------------------------
-    // Phase 1: stage-in — copy each destination chunk into scratch so
-    // every peer can read the slice destined for them in Phase 3.
+    // Phase 1: push — TPUT each destination chunk directly into the
+    // corresponding peer's scratch at offset my_rank*C. When dest == my_rank
+    // CommRemotePtr returns the local pointer, so TPUT is a local GM write.
     // ------------------------------------------------------------------
+    __gm__ float *scratch_dst_local = scratch + my_rank * COUNT_PER_RANK;
     for (int dest = 0; dest < nranks; ++dest) {
+        __gm__ float *scratch_dst_remote = CommRemotePtr(commCtx, scratch_dst_local, dest);
+
         Global inputChunkG(input + dest * COUNT_PER_RANK, shape, stride);
-        Global scratchChunkG(scratch + dest * COUNT_PER_RANK, shape, stride);
-        TLOAD(stageTile, inputChunkG);
-        set_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0);
-        wait_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0);
-        TSTORE(scratchChunkG, stageTile);
-        set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0);
-        wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0);
+        Global scratchChunkG(scratch_dst_remote, shape, stride);
+
+        pto::comm::TPUT(scratchChunkG, inputChunkG, pushTile);
     }
     pipe_barrier(PIPE_ALL);
 
     // ------------------------------------------------------------------
-    // Phase 2: device barrier — notify every peer that stage-in is done,
+    // Phase 2: device barrier — notify every peer that our pushes are done,
     // then wait until every peer has notified us.
     // ------------------------------------------------------------------
     for (int peer = 0; peer < nranks; ++peer) {
@@ -119,19 +128,17 @@ extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ in
     pipe_barrier(PIPE_ALL);
 
     // ------------------------------------------------------------------
-    // Phase 3: exchange — read chunk my_rank from every rank's scratch and
-    // write it into the corresponding slice of the output tensor.
-    // CommRemotePtr with pe==my_rank returns localPtr unchanged, so the
-    // self-read goes through the same code path as the remote reads.
+    // Phase 3: copy-out — every peer has pushed into our scratch at their
+    // rank's offset. Copy each slot into the output tensor. Purely local
+    // reads — no CommRemotePtr needed here.
     // ------------------------------------------------------------------
     for (int src = 0; src < nranks; ++src) {
-        __gm__ float *remote_chunk = CommRemotePtr(commCtx, scratch + my_rank * COUNT_PER_RANK, src);
-        Global remoteG(remote_chunk, shape, stride);
+        Global scratchSlotG(scratch + src * COUNT_PER_RANK, shape, stride);
         Global outputSlotG(output + src * COUNT_PER_RANK, shape, stride);
-        TLOAD(recvTile, remoteG);
+        TLOAD(pushTile, scratchSlotG);
         set_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0);
         wait_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0);
-        TSTORE(outputSlotG, recvTile);
+        TSTORE(outputSlotG, pushTile);
         set_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0);
         wait_flag(PIPE_MTE3, PIPE_MTE2, EVENT_ID0);
     }

examples/workers/l3/all_to_all_distributed/kernels/orchestration/all_to_all_orch.cpp

@@ -9,11 +9,11 @@
  * -----------------------------------------------------------------------------------------------------------
  */
 /**
- * AllToAll orchestration shim.
+ * AllToAll orchestration shim — push-based (2 kernel phases + local copy-out).
  *
  *   tensor(0) input   INPUT           (nranks*COUNT_PER_RANK floats)
  *   tensor(1) output  OUTPUT_EXISTING (nranks*COUNT_PER_RANK floats)
- *   tensor(2) scratch INOUT           (HCCL window slot)
+ *   tensor(2) scratch INOUT           (HCCL window slot — holds result after barrier)
  *   scalar(0) nranks
  *   scalar(1) CommContext device pointer
  */

examples/workers/l3/all_to_all_distributed/main.py

@@ -7,14 +7,14 @@
 # INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 # See LICENSE in the root of the software repository for the full text of the License.
 # -----------------------------------------------------------------------------------------------------------
-"""End-to-end distributed all-to-all — symmetric 3-phase pattern.
+"""End-to-end distributed all-to-all — symmetric 2-phase push-based pattern.
 
 Each rank owns nranks send chunks; chunk d is payload for rank d. After the
 exchange rank r holds chunk s from every source s in output[s*COUNT_PER_RANK]:
 
-  Phase 1 stage-in      for d in 0..N-1: input[d*C..) → scratch[d*C..)
+  Phase 1 push         for d in 0..N-1: TPUT input[d*C..) → peer d's scratch[my_rank*C..)
   Phase 2 device barrier signal matrix cross-rank sync via TNOTIFY/TWAIT
-  Phase 3 exchange      for s in 0..N-1: TLOAD(peer s scratch[my_rank*C]) → output[s*C..)
+  Phase 3 copy-out     for s in 0..N-1: TLOAD(local scratch[s*C..)) → output[s*C..)
 
 Run:
     python examples/workers/l3/all_to_all_distributed/main.py -p a2a3sim -d 0-1
@@ -55,7 +55,7 @@
 DTYPE_NBYTES = 4  # float32
 # Signal tail: one int32 slot per rank, bounded by kMaxSupportedRanks.
 SIGNAL_TAIL_NBYTES = 16 * 4  # 64 B
-# NOTE: the full scratch size depends on nranks (staging area = nranks * COUNT_PER_RANK floats).
+# NOTE: the full scratch size depends on nranks (result area = nranks * COUNT_PER_RANK floats).
 # It is computed inside run() once nranks is known.
 
 
@@ -119,7 +119,7 @@ def run(
 ) -> int:
     """Core logic — callable from both CLI and pytest."""
     nranks = len(device_ids)
-    # Scratch = nranks * COUNT_PER_RANK floats (staging) + signal tail.
+    # Scratch = nranks * COUNT_PER_RANK floats (result area) + signal tail.
     scratch_nbytes = nranks * COUNT_PER_RANK * DTYPE_NBYTES + SIGNAL_TAIL_NBYTES
     window_size = max(scratch_nbytes, 4 * 1024)
 

examples/workers/l3/ep_dispatch_combine/kernels/aiv/combine.cpp

@@ -162,15 +162,9 @@ extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ in
                 __gm__ bfloat16_t *dst_row_local = routed_y_buf_local + r * D;
                 __gm__ bfloat16_t *dst_row_remote = CommRemotePtr(comm_ctx, dst_row_local, dst);
 
-#if defined(__CPU_SIM)
-                for (int i = 0; i < D; ++i) {
-                    dst_row_remote[i] = src_row[i];
-                }
-#else
                 RowBfG src_g(src_row);
                 RowBfG dst_g(dst_row_remote);
                 pto::comm::TPUT(dst_g, src_g, push_tile);
-#endif
             }
         }
     }

examples/workers/l3/ep_dispatch_combine/kernels/aiv/dispatch.cpp

@@ -373,45 +373,27 @@ extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ in
         __gm__ bfloat16_t *x_dst_local = recv_x_local + row * D;
         __gm__ bfloat16_t *x_dst_remote = CommRemotePtr(comm_ctx, x_dst_local, dst);
 
-#if defined(__CPU_SIM)
-        for (int i = 0; i < D; ++i) {
-            x_dst_remote[i] = x_src[i];
-        }
-#else
         XGlobal x_src_g(x_src);
         XGlobal x_dst_g(x_dst_remote);
         pto::comm::TPUT(x_dst_g, x_src_g, x_tile);
-#endif
 
         // Channel 2: weight (FP32, 1xW_PAD); host pre-packed [w, 0, …, 0]
         __gm__ float *w_src = w_padded + r * W_PAD;
         __gm__ float *w_dst_local = recv_w_local + row * W_PAD;
         __gm__ float *w_dst_remote = CommRemotePtr(comm_ctx, w_dst_local, dst);
 
-#if defined(__CPU_SIM)
-        for (int i = 0; i < W_PAD; ++i) {
-            w_dst_remote[i] = w_src[i];
-        }
-#else
         WGlobal w_src_g(w_src);
         WGlobal w_dst_g(w_dst_remote);
         pto::comm::TPUT(w_dst_g, w_src_g, w_tile);
-#endif
 
         // Channel 3: idx (INT32, 1xIDX_PAD); host pre-packed [r, 0, …, 0]
         __gm__ int32_t *idx_src = idx_padded + r * IDX_PAD;
         __gm__ int32_t *idx_dst_local = recv_idx_local + row * IDX_PAD;
         __gm__ int32_t *idx_dst_remote = CommRemotePtr(comm_ctx, idx_dst_local, dst);
 
-#if defined(__CPU_SIM)
-        for (int i = 0; i < IDX_PAD; ++i) {
-            idx_dst_remote[i] = idx_src[i];
-        }
-#else
         IGlobal idx_src_g(idx_src);
         IGlobal idx_dst_g(idx_dst_remote);
         pto::comm::TPUT(idx_dst_g, idx_src_g, idx_tile);
-#endif
     }
     pipe_barrier(PIPE_ALL);
 

examples/workers/l3/ffn_tp_parallel/kernels/aiv/kernel_allreduce_sum.cpp

@@ -107,14 +107,8 @@ extern "C" __aicore__ __attribute__((always_inline)) void kernel_entry(__gm__ in
         }
         __gm__ float *remote_mailbox_base = CommRemotePtr(comm_ctx, mailbox_ptr, peer);
         __gm__ float *remote_slot_ptr = remote_mailbox_base + my_rank * kElemsPerPartial;
-#if defined(__CPU_SIM)
-        for (int i = 0; i < kElemsPerPartial; ++i) {
-            remote_slot_ptr[i] = partial_local_ptr[i];
-        }
-#else
         MatrixGlobal remote_slot(remote_slot_ptr);
         pto::comm::TPUT(remote_slot, partial_local_global, staging_tile);
-#endif
     }
     pipe_barrier(PIPE_ALL);