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 e4055045c..df2831d89 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 @@ -1069,7 +1069,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..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 @@ -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 @@ -101,10 +101,44 @@ 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); + } + + // 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; + } 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); + } + } + } + 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]; } @@ -126,7 +160,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 2892ce651..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 @@ -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; } @@ -554,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.completion_flags[my_id & ring.task_window_mask].store(1, std::memory_order_release); + 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) { @@ -580,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.completion_flags[next & ring.task_window_mask].load(std::memory_order_acquire) == 0) 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 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 3cc67f03b..5afb75178 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 @@ -240,8 +240,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]; @@ -1074,7 +1073,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);