Fix: port hbg host-overhead reduction (#1659) to a5 - #1759
Merged
ChaoZheng109 merged 2 commits intoAug 11, 2026
Merged
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: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughThe A5 host build graph now performs task-bounded uploads, initializes claimed slots during preparation, reports ready-queue overflow, and validates submission-time initialization with a poison test. ChangesA5 host build graph runtime
Estimated code review effort: 3 (Moderate) | ~25 minutes Possibly related issues
Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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.
Inline comments:
In `@tests/ut/cpp/a5/test_hbg_submit_poison.cpp`:
- Around line 135-165: Extend the validation loop around ring.task_descriptors,
ring.task_payloads, and ring.slot_states to assert that st.task and st.payload
are bound to the corresponding descriptor and payload objects for each slot.
Compare both bindings against the addresses of desc and pl, preserving the
existing state and field checks.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 831fbee6-b201-4791-a15e-c707ff1ee7fd
📒 Files selected for processing (12)
src/a5/runtime/host_build_graph/common/pto_runtime_status.hsrc/a5/runtime/host_build_graph/docs/RUNTIME_LOGIC.mdsrc/a5/runtime/host_build_graph/host/runtime_maker.cppsrc/a5/runtime/host_build_graph/runtime/orchestrator_core/pto_orchestrator.cppsrc/a5/runtime/host_build_graph/runtime/pto_runtime2_types.hsrc/a5/runtime/host_build_graph/runtime/pto_shared_memory.hsrc/a5/runtime/host_build_graph/runtime/scheduler/pto_scheduler.hsrc/a5/runtime/host_build_graph/runtime/shared/pto_runtime2_init.cppsrc/a5/runtime/host_build_graph/runtime/shared/pto_shared_memory.cppsrc/a5/runtime/host_build_graph/runtime/shared/pto_tensormap.cpptests/ut/cpp/CMakeLists.txttests/ut/cpp/a5/test_hbg_submit_poison.cpp
Comment on lines
+135
to
+165
| // Every claimed slot's device-read fields must carry real values, not poison. | ||
| for (int32_t local = 0; local < total; local++) { | ||
| SCOPED_TRACE(testing::Message() << "slot local_id=" << local); | ||
| const int32_t slot = ring.get_slot_by_task_id(local); | ||
| const PTO2TaskDescriptor &desc = ring.task_descriptors[slot]; | ||
| const PTO2TaskPayload &pl = ring.task_payloads[slot]; | ||
| const PTO2TaskSlotState &st = ring.slot_states[slot]; | ||
|
|
||
| // Descriptor: the task id is written to this exact local id. | ||
| EXPECT_EQ(desc.task_id.local(), static_cast<uint32_t>(local)); | ||
| // task_state is written at submit (reset_for_reuse skips it): PENDING for a | ||
| // dispatchable task, COMPLETED for a pre-completed hidden-alloc. Either way a | ||
| // real enum, never poison. | ||
| const PTO2TaskState state = st.task_state.load(std::memory_order_relaxed); | ||
| EXPECT_TRUE(state == PTO2_TASK_PENDING || state == PTO2_TASK_COMPLETED); | ||
| // Completion flag is written to a real 0/1 (pending vs pre-completed), not a | ||
| // poison byte (0xAA). | ||
| const uint8_t cflag = ring.completion_flags[slot].load(std::memory_order_relaxed); | ||
| EXPECT_LE(cflag, uint8_t{1}); | ||
| // Payload counts are real, not the poison bit pattern. | ||
| EXPECT_GE(pl.fanin_count, 0); | ||
| EXPECT_LE(pl.fanin_count, PTO2_MAX_FANIN); | ||
| EXPECT_GE(pl.tensor_count, 0); | ||
| EXPECT_GE(pl.scalar_count, 0); | ||
| // predicate.op is a dispatch-time field, read only for tasks the device | ||
| // actually dispatches. submit_task_common writes it (NONE when unset); a | ||
| // pre-completed hidden-alloc is never dispatched, so it does not. | ||
| if (state == PTO2_TASK_PENDING) { | ||
| EXPECT_LE(static_cast<uint8_t>(pl.predicate.op), static_cast<uint8_t>(PredicateOp::LE)); | ||
| } | ||
| } |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Assert the task and payload bindings.
The loop does not check st.task or st.payload. If bind_buffers() does not write either field, the poisoned pointer remains and this test passes. The device can then dereference an invalid slot pointer.
Proposed test addition
const PTO2TaskSlotState &st = ring.slot_states[slot];
+ EXPECT_EQ(st.task, &ring.task_descriptors[slot]);
+ EXPECT_EQ(st.payload, &ring.task_payloads[slot]);
+
// Descriptor: the task id is written to this exact local id.📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| // Every claimed slot's device-read fields must carry real values, not poison. | |
| for (int32_t local = 0; local < total; local++) { | |
| SCOPED_TRACE(testing::Message() << "slot local_id=" << local); | |
| const int32_t slot = ring.get_slot_by_task_id(local); | |
| const PTO2TaskDescriptor &desc = ring.task_descriptors[slot]; | |
| const PTO2TaskPayload &pl = ring.task_payloads[slot]; | |
| const PTO2TaskSlotState &st = ring.slot_states[slot]; | |
| // Descriptor: the task id is written to this exact local id. | |
| EXPECT_EQ(desc.task_id.local(), static_cast<uint32_t>(local)); | |
| // task_state is written at submit (reset_for_reuse skips it): PENDING for a | |
| // dispatchable task, COMPLETED for a pre-completed hidden-alloc. Either way a | |
| // real enum, never poison. | |
| const PTO2TaskState state = st.task_state.load(std::memory_order_relaxed); | |
| EXPECT_TRUE(state == PTO2_TASK_PENDING || state == PTO2_TASK_COMPLETED); | |
| // Completion flag is written to a real 0/1 (pending vs pre-completed), not a | |
| // poison byte (0xAA). | |
| const uint8_t cflag = ring.completion_flags[slot].load(std::memory_order_relaxed); | |
| EXPECT_LE(cflag, uint8_t{1}); | |
| // Payload counts are real, not the poison bit pattern. | |
| EXPECT_GE(pl.fanin_count, 0); | |
| EXPECT_LE(pl.fanin_count, PTO2_MAX_FANIN); | |
| EXPECT_GE(pl.tensor_count, 0); | |
| EXPECT_GE(pl.scalar_count, 0); | |
| // predicate.op is a dispatch-time field, read only for tasks the device | |
| // actually dispatches. submit_task_common writes it (NONE when unset); a | |
| // pre-completed hidden-alloc is never dispatched, so it does not. | |
| if (state == PTO2_TASK_PENDING) { | |
| EXPECT_LE(static_cast<uint8_t>(pl.predicate.op), static_cast<uint8_t>(PredicateOp::LE)); | |
| } | |
| } | |
| // Every claimed slot's device-read fields must carry real values, not poison. | |
| for (int32_t local = 0; local < total; local++) { | |
| SCOPED_TRACE(testing::Message() << "slot local_id=" << local); | |
| const int32_t slot = ring.get_slot_by_task_id(local); | |
| const PTO2TaskDescriptor &desc = ring.task_descriptors[slot]; | |
| const PTO2TaskPayload &pl = ring.task_payloads[slot]; | |
| const PTO2TaskSlotState &st = ring.slot_states[slot]; | |
| EXPECT_EQ(st.task, &ring.task_descriptors[slot]); | |
| EXPECT_EQ(st.payload, &ring.task_payloads[slot]); | |
| // Descriptor: the task id is written to this exact local id. | |
| EXPECT_EQ(desc.task_id.local(), static_cast<uint32_t>(local)); | |
| // task_state is written at submit (reset_for_reuse skips it): PENDING for a | |
| // dispatchable task, COMPLETED for a pre-completed hidden-alloc. Either way a | |
| // real enum, never poison. | |
| const PTO2TaskState state = st.task_state.load(std::memory_order_relaxed); | |
| EXPECT_TRUE(state == PTO2_TASK_PENDING || state == PTO2_TASK_COMPLETED); | |
| // Completion flag is written to a real 0/1 (pending vs pre-completed), not a | |
| // poison byte (0xAA). | |
| const uint8_t cflag = ring.completion_flags[slot].load(std::memory_order_relaxed); | |
| EXPECT_LE(cflag, uint8_t{1}); | |
| // Payload counts are real, not the poison bit pattern. | |
| EXPECT_GE(pl.fanin_count, 0); | |
| EXPECT_LE(pl.fanin_count, PTO2_MAX_FANIN); | |
| EXPECT_GE(pl.tensor_count, 0); | |
| EXPECT_GE(pl.scalar_count, 0); | |
| // predicate.op is a dispatch-time field, read only for tasks the device | |
| // actually dispatches. submit_task_common writes it (NONE when unset); a | |
| // pre-completed hidden-alloc is never dispatched, so it does not. | |
| if (state == PTO2_TASK_PENDING) { | |
| EXPECT_LE(static_cast<uint8_t>(pl.predicate.op), static_cast<uint8_t>(PredicateOp::LE)); | |
| } | |
| } |
🤖 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/a5/test_hbg_submit_poison.cpp` around lines 135 - 165, Extend
the validation loop around ring.task_descriptors, ring.task_payloads, and
ring.slot_states to assert that st.task and st.payload are bound to the
corresponding descriptor and payload objects for each slot. Compare both
bindings against the addresses of desc and pl, preserving the existing state and
field checks.
Fixes hw-native-sys#1716 only. HBG is one runtime kept in sync across two arch trees (hw-native-sys#1706), so a5 still paid the full pre-hw-native-sys#1659 host cost: a boot-time blanket slot reset, whole-window completion-flag zero, full sm_size SM upload and full arena_size arena upload, and no ready-queue overflow guard. Port the diff onto the a5 tree file-for-file: - Bounded SM H2D: allocate the host SM uninitialized, zero only the header, and upload each segment (descriptors / payloads / slot_states / completion_flags) bounded to [0, total_tasks). total_tasks is range-checked before it sizes the copies. - Init-on-write: per-slot reset_for_reuse() + completion-flag clear move from the boot-time whole-window loop into orch::prepare_task as each slot is claimed. The unclaimed tail is neither initialized, uploaded, nor read. - Skip the host-only orchestrator block (fanin_seen_epoch / scope_tasks / TensorMap) from the arena H2D; an always_assert(orch_start <= orch_end) guards the layout order before slicing it out. - Latch PTO2_ERROR_READY_QUEUE_OVERFLOW from push_ready_routed instead of dropping a ready task into an anonymous forward-progress stall. - Drop the redundant per-entry tensormap stores the preceding memset already zeroed. The ready queues ship in full, mirroring a2a3, so the two trees stay identical for hw-native-sys#1715 to build on — a5 does not bound them locally just because graph execution is not yet on a5. Adds tests/ut/cpp/a5/test_hbg_submit_poison.cpp pinning the "every device-read SM field is written at submit" contract, and RUNTIME_LOGIC.md §3.1 documenting the bounded-upload contract. Validated: 92/92 cpput, 8/8 a5sim hbg scene tests (vector_example, paged_attention, prepared_callable), a5 + a5sim runtimes compile clean.
char message[1024] -> std::array<char, 1024> (call sites use .data()/.size()); int32_t kernel_ids_capture[3] -> std::array<int32_t, PTO2_SUBTASK_SLOT_COUNT>, removing the hardcoded 3. After this a5/host_build_graph pto_orchestrator.cpp is byte-identical to the a2a3 build. Co-Authored-By: Claude <noreply@anthropic.com>
ChaoZheng109
force-pushed
the
fix/issue-1716-hbg-port-1659-a5
branch
from
August 11, 2026 08:13
cd5a110 to
fd58f9b
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.
Summary
Ports the host_build_graph per-dispatch host-overhead reduction from #1659 (a2a3-only) onto the a5 tree, file-for-file. HBG is one runtime kept in sync across two arch trees (#1706); a5 still carried the full pre-#1659 code, so it paid the same host
bindcost the a2a3 change eliminated.Fixes #1716.
What changed (mirrors #1659)
descriptors/payloads/slot_states/completion_flagsare each uploaded bounded to[0, total_tasks).total_tasksis range-checked before it sizes the copies.reset_for_reuse()+ completion-flag clear move from the boot-time whole-window loop intoorch::prepare_taskas each slot is claimed; the unclaimed tail is neither initialized, uploaded, nor read.fanin_seen_epoch/scope_tasks/ TensorMap) from the arena H2D, guarded byalways_assert(orch_start <= orch_end).push_ready_routedlatchesPTO2_ERROR_READY_QUEUE_OVERFLOWinstead of dropping a ready task into an anonymous stall.memsetalready zeroed.RUNTIME_LOGIC.md §3.1) and thereset_for_reusedoc/comment sites updated to the init-on-write model.Caveat resolved during the port
Per the issue and #1659, the ready queues ship in full (not bounded), mirroring a2a3 exactly, so the two trees stay identical for #1715 to build on. a5 does not bound them locally just because graph execution is not yet on a5.
Test / Validation
test_hbg_submit_poison.cppadded for a5 — pins the "every device-read SM field is written at submit" contract (fills the window with0xAApoison, submits a representative mix, asserts no claimed slot reads poison).vector_example,paged_attention,prepared_callable).host_build_graphcompile clean.🤖 Generated with Claude Code