Skip to content

Refactor: hbg unifies the task progress byte into a sequential state enum - #2130

Merged
poursoul merged 1 commit into
hw-native-sys:mainfrom
ChaoZheng109:refactor/hbg-unify-task-state-enum
Sep 7, 2026
Merged

Refactor: hbg unifies the task progress byte into a sequential state enum#2130
poursoul merged 1 commit into
hw-native-sys:mainfrom
ChaoZheng109:refactor/hbg-unify-task-state-enum

Conversation

@ChaoZheng109

Copy link
Copy Markdown
Collaborator

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-free fetch_or would double as a
monotone-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_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.

With the ordering guaranteed, the byte becomes a state. ChipTaskState gains
PUBLISHED between PENDING and COMPLETED, and the array is written with plain
stores. 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.

This also removes the redundancy of two spellings for one fact: progress_flags
is renamed task_states, 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 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 ChipTaskStorage
stride. The slot-resident task_state mirror is also unchanged — it stays
PENDING or COMPLETED and remains the in-graph readiness truth.

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.

Both arches and all three publish sites (dispatch_shape,
stage_consumer_blocks, stage_sync_start_cores) move together.

Testing

  • Simulation tests pass — a2a3 host_build_graph 12 passed / 7 skipped, a5
    host_build_graph 13 passed
  • C++ unit tests — 134/134 (ctest -LE requires_hardware)
  • Hardware A/B, a2a3 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%
  • Host bind phases, qwen, base/HEAD interleaved twice: the two repetitions
    disagree in sign (+5.8% / -12.7% on a ~0.3 ms control-plane min-of-sums),
    which per docs/dfx/hbg-bind-phases.md means no resolvable host movement —
    as expected, since the host side performs the same stores in the same places

Fixes #2106

@coderabbitai

coderabbitai Bot commented Sep 4, 2026

Copy link
Copy Markdown

Review Change Stack

Important

Review skipped

Auto incremental reviews are disabled on this repository.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Team

Run ID: 3a626e85-fe32-4bf3-8153-6e4b9dd8e392

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

  • 🔍 Trigger review

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Team

Run ID: 6c556d16-ab60-4846-a899-349be637636d

📥 Commits

Reviewing files that changed from the base of the PR and between 4e9d447 and 30e66bb.

📒 Files selected for processing (24)
  • src/a2a3/runtime/host_build_graph/docs/RUNTIME_LOGIC.md
  • src/a2a3/runtime/host_build_graph/host/runtime_maker.cpp
  • src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler.h
  • src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_cold_path.cpp
  • src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_completion.cpp
  • src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_context.h
  • src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_dispatch.cpp
  • src/a5/runtime/host_build_graph/docs/RUNTIME_LOGIC.md
  • src/a5/runtime/host_build_graph/host/runtime_maker.cpp
  • src/a5/runtime/host_build_graph/runtime/scheduler/scheduler.h
  • src/a5/runtime/host_build_graph/runtime/scheduler/scheduler_cold_path.cpp
  • src/a5/runtime/host_build_graph/runtime/scheduler/scheduler_completion.cpp
  • src/a5/runtime/host_build_graph/runtime/scheduler/scheduler_context.h
  • src/a5/runtime/host_build_graph/runtime/scheduler/scheduler_dispatch.cpp
  • src/common/host_build_graph/host/orchestrator.cpp
  • src/common/host_build_graph/runtime_types.h
  • src/common/host_build_graph/shared/runtime_init.cpp
  • src/common/host_build_graph/shared/shared_memory.cpp
  • src/common/host_build_graph/shared_memory.h
  • tests/ut/cpp/a2a3/test_hbg_submit_poison.cpp
  • tests/ut/cpp/a5/test_hbg_submit_poison.cpp
  • tests/ut/cpp/common/test_hbg_ed_qualification.cpp
  • tests/ut/cpp/common/test_hbg_slot_claim.cpp
  • tests/ut/cpp/common/test_hbg_sm_compaction.cpp

Included review availability: Your plan provides up to 1 included review per hour; 0 remain after this review.


📝 Walkthrough

Walkthrough

The runtime replaces progress_flags with a sequential task_states enum. Scheduler publication accounting now occurs before token writes, while publish-list sealing occurs afterward. Host paths, shared-memory layout code, documentation, and unit tests use the new state model.

