From 9f024590a1b253e8138c8af2f4071c3c8e969a90 Mon Sep 17 00:00:00 2001 From: ChaoZheng109 Date: Tue, 28 Jul 2026 18:40:39 -0700 Subject: [PATCH] host_build_graph: parallelize the boot initial-classify across AICPU threads MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The device boot classify — seed the whole graph's ready queues and wake lists (route roots to the ready queues, register every other task on its first unmet producer's wake list) — used to run O(total_tasks) serially on the boot leader while the other AICPU threads idle-waited on runtime_init_ready_. On a 65792-task paged_attention graph that is ~2.26 ms of one thread working and three spinning. Split it: on_orchestration_done keeps only the leader-only setup (attach, task count, queue allocation) and the classify moves to classify_partition(), which each thread runs over a disjoint contiguous slice of the submitted-task range. The boot leader publishes classify_ready_ once its setup is visible; every thread then classifies its slice and increments the classify_arrived_ barrier; the leader publishes runtime_init_ready_ only once all slices are done, so no thread dispatches against a half-seeded graph. push_ready_routed and register_wake are the same lock-free primitives the scheduler threads already use concurrently, and at boot no producer has completed (wake-list heads are nullptr, never SENTINEL), so registration never re-classifies — the parallel scan needs no new locking. Measured on a2a3, paged_attention batch=256 (65792 tasks), 4 AICPU threads: the classify phase drops from 2.26 ms (serial) to 1.58 ms (-30%). It does not reach the 4x of perfect scaling because concurrent register_wake on a high-fanout producer's wake list and push_ready_routed on the shared ready queues re-introduce the CAS contention the run-time 3S+1P split removes, plus the barrier waits on the slowest slice — but it still returns ~0.68 ms of otherwise-idle serial boot time to the run. matmul, bgemm, paged_attention, vector_example and predicated_dispatch all pass (the barrier introduces no deadlock and the seeded graph is correct). --- .../host_build_graph/aicpu/aicpu_executor.cpp | 42 +++++++++-- .../runtime/scheduler/scheduler_cold_path.cpp | 72 ++++++++++++------- .../runtime/scheduler/scheduler_context.h | 7 ++ 3 files changed, 89 insertions(+), 32 deletions(-) diff --git a/src/a2a3/runtime/host_build_graph/aicpu/aicpu_executor.cpp b/src/a2a3/runtime/host_build_graph/aicpu/aicpu_executor.cpp index 4275f0080..e9ce1a006 100644 --- a/src/a2a3/runtime/host_build_graph/aicpu/aicpu_executor.cpp +++ b/src/a2a3/runtime/host_build_graph/aicpu/aicpu_executor.cpp @@ -94,6 +94,14 @@ struct AicpuExecutor { std::atomic hs_arrived_{0}; std::atomic 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 classify_ready_{false}; + std::atomic 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); diff --git a/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_cold_path.cpp b/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_cold_path.cpp index 52ab48b28..5c20442d7 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_cold_path.cpp +++ b/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((static_cast(submitted) * thread_idx) / nthreads); + const int32_t hi = static_cast((static_cast(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); + } + } +} diff --git a/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_context.h b/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_context.h index 51e297e91..3aa699a44 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_context.h +++ b/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);