[Draft] [Performance] Decouple Orchestrator from handshaking#1345
[Draft] [Performance] Decouple Orchestrator from handshaking#1345SergioMartin86 wants to merge 2 commits into
Conversation
Collapse the AICore init handshake from a three-phase ping-pong
(aicore_regs_ready -> aicpu_regs_ready -> aicore_done) to a single round
trip. AICore now publishes {physical_core_id, core_type, aicore_done} in
one write, then polls its own DATA_MAIN_BASE SPR. The AICPU opens the
core's register window (platform_init_aicore_regs writes
DATA_MAIN_BASE=IDLE) only after observing aicore_done, so that
window-open write doubles as the release acknowledgement -- the dedicated
aicpu_regs_ready / aicore_regs_ready fields are removed from struct
Handshake.
a2a3 (already parallel via hw-native-sys#1279): within each thread's core slice,
sweep-poll on aicore_done -- a GM read, not the strictly-serial nGnRE
COND poll -- so per-core wakeups overlap (~max, not sum), and batch the
release barrier (one OUT_OF_ORDER_STORE_BARRIER per slice instead of one
per core).
a5: port the parallel-handshake split (pre_handshake_init /
handshake_partition / post_handshake_init) alongside the single-round-trip
protocol; a5 previously handshaked serially on one thread.
Reproduces the approach of upstream PR hw-native-sys#1214 on a fresh branch off main.
Verified a2a3 onboard (device 3): paged_attention Case1 200/200 golden
rounds, zero mismatches; preamble median 27us vs 50us baseline (-46%, device 3, CaseSmall1,
n=100). a5 compiles clean
(onboard + sim); a5 onboard testing deferred to CI.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01MqeALZTPEnDXTnbYfcnPfq
…RAFT] Let the orchestrator thread (tidx == nthreads-1) skip the AICore handshake and start building the (core-agnostic) task graph immediately, overlapping the schedulers' handshake instead of paying the ~40us preamble serially. It derives its AIC/AIV counts from cores_total_num_ (fixed 1:2 cluster ratio, known before handshake) rather than from post_handshake_init. The other nthreads-1 threads re-partition all cores; total_tasks_ moves to pre_handshake_init so the orchestrator's early SM reset can't race it. Measured (device 3, unlocked, 50 rounds, Device median): bgemm 692.8->661.5us (-4.5%), qwen 2179.6->2158.7us (-1.0%). Recovers most of the fixed handshake preamble hw-native-sys#1214 left on the table; bigger % on smaller workloads. A locality refinement (handshake_owned_clusters: each thread handshakes the cores it will manage, warming its own core_exec_states_) was measured NULL (within noise) and is kept off the active path for reference — remove before merge. DRAFT: a2a3 only; golden = paged_attention Case1 x10 only; measured unlocked. Stacks on hw-native-sys#1214 (single-round-trip handshake). Full record + resume checklist in docs/investigations/2026-07-orchestrator-decouple-preamble.md. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01MqeALZTPEnDXTnbYfcnPfq
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
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 |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request parallelizes the per-core AICore handshake during initialization across multiple AICPU threads to reduce preamble overhead. It splits the initialization process into pre-handshake, parallel partitioned handshake, and post-handshake phases for both the a2a3 and a5 architectures. Additionally, for a2a3, it introduces an optimization allowing the orchestrator thread to skip the handshake and build the task graph concurrently. The handshake protocol itself is simplified by combining reports and using register polling to verify that register windows are open. There are no review comments, so no further feedback is provided.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
1 participant
Human Summary
The idea is simple. Allow the orchestrator thread to start creating the task graph immediately, while the schedulers do the handshaking with the AI cores. There is an additional benefit of locality: the schedulers only handshake with the AIcores they will manage. Futhermore, this allows schedulers to start working immediately, even if others are still handshaking.
AI Writeup
[DRAFT] perf(runtime): decouple the orchestrator from the AICore handshake barrier
AI Summary
tl;dr for humans: After #1214, the AICore handshake is still a hard barrier —
all AICPU threads finish greeting the 72 cores (~40 µs) before any work starts.
But the orchestrator thread only builds the (core-agnostic) task graph; it
doesn't dispatch to cores, so it doesn't need the handshake. This PR lets it
skip the handshake and start building the graph immediately, hiding the ~40 µs
preamble behind graph-build. Result: −4.5% Device on small workloads (bgemm),
−1.0% on qwen decode — modest but real, recovering most of the fixed preamble
#1214 left on the table. A follow-on "locality" idea (each thread handshakes the
cores it will manage) was measured and showed no benefit — kept in the tree
for reference, to be removed before merge.
What changes
The orchestrator (
tidx == nthreads-1) returns frominit()immediately insteadof joining the handshake barrier, and derives its AIC/AIV core counts from
cores_total_num_(fixed 1 AIC : 2 AIV cluster ratio, known before handshake).The other
nthreads-1threads re-partition all cores among themselves (everyregister window still opens). The
total_tasks_read moves topre_handshake_initso the orchestrator's now-early SM reset can't race it.Timing model, before → after:
Measurements (device 3, unlocked, 50 rounds, --skip-golden, Device median)
Locality vs plain decouple is within noise → not worth the complexity.
Status / correctness
paged_attention Case1 ×10PASSES (both variants). Broader golden +locked re-measure still TODO (see below) — this touches the exact
COND/DATA_MAIN_BASE/dispatch machinery de-raced in perf(runtime): overlap AICore handshake wakeups; batch the release barrier #1214, so correctness needs
care before merge.
task-submitfor CI-grade numbers.Before merge / resume checklist
handshake_owned_clusters(locality, measured null) + its decl.task-submit); broaden golden (full case set, high rounds).emergency_shutdown/ handshake-failure path (orchestrator now skips handshake).Full record + rationale:
docs/investigations/2026-07-orchestrator-decouple-preamble.md(in-branch).🤖 Generated with Claude Code