From 247a38507cb2c9122efeb9c78189bf2755449c0b Mon Sep 17 00:00:00 2001 From: Chao Wang <26245345+ChaoWao@users.noreply.github.com> Date: Tue, 21 Jul 2026 07:09:49 -0700 Subject: [PATCH] perf(tmr): replace wiring with polling-based task readiness (~17% median device speedup) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace the tensormap_and_ringbuffer scheduler's fanin wiring (per-edge dep_pool + fanin_refcount/fanout_refcount atomics) with a polling completion design: a per-slot completion_flags byte array, a monotonic per-ring completed_watermark, and an intrusive wake_list for last-fanin notification. The orchestrator no longer wires every producer→consumer edge into a shared dep pool; consumers scan their inline fanin and either observe completion via the watermark or self-register on the one unmet producer's wake list. Fanin edges are stored inline in the payload (fanin_local_ids/fanin_ring_ids, PTO2_MAX_FANIN=128); heap reclamation is gated on last_consumer_local_id vs the per-ring watermark. Early-dispatch, predicated dispatch, and sync_start are preserved. Adopts #1345's barrier-free, batched per-thread scheduler init. Cold-path init, dispatch, cold-path, and completion logic are implemented in .cpp translation units (scheduler_{dispatch,cold_path,completion}.cpp, pto_orchestrator.cpp, shared/pto_*.cpp); only the per-task/per-completion hot paths stay inline in headers. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../aicpu/aicpu_executor.cpp | 86 +- .../common/intrinsic.h | 4 +- .../docs/MULTI_RING.md | 93 +- .../docs/RUNTIME_LOGIC.md | 22 +- .../docs/SCALAR_DATA_ACCESS.md | 2 +- .../docs/device_log_profiling.md | 4 +- .../docs/profiling_levels.md | 6 +- .../host/dep_gen_replay.cpp | 2 +- .../orchestration/common.cpp | 7 + .../runtime/aicore_completion_mailbox.h | 42 - .../runtime/aicore_completion_mailbox_types.h | 17 - .../backend/sdma/sdma_completion_kernel.h | 29 +- .../backend/sdma/sdma_completion_scheduler.h | 7 +- .../runtime/pto_async_kernel_api.h | 34 +- .../runtime/pto_async_wait.h | 78 +- .../runtime/pto_completion_token.h | 6 - .../runtime/pto_dep_compute.h | 113 +- .../runtime/pto_orchestrator.cpp | 1178 ++---------- .../runtime/pto_orchestrator.h | 171 +- .../runtime/pto_ring_buffer.cpp | 229 +-- .../runtime/pto_ring_buffer.h | 675 +------ .../runtime/pto_runtime2.cpp | 214 +-- .../runtime/pto_runtime2.h | 356 ++-- .../runtime/pto_runtime2_types.h | 452 ++--- .../runtime/pto_shared_memory.h | 188 +- .../runtime/pto_submit_types.h | 32 - .../runtime/pto_tensormap.h | 358 +--- .../runtime/scheduler/pto_scheduler.cpp | 103 +- .../runtime/scheduler/pto_scheduler.h | 1585 +++++----------- .../runtime/scheduler/scheduler_cold_path.cpp | 1607 ++++------------- .../scheduler/scheduler_completion.cpp | 735 ++------ .../runtime/scheduler/scheduler_context.h | 588 +++--- .../runtime/scheduler/scheduler_dispatch.cpp | 1545 +++------------- .../runtime/scheduler/scheduler_types.h | 179 +- .../runtime/shared/pto_runtime2_init.cpp | 668 ++----- .../runtime/shared/pto_shared_memory.cpp | 189 +- .../runtime/shared/pto_tensormap.cpp | 180 +- 37 files changed, 2462 insertions(+), 9322 deletions(-) diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/aicpu/aicpu_executor.cpp b/src/a2a3/runtime/tensormap_and_ringbuffer/aicpu/aicpu_executor.cpp index 1fef4668b9..6d251fa6a9 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/aicpu/aicpu_executor.cpp +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/aicpu/aicpu_executor.cpp @@ -257,74 +257,44 @@ int32_t AicpuExecutor::init(Runtime *runtime) { if (init_failed_.load(std::memory_order_acquire)) return -1; } - // The orchestrator (top thread, tidx == nthreads-1) does not dispatch to - // cores, so it skips the handshake entirely and returns to - // run() immediately to build the graph — overlapping the schedulers' handshake. - // Core counts it needs are derived from cores_total_num_ (fixed 1:2 ratio) in - // run(). The remaining (nthreads-1) threads re-partition ALL cores among - // themselves, so every core still gets its register window opened. - const bool decouple_orch = (nthreads > 1) && !serial_orch_sched_; - const bool is_orchestrator = (tidx == nthreads - 1); - if (decouple_orch && is_orchestrator) { - return 0; // do NOT touch the handshake or hs_arrived_ barrier - } - const int32_t hs_nthreads = decouple_orch ? (nthreads - 1) : nthreads; - - // Barrier-free scheduler init (the decoupled default). Each scheduler thread - // handshakes exactly the clusters it will dispatch to (blocked-layout - // ownership: cluster ci = {ci, N/3+2ci, N/3+2ci+1}, owned by ci % hs_nthreads) - // and self-assigns them, then returns straight to run(). With no all-thread - // barrier a thread starts dispatching to its own cores as soon as they come - // up, independent of peers still handshaking. hs_nthreads == active_sched_threads_ - // in this branch, so handshake ownership matches assign_own_clusters'. - if (decouple_orch) { - sched_ctx_.handshake_owned_clusters(runtime, tidx, hs_nthreads); - sched_ctx_.assign_own_clusters(tidx); -#if SIMPLER_DFX - // Profiling subsystems (pmu/dump/dep) need every core's physical_core_id, - // so gate their one-time leader init behind a barrier — DFX builds only. - hs_arrived_.fetch_add(1, std::memory_order_acq_rel); - if (is_leader) { - while (hs_arrived_.load(std::memory_order_acquire) < hs_nthreads) {} + // Multi-thread runs take the barrier-free init path: the orchestrator (last + // thread) skips the handshake and returns to build the graph immediately, and + // each scheduler thread handshakes + self-assigns only the clusters it will + // dispatch to (ci % sched_thread_num_ == tidx) — no all-thread barrier and no + // leader-serialized post_handshake_init. DFX-safe: polling inits PMU / dep-gen + // per-thread inside the dispatch loop (no global pmu_aicpu_init over all cores), + // so unlike #1345 upstream there is no leader-only profiling init needing a + // barrier. Single-thread runs keep the original path (no orch/scheduler split). + const bool barrier_free = (nthreads >= 2); + if (barrier_free) { + const bool is_orchestrator = (tidx == sched_thread_num_); + if (!is_orchestrator) { + sched_ctx_.handshake_owned_clusters(runtime, tidx, sched_thread_num_); + sched_ctx_.assign_own_clusters(tidx); if (sched_ctx_.handshake_failed()) { sched_ctx_.abort_and_shutdown(runtime); init_failed_.store(true, std::memory_order_release); init_done_.store(true, std::memory_order_release); return -1; } - sched_ctx_.post_handshake_profiling_init(); - init_done_.store(true, std::memory_order_release); - } else { - while (!init_done_.load(std::memory_order_acquire)) { - if (init_failed_.load(std::memory_order_acquire)) return -1; - } - if (init_failed_.load(std::memory_order_acquire)) return -1; } -#else - // Perf path: a scheduler that sees an invalid core report (its own or a - // peer's, observed so far) latches completed_ via abort_and_shutdown, which - // stops any peer still entering dispatch (run()'s is_completed() gate). A - // peer that already passed that gate is not joined here — its own cores are - // valid (it handshaked them), and the failure ends in the host device reset - // that reaps every core, so the residual overlap is bounded and - // non-corrupting. finished_count_ is reset per-run in deinit(), not here. - if (sched_ctx_.handshake_failed()) { - sched_ctx_.abort_and_shutdown(runtime); - init_failed_.store(true, std::memory_order_release); - return -1; + // No thread blocks: schedulers dispatch their own cores as soon as those + // cores are up; the orchestrator overlaps graph-build with peer handshakes. + if (is_leader) { + finished_count_.store(0, std::memory_order_release); + init_done_.store(true, std::memory_order_release); + LOG_INFO_V0("AicpuExecutor: Init complete (barrier-free)"); } -#endif return 0; } - // Serial / single-thread fallback: contiguous handshake + all-thread barrier + - // leader post_handshake_init (the orchestrator participates as a handshaker, - // and hs_nthreads != active_sched_threads_ makes per-thread self-assignment - // unsafe here). - sched_ctx_.handshake_partition(runtime, tidx, hs_nthreads); + // All threads: handshake this thread's slice of cores in parallel. + sched_ctx_.handshake_partition(runtime, tidx, nthreads); + + // Barrier: leader waits for every slice to finish, then completes init. hs_arrived_.fetch_add(1, std::memory_order_acq_rel); if (is_leader) { - while (hs_arrived_.load(std::memory_order_acquire) < hs_nthreads) {} + while (hs_arrived_.load(std::memory_order_acquire) < nthreads) {} finished_count_.store(0, std::memory_order_release); if (sched_ctx_.post_handshake_init(runtime) != 0) { init_failed_.store(true, std::memory_order_release); @@ -660,11 +630,7 @@ int32_t AicpuExecutor::run(Runtime *runtime) { // Fill ops / core counts (host can't resolve s_runtime_ops's // device address nor know the SchedulerContext's core fan-out). - // Core counts come from cores_total_num_ (fixed 1 AIC : 2 AIV - // cluster ratio). On the decoupled path aic_count()/aiv_count() are - // not populated until after this SM reset, so they cannot be read here. - int32_t spike_total = sched_ctx_.cores_total_num(); - runtime_finalize_after_wire(rt, spike_total / 3, (spike_total * 2) / 3); + runtime_finalize_after_wire(rt, sched_ctx_.aic_count(), sched_ctx_.aiv_count()); #if SIMPLER_DFX rt->orchestrator.l2_swimlane_level = get_l2_swimlane_level(); { diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/common/intrinsic.h b/src/a2a3/runtime/tensormap_and_ringbuffer/common/intrinsic.h index 768e6a6120..ba83a8b5c6 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/common/intrinsic.h +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/common/intrinsic.h @@ -63,7 +63,7 @@ * compiled, ran without error, and produced wrong output. Use * `get_sub_block_id(args)` instead, which reads from the runtime's * `GlobalContext.sub_block_id` that the scheduler initializes per - * AIV core in `scheduler_cold_path.cpp::SchedulerContext::init`. + * AIV core in `scheduler_context.h::SchedulerContext::init`. * * - `get_block_idx()` and `get_block_num()` are not redirected to * simpler's LocalContext either — use the `(args)` variants below @@ -97,7 +97,7 @@ static constexpr int32_t PTO2_EXT_PARAMS_COUNT = 2; /** * Args[] suffix indices for context pointers. - * Derived from MAX_TENSOR_ARGS(32) + MAX_SCALAR_ARGS(16). + * Derived from MAX_TENSOR_ARGS(16) + MAX_SCALAR_ARGS(32). * Users should not depend on these values; use the Get* functions below. */ static constexpr int32_t SPMD_LOCAL_CONTEXT_INDEX = 48; diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/docs/MULTI_RING.md b/src/a2a3/runtime/tensormap_and_ringbuffer/docs/MULTI_RING.md index 1f97121a89..9d8b520ba7 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/docs/MULTI_RING.md +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/docs/MULTI_RING.md @@ -178,8 +178,9 @@ Each ring's `last_task_alive` advances independently: ```text advance_ring_pointers(ring_id): // protected by per-ring advance_lock - la = ring->fc.last_task_alive - while ring->get_slot_state_by_task_id(la).task_state >= CONSUMED: + watermark = ring->completed_watermark + la = last_task_alive + while la <= watermark and watermark >= slot[la].last_consumer_local_id: reset slot for reuse la++ sync_to_sm() // release-store last_task_alive @@ -234,91 +235,25 @@ AICore uses `last_reg_val` to detect new dispatches — identical values cause s | `PTO2_HEAP_SIZE` | 256 MB | 1 GB | | `PTO2_DEP_LIST_POOL_SIZE` | 16384 | 65536 | -### 7.2 Runtime Overrides - -Each ring resource (`ring_task_window` / `ring_heap` / `ring_dep_pool`) is a -single `CallConfig.runtime_env` field that accepts **either** a scalar (broadcast -to every ring) **or** a list of four per-ring values. Precedence is resolved -independently for each resource and ring: - -```text -per-ring CallConfig entry (a scalar is broadcast to every entry) - > per-ring PTO2_RING_* env value - > scalar PTO2_RING_* env value - > compile-time default -``` - -`ring_id` is the scope-depth ring selected by the runtime: - -```text -scope depth 0 -> ring 0 -scope depth 1 -> ring 1 -scope depth 2 -> ring 2 -scope depth >=3 -> ring 3 -``` +### 7.2 Runtime Environment Overrides -Per-task via `CallConfig.runtime_env` — different L2 tasks in one launch can -each carry their own sizes. Invalid values raise at submit time (`validate()`). -Assign a scalar to size every ring the same: - -```python -cfg = CallConfig() -cfg.runtime_env.ring_task_window = 128 # power of 2, >= 4 -cfg.runtime_env.ring_heap = 262144 # bytes/ring, >= 1024 -cfg.runtime_env.ring_dep_pool = 256 # 4 .. INT32_MAX -orchestrator.submit_next_level(handle, args, cfg) -``` - -Assign a four-entry list to tune the scope-depth rings independently. The list -must contain exactly four entries; use `0` for an entry that should fall through -to the next precedence tier. All `CallConfig` values are integer byte/count -values, and each field always reads back as a four-entry list. - -```python -cfg = CallConfig() -cfg.runtime_env.ring_task_window = [8192, 16384, 131072, 524288] -cfg.runtime_env.ring_heap = [ - 128 * 1024 * 1024, - 256 * 1024 * 1024, - 384 * 1024 * 1024, - 512 * 1024 * 1024, -] -cfg.runtime_env.ring_dep_pool = [4096, 8192, 16384, 32768] -orchestrator.submit_next_level(handle, args, cfg) -``` - -Scene tests set the same keys under a nested `runtime_env` block in the -per-case `config` dict — each value is a scalar or a four-entry list: - -```python -"config": { - "runtime_env": { - "ring_task_window": [8192, 16384, 131072, 524288], - "ring_heap": [134217728, 268435456, 402653184, 536870912], - "ring_dep_pool": 256, # scalar broadcasts to every ring - } -} -``` - -Process-wide env fallback accepts either one scalar value or exactly four -comma-separated per-ring values. Invalid env values are logged and ignored, then -fall through to defaults. `PTO2_RING_HEAP` values are integer bytes: +Uniform (applies to all rings): ```bash -# Uniform, old behavior: PTO2_RING_TASK_WINDOW=1024 PTO2_RING_HEAP=1048576 PTO2_RING_DEP_POOL=1024 - -# Per-ring, indexed by ring_id 0..3: -PTO2_RING_TASK_WINDOW=8192,16384,131072,524288 -PTO2_RING_HEAP=134217728,268435456,402653184,536870912 -PTO2_RING_DEP_POOL=4096,8192,16384,32768 ``` -Use `--enable-scope-stats` to confirm the effective values for a real run. The -first line of `scope_stats/scope_stats.jsonl` includes `task_window_max`, -`heap_max`, and `dep_pool_max`, indexed by `ring`. +In `kernel_config.py`: + +```python +RUNTIME_ENV = { + "PTO2_RING_TASK_WINDOW": "128", + "PTO2_RING_HEAP": "262144", + "PTO2_RING_DEP_POOL": "256", +} +``` ### 7.3 Sizing Guidelines diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md b/src/a2a3/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md index 371d75ab45..971196d6e3 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md @@ -446,7 +446,7 @@ The orchestrator completes fanout wiring before publishing a task to the ready q - Checks `task_state >= COMPLETED` (early-finished optimization) - If not completed: prepends consumer to producer's `fanout_head` via `dep_pool.prepend` - **Releases** `fanout_lock` -3. Atomically releases the +1 redundance + completed-fanin count via `fanin_refcount.fetch_add` +3. Atomically releases the +1 redundance + early_finished count via `fanin_refcount.fetch_add` 4. If all deps satisfied: pushes task to the routed ready queue Zero-fanin tasks and tasks whose claimed producers are already completed skip dep_pool entry allocation and publish directly to the routed ready queue. @@ -586,7 +586,7 @@ This is protected by a per-ring try-lock (`advance_lock`) in `RingSchedState`, e ### 8.5 SchedulerContext -All scheduler-side state and methods live in `SchedulerContext` (`runtime/scheduler/scheduler_context.h`). It is held as a `sched_ctx_` member of `AicpuExecutor`; `AicpuExecutor` is a thin wrapper that owns the lifecycle atomics and the orchestration SO handle, and delegates everything else to `SchedulerContext`. +All scheduler-side state and methods live in `SchedulerContext` (`runtime/scheduler_context.h`). It is held as a `sched_ctx_` member of `AicpuExecutor`; `AicpuExecutor` is a thin wrapper that owns the lifecycle atomics and the orchestration SO handle, and delegates everything else to `SchedulerContext`. Public surface (called from `AicpuExecutor::init/run/deinit`): @@ -600,11 +600,7 @@ Public surface (called from `AicpuExecutor::init/run/deinit`): | `deinit()` | once per run | Reset every scheduler-owned field to its post-construction default | | Read-only accessors | various | `aic_count()` / `aiv_count()` / `is_completed()` / `completed_tasks_count()` | -Private internals are split across three .cpp files by responsibility: - -- `scheduler_completion.cpp` — completion polling, drain protocol -- `scheduler_dispatch.cpp` — task dispatch loop and helpers -- `scheduler_cold_path.cpp` — exit checks, stall diagnostics, profiling, lifecycle (`pre_handshake_init` / `handshake_partition` / `post_handshake_init` / `deinit`), core management (`assign_cores_to_threads` / `emergency_shutdown`), and `on_orchestration_done` +Private internals all live inline in `scheduler_context.h`, covering completion polling, drain protocol, task dispatch loop and helpers, exit checks, stall diagnostics, profiling, lifecycle (`init/deinit`), core management (`handshake_all_cores` / `assign_cores_to_threads` / `reassign_cores_for_all_threads` / `emergency_shutdown`), and `on_orchestration_done`. `AicpuExecutor` calls neither `handshake_*`, `assign_*`, `reassign_*`, nor `emergency_shutdown` directly — they are private, invoked only by `init` and `on_orchestration_done`. @@ -683,13 +679,11 @@ reserved core slot and a launch-visible payload before a consumer can pre-occupy `next_block_idx` records reservation progress; `published_block_count` independently establishes publication and early-candidate readiness. -The producer's slot-local dispatch-propagated bit in `lifecycle_flags` and fanout snapshot are -serialized under `fanout_lock` with consumer wiring. Wiring already locks and reads the -producer's 64-byte `PTO2TaskSlotState`, so testing this bit does not read a producer payload -cache line. An edge already in the snapshot is counted by scheduler propagation; wiring seeds -an edge added after the claim and enqueues the consumer if that seed completes -`dispatch_fanin`. This gives each eligible producer-consumer edge exactly one early-candidate -contribution regardless of which side acquires the lock first. +The producer's `dispatch_propagated` claim and fanout snapshot are serialized under +`fanout_lock` with consumer wiring. An edge already in the snapshot is counted by scheduler +propagation; wiring seeds an edge added after the claim and enqueues the consumer if that seed +completes `dispatch_fanin`. This gives each eligible producer-consumer edge exactly one +early-candidate contribution regardless of which side acquires the lock first. #### MIX per-core placement diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/docs/SCALAR_DATA_ACCESS.md b/src/a2a3/runtime/tensormap_and_ringbuffer/docs/SCALAR_DATA_ACCESS.md index ef1de83b40..94cc8a569a 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/docs/SCALAR_DATA_ACCESS.md +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/docs/SCALAR_DATA_ACCESS.md @@ -32,7 +32,7 @@ addr null-check → TensorMap lookup → spin-wait producer COMPLETED → comput - **addr null-check**: `buffer.addr == 0` means unallocated — log error, return 0 - **TensorMap lookup**: find producer task by `buffer.addr` -- **spin-wait**: wait until producer `task_state >= PTO2_TASK_COMPLETED` +- **spin-wait**: wait until producer's `completion_flags[local_id & mask] == 1` - **No producer** (lookup callback never fires): skip waiting, read immediately ### 3.2 set_tensor_data Flow diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/docs/device_log_profiling.md b/src/a2a3/runtime/tensormap_and_ringbuffer/docs/device_log_profiling.md index af661d440f..4b5cb362bf 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/docs/device_log_profiling.md +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/docs/device_log_profiling.md @@ -52,8 +52,8 @@ Thread 3: PTO2 total submitted tasks = 16704 ### Field Reference -| Field | Source (`pto_orchestrator.cpp`) | Description | -| ----- | ------------------------------- | ----------- | +| Field | Source (`pto_orchestrator.h`) | Description | +| ----- | ----------------------------- | ----------- | | **cost** | Wall-clock around `orch_func()` call | Total time including orchestration logic + scope overhead | | **total** | Sum of all sub-steps below | Accumulated time inside `submit_task` across all tasks | | **sync_tensormap** | `g_orch_sync_cycle` | TensorMap validity sync and optional cleanup before each submission | diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/docs/profiling_levels.md b/src/a2a3/runtime/tensormap_and_ringbuffer/docs/profiling_levels.md index fe82bcc009..3002565083 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/docs/profiling_levels.md +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/docs/profiling_levels.md @@ -48,7 +48,7 @@ Each sub-level macro requires `SIMPLER_DFX=1`: - Debug/diagnostic logs (always present) - Progress tracking (`PTO2 progress: completed=...`) -- Stall detection and dump (triggered after the `SCHEDULER_TIMEOUT_MS` wall-clock no-progress budget) +- Stall detection and dump (triggered only after `MAX_IDLE_ITERATIONS` idle loops) - Deadlock/livelock detection (`diagnose_stuck_state`, called on stall) **What's NOT compiled:** @@ -255,7 +255,7 @@ Identity fields the AICPU side used to write at level 1 (`func_id`, collector (`L2SwimlaneCollector::set_core_types`). AICore buffer rotation no longer piggy-backs on `complete_task`. AICPU -counts dispatches per core in the dispatch path (scheduler_dispatch in +counts dispatches per core in the dispatch path (scheduler_context in tensormap_and_ringbuffer; aicpu_executor in host_build_graph) and rotates the AICore buffer when the count is about to cross a `PLATFORM_AICORE_BUFFER_SIZE` boundary — strictly before @@ -428,7 +428,7 @@ definitions to runtime headers. ### Code Locations - Macro defaults and validation: `src/common/task_interface/profiling_config.h` -- Scheduler profiling: `src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_dispatch.cpp` and `scheduler_cold_path.cpp` +- Scheduler profiling: `src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler_context.h` - Orchestrator profiling: `src/a2a3/runtime/tensormap_and_ringbuffer/aicpu/aicpu_executor.cpp` - TensorMap profiling: `src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_tensormap.h` diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/host/dep_gen_replay.cpp b/src/a2a3/runtime/tensormap_and_ringbuffer/host/dep_gen_replay.cpp index d416931e3b..f89c3341e6 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/host/dep_gen_replay.cpp +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/host/dep_gen_replay.cpp @@ -558,7 +558,7 @@ dep_gen_replay_emit_deps_json(const DepGenRecord *records, size_t num_records, c // `explicit_dep_count` / `over->dep_count` originate from device // shared memory and are bounded by the writer to the array sizes, but // we clamp on read too so a corrupted record never drives an OOB read - // off the end of rec.explicit_deps[64] / over->deps[582]. + // off the end of rec.explicit_deps[64] / over->deps[326]. const uint64_t *deps_data; int32_t dc; if (rec.flags & DEP_GEN_FLAG_HAS_OVERFLOW) { diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/orchestration/common.cpp b/src/a2a3/runtime/tensormap_and_ringbuffer/orchestration/common.cpp index c4878a1c26..8768359def 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/orchestration/common.cpp +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/orchestration/common.cpp @@ -10,6 +10,13 @@ */ #include "common.h" +// LOG_ERROR can't be pulled from common/unified_log.h here because that header +// would re-#define LOG_INFO_V0..V9 already provided by pto_orchestration_api.h +// (orchestration routes them through the runtime ops table). For the limited +// use inside this file, write directly to stderr. +#include +#define LOG_ERROR(fmt, ...) std::fprintf(stderr, "[ERROR] " fmt "\n", ##__VA_ARGS__) + #ifdef __linux__ #include #include diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/aicore_completion_mailbox.h b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/aicore_completion_mailbox.h index 0f73a043a2..9d21ba10e3 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/aicore_completion_mailbox.h +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/aicore_completion_mailbox.h @@ -19,14 +19,6 @@ #include "pto_constants.h" #include "pto_task_id.h" -// AICPU-only MPSC ring used to convey deferred-completion observations from -// FIN-handling scheduler threads to the dispatch thread. Producers push under -// CAS on `head`; the single consumer (dispatch thread, under AsyncWaitList:: -// busy) drains in seq order. Kernel-side code never touches this struct — -// AICore writes go into DeferredCompletionSlab (see -// aicore_completion_mailbox_types.h), which the FIN thread reads, flattens -// into messages here, and forwards. - #define AICORE_COMPLETION_MAILBOX_CAPACITY 4096u #define AICORE_COMPLETION_MAILBOX_MASK (AICORE_COMPLETION_MAILBOX_CAPACITY - 1u) @@ -46,15 +38,8 @@ static_assert( #define MSG_KIND_TASK_NORMAL_DONE 1u struct AICoreCompletionMailboxMessage { - // Per-slot ready flag. Producer publishes `tail+1` after filling the rest - // of the slot with a release store; consumer waits for the matching seq - // value with an acquire load. The release-acquire pair publishes all - // other fields below as a side effect, so they stay plain. std::atomic seq; PTO2TaskId task_token; - // CONDITION: completion observation addr (counter / SDMA event record). - // TASK_NORMAL_DONE: PTO2TaskSlotState pointer carried over to the consumer - // so it can finalize the AsyncWaitEntry.slot_state binding. uint64_t addr; uint32_t expected_value; uint32_t engine; @@ -73,9 +58,6 @@ static_assert( "AICoreCompletionMailbox requires lock-free uint64_t atomics on every supported target" ); -// POD view of a drained message. `seq` is the ring's publication flag, not -// payload, so try_pop copies out only the fields below (and seq is not even -// copyable — it is a std::atomic). struct AICoreCompletionMsgView { PTO2TaskId task_token{PTO2TaskId::invalid()}; uint64_t addr{0}; @@ -98,21 +80,6 @@ struct AICoreCompletionMailbox { // consumer lock; a stale answer only over/under-triggers a drain attempt. bool has_pending() { return tail.load(std::memory_order_acquire) < head.load(std::memory_order_acquire); } - // MPSC push for a CONDITION message. Returns false when the ring is full - // (head - tail >= CAPACITY); caller should SPIN_WAIT_HINT and retry. - // Lock-free: CAS the shared head to claim a slot, write the fields, then - // release-store seq so the single consumer observes the publication. - // - // The head CAS is relaxed: head is a pure ticket counter and carries no - // data to the consumer — publication is solely the seq release-store, and - // slot-reuse safety rests on the acquire load of tail. The relaxed failure - // order is likewise sufficient since a lost CAS just re-reads head and - // retries. compare_exchange_weak is used because this loop already re-reads - // head and re-checks fullness, so masking LL/SC spurious failures (what - // _strong adds on aarch64) would only be a redundant inner retry. - // - // Safe to call concurrently from any number of producers; structurally - // independent of the AsyncWaitList::busy lock. bool try_push_condition( PTO2TaskId task_token, uint64_t addr, uint32_t expected_value, uint32_t engine, int32_t completion_type ) { @@ -136,9 +103,6 @@ struct AICoreCompletionMailbox { } } - // MPSC push for a TASK_NORMAL_DONE sentinel. Carries the PTO2TaskSlotState - // pointer in the `addr` field so the consumer can finish binding the - // AsyncWaitEntry.slot_state without going back to the FIN-handling thread. bool try_push_normal_done(PTO2TaskId task_token, uint64_t slot_state_addr) { while (true) { uint64_t h = head.load(std::memory_order_relaxed); @@ -159,12 +123,6 @@ struct AICoreCompletionMailbox { } } - // Single-consumer transport-level dequeue (caller holds the consumer lock). - // Returns false at the first not-yet-published slot (gap) or when empty; - // otherwise copies the next message in tail order into `out`, advances - // tail, and returns true. tail is consumer-only-written (relaxed read); - // head bounds the scan (relaxed); the seq acquire is the real publication - // gate; the tail release publishes "slot free" to reusing producers. bool try_pop(AICoreCompletionMsgView &out) { uint64_t t = tail.load(std::memory_order_relaxed); uint64_t h = head.load(std::memory_order_relaxed); diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/aicore_completion_mailbox_types.h b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/aicore_completion_mailbox_types.h index da0d89ad7e..aa8915a26b 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/aicore_completion_mailbox_types.h +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/aicore_completion_mailbox_types.h @@ -16,16 +16,6 @@ #include "pto_constants.h" -// Types shared across the AICore↔AICPU boundary. -// -// This header is reachable from AICore-side translation units (via -// pto_async_kernel_api.h / pto_completion_token.h / sdma_completion_kernel.h) -// and must stay parseable by every AICore toolchain configuration: no -// , no __atomic_* intrinsics, no MPSC ring buffer struct. -// -// The MPSC ring (AICoreCompletionMailbox) and its push/drain helpers live in -// aicore_completion_mailbox.h, which is AICPU-only. - inline constexpr int32_t MAX_COMPLETIONS_PER_TASK = 64; #define COMPLETION_ENGINE_SDMA 0u @@ -36,13 +26,6 @@ inline constexpr int32_t MAX_COMPLETIONS_PER_TASK = 64; #define COMPLETION_TYPE_COUNTER 0 #define COMPLETION_TYPE_SDMA_EVENT_RECORD 1 -// DeferredCompletionEntry / DeferredCompletionSlab back the per-task scratch -// area that AICore writes into to record "this completion has to be observed -// before the task can retire." The FIN-handling scheduler thread reads the -// slab, flattens entries into AICoreCompletionMailbox messages, and forwards -// them to the dispatch thread. `volatile` here is load-bearing: writers live -// on AICore and readers on AICPU, so the qualifier is the correct way to -// pin the compiler against caching / reordering on either side. struct DeferredCompletionEntry { uint64_t addr; uint32_t expected_value; diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/backend/sdma/sdma_completion_kernel.h b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/backend/sdma/sdma_completion_kernel.h index 0ff21908fe..94d76616ec 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/backend/sdma/sdma_completion_kernel.h +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/backend/sdma/sdma_completion_kernel.h @@ -36,17 +36,6 @@ enum class SdmaOp : uint8_t { TPUT = 1, }; -// SdmaRequestDescriptor bundles everything send_request_entry needs to drive -// one SDMA transfer + completion registration. It is a template because the -// destination / source / scratch types carry tensor shape & stride at compile -// time; the SdmaTget() / SdmaTput() helpers below let callers skip the -// template arguments. -// -// sync_id selects which event-record slot inside the workspace the engine -// writes into. Concurrent dispatches must use distinct sync_ids; today every -// caller submits one request per kernel invocation so passing 0 is safe. -// Future work (see .docs/25.comm-api-refactor/03.implementation-plan.md §5.2) -// will fold sync_id allocation into the adapter. template struct SdmaRequestDescriptor { SdmaOp op; @@ -91,9 +80,7 @@ register_pto_async_event(AsyncCtx &ctx, const PtoAsyncEvent &event, const PtoAsy (void)event.Wait(session); return; } - if (event.handle == 0) { - return; - } + if (event.handle == 0) return; const uint32_t engine = static_cast(event.engine); if (engine != static_cast(::pto::comm::DmaEngine::SDMA)) { @@ -111,17 +98,12 @@ register_pto_async_event(AsyncCtx &ctx, const PtoAsyncEvent &event, const PtoAsy defer_error(ctx, PTO2_ERROR_ASYNC_COMPLETION_INVALID); return; } - for (uint32_t queue_id = 0; queue_id < queue_num; ++queue_id) { + for (uint32_t queue_id = 0; queue_id < queue_num; ++queue_id) register_sdma_event_record(ctx, ::pto::comm::sdma::detail::GetEventRecord(recv_workspace, queue_id)); - } } } // namespace pto2::detail -// SDMA overload of the runtime's send_request_entry. Submits the descriptor -// to PTO-ISA, then registers the resulting AsyncEvent's GM flag(s) into the -// AsyncCtx deferred-wait slab and flushes. Returns false on submit/session -// failure (also records the error in ctx.completion_error_code). template inline __aicore__ bool send_request_entry(AsyncCtx &ctx, SdmaRequestDescriptor desc) { @@ -132,11 +114,8 @@ send_request_entry(AsyncCtx &ctx, SdmaRequestDescriptor(static_cast(record_addr)); cache_invalidate_range(reinterpret_cast(sdma_completion_cache_line(record)), PTO2_ALIGN_SIZE); diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_async_kernel_api.h b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_async_kernel_api.h index cf6eb47901..f44438c647 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_async_kernel_api.h +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_async_kernel_api.h @@ -29,10 +29,6 @@ #define __gm__ #endif -// Public surface: get_async_ctx, async_ctx_is_deferred, -// register_completion_condition, send_notification, -// save_expected_notification_counter. Everything else lives in -// pto2::detail and is reserved for backend adapters / internal use. namespace pto2::detail { inline __aicore__ void defer_load_slab(AsyncCtx &ctx) { @@ -46,9 +42,7 @@ inline __aicore__ void defer_load_slab(AsyncCtx &ctx) { } inline __aicore__ void defer_error(AsyncCtx &ctx, int32_t error_code) { - if (ctx.task_token.is_valid() && ctx.completion_error_code != nullptr) { - *ctx.completion_error_code = error_code; - } + if (ctx.task_token.is_valid() && ctx.completion_error_code != nullptr) *ctx.completion_error_code = error_code; } inline __aicore__ void defer_flush_range(volatile __gm__ void *addr, uint32_t size_bytes) { @@ -57,9 +51,8 @@ inline __aicore__ void defer_flush_range(volatile __gm__ void *addr, uint32_t si uintptr_t start = reinterpret_cast(addr) & ~(uintptr_t(PTO2_ALIGN_SIZE) - 1u); uintptr_t end = (reinterpret_cast(addr) + size_bytes + PTO2_ALIGN_SIZE - 1u) & ~(uintptr_t(PTO2_ALIGN_SIZE) - 1u); - for (uintptr_t p = start; p < end; p += PTO2_ALIGN_SIZE) { + for (uintptr_t p = start; p < end; p += PTO2_ALIGN_SIZE) dcci((__gm__ int32_t *)p, SINGLE_CACHE_LINE, CACHELINE_OUT); - } #else (void)addr; (void)size_bytes; @@ -70,16 +63,11 @@ inline __aicore__ void defer_flush(AsyncCtx &ctx) { if (ctx.task_token.is_invalid() || ctx.completion_count == nullptr) return; #if defined(__CCE_KT_TEST__) || defined(__CCE_AICORE__) || defined(__DAV_C220__) uint32_t count = *ctx.completion_count; - if (count > ctx.completion_capacity) { - count = ctx.completion_capacity; - } + if (count > ctx.completion_capacity) count = ctx.completion_capacity; uint32_t flush_bytes = static_cast(sizeof(*ctx.completion_count)); - if (ctx.completion_error_code != nullptr) { - flush_bytes += static_cast(sizeof(*ctx.completion_error_code)); - } - if (ctx.completion_entries != nullptr) { + if (ctx.completion_error_code != nullptr) flush_bytes += static_cast(sizeof(*ctx.completion_error_code)); + if (ctx.completion_entries != nullptr) flush_bytes += count * static_cast(sizeof(DeferredCompletionEntry)); - } defer_flush_range(ctx.completion_count, flush_bytes); #if defined(__CPU_SIM) dsb(0); @@ -110,21 +98,13 @@ inline __aicore__ AsyncCtx get_async_ctx(__gm__ int64_t *args) { inline __aicore__ bool async_ctx_is_deferred(const AsyncCtx &ctx) { return ctx.task_token.is_valid(); } -// Canonical writer: backend submit handlers build a CompletionToken and pass -// it here. Writes one DeferredCompletionEntry to the AsyncCtx slab and -// bumps completion_count. Returns false on overflow (also stores -// PTO2_ERROR_ASYNC_WAIT_OVERFLOW in ctx.completion_error_code) or when ctx is -// not currently a deferred context. inline __aicore__ bool register_completion_condition(AsyncCtx &ctx, const CompletionToken &token) { - if (ctx.task_token.is_invalid() || ctx.completion_count == nullptr || ctx.completion_entries == nullptr) { + if (ctx.task_token.is_invalid() || ctx.completion_count == nullptr || ctx.completion_entries == nullptr) return false; - } uint32_t idx = *ctx.completion_count; if (idx >= ctx.completion_capacity) { - if (ctx.completion_error_code != nullptr) { - *ctx.completion_error_code = PTO2_ERROR_ASYNC_WAIT_OVERFLOW; - } + if (ctx.completion_error_code != nullptr) *ctx.completion_error_code = PTO2_ERROR_ASYNC_WAIT_OVERFLOW; return false; } diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_async_wait.h b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_async_wait.h index 2f8f718fb5..b007c7b86c 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_async_wait.h +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_async_wait.h @@ -24,15 +24,11 @@ #include "pto_runtime2_types.h" struct PTO2SchedulerState; +struct PTO2LocalReadyBuffer; struct CompletionStats; inline constexpr int32_t MAX_ASYNC_WAITS = 64; -// The mailbox transport (has_pending / try_push_condition / -// try_push_normal_done / try_pop) lives as AICoreCompletionMailbox member -// functions in aicore_completion_mailbox.h. This file only holds the -// application layer: translating drained messages into wait-list state. - inline uintptr_t mailbox_cache_line(const volatile void *addr) { return reinterpret_cast(addr) & ~(uintptr_t(PTO2_ALIGN_SIZE) - 1u); } @@ -60,20 +56,15 @@ struct CompletionCondition { void retire(); }; -// Per-completion-type ops. SDMA_EVENT_RECORD detail lives in -// backend/sdma/sdma_completion_scheduler.h; the op wrappers below are thin -// glue mapping CompletionCondition.addr into the backend's raw-addr helpers. inline CompletionPollResult counter_poll_op(const CompletionCondition &cond) { - if (cond.counter_addr == nullptr) { - return {CompletionPollState::FAILED, PTO2_ERROR_ASYNC_COMPLETION_INVALID}; - } + if (cond.counter_addr == nullptr) return {CompletionPollState::FAILED, PTO2_ERROR_ASYNC_COMPLETION_INVALID}; return { *cond.counter_addr >= cond.expected_value ? CompletionPollState::READY : CompletionPollState::PENDING, PTO2_ERROR_NONE }; } -inline void counter_retire_op(CompletionCondition & /*cond*/) {} +inline void counter_retire_op(CompletionCondition &) {} inline CompletionPollResult sdma_event_record_poll_op(const CompletionCondition &cond) { return poll_sdma_event_record(cond.addr); @@ -92,22 +83,17 @@ inline const CompletionBackendOps *completion_backend_ops_for(int completion_typ } inline CompletionPollResult CompletionCondition::test() const { - if (satisfied) { - return {CompletionPollState::READY, PTO2_ERROR_NONE}; - } + if (satisfied) return {CompletionPollState::READY, PTO2_ERROR_NONE}; const CompletionBackendOps *ops = completion_backend_ops_for(completion_type); - if (ops == nullptr || ops->poll == nullptr) { + if (ops == nullptr || ops->poll == nullptr) return {CompletionPollState::FAILED, PTO2_ERROR_ASYNC_COMPLETION_INVALID}; - } return ops->poll(*this); } inline void CompletionCondition::retire() { if (retired) return; const CompletionBackendOps *ops = completion_backend_ops_for(completion_type); - if (ops != nullptr && ops->retire != nullptr) { - ops->retire(*this); - } + if (ops != nullptr && ops->retire != nullptr) ops->retire(*this); retired = true; } @@ -145,10 +131,6 @@ struct AsyncWaitList { std::atomic busy{0}; AsyncWaitEntry entries[MAX_ASYNC_WAITS]; int32_t count{0}; - // Diagnostic: counts every FIN-side try_push that hit a full mailbox. - // Expected to stay zero on real workloads (ring is 4096 entries); a - // non-zero value means consumers are too slow or the ring is undersized. - // Read by scheduler shutdown / l2 perf summary; not on the hot path. std::atomic mpsc_skipped_count{0}; void reset_for_reuse() { @@ -165,36 +147,21 @@ struct AsyncWaitList { void unlock() { busy.store(0, std::memory_order_release); } AsyncWaitEntry *find_entry_by_token(PTO2TaskId token) { - for (int32_t i = 0; i < count; i++) { + for (int32_t i = 0; i < count; i++) if (entries[i].task_token == token) return &entries[i]; - } return nullptr; } - // Captures the side-channel a scheduler-aware drain needs to complete - // NotDeferred tasks inline (without storing a transient entry in - // 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}; -#endif 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 - // translate it into wait-list state. An empty sink (sched == nullptr) just - // materializes entries; a sched-aware sink additionally inline-completes - // lonely NotDeferred NORMAL_DONEs without ever growing entries[]. int32_t drain_aicore_completion_mailbox_locked( AICoreCompletionMailbox *aicore_mailbox, DrainCompletionSink &sink, int32_t &error_code ) { @@ -203,17 +170,11 @@ struct AsyncWaitList { int32_t drained = 0; AICoreCompletionMsgView msg; - // try_pop is the transport layer (seq-gated, in-order dequeue); this - // loop is the application layer (translate each message into wait-list - // state). try_pop returns false at the first gap or when empty. while (aicore_mailbox->try_pop(msg)) { drained++; if (msg.kind == MSG_KIND_CONDITION) { AsyncWaitEntry *entry = find_entry_by_token(msg.task_token); if (entry == nullptr) { - // First message for this task — materialize the entry here. - // slot_state stays null until the matching TASK_NORMAL_DONE - // sentinel arrives. if (count >= MAX_ASYNC_WAITS) { error_code = PTO2_ERROR_ASYNC_WAIT_OVERFLOW; return drained; @@ -228,20 +189,13 @@ struct AsyncWaitList { if (!append_condition_locked( *entry, msg.addr, msg.expected_value, static_cast(msg.engine), msg.completion_type, error_code - )) { + )) return drained; - } } else if (msg.kind == MSG_KIND_TASK_NORMAL_DONE) { PTO2TaskSlotState *slot_state_ptr = reinterpret_cast(static_cast(msg.addr)); AsyncWaitEntry *entry = find_entry_by_token(msg.task_token); if (entry == nullptr) { - // Producers strictly order: all CONDITIONs for token T are - // pushed before the matching NORMAL_DONE (the acq_rel on - // on_subtask_complete enforces this across producers). So - // observing NORMAL_DONE first => the task registered no - // conditions => NotDeferred. Complete it inline when the - // sink allows; otherwise fall back to the entry-store path. if (sink.can_inline_complete()) { (void)try_inline_complete_locked(sink, *slot_state_ptr); continue; @@ -257,9 +211,7 @@ struct AsyncWaitList { entry->waiting_completion_count = 0; entry->normal_done = true; } else { - if (entry->slot_state == nullptr) { - entry->slot_state = slot_state_ptr; - } + if (entry->slot_state == nullptr) entry->slot_state = slot_state_ptr; entry->normal_done = true; } } else { @@ -293,15 +245,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 -#if SIMPLER_SCHED_PROFILING - , - int thread_idx -#endif - ); + AsyncPollResult poll_and_complete(AICoreCompletionMailbox *aicore_mailbox, PTO2SchedulerState *sched); }; #endif // PTO_ASYNC_WAIT_H diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_completion_token.h b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_completion_token.h index c5a8c345f3..61dae964a1 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_completion_token.h +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_completion_token.h @@ -17,12 +17,6 @@ #include "aicore_completion_mailbox_types.h" #include "pto_runtime_status.h" -// CompletionToken is the runtime-internal POD that backend submit handlers -// produce and the generic register_completion_condition() consumes. It is the -// ABI contract for "this is one completion to wait on" — independent of which -// backend (SDMA, RoCE, notification counter, ...) generated it. Each backend's -// (poll, retire) pair is registered in pto_async_wait.h's ops table, keyed by -// completion_type. struct CompletionToken { uint64_t addr; uint32_t expected_value; diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_dep_compute.h b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_dep_compute.h index f8392dfbf8..76ea8b11ad 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_dep_compute.h +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_dep_compute.h @@ -9,37 +9,6 @@ * ----------------------------------------------------------------------------------------------------------- */ -/** - * @file pto_dep_compute.h - * @brief Dependency computation primitives shared by runtime submit_task and dep_gen replay. - * - * Two header-only template entry points: - * - * compute_task_fanin — STEP 3 in submit_task: per-tensor creator retention (Step A) - * + tensormap.lookup for INPUT/INOUT (Step B). Calls back into - * user-supplied `emit` for each producer it identifies. - * - * register_task_outputs — STEP 4 in submit_task: tensormap.insert for INOUT and - * OUTPUT_EXISTING tensors. No callbacks. - * - * STEP 1 (explicit_deps) is intentionally left at the runtime call site because its - * `last_task_alive` shortcut + unchecked slot lookup is subtly different from the - * `slot_state->task->task_id == producer` reuse check in STEP 3. Unifying them would - * require two emit semantics or a marginal behavior change in transients — not worth - * the minor structural overlap. Replay handles STEP 1 with a one-line loop of its own. - * - * The Emit callback contract: - * bool emit(PTO2TaskId producer); - * - return true to continue (whether or not the producer was actually recorded — - * producer-not-alive / dedup-hit / etc. all return true silently) - * - return false to signal fatal (e.g. fanin spill overflow); caller bails - * - * Performance: Emit is a template parameter, not std::function. Both runtime - * (lambda capturing fanin_builder + sm_header) and replay (lambda capturing edge - * vector) instantiate at the call site and inline through. Do NOT replace with - * std::function — it would break the inlining and add ~5 ns/call to the orch hot path. - */ - #ifndef SRC_A2A3_RUNTIME_TENSORMAP_AND_RINGBUFFER_RUNTIME_PTO_DEP_COMPUTE_H_ #define SRC_A2A3_RUNTIME_TENSORMAP_AND_RINGBUFFER_RUNTIME_PTO_DEP_COMPUTE_H_ @@ -50,13 +19,6 @@ #include "pto_types.h" // TensorRef #include "tensor.h" -/** - * View struct for inputs to compute_task_fanin / register_task_outputs. - * - * Both runtime and replay assemble one of these from their own data sources - * (runtime: from Arg accessors; replay: from SubmitTraceEntry fields). All - * pointer arrays must remain valid for the duration of the call. - */ struct DepInputs { int32_t tensor_count; const TensorRef *tensors; // length = tensor_count (union; OUTPUT slots' .ptr is unused) @@ -65,24 +27,10 @@ struct DepInputs { const PTO2TaskId *explicit_deps; // length = explicit_dep_count (validity checked by caller) }; -/** - * Compute fanin for a task being submitted (STEP 3: Step A creator retention + - * Step B tensormap modifier lookup). - * - * For each non-OUTPUT tensor: - * - If owner_task_id is valid, emit(owner) - * - For INPUT/INOUT (and not manual_dep), tensor_map.lookup(*tensor) and emit - * each matching producer. INOUT+COVERED triggers tensor_map.remove_entry(entry). - * - * @return true on success (or producer-skipped-silently); false if emit signaled - * fatal — caller should propagate (after any fatal bookkeeping done by emit). - */ template [[nodiscard]] inline bool compute_task_fanin(const DepInputs &inputs, PTO2TensorMap &tensor_map, bool in_manual_scope, Emit emit) { - if (in_manual_scope) { - return true; - } + if (in_manual_scope) return true; for (int32_t i = 0; i < inputs.tensor_count; i++) { TensorArgType ptype = inputs.arg_types[i]; @@ -97,18 +45,12 @@ compute_task_fanin(const DepInputs &inputs, PTO2TensorMap &tensor_map, bool in_m // Step A: creator retention — all existing tensors extend their creator lifetime. PTO2TaskId owner = tensor->owner_task_id; if (owner.is_valid()) { - if (!emit(owner)) { - return false; - } + if (!emit(owner)) return false; } // Step B: only INPUT/INOUT need modifier dependency lookup. - if (ptype != TensorArgType::INPUT && ptype != TensorArgType::INOUT) { - continue; - } - if (tensor->manual_dep) { - continue; - } + if (ptype != TensorArgType::INPUT && ptype != TensorArgType::INOUT) continue; + if (tensor->manual_dep) continue; bool fatal = false; tensor_map.lookup(*tensor, [&](PTO2TensorMapEntry &entry, OverlapStatus overlap_status) -> bool { @@ -116,64 +58,25 @@ compute_task_fanin(const DepInputs &inputs, PTO2TensorMap &tensor_map, bool in_m fatal = true; return false; // stop iteration } - if (ptype == TensorArgType::INOUT && overlap_status == OverlapStatus::COVERED) { + if (ptype == TensorArgType::INOUT && overlap_status == OverlapStatus::COVERED) tensor_map.remove_entry(entry); - } return true; }); - if (fatal) { - return false; - } + if (fatal) return false; } return true; } -/** - * Register a task's outputs in the tensormap (STEP 4 in submit_task). - * - * For INOUT and OUTPUT_EXISTING tensors (excluding manual_dep), inserts the - * tensor into tensor_map keyed by its buffer.addr with `task_id` as producer. - * - * No-op when in_manual_scope. - */ inline void register_task_outputs(const DepInputs &inputs, PTO2TaskId task_id, PTO2TensorMap &tensor_map, bool in_manual_scope) { - if (in_manual_scope) { - return; - } + if (in_manual_scope) return; for (int32_t i = 0; i < inputs.tensor_count; i++) { TensorArgType ptype = inputs.arg_types[i]; if (ptype == TensorArgType::INOUT || ptype == TensorArgType::OUTPUT_EXISTING) { const Tensor *tensor = &inputs.tensors[i].ref(); - if (!tensor->manual_dep) { - tensor_map.insert(*tensor, task_id); - } - } - } -} - -/** - * Count the tensormap entries register_task_outputs() will insert for this task. - * - * Mirrors register_task_outputs()'s selection exactly (INOUT / OUTPUT_EXISTING, - * excluding manual_dep), so the returned value is the precise number of - * new_entry() calls that step makes. The orchestrator uses it to reserve pool - * capacity before inserting. Returns 0 in a manual scope (no registration). - */ -inline int32_t count_registrable_outputs(const DepInputs &inputs, bool in_manual_scope) { - if (in_manual_scope) { - return 0; - } - int32_t needed = 0; - for (int32_t i = 0; i < inputs.tensor_count; i++) { - TensorArgType ptype = inputs.arg_types[i]; - if (ptype == TensorArgType::INOUT || ptype == TensorArgType::OUTPUT_EXISTING) { - if (!inputs.tensors[i].ref().manual_dep) { - needed++; - } + if (!tensor->manual_dep) tensor_map.insert(*tensor, task_id); } } - return needed; } #endif // SRC_A2A3_RUNTIME_TENSORMAP_AND_RINGBUFFER_RUNTIME_PTO_DEP_COMPUTE_H_ diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.cpp b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.cpp index 5eb2c09580..7d62caa8fe 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.cpp +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.cpp @@ -12,23 +12,17 @@ /** * PTO Runtime2 - Orchestrator Implementation * - * Implements orchestrator state management, scope handling, and task submission. - * - * Based on: docs/RUNTIME_LOGIC.md + * Implements orchestrator state management, scope handling, and task submission + * for the polling completion design. */ #include "pto_orchestrator.h" #include -#include -#include -#include -#include -#include +#include #include "aicpu/dep_gen_collector_aicpu.h" #include "common/dep_gen.h" -#include "common/unified_log.h" #include "pto_dep_compute.h" #include "pto_runtime2_types.h" #include "pto_shared_memory.h" @@ -36,437 +30,65 @@ #include "pto_types.h" #include "tensor.h" -#if SIMPLER_DFX -#include "aicpu/scope_stats_collector_aicpu.h" -#include "aicpu/args_dump_aicpu.h" -#endif - -// Verify the captured Tensor blob size in DepGenRecord matches the runtime -// Tensor layout. The platform header defines DEP_GEN_TENSOR_SIZE without -// including runtime/tensor.h, so this check lives at the orch callsite. -static_assert(sizeof(Tensor) == DEP_GEN_TENSOR_SIZE, "DepGenRecord::tensors slot size out of sync with sizeof(Tensor)"); -// DEP_GEN_MAX_EXPLICIT_DEPS is a diagnostic-side capture cap only; the runtime -// imposes no hard cap on explicit dep count. If a submit exceeds this cap, -// dep_gen_aicpu_record_submit() logs and truncates — runtime correctness is -// unaffected, only the captured replay record is truncated. - -// Weak fallbacks: dep_gen_collector_aicpu.cpp provides the strong symbols in -// AICPU builds. Host builds (host_build_graph runtime, future dep_gen replay) -// link these no-op stubs so the runtime translation unit is self-contained. -// Visibility is hidden so the HOST .so doesn't export them into the global -// dynamic symbol table where they'd shadow the AICPU .so's strong symbols -// (same pattern as get_sys_cnt_aicpu / l2_swimlane_aicpu_record_orch_phase below). -extern "C" __attribute__((weak, visibility("hidden"))) bool is_dep_gen_enabled() { return false; } -__attribute__((weak, visibility("hidden"))) void dep_gen_aicpu_record_submit( - uint64_t, bool, bool, int, const void *const *, const uint8_t *, int, const uint64_t *, int, const int32_t[3] -) {} - -// Scope_stats enable gate, queried via the same predicate idiom as -// is_dep_gen_enabled above. The AICPU collector links the strong definition; -// host builds fall back to this weak `false`. Gating here still skips the -// cross-agent occupancy reads that feed the sample when scope_stats is disabled. -extern "C" __attribute__((weak, visibility("hidden"))) bool is_scope_stats_enabled() { return false; } - -// Heap-ring wrap report, called from the allocator (pto_ring_buffer.h) on each -// wrap. Strong definition lives in the AICPU collector; host builds fall back to -// this weak no-op so the runtime translation unit stays self-contained. -extern "C" __attribute__((weak, visibility("hidden"))) void scope_stats_note_heap_wrap(int) {} - -// ============================================================================= -// Orchestrator Profiling (compile-time toggle) -// ============================================================================= -#if SIMPLER_ORCH_PROFILING -#include "aicpu/device_time.h" -#include "aicpu/l2_swimlane_collector_aicpu.h" -// Weak fallback for builds that don't link device_time.cpp (e.g. host). -// The strong symbol from platform/.../device_time.cpp wins in the AICPU build. -// -// IMPORTANT: visibility("hidden") is required to prevent the HOST .so from -// exporting this weak fallback into the global dynamic symbol table via -// RTLD_GLOBAL. Without it, when the AICPU .so is loaded and its PLT entry -// for get_sys_cnt_aicpu is resolved, the dynamic linker finds the HOST .so's -// weak definition first (already in global table) and uses it — returning 0. -// With hidden visibility, the HOST .so does not export this symbol globally, -// so the AICPU .so's PLT resolves to its own strong definition from -// device_time.cpp. -__attribute__((weak, visibility("hidden"))) uint64_t get_sys_cnt_aicpu() { return 0; } -// Weak fallback for builds that don't link l2_swimlane_collector_aicpu.cpp. -// The strong symbol from the AICPU build wins when profiling is available. -// Also hidden to prevent HOST .so from polluting the global symbol table. -__attribute__((weak, visibility("hidden"))) void -l2_swimlane_aicpu_record_orch_phase(uint64_t, uint64_t, uint64_t, uint32_t) {} -// Accumulated cycles per sub-step (only needed for ORCH_PROFILING export) -static uint64_t g_orch_sync_cycle = 0; // tensormap sync -static uint64_t g_orch_alloc_cycle = 0; // unified task+heap alloc -static uint64_t g_orch_args_cycle = 0; // param copy -static uint64_t g_orch_lookup_cycle = 0; // tensormap lookup + dep building -static uint64_t g_orch_insert_cycle = 0; // tensormap insert -static uint64_t g_orch_fanin_cycle = 0; // fanin list + early-return check -static uint64_t g_orch_scope_end_cycle = 0; // scope_end overhead -static int64_t g_orch_submit_count = 0; -static uint32_t g_orch_submit_idx = 0; -uint64_t g_orch_alloc_wait_cycle = 0; -uint64_t g_orch_fanin_wait_cycle = 0; -uint64_t g_orch_alloc_atomic_count = 0; -uint64_t g_orch_args_atomic_count = 0; -uint64_t g_orch_scope_end_atomic_count = 0; -// Cycle accumulation is unconditional under SIMPLER_ORCH_PROFILING (that's what -// the flag is for) and feeds the per-sub-step `g_orch_*_cycle` cumulatives -// printed in the cold-path log. -// -// Per-submit ORCH_SUBMIT record is the only swim-lane emit on the orch -// path — one record per submit_task() / alloc_tensors() call spanning -// the entire [start, end] window. Per-sub-step phase records were dropped -// in favour of the cumulatives + per-submit envelope; the dispatcher -// already inserts one record at the end of each submit path via -// CYCLE_COUNT_ORCH_SUBMIT_RECORD. -#define CYCLE_COUNT_START() \ - bool _prof_active = (orch->l2_swimlane_level >= L2SwimlaneLevel::ORCH_PHASES); \ - uint64_t _t0 = get_sys_cnt_aicpu(), _t1; \ - uint64_t _submit_start_ts = _t0 -#define CYCLE_COUNT_LAP(acc) \ - do { \ - _t1 = get_sys_cnt_aicpu(); \ - acc += (_t1 - _t0); \ - _t0 = _t1; \ - } while (0) -#define CYCLE_COUNT_ORCH_SUBMIT_RECORD(tid) \ - do { \ - if (_prof_active) { \ - l2_swimlane_aicpu_record_orch_phase(_submit_start_ts, _t1, (tid), g_orch_submit_idx); \ - } \ - } while (0) -#elif SIMPLER_DFX -#include "aicpu/device_time.h" -#include "aicpu/l2_swimlane_collector_aicpu.h" -__attribute__((weak, visibility("hidden"))) uint64_t get_sys_cnt_aicpu() { return 0; } -__attribute__((weak, visibility("hidden"))) void -l2_swimlane_aicpu_record_orch_phase(uint64_t, uint64_t, uint64_t, uint32_t) {} -// submit_idx needed for swimlane task_id tagging (no cycle accumulation at this level) -static uint32_t g_orch_submit_idx = 0; -#define CYCLE_COUNT_START() \ - bool _prof_active = (orch->l2_swimlane_level >= L2SwimlaneLevel::ORCH_PHASES); \ - uint64_t _t0 = _prof_active ? get_sys_cnt_aicpu() : 0, _t1 = 0; \ - uint64_t _submit_start_ts = _t0 -#define CYCLE_COUNT_LAP(acc) \ - do { \ - } while (0) -#define CYCLE_COUNT_ORCH_SUBMIT_RECORD(tid) \ - do { \ - if (_prof_active) { \ - _t1 = get_sys_cnt_aicpu(); \ - l2_swimlane_aicpu_record_orch_phase(_submit_start_ts, _t1, (tid), g_orch_submit_idx); \ - } \ - } while (0) -#else -#define CYCLE_COUNT_START() -#define CYCLE_COUNT_LAP(acc) -#define CYCLE_COUNT_ORCH_SUBMIT_RECORD(tid) -#endif +// ----------------------------------------------------------------------------- +// File-local helpers +// ----------------------------------------------------------------------------- + +static int32_t orch_mark_fatal(PTO2OrchestratorState *orch, int32_t error_code); +static void orch_report_fatal_v(PTO2OrchestratorState *orch, int32_t error_code, const char *fmt, va_list args); +static void scope_tasks_push(PTO2OrchestratorState *orch, PTO2TaskSlotState *task_slot_state); +static bool prepare_task( + PTO2OrchestratorState *orch, const L0TaskArgs &args, int32_t total_output_size, ActiveMask active_mask, + PTO2PreparedTask *out +); +static PTO2OutputLayout calculate_output_layout(const L0TaskArgs &args); +static bool append_fanin_or_fail( + PTO2OrchestratorState *orch, PTO2TaskSlotState *prod_state, int32_t prod_local_id, PTO2FaninBuilder *fanin_builder +); +static bool check_scope_can_accept_task(PTO2OrchestratorState *orch, PTO2TaskAllocator &allocator); +static void prefetch_payload(PTO2TaskPayload *payload, int32_t tensor_count, int32_t scalar_count); +static TaskOutputTensors submit_task_common( + PTO2OrchestratorState *orch, const L0TaskArgs &args, ActiveMask active_mask, int32_t aic_kernel_id, + int32_t aiv0_kernel_id, int32_t aiv1_kernel_id +); static int32_t orch_mark_fatal(PTO2OrchestratorState *orch, int32_t error_code) { always_assert(orch != nullptr); orch->fatal = true; - if (error_code == PTO2_ERROR_NONE || orch->sm_header == nullptr) { - return PTO2_ERROR_NONE; - } + if (error_code == PTO2_ERROR_NONE || orch->sm_header == nullptr) return PTO2_ERROR_NONE; int32_t expected = PTO2_ERROR_NONE; std::atomic &orch_error_code = orch->sm_header->orch_error_code; - if (orch_error_code.compare_exchange_strong(expected, error_code, std::memory_order_acq_rel)) { - return error_code; - } + if (orch_error_code.compare_exchange_strong(expected, error_code, std::memory_order_acq_rel)) return error_code; return expected; } -static void -orch_report_fatal_v(PTO2OrchestratorState *orch, int32_t error_code, const char *func, const char *fmt, va_list args) { - int32_t latched_code = orch_mark_fatal(orch, error_code); - -#if SIMPLER_DFX - // Flush the current scope's peaks BEFORE the FATAL log line, so the - // diagnostic context (which pool/window filled up) appears right next to - // the failure reason. on_fatal is latched, so duplicate fatals from - // different layers don't print multiple stats lines. - scope_stats_on_fatal(); -#endif - - if (fmt == nullptr || fmt[0] == '\0') { - if (latched_code != PTO2_ERROR_NONE && latched_code != error_code) { - unified_log_error(func, "FATAL(code=%d, latched=%d)", error_code, latched_code); - } else { - unified_log_error(func, "FATAL(code=%d)", error_code); - } - return; - } - - char message[1024]; - vsnprintf(message, sizeof(message), fmt, args); - if (latched_code != PTO2_ERROR_NONE && latched_code != error_code) { - unified_log_error(func, "FATAL(code=%d, latched=%d): %s", error_code, latched_code, message); - return; - } - unified_log_error(func, "FATAL(code=%d): %s", error_code, message); -} - -void PTO2OrchestratorState::report_fatal(int32_t error_code, const char *func, const char *fmt, ...) { - auto *orch = this; - va_list args; - va_start(args, fmt); - orch_report_fatal_v(orch, error_code, func, fmt, args); - va_end(args); -} - -static uint32_t next_fanin_seen_epoch(PTO2OrchestratorState *orch) { - uint32_t next = orch->fanin_seen_current_epoch + 1; - if (next == 0) { - for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { - memset( - orch->fanin_seen_epoch[r], 0, - static_cast(orch->sm_header->rings[r].task_window_size) * sizeof(uint32_t) - ); - } - next = 1; - } - orch->fanin_seen_current_epoch = next; - return next; +static void orch_report_fatal_v(PTO2OrchestratorState *orch, int32_t error_code, const char *, va_list) { + // fmt + args are accepted for future logging-sink wiring but are not yet + // routed anywhere — the error_code is latched in shared memory via + // orch_mark_fatal and that's what callers actually observe. + orch_mark_fatal(orch, error_code); } -struct PTO2FaninBuilder { - PTO2FaninBuilder(PTO2OrchestratorState *orch, PTO2FaninPool &spill_pool, uint32_t seen_epoch) : - count(0), - spill_start(0), - orch(orch), - seen_epoch(seen_epoch), - spill_pool(spill_pool) {} - 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)); - } - - bool mark_seen(uint8_t prod_ring, int32_t prod_slot) { - if (prod_ring >= PTO2_MAX_RING_DEPTH || prod_slot < 0) { - return false; - } - uint32_t *seen = orch->fanin_seen_epoch[prod_ring]; - uint32_t slot = static_cast(prod_slot); - if (seen[slot] == seen_epoch) { - return true; - } - seen[slot] = seen_epoch; - return false; - } -}; - 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, uint8_t ring_id + PTO2OrchestratorState *orch, PTO2TaskSlotState *prod_state, int32_t prod_local_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(); - PTO2TaskState pstate = prod_state->task_state.load(std::memory_order_acquire); - bool gone = prod_state->task == nullptr || prod_state->task->task_id.local() != producer_task_id.local() || - pstate == PTO2_TASK_CONSUMED; - bool claim = !gone && !fanin_builder->mark_seen(prod_ring, prod_slot); - int32_t fanout_now = -1; - 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++; - fanout_now = static_cast(prod_state->fanout_count & ~PTO2_FANOUT_SCOPE_BIT); - } - 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 - // Dense-fanout diagnostic, emitted outside fanout_lock (no logging under the - // spinlock). The monotonic fanout_count++ crosses THRESHOLD+1 exactly once, - // so each dense producer warns exactly once and the AICPU hot path keeps a - // single predicted-not-taken branch on the common case. - if (fanout_now == PTO2_DEP_DEGREE_WARN_THRESHOLD + 1) { - LOG_WARN( - "dense dependency: task ring=%u id=%u fanout>%d [orch submit]", - static_cast(producer_task_id.ring()), producer_task_id.local(), PTO2_DEP_DEGREE_WARN_THRESHOLD - ); - } - // gone (stale/consumed) or an already-seen duplicate live producer: no new - // fanin edge either way. - if (!claim) { - return true; - } - - if (fanin_builder->count < PTO2_FANIN_INLINE_CAP) { - fanin_builder->inline_slots[fanin_builder->count++] = prod_state; - return true; - } - - PTO2FaninPool &fanin_pool = fanin_builder->spill_pool; - if (!fanin_pool.ensure_space(orch->sm_header->rings[ring_id], 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->contains(prod_state)) return true; + 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; - } - entry->slot_state = prod_state; - fanin_builder->count++; + int32_t idx = fanin_builder->count++; + fanin_builder->slots[idx] = prod_state; + fanin_builder->local_ids[idx] = prod_local_id; + fanin_builder->ring_ids[idx] = prod_state->ring_id; return true; } -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; - }); -} - -void PTO2OrchestratorState::mark_dep_pool_position(PTO2TaskSlotState &slot_state) { - PTO2SchedulerState *sched = scheduler; - auto &rss = sched->ring_sched_states[slot_state.ring_id]; - slot_state.dep_pool_mark = rss.dep_pool.top; -#if SIMPLER_DFX - if (is_scope_stats_enabled()) { - rss.publish_dep_pool_snapshot(); - } -#endif -} - -void PTO2OrchestratorState::wire_fanin_task(PTO2TaskSlotState &slot_state, int32_t wfanin) { - PTO2SchedulerState *sched = scheduler; - auto &rss = sched->ring_sched_states[slot_state.ring_id]; - PTO2TaskPayload *payload = slot_state.payload; - slot_state.fanin_count = wfanin + 1; - - int32_t completed_fanin = 0; - int32_t early_propagated = 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) { - completed_fanin++; - } else { - producer->fanout_head = rss.dep_pool.prepend(producer->fanout_head, &slot_state); - // The marker shares fanout_lock with propagation's snapshot. A set - // marker means this edge is outside that snapshot and needs a seed. - if (!early_disqualified && producer->has_dispatch_propagated()) { - early_propagated++; - } - } - producer->unlock_fanout(); - }); - - // Completed producers and edges outside a producer's one-shot propagation - // snapshot both need wiring-time contributions. Seed only when every - // producer is codegen-flagged; one unflagged producer disqualifies the task. - int32_t early_seed = completed_fanin + early_propagated; - if (!early_disqualified && early_seed != 0) { - int32_t dispatch_fanin = payload->dispatch_fanin.fetch_add(early_seed, std::memory_order_acq_rel) + early_seed; - // A fully pre-completed fanin routes normally. If any producer was live, - // the exact-full increment must enqueue the early candidate. - if (completed_fanin != payload->fanin_actual_count && dispatch_fanin == payload->fanin_actual_count) { - sched->try_enqueue_early_dispatch_candidate(slot_state); - } - } - - int32_t init_rc = completed_fanin + 1; - int32_t new_rc = slot_state.fanin_refcount.fetch_add(init_rc, std::memory_order_acq_rel) + init_rc; - mark_dep_pool_position(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) { - PTO2SchedulerState *sched = orch->scheduler; - auto &rss = sched->ring_sched_states[slot_state.ring_id]; - - // dep_pool is orchestrator-exclusive (no lock). ensure_space waits for the - // scheduler to advance last_task_alive and, on a wedged reclaim watermark, - // detects the deadlock with the same structural + wall-clock logic the - // heap/task-window allocator uses (all three share last_task_alive), latches - // PTO2_ERROR_DEP_POOL_OVERFLOW, and emits the structured report. 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(slot_state, wfanin); - return true; -} - -static void scope_tasks_push(PTO2OrchestratorState *orch, PTO2TaskSlotState *task_slot_state); - -struct PTO2PreparedTask { - PTO2TaskId task_id = PTO2TaskId::invalid(); - PTO2TaskAllocResult alloc_result = {-1, 0, nullptr, nullptr}; - PTO2TaskDescriptor *task = nullptr; - PTO2TaskPayload *payload = nullptr; - PTO2TaskSlotState *slot_state = nullptr; -}; - static PTO2OutputLayout calculate_output_layout(const L0TaskArgs &args) { PTO2OutputLayout layout; for (int32_t i = 0; i < args.tensor_count(); i++) { - if (args.tag(i) != TensorArgType::OUTPUT) { - continue; - } + if (args.tag(i) != TensorArgType::OUTPUT) continue; layout.offsets[i] = layout.total_output_size; layout.buffer_sizes[i] = PTO2_ALIGN_UP(args.tensor(i).create_info().buffer_size_bytes(), PTO2_PACKED_OUTPUT_ALIGN); @@ -475,39 +97,28 @@ static PTO2OutputLayout calculate_output_layout(const L0TaskArgs &args) { return layout; } -static bool check_scope_can_accept_task(PTO2OrchestratorState *orch, PTO2TaskAllocator &allocator, uint8_t ring_id) { +static bool check_scope_can_accept_task(PTO2OrchestratorState *orch, PTO2TaskAllocator &allocator) { always_assert(orch->scope_stack_top >= 0 && "Cannot submit task outside a scope"); int32_t scope_task_count = orch->scope_tasks_size - orch->scope_begins[orch->scope_stack_top]; - if (scope_task_count < allocator.window_size() - 1) { - return true; - } + if (scope_task_count < allocator.window_size() - 1) return true; - int32_t active_count = allocator.active_count(); - - LOG_ERROR("========================================"); - LOG_ERROR("FATAL: Scope Deadlock Detected! (ring %d)", ring_id); - LOG_ERROR("========================================"); - LOG_ERROR("Tasks in current scope (%d) >= task_window_size (%d).", scope_task_count, allocator.window_size()); - LOG_ERROR(" scope_depth: %d", orch->scope_stack_top + 1); - LOG_ERROR(" ring_id: %d", ring_id); - 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("Solution:"); - LOG_ERROR(" 1. Reduce tasks per scope (use batching/unroll)"); - LOG_ERROR(" 2. Increase task window (current: %d)", allocator.window_size()); - LOG_ERROR(" Compile-time: PTO2_TASK_WINDOW_SIZE in pto_runtime2_types.h"); - LOG_ERROR(" Runtime env: PTO2_RING_TASK_WINDOW="); - LOG_ERROR(" 3. Split work across multiple scopes"); - LOG_ERROR("========================================"); orch_mark_fatal(orch, PTO2_ERROR_SCOPE_DEADLOCK); return false; } +static void prefetch_payload(PTO2TaskPayload *payload, int32_t tensor_count, int32_t scalar_count) { + for (int32_t i = 0; i < tensor_count; i++) { + __builtin_prefetch(&payload->tensors[i], 1, 3); + __builtin_prefetch(reinterpret_cast(&payload->tensors[i]) + 64, 1, 3); + } + for (int32_t i = 0; i < scalar_count; i += 8) + __builtin_prefetch(&payload->scalars[i], 1, 3); + __builtin_prefetch(payload, 1, 3); + __builtin_prefetch(reinterpret_cast(payload) + 64, 1, 3); + __builtin_prefetch(reinterpret_cast(payload) + 128, 1, 3); +} + static bool prepare_task( PTO2OrchestratorState *orch, const L0TaskArgs &args, int32_t total_output_size, ActiveMask active_mask, PTO2PreparedTask *out @@ -515,9 +126,7 @@ static bool prepare_task( uint8_t ring_id = orch->current_ring_id(); auto &allocator = orch->rings[ring_id].task_allocator; - if (!check_scope_can_accept_task(orch, allocator, ring_id)) { - return false; - } + if (!check_scope_can_accept_task(orch, allocator)) return false; out->alloc_result = allocator.alloc(total_output_size); if (out->alloc_result.failed()) { @@ -530,67 +139,44 @@ static bool prepare_task( out->task = &orch->sm_header->rings[ring_id].task_descriptors[out->alloc_result.slot]; out->payload = &orch->sm_header->rings[ring_id].task_payloads[out->alloc_result.slot]; - // Reset the fanout/fanin bookkeeping for this reuse. The allocator only - // returns a slot whose previous occupant is CONSUMED and quiescent (alloc - // spins until last_task_alive passes it; in-order reclaim + acquire load), - // and the slot is not published to any scheduler thread until the - // Orch-side wiring publish at the end of submit_task_common — so this reset - // is race-free. Doing it here (not relying on the scheduler's eager - // reset-after-CONSUMED, which only covers the contiguously-reclaimed tail) - // makes every reused slot self-clean, which lets the per-boot SM init skip - // its O(window) per-slot loop. bind_ring is slot-invariant but cheap to - // re-assert on the already-dirtied cache line. + prefetch_payload(out->payload, args.tensor_count(), args.scalar_count()); + + // Reset the fanout/wake-list/subtask bookkeeping for this reuse. The allocator + // only returns a slot whose previous incarnation is fully consumed (alloc spins + // until completed_watermark passes its last_consumer_local_id), and the slot is + // not published to any scheduler thread until the wiring.queue.push at the end + // of submit_task_common — so this reset is race-free. Doing it here (not relying + // on the scheduler's eager reset-after-CONSUMED, which only covers the + // contiguously-reclaimed tail within a single run) makes every reused slot + // self-clean across runs, which lets the per-boot SM init skip its O(window) + // per-slot loop. bind_ring is slot-invariant but cheap to re-assert on the + // already-dirtied cache line. Mirrors upstream #1199. out->slot_state->bind_ring(ring_id); out->slot_state->reset_for_reuse(); - out->slot_state->fanin_count = 0; - out->slot_state->dep_pool_mark = 0; - - out->payload->prefetch(args.tensor_count(), args.scalar_count()); - - // Re-bind payload/task pointers each submit. Value is per-slot constant - // (same as &task_payloads[slot] / &task_descriptors[slot]), but writing - // here lets RingSchedState::init() skip the O(window_size) bind loop. - // Both writes hit the same 64B slot_state cache line we're about to - // dirty below, so the extra cost is two stores on an already-hot line. - // Must precede the Orch-side wiring publish at the end of - // submit_task_common — that publish is the first read of slot_state->task / - // slot_state->payload by scheduler threads. + out->slot_state->bind_buffers(out->payload, out->task); - // prepare_task does NO payload writes: all payload content (tensors/scalars + - // early-dispatch fields) is initialized in PTO2TaskPayload::init, the - // single payload-init point, which runs before Orch-side wiring publish. - - // Fields already reset by advance_ring_pointers (eager reset after CONSUMED): - // fanout_lock=0, fanout_count=PTO2_FANOUT_SCOPE_BIT, fanout_head=nullptr, - // fanin_refcount=0, fanout_refcount=0, completed_subtasks=0, next_block_idx=0 - // Fields immutable after RingSchedState::init(): - // ring_id - // task_state left as CONSUMED by eager reset (safe for stale wait_for_tensor - // observers); set to PENDING here when orchestrator actually reuses the slot. - out->slot_state->task_state.store(PTO2_TASK_PENDING, std::memory_order_relaxed); + // Clear the polling-fast completion byte for the newly-allocated slot. + // The previous incarnation's completer set this byte to 1; we publish 0 + // before this task can be added as a fanin to any consumer (single- + // orchestrator-thread guarantee) and before the wiring-queue push + // (release-acquire) makes the slot visible to thread 0. + orch->sm_header->rings[ring_id].completion_flags[out->alloc_result.slot].store(0, std::memory_order_relaxed); + // Seed last_consumer_local_id to self — with no consumers, the slot is + // safe to reclaim as soon as the watermark reaches this task itself. + out->slot_state->last_consumer_local_id = out->alloc_result.task_id; int16_t block_num = args.launch_spec.block_num(); out->slot_state->total_required_subtasks = 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 scope_tasks_push(orch, out->slot_state); return true; } -// ============================================================================= -// Scope Management -// ============================================================================= - static void scope_tasks_push(PTO2OrchestratorState *orch, PTO2TaskSlotState *task_slot_state) { if (orch->scope_tasks_size >= orch->scope_tasks_capacity) { - // scope_tasks lives in the per-Worker arena (single backing allocation), - // so realloc is not legal. Capacity is the total in-flight slot budget - // (sum of the per-ring task windows; see reserve_layout) — hitting it means - // every ring is saturated, so no further push could succeed regardless of - // buffer growth. orch->report_fatal( PTO2_ERROR_SCOPE_TASKS_OVERFLOW, __FUNCTION__, "scope_tasks buffer saturated at %d entries (all rings full)", orch->scope_tasks_capacity @@ -600,212 +186,14 @@ static void scope_tasks_push(PTO2OrchestratorState *orch, PTO2TaskSlotState *tas orch->scope_tasks[orch->scope_tasks_size++] = task_slot_state; } -void PTO2OrchestratorState::begin_scope(PTO2ScopeMode mode) { - auto *orch = this; - if (orch->fatal) { - return; - } - assert(orch->scope_stack_top < static_cast(orch->scope_stack_capacity - 1) && "Scope stack overflow"); - if (mode == PTO2ScopeMode::AUTO && orch->in_manual_scope()) { - report_fatal(PTO2_ERROR_INVALID_ARGS, __FUNCTION__, "auto scope nested inside manual scope is not supported"); - return; - } - - bool already_in_manual_scope = orch->in_manual_scope(); - ++orch->scope_stack_top; - orch->scope_begins[orch->scope_stack_top] = orch->scope_tasks_size; - if (mode == PTO2ScopeMode::MANUAL && !already_in_manual_scope) { - orch->manual_begin_depth = orch->scope_stack_top; - } -#if SIMPLER_DFX - // Gate via is_scope_stats_enabled() (weak-false in host builds) BEFORE the - // collector call: when disabled we pay nothing. Sample the current ring's - // task/heap start-end and tensormap usage at the scope boundary. - if (is_scope_stats_enabled()) { - uint8_t ring_id = orch->current_ring_id(); - auto &alloc = orch->rings[ring_id].task_allocator; - int32_t dep_pool_tail = 0; - int32_t dep_pool_top = 0; - if (orch->scheduler) { - orch->scheduler->ring_sched_states[ring_id].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() - ); - } -#endif -} - -void PTO2OrchestratorState::end_scope() { - auto *orch = this; - if (orch->fatal) { - return; - } - assert(orch->scope_stack_top >= 0 && "Scope stack underflow"); - - // Snapshot the ring start/end BEFORE the orchestrator drains pending tasks - // via scheduler->on_scope_end, so the end record reflects the scope's - // occupancy at close, not the residual after teardown. -#if SIMPLER_DFX - // Gate via is_scope_stats_enabled() (see begin_scope). One collector call - // emits the end-boundary record and tears down bookkeeping. - if (is_scope_stats_enabled()) { - uint8_t ring_id = orch->current_ring_id(); - auto &alloc = orch->rings[ring_id].task_allocator; - int32_t dep_pool_tail = 0; - int32_t dep_pool_top = 0; - if (orch->scheduler) { - orch->scheduler->ring_sched_states[ring_id].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() - ); - } -#endif - -#if SIMPLER_ORCH_PROFILING - uint64_t _se0 = get_sys_cnt_aicpu(); -#endif - - bool ending_manual_scope = orch->scope_stack_top == orch->manual_begin_depth; - int32_t begin = orch->scope_begins[orch->scope_stack_top--]; - int32_t count = orch->scope_tasks_size - begin; - if (ending_manual_scope) { - orch->manual_begin_depth = PTO2_MAX_SCOPE_DEPTH; - } - - if (orch->scheduler && count > 0) { - orch->scheduler->on_scope_end(&orch->scope_tasks[begin], count); - } - - // Rewind the task buffer — these entries are no longer needed - orch->scope_tasks_size = begin; - -#if SIMPLER_ORCH_PROFILING - uint64_t _se1 = get_sys_cnt_aicpu(); - g_orch_scope_end_cycle += (_se1 - _se0); -#endif -} - -// ============================================================================= -// Task Submission -// ============================================================================= - -// Ensure the tensormap entry pool has room for `needed` inserts before STEP 4 -// registers this task's outputs. The pool is watermark-reclaimed like the -// task/heap/fanin pools — retired tasks' entries free once last_task_alive -// advances — so an exhausted pool is back-pressure, not a hard error. Reclaim -// across all rings (entries from every ring share one pool); if still short, -// spin until reclaim actually frees entries, with the same 500 ms wall-clock -// backstop as the task allocator and fanin spill pool. A pool that stays full -// (no entry freed) is a genuine deadlock: latch PTO2_ERROR_TENSORMAP_OVERFLOW -// and bail. Returns false on deadlock or on a fatal already latched by another -// party. Cold path — the fast path returns immediately when the pool has room. -static bool ensure_tensormap_capacity(PTO2OrchestratorState *orch, int32_t needed) { - PTO2TensorMap &tm = orch->tensor_map; - if (tm.free_entries() >= needed) { - return true; - } - - int32_t alive[PTO2_MAX_RING_DEPTH]; - auto read_alive = [&]() { - for (int32_t r = 0; r < PTO2_MAX_RING_DEPTH; r++) { - // Relaxed: a self-correcting poll re-read every reclaim tick, so a stale - // watermark only defers reclaim one tick and never over-frees. - alive[r] = orch->sm_header->rings[r].fc.last_task_alive.load(std::memory_order_relaxed); - } - }; - - read_alive(); - int64_t cur_alive_sum = tm.reclaim_retired_all(alive); // kept for the deadlock diagnostic - int32_t prev_free = tm.free_entries(); - if (prev_free >= needed) { - return true; - } - - int spin_count = 0; - uint64_t block_cycle0 = 0; // wall-clock anchor for the deadlock backstop - bool block_timing = false; // false until the first no-reclaim-progress tick - while (tm.free_entries() < needed) { - spin_count++; - - // Reclaim (and the all-ring watermark reads it needs) is the costly part of - // this spin and the only path that frees entries; gate it to a periodic tick. - // Cold path, but the spin itself is tight. - if ((spin_count & 31) == 0) { - read_alive(); - cur_alive_sum = tm.reclaim_retired_all(alive); - int32_t cur_free = tm.free_entries(); - if (cur_free >= needed) { - return true; - } - // Progress is entries actually freed, NOT watermark movement: a ring can - // retire zero-output tasks (count_registrable_outputs == 0), advancing - // last_task_alive without freeing any entry. Gating the backstop on - // free_entries() keeps a wedged pool from dodging the timeout while some - // unrelated ring keeps draining. - if (cur_free > prev_free) { - spin_count = 0; - prev_free = cur_free; - block_timing = false; - } - } - - if ((spin_count & 1023) == 0) { - // A fatal latched elsewhere breaks this otherwise-unbounded spin. - if (orch->sm_header->orch_error_code.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 a cheap cntvct_el0 read; the 1024-spin gate keeps fatal-flag - // polling and timeout bookkeeping out of every entry-pool wait spin. - 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) { - LOG_ERROR("========================================"); - LOG_ERROR("FATAL: TensorMap Entry Pool Deadlock Detected!"); - LOG_ERROR("========================================"); - LOG_ERROR("TensorMap entry pool freed no entries for ~500 ms while a task waits."); - LOG_ERROR(" - Pool used: %d / %d", tm.current_used(), tm.pool_capacity()); - LOG_ERROR(" - Needed: %d entries", needed); - LOG_ERROR(" - last_task_alive (sum across rings): %" PRId64, cur_alive_sum); - LOG_ERROR("Diagnosis:"); - LOG_ERROR(" No retiring task is freeing tensormap entries (last_task_alive may"); - LOG_ERROR(" still move on rings with no registered outputs). Check TaskRing"); - LOG_ERROR(" diagnostics for the stalled producer."); - LOG_ERROR("Solution:"); - LOG_ERROR(" Increase PTO2_TENSORMAP_POOL_SIZE (current: %d).", tm.pool_capacity()); - LOG_ERROR("========================================"); - orch_mark_fatal(orch, PTO2_ERROR_TENSORMAP_OVERFLOW); - return false; - } - } - SPIN_WAIT_HINT(); - } - return true; -} - -// Shared body for submit_task / submit_dummy_task. Caller has already validated -// args.has_error, decided active_mask (empty for dummy), and resolved the per-slot -// kernel_ids (all INVALID_KERNEL_ID for dummy). Performs tensormap sync, fanin -// computation (explicit_deps + auto), output registration, slot init, and -// Orch-side wiring/ready publication. static TaskOutputTensors submit_task_common( PTO2OrchestratorState *orch, const L0TaskArgs &args, ActiveMask active_mask, int32_t aic_kernel_id, int32_t aiv0_kernel_id, int32_t aiv1_kernel_id ) { - CYCLE_COUNT_START(); TaskOutputTensors result; PTO2OutputLayout layout = calculate_output_layout(args); PTO2PreparedTask prepared; - if (!prepare_task(orch, args, layout.total_output_size, active_mask, &prepared)) { - return result; - } + if (!prepare_task(orch, args, layout.total_output_size, active_mask, &prepared)) return result; uint8_t ring_id = prepared.task_id.ring(); PTO2SchedulerState *sched = orch->scheduler; PTO2RingFlowControl &fc = orch->sm_header->rings[ring_id].fc; @@ -815,30 +203,12 @@ static TaskOutputTensors submit_task_common( PTO2TaskPayload &payload = *prepared.payload; result.set_task_id(task_id); - // dep_gen capture point: snapshot the orch submit_task inputs while the - // tensormap is still in its pre-lookup state for this task. Replay reads - // these records offline to reconstruct the complete dep graph — the sole - // source of truth for fanout now that the swimlane hot path no longer - // records it. -#if SIMPLER_DFX if (is_dep_gen_enabled()) { const void *tensor_ptrs[MAX_TENSOR_ARGS]; - // TensorArgType is `enum class : int32_t` (4 bytes); the on-disk record - // packs arg_types as uint8_t[16] (5-value enum fits in a byte). Narrow - // each tag here rather than letting the AICPU writer reinterpret a - // 4×-wider array as bytes — that path silently lost two of every three - // tags on little-endian and synthesized phantom self-edges in replay. uint8_t arg_types_u8[MAX_TENSOR_ARGS]; - // Clamp to MAX_TENSOR_ARGS even though the Arg builder caps adds at - // MAX_TENSOR_ARGS: defensive against any future builder bypass / - // shared-memory bit-flip that could otherwise overrun the two - // MAX_TENSOR_ARGS-sized stack buffers above. const int tc_raw = args.tensor_count(); const int tc = tc_raw > MAX_TENSOR_ARGS ? MAX_TENSOR_ARGS : tc_raw; for (int i = 0; i < tc; i++) { - // OUTPUT slots carry create_info (not yet a Tensor); skip them — - // they have no producer to look up and replay's per-tensor loop - // also skips OUTPUT. tensor_ptrs[i] = (args.tag(i) == TensorArgType::OUTPUT) ? nullptr : &args.tensor(i).ref(); arg_types_u8[i] = static_cast(args.tag(i)); } @@ -849,27 +219,12 @@ static TaskOutputTensors submit_task_common( args.launch_spec.block_num(), kernel_ids_capture ); } -#endif - - PTO2FaninBuilder fanin_builder(orch, orch->rings[ring_id].fanin_pool, next_fanin_seen_epoch(orch)); - CYCLE_COUNT_LAP(g_orch_alloc_cycle); + PTO2FaninBuilder fanin_builder; -#if SIMPLER_DFX - if (layout.total_output_size > 0) { - orch->buffers_allocated++; - orch->bytes_allocated += layout.total_output_size; - } -#endif - - // === STEP 2: Sync TensorMap validity and optional cleanup === - // Read current last_task_alive from shared memory for this ring int32_t sm_last_task_alive = fc.last_task_alive.load(std::memory_order_acquire); - orch->tensor_map.sync_tensormap(task_id, sm_last_task_alive); - CYCLE_COUNT_LAP(g_orch_sync_cycle); - for (uint32_t i = 0; i < args.explicit_dep_count(); i++) { PTO2TaskId dep_task_id = args.explicit_dep(i); if (!dep_task_id.is_valid()) { @@ -878,58 +233,30 @@ static TaskOutputTensors submit_task_common( ); return result; } - uint8_t dep_ring_id = dep_task_id.ring(); - PTO2SharedMemoryRingHeader &dep_ring = orch->sm_header->rings[dep_ring_id]; + PTO2SharedMemoryRingHeader &dep_ring = orch->sm_header->rings[dep_task_id.ring()]; int32_t dep_local_task_id = static_cast(dep_task_id.local()); int32_t dep_last_task_alive = dep_ring.fc.last_task_alive.load(std::memory_order_acquire); - if (dep_local_task_id < dep_last_task_alive) { - continue; - } - int32_t dep_slot = dep_ring.get_slot_by_task_id(dep_local_task_id); - PTO2TaskSlotState *producer_slot_state = &dep_ring.get_slot_state_by_slot(dep_slot); - if (!append_fanin_or_fail( - orch, dep_ring_id, dep_slot, producer_slot_state, dep_task_id, &fanin_builder, ring_id - )) { - return result; - } + if (dep_local_task_id < dep_last_task_alive) continue; + PTO2TaskSlotState *producer_slot_state = &dep_ring.get_slot_state_by_task_id(dep_local_task_id); + if (!append_fanin_or_fail(orch, producer_slot_state, dep_local_task_id, &fanin_builder)) return result; } - // === STEP 3: Lookup inputs (creator retention + tensormap modifier lookup) === DepInputs dep_inputs{ args.tensor_count(), args.tensor_data(), args.tag_data(), static_cast(args.explicit_dep_count()), args.explicit_deps_data(), }; auto runtime_emit = [&](PTO2TaskId producer_task_id) -> bool { - uint8_t prod_ring = producer_task_id.ring(); - PTO2SharedMemoryRingHeader &producer_ring = orch->sm_header->rings[prod_ring]; - int32_t prod_slot = producer_ring.get_slot_by_task_id(static_cast(producer_task_id.local())); - PTO2TaskSlotState *prod_state = &producer_ring.get_slot_state_by_slot(prod_slot); - return append_fanin_or_fail(orch, prod_ring, prod_slot, prod_state, producer_task_id, &fanin_builder, ring_id); + int32_t prod_local = static_cast(producer_task_id.local()); + PTO2TaskSlotState *prod_state = + &orch->sm_header->rings[producer_task_id.ring()].get_slot_state_by_task_id(prod_local); + return append_fanin_or_fail(orch, prod_state, prod_local, &fanin_builder); }; - if (!compute_task_fanin(dep_inputs, orch->tensor_map, orch->in_manual_scope(), runtime_emit)) { - return result; - } - - CYCLE_COUNT_LAP(g_orch_lookup_cycle); + if (!compute_task_fanin(dep_inputs, orch->tensor_map, orch->in_manual_scope(), runtime_emit)) return result; - // === STEP 4: Register outputs/inouts in TensorMap (must be separate from lookup) === - // Reserve pool capacity for this task's inserts before registering. The pool - // is shared across rings and reclaimed as last_task_alive advances; an - // exhausted pool back-pressures here (and detects a wedged watermark) rather - // than tripping new_entry()'s hard assert mid-registration. - int32_t tensormap_needed = count_registrable_outputs(dep_inputs, orch->in_manual_scope()); - if (tensormap_needed > 0 && !ensure_tensormap_capacity(orch, tensormap_needed)) { - return result; - } register_task_outputs(dep_inputs, task_id, orch->tensor_map, orch->in_manual_scope()); - CYCLE_COUNT_LAP(g_orch_insert_cycle); - - // === STEP 5: Batch-write to GM (single cache line burst) === - // Deferred from allocation phase to avoid scattered GM writes that get - // evicted by TensorMap lookup/insert cache pressure. __builtin_prefetch(&task, 1, 1); task.task_id = task_id; task.kernel_id[static_cast(PTO2SubtaskSlot::AIC)] = aic_kernel_id; @@ -939,125 +266,92 @@ 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; - // Dense-fanin diagnostic: fanin_builder.count is finalized here and submit - // runs once per task, so each dense consumer warns exactly once. The check - // is > THRESHOLD (not == THRESHOLD+1): the count lands at its total here, so - // an == crossing test would miss a task that jumps straight past THRESHOLD+1. - if (fanin_builder.count > PTO2_DEP_DEGREE_WARN_THRESHOLD) { - LOG_WARN( - "dense dependency: task ring=%u id=%u fanin>%d [orch submit]", static_cast(task_id.ring()), - task_id.local(), PTO2_DEP_DEGREE_WARN_THRESHOLD - ); - } - 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]; + // Single pass over fanin_builder: + // - Copy local_id/ring_id into payload so the scheduler can index the + // producer's ring's completion_flags from the consumer side. + // - Push this consumer's local_id into each same-ring producer's + // last_consumer high-water-mark, replacing the per-completion + // fanout_refcount notification. Reclamation gates on the per-ring + // completed_watermark reaching this value. Only update for same-ring + // fanin: cross-ring consumers live in a different local_id space, + // so their id is meaningless to the producer's ring's watermark. + // Cross-ring producer slots reclaim on scope_end / ring wrap instead + // — acceptable since cross-ring fanin (e.g. alloc_tensors output) + // is sparse. + // Use fanin_builder.ring_ids[i] (cache-warm SOA slice) for the same-ring + // check so cross-ring iters skip the slot_state dereference entirely. + const uint8_t self_ring = task_id.ring(); + const int32_t self_local = static_cast(task_id.local()); + payload.fanin_count = fanin_builder.count; + for (int32_t i = 0; i < fanin_builder.count; i++) { + const int32_t local = fanin_builder.local_ids[i]; + const uint8_t ring = fanin_builder.ring_ids[i]; + payload.fanin_local_ids[i] = local; + payload.fanin_ring_ids[i] = ring; + if (ring == self_ring) { + PTO2TaskSlotState *prod = fanin_builder.slots[i]; + if (self_local > prod->last_consumer_local_id) prod->last_consumer_local_id = self_local; + } } payload.init(args, result, prepared.alloc_result, layout); - cur_slot_state.set_allow_early_resolve(args.allow_early_resolve()); - - // Dispatch predicate: resolve the (tensor, indices) to an absolute GM address - // now so the scheduler can read it at the dispatch point with a single load, - // no Arg/Tensor access. Both branches write predicate.op explicitly because - // payload slots are ring-reused; op == NONE means "always dispatch". - { - const L0TaskPredicate &pred = args.predicate(); - if (pred.op != PredicateOp::NONE && pred.operand.tensor != nullptr && pred.operand.tensor->buffer.addr != 0) { - uint64_t elem_size = get_element_size(pred.operand.tensor->dtype); - uint64_t flat_offset = pred.operand.tensor->compute_flat_offset(pred.operand.indices, pred.operand.ndims); - payload.predicate.addr = pred.operand.tensor->buffer.addr + flat_offset * elem_size; - payload.predicate.target = pred.target; - payload.predicate.elem_size = static_cast(elem_size); - payload.predicate.op = pred.op; - } else { - payload.predicate.addr = 0; - payload.predicate.op = PredicateOp::NONE; - } - } -#if SIMPLER_DFX - if (is_dump_args_enabled()) { - if (args.scalar_count() > 0) { - set_dump_args_task_scalar_dtypes( - task_id.raw, static_cast(args.scalar_count()), args.scalar_dtypes() - ); - } - // Selective vs full dump is latched at dump_args_init from DumpDataHeader - // (host-decided before any dispatch), so it is race-free regardless of - // submission order. Here we only record each marked task's arg mask and - // metadata flags, which selective collection consults. - if (args.dump_arg_mask() != 0) { - set_dump_args_task_mask(task_id.raw, args.dump_arg_mask(), args.dump_arg_index_ambiguous_mask()); - } - } -#endif - - CYCLE_COUNT_LAP(g_orch_args_cycle); - - // === STEP 6: wire on the orchestrator side and publish readiness === - // Zero-fanin tasks and tasks whose claimed producers are already completed - // do not need fanout links or dep_pool entries. Tasks with live producers - // allocate fanout links here before any scheduler thread can dispatch them. - 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(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(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; - } - } - CYCLE_COUNT_LAP(g_orch_fanin_cycle); - CYCLE_COUNT_ORCH_SUBMIT_RECORD(task_id.raw); + while (!sched->wiring.queue.push(&cur_slot_state)) + SPIN_WAIT_HINT(); -#if SIMPLER_DFX - orch->tasks_submitted++; -#if SIMPLER_ORCH_PROFILING - g_orch_submit_count++; -#endif - g_orch_submit_idx++; -#endif return result; } +// ----------------------------------------------------------------------------- +// PTO2OrchestratorState members +// ----------------------------------------------------------------------------- + +void PTO2OrchestratorState::report_fatal(int32_t error_code, [[maybe_unused]] const char *func, const char *fmt, ...) { + auto *orch = this; + va_list args; + va_start(args, fmt); + orch_report_fatal_v(orch, error_code, fmt, args); + va_end(args); +} + +void PTO2OrchestratorState::begin_scope(PTO2ScopeMode mode) { + auto *orch = this; + if (orch->fatal) return; + assert(orch->scope_stack_top < static_cast(orch->scope_stack_capacity - 1) && "Scope stack overflow"); + if (mode == PTO2ScopeMode::AUTO && orch->in_manual_scope()) { + report_fatal(PTO2_ERROR_INVALID_ARGS, __FUNCTION__, "auto scope nested inside manual scope is not supported"); + return; + } + + bool already_in_manual_scope = orch->in_manual_scope(); + ++orch->scope_stack_top; + orch->scope_begins[orch->scope_stack_top] = orch->scope_tasks_size; + if (mode == PTO2ScopeMode::MANUAL && !already_in_manual_scope) orch->manual_begin_depth = orch->scope_stack_top; +} + +void PTO2OrchestratorState::end_scope() { + auto *orch = this; + if (orch->fatal) return; + assert(orch->scope_stack_top >= 0 && "Scope stack underflow"); + + bool ending_manual_scope = orch->scope_stack_top == orch->manual_begin_depth; + int32_t begin = orch->scope_begins[orch->scope_stack_top--]; + if (ending_manual_scope) orch->manual_begin_depth = PTO2_MAX_SCOPE_DEPTH; + + // Watermark-based reclamation: scope-end has no work to do — consumers + // no longer need to notify producers. + orch->scope_tasks_size = begin; +} + TaskOutputTensors PTO2OrchestratorState::submit_task(const MixedKernels &mixed_kernels, const L0TaskArgs &args) { auto *orch = this; // Orchestration API should short-circuit after fatal, but keep this entry // robust as a no-op in case a caller reaches it directly. - if (orch->fatal) { - return TaskOutputTensors{}; - } + if (orch->fatal) return TaskOutputTensors{}; // Validate Arg construction (errors recorded by add_input/add_output/etc.) if (args.has_error) { - LOG_ERROR("========================================"); - LOG_ERROR("FATAL: Invalid Arg Detected!"); - LOG_ERROR("========================================"); - LOG_ERROR("Error: %s", args.error_msg ? args.error_msg : "(unknown)"); - LOG_ERROR(" tensor_count: %d, scalar_count: %d", args.tensor_count(), args.scalar_count()); - LOG_ERROR("This is a bug in the orchestration code."); - LOG_ERROR("========================================"); orch_mark_fatal(orch, PTO2_ERROR_INVALID_ARGS); return TaskOutputTensors{}; } @@ -1069,11 +363,6 @@ TaskOutputTensors PTO2OrchestratorState::submit_task(const MixedKernels &mixed_k int16_t block_num = args.launch_spec.block_num(); always_assert(block_num >= 1 && "block_num must be >= 1"); - // Normalize single-AIV tasks: if only aiv1 is set (no aic, no aiv0), move - // it to the aiv0 slot. This guarantees the dispatch path can always use - // PTO2SubtaskSlot::AIV0 for single-AIV shapes without inspecting active_mask. - // Mixed tasks (AIC+AIV) keep their original AIV identity so the correct - // hardware channel (AIV0→AIC vs AIV1→AIC) is used at dispatch time. MixedKernels normalized = mixed_kernels; bool has_aic = active_mask.has_mask(PTO2_SUBTASK_MASK_AIC); bool has_aiv0 = active_mask.has_mask(PTO2_SUBTASK_MASK_AIV0); @@ -1086,9 +375,6 @@ TaskOutputTensors PTO2OrchestratorState::submit_task(const MixedKernels &mixed_k // Encode require_sync_start into active_mask bit 3 (only meaningful for tasks with block_num > 1) if (block_num > 1 && args.launch_spec.require_sync_start()) { - // Deadlock check: block_num >= total available slots of the required type. - // For MIX/AIC: limit is total_cluster_count (one AIC per cluster). - // For AIV: limit is total_aiv_count. PTO2ResourceShape shape = active_mask.to_shape(); int32_t limit = (shape == PTO2ResourceShape::AIV) ? orch->total_aiv_count : orch->total_cluster_count; if (limit > 0 && block_num > limit) { @@ -1101,34 +387,17 @@ TaskOutputTensors PTO2OrchestratorState::submit_task(const MixedKernels &mixed_k active_mask.set_sync_start(); } - if (args.predicate().op != PredicateOp::NONE) { - active_mask.set_has_predicate(); - } - return submit_task_common( orch, args, active_mask, normalized.aic_kernel_id, normalized.aiv0_kernel_id, normalized.aiv1_kernel_id ); } -// Submit a dependency-only task: full dependency graph participation -// (tensormap lookup/insert, explicit_deps, manual_dep, manual_scope) but no -// AICore dispatch. Empty active_mask routes the slot to the DUMMY ready -// bucket; dispatch loop short-circuits to completion. Accepts the same Arg -// shape as submit_task; scalars are permitted but never consumed. TaskOutputTensors PTO2OrchestratorState::submit_dummy_task(const L0TaskArgs &args) { auto *orch = this; - if (orch->fatal) { - return TaskOutputTensors{}; - } + if (orch->fatal) return TaskOutputTensors{}; if (args.has_error) { - LOG_ERROR("========================================"); - LOG_ERROR("FATAL: Invalid Arg in submit_dummy_task!"); - LOG_ERROR("========================================"); - LOG_ERROR("Error: %s", args.error_msg ? args.error_msg : "(unknown)"); - LOG_ERROR(" tensor_count: %d, scalar_count: %d", args.tensor_count(), args.scalar_count()); - LOG_ERROR("========================================"); orch_mark_fatal(orch, PTO2_ERROR_INVALID_ARGS); return TaskOutputTensors{}; } @@ -1141,9 +410,7 @@ TaskOutputTensors PTO2OrchestratorState::alloc_tensors(const L0TaskArgs &args) { auto *orch = this; // Orchestration API should short-circuit after fatal, but keep this entry // robust as a no-op in case a caller reaches it directly. - if (orch->fatal) { - return TaskOutputTensors{}; - } + if (orch->fatal) return TaskOutputTensors{}; if (args.tensor_count() <= 0) { report_fatal(PTO2_ERROR_INVALID_ARGS, __FUNCTION__, "alloc_tensors requires at least one TensorCreateInfo"); @@ -1162,8 +429,6 @@ TaskOutputTensors PTO2OrchestratorState::alloc_tensors(const L0TaskArgs &args) { } } - CYCLE_COUNT_START(); - if (args.has_error) { report_fatal( PTO2_ERROR_INVALID_ARGS, __FUNCTION__, "%s", @@ -1174,28 +439,17 @@ TaskOutputTensors PTO2OrchestratorState::alloc_tensors(const L0TaskArgs &args) { PTO2OutputLayout layout = calculate_output_layout(args); PTO2PreparedTask prepared; - if (!prepare_task(orch, args, layout.total_output_size, ActiveMask{}, &prepared)) { - return TaskOutputTensors{}; - } + if (!prepare_task(orch, args, layout.total_output_size, ActiveMask{}, &prepared)) return TaskOutputTensors{}; PTO2TaskDescriptor &task = *prepared.task; PTO2TaskPayload &payload = *prepared.payload; - CYCLE_COUNT_LAP(g_orch_alloc_cycle); - -#if SIMPLER_DFX - if (layout.total_output_size > 0) { - orch->buffers_allocated++; - orch->bytes_allocated += layout.total_output_size; - } -#endif - task.task_id = prepared.task_id; task.kernel_id[static_cast(PTO2SubtaskSlot::AIC)] = INVALID_KERNEL_ID; task.kernel_id[static_cast(PTO2SubtaskSlot::AIV0)] = INVALID_KERNEL_ID; task.kernel_id[static_cast(PTO2SubtaskSlot::AIV1)] = INVALID_KERNEL_ID; - // alloc_tensors builds a kernel-less descriptor that never dispatches; keep - // the slot untagged so a recycled ring slot cannot leak a stale tag. + // Kernel-less alloc descriptor never dispatches; keep the slot untagged so a + // recycled ring slot cannot leak a stale task-timing tag. task.task_timing_slot = TASK_TIMING_SLOT_NONE; task.packed_buffer_base = prepared.alloc_result.packed_base; task.packed_buffer_end = prepared.alloc_result.packed_end; @@ -1203,103 +457,41 @@ 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->rings[prepared.task_id.ring()].fanin_pool; - CYCLE_COUNT_LAP(g_orch_args_cycle); + payload.fanin_count = 0; if (prepared.slot_state != nullptr) { - // Hidden alloc tasks complete inline in the orchestrator before any - // consumer can exist, so they have no fanout to notify and no worker - // subtasks to retire. Running the full on_task_complete path - // would only pay unnecessary fanout_lock / traversal overhead here. - // The generic slot initialization done in prepare_task() is still - // required so scope_end can release the producer-side reference and - // drive the slot to CONSUMED, but worker dispatch fields are never - // observed for hidden alloc tasks. - // - // Flag the creator so it does NOT suppress its consumers' early-dispatch. - // Under the direct-only model an unflagged producer disqualifies its - // consumer, and a pre-completed producer only seeds dispatch_fanin when - // flagged. A buffer allocation is pure memory whose output is ready at - // creation — it should always be transparent, never a barrier. Unlike a - // 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(); + // (m) Inline completion uses completion_flags only. + uint8_t ring_id = prepared.task_id.ring(); + auto &ring = orch->sm_header->rings[ring_id]; + const int32_t my_id = static_cast(prepared.task_id.local()); + const int32_t mask = ring.task_window_mask; + ring.completion_flags[prepared.alloc_result.slot].store(1, std::memory_order_release); + // Inline-completed slots never reach on_mixed_task_complete, so + // CAS-advance the per-ring completed_watermark here. Without this, + // wait_for_tensor_ready(wait_for_consumers=true) on an alloc'd slot + // (e.g. set_tensor_data on its output) hangs because the watermark + // gate target (slot's own local_id) is never reached if no real + // task with local_id > my_id completes. + int32_t w = ring.completed_watermark.load(std::memory_order_acquire); + while (w < my_id) { + int32_t next = w + 1; + if (ring.completion_flags[next & 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; + } + } } orch->inline_completed_tasks++; - CYCLE_COUNT_LAP(g_orch_fanin_cycle); - CYCLE_COUNT_ORCH_SUBMIT_RECORD(prepared.task_id.raw); - -#if SIMPLER_DFX - orch->tasks_submitted++; -#if SIMPLER_ORCH_PROFILING - g_orch_submit_count++; -#endif - g_orch_submit_idx++; -#endif - return outputs; } -// ============================================================================= -// Flow Control -// ============================================================================= - void PTO2OrchestratorState::mark_done() { auto *orch = this; - for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { - int32_t total_tasks = orch->rings[r].task_allocator.active_count(); - if (total_tasks > 0) { - LOG_INFO_V0("=== [Orchestrator] ring %d: total_tasks=%d ===", r, total_tasks); - } - auto &fanin_pool = orch->rings[r].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; orch->scope_stack_top = -1; orch->manual_begin_depth = PTO2_MAX_SCOPE_DEPTH; -#if !SIMPLER_ORCH_PROFILING && SIMPLER_DFX - g_orch_submit_idx = 0; -#endif -} - -#if SIMPLER_ORCH_PROFILING -PTO2OrchProfilingData orchestrator_get_profiling() { - PTO2OrchProfilingData d; - d.sync_cycle = g_orch_sync_cycle; - d.alloc_cycle = g_orch_alloc_cycle; - d.args_cycle = g_orch_args_cycle; - d.lookup_cycle = g_orch_lookup_cycle; - d.insert_cycle = g_orch_insert_cycle; - d.fanin_cycle = g_orch_fanin_cycle; - d.scope_end_cycle = g_orch_scope_end_cycle; - d.submit_count = g_orch_submit_count; - d.alloc_wait_cycle = g_orch_alloc_wait_cycle; - d.fanin_wait_cycle = g_orch_fanin_wait_cycle; - d.alloc_atomic_count = g_orch_alloc_atomic_count; - d.args_atomic_count = g_orch_args_atomic_count; - d.scope_end_atomic_count = g_orch_scope_end_atomic_count; - - // Reset - g_orch_sync_cycle = g_orch_alloc_cycle = g_orch_args_cycle = 0; - g_orch_lookup_cycle = g_orch_insert_cycle = 0; - g_orch_fanin_cycle = g_orch_scope_end_cycle = 0; - g_orch_submit_count = 0; - g_orch_submit_idx = 0; - g_orch_alloc_wait_cycle = 0; - g_orch_fanin_wait_cycle = 0; - g_orch_alloc_atomic_count = 0; - g_orch_args_atomic_count = 0; - g_orch_scope_end_atomic_count = 0; - return d; } -#endif diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.h b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.h index 8b7f4cc320..23b7e57c25 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.h +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.h @@ -8,22 +8,6 @@ * See LICENSE in the root of the software repository for the full text of the License. * ----------------------------------------------------------------------------------------------------------- */ -/** - * PTO Runtime2 - Orchestrator Interface - * - * The Orchestrator is responsible for: - * 1. Executing the orchestration function (Turing-complete control flow) - * 2. Allocating intermediate buffers from the heap - * 3. Submitting tasks via async InCore function calls - * 4. Building the dependency graph using TensorMap - * 5. Managing buffer scopes for lifecycle control - * - * The Orchestrator can run on either: - * - Host CPU (lower latency for complex control, easier debugging) - * - Device AI_CPU (lower latency for task submission) - * - * Based on: docs/RUNTIME_LOGIC.md - */ #ifndef PTO_ORCHESTRATOR_H #define PTO_ORCHESTRATOR_H @@ -38,32 +22,55 @@ #include "pto_tensormap.h" #include "pto_types.h" -/** - * Layout descriptor produced by PTO2OrchestratorState::reserve_layout(). Holds - * arena offsets for every sub-region the orchestrator owns (per-ring fanin - * pools, scope arrays, plus the nested PTO2TensorMap layout). - */ +#include +#include +#include +#include "aicpu/dep_gen_collector_aicpu.h" +#include "common/dep_gen.h" +#include "pto_dep_compute.h" +#include "tensor.h" + +struct PTO2OrchestratorState; + +// Helper aggregate types that pto_orchestrator.cpp constructs by value while +// preparing and submitting tasks. +struct PTO2PreparedTask { + PTO2TaskId task_id = PTO2TaskId::invalid(); + PTO2TaskAllocResult alloc_result = {-1, 0, nullptr, nullptr}; + PTO2TaskDescriptor *task = nullptr; + PTO2TaskPayload *payload = nullptr; + PTO2TaskSlotState *slot_state = nullptr; +}; + +struct PTO2FaninBuilder { + int32_t count{0}; + PTO2TaskSlotState *slots[PTO2_MAX_FANIN]; + int32_t local_ids[PTO2_MAX_FANIN]; + uint8_t ring_ids[PTO2_MAX_FANIN]; + + bool contains(PTO2TaskSlotState *prod_state) const { + for (int32_t i = 0; i < count; i++) + if (slots[i] == prod_state) return true; + return false; + } +}; + struct PTO2OrchestratorLayout { - size_t off_fanin_pool[PTO2_MAX_RING_DEPTH]; - size_t off_fanin_seen_epoch[PTO2_MAX_RING_DEPTH]; size_t off_scope_tasks; size_t off_scope_begins; PTO2TensorMapLayout tensor_map; - int32_t dep_pool_capacities[PTO2_MAX_RING_DEPTH]; + int32_t dep_pool_capacity; int32_t scope_tasks_cap; uint64_t scope_stack_capacity; }; -// ============================================================================= -// Orchestrator State -// ============================================================================= - -/** - * Orchestrator state structure (private to Orchestrator) - * - * Contains all state needed for task graph construction and buffer management. - */ struct PTO2OrchestratorState { + // L2 swimlane profiling level — read by upstream aicpu_executor when + // bridging orchestrator init into the scheduler context. The polling + // design doesn't gate behavior on this directly, but the field must + // exist for the upstream code path to compile. + L2SwimlaneLevel l2_swimlane_level{L2SwimlaneLevel::DISABLED}; + // === SHARED MEMORY ACCESS === PTO2SharedMemoryHeader *sm_header; @@ -75,10 +82,6 @@ struct PTO2OrchestratorState { // === TENSOR MAP (Private) === PTO2TensorMap tensor_map; // Producer lookup - // === SCOPE STACK (Private) === - // Single contiguous buffer of task IDs, partitioned by scope level. - // scope_begins[i] is the index into scope_tasks where scope i starts. - // Tasks for the top scope occupy [scope_begins[top], scope_tasks_size). PTO2TaskSlotState **scope_tasks; // Flat buffer of taskSlotState (all scopes concatenated) int32_t scope_tasks_size; // Number of task IDs currently in the buffer int32_t scope_tasks_capacity; // Allocated capacity of scope_tasks @@ -87,45 +90,22 @@ struct PTO2OrchestratorState { uint64_t scope_stack_capacity; // Max nesting depth (PTO2_MAX_SCOPE_DEPTH) int32_t manual_begin_depth{PTO2_MAX_SCOPE_DEPTH}; - // === SCHEDULER STATE ACCESS === - // Same runtime-arena scheduler object; Orch-side wiring mutates dep_pool - // and publishes ready tasks through it before scheduler workers dispatch. - PTO2SchedulerState *scheduler; + PTO2SchedulerState *scheduler; // For simulated mode only // Total core counts set once at executor init; used for submit-time deadlock detection. int32_t total_cluster_count{0}; // AIC cores = MIX clusters int32_t total_aiv_count{0}; // AIV cores (= 2 × clusters on standard hardware) -#if SIMPLER_DFX - // L2 swimlane_level copied from get_l2_swimlane_level(). - L2SwimlaneLevel l2_swimlane_level{L2SwimlaneLevel::DISABLED}; -#endif // === GM HEAP (for output buffers) === void *gm_heap_base; // Base address of GM heap uint64_t gm_heap_size; // Total size of GM heap (all rings) - // === FATAL ERROR === - // Fatal error flag (single-thread access by orchestrator, no atomic needed) - // Cross-thread notification uses shared memory orch_error_code (atomic) bool fatal; - // Hidden alloc tasks complete synchronously inside the orchestrator and - // therefore bypass the executor's normal worker-completion counter path. - // The executor adds this count into its completed_tasks_ progress counter - // after orchestration finishes so shutdown/profiling totals remain closed. int64_t inline_completed_tasks{0}; // === STATISTICS === -#if SIMPLER_DFX - int64_t tasks_submitted; - int64_t buffers_allocated; - int64_t bytes_allocated; -#endif - - /** - * Get current ring index from scope depth. - * Maps scope depth to ring_id: min(scope_depth, PTO2_MAX_RING_DEPTH - 1) - */ + uint8_t current_ring_id() const { int32_t depth = scope_stack_top; if (depth < 0) depth = 0; @@ -134,51 +114,20 @@ struct PTO2OrchestratorState { bool in_manual_scope() const { return scope_stack_top >= manual_begin_depth; } - // === Cold-path API (defined in pto_orchestrator.cpp) === - - // Phase 1: declare every sub-region (per-ring fanin pool, scope arrays, - // 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, const int32_t task_window_sizes[PTO2_MAX_RING_DEPTH], - int32_t dep_pool_capacity = PTO2_DEP_LIST_POOL_SIZE - ); - static PTO2OrchestratorLayout reserve_layout( - DeviceArena &arena, const int32_t task_window_sizes[PTO2_MAX_RING_DEPTH], - const int32_t dep_pool_capacities[PTO2_MAX_RING_DEPTH] - ); + // === Cold-path API === - // Phase 3a: write everything *except* arena-internal pointer fields. - // sm_dev_base is the SM device address (only stored, never dereferenced); - // task_window_size feeds the per-ring SM address arithmetic. Safe to call - // on a host arena that holds the prebuilt image. + static PTO2OrchestratorLayout + reserve_layout(DeviceArena &arena, const int32_t task_window_sizes[PTO2_MAX_RING_DEPTH], int32_t dep_pool_capacity); bool init_data_from_layout( const PTO2OrchestratorLayout &layout, DeviceArena &arena, void *sm_dev_base, void *gm_heap, uint64_t heap_size, uint64_t task_window_size ); - bool init_data_from_layout( - const PTO2OrchestratorLayout &layout, DeviceArena &arena, void *sm_dev_base, void *gm_heap, - const uint64_t heap_sizes[PTO2_MAX_RING_DEPTH], const uint64_t task_window_sizes[PTO2_MAX_RING_DEPTH] - ); - bool reset_for_reuse( - const PTO2OrchestratorLayout &layout, void *sm_dev_base, void *gm_heap, - const uint64_t heap_sizes[PTO2_MAX_RING_DEPTH], const uint64_t task_window_sizes[PTO2_MAX_RING_DEPTH] - ); - - // Phase 3b: write the arena-internal pointer fields (scope_tasks, - // scope_begins, rings[].fanin_pool.base, tensor_map.{buckets,entry_pool, - // free_entry_list,task_entry_heads}, scheduler reference). - // Idempotent — host runs once on the image, AICPU runs once after attach. void wire_arena_pointers(const PTO2OrchestratorLayout &layout, DeviceArena &arena, PTO2SchedulerState *scheduler); - - // Forget pointers; arena owns the backing buffers. + void reset_for_reuse(); void destroy(); void set_scheduler(PTO2SchedulerState *scheduler); - void mark_dep_pool_position(PTO2TaskSlotState &slot_state); - void wire_fanin_task(PTO2TaskSlotState &slot_state, int32_t wfanin); void report_fatal(int32_t error_code, const char *func, const char *fmt, ...); - void begin_scope(PTO2ScopeMode mode = PTO2ScopeMode::AUTO); + void begin_scope(PTO2ScopeMode mode); void end_scope(); TaskOutputTensors submit_task(const MixedKernels &mixed_kernels, const L0TaskArgs &args); TaskOutputTensors submit_dummy_task(const L0TaskArgs &args); @@ -186,30 +135,4 @@ struct PTO2OrchestratorState { void mark_done(); }; -// ============================================================================= -// Orchestrator Profiling Data -// ============================================================================= - -#if SIMPLER_ORCH_PROFILING -struct PTO2OrchProfilingData { - uint64_t sync_cycle; - uint64_t alloc_cycle; // Combined task slot + heap allocation - uint64_t args_cycle; - uint64_t lookup_cycle; - uint64_t insert_cycle; - uint64_t fanin_cycle; - uint64_t scope_end_cycle; - 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 - // Atomic operation counts per phase - uint64_t alloc_atomic_count; - uint64_t args_atomic_count; - uint64_t scope_end_atomic_count; -}; - -PTO2OrchProfilingData orchestrator_get_profiling(); -#endif - #endif // PTO_ORCHESTRATOR_H diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_ring_buffer.cpp b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_ring_buffer.cpp index e9a4fc5c22..e3f2a30846 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_ring_buffer.cpp +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_ring_buffer.cpp @@ -8,229 +8,8 @@ * 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 a cheap cntvct_el0 read; the 1024-spin gate keeps fatal-flag - // polling and timeout bookkeeping out of every reclaim spin. - 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; - } -} - -void PTO2DepListPool::report_deadlock( - PTO2SharedMemoryRingHeader &ring, int32_t needed, int32_t last_alive, bool scope_gated -) { - int32_t current = ring.fc.current_task_index.load(std::memory_order_acquire); - LOG_ERROR("========================================"); - LOG_ERROR("FATAL: Dependency Pool Deadlock Detected!"); - LOG_ERROR("========================================"); - if (scope_gated) { - LOG_ERROR("Head task %d COMPLETED, all consumers released, scope still open ->", last_alive); - LOG_ERROR("only scope_end can free it and the orchestrator is blocked here."); - LOG_ERROR("Provable head-of-line deadlock."); - } else { - LOG_ERROR("DepListPool 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)", last_alive); - LOG_ERROR(" - current_task: %d", current); - LOG_ERROR(" - In-flight tasks: %d", current - last_alive); - // Head-task dump: what the shared reclaim watermark is actually waiting on. - if (ring.slot_states != nullptr) { - PTO2TaskSlotState &h = ring.get_slot_state_by_task_id(last_alive); - 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 - ); - } - 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("========================================"); -} - -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); - 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; - } - // (1) Structural, immediate: head task COMPLETED with every consumer - // released but its scope still open -> only scope_end frees it and a - // blocked orchestrator can never call it -> provable deadlock now. - // slot_states is null on a ring with no task-window backing. - bool head_scope_gated = false; - if (ring.slot_states != nullptr) { - PTO2TaskSlotState &head = ring.get_slot_state_by_task_id(cur_last_alive); - if (head.task_state.load(std::memory_order_acquire) == PTO2_TASK_COMPLETED) { - uint32_t rc = head.fanout_refcount.load(std::memory_order_acquire); - head_scope_gated = rc == (head.fanout_count & ~PTO2_FANOUT_SCOPE_BIT); - } - } - if (head_scope_gated) { - report_deadlock(ring, needed, cur_last_alive, /*scope_gated=*/true); - latch_pool_error(error_code_ptr, PTO2_ERROR_DEP_POOL_OVERFLOW); - return false; - } - // (2) Wall-clock backstop for the residual case the head test cannot - // prove. Absolute time (get_sys_cnt_aicpu is an MMIO read), sampled - // 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) { - report_deadlock(ring, needed, cur_last_alive, /*scope_gated=*/false); - latch_pool_error(error_code_ptr, PTO2_ERROR_DEP_POOL_OVERFLOW); - return false; - } - } - SPIN_WAIT_HINT(); - } - return true; -} +// PTO2TaskAllocator is a small, per-allocation hot-path class defined inline in +// pto_ring_buffer.h. The dependency-pool machinery that once lived here +// (PTO2FaninPool / PTO2DepListPool) is unused under the polling scheduler, so +// this translation unit carries no definitions. diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_ring_buffer.h b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_ring_buffer.h index 070074cee6..3d01f027c2 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_ring_buffer.h +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_ring_buffer.h @@ -8,28 +8,6 @@ * See LICENSE in the root of the software repository for the full text of the License. * ----------------------------------------------------------------------------------------------------------- */ -/** - * PTO Runtime2 - Ring Buffer Data Structures - * - * Implements ring buffer designs for zero-overhead memory management: - * - * 1. TaskAllocator - Unified task slot + output buffer allocation - * - Combines task ring (slot allocation) and heap ring (output buffer allocation) - * - Single spin-wait loop with unified back-pressure and deadlock detection - * - O(1) bump allocation for both task slots and heap buffers - * - * 2. FaninPool - Fanin spill entry allocation - * - Ring buffer for spilled fanin entries - * - O(1) append allocation - * - Implicit reclamation with task ring - * - * 3. DepListPool - Dependency list entry allocation - * - Ring buffer for linked list entries - * - O(1) prepend operation - * - Implicit reclamation with task ring - * - * Based on: docs/RUNTIME_LOGIC.md - */ #ifndef PTO_RING_BUFFER_H #define PTO_RING_BUFFER_H @@ -42,62 +20,34 @@ #include "pto_shared_memory.h" #include "aicpu/device_time.h" // get_sys_cnt_aicpu (deadlock wall-clock backstop) #include "common/platform_config.h" // PLATFORM_PROF_SYS_CNT_FREQ (deadlock wall-clock) -#include "common/unified_log.h" - -#if SIMPLER_DFX -// Heap-ring wrap reporting — the allocator is the only place each individual -// wrap is observable, so it notifies the scope_stats collector here. Gated: -// pays nothing (no include, no call) when profiling is compiled out. -#include "aicpu/scope_stats_collector_aicpu.h" -#endif -// Block notification interval (in spin counts) -#define PTO2_BLOCK_NOTIFY_INTERVAL 10000 -// Heap/task deadlock is detected structurally (head task COMPLETED + all -// consumers released + scope still open -> only scope_end can free it, which a -// blocked orchestrator can never reach). This wall-clock value is only a -// backstop for the residual case the structural test can't prove locally; it is -// an ABSOLUTE TIME (not a spin count), so it is stable across chips/contention. +// Heap/task deadlock backstop, expressed as an ABSOLUTE wall-clock budget of +// no-reclaim-progress (last_task_alive unchanged), NOT a spin count. The reclaim +// tail is a circular FIFO, so the oldest live task pins the whole live-set until +// its last consumer completes; when the heap sits near-full a single slow-to- +// complete task stalls reclaim for tens of ms while the scheduler catches up. A +// spin-count limit maps to wildly different wall-times under contention and +// false-trips that recoverable transient; a wall-clock deadline is stable across +// chips/load. 500 ms comfortably covers the transient (tens of ms) yet trips a +// genuine deadlock well before the 45 s op-execute timeout. Mirrors upstream/main. #define PTO2_ALLOC_DEADLOCK_TIMEOUT_CYCLES (PLATFORM_PROF_SYS_CNT_FREQ / 2) // 500 ms -// ============================================================================= -// Task Allocator (unified task slot + heap buffer allocation) -// ============================================================================= +// Check the deadlock wall clock only once per (mask+1) stalled spins: get_sys_cnt_aicpu() +// is an MMIO read, so polling it every spin would dominate the back-pressure loop. +// Must be 2^k - 1 (used as an AND mask). +static constexpr int PTO2_ALLOC_DEADLOCK_POLL_MASK = 1023; // ~every 1024 spins + +// Dep pool spin limit - if exceeded, dep pool capacity too small for workload +#define PTO2_DEP_POOL_SPIN_LIMIT 100000 -/** - * Unified task slot + heap buffer allocator. - * - * Since task and heap are always allocated together and the orchestrator is - * single-threaded, both pointers (task index, heap top) are tracked locally - * and published to shared memory via plain store — no fetch_add or CAS needed. - * - * The alloc() method checks both resources BEFORE committing to either, - * eliminating the need for rollback on partial failure. - */ class PTO2TaskAllocator { public: - /** - * Initialize the allocator with task ring and heap ring resources. - * - * All pointer arguments are device addresses (live in SM / GM heap); this - * function only stores them, no dereferences, so it is safe to invoke - * from host code that constructs a prebuilt arena image. - * - * Production callers leave `initial_local_task_id` at 0: the SM ring - * flow-control counters that current_index_ptr / last_alive_ptr point at - * start at zero (PTO2RingFlowControl::init() runs on the AICPU during SM - * reset), so we keep local_task_id_ aligned with that without reading the - * SM. Tests that drive SM state directly may pass a non-zero seed to - * exercise corner cases like task IDs near INT32_MAX. - */ void init( PTO2TaskDescriptor *descriptors, int32_t window_size, std::atomic *current_index_ptr, std::atomic *last_alive_ptr, void *heap_base, uint64_t heap_size, std::atomic *error_code_ptr, - PTO2TaskSlotState *slot_states = nullptr, int32_t initial_local_task_id = 0, uint8_t ring_id = 0 + int32_t initial_local_task_id = 0 ) { descriptors_ = descriptors; - slot_states_ = slot_states; - ring_id_ = ring_id; window_size_ = window_size; window_mask_ = window_size - 1; current_index_ptr_ = current_index_ptr; @@ -111,31 +61,28 @@ class PTO2TaskAllocator { last_alive_seen_ = 0; } - /** - * Allocate a task slot and its associated output buffer in one call. - * - * Both task index and heap top are maintained as local counters and - * published to shared memory only on success. Since the orchestrator is - * single-threaded, no CAS or fetch_add is needed — just check-then-commit. - * - * @param output_size Total packed output size in bytes (0 = no heap needed) - * @return Allocation result; check failed() for errors - */ + // Surgical reset for arena reuse: just the per-run counters. The + // arena-internal pointers (descriptors_, current_index_ptr_, etc.) are + // still valid, since wire_arena_pointers was called before this on the + // AICPU side. + void reset_for_reuse() { + local_task_id_ = 0; + heap_top_ = 0; + heap_tail_ = 0; + last_alive_seen_ = 0; + } + PTO2TaskAllocResult alloc(int32_t output_size) { uint64_t aligned_size = output_size > 0 ? PTO2_ALIGN_UP(static_cast(output_size), PTO2_ALIGN_SIZE) : 0; - int spin_count = 0; int32_t prev_last_alive = last_alive_ptr_->load(std::memory_order_acquire); int32_t last_alive = prev_last_alive; update_heap_tail(last_alive); bool blocked_on_heap = false; - uint64_t block_cycle0 = 0; // wall-clock anchor for the deadlock backstop - bool block_timing = false; // false until the first no-reclaim-progress spin -#if SIMPLER_ORCH_PROFILING - uint64_t wait_start = 0; - bool waiting = false; -#endif + uint64_t block_ts = 0; // wall-clock anchor at the first no-progress spin + bool block_timing = false; // false until reclaim stalls (last_alive frozen) + int poll_count = 0; // rate-limits the MMIO-costly sys-counter read while (true) { // Check both resources; commit only if both available @@ -143,9 +90,6 @@ class PTO2TaskAllocator { void *heap_ptr = try_bump_heap(aligned_size); if (heap_ptr) { int32_t task_id = commit_task(); -#if SIMPLER_ORCH_PROFILING - record_wait(spin_count, wait_start, waiting); -#endif return {task_id, task_id & window_mask_, heap_ptr, static_cast(heap_ptr) + aligned_size}; } blocked_on_heap = true; @@ -154,69 +98,32 @@ class PTO2TaskAllocator { } // Spin: wait for scheduler to advance last_task_alive - spin_count++; -#if SIMPLER_ORCH_PROFILING - if (!waiting) { - wait_start = get_sys_cnt_aicpu(); - waiting = true; - } -#endif last_alive = last_alive_ptr_->load(std::memory_order_acquire); update_heap_tail(last_alive); if (last_alive > prev_last_alive) { - // Reclaim advanced -> productive backpressure, not a deadlock. - spin_count = 0; + // Reclaim advanced -> productive back-pressure, not a deadlock. prev_last_alive = 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 alloc to orch_mark_fatal. Polled on the - // cold path only -- error_code_ptr_ is orch_error_code. - if (error_code_ptr_ != nullptr && error_code_ptr_->load(std::memory_order_acquire) != PTO2_ERROR_NONE) { - return {-1, -1, nullptr, nullptr}; - } - // Reclaim watermark is stuck. Run the deadlock checks only once - // per 1024 spins to keep the hot reclaim loop tight: - // get_sys_cnt_aicpu() is a cheap cntvct_el0 read, while this - // block polls the fatal flag and head_blocked_on_scope_end() - // walks the head slot (1024 spins is far below the wall-clock - // timeout, so detection latency is unaffected). - // (1) Structural, immediate: if the head task is COMPLETED with - // every consumer released but its scope still open, only - // scope_end can free it and a blocked orchestrator can never - // call it -> provable deadlock now. - if (head_blocked_on_scope_end(last_alive)) { - report_deadlock(output_size, blocked_on_heap, /*scope_gated=*/true); - return {-1, -1, nullptr, nullptr}; - } - // (2) Wall-clock backstop for the residual case the local head - // test can't prove (e.g. a closed sibling whose consumer is - // deferred). Absolute time, not a spin count. - uint64_t now = get_sys_cnt_aicpu(); + } else { + // No reclaim progress. Anchor the wall clock on the first such + // spin, then declare deadlock only after an ABSOLUTE time budget + // elapses (contention-independent). Read the sys counter once + // per 1024 spins to amortize its MMIO cost. if (!block_timing) { - block_cycle0 = now; + block_ts = get_sys_cnt_aicpu(); block_timing = true; - } else if (now - block_cycle0 >= PTO2_ALLOC_DEADLOCK_TIMEOUT_CYCLES) { - report_deadlock(output_size, blocked_on_heap, /*scope_gated=*/false); - return {-1, -1, nullptr, nullptr}; - } - if (spin_count % PTO2_BLOCK_NOTIFY_INTERVAL == 0) { - LOG_WARN( - "[TaskAllocator ring=%u] BLOCKED: tasks=%d/%d, heap_used=%" PRIu64 "/%" PRIu64 - ", heap_available=%" PRIu64 ", heap_cursor=%" PRIu64 ", on=%s, spins=%d", - static_cast(ring_id_), local_task_id_ - last_alive, window_size_, heap_used_bytes(), - heap_size_, heap_available(), heap_top_, blocked_on_heap ? "heap" : "task", spin_count - ); + poll_count = 0; + } else if (((++poll_count) & PTO2_ALLOC_DEADLOCK_POLL_MASK) == 0) { + if (get_sys_cnt_aicpu() - block_ts > PTO2_ALLOC_DEADLOCK_TIMEOUT_CYCLES) { + report_deadlock(blocked_on_heap); + return {-1, -1, nullptr, nullptr}; + } } } SPIN_WAIT_HINT(); } } - // ========================================================================= - // State queries - // ========================================================================= - int32_t active_count() const { int32_t last_alive = last_alive_ptr_->load(std::memory_order_acquire); return local_task_id_ - last_alive; @@ -252,10 +159,6 @@ class PTO2TaskAllocator { private: // --- Task Ring --- PTO2TaskDescriptor *descriptors_ = nullptr; - // Parallel to descriptors_, indexed by task_id & window_mask_. Read-only here, - // used by the deadlock detector to inspect the head task's state + fanout. - PTO2TaskSlotState *slot_states_ = nullptr; - uint8_t ring_id_ = 0; int32_t window_size_ = 0; int32_t window_mask_ = 0; std::atomic *current_index_ptr_ = nullptr; @@ -274,57 +177,24 @@ class PTO2TaskAllocator { // --- Shared --- std::atomic *error_code_ptr_ = nullptr; - // ========================================================================= - // Internal helpers - // ========================================================================= - - /** - * Commit a task slot: bump local counter and publish to shared memory. - * Must only be called after space check has passed. - */ int32_t commit_task() { int32_t task_id = local_task_id_++; current_index_ptr_->store(local_task_id_, std::memory_order_release); return task_id; } - /** - * Derive heap_tail_ from the last consumed task's packed_buffer_end. - * - * Every task has a valid packed_buffer_end (equal to packed_buffer_base - * for zero-size allocations), so the last consumed task always determines - * the correct heap_tail — no backward scan needed. - */ void update_heap_tail(int32_t last_alive) { if (last_alive <= last_alive_seen_) return; last_alive_seen_ = last_alive; PTO2TaskDescriptor &desc = descriptors_[(last_alive - 1) & window_mask_]; - uint64_t old_tail = heap_tail_; heap_tail_ = static_cast(static_cast(desc.packed_buffer_end) - static_cast(heap_base_)); -#if SIMPLER_DFX - // Reclaim pointer moves forward monotonically in ring order; a decrease - // means it wrapped past heap_size_ (occupancy < heap_size_ guarantees at - // most one wrap per call). Report it so scope_stats can unroll. - if (is_scope_stats_enabled() && heap_tail_ < old_tail) { - scope_stats_note_heap_wrap(SCOPE_STATS_HEAP_SIDE_RECLAIM); - } -#else - (void)old_tail; -#endif } - /** - * Bump the heap pointer for the given allocation size. - * Returns the allocated pointer, or nullptr if insufficient space. - * When alloc_size == 0, returns current position without advancing. - */ void *try_bump_heap(uint64_t alloc_size) { uint64_t top = heap_top_; - if (alloc_size == 0) { - return static_cast(heap_base_) + top; - } + if (alloc_size == 0) return static_cast(heap_base_) + top; uint64_t tail = heap_tail_; void *result; @@ -334,149 +204,22 @@ class PTO2TaskAllocator { result = static_cast(heap_base_) + top; heap_top_ = top + alloc_size; } else if (tail > alloc_size) { - LOG_DEBUG( - "try_bump_heap wrap-around alloc: top=%" PRIu64 ", tail=%" PRIu64 ", alloc=%" PRIu64, top, tail, - alloc_size - ); result = heap_base_; heap_top_ = alloc_size; -#if SIMPLER_DFX - // Allocation pointer just wrapped past heap_size_; report it so - // scope_stats can unroll the wrapping offset into a monotonic value. - // The collector attributes the wrap to the current scope's ring. - if (is_scope_stats_enabled()) scope_stats_note_heap_wrap(SCOPE_STATS_HEAP_SIDE_ALLOC); -#endif } else { - LOG_DEBUG( - "try_bump_heap failed (top>=tail): top=%" PRIu64 ", tail=%" PRIu64 ", alloc=%" PRIu64 - ", heap_size=%" PRIu64, - top, tail, alloc_size, heap_size_ - ); return nullptr; } + } else if (tail - top > alloc_size) { + result = static_cast(heap_base_) + top; + heap_top_ = top + alloc_size; } else { - if (tail - top > alloc_size) { - result = static_cast(heap_base_) + top; - heap_top_ = top + alloc_size; - } else { - LOG_DEBUG( - "try_bump_heap failed (top 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. - */ - 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); - } - - /** - * Report deadlock with targeted diagnostics. scope_gated == true means the - * head-of-line structural test proved it (waiting only on scope_end); - * false means the wall-clock backstop fired. - */ - void report_deadlock(int32_t requested_output_size, bool heap_blocked, bool scope_gated) { - int32_t last_alive = last_alive_ptr_->load(std::memory_order_acquire); - int32_t active_tasks = local_task_id_ - last_alive; - uint64_t htail = heap_tail_; - - LOG_ERROR("========================================"); - if (heap_blocked) { - LOG_ERROR("FATAL: Task Allocator Deadlock - Heap Exhausted! ring=%u", static_cast(ring_id_)); - } else { - LOG_ERROR("FATAL: Task Allocator Deadlock - Task Ring Full! ring=%u", static_cast(ring_id_)); - } - LOG_ERROR("========================================"); - if (scope_gated) { - LOG_ERROR("Head task %d COMPLETED, all consumers released, scope still open ->", last_alive); - LOG_ERROR("only scope_end can free it and the orchestrator is blocked here."); - LOG_ERROR("Provable head-of-line deadlock."); - } else { - LOG_ERROR( - "No reclaim progress for ~500 ms (%" PRIu64 " cycles wall clock).", - (uint64_t)PTO2_ALLOC_DEADLOCK_TIMEOUT_CYCLES - ); - } - LOG_ERROR( - " Task ring %u: current=%d, last_alive=%d, active=%d/%d (%.1f%%)", static_cast(ring_id_), - local_task_id_, last_alive, active_tasks, window_size_, 100.0 * active_tasks / window_size_ - ); - LOG_ERROR( - " Heap ring %u: top=%" PRIu64 ", tail=%" PRIu64 ", size=%" PRIu64 ", used=%" PRIu64 ", available=%" PRIu64, - static_cast(ring_id_), heap_top_, htail, heap_size_, heap_used_bytes(), heap_available() - ); - if (heap_blocked) { - LOG_ERROR(" Requested: %d bytes", requested_output_size); - } - // 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 - ); - } - LOG_ERROR("Solution:"); - if (scope_gated) { - LOG_ERROR(" The open scope's own allocation exceeds this ring. Either:"); - LOG_ERROR(" 1. Split the scope / reduce per-scope allocation (reclaim sooner), or"); - LOG_ERROR(" 2. Size the ring >= the scope's peak live-set (heap*2 may not be enough)."); - } else if (heap_blocked) { - LOG_ERROR( - " Increase heap (current: %" PRIu64 "); env PTO2_RING_HEAP= (e.g. %" PRIu64 ")", heap_size_, - heap_size_ * 2 - ); - LOG_ERROR( - " If one increase completes, it was under-provisioned; otherwise debug the stuck head consumer." - ); - } else { - LOG_ERROR( - " Increase task window (current: %d); env PTO2_RING_TASK_WINDOW= (e.g. %d)", window_size_, - active_tasks * 2 - ); - LOG_ERROR( - " If one increase completes, it was under-provisioned; otherwise debug the stuck head consumer." - ); - } - LOG_ERROR("========================================"); + void report_deadlock(bool heap_blocked) { if (error_code_ptr_) { int32_t code = heap_blocked ? PTO2_ERROR_HEAP_RING_DEADLOCK : PTO2_ERROR_FLOW_CONTROL_DEADLOCK; error_code_ptr_->store(code, std::memory_order_release); @@ -484,326 +227,8 @@ 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 reset_for_reuse(std::atomic *in_error_code_ptr) { - 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; - } - - void reset_for_reuse(std::atomic *in_error_code_ptr) { - top = 1; - tail = 1; - high_water = 0; - last_reclaimed = 0; - base[0].slot_state = nullptr; - base[0].next = nullptr; - error_code_ptr = in_error_code_ptr; - } - - /** - * Reclaim dead entries based on the 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 under pressure. The dep pool shares - * last_task_alive with the heap and task rings, so it detects a wedged - * reclaim watermark the same way PTO2TaskAllocator::alloc does: a structural - * head-of-line check plus a wall-clock backstop, each emitting report_deadlock. - */ - bool ensure_space(PTO2SharedMemoryRingHeader &ring, int32_t needed); - - /** - * Structured dep-pool deadlock report, mirroring PTO2TaskAllocator::report_deadlock. - * scope_gated marks the provable head-of-line case (head COMPLETED, all - * consumers released, scope still open) as opposed to the wall-clock backstop. - */ - void report_deadlock(PTO2SharedMemoryRingHeader &ring, int32_t needed, int32_t last_alive, bool scope_gated); - - /** - * 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) -// ============================================================================= - -/** - * Groups a TaskAllocator and DepPool into one per-depth unit. - * PTO2_MAX_RING_DEPTH instances provide independent reclamation per scope depth. - */ struct PTO2RingSet { PTO2TaskAllocator task_allocator; - PTO2FaninPool fanin_pool; }; #endif // PTO_RING_BUFFER_H diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2.cpp b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2.cpp index d48282410f..e89faddea9 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2.cpp +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2.cpp @@ -9,56 +9,19 @@ * ----------------------------------------------------------------------------------------------------------- */ -/** - * PTO Runtime2 - Main Implementation - * - * Implements the unified runtime API that combines orchestrator and scheduler. - * - * Based on: docs/RUNTIME_LOGIC.md - */ - #include "pto_runtime2.h" - #include #include #include #include - #include - #include "aicpu/device_time.h" #include "common/platform_config.h" // PLATFORM_PROF_SYS_CNT_FREQ (data-wait deadline) #include "common/unified_log.h" -#if SIMPLER_DFX #include "aicpu/scope_stats_collector_aicpu.h" -#endif - -// Weak fallback for HOST .so builds (never called, but satisfies linker). -// The AICPU build links the strong symbol from platform/.../device_time.cpp. -// Hidden visibility prevents HOST .so from polluting global symbol table. -__attribute__((weak, visibility("hidden"))) uint64_t get_sys_cnt_aicpu() { return 0; } - -// Derived here, not in pto_runtime2_types.h: that header is included by orchestrations -// that define PLATFORM_PROF_SYS_CNT_FREQ locally, so pulling the platform header into -// it caused a redefinition conflict (#1189). Scaling MS by the counter frequency (like -// SCHEDULER_TIMEOUT_CYCLES) keeps the data-wait wall-clock identical across arches. -static constexpr uint64_t PTO2_TENSOR_DATA_TIMEOUT_CYCLES = - (PTO2_TENSOR_DATA_TIMEOUT_MS * PLATFORM_PROF_SYS_CNT_FREQ) / 1000; - -// ============================================================================= -// Orchestration Ops Table (function-pointer dispatch for orchestration .so) -// ============================================================================= - -static TaskOutputTensors submit_task_impl(PTO2Runtime *rt, const MixedKernels &mixed_kernels, const L0TaskArgs &args) { - return rt->orchestrator.submit_task(mixed_kernels, args); -} - -static TaskOutputTensors alloc_tensors_impl(PTO2Runtime *rt, const L0TaskArgs &args) { - return rt->orchestrator.alloc_tensors(args); -} -static TaskOutputTensors submit_dummy_task_impl(PTO2Runtime *rt, const L0TaskArgs &args) { - return rt->orchestrator.submit_dummy_task(args); +void runtime_set_mode(PTO2Runtime *rt, PTO2RuntimeMode mode) { + if (rt) rt->mode = mode; } void rt_scope_begin(PTO2Runtime *rt) { @@ -67,12 +30,6 @@ void rt_scope_begin(PTO2Runtime *rt) { rt->orchestrator.begin_scope(mode); } -void rt_scope_end(PTO2Runtime *rt) { rt->orchestrator.end_scope(); } - -void rt_orchestration_done(PTO2Runtime *rt) { rt->orchestrator.mark_done(); } - -static bool is_fatal_impl(PTO2Runtime *rt) { return rt->orchestrator.fatal; } - void rt_report_fatal(PTO2Runtime *rt, int32_t error_code, const char *func, const char *fmt, ...) { va_list args; va_start(args, fmt); @@ -86,52 +43,35 @@ void rt_report_fatal(PTO2Runtime *rt, int32_t error_code, const char *func, cons va_end(args); } -// Wait for all producers of this tensor to be safe for data access. -// Checks owner metadata (lifecycle anchor) and OverlapMap (modifier writers). -// For reads: wait until each producer COMPLETED (done writing). -// For writes: also wait until all consumers done reading -// (consumer low bits of fanout_refcount >= consumer count, excluding the -// bit31 scope reference). -// Uses cycle-based timeout (checked every 1024 spins). -// Returns false on timeout (sets orch.fatal). MAYBE_UNINITIALIZED_BEGIN -static bool wait_for_tensor_ready(PTO2Runtime *rt, const Tensor &tensor, bool wait_for_consumers, const char *caller) { +inline bool wait_for_tensor_ready(PTO2Runtime *rt, const Tensor &tensor, bool wait_for_consumers, const char *caller) { PTO2TaskId owner = tensor.owner_task_id; PTO2OrchestratorState &orch = rt->orchestrator; - // Segmented wait: collect up to kSegmentCap producer slots, then flush by - // spinning on each. When the segment fills, we wait for the accumulated - // batch before continuing to gather more. Dedup is per-segment only; a - // producer that appears in two segments is waited on twice, which is - // idempotent (task_state is monotonic) and only adds one atomic load on - // the second encounter. constexpr int kSegmentCap = 64; const PTO2TaskSlotState *seg[kSegmentCap]; int seg_count = 0; + bool signaled = false; bool failed = false; auto wait_one_producer = [&](const PTO2TaskSlotState &slot) { uint8_t ring_id = slot.ring_id; int32_t local_id = static_cast(slot.task->task_id.local()); + auto &ring_hdr = orch.sm_header->rings[ring_id]; + const int32_t mask = ring_hdr.task_window_mask; uint64_t t0 = get_sys_cnt_aicpu(); int32_t spin_count = 0; - while (slot.task_state.load(std::memory_order_acquire) < PTO2_TASK_COMPLETED) { + // (m) Use completion_flags as the single completion signal. + while (ring_hdr.completion_flags[local_id & mask].load(std::memory_order_acquire) == 0) { SPIN_WAIT_HINT(); - if ((++spin_count & 1023) == 0) { - // A fatal latched elsewhere breaks this wait; cold path only. - if (orch.sm_header->orch_error_code.load(std::memory_order_acquire) != PTO2_ERROR_NONE) { - failed = true; - return; - } - if (get_sys_cnt_aicpu() - t0 > PTO2_TENSOR_DATA_TIMEOUT_CYCLES) { - orch.report_fatal( - PTO2_ERROR_TENSOR_WAIT_TIMEOUT, caller, - "Timeout (%llu cycles): producer (ring=%d, local=%d) not completed", - (unsigned long long)PTO2_TENSOR_DATA_TIMEOUT_CYCLES, ring_id, local_id - ); - failed = true; - return; - } + if ((++spin_count & 1023) == 0 && get_sys_cnt_aicpu() - t0 > PTO2_TENSOR_DATA_TIMEOUT_CYCLES) { + orch.report_fatal( + PTO2_ERROR_TENSOR_WAIT_TIMEOUT, caller, + "Timeout (%llu cycles): producer (ring=%d, local=%d) not completed", + (unsigned long long)PTO2_TENSOR_DATA_TIMEOUT_CYCLES, ring_id, local_id + ); + failed = true; + return; } } }; @@ -139,26 +79,23 @@ static bool wait_for_tensor_ready(PTO2Runtime *rt, const Tensor &tensor, bool wa auto wait_one_consumers = [&](const PTO2TaskSlotState &slot) { uint8_t ring_id = slot.ring_id; int32_t local_id = slot.task->task_id.local(); + // With watermark-based reclamation, "all consumers done" means the + // per-ring completed_watermark has reached this slot's recorded + // last_consumer_local_id. + PTO2SharedMemoryRingHeader &ring_hdr = rt->orchestrator.sm_header->rings[ring_id]; + int32_t target = slot.last_consumer_local_id; 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)) { + while (ring_hdr.completed_watermark.load(std::memory_order_acquire) < target) { SPIN_WAIT_HINT(); - if ((++spin_count & 1023) == 0) { - // A fatal latched elsewhere breaks this wait; cold path only. - if (orch.sm_header->orch_error_code.load(std::memory_order_acquire) != PTO2_ERROR_NONE) { - failed = true; - return; - } - if (get_sys_cnt_aicpu() - t0 > PTO2_TENSOR_DATA_TIMEOUT_CYCLES) { - orch.report_fatal( - PTO2_ERROR_TENSOR_WAIT_TIMEOUT, caller, - "Timeout (%llu cycles): consumers of producer (ring=%d, local=%d) not done", - (unsigned long long)PTO2_TENSOR_DATA_TIMEOUT_CYCLES, ring_id, local_id - ); - failed = true; - return; - } + if ((++spin_count & 1023) == 0 && get_sys_cnt_aicpu() - t0 > PTO2_TENSOR_DATA_TIMEOUT_CYCLES) { + orch.report_fatal( + PTO2_ERROR_TENSOR_WAIT_TIMEOUT, caller, + "Timeout (%llu cycles): consumers of producer (ring=%d, local=%d) not done", + (unsigned long long)PTO2_TENSOR_DATA_TIMEOUT_CYCLES, ring_id, local_id + ); + failed = true; + return; } } }; @@ -175,25 +112,26 @@ static bool wait_for_tensor_ready(PTO2Runtime *rt, const Tensor &tensor, bool wa }; auto try_push = [&](const PTO2TaskSlotState &s) { - for (int j = 0; j < seg_count; j++) { - if (seg[j] == &s) return; // per-segment dedup - } + for (int j = 0; j < seg_count; j++) + if (seg[j] == &s) return; if (seg_count == kSegmentCap) { flush_segment(); if (failed) return; } seg[seg_count++] = &s; + if (!signaled) { + orch.scheduler->wiring.orch_needs_drain.store(true, std::memory_order_release); + signaled = true; + } }; auto do_wait = [&]() { - // Step A: creator retention — read owner directly from tensor metadata if (owner.is_valid()) { auto &s = orch.sm_header->rings[owner.ring()].get_slot_state_by_task_id(owner.local()); try_push(s); if (failed) return; } - // Step B: modifier writer lookup (OverlapMap), direct callback orch.tensor_map.lookup(tensor, [&](PTO2TensorMapEntry &entry, OverlapStatus) -> bool { PTO2TaskId pid = entry.producer_task_id; auto &s = orch.sm_header->rings[pid.ring()].get_slot_state_by_task_id(pid.local()); @@ -205,22 +143,16 @@ static bool wait_for_tensor_ready(PTO2Runtime *rt, const Tensor &tensor, bool wa }; do_wait(); + if (signaled) orch.scheduler->wiring.orch_needs_drain.store(false, std::memory_order_release); return !failed; } + MAYBE_UNINITIALIZED_END -uint64_t get_tensor_data(PTO2Runtime *rt, const Tensor &tensor, uint32_t ndims, const uint32_t indices[]) { - if (tensor.buffer.addr == 0) { - unified_log_error( - __FUNCTION__, "get_tensor_data: buffer not allocated (addr=0). " - "Use the Tensor returned by add_output(TensorCreateInfo) after submit returns." - ); - return 0; - } +inline uint64_t get_tensor_data(PTO2Runtime *rt, const Tensor &tensor, uint32_t ndims, const uint32_t indices[]) { + if (tensor.buffer.addr == 0) return 0; - if (!wait_for_tensor_ready(rt, tensor, false, __FUNCTION__)) { - return 0; - } + if (!wait_for_tensor_ready(rt, tensor, false, __FUNCTION__)) return 0; uint64_t flat_offset = tensor.compute_flat_offset(indices, ndims); uint64_t elem_size = get_element_size(tensor.dtype); @@ -231,18 +163,10 @@ uint64_t get_tensor_data(PTO2Runtime *rt, const Tensor &tensor, uint32_t ndims, } void set_tensor_data(PTO2Runtime *rt, const Tensor &tensor, uint32_t ndims, const uint32_t indices[], uint64_t value) { - if (tensor.buffer.addr == 0) { - unified_log_error( - __FUNCTION__, "set_tensor_data: buffer not allocated (addr=0). " - "Use the Tensor returned by add_output(TensorCreateInfo) after submit returns." - ); - return; - } + if (tensor.buffer.addr == 0) return; // Wait for producer + all consumers before writing (WAW + WAR safety) - if (!wait_for_tensor_ready(rt, tensor, true, __FUNCTION__)) { - return; - } + if (!wait_for_tensor_ready(rt, tensor, true, __FUNCTION__)) return; uint64_t flat_offset = tensor.compute_flat_offset(indices, ndims); uint64_t elem_size = get_element_size(tensor.dtype); @@ -250,54 +174,20 @@ void set_tensor_data(PTO2Runtime *rt, const Tensor &tensor, uint32_t ndims, cons memcpy(ptr, &value, elem_size); } -// Ops-table entry that hands the call-site captured by PTO2ScopeGuard to the -// [ScopeStats] collector. The slot is always present in the struct to keep -// the layout stable; at SIMPLER_DFX=0 we fill nullptr so the orchestration -// .so's null-check skips it. -#if SIMPLER_DFX -static void scope_set_site_impl(const char *file, int line) { scope_stats_set_pending_site(file, line); } -#endif +TaskOutputTensors submit_task_impl(PTO2Runtime *rt, const MixedKernels &mixed_kernels, const L0TaskArgs &args) { + return rt->orchestrator.submit_task(mixed_kernels, args); +} -static const PTO2RuntimeOps s_runtime_ops = { - .submit_task = submit_task_impl, - .scope_begin = rt_scope_begin, - .scope_end = rt_scope_end, - .orchestration_done = rt_orchestration_done, - .is_fatal = is_fatal_impl, - .report_fatal = rt_report_fatal, - .log_error = unified_log_error, - .log_warn = unified_log_warn, - .log_debug = unified_log_debug, - .log_info_v = unified_log_info_v, - .get_tensor_data = get_tensor_data, - .set_tensor_data = set_tensor_data, - .alloc_tensors = alloc_tensors_impl, - .submit_dummy_task = submit_dummy_task_impl, -#if SIMPLER_DFX - .scope_set_site = scope_set_site_impl, -#else - .scope_set_site = nullptr, -#endif -}; +TaskOutputTensors alloc_tensors_impl(PTO2Runtime *rt, const L0TaskArgs &args) { + return rt->orchestrator.alloc_tensors(args); +} -// ============================================================================= -// Runtime Lifecycle (AICPU-only fixup) -// ============================================================================= -// -// Layout / init_data / wire / destroy live in -// runtime/shared/pto_runtime2_init.cpp so the host build can pre-populate the -// prebuilt arena image. The pieces below — wiring the ops table and the -// SPMD core counts — depend on the device-side s_runtime_ops global and the -// AICPU SchedulerContext respectively, so they remain in the AICPU build. +TaskOutputTensors submit_dummy_task_impl(PTO2Runtime *rt, const L0TaskArgs &args) { + return rt->orchestrator.submit_dummy_task(args); +} void runtime_finalize_after_wire(PTO2Runtime *rt, int32_t aic_count, int32_t aiv_count) { rt->ops = &s_runtime_ops; rt->orchestrator.total_cluster_count = aic_count; rt->orchestrator.total_aiv_count = aiv_count; } - -void runtime_set_mode(PTO2Runtime *rt, PTO2RuntimeMode mode) { - if (rt) { - rt->mode = mode; - } -} diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2.h b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2.h index a13acc4489..8c46c72160 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2.h +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2.h @@ -8,29 +8,6 @@ * See LICENSE in the root of the software repository for the full text of the License. * ----------------------------------------------------------------------------------------------------------- */ -/** - * PTO Runtime2 - Main Interface - * - * This is the main header for the PTO Runtime2 system. - * It provides a unified API for task graph construction and execution. - * - * Key Features: - * - Ring buffer based memory management (zero allocation overhead) - * - Lazy invalidation TensorMap for dependency discovery - * - Scope-based buffer lifecycle management - * - Per-task spinlocks for concurrent fanout updates - * - Orchestrator-Scheduler decoupling via shared memory - * - * Usage: - * 1. Create runtime: PTO2Runtime create methods - * 2. Build task graph in orchestration function: - * - begin_scope() / end_scope() - * - submit_task() - * 3. Mark orchestration complete: mark_done() - * 4. Destroy runtime - * - * Based on: docs/RUNTIME_LOGIC.md - */ #pragma once @@ -44,26 +21,28 @@ #include "pto_orchestrator.h" #include "aicore_completion_mailbox.h" -// ============================================================================= -// Runtime Context -// ============================================================================= +#include +#include +#include +#include "aicpu/device_time.h" +#include "common/platform_config.h" // PLATFORM_PROF_SYS_CNT_FREQ (data-wait deadline) +#include "common/unified_log.h" + +__attribute__((weak, visibility("hidden"))) uint64_t get_sys_cnt_aicpu(); + +// FREQ-scaled cycle count for the tensor-data wait timeout. Derived here, not +// in pto_runtime2_types.h: that header is included by orchestrations which +// define PLATFORM_PROF_SYS_CNT_FREQ locally, causing a redefinition conflict. +// Mirrors the upstream/main approach in pto_runtime2.cpp pre-polling-squash. +static constexpr uint64_t PTO2_TENSOR_DATA_TIMEOUT_CYCLES = + (PTO2_TENSOR_DATA_TIMEOUT_MS * PLATFORM_PROF_SYS_CNT_FREQ) / 1000; -/** - * Runtime execution mode - */ enum PTO2RuntimeMode { PTO2_MODE_EXECUTE = 0, // Execute tasks on workers PTO2_MODE_SIMULATE = 1, // Simulate task execution with cycle counting PTO2_MODE_GRAPH_ONLY = 2 // Build graph only, no execution }; -/** - * Function-pointer ops table for runtime operations. - * - * The orchestration .so calls runtime functions through this table - * (via pto_orchestration_api.h inline wrappers), so it has zero link - * dependencies on runtime .cpp files. - */ typedef struct PTO2Runtime PTO2Runtime; // forward declare for ops signatures struct PTO2RuntimeOps { @@ -74,11 +53,15 @@ struct PTO2RuntimeOps { bool (*is_fatal)(PTO2Runtime *rt); void (*report_fatal)(PTO2Runtime *rt, int32_t error_code, const char *func, const char *fmt, ...); - // Logging (populated by runtime, called by orchestration) + // Logging (populated by runtime, called by orchestration). + // ABI-aligned with pto_orchestration_api.h's PTO2RuntimeOps: log_error, + // log_warn, log_debug, log_info_v in this exact order. Mismatched layout + // here causes the orch SO to call wrong function pointers via rt->ops, + // which manifests as silent hangs in the dlopen'd orchestration code. void (*log_error)(const char *func, const char *fmt, ...); void (*log_warn)(const char *func, const char *fmt, ...); void (*log_debug)(const char *func, const char *fmt, ...); - // INFO with explicit verbosity tier (v ∈ [0,9]; gating done inside). + // INFO with explicit verbosity tier (v ∈ [0, 9]; gating done inside). void (*log_info_v)(const char *func, int v, const char *fmt, ...); // Cross-layer data access (orchestration reads/writes tensor values via runtime) @@ -129,21 +112,15 @@ struct ArenaOffsets { /** * Layout descriptor for the prebuilt runtime arena. Two named halves with * distinct lifetimes/semantics: `sizing` is the layout-defining input - * (capacities + scheduler timeout), `offsets` is the computed sub-region - * offsets + arena size. Produced once on the host by runtime_reserve_layout(); - * consumed by runtime_init_data_from_layout and runtime_wire_arena_pointers. + * (capacities), `offsets` is the computed sub-region offsets + arena size. + * Produced once on the host by runtime_reserve_layout(); consumed by + * runtime_init_data_from_layout and runtime_wire_arena_pointers. */ struct PTO2RuntimeArenaLayout { ArenaSizingKey sizing; ArenaOffsets offsets; }; -/** - * PTO Runtime2 context - * - * Contains all state for orchestration and scheduling. - * In simulated mode, runs in single process with shared address space. - */ struct PTO2Runtime { // Ops table (first field — used by orchestration .so via function pointers) const PTO2RuntimeOps *ops; @@ -166,142 +143,205 @@ struct PTO2Runtime { // Statistics int64_t total_cycles; - // Prebuilt-arena fast path metadata. Carries every offset - // wire_arena_pointers needs at AICPU boot so the AICPU can reconstruct - // all arena-internal pointer fields without re-running init_data. The - // device base of the runtime arena travels separately on the host-side - // Runtime (Runtime::prebuilt_arena_base_), since the AICPU needs it - // *before* dereferencing this image. Populated on host by - // runtime_init_data_from_layout + runtime_wire_arena_pointers; read by - // aicpu_executor.cpp. PTO2RuntimeArenaLayout prebuilt_layout; }; -// ============================================================================= -// Runtime Lifecycle API -// ============================================================================= - -/** - * Phase 1 — declare every sub-region (sm_handle wrapper, orchestrator / - * scheduler / tensor_map / mailbox / PTO2Runtime header) on the supplied - * arena. Pure arithmetic; does not touch device memory and may run on host. - * 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( +// Canonical per-ring form (matches upstream a5 signature). +inline 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] -); - -/** - * Phase 2 — write the data half of the runtime arena: standalone fields, - * memset'd arena regions, sub-structure initializers, and SM-side device - * pointers. The arena must already be committed (or attached); writes go - * into arena.base() + sub-region offsets. - * - * `sm_dev_base` / `gm_heap_dev_base` are device addresses; we only store - * them (never dereference). Safe to run on a host arena that owns a host - * mirror of the runtime image — the resulting buffer is rtMemcpy-ready. - * - * Returns the PTO2Runtime* that sits at layout.off_runtime within the arena. - * Caller must follow up with runtime_wire_arena_pointers; rt->ops and the - * AICore-side count fields are left untouched and must be filled by the - * AICPU at boot. - */ -PTO2Runtime *runtime_init_data_from_layout( - DeviceArena &arena, const PTO2RuntimeArenaLayout &layout, PTO2RuntimeMode mode, void *sm_dev_base, uint64_t sm_size, +) { + PTO2RuntimeArenaLayout layout{}; + for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { + layout.sizing.task_window_sizes[r] = task_window_sizes[r]; + layout.sizing.heap_sizes[r] = heap_sizes[r]; + layout.sizing.dep_pool_capacities[r] = dep_pool_capacities[r]; + } + + int32_t task_window_sizes_i32[PTO2_MAX_RING_DEPTH]; + for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) + task_window_sizes_i32[r] = static_cast(task_window_sizes[r]); + + layout.offsets.off_sm_handle = arena.reserve(sizeof(PTO2SharedMemoryHandle), alignof(PTO2SharedMemoryHandle)); + layout.offsets.orch = PTO2OrchestratorState::reserve_layout(arena, task_window_sizes_i32, dep_pool_capacities[0]); + layout.offsets.sched = PTO2SchedulerState::reserve_layout(arena, dep_pool_capacities[0]); + layout.offsets.off_runtime = arena.reserve(sizeof(PTO2Runtime), PTO2_ALIGN_SIZE); + layout.offsets.off_mailbox = arena.reserve(sizeof(AICoreCompletionMailbox), alignof(AICoreCompletionMailbox)); + + layout.offsets.arena_size = arena.total_size(); + return layout; +} + +// Single-size adapter: broadcasts the scalar to every ring. Defined after the +// per-ring overload so name lookup sees both at the call site. +inline PTO2RuntimeArenaLayout +runtime_reserve_layout(DeviceArena &arena, uint64_t task_window_size, int32_t dep_pool_capacity) { + uint64_t per_ring_task_window[PTO2_MAX_RING_DEPTH]; + uint64_t per_ring_heap[PTO2_MAX_RING_DEPTH]; + int32_t per_ring_dep_pool[PTO2_MAX_RING_DEPTH]; + for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { + per_ring_task_window[r] = task_window_size; + per_ring_heap[r] = 0; // Heap default; caller may set separately via runtime_init_data_from_layout. + per_ring_dep_pool[r] = dep_pool_capacity; + } + return runtime_reserve_layout(arena, per_ring_task_window, per_ring_heap, per_ring_dep_pool); +} + +inline PTO2Runtime *runtime_init_data_from_layout( + DeviceArena &arena, const PTO2RuntimeArenaLayout &layout, PTO2RuntimeMode mode, void *sm_dev_base, uint64_t, void *gm_heap_dev_base, uint64_t heap_size -); -PTO2Runtime *runtime_init_data_from_layout( +) { + PTO2Runtime *rt = static_cast(arena.region_ptr(layout.offsets.off_runtime)); + memset(rt, 0, sizeof(*rt)); + + auto *sm_wrap = static_cast(arena.region_ptr(layout.offsets.off_sm_handle)); + memset(sm_wrap, 0, sizeof(*sm_wrap)); + + // rt->ops is filled by the AICPU at boot. + rt->mode = mode; + rt->gm_heap = gm_heap_dev_base; + rt->gm_heap_size = heap_size > 0 ? heap_size * PTO2_MAX_RING_DEPTH : 0; + rt->gm_heap_owned = false; + rt->total_cycles = 0; + + if (!rt->orchestrator.init_data_from_layout( + layout.offsets.orch, arena, sm_dev_base, gm_heap_dev_base, heap_size, layout.sizing.task_window_sizes[0] + )) + return nullptr; + if (!rt->scheduler.init_data_from_layout(layout.offsets.sched, arena, sm_dev_base)) return nullptr; + + auto *mailbox = static_cast(arena.region_ptr(layout.offsets.off_mailbox)); + memset(mailbox, 0, sizeof(*mailbox)); + + return rt; +} + +// Per-ring overload (matches upstream a5 signature with sm_size + heap_sizes[]). +inline PTO2Runtime *runtime_init_data_from_layout( DeviceArena &arena, const PTO2RuntimeArenaLayout &layout, PTO2RuntimeMode mode, void *sm_dev_base, uint64_t sm_size, void *gm_heap_dev_base, const uint64_t heap_sizes[PTO2_MAX_RING_DEPTH] -); +) { + return runtime_init_data_from_layout(arena, layout, mode, sm_dev_base, sm_size, gm_heap_dev_base, heap_sizes[0]); +} -/** - * Phase 3 — wire every arena-internal pointer field (rt->sm_handle, - * rt->aicore_mailbox, orchestrator.{scope_tasks, scope_begins, scheduler, - * tensor_map.*, rings[].fanin_pool.base}, scheduler.{ready_queues, dep_pool}) - * so each holds arena.base() + offset. Idempotent — runs on - * both host (writing host-mirror addresses) and AICPU (writing device - * addresses) sides. - */ void runtime_wire_arena_pointers(DeviceArena &arena, const PTO2RuntimeArenaLayout &layout, PTO2Runtime *rt); -bool runtime_reset_for_reuse(DeviceArena &arena, const PTO2RuntimeArenaLayout &layout, PTO2Runtime *rt); -/** - * AICPU-only Phase 4 — fill in the few fields the host could not know at - * prebuilt-image build time: the ops table (s_runtime_ops is a device-side - * file-local global, host cannot resolve its device address) and the - * orchestrator's core counts (depend on the executor's scheduler context). - * Call once per boot after runtime_wire_arena_pointers. - */ -void runtime_finalize_after_wire(PTO2Runtime *rt, int32_t aic_count, int32_t aiv_count); +inline void runtime_destroy(PTO2Runtime *rt) { + // Arena buffer is pooled across runs by DeviceRunner — never freed here. + if (!rt) return; + rt->scheduler.destroy(); + rt->orchestrator.destroy(); + rt->aicore_mailbox = nullptr; + rt->sm_handle = nullptr; +} + +// Upstream-compatible overload: arena is ignored (arena lifetime is owned by +// the caller in the polling design too). +inline void runtime_destroy(PTO2Runtime *rt, DeviceArena & /*arena*/) { runtime_destroy(rt); } + +// Upstream arena-reuse path (#1234). On cache hits the host skips the +// arena re-upload, so the AICPU-side reset here is the only thing that +// scrubs the previous run's orchestrator/scheduler state. Currently +// re-runs init_data_from_layout on each sub-region followed by +// wire_arena_pointers (init_data_from_layout wipes the struct via +// *state = {}, so the wired pointers must be re-set). This adds ~2 ms of +// Device wall vs upstream's surgical reset_for_reuse; a fully surgical +// polling version is deferred as follow-up work (see the reset_for_reuse +// methods added on PTO2OrchestratorState / PTO2SchedulerState / +// PTO2TensorMap / PTO2TaskAllocator / PTO2ReadyQueue / PTO2SpscQueue for +// the scaffolding — the last-mile issue is that ready_queue's +// reset_for_reuse is a no-op and something in the surgical path leaves +// state that trips a scheduler stall on the second run). +void runtime_wire_arena_pointers(DeviceArena &arena, const PTO2RuntimeArenaLayout &layout, PTO2Runtime *rt); -/** - * Destroy runtime. With the prebuilt-arena fast path the arena buffer is - * pooled across runs by DeviceRunner, so we never call arena.release() - * here — the destructor only forgets sub-structure pointers (idempotent - * cleanup). - */ -void runtime_destroy(PTO2Runtime *rt, DeviceArena &arena); +bool runtime_reset_for_reuse(DeviceArena & /*arena*/, const PTO2RuntimeArenaLayout & /*layout*/, PTO2Runtime *rt); -/** - * Set execution mode - */ void runtime_set_mode(PTO2Runtime *rt, PTO2RuntimeMode mode); -// ============================================================================= -// Orchestration API (called by orchestration function) -// ============================================================================= - -/** - * Begin a new scope - * - * All tasks submitted within this scope will have their lifetime - * bounded by the scope. When scope_end() is called, the scope - * releases its reference to all enclosed tasks. - */ void rt_scope_begin(PTO2Runtime *rt); -/** - * End current scope - * - * Releases scope reference for all tasks submitted since scope_begin(). - * Tasks whose refcount reaches zero will have their buffers released. - */ -void rt_scope_end(PTO2Runtime *rt); +inline void rt_scope_end(PTO2Runtime *rt) { rt->orchestrator.end_scope(); } -/** - * Mark orchestration as complete - * - * Signals that no more tasks will be submitted. - */ -void rt_orchestration_done(PTO2Runtime *rt); +inline void rt_orchestration_done(PTO2Runtime *rt) { rt->orchestrator.mark_done(); } -/** - * Enter fatal state explicitly from orchestration. - */ void rt_report_fatal(PTO2Runtime *rt, int32_t error_code, const char *func, const char *fmt, ...); -/** - * Cross-layer data access: read a tensor value by waiting for its producer. - */ +// Orchestration-side logging dispatchers: orchestration .so calls +// LOG_*(fmt, ...) which routes through these ops into the unified log. +// Verbosity gates live inside the unified_log_* primitives. +inline void rt_log_error(const char *func, const char *fmt, ...) { + va_list args; + va_start(args, fmt); + char message[1024]; + vsnprintf(message, sizeof(message), fmt, args); + va_end(args); + unified_log_error(func, "%s", message); +} +inline void rt_log_warn(const char *func, const char *fmt, ...) { + va_list args; + va_start(args, fmt); + char message[1024]; + vsnprintf(message, sizeof(message), fmt, args); + va_end(args); + unified_log_warn(func, "%s", message); +} +inline void rt_log_debug(const char *func, const char *fmt, ...) { + va_list args; + va_start(args, fmt); + char message[1024]; + vsnprintf(message, sizeof(message), fmt, args); + va_end(args); + unified_log_debug(func, "%s", message); +} +inline void rt_log_info_v(const char *func, int v, const char *fmt, ...) { + va_list args; + va_start(args, fmt); + char message[1024]; + vsnprintf(message, sizeof(message), fmt, args); + va_end(args); + unified_log_info_v(func, v, "%s", message); +} + +MAYBE_UNINITIALIZED_BEGIN +bool wait_for_tensor_ready(PTO2Runtime *rt, const Tensor &tensor, bool wait_for_consumers, const char *caller); + +MAYBE_UNINITIALIZED_END + uint64_t get_tensor_data(PTO2Runtime *rt, const Tensor &tensor, uint32_t ndims, const uint32_t indices[]); -/** - * Cross-layer data access: write a value to a tensor at given indices. - * Waits for producer completion (WAW) and all consumers (WAR) via TensorMap. - * See set_tensor_data in pto_orchestration_api.h for full documentation. - */ void set_tensor_data(PTO2Runtime *rt, const Tensor &tensor, uint32_t ndims, const uint32_t indices[], uint64_t value); -/** - * Slim config struct exported by orchestration .so via aicpu_orchestration_config(). - * Shared definition with pto_orchestration_api.h (same layout, guarded). - */ +// Function-pointer ops table backing — moved from pto_runtime2.cpp so that +// the runtime_finalize_after_wire above can refer to it. + +TaskOutputTensors submit_task_impl(PTO2Runtime *rt, const MixedKernels &mixed_kernels, const L0TaskArgs &args); + +TaskOutputTensors alloc_tensors_impl(PTO2Runtime *rt, const L0TaskArgs &args); + +TaskOutputTensors submit_dummy_task_impl(PTO2Runtime *rt, const L0TaskArgs &args); + +inline bool is_fatal_impl(PTO2Runtime *rt) { return rt->orchestrator.fatal; } + +inline const PTO2RuntimeOps s_runtime_ops = { + .submit_task = submit_task_impl, + .scope_begin = rt_scope_begin, + .scope_end = rt_scope_end, + .orchestration_done = rt_orchestration_done, + .is_fatal = is_fatal_impl, + .report_fatal = rt_report_fatal, + .log_error = rt_log_error, + .log_warn = rt_log_warn, + .log_debug = rt_log_debug, + .log_info_v = rt_log_info_v, + .get_tensor_data = get_tensor_data, + .set_tensor_data = set_tensor_data, + .alloc_tensors = alloc_tensors_impl, + .submit_dummy_task = submit_dummy_task_impl, + .scope_set_site = nullptr, +}; + +void runtime_finalize_after_wire(PTO2Runtime *rt, int32_t aic_count, int32_t aiv_count); + #ifndef PTO2_ORCHESTRATION_CONFIG_DEFINED #define PTO2_ORCHESTRATION_CONFIG_DEFINED struct PTO2OrchestrationConfig { diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2_types.h b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2_types.h index ba004cf564..3a24cbef5e 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2_types.h +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2_types.h @@ -9,19 +9,6 @@ * ----------------------------------------------------------------------------------------------------------- */ -/** - * PTO Runtime2 - Core Type Definitions - * - * This header defines all fundamental types used by the PTO Runtime2 system: - * - Configuration constants - * - Worker types and task states - * - Tensor regions and task parameters - * - Task descriptors with fanin/fanout tracking - * - Dependency list entries - * - * Based on: docs/RUNTIME_LOGIC.md - */ - #ifndef SRC_A2A3_RUNTIME_TENSORMAP_AND_RINGBUFFER_RUNTIME_PTO_RUNTIME2_TYPES_H_ #define SRC_A2A3_RUNTIME_TENSORMAP_AND_RINGBUFFER_RUNTIME_PTO_RUNTIME2_TYPES_H_ @@ -40,11 +27,6 @@ #include "pto_task_id.h" #include "pto_types.h" -// Spin-wait hint for AICPU threads. On real hardware the AICPU has dedicated -// ARM A55 cores — no OS yield is needed, so the hint is a no-op. In simulation -// all threads share host CPU cores, so we yield to prevent starvation. -// This header is also compiled into the Host .so (for struct definitions only), -// where the hint is never called — the fallback no-op keeps Host builds clean. #if __has_include("spin_hint.h") #include "spin_hint.h" #else @@ -65,8 +47,7 @@ // Use pto2_task_slot(sched, task_id) for slot calculation. #define PTO2_TASK_WINDOW_SIZE 16384 // Default per-ring task window size (power of 2) -// Multi-ring: number of independent ring layers (HeapRing + TaskRing + DepPool per layer) -// Scope depth maps to ring index via: min(scope_depth, PTO2_MAX_RING_DEPTH - 1) +// Multi-ring layout: scope_depth → ring index (capped at PTO2_MAX_RING_DEPTH - 1). #define PTO2_MAX_RING_DEPTH 4 // Memory pools (per-ring defaults; total = value × PTO2_MAX_RING_DEPTH) @@ -77,11 +58,6 @@ // Scope management #define PTO2_MAX_SCOPE_DEPTH 64 // Maximum nesting depth -// Hard cap for the scope_tasks buffer. Equals the total in-flight ring slot -// budget (PTO2_TASK_WINDOW_SIZE × PTO2_MAX_RING_DEPTH): once every ring slot -// is in flight, no more tasks can ever be pushed regardless of buffer size. -// scope_tasks_push fatals on overflow rather than growing the arena-owned -// buffer (which would be UB on the arena's malloc'd backing). #define PTO2_SCOPE_TASKS_CAP (PTO2_TASK_WINDOW_SIZE * PTO2_MAX_RING_DEPTH) // Ready queue @@ -90,13 +66,14 @@ // Cross-thread early-dispatch work queue (power of two) #define PTO2_EARLY_DISPATCH_QUEUE_SIZE 64 -// Fanin storage -#define PTO2_FANIN_INLINE_CAP 64 +// Wiring queue +#define PTO2_WRIRING_QUEUE_SIZE 1024 // Per-shape queue size -// 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. -#define PTO2_DEP_DEGREE_WARN_THRESHOLD 16 +// Fanin storage — absolute max number of unique fanin dependencies per task. +// Matches the upstream/main PTO2_FANIN_INLINE_CAP so workloads that already +// fit there (qwen3_14b_decode, scalar_data_test, fanin_lookup_perf) keep +// fitting after the polling-design rewrite. +#define PTO2_MAX_FANIN 128 // TensorMap cleanup interval #define PTO2_TENSORMAP_CLEANUP_INTERVAL 64 // Cleanup every N retired tasks @@ -112,34 +89,13 @@ // a redefinition conflict. See issue #1189. constexpr uint64_t PTO2_TENSOR_DATA_TIMEOUT_MS = 15000; // 15 s -// ============================================================================= -// Task States -// ============================================================================= - -/** - * Task state enumeration - * - * State transitions: - * PENDING -> COMPLETED -> CONSUMED - * - * The slot stays in PENDING from submit through "ready in queue" and "running - * on a worker"; readiness and running-vs-idle are derived from fanin_refcount - * and per-core running_slot_state respectively, not from task_state itself. - * - * Conditions: - * PENDING->COMPLETED: all subtasks finish (set by scheduler) or task is a - * hidden alloc completed inline by the orchestrator - * COMPLETED->CONSUMED: fanout_refcount == fanout_count && state == COMPLETED - */ typedef enum { - PTO2_TASK_PENDING = 0, // Submitted; awaiting fanin, queued, or dispatched - PTO2_TASK_COMPLETED = 1, // Execution finished, output may still be in use - PTO2_TASK_CONSUMED = 2 // Output fully consumed, buffers can be released + PTO2_TASK_PENDING = 0, // Submitted; awaiting fanin, queued, or dispatched + PTO2_TASK_COMPLETED = 1 // Execution finished; per-ring completed_watermark + // advances past this slot's last_consumer_local_id + // to make its heap chunk reclaimable. } PTO2TaskState; -/** - * Result of a unified task allocation. - */ struct PTO2TaskAllocResult { int32_t task_id; // Absolute task ID (not wrapped) int32_t slot; // task_id & (window_size - 1) @@ -155,43 +111,8 @@ struct PTO2OutputLayout { int32_t total_output_size = 0; }; -// ============================================================================= -// 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 -}; - -// ============================================================================= -// Task Descriptor -// ============================================================================= - -/** - * Task descriptor structure (shared memory) - * - * Stored in the TaskDescriptor ring buffer in shared memory. - * Contains static identification and buffer pointers only. - * Dynamic scheduling state (fanin/fanout/task_state) is in PTO2TaskSlotState. - * - * Fields set by Orchestrator at submission, read by Scheduler for dispatch. - */ struct PTO2TaskDescriptor { // Mixed-task identification (encodes ring_id in upper 32 bits) PTO2TaskId task_id; // raw: (ring_id << 32) | local_id @@ -228,9 +149,8 @@ static_assert(offsetof(PTO2TaskDescriptor, packed_buffer_base) == 24, "packed_bu /** * Task payload data (cold path - only accessed during orchestration and dispatch) * - * Layout: metadata + inline fanin packed in the first 9 cache lines, followed - * by bulk tensor and scalar data. Small fanins stay fully inline; larger - * fanins spill into a per-ring ring buffer slice. + * Layout: metadata + flat fanin_local_ids[] in the first 2 cache lines, + * followed by bulk tensor and scalar data. */ // Early-dispatch claim states for PTO2TaskPayload::early_dispatch_state. enum PTO2EarlyDispatchState : uint8_t { @@ -264,82 +184,78 @@ enum PTO2EarlySyncDrainState : uint8_t { inline constexpr int PTO2_EARLY_DISPATCH_CORE_MASK_WORDS = 2; struct PTO2TaskPayload { - // === Cache lines 0-8 (576B) — metadata + inline fanin === + // === Cache lines 0-2 (192B) — metadata + fanin (wireless model) === 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]; - // 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 - // cache line 8 it shares only with the rarely-touched fanin tail, whereas in - // line 0 the early-dispatch atomics (written during staging) would false-share with - // tensor_count/scalar_count (read by build_payload at dispatch). Fits in the 40B - // between the fanin array (offset 536) and the 64B-aligned tensors[] (offset - // 576), so sizeof and tensors[] are unchanged. - // + // wireless: flat fanin_local_ids[] populated at submit. The thread-0 + // pending poll indexes a compact ring-level completion_flags byte array + // via these ids — avoids a pointer chase per fanin into a 128B-aligned + // slot_state. + int32_t fanin_count{0}; + // Real producer-edge count, the early-dispatch target (without the wiring + // +1 that fanin_count/ready_fanin carry). A consumer whose dispatch_fanin + // reaches fanin_actual_count has every producer flagged-and-fully-published + // or pre-completed => it is an early-dispatch candidate. + int32_t fanin_actual_count{0}; + int32_t fanin_local_ids[PTO2_MAX_FANIN]; + // Parallel array: producer's ring_id for each fanin edge. With multi-ring + // (PTO2_MAX_RING_DEPTH > 1), the consumer's pending poll must read the + // producer's ring's completion_flags — same-ring lookup is no longer a + // safe shortcut. Sized as bytes to stay cheap (PTO2_MAX_FANIN bytes). + uint8_t fanin_ring_ids[PTO2_MAX_FANIN]; + // === Early-dispatch metadata (AICPU-side only) === // Bitmask of global core_ids this consumer is pre-staged (gated) on. Concurrent // stagers publish bits with atomic fetch_or. A regular consumer destructively // splits them between release and late-stager owners; a sync_start drain keeps // the completed mask stable for its single cohort launch owner. std::atomic staged_core_mask[PTO2_EARLY_DISPATCH_CORE_MASK_WORDS]{}; - // Early-dispatch CANDIDATE detection (event-driven, dual of fanin_refcount): + // Early-dispatch CANDIDATE detection (event-driven, dual of the ready fanin): // seeded at wiring with producers already complete, then a flagged producer // bumps each consumer after all of its logical blocks are published. - // dispatch_fanin == fanin_actual_count <=> every producer is - // flagged-and-fully-published or was - // pre-completed => this task is an early-dispatch candidate (push early_dispatch_queues[shape]). + // dispatch_fanin == fanin_actual_count => early-dispatch candidate. std::atomic dispatch_fanin{0}; // CONSUMER side: fully-published + pre-completed producers // Number of logical blocks whose payloads and MMIO tokens are published. - // Claimed-but-unpublished blocks do not make a producer launch-visible. Its - // seq_cst updates pair with early_dispatch_state to avoid losing the final - // publish vs. release wakeup for a pre-staged producer. std::atomic published_block_count{0}; - // Lock-free claim state shared by the stagers (Hook 1, possibly several AICPU - // threads concurrently) and the completion-path release: 0=NONE, 1=STAGING, - // 3=DISPATCHED (2=STAGED is unused now). STAGING is the STABLE gated state — - // many threads stage blocks concurrently while it holds, each claiming a block - // via the atomic next_block_idx and OR-ing its cores into staged_core_mask. - // Release does STAGING->DISPATCHED. For a regular consumer it claims the current - // mask and a late stager rings only its remaining bits. A sync_start consumer - // preserves the mask for rendezvous counting and its single launch pass. + // Lock-free claim state shared by the stagers and the completion-path release: + // 0=NONE, 1=STAGING (stable gated), 3=DISPATCHED (2=STAGED unused). STAGING + // holds while many threads stage blocks concurrently, each claiming via the + // atomic next_block_idx and OR-ing its cores into staged_core_mask. std::atomic early_dispatch_state{0}; + std::atomic dispatch_propagated{0}; // PRODUCER side: once-guard for fanout propagation // The launch owner publishes COMPLETE only after all owned doorbells are // visible, keeping fanout private until every gated block has launched. std::atomic early_dispatch_launch_state{PTO2_EARLY_DISPATCH_LAUNCH_NONE}; - // sync_start early-dispatch rendezvous: count of this task's gated CORES currently - // occupying a RUNNING slot (staged directly to an idle core, or promoted from a - // gated pending slot). Counted per-core (not per-block) so it is shape-agnostic: a - // MIX block spans a cluster whose cores promote independently. A sync_start task's - // doorbells are rung only once this reaches popcount(staged_core_mask) AND the - // producer released, so all cores launch atomically. Unused (0) for non-sync_start. + // sync_start rendezvous: count of this task's gated CORES currently occupying a + // RUNNING slot. Doorbells ring only once this reaches popcount(staged_core_mask) + // AND the producer released, so all cores launch atomically. Counted per-core + // (shape-agnostic; a MIX block spans a cluster). Unused (0) for non-sync_start. std::atomic running_slot_count{0}; // Ownership handshake between the early sync queue and final ready routing. // A successful OWNER persists through ARMED and COMPLETE until payload - // reinitialization. READY records that producer release observed OWNER; - // only cancellation clears OWNER during the current task lifetime. + // reinitialization; only cancellation clears OWNER during a task lifetime. std::atomic early_sync_drain_state{PTO2_EARLY_SYNC_DRAIN_NONE}; - // === Cache line 9 (byte 576) — dispatch predicate (AICPU-only) === - // Offset is a fixed 576, independent of MAX_TENSOR_ARGS / MAX_SCALAR_ARGS. - // AICore never reads it — args are materialized from the tensor_count / tensors - // / scalars offsets only. Resolved at submit; evaluated by the scheduler at - // dispatch. + // Dispatch predicate (AICPU-only), cache-line aligned at its natural offset. + // With PTO2_MAX_FANIN=128 the inline wireless fanin arrays exceed upstream's + // fixed 576/640 metadata budget, so this payload does NOT preserve upstream's + // predicate@576 / tensors@640 offsets. That AICore fixed-offset arg- + // materialization contract only matters on the early-dispatch gated path; + // it is re-established (e.g. by moving fanin to a side array) when early + // dispatch is wired. Normal dispatch reads tensors via struct members only. alignas(64) DispatchPredicate predicate; - // === Cache lines 10-73 (4096B) — tensors (alignas(64) forces alignment) === + // === Tensors (alignas(64) forces alignment) === Tensor tensors[MAX_TENSOR_ARGS]; - // === Cache lines 74-75 (128B) — scalars === + // === Scalars === uint64_t scalars[MAX_SCALAR_ARGS]; - // Layout verification (size checks that don't need offsetof). static_assert(sizeof(Tensor) == 128, "Tensor must be 2 cache lines"); - static_assert(MAX_SCALAR_ARGS * sizeof(uint64_t) == 128, "scalar region must be 128B (2 cache lines)"); + static_assert( + MAX_SCALAR_ARGS * sizeof(uint64_t) == MAX_SCALAR_ARGS * 8, "scalar region size matches MAX_SCALAR_ARGS" + ); /** * Prefetch (for write) the regions init() is about to fill so the stores land * in warm cache. tensor_count/scalar_count come from the Arg — the payload's - * own counts are not set until init(). Warms the early-dispatch block at + * own counts are not set until init(). Warms the early-dispatch spec block at * offset 536 (cache line 8) too. A member fn lowers to the same prefetch * instructions as a free function (`this` is just a register), no cache impact. */ @@ -354,7 +270,6 @@ struct PTO2TaskPayload { __builtin_prefetch(this, 1, 3); __builtin_prefetch(reinterpret_cast(this) + 64, 1, 3); __builtin_prefetch(reinterpret_cast(this) + 128, 1, 3); - __builtin_prefetch(reinterpret_cast(this) + 512, 1, 3); // early-dispatch fields (cache line 8) } /** @@ -391,23 +306,17 @@ struct PTO2TaskPayload { // Eliminates branches; extra bytes within the same CL have zero additional cost. memcpy(scalars, args.scalars(), PTO2_ALIGN_UP(args.scalar_count() * sizeof(uint64_t), 64)); - // Early-dispatch metadata — the single init point for these - // fields. reset_for_reuse MUST NOT touch the payload (it runs on the - // scheduler's advance-ring path and would pull this cold cache line across - // structures); prepare_task only allocates/binds. prefetch() warms this - // line (offset 512) so these writes land in warm cache. - // - // early_dispatch_state / staged_core_mask / dispatch_fanin are all CONSUMER-side: a - // task whose own allow_early_resolve is false still has them touched when + // Early-dispatch metadata — the single init point for these fields. + // early_dispatch_state / staged_core_mask / dispatch_fanin are CONSUMER-side: + // a task whose own allow_early_resolve is false still has them touched when // one of ITS producers is flagged (propagate_dispatch_fanin bumps - // dispatch_fanin and may CAS early_dispatch_state on any consumer, independent of the - // consumer's own hint). So they MUST be zeroed here unconditionally. - // Publication and launch fields share this same per-submit lifetime - // and are reset here too. + // dispatch_fanin and may CAS early_dispatch_state on any consumer). So they + // MUST be zeroed here unconditionally, once per submit. early_dispatch_state.store(PTO2_EARLY_DISPATCH_NONE, std::memory_order_relaxed); for (int w = 0; w < PTO2_EARLY_DISPATCH_CORE_MASK_WORDS; w++) staged_core_mask[w].store(0, std::memory_order_relaxed); dispatch_fanin.store(0, std::memory_order_relaxed); + dispatch_propagated.store(0, std::memory_order_relaxed); published_block_count.store(0, std::memory_order_relaxed); early_dispatch_launch_state.store(PTO2_EARLY_DISPATCH_LAUNCH_NONE, std::memory_order_relaxed); running_slot_count.store(0, std::memory_order_relaxed); @@ -416,103 +325,51 @@ 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, predicate) == 576, - "dispatch predicate occupies cache line 9 at fixed byte 576 (before tensors, never moves)" -); +// Fused wireless-fanin + early-dispatch layout: flat fanin arrays (sized by +// PTO2_MAX_FANIN) + early-dispatch metadata + predicate, then bulk tensor/scalar +// data. Offsets are relative (not the upstream fixed 576/640): with +// PTO2_MAX_FANIN=128 the fanin arrays alone exceed 576B, so upstream's fixed +// AICore arg-materialization contract is re-established only when early dispatch +// is wired. tensors/scalars are naturally cache-line aligned and packed. static_assert( - offsetof(PTO2TaskPayload, tensors) == 640, "tensors must start at byte 640 (cache line 10, after predicate)" + offsetof(PTO2TaskPayload, fanin_local_ids) == 16, "fanin array must follow the metadata + fanin count words" ); static_assert( - offsetof(PTO2TaskPayload, scalars) == 640 + MAX_TENSOR_ARGS * sizeof(Tensor), + offsetof(PTO2TaskPayload, scalars) == offsetof(PTO2TaskPayload, tensors) + MAX_TENSOR_ARGS * sizeof(Tensor), "scalars must immediately follow tensors" ); static_assert( - sizeof(PTO2TaskPayload) == 640 + MAX_TENSOR_ARGS * sizeof(Tensor) + MAX_SCALAR_ARGS * sizeof(uint64_t), - "PTO2TaskPayload size = metadata(576) + predicate cache line(64) + tensors + scalars" + sizeof(PTO2TaskPayload) == offsetof(PTO2TaskPayload, scalars) + MAX_SCALAR_ARGS * sizeof(uint64_t), + "no trailing padding after scalars" ); -/** - * 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. - * - * 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 - */ - -// 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, - PTO2_DISPATCH_PROPAGATED = 1U << 3, -}; - -static_assert((PTO2_DISPATCH_PROPAGATED & (PTO2_READY_CLAIMED | PTO2_COMPLETION_DONE | PTO2_SUBTASK_DEFERRED)) == 0); - 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 + // Highest local task id among this slot's consumers. Set to this slot's + // own local_id in prepare_task; bumped via max() in submit_task_common for + // each consumer that has this slot as a fanin. The slot's heap chunk is + // safe to reclaim when the per-ring completed_watermark reaches at least + // this id (i.e. every task up to and including the last consumer has + // transitioned to COMPLETED). Single-writer (orchestrator) at submit time. + int32_t last_consumer_local_id; - // 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 (accessed with fanout_count in check_and_handle_consumed) - std::atomic fanout_refcount; // Dynamic: low bits = released consumers, bit31 = scope released - - // --- 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; + // --- (e) Wake-list: lightweight last-fanin notification --- + // When a pending consumer's fanin scan finds exactly ONE unmet fanin, + // it registers itself on the producer's wake list (CAS push). On producer + // completion, the producer atomic-exchanges wake_list_head to the + // SENTINEL value and pushes every waiter to the ready queues. Consumers + // that observe SENTINEL during registration push themselves directly + // (producer already completed). Reset to nullptr on slot reuse. + 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) uint8_t ring_id; // Ring layer (immutable after init) - // 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. - bool allow_early_resolve{false}; - // Concurrent lifecycle updates preserve unrelated bits. Slot reuse clears - // the byte only after the previous task lifetime is quiescent. - std::atomic lifecycle_flags{PTO2_LIFECYCLE_FLAGS_NONE}; - int32_t dep_pool_mark{0}; // Dep pool top after Orch-side wiring + 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) @@ -545,130 +402,27 @@ struct alignas(64) PTO2TaskSlotState { */ void bind_ring(uint8_t rid) { ring_id = rid; } - /** - * 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; } - // Lock-free callers use this only as a fast-path hint. A false result is - // rechecked by try_mark_dispatch_propagated() while holding fanout_lock. - bool has_dispatch_propagated() const { - return (lifecycle_flags.load(std::memory_order_acquire) & PTO2_DISPATCH_PROPAGATED) != 0; - } - - // The propagation owner holds fanout_lock from this claim through its - // fanout snapshot so wiring can classify late edges exactly once. - bool try_mark_dispatch_propagated() { - return (lifecycle_flags.fetch_or(PTO2_DISPATCH_PROPAGATED, std::memory_order_acq_rel) & - PTO2_DISPATCH_PROPAGATED) == 0; - } - - void mark_completed() { - task_state.store(PTO2_TASK_COMPLETED, std::memory_order_release); - lifecycle_flags.fetch_or(PTO2_COMPLETION_DONE, 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); } - - bool has_any_subtask_deferred() const { - return (lifecycle_flags.load(std::memory_order_acquire) & PTO2_SUBTASK_DEFERRED) != 0; - } - - void set_allow_early_resolve(bool v) { allow_early_resolve = v; } - - /** - * Reset dynamic scheduling fields for slot reuse. - * Called by advance_ring_pointers() after a slot transitions to CONSUMED - * and last_task_alive advances past it, but before sync_to_sm() publishes - * the new last_task_alive to the orchestrator. - * - * Skips payload, task, ring_id (immutable, bound once at init). - * Skips task_state: left as CONSUMED so that wait_for_tensor_ready() - * callers holding stale owner_task_id still observe a completed state. - * task_state is set to PENDING by the orchestrator when it reuses the slot. - */ 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); 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; - } - } + any_subtask_deferred.store(false, std::memory_order_relaxed); + // (e) Wake list: clear for the next incarnation. Previous incarnation + // left it at WAKE_LIST_SENTINEL (set by its on_mixed_task_complete). + wake_list_head.store(nullptr, std::memory_order_relaxed); + next_in_wake_list = nullptr; + // last_consumer_local_id is reset in prepare_task once the task_id is known. } - - void unlock_fanout() { fanout_lock.store(0, std::memory_order_release); } }; -static_assert(sizeof(PTO2TaskSlotState) == 64); +// (e) 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(uintptr_t{1}); + +static_assert(sizeof(PTO2TaskSlotState) <= 128, "slot state should fit in two cache lines"); #endif // SRC_A2A3_RUNTIME_TENSORMAP_AND_RINGBUFFER_RUNTIME_PTO_RUNTIME2_TYPES_H_ diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_shared_memory.h b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_shared_memory.h index f1058675de..415749f490 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_shared_memory.h +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_shared_memory.h @@ -8,49 +8,14 @@ * See LICENSE in the root of the software repository for the full text of the License. * ----------------------------------------------------------------------------------------------------------- */ -/** - * PTO Runtime2 - Shared Memory Layout - * - * Defines the shared memory structure for Orchestrator-Scheduler communication. - * - * Memory Layout (per-ring sections repeat for each ring 0..PTO2_MAX_RING_DEPTH-1): - * +---------------------------+ - * | SharedMemoryHeader | (per-ring flow control + sync) - * +---------------------------+ - * | Ring 0: TaskDescriptor[] | - * | Ring 0: TaskPayload[] | - * | Ring 0: TaskSlotState[] | - * +---------------------------+ - * | Ring 1: TaskDescriptor[] | - * | Ring 1: TaskPayload[] | - * | Ring 1: TaskSlotState[] | - * +---------------------------+ - * | ... | - * +---------------------------+ - * - * Design principles: - * - Only data needed for Orchestrator<->Scheduler communication is here - * - TensorMap, scope_stack, ready_queues, dep_pool are in private memory - * - Flow control via atomic counters/flags (no locks needed for single-word R/W) - * - * Based on: docs/RUNTIME_LOGIC.md - */ #pragma once #include "utils/device_arena.h" #include "pto_runtime2_types.h" -// ============================================================================= -// Shared Memory Header -// ============================================================================= - struct PTO2SharedMemoryHandle; -/** - * Per-ring flow control state in shared memory. - * Written/read by Orchestrator and Scheduler for synchronization. - */ struct alignas(64) PTO2RingFlowControl { // === Cache Line 0: Written by Orchestrator, Read by Scheduler === alignas(64) std::atomic current_task_index; // Task ring head (next to allocate) @@ -58,13 +23,6 @@ struct alignas(64) PTO2RingFlowControl { // === Cache Line 1: Written by Scheduler, Read by Orchestrator (for back-pressure) === alignas(64) std::atomic last_task_alive; // Task ring tail (oldest active task) - // Per-boot SM reset. PTO2TaskAllocator::init() seeds its private - // local_task_id_ from initial_local_task_id (default 0 in production) - // *without* dereferencing current_task_index — it relies on this reset - // running on every AICPU boot so 0 stays in sync. If you ever change - // the initial fc value or the boot ordering, update the default in - // PTO2TaskAllocator::init (pto_ring_buffer.h) in the same change, or - // submit IDs will be off by the divergence. void init() { current_task_index.store(0, std::memory_order_relaxed); last_task_alive.store(0, std::memory_order_relaxed); @@ -75,15 +33,15 @@ struct alignas(64) PTO2RingFlowControl { static_assert(sizeof(PTO2RingFlowControl) == 128, "PTO2RingFlowControl must be exactly 2 cache lines (128B)"); -/** - * Per-ring shared memory header section. - * - * Groups flow-control, layout info, and per-ring data pointers for a single ring. - * Pointers are host-side only (set by setup_pointers, invalid on device). - */ struct alignas(64) PTO2SharedMemoryRingHeader { PTO2RingFlowControl fc; + // Highest task_id such that every task with id in [0, completed_watermark] + // has reached COMPLETED. Maintained at task-completion time. Used to gate + // slot reclamation: a producer slot P is safe to retire when + // completed_watermark >= P.last_consumer_local_id. + alignas(64) std::atomic completed_watermark; + // Layout metadata (set once at init) uint64_t task_window_size; int32_t task_window_mask; @@ -95,30 +53,28 @@ struct alignas(64) PTO2SharedMemoryRingHeader { PTO2TaskPayload *task_payloads; PTO2TaskSlotState *slot_states; - int32_t get_slot_by_task_id(int32_t local_task_id) { return local_task_id & task_window_mask; } + // Compact contiguous array (one byte per slot) holding the polling-fast + // "task X completed?" flag. 0 = pending, 1 = completed. Indexed by + // local_id & task_window_mask. Writer: the task's completer at + // on_mixed_task_complete; Resetter: orchestrator in prepare_task for the + // newly-allocated slot. Reader: thread-0 fanin polling. Replaces a chain + // of 128B-aligned slot_state pointer derefs with byte reads into a single + // array — typically condenses 16 fanin checks into 1-2 cache lines. + std::atomic *completion_flags; PTO2TaskDescriptor &get_task_by_slot(int32_t slot) { return task_descriptors[slot]; } - PTO2TaskDescriptor &get_task_by_task_id(int32_t local_id) { - return task_descriptors[get_slot_by_task_id(local_id)]; - } + PTO2TaskDescriptor &get_task_by_task_id(int32_t local_id) { return task_descriptors[local_id & task_window_mask]; } PTO2TaskPayload &get_payload_by_slot(int32_t slot) { return task_payloads[slot]; } - PTO2TaskPayload &get_payload_by_task_id(int32_t local_id) { return task_payloads[get_slot_by_task_id(local_id)]; } + PTO2TaskPayload &get_payload_by_task_id(int32_t local_id) { return task_payloads[local_id & task_window_mask]; } PTO2TaskSlotState &get_slot_state_by_slot(int32_t slot) { return slot_states[slot]; } - PTO2TaskSlotState &get_slot_state_by_task_id(int32_t local_id) { - return slot_states[get_slot_by_task_id(local_id)]; - } + PTO2TaskSlotState &get_slot_state_by_task_id(int32_t local_id) { return slot_states[local_id & task_window_mask]; } }; -/** - * Shared memory header structure - * - * Contains per-ring flow control and global layout information. - */ struct alignas(PTO2_ALIGN_SIZE) PTO2SharedMemoryHeader { // === PER-RING FLOW CONTROL + LAYOUT INFO (set once at init) === PTO2SharedMemoryRingHeader rings[PTO2_MAX_RING_DEPTH]; @@ -167,14 +123,6 @@ static_assert( "PTO2SharedMemoryHeader should be reasonably sized" ); -// ============================================================================= -// Shared Memory Handle -// ============================================================================= - -/** - * Handle for shared memory lifecycle management (create/destroy). - * Runtime components (orchestrator, scheduler) use PTO2SharedMemoryHeader* directly. - */ struct PTO2SharedMemoryHandle { void *sm_base; // Base address of shared memory uint64_t sm_size; // Total size of shared memory @@ -182,58 +130,47 @@ struct PTO2SharedMemoryHandle { PTO2SharedMemoryHeader *header; // Ownership flag - bool is_owner; // True if this handle allocated the memory + bool is_owner; + // True if this handle allocated the memory // === Static helpers === static uint64_t calculate_size(uint64_t task_window_size); + static uint64_t calculate_size_per_ring(const uint64_t task_window_sizes[PTO2_MAX_RING_DEPTH]); - // UT convenience: reserve wrapper + sm_base on `arena`, commit, and init - // using default PTO2_TASK_WINDOW_SIZE / PTO2_HEAP_SIZE. Only valid when the - // arena is otherwise empty (the call performs the single commit). All - // memory is owned by the arena — caller must not call destroy(). static PTO2SharedMemoryHandle *create_and_init_default(DeviceArena &arena); // === Instance methods === - // In-place init for caller-provided wrapper storage (e.g. a region carved - // out of a DeviceArena). Sets is_owner = false, calls setup_pointers and - // init_header. Returns false when `sm_size` is too small for the requested - // `task_window_size`. - bool init(void *sm_base, uint64_t sm_size, uint64_t task_window_size, uint64_t heap_size); + bool init(void *sm_base_arg, uint64_t sm_size_arg, uint64_t task_window_size, uint64_t heap_size); + + // Per-ring init adapter (upstream signature). Polling-side init treats + // task_window_sizes[0] as canonical; rings 1..N inherit. heap_sizes[0] is + // passed to the per-ring header init below. bool init_per_ring( - void *sm_base, uint64_t sm_size, const uint64_t task_window_sizes[PTO2_MAX_RING_DEPTH], + void *sm_base_arg, uint64_t sm_size_arg, const uint64_t task_window_sizes[PTO2_MAX_RING_DEPTH], const uint64_t heap_sizes[PTO2_MAX_RING_DEPTH] ); void destroy(); + void print_layout(); + bool validate(); private: void init_header(uint64_t task_window_size, uint64_t heap_size); + void init_header_per_ring( const uint64_t task_window_sizes[PTO2_MAX_RING_DEPTH], const uint64_t heap_sizes[PTO2_MAX_RING_DEPTH] ); + void setup_pointers(uint64_t task_window_size); + void setup_pointers_per_ring(const uint64_t task_window_sizes[PTO2_MAX_RING_DEPTH]); }; -// ============================================================================= -// SM Device Layout Helpers -// ============================================================================= -// -// When the host pre-builds a runtime-arena image, it needs the device-side -// addresses of several SM sub-fields (ring flow-control counters, -// task_descriptors arrays, orch_error_code) so it can wire them into the -// orchestrator / scheduler init_data path without dereferencing the SM — -// the SM lives in device memory and cannot be touched from host. -// -// These helpers compute those addresses by offset arithmetic on the SM -// device base. Pure pointer math, no loads/stores; safe to call from host. -// The same arithmetic happens on AICPU too (via PTO2SharedMemoryHandle's -// own setup_pointers), so values are guaranteed consistent across sides. namespace pto2_sm_layout { inline std::atomic *orch_error_code_addr(void *sm_dev_base) noexcept { @@ -263,58 +200,23 @@ inline std::atomic *ring_last_task_alive_addr(void *sm_dev_base, int ri ); } -// Byte offsets (from the SM base) of one ring's three segments. The per-ring -// layout is: header, then for each ring descriptors -> payloads -> slot_states, -// every segment PTO2_ALIGN_UP-padded. -struct PTO2RingSegmentOffsets { - uint64_t descriptors; - uint64_t payloads; - uint64_t slot_states; - uint64_t end; // offset just past this ring's slot_states (next ring's start; total SM size for the last ring) -}; - -// Single source of truth for the per-ring SM layout. Returns offsets (not -// pointers), so it serves BOTH the host-side pointer setup -// (`setup_pointers_per_ring`, which adds `sm_base`) and the device-address -// helpers below (which add `sm_dev_base`). Adding or reordering a per-ring -// segment is a one-line edit here; every consumer follows automatically, so the -// layout walk can never silently disagree across call sites. -inline PTO2RingSegmentOffsets -ring_segment_offsets(const uint64_t task_window_sizes[PTO2_MAX_RING_DEPTH], int ring_id) noexcept { - assert(ring_id >= 0 && ring_id < PTO2_MAX_RING_DEPTH && "pto2_sm_layout: ring_id out of range"); - uint64_t off = PTO2_ALIGN_UP(sizeof(PTO2SharedMemoryHeader), PTO2_ALIGN_SIZE); - for (int r = 0; r < ring_id; r++) { - off += PTO2_ALIGN_UP(task_window_sizes[r] * sizeof(PTO2TaskDescriptor), PTO2_ALIGN_SIZE); - off += PTO2_ALIGN_UP(task_window_sizes[r] * sizeof(PTO2TaskPayload), PTO2_ALIGN_SIZE); - off += PTO2_ALIGN_UP(task_window_sizes[r] * sizeof(PTO2TaskSlotState), PTO2_ALIGN_SIZE); - } - PTO2RingSegmentOffsets o{}; - o.descriptors = off; - off += PTO2_ALIGN_UP(task_window_sizes[ring_id] * sizeof(PTO2TaskDescriptor), PTO2_ALIGN_SIZE); - o.payloads = off; - off += PTO2_ALIGN_UP(task_window_sizes[ring_id] * sizeof(PTO2TaskPayload), PTO2_ALIGN_SIZE); - o.slot_states = off; - off += PTO2_ALIGN_UP(task_window_sizes[ring_id] * sizeof(PTO2TaskSlotState), PTO2_ALIGN_SIZE); - o.end = off; - return o; -} - -// Device address of ring `ring_id`'s task_descriptors array. inline PTO2TaskDescriptor *ring_task_descriptors_addr( void *sm_dev_base, const uint64_t task_window_sizes[PTO2_MAX_RING_DEPTH], int ring_id ) noexcept { - return reinterpret_cast( - static_cast(sm_dev_base) + ring_segment_offsets(task_window_sizes, ring_id).descriptors - ); -} - -// Device address of ring `ring_id`'s slot_states array (used by the allocator's -// deadlock detector to inspect the head task's state/fanout). -inline PTO2TaskSlotState * -ring_slot_states_addr(void *sm_dev_base, const uint64_t task_window_sizes[PTO2_MAX_RING_DEPTH], int ring_id) noexcept { - return reinterpret_cast( - static_cast(sm_dev_base) + ring_segment_offsets(task_window_sizes, ring_id).slot_states - ); + assert(ring_id >= 0 && ring_id < PTO2_MAX_RING_DEPTH && "pto2_sm_layout: ring_id out of range"); + char *p = static_cast(sm_dev_base); + p += PTO2_ALIGN_UP(sizeof(PTO2SharedMemoryHeader), PTO2_ALIGN_SIZE); + for (int r = 0; r < ring_id; r++) { + p += PTO2_ALIGN_UP(task_window_sizes[r] * sizeof(PTO2TaskDescriptor), PTO2_ALIGN_SIZE); + p += PTO2_ALIGN_UP(task_window_sizes[r] * sizeof(PTO2TaskPayload), PTO2_ALIGN_SIZE); + p += PTO2_ALIGN_UP(task_window_sizes[r] * sizeof(PTO2TaskSlotState), PTO2_ALIGN_SIZE); + // Mirror setup_pointers_per_ring's per-ring stride exactly: completion_flags + // follows slot_states for every ring. Omitting it made this address diverge + // from where submit writes descriptors for rings >= 1, so the allocator read + // unwritten memory and update_heap_tail underflowed the reclaim tail. + p += PTO2_ALIGN_UP(task_window_sizes[r] * sizeof(std::atomic), PTO2_ALIGN_SIZE); + } + return reinterpret_cast(p); } } // namespace pto2_sm_layout diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_submit_types.h b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_submit_types.h index 56e6cab04f..923d7caa46 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_submit_types.h +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_submit_types.h @@ -9,27 +9,14 @@ * ----------------------------------------------------------------------------------------------------------- */ -/** - * PTO Submit Types - Shared submit-contract definitions - * - * Header-only definitions shared by orchestration-facing and runtime-facing - * headers. Keeps orchestration slim (no dependency on pto_runtime2_types.h). - */ - #pragma once #include inline constexpr int32_t INVALID_KERNEL_ID = -1; -/** - * Subtask slot count: AIC, AIV0, AIV1 - */ inline constexpr int32_t PTO2_SUBTASK_SLOT_COUNT = 3; -/** - * Subtask slot indices - */ enum class PTO2SubtaskSlot : uint8_t { AIC = 0, AIV0 = 1, @@ -114,14 +101,8 @@ enum class PTO2ResourceShape : uint8_t { DUMMY = 3, // Dependency-only (no AICore dispatch) }; -// Number of *dispatchable* resource shapes (AIC, AIV, MIX). DUMMY does not -// allocate a per-shape ready_queue entry / local buffer — it lives in a -// dedicated queue inside PTO2SchedulerState. inline constexpr int32_t PTO2_NUM_RESOURCE_SHAPES = 3; -/** - * Bitmask of active subtask slots + flags, sizeof == 1. - */ class ActiveMask { public: constexpr ActiveMask() = default; @@ -173,12 +154,6 @@ class ActiveMask { static_assert(sizeof(ActiveMask) == 1, "ActiveMask must be exactly 1 byte"); -/** - * Mixed-task submit contract. - * - * Each field holds either a valid kernel ID or INVALID_KERNEL_ID (inactive). - * At least one slot must be valid. - */ struct MixedKernels { int32_t aic_kernel_id{INVALID_KERNEL_ID}; int32_t aiv0_kernel_id{INVALID_KERNEL_ID}; @@ -193,13 +168,6 @@ struct MixedKernels { } }; -/** - * SPMD launch parameters carried inside Arg. - * - * Controls how many logical blocks (SPMD dimension) a single task - * is expanded into at dispatch time. Each block receives a unique - * block_idx in [0, block_num) via the per-dispatch LocalContext. - */ class PTO2LaunchSpec { public: constexpr PTO2LaunchSpec() = default; diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_tensormap.h b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_tensormap.h index 5c12c0296d..77903520e7 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_tensormap.h +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_tensormap.h @@ -9,37 +9,6 @@ * ----------------------------------------------------------------------------------------------------------- */ -/** - * PTO Runtime2 - TensorMap Interface - * - * TensorMap provides producer lookup for dependency discovery: - * - Maps Tensor -> producer task ID - * - Used by pto_submit_task() to find dependencies - * - * Key design features: - * 1. Ring buffer pool for entries (no malloc/free) - * 2. Lazy invalidation (entries become stale when producer retires) - * 3. Per-task per-ring entry tracking for efficient cleanup - * 4. OVERLAP DETECTION: Detects dependencies for overlapping sub-regions - * - * Hash table with chaining: - * - buckets[] array of head offsets - * - Entries linked via next_in_bucket - * - Insert at head (newest first) for sorted chains - * - * CRITICAL: Hash only by base_ptr - * ============================== - * For overlap detection to work, ALL sub-regions of the same base tensor - * MUST be in the SAME hash bucket. This allows lookup to compare all - * potentially overlapping regions. - * - * Overlap detection: Two regions create a dependency if: - * 1. Same base_ptr (raw tensor pointer) - * 2. Byte ranges [offset, offset+size) intersect - * - * Based on: docs/RUNTIME_LOGIC.md - */ - #pragma once #include "common.h" @@ -128,15 +97,15 @@ struct alignas(64) PTO2TensorMapEntry { // === Cache line 1 (64B) — lookup hot path; mirrors Tensor line 1 from byte 16 === uint64_t buffer_addr; // 8B [0, 8): tensor base address (hash key, mirrors Tensor::buffer.addr) PTO2TensorMapEntry *next_in_bucket; // 8B [8, 16): next entry in hash bucket chain (overlays Tensor::buffer.size) - PTO2TaskId producer_task_id; // 8B [16,24): mirrors Tensor::owner_task_id slot - uint64_t start_offset; // 8B [24,32): mirrors Tensor::start_offset (element offset) - int32_t version; // 4B [32,36): mirrors Tensor::version - uint32_t ndims; // 4B [36,40): mirrors Tensor::ndims - DataType dtype; // 1B [40,41): mirrors Tensor::dtype - bool manual_dep; // 1B [41,42): mirrors Tensor::manual_dep - bool is_contiguous; // 1B [42,43): mirrors Tensor::is_contiguous - uint8_t __padding1__; // 1B [43,44): mirrors Tensor padding - uint32_t shapes[MAX_TENSOR_DIMS]; // 20B [44,64): mirrors Tensor::shapes + PTO2TaskId producer_task_id; // 8B [16, 24): mirrors Tensor::owner_task_id slot + uint64_t start_offset; // 8B [24, 32): mirrors Tensor::start_offset (element offset) + int32_t version; // 4B [32, 36): mirrors Tensor::version + uint32_t ndims; // 4B [36, 40): mirrors Tensor::ndims + DataType dtype; // 1B [40, 41): mirrors Tensor::dtype + bool manual_dep; // 1B [41, 42): mirrors Tensor::manual_dep + bool is_contiguous; // 1B [42, 43): mirrors Tensor::is_contiguous + uint8_t __padding1__; // 1B [43, 44): mirrors Tensor padding + uint32_t shapes[MAX_TENSOR_DIMS]; // 20B [44, 64): mirrors Tensor::shapes // === Cache line 2 (64B) — chain manipulation + non-contiguous overlap data === PTO2TensorMapEntry *prev_in_bucket; // 8B [64, 72) @@ -144,24 +113,10 @@ struct alignas(64) PTO2TensorMapEntry { PTO2TensorMapEntry *prev_in_task; // 8B [80, 88) int32_t bucket_index; // 4B [88, 92): -1 when unlinked uint32_t __padding2__; // 4B [92, 96) - uint64_t extent_elem_cache; // 8B [96,104): non-contiguous extent (mirrors Tensor) - uint32_t strides[MAX_TENSOR_DIMS]; // 20B [104,124): element strides, mirrors Tensor::strides - uint8_t __padding3__[4]; // 4B [124,128) - - /** - * Copy overlap-relevant fields from a Tensor into this entry. - * - * 64B memcpy of Tensor cache line 1 populates buffer_addr (byte [0,8)), - * producer_task_id, start_offset, version, ndims, dtype, manual_dep, - * is_contiguous and shapes[]. Byte [8,16) holds Tensor::buffer.size in - * the source and gets written into next_in_bucket; that's harmless - * because link_entry() overwrites next_in_bucket immediately after. - * - * Cache line 2 (stride / extent_elem_cache) is derived from line 1 when - * the source is canonically contiguous (is_contiguous && start_offset==0), - * so the producer Tensor's cache line 2 stays cold during insert. Only - * non-contiguous producers pay one extra line 2 read. - */ + uint64_t extent_elem_cache; // 8B [96, 104): non-contiguous extent (mirrors Tensor) + uint32_t strides[MAX_TENSOR_DIMS]; // 20B [104, 124): element strides, mirrors Tensor::strides + uint8_t __padding3__[4]; // 4B [124, 128) + void copy_from_tensor(const Tensor &tensor) { memcpy(this, &tensor, 64); if (tensor.is_contiguous && tensor.start_offset == 0) { @@ -176,9 +131,8 @@ struct alignas(64) PTO2TensorMapEntry { } } else { extent_elem_cache = tensor.extent_elem_cache; - for (uint32_t i = 0; i < tensor.ndims; i++) { + for (uint32_t i = 0; i < tensor.ndims; i++) strides[i] = tensor.strides[i]; - } } } @@ -188,9 +142,8 @@ struct alignas(64) PTO2TensorMapEntry { // Create-info outputs are always contiguous with start_offset = 0; // extent_elem = prod(shapes); stride is row-major. uint64_t numel = 1; - for (uint32_t i = 0; i < tensor_create_info.ndims; i++) { + for (uint32_t i = 0; i < tensor_create_info.ndims; i++) numel *= tensor_create_info.shapes[i]; - } extent_elem_cache = numel; uint32_t s = 1; for (int32_t i = static_cast(tensor_create_info.ndims) - 1; i >= 0; i--) { @@ -199,11 +152,6 @@ struct alignas(64) PTO2TensorMapEntry { } } - /** - * Effective element extent of this entry. - * Contiguous-aligned views compute it from shapes alone (line 1 hit only); - * non-contiguous views read the cached value from line 2. - */ uint64_t effective_extent_elem() const { if (is_contiguous) { uint64_t n = 1; @@ -214,29 +162,10 @@ struct alignas(64) PTO2TensorMapEntry { return extent_elem_cache; } - /** - * Check overlap between input tensor and this entry (the producer output). - * - * Three-level cascade: - * L1 — O(1) byte-range intersection. Disjoint -> NO_OVERLAP. - * L2 — O(ndims) hyper-rectangle precise check, eligible only when both - * sides share the same canonical row-major axis layout (same - * dtype/ndims/strides[], stride descends as integer multiples, - * start_offset decomposes cleanly under the reference shape). - * Yields NO_OVERLAP / COVERED / OTHER per-dim. - * L3 — Non-hyper-rectangle pairs (transpose/permute mismatch, slice - * with step, etc): conservative OTHER. Exact enumeration via - * contiguous-segment merge is scheduled for a follow-up. - * - * COVERED is returned when `input` completely contains `entry` per-dim - * — dep_compute uses this to retire the now-redundant entry. - */ OverlapStatus check_overlap(const Tensor &input) const { debug_assert(input.buffer.addr == buffer_addr); debug_assert(input.version >= version); - if (input.version > version) { - return OverlapStatus::OTHER; - } + if (input.version > version) return OverlapStatus::OTHER; // -------- L1: byte-range intersection (O(1) fast reject) -------- const uint64_t in_begin = input.start_offset; @@ -245,27 +174,15 @@ struct alignas(64) PTO2TensorMapEntry { const uint64_t ent_end = start_offset + effective_extent_elem(); Segment in_range_bytes{in_begin, in_end}; Segment ent_range_bytes{ent_begin, ent_end}; - if (!in_range_bytes.line_segment_intersection(ent_range_bytes)) { - return OverlapStatus::NO_OVERLAP; - } + if (!in_range_bytes.line_segment_intersection(ent_range_bytes)) return OverlapStatus::NO_OVERLAP; // -------- L2 prereqs: same axis layout? -------- - if (input.dtype != dtype || input.ndims != ndims || ndims == 0) { - return OverlapStatus::OTHER; - } - for (uint32_t i = 0; i < ndims; i++) { + if (input.dtype != dtype || input.ndims != ndims || ndims == 0) return OverlapStatus::OTHER; + for (uint32_t i = 0; i < ndims; i++) if (input.strides[i] != strides[i]) return OverlapStatus::OTHER; - } - // strides[ndims-1] must be 1 and strides[i-1] must be an integer - // multiple of strides[i] for the row-major reference-shape derivation - // below to hold. This rejects slice-with-step (strides[d] != prev factor) - // and any view chain that scrambles the axis order. (strides is - // uint32_t with the > 0 invariant enforced at construction, so no - // sign check needed.) if (strides[ndims - 1] != 1) return OverlapStatus::OTHER; - for (uint32_t i = 1; i < ndims; i++) { + for (uint32_t i = 1; i < ndims; i++) if (strides[i - 1] % strides[i] != 0) return OverlapStatus::OTHER; - } // Derive reference shape A from stride. By construction stride is // row-major over A: strides[i] = prod(A[i+1..ndims-1]). So @@ -324,12 +241,8 @@ struct alignas(64) PTO2TensorMapEntry { for (uint32_t i = 0; i < ndims; i++) { Segment in_seg{in_offsets[i], static_cast(in_offsets[i]) + input.shapes[i]}; Segment ent_seg{ent_offsets[i], static_cast(ent_offsets[i]) + shapes[i]}; - if (!in_seg.line_segment_intersection(ent_seg)) { - return OverlapStatus::NO_OVERLAP; - } - if (!in_seg.contains(ent_seg)) { - input_contains_entry = false; - } + if (!in_seg.line_segment_intersection(ent_seg)) return OverlapStatus::NO_OVERLAP; + if (!in_seg.contains(ent_seg)) input_contains_entry = false; } return input_contains_entry ? OverlapStatus::COVERED : OverlapStatus::OTHER; } @@ -349,15 +262,6 @@ static_assert( offsetof(PTO2TensorMapEntry, prev_in_bucket) == 64, "TensorMapEntry must be exactly 2 cache lines (128 bytes)" ); -// ============================================================================= -// TensorMap Lookup Chain Length Statistics (compile-time toggle) -// ============================================================================= - -/** - * TensorMap structure - * - * Hash table with ring buffer entry pool and lazy invalidation. - */ struct PTO2TensorMap { // Hash table buckets (fixed size, power of 2) PTO2TensorMapEntry **buckets; // Array of offsets into entry_pool (-1 = empty) @@ -388,34 +292,8 @@ struct PTO2TensorMap { return task_local_id & (task_window_sizes[ring_id] - 1); } - // Accessors read by scope_stats_collector. Declared unconditionally so the - // collector .cpp compiles at SIMPLER_DFX=0 (collector is unconditional — - // setter symbols must export for host dlsym; the probe call sites that use - // these accessors stay gated by SIMPLER_DFX). int32_t current_used() const { return next_entry_idx - free_num; } int32_t pool_capacity() const { return pool_size; } - int32_t free_entries() const { return pool_size - current_used(); } - - // Reclaim retired entries across every ring, advancing each ring's cleanup - // cursor (last_cleanup[r]) to the supplied watermark. Returns the summed - // last_task_alive across rings — the monotone progress signal the - // orchestrator's exhaustion back-pressure loop watches to tell a transient - // shortage (some ring still retiring tasks) from a wedged pool (no ring - // advancing). Idempotent per watermark: a ring whose alive has not passed - // last_cleanup[r] is skipped, so it never double-frees. - int64_t reclaim_retired_all(const int32_t sm_last_task_alive[PTO2_MAX_RING_DEPTH]) { - int64_t alive_sum = 0; - for (int32_t r = 0; r < PTO2_MAX_RING_DEPTH; r++) { - int32_t alive = sm_last_task_alive[r]; - sync_validity(r, alive); - if (alive > last_cleanup[r]) { - cleanup_retired(r, last_cleanup[r], alive); - last_cleanup[r] = alive; - } - alive_sum += alive; - } - return alive_sum; - } // new_entry only allocates memory, does not assign attributes PTO2TensorMapEntry *new_entry() { @@ -442,9 +320,7 @@ struct PTO2TensorMap { } // Update successor's prev pointer - if (entry.next_in_bucket != nullptr) { - entry.next_in_bucket->prev_in_bucket = entry.prev_in_bucket; - } + if (entry.next_in_bucket != nullptr) entry.next_in_bucket->prev_in_bucket = entry.prev_in_bucket; free_entry_list[free_num++] = &entry; entry.bucket_index = -1; @@ -454,72 +330,30 @@ struct PTO2TensorMap { entry.prev_in_task = nullptr; } - // ============================================================================= - // TensorMap API - // ============================================================================= - - /** - * Phase 1: reserve every sub-region (buckets, entry_pool, free list, per-ring - * task_entry_heads) on the supplied arena. Records the resulting offsets in - * the returned layout descriptor. Must be called before the arena is - * committed. - */ static PTO2TensorMapLayout reserve_layout( - DeviceArena &arena, int32_t num_buckets, int32_t pool_size, const int32_t task_window_sizes[PTO2_MAX_RING_DEPTH] + DeviceArena &arena, int32_t new_num_buckets, int32_t new_pool_size, + const int32_t new_task_window_sizes[PTO2_MAX_RING_DEPTH] ); - /** - * Same as reserve_layout() with default sizes (PTO2_TENSORMAP_NUM_BUCKETS, - * PTO2_TENSORMAP_POOL_SIZE). - */ static PTO2TensorMapLayout - reserve_layout_default(DeviceArena &arena, const int32_t task_window_sizes[PTO2_MAX_RING_DEPTH]); - - /** - * Phase 3a: write everything *except* arena-internal pointer fields - * (buckets, entry_pool, free_entry_list, task_entry_heads[r]). - * Uses arena.region_ptr to address the arena regions for data writes, - * but does not store those addresses in struct fields. Safe to call on - * a host arena that holds the prebuilt image. - */ + reserve_layout_default(DeviceArena &arena, const int32_t new_task_window_sizes[PTO2_MAX_RING_DEPTH]); + bool init_data_from_layout(const PTO2TensorMapLayout &layout, DeviceArena &arena); - void reset_for_reuse(const PTO2TensorMapLayout &layout); - /** - * Phase 3b: write the arena-internal pointer fields. Idempotent; - * called once on the host arena and once on the AICPU after attach. - */ void wire_arena_pointers(const PTO2TensorMapLayout &layout, DeviceArena &arena); - /** - * Tear down state. Does not free memory — the arena owns the backing - * buffer. Pointers are set to nullptr so accidental reuse traps. - */ + // Surgical reset for arena reuse (#1234): O(1) epoch bump replaces the + // O(num_buckets + pool_size + Σ task_window_sizes) re-init of + // init_data_from_layout. bucket_epochs[i] and task_entry_head_epochs[r][i] + // are compared against current_epoch on every lookup/insert; bumping + // current_epoch invalidates all previous entries logically. Only on the + // rare wrap to 0 do we pay the O(num_buckets + Σ window) reset. + void reset_for_reuse(); + void destroy(); - /** - * Update validity threshold from shared memory - * Called periodically to refresh the lazy invalidation threshold. - * - * @param last_task_alive Current value from shared memory - */ void sync_validity(int32_t ring_id, int32_t last_task_alive) { this->last_task_alives[ring_id] = last_task_alive; } - /** - * Lookup producer for a tensor region - * - * Searches the hash table for matching regions and invokes the callback - * for each overlapping valid entry. - * Stale entries from different rings are skipped (not truncated). - * - * The callback receives (PTO2TensorMapEntry &, OverlapStatus) and should - * return true to continue iteration, false to stop early. It is safe for - * the callback to call remove_entry() on the current entry: next_in_bucket - * is latched before invocation. - * - * @param tensor Tensor to look up - * @param on_match Callback invoked for each overlapping entry - */ template void lookup(const Tensor &tensor, Fn &&on_match) { uint32_t bucket_index = hash(tensor.buffer.addr); @@ -528,80 +362,32 @@ struct PTO2TensorMap { } PTO2TensorMapEntry *cur_entry = buckets[bucket_index]; -#if SIMPLER_TENSORMAP_PROFILING - g_lookup_count++; - int32_t chain_len = 0; -#endif - while (cur_entry != nullptr) { PTO2TensorMapEntry *next_entry = cur_entry->next_in_bucket; -#if SIMPLER_TENSORMAP_PROFILING - chain_len++; -#endif - // Skip stale entries (no chain truncation — entries from different - // rings can be interleaved, so a stale entry from one ring does NOT - // imply subsequent entries from other rings are also stale) if (!entry_valid(*cur_entry)) { cur_entry = next_entry; continue; } - // Entry is valid - check if regions OVERLAP (not just exact match) - // Since we hash only by base_ptr, all entries in this bucket have - // potential to overlap. We must check actual byte-range overlap. if (tensor.buffer.addr == cur_entry->buffer_addr) { -#if SIMPLER_TENSORMAP_PROFILING - g_lookup_overlap_checks++; -#endif auto overlap_status = cur_entry->check_overlap(tensor); if (overlap_status != OverlapStatus::NO_OVERLAP) { -#if SIMPLER_TENSORMAP_PROFILING - g_lookup_overlap_hits++; -#endif - if (!on_match(*cur_entry, overlap_status)) { -#if SIMPLER_TENSORMAP_PROFILING - g_lookup_chain_total += chain_len; - if (chain_len > g_lookup_chain_max) g_lookup_chain_max = chain_len; -#endif - return; - } + if (!on_match(*cur_entry, overlap_status)) return; } } // Move to next entry cur_entry = next_entry; } -#if SIMPLER_TENSORMAP_PROFILING - g_lookup_chain_total += chain_len; - if (chain_len > g_lookup_chain_max) g_lookup_chain_max = chain_len; -#endif } - /** - * Insert a new entry (called when task produces output) - * - * Allocates from ring buffer pool, may overwrite stale entries. - * Inserts at head of hash bucket chain (maintains task_id ordering). - * - * @param tensor Tensor produced - * @param producer_task_id Task ID of producer - */ void insert(const Tensor &tensor, PTO2TaskId producer_task_id) { PTO2TensorMapEntry *entry = new_entry(); entry->copy_from_tensor(tensor); link_entry(entry, tensor.buffer.addr, producer_task_id); } - /** - * Cleanup stale entries for retired tasks - * - * Called periodically by Orchestrator when last_task_alive advances. - * Removes entries from bucket chains for tasks in [old, new) range. - * - * @param old_last_task_alive Previous threshold - * @param new_last_task_alive New threshold - */ void cleanup_retired(int32_t ring_id, int32_t old_last_task_alive, int32_t new_last_task_alive) { // Iterate through retired tasks on this ring and remove their entries for (int32_t local_id = old_last_task_alive; local_id < new_last_task_alive; local_id++) { @@ -628,30 +414,12 @@ struct PTO2TensorMap { } } - // ============================================================================= - // Internal Helpers (exposed for testing) - // ============================================================================= - - /** - * Compute hash for tensor addr - * - * Multiplicative hash using the golden-ratio constant. Multiplication - * mixes ALL input bits into the high bits of the product, so aligned - * addresses (low bits all-zero) still distribute evenly. We extract - * the top log2(num_buckets) bits which carry the most entropy. - */ uint32_t hash(uint64_t key) { key *= 0x9E3779B97F4A7C15ULL; return static_cast(key >> (64 - __builtin_ctz(num_buckets))); } - /** - * Link an initialized entry into bucket and task chains. - */ void link_entry(PTO2TensorMapEntry *entry, uint64_t addr, PTO2TaskId producer_task_id) { -#if SIMPLER_TENSORMAP_PROFILING - g_insert_count++; -#endif uint32_t bucket_index = hash(addr); auto ring_id = producer_task_id.ring(); auto local_id = producer_task_id.local(); @@ -666,9 +434,7 @@ struct PTO2TensorMap { } entry->bucket_index = bucket_index; entry->next_in_bucket = buckets[bucket_index]; - if (entry->next_in_bucket != nullptr) { - entry->next_in_bucket->prev_in_bucket = entry; - } + if (entry->next_in_bucket != nullptr) entry->next_in_bucket->prev_in_bucket = entry; buckets[bucket_index] = entry; entry->prev_in_bucket = nullptr; @@ -679,15 +445,10 @@ struct PTO2TensorMap { } entry->next_in_task = task_entry_heads[ring_id][task_slot]; entry->prev_in_task = nullptr; - if (entry->next_in_task != nullptr) { - entry->next_in_task->prev_in_task = entry; - } + if (entry->next_in_task != nullptr) entry->next_in_task->prev_in_task = entry; task_entry_heads[ring_id][task_slot] = entry; } - /** - * Check if entry is valid (producer has not retired) - */ bool entry_valid(const PTO2TensorMapEntry &entry) const { return static_cast(entry.producer_task_id.local()) >= last_task_alives[entry.producer_task_id.ring()]; } @@ -697,10 +458,6 @@ struct PTO2TensorMap { free_entry(entry); } - /** - * Remove entry from its task chain (O(1) with prev pointer) - * Called during pool wrap-around to unlink reused entries. - */ void remove_from_task(PTO2TensorMapEntry &entry) { always_assert(entry.bucket_index != -1); // must still be in a bucket // Update predecessor's next pointer (O(1) via prev_in_task) @@ -715,50 +472,13 @@ struct PTO2TensorMap { } // Update successor's prev pointer - if (entry.next_in_task != nullptr) { - entry.next_in_task->prev_in_task = entry.prev_in_task; - } + if (entry.next_in_task != nullptr) entry.next_in_task->prev_in_task = entry.prev_in_task; entry.next_in_task = nullptr; entry.prev_in_task = nullptr; } - // ============================================================================= - // Debug Utilities - // ============================================================================= - - /** - * Print TensorMap statistics - */ - void print_stats(); - - /** - * Get count of valid entries - */ int32_t valid_count(); - // ============================================================================= - // TensorMap Synchronization - // ============================================================================= - - /** - * Sync TensorMap validity threshold from shared memory - * - * Called periodically to refresh the lazy invalidation threshold. - * Also triggers cleanup if threshold has advanced significantly. - */ void sync_tensormap(PTO2TaskId task_id, int32_t sm_last_task_alive); }; - -#if SIMPLER_TENSORMAP_PROFILING -struct PTO2TensorMapProfilingData { - uint64_t lookup_chain_total; - uint64_t lookup_count; - int32_t lookup_chain_max; - uint64_t overlap_checks; - uint64_t overlap_hits; - uint64_t insert_count; -}; - -PTO2TensorMapProfilingData pto2_tensormap_get_profiling(); -#endif diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.cpp b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.cpp index 2f9d96d8a7..f2bcc9e437 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.cpp +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.cpp @@ -8,102 +8,9 @@ * See LICENSE in the root of the software repository for the full text of the License. * ----------------------------------------------------------------------------------------------------------- */ -/** - * PTO Runtime2 - Scheduler Implementation - * - * Implements scheduler state management, ready queues, and task lifecycle. - * - * Based on: docs/RUNTIME_LOGIC.md - */ - -#include "pto_scheduler.h" -#include -#include -#include "common/unified_log.h" - -#if SIMPLER_DFX -// Weak fallbacks for host/UT builds that don't link the scope_stats collector. -extern "C" __attribute__((weak, visibility("hidden"))) bool is_scope_stats_enabled() { return false; } -extern "C" __attribute__((weak, visibility("hidden"))) void scope_stats_note_heap_wrap(int) {} -#endif - -// ============================================================================= -// Scheduler Profiling Counters -// ============================================================================= - -#if SIMPLER_SCHED_PROFILING -#include "common/platform_config.h" - -uint64_t g_sched_lock_cycle[PLATFORM_MAX_AICPU_THREADS] = {}; -uint64_t g_sched_fanout_cycle[PLATFORM_MAX_AICPU_THREADS] = {}; -uint64_t g_sched_fanin_cycle[PLATFORM_MAX_AICPU_THREADS] = {}; -uint64_t g_sched_self_consumed_cycle[PLATFORM_MAX_AICPU_THREADS] = {}; -uint64_t g_sched_lock_wait_cycle[PLATFORM_MAX_AICPU_THREADS] = {}; -uint64_t g_sched_push_wait_cycle[PLATFORM_MAX_AICPU_THREADS] = {}; -uint64_t g_sched_pop_wait_cycle[PLATFORM_MAX_AICPU_THREADS] = {}; -uint64_t g_sched_lock_atomic_count[PLATFORM_MAX_AICPU_THREADS] = {}; -uint64_t g_sched_fanout_atomic_count[PLATFORM_MAX_AICPU_THREADS] = {}; -uint64_t g_sched_fanin_atomic_count[PLATFORM_MAX_AICPU_THREADS] = {}; -uint64_t g_sched_self_atomic_count[PLATFORM_MAX_AICPU_THREADS] = {}; -uint64_t g_sched_pop_atomic_count[PLATFORM_MAX_AICPU_THREADS] = {}; -uint64_t g_sched_complete_count[PLATFORM_MAX_AICPU_THREADS] = {}; - -PTO2SchedProfilingData scheduler_get_profiling(int thread_idx) { - PTO2SchedProfilingData d; - d.lock_cycle = std::exchange(g_sched_lock_cycle[thread_idx], 0); - d.fanout_cycle = std::exchange(g_sched_fanout_cycle[thread_idx], 0); - d.fanin_cycle = std::exchange(g_sched_fanin_cycle[thread_idx], 0); - d.self_consumed_cycle = std::exchange(g_sched_self_consumed_cycle[thread_idx], 0); - d.lock_wait_cycle = std::exchange(g_sched_lock_wait_cycle[thread_idx], 0); - d.push_wait_cycle = std::exchange(g_sched_push_wait_cycle[thread_idx], 0); - d.pop_wait_cycle = std::exchange(g_sched_pop_wait_cycle[thread_idx], 0); - d.lock_atomic_count = std::exchange(g_sched_lock_atomic_count[thread_idx], 0); - d.fanout_atomic_count = std::exchange(g_sched_fanout_atomic_count[thread_idx], 0); - d.fanin_atomic_count = std::exchange(g_sched_fanin_atomic_count[thread_idx], 0); - d.self_atomic_count = std::exchange(g_sched_self_atomic_count[thread_idx], 0); - d.pop_atomic_count = std::exchange(g_sched_pop_atomic_count[thread_idx], 0); - d.complete_count = std::exchange(g_sched_complete_count[thread_idx], 0); - return d; -} -#endif - -// ============================================================================= -// Debug Utilities -// ============================================================================= - -void PTO2SchedulerState::print_stats() { - PTO2SchedulerState *sched = this; - LOG_INFO_V0("=== Scheduler Statistics ==="); - for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { - if (sched->ring_sched_states[r].last_task_alive > 0) { - LOG_INFO_V0("Ring %d:", r); - LOG_INFO_V0(" last_task_alive: %d", sched->ring_sched_states[r].last_task_alive); - auto &dp = sched->ring_sched_states[r].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 - LOG_INFO_V0("tasks_completed: %lld", (long long)sched->tasks_completed.load(std::memory_order_relaxed)); - LOG_INFO_V0("tasks_consumed: %lld", (long long)sched->tasks_consumed.load(std::memory_order_relaxed)); -#endif - LOG_INFO_V0("============================"); -} - -void PTO2SchedulerState::print_queues() { - PTO2SchedulerState *sched = this; - LOG_INFO_V0("=== Ready Queues ==="); - - const char *shape_names[] = {"AIC", "AIV", "MIX"}; - - for (int i = 0; i < PTO2_NUM_RESOURCE_SHAPES; i++) { - LOG_INFO_V0(" %s: count=%" PRIu64, shape_names[i], sched->ready_queues[i].size()); - } - LOG_INFO_V0(" DUMMY: count=%" PRIu64, sched->dummy_ready_queue.size()); - LOG_INFO_V0("===================="); -} +// PTO2SchedulerState's members are per-task/per-completion hot-path routines +// (classify_fanin_state, register_wake, drain_wiring_queue, on_mixed_task_complete, +// …) and are defined inline in pto_scheduler.h. SchedulerContext's out-of-line +// methods live in scheduler_dispatch.cpp / scheduler_cold_path.cpp / +// scheduler_completion.cpp. This translation unit carries no definitions. diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.h b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.h index 87d74878c1..8fda94d22b 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.h +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.h @@ -9,71 +9,52 @@ * ----------------------------------------------------------------------------------------------------------- */ -/** - * PTO Runtime2 - Scheduler Interface - * - * 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) - * - * 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 - * - * Based on: docs/RUNTIME_LOGIC.md - */ - #pragma once #include #include "common/core_type.h" -#include "common/memory_barrier.h" #include "utils/device_arena.h" -#include "aicpu/platform_regs.h" // get_reg_ptr / RegId for the early-dispatch doorbell #include "pto_async_wait.h" #include "pto_ring_buffer.h" #include "pto_runtime2_types.h" #include "pto_shared_memory.h" -#include "aicpu/device_time.h" // get_sys_cnt_aicpu (weak; used by early-dispatch doorbell timing too) -#if SIMPLER_SCHED_PROFILING -#define PTO2_SCHED_CYCLE_START() uint64_t _st0 = get_sys_cnt_aicpu(), _st1 -#define PTO2_SCHED_CYCLE_LAP(acc) \ - do { \ - _st1 = get_sys_cnt_aicpu(); \ - acc += (_st1 - _st0); \ - _st0 = _st1; \ - } while (0) -#endif - -// ============================================================================= -// Ready Queue (Lock-free bounded MPMC — Vyukov design) -// ============================================================================= +// Forward declaration so this header can compile under both AICPU and host +// builds. The actual definition is provided by aicpu/device_time.cpp (AICPU) +// or a weak stub in pto_runtime2.h (host). Used only for sub-phase profiling. +uint64_t get_sys_cnt_aicpu(); -/** - * Per-slot entry: sequence counter for ABA safety + task payload - */ struct PTO2ReadyQueueSlot { std::atomic sequence; PTO2TaskSlotState *slot_state; - uint64_t task_id_snapshot; }; -/** - * Lock-free bounded MPMC queue (Dmitry Vyukov design) - * - * Key properties: - * - enqueue_pos and dequeue_pos on separate cache lines (no false sharing) - * - Per-slot sequence counter prevents ABA problem - * - Empty queue pop returns immediately (single atomic load, no lock) - * - CAS contention is split: producers only touch enqueue_pos, - * consumers only touch dequeue_pos - */ +// Number of CoreType values eligible for local dispatch (AIC=0, AIV=1) +static constexpr int PTO2_LOCAL_DISPATCH_TYPE_NUM = 2; + +struct PTO2LocalReadyBuffer { + PTO2TaskSlotState **slot_states = nullptr; + int count = 0; + int capacity = 0; + + void reset(PTO2TaskSlotState **buf, int cap) { + slot_states = buf; + count = 0; + capacity = cap; + } + + bool try_push(PTO2TaskSlotState *s) { + if (slot_states && count < capacity) { + slot_states[count++] = s; + return true; + } + return false; + } + + PTO2TaskSlotState *pop() { return (count > 0) ? slot_states[--count] : nullptr; } +}; + struct alignas(64) PTO2ReadyQueue { PTO2ReadyQueueSlot *slots; uint64_t capacity; @@ -92,11 +73,13 @@ struct alignas(64) PTO2ReadyQueue { return (e >= d) ? (e - d) : 0; } + // No-op: the sequence-based Vyukov MPMC queue is self-consistent across + // runs — every slot's sequence at end of run 1 equals the enqueue_pos + // where run 2's first push at that slot will land, so pushes/pops resume + // seamlessly without any reset. void reset_for_reuse() {} - bool push(PTO2TaskSlotState *slot_state) { return push_tagged(slot_state, 0); } - - bool push_tagged(PTO2TaskSlotState *slot_state, uint64_t task_id_snapshot) { + bool push(PTO2TaskSlotState *slot_state) { uint64_t pos; PTO2ReadyQueueSlot *slot; while (true) { @@ -107,25 +90,21 @@ struct alignas(64) PTO2ReadyQueue { if (diff == 0) { if (enqueue_pos.compare_exchange_weak( pos, pos + 1, std::memory_order_relaxed, std::memory_order_relaxed - )) { + )) break; - } } else if (diff < 0) { return false; // Queue full } } slot->slot_state = slot_state; - slot->task_id_snapshot = task_id_snapshot; slot->sequence.store(static_cast(pos + 1), std::memory_order_release); return true; } // Batch push: reserve count slots with a single CAS after confirming // every target slot is available under the usual Vyukov sequence check. - void push_batch(PTO2TaskSlotState **items, int count) { push_batch_tagged(items, nullptr, count); } - - void push_batch_tagged(PTO2TaskSlotState **items, const uint64_t *task_id_snapshots, int count) { + void push_batch(PTO2TaskSlotState **items, int count) { if (count == 0) return; uint64_t pos; @@ -141,73 +120,25 @@ struct alignas(64) PTO2ReadyQueue { break; } } - if (!ready) { - continue; - } + if (!ready) continue; if (enqueue_pos.compare_exchange_weak( pos, pos + count, std::memory_order_relaxed, std::memory_order_relaxed - )) { + )) break; - } } for (int i = 0; i < count; i++) { PTO2ReadyQueueSlot *slot = &slots[(pos + i) & mask]; slot->slot_state = items[i]; - slot->task_id_snapshot = task_id_snapshots == nullptr ? 0 : task_id_snapshots[i]; slot->sequence.store(static_cast(pos + i + 1), std::memory_order_release); } } -#if SIMPLER_ORCH_PROFILING || SIMPLER_SCHED_PROFILING - bool push(PTO2TaskSlotState *slot_state, uint64_t &atomic_count, uint64_t &wait_cycle) { - uint64_t pos; - PTO2ReadyQueueSlot *slot; - uint64_t t0 = get_sys_cnt_aicpu(); - bool contended = false; - uint32_t atomic_ops = 0; - while (true) { - pos = enqueue_pos.load(std::memory_order_relaxed); - slot = &slots[pos & mask]; - int64_t seq = slot->sequence.load(std::memory_order_acquire); - int64_t diff = seq - static_cast(pos); - atomic_ops += 2; // enqueue_pos.load + sequence.load - if (diff == 0) { - if (enqueue_pos.compare_exchange_weak( - pos, pos + 1, std::memory_order_relaxed, std::memory_order_relaxed - )) { - atomic_ops++; // successful CAS - break; - } - contended = true; - atomic_ops++; // failed CAS - } else if (diff < 0) { - return false; // Queue full - } else { - contended = true; // diff > 0: slot not yet released, spin - } - } - atomic_ops++; // final sequence.store - atomic_count += atomic_ops; - if (contended) { - wait_cycle += (get_sys_cnt_aicpu() - t0); - } - - slot->slot_state = slot_state; - slot->sequence.store(static_cast(pos + 1), std::memory_order_release); - return true; - } -#endif - - PTO2TaskSlotState *pop() { return pop_tagged(nullptr); } - - PTO2TaskSlotState *pop_tagged(uint64_t *task_id_snapshot) { + PTO2TaskSlotState *pop() { // Fast-path: skip slot load when queue is clearly empty uint64_t d = dequeue_pos.load(std::memory_order_relaxed); uint64_t e = enqueue_pos.load(std::memory_order_relaxed); - if (d >= e) { - return nullptr; - } + if (d >= e) return nullptr; uint64_t pos; PTO2ReadyQueueSlot *slot; @@ -226,66 +157,14 @@ struct alignas(64) PTO2ReadyQueue { } } - PTO2TaskSlotState *result = slot->slot_state; - if (task_id_snapshot != nullptr) *task_id_snapshot = slot->task_id_snapshot; - slot->sequence.store(static_cast(pos + mask + 1), std::memory_order_release); - return result; - } - -#if SIMPLER_SCHED_PROFILING - PTO2TaskSlotState *pop(uint64_t &atomic_count, uint64_t &wait_cycle) { - // Fast-path: skip slot load when queue is clearly empty - uint64_t d = dequeue_pos.load(std::memory_order_relaxed); - uint64_t e = enqueue_pos.load(std::memory_order_relaxed); - atomic_count += 2; // dequeue_pos.load + enqueue_pos.load - if (d >= e) { - return nullptr; - } - - uint64_t pos; - PTO2ReadyQueueSlot *slot; - uint64_t t0 = get_sys_cnt_aicpu(); - bool contended = false; - uint32_t atomic_ops = 0; - while (true) { - pos = dequeue_pos.load(std::memory_order_relaxed); - slot = &slots[pos & mask]; - int64_t seq = slot->sequence.load(std::memory_order_acquire); - int64_t diff = seq - static_cast(pos + 1); - atomic_ops += 2; // dequeue_pos.load + sequence.load - if (diff == 0) { - if (dequeue_pos.compare_exchange_weak( - pos, pos + 1, std::memory_order_relaxed, std::memory_order_relaxed - )) { - atomic_ops++; // successful CAS - break; - } - contended = true; - atomic_ops++; // failed CAS - } else if (diff < 0) { - atomic_count += atomic_ops; - return nullptr; // Queue empty - } else { - contended = true; - } - } - atomic_ops++; // final sequence.store - atomic_count += atomic_ops; - if (contended) { - wait_cycle += (get_sys_cnt_aicpu() - t0); - } - PTO2TaskSlotState *result = slot->slot_state; slot->sequence.store(static_cast(pos + mask + 1), std::memory_order_release); return result; } -#endif // Batch pop: reserve a contiguous run of ready slots with a single CAS. // Returns actual number of items popped (may be less than max_count). - int pop_batch(PTO2TaskSlotState **out, int max_count) { return pop_batch_tagged(out, nullptr, max_count); } - - int pop_batch_tagged(PTO2TaskSlotState **out, uint64_t *task_id_snapshots, int max_count) { + int pop_batch(PTO2TaskSlotState **out, int max_count) { uint64_t pos; int count; while (true) { @@ -299,9 +178,7 @@ struct alignas(64) PTO2ReadyQueue { count++; continue; } - if (diff < 0) { - break; - } + if (diff < 0) break; count = -1; break; } @@ -309,100 +186,119 @@ struct alignas(64) PTO2ReadyQueue { if (count < 0) continue; if (dequeue_pos.compare_exchange_weak( pos, pos + count, std::memory_order_relaxed, std::memory_order_relaxed - )) { + )) break; - } } for (int i = 0; i < count; i++) { PTO2ReadyQueueSlot *slot = &slots[(pos + i) & mask]; out[i] = slot->slot_state; - if (task_id_snapshots != nullptr) task_id_snapshots[i] = slot->task_id_snapshot; slot->sequence.store(static_cast(pos + i + mask + 1), std::memory_order_release); } return count; } +}; -#if SIMPLER_SCHED_PROFILING - int pop_batch(PTO2TaskSlotState **out, int max_count, uint64_t &atomic_count, uint64_t &wait_cycle) { - uint64_t pos; - int count; - uint64_t t0 = get_sys_cnt_aicpu(); - bool contended = false; - uint32_t atomic_ops = 0; - while (true) { - pos = dequeue_pos.load(std::memory_order_relaxed); - atomic_ops++; // dequeue_pos.load - count = 0; - while (count < max_count) { - PTO2ReadyQueueSlot *slot = &slots[(pos + count) & mask]; - int64_t seq = slot->sequence.load(std::memory_order_acquire); - int64_t diff = seq - static_cast(pos + count + 1); - atomic_ops++; // sequence.load - if (diff == 0) { - count++; - continue; - } - if (diff < 0) { - break; - } - contended = true; - count = -1; - break; - } - if (count == 0) { - atomic_count += atomic_ops; - return 0; - } - if (count < 0) { - continue; - } - if (dequeue_pos.compare_exchange_weak( - pos, pos + count, std::memory_order_relaxed, std::memory_order_relaxed - )) { - atomic_ops++; // successful CAS - break; - } - contended = true; - atomic_ops++; // failed CAS - } +size_t ready_queue_reserve_layout(DeviceArena &arena, uint64_t capacity); - for (int i = 0; i < count; i++) { - PTO2ReadyQueueSlot *slot = &slots[(pos + i) & mask]; - out[i] = slot->slot_state; - slot->sequence.store(static_cast(pos + i + mask + 1), std::memory_order_release); - atomic_ops++; // sequence.store +inline bool +ready_queue_init_data_from_layout(PTO2ReadyQueue *queue, DeviceArena &arena, size_t slots_off, uint64_t capacity); + +// Stores queue->slots = arena.region_ptr(slots_off). Idempotent. +void ready_queue_wire_arena_pointers(PTO2ReadyQueue *queue, DeviceArena &arena, size_t slots_off); + +void ready_queue_destroy(PTO2ReadyQueue *queue); + +struct alignas(64) PTO2SpscQueue { + // --- Producer cache lines (orchestrator thread) --- + alignas(64) std::atomic head_{0}; + alignas(64) uint64_t tail_cached_{0}; + + // --- Consumer cache lines (scheduler thread 0) --- + alignas(64) std::atomic tail_{0}; + alignas(64) uint64_t head_cached_{0}; + + // --- Shared Cacheline (read only) with mask and data ptr (immutable after init) --- + alignas(64) PTO2TaskSlotState **buffer_{nullptr}; + uint64_t mask_{0}; + + // Padding to exactly 5 cache lines + char padding[64 - sizeof(PTO2TaskSlotState **) - sizeof(uint64_t)]; + + static size_t reserve_layout(DeviceArena &arena, uint64_t capacity) { + return arena.reserve(capacity * sizeof(PTO2TaskSlotState *), PTO2_ALIGN_SIZE); + } + + bool init_data_from_layout(DeviceArena &arena, size_t buffer_off, uint64_t capacity) { + if (capacity == 0 || (capacity & (capacity - 1)) != 0) return false; + auto *buf = static_cast(arena.region_ptr(buffer_off)); + // calloc'd-equivalent: zero the slot pointers so spurious early pops + // observe nullptr. + for (uint64_t i = 0; i < capacity; i++) + buf[i] = nullptr; + mask_ = capacity - 1; + head_.store(0, std::memory_order_relaxed); + tail_.store(0, std::memory_order_relaxed); + tail_cached_ = 0; + head_cached_ = 0; + return true; + } + + // Wire the arena-internal pointer. Called by both host (with host arena) + // and AICPU (with device arena attached to the prebuilt image). + void wire_arena_pointers(DeviceArena &arena, size_t buffer_off) { + buffer_ = static_cast(arena.region_ptr(buffer_off)); + } + + void reset_for_reuse() { + uint64_t h = head_.load(std::memory_order_relaxed); + tail_.store(h, std::memory_order_relaxed); + tail_cached_ = h; + head_cached_ = h; + } + + // Arena owns the buffer; here we only forget our pointer. + void destroy() { buffer_ = nullptr; } + + bool push(PTO2TaskSlotState *item) { + uint64_t h = head_.load(std::memory_order_relaxed); + uint64_t next_h = h + 1; + if (next_h - tail_cached_ > mask_) { + tail_cached_ = tail_.load(std::memory_order_acquire); + if (next_h - tail_cached_ > mask_) return false; } - atomic_count += atomic_ops; - if (contended) { - wait_cycle += (get_sys_cnt_aicpu() - t0); + buffer_[h & mask_] = item; + head_.store(next_h, std::memory_order_release); + return true; + } + + // Pop up to max_count items (consumer only). Returns actual count. + int pop_batch(PTO2TaskSlotState **out, int max_count) { + uint64_t t = tail_.load(std::memory_order_relaxed); + uint64_t avail = head_cached_ - t; + if (avail < static_cast(max_count)) { + head_cached_ = head_.load(std::memory_order_acquire); + avail = head_cached_ - t; + if (avail == 0) return 0; } + int count = (avail < static_cast(max_count)) ? static_cast(avail) : max_count; + for (int i = 0; i < count; i++) + out[i] = buffer_[(t + i) & mask_]; + tail_.store(t + count, std::memory_order_release); return count; } -#endif + + // Approximate size (used for backoff decisions, not exact). + uint64_t size() const { + uint64_t h = head_.load(std::memory_order_acquire); + uint64_t t = tail_.load(std::memory_order_acquire); + return h - t; + } }; -// Cold-path ready queue operations (defined in pto_scheduler.cpp). Declared -// as non-member so PTO2ReadyQueue stays a POD-like struct with cache-line -// alignment. Storage is owned by the caller-supplied arena. -// reserve_layout: declare the slots[] region on the arena (must precede commit) -// init_from_layout: bind slots pointer from arena.region_ptr(off) and -// initialize sequence counters -// destroy: forget the slots pointer (arena owns the buffer) -size_t ready_queue_reserve_layout(DeviceArena &arena, uint64_t capacity); -// Writes everything *except* the arena-internal `slots` pointer field -// (sequences/positions on the slot array, capacity, mask). Uses -// arena.region_ptr(slots_off) only to address the slot array for writes; -// does NOT store the pointer in `queue->slots`. Call -// `ready_queue_wire_arena_pointers` afterwards to set the field itself. -bool ready_queue_init_data_from_layout(PTO2ReadyQueue *queue, DeviceArena &arena, size_t slots_off, uint64_t capacity); -// Stores queue->slots = arena.region_ptr(slots_off). Idempotent. -void ready_queue_wire_arena_pointers(PTO2ReadyQueue *queue, DeviceArena &arena, size_t slots_off); -void ready_queue_destroy(PTO2ReadyQueue *queue); +static_assert(sizeof(PTO2SpscQueue) == 5 * 64, "PTO2SpscQueue must be exactly 5 cache lines (320B)"); +// ============================================================================= -/** - * Statistics returned by mixed-task completion processing - */ struct CompletionStats { int32_t fanout_edges; // Number of fanout edges traversed (notify consumers) int32_t tasks_enqueued; // Number of consumers that became READY @@ -410,92 +306,48 @@ struct CompletionStats { bool mixed_task_completed; // True only when this callback completed a mixed task }; -/** - * Layout descriptor produced by PTO2SchedulerState::reserve_layout(). Holds - * the arena offsets of every sub-region the scheduler needs plus the - * capacities used at layout time (init_from_layout reuses them). - */ struct PTO2SchedulerLayout { size_t off_ready_queue_slots[PTO2_NUM_RESOURCE_SHAPES]; - size_t off_ready_sync_queue_slots[PTO2_NUM_RESOURCE_SHAPES]; 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[PTO2_MAX_RING_DEPTH]; + size_t off_pending_spsc_buffer; uint64_t ready_queue_capacity; - int32_t dep_pool_capacities[PTO2_MAX_RING_DEPTH]; + uint64_t spsc_capacity; }; -/** - * Scheduler state structure - * - * Contains dynamic state updated during task execution. - * Separated from shared memory for cache efficiency. - * Hot-path methods are defined inline (implicitly inline as member functions). - */ struct PTO2SchedulerState { // Shared memory access PTO2SharedMemoryHeader *sm_header; // Per-ring state struct alignas(64) RingSchedState { - // --- Cache Line 0: ring pointer (read-only) + hot path (read-write) --- PTO2SharedMemoryRingHeader *ring; 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. + std::atomic advance_lock; + // multi-thread CAS + bool init_data_from_layout(void *sm_dev_base, int32_t ring_id); - void reset_for_reuse(void *sm_dev_base, int32_t ring_id, std::atomic *orch_err); + void destroy(); void sync_to_sm() { ring->fc.last_task_alive.store(last_task_alive, std::memory_order_release); } -#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 - void advance_ring_pointers() { - int32_t current_task_index = ring->fc.current_task_index.load(std::memory_order_acquire); + const int32_t watermark = ring->completed_watermark.load(std::memory_order_acquire); int32_t old_last_task_alive = last_task_alive; - while (last_task_alive < current_task_index) { + // Retire any slot at the tail whose last consumer is at or below + // the global completed watermark — i.e. every consumer of this + // producer has reached COMPLETED. Implies this slot itself is + // COMPLETED because the seed value of last_consumer_local_id is + // the slot's own local_id. + while (last_task_alive <= watermark) { PTO2TaskSlotState &slot_state = ring->get_slot_state_by_task_id(last_task_alive); - if (slot_state.task_state.load(std::memory_order_acquire) != PTO2_TASK_CONSUMED) { - break; - } + if (watermark < slot_state.last_consumer_local_id) break; last_task_alive++; } - // Eager reset: prepare reclaimed slots for reuse while still hot in cache. - // Safe because last_task_alive has advanced past these slots but - // sync_to_sm has not yet published — the orchestrator cannot reuse - // them until the release store below. - // Skips payload, task, ring_id — immutable after RingSchedState::init(). - for (int32_t id = old_last_task_alive; id < last_task_alive; id++) { + for (int32_t id = old_last_task_alive; id < last_task_alive; id++) ring->get_slot_state_by_task_id(id).reset_for_reuse(); - } sync_to_sm(); } @@ -504,896 +356,320 @@ struct PTO2SchedulerState { // Ready queues remain global (scheduling is ring-agnostic) PTO2ReadyQueue ready_queues[PTO2_NUM_RESOURCE_SHAPES]; - // Ready sync_start queues, one per shape. A ready sync_start cohort parks here - // instead of ready_queues[] so the dispatch loop can drain it as a strict Tier-0 - // (sync_start > MIX > C/V) before any regular ready task takes a core, while - // reusing the same per-shape dispatch_shape machinery (fits-local inline vs - // stop-the-world drain, per-core MIX placement, head-start spacing). - PTO2ReadyQueue ready_sync_queues[PTO2_NUM_RESOURCE_SHAPES]; - // Dependency-only tasks (active_mask is empty, shape == DUMMY). Drained by // the dispatch loop and completed inline -- never goes to AICore. PTO2ReadyQueue dummy_ready_queue; + // Thread 0 exclusive: bounded SPSC drain → classify → route. The + // orchestrator pushes slot_states into the SPSC queue; thread 0 drains + // a batch per scheduler iter, classifies each task's fanin state, and + // routes terminally — either to a ready queue (all fanins met) or onto + // a producer's wake_list (first unmet). No intermediate FIFO: each + // drained task is classified once, never re-queued. The wake-list-only + // redesign made classify_fanin_state's decision terminal, so the + // previously-needed pending FIFO became dead weight on the critical + // path. + struct alignas(64) PendingState { + static constexpr int BACKOFF_LIMIT = 32; + static constexpr int DRAIN_BATCH = 30; + + // --- Thread 0 exclusive --- + int backoff_counter{0}; + PTO2TaskSlotState *drain_buf[DRAIN_BATCH]; + + // --- SPSC queue: orchestrator (push) ↔ thread 0 (pop) --- + PTO2SpscQueue queue; + + // --- Orchestrator write, thread 0 read --- + alignas(64) std::atomic orch_needs_drain{false}; + } wiring; + alignas(64) AsyncWaitList async_wait_list; - // Statistics (cold path, isolated from hot-path fields) -#if SIMPLER_SCHED_PROFILING - alignas(64) std::atomic tasks_completed; - std::atomic tasks_consumed; -#endif - // ========================================================================= - // Inline hot-path methods - // ========================================================================= - - // Route a ready slot to the right global queue. Dep-only tasks — DUMMY-shaped - // (empty active_mask) or a task whose dispatch predicate fails — live in - // dummy_ready_queue and are retired inline; a ready sync_start cohort goes to - // the per-shape ready_sync_queues[] (drained as Tier-0); everything else to - // ready_queues[]. void push_ready_routed(PTO2TaskSlotState *slot_state) { 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); - } - } - - 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; - } - } - } - - void check_and_handle_consumed(PTO2TaskSlotState &slot_state) { - // Read fanout_refcount/fanout_count and flip COMPLETED->CONSUMED under - // fanout_lock. The orchestrator claims producers (fanout_count++) under the - // same lock, so the consume decision is serialized against a concurrent - // claim: either the ++ lands first (count then exceeds refcount, so we do - // not consume and the producer stays pinned until released) or the consume - // lands first (the orchestrator then observes CONSUMED and skips the - // claim). Without this lock a claim racing the consume desyncs the slot's - // refcount and wedges in-order reclaim. - bool became_consumed = false; - slot_state.lock_fanout(); - if (slot_state.fanout_refcount.load(std::memory_order_acquire) == slot_state.fanout_count) { - PTO2TaskState expected = PTO2_TASK_COMPLETED; - became_consumed = slot_state.task_state.compare_exchange_strong( - expected, PTO2_TASK_CONSUMED, std::memory_order_acq_rel, std::memory_order_acquire - ); - } - slot_state.unlock_fanout(); - if (!became_consumed) return; - -#if SIMPLER_SCHED_PROFILING - tasks_consumed.fetch_add(1, std::memory_order_relaxed); -#endif - - int32_t ring_id = slot_state.ring_id; - // advance_ring_pointers (and the reset_for_reuse it triggers) MUST run - // outside fanout_lock: reset_for_reuse stores fanout_lock=0 and would - // clobber a held lock. Safe here — the slot is CONSUMED and quiescent. - // Try-lock — if another thread is advancing this ring, it will scan our CONSUMED task - int32_t expected_lock = 0; - if (ring_sched_states[ring_id].advance_lock.compare_exchange_strong( - expected_lock, 1, std::memory_order_acquire, std::memory_order_relaxed - )) { - ring_sched_states[ring_id].advance_ring_pointers(); - ring_sched_states[ring_id].advance_lock.store(0, std::memory_order_release); - } + if (shape == PTO2ResourceShape::DUMMY) dummy_ready_queue.push(slot_state); + else ready_queues[static_cast(shape)].push(slot_state); } -#if SIMPLER_ORCH_PROFILING || SIMPLER_SCHED_PROFILING - void check_and_handle_consumed(PTO2TaskSlotState &slot_state, uint64_t &atomic_count) { - // See the non-profiling overload for why the read + COMPLETED->CONSUMED - // flip is serialized against the orchestrator's claim under fanout_lock. - bool became_consumed = false; - slot_state.lock_fanout(); - atomic_count += 1; // lock CAS - uint32_t fc = slot_state.fanout_count; - uint32_t rc = slot_state.fanout_refcount.load(std::memory_order_acquire); - atomic_count += 1; // fanout_refcount.load (fanout_count is a plain read under lock) - if (rc == fc) { - PTO2TaskState expected = PTO2_TASK_COMPLETED; - became_consumed = slot_state.task_state.compare_exchange_strong( - expected, PTO2_TASK_CONSUMED, std::memory_order_acq_rel, std::memory_order_acquire - ); - atomic_count += 1; // CAS + bool fanin_satisfied(PTO2TaskSlotState *s) const { + const PTO2TaskPayload &p = *s->payload; + for (int32_t i = 0; i < p.fanin_count; i++) { + const auto &prod_ring = *ring_sched_states[p.fanin_ring_ids[i]].ring; + if (prod_ring.completion_flags[p.fanin_local_ids[i] & prod_ring.task_window_mask].load( + std::memory_order_acquire + ) == 0) + return false; } - slot_state.unlock_fanout(); - atomic_count += 1; // unlock store - if (!became_consumed) return; - -#if SIMPLER_SCHED_PROFILING - tasks_consumed.fetch_add(1, std::memory_order_relaxed); -#endif - - int32_t ring_id = slot_state.ring_id; - // advance_ring_pointers + reset_for_reuse run outside fanout_lock (reset - // stores fanout_lock=0). Safe — the slot is CONSUMED and quiescent. - // Try-lock — if another thread is advancing this ring, it will scan our CONSUMED task - int32_t expected_lock = 0; - if (ring_sched_states[ring_id].advance_lock.compare_exchange_strong( - expected_lock, 1, std::memory_order_acquire, std::memory_order_relaxed - )) { - ring_sched_states[ring_id].advance_ring_pointers(); - ring_sched_states[ring_id].advance_lock.store(0, std::memory_order_release); - atomic_count += 2; // try-lock CAS + unlock store - } else { - atomic_count += 1; // failed try-lock CAS - } - } -#endif - - void release_producer(PTO2TaskSlotState &slot_state) { - slot_state.fanout_refcount.fetch_add(1, std::memory_order_acq_rel); - check_and_handle_consumed(slot_state); - } - - // 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); - check_and_handle_consumed(slot_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 - check_and_handle_consumed(slot_state, atomic_count); - } - - 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 - check_and_handle_consumed(slot_state, atomic_count); - } -#endif - - // 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 - // the completion path — the moment its last producer's FIN satisfies fanin — - // instead of routing it through the ready queue and waiting for the dispatch - // pass to pop it. Returns true if the task is fully handled (caller must NOT - // push to the ready queue). Returns false when the caller must route C - // normally: either it was never pre-staged, OR it is a SPMD consumer only - // PARTIALLY pre-staged — the gated blocks are released by the doorbells rung - // here, and the remaining (next_block_idx .. logical_block_num) blocks - // dispatch normally off the ready queue. Lock-free claim shared with Hook 1 - // (the stager): CAS NONE->DISPATCHED wins => not pre-staged; otherwise flip - // STAGING->DISPATCHED and destructively claim the published doorbell bits. - - // Per-core early-dispatch doorbell table. Hook 1 records each gated core's - // (reg_addr, dispatch token) here at stage time; the completion-path release - // reads it back for the cores set in the consumer's staged_core_mask. One - // global table indexed by core_id (not per-task): gated cores in flight are - // bounded by the chip's core count (no two-level pre-dispatch), so this is the - // natural capacity and removes the old per-task 3-doorbell cap. - struct EarlyDispatchDoorbell { - uint64_t addr{0}; - uint32_t token{0}; - }; - EarlyDispatchDoorbell early_dispatch_doorbell_table[PTO2_EARLY_DISPATCH_CORE_MASK_WORDS * 64]{}; - - // Cross-thread early-dispatch work queues, one PTO2ReadyQueue MPMC instance per - // resource shape (AIC/AIV/MIX) — arena-backed, reserved/wired in pto_runtime2_init - // alongside the per-shape ready queues, and indexed the same way. A candidate is - // pushed to the queue for its own shape (active_mask.to_shape()) so the drain can - // pop per shape and size the pop to that shape's free cores, exactly as normal - // dispatch pops ready_queues[shape]. - // - // A consumer's SPMD blocks span cores owned by several AICPU threads, but only a - // thread RUNNING the consumer's producer discovers it (via the producer's - // fanout). When that producer is thread-local (e.g. a 16-block AIV op filling one - // thread's cores), the other threads never see the consumer and its blocks on - // their cores can't pre-stage. The first claimer pushes the partially-staged - // consumer here; every idle thread's early_dispatch pass pops one, stages a range - // onto ITS OWN cores (range-claim via next_block_idx), and re-pushes if blocks - // remain — exactly mirroring how a partially-dispatched SPMD task is re-pushed to - // the ready queue (scheduler_dispatch: pop -> claim -> re-push). A stale/released - // entry fails the STAGING check on pop and is dropped; a push that overflows is - // logged and the consumer's blocks fall back to normal dispatch. - PTO2ReadyQueue early_dispatch_queues[PTO2_NUM_RESOURCE_SHAPES]; - - // sync_start early-dispatch candidates park here instead of early_dispatch_queues[]: - // they need an atomic all-or-nothing stage via the drain barrier, not - // early_dispatch_shape's per-thread partial range-claim. Shape-agnostic (the - // rendezvous counts cores, not blocks), so a single queue serves all shapes; drained - // as the highest occupancy tier at the top of try_early_dispatch. - // - // Deliberately single, vs the normal source's per-shape ready_sync_queues[]: a READY - // sync cohort (producer done) can dispatch inline when it fits, so it reuses the - // per-shape dispatch_shape; an EARLY sync cohort always carries a non-zero - // src_payload gate and therefore always drains. The shape-agnostic rendezvous - // makes one queue sufficient. Same drain, two sources. - PTO2ReadyQueue early_sync_start_queue; - - static inline void ring_one_doorbell(uint64_t reg_addr, uint32_t token) { - volatile uint64_t *dmb = reinterpret_cast(get_reg_ptr(reg_addr, RegId::DATA_MAIN_BASE)); - uint64_t tk = static_cast(token); - *dmb = (tk << 32) | tk; // 64-bit STR: high=low=token releases the gated AICore - } - - inline void ring_staged_doorbell_bits(int word, uint64_t bits) { - while (bits != 0) { - int core_id = word * 64 + __builtin_ctzll(bits); - bits &= bits - 1; - ring_one_doorbell( - early_dispatch_doorbell_table[core_id].addr, early_dispatch_doorbell_table[core_id].token - ); - } - } - - static inline uint64_t claim_all_staged_doorbell_bits(std::atomic &mask) { - return mask.exchange(0, std::memory_order_seq_cst); - } - - static inline uint64_t claim_late_staged_doorbell_bits(std::atomic &mask, uint64_t candidates) { - return mask.fetch_and(~candidates, std::memory_order_seq_cst) & candidates; - } - - static inline bool should_gate_early_dispatch(bool force_gate, uint8_t early_dispatch_state) { - return force_gate || early_dispatch_state == PTO2_EARLY_DISPATCH_STAGING; - } - - static inline bool - ring_claimed_local_doorbell(uint64_t claimed_word, int core_id, uint64_t reg_addr, uint32_t token) { - if ((claimed_word & (1ULL << (core_id & 63))) == 0) return false; - ring_one_doorbell(reg_addr, token); return true; } - static inline bool try_claim_early_dispatch_launch(PTO2TaskPayload &payload) { - uint8_t expected = PTO2_EARLY_DISPATCH_LAUNCH_NONE; - return payload.early_dispatch_launch_state.compare_exchange_strong( - expected, PTO2_EARLY_DISPATCH_LAUNCH_RINGING, std::memory_order_seq_cst, std::memory_order_seq_cst - ); - } - - inline void record_published_blocks(PTO2TaskSlotState &slot_state, int32_t count) { - if (count <= 0 || !slot_state.allow_early_resolve) return; - slot_state.payload->published_block_count.fetch_add(static_cast(count), std::memory_order_seq_cst); - } - - // Ring one sync_start cohort from its stable staged_core_mask. The caller owns - // the NONE->RINGING launch latch and invokes this exactly once after drain - // staging completes, while the corresponding per-core table entries are live. - inline void ring_all_staged_doorbells(PTO2TaskSlotState &slot_state) { - for (int w = 0; w < PTO2_EARLY_DISPATCH_CORE_MASK_WORDS; w++) { - uint64_t bits = slot_state.payload->staged_core_mask[w].load(std::memory_order_seq_cst); - while (bits != 0) { - int core_id = w * 64 + __builtin_ctzll(bits); - bits &= bits - 1; - ring_one_doorbell( - early_dispatch_doorbell_table[core_id].addr, early_dispatch_doorbell_table[core_id].token - ); + // First-unmet classification used by the wiring-queue drain and the + // wake_list rescan. Returns: + // -1: all fanins met (route directly to ready) + // ≥0: index of the first unmet fanin (register on its producer's + // wake list). Decision is terminal — tasks are never re-queued + // for polling; rescans happen lazily on producer completion via + // on_mixed_task_complete's wake_list drain. + int classify_fanin_state(PTO2TaskSlotState *s) const { + const PTO2TaskPayload &p = *s->payload; + for (int32_t i = 0; i < p.fanin_count; i++) { + const auto &prod_ring = *ring_sched_states[p.fanin_ring_ids[i]].ring; + if (prod_ring.completion_flags[p.fanin_local_ids[i] & prod_ring.task_window_mask].load( + std::memory_order_acquire + ) == 0) { + return i; } } - } - - static inline bool try_claim_early_sync_drain(PTO2TaskPayload &payload) { - uint8_t expected = PTO2_EARLY_SYNC_DRAIN_NONE; - return payload.early_sync_drain_state.compare_exchange_strong( - expected, PTO2_EARLY_SYNC_DRAIN_OWNER, std::memory_order_seq_cst, std::memory_order_seq_cst - ); - } - - static inline bool owns_early_sync_drain(const PTO2TaskPayload &payload) { - return (payload.early_sync_drain_state.load(std::memory_order_acquire) & PTO2_EARLY_SYNC_DRAIN_OWNER) != 0; - } - - static inline void mark_early_sync_drain_armed(PTO2TaskPayload &payload) { - payload.early_sync_drain_state.fetch_or(PTO2_EARLY_SYNC_DRAIN_ARMED, std::memory_order_seq_cst); - } - - static inline bool publish_ready_to_early_sync_drain(PTO2TaskPayload &payload) { - uint8_t previous = - payload.early_sync_drain_state.fetch_or(PTO2_EARLY_SYNC_DRAIN_READY, std::memory_order_seq_cst); - return (previous & PTO2_EARLY_SYNC_DRAIN_OWNER) != 0; - } - - inline void cancel_early_sync_drain(PTO2TaskSlotState &slot_state) { - uint8_t previous = - slot_state.payload->early_sync_drain_state.exchange(PTO2_EARLY_SYNC_DRAIN_NONE, std::memory_order_seq_cst); - if ((previous & PTO2_EARLY_SYNC_DRAIN_OWNER) == 0) return; - if ((previous & PTO2_EARLY_SYNC_DRAIN_READY) != 0) { - push_ready_routed(&slot_state); - return; - } - if (slot_state.payload->early_dispatch_state.load(std::memory_order_seq_cst) == PTO2_EARLY_DISPATCH_STAGING) { - early_sync_start_queue.push_tagged(&slot_state, static_cast(slot_state.task->task_id.raw)); - } - } - - static inline void finish_early_sync_drain(PTO2TaskPayload &payload) { - uint8_t state = payload.early_sync_drain_state.load(std::memory_order_seq_cst); - while ((state & PTO2_EARLY_SYNC_DRAIN_OWNER) != 0 && (state & PTO2_EARLY_SYNC_DRAIN_COMPLETE) == 0) { - uint8_t desired = state | PTO2_EARLY_SYNC_DRAIN_COMPLETE; - if (payload.early_sync_drain_state.compare_exchange_weak( - state, desired, std::memory_order_seq_cst, std::memory_order_seq_cst - )) { - return; + return -1; + } + + // (e) Register `consumer` on `producer`'s wake list. If the producer has + // already completed (head == WAKE_LIST_SENTINEL) we must NOT assume the + // consumer is ready: classify_fanin_state() short-circuits at the FIRST + // unmet fanin, so the producer handed to us is only "an" unmet fanin, not + // necessarily the last one. A producer that completes in the window between + // that classify and this call would otherwise let us push a consumer whose + // *later* fanins are still pending — a premature dispatch that lets a + // dependent task run against a not-yet-produced input (e.g. the + // paged_attention online-softmax UP chain: UP_bn dispatched before UP_{bn-1} + // finishes writing the shared accumulator -> concurrent RMW -> wrong query). + // On the SENTINEL path we re-classify against ALL fanins and only route to + // ready when every fanin is satisfied; otherwise we re-target the next + // still-unmet producer and retry. Monotonic completion_flags guarantee + // termination. + void register_wake(PTO2TaskSlotState *producer, PTO2TaskSlotState *consumer) { + 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; // registered on a still-pending producer + } + // CAS failed: expected reloaded; retry (may now be SENTINEL). } - } - } - // sync_start rendezvous: a sync_start consumer's gated cores launch as an atomic - // cohort, so their doorbells are held until BOTH halves hold — every gated core - // occupies a running slot (running_slot_count == popcount(staged_core_mask)) AND the - // producer released (early_dispatch_state == DISPATCHED). Counting CORES (not blocks) makes - // this shape-agnostic: an AIC/AIV block is one core, a MIX block is a cluster whose - // cores promote pending->running independently. Called from both halves (the producer - // release and each pending->running promotion); whichever observes the second half - // wins the launch latch and rings exactly once. Returns true only to that winner, - // which may then expose the cohort to its fanout. - inline bool maybe_rendezvous_ring(PTO2TaskSlotState &slot_state) { - int32_t staged_cores = 0; - for (int w = 0; w < PTO2_EARLY_DISPATCH_CORE_MASK_WORDS; w++) - staged_cores += - __builtin_popcountll(slot_state.payload->staged_core_mask[w].load(std::memory_order_seq_cst)); - if (staged_cores == 0) return false; - if (slot_state.payload->running_slot_count.load(std::memory_order_seq_cst) != staged_cores) return false; - if (slot_state.payload->early_dispatch_state.load(std::memory_order_seq_cst) != PTO2_EARLY_DISPATCH_DISPATCHED) - return false; - if (!try_claim_early_dispatch_launch(*slot_state.payload)) return false; - ring_all_staged_doorbells(slot_state); - wmb(); - slot_state.payload->early_dispatch_launch_state.store( - PTO2_EARLY_DISPATCH_LAUNCH_COMPLETE, std::memory_order_release - ); - return true; - } - - inline bool retry_sync_start_rendezvous_after_drain(PTO2TaskSlotState &slot_state) { - if (!maybe_rendezvous_ring(slot_state)) return false; - propagate_dispatch_fanin(slot_state); - return true; - } - - // Attempt to claim and enqueue a consumer after every direct producer becomes - // launch-visible. Both propagation and wiring can complete dispatch_fanin. - inline void try_enqueue_early_dispatch_candidate(PTO2TaskSlotState &consumer) { - // Predicated tasks resolve at the ready point, never early-dispatch. The - // wiring caller has no separate predicate filter, so the gate lives here. - if (consumer.active_mask.has_predicate()) return; - PTO2ResourceShape shape = consumer.active_mask.to_shape(); - if (shape != PTO2ResourceShape::AIC && shape != PTO2ResourceShape::AIV && shape != PTO2ResourceShape::MIX) { - return; - } - - uint8_t expected = PTO2_EARLY_DISPATCH_NONE; - if (!consumer.payload->early_dispatch_state.compare_exchange_strong( - expected, PTO2_EARLY_DISPATCH_STAGING, std::memory_order_seq_cst, std::memory_order_seq_cst - )) { - return; - } - - uint64_t task_id = static_cast(consumer.task->task_id.raw); - // A sync-start cohort uses one shape-agnostic queue so only the - // all-or-nothing drain path can claim and stage it. - if (consumer.active_mask.requires_sync_start()) { - early_sync_start_queue.push_tagged(&consumer, task_id); - } else { - early_dispatch_queues[static_cast(shape)].push_tagged(&consumer, task_id); - } - } - - // 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; + // Producer completed. Re-check every fanin before committing. + int32_t state = classify_fanin_state(consumer); + if (state < 0) { + push_ready_routed(consumer); // all fanins now satisfied + return; } - if (was_pre_staged && launch_state != PTO2_EARLY_DISPATCH_LAUNCH_COMPLETE) return; - } - // The once-claim and fanout snapshot share fanout_lock with wiring, so a - // set marker tells wiring that its new edge is outside this snapshot. - // The unlocked precheck avoids fanout-lock traffic after propagation. - if (p.has_dispatch_propagated()) return; - p.lock_fanout(); - if (!p.try_mark_dispatch_propagated()) { - p.unlock_fanout(); - return; - } - // Nodes reachable from the captured head are immutable; wiring serializes - // later prepends under fanout_lock and seeds them separately. - PTO2DepListEntry *edge = p.fanout_head; - 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; - try_enqueue_early_dispatch_candidate(*c); + // Re-target the next still-unmet producer and retry. + const PTO2TaskPayload &p = *consumer->payload; + producer = + &ring_sched_states[p.fanin_ring_ids[state]].ring->get_slot_state_by_task_id(p.fanin_local_ids[state]); } } - // 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); + // Thread 0 entry point: drain a bounded batch from the orchestrator's + // SPSC queue, then classify+route each drained task terminally. Returns + // the count of routed tasks (also the drained count — each drained task + // is classified once and never re-queued). + // + // Sub-phase timing pointers (optional). If non-null, cumulative cycle/ + // iteration counters for Stage 1 (SPSC drain) and Stage 2 (classify+route) + // are accumulated into them. + int drain_wiring_queue(bool force_drain = false) { + // Stage 1: drain SPSC → drain_buf + int drained = wiring.queue.pop_batch(wiring.drain_buf, PendingState::DRAIN_BATCH); + + // Backoff when nothing to do and orchestrator isn't pressing + if (drained == 0) { + if (!force_drain && !wiring.orch_needs_drain.load(std::memory_order_acquire) && + wiring.backoff_counter < PendingState::BACKOFF_LIMIT) { + wiring.backoff_counter++; + return 0; } } - // 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; + wiring.backoff_counter = 0; + + // Stage 2: classify + route each drained task in-line. Each task's + // state is "all met → ready_queue" or "first unmet → register on that + // producer's wake_list". Tasks are scanned exactly once here; + // re-scans on producer completion happen via on_mixed_task_complete's + // wake_list drain. + for (int i = 0; i < drained; i++) { + PTO2TaskSlotState *s = wiring.drain_buf[i]; + int state = classify_fanin_state(s); + if (state < 0) { + push_ready_routed(s); + } else { + // Producer is in fanin_ring_ids[state] (may differ from + // the consumer's ring under multi-ring fanin). When the + // producer completes, its wake_list drain rescans this + // consumer and either pushes to ready or re-registers on + // the next unmet producer. + int32_t prod_local = s->payload->fanin_local_ids[state]; + uint8_t prod_ring = s->payload->fanin_ring_ids[state]; + auto &ring = *ring_sched_states[prod_ring].ring; + PTO2TaskSlotState *producer = &ring.get_slot_state_by_task_id(prod_local); + register_wake(producer, s); } - 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; + return drained; } -#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 - - int get_ready_tasks_batch(PTO2ReadyQueue *queues, PTO2ResourceShape shape, PTO2TaskSlotState **out, int max_count) { - return queues[static_cast(shape)].pop_batch(out, max_count); - } - -#if SIMPLER_SCHED_PROFILING int get_ready_tasks_batch( - PTO2ReadyQueue *queues, PTO2ResourceShape shape, PTO2TaskSlotState **out, int max_count, uint64_t &atomic_count, - uint64_t &wait_cycle + PTO2ResourceShape shape, PTO2LocalReadyBuffer &local_buf, PTO2TaskSlotState **out, int max_count ) { - return queues[static_cast(shape)].pop_batch(out, max_count, atomic_count, wait_cycle); - } -#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 + int count = 0; + while (count < max_count && local_buf.count > 0) + out[count++] = local_buf.slot_states[--local_buf.count]; + int remaining = max_count - count; + if (remaining > 0) count += ready_queues[static_cast(shape)].pop_batch(out + count, remaining); + return count; } - /** - * Subtask completion: atomic counter model. - * Called when a single subtask (AIC, AIV0, or AIV1) finishes on any block. - * Atomically increments completed_subtasks and checks whether all subtasks - * across all blocks are done. - * - * @return true if this was the last subtask, completing the entire task. - */ bool on_subtask_complete(PTO2TaskSlotState &slot_state) { - int16_t prev = slot_state.completed_subtasks.fetch_add(1, std::memory_order_acq_rel); + // Relaxed fetch_add: completed_subtasks is a pure counter with no + // other observers piggybacking state through it. The only readers + // are this fetch_add itself (per-subtask) and reset_for_reuse's + // relaxed init. Real publication of the producer's completion to + // consumer threads happens downstream in on_mixed_task_complete via + // completion_flag.store(release) + wake_list_head.exchange(acq_rel) + // — those are the AICPU↔AICPU sync edges. The producer→consumer + // GM data ordering is handled by AICore-side cache coherence + // independent of this counter's ordering. + int16_t prev = slot_state.completed_subtasks.fetch_add(1, std::memory_order_relaxed); return (prev + 1) == slot_state.total_required_subtasks; } - /** - * Two-stage completion: second stage. - * Called exactly once when all subtasks of a task are done (i.e., - * on_subtask_complete returned true). Walks the consumer (fanout) list, - * decrements each consumer's fanin, pushes newly-ready ones, and rings - * doorbells for early-dispatch hits. - * - * Non-PROFILING returns the consumer-walk count (= edges traversed). The - * Resolve swimlane bar reads it to label the bar with how many successors - * actually got resolved. PROFILING returns the richer CompletionStats - * whose `fanout_edges` carries the same number. - */ -#if SIMPLER_SCHED_PROFILING - CompletionStats -#else - uint32_t -#endif - on_task_complete( - PTO2TaskSlotState &slot_state -#if SIMPLER_SCHED_PROFILING - , - int thread_idx -#endif - ) { -#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); -#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++; + // Publish this slot as COMPLETED, then advance the per-ring monotonic + // completed_watermark — the highest local_id W such that every task + // 0..W has reached COMPLETED. Reclamation in advance_ring_pointers gates + // on watermark >= producer.last_consumer_local_id, so no consumer→producer + // notification edge is needed. + void on_mixed_task_complete(PTO2TaskSlotState &slot_state) { + // (m) Skip slot_state.task_state.store here; completion_flags below is + // the single source of truth. Saves one atomic release store per task. + const int32_t my_id = static_cast(slot_state.task->task_id.local()); + int32_t ring_id = slot_state.ring_id; + auto &rss = ring_sched_states[ring_id]; + auto &ring = *rss.ring; + + // Publish to the polling-fast completion array. Release ordering + // makes the producer's output writes visible to consumers that + // acquire-load this byte in fanin_satisfied. + ring.completion_flags[my_id & ring.task_window_mask].store(1, std::memory_order_release); + + // Drain the wake list. Each consumer registered on this slot was + // waiting on at least one unmet fanin (this one). After + // completion_flag is set above, atomic-exchange wake_list_head to + // SENTINEL (refusing any future registrations) and process each + // waiter: rescan its fanin, route to ready_queue if all met, else + // re-register on the new first-unmet producer. Ordering: + // completion_flag is set BEFORE the exchange, so any consumer that + // races a registration against our exchange and observes a SENTINEL + // during retry will see completion_flag=1 and either rescan-and-route + // or self-register on the next unmet. + 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; + // next_in_wake_list left as-is: every re-registration via + // register_wake() overwrites the field before the CAS publishes + // the consumer, and reset_for_reuse() clears it on slot reuse. + // No reader between here and the next overwrite/reset. + // Fast path: single-fanin waiters were waiting on *us* (the only + // possible fanin). No rescan needed — push straight to ready. + // Saves one classify_fanin_state call (a byte read in + // completion_flags) per waiter. Skips the cache-miss-prone + // multi-ring lookup for the common chain-task case where each + // task has exactly one predecessor. + if (waiter->payload->fanin_count == 1) { + push_ready_routed(waiter); + waiter = next; + continue; + } + int state = classify_fanin_state(waiter); + if (state < 0) { + push_ready_routed(waiter); + } else { + // Still some fanin unmet — re-register on the new first + // unmet producer's wake list. + int32_t prod_local = waiter->payload->fanin_local_ids[state]; + uint8_t prod_ring = waiter->payload->fanin_ring_ids[state]; + auto &prod_ring_hdr = *ring_sched_states[prod_ring].ring; + PTO2TaskSlotState *producer = &prod_ring_hdr.get_slot_state_by_task_id(prod_local); + register_wake(producer, waiter); + } + waiter = next; + } + + // CAS-advance the watermark, bounded by my_id (which we know is + // published since we just completed it). If a forward task we observe + // as COMPLETED is also published, but a gap remains, we stop — the + // task filling the gap will resume the walk when it completes. + int32_t w = ring.completed_watermark.load(std::memory_order_acquire); + while (w < my_id) { + 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; } -#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; -#endif + // Try to retire slots whose last consumer has reached COMPLETED. + // Gate the try-lock + advance walk on a lag threshold: most + // completions advance the watermark by 1 slot; firing the try-lock + // per completion costs ~10-30 ns × ~65K completions × N threads of + // wasted CAS attempts. With the gate, the try-lock fires ~32× less + // often. Empirically 32 is the sweet spot — bigger thresholds let + // the allocator stall more often waiting for reclamation. The lag + // read of last_task_alive is non-atomic but monotonic and only used + // as a hint — stale-but-OK. + if (w - rss.last_task_alive >= 32) { + int32_t expected_lock = 0; + if (rss.advance_lock.compare_exchange_strong( + expected_lock, 1, std::memory_order_acquire, std::memory_order_relaxed + )) { + rss.advance_ring_pointers(); + rss.advance_lock.store(0, std::memory_order_release); + } + } } - /** - * 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_self_atomic_count[]; - extern uint64_t g_sched_self_consumed_cycle[]; - 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 - - // Self consumed check -#if SIMPLER_SCHED_PROFILING - uint64_t self_atomics = 0; - check_and_handle_consumed(slot_state, self_atomics); - g_sched_self_atomic_count[thread_idx] += self_atomics; - PTO2_SCHED_CYCLE_LAP(g_sched_self_consumed_cycle[thread_idx]); - g_sched_complete_count[thread_idx]++; -#else - check_and_handle_consumed(slot_state); -#endif - return payload->fanin_actual_count; - } + // === Cold-path API === + + static PTO2SchedulerLayout reserve_layout(DeviceArena &arena, int32_t /*dep_pool_capacity*/); - // === Cold-path API (defined in pto_scheduler.cpp) === - - // Phase 1: declare every sub-region (ready_queue slots, dummy queue slots, - // 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, const int32_t dep_pool_capacities[PTO2_MAX_RING_DEPTH]); - static PTO2SchedulerLayout reserve_layout( - DeviceArena &arena, const int32_t dep_pool_capacities[PTO2_MAX_RING_DEPTH], - const int32_t task_window_sizes[PTO2_MAX_RING_DEPTH] - ); - - // Phase 3a: write everything *except* arena-internal pointer fields. - // `sm_dev_base` is the device address of the SM (only stored, never - // dereferenced here). Safe to call on a host arena that holds the - // prebuilt image buffer. (The orchestrator counterpart takes - // task_window_size for ring task_descriptors address arithmetic; the - // scheduler only needs the SM header / ring header base addresses, - // both window-size-independent.) bool init_data_from_layout(const PTO2SchedulerLayout &layout, DeviceArena &arena, void *sm_dev_base); - void reset_for_reuse(const PTO2SchedulerLayout &layout, void *sm_dev_base); - // Phase 3b: write the arena-internal pointer fields - // (ready_queues[].slots, dummy_ready_queue.slots, dep_pool.base for each - // ring). Called on both host and device sides. void wire_arena_pointers(const PTO2SchedulerLayout &layout, DeviceArena &arena); // Forget per-region pointers; arena owns the backing memory. void destroy(); - void print_stats(); - void print_queues(); + + // Surgical reset for arena reuse (#1234): resets per-run mutable state + // without redoing the O(ready_queue_capacity) buffer-zeroing that + // init_data_from_layout does. Ring pointer is re-set from sm_dev_base + // since we can't rely on the previous run's value being valid across + // arena reuse. + void reset_for_reuse(void *sm_dev_base); }; // Scheduler cold-path API is declared as PTO2SchedulerState member functions. -// See init()/destroy()/print_stats()/print_queues() below the struct definition. - -// try_inline_complete_locked: short-circuit NotDeferred completions seen during -// drain so they don't grow entries[]. Defined here (not in pto_async_wait.h) -// 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. +// See init()/destroy() below the struct definition. + inline bool AsyncWaitList::try_inline_complete_locked(AsyncWaitList::DrainCompletionSink &sink, PTO2TaskSlotState &slot_state) { - // Return value (CompletionStats / consumer-walk count) discarded: - // async-wait drain path has no Resolve swimlane bar attached. -#if SIMPLER_SCHED_PROFILING - (void)sink.sched->on_task_complete(slot_state, sink.thread_idx); -#else - (void)sink.sched->on_task_complete(slot_state); -#endif - 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.sched->on_mixed_task_complete(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 -#if SIMPLER_SCHED_PROFILING - , - int thread_idx -#endif -) { +inline AsyncPollResult +AsyncWaitList::poll_and_complete(AICoreCompletionMailbox *aicore_mailbox, PTO2SchedulerState *sched) { AsyncPollResult result; if (!try_lock()) return result; 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 int32_t drain_err = PTO2_ERROR_NONE; drain_aicore_completion_mailbox_locked(aicore_mailbox, sink, drain_err); @@ -1432,28 +708,7 @@ inline AsyncPollResult AsyncWaitList::poll_and_complete( } if (entry.normal_done && entry.waiting_completion_count <= 0) { - // Return value (CompletionStats / consumer-walk count) discarded: - // deferred-completion drain has no Resolve swimlane bar attached. -#if SIMPLER_SCHED_PROFILING - (void)sched->on_task_complete(*entry.slot_state, thread_idx); -#else - (void)sched->on_task_complete(*entry.slot_state); -#endif - // 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; + sched->on_mixed_task_complete(*entry.slot_state); result.completed++; int32_t last = count - 1; @@ -1465,37 +720,3 @@ inline AsyncPollResult AsyncWaitList::poll_and_complete( unlock(); return result; } - -// ============================================================================= -// Scheduler Profiling Data -// ============================================================================= - -#if SIMPLER_SCHED_PROFILING -struct PTO2SchedProfilingData { - // Sub-phase cycle breakdown within on_task_complete - uint64_t lock_cycle; // lock_fanout + state store + unlock - uint64_t fanout_cycle; // fanout traversal - uint64_t fanin_cycle; // fanin traversal - uint64_t self_consumed_cycle; // self check_and_handle_consumed - - // Wait times - uint64_t lock_wait_cycle; // spin-wait in fanout_lock - uint64_t push_wait_cycle; // CAS contention in push() - uint64_t pop_wait_cycle; // CAS contention in pop() - - // Atomic counts per sub-phase - uint64_t lock_atomic_count; - uint64_t fanout_atomic_count; - uint64_t fanin_atomic_count; - uint64_t self_atomic_count; - uint64_t pop_atomic_count; - - int64_t complete_count; -}; - -/** - * Get and reset scheduler profiling data for a specific thread. - * Returns accumulated profiling data and resets counters. - */ -PTO2SchedProfilingData scheduler_get_profiling(int thread_idx); -#endif diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_cold_path.cpp b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_cold_path.cpp index d15200ed01..c493ac691a 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_cold_path.cpp +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_cold_path.cpp @@ -8,161 +8,28 @@ * See LICENSE in the root of the software repository for the full text of the License. * ----------------------------------------------------------------------------------------------------------- */ + #include "scheduler_context.h" +#include #include -#include +#include -#include "common/unified_log.h" +#include "common.h" +#include "callable.h" +#include "aicpu/aicpu_device_config.h" #include "aicpu/dep_gen_collector_aicpu.h" -#include "aicpu/device_phase_aicpu.h" -#include "aicpu/device_time.h" -#include "aicpu/l2_swimlane_collector_aicpu.h" -#include "aicpu/platform_regs.h" -#include "aicpu/pmu_collector_aicpu.h" -#include "aicpu/args_dump_aicpu.h" -#include "common/memory_barrier.h" -#include "common/l2_swimlane_profiling.h" -#include "common/platform_config.h" -#include "pto_runtime2.h" -#include "pto_shared_memory.h" -#include "runtime.h" -#include "spin_hint.h" - -// ============================================================================= -// Cold-path helpers for the main dispatch loop (noinline to reduce hot-loop icache) -// ============================================================================= - -// Returns true iff this call won the first-writer CAS for sched_error_code — the -// caller may then write companion fields (e.g. the stall detail) knowing they -// describe the same observation that owns the latched code. -static bool latch_scheduler_error(PTO2SharedMemoryHeader *header, int32_t thread_idx, int32_t error_code) { - if (header == nullptr || error_code == PTO2_ERROR_NONE) { - return false; - } - // The first error code/thread pair wins; the bitmap cumulatively records all reporting threads. + +static void latch_scheduler_error(PTO2SharedMemoryHeader *header, int32_t thread_idx, int32_t error_code) { + if (header == nullptr || error_code == PTO2_ERROR_NONE) return; int32_t expected = PTO2_ERROR_NONE; - bool won = header->sched_error_code.compare_exchange_strong(expected, error_code, std::memory_order_acq_rel); - if (won) { + if (header->sched_error_code.compare_exchange_strong(expected, error_code, std::memory_order_acq_rel)) header->sched_error_thread.store(thread_idx, std::memory_order_release); - } - if (thread_idx >= 0 && thread_idx < 32) { + if (thread_idx >= 0 && thread_idx < 32) header->sched_error_bitmap.fetch_or(1U << static_cast(thread_idx), std::memory_order_acq_rel); - } - return won; -} - -LoopAction SchedulerContext::handle_orchestrator_exit( - int32_t thread_idx, PTO2SharedMemoryHeader *header, Runtime *runtime, int32_t &task_count -) { - if (completed_.load(std::memory_order_acquire)) { - return LoopAction::BREAK_LOOP; - } - int32_t orch_err = header->orch_error_code.load(std::memory_order_acquire); - if (orch_err != PTO2_ERROR_NONE) { - LOG_ERROR( - "Thread %d: Fatal error (code=%d), sending EXIT_SIGNAL to all cores. " - "completed_tasks=%d, total_tasks=%d", - thread_idx, orch_err, completed_tasks_.load(std::memory_order_relaxed), total_tasks_ - ); - if (!completed_.exchange(true, std::memory_order_acq_rel)) { - emergency_shutdown(runtime); - } - return LoopAction::BREAK_LOOP; - } - int32_t sched_err = header->sched_error_code.load(std::memory_order_acquire); - if (sched_err != PTO2_ERROR_NONE) { - LOG_ERROR("Thread %d: Scheduler fatal error detected (code=%d)", thread_idx, sched_err); - if (!completed_.exchange(true, std::memory_order_acq_rel)) { - emergency_shutdown(runtime); - } - return LoopAction::BREAK_LOOP; - } - - bool orch_done = orchestrator_done_.load(std::memory_order_acquire); - if (!orch_done) return LoopAction::NONE; - - task_count = total_tasks_; - if (task_count > 0 && completed_tasks_.load(std::memory_order_relaxed) >= task_count) { - completed_.store(true, std::memory_order_release); - LOG_INFO_V0( - "Thread %d: PTO2 completed tasks %d/%d", thread_idx, completed_tasks_.load(std::memory_order_relaxed), - task_count - ); - return LoopAction::BREAK_LOOP; - } - return LoopAction::NONE; -} - -LoopAction -SchedulerContext::check_idle_fatal_error(int32_t thread_idx, PTO2SharedMemoryHeader *header, Runtime *runtime) { - if (completed_.load(std::memory_order_acquire)) { - return LoopAction::BREAK_LOOP; - } - int32_t orch_err = header->orch_error_code.load(std::memory_order_acquire); - if (orch_err != PTO2_ERROR_NONE) { - LOG_ERROR("Thread %d: Fatal error detected (code=%d), sending EXIT_SIGNAL to all cores", thread_idx, orch_err); - if (!completed_.exchange(true, std::memory_order_acq_rel)) { - emergency_shutdown(runtime); - } - return LoopAction::BREAK_LOOP; - } - int32_t sched_err = header->sched_error_code.load(std::memory_order_acquire); - if (sched_err != PTO2_ERROR_NONE) { - LOG_ERROR("Thread %d: Scheduler fatal error detected (code=%d)", thread_idx, sched_err); - if (!completed_.exchange(true, std::memory_order_acq_rel)) { - emergency_shutdown(runtime); - } - return LoopAction::BREAK_LOOP; - } - return LoopAction::NONE; } -// ============================================================================= -// Stall diagnostic log format. -// -// Every line is self-contained — when scheduler threads emit concurrently and -// device_log interleaves their output, each line still carries enough context -// to identify which thread / iteration / object it belongs to. -// -// Prefix on every line: -// [STALL thread=N idle_iterations=K] CATEGORY ... -// -// All scheduler threads spinning at the same idle rate hit STALL_LOG_INTERVAL -// together, so lines with the same idle_iterations belong to one diagnostic -// round; grep "idle_iterations=N" groups one round's output. -// -// Categories (and which thread emits them): -// SUMMARY — completed / total counts and scan totals (thread 0 only) -// TASK — one per non-completed task scanned from shared rings (thread 0 only) -// - state=RUNNING: includes running_on=[...] cross-ref -// - state=READY: fanin satisfied but no idle core yet -// - state=WAIT: includes missing_deps=N -// CLUSTER — one per cluster owned by this thread (every thread) -// - busy slot shows kernel + task_id + cond_reg_state; -// ANOMALY suffix when COND register is fin while software -// still has the slot marked busy. -// -// Reader workflow: -// 1. grep SUMMARY -> overall completion status -// 2. grep "idle_iterations=N TASK" -> stuck RUNNING task and which -// core/thread it is on -// 3. grep "idle_iterations=N CLUSTER.*task=" -> cross-check via the -// cluster line (or just -// read running_on in step 2) -// ============================================================================= - -namespace { - -// Format a core's idle/busy state into a fixed buffer. Used inside CLUSTER lines. -// Layout (idle): coreN(idle) -// Layout (busy): coreN(busy kernel=K task=T cond_reg_state=ack) -// Layout (anomaly): coreN(busy kernel=K task=T cond_reg_state=fin ANOMALY) -// -// Healthy busy: COND register reports ack (AICore still executing). fin means -// AICore wrote completion but AICPU hasn't recycled the running slot yet — -// either a completion-poll bug or the diagnostic raced the recycle. -void format_core_status( +static void format_core_status( char *buf, size_t buf_size, int32_t core_id, bool idle, const CoreExecState *core_state, uint64_t reg_addr_for_cond ) { if (idle) { @@ -179,887 +46,18 @@ void format_core_status( uint64_t cond_reg = read_reg(reg_addr_for_cond, RegId::COND); int32_t hw_state = EXTRACT_TASK_STATE(cond_reg); const char *cond_reg_state_str = (hw_state == TASK_ACK_STATE) ? "ack" : "fin"; - if (hw_state == TASK_ACK_STATE) { + if (hw_state == TASK_ACK_STATE) snprintf( buf, buf_size, "core%d(busy kernel=%d task=%" PRId64 " cond_reg_state=%s)", core_id, kernel, task_id_raw, - cond_reg_state_str - ); - } else { - snprintf( - buf, buf_size, - "core%d(busy kernel=%d task=%" PRId64 - " cond_reg_state=%s ANOMALY cond_tok=%d running_tok=%d pending_tok=%d)", - core_id, kernel, task_id_raw, cond_reg_state_str, EXTRACT_TASK_ID(cond_reg), - core_state->running_reg_task_id, core_state->pending_reg_task_id - ); - } -} - -} // namespace - -int32_t SchedulerContext::find_core_owner_thread(int32_t core_id) const { - for (int32_t t = 0; t < aicpu_thread_num_; t++) { - const int32_t *ids = core_trackers_[t].core_ids(); - int32_t n = core_trackers_[t].core_num(); - for (int32_t i = 0; i < n; i++) { - if (ids[i] == core_id) return t; - } - } - return -1; -} - -bool SchedulerContext::self_owns_running_task(int32_t thread_idx) const { - const int32_t *cores = core_trackers_[thread_idx].core_ids(); - int32_t core_num = core_trackers_[thread_idx].core_num(); - for (int32_t i = 0; i < core_num; i++) { - if (core_exec_states_[cores[i]].running_slot_state != nullptr) { - return true; - } - } - return false; -} - -bool SchedulerContext::no_thread_owns_running_task() const { - for (int32_t t = 0; t < aicpu_thread_num_; t++) { - if (self_owns_running_task(t)) return false; - } - return true; -} - -void SchedulerContext::log_stall_diagnostics( - int32_t thread_idx, int32_t task_count, int32_t idle_iterations, int32_t last_progress_count -) { - CoreTracker &tracker = core_trackers_[thread_idx]; - - // T0 owns the shared-ring scan; printing it from other threads would - // produce identical TASK lines once per scheduler thread. - if (thread_idx == 0) { - int32_t cnt_ready = 0, cnt_waiting = 0, cnt_running = 0, submitted_in_ring = 0; - for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { - PTO2SharedMemoryRingHeader &ring = *sched_->ring_sched_states[r].ring; - int32_t ring_task_count = ring.fc.current_task_index.load(std::memory_order_relaxed); - submitted_in_ring += ring_task_count; - // Scan only live task_ids [last_task_alive, current_task_index): slots - // wrap (slot = task_id % window), so starting at 0 re-reads each live - // slot once per earlier task_id and inflates the scan_* counts. - int32_t ring_task_start = ring.fc.last_task_alive.load(std::memory_order_relaxed); - for (int32_t si = ring_task_start; 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; - 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]; - int64_t task_id = static_cast(slot_state.task->task_id.raw); - if (st >= PTO2_TASK_COMPLETED) continue; - // task_state has no intermediate ready/running value — it - // stays PENDING until the worker stores COMPLETED. Classify - // by the ground truth instead: a slot is RUNNING iff some - // core has it as running_slot_state. A task occupies at most - // 3 cores (one cluster), all under the same owner thread by - // construction of assign_cores_to_threads. - char running_on[192] = {0}; - int32_t owner = -1; - int32_t pos = 0; - bool is_running = false; - for (int32_t cid = 0; cid < cores_total_num_ && pos + 32 < (int32_t)sizeof(running_on); cid++) { - if (core_exec_states_[cid].running_slot_state != &slot_state) continue; - is_running = true; - if (owner < 0) owner = find_core_owner_thread(cid); - const char *sname = subslot_name(core_exec_states_[cid].running_subslot); - int32_t written = snprintf( - running_on + pos, sizeof(running_on) - pos, "%score=%d(%s)", pos == 0 ? "" : " ", cid, sname - ); - if (written > 0) pos += written; - } - - if (is_running) { - cnt_running++; - 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] " - "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 - ); - continue; - } - if (rc >= fi) { - cnt_ready++; - 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]", - thread_idx, idle_iterations, r, task_id, rc, fi, kid_aic, kid_aiv0, kid_aiv1 - ); - continue; - } - cnt_waiting++; - 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", - thread_idx, idle_iterations, r, task_id, rc, fi, kid_aic, kid_aiv0, kid_aiv1, fi - rc - ); - } - } - int32_t effective_total = task_count > 0 ? task_count : submitted_in_ring; - int32_t c = completed_tasks_.load(std::memory_order_relaxed); - LOG_INFO_V9( - "[STALL thread=%d idle_iterations=%d] SUMMARY completed=%d/%d last_progress_iteration=%d " - "scan_ready=%d scan_waiting=%d scan_running=%d", - thread_idx, idle_iterations, c, effective_total, last_progress_count, cnt_ready, cnt_waiting, cnt_running - ); - } - - // CLUSTER lines: one per cluster this thread owns. - // cluster_id = local_cluster_idx * active_sched_threads_ + thread_idx, matching the - // round-robin assignment in assign_cores_to_threads. - int32_t ast = active_sched_threads_ > 0 ? active_sched_threads_ : aicpu_thread_num_; - for (int32_t cli = 0; cli < tracker.get_cluster_count() && cli < STALL_DUMP_CORE_MAX; cli++) { - int32_t offset = cli * 3; - int32_t aic_id = tracker.get_aic_core_id(offset); - int32_t aiv0_id = tracker.get_aiv0_core_id(offset); - int32_t aiv1_id = tracker.get_aiv1_core_id(offset); - bool aic_idle = tracker.is_aic_core_idle(offset); - bool aiv0_idle = tracker.is_aiv0_core_idle(offset); - bool aiv1_idle = tracker.is_aiv1_core_idle(offset); - int32_t cluster_id = cli * ast + thread_idx; - char aic_buf[192], aiv0_buf[192], aiv1_buf[192]; - format_core_status( - aic_buf, sizeof(aic_buf), aic_id, aic_idle, &core_exec_states_[aic_id], core_exec_states_[aic_id].reg_addr - ); - format_core_status( - aiv0_buf, sizeof(aiv0_buf), aiv0_id, aiv0_idle, &core_exec_states_[aiv0_id], - core_exec_states_[aiv0_id].reg_addr - ); - format_core_status( - aiv1_buf, sizeof(aiv1_buf), aiv1_id, aiv1_idle, &core_exec_states_[aiv1_id], - core_exec_states_[aiv1_id].reg_addr - ); - LOG_INFO_V9( - "[STALL thread=%d idle_iterations=%d] CLUSTER cluster_id=%d aic=%s aiv0=%s aiv1=%s", thread_idx, - idle_iterations, cluster_id, aic_buf, aiv0_buf, aiv1_buf - ); - } -} - -void SchedulerContext::log_shutdown_stall_snapshot( - int32_t trigger_thread_idx, int32_t trigger_idle_iterations, int32_t trigger_last_progress_count -) { - LOG_WARN( - "[SHUTDOWN_SNAPSHOT trigger_thread=%d reason=scheduler_timeout idle_iterations=%d] " - "dumping all scheduler threads before emergency shutdown", - trigger_thread_idx, trigger_idle_iterations - ); - int32_t thread_count = active_sched_threads_ > 0 ? active_sched_threads_ : aicpu_thread_num_; - if (thread_count < 0 || thread_count > MAX_AICPU_THREADS) { - LOG_ERROR( - "[SHUTDOWN_SNAPSHOT trigger_thread=%d] invalid thread_count=%d, clamping to [0,%d]", trigger_thread_idx, - thread_count, MAX_AICPU_THREADS - ); - thread_count = thread_count < 0 ? 0 : MAX_AICPU_THREADS; - } - for (int32_t t = 0; t < thread_count; t++) { - log_stall_diagnostics(t, total_tasks_, trigger_idle_iterations, trigger_last_progress_count); - } -} - -SchedulerContext::StallClassification SchedulerContext::classify_stall_reason() const { - StallClassification cls{}; - cls.stuck_task_id = -1; - cls.stuck_core = -1; - int32_t cnt_running = 0, cnt_ready = 0, cnt_waiting = 0; - for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { - PTO2SharedMemoryRingHeader &ring = *sched_->ring_sched_states[r].ring; - int32_t ring_task_count = ring.fc.current_task_index.load(std::memory_order_relaxed); - // Active task_ids live in [last_task_alive, current_task_index); slots wrap - // (slot = task_id % window), so scanning from 0 re-reads each live slot once - // per earlier task_id that mapped to it -- inflating the counts to O(history). - // Start at the tail so each live slot is visited exactly once (O(window)). - int32_t ring_task_start = ring.fc.last_task_alive.load(std::memory_order_relaxed); - for (int32_t si = ring_task_start; 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); - if (st >= PTO2_TASK_COMPLETED) continue; - // Same ground truth as log_stall_diagnostics: task_state stays PENDING - // until COMPLETED, so RUNNING is read from core ownership, not the slot. - int32_t run_core = -1; - for (int32_t cid = 0; cid < cores_total_num_; cid++) { - if (core_exec_states_[cid].running_slot_state == &slot_state) { - run_core = cid; - break; - } - } - if (run_core >= 0) { - if (cnt_running == 0) { - // Snapshot the non-atomic task pointer once: it can be null on a - // torn slot, and a concurrent writer may flip it mid-read. - PTO2TaskDescriptor *task_ptr = slot_state.task; - cls.stuck_task_id = (task_ptr != nullptr) ? static_cast(task_ptr->task_id.raw) : -1; - cls.stuck_core = run_core; - } - cnt_running++; - continue; - } - int32_t rc = slot_state.fanin_refcount.load(std::memory_order_relaxed); - int32_t fi = slot_state.fanin_count; - if (rc >= fi) { - cnt_ready++; - continue; - } - cnt_waiting++; - } - } - cls.cnt_running = cnt_running; - cls.cnt_ready = cnt_ready; - cls.cnt_waiting = cnt_waiting; - cls.completed = completed_tasks_.load(std::memory_order_relaxed); - cls.total = total_tasks_; - cls.orch_done = orchestrator_done_ ? 1 : 0; - cls.detail = classify_stall_detail(cnt_running, cnt_ready, cnt_waiting, cls.orch_done); - return cls; -} - -int32_t SchedulerContext::handle_timeout_exit( - int32_t thread_idx, PTO2SharedMemoryHeader *header, Runtime *runtime, int32_t idle_iterations, - int32_t last_progress_count -#if SIMPLER_DFX - , - uint64_t sched_start_ts -#endif -) { - StallClassification cls = classify_stall_reason(); - LOG_ERROR( - "[STALL thread=%d idle_iterations=%d] TIMEOUT_EXIT after_idle_iterations=%d sub_class=%s " - "completed=%d/%d running=%d ready=%d waiting=%d orch_done=%d stuck_task_id=%" PRId64 " stuck_core=%d", - thread_idx, idle_iterations, idle_iterations, stall_detail_name(cls.detail), cls.completed, cls.total, - cls.cnt_running, cls.cnt_ready, cls.cnt_waiting, cls.orch_done, cls.stuck_task_id, cls.stuck_core - ); - // Only the thread that wins the code-100 latch publishes the detail/locators, - // keeping the host-visible sub-class consistent with the latched code. - if (latch_scheduler_error(header, thread_idx, PTO2_ERROR_SCHEDULER_TIMEOUT) && header != nullptr) { - header->sched_stall_completed.store(cls.completed, std::memory_order_relaxed); - header->sched_stall_total.store(cls.total, std::memory_order_relaxed); - header->sched_stall_cnt_running.store(cls.cnt_running, std::memory_order_relaxed); - header->sched_stall_cnt_ready.store(cls.cnt_ready, std::memory_order_relaxed); - header->sched_stall_cnt_waiting.store(cls.cnt_waiting, std::memory_order_relaxed); - header->sched_stall_orch_done.store(cls.orch_done, std::memory_order_relaxed); - header->sched_stall_task_id.store(cls.stuck_task_id, std::memory_order_relaxed); - header->sched_stall_core.store(cls.stuck_core, std::memory_order_relaxed); - // detail published last (release) so a host reading a non-NONE detail - // sees the locators above already settled. - header->sched_stall_detail.store(cls.detail, std::memory_order_release); - } - if (!completed_.exchange(true, std::memory_order_acq_rel)) { - log_shutdown_stall_snapshot(thread_idx, idle_iterations, last_progress_count); -#if SIMPLER_DFX - // Capture the in-flight kernels' partial output before signalling the - // cores to exit, so the dump reflects the live stuck state. - if (is_dump_args_enabled()) { - dump_running_task_outputs( - thread_idx, cores_total_num_, - [this](int32_t cid) { - return core_exec_states_[cid].running_slot_state; - }, - [](ActiveMask active_mask, int raw_subtask_id) { - return active_mask.subtask_active(static_cast(raw_subtask_id)); - }, - [this](int32_t func_id) { - return get_function_bin_addr(func_id); - } - ); - } -#endif - emergency_shutdown(runtime); - } -#if SIMPLER_DFX - uint64_t sched_timeout_ts = get_sys_cnt_aicpu(); - aicpu_phase_set_window( - AicpuPhase::SchedWindow, static_cast(sched_start_ts), static_cast(sched_timeout_ts) - ); -#if SIMPLER_SCHED_PROFILING - LOG_INFO_V9( - "Thread %d: sched_start=%" PRIu64 " sched_end(timeout)=%" PRIu64 " sched_cost=%.3fus", thread_idx, - static_cast(sched_start_ts), static_cast(sched_timeout_ts), - cycles_to_us(sched_timeout_ts - sched_start_ts) - ); -#endif -#endif - return -PTO2_ERROR_SCHEDULER_TIMEOUT; -} - -#if SIMPLER_DFX -void SchedulerContext::log_l2_swimlane_summary(int32_t thread_idx, [[maybe_unused]] int32_t cur_thread_completed) { - auto &l2_swimlane = sched_l2_swimlane_[thread_idx]; - uint64_t sched_end_ts = get_sys_cnt_aicpu(); - // Ride the sched window home to the host phase buffer (the host reduces - // across sched threads → the `Sched` [STRACE] marker). The verbose - // per-thread device-log line below is now opt-in deep-dive. - aicpu_phase_set_window( - AicpuPhase::SchedWindow, static_cast(l2_swimlane.sched_start_ts), static_cast(sched_end_ts) - ); -#if SIMPLER_SCHED_PROFILING - LOG_INFO_V9( - "Thread %d: sched_start=%" PRIu64 " sched_end=%" PRIu64 " sched_cost=%.3fus", thread_idx, - static_cast(l2_swimlane.sched_start_ts), static_cast(sched_end_ts), - cycles_to_us(sched_end_ts - l2_swimlane.sched_start_ts) - ); - - uint64_t sched_total = l2_swimlane.sched_complete_cycle + l2_swimlane.sched_async_cycle + - l2_swimlane.sched_dispatch_cycle + l2_swimlane.sched_idle_cycle; - if (sched_total == 0) sched_total = 1; - - { - PTO2SchedProfilingData sp = scheduler_get_profiling(thread_idx); - uint64_t otc_total = sp.lock_cycle + sp.fanout_cycle + sp.fanin_cycle + sp.self_consumed_cycle; - uint64_t complete_poll = - (l2_swimlane.sched_complete_cycle > otc_total + l2_swimlane.sched_complete_perf_cycle) ? - (l2_swimlane.sched_complete_cycle - otc_total - l2_swimlane.sched_complete_perf_cycle) : - 0; - uint64_t dispatch_poll = (l2_swimlane.sched_dispatch_cycle > - l2_swimlane.sched_dispatch_pop_cycle + l2_swimlane.sched_dispatch_setup_cycle) ? - (l2_swimlane.sched_dispatch_cycle - l2_swimlane.sched_dispatch_pop_cycle - - l2_swimlane.sched_dispatch_setup_cycle) : - 0; - - LOG_INFO_V9( - "Thread %d: === Scheduler Phase Breakdown: total=%.3fus, %d tasks ===", thread_idx, - cycles_to_us(sched_total), cur_thread_completed - ); - - // fanout / fanin per-thread aggregates live in - // sched_overhead_analysis.compute_dag_stats_from_deps (deps.json edges - // × core_to_thread). - LOG_INFO_V9( - "Thread %d: complete : %.3fus (%.1f%%)", thread_idx, cycles_to_us(l2_swimlane.sched_complete_cycle), - l2_swimlane.sched_complete_cycle * 100.0 / sched_total - ); - - uint64_t c_parent = l2_swimlane.sched_complete_cycle > 0 ? l2_swimlane.sched_complete_cycle : 1; - uint64_t complete_miss_count = (l2_swimlane.complete_probe_count > l2_swimlane.complete_hit_count) ? - (l2_swimlane.complete_probe_count - l2_swimlane.complete_hit_count) : - 0; - double complete_hit_rate = l2_swimlane.complete_probe_count > 0 ? - l2_swimlane.complete_hit_count * 100.0 / l2_swimlane.complete_probe_count : - 0.0; - LOG_INFO_V9( - "Thread %d: poll : %.3fus (%.1f%%) hit=%" PRIu64 ", miss=%" PRIu64 ", hit_rate=%.1f%%", - thread_idx, cycles_to_us(complete_poll), complete_poll * 100.0 / c_parent, - static_cast(l2_swimlane.complete_hit_count), static_cast(complete_miss_count), - complete_hit_rate - ); - LOG_INFO_V9( - "Thread %d: otc_lock : %.3fus (%.1f%%) work=%.3fus wait=%.3fus atomics=%" PRIu64 "", thread_idx, - cycles_to_us(sp.lock_cycle), sp.lock_cycle * 100.0 / c_parent, - cycles_to_us(sp.lock_cycle - sp.lock_wait_cycle), cycles_to_us(sp.lock_wait_cycle), - static_cast(sp.lock_atomic_count) - ); - LOG_INFO_V9( - "Thread %d: otc_fanout : %.3fus (%.1f%%) work=%.3fus wait=%.3fus atomics=%" PRIu64 "", thread_idx, - cycles_to_us(sp.fanout_cycle), sp.fanout_cycle * 100.0 / c_parent, - cycles_to_us(sp.fanout_cycle - sp.push_wait_cycle), cycles_to_us(sp.push_wait_cycle), - static_cast(sp.fanout_atomic_count) - ); - LOG_INFO_V9( - "Thread %d: otc_fanin : %.3fus (%.1f%%) atomics=%" PRIu64 "", thread_idx, - cycles_to_us(sp.fanin_cycle), sp.fanin_cycle * 100.0 / c_parent, - static_cast(sp.fanin_atomic_count) - ); - LOG_INFO_V9( - "Thread %d: otc_self : %.3fus (%.1f%%) atomics=%" PRIu64 "", thread_idx, - cycles_to_us(sp.self_consumed_cycle), sp.self_consumed_cycle * 100.0 / c_parent, - static_cast(sp.self_atomic_count) - ); - LOG_INFO_V9( - "Thread %d: perf : %.3fus (%.1f%%)", thread_idx, - cycles_to_us(l2_swimlane.sched_complete_perf_cycle), - l2_swimlane.sched_complete_perf_cycle * 100.0 / c_parent - ); - - LOG_INFO_V9( - "Thread %d: async_poll : %.3fus (%.1f%%)", thread_idx, cycles_to_us(l2_swimlane.sched_async_cycle), - l2_swimlane.sched_async_cycle * 100.0 / sched_total - ); - - LOG_INFO_V9( - "Thread %d: dispatch : %.3fus (%.1f%%)", thread_idx, cycles_to_us(l2_swimlane.sched_dispatch_cycle), - l2_swimlane.sched_dispatch_cycle * 100.0 / sched_total - ); - - uint64_t d_parent = l2_swimlane.sched_dispatch_cycle > 0 ? l2_swimlane.sched_dispatch_cycle : 1; - LOG_INFO_V9( - "Thread %d: poll : %.3fus (%.1f%%)", thread_idx, cycles_to_us(dispatch_poll), - dispatch_poll * 100.0 / d_parent - ); - LOG_INFO_V9( - "Thread %d: pop : %.3fus (%.1f%%) work=%.3fus wait=%.3fus atomics=%" PRIu64 "", thread_idx, - cycles_to_us(l2_swimlane.sched_dispatch_pop_cycle), l2_swimlane.sched_dispatch_pop_cycle * 100.0 / d_parent, - cycles_to_us(l2_swimlane.sched_dispatch_pop_cycle - sp.pop_wait_cycle), cycles_to_us(sp.pop_wait_cycle), - static_cast(sp.pop_atomic_count) - ); - LOG_INFO_V9( - "Thread %d: setup : %.3fus (%.1f%%)", thread_idx, - cycles_to_us(l2_swimlane.sched_dispatch_setup_cycle), - l2_swimlane.sched_dispatch_setup_cycle * 100.0 / d_parent - ); - - LOG_INFO_V9( - "Thread %d: idle : %.3fus (%.1f%%)", thread_idx, cycles_to_us(l2_swimlane.sched_idle_cycle), - l2_swimlane.sched_idle_cycle * 100.0 / sched_total - ); - - if (cur_thread_completed > 0) { - LOG_INFO_V9( - "Thread %d: avg/complete : %.3fus", thread_idx, - cycles_to_us(l2_swimlane.sched_complete_cycle) / cur_thread_completed - ); - } - } - LOG_INFO_V9( - "Thread %d: Scheduler summary: total_time=%.3fus, loops=%" PRIu64 ", tasks_scheduled=%d", thread_idx, - cycles_to_us(sched_total), static_cast(l2_swimlane.sched_loop_count), cur_thread_completed - ); -#endif -} -#endif - -// ============================================================================= -// Shutdown: deinit AICore regs for this thread's cores (and PMU finalize if enabled). -// Orchestrator threads have core_trackers_[thread_idx].core_num() == 0 -> no-op. -// platform_deinit_aicore_regs is idempotent; safe to call after early completion. -// ============================================================================= -int32_t SchedulerContext::shutdown(int32_t thread_idx) { - const int32_t *cores = core_trackers_[thread_idx].core_ids(); - int32_t core_num = core_trackers_[thread_idx].core_num(); - if (core_num == 0) return 0; - -#if SIMPLER_DFX - if (is_pmu_enabled()) { - pmu_aicpu_finalize(cores, core_num); - } -#endif - - LOG_INFO_V0("Thread %d: Shutting down %d cores", thread_idx, core_num); - int32_t rc = 0; - for (int32_t i = 0; i < core_num; i++) { - int32_t core_id = cores[i]; - uint64_t reg_addr = core_exec_states_[core_id].reg_addr; - if (reg_addr != 0) { - // Timeout means AICore is unresponsive. Log and continue deiniting remaining cores. - if (platform_deinit_aicore_regs(reg_addr) != 0) { - LOG_ERROR("Thread %d: Core %d deinit timed out", thread_idx, core_id); - rc = -1; - } - } else { - LOG_ERROR("Thread %d: Core %d has invalid register address", thread_idx, core_id); - } - } - LOG_INFO_V0("Thread %d: Shutdown complete", thread_idx); - return rc; -} - -// ============================================================================= -// Handshake a contiguous slice of AICore workers. Runs on every AICPU thread in -// parallel (partitioned by tidx/nthreads); the leader's pre_handshake_init has -// already zeroed state, set cores_total_num_, and reset the counts/flag. The -// per-core work here — releasing the core, then opening its register window over -// serial MMIO — is what dominates preamble, so splitting the slice across -// threads is the whole point. Within a slice we still sweep (poll every -// outstanding core per pass, service whichever reported) so one slow core's -// wakeup overlaps its neighbours' instead of blocking them. Worker-id lists are -// built serially in post_handshake_init (core-index order) once every slice has -// landed, so the shared aic_count_/aiv_count_ are written by one thread only. -// ============================================================================= -void SchedulerContext::handshake_partition(Runtime *runtime, int32_t tidx, int32_t nthreads) { - Handshake *all_handshakes = reinterpret_cast(runtime->dev.workers); - const int32_t total = cores_total_num_; - const int32_t lo = static_cast((static_cast(tidx) * total) / nthreads); - const int32_t hi = static_cast((static_cast(tidx + 1) * total) / nthreads); - - // The AICore publishes {physical_core_id, core_type, aicore_done} on launch, - // gated by nothing. task is not published here: the AICore's aicore_done - // report flushes its whole handshake cache line, so a task stored before the - // report would be clobbered. task is written per core in the sweep below, - // after that core's aicore_done is observed and before its window opens (the - // point the AICore reads task). - - // Get platform physical cores count for validation - uint32_t max_physical_cores_count = platform_get_physical_cores_count(); - - // Step 2: collect responses from this slice. Each core reports - // {physical_core_id, core_type, aicore_done} in one write, then waits — by - // polling its own DATA_MAIN_BASE SPR — for us to open its register window. - // We sweep the slice: poll every outstanding core per pass and service - // whichever have reported, rather than blocking on core i before looking at - // core i+1, so per-core wakeups overlap (≈ max, not Σ). aicore_done is a GM - // read (not the nGnRE MMIO reg window), so sweeping is not forced serial the - // way RegId::COND polling is. - // - // Servicing a core = validate its physical_core_id, then open its register - // window (platform_init_aicore_regs: FAST_PATH + DATA_MAIN_BASE=IDLE). That - // IDLE write is *also* the signal the core polls for to leave its - // post-report wait — so opening the window IS the acknowledgement. There is - // no separate aicpu_regs_ready ack and no second round-trip. AIC/AIV - // classification is deferred to post_handshake_init (serial) so aic_count_/ - // aiv_count_ are never incremented from more than one thread. - uint64_t *regs = reinterpret_cast(regs_); - bool core_serviced[RUNTIME_MAX_WORKER] = {false}; - - // Every core publishes aicore_done on launch, so the whole slice is already - // reported when the AICPU sweeps it. The reported cores are collected first, - // then serviced in batched phases (publish tasks, open windows, store - // CoreExecStates); each phase issues its stores without interleaving another - // phase's, so posted MMIO STRs and write-through GM stores do not serialize. - struct ReadyCore { - int32_t i; - uint32_t pcid; - uint64_t reg_addr; - CoreType core_type; - }; - ReadyCore ready[RUNTIME_MAX_WORKER]; - int32_t n_ready = 0; - - // Phase 1: collect every reported core in this slice and prefetch its - // CoreExecState line for write, so the Phase 4 struct store hits a warm line. - for (int32_t remaining = hi - lo; remaining > 0;) { - for (int32_t i = lo; i < hi; i++) { - if (core_serviced[i]) continue; - Handshake *hank = &all_handshakes[i]; - if (hank->aicore_done == 0) { - SPIN_WAIT_HINT(); - continue; - } - uint32_t physical_core_id = hank->physical_core_id; - if (physical_core_id >= max_physical_cores_count) { - LOG_ERROR( - "Core %d reported invalid physical_core_id=%u (platform max=%u)", i, physical_core_id, - max_physical_cores_count - ); - handshake_failed_.store(true, std::memory_order_release); - core_serviced[i] = true; - remaining--; - continue; - } - __builtin_prefetch(&core_exec_states_[i], 1, 3); - ready[n_ready++] = {i, physical_core_id, regs[physical_core_id], hank->core_type}; - core_serviced[i] = true; - remaining--; - } - } - - // Phase 2: publish every task pointer, then ONE barrier. The core reads task - // only after its window opens (Phase 3); a single barrier orders all task - // stores before any window STR. Writing task now (after the report) also - // keeps the core's CACHELINE_OUT report flush from clobbering it. - for (int32_t r = 0; r < n_ready; r++) { - all_handshakes[ready[r].i].task = reinterpret_cast(&payload_per_core_[ready[r].i][0]); - } - OUT_OF_ORDER_STORE_BARRIER(); - - // Phase 3: open every window. platform_init_aicore_regs' STRs are posted - // Device-nGnRE writes, issued back-to-back with no interleaved GM stores. - for (int32_t r = 0; r < n_ready; r++) { - platform_init_aicore_regs(ready[r].reg_addr); - } - - // Phase 4: publish each CoreExecState with a single (prefetched) struct store. - // core_exec_states_ is AICPU-private (the scheduler reads it, never the core), - // so it may be written after the windows open. - for (int32_t r = 0; r < n_ready; r++) { - int32_t i = ready[r].i; - CoreExecState st{}; - st.reg_addr = ready[r].reg_addr; - st.cond_ptr = get_reg_ptr(ready[r].reg_addr, RegId::COND); - st.running_reg_task_id = AICPU_TASK_INVALID; - st.pending_reg_task_id = AICPU_TASK_INVALID; -#if !SIMPLER_DFX - st.worker_id = i; - st.physical_core_id = ready[r].pcid; - st.core_type = ready[r].core_type; -#endif - core_exec_states_[i] = st; - core_type_compact_[i] = static_cast(ready[r].core_type); -#if SIMPLER_DFX - physical_core_ids_[i] = ready[r].pcid; -#endif - } - OUT_OF_ORDER_STORE_BARRIER(); -} - -// Handshake exactly the cores this scheduler thread will later manage. Blocked -// core layout ([0,N/3) AIC, [N/3,N) AIV) makes ownership predictable before -// handshake: cluster ci = {ci, N/3+2ci, N/3+2ci+1}, assigned to thread -// ci % active_threads. Same protocol as handshake_partition, but over the owned -// set instead of a contiguous slice. -void SchedulerContext::handshake_owned_clusters(Runtime *runtime, int32_t tidx, int32_t active_threads) { - Handshake *all_handshakes = reinterpret_cast(runtime->dev.workers); - const int32_t aic_n = cores_total_num_ / 3; - - int32_t owned[RUNTIME_MAX_WORKER]; - int32_t own_n = 0; - for (int32_t ci = tidx; ci < aic_n; ci += active_threads) { - owned[own_n++] = ci; // AIC - owned[own_n++] = aic_n + 2 * ci; // AIV0 - owned[own_n++] = aic_n + 2 * ci + 1; // AIV1 - } - - uint32_t max_physical_cores_count = platform_get_physical_cores_count(); - uint64_t *regs = reinterpret_cast(regs_); - bool core_serviced[RUNTIME_MAX_WORKER] = {false}; - - // Batched 4-phase handshake (mirrors handshake_partition over the owned set): - // collect reports, then publish tasks / open windows / store CoreExecStates in - // separate passes so posted MMIO STRs and GM stores don't serialize, and only - // two barriers fire for the whole owned set (not one per core). - struct ReadyCore { - int32_t i; - uint32_t pcid; - uint64_t reg_addr; - CoreType core_type; - }; - ReadyCore ready[RUNTIME_MAX_WORKER]; - int32_t n_ready = 0; - - // Phase 1: collect every reported owned core, prefetch its CoreExecState line. - for (int32_t remaining = own_n; remaining > 0;) { - for (int32_t k = 0; k < own_n; k++) { - int32_t i = owned[k]; - if (core_serviced[i]) continue; - Handshake *hank = &all_handshakes[i]; - if (hank->aicore_done == 0) { - SPIN_WAIT_HINT(); - continue; - } - uint32_t physical_core_id = hank->physical_core_id; - if (physical_core_id >= max_physical_cores_count) { - LOG_ERROR( - "Core %d reported invalid physical_core_id=%u (platform max=%u)", i, physical_core_id, - max_physical_cores_count - ); - handshake_failed_.store(true, std::memory_order_release); - core_serviced[i] = true; - remaining--; - continue; - } - __builtin_prefetch(&core_exec_states_[i], 1, 3); - ready[n_ready++] = {i, physical_core_id, regs[physical_core_id], hank->core_type}; - core_serviced[i] = true; - remaining--; - } - } - - // Phase 2: publish every task pointer, then ONE barrier. The core reads task - // only after its window opens (Phase 3), so a single barrier orders all task - // stores before any window STR. - for (int32_t r = 0; r < n_ready; r++) { - all_handshakes[ready[r].i].task = reinterpret_cast(&payload_per_core_[ready[r].i][0]); - } - OUT_OF_ORDER_STORE_BARRIER(); - - // Phase 3: open every window (the IDLE write is also the core's ack). - for (int32_t r = 0; r < n_ready; r++) { - platform_init_aicore_regs(ready[r].reg_addr); - } - - // Phase 4: publish each CoreExecState (AICPU-private, may follow the windows). - // reg_task_id fields start INVALID (pre_handshake_init memset zeroed them, and - // 0 is a valid task id); core_type_compact_ is filled for parity. - for (int32_t r = 0; r < n_ready; r++) { - int32_t i = ready[r].i; - CoreExecState st{}; - st.reg_addr = ready[r].reg_addr; - st.cond_ptr = get_reg_ptr(ready[r].reg_addr, RegId::COND); - st.running_reg_task_id = AICPU_TASK_INVALID; - st.pending_reg_task_id = AICPU_TASK_INVALID; -#if !SIMPLER_DFX - st.worker_id = i; - st.physical_core_id = ready[r].pcid; - st.core_type = ready[r].core_type; -#endif - core_exec_states_[i] = st; - core_type_compact_[i] = static_cast(ready[r].core_type); -#if SIMPLER_DFX - physical_core_ids_[i] = ready[r].pcid; -#endif - } - OUT_OF_ORDER_STORE_BARRIER(); -} - -// ============================================================================= -// Per-thread self-assignment (barrier-free init). Thread tidx owns the clusters -// ci with ci % active_sched_threads_ == tidx (same round-robin as -// assign_cores_to_threads), and the blocked layout gives their worker ids -// directly, so a thread populates its own CoreTracker + per-core payload state -// right after handshaking its own clusters, with no all-thread barrier. -// ============================================================================= -void SchedulerContext::assign_own_clusters(int32_t tidx) { - const int32_t aic_n = cores_total_num_ / 3; - const int32_t active = active_sched_threads_; - - CoreTracker &tracker = core_trackers_[tidx]; - int32_t own_n = 0; - for (int32_t ci = tidx; ci < aic_n; ci += active) - own_n++; - tracker.init(own_n); - - int32_t local = 0; - for (int32_t ci = tidx; ci < aic_n; ci += active) { - tracker.set_cluster(local++, ci, aic_n + 2 * ci, aic_n + 2 * ci + 1); - } - - // Per-cluster GlobalContext sub_block_id (mirrors post_handshake_init) for - // this thread's owned cores only — a thread only ever dispatches to its own. - for (int32_t c = 0; c < tracker.get_cluster_count(); c++) { - int32_t cluster_offset = c * 3; - int32_t aiv0_id = tracker.get_core_id_by_offset(tracker.get_aiv0_core_offset(cluster_offset)); - int32_t aiv1_id = tracker.get_core_id_by_offset(tracker.get_aiv1_core_offset(cluster_offset)); - payload_per_core_[aiv0_id][0].global_context.sub_block_id = 0; - payload_per_core_[aiv0_id][1].global_context.sub_block_id = 0; - payload_per_core_[aiv1_id][0].global_context.sub_block_id = 1; - payload_per_core_[aiv1_id][1].global_context.sub_block_id = 1; - } - - // Per-dispatch AsyncCtx constant prefill + one-time slab clear for owned cores - // (mirrors post_handshake_init's all-core loop, restricted to this thread's). - for (int32_t c = 0; c < tracker.get_cluster_count(); c++) { - for (int32_t sub = 0; sub < 3; sub++) { - int32_t core_id = tracker.get_core_id_by_offset(c * 3 + sub); - for (int32_t buf = 0; buf < 2; buf++) { - PTO2DispatchPayload &dp = payload_per_core_[core_id][buf]; - AsyncCtx &ac = dp.local_context.async_ctx; - volatile DeferredCompletionSlab *slab = &deferred_slab_per_core_[core_id][buf]; - ac.completion_count = &slab->count; - ac.completion_error_code = &slab->error_code; - ac.completion_entries = &slab->entries[0]; - ac.completion_capacity = MAX_COMPLETIONS_PER_TASK; - slab->count = 0; - slab->error_code = PTO2_ERROR_NONE; - dp.args[PAYLOAD_LOCAL_CONTEXT_INDEX] = reinterpret_cast(&dp.local_context); - dp.args[PAYLOAD_GLOBAL_CONTEXT_INDEX] = reinterpret_cast(&dp.global_context); - } - } - } -} - -// Abort the run on a handshake failure discovered without the all-thread barrier -// (non-DFX path): latch completion so every scheduler thread exits its dispatch -// loop, and broadcast exit to whatever cores did come up. Idempotent. -void SchedulerContext::abort_and_shutdown(Runtime *runtime) { - if (!completed_.exchange(true, std::memory_order_acq_rel)) { - emergency_shutdown(runtime); - } -} - -// Profiling-subsystem init (leader-only). pmu_aicpu_init needs every core's -// physical_core_id, so the barrier-free init path calls this behind an -// all-thread barrier compiled only into DFX builds. No-op otherwise. -void SchedulerContext::post_handshake_profiling_init() { -#if SIMPLER_DFX - if (is_dump_args_enabled()) { - dump_args_init(active_sched_threads_); - } - if (is_pmu_enabled()) { - pmu_aicpu_init(physical_core_ids_, cores_total_num_); - LOG_INFO_V0("PMU profiling started on %d cores", cores_total_num_); - } - if (is_dep_gen_enabled()) { - dep_gen_aicpu_init(); - } -#endif -} - -// ============================================================================= -// Assign discovered cores to scheduler threads (cluster-aligned round-robin). -// ============================================================================= -bool SchedulerContext::assign_cores_to_threads() { - // Cluster-aligned round-robin assignment: cluster ci -> sched thread ci % active_sched_threads_. - // Each cluster = 1 AIC + 2 adjacent AIV; the triple is always kept together. - active_sched_threads_ = (sched_thread_num_ > 0) ? sched_thread_num_ : aicpu_thread_num_; - int32_t cluster_count = aic_count_; - - // Max clusters any single sched thread can hold: ceil(cluster_count / active_sched_threads_). - int32_t max_clusters_per_thread = (cluster_count + active_sched_threads_ - 1) / active_sched_threads_; - int32_t thread_cores_num = max_clusters_per_thread * 3; - - if (thread_cores_num > CoreTracker::MAX_CORE_PER_THREAD) { - LOG_ERROR("Can't assign more then 64 cores in per scheduler"); - return false; - } - - LOG_INFO_V0( - "Assigning cores (round-robin): %d clusters across %d sched threads (%d AIC, %d AIV)", cluster_count, - active_sched_threads_, aic_count_, aiv_count_ - ); - - // running_reg_task_id / pending_reg_task_id for every serviced core are reset - // in handshake_partition's sweep. - - // Count clusters per thread first (round-robin may distribute unevenly) - int32_t clusters_per_thread[MAX_AICPU_THREADS] = {}; - for (int32_t ci = 0; ci < cluster_count; ci++) { - clusters_per_thread[ci % active_sched_threads_]++; - } - for (int32_t i = 0; i < active_sched_threads_; i++) { - core_trackers_[i].init(clusters_per_thread[i]); - } - - int32_t cluster_idx_per_thread[MAX_AICPU_THREADS] = {}; - - for (int32_t ci = 0; ci < cluster_count; ci++) { - int32_t t = ci % active_sched_threads_; - - int32_t aic_wid = aic_worker_ids_[ci]; - int32_t aiv0_wid = aiv_worker_ids_[2 * ci]; - int32_t aiv1_wid = aiv_worker_ids_[2 * ci + 1]; - - core_trackers_[t].set_cluster(cluster_idx_per_thread[t]++, aic_wid, aiv0_wid, aiv1_wid); - - LOG_INFO_V0("Thread %d: cluster %d (AIC=%d, AIV0=%d, AIV1=%d)", t, ci, aic_wid, aiv0_wid, aiv1_wid); - } - - for (int32_t t = 0; t < aicpu_thread_num_; t++) { - LOG_INFO_V0( - "Thread %d: total %d cores (%d clusters)", t, core_trackers_[t].core_num(), - core_trackers_[t].get_cluster_count() - ); - } - - LOG_INFO_V0( - "Config: threads=%d, cores=%d, cores_per_thread=%d", aicpu_thread_num_, cores_total_num_, thread_cores_num - ); - return true; -} - -// ============================================================================= -// Emergency shutdown: broadcast exit signal to every handshake'd core and -// deinit their AICore register blocks. Idempotent. -// ============================================================================= -void SchedulerContext::emergency_shutdown(Runtime *runtime) { - (void)runtime; // exit is now delivered via each core's register block, not GM - LOG_WARN("Emergency shutdown: sending exit signal to all initialized cores"); - int32_t timeout_count = 0; - for (int32_t i = 0; i < cores_total_num_; i++) { - // platform_deinit_aicore_regs writes DATA_MAIN_BASE=EXIT, which both - // releases a core still polling for its window to open and signals it to - // exit. Cores never opened (reg_addr==0) are reaped by the host device - // reset that follows a handshake failure. - if (core_exec_states_[i].reg_addr != 0) { - if (platform_deinit_aicore_regs(core_exec_states_[i].reg_addr) != 0) { - timeout_count++; - } - } - } - if (timeout_count > 0) { - LOG_ERROR("Emergency shutdown: %d cores did not acknowledge exit", timeout_count); - } - LOG_WARN("Emergency shutdown complete"); + cond_reg_state_str + ); + else + snprintf( + buf, buf_size, "core%d(busy kernel=%d task=%" PRId64 " cond_reg_state=%s ANOMALY)", core_id, kernel, + task_id_raw, cond_reg_state_str + ); } -// ============================================================================= -// Lifecycle: init / deinit -// ============================================================================= int32_t SchedulerContext::pre_handshake_init( Runtime *runtime, int32_t aicpu_thread_num, int32_t sched_thread_num, uint64_t regs_base ) { @@ -1073,86 +71,101 @@ int32_t SchedulerContext::pre_handshake_init( sched_thread_num_ = sched_thread_num; regs_ = regs_base; -#if SIMPLER_DFX - // l2_swimlane_aicpu_init promotes g_l2_swimlane_level from the shared-memory - // header — must be called BEFORE caching the level, otherwise the cached - // value would still be 0 (only the binary enable bit has been seeded by - // kernel.cpp at this point). Reset the cached level on disabled runs so a - // prior enabled launch's level can't leak into the phase-record gates in - // scheduler_dispatch. This runs on the leader before it publishes - // hs_setup_done_, so it happens-before every thread's handshake_partition - // (and therefore before any aicpu_ready=1 write). + // Initialize l2-swimlane buffers BEFORE any thread writes aicpu_ready in + // handshake_partition so the AICore-side rotation table slots are + // populated when AICore reads them post-handshake. AICore stashes + // &rotation_table[block_idx] at entry; the slot CONTENTS (the actual + // record buffer pointer it later dereferences) are written here. + // aicpu_ready=1 is AICore's signal to proceed past Phase 1 — once it has + // the green light, it expects the slot to be initialized. This runs on + // the leader before it publishes hs_setup_done_, so it happens-before + // every thread's handshake. See the contract comment in + // aicore/aicore_executor.cpp:105-110 and the parallel call in + // host_build_graph/aicpu/aicpu_executor.cpp:341. Without this call, + // --enable-l2-swimlane runs hit AICore-side memory corruption that + // surfaces as orch FLOW_CONTROL_DEADLOCK (paged_attention C1) or + // sched SCHEDULER_TIMEOUT (multi_round_paged_attention C1) depending + // on which AICore op first touches the uninitialized slot. if (is_l2_swimlane_enabled()) { l2_swimlane_aicpu_init(runtime->dev.worker_count); - l2_swimlane_level_ = get_l2_swimlane_level(); - if (l2_swimlane_level_ >= L2SwimlaneLevel::SCHED_PHASES) { - // Sched-phase pool count must match the dump_args_init thread count - // below. This block runs before assign_cores_to_threads, so the - // active_sched_threads_ member isn't set yet — recompute the same - // normalization locally: sched_thread_num_ <= 0 means "use all AICPU - // threads as scheduler threads" (see assign_cores_to_threads' - // active_sched_threads_). Without it, init_phase would prime zero - // sched pools and all sched_phase emits would silently drop. - const int sched_phase_threads = (sched_thread_num_ > 0) ? sched_thread_num_ : aicpu_thread_num_; - // Orchestration is always single-threaded, so orch-phase is one pool - // (ordinal 0) — see record_orch_phase. - const int orch_phase_threads = 1; - l2_swimlane_aicpu_init_phase(runtime->dev.worker_count, sched_phase_threads, orch_phase_threads); - } - } else { - l2_swimlane_level_ = L2SwimlaneLevel::DISABLED; } -#endif - // Core count is needed by every thread to compute its handshake slice. cores_total_num_ = runtime->dev.worker_count; - if (cores_total_num_ == 0 || cores_total_num_ > RUNTIME_MAX_WORKER) { - LOG_ERROR("Invalid cores_total_num %d (expected 1-%d)", cores_total_num_, RUNTIME_MAX_WORKER); - return -1; - } - // The blocked 1 AIC : 2 AIV layout requires an exact multiple of 3: cluster ci - // owns cores {ci, N/3+2ci, N/3+2ci+1}, so a non-zero remainder leaves the tail - // AIV cores [3*(N/3), N) in no cluster — unhandshaked, their windows never open, - // and the run hangs at the op-execute timeout. assign_cores_to_threads pairs - // aiv_worker_ids_[2*ci]/[2*ci+1] on the serial path too, so this holds for both. - if (cores_total_num_ % 3 != 0) { - LOG_ERROR("cores_total_num %d is not a multiple of 3 (blocked 1 AIC : 2 AIV layout)", cores_total_num_); - return -1; - } - // Blocked core layout ([0,N/3) AIC, [N/3,N) AIV) with a fixed 1:2 AIC:AIV - // ratio makes these exact pre-handshake, so scheduler threads can self-assign - // their owned clusters (assign_own_clusters) without the post-handshake - // discovery pass and its all-thread barrier. - aic_count_ = cores_total_num_ / 3; - aiv_count_ = (cores_total_num_ * 2) / 3; - active_sched_threads_ = (sched_thread_num_ > 0) ? sched_thread_num_ : aicpu_thread_num_; + if (cores_total_num_ == 0 || cores_total_num_ > RUNTIME_MAX_WORKER) return -1; + aic_count_ = 0; + aiv_count_ = 0; handshake_failed_.store(false, std::memory_order_release); - // State the barrier-free per-thread init path no longer reaches via - // post_handshake_init; reset on the leader before any scheduler thread is - // released to dispatch. - completed_tasks_.store(0, std::memory_order_release); - orchestrator_done_.store(false, std::memory_order_release); + // State the barrier-free init path (handshake_owned_clusters / + // assign_own_clusters) reads without a leader post_handshake_init: + // scheduler-thread count and func table. The leader sets these before + // publishing hs_setup_done_, so they happen-before any thread's + // self-assignment. The barrier path re-derives them in post_handshake_init, + // so this is redundant (not harmful) there. + // + // payload_per_core_ / deferred_slab_per_core_ are deliberately NOT memset: + // build_payload() overwrites every dispatched payload field and dispatch + // resets slab count/error_code before the slab can be read (the same + // invariant deinit() relies on to skip the ~300 KB zeroing). The global + // memset here was ~37 us of preamble on the critical path — upstream skips + // it and that was the bulk of polling's small-kernel preamble gap. + // sub_block_id is set per owned core in assign_own_clusters and persists + // across runs, so it survives without the memset. + active_sched_threads_ = (sched_thread_num > 0) ? sched_thread_num : aicpu_thread_num; func_id_to_addr_ = runtime->dev.func_id_to_addr_; + return 0; +} - // total_tasks_ must be read before hs_setup_done_ is published: on the - // decoupled path the orchestrator resets the SM as soon as it observes - // hs_setup_done_, which zeroes these ring counters, so the read completes here - // (on the leader, before any thread is released) rather than post-handshake. - if (runtime->get_gm_sm_ptr()) { - auto *header = static_cast(runtime->get_gm_sm_ptr()); - int64_t pto2_count = 0; - for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { - int32_t ring_tasks = header->rings[r].fc.current_task_index.load(std::memory_order_acquire); - if (ring_tasks > 0 && ring_tasks <= PTO2_SCOPE_TASKS_CAP) pto2_count += ring_tasks; - } - total_tasks_ = static_cast(pto2_count); - } else { - total_tasks_ = 0; +void SchedulerContext::handshake_partition(Runtime *runtime, int32_t tidx, int32_t nthreads) { + Handshake *all_handshakes = reinterpret_cast(runtime->dev.workers); + const int32_t total = cores_total_num_; + const int32_t lo = static_cast((static_cast(tidx) * total) / nthreads); + const int32_t hi = static_cast((static_cast(tidx + 1) * total) / nthreads); + + // Step 1: signal this slice's cores to proceed past Phase 1. + for (int32_t i = lo; i < hi; i++) { + all_handshakes[i].task = reinterpret_cast(&payload_per_core_[i][0]); + OUT_OF_ORDER_STORE_BARRIER(); + all_handshakes[i].aicpu_ready = 1; } + OUT_OF_ORDER_STORE_BARRIER(); - LOG_INFO_V0("Handshaking with %d cores", cores_total_num_); - return 0; + uint32_t max_physical_cores_count = platform_get_physical_cores_count(); + + // Step 2: wait for this slice's cores, then init their registers. + // Single-round-trip (#1310): the AICore publishes {physical_core_id, + // core_type, aicore_done} in one write on launch, so aicore_done alone + // gates discovery — the separate aicore_regs_ready / aicpu_regs_ready + // round was removed. Registers are opened after aicore_done is observed. + for (int32_t i = lo; i < hi; i++) { + Handshake *hank = &all_handshakes[i]; + + while (hank->aicore_done == 0) + SPIN_WAIT_HINT(); + + uint32_t physical_core_id = hank->physical_core_id; + + if (physical_core_id >= max_physical_cores_count) { + handshake_failed_.store(true, std::memory_order_release); + continue; + } + + uint64_t *regs = reinterpret_cast(regs_); + uint64_t reg_addr = regs[physical_core_id]; + + CoreType type = hank->core_type; + + // Open this core's window after discovery. + platform_init_aicore_regs(reg_addr); + OUT_OF_ORDER_STORE_BARRIER(); + + core_exec_states_[i].reg_addr = reg_addr; + core_exec_states_[i].cond_ptr = get_reg_ptr(reg_addr, RegId::COND); + + core_exec_states_[i].worker_id = i; + core_exec_states_[i].physical_core_id = physical_core_id; + core_exec_states_[i].core_type = type; + } } int32_t SchedulerContext::post_handshake_init(Runtime *runtime) { @@ -1161,69 +174,37 @@ int32_t SchedulerContext::post_handshake_init(Runtime *runtime) { return -1; } - // Build the AIC/AIV worker-id lists in core-index order, which - // assign_cores_to_threads pairs into clusters. core_type is read from the - // contiguously packed core_type_compact_ the sweep filled, not the 64B-aligned - // per-core Handshake struct. aic_worker_ids_/aiv_worker_ids_ store through to - // HBM, so the lists are built in local (cached) buffers and published with two - // wide memcpys rather than element by element. - int32_t local_aic[RUNTIME_MAX_WORKER]; - int32_t local_aiv[RUNTIME_MAX_WORKER]; - int32_t la = 0, lv = 0; + // Build cluster-ordered AIC/AIV worker-id lists from the discovered + // cores. Serial and MMIO-free — the expensive per-core handshake already + // ran in parallel. Core-index order matches the original single-thread + // handshake so assign_cores_to_threads forms identical clusters. for (int32_t i = 0; i < cores_total_num_; i++) { - if (static_cast(core_type_compact_[i]) == CoreType::AIC) { - local_aic[la++] = i; - } else { - local_aiv[lv++] = i; - } + if (core_exec_states_[i].core_type == CoreType::AIC) aic_worker_ids_[aic_count_++] = i; + else aiv_worker_ids_[aiv_count_++] = i; } - memcpy(aic_worker_ids_, local_aic, static_cast(la) * sizeof(int32_t)); - memcpy(aiv_worker_ids_, local_aiv, static_cast(lv) * sizeof(int32_t)); - aic_count_ = la; - aiv_count_ = lv; - LOG_INFO_V0("Core discovery complete: %d AIC, %d AIV", aic_count_, aiv_count_); - if (!assign_cores_to_threads()) { - return -1; - } + if (!assign_cores_to_threads()) return -1; - // Profiling-subsystem buffer/state init: single-threaded cold path (leader - // only), so the "do it once" guarantee is structural (no CAS needed). Runs - // after the handshake / assign_cores_to_threads because pmu_aicpu_init needs - // physical_core_ids_ / cores_total_num_. Mirrors the l2_swimlane_aicpu_init - // convention above; the per-thread *_set_orch_thread_idx setters stay on the - // orchestrator thread (see aicpu_executor.cpp). -#if SIMPLER_DFX - if (is_dump_args_enabled()) { - dump_args_init(active_sched_threads_); - } - if (is_pmu_enabled()) { - pmu_aicpu_init(physical_core_ids_, cores_total_num_); - LOG_INFO_V0("PMU profiling started on %d cores", cores_total_num_); - } - // dep_gen is host-driven (SubmitTrace) — runtime-gated by the host flag — - // and compiles out with the other profiling subsystems at SIMPLER_DFX=0. - // init() only pops the initial buffer from instance 0's free_queue; the - // orchestrator thread still records its idx via - // dep_gen_aicpu_set_orch_thread_idx() before the first record_submit. - if (is_dep_gen_enabled()) { - dep_gen_aicpu_init(); + // Initialize task counters. Task count comes from PTO2 shared memory. + if (runtime->get_gm_sm_ptr()) { + auto *header = static_cast(runtime->get_gm_sm_ptr()); + int64_t pto2_count = 0; + for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { + int32_t ring_tasks = header->rings[r].fc.current_task_index.load(std::memory_order_acquire); + if (ring_tasks > 0 && ring_tasks <= PTO2_SCOPE_TASKS_CAP) pto2_count += ring_tasks; + } + total_tasks_ = static_cast(pto2_count); + } else { + total_tasks_ = 0; } -#endif - - // total_tasks_ is read in pre_handshake_init (before the orchestrator's early - // SM reset on the decoupled path can zero the ring counters). completed_tasks_.store(0, std::memory_order_release); // Device orchestration: the orchestrator thread flips this when the graph is built. - orchestrator_done_.store(false, std::memory_order_release); + orchestrator_done_ = false; - // prepare_subtask_to_core fully writes a per-core payload / deferred-slab slot - // before the AICore is told to read it: build_payload sets - // function_bin_addr/args/local_context/not_ready, and deferred_slab->count/ - // error_code are reset inline on every dispatch. An AICore reads a slot only - // after a dispatch targets it (DATA_MAIN_BASE), so a prior round's bytes in an - // untouched slot are never observed. + // Clear per-core dispatch payloads + memset(payload_per_core_, 0, sizeof(payload_per_core_)); + memset(deferred_slab_per_core_, 0, sizeof(deferred_slab_per_core_)); // Initialize per-core GlobalContext (sub_block_id) based on cluster position. // This is done once at startup and never modified afterwards. @@ -1240,35 +221,6 @@ int32_t SchedulerContext::post_handshake_init(Runtime *runtime) { } } - // Prefill the per-dispatch AsyncCtx constant fields once. Of AsyncCtx's five - // fields, four are constant for a given (core, buf_idx): the three pointers - // target the fixed deferred_slab_per_core_[core][buf] members, and capacity is - // MAX_COMPLETIONS_PER_TASK. Only task_token varies per dispatch, so build_payload - // writes just that; these constants survive across dispatches because the - // payload buffer is never zeroed between them. - // The two context-pointer args are also per-(core, buf_idx) constants — they - // target this buffer's own local_context / global_context — so prefill them - // here too and drop them from the per-dispatch build_payload writes. This keeps - // the per-dispatch write footprint on the CL0 control block only (args[48]/[49] - // live on a later line). - for (int32_t core_id = 0; core_id < RUNTIME_MAX_WORKER; core_id++) { - for (int32_t buf = 0; buf < 2; buf++) { - PTO2DispatchPayload &dp = payload_per_core_[core_id][buf]; - AsyncCtx &ac = dp.local_context.async_ctx; - volatile DeferredCompletionSlab *slab = &deferred_slab_per_core_[core_id][buf]; - ac.completion_count = &slab->count; - ac.completion_error_code = &slab->error_code; - ac.completion_entries = &slab->entries[0]; - ac.completion_capacity = MAX_COMPLETIONS_PER_TASK; - // Clear the slab once here; thereafter only the completion path re-clears - // count (and only when a deferred task dirtied it), never per dispatch. - slab->count = 0; - slab->error_code = PTO2_ERROR_NONE; - dp.args[PAYLOAD_LOCAL_CONTEXT_INDEX] = reinterpret_cast(&dp.local_context); - dp.args[PAYLOAD_GLOBAL_CONTEXT_INDEX] = reinterpret_cast(&dp.global_context); - } - } - func_id_to_addr_ = runtime->dev.func_id_to_addr_; return 0; @@ -1282,34 +234,28 @@ void SchedulerContext::deinit() { core_exec_states_[i].pending_reg_task_id = AICPU_TASK_INVALID; } - // No per-core memset of payload_per_core_ / deferred_slab_per_core_ here - // (~300 KB across all cores). They are re-initialized before they can be read: - // build_payload() overwrites the per-dispatch payload fields (function addr, - // args[0..num_args) or src_payload, block_idx/block_num, async_ctx.task_token) - // on the exact [core][buf_idx] about to run; the async_ctx slab pointers + - // capacity, the two context-pointer args, and the deferred slab (count = 0 / - // error_code = NONE) are all cleared once per run in init() — the slab is - // thereafter re-cleared only by the completion path after a deferred task. - // The consumer side cannot reach a stale slot either: the - // drain only services a core's running_reg_task_id, and the loop above - // already reset every core_exec_states_[].running/pending_reg_task_id to - // AICPU_TASK_INVALID — so no FIN for an undispatched slot is processed, and - // the count-gated consumer never reads entries[] past the fresh count. + // Neither per-core array is zeroed here (mirrors upstream, ~300 KB + // saved). payload_per_core_: build_payload() overwrites every dispatched + // field including not_ready=0. deferred_slab_per_core_: dispatch resets + // count=0/error_code=NONE before the slab can be read, the loop above + // reset every running/pending_reg_task_id to INVALID so no undispatched + // slot is ever drained, and the consumer is count-gated so it never + // reads entries[] past the fresh count. // Reset sync-start drain coordination — a previous run that aborted mid-drain // would otherwise leave dirty pending/elected/ack state for the next reuse. drain_state_.sync_start_pending.store(0, std::memory_order_release); drain_state_.drain_worker_elected.store(0, std::memory_order_release); drain_state_.drain_ack_mask.store(0, std::memory_order_release); - drain_state_.drain_stage_go.store(0, std::memory_order_release); - drain_state_.drain_stage_done_mask.store(0, std::memory_order_release); - drain_state_.drain_running_staged.store(0, std::memory_order_release); drain_state_.pending_task.store(nullptr, std::memory_order_release); // Reset task counters and orchestrator state completed_tasks_.store(0, std::memory_order_release); total_tasks_ = 0; - orchestrator_done_.store(false, std::memory_order_release); + orchestrator_done_ = false; + pto2_init_done_.store(false, std::memory_order_release); + pto2_init_complete_.store(false, std::memory_order_release); + completed_.store(false, std::memory_order_release); // Reset core discovery and assignment state @@ -1319,9 +265,8 @@ void SchedulerContext::deinit() { aicpu_thread_num_ = 0; sched_thread_num_ = 0; active_sched_threads_ = 0; - for (int32_t t = 0; t < MAX_AICPU_THREADS; t++) { + for (int32_t t = 0; t < MAX_AICPU_THREADS; t++) core_trackers_[t] = CoreTracker{}; - } regs_ = 0; sched_ = nullptr; @@ -1329,71 +274,255 @@ void SchedulerContext::deinit() { func_id_to_addr_ = nullptr; } +int32_t SchedulerContext::shutdown(int32_t thread_idx) { + const int32_t *cores = core_trackers_[thread_idx].core_ids(); + int32_t core_num = core_trackers_[thread_idx].core_num(); + if (core_num == 0) return 0; + + int32_t rc = 0; + for (int32_t i = 0; i < core_num; i++) { + int32_t core_id = cores[i]; + uint64_t reg_addr = core_exec_states_[core_id].reg_addr; + if (reg_addr != 0) { + // Timeout means AICore is unresponsive. Log and continue deiniting remaining cores. + if (platform_deinit_aicore_regs(reg_addr) != 0) rc = -1; + } else { + } + } + return rc; +} + +void SchedulerContext::on_orchestration_done(Runtime *runtime, PTO2Runtime *rt, int32_t, int32_t total_tasks) { + on_orchestration_done(runtime, rt, total_tasks); +} + +void SchedulerContext::on_orchestration_done(Runtime *runtime, PTO2Runtime *rt, int32_t total_tasks) { + total_tasks_ = total_tasks; + + // Fold tasks completed inline during orchestration + int32_t inline_completed = static_cast(rt->orchestrator.inline_completed_tasks); + if (inline_completed > 0) completed_tasks_.fetch_add(inline_completed, std::memory_order_relaxed); + orchestrator_done_ = true; + + // Check for fatal error from orchestration; if so, shut down immediately. + int32_t orch_err = 0; + if (sched_->sm_header) orch_err = sched_->sm_header->orch_error_code.load(std::memory_order_relaxed); + if (orch_err != PTO2_ERROR_NONE) { + if (!completed_.exchange(true, std::memory_order_acq_rel)) emergency_shutdown(runtime); + } +} + void SchedulerContext::bind_runtime(PTO2Runtime *rt) { rt_ = rt; sched_ = &rt->scheduler; } -void SchedulerContext::wait_for_orchestration_done_before_dispatch(Runtime *runtime, int32_t thread_idx) { - while (!orchestration_done() && !completed_.load(std::memory_order_acquire)) { - if (sched_ != nullptr && sched_->sm_header != nullptr && - check_idle_fatal_error(thread_idx, sched_->sm_header, runtime) == LoopAction::BREAK_LOOP) { - break; +void SchedulerContext::wait_for_orchestration_done_before_dispatch(Runtime *, int32_t thread_idx) { + while (!orchestrator_done_) { + if (thread_idx == 0 && sched_ != nullptr) { + sched_->drain_wiring_queue(false); } SPIN_WAIT_HINT(); } } -// ============================================================================= -// Post-orchestration bookkeeping. Runs on the orchestrator thread once the -// build phase finishes; folds inline-completed tasks, flips orchestrator_done_, -// and drives the orchestrator → scheduler core transition (or fatal shutdown). -// ============================================================================= -void SchedulerContext::on_orchestration_done( - Runtime *runtime, PTO2Runtime *rt, [[maybe_unused]] int32_t thread_idx, int32_t total_tasks -) { -#if SIMPLER_DFX - if (l2_swimlane_level_ >= L2SwimlaneLevel::ORCH_PHASES) { - // Flush the orchestrator's orch-phase buffer (single instance, pool 0). - // The orchestrator has no scheduler-phase pool of its own — those belong - // to the scheduler threads and are flushed in scheduler_dispatch. - l2_swimlane_aicpu_flush_orch_phase_buffer(thread_idx); +bool SchedulerContext::assign_cores_to_threads() { + // Cluster-aligned round-robin assignment: cluster ci -> sched thread ci % active_sched_threads_. + // Each cluster = 1 AIC + 2 adjacent AIV; the triple is always kept together. + active_sched_threads_ = (sched_thread_num_ > 0) ? sched_thread_num_ : aicpu_thread_num_; + int32_t cluster_count = aic_count_; + + // Max clusters any single sched thread can hold: ceil(cluster_count / active_sched_threads_). + int32_t max_clusters_per_thread = (cluster_count + active_sched_threads_ - 1) / active_sched_threads_; + int32_t thread_cores_num = max_clusters_per_thread * 3; + + if (thread_cores_num > CoreTracker::MAX_CORE_PER_THREAD) return false; + + for (int32_t i = 0; i < RUNTIME_MAX_WORKER; i++) { + core_exec_states_[i].running_reg_task_id = AICPU_TASK_INVALID; + core_exec_states_[i].pending_reg_task_id = AICPU_TASK_INVALID; } -#endif - total_tasks_ = total_tasks; + // Count clusters per thread first (round-robin may distribute unevenly) + int32_t clusters_per_thread[MAX_AICPU_THREADS] = {}; + for (int32_t ci = 0; ci < cluster_count; ci++) + clusters_per_thread[ci % active_sched_threads_]++; + for (int32_t i = 0; i < active_sched_threads_; i++) + core_trackers_[i].init(clusters_per_thread[i]); - // Fold tasks completed inline during orchestration - int32_t inline_completed = static_cast(rt->orchestrator.inline_completed_tasks); - if (inline_completed > 0) { - completed_tasks_.fetch_add(inline_completed, std::memory_order_relaxed); -#if SIMPLER_SCHED_PROFILING - rt->scheduler.tasks_completed.fetch_add(inline_completed, std::memory_order_relaxed); -#endif + int32_t cluster_idx_per_thread[MAX_AICPU_THREADS] = {}; + + for (int32_t ci = 0; ci < cluster_count; ci++) { + int32_t t = ci % active_sched_threads_; + + int32_t aic_wid = aic_worker_ids_[ci]; + int32_t aiv0_wid = aiv_worker_ids_[2 * ci]; + int32_t aiv1_wid = aiv_worker_ids_[2 * ci + 1]; + + core_trackers_[t].set_cluster(cluster_idx_per_thread[t]++, aic_wid, aiv0_wid, aiv1_wid); } - orchestrator_done_.store(true, std::memory_order_release); - // Check for fatal error from orchestration; if so, shut down immediately. - int32_t orch_err = 0; - if (sched_->sm_header) { - orch_err = sched_->sm_header->orch_error_code.load(std::memory_order_relaxed); + for (int32_t t = 0; t < aicpu_thread_num_; t++) {} + + return true; +} + +void SchedulerContext::emergency_shutdown(Runtime *runtime) { + Handshake *all_handshakes = reinterpret_cast(runtime->dev.workers); + int32_t timeout_count = 0; + for (int32_t i = 0; i < cores_total_num_; i++) { + Handshake *hank = &all_handshakes[i]; + OUT_OF_ORDER_STORE_BARRIER(); + (void)hank; // single-round-trip: no aicpu_regs_ready round; deinit forces exit + if (core_exec_states_[i].reg_addr != 0) { + if (platform_deinit_aicore_regs(core_exec_states_[i].reg_addr) != 0) timeout_count++; + } } + if (timeout_count > 0) {} +} + +LoopAction SchedulerContext::handle_orchestrator_exit(PTO2SharedMemoryHeader *header, Runtime *runtime) { + if (completed_.load(std::memory_order_acquire)) return LoopAction::BREAK_LOOP; + int32_t orch_err = header->orch_error_code.load(std::memory_order_acquire); if (orch_err != PTO2_ERROR_NONE) { - if (!completed_.exchange(true, std::memory_order_acq_rel)) { - emergency_shutdown(runtime); - } + if (!completed_.exchange(true, std::memory_order_acq_rel)) emergency_shutdown(runtime); + return LoopAction::BREAK_LOOP; + } + int32_t sched_err = header->sched_error_code.load(std::memory_order_acquire); + if (sched_err != PTO2_ERROR_NONE) { + if (!completed_.exchange(true, std::memory_order_acq_rel)) emergency_shutdown(runtime); + return LoopAction::BREAK_LOOP; + } + + if (!orchestrator_done_) return LoopAction::NONE; + + if (total_tasks_ > 0 && completed_tasks_.load(std::memory_order_relaxed) >= total_tasks_) { + completed_.store(true, std::memory_order_release); + return LoopAction::BREAK_LOOP; + } + return LoopAction::NONE; +} + +LoopAction SchedulerContext::check_idle_fatal_error(PTO2SharedMemoryHeader *header, Runtime *runtime) { + if (completed_.load(std::memory_order_acquire)) return LoopAction::BREAK_LOOP; + int32_t orch_err = header->orch_error_code.load(std::memory_order_acquire); + if (orch_err != PTO2_ERROR_NONE) { + if (!completed_.exchange(true, std::memory_order_acq_rel)) emergency_shutdown(runtime); + return LoopAction::BREAK_LOOP; + } + int32_t sched_err = header->sched_error_code.load(std::memory_order_acquire); + if (sched_err != PTO2_ERROR_NONE) { + if (!completed_.exchange(true, std::memory_order_acq_rel)) emergency_shutdown(runtime); + return LoopAction::BREAK_LOOP; } + return LoopAction::NONE; +} + +void SchedulerContext::log_stall_diagnostics(int32_t thread_idx) { + CoreTracker &tracker = core_trackers_[thread_idx]; + + // T0 owns the shared-ring scan; printing it from other threads would + // produce identical TASK lines once per scheduler thread. + if (thread_idx == 0) { + int32_t cnt_ready = 0, cnt_waiting = 0, cnt_running = 0; + for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { + PTO2SharedMemoryRingHeader &ring = *sched_->ring_sched_states[r].ring; + int32_t ring_task_count = ring.fc.current_task_index.load(std::memory_order_relaxed); + for (int32_t si = 0; si < ring_task_count; si++) { + PTO2TaskSlotState &slot_state = ring.get_slot_state_by_task_id(si); + // (m) task_state retired; use completion_flags directly. + bool fanin_ready = sched_->fanin_satisfied(&slot_state); + if (ring.completion_flags[si & ring.task_window_mask].load(std::memory_order_relaxed) != 0) continue; + char running_on[192] = {0}; + int32_t owner = -1; + int32_t pos = 0; + bool is_running = false; + for (int32_t cid = 0; cid < cores_total_num_ && pos + 32 < (int32_t)sizeof(running_on); cid++) { + if (core_exec_states_[cid].running_slot_state != &slot_state) continue; + is_running = true; + if (owner < 0) owner = find_core_owner_thread(cid); + const char *sname = subslot_name(core_exec_states_[cid].running_subslot); + int32_t written = snprintf( + running_on + pos, sizeof(running_on) - pos, "%score=%d(%s)", pos == 0 ? "" : " ", cid, sname + ); + if (written > 0) pos += written; + } -#if SIMPLER_DFX - // Write the core-to-thread mapping so the profiling data reflects the - // scheduler threads' final core distribution. - if (l2_swimlane_level_ >= L2SwimlaneLevel::SCHED_PHASES) { - l2_swimlane_aicpu_init_core_assignments(cores_total_num_); - for (int32_t t = 0; t < active_sched_threads_; t++) { - l2_swimlane_aicpu_write_core_assignments_for_thread( - t, core_trackers_[t].core_ids(), core_trackers_[t].core_num() - ); + if (is_running) { + cnt_running++; + if (cnt_running > STALL_DUMP_READY_MAX) continue; + continue; + } + if (fanin_ready) { + cnt_ready++; + if (cnt_ready > STALL_DUMP_READY_MAX) continue; + continue; + } + cnt_waiting++; + if (cnt_waiting > STALL_DUMP_WAIT_MAX) continue; + } } } -#endif + + for (int32_t cli = 0; cli < tracker.get_cluster_count() && cli < STALL_DUMP_CORE_MAX; cli++) { + int32_t offset = cli * 3; + int32_t aic_id = tracker.get_aic_core_id(offset); + int32_t aiv0_id = tracker.get_aiv0_core_id(offset); + int32_t aiv1_id = tracker.get_aiv1_core_id(offset); + bool aic_idle = tracker.is_aic_core_idle(offset); + bool aiv0_idle = tracker.is_aiv0_core_idle(offset); + bool aiv1_idle = tracker.is_aiv1_core_idle(offset); + char aic_buf[128], aiv0_buf[128], aiv1_buf[128]; + format_core_status( + aic_buf, sizeof(aic_buf), aic_id, aic_idle, &core_exec_states_[aic_id], core_exec_states_[aic_id].reg_addr + ); + format_core_status( + aiv0_buf, sizeof(aiv0_buf), aiv0_id, aiv0_idle, &core_exec_states_[aiv0_id], + core_exec_states_[aiv0_id].reg_addr + ); + format_core_status( + aiv1_buf, sizeof(aiv1_buf), aiv1_id, aiv1_idle, &core_exec_states_[aiv1_id], + core_exec_states_[aiv1_id].reg_addr + ); + } +} + +void SchedulerContext::log_shutdown_stall_snapshot() { + int32_t thread_count = active_sched_threads_ > 0 ? active_sched_threads_ : aicpu_thread_num_; + if (thread_count < 0 || thread_count > MAX_AICPU_THREADS) thread_count = thread_count < 0 ? 0 : MAX_AICPU_THREADS; + for (int32_t t = 0; t < thread_count; t++) + log_stall_diagnostics(t); +} + +int32_t SchedulerContext::find_core_owner_thread(int32_t core_id) const { + for (int32_t t = 0; t < aicpu_thread_num_; t++) { + const int32_t *ids = core_trackers_[t].core_ids(); + int32_t n = core_trackers_[t].core_num(); + for (int32_t i = 0; i < n; i++) + if (ids[i] == core_id) return t; + } + return -1; +} + +bool SchedulerContext::self_owns_running_task(int32_t thread_idx) const { + const int32_t *cores = core_trackers_[thread_idx].core_ids(); + int32_t core_num = core_trackers_[thread_idx].core_num(); + for (int32_t i = 0; i < core_num; i++) + if (core_exec_states_[cores[i]].running_slot_state != nullptr) return true; + return false; +} + +bool SchedulerContext::no_thread_owns_running_task() const { + for (int32_t t = 0; t < aicpu_thread_num_; t++) + if (self_owns_running_task(t)) return false; + return true; +} + +int32_t SchedulerContext::handle_timeout_exit(int32_t thread_idx, PTO2SharedMemoryHeader *header, Runtime *runtime) { + latch_scheduler_error(header, thread_idx, PTO2_ERROR_SCHEDULER_TIMEOUT); + if (!completed_.exchange(true, std::memory_order_acq_rel)) { + log_shutdown_stall_snapshot(); + emergency_shutdown(runtime); + } + return -PTO2_ERROR_SCHEDULER_TIMEOUT; } diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_completion.cpp b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_completion.cpp index 98ee54dc98..6915403bdf 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_completion.cpp +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_completion.cpp @@ -8,36 +8,20 @@ * See LICENSE in the root of the software repository for the full text of the License. * ----------------------------------------------------------------------------------------------------------- */ + #include "scheduler_context.h" #include +#include +#include -#include "common/unified_log.h" -#include "aicpu/device_time.h" -#include "aicpu/platform_regs.h" -#include "common/l2_swimlane_profiling.h" -#include "common/memory_barrier.h" -#include "common/platform_config.h" -#include "pto_runtime2.h" -#include "runtime.h" -#include "spin_hint.h" - -// Performance profiling headers -#include "aicpu/l2_swimlane_collector_aicpu.h" -#include "aicpu/pmu_collector_aicpu.h" -#include "aicpu/args_dump_aicpu.h" - -// ============================================================================= -// Dual-slot state machine helpers -// ============================================================================= - -namespace { -inline constexpr int32_t PTO2_DEFERRED_RELEASE_CAP = 256; -} +#include "common.h" +#include "callable.h" +#include "aicpu/aicpu_device_config.h" +#include "aicpu/dep_gen_collector_aicpu.h" -// Pure function: read register result -> SlotTransition (no side effects). SlotTransition SchedulerContext::decide_slot_transition( - int32_t reg_task_id, int32_t reg_state, int32_t running_id, int32_t pending_id, bool pending_gated + int32_t reg_task_id, int32_t reg_state, int32_t running_id, int32_t pending_id ) { SlotTransition t; if (pending_id != AICPU_TASK_INVALID && reg_task_id == pending_id) { @@ -45,9 +29,7 @@ SlotTransition SchedulerContext::decide_slot_transition( t.running_done = true; // Serial execution: pending event implies running done t.running_freed = true; t.pending_freed = true; - if (reg_state == TASK_FIN_STATE) { - t.pending_done = true; // Case 1: pending FIN - } + if (reg_state == TASK_FIN_STATE) t.pending_done = true; // Case 1: pending FIN // else: Case 2: pending ACK (pending_done stays false) } else if (reg_task_id == running_id) { if (reg_state == TASK_FIN_STATE) { @@ -56,22 +38,9 @@ SlotTransition SchedulerContext::decide_slot_transition( t.matched = true; t.running_done = true; t.running_freed = true; - } else if (pending_gated) { - // Case 3.3: running FIN, pending is a EARLY-DISPATCH GATED task. The - // Case 3.1 "wait for the pending's ack" shortcut assumes the AICore - // immediately runs the pending task; a gated task instead spins on - // its doorbell and never acks until its producer completes — and - // that producer's completion depends on collecting THIS running FIN. - // Waiting would deadlock. Complete the running FIN now and promote - // the gated task (it then skip-gates until its doorbell). pending is - // NOT freed (it promotes, not retires) so the bitmap update keeps the - // core off-limits — no second gated block, no doorbell overwrite. - t.matched = true; - t.running_done = true; - t.running_freed = true; } - // Case 3.1: running FIN, NON-gated pending exists -> skip (transient - // state). Case 1/2 (pending ack/FIN) completes running implicitly. + // Case 3.1: running FIN, pending exists -> skip (transient state). + // Case 1/2 (pending ACK/FIN) will complete running implicitly via running_done=true. } else { // Case 4: running ACK -- only pending_freed (slot now hardware-latched) t.matched = true; @@ -81,60 +50,40 @@ SlotTransition SchedulerContext::decide_slot_transition( return t; } -// 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 -#if SIMPLER_DFX - , - uint64_t dispatch_ts, uint64_t finish_ts -#endif + PTO2TaskSlotState &slot_state, int32_t expected_reg_task_id, int32_t core_id, int32_t &completed_this_turn ) { -#if SIMPLER_DFX - auto &l2_swimlane = sched_l2_swimlane_[thread_idx]; -#else - (void)hank; -#endif - // MPSC fast-path is opt-in per task: only tasks with at least one subtask - // that registered a deferred condition route through the mailbox. Pure - // non-deferred tasks complete inline on this thread (matching pre-MPSC - // behavior — keeps the common case parallelized across scheduler threads - // instead of serializing through the single consumer). The - // deferred-completion flag on slot_state is the discriminator; it's set - // (release) before on_subtask_complete and read (acquire) after, so the - // last subtask sees flag writes from any earlier subtask of the same task. AICoreCompletionMailbox *mailbox = rt_ != nullptr ? rt_->aicore_mailbox : nullptr; bool defer_completion_to_consumer = false; if (slot_state.payload != nullptr) { volatile DeferredCompletionSlab *deferred_slab = &deferred_slab_per_core_[core_id][expected_reg_task_id & 1]; - int32_t slab_err = deferred_slab->error_code; - if (slab_err != PTO2_ERROR_NONE) { - int32_t expected = PTO2_ERROR_NONE; - sched_->sm_header->sched_error_code.compare_exchange_strong( - expected, slab_err, std::memory_order_acq_rel, std::memory_order_acquire - ); - completed_.store(true, std::memory_order_release); - return; - } - + // (q) Read count first. AICore only writes error_code as part of a + // condition-registration attempt that also increments count, so + // count == 0 ⇒ no error and no conditions to forward. This is the + // common path for kernels that don't use async waits (paged + // attention, GEMM, etc.) and saves an L1 load + branch per call. uint32_t cond_count = deferred_slab->count; - if (cond_count > MAX_COMPLETIONS_PER_TASK) { - int32_t expected = PTO2_ERROR_NONE; - sched_->sm_header->sched_error_code.compare_exchange_strong( - expected, PTO2_ERROR_ASYNC_REGISTRATION_FAILED, std::memory_order_acq_rel, std::memory_order_acquire - ); - completed_.store(true, std::memory_order_release); - return; - } + if (cond_count != 0) { + int32_t slab_err = deferred_slab->error_code; + if (slab_err != PTO2_ERROR_NONE) { + int32_t expected = PTO2_ERROR_NONE; + sched_->sm_header->sched_error_code.compare_exchange_strong( + expected, slab_err, std::memory_order_acq_rel, std::memory_order_acquire + ); + completed_.store(true, std::memory_order_release); + return; + } + if (cond_count > MAX_COMPLETIONS_PER_TASK) { + int32_t expected = PTO2_ERROR_NONE; + sched_->sm_header->sched_error_code.compare_exchange_strong( + expected, PTO2_ERROR_ASYNC_REGISTRATION_FAILED, std::memory_order_acq_rel, std::memory_order_acquire + ); + completed_.store(true, std::memory_order_release); + return; + } - if (cond_count > 0) { - // Publish "this task is deferred" before on_subtask_complete so the - // acq_rel fetch_add inside on_subtask_complete makes the flag - // visible to whichever subtask sees task_complete=true (which may - // be this thread or a later one). - slot_state.mark_any_subtask_deferred(); + slot_state.any_subtask_deferred.store(true, std::memory_order_release); const PTO2TaskId token = slot_state.task->task_id; for (uint32_t i = 0; i < cond_count; ++i) { @@ -144,25 +93,13 @@ void SchedulerContext::complete_slot_task( SPIN_WAIT_HINT(); } } - // Re-clear for the next reuse of this (core, buf) slot. Done here — on - // the hot cache line we just read — instead of on every dispatch, since - // only a deferred task (count > 0) ever dirties it. error_code needs no - // reset: a non-NONE code aborted the run above. - deferred_slab->count = 0; } } - bool task_complete = sched_->on_subtask_complete(slot_state); + bool mixed_complete = sched_->on_subtask_complete(slot_state); -#if SIMPLER_DFX - // Sub-block retire that did not finish the slot: record it so the poll - // iteration becomes visible on the scheduler lane (the SPMD harvest tail). - if (!task_complete && l2_swimlane_level_ >= L2SwimlaneLevel::SCHED_PHASES) { - l2_swimlane.phase_subretire_count++; - } -#endif - - if (task_complete && slot_state.payload != nullptr && slot_state.has_any_subtask_deferred()) { + if (mixed_complete && slot_state.payload != nullptr && + slot_state.any_subtask_deferred.load(std::memory_order_acquire)) { // Some subtask of this task registered conditions; finish the // registration by handing the slot_state off to the consumer. while (!mailbox->try_push_normal_done(slot_state.task->task_id, reinterpret_cast(&slot_state))) { @@ -172,137 +109,28 @@ void SchedulerContext::complete_slot_task( defer_completion_to_consumer = true; } - if (task_complete && !defer_completion_to_consumer) { -#if SIMPLER_DFX - if (is_dump_args_enabled()) { - dump_args_for_task( - thread_idx, slot_state, ArgsDumpStage::AFTER_COMPLETION, - [](ActiveMask active_mask, int raw_subtask_id) { - return active_mask.subtask_active(static_cast(raw_subtask_id)); - }, - [this](int32_t func_id) { - return get_function_bin_addr(func_id); - } - ); - } -#endif -#if SIMPLER_DFX - // Time Resolve (walk the consumer list, decrement each consumer's - // fanin, push the newly-ready ones, ring doorbells for early-dispatch - // hits) so it renders as a child bar nested inside this iteration's - // Complete bar. The 1 µs floor below filters out the ~88% of tasks - // with 1-2 consumers (~500 ns Resolve) so only the long broadcast / - // reduction walks stand out on the lane. - uint64_t resolve_t0 = (l2_swimlane_level_ >= L2SwimlaneLevel::SCHED_PHASES) ? get_sys_cnt_aicpu() : 0; -#endif - // [[maybe_unused]] silences -Werror=unused-but-set-variable on the - // profiling-flags-smoke build path where SIMPLER_DFX is OFF and - // the Resolve emit below is excluded. - [[maybe_unused]] uint32_t consumers_resolved = 0; -#if SIMPLER_SCHED_PROFILING - // SCHED_PROFILING variant takes thread_idx for its per-thread atomic - // counter side-effects (g_sched_*_atomic_count[thread_idx], consumed - // by the otc_* log lines). It returns CompletionStats whose - // `fanout_edges` is the consumer-walk count. - consumers_resolved = sched_->on_task_complete(slot_state, thread_idx).fanout_edges; -#else - consumers_resolved = sched_->on_task_complete(slot_state); -#endif -#if SIMPLER_DFX - if (resolve_t0 != 0) { - uint64_t resolve_t1 = get_sys_cnt_aicpu(); - // Filter: drop Resolve bars under 1 µs so the lane shows only - // resolves that did meaningful work (high consumer counts or - // doorbells). 50 cycles @ 50 MHz = 1 µs (PLATFORM_PROF_SYS_CNT_FREQ - // is the device sys-cnt frequency). - constexpr uint64_t RESOLVE_EMIT_MIN_CYCLES = PLATFORM_PROF_SYS_CNT_FREQ / 1'000'000; // 1 µs - if (resolve_t1 - resolve_t0 >= RESOLVE_EMIT_MIN_CYCLES) { - l2_swimlane_aicpu_record_sched_phase( - thread_idx, L2SwimlaneSchedPhaseKind::Resolve, resolve_t0, resolve_t1, l2_swimlane.sched_loop_count, - consumers_resolved - ); - } - } - l2_swimlane.phase_complete_count++; -#endif - 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; - } + if (mixed_complete && !defer_completion_to_consumer) { + sched_->on_mixed_task_complete(slot_state); completed_this_turn++; } - -#if SIMPLER_DFX - // Level gate: at AICORE_TIMING (level=1) the AICore record alone carries - // {start, end, task_token_raw}, host resolves func_id/core_type from - // dep_gen / per-core mapping, and AICPU has nothing to write. Only at - // AICPU_TIMING (level=2) and above does AICPU contribute dispatch/finish - // timestamps via complete_task. Bypassing here saves the per-completion - // hot-path cost (counter inc + ring lookup + record store + wmb + buffer - // rotation bookkeeping) for runs that only want AICore timing. - if (l2_swimlane.l2_swimlane_enabled && l2_swimlane_level_ >= L2SwimlaneLevel::AICPU_TIMING) { -#if SIMPLER_SCHED_PROFILING - uint64_t t_perf_start = get_sys_cnt_aicpu(); -#endif - - if (l2_swimlane_aicpu_complete_task( - core_id, thread_idx, static_cast(expected_reg_task_id), dispatch_ts, finish_ts - ) != 0) { - LOG_ERROR( - "Core %d: l2_swimlane_aicpu_complete_task failed for task 0x%" PRIx64, core_id, - static_cast(slot_state.task->task_id.raw) - ); - } -#if SIMPLER_SCHED_PROFILING - l2_swimlane.sched_complete_perf_cycle += (get_sys_cnt_aicpu() - t_perf_start); -#endif - } - - if (is_pmu_enabled()) { - pmu_aicpu_record_task( - core_id, thread_idx, slot_state.task->task_id.raw, - slot_state.task->kernel_id[static_cast(subslot)], hank[core_id].core_type - ); - } -#endif } -// Promote pending slot data to running slot. Clears pending fields. void SchedulerContext::promote_pending_to_running(CoreExecState &core) { core.running_slot_state = core.pending_slot_state; core.running_reg_task_id = core.pending_reg_task_id; core.running_subslot = core.pending_subslot; -#if SIMPLER_DFX - core.running_dispatch_timestamp = core.pending_dispatch_timestamp; -#endif core.pending_slot_state = nullptr; core.pending_reg_task_id = AICPU_TASK_INVALID; } -// Clear running slot (core becomes idle). void SchedulerContext::clear_running_slot(CoreExecState &core) { core.running_slot_state = nullptr; core.running_reg_task_id = AICPU_TASK_INVALID; } 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 + int32_t thread_idx, int32_t &completed_this_turn, int32_t &cur_thread_completed, bool &made_progress ) { -#if SIMPLER_SCHED_PROFILING - auto &l2_swimlane = sched_l2_swimlane_[thread_idx]; -#endif CoreTracker &tracker = core_trackers_[thread_idx]; auto running_core_states = tracker.get_all_running_cores(); while (running_core_states.has_value()) { @@ -310,159 +138,40 @@ void SchedulerContext::check_running_cores_for_completion( int32_t core_id = tracker.get_core_id_by_offset(bit_pos); CoreExecState &core = core_exec_states_[core_id]; - // Skip gated early-dispatch cores. A STAGED task is parked on this core - // waiting for its doorbell — it physically cannot ACK/FIN yet, so - // reading its COND (MMIO, and the core is hot-spinning on its own SPR) - // every poll is pure waste that drags out the completion phase. The - // doorbell (try_early_dispatch_release) flips early_dispatch_state to DISPATCHED, at - // which point the core becomes pollable again and its FIN is caught. - // Cheap cacheable load; no MMIO. Pending slot is empty while gated. - { - PTO2TaskSlotState *rs = core.running_slot_state; - if (rs != nullptr && rs->payload != nullptr && - rs->payload->early_dispatch_state.load(std::memory_order_relaxed) == PTO2_EARLY_DISPATCH_STAGING) { - continue; - } - } - - // --- Judgment phase: read register, derive transition --- - // Use the precomputed cond_ptr (resolved once in handshake) to skip - // the reg_offset switch and reg_addr addition on every poll. - // reg_load_acquire makes this an atomic acquire under __CPU_SIM so it - // pairs with the AICore's release store of the FIN (without it the sim - // poll races the FIN publish and can miss it); on hardware it is the - // same plain volatile load the bare deref used to be. - uint64_t reg_val = static_cast(reg_load_acquire(core.cond_ptr)); - // ARM64 allows Device-nGnRnE -> Normal-cacheable load reorder; the - // rmb() pins any AICore-published cacheable reads downstream of the - // FIN observation. Replaces the post-`__sync_synchronize` that the - // old read_reg() helper carried implicitly. + uint64_t reg_val = static_cast(*core.cond_ptr); rmb(); int32_t reg_task_id = EXTRACT_TASK_ID(reg_val); int32_t reg_state = EXTRACT_TASK_STATE(reg_val); -#if SIMPLER_SCHED_PROFILING - if (l2_swimlane.l2_swimlane_enabled) { - l2_swimlane.complete_probe_count++; - } -#endif - - // A pending task is "gated" when it is an early-dispatch pre-stage still parked on - // its doorbell: it will not ack on the producer's FIN, so the Case 3.1 wait-for- - // pending-ack shortcut would deadlock. Detect it so decide_slot_transition completes - // the running FIN and PROMOTES it (Case 3.3) instead. - // - // "Gated" is "not yet launched (rung)", which is NOT the same as - // early_dispatch_state. STAGING covers the pre-release window. But a sync_start block - // stays gated even AFTER its producer releases (early_dispatch_state STAGING -> - // DISPATCHED): the cohort is not rung until the rendezvous has assembled EVERY core - // into a running slot, and this pending block (being promoted now) is by definition - // not yet counted, so the ring has not fired and it is still gated. Classifying it by - // STAGING alone would, once the producer's release beats the last promotion, treat it - // as a normal task and wait for an ack that never comes -> deadlock (the - // nondeterministic sync_start stall). - uint8_t pending_ss = - (core.pending_slot_state != nullptr && core.pending_slot_state->payload != nullptr) ? - core.pending_slot_state->payload->early_dispatch_state.load(std::memory_order_relaxed) : - static_cast(PTO2_EARLY_DISPATCH_NONE); - bool pending_gated = - (core.pending_slot_state != nullptr && core.pending_slot_state->payload != nullptr && - (pending_ss == PTO2_EARLY_DISPATCH_STAGING || - (pending_ss == PTO2_EARLY_DISPATCH_DISPATCHED && - core.pending_slot_state->active_mask.requires_sync_start()))); - SlotTransition t = decide_slot_transition( - reg_task_id, reg_state, core.running_reg_task_id, core.pending_reg_task_id, pending_gated - ); + SlotTransition t = + decide_slot_transition(reg_task_id, reg_state, core.running_reg_task_id, core.pending_reg_task_id); if (!t.matched) continue; -#if SIMPLER_DFX - // Release an ACK-gated AICore swimlane buffer if this matched ACK/FIN is - // the one a rotation is waiting on: it proves the core advanced past the - // just-rotated buffer's tail record (FIN precedes the record write on - // this runtime, so rotation could not release the buffer itself). - if (l2_swimlane_level_ != L2SwimlaneLevel::DISABLED) { - l2_swimlane_aicpu_on_aicore_ack(core_id, thread_idx, static_cast(reg_task_id)); - } -#endif - -#if SIMPLER_SCHED_PROFILING - if (l2_swimlane.l2_swimlane_enabled && (t.running_done || t.pending_done)) { - l2_swimlane.complete_hit_count++; - } -#endif - -#if SIMPLER_DFX - // Capture finish_ts at the FIN observation point — right after rmb() - // above pinned the cacheable AICore reads downstream of the register - // load, and BEFORE any fanin / deferred-release work. Anything later - // (slot transition apply, complete_slot_task fanin processing) would - // charge AICPU completion-processing cost to the (end → finish) - // span, masking the actual FIN-delivery latency. - uint64_t finish_ts = 0; - if (l2_swimlane_level_ >= L2SwimlaneLevel::AICPU_TIMING && (t.pending_done || t.running_done)) { - finish_ts = get_sys_cnt_aicpu(); - } -#endif - // --- Apply phase: execute actions based on transition --- // 1. Complete finished tasks (capture pointers before modifying core state) if (t.pending_done) { // Task-timing finish: latest FIN observation for a tagged task, folded - // as max. Sampled after the rmb above and before complete_slot_task runs - // fanin / deferred-completion (which may also clear pending_slot_state), - // matching L2's finish_time point. Independent of L2 swimlane level, so - // it works in SIMPLER_DFX=0 builds; untagged tasks pay only the compare. - if (core.pending_slot_state->task->task_timing_slot != TASK_TIMING_SLOT_NONE) { + // as max. Sampled before complete_slot_task clears pending_slot_state. + if (core.pending_slot_state->task->task_timing_slot != TASK_TIMING_SLOT_NONE) aicpu_task_timing_finish(core.pending_slot_state->task->task_timing_slot, thread_idx); - } - 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 -#if SIMPLER_DFX - , - core.pending_dispatch_timestamp, finish_ts -#endif - ); + complete_slot_task(*core.pending_slot_state, core.pending_reg_task_id, core_id, completed_this_turn); cur_thread_completed++; } if (t.running_done) { - if (core.running_slot_state->task->task_timing_slot != TASK_TIMING_SLOT_NONE) { + if (core.running_slot_state->task->task_timing_slot != TASK_TIMING_SLOT_NONE) aicpu_task_timing_finish(core.running_slot_state->task->task_timing_slot, thread_idx); - } - 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 -#if SIMPLER_DFX - , - core.running_dispatch_timestamp, finish_ts -#endif - ); + complete_slot_task(*core.running_slot_state, core.running_reg_task_id, core_id, completed_this_turn); cur_thread_completed++; } // 2. Update slot data if (t.running_freed) { if (core.pending_slot_state != nullptr && !t.pending_done) { - // A gated sync_start block promoting into the running slot advances the - // rendezvous. Capture that BEFORE promote nulls the pending fields; after - // it lands, bump running_slot_count and ring iff this was the block that - // completed the cohort (and the producer already released). - PTO2TaskSlotState *promoted = core.pending_slot_state; - bool sync_start_promote = pending_gated && promoted->active_mask.requires_sync_start(); promote_pending_to_running(core); // Case 2 or Case 3 (with pending) - if (sync_start_promote) { - promoted->payload->running_slot_count.fetch_add(1, std::memory_order_seq_cst); - if (sched_->maybe_rendezvous_ring(*promoted)) { - sched_->propagate_dispatch_fanin(*promoted); - } - } } else { clear_running_slot(core); // Case 1 or Case 3 (no pending) if (t.pending_done) { - // Case 1: pending FIN observed directly -- clear stale pending fields. - // Without this, pending_reg_task_id retains a stale value that blocks - // clear_pending_occupied and permanently degrades pipelining. core.pending_slot_state = nullptr; core.pending_reg_task_id = AICPU_TASK_INVALID; } @@ -475,230 +184,76 @@ void SchedulerContext::check_running_cores_for_completion( tracker.change_core_state(bit_pos); // Mark idle tracker.clear_pending_occupied(bit_pos); // Idle safeguard: no payload to protect } else if (t.pending_freed && core.pending_reg_task_id == AICPU_TASK_INVALID) { - // Case 4 (running ACK) or Case 2 (pending ACK): clear pending_occupied only - // when no pending task is currently held. Otherwise pending slot is occupied - // by a pre-loaded task and must stay protected. tracker.clear_pending_occupied(bit_pos); } // 4. Progress signal (only when running task completes) - if (t.running_done) { - made_progress = true; - } + if (t.running_done) made_progress = true; } } -// ============================================================================= -// sync_start drain protocol -// ============================================================================= - -// Take ownership of slot_state and signal all threads to enter drain mode. -// Returns true if this thread won the CAS and owns the drain slot. -// Returns false if another thread already holds drain; caller must re-push slot_state. -// -// Two-phase protocol: CAS 0 -> -1 (sentinel) to claim ownership, store task and -// reset election flag, then release-store block_num. Other threads acquire-load -// sync_start_pending; seeing block_num > 0 ensures all relaxed stores are visible. bool SchedulerContext::enter_drain_mode(PTO2TaskSlotState *slot_state, int32_t block_num) { int32_t expected = 0; if (!drain_state_.sync_start_pending.compare_exchange_strong( expected, -1, std::memory_order_relaxed, std::memory_order_relaxed - )) { + )) return false; // Another thread already holds the drain slot. - } - // We own the drain slot. Store the task and reset the coordination flags before making - // it visible. + // We own the drain slot. Store the task and reset election flag before making it visible. drain_state_.pending_task.store(slot_state, std::memory_order_release); drain_state_.drain_ack_mask.store(0, std::memory_order_relaxed); drain_state_.drain_worker_elected.store(0, std::memory_order_relaxed); - drain_state_.drain_stage_go.store(0, std::memory_order_relaxed); - drain_state_.drain_stage_done_mask.store(0, std::memory_order_relaxed); - drain_state_.drain_running_staged.store(0, std::memory_order_relaxed); // Release store: all stores above are now visible to any thread that // acquire-loads sync_start_pending and sees block_num > 0. drain_state_.sync_start_pending.store(block_num, std::memory_order_release); return true; } -// Count total available resources across all scheduler threads for a given shape. -// include_pending adds each thread's pending-capable cores/clusters — used by the -// gated (early) sync_start drain, which pre-stages onto idle running slots AND onto -// busy cores' pending slots. The ready drain (include_pending=false) counts idle only. -int32_t SchedulerContext::count_global_available(PTO2ResourceShape shape, uint8_t core_mask, bool include_pending) { +int32_t SchedulerContext::count_global_available(PTO2ResourceShape shape) { int32_t total = 0; - for (int32_t t = 0; t < active_sched_threads_; t++) { - if (shape == PTO2ResourceShape::MIX) { - // Gated MIX uses split placement (each core to running-if-idle / pending-if-busy), - // so a cluster is available iff every used core has some free slot. The ready - // path (include_pending=false) still needs whole-cluster idle placement. - total += include_pending ? core_trackers_[t].count_mix_split_clusters(core_mask) : - core_trackers_[t].count_mix_running_clusters(core_mask); - } else { - total += core_trackers_[t].get_idle_core_offset_states(shape).count(); - if (include_pending) { - total += core_trackers_[t].get_pending_core_offset_states(shape).count(); - } - } - } + for (int32_t t = 0; t < active_sched_threads_; t++) + total += core_trackers_[t].get_idle_core_offset_states(shape).count(); return total; } -// One thread's share of the drain staging: CAS-claim block indices and publish them onto -// THIS thread's own cores, concurrently with peers. Returns the number of cores placed on a -// RUNNING slot (the rendezvous seed contribution). Each thread touches only its own tracker -// and its own cores' doorbell-table entries; the CAS on next_block_idx and the fetch_or into -// staged_core_mask are the only cross-thread points. -// -// A gated (early) sync_start drain pre-stages every block behind its doorbell -// (prepare_block_for_dispatch is force-gated for the claimed drain range) and defers the launch -// to the rendezvous: idle cores take a gated RUNNING slot; busy cores take a gated PENDING -// slot (promoted by Case 3.3 as those cores' running tasks FIN). A non-gated (ready) drain -// leaves early_dispatch_state==NONE, so every block launches immediately on an idle running slot and -// the pending pass is skipped. For MIX, a gated block uses SPLIT placement (each core -// independently: idle->running, busy->pending) — safe only because gated. -int32_t -SchedulerContext::drain_stage_cores(PTO2TaskSlotState *slot_state, int32_t block_num, int32_t thread_idx, bool gated) { - CoreTracker &tracker = core_trackers_[thread_idx]; +void SchedulerContext::drain_worker_dispatch(int32_t block_num) { + PTO2TaskSlotState *slot_state = drain_state_.pending_task.load(std::memory_order_acquire); + if (!slot_state) { + drain_state_.sync_start_pending.store(0, std::memory_order_release); + return; + } PTO2ResourceShape shape = slot_state->active_mask.to_shape(); - uint8_t core_mask = slot_state->active_mask.core_mask(); - bool mix_split = gated && shape == PTO2ResourceShape::MIX; - int32_t running_staged = 0; - - // Stage from this thread's `valid` cores/clusters: CAS-claim a block-index range sized to - // what this thread can place (against peers claiming concurrently), then publish those - // blocks onto valid cores. prepare_block_for_dispatch decides each MIX core's slot per-core - // (idle -> running, busy -> pending when to_pending); a MIX cluster's idle cores are the - // running-slot cores, counted BEFORE staging mutates the tracker (rendezvous seed). - auto stage = [&](CoreTracker::BitStates valid, bool to_pending) { - while (valid.has_value()) { - int32_t avail = valid.count(); - int32_t start = 0; - int32_t claim = slot_state->claim_block_range(block_num, avail, start); - if (claim == 0) return; -#if SIMPLER_DFX - bool sub_prof = l2_swimlane_level_ >= L2SwimlaneLevel::SCHED_PHASES; - uint64_t prep_t0 = sub_prof ? get_sys_cnt_aicpu() : 0; -#endif - PublishHandle handles[CoreTracker::MAX_CLUSTERS * 3]; - int handle_count = 0; - int32_t claimed[CoreTracker::MAX_CLUSTERS * 3]; - for (int32_t b = 0; b < claim; b++) - claimed[b] = valid.pop_first(); - bool is_mix = (shape == PTO2ResourceShape::MIX); - if (claim > 0) prefetch_block_dst(thread_idx, claimed[0], is_mix); - for (int32_t b = 0; b < claim; b++) { - if (b + 1 < claim) prefetch_block_dst(thread_idx, claimed[b + 1], is_mix); - auto core_offset = claimed[b]; - if (shape == PTO2ResourceShape::MIX) { - running_staged += tracker.mix_cluster_idle_core_count(core_offset, core_mask); - } - handle_count += prepare_block_for_dispatch( - thread_idx, core_offset, *slot_state, shape, to_pending, start + b, &handles[handle_count], gated - ); - } - wmb(); - uint64_t dispatch_ts = 0; -#if SIMPLER_DFX - uint64_t pub_t0 = 0; - if (sub_prof) { - pub_t0 = get_sys_cnt_aicpu(); - // DrainPrepare bar: cluster scan happened before this lambda, so this covers the - // build_payload work for `claim` blocks (handle_count subtasks). - l2_swimlane_aicpu_record_sched_phase( - thread_idx, L2SwimlaneSchedPhaseKind::DrainPrepare, prep_t0, pub_t0, - sched_l2_swimlane_[thread_idx].sched_loop_count, static_cast(handle_count) - ); - } - if (l2_swimlane_level_ >= L2SwimlaneLevel::AICPU_TIMING) { - dispatch_ts = pub_t0 != 0 ? pub_t0 : get_sys_cnt_aicpu(); - } -#endif - // Accumulate this batch's gated cores into a LOCAL mask and OR it into the shared - // staged_core_mask ONCE below, instead of a seq_cst fetch_or per subtask — that - // per-write atomic contends across all drain threads on the same 2 words and was - // ~half the publish cost. The doorbell-table writes stay per-core (unique cid, no - // contention). - uint64_t my_mask[PTO2_EARLY_DISPATCH_CORE_MASK_WORDS] = {0}; - for (int i = 0; i < handle_count; i++) { - publish_subtask_to_core(handles[i], dispatch_ts, thread_idx); - if (gated) { - int32_t cid = tracker.get_core_id_by_offset(handles[i].core_offset); - sched_->early_dispatch_doorbell_table[cid].addr = handles[i].reg_addr; - sched_->early_dispatch_doorbell_table[cid].token = handles[i].reg_task_id; - my_mask[cid >> 6] |= 1ULL << (cid & 63); - } - } - if (gated) { - for (int w = 0; w < PTO2_EARLY_DISPATCH_CORE_MASK_WORDS; w++) { - if (my_mask[w] != 0) { - slot_state->payload->staged_core_mask[w].fetch_or(my_mask[w], std::memory_order_seq_cst); - } - } - } -#if SIMPLER_DFX - if (sub_prof) { - // DrainPublish bar: the MMIO write_reg per subtask (+ gated doorbell/mask record). - l2_swimlane_aicpu_record_sched_phase( - thread_idx, L2SwimlaneSchedPhaseKind::DrainPublish, pub_t0, get_sys_cnt_aicpu(), - sched_l2_swimlane_[thread_idx].sched_loop_count, static_cast(handle_count) - ); - } -#endif - sched_->record_published_blocks(*slot_state, claim); - // AIC/AIV running placement (whole block on idle cores); MIX running cores are - // counted per-cluster above (mix_cluster_idle_core_count). - if (gated && shape != PTO2ResourceShape::MIX && !to_pending) running_staged += handle_count; - } - }; - - if (mix_split) { - // Gated MIX: to_pending=true opts every BUSY used core into its pending slot while idle - // used cores take running slots (prepare_block_for_dispatch: to_pending && !is_core_idle). - stage(tracker.get_mix_split_cluster_offset_states(core_mask), /*to_pending=*/true); - } else { - auto idle = (shape == PTO2ResourceShape::MIX) ? tracker.get_mix_running_cluster_offset_states(core_mask) : - tracker.get_idle_core_offset_states(shape); - stage(idle, /*to_pending=*/false); // idle -> running (ready launch + gated pre-stage) - if (gated) { - stage(tracker.get_pending_core_offset_states(shape), /*to_pending=*/true); + + for (int32_t t = 0; t < active_sched_threads_ && slot_state->next_block_idx < block_num; t++) { + auto valid = core_trackers_[t].get_idle_core_offset_states(shape); + int32_t remaining = slot_state->logical_block_num - slot_state->next_block_idx; + int32_t claim = std::min(valid.count(), remaining); + int32_t start = slot_state->next_block_idx; + slot_state->next_block_idx += claim; + PublishHandle handles[CoreTracker::MAX_CLUSTERS * 3]; + int handle_count = 0; + for (int32_t b = 0; b < claim; b++) { + auto core_offset = valid.pop_first(); + handle_count += prepare_block_for_dispatch( + t, core_offset, *slot_state, shape, false, start + b, &handles[handle_count] + ); } + wmb(); + uint64_t dispatch_ts = 0; + for (int i = 0; i < handle_count; i++) + publish_subtask_to_core(handles[i], dispatch_ts, t); } - return running_staged; + + std::atomic_thread_fence(std::memory_order_release); + drain_state_.pending_task.store(nullptr, std::memory_order_release); + drain_state_.drain_ack_mask.store(0, std::memory_order_relaxed); + drain_state_.drain_worker_elected.store(0, std::memory_order_relaxed); + drain_state_.sync_start_pending.store(0, std::memory_order_release); } -// Called by each scheduler thread when drain_state_.sync_start_pending != 0. -// -// Protocol: -// 1. Ack barrier: all threads signal they've stopped dispatch, spin until all acked. -// If this thread's ack bit gets cleared while waiting, a reset occurred -- return. -// 2. Election + availability: one thread wins the CAS. It checks global resources; if -// insufficient it resets ack/election so all threads resume completion polling to free -// cores, then retry. If sufficient it releases parallel staging (stage_go). -// 3. Parallel stage: EVERY thread stages its OWN cores concurrently (CAS-claimed block -// indices), accumulates its running-slot cores, and marks its stage_done bit. -// 4. Finalize: the elected thread waits for all stage_done bits, seeds the rendezvous -// (running_slot_count) for a gated drain, and reopens the gate -// (a release-store the non-elected threads acquire, so the seed is visible before any -// completion promotes a pending block). Non-elected threads spin until the gate reopens. -void SchedulerContext::handle_drain_mode(int32_t thread_idx, [[maybe_unused]] uint64_t *out_stage_wall_cycles) { -#if SIMPLER_DFX - bool drain_prof = (l2_swimlane_level_ >= L2SwimlaneLevel::SCHED_PHASES && out_stage_wall_cycles != nullptr); - uint64_t drain_acked_ts = 0; // set at ack-barrier end; used to measure the stage wall -#endif - // Every spin in this function honors is_completed(): once the run latches - // completed_ (all tasks done, or a fatal error raised elsewhere), peers leave - // the dispatch loop and stop participating in the drain. A thread parked in a - // drain spin would then wait forever for acks / a gate-open that can no longer - // arrive -- the AICPU watchdog never fires here because these spins live - // outside the dispatch loop's wall-clock budget, so the hang escalates straight - // to the 3 s STARS op-exec timeout (507018) and poisons the device. Bailing on - // completed_ is always safe: any pending sync_start task is either already - // dispatched (a stale re-popped slot) or moot under teardown, and deinit() - // resets drain_state_ before the next run, so leaving it dirty is harmless. +void SchedulerContext::handle_drain_mode(int32_t thread_idx) { // Spin until drain is fully initialized (sentinel -1 -> block_num > 0). int32_t block_num; do { - if (is_completed()) return; block_num = drain_state_.sync_start_pending.load(std::memory_order_acquire); } while (block_num < 0); if (block_num == 0) return; @@ -711,122 +266,44 @@ void SchedulerContext::handle_drain_mode(int32_t thread_idx, [[maybe_unused]] ui // Spin until all threads have acked. // If our bit is cleared while waiting, elected reset due to insufficient resources. while (true) { - if (is_completed()) return; uint32_t ack = drain_state_.drain_ack_mask.load(std::memory_order_acquire); if ((ack & all_acked) == all_acked) break; if ((ack & (1u << thread_idx)) == 0) return; SPIN_WAIT_HINT(); } + // Election -- exactly one thread wins the CAS. int32_t expected = 0; drain_state_.drain_worker_elected.compare_exchange_strong( expected, thread_idx + 1, std::memory_order_acquire, std::memory_order_relaxed ); - bool elected = drain_state_.drain_worker_elected.load(std::memory_order_relaxed) == thread_idx + 1; - - PTO2TaskSlotState *slot_state = drain_state_.pending_task.load(std::memory_order_acquire); - // OWNER is acquired before the drain is published and persists through - // completion, so every staging thread makes the same gate decision even if - // producer release changes early_dispatch_state during the barrier. - bool gated = slot_state != nullptr && slot_state->payload != nullptr && - PTO2SchedulerState::owns_early_sync_drain(*slot_state->payload); - - if (elected) { - if (slot_state == nullptr) { - // pending_task observed null only when a concurrent drain completion already cleared - // it. Stale-elected: release the election lock and return. Do NOT clear drain_ack_mask - // / sync_start_pending -- a *new* drain run may already be accumulating acks. - drain_state_.drain_worker_elected.store(0, std::memory_order_release); - return; - } - PTO2ResourceShape shape = slot_state->active_mask.to_shape(); - // A gated drain may pre-stage onto pending slots too (idle+pending); the ready drain - // needs block_num idle cores/clusters. - int32_t available = - count_global_available(shape, slot_state->active_mask.core_mask(), /*include_pending=*/gated); - if (available < block_num) { - // Insufficient -- reset so all threads resume completion polling to free cores, then retry. - drain_state_.drain_ack_mask.store(0, std::memory_order_release); - drain_state_.drain_worker_elected.store(0, std::memory_order_release); - return; - } - // Release parallel staging: every thread (this one included) now stages its own cores. - drain_state_.drain_running_staged.store(0, std::memory_order_relaxed); - drain_state_.drain_stage_done_mask.store(0, std::memory_order_relaxed); - drain_state_.drain_stage_go.store(1, std::memory_order_release); - } else { - // Non-elected: wait for the go signal, or bail if the elected thread reset (stale / - // insufficient resources). - while (drain_state_.drain_stage_go.load(std::memory_order_acquire) == 0) { - if (is_completed()) return; - if (drain_state_.drain_worker_elected.load(std::memory_order_acquire) == 0) return; - SPIN_WAIT_HINT(); - } - slot_state = drain_state_.pending_task.load(std::memory_order_acquire); - if (slot_state == nullptr) return; - gated = slot_state->payload != nullptr && PTO2SchedulerState::owns_early_sync_drain(*slot_state->payload); - } - // Parallel stage this thread's own cores (CAS-claimed block indices), then mark done. -#if SIMPLER_DFX - if (drain_prof) drain_acked_ts = get_sys_cnt_aicpu(); // pre-stage -#endif - int32_t my_running = drain_stage_cores(slot_state, block_num, thread_idx, gated); -#if SIMPLER_DFX - // out param carries the PURE drain_stage_cores wall (build_payload + MMIO publish of - // this thread's cores), isolating it from availability + stage_go handshake. - if (drain_prof && drain_acked_ts != 0) *out_stage_wall_cycles = get_sys_cnt_aicpu() - drain_acked_ts; -#endif - drain_state_.drain_running_staged.fetch_add(my_running, std::memory_order_acq_rel); - drain_state_.drain_stage_done_mask.fetch_or(1u << thread_idx, std::memory_order_release); - - if (!elected) { - // Non-elected: staging done; wait for the elected thread to reopen the gate. Exiting via - // sync_start_pending==0 (release/acquire) or drain_worker_elected==0 both synchronize - // with the elected's finalize (its release fence sequences the seed before both stores), - // so the running_slot_count seed is visible before this thread resumes completions. + if (drain_state_.drain_worker_elected.load(std::memory_order_relaxed) != thread_idx + 1) { + // Non-elected: spin-wait for drain completion or resource-insufficient reset. while (drain_state_.sync_start_pending.load(std::memory_order_acquire) != 0) { - if (is_completed()) return; if (drain_state_.drain_worker_elected.load(std::memory_order_acquire) == 0) return; SPIN_WAIT_HINT(); } return; } - // Elected: wait for all threads to finish staging, then seed the rendezvous and reopen. - while ((drain_state_.drain_stage_done_mask.load(std::memory_order_acquire) & all_acked) != all_acked) { - if (is_completed()) return; - SPIN_WAIT_HINT(); - } - if (gated) { - // Seed the rendezvous with the running-slot cores staged across all threads; pending - // cores advance it as they promote. maybe_rendezvous_ring (producer release) rings iff - // this already equals popcount(staged_core_mask) — i.e. no pending spill. - slot_state->payload->running_slot_count.store( - static_cast(drain_state_.drain_running_staged.load(std::memory_order_acquire)), - std::memory_order_seq_cst - ); + // Elected: check if global resources are sufficient. + PTO2TaskSlotState *slot_state = drain_state_.pending_task.load(std::memory_order_acquire); + if (slot_state == nullptr) { + drain_state_.drain_worker_elected.store(0, std::memory_order_release); + return; } - // Clear drain state and reopen the gate FIRST, so the other threads resume immediately. - // Release fence sequences the seed + tracker mutations before every clear, so any thread - // that acquire-observes one of them (sync_start_pending==0 / drain_worker_elected==0) sees - // the seed. `slot_state` is a local holding the fa_fused slot (not drain_state_), so it stays - // valid for the propagate below even if a new drain reuses pending_task after reopen. - std::atomic_thread_fence(std::memory_order_release); - drain_state_.pending_task.store(nullptr, std::memory_order_release); - drain_state_.drain_stage_go.store(0, std::memory_order_relaxed); - drain_state_.drain_stage_done_mask.store(0, std::memory_order_relaxed); - drain_state_.drain_ack_mask.store(0, std::memory_order_relaxed); - drain_state_.drain_worker_elected.store(0, std::memory_order_relaxed); - drain_state_.sync_start_pending.store(0, std::memory_order_release); + PTO2ResourceShape shape = slot_state->active_mask.to_shape(); + int32_t available = count_global_available(shape); - // Recheck after publishing the drain seed. The producer-side rendezvous check can race - // ahead of drain completion and fail while running_slot_count is still incomplete. When - // every block landed directly in a running slot, no pending promotion remains to retry it. - if (gated) { - sched_->retry_sync_start_rendezvous_after_drain(*slot_state); - } else { - sched_->propagate_dispatch_fanin(*slot_state); + if (available < block_num) { + // Insufficient resources -- reset drain fields so threads can resume + // completion polling to free running cores, then retry. + drain_state_.drain_ack_mask.store(0, std::memory_order_release); + drain_state_.drain_worker_elected.store(0, std::memory_order_release); + return; } - PTO2SchedulerState::finish_early_sync_drain(*slot_state->payload); + + // Dispatch -- all other threads are spinning, elected thread has exclusive tracker access. + drain_worker_dispatch(block_num); } diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_context.h b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_context.h index a6f316173f..38a6f7eb1b 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_context.h +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_context.h @@ -11,20 +11,34 @@ #ifndef SCHEDULER_CONTEXT_H #define SCHEDULER_CONTEXT_H -#include "aicpu/device_phase_aicpu.h" #include "aicpu/platform_regs.h" +#include "aicpu/l2_swimlane_collector_aicpu.h" #include "common/l2_swimlane_profiling.h" -#include "common/unified_log.h" -#include "scheduler_types.h" +#include "scheduler/scheduler_types.h" #include "scheduler/pto_scheduler.h" #include "aicore_completion_mailbox.h" #include "pto2_dispatch_payload.h" -// These macros are defined in runtime.h, but we cannot include it here -// (it pulls in Handshake which we only forward-declare). Mirror the -// authoritative values so the class layout compiles standalone. +#include +#include +#include "runtime.h" +#include "pto_runtime2.h" +#include "pto_shared_memory.h" +#include "aicpu/device_time.h" +#include "aicpu/device_phase_aicpu.h" +#include "aicpu/pmu_collector_aicpu.h" +#include "aicpu/args_dump_aicpu.h" +#include "common/memory_barrier.h" +#include "common/platform_config.h" +#include "common/unified_log.h" +#include "spin_hint.h" + +#ifndef unlikely +#define unlikely(x) __builtin_expect(!!(x), 0) +#endif + #ifndef RUNTIME_MAX_WORKER #define RUNTIME_MAX_WORKER 72 #endif @@ -37,111 +51,207 @@ class Runtime; struct Handshake; struct PTO2Runtime; -/** - * SchedulerContext: owns all scheduler-side state and methods. - * - * Held as a member of AicpuExecutor (sched_ctx_). The single public entry - * point is resolve_and_dispatch(), called once per scheduler thread. - * - * All dispatch/completion/drain/cold-path logic is implemented as private - * member methods, split across three .cpp files by responsibility: - * - scheduler_completion.cpp (completion polling, drain protocol) - * - scheduler_cold_path.cpp (exit checks, stall diagnostics, profiling) - * - scheduler_dispatch.cpp (task dispatch loop and helpers) - */ class SchedulerContext { public: - // ========================================================================= - // Lifecycle - // ========================================================================= - - // Initialize scheduler state from the given runtime and thread layout. Split - // into three parts so the per-core AICore handshake — a serial, MMIO-bound - // loop that dominates preamble (~217 µs of ~283 µs for 72 cores) — can run in - // parallel across all AICPU threads. Orchestrated by AicpuExecutor::init: - // the leader runs pre_handshake_init, every thread handshakes a disjoint - // slice of cores via handshake_partition, then the leader runs - // post_handshake_init after a barrier. + // Init is split into three parts so the per-core AICore handshake — a + // serial, MMIO-bound loop that dominates preamble (~217 µs of ~283 µs for + // 72 cores) — can run in parallel across all AICPU threads. The leader + // (exec_idx 0) runs pre_handshake_init, then every thread handshakes a + // disjoint slice of cores via handshake_partition, then the leader runs + // post_handshake_init after a barrier. See AicpuExecutor::init. // - // Leader-only: per-core state + config + swimlane buffers + core count. Must - // be published before any thread enters handshake_partition. Returns 0 on - // success, negative on failure. + // Leader-only: per-core state + config + swimlane buffers + core count. + // Must be published before any thread enters handshake_partition. int32_t pre_handshake_init(Runtime *runtime, int32_t aicpu_thread_num, int32_t sched_thread_num, uint64_t regs_base); - // All threads: handshake this thread's contiguous slice [lo, hi) of cores - // (partitioned by tidx/nthreads). Each core is touched by exactly one thread. + + // All threads: handshake this thread's contiguous slice [lo, hi) of cores. + // Each core is touched by exactly one thread (contiguous, gap-free + // partition), so core_exec_states_ writes are race-free. The AIC/AIV + // worker-id lists are built serially in post_handshake_init to preserve the + // core-index ordering that assign_cores_to_threads relies on. void handshake_partition(Runtime *runtime, int32_t tidx, int32_t nthreads); - // Handshake exactly the cores this scheduler thread will later manage: - // clusters {tidx, tidx+active, ...}, cluster ci = - // {ci, N/3+2ci, N/3+2ci+1} (blocked layout: [0,N/3) AIC, [N/3,N) AIV). Matches - // assign_cores_to_threads' round-robin so handshake warms the same - // core_exec_states_ the thread later dispatches from. - void handshake_owned_clusters(Runtime *runtime, int32_t tidx, int32_t active_threads); - // Barrier-free counterpart of assign_cores_to_threads: thread tidx populates - // its own CoreTracker + per-core payload state for the clusters it owns, right - // after handshaking them — no all-thread barrier or leader post_handshake_init. - void assign_own_clusters(int32_t tidx); - // Latch completion + shutdown cores on a handshake failure seen without the - // barrier (non-DFX path). Idempotent. - void abort_and_shutdown(Runtime *runtime); - // Leader-only profiling-subsystem init (DFX builds); called behind a barrier - // in the barrier-free path since pmu_aicpu_init needs all physical_core_ids_. - void post_handshake_profiling_init(); + + // Barrier-free init (multi-thread): a scheduler thread handshakes only + // the cores it will later dispatch to — clusters ci with ci % active_threads == + // tidx, cluster ci = {ci, aic_n+2ci, aic_n+2ci+1} in the blocked layout + // ([0,aic_n) AIC, [aic_n,3*aic_n) AIV). Same per-core handshake as + // handshake_partition, over the owned set instead of a contiguous slice, so + // core_exec_states_ writes stay race-free (each core owned by one thread). + void handshake_owned_clusters(Runtime *runtime, int32_t tidx, int32_t active_threads) { + Handshake *all_handshakes = reinterpret_cast(runtime->dev.workers); + const int32_t aic_n = cores_total_num_ / 3; + + int32_t owned[RUNTIME_MAX_WORKER]; + int32_t own_n = 0; + for (int32_t ci = tidx; ci < aic_n; ci += active_threads) { + owned[own_n++] = ci; // AIC + owned[own_n++] = aic_n + 2 * ci; // AIV0 + owned[own_n++] = aic_n + 2 * ci + 1; // AIV1 + } + + // Batched 4-phase handshake (adopts #1345's batching, keeping polling's + // aicpu_ready release since the AICore waits on it before reporting): tasks + // / windows / states are each published in one pass so posted MMIO STRs and + // GM stores don't serialize, and only ~4 barriers fire for the whole owned + // set instead of ~2 per core. + + // Phase 1: publish every task pointer (one barrier so all are visible before + // any release), then release every owned core (aicpu_ready=1). The AICore + // reads its task only after observing aicpu_ready. + for (int32_t k = 0; k < own_n; k++) + all_handshakes[owned[k]].task = reinterpret_cast(&payload_per_core_[owned[k]][0]); + OUT_OF_ORDER_STORE_BARRIER(); + for (int32_t k = 0; k < own_n; k++) + all_handshakes[owned[k]].aicpu_ready = 1; + OUT_OF_ORDER_STORE_BARRIER(); + + uint32_t max_physical_cores_count = platform_get_physical_cores_count(); + uint64_t *regs = reinterpret_cast(regs_); + + struct ReadyCore { + int32_t i; + uint32_t pcid; + uint64_t reg_addr; + CoreType core_type; + }; + ReadyCore ready[RUNTIME_MAX_WORKER]; + int32_t n_ready = 0; + bool core_serviced[RUNTIME_MAX_WORKER] = {false}; + + // Phase 2: collect every owned core's report (spin until all done), + // prefetching each CoreExecState line for the write pass. + for (int32_t remaining = own_n; remaining > 0;) { + for (int32_t k = 0; k < own_n; k++) { + int32_t i = owned[k]; + if (core_serviced[i]) continue; + Handshake *hank = &all_handshakes[i]; + if (hank->aicore_done == 0) { + SPIN_WAIT_HINT(); + continue; + } + uint32_t physical_core_id = hank->physical_core_id; + if (physical_core_id >= max_physical_cores_count) { + handshake_failed_.store(true, std::memory_order_release); + core_serviced[i] = true; + remaining--; + continue; + } + __builtin_prefetch(&core_exec_states_[i], 1, 3); + ready[n_ready++] = {i, physical_core_id, regs[physical_core_id], hank->core_type}; + core_serviced[i] = true; + remaining--; + } + } + + // Phase 3: open every core's register window, then ONE barrier. + for (int32_t r = 0; r < n_ready; r++) + platform_init_aicore_regs(ready[r].reg_addr); + OUT_OF_ORDER_STORE_BARRIER(); + + // Phase 4: publish each CoreExecState (AICPU-private, may follow the + // windows). running/pending_reg_task_id start INVALID: pre_handshake_init + // memset them to 0, which is a valid task id (AICPU_TASK_INVALID, the idle + // sentinel, is not 0), so without this reset the scheduler reads every core + // as "running task 0", never dispatches, and trips SCHEDULER_TIMEOUT. + for (int32_t r = 0; r < n_ready; r++) { + int32_t i = ready[r].i; + core_exec_states_[i].reg_addr = ready[r].reg_addr; + core_exec_states_[i].cond_ptr = get_reg_ptr(ready[r].reg_addr, RegId::COND); + core_exec_states_[i].worker_id = i; + core_exec_states_[i].physical_core_id = ready[r].pcid; + core_exec_states_[i].core_type = ready[r].core_type; + core_exec_states_[i].running_reg_task_id = AICPU_TASK_INVALID; + core_exec_states_[i].pending_reg_task_id = AICPU_TASK_INVALID; + } + OUT_OF_ORDER_STORE_BARRIER(); + } + + // Barrier-free counterpart of post_handshake_init's assignment: thread tidx + // populates its own CoreTracker + per-owned-core sub_block_id right after + // handshaking its clusters — no all-thread barrier, no leader serialization. + // The blocked layout gives the owned clusters' worker ids directly, so no + // aic_worker_ids_ discovery is needed. AsyncCtx/slab pointers are set per + // dispatch by build_payload (as on the barrier path), so only the tracker and + // the one-time sub_block_id are set here. + void assign_own_clusters(int32_t tidx) { + const int32_t aic_n = cores_total_num_ / 3; + const int32_t active = active_sched_threads_; + + CoreTracker &tracker = core_trackers_[tidx]; + int32_t own_n = 0; + for (int32_t ci = tidx; ci < aic_n; ci += active) + own_n++; + tracker.init(own_n); + + int32_t local = 0; + for (int32_t ci = tidx; ci < aic_n; ci += active) + tracker.set_cluster(local++, ci, aic_n + 2 * ci, aic_n + 2 * ci + 1); + + for (int32_t c = 0; c < tracker.get_cluster_count(); c++) { + int32_t cluster_offset = c * 3; + int32_t aiv0_id = tracker.get_core_id_by_offset(tracker.get_aiv0_core_offset(cluster_offset)); + int32_t aiv1_id = tracker.get_core_id_by_offset(tracker.get_aiv1_core_offset(cluster_offset)); + payload_per_core_[aiv0_id][0].global_context.sub_block_id = 0; + payload_per_core_[aiv0_id][1].global_context.sub_block_id = 0; + payload_per_core_[aiv1_id][0].global_context.sub_block_id = 1; + payload_per_core_[aiv1_id][1].global_context.sub_block_id = 1; + } + } + + // Latch completion + broadcast exit on a handshake failure seen without the + // all-thread barrier (barrier-free path). Idempotent. + void abort_and_shutdown(Runtime *runtime) { + if (!completed_.exchange(true, std::memory_order_acq_rel)) { + emergency_shutdown(runtime); + } + } + bool handshake_failed() const { return handshake_failed_.load(std::memory_order_acquire); } + // Leader-only, after the handshake barrier: build worker-id lists, assign - // cores, init profiling subsystems, read task counts, init payloads. + // cores to threads, read task counts, init dispatch payloads. int32_t post_handshake_init(Runtime *runtime); // Reset all SchedulerContext-owned state to its post-construction defaults. // Called by AicpuExecutor::deinit() during per-run teardown. void deinit(); - // ========================================================================= - // Per-thread execution entry points (called by AicpuExecutor::run) - // ========================================================================= - // Main scheduler thread entry: poll completion + dispatch ready tasks. int32_t resolve_and_dispatch(Runtime *runtime, int32_t thread_idx); - // Shutdown AICore registers for this thread's assigned cores. - // Also runs PMU finalize (SIMPLER_DFX) before deinit when enabled. - // Orchestrator threads (core_trackers_[thread_idx].core_num() == 0) are a no-op. int32_t shutdown(int32_t thread_idx); - // Run all post-orchestration scheduler bookkeeping: - // - publishes core assignments to the perf collector (SIMPLER_DFX) - // - latches submitted task count from PTO2 shared memory - // - folds inline_completed_tasks into completed_tasks_ - // - flips orchestrator_done_ and triggers core transition - // (skipped on fatal error — emergency_shutdown runs instead) - // Callers must invoke rt_orchestration_done(rt) before this — that - // step belongs to the orchestrator lifecycle, not the scheduler. - void on_orchestration_done(Runtime *runtime, PTO2Runtime *rt, int32_t thread_idx, int32_t total_tasks); + // Upstream-compatible overload: signature is (runtime, rt, thread_idx, total_tasks). + // thread_idx is ignored — polling scheduler's bookkeeping is thread-agnostic at + // this point. + void on_orchestration_done(Runtime *runtime, PTO2Runtime *rt, int32_t /*thread_idx*/, int32_t total_tasks); + + void on_orchestration_done(Runtime *runtime, PTO2Runtime *rt, int32_t total_tasks); // Bind the PTO2Runtime scheduler pointer. Required in device-orchestration // mode where rt is created by the orchestrator thread after init(). void bind_runtime(PTO2Runtime *rt); - // Serial orch->sched mode pre-dispatch wait. No AICore dispatch happens - // before orchestrator_done_. - void wait_for_orchestration_done_before_dispatch(Runtime *runtime, int32_t thread_idx); - - // ========================================================================= - // State queries / external synchronization points - // ========================================================================= + // Serial orch->sched mode pre-dispatch gate. Spin until the orchestrator + // marks itself done; thread 0 may drain the wiring SPSC in the meantime + // so the orchestrator's submit_task pushes don't back-pressure. Other + // threads idle on the orchestrator_done_ flag. + void wait_for_orchestration_done_before_dispatch(Runtime * /*runtime*/, int32_t thread_idx); int32_t aic_count() const { return aic_count_; } int32_t aiv_count() const { return aiv_count_; } - int32_t cores_total_num() const { return cores_total_num_; } bool is_completed() const { return completed_.load(std::memory_order_acquire); } int32_t completed_tasks_count() const { return completed_tasks_.load(std::memory_order_acquire); } - bool orchestration_done() const { return orchestrator_done_.load(std::memory_order_relaxed); } -private: - // ========================================================================= - // State - // ========================================================================= + // Block until the first scheduler thread has finished one-time PTO2 init. + // Called by the orchestrator thread in device-orch mode. + void wait_pto2_init_complete() const { + while (!pto2_init_complete_.load(std::memory_order_acquire)) + SPIN_WAIT_HINT(); + } +private: // --- Scheduler binding & per-core runtime state --- alignas(64) PTO2SchedulerState *sched_{nullptr}; PTO2Runtime *rt_{nullptr}; @@ -156,27 +266,17 @@ class SchedulerContext { // buf_idx = reg_task_id & 1; adjacent dispatches alternate automatically. PTO2DispatchPayload payload_per_core_[RUNTIME_MAX_WORKER][2]; - // Per-core deferred-completion software registration storage. This has - // the same runtime lifetime as payload_per_core_, but is kept out of the - // dispatch payload so normal task dispatch layout and cache footprint stay - // unchanged. DeferredCompletionSlab deferred_slab_per_core_[RUNTIME_MAX_WORKER][2]; // sync_start drain coordination SyncStartDrainState drain_state_; -#if SIMPLER_DFX - SchedL2SwimlaneCounters sched_l2_swimlane_[MAX_AICPU_THREADS]; - // Cached once at init() from get_l2_swimlane_level(), AFTER - // l2_swimlane_aicpu_init has promoted the level from the shared-memory header. - L2SwimlaneLevel l2_swimlane_level_{L2SwimlaneLevel::DISABLED}; -#endif - // --- Task-execution tracking --- std::atomic completed_tasks_{0}; int32_t total_tasks_{0}; // Device orchestration: set by last orchestrator when graph is built; schedulers poll it. - std::atomic orchestrator_done_{false}; + // volatile prevents the compiler from hoisting the load out of spin loops. + volatile bool orchestrator_done_{false}; std::atomic completed_{false}; uint64_t *func_id_to_addr_{nullptr}; @@ -192,29 +292,16 @@ class SchedulerContext { int32_t aic_count_{0}; int32_t aiv_count_{0}; - // Compact per-core CoreType, packed contiguously (~2 cache lines total) so - // post_handshake_init's ordered discovery scan reads it instead of taking a - // per-core volatile GM load from the 64B-aligned Handshake struct. Filled by - // each handshake thread for its own [lo,hi) slice during the parallel sweep. - uint8_t core_type_compact_[RUNTIME_MAX_WORKER]{}; - - // Set by any thread whose slice hits an invalid physical_core_id in - // handshake_partition; checked by the leader in post_handshake_init. - std::atomic handshake_failed_{false}; - // Platform AICore-register base array (set by AicpuExecutor before init()). uint64_t regs_{0}; -#if SIMPLER_DFX - // PMU profiling: physical core IDs for PMU MMIO base resolution. - // Separate storage because CoreExecState's 64-byte budget has no room for - // physical_core_id when SIMPLER_DFX=1. - uint32_t physical_core_ids_[RUNTIME_MAX_WORKER]{}; -#endif + // --- One-time init coordination --- + std::atomic pto2_init_done_{false}; + std::atomic pto2_init_complete_{false}; - // ========================================================================= - // Core management (scheduler_cold_path.cpp) - // ========================================================================= + // Set by any thread whose slice hits an invalid physical_core_id in + // handshake_partition; checked by the leader in post_handshake_init. + std::atomic handshake_failed_{false}; // Assign discovered cores (cluster = 1 AIC + 2 AIV) round-robin across scheduler threads. bool assign_cores_to_threads(); @@ -223,15 +310,8 @@ class SchedulerContext { // deinit their AICore register blocks. Idempotent. void emergency_shutdown(Runtime *runtime); - // ========================================================================= - // Dispatch (scheduler_dispatch.cpp) - // ========================================================================= - static const char *shape_name(PTO2ResourceShape shape); - // Lower-case rendering of PTO2SubtaskSlot, used by dispatch and stall logs. - // Kept lower-case to match the `kernels=[aic:N aiv0:N aiv1:N]` field - // convention already established in the stall log family. static inline const char *subslot_name(PTO2SubtaskSlot s) { switch (s) { case PTO2SubtaskSlot::AIC: @@ -245,23 +325,14 @@ class SchedulerContext { } int pop_ready_tasks_batch( - PTO2ReadyQueue *queues, PTO2ResourceShape shape, int32_t thread_idx, PTO2TaskSlotState **out, int max_count + PTO2ResourceShape shape, PTO2LocalReadyBuffer &local_buf, PTO2TaskSlotState **out, int max_count ); void build_payload( PTO2DispatchPayload &dispatch_payload, PTO2TaskSlotState &slot_state, PTO2SubtaskSlot subslot, - int32_t block_idx, bool force_gate + const AsyncCtx &async_ctx, int32_t block_idx ); - // Batched-dispatch primitives. prepare_* builds the payload and per-core - // state; publish_* issues the MMIO register write. Callers must wmb() - // between the prepare batch and the publish batch, then sample - // get_sys_cnt_aicpu() once and pass it to publish_* for every handle. - // - // dispatch_timestamp_slot points to the CoreExecState slot - // (pending_dispatch_timestamp / running_dispatch_timestamp) selected at - // prepare time, or nullptr when L2 swimlane is below AICPU_TIMING and no - // dispatch timestamp is being recorded. struct PublishHandle { uint64_t reg_addr; uint32_t reg_task_id; @@ -270,274 +341,85 @@ class SchedulerContext { int32_t task_timing_slot; // TASK_TIMING_SLOT_NONE unless the task is tagged }; - PublishHandle prepare_subtask_to_core( + SchedulerContext::PublishHandle prepare_subtask_to_core( int32_t thread_idx, int32_t core_offset, PTO2TaskSlotState &slot_state, PTO2SubtaskSlot subslot, - bool to_pending, int32_t block_idx, bool force_gate + bool to_pending, int32_t block_idx ); - // `thread_idx` is the publishing Scheduler thread's index, used to select the - // per-thread task-timing record; every call site already has it in scope. + // `thread_idx` selects the publishing Scheduler thread's per-thread task-timing + // record; every call site already has it in scope. inline void publish_subtask_to_core(const PublishHandle &h, uint64_t dispatch_ts, int32_t thread_idx) { - if (h.dispatch_timestamp_slot != nullptr) { - *h.dispatch_timestamp_slot = dispatch_ts; - } + if (h.dispatch_timestamp_slot != nullptr) *h.dispatch_timestamp_slot = dispatch_ts; // Task-timing dispatch: earliest DATA_MAIN_BASE publication for a tagged - // task, folded as min. Untagged tasks pay only this cache-hot compare and - // never read the sys counter. Independent of L2 swimlane level. - if (h.task_timing_slot != TASK_TIMING_SLOT_NONE) { - aicpu_task_timing_dispatch(h.task_timing_slot, thread_idx); - } + // task, folded as min. Untagged tasks pay only this cache-hot compare. + if (h.task_timing_slot != TASK_TIMING_SLOT_NONE) aicpu_task_timing_dispatch(h.task_timing_slot, thread_idx); write_reg(h.reg_addr, RegId::DATA_MAIN_BASE, static_cast(h.reg_task_id)); } - // Prefetch the cold per-core structures the next block's prepare touches. - // Ordering is load-bearing: issue the STALLING LOAD first — CoreExecState, - // read by dispatch_seq++ whose value feeds reg_task_id -> buf_idx -> the whole - // dispatch — for every core of the block, BEFORE the store-target prefetches. - // MSHRs saturate (a MIX block warms 3 cores); issuing the read prefetches - // first keeps them from being the ones dropped. The dispatch-buffer writes - // still get prefetched (measured to help ~30% on this shallow-store-buffer - // control core), just after the reads. rw=1 on CoreExecState (read AND - // written) gives Exclusive, serving both without a Shared->Exclusive upgrade. - inline void prefetch_block_dst(int32_t thread_idx, int32_t core_offset, bool is_mix) { - CoreTracker &tracker = core_trackers_[thread_idx]; - int32_t cids[3] = {}; - int32_t nc = 0; - if (is_mix) { - cids[nc++] = tracker.get_core_id_by_offset(tracker.get_aic_core_offset(core_offset)); - cids[nc++] = tracker.get_core_id_by_offset(tracker.get_aiv0_core_offset(core_offset)); - cids[nc++] = tracker.get_core_id_by_offset(tracker.get_aiv1_core_offset(core_offset)); - } else { - cids[nc++] = tracker.get_core_id_by_offset(core_offset); - } - // Stalling loads first. - for (int32_t i = 0; i < nc; i++) - __builtin_prefetch(&core_exec_states_[cids[i]], 1, 3); - // Store targets after (dispatch buffer CL0 control + CL1 args, both bufs). - for (int32_t i = 0; i < nc; i++) { - for (int32_t buf = 0; buf < 2; buf++) { - const char *dp = reinterpret_cast(&payload_per_core_[cids[i]][buf]); - __builtin_prefetch(dp, 1, 3); - __builtin_prefetch(dp + 64, 1, 3); - } - } - } - // Fan out one block's subtasks (1 for AIC/AIV, 1-3 for MIX) into the // caller-supplied handles buffer. Returns the number of handles written. int prepare_block_for_dispatch( int32_t thread_idx, int32_t core_offset, PTO2TaskSlotState &slot_state, PTO2ResourceShape shape, - bool to_pending, int32_t block_idx, PublishHandle *out_handles, bool force_gate = false + bool to_pending, int32_t block_idx, PublishHandle *out_handles ); void dispatch_shape( - int32_t thread_idx, PTO2ReadyQueue *disp_queues, PTO2ResourceShape shape, CoreTracker::DispatchPhase phase, - CoreTracker &tracker, bool &entered_drain, bool &made_progress, bool &try_pushed - ); - - // Early-dispatch (Hook 1). Mirrors dispatch_ready_tasks: owns its - // own gating (off-PMU, this thread has a spare slot, and no normal ready work - // is queued) and sets made_progress / try_pushed when it stages, so the caller - // is a single unconditional call like normal dispatch. After normal dispatch - // leaves idle cores spare, pre-stage the consumers of any RUNNING flagged - // producer onto those cores with a non-zero src_payload (gated). Touches no dependency - // state — the task is released by the doorbell at its normal ready-pop (Hook 2). - int32_t try_early_dispatch( - int32_t thread_idx, CoreTracker &tracker, bool pmu_active, bool &made_progress, bool &try_pushed - ); - - // Stage the already-claimed range [start, start+count) of consumer `c` onto - // thread_idx's idle (RUNNING slot) then pending (gated-pending, promote-on-FIN) - // cores from the provided free-core sets. The caller claims next_block_idx and - // re-pushes `c` BEFORE calling, so this expensive prepare+publish runs - // concurrently with peers (mirrors the normal SPMD dispatch path). Returns the - // number of blocks staged. - int32_t stage_consumer_blocks( - int32_t thread_idx, PTO2TaskSlotState *c, PTO2ResourceShape shape, int32_t start, int32_t count, - CoreTracker::BitStates &idle, CoreTracker::BitStates &pend + int32_t thread_idx, PTO2ResourceShape shape, CoreTracker::DispatchPhase phase, PTO2LocalReadyBuffer &local_buf, + CoreTracker &tracker, bool &entered_drain, bool &made_progress ); - // Early-dispatch analog of dispatch_shape: drain early_dispatch_queues[shape] and - // pre-stage claimed block ranges onto this thread's free cores of `shape` for the - // given phase (IDLE -> onto idle cores in the RUNNING slot; PENDING -> onto a - // running core's gated pending slot). Pop is sized to the shape's capacity exactly - // as dispatch_shape sizes normal dispatch. Returns the number of blocks staged. - int32_t early_dispatch_shape(int32_t thread_idx, PTO2ResourceShape shape, CoreTracker::DispatchPhase phase); - - // One pass of "Phase 4" in the resolve_and_dispatch loop: IDLE-stage dispatch - // for MIX then (if no mix residual) AIC/AIV; then PENDING-stage dispatch with - // cross-thread idle gating. MIX is strictly prioritized — when mix residual is - // detected after MIX-IDLE, AIC/AIV are skipped for the whole pass but - // MIX-PENDING still runs. - // - // Forward-progress argument for AIC/AIV: skip_aic_aiv is sticky for the - // current pass only. The next loop iteration re-evaluates after Phase 1 - // completion polling and the global MIX queue draining (here or on any - // peer thread). AIC/AIV starvation is therefore bounded by MIX throughput, - // not unbounded — once mix completes on at least one cluster, the next - // pass either drains the residual or admits AIC/AIV. void dispatch_ready_tasks( - int32_t thread_idx, CoreTracker &tracker, bool pmu_active, bool &made_progress, bool &try_pushed + int32_t thread_idx, CoreTracker &tracker, PTO2LocalReadyBuffer (&local_bufs)[PTO2_NUM_RESOURCE_SHAPES], + bool pmu_active, bool &made_progress ); - // Shared staging order for both dispatch sources (normal ready + speculative early): - // MIX strict priority, IDLE stage before PENDING stage, cross-thread idle gating - // (MIX-IDLE ▶ c/v-IDLE ▶ MIX-PEND ▶ c/v-PEND). `stage(shape, phase)` stages that - // shape+phase bucket for the source and returns true to STOP the pass (normal returns - // true when it enters drain mode; early always returns false). `residual_mix()` reports - // whether MIX work remains queued for the source (normal reads ready_queues[MIX], early - // reads early_dispatch_queues[MIX]). IDLE runs under PMU; PENDING is withheld under PMU. - template - void run_staging_order(int32_t thread_idx, bool pmu_active, StageFn &&stage, ResidualMixFn &&residual_mix); - - // Returns true if any *other* scheduler thread currently has an idle core - // matching `shape`. Used as a scheduling hint on the PENDING dispatch path - // — see the implementation in scheduler_dispatch.cpp for the hint-semantics - // rationale and the safety argument against the drain worker. bool has_idle_in_other_threads(int32_t self_thread_idx, PTO2ResourceShape shape) const; - // True if mix tasks remain in the global MIX ready queue. Approximate — - // PTO2ReadyQueue::size() (see pto_scheduler.h) snapshots its enqueue/dequeue - // positions with std::memory_order_relaxed and may interleave with concurrent - // push/pop. A stale read here causes at most one - // extra/missed AIC/AIV skip and self-corrects on the next loop iteration. - bool has_residual_mix() const { - return sched_->ready_queues[static_cast(PTO2ResourceShape::MIX)].size() > 0; - } - - // Tier-0 analog of has_residual_mix for the ready sync_start lane: true if MIX - // sync_start cohorts remain queued, so the Tier-0 pass keeps MIX strict priority - // over its own AIC/AIV sync work. Same relaxed-size snapshot caveat. - bool has_residual_sync_mix() const { - return sched_->ready_sync_queues[static_cast(PTO2ResourceShape::MIX)].size() > 0; - } - - // Early-dispatch analog of has_residual_mix: true if MIX early-dispatch candidates - // remain queued. has_residual_mix reads the normal MIX ready queue, which is empty - // whenever the Phase-4b early pass runs (it is gated on all ready_queues being - // empty), so early-dispatch MIX priority needs its own residual check against - // early_dispatch_queues[MIX]. Same relaxed-size snapshot caveat as has_residual_mix. - bool has_residual_early_mix() const { - return sched_->early_dispatch_queues[static_cast(PTO2ResourceShape::MIX)].size() > 0; + bool has_residual_mix(const PTO2LocalReadyBuffer &mix_local_buf) const { + return mix_local_buf.count > 0 || sched_->ready_queues[static_cast(PTO2ResourceShape::MIX)].size() > 0; } - // ========================================================================= - // Completion & drain (scheduler_completion.cpp) - // ========================================================================= - - static SlotTransition decide_slot_transition( - int32_t reg_task_id, int32_t reg_state, int32_t running_id, int32_t pending_id, bool pending_gated = false - ); + static SlotTransition + decide_slot_transition(int32_t reg_task_id, int32_t reg_state, int32_t running_id, int32_t pending_id); 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 -#if SIMPLER_DFX - , - uint64_t dispatch_ts, uint64_t finish_ts -#endif + PTO2TaskSlotState &slot_state, int32_t expected_reg_task_id, int32_t core_id, int32_t &completed_this_turn ); static void promote_pending_to_running(CoreExecState &core); + static void clear_running_slot(CoreExecState &core); 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 + int32_t thread_idx, int32_t &completed_this_turn, int32_t &cur_thread_completed, bool &made_progress ); bool enter_drain_mode(PTO2TaskSlotState *slot_state, int32_t block_num); - int32_t count_global_available(PTO2ResourceShape shape, uint8_t core_mask, bool include_pending = false); - // One thread's share of the drain: CAS-claim block indices and stage them onto THIS - // thread's own cores (parallel with peers), returning the number of running-slot cores - // staged (the rendezvous seed contribution). - int32_t drain_stage_cores(PTO2TaskSlotState *slot_state, int32_t block_num, int32_t thread_idx, bool gated); - // out_stage_wall_cycles (profiling only): cycles this thread spent in drain_stage_cores - // (prepare + publish), set ONLY on threads that actually staged. Lets the caller isolate - // the pure stage wall from the ack-barrier + finalize spans in the Drain bar. - void handle_drain_mode(int32_t thread_idx, uint64_t *out_stage_wall_cycles = nullptr); - - // ========================================================================= - // Cold path: exit checks, stall diagnostics, profiling (scheduler_cold_path.cpp) - // ========================================================================= - - __attribute__((noinline, cold)) LoopAction - handle_orchestrator_exit(int32_t thread_idx, PTO2SharedMemoryHeader *header, Runtime *runtime, int32_t &task_count); - - __attribute__((noinline, cold)) LoopAction - check_idle_fatal_error(int32_t thread_idx, PTO2SharedMemoryHeader *header, Runtime *runtime); - - __attribute__((noinline, cold)) void - log_stall_diagnostics(int32_t thread_idx, int32_t task_count, int32_t idle_iterations, int32_t last_progress_count); - - __attribute__((noinline, cold)) void log_shutdown_stall_snapshot( - int32_t trigger_thread_idx, int32_t trigger_idle_iterations, int32_t trigger_last_progress_count - ); - // Reverse lookup: given a global core_id, find which scheduler thread's - // tracker owns it. Returns -1 if not found. Linear scan — only used on - // the cold diagnostic path. - int32_t find_core_owner_thread(int32_t core_id) const; + int32_t count_global_available(PTO2ResourceShape shape); - // Does this thread own any core with a RUNNING task (running_slot_state set)? - // Gates the scheduler timeout fatal latch: a thread without an owned - // RUNNING task has no first-hand evidence of a stuck dispatch and must - // not declare global fatal on its own idle observation. The thread that - // does own the stuck task will reach the budget on its own polls and - // latch with valid evidence (or recover when the COND register flips). - bool self_owns_running_task(int32_t thread_idx) const; + void drain_worker_dispatch(int32_t block_num); - // Does *any* scheduler thread own a RUNNING task? Used as the second - // fatal-latch condition: if the wall-clock budget elapsed AND no thread - // owns RUNNING work AND tasks remain incomplete, the system is in a - // pre-dispatch / WAIT-only deadlock (e.g. dependency cycle) and the - // ownerless idle threads are the only observers — let one of them latch. - bool no_thread_owns_running_task() const; + void handle_drain_mode(int32_t thread_idx); - // One-glance classification of a no-progress timeout, derived from state the - // scheduler already holds at the stall. Reduces the multi-state snapshot to a - // dominant PTO2_STALL_DETAIL_* sub-class plus a few locator fields, which - // handle_timeout_exit propagates to host alongside the unchanged code 100. - struct StallClassification { - int32_t detail; // PTO2_STALL_DETAIL_* - int32_t cnt_running; // tasks observed RUNNING (on a core) - int32_t cnt_ready; // fanin-satisfied but not dispatched - int32_t cnt_waiting; // still waiting on fanin - int32_t completed; // completed_tasks_ snapshot - int32_t total; // total_tasks_ snapshot - int32_t orch_done; // orchestrator_done flag (0/1) - int64_t stuck_task_id; // S1: first RUNNING task's id (-1 if none) - int32_t stuck_core; // S1: core hosting it (-1 if none) - }; + LoopAction handle_orchestrator_exit(PTO2SharedMemoryHeader *header, Runtime *runtime); - // Scan the rings once (same ground truth as log_stall_diagnostics: a slot is - // RUNNING iff a core holds it as running_slot_state) and reduce to a - // StallClassification. Pure reads — safe to call from any scheduler thread. - __attribute__((noinline, cold)) StallClassification classify_stall_reason() const; - - __attribute__((noinline, cold)) int32_t handle_timeout_exit( - int32_t thread_idx, PTO2SharedMemoryHeader *header, Runtime *runtime, int32_t idle_iterations, - int32_t last_progress_count -#if SIMPLER_DFX - , - uint64_t sched_start_ts -#endif - ); + LoopAction check_idle_fatal_error(PTO2SharedMemoryHeader *header, Runtime *runtime); -#if SIMPLER_DFX - __attribute__((noinline, cold)) void log_l2_swimlane_summary(int32_t thread_idx, int32_t cur_thread_completed); -#endif + void log_stall_diagnostics(int32_t thread_idx); + + void log_shutdown_stall_snapshot(); + + int32_t find_core_owner_thread(int32_t core_id) const; - // ========================================================================= - // Small inline helpers - // ========================================================================= + bool self_owns_running_task(int32_t thread_idx) const; + + bool no_thread_owns_running_task() const; + + int32_t handle_timeout_exit(int32_t thread_idx, PTO2SharedMemoryHeader *header, Runtime *runtime); uint64_t get_function_bin_addr(int func_id) const { - if (!func_id_to_addr_ || func_id < 0 || func_id >= RUNTIME_MAX_FUNC_ID) { - LOG_ERROR("func_id=%d is out of range [0, %d) or map is null", func_id, RUNTIME_MAX_FUNC_ID); - return 0; - } + if (!func_id_to_addr_ || func_id < 0 || func_id >= RUNTIME_MAX_FUNC_ID) return 0; return func_id_to_addr_[func_id]; } }; diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_dispatch.cpp b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_dispatch.cpp index 4ae92944d0..5a1a6cfb24 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_dispatch.cpp +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_dispatch.cpp @@ -8,59 +8,154 @@ * See LICENSE in the root of the software repository for the full text of the License. * ----------------------------------------------------------------------------------------------------------- */ + #include "scheduler_context.h" #include #include #include -#include "common.h" // debug_assert - -#include "common/unified_log.h" -#include "aicpu/aicpu_device_config.h" -#include "aicpu/device_time.h" -#include "aicpu/platform_regs.h" +#include "common.h" #include "callable.h" -#include "common/l2_swimlane_profiling.h" -#include "common/memory_barrier.h" -#include "common/platform_config.h" -#include "pto_runtime2.h" -#include "runtime.h" -#include "spin_hint.h" - -// Performance profiling headers -#include "aicpu/l2_swimlane_collector_aicpu.h" -#include "aicpu/pmu_collector_aicpu.h" -#include "aicpu/args_dump_aicpu.h" - -#ifndef unlikely -#define unlikely(x) __builtin_expect(!!(x), 0) +#include "aicpu/aicpu_device_config.h" +#include "aicpu/dep_gen_collector_aicpu.h" + +int32_t SchedulerContext::resolve_and_dispatch(Runtime *runtime, int32_t thread_idx) { + always_assert(sched_ != nullptr); + CoreTracker &tracker = core_trackers_[thread_idx]; + + PTO2SharedMemoryHeader *header = sched_->sm_header; + if (!header) return -1; + + // One-time init: assign perf buffers (one thread does it; others wait) + if (!pto2_init_done_.exchange(true, std::memory_order_acq_rel)) + pto2_init_complete_.store(true, std::memory_order_release); + else + while (!pto2_init_complete_.load(std::memory_order_acquire)) + SPIN_WAIT_HINT(); + + int32_t cur_thread_completed = 0; + int32_t idle_iterations = 0; + + constexpr int LOCAL_READY_CAP_PER_TYPE = 64; + PTO2TaskSlotState *local_ptrs[PTO2_NUM_RESOURCE_SHAPES][LOCAL_READY_CAP_PER_TYPE]; + PTO2LocalReadyBuffer local_bufs[PTO2_NUM_RESOURCE_SHAPES]; + for (int32_t i = 0; i < PTO2_NUM_RESOURCE_SHAPES; i++) + local_bufs[i].reset(local_ptrs[i], LOCAL_READY_CAP_PER_TYPE); + + const bool pmu_active = is_pmu_enabled(); + + uint64_t last_progress_ts = get_sys_cnt_aicpu(); + + // Dispatch-loop start timestamp for the SchedWindow phase marker (the + // host reduces min(start)/max(end) across sched threads → the `Sched` + // span). This one call ≈ one kernel launch. + [[maybe_unused]] const uint64_t sched_loop_start_ts = get_sys_cnt_aicpu(); + + while (true) { + if (completed_.load(std::memory_order_acquire)) break; + bool made_progress = false; + if (!tracker.has_any_running_cores()) { + LoopAction action = handle_orchestrator_exit(header, runtime); + if (action == LoopAction::BREAK_LOOP) break; + } + + // Phase 1: Check running cores for completion + int32_t completed_this_turn = 0; + + if (tracker.has_any_running_cores()) { + check_running_cores_for_completion(thread_idx, completed_this_turn, cur_thread_completed, made_progress); + } + if (completed_this_turn > 0) { + completed_tasks_.fetch_add(completed_this_turn, std::memory_order_relaxed); + } + + if (rt_ != nullptr && rt_->aicore_mailbox != nullptr && + (sched_->async_wait_list.count > 0 || rt_->aicore_mailbox->has_pending())) { + AsyncPollResult poll_result = sched_->async_wait_list.poll_and_complete(rt_->aicore_mailbox, sched_); + if (poll_result.error_code != PTO2_ERROR_NONE) { + int32_t expected = PTO2_ERROR_NONE; + header->sched_error_code.compare_exchange_strong( + expected, poll_result.error_code, std::memory_order_acq_rel, std::memory_order_acquire + ); + completed_.store(true, std::memory_order_release); + break; + } + if (poll_result.completed > 0) { + completed_tasks_.fetch_add(poll_result.completed, std::memory_order_relaxed); + made_progress = true; + } + } + + // Phase 2 drain check + if (drain_state_.sync_start_pending.load(std::memory_order_acquire) != 0) { + handle_drain_mode(thread_idx); + continue; + } + + // Phase 3: Drain wiring queue (thread 0 only). + if (thread_idx == 0) { + int wired = sched_->drain_wiring_queue(orchestrator_done_); + if (wired > 0) made_progress = true; + } + + if (thread_idx == 0) { + constexpr int DUMMY_DRAIN_BATCH = 16; + PTO2TaskSlotState *dummy_batch[DUMMY_DRAIN_BATCH]; + int dummy_got = sched_->dummy_ready_queue.pop_batch(dummy_batch, DUMMY_DRAIN_BATCH); + for (int di = 0; di < dummy_got; di++) { + PTO2TaskSlotState &dummy_slot = *dummy_batch[di]; + sched_->on_mixed_task_complete(dummy_slot); + completed_tasks_.fetch_add(1, std::memory_order_relaxed); + cur_thread_completed++; + } + if (dummy_got > 0) made_progress = true; + } + + // Phase 4: MIX-strict-priority dispatch with phase-split and + // cross-thread idle gating. See dispatch_ready_tasks for the policy. + dispatch_ready_tasks(thread_idx, tracker, local_bufs, pmu_active, made_progress); + + if (made_progress) { + idle_iterations = 0; + last_progress_ts = get_sys_cnt_aicpu(); + } else { + idle_iterations++; + + if (idle_iterations % FATAL_ERROR_CHECK_INTERVAL == 0) { + LoopAction action = check_idle_fatal_error(header, runtime); + if (action == LoopAction::BREAK_LOOP) break; + } + + if (idle_iterations % STALL_LOG_INTERVAL == 0) log_stall_diagnostics(thread_idx); + if (get_sys_cnt_aicpu() - last_progress_ts > SCHEDULER_TIMEOUT_CYCLES) { + bool self_owns = self_owns_running_task(thread_idx); + bool global_stuck = !self_owns && total_tasks_ > 0 && + completed_tasks_.load(std::memory_order_relaxed) < total_tasks_ && + no_thread_owns_running_task(); + if (self_owns || global_stuck) return handle_timeout_exit(thread_idx, header, runtime); + last_progress_ts = get_sys_cnt_aicpu(); + } + SPIN_WAIT_HINT(); + } + } + +#if PTO2_PROFILING + // Ride this scheduler thread's dispatch window home to the per-thread + // phase buffer. The host reduces min(start)/max(end) across the sched + // threads into the `Sched` marker, so Effective = + // max(orch_end, sched_end) - min(orch_start, sched_start) + // ends when the LAST scheduler finishes all tasks — not when the + // orchestrator finished submitting. Without this the sched span is + // absent and Effective collapses to the orch-submit window. sched_end_ts + // is the loop-exit time (completed_ observed = all tasks done). + const uint64_t sched_end_ts = get_sys_cnt_aicpu(); + aicpu_phase_set_window(AicpuPhase::SchedWindow, sched_loop_start_ts, sched_end_ts); #endif -// AICore materializes args[] from src_payload on the gated path using the -// byte offsets in pto2_dispatch_payload.h (the AICore .o cannot see PTO2TaskPayload). -// Pin those constants to the real layout here, where the struct is fully visible. -static_assert(offsetof(PTO2TaskPayload, tensor_count) == PTO2_TASKPAYLOAD_TENSOR_COUNT_OFFSET); -static_assert(offsetof(PTO2TaskPayload, scalar_count) == PTO2_TASKPAYLOAD_SCALAR_COUNT_OFFSET); -static_assert(offsetof(PTO2TaskPayload, tensors) == PTO2_TASKPAYLOAD_TENSORS_OFFSET); -static_assert(offsetof(PTO2TaskPayload, scalars) == PTO2_TASKPAYLOAD_SCALARS_OFFSET); -static_assert(sizeof(Tensor) == PTO2_TASKPAYLOAD_TENSOR_STRIDE); - -// ============================================================================= -// Dispatch helpers -// ============================================================================= - -namespace { -inline constexpr int32_t PTO2_DEFERRED_RELEASE_CAP = 256; + return cur_thread_completed; } -// 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. -static_assert( - RUNTIME_MAX_WORKER <= PTO2_EARLY_DISPATCH_CORE_MASK_WORDS * 64, - "staged_core_mask too small for RUNTIME_MAX_WORKER cores" -); - const char *SchedulerContext::shape_name(PTO2ResourceShape shape) { switch (shape) { case PTO2ResourceShape::AIC: @@ -75,97 +170,43 @@ const char *SchedulerContext::shape_name(PTO2ResourceShape shape) { return "UNKNOWN"; } -bool SchedulerContext::has_idle_in_other_threads(int32_t self_thread_idx, PTO2ResourceShape shape) const { - // Cross-thread read of peer trackers without explicit synchronization. The - // backing `core_states_` is a naturally aligned uint64_t; aarch64 guarantees - // single-copy atomicity for an 8-byte aligned load, so no torn read. The - // value is consumed only as a scheduling *hint* — a stale read at worst - // causes one missed/extra pending dispatch, corrected on the next iteration. - // Drain-mode cross-thread writes are serialized by handle_drain_mode's ack - // barrier (all peers spin out of the dispatch path before any tracker - // mutation), so this routine is never racing the drain worker. - for (int32_t t = 0; t < active_sched_threads_; t++) { - if (t == self_thread_idx) continue; - if (core_trackers_[t].get_idle_core_offset_states(shape).has_value()) { - return true; - } - } - return false; -} - int SchedulerContext::pop_ready_tasks_batch( - PTO2ReadyQueue *queues, PTO2ResourceShape shape, int32_t thread_idx, PTO2TaskSlotState **out, int max_count + PTO2ResourceShape shape, PTO2LocalReadyBuffer &local_buf, PTO2TaskSlotState **out, int max_count ) { -#if SIMPLER_DFX - auto &l2_swimlane = sched_l2_swimlane_[thread_idx]; -#if SIMPLER_SCHED_PROFILING - extern uint64_t g_sched_pop_atomic_count[], g_sched_pop_wait_cycle[]; - uint64_t t_pop_start = get_sys_cnt_aicpu(); - int count = sched_->get_ready_tasks_batch( - queues, shape, out, max_count, g_sched_pop_atomic_count[thread_idx], g_sched_pop_wait_cycle[thread_idx] - ); - l2_swimlane.sched_dispatch_pop_cycle += (get_sys_cnt_aicpu() - t_pop_start); -#else - int count = sched_->get_ready_tasks_batch(queues, shape, out, max_count); -#endif - if (l2_swimlane_level_ >= L2SwimlaneLevel::SCHED_PHASES) { - if (count > 0) { - l2_swimlane.pop_hit += count; - } else { - l2_swimlane.pop_miss++; - } - } -#else - (void)thread_idx; - int count = sched_->get_ready_tasks_batch(queues, shape, out, max_count); -#endif - return count; + return sched_->get_ready_tasks_batch(shape, local_buf, out, max_count); } void SchedulerContext::build_payload( - PTO2DispatchPayload &dispatch_payload, PTO2TaskSlotState &slot_state, PTO2SubtaskSlot subslot, int32_t block_idx, - bool force_gate + PTO2DispatchPayload &dispatch_payload, PTO2TaskSlotState &slot_state, PTO2SubtaskSlot subslot, + const AsyncCtx &async_ctx, int32_t block_idx ) { int32_t slot_idx = static_cast(subslot); uint64_t callable_addr = get_function_bin_addr(slot_state.task->kernel_id[slot_idx]); const CoreCallable *callable = reinterpret_cast(callable_addr); dispatch_payload.function_bin_addr = callable->resolved_addr(); auto &payload = *slot_state.payload; - // A claimed early-stage range stays gated even if producer completion flips - // the shared state before this payload is built. All other dispatches run on - // pickup. - if (PTO2SchedulerState::should_gate_early_dispatch( - force_gate, payload.early_dispatch_state.load(std::memory_order_relaxed) - )) { - // Gated task: hand the idle AICore the source payload (non-zero = gate) and - // let it fill args[0..num_args) itself during its doorbell wait, instead of - // paying the arg-vector write on this scheduler thread. - dispatch_payload.src_payload = reinterpret_cast(&payload); - } else { - // Ready task: fill args here; src_payload = 0 signals AICore to run on pickup. - dispatch_payload.src_payload = 0; - int n = 0; - for (int32_t i = 0; i < payload.tensor_count; i++) { - dispatch_payload.args[n++] = reinterpret_cast(&payload.tensors[i]); - } - for (int32_t i = 0; i < payload.scalar_count; i++) { - dispatch_payload.args[n++] = payload.scalars[i]; - } - } + int n = 0; + for (int32_t i = 0; i < payload.tensor_count; i++) + dispatch_payload.args[n++] = reinterpret_cast(&payload.tensors[i]); + for (int32_t i = 0; i < payload.scalar_count; i++) + dispatch_payload.args[n++] = payload.scalars[i]; dispatch_payload.local_context.block_idx = block_idx; dispatch_payload.local_context.block_num = slot_state.logical_block_num; - // AsyncCtx's slab pointers + capacity are prefilled once per (core, buf_idx) - // in init(); only task_token varies per dispatch. deferred_slab is never null - // on this path, so task_token is always the live task id (matches the old - // AsyncCtx::make(non-null) result). args[PAYLOAD_LOCAL_CONTEXT_INDEX] / - // [PAYLOAD_GLOBAL_CONTEXT_INDEX] are per-(core, buf_idx) constants, also - // prefilled in init(). - dispatch_payload.local_context.async_ctx.task_token = slot_state.task->task_id; + dispatch_payload.local_context.async_ctx = async_ctx; + dispatch_payload.args[PAYLOAD_LOCAL_CONTEXT_INDEX] = reinterpret_cast(&dispatch_payload.local_context); + dispatch_payload.args[PAYLOAD_GLOBAL_CONTEXT_INDEX] = reinterpret_cast(&dispatch_payload.global_context); + // Speculative early-dispatch gate. Polling has no staging path, so every + // dispatch is execute-on-pickup. Writing this per dispatch (as upstream + // does) is what lets deinit skip the ~72 KB payload_per_core_ memset: + // AICore reads src_payload each pickup and a stale non-zero would hang it + // on the doorbell wait. src_payload==0 means "ready" (#1328 folded the old + // not_ready flag into src_payload: 0=ready, non-zero=gated source pointer). + dispatch_payload.src_payload = 0; } SchedulerContext::PublishHandle SchedulerContext::prepare_subtask_to_core( int32_t thread_idx, int32_t core_offset, PTO2TaskSlotState &slot_state, PTO2SubtaskSlot subslot, bool to_pending, - int32_t block_idx, bool force_gate + int32_t block_idx ) { CoreTracker &tracker = core_trackers_[thread_idx]; auto core_id = tracker.get_core_id_by_offset(core_offset); @@ -183,12 +224,11 @@ SchedulerContext::PublishHandle SchedulerContext::prepare_subtask_to_core( uint32_t buf_idx = reg_task_id & 1u; PTO2DispatchPayload &payload = payload_per_core_[core_id][buf_idx]; - // The deferred_slab is NOT reset here: it is init-cleared once and the - // completion path re-clears count only after a task actually recorded a - // deferred completion (count > 0), which is rare. A non-deferred task never - // touches count/error_code, so the slab stays clean without a per-dispatch - // write — keeping this cold per-core line off the dispatch path. - build_payload(payload, slot_state, subslot, block_idx, force_gate); + DeferredCompletionSlab *deferred_slab = &deferred_slab_per_core_[core_id][buf_idx]; + deferred_slab->count = 0; + deferred_slab->error_code = PTO2_ERROR_NONE; + AsyncCtx async_ctx = AsyncCtx::make(slot_state.task->task_id, deferred_slab); + build_payload(payload, slot_state, subslot, async_ctx, block_idx); if (to_pending) { core_exec_state.pending_subslot = subslot; @@ -202,35 +242,7 @@ SchedulerContext::PublishHandle SchedulerContext::prepare_subtask_to_core( } tracker.set_pending_occupied(core_offset); - LOG_DEBUG( - "Thread %d: Dispatched %s %s task %" PRId64 " kernel_id=[%d,%d,%d] block_idx=%d/total_blocks=%d to" - " core_offset=%d core_id=%d reg_task_id=%u", - thread_idx, to_pending ? "pending" : "idle", subslot_name(subslot), - static_cast(slot_state.task->task_id.raw), slot_state.task->kernel_id[0], - slot_state.task->kernel_id[1], slot_state.task->kernel_id[2], block_idx, slot_state.logical_block_num, - core_offset, core_id, reg_task_id - ); - - // AICore buffer rotation lives on the dispatch path: count this dispatch - // and rotate before write_reg when we're about to cross a BUFFER_SIZE - // boundary. The just-filled buffer is not handed to the host here — it is - // released once AICore ACKs this boundary dispatch (see the ACK hook in - // check_running_cores_for_completion), because FIN precedes the swimlane - // record on this runtime. `reg_task_id` is passed as that ACK gate. Gated on - // the same enable bit as flush so level=1 (AICORE_TIMING-only) participates. -#if SIMPLER_DFX - if (l2_swimlane_level_ != L2SwimlaneLevel::DISABLED) { - l2_swimlane_aicpu_on_aicore_dispatch(core_id, thread_idx, reg_task_id); - } -#endif - uint64_t *dispatch_timestamp_slot = nullptr; -#if SIMPLER_DFX - if (l2_swimlane_level_ >= L2SwimlaneLevel::AICPU_TIMING) { - dispatch_timestamp_slot = - to_pending ? &core_exec_state.pending_dispatch_timestamp : &core_exec_state.running_dispatch_timestamp; - } -#endif return PublishHandle{ core_exec_state.reg_addr, reg_task_id, core_offset, dispatch_timestamp_slot, slot_state.task->task_timing_slot @@ -239,115 +251,58 @@ SchedulerContext::PublishHandle SchedulerContext::prepare_subtask_to_core( int SchedulerContext::prepare_block_for_dispatch( int32_t thread_idx, int32_t core_offset, PTO2TaskSlotState &slot_state, PTO2ResourceShape shape, bool to_pending, - int32_t block_idx, PublishHandle *out_handles, bool force_gate + int32_t block_idx, PublishHandle *out_handles ) { -#if SIMPLER_DFX - if (is_dump_args_enabled()) { - dump_args_for_task( - thread_idx, slot_state, ArgsDumpStage::BEFORE_DISPATCH, - [](ActiveMask active_mask, int raw_subtask_id) { - return active_mask.subtask_active(static_cast(raw_subtask_id)); - }, - [this](int32_t func_id) { - return get_function_bin_addr(func_id); - } - ); - } -#endif CoreTracker &tracker = core_trackers_[thread_idx]; if (shape == PTO2ResourceShape::MIX) { uint8_t cmask = slot_state.active_mask.core_mask(); int n = 0; - // Per-core slot placement (#1308): an idle used core takes its running slot - // (tracked by the completion poller), a busy used core takes its gated pending slot - // (promoted on completion). The sync_start drain relies on this — it passes - // to_pending=true so every busy core opts into a pending slot; the non-zero - // src_payload gate keeps the whole cohort waiting for the rendezvous. if (cmask & PTO2_SUBTASK_MASK_AIC) { bool p = to_pending && !tracker.is_aic_core_idle(core_offset); out_handles[n++] = prepare_subtask_to_core( - thread_idx, tracker.get_aic_core_offset(core_offset), slot_state, PTO2SubtaskSlot::AIC, p, block_idx, - force_gate + thread_idx, tracker.get_aic_core_offset(core_offset), slot_state, PTO2SubtaskSlot::AIC, p, block_idx ); } if (cmask & PTO2_SUBTASK_MASK_AIV0) { bool p = to_pending && !tracker.is_aiv0_core_idle(core_offset); out_handles[n++] = prepare_subtask_to_core( - thread_idx, tracker.get_aiv0_core_offset(core_offset), slot_state, PTO2SubtaskSlot::AIV0, p, block_idx, - force_gate + thread_idx, tracker.get_aiv0_core_offset(core_offset), slot_state, PTO2SubtaskSlot::AIV0, p, block_idx ); } if (cmask & PTO2_SUBTASK_MASK_AIV1) { bool p = to_pending && !tracker.is_aiv1_core_idle(core_offset); out_handles[n++] = prepare_subtask_to_core( - thread_idx, tracker.get_aiv1_core_offset(core_offset), slot_state, PTO2SubtaskSlot::AIV1, p, block_idx, - force_gate + thread_idx, tracker.get_aiv1_core_offset(core_offset), slot_state, PTO2SubtaskSlot::AIV1, p, block_idx ); } -#if SIMPLER_DFX - sched_l2_swimlane_[thread_idx].phase_dispatch_count += __builtin_popcount(cmask); -#endif return n; } else if (shape == PTO2ResourceShape::AIC) { - out_handles[0] = prepare_subtask_to_core( - thread_idx, core_offset, slot_state, PTO2SubtaskSlot::AIC, to_pending, block_idx, force_gate - ); -#if SIMPLER_DFX - sched_l2_swimlane_[thread_idx].phase_dispatch_count += 1; -#endif + out_handles[0] = + prepare_subtask_to_core(thread_idx, core_offset, slot_state, PTO2SubtaskSlot::AIC, to_pending, block_idx); return 1; } else { - out_handles[0] = prepare_subtask_to_core( - thread_idx, core_offset, slot_state, PTO2SubtaskSlot::AIV0, to_pending, block_idx, force_gate - ); -#if SIMPLER_DFX - sched_l2_swimlane_[thread_idx].phase_dispatch_count += 1; -#endif + out_handles[0] = + prepare_subtask_to_core(thread_idx, core_offset, slot_state, PTO2SubtaskSlot::AIV0, to_pending, block_idx); return 1; } } void SchedulerContext::dispatch_shape( - int32_t thread_idx, PTO2ReadyQueue *disp_queues, PTO2ResourceShape shape, CoreTracker::DispatchPhase phase, - CoreTracker &tracker, bool &entered_drain, bool &made_progress, bool &try_pushed + int32_t thread_idx, PTO2ResourceShape shape, CoreTracker::DispatchPhase phase, PTO2LocalReadyBuffer &local_buf, + CoreTracker &tracker, bool &entered_drain, bool &made_progress ) { -#if SIMPLER_SCHED_PROFILING - auto &l2_swimlane = sched_l2_swimlane_[thread_idx]; -#endif if (entered_drain) return; bool is_pending = (phase == CoreTracker::DispatchPhase::PENDING); - bool is_mix = (shape == PTO2ResourceShape::MIX); - auto cores = is_mix ? tracker.get_cluster_offset_states() : tracker.get_dispatchable_cores(shape, phase); + auto cores = tracker.get_dispatchable_cores(shape, phase); if (!cores.has_value()) return; while (cores.has_value() && !entered_drain) { int want = cores.count(); PTO2TaskSlotState *batch[CoreTracker::MAX_CLUSTERS * 3]; - int got = pop_ready_tasks_batch(disp_queues, shape, thread_idx, batch, want); + int got = pop_ready_tasks_batch(shape, local_buf, batch, want); if (got == 0) break; - // sync_start exclusion gate. - // - // When the popped batch contains a sync_start task we MUST publish each - // prior task with its own wmb so AICore receives them with time - // separation. The drain coordinator's `count_global_available()` check - // reads the per-thread CoreTracker, and although `prepare_block_for_dispatch` - // marks cores occupied synchronously, the head-start between successive - // tasks is what lets the surrounding completion loop catch up on FINs in - // the retry window when the sync_start task hits insufficient resources. - // Bursting all prior tasks at the end of the pop (cross-task batching) - // collapses that head-start and causes spmd_sync_start_stress to time - // out via 507018 on ~40% of runs — see - // docs/investigations/2026-06-cross-task-batched-publish.md. - // - // When the batch carries no sync_start task, no drain entry can happen - // in this pop, so we hoist `handles[]`, `wmb()`, and the publish loop - // out of the per-task body. One wmb amortizes across all tasks and one - // dispatch_ts is shared, which restores ~60 ns first-to-last AICore - // start span for single-block decode kernels (out_proj, q_proj, ...). - // Detection is a single mask check per task — cheap relative to even - // one register write. bool any_sync_start = false; for (int bi = 0; bi < got; bi++) { if (batch[bi]->active_mask.requires_sync_start()) { @@ -356,78 +311,35 @@ void SchedulerContext::dispatch_shape( } } - // handles[] is sized for the MIX worst case: total claims across the - // pop bounded by `cores.count() ≤ MAX_CLUSTERS`, and each block - // contributes ≤ 3 subtasks for MIX. PublishHandle handles[CoreTracker::MAX_CLUSTERS * 3]; int handle_count = 0; bool dispatched_any = false; - // Logical block ranges published by this pop. AIV can pop two entries per - // cluster, so this ledger has the same worst-case bound as handles[]. - PTO2TaskSlotState *published_list[CoreTracker::MAX_CLUSTERS * 3]; - int16_t published_counts[CoreTracker::MAX_CLUSTERS * 3]; - int published_n = 0; -#if SIMPLER_SCHED_PROFILING - uint64_t t_setup_start = get_sys_cnt_aicpu(); -#endif - // Flush prepared-but-unpublished handles. Required before - // `enter_drain_mode` so the drain coordinator sees cores as occupied, - // and at the per-task boundary when `any_sync_start` is true. auto flush_publish = [&]() { if (handle_count == 0) return; wmb(); uint64_t dispatch_ts = 0; -#if SIMPLER_DFX - if (l2_swimlane_level_ >= L2SwimlaneLevel::AICPU_TIMING) { - dispatch_ts = get_sys_cnt_aicpu(); - } -#endif - for (int i = 0; i < handle_count; i++) { + for (int i = 0; i < handle_count; i++) publish_subtask_to_core(handles[i], dispatch_ts, thread_idx); - } handle_count = 0; made_progress = true; }; for (int bi = 0; bi < got; bi++) { PTO2TaskSlotState *slot_state = batch[bi]; - CoreTracker::BitStates selected_mix_clusters(0ULL); - - if (is_mix) { - auto candidates = cores; - uint8_t cmask = slot_state->active_mask.core_mask(); - auto wanted = is_pending ? CoreTracker::MixPlacement::PENDING : CoreTracker::MixPlacement::RUNNING; - while (candidates.has_value()) { - int32_t cluster_offset = candidates.pop_first(); - if (tracker.classify_mix_cluster(cluster_offset, cmask) == wanted) { - selected_mix_clusters |= CoreTracker::BitStates(1ULL << cluster_offset); - } - } - if (!selected_mix_clusters.has_value()) { - disp_queues[static_cast(shape)].push(slot_state); - continue; - } - } - - // (Early-dispatch pre-staged tasks never reach this ready-pop: they are - // released by their doorbell in release_fanin_and_check_ready the - // instant their last producer completes — see try_early_dispatch_release.) if (slot_state->active_mask.requires_sync_start()) { if (is_pending) { - disp_queues[static_cast(shape)].push(slot_state); + sched_->ready_queues[static_cast(shape)].push(slot_state); continue; } - int32_t available = is_mix ? selected_mix_clusters.count() : cores.count(); + int32_t available = cores.count(); if (available < slot_state->logical_block_num) { flush_publish(); - if (!enter_drain_mode(slot_state, slot_state->logical_block_num)) { - disp_queues[static_cast(shape)].push(slot_state); - } - for (int rem = bi + 1; rem < got; rem++) { - disp_queues[static_cast(shape)].push(batch[rem]); - } + if (!enter_drain_mode(slot_state, slot_state->logical_block_num)) + sched_->ready_queues[static_cast(shape)].push(slot_state); + for (int rem = bi + 1; rem < got; rem++) + sched_->ready_queues[static_cast(shape)].push(batch[rem]); entered_drain = true; break; } @@ -435,117 +347,115 @@ void SchedulerContext::dispatch_shape( if (!cores.has_value()) { flush_publish(); - disp_queues[static_cast(shape)].push_batch(&batch[bi], got - bi); + sched_->ready_queues[static_cast(shape)].push_batch(&batch[bi], got - bi); break; } - // Claim a contiguous range of blocks, hand the slot back to the - // ready queue immediately, then perform the expensive dispatches. - // This lets other schedulers concurrently claim and dispatch the - // remaining blocks of the same SPMD task instead of spinning while - // this thread fills all its own cores. Only local `start + b` is - // read after the push — `next_block_idx` may already be advanced - // by another scheduler that popped the slot. - int32_t available = is_mix ? selected_mix_clusters.count() : cores.count(); - int32_t start = 0; - int32_t claim = slot_state->claim_block_range(slot_state->logical_block_num, available, start); - if (claim == 0) continue; dispatched_any = true; - try_pushed = true; - - published_list[published_n] = slot_state; - published_counts[published_n] = static_cast(claim); - published_n++; + int32_t remaining = slot_state->logical_block_num - slot_state->next_block_idx; + int32_t claim = std::min(cores.count(), remaining); + int32_t start = slot_state->next_block_idx; + slot_state->next_block_idx += claim; - if (start + claim < slot_state->logical_block_num) { - disp_queues[static_cast(shape)].push(slot_state); - } + if (slot_state->next_block_idx < slot_state->logical_block_num) + sched_->ready_queues[static_cast(shape)].push(slot_state); for (int32_t b = 0; b < claim; b++) { - auto core_offset = is_mix ? selected_mix_clusters.pop_first() : cores.pop_first(); - if (is_mix) { - cores.clear_bit(core_offset); - } + auto core_offset = cores.pop_first(); handle_count += prepare_block_for_dispatch( thread_idx, core_offset, *slot_state, shape, is_pending, start + b, &handles[handle_count] ); } - // Sync_start exclusion: flush per task so prior tasks have head- - // start time before any sync_start drain check. Normal batches - // fall through and accumulate for one cross-task flush at the - // end of the pop. - if (any_sync_start) { - flush_publish(); - } + if (any_sync_start) flush_publish(); } flush_publish(); - for (int i = 0; i < published_n; i++) { - sched_->record_published_blocks(*published_list[i], published_counts[i]); - sched_->propagate_dispatch_fanin(*published_list[i]); - } -#if SIMPLER_SCHED_PROFILING - l2_swimlane.sched_dispatch_setup_cycle += (get_sys_cnt_aicpu() - t_setup_start); -#endif if (!dispatched_any) break; - if (!cores.has_value()) { - cores = is_mix ? tracker.get_cluster_offset_states() : tracker.get_dispatchable_cores(shape, phase); - } + if (!cores.has_value()) cores = tracker.get_dispatchable_cores(shape, phase); } } -template -void SchedulerContext::run_staging_order( - int32_t thread_idx, bool pmu_active, StageFn &&stage, ResidualMixFn &&residual_mix +void SchedulerContext::dispatch_ready_tasks( + int32_t thread_idx, CoreTracker &tracker, PTO2LocalReadyBuffer (&local_bufs)[PTO2_NUM_RESOURCE_SHAPES], + bool pmu_active, bool &made_progress ) { using Phase = CoreTracker::DispatchPhase; + constexpr int32_t MIX_I = static_cast(PTO2ResourceShape::MIX); - // MIX is handled explicitly at the top of each stage; only AIC/AIV cycle - // through this 2-elem array, with order toggled by thread parity for - // shape-level load balancing across threads. static constexpr PTO2ResourceShape kAicAivOrder[2][2] = { {PTO2ResourceShape::AIC, PTO2ResourceShape::AIV}, {PTO2ResourceShape::AIV, PTO2ResourceShape::AIC}, }; const PTO2ResourceShape *aic_aiv = kAicAivOrder[thread_idx & 1]; + const int32_t bd_per_thread = PLATFORM_MAX_BLOCKDIM / active_sched_threads_; + const int32_t thread_capacity[PTO2_NUM_RESOURCE_SHAPES] = { + bd_per_thread * PLATFORM_AIC_CORES_PER_BLOCKDIM, + bd_per_thread * PLATFORM_AIV_CORES_PER_BLOCKDIM, + bd_per_thread, + }; + for (int32_t s = 0; s < PTO2_NUM_RESOURCE_SHAPES; s++) { + auto &lb = local_bufs[s]; + int32_t excess = lb.count - thread_capacity[s]; + if (excess <= 0) continue; + if (!has_idle_in_other_threads(thread_idx, static_cast(s))) continue; + sched_->ready_queues[s].push_batch(&lb.slot_states[lb.count - excess], excess); + lb.count -= excess; + } + + auto flush_local_bufs = [&]() { + for (int32_t s = 0; s < PTO2_NUM_RESOURCE_SHAPES; s++) { + auto &lb = local_bufs[s]; + if (lb.count > 0) { + sched_->ready_queues[s].push_batch(lb.slot_states, lb.count); + lb.count = 0; + } + } + }; + struct FlushGuard { + decltype(flush_local_bufs) &flush_fn; + ~FlushGuard() { flush_fn(); } + } flush_guard{flush_local_bufs}; + + bool entered_drain = false; + // ===== IDLE stage ===== - if (stage(PTO2ResourceShape::MIX, Phase::IDLE)) return; + dispatch_shape( + thread_idx, PTO2ResourceShape::MIX, Phase::IDLE, local_bufs[MIX_I], tracker, entered_drain, made_progress + ); + if (entered_drain) return; - // MIX-IDLE residual: AIC/AIV (both IDLE and PENDING) yield for this pass. - // MIX-PENDING below still runs — that is the core of "mix strict priority": - // pending slots are spent on mix before AIC/AIV get any chance. - bool skip_aic_aiv = residual_mix(); + bool skip_aic_aiv = has_residual_mix(local_bufs[MIX_I]); if (!skip_aic_aiv) { for (int i = 0; i < 2; i++) { - if (stage(aic_aiv[i], Phase::IDLE)) return; + PTO2ResourceShape s = aic_aiv[i]; + dispatch_shape( + thread_idx, s, Phase::IDLE, local_bufs[static_cast(s)], tracker, entered_drain, made_progress + ); + if (entered_drain) return; } } + // Flush between IDLE and PENDING so PENDING-stage queue-size checks and any + // peer-thread reads see the IDLE-stage release_fanin output. + flush_local_bufs(); + if (pmu_active) return; - // ===== PENDING stage ===== - // MIX-PENDING gate: skip when a peer has an idle MIX-capable cluster — that - // peer's next IDLE-MIX iteration will pull the mix task from the global - // queue at lower latency than us pre-loading a pending slot here. Forward - // progress for MIX is preserved: at least one thread will run MIX-IDLE next - // pass and consume the residual. - // - // The gate is NOT subject to skip_aic_aiv — residual mix continues to drain - // via pending slots on this thread when no peer is idle. if (!has_idle_in_other_threads(thread_idx, PTO2ResourceShape::MIX)) { - if (stage(PTO2ResourceShape::MIX, Phase::PENDING)) return; + dispatch_shape( + thread_idx, PTO2ResourceShape::MIX, Phase::PENDING, local_bufs[MIX_I], tracker, entered_drain, made_progress + ); + if (entered_drain) return; } // Re-check after MIX-PENDING. If MIX-IDLE already set skip_aic_aiv, leave // it set; otherwise, escalate iff PENDING-MIX left residual. - if (!skip_aic_aiv && residual_mix()) { - skip_aic_aiv = true; - } + if (!skip_aic_aiv && has_residual_mix(local_bufs[MIX_I])) skip_aic_aiv = true; if (skip_aic_aiv) return; @@ -554,964 +464,17 @@ void SchedulerContext::run_staging_order( for (int i = 0; i < 2; i++) { PTO2ResourceShape s = aic_aiv[i]; if (has_idle_in_other_threads(thread_idx, s)) continue; - if (stage(s, Phase::PENDING)) return; - } -} - -void SchedulerContext::dispatch_ready_tasks( - int32_t thread_idx, CoreTracker &tracker, bool pmu_active, bool &made_progress, bool &try_pushed -) { - // Normal ready dispatch (is_ready): dispatch_shape places each block on pickup and - // signals a stop by setting entered_drain when it enters a sync_start drain. - bool entered_drain = false; - - // Tier 0: ready sync_start cohorts take cores before any regular ready task - // (sync_start > MIX > C/V within the normal source). Same order and machinery, - // fed from ready_sync_queues; an oversized cohort arms the stop-the-world drain - // (entered_drain), which also short-circuits the regular tier below. - run_staging_order( - thread_idx, pmu_active, - [&](PTO2ResourceShape shape, CoreTracker::DispatchPhase phase) { - dispatch_shape( - thread_idx, sched_->ready_sync_queues, shape, phase, tracker, entered_drain, made_progress, try_pushed - ); - return entered_drain; - }, - [&] { - return has_residual_sync_mix(); - } - ); - if (entered_drain) return; - - // Tier 1: regular ready work. - run_staging_order( - thread_idx, pmu_active, - [&](PTO2ResourceShape shape, CoreTracker::DispatchPhase phase) { - dispatch_shape( - thread_idx, sched_->ready_queues, shape, phase, tracker, entered_drain, made_progress, try_pushed - ); - return entered_drain; - }, - [&] { - return has_residual_mix(); - } - ); -} - -// Stage the ALREADY-CLAIMED range [start, start+count) of consumer `c` onto -// thread_idx's idle then pending cores. The caller has atomically advanced -// next_block_idx by `count` AND re-pushed `c` for peers -// BEFORE calling this — so this, the expensive prepare+publish, runs CONCURRENTLY -// with peers staging other ranges of the same consumer. This mirrors the normal -// SPMD dispatch path (claim range -> re-push -> dispatch). -// `idle`/`pend` are this thread's free-core sets, sized so idle.count+pend.count >= -// count (the caller clamped the claim to them), so all `count` blocks get a core. -// -// Rule 1: idle cores -> gated task in the RUNNING slot. Rule 2: PENDING slot of -// cores running a real task -> promoted in when that task FINs (gated-pending Case -// 3.3 in decide_slot_transition completes the running FIN + promotes instead of -// waiting for an ack the gated task never sends). Each staged core stays -// pending_occupied while gated, so no second gated block stacks on it. -// -// Doorbell ownership: release flips STAGING->DISPATCHED and exchanges the shared -// mask to claim its bits. A late stager ORs its bits, then fetch-and-clears only -// those release did not take and rings them from its immutable local handles. -// The seq_cst order guarantees every gated core has exactly one writer. -int32_t SchedulerContext::stage_consumer_blocks( - int32_t thread_idx, PTO2TaskSlotState *c, PTO2ResourceShape shape, int32_t start, int32_t count, - CoreTracker::BitStates &idle, CoreTracker::BitStates &pend -) { - CoreTracker &tracker = core_trackers_[thread_idx]; - // Stamp the real pre-stage time (NOT 0) so the swimlane shows these blocks - // dispatched during the producer's run, not at trace start. - uint64_t early_dispatch_ts = get_sys_cnt_aicpu(); - uint64_t my_cores[PTO2_EARLY_DISPATCH_CORE_MASK_WORDS] = {0}; // cores gated by this staging pass - int32_t staged = 0; - int32_t block = start; - // Mirror the normal flush_publish (scheduler_dispatch.cpp wmb()+publish loop): - // prepare ALL claimed blocks' payloads (idle bucket -> running slot, pend bucket - // -> gated pending), then ONE wmb(), then publish. The wmb guarantees the - // src_payload gate + source args are globally visible before any DATA_MAIN_BASE token — - // without it a gated core can pick up the token and dcci a stale payload. The - // shared `count` budget bounds total blocks <= free clusters/cores, so both - // buckets fit one handles[] buffer. - PublishHandle handles[CoreTracker::MAX_CLUSTERS * 3]; - int n = 0; - auto prepare_from = [&](CoreTracker::BitStates &avail, bool to_pending) { - while (count > 0 && avail.has_value()) { - int32_t core_offset = avail.pop_first(); - n += prepare_block_for_dispatch( - thread_idx, core_offset, *c, shape, to_pending, block, &handles[n], /*force_gate=*/true - ); - block++; - count--; - staged++; - } - }; - if (idle.has_value()) prepare_from(idle, /*to_pending=*/false); - if (pend.has_value()) prepare_from(pend, /*to_pending=*/true); - if (n > 0) { - wmb(); - for (int i = 0; i < n; i++) { - publish_subtask_to_core(handles[i], early_dispatch_ts, thread_idx); - int32_t cid = tracker.get_core_id_by_offset(handles[i].core_offset); - sched_->early_dispatch_doorbell_table[cid].addr = handles[i].reg_addr; - sched_->early_dispatch_doorbell_table[cid].token = handles[i].reg_task_id; - my_cores[cid >> 6] |= (1ULL << (cid & 63)); - } - } - // Publish all this thread's gated cores into the shared mask in one OR per word - // (vs one per subtask) so release sees them; seq_cst keeps the self-ring order. - for (int w = 0; w < PTO2_EARLY_DISPATCH_CORE_MASK_WORDS; w++) - if (my_cores[w] != 0) c->payload->staged_core_mask[w].fetch_or(my_cores[w], std::memory_order_seq_cst); - - // Full publication and release are independent events. The seq_cst - // state/launch/count operations form a two-sided handshake. A released - // block must ring before contributing to the publication count. - bool released = staged > 0 && - c->payload->early_dispatch_state.load(std::memory_order_seq_cst) == PTO2_EARLY_DISPATCH_DISPATCHED; - - // Claim only bits the release path did not take. Local handles remain valid - // even if the shared per-core table is reused before this thread resumes. - if (released) { - uint64_t owned[PTO2_EARLY_DISPATCH_CORE_MASK_WORDS] = {0}; - for (int w = 0; w < PTO2_EARLY_DISPATCH_CORE_MASK_WORDS; w++) { - if (my_cores[w] != 0) { - owned[w] = - PTO2SchedulerState::claim_late_staged_doorbell_bits(c->payload->staged_core_mask[w], my_cores[w]); - } - } - for (int i = 0; i < n; i++) { - int32_t cid = tracker.get_core_id_by_offset(handles[i].core_offset); - PTO2SchedulerState::ring_claimed_local_doorbell( - owned[cid >> 6], cid, handles[i].reg_addr, handles[i].reg_task_id - ); - } - wmb(); - } - sched_->record_published_blocks(*c, staged); - // Retry unconditionally after publication. The guards are cheap, and a - // pre-ring state read can become stale if release completes before this - // count update. - sched_->propagate_dispatch_fanin(*c); - return staged; -} - -// Early-dispatch analog of dispatch_shape: drain early_dispatch_queues[shape] and -// pre-stage claimed block ranges onto this thread's `shape` cores for `phase`. IDLE -// stages onto idle cores (RUNNING slot, gated); PENDING stages onto a running core's -// gated pending slot. Producer propagation and late wiring push candidates to the -// shape's queue when dispatch_fanin becomes complete, so the shape is the queue -// index (no per-consumer to_shape()). Returns the number of blocks staged. -int32_t -SchedulerContext::early_dispatch_shape(int32_t thread_idx, PTO2ResourceShape shape, CoreTracker::DispatchPhase phase) { - CoreTracker &tracker = core_trackers_[thread_idx]; - int32_t s = static_cast(shape); - bool is_mix = (shape == PTO2ResourceShape::MIX); - bool is_idle = (phase == CoreTracker::DispatchPhase::IDLE); - - // Size the pop exactly as dispatch_shape does: MIX to the cluster count, else the - // phase's dispatchable-core count. Skip the queue entirely when no core is free - // for this shape+phase (avoids a pointless pop + immediate push-back). - CoreTracker::BitStates cores = - is_mix ? tracker.get_cluster_offset_states() : tracker.get_dispatchable_cores(shape, phase); - if (!cores.has_value()) return 0; - - int32_t total_staged = 0; - PTO2TaskSlotState *batch[CoreTracker::MAX_CLUSTERS * 3]; - uint64_t task_id_snapshots[CoreTracker::MAX_CLUSTERS * 3]; - // Batch-pop in one queue op (fewer CAS than one pop per consumer); the pop is - // bounded by the shape's capacity so the stack buffer always holds it. Then for - // each consumer: CLAIM a range sized to THIS thread's free cores by advancing - // next_block_idx with a CAS (atomic — next_block_idx is shared with normal - // dispatch, which also claims it if release routes the consumer to the ready - // queue, so a plain store could double-dispatch), RE-PUSH it for peers, THEN do - // the expensive prepare+publish. Re-pushing before staging lets peers claim the - // next range and stage CONCURRENTLY — a wide consumer (online_softmax, 48 blocks) - // is filled by all idle threads in parallel. When cores run out mid-batch the - // unprocessed remainder is pushed back for peers (mirrors normal's push_batch of - // the unconsumed tail). - int got = sched_->early_dispatch_queues[s].pop_batch_tagged(batch, task_id_snapshots, cores.count()); - for (int bi = 0; bi < got; bi++) { - PTO2TaskSlotState *c = batch[bi]; - if (static_cast(c->task->task_id.raw) != task_id_snapshots[bi]) continue; - if (c->payload->early_dispatch_state.load(std::memory_order_acquire) != PTO2_EARLY_DISPATCH_STAGING) - continue; // released - - // The single free-core bucket for this phase. For MIX, an active-mask-aware - // whole-cluster scan keeps only the clusters whose placement matches the phase - // (RUNNING placement for IDLE, PENDING placement for PENDING), matching normal - // dispatch's classify_mix_cluster — unused cores in the cluster are ignored, so - // a MIX whose unused AIV is busy is not stranded. For AIC/AIV it is just the - // phase's dispatchable cores. - CoreTracker::BitStates bucket; - if (is_mix) { - auto wanted = is_idle ? CoreTracker::MixPlacement::RUNNING : CoreTracker::MixPlacement::PENDING; - uint8_t cmask = c->active_mask.core_mask(); - CoreTracker::BitStates candidates = tracker.get_cluster_offset_states(); - while (candidates.has_value()) { - int32_t cluster_offset = candidates.pop_first(); - if (tracker.classify_mix_cluster(cluster_offset, cmask) == wanted) { - bucket |= CoreTracker::BitStates(1ULL << cluster_offset); - } - } - } else { - bucket = tracker.get_dispatchable_cores(shape, phase); - } - int32_t freecores = bucket.has_value() ? bucket.count() : 0; - if (freecores == 0) { // no cores for this shape+phase — give this + the unprocessed rest back - sched_->early_dispatch_queues[s].push_batch_tagged(&batch[bi], &task_id_snapshots[bi], got - bi); - break; - } - int32_t start = 0; - int32_t claim = c->claim_block_range(c->logical_block_num, freecores, start); - if (claim == 0) continue; // nothing left to claim -> drop (no re-push) - // Re-push for concurrent peers BEFORE the expensive staging. - if (start + claim < c->logical_block_num) { - if (!sched_->early_dispatch_queues[s].push_tagged(c, task_id_snapshots[bi])) - LOG_INFO_V9( - "[EARLY_DISPATCH] queue full on re-push, consumer=%" PRId64, - static_cast(c->task->task_id.raw) - ); - } - // stage_consumer_blocks fills the idle bucket (RUNNING slot) then the pend - // bucket (gated pending); pass this phase's bucket in the matching slot and an - // empty other so only the phase's cores are staged. - CoreTracker::BitStates empty(0ULL); - total_staged += is_idle ? stage_consumer_blocks(thread_idx, c, shape, start, claim, bucket, empty) : - stage_consumer_blocks(thread_idx, c, shape, start, claim, empty, bucket); - } - return total_staged; -} - -// Early-dispatch drain (idle pass) — the EARLY source's analog of dispatch_ready_tasks. -// Both sources share run_staging_order for the shape order (MIX strict priority, IDLE -// before PENDING, cross-thread idle gating: MIX-IDLE ▶ c/v-IDLE ▶ MIX-PEND ▶ c/v-PEND) and -// both drain their sync_start cohort FIRST as the highest occupancy tier (Tier 0), via the -// same all-or-nothing drain barrier. This one owns its own gating and progress flags. -// Returns the number of blocks staged this pass (for the EarlyDispatch swimlane bar). -int32_t SchedulerContext::try_early_dispatch( - int32_t thread_idx, CoreTracker &tracker, bool pmu_active, bool &made_progress, bool &try_pushed -) { - // Gate, owned here rather than by the caller (mirrors dispatch_ready_tasks - // withholding PENDING under PMU internally): - // - pmu_active: staging gated work perturbs the single-issue PMU windows the - // same way dual-issue PENDING dispatch does, so early dispatch is off. - // - has_any_free_slot: this thread has no spare capacity to stage onto (a - // purely local read; a fully-occupied thread bails before touching shared - // queues). - // - ready queues empty: normal dispatch (both the ready sync_start lane and the - // regular ready_queues) strictly precedes early — there is no real ready task - // to delay only when every normal queue is drained. - if (pmu_active || !tracker.has_any_free_slot()) return 0; - for (int s = 0; s < PTO2_NUM_RESOURCE_SHAPES; s++) { - if (sched_->ready_sync_queues[s].size() > 0 || sched_->ready_queues[s].size() > 0) return 0; - } - - // ===== Tier 0: sync_start cohorts (highest occupancy tier, all-or-nothing) ===== - // sync_start candidates park in their own shape-agnostic queue. They cannot ride - // early_dispatch_shape's per-thread partial range-claim (a partial claim strands gated - // cohort blocks nobody rings, so the rendezvous never reaches block_num). Instead arm the - // drain barrier: it takes exclusive tracker access and only stages when global - // idle+pending >= block_num, guaranteeing all-or-nothing. Win => the dispatch loop runs - // the gated drain next iteration; lose (a drain is already armed, or capacity < - // block_num) => cancel the owner and re-push or transfer its final ready route. A - // non-STAGING pop was already released and is dropped. Staging happens inside the drain, - // so this arms at most one drain and adds no blocks to total_staged here. - uint64_t sync_task_id_snapshot = 0; - if (PTO2TaskSlotState *c = sched_->early_sync_start_queue.pop_tagged(&sync_task_id_snapshot)) { - bool current_sync_task = static_cast(c->task->task_id.raw) == sync_task_id_snapshot && - c->active_mask.requires_sync_start(); - if (current_sync_task && PTO2SchedulerState::try_claim_early_sync_drain(*c->payload)) { - if (c->payload->early_dispatch_state.load(std::memory_order_seq_cst) != PTO2_EARLY_DISPATCH_STAGING) { - sched_->cancel_early_sync_drain(*c); - } else if (enter_drain_mode(c, c->logical_block_num)) { - PTO2SchedulerState::mark_early_sync_drain_armed(*c->payload); - } else { - sched_->cancel_early_sync_drain(*c); - } - } - } - - // Regular early staging (NOT is_ready): same MIX/idle/pending order as normal dispatch, - // via the shared skeleton. early_dispatch_shape stages a gated block range and never - // enters drain, so the stage callback always reports "no stop". - int32_t total_staged = 0; - run_staging_order( - thread_idx, pmu_active, - [&](PTO2ResourceShape shape, CoreTracker::DispatchPhase phase) { - total_staged += early_dispatch_shape(thread_idx, shape, phase); - return false; - }, - [&] { - return has_residual_early_mix(); - } - ); - - // Staging is dispatch work: reset the idle/stall clock and route this iter's tail - // cycles to the dispatch accumulator, exactly as normal dispatch does. - if (total_staged > 0) { - made_progress = true; - try_pushed = true; + dispatch_shape( + thread_idx, s, Phase::PENDING, local_bufs[static_cast(s)], tracker, entered_drain, made_progress + ); + if (entered_drain) return; } - return total_staged; } -// ============================================================================= -// Main scheduler dispatch loop -// ============================================================================= - -int32_t SchedulerContext::resolve_and_dispatch(Runtime *runtime, int32_t thread_idx) { - always_assert(sched_ != nullptr); - CoreTracker &tracker = core_trackers_[thread_idx]; - LOG_INFO_V0("Thread %d: resolve_and_dispatch entry", thread_idx); - - PTO2SharedMemoryHeader *header = sched_->sm_header; - if (!header) { - LOG_ERROR("PTO2 dispatch: header is null"); - return -1; - } - LOG_INFO_V0( - "Thread %d: header=%p, task_desc_offset[0]=%lu, window_size=%lu", thread_idx, static_cast(header), - static_cast(header->rings[0].task_descriptors_offset), - static_cast(header->rings[0].task_window_size) - ); - - Handshake *hank = static_cast(runtime->dev.workers); - LOG_INFO_V0( - "Thread %d: hank=%p, window_size=%lu", thread_idx, static_cast(hank), - static_cast(header->rings[0].task_window_size) - ); - - LOG_INFO_V0("Thread %d: PTO2 dispatch starting with %d cores", thread_idx, core_trackers_[thread_idx].core_num()); - int32_t cur_thread_completed = 0; - // Non-zero once a scheduler-hang timeout latches; returned in place of the - // completed count so the caller still sees the negative error rc while the - // shared end-of-loop flush below runs. - int32_t timeout_rc = 0; - int32_t idle_iterations = 0; - int32_t last_progress_count = 0; -#if SIMPLER_DFX - auto &l2_swimlane = sched_l2_swimlane_[thread_idx]; - l2_swimlane.reset(); - l2_swimlane.l2_swimlane_enabled = (l2_swimlane_level_ != L2SwimlaneLevel::DISABLED); -#endif - - 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 - // compiler cannot hoist it across the dispatch loop on its own. -#if SIMPLER_DFX - const bool pmu_active = is_pmu_enabled(); -#else - // PMU is definitionally off when profiling is compiled out; hard-set false - // so dispatch keeps its overlapping (non-single-issue) fast path. - constexpr bool pmu_active = false; -#endif - -#if SIMPLER_DFX - l2_swimlane.sched_start_ts = get_sys_cnt_aicpu(); -#endif - -#if SIMPLER_DFX - // Queue-depth snapshot carried across the iteration boundary: each phase - // emit consumes (phase_start_shared) and refreshes it with its own end - // snapshot so the next phase's "at_start" equals the previous phase's - // "at_end". - // - // L2SWIMLANE_NUM_QUEUE_SHAPES (3) matches PTO2_NUM_RESOURCE_SHAPES: AIC/AIV/MIX. - // - // **Hot-path cost discipline.** Shared depth (PTO2ReadyQueue::size) is two - // atomic relaxed loads against cache lines that all peer sched threads also - // write to (enqueue_pos and dequeue_pos bounce on every push + every pop). - // With both phases emitting per iter that's cross-core loads × thousands of - // iters per run, a measurable AICPU slowdown. Mitigation: lazy + per-iter - // cached shared snapshot, refreshed at most once per iteration. The - // complete-emit and dispatch-emit in the same iter both reuse the same - // shared sample. - static_assert( - L2SWIMLANE_NUM_QUEUE_SHAPES == PTO2_NUM_RESOURCE_SHAPES, - "queue snapshot width must match runtime resource shape count" - ); - int16_t phase_start_shared[L2SWIMLANE_NUM_QUEUE_SHAPES] = {0}; - int16_t iter_shared_snapshot[L2SWIMLANE_NUM_QUEUE_SHAPES] = {0}; - bool iter_shared_sampled = false; - auto get_or_sample_shared = [&]() -> const int16_t * { - if (!iter_shared_sampled) { - // Clamp to int16_t max before narrowing. PTO2_PROF_READYQUEUE_SIZE - // is in the low thousands today but could grow with platform - // scaling — without clamp, sizes above 32767 wrap to negatives - // and silently corrupt the snapshot. - constexpr size_t kMax = static_cast(std::numeric_limits::max()); - for (int s = 0; s < L2SWIMLANE_NUM_QUEUE_SHAPES; s++) { - // Total normal-source ready depth of shape `s` = regular ready lane + the - // sync_start Tier-0 lane; both feed dispatch_ready_tasks for this shape. - const size_t qsize = sched_->ready_queues[s].size() + sched_->ready_sync_queues[s].size(); - iter_shared_snapshot[s] = static_cast(std::min(qsize, kMax)); - } - iter_shared_sampled = true; - } - return iter_shared_snapshot; - }; - auto capture_phase_end = [&](int16_t shared_out[L2SWIMLANE_NUM_QUEUE_SHAPES]) { - const int16_t *shared_cached = get_or_sample_shared(); - for (int s = 0; s < L2SWIMLANE_NUM_QUEUE_SHAPES; s++) - shared_out[s] = shared_cached[s]; - }; - // Queue-mutating phases (Complete / Dummy) push newly-ready consumers - // straight into the shared ready_queues[] (the local-first buffer is gone), - // so their end-of-phase shared depth differs from their start. Force a fresh - // re-sample for those emits — this also refreshes the per-iter cache so the - // next phase's start snapshot is not stale. - auto capture_phase_end_fresh = [&](int16_t shared_out[L2SWIMLANE_NUM_QUEUE_SHAPES]) { - iter_shared_sampled = false; - capture_phase_end(shared_out); - }; - if (l2_swimlane_level_ >= L2SwimlaneLevel::SCHED_PHASES) { - capture_phase_end(phase_start_shared); - } -#endif - - // Wall-clock timestamp of the last completed task on this thread. - // Updated on made_progress; consulted to decide whether the wall-clock - // budget for declaring a scheduler hang has elapsed. Initialized to - // "now" so the first budget cycle starts when this thread does, not at - // an undefined value. - uint64_t last_progress_ts = get_sys_cnt_aicpu(); - // Per-device override latched once at worker init by simpler_aicpu_init - // (InitArgs.scheduler_timeout_ms -> resident-SO global). 0 means no - // override; fall back to the compile-time SCHEDULER_TIMEOUT_CYCLES. - uint64_t scheduler_timeout_cycles = SCHEDULER_TIMEOUT_CYCLES; - const int32_t scheduler_timeout_ms_override = get_scheduler_timeout_ms(); - if (scheduler_timeout_ms_override > 0) { - scheduler_timeout_cycles = - static_cast(scheduler_timeout_ms_override) * (PLATFORM_PROF_SYS_CNT_FREQ / 1000); - } - - while (true) { - if (completed_.load(std::memory_order_acquire)) { - break; - } - bool made_progress = false; -#if SIMPLER_DFX - CYCLE_COUNT_START(); - l2_swimlane.sched_loop_count++; - uint64_t _t0_phase = _t0; - // Release is the only "no Complete/Dispatch bar" attribution we keep — - // emitted with its own span in the idle branch below. Iterations that - // only scan/poll show as blank gaps; the per-loop Poll/Scan bars (PR - // #1079 debug overlay) were removed since "scheduler is polling when - // there's nothing to do" carries no actionable signal. - // Per-iter lazy shared-queue snapshot: first phase emit in this iter - // pays the atomic-load cost, subsequent emits in the same iter reuse - // the cached value. Reset here so we re-sample exactly once per iter - // (or skip entirely on iters with no phase emit). - iter_shared_sampled = false; -#endif - int32_t task_count = 0; - if (!tracker.has_any_running_cores()) { - LoopAction action = handle_orchestrator_exit(thread_idx, header, runtime, task_count); - if (action == LoopAction::BREAK_LOOP) break; - } - -#if SIMPLER_DFX - CYCLE_COUNT_LAP(l2_swimlane.sched_idle_cycle); -#endif - - // Phase 1: Check running cores for completion - int32_t completed_this_turn = 0; - - 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 - ); - } - if (completed_this_turn > 0) { -#if SIMPLER_SCHED_PROFILING - sched_->tasks_completed.fetch_add(completed_this_turn, std::memory_order_relaxed); -#endif - int32_t prev = completed_tasks_.fetch_add(completed_this_turn, std::memory_order_relaxed); - int32_t new_total = prev + completed_this_turn; - last_progress_count = new_total; - if (thread_idx == 0 && task_count > 0) { - if (new_total <= PROGRESS_VERBOSE_THRESHOLD || - new_total / PROGRESS_LOG_INTERVAL != prev / PROGRESS_LOG_INTERVAL || new_total >= task_count) { - LOG_INFO_V9( - "PTO2 progress: completed=%d total=%d (%.1f%%)", new_total, task_count, - 100.0 * new_total / task_count - ); - } - } - } - -#if SIMPLER_DFX - // Close the Complete phase BEFORE the async-wait poll so async-engine - // (SDMA/RoCE/URMA/CCU) wait time lands in its own AsyncPoll bar instead - // of folding into the Complete span. A finished slot OR a sub-block - // retire that finished no slot both count as completion work (the latter - // surfaces the SPMD harvest tail; on a pure-retire iteration - // phase_complete_count is 0). - if (!try_completed) { - CYCLE_COUNT_LAP(l2_swimlane.sched_idle_cycle); - } else { - CYCLE_COUNT_LAP(l2_swimlane.sched_complete_cycle); - if (l2_swimlane_level_ >= L2SwimlaneLevel::SCHED_PHASES && - (l2_swimlane.phase_complete_count > 0 || l2_swimlane.phase_subretire_count > 0)) { - // Complete's release_fanin pushes newly-ready consumers into the - // shared ready_queues[], so the end depth differs from the start. - int16_t phase_end_shared[L2SWIMLANE_NUM_QUEUE_SHAPES]; - capture_phase_end_fresh(phase_end_shared); - l2_swimlane_aicpu_record_sched_phase( - thread_idx, L2SwimlaneSchedPhaseKind::Complete, _t0_phase, _t1, l2_swimlane.sched_loop_count, - l2_swimlane.phase_complete_count + l2_swimlane.phase_subretire_count, /*pop_hit=*/0, - /*pop_miss=*/0, phase_start_shared, phase_end_shared - ); - for (int s = 0; s < L2SWIMLANE_NUM_QUEUE_SHAPES; s++) - phase_start_shared[s] = phase_end_shared[s]; - l2_swimlane.phase_complete_count = 0; - l2_swimlane.phase_subretire_count = 0; - } - // Advance past the completion check even when no Complete bar was - // emitted (both counts 0). Otherwise its wall time folds into the - // next bar (AsyncPoll, or Dispatch when the poll is skipped). - _t0_phase = _t1; - } -#endif - - 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 -#if SIMPLER_SCHED_PROFILING - , - thread_idx -#endif - ); - if (poll_result.error_code != PTO2_ERROR_NONE) { - int32_t expected = PTO2_ERROR_NONE; - header->sched_error_code.compare_exchange_strong( - expected, poll_result.error_code, std::memory_order_acq_rel, std::memory_order_acquire - ); - completed_.store(true, std::memory_order_release); - break; - } - if (poll_result.completed > 0) { -#if SIMPLER_SCHED_PROFILING - sched_->tasks_completed.fetch_add(poll_result.completed, std::memory_order_relaxed); -#endif - int32_t prev = completed_tasks_.fetch_add(poll_result.completed, std::memory_order_relaxed); - int32_t new_total = prev + poll_result.completed; - last_progress_count = new_total; - made_progress = true; - } -#if SIMPLER_DFX - // AsyncPoll phase: the async-wait completion poll, split out of - // Complete. Recorded here inside the poll branch, so "did the poll - // run" needs no separate flag. sched_async_cycle accrues only on - // iterations that actually poll. The bar is emitted even on a - // zero-completion poll so its polling cost stays visible rather than - // folding into the next bar. tasks_processed = async subtasks - // completed this iter. - CYCLE_COUNT_LAP(l2_swimlane.sched_async_cycle); - if (l2_swimlane_level_ >= L2SwimlaneLevel::SCHED_PHASES) { - // A completing poll runs on_task_complete, which pushes - // newly-ready consumers into the shared ready_queues[] — so the - // end depth then differs from the start; a zero-completion poll - // leaves the queues untouched and the cached start sample holds. - int16_t phase_end_shared[L2SWIMLANE_NUM_QUEUE_SHAPES]; - if (poll_result.completed > 0) { - capture_phase_end_fresh(phase_end_shared); - } else { - capture_phase_end(phase_end_shared); - } - l2_swimlane_aicpu_record_sched_phase( - thread_idx, L2SwimlaneSchedPhaseKind::AsyncPoll, _t0_phase, _t1, l2_swimlane.sched_loop_count, - static_cast(poll_result.completed), /*pop_hit=*/0, /*pop_miss=*/0, phase_start_shared, - phase_end_shared - ); - for (int s = 0; s < L2SWIMLANE_NUM_QUEUE_SHAPES; s++) - phase_start_shared[s] = phase_end_shared[s]; - _t0_phase = _t1; - } -#endif - } - - bool try_pushed = false; - - // Phase 2 drain check - if (drain_state_.sync_start_pending.load(std::memory_order_acquire) != 0) { -#if SIMPLER_DFX - // The drain is otherwise a swimlane blind spot: the `continue` below skips - // every phase record, and handle_drain_mode is uninstrumented. Time it here so - // the sync_start stop-the-world window shows on the scheduler lane (one bar per - // iteration that enters the drain; retries appear as multiple bars). - uint64_t drain_t0 = (l2_swimlane_level_ >= L2SwimlaneLevel::SCHED_PHASES) ? get_sys_cnt_aicpu() : 0; - uint64_t drain_stage_wall = 0; // set by handle_drain_mode ONLY if this thread staged - handle_drain_mode(thread_idx, &drain_stage_wall); - // Record a Drain bar only when this thread actually did drain work (reached - // drain_stage_cores). The many no-op entries — ack + availability-insufficient - // reset, stale-elected, non-elected bail before stage_go — never stage, so they - // would otherwise clutter the lane with zero-work drain(0) bars. - if (l2_swimlane_level_ >= L2SwimlaneLevel::SCHED_PHASES && drain_stage_wall != 0) { - l2_swimlane_aicpu_record_sched_phase( - thread_idx, L2SwimlaneSchedPhaseKind::Drain, drain_t0, get_sys_cnt_aicpu(), - l2_swimlane.sched_loop_count, static_cast(drain_stage_wall) - ); - } -#else - handle_drain_mode(thread_idx); -#endif - continue; - } - - // Phase 3: Drain dummy ready queue (S0/S1/S2). - // - // Dependency-only tasks bypass AICore dispatch: they go through the - // scheduler so fanin/fanout edges stay consistent, but completion is - // signalled inline here. The ready queue is MPMC, and the fanout path - // uses per-slot locks/atomics, so multiple scheduler threads can share - // the dependency-only resolve work. - if (thread_idx < 3) { - constexpr int DUMMY_DRAIN_BATCH = 8; - PTO2TaskSlotState *dummy_batch[DUMMY_DRAIN_BATCH]; - int dummy_got = sched_->dummy_ready_queue.pop_batch(dummy_batch, DUMMY_DRAIN_BATCH); -#if SIMPLER_DFX - // Dummy outer phase: covers handling of all dummies popped this - // iter. Per-dummy DummyTask markers are emitted to a SEPARATE lane - // (Worker View AICPU_N) by the converter, so they do not nest - // under this bar. Resolve emits below DO land on the sched lane - // and nest under this Dummy outer by time containment. - uint64_t dummy_outer_t0 = - (dummy_got > 0 && l2_swimlane_level_ >= L2SwimlaneLevel::SCHED_PHASES) ? get_sys_cnt_aicpu() : 0; -#endif - for (int di = 0; di < dummy_got; di++) { - PTO2TaskSlotState &dummy_slot = *dummy_batch[di]; - - // ----- Resolve work: walk this dummy's consumer list. ------ - // Same 1 µs filter as the main-path Resolve emit suppresses - // dummies whose consumer release runs sub-microsecond. -#if SIMPLER_DFX - uint64_t dummy_resolve_t0 = - (l2_swimlane_level_ >= L2SwimlaneLevel::SCHED_PHASES) ? get_sys_cnt_aicpu() : 0; -#endif - // [[maybe_unused]] silences -Werror=unused-but-set-variable on - // the profiling-flags-smoke build path where SIMPLER_DFX is - // OFF and the Resolve emit below is excluded. - [[maybe_unused]] uint32_t dummy_consumers = 0; -#if SIMPLER_SCHED_PROFILING - dummy_consumers = sched_->on_task_complete(dummy_slot, thread_idx).fanout_edges; -#else - dummy_consumers = sched_->on_task_complete(dummy_slot); -#endif -#if SIMPLER_DFX - if (dummy_resolve_t0 != 0) { - uint64_t dummy_resolve_t1 = get_sys_cnt_aicpu(); - constexpr uint64_t RESOLVE_EMIT_MIN_CYCLES = PLATFORM_PROF_SYS_CNT_FREQ / 1'000'000; // 1 µs - if (dummy_resolve_t1 - dummy_resolve_t0 >= RESOLVE_EMIT_MIN_CYCLES) { - l2_swimlane_aicpu_record_sched_phase( - thread_idx, L2SwimlaneSchedPhaseKind::Resolve, dummy_resolve_t0, dummy_resolve_t1, - sched_l2_swimlane_[thread_idx].sched_loop_count, dummy_consumers - ); - } - l2_swimlane_aicpu_record_dummy_task( - thread_idx, dummy_resolve_t0, sched_l2_swimlane_[thread_idx].sched_loop_count, - dummy_slot.task->task_id.raw - ); - } -#endif - // 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 - } - } - int32_t prev = completed_tasks_.fetch_add(1, std::memory_order_relaxed); - last_progress_count = prev + 1; - cur_thread_completed++; - } - if (dummy_got > 0) { - made_progress = true; - } -#if SIMPLER_DFX - // Emit Dummy outer over the whole dummy_drain pass. Span starts at - // dummy_outer_t0 (captured after pop_batch) and ends at "now". - // tasks_processed = dummy_got. Advancing _t0_phase here makes the - // following Dispatch / EarlyDispatch / second-Complete bars start - // at this end. - if (dummy_outer_t0 != 0) { - uint64_t dummy_outer_t1 = get_sys_cnt_aicpu(); - int16_t phase_end_shared[L2SWIMLANE_NUM_QUEUE_SHAPES]; - capture_phase_end_fresh(phase_end_shared); - l2_swimlane_aicpu_record_sched_phase( - thread_idx, L2SwimlaneSchedPhaseKind::Dummy, dummy_outer_t0, dummy_outer_t1, - l2_swimlane.sched_loop_count, static_cast(dummy_got), /*pop_hit=*/0, - /*pop_miss=*/0, phase_start_shared, phase_end_shared - ); - for (int s = 0; s < L2SWIMLANE_NUM_QUEUE_SHAPES; s++) - phase_start_shared[s] = phase_end_shared[s]; - _t0_phase = dummy_outer_t1; - // We do NOT re-sync _t0/_t1 — the dummy span will be absorbed - // into the next CYCLE_COUNT_LAP accumulator. The phase-model - // anchor (_t0_phase) is the authoritative source for bar spans - // on the swimlane; the cycle accumulators are coarse aggregates. - } -#endif - } - - // Phase 4: MIX-strict-priority dispatch with phase-split and - // cross-thread idle gating. See dispatch_ready_tasks for the policy. -#if SIMPLER_DFX - uint64_t dispatch_t0 = (l2_swimlane_level_ >= L2SwimlaneLevel::SCHED_PHASES) ? get_sys_cnt_aicpu() : 0; -#endif - dispatch_ready_tasks(thread_idx, tracker, pmu_active, made_progress, try_pushed); -#if SIMPLER_DFX - // Emit Dispatch IMMEDIATELY after dispatch_ready_tasks so its span - // covers the actual publish work — not the trailing second-poll / - // early-dispatch time. (Pre-redesign the Dispatch emit lived at iter - // end with span extending past the second poll, which made finish_time - // events from the second poll fall under the Dispatch bar rather than - // a Complete bar of their own — confusing for trace consumers.) - if (dispatch_t0 != 0 && l2_swimlane.phase_dispatch_count > 0) { - uint64_t dispatch_t1 = get_sys_cnt_aicpu(); - uint64_t pop_hit_delta = l2_swimlane.pop_hit - l2_swimlane.pop_hit_at_last_emit; - uint64_t pop_miss_delta = l2_swimlane.pop_miss - l2_swimlane.pop_miss_at_last_emit; - debug_assert(pop_hit_delta < (1ULL << 32)); - debug_assert(pop_miss_delta < (1ULL << 32)); - int16_t phase_end_shared[L2SWIMLANE_NUM_QUEUE_SHAPES]; - capture_phase_end(phase_end_shared); - l2_swimlane_aicpu_record_sched_phase( - thread_idx, L2SwimlaneSchedPhaseKind::Dispatch, _t0_phase, dispatch_t1, l2_swimlane.sched_loop_count, - l2_swimlane.phase_dispatch_count, static_cast(pop_hit_delta), - static_cast(pop_miss_delta), phase_start_shared, phase_end_shared - ); - for (int s = 0; s < L2SWIMLANE_NUM_QUEUE_SHAPES; s++) { - phase_start_shared[s] = phase_end_shared[s]; - } - _t0_phase = dispatch_t1; - l2_swimlane.phase_dispatch_count = 0; - l2_swimlane.pop_hit_at_last_emit = l2_swimlane.pop_hit; - l2_swimlane.pop_miss_at_last_emit = l2_swimlane.pop_miss; - } -#endif - - // Phase 4b: early-dispatch onto spare cores, mirroring the Phase 4 call - // shape. try_early_dispatch owns its own gating (off-PMU, this - // thread has a spare slot, no normal ready work queued) and updates - // made_progress / try_pushed, so this is a single unconditional call — it - // returns 0 without staging when gated out. -#if SIMPLER_DFX - bool early_dispatch_record = l2_swimlane_level_ >= L2SwimlaneLevel::SCHED_PHASES; - uint64_t early_dispatch_t0 = early_dispatch_record ? get_sys_cnt_aicpu() : 0; -#endif - [[maybe_unused]] int32_t staged_count = - try_early_dispatch(thread_idx, tracker, pmu_active, made_progress, try_pushed); -#if SIMPLER_DFX - // Emit an EarlyDispatch bar so a staging-dominated iteration is attributed - // to early-dispatch rather than disappearing into a blank gap. - if (early_dispatch_record && staged_count > 0) { - uint64_t early_dispatch_t1 = get_sys_cnt_aicpu(); - l2_swimlane_aicpu_record_sched_phase( - thread_idx, L2SwimlaneSchedPhaseKind::EarlyDispatch, early_dispatch_t0, early_dispatch_t1, - sched_l2_swimlane_[thread_idx].sched_loop_count, static_cast(staged_count) - ); - // prepare_block_for_dispatch bumped phase_dispatch_count while staging; - // those blocks belong to this EarlyDispatch bar, so clear the counter - // before it leaks into the next Dispatch bar. - sched_l2_swimlane_[thread_idx].phase_dispatch_count = 0; - // Advance _t0_phase so the next phase bar starts at the EarlyDispatch - // end, not before it (otherwise their spans overlap and the - // outer-phase mutual-exclusion breaks). - _t0_phase = early_dispatch_t1; - } -#endif - -#if SIMPLER_DFX - // Cycle-counter LAP for the iter tail. Dispatch's emit moved earlier - // (see Phase 4 above) so this branch only routes the time accumulator. - if (!try_pushed) { - CYCLE_COUNT_LAP(l2_swimlane.sched_idle_cycle); - } else { - CYCLE_COUNT_LAP(l2_swimlane.sched_dispatch_cycle); - } -#endif - -#if !SIMPLER_DFX - (void)try_completed; - (void)try_pushed; -#endif - - if (made_progress) { - 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 - idle_iterations++; - - if (idle_iterations % FATAL_ERROR_CHECK_INTERVAL == 0) { - LoopAction action = check_idle_fatal_error(thread_idx, header, runtime); - if (action == LoopAction::BREAK_LOOP) break; - } - - if (idle_iterations % STALL_LOG_INTERVAL == 0) { - log_stall_diagnostics(thread_idx, total_tasks_, idle_iterations, last_progress_count); - } - // Wall-clock budget gate, with two fatal-latch branches: - // - // 1. Self owns a RUNNING task — first-hand evidence the - // dispatch is stuck. Latch. - // 2. No thread anywhere owns a RUNNING task AND tasks remain - // unfinished — the system is in a pre-dispatch / WAIT-only - // deadlock (e.g. dependency cycle). Ownerless idle threads - // are the only observers; let this one latch on the global - // evidence (`completed_tasks_ < total_tasks_` and - // `no_thread_owns_running_task()`). - // - // Otherwise: a sibling thread owns a RUNNING task but hasn't - // hit its own budget yet (typical distributed startup-skew - // case) — refresh last_progress_ts and keep spinning. The - // STALL diagnostic above still fires periodically so - // observability is preserved. - if (get_sys_cnt_aicpu() - last_progress_ts > scheduler_timeout_cycles) { - bool self_owns = self_owns_running_task(thread_idx); - bool global_stuck = !self_owns && total_tasks_ > 0 && - completed_tasks_.load(std::memory_order_relaxed) < total_tasks_ && - no_thread_owns_running_task(); - if (self_owns || global_stuck) { - // Latch the error + emergency_shutdown, then break to the - // shared end-of-loop cleanup so the diagnostic buffers get - // flushed to the host. An early return here would strand the - // stuck task's already-dumped inputs and every completed - // task's in/out records in the unflushed per-thread dump - // buffer — exactly the state we need to triage the hang. - timeout_rc = handle_timeout_exit( - thread_idx, header, runtime, idle_iterations, last_progress_count -#if SIMPLER_DFX - , - l2_swimlane.sched_start_ts -#endif - ); - break; - } - last_progress_ts = get_sys_cnt_aicpu(); - } - SPIN_WAIT_HINT(); -#if SIMPLER_DFX - CYCLE_COUNT_LAP(l2_swimlane.sched_idle_cycle); - // _t0_phase advances through idle laps so the next emitted - // COMPLETE/DISPATCH bar starts at the iter it actually ran in, not - // at the start of the preceding idle stretch. The idle/poll time - // itself is attributed by the activity-fill below — no blanks. - if (l2_swimlane_level_ >= L2SwimlaneLevel::SCHED_PHASES) { - _t0_phase = _t1; - } -#endif - } - } - - // 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 - } - -#if SIMPLER_DFX - // Final-drain: emit any pop_hit / pop_miss accrued since the last - // dispatch emit (typically the trailing idle loops while waiting for - // orchestrator_done_) as a zero-duration synthetic dispatch record so - // sum(record.pop_*) reconciles with the run-cumulative counter. - // Gate on SCHED_PHASES — at lower levels the phase buffer is never - // flushed (see below), so writing this record would be wasted work. - if (l2_swimlane_level_ >= L2SwimlaneLevel::SCHED_PHASES) { - uint64_t final_pop_hit_delta = l2_swimlane.pop_hit - l2_swimlane.pop_hit_at_last_emit; - uint64_t final_pop_miss_delta = l2_swimlane.pop_miss - l2_swimlane.pop_miss_at_last_emit; - debug_assert(final_pop_hit_delta < (1ULL << 32)); - debug_assert(final_pop_miss_delta < (1ULL << 32)); - if (final_pop_hit_delta != 0 || final_pop_miss_delta != 0) { - uint64_t t_now = get_sys_cnt_aicpu(); - int16_t phase_end_shared[L2SWIMLANE_NUM_QUEUE_SHAPES]; - capture_phase_end(phase_end_shared); - l2_swimlane_aicpu_record_sched_phase( - thread_idx, L2SwimlaneSchedPhaseKind::Dispatch, t_now, t_now, l2_swimlane.sched_loop_count, 0, - static_cast(final_pop_hit_delta), static_cast(final_pop_miss_delta), - phase_end_shared, phase_end_shared - ); - l2_swimlane.pop_hit_at_last_emit = l2_swimlane.pop_hit; - l2_swimlane.pop_miss_at_last_emit = l2_swimlane.pop_miss; - } - } - log_l2_swimlane_summary(thread_idx, cur_thread_completed); -#endif - -#if SIMPLER_DFX - if (l2_swimlane.l2_swimlane_enabled) { - l2_swimlane_aicpu_flush( - thread_idx, core_trackers_[thread_idx].core_ids(), core_trackers_[thread_idx].core_num() - ); - if (l2_swimlane_level_ >= L2SwimlaneLevel::SCHED_PHASES) { - l2_swimlane_aicpu_flush_sched_phase_buffer(thread_idx); - } - } -#endif -#if SIMPLER_DFX - if (is_dump_args_enabled()) { - dump_args_flush(thread_idx); - } -#endif -#if SIMPLER_DFX - if (is_pmu_enabled()) { - pmu_aicpu_flush_buffers( - thread_idx, core_trackers_[thread_idx].core_ids(), core_trackers_[thread_idx].core_num() - ); +bool SchedulerContext::has_idle_in_other_threads(int32_t self_thread_idx, PTO2ResourceShape shape) const { + for (int32_t t = 0; t < active_sched_threads_; t++) { + if (t == self_thread_idx) continue; + if (core_trackers_[t].get_idle_core_offset_states(shape).has_value()) return true; } -#endif - - return timeout_rc != 0 ? timeout_rc : cur_thread_completed; + return false; } diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_types.h b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_types.h index ff60af0f67..33c4c47977 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_types.h +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_types.h @@ -19,55 +19,14 @@ #include "pto_runtime2_types.h" #include "spin_hint.h" -// ============================================================================= -// Profiling macros (compile-time gated) -// ============================================================================= - -#if SIMPLER_DFX -#include "aicpu/device_time.h" -// Accumulated nanoseconds per sub-step -#define CYCLE_COUNT_START() uint64_t _t0 = get_sys_cnt_aicpu(), _t1 -#define CYCLE_COUNT_LAP(acc) \ - do { \ - _t1 = get_sys_cnt_aicpu(); \ - acc += (_t1 - _t0); \ - _t0 = _t1; \ - } while (0) -#else -#define CYCLE_COUNT_START() -#define CYCLE_COUNT_LAP(acc) -#endif - -// ============================================================================= -// Scheduler constants -// ============================================================================= - constexpr int32_t MAX_AICPU_THREADS = PLATFORM_MAX_AICPU_THREADS; -// Periodic cadence (in idle iterations) for emitting the per-thread STALL -// diagnostic while no progress is being made. Purely an observability knob, -// independent of the wall-clock timeout below: small enough to fire a few times -// before the budget expires, large enough not to flood device_log. +// PLATFORM_MAX_IDLE_ITERATIONS was removed upstream; fixed cadence matches a5's +// equivalent (used only for per-thread diagnostic logging, not for the fatal- +// timeout path which uses wall-clock). constexpr int32_t STALL_LOG_INTERVAL = 480000; constexpr int32_t FATAL_ERROR_CHECK_INTERVAL = 1024; // Check orchestrator error every N idle iters -// Wall-clock budget for declaring "no progress = scheduler timeout". Replaces -// the per-thread iteration-count cap that once lived here as MAX_IDLE_ITERATIONS -// for the fatal-latch decision; STALL_LOG_INTERVAL above keeps the per-thread -// diagnostic cadence. -// -// Using wall-clock here is load-bearing for distributed runs: with per-thread -// iteration counts, a pure-idle thread spinning ~115 ns/iter hits the cap in -// ~92 ms while a sibling thread polling a RUNNING task takes ~200 ms for the -// same iteration count. The fast spinner racing ahead and latching fatal -// kills the slower-but-correct poller mid-poll — see the distributed -// startup-skew scenario in issue #897. -// -// The budget is platform-defined (PLATFORM_SCHEDULER_TIMEOUT_MS in spin_hint.h). -// Onboard keeps it below the STARS op-execute and host stream-sync budgets so -// the AICPU can flush diagnostics before the host-visible timeout chain fires. -// Sim has no STARS or ACL stream-sync timeout, but uses the same no-progress -// watchdog shape. See spin_hint.h for the per-variant rationale. constexpr int32_t SCHEDULER_TIMEOUT_MS = PLATFORM_SCHEDULER_TIMEOUT_MS; constexpr uint64_t SCHEDULER_TIMEOUT_CYCLES = static_cast(SCHEDULER_TIMEOUT_MS) * (PLATFORM_PROF_SYS_CNT_FREQ / 1000); @@ -77,21 +36,11 @@ constexpr int32_t STALL_DUMP_CORE_MAX = 8; constexpr int32_t PROGRESS_VERBOSE_THRESHOLD = 10; // log every completion for the first N tasks constexpr int32_t PROGRESS_LOG_INTERVAL = 250; // log every N completions after threshold -// ============================================================================= -// Control flow signal from cold-path helpers back to the main dispatch loop. -// ============================================================================= - enum class LoopAction : int8_t { NONE, // cold path did not trigger; proceed normally BREAK_LOOP, // equivalent to 'break' from the while(true) loop }; -// ============================================================================= -// Per-core state: one cache line per core to eliminate false sharing -// and co-locate all hot-path fields for minimal cache misses. -// Dual-slot layout: running (currently executing) + pending (pre-loaded, awaiting hardware pickup). -// ============================================================================= - struct alignas(64) CoreExecState { // --- Hot fields (completion + dispatch, every iteration) --- uint64_t reg_addr; // offset 0: register base address (set once in handshake) @@ -103,34 +52,15 @@ struct alignas(64) CoreExecState { PTO2SubtaskSlot running_subslot; // offset 36: which subtask slot is running PTO2SubtaskSlot pending_subslot; // offset 37: which subtask slot is pending uint8_t pad0_[2]; // offset 38: alignment padding - // Precomputed COND register pointer; resolved once in handshake so the - // hot completion poll does a single volatile load instead of recomputing - // reg_base + reg_offset(COND) on every iteration. - volatile uint32_t *cond_ptr; // offset 40: precomputed pointer to COND register -#if SIMPLER_DFX - // --- Profiling fields (dispatch path, compile-time gated) --- - uint64_t running_dispatch_timestamp; // offset 48: AICPU dispatch timestamp for running task - uint64_t pending_dispatch_timestamp; // offset 56: AICPU dispatch timestamp for pending task -#else + volatile uint32_t *cond_ptr; // offset 40: precomputed pointer to COND register // --- Cold fields (init/diagnostics only, never in hot path) --- int32_t worker_id; // offset 48: index in runtime.dev.workers[] uint32_t physical_core_id; // offset 52: hardware physical core ID CoreType core_type; // offset 56: AIC or AIV (enum class : int32_t) uint8_t pad2_[4]; // offset 60: pad to 64 bytes -#endif }; static_assert(sizeof(CoreExecState) == 64, "CoreExecState must occupy exactly one cache line"); -// ============================================================================= -// CoreTracker: cluster-based bitmask tracker for idle/running core state. -// -// core_states_ encodes per-cluster core idle/running in 3 bits per cluster: -// bit i*3 = AIC of cluster i (1 = idle, 0 = running) -// bit i*3+1 = AIV0 of cluster i -// bit i*3+2 = AIV1 of cluster i -// Max 21 clusters per tracker (63 bits in uint64_t). -// ============================================================================= - class alignas(64) CoreTracker { public: static inline int32_t MAX_CORE_PER_THREAD = 63; @@ -159,7 +89,6 @@ class alignas(64) CoreTracker { bool has_value() const { return states_ > 0; } int32_t count() const { return __builtin_popcountll(states_); } - void clear_bit(int32_t offset) { states_ &= ~(1ULL << offset); } // Extract the lowest set bit from mask, clear it, and return its position. // Returns -1 if mask is empty. @@ -199,11 +128,8 @@ class alignas(64) CoreTracker { template bool has_running_cores() const { - if constexpr (CT == CoreType::AIC) { - return ((~core_states_) & aic_mask_).has_value(); - } else { - return ((~core_states_) & aiv_mask_).has_value(); - } + if constexpr (CT == CoreType::AIC) return ((~core_states_) & aic_mask_).has_value(); + else return ((~core_states_) & aiv_mask_).has_value(); } bool has_any_running_cores() const { return ((~core_states_) & (aic_mask_ | aiv_mask_)).has_value(); } @@ -229,29 +155,26 @@ class alignas(64) CoreTracker { template int32_t get_running_count() const { - if constexpr (CT == CoreType::AIC) { - return ((~core_states_) & aic_mask_).count(); - } else { - return ((~core_states_) & aiv_mask_).count(); - } + if constexpr (CT == CoreType::AIC) return ((~core_states_) & aic_mask_).count(); + else return ((~core_states_) & aiv_mask_).count(); } // Return an opaque bitmask for iterating running cores of a given type. // Use pop_first() to extract core bit offsets one at a time. template BitStates get_running_cores() const { - if constexpr (CT == CoreType::AIC) { - return (~core_states_) & aic_mask_; - } else { - return (~core_states_) & aiv_mask_; - } + if constexpr (CT == CoreType::AIC) return (~core_states_) & aic_mask_; + else return (~core_states_) & aiv_mask_; } BitStates get_all_running_cores() const { return (~core_states_) & (aic_mask_ | aiv_mask_); } - BitStates get_cluster_offset_states() const { return aic_mask_; } // --- Cluster matching --- + // All cluster base offsets (one bit per cluster). Used by the gated MIX + // split-placement helpers, which iterate every cluster regardless of shape. + BitStates get_cluster_offset_states() const { return aic_mask_; } + BitStates get_valid_cluster_offset_states(PTO2ResourceShape shape) const { switch (shape) { case PTO2ResourceShape::AIC: @@ -291,10 +214,6 @@ class alignas(64) CoreTracker { // Toggle bit at the given bit offset (running <-> idle) void change_core_state(int32_t bit_offset) { core_states_ ^= BitStates(1ULL << bit_offset); } - // --- Pending-occupied tracking --- - // Tracks whether a core's pending payload slot is occupied (awaiting hardware ACK). - // SET on dispatch (both running-first and pending), CLEAR on idle or pending_freed. - void set_pending_occupied(int32_t bit_offset) { pending_occupied_ |= BitStates(1ULL << bit_offset); } void clear_pending_occupied(int32_t bit_offset) { pending_occupied_ ^= (pending_occupied_ & BitStates(1ULL << bit_offset)); @@ -302,20 +221,10 @@ class alignas(64) CoreTracker { // --- Two-phase dispatch queries --- - // Idle dispatch: returns bit offsets of idle cores for the given shape. - // For AIC: 1 bit per cluster (core offset == cluster offset). - // For AIV: 1 bit per AIV core (2 bits per cluster at aiv_mask_ positions). - // Only AIC needs pending_occupied filtering: by invariant, idle cores (core_states_ bit=1) - // always have pending_occupied=0, so AIV/MIX need no extra filtering. - // Skipping the AIC-centric filter also fixes a latent bug where a running+pending AIC core - // would incorrectly block AIV idle dispatch on the same cluster. BitStates get_idle_core_offset_states(PTO2ResourceShape shape) const { - if (shape == PTO2ResourceShape::AIC) { + if (shape == PTO2ResourceShape::AIC) return get_valid_cluster_offset_states(shape) & ~(pending_occupied_ & aic_mask_); - } - if (shape == PTO2ResourceShape::AIV) { - return core_states_ & aiv_mask_; - } + if (shape == PTO2ResourceShape::AIV) return core_states_ & aiv_mask_; return get_valid_cluster_offset_states(shape); // MIX: cluster-level } @@ -425,16 +334,13 @@ class alignas(64) CoreTracker { BitStates available = ~pending_occupied_; BitStates mix_available = (available & aic_mask_) & ((available >> 1) & aic_mask_) & ((available >> 2) & aic_mask_); - // Pending MIX can only reuse a fully-running cluster. Partially-running clusters - // could split one MIX block across immediate and pending placement. + // Exclude fully-idle clusters (handled by IDLE phase) to prevent double-dispatch. BitStates running = ~core_states_; - BitStates cluster_all_running = - (running & aic_mask_) & ((running >> 1) & aic_mask_) & ((running >> 2) & aic_mask_); - return mix_available & cluster_all_running; - } - if (shape == PTO2ResourceShape::AIC) { - return (~core_states_) & aic_mask_ & ~(pending_occupied_ & aic_mask_); + BitStates cluster_has_running = + (running & aic_mask_) | ((running >> 1) & aic_mask_) | ((running >> 2) & aic_mask_); + return mix_available & cluster_has_running; } + if (shape == PTO2ResourceShape::AIC) return (~core_states_) & aic_mask_ & ~(pending_occupied_ & aic_mask_); // AIV return (~core_states_) & aiv_mask_ & ~pending_occupied_; } @@ -464,11 +370,6 @@ class alignas(64) CoreTracker { int32_t core_id_map_[63]; // bit_position -> worker_id, max 21 clusters * 3 }; -// ============================================================================= -// SlotTransition: pure event signals from a single register poll. -// true = event occurred, false = no-op (maintain current state). -// ============================================================================= - struct SlotTransition { bool running_done = false; // running task completed bool pending_done = false; // pending task completed @@ -477,44 +378,6 @@ struct SlotTransition { bool matched = false; // some case was hit (otherwise skip apply) }; -// ============================================================================= -// Profiling counters (compile-time gated) -// ============================================================================= - -#if SIMPLER_DFX -struct alignas(64) SchedL2SwimlaneCounters { - bool l2_swimlane_enabled{false}; - uint64_t sched_start_ts{0}; - uint64_t sched_complete_cycle{0}; - uint64_t sched_dispatch_cycle{0}; - uint64_t sched_idle_cycle{0}; - uint64_t sched_async_cycle{0}; - uint64_t sched_loop_count{0}; - uint32_t phase_complete_count{0}; - // Sub-block retires that did NOT finish a slot (SPMD blocks of a multi-block - // task retiring one at a time). Counted separately so the Complete-phase - // emit can fire on poll iterations that only retired sub-blocks — otherwise - // the serial-harvest tail of an SPMD slot is invisible (no slot completes - // until the last block, leaving the scheduler lane blank for that window). - uint32_t phase_subretire_count{0}; - uint32_t phase_dispatch_count{0}; - // Per-emit delta is (current - *_at_last_emit). Accumulated only when - // l2_swimlane_level_ >= SCHED_PHASES. - uint64_t pop_hit{0}; - uint64_t pop_miss{0}; - uint64_t pop_hit_at_last_emit{0}; - uint64_t pop_miss_at_last_emit{0}; -#if SIMPLER_SCHED_PROFILING - uint64_t complete_probe_count{0}; - uint64_t complete_hit_count{0}; - uint64_t sched_complete_perf_cycle{0}; - uint64_t sched_dispatch_pop_cycle{0}; - uint64_t sched_dispatch_setup_cycle{0}; -#endif - void reset() { *this = SchedL2SwimlaneCounters{}; } -}; -#endif - // ============================================================================= // sync_start drain coordination // ============================================================================= diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/shared/pto_runtime2_init.cpp b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/shared/pto_runtime2_init.cpp index fda19e1074..a8bfaced6f 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/shared/pto_runtime2_init.cpp +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/shared/pto_runtime2_init.cpp @@ -8,50 +8,131 @@ * See LICENSE in the root of the software repository for the full text of the License. * ----------------------------------------------------------------------------------------------------------- */ + /** - * Host/AICPU shared runtime-arena layout, init_data and wire implementations. - * - * Lives under runtime/shared/ so it is included in both the host_runtime.so - * build (host pre-populates the prebuilt arena image) and the aicpu_runtime - * build (AICPU runs wire_arena_pointers + reset_for_reuse after attach). The - * device-only parts of pto_runtime2.cpp / pto_orchestrator.cpp / pto_scheduler.cpp - * (ops table, scope/submit/dispatch business logic, profiling) stay in their - * original files and the aicpu build only. + * PTO Runtime2 - cold-path layout/init/wire/reset for the orchestrator and + * scheduler. Compiled into both the host and AICPU runtimes (runtime/shared), + * so host-side arena setup and AICPU-side execution share one definition. */ -#include -#include - #include #include "pto_orchestrator.h" #include "pto_runtime2.h" -#include "pto_ring_buffer.h" #include "pto_shared_memory.h" #include "pto_tensormap.h" #include "scheduler/pto_scheduler.h" -static bool sum_ring_heap_sizes(const uint64_t heap_sizes[PTO2_MAX_RING_DEPTH], uint64_t *total) { - uint64_t sum = 0; +PTO2OrchestratorLayout PTO2OrchestratorState::reserve_layout( + DeviceArena &arena, const int32_t task_window_sizes[PTO2_MAX_RING_DEPTH], int32_t dep_pool_capacity +) { + PTO2OrchestratorLayout layout{}; + layout.dep_pool_capacity = dep_pool_capacity; + // scope_tasks holds every task in the open scope across all rings, so its cap + // is the real in-flight budget = sum of the (runtime) per-ring windows. + // Accumulate in int64; each window is validated <= INT32_MAX individually but + // their sum can exceed it. See upstream #1192. + int64_t scope_tasks_cap = 0; for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { - if (heap_sizes[r] > std::numeric_limits::max() - sum) { - LOG_ERROR("Total ring heap size overflows uint64_t"); - return false; - } - sum += heap_sizes[r]; + always_assert(task_window_sizes[r] > 0); + scope_tasks_cap += task_window_sizes[r]; } - *total = sum; + always_assert(scope_tasks_cap <= std::numeric_limits::max()); + layout.scope_tasks_cap = static_cast(scope_tasks_cap); + layout.scope_stack_capacity = PTO2_MAX_SCOPE_DEPTH; + + layout.off_scope_tasks = arena.reserve( + static_cast(layout.scope_tasks_cap) * sizeof(PTO2TaskSlotState *), alignof(PTO2TaskSlotState *) + ); + layout.off_scope_begins = + arena.reserve(static_cast(layout.scope_stack_capacity) * sizeof(int32_t), alignof(int32_t)); + layout.tensor_map = PTO2TensorMap::reserve_layout_default(arena, task_window_sizes); + return layout; +} + +bool PTO2OrchestratorState::init_data_from_layout( + const PTO2OrchestratorLayout &layout, DeviceArena &arena, void *sm_dev_base, void *gm_heap, uint64_t heap_size, + uint64_t task_window_size +) { + auto *orch = this; + *orch = PTO2OrchestratorState{}; + + orch->sm_header = reinterpret_cast(sm_dev_base); + orch->gm_heap_base = gm_heap; + orch->gm_heap_size = heap_size * PTO2_MAX_RING_DEPTH; + orch->fatal = false; + + // Mirror the SM API's per-ring window-size shape so a future per-ring + // SM layout cannot silently disagree with the addresses we compute here. + uint64_t task_window_sizes[PTO2_MAX_RING_DEPTH]; + for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) + task_window_sizes[r] = task_window_size; + + auto *orch_err = pto2_sm_layout::orch_error_code_addr(sm_dev_base); + for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { + void *ring_heap_base = reinterpret_cast(gm_heap) + r * heap_size; + auto *task_descs_dev = pto2_sm_layout::ring_task_descriptors_addr(sm_dev_base, task_window_sizes, r); + auto *cur_idx_dev = pto2_sm_layout::ring_current_task_index_addr(sm_dev_base, r); + auto *last_alive_dev = pto2_sm_layout::ring_last_task_alive_addr(sm_dev_base, r); + + orch->rings[r].task_allocator.init( + task_descs_dev, static_cast(task_window_size), cur_idx_dev, last_alive_dev, ring_heap_base, + heap_size, orch_err + ); + } + + if (!orch->tensor_map.init_data_from_layout(layout.tensor_map, arena)) return false; + + orch->scope_tasks_size = 0; + orch->scope_tasks_capacity = layout.scope_tasks_cap; + orch->scope_stack_top = -1; + orch->scope_stack_capacity = layout.scope_stack_capacity; + orch->manual_begin_depth = PTO2_MAX_SCOPE_DEPTH; + return true; } -// ============================================================================= -// Ready queue -// ============================================================================= +void PTO2OrchestratorState::wire_arena_pointers( + const PTO2OrchestratorLayout &layout, DeviceArena &arena, PTO2SchedulerState *scheduler_arg +) { + auto *orch = this; + orch->tensor_map.wire_arena_pointers(layout.tensor_map, arena); + orch->scope_tasks = static_cast(arena.region_ptr(layout.off_scope_tasks)); + orch->scope_begins = static_cast(arena.region_ptr(layout.off_scope_begins)); + orch->scheduler = scheduler_arg; +} + +// Surgical reset for the arena-reuse path (#1234). Only touches state that +// mutates across runs — leaves the arena-internal pointers wired by +// wire_arena_pointers alone, and skips the O(pool_size + num_buckets) +// tensor_map re-init in favour of an epoch bump (bucket_epochs and +// task_entry_head_epochs are compared against current_epoch on every +// lookup; a bump invalidates all stale entries in O(1)). +void PTO2OrchestratorState::reset_for_reuse() { + auto *orch = this; + for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { + orch->rings[r].task_allocator.reset_for_reuse(); + } + orch->tensor_map.reset_for_reuse(); + orch->scope_tasks_size = 0; + orch->scope_stack_top = -1; + orch->manual_begin_depth = PTO2_MAX_SCOPE_DEPTH; + orch->fatal = false; + orch->inline_completed_tasks = 0; + orch->fanin_seen_current_epoch++; + if (orch->fanin_seen_current_epoch == 0) orch->fanin_seen_current_epoch = 1; +} + +// Forget pointers; arena owns the backing buffers. +void PTO2OrchestratorState::destroy() { + auto *orch = this; + orch->tensor_map.destroy(); + orch->scope_tasks = nullptr; + orch->scope_begins = nullptr; +} +void PTO2OrchestratorState::set_scheduler(PTO2SchedulerState *scheduler) { this->scheduler = scheduler; } size_t ready_queue_reserve_layout(DeviceArena &arena, uint64_t capacity) { - // Align the slots[] base to a full cache line so MPMC CAS traffic on the - // first slot cannot false-share with whatever region sits in front of us - // (e.g. orchestrator tensormap heads written by the orch thread). return arena.reserve(capacity * sizeof(PTO2ReadyQueueSlot), PTO2_ALIGN_SIZE); } @@ -81,91 +162,24 @@ void ready_queue_destroy(PTO2ReadyQueue *queue) { queue->slots = nullptr; } -// ============================================================================= -// Scheduler -// ============================================================================= - bool PTO2SchedulerState::RingSchedState::init_data_from_layout(void *sm_dev_base, int32_t ring_id) { - // ring stores the device address of the SM ring header — pure offset - // arithmetic, no SM load. ring = pto2_sm_layout::ring_header_addr(sm_dev_base, ring_id); 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 (bind_ring + reset_for_reuse + - // fanin_count/active_mask zero) lives in PTO2SharedMemoryHandle:: - // init_header_per_ring so the AICPU performs it during SM reset; host - // prebuilt-arena init skips SM access here. - return true; } -void PTO2SchedulerState::RingSchedState::reset_for_reuse( - void *sm_dev_base, int32_t ring_id, std::atomic *orch_err -) { - ring = pto2_sm_layout::ring_header_addr(sm_dev_base, ring_id); - last_task_alive = 0; - advance_lock.store(0, std::memory_order_relaxed); - dep_pool.reset_for_reuse(orch_err); -#if SIMPLER_DFX - dep_pool_snapshot_tail.store(1, std::memory_order_relaxed); - dep_pool_snapshot_top.store(1, std::memory_order_relaxed); -#endif -} - void PTO2SchedulerState::RingSchedState::destroy() { ring = nullptr; } -PTO2SchedulerLayout PTO2SchedulerState::reserve_layout(DeviceArena &arena, int32_t dep_pool_capacity) { - int32_t dep_pool_capacities[PTO2_MAX_RING_DEPTH]; - int32_t task_window_sizes[PTO2_MAX_RING_DEPTH]; - for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { - dep_pool_capacities[r] = dep_pool_capacity; - task_window_sizes[r] = PTO2_TASK_WINDOW_SIZE; - } - return reserve_layout(arena, dep_pool_capacities, task_window_sizes); -} - -PTO2SchedulerLayout -PTO2SchedulerState::reserve_layout(DeviceArena &arena, const int32_t dep_pool_capacities[PTO2_MAX_RING_DEPTH]) { - int32_t task_window_sizes[PTO2_MAX_RING_DEPTH]; - for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { - task_window_sizes[r] = PTO2_TASK_WINDOW_SIZE; - } - return reserve_layout(arena, dep_pool_capacities, task_window_sizes); -} - -PTO2SchedulerLayout PTO2SchedulerState::reserve_layout( - DeviceArena &arena, const int32_t dep_pool_capacities[PTO2_MAX_RING_DEPTH], - const int32_t task_window_sizes[PTO2_MAX_RING_DEPTH] -) { - (void)task_window_sizes; +PTO2SchedulerLayout PTO2SchedulerState::reserve_layout(DeviceArena &arena, int32_t) { PTO2SchedulerLayout layout{}; layout.ready_queue_capacity = PTO2_READY_QUEUE_SIZE; - for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { - layout.dep_pool_capacities[r] = dep_pool_capacities[r]; - } + layout.spsc_capacity = PTO2_WRIRING_QUEUE_SIZE; - for (int i = 0; i < PTO2_NUM_RESOURCE_SHAPES; i++) { + 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); - } - for (int i = 0; i < PTO2_NUM_RESOURCE_SHAPES; i++) { - layout.off_ready_sync_queue_slots[i] = ready_queue_reserve_layout(arena, PTO2_READY_QUEUE_SIZE); - } layout.off_dummy_ready_queue_slots = ready_queue_reserve_layout(arena, PTO2_READY_QUEUE_SIZE); - for (int i = 0; i < PTO2_NUM_RESOURCE_SHAPES; i++) { - 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); - for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { - // Force a cache-line base so Orch-side dep_pool writes do not invalidate - // adjacent multi-threaded regions like ready_queue.slots. - layout.off_dep_pool_entries[r] = - arena.reserve(static_cast(dep_pool_capacities[r]) * sizeof(PTO2DepListEntry), PTO2_ALIGN_SIZE); - } + layout.off_pending_spsc_buffer = PTO2SpscQueue::reserve_layout(arena, PTO2_WRIRING_QUEUE_SIZE); return layout; } @@ -174,437 +188,62 @@ bool PTO2SchedulerState::init_data_from_layout( ) { PTO2SchedulerState *sched = this; sched->sm_header = reinterpret_cast(sm_dev_base); -#if SIMPLER_SCHED_PROFILING - sched->tasks_completed.store(0, std::memory_order_relaxed); - sched->tasks_consumed.store(0, std::memory_order_relaxed); -#endif - for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { - if (!sched->ring_sched_states[r].init_data_from_layout(sm_dev_base, r)) { - return false; - } - } + for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) + if (!sched->ring_sched_states[r].init_data_from_layout(sm_dev_base, r)) return false; - for (int i = 0; i < PTO2_NUM_RESOURCE_SHAPES; i++) { + for (int i = 0; i < PTO2_NUM_RESOURCE_SHAPES; i++) if (!ready_queue_init_data_from_layout( &sched->ready_queues[i], arena, layout.off_ready_queue_slots[i], layout.ready_queue_capacity - )) { + )) return false; - } - } - for (int i = 0; i < PTO2_NUM_RESOURCE_SHAPES; i++) { - if (!ready_queue_init_data_from_layout( - &sched->ready_sync_queues[i], arena, layout.off_ready_sync_queue_slots[i], layout.ready_queue_capacity - )) { - return false; - } - } if (!ready_queue_init_data_from_layout( &sched->dummy_ready_queue, arena, layout.off_dummy_ready_queue_slots, layout.ready_queue_capacity - )) { + )) return false; - } - for (int i = 0; i < PTO2_NUM_RESOURCE_SHAPES; i++) { - if (!ready_queue_init_data_from_layout( - &sched->early_dispatch_queues[i], arena, layout.off_early_dispatch_queue_slots[i], - PTO2_EARLY_DISPATCH_QUEUE_SIZE - )) { - return false; - } - } - if (!ready_queue_init_data_from_layout( - &sched->early_sync_start_queue, arena, layout.off_early_sync_start_queue_slots, - PTO2_EARLY_DISPATCH_QUEUE_SIZE - )) { + + if (!sched->wiring.queue.init_data_from_layout(arena, layout.off_pending_spsc_buffer, layout.spsc_capacity)) return false; - } - auto *orch_err = pto2_sm_layout::orch_error_code_addr(sm_dev_base); - for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { - auto *dep_entries = static_cast(arena.region_ptr(layout.off_dep_pool_entries[r])); - memset(dep_entries, 0, static_cast(layout.dep_pool_capacities[r]) * sizeof(PTO2DepListEntry)); - sched->ring_sched_states[r].dep_pool.init(dep_entries, layout.dep_pool_capacities[r], orch_err); - } + sched->wiring.backoff_counter = 0; return true; } -void PTO2SchedulerState::reset_for_reuse(const PTO2SchedulerLayout &layout, void *sm_dev_base) { - PTO2SchedulerState *sched = this; - sched->sm_header = reinterpret_cast(sm_dev_base); -#if SIMPLER_SCHED_PROFILING - sched->tasks_completed.store(0, std::memory_order_relaxed); - sched->tasks_consumed.store(0, std::memory_order_relaxed); -#endif - - auto *orch_err = pto2_sm_layout::orch_error_code_addr(sm_dev_base); - for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { - sched->ring_sched_states[r].reset_for_reuse(sm_dev_base, r, orch_err); - } - - for (int i = 0; i < PTO2_NUM_RESOURCE_SHAPES; i++) { - sched->ready_queues[i].reset_for_reuse(); - } - for (int i = 0; i < PTO2_NUM_RESOURCE_SHAPES; i++) { - sched->ready_sync_queues[i].reset_for_reuse(); - } - sched->dummy_ready_queue.reset_for_reuse(); - for (int i = 0; i < PTO2_NUM_RESOURCE_SHAPES; i++) { - sched->early_dispatch_queues[i].reset_for_reuse(); - } - sched->early_sync_start_queue.reset_for_reuse(); - - sched->async_wait_list.reset_for_reuse(); - (void)layout; -} - void PTO2SchedulerState::wire_arena_pointers(const PTO2SchedulerLayout &layout, DeviceArena &arena) { PTO2SchedulerState *sched = this; - for (int i = 0; i < PTO2_NUM_RESOURCE_SHAPES; i++) { + for (int i = 0; i < PTO2_NUM_RESOURCE_SHAPES; i++) ready_queue_wire_arena_pointers(&sched->ready_queues[i], arena, layout.off_ready_queue_slots[i]); - } - for (int i = 0; i < PTO2_NUM_RESOURCE_SHAPES; i++) { - ready_queue_wire_arena_pointers(&sched->ready_sync_queues[i], arena, layout.off_ready_sync_queue_slots[i]); - } ready_queue_wire_arena_pointers(&sched->dummy_ready_queue, arena, layout.off_dummy_ready_queue_slots); - for (int i = 0; i < PTO2_NUM_RESOURCE_SHAPES; i++) { - ready_queue_wire_arena_pointers( - &sched->early_dispatch_queues[i], arena, layout.off_early_dispatch_queue_slots[i] - ); - } - ready_queue_wire_arena_pointers(&sched->early_sync_start_queue, arena, layout.off_early_sync_start_queue_slots); - for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { - sched->ring_sched_states[r].dep_pool.base = - static_cast(arena.region_ptr(layout.off_dep_pool_entries[r])); - } + sched->wiring.queue.wire_arena_pointers(arena, layout.off_pending_spsc_buffer); } void PTO2SchedulerState::destroy() { PTO2SchedulerState *sched = this; - for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { + for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) sched->ring_sched_states[r].destroy(); - sched->ring_sched_states[r].dep_pool.base = nullptr; - } - for (int i = 0; i < PTO2_NUM_RESOURCE_SHAPES; i++) { + sched->wiring.queue.destroy(); + for (int i = 0; i < PTO2_NUM_RESOURCE_SHAPES; i++) ready_queue_destroy(&sched->ready_queues[i]); - } - for (int i = 0; i < PTO2_NUM_RESOURCE_SHAPES; i++) { - ready_queue_destroy(&sched->ready_sync_queues[i]); - } ready_queue_destroy(&sched->dummy_ready_queue); - for (int i = 0; i < PTO2_NUM_RESOURCE_SHAPES; i++) { - ready_queue_destroy(&sched->early_dispatch_queues[i]); - } - ready_queue_destroy(&sched->early_sync_start_queue); -} - -// ============================================================================= -// Orchestrator -// ============================================================================= - -PTO2OrchestratorLayout PTO2OrchestratorState::reserve_layout( - DeviceArena &arena, const int32_t task_window_sizes[PTO2_MAX_RING_DEPTH], int32_t dep_pool_capacity -) { - int32_t dep_pool_capacities[PTO2_MAX_RING_DEPTH]; - for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { - dep_pool_capacities[r] = dep_pool_capacity; - } - return reserve_layout(arena, task_window_sizes, dep_pool_capacities); } -PTO2OrchestratorLayout PTO2OrchestratorState::reserve_layout( - DeviceArena &arena, const int32_t task_window_sizes[PTO2_MAX_RING_DEPTH], - const int32_t dep_pool_capacities[PTO2_MAX_RING_DEPTH] -) { - PTO2OrchestratorLayout layout{}; - // scope_tasks holds every task in the open scope across all rings, so its cap - // is the real in-flight budget = sum of the (runtime) per-ring windows. Using - // the compile-time PTO2_SCOPE_TASKS_CAP instead under-sized the buffer when - // ring_task_window was enlarged past the default (premature SCOPE_TASKS_OVERFLOW) - // and over-allocated it when shrunk. See issue #1188. - // - // Accumulate in int64: each window is validated <= INT32_MAX individually, but - // the sum of PTO2_MAX_RING_DEPTH windows can exceed it — a bare int32 sum would - // wrap to a negative/undersized cap. Bound the result before narrowing. - int64_t scope_tasks_cap = 0; - for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { - always_assert(task_window_sizes[r] > 0); - scope_tasks_cap += task_window_sizes[r]; - } - always_assert(scope_tasks_cap <= std::numeric_limits::max()); - layout.scope_tasks_cap = static_cast(scope_tasks_cap); - layout.scope_stack_capacity = PTO2_MAX_SCOPE_DEPTH; - for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { - layout.dep_pool_capacities[r] = dep_pool_capacities[r]; - } - - for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { - const size_t fanin_pool_bytes = - PTO2_ALIGN_UP(static_cast(dep_pool_capacities[r]) * sizeof(PTO2FaninSpillEntry), PTO2_ALIGN_SIZE); - layout.off_fanin_pool[r] = arena.reserve(fanin_pool_bytes, PTO2_ALIGN_SIZE); - - always_assert(task_window_sizes[r] > 0 && (task_window_sizes[r] & (task_window_sizes[r] - 1)) == 0); - const size_t seen_epoch_bytes = - PTO2_ALIGN_UP(static_cast(task_window_sizes[r]) * sizeof(uint32_t), PTO2_ALIGN_SIZE); - layout.off_fanin_seen_epoch[r] = arena.reserve(seen_epoch_bytes, PTO2_ALIGN_SIZE); - } - layout.off_scope_tasks = - arena.reserve(static_cast(layout.scope_tasks_cap) * sizeof(uintptr_t), alignof(PTO2TaskSlotState *)); - layout.off_scope_begins = - arena.reserve(static_cast(layout.scope_stack_capacity) * sizeof(int32_t), alignof(int32_t)); - layout.tensor_map = PTO2TensorMap::reserve_layout_default(arena, task_window_sizes); - return layout; -} - -bool PTO2OrchestratorState::init_data_from_layout( - const PTO2OrchestratorLayout &layout, DeviceArena &arena, void *sm_dev_base, void *gm_heap, uint64_t heap_size, - uint64_t task_window_size -) { - uint64_t heap_sizes[PTO2_MAX_RING_DEPTH]; - uint64_t task_window_sizes[PTO2_MAX_RING_DEPTH]; - for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { - heap_sizes[r] = heap_size; - task_window_sizes[r] = task_window_size; - } - return init_data_from_layout(layout, arena, sm_dev_base, gm_heap, heap_sizes, task_window_sizes); -} - -bool PTO2OrchestratorState::init_data_from_layout( - const PTO2OrchestratorLayout &layout, DeviceArena &arena, void *sm_dev_base, void *gm_heap, - const uint64_t heap_sizes[PTO2_MAX_RING_DEPTH], const uint64_t task_window_sizes[PTO2_MAX_RING_DEPTH] -) { - auto *orch = this; - *orch = PTO2OrchestratorState{}; - - orch->sm_header = reinterpret_cast(sm_dev_base); - orch->gm_heap_base = gm_heap; - uint64_t total_heap_size = 0; - if (!sum_ring_heap_sizes(heap_sizes, &total_heap_size)) { - return false; - } - orch->gm_heap_size = total_heap_size; - orch->fatal = false; - - auto *orch_err = pto2_sm_layout::orch_error_code_addr(sm_dev_base); - uint64_t heap_offset = 0; - for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { - void *ring_heap_base = reinterpret_cast(gm_heap) + heap_offset; - auto *task_descs_dev = pto2_sm_layout::ring_task_descriptors_addr(sm_dev_base, task_window_sizes, r); - auto *slot_states_dev = pto2_sm_layout::ring_slot_states_addr(sm_dev_base, task_window_sizes, r); - auto *cur_idx_dev = pto2_sm_layout::ring_current_task_index_addr(sm_dev_base, r); - auto *last_alive_dev = pto2_sm_layout::ring_last_task_alive_addr(sm_dev_base, r); - - orch->rings[r].task_allocator.init( - task_descs_dev, static_cast(task_window_sizes[r]), cur_idx_dev, last_alive_dev, ring_heap_base, - heap_sizes[r], orch_err, slot_states_dev, 0, static_cast(r) - ); - heap_offset += heap_sizes[r]; - - const size_t fanin_pool_bytes = PTO2_ALIGN_UP( - static_cast(layout.dep_pool_capacities[r]) * sizeof(PTO2FaninSpillEntry), PTO2_ALIGN_SIZE - ); - auto *fanin_entries = static_cast(arena.region_ptr(layout.off_fanin_pool[r])); - memset(fanin_entries, 0, fanin_pool_bytes); - orch->rings[r].fanin_pool.init(fanin_entries, layout.dep_pool_capacities[r], orch_err); - - const size_t seen_epoch_bytes = PTO2_ALIGN_UP( - static_cast(layout.tensor_map.task_window_sizes[r]) * sizeof(uint32_t), PTO2_ALIGN_SIZE - ); - auto *seen_epoch = static_cast(arena.region_ptr(layout.off_fanin_seen_epoch[r])); - memset(seen_epoch, 0, seen_epoch_bytes); - orch->fanin_seen_epoch[r] = seen_epoch; - } - - if (!orch->tensor_map.init_data_from_layout(layout.tensor_map, arena)) { - return false; - } - - orch->scope_tasks_size = 0; - orch->scope_tasks_capacity = layout.scope_tasks_cap; - orch->scope_stack_top = -1; - orch->scope_stack_capacity = layout.scope_stack_capacity; - orch->manual_begin_depth = PTO2_MAX_SCOPE_DEPTH; - - return true; -} - -bool PTO2OrchestratorState::reset_for_reuse( - const PTO2OrchestratorLayout &layout, void *sm_dev_base, void *gm_heap, - const uint64_t heap_sizes[PTO2_MAX_RING_DEPTH], const uint64_t task_window_sizes[PTO2_MAX_RING_DEPTH] -) { - auto *orch = this; - orch->sm_header = reinterpret_cast(sm_dev_base); - orch->gm_heap_base = gm_heap; - uint64_t total_heap_size = 0; - if (!sum_ring_heap_sizes(heap_sizes, &total_heap_size)) { - return false; - } - orch->gm_heap_size = total_heap_size; - orch->fatal = false; - orch->inline_completed_tasks = 0; - - uint32_t next_epoch = orch->fanin_seen_current_epoch + 1; - if (next_epoch == 0) { - next_epoch = 1; - for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { - memset( - orch->fanin_seen_epoch[r], 0, - static_cast(layout.tensor_map.task_window_sizes[r]) * sizeof(uint32_t) - ); - } - } - orch->fanin_seen_current_epoch = next_epoch; - - auto *orch_err = pto2_sm_layout::orch_error_code_addr(sm_dev_base); - uint64_t heap_offset = 0; - for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { - void *ring_heap_base = reinterpret_cast(gm_heap) + heap_offset; - auto *task_descs_dev = pto2_sm_layout::ring_task_descriptors_addr(sm_dev_base, task_window_sizes, r); - auto *slot_states_dev = pto2_sm_layout::ring_slot_states_addr(sm_dev_base, task_window_sizes, r); - auto *cur_idx_dev = pto2_sm_layout::ring_current_task_index_addr(sm_dev_base, r); - auto *last_alive_dev = pto2_sm_layout::ring_last_task_alive_addr(sm_dev_base, r); - - orch->rings[r].task_allocator.init( - task_descs_dev, static_cast(task_window_sizes[r]), cur_idx_dev, last_alive_dev, ring_heap_base, - heap_sizes[r], orch_err, slot_states_dev, 0, static_cast(r) - ); - heap_offset += heap_sizes[r]; - orch->rings[r].fanin_pool.reset_for_reuse(orch_err); - } - - orch->tensor_map.reset_for_reuse(layout.tensor_map); - orch->scope_tasks_size = 0; - orch->scope_tasks_capacity = layout.scope_tasks_cap; - orch->scope_stack_top = -1; - orch->scope_stack_capacity = layout.scope_stack_capacity; - orch->manual_begin_depth = PTO2_MAX_SCOPE_DEPTH; - orch->total_cluster_count = 0; - orch->total_aiv_count = 0; -#if SIMPLER_DFX - orch->tasks_submitted = 0; - orch->buffers_allocated = 0; - orch->bytes_allocated = 0; -#endif - return true; -} - -void PTO2OrchestratorState::wire_arena_pointers( - const PTO2OrchestratorLayout &layout, DeviceArena &arena, PTO2SchedulerState *scheduler_arg -) { - auto *orch = this; - for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { - orch->rings[r].fanin_pool.base = static_cast(arena.region_ptr(layout.off_fanin_pool[r])); - orch->fanin_seen_epoch[r] = static_cast(arena.region_ptr(layout.off_fanin_seen_epoch[r])); - } - orch->tensor_map.wire_arena_pointers(layout.tensor_map, arena); - orch->scope_tasks = static_cast(arena.region_ptr(layout.off_scope_tasks)); - orch->scope_begins = static_cast(arena.region_ptr(layout.off_scope_begins)); - orch->scheduler = scheduler_arg; -} - -void PTO2OrchestratorState::destroy() { - auto *orch = this; - orch->tensor_map.destroy(); - for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { - orch->rings[r].fanin_pool.base = nullptr; - orch->fanin_seen_epoch[r] = nullptr; - } - orch->scope_tasks = nullptr; - orch->scope_begins = nullptr; -} - -void PTO2OrchestratorState::set_scheduler(PTO2SchedulerState *scheduler) { this->scheduler = scheduler; } - -// ============================================================================= -// Top-level runtime arena -// ============================================================================= - -PTO2RuntimeArenaLayout -runtime_reserve_layout(DeviceArena &arena, uint64_t task_window_size, int32_t dep_pool_capacity) { - 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); -} - -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] -) { - PTO2RuntimeArenaLayout layout{}; - - for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { - layout.sizing.task_window_sizes[r] = task_window_sizes[r]; - layout.sizing.heap_sizes[r] = heap_sizes[r]; - layout.sizing.dep_pool_capacities[r] = dep_pool_capacities[r]; - } - - layout.offsets.off_sm_handle = arena.reserve(sizeof(PTO2SharedMemoryHandle), alignof(PTO2SharedMemoryHandle)); - int32_t task_window_sizes_i32[PTO2_MAX_RING_DEPTH]; - for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { - task_window_sizes_i32[r] = static_cast(task_window_sizes[r]); - } - layout.offsets.orch = PTO2OrchestratorState::reserve_layout(arena, task_window_sizes_i32, dep_pool_capacities); - layout.offsets.sched = PTO2SchedulerState::reserve_layout(arena, dep_pool_capacities, task_window_sizes_i32); - layout.offsets.off_runtime = arena.reserve(sizeof(PTO2Runtime), PTO2_ALIGN_SIZE); - layout.offsets.off_mailbox = arena.reserve(sizeof(AICoreCompletionMailbox), alignof(AICoreCompletionMailbox)); - - layout.offsets.arena_size = arena.total_size(); - return layout; -} - -PTO2Runtime *runtime_init_data_from_layout( - DeviceArena &arena, const PTO2RuntimeArenaLayout &layout, PTO2RuntimeMode mode, void *sm_dev_base, - uint64_t /*sm_size*/, void *gm_heap_dev_base, uint64_t heap_size -) { - uint64_t heap_sizes[PTO2_MAX_RING_DEPTH]; +void PTO2SchedulerState::reset_for_reuse(void *sm_dev_base) { + PTO2SchedulerState *sched = this; + sched->sm_header = reinterpret_cast(sm_dev_base); for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { - heap_sizes[r] = heap_size; - } - return runtime_init_data_from_layout(arena, layout, mode, sm_dev_base, 0, gm_heap_dev_base, heap_sizes); -} - -PTO2Runtime *runtime_init_data_from_layout( - DeviceArena &arena, const PTO2RuntimeArenaLayout &layout, PTO2RuntimeMode mode, void *sm_dev_base, - uint64_t /*sm_size*/, void *gm_heap_dev_base, const uint64_t heap_sizes[PTO2_MAX_RING_DEPTH] -) { - PTO2Runtime *rt = static_cast(arena.region_ptr(layout.offsets.off_runtime)); - memset(rt, 0, sizeof(*rt)); - - auto *sm_wrap = static_cast(arena.region_ptr(layout.offsets.off_sm_handle)); - memset(sm_wrap, 0, sizeof(*sm_wrap)); - - // rt->ops is filled by the AICPU at boot. - rt->mode = mode; - rt->gm_heap = gm_heap_dev_base; - uint64_t total_heap_size = 0; - if (!sum_ring_heap_sizes(heap_sizes, &total_heap_size)) { - return nullptr; + sched->ring_sched_states[r].ring = pto2_sm_layout::ring_header_addr(sm_dev_base, r); + sched->ring_sched_states[r].last_task_alive = 0; + sched->ring_sched_states[r].advance_lock.store(0, std::memory_order_relaxed); } - rt->gm_heap_size = total_heap_size; - rt->gm_heap_owned = false; - rt->total_cycles = 0; - - if (!rt->orchestrator.init_data_from_layout( - layout.offsets.orch, arena, sm_dev_base, gm_heap_dev_base, heap_sizes, layout.sizing.task_window_sizes - )) { - return nullptr; - } - if (!rt->scheduler.init_data_from_layout(layout.offsets.sched, arena, sm_dev_base)) { - return nullptr; - } - - auto *mailbox = static_cast(arena.region_ptr(layout.offsets.off_mailbox)); - memset(mailbox, 0, sizeof(*mailbox)); - - return rt; + for (int i = 0; i < PTO2_NUM_RESOURCE_SHAPES; i++) + sched->ready_queues[i].reset_for_reuse(); + sched->dummy_ready_queue.reset_for_reuse(); + sched->wiring.queue.reset_for_reuse(); + sched->wiring.backoff_counter = 0; + sched->wiring.orch_needs_drain.store(false, std::memory_order_relaxed); + sched->async_wait_list.reset_for_reuse(); } - void runtime_wire_arena_pointers(DeviceArena &arena, const PTO2RuntimeArenaLayout &layout, PTO2Runtime *rt) { rt->sm_handle = static_cast(arena.region_ptr(layout.offsets.off_sm_handle)); rt->aicore_mailbox = static_cast(arena.region_ptr(layout.offsets.off_mailbox)); @@ -612,37 +251,18 @@ void runtime_wire_arena_pointers(DeviceArena &arena, const PTO2RuntimeArenaLayou rt->scheduler.wire_arena_pointers(layout.offsets.sched, arena); } -bool runtime_reset_for_reuse(DeviceArena &arena, const PTO2RuntimeArenaLayout &layout, PTO2Runtime *rt) { - (void)arena; - if (rt == nullptr) { - return false; - } +bool runtime_reset_for_reuse(DeviceArena &, const PTO2RuntimeArenaLayout &, PTO2Runtime *rt) { + if (rt == nullptr) return false; rt->pending_scope_mode = PTO2ScopeMode::AUTO; rt->total_cycles = 0; rt->gm_heap_owned = false; - uint64_t total_heap_size = 0; - if (!sum_ring_heap_sizes(layout.sizing.heap_sizes, &total_heap_size)) { - return false; - } - rt->gm_heap_size = total_heap_size; + void *sm_dev_base = rt->sm_handle ? rt->sm_handle->sm_base : nullptr; + if (sm_dev_base == nullptr) return false; - if (!rt->orchestrator.reset_for_reuse( - layout.offsets.orch, rt->sm_handle->sm_base, rt->gm_heap, layout.sizing.heap_sizes, - layout.sizing.task_window_sizes - )) { - return false; - } - rt->scheduler.reset_for_reuse(layout.offsets.sched, rt->sm_handle->sm_base); - return true; -} + rt->orchestrator.reset_for_reuse(); + rt->scheduler.reset_for_reuse(sm_dev_base); -void runtime_destroy(PTO2Runtime *rt, DeviceArena & /*arena*/) { - // Arena buffer is pooled across runs by DeviceRunner — never freed here. - if (!rt) return; - rt->scheduler.destroy(); - rt->orchestrator.destroy(); - rt->aicore_mailbox = nullptr; - rt->sm_handle = nullptr; + return true; } diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/shared/pto_shared_memory.cpp b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/shared/pto_shared_memory.cpp index 2ebeb42edc..41914162f3 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/shared/pto_shared_memory.cpp +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/shared/pto_shared_memory.cpp @@ -8,77 +8,65 @@ * See LICENSE in the root of the software repository for the full text of the License. * ----------------------------------------------------------------------------------------------------------- */ -/** - * PTO Runtime2 - Shared Memory Implementation - * - * Implements shared memory allocation, initialization, and management - * for Orchestrator-Scheduler communication. - * - * Based on: docs/RUNTIME_LOGIC.md - */ #include "pto_shared_memory.h" + #include #include #include -#include "common/unified_log.h" -// ============================================================================= -// Size Calculation -// ============================================================================= +#include "common/unified_log.h" uint64_t PTO2SharedMemoryHandle::calculate_size(uint64_t task_window_size) { uint64_t task_window_sizes[PTO2_MAX_RING_DEPTH]; - for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { + for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) task_window_sizes[r] = task_window_size; - } return calculate_size_per_ring(task_window_sizes); } uint64_t PTO2SharedMemoryHandle::calculate_size_per_ring(const uint64_t task_window_sizes[PTO2_MAX_RING_DEPTH]) { - // Total SM size = offset just past the last ring, from the single source of - // truth for the layout (pto2_sm_layout::ring_segment_offsets). - return pto2_sm_layout::ring_segment_offsets(task_window_sizes, PTO2_MAX_RING_DEPTH - 1).end; -} + uint64_t size = 0; -// ============================================================================= -// Creation and Destruction -// ============================================================================= + // Header (aligned to cache line) + size += PTO2_ALIGN_UP(sizeof(PTO2SharedMemoryHeader), PTO2_ALIGN_SIZE); -void PTO2SharedMemoryHandle::setup_pointers_per_ring(const uint64_t task_window_sizes[PTO2_MAX_RING_DEPTH]) { - char *base = (char *)sm_base; - header = (PTO2SharedMemoryHeader *)base; - - // Per-ring descriptors / payloads / slot_states — offsets from the single - // source of truth (pto2_sm_layout::ring_segment_offsets), so this setup and - // the device-address helpers cannot drift. + // Per-ring task descriptors and payloads for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { - auto off = pto2_sm_layout::ring_segment_offsets(task_window_sizes, r); - auto &ring = header->rings[r]; - ring.task_descriptors = (PTO2TaskDescriptor *)(base + off.descriptors); - ring.task_payloads = (PTO2TaskPayload *)(base + off.payloads); - ring.slot_states = (PTO2TaskSlotState *)(base + off.slot_states); + size += PTO2_ALIGN_UP(task_window_sizes[r] * sizeof(PTO2TaskDescriptor), PTO2_ALIGN_SIZE); + size += PTO2_ALIGN_UP(task_window_sizes[r] * sizeof(PTO2TaskPayload), PTO2_ALIGN_SIZE); + size += PTO2_ALIGN_UP(task_window_sizes[r] * sizeof(PTO2TaskSlotState), PTO2_ALIGN_SIZE); + size += PTO2_ALIGN_UP(task_window_sizes[r] * sizeof(std::atomic), PTO2_ALIGN_SIZE); } + + return size; } -void PTO2SharedMemoryHandle::setup_pointers(uint64_t task_window_size) { - uint64_t task_window_sizes[PTO2_MAX_RING_DEPTH]; - for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { - task_window_sizes[r] = task_window_size; - } - setup_pointers_per_ring(task_window_sizes); +PTO2SharedMemoryHandle *PTO2SharedMemoryHandle::create_and_init_default(DeviceArena &arena) { + const uint64_t buffer_size = calculate_size(PTO2_TASK_WINDOW_SIZE); + const size_t off_handle = arena.reserve(sizeof(PTO2SharedMemoryHandle), alignof(PTO2SharedMemoryHandle)); + const size_t off_buffer = arena.reserve(static_cast(buffer_size), PTO2_ALIGN_SIZE); + if (arena.commit() == nullptr) return nullptr; + + auto *handle = static_cast(arena.region_ptr(off_handle)); + memset(handle, 0, sizeof(*handle)); + void *buffer = arena.region_ptr(off_buffer); + memset(buffer, 0, static_cast(buffer_size)); + if (!handle->init(buffer, buffer_size, PTO2_TASK_WINDOW_SIZE, PTO2_HEAP_SIZE)) return nullptr; + return handle; } bool PTO2SharedMemoryHandle::init( void *sm_base_arg, uint64_t sm_size_arg, uint64_t task_window_size, uint64_t heap_size ) { - uint64_t task_window_sizes[PTO2_MAX_RING_DEPTH]; - uint64_t heap_sizes[PTO2_MAX_RING_DEPTH]; - for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { - task_window_sizes[r] = task_window_size; - heap_sizes[r] = heap_size; - } - return init_per_ring(sm_base_arg, sm_size_arg, task_window_sizes, heap_sizes); + if (!sm_base_arg || sm_size_arg == 0) return false; + if (sm_size_arg < calculate_size(task_window_size)) return false; + + sm_base = sm_base_arg; + sm_size = sm_size_arg; + is_owner = false; + setup_pointers(task_window_size); + init_header(task_window_size, heap_size); + return true; } bool PTO2SharedMemoryHandle::init_per_ring( @@ -96,20 +84,6 @@ bool PTO2SharedMemoryHandle::init_per_ring( return true; } -PTO2SharedMemoryHandle *PTO2SharedMemoryHandle::create_and_init_default(DeviceArena &arena) { - const uint64_t buffer_size = calculate_size(PTO2_TASK_WINDOW_SIZE); - const size_t off_handle = arena.reserve(sizeof(PTO2SharedMemoryHandle), alignof(PTO2SharedMemoryHandle)); - const size_t off_buffer = arena.reserve(static_cast(buffer_size), PTO2_ALIGN_SIZE); - if (arena.commit() == nullptr) return nullptr; - - auto *handle = static_cast(arena.region_ptr(off_handle)); - memset(handle, 0, sizeof(*handle)); - void *buffer = arena.region_ptr(off_buffer); - memset(buffer, 0, static_cast(buffer_size)); - if (!handle->init(buffer, buffer_size, PTO2_TASK_WINDOW_SIZE, PTO2_HEAP_SIZE)) return nullptr; - return handle; -} - void PTO2SharedMemoryHandle::destroy() { // Arena-owned wrappers (is_owner == false) are reclaimed by arena.release(); // calling destroy on them is a no-op so existing callers stay safe. @@ -119,11 +93,12 @@ void PTO2SharedMemoryHandle::destroy() { } } -// ============================================================================= -// Initialization -// ============================================================================= -// -// no need init data in pool, init pool data when used +void PTO2SharedMemoryHandle::print_layout() { + if (!header) return; + + for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) {} +} + void PTO2SharedMemoryHandle::init_header(uint64_t task_window_size, uint64_t heap_size) { uint64_t task_window_sizes[PTO2_MAX_RING_DEPTH]; uint64_t heap_sizes[PTO2_MAX_RING_DEPTH]; @@ -140,6 +115,16 @@ void PTO2SharedMemoryHandle::init_header_per_ring( // Per-ring flow control (start at 0) for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { header->rings[r].fc.init(); + // -1 = "no task completed yet"; first task to complete (local_id 0) + // will advance the watermark to 0. + header->rings[r].completed_watermark.store(-1, std::memory_order_relaxed); + // Shared memory is not guaranteed zero on device. The watermark-advance + // loop reads completion_flags for not-yet-completed slots, so stale + // non-zero bytes would prematurely advance completed_watermark and retire + // live slots. Zero the whole per-ring flag block (1 byte/slot, cheap). + __builtin_memset( + (void *)header->rings[r].completion_flags, 0, task_window_sizes[r] * sizeof(std::atomic) + ); } header->orchestrator_done.store(0, std::memory_order_relaxed); @@ -165,53 +150,44 @@ void PTO2SharedMemoryHandle::init_header_per_ring( header->sched_error_bitmap.store(0, std::memory_order_relaxed); header->sched_error_code.store(PTO2_ERROR_NONE, std::memory_order_relaxed); header->sched_error_thread.store(-1, std::memory_order_relaxed); - header->sched_stall_detail.store(PTO2_STALL_DETAIL_NONE, std::memory_order_relaxed); - header->sched_stall_completed.store(0, std::memory_order_relaxed); - header->sched_stall_total.store(0, std::memory_order_relaxed); - header->sched_stall_cnt_running.store(0, std::memory_order_relaxed); - header->sched_stall_cnt_ready.store(0, std::memory_order_relaxed); - header->sched_stall_cnt_waiting.store(0, std::memory_order_relaxed); - header->sched_stall_orch_done.store(0, std::memory_order_relaxed); - header->sched_stall_task_id.store(-1, std::memory_order_relaxed); - header->sched_stall_core.store(-1, std::memory_order_relaxed); - - // No per-slot loop: prepare_task resets each slot when it allocates it, and - // the scheduler only scans submitted task_ids [last_task_alive, - // current_task_index), so unsubmitted slots are never read. Per-boot reset - // is just the header fields above; per-slot state is set lazily at submit. + + // No per-slot loop: prepare_task() resets each slot when the allocator + // hands it out (bind_ring + reset_for_reuse + per-submit fields). The + // scheduler only scans submitted task_ids [last_task_alive, + // current_task_index), so unsubmitted slots are never read. Cost moves + // from O(sum(task_window_sizes)) every run to O(tasks actually + // submitted) — and stays on the device. Mirrors upstream #1199. } -// ============================================================================= -// Debug Utilities -// ============================================================================= +void PTO2SharedMemoryHandle::setup_pointers(uint64_t task_window_size) { + uint64_t task_window_sizes[PTO2_MAX_RING_DEPTH]; + for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) + task_window_sizes[r] = task_window_size; + setup_pointers_per_ring(task_window_sizes); +} -void PTO2SharedMemoryHandle::print_layout() { - if (!header) return; +void PTO2SharedMemoryHandle::setup_pointers_per_ring(const uint64_t task_window_sizes[PTO2_MAX_RING_DEPTH]) { + char *ptr = (char *)sm_base; - PTO2SharedMemoryHeader *h = header; + // Header + header = (PTO2SharedMemoryHeader *)ptr; + ptr += PTO2_ALIGN_UP(sizeof(PTO2SharedMemoryHeader), PTO2_ALIGN_SIZE); - LOG_INFO_V0("=== PTO2 Shared Memory Layout ==="); - LOG_INFO_V0("Base address: %p", sm_base); - LOG_INFO_V0("Total size: %" PRIu64 " bytes", h->total_size); - LOG_INFO_V0("Ring depth: %d", PTO2_MAX_RING_DEPTH); + // Per-ring task descriptors, payloads, and slot states for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { - LOG_INFO_V0("Ring %d:", r); - LOG_INFO_V0(" task_window_size: %" PRIu64, h->rings[r].task_window_size); - LOG_INFO_V0(" heap_size: %" PRIu64 " bytes", h->rings[r].heap_size); - LOG_INFO_V0( - " descriptors_off: %" PRIu64 " (0x%" PRIx64 ")", h->rings[r].task_descriptors_offset, - h->rings[r].task_descriptors_offset - ); - LOG_INFO_V0(" current_task_idx: %d", h->rings[r].fc.current_task_index.load(std::memory_order_acquire)); - LOG_INFO_V0(" last_task_alive: %d", h->rings[r].fc.last_task_alive.load(std::memory_order_acquire)); + auto &ring = header->rings[r]; + ring.task_descriptors = (PTO2TaskDescriptor *)ptr; + ptr += PTO2_ALIGN_UP(task_window_sizes[r] * sizeof(PTO2TaskDescriptor), PTO2_ALIGN_SIZE); + + ring.task_payloads = (PTO2TaskPayload *)ptr; + ptr += PTO2_ALIGN_UP(task_window_sizes[r] * sizeof(PTO2TaskPayload), PTO2_ALIGN_SIZE); + + ring.slot_states = (PTO2TaskSlotState *)ptr; + ptr += PTO2_ALIGN_UP(task_window_sizes[r] * sizeof(PTO2TaskSlotState), PTO2_ALIGN_SIZE); + + ring.completion_flags = (std::atomic *)ptr; + ptr += PTO2_ALIGN_UP(task_window_sizes[r] * sizeof(std::atomic), PTO2_ALIGN_SIZE); } - LOG_INFO_V0("orchestrator_done: %d", h->orchestrator_done.load(std::memory_order_acquire)); - LOG_INFO_V0("Error state:"); - LOG_INFO_V0(" orch_error_code: %d", h->orch_error_code.load(std::memory_order_relaxed)); - LOG_INFO_V0(" sched_error_bitmap: 0x%x", h->sched_error_bitmap.load(std::memory_order_relaxed)); - LOG_INFO_V0(" sched_error_code: %d", h->sched_error_code.load(std::memory_order_relaxed)); - LOG_INFO_V0(" sched_error_thread: %d", h->sched_error_thread.load(std::memory_order_relaxed)); - LOG_INFO_V0("================================"); } bool PTO2SharedMemoryHandle::validate() { @@ -220,9 +196,8 @@ bool PTO2SharedMemoryHandle::validate() { PTO2SharedMemoryHeader *h = header; - for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { + for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) if (!h->rings[r].fc.validate(this, r)) return false; - } return true; } diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/shared/pto_tensormap.cpp b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/shared/pto_tensormap.cpp index 3854099d0e..63a669b028 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/shared/pto_tensormap.cpp +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/shared/pto_tensormap.cpp @@ -8,20 +8,6 @@ * See LICENSE in the root of the software repository for the full text of the License. * ----------------------------------------------------------------------------------------------------------- */ -/** - * PTO Runtime2 - TensorMap Implementation - * - * Implements TensorMap with ring buffer pool, lazy invalidation, - * and chain truncation optimization. - * - * Key features: - * 1. O(1) insert at bucket head - * 2. O(valid_entries) lookup with chain truncation - * 3. Automatic stale entry cleanup during lookup - * 4. Periodic explicit cleanup for long chains - * - * Based on: docs/RUNTIME_LOGIC.md - */ #include "pto_tensormap.h" @@ -31,35 +17,19 @@ #include "common.h" #include "common/unified_log.h" -// ============================================================================= -// TensorMap Lookup Chain Length Statistics (compile-time toggle) -// ============================================================================= -#if SIMPLER_TENSORMAP_PROFILING -uint64_t g_lookup_chain_total = 0; -uint64_t g_lookup_count = 0; -int32_t g_lookup_chain_max = 0; -uint64_t g_lookup_overlap_checks = 0; -uint64_t g_lookup_overlap_hits = 0; -uint64_t g_insert_count = 0; -#endif - -// ============================================================================= -// Initialization and Destruction -// ============================================================================= - PTO2TensorMapLayout PTO2TensorMap::reserve_layout( DeviceArena &arena, int32_t new_num_buckets, int32_t new_pool_size, const int32_t new_task_window_sizes[PTO2_MAX_RING_DEPTH] ) { - // num_buckets must be a power of two for the hash truncation to work. - always_assert((new_num_buckets & (new_num_buckets - 1)) == 0); + // num_buckets must be a positive power of two for the hash truncation to + // work (0 passes the power-of-two test but makes __builtin_ctz UB). + always_assert(new_num_buckets > 0 && (new_num_buckets & (new_num_buckets - 1)) == 0); PTO2TensorMapLayout layout{}; layout.num_buckets = new_num_buckets; layout.pool_size = new_pool_size; - for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { + for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) layout.task_window_sizes[r] = new_task_window_sizes[r]; - } layout.off_buckets = arena.reserve( static_cast(new_num_buckets) * sizeof(PTO2TensorMapEntry *), alignof(PTO2TensorMapEntry *) @@ -102,9 +72,6 @@ bool PTO2TensorMap::init_data_from_layout(const PTO2TensorMapLayout &layout, Dev bucket_epochs_arena[i] = 0; } - // entry_pool: zero-init equivalent to the previous calloc(entry_pool, ...). - // The pool's persistent invariant after init is "bucket_index == -1 means - // not linked", set explicitly below. memset(entry_pool_arena, 0, static_cast(pool_size) * sizeof(PTO2TensorMapEntry)); for (int32_t i = 0; i < pool_size; i++) { entry_pool_arena[i].bucket_index = -1; @@ -137,27 +104,6 @@ bool PTO2TensorMap::init_data_from_layout(const PTO2TensorMapLayout &layout, Dev return true; } -void PTO2TensorMap::reset_for_reuse(const PTO2TensorMapLayout &layout) { - num_buckets = layout.num_buckets; - pool_size = layout.pool_size; - next_entry_idx = 0; - free_num = 0; - current_epoch++; - if (current_epoch == 0) { - current_epoch = 1; - memset(bucket_epochs, 0, static_cast(layout.num_buckets) * sizeof(uint32_t)); - for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { - memset(task_entry_head_epochs[r], 0, static_cast(layout.task_window_sizes[r]) * sizeof(uint32_t)); - } - } - - for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { - task_window_sizes[r] = layout.task_window_sizes[r]; - last_task_alives[r] = 0; - last_cleanup[r] = 0; - } -} - void PTO2TensorMap::wire_arena_pointers(const PTO2TensorMapLayout &layout, DeviceArena &arena) { buckets = static_cast(arena.region_ptr(layout.off_buckets)); bucket_epochs = static_cast(arena.region_ptr(layout.off_bucket_epochs)); @@ -169,88 +115,38 @@ void PTO2TensorMap::wire_arena_pointers(const PTO2TensorMapLayout &layout, Devic } } -void PTO2TensorMap::destroy() { - // Arena owns the backing memory; here we only forget our pointers so any - // stray post-destroy access trips a nullptr dereference instead of reading - // a recycled allocation. - buckets = nullptr; - bucket_epochs = nullptr; - entry_pool = nullptr; - free_entry_list = nullptr; +void PTO2TensorMap::reset_for_reuse() { + next_entry_idx = 0; + free_num = 0; for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { - task_entry_heads[r] = nullptr; - task_entry_head_epochs[r] = nullptr; - } -} - -// ============================================================================= -// Debug Utilities -// ============================================================================= - -void PTO2TensorMap::print_stats() { - int32_t valid = 0; - int32_t stale = 0; - int32_t empty_buckets = 0; - int32_t max_chain = 0; - int64_t total_chain = 0; - int32_t non_empty_buckets = 0; - - // Count entries - for (int32_t i = 0; i < pool_size; i++) { - if (entry_pool[i].bucket_index != -1) { - if (entry_valid(entry_pool[i])) { - valid++; - } else { - stale++; - } - } + last_task_alives[r] = 0; + last_cleanup[r] = 0; } - - // Count bucket stats - for (int32_t b = 0; b < num_buckets; b++) { - int32_t chain_len = 0; - auto cur_entry = buckets[b]; - - while (cur_entry != nullptr) { - chain_len++; - cur_entry = cur_entry->next_in_bucket; - } - - if (chain_len == 0) { - empty_buckets++; - } else { - non_empty_buckets++; - total_chain += chain_len; - if (chain_len > max_chain) { - max_chain = chain_len; - } + current_epoch++; + if (current_epoch == 0) { + current_epoch = 1; + for (int32_t i = 0; i < num_buckets; i++) + bucket_epochs[i] = 0; + for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { + for (int32_t i = 0; i < task_window_sizes[r]; i++) + task_entry_head_epochs[r][i] = 0; } } +} - LOG_INFO_V0("=== TensorMap Statistics ==="); - LOG_INFO_V0("Pool size: %d", pool_size); - LOG_INFO_V0("Pool next entry idx: %d", next_entry_idx); - LOG_INFO_V0("Pool free_num: %d", free_num); - LOG_INFO_V0("Num buckets: %d", num_buckets); - LOG_INFO_V0("Valid entries: %d", valid); - LOG_INFO_V0("Stale entries: %d", stale); - LOG_INFO_V0("Empty buckets: %d", empty_buckets); - LOG_INFO_V0("Max chain len: %d", max_chain); - LOG_INFO_V0("Avg chain len: %.2f", non_empty_buckets > 0 ? (float)total_chain / non_empty_buckets : 0); - for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { - LOG_INFO_V0("Last task alive[%d]: %d", r, last_task_alives[r]); - } - LOG_INFO_V0("============================"); +void PTO2TensorMap::destroy() { + buckets = nullptr; + entry_pool = nullptr; + free_entry_list = nullptr; + for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) + task_entry_heads[r] = nullptr; } int32_t PTO2TensorMap::valid_count() { int32_t count = 0; - for (int32_t i = 0; i < pool_size; i++) { - if (entry_pool[i].bucket_index != -1 && entry_valid(entry_pool[i])) { - count++; - } - } + for (int32_t i = 0; i < pool_size; i++) + if (entry_pool[i].bucket_index != -1 && entry_valid(entry_pool[i])) count++; return count; } @@ -268,27 +164,3 @@ void PTO2TensorMap::sync_tensormap(PTO2TaskId task_id, int32_t sm_last_task_aliv last_cleanup[ring_id] = sm_last_task_alive; } } - -// ============================================================================= -// TensorMap Lookup Profiling -// ============================================================================= -#if SIMPLER_TENSORMAP_PROFILING -PTO2TensorMapProfilingData pto2_tensormap_get_profiling() { - PTO2TensorMapProfilingData d; - d.lookup_chain_total = g_lookup_chain_total; - d.lookup_count = g_lookup_count; - d.lookup_chain_max = g_lookup_chain_max; - d.overlap_checks = g_lookup_overlap_checks; - d.overlap_hits = g_lookup_overlap_hits; - d.insert_count = g_insert_count; - - // Reset - g_lookup_chain_total = 0; - g_lookup_count = 0; - g_lookup_chain_max = 0; - g_lookup_overlap_checks = 0; - g_lookup_overlap_hits = 0; - g_insert_count = 0; - return d; -} -#endif