Conversation
📝 WalkthroughWalkthroughThe change adds per-submit early-dispatch launch tracking, atomic staged-doorbell claiming, forced gating for consumer staging, guarded fan-in propagation, and wiring tests for ownership, ordering, release, and register writes. ChangesEarly-dispatch scheduling
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant SchedulerContext
participant PTO2SchedulerState
participant PTO2TaskPayload
participant DoorbellRegisters
SchedulerContext->>PTO2TaskPayload: prepare gated consumer payloads
SchedulerContext->>PTO2SchedulerState: publish and propagate fan-in
PTO2SchedulerState->>PTO2TaskPayload: claim launch and staged bits
PTO2SchedulerState->>DoorbellRegisters: ring claimed doorbells
PTO2SchedulerState->>PTO2TaskPayload: mark launch complete
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 refactors the early dispatch synchronization mechanism to ensure strict doorbell ownership and prevent race conditions between concurrent stagers and the release path. It introduces a new PTO2EarlyDispatchLaunchState state machine, destructively splits the staged core mask between release and late-stager paths, and updates payload building to keep claimed early-stage ranges gated. Additionally, comprehensive unit tests and test stubs have been added to verify the correctness of the new synchronization handshake. No review comments were provided, so there is no feedback to address.
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.
- Atomically split staged-core bits between release and late stagers - Serialize release with a NONE/RINGING/COMPLETE launch state - Ring captured local handles before publishing late-staged blocks - Retry fanout after both publication and release race orders - Keep preclaimed ranges gated across concurrent producer release
ff87722 to
60e59a4
Compare
…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 #1326 (full-publication gate) and #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 (#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.
…1329) - Atomically split staged-core bits between release and late stagers - Serialize release with a NONE/RINGING/COMPLETE launch state - Ring captured local handles before publishing late-staged blocks - Retry fanout after both publication and release race orders - Keep preclaimed ranges gated across concurrent producer release
…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.
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.
…enum The per-task progress byte encoded a linear progression through bit-containment values (0x0 -> 0x2 -> 0x3) so a lock-free fetch_or would double as a monotone-max. That was necessary because the PUBLISHED bookkeeping ran after the final MMIO token write while the FIN -> completion chain forks off that same token write, leaving the two writes causally unordered: a plain store of a sequential value could let a late publish regress an already-completed byte and livelock the wake-list sentinel protocol. Establish the order by construction instead. record_published_blocks splits into account_published_blocks, called before a batch's token writes and returning whether this caller reached the task's total, and the existing seal_ed_publish_list, called after the flush. The PUBLISHED store therefore precedes every token the task emits, and those tokens' FINs gate the all-FIN completion, so PUBLISHED < token < FIN < COMPLETED holds regardless of how sibling publishers interleave. Cores are claimed at prepare time, before any bookkeeping, so a consumer staged off the earlier PUBLISHED cannot occupy cores the producer's in-flight blocks still need. This retires the ordering constraint hw-native-sys#1326 and hw-native-sys#1329 recorded in stage_consumer_blocks, which required a released block to ring before contributing to the publication count. Doorbell ownership does not depend on that count: it is a two-sided seq_cst handshake between the staged_core_mask fetch_or and the release path's early_dispatch_state store, and every staged bit is claimed exactly once by whichever side observes the other. A consumer released by the earlier publication cannot take cores this task's blocks already hold, and holds nothing the rings wait on. The comment there now states that as the current invariant. With the ordering guaranteed, the byte becomes a ChipTaskState written with plain stores: PENDING -> PUBLISHED -> COMPLETED. Readers use ordered comparisons, and COMPLETED > PUBLISHED is what lets a tracked producer that never publishes (DUMMY, predicate-retired) release its publish-list waiters through the completion store alone. Because every value but PENDING and PUBLISHED now reads as completed, the array's comment states the init-on-write discipline the reads depend on. The array is renamed progress_flags -> task_states to match, and TASK_FLAG_* / is_completion_flag_set / is_publish_flag_set give way to is_completed / is_published / store_completed / store_published / reset_task_state. The slot-resident task_state mirror is unchanged: it stays PENDING or COMPLETED and remains the in-graph readiness truth. tensormap_and_ringbuffer is untouched and keeps its own record_published_blocks: it has no such byte array, publishes no state, and derives readiness from fanin_refcount under the push model. Also corrects the ChipTaskSlotState header comment, which claimed the struct is "NOT in shared memory" while a GLOBAL task's slot lives in the SM image's storage segment.
…enum (#2130) The per-task progress byte encoded a linear progression through bit-containment values (0x0 -> 0x2 -> 0x3) so a lock-free fetch_or would double as a monotone-max. That was necessary because the PUBLISHED bookkeeping ran after the final MMIO token write while the FIN -> completion chain forks off that same token write, leaving the two writes causally unordered: a plain store of a sequential value could let a late publish regress an already-completed byte and livelock the wake-list sentinel protocol. Establish the order by construction instead. record_published_blocks splits into account_published_blocks, called before a batch's token writes and returning whether this caller reached the task's total, and the existing seal_ed_publish_list, called after the flush. The PUBLISHED store therefore precedes every token the task emits, and those tokens' FINs gate the all-FIN completion, so PUBLISHED < token < FIN < COMPLETED holds regardless of how sibling publishers interleave. Cores are claimed at prepare time, before any bookkeeping, so a consumer staged off the earlier PUBLISHED cannot occupy cores the producer's in-flight blocks still need. This retires the ordering constraint #1326 and #1329 recorded in stage_consumer_blocks, which required a released block to ring before contributing to the publication count. Doorbell ownership does not depend on that count: it is a two-sided seq_cst handshake between the staged_core_mask fetch_or and the release path's early_dispatch_state store, and every staged bit is claimed exactly once by whichever side observes the other. A consumer released by the earlier publication cannot take cores this task's blocks already hold, and holds nothing the rings wait on. The comment there now states that as the current invariant. With the ordering guaranteed, the byte becomes a ChipTaskState written with plain stores: PENDING -> PUBLISHED -> COMPLETED. Readers use ordered comparisons, and COMPLETED > PUBLISHED is what lets a tracked producer that never publishes (DUMMY, predicate-retired) release its publish-list waiters through the completion store alone. Because every value but PENDING and PUBLISHED now reads as completed, the array's comment states the init-on-write discipline the reads depend on. The array is renamed progress_flags -> task_states to match, and TASK_FLAG_* / is_completion_flag_set / is_publish_flag_set give way to is_completed / is_published / store_completed / store_published / reset_task_state. The slot-resident task_state mirror is unchanged: it stays PENDING or COMPLETED and remains the in-graph readiness truth. tensormap_and_ringbuffer is untouched and keeps its own record_published_blocks: it has no such byte array, publishes no state, and derives readiness from fanin_refcount under the push model. Also corrects the ChipTaskSlotState header comment, which claimed the struct is "NOT in shared memory" while a GLOBAL task's slot lives in the SM image's storage segment.
Summary
NONE -> RINGING -> COMPLETE, and expose fanout only after release-owned doorbells and full block publication are visible.Context
The non-sync early-dispatch path on
mainallowed release and a late stager to ring the same bit. The late stager then reread a shared per-core table that could already belong to a later task.This is an independent main bug found while auditing #1304. It was not required for the clean Qwen 3/3 pass: that minimal run used merged #1326 plus the launch serialization commit on #1304. PR #1304 should reuse this launch state when it rebases.
Testing
pre-commit run --files ...tensormap_and_ringbuffer: 43/43 applicable sim cases passedasync_notify_demo: 1/1 passed undertask-submit --device 0,1(task_20260711_034229_20976230978)The sim split avoids an existing harness mismatch:
async_notify_demoignores the CLI platform and hardcodes onboardrun("a2a3", [0,1]), so it was validated separately under the device lock.