From 963582701855e9fc2dc46219a788275371de9e03 Mon Sep 17 00:00:00 2001 From: Chao Wang <26245345+ChaoWao@users.noreply.github.com> Date: Mon, 20 Jul 2026 01:58:17 -0700 Subject: [PATCH] Fix: preserve early dispatch for late-wired consumers Serialize the dispatch-propagation claim with its fanout snapshot so wiring can distinguish included edges from late edges. Seed late fully-published producers during wiring, enqueue whichever contribution completes dispatch fanin, and cover both serialized and concurrent orderings. --- .../docs/RUNTIME_LOGIC.md | 16 +++-- .../runtime/pto_orchestrator.cpp | 22 ++++-- .../runtime/scheduler/pto_scheduler.h | 64 ++++++++++------- .../runtime/scheduler/scheduler_dispatch.cpp | 6 +- tests/ut/cpp/a2a3/test_wiring.cpp | 71 +++++++++++++++++++ 5 files changed, 142 insertions(+), 37 deletions(-) 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 3af12bfc36..7b49a7e1d4 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md @@ -668,11 +668,11 @@ block is not mistaken for a normal task). #### Early-candidate gate: producer must publish every block (deadlock avoidance) -`propagate_dispatch_fanin` (the EARLY-source candidate trigger) no-ops until the producer is -**fully published**: `published_block_count == logical_block_num`. Normal dispatch, regular -early staging, and the sync drain increment this counter only after the claimed range's payloads -and MMIO dispatch tokens are visible. A staged producer also waits for release and completion of -its owned doorbell pass before exposing fanout. +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. 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 @@ -683,6 +683,12 @@ reserved core slot and a launch-visible payload before a consumer can pre-occupy `next_block_idx` records reservation progress; `published_block_count` independently establishes publication and early-candidate readiness. +The producer's `dispatch_propagated` claim and fanout snapshot are serialized under +`fanout_lock` with consumer wiring. An edge already in the snapshot is counted by scheduler +propagation; wiring seeds an edge added after the claim and enqueues the consumer if that seed +completes `dispatch_fanin`. This gives each eligible producer-consumer edge exactly one +early-candidate contribution regardless of which side acquires the lock first. + #### MIX per-core placement A MIX task spans a cluster (1 AIC + 2 AIV). `classify_mix_cluster` admits a cluster whenever diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.cpp b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.cpp index 0fb1f76d28..5aa82dadfc 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.cpp +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.cpp @@ -390,6 +390,7 @@ void PTO2OrchestratorState::wire_fanin_task(PTO2TaskSlotState &slot_state, int32 slot_state.fanin_count = wfanin + 1; int32_t early_finished = 0; + int32_t early_propagated = 0; bool early_disqualified = false; for_each_fanin_slot_state(*payload, [&](PTO2TaskSlotState *producer) { producer->lock_fanout(); @@ -401,15 +402,26 @@ void PTO2OrchestratorState::wire_fanin_task(PTO2TaskSlotState &slot_state, int32 early_finished++; } else { producer->fanout_head = rss.dep_pool.prepend(producer->fanout_head, &slot_state); + // The marker shares fanout_lock with propagation's snapshot. A set + // marker means this edge is outside that snapshot and needs a seed. + if (producer->payload->dispatch_propagated.load(std::memory_order_acquire) != 0) { + early_propagated++; + } } producer->unlock_fanout(); }); - // Pre-completed producers will not dispatch again. Seed dispatch_fanin only - // when every producer is codegen-flagged; one unflagged producer makes the - // direct-only early-dispatch candidate count unreachable by design. - if (!early_disqualified && early_finished != 0) { - payload->dispatch_fanin.fetch_add(early_finished, std::memory_order_acq_rel); + // Completed producers and edges outside a producer's one-shot propagation + // snapshot both need wiring-time contributions. Seed only when every + // producer is codegen-flagged; one unflagged producer disqualifies the task. + int32_t early_seed = early_finished + early_propagated; + if (!early_disqualified && early_seed != 0) { + int32_t dispatch_fanin = payload->dispatch_fanin.fetch_add(early_seed, std::memory_order_acq_rel) + early_seed; + // A fully pre-completed fanin routes normally. If any producer was live, + // the exact-full increment must enqueue the early candidate. + if (early_finished != payload->fanin_actual_count && dispatch_fanin == payload->fanin_actual_count) { + sched->try_enqueue_early_dispatch_candidate(slot_state); + } } int32_t init_rc = early_finished + 1; 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 93bddf21ca..1d02c32aac 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 @@ -864,6 +864,34 @@ struct PTO2SchedulerState { return true; } + // Attempt to claim and enqueue a consumer after every direct producer becomes + // launch-visible. Both propagation and wiring can complete dispatch_fanin. + inline void try_enqueue_early_dispatch_candidate(PTO2TaskSlotState &consumer) { + // Predicated tasks resolve at the ready point, never early-dispatch. The + // wiring caller has no separate predicate filter, so the gate lives here. + if (consumer.active_mask.has_predicate()) return; + PTO2ResourceShape shape = consumer.active_mask.to_shape(); + if (shape != PTO2ResourceShape::AIC && shape != PTO2ResourceShape::AIV && shape != PTO2ResourceShape::MIX) { + return; + } + + uint8_t expected = PTO2_EARLY_DISPATCH_NONE; + if (!consumer.payload->early_dispatch_state.compare_exchange_strong( + expected, PTO2_EARLY_DISPATCH_STAGING, std::memory_order_seq_cst, std::memory_order_seq_cst + )) { + return; + } + + 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. + if (consumer.active_mask.requires_sync_start()) { + early_sync_start_queue.push_tagged(&consumer, task_id); + } else { + early_dispatch_queues[static_cast(shape)].push_tagged(&consumer, task_id); + } + } + // 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 @@ -892,10 +920,18 @@ struct PTO2SchedulerState { } if (was_pre_staged && launch_state != PTO2_EARLY_DISPATCH_LAUNCH_COMPLETE) return; } - if (p.payload->dispatch_propagated.exchange(1, std::memory_order_acq_rel) != 0) - return; // already propagated once + // The once-claim and fanout snapshot share fanout_lock with wiring, so a + // set marker tells wiring that its new edge is outside this snapshot. + // The unlocked precheck avoids fanout-lock traffic after propagation. + if (p.payload->dispatch_propagated.load(std::memory_order_acquire) != 0) return; p.lock_fanout(); - PTO2DepListEntry *edge = p.fanout_head; // snapshot head, walk lock-free (fanout stable by dispatch) + if (p.payload->dispatch_propagated.exchange(1, std::memory_order_acq_rel) != 0) { + p.unlock_fanout(); + return; + } + // Nodes reachable from the captured head are immutable; wiring serializes + // later prepends under fanout_lock and seeds them separately. + PTO2DepListEntry *edge = p.fanout_head; p.unlock_fanout(); for (; edge != nullptr; edge = edge->next) { PTO2TaskSlotState *c = edge->slot_state; @@ -909,27 +945,7 @@ struct PTO2SchedulerState { // seed short and never bumps, so this stays unreachable for that consumer. int32_t nf = c->payload->dispatch_fanin.fetch_add(1, std::memory_order_acq_rel) + 1; if (nf != c->payload->fanin_actual_count) continue; - PTO2ResourceShape shape = c->active_mask.to_shape(); - if (shape != PTO2ResourceShape::AIC && shape != PTO2ResourceShape::AIV && shape != PTO2ResourceShape::MIX) - continue; - // sync_start blocks launch as an atomic cohort. They ride the early-dispatch path - // too, but through a dedicated shape-agnostic queue (early_sync_start_queue), - // drained as the highest tier in try_early_dispatch via the gated drain + - // running-slot rendezvous (enter_drain_mode pre-stages every block gated — MIX - // with per-core split placement — and the rendezvous rings them together once - // every gated core occupies a running slot and the producer released). They must - // NOT enter early_dispatch_queues[shape]: that path's partial range-claim would - // strand gated cohort blocks nobody rings, so the rendezvous never reaches - // block_num. The rendezvous counts CORES (a MIX block spans a cluster), so one - // shape-agnostic queue suffices. - uint8_t expect = PTO2_EARLY_DISPATCH_NONE; // exactly-once: only the CAS winner enqueues - if (!c->payload->early_dispatch_state.compare_exchange_strong( - expect, PTO2_EARLY_DISPATCH_STAGING, std::memory_order_seq_cst, std::memory_order_seq_cst - )) - continue; - uint64_t task_id = static_cast(c->task->task_id.raw); - if (c->active_mask.requires_sync_start()) early_sync_start_queue.push_tagged(c, task_id); - else early_dispatch_queues[static_cast(shape)].push_tagged(c, task_id); + try_enqueue_early_dispatch_candidate(*c); } } 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 d933d86fb9..4ae92944d0 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 @@ -700,9 +700,9 @@ int32_t SchedulerContext::stage_consumer_blocks( // Early-dispatch analog of dispatch_shape: drain early_dispatch_queues[shape] and // pre-stage claimed block ranges onto this thread's `shape` cores for `phase`. IDLE // stages onto idle cores (RUNNING slot, gated); PENDING stages onto a running core's -// gated pending slot. Candidates are pushed to the shape's queue EVENT-DRIVEN by -// propagate_dispatch_fanin, so the shape is the queue index (no per-consumer -// to_shape()). Returns the number of blocks staged. +// gated pending slot. Producer propagation and late wiring push candidates to the +// shape's queue when dispatch_fanin becomes complete, so the shape is the queue +// index (no per-consumer to_shape()). Returns the number of blocks staged. int32_t SchedulerContext::early_dispatch_shape(int32_t thread_idx, PTO2ResourceShape shape, CoreTracker::DispatchPhase phase) { CoreTracker &tracker = core_trackers_[thread_idx]; diff --git a/tests/ut/cpp/a2a3/test_wiring.cpp b/tests/ut/cpp/a2a3/test_wiring.cpp index 5df05d849e..a085e2bdf9 100644 --- a/tests/ut/cpp/a2a3/test_wiring.cpp +++ b/tests/ut/cpp/a2a3/test_wiring.cpp @@ -386,6 +386,77 @@ TEST_F(WiringTest, EarlyDispatchWaitsForAllProducerBlocksPublished) { EXPECT_EQ(payload.dispatch_fanin.load(), payload.fanin_actual_count); } +TEST_F(WiringTest, LateWiredFullyPublishedProducerStillSeedsEarlyDispatch) { + alignas(64) PTO2TaskSlotState producer, consumer; + alignas(64) PTO2TaskPayload consumer_payload; + memset(&consumer_payload, 0, sizeof(consumer_payload)); + PTO2TaskDescriptor consumer_desc{}; + + init_slot(producer, PTO2_TASK_PENDING, 1, 1); + producer.allow_early_resolve = true; + sched.record_published_blocks(producer, producer.logical_block_num); + sched.propagate_dispatch_fanin(producer); + ASSERT_EQ(producer.payload->dispatch_propagated.load(), 1); + + init_slot(consumer, PTO2_TASK_PENDING, 0, 1); + consumer_payload.fanin_actual_count = 1; + consumer_payload.fanin_inline_slot_states[0] = &producer; + consumer.payload = &consumer_payload; + consumer.task = &consumer_desc; + + wire_fanin(consumer, 1); + + EXPECT_EQ(consumer_payload.dispatch_fanin.load(), consumer_payload.fanin_actual_count); + EXPECT_EQ(consumer_payload.early_dispatch_state.load(), PTO2_EARLY_DISPATCH_STAGING); + auto shape = static_cast(consumer.active_mask.to_shape()); + EXPECT_EQ(sched.early_dispatch_queues[shape].pop(), &consumer); + EXPECT_NE(producer.fanout_head, nullptr); +} + +TEST_F(WiringTest, WiringSeedEnqueuesAfterConcurrentPropagation) { + alignas(64) PTO2TaskSlotState producers[3], consumer; + alignas(64) PTO2TaskPayload consumer_payload; + memset(&consumer_payload, 0, sizeof(consumer_payload)); + PTO2TaskDescriptor consumer_desc{}; + + init_slot(producers[0], PTO2_TASK_COMPLETED, 1, 1); + init_slot(producers[1], PTO2_TASK_PENDING, 1, 1); + init_slot(producers[2], PTO2_TASK_COMPLETED, 1, 1); + for (auto &producer : producers) + producer.allow_early_resolve = true; + sched.record_published_blocks(producers[1], producers[1].logical_block_num); + + init_slot(consumer, PTO2_TASK_PENDING, 0, 1); + consumer_payload.fanin_actual_count = 3; + for (int i = 0; i < 3; i++) + consumer_payload.fanin_inline_slot_states[i] = &producers[i]; + consumer.payload = &consumer_payload; + consumer.task = &consumer_desc; + + producers[2].lock_fanout(); + std::thread wiring([&]() { + wire_fanin(consumer, 3); + }); + + bool live_edge_wired = false; + while (!live_edge_wired) { + producers[1].lock_fanout(); + live_edge_wired = producers[1].fanout_head != nullptr; + producers[1].unlock_fanout(); + if (!live_edge_wired) std::this_thread::yield(); + } + + sched.propagate_dispatch_fanin(producers[1]); + EXPECT_EQ(consumer_payload.dispatch_fanin.load(), 1); + producers[2].unlock_fanout(); + wiring.join(); + + EXPECT_EQ(consumer_payload.dispatch_fanin.load(), consumer_payload.fanin_actual_count); + EXPECT_EQ(consumer_payload.early_dispatch_state.load(), PTO2_EARLY_DISPATCH_STAGING); + auto shape = static_cast(consumer.active_mask.to_shape()); + EXPECT_EQ(sched.early_dispatch_queues[shape].pop(), &consumer); +} + TEST_F(WiringTest, ConcurrentBlockRangeClaimsDoNotOverlap) { alignas(64) PTO2TaskSlotState task_slot; init_slot(task_slot, PTO2_TASK_PENDING, 1, 1);