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 6f3d21a7f..7dcfb9f7b 100644 --- a/src/a2a3/runtime/host_build_graph/docs/RUNTIME_LOGIC.md +++ b/src/a2a3/runtime/host_build_graph/docs/RUNTIME_LOGIC.md @@ -30,7 +30,7 @@ four layers of execution: │ │ │ ┌─────────────────────────────────────────────────────────────────┐ │ │ │ Shared Memory (GM) │ │ -│ │ SharedMemoryHeader │ TaskDescriptors[] │ DepListPool │ │ +│ │ SharedMemoryHeader │ TaskDescriptors[] │ Payloads[] │ SlotStates[] │ │ │ GM Heap (output buffers) │ │ │ └─────────────────────────────────────────────────────────────────┘ │ │ │ @@ -46,20 +46,19 @@ four layers of execution: > > - **host_build_graph (this runtime):** the **host** dlopens the orchestration > `.so` and runs it to completion, populating shared memory + the prebuilt -> arena. It **wires the fanin/fanout adjacency inline during submit** — -> allocates `dep_pool` entries under each live producer's `fanout_lock`, links -> the consumer onto the producer's `fanout_head`, and seeds readiness directly -> into the ready queue (`push_ready_routed` for zero-fanin tasks, -> `route_ready_once` once a producer's fanin resolves). The host then -> relocates every cross-task pointer to its final device address -> (`relocate_host_orch_image`) **before** the H2D copy — pointers into the SM -> and pointers into the arena shift by independent deltas — and ships the -> image. The device boots **scheduler-only**: no on-device orchestrator thread -> and no on-device pointer fixup; it attaches the already-device-addressed -> SM/arena, whose ready queue is already seeded, and dispatches. +> arena. Dependencies are recorded as **position-independent integer fanin**: +> each task's payload stores its producers' local-ids (`fanin_local_ids`) and +> each producer records its highest consumer id (`last_consumer_local_id`). +> There is no fanout adjacency, dep-pool, or per-task lock. Because fanin is +> integer ids, the host→device relocation (`relocate_host_orch_image`) +> collapses to fixing up only the per-slot `task` / `payload` pointers before +> the H2D copy. The device boots **scheduler-only**: no on-device orchestrator +> thread and no fanout wiring; on boot it scans the submitted tasks and +> classifies each (fanin-free → ready queue; otherwise register on its first +> unmet producer's wake list — see §8), then dispatches. > - **tensormap_and_ringbuffer:** the orchestrator runs **on-device** on AICPU -> thread N-1, concurrently with the scheduler threads, wiring the same -> fanin/fanout adjacency inline during submit. +> thread N-1, concurrently with the scheduler threads, recording the same +> integer fanin during submit. > > Where a section below says "the orchestrator runs on AICPU Thread 3" or > "Thread 3 dlopens the SO", read that as the **tensormap (device-orch) @@ -93,7 +92,7 @@ The primary production runtime. Uses ring buffers for task slots and output memo - **Dependencies**: automatically derived from tensor read/write patterns via TensorMap - **Thread model**: 3 scheduler threads + 1 orchestrator thread on AICPU - **Single ring**: host_build_graph builds the whole graph on the host with no - execution-time reclaim, so HeapRing, TaskRing, and DepPool are single + execution-time reclaim, so HeapRing and TaskRing are single whole-graph-resident instances (`PTO2_MAX_RING_DEPTH == 1`); all scope depths map to ring 0. - **Use case**: production workloads; supports streaming, flow control, and large batch sizes @@ -137,7 +136,7 @@ Two platform implementations exist under `src/platform/`, sharing a common inter ## 3. Shared Memory Layout -The orchestrator and schedulers communicate through a contiguous shared memory region in Global Memory (GM). The single ring's TaskDescriptor and DepListPool sections are laid out by `pto2_sm_layout::ring_segment_offsets`. +The orchestrator and schedulers communicate through a contiguous shared memory region in Global Memory (GM). The single ring's TaskDescriptor, TaskPayload, and TaskSlotState sections (plus the per-slot `completion_flags`) are laid out by `pto2_sm_layout::ring_segment_offsets`. ```text ┌─────────────────────────────┐ offset 0 @@ -179,11 +178,10 @@ Alignment is 64 bytes (`PTO2_ALIGN_SIZE`). ## 4. Ring Buffer Mechanisms -> **Single ring**: TaskRing, HeapRing, and DepPool are single whole-graph -> instances (`PTO2_MAX_RING_DEPTH == 1`). The host builds the entire graph -> before the device boots and there is no execution-time reclaim, so a -> per-scope-depth ring split would buy nothing — every scope depth maps to -> ring 0. +> **Single ring**: TaskRing and HeapRing are single whole-graph instances +> (`PTO2_MAX_RING_DEPTH == 1`). The host builds the entire graph before the +> device boots and there is no execution-time reclaim, so a per-scope-depth +> ring split would buy nothing — every scope depth maps to ring 0. ### 4.1 Task Ring @@ -227,13 +225,21 @@ The heap ring manages output buffer allocation from a circular GM heap. **Reclamation**: When `last_task_alive` advances past a task, its `packed_buffer_end` is used to advance `heap_tail`, freeing the memory region. -### 4.3 Dependency List Pool +### 4.3 Dependency Representation (polling completion) -A simple bump allocator for `PTO2DepListEntry` nodes used in fanin/fanout linked lists. +There is no dependency-list pool or fanout adjacency. A task's dependencies are +stored inline on its payload as a flat array of position-independent producer +local-ids: -- **Entry 0**: NULL sentinel (`task_id=-1, next_offset=0`) -- **Allocation**: `pool->top++`, wraps around when full -- **Reclamation**: implicit — old entries become unreachable as `last_task_alive` advances +- `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). + +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` +pointers are relocated (see §7.3). ### 4.4 Flow Control and Back-Pressure @@ -371,13 +377,13 @@ When `PTO2OrchestratorState::submit_task` processes parameters: | `kernel_id[3]` | Per-slot kernel IDs: `[AIC, AIV0, AIV1]`; `INVALID_KERNEL_ID` = inactive | | `active_mask` | Bitmask of active subtask slots: `bit0=AIC`, `bit1=AIV0`, `bit2=AIV1` | | `completed_subtasks` | Atomic counter; each subtask increments on completion. Trigger condition: `completed_subtasks == total_required_subtasks` | -| `fanin_count` | Number of producer dependencies (set inline during submit wiring) | -| `fanout_lock` | Per-task spinlock for fanout modification (used by submit-time wiring + scheduler completion) | -| `fanout_head` | Head of fanout consumer list (pointer, protected by `fanout_lock`) | -| `fanout_count` | 1 (scope ref) + number of consumers | | `packed_buffer_base` | Start of packed buffer in GM Heap | | `packed_buffer_end` | End of packed buffer (for heap reclamation) | +Fanin/fanout are not on the descriptor: producer ids live on the payload +(`fanin_local_ids`, §6.1b) and consumers are reached at completion via the +per-slot wake list (§8.2), not a fanout adjacency list. + ### 6.1b PTO2TaskPayload (Cold Path) | Field | Description | @@ -386,25 +392,32 @@ When `PTO2OrchestratorState::submit_task` processes parameters: | `scalar_value[16]` | Scalar parameter values | | `is_tensor[16]` | Whether each parameter is tensor or scalar | | `param_count` | Number of valid parameters | -| `fanin_slot_states[]` | Producer slot state pointers (used by `on_task_release`) | -| `fanin_actual_count` | Actual fanin count | +| `fanin_local_ids[fanin_count]` | Producer local-ids (position-independent; readiness = all their `completion_flags` set) | +| `fanin_count` | Number of producer dependencies (`<= PTO2_MAX_FANIN`, no spill) | +| `predicate` | Dispatch predicate (`DispatchPredicate`, cache line 9 / byte 576); evaluated by the scheduler at the ready point (see §8.3) | + +`last_consumer_local_id` (the highest consumer id of this task) lives on the +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 +`task_state` in `wait_for_tensor_ready`, the allocator deadlock detector, and +the cold-path stall dump. + ```text - [0] PENDING ──worker(s) done──► [1] COMPLETED ──fanout done──► [2] CONSUMED - ▲ │ - │ ▼ - └──────────────────── slot recycled ◄───────────────────────────┘ + [0] PENDING ──worker(s) done, on_mixed_task_complete──► [1] COMPLETED ``` -In the scheduler's `task_state[]` array (`std::atomic`): +- **0 (PENDING)**: slot allocated; remains PENDING while waiting on producers, + queued, or dispatched. +- **1 (COMPLETED)**: all subtasks done. `on_mixed_task_complete` sets this + (host mirror) and, in the same step, sets `completion_flags[id]` (device + readiness) and advances `completed_watermark`. -- **0 (PENDING)**: slot is allocated and remains PENDING through "waiting on - producers", "queued in ready queue", and "dispatched to a worker"; ready vs - running is derived from `fanin_refcount` and per-core `running_slot_state` -- **1 (COMPLETED)**: hardware execution complete, output may still be in use -- **2 (CONSUMED)**: output fully consumed, buffers can be released +There is no runtime CONSUMED flip and no slot recycle: host_build_graph is +whole-graph-resident (§8.4). --- @@ -419,10 +432,10 @@ structure tensormap drives on AICPU thread N-1. Key members: -- `ring`: the single `PTO2RingSet` (HeapRing + TaskRing + FaninPool). +- `ring`: the single `PTO2RingSet` (HeapRing + TaskRing). - `tensor_map`, `tensor_pool`: dependency tracking - `scope_tasks[]`, `scope_begins[]`, `scope_stack_top`: scope nesting stack (flat buffer partitioned by level) -- `scheduler`: pointer to scheduler state (for inline fanout wiring and ready queue access) +- `scheduler`: pointer to scheduler state (for seeding zero-fanin tasks into the ready queue) - `gm_heap_base`, `gm_heap_size`: GM heap for output buffers ### 7.2 Task Submission Flow (`PTO2OrchestratorState::submit_task`) @@ -432,65 +445,42 @@ Key members: | 0 | `PTO2TensorMap::sync_tensormap` — prune stale TensorMap entries | | 1 | `PTO2TaskAllocator::alloc` — allocate task slot (may block on flow control) | | 2 | Initialize task descriptor + slot state, copy parameters | -| 3 | **Lookup**: for each INPUT/INOUT param, search TensorMap for producers; collect producer pointers in `PTO2FaninBuilder` | +| 3 | **Lookup**: for each INPUT/INOUT param, search TensorMap for producers; collect producer ids in `PTO2FaninBuilder` | | 4 | **Insert**: register OUTPUT/INOUT args in TensorMap | -| 5 | **Record fanin metadata**: store producer pointers in `payload->fanin_inline_slot_states[]` (+ spill pool if >64); increment each producer's `fanout_count` (no lock needed — single writer). This step runs **before** `payload.init()`. | -| 6 | **Wire fanout inline**: for a task with live producers, lock each producer, allocate `dep_pool` entries, prepend the consumer to each producer's `fanout_head`, then seed readiness — `push_ready_routed` for zero-fanin tasks, `route_ready_once` once the fanin refcount is already satisfied. | - -> **Note**: The orchestrator wires fanout **inline in submit** (Step 6) — -> there is no device-side wiring queue. The `dep_pool` is sized for the whole -> graph — there is no reclaim during host orchestration, exactly like the task -> window and GM heap — so an exhausted pool latches -> `PTO2_ERROR_DEP_POOL_OVERFLOW` and aborts the run. The `fanout_head` / -> dep-entry / ready-queue pointers this produces are host-DDR addresses; -> `relocate_host_orch_image` shifts them to device addresses before H2D (SM -> pointers and arena pointers by independent deltas). - -### 7.3 Fanout Wiring - -The orchestrator wires each task's fanin/fanout adjacency inline during submit -(Step 6). Three cases, by the state of the claimed producers: - -1. **Zero-fanin** (`fanin_builder.count == 0`): no producers to link. Sets - `fanin_count = 1`, releases the +1 self-reference, records the `dep_pool` - position (`orch_mark_dep_pool_position`), and seeds the task directly via - `push_ready_routed`. -2. **All claimed producers already completed**: no live fanout links are - needed. Sets `fanin_count = N + 1`, primes `dispatch_fanin` when every - producer is codegen-flagged, releases the refcount, and pushes via - `push_ready_routed`. -3. **At least one live producer** (`orch_wire_live_fanin_task` → - `orch_wire_fanin_task`): - - Reserves `dep_pool` space (`ensure_space`); an exhausted pool latches - `PTO2_ERROR_DEP_POOL_OVERFLOW`. - - Sets `fanin_count = N + 1` (+1 self-reference prevents premature - readiness). - - For each producer, under its `fanout_lock`: if `task_state >= COMPLETED`, - count it as early-finished; otherwise prepend the consumer to the - producer's `fanout_head` via `dep_pool.prepend`. - - Seeds `dispatch_fanin` by the early-finished count when every producer is - codegen-flagged. - - Releases the +1 self-reference plus the early-finished count via - `fanin_refcount.fetch_add`; if that already satisfies `fanin_count`, - publishes readiness via `route_ready_once`. - -`push_ready_routed` pushes a ready slot straight to its shape's queue; -`route_ready_once` claims the slot exactly once (CAS), takes the early-dispatch -doorbell path if the task was pre-staged, and otherwise routes it to the queue. - -The scheduler's completion handler (`on_task_complete`) mirrors the wiring: -acquire `fanout_lock`, mark the task COMPLETED, read `fanout_head`, release the -lock, then traverse the fanout list releasing each consumer's fanin -(`release_fanin_and_check_ready`) and pushing newly-ready consumers via -`route_ready_once`. This protocol guarantees every consumer is accounted for -exactly once. (host-orch does not flip tasks to CONSUMED — see §8.4.) +| 5 | **Record fanin**: `append_fanin_or_fail` dedups producers and writes their local-ids into `payload->fanin_local_ids[]`; each producer's `last_consumer_local_id` is bumped to this task's id. `payload.fanin_count` is set from the builder. | +| 6 | **Seed readiness**: a task with zero fanin (`fanin_count == 0`) is pushed straight to the ready queue via `push_ready_routed`. A task with fanin is left for the device boot classify (§8) — the host does not pre-register wake lists. | + +> **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. + +### 7.3 Dependency Recording and Relocation + +`append_fanin_or_fail` records, per consumer, the deduped list of producer +local-ids into `payload->fanin_local_ids[]` and bumps `fanin_count`; it also +raises each producer's `last_consumer_local_id` to the consumer's id. An +overflow past `PTO2_MAX_FANIN` is fatal (`PTO2_ERROR_...`), since there is no +spill. + +`relocate_host_orch_image` runs on the host before H2D. Because fanin is +position-independent integer ids, the only pointers needing fixup are the +per-slot `task` and `payload` pointers (SM-region delta); the fanout adjacency, +dep-pool, and ready-queue pointer relocation of the wiring model are gone. + +Readiness and completion are handled entirely device-side by the scheduler: +the boot classify seeds the ready queue and registers wake lists (§8.2), and +`on_mixed_task_complete` publishes each producer's `completion_flags` and drains +its wake list. See §8.2 for the completion protocol. ### 7.4 Scope Mechanism (`PTO2_SCOPE`) Scopes control the lifetime of intermediate buffers. Each scope: - Tracks tasks submitted within it via a flat `scope_tasks[]` buffer partitioned by `scope_begins[]` -- On `scope_end`: increments `fanout_refcount` for scope tasks; when it reaches `fanout_count`, the task's packed buffer can be reclaimed +- Scopes bound intermediate-buffer lifetime **structurally** (the orchestration function that built the graph). host_build_graph is whole-graph-resident, so `scope_end` performs no runtime buffer reclaim — there is no `fanout_refcount`. ```cpp PTO2_SCOPE(rt) { @@ -568,10 +558,31 @@ Core assignment: AICs and AIVs are divided equally among the 3 scheduler threads Each scheduler thread runs a tight loop with two main phases: -**Phase 1 — Completion Handling**: - -- Poll register `COND` on each managed core -- When `TASK_FIN_STATE` detected: record completion timestamps, call `on_subtask_complete(task_id, subslot)` to increment the completion counter; when `completed_subtasks == total_required_subtasks`, trigger `on_task_complete(task_id)` which marks `task_state[slot] = COMPLETED`, acquires fanout lock, traverses fanout list (incrementing consumers' `fanin_refcount`), marks `task_state[slot] = CONSUMED`, and advances `last_task_alive` watermark +**Phase 1 — Completion Handling (polling)**: + +- Poll register `COND` on each managed core. +- 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). + +**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. **Phase 2 — Dispatch**: @@ -591,20 +602,25 @@ Ready queues use a lock-free bounded MPMC (Vyukov) design: - Per-slot sequence counters prevent ABA problems - `enqueue_pos` and `dequeue_pos` are on separate cache lines to avoid false sharing -### 8.4 No Runtime Watermark Advancement (host-orch) - -host_build_graph is whole-graph-resident: the host builds the entire task -graph into a single ring and H2Ds it once, and the device runs it with **no -execution-time reclaim**. `last_task_alive` is initialized to 0 and is **not -advanced at runtime** — slots are never recycled within a run, so there is no -`advance_ring_pointers` / per-ring `advance_lock` step (both removed; they -existed only for the device-orch reclaim path in `tensormap_and_ringbuffer`). -Completion is tracked by `completed_tasks_`; consumer waits key on -`fanout_refcount` rather than on a watermark. - -`reset_for_reuse()` survives but runs **once at init** -(`pto_shared_memory.cpp`) to zero each slot's dynamic scheduling fields before -the host orchestrator populates them — it is not a runtime recycle hook. +### 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. + +It is **load-bearing**: the host `wait_for_tensor_ready(..., wait_for_consumers)` +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. + +Slot reclaim is inert: host_build_graph is whole-graph-resident, so +`last_task_alive` is never advanced at runtime and there is no +`advance_ring_pointers` step. `reset_for_reuse()` runs **once at init** +(`pto_shared_memory.cpp`) to zero each slot before the host orchestrator +populates it — it is not a runtime recycle hook. ### 8.5 SchedulerContext diff --git a/src/a2a3/runtime/host_build_graph/host/runtime_maker.cpp b/src/a2a3/runtime/host_build_graph/host/runtime_maker.cpp index 4d4a59dad..59f791686 100644 --- a/src/a2a3/runtime/host_build_graph/host/runtime_maker.cpp +++ b/src/a2a3/runtime/host_build_graph/host/runtime_maker.cpp @@ -201,39 +201,32 @@ static uint64_t read_ring_override(const uint64_t *base, int idx) { return value; } -// Each of ring_task_window / ring_heap / ring_dep_pool is a per-ring array of -// PTO2_MAX_RING_DEPTH entries (0 = unset). Precedence per ring: per-task entry > -// PTO2_RING_* env value > compile-time default. A "size all rings the same" -// request arrives already broadcast to every entry by the caller. +// Each of ring_task_window / ring_heap is a per-ring array of PTO2_MAX_RING_DEPTH +// entries (0 = unset). Precedence per ring: per-task entry > PTO2_RING_* env value +// > compile-time default. A "size all rings the same" request arrives already +// broadcast to every entry by the caller. (Polling has no dep_pool, so the former +// PTO2_RING_DEP_POOL knob is gone.) static bool resolve_ring_config( - const uint64_t *ring_task_window, const uint64_t *ring_heap, const uint64_t *ring_dep_pool, - uint64_t eff_task_window_sizes[PTO2_MAX_RING_DEPTH], uint64_t eff_heap_sizes[PTO2_MAX_RING_DEPTH], - int32_t eff_dep_pool_capacities[PTO2_MAX_RING_DEPTH] + const uint64_t *ring_task_window, const uint64_t *ring_heap, uint64_t eff_task_window_sizes[PTO2_MAX_RING_DEPTH], + uint64_t eff_heap_sizes[PTO2_MAX_RING_DEPTH] ) { - uint64_t dep_pool_values[PTO2_MAX_RING_DEPTH]; for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { eff_task_window_sizes[r] = PTO2_TASK_WINDOW_SIZE; eff_heap_sizes[r] = PTO2_HEAP_SIZE; - dep_pool_values[r] = PTO2_DEP_LIST_POOL_SIZE; } apply_env_ring_values("PTO2_RING_TASK_WINDOW", 4, static_cast(INT32_MAX), true, eff_task_window_sizes); apply_env_ring_values("PTO2_RING_HEAP", 1024, std::numeric_limits::max(), false, eff_heap_sizes); - apply_env_ring_values("PTO2_RING_DEP_POOL", 4, static_cast(INT32_MAX), false, dep_pool_values); for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { const uint64_t task_window_override = read_ring_override(ring_task_window, r); const uint64_t heap_override = read_ring_override(ring_heap, r); - const uint64_t dep_pool_override = read_ring_override(ring_dep_pool, r); if (task_window_override != 0) { eff_task_window_sizes[r] = task_window_override; } if (heap_override != 0) { eff_heap_sizes[r] = heap_override; } - if (dep_pool_override != 0) { - dep_pool_values[r] = dep_pool_override; - } if (eff_task_window_sizes[r] < 4 || eff_task_window_sizes[r] > static_cast(INT32_MAX) || !is_power_of_2_u64(eff_task_window_sizes[r])) { @@ -246,11 +239,6 @@ static bool resolve_ring_config( LOG_ERROR("ring_heap[%d]=%" PRIu64 " must be >= 1024", r, eff_heap_sizes[r]); return false; } - if (dep_pool_values[r] < 4 || dep_pool_values[r] > static_cast(INT32_MAX)) { - LOG_ERROR("ring_dep_pool[%d]=%" PRIu64 " must be in [4, INT32_MAX]", r, dep_pool_values[r]); - return false; - } - eff_dep_pool_capacities[r] = static_cast(dep_pool_values[r]); } return true; @@ -369,8 +357,8 @@ struct HostOrchEntryPoints { // (returns false) rather than shipping un-relocated host pointers to the device. // Returns false on any unrelocatable pointer so the caller can fail the prepare. static bool relocate_host_orch_image( - PTO2SharedMemoryHandle &host_sm_handle, PTO2Runtime *rt, uint64_t host_sm, uint64_t sm_size, int64_t sm_delta, - uint64_t host_arena, uint64_t arena_size, int64_t arena_delta + PTO2SharedMemoryHandle &host_sm_handle, [[maybe_unused]] PTO2Runtime *rt, uint64_t host_sm, uint64_t sm_size, + int64_t sm_delta, uint64_t host_arena, uint64_t arena_size, int64_t arena_delta ) { // host_build_graph is single-ring; the loops below iterate the lone ring and // index header->ring (singular). If the ring depth ever grows, those loops @@ -418,69 +406,16 @@ static bool relocate_host_orch_image( int32_t count = ring.fc.current_task_index.load(std::memory_order_acquire); for (int32_t slot = 0; slot < count; slot++) { PTO2TaskSlotState *ss = &ring.slot_states[slot]; + // Polling: fanin is a flat array of position-independent local-id + // integers on the payload, so only the two per-slot arena/SM + // pointers need relocating. There is no fanout_head/dep_pool graph + // and no host-seeded ready queue (the device boot scan classifies), + // so those relocation passes are gone. reloc(ss->task); reloc(ss->payload); - reloc(ss->fanout_head); - - PTO2TaskPayload *payload = &ring.task_payloads[slot]; - int32_t nf = payload->fanin_actual_count; - if (nf > PTO2_FANIN_INLINE_CAP) { - // host-orch does not yet relocate the multi-fanin spill pool, - // so a task with more than the inline cap of producers would - // ship un-relocated host pointers to the device. Fail loud - // rather than silently clamp (tensormap handles this via the - // on-device spill pool; that path is not wired here yet). - LOG_ERROR( - "host-orch: task slot %d has fanin %d > inline cap %d; multi-fanin spill relocation is not " - "implemented", - slot, nf, PTO2_FANIN_INLINE_CAP - ); - ok = false; - nf = PTO2_FANIN_INLINE_CAP; - } - for (int32_t i = 0; i < nf; i++) { - reloc(payload->fanin_inline_slot_states[i]); - } } } } - - // Per-ring fanout adjacency (dep_pool entries) built inline during host - // submit (orch_wire_fanin_task). Each live entry [tail, top) carries a - // consumer slot_state pointer (into the SM) and a next pointer (into the - // arena). - for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { - PTO2DepListPool &dp = rt->scheduler.ring_sched_state.dep_pool; - if (dp.base == nullptr || dp.capacity == 0) { - continue; - } - for (int32_t i = dp.tail; i < dp.top; i++) { - PTO2DepListEntry &e = dp.base[i % dp.capacity]; - reloc(e.slot_state); - reloc(e.next); - } - } - - // Ready queues seeded by push_ready_routed on the host. Each populated slot - // [dequeue_pos, enqueue_pos) holds a slot_state pointer into the SM. - auto reloc_ready = [&](PTO2ReadyQueue &q) { - if (q.slots == nullptr) { - return; - } - uint64_t enq = q.enqueue_pos.load(std::memory_order_relaxed); - uint64_t deq = q.dequeue_pos.load(std::memory_order_relaxed); - for (uint64_t pos = deq; pos < enq; pos++) { - reloc(q.slots[pos & q.mask].slot_state); - } - }; - for (int i = 0; i < PTO2_NUM_RESOURCE_SHAPES; i++) { - reloc_ready(rt->scheduler.ready_queues[i]); - reloc_ready(rt->scheduler.ready_sync_queues[i]); - } - reloc_ready(rt->scheduler.dummy_ready_queue); - for (int i = 0; i < PTO2_NUM_RESOURCE_SHAPES; i++) { - reloc_ready(rt->scheduler.early_dispatch_queues[i]); - } return ok; } @@ -685,7 +620,7 @@ register_callable_impl(const ChipCallable *callable, uint64_t (*upload_fn)(const extern "C" int bind_callable_to_runtime_impl( Runtime *runtime, const HostApi *api, const ChipStorageTaskArgs *orch_args, void *host_orch_func_ptr, const ArgDirection *signature, int sig_count, const uint64_t *ring_task_window, const uint64_t *ring_heap, - const uint64_t *ring_dep_pool + [[maybe_unused]] const uint64_t *ring_dep_pool // polling has no dep_pool; kept for ABI stability ) { if (runtime == nullptr) { LOG_ERROR("Runtime pointer is null"); @@ -710,19 +645,12 @@ extern "C" int bind_callable_to_runtime_impl( uint64_t eff_task_window_sizes[PTO2_MAX_RING_DEPTH]; uint64_t eff_heap_sizes[PTO2_MAX_RING_DEPTH]; - int32_t eff_dep_pool_capacities[PTO2_MAX_RING_DEPTH]; - if (!resolve_ring_config( - ring_task_window, ring_heap, ring_dep_pool, eff_task_window_sizes, eff_heap_sizes, eff_dep_pool_capacities - )) { + if (!resolve_ring_config(ring_task_window, ring_heap, eff_task_window_sizes, eff_heap_sizes)) { return -1; } const std::string task_window_log = format_ring_array(eff_task_window_sizes); const std::string heap_log = format_ring_array(eff_heap_sizes); - const std::string dep_pool_log = format_ring_array(eff_dep_pool_capacities); - LOG_INFO_V0( - "Ring buffer sizes: task_window=%s heap=%s dep_pool=%s", task_window_log.c_str(), heap_log.c_str(), - dep_pool_log.c_str() - ); + LOG_INFO_V0("Ring buffer sizes: task_window=%s heap=%s", task_window_log.c_str(), heap_log.c_str()); // Build device args: copy from input, replace host tensor pointers with device pointers ChipStorageTaskArgs device_args; @@ -826,8 +754,7 @@ extern "C" int bind_callable_to_runtime_impl( int64_t t_prebuilt_start = _now_ms(); DeviceArena host_arena; // libc malloc backend by default - PTO2RuntimeArenaLayout layout = - runtime_reserve_layout(host_arena, eff_task_window_sizes, eff_heap_sizes, eff_dep_pool_capacities); + PTO2RuntimeArenaLayout layout = runtime_reserve_layout(host_arena, eff_task_window_sizes, eff_heap_sizes); if (host_arena.commit(DeviceArena::kDefaultBaseAlign) == nullptr) { LOG_ERROR("Failed to commit host arena for prebuilt runtime image"); return -1; 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 235b7e9cb..960ef9396 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 @@ -264,24 +264,24 @@ static uint32_t next_fanin_seen_epoch(PTO2OrchestratorState *orch) { return next; } +// Polling: fanin is a flat array of position-independent producer local ids on +// the payload (no dep-pool spill, no producer pointers). The builder writes them +// directly into payload->fanin_local_ids as producers are appended, deduping by +// slot and hard-capping at PTO2_MAX_FANIN. self_local is this task's own local id +// (the consumer), used to bump each producer's last_consumer_local_id (the +// reclaim gate the host wait_for_consumers polls via completed_watermark). struct PTO2FaninBuilder { - PTO2FaninBuilder(PTO2OrchestratorState *orch, PTO2FaninPool &spill_pool, uint32_t seen_epoch) : + PTO2FaninBuilder(PTO2OrchestratorState *orch, PTO2TaskPayload *payload, int32_t self_local, uint32_t seen_epoch) : count(0), - spill_start(0), orch(orch), seen_epoch(seen_epoch), - spill_pool(spill_pool) {} + self_local(self_local), + payload(payload) {} int32_t count{0}; - int32_t spill_start{0}; PTO2OrchestratorState *orch{nullptr}; uint32_t seen_epoch{0}; - PTO2FaninPool &spill_pool; - PTO2TaskSlotState *inline_slots[PTO2_FANIN_INLINE_CAP]; - - template - PTO2FaninForEachReturn for_each(Fn &&fn) const { - return for_each_fanin_storage(inline_slots, count, spill_start, spill_pool, static_cast(fn)); - } + int32_t self_local{0}; + PTO2TaskPayload *payload{nullptr}; bool mark_seen(uint8_t prod_ring, int32_t prod_slot) { if (prod_ring >= PTO2_MAX_RING_DEPTH || prod_slot < 0) { @@ -301,169 +301,33 @@ static bool append_fanin_or_fail( PTO2OrchestratorState *orch, uint8_t prod_ring, int32_t prod_slot, PTO2TaskSlotState *prod_state, PTO2TaskId producer_task_id, PTO2FaninBuilder *fanin_builder ) { - // Decide-and-claim under the producer's fanout_lock. Two conditions make this - // resolved slot a non-dependency, and both must be checked together with the - // fanout_count++ so the producer cannot slip from live to consumed/reused in - // between: - // (1) Generation mismatch — the producer was CONSUMED, its slot - // reset_for_reuse'd and rebound to a newer task. The cached - // owner_task_id still resolves to this slot, but it no longer holds our - // producer; ++'ing it would corrupt an unrelated task. - // (2) Already CONSUMED in place — finished, output ready, no real edge. - // In either case, adding it to the fanin and bumping fanout_count would leave - // a stale ++/release pair (Orch-side wiring drops the fanout edge but keeps - // the fanin slot, so on_task_release still release_producer()'s it) that - // desyncs the slot's refcount (rc != fc) and wedges in-order reclaim. Claiming a live - // producer under the lock pins it: fanout_count now counts us, so it cannot - // reach CONSUMED (rc == fc) until we release it in on_task_release, keeping the - // slot's generation stable until then. check_and_handle_consumed flips - // COMPLETED->CONSUMED under the same lock, so the check and the ++ are atomic - // against the consume. fanout_count is lock-protected per the - // PTO2TaskSlotState contract. - // - // Dedup (mark_seen) happens HERE, gated on a live producer — NOT before the - // gone check. mark_seen keys only on (ring, slot); a stale owner that resolves - // to a reused slot must not record it as seen, or a later dependency on the - // live generation in the same submission would hit mark_seen and be skipped - // without claiming it (dropped edge). Marking only when !gone keeps the dedup - // keyed to the live producer, and doing it before the ++ still suppresses a - // double-count for a producer named twice in one submission. - prod_state->lock_fanout(); - bool gone = prod_state->task == nullptr || prod_state->task->task_id.local() != producer_task_id.local() || - prod_state->task_state.load(std::memory_order_acquire) == PTO2_TASK_CONSUMED; - bool claim = !gone && !fanin_builder->mark_seen(prod_ring, prod_slot); - if (claim) { - // Low bits hold the consumer count; bit31 is the scope ref. The consumer - // count must never carry into bit31 (would corrupt the scope-release - // flag) — true for any sane fanout (<< 2^31). - assert( - (prod_state->fanout_count & ~PTO2_FANOUT_SCOPE_BIT) < (PTO2_FANOUT_SCOPE_BIT - 1) && - "fanout consumer count overflow into scope bit" - ); - prod_state->fanout_count++; - } - prod_state->unlock_fanout(); -#if SIMPLER_ORCH_PROFILING - // lock + unlock always; one fanout_count store when we actually claim. - g_orch_args_atomic_count += claim ? 3 : 2; -#endif - // gone (stale/consumed) or an already-seen duplicate live producer: no new - // fanin edge either way. - if (!claim) { + // 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. + if (prod_state->task == nullptr || prod_state->task->task_id.local() != producer_task_id.local()) { return true; } - - if (fanin_builder->count < PTO2_FANIN_INLINE_CAP) { - fanin_builder->inline_slots[fanin_builder->count++] = prod_state; + // Dedup by (ring, slot). Single-ring hbg: prod_ring is always 0. + if (fanin_builder->mark_seen(prod_ring, prod_slot)) { return true; } - - PTO2FaninPool &fanin_pool = fanin_builder->spill_pool; - if (!fanin_pool.ensure_space(orch->sm_header->ring, 1)) { - orch_mark_fatal(orch, PTO2_ERROR_DEP_POOL_OVERFLOW); - return false; - } - int32_t spill_idx = fanin_pool.top; - PTO2FaninSpillEntry *entry = fanin_pool.alloc(); - if (entry == nullptr) { + if (fanin_builder->count >= PTO2_MAX_FANIN) { orch_mark_fatal(orch, PTO2_ERROR_DEP_POOL_OVERFLOW); return false; } - if (fanin_builder->count == PTO2_FANIN_INLINE_CAP) { - fanin_builder->spill_start = spill_idx; + 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. + if (fanin_builder->self_local > prod_state->last_consumer_local_id) { + prod_state->last_consumer_local_id = fanin_builder->self_local; } - entry->slot_state = prod_state; - fanin_builder->count++; return true; } static void scope_tasks_push(PTO2OrchestratorState *orch, PTO2TaskSlotState *task_slot_state); -static bool all_claimed_fanin_completed(const PTO2FaninBuilder &fanin_builder) { - if (fanin_builder.count == 0) return true; - return fanin_builder.for_each([](PTO2TaskSlotState *producer) -> bool { - return producer != nullptr && producer->task_state.load(std::memory_order_acquire) >= PTO2_TASK_COMPLETED; - }); -} - -static bool all_claimed_fanin_allow_early_resolve(const PTO2FaninBuilder &fanin_builder) { - if (fanin_builder.count == 0) return true; - return fanin_builder.for_each([](PTO2TaskSlotState *producer) -> bool { - return producer != nullptr && producer->allow_early_resolve; - }); -} - -// Record the dep_pool top reached after this slot's Orch-side wiring, so the -// scheduler can bound its reclaim scan to the entries this task allocated. -static void orch_mark_dep_pool_position(PTO2OrchestratorState *orch, PTO2TaskSlotState &slot_state) { - auto &rss = orch->scheduler->ring_sched_state; - slot_state.dep_pool_mark = rss.dep_pool.top; -#if SIMPLER_DFX - if (is_scope_stats_enabled()) { - rss.publish_dep_pool_snapshot(); - } -#endif -} - -// Wire a task with at least one live producer: prepend a fanout link on every -// live producer under its fanout_lock, seed dispatch_fanin for pre-completed -// codegen-flagged producers, then publish readiness via route_ready_once once -// the fanin refcount is already satisfied. -static void orch_wire_fanin_task(PTO2OrchestratorState *orch, PTO2TaskSlotState &slot_state, int32_t wfanin) { - PTO2SchedulerState *sched = orch->scheduler; - auto &rss = sched->ring_sched_state; - PTO2TaskPayload *payload = slot_state.payload; - slot_state.fanin_count = wfanin + 1; - - int32_t early_finished = 0; - bool early_disqualified = false; - for_each_fanin_slot_state(*payload, [&](PTO2TaskSlotState *producer) { - producer->lock_fanout(); - int32_t pstate = producer->task_state.load(std::memory_order_acquire); - if (!early_disqualified && !producer->allow_early_resolve) { - early_disqualified = true; - } - if (pstate >= PTO2_TASK_COMPLETED) { - early_finished++; - } else { - producer->fanout_head = rss.dep_pool.prepend(producer->fanout_head, &slot_state); - } - producer->unlock_fanout(); - }); - - // Pre-completed producers will not dispatch again. Seed dispatch_fanin only - // when every producer is codegen-flagged; one unflagged producer makes the - // direct-only early-dispatch candidate count unreachable by design. - if (!early_disqualified && early_finished != 0) { - payload->dispatch_fanin.fetch_add(early_finished, std::memory_order_acq_rel); - } - - int32_t init_rc = early_finished + 1; - int32_t new_rc = slot_state.fanin_refcount.fetch_add(init_rc, std::memory_order_acq_rel) + init_rc; - orch_mark_dep_pool_position(orch, slot_state); - if (new_rc >= slot_state.fanin_count) { - sched->route_ready_once(slot_state); - } -} - -static bool orch_wire_live_fanin_task(PTO2OrchestratorState *orch, PTO2TaskSlotState &slot_state, int32_t wfanin) { - auto &rss = orch->scheduler->ring_sched_state; - - // dep_pool is orchestrator-exclusive during host orchestration (no lock). - // The pool is sized for the whole graph (no reclaim during host - // orchestration, exactly like the task window and GM heap), so an exhausted - // pool latches PTO2_ERROR_DEP_POOL_OVERFLOW and reports the deadlock with the - // same structural + wall-clock logic the heap/task-window allocator uses. A - // false return also covers a fatal already latched elsewhere. - if (!rss.dep_pool.ensure_space(*rss.ring, wfanin)) { - orch->fatal = true; - return false; - } - - orch_wire_fanin_task(orch, slot_state, wfanin); - return true; -} - struct PTO2PreparedTask { PTO2TaskId task_id = PTO2TaskId::invalid(); PTO2TaskAllocResult alloc_result = {-1, 0, nullptr, nullptr}; @@ -505,9 +369,9 @@ static bool check_scope_can_accept_task(PTO2OrchestratorState *orch, PTO2TaskAll LOG_ERROR(" scope_task_count: %d", scope_task_count); LOG_ERROR(" active_tasks: %d / %d", active_count, allocator.window_size()); LOG_ERROR("Root Cause:"); - LOG_ERROR(" Tasks within a scope hold a fanout_count reference that is only"); - LOG_ERROR(" released at scope_end. When scope task count >= window_size,"); - LOG_ERROR(" no slots can be reclaimed -> deadlock."); + LOG_ERROR(" host_build_graph is whole-graph-resident: the host builds the entire"); + LOG_ERROR(" scope before the device runs, so no slots reclaim during the build."); + LOG_ERROR(" When scope task count >= window_size the ring overflows."); LOG_ERROR("Solution:"); LOG_ERROR(" 1. Reduce tasks per scope (use batching/unroll)"); LOG_ERROR(" 2. Increase task window (current: %d)", allocator.window_size()); @@ -571,7 +435,12 @@ static bool prepare_task( static_cast(block_num * __builtin_popcount(active_mask.core_mask())); out->slot_state->logical_block_num = block_num; out->slot_state->active_mask = active_mask; - // fanin_count is set during Orch-side wiring + // 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). + 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); return true; @@ -621,11 +490,9 @@ void PTO2OrchestratorState::begin_scope(PTO2ScopeMode mode) { if (is_scope_stats_enabled()) { uint8_t ring_id = 0; auto &alloc = orch->ring.task_allocator; + // Polling: no dep_pool to report (readiness is via completion_flags). int32_t dep_pool_tail = 0; int32_t dep_pool_top = 0; - if (orch->scheduler) { - orch->scheduler->ring_sched_state.read_dep_pool_snapshot(dep_pool_tail, dep_pool_top); - } scope_stats_begin( ring_id, alloc.task_tail(), alloc.task_head(), alloc.heap_tail(), alloc.heap_top(), dep_pool_tail, dep_pool_top, orch->tensor_map.current_used() @@ -650,11 +517,9 @@ void PTO2OrchestratorState::end_scope() { if (is_scope_stats_enabled()) { uint8_t ring_id = 0; auto &alloc = orch->ring.task_allocator; + // Polling: no dep_pool to report (readiness is via completion_flags). int32_t dep_pool_tail = 0; int32_t dep_pool_top = 0; - if (orch->scheduler) { - orch->scheduler->ring_sched_state.read_dep_pool_snapshot(dep_pool_tail, dep_pool_top); - } scope_stats_end( ring_id, alloc.task_tail(), alloc.task_head(), alloc.heap_tail(), alloc.heap_top(), dep_pool_tail, dep_pool_top, orch->tensor_map.current_used() @@ -842,7 +707,7 @@ static TaskOutputTensors submit_task_common( ); } - PTO2FaninBuilder fanin_builder(orch, orch->ring.fanin_pool, next_fanin_seen_epoch(orch)); + PTO2FaninBuilder fanin_builder(orch, &payload, static_cast(task_id.local()), next_fanin_seen_epoch(orch)); CYCLE_COUNT_LAP(g_orch_alloc_cycle); @@ -928,20 +793,9 @@ static TaskOutputTensors submit_task_common( task.packed_buffer_base = prepared.alloc_result.packed_base; task.packed_buffer_end = prepared.alloc_result.packed_end; - // fanout_count was already incremented per live producer inside - // append_fanin_or_fail, atomically with the consumed/generation check under - // the producer's fanout_lock. Doing it there (rather than a separate pass - // here) is what prevents a producer from transitioning to CONSUMED between - // the dependency decision and the claim. - int32_t inline_count = std::min(fanin_builder.count, PTO2_FANIN_INLINE_CAP); - // Store fanin metadata in payload for scheduler to iterate - payload.fanin_actual_count = fanin_builder.count; - payload.fanin_spill_start = fanin_builder.spill_start; - payload.fanin_spill_pool = &fanin_builder.spill_pool; - for (int i = 0; i < inline_count; i++) { - payload.fanin_inline_slot_states[i] = fanin_builder.inline_slots[i]; - } - + // append_fanin_or_fail wrote each producer's local id straight into + // payload.fanin_local_ids and bumped its last_consumer_local_id; the count is + // published in STEP 6 below. payload.init does not touch the fanin region. payload.init(args, result, prepared.alloc_result, layout); cur_slot_state.set_allow_early_resolve(args.allow_early_resolve()); @@ -982,39 +836,17 @@ static TaskOutputTensors submit_task_common( CYCLE_COUNT_LAP(g_orch_args_cycle); - // === STEP 6: wire on the orchestrator side and publish readiness === - // host_build_graph host-orch: the orchestrator runs to completion on the host - // before the device boots scheduler-only, so wire the fanout adjacency here — - // lock each producer, allocate dep_pool entries, and publish readiness — - // directly in submit instead of deferring it to a device-side wiring-queue - // drain. Zero-fanin tasks and tasks whose claimed producers are already - // completed need no fanout links or dep_pool entries and go straight to the - // ready queue; tasks with live producers allocate fanout links first. The - // dep_pool is sized for the whole graph (no reclaim during host - // orchestration, exactly like the task window and GM heap), so an exhausted - // pool latches PTO2_ERROR_DEP_POOL_OVERFLOW via dep_pool.ensure_space() — the - // same failure class as a task-window/heap overflow. The resulting - // fanout_head / dep-entry / ready-queue pointers are host-DDR addresses; - // runtime_maker relocates them to device addresses before H2D. - if (fanin_builder.count == 0) { - cur_slot_state.fanin_count = 1; - cur_slot_state.fanin_refcount.store(1, std::memory_order_release); - orch_mark_dep_pool_position(orch, cur_slot_state); - sched->push_ready_routed(&cur_slot_state); - } else if (all_claimed_fanin_completed(fanin_builder)) { - int32_t ready_seed = fanin_builder.count + 1; - cur_slot_state.fanin_count = ready_seed; - if (all_claimed_fanin_allow_early_resolve(fanin_builder)) { - payload.dispatch_fanin.store(fanin_builder.count, std::memory_order_release); - } - cur_slot_state.fanin_refcount.store(ready_seed, std::memory_order_release); - orch_mark_dep_pool_position(orch, cur_slot_state); - sched->push_ready_routed(&cur_slot_state); - } else { - if (!orch_wire_live_fanin_task(orch, cur_slot_state, fanin_builder.count)) { - return result; - } - } + // === STEP 6: publish the inline fanin count (device boot classifies) === + // Polling + host-orch: append_fanin_or_fail already wrote each producer's + // local id into payload.fanin_local_ids and bumped its last_consumer_local_id. + // All that remains is to record how many. There is NO fanout adjacency, NO + // dep_pool, and NO ready routing here — the device boot scan classifies every + // task exactly once (fanin_satisfied -> push_ready_routed, else register_wake) + // before the scheduler dispatch loop starts. Because fanin is now a flat array + // of position-independent integers, none of this needs host->device pointer + // relocation. + payload.fanin_count = fanin_builder.count; + (void)sched; CYCLE_COUNT_LAP(g_orch_fanin_cycle); CYCLE_COUNT_ORCH_SUBMIT_RECORD(task_id.raw); @@ -1192,9 +1024,7 @@ TaskOutputTensors PTO2OrchestratorState::alloc_tensors(const L0TaskArgs &args) { TaskOutputTensors outputs; outputs.set_task_id(prepared.task_id); payload.init(args, outputs, prepared.alloc_result, layout); - payload.fanin_actual_count = 0; - payload.fanin_spill_start = 0; - payload.fanin_spill_pool = &orch->ring.fanin_pool; + payload.fanin_count = 0; // hidden-alloc tasks have no producer dependencies CYCLE_COUNT_LAP(g_orch_args_cycle); if (prepared.slot_state != nullptr) { @@ -1215,7 +1045,16 @@ TaskOutputTensors PTO2OrchestratorState::alloc_tensors(const L0TaskArgs &args) { // codegen task there is no Arg-driven hint to honor here, so mark it // unconditionally. prepared.slot_state->allow_early_resolve = true; - prepared.slot_state->mark_completed(); + prepared.slot_state->mark_completed(); // host-visible task_state mirror + // Polling: pre-set the device-visible completion_flags byte 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.) + PTO2SharedMemoryRingHeader &done_ring = orch->sm_header->ring; + int32_t done_local = static_cast(prepared.task_id.local()); + done_ring.completion_flags[done_local & done_ring.task_window_mask].store(1, std::memory_order_release); } orch->inline_completed_tasks++; @@ -1244,13 +1083,6 @@ void PTO2OrchestratorState::mark_done() { if (total_tasks > 0) { LOG_INFO_V0("=== [Orchestrator] ring %d: total_tasks=%d ===", r, total_tasks); } - auto &fanin_pool = orch->ring.fanin_pool; - if (fanin_pool.top > 1) { - LOG_INFO_V0( - "=== [FaninPool %d] top=%d tail=%d used=%d high_water=%d capacity=%d ===", r, fanin_pool.top, - fanin_pool.tail, fanin_pool.top - fanin_pool.tail, fanin_pool.high_water, fanin_pool.capacity - ); - } } orch->sm_header->orchestrator_done.store(1, std::memory_order_release); orch->scope_tasks_size = 0; diff --git a/src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_ring_buffer.cpp b/src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_ring_buffer.cpp index c2d7e7660..21b9e6ff8 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_ring_buffer.cpp +++ b/src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_ring_buffer.cpp @@ -8,178 +8,10 @@ * See LICENSE in the root of the software repository for the full text of the License. * ----------------------------------------------------------------------------------------------------------- */ -/** - * PTO Runtime2 - Ring Buffer Implementation - * - * Implements DepListPool ring buffer for zero-overhead dependency management. - * TaskAllocator methods are defined inline in pto_ring_buffer.h. - * - * Based on: docs/RUNTIME_LOGIC.md - */ - -#include "pto_ring_buffer.h" -#include -#include -#include "common/unified_log.h" -#include "scheduler/pto_scheduler.h" - -static void latch_pool_error(std::atomic *error_code_ptr, int32_t error_code) { - if (error_code_ptr == nullptr) { - return; - } - int32_t expected = PTO2_ERROR_NONE; - error_code_ptr->compare_exchange_strong(expected, error_code, std::memory_order_acq_rel); -} - -// ============================================================================= -// Fanin Spill Pool Implementation -// ============================================================================= -void PTO2FaninPool::reclaim(PTO2SharedMemoryRingHeader &ring, int32_t sm_last_task_alive) { - if (sm_last_task_alive <= reclaim_task_cursor) return; - - int32_t scan_end = sm_last_task_alive; - for (int32_t task_id = reclaim_task_cursor; task_id < scan_end; ++task_id) { - PTO2TaskPayload &payload = ring.get_payload_by_task_id(task_id); - if (payload.fanin_spill_pool != this) { - continue; - } - - int32_t inline_count = std::min(payload.fanin_actual_count, PTO2_FANIN_INLINE_CAP); - int32_t spill_edge_count = payload.fanin_actual_count - inline_count; - if (spill_edge_count > 0) { - advance_tail(payload.fanin_spill_start + spill_edge_count); - } - } - reclaim_task_cursor = scan_end; -} - -bool PTO2FaninPool::ensure_space(PTO2SharedMemoryRingHeader &ring, int32_t needed) { - if (available() >= needed) return true; - - int spin_count = 0; - int32_t prev_last_alive = ring.fc.last_task_alive.load(std::memory_order_acquire); - uint64_t block_cycle0 = 0; // wall-clock anchor for the deadlock backstop - bool block_timing = false; // false until the first no-reclaim-progress spin - while (available() < needed) { - reclaim(ring, prev_last_alive); - if (available() >= needed) return true; - - spin_count++; - - int32_t cur_last_alive = ring.fc.last_task_alive.load(std::memory_order_acquire); - if (cur_last_alive > prev_last_alive) { - spin_count = 0; - prev_last_alive = cur_last_alive; - block_timing = false; - } else if ((spin_count & 1023) == 0) { - // A fatal latched elsewhere breaks this otherwise-unbounded spin; the - // caller maps the failed ensure_space to orch_mark_fatal. Cold path. - if (error_code_ptr != nullptr && error_code_ptr->load(std::memory_order_acquire) != PTO2_ERROR_NONE) { - return false; - } - // Absolute-time backstop, matching the task allocator: stable across - // chips/contention, unlike a fixed spin count. get_sys_cnt_aicpu() - // is an MMIO read, so sample it only once per 1024 spins. - uint64_t now = get_sys_cnt_aicpu(); - if (!block_timing) { - block_cycle0 = now; - block_timing = true; - } else if (now - block_cycle0 >= PTO2_ALLOC_DEADLOCK_TIMEOUT_CYCLES) { - int32_t current = ring.fc.current_task_index.load(std::memory_order_acquire); - LOG_ERROR("========================================"); - LOG_ERROR("FATAL: Fanin Spill Pool Deadlock Detected!"); - LOG_ERROR("========================================"); - LOG_ERROR("Fanin spill pool cannot reclaim space after ~500 ms (no progress)."); - LOG_ERROR( - " - Pool used: %d / %d (%.1f%%)", used(), capacity, - (capacity > 0) ? (100.0 * used() / capacity) : 0.0 - ); - LOG_ERROR(" - Pool top: %d (linear)", top); - LOG_ERROR(" - Pool tail: %d (linear)", tail); - LOG_ERROR(" - High water: %d", high_water); - LOG_ERROR(" - Needed: %d entries", needed); - LOG_ERROR(" - last_task_alive: %d (stuck here)", cur_last_alive); - LOG_ERROR(" - current_task: %d", current); - LOG_ERROR(" - In-flight tasks: %d", current - cur_last_alive); - LOG_ERROR("Diagnosis:"); - LOG_ERROR(" last_task_alive is not advancing, so fanin spill pool tail"); - LOG_ERROR(" cannot reclaim. Check TaskRing diagnostics for root cause."); - LOG_ERROR("Solution:"); - LOG_ERROR( - " Increase fanin spill pool capacity (current: %d, recommended: %d)", capacity, high_water * 2 - ); - LOG_ERROR(" Compile-time: PTO2_DEP_LIST_POOL_SIZE in pto_runtime2_types.h"); - LOG_ERROR(" Runtime env: PTO2_RING_DEP_POOL=%d", high_water * 2); - LOG_ERROR("========================================"); - latch_pool_error(error_code_ptr, PTO2_ERROR_DEP_POOL_OVERFLOW); - return false; - } - } - SPIN_WAIT_HINT(); - } - return true; -} - -// ============================================================================= -// Dependency List Pool Implementation -// ============================================================================= -void PTO2DepListPool::reclaim(PTO2SharedMemoryRingHeader &ring, int32_t sm_last_task_alive) { - if (sm_last_task_alive >= last_reclaimed + PTO2_DEP_POOL_CLEANUP_INTERVAL && sm_last_task_alive > 0) { - int32_t mark = ring.get_slot_state_by_task_id(sm_last_task_alive - 1).dep_pool_mark; - if (mark > 0) { - advance_tail(mark); - } - last_reclaimed = sm_last_task_alive; - } -} - -bool PTO2DepListPool::ensure_space(PTO2SharedMemoryRingHeader &ring, int32_t needed) { - if (available() >= needed) return true; - - int spin_count = 0; - int32_t prev_last_alive = ring.fc.last_task_alive.load(std::memory_order_acquire); - while (available() < needed) { - reclaim(ring, prev_last_alive); - if (available() >= needed) return true; - - spin_count++; - - // Progress detection: reset spin counter if last_task_alive advances - int32_t cur_last_alive = ring.fc.last_task_alive.load(std::memory_order_acquire); - if (cur_last_alive > prev_last_alive) { - spin_count = 0; - prev_last_alive = cur_last_alive; - } - if (spin_count >= PTO2_DEP_POOL_SPIN_LIMIT) { - int32_t current = ring.fc.current_task_index.load(std::memory_order_acquire); - LOG_ERROR("========================================"); - LOG_ERROR("FATAL: Dependency Pool Deadlock Detected!"); - LOG_ERROR("========================================"); - LOG_ERROR("DepListPool cannot reclaim space after %d spins (no progress).", spin_count); - LOG_ERROR( - " - Pool used: %d / %d (%.1f%%)", used(), capacity, - (capacity > 0) ? (100.0 * used() / capacity) : 0.0 - ); - LOG_ERROR(" - Pool top: %d (linear)", top); - LOG_ERROR(" - Pool tail: %d (linear)", tail); - LOG_ERROR(" - High water: %d", high_water); - LOG_ERROR(" - Needed: %d entries", needed); - LOG_ERROR(" - last_task_alive: %d (stuck here)", cur_last_alive); - LOG_ERROR(" - current_task: %d", current); - LOG_ERROR(" - In-flight tasks: %d", current - cur_last_alive); - LOG_ERROR("Diagnosis:"); - LOG_ERROR(" last_task_alive is not advancing, so dep pool tail"); - LOG_ERROR(" cannot reclaim. Check TaskRing diagnostics for root cause."); - LOG_ERROR("Solution:"); - LOG_ERROR(" Increase dep pool capacity (current: %d, recommended: %d)", capacity, high_water * 2); - LOG_ERROR(" Compile-time: PTO2_DEP_LIST_POOL_SIZE in pto_runtime2_types.h"); - LOG_ERROR(" Runtime env: PTO2_RING_DEP_POOL=%d", high_water * 2); - LOG_ERROR("========================================"); - latch_pool_error(error_code_ptr, PTO2_ERROR_DEP_POOL_OVERFLOW); - return false; - } - SPIN_WAIT_HINT(); - } - return true; -} +// The polling completion scheduler has no dep-pool / fanin-spill pool: producer +// dependencies are position-independent local-id integers on the payload and +// readiness is derived from the ring completion_flags. The pool ring buffers +// (PTO2FaninPool / PTO2DepListPool) that once lived here are gone, so this +// translation unit carries no definitions. PTO2TaskAllocator's methods are +// defined inline in pto_ring_buffer.h. 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 2f1bcae12..0ac79ba90 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 @@ -159,8 +159,12 @@ static bool wait_for_tensor_ready(PTO2Runtime *rt, const Tensor &tensor, bool wa int32_t local_id = slot.task->task_id.local(); uint64_t t0 = get_sys_cnt_aicpu(); int32_t spin_count = 0; - while ((slot.fanout_refcount.load(std::memory_order_acquire) & ~PTO2_FANOUT_SCOPE_BIT) < - (slot.fanout_count & ~PTO2_FANOUT_SCOPE_BIT)) { + // Polling: all consumers of this producer have retired once the per-ring + // completed_watermark reaches 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) { 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 4c9f37a18..4a7cf9ffc 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 @@ -170,9 +170,6 @@ struct AsyncWaitList { // entries[]). struct DrainCompletionSink { PTO2SchedulerState *sched{nullptr}; - PTO2TaskSlotState **deferred_release_slot_states{nullptr}; - int32_t *deferred_release_count{nullptr}; - int32_t deferred_release_capacity{0}; int32_t inline_completed{0}; #if SIMPLER_SCHED_PROFILING int32_t thread_idx{0}; @@ -181,8 +178,7 @@ struct AsyncWaitList { bool can_inline_complete() const { return sched != nullptr; } }; - // Inline-complete a NotDeferred task during drain. Returns false on - // deferred_release_slot_states overflow. + // Inline-complete a NotDeferred task during drain. bool try_inline_complete_locked(DrainCompletionSink &sink, PTO2TaskSlotState &slot_state); // Single-consumer drain: pop each published message in tail order and @@ -288,9 +284,7 @@ struct AsyncWaitList { template AsyncPollResult poll_and_complete( - AICoreCompletionMailbox *aicore_mailbox, PTO2SchedulerState *sched, - PTO2TaskSlotState **deferred_release_slot_states, int32_t &deferred_release_count, - int32_t deferred_release_capacity + AICoreCompletionMailbox *aicore_mailbox, PTO2SchedulerState *sched #if SIMPLER_SCHED_PROFILING , int thread_idx diff --git a/src/a2a3/runtime/host_build_graph/runtime/pto_orchestrator.h b/src/a2a3/runtime/host_build_graph/runtime/pto_orchestrator.h index b252f3a19..3582a8875 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/pto_orchestrator.h +++ b/src/a2a3/runtime/host_build_graph/runtime/pto_orchestrator.h @@ -44,12 +44,10 @@ * pools, scope arrays, plus the nested PTO2TensorMap layout). */ struct PTO2OrchestratorLayout { - size_t off_fanin_pool; size_t off_fanin_seen_epoch; size_t off_scope_tasks; size_t off_scope_begins; PTO2TensorMapLayout tensor_map; - int32_t dep_pool_capacity; int32_t scope_tasks_cap; uint64_t scope_stack_capacity; }; @@ -130,8 +128,7 @@ struct PTO2OrchestratorState { // tensor_map sub-layout) on the supplied arena. task_window_sizes feeds // the nested tensor_map layout. Returned layout is consumed by // init_from_layout. - static PTO2OrchestratorLayout - reserve_layout(DeviceArena &arena, int32_t task_window_size, int32_t dep_pool_capacity = PTO2_DEP_LIST_POOL_SIZE); + static PTO2OrchestratorLayout reserve_layout(DeviceArena &arena, int32_t task_window_size); // Phase 3a: write everything *except* arena-internal pointer fields. // sm_dev_base is the SM device address (only stored, never dereferenced); @@ -176,7 +173,7 @@ struct PTO2OrchProfilingData { int64_t submit_count; // Wait time tracking for blocking phases uint64_t alloc_wait_cycle; // Cycles spent waiting in unified alloc - uint64_t fanin_wait_cycle; // Cycles spent waiting in fanout_lock + uint64_t fanin_wait_cycle; // Legacy (wiring): fanout_lock wait; polling has no such lock // Atomic operation counts per phase uint64_t alloc_atomic_count; uint64_t args_atomic_count; diff --git a/src/a2a3/runtime/host_build_graph/runtime/pto_ring_buffer.h b/src/a2a3/runtime/host_build_graph/runtime/pto_ring_buffer.h index 49a848e2b..a39ab68dc 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/pto_ring_buffer.h +++ b/src/a2a3/runtime/host_build_graph/runtime/pto_ring_buffer.h @@ -385,25 +385,16 @@ class PTO2TaskAllocator { #endif /** - * Structural deadlock test on the reclaim head. + * Structural deadlock test on the reclaim head — inert under polling. * - * The head (oldest un-CONSUMED task, at last_task_alive) gates all - * reclamation. If it is COMPLETED and every consumer reference is released - * (low bits of fanout_refcount == consumer count) but the scope reference - * (bit31) is still unset, the only release left is its scope_end. Because - * this is evaluated while the orchestrator is blocked in alloc(), scope_end - * can never be reached -> provable deadlock, no timeout required. - * - * The COMPLETED guard is mandatory: a zero-consumer task has - * refcount == 0 == (count & ~SCOPE_BIT) from birth, before it has run. + * The wiring model used a per-task scope refcount (fanout_count/refcount) to + * prove a head-of-line deadlock without a timeout. Polling removes those + * fields, and host_build_graph is whole-graph-resident host-orchestrated: no + * task completes during host build (the device runs afterward), so the head is + * never COMPLETED here and the structural test cannot apply. A genuine + * ring/heap overflow during build is caught by the wall-clock backstop. */ - bool head_blocked_on_scope_end(int32_t head_task_id) const { - if (slot_states_ == nullptr) return false; - PTO2TaskSlotState &h = slot_states_[head_task_id & window_mask_]; - if (h.task_state.load(std::memory_order_acquire) != PTO2_TASK_COMPLETED) return false; - uint32_t rc = h.fanout_refcount.load(std::memory_order_acquire); - return rc == (h.fanout_count & ~PTO2_FANOUT_SCOPE_BIT); - } + bool head_blocked_on_scope_end(int32_t /*head_task_id*/) const { return false; } /** * Report deadlock with targeted diagnostics. scope_gated == true means the @@ -446,12 +437,9 @@ class PTO2TaskAllocator { // Head-task state dump: what the reclaim watermark is actually waiting on. if (slot_states_ != nullptr) { PTO2TaskSlotState &h = slot_states_[last_alive & window_mask_]; - uint32_t fc = h.fanout_count; - uint32_t rc = h.fanout_refcount.load(std::memory_order_acquire); LOG_ERROR( - " Head task %d: state=%d, consumers=%u/%u, scope_released=%d", last_alive, - static_cast(h.task_state.load(std::memory_order_acquire)), rc & ~PTO2_FANOUT_SCOPE_BIT, - fc & ~PTO2_FANOUT_SCOPE_BIT, (rc & PTO2_FANOUT_SCOPE_BIT) ? 1 : 0 + " Head task %d: state=%d, last_consumer=%d", last_alive, + static_cast(h.task_state.load(std::memory_order_acquire)), h.last_consumer_local_id ); } LOG_ERROR("Solution:"); @@ -478,286 +466,6 @@ class PTO2TaskAllocator { } }; -// ============================================================================= -// Fanin Spill Pool -// ============================================================================= - -/** - * Fanin spill pool structure - * - * True ring buffer for allocating spilled fanin entries. - * Entries are reclaimed when their consumer tasks become CONSUMED. - * - * Linear counters (top, tail) grow monotonically; the physical index - * is obtained via modulo: base[linear_index % capacity]. - */ -struct PTO2FaninPool { - PTO2FaninSpillEntry *base; // Pool base address - int32_t capacity; // Total number of entries - int32_t top; // Linear next-allocation counter (starts from 1) - int32_t tail; // Linear first-alive counter (entries before this are dead) - int32_t high_water; // Peak concurrent usage (top - tail) - int32_t reclaim_task_cursor{0}; // Last task id scanned for reclaim on this pool - - std::atomic *error_code_ptr = nullptr; - - void init(PTO2FaninSpillEntry *in_base, int32_t in_capacity, std::atomic *in_error_code_ptr) { - base = in_base; - capacity = in_capacity; - top = 1; - tail = 1; - high_water = 0; - reclaim_task_cursor = 0; - base[0].slot_state = nullptr; - error_code_ptr = in_error_code_ptr; - } - - void reclaim(PTO2SharedMemoryRingHeader &ring, int32_t sm_last_task_alive); - - bool ensure_space(PTO2SharedMemoryRingHeader &ring, int32_t needed); - - PTO2FaninSpillEntry *alloc() { - int32_t used = top - tail; - if (used >= capacity) { - LOG_ERROR("========================================"); - LOG_ERROR("FATAL: Fanin Spill Pool Overflow!"); - LOG_ERROR("========================================"); - LOG_ERROR("Fanin spill pool exhausted: %d entries alive (capacity=%d).", used, capacity); - LOG_ERROR(" - Pool top: %d (linear)", top); - LOG_ERROR(" - Pool tail: %d (linear)", tail); - LOG_ERROR(" - High water: %d", high_water); - LOG_ERROR("Solution:"); - LOG_ERROR(" Increase fanin spill pool capacity (current: %d, recommended: %d).", capacity, capacity * 2); - LOG_ERROR(" Compile-time: PTO2_DEP_LIST_POOL_SIZE in pto_runtime2_types.h"); - LOG_ERROR(" Runtime env: PTO2_RING_DEP_POOL=%d", capacity * 2); - LOG_ERROR("========================================"); - if (error_code_ptr) { - error_code_ptr->store(PTO2_ERROR_DEP_POOL_OVERFLOW, std::memory_order_release); - } - return nullptr; - } - int32_t idx = top % capacity; - top++; - used++; - if (used > high_water) high_water = used; - return &base[idx]; - } - - void advance_tail(int32_t new_tail) { - if (new_tail > tail) { - tail = new_tail; - } - } - - int32_t used() const { return top - tail; } - - int32_t available() const { return capacity - used(); } -}; - -template -using PTO2FaninCallbackResult = std::invoke_result_t; - -template -using PTO2FaninForEachReturn = std::conditional_t, void>, void, bool>; - -template -inline PTO2FaninForEachReturn for_each_fanin_storage( - InlineSlots &&inline_slot_states, int32_t fanin_count, int32_t spill_start, PTO2FaninPool &spill_pool, Fn &&fn -) { - using FaninCallbackResult = PTO2FaninCallbackResult; - static_assert( - std::is_same_v || std::is_same_v, - "fanin callback must return void or bool" - ); - - if constexpr (std::is_void_v) { - int32_t inline_count = std::min(fanin_count, PTO2_FANIN_INLINE_CAP); - for (int32_t i = 0; i < inline_count; i++) { - fn(inline_slot_states[i]); - } - - int32_t spill_count = fanin_count - inline_count; - if (spill_count <= 0) { - return; - } - - int32_t start_idx = spill_start % spill_pool.capacity; - int32_t first_count = std::min(spill_count, spill_pool.capacity - start_idx); - PTO2FaninSpillEntry *first = spill_pool.base + start_idx; - for (int32_t i = 0; i < first_count; i++) { - fn(first[i].slot_state); - } - - int32_t second_count = spill_count - first_count; - for (int32_t i = 0; i < second_count; i++) { - fn(spill_pool.base[i].slot_state); - } - return; - } else { - int32_t inline_count = std::min(fanin_count, PTO2_FANIN_INLINE_CAP); - for (int32_t i = 0; i < inline_count; i++) { - if (!fn(inline_slot_states[i])) { - return false; - } - } - - int32_t spill_count = fanin_count - inline_count; - if (spill_count <= 0) { - return true; - } - - int32_t start_idx = spill_start % spill_pool.capacity; - int32_t first_count = std::min(spill_count, spill_pool.capacity - start_idx); - PTO2FaninSpillEntry *first = spill_pool.base + start_idx; - for (int32_t i = 0; i < first_count; i++) { - if (!fn(first[i].slot_state)) { - return false; - } - } - - int32_t second_count = spill_count - first_count; - for (int32_t i = 0; i < second_count; i++) { - if (!fn(spill_pool.base[i].slot_state)) { - return false; - } - } - return true; - } -} - -template -inline PTO2FaninForEachReturn for_each_fanin_slot_state(const PTO2TaskPayload &payload, Fn &&fn) { - return for_each_fanin_storage( - payload.fanin_inline_slot_states, payload.fanin_actual_count, payload.fanin_spill_start, - *payload.fanin_spill_pool, static_cast(fn) - ); -} - -// ============================================================================= -// Dependency List Pool -// ============================================================================= - -/** - * Dependency list pool structure - * - * True ring buffer for allocating linked list entries. - * Entries are reclaimed when their producer tasks become CONSUMED, - * as tracked by the orchestrator via dep_pool_mark per task. - * - * Linear counters (top, tail) grow monotonically; the physical index - * is obtained via modulo: base[linear_index % capacity]. - */ -struct PTO2DepListPool { - PTO2DepListEntry *base; // Pool base address - int32_t capacity; // Total number of entries - int32_t top; // Linear next-allocation counter (starts from 1) - int32_t tail; // Linear first-alive counter (entries before this are dead) - int32_t high_water; // Peak concurrent usage (top - tail) - int32_t last_reclaimed{0}; // last_task_alive at last successful reclamation - - // Error code pointer for fatal error reporting (→ sm_header->orch_error_code) - std::atomic *error_code_ptr = nullptr; - - /** - * - * Initialize dependency list pool - * @param base Pool base address from shared memory - * @param capacity Total number of entries - */ - void init(PTO2DepListEntry *in_base, int32_t in_capacity, std::atomic *in_error_code_ptr) { - base = in_base; - capacity = in_capacity; - top = 1; // Start from 1, 0 means NULL/empty - tail = 1; // Match initial top (no reclaimable entries yet) - high_water = 0; - last_reclaimed = 0; - - // Initialize entry 0 as NULL marker - base[0].slot_state = nullptr; - base[0].next = nullptr; - - error_code_ptr = in_error_code_ptr; - } - - /** - * Reclaim dead entries based on scheduler's slot state dep_pool_mark. - * Safe to call multiple times — only advances tail forward. - * - * @param ring Ring header (for reading slot dep_pool_mark) - * @param sm_last_task_alive Current last_task_alive from shared memory - */ - void reclaim(PTO2SharedMemoryRingHeader &ring, int32_t sm_last_task_alive); - - /** - * Ensure dep pool for a specific ring has at least `needed` entries available. - * Spin-waits for reclamation if under pressure. Detects deadlock if no progress. - */ - bool ensure_space(PTO2SharedMemoryRingHeader &ring, int32_t needed); - - /** - * Allocate a single entry from the pool (single-thread per pool instance) - * - * @return Pointer to allocated entry, or nullptr on fatal error - */ - PTO2DepListEntry *alloc() { - int32_t used = top - tail; - if (used >= capacity) { - LOG_ERROR("========================================"); - LOG_ERROR("FATAL: Dependency Pool Overflow!"); - LOG_ERROR("========================================"); - LOG_ERROR("DepListPool exhausted: %d entries alive (capacity=%d).", used, capacity); - LOG_ERROR(" - Pool top: %d (linear)", top); - LOG_ERROR(" - Pool tail: %d (linear)", tail); - LOG_ERROR(" - High water: %d", high_water); - LOG_ERROR("Solution:"); - LOG_ERROR(" Increase dep pool capacity (current: %d, recommended: %d).", capacity, capacity * 2); - LOG_ERROR(" Compile-time: PTO2_DEP_LIST_POOL_SIZE in pto_runtime2_types.h"); - LOG_ERROR(" Runtime env: PTO2_RING_DEP_POOL=%d", capacity * 2); - LOG_ERROR("========================================"); - if (error_code_ptr) { - error_code_ptr->store(PTO2_ERROR_DEP_POOL_OVERFLOW, std::memory_order_release); - } - return nullptr; - } - int32_t idx = top % capacity; - top++; - used++; - if (used > high_water) high_water = used; - return &base[idx]; - } - - /** - * Advance the tail pointer, reclaiming dead entries. - * Called by the orchestrator based on last_task_alive advancement. - */ - void advance_tail(int32_t new_tail) { - if (new_tail > tail) { - tail = new_tail; - } - } - - /** - * Prepend a task ID to a dependency list - * - * O(1) operation: allocates new entry and links to current head. - * - * @param current_head Current list head offset (0 = empty list) - * @param task_slot Task slot to prepend - * @return New head offset - */ - PTO2DepListEntry *prepend(PTO2DepListEntry *cur, PTO2TaskSlotState *slot_state) { - PTO2DepListEntry *new_entry = alloc(); - if (!new_entry) return nullptr; - new_entry->slot_state = slot_state; - new_entry->next = cur; - return new_entry; - } - - int32_t used() const { return top - tail; } - - int32_t available() const { return capacity - used(); } -}; - // ============================================================================= // Ring Set (per-depth aggregate) // ============================================================================= @@ -768,7 +476,6 @@ struct PTO2DepListPool { */ struct PTO2RingSet { PTO2TaskAllocator task_allocator; - PTO2FaninPool fanin_pool; }; #endif // PTO_RING_BUFFER_H diff --git a/src/a2a3/runtime/host_build_graph/runtime/pto_runtime2.h b/src/a2a3/runtime/host_build_graph/runtime/pto_runtime2.h index 4ecd2a14a..ad7989e2f 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/pto_runtime2.h +++ b/src/a2a3/runtime/host_build_graph/runtime/pto_runtime2.h @@ -112,7 +112,6 @@ struct PTO2RuntimeArenaLayout { // Cached parameters (re-used by init_data + wire stages). uint64_t task_window_sizes[PTO2_MAX_RING_DEPTH]{}; uint64_t heap_sizes[PTO2_MAX_RING_DEPTH]{}; - int32_t dep_pool_capacities[PTO2_MAX_RING_DEPTH]{}; // Total arena byte size post-commit. Used by host to size the prebuilt // image buffer and as the rtMemcpy length. @@ -169,12 +168,10 @@ struct PTO2Runtime { * Returns the layout descriptor; caller commits/attaches the arena before * Phase 2/3. */ -PTO2RuntimeArenaLayout runtime_reserve_layout( - DeviceArena &arena, uint64_t task_window_size, int32_t dep_pool_capacity = PTO2_DEP_LIST_POOL_SIZE -); +PTO2RuntimeArenaLayout runtime_reserve_layout(DeviceArena &arena, uint64_t task_window_size); PTO2RuntimeArenaLayout runtime_reserve_layout( DeviceArena &arena, const uint64_t task_window_sizes[PTO2_MAX_RING_DEPTH], - const uint64_t heap_sizes[PTO2_MAX_RING_DEPTH], const int32_t dep_pool_capacities[PTO2_MAX_RING_DEPTH] + const uint64_t heap_sizes[PTO2_MAX_RING_DEPTH] ); /** 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 0548bd739..d9a3e65d8 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 @@ -82,7 +82,6 @@ // Memory pools (total = value, single ring) #define PTO2_HEAP_SIZE (256 * 1024 * 1024) // 256MB -#define PTO2_DEP_LIST_POOL_SIZE 16384 // Per-ring dependency list pool entries #define PTO2_TENSORMAP_POOL_SIZE (65536) // TensorMap entry pool #define PTO2_TENSORMAP_NUM_BUCKETS 4096 // Power of 2 for fast hash (4096×8B=32KB fits L1) @@ -104,6 +103,12 @@ // Fanin storage #define PTO2_FANIN_INLINE_CAP 64 +// Polling-scheduler inline fanin cap. The polling model stores producer +// dependencies as flat position-independent local-id integers on the payload +// (no dep-pool spill), so a task's fanin degree is hard-capped here. Must cover +// the worst-case fanin of any workload (paged_attention is the densest). +#define PTO2_MAX_FANIN 128 + // Dependency-degree diagnostic: warn once when a task's fanin or a producer's // fanout first exceeds this degree, so dense dependency graphs surface without // flooding the AICPU hot-path device log. @@ -170,25 +175,7 @@ struct PTO2OutputLayout { // Dependency List Entry // ============================================================================= -/** - * Fanin spill entry - * Stored in the dedicated fanin spill ring buffer. - */ -struct PTO2TaskSlotState; // Forward declaration -struct PTO2FaninPool; // Forward declaration -struct PTO2FaninSpillEntry { - PTO2TaskSlotState *slot_state; -}; -static_assert(sizeof(PTO2FaninSpillEntry) == sizeof(uintptr_t)); - -/** - * Dependency list entry (singly-linked list node) - * Stored in DepListPool ring buffer. - */ -struct PTO2DepListEntry { - PTO2TaskSlotState *slot_state; // Consumer slot state (direct pointer) - PTO2DepListEntry *next; // next entry -}; +struct PTO2TaskSlotState; // Forward declaration (defined below) // ============================================================================= // Task Descriptor @@ -278,10 +265,16 @@ struct PTO2TaskPayload { // === Cache lines 0-8 (576B) — metadata + inline fanin === int32_t tensor_count{0}; int32_t scalar_count{0}; - int32_t fanin_actual_count{0}; // Actual fanin count (without the +1 redundance) - int32_t fanin_spill_start{0}; // Linear start index in fanin spill pool (0 = no spill) - PTO2FaninPool *fanin_spill_pool{nullptr}; - PTO2TaskSlotState *fanin_inline_slot_states[PTO2_FANIN_INLINE_CAP]; + int32_t fanin_count{0}; // Producer dependency count (raw, no +1 redundance) + // Producer dependencies as position-independent local task ids. Single-ring + // hbg: every producer is ring 0, so no per-edge ring id is stored. Scanned + // by fanin_satisfied / classify_fanin_state against the ring completion_flags. + // Hard-capped at PTO2_MAX_FANIN (no dep-pool spill). + int32_t fanin_local_ids[PTO2_MAX_FANIN]; + // Reserved: preserves the early-dispatch block and tensors[] offsets. tensors + // must stay at byte 576 (AICore arg-materialization contract), so this fanin + // region keeps its original 528-byte footprint. + int32_t _fanin_reserved[3]; // Early-dispatch metadata (AICPU-side only). Ordered by descending // alignment (8B mask, 4B fanin, then 2B/1B counters and flags) so the block packs with no // internal padding. Kept here after the fanin array (not moved up front): on @@ -429,10 +422,7 @@ struct PTO2TaskPayload { }; // PTO2TaskPayload layout verification (offsetof requires complete type). -static_assert(offsetof(PTO2TaskPayload, fanin_spill_pool) == 16, "spill pool pointer layout drift"); -static_assert( - offsetof(PTO2TaskPayload, fanin_inline_slot_states) == 24, "inline fanin array must follow spill metadata" -); +static_assert(offsetof(PTO2TaskPayload, fanin_local_ids) == 12, "inline fanin id array must follow fanin_count"); static_assert( offsetof(PTO2TaskPayload, predicate) == 576, "dispatch predicate occupies cache line 9 at fixed byte 576 (before tensors, never moves)" @@ -452,74 +442,53 @@ static_assert( /** * Per-task slot scheduling state (scheduler-private, NOT in shared memory) * - * Consolidates all hot-path scheduling fields into a single cache-friendly - * structure (64 bytes = one cache line). Accessing any field of a task's - * slot state brings all related fields into the same cache line. + * 64 bytes = one cache line. Under the polling completion model a task's + * readiness is derived from its producers' completion_flags (in the ring + * header); producer completion is published by setting this task's own + * completion_flag + draining its wake list. There is no fanout adjacency, + * refcount, or per-task lock here. * - * Concurrency notes: - * - fanout_head, fanout_count protected by fanout_lock (per-task spinlock) - * - fanin_count set once at submission, read-only after (hot path for ready check) - * - task_state, fanin_refcount, fanout_refcount updated atomically + * task_state is retained (a COMPLETED store on completion) because the HOST + * still polls it: the completion-wait in pto_runtime2.cpp, the allocator + * deadlock detector, and the cold-path stall dump. completion_flags is the + * device-side readiness truth; task_state is the host-visible mirror. */ - -// fanout_count / fanout_refcount bit encoding (both uint32): -// bits [30:0] = consumer references (count: # consumers; refcount: # released) -// bit [31] = the owning scope's reference (PTO2_FANOUT_SCOPE_BIT) -// fanout_count is seeded to PTO2_FANOUT_SCOPE_BIT and ++'d per consumer, so it -// ends as (SCOPE_BIT | num_consumers). release adds 1 (consumer completion) or -// SCOPE_BIT (scope_end). CONSUMED iff fanout_refcount == fanout_count (every -// consumer released AND scope bit set). Keeping the scope ref in a distinct bit -// (rather than folding scope + consumers into one count) lets a consumer reach -// fanout_refcount == (fanout_count & ~PTO2_FANOUT_SCOPE_BIT) while the scope bit -// is still unset -- i.e. "all consumers done but scope still open" stays -// distinguishable from "fully consumed". The heap/task deadlock detector keys -// off exactly that complement: that condition with state==COMPLETED means the -// head can only be released by scope_end, which a blocked orchestrator can -// never reach -> provable deadlock. -static constexpr uint32_t PTO2_FANOUT_SCOPE_BIT = 0x80000000u; - -enum PTO2TaskLifecycleFlag : uint8_t { - PTO2_LIFECYCLE_FLAGS_NONE = 0, - PTO2_READY_CLAIMED = 1U << 0, - PTO2_COMPLETION_DONE = 1U << 1, - PTO2_SUBTASK_DEFERRED = 1U << 2, -}; - struct alignas(64) PTO2TaskSlotState { - // Fanout lock + list (accessed together under lock in on_task_complete) - std::atomic fanout_lock; // Per-task spinlock (0=unlocked, 1=locked) - uint32_t fanout_count; // SCOPE_BIT (owning scope) | number of consumers - - PTO2DepListEntry *fanout_head; // Pointer to first fanout entry (nullptr = empty) - - // Task state (completion, consumed check, ready check) - std::atomic task_state; // PENDING/COMPLETED/CONSUMED - - // Fanin (accessed together in release_fanin_and_check_ready) - std::atomic fanin_refcount; // Dynamic: counts completed producers - int32_t fanin_count; // Number of producer dependencies (set once by wiring) - - // Fanout refcount (read alongside fanout_count by consumer-wait checks) - std::atomic fanout_refcount; // Dynamic: low bits = released consumers, bit31 = scope released + // 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. + // 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. + int32_t last_consumer_local_id; + + // Host-visible completion mirror. PENDING at submit; COMPLETED at + // on_mixed_task_complete. Read by the host completion-wait / deadlock + // detector / cold-path dump; the device readiness path uses completion_flags. + std::atomic task_state; // --- Per-slot constant, re-bound by orch::prepare_task each submit --- - // Value is the same on every reuse (&task_payloads[slot] / &task_descriptors[slot]), - // but written here per-submit instead of in an O(window_size) init loop — - // these are the only "scale-dependent" pointers in this struct, so moving - // them out of init makes startup cost independent of task_window_size. PTO2TaskPayload *payload; PTO2TaskDescriptor *task; + // --- Wake list: last-fanin notification (intrusive, lock-free) --- + // A pending consumer whose fanin scan finds an unmet producer registers on + // that producer's wake list (CAS push through next_in_wake_list). On + // completion the producer atomic-exchanges wake_list_head to + // WAKE_LIST_SENTINEL and routes every waiter. Reset to nullptr at init. + std::atomic wake_list_head{nullptr}; + PTO2TaskSlotState *next_in_wake_list{nullptr}; + // --- Set per-submit (depend on task inputs) --- ActiveMask active_mask; // Bitmask of active subtask slots (set once) - // These one-byte flags live in the padding before dep_pool_mark to keep - // PTO2TaskSlotState at 64 bytes. - // Codegen early-dispatch hint, copied from Arg at submit. Lives on - // slot_state (not payload) so fanin walks read the already-hot producer - // slot_state cache line. + // Codegen early-dispatch hint, copied from Arg at submit. Stored for the + // Milestone-2 publish-status early-dispatch; no Milestone-1 hot-path reader. bool allow_early_resolve{false}; - std::atomic lifecycle_flags{PTO2_LIFECYCLE_FLAGS_NONE}; - int32_t dep_pool_mark{0}; // Dep pool top after Orch-side wiring + // Set by any subtask FIN that pushed a deferred-completion CONDITION to the + // runtime mailbox; read by the last subtask FIN to decide inline vs + // MPSC-deferred completion. The release write is sequenced before + // on_subtask_complete's acq_rel fetch_add and the acquire read after. + std::atomic any_subtask_deferred{false}; + uint8_t _async_pad{0}; std::atomic completed_subtasks{0}; // Each core completion increments by 1 int16_t total_required_subtasks{0}; // = logical_block_num * popcount(active_mask) @@ -546,117 +515,47 @@ struct alignas(64) PTO2TaskSlotState { return 0; } - /** - * Re-bind the per-slot payload/task pointers. Called by - * orch::prepare_task on every submit. Value is constant for a given - * slot, but we pay the cheap re-write each submit (both fields land on - * the same 64B slot_state cache line that prepare_task is already - * dirtying) to avoid the init-time per-slot loop. - */ void bind_buffers(PTO2TaskPayload *p, PTO2TaskDescriptor *t) { payload = p; task = t; } - void mark_completed() { - task_state.store(PTO2_TASK_COMPLETED, std::memory_order_release); - lifecycle_flags.fetch_or(PTO2_COMPLETION_DONE, std::memory_order_release); - } + // Host-visible completion mirror. The device readiness truth + // (completion_flags[slot]) is published by the scheduler's + // on_mixed_task_complete; this store makes the same fact visible to the + // host completion-wait / deadlock detector. + void mark_completed() { task_state.store(PTO2_TASK_COMPLETED, std::memory_order_release); } - bool is_completion_flag_set() const { - return (lifecycle_flags.load(std::memory_order_acquire) & PTO2_COMPLETION_DONE) != 0; - } - - // Set by any subtask FIN that pushed deferred-completion CONDITIONs to the - // runtime mailbox; read by the last subtask FIN to decide whether the task - // needs MPSC-deferred completion or can complete inline on this thread. The - // release write is sequenced before on_subtask_complete's acq_rel fetch_add - // and the acquire read after, so all earlier subtasks' writes are visible to - // the last subtask. - void mark_any_subtask_deferred() { lifecycle_flags.fetch_or(PTO2_SUBTASK_DEFERRED, std::memory_order_release); } + void mark_any_subtask_deferred() { any_subtask_deferred.store(true, std::memory_order_release); } - bool has_any_subtask_deferred() const { - return (lifecycle_flags.load(std::memory_order_acquire) & PTO2_SUBTASK_DEFERRED) != 0; - } + bool has_any_subtask_deferred() const { return any_subtask_deferred.load(std::memory_order_acquire); } void set_allow_early_resolve(bool v) { allow_early_resolve = v; } /** - * Reset dynamic scheduling fields to their pristine values. - * In host_build_graph this runs once per slot at init (pto_shared_memory.cpp) - * to zero the scheduling state before the host orchestrator populates it — - * there is no execution-time slot recycle (whole-graph-resident, no reclaim), - * so unlike the device-orch path this is not re-invoked after CONSUMED. - * - * Skips payload, task (immutable, bound once at init). - * Skips task_state: the orchestrator sets it to PENDING when it populates - * the slot. + * Reset dynamic scheduling fields to their pristine values. Runs once per + * slot at init (pto_shared_memory.cpp) — whole-graph-resident hbg has no + * execution-time slot recycle. Skips payload/task (bound once) and + * task_state (the orchestrator sets PENDING when it populates the slot). + * wake_list_head starts nullptr (open for registration), NOT SENTINEL. */ void reset_for_reuse() { - fanout_lock.store(0, std::memory_order_relaxed); - fanout_count = PTO2_FANOUT_SCOPE_BIT; // bit31 = owning-scope ref; consumers ++ into low bits - fanout_head = nullptr; - fanin_refcount.store(0, std::memory_order_relaxed); - fanout_refcount.store(0, std::memory_order_relaxed); + wake_list_head.store(nullptr, std::memory_order_relaxed); + next_in_wake_list = nullptr; + any_subtask_deferred.store(false, std::memory_order_relaxed); completed_subtasks.store(0, std::memory_order_relaxed); next_block_idx.store(0, std::memory_order_relaxed); - lifecycle_flags.store(PTO2_LIFECYCLE_FLAGS_NONE, std::memory_order_relaxed); allow_early_resolve = false; - // Note: payload early-dispatch fields (state, masks, fanin, publication count) - // are NOT reset here — this method skips the payload by contract. They are - // (re)initialized in PTO2TaskPayload::init on every submit, before the slot - // becomes visible to the scheduler. - } - - // === Per-task fanout spinlock === - // - // Used by BOTH the orchestrator and the scheduler. The fanout_lock MUST - // be held whenever reading or writing fanout_head / fanout_count, because - // the orchestrator adds consumers concurrently with the scheduler - // traversing the list after task completion. - -#if SIMPLER_ORCH_PROFILING || SIMPLER_SCHED_PROFILING - void lock_fanout(uint64_t &atomic_count, uint64_t &wait_cycle) { - uint64_t t0 = get_sys_cnt_aicpu(); - bool contended = false; - uint32_t atomic_ops = 0; - - for (;;) { - while (fanout_lock.load(std::memory_order_acquire) != 0) { - contended = true; - atomic_ops++; - SPIN_WAIT_HINT(); - } - int32_t expected = 0; - if (fanout_lock.compare_exchange_weak(expected, 1, std::memory_order_acquire, std::memory_order_relaxed)) { - atomic_ops++; - atomic_count += atomic_ops; - if (contended) { - wait_cycle += (get_sys_cnt_aicpu() - t0); - } - return; - } - contended = true; - atomic_ops++; - } - } -#endif - - void lock_fanout() { - for (;;) { - while (fanout_lock.load(std::memory_order_acquire) != 0) { - SPIN_WAIT_HINT(); - } - int32_t expected = 0; - if (fanout_lock.compare_exchange_weak(expected, 1, std::memory_order_acquire, std::memory_order_relaxed)) { - return; - } - } + // last_consumer_local_id is seeded in prepare_task once the id is known. + // Payload early-dispatch/fanin fields are (re)initialized in + // PTO2TaskPayload::init on every submit, before the slot is visible. } - - void unlock_fanout() { fanout_lock.store(0, std::memory_order_release); } }; static_assert(sizeof(PTO2TaskSlotState) == 64); +// Sentinel marking a wake list as "owner already completed; no more +// registrations accepted". Distinct from any real slot_state pointer. +inline PTO2TaskSlotState *const WAKE_LIST_SENTINEL = reinterpret_cast(static_cast(0x1)); + #endif // SRC_A2A3_RUNTIME_TENSORMAP_AND_RINGBUFFER_RUNTIME_PTO_RUNTIME2_TYPES_H_ 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 7d240f7f3..1d861b6cf 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 @@ -80,6 +80,14 @@ 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). + alignas(64) std::atomic completed_watermark; + // Layout metadata (set once at init) uint64_t task_window_size; int32_t task_window_mask; @@ -91,6 +99,12 @@ 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 (fanin_satisfied). + // Zeroed host-side at init. Indexed by local_id & task_window_mask. + std::atomic *completion_flags; + int32_t get_slot_by_task_id(int32_t local_task_id) { return local_task_id & task_window_mask; } PTO2TaskDescriptor &get_task_by_slot(int32_t slot) { return task_descriptors[slot]; } @@ -110,9 +124,9 @@ struct alignas(64) PTO2SharedMemoryRingHeader { } }; -static_assert(sizeof(PTO2SharedMemoryRingHeader) == 192, "PTO2SharedMemoryRingHeader layout drift"); +static_assert(sizeof(PTO2SharedMemoryRingHeader) == 256, "PTO2SharedMemoryRingHeader layout drift"); static_assert( - offsetof(PTO2SharedMemoryRingHeader, task_descriptors_offset) == 152, + offsetof(PTO2SharedMemoryRingHeader, task_descriptors_offset) == 160, "PTO2SharedMemoryRingHeader task_descriptors_offset layout drift" ); @@ -144,10 +158,10 @@ struct alignas(PTO2_ALIGN_SIZE) PTO2SharedMemoryHeader { std::atomic sched_error_thread; // Thread index of last error writer }; -static_assert(sizeof(PTO2SharedMemoryHeader) == 256, "PTO2SharedMemoryHeader layout drift"); -static_assert(offsetof(PTO2SharedMemoryHeader, total_size) == 200, "PTO2SharedMemoryHeader total_size layout drift"); +static_assert(sizeof(PTO2SharedMemoryHeader) == 320, "PTO2SharedMemoryHeader layout drift"); +static_assert(offsetof(PTO2SharedMemoryHeader, total_size) == 264, "PTO2SharedMemoryHeader total_size layout drift"); static_assert( - offsetof(PTO2SharedMemoryHeader, orch_error_code) == 208, "PTO2SharedMemoryHeader orch_error_code layout drift" + offsetof(PTO2SharedMemoryHeader, orch_error_code) == 272, "PTO2SharedMemoryHeader orch_error_code layout drift" ); // ============================================================================= @@ -260,7 +274,8 @@ struct PTO2RingSegmentOffsets { uint64_t descriptors; uint64_t payloads; uint64_t slot_states; - uint64_t end; // offset just past slot_states (total SM size) + uint64_t completion_flags; // polling-completion byte array (1 byte/slot) + uint64_t end; // offset just past completion_flags (total SM size) }; // Single source of truth for the SM segment layout. Returns offsets (not @@ -278,6 +293,8 @@ inline PTO2RingSegmentOffsets ring_segment_offsets(uint64_t task_window_size) no off += PTO2_ALIGN_UP(task_window_size * sizeof(PTO2TaskPayload), PTO2_ALIGN_SIZE); 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); o.end = off; return o; } @@ -297,4 +314,11 @@ 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 *>( + static_cast(sm_dev_base) + ring_segment_offsets(task_window_size).completion_flags + ); +} + } // namespace pto2_sm_layout diff --git a/src/a2a3/runtime/host_build_graph/runtime/scheduler/pto_scheduler.cpp b/src/a2a3/runtime/host_build_graph/runtime/scheduler/pto_scheduler.cpp index 72f38d299..4c4cd838c 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/scheduler/pto_scheduler.cpp +++ b/src/a2a3/runtime/host_build_graph/runtime/scheduler/pto_scheduler.cpp @@ -78,13 +78,6 @@ void PTO2SchedulerState::print_stats() { if (sched->ring_sched_state.last_task_alive > 0) { LOG_INFO_V0("Ring %d:", r); LOG_INFO_V0(" last_task_alive: %d", sched->ring_sched_state.last_task_alive); - auto &dp = sched->ring_sched_state.dep_pool; - if (dp.top > 0) { - LOG_INFO_V0( - " dep_pool: top=%d tail=%d used=%d high_water=%d capacity=%d", dp.top, dp.tail, dp.top - dp.tail, - dp.high_water, dp.capacity - ); - } } } #if SIMPLER_SCHED_PROFILING 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 a46f484df..c0b61c02d 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 @@ -14,15 +14,16 @@ * * The Scheduler is responsible for: * 1. Maintaining per-resource-shape ready queues - * 2. Tracking task state (PENDING -> COMPLETED -> CONSUMED) - * 3. Managing fanin/fanout refcounts for dependency resolution - * 4. Advancing last_task_alive for heap reclamation - * 5. Two-stage mixed-task completion (subtask done bits → mixed-task complete) + * 2. Polling-completion dependency resolution: a task is ready when every + * producer named in its inline fanin has set its completion_flags byte; + * 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) + * 4. Two-stage mixed-task completion (subtask done bits -> mixed-task complete) * - * The Scheduler runs on Device AI_CPU and processes: - * - Task state transitions based on fanin_refcount - * - Buffer lifecycle based on fanout_refcount - * - Ring pointer advancement for flow control + * The Scheduler runs on Device AI_CPU. host_build_graph is scheduler-only (the + * orchestrator runs to completion on the host); there is no on-device slot + * reclaim (whole-graph-resident), so last_task_alive is not advanced here. * * Based on: docs/RUNTIME_LOGIC.md */ @@ -419,9 +420,7 @@ struct PTO2SchedulerLayout { size_t off_dummy_ready_queue_slots; size_t off_early_dispatch_queue_slots[PTO2_NUM_RESOURCE_SHAPES]; size_t off_early_sync_start_queue_slots; - size_t off_dep_pool_entries; uint64_t ready_queue_capacity; - int32_t dep_pool_capacity; }; /** @@ -442,34 +441,12 @@ struct PTO2SchedulerState { int32_t last_task_alive; std::atomic advance_lock; // multi-thread CAS - // --- Cache Line 1+: Orch-side wiring dep_pool --- - alignas(64) PTO2DepListPool dep_pool; -#if SIMPLER_DFX - // Published only for scope_stats; orchestrator must not read dep_pool's non-atomic counters directly. - alignas(64) std::atomic dep_pool_snapshot_tail; - std::atomic dep_pool_snapshot_top; -#endif - - // Initialize arena-internal data + arena-external pointers; does NOT - // store dep_pool.base (that lives in the runtime arena and is wired - // by SchedulerState::wire_arena_pointers). The `ring` field stores - // the device address of the SM ring header — computed via offset - // arithmetic, no SM dereference. + // Polling: no per-ring dep_pool. Readiness is derived from the SM ring's + // completion_flags; there is no arena-side wiring pool to reserve or wire. + // The `ring` field stores the device address of the SM ring header — + // computed via offset arithmetic, no SM dereference. bool init_data_from_layout(void *sm_dev_base); void destroy(); - -#if SIMPLER_DFX - void publish_dep_pool_snapshot() { - dep_pool_snapshot_tail.store(dep_pool.tail, std::memory_order_release); - dep_pool_snapshot_top.store(dep_pool.top, std::memory_order_release); - } - - void read_dep_pool_snapshot(int32_t &tail, int32_t &top) const { - top = dep_pool_snapshot_top.load(std::memory_order_acquire); - tail = dep_pool_snapshot_tail.load(std::memory_order_acquire); - if (tail > top) tail = top; - } -#endif } ring_sched_state; // Ready queues remain global (scheduling is ring-agnostic) @@ -514,53 +491,116 @@ struct PTO2SchedulerState { } } - bool try_claim_ready_once(PTO2TaskSlotState &slot_state) { - uint8_t flags = slot_state.lifecycle_flags.load(std::memory_order_acquire); - for (;;) { - if ((flags & PTO2_READY_CLAIMED) != 0) return false; - uint8_t desired_flags = flags | PTO2_READY_CLAIMED; - if (slot_state.lifecycle_flags.compare_exchange_weak( - flags, desired_flags, std::memory_order_acq_rel, std::memory_order_acquire - )) { - return true; - } + // ---- 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 + // there is no per-edge ring indirection. + + bool fanin_satisfied(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.completion_flags[p.fanin_local_ids[i] & ring.task_window_mask].load(std::memory_order_acquire) == + 0) + return false; } + return true; } - // host_build_graph host-orch: the COMPLETED->CONSUMED flip existed only to - // gate execution-time reclaim (now removed) and to serialize against the - // orchestrator's concurrent producer claim — which cannot happen here, since - // the orchestrator runs to completion on the host before the device - // scheduler starts. Nothing reads CONSUMED during device execution - // (completion uses completed_tasks_; wait_for_tensor_ready's consumer wait - // keys on fanout_refcount), so the flip is gone. fanout_refcount is still - // bumped by release_producer / release_producer_scope so the host-side - // wait_for_consumers sees consumers retire. - - void release_producer(PTO2TaskSlotState &slot_state) { - slot_state.fanout_refcount.fetch_add(1, std::memory_order_acq_rel); + // First-unmet classification. Returns -1 (all fanins met -> route to ready) + // 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 { + const PTO2TaskPayload &p = *s->payload; + const PTO2SharedMemoryRingHeader &ring = *ring_sched_state.ring; + for (int32_t i = 0; i < p.fanin_count; i++) { + if (ring.completion_flags[p.fanin_local_ids[i] & ring.task_window_mask].load(std::memory_order_acquire) == + 0) + return i; + } + return -1; } - // Scope-end release: sets bit31 (PTO2_FANOUT_SCOPE_BIT) instead of bumping a - // consumer ref. Called exactly once per task from on_scope_end. Keeping it a - // distinct add lets the deadlock detector tell "waiting only on scope_end" - // (head COMPLETED, refcount == fanout_count & ~SCOPE_BIT) apart from - // "waiting on a consumer". - void release_producer_scope(PTO2TaskSlotState &slot_state) { - slot_state.fanout_refcount.fetch_add(PTO2_FANOUT_SCOPE_BIT, std::memory_order_acq_rel); + // Register `consumer` on `producer`'s wake list. If the producer already + // 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) { + PTO2SharedMemoryRingHeader &ring = *ring_sched_state.ring; + while (true) { + PTO2TaskSlotState *expected = producer->wake_list_head.load(std::memory_order_relaxed); + while (expected != WAKE_LIST_SENTINEL) { + consumer->next_in_wake_list = expected; + if (producer->wake_list_head.compare_exchange_weak( + expected, consumer, std::memory_order_acq_rel, std::memory_order_relaxed + )) { + return; + } + } + int32_t state = classify_fanin_state(consumer); + if (state < 0) { + push_ready_routed(consumer); + return; + } + producer = &ring.get_slot_state_by_task_id(consumer->payload->fanin_local_ids[state]); + } } -#if SIMPLER_ORCH_PROFILING || SIMPLER_SCHED_PROFILING - void release_producer(PTO2TaskSlotState &slot_state, uint64_t &atomic_count) { - slot_state.fanout_refcount.fetch_add(1, std::memory_order_acq_rel); - atomic_count += 1; // fanout_refcount.fetch_add - } + // 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 my_id = static_cast(slot_state.task->task_id.local()); + PTO2SharedMemoryRingHeader &ring = *ring_sched_state.ring; + + slot_state.mark_completed(); // host-visible mirror (task_state = COMPLETED) + ring.completion_flags[my_id & ring.task_window_mask].store(1, std::memory_order_release); + + 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; + if (waiter->payload->fanin_count == 1) { + push_ready_routed(waiter); // single-fanin waiter was waiting only on us + waiter = next; + continue; + } + int state = classify_fanin_state(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); + } + waiter = next; + } - void release_producer_scope(PTO2TaskSlotState &slot_state, uint64_t &atomic_count) { - slot_state.fanout_refcount.fetch_add(PTO2_FANOUT_SCOPE_BIT, std::memory_order_acq_rel); - atomic_count += 1; // fanout_refcount.fetch_add + // completed_watermark = highest id such that every task in [0, watermark] + // has its completion_flags byte set. The host wait_for_consumers gates on + // watermark >= producer.last_consumer_local_id, so the walk must extend to + // the full contiguous completed prefix — NOT cap at my_id. Capping at my_id + // makes the final value order-dependent: a low-id task completing after a + // higher one would leave the watermark stuck below the true prefix, hanging + // any wait_for_consumers whose last_consumer sits in the gap. + const int32_t submitted = ring.fc.current_task_index.load(std::memory_order_acquire); + int32_t w = ring.completed_watermark.load(std::memory_order_acquire); + while (w + 1 < submitted) { + int32_t next = w + 1; + if (ring.completion_flags[next & ring.task_window_mask].load(std::memory_order_acquire) == 0) break; + if (ring.completed_watermark.compare_exchange_weak( + w, next, std::memory_order_acq_rel, std::memory_order_acquire + )) { + w = next; + } + } } -#endif + + // 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 >= + // producer.last_consumer_local_id, not by bumping a producer refcount. // Early-dispatch release. If the now-ready task was pre-staged // (gated on a core), ring its DATA_MAIN_BASE high-32 doorbell RIGHT HERE in @@ -761,242 +801,16 @@ struct PTO2SchedulerState { return true; } - // Event-driven candidate detection (the dual of fanin_refcount/ready). Called after a - // FLAGGED producer `p` publishes blocks (normal dispatch, early-dispatch release, or the - // sync_start drain), but no-ops until every logical block is launch-visible. Only then - // does it walk p's fanout and bump each consumer's - // dispatch_fanin. A consumer whose dispatch_fanin reaches fanin_actual_count (= every - // producer is flagged-and-fully-dispatched, or was already complete when the consumer was - // wired) is an early-dispatch candidate: CAS NONE->STAGING (exactly-once) and push to - // early_dispatch_queues[shape] (or early_sync_start_queue for a require_sync_start cohort) - // for the drain to pre-stage. The full-publication gate is load-bearing: a consumer that - // pre-occupies cores off a producer whose later blocks are not yet launch-visible would gate the - // very cores those blocks need, starving the producer so it never completes and the - // rendezvous never rings — a resource deadlock. published_block_count advances after - // every placement path publishes its claimed range. Once-guarded per producer. - // Only codegen-flagged producers propagate: a task's successors early-dispatch off its - // DIRECT producers' marks, never an inherited chain. - void propagate_dispatch_fanin(PTO2TaskSlotState &p) { - if (!p.allow_early_resolve) return; // only codegen-flagged (direct) producers propagate - if (p.payload->published_block_count.load(std::memory_order_seq_cst) < p.logical_block_num) return; - if (p.payload->early_dispatch_state.load(std::memory_order_acquire) == PTO2_EARLY_DISPATCH_STAGING) return; - uint8_t launch_state = p.payload->early_dispatch_launch_state.load(std::memory_order_acquire); - if (launch_state == PTO2_EARLY_DISPATCH_LAUNCH_RINGING) return; - if (p.active_mask.requires_sync_start()) { - bool was_pre_staged = false; - for (int w = 0; w < PTO2_EARLY_DISPATCH_CORE_MASK_WORDS; w++) { - was_pre_staged |= p.payload->staged_core_mask[w].load(std::memory_order_acquire) != 0; - } - if (was_pre_staged && launch_state != PTO2_EARLY_DISPATCH_LAUNCH_COMPLETE) return; - } - if (p.payload->dispatch_propagated.exchange(1, std::memory_order_acq_rel) != 0) - return; // already propagated once - p.lock_fanout(); - PTO2DepListEntry *edge = p.fanout_head; // snapshot head, walk lock-free (fanout stable by dispatch) - p.unlock_fanout(); - for (; edge != nullptr; edge = edge->next) { - PTO2TaskSlotState *c = edge->slot_state; - if (c->active_mask.has_predicate()) continue; // predicated consumers never early-dispatch - // Compare to fanin_actual_count (the real producer-edge count), NOT - // fanin_count: fanin_count = fanin_actual_count + 1 (a self/wiring +1 that - // ready_fanin gets but dispatch_fanin does not). dispatch_fanin starts at - // the wiring-time flagged-pre-completed seed and is bumped here by flagged - // producers; reaching fanin_actual_count means every producer is - // flagged-and-fully-published or was pre-completed. An unflagged producer leaves the - // seed short and never bumps, so this stays unreachable for that consumer. - int32_t nf = c->payload->dispatch_fanin.fetch_add(1, std::memory_order_acq_rel) + 1; - if (nf != c->payload->fanin_actual_count) continue; - PTO2ResourceShape shape = c->active_mask.to_shape(); - if (shape != PTO2ResourceShape::AIC && shape != PTO2ResourceShape::AIV && shape != PTO2ResourceShape::MIX) - continue; - // sync_start blocks launch as an atomic cohort. They ride the early-dispatch path - // too, but through a dedicated shape-agnostic queue (early_sync_start_queue), - // drained as the highest tier in try_early_dispatch via the gated drain + - // running-slot rendezvous (enter_drain_mode pre-stages every block gated — MIX - // with per-core split placement — and the rendezvous rings them together once - // every gated core occupies a running slot and the producer released). They must - // NOT enter early_dispatch_queues[shape]: that path's partial range-claim would - // strand gated cohort blocks nobody rings, so the rendezvous never reaches - // block_num. The rendezvous counts CORES (a MIX block spans a cluster), so one - // shape-agnostic queue suffices. - uint8_t expect = PTO2_EARLY_DISPATCH_NONE; // exactly-once: only the CAS winner enqueues - if (!c->payload->early_dispatch_state.compare_exchange_strong( - expect, PTO2_EARLY_DISPATCH_STAGING, std::memory_order_seq_cst, std::memory_order_seq_cst - )) - continue; - uint64_t task_id = static_cast(c->task->task_id.raw); - if (c->active_mask.requires_sync_start()) early_sync_start_queue.push_tagged(c, task_id); - else early_dispatch_queues[static_cast(shape)].push_tagged(c, task_id); - } - } - - // Collects consumers released via the early-dispatch-doorbell path during a - // single on_task_complete fanout walk, so their dispatch_fanin - // propagation runs AFTER the walk — never between two siblings' doorbells. - struct EarlyDispatchReleaseSink { - static constexpr int CAP = 32; - PTO2TaskSlotState *items[CAP]; - int n = 0; - inline bool push(PTO2TaskSlotState *s) { - if (n >= CAP) return false; - items[n++] = s; - return true; - } - }; - - inline bool try_early_dispatch_release(PTO2TaskSlotState &slot_state, EarlyDispatchReleaseSink *sink = nullptr) { - // A predicated task is evaluated at the ready point, never early-dispatched. - if (slot_state.active_mask.has_predicate()) return false; - // Never staged => CAS NONE->DISPATCHED wins => dispatch normally. - uint8_t expect = PTO2_EARLY_DISPATCH_NONE; - if (slot_state.payload->early_dispatch_state.compare_exchange_strong( - expect, PTO2_EARLY_DISPATCH_DISPATCHED, std::memory_order_seq_cst, std::memory_order_seq_cst - )) { - return false; - } - // route_ready_once admits one caller, but keep the helper defensive: a - // duplicate that observes or loses an in-progress launch must never - // route the same partial task a second time. - if (expect != PTO2_EARLY_DISPATCH_STAGING) return true; - bool sync_start = slot_state.active_mask.requires_sync_start(); - if (!sync_start && !try_claim_early_dispatch_launch(*slot_state.payload)) return true; - expect = PTO2_EARLY_DISPATCH_STAGING; - slot_state.payload->early_dispatch_state.compare_exchange_strong( - expect, PTO2_EARLY_DISPATCH_DISPATCHED, std::memory_order_seq_cst, std::memory_order_seq_cst - ); - // Producer released: ring the gated cores. A non-sync_start consumer launches - // each block the instant its doorbell fires. A sync_start consumer instead holds - // for the rendezvous — every gated core must occupy a running slot first — so the - // flip to DISPATCHED above is only the producer-released half; maybe_rendezvous_ring - // rings now iff running_slot_count already reached popcount(staged_core_mask) (all - // gated cores took idle running slots), else the last pending->running promotion rings. - bool launched = true; - if (sync_start) { - launched = maybe_rendezvous_ring(slot_state); - } else { - // Destructively claim every published bit. A stager racing this pass - // can claim only bits that land after the exchange, so each gated core - // has exactly one doorbell writer. - for (int w = 0; w < PTO2_EARLY_DISPATCH_CORE_MASK_WORDS; w++) { - uint64_t owned = claim_all_staged_doorbell_bits(slot_state.payload->staged_core_mask[w]); - ring_staged_doorbell_bits(w, owned); - } - wmb(); - slot_state.payload->early_dispatch_launch_state.store( - PTO2_EARLY_DISPATCH_LAUNCH_COMPLETE, std::memory_order_seq_cst - ); - } - // This pre-staged consumer was just released by its doorbell — it starts - // running NOW, so propagate dispatch_fanin to ITS consumers (only if it is - // itself codegen-flagged; the gate inside no-ops otherwise). Defer it via the - // sink so it runs after the whole fanout walk: doing it inline here would - // delay the doorbells of later consumers in the same producer's fanout. - // Fallback to inline if no sink / sink full. - if (launched) { - if (sink == nullptr || !sink->push(&slot_state)) { - propagate_dispatch_fanin(slot_state); - } - } - // No explicit removal from the cross-thread queue: a still-queued entry for - // this consumer is now DISPATCHED and is dropped when a peer pops it. - // Fully pre-staged => skip the ready queue. Partially staged SPMD consumer => - // fall through so the caller pushes C; dispatch resumes from next_block_idx. - return slot_state.next_block_idx.load(std::memory_order_seq_cst) >= slot_state.logical_block_num; - } - - bool route_ready_once(PTO2TaskSlotState &slot_state, EarlyDispatchReleaseSink *sink = nullptr) { - if (!try_claim_ready_once(slot_state)) return false; - - // Early-dispatch: pre-staged tasks are released by doorbell - // here, skipping the ready-queue round-trip entirely. - bool early_handled = try_early_dispatch_release(slot_state, sink); - if (slot_state.active_mask.requires_sync_start()) { - bool drain_owned = publish_ready_to_early_sync_drain(*slot_state.payload); - if (early_handled || drain_owned) return true; - } else if (early_handled) { - return true; - } - - PTO2ResourceShape shape = slot_state.active_mask.to_shape(); - if (shape == PTO2ResourceShape::DUMMY || - (slot_state.active_mask.has_predicate() && !slot_state.payload->predicate.pass())) { - dummy_ready_queue.push(&slot_state); - } else if (slot_state.active_mask.requires_sync_start()) { - ready_sync_queues[static_cast(shape)].push(&slot_state); - } else { - ready_queues[static_cast(shape)].push(&slot_state); - } - return true; - } - -#if SIMPLER_ORCH_PROFILING || SIMPLER_SCHED_PROFILING - bool route_ready_once( - PTO2TaskSlotState &slot_state, uint64_t &atomic_count, uint64_t &push_wait, - EarlyDispatchReleaseSink *sink = nullptr - ) { - uint8_t flags = slot_state.lifecycle_flags.load(std::memory_order_acquire); - atomic_count += 1; // lifecycle_flags load - for (;;) { - if ((flags & PTO2_READY_CLAIMED) != 0) return false; - uint8_t desired_flags = flags | PTO2_READY_CLAIMED; - if (slot_state.lifecycle_flags.compare_exchange_weak( - flags, desired_flags, std::memory_order_acq_rel, std::memory_order_acquire - )) { - atomic_count += 1; // lifecycle_flags CAS - break; - } - atomic_count += 1; // failed lifecycle_flags CAS - } - - // Early-dispatch: pre-staged tasks are released by doorbell - // here, skipping the ready-queue round-trip entirely. - bool early_handled = try_early_dispatch_release(slot_state, sink); - if (slot_state.active_mask.requires_sync_start()) { - bool drain_owned = publish_ready_to_early_sync_drain(*slot_state.payload); - if (early_handled || drain_owned) return true; - } else if (early_handled) { - return true; - } - - PTO2ResourceShape shape = slot_state.active_mask.to_shape(); - if (shape == PTO2ResourceShape::DUMMY || - (slot_state.active_mask.has_predicate() && !slot_state.payload->predicate.pass())) { - dummy_ready_queue.push(&slot_state, atomic_count, push_wait); - } else if (slot_state.active_mask.requires_sync_start()) { - ready_sync_queues[static_cast(shape)].push(&slot_state, atomic_count, push_wait); - } else { - ready_queues[static_cast(shape)].push(&slot_state, atomic_count, push_wait); - } - return true; - } -#endif - - bool release_fanin_and_check_ready(PTO2TaskSlotState &slot_state, EarlyDispatchReleaseSink *sink = nullptr) { - // Atomically increment fanin_refcount and check if all producers are done - // ACQ_REL on fanin_refcount already synchronizes with the orchestrator's - // init release, making fanin_count visible — plain load suffices. - int32_t new_refcount = slot_state.fanin_refcount.fetch_add(1, std::memory_order_acq_rel) + 1; - - if (new_refcount == slot_state.fanin_count) { - return route_ready_once(slot_state, sink); - } - return false; - } - -#if SIMPLER_ORCH_PROFILING || SIMPLER_SCHED_PROFILING - bool release_fanin_and_check_ready( - PTO2TaskSlotState &slot_state, uint64_t &atomic_count, uint64_t &push_wait, - EarlyDispatchReleaseSink *sink = nullptr - ) { - int32_t new_refcount = slot_state.fanin_refcount.fetch_add(1, std::memory_order_acq_rel) + 1; - atomic_count += 1; // fanin_refcount.fetch_add - - if (new_refcount == slot_state.fanin_count) { - return route_ready_once(slot_state, atomic_count, push_wait, sink); - } - return false; - } -#endif + // Milestone 1: early-dispatch (predicated / allow_early_resolve) is stubbed. + // This producer-push propagation walked the wiring fanout list bumping each + // consumer's dispatch_fanin to pre-stage early-dispatch candidates — all of + // which (fanout_head, dispatch_fanin, fanin_actual_count, dispatch_propagated) + // are gone under polling. Nothing pre-stages into early_dispatch_queues / + // early_sync_start_queue, so tasks reach cores only through the normal ready + // path (wake drain -> push_ready_routed). Milestone 2 replaces this with the + // consumer-pull publish_flags design. sync_start cohorts still launch via + // ready_sync_queues (unaffected). + void propagate_dispatch_fanin(PTO2TaskSlotState & /*p*/) {} int get_ready_tasks_batch(PTO2ReadyQueue *queues, PTO2ResourceShape shape, PTO2TaskSlotState **out, int max_count) { return queues[static_cast(shape)].pop_batch(out, max_count); @@ -1011,22 +825,11 @@ struct PTO2SchedulerState { } #endif - void on_scope_end(PTO2TaskSlotState **task_slot_states, int32_t count) { -#if SIMPLER_ORCH_PROFILING - extern uint64_t g_orch_scope_end_atomic_count; - if (count > 0) __builtin_prefetch(task_slot_states[0], 1, 0); - for (int32_t i = 0; i < count; i++) { - if (i + 1 < count) __builtin_prefetch(task_slot_states[i + 1], 1, 0); - release_producer_scope(*task_slot_states[i], g_orch_scope_end_atomic_count); - } -#else - if (count > 0) __builtin_prefetch(task_slot_states[0], 1, 0); - for (int32_t i = 0; i < count; i++) { - if (i + 1 < count) __builtin_prefetch(task_slot_states[i + 1], 1, 0); - release_producer_scope(*task_slot_states[i]); - } -#endif - } + // 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 + // 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*/) {} /** * Subtask completion: atomic counter model. @@ -1065,117 +868,26 @@ struct PTO2SchedulerState { int thread_idx #endif ) { + // Polling completion: publish the host-visible task_state mirror + the + // device-visible completion_flags byte, 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); #if SIMPLER_SCHED_PROFILING - CompletionStats stats = {0, 0, 0, true}; -#else - uint32_t consumer_walk_count = 0; -#endif -#if SIMPLER_SCHED_PROFILING - extern uint64_t g_sched_lock_cycle[], g_sched_fanout_cycle[]; - extern uint64_t g_sched_lock_atomic_count[], g_sched_lock_wait_cycle[]; - extern uint64_t g_sched_fanout_atomic_count[], g_sched_push_wait_cycle[]; - uint64_t lock_atomics = 0, lock_wait = 0; - PTO2_SCHED_CYCLE_START(); -#endif - -#if SIMPLER_SCHED_PROFILING - slot_state.lock_fanout(lock_atomics, lock_wait); + (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}; #else - slot_state.lock_fanout(); -#endif - slot_state.mark_completed(); - PTO2DepListEntry *current = slot_state.fanout_head; // Protected by fanout_lock - slot_state.unlock_fanout(); - -#if SIMPLER_SCHED_PROFILING - lock_atomics += 3; // task_state.store + lifecycle_flags.fetch_or + unlock.store - g_sched_lock_atomic_count[thread_idx] += lock_atomics; - g_sched_lock_wait_cycle[thread_idx] += lock_wait; - PTO2_SCHED_CYCLE_LAP(g_sched_lock_cycle[thread_idx]); -#endif - - // Fanout: notify consumers. A pre-staged consumer that becomes ready has - // its doorbell rung INLINE (db = nullptr) the moment its node is reached, - // not batched to after the whole walk — so a flagged consumer near the - // front of the list starts immediately and overlaps the remaining - // release_fanin work for the other consumers, instead of waiting for the - // full O(fanout-degree) walk (~5us for a 50-consumer producer). - // - // Safe on silicon: the producer's slot is already COMPLETED here — every - // SPMD block has FIN'd AND dcci-flushed its output to HBM before - // on_task_complete runs — so a released consumer never reads stale - // producer output. (Batching used to align the released wave, but pushed - // every doorbell to the end of the walk, defeating the whole point of - // early-dispatch: minimal producer-end -> consumer-start.) -#if SIMPLER_SCHED_PROFILING - uint64_t fanout_atomics = 0, push_wait = 0; -#endif - // Doorbells for released pre-staged consumers fire INLINE in the walk - // below; their dispatch_fanin propagation is collected here and replayed - // after the walk, so no consumer's doorbell waits on a sibling's propagate. - EarlyDispatchReleaseSink rel_sink; - while (current != nullptr) { - PTO2TaskSlotState &consumer_slot = *current->slot_state; -#if SIMPLER_SCHED_PROFILING - stats.fanout_edges++; - if (release_fanin_and_check_ready(consumer_slot, fanout_atomics, push_wait, &rel_sink)) { - stats.tasks_enqueued++; - } -#else - consumer_walk_count++; - release_fanin_and_check_ready(consumer_slot, &rel_sink); -#endif - current = current->next; - } - for (int i = 0; i < rel_sink.n; i++) { - propagate_dispatch_fanin(*rel_sink.items[i]); - } - -#if SIMPLER_SCHED_PROFILING - g_sched_fanout_atomic_count[thread_idx] += fanout_atomics; - g_sched_push_wait_cycle[thread_idx] += push_wait; - PTO2_SCHED_CYCLE_LAP(g_sched_fanout_cycle[thread_idx]); - return stats; -#else - return consumer_walk_count; + return 0; #endif } - /** - * Cold path: release producers (fanin traversal) + check self for CONSUMED. - * Returns fanin edge count for profiling. - */ - -#if SIMPLER_SCHED_PROFILING - int32_t on_task_release(PTO2TaskSlotState &slot_state, int32_t thread_idx) { - PTO2_SCHED_CYCLE_START(); - extern uint64_t g_sched_fanin_cycle[], g_sched_fanin_atomic_count[]; - extern uint64_t g_sched_complete_count[]; - uint64_t fanin_atomics = 0; -#else - int32_t on_task_release(PTO2TaskSlotState &slot_state) { -#endif - PTO2TaskPayload *payload = slot_state.payload; - for_each_fanin_slot_state(*payload, [&](PTO2TaskSlotState *producer_slot_state) { -#if SIMPLER_SCHED_PROFILING - release_producer(*producer_slot_state, fanin_atomics); -#else - release_producer(*producer_slot_state); -#endif - }); -#if SIMPLER_SCHED_PROFILING - g_sched_fanin_atomic_count[thread_idx] += fanin_atomics; - PTO2_SCHED_CYCLE_LAP(g_sched_fanin_cycle[thread_idx]); -#endif - - // host-orch: no self CONSUMED flip — see release_producer. The task's - // own fanout_refcount is already complete via the consumer/scope - // releases above; nothing reads CONSUMED during device execution. -#if SIMPLER_SCHED_PROFILING - g_sched_complete_count[thread_idx]++; -#endif - return payload->fanin_actual_count; - } + // on_task_release is gone under polling. It existed to bump each producer's + // fanout_refcount so the host wait_for_consumers could observe consumer + // retirement; that signal is now the per-ring completed_watermark advanced by + // on_mixed_task_complete. There is likewise no self CONSUMED flip (host-orch + // never reclaimed slots on device). // === Cold-path API (defined in pto_scheduler.cpp) === @@ -1183,7 +895,7 @@ struct PTO2SchedulerState { // per-ring dep_pool entries) on the supplied arena. // Capacities are baked into the returned layout; init_data_from_layout uses // the same values. - static PTO2SchedulerLayout reserve_layout(DeviceArena &arena, int32_t dep_pool_capacity = PTO2_DEP_LIST_POOL_SIZE); + static PTO2SchedulerLayout reserve_layout(DeviceArena &arena); // Phase 3a: write everything *except* arena-internal pointer fields. // `sm_dev_base` is the device address of the SM (only stored, never @@ -1213,10 +925,8 @@ struct PTO2SchedulerState { // because PTO2SchedulerState's on_task_complete signature is only known // after its full definition above. // -// When the deferred_release_slot_states[] buffer is full, drain it via -// on_task_release before appending — mirrors the same overflow-drain idiom -// that scheduler_completion.cpp's inline NotDeferred path uses, so high task -// rates don't surface as ASYNC_WAIT_OVERFLOW errors. +// Polling: on_task_complete publishes completion + drains the wake list inline, +// so the async-drain path no longer buffers producer releases. inline bool AsyncWaitList::try_inline_complete_locked(AsyncWaitList::DrainCompletionSink &sink, PTO2TaskSlotState &slot_state) { // Return value (CompletionStats / consumer-walk count) discarded: @@ -1226,26 +936,13 @@ AsyncWaitList::try_inline_complete_locked(AsyncWaitList::DrainCompletionSink &si #else (void)sink.sched->on_task_complete(slot_state); #endif - if (*sink.deferred_release_count >= sink.deferred_release_capacity) { - while (*sink.deferred_release_count > 0) { -#if SIMPLER_SCHED_PROFILING - (void)sink.sched->on_task_release( - *sink.deferred_release_slot_states[--(*sink.deferred_release_count)], sink.thread_idx - ); -#else - sink.sched->on_task_release(*sink.deferred_release_slot_states[--(*sink.deferred_release_count)]); -#endif - } - } - sink.deferred_release_slot_states[(*sink.deferred_release_count)++] = &slot_state; sink.inline_completed++; return true; } template inline AsyncPollResult AsyncWaitList::poll_and_complete( - AICoreCompletionMailbox *aicore_mailbox, PTO2SchedulerState *sched, - PTO2TaskSlotState **deferred_release_slot_states, int32_t &deferred_release_count, int32_t deferred_release_capacity + AICoreCompletionMailbox *aicore_mailbox, PTO2SchedulerState *sched #if SIMPLER_SCHED_PROFILING , int thread_idx @@ -1256,9 +953,6 @@ inline AsyncPollResult AsyncWaitList::poll_and_complete( AsyncWaitList::DrainCompletionSink sink{}; sink.sched = sched; - sink.deferred_release_slot_states = deferred_release_slot_states; - sink.deferred_release_count = &deferred_release_count; - sink.deferred_release_capacity = deferred_release_capacity; #if SIMPLER_SCHED_PROFILING sink.thread_idx = thread_idx; #endif @@ -1307,21 +1001,7 @@ inline AsyncPollResult AsyncWaitList::poll_and_complete( #else (void)sched->on_task_complete(*entry.slot_state); #endif - // Drain deferred_release in place when the buffer fills — same - // overflow-drain idiom used by complete_slot_task's inline path - // and by try_inline_complete_locked. Without this, large bursts - // of completable wait_list entries in a single poll surfaced as - // ASYNC_WAIT_OVERFLOW under the MPSC model. - if (deferred_release_count >= deferred_release_capacity) { - while (deferred_release_count > 0) { -#if SIMPLER_SCHED_PROFILING - (void)sched->on_task_release(*deferred_release_slot_states[--deferred_release_count], thread_idx); -#else - sched->on_task_release(*deferred_release_slot_states[--deferred_release_count]); -#endif - } - } - deferred_release_slot_states[deferred_release_count++] = entry.slot_state; + // Polling: completion is fully published inline; no deferred release. result.completed++; int32_t last = count - 1; @@ -1347,7 +1027,7 @@ struct PTO2SchedProfilingData { uint64_t self_consumed_cycle; // self check_and_handle_consumed // Wait times - uint64_t lock_wait_cycle; // spin-wait in fanout_lock + uint64_t lock_wait_cycle; // Legacy (wiring): fanout_lock spin-wait; polling has no such lock uint64_t push_wait_cycle; // CAS contention in push() uint64_t pop_wait_cycle; // CAS contention in pop() 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 25afe53af..c7bc14c22 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 @@ -233,8 +233,18 @@ void SchedulerContext::log_stall_diagnostics( for (int32_t si = 0; si < ring_task_count; si++) { PTO2TaskSlotState &slot_state = ring.get_slot_state_by_task_id(si); PTO2TaskState st = slot_state.task_state.load(std::memory_order_relaxed); - int32_t rc = slot_state.fanin_refcount.load(std::memory_order_relaxed); - int32_t fi = slot_state.fanin_count; + // Polling: no fanin_refcount. Recompute met/total from the inline + // fanin ids vs the ring completion_flags (rc = satisfied producers, + // fi = raw producer count) so the stall dump still shows readiness. + int32_t fi = slot_state.payload != nullptr ? slot_state.payload->fanin_count : 0; + int32_t rc = 0; + if (slot_state.payload != nullptr) { + for (int32_t k = 0; k < fi; k++) { + int32_t pid = slot_state.payload->fanin_local_ids[k]; + if (ring.completion_flags[pid & ring.task_window_mask].load(std::memory_order_relaxed) != 0) + rc++; + } + } int32_t kid_aic = slot_state.task->kernel_id[0]; int32_t kid_aiv0 = slot_state.task->kernel_id[1]; int32_t kid_aiv1 = slot_state.task->kernel_id[2]; @@ -266,7 +276,7 @@ void SchedulerContext::log_stall_diagnostics( if (cnt_running > STALL_DUMP_READY_MAX) continue; LOG_INFO_V9( "[STALL thread=%d idle_iterations=%d] TASK ring=%d task_id=%" PRId64 - " state=RUNNING fanin_refcount=%d/%d kernels=[aic:%d aiv0:%d aiv1:%d] " + " state=RUNNING fanin_met=%d/%d kernels=[aic:%d aiv0:%d aiv1:%d] " "running_on=[owner_thread=%d cores=[%s]]", thread_idx, idle_iterations, r, task_id, rc, fi, kid_aic, kid_aiv0, kid_aiv1, owner, running_on ); @@ -277,7 +287,7 @@ void SchedulerContext::log_stall_diagnostics( if (cnt_ready > STALL_DUMP_READY_MAX) continue; LOG_INFO_V9( "[STALL thread=%d idle_iterations=%d] TASK ring=%d task_id=%" PRId64 - " state=READY fanin_refcount=%d/%d kernels=[aic:%d aiv0:%d aiv1:%d]", + " state=READY fanin_met=%d/%d kernels=[aic:%d aiv0:%d aiv1:%d]", thread_idx, idle_iterations, r, task_id, rc, fi, kid_aic, kid_aiv0, kid_aiv1 ); continue; @@ -286,7 +296,7 @@ void SchedulerContext::log_stall_diagnostics( if (cnt_waiting > STALL_DUMP_WAIT_MAX) continue; LOG_INFO_V9( "[STALL thread=%d idle_iterations=%d] TASK ring=%d task_id=%" PRId64 - " state=WAIT fanin_refcount=%d/%d kernels=[aic:%d aiv0:%d aiv1:%d] missing_deps=%d", + " state=WAIT fanin_met=%d/%d kernels=[aic:%d aiv0:%d aiv1:%d] missing_deps=%d", thread_idx, idle_iterations, r, task_id, rc, fi, kid_aic, kid_aiv0, kid_aiv1, fi - rc ); } @@ -1068,6 +1078,34 @@ void SchedulerContext::on_orchestration_done( } } + // Polling initial classify (device boot): the host built the whole graph and + // no producer has executed yet — every completion_flags byte is 0 except the + // hidden-alloc tasks the host completed inline (pre-set to 1). Classify each + // submitted task exactly once: route roots (all fanin met) to the ready queues + // and register the rest on their first unmet producer's wake list. This + // replaces the host-side wiring drain the wiring model deferred to a device + // queue. Runs on the boot leader BEFORE the caller publishes + // runtime_init_ready_ (release), so it is race-free against the scheduler + // threads (they acquire that store before dispatching) and nothing completes + // during the scan. + if (orch_err == PTO2_ERROR_NONE && sched_->ring_sched_state.ring != nullptr) { + PTO2SharedMemoryRingHeader &ring = *sched_->ring_sched_state.ring; + const int32_t submitted = ring.fc.current_task_index.load(std::memory_order_acquire); + for (int32_t id = 0; id < submitted; id++) { + if (ring.completion_flags[id & ring.task_window_mask].load(std::memory_order_acquire) != 0) { + 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); + 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); + } + } + } + #if SIMPLER_DFX // Write the core-to-thread mapping so the profiling data reflects the // scheduler threads' final core distribution. 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 dc39400a0..ba5e8de69 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 @@ -31,9 +31,7 @@ // Dual-slot state machine helpers // ============================================================================= -namespace { -inline constexpr int32_t PTO2_DEFERRED_RELEASE_CAP = 256; -} +namespace {} // Pure function: read register result -> SlotTransition (no side effects). SlotTransition SchedulerContext::decide_slot_transition( @@ -84,8 +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, int32_t &completed_this_turn, - PTO2TaskSlotState *deferred_release_slot_states[], int32_t &deferred_release_count + [[maybe_unused]] 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 @@ -225,21 +222,8 @@ void SchedulerContext::complete_slot_task( } l2_swimlane.phase_complete_count++; #endif - if (deferred_release_count < PTO2_DEFERRED_RELEASE_CAP) { - deferred_release_slot_states[deferred_release_count++] = &slot_state; - } else { - LOG_INFO_V9("Thread %d: release", thread_idx); - while (deferred_release_count > 0) { -#if SIMPLER_SCHED_PROFILING - // SCHED_PROFILING variant takes thread_idx for the per-thread - // atomic counter side-effects. The return value is unused. - (void)sched_->on_task_release(*deferred_release_slot_states[--deferred_release_count], thread_idx); -#else - sched_->on_task_release(*deferred_release_slot_states[--deferred_release_count]); -#endif - } - deferred_release_slot_states[deferred_release_count++] = &slot_state; - } + // Polling: on_task_complete published completion + drained the wake list + // inline; no deferred producer-release step. completed_this_turn++; } @@ -298,7 +282,7 @@ void SchedulerContext::clear_running_slot(CoreExecState &core) { void SchedulerContext::check_running_cores_for_completion( int32_t thread_idx, Handshake *hank, int32_t &completed_this_turn, int32_t &cur_thread_completed, - bool &made_progress, PTO2TaskSlotState *deferred_release_slot_states[], int32_t &deferred_release_count + bool &made_progress ) { #if SIMPLER_SCHED_PROFILING auto &l2_swimlane = sched_l2_swimlane_[thread_idx]; @@ -404,7 +388,7 @@ void SchedulerContext::check_running_cores_for_completion( } complete_slot_task( *core.pending_slot_state, core.pending_reg_task_id, core.pending_subslot, thread_idx, core_id, hank, - completed_this_turn, deferred_release_slot_states, deferred_release_count + completed_this_turn #if SIMPLER_DFX , core.pending_dispatch_timestamp, finish_ts @@ -418,7 +402,7 @@ void SchedulerContext::check_running_cores_for_completion( } complete_slot_task( *core.running_slot_state, core.running_reg_task_id, core.running_subslot, thread_idx, core_id, hank, - completed_this_turn, deferred_release_slot_states, deferred_release_count + completed_this_turn #if SIMPLER_DFX , core.running_dispatch_timestamp, finish_ts 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 7ff03fa04..6c5e3200b 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 @@ -404,8 +404,7 @@ class SchedulerContext { void complete_slot_task( PTO2TaskSlotState &slot_state, int32_t expected_reg_task_id, PTO2SubtaskSlot subslot, int32_t thread_idx, - int32_t core_id, Handshake *hank, int32_t &completed_this_turn, - PTO2TaskSlotState *deferred_release_slot_states[], int32_t &deferred_release_count + int32_t core_id, Handshake *hank, int32_t &completed_this_turn #if SIMPLER_DFX , uint64_t dispatch_ts, uint64_t finish_ts @@ -417,7 +416,7 @@ class SchedulerContext { void check_running_cores_for_completion( int32_t thread_idx, Handshake *hank, int32_t &completed_this_turn, int32_t &cur_thread_completed, - bool &made_progress, PTO2TaskSlotState *deferred_release_slot_states[], int32_t &deferred_release_count + bool &made_progress ); bool enter_drain_mode(PTO2TaskSlotState *slot_state, int32_t block_num); 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 a3353a1ab..b86000628 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 @@ -50,9 +50,7 @@ static_assert(sizeof(Tensor) == PTO2_TASKPAYLOAD_TENSOR_STRIDE); // Dispatch helpers // ============================================================================= -namespace { -inline constexpr int32_t PTO2_DEFERRED_RELEASE_CAP = 256; -} +namespace {} // The early-dispatch core bitmask (PTO2_EARLY_DISPATCH_CORE_MASK_WORDS * 64 bits) must cover // every global core_id, and the per-core doorbell table is sized to match. @@ -896,9 +894,6 @@ int32_t SchedulerContext::resolve_and_dispatch(Runtime *runtime, int32_t thread_ l2_swimlane.l2_swimlane_enabled = (l2_swimlane_level_ != L2SwimlaneLevel::DISABLED); #endif - PTO2TaskSlotState *deferred_release_slot_states[PTO2_DEFERRED_RELEASE_CAP]; - int32_t deferred_release_count = 0; - // PMU runs require single-issue dispatch — overlapping in-flight tasks // pollute per-task PMU counters, so skip the PENDING pre-load phase. // Cached at function scope: is_pmu_enabled() is extern "C" and the @@ -1025,8 +1020,7 @@ int32_t SchedulerContext::resolve_and_dispatch(Runtime *runtime, int32_t thread_ bool try_completed = tracker.has_any_running_cores(); if (try_completed) { check_running_cores_for_completion( - thread_idx, hank, completed_this_turn, cur_thread_completed, made_progress, - deferred_release_slot_states, deferred_release_count + thread_idx, hank, completed_this_turn, cur_thread_completed, made_progress ); } if (completed_this_turn > 0) { @@ -1050,8 +1044,7 @@ int32_t SchedulerContext::resolve_and_dispatch(Runtime *runtime, int32_t thread_ 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_, deferred_release_slot_states, deferred_release_count, - PTO2_DEFERRED_RELEASE_CAP + rt_->aicore_mailbox, sched_ #if SIMPLER_SCHED_PROFILING , thread_idx @@ -1188,21 +1181,10 @@ int32_t SchedulerContext::resolve_and_dispatch(Runtime *runtime, int32_t thread_ ); } #endif - // Dummy tasks have no subtasks to retire and no fanout pre-conditions - // beyond their own producers; release self-reference so the slot can - // reach CONSUMED once all consumers drain. - deferred_release_slot_states[deferred_release_count++] = &dummy_slot; - if (deferred_release_count >= PTO2_DEFERRED_RELEASE_CAP) { - while (deferred_release_count > 0) { -#if SIMPLER_SCHED_PROFILING - (void)sched_->on_task_release( - *deferred_release_slot_states[--deferred_release_count], thread_idx - ); -#else - sched_->on_task_release(*deferred_release_slot_states[--deferred_release_count]); -#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++; @@ -1322,33 +1304,7 @@ int32_t SchedulerContext::resolve_and_dispatch(Runtime *runtime, int32_t thread_ idle_iterations = 0; last_progress_ts = get_sys_cnt_aicpu(); } else { -#if SIMPLER_DFX - uint64_t rel_t0 = (l2_swimlane_level_ >= L2SwimlaneLevel::SCHED_PHASES && deferred_release_count > 0) ? - get_sys_cnt_aicpu() : - 0; - // Snapshot the slot count before the drain loop decrements it to 0, - // so the Release bar can report how many slots it drained. - uint32_t released_count = static_cast(deferred_release_count); -#endif - while (deferred_release_count > 0) { -#if SIMPLER_SCHED_PROFILING - (void)sched_->on_task_release(*deferred_release_slot_states[--deferred_release_count], thread_idx); -#else - sched_->on_task_release(*deferred_release_slot_states[--deferred_release_count]); -#endif - } -#if SIMPLER_DFX - // Release is a distinct operation from the poll scan — emit it with - // its own span (Perfetto nests it inside the surrounding poll/idle - // run by time-containment) rather than competing with poll for one - // per-iteration label. - if (rel_t0 != 0) { - l2_swimlane_aicpu_record_sched_phase( - thread_idx, L2SwimlaneSchedPhaseKind::Release, rel_t0, get_sys_cnt_aicpu(), - l2_swimlane.sched_loop_count, released_count - ); - } -#endif + // Polling: no deferred producer-release phase to drain on an idle pass. idle_iterations++; if (idle_iterations % FATAL_ERROR_CHECK_INTERVAL == 0) { @@ -1412,18 +1368,7 @@ int32_t SchedulerContext::resolve_and_dispatch(Runtime *runtime, int32_t thread_ } } - // Drain any entries left in the deferred-release batch. The in-loop flush - // only fires on idle iterations and on buffer-full; a loop exit while the - // last iteration made progress can leave entries un-released. Drop them - // here so every consumed producer slot completes its on_task_release - // regardless of which loop-exit path fired. - while (deferred_release_count > 0) { -#if SIMPLER_SCHED_PROFILING - (void)sched_->on_task_release(*deferred_release_slot_states[--deferred_release_count], thread_idx); -#else - sched_->on_task_release(*deferred_release_slot_states[--deferred_release_count]); -#endif - } + // Polling: no deferred producer-release batch to drain at loop exit. #if SIMPLER_DFX // Final-drain: emit any pop_hit / pop_miss accrued since the last 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 9983b2bab..61e61e515 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 @@ -91,10 +91,6 @@ bool PTO2SchedulerState::RingSchedState::init_data_from_layout(void *sm_dev_base ring = pto2_sm_layout::ring_header_addr(sm_dev_base); last_task_alive = 0; advance_lock.store(0, std::memory_order_relaxed); -#if SIMPLER_DFX - dep_pool_snapshot_tail.store(1, std::memory_order_relaxed); - dep_pool_snapshot_top.store(1, std::memory_order_relaxed); -#endif // Per-slot SM-side initialization (reset_for_reuse + fanin_count/active_mask // zero) lives in PTO2SharedMemoryHandle::init_header_per_ring so the AICPU @@ -105,10 +101,9 @@ bool PTO2SchedulerState::RingSchedState::init_data_from_layout(void *sm_dev_base void PTO2SchedulerState::RingSchedState::destroy() { ring = nullptr; } -PTO2SchedulerLayout PTO2SchedulerState::reserve_layout(DeviceArena &arena, int32_t dep_pool_capacity) { +PTO2SchedulerLayout PTO2SchedulerState::reserve_layout(DeviceArena &arena) { PTO2SchedulerLayout layout{}; layout.ready_queue_capacity = PTO2_READY_QUEUE_SIZE; - layout.dep_pool_capacity = dep_pool_capacity; for (int i = 0; i < PTO2_NUM_RESOURCE_SHAPES; i++) { layout.off_ready_queue_slots[i] = ready_queue_reserve_layout(arena, PTO2_READY_QUEUE_SIZE); @@ -121,11 +116,8 @@ PTO2SchedulerLayout PTO2SchedulerState::reserve_layout(DeviceArena &arena, int32 layout.off_early_dispatch_queue_slots[i] = ready_queue_reserve_layout(arena, PTO2_EARLY_DISPATCH_QUEUE_SIZE); } layout.off_early_sync_start_queue_slots = ready_queue_reserve_layout(arena, PTO2_EARLY_DISPATCH_QUEUE_SIZE); - // Force a cache-line base so writes from scheduler thread 0 (sole writer of - // the dep_pool) do not invalidate adjacent multi-threaded regions like - // ready_queue.slots. - layout.off_dep_pool_entries = - arena.reserve(static_cast(dep_pool_capacity) * sizeof(PTO2DepListEntry), PTO2_ALIGN_SIZE); + // Polling: no dep_pool arena region — producer dependencies are inline ids on + // the payload and readiness is via completion_flags. return layout; } @@ -177,11 +169,7 @@ bool PTO2SchedulerState::init_data_from_layout( return false; } - auto *orch_err = pto2_sm_layout::orch_error_code_addr(sm_dev_base); - auto *dep_entries = static_cast(arena.region_ptr(layout.off_dep_pool_entries)); - memset(dep_entries, 0, static_cast(layout.dep_pool_capacity) * sizeof(PTO2DepListEntry)); - sched->ring_sched_state.dep_pool.init(dep_entries, layout.dep_pool_capacity, orch_err); - + // Polling: no dep_pool arena region to initialize. return true; } @@ -200,14 +188,11 @@ void PTO2SchedulerState::wire_arena_pointers(const PTO2SchedulerLayout &layout, ); } ready_queue_wire_arena_pointers(&sched->early_sync_start_queue, arena, layout.off_early_sync_start_queue_slots); - sched->ring_sched_state.dep_pool.base = - static_cast(arena.region_ptr(layout.off_dep_pool_entries)); } void PTO2SchedulerState::destroy() { PTO2SchedulerState *sched = this; sched->ring_sched_state.destroy(); - sched->ring_sched_state.dep_pool.base = nullptr; for (int i = 0; i < PTO2_NUM_RESOURCE_SHAPES; i++) { ready_queue_destroy(&sched->ready_queues[i]); } @@ -225,8 +210,7 @@ void PTO2SchedulerState::destroy() { // Orchestrator // ============================================================================= -PTO2OrchestratorLayout -PTO2OrchestratorState::reserve_layout(DeviceArena &arena, int32_t task_window_size, int32_t dep_pool_capacity) { +PTO2OrchestratorLayout PTO2OrchestratorState::reserve_layout(DeviceArena &arena, int32_t task_window_size) { PTO2OrchestratorLayout layout{}; // scope_tasks holds every task in the open scope, so its cap is the real // in-flight budget = the (runtime) task window. Using the compile-time @@ -236,12 +220,8 @@ PTO2OrchestratorState::reserve_layout(DeviceArena &arena, int32_t task_window_si always_assert(task_window_size > 0); layout.scope_tasks_cap = task_window_size; layout.scope_stack_capacity = PTO2_MAX_SCOPE_DEPTH; - layout.dep_pool_capacity = dep_pool_capacity; - - const size_t fanin_pool_bytes = - PTO2_ALIGN_UP(static_cast(dep_pool_capacity) * sizeof(PTO2FaninSpillEntry), PTO2_ALIGN_SIZE); - layout.off_fanin_pool = arena.reserve(fanin_pool_bytes, PTO2_ALIGN_SIZE); + // Polling: no fanin-spill pool — producer ids are inline on the payload. always_assert(task_window_size > 0 && (task_window_size & (task_window_size - 1)) == 0); const size_t seen_epoch_bytes = PTO2_ALIGN_UP(static_cast(task_window_size) * sizeof(uint32_t), PTO2_ALIGN_SIZE); @@ -278,12 +258,6 @@ bool PTO2OrchestratorState::init_data_from_layout( orch_err, slot_states_dev ); - const size_t fanin_pool_bytes = - PTO2_ALIGN_UP(static_cast(layout.dep_pool_capacity) * sizeof(PTO2FaninSpillEntry), PTO2_ALIGN_SIZE); - auto *fanin_entries = static_cast(arena.region_ptr(layout.off_fanin_pool)); - memset(fanin_entries, 0, fanin_pool_bytes); - orch->ring.fanin_pool.init(fanin_entries, layout.dep_pool_capacity, orch_err); - const size_t seen_epoch_bytes = PTO2_ALIGN_UP(static_cast(layout.tensor_map.task_window_size) * sizeof(uint32_t), PTO2_ALIGN_SIZE); auto *seen_epoch = static_cast(arena.region_ptr(layout.off_fanin_seen_epoch)); @@ -307,7 +281,6 @@ void PTO2OrchestratorState::wire_arena_pointers( const PTO2OrchestratorLayout &layout, DeviceArena &arena, PTO2SchedulerState *scheduler_arg ) { auto *orch = this; - orch->ring.fanin_pool.base = static_cast(arena.region_ptr(layout.off_fanin_pool)); orch->fanin_seen_epoch = static_cast(arena.region_ptr(layout.off_fanin_seen_epoch)); orch->tensor_map.wire_arena_pointers(layout.tensor_map, arena); orch->scope_tasks = static_cast(arena.region_ptr(layout.off_scope_tasks)); @@ -318,7 +291,6 @@ void PTO2OrchestratorState::wire_arena_pointers( void PTO2OrchestratorState::destroy() { auto *orch = this; orch->tensor_map.destroy(); - orch->ring.fanin_pool.base = nullptr; orch->fanin_seen_epoch = nullptr; orch->scope_tasks = nullptr; orch->scope_begins = nullptr; @@ -330,36 +302,30 @@ void PTO2OrchestratorState::set_scheduler(PTO2SchedulerState *scheduler) { this- // Top-level runtime arena // ============================================================================= -PTO2RuntimeArenaLayout -runtime_reserve_layout(DeviceArena &arena, uint64_t task_window_size, int32_t dep_pool_capacity) { +PTO2RuntimeArenaLayout runtime_reserve_layout(DeviceArena &arena, uint64_t task_window_size) { uint64_t task_window_sizes[PTO2_MAX_RING_DEPTH]; uint64_t heap_sizes[PTO2_MAX_RING_DEPTH]; - int32_t dep_pool_capacities[PTO2_MAX_RING_DEPTH]; for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { task_window_sizes[r] = task_window_size; heap_sizes[r] = 0; - dep_pool_capacities[r] = dep_pool_capacity; } - return runtime_reserve_layout(arena, task_window_sizes, heap_sizes, dep_pool_capacities); + return runtime_reserve_layout(arena, task_window_sizes, heap_sizes); } PTO2RuntimeArenaLayout runtime_reserve_layout( DeviceArena &arena, const uint64_t task_window_sizes[PTO2_MAX_RING_DEPTH], - const uint64_t heap_sizes[PTO2_MAX_RING_DEPTH], const int32_t dep_pool_capacities[PTO2_MAX_RING_DEPTH] + const uint64_t heap_sizes[PTO2_MAX_RING_DEPTH] ) { PTO2RuntimeArenaLayout layout{}; for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { layout.task_window_sizes[r] = task_window_sizes[r]; layout.heap_sizes[r] = heap_sizes[r]; - layout.dep_pool_capacities[r] = dep_pool_capacities[r]; } layout.off_sm_handle = arena.reserve(sizeof(PTO2SharedMemoryHandle), alignof(PTO2SharedMemoryHandle)); - layout.orch = PTO2OrchestratorState::reserve_layout( - arena, static_cast(task_window_sizes[0]), dep_pool_capacities[0] - ); - layout.sched = PTO2SchedulerState::reserve_layout(arena, dep_pool_capacities[0]); + layout.orch = PTO2OrchestratorState::reserve_layout(arena, static_cast(task_window_sizes[0])); + layout.sched = PTO2SchedulerState::reserve_layout(arena); layout.off_runtime = arena.reserve(sizeof(PTO2Runtime), PTO2_ALIGN_SIZE); layout.off_mailbox = arena.reserve(sizeof(AICoreCompletionMailbox), alignof(AICoreCompletionMailbox)); 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 a0798a65e..9b33075a7 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,6 +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); } void PTO2SharedMemoryHandle::setup_pointers(uint64_t task_window_size) { @@ -153,6 +154,10 @@ void PTO2SharedMemoryHandle::init_header_per_ring( // 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); + header->orchestrator_done.store(0, std::memory_order_relaxed); // Ring layout info @@ -182,9 +187,13 @@ void PTO2SharedMemoryHandle::init_header_per_ring( auto &ring = header->ring; for (uint64_t i = 0; i < task_window_sizes[0]; i++) { ring.slot_states[i].reset_for_reuse(); - ring.slot_states[i].fanin_count = 0; 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)); } // =============================================================================