diff --git a/src/a2a3/platform/include/common/l2_swimlane_profiling.h b/src/a2a3/platform/include/common/l2_swimlane_profiling.h index b552df3f35..61dc7a1da3 100644 --- a/src/a2a3/platform/include/common/l2_swimlane_profiling.h +++ b/src/a2a3/platform/include/common/l2_swimlane_profiling.h @@ -514,9 +514,18 @@ enum class L2SwimlaneSchedPhaseKind : uint32_t { // push newly-ready successors, ring doorbells for // early-dispatch hits. tasks_processed = # consumers visited. // Separate-lane (Worker View pid=4 AICPU_N) - DummyTask = 7, // Per-dummy identity marker (zero-width). tasks_processed - // = task_token_raw low 32 bits so deps.json flow arrows - // can land on it. + DummyTask = 7, // Per-dummy identity marker (zero-width). tasks_processed + // = task_token_raw low 32 bits so deps.json flow arrows + // can land on it. + Drain = 8, // handle_drain_mode outer: the sync_start stop-the-world drain + // (ack barrier + availability + parallel stage + finalize). + // One bar per dispatch-loop iteration that enters the drain, + // so retries show as multiple bars. Otherwise this time is a + // swimlane blind spot (the loop `continue`s past all records). + DrainPrepare = 9, // inner: this thread's drain_stage_cores prepare pass + // (cluster scan + build_payload). tasks_processed = subtasks. + DrainPublish = 10, // inner: this thread's drain_stage_cores publish pass + // (MMIO write_reg per subtask). tasks_processed = subtasks. }; /** Index layout of the queue-depth snapshot arrays below: AIC=0, AIV=1, MIX=2. @@ -546,7 +555,7 @@ struct L2SwimlaneAicpuSchedPhaseRecord { uint32_t tasks_processed; // Tasks processed in this phase batch uint32_t pop_hit; // SCHED_DISPATCH delta since last emit (0 for Complete) uint32_t pop_miss; // SCHED_DISPATCH delta since last emit (0 for Complete) - int16_t shared_depth_at_start[L2SWIMLANE_NUM_QUEUE_SHAPES]; // sched->ready_queues[shape].size() + int16_t shared_depth_at_start[L2SWIMLANE_NUM_QUEUE_SHAPES]; // ready_queues[shape] + ready_sync_queues[shape] int16_t shared_depth_at_end[L2SWIMLANE_NUM_QUEUE_SHAPES]; uint32_t _pad[4]; // 64B alignment padding }; diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md b/src/a2a3/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md index 459ec5c492..3af12bfc36 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md @@ -546,9 +546,11 @@ Each scheduler thread runs a tight loop with two main phases: - Poll register `COND` on each managed core - When `TASK_FIN_STATE` detected: record completion timestamps, call `on_subtask_complete(task_id, subslot)` to increment the completion counter; when `completed_subtasks == total_required_subtasks`, trigger `on_task_complete(task_id)` which marks `task_state[slot] = COMPLETED`, acquires fanout lock, traverses fanout list (incrementing consumers' `fanin_refcount`), marks `task_state[slot] = CONSUMED`, and advances `last_task_alive` watermark -**Phase 2 — Dispatch**: +**Phase 2 — Dispatch** (full model in §8.6): -- For each idle core: pop a task from the matching shape-based ready queue (lock-free MPMC Vyukov queue, one per resource shape) +- Drain each source (normal ready ▸ speculative early) in occupancy order — `sync_start` + Tier-0 ▸ MIX ▸ AIC/AIV, idle ▸ pending — popping from the matching shape-based ready queue + (lock-free MPMC Vyukov queue, one per resource shape) - Build `PTO2DispatchPayload` from `TaskDescriptor` with `task_id`, `subslot`, `kernel_id`, and `core_type` - Write task pointer to `Handshake.task`, signal AICore via register `DATA_MAIN_BASE` @@ -558,8 +560,11 @@ After these phases, the scheduler updates profiling headers and checks for termi Ready queues use a lock-free bounded MPMC (Vyukov) design: -- One `PTO2ReadyQueue` per resource shape (5 shapes: `AIC_ONLY`, `AIV_X1`, `AIV_X2`, `AIC_AIV_X1`, `AIC_AIV_X2`) -- **Push**: any thread (orchestrator via `init_task`, or scheduler on completion) pushes newly-ready tasks to the queue matching `task->active_mask.to_shape()` +- One `PTO2ReadyQueue` per resource shape — 3 shapes (`PTO2_NUM_RESOURCE_SHAPES`): `MIX` + (AIC+AIV cluster), `AIC`, `AIV`. Alongside `ready_queues[]` there is a per-shape + `ready_sync_queues[]` (sync_start Tier-0) and the speculative `early_dispatch_queues[]` / + `early_sync_start_queue` — see §8.6 for the full source × tier model. +- **Push**: any thread (orchestrator via `init_task`, or scheduler on completion) pushes newly-ready tasks to the queue matching `task->active_mask.to_shape()` (sync_start cohorts to the sync lane) - **Pop**: scheduler threads pop from the queue matching the idle core's resource shape - Per-slot sequence counters prevent ABA problems - `enqueue_pos` and `dequeue_pos` are on separate cache lines to avoid false sharing @@ -603,6 +608,88 @@ Private internals are split across three .cpp files by responsibility: `AicpuExecutor` calls neither `handshake_*`, `assign_*`, `reassign_*`, nor `emergency_shutdown` directly — they are private, invoked only by `init` and `on_orchestration_done`. +### 8.6 Dispatch model — two sources, sync tiers, occupancy order + +`resolve_and_dispatch` places ready and speculative work onto AICore cores under one +occupancy model. Two orthogonal axes decide *what* runs and *where*: + +- **Source** — `NORMAL` (all producers done; the task sits in a ready queue and launches on + pickup) vs `EARLY` (a *speculative* pre-stage of a not-yet-released task; its dispatch + payload carries a non-zero `src_payload` gate and launches later by a doorbell). Normal + strictly precedes early. +- **Cohort** — `SYNC_START` (an SPMD cohort that must launch atomically) vs `REGULAR` (each + block launches independently). "is it ready" (source) and "does it need a rendezvous" + (cohort) are orthogonal. + +Within each source the occupancy order is **`sync_start` ▸ MIX ▸ AIC/AIV** (shape), and per +shape **idle ▸ pending** (an idle core takes its running slot; a busy core takes its gated +pending slot, promoted on completion). This order lives in one shared skeleton, +`run_staging_order`; the normal and early sources differ only in the per-shape stage callback +(pickup vs gated). + +#### Queues + +| Source | Regular lanes | sync_start lane | +| ------ | ------------- | --------------- | +| NORMAL (ready) | `ready_queues[MIX\|AIC\|AIV]` | `ready_sync_queues[MIX\|AIC\|AIV]` (per-shape) | +| EARLY (speculative) | `early_dispatch_queues[MIX\|AIC\|AIV]` | `early_sync_start_queue` (single) | + +A task routes to the sync lane iff `active_mask.requires_sync_start()`. In each source the +sync lane is drained as a strict **Tier-0** before the regular lane (`sync_start > MIX > C/V`), +and early dispatch runs only once *both* normal lanes are empty (normal ▸ early). + +**Asymmetry (deliberate):** the normal sync lane is per-shape (3 queues) because a ready sync +cohort can dispatch *inline* when it fits, reusing the per-shape `dispatch_shape`; the early +sync lane is a single, shape-agnostic queue because an early cohort is *always* gated → always +takes the drain path, whose rendezvous counts cores (not blocks) and is shape-agnostic. Both +feed the same drain. + +#### sync_start drain + rendezvous + +A sync_start cohort of `block_num` cores must occupy all its cores before any of them run. +When it cannot fit inline, `enter_drain_mode` arms a stop-the-world drain: + +1. **Single election** — a CAS on `sync_start_pending` (0 → −1) makes drains mutually + exclusive; only one cohort drains at a time, regardless of source. +2. **All-or-nothing** — the elected thread checks `count_global_available >= block_num` + *before* staging; if short it aborts (stages nothing) and retries after completions free + cores. A cohort is fully staged or not at all — never partial. +3. **Parallel stage** — all threads barrier, then each CAS-claims a block range and stages + its own cores with a non-zero `src_payload` gate: idle cores → running slots, busy cores → + pending slots. +4. **Rendezvous launch** — `running_slot_count` counts staged running-slot cores; when it + reaches `popcount(staged_core_mask)` **and** the producer has released, + `maybe_rendezvous_ring` rings every gated core's doorbell together — the cohort starts as one. + +Single-election + all-or-nothing make the drain deadlock-free across multiple cohorts: at most +one drains, and it fully stages or waits, so two cohorts can never each half-occupy the cluster +set (see the completion path's `pending_gated` classification for why a promoted-but-still-gated +block is not mistaken for a normal task). + +#### Early-candidate gate: producer must publish every block (deadlock avoidance) + +`propagate_dispatch_fanin` (the EARLY-source candidate trigger) no-ops until the producer is +**fully published**: `published_block_count == logical_block_num`. Normal dispatch, regular +early staging, and the sync drain increment this counter only after the claimed range's payloads +and MMIO dispatch tokens are visible. A staged producer also waits for release and completion of +its owned doorbell pass before exposing fanout. + +This is load-bearing: a flagged SPMD producer with more blocks than cores (for example, a +50-block AIC projection on 24 AIC cores) dispatches in waves. If its first wave triggered a +downstream MIX cohort to gate every running and pending slot, the remaining producer blocks +would find no core, never complete, and the cohort rendezvous waiting for producer release would +never ring. Full publication is stronger than full reservation: every producer block has both a +reserved core slot and a launch-visible payload before a consumer can pre-occupy resources. +`next_block_idx` records reservation progress; `published_block_count` independently establishes +publication and early-candidate readiness. + +#### MIX per-core placement + +A MIX task spans a cluster (1 AIC + 2 AIV). `classify_mix_cluster` admits a cluster whenever +every used core has a free slot; `prepare_block_for_dispatch` then places **per core** +(`to_pending && !is_core_idle`): idle cores → running, busy cores → pending. Cross-core start +skew within a block is tolerated by AICore incore synchronization. + --- ## 9. AICore Worker Interaction diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2_types.h b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2_types.h index e5730d8717..77c1672d3f 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2_types.h +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2_types.h @@ -229,6 +229,14 @@ enum PTO2EarlyDispatchLaunchState : uint8_t { 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 // it occupies is recorded as a bitmask (staged_core_mask, 1 bit per global // core_id); the completion-path release iterates the set bits and rings each @@ -255,9 +263,10 @@ 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, then destructively split between the - // release and late-stager paths. (Re)initialized in PTO2TaskPayload::init. + // 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 @@ -276,14 +285,26 @@ struct PTO2TaskPayload { // 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 and claims the current mask; a thread that - // stages a block after that flip claims and rings only its remaining bits. + // 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 release owner publishes COMPLETE only after all doorbells it claimed - // are visible. Combined with published_block_count, this keeps fanout - // private until release-owned and late-owned blocks have both launched. + // 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 === @@ -368,6 +389,8 @@ struct PTO2TaskPayload { 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); } }; diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.h b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.h index c32efcb785..baa2d7b76d 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.h +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.h @@ -407,6 +407,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[PTO2_MAX_RING_DEPTH]; uint64_t ready_queue_capacity; int32_t dep_pool_capacities[PTO2_MAX_RING_DEPTH]; @@ -693,6 +694,19 @@ 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); @@ -740,23 +754,128 @@ struct PTO2SchedulerState { slot_state.payload->published_block_count.fetch_add(static_cast(count), std::memory_order_seq_cst); } - // Event-driven candidate detection (the dual of fanin_refcount/ready). Call after - // publishing a FLAGGED producer's blocks: once every logical block is launch-visible, - // 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-fully-published, 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. + // 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_seq_cst) == PTO2_EARLY_DISPATCH_STAGING) return; - if (p.payload->early_dispatch_launch_state.load(std::memory_order_seq_cst) == - PTO2_EARLY_DISPATCH_LAUNCH_RINGING) - 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(); @@ -773,16 +892,26 @@ struct PTO2SchedulerState { // 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); } } @@ -811,30 +940,45 @@ struct PTO2SchedulerState { // 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 || !try_claim_early_dispatch_launch(*slot_state.payload)) return true; + 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 ); - // 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); + // 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 + ); } - 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 // itself codegen-flagged; the gate inside no-ops otherwise). Defer it via the // 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. @@ -848,7 +992,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) { @@ -882,7 +1032,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/tensormap_and_ringbuffer/runtime/scheduler/scheduler_cold_path.cpp b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_cold_path.cpp index e00144d437..8fc84f5d7b 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_cold_path.cpp +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_cold_path.cpp @@ -1094,6 +1094,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/tensormap_and_ringbuffer/runtime/scheduler/scheduler_completion.cpp b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_completion.cpp index f455c17625..0ce29e9bb9 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_completion.cpp +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_completion.cpp @@ -343,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 ); @@ -414,7 +429,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) { @@ -464,10 +491,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); @@ -475,98 +506,182 @@ 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; - // Gather the claimed offsets so the per-core destination structures of block - // b+1 can be prefetched while block b's prepare runs (software pipeline — - // this drain launches every block on a single elected thread). - 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(t, claimed[0], is_mix); - for (int32_t b = 0; b < claim; b++) { - if (b + 1 < claim) prefetch_block_dst(t, claimed[b + 1], is_mix); - handle_count += - prepare_block_for_dispatch(t, claimed[b], *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, claim = 0; + while (true) { + int16_t cur = slot_state->next_block_idx.load(std::memory_order_relaxed); + if (cur >= block_num) return; // all blocks claimed + int32_t cnt = block_num - cur; + if (cnt > avail) cnt = avail; + if (slot_state->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; + } + } + 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 - for (int i = 0; i < handle_count; i++) { - publish_subtask_to_core(handles[i], dispatch_ts); + 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 + // 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); } - sched_->record_published_blocks(*slot_state, claim); } - - // 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 @@ -599,46 +714,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/tensormap_and_ringbuffer/runtime/scheduler/scheduler_context.h b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_context.h index b8023c6c56..e93e47ba89 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_context.h +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_context.h @@ -273,7 +273,7 @@ class SchedulerContext { // 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 cids[3] = {}; int32_t nc = 0; if (is_mix) { cids[nc++] = tracker.get_core_id_by_offset(tracker.get_aic_core_offset(core_offset)); @@ -420,9 +420,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/tensormap_and_ringbuffer/runtime/scheduler/scheduler_dispatch.cpp b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_dispatch.cpp index d4153496fd..50f8b8bbb9 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_dispatch.cpp +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_dispatch.cpp @@ -256,6 +256,11 @@ 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( @@ -790,11 +795,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 @@ -814,9 +819,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, @@ -928,7 +955,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; @@ -1090,7 +1119,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/tensormap_and_ringbuffer/runtime/scheduler/scheduler_types.h b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_types.h index 76a79952ae..e4519cd892 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_types.h +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_types.h @@ -354,6 +354,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(); @@ -478,7 +525,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/tensormap_and_ringbuffer/runtime/shared/pto_runtime2_init.cpp b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/shared/pto_runtime2_init.cpp index b8e3785ba6..fda19e1074 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/shared/pto_runtime2_init.cpp +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/shared/pto_runtime2_init.cpp @@ -159,6 +159,7 @@ PTO2SchedulerLayout PTO2SchedulerState::reserve_layout( 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); for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { // Force a cache-line base so Orch-side dep_pool writes do not invalidate // adjacent multi-threaded regions like ready_queue.slots. @@ -211,6 +212,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); for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { @@ -245,6 +252,7 @@ void PTO2SchedulerState::reset_for_reuse(const PTO2SchedulerLayout &layout, void for (int i = 0; i < PTO2_NUM_RESOURCE_SHAPES; i++) { sched->early_dispatch_queues[i].reset_for_reuse(); } + sched->early_sync_start_queue.reset_for_reuse(); sched->async_wait_list.reset_for_reuse(); (void)layout; @@ -264,6 +272,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); for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { sched->ring_sched_states[r].dep_pool.base = static_cast(arena.region_ptr(layout.off_dep_pool_entries[r])); @@ -286,6 +295,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); } // ============================================================================= diff --git a/src/a5/platform/include/common/l2_swimlane_profiling.h b/src/a5/platform/include/common/l2_swimlane_profiling.h index a9cbffa6f9..ce38f3460d 100644 --- a/src/a5/platform/include/common/l2_swimlane_profiling.h +++ b/src/a5/platform/include/common/l2_swimlane_profiling.h @@ -514,9 +514,18 @@ enum class L2SwimlaneSchedPhaseKind : uint32_t { // push newly-ready successors, ring doorbells for // early-dispatch hits. tasks_processed = # consumers visited. // Separate-lane (Worker View pid=4 AICPU_N) - DummyTask = 7, // Per-dummy identity marker (zero-width). tasks_processed - // = task_token_raw low 32 bits so deps.json flow arrows - // can land on it. + DummyTask = 7, // Per-dummy identity marker (zero-width). tasks_processed + // = task_token_raw low 32 bits so deps.json flow arrows + // can land on it. + Drain = 8, // handle_drain_mode outer: the sync_start stop-the-world drain + // (ack barrier + availability + parallel stage + finalize). + // One bar per dispatch-loop iteration that enters the drain, + // so retries show as multiple bars. Otherwise this time is a + // swimlane blind spot (the loop `continue`s past all records). + DrainPrepare = 9, // inner: this thread's drain_stage_cores prepare pass + // (cluster scan + build_payload). tasks_processed = subtasks. + DrainPublish = 10, // inner: this thread's drain_stage_cores publish pass + // (MMIO write_reg per subtask). tasks_processed = subtasks. }; /** Index layout of the queue-depth snapshot arrays below: AIC=0, AIV=1, MIX=2. diff --git a/src/common/platform/shared/host/l2_swimlane_collector.cpp b/src/common/platform/shared/host/l2_swimlane_collector.cpp index 71f716c0a0..f8ec75fd0f 100644 --- a/src/common/platform/shared/host/l2_swimlane_collector.cpp +++ b/src/common/platform/shared/host/l2_swimlane_collector.cpp @@ -1028,6 +1028,12 @@ int L2SwimlaneCollector::export_swimlane_json() { return "resolve"; case L2SwimlaneSchedPhaseKind::DummyTask: return "dummy_task"; + case L2SwimlaneSchedPhaseKind::Drain: + return "drain"; + case L2SwimlaneSchedPhaseKind::DrainPrepare: + return "drain_prepare"; + case L2SwimlaneSchedPhaseKind::DrainPublish: + return "drain_publish"; } return "unknown"; }; diff --git a/tests/st/a2a3/tensormap_and_ringbuffer/spmd_sync_start_early_dispatch/kernels/aiv/kernel_spmd_write_slow.cpp b/tests/st/a2a3/tensormap_and_ringbuffer/spmd_sync_start_early_dispatch/kernels/aiv/kernel_spmd_write_slow.cpp new file mode 100644 index 0000000000..9b0cae98ea --- /dev/null +++ b/tests/st/a2a3/tensormap_and_ringbuffer/spmd_sync_start_early_dispatch/kernels/aiv/kernel_spmd_write_slow.cpp @@ -0,0 +1,76 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ + +/** + * SPMD Multi-Block Write Kernel with a spin delay (AIV) + * + * Same result as the plain write kernel, but spins for `spin_iters` before the + * store. A slow producer keeps its cores occupied long enough that a dependent + * consumer is processed as an early-dispatch candidate WHILE the producer is still + * running — the only way a trivial-kernel scene exercises the speculative + * gated-dispatch path (a fast producer finishes before the consumer is ever staged). + * + * out[(base_cl + block_idx) * FLOATS_PER_CACHE_LINE] = float(block_idx) + * + * Args: + * args[0] = output Tensor* (INOUT) + * args[1] = scalar: base_cl (starting cache line index for this task) + * args[2] = scalar: spin_iters (0 = no delay) + */ + +#include +#include + +#include "tensor.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] // NOLINT(whitespace/braces) +#endif + +#include "intrinsic.h" + +static constexpr int32_t FLOATS_PER_CACHE_LINE = 16; + +#ifdef PTO_CPUSTUB_HPP +#define dcci(...) \ + do { \ + } while (0) +#endif +#ifndef SINGLE_CACHE_LINE +#define SINGLE_CACHE_LINE 0 +#endif +#ifndef CACHELINE_OUT +#define CACHELINE_OUT 0 +#endif + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *out_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ float *out = reinterpret_cast<__gm__ float *>(out_tensor->buffer.addr) + out_tensor->start_offset; + + int32_t base_cl = static_cast(args[1]); + int32_t spin_iters = static_cast(args[2]); + int32_t block_idx = get_block_idx(args); + int32_t offset = (base_cl + block_idx) * FLOATS_PER_CACHE_LINE; + + volatile int32_t acc = 0; + for (int32_t i = 0; i < spin_iters; i++) { + acc++; // ++ not += i: += i overflows int32 for large spin_iters (UB); volatile keeps the loop + } + (void)acc; + + out[offset] = static_cast(block_idx); + + dcci(&out[offset], SINGLE_CACHE_LINE, CACHELINE_OUT); +} diff --git a/tests/st/a2a3/tensormap_and_ringbuffer/spmd_sync_start_early_dispatch/kernels/orchestration/spmd_sync_start_early_dispatch_orch.cpp b/tests/st/a2a3/tensormap_and_ringbuffer/spmd_sync_start_early_dispatch/kernels/orchestration/spmd_sync_start_early_dispatch_orch.cpp new file mode 100644 index 0000000000..316fedeec2 --- /dev/null +++ b/tests/st/a2a3/tensormap_and_ringbuffer/spmd_sync_start_early_dispatch/kernels/orchestration/spmd_sync_start_early_dispatch_orch.cpp @@ -0,0 +1,90 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ + +/** + * SPMD Sync-Start Early-Dispatch Orchestration + * + * Exercises the sync_start early-dispatch path (gated drain + running-slot + * rendezvous): a FLAGGED producer's dispatch makes a require_sync_start consumer an + * early-dispatch candidate. The consumer is pre-staged gated across idle running + * slots AND (once idle runs out) busy cores' pending slots, then all blocks are + * released together at the rendezvous once every block occupies a running slot and + * the producer completes. + * + * Tasks (the producer writes one cache line per block; the MIX consumer writes + * one per participating core, so three cache lines per block): + * P: AIC block_num=50, base_cl=0, allow_early_resolve=true + * C: MIX block_num=24, base_cl=50, require_sync_start=true, dep=[P] + * + * Args layout: [output] + */ + +#include +#include + +#include "pto_orchestration_api.h" // NOLINT(build/include_subdir) +#include "pto_arg_with_deps.h" // NOLINT(build/include_subdir) + +#define FUNC_SPMD_WRITE_AIC 0 +#define FUNC_SPMD_MIX_AIC 1 +#define FUNC_SPMD_MIX_AIV0 2 +#define FUNC_SPMD_MIX_AIV1 3 + +extern "C" { + +__attribute__((visibility("default"))) PTO2OrchestrationConfig aicpu_orchestration_config(const L2TaskArgs &orch_args) { + (void)orch_args; // NOLINT(readability/casting) + return PTO2OrchestrationConfig{ + .expected_arg_count = 1, + }; +} + +// spin_iters chosen so the producer stays on-core long enough for the scheduler to +// process the consumer as an early-dispatch candidate WHILE the producer is running +// (a fast producer would finish first and route the consumer through the ready path). +static constexpr int64_t PRODUCER_SPIN_ITERS = 10000000; + +static PTO2TaskId submit_producer(const Tensor &out, int16_t block_num, int64_t base_cl) { + L0TaskArgs args; + args.add_inout(out); + args.add_scalar(base_cl); + args.add_scalar(PRODUCER_SPIN_ITERS); + args.launch_spec.set_block_num(block_num); + args.set_allow_early_resolve(true); // flagged: consumers may early-dispatch off it + return rt_submit_aic_task(FUNC_SPMD_WRITE_AIC, args).task_id(); +} + +static void submit_sync_consumer(const Tensor &out, int16_t block_num, int64_t base_cl, PTO2TaskId dep) { + MixedKernels kernels; + kernels.aic_kernel_id = FUNC_SPMD_MIX_AIC; + kernels.aiv0_kernel_id = FUNC_SPMD_MIX_AIV0; + kernels.aiv1_kernel_id = FUNC_SPMD_MIX_AIV1; + L0TaskArgsWithDeps<4> args; + args.add_inout(out); + args.add_scalar(base_cl); + args.launch_spec.set_block_num(block_num); + args.launch_spec.set_require_sync_start(true); // atomic cohort launch + args.add_dep(dep); // sole producer, flagged -> early-dispatch candidate + rt_submit_task(kernels, args); +} + +__attribute__((visibility("default"))) void aicpu_orchestration_entry(const L2TaskArgs &orch_args) { + const Tensor &ext_output = orch_args.tensor(0).ref(); + + rt_scope_begin(PTO2ScopeMode::MANUAL); + PTO2TaskId prod = submit_producer(ext_output, 50, 0); + submit_sync_consumer(ext_output, 24, 50, prod); + rt_scope_end(); + + LOG_INFO_V9("[spmd_sync_start_early_dispatch] wide producer + MIX sync_start consumer submitted"); +} + +} // extern "C" diff --git a/tests/st/a2a3/tensormap_and_ringbuffer/spmd_sync_start_early_dispatch/test_spmd_sync_start_early_dispatch.py b/tests/st/a2a3/tensormap_and_ringbuffer/spmd_sync_start_early_dispatch/test_spmd_sync_start_early_dispatch.py new file mode 100644 index 0000000000..87ecc1cdb8 --- /dev/null +++ b/tests/st/a2a3/tensormap_and_ringbuffer/spmd_sync_start_early_dispatch/test_spmd_sync_start_early_dispatch.py @@ -0,0 +1,93 @@ +#!/usr/bin/env python3 +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +"""A wide flagged producer feeds a MIX sync_start early-dispatch consumer. + +The consumer must not stage until every producer block range has reserved a core slot; +otherwise it occupies every AIC slot and strands the producer's unclaimed remainder. +""" + +import torch +from simpler.task_interface import ArgDirection as D + +from simpler_setup import SceneTestCase, TaskArgsBuilder, Tensor, scene_test + +FLOATS_PER_CACHE_LINE = 16 +PRODUCER_BLOCKS = 50 +SYNC_BLOCKS = 24 +SYNC_BASE_CL = PRODUCER_BLOCKS +TOTAL_CL = SYNC_BASE_CL + SYNC_BLOCKS * 3 + + +@scene_test(level=2, runtime="tensormap_and_ringbuffer") +class TestSpmdSyncStartEarlyDispatch(SceneTestCase): + RTOL = 0 + ATOL = 0 + + CALLABLE = { + "orchestration": { + "source": "kernels/orchestration/spmd_sync_start_early_dispatch_orch.cpp", + "function_name": "aicpu_orchestration_entry", + "signature": [D.INOUT], + }, + "incores": [ + { + "func_id": 0, + "name": "SPMD_WRITE_AIC", + "source": "kernels/aiv/kernel_spmd_write_slow.cpp", + "core_type": "aic", + "signature": [D.INOUT], + }, + { + "func_id": 1, + "name": "SPMD_MIX_AIC", + "source": "../spmd_multiblock_mix/kernels/aic/kernel_spmd_mix.cpp", + "core_type": "aic", + "signature": [D.INOUT], + }, + { + "func_id": 2, + "name": "SPMD_MIX_AIV0", + "source": "../spmd_multiblock_mix/kernels/aiv/kernel_spmd_mix.cpp", + "core_type": "aiv", + "signature": [D.INOUT], + }, + { + "func_id": 3, + "name": "SPMD_MIX_AIV1", + "source": "../spmd_multiblock_mix/kernels/aiv/kernel_spmd_mix.cpp", + "core_type": "aiv", + "signature": [D.INOUT], + }, + ], + } + + CASES = [ + { + "name": "Case1", + "platforms": ["a2a3sim", "a2a3"], + "config": {"aicpu_thread_num": 4, "block_dim": 24}, + "params": {}, + } + ] + + def generate_args(self, params): + return TaskArgsBuilder(Tensor("output", torch.zeros(TOTAL_CL * FLOATS_PER_CACHE_LINE, dtype=torch.float32))) + + def compute_golden(self, args, params): + out = args.output + for block_idx in range(PRODUCER_BLOCKS): + out[block_idx * FLOATS_PER_CACHE_LINE] = float(block_idx) + for block_idx in range(SYNC_BLOCKS): + for slot in range(3): + out[(SYNC_BASE_CL + block_idx * 3 + slot) * FLOATS_PER_CACHE_LINE] = float(block_idx) + + +if __name__ == "__main__": + SceneTestCase.run_module(__name__) diff --git a/tests/st/a2a3/tensormap_and_ringbuffer/spmd_sync_start_mix_spill/kernels/aic/kernel_spmd_mix_slow.cpp b/tests/st/a2a3/tensormap_and_ringbuffer/spmd_sync_start_mix_spill/kernels/aic/kernel_spmd_mix_slow.cpp new file mode 100644 index 0000000000..0485d14b09 --- /dev/null +++ b/tests/st/a2a3/tensormap_and_ringbuffer/spmd_sync_start_mix_spill/kernels/aic/kernel_spmd_mix_slow.cpp @@ -0,0 +1,68 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ + +/** + * SPMD MIX slow kernel (AIC): writes float(block_idx) at cache line + * (base_cl + block_idx * 3 + 0), then optionally spins so the task holds its + * cluster long enough for a dependent sync_start cohort to pre-stage while it runs. + * + * Args: args[0] = output Tensor* (INOUT), args[1] = base_cl, args[2] = spin_iters. + */ + +#include +#include + +#include "tensor.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] // NOLINT(whitespace/braces) +#endif + +#include "intrinsic.h" + +static constexpr int32_t FLOATS_PER_CACHE_LINE = 16; +static constexpr int32_t SLOTS_PER_BLOCK = 3; + +#ifdef PTO_CPUSTUB_HPP +#define dcci(...) \ + do { \ + } while (0) +#endif +#ifndef SINGLE_CACHE_LINE +#define SINGLE_CACHE_LINE 0 +#endif +#ifndef CACHELINE_OUT +#define CACHELINE_OUT 0 +#endif + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *out_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ float *out = reinterpret_cast<__gm__ float *>(out_tensor->buffer.addr) + out_tensor->start_offset; + + int32_t base_cl = static_cast(args[1]); + int32_t spin_iters = static_cast(args[2]); + int32_t block_idx = get_block_idx(args); + int32_t offset = (base_cl + block_idx * SLOTS_PER_BLOCK + 0) * FLOATS_PER_CACHE_LINE; + + volatile int32_t acc = 0; + for (int32_t i = 0; i < spin_iters; i++) { + acc++; // ++ not += i: += i overflows int32 for large spin_iters (UB); volatile keeps the loop + } + (void)acc; + + out[offset] = static_cast(block_idx); + + dcci(&out[offset], SINGLE_CACHE_LINE, CACHELINE_OUT); +} diff --git a/tests/st/a2a3/tensormap_and_ringbuffer/spmd_sync_start_mix_spill/kernels/aiv/kernel_spmd_mix_slow.cpp b/tests/st/a2a3/tensormap_and_ringbuffer/spmd_sync_start_mix_spill/kernels/aiv/kernel_spmd_mix_slow.cpp new file mode 100644 index 0000000000..6d7fcb73dc --- /dev/null +++ b/tests/st/a2a3/tensormap_and_ringbuffer/spmd_sync_start_mix_spill/kernels/aiv/kernel_spmd_mix_slow.cpp @@ -0,0 +1,69 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ + +/** + * SPMD MIX slow kernel (AIV): writes float(block_idx) at cache line + * (base_cl + block_idx * 3 + 1 + sub_block_id), then optionally spins so the task + * holds its cluster while a dependent sync_start cohort pre-stages. + * + * Args: args[0] = output Tensor* (INOUT), args[1] = base_cl, args[2] = spin_iters. + */ + +#include +#include + +#include "tensor.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] // NOLINT(whitespace/braces) +#endif + +#include "intrinsic.h" + +static constexpr int32_t FLOATS_PER_CACHE_LINE = 16; +static constexpr int32_t SLOTS_PER_BLOCK = 3; + +#ifdef PTO_CPUSTUB_HPP +#define dcci(...) \ + do { \ + } while (0) +#endif +#ifndef SINGLE_CACHE_LINE +#define SINGLE_CACHE_LINE 0 +#endif +#ifndef CACHELINE_OUT +#define CACHELINE_OUT 0 +#endif + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *out_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ float *out = reinterpret_cast<__gm__ float *>(out_tensor->buffer.addr) + out_tensor->start_offset; + + int32_t base_cl = static_cast(args[1]); + int32_t spin_iters = static_cast(args[2]); + int32_t block_idx = get_block_idx(args); + int32_t sub_block_id = get_sub_block_id(args); + int32_t offset = (base_cl + block_idx * SLOTS_PER_BLOCK + 1 + sub_block_id) * FLOATS_PER_CACHE_LINE; + + volatile int32_t acc = 0; + for (int32_t i = 0; i < spin_iters; i++) { + acc++; // ++ not += i: += i overflows int32 for large spin_iters (UB); volatile keeps the loop + } + (void)acc; + + out[offset] = static_cast(block_idx); + + dcci(&out[offset], SINGLE_CACHE_LINE, CACHELINE_OUT); +} diff --git a/tests/st/a2a3/tensormap_and_ringbuffer/spmd_sync_start_mix_spill/kernels/aiv/kernel_spmd_write_slow.cpp b/tests/st/a2a3/tensormap_and_ringbuffer/spmd_sync_start_mix_spill/kernels/aiv/kernel_spmd_write_slow.cpp new file mode 100644 index 0000000000..9b0cae98ea --- /dev/null +++ b/tests/st/a2a3/tensormap_and_ringbuffer/spmd_sync_start_mix_spill/kernels/aiv/kernel_spmd_write_slow.cpp @@ -0,0 +1,76 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ + +/** + * SPMD Multi-Block Write Kernel with a spin delay (AIV) + * + * Same result as the plain write kernel, but spins for `spin_iters` before the + * store. A slow producer keeps its cores occupied long enough that a dependent + * consumer is processed as an early-dispatch candidate WHILE the producer is still + * running — the only way a trivial-kernel scene exercises the speculative + * gated-dispatch path (a fast producer finishes before the consumer is ever staged). + * + * out[(base_cl + block_idx) * FLOATS_PER_CACHE_LINE] = float(block_idx) + * + * Args: + * args[0] = output Tensor* (INOUT) + * args[1] = scalar: base_cl (starting cache line index for this task) + * args[2] = scalar: spin_iters (0 = no delay) + */ + +#include +#include + +#include "tensor.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] // NOLINT(whitespace/braces) +#endif + +#include "intrinsic.h" + +static constexpr int32_t FLOATS_PER_CACHE_LINE = 16; + +#ifdef PTO_CPUSTUB_HPP +#define dcci(...) \ + do { \ + } while (0) +#endif +#ifndef SINGLE_CACHE_LINE +#define SINGLE_CACHE_LINE 0 +#endif +#ifndef CACHELINE_OUT +#define CACHELINE_OUT 0 +#endif + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *out_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ float *out = reinterpret_cast<__gm__ float *>(out_tensor->buffer.addr) + out_tensor->start_offset; + + int32_t base_cl = static_cast(args[1]); + int32_t spin_iters = static_cast(args[2]); + int32_t block_idx = get_block_idx(args); + int32_t offset = (base_cl + block_idx) * FLOATS_PER_CACHE_LINE; + + volatile int32_t acc = 0; + for (int32_t i = 0; i < spin_iters; i++) { + acc++; // ++ not += i: += i overflows int32 for large spin_iters (UB); volatile keeps the loop + } + (void)acc; + + out[offset] = static_cast(block_idx); + + dcci(&out[offset], SINGLE_CACHE_LINE, CACHELINE_OUT); +} diff --git a/tests/st/a2a3/tensormap_and_ringbuffer/spmd_sync_start_mix_spill/kernels/orchestration/spmd_sync_start_mix_spill_orch.cpp b/tests/st/a2a3/tensormap_and_ringbuffer/spmd_sync_start_mix_spill/kernels/orchestration/spmd_sync_start_mix_spill_orch.cpp new file mode 100644 index 0000000000..a2919608f0 --- /dev/null +++ b/tests/st/a2a3/tensormap_and_ringbuffer/spmd_sync_start_mix_spill/kernels/orchestration/spmd_sync_start_mix_spill_orch.cpp @@ -0,0 +1,102 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ + +/** + * SPMD sync_start MIX pending-spill Orchestration + * + * Reproduces the gated-cohort rendezvous deadlock on a MIX cohort whose clusters + * spill PER-CORE to pending slots: a FLAGGED AIV producer occupies ALL 48 AIV cores + * (and spins), leaving the 24 AIC cores idle. When the require_sync_start MIX + * consumer pre-stages as an early-dispatch candidate, EVERY one of its 24 clusters + * is mixed — the AIC lands on an idle running slot, both AIVs on the producer's busy + * cores' gated pending slots (drain_stage_cores takes the to_pending=true split path, + * mix_cluster_idle_core_count = 1 per cluster). The rendezvous seed then counts only + * the 24 running AICs while staged_core_mask counts all 72 cores; the 48 pending AIVs + * must promote to close the gap. If the seed/mask counting diverges on this MIX + * per-core split, the doorbells never fire -> the cohort never launches -> its + * consumers never complete -> allocator deadlock. + * + * The MIX consumer writes float(block_idx) at 3 cache lines + * (base_cl + block_idx*3 + {0,1,2}) — AIC slot 0, AIV0 slot 1, AIV1 slot 2. The AIV + * producer writes float(block_idx) at cache line (base_cl + block_idx). + * + * Tasks (deps explicit-only): + * P: AIV block_num=48, base_cl=0, allow_early_resolve=true, spins (occupies all 48 AIV cores) + * C: MIX block_num=24, base_cl=48, require_sync_start=true, dep=[P] (24 AIC idle->running, + * 48 AIV busy->pending) + * + * Args layout: [output] + */ + +#include +#include + +#include "pto_orchestration_api.h" // NOLINT(build/include_subdir) +#include "pto_arg_with_deps.h" // NOLINT(build/include_subdir) + +#define FUNC_SPMD_MIX_AIC 0 +#define FUNC_SPMD_MIX_AIV0 1 +#define FUNC_SPMD_MIX_AIV1 2 +#define FUNC_SPMD_WRITE_AIV 3 + +extern "C" { + +__attribute__((visibility("default"))) PTO2OrchestrationConfig aicpu_orchestration_config(const L2TaskArgs &orch_args) { + (void)orch_args; // NOLINT(readability/casting) + return PTO2OrchestrationConfig{ + .expected_arg_count = 1, + }; +} + +// The producer must stay on-core long enough for the scheduler to pre-stage the +// consumer as an early-dispatch candidate WHILE the producer's blocks are running, +// so the producer's clusters are busy and the consumer spills them to pending slots. +static constexpr int64_t PRODUCER_SPIN_ITERS = 2000000; + +static MixedKernels mix_kernels() { + MixedKernels mk; + mk.aic_kernel_id = FUNC_SPMD_MIX_AIC; + mk.aiv0_kernel_id = FUNC_SPMD_MIX_AIV0; + mk.aiv1_kernel_id = FUNC_SPMD_MIX_AIV1; + return mk; +} + +static PTO2TaskId submit_aiv_producer(const Tensor &out, int16_t block_num, int64_t base_cl) { + L0TaskArgs args; + args.add_inout(out); + args.add_scalar(base_cl); + args.add_scalar(PRODUCER_SPIN_ITERS); + args.launch_spec.set_block_num(block_num); + args.set_allow_early_resolve(true); // flagged: the sync_start consumer may early-dispatch off it + return rt_submit_aiv_task(FUNC_SPMD_WRITE_AIV, args).task_id(); +} + +static void submit_mix_sync_consumer(const Tensor &out, int16_t block_num, int64_t base_cl, PTO2TaskId dep) { + L0TaskArgsWithDeps<4> args; + args.add_inout(out); + args.add_scalar(base_cl); + args.add_scalar(0); // consumer does not spin + args.launch_spec.set_block_num(block_num); + args.launch_spec.set_require_sync_start(true); // atomic cohort launch via the drain + rendezvous + args.add_dep(dep); // sole flagged producer -> early-dispatch candidate + rt_submit_task(mix_kernels(), args); +} + +__attribute__((visibility("default"))) void aicpu_orchestration_entry(const L2TaskArgs &orch_args) { + const Tensor &ext_output = orch_args.tensor(0).ref(); + + PTO2TaskId prod = submit_aiv_producer(ext_output, 48, 0); + submit_mix_sync_consumer(ext_output, 24, 48, prod); + + LOG_INFO_V9("[spmd_sync_start_mix_spill] flagged AIV producer (48) + sync_start MIX consumer (24) submitted"); +} + +} // extern "C" diff --git a/tests/st/a2a3/tensormap_and_ringbuffer/spmd_sync_start_mix_spill/test_spmd_sync_start_mix_spill.py b/tests/st/a2a3/tensormap_and_ringbuffer/spmd_sync_start_mix_spill/test_spmd_sync_start_mix_spill.py new file mode 100644 index 0000000000..a8ef74eb6b --- /dev/null +++ b/tests/st/a2a3/tensormap_and_ringbuffer/spmd_sync_start_mix_spill/test_spmd_sync_start_mix_spill.py @@ -0,0 +1,99 @@ +#!/usr/bin/env python3 +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +"""sync_start MIX per-core pending-spill: a flagged AIV producer occupies all 48 AIV cores (and +spins), leaving the 24 AIC cores idle. The require_sync_start MIX consumer then pre-stages with +EVERY cluster mixed — AIC on an idle running slot, both AIVs on the producer's busy cores' gated +pending slots. Exercises the rendezvous seed/mask counting on the MIX per-core split path +(drain_stage_cores to_pending=true, mix_cluster_idle_core_count=1/cluster + Case 3.3 promote for +the 48 pending AIVs). A counting mismatch stalls the rendezvous -> gated cores never launch -> +allocator deadlock.""" + +import torch +from simpler.task_interface import ArgDirection as D + +from simpler_setup import SceneTestCase, TaskArgsBuilder, Tensor, scene_test + +FLOATS_PER_CACHE_LINE = 16 +SLOTS_PER_BLOCK = 3 # MIX consumer block writes 3 cache lines: AIC slot 0, AIV0 slot 1, AIV1 slot 2 +PRODUCER_BLOCKS = 48 # AIV producer: 1 cache line per block, base_cl 0 +PRODUCER_BASE_CL = 0 +CONSUMER_BLOCKS = 24 # MIX consumer: 3 cache lines per block +CONSUMER_BASE_CL = 48 +TOTAL_CL = 120 # 48 (producer) + 24*3 (consumer) + + +@scene_test(level=2, runtime="tensormap_and_ringbuffer") +class TestSpmdSyncStartMixSpill(SceneTestCase): + RTOL = 0 + ATOL = 0 + + CALLABLE = { + "orchestration": { + "source": "kernels/orchestration/spmd_sync_start_mix_spill_orch.cpp", + "function_name": "aicpu_orchestration_entry", + "signature": [D.INOUT], + }, + "incores": [ + { + "func_id": 0, + "name": "SPMD_MIX_AIC", + "source": "kernels/aic/kernel_spmd_mix_slow.cpp", + "core_type": "aic", + "signature": [D.INOUT], + }, + { + "func_id": 1, + "name": "SPMD_MIX_AIV0", + "source": "kernels/aiv/kernel_spmd_mix_slow.cpp", + "core_type": "aiv", + "signature": [D.INOUT], + }, + { + "func_id": 2, + "name": "SPMD_MIX_AIV1", + "source": "kernels/aiv/kernel_spmd_mix_slow.cpp", + "core_type": "aiv", + "signature": [D.INOUT], + }, + { + "func_id": 3, + "name": "SPMD_WRITE_AIV", + "source": "kernels/aiv/kernel_spmd_write_slow.cpp", + "core_type": "aiv", + "signature": [D.INOUT], + }, + ], + } + + CASES = [ + { + "name": "Case1", + "platforms": ["a2a3sim", "a2a3"], + "config": {"aicpu_thread_num": 3, "block_dim": 24}, + "params": {}, + } + ] + + def generate_args(self, params): + return TaskArgsBuilder(Tensor("output", torch.zeros(TOTAL_CL * FLOATS_PER_CACHE_LINE, dtype=torch.float32))) + + def compute_golden(self, args, params): + out = args.output + # AIV producer: 1 cache line per block. + for block_idx in range(PRODUCER_BLOCKS): + out[(PRODUCER_BASE_CL + block_idx) * FLOATS_PER_CACHE_LINE] = float(block_idx) + # MIX consumer: 3 cache lines per block (AIC slot 0, AIV0 slot 1, AIV1 slot 2). + for block_idx in range(CONSUMER_BLOCKS): + for slot in range(SLOTS_PER_BLOCK): + out[(CONSUMER_BASE_CL + block_idx * SLOTS_PER_BLOCK + slot) * FLOATS_PER_CACHE_LINE] = float(block_idx) + + +if __name__ == "__main__": + SceneTestCase.run_module(__name__) diff --git a/tests/ut/cpp/a2a3/test_wiring.cpp b/tests/ut/cpp/a2a3/test_wiring.cpp index c491fca56d..dec905f18f 100644 --- a/tests/ut/cpp/a2a3/test_wiring.cpp +++ b/tests/ut/cpp/a2a3/test_wiring.cpp @@ -522,6 +522,202 @@ TEST_F(WiringTest, EarlyDispatchReleaseConsumesDoorbellMask) { EXPECT_EQ(get_test_reg_stub_value(), 0); } +TEST_F(WiringTest, SyncStartDoorbellPassHasOneOwner) { + PTO2TaskPayload payload{}; + + for (int iteration = 0; iteration < 1000; iteration++) { + payload.early_dispatch_launch_state.store(PTO2_EARLY_DISPATCH_LAUNCH_NONE, std::memory_order_relaxed); + std::atomic start{false}; + bool first_won = false; + bool second_won = false; + + std::thread first([&] { + while (!start.load(std::memory_order_acquire)) {} + first_won = PTO2SchedulerState::try_claim_early_dispatch_launch(payload); + }); + std::thread second([&] { + while (!start.load(std::memory_order_acquire)) {} + second_won = PTO2SchedulerState::try_claim_early_dispatch_launch(payload); + }); + + start.store(true, std::memory_order_release); + first.join(); + second.join(); + + EXPECT_NE(first_won, second_won); + EXPECT_EQ( + payload.early_dispatch_launch_state.load(std::memory_order_acquire), PTO2_EARLY_DISPATCH_LAUNCH_RINGING + ); + } +} + +TEST_F(WiringTest, SyncStartDrainFinalizeRetriesProducerFirstRendezvous) { + alignas(64) PTO2TaskSlotState sync_consumer, downstream; + init_slot(sync_consumer, PTO2_TASK_PENDING, 1, 1); + init_slot(downstream, PTO2_TASK_PENDING, 1, 1); + + sync_consumer.active_mask = ActiveMask(PTO2_SUBTASK_MASK_AIV0); + sync_consumer.active_mask.set_sync_start(); + sync_consumer.allow_early_resolve = true; + sync_consumer.logical_block_num = 2; + sync_consumer.next_block_idx.store(2, std::memory_order_relaxed); + sched.record_published_blocks(sync_consumer, sync_consumer.logical_block_num); + sync_consumer.payload->staged_core_mask[0].store(0b11, std::memory_order_relaxed); + sync_consumer.payload->running_slot_count.store(0, std::memory_order_relaxed); + sync_consumer.payload->early_dispatch_state.store(PTO2_EARLY_DISPATCH_STAGING, std::memory_order_relaxed); + + downstream.payload->fanin_actual_count = 1; + PTO2DepListEntry dep{}; + dep.slot_state = &downstream; + sync_consumer.fanout_head = &dep; + + // Producer release wins before drain publishes its running-slot seed. With no pending + // promotions, drain finalize is the only remaining rendezvous retry. + EXPECT_TRUE(sched.try_early_dispatch_release(sync_consumer)); + EXPECT_EQ(sync_consumer.payload->early_dispatch_state.load(), PTO2_EARLY_DISPATCH_DISPATCHED); + EXPECT_EQ(sync_consumer.payload->early_dispatch_launch_state.load(), PTO2_EARLY_DISPATCH_LAUNCH_NONE); + EXPECT_EQ(downstream.payload->dispatch_fanin.load(), 0); + + sync_consumer.payload->running_slot_count.store(2, std::memory_order_seq_cst); + EXPECT_TRUE(sched.retry_sync_start_rendezvous_after_drain(sync_consumer)); + EXPECT_EQ(sync_consumer.payload->early_dispatch_launch_state.load(), PTO2_EARLY_DISPATCH_LAUNCH_COMPLETE); + EXPECT_EQ(sync_consumer.payload->dispatch_propagated.load(), 1); + EXPECT_EQ(downstream.payload->dispatch_fanin.load(), 1); + + EXPECT_FALSE(sched.retry_sync_start_rendezvous_after_drain(sync_consumer)); + EXPECT_EQ(downstream.payload->dispatch_fanin.load(), 1); +} + +TEST_F(WiringTest, ArmedEarlySyncDrainOwnsFinalReadyRoute) { + alignas(64) PTO2TaskSlotState sync_consumer; + init_slot(sync_consumer, PTO2_TASK_PENDING, 1, 1); + sync_consumer.active_mask = ActiveMask(PTO2_SUBTASK_MASK_AIV0); + sync_consumer.active_mask.set_sync_start(); + sync_consumer.logical_block_num = 2; + sync_consumer.payload->early_dispatch_state.store(PTO2_EARLY_DISPATCH_STAGING, std::memory_order_relaxed); + + ASSERT_TRUE(PTO2SchedulerState::try_claim_early_sync_drain(*sync_consumer.payload)); + PTO2SchedulerState::mark_early_sync_drain_armed(*sync_consumer.payload); + EXPECT_TRUE(sched.release_fanin_and_check_ready(sync_consumer)); + auto shape = static_cast(sync_consumer.active_mask.to_shape()); + EXPECT_EQ(sched.ready_sync_queues[shape].pop(), nullptr); + EXPECT_EQ(sync_consumer.payload->early_dispatch_state.load(), PTO2_EARLY_DISPATCH_DISPATCHED); + EXPECT_EQ( + sync_consumer.payload->early_sync_drain_state.load(), + PTO2_EARLY_SYNC_DRAIN_OWNER | PTO2_EARLY_SYNC_DRAIN_ARMED | PTO2_EARLY_SYNC_DRAIN_READY + ); +} + +TEST_F(WiringTest, ArmedEarlySyncDrainKeepsEveryStagerGatedAfterReady) { + alignas(64) PTO2TaskPayload payload{}; + payload.early_dispatch_state.store(PTO2_EARLY_DISPATCH_STAGING, std::memory_order_relaxed); + + ASSERT_TRUE(PTO2SchedulerState::try_claim_early_sync_drain(payload)); + PTO2SchedulerState::mark_early_sync_drain_armed(payload); + + std::atomic start{false}; + std::atomic before_ready{0}; + std::atomic ready_published{false}; + bool gated[4] = {false, false, false, false}; + std::thread stagers[4]; + for (int i = 0; i < 4; i++) { + stagers[i] = std::thread([&, i] { + while (!start.load(std::memory_order_acquire)) {} + if (i >= 2) { + while (!ready_published.load(std::memory_order_acquire)) {} + } + bool force_gate = PTO2SchedulerState::owns_early_sync_drain(payload); + gated[i] = PTO2SchedulerState::should_gate_early_dispatch( + force_gate, payload.early_dispatch_state.load(std::memory_order_relaxed) + ); + if (i < 2) before_ready.fetch_add(1, std::memory_order_release); + }); + } + start.store(true, std::memory_order_release); + while (before_ready.load(std::memory_order_acquire) != 2) {} + payload.early_dispatch_state.store(PTO2_EARLY_DISPATCH_DISPATCHED, std::memory_order_seq_cst); + bool ready_has_owner = PTO2SchedulerState::publish_ready_to_early_sync_drain(payload); + ready_published.store(true, std::memory_order_release); + for (auto &stager : stagers) + stager.join(); + + EXPECT_TRUE(ready_has_owner); + for (bool was_gated : gated) + EXPECT_TRUE(was_gated); + + alignas(64) PTO2TaskPayload normal_ready{}; + normal_ready.early_dispatch_state.store(PTO2_EARLY_DISPATCH_DISPATCHED, std::memory_order_relaxed); + EXPECT_FALSE(PTO2SchedulerState::publish_ready_to_early_sync_drain(normal_ready)); + EXPECT_FALSE(PTO2SchedulerState::owns_early_sync_drain(normal_ready)); + EXPECT_FALSE( + PTO2SchedulerState::should_gate_early_dispatch( + PTO2SchedulerState::owns_early_sync_drain(normal_ready), + normal_ready.early_dispatch_state.load(std::memory_order_relaxed) + ) + ); +} + +TEST_F(WiringTest, DrainFinishBetweenReleasePhasesRetainsOwner) { + alignas(64) PTO2TaskSlotState sync_consumer; + init_slot(sync_consumer, PTO2_TASK_PENDING, 1, 1); + sync_consumer.active_mask = ActiveMask(PTO2_SUBTASK_MASK_AIV0); + sync_consumer.active_mask.set_sync_start(); + sync_consumer.logical_block_num = 2; + sync_consumer.payload->early_dispatch_state.store(PTO2_EARLY_DISPATCH_STAGING, std::memory_order_relaxed); + + ASSERT_TRUE(PTO2SchedulerState::try_claim_early_sync_drain(*sync_consumer.payload)); + PTO2SchedulerState::mark_early_sync_drain_armed(*sync_consumer.payload); + + // READY publication may lag the STAGING-to-DISPATCHED release transition. + EXPECT_FALSE(sched.try_early_dispatch_release(sync_consumer)); + sync_consumer.next_block_idx.store(sync_consumer.logical_block_num, std::memory_order_seq_cst); + PTO2SchedulerState::finish_early_sync_drain(*sync_consumer.payload); + + EXPECT_TRUE(PTO2SchedulerState::publish_ready_to_early_sync_drain(*sync_consumer.payload)); + EXPECT_EQ( + sync_consumer.payload->early_sync_drain_state.load(), + PTO2_EARLY_SYNC_DRAIN_OWNER | PTO2_EARLY_SYNC_DRAIN_ARMED | PTO2_EARLY_SYNC_DRAIN_READY | + PTO2_EARLY_SYNC_DRAIN_COMPLETE + ); +} + +TEST_F(WiringTest, CancelledEarlySyncDrainRoutesProducerRelease) { + alignas(64) PTO2TaskSlotState sync_consumer; + init_slot(sync_consumer, PTO2_TASK_PENDING, 1, 1); + sync_consumer.active_mask = ActiveMask(PTO2_SUBTASK_MASK_AIV0); + sync_consumer.active_mask.set_sync_start(); + sync_consumer.logical_block_num = 2; + sync_consumer.payload->early_dispatch_state.store(PTO2_EARLY_DISPATCH_STAGING, std::memory_order_relaxed); + + ASSERT_TRUE(PTO2SchedulerState::try_claim_early_sync_drain(*sync_consumer.payload)); + sched.cancel_early_sync_drain(sync_consumer); + EXPECT_TRUE(sched.release_fanin_and_check_ready(sync_consumer)); + + auto shape = static_cast(sync_consumer.active_mask.to_shape()); + EXPECT_EQ(sched.ready_sync_queues[shape].pop(), &sync_consumer); + EXPECT_EQ(sched.ready_sync_queues[shape].pop(), nullptr); + EXPECT_EQ(sched.early_sync_start_queue.pop(), &sync_consumer); +} + +TEST_F(WiringTest, ProducerReleaseTransfersReadyRouteToCancellingDrain) { + alignas(64) PTO2TaskSlotState sync_consumer; + init_slot(sync_consumer, PTO2_TASK_PENDING, 1, 1); + sync_consumer.active_mask = ActiveMask(PTO2_SUBTASK_MASK_AIV0); + sync_consumer.active_mask.set_sync_start(); + sync_consumer.logical_block_num = 2; + sync_consumer.payload->early_dispatch_state.store(PTO2_EARLY_DISPATCH_STAGING, std::memory_order_relaxed); + + ASSERT_TRUE(PTO2SchedulerState::try_claim_early_sync_drain(*sync_consumer.payload)); + EXPECT_TRUE(sched.release_fanin_and_check_ready(sync_consumer)); + auto shape = static_cast(sync_consumer.active_mask.to_shape()); + EXPECT_EQ(sched.ready_sync_queues[shape].pop(), nullptr); + + sched.cancel_early_sync_drain(sync_consumer); + EXPECT_EQ(sched.ready_sync_queues[shape].pop(), &sync_consumer); + EXPECT_EQ(sched.ready_sync_queues[shape].pop(), nullptr); + EXPECT_EQ(sched.early_sync_start_queue.pop(), nullptr); +} + TEST_F(WiringTest, EarlyDispatchBlockedByUnflaggedProducer) { alignas(64) PTO2TaskSlotState task_slot; alignas(64) PTO2TaskSlotState p_flagged, q_unflagged;