fix(runtime): avoid global drain for local early sync-start

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)

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<uint64_t> 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

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;

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);
     }

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