Fix: defer early dispatch until producer publication - #1326
Merged
ChaoWao merged 1 commit intoJul 11, 2026
Merged
Conversation
Track logical block publication after payload and MMIO visibility across normal dispatch, early staging, and sync-start drain paths. Gate dispatch-fanin propagation on the full published block count and retry it from the final stager when producer release wins the race. Add wiring coverage for partial publication and once-only propagation.
📝 WalkthroughWalkthroughEarly-dispatch metadata now tracks published logical blocks. Scheduler publication paths update this counter, and fanin propagation waits until flagged producers publish all blocks. Wiring tests cover the new gating and existing producer-flag conditions. ChangesEarly-dispatch publication gating
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant SchedulerContext
participant PTO2TaskPayload
participant propagate_dispatch_fanin
participant ConsumerSlot
SchedulerContext->>PTO2TaskPayload: record published block count
SchedulerContext->>propagate_dispatch_fanin: propagate dispatch fanin
propagate_dispatch_fanin->>PTO2TaskPayload: check published_block_count
propagate_dispatch_fanin->>ConsumerSlot: update early-dispatch fanin when fully published
Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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 refactors the early-dispatch mechanism to ensure that propagation only occurs after all logical blocks of a flagged producer are fully published. It introduces an atomic published_block_count to track publication progress and updates the scheduler's dispatch and completion paths to record published blocks. The unit tests have been updated to verify this new behavior. Feedback on the changes highlights a potential buffer overflow risk in scheduler_dispatch.cpp, where the newly introduced published_list and published_counts arrays are populated without defensive bounds checks.
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[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tests/ut/cpp/a2a3/test_wiring.cpp (1)
367-375: 🩺 Stability & Availability | 🔵 Trivial | 🏗️ Heavy liftAdd coverage for the release-first handshake path. This test covers count gating and the once-only fanout, but it still misses the opposite ordering where release flips
DISPATCHEDbefore the last stager publishes. Add a deterministic case for that interleaving so both sides of the handshake stay covered.🤖 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 `@tests/ut/cpp/a2a3/test_wiring.cpp` around lines 367 - 375, Add a deterministic release-first handshake test near the existing count-gating assertions in the wiring test: exercise the path where release sets DISPATCHED before the final stager publishes, then publish the remaining block and invoke propagation, asserting fan-in completion and exactly-once dispatch. Reuse the existing scheduler, producer, payload, and handshake APIs while preserving the current coverage for the stager-first ordering.
🤖 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 `@tests/ut/cpp/a2a3/test_wiring.cpp`:
- Around line 367-375: Add a deterministic release-first handshake test near the
existing count-gating assertions in the wiring test: exercise the path where
release sets DISPATCHED before the final stager publishes, then publish the
remaining block and invoke propagation, asserting fan-in completion and
exactly-once dispatch. Reuse the existing scheduler, producer, payload, and
handshake APIs while preserving the current coverage for the stager-first
ordering.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 3e97c185-6004-402d-843c-ba7be128626c
📒 Files selected for processing (5)
src/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_completion.cppsrc/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_dispatch.cpptests/ut/cpp/a2a3/test_wiring.cpp
7 tasks
4 tasks
ChaoWao
added a commit
that referenced
this pull request
Jul 11, 2026
…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.
3 tasks
ChaoWao
added a commit
that referenced
this pull request
Jul 11, 2026
…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.
doraemonmj
pushed a commit
to doraemonmj/simpler_wc
that referenced
this pull request
Jul 14, 2026
) Track logical block publication after payload and MMIO visibility across normal dispatch, early staging, and sync-start drain paths. Gate dispatch-fanin propagation on the full published block count and retry it from the final stager when producer release wins the race. Add wiring coverage for partial publication and once-only propagation.
doraemonmj
pushed a commit
to doraemonmj/simpler_wc
that referenced
this pull request
Jul 14, 2026
…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.
doraemonmj
pushed a commit
to doraemonmj/simpler_wc
that referenced
this pull request
Jul 14, 2026
…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.
yanghaoran29
added a commit
to yanghaoran29/simpler
that referenced
this pull request
Jul 21, 2026
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.
ChaoWao
pushed a commit
that referenced
this pull request
Jul 21, 2026
) * 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.
ChaoZheng109
added a commit
to ChaoZheng109/simpler
that referenced
this pull request
Sep 7, 2026
…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.
4 tasks
poursoul
pushed a commit
that referenced
this pull request
Sep 7, 2026
…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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Testing
b9564b7e, the exact Qwen3 PR691-kernel workloadcompletes 3/3 at a 2GB ring heap
deadlock at
current=7789, last_alive=5255, with the head task at 82/109consumers
This is separate from PR #1304 because the partial-publication race is already
present on main. PR #1304 carries the independent sync-start cohort launch fix.