host_build_graph: dedicated resolution thread (3S+1P) - #1544
Conversation
Split the AICPU threads into (aicpu_thread_num - 1) core-owning schedulers (S) plus one core-less resolution thread (P); enabled whenever there are at least two threads. Each S thread detects FIN on its own cores, frees the core immediately, and hands the finished task to P over a per-S SPSC queue. P runs on_task_complete for every finished task — publish completion_flags, drain the wake list, advance the watermark — and owns completed_tasks_ and the terminal completed_ flip, so the schedulers keep dispatching until P has resolved the whole graph and the host's wait_for_consumers never sees a stranded prefix. Because P is the sole producer of the ready queues, the enqueue CAS never contends and the wake-list drain is single-threaded. Async deferred completions and dependency-only retirement (dummy tasks and predicate-failed tasks, both routed to dummy_ready_queue) also move to P, so the scheduler loop stays purely core-local — poll own COND, dispatch own cores — with no per-iteration polling of the shared mailbox or dummy queue. Measured on a2a3, paged_attention batch=256 (65536 tasks) on one die: 29.5ms (4S) -> 25.4ms (3S+1P), -14%. The scheduler-thread completion phase roughly halves as the resolution contention leaves the hot path. The dependency-bound A/B (matmul/add, tile size held constant) confirms the direction splits with graph shape: a wide parallel graph speeds up, a single serial chain costs a few percent from the extra S->P hop.
|
Important Review skippedAuto incremental reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughThe scheduler adds a dedicated resolution thread for configurations with at least two AICPU threads. Scheduler threads enqueue completed tasks into SPSC queues, while the resolution thread processes task completions, deferred completions, and dummy-ready tasks. ChangesDedicated resolution thread
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant AICPUExecutor
participant SchedulerThreads
participant SPQueues
participant ResolutionThread
participant SchedulerContext
AICPUExecutor->>SchedulerThreads: run resolve_and_dispatch
SchedulerThreads->>SPQueues: enqueue completed slot
AICPUExecutor->>ResolutionThread: run run_resolution_thread
ResolutionThread->>SPQueues: dequeue completed slot
ResolutionThread->>SchedulerContext: call on_task_complete
ResolutionThread->>SchedulerContext: resolve deferred and dummy tasks
Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (3)
src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_cold_path.cpp (1)
1050-1063: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winCapacity is sized by
total_tasks, but the queue only needs to hold in-flight completions.At the 1<<20 ceiling this allocates 8 MB per scheduler thread (24 MB for 3S) even though the number of unresolved completions in flight is bounded by the ring's task window, not by the total task count. Sizing from
ring.task_window_size(rounded up, floored at 256) would keep the same no-spin property with far less memory.Also, when
total_tasks > (1u << 20)the capacity is clamped belowtotal_tasks, so thepush()backpressure spin can engage — that contradicts the "never fires in practice" claim in theCompletedTaskQueueheader comment. The spin is safe (P drains independently), but the comment should be softened.♻️ Suggested sizing
if (p_thread_mode_) { uint64_t cap = 256; - while (cap < static_cast<uint64_t>(total_tasks) && cap < (1u << 20)) { + // In-flight completions are bounded by the task window, not by the + // total task count. + const uint64_t needed = static_cast<uint64_t>(total_tasks); + while (cap < needed && cap < (1u << 20)) { cap <<= 1; }#!/bin/bash rg -nP --type=cpp -C3 '\btask_window_size' src/a2a3/runtime/host_build_graph | head -60🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_cold_path.cpp` around lines 1050 - 1063, Size each per-thread CompletedTaskQueue in the p_thread_mode_ initialization block from ring.task_window_size instead of total_tasks, rounding up to a power of two with a minimum capacity of 256 and no total-task-based ceiling. Update the CompletedTaskQueue header comment to avoid claiming push backpressure never occurs, while preserving that its spin remains safe because P drains independently.src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_context.h (1)
54-64: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider guarding
initagainst a non-power-of-two capacity and deleting copy/move.
mask = capacity_pow2 - 1silently corrupts the ring if a future caller passes a non-power-of-two. Also, the raw owningbufwith no rule-of-5 handling means an accidental copy ofSchedulerContextwould double-free.♻️ Suggested hardening
void init(uint64_t capacity_pow2) { + debug_assert(capacity_pow2 != 0 && (capacity_pow2 & (capacity_pow2 - 1)) == 0); cap = capacity_pow2;🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_context.h` around lines 54 - 64, Harden SchedulerContext::init by validating that capacity_pow2 is nonzero and a power of two before computing mask, using the existing project error-handling convention for invalid input. Make SchedulerContext non-copyable and non-movable (or otherwise implement complete ownership-safe rule-of-five behavior) so buf cannot be double-freed through accidental copies or moves; preserve destroy’s ownership cleanup.src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_dispatch.cpp (1)
986-991: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winP returns without the DFX/PMU teardown that
resolve_and_dispatchperforms.
resolve_and_dispatchends withlog_l2_swimlane_summary,l2_swimlane_aicpu_flush,l2_swimlane_aicpu_flush_sched_phase_buffer,dump_args_flushandpmu_aicpu_flush_buffers(Lines 1530-1554), and resetssched_l2_swimlane_[thread_idx]at entry. P does neither, so any per-thread DFX records written underthread_idx == p_thread_idx_(e.g. fromon_task_complete/poll_and_completein profiling builds) are dropped. Worth adding at least the swimlane/dump flushes for P.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_dispatch.cpp` around lines 986 - 991, Update P to perform the same per-thread DFX/PMU teardown as resolve_and_dispatch before returning completed_tasks_: call log_l2_swimlane_summary, l2_swimlane_aicpu_flush, l2_swimlane_aicpu_flush_sched_phase_buffer, dump_args_flush, and pmu_aicpu_flush_buffers for p_thread_idx_. Also reset sched_l2_swimlane_[thread_idx] on entry so profiling records are retained and flushed consistently.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In
`@src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_completion.cpp`:
- Around line 186-194: Update the p_thread_mode_ branch in the scheduler
completion path so sched_->tasks_completed is incremented when the finished task
is handed to P, matching the completed_tasks_ update performed by P. Use the
existing atomic fetch_add accounting mechanism and ensure this applies when
task_complete and t.running_done are observed without changing non-P completion
behavior.
In `@src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_dispatch.cpp`:
- Around line 969-986: Update P’s stall-detection block in the scheduler
dispatch flow to mirror the running-task guard used by resolve_and_dispatch:
only latch PTO2_ERROR_SCHEDULER_TIMEOUT and call emergency_shutdown when
self_owns_running_task(thread_idx) is true, or when no thread owns a running
task and the existing outstanding-work conditions hold. Otherwise refresh
last_progress_ts and continue spinning, preventing long-running or
startup-skewed tasks from triggering a false timeout.
---
Nitpick comments:
In `@src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_cold_path.cpp`:
- Around line 1050-1063: Size each per-thread CompletedTaskQueue in the
p_thread_mode_ initialization block from ring.task_window_size instead of
total_tasks, rounding up to a power of two with a minimum capacity of 256 and no
total-task-based ceiling. Update the CompletedTaskQueue header comment to avoid
claiming push backpressure never occurs, while preserving that its spin remains
safe because P drains independently.
In `@src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_context.h`:
- Around line 54-64: Harden SchedulerContext::init by validating that
capacity_pow2 is nonzero and a power of two before computing mask, using the
existing project error-handling convention for invalid input. Make
SchedulerContext non-copyable and non-movable (or otherwise implement complete
ownership-safe rule-of-five behavior) so buf cannot be double-freed through
accidental copies or moves; preserve destroy’s ownership cleanup.
In `@src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_dispatch.cpp`:
- Around line 986-991: Update P to perform the same per-thread DFX/PMU teardown
as resolve_and_dispatch before returning completed_tasks_: call
log_l2_swimlane_summary, l2_swimlane_aicpu_flush,
l2_swimlane_aicpu_flush_sched_phase_buffer, dump_args_flush, and
pmu_aicpu_flush_buffers for p_thread_idx_. Also reset
sched_l2_swimlane_[thread_idx] on entry so profiling records are retained and
flushed consistently.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: d0492dc7-e40b-495c-a248-5e1c41c361f6
📒 Files selected for processing (5)
src/a2a3/runtime/host_build_graph/aicpu/aicpu_executor.cppsrc/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_cold_path.cppsrc/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_completion.cppsrc/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_context.hsrc/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_dispatch.cpp
…eue sizing - P's idle timeout now only latches when work is outstanding AND no thread owns a running task (mirrors resolve_and_dispatch); a task that legitimately runs longer than the budget refreshes last_progress_ts instead of false-latching PTO2_ERROR_SCHEDULER_TIMEOUT. - P resets its per-thread swimlane state on entry and flushes swimlane / dump-args / PMU buffers on exit, so DFX records written under p_thread_idx_ are not dropped. - P mirrors completed_tasks_ into the SIMPLER_SCHED_PROFILING tasks_completed counter, which the scheduler threads' completed_this_turn no longer feeds. - Size each CompletedTaskQueue to min(total_tasks, ring task window) instead of total_tasks capped at 1<<20 — the window already bounds in-flight completions, so there is no ceiling that could clamp capacity below the working set and engage the backpressure spin. Assert power-of-two capacity; note the atomic cursors already make the queue non-copyable.
|
Thanks for the review. Addressed in Actionable
Nitpicks
Performance (this fixed commit
Direction is unchanged from the original write-up: scheduler-bound wide graphs speed up (−17.6% on pa here), compute-bound is neutral (no headroom), and a pure serial chain pays a few percent for the extra S→P hop that it cannot hide behind parallel work. All five pass ( 🤖 Generated with Claude Code |
…back branches The dedicated resolution (P) thread is not optional: like tmr's scheduler + orchestrator split, host_build_graph needs at least two AICPU threads (one core-owning scheduler + one resolution thread). A single thread cannot both own cores and resolve on a dedicated thread, and no config runs fewer than two threads (scene tests use 3 or 4; the default is 3). Require aicpu_thread_num >= 2 in assign_cores_to_threads and remove the `p_thread_mode_` flag and every branch on it. This deletes the retained inline completion path from complete_slot_task and the per-iteration async-poll and dummy/predicate-drain blocks from resolve_and_dispatch — those were dead in the only supported configuration and only added predicted branches to the scheduler's hot loop. Async, dummy and predicate resolution live solely in run_resolution_thread now; the scheduler loop goes straight from completion detection to dispatch. Behavior is unchanged for the >= 2 thread configs (all of them). matmul, bgemm, paged_attention, vector_example and predicated_dispatch (the dummy / predicate path, now P-only) all pass.
Follow-up: rigorous 100-round re-measurement (corrects the single-run numbers above)The per-workload percentages I posted earlier were single-run and, for the Method — same physical card (a2a3 silicon, device 6,
The earlier small-case results were measurement artifacts. At 1–2 ms the same Why pa is the only real win, and it is not diluted: the win requires the Net: 3S+1P is a scenario-specific optimization — a solid, reproducible |
1 participant
What
Split the
host_build_graphAICPU threads into 3 core-owning schedulers (S) + 1 core-less resolution thread (P) (generally:aicpu_thread_num - 1S threads + 1 P, enabled at ≥2 threads).on_task_completefor every finished task — publishcompletion_flags, drain the wake list, advancecompleted_watermark. P ownscompleted_tasks_and the terminalcompleted_flip, so the S threads keep dispatching until P has resolved the whole graph (the host'swait_for_consumersnever sees a stranded prefix).Because P is the sole producer of the ready queues, the enqueue CAS never contends and the wake-list drain is single-threaded. Async deferred completions and dependency-only retirement (dummy + predicate-failed tasks) also move to P, so the S loop stays purely core-local — poll own COND, dispatch own cores — with no per-iteration polling of the shared mailbox / dummy queue.
Scope: a2a3
host_build_graphonly.Why
Under the polling-completion model (#1435), the 4 symmetric scheduler threads each do completion and resolution and dispatch, so the wake-list CAS and the MPMC
ready_queues[]enqueue contend across all 4. Concentrating resolution on one thread removes that contention and keeps the dispatch threads' hot loop core-local.Performance (a2a3, single die, pinned device, L2-swimlane device-execution span)
Same binary A/B (4S = this change reverted, 3S+1P = this change):
paged_attentionCase1 (65536 tasks)deep_chain(512-deep add-scalar)Reading:
pa−14%): the scheduler-thread completion phase roughly halves as resolution contention leaves the hot path (3.55M→1.74M cyc).The dep-bound A/B holds kernels and tile size constant and varies only the dependency structure, isolating the direction split to the resolution path.
Correctness:
matmul,bgemm,paged_attention(Case1 + Case2 SPMD),vector_example,predicated_dispatch, and the 2-layer qwen3 decode graph all pass on 3S+1P.Notes / open questions
on_task_completeis internally CAS-safe).🤖 Generated with Claude Code