From 4998ecd3098c4c6e4c2636526d33df75e71a39ad Mon Sep 17 00:00:00 2001 From: vegetabledoww <17863935975@163.com> Date: Tue, 28 Jul 2026 19:22:36 -0700 Subject: [PATCH] fix(runtime): avoid global drain for local early sync-start Stage an early sync-start cohort directly on its owning scheduler when that scheduler's running and pending slots can hold every block. Keep the stop-the-world drain as the fallback when capacity must be gathered across scheduler threads, and preserve the one-shot producer/stager rendezvous. Share available-block accounting across AIC, AIV, and MIX placement, cover launch ownership and pending-slot cases with unit tests, and add A2A3/A5 regression scenes that require early_dispatch without a drain when the cohort fits locally. Fixes #1548 --- .../host_build_graph/docs/RUNTIME_LOGIC.md | 9 + .../runtime/pto_runtime2_types.h | 5 +- .../runtime/scheduler/pto_scheduler.h | 27 +-- .../scheduler/scheduler_completion.cpp | 63 +++---- .../runtime/scheduler/scheduler_context.h | 17 +- .../runtime/scheduler/scheduler_dispatch.cpp | 45 +++-- .../runtime/scheduler/scheduler_types.h | 17 ++ .../docs/RUNTIME_LOGIC.md | 72 ++++---- .../runtime/pto_runtime2_types.h | 5 +- .../runtime/scheduler/pto_scheduler.h | 38 ++-- .../scheduler/scheduler_completion.cpp | 63 +++---- .../runtime/scheduler/scheduler_context.h | 17 +- .../runtime/scheduler/scheduler_dispatch.cpp | 45 +++-- .../runtime/scheduler/scheduler_types.h | 17 ++ .../docs/RUNTIME_LOGIC.md | 49 +++-- .../runtime/scheduler/pto_scheduler.h | 12 +- .../scheduler/scheduler_completion.cpp | 37 ++-- .../runtime/scheduler/scheduler_context.h | 8 +- .../runtime/scheduler/scheduler_dispatch.cpp | 107 +++++++---- .../runtime/scheduler/scheduler_types.h | 15 ++ .../include/common/l2_swimlane_profiling.h | 4 +- .../kernels/aiv/kernel_spmd_spin.cpp | 102 +++++++++++ .../sync_start_early_local_owner_orch.cpp | 150 ++++++++++++++++ .../test_sync_start_early_local_owner.py | 169 ++++++++++++++++++ .../spmd_sync_start_mix_spill_orch.cpp | 2 +- .../test_spmd_sync_start_mix_spill.py | 2 +- .../kernels/aiv/kernel_spmd_spin.cpp | 102 +++++++++++ .../sync_start_early_local_owner_orch.cpp | 142 +++++++++++++++ .../test_sync_start_early_local_owner.py | 163 +++++++++++++++++ tests/ut/cpp/a2a3/test_scheduler_state.cpp | 57 ++++++ tests/ut/cpp/a2a3/test_wiring.cpp | 43 ++++- tests/ut/cpp/a5/test_scheduler_state.cpp | 57 ++++++ tests/ut/cpp/a5/test_wiring.cpp | 118 ++++++++++++ 33 files changed, 1533 insertions(+), 246 deletions(-) create mode 100644 tests/st/a2a3/tensormap_and_ringbuffer/dfx/l2_swimlane/kernels/aiv/kernel_spmd_spin.cpp create mode 100644 tests/st/a2a3/tensormap_and_ringbuffer/dfx/l2_swimlane/kernels/orchestration/sync_start_early_local_owner_orch.cpp create mode 100644 tests/st/a2a3/tensormap_and_ringbuffer/dfx/l2_swimlane/test_sync_start_early_local_owner.py create mode 100644 tests/st/a5/tensormap_and_ringbuffer/dfx/l2_swimlane/kernels/aiv/kernel_spmd_spin.cpp create mode 100644 tests/st/a5/tensormap_and_ringbuffer/dfx/l2_swimlane/kernels/orchestration/sync_start_early_local_owner_orch.cpp create mode 100644 tests/st/a5/tensormap_and_ringbuffer/dfx/l2_swimlane/test_sync_start_early_local_owner.py diff --git a/src/a2a3/runtime/host_build_graph/docs/RUNTIME_LOGIC.md b/src/a2a3/runtime/host_build_graph/docs/RUNTIME_LOGIC.md index 6bf96777e7..cd10a5409f 100644 --- a/src/a2a3/runtime/host_build_graph/docs/RUNTIME_LOGIC.md +++ b/src/a2a3/runtime/host_build_graph/docs/RUNTIME_LOGIC.md @@ -592,6 +592,15 @@ is seeded by the device **boot classify** (`on_orchestration_done`), which scans the submitted tasks once and either pushes the fanin-free ones to the ready queue or registers each remaining task on its first unmet producer. +**Early staging status.** Early producer propagation is currently disabled in +HBG: `propagate_dispatch_fanin` is a stub, so the polling path does not populate +`early_dispatch_queues` or `early_sync_start_queue`. The scheduler retains the +early-staging machinery as a dormant mirror of the TMR runtime for a future +consumer-pull publication path. If that path populates the queues, a +`sync_start` cohort that fits one scheduler's idle and pending slots stages +locally; larger cohorts use the generation-tagged global drain. This is not an +active end-to-end HBG fast path today. + **Phase 2 — Dispatch**: - For each idle core: pop a task from the matching shape-based ready queue (lock-free MPMC Vyukov queue, one per resource shape) diff --git a/src/a2a3/runtime/host_build_graph/runtime/pto_runtime2_types.h b/src/a2a3/runtime/host_build_graph/runtime/pto_runtime2_types.h index 6fb4ed5294..933c14b1a4 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/pto_runtime2_types.h +++ b/src/a2a3/runtime/host_build_graph/runtime/pto_runtime2_types.h @@ -274,8 +274,9 @@ struct PTO2TaskPayload { // // 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. + // splits them between release and late-stager owners; a sync_start cohort keeps + // the completed mask stable for its single launch owner, whether staging is local + // or uses the global drain fallback. 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 diff --git a/src/a2a3/runtime/host_build_graph/runtime/scheduler/pto_scheduler.h b/src/a2a3/runtime/host_build_graph/runtime/scheduler/pto_scheduler.h index 720b38f4f9..3f20a188fd 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/scheduler/pto_scheduler.h +++ b/src/a2a3/runtime/host_build_graph/runtime/scheduler/pto_scheduler.h @@ -634,16 +634,16 @@ struct PTO2SchedulerState { 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. + // they need an atomic all-or-nothing stage, 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 and one owner selects the + // local stage or global-drain fallback. // // 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. + // per-shape dispatch_shape. An EARLY sync cohort carries a non-zero src_payload + // gate; its owner stages locally when one tracker fits and uses the global drain + // otherwise. HBG producer propagation currently leaves this queue dormant. PTO2ReadyQueue early_sync_start_queue; static inline void ring_one_doorbell(uint64_t reg_addr, uint32_t token) { @@ -694,8 +694,8 @@ struct PTO2SchedulerState { } // 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. + // the NONE->RINGING launch latch and invokes this exactly once after local or + // global 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); @@ -764,12 +764,17 @@ struct PTO2SchedulerState { // 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) { + // running_slot_count is the publication seed: every staged_core_mask OR + // happens-before its final store. Read the seed first, then the mask, so + // observing the final count cannot be paired with a partially published + // mask when producer release races staging. + int32_t running_cores = slot_state.payload->running_slot_count.load(std::memory_order_seq_cst); 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 (running_cores != 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; @@ -781,7 +786,7 @@ struct PTO2SchedulerState { return true; } - inline bool retry_sync_start_rendezvous_after_drain(PTO2TaskSlotState &slot_state) { + inline bool retry_sync_start_rendezvous_after_staging(PTO2TaskSlotState &slot_state) { if (!maybe_rendezvous_ring(slot_state)) return false; propagate_dispatch_fanin(slot_state); return true; diff --git a/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_completion.cpp b/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_completion.cpp index 2bb0e109b9..bb46493ad3 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_completion.cpp +++ b/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_completion.cpp @@ -464,42 +464,35 @@ bool SchedulerContext::enter_drain_mode(PTO2TaskSlotState *slot_state, int32_t b 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) { - // 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(); - } - } + total += core_trackers_[t].count_available_blocks(shape, core_mask, include_pending); } return total; } -// 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 +// Stage one thread's share of a sync_start cohort: CAS-claim block indices and +// publish them onto THIS thread's own cores. The global drain calls this +// concurrently on peers; the local fast path calls it only after proving this +// tracker can hold the entire cohort. Each thread touches only its own tracker +// and doorbell-table entries; the CAS on next_block_idx and 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) { +// Gated (early) sync_start staging, whether local or through the global drain, pre-stages +// every block behind its doorbell (prepare_block_for_dispatch is force-gated for the claimed +// range) and defers 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. +SchedulerContext::SyncStartStageResult SchedulerContext::stage_sync_start_cores( + PTO2TaskSlotState *slot_state, int32_t block_num, int32_t thread_idx, bool gated, + [[maybe_unused]] bool record_drain_phases +) { CoreTracker &tracker = core_trackers_[thread_idx]; PTO2ResourceShape shape = slot_state->active_mask.to_shape(); uint8_t core_mask = slot_state->active_mask.core_mask(); bool mix_split = gated && shape == PTO2ResourceShape::MIX; - int32_t running_staged = 0; + SyncStartStageResult result; // 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 @@ -512,8 +505,9 @@ SchedulerContext::drain_stage_cores(PTO2TaskSlotState *slot_state, int32_t block int32_t start = 0; int32_t claim = slot_state->claim_block_range(block_num, avail, start); if (claim == 0) return; + result.staged_blocks += claim; #if SIMPLER_DFX - bool sub_prof = l2_swimlane_level_ >= L2SwimlaneLevel::SCHED_PHASES; + bool sub_prof = record_drain_phases && l2_swimlane_level_ >= L2SwimlaneLevel::SCHED_PHASES; uint64_t prep_t0 = sub_prof ? get_sys_cnt_aicpu() : 0; #endif PublishHandle handles[CoreTracker::MAX_CLUSTERS * 3]; @@ -527,7 +521,7 @@ SchedulerContext::drain_stage_cores(PTO2TaskSlotState *slot_state, int32_t block 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); + result.running_cores += 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 @@ -584,7 +578,7 @@ SchedulerContext::drain_stage_cores(PTO2TaskSlotState *slot_state, int32_t block 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 (gated && shape != PTO2ResourceShape::MIX && !to_pending) result.running_cores += handle_count; } }; @@ -600,7 +594,7 @@ SchedulerContext::drain_stage_cores(PTO2TaskSlotState *slot_state, int32_t block stage(tracker.get_pending_core_offset_states(shape), /*to_pending=*/true); } } - return running_staged; + return result; } // Called by each scheduler thread when drain_state_.sync_start_pending != 0. @@ -699,13 +693,14 @@ void SchedulerContext::handle_drain_mode(int32_t thread_idx, [[maybe_unused]] ui #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); + SyncStartStageResult staged = + stage_sync_start_cores(slot_state, block_num, thread_idx, gated, /*record_drain_phases=*/true); #if SIMPLER_DFX - // out param carries the PURE drain_stage_cores wall (build_payload + MMIO publish of + // out param carries the PURE stage_sync_start_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_running_staged.fetch_add(staged.running_cores, std::memory_order_acq_rel); drain_state_.drain_stage_done_mask.fetch_or(1u << thread_idx, std::memory_order_release); if (!coordinator) { @@ -750,7 +745,7 @@ void SchedulerContext::handle_drain_mode(int32_t thread_idx, [[maybe_unused]] ui // 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); + sched_->retry_sync_start_rendezvous_after_staging(*slot_state); } else { sched_->propagate_dispatch_fanin(*slot_state); } diff --git a/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_context.h b/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_context.h index 3aa699a446..680b3046d7 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_context.h +++ b/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_context.h @@ -497,11 +497,18 @@ class SchedulerContext { bool enter_drain_mode(PTO2TaskSlotState *slot_state, int32_t block_num); 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 + struct SyncStartStageResult { + int32_t staged_blocks{0}; + int32_t running_cores{0}; + }; + // CAS-claim sync_start block indices and stage them onto THIS thread's own + // cores. The global drain invokes this in parallel on every scheduler; + // the local fast path invokes it once after proving this tracker has enough + // capacity. record_drain_phases keeps local work attributed to EarlyDispatch. + SyncStartStageResult stage_sync_start_cores( + PTO2TaskSlotState *slot_state, int32_t block_num, int32_t thread_idx, bool gated, bool record_drain_phases + ); + // out_stage_wall_cycles (profiling only): cycles this thread spent in stage_sync_start_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); diff --git a/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_dispatch.cpp b/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_dispatch.cpp index 24fd8491ea..4a8c467fab 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_dispatch.cpp +++ b/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_dispatch.cpp @@ -257,7 +257,7 @@ int SchedulerContext::prepare_block_for_dispatch( 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 + // (promoted on completion). Gated sync_start staging 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) { @@ -783,9 +783,9 @@ SchedulerContext::early_dispatch_shape(int32_t thread_idx, PTO2ResourceShape sha // 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. +// before PENDING, cross-thread idle gating: MIX-IDLE ▶ c/v-IDLE ▶ MIX-PEND ▶ c/v-PEND). +// Each handles its sync_start cohort FIRST as Tier 0: an exact local fit stages on +// one owner, while a capacity-short cohort falls back to the global drain. // 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 @@ -805,16 +805,15 @@ int32_t SchedulerContext::try_early_dispatch( if (sched_->ready_sync_queues[s].size() > 0 || sched_->ready_queues[s].size() > 0) return 0; } + int32_t total_staged = 0; + // ===== 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. + // early_dispatch_shape's per-thread partial range-claim: a partial cohort would strand + // gated blocks nobody can ring. After claiming the per-task owner, first use the local + // tracker when it can hold the entire cohort; only the capacity-short case arms the + // stop-the-world drain. Both paths force-gate every block even if producer release races + // STAGING -> DISPATCHED. A non-STAGING pop was already released and is dropped. uint64_t sync_task_id_snapshot = 0; if (PTO2TaskSlotState *c = sched_->early_sync_start_queue.pop_tagged(&sync_task_id_snapshot)) { bool current_sync_task = @@ -822,6 +821,25 @@ int32_t SchedulerContext::try_early_dispatch( if (current_sync_task && 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 (drain_state_.sync_start_pending.load(std::memory_order_acquire) == 0 && + tracker.count_available_blocks( + c->active_mask.to_shape(), c->active_mask.core_mask(), /*include_pending=*/true + ) >= c->logical_block_num) { + // From this point onward the operation is all-or-nothing. Only this + // scheduler mutates its tracker, and global drain coordinators must + // wait for this scheduler's generation-tagged ack before inspecting it. + PTO2SchedulerState::mark_early_sync_drain_armed(*c->payload); + always_assert(c->next_block_idx.load(std::memory_order_seq_cst) == 0); + SyncStartStageResult staged = stage_sync_start_cores( + c, c->logical_block_num, thread_idx, /*gated=*/true, /*record_drain_phases=*/false + ); + always_assert(staged.staged_blocks == c->logical_block_num); + c->payload->running_slot_count.store( + static_cast(staged.running_cores), std::memory_order_seq_cst + ); + sched_->retry_sync_start_rendezvous_after_staging(*c); + PTO2SchedulerState::finish_early_sync_drain(*c->payload); + total_staged += staged.staged_blocks; } else if (enter_drain_mode(c, c->logical_block_num)) { PTO2SchedulerState::mark_early_sync_drain_armed(*c->payload); } else { @@ -833,7 +851,6 @@ int32_t SchedulerContext::try_early_dispatch( // 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, [&](PTO2ResourceShape shape, CoreTracker::DispatchPhase phase) { @@ -1256,7 +1273,7 @@ int32_t SchedulerContext::resolve_and_dispatch(Runtime *runtime, int32_t thread_ 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 + // stage_sync_start_cores). The many no-op entries — ack + availability-insufficient // reset or follower 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) { diff --git a/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_types.h b/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_types.h index 91d2d226c5..f303a01cee 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_types.h +++ b/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_types.h @@ -454,6 +454,23 @@ class alignas(64) CoreTracker { return get_mix_running_cluster_offset_states(core_mask).count(); } + // Number of whole logical blocks this tracker can accept for sync_start + // staging. AIC/AIV count cores; MIX counts clusters because one logical MIX + // block may occupy multiple cores in the same cluster. Gated early staging + // includes pending slots, while ready staging is restricted to idle slots. + int32_t count_available_blocks(PTO2ResourceShape shape, uint8_t core_mask, bool include_pending) const { + if (shape == PTO2ResourceShape::MIX) { + return include_pending ? count_mix_split_clusters(core_mask) : count_mix_running_clusters(core_mask); + } + if (shape == PTO2ResourceShape::DUMMY) return 0; + + int32_t available = get_idle_core_offset_states(shape).count(); + if (include_pending) { + available += get_pending_core_offset_states(shape).count(); + } + return available; + } + BitStates get_pending_core_offset_states(PTO2ResourceShape shape) const { if (shape == PTO2ResourceShape::MIX) { // Shape-level query kept conservative for legacy callers/tests. 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 d9eab284f4..2a4cb56be4 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md @@ -555,7 +555,7 @@ Each scheduler thread runs a tight loop with two main phases: **Phase 2 — Dispatch** (full model in §8.6): -- Drain each source (normal ready ▸ speculative early) in occupancy order — `sync_start` +- Service 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` @@ -644,45 +644,53 @@ pending slot, promoted on completion). This order lives in one shared skeleton, | EARLY (speculative) | `early_dispatch_queues[MIX\|AIC\|AIV]` | `early_sync_start_queue` (single) | A task routes to the sync lane iff `task_attrs.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`), +sync lane is serviced 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 ownership** — 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** — scheduler thread 0 coordinates the generation-tagged ack tree and - checks `count_global_available >= block_num` *before* staging. If short, it advances the - attempt and retries after completions free cores. A cohort is fully staged or not at all — - never partial. -3. **Parallel stage** — after thread 0 broadcasts the completed root token, each scheduler - thread 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 ownership + 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). +sync lane is a single, shape-agnostic queue because its per-task owner must make one +all-or-nothing decision. The owner stages locally when its tracker can hold the complete +cohort and falls back to the global drain only when local capacity is short. Both paths use +the same gated, core-counting rendezvous. + +#### sync_start local fast path, drain fallback, and rendezvous + +A sync_start cohort of `block_num` logical blocks must occupy all of its required core slots +before any core runs. The early Tier-0 owner first asks its own `CoreTracker` for complete +idle+pending capacity. AIC/AIV capacity is counted in cores; MIX capacity is counted in +logical clusters even though each block may stage multiple cores. + +1. **Local Case A** — if no global drain is already published and one owner tracker has at + least `block_num` slots, that scheduler force-gates and stages the entire cohort + synchronously. It does not modify `sync_start_pending`, the drain generation, or ack + tokens. Once staging starts it cannot cancel or partially fall back. +2. **Global Case B** — otherwise `enter_drain_mode` uses a CAS on `sync_start_pending` + (0 → −1) to select one global drain. Scheduler thread 0 coordinates the + generation-tagged ack tree and checks global capacity before releasing parallel staging. + If short, it advances the attempt and retries after completions free cores. +3. **Parallel fallback stage** — after thread 0 broadcasts the completed root token, each + scheduler CAS-claims a block range and stages only its own cores. The generation-tagged + barrier also serializes a global coordinator with any local staging that raced drain + publication: the local scheduler cannot ack until its tracker mutations are complete. +4. **Rendezvous launch** — both paths publish every `staged_core_mask` word before storing + the final `running_slot_count` seed. `maybe_rendezvous_ring` reads the seed first and then + the mask; when the counts match **and** the producer has released, one launch-latch winner + rings every gated core's doorbell together. + +Per-task ownership plus all-or-nothing admission prevents two cohorts from each partially +occupying the available cluster set. The global drain remains single-owner; independent local +cohorts may proceed on disjoint scheduler-owned trackers. 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) Producer-side `propagate_dispatch_fanin` 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. +`published_block_count == logical_block_num`. Normal dispatch, regular early staging, and +both local and global sync staging 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 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 d10dc107b6..9e50b94819 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 @@ -270,8 +270,9 @@ struct PTO2TaskPayload { // // 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. + // splits them between release and late-stager owners; a sync_start cohort keeps + // the completed mask stable for its single launch owner, whether staging is local + // or uses the global drain fallback. 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 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 72cd0f096c..b790e3aab3 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 @@ -749,16 +749,17 @@ struct PTO2SchedulerState { 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. + // they need an atomic all-or-nothing owner stage, 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. The owner stages locally + // when its tracker fits and falls back to the global drain otherwise. // // 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. + // per-shape dispatch_shape. An EARLY sync cohort always carries a non-zero + // src_payload gate; it stages directly when one owner tracker can hold the whole + // cohort and falls back to the global drain otherwise. The shape-agnostic + // rendezvous makes one queue sufficient. PTO2ReadyQueue early_sync_start_queue; static inline void ring_one_doorbell(uint64_t reg_addr, uint32_t token) { @@ -809,8 +810,8 @@ struct PTO2SchedulerState { } // 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. + // the NONE->RINGING launch latch and invokes this exactly once after local or + // global 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); @@ -879,12 +880,17 @@ struct PTO2SchedulerState { // 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) { + // running_slot_count is the publication seed: every staged_core_mask OR + // happens-before its final store. Read the seed first, then the mask, so + // observing the final count cannot be paired with a partially published + // mask when producer release races staging. + int32_t running_cores = slot_state.payload->running_slot_count.load(std::memory_order_seq_cst); 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 (running_cores != 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; @@ -896,7 +902,7 @@ struct PTO2SchedulerState { return true; } - inline bool retry_sync_start_rendezvous_after_drain(PTO2TaskSlotState &slot_state) { + inline bool retry_sync_start_rendezvous_after_staging(PTO2TaskSlotState &slot_state) { if (!maybe_rendezvous_ring(slot_state)) return false; propagate_dispatch_fanin(slot_state); return true; @@ -921,8 +927,8 @@ struct PTO2SchedulerState { } uint64_t task_id = static_cast(consumer.task->task_id.raw); - // A sync-start cohort uses one shape-agnostic queue so only the - // all-or-nothing drain path can claim and stage it. + // A sync-start cohort uses one shape-agnostic queue so one owner can + // choose an all-or-nothing local stage or the global-drain fallback. if (consumer.task_attrs.requires_sync_start()) { early_sync_start_queue.push_tagged(&consumer, task_id); } else { @@ -931,14 +937,14 @@ struct PTO2SchedulerState { } // 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 + // FLAGGED producer `p` publishes blocks (normal dispatch, early-dispatch release, or + // sync_start staging), 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 + // for its owner 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 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 1eab9dde5f..d9424506ad 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 @@ -527,42 +527,35 @@ bool SchedulerContext::enter_drain_mode(PTO2TaskSlotState *slot_state, int32_t b 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) { - // 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(); - } - } + total += core_trackers_[t].count_available_blocks(shape, core_mask, include_pending); } return total; } -// 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 +// Stage one thread's share of a sync_start cohort: CAS-claim block indices and +// publish them onto THIS thread's own cores. The global drain calls this +// concurrently on peers; the local fast path calls it only after proving this +// tracker can hold the entire cohort. Each thread touches only its own tracker +// and doorbell-table entries; the CAS on next_block_idx and 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) { +// Gated (early) sync_start staging, whether local or through the global drain, pre-stages +// every block behind its doorbell (prepare_block_for_dispatch is force-gated for the claimed +// range) and defers 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. +SchedulerContext::SyncStartStageResult SchedulerContext::stage_sync_start_cores( + PTO2TaskSlotState *slot_state, int32_t block_num, int32_t thread_idx, bool gated, + [[maybe_unused]] bool record_drain_phases +) { CoreTracker &tracker = core_trackers_[thread_idx]; PTO2ResourceShape shape = slot_state->active_mask.to_shape(); uint8_t core_mask = slot_state->active_mask.core_mask(); bool mix_split = gated && shape == PTO2ResourceShape::MIX; - int32_t running_staged = 0; + SyncStartStageResult result; // 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 @@ -575,8 +568,9 @@ SchedulerContext::drain_stage_cores(PTO2TaskSlotState *slot_state, int32_t block int32_t start = 0; int32_t claim = slot_state->claim_block_range(block_num, avail, start); if (claim == 0) return; + result.staged_blocks += claim; #if SIMPLER_DFX - bool sub_prof = l2_swimlane_level_ >= L2SwimlaneLevel::SCHED_PHASES; + bool sub_prof = record_drain_phases && l2_swimlane_level_ >= L2SwimlaneLevel::SCHED_PHASES; uint64_t prep_t0 = sub_prof ? get_sys_cnt_aicpu() : 0; #endif PublishHandle handles[CoreTracker::MAX_CLUSTERS * 3]; @@ -590,7 +584,7 @@ SchedulerContext::drain_stage_cores(PTO2TaskSlotState *slot_state, int32_t block 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); + result.running_cores += 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 @@ -647,7 +641,7 @@ SchedulerContext::drain_stage_cores(PTO2TaskSlotState *slot_state, int32_t block 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 (gated && shape != PTO2ResourceShape::MIX && !to_pending) result.running_cores += handle_count; } }; @@ -663,7 +657,7 @@ SchedulerContext::drain_stage_cores(PTO2TaskSlotState *slot_state, int32_t block stage(tracker.get_pending_core_offset_states(shape), /*to_pending=*/true); } } - return running_staged; + return result; } // Called by each scheduler thread when drain_state_.sync_start_pending != 0. @@ -764,13 +758,14 @@ void SchedulerContext::handle_drain_mode(int32_t thread_idx, [[maybe_unused]] ui #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); + SyncStartStageResult staged = + stage_sync_start_cores(slot_state, block_num, thread_idx, gated, /*record_drain_phases=*/true); #if SIMPLER_DFX - // out param carries the PURE drain_stage_cores wall (build_payload + MMIO publish of + // out param carries the PURE stage_sync_start_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_running_staged.fetch_add(staged.running_cores, std::memory_order_acq_rel); drain_state_.drain_stage_done_mask.fetch_or(1u << thread_idx, std::memory_order_release); if (!coordinator) { @@ -815,7 +810,7 @@ void SchedulerContext::handle_drain_mode(int32_t thread_idx, [[maybe_unused]] ui // 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); + sched_->retry_sync_start_rendezvous_after_staging(*slot_state); } else { sched_->propagate_dispatch_fanin(*slot_state); } 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 cc1f27ae6e..9a25386ccf 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 @@ -450,11 +450,18 @@ class SchedulerContext { bool enter_drain_mode(PTO2TaskSlotState *slot_state, int32_t block_num); 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 + struct SyncStartStageResult { + int32_t staged_blocks{0}; + int32_t running_cores{0}; + }; + // CAS-claim sync_start block indices and stage them onto THIS thread's own + // cores. The global drain invokes this in parallel on every scheduler; + // the local fast path invokes it once after proving this tracker has enough + // capacity. record_drain_phases keeps local work attributed to EarlyDispatch. + SyncStartStageResult stage_sync_start_cores( + PTO2TaskSlotState *slot_state, int32_t block_num, int32_t thread_idx, bool gated, bool record_drain_phases + ); + // out_stage_wall_cycles (profiling only): cycles this thread spent in stage_sync_start_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); 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 d0849db3b2..659d699fc7 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 @@ -260,7 +260,7 @@ int SchedulerContext::prepare_block_for_dispatch( 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 + // (promoted on completion). Gated sync_start staging 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) { @@ -786,9 +786,9 @@ SchedulerContext::early_dispatch_shape(int32_t thread_idx, PTO2ResourceShape sha // 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. +// before PENDING, cross-thread idle gating: MIX-IDLE ▶ c/v-IDLE ▶ MIX-PEND ▶ c/v-PEND). +// Each handles its sync_start cohort FIRST as Tier 0: an exact local fit stages on +// one owner, while a capacity-short cohort falls back to the global drain. // 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 @@ -808,16 +808,15 @@ int32_t SchedulerContext::try_early_dispatch( if (sched_->ready_sync_queues[s].size() > 0 || sched_->ready_queues[s].size() > 0) return 0; } + int32_t total_staged = 0; + // ===== 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. + // early_dispatch_shape's per-thread partial range-claim: a partial cohort would strand + // gated blocks nobody can ring. After claiming the per-task owner, first use the local + // tracker when it can hold the entire cohort; only the capacity-short case arms the + // stop-the-world drain. Both paths force-gate every block even if producer release races + // STAGING -> DISPATCHED. A non-STAGING pop was already released and is dropped. uint64_t sync_task_id_snapshot = 0; if (PTO2TaskSlotState *c = sched_->early_sync_start_queue.pop_tagged(&sync_task_id_snapshot)) { bool current_sync_task = @@ -825,6 +824,25 @@ int32_t SchedulerContext::try_early_dispatch( if (current_sync_task && 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 (drain_state_.sync_start_pending.load(std::memory_order_acquire) == 0 && + tracker.count_available_blocks( + c->active_mask.to_shape(), c->active_mask.core_mask(), /*include_pending=*/true + ) >= c->logical_block_num) { + // From this point onward the operation is all-or-nothing. Only this + // scheduler mutates its tracker, and global drain coordinators must + // wait for this scheduler's generation-tagged ack before inspecting it. + PTO2SchedulerState::mark_early_sync_drain_armed(*c->payload); + always_assert(c->next_block_idx.load(std::memory_order_seq_cst) == 0); + SyncStartStageResult staged = stage_sync_start_cores( + c, c->logical_block_num, thread_idx, /*gated=*/true, /*record_drain_phases=*/false + ); + always_assert(staged.staged_blocks == c->logical_block_num); + c->payload->running_slot_count.store( + static_cast(staged.running_cores), std::memory_order_seq_cst + ); + sched_->retry_sync_start_rendezvous_after_staging(*c); + PTO2SchedulerState::finish_early_sync_drain(*c->payload); + total_staged += staged.staged_blocks; } else if (enter_drain_mode(c, c->logical_block_num)) { PTO2SchedulerState::mark_early_sync_drain_armed(*c->payload); } else { @@ -836,7 +854,6 @@ int32_t SchedulerContext::try_early_dispatch( // 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, [&](PTO2ResourceShape shape, CoreTracker::DispatchPhase phase) { @@ -1145,7 +1162,7 @@ int32_t SchedulerContext::resolve_and_dispatch(Runtime *runtime, int32_t thread_ 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 + // stage_sync_start_cores). The many no-op entries — ack + availability-insufficient // reset or follower 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) { 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 a4ee4493ca..83e1c09338 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 @@ -443,6 +443,23 @@ class alignas(64) CoreTracker { return get_mix_running_cluster_offset_states(core_mask).count(); } + // Number of whole logical blocks this tracker can accept for sync_start + // staging. AIC/AIV count cores; MIX counts clusters because one logical MIX + // block may occupy multiple cores in the same cluster. Gated early staging + // includes pending slots, while ready staging is restricted to idle slots. + int32_t count_available_blocks(PTO2ResourceShape shape, uint8_t core_mask, bool include_pending) const { + if (shape == PTO2ResourceShape::MIX) { + return include_pending ? count_mix_split_clusters(core_mask) : count_mix_running_clusters(core_mask); + } + if (shape == PTO2ResourceShape::DUMMY) return 0; + + int32_t available = get_idle_core_offset_states(shape).count(); + if (include_pending) { + available += get_pending_core_offset_states(shape).count(); + } + return available; + } + BitStates get_pending_core_offset_states(PTO2ResourceShape shape) const { if (shape == PTO2ResourceShape::MIX) { // Shape-level query kept conservative for legacy callers/tests. diff --git a/src/a5/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md b/src/a5/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md index 1b0e37a45b..a6f486e989 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md +++ b/src/a5/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md @@ -651,22 +651,39 @@ dispatch runs only once normal `ready_queues[]` are empty **and** the local every *direct* producer is flagged `allow_early_resolve` (slot-state hint from Arg, or unconditional true for hidden alloc creators). There is no auto-chain inheritance. -#### 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 ownership** — a CAS on `sync_start_pending` makes drains mutually exclusive. -2. **All-or-nothing** — scheduler thread 0 coordinates the generation-tagged ack tree and - checks global available capacity ≥ `block_num` before staging; if short it advances the - attempt and retries after completions free cores. -3. **Parallel stage** — after thread 0 broadcasts the completed root token, each scheduler - thread CAS-claims a block range and stages its own cores with a non-zero `src_payload` - gate (`drain_stage_cores`): idle → running, busy → pending (`pending_gated` when still - waiting for the doorbell). -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. +#### sync_start local stage, drain fallback, and rendezvous + +A sync_start cohort of `block_num` blocks must occupy all its cores before any of them run. +An early Tier-0 pop first claims the cohort's task-local staging owner, then chooses one of +two all-or-nothing cases: + +1. **Case A — single-owner local stage** — when the popping scheduler's own + `CoreTracker` has at least `block_num` compatible idle + pending slots and no global + drain is already published, it force-gates and stages the complete cohort locally. + AIC/AIV capacity is the number of idle plus pending-capable cores; MIX capacity is the + number of active-mask-compatible split-placement clusters. This path never writes + `sync_start_pending`, `drain_attempt`, or the drain ack tokens. +2. **Case B — global drain fallback** — otherwise `enter_drain_mode` CAS-claims + `sync_start_pending`. Scheduler thread 0 coordinates the generation-tagged ack tree, + checks global capacity, and retries after completions when capacity is short. Once + released, every scheduler thread CAS-claims block ranges and stages only its own cores. +3. **Shared staging** — `stage_sync_start_cores` publishes every claimed block with a + non-zero `src_payload` gate: idle cores → running slots, busy cores → pending slots + (`pending_gated` while waiting for the doorbell). It returns both logical blocks staged + and running-slot cores, because one MIX block may use multiple cores. +4. **Rendezvous launch** — after the complete `staged_core_mask` is published, + `running_slot_count` is seeded with the staged running-slot cores. The rendezvous reads + that seed before the full mask; when it equals `popcount(staged_core_mask)` **and** the + producer has released, `maybe_rendezvous_ring` rings every gated core's doorbell + together — the cohort starts as one. + +If a global drain is published concurrently after Case A's zero check, its ack barrier +waits until the local scheduler returns to the loop; the coordinator then observes the +locally occupied tracker slots before making its global capacity decision. + +At scheduler-phase profiling level, local speculative staging is emitted as +`EarlyDispatch` rather than being folded into the normal `Dispatch` bar. This lets the +single-owner regression prove that the cohort was staged by Case A. Doorbell ownership is exclusive (`claim_all_staged_doorbell_bits` / `claim_late_staged_doorbell_bits`); launch is latched via `early_dispatch_launch_state`. diff --git a/src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.h b/src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.h index 091e40753e..5c0fb352b2 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.h +++ b/src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.h @@ -689,6 +689,9 @@ struct PTO2SchedulerState { }; EarlyDispatchDoorbell early_dispatch_doorbell_table[PTO2_EARLY_DISPATCH_CORE_MASK_WORDS * 64]{}; PTO2ReadyQueue early_dispatch_queues[PTO2_NUM_RESOURCE_SHAPES]; + // Shape-agnostic sync_start candidates require an all-or-nothing owner + // decision. The owner stages locally when one tracker can hold the complete + // cohort and falls back to the generation-tagged global drain otherwise. PTO2ReadyQueue early_sync_start_queue; static inline void ring_one_doorbell(uint64_t reg_addr, uint32_t token) { @@ -798,12 +801,17 @@ struct PTO2SchedulerState { } inline bool maybe_rendezvous_ring(PTO2TaskSlotState &slot_state) { + // Staging publishes the complete mask before seeding running_slot_count. + // Read the seed first: observing the final seq_cst seed then orders every + // mask read after all of the stager's mask updates. Reading the mask first + // could pair a partial mask with the later final seed and ring too early. + int32_t running_cores = slot_state.payload->running_slot_count.load(std::memory_order_seq_cst); 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 (running_cores != 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; @@ -815,7 +823,7 @@ struct PTO2SchedulerState { return true; } - inline bool retry_sync_start_rendezvous_after_drain(PTO2TaskSlotState &slot_state) { + inline bool retry_sync_start_rendezvous_after_staging(PTO2TaskSlotState &slot_state) { if (!maybe_rendezvous_ring(slot_state)) return false; propagate_dispatch_fanin(slot_state); return true; diff --git a/src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_completion.cpp b/src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_completion.cpp index dbc88af65d..f4ac4c6b2a 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_completion.cpp +++ b/src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_completion.cpp @@ -437,26 +437,25 @@ bool SchedulerContext::enter_drain_mode(PTO2TaskSlotState *slot_state, int32_t b 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 += 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(); - } - } + total += core_trackers_[t].count_available_blocks(shape, core_mask, include_pending); } return total; } -int32_t -SchedulerContext::drain_stage_cores(PTO2TaskSlotState *slot_state, int32_t block_num, int32_t thread_idx, bool gated) { +// Stage one scheduler thread's share of a sync_start cohort. The global drain +// calls this concurrently on every tracker; the Case-A fast path calls it once +// after proving this tracker can hold the whole cohort. staged_blocks is logical +// cohort progress, while running_cores seeds the core-count rendezvous (MIX may +// contribute multiple running cores per logical block). +SchedulerContext::SyncStartStageResult SchedulerContext::stage_sync_start_cores( + PTO2TaskSlotState *slot_state, int32_t block_num, int32_t thread_idx, bool gated, + [[maybe_unused]] bool record_drain_phases +) { CoreTracker &tracker = core_trackers_[thread_idx]; PTO2ResourceShape shape = slot_state->active_mask.to_shape(); uint8_t core_mask = slot_state->active_mask.core_mask(); bool mix_split = gated && shape == PTO2ResourceShape::MIX; - int32_t running_staged = 0; + SyncStartStageResult result; auto stage = [&](CoreTracker::BitStates valid, bool to_pending) { while (valid.has_value()) { @@ -464,6 +463,7 @@ SchedulerContext::drain_stage_cores(PTO2TaskSlotState *slot_state, int32_t block int32_t start = 0; int32_t claim = slot_state->claim_block_range(block_num, avail, start); if (claim == 0) return; + result.staged_blocks += claim; PublishHandle handles[CoreTracker::MAX_CLUSTERS * 3]; int handle_count = 0; int32_t claimed[CoreTracker::MAX_CLUSTERS * 3]; @@ -472,7 +472,7 @@ SchedulerContext::drain_stage_cores(PTO2TaskSlotState *slot_state, int32_t block for (int32_t b = 0; b < claim; b++) { auto core_offset = claimed[b]; if (shape == PTO2ResourceShape::MIX) { - running_staged += tracker.mix_cluster_idle_core_count(core_offset, core_mask); + result.running_cores += 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 @@ -503,7 +503,7 @@ SchedulerContext::drain_stage_cores(PTO2TaskSlotState *slot_state, int32_t block } } sched_->record_published_blocks(*slot_state, claim); - if (gated && shape != PTO2ResourceShape::MIX && !to_pending) running_staged += handle_count; + if (gated && shape != PTO2ResourceShape::MIX && !to_pending) result.running_cores += handle_count; } }; @@ -517,7 +517,7 @@ SchedulerContext::drain_stage_cores(PTO2TaskSlotState *slot_state, int32_t block stage(tracker.get_pending_core_offset_states(shape), /*to_pending=*/true); } } - return running_staged; + return result; } void SchedulerContext::handle_drain_mode(int32_t thread_idx, [[maybe_unused]] uint64_t *out_stage_wall_cycles) { @@ -580,11 +580,12 @@ void SchedulerContext::handle_drain_mode(int32_t thread_idx, [[maybe_unused]] ui #if SIMPLER_DFX if (drain_prof) drain_acked_ts = get_sys_cnt_aicpu(); #endif - int32_t my_running = drain_stage_cores(slot_state, block_num, thread_idx, gated); + SyncStartStageResult staged = + stage_sync_start_cores(slot_state, block_num, thread_idx, gated, /*record_drain_phases=*/true); #if SIMPLER_DFX 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_running_staged.fetch_add(staged.running_cores, std::memory_order_acq_rel); drain_state_.drain_stage_done_mask.fetch_or(1u << thread_idx, std::memory_order_release); if (!coordinator) { @@ -613,7 +614,7 @@ void SchedulerContext::handle_drain_mode(int32_t thread_idx, [[maybe_unused]] ui drain_state_.sync_start_pending.store(0, std::memory_order_release); if (gated) { - sched_->retry_sync_start_rendezvous_after_drain(*slot_state); + sched_->retry_sync_start_rendezvous_after_staging(*slot_state); } else { sched_->propagate_dispatch_fanin(*slot_state); } diff --git a/src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_context.h b/src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_context.h index 2fed49b305..243b77755e 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_context.h +++ b/src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_context.h @@ -376,7 +376,13 @@ class SchedulerContext { bool enter_drain_mode(PTO2TaskSlotState *slot_state, int32_t block_num); int32_t count_global_available(PTO2ResourceShape shape, uint8_t core_mask, bool include_pending = false); - int32_t drain_stage_cores(PTO2TaskSlotState *slot_state, int32_t block_num, int32_t thread_idx, bool gated); + struct SyncStartStageResult { + int32_t staged_blocks{0}; + int32_t running_cores{0}; + }; + SyncStartStageResult stage_sync_start_cores( + PTO2TaskSlotState *slot_state, int32_t block_num, int32_t thread_idx, bool gated, bool record_drain_phases + ); void handle_drain_mode(int32_t thread_idx, uint64_t *out_stage_wall_cycles = nullptr); // ========================================================================= diff --git a/src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_dispatch.cpp b/src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_dispatch.cpp index 9c1d882b78..b28b53118e 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_dispatch.cpp +++ b/src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_dispatch.cpp @@ -595,7 +595,13 @@ int32_t SchedulerContext::try_early_dispatch( if (sched_->ready_queues[s].size() > 0) return 0; } - // Tier 0: sync_start early cohorts (shape-agnostic, all-or-nothing drain). + int32_t total_staged = 0; + + // Tier 0: sync_start early cohorts (shape-agnostic, all-or-nothing). + // Case A stages the entire cohort on this scheduler when its own idle + + // pending capacity is sufficient and no global drain is already published. + // Case B preserves the global drain fallback for cohorts that need cores + // owned by multiple schedulers. Neither case permits a partial cohort. uint64_t sync_task_id_snapshot = 0; if (PTO2TaskSlotState *c = sched_->early_sync_start_queue.pop_tagged(&sync_task_id_snapshot)) { bool current_sync_task = @@ -603,6 +609,27 @@ int32_t SchedulerContext::try_early_dispatch( if (current_sync_task && 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 (drain_state_.sync_start_pending.load(std::memory_order_acquire) == 0 && + tracker.count_available_blocks( + c->active_mask.to_shape(), c->active_mask.core_mask(), /*include_pending=*/true + ) >= c->logical_block_num) { + // OWNER protects the producer-ready handoff. ARMED makes the + // ownership state match the global path before any gated + // payload becomes visible. Once staging starts, the local + // capacity invariant guarantees completion; never fall back + // after publishing a partial cohort. + PTO2SchedulerState::mark_early_sync_drain_armed(*c->payload); + always_assert(c->next_block_idx.load(std::memory_order_seq_cst) == 0); + SyncStartStageResult staged = stage_sync_start_cores( + c, c->logical_block_num, thread_idx, /*gated=*/true, /*record_drain_phases=*/false + ); + always_assert(staged.staged_blocks == c->logical_block_num); + c->payload->running_slot_count.store( + static_cast(staged.running_cores), std::memory_order_seq_cst + ); + sched_->retry_sync_start_rendezvous_after_staging(*c); + PTO2SchedulerState::finish_early_sync_drain(*c->payload); + total_staged += staged.staged_blocks; } else if (enter_drain_mode(c, c->logical_block_num)) { PTO2SchedulerState::mark_early_sync_drain_armed(*c->payload); } else { @@ -618,7 +645,6 @@ int32_t SchedulerContext::try_early_dispatch( }; const PTO2ResourceShape *aic_aiv = kAicAivOrder[thread_idx & 1]; - int32_t total_staged = 0; total_staged += early_dispatch_shape(thread_idx, PTO2ResourceShape::MIX, Phase::IDLE); bool skip_aic_aiv = has_residual_early_mix(); if (!skip_aic_aiv) { @@ -1073,41 +1099,62 @@ int32_t SchedulerContext::resolve_and_dispatch(Runtime *runtime, int32_t thread_ // Phase 4: MIX-strict-priority dispatch with phase-split and // cross-thread idle gating. See dispatch_ready_tasks for the policy. // pmu_active is cached at function scope above (loop-invariant). +#if SIMPLER_DFX + uint64_t dispatch_t0 = (l2_swimlane_level_ >= L2SwimlaneLevel::SCHED_PHASES) ? get_sys_cnt_aicpu() : 0; +#endif dispatch_ready_tasks(runtime, thread_idx, tracker, pmu_active, made_progress, try_pushed); +#if SIMPLER_DFX + // Close normal Dispatch before speculative staging so the two sources + // remain distinguishable in scheduler-phase traces. + if (dispatch_t0 != 0 && l2_swimlane.phase_dispatch_count > 0) { + uint64_t dispatch_t1 = get_sys_cnt_aicpu(); + uint64_t pop_hit_delta = l2_swimlane.pop_hit - l2_swimlane.pop_hit_at_last_emit; + uint64_t pop_miss_delta = l2_swimlane.pop_miss - l2_swimlane.pop_miss_at_last_emit; + debug_assert(pop_hit_delta < (1ULL << 32)); + debug_assert(pop_miss_delta < (1ULL << 32)); + int16_t phase_end_shared[L2SWIMLANE_NUM_QUEUE_SHAPES]; + capture_phase_end(phase_end_shared); + l2_swimlane_aicpu_record_sched_phase( + thread_idx, L2SwimlaneSchedPhaseKind::Dispatch, _t0_phase, dispatch_t1, l2_swimlane.sched_loop_count, + l2_swimlane.phase_dispatch_count, static_cast(pop_hit_delta), + static_cast(pop_miss_delta), phase_start_shared, phase_end_shared + ); + for (int s = 0; s < L2SWIMLANE_NUM_QUEUE_SHAPES; s++) { + phase_start_shared[s] = phase_end_shared[s]; + } + _t0_phase = dispatch_t1; + l2_swimlane.phase_dispatch_count = 0; + l2_swimlane.pop_hit_at_last_emit = l2_swimlane.pop_hit; + l2_swimlane.pop_miss_at_last_emit = l2_swimlane.pop_miss; + } +#endif // Phase 4b: early-dispatch onto spare cores after normal dispatch. - (void)try_early_dispatch(thread_idx, tracker, pmu_active, made_progress, try_pushed); +#if SIMPLER_DFX + bool early_dispatch_record = l2_swimlane_level_ >= L2SwimlaneLevel::SCHED_PHASES; + uint64_t early_dispatch_t0 = early_dispatch_record ? get_sys_cnt_aicpu() : 0; +#endif + [[maybe_unused]] int32_t staged_count = + try_early_dispatch(thread_idx, tracker, pmu_active, made_progress, try_pushed); +#if SIMPLER_DFX + if (early_dispatch_record && staged_count > 0) { + uint64_t early_dispatch_t1 = get_sys_cnt_aicpu(); + l2_swimlane_aicpu_record_sched_phase( + thread_idx, L2SwimlaneSchedPhaseKind::EarlyDispatch, early_dispatch_t0, early_dispatch_t1, + l2_swimlane.sched_loop_count, static_cast(staged_count) + ); + // prepare_block_for_dispatch accounts every publish in the shared + // dispatch counter; these blocks belong to EarlyDispatch instead. + l2_swimlane.phase_dispatch_count = 0; + _t0_phase = early_dispatch_t1; + } +#endif #if SIMPLER_DFX if (!try_pushed) { CYCLE_COUNT_LAP(l2_swimlane.sched_idle_cycle); } else { CYCLE_COUNT_LAP(l2_swimlane.sched_dispatch_cycle); - if (l2_swimlane_level_ >= L2SwimlaneLevel::SCHED_PHASES && l2_swimlane.phase_dispatch_count > 0) { - // Final-drain at loop end emits the trailing-idle tail so - // sum-of-deltas == run-cumulative. - uint64_t pop_hit_delta = l2_swimlane.pop_hit - l2_swimlane.pop_hit_at_last_emit; - uint64_t pop_miss_delta = l2_swimlane.pop_miss - l2_swimlane.pop_miss_at_last_emit; - // L2SwimlaneAicpuSchedPhaseRecord's dispatch counters are uint32 — an overflow means - // an emit was missed for ~4 billion pops, which is well outside any - // realistic dispatch cadence and silently truncates without this guard. - debug_assert(pop_hit_delta < (1ULL << 32)); - debug_assert(pop_miss_delta < (1ULL << 32)); - int16_t phase_end_shared[L2SWIMLANE_NUM_QUEUE_SHAPES]; - capture_phase_end(phase_end_shared); - l2_swimlane_aicpu_record_sched_phase( - thread_idx, L2SwimlaneSchedPhaseKind::Dispatch, _t0_phase, _t1, l2_swimlane.sched_loop_count, - l2_swimlane.phase_dispatch_count, static_cast(pop_hit_delta), - static_cast(pop_miss_delta), phase_start_shared, phase_end_shared - ); - for (int s = 0; s < L2SWIMLANE_NUM_QUEUE_SHAPES; s++) { - phase_start_shared[s] = phase_end_shared[s]; - } - _t0_phase = _t1; - l2_swimlane.phase_dispatch_count = 0; - l2_swimlane.pop_hit_at_last_emit = l2_swimlane.pop_hit; - l2_swimlane.pop_miss_at_last_emit = l2_swimlane.pop_miss; - } } #endif @@ -1190,8 +1237,8 @@ int32_t SchedulerContext::resolve_and_dispatch(Runtime *runtime, int32_t thread_ SPIN_WAIT_HINT(); #if SIMPLER_DFX CYCLE_COUNT_LAP(l2_swimlane.sched_idle_cycle); - // a2a3 design has Complete + Dispatch sched phases only; idle gaps - // are reconstructed at post-process time from sched record spacing. + // Idle gaps are reconstructed at post-process time from scheduler + // phase-record spacing. (void)_t0_phase; #endif } diff --git a/src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_types.h b/src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_types.h index d634b32a4c..a82b886861 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_types.h +++ b/src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_types.h @@ -443,6 +443,21 @@ class alignas(64) CoreTracker { return (~core_states_) & aiv_mask_ & ~pending_occupied_; } + // Number of logical sync_start blocks this tracker can place. AIC/AIV use + // one core per block; MIX uses one active-mask-compatible cluster per block. + // Gated early staging may use both idle running slots and free pending slots, + // while a ready cohort is restricted to idle running slots. + int32_t count_available_blocks(PTO2ResourceShape shape, uint8_t core_mask, bool include_pending) const { + if (shape == PTO2ResourceShape::MIX) { + return include_pending ? count_mix_split_clusters(core_mask) : count_mix_running_clusters(core_mask); + } + if (shape != PTO2ResourceShape::AIC && shape != PTO2ResourceShape::AIV) return 0; + + int32_t available = get_idle_core_offset_states(shape).count(); + if (include_pending) available += get_pending_core_offset_states(shape).count(); + return available; + } + // --- Two-phase dispatch unified query --- enum class DispatchPhase : uint8_t { IDLE, PENDING }; diff --git a/src/common/platform/include/common/l2_swimlane_profiling.h b/src/common/platform/include/common/l2_swimlane_profiling.h index 53344d93a4..e1c85536f3 100644 --- a/src/common/platform/include/common/l2_swimlane_profiling.h +++ b/src/common/platform/include/common/l2_swimlane_profiling.h @@ -527,9 +527,9 @@ enum class L2SwimlaneSchedPhaseKind : uint32_t { // 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 + DrainPrepare = 9, // inner: this thread's global sync_start staging prepare pass // (cluster scan + build_payload). tasks_processed = subtasks. - DrainPublish = 10, // inner: this thread's drain_stage_cores publish pass + DrainPublish = 10, // inner: this thread's global sync_start staging publish pass // (MMIO write_reg per subtask). tasks_processed = subtasks. // Outer (sched lane): async-wait completion polling, split out of Complete // so async-engine (SDMA/RoCE/URMA/CCU) wait time is attributed to its own diff --git a/tests/st/a2a3/tensormap_and_ringbuffer/dfx/l2_swimlane/kernels/aiv/kernel_spmd_spin.cpp b/tests/st/a2a3/tensormap_and_ringbuffer/dfx/l2_swimlane/kernels/aiv/kernel_spmd_spin.cpp new file mode 100644 index 0000000000..a34b6b63cf --- /dev/null +++ b/tests/st/a2a3/tensormap_and_ringbuffer/dfx/l2_swimlane/kernels/aiv/kernel_spmd_spin.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. + * ----------------------------------------------------------------------------------------------------------- + */ + +/** + * An AIV blocker that occupies every running slot for one measured second. + * The wall-clock hold is independent of core execution speed and leaves a + * large margin for the orchestration's scheduler-loop fences. + * + * Each block owns one status cache line: + * 0 = released, 1 = running. + * + * Args: + * args[0] = inout Tensor* + * args[1] = scalar status base cache line + * args[2] = scalar status count + */ + +#include +#include + +#ifdef PTO_CPUSTUB_HPP +#include +#include +#endif + +#include "common/platform_config.h" +#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 float BLOCKER_STARTED = 1.0F; +static constexpr uint64_t BLOCKER_HOLD_TICKS = PLATFORM_PROF_SYS_CNT_FREQ; + +#ifdef PTO_CPUSTUB_HPP +#ifndef dcci +#define dcci(...) \ + do { \ + } while (0) +#endif +#define SPIN_WAIT_HINT() std::this_thread::yield() + +static uint64_t blocker_now_ticks() { + return static_cast( + std::chrono::duration_cast(std::chrono::steady_clock::now().time_since_epoch()) + .count() + ); +} +#else +#define SPIN_WAIT_HINT() ((void)0) + +static __aicore__ uint64_t blocker_now_ticks() { return get_sys_cnt(); } +#endif +#ifndef SINGLE_CACHE_LINE +#define SINGLE_CACHE_LINE 0 +#endif +#ifndef CACHELINE_OUT +#define CACHELINE_OUT 0 +#endif + +static __aicore__ void publish_value(__gm__ float *value, float state) { + *value = state; + dcci(value, SINGLE_CACHE_LINE, CACHELINE_OUT); +} + +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 status_base_cl = static_cast(args[1]); + int32_t status_count = static_cast(args[2]); + int32_t block_idx = get_block_idx(args); + if (block_idx < 0 || block_idx >= status_count) { + return; + } + + __gm__ float *status = &out[(status_base_cl + block_idx) * FLOATS_PER_CACHE_LINE]; + uint64_t wait_start = blocker_now_ticks(); + + publish_value(status, BLOCKER_STARTED); + + while (blocker_now_ticks() - wait_start < BLOCKER_HOLD_TICKS) { + SPIN_WAIT_HINT(); + } + + publish_value(status, 0.0F); +} diff --git a/tests/st/a2a3/tensormap_and_ringbuffer/dfx/l2_swimlane/kernels/orchestration/sync_start_early_local_owner_orch.cpp b/tests/st/a2a3/tensormap_and_ringbuffer/dfx/l2_swimlane/kernels/orchestration/sync_start_early_local_owner_orch.cpp new file mode 100644 index 0000000000..4d6120cc23 --- /dev/null +++ b/tests/st/a2a3/tensormap_and_ringbuffer/dfx/l2_swimlane/kernels/orchestration/sync_start_early_local_owner_orch.cpp @@ -0,0 +1,150 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +/** + * Issue #1548: early sync-start cohort that fits one scheduler. + * + * Mode 0 uses a one-block flagged AIC producer. Its dependent AIV + * require_sync_start consumer fits one scheduler's local idle AIV capacity, + * including a four-block cohort when two peer schedulers are active. + * + * Mode 1 fills every AIV running slot with a blocker. Scheduler-loop fences + * first retire every blocker dispatch ACK, then give the speculative consumer + * one full loop to stage while the one-second blockers remain active. + * + * The slow producer and the consumer write disjoint cache lines: + * P: AIC block_num=1, base_cl=0, allow_early_resolve=true + * C: AIV block_num=consumer_blocks, base_cl=1, require_sync_start=true, dep=[P] + */ + +#include +#include + +#include "aicpu/cache_maintenance.h" +#include "pto_arg_with_deps.h" // NOLINT(build/include_subdir) +#include "pto_orchestration_api.h" // NOLINT(build/include_subdir) + +#define FUNC_SLOW_PRODUCER_AIC 0 +#define FUNC_SYNC_CONSUMER_AIV 1 +#define FUNC_SPIN_BLOCKER_AIV 2 + +extern "C" { + +__attribute__((visibility("default"))) PTO2OrchestrationConfig aicpu_orchestration_config(const L2TaskArgs &orch_args) { + (void)orch_args; // NOLINT(readability/casting) + return PTO2OrchestrationConfig{ + .expected_arg_count = 3, + }; +} + +static constexpr int16_t PRODUCER_BLOCKS = 1; +static constexpr int32_t FLOATS_PER_CACHE_LINE = 16; +static constexpr int32_t BLOCKER_STATUS_BASE_CL = 9; +static constexpr int32_t BLOCKER_STATUS_CAPACITY = 72; +// Matches the existing early-dispatch regression's delay. It is long enough +// for the scheduler to pop the consumer before the producer completes. +static constexpr int64_t PRODUCER_SPIN_ITERS = 10000000; + +static bool wait_for_blockers_started(const Tensor &out, int32_t blocker_count) { + volatile float *data = out.data_as() + out.start_offset; + volatile float *statuses = data + BLOCKER_STATUS_BASE_CL * FLOATS_PER_CACHE_LINE; + while (!rt_is_fatal()) { + cache_invalidate_range( + const_cast(statuses), static_cast(blocker_count) * FLOATS_PER_CACHE_LINE * sizeof(float) + ); + bool all_started = true; + for (int32_t block_idx = 0; block_idx < blocker_count; block_idx++) { + float status = statuses[block_idx * FLOATS_PER_CACHE_LINE]; + if (status != 1.0F) { + all_started = false; + break; + } + } + if (all_started) break; + SPIN_WAIT_HINT(); + } + return !rt_is_fatal(); +} + +static bool wait_for_scheduler_loop_fence() { + L0TaskArgs first_args; + PTO2TaskId first = rt_submit_dummy_task(first_args).task_id(); + + uint32_t fence_shape[1] = {1}; + TensorCreateInfo fence_info(fence_shape, 1, DataType::INT32); + L0TaskArgs second_args; + second_args.add_output(fence_info); + PTO2TaskId deps[1] = {first}; + second_args.set_dependencies(deps, 1); + TaskOutputTensors outputs = rt_submit_dummy_task(second_args); + + // D2 cannot be popped in the batch that completes D1. Observing D2's + // output therefore proves that one complete scheduler loop elapsed after + // this fence was submitted. + const Tensor &fence = outputs.get_ref(0); + uint32_t index[1] = {0}; + (void)get_tensor_data(fence, 1, index); + return !rt_is_fatal(); +} + +static PTO2TaskId submit_producer(const Tensor &out) { + L0TaskArgs args; + args.add_inout(out); + args.add_scalar(0); // base cache line + args.add_scalar(PRODUCER_SPIN_ITERS); + args.launch_spec.set_block_num(PRODUCER_BLOCKS); + args.set_allow_early_resolve(true); + return rt_submit_aic_task(FUNC_SLOW_PRODUCER_AIC, args).task_id(); +} + +static PTO2TaskId submit_aiv_blocker(const Tensor &out, int32_t blocker_count) { + L0TaskArgs args; + args.add_inout(out); + args.add_scalar(BLOCKER_STATUS_BASE_CL); + args.add_scalar(blocker_count); + args.launch_spec.set_block_num(static_cast(blocker_count)); + args.set_allow_early_resolve(true); + return rt_submit_aiv_task(FUNC_SPIN_BLOCKER_AIV, args).task_id(); +} + +static void submit_consumer(const Tensor &out, PTO2TaskId producer, int16_t consumer_blocks) { + L0TaskArgsWithDeps<1> args; + args.add_inout(out); + args.add_scalar(PRODUCER_BLOCKS); // base cache line + args.launch_spec.set_block_num(consumer_blocks); + args.launch_spec.set_require_sync_start(true); + args.add_dep(producer); + rt_submit_aiv_task(FUNC_SYNC_CONSUMER_AIV, args); +} + +__attribute__((visibility("default"))) void aicpu_orchestration_entry(const L2TaskArgs &orch_args) { + const Tensor &output = orch_args.tensor(0).ref(); + const bool use_pending = orch_args.scalar(0) != 0; + const int16_t consumer_blocks = static_cast(orch_args.scalar(1)); + const int32_t blocker_count = static_cast(rt_available_aiv_count()); + + if (use_pending && (blocker_count <= 0 || blocker_count > BLOCKER_STATUS_CAPACITY || consumer_blocks <= 0 || + consumer_blocks > blocker_count)) { + rt_report_fatal( + PTO2_ERROR_INVALID_ARGS, "invalid pending geometry: blockers=%d consumer=%d", blocker_count, consumer_blocks + ); + return; + } + + rt_scope_begin(PTO2ScopeMode::MANUAL); + PTO2TaskId producer = use_pending ? submit_aiv_blocker(output, blocker_count) : submit_producer(output); + if (use_pending && (!wait_for_blockers_started(output, blocker_count) || !wait_for_scheduler_loop_fence())) return; + submit_consumer(output, producer, consumer_blocks); + if (use_pending && !wait_for_scheduler_loop_fence()) return; + rt_scope_end(); +} + +} // extern "C" diff --git a/tests/st/a2a3/tensormap_and_ringbuffer/dfx/l2_swimlane/test_sync_start_early_local_owner.py b/tests/st/a2a3/tensormap_and_ringbuffer/dfx/l2_swimlane/test_sync_start_early_local_owner.py new file mode 100644 index 0000000000..3dc86ae5ed --- /dev/null +++ b/tests/st/a2a3/tensormap_and_ringbuffer/dfx/l2_swimlane/test_sync_start_early_local_owner.py @@ -0,0 +1,169 @@ +#!/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. +# ----------------------------------------------------------------------------------------------------------- +"""Regression for #1548: an early sync-start cohort fits one scheduler. + +The idle cases use a slow flagged AIC producer and exercise both one and three +scheduler threads. The dependent AIV sync-start consumer fits one owner's local +tracker and must not stop its peers. + +The pending case instead keeps every AIV core busy with a blocker. A +pre-submit scheduler-loop fence retires all blocker ACKs. A second fence after +consumer submission completes only after an early-dispatch opportunity while +the measured one-second hold keeps every running slot occupied. Phase capture +verifies all eight blocks staged locally. + +The phase assertion runs only with L2 swimlane level >= 3. The ordinary scene +test still checks workload correctness when diagnostics are disabled. +""" + +from __future__ import annotations + +from collections import Counter +from pathlib import Path + +import torch +from simpler.task_interface import ArgDirection as D + +from simpler_setup import Scalar, SceneTestCase, TaskArgsBuilder, Tensor, scene_test +from simpler_setup.tools.swimlane_converter import read_perf_data + +FLOATS_PER_CACHE_LINE = 16 +PRODUCER_BLOCKS = 1 +MAX_CONSUMER_BLOCKS = 8 +TOTAL_BLOCKS = PRODUCER_BLOCKS + MAX_CONSUMER_BLOCKS +BLOCKER_STATUS_CAPACITY = 72 +OUTPUT_CACHE_LINES = TOTAL_BLOCKS + BLOCKER_STATUS_CAPACITY + + +@scene_test(level=2, runtime="tensormap_and_ringbuffer") +class TestSyncStartEarlyLocalOwner(SceneTestCase): + RTOL = 0 + ATOL = 0 + + CALLABLE = { + "orchestration": { + "source": "kernels/orchestration/sync_start_early_local_owner_orch.cpp", + "function_name": "aicpu_orchestration_entry", + "signature": [D.INOUT], + }, + "incores": [ + { + "func_id": 0, + "name": "SLOW_FLAGGED_PRODUCER_AIC", + "source": "../../spmd_sync_start_early_dispatch/kernels/aiv/kernel_spmd_write_slow.cpp", + "core_type": "aic", + "signature": [D.INOUT], + }, + { + "func_id": 1, + "name": "SYNC_START_CONSUMER_AIV", + "source": "../../spmd_multiblock_aiv/kernels/aiv/kernel_spmd_write.cpp", + "core_type": "aiv", + "signature": [D.INOUT], + }, + { + "func_id": 2, + "name": "SPIN_BLOCKER_AIV", + "source": "kernels/aiv/kernel_spmd_spin.cpp", + "core_type": "aiv", + "signature": [D.INOUT], + }, + ], + } + + CASES = [ + { + "name": "single_scheduler_idle", + "platforms": ["a2a3sim", "a2a3"], + # One orchestrator plus one scheduler. That scheduler owns all AIV + # cores, so the eight-block consumer fits its local idle capacity. + "config": {"aicpu_thread_num": 2}, + "params": {"mode": 0, "consumer_blocks": 8}, + }, + { + "name": "three_schedulers_idle", + "platforms": ["a2a3sim", "a2a3"], + # One orchestrator plus three schedulers. Whichever scheduler pops + # the four-block cohort owns enough AIV cores to stage it locally, + # so no non-owner scheduler should participate in a global drain. + "config": {"aicpu_thread_num": 4}, + "params": {"mode": 0, "consumer_blocks": 4}, + }, + { + "name": "single_scheduler_pending_only", + "platforms": ["a2a3sim", "a2a3"], + # Every blocker reports started before an ACK-sweep fence. A second + # scheduler-loop fence completes during their measured one-second + # hold, proving the consumer used pending rather than idle slots. + "config": {"aicpu_thread_num": 2}, + "params": {"mode": 1, "consumer_blocks": 8}, + }, + ] + + def generate_args(self, params): + return TaskArgsBuilder( + Tensor("output", torch.zeros(OUTPUT_CACHE_LINES * FLOATS_PER_CACHE_LINE, dtype=torch.float32)), + Scalar("mode", int(params["mode"])), + Scalar("consumer_blocks", int(params["consumer_blocks"])), + ) + + def compute_golden(self, args, params): + for block_idx in range(PRODUCER_BLOCKS): + args.output[block_idx * FLOATS_PER_CACHE_LINE] = float(block_idx) + for block_idx in range(int(params["consumer_blocks"])): + args.output[(PRODUCER_BLOCKS + block_idx) * FLOATS_PER_CACHE_LINE] = float(block_idx) + + def _build_config(self, config_dict, *args, **kwargs): + config = super()._build_config(config_dict, *args, **kwargs) + self._trace_perf_level = int(kwargs.get("enable_l2_swimlane", args[0] if args else 0)) + output_prefix = kwargs.get("output_prefix", "") + self._trace_perf_path = Path(output_prefix) / "l2_swimlane_records.json" if output_prefix else None + return config + + def compare_outputs(self, test_args, golden_args, output_names, params): + super().compare_outputs(test_args, golden_args, output_names, params) + if getattr(self, "_trace_perf_level", 0) < 3: + return + + perf_path = self._trace_perf_path + assert perf_path is not None, "L2 swimlane enabled without an output prefix" + assert perf_path.exists(), f"l2_swimlane_records.json missing under {perf_path.parent}" + + perf = read_perf_data(perf_path) + assert int(perf.get("l2_swimlane_level", 0)) >= 3, f"scheduler phases missing from {perf_path}" + phase_records = [record for thread in perf.get("aicpu_scheduler_phases", []) for record in thread] + assert phase_records, f"scheduler phase capture is empty under {perf_path}" + phase_counts = Counter(record.get("phase") for record in phase_records) + expected_blocks = int(params["consumer_blocks"]) + drain_phase_names = {"drain", "drain_prepare", "drain_publish"} + drain_records = [record for record in phase_records if record.get("phase") in drain_phase_names] + drain_published = sum( + int(record.get("tasks_processed", 0)) for record in phase_records if record.get("phase") == "drain_publish" + ) + + assert not drain_records, ( + "#1548: the early sync-start consumer fits one scheduler's local AIV " + "running/pending capacity but entered global drain; " + f"drain_publish staged {drain_published}/{expected_blocks} blocks, " + f"mode={params['mode']}, phase_counts={dict(phase_counts)}, artifact={perf_path}" + ) + + early_staged = sum( + int(record.get("tasks_processed", 0)) for record in phase_records if record.get("phase") == "early_dispatch" + ) + assert early_staged == expected_blocks, ( + "the local-owner fast path was not observed: " + f"early_dispatch staged {early_staged}, expected {expected_blocks}; " + f"mode={params['mode']}, phase_counts={dict(phase_counts)}, artifact={perf_path}" + ) + + +if __name__ == "__main__": + SceneTestCase.run_module(__name__) 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 index 80ce9f481a..55b6944129 100644 --- 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 @@ -17,7 +17,7 @@ * (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, + * cores' gated pending slots (stage_sync_start_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 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 index 744ea2ffb3..f4e14fcc0f 100644 --- 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 @@ -11,7 +11,7 @@ (and spins), leaving every AIC core 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 +(stage_sync_start_cores to_pending=true, mix_cluster_idle_core_count=1/cluster + Case 3.3 promote for every pending AIV). A counting mismatch stalls the rendezvous -> gated cores never launch -> allocator deadlock. diff --git a/tests/st/a5/tensormap_and_ringbuffer/dfx/l2_swimlane/kernels/aiv/kernel_spmd_spin.cpp b/tests/st/a5/tensormap_and_ringbuffer/dfx/l2_swimlane/kernels/aiv/kernel_spmd_spin.cpp new file mode 100644 index 0000000000..a34b6b63cf --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/dfx/l2_swimlane/kernels/aiv/kernel_spmd_spin.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. + * ----------------------------------------------------------------------------------------------------------- + */ + +/** + * An AIV blocker that occupies every running slot for one measured second. + * The wall-clock hold is independent of core execution speed and leaves a + * large margin for the orchestration's scheduler-loop fences. + * + * Each block owns one status cache line: + * 0 = released, 1 = running. + * + * Args: + * args[0] = inout Tensor* + * args[1] = scalar status base cache line + * args[2] = scalar status count + */ + +#include +#include + +#ifdef PTO_CPUSTUB_HPP +#include +#include +#endif + +#include "common/platform_config.h" +#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 float BLOCKER_STARTED = 1.0F; +static constexpr uint64_t BLOCKER_HOLD_TICKS = PLATFORM_PROF_SYS_CNT_FREQ; + +#ifdef PTO_CPUSTUB_HPP +#ifndef dcci +#define dcci(...) \ + do { \ + } while (0) +#endif +#define SPIN_WAIT_HINT() std::this_thread::yield() + +static uint64_t blocker_now_ticks() { + return static_cast( + std::chrono::duration_cast(std::chrono::steady_clock::now().time_since_epoch()) + .count() + ); +} +#else +#define SPIN_WAIT_HINT() ((void)0) + +static __aicore__ uint64_t blocker_now_ticks() { return get_sys_cnt(); } +#endif +#ifndef SINGLE_CACHE_LINE +#define SINGLE_CACHE_LINE 0 +#endif +#ifndef CACHELINE_OUT +#define CACHELINE_OUT 0 +#endif + +static __aicore__ void publish_value(__gm__ float *value, float state) { + *value = state; + dcci(value, SINGLE_CACHE_LINE, CACHELINE_OUT); +} + +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 status_base_cl = static_cast(args[1]); + int32_t status_count = static_cast(args[2]); + int32_t block_idx = get_block_idx(args); + if (block_idx < 0 || block_idx >= status_count) { + return; + } + + __gm__ float *status = &out[(status_base_cl + block_idx) * FLOATS_PER_CACHE_LINE]; + uint64_t wait_start = blocker_now_ticks(); + + publish_value(status, BLOCKER_STARTED); + + while (blocker_now_ticks() - wait_start < BLOCKER_HOLD_TICKS) { + SPIN_WAIT_HINT(); + } + + publish_value(status, 0.0F); +} diff --git a/tests/st/a5/tensormap_and_ringbuffer/dfx/l2_swimlane/kernels/orchestration/sync_start_early_local_owner_orch.cpp b/tests/st/a5/tensormap_and_ringbuffer/dfx/l2_swimlane/kernels/orchestration/sync_start_early_local_owner_orch.cpp new file mode 100644 index 0000000000..39dd23a774 --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/dfx/l2_swimlane/kernels/orchestration/sync_start_early_local_owner_orch.cpp @@ -0,0 +1,142 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +/** + * Issue #1548: an early sync-start cohort that fits one a5 scheduler. + * + * Mode 0 uses a slow, flagged one-block AIC producer and covers both one and + * three scheduler threads, with the consumer cohort sized to fit one local + * owner. Mode 1 fills every AIV running slot with a blocker. Scheduler-loop + * fences first retire every blocker dispatch ACK, then give the speculative + * consumer one full loop to stage while the one-second blockers remain active. + */ + +#include +#include + +#include "aicpu/cache_maintenance.h" +#include "pto_arg_with_deps.h" // NOLINT(build/include_subdir) +#include "pto_orchestration_api.h" // NOLINT(build/include_subdir) + +#define FUNC_SLOW_PRODUCER_AIC 0 +#define FUNC_SYNC_CONSUMER_AIV 1 +#define FUNC_SPIN_BLOCKER_AIV 2 + +extern "C" { + +__attribute__((visibility("default"))) PTO2OrchestrationConfig aicpu_orchestration_config(const L2TaskArgs &orch_args) { + (void)orch_args; // NOLINT(readability/casting) + return PTO2OrchestrationConfig{ + .expected_arg_count = 3, + }; +} + +static constexpr int16_t PRODUCER_BLOCKS = 1; +static constexpr int32_t FLOATS_PER_CACHE_LINE = 16; +static constexpr int32_t BLOCKER_STATUS_BASE_CL = 9; +static constexpr int32_t BLOCKER_STATUS_CAPACITY = 72; +static constexpr int64_t PRODUCER_SPIN_ITERS = 10000000; + +static bool wait_for_blockers_started(const Tensor &out, int32_t blocker_count) { + volatile float *data = out.data_as() + out.start_offset; + volatile float *statuses = data + BLOCKER_STATUS_BASE_CL * FLOATS_PER_CACHE_LINE; + while (!rt_is_fatal()) { + cache_invalidate_range( + const_cast(statuses), static_cast(blocker_count) * FLOATS_PER_CACHE_LINE * sizeof(float) + ); + bool all_started = true; + for (int32_t block_idx = 0; block_idx < blocker_count; block_idx++) { + float status = statuses[block_idx * FLOATS_PER_CACHE_LINE]; + if (status != 1.0F) { + all_started = false; + break; + } + } + if (all_started) break; + SPIN_WAIT_HINT(); + } + return !rt_is_fatal(); +} + +static bool wait_for_scheduler_loop_fence() { + L0TaskArgs first_args; + PTO2TaskId first = rt_submit_dummy_task(first_args).task_id(); + + uint32_t fence_shape[1] = {1}; + TensorCreateInfo fence_info(fence_shape, 1, DataType::INT32); + L0TaskArgs second_args; + second_args.add_output(fence_info); + PTO2TaskId deps[1] = {first}; + second_args.set_dependencies(deps, 1); + TaskOutputTensors outputs = rt_submit_dummy_task(second_args); + + // D2 cannot be popped in the batch that completes D1. Observing D2's + // output therefore proves that one complete scheduler loop elapsed after + // this fence was submitted. + const Tensor &fence = outputs.get_ref(0); + uint32_t index[1] = {0}; + (void)get_tensor_data(fence, 1, index); + return !rt_is_fatal(); +} + +static PTO2TaskId submit_producer(const Tensor &out) { + L0TaskArgs args; + args.add_inout(out); + args.add_scalar(0); + args.add_scalar(PRODUCER_SPIN_ITERS); + args.launch_spec.set_core_num(PRODUCER_BLOCKS); + args.set_allow_early_resolve(true); + return rt_submit_aic_task(FUNC_SLOW_PRODUCER_AIC, args).task_id(); +} + +static PTO2TaskId submit_aiv_blocker(const Tensor &out, int32_t blocker_count) { + L0TaskArgs args; + args.add_inout(out); + args.add_scalar(BLOCKER_STATUS_BASE_CL); + args.add_scalar(blocker_count); + args.launch_spec.set_core_num(static_cast(blocker_count)); + args.set_allow_early_resolve(true); + return rt_submit_aiv_task(FUNC_SPIN_BLOCKER_AIV, args).task_id(); +} + +static void submit_consumer(const Tensor &out, PTO2TaskId producer, int16_t consumer_blocks) { + L0TaskArgsWithDeps<1> args; + args.add_inout(out); + args.add_scalar(PRODUCER_BLOCKS); + args.launch_spec.set_core_num(consumer_blocks); + args.launch_spec.set_require_sync_start(true); + args.add_dep(producer); + rt_submit_aiv_task(FUNC_SYNC_CONSUMER_AIV, args); +} + +__attribute__((visibility("default"))) void aicpu_orchestration_entry(const L2TaskArgs &orch_args) { + const Tensor &output = orch_args.tensor(0).ref(); + const bool use_pending = orch_args.scalar(0) != 0; + const int16_t consumer_blocks = static_cast(orch_args.scalar(1)); + const int32_t blocker_count = static_cast(rt_available_aiv_count()); + + if (use_pending && (blocker_count <= 0 || blocker_count > BLOCKER_STATUS_CAPACITY || consumer_blocks <= 0 || + consumer_blocks > blocker_count)) { + rt_report_fatal( + PTO2_ERROR_INVALID_ARGS, "invalid pending geometry: blockers=%d consumer=%d", blocker_count, consumer_blocks + ); + return; + } + + rt_scope_begin(PTO2ScopeMode::MANUAL); + PTO2TaskId producer = use_pending ? submit_aiv_blocker(output, blocker_count) : submit_producer(output); + if (use_pending && (!wait_for_blockers_started(output, blocker_count) || !wait_for_scheduler_loop_fence())) return; + submit_consumer(output, producer, consumer_blocks); + if (use_pending && !wait_for_scheduler_loop_fence()) return; + rt_scope_end(); +} + +} // extern "C" diff --git a/tests/st/a5/tensormap_and_ringbuffer/dfx/l2_swimlane/test_sync_start_early_local_owner.py b/tests/st/a5/tensormap_and_ringbuffer/dfx/l2_swimlane/test_sync_start_early_local_owner.py new file mode 100644 index 0000000000..754b0a0b4a --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/dfx/l2_swimlane/test_sync_start_early_local_owner.py @@ -0,0 +1,163 @@ +#!/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. +# ----------------------------------------------------------------------------------------------------------- +"""Regression for #1548: an a5 early sync-start cohort fits one scheduler. + +Idle-capacity cases run with one and three schedulers, using a cohort sized to +fit any local owner in that topology. A pending-capacity case fills every AIV +running slot first, forcing the local owner to stage the eight-block consumer +entirely into free pending slots. A scheduler-loop fence makes the setup +deterministic; a second fence keeps the blockers running through the consumer's +early-dispatch opportunity during their measured one-second hold. Phase capture +then verifies all eight local stages. +""" + +from __future__ import annotations + +from collections import Counter +from pathlib import Path + +import torch +from simpler.task_interface import ArgDirection as D + +from simpler_setup import Scalar, SceneTestCase, TaskArgsBuilder, Tensor, scene_test +from simpler_setup.tools.swimlane_converter import read_perf_data + +FLOATS_PER_CACHE_LINE = 16 +PRODUCER_BLOCKS = 1 +MAX_CONSUMER_BLOCKS = 8 +TOTAL_BLOCKS = PRODUCER_BLOCKS + MAX_CONSUMER_BLOCKS +BLOCKER_STATUS_CAPACITY = 72 +OUTPUT_CACHE_LINES = TOTAL_BLOCKS + BLOCKER_STATUS_CAPACITY + + +@scene_test(level=2, runtime="tensormap_and_ringbuffer") +class TestSyncStartEarlyLocalOwner(SceneTestCase): + RTOL = 0 + ATOL = 0 + + CALLABLE = { + "orchestration": { + "source": "kernels/orchestration/sync_start_early_local_owner_orch.cpp", + "function_name": "aicpu_orchestration_entry", + "signature": [D.INOUT], + }, + "incores": [ + { + "func_id": 0, + "name": "SLOW_FLAGGED_PRODUCER_AIC", + "source": "../../spmd_sync_start_early_dispatch/kernels/aiv/kernel_spmd_write_slow.cpp", + "core_type": "aic", + "signature": [D.INOUT], + }, + { + "func_id": 1, + "name": "SYNC_START_CONSUMER_AIV", + "source": "../../spmd_multiblock_aiv/kernels/aiv/kernel_spmd_write.cpp", + "core_type": "aiv", + "signature": [D.INOUT], + }, + { + "func_id": 2, + "name": "SPIN_BLOCKER_AIV", + "source": "kernels/aiv/kernel_spmd_spin.cpp", + "core_type": "aiv", + "signature": [D.INOUT], + }, + ], + } + + CASES = [ + { + "name": "single_scheduler_idle", + "platforms": ["a5sim", "a5"], + # One orchestrator plus one scheduler. The scheduler owns every AIV + # core, so this eight-block cohort fits its local idle capacity. + "config": {"aicpu_thread_num": 2}, + "params": {"mode": 0, "consumer_blocks": 8}, + }, + { + "name": "three_schedulers_idle", + "platforms": ["a5sim", "a5"], + # One orchestrator plus three schedulers. The popping scheduler can + # stage all four blocks without stopping either peer. + "config": {"aicpu_thread_num": 4}, + "params": {"mode": 0, "consumer_blocks": 4}, + }, + { + "name": "single_scheduler_pending_only", + "platforms": ["a5sim", "a5"], + # Every blocker reports started before an ACK-sweep fence. A second + # scheduler-loop fence completes during their measured one-second + # hold, proving the consumer used pending rather than idle slots. + "config": {"aicpu_thread_num": 2}, + "params": {"mode": 1, "consumer_blocks": 8}, + }, + ] + + def generate_args(self, params): + return TaskArgsBuilder( + Tensor("output", torch.zeros(OUTPUT_CACHE_LINES * FLOATS_PER_CACHE_LINE, dtype=torch.float32)), + Scalar("mode", int(params["mode"])), + Scalar("consumer_blocks", int(params["consumer_blocks"])), + ) + + def compute_golden(self, args, params): + for block_idx in range(PRODUCER_BLOCKS): + args.output[block_idx * FLOATS_PER_CACHE_LINE] = float(block_idx) + for block_idx in range(int(params["consumer_blocks"])): + args.output[(PRODUCER_BLOCKS + block_idx) * FLOATS_PER_CACHE_LINE] = float(block_idx) + + def _build_config(self, config_dict, *args, **kwargs): + config = super()._build_config(config_dict, *args, **kwargs) + self._trace_perf_level = int(kwargs.get("enable_l2_swimlane", args[0] if args else 0)) + output_prefix = kwargs.get("output_prefix", "") + self._trace_perf_path = Path(output_prefix) / "l2_swimlane_records.json" if output_prefix else None + return config + + def compare_outputs(self, test_args, golden_args, output_names, params): + super().compare_outputs(test_args, golden_args, output_names, params) + if getattr(self, "_trace_perf_level", 0) < 3: + return + + perf_path = self._trace_perf_path + assert perf_path is not None, "L2 swimlane enabled without an output prefix" + assert perf_path.exists(), f"l2_swimlane_records.json missing under {perf_path.parent}" + + perf = read_perf_data(perf_path) + assert int(perf.get("l2_swimlane_level", 0)) >= 3, f"scheduler phases missing from {perf_path}" + phase_records = [record for thread in perf.get("aicpu_scheduler_phases", []) for record in thread] + assert phase_records, f"scheduler phase capture is empty under {perf_path}" + phase_counts = Counter(record.get("phase") for record in phase_records) + expected_blocks = int(params["consumer_blocks"]) + drain_phase_names = {"drain", "drain_prepare", "drain_publish"} + drain_records = [record for record in phase_records if record.get("phase") in drain_phase_names] + drain_published = sum( + int(record.get("tasks_processed", 0)) for record in phase_records if record.get("phase") == "drain_publish" + ) + + assert not drain_records, ( + "#1548: the a5 early sync-start consumer fits one scheduler's local " + "AIV running/pending capacity but entered global drain; " + f"drain_publish staged {drain_published}/{expected_blocks} blocks, " + f"mode={params['mode']}, phase_counts={dict(phase_counts)}, artifact={perf_path}" + ) + + early_staged = sum( + int(record.get("tasks_processed", 0)) for record in phase_records if record.get("phase") == "early_dispatch" + ) + assert early_staged == expected_blocks, ( + "the a5 local-owner fast path was not observed: " + f"early_dispatch staged {early_staged}, expected {expected_blocks}; " + f"mode={params['mode']}, phase_counts={dict(phase_counts)}, artifact={perf_path}" + ) + + +if __name__ == "__main__": + SceneTestCase.run_module(__name__) diff --git a/tests/ut/cpp/a2a3/test_scheduler_state.cpp b/tests/ut/cpp/a2a3/test_scheduler_state.cpp index fce152b02a..7711d227c8 100644 --- a/tests/ut/cpp/a2a3/test_scheduler_state.cpp +++ b/tests/ut/cpp/a2a3/test_scheduler_state.cpp @@ -524,3 +524,60 @@ TEST(CoreTrackerTest, MixRunningClusterHelpersRejectOccupiedUsedPendingSlot) { EXPECT_EQ(tracker.count_mix_running_clusters(used_mask), 0); } + +TEST(CoreTrackerTest, CountAvailableBlocksAivIncludesFreePendingSlots) { + CoreTracker tracker; + tracker.init(2); + tracker.set_cluster(0, 0, 1, 2); + tracker.set_cluster(1, 3, 4, 5); + + tracker.change_core_state(1); + tracker.change_core_state(4); + + constexpr uint8_t aiv_mask = PTO2_SUBTASK_MASK_AIV0; + constexpr int32_t exact_fit = 4; + EXPECT_EQ(tracker.count_available_blocks(PTO2ResourceShape::AIV, aiv_mask, false), 2); + EXPECT_EQ(tracker.count_available_blocks(PTO2ResourceShape::AIV, aiv_mask, true), exact_fit); + + tracker.set_pending_occupied(4); + EXPECT_EQ(tracker.count_available_blocks(PTO2ResourceShape::AIV, aiv_mask, true), exact_fit - 1); +} + +TEST(CoreTrackerTest, CountAvailableBlocksMixCountsLogicalClusters) { + CoreTracker tracker; + tracker.init(2); + tracker.set_cluster(0, 0, 1, 2); + tracker.set_cluster(1, 3, 4, 5); + + tracker.change_core_state(4); + + constexpr uint8_t used_mask = PTO2_SUBTASK_MASK_AIC | PTO2_SUBTASK_MASK_AIV0 | PTO2_SUBTASK_MASK_AIV1; + EXPECT_EQ(tracker.count_available_blocks(PTO2ResourceShape::MIX, used_mask, false), 1); + EXPECT_EQ(tracker.count_available_blocks(PTO2ResourceShape::MIX, used_mask, true), 2); +} + +TEST(CoreTrackerTest, CountAvailableBlocksMixRejectsOccupiedUsedPendingSlot) { + CoreTracker tracker; + tracker.init(1); + tracker.set_cluster(0, 0, 1, 2); + + constexpr uint8_t used_mask = PTO2_SUBTASK_MASK_AIC | PTO2_SUBTASK_MASK_AIV0; + tracker.change_core_state(1); + tracker.set_pending_occupied(1); + + EXPECT_EQ(tracker.count_available_blocks(PTO2ResourceShape::MIX, used_mask, false), 0); + EXPECT_EQ(tracker.count_available_blocks(PTO2ResourceShape::MIX, used_mask, true), 0); +} + +TEST(CoreTrackerTest, CountAvailableBlocksMixIgnoresUnavailableUnusedCore) { + CoreTracker tracker; + tracker.init(1); + tracker.set_cluster(0, 0, 1, 2); + + constexpr uint8_t used_mask = PTO2_SUBTASK_MASK_AIC | PTO2_SUBTASK_MASK_AIV0; + tracker.change_core_state(2); + tracker.set_pending_occupied(2); + + EXPECT_EQ(tracker.count_available_blocks(PTO2ResourceShape::MIX, used_mask, false), 1); + EXPECT_EQ(tracker.count_available_blocks(PTO2ResourceShape::MIX, used_mask, true), 1); +} diff --git a/tests/ut/cpp/a2a3/test_wiring.cpp b/tests/ut/cpp/a2a3/test_wiring.cpp index a13fe75f32..8f781c914a 100644 --- a/tests/ut/cpp/a2a3/test_wiring.cpp +++ b/tests/ut/cpp/a2a3/test_wiring.cpp @@ -701,7 +701,7 @@ TEST_F(WiringTest, SyncStartDoorbellPassHasOneOwner) { } } -TEST_F(WiringTest, SyncStartDrainFinalizeRetriesProducerFirstRendezvous) { +TEST_F(WiringTest, SyncStartStagingFinalizeRetriesProducerFirstRendezvous) { alignas(64) PTO2TaskSlotState sync_consumer, downstream; init_slot(sync_consumer, PTO2_TASK_PENDING, 1, 1); init_slot(downstream, PTO2_TASK_PENDING, 1, 1); @@ -721,20 +721,51 @@ TEST_F(WiringTest, SyncStartDrainFinalizeRetriesProducerFirstRendezvous) { 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. + // Producer release wins before staging publishes its running-slot seed. With no pending + // promotions, staging 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_TRUE(sched.retry_sync_start_rendezvous_after_staging(sync_consumer)); EXPECT_EQ(sync_consumer.payload->early_dispatch_launch_state.load(), PTO2_EARLY_DISPATCH_LAUNCH_COMPLETE); EXPECT_TRUE(sync_consumer.has_dispatch_propagated()); EXPECT_EQ(downstream.payload->dispatch_fanin.load(), 1); - EXPECT_FALSE(sched.retry_sync_start_rendezvous_after_drain(sync_consumer)); + EXPECT_FALSE(sched.retry_sync_start_rendezvous_after_staging(sync_consumer)); + EXPECT_EQ(downstream.payload->dispatch_fanin.load(), 1); +} + +TEST_F(WiringTest, SyncStartProducerReleaseCompletesStagerFirstRendezvous) { + 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.task_attrs.set_sync_start(); + sync_consumer.task_attrs.set_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(2, 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; + + EXPECT_FALSE(sched.retry_sync_start_rendezvous_after_staging(sync_consumer)); + 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_COMPLETE); + EXPECT_TRUE(sync_consumer.has_dispatch_propagated()); + EXPECT_EQ(downstream.payload->dispatch_fanin.load(), 1); + + EXPECT_FALSE(sched.retry_sync_start_rendezvous_after_staging(sync_consumer)); EXPECT_EQ(downstream.payload->dispatch_fanin.load(), 1); } @@ -807,7 +838,7 @@ TEST_F(WiringTest, ArmedEarlySyncDrainKeepsEveryStagerGatedAfterReady) { ); } -TEST_F(WiringTest, DrainFinishBetweenReleasePhasesRetainsOwner) { +TEST_F(WiringTest, EarlySyncFinishBetweenReleasePhasesRetainsOwnerCompleteState) { alignas(64) PTO2TaskSlotState sync_consumer; init_slot(sync_consumer, PTO2_TASK_PENDING, 1, 1); sync_consumer.active_mask = ActiveMask(PTO2_SUBTASK_MASK_AIV0); diff --git a/tests/ut/cpp/a5/test_scheduler_state.cpp b/tests/ut/cpp/a5/test_scheduler_state.cpp index f44ef78cf8..784d09d6c5 100644 --- a/tests/ut/cpp/a5/test_scheduler_state.cpp +++ b/tests/ut/cpp/a5/test_scheduler_state.cpp @@ -524,3 +524,60 @@ TEST(CoreTrackerTest, MixRunningClusterHelpersRejectOccupiedUsedPendingSlot) { EXPECT_EQ(tracker.count_mix_running_clusters(used_mask), 0); } + +TEST(CoreTrackerTest, CountAvailableBlocksAivIncludesFreePendingSlots) { + CoreTracker tracker; + tracker.init(2); + tracker.set_cluster(0, 0, 1, 2); + tracker.set_cluster(1, 3, 4, 5); + + tracker.change_core_state(1); + tracker.change_core_state(4); + + constexpr uint8_t aiv_mask = PTO2_SUBTASK_MASK_AIV0; + constexpr int32_t exact_fit = 4; + EXPECT_EQ(tracker.count_available_blocks(PTO2ResourceShape::AIV, aiv_mask, false), 2); + EXPECT_EQ(tracker.count_available_blocks(PTO2ResourceShape::AIV, aiv_mask, true), exact_fit); + + tracker.set_pending_occupied(4); + EXPECT_EQ(tracker.count_available_blocks(PTO2ResourceShape::AIV, aiv_mask, true), exact_fit - 1); +} + +TEST(CoreTrackerTest, CountAvailableBlocksMixCountsLogicalClusters) { + CoreTracker tracker; + tracker.init(2); + tracker.set_cluster(0, 0, 1, 2); + tracker.set_cluster(1, 3, 4, 5); + + tracker.change_core_state(4); + + constexpr uint8_t used_mask = PTO2_SUBTASK_MASK_AIC | PTO2_SUBTASK_MASK_AIV0 | PTO2_SUBTASK_MASK_AIV1; + EXPECT_EQ(tracker.count_available_blocks(PTO2ResourceShape::MIX, used_mask, false), 1); + EXPECT_EQ(tracker.count_available_blocks(PTO2ResourceShape::MIX, used_mask, true), 2); +} + +TEST(CoreTrackerTest, CountAvailableBlocksMixRejectsOccupiedUsedPendingSlot) { + CoreTracker tracker; + tracker.init(1); + tracker.set_cluster(0, 0, 1, 2); + + constexpr uint8_t used_mask = PTO2_SUBTASK_MASK_AIC | PTO2_SUBTASK_MASK_AIV0; + tracker.change_core_state(1); + tracker.set_pending_occupied(1); + + EXPECT_EQ(tracker.count_available_blocks(PTO2ResourceShape::MIX, used_mask, false), 0); + EXPECT_EQ(tracker.count_available_blocks(PTO2ResourceShape::MIX, used_mask, true), 0); +} + +TEST(CoreTrackerTest, CountAvailableBlocksMixIgnoresUnavailableUnusedCore) { + CoreTracker tracker; + tracker.init(1); + tracker.set_cluster(0, 0, 1, 2); + + constexpr uint8_t used_mask = PTO2_SUBTASK_MASK_AIC | PTO2_SUBTASK_MASK_AIV0; + tracker.change_core_state(2); + tracker.set_pending_occupied(2); + + EXPECT_EQ(tracker.count_available_blocks(PTO2ResourceShape::MIX, used_mask, false), 1); + EXPECT_EQ(tracker.count_available_blocks(PTO2ResourceShape::MIX, used_mask, true), 1); +} diff --git a/tests/ut/cpp/a5/test_wiring.cpp b/tests/ut/cpp/a5/test_wiring.cpp index 24b25a663c..323bad1baf 100644 --- a/tests/ut/cpp/a5/test_wiring.cpp +++ b/tests/ut/cpp/a5/test_wiring.cpp @@ -259,6 +259,124 @@ TEST_F(WiringTest, WireTaskMixedProducerStates) { EXPECT_EQ(producers[2].fanout_head, nullptr); // early finished } +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, SyncStartStagingFinalizeRetriesProducerFirstRendezvous) { + 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.task_attrs.set_sync_start(); + sync_consumer.task_attrs.set_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; + + 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_staging(sync_consumer)); + EXPECT_EQ(sync_consumer.payload->early_dispatch_launch_state.load(), PTO2_EARLY_DISPATCH_LAUNCH_COMPLETE); + EXPECT_TRUE(sync_consumer.has_dispatch_propagated()); + EXPECT_EQ(downstream.payload->dispatch_fanin.load(), 1); + + EXPECT_FALSE(sched.retry_sync_start_rendezvous_after_staging(sync_consumer)); + EXPECT_EQ(downstream.payload->dispatch_fanin.load(), 1); +} + +TEST_F(WiringTest, SyncStartProducerReleaseCompletesStagerFirstRendezvous) { + 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.task_attrs.set_sync_start(); + sync_consumer.task_attrs.set_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(2, 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; + + EXPECT_FALSE(sched.retry_sync_start_rendezvous_after_staging(sync_consumer)); + 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_COMPLETE); + EXPECT_TRUE(sync_consumer.has_dispatch_propagated()); + EXPECT_EQ(downstream.payload->dispatch_fanin.load(), 1); + + EXPECT_FALSE(sched.retry_sync_start_rendezvous_after_staging(sync_consumer)); + EXPECT_EQ(downstream.payload->dispatch_fanin.load(), 1); +} + +TEST_F(WiringTest, EarlySyncFinishBetweenReleasePhasesRetainsOwnerCompleteState) { + 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.task_attrs.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_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 + ); +} + // ============================================================================= // on_task_complete: notifies consumers via fanout chain // =============================================================================