From 0b39d2976bb8510f9f927e2e6f67d0d0b8e0f0e1 Mon Sep 17 00:00:00 2001 From: raphaelsteiner Date: Fri, 24 Jul 2026 10:34:42 +0200 Subject: [PATCH 1/5] modularised completion flags Signed-off-by: raphaelsteiner --- .../runtime/orchestrator_core/pto_orchestrator.cpp | 2 +- .../host_build_graph/runtime/pto_shared_memory.h | 10 +++++++++- .../runtime/scheduler/pto_scheduler.h | 12 ++++-------- .../runtime/scheduler/scheduler_cold_path.cpp | 5 ++--- 4 files changed, 16 insertions(+), 13 deletions(-) diff --git a/src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_orchestrator.cpp b/src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_orchestrator.cpp index df383b864..32e8ebe55 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_orchestrator.cpp +++ b/src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_orchestrator.cpp @@ -1064,7 +1064,7 @@ TaskOutputTensors PTO2OrchestratorState::alloc_tensors(const L0TaskArgs &args) { // pre-set flag when a later on-device task completes.) PTO2SharedMemoryRingHeader &done_ring = orch->sm_header->ring; int32_t done_local = static_cast(prepared.task_id.local()); - done_ring.completion_flags[done_local & done_ring.task_window_mask].store(1, std::memory_order_release); + done_ring.set_completion_flag(done_local); } orch->inline_completed_tasks++; diff --git a/src/a2a3/runtime/host_build_graph/runtime/pto_shared_memory.h b/src/a2a3/runtime/host_build_graph/runtime/pto_shared_memory.h index 1d861b6cf..17f7d7da5 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/pto_shared_memory.h +++ b/src/a2a3/runtime/host_build_graph/runtime/pto_shared_memory.h @@ -101,10 +101,18 @@ struct alignas(64) PTO2SharedMemoryRingHeader { // Polling-completion state (device-addressed array, one byte per slot). // 0 = pending, 1 = task fully COMPLETED. Writer = the task's completer at - // on_mixed_task_complete; reader = consumer fanin polling (fanin_satisfied). + // on_mixed_task_complete; reader = consumer fanin polling (is_completion_flag_set). // Zeroed host-side at init. Indexed by local_id & task_window_mask. std::atomic *completion_flags; + bool is_completion_flag_set(int32_t local_id, std::memory_order order = std::memory_order_acquire) const { + return completion_flags[local_id & task_window_mask].load(order) != 0; + } + + void set_completion_flag(int32_t local_id, std::memory_order order = std::memory_order_release) const { + completion_flags[local_id & task_window_mask].store(1, order); + } + int32_t get_slot_by_task_id(int32_t local_task_id) { return local_task_id & task_window_mask; } PTO2TaskDescriptor &get_task_by_slot(int32_t slot) { return task_descriptors[slot]; } diff --git a/src/a2a3/runtime/host_build_graph/runtime/scheduler/pto_scheduler.h b/src/a2a3/runtime/host_build_graph/runtime/scheduler/pto_scheduler.h index 2892ce651..c8bbd11bb 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/scheduler/pto_scheduler.h +++ b/src/a2a3/runtime/host_build_graph/runtime/scheduler/pto_scheduler.h @@ -500,9 +500,7 @@ struct PTO2SchedulerState { const PTO2TaskPayload &p = *s->payload; const PTO2SharedMemoryRingHeader &ring = *ring_sched_state.ring; for (int32_t i = 0; i < p.fanin_count; i++) { - if (ring.completion_flags[p.fanin_local_ids[i] & ring.task_window_mask].load(std::memory_order_acquire) == - 0) - return false; + if (!ring.is_completion_flag_set(p.fanin_local_ids[i])) return false; } return true; } @@ -515,9 +513,7 @@ struct PTO2SchedulerState { const PTO2TaskPayload &p = *s->payload; const PTO2SharedMemoryRingHeader &ring = *ring_sched_state.ring; for (int32_t i = 0; i < p.fanin_count; i++) { - if (ring.completion_flags[p.fanin_local_ids[i] & ring.task_window_mask].load(std::memory_order_acquire) == - 0) - return i; + if (!ring.is_completion_flag_set(p.fanin_local_ids[i])) return i; } return -1; } @@ -558,7 +554,7 @@ struct PTO2SchedulerState { PTO2SharedMemoryRingHeader &ring = *ring_sched_state.ring; slot_state.mark_completed(); // host-visible mirror (task_state = COMPLETED) - ring.completion_flags[my_id & ring.task_window_mask].store(1, std::memory_order_release); + ring.set_completion_flag(my_id); PTO2TaskSlotState *waiter = slot_state.wake_list_head.exchange(WAKE_LIST_SENTINEL, std::memory_order_acq_rel); while (waiter != nullptr && waiter != WAKE_LIST_SENTINEL) { @@ -588,7 +584,7 @@ struct PTO2SchedulerState { int32_t w = ring.completed_watermark.load(std::memory_order_acquire); while (w + 1 < submitted) { int32_t next = w + 1; - if (ring.completion_flags[next & ring.task_window_mask].load(std::memory_order_acquire) == 0) break; + if (!ring.is_completion_flag_set(next)) break; if (ring.completed_watermark.compare_exchange_weak( w, next, std::memory_order_acq_rel, std::memory_order_acquire )) { diff --git a/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_cold_path.cpp b/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_cold_path.cpp index c7bc14c22..f8b17ecbc 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_cold_path.cpp +++ b/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_cold_path.cpp @@ -241,8 +241,7 @@ void SchedulerContext::log_stall_diagnostics( if (slot_state.payload != nullptr) { for (int32_t k = 0; k < fi; k++) { int32_t pid = slot_state.payload->fanin_local_ids[k]; - if (ring.completion_flags[pid & ring.task_window_mask].load(std::memory_order_relaxed) != 0) - rc++; + if (ring.is_completion_flag_set(pid, std::memory_order_relaxed)) rc++; } } int32_t kid_aic = slot_state.task->kernel_id[0]; @@ -1092,7 +1091,7 @@ void SchedulerContext::on_orchestration_done( PTO2SharedMemoryRingHeader &ring = *sched_->ring_sched_state.ring; const int32_t submitted = ring.fc.current_task_index.load(std::memory_order_acquire); for (int32_t id = 0; id < submitted; id++) { - if (ring.completion_flags[id & ring.task_window_mask].load(std::memory_order_acquire) != 0) { + if (ring.is_completion_flag_set(id)) { continue; // completed on the host (hidden alloc); nothing to dispatch } PTO2TaskSlotState &s = ring.get_slot_state_by_task_id(id); From 58e0317b8e234b6bfa78b5d3fd0fb74f62ea1c4f Mon Sep 17 00:00:00 2001 From: raphaelsteiner Date: Fri, 24 Jul 2026 14:19:05 +0200 Subject: [PATCH 2/5] padding Signed-off-by: raphaelsteiner --- src/a2a3/runtime/host_build_graph/runtime/pto_shared_memory.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/a2a3/runtime/host_build_graph/runtime/pto_shared_memory.h b/src/a2a3/runtime/host_build_graph/runtime/pto_shared_memory.h index 17f7d7da5..d4ccfb0d7 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/pto_shared_memory.h +++ b/src/a2a3/runtime/host_build_graph/runtime/pto_shared_memory.h @@ -89,7 +89,7 @@ struct alignas(64) PTO2SharedMemoryRingHeader { alignas(64) std::atomic completed_watermark; // Layout metadata (set once at init) - uint64_t task_window_size; + alignas(64) uint64_t task_window_size; int32_t task_window_mask; uint64_t heap_size; uint64_t task_descriptors_offset; // Offset from SM base, in bytes From 4ffdfed318d20e0475b06ea6d0b852ae1b4d3c2a Mon Sep 17 00:00:00 2001 From: raphaelsteiner Date: Mon, 27 Jul 2026 10:41:08 +0200 Subject: [PATCH 3/5] batched watermark update Signed-off-by: raphaelsteiner --- .../runtime/pto_shared_memory.h | 20 ++++++++++++++++++- .../runtime/scheduler/pto_scheduler.h | 18 ++++------------- 2 files changed, 23 insertions(+), 15 deletions(-) diff --git a/src/a2a3/runtime/host_build_graph/runtime/pto_shared_memory.h b/src/a2a3/runtime/host_build_graph/runtime/pto_shared_memory.h index d4ccfb0d7..9b13b2c8a 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/pto_shared_memory.h +++ b/src/a2a3/runtime/host_build_graph/runtime/pto_shared_memory.h @@ -113,6 +113,24 @@ struct alignas(64) PTO2SharedMemoryRingHeader { completion_flags[local_id & task_window_mask].store(1, order); } + // set completion flag first before updating the watermark (logic requirement) + void update_completed_watermark() { + int32_t curr_watermark = completed_watermark.load(std::memory_order_acquire); + const int32_t submitted = fc.current_task_index.load(std::memory_order_acquire); + + int32_t next = curr_watermark; + while (true) { + while (next + 1 < submitted && is_completion_flag_set(next + 1)) { + ++next; + } + if (next == curr_watermark) { return; } + + if (completed_watermark.compare_exchange_strong(curr_watermark, next, std::memory_order_acq_rel, std::memory_order_acquire)) { + curr_watermark = next; + } + } + } + int32_t get_slot_by_task_id(int32_t local_task_id) { return local_task_id & task_window_mask; } PTO2TaskDescriptor &get_task_by_slot(int32_t slot) { return task_descriptors[slot]; } @@ -134,7 +152,7 @@ struct alignas(64) PTO2SharedMemoryRingHeader { static_assert(sizeof(PTO2SharedMemoryRingHeader) == 256, "PTO2SharedMemoryRingHeader layout drift"); static_assert( - offsetof(PTO2SharedMemoryRingHeader, task_descriptors_offset) == 160, + offsetof(PTO2SharedMemoryRingHeader, task_descriptors_offset) == 216, "PTO2SharedMemoryRingHeader task_descriptors_offset layout drift" ); diff --git a/src/a2a3/runtime/host_build_graph/runtime/scheduler/pto_scheduler.h b/src/a2a3/runtime/host_build_graph/runtime/scheduler/pto_scheduler.h index c8bbd11bb..720b38f4f 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/scheduler/pto_scheduler.h +++ b/src/a2a3/runtime/host_build_graph/runtime/scheduler/pto_scheduler.h @@ -550,11 +550,11 @@ struct PTO2SchedulerState { // watermark >= producer.last_consumer_local_id). Whole-graph-resident hbg // has no device slot reclaim, so no advance_ring_pointers here. void on_mixed_task_complete(PTO2TaskSlotState &slot_state) { - const int32_t my_id = static_cast(slot_state.task->task_id.local()); + const int32_t task_id = static_cast(slot_state.task->task_id.local()); PTO2SharedMemoryRingHeader &ring = *ring_sched_state.ring; slot_state.mark_completed(); // host-visible mirror (task_state = COMPLETED) - ring.set_completion_flag(my_id); + ring.set_completion_flag(task_id); PTO2TaskSlotState *waiter = slot_state.wake_list_head.exchange(WAKE_LIST_SENTINEL, std::memory_order_acq_rel); while (waiter != nullptr && waiter != WAKE_LIST_SENTINEL) { @@ -576,21 +576,11 @@ struct PTO2SchedulerState { // completed_watermark = highest id such that every task in [0, watermark] // has its completion_flags byte set. The host wait_for_consumers gates on // watermark >= producer.last_consumer_local_id, so the walk must extend to - // the full contiguous completed prefix — NOT cap at my_id. Capping at my_id + // the full contiguous completed prefix — NOT cap at task_id. Capping at task_id // makes the final value order-dependent: a low-id task completing after a // higher one would leave the watermark stuck below the true prefix, hanging // any wait_for_consumers whose last_consumer sits in the gap. - const int32_t submitted = ring.fc.current_task_index.load(std::memory_order_acquire); - int32_t w = ring.completed_watermark.load(std::memory_order_acquire); - while (w + 1 < submitted) { - int32_t next = w + 1; - if (!ring.is_completion_flag_set(next)) break; - if (ring.completed_watermark.compare_exchange_weak( - w, next, std::memory_order_acq_rel, std::memory_order_acquire - )) { - w = next; - } - } + ring.update_completed_watermark(); } // Polling: there is no ready-claim CAS (a producer routes each waiter exactly From f73e19ee52f65103ad159e33c6e7ef550bc219b8 Mon Sep 17 00:00:00 2001 From: raphaelsteiner Date: Mon, 27 Jul 2026 13:00:47 +0200 Subject: [PATCH 4/5] precommit Signed-off-by: raphaelsteiner --- .../runtime/host_build_graph/runtime/pto_shared_memory.h | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/src/a2a3/runtime/host_build_graph/runtime/pto_shared_memory.h b/src/a2a3/runtime/host_build_graph/runtime/pto_shared_memory.h index 9b13b2c8a..5645f2ded 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/pto_shared_memory.h +++ b/src/a2a3/runtime/host_build_graph/runtime/pto_shared_memory.h @@ -123,9 +123,13 @@ struct alignas(64) PTO2SharedMemoryRingHeader { while (next + 1 < submitted && is_completion_flag_set(next + 1)) { ++next; } - if (next == curr_watermark) { return; } + if (next == curr_watermark) { + return; + } - if (completed_watermark.compare_exchange_strong(curr_watermark, next, std::memory_order_acq_rel, std::memory_order_acquire)) { + if (completed_watermark.compare_exchange_strong( + curr_watermark, next, std::memory_order_acq_rel, std::memory_order_acquire + )) { curr_watermark = next; } } From 825f93c7707df9fba0e9d9b18c0b756e305fbaf9 Mon Sep 17 00:00:00 2001 From: raphaelsteiner Date: Mon, 27 Jul 2026 13:14:25 +0200 Subject: [PATCH 5/5] micro optimisation Signed-off-by: raphaelsteiner --- src/a2a3/runtime/host_build_graph/runtime/pto_shared_memory.h | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/a2a3/runtime/host_build_graph/runtime/pto_shared_memory.h b/src/a2a3/runtime/host_build_graph/runtime/pto_shared_memory.h index 5645f2ded..bd404c3bc 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/pto_shared_memory.h +++ b/src/a2a3/runtime/host_build_graph/runtime/pto_shared_memory.h @@ -131,6 +131,10 @@ struct alignas(64) PTO2SharedMemoryRingHeader { curr_watermark, next, std::memory_order_acq_rel, std::memory_order_acquire )) { curr_watermark = next; + } else { + // The acquire release semantics of the successful CAS guarantee that in the case of failure this thread + // also synchronises with the thread reporting the completion through the intermediary thread(s). + next = std::max(next, curr_watermark); } } }