perf(runtime): remake of #1214 -- Single-round-trip AICore handshake

[SergioMartin86]

Human Summary

Currently, the AICPU cores are initializing AI vector and cube cores by performing a 3-step handshaking procedure. This PR takes the ideas of #1214 by @ChaoWao and adapts it to the current main, where the handshaking is already parallelized (see: #1279). The main benefits:

• The handshake procedure is now a 1-step (single-way) procedure
• The handshake parallelization is also now applies to A5

As with #1279, this optimization benefits the short/coarsed-grained-task kernels the most, reporting from 1.7% to 8.7% speedup on simpler benchmark set.

Note: I don't have a A5 machine for development, so I'm counting on the CI to check this works.

AI Summary

tl;dr for humans: The AICore init handshake currently costs three
round-trips per core between the AICPU and each of the 72 AICores. This PR
collapses it to one by folding the "registers are ready" acknowledgement
into the register-window-open write that already had to happen. Two handshake
fields are deleted, a5 gets the parallel-handshake split it was missing, and
the per-core preamble shrinks. Net effect: every benchmark case runs faster
on Device (−1.7% to −8.7%, median ≈ −5%) with zero correctness regressions
(200/200 golden rounds). It's a pure preamble/latency win — orchestration and
scheduling are untouched.

Re-does the approach of #1214 on a fresh branch off current main.

---

What changes

Collapse the AICore init handshake from a three-phase ping-pong
(aicore_regs_ready → aicpu_regs_ready → aicore_done) to a single round
trip:

1. Single-round-trip protocol. AICore publishes
   {physical_core_id, core_type, aicore_done} in one write, then polls its own
   DATA_MAIN_BASE SPR. The AICPU opens that core's register window
   (platform_init_aicore_regs writes DATA_MAIN_BASE = AICPU_IDLE_TASK_ID,
   non-zero) only after observing aicore_done — so the window-open write
   doubles as the release acknowledgement. The dedicated aicpu_regs_ready /
   aicore_regs_ready fields are removed from struct Handshake.

2. Sweep polling (a2a3, already parallel across AICPU threads). Within each
   thread's core slice, poll every outstanding core per pass on aicore_done
   (a GM read — not the strictly-serial nGnRE COND MMIO poll) and service
   whichever have reported, so per-core wakeups overlap (≈ max, not Σ) instead
   of blocking core i before looking at core i+1.

3. Batched release barrier. One OUT_OF_ORDER_STORE_BARRIER per slice
   instead of one per core.

4. a5. Gains the parallel-handshake split (pre_handshake_init /
   handshake_partition / post_handshake_init) for the first time — a5
   previously handshaked serially on one thread — plus the single-round-trip
   protocol.

This is not work-pipelining: the handshake is still a full barrier before
dispatch. It just makes that barrier cheaper.

Why it's safe

The single-round-trip rests on a silicon-timing invariant: a kernel launch
resets DATA_MAIN_BASE to 0; the AICPU writes it non-zero to open FAST_PATH;
and the AICPU spends microseconds in assign_cores_to_threads between opening
the window and the first dispatch, so the core's poll observes the IDLE value
long before any real task_id — it cannot miss it. Verified below with a heavy
golden run.

Verification (a2a3 onboard)

Correctness — paged_attention Case1, 200/200 golden rounds, 0
mismatches (golden checked per round).

Performance — 8-case set, 100 rounds --skip-golden, back-to-back on the
same device. Device µs (the metric; Effective/Orch/Sched stay flat
within noise because they exclude the preamble):

Case	main	this PR	Δ
alternating_matmul_add C1	801.1	752.5	−6.1%
benchmark_bgemm C0	763.2	696.7	−8.7%
paged_attention_unroll C1	1242.9	1197.6	−3.6%
paged_attention_unroll C2	660.4	612.3	−7.3%
pau_manual_scope C1	1245.4	1187.7	−4.6%
pau_manual_scope C2	661.5	606.4	−8.3%
batch_paged_attention C1	3184.8	3131.0	−1.7%
qwen3_14b_decode	2244.5	2170.2	−3.3%

The delta is a near-constant ~45–74µs across all workloads — the signature of a
fixed per-round preamble reduction (larger % on smaller workloads). Preamble
span itself: median 50µs → 27µs (−46%).

