[Optimization] Parallelize AICore handshake across AICPU threads -- 3~22% Device Speedup #1279
[Optimization] Parallelize AICore handshake across AICPU threads -- 3~22% Device Speedup #1279SergioMartin86 wants to merge 1 commit into
Conversation
The AICPU->AICore handshake in scheduler init discovered all cores serially on a single thread while the other AICPU threads spun idle waiting on init_done_. On a2a3 (72 cores) this serial per-core MMIO (reg init + two spin round-trips per core) measured ~217 us of the ~283 us preamble -- a fixed per-round cost that dominates device time for small workloads and is a flat tax on every run. Parallelize it: split SchedulerContext::init into three parts so every AICPU thread handshakes a disjoint slice of cores. - pre_handshake_init (leader): zero per-core state, wire config, init l2-swimlane buffers, set core count -- published before any thread handshakes so the swimlane-before-aicpu_ready ordering still holds. - handshake_partition (all threads): each thread handshakes its contiguous [lo, hi) slice; core writes are race-free (disjoint), and an invalid physical_core_id sets an atomic flag. - post_handshake_init (leader, after a barrier): build the AIC/AIV worker-id lists in core-index order (identical clusters to before, read from the Handshake struct so it stays correct under PTO2_PROFILING), assign cores, init profiling subsystems, payloads. AicpuExecutor::init orchestrates it: the leader (affinity exec_idx 0) publishes hs_setup_done_, all threads handshake their slice, then an hs_arrived_ barrier gates the leader's finish. init_failed_ is stored before hs_setup_done_ (both release) so a pre-handshake failure is always visible before any thread proceeds. Drops the now-dead initialized_ CAS and replaces the preamble spin with a direct init-return check. Measured on a2a3: preamble ~283 us -> ~85 us, a ~187 us/round device reduction (-7.5% on qwen3-14b decode, up to -22% on small workloads where the fixed handshake cost dominated). Compute window unchanged. Golden-verified across vector_example, benchmark_bgemm, paged_attention (65k tasks), and qwen3_14b_decode. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01MqeALZTPEnDXTnbYfcnPfq
📝 WalkthroughWalkthroughThe PR refactors AICPU/scheduler initialization from a single-threaded ChangesScheduler Handshake Refactor
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Thread as AICPU Thread
participant SchedulerContext
participant Handshake as Handshake Records
Thread->>SchedulerContext: pre_handshake_init(...)
SchedulerContext-->>Thread: clears handshake_failed_, sets cores_total_num_
Thread->>SchedulerContext: handshake_partition(tidx, nthreads)
SchedulerContext->>Handshake: write payload pointer, set aicpu_ready=1 for slice
SchedulerContext->>Handshake: wait aicore_regs_ready/aicore_done for slice
alt invalid physical_core_id
SchedulerContext->>SchedulerContext: handshake_failed_.store(true)
end
Thread->>SchedulerContext: post_handshake_init()
alt handshake_failed_
SchedulerContext->>SchedulerContext: emergency_shutdown
else success
SchedulerContext->>Handshake: build AIC/AIV worker-id lists
SchedulerContext->>SchedulerContext: assign_cores_to_threads()
SchedulerContext->>SchedulerContext: profiling/task-count init
end
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 |
![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 to optimize the preamble execution time. It splits the previous single-threaded initialization into three phases: a leader-only pre-handshake setup, a parallelized handshake partition where each thread handles a disjoint slice of cores, and a leader-only post-handshake phase. Feedback suggests that instead of silently falling back to a default index when platform_aicpu_affinity_thread_idx() returns an invalid value, the code should validate the thread index and fail loudly with an error code.
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.
| int32_t tidx = platform_aicpu_affinity_thread_idx(); | ||
| if (tidx < 0) tidx = 0; | ||
| const bool is_leader = (tidx == 0); |
There was a problem hiding this comment.
When implementing thread affinity gates or filters, we must validate input parameters and thread indices rather than silently bypassing the gate or falling back to dynamic index assignment. If platform_aicpu_affinity_thread_idx() returns an invalid index (such as a negative value or one exceeding the thread count), we should log an error and return an error code to fail loudly.
int32_t tidx = platform_aicpu_affinity_thread_idx();
if (tidx < 0 || tidx >= nthreads) {
AICPU_LOGE("Invalid thread affinity index: %d", tidx);
return AICPU_ERROR_INVALID_PARAMETER;
}
const bool is_leader = (tidx == 0);References
- When implementing thread affinity gates or filters, validate input parameters (such as allowed CPU count or launch count) and fail loudly (e.g., log an error and return an error code) if they are invalid, rather than silently dropping threads or bypassing the gate.
1 participant
Human Summary
Currently, upstream/main has a single thread handle the handshake for all compute cores. This is a fixed cost that represents a significant part of the kernel running time. This PR parallelizes the process among the other AICPU cores, obtaining an up to 22% speedup.
The effect is mostly notable on small kernels, whose fixed runtime component is more significant. This improvements is complementary to that of #1137 which improves the dynamic aspect, favoring large kernels with many small tasks.
AI Summary
Every time the runtime starts a run, it does a "roll call" of all 72 AICore
compute units — greeting each one and setting up its registers before any work
can begin. Today one AICPU thread does this roll call one core at a time,
while the other three threads sit idle waiting. That serial roll call takes
~217 µs of the ~283 µs startup on a2a3, and because it's a fixed cost paid
every single run, it hurts short jobs the most.
This PR has all four AICPU threads greet the cores in parallel, each handling
a quarter of them. Startup drops from ~283 µs to ~85 µs, which takes ~172 µs
off the device time of every run — −7.5% on qwen3-14B decode and up to
−22% on small workloads. The actual compute is untouched; only the startup gets
faster. Correctness is unchanged (golden-verified on 8 cases).
The problem
Before any tasks run, the scheduler has to discover the AICore units it will
dispatch to. For each core it: waits for the core to boot and report in, reads
its physical ID, initializes its control registers over MMIO, and waits for the
core to acknowledge. This "handshake" happens fresh on every round (the cores
exit at the end of each round and are relaunched).
The catch: this loop ran serially on a single thread. On the Ascend AICPU,
MMIO register access is strictly serial (~95 ns per access, no pipelining — a
documented hardware constraint), so greeting 72 cores one-by-one is slow by
construction. Meanwhile the other three AICPU threads spun idle, doing nothing
but waiting for the one thread to finish.
We profiled it directly: of the ~283 µs startup ("preamble") on qwen3-14B decode,
~217 µs (82%) was the serial per-core handshake work. Only ~11 µs was
unavoidable (waiting for the first core to physically boot).
Because this cost is fixed per round, it's a small slice of a big job but a
huge slice of a small one — e.g. ~26% of a simple matmul's device time, and it
was a major contributor to why small workloads looked slow.
The fix
The hardware guidance is explicit: the only way to speed up cross-core polling
is to use more threads. So we do exactly that.
SchedulerContext::initis split into three phases:pre_handshake_init(leader thread only): the shared setup — zero theper-core state, wire configuration, initialize profiling/swimlane buffers,
and record the core count. Published to the other threads via a flag.
handshake_partition(all four threads): each thread greets a disjoint,contiguous slice of the cores (18 each for 72). Because the slices don't
overlap, there's no contention — the expensive MMIO work now runs four-wide.
post_handshake_init(leader thread only, after a barrier): the cheapbookkeeping — build the ordered AIC/AIV worker lists, assign cores to
scheduler threads, initialize dispatch payloads.
AicpuExecutor::initorchestrates the three phases with a simple barrier: theleader publishes "setup done", all threads handshake their slice and report in,
and the leader waits for all of them before finishing. Startup timing:
Results
Direct A/B of this PR vs
upstream/main(both the same upstream runtime, sothis isolates the handshake change), back-to-back on the same NPU (device 2),
100 rounds each, trimmed-80. Device time is the metric of record (on-NPU
wall clock):
Every workload improves by a near-constant ~172 µs — the tell-tale signature
of removing a fixed overhead. The relative win scales inversely with job size:
biggest on the small workloads that the fixed cost dominated, smallest on the
long-running ones. The compute window itself is unchanged (−0.3% to +2.0%,
noise), because this only touches startup — not scheduling or kernel execution.
Correctness
Golden-verified (output compared against a PyTorch reference) on a deliberately
diverse spread, all PASSED:
vector_example— tiny, simple graphbenchmark_bgemm— AIC matrix-multiply (Case0, Bgemm64)paged_attention— 65,792-task orchestration-heavy run (Case1) + a small caseqwen3_14b_decode— dense two-layer transformer decodeThe 65k-task and multi-scope cases are the strongest exercise of the new
multi-thread partition + barrier path.
Why it's safe
correct, so it's unconditional.
gap-free partition), so the per-core writes never overlap.
core-index order as before, so core-to-thread assignment is identical.
(both release/acquire), so a failing init can never let a thread proceed. The
swimlane-buffer init still happens before any core is signaled.
PTO2_PROFILINGthe post-pass reads core typeand physical ID from the Handshake struct, so it works whether or not
profiling is compiled in.
Files changed (5)
All in
src/a2a3/runtime/tensormap_and_ringbuffer/:runtime/scheduler/scheduler_cold_path.cpp— splitinitintopre_handshake_init/post_handshake_init; replacehandshake_all_coreswith the sliced
handshake_partition.runtime/scheduler/scheduler_context.h— updated declarations; added thehandshake_failed_flag.aicpu/aicpu_executor.cpp— parallel init orchestration + barrier; removedthe dead single-leader CAS.
aicore/aicore_executor.cpp— comment updated to reference the new function.docs/RUNTIME_LOGIC.md— doc updated to reference the new functions.Notes for reviewers
upstream/main— bothbuilt from the same upstream runtime and run back-to-back on the same NPU — so
it isolates the handshake change with nothing else varying. The ~85 µs startup
and per-phase handshake breakdown come from direct profiling of the same
optimization.
collector produces no data on the
tensormap_and_ringbufferruntime. It isunrelated to this change and tracked separately.