Update: modularisation of completion watermark together with a false-sharing alignment fix and batch updates.

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<int32_t>(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++;
 

src/a2a3/runtime/host_build_graph/runtime/pto_shared_memory.h

@@ -89,7 +89,7 @@ struct alignas(64) PTO2SharedMemoryRingHeader {
     alignas(64) std::atomic<int32_t> 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<uint8_t> *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"
 );
 

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<int32_t>(slot_state.task->task_id.local());
+        const int32_t task_id = static_cast<int32_t>(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

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];
@@ -1076,7 +1075,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);