host_build_graph: parallelize the boot initial-classify across AICPU threads

src/a2a3/runtime/host_build_graph/aicpu/aicpu_executor.cpp

@@ -94,6 +94,14 @@ struct AicpuExecutor {
     std::atomic<int32_t> hs_arrived_{0};
     std::atomic<int32_t> hs_thread_seq_{0};
 
+    // Parallel-boot-classify coordination (see AicpuExecutor::run). classify_ready_
+    // is published by the boot leader once its leader-only orchestration setup is
+    // visible; classify_arrived_ is the barrier counting threads that finished
+    // their slice of the initial classify. Both are one-shot per run and reset in
+    // deinit().
+    std::atomic<bool> classify_ready_{false};
+    std::atomic<int32_t> classify_arrived_{0};
+
     int32_t aicpu_thread_num_{0};
 
     // ===== Task queue state (managed by scheduler ready queues) =====
@@ -291,16 +299,38 @@ int32_t AicpuExecutor::run(Runtime *runtime) {
             LOG_INFO_V0("Thread %d: host-orch boot complete (%d tasks)", thread_idx, runtime->host_total_tasks);
         }
 
-        runtime_init_ready_.store(true, std::memory_order_release);
+        // Publish "leader setup done" (SM attached, task count latched, queues
+        // allocated). Every thread then classifies its slice below before any of
+        // them may dispatch — the leader holds runtime_init_ready_ until then.
+        classify_ready_.store(true, std::memory_order_release);
     }
 
