Add: early-dispatch for require_sync_start SPMD cohorts (a2a3 tmr) - #1304
Conversation
|
Important Review skippedAuto incremental reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughThis PR adds a gated sync_start drain protocol with a doorbell rendezvous mechanism to the a2a3 scheduler: a new ChangesSync-start drain rendezvous feature
Estimated code review effort: 4 (Complex) | ~75 minutes Sequence Diagram(s)sequenceDiagram
participant Dispatcher as scheduler_dispatch
participant Completion as SchedulerContext (drain)
participant Tracker as CoreTracker
participant Scheduler as PTO2SchedulerState
participant Payload as PTO2TaskPayload
Dispatcher->>Completion: enter_drain_mode(consumer, logical_block_num)
Completion->>Completion: reset drain_stage_go, drain_stage_done_mask, drain_running_staged
Completion->>Completion: elect drain worker, publish sync_start_pending
Completion->>Completion: handle_drain_mode(thread_idx, out_ack_wait_cycles)
Completion->>Tracker: get_mix_split_cluster_offset_states / core_offset_to_pending
Completion->>Completion: drain_stage_cores(...) parallel staging
Completion->>Payload: seed running_slot_count for gated drain
Completion->>Completion: sync on drain_stage_done_mask, reopen gate
Completion->>Scheduler: propagate_dispatch_fanin
Scheduler->>Scheduler: maybe_rendezvous_ring(slot_state)
Scheduler->>Scheduler: ring doorbells when running_slot_count matches staged core count
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 |
There was a problem hiding this comment.
Code Review
This pull request parallelizes the sync_start drain staging and introduces a rendezvous mechanism to prevent deadlocks during early dispatch. It also adds split placement for MIX shape blocks, updates completion and dispatch logic, and introduces profiling support for the drain phases alongside a new system test. The review feedback highlights two important issues: a potential out-of-bounds memory write in drain_stage_cores due to a missing bounds check on the core ID, and undefined behavior from signed integer overflow in the test kernel's spin loop.
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.
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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
`@src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_context.h`:
- Around line 351-353: The `handle_drain_mode` profiling output is mislabeled:
`out_ack_wait_cycles` in `scheduler_context.h` still reads like ack-barrier
time, but the implementation records stage/finalize wall time. Update the
parameter name and surrounding comments to match the actual value, or change the
implementation if ack-wait was intended, so the meaning is consistent across
`handle_drain_mode` and its caller-facing profiling data.
🪄 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: 49a3f1c5-87c4-4fb6-9f7d-19acdf8ffe44
📒 Files selected for processing (13)
src/a2a3/platform/include/common/l2_swimlane_profiling.hsrc/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2_types.hsrc/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.hsrc/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_cold_path.cppsrc/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/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_types.hsrc/a5/platform/include/common/l2_swimlane_profiling.hsrc/common/platform/shared/host/l2_swimlane_collector.cpptests/st/a2a3/tensormap_and_ringbuffer/spmd_sync_start_early_dispatch/kernels/aiv/kernel_spmd_write_slow.cpptests/st/a2a3/tensormap_and_ringbuffer/spmd_sync_start_early_dispatch/kernels/orchestration/spmd_sync_start_early_dispatch_orch.cpptests/st/a2a3/tensormap_and_ringbuffer/spmd_sync_start_early_dispatch/test_spmd_sync_start_early_dispatch.py
0d08ed3 to
efbaed6
Compare
efbaed6 to
0ef84c8
Compare
Aligns the a2a3 tensormap_and_ringbuffer scheduler with one occupancy model: two task sources (normal ready ▸ speculative early), and within each source sync_start > MIX > C/V, idle > pending. - Share one staging-order skeleton. dispatch_ready_tasks and try_early_dispatch each carried a hand-mirrored copy of the same order (MIX strict priority, IDLE before PENDING, cross-thread idle gating). Extract it into one run_staging_order helper parameterized by the per-shape stage callback; the two sources now differ only in that callback (normal = pickup/is_ready, early = gated). Behavior-identical. - Ready sync_start as strict Tier-0 in normal dispatch. Normal sync_start cohorts previously rode ready_queues[shape] mixed with regular tasks, so an AIC/AIV sync waited behind regular MIX. Park them in dedicated per-shape ready_sync_queues[] and drain as a Tier-0 pass before regular ready work. The pass reuses dispatch_shape unchanged (now parameterized by the queue array): same fits-local-inline vs stop-the-world drain decision, per-core MIX placement, and per-task head-start spacing. The regular tier passes ready_queues exactly as before, so regular dispatch is byte-identical; the early-dispatch gate now also drains the sync lanes so early never preempts a queued normal sync. Scope: a2a3 tensormap_and_ringbuffer only; regular dispatch path unchanged. Independent of the sync_start early-dispatch work (#1304) — this adds only the normal-source sync tier plus the order unification. Testing: a2a3sim 13 scene tests; a2a3 onboard 8 scene tests + spmd_sync_start_ stress x30 (0 failures, no 507018); full CI green.
642025f to
c4cf46f
Compare
Let require_sync_start SPMD cohorts pre-stage onto their cores while the producer still runs, then launch together via a unified doorbell, instead of waiting idle for the producer to fully retire. - Pre-stage gated: a flagged (allow_early_resolve) producer's dispatch bumps consumers' dispatch_fanin; a sync_start block that reaches fanin_actual_count stages onto its cores with not_ready=1 (spec_state STAGING) so it waits on a DATA_MAIN_BASE doorbell rather than launching early. - Rendezvous launch: running_slot_count counts staged gated CORES (shape- agnostic); when it reaches popcount(staged_core_mask) AND the producer has released, maybe_rendezvous_ring() rings every staged doorbell atomically so the cohort starts as one. - MIX split placement: a MIX cluster (1 AIC + 2 AIV) stages per-core into whatever running/pending slot each core has free, gated so partial placement is safe; count_mix_split_clusters gates availability. - Parallel drain: handle_drain_mode runs the stop-the-world sync_start drain across the dispatch threads (CAS-claimed block indices, per-thread local staged_core_mask merged with one fetch_or, reopen before propagate_fanin), replacing the serial stage that spread the cohort launch across the round. - pending_gated fix: a sync_start block stays gated after the producer releases (DISPATCHED), so pending_gated = STAGING || (DISPATCHED && sync_start); the STAGING-only test misclassified a released-but-still-gated pending block as a normal task and deadlocked waiting for an ack that never came. - Deadlock avoidance (early-candidate gate): propagate_dispatch_fanin no-ops until the producer is FULLY dispatched (next_block_idx == logical_block_num). A flagged SPMD producer with more blocks than cores (e.g. a 50-block AIC proj on 24 cores) dispatches in waves; if its first wave triggered a downstream MIX cohort to gate every core, the producer's remaining blocks would find no core, never complete, and the cohort's rendezvous -- which waits for that producer to release -- would never ring. Gating on full dispatch keeps the producer's cores committed before any consumer pre-occupies. next_block_idx is the gate because it advances on every placement path (normal dispatch, early-release, drain), covering all propagate_dispatch_fanin call sites. - Swimlane: Drain / DrainPrepare / DrainPublish sched phases (gated on PTO2_PROFILING + SCHED_PHASES) make the drain visible instead of a blind spot. - Docs: RUNTIME_LOGIC.md 8.6 (two-source x sync-tier dispatch model, drain + rendezvous, early-candidate gate, MIX per-core placement). - Tests: spmd_sync_start_early_dispatch drives a wide flagged producer into a gated sync_start consumer; spmd_sync_start_mix_spill exercises the MIX per-core pending-spill rendezvous path. Verified: qwen3-14B batch-16 decode completes at 1GB and 2GB ring heap (previously deadlocked every run -- Heap Exhausted at 1GB, running-stalled at 2GB); 14 a2a3 st tests pass (spmd / sync_start / mix / paged / stress). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Update wiring tests to model partial and complete block claims, and turn the sync-start scene into the AIC50-to-MIX24 resource-cycle regression. Keep every dispatched batch item in the deferred propagation list so a final AIV claim beyond MAX_CLUSTERS cannot miss the only post-gate retry. Document next_block_idx as a reservation counter rather than a publication counter.
- Claim one rendezvous owner for the complete staged doorbell pass - Expose fanout only after the cohort launch reaches COMPLETE - Retry producer-first rendezvous after gated drain finalization - Cover launch ownership and the no-promotion lost-wakeup ordering
Claim a payload-local ownership tombstone before an early sync task arms the global drain. Final producer release completes the STAGING-to-DISPATCHED transition, then transfers ready routing through an atomic READY handshake instead of enqueueing a second copy. Cancellation routes exactly once for both release orderings. Successful drains retain OWNER through COMPLETE until payload reuse so a delayed READY publisher cannot race past drain finalization. Add deterministic coverage for arm, cancellation, and split release/finalize interleavings.
b9564b7 to
a60fce0
Compare
…gbuffer (#1335) Follows PR #1327 (which synced hbg's scheduler to tmr's #1319 state). tmr then landed five more early-dispatch changes on the same subsystem; this ports them into host_build_graph's copies, preserving hbg's host-orchestration invariants (single ring, no execution-time reclaim, flat Runtime, no ACK-gate, host-side orchestrator). Absorbed from tmr: - #1304 Add early-dispatch for require_sync_start SPMD cohorts, incl. the drain_stage_cores parallel-drain refactor (handle_drain_mode two-stage). - #1326 Defer early dispatch until producer publication (correctness/race): published_block_count + record_published_blocks; propagate_dispatch_fanin is gated on published_block_count >= logical_block_num; the published_list ledger in dispatch and the stage_consumer_blocks two-sided publication/release handshake. - #1328 Shrink dispatch cold-write cost via a folded-gate src_payload (PTO2DispatchPayload: not_ready -> volatile uint64_t src_payload, 0 = ready, non-zero = gated source PTO2TaskPayload address; global_context to the tail; the PTO2_TASKPAYLOAD_*_OFFSET constants + static_asserts). build_payload writes args on the ready path and only src_payload on the gated path; the idle AICore fills args[] from src_payload during its doorbell wait; init-time prefill of async_ctx slab pointers/capacity + context-ptr args; deferred_slab reset moved off the dispatch path; software-pipelined prefetch of the next block. - #1329 Make early-dispatch doorbell ownership exclusive (correctness). - #1331 Serialize block range claims (correctness): every block-range claim now goes through slot_state->claim_block_range(...) instead of an inline next_block_idx CAS. Preserved host-orch divergences (not overwritten by tmr): PTO2_MAX_RING_DEPTH==1 with rings[]/ring_id removed; no advance_ring_pointers / reset_for_reuse / check_and_handle_consumed (completion via completed_tasks_, consumer wait keys on fanout_refcount); flat runtime->workers / runtime->func_id_to_addr_ (not runtime->dev.*); pto_runtime2_types.h keeps its dropped pto2_dispatch_payload.h include; hbg's AICore FIN-vs-swimlane-record ordering and its no-ACK-gate completion path (drains via completion-before-dispatch + backstop flushes); no device orchestrator (orchestrator_done_ absent, scheduler threads gate on runtime_init_ready_). Out of scope (unrelated tmr changes hbg still lacks, tracked for a follow-up): per-device scheduler timeout (get_scheduler_timeout_ms), PTO2_TENSOR_DATA_TIMEOUT MS-vs-cycles (#1189), stall sub-classification (#1182), runtime scope_tasks_cap (#1188). Verified: all 8 runtime targets build -Werror; a2a3sim host_build_graph scene suite 10/10; a2a3 onboard host_build_graph suite 10 passed / 1 skipped (paged_attention exercises the new claim_block_range serialization, drain_stage_cores drain, src_payload gated arg-fill, and published-block gating). tensormap_and_ringbuffer and the a5 runtimes are untouched.
…ve-sys#1319) Aligns the a2a3 tensormap_and_ringbuffer scheduler with one occupancy model: two task sources (normal ready ▸ speculative early), and within each source sync_start > MIX > C/V, idle > pending. - Share one staging-order skeleton. dispatch_ready_tasks and try_early_dispatch each carried a hand-mirrored copy of the same order (MIX strict priority, IDLE before PENDING, cross-thread idle gating). Extract it into one run_staging_order helper parameterized by the per-shape stage callback; the two sources now differ only in that callback (normal = pickup/is_ready, early = gated). Behavior-identical. - Ready sync_start as strict Tier-0 in normal dispatch. Normal sync_start cohorts previously rode ready_queues[shape] mixed with regular tasks, so an AIC/AIV sync waited behind regular MIX. Park them in dedicated per-shape ready_sync_queues[] and drain as a Tier-0 pass before regular ready work. The pass reuses dispatch_shape unchanged (now parameterized by the queue array): same fits-local-inline vs stop-the-world drain decision, per-core MIX placement, and per-task head-start spacing. The regular tier passes ready_queues exactly as before, so regular dispatch is byte-identical; the early-dispatch gate now also drains the sync lanes so early never preempts a queued normal sync. Scope: a2a3 tensormap_and_ringbuffer only; regular dispatch path unchanged. Independent of the sync_start early-dispatch work (hw-native-sys#1304) — this adds only the normal-source sync tier plus the order unification. Testing: a2a3sim 13 scene tests; a2a3 onboard 8 scene tests + spmd_sync_start_ stress x30 (0 failures, no 507018); full CI green.
…w-native-sys#1304) Let require_sync_start SPMD cohorts pre-stage while their producer is still running, then launch the whole cohort atomically through one rendezvous doorbell pass. Cohorts ride a shape-agnostic early_sync_start_queue drained as an all-or-nothing Tier-0; the rendezvous counts running cores (not blocks) so MIX pending-to-running promotion joins the same launch. - Full-reservation gate: propagate early dispatch only after the producer publishes every block (published_block_count >= logical_block_num), so a flagged SPMD producer commits all its slots before a downstream cohort can pre-occupy them. - Single-owner cohort launch: keep the sync cohort's staged_core_mask non-destructive and select exactly one ringer with a NONE -> RINGING -> COMPLETE launch latch. Publish COMPLETE (and expose fanout) only after every doorbell MMIO write plus a write barrier. - Serialized early-vs-ready routing: a payload-local OWNER|ARMED|READY|COMPLETE state gives the early drain and the final ready route exactly one owner. A ready event that observes OWNER does not enqueue a second copy; cancellation requeues the early candidate or transfers the recorded ready event; OWNER persists through COMPLETE so a delayed ready publication cannot race drain finalization. The owner also fixes the cohort's gated decision across a concurrent producer release. - Retry the rendezvous after drain finalization for producer-first ordering. Rebased onto main including hw-native-sys#1326 (full-publication gate) and hw-native-sys#1329 (destructive exactly-one doorbell ownership for ordinary non-sync early dispatch; the sync-start rendezvous instead uses its stable mask plus the launch latch). Adds 7 deterministic concurrency regression tests (arm, cancellation, split release/finalize interleavings, doorbell single-owner, drain-finalize retry) in tests/ut/cpp/a2a3/test_wiring.cpp, plus scheduler phase tracing and runtime-logic docs.
…gbuffer (hw-native-sys#1335) Follows PR hw-native-sys#1327 (which synced hbg's scheduler to tmr's hw-native-sys#1319 state). tmr then landed five more early-dispatch changes on the same subsystem; this ports them into host_build_graph's copies, preserving hbg's host-orchestration invariants (single ring, no execution-time reclaim, flat Runtime, no ACK-gate, host-side orchestrator). Absorbed from tmr: - hw-native-sys#1304 Add early-dispatch for require_sync_start SPMD cohorts, incl. the drain_stage_cores parallel-drain refactor (handle_drain_mode two-stage). - hw-native-sys#1326 Defer early dispatch until producer publication (correctness/race): published_block_count + record_published_blocks; propagate_dispatch_fanin is gated on published_block_count >= logical_block_num; the published_list ledger in dispatch and the stage_consumer_blocks two-sided publication/release handshake. - hw-native-sys#1328 Shrink dispatch cold-write cost via a folded-gate src_payload (PTO2DispatchPayload: not_ready -> volatile uint64_t src_payload, 0 = ready, non-zero = gated source PTO2TaskPayload address; global_context to the tail; the PTO2_TASKPAYLOAD_*_OFFSET constants + static_asserts). build_payload writes args on the ready path and only src_payload on the gated path; the idle AICore fills args[] from src_payload during its doorbell wait; init-time prefill of async_ctx slab pointers/capacity + context-ptr args; deferred_slab reset moved off the dispatch path; software-pipelined prefetch of the next block. - hw-native-sys#1329 Make early-dispatch doorbell ownership exclusive (correctness). - hw-native-sys#1331 Serialize block range claims (correctness): every block-range claim now goes through slot_state->claim_block_range(...) instead of an inline next_block_idx CAS. Preserved host-orch divergences (not overwritten by tmr): PTO2_MAX_RING_DEPTH==1 with rings[]/ring_id removed; no advance_ring_pointers / reset_for_reuse / check_and_handle_consumed (completion via completed_tasks_, consumer wait keys on fanout_refcount); flat runtime->workers / runtime->func_id_to_addr_ (not runtime->dev.*); pto_runtime2_types.h keeps its dropped pto2_dispatch_payload.h include; hbg's AICore FIN-vs-swimlane-record ordering and its no-ACK-gate completion path (drains via completion-before-dispatch + backstop flushes); no device orchestrator (orchestrator_done_ absent, scheduler threads gate on runtime_init_ready_). Out of scope (unrelated tmr changes hbg still lacks, tracked for a follow-up): per-device scheduler timeout (get_scheduler_timeout_ms), PTO2_TENSOR_DATA_TIMEOUT MS-vs-cycles (hw-native-sys#1189), stall sub-classification (hw-native-sys#1182), runtime scope_tasks_cap (hw-native-sys#1188). Verified: all 8 runtime targets build -Werror; a2a3sim host_build_graph scene suite 10/10; a2a3 onboard host_build_graph suite 10 passed / 1 skipped (paged_attention exercises the new claim_block_range serialization, drain_stage_cores drain, src_payload gated arg-fill, and published-block gating). tensormap_and_ringbuffer and the a5 runtimes are untouched.
#1304 added the sync_start stop-the-world drain phase kinds (Drain, DrainPrepare, DrainPublish) to the device and host emit sides (l2_swimlane_profiling.h + host l2_swimlane_collector.cpp) but never taught the Python renderer about them. swimlane_converter.py drops any phase outside its sched-lane whitelist, so a sync_start cohort's launch — emitted as drain_publish, not dispatch — was silently filtered out and the scheduler dispatch work for sync_start SPMD cohorts vanished from the trace. - Add drain/drain_prepare/drain_publish to phase_colors (cq_build family; the outer drain bar time-contains the two inner staging passes so Perfetto nests them by depth) - Add the three phases to the sched-lane render whitelist Verified on a2a3 silicon at l2_swimlane_level=4: spmd_sync_start, spmd_sync_start_early_dispatch, and spmd_sync_start_stress now render the drain phases (rendered phase histogram matches the raw record kinds).
Bring a5 tensormap_and_ringbuffer in line with a2a3 early-dispatch (hw-native-sys#989/hw-native-sys#1079/hw-native-sys#1285/hw-native-sys#1288/hw-native-sys#1297/hw-native-sys#1304/hw-native-sys#1326/hw-native-sys#1329/hw-native-sys#1336/hw-native-sys#1405/hw-native-sys#1340): prepare/publish, src_payload gate + DMB high32 doorbell (sim uses __atomic_load_n ACQUIRE), direct-only eligibility, spare-slot Phase 4b, sync_start early rendezvous, and dep_gen early_dispatch truth. Document the a5 model in RUNTIME_LOGIC §8.6. Add a5 STs: plain early_dispatch plus sync_start early_dispatch and mix_spill ports (EarlyOn/Off) for board validation.
) * feat(a5/tmr): port allow_early_resolve / early-dispatch from a2a3 Bring a5 tensormap_and_ringbuffer in line with a2a3 early-dispatch (#989/#1079/#1285/#1288/#1297/#1304/#1326/#1329/#1336/#1405/#1340): prepare/publish, src_payload gate + DMB high32 doorbell (sim uses __atomic_load_n ACQUIRE), direct-only eligibility, spare-slot Phase 4b, sync_start early rendezvous, and dep_gen early_dispatch truth. Document the a5 model in RUNTIME_LOGIC §8.6. Add a5 STs: plain early_dispatch plus sync_start early_dispatch and mix_spill ports (EarlyOn/Off) for board validation. * test(a5): drop local early-dispatch ST ports from this PR Remove the a5-only early_dispatch / sync_start early ST scenes added for board experiments; keep the runtime port focused for CI. * chore(a5): drop unrelated files that broke pre-commit on the PR Remove accidental local helpers, migrate patches, and the acc_c2v example that were committed with the UT payload-pool fix. Keep the early-dispatch runtime port, a2a3-ported sync_start STs, and a5 UT payload/task pool binding.
Summary
Let
require_sync_startSPMD cohorts pre-stage while their producer isrunning, then launch the whole cohort through one rendezvous doorbell pass.
the same rendezvous.
staged_core_masknon-destructive. ANONE -> RINGING -> COMPLETElaunch latch selects exactly one ringer.COMPLETEonly after all doorbell MMIO writes and a write barrier;downstream early dispatch stays hidden until then.
OWNER|ARMED|READY|COMPLETEstate to give the earlydrain and final ready routing exactly one owner. The owner also fixes the
cohort's gated decision across a concurrent producer release.
post-publish list.
This branch is rebased onto current
main, including #1326 and #1329.#1326 supplies the full-publication gate. #1329 retains destructive
exactly-one doorbell ownership for ordinary non-sync early dispatch; the
sync-start rendezvous deliberately uses its stable mask plus the launch latch.
The independent normal-dispatch/late-stager block-range claim race is not
included here and is being submitted separately because it exists on
mainwithout #1304.Why the final ownership state is needed
Fixing the launch latch alone leaves a second race: an early-sync queue worker
can pop and arm a task while producer completion simultaneously routes that
same task into
ready_sync_queues. Both paths then own one slot.The payload-local state transfers the route once:
A ready event that observes
OWNERdoes not enqueue the task. Cancellationeither requeues the early candidate or transfers the recorded ready event.
OWNERpersists through completion until payload reuse, covering a delayedready publication.
Testing
pre-commit run --from-ref main --to-ref HEAD: passed20 generated tokens, reference four-scope prefill kernel:
4/4 passed on the PR-only scheduler tree; one run is
provenance-locked to final head
a60fce05block-range CAS: 10/10 passed
14/14 passed on the preceding head
The two-scope cached prefill kernel is not scheduler evidence: its known manual
scope regression exhausts the ring heap. Hardware acceptance above uses the
four-scope reference kernel.