Changes

Task-state publication protocol

Layer / File(s) Summary
Task-state and shared-memory contract
src/common/host_build_graph/runtime_types.h, src/common/host_build_graph/shared_memory.h, src/common/host_build_graph/shared_memory.cpp, src/common/host_build_graph/shared/runtime_init.cpp
ChipTaskState now includes PUBLISHED. Shared-memory accessors, offsets, pointer setup, and compaction use task_states.
Scheduler publication and completion flow
src/{a2a3,a5}/runtime/host_build_graph/runtime/scheduler/*
Scheduler completion stores COMPLETED. Publication accounting stores PUBLISHED before token writes and seals publish lists afterward. Readiness checks use ordered task-state accessors.
Host task-state initialization and completion
src/common/host_build_graph/host/orchestrator.cpp, src/{a2a3,a5}/runtime/host_build_graph/host/runtime_maker.cpp
Host submission resets task states. Host-completed tasks store COMPLETED. Related comments use the new terminology.
Documentation and state-model validation
src/{a2a3,a5}/runtime/host_build_graph/docs/RUNTIME_LOGIC.md, tests/ut/cpp/**
Documentation and tests now cover task_states, enum values, shared-memory compaction, slot poisoning, and account-then-seal publication.

Estimated code review effort: 4 (Complex) | ~45 minutes

Merge Risk: ⚪ Minimal · up to 30e66

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
Loading

Poem

A rabbit reads each line,
The patch grows clear beneath the moon,
Small changes hop in place,
Tests guard the garden path,
Reviews bloom before the dawn.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning 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:… Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Title check ✅ Passed The title clearly summarizes the main change: replacing the HBG progress byte encoding with a sequential task-state enum.
Description check ✅ Passed The description directly explains the state refactor, publish-ordering changes, affected publish sites, compatibility model, and reported testing.
Linked Issues check ✅ Passed The implementation satisfies the coding objectives in [#2106]: it introduces sequential PENDING/PUBLISHED/COMPLETED states, moves publication accounting before token writes at all listed sites and arc…
Out of Scope Changes check ✅ Passed The changes remain within [#2106]. Documentation, comments, shared-memory code, scheduler code, and unit tests directly support the task-state refactor and publication-ordering changes.
Full details: Docstring Coverage

Explanation

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.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

@ChaoZheng109
ChaoZheng109 force-pushed the refactor/hbg-unify-task-state-enum branch from 30e66bb to 6a6079c Compare September 7, 2026 03:33
@ChaoZheng109

Copy link
Copy Markdown
Collaborator Author

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 published_block_count. It is a two-sided seq_cst handshake between the late stager's staged_core_mask.fetch_or and the release path's early_dispatch_state store, with each bit claimed exactly once — exchange(0) on the releaser side, fetch_and(~candidates) on the late stager's. Whichever side observes the other's write rings the bits the other did not take. A consumer released by the earlier publication cannot take cores this task's blocks already occupy (prepare_block_for_dispatch claimed them before any accounting) and holds nothing the rings wait on, so the count moving ahead of them costs a few instructions of skew and nothing else. Accounting stays before the tokens; the comment now states the handshake as the current invariant.

2. Rename drift — valid, all fixed. The layout diagram in shared_memory.h (stale type and stale name), both "bits are monotonic" phrasings in runtime_types.h, and every "completion flag(s)" naming the renamed array: shared_memory.h:202, shared/shared_memory.cpp, shared/runtime_init.cpp, host/orchestrator.cpp, task_id.h, both RUNTIME_LOGIC.md, GRAPH_EXECUTION.md. The two remaining matches are scheduler_completion.cpp:99's "deferred-completion flag", which is any_subtask_deferred — a different flag, deliberately untouched.

3. tmr divergence — the premise is wrong, but a note was still worth adding. tmr has no progress_flags: git grep progress_flags src/*/runtime/tensormap_and_ringbuffer/ is empty. Its is_completion_flag_set reads lifecycle_flags & COMPLETION_DONE, a per-slot byte inside ChipTaskSlotState, and its readiness is fanin_refcount under the push model — there is no byte array to scan, no publish byte, no seal, and no publish list. Its record_published_blocks only does published_block_count.fetch_add; it writes no state at all. So there is no fetch_or encoding in tmr to diverge from, and #2106's hbg-only Location is not an oversight. What genuinely differs after this PR is the function name, so the commit message now says tmr is untouched and why.

