From 60e59a4e5682f18b707fd8adf01d3b3233bb34ab Mon Sep 17 00:00:00 2001 From: Chao Wang <26245345+ChaoWao@users.noreply.github.com> Date: Sat, 11 Jul 2026 03:15:51 -0700 Subject: [PATCH] Fix: make early-dispatch doorbell ownership exclusive - Atomically split staged-core bits between release and late stagers - Serialize release with a NONE/RINGING/COMPLETE launch state - Ring captured local handles before publishing late-staged blocks - Retry fanout after both publication and release race orders - Keep preclaimed ranges gated across concurrent producer release --- .../runtime/pto_runtime2_types.h | 28 ++-- .../runtime/scheduler/pto_scheduler.h | 74 ++++++--- .../runtime/scheduler/scheduler_context.h | 9 +- .../runtime/scheduler/scheduler_dispatch.cpp | 89 ++++++----- tests/ut/cpp/a2a3/test_wiring.cpp | 144 ++++++++++++++++++ tests/ut/cpp/stubs/test_stubs.cpp | 20 ++- 6 files changed, 293 insertions(+), 71 deletions(-) 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 b5a44d0fc4..e5730d8717 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 @@ -219,8 +219,14 @@ struct PTO2TaskDescriptor { enum PTO2EarlyDispatchState : uint8_t { PTO2_EARLY_DISPATCH_NONE = 0, // not pre-staged PTO2_EARLY_DISPATCH_STAGING = 1, // Hook 1 claimed it; staging in progress - PTO2_EARLY_DISPATCH_STAGED = 2, // staged on a core, gated; staged_* fields valid - PTO2_EARLY_DISPATCH_DISPATCHED = 3 // routed via the normal dispatch path (no pre-stage) + PTO2_EARLY_DISPATCH_STAGED = 2, // reserved + PTO2_EARLY_DISPATCH_DISPATCHED = 3 // producers released; staged blocks may still be gated +}; + +enum PTO2EarlyDispatchLaunchState : uint8_t { + PTO2_EARLY_DISPATCH_LAUNCH_NONE = 0, + PTO2_EARLY_DISPATCH_LAUNCH_RINGING = 1, + PTO2_EARLY_DISPATCH_LAUNCH_COMPLETE = 2, }; // A pre-staged consumer occupies one core per gated subtask block. WHICH cores @@ -250,8 +256,8 @@ struct PTO2TaskPayload { // 576), so sizeof and tensors[] are unchanged. // // Bitmask of global core_ids this consumer is pre-staged (gated) on. Set with - // atomic fetch_or by concurrent stagers; read by release. (Re)initialized in - // PTO2TaskPayload::init before the slot can be staged again. + // atomic fetch_or by concurrent stagers, then destructively split between the + // release and late-stager paths. (Re)initialized in PTO2TaskPayload::init. 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 @@ -270,11 +276,14 @@ struct PTO2TaskPayload { // 3=DISPATCHED (2=STAGED is unused now). STAGING is the STABLE gated state — // many threads stage blocks concurrently while it holds, each claiming a block // via the atomic next_block_idx and OR-ing its cores into staged_core_mask. - // Release does STAGING->DISPATCHED then rings the mask; a thread that stages a - // block AFTER release flipped DISPATCHED rings that block's doorbell itself - // (self-ring), so no doorbell is ever missed. + // Release does STAGING->DISPATCHED and claims the current mask; a thread that + // stages a block after that flip claims and rings only its remaining bits. std::atomic early_dispatch_state{0}; std::atomic dispatch_propagated{0}; // PRODUCER side: once-guard for fanout propagation + // The release owner publishes COMPLETE only after all doorbells it claimed + // are visible. Combined with published_block_count, this keeps fanout + // private until release-owned and late-owned blocks have both launched. + std::atomic early_dispatch_launch_state{PTO2_EARLY_DISPATCH_LAUNCH_NONE}; // === Cache lines 9-72 (4096B) — tensors (alignas(64) forces alignment) === Tensor tensors[MAX_TENSOR_ARGS]; // === Cache lines 73-74 (128B) — scalars === @@ -350,14 +359,15 @@ struct PTO2TaskPayload { // one of ITS producers is flagged (propagate_dispatch_fanin bumps // dispatch_fanin and may CAS early_dispatch_state on any consumer, independent of the // consumer's own hint). So they MUST be zeroed here unconditionally. - // published_block_count and dispatch_propagated are producer-side, but share - // this same per-submit lifetime and are reset here too. + // Publication, propagation, and launch fields share this same + // per-submit lifetime and are reset here too. early_dispatch_state.store(PTO2_EARLY_DISPATCH_NONE, std::memory_order_relaxed); for (int w = 0; w < PTO2_EARLY_DISPATCH_CORE_MASK_WORDS; w++) staged_core_mask[w].store(0, std::memory_order_relaxed); dispatch_fanin.store(0, std::memory_order_relaxed); dispatch_propagated.store(0, std::memory_order_relaxed); published_block_count.store(0, std::memory_order_relaxed); + early_dispatch_launch_state.store(PTO2_EARLY_DISPATCH_LAUNCH_NONE, std::memory_order_relaxed); } }; 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 4e1dcb1e31..c32efcb785 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 @@ -32,6 +32,7 @@ #include #include "common/core_type.h" +#include "common/memory_barrier.h" #include "utils/device_arena.h" #include "aicpu/platform_regs.h" // get_reg_ptr / RegId for the early-dispatch doorbell #include "pto_async_wait.h" @@ -657,8 +658,8 @@ struct PTO2SchedulerState { // PARTIALLY pre-staged — the gated blocks are released by the doorbells rung // here, and the remaining (next_block_idx .. logical_block_num) blocks // dispatch normally off the ready queue. Lock-free claim shared with Hook 1 - // (the stager): CAS NONE->DISPATCHED wins => not pre-staged; lose => STAGED - // (spin past the brief STAGING window so the mask is visible), then ring. + // (the stager): CAS NONE->DISPATCHED wins => not pre-staged; otherwise flip + // STAGING->DISPATCHED and destructively claim the published doorbell bits. // Per-core early-dispatch doorbell table. Hook 1 records each gated core's // (reg_addr, dispatch token) here at stage time; the completion-path release @@ -698,6 +699,42 @@ struct PTO2SchedulerState { *dmb = (tk << 32) | tk; // 64-bit STR: high=low=token releases the gated AICore } + inline void ring_staged_doorbell_bits(int word, uint64_t bits) { + while (bits != 0) { + int core_id = word * 64 + __builtin_ctzll(bits); + bits &= bits - 1; + ring_one_doorbell( + early_dispatch_doorbell_table[core_id].addr, early_dispatch_doorbell_table[core_id].token + ); + } + } + + static inline uint64_t claim_all_staged_doorbell_bits(std::atomic &mask) { + return mask.exchange(0, std::memory_order_seq_cst); + } + + static inline uint64_t claim_late_staged_doorbell_bits(std::atomic &mask, uint64_t candidates) { + return mask.fetch_and(~candidates, std::memory_order_seq_cst) & candidates; + } + + static inline bool should_gate_early_dispatch(bool force_gate, uint8_t early_dispatch_state) { + return force_gate || early_dispatch_state == PTO2_EARLY_DISPATCH_STAGING; + } + + static inline bool + ring_claimed_local_doorbell(uint64_t claimed_word, int core_id, uint64_t reg_addr, uint32_t token) { + if ((claimed_word & (1ULL << (core_id & 63))) == 0) return false; + ring_one_doorbell(reg_addr, token); + return true; + } + + static inline bool try_claim_early_dispatch_launch(PTO2TaskPayload &payload) { + uint8_t expected = PTO2_EARLY_DISPATCH_LAUNCH_NONE; + return payload.early_dispatch_launch_state.compare_exchange_strong( + expected, PTO2_EARLY_DISPATCH_LAUNCH_RINGING, std::memory_order_seq_cst, std::memory_order_seq_cst + ); + } + inline void record_published_blocks(PTO2TaskSlotState &slot_state, int32_t count) { if (count <= 0 || !slot_state.allow_early_resolve) return; slot_state.payload->published_block_count.fetch_add(static_cast(count), std::memory_order_seq_cst); @@ -716,6 +753,10 @@ struct PTO2SchedulerState { void propagate_dispatch_fanin(PTO2TaskSlotState &p) { if (!p.allow_early_resolve) return; // only codegen-flagged (direct) producers propagate if (p.payload->published_block_count.load(std::memory_order_seq_cst) < p.logical_block_num) return; + if (p.payload->early_dispatch_state.load(std::memory_order_seq_cst) == PTO2_EARLY_DISPATCH_STAGING) return; + if (p.payload->early_dispatch_launch_state.load(std::memory_order_seq_cst) == + PTO2_EARLY_DISPATCH_LAUNCH_RINGING) + return; if (p.payload->dispatch_propagated.exchange(1, std::memory_order_acq_rel) != 0) return; // already propagated once p.lock_fanout(); @@ -767,28 +808,25 @@ struct PTO2SchedulerState { )) { return false; } - // Staged (STAGING). Flip STAGING->DISPATCHED, THEN read the mask. seq_cst - // gives a total order with the concurrent stagers, each of which OR-s its - // core into the mask and THEN loads early_dispatch_state: a stager whose bit lands - // before this CAS is read here and rung; a stager whose bit lands after - // sees DISPATCHED and rings that core itself (self-ring in - // stage_consumer_blocks). Either way every gated core's doorbell fires once - // (a double-ring is harmless — the AICore already matched). This replaces - // the old transient-STAGING spin: STAGING is now the stable gated state. + // route_ready_once admits one caller, but keep the helper defensive: a + // duplicate that observes or loses an in-progress launch must never + // route the same partial task a second time. + if (expect != PTO2_EARLY_DISPATCH_STAGING || !try_claim_early_dispatch_launch(*slot_state.payload)) return true; expect = PTO2_EARLY_DISPATCH_STAGING; slot_state.payload->early_dispatch_state.compare_exchange_strong( expect, PTO2_EARLY_DISPATCH_DISPATCHED, std::memory_order_seq_cst, std::memory_order_seq_cst ); + // Destructively claim every published bit. A stager racing this pass + // can claim only bits that land after the exchange, so each gated core + // has exactly one doorbell writer. for (int w = 0; w < PTO2_EARLY_DISPATCH_CORE_MASK_WORDS; w++) { - uint64_t bits = slot_state.payload->staged_core_mask[w].load(std::memory_order_seq_cst); - while (bits != 0) { - int core_id = w * 64 + __builtin_ctzll(bits); - bits &= bits - 1; - ring_one_doorbell( - early_dispatch_doorbell_table[core_id].addr, early_dispatch_doorbell_table[core_id].token - ); - } + uint64_t owned = claim_all_staged_doorbell_bits(slot_state.payload->staged_core_mask[w]); + ring_staged_doorbell_bits(w, owned); } + wmb(); + slot_state.payload->early_dispatch_launch_state.store( + PTO2_EARLY_DISPATCH_LAUNCH_COMPLETE, std::memory_order_seq_cst + ); // This pre-staged consumer was just released by its doorbell — it starts // running NOW, so propagate dispatch_fanin to ITS consumers (only if it is // itself codegen-flagged; the gate inside no-ops otherwise). Defer it via the 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 4143242db5..b8023c6c56 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 @@ -230,7 +230,8 @@ class SchedulerContext { ); void build_payload( - PTO2DispatchPayload &dispatch_payload, PTO2TaskSlotState &slot_state, PTO2SubtaskSlot subslot, int32_t block_idx + PTO2DispatchPayload &dispatch_payload, PTO2TaskSlotState &slot_state, PTO2SubtaskSlot subslot, + int32_t block_idx, bool force_gate ); // Batched-dispatch primitives. prepare_* builds the payload and per-core @@ -251,7 +252,7 @@ class SchedulerContext { PublishHandle prepare_subtask_to_core( int32_t thread_idx, int32_t core_offset, PTO2TaskSlotState &slot_state, PTO2SubtaskSlot subslot, - bool to_pending, int32_t block_idx + bool to_pending, int32_t block_idx, bool force_gate ); inline void publish_subtask_to_core(const PublishHandle &h, uint64_t dispatch_ts) { @@ -298,7 +299,7 @@ class SchedulerContext { // caller-supplied handles buffer. Returns the number of handles written. int prepare_block_for_dispatch( int32_t thread_idx, int32_t core_offset, PTO2TaskSlotState &slot_state, PTO2ResourceShape shape, - bool to_pending, int32_t block_idx, PublishHandle *out_handles + bool to_pending, int32_t block_idx, PublishHandle *out_handles, bool force_gate = false ); void dispatch_shape( @@ -311,7 +312,7 @@ class SchedulerContext { // is queued) and sets made_progress / try_pushed when it stages, so the caller // is a single unconditional call like normal dispatch. After normal dispatch // leaves idle cores spare, pre-stage the consumers of any RUNNING flagged - // producer onto those cores with not_ready=1 (gated). Touches no dependency + // producer onto those cores with a non-zero src_payload (gated). Touches no dependency // state — the task is released by the doorbell at its normal ready-pop (Hook 2). int32_t try_early_dispatch( int32_t thread_idx, CoreTracker &tracker, bool pmu_active, bool &made_progress, bool &try_pushed 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 4dac1c4215..d4153496fd 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 @@ -37,7 +37,7 @@ #define unlikely(x) __builtin_expect(!!(x), 0) #endif -// AICore materializes args[] from src_payload on the not_ready path using the +// AICore materializes args[] from src_payload on the gated path using the // byte offsets in pto2_dispatch_payload.h (the AICore .o cannot see PTO2TaskPayload). // Pin those constants to the real layout here, where the struct is fully visible. static_assert(offsetof(PTO2TaskPayload, tensor_count) == PTO2_TASKPAYLOAD_TENSOR_COUNT_OFFSET); @@ -123,17 +123,20 @@ int SchedulerContext::pop_ready_tasks_batch( } void SchedulerContext::build_payload( - PTO2DispatchPayload &dispatch_payload, PTO2TaskSlotState &slot_state, PTO2SubtaskSlot subslot, int32_t block_idx + PTO2DispatchPayload &dispatch_payload, PTO2TaskSlotState &slot_state, PTO2SubtaskSlot subslot, int32_t block_idx, + bool force_gate ) { int32_t slot_idx = static_cast(subslot); uint64_t callable_addr = get_function_bin_addr(slot_state.task->kernel_id[slot_idx]); const CoreCallable *callable = reinterpret_cast(callable_addr); dispatch_payload.function_bin_addr = callable->resolved_addr(); auto &payload = *slot_state.payload; - // Early-dispatch: a task being staged (Hook 1 set early_dispatch_state to - // STAGING before this call) is gated — the AICore must wait for the - // DATA_MAIN_BASE high-32 doorbell. All other dispatches run on pickup. - if (payload.early_dispatch_state.load(std::memory_order_relaxed) == PTO2_EARLY_DISPATCH_STAGING) { + // A claimed early-stage range stays gated even if producer completion flips + // the shared state before this payload is built. All other dispatches run on + // pickup. + if (PTO2SchedulerState::should_gate_early_dispatch( + force_gate, payload.early_dispatch_state.load(std::memory_order_relaxed) + )) { // Gated task: hand the idle AICore the source payload (non-zero = gate) and // let it fill args[0..num_args) itself during its doorbell wait, instead of // paying the arg-vector write on this scheduler thread. @@ -162,7 +165,7 @@ void SchedulerContext::build_payload( SchedulerContext::PublishHandle SchedulerContext::prepare_subtask_to_core( int32_t thread_idx, int32_t core_offset, PTO2TaskSlotState &slot_state, PTO2SubtaskSlot subslot, bool to_pending, - int32_t block_idx + int32_t block_idx, bool force_gate ) { CoreTracker &tracker = core_trackers_[thread_idx]; auto core_id = tracker.get_core_id_by_offset(core_offset); @@ -185,7 +188,7 @@ SchedulerContext::PublishHandle SchedulerContext::prepare_subtask_to_core( // deferred completion (count > 0), which is rare. A non-deferred task never // touches count/error_code, so the slab stays clean without a per-dispatch // write — keeping this cold per-core line off the dispatch path. - build_payload(payload, slot_state, subslot, block_idx); + build_payload(payload, slot_state, subslot, block_idx, force_gate); if (to_pending) { core_exec_state.pending_subslot = subslot; @@ -234,7 +237,7 @@ SchedulerContext::PublishHandle SchedulerContext::prepare_subtask_to_core( int SchedulerContext::prepare_block_for_dispatch( int32_t thread_idx, int32_t core_offset, PTO2TaskSlotState &slot_state, PTO2ResourceShape shape, bool to_pending, - int32_t block_idx, PublishHandle *out_handles + int32_t block_idx, PublishHandle *out_handles, bool force_gate ) { #if SIMPLER_DFX if (is_dump_args_enabled()) { @@ -256,19 +259,22 @@ int SchedulerContext::prepare_block_for_dispatch( if (cmask & PTO2_SUBTASK_MASK_AIC) { bool p = to_pending && !tracker.is_aic_core_idle(core_offset); out_handles[n++] = prepare_subtask_to_core( - thread_idx, tracker.get_aic_core_offset(core_offset), slot_state, PTO2SubtaskSlot::AIC, p, block_idx + thread_idx, tracker.get_aic_core_offset(core_offset), slot_state, PTO2SubtaskSlot::AIC, p, block_idx, + force_gate ); } if (cmask & PTO2_SUBTASK_MASK_AIV0) { bool p = to_pending && !tracker.is_aiv0_core_idle(core_offset); out_handles[n++] = prepare_subtask_to_core( - thread_idx, tracker.get_aiv0_core_offset(core_offset), slot_state, PTO2SubtaskSlot::AIV0, p, block_idx + thread_idx, tracker.get_aiv0_core_offset(core_offset), slot_state, PTO2SubtaskSlot::AIV0, p, block_idx, + force_gate ); } if (cmask & PTO2_SUBTASK_MASK_AIV1) { bool p = to_pending && !tracker.is_aiv1_core_idle(core_offset); out_handles[n++] = prepare_subtask_to_core( - thread_idx, tracker.get_aiv1_core_offset(core_offset), slot_state, PTO2SubtaskSlot::AIV1, p, block_idx + thread_idx, tracker.get_aiv1_core_offset(core_offset), slot_state, PTO2SubtaskSlot::AIV1, p, block_idx, + force_gate ); } #if SIMPLER_DFX @@ -276,15 +282,17 @@ int SchedulerContext::prepare_block_for_dispatch( #endif return n; } else if (shape == PTO2ResourceShape::AIC) { - out_handles[0] = - prepare_subtask_to_core(thread_idx, core_offset, slot_state, PTO2SubtaskSlot::AIC, to_pending, block_idx); + out_handles[0] = prepare_subtask_to_core( + thread_idx, core_offset, slot_state, PTO2SubtaskSlot::AIC, to_pending, block_idx, force_gate + ); #if SIMPLER_DFX sched_l2_swimlane_[thread_idx].phase_dispatch_count += 1; #endif return 1; } else { - out_handles[0] = - prepare_subtask_to_core(thread_idx, core_offset, slot_state, PTO2SubtaskSlot::AIV0, to_pending, block_idx); + out_handles[0] = prepare_subtask_to_core( + thread_idx, core_offset, slot_state, PTO2SubtaskSlot::AIV0, to_pending, block_idx, force_gate + ); #if SIMPLER_DFX sched_l2_swimlane_[thread_idx].phase_dispatch_count += 1; #endif @@ -599,10 +607,10 @@ void SchedulerContext::dispatch_ready_tasks( // waiting for an ack the gated task never sends). Each staged core stays // pending_occupied while gated, so no second gated block stacks on it. // -// Self-ring: release flips STAGING->DISPATCHED then rings the mask. A block staged -// after that flip isn't in the mask release read, so this thread rings it here. The -// seq_cst order between "OR mask then load early_dispatch_state" (here) and "store DISPATCHED -// then read mask" (release) guarantees every gated core's doorbell fires. +// Doorbell ownership: release flips STAGING->DISPATCHED and exchanges the shared +// mask to claim its bits. A late stager ORs its bits, then fetch-and-clears only +// those release did not take and rings them from its immutable local handles. +// The seq_cst order guarantees every gated core has exactly one writer. int32_t SchedulerContext::stage_consumer_blocks( int32_t thread_idx, PTO2TaskSlotState *c, PTO2ResourceShape shape, int32_t start, int32_t count, CoreTracker::BitStates &idle, CoreTracker::BitStates &pend @@ -611,13 +619,13 @@ int32_t SchedulerContext::stage_consumer_blocks( // Stamp the real pre-stage time (NOT 0) so the swimlane shows these blocks // dispatched during the producer's run, not at trace start. uint64_t early_dispatch_ts = get_sys_cnt_aicpu(); - uint64_t my_cores[PTO2_EARLY_DISPATCH_CORE_MASK_WORDS] = {0}; // cores this thread gated (for self-ring) + uint64_t my_cores[PTO2_EARLY_DISPATCH_CORE_MASK_WORDS] = {0}; // cores gated by this staging pass int32_t staged = 0; int32_t block = start; // Mirror the normal flush_publish (scheduler_dispatch.cpp wmb()+publish loop): // prepare ALL claimed blocks' payloads (idle bucket -> running slot, pend bucket // -> gated pending), then ONE wmb(), then publish. The wmb guarantees the - // not_ready gate + args are globally visible before any DATA_MAIN_BASE token — + // src_payload gate + source args are globally visible before any DATA_MAIN_BASE token — // without it a gated core can pick up the token and dcci a stale payload. The // shared `count` budget bounds total blocks <= free clusters/cores, so both // buckets fit one handles[] buffer. @@ -626,7 +634,9 @@ int32_t SchedulerContext::stage_consumer_blocks( auto prepare_from = [&](CoreTracker::BitStates &avail, bool to_pending) { while (count > 0 && avail.has_value()) { int32_t core_offset = avail.pop_first(); - n += prepare_block_for_dispatch(thread_idx, core_offset, *c, shape, to_pending, block, &handles[n]); + n += prepare_block_for_dispatch( + thread_idx, core_offset, *c, shape, to_pending, block, &handles[n], /*force_gate=*/true + ); block++; count--; staged++; @@ -650,27 +660,34 @@ int32_t SchedulerContext::stage_consumer_blocks( if (my_cores[w] != 0) c->payload->staged_core_mask[w].fetch_or(my_cores[w], std::memory_order_seq_cst); // Full publication and release are independent events. The seq_cst - // count/state rechecks form a two-sided handshake: release propagates if - // publication won, and the final stager propagates if release won. - sched_->record_published_blocks(*c, staged); + // state/launch/count operations form a two-sided handshake. A released + // block must ring before contributing to the publication count. bool released = staged > 0 && c->payload->early_dispatch_state.load(std::memory_order_seq_cst) == PTO2_EARLY_DISPATCH_DISPATCHED; - // If release already flipped DISPATCHED, it may have read the mask before our - // bits landed — ring our own cores so none is left gated forever. + // Claim only bits the release path did not take. Local handles remain valid + // even if the shared per-core table is reused before this thread resumes. if (released) { + uint64_t owned[PTO2_EARLY_DISPATCH_CORE_MASK_WORDS] = {0}; for (int w = 0; w < PTO2_EARLY_DISPATCH_CORE_MASK_WORDS; w++) { - uint64_t bits = my_cores[w]; - while (bits != 0) { - int cid = w * 64 + __builtin_ctzll(bits); - bits &= bits - 1; - PTO2SchedulerState::ring_one_doorbell( - sched_->early_dispatch_doorbell_table[cid].addr, sched_->early_dispatch_doorbell_table[cid].token - ); + if (my_cores[w] != 0) { + owned[w] = + PTO2SchedulerState::claim_late_staged_doorbell_bits(c->payload->staged_core_mask[w], my_cores[w]); } } + for (int i = 0; i < n; i++) { + int32_t cid = tracker.get_core_id_by_offset(handles[i].core_offset); + PTO2SchedulerState::ring_claimed_local_doorbell( + owned[cid >> 6], cid, handles[i].reg_addr, handles[i].reg_task_id + ); + } + wmb(); } - if (released) sched_->propagate_dispatch_fanin(*c); + sched_->record_published_blocks(*c, staged); + // Retry unconditionally after publication. The guards are cheap, and a + // pre-ring state read can become stale if release completes before this + // count update. + sched_->propagate_dispatch_fanin(*c); return staged; } diff --git a/tests/ut/cpp/a2a3/test_wiring.cpp b/tests/ut/cpp/a2a3/test_wiring.cpp index ad67508a7d..c491fca56d 100644 --- a/tests/ut/cpp/a2a3/test_wiring.cpp +++ b/tests/ut/cpp/a2a3/test_wiring.cpp @@ -33,6 +33,10 @@ #include "utils/device_arena.h" #include "scheduler/pto_scheduler.h" +void reset_test_reg_stub(); +uint64_t get_test_reg_stub_value(); +uint64_t get_test_reg_stub_base_addr(); + // ============================================================================= // Fixture: sets up runtime state with shared memory and provides helpers // ============================================================================= @@ -378,6 +382,146 @@ TEST_F(WiringTest, EarlyDispatchWaitsForAllProducerBlocksPublished) { EXPECT_EQ(payload.dispatch_fanin.load(), payload.fanin_actual_count); } +TEST_F(WiringTest, EarlyDispatchDoorbellBitsHaveOneOwner) { + constexpr uint64_t all_bits = 0b1111; + constexpr uint64_t late_bits = 0b1010; + + std::atomic release_first{all_bits}; + uint64_t release_owned = PTO2SchedulerState::claim_all_staged_doorbell_bits(release_first); + uint64_t late_owned = PTO2SchedulerState::claim_late_staged_doorbell_bits(release_first, late_bits); + EXPECT_EQ(release_owned, all_bits); + EXPECT_EQ(late_owned, 0); + + std::atomic late_first{all_bits}; + late_owned = PTO2SchedulerState::claim_late_staged_doorbell_bits(late_first, late_bits); + release_owned = PTO2SchedulerState::claim_all_staged_doorbell_bits(late_first); + EXPECT_EQ(late_owned, late_bits); + EXPECT_EQ(release_owned, all_bits & ~late_bits); + EXPECT_EQ(release_owned & late_owned, 0); + EXPECT_EQ(release_owned | late_owned, all_bits); + EXPECT_EQ(late_first.load(std::memory_order_acquire), 0); + + std::atomic published_after_release{0}; + release_owned = PTO2SchedulerState::claim_all_staged_doorbell_bits(published_after_release); + published_after_release.fetch_or(late_bits, std::memory_order_seq_cst); + late_owned = PTO2SchedulerState::claim_late_staged_doorbell_bits(published_after_release, late_bits); + EXPECT_EQ(release_owned, 0); + EXPECT_EQ(late_owned, late_bits); + EXPECT_EQ(published_after_release.load(std::memory_order_acquire), 0); +} + +TEST_F(WiringTest, EarlyDispatchClaimStaysGatedAfterRelease) { + EXPECT_TRUE(PTO2SchedulerState::should_gate_early_dispatch(true, PTO2_EARLY_DISPATCH_DISPATCHED)); + EXPECT_TRUE(PTO2SchedulerState::should_gate_early_dispatch(false, PTO2_EARLY_DISPATCH_STAGING)); + EXPECT_FALSE(PTO2SchedulerState::should_gate_early_dispatch(false, PTO2_EARLY_DISPATCH_DISPATCHED)); +} + +TEST_F(WiringTest, EarlyDispatchLaunchHasSingleOwner) { + alignas(64) PTO2TaskPayload payload{}; + std::atomic start{false}; + bool won[2] = {false, false}; + + std::thread contenders[2]; + for (int i = 0; i < 2; i++) { + contenders[i] = std::thread([&, i] { + while (!start.load(std::memory_order_acquire)) {} + won[i] = PTO2SchedulerState::try_claim_early_dispatch_launch(payload); + }); + } + start.store(true, std::memory_order_release); + for (auto &contender : contenders) + contender.join(); + + EXPECT_NE(won[0], won[1]); + EXPECT_EQ(payload.early_dispatch_launch_state.load(), PTO2_EARLY_DISPATCH_LAUNCH_RINGING); +} + +TEST_F(WiringTest, EarlyDispatchFanoutWaitsForDoorbellPass) { + alignas(64) PTO2TaskSlotState producer, consumer; + init_slot(producer, PTO2_TASK_PENDING, 1, 1); + init_slot(consumer, PTO2_TASK_PENDING, 1, 1); + + producer.allow_early_resolve = true; + producer.payload->published_block_count.store(1, std::memory_order_relaxed); + consumer.payload->fanin_actual_count = 1; + + PTO2DepListEntry dep{}; + dep.slot_state = &consumer; + producer.fanout_head = &dep; + + producer.payload->early_dispatch_state.store(PTO2_EARLY_DISPATCH_STAGING, std::memory_order_relaxed); + sched.propagate_dispatch_fanin(producer); + EXPECT_EQ(consumer.payload->dispatch_fanin.load(), 0); + EXPECT_EQ(producer.payload->dispatch_propagated.load(), 0); + + producer.payload->early_dispatch_launch_state.store(PTO2_EARLY_DISPATCH_LAUNCH_RINGING); + producer.payload->early_dispatch_state.store(PTO2_EARLY_DISPATCH_DISPATCHED, std::memory_order_release); + sched.propagate_dispatch_fanin(producer); + EXPECT_EQ(consumer.payload->dispatch_fanin.load(), 0); + EXPECT_EQ(producer.payload->dispatch_propagated.load(), 0); + + producer.payload->early_dispatch_launch_state.store(PTO2_EARLY_DISPATCH_LAUNCH_COMPLETE, std::memory_order_seq_cst); + sched.propagate_dispatch_fanin(producer); + EXPECT_EQ(consumer.payload->dispatch_fanin.load(), 1); + EXPECT_EQ(producer.payload->dispatch_propagated.load(), 1); +} + +TEST_F(WiringTest, LateStagerRingsCapturedDoorbellAfterTableReuse) { + constexpr int core_id = 3; + constexpr uint64_t captured_addr = 0x12340000; + constexpr uint64_t reused_addr = 0x56780000; + constexpr uint32_t captured_token = 7; + constexpr uint32_t reused_token = 9; + + reset_test_reg_stub(); + sched.early_dispatch_doorbell_table[core_id].addr = reused_addr; + sched.early_dispatch_doorbell_table[core_id].token = reused_token; + + uint64_t claimed = 1ULL << core_id; + EXPECT_TRUE(PTO2SchedulerState::ring_claimed_local_doorbell(claimed, core_id, captured_addr, captured_token)); + EXPECT_EQ(get_test_reg_stub_base_addr(), captured_addr); + EXPECT_EQ(get_test_reg_stub_value(), (static_cast(captured_token) << 32) | captured_token); +} + +TEST_F(WiringTest, EarlyDispatchReleaseConsumesDoorbellMask) { + constexpr int core_id = 5; + constexpr uint64_t reg_addr = 0x98760000; + constexpr uint32_t token = 11; + alignas(64) PTO2TaskSlotState task; + init_slot(task, PTO2_TASK_PENDING, 1, 1); + task.allow_early_resolve = true; + task.next_block_idx.store(1, std::memory_order_relaxed); + task.payload->early_dispatch_state.store(PTO2_EARLY_DISPATCH_STAGING, std::memory_order_relaxed); + task.payload->staged_core_mask[0].store(1ULL << core_id, std::memory_order_relaxed); + sched.early_dispatch_doorbell_table[core_id].addr = reg_addr; + sched.early_dispatch_doorbell_table[core_id].token = token; + bool released_before = + task.payload->early_dispatch_state.load(std::memory_order_seq_cst) == PTO2_EARLY_DISPATCH_DISPATCHED; + + reset_test_reg_stub(); + EXPECT_FALSE(released_before); + EXPECT_TRUE(sched.try_early_dispatch_release(task)); + EXPECT_EQ(task.payload->early_dispatch_state.load(), PTO2_EARLY_DISPATCH_DISPATCHED); + EXPECT_EQ(task.payload->early_dispatch_launch_state.load(), PTO2_EARLY_DISPATCH_LAUNCH_COMPLETE); + EXPECT_EQ(task.payload->staged_core_mask[0].load(), 0); + EXPECT_EQ(get_test_reg_stub_base_addr(), reg_addr); + EXPECT_EQ(get_test_reg_stub_value(), (static_cast(token) << 32) | token); + EXPECT_EQ(task.payload->dispatch_propagated.load(), 0); + + sched.record_published_blocks(task, 1); + // The staging path must retry even though its earlier state snapshot was + // STAGING; release already missed this final publication count. + sched.propagate_dispatch_fanin(task); + EXPECT_EQ(task.payload->dispatch_propagated.load(), 1); + + sched.early_dispatch_doorbell_table[core_id].addr = 0x11110000; + sched.early_dispatch_doorbell_table[core_id].token = 12; + reset_test_reg_stub(); + EXPECT_TRUE(sched.try_early_dispatch_release(task)); + EXPECT_EQ(get_test_reg_stub_base_addr(), 0); + EXPECT_EQ(get_test_reg_stub_value(), 0); +} + TEST_F(WiringTest, EarlyDispatchBlockedByUnflaggedProducer) { alignas(64) PTO2TaskSlotState task_slot; alignas(64) PTO2TaskSlotState p_flagged, q_unflagged; diff --git a/tests/ut/cpp/stubs/test_stubs.cpp b/tests/ut/cpp/stubs/test_stubs.cpp index eea59ec6d9..f518ce81d5 100644 --- a/tests/ut/cpp/stubs/test_stubs.cpp +++ b/tests/ut/cpp/stubs/test_stubs.cpp @@ -91,12 +91,24 @@ void cache_flush_range(const void * /* addr */, size_t /* size */) {} // early-dispatch) is an inline that resolves a register id to its MMIO pointer // via get_reg_ptr and writes a 64-bit token through it. There is no MMIO on the // host UT runner; hand back writable static storage (8 bytes — the doorbell is a -// 64-bit store) so the inline links and any write is harmless. -volatile uint32_t *get_reg_ptr(uint64_t /* reg_base_addr */, RegId /* reg */) { - static volatile uint64_t dummy_reg = 0; - return reinterpret_cast(&dummy_reg); +// 64-bit store) and retain the requested base address for ownership tests. +static volatile uint64_t g_test_reg = 0; +static uint64_t g_test_reg_base_addr = 0; + +volatile uint32_t *get_reg_ptr(uint64_t reg_base_addr, RegId /* reg */) { + g_test_reg_base_addr = reg_base_addr; + return reinterpret_cast(&g_test_reg); +} + +void reset_test_reg_stub() { + g_test_reg = 0; + g_test_reg_base_addr = 0; } +uint64_t get_test_reg_stub_value() { return g_test_reg; } + +uint64_t get_test_reg_stub_base_addr() { return g_test_reg_base_addr; } + // ============================================================================= // runtime_maker.cpp stub (bind_callable_to_runtime_impl) // =============================================================================