perf: prefetch dispatch buffers and inline timestamp reads#1243
Conversation
📝 WalkthroughWalkthroughThis PR adds ChangesScheduler Dispatch Prefetching
Estimated code review effort: 1 (Trivial) | ~3 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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 introduces __builtin_prefetch instructions in the scheduler dispatch code for both a2a3 and a5 runtimes to pull ownership of per-core dispatch buffers (payload and deferred_slab) before performing hot stores. The reviewer suggests extending the prefetch range to cover tail cache lines (offsets 384+ and 504) that are also written to, in order to prevent RFO cache misses. Additionally, the reviewer recommends reordering the construction of async_ctx before the writes to provide an execution window that overlaps and hides the prefetch latency.
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.
| // Pull ownership of the per-core dispatch buffers before the hot stores below. | ||
| __builtin_prefetch(reinterpret_cast<const char *>(&payload), 1, 3); | ||
| __builtin_prefetch(reinterpret_cast<const char *>(&payload) + 64, 1, 3); | ||
| __builtin_prefetch(reinterpret_cast<const char *>(deferred_slab), 1, 3); | ||
| deferred_slab->count = 0; | ||
| deferred_slab->error_code = PTO2_ERROR_NONE; | ||
| AsyncCtx async_ctx = AsyncCtx::make(slot_state.task->task_id, deferred_slab); |
There was a problem hiding this comment.
The prefetch for payload only covers the first two cache lines (offsets 0 and 64). However, build_payload unconditionally writes to local_context (offset 384+) and not_ready (offset 504), which reside in cache lines 6 and 7. Prefetching these tail cache lines prevents RFO cache misses on these hot stores.
Additionally, deferred_slab is written to immediately after its prefetch instruction, leaving no execution cycles to hide memory latency. Moving the construction of async_ctx before the writes provides a small execution window to overlap the prefetch.
// Pull ownership of the per-core dispatch buffers before the hot stores below.
__builtin_prefetch(reinterpret_cast<const char *>(&payload), 1, 3);
__builtin_prefetch(reinterpret_cast<const char *>(&payload) + 64, 1, 3);
__builtin_prefetch(reinterpret_cast<const char *>(&payload) + 384, 1, 3);
__builtin_prefetch(reinterpret_cast<const char *>(&payload) + 448, 1, 3);
__builtin_prefetch(reinterpret_cast<const char *>(deferred_slab), 1, 3);
// Construct AsyncCtx first to give prefetch instructions some cycles to run
AsyncCtx async_ctx = AsyncCtx::make(slot_state.task->task_id, deferred_slab);
deferred_slab->count = 0;
deferred_slab->error_code = PTO2_ERROR_NONE;| // Pull ownership of the per-core dispatch buffers before the hot stores below. | ||
| __builtin_prefetch(reinterpret_cast<const char *>(&payload), 1, 3); | ||
| __builtin_prefetch(reinterpret_cast<const char *>(&payload) + 64, 1, 3); | ||
| __builtin_prefetch(reinterpret_cast<const char *>(deferred_slab), 1, 3); | ||
| deferred_slab->count = 0; | ||
| deferred_slab->error_code = PTO2_ERROR_NONE; | ||
| AsyncCtx async_ctx = AsyncCtx::make(slot_state.task->task_id, deferred_slab); |
There was a problem hiding this comment.
The prefetch for payload only covers the first two cache lines (offsets 0 and 64). However, build_payload unconditionally writes to local_context (offset 384+), which resides in cache line 6. Prefetching this tail cache line prevents RFO cache misses on these hot stores.
Additionally, deferred_slab is written to immediately after its prefetch instruction, leaving no execution cycles to hide memory latency. Moving the construction of async_ctx before the writes provides a small execution window to overlap the prefetch.
// Pull ownership of the per-core dispatch buffers before the hot stores below.
__builtin_prefetch(reinterpret_cast<const char *>(&payload), 1, 3);
__builtin_prefetch(reinterpret_cast<const char *>(&payload) + 64, 1, 3);
__builtin_prefetch(reinterpret_cast<const char *>(&payload) + 384, 1, 3);
__builtin_prefetch(reinterpret_cast<const char *>(deferred_slab), 1, 3);
// Construct AsyncCtx first to give prefetch instructions some cycles to run
AsyncCtx async_ctx = AsyncCtx::make(slot_state.task->task_id, deferred_slab);
deferred_slab->count = 0;
deferred_slab->error_code = PTO2_ERROR_NONE;![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_dispatch.cpp (1)
165-167: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueCache-line stride is a hardcoded magic number.
64(cache-line size) appears as a bare literal here and is duplicated in the a5 variant. Extracting a named constant (e.g.,CACHE_LINE_SIZE) in a shared header would make intent clearer and avoid silent drift if line size assumptions ever change per target.♻️ Example
+static constexpr size_t CACHE_LINE_SIZE = 64; ... - __builtin_prefetch(reinterpret_cast<const char *>(&payload), 1, 3); - __builtin_prefetch(reinterpret_cast<const char *>(&payload) + 64, 1, 3); + __builtin_prefetch(reinterpret_cast<const char *>(&payload), 1, 3); + __builtin_prefetch(reinterpret_cast<const char *>(&payload) + CACHE_LINE_SIZE, 1, 3);🤖 Prompt for 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. In `@src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_dispatch.cpp` around lines 165 - 167, The prefetch logic in scheduler_dispatch.cpp uses a hardcoded cache-line stride literal, which is duplicated across the a5 variant and should be centralized. Introduce a shared named constant such as CACHE_LINE_SIZE in the relevant common header, then update the prefetch calls in the scheduler dispatch code to use that symbol instead of the bare 64 literal so the intent stays clear and target-specific assumptions do not drift.
🤖 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.
Nitpick comments:
In
`@src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_dispatch.cpp`:
- Around line 165-167: The prefetch logic in scheduler_dispatch.cpp uses a
hardcoded cache-line stride literal, which is duplicated across the a5 variant
and should be centralized. Introduce a shared named constant such as
CACHE_LINE_SIZE in the relevant common header, then update the prefetch calls in
the scheduler dispatch code to use that symbol instead of the bare 64 literal so
the intent stays clear and target-specific assumptions do not drift.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 10dc6104-131f-42de-ad43-0ea82debb913
📒 Files selected for processing (2)
src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_dispatch.cppsrc/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_dispatch.cpp
TaoZQY
changed the title
1 participant
Summary
Add two small dispatch-path optimizations from issue #1103:
The prefetch change warms:
PTO2DispatchPayloadPTO2DispatchPayloadDeferredCompletionSlabThe fast-counter change adds platform-selected
device_time_fast.hhelpers:mrs cntvct_el0get_sys_cnt_aicpu()so CPU-sim keeps its chrono-backed tick scaleThe changes are applied to both a2a3 and a5 dispatch paths. They do not change ABI layout, task state transitions, ring-buffer publishing, or completion semantics.
Motivation
Issue #1103 tracks remaining dispatch-path optimizations from the scheduler optimization branch. The AICPU scheduler writes per-core dispatch payloads and deferred completion slabs on every subtask dispatch. These buffers can be cold, so the first writes may pay cache-line ownership/RFO latency. Pulling ownership slightly earlier can hide part of that cost behind the existing dispatch setup work.
For L2 swimlane AICPU timing, dispatch timestamps are sampled immediately before publishing the DATA_MAIN_BASE doorbell. The onboard implementation of
get_sys_cnt_aicpu()already readscntvct_el0, but it is an out-of-line function call. Inlining this read at the dispatch timestamp sites reduces profiling overhead and keeps the timestamp closer to the doorbell write.Test
Built and installed the branch:
PTO_ISA_ROOT=/data/pyptouser/zhangtao/zt/simpler/build/pto-isa python -m pip install --no-build-isolation .Hardware tests submitted through
task-submit:Task ids:
All tested cases passed.
Default performance comparison, baseline -> prefetch:
spmd_multiblock_mixspmd_sync_start_stressqwen3_14b_decodeDefault performance comparison, prefetch -> prefetch + fast counter:
spmd_multiblock_mixspmd_sync_start_stressqwen3_14b_decodeL2 swimlane AICPU timing smoke tests with
--enable-l2-swimlane 2:spmd_multiblock_mixspmd_sync_start_stressqwen3_14b_decodeNote:
spmd_serial_chain_mixwas not present in the latest upstream tree used for this branch, sospmd_multiblock_mixwas used as the available SPMD mix coverage.