4. Unwritten-byte reads — valid, documented. 0xAA & 1 == 0 read as pending; 0xAA >= 2 reads as completed. Unreachable today (init-on-write in orch::prepare_task, total_tasks bounding every walk), but the failure mode did get stricter, so the array's comment now states the discipline the reads depend on.

5. No direct test of the split — valid, added. PublishedStateIsVisibleBeforeTheChainSeals enters the window on purpose: after account_published_blocks returns true, is_published is already true while ed_publish_list_head is still open and the drain queue still empty; the seal then closes the chain and hands the waiter over. Added CompletionAloneCountsAsPublished alongside it, since COMPLETED > PUBLISHED is the other half of the new contract and the DUMMY / predicate-retired path depends on it.

6. Missing premise — valid, added. account_published_blocks now names the precondition the call site owes: every block counted has its token written by the same caller before it returns. Flagged as the thing to re-check when adding a publish site.

7. API asymmetry — fixed. store_published takes a memory_order like the other three.

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.
@ChaoZheng109

Copy link
Copy Markdown
Collaborator Author

Device A/B re-run at 100 rounds — no regression

The 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.

Merge-base: 4e9d4475  ->  HEAD: 6a6079c1
Args: -n 100 -r host_build_graph      (pto_isa.pin a8040450)
Device: baseline=4, current=4         (one pinned die, arms sequential)
Example / Case Base (us) HEAD (us) Delta Change
alternating_matmul_add/Case1
(host) 32114.2 31790.7 -323.5 -1.01%
(device) 838.4 836.4 -2.0 -0.24%
benchmark_bgemm/Case0
(host) 15246.7 14764.6 -482.1 -3.16%
(device) 710.8 705.9 -4.9 -0.69%
paged_attention_unroll/Case1
(host) 32676.9 32140.7 -536.2 -1.64%
(device) 1256.2 1259.7 +3.5 +0.28%
paged_attention_unroll/Case2
(host) 9747.1 10176.7 +429.6 +4.41%
(device) 663.3 661.8 -1.5 -0.23%
paged_attention_unroll_manual_scope/Case1
(host) 32188.5 31770.5 -418.0 -1.30%
(device) 1258.1 1257.1 -1.0 -0.08%
paged_attention_unroll_manual_scope/Case2
(host) 10094.9 10009.9 -85.0 -0.84%
(device) 658.0 660.8 +2.8 +0.43%
batch_paged_attention/Case1
(host) 35040.6 34387.2 -653.4 -1.86%
(device) 2811.9 2803.5 -8.4 -0.30%
qwen3_14b_decode/GraphExecutionBatch16Seq3500
(host) 37769.4 38215.6 +446.2 +1.18%
(device) 37242.1 37686.8 +444.7 +1.19%

Device: 5 improved, 3 regressed, of 8. No regression above 2%.

Reading the numbers:

  • The device spread is -0.69% .. +1.19%, entirely within the noise band — which is what this change should produce, since it reorders a few instructions on the publish path without adding or removing work.
  • qwen's +1.19% is the largest absolute move (+445 us), but it sits inside twice the same-binary drift this box shows across pairs (+/-273 us, recorded while measuring Add: host_build_graph early dispatch via an ED publish list #2095), and host and device moved by the same amount in the same direction — a whole-run drift, not a scheduler-side change.
  • The host column (-3.16% .. +4.41%) is not interpretable on its own: it includes Python overhead and scatters about four times as widely as the device column. The clean host-side answer is the interleaved hbg-bind-phases measurement reported earlier, whose two repetitions disagreed in sign and therefore resolve no movement.

Both arms ran under one task-submit allocation on die 4, baseline built in its own worktree venv at the merge-base. tensormap_and_ringbuffer was not benchmarked: this branch does not touch it, so that arm would be a zero measurement by construction.

@ChaoZheng109

Copy link
Copy Markdown
Collaborator Author

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.

@poursoul
poursoul merged commit f6c8621 into hw-native-sys:main Sep 7, 2026
20 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Code Health] hbg: make PUBLISHED precede COMPLETED by construction and turn the progress byte into a sequential enum

2 participants