refactor(a2a3/early-dispatch): align acquisition and ordering with normal dispatch#1297
Open
ChaoZheng109 wants to merge 1 commit into
Open
Conversation
📝 WalkthroughWalkthroughThis PR converts the scheduler's single shared early-dispatch queue into an array of queues partitioned by ChangesPer-shape early-dispatch queue refactor
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Scheduler as SchedulerContext
participant EarlyQueue as early_dispatch_queues[shape]
participant CoreTracker
participant Consumer
Scheduler->>EarlyQueue: drain queue for shape
EarlyQueue-->>Scheduler: staged consumer
Scheduler->>Consumer: check staging state (PTO2_SPEC_STAGING)
Scheduler->>CoreTracker: get_free_slot_states() for phase
CoreTracker-->>Scheduler: free core bitmask
alt free cores available
Scheduler->>Consumer: CAS-claim block range
Scheduler->>CoreTracker: stage_consumer_blocks
else no free cores
Scheduler->>EarlyQueue: re-push consumer
end
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 |
|
Warning Gemini encountered an error creating the review. You can try again by commenting |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.h (1)
1060-1065: 🩺 Stability & Availability | 🟡 Minor | ⚡ Quick winSilent drop on early-dispatch queue overflow.
push(c)'s return value is discarded here, unlike the analogous re-push inearly_dispatch_shape(scheduler_dispatch.cpp) which checks and logs on failure. Functionally this is safe —try_speculative_releasestill falls back to normal ready-queue dispatch when the consumer was never actually staged — but a full queue here silently loses the early-dispatch optimization with no diagnostic trail.🔍 Proposed fix to log on overflow
- early_dispatch_queues[static_cast<int32_t>(shape)].push(c); + if (!early_dispatch_queues[static_cast<int32_t>(shape)].push(c)) { + LOG_INFO_V9( + "[SPEC] early-dispatch queue full on initial push, consumer=%" PRId64, + static_cast<int64_t>(c->task->task_id.raw) + ); + }🤖 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/pto_scheduler.h` around lines 1060 - 1065, The early-dispatch enqueue in try_speculative_release silently ignores a failed push on early_dispatch_queues, so add a check for the push(c) result and log a diagnostic when the queue is full. Use the existing early-dispatch failure handling in early_dispatch_shape and the scheduler logging path as the model, and keep the current fallback behavior to normal ready-queue dispatch unchanged.
🤖 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.
Outside diff comments:
In `@src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.h`:
- Around line 1060-1065: The early-dispatch enqueue in try_speculative_release
silently ignores a failed push on early_dispatch_queues, so add a check for the
push(c) result and log a diagnostic when the queue is full. Use the existing
early-dispatch failure handling in early_dispatch_shape and the scheduler
logging path as the model, and keep the current fallback behavior to normal
ready-queue dispatch unchanged.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: f0b07d6f-b667-46f8-a95d-89fe0add4984
📒 Files selected for processing (6)
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_context.hsrc/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_dispatch.cppsrc/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_types.hsrc/a2a3/runtime/tensormap_and_ringbuffer/runtime/shared/pto_runtime2_init.cpp
…rmal dispatch Unify the early-dispatch (ed) path in the a2a3 tensormap_and_ringbuffer scheduler with the normal-dispatch (nd) path so the scheduler's two task sources share the same primitives — pop sizing, gate helpers, staging, ordering, and progress accounting. - Size the drain pop to this thread's free-slot count instead of a hardcoded 8, matching nd's cores.count() sizing. New CoreTracker::get_free_slot_states() primitive; has_any_free_slot() re-expressed on top of it. - Split the single shape-mixed early_dispatch_queue into per-shape early_dispatch_queues[3], mirroring ready_queues[]; candidates bucket by active_mask.to_shape() at push time, so the drain pops per shape with the shape known as the queue index (reserve/init/wire/reset/destroy loop over the three queues, symmetric with ready_queues[]). - Order ed like nd: MIX strict priority, IDLE before PENDING, cross-thread idle gating (MIX-IDLE > c/v-IDLE > MIX-PEND > c/v-PEND). New early_dispatch_shape() (the ed analog of dispatch_shape) reuses has_idle_in_other_threads and stages onto only the phase's bucket; new has_residual_early_mix() for the MIX gate (has_residual_mix reads the normal MIX ready queue, empty when ed runs). - Move ed's gate (off-PMU, has_any_free_slot, all ready_queues empty) inside try_early_dispatch, mirroring how dispatch_ready_tasks owns its own PMU/gating, so the caller is a single unconditional call; set made_progress / try_pushed when staging, so an iteration that only early-stages is no longer miscounted as idle. - Rename try_speculative_early_dispatch -> try_early_dispatch. a2a3 only (a5 has no early-dispatch path). ed stays disabled under PMU. Tested on a2a3sim across ~15 SPMD/MIX/dependency-heavy scene tests — no deadlock, golden drift, or timeout/stall regression.
ChaoZheng109
force-pushed
the
refactor/unify-scheduler-dispatch
branch
from
f886b58 to
d798a6e
Compare
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.
What
Unifies the early-dispatch (ed) path in the a2a3
tensormap_and_ringbufferscheduler with the normal-dispatch (nd) path so the scheduler's two task sources
share the same primitives — pop sizing, gate helpers, staging, ordering, and
progress accounting. This change:
Sizes the drain pop to free-slot count, not a fixed 8 — ed drained a
hardcoded 8 consumers per pass while nd sizes its pop to available capacity.
Both now use the same
BitStates.count()sizing (newCoreTracker::get_free_slot_states()primitive;has_any_free_slot()re-expressed on top of it).
Splits the early-dispatch queue per resource shape — the single
shape-mixed
early_dispatch_queuebecomesearly_dispatch_queues[3],mirroring the per-shape
ready_queues[]. Candidates bucket byactive_mask.to_shape()at push time; the drain pops per shape with the shapeknown as the queue index (reserve/init/wire/reset/destroy loop over the three
queues, symmetric with
ready_queues[]).Gives ed MIX-priority two-stage ordering — the same ordering nd uses: MIX
strict priority, IDLE stage before PENDING stage, cross-thread idle gating —
MIX-IDLE ▶ c/v-IDLE ▶ MIX-PEND ▶ c/v-PEND. New
early_dispatch_shape()(the edanalog of
dispatch_shape) reuseshas_idle_in_other_threadsand stages ontoonly the phase's bucket; new
has_residual_early_mix()for the MIX gate(
has_residual_mixreads the normal MIX ready queue, empty when ed runs).Encapsulates gating + updates progress flags — ed's gate (off-PMU,
has_any_free_slot, all ready_queues empty) moves insidetry_early_dispatch(mirroring how
dispatch_ready_tasksowns its own PMU/gating), so the calleris a single unconditional call;
made_progress/try_pushedare set whenstaging, so an iteration that only early-stages is no longer miscounted as
idle.
Renames
try_speculative_early_dispatch→try_early_dispatch.Why
The two dispatch sources had diverged: nd sized pops to capacity, ordered by
MIX-priority / IDLE-then-PENDING, and updated the loop's progress flags from a
single encapsulated call; ed used a flat fixed-8 drain over a mixed queue with
caller-side gating and no progress-flag updates. Aligning them makes the
scheduler's two task sources behave consistently and share the same primitives.
Scope / notes
tensormap_and_ringbufferruntime has no early-dispatchpath.
IDLE-only under PMU. Aligning ed to IDLE-only under PMU was considered and
dropped — a gated core polls DMB (MMIO) while waiting for its doorbell, which
would perturb the single-issue measurement windows PMU relies on.
dispatch_shape): early staging isgated pre-stage + doorbell, distinct from normal immediate dispatch.
early_dispatchL2-swimlane bar (name and semantics) is unchanged, so no docupdate is needed.
Testing
a2a3sim: ~15 SPMD / MIX / dependency-heavy scene tests pass
(
spmd_multiblock_mix/_aiv,spmd_starvation,spmd_sync_start/_stress,paged_attention family,
alternating_matmul_add,fanin_lookup_perf,mixed_example,dummy_task) — no deadlock, no golden drift, no timeout/stallregression. Built + resynced + tested against latest upstream/main. a2a3 onboard
not run (box is a5 silicon); sim is the silicon-agnostic path — onboard left to
CI.