Refactor: host build graph, eager completion watermark, 4 AICPU schedulers

src/a2a3/runtime/host_build_graph/aicpu/aicpu_executor.cpp

@@ -335,11 +335,7 @@ int32_t AicpuExecutor::run(Runtime *runtime) {
             LOG_ERROR("Thread %d: rt is null after orchestrator error, skipping dispatch", thread_idx);
         } else {
             sched_ctx_.bind_runtime(rt);
-            // 3S+1P: the last thread is the core-less resolution (P) thread; the
-            // rest are core-owning schedulers (S).
-            int32_t completed = (thread_idx == sched_ctx_.p_thread_idx()) ?
-                                    sched_ctx_.run_resolution_thread(runtime, thread_idx) :
-                                    sched_ctx_.resolve_and_dispatch(runtime, thread_idx);
+            int32_t completed = sched_ctx_.resolve_and_dispatch(runtime, thread_idx);
             if (completed < 0) {
                 LOG_ERROR("Thread %d: Scheduler failed with rc=%d", thread_idx, completed);
                 run_rc = completed;

src/a2a3/runtime/host_build_graph/docs/RUNTIME_LOGIC.md

@@ -238,9 +238,10 @@ local-ids:
 
 - `fanin_local_ids[fanin_count]` — the local-ids of this task's direct producers
   (`fanin_count <= PTO2_MAX_FANIN`; there is no spill, so an overflow is fatal).
-- A per-slot `completion_flags[id]` byte in the ring header is the device-side
-  readiness truth: a task is ready iff every id in its `fanin_local_ids` has its
-  `completion_flags` byte set (see §8.2).
+- A per-slot `completion_flags[id]` entry (`int32`: `-1` pending, else the id
+  itself) in the ring header is the device-side readiness truth: a task is
+  ready iff every id in its `fanin_local_ids` has its `completion_flags` entry
+  set to its own id (see §8.2).
 
 Because the fanin is integer ids rather than pointers, the host→device image
 needs no fanout/dep-pool relocation — only the per-slot `task` / `payload`
@@ -407,7 +408,7 @@ slot state, not the payload; the host consumer-wait gates on it (§8.4).
 ### 6.2 Task State Machine
 
 `task_state` is the **host-visible mirror** of completion; the device-side
-readiness truth is the per-slot `completion_flags` byte (§8.2). The host polls
+readiness truth is the per-slot `completion_flags` entry (§8.2). The host polls
 `task_state` in `wait_for_tensor_ready`, the allocator deadlock detector, and
 the cold-path stall dump.
 
@@ -457,10 +458,12 @@ Key members:
 
 > **Note**: There is no fanout adjacency, dep-pool, or per-producer lock. A
 > producer inline-completed on the host (e.g. a hidden-alloc task) pre-sets its
-> own `completion_flags[id] = 1` in the H2D image so device consumers see it as
-> already satisfied. The only cross-task pointers the image carries are the
-> per-slot `task` / `payload` pointers, which `relocate_host_orch_image` shifts
-> to device addresses before H2D.
+> own `completion_flags[id] = id` in the H2D image so device consumers see it as
+> already satisfied, and calls `update_completed_watermark` itself right after —
+> the frontier-gating in §8.4 means no later on-device completion can advance the
+> watermark past this task on its behalf. The only cross-task pointers the image
+> carries are the per-slot `task` / `payload` pointers, which
+> `relocate_host_orch_image` shifts to device addresses before H2D.
 
 ### 7.3 Dependency Recording and Relocation
 
@@ -572,25 +575,78 @@ Each scheduler thread runs a tight loop with two main phases:
 - When `TASK_FIN_STATE` detected: call `on_subtask_complete`; when
   `completed_subtasks == total_required_subtasks`, call `on_mixed_task_complete`,
   which:
-  1. mirrors `task_state = COMPLETED` (host-visible) and sets the device
-     readiness truth `completion_flags[my_id] = 1` (release);
-  2. drains this task's intrusive **wake list** — `wake_list_head.exchange(SENTINEL)`
-     — reclassifying each waiter: a waiter whose remaining fanin is now all met
-     is pushed via `push_ready_routed`; otherwise it re-registers on its next
-     unmet producer. After the exchange the head is `SENTINEL`, so a consumer
-     registering concurrently re-checks the flags instead of being lost;
-  3. CAS-advances `completed_watermark` over the contiguous completed prefix
-     (§8.4).
+  1. mirrors `task_state = COMPLETED` (host-visible);
+  2. calls `try_set_completion_flag` to set the device readiness truth
+     `completion_flags[my_id] = my_id` (release). On failure (the previous
+     occupant of `my_id`'s flag-slot isn't yet certified, see **Flag-slot
+     reuse** below) it pushes `my_id` onto
+     `failed_heap_of_set_completion_flag[thread_idx]` and returns immediately
+     — steps 3 and 4 below are skipped and become `my_id`'s one deferred
+     chance, taken later by `retry_set_completion_flags`;
+  3. on success, drains this task's intrusive **wake list** —
+     `wake_list_head.exchange(SENTINEL)` — via the shared `drain_wake_list`
+     helper, reclassifying each waiter: a waiter whose remaining fanin is now
+     all met is pushed via `push_ready_routed`; otherwise it re-registers on
+     its next unmet producer. After the exchange the head is `SENTINEL`, so a
+     consumer registering concurrently re-checks the flags instead of being
+     lost;
+  4. calls `update_completed_watermark(thread_idx, my_id)`, which CAS-advances
+     `completed_watermark` over the contiguous completed prefix if — and only
+     if — `my_id` is exactly the current watermark (§8.4).
 
 **Readiness / wake registration.** A task is ready iff every id in its
-`fanin_local_ids` has its `completion_flags` byte set (`fanin_satisfied` /
-`classify_fanin_state`, acquire loads). A not-yet-ready task registers itself on
-its **first unmet** producer's wake list (`register_wake`); that producer's
-completion re-drives the classification. The decision is terminal — tasks are
-never re-polled — because `completion_flags` are monotonic. This wake machinery
-is seeded by the device **boot classify** (`on_orchestration_done`), which scans
-the submitted tasks once and either pushes the fanin-free ones to the ready
-queue or registers each remaining task on its first unmet producer.
+`fanin_local_ids` has its `completion_flags` entry set to its own id
+(`fanin_satisfied` / `classify_fanin_state`, acquire loads). A not-yet-ready
+task registers itself on its **first unmet** producer's wake list
+(`register_wake`); that producer's completion re-drives the classification.
+The decision is terminal — tasks are never re-polled — because
+`completion_flags` are monotonic. This wake machinery is seeded by the device
+**boot classify** (`on_orchestration_done`), which scans the submitted tasks
+once and either pushes the fanin-free ones to the ready queue or registers
+each remaining task on its first unmet producer.
+
+**Flag-slot reuse.** `completion_flags` holds `task_window_size` entries,
+indexed by `flag_index(local_id)` rather than the raw `local_id &
+task_window_mask`: `flag_index` swaps the low `shuffle_lower_bits` bits of
+`local_id & task_window_mask` into the high position (and the remaining high
+bits down into the low position), so consecutive `local_id`s land on
+different cachelines instead of packing `64 / sizeof(int32_t)` consecutive
+ids onto the same one — the layout `update_completed_watermark` scans
+linearly. This is a bijection on `[0, task_window_size)`, so the reuse
+period is unchanged: the array does not grow with the number of tasks a run
+submits, and the slot written for `local_id` is later reused for `local_id +
+task_window_size`. The number of tasks a run may submit is therefore not
+bounded by the array's size — only the reuse of a given slot is ordered: a
+write for `local_id` must not land before `completed_watermark` certifies
+`local_id - task_window_size` (the slot's previous occupant) complete.
+
+The device scheduler enforces this without blocking. `on_mixed_task_complete`
+calls `try_set_completion_flag`: it stores the flag and returns `true` when
+the previous occupant is already certified, or leaves the slot untouched and
+returns `false` otherwise. On `false` the thread pushes `task_id` onto
+`failed_heap_of_set_completion_flag[thread_idx]` (a per-thread min-heap)
+instead of spinning, and returns without draining `task_id`'s wake list or
+touching the watermark — both stay `task_id`'s one deferred chance until
+the dispatch loop's per-iteration call to `retry_set_completion_flags`
+retries the smallest pending id. Once that retry succeeds for `min_task`,
+`retry_set_completion_flags` drains `min_task`'s wake list (the same
+`drain_wake_list` routing `on_mixed_task_complete` uses on its success path)
+and then calls `update_completed_watermark(thread_idx, min_task)` (§8.4 — so
+each id still gets exactly one wake-list drain and one
+`update_completed_watermark` call, made right after its flag is actually
+set, whether that happens inline or later via retry). This replaces the
+older requirement — task `t` must finish before task `t + task_window_size`
+— with a weaker one: task `t` must not depend on a task with id `>= t +
+task_window_size`. That holds automatically whenever task ids follow a
+topological order of the dependency graph (the normal case — a task's fanins
+always carry smaller ids than the task itself), so the heap is expected to
+drain on its very next retry and exists only as a failsafe.
+
+`is_completion_flag_set` stays correct across reuse — whether or not an
+entry is currently sitting unresolved in the failed-heap — by falling back
+to `completed_watermark`: once a task's id is behind the watermark it is
+reported complete regardless of what a later lap has since written into its
+slot.
 
 **Early staging status.** Early producer propagation is currently disabled in
 HBG: `propagate_dispatch_fanin` is a stub, so the polling path does not populate
@@ -621,15 +677,47 @@ Ready queues use a lock-free bounded MPMC (Vyukov) design:
 
 ### 8.4 Completion Watermark (host consumer-wait gate)
 
-`completed_watermark` is the highest id such that every task in
-`[0, completed_watermark]` has its `completion_flags` byte set. The tail of
-`on_mixed_task_complete` CAS-advances it over the **full contiguous completed
-prefix** (bounded by `current_task_index`, not by the completing task's own id)
-— capping at `my_id` would make the final value completion-order-dependent and
-strand it below the true prefix.
+`completed_watermark` is the lowest id not yet guaranteed complete: every task
+in `[0, completed_watermark)` has its `completion_flags` entry set.
+
+Every completer calls `update_completed_watermark(thread_idx, my_id)` exactly
+once, and only *after* `my_id`'s own `completion_flags` entry is actually
+set — never before. The host orchestrator's inline-completed hidden-alloc
+task (§7.2) gets this for free: it calls the blocking `set_completion_flag`,
+which cannot return before the store lands, immediately followed by
+`update_completed_watermark`. The device path (`on_mixed_task_complete`)
+calls `try_set_completion_flag` (**Flag-slot reuse**, §8.2); when that
+returns `true` it calls `update_completed_watermark` right after, same as
+the host path. When it returns `false` (deferred to
+`failed_heap_of_set_completion_flag`), `on_mixed_task_complete` skips the
+`update_completed_watermark` call entirely — `task_id`'s one chance moves to
+`retry_set_completion_flags`, which calls
+`update_completed_watermark(thread_idx, min_task)` right after the retried
+`try_set_completion_flag` finally succeeds. Either way, the call that
+"belongs" to an id is made exactly once, and only once that id's flag is
+genuinely visible — a task's `update_completed_watermark` call must never
+run before the store its own walk trusts as already-set, since the walk
+skips re-checking `my_id`'s own `completion_flags` entry, so calling too
+early could advance the watermark past `my_id` before that write is
+visible.
+
+The call is **eager but frontier-gated**: every completer makes it (promptly
+on the host and the common device case, or after a deferred retry in the
+failed-heap case — §8.2), but `update_completed_watermark` is a no-op unless `my_id` equals the watermark
+it currently observes. Only the completer landing exactly at the frontier
+does the work, CAS-advancing over the **full contiguous completed prefix**
+(bounded by `current_task_index`, not by `my_id`) — capping at `my_id` would
+make the final value completion-order-dependent and strand it below the true
+prefix. A completer that finishes out of order (its id ahead of the
+watermark) defers: it does not retry or block, it relies on whichever thread
+later completes the frontier task to walk forward through its already-set
+flag. This is why the call cannot be skipped or deferred for *any* completion
+path, including host-side ones — a task sitting at the frontier with no
+completer ever calling `update_completed_watermark` for it stalls the
+watermark permanently.
 
 It is **load-bearing**: the host `wait_for_tensor_ready(..., wait_for_consumers)`
-gates on `completed_watermark >= producer.last_consumer_local_id` to observe
+gates on `completed_watermark > producer.last_consumer_local_id` to observe
 "every consumer of this producer has retired" — replacing the wiring model's
 `fanout_refcount == fanout_count` check.
 

src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_orchestrator.cpp

@@ -301,7 +301,7 @@ static bool append_fanin_or_fail(
     // Skip a stale/reused producer slot: the cached owner id no longer resolves
     // to this producer (defensive — whole-graph-resident hbg does not reuse slots
     // at build time). A COMPLETED producer IS a real fanin edge under polling (its
-    // completion_flags byte is set), so it is not skipped.
+    // completion_flags entry is set), so it is not skipped.
     if (prod_state->task == nullptr || prod_state->task->task_id.local() != producer_task_id.local()) {
         return true;
     }
@@ -316,7 +316,7 @@ static bool append_fanin_or_fail(
     fanin_builder->payload->fanin_local_ids[fanin_builder->count++] = static_cast<int32_t>(producer_task_id.local());
 
     // Reclaim gate: record this task as a consumer of the producer. The producer
-    // slot retires once the per-ring completed_watermark reaches this consumer id.
+    // slot retires once the per-ring completed_watermark exceeds this consumer id.
     if (fanin_builder->self_local > prod_state->last_consumer_local_id) {
         prod_state->last_consumer_local_id = fanin_builder->self_local;
     }
@@ -434,9 +434,9 @@ static bool prepare_task(
     out->slot_state->active_mask = active_mask;
     out->slot_state->task_attrs = task_attrs;
     // Reclaim gate: seed last_consumer to self, so a producer with no consumers
-    // is retirable once completed_watermark >= its own id. Each fanin edge bumps
-    // it in append_fanin_or_fail. completion_flags for this slot are already 0
-    // (zeroed once at init; whole-graph-resident hbg never reuses a slot).
+    // is retirable once completed_watermark > its own id. Each fanin edge bumps
+    // it in append_fanin_or_fail. completion_flags for this slot are already -1
+    // (set once at init; whole-graph-resident hbg never reuses a slot).
     out->slot_state->last_consumer_local_id = static_cast<int32_t>(out->task_id.local());
     // payload.fanin_count is set in submit_task_common's STEP 6.
     scope_tasks_push(orch, out->slot_state);
@@ -1061,15 +1061,22 @@ TaskOutputTensors PTO2OrchestratorState::alloc_tensors(const L0TaskArgs &args) {
         // unconditionally.
         prepared.slot_state->task_attrs.set_early_resolve(true);
         prepared.slot_state->mark_completed();  // host-visible task_state mirror
-        // Polling: pre-set the device-visible completion_flags byte in the H2D
+        // Polling: pre-set the device-visible completion_flags entry in the H2D
         // image. Consumers poll completion_flags (not task_state), so a hidden-alloc
         // producer completed here on the host must publish its flag too — otherwise
         // every consumer register_wakes on a producer that never runs on device and
-        // the run hangs. (The device watermark walk transparently steps past this
-        // pre-set flag when a later on-device task completes.)
+        // 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.
+        // Runs on the host, before any AICPU scheduler thread exists — pass
+        // PLATFORM_MAX_AICPU_THREADS (one past the last valid device thread
+        // index) as the thread_idx, marking this store as host-originated.
         PTO2SharedMemoryRingHeader &done_ring = orch->sm_header->ring;
         int32_t done_local = static_cast<int32_t>(prepared.task_id.local());
-        done_ring.set_completion_flag(done_local);
+        done_ring.set_completion_flag(PLATFORM_MAX_AICPU_THREADS, done_local);
+        done_ring.update_completed_watermark(PLATFORM_MAX_AICPU_THREADS, done_local);
     }
     orch->inline_completed_tasks++;
 

src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_runtime2.cpp

@@ -160,11 +160,11 @@ static bool wait_for_tensor_ready(PTO2Runtime *rt, const Tensor &tensor, bool wa
         uint64_t t0 = get_sys_cnt_aicpu();
         int32_t spin_count = 0;
         // Polling: all consumers of this producer have retired once the per-ring
-        // completed_watermark reaches the producer's highest consumer id (set at
+        // completed_watermark exceeds the producer's highest consumer id (set at
         // submit in append_fanin_or_fail). Replaces the fanout_refcount ==
         // fanout_count wiring check, which polling removes.
         PTO2SharedMemoryRingHeader &cons_ring = orch.sm_header->ring;
-        while (cons_ring.completed_watermark.load(std::memory_order_acquire) < slot.last_consumer_local_id) {
+        while (cons_ring.completed_watermark.load(std::memory_order_acquire) <= slot.last_consumer_local_id) {
             SPIN_WAIT_HINT();
             if ((++spin_count & 1023) == 0) {
                 // A fatal latched elsewhere (e.g. the scheduler-side wiring

src/a2a3/runtime/host_build_graph/runtime/pto_async_wait.h

@@ -171,9 +171,7 @@ struct AsyncWaitList {
     struct DrainCompletionSink {
         PTO2SchedulerState *sched{nullptr};
         int32_t inline_completed{0};
-#if SIMPLER_SCHED_PROFILING
         int32_t thread_idx{0};
-#endif
 
         bool can_inline_complete() const { return sched != nullptr; }
     };
@@ -283,13 +281,8 @@ struct AsyncWaitList {
     }
 
     template <bool Profiling>
-    AsyncPollResult poll_and_complete(
-        AICoreCompletionMailbox *aicore_mailbox, PTO2SchedulerState *sched
-#if SIMPLER_SCHED_PROFILING
-        ,
-        int thread_idx
-#endif
-    );
+    AsyncPollResult
+    poll_and_complete(AICoreCompletionMailbox *aicore_mailbox, PTO2SchedulerState *sched, int32_t thread_idx);
 };
 
 #endif  // PTO_ASYNC_WAIT_H

src/a2a3/runtime/host_build_graph/runtime/pto_runtime2_types.h

@@ -444,7 +444,7 @@ static_assert(
  */
 struct alignas(64) PTO2TaskSlotState {
     // Highest local task id among this slot's consumers. Reclaim gate: the slot
-    // is safe to retire once the per-ring completed_watermark reaches this id.
+    // is safe to retire once the per-ring completed_watermark exceeds this id.
     // Whole-graph-resident hbg never reclaims at runtime, so this is
     // inert-but-scaffolded for parity. Seeded to own local_id in prepare_task;
     // bumped via max() at submit for each consumer.