Sync: host_build_graph scheduler up a generation to tensormap_and_ringbuffer#1327
Conversation
|
Warning Review limit reached
Next review available in: 32 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (18)
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 parallelizes the per-core AICore handshake across AICPU threads to reduce initialization latency, restructures the dispatch mechanism with per-shape early-dispatch and sync-start queues, and moves fanout wiring inline to the host submit path. The review feedback highlights critical issues, including potential deadlocks in the parallel initialization barrier when threads fail or have out-of-range affinity indices, a potential null pointer dereference and use-after-free in the fanin wiring logic, and a logic error in classify_mix_cluster that could overwrite occupied pending slots.
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.
ChaoWao
force-pushed
the
sync-hbg-scheduler-from-tmr
branch
from
4da20c4 to
d6c4513
Compare
…gbuffer hbg's scheduler was a copy of tmr@fb28fd71 and had fallen a full generation behind. This absorbs tmr's post-fork scheduler evolution while preserving hbg's host-orchestration invariants (single ring, no execution-time reclaim, flat Runtime, host-side orchestrator). Absorbed from tmr: - Remove PTO2LocalReadyBuffer local-first dispatch and its local_bufs threading (hw-native-sys#1245); pop/dispatch now take PTO2ReadyQueue* directly. - Ready-claim completion protocol: ready_state / mark_completed / mark_any_subtask_deferred (replaces the any_subtask_deferred bool). - Rename speculative -> early_dispatch and remove the auto-chain (spec_chain_active/depth, PTO2_SPEC_CHAIN_MAX); move allow_early_resolve from payload onto slot_state. - Per-shape early_dispatch_queues[] (was a single queue) + ready_sync_queues[] Tier-0 sync_start staging (hw-native-sys#1319); run_staging_order template; try_early_dispatch / early_dispatch_shape; has_any_free_slot predicates. - MIX subtasks placed per-core instead of uniform cluster placement (hw-native-sys#1308). - Remove the on-device wiring subsystem (PTO2SpscQueue / WiringState / drain_wiring_queue / wire_task, hw-native-sys#1263): the host orchestrator now wires fanout adjacency inline during submit and seeds readiness via the scheduler's push_ready_routed / route_ready_once, with no device-side wiring queue. - Single-round-trip / parallel AICore handshake (hw-native-sys#1279/hw-native-sys#1310): handshake_all_cores -> pre_handshake_init + per-thread handshake_partition + post_handshake_init, with the matching aicpu_executor boot restructure and AICore-worker protocol. Drop the now-dead two-round-trip Handshake fields aicpu_regs_ready / aicore_regs_ready (aligning the struct to tmr's). - Drop the per-run per-core memset in scheduler init/deinit (hw-native-sys#1212); gate dep_gen init under SIMPLER_DFX; compile out pmu_active when SIMPLER_DFX is off; remove the redundant second completion poll; remove the dead device-orch core transition (handle_core_transition / reassign_cores_for_all_threads / orch_to_sched_). Remove vestigial device-orch synchronization: with host orchestration the orchestrator runs to completion on the host before the device boots, so the device-side orchestrator_done_ flag is always set before the scheduler threads start (they gate on runtime_init_ready_, whose release/acquire already publishes total_tasks_). Removed orchestrator_done_, the always-true completion guard, and the never-called wait_for_orchestration_done_before_dispatch / orchestration_done() carried over from tmr's concurrent device-orch model. 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; hbg's swimlane behavior (no on_aicore_ack hook — the just-filled buffer drains via the completion-before-dispatch invariant and the next-rotation / run-end backstop flushes). Skipped as not applicable to hbg: stall-classification publishing (shared-memory lacks the sched_stall_* fields; last_task_alive semantics differ); per-device scheduler timeout (aicpu_device_config.h / get_scheduler_timeout_ms absent); runtime->dev.* Runtime nesting. get_reg_ptr weak host fallback: route_ready_once transitively ODR-uses the early-dispatch doorbell inline (ring_one_doorbell -> get_reg_ptr) in the host lib, but no core is gated during host graph-build so it never fires; a weak hidden fallback satisfies the loader (same pattern as get_sys_cnt_aicpu). 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 handshake, route_ready_once wiring, and ready-claim path). tensormap_and_ringbuffer and the a5 runtimes are untouched.
…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.
1 participant
What
hbg's scheduler was a copy of
tmr@fb28fd71and had fallen a full generationbehind. This PR absorbs tmr's post-fork scheduler evolution into
host_build_graphwhile preserving hbg's host-orchestration invariants.Absorbed from tmr
PTO2LocalReadyBufferlocal-first dispatch +local_bufsthreading.ready_state/mark_completed/mark_any_subtask_deferred.allow_early_resolvepayload→slot_state.early_dispatch_queues[]+ready_sync_queues[]Tier-0 staging;run_staging_order;try_early_dispatch.push_ready_routed/route_ready_once(no device wiring queue).handshake_partition+ executor boot restructure + AICore-worker protocol).SIMPLER_DFX; compile outpmu_active; remove the second completion poll and the dead device-orch core transition.Preserved host-orch divergences (not overwritten)
PTO2_MAX_RING_DEPTH==1(norings[]/ring_id); no execution-time reclaim(
completed_tasks_+fanout_refcount, noadvance_ring_pointers/check_and_handle_consumed);flat
runtime->workers; hbg swimlane behavior (noon_aicore_ackhook — drains viacompletion-before-dispatch + backstop flushes).
Skipped (not applicable to hbg)
Stall-classification publishing (no
sched_stall_*fields); per-device schedulertimeout (
aicpu_device_config.habsent);runtime->dev.*nesting.Verification
-Werror.host_build_graphscene suite: 10/10 pass.host_build_graphsuite: 10 passed / 1 skipped (paged_attention exercises the new handshake,route_ready_oncewiring, and ready-claim path).tensormap_and_ringbufferand the a5 runtimes are untouched.