Fix: preserve early dispatch for late-wired consumers

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

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;

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<uint64_t>(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<int32_t>(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<uint64_t>(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<int32_t>(shape)].push_tagged(c, task_id);
+            try_enqueue_early_dispatch_candidate(*c);
         }
     }
 

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];

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<int32_t>(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<int32_t>(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);