Refactor: host build graph, eager completion watermark, 4 AICPU schedulers - #1618
Open
raphael-s-steiner wants to merge 1 commit into
Open
Refactor: host build graph, eager completion watermark, 4 AICPU schedulers#1618raphael-s-steiner wants to merge 1 commit into
raphael-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:
📝 WalkthroughWalkthroughThis PR removes the dedicated resolution thread from the AICPU scheduler, unifying all threads under a single scheduling model. Completion flags change from byte-based to int32 identity stamps with stricter watermark semantics, requiring updates across shared memory, orchestrator, runtime, and scheduler code, plus a per-thread retry mechanism, documentation, and a new test. ChangesScheduler unification and completion-flag rework
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant SchedulerThread
participant CompletionFlags
participant WakeList
participant Watermark
SchedulerThread->>CompletionFlags: try_set_completion_flag(thread_idx, local_id)
alt flag set successfully
CompletionFlags-->>SchedulerThread: success
SchedulerThread->>WakeList: drain_wake_list(thread_idx)
SchedulerThread->>Watermark: update_completed_watermark(thread_idx, local_id)
else reuse not yet certified
CompletionFlags-->>SchedulerThread: failure
SchedulerThread->>SchedulerThread: push to failed_heap_of_set_completion_flag
SchedulerThread->>SchedulerThread: retry_set_completion_flags(thread_idx) later
end
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 |
raphael-s-steiner
force-pushed
the
refactor/eager-completion-mark
branch
from
July 31, 2026 12:24
2da839f to
85aec90
Compare
raphael-s-steiner
marked this pull request as ready for review
July 31, 2026 12:39
raphael-s-steiner
force-pushed
the
refactor/eager-completion-mark
branch
from
July 31, 2026 15:08
85aec90 to
7eb4bda
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: 1
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/scheduler/scheduler_cold_path.cpp (1)
699-712: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winLower the assigned thread count instead of returning 0 for the all-active case.
nthreads == 1becomesaicpu_thread_num_ = 1, which makesassign_cores_to_threads()returnfalsebecause scheduler threads are configured to be fewer thannthreads. Also avoid assigningaicpu_thread_num_ == MAX_AICPU_THREADS:assign_cores_to_threads()then loops over allcore_trackers_/array entries whileaic_count_ == 0, soaic_count_ / active_sched_threads_yields 0 and no cores are registered (same as returning 0 early withaicpu_thread_num_ = 2).🤖 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_cold_path.cpp` around lines 699 - 712, Update the scheduler-thread count calculation that feeds SchedulerContext::assign_cores_to_threads() so the all-active case lowers the count instead of returning 0: when nthreads == 1, set aicpu_thread_num_ to 1 only if that satisfies the configured constraint, otherwise reduce it to a valid value; never assign MAX_AICPU_THREADS, particularly when aic_count_ == 0. Preserve a positive thread count that lets assign_cores_to_threads() complete without zero-cluster division or empty core registration.
🧹 Nitpick comments (3)
tests/ut/cpp/CMakeLists.txt (1)
670-697: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicate boilerplate from
add_a2a3_hbg_runtime_test.Lines 674-697 repeat the include directories, link libraries, and test-registration code from
add_a2a3_hbg_runtime_test(lines 122-147). Only the extra compiled sources differ:pto_shared_memory.cpphere versusscope_stats_collector_aicpu.cppin the function.Generalize the function to accept an extra-sources list. This removes the duplicate block and keeps future host-build-graph test targets consistent.
♻️ Proposed refactor
-function(add_a2a3_hbg_runtime_test name src) +function(add_a2a3_hbg_runtime_test name src) + set(extra_srcs ${ARGN}) add_executable(${name} ${src} ${CMAKE_SOURCE_DIR}/stubs/test_stubs.cpp - ${CMAKE_SOURCE_DIR}/../../../src/common/platform/shared/aicpu/scope_stats_collector_aicpu.cpp + ${extra_srcs} ) ... endfunction()Then define the new test as:
add_a2a3_hbg_runtime_test(test_hbg_shared_memory a2a3/test_hbg_shared_memory.cpp ${CMAKE_SOURCE_DIR}/../../../src/a2a3/runtime/host_build_graph/runtime/shared/pto_shared_memory.cpp )🤖 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/CMakeLists.txt` around lines 670 - 697, Update add_a2a3_hbg_runtime_test to accept and append an extra-sources list when creating the executable, while retaining its existing include directories, link libraries, test registration, and labels. Replace the standalone test_hbg_shared_memory target block with an add_a2a3_hbg_runtime_test call passing its test source and pto_shared_memory.cpp.src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_orchestrator.cpp (1)
1064-1079: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winCorrect the watermark comment.
The code is right. The explanation is not.
update_completed_watermarkwalks forward withis_completion_flag_set(next), so a later device completer that lands exactly on the frontier does walk past this pre-set flag. The precise reason for the explicit host call is narrower: only a completer whoselocal_idequals the current watermark advances it, so if the frontier already sits at this task's id, no other completer will ever call with that id.📝 Proposed comment fix
- // every consumer register_wakes on a producer that never runs on device and - // the run hangs. update_completed_watermark only advances when called with - // local_id equal to the current watermark, so this task's own call is the - // only chance to move the watermark past it — a later on-device completer - // whose local_id no longer matches the (still-stuck) watermark will no-op, - // not walk past this pre-set flag on our behalf. + // every consumer register_wakes on a producer that never runs on device and + // the run hangs. update_completed_watermark advances only when its local_id + // equals the current watermark. No device thread ever calls it with THIS + // task's local_id, so if the frontier already sits at this id, only this + // call can move it forward. (A later device completer that does land on the + // frontier walks over this pre-set flag as part of its prefix walk.)🤖 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/orchestrator_core/pto_orchestrator.cpp` around lines 1064 - 1079, Update the multi-line comment block preceding the set_completion_flag and update_completed_watermark calls to correct the explanation of watermark advancement behavior. Replace the incorrect statement that a later device completer will no-op and not walk past the pre-set flag with the accurate explanation that update_completed_watermark walks forward using is_completion_flag_set, so device completers landing on the frontier do walk past pre-set flags. Clarify the narrower and actual reason for the explicit host call: only a completer whose local_id equals the current watermark advances it, so if the frontier already sits at this task's local_id (done_local), no other completer will ever call with that matching id to move the watermark forward.src/a2a3/runtime/host_build_graph/runtime/scheduler/pto_scheduler.h (1)
474-515: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winReplace
reallocwith arena or fixed-storage for the failed-completion heap.
FailedCompletionFlagHeapis allocated from the AICPU scheduler state but still calls stdlibcrealloc/freeand aborts on allocation failure. Since this heap is only needed in rare completion-flag CAS contention, use a small fixed capacity or the scheduler arena instead.🤖 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` around lines 474 - 515, Update FailedCompletionFlagHeap to avoid stdlibc realloc/free and allocation-failure aborts by using fixed-capacity storage or allocation from the scheduler arena. Preserve push/pop heap behavior and ensure destroy performs only the corresponding non-stdlib cleanup, with capacity sized for the rare completion-flag contention use case.
🤖 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/runtime/shared/pto_shared_memory.cpp`:
- Around line 153-160: Update the task_window_sizes validation loop in the
shared-memory initialization path to reject zero values before calling
__builtin_ctzll and reject any non-power-of-two value. Keep the existing
shuffle_lower_bits constraint, ensuring every accepted size is a nonzero power
of two with sufficient trailing zero bits before task_window_size is assigned.
---
Outside diff comments:
In `@src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_cold_path.cpp`:
- Around line 699-712: Update the scheduler-thread count calculation that feeds
SchedulerContext::assign_cores_to_threads() so the all-active case lowers the
count instead of returning 0: when nthreads == 1, set aicpu_thread_num_ to 1
only if that satisfies the configured constraint, otherwise reduce it to a valid
value; never assign MAX_AICPU_THREADS, particularly when aic_count_ == 0.
Preserve a positive thread count that lets assign_cores_to_threads() complete
without zero-cluster division or empty core registration.
---
Nitpick comments:
In
`@src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_orchestrator.cpp`:
- Around line 1064-1079: Update the multi-line comment block preceding the
set_completion_flag and update_completed_watermark calls to correct the
explanation of watermark advancement behavior. Replace the incorrect statement
that a later device completer will no-op and not walk past the pre-set flag with
the accurate explanation that update_completed_watermark walks forward using
is_completion_flag_set, so device completers landing on the frontier do walk
past pre-set flags. Clarify the narrower and actual reason for the explicit host
call: only a completer whose local_id equals the current watermark advances it,
so if the frontier already sits at this task's local_id (done_local), no other
completer will ever call with that matching id to move the watermark forward.
In `@src/a2a3/runtime/host_build_graph/runtime/scheduler/pto_scheduler.h`:
- Around line 474-515: Update FailedCompletionFlagHeap to avoid stdlibc
realloc/free and allocation-failure aborts by using fixed-capacity storage or
allocation from the scheduler arena. Preserve push/pop heap behavior and ensure
destroy performs only the corresponding non-stdlib cleanup, with capacity sized
for the rare completion-flag contention use case.
In `@tests/ut/cpp/CMakeLists.txt`:
- Around line 670-697: Update add_a2a3_hbg_runtime_test to accept and append an
extra-sources list when creating the executable, while retaining its existing
include directories, link libraries, test registration, and labels. Replace the
standalone test_hbg_shared_memory target block with an add_a2a3_hbg_runtime_test
call passing its test source and pto_shared_memory.cpp.
🪄 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: 44dad46c-6e3c-47e9-97aa-96c14a2fab67
📒 Files selected for processing (16)
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_runtime2_init.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
…ulers Co-authored-by: noabauma <noah.baumann@h-partners.com> Co-authored-by: Sergio Martin <eienburuu@gmail.com>
raphael-s-steiner
force-pushed
the
refactor/eager-completion-mark
branch
from
July 31, 2026 15:22
7eb4bda to
a3a637f
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.
Refactor: host build graph, eager completion watermark, 4 AICPU schedulers
Base:
main· Branch:refactor/eager-completion-mark· 1 commit, 15 files (+651/-466)One of two possible versions of
host-build-graphcompletion-watermark refactor the other one being #1619.This eager one is more performant due to minimal work, but has delicate synchronization logic.
Summary
Reverts the 3S+1P scheduler split back to 4 uniform AICPU schedulers, and reworks
completed_watermarkmaintenance to be eager rather than lazy. Eager updates requirecompletion_flagsto become a reusable, thread-safe structure (previously a singlebyte per slot, implicitly host-only), which in turn needed a failsafe against the
flag-slot-reuse deadlock the old design couldn't hit.
1. Revert: 3 schedulers + 1 dedicated resolution thread → 4 AICPU schedulers
The earlier design split AICPU threads into 3 core-owning scheduler (S) threads plus
1 core-less resolution (P) thread that alone drained completions, published
completion_flags, drained wake lists, and advanced the watermark — funneling allcompletion resolution through a single thread.
This PR removes that split entirely:
CompletedTaskQueue(the per-S → P SPSC handoff ring) andrun_resolution_threadare deleted (
scheduler_context.h,scheduler_dispatch.cpp).p_thread_idx()/p_thread_idx_are gone;assign_cores_to_threadsno longerreserves the last thread as core-less, and the
aicpu_thread_num >= 2floor (1 S +1 P) is dropped —
active_sched_threads_ = aicpu_thread_num_again, so all 4threads own cores and resolve their own completions.
aicpu_executor.cppcallsresolve_and_dispatchuniformly instead of branching onwhether a thread is the P thread.
job) move back into each scheduler thread's own
resolve_and_dispatchloop.2. Completion watermark: eager updates
completed_watermarkis now advanced eagerly by every completer, not justopportunistically:
on_mixed_task_complete), its deferred retry(
retry_set_completion_flags), and the host orchestrator's inline hidden-alloccompletion — calls
update_completed_watermark(thread_idx, my_id)exactly once,immediately after that id's own
completion_flagsentry is actually visible.my_idis exactly the current watermark frontier; onlythe completer landing at the frontier does the CAS-advance walk over the full
contiguous completed prefix. Out-of-order completers defer to whoever completes the
frontier task later.
cached_completed_watermarkavoids re-reading the atomic on everyis_completion_flag_setcheck by falling back to a cached watermark value.completed_watermarkis now "lowest id not yetguaranteed complete" (was "highest id guaranteed complete"), so comparisons flip
from
>=to>at call sites (wait_for_tensor_ready, reclaim gates, etc.).3.
completion_flags: fewer flags than tasks, slots reusedcompletion_flagschanges from auint8_t[task_window_size]byte array (host-onlywriter, implicitly one-shot) to an
int32_t[task_window_size]array where each entrystores either
-1(pending) or thelocal_idthat owns it. Because the stored valueis the id itself rather than a boolean, a slot can be safely reused across laps:
the array no longer needs to be sized to the total task count, only to the ring's
task window —
local_idandlocal_id + task_window_sizeshare a slot, and reuse isgated on
completed_watermarkhaving certified the slot's previous occupant first.flag_index()bit-reindexeslocal_id & task_window_mask(swaps lowshuffle_lower_bitsbits into the high position) so consecutive ids land ondifferent cachelines, keeping
update_completed_watermark's linear scancache-friendly.
set_completion_flag(host, blocking) and the newtry_set_completion_flag(device, non-blocking) both gate the store on the previous occupant being
certified;
try_set_completion_flagreturnsfalseinstead of spinning when itisn't.
is_completion_flag_setfalls back tocompleted_watermarkso a slot that's beenoverwritten by a later lap still reports the earlier id as complete.
4. Deadlock failsafe for flag-slot reuse
The correctness argument for reuse is: task
tmust not depend on a task with id>= t + task_window_size, which holds automatically since task ids follow thedependency graph's topological order. But rather than assume that invariant always
holds, a failsafe absorbs a violation instead of deadlocking:
try_set_completion_flagfails insideon_mixed_task_complete, the task id ispushed onto a per-thread min-heap (
failed_heap_of_set_completion_flag) instead ofthe thread spinning or blocking. Wake-list drain and the watermark update for that
id are skipped and deferred.
retry_set_completion_flags, which retries thesmallest pending id in the heap; on success it drains that task's wake list and
advances the watermark — its one deferred chance, taken later instead of never.
isn't — it exists purely as a backstop, not a normal-path mechanism.
Also in this diff
docs/RUNTIME_LOGIC.md(§6.2, §7.2, §8.2, §8.4) rewritten to match the above.test_hbg_shared_memory.cppcovers theshuffle_higher_bitsinvariant
flag_index()depends on (rejects atask_window_sizetoo small toprovide
shuffle_lower_bitsof headroom, which would otherwise be a negativeshift / UB).
PTO2SharedMemoryRingHeadergrows from 256 → 576 bytes andPTO2SharedMemoryHeaderfrom 320 → 640 bytes (newcached_completed_watermarkarray + wider
completion_flagsentries); layoutstatic_asserts updatedaccordingly.
Performance
Comparison against main and #1619 for both device wall-clock and kernel only (as measured by tracr)