diff --git a/src/a2a3/runtime/host_build_graph/aicore/aicore_executor.cpp b/src/a2a3/runtime/host_build_graph/aicore/aicore_executor.cpp index 0cb67144b..c5f09a56b 100644 --- a/src/a2a3/runtime/host_build_graph/aicore/aicore_executor.cpp +++ b/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( + 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 diff --git a/src/a2a3/runtime/host_build_graph/runtime/pto2_dispatch_payload.h b/src/a2a3/runtime/host_build_graph/runtime/pto2_dispatch_payload.h index e1bb3465e..899de63f0 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/pto2_dispatch_payload.h +++ b/src/a2a3/runtime/host_build_graph/runtime/pto2_dispatch_payload.h @@ -34,6 +34,7 @@ #pragma once +#include #include #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"); diff --git a/src/a2a3/runtime/host_build_graph/runtime/pto_runtime2_types.h b/src/a2a3/runtime/host_build_graph/runtime/pto_runtime2_types.h index 94f510656..5bd2fd598 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/pto_runtime2_types.h +++ b/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,7 +260,7 @@ 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 @@ -254,26 +268,48 @@ struct PTO2TaskPayload { // 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 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 dispatch_fanin{0}; // CONSUMER side: flagged-dispatched + pre-completed producers + std::atomic 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 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 early_dispatch_state{0}; std::atomic 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 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 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 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 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 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(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. diff --git a/src/a2a3/runtime/host_build_graph/runtime/scheduler/pto_scheduler.h b/src/a2a3/runtime/host_build_graph/runtime/scheduler/pto_scheduler.h index 00fd72395..6ac8d7f59 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/scheduler/pto_scheduler.h +++ b/src/a2a3/runtime/host_build_graph/runtime/scheduler/pto_scheduler.h @@ -32,6 +32,7 @@ #include #include "common/core_type.h" +#include "common/memory_barrier.h" #include "utils/device_arena.h" #include "aicpu/platform_regs.h" // get_reg_ptr / RegId for the early-dispatch doorbell #include "pto_async_wait.h" @@ -404,6 +405,7 @@ struct PTO2SchedulerLayout { size_t off_ready_sync_queue_slots[PTO2_NUM_RESOURCE_SHAPES]; size_t off_dummy_ready_queue_slots; size_t off_early_dispatch_queue_slots[PTO2_NUM_RESOURCE_SHAPES]; + size_t off_early_sync_start_queue_slots; size_t off_dep_pool_entries; uint64_t ready_queue_capacity; int32_t dep_pool_capacity; @@ -554,8 +556,8 @@ struct PTO2SchedulerState { // PARTIALLY pre-staged — the gated blocks are released by the doorbells rung // here, and the remaining (next_block_idx .. logical_block_num) blocks // dispatch normally off the ready queue. Lock-free claim shared with Hook 1 - // (the stager): CAS NONE->DISPATCHED wins => not pre-staged; lose => STAGED - // (spin past the brief STAGING window so the mask is visible), then ring. + // (the stager): CAS NONE->DISPATCHED wins => not pre-staged; otherwise flip + // STAGING->DISPATCHED and destructively claim the published doorbell bits. // Per-core early-dispatch doorbell table. Hook 1 records each gated core's // (reg_addr, dispatch token) here at stage time; the completion-path release @@ -589,23 +591,188 @@ struct PTO2SchedulerState { // logged and the consumer's blocks fall back to normal dispatch. PTO2ReadyQueue early_dispatch_queues[PTO2_NUM_RESOURCE_SHAPES]; + // sync_start early-dispatch candidates park here instead of early_dispatch_queues[]: + // they need an atomic all-or-nothing stage via the drain barrier, not + // early_dispatch_shape's per-thread partial range-claim. Shape-agnostic (the + // rendezvous counts cores, not blocks), so a single queue serves all shapes; drained + // as the highest occupancy tier at the top of try_early_dispatch. + // + // Deliberately single, vs the normal source's per-shape ready_sync_queues[]: a READY + // sync cohort (producer done) can dispatch inline when it fits, so it reuses the + // per-shape dispatch_shape; an EARLY sync cohort always carries a non-zero + // src_payload gate and therefore always drains. The shape-agnostic rendezvous + // makes one queue sufficient. Same drain, two sources. + PTO2ReadyQueue early_sync_start_queue; + static inline void ring_one_doorbell(uint64_t reg_addr, uint32_t token) { volatile uint64_t *dmb = reinterpret_cast(get_reg_ptr(reg_addr, RegId::DATA_MAIN_BASE)); uint64_t tk = static_cast(token); *dmb = (tk << 32) | tk; // 64-bit STR: high=low=token releases the gated AICore } - // Event-driven candidate detection (the dual of fanin_refcount/ready). Call when a - // FLAGGED producer `p` DISPATCHES (starts running): walk its fanout and bump each - // consumer's dispatch_fanin. A consumer whose dispatch_fanin reaches - // fanin_actual_count (= every producer is either flagged-and-dispatched, or was - // already complete when the consumer was wired) is an early-dispatch candidate: - // CAS NONE->STAGING (exactly-once) and push to early_dispatch_queues[shape] for the idle drain to - // pre-stage. Once-guarded per producer so an SPMD producer's block-by-block - // dispatch propagates once. Only codegen-flagged producers propagate: a task's - // successors early-dispatch off its DIRECT producers' marks, never an inherited chain. + inline void ring_staged_doorbell_bits(int word, uint64_t bits) { + while (bits != 0) { + int core_id = word * 64 + __builtin_ctzll(bits); + bits &= bits - 1; + ring_one_doorbell( + early_dispatch_doorbell_table[core_id].addr, early_dispatch_doorbell_table[core_id].token + ); + } + } + + static inline uint64_t claim_all_staged_doorbell_bits(std::atomic &mask) { + return mask.exchange(0, std::memory_order_seq_cst); + } + + static inline uint64_t claim_late_staged_doorbell_bits(std::atomic &mask, uint64_t candidates) { + return mask.fetch_and(~candidates, std::memory_order_seq_cst) & candidates; + } + + static inline bool should_gate_early_dispatch(bool force_gate, uint8_t early_dispatch_state) { + return force_gate || early_dispatch_state == PTO2_EARLY_DISPATCH_STAGING; + } + + static inline bool + ring_claimed_local_doorbell(uint64_t claimed_word, int core_id, uint64_t reg_addr, uint32_t token) { + if ((claimed_word & (1ULL << (core_id & 63))) == 0) return false; + ring_one_doorbell(reg_addr, token); + return true; + } + + static inline bool try_claim_early_dispatch_launch(PTO2TaskPayload &payload) { + uint8_t expected = PTO2_EARLY_DISPATCH_LAUNCH_NONE; + return payload.early_dispatch_launch_state.compare_exchange_strong( + expected, PTO2_EARLY_DISPATCH_LAUNCH_RINGING, std::memory_order_seq_cst, std::memory_order_seq_cst + ); + } + + inline void record_published_blocks(PTO2TaskSlotState &slot_state, int32_t count) { + if (count <= 0 || !slot_state.allow_early_resolve) return; + slot_state.payload->published_block_count.fetch_add(static_cast(count), std::memory_order_seq_cst); + } + + // Ring one sync_start cohort from its stable staged_core_mask. The caller owns + // the NONE->RINGING launch latch and invokes this exactly once after drain + // staging completes, while the corresponding per-core table entries are live. + inline void ring_all_staged_doorbells(PTO2TaskSlotState &slot_state) { + for (int w = 0; w < PTO2_EARLY_DISPATCH_CORE_MASK_WORDS; w++) { + uint64_t bits = slot_state.payload->staged_core_mask[w].load(std::memory_order_seq_cst); + while (bits != 0) { + int core_id = w * 64 + __builtin_ctzll(bits); + bits &= bits - 1; + ring_one_doorbell( + early_dispatch_doorbell_table[core_id].addr, early_dispatch_doorbell_table[core_id].token + ); + } + } + } + + static inline bool try_claim_early_sync_drain(PTO2TaskPayload &payload) { + uint8_t expected = PTO2_EARLY_SYNC_DRAIN_NONE; + return payload.early_sync_drain_state.compare_exchange_strong( + expected, PTO2_EARLY_SYNC_DRAIN_OWNER, std::memory_order_seq_cst, std::memory_order_seq_cst + ); + } + + static inline bool owns_early_sync_drain(const PTO2TaskPayload &payload) { + return (payload.early_sync_drain_state.load(std::memory_order_acquire) & PTO2_EARLY_SYNC_DRAIN_OWNER) != 0; + } + + static inline void mark_early_sync_drain_armed(PTO2TaskPayload &payload) { + payload.early_sync_drain_state.fetch_or(PTO2_EARLY_SYNC_DRAIN_ARMED, std::memory_order_seq_cst); + } + + static inline bool publish_ready_to_early_sync_drain(PTO2TaskPayload &payload) { + uint8_t previous = + payload.early_sync_drain_state.fetch_or(PTO2_EARLY_SYNC_DRAIN_READY, std::memory_order_seq_cst); + return (previous & PTO2_EARLY_SYNC_DRAIN_OWNER) != 0; + } + + inline void cancel_early_sync_drain(PTO2TaskSlotState &slot_state) { + uint8_t previous = + slot_state.payload->early_sync_drain_state.exchange(PTO2_EARLY_SYNC_DRAIN_NONE, std::memory_order_seq_cst); + if ((previous & PTO2_EARLY_SYNC_DRAIN_OWNER) == 0) return; + if ((previous & PTO2_EARLY_SYNC_DRAIN_READY) != 0) { + push_ready_routed(&slot_state); + return; + } + if (slot_state.payload->early_dispatch_state.load(std::memory_order_seq_cst) == PTO2_EARLY_DISPATCH_STAGING) { + early_sync_start_queue.push(&slot_state); + } + } + + static inline void finish_early_sync_drain(PTO2TaskPayload &payload) { + uint8_t state = payload.early_sync_drain_state.load(std::memory_order_seq_cst); + while ((state & PTO2_EARLY_SYNC_DRAIN_OWNER) != 0 && (state & PTO2_EARLY_SYNC_DRAIN_COMPLETE) == 0) { + uint8_t desired = state | PTO2_EARLY_SYNC_DRAIN_COMPLETE; + if (payload.early_sync_drain_state.compare_exchange_weak( + state, desired, std::memory_order_seq_cst, std::memory_order_seq_cst + )) { + return; + } + } + } + // sync_start rendezvous: a sync_start consumer's gated cores launch as an atomic + // cohort, so their doorbells are held until BOTH halves hold — every gated core + // occupies a running slot (running_slot_count == popcount(staged_core_mask)) AND the + // producer released (early_dispatch_state == DISPATCHED). Counting CORES (not blocks) makes + // this shape-agnostic: an AIC/AIV block is one core, a MIX block is a cluster whose + // cores promote pending->running independently. Called from both halves (the producer + // release and each pending->running promotion); whichever observes the second half + // wins the launch latch and rings exactly once. Returns true only to that winner, + // which may then expose the cohort to its fanout. + inline bool maybe_rendezvous_ring(PTO2TaskSlotState &slot_state) { + int32_t staged_cores = 0; + for (int w = 0; w < PTO2_EARLY_DISPATCH_CORE_MASK_WORDS; w++) + staged_cores += + __builtin_popcountll(slot_state.payload->staged_core_mask[w].load(std::memory_order_seq_cst)); + if (staged_cores == 0) return false; + if (slot_state.payload->running_slot_count.load(std::memory_order_seq_cst) != staged_cores) return false; + if (slot_state.payload->early_dispatch_state.load(std::memory_order_seq_cst) != PTO2_EARLY_DISPATCH_DISPATCHED) + return false; + if (!try_claim_early_dispatch_launch(*slot_state.payload)) return false; + ring_all_staged_doorbells(slot_state); + wmb(); + slot_state.payload->early_dispatch_launch_state.store( + PTO2_EARLY_DISPATCH_LAUNCH_COMPLETE, std::memory_order_release + ); + return true; + } + + inline bool retry_sync_start_rendezvous_after_drain(PTO2TaskSlotState &slot_state) { + if (!maybe_rendezvous_ring(slot_state)) return false; + propagate_dispatch_fanin(slot_state); + return true; + } + + // Event-driven candidate detection (the dual of fanin_refcount/ready). Called after a + // FLAGGED producer `p` publishes blocks (normal dispatch, early-dispatch release, or the + // sync_start drain), but no-ops until every logical block is launch-visible. Only then + // does it walk p's fanout and bump each consumer's + // dispatch_fanin. A consumer whose dispatch_fanin reaches fanin_actual_count (= every + // producer is flagged-and-fully-dispatched, or was already complete when the consumer was + // wired) is an early-dispatch candidate: CAS NONE->STAGING (exactly-once) and push to + // early_dispatch_queues[shape] (or early_sync_start_queue for a require_sync_start cohort) + // for the drain to pre-stage. The full-publication gate is load-bearing: a consumer that + // pre-occupies cores off a producer whose later blocks are not yet launch-visible would gate the + // very cores those blocks need, starving the producer so it never completes and the + // rendezvous never rings — a resource deadlock. published_block_count advances after + // every placement path publishes its claimed range. Once-guarded per producer. + // Only codegen-flagged producers propagate: a task's successors early-dispatch off its + // DIRECT producers' marks, never an inherited chain. void propagate_dispatch_fanin(PTO2TaskSlotState &p) { if (!p.allow_early_resolve) return; // only codegen-flagged (direct) producers propagate + if (p.payload->published_block_count.load(std::memory_order_seq_cst) < p.logical_block_num) return; + if (p.payload->early_dispatch_state.load(std::memory_order_acquire) == PTO2_EARLY_DISPATCH_STAGING) return; + uint8_t launch_state = p.payload->early_dispatch_launch_state.load(std::memory_order_acquire); + if (launch_state == PTO2_EARLY_DISPATCH_LAUNCH_RINGING) return; + if (p.active_mask.requires_sync_start()) { + bool was_pre_staged = false; + for (int w = 0; w < PTO2_EARLY_DISPATCH_CORE_MASK_WORDS; w++) { + was_pre_staged |= p.payload->staged_core_mask[w].load(std::memory_order_acquire) != 0; + } + if (was_pre_staged && launch_state != PTO2_EARLY_DISPATCH_LAUNCH_COMPLETE) return; + } if (p.payload->dispatch_propagated.exchange(1, std::memory_order_acq_rel) != 0) return; // already propagated once p.lock_fanout(); @@ -618,20 +785,30 @@ struct PTO2SchedulerState { // ready_fanin gets but dispatch_fanin does not). dispatch_fanin starts at // the wiring-time flagged-pre-completed seed and is bumped here by flagged // producers; reaching fanin_actual_count means every producer is - // flagged-dispatched or was pre-completed. An unflagged producer leaves the + // flagged-and-fully-published or was pre-completed. An unflagged producer leaves the // seed short and never bumps, so this stays unreachable for that consumer. int32_t nf = c->payload->dispatch_fanin.fetch_add(1, std::memory_order_acq_rel) + 1; if (nf != c->payload->fanin_actual_count) continue; - if (c->active_mask.requires_sync_start()) continue; // sync_start can't be block-by-block pre-staged PTO2ResourceShape shape = c->active_mask.to_shape(); if (shape != PTO2ResourceShape::AIC && shape != PTO2ResourceShape::AIV && shape != PTO2ResourceShape::MIX) continue; + // sync_start blocks launch as an atomic cohort. They ride the early-dispatch path + // too, but through a dedicated shape-agnostic queue (early_sync_start_queue), + // drained as the highest tier in try_early_dispatch via the gated drain + + // running-slot rendezvous (enter_drain_mode pre-stages every block gated — MIX + // with per-core split placement — and the rendezvous rings them together once + // every gated core occupies a running slot and the producer released). They must + // NOT enter early_dispatch_queues[shape]: that path's partial range-claim would + // strand gated cohort blocks nobody rings, so the rendezvous never reaches + // block_num. The rendezvous counts CORES (a MIX block spans a cluster), so one + // shape-agnostic queue suffices. uint8_t expect = PTO2_EARLY_DISPATCH_NONE; // exactly-once: only the CAS winner enqueues if (!c->payload->early_dispatch_state.compare_exchange_strong( expect, PTO2_EARLY_DISPATCH_STAGING, std::memory_order_seq_cst, std::memory_order_seq_cst )) continue; - early_dispatch_queues[static_cast(shape)].push(c); + if (c->active_mask.requires_sync_start()) early_sync_start_queue.push(c); + else early_dispatch_queues[static_cast(shape)].push(c); } } @@ -657,27 +834,37 @@ struct PTO2SchedulerState { )) { return false; } - // Staged (STAGING). Flip STAGING->DISPATCHED, THEN read the mask. seq_cst - // gives a total order with the concurrent stagers, each of which OR-s its - // core into the mask and THEN loads early_dispatch_state: a stager whose bit lands - // before this CAS is read here and rung; a stager whose bit lands after - // sees DISPATCHED and rings that core itself (self-ring in - // stage_consumer_blocks). Either way every gated core's doorbell fires once - // (a double-ring is harmless — the AICore already matched). This replaces - // the old transient-STAGING spin: STAGING is now the stable gated state. + // route_ready_once admits one caller, but keep the helper defensive: a + // duplicate that observes or loses an in-progress launch must never + // route the same partial task a second time. + if (expect != PTO2_EARLY_DISPATCH_STAGING) return true; + bool sync_start = slot_state.active_mask.requires_sync_start(); + if (!sync_start && !try_claim_early_dispatch_launch(*slot_state.payload)) return true; expect = PTO2_EARLY_DISPATCH_STAGING; slot_state.payload->early_dispatch_state.compare_exchange_strong( expect, PTO2_EARLY_DISPATCH_DISPATCHED, std::memory_order_seq_cst, std::memory_order_seq_cst ); - for (int w = 0; w < PTO2_EARLY_DISPATCH_CORE_MASK_WORDS; w++) { - uint64_t bits = slot_state.payload->staged_core_mask[w].load(std::memory_order_seq_cst); - while (bits != 0) { - int core_id = w * 64 + __builtin_ctzll(bits); - bits &= bits - 1; - ring_one_doorbell( - early_dispatch_doorbell_table[core_id].addr, early_dispatch_doorbell_table[core_id].token - ); + // Producer released: ring the gated cores. A non-sync_start consumer launches + // each block the instant its doorbell fires. A sync_start consumer instead holds + // for the rendezvous — every gated core must occupy a running slot first — so the + // flip to DISPATCHED above is only the producer-released half; maybe_rendezvous_ring + // rings now iff running_slot_count already reached popcount(staged_core_mask) (all + // gated cores took idle running slots), else the last pending->running promotion rings. + bool launched = true; + if (sync_start) { + launched = maybe_rendezvous_ring(slot_state); + } else { + // Destructively claim every published bit. A stager racing this pass + // can claim only bits that land after the exchange, so each gated core + // has exactly one doorbell writer. + for (int w = 0; w < PTO2_EARLY_DISPATCH_CORE_MASK_WORDS; w++) { + uint64_t owned = claim_all_staged_doorbell_bits(slot_state.payload->staged_core_mask[w]); + ring_staged_doorbell_bits(w, owned); } + wmb(); + slot_state.payload->early_dispatch_launch_state.store( + PTO2_EARLY_DISPATCH_LAUNCH_COMPLETE, std::memory_order_seq_cst + ); } // This pre-staged consumer was just released by its doorbell — it starts // running NOW, so propagate dispatch_fanin to ITS consumers (only if it is @@ -685,8 +872,10 @@ struct PTO2SchedulerState { // sink so it runs after the whole fanout walk: doing it inline here would // delay the doorbells of later consumers in the same producer's fanout. // Fallback to inline if no sink / sink full. - if (sink == nullptr || !sink->push(&slot_state)) { - propagate_dispatch_fanin(slot_state); + if (launched) { + if (sink == nullptr || !sink->push(&slot_state)) { + propagate_dispatch_fanin(slot_state); + } } // No explicit removal from the cross-thread queue: a still-queued entry for // this consumer is now DISPATCHED and is dropped when a peer pops it. @@ -700,7 +889,13 @@ struct PTO2SchedulerState { // Early-dispatch: pre-staged tasks are released by doorbell // here, skipping the ready-queue round-trip entirely. - if (try_early_dispatch_release(slot_state, sink)) return true; + bool early_handled = try_early_dispatch_release(slot_state, sink); + if (slot_state.active_mask.requires_sync_start()) { + bool drain_owned = publish_ready_to_early_sync_drain(*slot_state.payload); + if (early_handled || drain_owned) return true; + } else if (early_handled) { + return true; + } PTO2ResourceShape shape = slot_state.active_mask.to_shape(); if (shape == PTO2ResourceShape::DUMMY) { @@ -734,7 +929,13 @@ struct PTO2SchedulerState { // Early-dispatch: pre-staged tasks are released by doorbell // here, skipping the ready-queue round-trip entirely. - if (try_early_dispatch_release(slot_state, sink)) return true; + bool early_handled = try_early_dispatch_release(slot_state, sink); + if (slot_state.active_mask.requires_sync_start()) { + bool drain_owned = publish_ready_to_early_sync_drain(*slot_state.payload); + if (early_handled || drain_owned) return true; + } else if (early_handled) { + return true; + } PTO2ResourceShape shape = slot_state.active_mask.to_shape(); if (shape == PTO2ResourceShape::DUMMY) { diff --git a/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_cold_path.cpp b/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_cold_path.cpp index 594a70e7e..918305a2c 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_cold_path.cpp +++ b/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_cold_path.cpp @@ -933,6 +933,35 @@ int32_t SchedulerContext::post_handshake_init(Runtime *runtime) { } } + // Prefill the per-dispatch AsyncCtx constant fields once. Of AsyncCtx's five + // fields, four are constant for a given (core, buf_idx): the three pointers + // target the fixed deferred_slab_per_core_[core][buf] members, and capacity is + // MAX_COMPLETIONS_PER_TASK. Only task_token varies per dispatch, so build_payload + // writes just that; these constants survive across dispatches because the + // payload buffer is never zeroed between them. + // The two context-pointer args are also per-(core, buf_idx) constants — they + // target this buffer's own local_context / global_context — so prefill them + // here too and drop them from the per-dispatch build_payload writes. This keeps + // the per-dispatch write footprint on the CL0 control block only (args[48]/[49] + // live on a later line). + for (int32_t core_id = 0; core_id < RUNTIME_MAX_WORKER; core_id++) { + for (int32_t buf = 0; buf < 2; buf++) { + PTO2DispatchPayload &dp = payload_per_core_[core_id][buf]; + AsyncCtx &ac = dp.local_context.async_ctx; + volatile DeferredCompletionSlab *slab = &deferred_slab_per_core_[core_id][buf]; + ac.completion_count = &slab->count; + ac.completion_error_code = &slab->error_code; + ac.completion_entries = &slab->entries[0]; + ac.completion_capacity = MAX_COMPLETIONS_PER_TASK; + // Clear the slab once here; thereafter only the completion path re-clears + // count (and only when a deferred task dirtied it), never per dispatch. + slab->count = 0; + slab->error_code = PTO2_ERROR_NONE; + dp.args[PAYLOAD_LOCAL_CONTEXT_INDEX] = reinterpret_cast(&dp.local_context); + dp.args[PAYLOAD_GLOBAL_CONTEXT_INDEX] = reinterpret_cast(&dp.global_context); + } + } + func_id_to_addr_ = runtime->func_id_to_addr_; return 0; @@ -947,11 +976,14 @@ void SchedulerContext::deinit() { } // No per-core memset of payload_per_core_ / deferred_slab_per_core_ here - // (~300 KB across all cores). Both are fully re-initialized at dispatch - // before they can be read: dispatch_task sets deferred_slab->count = 0 / - // error_code = NONE and build_payload() overwrites every payload field - // (function addr, args[], contexts, not_ready) on the exact [core][buf_idx] - // about to run. The consumer side cannot reach a stale slot either: the + // (~300 KB across all cores). They are re-initialized before they can be read: + // build_payload() overwrites the per-dispatch payload fields (function addr, + // args[0..num_args) or src_payload, block_idx/block_num, async_ctx.task_token) + // on the exact [core][buf_idx] about to run; the async_ctx slab pointers + + // capacity, the two context-pointer args, and the deferred slab (count = 0 / + // error_code = NONE) are all cleared once per run in init() — the slab is + // thereafter re-cleared only by the completion path after a deferred task. + // The consumer side cannot reach a stale slot either: the // drain only services a core's running_reg_task_id, and the loop above // already reset every core_exec_states_[].running/pending_reg_task_id to // AICPU_TASK_INVALID — so no FIN for an undispatched slot is processed, and @@ -962,6 +994,9 @@ void SchedulerContext::deinit() { drain_state_.sync_start_pending.store(0, std::memory_order_release); drain_state_.drain_worker_elected.store(0, std::memory_order_release); drain_state_.drain_ack_mask.store(0, std::memory_order_release); + drain_state_.drain_stage_go.store(0, std::memory_order_release); + drain_state_.drain_stage_done_mask.store(0, std::memory_order_release); + drain_state_.drain_running_staged.store(0, std::memory_order_release); drain_state_.pending_task.store(nullptr, std::memory_order_release); // Reset task counters and orchestrator state diff --git a/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_completion.cpp b/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_completion.cpp index cfd9007b8..59b0b7b0d 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_completion.cpp +++ b/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_completion.cpp @@ -144,6 +144,11 @@ void SchedulerContext::complete_slot_task( SPIN_WAIT_HINT(); } } + // Re-clear for the next reuse of this (core, buf) slot. Done here — on + // the hot cache line we just read — instead of on every dispatch, since + // only a deferred task (count > 0) ever dirties it. error_code needs no + // reset: a non-NONE code aborted the run above. + deferred_slab->count = 0; } } @@ -338,14 +343,29 @@ void SchedulerContext::check_running_cores_for_completion( } #endif - // A pending task is "gated" when it is a early-dispatch pre-stage still - // waiting on its doorbell (STAGED): it will not ack on the producer's FIN, - // so the Case 3.1 wait-for-pending-ack shortcut would deadlock. Detect it - // so decide_slot_transition completes the running FIN and promotes it. + // A pending task is "gated" when it is an early-dispatch pre-stage still parked on + // its doorbell: it will not ack on the producer's FIN, so the Case 3.1 wait-for- + // pending-ack shortcut would deadlock. Detect it so decide_slot_transition completes + // the running FIN and PROMOTES it (Case 3.3) instead. + // + // "Gated" is "not yet launched (rung)", which is NOT the same as + // early_dispatch_state. STAGING covers the pre-release window. But a sync_start block + // stays gated even AFTER its producer releases (early_dispatch_state STAGING -> + // DISPATCHED): the cohort is not rung until the rendezvous has assembled EVERY core + // into a running slot, and this pending block (being promoted now) is by definition + // not yet counted, so the ring has not fired and it is still gated. Classifying it by + // STAGING alone would, once the producer's release beats the last promotion, treat it + // as a normal task and wait for an ack that never comes -> deadlock (the + // nondeterministic sync_start stall). + uint8_t pending_ss = + (core.pending_slot_state != nullptr && core.pending_slot_state->payload != nullptr) ? + core.pending_slot_state->payload->early_dispatch_state.load(std::memory_order_relaxed) : + static_cast(PTO2_EARLY_DISPATCH_NONE); bool pending_gated = (core.pending_slot_state != nullptr && core.pending_slot_state->payload != nullptr && - core.pending_slot_state->payload->early_dispatch_state.load(std::memory_order_relaxed) == - PTO2_EARLY_DISPATCH_STAGING); + (pending_ss == PTO2_EARLY_DISPATCH_STAGING || + (pending_ss == PTO2_EARLY_DISPATCH_DISPATCHED && + core.pending_slot_state->active_mask.requires_sync_start()))); SlotTransition t = decide_slot_transition( reg_task_id, reg_state, core.running_reg_task_id, core.pending_reg_task_id, pending_gated ); @@ -399,7 +419,19 @@ void SchedulerContext::check_running_cores_for_completion( // 2. Update slot data if (t.running_freed) { if (core.pending_slot_state != nullptr && !t.pending_done) { + // A gated sync_start block promoting into the running slot advances the + // rendezvous. Capture that BEFORE promote nulls the pending fields; after + // it lands, bump running_slot_count and ring iff this was the block that + // completed the cohort (and the producer already released). + PTO2TaskSlotState *promoted = core.pending_slot_state; + bool sync_start_promote = pending_gated && promoted->active_mask.requires_sync_start(); promote_pending_to_running(core); // Case 2 or Case 3 (with pending) + if (sync_start_promote) { + promoted->payload->running_slot_count.fetch_add(1, std::memory_order_seq_cst); + if (sched_->maybe_rendezvous_ring(*promoted)) { + sched_->propagate_dispatch_fanin(*promoted); + } + } } else { clear_running_slot(core); // Case 1 or Case 3 (no pending) if (t.pending_done) { @@ -449,10 +481,14 @@ bool SchedulerContext::enter_drain_mode(PTO2TaskSlotState *slot_state, int32_t b )) { return false; // Another thread already holds the drain slot. } - // We own the drain slot. Store the task and reset election flag before making it visible. + // We own the drain slot. Store the task and reset the coordination flags before making + // it visible. drain_state_.pending_task.store(slot_state, std::memory_order_release); drain_state_.drain_ack_mask.store(0, std::memory_order_relaxed); drain_state_.drain_worker_elected.store(0, std::memory_order_relaxed); + drain_state_.drain_stage_go.store(0, std::memory_order_relaxed); + drain_state_.drain_stage_done_mask.store(0, std::memory_order_relaxed); + drain_state_.drain_running_staged.store(0, std::memory_order_relaxed); // Release store: all stores above are now visible to any thread that // acquire-loads sync_start_pending and sees block_num > 0. drain_state_.sync_start_pending.store(block_num, std::memory_order_release); @@ -460,90 +496,170 @@ bool SchedulerContext::enter_drain_mode(PTO2TaskSlotState *slot_state, int32_t b } // Count total available resources across all scheduler threads for a given shape. -int32_t SchedulerContext::count_global_available(PTO2ResourceShape shape, uint8_t core_mask) { +// include_pending adds each thread's pending-capable cores/clusters — used by the +// gated (early) sync_start drain, which pre-stages onto idle running slots AND onto +// busy cores' pending slots. The ready drain (include_pending=false) counts idle only. +int32_t SchedulerContext::count_global_available(PTO2ResourceShape shape, uint8_t core_mask, bool include_pending) { int32_t total = 0; for (int32_t t = 0; t < active_sched_threads_; t++) { if (shape == PTO2ResourceShape::MIX) { - total += core_trackers_[t].count_mix_running_clusters(core_mask); + // Gated MIX uses split placement (each core to running-if-idle / pending-if-busy), + // so a cluster is available iff every used core has some free slot. The ready + // path (include_pending=false) still needs whole-cluster idle placement. + total += include_pending ? core_trackers_[t].count_mix_split_clusters(core_mask) : + core_trackers_[t].count_mix_running_clusters(core_mask); } else { total += core_trackers_[t].get_idle_core_offset_states(shape).count(); + if (include_pending) { + total += core_trackers_[t].get_pending_core_offset_states(shape).count(); + } } } return total; } -// Drain worker: dispatch all blocks in one pass across all threads' trackers. -// Called only when global resources >= block_num, so one pass always suffices. -// All other threads are spinning -- the drain worker has exclusive tracker access. -void SchedulerContext::drain_worker_dispatch(int32_t block_num) { - PTO2TaskSlotState *slot_state = drain_state_.pending_task.load(std::memory_order_acquire); - if (!slot_state) { - drain_state_.sync_start_pending.store(0, std::memory_order_release); - return; - } +// One thread's share of the drain staging: CAS-claim block indices and publish them onto +// THIS thread's own cores, concurrently with peers. Returns the number of cores placed on a +// RUNNING slot (the rendezvous seed contribution). Each thread touches only its own tracker +// and its own cores' doorbell-table entries; the CAS on next_block_idx and the fetch_or into +// staged_core_mask are the only cross-thread points. +// +// A gated (early) sync_start drain pre-stages every block behind its doorbell +// (prepare_block_for_dispatch is force-gated for the claimed drain range) and defers the launch +// to the rendezvous: idle cores take a gated RUNNING slot; busy cores take a gated PENDING +// slot (promoted by Case 3.3 as those cores' running tasks FIN). A non-gated (ready) drain +// leaves early_dispatch_state==NONE, so every block launches immediately on an idle running slot and +// the pending pass is skipped. For MIX, a gated block uses SPLIT placement (each core +// independently: idle->running, busy->pending) — safe only because gated. +int32_t +SchedulerContext::drain_stage_cores(PTO2TaskSlotState *slot_state, int32_t block_num, int32_t thread_idx, bool gated) { + CoreTracker &tracker = core_trackers_[thread_idx]; PTO2ResourceShape shape = slot_state->active_mask.to_shape(); uint8_t core_mask = slot_state->active_mask.core_mask(); - - for (int32_t t = 0; - t < active_sched_threads_ && slot_state->next_block_idx.load(std::memory_order_relaxed) < block_num; t++) { - auto valid = (shape == PTO2ResourceShape::MIX) ? - core_trackers_[t].get_mix_running_cluster_offset_states(core_mask) : - core_trackers_[t].get_idle_core_offset_states(shape); - int32_t start = slot_state->next_block_idx.load(std::memory_order_relaxed); - int32_t remaining = slot_state->logical_block_num - start; - int32_t claim = std::min(valid.count(), remaining); - slot_state->next_block_idx.store(static_cast(start + claim), std::memory_order_relaxed); - PublishHandle handles[CoreTracker::MAX_CLUSTERS * 3]; - int handle_count = 0; - for (int32_t b = 0; b < claim; b++) { - auto core_offset = valid.pop_first(); - handle_count += prepare_block_for_dispatch( - t, core_offset, *slot_state, shape, false, start + b, &handles[handle_count] - ); - } - wmb(); - uint64_t dispatch_ts = 0; + bool mix_split = gated && shape == PTO2ResourceShape::MIX; + int32_t running_staged = 0; + + // Stage from this thread's `valid` cores/clusters: CAS-claim a block-index range sized to + // what this thread can place (against peers claiming concurrently), then publish those + // blocks onto valid cores. prepare_block_for_dispatch decides each MIX core's slot per-core + // (idle -> running, busy -> pending when to_pending); a MIX cluster's idle cores are the + // running-slot cores, counted BEFORE staging mutates the tracker (rendezvous seed). + auto stage = [&](CoreTracker::BitStates valid, bool to_pending) { + while (valid.has_value()) { + int32_t avail = valid.count(); + int32_t start = 0; + int32_t claim = slot_state->claim_block_range(block_num, avail, start); + if (claim == 0) return; #if SIMPLER_DFX - if (l2_swimlane_level_ >= L2SwimlaneLevel::AICPU_TIMING) { - dispatch_ts = get_sys_cnt_aicpu(); - } + bool sub_prof = l2_swimlane_level_ >= L2SwimlaneLevel::SCHED_PHASES; + uint64_t prep_t0 = sub_prof ? get_sys_cnt_aicpu() : 0; +#endif + PublishHandle handles[CoreTracker::MAX_CLUSTERS * 3]; + int handle_count = 0; + int32_t claimed[CoreTracker::MAX_CLUSTERS * 3]; + for (int32_t b = 0; b < claim; b++) + claimed[b] = valid.pop_first(); + bool is_mix = (shape == PTO2ResourceShape::MIX); + if (claim > 0) prefetch_block_dst(thread_idx, claimed[0], is_mix); + for (int32_t b = 0; b < claim; b++) { + if (b + 1 < claim) prefetch_block_dst(thread_idx, claimed[b + 1], is_mix); + auto core_offset = claimed[b]; + if (shape == PTO2ResourceShape::MIX) { + running_staged += tracker.mix_cluster_idle_core_count(core_offset, core_mask); + } + handle_count += prepare_block_for_dispatch( + thread_idx, core_offset, *slot_state, shape, to_pending, start + b, &handles[handle_count], gated + ); + } + wmb(); + uint64_t dispatch_ts = 0; +#if SIMPLER_DFX + uint64_t pub_t0 = 0; + if (sub_prof) { + pub_t0 = get_sys_cnt_aicpu(); + // DrainPrepare bar: cluster scan happened before this lambda, so this covers the + // build_payload work for `claim` blocks (handle_count subtasks). + l2_swimlane_aicpu_record_sched_phase( + thread_idx, L2SwimlaneSchedPhaseKind::DrainPrepare, prep_t0, pub_t0, + sched_l2_swimlane_[thread_idx].sched_loop_count, static_cast(handle_count) + ); + } + if (l2_swimlane_level_ >= L2SwimlaneLevel::AICPU_TIMING) { + dispatch_ts = pub_t0 != 0 ? pub_t0 : get_sys_cnt_aicpu(); + } #endif - for (int i = 0; i < handle_count; i++) { - publish_subtask_to_core(handles[i], dispatch_ts); + // Accumulate this batch's gated cores into a LOCAL mask and OR it into the shared + // staged_core_mask ONCE below, instead of a seq_cst fetch_or per subtask — that + // per-write atomic contends across all drain threads on the same 2 words and was + // ~half the publish cost. The doorbell-table writes stay per-core (unique cid, no + // contention). + uint64_t my_mask[PTO2_EARLY_DISPATCH_CORE_MASK_WORDS] = {0}; + for (int i = 0; i < handle_count; i++) { + publish_subtask_to_core(handles[i], dispatch_ts); + if (gated) { + int32_t cid = tracker.get_core_id_by_offset(handles[i].core_offset); + sched_->early_dispatch_doorbell_table[cid].addr = handles[i].reg_addr; + sched_->early_dispatch_doorbell_table[cid].token = handles[i].reg_task_id; + my_mask[cid >> 6] |= 1ULL << (cid & 63); + } + } + if (gated) { + for (int w = 0; w < PTO2_EARLY_DISPATCH_CORE_MASK_WORDS; w++) { + if (my_mask[w] != 0) { + slot_state->payload->staged_core_mask[w].fetch_or(my_mask[w], std::memory_order_seq_cst); + } + } + } +#if SIMPLER_DFX + if (sub_prof) { + // DrainPublish bar: the MMIO write_reg per subtask (+ gated doorbell/mask record). + l2_swimlane_aicpu_record_sched_phase( + thread_idx, L2SwimlaneSchedPhaseKind::DrainPublish, pub_t0, get_sys_cnt_aicpu(), + sched_l2_swimlane_[thread_idx].sched_loop_count, static_cast(handle_count) + ); + } +#endif + sched_->record_published_blocks(*slot_state, claim); + // AIC/AIV running placement (whole block on idle cores); MIX running cores are + // counted per-cluster above (mix_cluster_idle_core_count). + if (gated && shape != PTO2ResourceShape::MIX && !to_pending) running_staged += handle_count; + } + }; + + if (mix_split) { + // Gated MIX: to_pending=true opts every BUSY used core into its pending slot while idle + // used cores take running slots (prepare_block_for_dispatch: to_pending && !is_core_idle). + stage(tracker.get_mix_split_cluster_offset_states(core_mask), /*to_pending=*/true); + } else { + auto idle = (shape == PTO2ResourceShape::MIX) ? tracker.get_mix_running_cluster_offset_states(core_mask) : + tracker.get_idle_core_offset_states(shape); + stage(idle, /*to_pending=*/false); // idle -> running (ready launch + gated pre-stage) + if (gated) { + stage(tracker.get_pending_core_offset_states(shape), /*to_pending=*/true); } } - - // The drain path IS this sync_start producer's dispatch, so it must bump its - // consumers' dispatch_fanin like the normal dispatch path - // (scheduler_dispatch.cpp, post-publish) -- otherwise a consumer whose only - // flagged producer is a sync_start (drain-dispatched) task never becomes an - // early-dispatch candidate. Idempotent via propagate's dispatch_propagated - // once-guard; the internal gate no-ops for an unflagged producer. - sched_->propagate_dispatch_fanin(*slot_state); - - // All blocks dispatched -- clear drain state. - // Release fence ensures tracker mutations are visible to threads that - // acquire-load sync_start_pending == 0 and resume normal operation. - std::atomic_thread_fence(std::memory_order_release); - drain_state_.pending_task.store(nullptr, std::memory_order_release); - drain_state_.drain_ack_mask.store(0, std::memory_order_relaxed); - drain_state_.drain_worker_elected.store(0, std::memory_order_relaxed); - drain_state_.sync_start_pending.store(0, std::memory_order_release); + return running_staged; } // Called by each scheduler thread when drain_state_.sync_start_pending != 0. // -// Protocol (single-stage ack barrier): -// 1. Ack barrier: all threads signal they've stopped dispatch, then spin -// until all ack bits are set. -// If this thread's bit gets cleared while waiting, a reset occurred -- return. -// 2. Election: one thread wins the CAS and becomes the drain worker. -// If resources are insufficient, reset ack/election fields and return -- -// all threads resume completion polling to free running cores, then retry. -// 3. Dispatch: elected thread dispatches all blocks (one pass, resources guaranteed). -// Non-elected threads spin-wait until sync_start_pending == 0. -// During dispatch the elected thread has exclusive tracker access. -void SchedulerContext::handle_drain_mode(int32_t thread_idx) { +// Protocol: +// 1. Ack barrier: all threads signal they've stopped dispatch, spin until all acked. +// If this thread's ack bit gets cleared while waiting, a reset occurred -- return. +// 2. Election + availability: one thread wins the CAS. It checks global resources; if +// insufficient it resets ack/election so all threads resume completion polling to free +// cores, then retry. If sufficient it releases parallel staging (stage_go). +// 3. Parallel stage: EVERY thread stages its OWN cores concurrently (CAS-claimed block +// indices), accumulates its running-slot cores, and marks its stage_done bit. +// 4. Finalize: the elected thread waits for all stage_done bits, seeds the rendezvous +// (running_slot_count) for a gated drain, and reopens the gate +// (a release-store the non-elected threads acquire, so the seed is visible before any +// completion promotes a pending block). Non-elected threads spin until the gate reopens. +void SchedulerContext::handle_drain_mode(int32_t thread_idx, [[maybe_unused]] uint64_t *out_stage_wall_cycles) { +#if SIMPLER_DFX + bool drain_prof = (l2_swimlane_level_ >= L2SwimlaneLevel::SCHED_PHASES && out_stage_wall_cycles != nullptr); + uint64_t drain_acked_ts = 0; // set at ack-barrier end; used to measure the stage wall +#endif // Every spin in this function honors is_completed(): once the run latches // completed_ (all tasks done, or a fatal error raised elsewhere), peers leave // the dispatch loop and stop participating in the drain. A thread parked in a @@ -576,46 +692,116 @@ void SchedulerContext::handle_drain_mode(int32_t thread_idx) { if ((ack & (1u << thread_idx)) == 0) return; SPIN_WAIT_HINT(); } - // Election -- exactly one thread wins the CAS. int32_t expected = 0; drain_state_.drain_worker_elected.compare_exchange_strong( expected, thread_idx + 1, std::memory_order_acquire, std::memory_order_relaxed ); + bool elected = drain_state_.drain_worker_elected.load(std::memory_order_relaxed) == thread_idx + 1; - if (drain_state_.drain_worker_elected.load(std::memory_order_relaxed) != thread_idx + 1) { - // Non-elected: spin-wait for drain completion or resource-insufficient reset. - while (drain_state_.sync_start_pending.load(std::memory_order_acquire) != 0) { + PTO2TaskSlotState *slot_state = drain_state_.pending_task.load(std::memory_order_acquire); + // OWNER is acquired before the drain is published and persists through + // completion, so every staging thread makes the same gate decision even if + // producer release changes early_dispatch_state during the barrier. + bool gated = slot_state != nullptr && slot_state->payload != nullptr && + PTO2SchedulerState::owns_early_sync_drain(*slot_state->payload); + + if (elected) { + if (slot_state == nullptr) { + // pending_task observed null only when a concurrent drain completion already cleared + // it. Stale-elected: release the election lock and return. Do NOT clear drain_ack_mask + // / sync_start_pending -- a *new* drain run may already be accumulating acks. + drain_state_.drain_worker_elected.store(0, std::memory_order_release); + return; + } + PTO2ResourceShape shape = slot_state->active_mask.to_shape(); + // A gated drain may pre-stage onto pending slots too (idle+pending); the ready drain + // needs block_num idle cores/clusters. + int32_t available = + count_global_available(shape, slot_state->active_mask.core_mask(), /*include_pending=*/gated); + if (available < block_num) { + // Insufficient -- reset so all threads resume completion polling to free cores, then retry. + drain_state_.drain_ack_mask.store(0, std::memory_order_release); + drain_state_.drain_worker_elected.store(0, std::memory_order_release); + return; + } + // Release parallel staging: every thread (this one included) now stages its own cores. + drain_state_.drain_running_staged.store(0, std::memory_order_relaxed); + drain_state_.drain_stage_done_mask.store(0, std::memory_order_relaxed); + drain_state_.drain_stage_go.store(1, std::memory_order_release); + } else { + // Non-elected: wait for the go signal, or bail if the elected thread reset (stale / + // insufficient resources). + while (drain_state_.drain_stage_go.load(std::memory_order_acquire) == 0) { if (is_completed()) return; if (drain_state_.drain_worker_elected.load(std::memory_order_acquire) == 0) return; SPIN_WAIT_HINT(); } - return; + slot_state = drain_state_.pending_task.load(std::memory_order_acquire); + if (slot_state == nullptr) return; + gated = slot_state->payload != nullptr && PTO2SchedulerState::owns_early_sync_drain(*slot_state->payload); } - // Elected: check if global resources are sufficient. - PTO2TaskSlotState *slot_state = drain_state_.pending_task.load(std::memory_order_acquire); - if (slot_state == nullptr) { - // pending_task is observed null only when a concurrent drain completion - // already cleared it (drain_worker_dispatch nulls it before reopening the - // gate). That drain is done and this is a stale-elected thread, so just - // release the election lock and return. Do NOT clear drain_ack_mask or - // sync_start_pending: a *new* drain run may already be active and - // accumulating acks, and zeroing them would corrupt it into a hang. - drain_state_.drain_worker_elected.store(0, std::memory_order_release); + // Parallel stage this thread's own cores (CAS-claimed block indices), then mark done. +#if SIMPLER_DFX + if (drain_prof) drain_acked_ts = get_sys_cnt_aicpu(); // pre-stage +#endif + int32_t my_running = drain_stage_cores(slot_state, block_num, thread_idx, gated); +#if SIMPLER_DFX + // out param carries the PURE drain_stage_cores wall (build_payload + MMIO publish of + // this thread's cores), isolating it from availability + stage_go handshake. + if (drain_prof && drain_acked_ts != 0) *out_stage_wall_cycles = get_sys_cnt_aicpu() - drain_acked_ts; +#endif + drain_state_.drain_running_staged.fetch_add(my_running, std::memory_order_acq_rel); + drain_state_.drain_stage_done_mask.fetch_or(1u << thread_idx, std::memory_order_release); + + if (!elected) { + // Non-elected: staging done; wait for the elected thread to reopen the gate. Exiting via + // sync_start_pending==0 (release/acquire) or drain_worker_elected==0 both synchronize + // with the elected's finalize (its release fence sequences the seed before both stores), + // so the running_slot_count seed is visible before this thread resumes completions. + while (drain_state_.sync_start_pending.load(std::memory_order_acquire) != 0) { + if (is_completed()) return; + if (drain_state_.drain_worker_elected.load(std::memory_order_acquire) == 0) return; + SPIN_WAIT_HINT(); + } return; } - PTO2ResourceShape shape = slot_state->active_mask.to_shape(); - int32_t available = count_global_available(shape, slot_state->active_mask.core_mask()); - if (available < block_num) { - // Insufficient resources -- reset drain fields so threads can resume - // completion polling to free running cores, then retry. - drain_state_.drain_ack_mask.store(0, std::memory_order_release); - drain_state_.drain_worker_elected.store(0, std::memory_order_release); - return; + // Elected: wait for all threads to finish staging, then seed the rendezvous and reopen. + while ((drain_state_.drain_stage_done_mask.load(std::memory_order_acquire) & all_acked) != all_acked) { + if (is_completed()) return; + SPIN_WAIT_HINT(); } + if (gated) { + // Seed the rendezvous with the running-slot cores staged across all threads; pending + // cores advance it as they promote. maybe_rendezvous_ring (producer release) rings iff + // this already equals popcount(staged_core_mask) — i.e. no pending spill. + slot_state->payload->running_slot_count.store( + static_cast(drain_state_.drain_running_staged.load(std::memory_order_acquire)), + std::memory_order_seq_cst + ); + } + // Clear drain state and reopen the gate FIRST, so the other threads resume immediately. + // Release fence sequences the seed + tracker mutations before every clear, so any thread + // that acquire-observes one of them (sync_start_pending==0 / drain_worker_elected==0) sees + // the seed. `slot_state` is a local holding the fa_fused slot (not drain_state_), so it stays + // valid for the propagate below even if a new drain reuses pending_task after reopen. + std::atomic_thread_fence(std::memory_order_release); + drain_state_.pending_task.store(nullptr, std::memory_order_release); + drain_state_.drain_stage_go.store(0, std::memory_order_relaxed); + drain_state_.drain_stage_done_mask.store(0, std::memory_order_relaxed); + drain_state_.drain_ack_mask.store(0, std::memory_order_relaxed); + drain_state_.drain_worker_elected.store(0, std::memory_order_relaxed); + drain_state_.sync_start_pending.store(0, std::memory_order_release); - // Dispatch -- all other threads are spinning, elected thread has exclusive tracker access. - drain_worker_dispatch(block_num); + // Recheck after publishing the drain seed. The producer-side rendezvous check can race + // ahead of drain completion and fail while running_slot_count is still incomplete. When + // every block landed directly in a running slot, no pending promotion remains to retry it. + if (gated) { + sched_->retry_sync_start_rendezvous_after_drain(*slot_state); + } else { + sched_->propagate_dispatch_fanin(*slot_state); + } + PTO2SchedulerState::finish_early_sync_drain(*slot_state->payload); } diff --git a/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_context.h b/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_context.h index 96b0c14d2..079706e40 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_context.h +++ b/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_context.h @@ -222,7 +222,7 @@ class SchedulerContext { void build_payload( PTO2DispatchPayload &dispatch_payload, PTO2TaskSlotState &slot_state, PTO2SubtaskSlot subslot, - const AsyncCtx &async_ctx, int32_t block_idx + int32_t block_idx, bool force_gate ); // Batched-dispatch primitives. prepare_* builds the payload and per-core @@ -243,7 +243,7 @@ class SchedulerContext { PublishHandle prepare_subtask_to_core( int32_t thread_idx, int32_t core_offset, PTO2TaskSlotState &slot_state, PTO2SubtaskSlot subslot, - bool to_pending, int32_t block_idx + bool to_pending, int32_t block_idx, bool force_gate ); inline void publish_subtask_to_core(const PublishHandle &h, uint64_t dispatch_ts) { @@ -253,11 +253,44 @@ class SchedulerContext { write_reg(h.reg_addr, RegId::DATA_MAIN_BASE, static_cast(h.reg_task_id)); } + // Prefetch the cold per-core structures the next block's prepare touches. + // Ordering is load-bearing: issue the STALLING LOAD first — CoreExecState, + // read by dispatch_seq++ whose value feeds reg_task_id -> buf_idx -> the whole + // dispatch — for every core of the block, BEFORE the store-target prefetches. + // MSHRs saturate (a MIX block warms 3 cores); issuing the read prefetches + // first keeps them from being the ones dropped. The dispatch-buffer writes + // still get prefetched (measured to help ~30% on this shallow-store-buffer + // control core), just after the reads. rw=1 on CoreExecState (read AND + // written) gives Exclusive, serving both without a Shared->Exclusive upgrade. + inline void prefetch_block_dst(int32_t thread_idx, int32_t core_offset, bool is_mix) { + CoreTracker &tracker = core_trackers_[thread_idx]; + int32_t cids[3] = {}; + int32_t nc = 0; + if (is_mix) { + cids[nc++] = tracker.get_core_id_by_offset(tracker.get_aic_core_offset(core_offset)); + cids[nc++] = tracker.get_core_id_by_offset(tracker.get_aiv0_core_offset(core_offset)); + cids[nc++] = tracker.get_core_id_by_offset(tracker.get_aiv1_core_offset(core_offset)); + } else { + cids[nc++] = tracker.get_core_id_by_offset(core_offset); + } + // Stalling loads first. + for (int32_t i = 0; i < nc; i++) + __builtin_prefetch(&core_exec_states_[cids[i]], 1, 3); + // Store targets after (dispatch buffer CL0 control + CL1 args, both bufs). + for (int32_t i = 0; i < nc; i++) { + for (int32_t buf = 0; buf < 2; buf++) { + const char *dp = reinterpret_cast(&payload_per_core_[cids[i]][buf]); + __builtin_prefetch(dp, 1, 3); + __builtin_prefetch(dp + 64, 1, 3); + } + } + } + // Fan out one block's subtasks (1 for AIC/AIV, 1-3 for MIX) into the // caller-supplied handles buffer. Returns the number of handles written. int prepare_block_for_dispatch( int32_t thread_idx, int32_t core_offset, PTO2TaskSlotState &slot_state, PTO2ResourceShape shape, - bool to_pending, int32_t block_idx, PublishHandle *out_handles + bool to_pending, int32_t block_idx, PublishHandle *out_handles, bool force_gate = false ); void dispatch_shape( @@ -270,7 +303,7 @@ class SchedulerContext { // is queued) and sets made_progress / try_pushed when it stages, so the caller // is a single unconditional call like normal dispatch. After normal dispatch // leaves idle cores spare, pre-stage the consumers of any RUNNING flagged - // producer onto those cores with not_ready=1 (gated). Touches no dependency + // producer onto those cores with a non-zero src_payload (gated). Touches no dependency // state — the task is released by the doorbell at its normal ready-pop (Hook 2). int32_t try_early_dispatch( int32_t thread_idx, CoreTracker &tracker, bool pmu_active, bool &made_progress, bool &try_pushed @@ -278,7 +311,7 @@ class SchedulerContext { // Stage the already-claimed range [start, start+count) of consumer `c` onto // thread_idx's idle (RUNNING slot) then pending (gated-pending, promote-on-FIN) - // cores from the provided free-core sets. The caller advances next_block_idx and + // cores from the provided free-core sets. The caller claims next_block_idx and // re-pushes `c` BEFORE calling, so this expensive prepare+publish runs // concurrently with peers (mirrors the normal SPMD dispatch path). Returns the // number of blocks staged. @@ -378,9 +411,15 @@ class SchedulerContext { ); bool enter_drain_mode(PTO2TaskSlotState *slot_state, int32_t block_num); - int32_t count_global_available(PTO2ResourceShape shape, uint8_t core_mask); - void drain_worker_dispatch(int32_t block_num); - void handle_drain_mode(int32_t thread_idx); + int32_t count_global_available(PTO2ResourceShape shape, uint8_t core_mask, bool include_pending = false); + // One thread's share of the drain: CAS-claim block indices and stage them onto THIS + // thread's own cores (parallel with peers), returning the number of running-slot cores + // staged (the rendezvous seed contribution). + int32_t drain_stage_cores(PTO2TaskSlotState *slot_state, int32_t block_num, int32_t thread_idx, bool gated); + // out_stage_wall_cycles (profiling only): cycles this thread spent in drain_stage_cores + // (prepare + publish), set ONLY on threads that actually staged. Lets the caller isolate + // the pure stage wall from the ack-barrier + finalize spans in the Drain bar. + void handle_drain_mode(int32_t thread_idx, uint64_t *out_stage_wall_cycles = nullptr); // ========================================================================= // Cold path: exit checks, stall diagnostics, profiling (scheduler_cold_path.cpp) diff --git a/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_dispatch.cpp b/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_dispatch.cpp index 70f3b6fd6..ee0dfb844 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_dispatch.cpp +++ b/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_dispatch.cpp @@ -36,6 +36,15 @@ #define unlikely(x) __builtin_expect(!!(x), 0) #endif +// AICore materializes args[] from src_payload on the gated path using the +// byte offsets in pto2_dispatch_payload.h (the AICore .o cannot see PTO2TaskPayload). +// Pin those constants to the real layout here, where the struct is fully visible. +static_assert(offsetof(PTO2TaskPayload, tensor_count) == PTO2_TASKPAYLOAD_TENSOR_COUNT_OFFSET); +static_assert(offsetof(PTO2TaskPayload, scalar_count) == PTO2_TASKPAYLOAD_SCALAR_COUNT_OFFSET); +static_assert(offsetof(PTO2TaskPayload, tensors) == PTO2_TASKPAYLOAD_TENSORS_OFFSET); +static_assert(offsetof(PTO2TaskPayload, scalars) == PTO2_TASKPAYLOAD_SCALARS_OFFSET); +static_assert(sizeof(Tensor) == PTO2_TASKPAYLOAD_TENSOR_STRIDE); + // ============================================================================= // Dispatch helpers // ============================================================================= @@ -113,37 +122,49 @@ int SchedulerContext::pop_ready_tasks_batch( } void SchedulerContext::build_payload( - PTO2DispatchPayload &dispatch_payload, PTO2TaskSlotState &slot_state, PTO2SubtaskSlot subslot, - const AsyncCtx &async_ctx, int32_t block_idx + PTO2DispatchPayload &dispatch_payload, PTO2TaskSlotState &slot_state, PTO2SubtaskSlot subslot, int32_t block_idx, + bool force_gate ) { int32_t slot_idx = static_cast(subslot); uint64_t callable_addr = get_function_bin_addr(slot_state.task->kernel_id[slot_idx]); const CoreCallable *callable = reinterpret_cast(callable_addr); dispatch_payload.function_bin_addr = callable->resolved_addr(); auto &payload = *slot_state.payload; - int n = 0; - for (int32_t i = 0; i < payload.tensor_count; i++) { - dispatch_payload.args[n++] = reinterpret_cast(&payload.tensors[i]); - } - for (int32_t i = 0; i < payload.scalar_count; i++) { - dispatch_payload.args[n++] = payload.scalars[i]; + // A claimed early-stage range stays gated even if producer completion flips + // the shared state before this payload is built. All other dispatches run on + // pickup. + if (PTO2SchedulerState::should_gate_early_dispatch( + force_gate, payload.early_dispatch_state.load(std::memory_order_relaxed) + )) { + // Gated task: hand the idle AICore the source payload (non-zero = gate) and + // let it fill args[0..num_args) itself during its doorbell wait, instead of + // paying the arg-vector write on this scheduler thread. + dispatch_payload.src_payload = reinterpret_cast(&payload); + } else { + // Ready task: fill args here; src_payload = 0 signals AICore to run on pickup. + dispatch_payload.src_payload = 0; + int n = 0; + for (int32_t i = 0; i < payload.tensor_count; i++) { + dispatch_payload.args[n++] = reinterpret_cast(&payload.tensors[i]); + } + for (int32_t i = 0; i < payload.scalar_count; i++) { + dispatch_payload.args[n++] = payload.scalars[i]; + } } dispatch_payload.local_context.block_idx = block_idx; dispatch_payload.local_context.block_num = slot_state.logical_block_num; - dispatch_payload.local_context.async_ctx = async_ctx; - dispatch_payload.args[PAYLOAD_LOCAL_CONTEXT_INDEX] = reinterpret_cast(&dispatch_payload.local_context); - dispatch_payload.args[PAYLOAD_GLOBAL_CONTEXT_INDEX] = reinterpret_cast(&dispatch_payload.global_context); - // Early-dispatch: a task being staged (Hook 1 set early_dispatch_state to - // STAGING before this call) is gated — the AICore must wait for the - // DATA_MAIN_BASE high-32 doorbell. All other dispatches run on pickup. - dispatch_payload.not_ready = - (slot_state.payload->early_dispatch_state.load(std::memory_order_relaxed) == PTO2_EARLY_DISPATCH_STAGING) ? 1 : - 0; + // AsyncCtx's slab pointers + capacity are prefilled once per (core, buf_idx) + // in init(); only task_token varies per dispatch. deferred_slab is never null + // on this path, so task_token is always the live task id (matches the old + // AsyncCtx::make(non-null) result). args[PAYLOAD_LOCAL_CONTEXT_INDEX] / + // [PAYLOAD_GLOBAL_CONTEXT_INDEX] are per-(core, buf_idx) constants, also + // prefilled in init(). + dispatch_payload.local_context.async_ctx.task_token = slot_state.task->task_id; } SchedulerContext::PublishHandle SchedulerContext::prepare_subtask_to_core( int32_t thread_idx, int32_t core_offset, PTO2TaskSlotState &slot_state, PTO2SubtaskSlot subslot, bool to_pending, - int32_t block_idx + int32_t block_idx, bool force_gate ) { CoreTracker &tracker = core_trackers_[thread_idx]; auto core_id = tracker.get_core_id_by_offset(core_offset); @@ -161,11 +182,12 @@ SchedulerContext::PublishHandle SchedulerContext::prepare_subtask_to_core( uint32_t buf_idx = reg_task_id & 1u; PTO2DispatchPayload &payload = payload_per_core_[core_id][buf_idx]; - DeferredCompletionSlab *deferred_slab = &deferred_slab_per_core_[core_id][buf_idx]; - deferred_slab->count = 0; - deferred_slab->error_code = PTO2_ERROR_NONE; - AsyncCtx async_ctx = AsyncCtx::make(slot_state.task->task_id, deferred_slab); - build_payload(payload, slot_state, subslot, async_ctx, block_idx); + // The deferred_slab is NOT reset here: it is init-cleared once and the + // completion path re-clears count only after a task actually recorded a + // deferred completion (count > 0), which is rare. A non-deferred task never + // touches count/error_code, so the slab stays clean without a per-dispatch + // write — keeping this cold per-core line off the dispatch path. + build_payload(payload, slot_state, subslot, block_idx, force_gate); if (to_pending) { core_exec_state.pending_subslot = subslot; @@ -213,7 +235,7 @@ SchedulerContext::PublishHandle SchedulerContext::prepare_subtask_to_core( int SchedulerContext::prepare_block_for_dispatch( int32_t thread_idx, int32_t core_offset, PTO2TaskSlotState &slot_state, PTO2ResourceShape shape, bool to_pending, - int32_t block_idx, PublishHandle *out_handles + int32_t block_idx, PublishHandle *out_handles, bool force_gate ) { #if SIMPLER_DFX if (is_dump_args_enabled()) { @@ -232,22 +254,30 @@ int SchedulerContext::prepare_block_for_dispatch( if (shape == PTO2ResourceShape::MIX) { uint8_t cmask = slot_state.active_mask.core_mask(); int n = 0; + // Per-core slot placement (#1308): an idle used core takes its running slot + // (tracked by the completion poller), a busy used core takes its gated pending slot + // (promoted on completion). The sync_start drain relies on this — it passes + // to_pending=true so every busy core opts into a pending slot; the non-zero + // src_payload gate keeps the whole cohort waiting for the rendezvous. if (cmask & PTO2_SUBTASK_MASK_AIC) { bool p = to_pending && !tracker.is_aic_core_idle(core_offset); out_handles[n++] = prepare_subtask_to_core( - thread_idx, tracker.get_aic_core_offset(core_offset), slot_state, PTO2SubtaskSlot::AIC, p, block_idx + thread_idx, tracker.get_aic_core_offset(core_offset), slot_state, PTO2SubtaskSlot::AIC, p, block_idx, + force_gate ); } if (cmask & PTO2_SUBTASK_MASK_AIV0) { bool p = to_pending && !tracker.is_aiv0_core_idle(core_offset); out_handles[n++] = prepare_subtask_to_core( - thread_idx, tracker.get_aiv0_core_offset(core_offset), slot_state, PTO2SubtaskSlot::AIV0, p, block_idx + thread_idx, tracker.get_aiv0_core_offset(core_offset), slot_state, PTO2SubtaskSlot::AIV0, p, block_idx, + force_gate ); } if (cmask & PTO2_SUBTASK_MASK_AIV1) { bool p = to_pending && !tracker.is_aiv1_core_idle(core_offset); out_handles[n++] = prepare_subtask_to_core( - thread_idx, tracker.get_aiv1_core_offset(core_offset), slot_state, PTO2SubtaskSlot::AIV1, p, block_idx + thread_idx, tracker.get_aiv1_core_offset(core_offset), slot_state, PTO2SubtaskSlot::AIV1, p, block_idx, + force_gate ); } #if SIMPLER_DFX @@ -255,15 +285,17 @@ int SchedulerContext::prepare_block_for_dispatch( #endif return n; } else if (shape == PTO2ResourceShape::AIC) { - out_handles[0] = - prepare_subtask_to_core(thread_idx, core_offset, slot_state, PTO2SubtaskSlot::AIC, to_pending, block_idx); + out_handles[0] = prepare_subtask_to_core( + thread_idx, core_offset, slot_state, PTO2SubtaskSlot::AIC, to_pending, block_idx, force_gate + ); #if SIMPLER_DFX sched_l2_swimlane_[thread_idx].phase_dispatch_count += 1; #endif return 1; } else { - out_handles[0] = - prepare_subtask_to_core(thread_idx, core_offset, slot_state, PTO2SubtaskSlot::AIV0, to_pending, block_idx); + out_handles[0] = prepare_subtask_to_core( + thread_idx, core_offset, slot_state, PTO2SubtaskSlot::AIV0, to_pending, block_idx, force_gate + ); #if SIMPLER_DFX sched_l2_swimlane_[thread_idx].phase_dispatch_count += 1; #endif @@ -326,14 +358,11 @@ void SchedulerContext::dispatch_shape( PublishHandle handles[CoreTracker::MAX_CLUSTERS * 3]; int handle_count = 0; bool dispatched_any = false; - // Slots dispatched this pop whose dispatch_fanin must be propagated to - // consumers. Deferred until AFTER publish (below) so a flagged producer's - // fanout walk never sits between claiming cores and publishing its own - // blocks — doing it inline delays this thread's blocks while peer threads - // co-dispatching the same SPMD task publish immediately, misaligning the - // task's block starts. Bounded by cores.count() ≤ MAX_CLUSTERS dispatches. - PTO2TaskSlotState *prop_list[CoreTracker::MAX_CLUSTERS]; - int prop_n = 0; + // Logical block ranges published by this pop. AIV can pop two entries per + // cluster, so this ledger has the same worst-case bound as handles[]. + PTO2TaskSlotState *published_list[CoreTracker::MAX_CLUSTERS * 3]; + int16_t published_counts[CoreTracker::MAX_CLUSTERS * 3]; + int published_n = 0; #if SIMPLER_SCHED_PROFILING uint64_t t_setup_start = get_sys_cnt_aicpu(); #endif @@ -406,14 +435,6 @@ void SchedulerContext::dispatch_shape( break; } - dispatched_any = true; - try_pushed = true; - // Record for deferred dispatch_fanin propagation after this pop's - // blocks are published (see after the loop). propagate's own guard - // filters non-flagged slots, so recording unconditionally is cheap. - if (prop_n < static_cast(sizeof(prop_list) / sizeof(prop_list[0]))) { - prop_list[prop_n++] = slot_state; - } // Claim a contiguous range of blocks, hand the slot back to the // ready queue immediately, then perform the expensive dispatches. // This lets other schedulers concurrently claim and dispatch the @@ -421,11 +442,16 @@ void SchedulerContext::dispatch_shape( // this thread fills all its own cores. Only local `start + b` is // read after the push — `next_block_idx` may already be advanced // by another scheduler that popped the slot. - int32_t start = slot_state->next_block_idx.load(std::memory_order_relaxed); - int32_t remaining = slot_state->logical_block_num - start; int32_t available = is_mix ? selected_mix_clusters.count() : cores.count(); - int32_t claim = std::min(available, remaining); - slot_state->next_block_idx.store(static_cast(start + claim), std::memory_order_relaxed); + int32_t start = 0; + int32_t claim = slot_state->claim_block_range(slot_state->logical_block_num, available, start); + if (claim == 0) continue; + dispatched_any = true; + try_pushed = true; + + published_list[published_n] = slot_state; + published_counts[published_n] = static_cast(claim); + published_n++; if (start + claim < slot_state->logical_block_num) { disp_queues[static_cast(shape)].push(slot_state); @@ -451,11 +477,9 @@ void SchedulerContext::dispatch_shape( } flush_publish(); - // Blocks are published; now propagate dispatch_fanin for any flagged - // producers dispatched above (knob A: producer is running). Off the - // pre-publish path so it cannot delay or misalign their blocks. - for (int i = 0; i < prop_n; i++) { - sched_->propagate_dispatch_fanin(*prop_list[i]); + for (int i = 0; i < published_n; i++) { + sched_->record_published_blocks(*published_list[i], published_counts[i]); + sched_->propagate_dispatch_fanin(*published_list[i]); } #if SIMPLER_SCHED_PROFILING l2_swimlane.sched_dispatch_setup_cycle += (get_sys_cnt_aicpu() - t_setup_start); @@ -571,11 +595,11 @@ void SchedulerContext::dispatch_ready_tasks( } // Stage the ALREADY-CLAIMED range [start, start+count) of consumer `c` onto -// thread_idx's idle then pending cores. The caller (the queue drain) has advanced -// next_block_idx by `count` under pop-exclusivity AND re-pushed `c` for peers +// thread_idx's idle then pending cores. The caller has atomically advanced +// next_block_idx by `count` AND re-pushed `c` for peers // BEFORE calling this — so this, the expensive prepare+publish, runs CONCURRENTLY // with peers staging other ranges of the same consumer. This mirrors the normal -// SPMD dispatch path (claim range -> store next_block_idx -> re-push -> dispatch). +// SPMD dispatch path (claim range -> re-push -> dispatch). // `idle`/`pend` are this thread's free-core sets, sized so idle.count+pend.count >= // count (the caller clamped the claim to them), so all `count` blocks get a core. // @@ -585,10 +609,10 @@ void SchedulerContext::dispatch_ready_tasks( // waiting for an ack the gated task never sends). Each staged core stays // pending_occupied while gated, so no second gated block stacks on it. // -// Self-ring: release flips STAGING->DISPATCHED then rings the mask. A block staged -// after that flip isn't in the mask release read, so this thread rings it here. The -// seq_cst order between "OR mask then load early_dispatch_state" (here) and "store DISPATCHED -// then read mask" (release) guarantees every gated core's doorbell fires. +// Doorbell ownership: release flips STAGING->DISPATCHED and exchanges the shared +// mask to claim its bits. A late stager ORs its bits, then fetch-and-clears only +// those release did not take and rings them from its immutable local handles. +// The seq_cst order guarantees every gated core has exactly one writer. int32_t SchedulerContext::stage_consumer_blocks( int32_t thread_idx, PTO2TaskSlotState *c, PTO2ResourceShape shape, int32_t start, int32_t count, CoreTracker::BitStates &idle, CoreTracker::BitStates &pend @@ -597,13 +621,13 @@ int32_t SchedulerContext::stage_consumer_blocks( // Stamp the real pre-stage time (NOT 0) so the swimlane shows these blocks // dispatched during the producer's run, not at trace start. uint64_t early_dispatch_ts = get_sys_cnt_aicpu(); - uint64_t my_cores[PTO2_EARLY_DISPATCH_CORE_MASK_WORDS] = {0}; // cores this thread gated (for self-ring) + uint64_t my_cores[PTO2_EARLY_DISPATCH_CORE_MASK_WORDS] = {0}; // cores gated by this staging pass int32_t staged = 0; int32_t block = start; // Mirror the normal flush_publish (scheduler_dispatch.cpp wmb()+publish loop): // prepare ALL claimed blocks' payloads (idle bucket -> running slot, pend bucket // -> gated pending), then ONE wmb(), then publish. The wmb guarantees the - // not_ready gate + args are globally visible before any DATA_MAIN_BASE token — + // src_payload gate + source args are globally visible before any DATA_MAIN_BASE token — // without it a gated core can pick up the token and dcci a stale payload. The // shared `count` budget bounds total blocks <= free clusters/cores, so both // buckets fit one handles[] buffer. @@ -612,7 +636,9 @@ int32_t SchedulerContext::stage_consumer_blocks( auto prepare_from = [&](CoreTracker::BitStates &avail, bool to_pending) { while (count > 0 && avail.has_value()) { int32_t core_offset = avail.pop_first(); - n += prepare_block_for_dispatch(thread_idx, core_offset, *c, shape, to_pending, block, &handles[n]); + n += prepare_block_for_dispatch( + thread_idx, core_offset, *c, shape, to_pending, block, &handles[n], /*force_gate=*/true + ); block++; count--; staged++; @@ -635,21 +661,35 @@ int32_t SchedulerContext::stage_consumer_blocks( for (int w = 0; w < PTO2_EARLY_DISPATCH_CORE_MASK_WORDS; w++) if (my_cores[w] != 0) c->payload->staged_core_mask[w].fetch_or(my_cores[w], std::memory_order_seq_cst); - // If release already flipped DISPATCHED, it may have read the mask before our - // bits landed — ring our own cores so none is left gated forever. - if (staged > 0 && - c->payload->early_dispatch_state.load(std::memory_order_seq_cst) == PTO2_EARLY_DISPATCH_DISPATCHED) { + // Full publication and release are independent events. The seq_cst + // state/launch/count operations form a two-sided handshake. A released + // block must ring before contributing to the publication count. + bool released = staged > 0 && + c->payload->early_dispatch_state.load(std::memory_order_seq_cst) == PTO2_EARLY_DISPATCH_DISPATCHED; + + // Claim only bits the release path did not take. Local handles remain valid + // even if the shared per-core table is reused before this thread resumes. + if (released) { + uint64_t owned[PTO2_EARLY_DISPATCH_CORE_MASK_WORDS] = {0}; for (int w = 0; w < PTO2_EARLY_DISPATCH_CORE_MASK_WORDS; w++) { - uint64_t bits = my_cores[w]; - while (bits != 0) { - int cid = w * 64 + __builtin_ctzll(bits); - bits &= bits - 1; - PTO2SchedulerState::ring_one_doorbell( - sched_->early_dispatch_doorbell_table[cid].addr, sched_->early_dispatch_doorbell_table[cid].token - ); + if (my_cores[w] != 0) { + owned[w] = + PTO2SchedulerState::claim_late_staged_doorbell_bits(c->payload->staged_core_mask[w], my_cores[w]); } } + for (int i = 0; i < n; i++) { + int32_t cid = tracker.get_core_id_by_offset(handles[i].core_offset); + PTO2SchedulerState::ring_claimed_local_doorbell( + owned[cid >> 6], cid, handles[i].reg_addr, handles[i].reg_task_id + ); + } + wmb(); } + sched_->record_published_blocks(*c, staged); + // Retry unconditionally after publication. The guards are cheap, and a + // pre-ring state read can become stale if release completes before this + // count update. + sched_->propagate_dispatch_fanin(*c); return staged; } @@ -717,22 +757,8 @@ SchedulerContext::early_dispatch_shape(int32_t thread_idx, PTO2ResourceShape sha sched_->early_dispatch_queues[s].push_batch(&batch[bi], got - bi); break; } - // CAS-claim a contiguous range [start, start+claim) sized to this thread's - // free cores; CAS keeps it atomic against peers AND normal dispatch. - int32_t start = 0, claim = 0; - while (true) { - int16_t cur = c->next_block_idx.load(std::memory_order_relaxed); - if (cur >= c->logical_block_num) break; // fully claimed - int32_t cnt = c->logical_block_num - cur; - if (cnt > freecores) cnt = freecores; - if (c->next_block_idx.compare_exchange_weak( - cur, static_cast(cur + cnt), std::memory_order_seq_cst, std::memory_order_relaxed - )) { - start = cur; - claim = cnt; - break; - } - } + int32_t start = 0; + int32_t claim = c->claim_block_range(c->logical_block_num, freecores, start); if (claim == 0) continue; // nothing left to claim -> drop (no re-push) // Re-push for concurrent peers BEFORE the expensive staging. if (start + claim < c->logical_block_num) { @@ -752,11 +778,11 @@ SchedulerContext::early_dispatch_shape(int32_t thread_idx, PTO2ResourceShape sha return total_staged; } -// Early-dispatch drain (idle pass), mirroring dispatch_ready_tasks: owns its own -// gating and progress-flag updates, and orders staging the same way — MIX strict -// priority, IDLE stage before PENDING stage, cross-thread idle gating -// (MIX-IDLE ▶ c/v-IDLE ▶ MIX-PEND ▶ c/v-PEND). sync_start doesn't apply here (those -// tasks are excluded from early dispatch at push time, propagate_dispatch_fanin). +// Early-dispatch drain (idle pass) — the EARLY source's analog of dispatch_ready_tasks. +// Both sources share run_staging_order for the shape order (MIX strict priority, IDLE +// before PENDING, cross-thread idle gating: MIX-IDLE ▶ c/v-IDLE ▶ MIX-PEND ▶ c/v-PEND) and +// both drain their sync_start cohort FIRST as the highest occupancy tier (Tier 0), via the +// same all-or-nothing drain barrier. This one owns its own gating and progress flags. // Returns the number of blocks staged this pass (for the EarlyDispatch swimlane bar). int32_t SchedulerContext::try_early_dispatch( int32_t thread_idx, CoreTracker &tracker, bool pmu_active, bool &made_progress, bool &try_pushed @@ -776,9 +802,31 @@ int32_t SchedulerContext::try_early_dispatch( if (sched_->ready_sync_queues[s].size() > 0 || sched_->ready_queues[s].size() > 0) return 0; } - // Early staging (NOT is_ready): same MIX/idle/pending order as normal dispatch, via the - // shared skeleton. early_dispatch_shape stages a gated block range and never enters drain, - // so the stage callback always reports "no stop". + // ===== Tier 0: sync_start cohorts (highest occupancy tier, all-or-nothing) ===== + // sync_start candidates park in their own shape-agnostic queue. They cannot ride + // early_dispatch_shape's per-thread partial range-claim (a partial claim strands gated + // cohort blocks nobody rings, so the rendezvous never reaches block_num). Instead arm the + // drain barrier: it takes exclusive tracker access and only stages when global + // idle+pending >= block_num, guaranteeing all-or-nothing. Win => the dispatch loop runs + // the gated drain next iteration; lose (a drain is already armed, or capacity < + // block_num) => cancel the owner and re-push or transfer its final ready route. A + // non-STAGING pop was already released and is dropped. Staging happens inside the drain, + // so this arms at most one drain and adds no blocks to total_staged here. + if (PTO2TaskSlotState *c = sched_->early_sync_start_queue.pop()) { + if (PTO2SchedulerState::try_claim_early_sync_drain(*c->payload)) { + if (c->payload->early_dispatch_state.load(std::memory_order_seq_cst) != PTO2_EARLY_DISPATCH_STAGING) { + sched_->cancel_early_sync_drain(*c); + } else if (enter_drain_mode(c, c->logical_block_num)) { + PTO2SchedulerState::mark_early_sync_drain_armed(*c->payload); + } else { + sched_->cancel_early_sync_drain(*c); + } + } + } + + // Regular early staging (NOT is_ready): same MIX/idle/pending order as normal dispatch, + // via the shared skeleton. early_dispatch_shape stages a gated block range and never + // enters drain, so the stage callback always reports "no stop". int32_t total_staged = 0; run_staging_order( thread_idx, pmu_active, @@ -890,7 +938,9 @@ int32_t SchedulerContext::resolve_and_dispatch(Runtime *runtime, int32_t thread_ // and silently corrupt the snapshot. constexpr size_t kMax = static_cast(std::numeric_limits::max()); for (int s = 0; s < L2SWIMLANE_NUM_QUEUE_SHAPES; s++) { - const size_t qsize = sched_->ready_queues[s].size(); + // Total normal-source ready depth of shape `s` = regular ready lane + the + // sync_start Tier-0 lane; both feed dispatch_ready_tasks for this shape. + const size_t qsize = sched_->ready_queues[s].size() + sched_->ready_sync_queues[s].size(); iter_shared_snapshot[s] = static_cast(std::min(qsize, kMax)); } iter_shared_sampled = true; @@ -922,7 +972,6 @@ int32_t SchedulerContext::resolve_and_dispatch(Runtime *runtime, int32_t thread_ // "now" so the first budget cycle starts when this thread does, not at // an undefined value. uint64_t last_progress_ts = get_sys_cnt_aicpu(); - while (true) { if (completed_.load(std::memory_order_acquire)) { break; @@ -1043,7 +1092,27 @@ int32_t SchedulerContext::resolve_and_dispatch(Runtime *runtime, int32_t thread_ // Phase 2 drain check if (drain_state_.sync_start_pending.load(std::memory_order_acquire) != 0) { +#if SIMPLER_DFX + // The drain is otherwise a swimlane blind spot: the `continue` below skips + // every phase record, and handle_drain_mode is uninstrumented. Time it here so + // the sync_start stop-the-world window shows on the scheduler lane (one bar per + // iteration that enters the drain; retries appear as multiple bars). + uint64_t drain_t0 = (l2_swimlane_level_ >= L2SwimlaneLevel::SCHED_PHASES) ? get_sys_cnt_aicpu() : 0; + uint64_t drain_stage_wall = 0; // set by handle_drain_mode ONLY if this thread staged + handle_drain_mode(thread_idx, &drain_stage_wall); + // Record a Drain bar only when this thread actually did drain work (reached + // drain_stage_cores). The many no-op entries — ack + availability-insufficient + // reset, stale-elected, non-elected bail before stage_go — never stage, so they + // would otherwise clutter the lane with zero-work drain(0) bars. + if (l2_swimlane_level_ >= L2SwimlaneLevel::SCHED_PHASES && drain_stage_wall != 0) { + l2_swimlane_aicpu_record_sched_phase( + thread_idx, L2SwimlaneSchedPhaseKind::Drain, drain_t0, get_sys_cnt_aicpu(), + l2_swimlane.sched_loop_count, static_cast(drain_stage_wall) + ); + } +#else handle_drain_mode(thread_idx); +#endif continue; } diff --git a/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_types.h b/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_types.h index 757ea3ad5..82e73429e 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_types.h +++ b/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_types.h @@ -365,6 +365,53 @@ class alignas(64) CoreTracker { return MixPlacement::PENDING; } + // --- Gated MIX split placement --- + // A gated MIX block can place each of its cores INDEPENDENTLY (idle core -> gated + // running slot; busy core with a free pending slot -> gated pending slot), because + // every core waits on the doorbell and nothing executes until the rendezvous rings. + // This is unsafe for immediate (non-gated) dispatch — hence separate from + // classify_mix_cluster, which forces a single placement for the whole cluster. + + // Cores of `cluster_offset` named by core_mask. + BitStates mix_used_cores(int32_t cluster_offset, uint8_t core_mask) const { + BitStates used(0ULL); + if (core_mask & PTO2_SUBTASK_MASK_AIC) used |= BitStates(1ULL << cluster_offset); + if (core_mask & PTO2_SUBTASK_MASK_AIV0) used |= BitStates(1ULL << (cluster_offset + 1)); + if (core_mask & PTO2_SUBTASK_MASK_AIV1) used |= BitStates(1ULL << (cluster_offset + 2)); + return used; + } + + // Every used core has SOME free slot: a core lacks one only when it is running AND + // its pending slot is occupied (both slots taken). + bool mix_cluster_all_slots(int32_t cluster_offset, uint8_t core_mask) const { + BitStates used = mix_used_cores(cluster_offset, core_mask); + if (!used.has_value()) return false; + BitStates no_slot = (~core_states_) & pending_occupied_; // running AND pending taken + return !(used & no_slot).has_value(); + } + + // Used cores that are idle -> will take a running slot (rendezvous seed count). + int32_t mix_cluster_idle_core_count(int32_t cluster_offset, uint8_t core_mask) const { + return (mix_used_cores(cluster_offset, core_mask) & core_states_).count(); + } + + // Clusters where every used core has a free slot (gated MIX split gate/iteration). + BitStates get_mix_split_cluster_offset_states(uint8_t core_mask) const { + BitStates result(0ULL); + BitStates candidates = get_cluster_offset_states(); + while (candidates.has_value()) { + int32_t off = candidates.pop_first(); + if (mix_cluster_all_slots(off, core_mask)) { + result |= BitStates(1ULL << off); + } + } + return result; + } + + int32_t count_mix_split_clusters(uint8_t core_mask) const { + return get_mix_split_cluster_offset_states(core_mask).count(); + } + BitStates get_mix_running_cluster_offset_states(uint8_t core_mask) const { BitStates result(0ULL); BitStates candidates = get_cluster_offset_states(); @@ -489,7 +536,15 @@ struct alignas(64) SyncStartDrainState { std::atomic drain_worker_elected{0}; // 0=none; >0: elected thread's (thread_idx+1) std::atomic drain_ack_mask{0}; // bit per thread; all-set = all threads reached ack barrier std::atomic pending_task{nullptr}; // held task (not re-queued) - int32_t _pad[10]; + // Parallel staging: after the elected thread confirms global availability it sets + // stage_go, releasing every thread to stage its OWN cores concurrently (vs the old + // single-thread serial fill). Each thread ORs its bit into stage_done_mask when it + // finishes and accumulates its running-slot cores into running_staged; the elected + // thread waits for all bits, seeds the rendezvous, and reopens the gate. + std::atomic drain_stage_go{0}; // 0=hold; 1=elected released parallel staging + std::atomic drain_stage_done_mask{0}; // bit per thread; all-set = all threads done staging + std::atomic drain_running_staged{0}; // sum of running-slot cores staged (rendezvous seed) + int32_t _pad[7]; }; static_assert(sizeof(SyncStartDrainState) == 64); diff --git a/src/a2a3/runtime/host_build_graph/runtime/shared/pto_runtime2_init.cpp b/src/a2a3/runtime/host_build_graph/runtime/shared/pto_runtime2_init.cpp index fdddade21..b1b5bb47e 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/shared/pto_runtime2_init.cpp +++ b/src/a2a3/runtime/host_build_graph/runtime/shared/pto_runtime2_init.cpp @@ -120,6 +120,7 @@ PTO2SchedulerLayout PTO2SchedulerState::reserve_layout(DeviceArena &arena, int32 for (int i = 0; i < PTO2_NUM_RESOURCE_SHAPES; i++) { layout.off_early_dispatch_queue_slots[i] = ready_queue_reserve_layout(arena, PTO2_EARLY_DISPATCH_QUEUE_SIZE); } + layout.off_early_sync_start_queue_slots = ready_queue_reserve_layout(arena, PTO2_EARLY_DISPATCH_QUEUE_SIZE); // Force a cache-line base so writes from scheduler thread 0 (sole writer of // the dep_pool) do not invalidate adjacent multi-threaded regions like // ready_queue.slots. @@ -169,6 +170,12 @@ bool PTO2SchedulerState::init_data_from_layout( return false; } } + if (!ready_queue_init_data_from_layout( + &sched->early_sync_start_queue, arena, layout.off_early_sync_start_queue_slots, + PTO2_EARLY_DISPATCH_QUEUE_SIZE + )) { + return false; + } auto *orch_err = pto2_sm_layout::orch_error_code_addr(sm_dev_base); auto *dep_entries = static_cast(arena.region_ptr(layout.off_dep_pool_entries)); @@ -192,6 +199,7 @@ void PTO2SchedulerState::wire_arena_pointers(const PTO2SchedulerLayout &layout, &sched->early_dispatch_queues[i], arena, layout.off_early_dispatch_queue_slots[i] ); } + ready_queue_wire_arena_pointers(&sched->early_sync_start_queue, arena, layout.off_early_sync_start_queue_slots); sched->ring_sched_state.dep_pool.base = static_cast(arena.region_ptr(layout.off_dep_pool_entries)); } @@ -210,6 +218,7 @@ void PTO2SchedulerState::destroy() { for (int i = 0; i < PTO2_NUM_RESOURCE_SHAPES; i++) { ready_queue_destroy(&sched->early_dispatch_queues[i]); } + ready_queue_destroy(&sched->early_sync_start_queue); } // =============================================================================