a5: compiles clean (onboard + sim); a5 onboard testing deferred to CI
(no a5 silicon on the dev box used).

Files

• a2a3: runtime.h (struct), aicore/aicore_executor.cpp,
  runtime/scheduler/scheduler_cold_path.cpp
• a5: runtime.h, aicore/aicore_executor.cpp, aicpu/aicpu_executor.cpp,
  runtime/scheduler/scheduler_cold_path.cpp, scheduler_context.h, docs
• Shared docs describing the old 3-phase protocol updated to match.

Notes for reviewers

• Benchmarks above were run unlocked (no task-submit on the box), so the
  sub-2% cases (batch) are "improved but within noise"; the overall direction —
  every case faster, constant-offset — is unambiguous. Worth a CI-locked re-run.
• emergency_shutdown no longer broadcasts aicpu_regs_ready; exit is now
  delivered via each core's register block (platform_deinit_aicore_regs writes
  DATA_MAIN_BASE = EXIT), which also releases a core still polling for its
  window. This path is only hit on handshake failure and isn't covered by the
  golden run — reviewed by inspection.

🤖 Generated with Claude Code

[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: bda8e87f-eecb-48d9-9832-1cd79eafcb05

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

This PR removes the aicpu_regs_ready/aicore_regs_ready fields from the Handshake struct in both the a2a3 and a5 runtimes, replacing the register-ready round-trip with an aicore_done write and DATA_MAIN_BASE polling. The a5 scheduler cold-path init is split into pre_handshake_init, partitioned handshake_partition, and post_handshake_init, and AicpuExecutor init/deinit/execute are reworked into a multi-thread barrier-based flow. Documentation is updated to match.

Changes

a2a3 handshake protocol update

Layer / File(s)	Summary
Handshake struct field removal src/a2a3/runtime/tensormap_and_ringbuffer/runtime/runtime.h	Removes aicpu_regs_ready/aicore_regs_ready fields and updates field documentation.
AICore executor handshake rewrite src/a2a3/runtime/tensormap_and_ringbuffer/aicore/aicore_executor.cpp	Combines physical_core_id/core_type reporting with aicore_done, then polls DATA_MAIN_BASE before writing initial COND.
Scheduler handshake_partition/emergency_shutdown rewrite src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_cold_path.cpp	Refactors handshake into slice-wide publish/release phases polling aicore_done with core_serviced tracking; removes per-core aicpu_regs_ready from emergency shutdown.
Related documentation updates docs/aicpu-kernel-launch-mechanisms.md, docs/hardware/mmio-performance.md, docs/investigations/..., docs/troubleshooting/sim-oversubscription-hang.md	Updates docs to reference aicore_done/handshake_partition/DATA_MAIN_BASE instead of removed fields/functions.

a5 handshake protocol and init lifecycle refactor

Layer / File(s)	Summary
Handshake struct field removal src/a5/runtime/tensormap_and_ringbuffer/runtime/runtime.h	Removes aicpu_regs_ready/aicore_regs_ready fields, updates comments for remaining fields.
AICore executor handshake rewrite src/a5/runtime/tensormap_and_ringbuffer/aicore/aicore_executor.cpp	Reorders handshake to report identity fields and aicore_done earlier, waits on DATA_MAIN_BASE before writing COND.
SchedulerContext API split src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_context.h	Replaces init with pre_handshake_init/handshake_partition/post_handshake_init, adds handshake_failed_ flag.
Scheduler cold-path partitioned handshake implementation src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_cold_path.cpp	Implements per-thread partitioned handshake polling aicore_done, updates emergency_shutdown, splits init into pre/post phases rebuilding worker-id lists.
AicpuExecutor multi-thread init/deinit/execute rewrite src/a5/runtime/tensormap_and_ringbuffer/aicpu/aicpu_executor.cpp	Replaces single-entry guard with barrier-based multi-thread handshake coordination atomics; aicpu_execute() checks init() return value directly.
RUNTIME_LOGIC documentation update src/a5/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md	Updates docs to describe new cold-path lifecycle naming.

Estimated code review effort: 4 (Complex) | ~60 minutes

Sequence Diagram(s)

sequenceDiagram
  participant AicpuExecutor
  participant SchedulerContext
  participant AICore
  AicpuExecutor->>SchedulerContext: pre_handshake_init(runtime, ...)
  AicpuExecutor->>SchedulerContext: handshake_partition(tidx, nthreads)
  SchedulerContext->>AICore: publish payload pointers, set aicpu_ready
  AICore->>SchedulerContext: report physical_core_id, core_type, aicore_done
  SchedulerContext->>SchedulerContext: platform_init_aicore_regs(reg_addr)
  AicpuExecutor->>SchedulerContext: post_handshake_init(runtime)
  SchedulerContext->>SchedulerContext: rebuild worker-id lists, assign_cores_to_threads

Loading

Possibly related PRs

• hw-native-sys/simpler#1279: Both PRs refactor the scheduler cold-path handshake by splitting SchedulerContext::init into pre_handshake_init/parallel handshake_partition/post_handshake_init and adjust waits from *_regs_ready fields to aicore_done.

Poem

A rabbit hopped through registers deep,
Found stale old flags that would not sleep. 🐇
Traded regs_ready for aicore_done,
Split the init phases, one by one.
Now handshakes hop in partitioned grace—
No more deadlock in this burrow space! 🥕

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 52.94% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly names the main change: a single-round-trip AICore handshake.
Description check	✅ Passed	The description directly matches the handshake refactor, A5 parallelization, and docs updates in the changeset.
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.}

