diff --git a/src/a2a3/runtime/host_build_graph/aicpu/aicpu_executor.cpp b/src/a2a3/runtime/host_build_graph/aicpu/aicpu_executor.cpp index ccc76b6b4c..442eb4584e 100644 --- a/src/a2a3/runtime/host_build_graph/aicpu/aicpu_executor.cpp +++ b/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; diff --git a/src/a2a3/runtime/host_build_graph/docs/RUNTIME_LOGIC.md b/src/a2a3/runtime/host_build_graph/docs/RUNTIME_LOGIC.md index cd10a5409f..07be2addd0 100644 --- a/src/a2a3/runtime/host_build_graph/docs/RUNTIME_LOGIC.md +++ b/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. 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 df2831d891..67fe7a43ea 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 @@ -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(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(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(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++; diff --git a/src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_runtime2.cpp b/src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_runtime2.cpp index 41e7396cdf..e028edc707 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_runtime2.cpp +++ b/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 diff --git a/src/a2a3/runtime/host_build_graph/runtime/pto_async_wait.h b/src/a2a3/runtime/host_build_graph/runtime/pto_async_wait.h index 4a7cf9ffcf..15ddf8f070 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/pto_async_wait.h +++ b/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 - 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 diff --git a/src/a2a3/runtime/host_build_graph/runtime/pto_runtime2_types.h b/src/a2a3/runtime/host_build_graph/runtime/pto_runtime2_types.h index 933c14b1a4..2bd8247228 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/pto_runtime2_types.h +++ b/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. 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 bd404c3bc2..97c2575d4a 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 @@ -33,7 +33,9 @@ #pragma once #include +#include +#include "common/platform_config.h" #include "utils/device_arena.h" #include "pto_runtime2_types.h" @@ -80,17 +82,36 @@ static_assert(sizeof(PTO2RingFlowControl) == 128, "PTO2RingFlowControl must be e struct alignas(64) PTO2SharedMemoryRingHeader { PTO2RingFlowControl fc; - // Highest task_id such that every task with id in [0, completed_watermark] - // has its completion_flags byte set. Advanced over the full contiguous - // completed prefix at task-completion time (on_mixed_task_complete). The host - // consumer-wait gates on it: a producer slot P's consumers have all retired - // once completed_watermark >= P.last_consumer_local_id. On its own cache line - // (concurrent CAS-advance by completing threads). + // Lowest task_id not yet guaranteed complete: every task with id in + // [0, completed_watermark) has its completion_flags entry set. Every completer + // (device on_mixed_task_complete, or the host orchestrator for an inline-completed + // hidden-alloc task) calls update_completed_watermark exactly once, right after its + // own flag is actually set — never before. When on_mixed_task_complete's + // try_set_completion_flag fails, the call moves to retry_set_completion_flags, which + // makes it once the retried try_set_completion_flag succeeds instead (docs/ + // RUNTIME_LOGIC.md §8.4); only the one landing exactly at the current watermark walks + // it forward over the contiguous completed prefix. The host consumer-wait gates on it: a producer + // slot P's consumers have all retired once completed_watermark > + // P.last_consumer_local_id. On its own cache line (concurrent CAS-advance by + // completing threads). alignas(64) std::atomic completed_watermark; + struct alignas(64) AlignedInt32 { + int32_t val_{0u}; + int8_t pad_[64 - sizeof(int32_t)]; + }; + + // Embedded by value, not behind a unique_ptr: this header is placed by + // memset over raw arena/shared memory (PTO2SharedMemoryHandle::init_header, + // create_and_init_default), never via a constructor call, so a smart-pointer + // member's default initializer would never run and would dereference null. + alignas(64) mutable std::array cached_completed_watermark; + // Layout metadata (set once at init) alignas(64) uint64_t task_window_size; int32_t task_window_mask; + static constexpr int32_t shuffle_lower_bits{4}; // At least 0 and at most log_2(task_window_size) + int32_t shuffle_higher_bits; // log_2(task_window_size) - shuffle_lower_bits uint64_t heap_size; uint64_t task_descriptors_offset; // Offset from SM base, in bytes @@ -99,43 +120,135 @@ struct alignas(64) PTO2SharedMemoryRingHeader { PTO2TaskPayload *task_payloads; PTO2TaskSlotState *slot_states; - // 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 (is_completion_flag_set). - // Zeroed host-side at init. Indexed by local_id & task_window_mask. - std::atomic *completion_flags; + // Polling-completion state (device-addressed array, one int32 per slot, + // task_window_size slots total). -1 = pending, local_id = task fully + // COMPLETED. -1 (not 0) is the pending sentinel because local_id 0 is a + // valid completion stamp (the first task ever submitted) and must stay + // distinguishable from "not yet completed". Writer = the task's completer + // at on_mixed_task_complete; reader = consumer fanin polling + // (is_completion_flag_set). Set to -1 host-side at init. Indexed by + // flag_index(local_id), a bit-reindexing of local_id & task_window_mask + // (same task_window_size-entry period, different physical slot — see + // flag_index), so the entry for local_id is reused by local_id + + // task_window_size once the run submits more tasks than the array has + // slots for — see set_completion_flag / try_set_completion_flag for the + // ordering this reuse requires. + std::atomic *completion_flags; + + // To have subsequent tasks on different cachelines one needs shuffle_higher_bits >= 4 (because 64 / sizeof(int32_t) + // = 2^4) and shuffle_lower_bits >= 1. Every 2^shuffle_lower_bits will typically be on the same cacheline which is + // good for update_completion_watermark + constexpr int32_t flag_index(const int32_t local_id) const { + int32_t low_bits = local_id & ((static_cast(1) << shuffle_lower_bits) - static_cast(1)); + low_bits <<= shuffle_higher_bits; + + int32_t high_bits = local_id & task_window_mask; + high_bits >>= shuffle_lower_bits; + + return low_bits | high_bits; + } - 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; + // Once local_id falls behind completed_watermark it is reported set for + // the rest of the run, even after a later lap overwrites its slot: the + // watermark fallback is itself monotonic and outlives the raw slot value. + bool is_completion_flag_set( + const int32_t thread_idx, const int32_t local_id, std::memory_order order = std::memory_order_acquire + ) const { + int32_t &cached_cw = cached_completed_watermark[thread_idx].val_; + return (local_id < cached_cw) || (completion_flags[flag_index(local_id)].load(order) == local_id) || + (local_id < (cached_cw = completed_watermark.load(std::memory_order_acquire))); } - 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); + // local_id and local_id - task_window_size share a slot (flag_index wraps + // modulo task_window_size). earlier_task is one past that slot's previous + // occupant, so the wait below blocks the store until completed_watermark + // certifies the previous occupant complete: task t must complete before + // task t + task_window_size may publish its own entry into the shared + // slot. This bounds how far a slot can be reused ahead of retirement, not + // how many tasks a run may submit in total. + // Must be followed by a call to update_completed_watermark with same thread_idx and local_id (logical requirement) + // Only the host orchestrator's inline hidden-alloc completion path (submit_task, + // §7.2) still calls this blocking form; the device scheduler uses + // try_set_completion_flag below instead. + void set_completion_flag( + const int32_t thread_idx, const int32_t local_id, std::memory_order order = std::memory_order_release + ) const { + const int32_t earlier_task = local_id - task_window_mask; + int32_t &cached_cw = cached_completed_watermark[thread_idx].val_; + while ((cached_cw < earlier_task) && + ((cached_cw = completed_watermark.load(std::memory_order_acquire)) < earlier_task)) { + SPIN_WAIT_HINT(); + } + completion_flags[flag_index(local_id)].store(local_id, order); + } + + // Non-blocking counterpart to set_completion_flag: same slot-reuse condition + // (the previous occupant of local_id's aliased slot must be certified by + // completed_watermark first), but returns false instead of spinning when it + // is not yet certified, leaving completion_flags untouched. The device + // scheduler defers a false return to failed_heap_of_set_completion_flag and + // retries later rather than blocking the dispatch loop — see + // docs/RUNTIME_LOGIC.md "Flag-slot reuse" (§8.2) for the relaxed + // correctness argument this relies on. Must be followed by a call to + // update_completed_watermark with the same thread_idx and local_id only when + // this returns true — on false, the caller must defer that call to whichever + // retry finally succeeds (never call it against a local_id whose flag isn't + // set yet, §8.4). + bool try_set_completion_flag( + const int32_t thread_idx, const int32_t local_id, std::memory_order order = std::memory_order_release + ) { + const int32_t earlier_task = local_id - task_window_mask; + int32_t &cached_cw = cached_completed_watermark[thread_idx].val_; + if (cached_cw < earlier_task) { + cached_cw = completed_watermark.load(std::memory_order_acquire); + if (cached_cw < earlier_task) { + return false; + } + } + completion_flags[flag_index(local_id)].store(local_id, order); + return true; } - // set completion flag first before updating the watermark (logic requirement) - void update_completed_watermark() { + // Must be called for `local_id` only after set_completion_flag(thread_idx, local_id): + // the walk below never checks local_id's own completion_flags entry, it trusts the + // caller already set it, so calling out of order could advance the watermark past + // local_id before that write is actually visible. + void update_completed_watermark(const int32_t thread_idx, const int32_t local_id) { + // Thread fence required such that setting the flag of task local_id is written to GM before completed_watermark + // is read + std::atomic_thread_fence(std::memory_order_seq_cst); + int32_t curr_watermark = completed_watermark.load(std::memory_order_acquire); + // No-op unless this call sits exactly at the frontier: every completer calls this + // eagerly, but only the one whose own local_id equals the current watermark performs + // the advance walk below. A completer that finishes out of order (local_id ahead of + // the watermark) defers to whichever thread later completes the frontier task. + if (curr_watermark != local_id) { + cached_completed_watermark[thread_idx].val_ = curr_watermark; + return; + } + const int32_t submitted = fc.current_task_index.load(std::memory_order_acquire); - int32_t next = curr_watermark; + int32_t next = local_id + 1; while (true) { - while (next + 1 < submitted && is_completion_flag_set(next + 1)) { + while (next < submitted && is_completion_flag_set(thread_idx, next)) { ++next; } - if (next == curr_watermark) { - return; - } - if (completed_watermark.compare_exchange_strong( + if (not 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); + // In case of failure, we hand off the update responsibility to the thread who succeeded + cached_completed_watermark[thread_idx].val_ = curr_watermark; + break; + } + // Thread fence required such that completed_watermark is written to GM before the next flag is read + std::atomic_thread_fence(std::memory_order_seq_cst); + if (not(next < submitted && is_completion_flag_set(thread_idx, next))) { + break; } + curr_watermark = next++; } } @@ -158,9 +271,9 @@ struct alignas(64) PTO2SharedMemoryRingHeader { } }; -static_assert(sizeof(PTO2SharedMemoryRingHeader) == 256, "PTO2SharedMemoryRingHeader layout drift"); +static_assert(sizeof(PTO2SharedMemoryRingHeader) == 576, "PTO2SharedMemoryRingHeader layout drift"); static_assert( - offsetof(PTO2SharedMemoryRingHeader, task_descriptors_offset) == 216, + offsetof(PTO2SharedMemoryRingHeader, task_descriptors_offset) == 536, "PTO2SharedMemoryRingHeader task_descriptors_offset layout drift" ); @@ -192,10 +305,10 @@ struct alignas(PTO2_ALIGN_SIZE) PTO2SharedMemoryHeader { std::atomic sched_error_thread; // Thread index of last error writer }; -static_assert(sizeof(PTO2SharedMemoryHeader) == 320, "PTO2SharedMemoryHeader layout drift"); -static_assert(offsetof(PTO2SharedMemoryHeader, total_size) == 264, "PTO2SharedMemoryHeader total_size layout drift"); +static_assert(sizeof(PTO2SharedMemoryHeader) == 640, "PTO2SharedMemoryHeader layout drift"); +static_assert(offsetof(PTO2SharedMemoryHeader, total_size) == 584, "PTO2SharedMemoryHeader total_size layout drift"); static_assert( - offsetof(PTO2SharedMemoryHeader, orch_error_code) == 272, "PTO2SharedMemoryHeader orch_error_code layout drift" + offsetof(PTO2SharedMemoryHeader, orch_error_code) == 592, "PTO2SharedMemoryHeader orch_error_code layout drift" ); // ============================================================================= @@ -251,8 +364,8 @@ struct PTO2SharedMemoryHandle { bool validate(); private: - void init_header(uint64_t task_window_size, uint64_t heap_size); - void init_header_per_ring( + bool init_header(uint64_t task_window_size, uint64_t heap_size); + bool init_header_per_ring( const uint64_t task_window_sizes[PTO2_MAX_RING_DEPTH], const uint64_t heap_sizes[PTO2_MAX_RING_DEPTH] ); void setup_pointers(uint64_t task_window_size); @@ -308,7 +421,7 @@ struct PTO2RingSegmentOffsets { uint64_t descriptors; uint64_t payloads; uint64_t slot_states; - uint64_t completion_flags; // polling-completion byte array (1 byte/slot) + uint64_t completion_flags; // polling-completion flag array (1 int32/slot) uint64_t end; // offset just past completion_flags (total SM size) }; @@ -328,7 +441,7 @@ inline PTO2RingSegmentOffsets ring_segment_offsets(uint64_t task_window_size) no o.slot_states = off; off += PTO2_ALIGN_UP(task_window_size * sizeof(PTO2TaskSlotState), PTO2_ALIGN_SIZE); o.completion_flags = off; - off += PTO2_ALIGN_UP(task_window_size * sizeof(std::atomic), PTO2_ALIGN_SIZE); + off += PTO2_ALIGN_UP(task_window_size * sizeof(std::atomic), PTO2_ALIGN_SIZE); o.end = off; return o; } @@ -348,9 +461,9 @@ inline PTO2TaskSlotState *ring_slot_states_addr(void *sm_dev_base, uint64_t task ); } -// Device address of the polling completion_flags byte array. -inline std::atomic *ring_completion_flags_addr(void *sm_dev_base, uint64_t task_window_size) noexcept { - return reinterpret_cast *>( +// Device address of the polling completion_flags array. +inline std::atomic *ring_completion_flags_addr(void *sm_dev_base, uint64_t task_window_size) noexcept { + return reinterpret_cast *>( static_cast(sm_dev_base) + ring_segment_offsets(task_window_size).completion_flags ); } 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 3f20a188fd..705acd382c 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 @@ -15,7 +15,7 @@ * The Scheduler is responsible for: * 1. Maintaining per-resource-shape ready queues * 2. Polling-completion dependency resolution: a task is ready when every - * producer named in its inline fanin has set its completion_flags byte; + * producer named in its inline fanin has set its completion_flags entry; * a producer publishes completion + drains its wake list on finish * 3. Publishing the host-visible task_state mirror (PENDING -> COMPLETED) and * advancing the per-ring completed_watermark (consumer-retirement signal) @@ -30,8 +30,12 @@ #pragma once +#include #include +#include +#include +#include "common.h" // always_assert #include "common/core_type.h" #include "common/memory_barrier.h" #include "utils/device_arena.h" @@ -52,6 +56,10 @@ } while (0) #endif +#ifndef unlikely +#define unlikely(x) __builtin_expect(!!(x), 0) +#endif + // ============================================================================= // Ready Queue (Lock-free bounded MPMC — Vyukov design) // ============================================================================= @@ -463,6 +471,48 @@ struct PTO2SchedulerState { // the dispatch loop and completed inline -- never goes to AICore. PTO2ReadyQueue dummy_ready_queue; + // POD dynamic min-heap: entries/count/capacity all start at 0 under the + // arena's memset — a valid empty state, unlike std::vector's non-trivial + // constructor. Grows via realloc on push instead of a fixed cap. + // destroy() frees the buffer; PTO2SchedulerState::destroy() must reach + // every thread's instance or the buffer leaks. + struct FailedCompletionFlagHeap { + int32_t *entries = nullptr; + int32_t count = 0; + int32_t capacity = 0; + + void push(int32_t task_id) { + if (count == capacity) { + int32_t new_capacity = capacity == 0 ? 8 : capacity * 2; + int32_t *grown = + static_cast(realloc(entries, static_cast(new_capacity) * sizeof(int32_t))); + always_assert(grown != nullptr); + entries = grown; + capacity = new_capacity; + } + entries[count++] = task_id; + std::push_heap(entries, entries + count, std::greater<>{}); + } + + int32_t pop_min() { + std::pop_heap(entries, entries + count, std::greater<>{}); + return entries[--count]; + } + + bool empty() const { return count == 0; } + + void destroy() { + free(entries); + entries = nullptr; + count = 0; + capacity = 0; + } + }; + + // Heap of tasks for which setting the completion flag failed. Should always + // be empty. This only exists as a failsafe. + FailedCompletionFlagHeap failed_heap_of_set_completion_flag[PLATFORM_MAX_AICPU_THREADS]; + alignas(64) AsyncWaitList async_wait_list; // Statistics (cold path, isolated from hot-path fields) @@ -493,14 +543,14 @@ struct PTO2SchedulerState { // ---- Polling completion primitives (single-ring hbg) ---------------------- // Readiness: a task is ready iff every producer named in its inline fanin has - // set its completion_flags byte. Single-ring: all producers are ring 0, so + // set its completion_flags entry. Single-ring: all producers are ring 0, so // there is no per-edge ring indirection. - bool fanin_satisfied(const PTO2TaskSlotState *s) const { + bool fanin_satisfied(int32_t thread_idx, const PTO2TaskSlotState *s) const { const PTO2TaskPayload &p = *s->payload; const PTO2SharedMemoryRingHeader &ring = *ring_sched_state.ring; for (int32_t i = 0; i < p.fanin_count; i++) { - if (!ring.is_completion_flag_set(p.fanin_local_ids[i])) return false; + if (!ring.is_completion_flag_set(thread_idx, p.fanin_local_ids[i])) return false; } return true; } @@ -509,11 +559,11 @@ struct PTO2SchedulerState { // or the index of the first unmet fanin (register on that producer's wake // list). The decision is terminal: tasks are never re-polled; a producer's // completion re-scans its waiters via on_mixed_task_complete's wake drain. - int classify_fanin_state(const PTO2TaskSlotState *s) const { + int classify_fanin_state(int32_t thread_idx, const PTO2TaskSlotState *s) const { const PTO2TaskPayload &p = *s->payload; const PTO2SharedMemoryRingHeader &ring = *ring_sched_state.ring; for (int32_t i = 0; i < p.fanin_count; i++) { - if (!ring.is_completion_flag_set(p.fanin_local_ids[i])) return i; + if (!ring.is_completion_flag_set(thread_idx, p.fanin_local_ids[i])) return i; } return -1; } @@ -522,7 +572,7 @@ struct PTO2SchedulerState { // completed (head == SENTINEL), re-classify against ALL fanins: route to // ready only when every fanin is met, else re-target the next unmet producer // and retry. Monotonic completion_flags guarantee termination. - void register_wake(PTO2TaskSlotState *producer, PTO2TaskSlotState *consumer) { + void register_wake(int32_t thread_idx, PTO2TaskSlotState *producer, PTO2TaskSlotState *consumer) { PTO2SharedMemoryRingHeader &ring = *ring_sched_state.ring; while (true) { PTO2TaskSlotState *expected = producer->wake_list_head.load(std::memory_order_relaxed); @@ -534,7 +584,7 @@ struct PTO2SchedulerState { return; } } - int32_t state = classify_fanin_state(consumer); + int32_t state = classify_fanin_state(thread_idx, consumer); if (state < 0) { push_ready_routed(consumer); return; @@ -543,19 +593,13 @@ struct PTO2SchedulerState { } } - // Producer completion under polling: publish the host-visible task_state - // mirror + the device-visible completion_flags byte, drain the wake list - // (route/re-register each waiter), then CAS-advance the monotonic - // completed_watermark (load-bearing: the host wait_for_consumers gates on - // 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 task_id = static_cast(slot_state.task->task_id.local()); + // Drains slot_state's wake list: each waiter is routed to ready (single + // fanin, or all fanins now met) or re-registered on its next unmet + // producer. Shared by on_mixed_task_complete and + // retry_set_completion_flags, both of which drain a producer's wake list + // immediately after that producer's completion_flags entry becomes set. + void drain_wake_list(int32_t thread_idx, PTO2TaskSlotState &slot_state) { PTO2SharedMemoryRingHeader &ring = *ring_sched_state.ring; - - slot_state.mark_completed(); // host-visible mirror (task_state = COMPLETED) - 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) { PTO2TaskSlotState *next = waiter->next_in_wake_list; @@ -564,28 +608,80 @@ struct PTO2SchedulerState { waiter = next; continue; } - int state = classify_fanin_state(waiter); + int state = classify_fanin_state(thread_idx, waiter); if (state < 0) { push_ready_routed(waiter); } else { - register_wake(&ring.get_slot_state_by_task_id(waiter->payload->fanin_local_ids[state]), waiter); + register_wake( + thread_idx, &ring.get_slot_state_by_task_id(waiter->payload->fanin_local_ids[state]), waiter + ); } waiter = next; } + } + + // Producer completion under polling: publish the host-visible task_state + // mirror + the device-visible completion_flags entry, drain the wake list + // (route/re-register each waiter), then CAS-advance the monotonic + // completed_watermark (load-bearing: the host wait_for_consumers gates on + // 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(int32_t thread_idx, PTO2TaskSlotState &slot_state) { + 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) + if (unlikely(not ring.try_set_completion_flag(thread_idx, task_id))) { + failed_heap_of_set_completion_flag[thread_idx].push(task_id); + return; + } + + drain_wake_list(thread_idx, slot_state); + + // completed_watermark = lowest id not yet guaranteed complete: every task + // in [0, watermark) has its completion_flags entry 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 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. + // + // task_id gets exactly one update_completed_watermark call, made only + // after task_id's own completion_flags entry is actually visible: the + // early return above guarantees try_set_completion_flag succeeded, so + // the flag is already set by the time this line runs. When + // try_set_completion_flag fails instead, task_id's one chance to call + // update_completed_watermark moves to retry_set_completion_flags, + // made right after that retry succeeds. + ring.update_completed_watermark(thread_idx, task_id); + } + + // min_task's update_completed_watermark call was skipped in on_mixed_task_complete + // (its try_set_completion_flag failed there), so this is that task's one and only + // remaining chance to make it -- must fire here, immediately after the retried + // try_set_completion_flag succeeds, or the watermark could stall on min_task forever + // if some other thread's walk reaches it first (see docs/RUNTIME_LOGIC.md §8.4). + void retry_set_completion_flags(int32_t thread_idx) { + PTO2SharedMemoryRingHeader &ring = *ring_sched_state.ring; + FailedCompletionFlagHeap &heap = failed_heap_of_set_completion_flag[thread_idx]; + + while (unlikely(not heap.empty())) { + int32_t min_task = heap.entries[0]; + if (not ring.try_set_completion_flag(thread_idx, min_task)) { + break; + } + heap.pop_min(); + + drain_wake_list(thread_idx, ring.get_slot_state_by_task_id(min_task)); - // 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 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. - ring.update_completed_watermark(); + ring.update_completed_watermark(thread_idx, min_task); + } } // Polling: there is no ready-claim CAS (a producer routes each waiter exactly // once via the wake-list drain) and no per-producer consumer/scope refcount. - // Consumer retirement is observed by the host through completed_watermark >= + // Consumer retirement is observed by the host through completed_watermark > // producer.last_consumer_local_id, not by bumping a producer refcount. // Early-dispatch release. If the now-ready task was pre-staged @@ -818,7 +914,7 @@ struct PTO2SchedulerState { // Polling: scope-end takes no per-producer action. Under the wiring model // this bumped each task's scope refcount (PTO2_FANOUT_SCOPE_BIT); reclaim now - // gates on completed_watermark >= last_consumer_local_id, which needs no + // gates on completed_watermark > last_consumer_local_id, which needs no // scope reference. Kept as a no-op so the orchestrator call site is unchanged. void on_scope_end(PTO2TaskSlotState ** /*task_slot_states*/, int32_t /*count*/) {} @@ -852,20 +948,13 @@ struct PTO2SchedulerState { #else uint32_t #endif - on_task_complete( - PTO2TaskSlotState &slot_state -#if SIMPLER_SCHED_PROFILING - , - int thread_idx -#endif - ) { + on_task_complete(PTO2TaskSlotState &slot_state, int32_t thread_idx) { // Polling completion: publish the host-visible task_state mirror + the - // device-visible completion_flags byte, drain the wake list (route or + // device-visible completion_flags entry, drain the wake list (route or // re-register each waiter), and advance the watermark. Replaces the // fanout-list walk + fanin_refcount decrements of the wiring model. - on_mixed_task_complete(slot_state); + on_mixed_task_complete(thread_idx, slot_state); #if SIMPLER_SCHED_PROFILING - (void)thread_idx; // Resolved-successor accounting is not tracked on the polling path (the // producer no longer enumerates its consumers); report 0 for the DFX bar. return CompletionStats{0, 0, 0, true}; @@ -922,31 +1011,21 @@ inline bool AsyncWaitList::try_inline_complete_locked(AsyncWaitList::DrainCompletionSink &sink, PTO2TaskSlotState &slot_state) { // Return value (CompletionStats / consumer-walk count) discarded: // async-wait drain path has no Resolve swimlane bar attached. -#if SIMPLER_SCHED_PROFILING (void)sink.sched->on_task_complete(slot_state, sink.thread_idx); -#else - (void)sink.sched->on_task_complete(slot_state); -#endif sink.inline_completed++; return true; } template inline AsyncPollResult AsyncWaitList::poll_and_complete( - AICoreCompletionMailbox *aicore_mailbox, PTO2SchedulerState *sched -#if SIMPLER_SCHED_PROFILING - , - int thread_idx -#endif + AICoreCompletionMailbox *aicore_mailbox, PTO2SchedulerState *sched, int32_t thread_idx ) { AsyncPollResult result; if (!try_lock()) return result; AsyncWaitList::DrainCompletionSink sink{}; sink.sched = sched; -#if SIMPLER_SCHED_PROFILING sink.thread_idx = thread_idx; -#endif int32_t drain_err = PTO2_ERROR_NONE; drain_aicore_completion_mailbox_locked(aicore_mailbox, sink, drain_err); @@ -987,11 +1066,7 @@ inline AsyncPollResult AsyncWaitList::poll_and_complete( if (entry.normal_done && entry.waiting_completion_count <= 0) { // Return value (CompletionStats / consumer-walk count) discarded: // deferred-completion drain has no Resolve swimlane bar attached. -#if SIMPLER_SCHED_PROFILING (void)sched->on_task_complete(*entry.slot_state, thread_idx); -#else - (void)sched->on_task_complete(*entry.slot_state); -#endif // Polling: completion is fully published inline; no deferred release. result.completed++; 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 a28b26141c..e462828e73 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,7 +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.is_completion_flag_set(pid, std::memory_order_relaxed)) rc++; + if (ring.is_completion_flag_set(thread_idx, pid, std::memory_order_relaxed)) rc++; } } int32_t kid_aic = slot_state.task->kernel_id[0]; @@ -699,21 +699,7 @@ void SchedulerContext::handshake_partition(Runtime *runtime, int32_t tidx, int32 bool SchedulerContext::assign_cores_to_threads() { // Cluster-aligned round-robin assignment: cluster ci -> sched thread ci % active_sched_threads_. // Each cluster = 1 AIC + 2 adjacent AIV; the triple is always kept together. - // - // 3S+1P: the last AICPU thread is the core-less resolution thread (P); cores - // partition across the remaining (aicpu_thread_num_ - 1) scheduler threads - // only, so P never owns a cluster and never polls a COND register. P is - // mandatory — like tmr's scheduler + orchestrator split, host_build_graph - // needs at least two AICPU threads (one S + one P); one thread cannot own - // cores and resolve on a dedicated thread at once. - if (aicpu_thread_num_ < 2) { - LOG_ERROR( - "host_build_graph requires aicpu_thread_num >= 2 (1 scheduler + 1 resolution); got %d", aicpu_thread_num_ - ); - return false; - } - p_thread_idx_ = aicpu_thread_num_ - 1; - active_sched_threads_ = aicpu_thread_num_ - 1; + active_sched_threads_ = aicpu_thread_num_; int32_t cluster_count = aic_count_; // Max clusters any single sched thread can hold: ceil(cluster_count / active_sched_threads_). @@ -1055,29 +1041,6 @@ void SchedulerContext::on_orchestration_done( total_tasks_ = total_tasks; - // Allocate the per-S CompletedTaskQueues here on the boot leader, before it - // releases runtime_init_ready_ — no scheduler thread can push until then. - // Completed-but-unresolved tasks in flight are bounded by BOTH the total task - // count and the ring's task window (a task must occupy a ring slot to run and - // complete), so size to the tighter of the two, rounded up to a power of two - // and floored at 256. The window already caps this, so there is no artificial - // ceiling and a producer never has to spin on a full queue. - uint64_t sp_bound = static_cast(total_tasks); - if (sched_->ring_sched_state.ring != nullptr) { - uint64_t window = static_cast(sched_->ring_sched_state.ring->task_window_mask) + 1; - if (window < sp_bound) { - sp_bound = window; - } - } - uint64_t sp_cap = 256; - while (sp_cap < sp_bound) { - sp_cap <<= 1; - } - for (int32_t t = 0; t < active_sched_threads_; t++) { - sp_queues_[t].destroy(); // free a prior run's buffer before re-alloc - sp_queues_[t].init(sp_cap); - } - // Fold tasks completed inline during orchestration int32_t inline_completed = static_cast(rt->orchestrator.inline_completed_tasks); if (inline_completed > 0) { @@ -1144,16 +1107,16 @@ void SchedulerContext::classify_partition(int32_t thread_idx, int32_t nthreads) const int32_t lo = static_cast((static_cast(submitted) * thread_idx) / nthreads); const int32_t hi = static_cast((static_cast(submitted) * (thread_idx + 1)) / nthreads); for (int32_t id = lo; id < hi; id++) { - if (ring.is_completion_flag_set(id)) { + if (ring.is_completion_flag_set(thread_idx, id)) { continue; // completed on the host (hidden alloc); nothing to dispatch } PTO2TaskSlotState &s = ring.get_slot_state_by_task_id(id); - int32_t state = sched_->classify_fanin_state(&s); + int32_t state = sched_->classify_fanin_state(thread_idx, &s); if (state < 0) { sched_->push_ready_routed(&s); } else { int32_t prod_local = s.payload->fanin_local_ids[state]; - sched_->register_wake(&ring.get_slot_state_by_task_id(prod_local), &s); + sched_->register_wake(thread_idx, &ring.get_slot_state_by_task_id(prod_local), &s); } } } diff --git a/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_completion.cpp b/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_completion.cpp index bb46493ad3..8cfe85d515 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_completion.cpp +++ b/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_completion.cpp @@ -82,7 +82,7 @@ SlotTransition SchedulerContext::decide_slot_transition( // Complete one slot's task: subtask counting, mixed completion, deferred release, profiling. void SchedulerContext::complete_slot_task( PTO2TaskSlotState &slot_state, int32_t expected_reg_task_id, [[maybe_unused]] PTO2SubtaskSlot subslot, - int32_t thread_idx, int32_t core_id, Handshake *hank, [[maybe_unused]] int32_t &completed_this_turn + int32_t thread_idx, int32_t core_id, Handshake *hank, int32_t &completed_this_turn #if SIMPLER_DFX , uint64_t dispatch_ts, uint64_t finish_ts @@ -183,15 +183,46 @@ void SchedulerContext::complete_slot_task( ); } #endif - // 3S+1P: hand the finished task to the dedicated resolution (P) thread. - // P publishes completion_flags, drains the wake list, and advances the - // watermark — and owns completed_tasks_, so this scheduler thread neither - // resolves nor bumps completed_this_turn. (The Resolve swimlane bar is - // emitted by P, not here.) - sp_queues_[thread_idx].push(&slot_state); #if SIMPLER_DFX + // Time Resolve (walk the consumer list, decrement each consumer's + // fanin, push the newly-ready ones, ring doorbells for early-dispatch + // hits) so it renders as a child bar nested inside this iteration's + // Complete bar. The 1 µs floor below filters out the ~88% of tasks + // with 1-2 consumers (~500 ns Resolve) so only the long broadcast / + // reduction walks stand out on the lane. + uint64_t resolve_t0 = (l2_swimlane_level_ >= L2SwimlaneLevel::SCHED_PHASES) ? get_sys_cnt_aicpu() : 0; +#endif + // [[maybe_unused]] silences -Werror=unused-but-set-variable on the + // profiling-flags-smoke build path where SIMPLER_DFX is OFF and + // the Resolve emit below is excluded. + [[maybe_unused]] uint32_t consumers_resolved = 0; +#if SIMPLER_SCHED_PROFILING + // SCHED_PROFILING variant returns CompletionStats whose `fanout_edges` + // is the consumer-walk count. + consumers_resolved = sched_->on_task_complete(slot_state, thread_idx).fanout_edges; +#else + consumers_resolved = sched_->on_task_complete(slot_state, thread_idx); +#endif +#if SIMPLER_DFX + if (resolve_t0 != 0) { + uint64_t resolve_t1 = get_sys_cnt_aicpu(); + // Filter: drop Resolve bars under 1 µs so the lane shows only + // resolves that did meaningful work (high consumer counts or + // doorbells). 50 cycles @ 50 MHz = 1 µs (PLATFORM_PROF_SYS_CNT_FREQ + // is the device sys-cnt frequency). + constexpr uint64_t RESOLVE_EMIT_MIN_CYCLES = PLATFORM_PROF_SYS_CNT_FREQ / 1'000'000; // 1 µs + if (resolve_t1 - resolve_t0 >= RESOLVE_EMIT_MIN_CYCLES) { + l2_swimlane_aicpu_record_sched_phase( + thread_idx, L2SwimlaneSchedPhaseKind::Resolve, resolve_t0, resolve_t1, l2_swimlane.sched_loop_count, + consumers_resolved + ); + } + } l2_swimlane.phase_complete_count++; #endif + // Polling: on_task_complete published completion + drained the wake list + // inline; no deferred producer-release step. + completed_this_turn++; } #if SIMPLER_DFX diff --git a/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_context.h b/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_context.h index 680b3046d7..f92f0a6e26 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_context.h +++ b/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_context.h @@ -37,57 +37,6 @@ class Runtime; struct Handshake; struct PTO2Runtime; -// SPSC ring carrying completed-but-unresolved task slots from one scheduler (S) -// thread to the dedicated resolution (P) thread. Whole-graph-resident hbg never -// reclaims a slot, so the queued PTO2TaskSlotState* stays valid until P reads it. -// Single producer (the owning S thread) / single consumer (P): plain -// acquire/release on head/tail, no CAS. Capacity is a power of two sized to the -// in-flight bound (min of total tasks and the ring window), so a producer does -// not spin on a full queue in practice; the spin is a correctness backstop only -// (P drains independently, so it always makes progress). The std::atomic cursors -// make this type non-copyable / non-movable, so `buf` cannot be double-freed -// through an accidental copy. -struct CompletedTaskQueue { - PTO2TaskSlotState **buf{nullptr}; - uint64_t cap{0}; - uint64_t mask{0}; - alignas(64) std::atomic head{0}; // consumer (P) cursor - alignas(64) std::atomic tail{0}; // producer (S) cursor - - void init(uint64_t capacity_pow2) { - debug_assert(capacity_pow2 != 0 && (capacity_pow2 & (capacity_pow2 - 1)) == 0); - cap = capacity_pow2; - mask = capacity_pow2 - 1; - buf = new PTO2TaskSlotState *[capacity_pow2]; - head.store(0, std::memory_order_relaxed); - tail.store(0, std::memory_order_relaxed); - } - void destroy() { - delete[] buf; - buf = nullptr; - } - // Producer (S). Spins only if the ring is full — sized so this never fires. - void push(PTO2TaskSlotState *s) { - uint64_t t = tail.load(std::memory_order_relaxed); - while (t - head.load(std::memory_order_acquire) >= cap) { - SPIN_WAIT_HINT(); - } - buf[t & mask] = s; - tail.store(t + 1, std::memory_order_release); - } - // Consumer (P). Returns nullptr when empty. - PTO2TaskSlotState *pop() { - uint64_t h = head.load(std::memory_order_relaxed); - if (h == tail.load(std::memory_order_acquire)) { - return nullptr; - } - PTO2TaskSlotState *s = buf[h & mask]; - head.store(h + 1, std::memory_order_release); - return s; - } - uint64_t size() const { return tail.load(std::memory_order_acquire) - head.load(std::memory_order_acquire); } -}; - /** * SchedulerContext: owns all scheduler-side state and methods. * @@ -136,16 +85,6 @@ class SchedulerContext { // Main scheduler thread entry: poll completion + dispatch ready tasks. int32_t resolve_and_dispatch(Runtime *runtime, int32_t thread_idx); - // Dedicated resolution (P) thread entry (3S+1P). Owns no cores: drains the - // per-S CompletedTaskQueues and runs on_task_complete for each finished task - // (completion_flags publish + wake-list drain + watermark advance), making P - // the sole producer of the ready queues. Owns completed_tasks_ / termination. - int32_t run_resolution_thread(Runtime *runtime, int32_t thread_idx); - - // Index of the dedicated resolution (P) thread — the last AICPU thread. - // host_build_graph always reserves it (aicpu_thread_num >= 2 is required). - int32_t p_thread_idx() const { return p_thread_idx_; } - // Shutdown AICore registers for this thread's assigned cores. // Also runs PMU finalize (SIMPLER_DFX) before deinit when enabled. // Orchestrator threads (core_trackers_[thread_idx].core_num() == 0) are a no-op. @@ -226,15 +165,6 @@ class SchedulerContext { int32_t aicpu_thread_num_{0}; int32_t cores_total_num_{0}; - // --- 3S+1P dedicated resolution thread --- - // The AICPU threads split into (aicpu_thread_num_ - 1) core-owning schedulers - // (S) plus one core-less resolution thread (P) at index p_thread_idx_ = - // aicpu_thread_num_ - 1. host_build_graph requires aicpu_thread_num_ >= 2 (one - // S + one P). Each S thread hands its finished tasks to P through - // sp_queues_[its_thread_idx]. - int32_t p_thread_idx_{-1}; - CompletedTaskQueue sp_queues_[MAX_AICPU_THREADS]; - // Cluster-ordered worker_id lists, populated by post_handshake_init(). int32_t aic_worker_ids_[RUNTIME_MAX_WORKER]{}; int32_t aiv_worker_ids_[RUNTIME_MAX_WORKER]{}; diff --git a/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_dispatch.cpp b/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_dispatch.cpp index 4a8c467fab..1aeb8c2ca3 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_dispatch.cpp +++ b/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_dispatch.cpp @@ -871,183 +871,6 @@ int32_t SchedulerContext::try_early_dispatch( return total_staged; } -// ============================================================================= -// Dedicated resolution (P) thread — 3S+1P -// ============================================================================= - -// P owns no AICore cores. It drains the per-S CompletedTaskQueues and runs -// on_task_complete for every finished task: publish completion_flags, drain the -// wake list (route/re-register waiters into the ready queues), advance the -// watermark. As the sole producer of the ready queues its enqueues never -// contend. P owns completed_tasks_ and the terminal completed_ flip, so the S -// threads keep dispatching until P has resolved the whole graph (watermark fully -// advanced) — the host's wait_for_consumers never observes a stranded prefix. -int32_t SchedulerContext::run_resolution_thread(Runtime *runtime, int32_t thread_idx) { - always_assert(sched_ != nullptr); - PTO2SharedMemoryHeader *header = sched_->sm_header; - if (!header) { - LOG_ERROR("PTO2 resolution: header is null"); - return -1; - } - LOG_INFO("Thread %d: resolution (P) thread starting, serving %d schedulers", thread_idx, active_sched_threads_); - -#if SIMPLER_DFX - auto &l2_swimlane = sched_l2_swimlane_[thread_idx]; - l2_swimlane.reset(); - l2_swimlane.l2_swimlane_enabled = (l2_swimlane_level_ != L2SwimlaneLevel::DISABLED); -#endif - - uint64_t last_progress_ts = get_sys_cnt_aicpu(); - uint64_t scheduler_timeout_cycles = SCHEDULER_TIMEOUT_CYCLES; - const int32_t scheduler_timeout_ms_override = get_scheduler_timeout_ms(); - if (scheduler_timeout_ms_override > 0) { - scheduler_timeout_cycles = - static_cast(scheduler_timeout_ms_override) * PLATFORM_PROF_SYS_CNT_FREQ / 1000; - } - - while (true) { - if (completed_.load(std::memory_order_acquire)) break; - - // Propagate a fatal error latched by the orchestrator (host) or a - // scheduler thread; mirror resolve_and_dispatch's exit behavior. - if (header->orch_error_code.load(std::memory_order_acquire) != PTO2_ERROR_NONE || - header->sched_error_code.load(std::memory_order_acquire) != PTO2_ERROR_NONE) { - if (!completed_.exchange(true, std::memory_order_acq_rel)) { - emergency_shutdown(runtime); - } - break; - } - - int32_t resolved_this_pass = 0; - for (int32_t s = 0; s < active_sched_threads_; s++) { - PTO2TaskSlotState *slot; - while ((slot = sp_queues_[s].pop()) != nullptr) { -#if SIMPLER_SCHED_PROFILING - (void)sched_->on_task_complete(*slot, thread_idx); -#else - (void)sched_->on_task_complete(*slot); -#endif - resolved_this_pass++; - } - } - - // Async deferred completions, moved off the scheduler threads. Every - // condition that fires resolves via on_task_complete inside - // poll_and_complete, so async ready tasks also enter the ready queues - // through P alone. - if (rt_ != nullptr && rt_->aicore_mailbox != nullptr && - (sched_->async_wait_list.count > 0 || rt_->aicore_mailbox->has_pending())) { - AsyncPollResult poll_result = sched_->async_wait_list.poll_and_complete( - rt_->aicore_mailbox, sched_ -#if SIMPLER_SCHED_PROFILING - , - thread_idx -#endif - ); - if (poll_result.error_code != PTO2_ERROR_NONE) { - int32_t expected = PTO2_ERROR_NONE; - header->sched_error_code.compare_exchange_strong( - expected, poll_result.error_code, std::memory_order_acq_rel, std::memory_order_acquire - ); - if (!completed_.exchange(true, std::memory_order_acq_rel)) { - emergency_shutdown(runtime); - } - break; - } - resolved_this_pass += poll_result.completed; - } - - // Dependency-only tasks (empty active_mask, or a predicate that failed) - // route to dummy_ready_queue during resolution; P produces and drains it, - // so the queue is single-threaded end to end. Loop until empty — a dummy's - // resolution can make further dummies ready in the same pass. - { - constexpr int DUMMY_DRAIN_BATCH = 8; - PTO2TaskSlotState *dummy_batch[DUMMY_DRAIN_BATCH]; - int dummy_got; - while ((dummy_got = sched_->dummy_ready_queue.pop_batch(dummy_batch, DUMMY_DRAIN_BATCH)) > 0) { - for (int di = 0; di < dummy_got; di++) { -#if SIMPLER_SCHED_PROFILING - (void)sched_->on_task_complete(*dummy_batch[di], thread_idx); -#else - (void)sched_->on_task_complete(*dummy_batch[di]); -#endif - resolved_this_pass++; - } - } - } - - if (resolved_this_pass > 0) { - int32_t new_total = - completed_tasks_.fetch_add(resolved_this_pass, std::memory_order_relaxed) + resolved_this_pass; -#if SIMPLER_SCHED_PROFILING - // P owns the completion accounting, so it owns the profiling mirror too - // (the S threads' completed_this_turn no longer feeds it in P mode). - sched_->tasks_completed.fetch_add(resolved_this_pass, std::memory_order_relaxed); -#endif - last_progress_ts = get_sys_cnt_aicpu(); - if (total_tasks_ > 0 && new_total >= total_tasks_) { - completed_.store(true, std::memory_order_release); - LOG_INFO("Thread %d: P resolved all tasks %d/%d", thread_idx, new_total, total_tasks_); - break; - } - continue; // fast re-drain while work keeps arriving - } - - // Idle: nothing to resolve this pass. A task legitimately in flight — some - // thread still owns a RUNNING core — means P is merely waiting for that - // task to finish, not stalled: refresh the budget and keep spinning - // (mirrors resolve_and_dispatch's sibling-owns-running guard, so a task - // that runs longer than the timeout does not false-latch here). Only latch - // a hang when work is outstanding AND no thread anywhere owns a running - // task — a genuine forward-progress stall / pre-dispatch deadlock. - uint64_t now = get_sys_cnt_aicpu(); - if (now - last_progress_ts > scheduler_timeout_cycles) { - bool outstanding = total_tasks_ > 0 && completed_tasks_.load(std::memory_order_relaxed) < total_tasks_; - if (outstanding && no_thread_owns_running_task()) { - LOG_ERROR( - "Thread %d: P resolution stall (%d/%d resolved)", thread_idx, - completed_tasks_.load(std::memory_order_relaxed), total_tasks_ - ); - int32_t expected = PTO2_ERROR_NONE; - header->sched_error_code.compare_exchange_strong( - expected, PTO2_ERROR_SCHEDULER_TIMEOUT, std::memory_order_acq_rel, std::memory_order_acquire - ); - if (!completed_.exchange(true, std::memory_order_acq_rel)) { - emergency_shutdown(runtime); - } - break; - } - last_progress_ts = now; // a task is still running (or none outstanding): not a stall - } - SPIN_WAIT_HINT(); - } - -#if SIMPLER_DFX - // P owns no cores, so the AICore-keyed flushes below iterate an empty core - // list; the sched-phase-buffer flush is the one that matters — it drains any - // per-thread records P wrote (e.g. under SCHED_PROFILING) so they are not lost. - if (l2_swimlane.l2_swimlane_enabled) { - l2_swimlane_aicpu_flush( - thread_idx, core_trackers_[thread_idx].core_ids(), core_trackers_[thread_idx].core_num() - ); - if (l2_swimlane_level_ >= L2SwimlaneLevel::SCHED_PHASES) { - l2_swimlane_aicpu_flush_sched_phase_buffer(thread_idx); - } - } - if (is_dump_args_enabled()) { - dump_args_flush(thread_idx); - } - if (is_pmu_enabled()) { - pmu_aicpu_flush_buffers( - thread_idx, core_trackers_[thread_idx].core_ids(), core_trackers_[thread_idx].core_num() - ); - } -#endif - - return completed_tasks_.load(std::memory_order_relaxed); -} - // ============================================================================= // Main scheduler dispatch loop // ============================================================================= @@ -1170,6 +993,7 @@ int32_t SchedulerContext::resolve_and_dispatch(Runtime *runtime, int32_t thread_ } while (true) { if (completed_.load(std::memory_order_acquire)) { + sched_->retry_set_completion_flags(thread_idx); break; } bool made_progress = false; @@ -1191,7 +1015,10 @@ int32_t SchedulerContext::resolve_and_dispatch(Runtime *runtime, int32_t thread_ int32_t task_count = 0; if (!tracker.has_any_running_cores()) { LoopAction action = handle_orchestrator_exit(thread_idx, header, runtime, task_count); - if (action == LoopAction::BREAK_LOOP) break; + if (action == LoopAction::BREAK_LOOP) { + sched_->retry_set_completion_flags(thread_idx); + break; + } } #if SIMPLER_DFX @@ -1225,11 +1052,28 @@ int32_t SchedulerContext::resolve_and_dispatch(Runtime *runtime, int32_t thread_ } } - // Async deferred-completion polling and dependency-only (dummy / - // predicate-failed) retirement both run on P, which owns every - // completion→ready transition — the scheduler threads' loop stays purely - // core-local (poll own COND, dispatch own cores) and never touches the - // shared mailbox or dummy queue. + if (rt_ != nullptr && rt_->aicore_mailbox != nullptr && + (sched_->async_wait_list.count > 0 || rt_->aicore_mailbox->has_pending())) { + AsyncPollResult poll_result = + sched_->async_wait_list.poll_and_complete(rt_->aicore_mailbox, sched_, thread_idx); + if (poll_result.error_code != PTO2_ERROR_NONE) { + int32_t expected = PTO2_ERROR_NONE; + header->sched_error_code.compare_exchange_strong( + expected, poll_result.error_code, std::memory_order_acq_rel, std::memory_order_acquire + ); + completed_.store(true, std::memory_order_release); + break; + } + if (poll_result.completed > 0) { +#if SIMPLER_SCHED_PROFILING + sched_->tasks_completed.fetch_add(poll_result.completed, std::memory_order_relaxed); +#endif + int32_t prev = completed_tasks_.fetch_add(poll_result.completed, std::memory_order_relaxed); + int32_t new_total = prev + poll_result.completed; + last_progress_count = new_total; + made_progress = true; + } + } #if SIMPLER_DFX if (!try_completed) { @@ -1260,6 +1104,8 @@ int32_t SchedulerContext::resolve_and_dispatch(Runtime *runtime, int32_t thread_ } #endif + sched_->retry_set_completion_flags(thread_idx); + bool try_pushed = false; // Phase 2 drain check @@ -1288,9 +1134,97 @@ int32_t SchedulerContext::resolve_and_dispatch(Runtime *runtime, int32_t thread_ continue; } - // Phase 3 (dependency-only dummy / predicate-failed retirement) runs on - // the resolution thread P, not here — see run_resolution_thread. The - // scheduler loop goes straight from completion detection to dispatch. + // Phase 3: Drain dummy ready queue (every scheduler thread). + // + // Dependency-only tasks bypass AICore dispatch: they go through the + // scheduler so fanin/fanout edges stay consistent, but completion is + // signalled inline here. The ready queue is MPMC, and the fanout path + // uses per-slot locks/atomics, so multiple scheduler threads can share + // the dependency-only resolve work. + if (thread_idx < aicpu_thread_num_) { + constexpr int DUMMY_DRAIN_BATCH = 8; + PTO2TaskSlotState *dummy_batch[DUMMY_DRAIN_BATCH]; + int dummy_got = sched_->dummy_ready_queue.pop_batch(dummy_batch, DUMMY_DRAIN_BATCH); +#if SIMPLER_DFX + // Dummy outer phase: covers handling of all dummies popped this + // iter. Per-dummy DummyTask markers are emitted to a SEPARATE lane + // (Worker View AICPU_N) by the converter, so they do not nest + // under this bar. Resolve emits below DO land on the sched lane + // and nest under this Dummy outer by time containment. + uint64_t dummy_outer_t0 = + (dummy_got > 0 && l2_swimlane_level_ >= L2SwimlaneLevel::SCHED_PHASES) ? get_sys_cnt_aicpu() : 0; +#endif + for (int di = 0; di < dummy_got; di++) { + PTO2TaskSlotState &dummy_slot = *dummy_batch[di]; + + // ----- Resolve work: walk this dummy's consumer list. ------ + // Same 1 µs filter as the main-path Resolve emit suppresses + // dummies whose consumer release runs sub-microsecond. +#if SIMPLER_DFX + uint64_t dummy_resolve_t0 = + (l2_swimlane_level_ >= L2SwimlaneLevel::SCHED_PHASES) ? get_sys_cnt_aicpu() : 0; +#endif + // [[maybe_unused]] silences -Werror=unused-but-set-variable on + // the profiling-flags-smoke build path where SIMPLER_DFX is + // OFF and the Resolve emit below is excluded. + [[maybe_unused]] uint32_t dummy_consumers = 0; +#if SIMPLER_SCHED_PROFILING + dummy_consumers = sched_->on_task_complete(dummy_slot, thread_idx).fanout_edges; +#else + dummy_consumers = sched_->on_task_complete(dummy_slot, thread_idx); +#endif +#if SIMPLER_DFX + if (dummy_resolve_t0 != 0) { + uint64_t dummy_resolve_t1 = get_sys_cnt_aicpu(); + constexpr uint64_t RESOLVE_EMIT_MIN_CYCLES = PLATFORM_PROF_SYS_CNT_FREQ / 1'000'000; // 1 µs + if (dummy_resolve_t1 - dummy_resolve_t0 >= RESOLVE_EMIT_MIN_CYCLES) { + l2_swimlane_aicpu_record_sched_phase( + thread_idx, L2SwimlaneSchedPhaseKind::Resolve, dummy_resolve_t0, dummy_resolve_t1, + sched_l2_swimlane_[thread_idx].sched_loop_count, dummy_consumers + ); + } + l2_swimlane_aicpu_record_dummy_task( + thread_idx, dummy_resolve_t0, sched_l2_swimlane_[thread_idx].sched_loop_count, + dummy_slot.task->task_id.raw + ); + } +#endif + // Polling: on_task_complete already published this slot's + // completion + drained its wake list inline. There is no deferred + // producer-release phase — consumer retirement is observed via the + // per-ring completed_watermark, not by bumping producer refcounts. + int32_t prev = completed_tasks_.fetch_add(1, std::memory_order_relaxed); + last_progress_count = prev + 1; + cur_thread_completed++; + } + if (dummy_got > 0) { + made_progress = true; + } +#if SIMPLER_DFX + // Emit Dummy outer over the whole dummy_drain pass. Span starts at + // dummy_outer_t0 (captured after pop_batch) and ends at "now". + // tasks_processed = dummy_got. Advancing _t0_phase here makes the + // following Dispatch / EarlyDispatch / second-Complete bars start + // at this end. + if (dummy_outer_t0 != 0) { + uint64_t dummy_outer_t1 = get_sys_cnt_aicpu(); + int16_t phase_end_shared[L2SWIMLANE_NUM_QUEUE_SHAPES]; + capture_phase_end_fresh(phase_end_shared); + l2_swimlane_aicpu_record_sched_phase( + thread_idx, L2SwimlaneSchedPhaseKind::Dummy, dummy_outer_t0, dummy_outer_t1, + l2_swimlane.sched_loop_count, static_cast(dummy_got), /*pop_hit=*/0, + /*pop_miss=*/0, phase_start_shared, phase_end_shared + ); + for (int s = 0; s < L2SWIMLANE_NUM_QUEUE_SHAPES; s++) + phase_start_shared[s] = phase_end_shared[s]; + _t0_phase = dummy_outer_t1; + // We do NOT re-sync _t0/_t1 — the dummy span will be absorbed + // into the next CYCLE_COUNT_LAP accumulator. The phase-model + // anchor (_t0_phase) is the authoritative source for bar spans + // on the swimlane; the cycle accumulators are coarse aggregates. + } +#endif + } // Phase 4: MIX-strict-priority dispatch with phase-split and // cross-thread idle gating. See dispatch_ready_tasks for the policy. diff --git a/src/a2a3/runtime/host_build_graph/runtime/shared/pto_runtime2_init.cpp b/src/a2a3/runtime/host_build_graph/runtime/shared/pto_runtime2_init.cpp index 61e61e5152..65633bdaba 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/shared/pto_runtime2_init.cpp +++ b/src/a2a3/runtime/host_build_graph/runtime/shared/pto_runtime2_init.cpp @@ -193,6 +193,9 @@ void PTO2SchedulerState::wire_arena_pointers(const PTO2SchedulerLayout &layout, void PTO2SchedulerState::destroy() { PTO2SchedulerState *sched = this; sched->ring_sched_state.destroy(); + for (int i = 0; i < PLATFORM_MAX_AICPU_THREADS; i++) { + sched->failed_heap_of_set_completion_flag[i].destroy(); + } for (int i = 0; i < PTO2_NUM_RESOURCE_SHAPES; i++) { ready_queue_destroy(&sched->ready_queues[i]); } diff --git a/src/a2a3/runtime/host_build_graph/runtime/shared/pto_shared_memory.cpp b/src/a2a3/runtime/host_build_graph/runtime/shared/pto_shared_memory.cpp index fcff0a154c..0afc362937 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/shared/pto_shared_memory.cpp +++ b/src/a2a3/runtime/host_build_graph/runtime/shared/pto_shared_memory.cpp @@ -57,7 +57,7 @@ void PTO2SharedMemoryHandle::setup_pointers_per_ring(const uint64_t task_window_ ring.task_descriptors = (PTO2TaskDescriptor *)(base + off.descriptors); ring.task_payloads = (PTO2TaskPayload *)(base + off.payloads); ring.slot_states = (PTO2TaskSlotState *)(base + off.slot_states); - ring.completion_flags = (std::atomic *)(base + off.completion_flags); + ring.completion_flags = (std::atomic *)(base + off.completion_flags); } void PTO2SharedMemoryHandle::setup_pointers(uint64_t task_window_size) { @@ -91,8 +91,7 @@ bool PTO2SharedMemoryHandle::init_per_ring( sm_size = sm_size_arg; is_owner = false; setup_pointers_per_ring(task_window_sizes); - init_header_per_ring(task_window_sizes, heap_sizes); - return true; + return init_header_per_ring(task_window_sizes, heap_sizes); } bool PTO2SharedMemoryHandle::attach_populated( @@ -138,25 +137,43 @@ void PTO2SharedMemoryHandle::destroy() { // ============================================================================= // // no need init data in pool, init pool data when used -void PTO2SharedMemoryHandle::init_header(uint64_t task_window_size, uint64_t heap_size) { +bool PTO2SharedMemoryHandle::init_header(uint64_t task_window_size, uint64_t heap_size) { uint64_t task_window_sizes[PTO2_MAX_RING_DEPTH]; uint64_t heap_sizes[PTO2_MAX_RING_DEPTH]; for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { task_window_sizes[r] = task_window_size; heap_sizes[r] = heap_size; } - init_header_per_ring(task_window_sizes, heap_sizes); + return init_header_per_ring(task_window_sizes, heap_sizes); } -void PTO2SharedMemoryHandle::init_header_per_ring( +bool PTO2SharedMemoryHandle::init_header_per_ring( const uint64_t task_window_sizes[PTO2_MAX_RING_DEPTH], const uint64_t heap_sizes[PTO2_MAX_RING_DEPTH] ) { + // flag_index() shifts a value left by shuffle_higher_bits; task_window_size must + // contribute at least shuffle_lower_bits trailing zero bits or that shift amount + // goes negative, which is undefined behavior. Checked here, at the single + // point that assigns header->ring.task_window_size below, so every caller + // (init_per_ring, init_header) is covered without duplicating the guard. + for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { + if (__builtin_ctzll(task_window_sizes[r]) < PTO2SharedMemoryRingHeader::shuffle_lower_bits) return false; + } + // Flow control (starts at 0) header->ring.fc.init(); - // Polling completion: -1 = "no task completed yet"; the first task to - // complete (local_id 0) advances the watermark to 0. - header->ring.completed_watermark.store(-1, std::memory_order_relaxed); + // Polling completion: 0 = "no task completed yet"; the first task to + // complete (local_id 0) advances the watermark to 1. + header->ring.completed_watermark.store(0, std::memory_order_relaxed); + + // Per-thread cache of the last-observed completed_watermark (including the + // host's own slot at index PLATFORM_MAX_AICPU_THREADS). Shared memory is + // not guaranteed zero on device; a stale nonzero entry would make + // is_completion_flag_set's cache short-circuit report a not-yet-completed + // local_id as done. + for (auto &cached : header->ring.cached_completed_watermark) { + cached.val_ = 0; + } header->orchestrator_done.store(0, std::memory_order_relaxed); @@ -164,6 +181,7 @@ void PTO2SharedMemoryHandle::init_header_per_ring( uint64_t offset = PTO2_ALIGN_UP(sizeof(PTO2SharedMemoryHeader), PTO2_ALIGN_SIZE); header->ring.task_window_size = task_window_sizes[0]; header->ring.task_window_mask = static_cast(task_window_sizes[0] - 1); + header->ring.shuffle_higher_bits = __builtin_ctzll(task_window_sizes[0]) - header->ring.shuffle_lower_bits; header->ring.heap_size = heap_sizes[0]; header->ring.task_descriptors_offset = offset; offset += PTO2_ALIGN_UP(task_window_sizes[0] * sizeof(PTO2TaskDescriptor), PTO2_ALIGN_SIZE); @@ -190,10 +208,13 @@ void PTO2SharedMemoryHandle::init_header_per_ring( ring.slot_states[i].active_mask = ActiveMask{}; } - // Polling completion flags: 0 = pending. Shared memory is not guaranteed - // zero on device; stale non-zero bytes would make consumers observe a - // producer as already completed. Zero the whole per-ring array once. - __builtin_memset((void *)ring.completion_flags, 0, task_window_sizes[0] * sizeof(std::atomic)); + // Polling completion flags: -1 = pending (0 is a valid completion stamp — + // local_id 0 — so it cannot double as the pending sentinel). Shared memory + // is not guaranteed zero on device; stale bytes would make consumers + // observe a producer as already completed. 0xFF fills every int32 slot + // with -1 regardless of endianness. Set the whole per-ring array once. + __builtin_memset((void *)ring.completion_flags, 0xFF, task_window_sizes[0] * sizeof(std::atomic)); + return true; } // ============================================================================= diff --git a/tests/ut/cpp/CMakeLists.txt b/tests/ut/cpp/CMakeLists.txt index e4d0a123d5..e0fa789dbd 100644 --- a/tests/ut/cpp/CMakeLists.txt +++ b/tests/ut/cpp/CMakeLists.txt @@ -667,6 +667,34 @@ add_a2a3_test(test_task_timing_slots a2a3/test_task_timing_slots.cpp) add_a2a3_runtime_test(test_task_allocator a2a3/test_task_allocator.cpp) add_a2a3_runtime_test(test_scope_deadlock_detection common/test_scope_deadlock_detection.cpp) add_a2a3_hbg_runtime_test(test_hbg_task_allocator a2a3/test_task_allocator.cpp) + +# host_build_graph's pto_shared_memory.cpp is not part of add_a2a3_hbg_runtime_test's +# sources (only test_task_allocator needs that helper today, and pto_ring_buffer.h is +# header-only), so this test compiles it directly alongside the gtest source. +add_executable(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 + ${CMAKE_SOURCE_DIR}/stubs/test_stubs.cpp +) +target_include_directories(test_hbg_shared_memory PRIVATE + ${GTEST_INCLUDE_DIRS} + ${CMAKE_SOURCE_DIR}/../../../src/a2a3/runtime/host_build_graph/orchestration + ${CMAKE_SOURCE_DIR}/../../../src/a2a3/runtime/host_build_graph/runtime + ${CMAKE_SOURCE_DIR}/../../../src/a2a3/runtime/host_build_graph/common + ${CMAKE_SOURCE_DIR}/../../../src/a2a3/platform/include + ${CMAKE_SOURCE_DIR}/../../../src/common/platform/sim/aicpu + ${CMAKE_SOURCE_DIR}/../../../src/common/platform/include + ${CMAKE_SOURCE_DIR}/../../../src/common/task_interface + ${CMAKE_SOURCE_DIR}/../../../src/common/log/include + ${CMAKE_SOURCE_DIR}/../../../src/common +) +target_link_libraries(test_hbg_shared_memory PRIVATE + ${GTEST_MAIN_LIB} + ${GTEST_LIB} + pthread +) +add_test(NAME test_hbg_shared_memory COMMAND test_hbg_shared_memory) +set_tests_properties(test_hbg_shared_memory PROPERTIES LABELS "no_hardware") add_a2a3_hbg_runtime_test(test_hbg_tensormap a2a3/test_hbg_tensormap.cpp) # PTO2TensorMap's out-of-line members (reserve_layout / init / valid_count) live # in this .cpp; no other add_a2a3_hbg_runtime_test target needs them. diff --git a/tests/ut/cpp/a2a3/test_hbg_shared_memory.cpp b/tests/ut/cpp/a2a3/test_hbg_shared_memory.cpp new file mode 100644 index 0000000000..445e1d72a7 --- /dev/null +++ b/tests/ut/cpp/a2a3/test_hbg_shared_memory.cpp @@ -0,0 +1,53 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ +/** + * Unit tests for the shuffle_higher_bits invariant in + * host_build_graph's pto_shared_memory.cpp: flag_index() shifts by + * shuffle_higher_bits, so init must reject a task_window_size that would + * make that shift amount negative rather than let it become undefined + * behavior. + */ + +#include + +#include + +#include "pto_shared_memory.h" + +namespace { + +TEST(HbgSharedMemoryShuffleBits, CreateAndInitDefaultProducesNonNegativeShuffleHigherBits) { + DeviceArena arena; + PTO2SharedMemoryHandle *handle = PTO2SharedMemoryHandle::create_and_init_default(arena); + ASSERT_NE(handle, nullptr); + EXPECT_GE(handle->header->ring.shuffle_higher_bits, 0); +} + +TEST(HbgSharedMemoryShuffleBits, InitPerRingRejectsTaskWindowBelowShuffleFloor) { + // ctzll(8) == 3 < shuffle_lower_bits (4), so shuffle_higher_bits would be -1. + constexpr uint64_t kTaskWindowSize = 8; + static_assert(kTaskWindowSize < (uint64_t{1} << PTO2SharedMemoryRingHeader::shuffle_lower_bits), ""); + + uint64_t task_window_sizes[PTO2_MAX_RING_DEPTH]; + uint64_t heap_sizes[PTO2_MAX_RING_DEPTH]; + for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { + task_window_sizes[r] = kTaskWindowSize; + heap_sizes[r] = 4096; + } + + const uint64_t sm_size = PTO2SharedMemoryHandle::calculate_size_per_ring(task_window_sizes); + std::vector buf(sm_size); + + PTO2SharedMemoryHandle handle{}; + EXPECT_FALSE(handle.init_per_ring(buf.data(), sm_size, task_window_sizes, heap_sizes)); +} + +} // namespace