-    // Every AICPU thread schedules its assigned cores; the boot thread above
-    // falls through to here after publishing runtime_init_ready_.
-    if (!sched_ctx_.is_completed()) {
-        // Wait for the boot thread to attach the SM header and publish the task count.
+    // Parallel initial classify. Every AICPU thread waits for the leader's
+    // orchestration setup, seeds its disjoint slice of the whole graph's ready
+    // set + wake lists, then barriers. Only once all slices are done does the
+    // leader publish runtime_init_ready_, so no thread dispatches against a
+    // half-seeded graph. This replaces the O(total_tasks) serial classify the
+    // leader used to run alone while the others idle-waited.
+    while (!classify_ready_.load(std::memory_order_acquire)) {
+        SPIN_WAIT_HINT();
+    }
+    if (!sched_ctx_.is_completed() && rt != nullptr) {
+        sched_ctx_.classify_partition(thread_idx, aicpu_thread_num_);
+    }
+    classify_arrived_.fetch_add(1, std::memory_order_acq_rel);
+    if (thread_idx == aicpu_thread_num_ - 1) {
+        while (classify_arrived_.load(std::memory_order_acquire) < aicpu_thread_num_) {
+            SPIN_WAIT_HINT();
+        }
+        runtime_init_ready_.store(true, std::memory_order_release);
+    } else {
         while (!runtime_init_ready_.load(std::memory_order_acquire)) {
             SPIN_WAIT_HINT();
         }
+    }
+
+    // Every AICPU thread schedules its assigned cores.
+    if (!sched_ctx_.is_completed()) {
         if (rt == nullptr) {
             LOG_ERROR("Thread %d: rt is null after orchestrator error, skipping dispatch", thread_idx);
         } else {
@@ -369,6 +399,8 @@ void AicpuExecutor::deinit(Runtime *runtime) {
     hs_setup_done_.store(false, std::memory_order_release);
     hs_arrived_.store(0, std::memory_order_release);
     hs_thread_seq_.store(0, std::memory_order_release);
+    classify_ready_.store(false, std::memory_order_release);
+    classify_arrived_.store(0, std::memory_order_release);
     thread_idx_.store(0, std::memory_order_release);
     finished_.store(false, std::memory_order_release);
 

src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_cold_path.cpp

@@ -1098,33 +1098,11 @@ void SchedulerContext::on_orchestration_done(
         }
     }
 
-    // Polling initial classify (device boot): the host built the whole graph and
-    // no producer has executed yet — every completion_flags byte is 0 except the
-    // hidden-alloc tasks the host completed inline (pre-set to 1). Classify each
-    // submitted task exactly once: route roots (all fanin met) to the ready queues
-    // and register the rest on their first unmet producer's wake list. This
-    // replaces the host-side wiring drain the wiring model deferred to a device
-    // queue. Runs on the boot leader BEFORE the caller publishes
-    // runtime_init_ready_ (release), so it is race-free against the scheduler
-    // threads (they acquire that store before dispatching) and nothing completes
-    // during the scan.
-    if (orch_err == PTO2_ERROR_NONE && sched_->ring_sched_state.ring != nullptr) {
-        PTO2SharedMemoryRingHeader &ring = *sched_->ring_sched_state.ring;
-        const int32_t submitted = ring.fc.current_task_index.load(std::memory_order_acquire);
-        for (int32_t id = 0; id < submitted; id++) {
-            if (ring.is_completion_flag_set(id)) {
-                continue;  // completed on the host (hidden alloc); nothing to dispatch
-            }
-            PTO2TaskSlotState &s = ring.get_slot_state_by_task_id(id);
-            int32_t state = sched_->classify_fanin_state(&s);
-            if (state < 0) {
-                sched_->push_ready_routed(&s);
-            } else {
-                int32_t prod_local = s.payload->fanin_local_ids[state];
-                sched_->register_wake(&ring.get_slot_state_by_task_id(prod_local), &s);
-            }
-        }
-    }
+    // The polling initial classify (seed the ready queues + wake lists for the
+    // whole graph) runs AFTER this, partitioned across all AICPU threads in
+    // classify_partition() — see AicpuExecutor::run. It is kept out of this
+    // leader-only setup so the O(total_tasks) scan is not serial on one thread
+    // while the others idle-wait for runtime_init_ready_.
 
 #if SIMPLER_DFX
     // Write the core-to-thread mapping so the profiling data reflects the
@@ -1139,3 +1117,43 @@ void SchedulerContext::on_orchestration_done(
     }
 #endif
 }
+
+// Polling initial classify (device boot), partitioned across all AICPU threads.
+// The host built the whole graph and no producer has executed yet — every
+// completion_flags byte is 0 except the hidden-alloc tasks the host completed
+// inline (pre-set to 1). Each thread classifies its contiguous slice of the
+// submitted-task range exactly once: route roots (all fanin met) to the ready
+// queues and register the rest on their first unmet producer's wake list.
+//
+// This is the same work the wiring model deferred to a device queue, now run
+// N-way parallel. push_ready_routed (MPMC ready queues) and register_wake
+// (lock-free wake-list CAS) are the same concurrency-safe primitives the
+// scheduler threads use during the run, and at boot no producer has completed
+// (wake_list heads are nullptr, never SENTINEL), so registration never
+// re-classifies. The caller barriers all threads here BEFORE any of them
+// publishes runtime_init_ready_, so the whole ready-set / wake-list graph is
+// fully seeded before the first dispatch.
+void SchedulerContext::classify_partition(int32_t thread_idx, int32_t nthreads) {
+    if (completed_.load(std::memory_order_acquire) || sched_->ring_sched_state.ring == nullptr) {
+        return;
+    }
+    PTO2SharedMemoryRingHeader &ring = *sched_->ring_sched_state.ring;
+    const int32_t submitted = ring.fc.current_task_index.load(std::memory_order_acquire);
+    // Disjoint contiguous slices covering [0, submitted): thread t owns
+    // [submitted*t/nthreads, submitted*(t+1)/nthreads). int64 math avoids overflow.
+    const int32_t lo = static_cast<int32_t>((static_cast<int64_t>(submitted) * thread_idx) / nthreads);
+    const int32_t hi = static_cast<int32_t>((static_cast<int64_t>(submitted) * (thread_idx + 1)) / nthreads);
+    for (int32_t id = lo; id < hi; id++) {
+        if (ring.is_completion_flag_set(id)) {
+            continue;  // completed on the host (hidden alloc); nothing to dispatch
+        }
+        PTO2TaskSlotState &s = ring.get_slot_state_by_task_id(id);
+        int32_t state = sched_->classify_fanin_state(&s);
+        if (state < 0) {
+            sched_->push_ready_routed(&s);
+        } else {
+            int32_t prod_local = s.payload->fanin_local_ids[state];
+            sched_->register_wake(&ring.get_slot_state_by_task_id(prod_local), &s);
+        }
+    }
+}

src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_context.h

@@ -160,6 +160,13 @@ class SchedulerContext {
     // step belongs to the orchestrator lifecycle, not the scheduler.
     void on_orchestration_done(Runtime *runtime, PTO2Runtime *rt, int32_t thread_idx, int32_t total_tasks);
 
+    // Seed the ready queues + wake lists for the whole graph at boot. Called by
+    // every AICPU thread on a disjoint slice of the submitted-task range, after
+    // on_orchestration_done and before runtime_init_ready_ (the caller barriers
+    // all threads between the two). Concurrency-safe: push_ready_routed and
+    // register_wake are the same lock-free primitives used during the run.
+    void classify_partition(int32_t thread_idx, int32_t nthreads);
+
     // Bind the PTO2Runtime scheduler pointer.
     void bind_runtime(PTO2Runtime *rt);