Refactor: host build graph, lazy completion watermark, 4 AICPU schedulers - #1619
Refactor: host build graph, lazy completion watermark, 4 AICPU schedulers#1619raphael-s-steiner wants to merge 1 commit into
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 scheduler replaces byte completion flags and a dedicated resolution thread with task-id stamps, per-thread watermark tracking, deferred publication retries, and inline task resolution across all AICPU scheduler threads. ChangesCompletion and scheduler execution
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant SchedulerThread
participant PTO2SchedulerState
participant PTO2SharedMemoryRingHeader
participant AsyncWaitList
SchedulerThread->>PTO2SchedulerState: complete task with thread_idx
PTO2SchedulerState->>PTO2SharedMemoryRingHeader: publish or queue completion flag
PTO2SchedulerState->>PTO2SharedMemoryRingHeader: retry queued flags
PTO2SchedulerState->>AsyncWaitList: drain wake list after publication
PTO2SharedMemoryRingHeader-->>SchedulerThread: advance completed watermark
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 |
e597958 to
5a13e02
Compare
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
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/host_build_graph/runtime/shared/pto_shared_memory.cpp (1)
83-102: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winGuard the shuffle floor inside
init_header_per_ring.
init_per_ringrejectstask_window_sizes[r]whose__builtin_ctzll(...)is less thanshuffle_lower_bits, butinit_header()callsinit_header_per_ring()directly and can setshuffle_higher_bitsto a negative value.flag_index()then shifts by that negative amount, causing undefined behavior if a smalltask_window_sizeis passed there. Put this validation ininit_header_per_ringso every caller is protected.🤖 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/host_build_graph/runtime/shared/pto_shared_memory.cpp` around lines 83 - 102, Move the task-window alignment validation from PTO2SharedMemoryHandle::init_per_ring into PTO2SharedMemoryHandle::init_header_per_ring, checking every task_window_sizes entry before calculating or storing shuffle_higher_bits. Ensure direct callers such as init_header() reject values with insufficient trailing zero bits, preventing flag_index() from receiving a negative shift amount.
🧹 Nitpick comments (5)
src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_dispatch.cpp (2)
1141-1144: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDrop the always-true
thread_idxguard, or state the intent.
AicpuExecutor::runrejects anythread_idx >= aicpu_thread_num_before it callsresolve_and_dispatch(aicpu_executor.cpplines 220-226). The condition at line 1141 is therefore always true, and every thread drains the dummy queue. If the guard is a leftover from the removed resolution-thread partition, remove it. If it encodes a future restriction, replace it with an assertion.🤖 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/host_build_graph/runtime/scheduler/scheduler_dispatch.cpp` around lines 1141 - 1144, Remove the redundant thread_idx < aicpu_thread_num_ guard around the dummy_ready_queue drain in resolve_and_dispatch, since AicpuExecutor::run already validates the index; keep the dummy batch pop behavior unchanged for every valid executor thread.
1293-1301: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueResolve the "Can be made less frequent" note before merge.
Line 1300 carries an open tuning note inside the hot loop.
update_completed_watermarkscans the completed prefix once per iteration on every thread, so the cost scales withtask_window_sizetimes the iteration count. Either record the measured cost that justifies the current frequency, or convert the note into a tracked item. The PR already contains measurements for this branch, so a short statement of the observed cost is enough.🤖 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/host_build_graph/runtime/scheduler/scheduler_dispatch.cpp` around lines 1293 - 1301, The hot-loop comment near header->ring.update_completed_watermark must resolve the “Can be made less frequent” note. Replace it with a brief statement documenting the measured cost of the per-iteration completed-prefix scan, using the existing branch measurements, or convert it into a tracked item with an issue reference; preserve the watermark update behavior.src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_completion.cpp (1)
186-225: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winExtract the duplicated Resolve phase emit. Both sites now time
on_task_complete, redefine the sameRESOLVE_EMIT_MIN_CYCLESconstant, apply the same 1 µs filter, and calll2_swimlane_aicpu_record_sched_phasewithL2SwimlaneSchedPhaseKind::Resolve. The shared root cause is the absence of one helper for this emit, so the constant and the filter can drift apart.
src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_completion.cpp#L186-L225: replace the inline timing, constant, and emit with a call to a single helper, for exampleemit_resolve_phase(thread_idx, t0, t1, consumers_resolved), and keep theon_task_completecall at this site.src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_dispatch.cpp#L1160-L1188: call the same helper for the dummy path and delete the secondRESOLVE_EMIT_MIN_CYCLESdefinition.🤖 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/host_build_graph/runtime/scheduler/scheduler_completion.cpp` around lines 186 - 225, The Resolve phase emission logic is duplicated across both scheduler paths. In src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_completion.cpp:186-225, keep the on_task_complete call but replace the inline timing/filter/recording logic with a shared emit_resolve_phase helper; in src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_dispatch.cpp:1160-1188, use the same helper for the dummy path and remove its duplicate RESOLVE_EMIT_MIN_CYCLES definition, centralizing the 1 µs threshold and Resolve recording behavior.src/a2a3/runtime/host_build_graph/runtime/scheduler/pto_scheduler.h (1)
34-34: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winAdd explicit headers for the heap algorithms and comparator.
pto_scheduler.husesstd::push_heap,std::pop_heap, andstd::greater<>, but only includes<atomic>and<vector>among the standard headers. Add<algorithm>for the heap operations and<functional>forstd::greater<>.♻️ Proposed include addition
+#include <algorithm> +#include <functional> `#include` <vector>🤖 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/host_build_graph/runtime/scheduler/pto_scheduler.h` at line 34, Update the includes in pto_scheduler.h to add the standard headers algorithm and functional, supporting the existing std::push_heap/std::pop_heap and std::greater<> usages while retaining the current includes.src/a2a3/runtime/host_build_graph/runtime/pto_async_wait.h (1)
283-285: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRemove the unused
Profilingtemplate parameter.The declaration, definition, call site, and body use
<false>only and do not readProfiling, so make this a plain member function unless a future specialization adds profiling behavior. The matching out-of-class definition usestemplate <bool Profiling>.🤖 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/host_build_graph/runtime/pto_async_wait.h` around lines 283 - 285, Remove the unused Profiling template parameter from poll_and_complete, converting its declaration and matching out-of-class definition to a non-template member function. Update every call site currently invoking poll_and_complete<false> to call poll_and_complete directly, preserving the existing behavior.
🤖 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 `@src/a2a3/runtime/host_build_graph/docs/RUNTIME_LOGIC.md`:
- Around line 664-667: Update the RUNTIME_LOGIC.md description of
retry_set_completion_flags to state that it runs at the Phase 1 → Phase 2
boundary, replacing the “top of the loop iteration” wording while preserving the
existing behavior details.
- Around line 686-694: Correct the RUNTIME_LOGIC.md description in §§8.2 and 8.4
to distinguish the two watermark advances: on_mixed_task_complete uses
weak_update_completed_watermark for a bounded advance through task_id + 1, while
resolve_and_dispatch performs the full contiguous-prefix advance once per
scheduler-loop iteration. Remove the claim that the completion callback advances
the full prefix or that capping at my_id is incorrect, and update §8.2 step 3 to
reference both mechanisms accurately.
In `@src/a2a3/runtime/host_build_graph/runtime/scheduler/pto_scheduler.h`:
- Around line 471-473: Replace
PTO2SchedulerState::failed_heap_of_set_completion_flag with a fixed int32_t
storage array sized by a PTO2_MAX_... capacity and a count field, avoiding
non-trivial std::vector state. Update all accesses to use the count and bounded
array, and report an overflow when the failsafe cannot record another entry
instead of silently exceeding capacity.
In `@src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_dispatch.cpp`:
- Around line 994-998: Update the scheduler loop in the relevant dispatch
function so every exit path—completed_, orchestrator or scheduler errors, async
polling errors, and timeout—calls sched_->retry_set_completion_flags(thread_idx)
before header->ring.update_completed_watermark(thread_idx). Ensure deferred
completion entries are retried consistently before leaving the loop, without
changing the existing exit behavior.
---
Outside diff comments:
In `@src/a2a3/runtime/host_build_graph/runtime/shared/pto_shared_memory.cpp`:
- Around line 83-102: Move the task-window alignment validation from
PTO2SharedMemoryHandle::init_per_ring into
PTO2SharedMemoryHandle::init_header_per_ring, checking every task_window_sizes
entry before calculating or storing shuffle_higher_bits. Ensure direct callers
such as init_header() reject values with insufficient trailing zero bits,
preventing flag_index() from receiving a negative shift amount.
---
Nitpick comments:
In `@src/a2a3/runtime/host_build_graph/runtime/pto_async_wait.h`:
- Around line 283-285: Remove the unused Profiling template parameter from
poll_and_complete, converting its declaration and matching out-of-class
definition to a non-template member function. Update every call site currently
invoking poll_and_complete<false> to call poll_and_complete directly, preserving
the existing behavior.
In `@src/a2a3/runtime/host_build_graph/runtime/scheduler/pto_scheduler.h`:
- Line 34: Update the includes in pto_scheduler.h to add the standard headers
algorithm and functional, supporting the existing std::push_heap/std::pop_heap
and std::greater<> usages while retaining the current includes.
In
`@src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_completion.cpp`:
- Around line 186-225: The Resolve phase emission logic is duplicated across
both scheduler paths. In
src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_completion.cpp:186-225,
keep the on_task_complete call but replace the inline timing/filter/recording
logic with a shared emit_resolve_phase helper; in
src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_dispatch.cpp:1160-1188,
use the same helper for the dummy path and remove its duplicate
RESOLVE_EMIT_MIN_CYCLES definition, centralizing the 1 µs threshold and Resolve
recording behavior.
In `@src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_dispatch.cpp`:
- Around line 1141-1144: Remove the redundant thread_idx < aicpu_thread_num_
guard around the dummy_ready_queue drain in resolve_and_dispatch, since
AicpuExecutor::run already validates the index; keep the dummy batch pop
behavior unchanged for every valid executor thread.
- Around line 1293-1301: The hot-loop comment near
header->ring.update_completed_watermark must resolve the “Can be made less
frequent” note. Replace it with a brief statement documenting the measured cost
of the per-iteration completed-prefix scan, using the existing branch
measurements, or convert it into a tracked item with an issue reference;
preserve the watermark update behavior.
🪄 Autofix (Beta)
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: 6a325e04-21b6-4075-8a13-25d7dd3f5c57
📒 Files selected for processing (15)
src/a2a3/runtime/host_build_graph/aicpu/aicpu_executor.cppsrc/a2a3/runtime/host_build_graph/docs/RUNTIME_LOGIC.mdsrc/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_orchestrator.cppsrc/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_runtime2.cppsrc/a2a3/runtime/host_build_graph/runtime/pto_async_wait.hsrc/a2a3/runtime/host_build_graph/runtime/pto_runtime2_types.hsrc/a2a3/runtime/host_build_graph/runtime/pto_shared_memory.hsrc/a2a3/runtime/host_build_graph/runtime/scheduler/pto_scheduler.hsrc/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_cold_path.cppsrc/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_completion.cppsrc/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_context.hsrc/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_dispatch.cppsrc/a2a3/runtime/host_build_graph/runtime/shared/pto_shared_memory.cpptests/ut/cpp/CMakeLists.txttests/ut/cpp/a2a3/test_hbg_shared_memory.cpp
💤 Files with no reviewable changes (1)
- src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_context.h
5a13e02 to
d792027
Compare
…lers Co-authored-by: noabauma <noah.baumann@h-partners.com> Co-authored-by: Sergio Martin <eienburuu@gmail.com>
d792027 to
f2a725e
Compare
1 participant
Refactor: host build graph — 4 AICPU schedulers, lazy completion watermark
Branch
refactor/lazy-completion-mark→main(single commit,e5979580),simplerrepo.One of two possible versions of
host-build-graphcompletion-watermark refactor the other one being #1618.This lazy one is somewhat less performant due to duplicated work, but has easier synchronization.
Revert: 3S+1P → back to 4 full schedulers
The prior model split the
aicpu_thread_num_AICPU threads into(N-1)core-owning scheduler threads (S) plus one dedicated, core-less resolution thread (P). Finished tasks were handed off through a per-S SPSCCompletedTaskQueue, and only P calledon_task_complete— publishingcompletion_flags, draining wake lists, and advancing the watermark.This branch reverts that split: all
aicpu_thread_num_threads are now full schedulers again (scheduler_cold_path.cpp—active_sched_threads_ = aicpu_thread_num_, theaicpu_thread_num_ >= 2gate is gone). Each thread's loop (scheduler_dispatch.cpp) now does completion (Phase 1), async/mailbox polling, dummy/dependency-only drain (Phase 3), and dispatch (Phase 4) itself. Removed entirely:run_resolution_thread,CompletedTaskQueue,p_thread_idx(), and thesp_queues_array.Completion watermark: lazy updates
completed_flagschanged from a 1-byte-per-slot array (0/1) toint32stamps:-1= pending, else the completing task's ownlocal_id. This is what makes flag-slot reuse (below) detectable.weak_update_completed_watermark, capped attask_id + 1) done inline inon_mixed_task_complete— instead of always walking the full contiguous completed prefix on every single completion.update_completed_watermark) still exists but now runs at most once per scheduler-loop iteration, after that iteration's completions and dispatch, plus at the loop's exit points (scheduler_cold_path.cpp'shandle_orchestrator_exit) — so a low-id straggler completing after a higher one doesn't strand the watermark.completed_watermarkis now "lowest id not yet guaranteed complete" (was "highest id guaranteed complete"), so every gate changed from>=to>(wait_for_tensor_ready,on_scope_end,append_fanin_or_fail, etc. inpto_runtime2.cpp,pto_orchestrator.cpp).Flag-slot reuse (fewer flags than tasks)
completion_flagsis sized totask_window_size(the ring window), not to the total task count, so a slot is reused bylocal_id + task_window_sizeonce a run submits more tasks than the array has room for:flag_index()bit-shuffle (swapsshuffle_lower_bits/shuffle_higher_bitsoflocal_id & task_window_mask) so consecutive ids land on different cachelines — good for the linear watermark scan — while remaining a bijection over[0, task_window_size), so the reuse period is unchanged.is_completion_flag_setstays correct across reuse via acompleted_watermarkfallback (cached per-thread in the newcached_completed_watermarkarray): once an id falls behind the watermark it reads as complete forever, even after its slot is later overwritten by a subsequent lap.try_set_completion_flag(non-blocking) only writes once the slot's previous occupant (local_id - task_window_mask) is watermark-certified; otherwise it returnsfalsewithout writing or spinning.set_completion_flag(blocking/spinning) survives only for the host-side early-resolve publish inalloc_tensors(pto_orchestrator.cpp) — there's no AICPU scheduler loop on the host to retry from, so it has to spin instead.Deadlock failsafe
A failed
try_set_completion_flagis pushed onto a per-AICPU-thread min-heap,failed_heap_of_set_completion_flag(pto_scheduler.h), and retried once per loop iteration viaretry_set_completion_flags, which drains it in ascending task-id order and stops at the first still-blocked entry (the heap min always has the loosest watermark requirement). Critically, the deferred store's wake-list drain is deferred right along with it — draining a producer's wake list before its flag is durably set would let a reclassified waiter re-register on an already-SENTINELproducer and spin forever, since that producer will never runon_mixed_task_completeagain.Correctness (no deadlock) rests on one invariant: no task may depend, directly or transitively, on a task whose id is
>= its own id + task_window_size— a topologically-submitted graph satisfies this for free, since every fanin producer's id is already less than its consumer's. The heap/retry machinery is a pure failsafe on top of that — the code comment onfailed_heap_of_set_completion_flagnotes it "should always be empty."Also included
test_hbg_shared_memory.cpp(+ CMake target) validating theshuffle_higher_bitsinvariant:init_per_ringnow rejects atask_window_sizetoo small to provideshuffle_lower_bits(4) trailing zero bits, which would otherwise makeflag_index's shift amount negative (UB).PTO2SharedMemoryRingHeader/PTO2SharedMemoryHeadergrew (256→576 / 320→640 bytes) to hold the per-thread cached-watermark array and shuffle metadata; layoutstatic_asserts updated accordingly.Performance
Comparison against main and #1618 for both device wall-clock and kernel only (as measured by tracr)