Add selective task-timing slots for lightweight dispatch/finish timing#1359
Add selective task-timing slots for lightweight dispatch/finish timing#1359doraemonmj wants to merge 1 commit into
Conversation
![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 implements selective task-timing slots (0–15) to measure specific task dispatch-to-finish windows without enabling the full L2 swimlane. This feature is supported across both A2A3 and A5 runtimes, including host_build_graph and tensormap_and_ringbuffer. The implementation reuses existing padding in the task descriptor to avoid size growth, hooks into scheduler dispatch and completion to record cycles, and updates host-side runners to read back and emit these timing slots as [STRACE] spans. Comprehensive system and unit tests are also added. Feedback on the changes highlights an ABI compatibility issue in OrchestrationRuntimeOps, where inserting the new set_task_timing_slot function pointer in the middle of the struct shifts subsequent fields; it is recommended to append new pointers to the end of the struct to maintain ABI stability.
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.
📝 WalkthroughWalkthroughSelective task-timing slots add 16 opt-in dispatch-to-finish windows for tagged tasks. The implementation carries slot metadata through orchestration and schedulers, aggregates records across threads, resolves host timestamps, emits STRACE spans, and adds examples, tests, and documentation. ChangesSelective task-timing slots
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Orchestrator
participant Scheduler
participant DevicePhaseBuffer
participant HostRunner
participant STRACE
Orchestrator->>Scheduler: submit tagged task
Scheduler->>DevicePhaseBuffer: fold dispatch and finish cycles
HostRunner->>DevicePhaseBuffer: reset before run and read after synchronization
HostRunner->>STRACE: emit task_slot_N device-clock span
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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (2)
src/a5/runtime/host_build_graph/aicpu/aicpu_executor.cpp (1)
770-774: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winComplete the older task before the newer task for chronological accuracy.
Reorder the calls to ensure the implicitly completed older task (
prev_running_id) is recorded before the newly finished task (completed_task_id). Sincefold_task_finishreads the current system cycle for the timestamp, calling it in reverse chronological order can result in the task that finished first receiving a slightly later timestamp.♻️ Proposed refactor
// Both the pending task (FIN observed here) and the running task // (implicitly done — AICore overwrote COND before its FIN) complete. - fold_task_finish(runtime, completed_task_id, thread_idx); fold_task_finish(runtime, prev_running_id, thread_idx); + fold_task_finish(runtime, completed_task_id, thread_idx);🤖 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/a5/runtime/host_build_graph/aicpu/aicpu_executor.cpp` around lines 770 - 774, In the completion handling block, reorder the fold_task_finish calls so the older implicitly completed task prev_running_id is recorded before the newer completed_task_id task. Keep both calls and their existing arguments unchanged.src/common/platform/onboard/host/c_api_shared.cpp (1)
537-546: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winGuard
kTaskSlotNamessize againstNUM_TASK_TIMING_SLOTSdrift, like the neighboringkPhasestable does.Both
emit_device_phase_markersimplementations declarekTaskSlotNames[NUM_TASK_TIMING_SLOTS]with exactly 16 literal entries but no compile-time check tying the two together, unlike the adjacentkPhases[]table which hasstatic_assert(sizeof(kPhases)/sizeof(kPhases[0]) == NUM_AICPU_PHASES - 1, ...). IfNUM_TASK_TIMING_SLOTSis ever changed without updating this literal list, the extra array slots silently zero-fill tonullptr, and the loop would pass a null name intoSTRACE_DEV_SPAN_ATfor any newly-added slot index that reports a valid (dispatch < finish) interval.
src/common/platform/onboard/host/c_api_shared.cpp#L537-L546: add astatic_assert(sizeof(kTaskSlotNames) / sizeof(kTaskSlotNames[0]) == NUM_TASK_TIMING_SLOTS, ...)next to the array.src/common/platform/sim/host/c_api_shared.cpp#L494-L503: add the identical static_assert here as well.🛡️ Proposed fix (apply to both files)
static const char *const kTaskSlotNames[NUM_TASK_TIMING_SLOTS] = { "simpler_run.runner_run.device_wall.task_slot_0", "simpler_run.runner_run.device_wall.task_slot_1", "simpler_run.runner_run.device_wall.task_slot_2", "simpler_run.runner_run.device_wall.task_slot_3", "simpler_run.runner_run.device_wall.task_slot_4", "simpler_run.runner_run.device_wall.task_slot_5", "simpler_run.runner_run.device_wall.task_slot_6", "simpler_run.runner_run.device_wall.task_slot_7", "simpler_run.runner_run.device_wall.task_slot_8", "simpler_run.runner_run.device_wall.task_slot_9", "simpler_run.runner_run.device_wall.task_slot_10", "simpler_run.runner_run.device_wall.task_slot_11", "simpler_run.runner_run.device_wall.task_slot_12", "simpler_run.runner_run.device_wall.task_slot_13", "simpler_run.runner_run.device_wall.task_slot_14", "simpler_run.runner_run.device_wall.task_slot_15", }; + static_assert( + sizeof(kTaskSlotNames) / sizeof(kTaskSlotNames[0]) == NUM_TASK_TIMING_SLOTS, + "kTaskSlotNames[] must list every task-timing slot — add the new slot name here" + );🤖 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/common/platform/onboard/host/c_api_shared.cpp` around lines 537 - 546, Add a compile-time size check immediately after kTaskSlotNames in both src/common/platform/onboard/host/c_api_shared.cpp lines 537-546 and src/common/platform/sim/host/c_api_shared.cpp lines 494-503, asserting its element count equals NUM_TASK_TIMING_SLOTS; use the same static_assert pattern as the neighboring kPhases table.
🤖 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 `@examples/workers/l2/task_timing_slots/test_task_timing_e2e.py`:
- Around line 187-211: Guard the baseline slot lookups in
test_duplicate_slot_merges_window by storing the results of
_slot_spans(base_err, 0) and _slot_spans(base_err, 1), asserting each list is
non-empty with a clear failure message, and only then indexing element zero to
compute base_single.
In `@src/a2a3/runtime/host_build_graph/runtime/pto_types.h`:
- Around line 235-249: Update set_task_timing_slot in the shown pto_types.h and
its sibling tensormap_and_ringbuffer/runtime/pto_types.h to preserve the sticky
first-error invariant: when the slot is out of range, record has_error and
error_msg only if no prior error exists, matching set_error(). Keep valid-slot
assignment and existing validation behavior unchanged.
In `@src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_types.h`:
- Around line 266-280: Update set_task_timing_slot in both the shared
pto_types.h implementation and its host_build_graph fork to report out-of-range
slots through the existing sticky first-error mechanism, such as set_error(),
instead of assigning has_error and error_msg directly. Preserve the current
range validation and return behavior while ensuring an earlier diagnostic cannot
be overwritten.
---
Nitpick comments:
In `@src/a5/runtime/host_build_graph/aicpu/aicpu_executor.cpp`:
- Around line 770-774: In the completion handling block, reorder the
fold_task_finish calls so the older implicitly completed task prev_running_id is
recorded before the newer completed_task_id task. Keep both calls and their
existing arguments unchanged.
In `@src/common/platform/onboard/host/c_api_shared.cpp`:
- Around line 537-546: Add a compile-time size check immediately after
kTaskSlotNames in both src/common/platform/onboard/host/c_api_shared.cpp lines
537-546 and src/common/platform/sim/host/c_api_shared.cpp lines 494-503,
asserting its element count equals NUM_TASK_TIMING_SLOTS; use the same
static_assert pattern as the neighboring kPhases table.
🪄 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: c4d9be09-87c4-4c65-8f00-fdb155ca2c3c
📒 Files selected for processing (45)
docs/dfx/device-phases.mddocs/dfx/l2-swimlane-profiling.mddocs/dfx/l2-timing.mdexamples/workers/l2/task_timing_slots/__init__.pyexamples/workers/l2/task_timing_slots/kernels/orchestration/task_timing_orch.cppexamples/workers/l2/task_timing_slots/test_task_timing_e2e.pysrc/a2a3/platform/sim/host/device_runner.cppsrc/a2a3/runtime/host_build_graph/orchestration/pto_arg_with_deps.hsrc/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_orchestrator.cppsrc/a2a3/runtime/host_build_graph/runtime/pto_runtime2_types.hsrc/a2a3/runtime/host_build_graph/runtime/pto_types.hsrc/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_completion.cppsrc/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_context.hsrc/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_dispatch.cppsrc/a2a3/runtime/tensormap_and_ringbuffer/orchestration/pto_arg_with_deps.hsrc/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.cppsrc/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2_types.hsrc/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_types.hsrc/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_completion.cppsrc/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_context.hsrc/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_dispatch.cppsrc/a5/platform/sim/host/device_runner.cppsrc/a5/runtime/host_build_graph/aicpu/aicpu_executor.cppsrc/a5/runtime/host_build_graph/host/runtime_maker.cppsrc/a5/runtime/host_build_graph/orchestration/orchestration_api.hsrc/a5/runtime/host_build_graph/runtime/runtime.cppsrc/a5/runtime/host_build_graph/runtime/runtime.hsrc/a5/runtime/tensormap_and_ringbuffer/orchestration/pto_arg_with_deps.hsrc/a5/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.cppsrc/a5/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2_types.hsrc/a5/runtime/tensormap_and_ringbuffer/runtime/pto_types.hsrc/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_completion.cppsrc/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_context.hsrc/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_dispatch.cppsrc/common/platform/include/aicpu/device_phase_aicpu.hsrc/common/platform/include/common/device_phase.hsrc/common/platform/onboard/host/c_api_shared.cppsrc/common/platform/onboard/host/device_runner_base.cppsrc/common/platform/onboard/host/device_runner_base.hsrc/common/platform/sim/host/c_api_shared.cppsrc/common/platform/sim/host/device_runner_base.htests/st/a5/host_build_graph/task_timing/kernels/orchestration/task_timing_a5hbg_orch.cpptests/st/a5/host_build_graph/task_timing/test_task_timing_a5hbg.pytests/ut/cpp/CMakeLists.txttests/ut/cpp/a2a3/test_task_timing_slots.cpp
doraemonmj
force-pushed
the
pr/task-timing-slots
branch
4 times, most recently
from
852a28b to
dd4267a
Compare
Adds 16 fixed per-run task-timing slots on the existing device-phase transport. Orchestration tags selected tasks with a slot id (0..15) and the Scheduler folds each tagged task's AICPU dispatch/finish cycles into that slot; the host resets slots per run, reads them back after stream sync, and emits each complete slot as a device-clock task_slot_N [STRACE] span. Works with L2 swimlane disabled and in SIMPLER_DFX=0 builds; untagged tasks pay only a cache-hot sentinel check. Covers a2a3/a5, onboard/sim, both tensormap_and_ringbuffer and host_build_graph, including the legacy a5 host_build_graph add_task path. Adds unit, sim e2e, and a5 onboard tests plus device-phase/L2-timing docs. Relocated tests to tests/st/task_timing/ for better organization: - task_timing_slots/ - RT2 e2e tests (tensormap_and_ringbuffer) - task_timing_a5hbg/ - a5 host_build_graph legacy API tests
1 participant
Summary
Adds 16 fixed, per-run task-timing slots on the existing lightweight
device-phase transport. Orchestration tags selected tasks with a slot id
(
0..15); the Scheduler folds each tagged task's AICPUdispatch/finishcycles into that slot. The host resets slots before each run, reads them
back once after stream sync, and emits each complete slot as a device-clock
simpler_run.runner_run.device_wall.task_slot_N[STRACE]span.Works with L2 swimlane disabled and in
SIMPLER_DFX=0builds — it startsno collector threads and writes no per-task AICore records. Untagged tasks
pay only a cache-hot sentinel check (
!= TASK_TIMING_SLOT_NONE); they neverread the sys counter.
Closes #1325.
Design
TaskTimingRecord[16]tail appended after theAicpuPhaseRecordregion in the same device buffer — same base pointer,same per-run H2D reset and post-sync D2H copy.
PTO2TaskDescriptor::kernel_id[3], so the descriptor does not grow;static_asserts pin its size andpacked_buffer_baseoffset.min(dispatch)/max(finish)across Schedulerthreads and across a MIX task's subtasks / an SPMD task's blocks. Reusing a
slot merges into one window; distinct slots recover
finish(B) - dispatch(A).dispatch= earliest gatedDATA_MAIN_BASEpublication;finish= latest FIN observation (afterrmb(), before fanin/deferred-completion). Incomplete/unset slots areskipped, so a short/failed run leaks no stale data.
Runtime::Taskpath, so it gets asmall
set_task_timing_slot(runtime, task_id, slot)compatibility setterrather than the RT2
L0TaskArgsfield.Coverage
Supports a2a3/a5, onboard/sim, and both
tensormap_and_ringbufferandhost_build_graph, plus unit / sim-e2e / onboard tests and device-phase /L2-timing docs.
Verification
Unit — a2a3 (
tests/ut/cpp/a2a3/test_task_timing_slots.cpp, 9/9 pass):arg->descriptor propagation, slot bounds/sentinel, fixed-buffer layout +
offsets, cross-thread min/max reduction, incomplete-slot skip, per-run reset.
Build: a2a3sim + a5sim + a2a3 (onboard) all build; a2a3 cpput subset plus
the 9 new cases pass 100%.
e2e — a2a3, sim + onboard (
tests/st/task_timing/task_timing_slots/, L2swimlane OFF): 7 instances on sim and 7 on onboard, all pass, across both
tensormap_and_ringbufferandhost_build_graph:test_distinct_slots_emit_markersdispatch(slot1) >= finish(slot0)test_duplicate_slot_merges_windowtest_spmd_task_aggregates_across_threadsblock_num=8-> cross-thread min/max -> one complete slottest_mix_task_aggregates_across_subtasksMeasured
[STRACE](a2a3, ns, device clock):task_slot_0 ts=14600 dur=428500,task_slot_1 ts=444240 dur=49560->
dispatch(0) < finish(0) < dispatch(1)(dependency respected),whole-chain
finish(1) - dispatch(0) = 479200.task-submit, real H2D/D2H):task_slot_0 ts=270620 dur=8280,task_slot_1 ts=284260 dur=7979,whole-chain
21619— confirms H2D reset -> AICPU fold -> D2H tail readback-> emit on real silicon.
e2e — a5 host_build_graph, onboard (passing):
tests/st/task_timing/task_timing_a5hbg/test_task_timing_a5hbg.py::test_a5hbg_task_timing_slots_emit_markers,~17.5s — the legacy
add_taskpath end to end: bothtask_slot_0/task_slot_1markers with positive duration,
dispatch(slot1) >= finish(slot0), goldencorrect, swimlane off.
Static / lint — local pre-commit (all pass): clang-format, clang-tidy,
cpplint, ruff (check + format), pyright, plus header / English-only /
trailing-whitespace / large-file gates.
Integration: rebases cleanly onto current
main. The only conflict wasthe
Argstruct gaining both the upstream dispatchpredicate_and this PR'stask_timing_slot_member — resolved as a union; thePTO2TaskDescriptorsize /
packed_buffer_base-offsetstatic_asserts still hold.Not run locally — relies on CI: a5
tensormap_and_ringbuffer(its a5sime2e needs g++-15; this box has g++-12).