host_build_graph: parallelize the boot initial-classify across AICPU threads

[ChaoZheng109]

What

Parallelize the host_build_graph device boot initial-classify across the AICPU threads.

The boot classify seeds the whole graph's ready queues and wake lists — route roots (all fanin already met) to the ready queues, register every other task on its first unmet producer's wake list. It ran O(total_tasks) serially on the boot leader while the other AICPU threads idle-waited on runtime_init_ready_.

This PR splits it:

• on_orchestration_done keeps only the leader-only setup (attach, task-count latch, SPSC queue allocation).
• 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 classifies its slice and increments the classify_arrived_ barrier; the leader publishes runtime_init_ready_ only once all slices finish — 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 at run time, 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, only the boot barrier (modeled on the existing handshake barrier).

Scope: a2a3 host_build_graph only. Follow-up to #1544.

Why

On a 65792-task paged_attention graph the serial classify is ~2.26 ms of one thread working while three spin — pure serial boot latency on the critical path before any task dispatches. The other threads are right there, idle-waiting, and the classify is embarrassingly parallel (each task is independent).

Performance (a2a3, paged_attention Case1, 65792 tasks, 4 AICPU threads)

Classify-phase wall (boot leader, from setup-done to runtime_init_ready_):

	classify phase
serial (before)	2.26 ms
parallel (this PR)	1.58 ms
Δ	−30% (~0.68 ms returned to the run)

It does not reach the 4× of perfect scaling: 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, and the barrier waits on the slowest slice. The gain is workload-dependent (largest on scheduler-bound graphs with many small tasks) — ~0.68 ms is ~2.7% of this run's ~25 ms device execution, and negligible for compute-bound long runs. Still a strict improvement: otherwise-idle serial boot time handed back, off the critical-path start.

Correctness

matmul, bgemm, paged_attention (Case1), vector_example and predicated_dispatch all pass on a2a3 — the barrier introduces no deadlock and the parallel-seeded ready-set / wake-list graph is correct.

🤖 Generated with Claude Code

[coderabbitai]

📝 Walkthrough

Walkthrough

AICPU startup now partitions initial task classification across threads. Each thread seeds ready queues or wake lists for its task range, and runtime initialization becomes ready only after all classification partitions complete.

Changes

Parallel initial classification

Layer / File(s)	Summary
Partitioned scheduler classification src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_context.h, src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_cold_path.cpp	Adds classify_partition, which classifies disjoint submitted-task ranges and routes tasks to ready queues or wake lists.
AICPU classification barrier src/a2a3/runtime/host_build_graph/aicpu/aicpu_executor.cpp	Coordinates leader setup, per-thread classification, arrival synchronization, delayed runtime readiness, and barrier reset during deinitialization.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Possibly related PRs

• hw-native-sys/simpler#1176: Changes synchronization at the same orchestration-completion hook.
• hw-native-sys/simpler#1435: Introduces the earlier leader-only initial classification path refactored here.
• hw-native-sys/simpler#1518: Also modifies runtime_init_ready_ publication in AicpuExecutor::run.

Sequence Diagram(s)

sequenceDiagram
  participant BootLeader
  participant AICPUThreads
  participant SchedulerContext
  BootLeader->>AICPUThreads: Publish classify_ready_
  AICPUThreads->>SchedulerContext: Classify each thread partition
  SchedulerContext-->>AICPUThreads: Seed ready queues or wake lists
  AICPUThreads->>AICPUThreads: Wait for all classify arrivals
  AICPUThreads-->>BootLeader: Publish runtime_init_ready_

Loading

Poem

A rabbit watched the threads divide,
While ready queues grew side by side.
Each wake list found its proper place,
Then readiness sprang from the race.
“Hop!” said Bun, “the graph’s in stride!”

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly states the main change: parallelizing boot initial-classify across AICPU threads.
Description check	✅ Passed	The description is directly aligned with the changeset and explains the parallel boot-classify behavior, rationale, and validation.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 1

🤖 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/aicpu/aicpu_executor.cpp`:
- Around line 320-329: Update the worker startup/classification barrier around
classify_arrived_ and runtime_init_ready_ to use a shared abort state when an
out-of-range worker exits before arrival. Ensure every boot wait checks this
state, and change the early invalid-affinity return to signal the abort and
follow normal shutdown so valid workers cannot spin indefinitely.

🪄 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: 3afc291e-b710-4e58-98d3-fdbbf301df33

📥 Commits

Reviewing files that changed from the base of the PR and between 453ff71 and 9f02459.

📒 Files selected for processing (3)

• src/a2a3/runtime/host_build_graph/aicpu/aicpu_executor.cpp
• src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_cold_path.cpp
• src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_context.h