[Optimization] Parallelize AICore handshake across AICPU threads -- 3~22% Device Speedup

[SergioMartin86]

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::init is split into three phases:

1. pre_handshake_init (leader thread only): the shared setup — zero the
   per-core state, wire configuration, initialize profiling/swimlane buffers,
   and record the core count. Published to the other threads via a flag.
2. 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.
3. post_handshake_init (leader thread only, after a barrier): the cheap
   bookkeeping — build the ordered AIC/AIV worker lists, assign cores to
   scheduler threads, initialize dispatch payloads.

AicpuExecutor::init orchestrates the three phases with a simple barrier: the
leader publishes "setup done", all threads handshake their slice and report in,
and the leader waits for all of them before finishing. Startup timing:

                serial (before)            parallel (after)
  broadcast     ~35 us                     ~9 us  (split 4 ways)
  boot wait     ~11 us                     ~11 us (unavoidable, physical)
  per-core work ~217 us  <-- the tax       ~55 us (split 4 ways)
  ---------------------------------------------------------------
  handshake     ~263 us                    ~75 us

Results

Direct A/B of this PR vs upstream/main (both the same upstream runtime, so
this 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):

Workload	upstream/main	this PR	change
paged_attention_unroll C2	866.5 µs	675.3 µs	−22.1%
pa_unroll_manual_scope C2	854.3 µs	678.9 µs	−20.5%
benchmark_bgemm	962.5 µs	790.3 µs	−17.9%
alternating_matmul_add	987.2 µs	815.9 µs	−17.4%
paged_attention_unroll C1	1448.2 µs	1261.0 µs	−12.9%
pa_unroll_manual_scope C1	1430.9 µs	1261.2 µs	−11.9%
qwen3_14b_decode	2409.8 µs	2228.3 µs	−7.5%
batch_paged_attention	3400.9 µs	3270.6 µs	−3.8%

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 graph
• benchmark_bgemm — AIC matrix-multiply (Case0, Bgemm64)
• paged_attention — 65,792-task orchestration-heavy run (Case1) + a small case
• qwen3_14b_decode — dense two-layer transformer decode

The 65k-task and multi-scope cases are the strongest exercise of the new
multi-thread partition + barrier path.

Why it's safe

• No new config or feature flag — parallelizing the handshake is always
  correct, so it's unconditional.
• No data races — each core is handshaked by exactly one thread (contiguous,
  gap-free partition), so the per-core writes never overlap.
• Ordering preserved — the AIC/AIV worker lists are rebuilt in the same
  core-index order as before, so core-to-thread assignment is identical.
• Memory ordering — a setup failure is stored before the "setup done" flag
  (both release/acquire), so a failing init can never let a thread proceed. The
  swimlane-buffer init still happens before any core is signaled.
• Profiling-agnostic — under PTO2_PROFILING the post-pass reads core type
  and 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 — split init into
  pre_handshake_init / post_handshake_init; replace handshake_all_cores
  with the sliced handshake_partition.
• runtime/scheduler/scheduler_context.h — updated declarations; added the
  handshake_failed_ flag.
• aicpu/aicpu_executor.cpp — parallel init orchestration + barrier; removed
  the 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

• The Results table is a direct A/B of this branch against upstream/main — both
  built 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.
• A separate, pre-existing defect was noticed during this work: the L2 swimlane
  collector produces no data on the tensormap_and_ringbuffer runtime. It is
  unrelated to this change and tracked separately.

[coderabbitai]

Important

Review skipped

Auto incremental reviews are disabled on this repository.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: 78bfbb69-f716-406a-b775-f8d9304d3fcf

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

📝 Walkthrough

Walkthrough

The PR refactors AICPU/scheduler initialization from a single-threaded init()/handshake_all_cores() flow into a three-phase model: pre_handshake_init, parallelized handshake_partition, and post_handshake_init. AicpuExecutor is rewritten to coordinate all threads through this flow via new atomics, and related comments/docs are updated accordingly.

Changes

Scheduler Handshake Refactor

Layer / File(s)	Summary
Scheduler init API and state contract src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_context.h	Replaces the single init() declaration with pre_handshake_init, handshake_partition, and post_handshake_init; removes handshake_all_cores; adds handshake_failed_ atomic flag.
Partitioned handshake and split init implementation src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_cold_path.cpp	Implements handshake_partition to process only a per-thread worker-id slice, latches failures into a shared handshake_failed_ flag instead of returning immediately, and splits init logic into pre_handshake_init (setup) and post_handshake_init (worker-list building, assign_cores_to_threads, profiling/task-count init, emergency shutdown on failure).
AICPU executor multi-thread init coordination src/a2a3/runtime/tensormap_and_ringbuffer/aicpu/aicpu_executor.cpp	Adds hs_setup_done_/hs_arrived_ atomics; rewrites init() so a leader thread runs pre_handshake_init/post_handshake_init while all threads run handshake_partition via a barrier; updates deinit() and simplifies aicpu_execute() to use init()'s return value.
Comment and documentation updates src/a2a3/runtime/tensormap_and_ringbuffer/aicore/aicore_executor.cpp, src/a2a3/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md	Updates inline comment and RUNTIME_LOGIC.md to describe the new handshake phase names and ordering.

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

Loading

Possibly related PRs

• hw-native-sys/simpler#1060: Addresses the same AICPU↔AICore handshake coupling by changing host launch order around the old handshake_all_cores() spin-wait race that this PR replaces.
• hw-native-sys/simpler#1217: Also directly modifies the SchedulerContext::init machinery and core-transition flow in the same components.

Poem

A hop, a skip, three phases neat,
pre, partition, post — complete!
No more waiting, all threads race,
each core finds its rightful place. 🐇
Handshake done, the burrow cheers,
init flows clean, away with fears!

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
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.
Title check	✅ Passed	The title clearly summarizes the main change: parallelizing the AICore handshake across AICPU threads.
Description check	✅ Passed	The description is directly about the handshake parallelization and its implementation and results.

---

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.}

[gemini-code-assist]

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.