Sync: host_build_graph early-dispatch up to current tensormap_and_ringbuffer

src/a2a3/runtime/host_build_graph/aicore/aicore_executor.cpp

@@ -147,11 +147,11 @@ __aicore__ __attribute__((weak)) void aicore_execute(__gm__ Runtime *runtime, in
             // critical-path prep (receive_time → start_time, software-tunable).
             // Stored in the record as a 32-bit delta `start_time - receive_time`.
             //
-            // For the common path (not_ready == 0) the new task_id on
+            // For the common path (src_payload == 0) the new task_id on
             // DATA_MAIN_BASE is itself the ready signal, so receive_time is
             // stamped immediately and local_setup covers dcci + ack.
             //
-            // For the speculative early-dispatch path (not_ready == 1) the
+            // For the speculative early-dispatch path (src_payload != 0) the
             // dcci ran BEFORE the dependency-wait spin, so its cost is hidden
             // behind the doorbell-wait — not on the critical path between
             // "task genuinely ready" and "kernel begins". receive_time is
@@ -178,8 +178,32 @@ __aicore__ __attribute__((weak)) void aicore_execute(__gm__ Runtime *runtime, in
             // while the task is gated — preventing a real task from being
             // dual-issued behind it. The kernel's own input dcci runs inside
             // execute_task() below — strictly AFTER this gate — so predecessor
-            // outputs are visible. not_ready == 0 (the common path) skips this.
-            if (exec_payload->not_ready) {
+            // outputs are visible. src_payload == 0 (the common path) skips this;
+            // a non-zero src_payload is both the gate flag and the source
+            // PTO2TaskPayload.
+            if (exec_payload->src_payload != 0) {
+                // AICPU staged only src_payload, not the arg vector — fill
+                // args[0..num_args) ourselves now, while we are idle waiting for
+                // the doorbell. The whole-cache dcci(ENTIRE_DATA_CACHE) above
+                // already invalidated src's lines, so tensor_count/scalar_count/
+                // scalars read coherently with the orchestrator's submit writes.
+                // args[SPMD_LOCAL_CONTEXT_INDEX]/[SPMD_GLOBAL_CONTEXT_INDEX] are
+                // still written by the AICPU (num_args <= 48 never reaches them).
+                __gm__ char *src = reinterpret_cast<__gm__ char *>(exec_payload->src_payload);
+                int32_t tensor_count = *reinterpret_cast<__gm__ int32_t *>(src + PTO2_TASKPAYLOAD_TENSOR_COUNT_OFFSET);
+                int32_t scalar_count = *reinterpret_cast<__gm__ int32_t *>(src + PTO2_TASKPAYLOAD_SCALAR_COUNT_OFFSET);
+                __gm__ uint64_t *src_scalars =
+                    reinterpret_cast<__gm__ uint64_t *>(src + PTO2_TASKPAYLOAD_SCALARS_OFFSET);
+                int n = 0;
+                for (int32_t i = 0; i < tensor_count; i++) {
+                    exec_payload->args[n++] = reinterpret_cast<uint64_t>(
+                        src + PTO2_TASKPAYLOAD_TENSORS_OFFSET + i * PTO2_TASKPAYLOAD_TENSOR_STRIDE
+                    );
+                }
+                for (int32_t i = 0; i < scalar_count; i++) {
+                    exec_payload->args[n++] = src_scalars[i];
+                }
+                OUT_OF_ORDER_STORE_BARRIER();
                 while (true) {
                     // Honor teardown: shutdown overwrites the low half with EXIT.
                     // Check it on the doorbell-match iteration too, so an EXIT that

src/a2a3/runtime/host_build_graph/runtime/pto2_dispatch_payload.h

@@ -34,6 +34,7 @@
 
 #pragma once
 
+#include <stddef.h>
 #include <stdint.h>
 
 #include "arg_direction.h"
@@ -57,6 +58,18 @@ static_assert(
     "GLOBAL_CONTEXT_INDEX out of sync with intrinsic.h"
 );
 
+// Byte offsets into PTO2TaskPayload used by AICore to materialize args[] on the
+// not_ready (early-dispatch) path. The AICore .o does not include
+// pto_runtime2_types.h, so it reads tensor_count / scalar_count / tensors[] /
+// scalars[] through these constants. The AICPU side (scheduler_dispatch.cpp)
+// static_asserts them against offsetof where the full struct is visible, so any
+// PTO2TaskPayload layout drift fails the build rather than corrupting args[].
+constexpr uint32_t PTO2_TASKPAYLOAD_TENSOR_COUNT_OFFSET = 0;
+constexpr uint32_t PTO2_TASKPAYLOAD_SCALAR_COUNT_OFFSET = 4;
+constexpr uint32_t PTO2_TASKPAYLOAD_TENSORS_OFFSET = 576;
+constexpr uint32_t PTO2_TASKPAYLOAD_SCALARS_OFFSET = 4672;
+constexpr uint32_t PTO2_TASKPAYLOAD_TENSOR_STRIDE = 128;  // sizeof(Tensor)
+
 /**
  * Per-core dispatch payload: function address + args[] + SPMD context.
  *
@@ -68,24 +81,40 @@ static_assert(
  * concurrently dispatched cores.
  */
 struct alignas(64) PTO2DispatchPayload {
-    uint64_t function_bin_addr;            /**< Kernel entry address in GM (set by Scheduler) */
-    uint64_t args[PTO2_DISPATCH_MAX_ARGS]; /**< Kernel arguments (GM pointers + scalars + ext params) */
+    // === Cache line 0 (64B): control block, the only line written per dispatch ===
+    // function_bin_addr, local_context.{block_idx,block_num,async_ctx.task_token}
+    // and src_payload are the per-dispatch writes; async_ctx's slab pointers +
+    // capacity are cold (prefilled once at init) but ride this hot line for free.
+    // Sized to exactly 64B so both dispatch paths write one control line: the
+    // ready path (src_payload = 0) then also fills args[0..num_args); the gated
+    // path (src_payload = &PTO2TaskPayload) leaves args[] to the idle AICore.
+    uint64_t function_bin_addr; /**< Kernel entry address in GM (set by Scheduler). */
 
-    /** Per-dispatch context: block_idx and block_num.
-     *  Written by build_payload() before each dispatch.
-     *  args[SPMD_LOCAL_CONTEXT_INDEX] points here. */
+    /** Per-dispatch context: block_idx/block_num (hot) + async_ctx (task_token hot,
+     *  slab pointers + capacity prefilled once at init). args[SPMD_LOCAL_CONTEXT_INDEX]
+     *  points here. */
     LocalContext local_context;
 
-    /** Per-core global context: sub_block_id (AIV lane identity).
-     *  Initialized once by init_global_context() at runtime startup.
-     *  args[SPMD_GLOBAL_CONTEXT_INDEX] points here. */
-    GlobalContext global_context;
+    /** Early-dispatch gate AND source pointer, folded into one field. 0 = ready:
+     *  AICore executes on pickup (args[] already filled by the AICPU). Non-zero =
+     *  gated: the value is the source PTO2TaskPayload address; the AICore fills
+     *  args[0..num_args) from it, then waits for the doorbell (DATA_MAIN_BASE
+     *  high 32 == this dispatch's reg_task_id) before executing. A payload address
+     *  is never 0, so the gate flag is lossless. */
+    volatile uint64_t src_payload;
 
-    /** Speculative early-dispatch gate. 0 = ready: AICore executes on pickup.
-     *  1 = not-ready: AICore waits until AICPU rings the doorbell
-     *  (DATA_MAIN_BASE high 32 == this dispatch's reg_task_id) before executing. */
-    volatile uint32_t not_ready;
-    uint8_t reserved_payload_abi_pad[4];
+    // === Cache lines 1..7: kernel argument vector ===
+    /** [0..num_args) = GM tensor pointers + scalar values; [SPMD_LOCAL_CONTEXT_INDEX]
+     *  = &local_context, [SPMD_GLOBAL_CONTEXT_INDEX] = &global_context (both prefilled
+     *  once at init). On the gated path the AICPU leaves [0..num_args) unwritten and
+     *  the idle AICore fills them from src_payload during its gate wait. */
+    uint64_t args[PTO2_DISPATCH_MAX_ARGS];
+
+    /** Per-core global context: sub_block_id (AIV lane identity). Cold: written once
+     *  at init, never per dispatch, so it lives in the tail (not the CL0 control
+     *  block). args[SPMD_GLOBAL_CONTEXT_INDEX] points here. */
+    GlobalContext global_context;
+    // No explicit tail padding: alignas(64) rounds sizeof up to 512 (8 cache lines).
 
     static_assert(sizeof(args[0]) == 8);
     static_assert(
@@ -95,3 +124,5 @@ struct alignas(64) PTO2DispatchPayload {
 };
 
 static_assert(sizeof(PTO2DispatchPayload) == 512, "PTO2DispatchPayload hardware ABI size drift");
+static_assert(offsetof(PTO2DispatchPayload, args) == 64, "args[] must start at cache line 1 (control block = CL0)");
+static_assert(offsetof(PTO2DispatchPayload, src_payload) < 64, "src_payload (gate) must live on the CL0 control block");

src/a2a3/runtime/host_build_graph/runtime/pto_runtime2_types.h

@@ -224,8 +224,22 @@ struct PTO2TaskDescriptor {
 enum PTO2EarlyDispatchState : uint8_t {
     PTO2_EARLY_DISPATCH_NONE = 0,       // not pre-staged
     PTO2_EARLY_DISPATCH_STAGING = 1,    // Hook 1 claimed it; staging in progress
-    PTO2_EARLY_DISPATCH_STAGED = 2,     // staged on a core, gated; staged_* fields valid
-    PTO2_EARLY_DISPATCH_DISPATCHED = 3  // routed via the normal dispatch path (no pre-stage)
+    PTO2_EARLY_DISPATCH_STAGED = 2,     // reserved
+    PTO2_EARLY_DISPATCH_DISPATCHED = 3  // producers released; staged blocks may still be gated
+};
+
+enum PTO2EarlyDispatchLaunchState : uint8_t {
+    PTO2_EARLY_DISPATCH_LAUNCH_NONE = 0,
+    PTO2_EARLY_DISPATCH_LAUNCH_RINGING = 1,
+    PTO2_EARLY_DISPATCH_LAUNCH_COMPLETE = 2,
+};
+
+enum PTO2EarlySyncDrainState : uint8_t {
+    PTO2_EARLY_SYNC_DRAIN_NONE = 0,
+    PTO2_EARLY_SYNC_DRAIN_OWNER = 1 << 0,
+    PTO2_EARLY_SYNC_DRAIN_ARMED = 1 << 1,
+    PTO2_EARLY_SYNC_DRAIN_READY = 1 << 2,
+    PTO2_EARLY_SYNC_DRAIN_COMPLETE = 1 << 3,
 };
 
 // A pre-staged consumer occupies one core per gated subtask block. WHICH cores
@@ -246,34 +260,56 @@ struct PTO2TaskPayload {
     PTO2FaninPool *fanin_spill_pool{nullptr};
     PTO2TaskSlotState *fanin_inline_slot_states[PTO2_FANIN_INLINE_CAP];
     // Early-dispatch metadata (AICPU-side only). Ordered by descending
-    // alignment (8B mask, 4B fanin, then 1B flags) so the block packs with no
+    // alignment (8B mask, 4B fanin, then 2B/1B counters and flags) so the block packs with no
     // internal padding. Kept here after the fanin array (not moved up front): on
     // cache line 8 it shares only with the rarely-touched fanin tail, whereas in
     // line 0 the early-dispatch atomics (written during staging) would false-share with
     // tensor_count/scalar_count (read by build_payload at dispatch). Fits in the 40B
     // between the fanin array (offset 536) and the 64B-aligned tensors[] (offset
     // 576), so sizeof and tensors[] are unchanged.
     //
-    // Bitmask of global core_ids this consumer is pre-staged (gated) on. Set with
-    // atomic fetch_or by concurrent stagers; read by release. (Re)initialized in
-    // PTO2TaskPayload::init before the slot can be staged again.
+    // Bitmask of global core_ids this consumer is pre-staged (gated) on. Concurrent
+    // stagers publish bits with atomic fetch_or. A regular consumer destructively
+    // splits them between release and late-stager owners; a sync_start drain keeps
+    // the completed mask stable for its single cohort launch owner.
     std::atomic<uint64_t> staged_core_mask[PTO2_EARLY_DISPATCH_CORE_MASK_WORDS]{};
     // Early-dispatch CANDIDATE detection (event-driven, dual of fanin_refcount):
-    // seeded at wiring with producers already complete, then a flagged producer's
-    // DISPATCH bumps each consumer's dispatch_fanin. dispatch_fanin ==
-    // fanin_actual_count  <=>  every producer is flagged-and-dispatched or was
+    // seeded at wiring with producers already complete, then a flagged producer
+    // bumps each consumer after all of its logical blocks are published.
+    // dispatch_fanin == fanin_actual_count  <=>  every producer is
+    // flagged-and-fully-published or was
     // pre-completed  =>  this task is an early-dispatch candidate (push early_dispatch_queues[shape]).
-    std::atomic<int32_t> dispatch_fanin{0};  // CONSUMER side: flagged-dispatched + pre-completed producers
+    std::atomic<int32_t> dispatch_fanin{0};  // CONSUMER side: fully-published + pre-completed producers
+    // Number of logical blocks whose payloads and MMIO tokens are published.
+    // Claimed-but-unpublished blocks do not make a producer launch-visible. Its
+    // seq_cst updates pair with early_dispatch_state to avoid losing the final
+    // publish vs. release wakeup for a pre-staged producer.
+    std::atomic<int16_t> published_block_count{0};
     // Lock-free claim state shared by the stagers (Hook 1, possibly several AICPU
     // threads concurrently) and the completion-path release: 0=NONE, 1=STAGING,
     // 3=DISPATCHED (2=STAGED is unused now). STAGING is the STABLE gated state —
     // many threads stage blocks concurrently while it holds, each claiming a block
     // via the atomic next_block_idx and OR-ing its cores into staged_core_mask.
-    // Release does STAGING->DISPATCHED then rings the mask; a thread that stages a
-    // block AFTER release flipped DISPATCHED rings that block's doorbell itself
-    // (self-ring), so no doorbell is ever missed.
+    // Release does STAGING->DISPATCHED. For a regular consumer it claims the current
+    // mask and a late stager rings only its remaining bits. A sync_start consumer
+    // preserves the mask for rendezvous counting and its single launch pass.
     std::atomic<uint8_t> early_dispatch_state{0};
     std::atomic<uint8_t> dispatch_propagated{0};  // PRODUCER side: once-guard for fanout propagation
+    // The launch owner publishes COMPLETE only after all owned doorbells are
+    // visible, keeping fanout private until every gated block has launched.
+    std::atomic<uint8_t> early_dispatch_launch_state{PTO2_EARLY_DISPATCH_LAUNCH_NONE};
+    // sync_start early-dispatch rendezvous: count of this task's gated CORES currently
+    // occupying a RUNNING slot (staged directly to an idle core, or promoted from a
+    // gated pending slot). Counted per-core (not per-block) so it is shape-agnostic: a
+    // MIX block spans a cluster whose cores promote independently. A sync_start task's
+    // doorbells are rung only once this reaches popcount(staged_core_mask) AND the
+    // producer released, so all cores launch atomically. Unused (0) for non-sync_start.
+    std::atomic<int16_t> running_slot_count{0};
+    // Ownership handshake between the early sync queue and final ready routing.
+    // A successful OWNER persists through ARMED and COMPLETE until payload
+    // reinitialization. READY records that producer release observed OWNER;
+    // only cancellation clears OWNER during the current task lifetime.
+    std::atomic<uint8_t> early_sync_drain_state{PTO2_EARLY_SYNC_DRAIN_NONE};
     // === Cache lines 9-72 (4096B) — tensors (alignas(64) forces alignment) ===
     Tensor tensors[MAX_TENSOR_ARGS];
     // === Cache lines 73-74 (128B) — scalars ===
@@ -349,11 +385,17 @@ struct PTO2TaskPayload {
         // one of ITS producers is flagged (propagate_dispatch_fanin bumps
         // dispatch_fanin and may CAS early_dispatch_state on any consumer, independent of the
         // consumer's own hint). So they MUST be zeroed here unconditionally.
+        // Publication, propagation, and launch fields share this same
+        // per-submit lifetime and are reset here too.
         early_dispatch_state.store(PTO2_EARLY_DISPATCH_NONE, std::memory_order_relaxed);
         for (int w = 0; w < PTO2_EARLY_DISPATCH_CORE_MASK_WORDS; w++)
             staged_core_mask[w].store(0, std::memory_order_relaxed);
         dispatch_fanin.store(0, std::memory_order_relaxed);
         dispatch_propagated.store(0, std::memory_order_relaxed);
+        published_block_count.store(0, std::memory_order_relaxed);
+        early_dispatch_launch_state.store(PTO2_EARLY_DISPATCH_LAUNCH_NONE, std::memory_order_relaxed);
+        running_slot_count.store(0, std::memory_order_relaxed);
+        early_sync_drain_state.store(PTO2_EARLY_SYNC_DRAIN_NONE, std::memory_order_relaxed);
     }
 };
 
@@ -453,12 +495,28 @@ struct alignas(64) PTO2TaskSlotState {
     std::atomic<int16_t> completed_subtasks{0};  // Each core completion increments by 1
     int16_t total_required_subtasks{0};          // = logical_block_num * popcount(active_mask)
     int16_t logical_block_num{1};                // Total logical blocks (set by orchestrator)
-    // Next block to dispatch. Atomic so concurrent early-dispatch stagers can each
-    // claim a distinct block via CAS; normal dispatch (ready-queue serialized)
-    // uses plain relaxed load/store. The two phases never overlap in time (staging
-    // happens before release; normal dispatch of the remainder happens after).
+    // Next block to dispatch. Normal dispatch and late early-dispatch stagers
+    // can run concurrently after a partial staged release. All paths claim
+    // ranges through claim_block_range().
     std::atomic<int16_t> next_block_idx{0};
 
+    int32_t claim_block_range(int32_t block_limit, int32_t max_count, int32_t &start) {
+        int16_t current = next_block_idx.load(std::memory_order_relaxed);
+        while (current < block_limit && max_count > 0) {
+            int32_t count = block_limit - current;
+            if (count > max_count) count = max_count;
+            int16_t desired = static_cast<int16_t>(current + count);
+            if (next_block_idx.compare_exchange_weak(
+                    current, desired, std::memory_order_seq_cst, std::memory_order_relaxed
+                )) {
+                start = current;
+                return count;
+            }
+        }
+        start = current;
+        return 0;
+    }
+
     /**
      * Re-bind the per-slot payload/task pointers. Called by
      * orch::prepare_task on every submit. Value is constant for a given
@@ -515,7 +573,7 @@ struct alignas(64) PTO2TaskSlotState {
         next_block_idx.store(0, std::memory_order_relaxed);
         ready_state.store(PTO2_READY_UNCLAIMED, std::memory_order_relaxed);
         allow_early_resolve = false;
-        // Note: payload early-dispatch fields (early_dispatch_state / staged_core_mask / dispatch_fanin)
+        // Note: payload early-dispatch fields (state, masks, fanin, publication count)
         // are NOT reset here — this method skips the payload by contract. They are
         // (re)initialized in PTO2TaskPayload::init on every submit, before the slot
         // becomes visible to the scheduler.