Refactor: hbg unifies the task progress byte into a sequential state enum - #2130
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: Team Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Team Run ID: 📒 Files selected for processing (24)
Included review availability: Your plan provides up to 1 included review per hour; 0 remain after this review. 📝 WalkthroughWalkthroughThe runtime replaces ChangesTask-state publication protocol
Estimated code review effort: 4 (Complex) | ~45 minutes Merge Risk: ⚪ Minimal · up to The sequential task-state migration preserves the intended publication ordering across the reviewed scheduler and host paths, with no unresolved merge-blocking risk identified. Sequence Diagram(s)sequenceDiagram
participant SchedulerDispatch
participant SchedulerState
participant SharedMemoryTaskHeader
participant DeviceCores
SchedulerDispatch->>SchedulerState: account_published_blocks
SchedulerState->>SharedMemoryTaskHeader: store_published
SchedulerDispatch->>DeviceCores: write payload and MMIO tokens
SchedulerDispatch->>SchedulerState: seal_ed_publish_list
DeviceCores->>SharedMemoryTaskHeader: complete task
SharedMemoryTaskHeader->>SchedulerState: store_completed
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Full details: Docstring CoverageExplanation Docstring coverage is 54.29% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 70 functions across 22 files. (2 skipped: 2 unsupported.) 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 |
30e66bb to
6a6079c
Compare
|
Thanks — six of the seven were real and are fixed in 6a6079c. One has a factually wrong premise; details below. 1. The #1326/#1329 ordering comment — valid, and the analysis holds. You are right that the PR silently retired a constraint two fix commits wrote down. It is now argued explicitly, in the commit message and at the site. The reason it is safe: doorbell ownership never involved 2. Rename drift — valid, all fixed. The layout diagram in 3. tmr divergence — the premise is wrong, but a note was still worth adding. tmr has no 4. Unwritten-byte reads — valid, documented. 5. No direct test of the split — valid, added. 6. Missing premise — valid, added. 7. API asymmetry — fixed. Re-verified after the changes: a2a3 hbg sim 12 passed / 7 skipped, a5 hbg 13 passed, 134/134 cpp UTs. The hardware A/B is unchanged by this round — it only touched comments, one test file, and a defaulted parameter. |
…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.
Device A/B re-run at 100 rounds — no regressionThe A/B in the description was 30 rounds; this is the same protocol at 100, on the post-review head. No device delta reaches the 2% review threshold.
Device: 5 improved, 3 regressed, of 8. No regression above 2%. Reading the numbers:
Both arms ran under one |
|
Follow-up #2144 (in-graph early-dispatch prerequisites: recording-time verdicts, sorted candidate CSR rows, backward in-graph fanin scan) is stacked on this branch and should merge after it. |
Summary
Implements #2106, plus the re-encode it enables.
The per-task progress byte encoded a linear progression through bit-containment
values (
0x0 -> 0x2 -> 0x3) so a lock-freefetch_orwould double as amonotone-max. That trick was load-bearing: 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_blockssplitsinto
account_published_blocks, called before a batch's token writes andreturning whether this caller reached the task's total, and the existing
seal_ed_publish_list, called after the flush. The PUBLISHED store thereforeprecedes every token the task emits, and those tokens' FINs gate the all-FIN
completion, so
PUBLISHED < token < FIN < COMPLETEDholds regardless of howsibling 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.
With the ordering guaranteed, the byte becomes a state.
ChipTaskStategainsPUBLISHEDbetweenPENDINGandCOMPLETED, and the array is written with plainstores. Readers use ordered comparisons, and
COMPLETED > PUBLISHEDis what letsa tracked producer that never publishes (DUMMY, predicate-retired) release its
publish-list waiters through the completion store alone.
This also removes the redundancy of two spellings for one fact:
progress_flagsis renamed
task_states, andTASK_FLAG_*/is_completion_flag_set/is_publish_flag_setgive way tois_completed/is_published/store_completed/store_published/reset_task_state.The columnar byte-array layout is unchanged and deliberate: under the polling
model a fanin scan reads many producers' states at once, which one cache line
answers here and would take one line per producer inside the
ChipTaskStoragestride. The slot-resident
task_statemirror is also unchanged — it staysPENDINGorCOMPLETEDand remains the in-graph readiness truth.Also corrects the
ChipTaskSlotStateheader comment, which claimed the struct is"NOT in shared memory" while a GLOBAL task's slot lives in the SM image's storage
segment.
Both arches and all three publish sites (
dispatch_shape,stage_consumer_blocks,stage_sync_start_cores) move together.Testing
host_build_graph12 passed / 7 skipped, a5host_build_graph13 passedctest -LE requires_hardware)host_build_graph, 8 cases x 30 rounds, baseline(merge-base) built in its own worktree venv, both arms sequential on one
pinned die: every device delta within +/-1.1%, none above the 2% review
threshold;
qwen3_14b_decode+1.06%bindphases, qwen, base/HEAD interleaved twice: the two repetitionsdisagree in sign (+5.8% / -12.7% on a ~0.3 ms control-plane min-of-sums),
which per
docs/dfx/hbg-bind-phases.mdmeans no resolvable host movement —as expected, since the host side performs the same stores in the same places
Fixes #2106