[gemini-code-assist]

Code Review

This pull request parallelizes the AICore handshake process across multiple AICPU threads to reduce preamble latency, splitting the initialization into pre-handshake, partitioned handshake, and post-handshake phases, while simplifying the handshake protocol by removing redundant flags. The review feedback highlights critical issues regarding bare busy-spins in the thread synchronization loops of AicpuExecutor::init that could lead to CPU starvation or livelocks in simulation environments, recommending the use of SPIN_WAIT_HINT() and relaxed memory ordering. Additionally, it points out a potential out-of-bounds vulnerability in SchedulerContext::pre_handshake_init due to an incomplete bounds check on the signed cores_total_num_ variable, which should be updated to reject negative values.

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.

[coderabbitai]

Actionable comments posted: 2

🤖 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 `@docs/aicpu-kernel-launch-mechanisms.md`:
- Around line 195-205: Update the remaining references from
SchedulerContext::handshake_all_cores to SchedulerContext::handshake_partition
in the affected documentation snippet so the naming matches the renamed
scheduler function and the already-updated hardware doc. Keep the surrounding
explanation intact, and ensure any mention of the handshake flow in this section
consistently uses handshake_partition alongside the updated aicore_done field
name.

In `@src/a5/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md`:
- Line 602: The “Public surface” table is stale because it still describes
SchedulerContext::init(runtime, aicpu_thread_num, sched_thread_num, regs_base)
even though the lifecycle was split into pre_handshake_init,
handshake_partition, and post_handshake_init. Update that table entry to reflect
the new public methods and remove or replace the old init() reference so it
matches the current SchedulerContext surface and the scheduler_cold_path.cpp
lifecycle split.

🪄 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

Run ID: efb67873-2118-4fc4-81a0-b617be4c3327

📥 Commits

Reviewing files that changed from the base of the PR and between 29595b1 and 9bcfca8.

📒 Files selected for processing (13)

• docs/aicpu-kernel-launch-mechanisms.md
• docs/hardware/mmio-performance.md
• docs/investigations/2026-06-pa-unroll-207001-optimeout-window.md
• docs/troubleshooting/sim-oversubscription-hang.md
• src/a2a3/runtime/tensormap_and_ringbuffer/aicore/aicore_executor.cpp
• src/a2a3/runtime/tensormap_and_ringbuffer/runtime/runtime.h
• src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_cold_path.cpp
• src/a5/runtime/tensormap_and_ringbuffer/aicore/aicore_executor.cpp
• src/a5/runtime/tensormap_and_ringbuffer/aicpu/aicpu_executor.cpp
• src/a5/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md
• src/a5/runtime/tensormap_and_ringbuffer/runtime/runtime.h
• src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_cold_path.cpp
• src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_context.h