From d6c45130e9ca15038dd0154e6dd1e0a2987db8e1 Mon Sep 17 00:00:00 2001 From: Chao Wang <26245345+ChaoWao@users.noreply.github.com> Date: Sat, 11 Jul 2026 03:06:56 -0700 Subject: [PATCH] Sync: host_build_graph scheduler up a generation to tensormap_and_ringbuffer MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit hbg's scheduler was a copy of tmr@fb28fd71 and had fallen a full generation behind. This absorbs tmr's post-fork scheduler evolution while preserving hbg's host-orchestration invariants (single ring, no execution-time reclaim, flat Runtime, host-side orchestrator). Absorbed from tmr: - Remove PTO2LocalReadyBuffer local-first dispatch and its local_bufs threading (#1245); pop/dispatch now take PTO2ReadyQueue* directly. - Ready-claim completion protocol: ready_state / mark_completed / mark_any_subtask_deferred (replaces the any_subtask_deferred bool). - Rename speculative -> early_dispatch and remove the auto-chain (spec_chain_active/depth, PTO2_SPEC_CHAIN_MAX); move allow_early_resolve from payload onto slot_state. - Per-shape early_dispatch_queues[] (was a single queue) + ready_sync_queues[] Tier-0 sync_start staging (#1319); run_staging_order template; try_early_dispatch / early_dispatch_shape; has_any_free_slot predicates. - MIX subtasks placed per-core instead of uniform cluster placement (#1308). - Remove the on-device wiring subsystem (PTO2SpscQueue / WiringState / drain_wiring_queue / wire_task, #1263): the host orchestrator now wires fanout adjacency inline during submit and seeds readiness via the scheduler's push_ready_routed / route_ready_once, with no device-side wiring queue. - Single-round-trip / parallel AICore handshake (#1279/#1310): handshake_all_cores -> pre_handshake_init + per-thread handshake_partition + post_handshake_init, with the matching aicpu_executor boot restructure and AICore-worker protocol. Drop the now-dead two-round-trip Handshake fields aicpu_regs_ready / aicore_regs_ready (aligning the struct to tmr's). - Drop the per-run per-core memset in scheduler init/deinit (#1212); gate dep_gen init under SIMPLER_DFX; compile out pmu_active when SIMPLER_DFX is off; remove the redundant second completion poll; remove the dead device-orch core transition (handle_core_transition / reassign_cores_for_all_threads / orch_to_sched_). Remove vestigial device-orch synchronization: with host orchestration the orchestrator runs to completion on the host before the device boots, so the device-side orchestrator_done_ flag is always set before the scheduler threads start (they gate on runtime_init_ready_, whose release/acquire already publishes total_tasks_). Removed orchestrator_done_, the always-true completion guard, and the never-called wait_for_orchestration_done_before_dispatch / orchestration_done() carried over from tmr's concurrent device-orch model. Preserved host-orch divergences (not overwritten by tmr): PTO2_MAX_RING_DEPTH==1 with rings[]/ring_id removed; no advance_ring_pointers / reset_for_reuse / check_and_handle_consumed (completion via completed_tasks_, consumer wait keys on fanout_refcount); flat runtime->workers; hbg's swimlane behavior (no on_aicore_ack hook — the just-filled buffer drains via the completion-before-dispatch invariant and the next-rotation / run-end backstop flushes). Skipped as not applicable to hbg: stall-classification publishing (shared-memory lacks the sched_stall_* fields; last_task_alive semantics differ); per-device scheduler timeout (aicpu_device_config.h / get_scheduler_timeout_ms absent); runtime->dev.* Runtime nesting. get_reg_ptr weak host fallback: route_ready_once transitively ODR-uses the early-dispatch doorbell inline (ring_one_doorbell -> get_reg_ptr) in the host lib, but no core is gated during host graph-build so it never fires; a weak hidden fallback satisfies the loader (same pattern as get_sys_cnt_aicpu). Verified: all 8 runtime targets build -Werror; a2a3sim host_build_graph scene suite 10/10; a2a3 onboard host_build_graph suite 10 passed / 1 skipped (paged_attention exercises the new handshake, route_ready_once wiring, and ready-claim path). tensormap_and_ringbuffer and the a5 runtimes are untouched. --- .../aicore/aicore_executor.cpp | 64 +- .../host_build_graph/aicpu/aicpu_executor.cpp | 142 ++-- .../host_build_graph/docs/RUNTIME_LOGIC.md | 118 ++-- .../docs/SUBMIT_BY_CLUSTER.md | 2 +- .../host_build_graph/host/runtime_maker.cpp | 31 +- .../orchestrator_core/pto_orchestrator.cpp | 165 ++++- .../orchestrator_core/pto_runtime2.cpp | 8 - .../host_build_graph/runtime/pto_async_wait.h | 4 +- .../host_build_graph/runtime/pto_runtime2.h | 5 +- .../runtime/pto_runtime2_types.h | 127 ++-- .../host_build_graph/runtime/runtime.h | 18 +- .../runtime/scheduler/pto_scheduler.h | 651 ++++-------------- .../runtime/scheduler/scheduler_cold_path.cpp | 463 ++++++------- .../scheduler/scheduler_completion.cpp | 43 +- .../runtime/scheduler/scheduler_context.h | 139 ++-- .../runtime/scheduler/scheduler_dispatch.cpp | 632 ++++++++--------- .../runtime/scheduler/scheduler_types.h | 49 +- .../runtime/shared/pto_runtime2_init.cpp | 52 +- 18 files changed, 1264 insertions(+), 1449 deletions(-) diff --git a/src/a2a3/runtime/host_build_graph/aicore/aicore_executor.cpp b/src/a2a3/runtime/host_build_graph/aicore/aicore_executor.cpp index bbe9cf839..0cb67144b 100644 --- a/src/a2a3/runtime/host_build_graph/aicore/aicore_executor.cpp +++ b/src/a2a3/runtime/host_build_graph/aicore/aicore_executor.cpp @@ -47,15 +47,17 @@ __aicore__ __attribute__((always_inline)) static void execute_task(__gm__ PTO2Di * AICore main execution loop * * Implements the AICPU-AICore register-based dispatch protocol: - * 1. Wait for AICPU ready signal via handshake buffer - * 2. Report physical core ID and core type, signal AICore ready + * 1. Report physical core ID and core type, signal aicore_done (no AICPU wait) + * 2. Wait for the AICPU to open our register window (DATA_MAIN_BASE != 0) * 3. Cache per-core PTO2DispatchPayload pointer from hank->task * 4. Poll DATA_MAIN_BASE register for task dispatch until exit signal * - * AICPU writes &s_payload_per_core[i] to hank->task before setting - * aicpu_ready=1. AICore caches this pointer and reads function_bin_addr + - * args pointer from it on each dispatch. reg_val is a monotonically - * increasing task ID used only for dispatch signaling and ACK/FIN protocol. + * AICore reports on launch; the AICPU writes &s_payload_per_core[i] to + * hank->task and then opens the register window (DATA_MAIN_BASE = IDLE), which + * is itself the acknowledgement. AICore caches this pointer and reads + * function_bin_addr + args pointer from it on each dispatch. reg_val is a + * monotonically increasing task ID used only for dispatch signaling and + * ACK/FIN protocol. * * Profiling state (enable flag, L2 swimlane rotation channel) is published into the platform * via set_aicore_profiling_flag / set_aicore_l2_swimlane_ring at kernel entry — @@ -69,32 +71,36 @@ __aicore__ __attribute__((always_inline)) static void execute_task(__gm__ PTO2Di __aicore__ __attribute__((weak)) void aicore_execute(__gm__ Runtime *runtime, int block_idx, CoreType core_type) { __gm__ Handshake *my_hank = (__gm__ Handshake *)(&runtime->workers[block_idx]); - // Phase 1: Wait for AICPU initialization signal - while (my_hank->aicpu_ready == 0) { - dcci(my_hank, SINGLE_CACHE_LINE); - SPIN_WAIT_HINT(); - } - - // Phase 2: Report physical core ID, signal ready + // Phase 1: report physical core ID + core type and signal done in one write, + // with no wait for the AICPU — both fields are self-known. The AICPU opens + // this core's register window only after it observes aicore_done, so a single + // report suffices. The host clears aicore_done before this kernel launches, + // so the value the AICPU reads is this run's report, never a stale prior one. my_hank->physical_core_id = get_physical_core_id(); - OUT_OF_ORDER_STORE_BARRIER(); - my_hank->aicore_regs_ready = 1; - dcci(&my_hank->aicore_regs_ready, SINGLE_CACHE_LINE, CACHELINE_OUT); - while (my_hank->aicpu_regs_ready == 0) { - dcci(&my_hank->aicpu_regs_ready, SINGLE_CACHE_LINE); - SPIN_WAIT_HINT(); - } - // Report initial idle status via register - write_reg(RegId::COND, AICORE_IDLE_VALUE); - - // Phase 3: Report core type, signal ready my_hank->core_type = core_type; OUT_OF_ORDER_STORE_BARRIER(); my_hank->aicore_done = block_idx + 1; // Signal ready (use block_idx + 1 to avoid 0) - dcci(my_hank, SINGLE_CACHE_LINE, CACHELINE_OUT); - // Cache per-core dispatch payload pointer (set by AICPU before aicpu_ready) + // Phase 2: Wait for the AICPU to open our register window. A kernel launch + // resets DATA_MAIN_BASE to 0 (verified on a2a3 silicon); the AICPU writes + // DATA_MAIN_BASE = AICPU_IDLE_TASK_ID (non-zero) as it opens FAST_PATH, so a + // non-zero read means the window is open and reads/writes are valid. The + // AICPU runs assign_cores_to_threads (µs) between opening the window and the + // first dispatch, so this IDLE is observed long before any task_id lands — + // the poll cannot miss it and mistake a later task for the reset value. + // Window-open is the sync point for everything the AICPU publishes (task + // pointer, swimlane head): the AICPU writes those before opening the window. + while (read_reg(RegId::DATA_MAIN_BASE) == 0) { + SPIN_WAIT_HINT(); + } + // Report initial idle status via register (FAST_PATH is now open). + write_reg(RegId::COND, AICORE_IDLE_VALUE); + + // The AICPU writes task after observing our report (so our CACHELINE_OUT flush + // above cannot clobber it) and before opening the window; dcci to read its + // fresh value here. + dcci(my_hank, SINGLE_CACHE_LINE); __gm__ PTO2DispatchPayload *payload = reinterpret_cast<__gm__ PTO2DispatchPayload *>(my_hank->task); uint32_t enable_profiling_flag = get_aicore_profiling_flag(); @@ -103,9 +109,9 @@ __aicore__ __attribute__((weak)) void aicore_execute(__gm__ Runtime *runtime, in bool pmu_enabled = SIMPLER_GET_DFX_FLAG(enable_profiling_flag, SIMPLER_DFX_FLAG_PMU); // Per-core L2SwimlaneActiveHead channel. AICPU completes - // `l2_swimlane_aicpu_init` before writing `aicpu_ready = 1` in - // `handshake_all_cores`, and Phase 1 above has already observed - // `aicpu_ready == 1`, so the rotation-table slot is populated and the + // `l2_swimlane_aicpu_init` (in pre_handshake_init) before any thread opens a + // register window in `handshake_partition`, and Phase 2 above has already + // observed our window open, so the rotation-table slot is populated and the // first deref is safe here — off the dispatch→start critical path. __gm__ L2SwimlaneActiveHead *l2_swimlane_head = l2_swimlane_enabled ? get_l2_swimlane_aicore_head() : nullptr; // cached_buf_seq must start != AICPU's initial head.current_buf_seq (0) diff --git a/src/a2a3/runtime/host_build_graph/aicpu/aicpu_executor.cpp b/src/a2a3/runtime/host_build_graph/aicpu/aicpu_executor.cpp index e290edfa3..8e233fd96 100644 --- a/src/a2a3/runtime/host_build_graph/aicpu/aicpu_executor.cpp +++ b/src/a2a3/runtime/host_build_graph/aicpu/aicpu_executor.cpp @@ -84,11 +84,20 @@ struct AicpuExecutor { // ===== Thread management state ===== std::atomic thread_idx_{0}; - std::atomic initialized_{false}; std::atomic init_done_{false}; std::atomic init_failed_{false}; std::atomic finished_{false}; + // Parallel-handshake coordination (see AicpuExecutor::init). hs_setup_done_ + // is published by the leader once the shared pre-handshake setup is visible; + // hs_arrived_ is the barrier counting threads that finished their core slice. + // hs_thread_seq_ hands out a distinct [0, nthreads) index when the platform + // exposes no affinity idx (sim, where platform_aicpu_affinity_thread_idx() + // is -1 during init) so the threads don't all collapse to leader 0. + std::atomic hs_setup_done_{false}; + std::atomic hs_arrived_{0}; + std::atomic hs_thread_seq_{0}; + int32_t aicpu_thread_num_{0}; // ===== Task queue state (managed by scheduler ready queues) ===== @@ -115,42 +124,88 @@ static AicpuExecutor g_aicpu_executor; // ===== AicpuExecutor Method Implementations ===== int32_t AicpuExecutor::init(Runtime *runtime) { - bool expected = false; - if (!initialized_.compare_exchange_strong(expected, true, std::memory_order_acq_rel, std::memory_order_acquire)) { - return 0; - } - - LOG_INFO_V0("AicpuExecutor: Initializing"); - if (runtime == nullptr) { LOG_ERROR("runtime is nullptr"); init_failed_.store(true, std::memory_order_release); return -1; } - // Read execution parameters from runtime. The 0 → 1 fixup runs before the - // sched_thread_num_ derivation so a zero input doesn't leave the scheduler - // count at -1. - aicpu_thread_num_ = runtime->aicpu_thread_num; - if (aicpu_thread_num_ == 0) aicpu_thread_num_ = 1; - sched_thread_num_ = aicpu_thread_num_ - 1; - orch_to_sched_ = runtime->orch_to_sched; - - if (aicpu_thread_num_ < 1 || aicpu_thread_num_ > MAX_AICPU_THREADS) { - LOG_ERROR("Invalid aicpu_thread_num: %d", aicpu_thread_num_); + // All AICPU threads enter init. The per-core AICore handshake is the + // dominant preamble cost (serial MMIO, ~217 µs of ~283 µs for 72 cores), so + // it is parallelized: the leader (tidx 0) does the shared setup, every + // thread handshakes a disjoint slice of cores, then the leader finishes init + // after a barrier. Non-leaders spin on init_done_. + int32_t nthreads = runtime->aicpu_thread_num; + if (nthreads == 0) nthreads = 1; + if (nthreads < 1 || nthreads > MAX_AICPU_THREADS) { + LOG_ERROR("Invalid aicpu_thread_num: %d", nthreads); init_failed_.store(true, std::memory_order_release); return -1; } - - if (sched_ctx_.init(runtime, aicpu_thread_num_, sched_thread_num_, orch_to_sched_, get_platform_regs()) != 0) { - init_failed_.store(true, std::memory_order_release); + // Each thread needs a distinct index in [0, nthreads) to pick the leader and + // partition the cores. Onboard the gate filter assigns it (exec_idx); sim's + // gate does not, so platform_aicpu_affinity_thread_idx() is -1 here for every + // thread — hand those a distinct index from a counter (mirrors run()'s + // thread_idx_++ fallback) instead of collapsing them all to leader 0, which + // would run pre_/post_handshake_init on every thread and race the shared + // scheduler state. Exactly nthreads threads reach init (the gate drops the + // rest), so the counter yields a gap-free [0, nthreads). + int32_t tidx = platform_aicpu_affinity_thread_idx(); + if (tidx < 0) tidx = hs_thread_seq_.fetch_add(1, std::memory_order_acq_rel); + // A thread whose index still falls outside [0, nthreads) owns no core slice: + // handshake_partition would compute lo/hi past cores_total_num_ and index + // all_handshakes[]/core_exec_states_ out of bounds. Reject it here (mirrors + // the bounds guard already in run()). Fail only this thread and do NOT set + // init_failed_ — that would make the valid peers abort before their + // hs_arrived_ increment and hang the leader at the barrier below. + if (tidx >= nthreads) { + LOG_ERROR("AICPU affinity thread idx %d out of range [0,%d) in init", tidx, nthreads); return -1; } + const bool is_leader = (tidx == 0); + + if (is_leader) { + LOG_INFO_V0("AicpuExecutor: Initializing"); + // The 0 → 1 fixup already applied above; derive scheduler count from it. + aicpu_thread_num_ = nthreads; + sched_thread_num_ = nthreads - 1; + orch_to_sched_ = runtime->orch_to_sched; + + hs_arrived_.store(0, std::memory_order_relaxed); + if (sched_ctx_.pre_handshake_init(runtime, aicpu_thread_num_, sched_thread_num_, get_platform_regs()) != 0) { + init_failed_.store(true, std::memory_order_release); + hs_setup_done_.store(true, std::memory_order_release); + return -1; + } + hs_setup_done_.store(true, std::memory_order_release); + } else { + while (!hs_setup_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; + } - finished_count_.store(0, std::memory_order_release); - - init_done_.store(true, std::memory_order_release); - LOG_INFO_V0("AicpuExecutor: Init complete"); + // 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) < 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); + init_done_.store(true, std::memory_order_release); + return -1; + } + init_done_.store(true, std::memory_order_release); + LOG_INFO_V0("AicpuExecutor: Init complete"); + } 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; + } return 0; } @@ -206,11 +261,12 @@ int32_t AicpuExecutor::run(Runtime *runtime) { sched_ctx_.bind_runtime(rt); - // Publish orchestration completion (sets total_tasks_ + orchestrator_done_) - // BEFORE releasing the scheduler threads. Otherwise they unblock with - // total_tasks_=0 and can race to an early exit before the host task - // count is published (host-orch has no concurrent orchestrator to - // keep them alive). + // Latch the host-built task count (on_orchestration_done sets total_tasks_) + // BEFORE the runtime_init_ready_ release below — that store is the barrier + // that unblocks the scheduler threads. Otherwise they would acquire + // runtime_init_ready_ with total_tasks_=0 and race to an early exit before + // the host task count is visible (host-orch has no concurrent orchestrator + // to keep them alive). // NOTE: do NOT call rt_orchestration_done(rt) here. The HOST already // called it in run_host_orchestration; the orchestrator's own // task-allocator pointers are intentionally NOT relocated (only the @@ -297,9 +353,11 @@ void AicpuExecutor::deinit(Runtime *runtime) { LOG_INFO_V0("DeInit: Runtime execution state reset"); - initialized_.store(false, std::memory_order_release); init_done_.store(false, std::memory_order_release); init_failed_.store(false, std::memory_order_release); + hs_setup_done_.store(false, std::memory_order_release); + hs_arrived_.store(0, std::memory_order_release); + hs_thread_seq_.store(0, std::memory_order_release); thread_idx_.store(0, std::memory_order_release); finished_.store(false, std::memory_order_release); @@ -324,10 +382,11 @@ extern "C" int32_t aicpu_prewarm_callable(Runtime *runtime) { * * This is called by DynTileFwkBackendKernelServer in kernel.cpp. * Orchestrates the complete task runtime execution: - * 1. Initialize executor (thread-safe, first thread only) - * 2. Wait for initialization to complete - * 3. Execute tasks on managed cores - * 4. Cleanup when last thread finishes + * 1. Initialize executor: all threads enter init(), which handshakes the cores + * in parallel and barriers internally until init is complete (or a thread + * failed); its return value is authoritative on every thread. + * 2. Execute tasks on managed cores + * 3. Cleanup when last thread finishes * * @param runtime Pointer to Runtime structure * @return 0 on success, non-zero on error @@ -340,13 +399,12 @@ extern "C" int32_t aicpu_execute(Runtime *runtime) { LOG_INFO_V0("%s", "aicpu_execute: Starting AICPU kernel execution"); - g_aicpu_executor.init(runtime); - - while (!g_aicpu_executor.init_done_.load(std::memory_order_acquire)) { - if (g_aicpu_executor.init_failed_.load(std::memory_order_acquire)) { - LOG_ERROR("%s", "aicpu_execute: Initialization failed, aborting execution"); - return -1; - } + // init() barriers every thread internally until init is complete on the + // leader (or a thread failed), then returns the status — so a non-zero + // return is authoritative on all threads and no extra spin is needed. + if (g_aicpu_executor.init(runtime) != 0) { + LOG_ERROR("%s", "aicpu_execute: Initialization failed, aborting execution"); + return -1; } int32_t rc = g_aicpu_executor.run(runtime); diff --git a/src/a2a3/runtime/host_build_graph/docs/RUNTIME_LOGIC.md b/src/a2a3/runtime/host_build_graph/docs/RUNTIME_LOGIC.md index 5c7e53af4..6b0f6fe2e 100644 --- a/src/a2a3/runtime/host_build_graph/docs/RUNTIME_LOGIC.md +++ b/src/a2a3/runtime/host_build_graph/docs/RUNTIME_LOGIC.md @@ -46,19 +46,20 @@ four layers of execution: > > - **host_build_graph (this runtime):** the **host** dlopens the orchestration > `.so` and runs it to completion, populating shared memory + the prebuilt -> arena. Because the orchestrator runs to completion before any scheduler -> exists, it **wires the fanout adjacency inline during submit** (lock -> producers, allocate `dep_pool` entries, seed the ready queue) instead of -> deferring it to a device-side wiring queue. The host then relocates every -> cross-task pointer to its final device address (`relocate_host_orch_image`) -> **before** the H2D copy — pointers into the SM and pointers into the arena -> shift by independent deltas — and ships the image. The device boots -> **scheduler-only**: no on-device orchestrator thread, no on-device pointer -> fixup, no wiring drain (the ready queue is already seeded); it attaches the -> already-device-addressed SM/arena and dispatches. +> arena. It **wires the fanin/fanout adjacency inline during submit** — +> allocates `dep_pool` entries under each live producer's `fanout_lock`, links +> the consumer onto the producer's `fanout_head`, and seeds readiness directly +> into the ready queue (`push_ready_routed` for zero-fanin tasks, +> `route_ready_once` once a producer's fanin resolves). The host then +> relocates every cross-task pointer to its final device address +> (`relocate_host_orch_image`) **before** the H2D copy — pointers into the SM +> and pointers into the arena shift by independent deltas — and ships the +> image. The device boots **scheduler-only**: no on-device orchestrator thread +> and no on-device pointer fixup; it attaches the already-device-addressed +> SM/arena, whose ready queue is already seeded, and dispatches. > - **tensormap_and_ringbuffer:** the orchestrator runs **on-device** on AICPU -> thread N-1, concurrently with the scheduler threads, and defers fanout -> wiring to the scheduler's global wiring queue (drained by thread 0). +> thread N-1, concurrently with the scheduler threads, wiring the same +> fanin/fanout adjacency inline during submit. > > Where a section below says "the orchestrator runs on AICPU Thread 3" or > "Thread 3 dlopens the SO", read that as the **tensormap (device-orch) @@ -372,8 +373,8 @@ When `PTO2OrchestratorState::submit_task` processes parameters: | `kernel_id[3]` | Per-slot kernel IDs: `[AIC, AIV0, AIV1]`; `INVALID_KERNEL_ID` = inactive | | `active_mask` | Bitmask of active subtask slots: `bit0=AIC`, `bit1=AIV0`, `bit2=AIV1` | | `completed_subtasks` | Atomic counter; each subtask increments on completion. Trigger condition: `completed_subtasks == total_required_subtasks` | -| `fanin_count` | Number of producer dependencies (set by scheduler during wiring) | -| `fanout_lock` | Per-task spinlock for concurrent fanout modification (used by scheduler wiring + completion) | +| `fanin_count` | Number of producer dependencies (set inline during submit wiring) | +| `fanout_lock` | Per-task spinlock for fanout modification (used by submit-time wiring + scheduler completion) | | `fanout_head` | Head of fanout consumer list (pointer, protected by `fanout_lock`) | | `fanout_count` | 1 (scope ref) + number of consumers | | `packed_buffer_base` | Start of packed buffer in GM Heap | @@ -436,46 +437,55 @@ Key members: | 3 | **Lookup**: for each INPUT/INOUT param, search TensorMap for producers; collect producer pointers in `PTO2FaninBuilder` | | 4 | **Insert**: register OUTPUT/INOUT args in TensorMap | | 5 | **Record fanin metadata**: store producer pointers in `payload->fanin_inline_slot_states[]` (+ spill pool if >64); increment each producer's `fanout_count` (no lock needed — single writer). This step runs **before** `payload.init()`. | -| 6 | **Wire fanout inline** (`PTO2SchedulerState::wire_task`): lock each producer, allocate `dep_pool` entries, prepend the consumer to each producer's `fanout_head`, and seed the ready queue when all deps are already satisfied. | - -> **Note**: Under host_build_graph the orchestrator runs to completion on the -> host before any scheduler exists, so fanout wiring is done **inline in -> submit** (Step 6) rather than deferred to a device-side wiring queue. The -> `dep_pool` is sized for the whole graph — there is no reclaim during host -> orchestration, exactly like the task window and GM heap — so an exhausted -> pool latches `PTO2_ERROR_DEP_POOL_OVERFLOW` and aborts the run. The -> `fanout_head` / dep-entry / ready-queue pointers this produces are host-DDR -> addresses; `relocate_host_orch_image` shifts them to device addresses before -> H2D (SM pointers and arena pointers by independent deltas). -> -> The inherited tensormap (device-orch) variant instead defers this to the -> scheduler's global `wiring_queue` (SPSC, drained by thread 0) to keep the -> on-device orchestrator's submit path off `fanout_lock` / `dep_pool`. That -> path is described below for the shared mechanics; host_build_graph performs -> the same steps synchronously in Step 6. - -### 7.3 Fanout Wiring (`wire_task`) - -Whether wired inline on the host (host_build_graph) or drained from the -device-side wiring queue (tensormap), each task is wired by the same -`wire_task` logic: - -1. Sets `fanin_count = N + 1` (+1 redundance to prevent premature readiness) -2. For each producer in `payload->fanin_slot_states[]`: - - **Acquires** the producer's `fanout_lock` - - 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 + early_finished count via `fanin_refcount.fetch_add` -4. If all deps satisfied: pushes task to ready queue - -The scheduler's completion handler mirrors this: - -1. **Acquire** `fanout_lock`, mark `task_state = COMPLETED`, read `fanout_head`, **release** lock -2. Traverse fanout list, incrementing each consumer's `fanin_refcount` -3. Mark `task_state = CONSUMED` when `fanout_refcount` reaches `fanout_count` - -This protocol guarantees every consumer is accounted for exactly once. +| 6 | **Wire fanout inline**: for a task with live producers, lock each producer, allocate `dep_pool` entries, prepend the consumer to each producer's `fanout_head`, then seed readiness — `push_ready_routed` for zero-fanin tasks, `route_ready_once` once the fanin refcount is already satisfied. | + +> **Note**: The orchestrator wires fanout **inline in submit** (Step 6) — +> there is no device-side wiring queue. The `dep_pool` is sized for the whole +> graph — there is no reclaim during host orchestration, exactly like the task +> window and GM heap — so an exhausted pool latches +> `PTO2_ERROR_DEP_POOL_OVERFLOW` and aborts the run. The `fanout_head` / +> dep-entry / ready-queue pointers this produces are host-DDR addresses; +> `relocate_host_orch_image` shifts them to device addresses before H2D (SM +> pointers and arena pointers by independent deltas). + +### 7.3 Fanout Wiring + +The orchestrator wires each task's fanin/fanout adjacency inline during submit +(Step 6). Three cases, by the state of the claimed producers: + +1. **Zero-fanin** (`fanin_builder.count == 0`): no producers to link. Sets + `fanin_count = 1`, releases the +1 self-reference, records the `dep_pool` + position (`orch_mark_dep_pool_position`), and seeds the task directly via + `push_ready_routed`. +2. **All claimed producers already completed**: no live fanout links are + needed. Sets `fanin_count = N + 1`, primes `dispatch_fanin` when every + producer is codegen-flagged, releases the refcount, and pushes via + `push_ready_routed`. +3. **At least one live producer** (`orch_wire_live_fanin_task` → + `orch_wire_fanin_task`): + - Reserves `dep_pool` space (`ensure_space`); an exhausted pool latches + `PTO2_ERROR_DEP_POOL_OVERFLOW`. + - Sets `fanin_count = N + 1` (+1 self-reference prevents premature + readiness). + - For each producer, under its `fanout_lock`: if `task_state >= COMPLETED`, + count it as early-finished; otherwise prepend the consumer to the + producer's `fanout_head` via `dep_pool.prepend`. + - Seeds `dispatch_fanin` by the early-finished count when every producer is + codegen-flagged. + - Releases the +1 self-reference plus the early-finished count via + `fanin_refcount.fetch_add`; if that already satisfies `fanin_count`, + publishes readiness via `route_ready_once`. + +`push_ready_routed` pushes a ready slot straight to its shape's queue; +`route_ready_once` claims the slot exactly once (CAS), takes the early-dispatch +doorbell path if the task was pre-staged, and otherwise routes it to the queue. + +The scheduler's completion handler (`on_task_complete`) mirrors the wiring: +acquire `fanout_lock`, mark the task COMPLETED, read `fanout_head`, release the +lock, then traverse the fanout list releasing each consumer's fanin +(`release_fanin_and_check_ready`) and pushing newly-ready consumers via +`route_ready_once`. This protocol guarantees every consumer is accounted for +exactly once. (host-orch does not flip tasks to CONSUMED — see §8.4.) ### 7.4 Scope Mechanism (`PTO2_SCOPE`) diff --git a/src/a2a3/runtime/host_build_graph/docs/SUBMIT_BY_CLUSTER.md b/src/a2a3/runtime/host_build_graph/docs/SUBMIT_BY_CLUSTER.md index 71debb93e..7782020b4 100644 --- a/src/a2a3/runtime/host_build_graph/docs/SUBMIT_BY_CLUSTER.md +++ b/src/a2a3/runtime/host_build_graph/docs/SUBMIT_BY_CLUSTER.md @@ -165,7 +165,7 @@ This project-defined flattened numbering is kept unchanged. 1. Validate submit arguments. 2. Allocate mixed-task ID and initialize descriptor/payload/slot_state once. 3. Lookup producers via TensorMap; collect fanin metadata and increment producers' `fanout_count`. -4. Wire fanout edges and determine readiness. Under host_build_graph (host-orch) this is done **inline in submit** via `wire_task`; the inherited tensormap (device-orch) path instead pushes the task to the scheduler's wiring queue, which scheduler thread 0 drains asynchronously. +4. Wire fanout edges and determine readiness **inline in submit**: link each live producer's `fanout_head` via `dep_pool` entries and seed readiness directly — `push_ready_routed` for zero-fanin tasks, `route_ready_once` once a producer's fanin resolves. No device-side wiring queue is involved. 5. Dispatch all active lanes atomically when resources allow. 6. Aggregate completion and release downstream once. diff --git a/src/a2a3/runtime/host_build_graph/host/runtime_maker.cpp b/src/a2a3/runtime/host_build_graph/host/runtime_maker.cpp index c0c9fbda3..0f11c632d 100644 --- a/src/a2a3/runtime/host_build_graph/host/runtime_maker.cpp +++ b/src/a2a3/runtime/host_build_graph/host/runtime_maker.cpp @@ -281,7 +281,7 @@ namespace { // host_build_graph is host-orchestration-first: the HOST dlopens the // orchestration .so and runs it to completion. The shared memory + arena carry // host-DDR cross-task pointers (slot_state.task/payload, -// payload.fanin_inline_slot_states[], wiring queue); the host relocates them to +// payload.fanin_inline_slot_states[], dep_pool/ready queues); the host relocates them to // their final device addresses (relocate_host_orch_image, below) BEFORE the H2D // copy, so the device receives a fully device-addressed image and schedules // only — no on-device pointer fixup. @@ -355,13 +355,11 @@ struct HostOrchEntryPoints { // Relocated pointers span TWO regions with DIFFERENT deltas: the SM block // (slot_state.task/.payload, fanin_inline_slot_states[], dep-entry.slot_state, // ready-queue slot.slot_state) and the arena block (slot_state.fanout_head, -// dep-entry.next, wiring-queue entries point into the SM but live in the arena). +// dep-entry.next point into the SM but live in the arena). // Rather than track which delta each field needs, reloc() classifies every // pointer by the region it points INTO and applies that region's delta; foreign -// and null pointers pass through untouched. This makes the same walk correct -// whether wiring is still drained on the device (dep_pool/ready empty here, only -// the wiring queue populated) or folded into the host submit (dep_pool/ready -// populated, wiring queue empty). +// and null pointers pass through untouched. The fanout adjacency is wired inline +// during host submit, so dep_pool/ready are already populated here. // // The orchestrator's own task-allocator pointers are intentionally NOT relocated // (the device runs scheduler-only and never dereferences them, and must not call @@ -446,10 +444,10 @@ static bool relocate_host_orch_image( } } - // Per-ring fanout adjacency (dep_pool entries) built by wire_task on the - // host. Each live entry [tail, top) carries a consumer slot_state pointer - // (into the SM) and a next pointer (into the arena). Empty when wiring is - // still drained on the device. + // Per-ring fanout adjacency (dep_pool entries) built inline during host + // submit (orch_wire_fanin_task). Each live entry [tail, top) carries a + // consumer slot_state pointer (into the SM) and a next pointer (into the + // arena). for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { PTO2DepListPool &dp = rt->scheduler.ring_sched_state.dep_pool; if (dp.base == nullptr || dp.capacity == 0) { @@ -476,18 +474,11 @@ static bool relocate_host_orch_image( }; for (int i = 0; i < PTO2_NUM_RESOURCE_SHAPES; i++) { reloc_ready(rt->scheduler.ready_queues[i]); + reloc_ready(rt->scheduler.ready_sync_queues[i]); } reloc_ready(rt->scheduler.dummy_ready_queue); - reloc_ready(rt->scheduler.early_dispatch_queue); - - // Wiring queue (in the arena) — populated only while wiring is still drained - // on the device; its entries are slot_state pointers into the SM. Empty once - // wiring is folded into the host submit. - PTO2SpscQueue &wq = rt->scheduler.wiring.queue; - uint64_t head = wq.head_.load(std::memory_order_relaxed); - uint64_t tail = wq.tail_.load(std::memory_order_relaxed); - for (uint64_t pos = tail; pos < head; pos++) { - reloc(wq.buffer_[pos & wq.mask_]); + for (int i = 0; i < PTO2_NUM_RESOURCE_SHAPES; i++) { + reloc_ready(rt->scheduler.early_dispatch_queues[i]); } return ok; } diff --git a/src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_orchestrator.cpp b/src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_orchestrator.cpp index d4e7dc2e1..eee5c5151 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_orchestrator.cpp +++ b/src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_orchestrator.cpp @@ -72,6 +72,17 @@ extern "C" __attribute__((weak, visibility("hidden"))) bool is_scope_stats_enabl // 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) {} +// AICore register accessor (aicpu/platform_regs.h). The host orchestrator's +// route_ready_once path transitively ODR-uses the early-dispatch doorbell inline +// (pto_scheduler.h ring_one_doorbell), but no core is gated during host +// graph-build, so the doorbell never fires and this weak host fallback only +// satisfies the linker. The AICPU build links the strong definition from +// platform/.../platform_regs.cpp; hidden so the HOST .so does not shadow it. +__attribute__((weak, visibility("hidden"))) volatile uint32_t *get_reg_ptr(uint64_t, RegId) { + static volatile uint32_t sink = 0; + return &sink; +} + // ============================================================================= // Orchestrator Profiling (compile-time toggle) // ============================================================================= @@ -276,9 +287,9 @@ static bool append_fanin_or_fail( // 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 (wire_task 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 + // 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 @@ -344,6 +355,91 @@ static bool append_fanin_or_fail( static void scope_tasks_push(PTO2OrchestratorState *orch, PTO2TaskSlotState *task_slot_state); +static bool all_claimed_fanin_completed(const PTO2FaninBuilder &fanin_builder) { + if (fanin_builder.count == 0) return true; + return fanin_builder.for_each([](PTO2TaskSlotState *producer) -> bool { + return producer != nullptr && producer->task_state.load(std::memory_order_acquire) >= PTO2_TASK_COMPLETED; + }); +} + +static bool all_claimed_fanin_allow_early_resolve(const PTO2FaninBuilder &fanin_builder) { + if (fanin_builder.count == 0) return true; + return fanin_builder.for_each([](PTO2TaskSlotState *producer) -> bool { + return producer != nullptr && producer->allow_early_resolve; + }); +} + +// Record the dep_pool top reached after this slot's Orch-side wiring, so the +// scheduler can bound its reclaim scan to the entries this task allocated. +static void orch_mark_dep_pool_position(PTO2OrchestratorState *orch, PTO2TaskSlotState &slot_state) { + auto &rss = orch->scheduler->ring_sched_state; + slot_state.dep_pool_mark = rss.dep_pool.top; +#if SIMPLER_DFX + if (is_scope_stats_enabled()) { + rss.publish_dep_pool_snapshot(); + } +#endif +} + +// Wire a task with at least one live producer: prepend a fanout link on every +// live producer under its fanout_lock, seed dispatch_fanin for pre-completed +// codegen-flagged producers, then publish readiness via route_ready_once once +// the fanin refcount is already satisfied. +static void orch_wire_fanin_task(PTO2OrchestratorState *orch, PTO2TaskSlotState &slot_state, int32_t wfanin) { + PTO2SchedulerState *sched = orch->scheduler; + auto &rss = sched->ring_sched_state; + PTO2TaskPayload *payload = slot_state.payload; + slot_state.fanin_count = wfanin + 1; + + int32_t early_finished = 0; + bool early_disqualified = false; + for_each_fanin_slot_state(*payload, [&](PTO2TaskSlotState *producer) { + producer->lock_fanout(); + int32_t pstate = producer->task_state.load(std::memory_order_acquire); + if (!early_disqualified && !producer->allow_early_resolve) { + early_disqualified = true; + } + if (pstate >= PTO2_TASK_COMPLETED) { + early_finished++; + } else { + producer->fanout_head = rss.dep_pool.prepend(producer->fanout_head, &slot_state); + } + producer->unlock_fanout(); + }); + + // Pre-completed producers will not dispatch again. Seed dispatch_fanin only + // when every producer is codegen-flagged; one unflagged producer makes the + // direct-only early-dispatch candidate count unreachable by design. + if (!early_disqualified && early_finished != 0) { + payload->dispatch_fanin.fetch_add(early_finished, std::memory_order_acq_rel); + } + + int32_t init_rc = early_finished + 1; + int32_t new_rc = slot_state.fanin_refcount.fetch_add(init_rc, std::memory_order_acq_rel) + init_rc; + orch_mark_dep_pool_position(orch, slot_state); + if (new_rc >= slot_state.fanin_count) { + sched->route_ready_once(slot_state); + } +} + +static bool orch_wire_live_fanin_task(PTO2OrchestratorState *orch, PTO2TaskSlotState &slot_state, int32_t wfanin) { + auto &rss = orch->scheduler->ring_sched_state; + + // dep_pool is orchestrator-exclusive during host orchestration (no lock). + // The pool is sized for the whole graph (no reclaim during host + // orchestration, exactly like the task window and GM heap), so an exhausted + // pool latches PTO2_ERROR_DEP_POOL_OVERFLOW and reports the deadlock with the + // same structural + wall-clock logic the heap/task-window allocator uses. A + // false return also covers a fatal already latched elsewhere. + if (!rss.dep_pool.ensure_space(*rss.ring, wfanin)) { + orch->fatal = true; + return false; + } + + orch_wire_fanin_task(orch, slot_state, wfanin); + return true; +} + struct PTO2PreparedTask { PTO2TaskId task_id = PTO2TaskId::invalid(); PTO2TaskAllocResult alloc_result = {-1, 0, nullptr, nullptr}; @@ -428,14 +524,14 @@ static bool prepare_task( // 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 scheduler wiring.queue.push at the end of - // submit_task_common — that push is the first read of slot_state->task / - // slot_state->payload by another thread. + // 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 spec fields) is initialized in PTO2TaskPayload::init, the - // single payload-init point, which runs before the scheduler wiring push. + // early-dispatch fields) is initialized in PTO2TaskPayload::init, the + // single payload-init point, which runs before Orch-side wiring publish. // Fields already zeroed by reset_for_reuse() at slot init: // fanout_lock=0, fanout_count=PTO2_FANOUT_SCOPE_BIT, fanout_head=nullptr, @@ -451,7 +547,7 @@ static bool prepare_task( static_cast(block_num * __builtin_popcount(active_mask.core_mask())); out->slot_state->logical_block_num = block_num; out->slot_state->active_mask = active_mask; - // fanin_count is set by scheduler during wiring + // fanin_count is set during Orch-side wiring scope_tasks_push(orch, out->slot_state); return true; @@ -667,8 +763,8 @@ static bool ensure_tensormap_capacity(PTO2OrchestratorState *orch, int32_t neede // 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 pushes -// to the scheduler wiring queue. +// 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 @@ -822,6 +918,7 @@ static TaskOutputTensors submit_task_common( } payload.init(args, result, prepared.alloc_result, layout); + cur_slot_state.set_allow_early_resolve(args.allow_early_resolve()); #if SIMPLER_DFX if (is_dump_args_enabled()) { if (args.scalar_count() > 0) { @@ -841,21 +938,38 @@ static TaskOutputTensors submit_task_common( CYCLE_COUNT_LAP(g_orch_args_cycle); - // === STEP 6: wire fanout inline === + // === STEP 6: wire on the orchestrator side and publish readiness === // host_build_graph host-orch: the orchestrator runs to completion on the host // before the device boots scheduler-only, so wire the fanout adjacency here — - // lock each producer, allocate dep_pool entries, and seed the ready queue — + // lock each producer, allocate dep_pool entries, and publish readiness — // directly in submit instead of deferring it to a device-side wiring-queue - // drain. The dep_pool is sized for the whole graph (no reclaim during host + // drain. Zero-fanin tasks and tasks whose claimed producers are already + // completed need no fanout links or dep_pool entries and go straight to the + // ready queue; tasks with live producers allocate fanout links first. The + // dep_pool is sized for the whole graph (no reclaim during host // orchestration, exactly like the task window and GM heap), so an exhausted - // pool latches PTO2_ERROR_DEP_POOL_OVERFLOW via dep_pool.alloc() — the same - // failure class as a task-window/heap overflow. The resulting + // pool latches PTO2_ERROR_DEP_POOL_OVERFLOW via dep_pool.ensure_space() — the + // same failure class as a task-window/heap overflow. The resulting // fanout_head / dep-entry / ready-queue pointers are host-DDR addresses; // runtime_maker relocates them to device addresses before H2D. - sched->wire_task(sched->ring_sched_state, &cur_slot_state, payload.fanin_actual_count); - if (orch->sm_header->orch_error_code.load(std::memory_order_acquire) != PTO2_ERROR_NONE) { - orch->fatal = true; - return result; + if (fanin_builder.count == 0) { + cur_slot_state.fanin_count = 1; + cur_slot_state.fanin_refcount.store(1, std::memory_order_release); + orch_mark_dep_pool_position(orch, cur_slot_state); + sched->push_ready_routed(&cur_slot_state); + } else if (all_claimed_fanin_completed(fanin_builder)) { + int32_t ready_seed = fanin_builder.count + 1; + cur_slot_state.fanin_count = ready_seed; + if (all_claimed_fanin_allow_early_resolve(fanin_builder)) { + payload.dispatch_fanin.store(fanin_builder.count, std::memory_order_release); + } + cur_slot_state.fanin_refcount.store(ready_seed, std::memory_order_release); + orch_mark_dep_pool_position(orch, cur_slot_state); + sched->push_ready_routed(&cur_slot_state); + } else { + if (!orch_wire_live_fanin_task(orch, cur_slot_state, fanin_builder.count)) { + return result; + } } CYCLE_COUNT_LAP(g_orch_fanin_cycle); @@ -1041,7 +1155,16 @@ TaskOutputTensors PTO2OrchestratorState::alloc_tensors(const L0TaskArgs &args) { // 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. - prepared.slot_state->task_state.store(PTO2_TASK_COMPLETED, std::memory_order_release); + // + // 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(); } orch->inline_completed_tasks++; diff --git a/src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_runtime2.cpp b/src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_runtime2.cpp index f22c7f5d5..07c7a76c0 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_runtime2.cpp +++ b/src/a2a3/runtime/host_build_graph/runtime/orchestrator_core/pto_runtime2.cpp @@ -100,7 +100,6 @@ static bool wait_for_tensor_ready(PTO2Runtime *rt, const Tensor &tensor, bool wa 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) { @@ -178,10 +177,6 @@ static bool wait_for_tensor_ready(PTO2Runtime *rt, const Tensor &tensor, bool wa 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 = [&]() { @@ -204,9 +199,6 @@ 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 diff --git a/src/a2a3/runtime/host_build_graph/runtime/pto_async_wait.h b/src/a2a3/runtime/host_build_graph/runtime/pto_async_wait.h index 1bd32b28a..4c9f37a18 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/pto_async_wait.h +++ b/src/a2a3/runtime/host_build_graph/runtime/pto_async_wait.h @@ -24,7 +24,6 @@ #include "pto_runtime2_types.h" struct PTO2SchedulerState; -struct PTO2LocalReadyBuffer; struct CompletionStats; inline constexpr int32_t MAX_ASYNC_WAITS = 64; @@ -171,7 +170,6 @@ struct AsyncWaitList { // entries[]). struct DrainCompletionSink { PTO2SchedulerState *sched{nullptr}; - PTO2LocalReadyBuffer *local_bufs{nullptr}; PTO2TaskSlotState **deferred_release_slot_states{nullptr}; int32_t *deferred_release_count{nullptr}; int32_t deferred_release_capacity{0}; @@ -290,7 +288,7 @@ struct AsyncWaitList { template AsyncPollResult poll_and_complete( - AICoreCompletionMailbox *aicore_mailbox, PTO2SchedulerState *sched, PTO2LocalReadyBuffer *local_bufs, + AICoreCompletionMailbox *aicore_mailbox, PTO2SchedulerState *sched, PTO2TaskSlotState **deferred_release_slot_states, int32_t &deferred_release_count, int32_t deferred_release_capacity #if SIMPLER_SCHED_PROFILING diff --git a/src/a2a3/runtime/host_build_graph/runtime/pto_runtime2.h b/src/a2a3/runtime/host_build_graph/runtime/pto_runtime2.h index c8d798b89..4ecd2a14a 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/pto_runtime2.h +++ b/src/a2a3/runtime/host_build_graph/runtime/pto_runtime2.h @@ -204,8 +204,9 @@ PTO2Runtime *runtime_init_data_from_layout( /** * Phase 3 — wire every arena-internal pointer field (rt->sm_handle, * rt->aicore_mailbox, orchestrator.{scope_tasks, scope_begins, scheduler, - * tensor_map.*, ring.fanin_pool.base}, scheduler.{ready_queues, dep_pool, - * wiring.queue}) so each holds arena.base() + offset. Idempotent — runs on + * tensor_map.*, ring.fanin_pool.base}, scheduler.{ready_queues, + * ready_sync_queues, early_dispatch_queues, dep_pool}) so each holds + * arena.base() + offset. Idempotent — runs on * both host (writing host-mirror addresses) and AICPU (writing device * addresses) sides. */ diff --git a/src/a2a3/runtime/host_build_graph/runtime/pto_runtime2_types.h b/src/a2a3/runtime/host_build_graph/runtime/pto_runtime2_types.h index 7e17f5a85..94f510656 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/pto_runtime2_types.h +++ b/src/a2a3/runtime/host_build_graph/runtime/pto_runtime2_types.h @@ -101,12 +101,14 @@ // Cross-thread early-dispatch work queue (power of two) #define PTO2_EARLY_DISPATCH_QUEUE_SIZE 64 -// Wiring queue -#define PTO2_WRIRING_QUEUE_SIZE 1024 // Per-shape queue size - // Fanin storage #define PTO2_FANIN_INLINE_CAP 64 +// 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 + // TensorMap cleanup interval #define PTO2_TENSORMAP_CLEANUP_INTERVAL 64 // Cleanup every N retired tasks #define PTO2_DEP_POOL_CLEANUP_INTERVAL 64 // Cleanup every N retired tasks @@ -218,12 +220,12 @@ struct PTO2TaskDescriptor { * by bulk tensor and scalar data. Small fanins stay fully inline; larger * fanins spill into a per-ring ring buffer slice. */ -// Speculative early-dispatch claim states for PTO2TaskPayload::spec_state. -enum PTO2SpecState : uint8_t { - PTO2_SPEC_NONE = 0, // not pre-staged - PTO2_SPEC_STAGING = 1, // Hook 1 claimed it; staging in progress - PTO2_SPEC_STAGED = 2, // staged on a core, gated; staged_* fields valid - PTO2_SPEC_DISPATCHED = 3 // routed via the normal dispatch path (no pre-stage) +// Early-dispatch claim states for PTO2TaskPayload::early_dispatch_state. +enum PTO2EarlyDispatchState : uint8_t { + PTO2_EARLY_DISPATCH_NONE = 0, // not pre-staged + PTO2_EARLY_DISPATCH_STAGING = 1, // Hook 1 claimed it; staging in progress + PTO2_EARLY_DISPATCH_STAGED = 2, // staged on a core, gated; staged_* fields valid + PTO2_EARLY_DISPATCH_DISPATCHED = 3 // routed via the normal dispatch path (no pre-stage) }; // A pre-staged consumer occupies one core per gated subtask block. WHICH cores @@ -233,7 +235,7 @@ enum PTO2SpecState : uint8_t { // chip's core count (RUNTIME_MAX_WORKER = 72; no two-level pre-dispatch means // gated cores in flight <= core count), NOT by block_num — so a wide SPMD // consumer can pre-stage all its idle cores. 2 words = 128 bits >= 72. -inline constexpr int PTO2_SPEC_CORE_MASK_WORDS = 2; +inline constexpr int PTO2_EARLY_DISPATCH_CORE_MASK_WORDS = 2; struct PTO2TaskPayload { // === Cache lines 0-8 (576B) — metadata + inline fanin === @@ -243,11 +245,11 @@ struct PTO2TaskPayload { 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]; - // Speculative early-dispatch metadata (AICPU-side only). Ordered by descending + // Early-dispatch metadata (AICPU-side only). Ordered by descending // alignment (8B mask, 4B fanin, then 1B 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 spec atomics (written during staging) would false-share with + // 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. @@ -255,14 +257,13 @@ struct PTO2TaskPayload { // Bitmask of global core_ids this consumer is pre-staged (gated) on. Set with // atomic fetch_or by concurrent stagers; read by release. (Re)initialized in // PTO2TaskPayload::init before the slot can be staged again. - std::atomic staged_core_mask[PTO2_SPEC_CORE_MASK_WORDS]{}; + std::atomic staged_core_mask[PTO2_EARLY_DISPATCH_CORE_MASK_WORDS]{}; // Early-dispatch CANDIDATE detection (event-driven, dual of fanin_refcount): // seeded at wiring with producers already complete, then a flagged producer's // DISPATCH bumps each consumer's dispatch_fanin. dispatch_fanin == // fanin_actual_count <=> every producer is flagged-and-dispatched or was - // pre-completed => this task is an early-dispatch candidate (push early_dispatch_queue). + // pre-completed => this task is an early-dispatch candidate (push early_dispatch_queues[shape]). std::atomic dispatch_fanin{0}; // CONSUMER side: flagged-dispatched + pre-completed producers - bool allow_early_resolve{false}; // codegen hint copied from Arg in PTO2TaskPayload::init // 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 — @@ -271,10 +272,8 @@ struct PTO2TaskPayload { // Release does STAGING->DISPATCHED then rings the mask; a thread that stages a // block AFTER release flipped DISPATCHED rings that block's doorbell itself // (self-ring), so no doorbell is ever missed. - std::atomic spec_state{0}; + std::atomic early_dispatch_state{0}; std::atomic dispatch_propagated{0}; // PRODUCER side: once-guard for fanout propagation - std::atomic spec_chain_active{0}; // inherited early-dispatch flag (auto-chain past codegen flag) - uint8_t spec_chain_depth{0}; // auto-chain depth; inherited = parent+1, capped // === Cache lines 9-72 (4096B) — tensors (alignas(64) forces alignment) === Tensor tensors[MAX_TENSOR_ARGS]; // === Cache lines 73-74 (128B) — scalars === @@ -287,7 +286,7 @@ struct PTO2TaskPayload { /** * 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 spec block at + * own counts are not set until init(). Warms the early-dispatch 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. */ @@ -302,7 +301,7 @@ 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); // spec fields (cache line 8) + __builtin_prefetch(reinterpret_cast(this) + 512, 1, 3); // early-dispatch fields (cache line 8) } /** @@ -339,26 +338,22 @@ 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)); - // Speculative early-dispatch metadata — the single init point for these + // Early-dispatch metadata — the single init point for these // fields. reset_for_reuse MUST NOT touch the payload (it runs at slot // init 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. // - // spec_state / staged_core_mask / dispatch_fanin / spec_chain_* are all - // CONSUMER-side: a task with allow_early_resolve == false still has them - // touched when one of ITS producers is flagged (propagate_dispatch_fanin - // bumps dispatch_fanin and may CAS spec_state / set the auto-chain flag on - // any consumer, independent of the consumer's own hint). So they MUST be - // zeroed here unconditionally — no per-task allow_early_resolve gating. - allow_early_resolve = args.allow_early_resolve(); - spec_state.store(PTO2_SPEC_NONE, std::memory_order_relaxed); - for (int w = 0; w < PTO2_SPEC_CORE_MASK_WORDS; w++) + // 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 + // 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. + 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); - spec_chain_active.store(0, std::memory_order_relaxed); - spec_chain_depth = 0; } }; @@ -406,6 +401,19 @@ static_assert( // never reach -> provable deadlock. static constexpr uint32_t PTO2_FANOUT_SCOPE_BIT = 0x80000000u; +enum PTO2ReadyState : uint8_t { + PTO2_READY_UNCLAIMED = 0, + PTO2_READY_CLAIMED = 1, +}; + +enum PTO2CompletionFlag : uint8_t { + PTO2_COMPLETION_DONE = 2, +}; + +enum PTO2DeferredCompletionFlag : uint8_t { + PTO2_SUBTASK_DEFERRED = 4, +}; + 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) @@ -433,20 +441,19 @@ struct alignas(64) PTO2TaskSlotState { // --- Set per-submit (depend on task inputs) --- ActiveMask active_mask; // Bitmask of active subtask slots (set once) - // 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 write is sequenced before on_subtask_complete's acq_rel - // fetch_add and the read after, so all earlier subtasks' writes are visible - // to the last subtask. - std::atomic any_subtask_deferred{false}; - uint8_t _async_pad{0}; - int32_t dep_pool_mark{0}; // Dep pool top after wiring (thread-0-only) + // 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}; + std::atomic ready_state{PTO2_READY_UNCLAIMED}; + int32_t dep_pool_mark{0}; // Dep pool top after Orch-side wiring std::atomic completed_subtasks{0}; // Each core completion increments by 1 int16_t total_required_subtasks{0}; // = logical_block_num * popcount(active_mask) int16_t logical_block_num{1}; // Total logical blocks (set by orchestrator) - // Next block to dispatch. Atomic so concurrent speculative stagers can each + // Next block to dispatch. Atomic so concurrent early-dispatch stagers can each // claim a distinct block via CAS; normal dispatch (ready-queue serialized) // uses plain relaxed load/store. The two phases never overlap in time (staging // happens before release; normal dispatch of the remainder happens after). @@ -464,6 +471,29 @@ struct alignas(64) PTO2TaskSlotState { task = t; } + void mark_completed() { + task_state.store(PTO2_TASK_COMPLETED, std::memory_order_release); + ready_state.fetch_or(PTO2_COMPLETION_DONE, std::memory_order_release); + } + + bool is_completion_flag_set() const { + return (ready_state.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() { ready_state.fetch_or(PTO2_SUBTASK_DEFERRED, std::memory_order_release); } + + bool has_any_subtask_deferred() const { + return (ready_state.load(std::memory_order_acquire) & PTO2_SUBTASK_DEFERRED) != 0; + } + + void set_allow_early_resolve(bool v) { allow_early_resolve = v; } + /** * Reset dynamic scheduling fields to their pristine values. * In host_build_graph this runs once per slot at init (pto_shared_memory.cpp) @@ -471,7 +501,7 @@ struct alignas(64) PTO2TaskSlotState { * there is no execution-time slot recycle (whole-graph-resident, no reclaim), * so unlike the device-orch path this is not re-invoked after CONSUMED. * - * Skips payload, task, ring_id (immutable, bound once at init). + * Skips payload, task (immutable, bound once at init). * Skips task_state: the orchestrator sets it to PENDING when it populates * the slot. */ @@ -483,11 +513,12 @@ struct alignas(64) PTO2TaskSlotState { 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); - any_subtask_deferred.store(false, std::memory_order_relaxed); - // Note: payload spec fields (spec_state / staged_core_mask / dispatch_fanin / - // spec_chain_*) 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. + ready_state.store(PTO2_READY_UNCLAIMED, std::memory_order_relaxed); + allow_early_resolve = false; + // Note: payload early-dispatch fields (early_dispatch_state / staged_core_mask / dispatch_fanin) + // 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 === diff --git a/src/a2a3/runtime/host_build_graph/runtime/runtime.h b/src/a2a3/runtime/host_build_graph/runtime/runtime.h index 500315d8e..d49e6d724 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/runtime.h +++ b/src/a2a3/runtime/host_build_graph/runtime/runtime.h @@ -88,18 +88,18 @@ constexpr int RUNTIME_DEFAULT_READY_QUEUE_SHARDS = PLATFORM_MAX_AICPU_THREADS - * * Field Access Patterns: * - aicpu_ready: Written by AICPU, read by AICore - * - aicore_done: Written by AICore, read by AICPU + * - aicore_done: Written by AICore, read by AICPU (final report; physical_core_id + * and core_type are published alongside it in the same write) * - task: Written by AICPU, read by AICore (0 = not ready, non-zero = PTO2DispatchPayload*) - * - core_type: Written by AICPU, read by AICore (CoreType::AIC or CoreType::AIV) + * - core_type: Written by AICore (with aicore_done), read by AICPU (CoreType::AIC or CoreType::AIV) + * - physical_core_id: Written by AICore (with aicore_done), read by AICPU */ struct Handshake { - volatile uint32_t aicpu_ready; // AICPU ready signal: 0=not ready, 1=ready - volatile uint32_t aicore_done; // AICore ready signal: 0=not ready, core_id+1=ready - volatile uint64_t task; // Init: PTO2DispatchPayload* (set before aicpu_ready); runtime: unused - volatile CoreType core_type; // Core type: CoreType::AIC or CoreType::AIV - volatile uint32_t physical_core_id; // Physical core ID - volatile uint32_t aicpu_regs_ready; // AICPU register init done: 0=pending, 1=done - volatile uint32_t aicore_regs_ready; // AICore ID reported: 0=pending, 1=done + volatile uint32_t aicpu_ready; // AICPU ready signal: 0=not ready, 1=ready + volatile uint32_t aicore_done; // AICore ready signal: 0=not ready, core_id+1=ready + volatile uint64_t task; // Init: PTO2DispatchPayload* (set before aicpu_ready); runtime: unused + volatile CoreType core_type; // Core type: CoreType::AIC or CoreType::AIV (reported by AICore with aicore_done) + volatile uint32_t physical_core_id; // Physical core ID (reported by AICore with aicore_done) } __attribute__((aligned(64))); /** diff --git a/src/a2a3/runtime/host_build_graph/runtime/scheduler/pto_scheduler.h b/src/a2a3/runtime/host_build_graph/runtime/scheduler/pto_scheduler.h index b1d66dbc0..00fd72395 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/scheduler/pto_scheduler.h +++ b/src/a2a3/runtime/host_build_graph/runtime/scheduler/pto_scheduler.h @@ -33,13 +33,13 @@ #include "common/core_type.h" #include "utils/device_arena.h" -#include "aicpu/platform_regs.h" // get_reg_ptr / RegId for the speculative doorbell +#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 spec doorbell timing too) +#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) \ @@ -62,42 +62,6 @@ struct PTO2ReadyQueueSlot { PTO2TaskSlotState *slot_state; }; -/** - * Thread-local ready buffer for local-first dispatch optimization. - * - * Two buffers per scheduling thread, one per CoreType (AIC=0, AIV=1). - * Initialized once before the scheduling loop; must be empty at - * the start of each iteration (verified by always_assert). - * - * Phase 1 fills per-CoreType buffers via on_task_complete(). - * The dispatch stage drains them local-first via get_ready_tasks_batch, - * with any remaining tasks pushed to the global ready queue. - */ -// 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; } -}; - /** * Lock-free bounded MPMC queue (Dmitry Vyukov design) * @@ -420,128 +384,6 @@ bool ready_queue_init_data_from_layout(PTO2ReadyQueue *queue, DeviceArena &arena void ready_queue_wire_arena_pointers(PTO2ReadyQueue *queue, DeviceArena &arena, size_t slots_off); void ready_queue_destroy(PTO2ReadyQueue *queue); -// ============================================================================= -// SPSC Queue (Single-Producer Single-Consumer, wait-free) -// ============================================================================= -// -// Bounded ring buffer optimized for the wiring queue use case: -// - Producer: orchestrator thread (push) -// - Consumer: scheduler thread 0 (pop_batch) -// -// Design based on Rigtorp's cached-index technique: each side caches -// the other's index locally, avoiding cross-core cache line bouncing -// on the hot path. Only when the local cache says "full" or "empty" -// does the thread issue an acquire load on the remote index. -// -// Memory layout: 5 cache-line-aligned fields ensure zero false sharing. - -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)]; - - // Reserve the backing buffer region on the supplied arena. Returns the - // region offset, to be passed to init_from_layout() after the arena is - // committed. Cache-line aligned: the buffer is shared between the - // orchestrator (push) and scheduler thread 0 (pop_batch), so its base - // must not false-share with neighboring regions. - static size_t reserve_layout(DeviceArena &arena, uint64_t capacity) { - return arena.reserve(capacity * sizeof(uintptr_t), PTO2_ALIGN_SIZE); - } - - // Writes everything except the arena-internal `buffer_` pointer field - // (zeros the slot pointer array, mask/head/tail). The host pre-builds the - // image without storing a host address in buffer_; the AICPU wires - // buffer_ at boot via wire_arena_pointers(). - 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)); - } - - // Arena owns the buffer; here we only forget our pointer. - void destroy() { buffer_ = nullptr; } - - // Push one item (producer only). Returns false if queue is full. - // Full condition: next_h - tail > mask_ (i.e. > capacity-1), so the - // effective usable capacity is capacity-1 (one slot is wasted as a - // sentinel to distinguish full from empty). uint64_t wrapping is safe - // since head and tail are monotonically increasing and subtraction - // wraps correctly. - 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; - } - } - 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; - } - - // 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; - } - - // Full ⟺ the producer's next push() would fail: size has reached the - // usable capacity (mask_ = capacity - 1, one slot reserved as sentinel). - // Used by the wiring-queue deadlock detector to prove the orchestrator is - // blocked in push(). - bool full() const { return size() >= mask_; } -}; - -static_assert(sizeof(PTO2SpscQueue) == 5 * 64, "PTO2SpscQueue must be exactly 5 cache lines (320B)"); -// ============================================================================= - /** * Statistics returned by mixed-task completion processing */ @@ -559,12 +401,11 @@ struct CompletionStats { */ 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; + size_t off_early_dispatch_queue_slots[PTO2_NUM_RESOURCE_SHAPES]; size_t off_dep_pool_entries; - size_t off_wiring_spsc_buffer; uint64_t ready_queue_capacity; - uint64_t spsc_capacity; int32_t dep_pool_capacity; }; @@ -586,12 +427,8 @@ struct PTO2SchedulerState { int32_t last_task_alive; std::atomic advance_lock; // multi-thread CAS - // --- Cache Line 1+: Thread 0 only (wiring dep_pool) --- + // --- Cache Line 1+: Orch-side wiring dep_pool --- alignas(64) PTO2DepListPool dep_pool; - // One-shot latch for the wiring-queue deadlock report (thread 0 only): - // the drain breaks on dep_pool exhaustion every call while wedged, so - // the tier-1 structural diagnostic is emitted once, not per call. - bool dep_deadlock_reported = false; #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; @@ -623,45 +460,17 @@ 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; - // Wiring subsystem — groups all wiring-related state for cache-line isolation. - // - // Three cache-line regions by writer: - // 1. batch_* / backoff — thread 0 exclusive (local batch buffer) - // 2. queue — SPSC: orchestrator push, thread 0 pop - // 3. orch_needs_drain — orchestrator write, thread 0 read - struct alignas(64) WiringState { - static constexpr uint64_t BATCH_SIZE = 30; - static constexpr int BACKOFF_LIMIT = 32; - - // --- Thread 0 exclusive: local batch buffer + backoff --- - int batch_count = 0; - int batch_index = 0; - int backoff_counter = 0; - PTO2TaskSlotState *batch[BATCH_SIZE]; - - // --- SPSC queue: orchestrator (push) ↔ thread 0 (pop) --- - PTO2SpscQueue queue; - - // --- Orchestrator write, thread 0 read --- - alignas(64) std::atomic orch_needs_drain{false}; - // Set to 1 only while the orchestrator is actually spinning in - // queue.push() (queue full), cleared on a successful push. The wiring - // deadlock detector reads this as the producer-blocked observable: it - // proves the orchestrator is stuck BEFORE its scope_end, as opposed to - // having just filled the queue with its last in-scope push and being - // about to call scope_end (which would release the head -> no deadlock). - std::atomic producer_blocked{0}; - } wiring; - - static_assert( - offsetof(WiringState, queue) == 256, "WiringState: batch region must be exactly 4 cache lines before queue" - ); - static_assert(sizeof(WiringState) == 640, "WiringState must be exactly 10 cache lines (640B)"); - alignas(64) AsyncWaitList async_wait_list; // Statistics (cold path, isolated from hot-path fields) @@ -673,192 +482,31 @@ struct PTO2SchedulerState { // Inline hot-path methods // ========================================================================= - /** - * Drain wiring queue: pop submitted tasks and wire their fanout edges. - * Called by scheduler thread 0 each loop iteration. Sets fanin_count, - * acquires fanout_lock per producer, allocates dep_pool entries, and - * pushes ready tasks to the appropriate ready queue. - * - * @return Number of tasks wired this call. - */ - - int drain_wiring_queue(bool force_drain = false) { - int wired = 0; - - // Refill local batch buffer when exhausted. - if (wiring.batch_index >= wiring.batch_count) { - // Backoff: defer pop when queue holds fewer than a full batch, - // unless force_drain, orch_needs_drain, or backoff limit reached. - if (!force_drain && wiring.queue.size() < WiringState::BATCH_SIZE) { - if (!wiring.orch_needs_drain.load(std::memory_order_acquire) && - wiring.backoff_counter < WiringState::BACKOFF_LIMIT) { - wiring.backoff_counter++; - return 0; - } - } - wiring.backoff_counter = 0; - wiring.batch_count = wiring.queue.pop_batch(wiring.batch, WiringState::BATCH_SIZE); - wiring.batch_index = 0; - if (wiring.batch_count == 0) return 0; - } - - // Process tasks from local buffer in strict FIFO order. - while (wiring.batch_index < wiring.batch_count) { - PTO2TaskSlotState *ws = wiring.batch[wiring.batch_index]; - auto &rss = ring_sched_state; - int32_t wfanin = ws->payload->fanin_actual_count; - - if (wfanin > 0 && rss.dep_pool.available() < wfanin) { - rss.dep_pool.reclaim(*rss.ring, rss.last_task_alive); - if (rss.dep_pool.available() < wfanin) { -#if SIMPLER_DFX - if (is_scope_stats_enabled()) { - rss.publish_dep_pool_snapshot(); - } -#endif - // dep_pool can't reclaim because the reclaim watermark is - // wedged. This runs on the scheduler thread, so unlike - // alloc()'s detector it cannot self-observe that the - // orchestrator is blocked; wiring.producer_blocked is the - // external certificate -- the orchestrator sets it ONLY while - // it is actually spinning in queue.push() (cleared on a - // successful push), so the "just filled the queue then called - // scope_end" case (push succeeded -> flag stays 0) cannot trip - // a false report. With the producer provably stuck in push - // (program-order before its scope_end) AND the head COMPLETED, - // all consumers released, scope still open (only scope_end - // frees it), scope_end can never run -> provable head-of-line - // deadlock. The producer-blocked gate also pins the head: - // scope_end has not run, so the scope-gated head cannot be - // CONSUMED/reset concurrently while we read it. - if (!rss.dep_deadlock_reported && wiring.producer_blocked.load(std::memory_order_acquire) != 0) { - int32_t last_alive = rss.last_task_alive; - PTO2TaskSlotState &h = rss.ring->get_slot_state_by_task_id(last_alive); - // Read the head under its fanout_lock: fanout_count is a - // lock-protected field, and one snapshot keeps the check - // and the report consistent. - h.lock_fanout(); - int32_t state = h.task_state.load(std::memory_order_acquire); - uint32_t fc = h.fanout_count; - uint32_t rc = h.fanout_refcount.load(std::memory_order_acquire); - h.unlock_fanout(); - bool head_scope_gated = (state == PTO2_TASK_COMPLETED) && (rc == (fc & ~PTO2_FANOUT_SCOPE_BIT)); - if (head_scope_gated) { - rss.dep_deadlock_reported = true; - report_wiring_deadlock(rss, wfanin, last_alive, state, fc, rc); - // Latch the shared fatal so both sides exit fast off - // one error code: the scheduler cold-path poll - // (handle_orchestrator_exit) emergency_shutdowns, and - // the orchestrator's push spin breaks out and unwinds. - if (rss.dep_pool.error_code_ptr != nullptr) { - int32_t expected = PTO2_ERROR_NONE; - rss.dep_pool.error_code_ptr->compare_exchange_strong( - expected, PTO2_ERROR_DEP_POOL_OVERFLOW, std::memory_order_acq_rel - ); - } - } - } - break; // not enough dep_pool space — keep remainder for next call - } - } - - wiring.batch_index++; - wire_task(rss, ws, wfanin); - wired++; - } - - return wired; - } - - // Tier-1 structural diagnostic for a provable wiring-queue deadlock (head - // COMPLETED + all consumers released + scope still open, dep_pool exhausted, - // orchestrator provably blocked in push). The head snapshot (state/fc/rc) is - // taken under fanout_lock by the caller and passed in, so the report agrees - // with the check and reads no lock-protected field unlocked. - void report_wiring_deadlock( - RingSchedState &rss, int32_t wfanin, int32_t last_alive, int32_t state, uint32_t fc, uint32_t rc - ) { - LOG_ERROR("========================================"); - LOG_ERROR("FATAL: Wiring-Queue Deadlock - Dep Pool Exhausted!"); - LOG_ERROR("========================================"); - LOG_ERROR("Head task %d COMPLETED, all consumers released, scope still open ->", last_alive); - LOG_ERROR("only scope_end can free it, but the orchestrator is blocked on a full wiring"); - LOG_ERROR("queue (in push, before its scope_end). Provable head-of-line deadlock."); - LOG_ERROR( - " Head task %d: state=%d, consumers=%u/%u, scope_released=%d", last_alive, state, - rc & ~PTO2_FANOUT_SCOPE_BIT, fc & ~PTO2_FANOUT_SCOPE_BIT, (rc & PTO2_FANOUT_SCOPE_BIT) ? 1 : 0 - ); - LOG_ERROR(" Dep pool: used=%d/%d, needed=%d entries", rss.dep_pool.used(), rss.dep_pool.capacity, wfanin); - LOG_ERROR("Solution:"); - LOG_ERROR(" The open scope's fanout exceeds the dep pool. Either split the scope, or"); - LOG_ERROR(" raise PTO2_RING_DEP_POOL (compile-time PTO2_DEP_LIST_POOL_SIZE)."); - LOG_ERROR("========================================"); - } - - // Route a ready slot to the right global queue. Dummy tasks (empty - // active_mask) live in dummy_ready_queue; everything else goes to the - // per-shape ready_queues[]. Used by paths that do not have a thread-local - // ready buffer (e.g. wiring). See push_ready_routed_local for the - // dispatch-time fast path. + // Route a ready slot to the right global queue. Dummy tasks (empty active_mask) + // live in dummy_ready_queue; 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) { 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); } } - /** - * Wire fanout edges for a single task. Sets fanin_count, acquires each - * producer's fanout_lock, allocates dep_pool entries for live producers, - * pushes the task to the ready queue once its fanin refcount is satisfied. - */ - void wire_task(RingSchedState &rss, PTO2TaskSlotState *ws, int32_t wfanin) { - PTO2TaskPayload *wp = ws->payload; - ws->fanin_count = wfanin + 1; - - if (wfanin != 0) { - int32_t early_finished = 0; - for_each_fanin_slot_state(*wp, [&](PTO2TaskSlotState *producer) { - producer->lock_fanout(); - int32_t pstate = producer->task_state.load(std::memory_order_acquire); - if (pstate >= PTO2_TASK_COMPLETED) { - early_finished++; - } else { - producer->fanout_head = rss.dep_pool.prepend(producer->fanout_head, ws); - } - producer->unlock_fanout(); - }); - - // Seed dispatch_fanin with producers already complete at wiring - // time (e.g. buffer-creator tasks that finished before this - // consumer entered the graph). Such producers never dispatch at - // runtime, so they can never bump dispatch_fanin via the fanout - // walk; without this seed the candidate compare - // (dispatch_fanin == fanin_actual_count) would be unreachable - // whenever any producer is pre-completed. Mirrors the - // early_finished seed that ready_fanin gets via init_rc. - if (early_finished != 0) { - wp->dispatch_fanin.fetch_add(early_finished, std::memory_order_acq_rel); - } - - int32_t init_rc = early_finished + 1; - int32_t new_rc = ws->fanin_refcount.fetch_add(init_rc, std::memory_order_acq_rel) + init_rc; - if (new_rc >= ws->fanin_count) { - push_ready_routed(ws); + bool try_claim_ready_once(PTO2TaskSlotState &slot_state) { + uint8_t state = slot_state.ready_state.load(std::memory_order_acquire); + for (;;) { + if ((state & PTO2_READY_CLAIMED) != 0) return false; + uint8_t desired = state | PTO2_READY_CLAIMED; + if (slot_state.ready_state.compare_exchange_weak( + state, desired, std::memory_order_acq_rel, std::memory_order_acquire + )) { + return true; } - } else { - ws->fanin_refcount.fetch_add(1, std::memory_order_acq_rel); - push_ready_routed(ws); } - - ws->dep_pool_mark = rss.dep_pool.top; -#if SIMPLER_DFX - if (is_scope_stats_enabled()) { - rss.publish_dep_pool_snapshot(); - } -#endif } // host_build_graph host-orch: the COMPLETED->CONSUMED flip existed only to @@ -896,7 +544,7 @@ struct PTO2SchedulerState { } #endif - // Speculative early-dispatch release. If the now-ready task was pre-staged + // 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 @@ -909,32 +557,37 @@ struct PTO2SchedulerState { // (the stager): CAS NONE->DISPATCHED wins => not pre-staged; lose => STAGED // (spin past the brief STAGING window so the mask is visible), then ring. - // Per-core speculative doorbell table. Hook 1 records each gated core's + // 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 SpecDoorbell { + struct EarlyDispatchDoorbell { uint64_t addr{0}; uint32_t token{0}; }; - SpecDoorbell spec_doorbell_table[PTO2_SPEC_CORE_MASK_WORDS * 64]{}; - - // Cross-thread early-dispatch work queue (a PTO2ReadyQueue MPMC instance, - // arena-backed — reserved/wired in pto_runtime2_init alongside the ready queues). + 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 + // 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_queue; + PTO2ReadyQueue early_dispatch_queues[PTO2_NUM_RESOURCE_SHAPES]; 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)); @@ -942,23 +595,19 @@ struct PTO2SchedulerState { *dmb = (tk << 32) | tk; // 64-bit STR: high=low=token releases the gated AICore } - // auto-chain depth cap: a candidate inherits the flag only while depth < this. - static constexpr uint8_t PTO2_SPEC_CHAIN_MAX = 4; - // Event-driven candidate detection (the dual of fanin_refcount/ready). Call when a // FLAGGED producer `p` DISPATCHES (starts running): walk its fanout and bump each // consumer's dispatch_fanin. A consumer whose dispatch_fanin reaches // fanin_actual_count (= every producer is either flagged-and-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_queue for the idle drain to + // CAS NONE->STAGING (exactly-once) and push to early_dispatch_queues[shape] for the idle drain to // pre-stage. Once-guarded per producer so an SPMD producer's block-by-block - // dispatch propagates once. Replaces the old per-iteration pass-1 PULL scan. + // dispatch propagates once. 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.payload->allow_early_resolve || p.payload->spec_chain_active.load(std::memory_order_acquire))) - return; // only flagged (codegen or inherited) producers propagate + if (!p.allow_early_resolve) return; // only codegen-flagged (direct) producers propagate if (p.payload->dispatch_propagated.exchange(1, std::memory_order_acq_rel) != 0) return; // already propagated once - uint8_t child_depth = static_cast(p.payload->spec_chain_depth + 1); p.lock_fanout(); PTO2DepListEntry *edge = p.fanout_head; // snapshot head, walk lock-free (fanout stable by dispatch) p.unlock_fanout(); @@ -967,32 +616,29 @@ struct PTO2SchedulerState { // 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 early_finished seed (producers already complete) and is - // bumped here by flagged producers; reaching fanin_actual_count means every - // producer is flagged-dispatched or was pre-completed. + // the wiring-time flagged-pre-completed seed and is bumped here by flagged + // producers; reaching fanin_actual_count means every producer is + // flagged-dispatched 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; if (c->active_mask.requires_sync_start()) continue; // sync_start can't be block-by-block pre-staged PTO2ResourceShape shape = c->active_mask.to_shape(); if (shape != PTO2ResourceShape::AIC && shape != PTO2ResourceShape::AIV && shape != PTO2ResourceShape::MIX) continue; - uint8_t expect = PTO2_SPEC_NONE; // exactly-once: only the CAS winner enqueues - if (!c->payload->spec_state.compare_exchange_strong( - expect, PTO2_SPEC_STAGING, std::memory_order_seq_cst, std::memory_order_seq_cst + uint8_t expect = PTO2_EARLY_DISPATCH_NONE; // exactly-once: only the CAS winner enqueues + if (!c->payload->early_dispatch_state.compare_exchange_strong( + expect, PTO2_EARLY_DISPATCH_STAGING, std::memory_order_seq_cst, std::memory_order_seq_cst )) continue; - if (child_depth < PTO2_SPEC_CHAIN_MAX) { // auto-chain: C propagates to ITS consumers - c->payload->spec_chain_depth = child_depth; - c->payload->spec_chain_active.store(1, std::memory_order_release); - } - early_dispatch_queue.push(c); + early_dispatch_queues[static_cast(shape)].push(c); } } - // Collects consumers released via the speculative-doorbell path during a + // 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 SpecReleaseSink { + struct EarlyDispatchReleaseSink { static constexpr int CAP = 32; PTO2TaskSlotState *items[CAP]; int n = 0; @@ -1003,39 +649,42 @@ struct PTO2SchedulerState { } }; - inline bool try_speculative_release(PTO2TaskSlotState &slot_state, SpecReleaseSink *sink = nullptr) { + inline bool try_early_dispatch_release(PTO2TaskSlotState &slot_state, EarlyDispatchReleaseSink *sink = nullptr) { // Never staged => CAS NONE->DISPATCHED wins => dispatch normally. - uint8_t expect = PTO2_SPEC_NONE; - if (slot_state.payload->spec_state.compare_exchange_strong( - expect, PTO2_SPEC_DISPATCHED, std::memory_order_seq_cst, std::memory_order_seq_cst + 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; } // Staged (STAGING). Flip STAGING->DISPATCHED, THEN read the mask. seq_cst // gives a total order with the concurrent stagers, each of which OR-s its - // core into the mask and THEN loads spec_state: a stager whose bit lands + // core into the mask and THEN loads early_dispatch_state: a stager whose bit lands // before this CAS is read here and rung; a stager whose bit lands after // sees DISPATCHED and rings that core itself (self-ring in // stage_consumer_blocks). Either way every gated core's doorbell fires once // (a double-ring is harmless — the AICore already matched). This replaces // the old transient-STAGING spin: STAGING is now the stable gated state. - expect = PTO2_SPEC_STAGING; - slot_state.payload->spec_state.compare_exchange_strong( - expect, PTO2_SPEC_DISPATCHED, std::memory_order_seq_cst, std::memory_order_seq_cst + 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 ); - for (int w = 0; w < PTO2_SPEC_CORE_MASK_WORDS; w++) { + 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(spec_doorbell_table[core_id].addr, spec_doorbell_table[core_id].token); + ring_one_doorbell( + early_dispatch_doorbell_table[core_id].addr, early_dispatch_doorbell_table[core_id].token + ); } } // This pre-staged consumer was just released by its doorbell — it starts - // running NOW, so propagate dispatch_fanin to ITS consumers (auto-chain, - // knob A). 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. + // 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 (sink == nullptr || !sink->push(&slot_state)) { propagate_dispatch_fanin(slot_state); } @@ -1046,29 +695,67 @@ struct PTO2SchedulerState { return slot_state.next_block_idx.load(std::memory_order_seq_cst) >= slot_state.logical_block_num; } - bool release_fanin_and_check_ready( - PTO2TaskSlotState &slot_state, PTO2LocalReadyBuffer *local_bufs = nullptr, SpecReleaseSink *sink = nullptr + 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. + if (try_early_dispatch_release(slot_state, sink)) return true; + + PTO2ResourceShape shape = slot_state.active_mask.to_shape(); + if (shape == PTO2ResourceShape::DUMMY) { + 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 state = slot_state.ready_state.load(std::memory_order_acquire); + atomic_count += 1; // ready_state load + for (;;) { + if ((state & PTO2_READY_CLAIMED) != 0) return false; + uint8_t desired = state | PTO2_READY_CLAIMED; + if (slot_state.ready_state.compare_exchange_weak( + state, desired, std::memory_order_acq_rel, std::memory_order_acquire + )) { + atomic_count += 1; // ready_state CAS + break; + } + atomic_count += 1; // failed ready_state CAS + } + + // Early-dispatch: pre-staged tasks are released by doorbell + // here, skipping the ready-queue round-trip entirely. + if (try_early_dispatch_release(slot_state, sink)) return true; + + PTO2ResourceShape shape = slot_state.active_mask.to_shape(); + if (shape == PTO2ResourceShape::DUMMY) { + 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) { - // Speculative early-dispatch: pre-staged tasks are released by doorbell - // here, skipping the ready-queue round-trip entirely. - if (try_speculative_release(slot_state, sink)) return true; - // Local-first: try per-CoreType thread-local buffer before global queue - // Route by active_mask: AIC-containing tasks → buf[0], AIV-only → buf[1] - // DUMMY shape is out of range for local_bufs (sized PTO2_NUM_RESOURCE_SHAPES); - // dummy slots bypass the local fast path and go straight to dummy_ready_queue. - PTO2ResourceShape shape = slot_state.active_mask.to_shape(); - if (shape == PTO2ResourceShape::DUMMY) { - dummy_ready_queue.push(&slot_state); - } else if (!local_bufs || !local_bufs[static_cast(shape)].try_push(&slot_state)) { - ready_queues[static_cast(shape)].push(&slot_state); - } - return true; + return route_ready_once(slot_state, sink); } return false; } @@ -1076,60 +763,28 @@ struct PTO2SchedulerState { #if SIMPLER_ORCH_PROFILING || SIMPLER_SCHED_PROFILING bool release_fanin_and_check_ready( PTO2TaskSlotState &slot_state, uint64_t &atomic_count, uint64_t &push_wait, - PTO2LocalReadyBuffer *local_bufs = nullptr, SpecReleaseSink *sink = nullptr + 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) { - // Speculative early-dispatch: pre-staged tasks are released by doorbell - // here, skipping the ready-queue round-trip entirely. - if (try_speculative_release(slot_state, sink)) return true; - // Local-first: try per-CoreType thread-local buffer before global queue. - // Dummy slots bypass local_bufs (out-of-range for PTO2_NUM_RESOURCE_SHAPES) - // and go straight to dummy_ready_queue; use the profiling-aware push so - // atomic_count / push_wait stay consistent with the non-dummy path. - PTO2ResourceShape shape = slot_state.active_mask.to_shape(); - if (shape == PTO2ResourceShape::DUMMY) { - dummy_ready_queue.push(&slot_state, atomic_count, push_wait); - } else if (!local_bufs || !local_bufs[static_cast(shape)].try_push(&slot_state)) { - ready_queues[static_cast(shape)].push(&slot_state, atomic_count, push_wait); - } - return true; + return route_ready_once(slot_state, atomic_count, push_wait, sink); } return false; } #endif - int get_ready_tasks_batch( - PTO2ResourceShape shape, PTO2LocalReadyBuffer &local_buf, PTO2TaskSlotState **out, int max_count - ) { - 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; + 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( - PTO2ResourceShape shape, PTO2LocalReadyBuffer &local_buf, PTO2TaskSlotState **out, int max_count, - uint64_t &atomic_count, uint64_t &wait_cycle + PTO2ReadyQueue *queues, PTO2ResourceShape shape, PTO2TaskSlotState **out, int max_count, uint64_t &atomic_count, + uint64_t &wait_cycle ) { - 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, atomic_count, wait_cycle); - } - return count; + return queues[static_cast(shape)].pop_batch(out, max_count, atomic_count, wait_cycle); } #endif @@ -1168,7 +823,7 @@ struct PTO2SchedulerState { * 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 speculative hits. + * 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 @@ -1181,12 +836,11 @@ struct PTO2SchedulerState { uint32_t #endif on_task_complete( - PTO2TaskSlotState &slot_state, + PTO2TaskSlotState &slot_state #if SIMPLER_SCHED_PROFILING - int thread_idx, + , + int thread_idx #endif - - PTO2LocalReadyBuffer *local_bufs = nullptr ) { #if SIMPLER_SCHED_PROFILING CompletionStats stats = {0, 0, 0, true}; @@ -1206,12 +860,12 @@ struct PTO2SchedulerState { #else slot_state.lock_fanout(); #endif - slot_state.task_state.store(PTO2_TASK_COMPLETED, std::memory_order_release); + slot_state.mark_completed(); PTO2DepListEntry *current = slot_state.fanout_head; // Protected by fanout_lock slot_state.unlock_fanout(); #if SIMPLER_SCHED_PROFILING - lock_atomics += 2; // state.store + unlock.store + lock_atomics += 3; // task_state.store + ready_state.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]); @@ -1229,24 +883,24 @@ struct PTO2SchedulerState { // 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 - // speculative early-dispatch: minimal producer-end -> consumer-start.) + // 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. - SpecReleaseSink rel_sink; + 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, local_bufs, &rel_sink)) { + if (release_fanin_and_check_ready(consumer_slot, fanout_atomics, push_wait, &rel_sink)) { stats.tasks_enqueued++; } #else consumer_walk_count++; - release_fanin_and_check_ready(consumer_slot, local_bufs, &rel_sink); + release_fanin_and_check_ready(consumer_slot, &rel_sink); #endif current = current->next; } @@ -1303,7 +957,7 @@ struct PTO2SchedulerState { // === 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, wiring SPSC buffer) on the supplied arena. + // 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); @@ -1319,7 +973,7 @@ struct PTO2SchedulerState { // Phase 3b: write the arena-internal pointer fields // (ready_queues[].slots, dummy_ready_queue.slots, dep_pool.base for each - // ring, wiring.queue.buffer_). Called on both host and device sides. + // 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. @@ -1345,9 +999,9 @@ AsyncWaitList::try_inline_complete_locked(AsyncWaitList::DrainCompletionSink &si // 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, sink.local_bufs); + (void)sink.sched->on_task_complete(slot_state, sink.thread_idx); #else - (void)sink.sched->on_task_complete(slot_state, sink.local_bufs); + (void)sink.sched->on_task_complete(slot_state); #endif if (*sink.deferred_release_count >= sink.deferred_release_capacity) { while (*sink.deferred_release_count > 0) { @@ -1367,7 +1021,7 @@ AsyncWaitList::try_inline_complete_locked(AsyncWaitList::DrainCompletionSink &si template inline AsyncPollResult AsyncWaitList::poll_and_complete( - AICoreCompletionMailbox *aicore_mailbox, PTO2SchedulerState *sched, PTO2LocalReadyBuffer *local_bufs, + AICoreCompletionMailbox *aicore_mailbox, PTO2SchedulerState *sched, PTO2TaskSlotState **deferred_release_slot_states, int32_t &deferred_release_count, int32_t deferred_release_capacity #if SIMPLER_SCHED_PROFILING , @@ -1379,7 +1033,6 @@ inline AsyncPollResult AsyncWaitList::poll_and_complete( AsyncWaitList::DrainCompletionSink sink{}; sink.sched = sched; - sink.local_bufs = local_bufs; sink.deferred_release_slot_states = deferred_release_slot_states; sink.deferred_release_count = &deferred_release_count; sink.deferred_release_capacity = deferred_release_capacity; @@ -1427,9 +1080,9 @@ inline AsyncPollResult AsyncWaitList::poll_and_complete( // 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, local_bufs); + (void)sched->on_task_complete(*entry.slot_state, thread_idx); #else - (void)sched->on_task_complete(*entry.slot_state, local_bufs); + (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 diff --git a/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_cold_path.cpp b/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_cold_path.cpp index 7b612163e..594a70e7e 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_cold_path.cpp +++ b/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_cold_path.cpp @@ -73,9 +73,6 @@ LoopAction SchedulerContext::handle_orchestrator_exit( return LoopAction::BREAK_LOOP; } - bool orch_done = orchestrator_done_; - 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); @@ -88,21 +85,6 @@ LoopAction SchedulerContext::handle_orchestrator_exit( return LoopAction::NONE; } -LoopAction SchedulerContext::handle_core_transition(bool &cores_released) { - if (!transition_requested_.load(std::memory_order_acquire)) return LoopAction::NONE; - if (!reassigned_.load(std::memory_order_acquire)) { - wait_reassign_.fetch_add(1, std::memory_order_release); - while (!reassigned_.load(std::memory_order_acquire)) { - if (completed_.load(std::memory_order_acquire)) { - return LoopAction::BREAK_LOOP; - } - SPIN_WAIT_HINT(); - } - } - cores_released = true; - return LoopAction::NONE; -} - LoopAction SchedulerContext::check_idle_fatal_error(int32_t thread_idx, PTO2SharedMemoryHeader *header, Runtime *runtime) { if (completed_.load(std::memory_order_acquire)) { @@ -320,7 +302,7 @@ void SchedulerContext::log_stall_diagnostics( // 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 / reassign_cores_for_all_threads. + // 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; @@ -427,8 +409,8 @@ void SchedulerContext::log_l2_swimlane_summary(int32_t thread_idx, int32_t cur_t cycles_to_us(sched_end_ts - l2_swimlane.sched_start_ts) ); - uint64_t sched_total = l2_swimlane.sched_wiring_cycle + l2_swimlane.sched_complete_cycle + - l2_swimlane.sched_dispatch_cycle + l2_swimlane.sched_idle_cycle; + uint64_t sched_total = + l2_swimlane.sched_complete_cycle + l2_swimlane.sched_dispatch_cycle + l2_swimlane.sched_idle_cycle; if (sched_total == 0) sched_total = 1; #if SIMPLER_SCHED_PROFILING @@ -521,19 +503,6 @@ void SchedulerContext::log_l2_swimlane_summary(int32_t thread_idx, int32_t cur_t l2_swimlane.sched_dispatch_setup_cycle * 100.0 / d_parent ); -#if SIMPLER_SCHED_PROFILING - LOG_INFO_V9( - "Thread %d: wiring : %.3fus (%.1f%%) tasks=%d", thread_idx, - cycles_to_us(l2_swimlane.sched_wiring_cycle), l2_swimlane.sched_wiring_cycle * 100.0 / sched_total, - l2_swimlane.phase_wiring_count - ); -#else - LOG_INFO_V9( - "Thread %d: wiring : %.3fus (%.1f%%)", thread_idx, cycles_to_us(l2_swimlane.sched_wiring_cycle), - l2_swimlane.sched_wiring_cycle * 100.0 / sched_total - ); -#endif - 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 @@ -590,102 +559,131 @@ int32_t SchedulerContext::shutdown(int32_t thread_idx) { } // ============================================================================= -// Handshake with all AICore workers; discover core type and reg address. +// 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. // ============================================================================= -int32_t SchedulerContext::handshake_all_cores(Runtime *runtime) { +void SchedulerContext::handshake_partition(Runtime *runtime, int32_t tidx, int32_t nthreads) { Handshake *all_handshakes = reinterpret_cast(runtime->workers); - cores_total_num_ = runtime->worker_count; + 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); - // Validate cores_total_num_ before using as array index - 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; - } - - aic_count_ = 0; - aiv_count_ = 0; - - LOG_INFO_V0("Handshaking with %d cores", cores_total_num_); - - // Step 1: Write per-core payload addresses and send handshake signal. - // OUT_OF_ORDER_STORE_BARRIER() ensures task is globally visible before - // aicpu_ready=1, so AICore reads the correct payload pointer after waking up. - for (int32_t i = 0; i < cores_total_num_; 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(); + // 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: Wait for all cores to respond, collect core type and register addresses - bool handshake_failed = false; - for (int32_t i = 0; i < cores_total_num_; i++) { - Handshake *hank = &all_handshakes[i]; - - while (hank->aicore_regs_ready == 0) { - SPIN_WAIT_HINT(); - } - - 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 = true; - continue; - } - - uint64_t *regs = reinterpret_cast(regs_); - uint64_t reg_addr = regs[physical_core_id]; - - // Initialize AICore registers after discovery (first round) - platform_init_aicore_regs(reg_addr); - OUT_OF_ORDER_STORE_BARRIER(); - hank->aicpu_regs_ready = 1; - - OUT_OF_ORDER_STORE_BARRIER(); - - while (hank->aicore_done == 0) { - SPIN_WAIT_HINT(); + // 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--; } + } - CoreType type = hank->core_type; + // 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(); - core_exec_states_[i].reg_addr = reg_addr; - core_exec_states_[i].cond_ptr = get_reg_ptr(reg_addr, RegId::COND); + // 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); + } -#if SIMPLER_DFX - // Record physical_core_id for PMU init later (CoreExecState has no room - // for this field under SIMPLER_DFX). - physical_core_ids_[i] = physical_core_id; -#endif + // 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 - core_exec_states_[i].worker_id = i; - core_exec_states_[i].physical_core_id = physical_core_id; - core_exec_states_[i].core_type = type; + 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 - - if (type == CoreType::AIC) { - aic_worker_ids_[aic_count_++] = i; - LOG_INFO_V0("Core %d: AIC, physical_id=%u, reg_addr=0x%lx", i, physical_core_id, reg_addr); - } else { - aiv_worker_ids_[aiv_count_++] = i; - LOG_INFO_V0("Core %d: AIV, physical_id=%u, reg_addr=0x%lx", i, physical_core_id, reg_addr); - } - } - - if (handshake_failed) { - emergency_shutdown(runtime); - return -1; } - - LOG_INFO_V0("Core discovery complete: %d AIC, %d AIV", aic_count_, aiv_count_); - return 0; + OUT_OF_ORDER_STORE_BARRIER(); } // ============================================================================= @@ -711,10 +709,8 @@ bool SchedulerContext::assign_cores_to_threads() { active_sched_threads_, aic_count_, aiv_count_ ); - 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; - } + // 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] = {}; @@ -752,88 +748,19 @@ bool SchedulerContext::assign_cores_to_threads() { return true; } -// ============================================================================= -// Reassign all cores across all threads (sched + orchestrator) after orchestration. -// ============================================================================= -void SchedulerContext::reassign_cores_for_all_threads() { - LOG_INFO_V0( - "Reassigning cores (cluster-aligned) for %d threads: %d AIC, %d AIV", aicpu_thread_num_, aic_count_, aiv_count_ - ); - - // Collect running worker_ids from all current trackers - bool running_cores[RUNTIME_MAX_WORKER] = {}; - for (int32_t i = 0; i < aicpu_thread_num_; i++) { - auto all_running = core_trackers_[i].get_all_running_cores(); - int32_t bp; - while ((bp = all_running.pop_first()) >= 0) { - running_cores[core_trackers_[i].get_core_id_by_offset(bp)] = true; - } - } - - // Count clusters per thread (round-robin across all threads) - int32_t cluster_count = aic_count_; - int32_t clusters_per_thread[MAX_AICPU_THREADS] = {}; - for (int32_t ci = 0; ci < cluster_count; ci++) { - clusters_per_thread[ci % aicpu_thread_num_]++; - } - - // Re-init all trackers and reset core counts - for (int32_t i = 0; i < aicpu_thread_num_; i++) { - core_trackers_[i].init(clusters_per_thread[i]); - } - - // Assign clusters round-robin and restore running state - int32_t cluster_idx_per_thread[MAX_AICPU_THREADS] = {}; - for (int32_t ci = 0; ci < cluster_count; ci++) { - int32_t t = ci % aicpu_thread_num_; - - 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]; - - int32_t cl_idx = cluster_idx_per_thread[t]++; - core_trackers_[t].set_cluster(cl_idx, aic_wid, aiv0_wid, aiv1_wid); - - // init() marks all idle; toggle cores that were running and restore pending_occupied - if (running_cores[aic_wid]) { - core_trackers_[t].change_core_state(cl_idx * 3); - core_trackers_[t].set_pending_occupied(cl_idx * 3); - } - if (running_cores[aiv0_wid]) { - core_trackers_[t].change_core_state(cl_idx * 3 + 1); - core_trackers_[t].set_pending_occupied(cl_idx * 3 + 1); - } - if (running_cores[aiv1_wid]) { - core_trackers_[t].change_core_state(cl_idx * 3 + 2); - core_trackers_[t].set_pending_occupied(cl_idx * 3 + 2); - } - } - - // Log final distribution - LOG_INFO_V0("Core reassignment complete:"); - for (int32_t t = 0; t < aicpu_thread_num_; t++) { - int32_t aic_running = core_trackers_[t].get_running_count(); - int32_t aiv_running = core_trackers_[t].get_running_count(); - LOG_INFO_V0( - " Thread %d: %d cores, %d clusters (AIC running=%d, AIV running=%d)", t, core_trackers_[t].core_num(), - core_trackers_[t].get_cluster_count(), aic_running, aiv_running - ); - } - active_sched_threads_ = aicpu_thread_num_; -} - // ============================================================================= // 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"); - Handshake *all_handshakes = reinterpret_cast(runtime->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(); - hank->aicpu_regs_ready = 1; + // 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++; @@ -849,8 +776,8 @@ void SchedulerContext::emergency_shutdown(Runtime *runtime) { // ============================================================================= // Lifecycle: init / deinit // ============================================================================= -int32_t SchedulerContext::init( - Runtime *runtime, int32_t aicpu_thread_num, int32_t sched_thread_num, bool orch_to_sched, uint64_t regs_base +int32_t SchedulerContext::pre_handshake_init( + Runtime *runtime, int32_t aicpu_thread_num, int32_t sched_thread_num, uint64_t regs_base ) { always_assert(runtime != nullptr); @@ -860,7 +787,6 @@ int32_t SchedulerContext::init( // Wire thread/transition configuration that handshake/assign need to read. aicpu_thread_num_ = aicpu_thread_num; sched_thread_num_ = sched_thread_num; - orch_to_sched_ = orch_to_sched; regs_ = regs_base; #if SIMPLER_DFX @@ -869,7 +795,9 @@ int32_t SchedulerContext::init( // 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. + // 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). if (is_l2_swimlane_enabled()) { l2_swimlane_aicpu_init(runtime->worker_count); l2_swimlane_level_ = get_l2_swimlane_level(); @@ -881,10 +809,9 @@ int32_t SchedulerContext::init( // 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 active_sched = (sched_thread_num_ > 0) ? sched_thread_num_ : aicpu_thread_num_; - const int sched_phase_threads = orch_to_sched_ ? aicpu_thread_num_ : active_sched; + 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) in both modes — see record_orch_phase. + // (ordinal 0) — see record_orch_phase. const int orch_phase_threads = 1; l2_swimlane_aicpu_init_phase(runtime->worker_count, sched_phase_threads, orch_phase_threads); } @@ -893,38 +820,75 @@ int32_t SchedulerContext::init( } #endif - // Discover cores and assign to scheduler threads. - int32_t rc = handshake_all_cores(runtime); - if (rc != 0) { - LOG_ERROR("handshake_all_cores failed"); - return rc; + // Core count is needed by every thread to compute its handshake slice. + cores_total_num_ = runtime->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; } + aic_count_ = 0; + aiv_count_ = 0; + handshake_failed_.store(false, std::memory_order_release); + + LOG_INFO_V0("Handshaking with %d cores", cores_total_num_); + return 0; +} + +int32_t SchedulerContext::post_handshake_init(Runtime *runtime) { + if (handshake_failed_.load(std::memory_order_acquire)) { + emergency_shutdown(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; + 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; + } + } + 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; } - // Profiling-subsystem buffer/state init: single-threaded cold path, so the - // "do it once" guarantee is structural (no CAS needed). Runs after - // handshake_all_cores / assign_cores_to_threads because pmu_aicpu_init needs + // 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(orch_to_sched_ ? aicpu_thread_num_ : active_sched_threads_); + 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_); } -#endif - // dep_gen is host-driven (SubmitTrace) and gated independently of - // SIMPLER_DFX. init() only pops the initial buffer from instance 0's - // free_queue; the orchestrator thread still records its idx via + // 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(); } +#endif // Initialize task counters. Task count comes from PTO2 shared memory. if (runtime->get_gm_sm_ptr()) { @@ -947,12 +911,12 @@ int32_t SchedulerContext::init( } completed_tasks_.store(0, std::memory_order_release); - // Device orchestration: the orchestrator thread flips this when the graph is built. - orchestrator_done_ = false; - - // 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_)); + // 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. // Initialize per-core GlobalContext (sub_block_id) based on cluster position. // This is done once at startup and never modified afterwards. @@ -982,9 +946,16 @@ void SchedulerContext::deinit() { core_exec_states_[i].pending_reg_task_id = AICPU_TASK_INVALID; } - // 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_)); + // No per-core memset of payload_per_core_ / deferred_slab_per_core_ here + // (~300 KB across all cores). Both are fully re-initialized at dispatch + // before they can be read: dispatch_task sets deferred_slab->count = 0 / + // error_code = NONE and build_payload() overwrites every payload field + // (function addr, args[], contexts, not_ready) on the exact [core][buf_idx] + // about to run. 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. // Reset sync-start drain coordination — a previous run that aborted mid-drain // would otherwise leave dirty pending/elected/ack state for the next reuse. @@ -996,12 +967,6 @@ void SchedulerContext::deinit() { // Reset task counters and orchestrator state completed_tasks_.store(0, std::memory_order_release); total_tasks_ = 0; - orchestrator_done_ = false; - - // Reset core transition state - transition_requested_.store(false, std::memory_order_release); - wait_reassign_.store(0, std::memory_order_release); - reassigned_.store(false, std::memory_order_release); completed_.store(false, std::memory_order_release); // Reset core discovery and assignment state @@ -1010,7 +975,6 @@ void SchedulerContext::deinit() { cores_total_num_ = 0; aicpu_thread_num_ = 0; sched_thread_num_ = 0; - orch_to_sched_ = false; active_sched_threads_ = 0; for (int32_t t = 0; t < MAX_AICPU_THREADS; t++) { core_trackers_[t] = CoreTracker{}; @@ -1028,12 +992,15 @@ void SchedulerContext::bind_runtime(PTO2Runtime *rt) { } // ============================================================================= -// 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). +// Post-orchestration bookkeeping. Runs once on the boot leader after the +// host-built image is attached; latches total_tasks_ and folds inline-completed +// tasks (or shuts down on a fatal orchestration error). The caller publishes +// runtime_init_ready_ (release) after this returns — that store is what makes +// total_tasks_ visible to the scheduler threads, which acquire it before +// dispatching. // ============================================================================= void SchedulerContext::on_orchestration_done( - Runtime *runtime, PTO2Runtime *rt, int32_t thread_idx, int32_t total_tasks + Runtime *runtime, PTO2Runtime *rt, [[maybe_unused]] int32_t thread_idx, int32_t total_tasks ) { #if SIMPLER_DFX if (l2_swimlane_level_ >= L2SwimlaneLevel::ORCH_PHASES) { @@ -1054,7 +1021,6 @@ void SchedulerContext::on_orchestration_done( rt->scheduler.tasks_completed.fetch_add(inline_completed, std::memory_order_relaxed); #endif } - orchestrator_done_ = true; // Check for fatal error from orchestration; if so, shut down immediately. int32_t orch_err = 0; @@ -1067,32 +1033,9 @@ void SchedulerContext::on_orchestration_done( } } - // Skip core transition on fatal error — cores already shut down above. - if (completed_.load(std::memory_order_acquire)) { - // Signal transition to unblock scheduler threads waiting at core transition - transition_requested_.store(true, std::memory_order_release); - reassigned_.store(true, std::memory_order_release); - } else if (orch_to_sched_) { - LOG_INFO_V0("Thread %d: Set orchestrator_done=true, requesting core transition", thread_idx); - transition_requested_.store(true, std::memory_order_release); - - // Wait for scheduler threads to acknowledge transition request - while (wait_reassign_.load(std::memory_order_acquire) != sched_thread_num_) { - if (completed_.load(std::memory_order_acquire)) { - break; - } - SPIN_WAIT_HINT(); - } - if (!completed_.load(std::memory_order_acquire)) { - reassign_cores_for_all_threads(); - reassigned_.store(true, std::memory_order_release); - } - } - #if SIMPLER_DFX - // Write core-to-thread mapping AFTER reassignment so the profiling data - // reflects the final distribution (all active_sched_threads_, including - // former orchestrator threads when orch_to_sched_ is enabled). + // 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++) { diff --git a/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_completion.cpp b/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_completion.cpp index 3c35d688b..cfd9007b8 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_completion.cpp +++ b/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_completion.cpp @@ -57,7 +57,7 @@ SlotTransition SchedulerContext::decide_slot_transition( t.running_done = true; t.running_freed = true; } else if (pending_gated) { - // Case 3.3: running FIN, pending is a SPECULATIVE GATED task. The + // 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 @@ -85,7 +85,7 @@ SlotTransition SchedulerContext::decide_slot_transition( 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, PTO2LocalReadyBuffer *local_bufs + PTO2TaskSlotState *deferred_release_slot_states[], int32_t &deferred_release_count #if SIMPLER_DFX , uint64_t dispatch_ts, uint64_t finish_ts @@ -101,7 +101,7 @@ void SchedulerContext::complete_slot_task( // 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 - // any_subtask_deferred flag on slot_state is the discriminator; it's set + // 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; @@ -134,7 +134,7 @@ void SchedulerContext::complete_slot_task( // 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.any_subtask_deferred.store(true, std::memory_order_release); + slot_state.mark_any_subtask_deferred(); const PTO2TaskId token = slot_state.task->task_id; for (uint32_t i = 0; i < cond_count; ++i) { @@ -157,8 +157,7 @@ void SchedulerContext::complete_slot_task( } #endif - if (task_complete && slot_state.payload != nullptr && - slot_state.any_subtask_deferred.load(std::memory_order_acquire)) { + if (task_complete && slot_state.payload != nullptr && slot_state.has_any_subtask_deferred()) { // 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))) { @@ -184,7 +183,7 @@ void SchedulerContext::complete_slot_task( #endif #if SIMPLER_DFX // Time Resolve (walk the consumer list, decrement each consumer's - // fanin, push the newly-ready ones, ring doorbells for speculative + // 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 / @@ -200,9 +199,9 @@ void SchedulerContext::complete_slot_task( // 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, local_bufs).fanout_edges; + consumers_resolved = sched_->on_task_complete(slot_state, thread_idx).fanout_edges; #else - consumers_resolved = sched_->on_task_complete(slot_state, local_bufs); + consumers_resolved = sched_->on_task_complete(slot_state); #endif #if SIMPLER_DFX if (resolve_t0 != 0) { @@ -294,8 +293,7 @@ void SchedulerContext::clear_running_slot(CoreExecState &core) { void SchedulerContext::check_running_cores_for_completion( int32_t thread_idx, Handshake *hank, int32_t &completed_this_turn, int32_t &cur_thread_completed, - bool &made_progress, PTO2TaskSlotState *deferred_release_slot_states[], int32_t &deferred_release_count, - PTO2LocalReadyBuffer *local_bufs + bool &made_progress, PTO2TaskSlotState *deferred_release_slot_states[], int32_t &deferred_release_count ) { #if SIMPLER_SCHED_PROFILING auto &l2_swimlane = sched_l2_swimlane_[thread_idx]; @@ -307,17 +305,17 @@ 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 speculative cores. A STAGED task is parked on this core + // 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_speculative_release) flips spec_state to DISPATCHED, at + // 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->spec_state.load(std::memory_order_relaxed) == PTO2_SPEC_STAGING) { + rs->payload->early_dispatch_state.load(std::memory_order_relaxed) == PTO2_EARLY_DISPATCH_STAGING) { continue; } } @@ -340,13 +338,14 @@ void SchedulerContext::check_running_cores_for_completion( } #endif - // A pending task is "gated" when it is a speculative pre-stage still + // A pending task is "gated" when it is a early-dispatch pre-stage still // waiting on its doorbell (STAGED): 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. bool pending_gated = (core.pending_slot_state != nullptr && core.pending_slot_state->payload != nullptr && - core.pending_slot_state->payload->spec_state.load(std::memory_order_relaxed) == PTO2_SPEC_STAGING); + core.pending_slot_state->payload->early_dispatch_state.load(std::memory_order_relaxed) == + PTO2_EARLY_DISPATCH_STAGING); SlotTransition t = decide_slot_transition( reg_task_id, reg_state, core.running_reg_task_id, core.pending_reg_task_id, pending_gated ); @@ -377,7 +376,7 @@ void SchedulerContext::check_running_cores_for_completion( if (t.pending_done) { 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, local_bufs + completed_this_turn, deferred_release_slot_states, deferred_release_count #if SIMPLER_DFX , core.pending_dispatch_timestamp, finish_ts @@ -388,7 +387,7 @@ void SchedulerContext::check_running_cores_for_completion( if (t.running_done) { 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, local_bufs + completed_this_turn, deferred_release_slot_states, deferred_release_count #if SIMPLER_DFX , core.running_dispatch_timestamp, finish_ts @@ -514,6 +513,14 @@ void SchedulerContext::drain_worker_dispatch(int32_t block_num) { } } + // The drain path IS this sync_start producer's dispatch, so it must bump its + // consumers' dispatch_fanin like the normal dispatch path + // (scheduler_dispatch.cpp, post-publish) -- otherwise a consumer whose only + // flagged producer is a sync_start (drain-dispatched) task never becomes an + // early-dispatch candidate. Idempotent via propagate's dispatch_propagated + // once-guard; the internal gate no-ops for an unflagged producer. + sched_->propagate_dispatch_fanin(*slot_state); + // All blocks dispatched -- clear drain state. // Release fence ensures tracker mutations are visible to threads that // acquire-load sync_start_pending == 0 and resume normal operation. diff --git a/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_context.h b/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_context.h index 54f14ab7e..96b0c14d2 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_context.h +++ b/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_context.h @@ -54,15 +54,25 @@ class SchedulerContext { // Lifecycle // ========================================================================= - // Initialize scheduler state from the given runtime and thread layout. - // - Discovers cores via handshake_all_cores() - // - Assigns cores to scheduler threads - // - Resets task counters, payloads, per-core GlobalContext - // - Binds func_id_to_addr_ / initial sched_ (if rt is already known) - // - Captures AICore-register base (consumed by handshake_all_cores()) - // Returns 0 on success, negative on failure (handshake / assignment error). + // 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. + // + // 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. int32_t - init(Runtime *runtime, int32_t aicpu_thread_num, int32_t sched_thread_num, bool orch_to_sched, uint64_t regs_base); + 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. + void handshake_partition(Runtime *runtime, int32_t tidx, int32_t nthreads); + // Leader-only, after the handshake barrier: build worker-id lists, assign + // cores, init profiling subsystems, read task counts, init 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. @@ -84,14 +94,12 @@ class SchedulerContext { // - 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); - // Bind the PTO2Runtime scheduler pointer. Required in device-orchestration - // mode where rt is created by the orchestrator thread after init(). + // Bind the PTO2Runtime scheduler pointer. void bind_runtime(PTO2Runtime *rt); // ========================================================================= @@ -141,30 +149,31 @@ class SchedulerContext { // --- 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. - // 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}; - // --- Core-transition coordination --- - std::atomic transition_requested_{false}; - std::atomic wait_reassign_{0}; - std::atomic reassigned_{false}; - // --- Thread/core configuration --- int32_t active_sched_threads_{0}; int32_t sched_thread_num_{0}; - bool orch_to_sched_{false}; int32_t aicpu_thread_num_{0}; int32_t cores_total_num_{0}; - // Cluster-ordered worker_id lists, populated by handshake_all_cores(). + // Cluster-ordered worker_id lists, populated by post_handshake_init(). int32_t aic_worker_ids_[RUNTIME_MAX_WORKER]{}; int32_t aiv_worker_ids_[RUNTIME_MAX_WORKER]{}; 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}; @@ -179,15 +188,9 @@ class SchedulerContext { // Core management (scheduler_cold_path.cpp) // ========================================================================= - // Handshake with all AICore workers; populates core_exec_states_, worker id lists. - int32_t handshake_all_cores(Runtime *runtime); - // Assign discovered cores (cluster = 1 AIC + 2 AIV) round-robin across scheduler threads. bool assign_cores_to_threads(); - // Re-distribute all cores across all threads after orchestration completes. - void reassign_cores_for_all_threads(); - // Emergency shutdown: broadcast exit signal to every handshake'd core and // deinit their AICore register blocks. Idempotent. void emergency_shutdown(Runtime *runtime); @@ -214,8 +217,7 @@ class SchedulerContext { } int pop_ready_tasks_batch( - PTO2ResourceShape shape, int32_t thread_idx, PTO2LocalReadyBuffer &local_buf, PTO2TaskSlotState **out, - int max_count + PTO2ReadyQueue *queues, PTO2ResourceShape shape, int32_t thread_idx, PTO2TaskSlotState **out, int max_count ); void build_payload( @@ -259,15 +261,20 @@ class SchedulerContext { ); void dispatch_shape( - int32_t thread_idx, PTO2ResourceShape shape, CoreTracker::DispatchPhase phase, PTO2LocalReadyBuffer &local_buf, + int32_t thread_idx, PTO2ReadyQueue *disp_queues, PTO2ResourceShape shape, CoreTracker::DispatchPhase phase, CoreTracker &tracker, bool &entered_drain, bool &made_progress, bool &try_pushed ); - // Speculative early-dispatch (Hook 1). After normal dispatch leaves idle - // cores spare, pre-stage the consumers of any RUNNING flagged producer onto - // those cores with not_ready=1 (gated). Touches no dependency state — the - // task is released by the doorbell at its normal ready-pop (Hook 2). - int32_t try_speculative_early_dispatch(int32_t thread_idx); + // 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 not_ready=1 (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) @@ -280,11 +287,18 @@ class SchedulerContext { CoreTracker::BitStates &idle, CoreTracker::BitStates &pend ); + // 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; mid-flush of local buffers; 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. + // 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 @@ -293,25 +307,48 @@ class SchedulerContext { // 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, PTO2LocalReadyBuffer (&local_bufs)[PTO2_NUM_RESOURCE_SHAPES], - bool pmu_active, bool &made_progress, bool &try_pushed + int32_t thread_idx, CoreTracker &tracker, bool pmu_active, bool &made_progress, bool &try_pushed ); + // 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 anywhere this thread could see them: the caller's - // MIX local LIFO stack or the global MIX ready queue. Approximate — + // 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. Don't confuse with PTO2SpscQueue::size(), which uses acquire - // loads — that one isn't on this path. A stale read here causes at most one + // 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 PTO2LocalReadyBuffer &mix_local_buf) const { - return mix_local_buf.count > 0 || sched_->ready_queues[static_cast(PTO2ResourceShape::MIX)].size() > 0; + 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; } // ========================================================================= @@ -325,8 +362,7 @@ class SchedulerContext { void complete_slot_task( PTO2TaskSlotState &slot_state, int32_t expected_reg_task_id, PTO2SubtaskSlot subslot, int32_t thread_idx, int32_t core_id, Handshake *hank, int32_t &completed_this_turn, - PTO2TaskSlotState *deferred_release_slot_states[], int32_t &deferred_release_count, - PTO2LocalReadyBuffer *local_bufs + PTO2TaskSlotState *deferred_release_slot_states[], int32_t &deferred_release_count #if SIMPLER_DFX , uint64_t dispatch_ts, uint64_t finish_ts @@ -338,8 +374,7 @@ class SchedulerContext { void check_running_cores_for_completion( int32_t thread_idx, Handshake *hank, int32_t &completed_this_turn, int32_t &cur_thread_completed, - bool &made_progress, PTO2TaskSlotState *deferred_release_slot_states[], int32_t &deferred_release_count, - PTO2LocalReadyBuffer *local_bufs + bool &made_progress, PTO2TaskSlotState *deferred_release_slot_states[], int32_t &deferred_release_count ); bool enter_drain_mode(PTO2TaskSlotState *slot_state, int32_t block_num); @@ -354,8 +389,6 @@ class SchedulerContext { __attribute__((noinline, cold)) LoopAction handle_orchestrator_exit(int32_t thread_idx, PTO2SharedMemoryHeader *header, Runtime *runtime, int32_t &task_count); - __attribute__((noinline, cold)) LoopAction handle_core_transition(bool &cores_released); - __attribute__((noinline, cold)) LoopAction check_idle_fatal_error(int32_t thread_idx, PTO2SharedMemoryHeader *header, Runtime *runtime); diff --git a/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_dispatch.cpp b/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_dispatch.cpp index ef2700db6..70f3b6fd6 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_dispatch.cpp +++ b/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_dispatch.cpp @@ -44,10 +44,11 @@ namespace { inline constexpr int32_t PTO2_DEFERRED_RELEASE_CAP = 256; } -// The speculative core bitmask (PTO2_SPEC_CORE_MASK_WORDS * 64 bits) must cover +// 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_SPEC_CORE_MASK_WORDS * 64, "staged_core_mask too small for RUNTIME_MAX_WORKER cores" + 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) { @@ -83,7 +84,7 @@ bool SchedulerContext::has_idle_in_other_threads(int32_t self_thread_idx, PTO2Re } int SchedulerContext::pop_ready_tasks_batch( - PTO2ResourceShape shape, int32_t thread_idx, PTO2LocalReadyBuffer &local_buf, PTO2TaskSlotState **out, int max_count + PTO2ReadyQueue *queues, PTO2ResourceShape shape, int32_t thread_idx, PTO2TaskSlotState **out, int max_count ) { #if SIMPLER_DFX auto &l2_swimlane = sched_l2_swimlane_[thread_idx]; @@ -91,11 +92,11 @@ int SchedulerContext::pop_ready_tasks_batch( 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( - shape, local_buf, out, max_count, g_sched_pop_atomic_count[thread_idx], g_sched_pop_wait_cycle[thread_idx] + 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(shape, local_buf, out, max_count); + int count = sched_->get_ready_tasks_batch(queues, shape, out, max_count); #endif if (l2_swimlane_level_ >= L2SwimlaneLevel::SCHED_PHASES) { if (count > 0) { @@ -106,7 +107,7 @@ int SchedulerContext::pop_ready_tasks_batch( } #else (void)thread_idx; - int count = sched_->get_ready_tasks_batch(shape, local_buf, out, max_count); + int count = sched_->get_ready_tasks_batch(queues, shape, out, max_count); #endif return count; } @@ -132,11 +133,12 @@ void SchedulerContext::build_payload( 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: a task being staged (Hook 1 set spec_state to + // Early-dispatch: a task being staged (Hook 1 set early_dispatch_state to // STAGING before this call) is gated — the AICore must wait for the // DATA_MAIN_BASE high-32 doorbell. All other dispatches run on pickup. dispatch_payload.not_ready = - (slot_state.payload->spec_state.load(std::memory_order_relaxed) == PTO2_SPEC_STAGING) ? 1 : 0; + (slot_state.payload->early_dispatch_state.load(std::memory_order_relaxed) == PTO2_EARLY_DISPATCH_STAGING) ? 1 : + 0; } SchedulerContext::PublishHandle SchedulerContext::prepare_subtask_to_core( @@ -231,21 +233,21 @@ int SchedulerContext::prepare_block_for_dispatch( uint8_t cmask = slot_state.active_mask.core_mask(); int n = 0; 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, to_pending, - block_idx + 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, to_pending, - block_idx + 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, to_pending, - block_idx + thread_idx, tracker.get_aiv1_core_offset(core_offset), slot_state, PTO2SubtaskSlot::AIV1, p, block_idx ); } #if SIMPLER_DFX @@ -270,7 +272,7 @@ int SchedulerContext::prepare_block_for_dispatch( } void SchedulerContext::dispatch_shape( - int32_t thread_idx, PTO2ResourceShape shape, CoreTracker::DispatchPhase phase, PTO2LocalReadyBuffer &local_buf, + int32_t thread_idx, PTO2ReadyQueue *disp_queues, PTO2ResourceShape shape, CoreTracker::DispatchPhase phase, CoreTracker &tracker, bool &entered_drain, bool &made_progress, bool &try_pushed ) { #if SIMPLER_SCHED_PROFILING @@ -286,7 +288,7 @@ void SchedulerContext::dispatch_shape( while (cores.has_value() && !entered_drain) { int want = cores.count(); PTO2TaskSlotState *batch[CoreTracker::MAX_CLUSTERS * 3]; - int got = pop_ready_tasks_batch(shape, thread_idx, local_buf, batch, want); + int got = pop_ready_tasks_batch(disp_queues, shape, thread_idx, batch, want); if (got == 0) break; // sync_start exclusion gate. @@ -370,28 +372,28 @@ void SchedulerContext::dispatch_shape( } } if (!selected_mix_clusters.has_value()) { - sched_->ready_queues[static_cast(shape)].push(slot_state); + disp_queues[static_cast(shape)].push(slot_state); continue; } } - // (Speculative pre-staged tasks never reach this ready-pop: they are + // (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_speculative_release.) + // instant their last producer completes — see try_early_dispatch_release.) if (slot_state->active_mask.requires_sync_start()) { if (is_pending) { - sched_->ready_queues[static_cast(shape)].push(slot_state); + disp_queues[static_cast(shape)].push(slot_state); continue; } int32_t available = is_mix ? selected_mix_clusters.count() : cores.count(); if (available < slot_state->logical_block_num) { flush_publish(); if (!enter_drain_mode(slot_state, slot_state->logical_block_num)) { - sched_->ready_queues[static_cast(shape)].push(slot_state); + disp_queues[static_cast(shape)].push(slot_state); } for (int rem = bi + 1; rem < got; rem++) { - sched_->ready_queues[static_cast(shape)].push(batch[rem]); + disp_queues[static_cast(shape)].push(batch[rem]); } entered_drain = true; break; @@ -400,7 +402,7 @@ void SchedulerContext::dispatch_shape( if (!cores.has_value()) { flush_publish(); - sched_->ready_queues[static_cast(shape)].push_batch(&batch[bi], got - bi); + disp_queues[static_cast(shape)].push_batch(&batch[bi], got - bi); break; } @@ -426,7 +428,7 @@ void SchedulerContext::dispatch_shape( slot_state->next_block_idx.store(static_cast(start + claim), std::memory_order_relaxed); if (start + claim < slot_state->logical_block_num) { - sched_->ready_queues[static_cast(shape)].push(slot_state); + disp_queues[static_cast(shape)].push(slot_state); } for (int32_t b = 0; b < claim; b++) { @@ -467,12 +469,11 @@ void SchedulerContext::dispatch_shape( } } -void SchedulerContext::dispatch_ready_tasks( - int32_t thread_idx, CoreTracker &tracker, PTO2LocalReadyBuffer (&local_bufs)[PTO2_NUM_RESOURCE_SHAPES], - bool pmu_active, bool &made_progress, bool &try_pushed +template +void SchedulerContext::run_staging_order( + int32_t thread_idx, bool pmu_active, StageFn &&stage, ResidualMixFn &&residual_mix ) { 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 @@ -483,124 +484,41 @@ void SchedulerContext::dispatch_ready_tasks( }; const PTO2ResourceShape *aic_aiv = kAicAivOrder[thread_idx & 1]; - // Spill overflow from local_bufs to the shared ready queue BEFORE we start - // dispatching. release_fanin's fast path packs all newly-ready consumers - // into the producing thread's local_bufs (zero atomic, peer-invisible). For - // batch releases (e.g. attn_fence → 50 out_proj consumers) that - // overshoots this thread's slot budget so peers are starving while we - // hoard. The cross-thread invisibility window between "complete pushes 50 - // to local" and "IDLE-AIC's mid-phase flush exposes overflow to shared" - // is what shows up in the swimlane as the multi-microsecond inter-thread - // stagger on out_proj's first wave. - // - // Gate conditions: - // (a) local count exceeds this thread's per-shape block budget — we - // can't dispatch them all even with both RUNNING+PENDING slots; - // (b) at least one peer has idle cores in this shape — they want work. - // Both must hold to avoid wasting a CAS push when we could profitably - // self-dispatch the overflow. Condition (b) reads peer CoreTracker - // (plain 8-byte load on a rarely-contended cache line, ~5 ns) — we - // deliberately avoid ready_queues[s].size() here, which is two atomic - // loads on lines pushers + poppers actively bounce. - // - // Capacity derives from how cores are partitioned across sched threads: - // per-shape budget = (PLATFORM_MAX_BLOCKDIM / active_sched_threads_) - // × cores_per_blockdim_for_that_shape - // MIX is 1 cluster per block dim, so its budget equals the block-dim - // share without multiplying. - // - // Push the trailing `excess` slot pointers — O(1) count decrement, no - // memmove. push_batch is one CAS for the whole excess; peers see the - // batch immediately and can race for them. - const int32_t bd_per_thread = PLATFORM_MAX_BLOCKDIM / active_sched_threads_; - const int32_t thread_capacity[PTO2_NUM_RESOURCE_SHAPES] = { - /*AIC=*/bd_per_thread * PLATFORM_AIC_CORES_PER_BLOCKDIM, - /*AIV=*/bd_per_thread * PLATFORM_AIV_CORES_PER_BLOCKDIM, - /*MIX=*/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; - } - } - }; - // Every return path below must flush; wrap in RAII so we cannot forget. - // The mid-function flush between IDLE and PENDING is still called - // explicitly — guard only covers exit. - struct FlushGuard { - decltype(flush_local_bufs) &flush_fn; - ~FlushGuard() { flush_fn(); } - } flush_guard{flush_local_bufs}; - - bool entered_drain = false; - // ===== IDLE stage ===== - dispatch_shape( - thread_idx, PTO2ResourceShape::MIX, Phase::IDLE, local_bufs[MIX_I], tracker, entered_drain, made_progress, - try_pushed - ); - if (entered_drain) return; + if (stage(PTO2ResourceShape::MIX, Phase::IDLE)) 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 = has_residual_mix(local_bufs[MIX_I]); + bool skip_aic_aiv = residual_mix(); if (!skip_aic_aiv) { for (int i = 0; i < 2; i++) { - PTO2ResourceShape s = aic_aiv[i]; - dispatch_shape( - thread_idx, s, Phase::IDLE, local_bufs[static_cast(s)], tracker, entered_drain, made_progress, - try_pushed - ); - if (entered_drain) return; + if (stage(aic_aiv[i], Phase::IDLE)) 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 (already flushed above) 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. + // 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)) { - dispatch_shape( - thread_idx, PTO2ResourceShape::MIX, Phase::PENDING, local_bufs[MIX_I], tracker, entered_drain, - made_progress, try_pushed - ); - if (entered_drain) return; + if (stage(PTO2ResourceShape::MIX, Phase::PENDING)) 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 && has_residual_mix(local_bufs[MIX_I])) { + if (!skip_aic_aiv && residual_mix()) { skip_aic_aiv = true; } - // PENDING-MIX may have re-populated AIC/AIV local_bufs via release_fanin - // during in-flight completions; flush_guard ensures these don't carry - // across to the next iteration's IDLE stage. if (skip_aic_aiv) return; // AIC/AIV-PENDING gate: a peer-idle skip is a delay, not a loss — the peer @@ -608,14 +526,50 @@ void SchedulerContext::dispatch_ready_tasks( for (int i = 0; i < 2; i++) { PTO2ResourceShape s = aic_aiv[i]; if (has_idle_in_other_threads(thread_idx, s)) continue; - dispatch_shape( - thread_idx, s, Phase::PENDING, local_bufs[static_cast(s)], tracker, entered_drain, made_progress, - try_pushed - ); - if (entered_drain) return; + 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 (the queue drain) has advanced // next_block_idx by `count` under pop-exclusivity AND re-pushed `c` for peers @@ -633,7 +587,7 @@ void SchedulerContext::dispatch_ready_tasks( // // Self-ring: release flips STAGING->DISPATCHED then rings the mask. A block staged // after that flip isn't in the mask release read, so this thread rings it here. The -// seq_cst order between "OR mask then load spec_state" (here) and "store DISPATCHED +// seq_cst order between "OR mask then load early_dispatch_state" (here) and "store DISPATCHED // then read mask" (release) guarantees every gated core's doorbell fires. int32_t SchedulerContext::stage_consumer_blocks( int32_t thread_idx, PTO2TaskSlotState *c, PTO2ResourceShape shape, int32_t start, int32_t count, @@ -643,17 +597,19 @@ int32_t SchedulerContext::stage_consumer_blocks( // 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_SPEC_CORE_MASK_WORDS] = {0}; // cores this thread gated (for self-ring) + uint64_t my_cores[PTO2_EARLY_DISPATCH_CORE_MASK_WORDS] = {0}; // cores this thread gated (for self-ring) int32_t staged = 0; int32_t block = start; - auto stage_from = [&](CoreTracker::BitStates &avail, bool to_pending) { - // Mirror the normal flush_publish (scheduler_dispatch.cpp wmb()+publish loop): - // prepare all claimed blocks' payloads, one wmb(), then publish. The wmb - // guarantees the not_ready gate + 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 doorbell/release path mirrors normal dispatch). - PublishHandle handles[CoreTracker::MAX_CLUSTERS * 3]; - int n = 0; + // 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 + // not_ready gate + 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]); @@ -661,33 +617,35 @@ int32_t SchedulerContext::stage_consumer_blocks( count--; staged++; } - if (n == 0) return; + }; + 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); int32_t cid = tracker.get_core_id_by_offset(handles[i].core_offset); - sched_->spec_doorbell_table[cid].addr = handles[i].reg_addr; - sched_->spec_doorbell_table[cid].token = handles[i].reg_task_id; + 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)); } - }; - if (idle.has_value()) stage_from(idle, /*to_pending=*/false); - if (pend.has_value()) stage_from(pend, /*to_pending=*/true); + } // 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_SPEC_CORE_MASK_WORDS; w++) + 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); // If release already flipped DISPATCHED, it may have read the mask before our // bits landed — ring our own cores so none is left gated forever. - if (staged > 0 && c->payload->spec_state.load(std::memory_order_seq_cst) == PTO2_SPEC_DISPATCHED) { - for (int w = 0; w < PTO2_SPEC_CORE_MASK_WORDS; w++) { + if (staged > 0 && + c->payload->early_dispatch_state.load(std::memory_order_seq_cst) == PTO2_EARLY_DISPATCH_DISPATCHED) { + for (int w = 0; w < PTO2_EARLY_DISPATCH_CORE_MASK_WORDS; w++) { uint64_t bits = my_cores[w]; while (bits != 0) { int cid = w * 64 + __builtin_ctzll(bits); bits &= bits - 1; PTO2SchedulerState::ring_one_doorbell( - sched_->spec_doorbell_table[cid].addr, sched_->spec_doorbell_table[cid].token + sched_->early_dispatch_doorbell_table[cid].addr, sched_->early_dispatch_doorbell_table[cid].token ); } } @@ -695,34 +653,68 @@ int32_t SchedulerContext::stage_consumer_blocks( return staged; } -// Early-dispatch drain (idle pass). Candidates are pushed to early_dispatch_queue -// EVENT-DRIVEN by propagate_dispatch_fanin (a flagged producer's dispatch bumps its -// consumers' dispatch_fanin; reaching fanin_count enqueues the consumer) — there is -// no per-iteration PULL scan here anymore. This pass only DRAINS the queue. -// Returns the number of blocks staged this pass (for the EarlyDispatch swimlane bar). -int32_t SchedulerContext::try_speculative_early_dispatch(int32_t thread_idx) { - constexpr int PTO2_EARLY_DISPATCH_DRAIN_MAX = 8; // bounded pops per pass +// 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. Candidates are pushed to the shape's queue EVENT-DRIVEN by +// propagate_dispatch_fanin, 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 total_staged = 0; + 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; - // Drain the queue — mirrors the normal SPMD dispatch path. Pop a 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 instead of a serial winner-then-peer daisy chain. Bounded pops/pass. - for (int n = 0; n < PTO2_EARLY_DISPATCH_DRAIN_MAX; n++) { - PTO2TaskSlotState *c = sched_->early_dispatch_queue.pop(); - if (c == nullptr) break; - if (c->payload->spec_state.load(std::memory_order_acquire) != PTO2_SPEC_STAGING) continue; // released - PTO2ResourceShape shape = c->active_mask.to_shape(); - auto idle = tracker.get_idle_core_offset_states(shape); - auto pend = tracker.get_pending_core_offset_states(shape); - int32_t freecores = (idle.has_value() ? idle.count() : 0) + (pend.has_value() ? pend.count() : 0); - if (freecores == 0) { // no free cores of this shape — give it back for peers and stop - sched_->early_dispatch_queue.push(c); + int32_t total_staged = 0; + PTO2TaskSlotState *batch[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(batch, cores.count()); + for (int bi = 0; bi < got; bi++) { + PTO2TaskSlotState *c = batch[bi]; + 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(&batch[bi], got - bi); break; } // CAS-claim a contiguous range [start, start+claim) sized to this thread's @@ -744,12 +736,66 @@ int32_t SchedulerContext::try_speculative_early_dispatch(int32_t thread_idx) { 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_queue.push(c)) + if (!sched_->early_dispatch_queues[s].push(c)) LOG_INFO_V9( - "[SPEC] queue full on re-push, consumer=%" PRId64, static_cast(c->task->task_id.raw) + "[EARLY_DISPATCH] queue full on re-push, consumer=%" PRId64, + static_cast(c->task->task_id.raw) ); } - total_staged += stage_consumer_blocks(thread_idx, c, shape, start, claim, idle, pend); + // 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), mirroring dispatch_ready_tasks: owns its own +// gating and progress-flag updates, and orders staging the same way — MIX strict +// priority, IDLE stage before PENDING stage, cross-thread idle gating +// (MIX-IDLE ▶ c/v-IDLE ▶ MIX-PEND ▶ c/v-PEND). sync_start doesn't apply here (those +// tasks are excluded from early dispatch at push time, propagate_dispatch_fanin). +// 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; + } + + // 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; } return total_staged; } @@ -794,22 +840,20 @@ int32_t SchedulerContext::resolve_and_dispatch(Runtime *runtime, int32_t thread_ l2_swimlane.l2_swimlane_enabled = (l2_swimlane_level_ != L2SwimlaneLevel::DISABLED); #endif - 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); - } PTO2TaskSlotState *deferred_release_slot_states[PTO2_DEFERRED_RELEASE_CAP]; int32_t deferred_release_count = 0; - bool cores_released = false; - // 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(); @@ -817,34 +861,27 @@ int32_t SchedulerContext::resolve_and_dispatch(Runtime *runtime, int32_t thread_ #if SIMPLER_DFX // Queue-depth snapshot carried across the iteration boundary: each phase - // emit consumes (phase_start_*) and refreshes them with its own end snapshot - // so the next phase's "at_start" equals the previous phase's "at_end". + // 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.** Local depth (this thread's PTO2LocalReadyBuffer) - // is a single int read on a register-cached stack — free. 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 flush_local_bufs + every pop). With both phases emitting - // per iter that's 12 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; the - // big transitions (local→shared flush) still show up across iter boundaries. + // **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_local[L2SWIMLANE_NUM_QUEUE_SHAPES] = {0}; 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 capture_local_snapshot = [&](int16_t local_out[L2SWIMLANE_NUM_QUEUE_SHAPES]) { - for (int s = 0; s < L2SWIMLANE_NUM_QUEUE_SHAPES; s++) { - local_out[s] = static_cast(local_bufs[s].count); - } - }; auto get_or_sample_shared = [&]() -> const int16_t * { if (!iter_shared_sampled) { // Clamp to int16_t max before narrowing. PTO2_PROF_READYQUEUE_SIZE @@ -860,15 +897,22 @@ int32_t SchedulerContext::resolve_and_dispatch(Runtime *runtime, int32_t thread_ } return iter_shared_snapshot; }; - auto capture_phase_end = [&](int16_t local_out[L2SWIMLANE_NUM_QUEUE_SHAPES], - int16_t shared_out[L2SWIMLANE_NUM_QUEUE_SHAPES]) { - capture_local_snapshot(local_out); + 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_local, phase_start_shared); + capture_phase_end(phase_start_shared); } #endif @@ -905,11 +949,6 @@ int32_t SchedulerContext::resolve_and_dispatch(Runtime *runtime, int32_t thread_ if (action == LoopAction::BREAK_LOOP) break; } - if (!cores_released && orch_to_sched_) { - LoopAction action = handle_core_transition(cores_released); - if (action == LoopAction::BREAK_LOOP) break; - } - #if SIMPLER_DFX CYCLE_COUNT_LAP(l2_swimlane.sched_idle_cycle); #endif @@ -921,7 +960,7 @@ int32_t SchedulerContext::resolve_and_dispatch(Runtime *runtime, int32_t thread_ 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, local_bufs + deferred_release_slot_states, deferred_release_count ); } if (completed_this_turn > 0) { @@ -945,7 +984,7 @@ int32_t SchedulerContext::resolve_and_dispatch(Runtime *runtime, int32_t thread_ if (rt_ != nullptr && rt_->aicore_mailbox != nullptr && (sched_->async_wait_list.count > 0 || rt_->aicore_mailbox->has_pending())) { AsyncPollResult poll_result = sched_->async_wait_list.poll_and_complete( - rt_->aicore_mailbox, sched_, local_bufs, deferred_release_slot_states, deferred_release_count, + rt_->aicore_mailbox, sched_, deferred_release_slot_states, deferred_release_count, PTO2_DEFERRED_RELEASE_CAP #if SIMPLER_SCHED_PROFILING , @@ -982,24 +1021,17 @@ int32_t SchedulerContext::resolve_and_dispatch(Runtime *runtime, int32_t thread_ // iteration; on a pure-retire iteration phase_complete_count is 0). if (l2_swimlane_level_ >= L2SwimlaneLevel::SCHED_PHASES && (l2_swimlane.phase_complete_count > 0 || l2_swimlane.phase_subretire_count > 0)) { - // Local depth is cheap (this thread's own buffer counter). - // Shared depth is NOT sampled here: complete's release_fanin - // pushes to local_bufs in the fast path (try_push succeeds - // until cap=64). Shared only changes on dispatch's flush - // path. Carrying phase_start_shared forward as end_shared - // is the right answer 99% of the time AND skips three - // contended atomic loads per emit. - int16_t phase_end_local[L2SWIMLANE_NUM_QUEUE_SHAPES]; - capture_local_snapshot(phase_end_local); + // 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_start_shared + /*pop_miss=*/0, phase_start_shared, phase_end_shared ); - for (int s = 0; s < L2SWIMLANE_NUM_QUEUE_SHAPES; s++) { - phase_start_local[s] = phase_end_local[s]; - // phase_start_shared unchanged — carried forward - } + for (int s = 0; s < L2SWIMLANE_NUM_QUEUE_SHAPES; s++) + phase_start_shared[s] = phase_end_shared[s]; _t0_phase = _t1; l2_swimlane.phase_complete_count = 0; l2_swimlane.phase_subretire_count = 0; @@ -1015,41 +1047,15 @@ int32_t SchedulerContext::resolve_and_dispatch(Runtime *runtime, int32_t thread_ continue; } - // Phase 3: Drain wiring queue (thread 0 only) - int wired = 0; - if (thread_idx == 0) { - wired = sched_->drain_wiring_queue(orchestrator_done_); - if (wired > 0) { - made_progress = true; -#if SIMPLER_SCHED_PROFILING - l2_swimlane.phase_wiring_count += wired; -#endif - } - } -#if SIMPLER_DFX - CYCLE_COUNT_LAP(l2_swimlane.sched_wiring_cycle); - // Wire outer phase: emit one bar covering this iter's drain_wiring_queue - // pass when it wired any tasks. tasks_processed = wired count. Resolve - // does NOT nest under Wire — wiring only enqueues, the consumer release - // happens later in Complete/Dummy. - if (l2_swimlane_level_ >= L2SwimlaneLevel::SCHED_PHASES && wired > 0) { - int16_t phase_end_local[L2SWIMLANE_NUM_QUEUE_SHAPES]; - capture_local_snapshot(phase_end_local); - for (int s = 0; s < L2SWIMLANE_NUM_QUEUE_SHAPES; s++) { - phase_start_local[s] = phase_end_local[s]; - } - _t0_phase = _t1; - } -#endif - - // Phase 3b: Drain dummy ready queue (thread 0 only). + // 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. Pinned to thread 0 to avoid cross-thread - // races and to keep cache hot near the wiring drain above. - if (thread_idx == 0) { - constexpr int DUMMY_DRAIN_BATCH = 16; + // 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 @@ -1094,9 +1100,9 @@ int32_t SchedulerContext::resolve_and_dispatch(Runtime *runtime, int32_t thread_ // 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, local_bufs).fanout_edges; + dummy_consumers = sched_->on_task_complete(dummy_slot, thread_idx).fanout_edges; #else - dummy_consumers = sched_->on_task_complete(dummy_slot, local_bufs); + dummy_consumers = sched_->on_task_complete(dummy_slot); #endif #if SIMPLER_DFX if (dummy_resolve_t0 != 0) { @@ -1139,17 +1145,16 @@ int32_t SchedulerContext::resolve_and_dispatch(Runtime *runtime, int32_t thread_ // following Dispatch / EarlyDispatch / second-Complete bars start // at this end. if (dummy_outer_t0 != 0) { - int16_t phase_end_local[L2SWIMLANE_NUM_QUEUE_SHAPES]; - capture_local_snapshot(phase_end_local); 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_start_shared + /*pop_miss=*/0, phase_start_shared, phase_end_shared ); - for (int s = 0; s < L2SWIMLANE_NUM_QUEUE_SHAPES; s++) { - phase_start_local[s] = phase_end_local[s]; - } + 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 @@ -1164,7 +1169,7 @@ int32_t SchedulerContext::resolve_and_dispatch(Runtime *runtime, int32_t thread_ #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, local_bufs, pmu_active, made_progress, try_pushed); + 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 / @@ -1172,22 +1177,20 @@ int32_t SchedulerContext::resolve_and_dispatch(Runtime *runtime, int32_t thread_ // 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 && try_pushed && l2_swimlane.phase_dispatch_count > 0) { + 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_local[L2SWIMLANE_NUM_QUEUE_SHAPES]; int16_t phase_end_shared[L2SWIMLANE_NUM_QUEUE_SHAPES]; - capture_phase_end(phase_end_local, phase_end_shared); + 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_local[s] = phase_end_local[s]; phase_start_shared[s] = phase_end_shared[s]; } _t0_phase = dispatch_t1; @@ -1197,26 +1200,17 @@ int32_t SchedulerContext::resolve_and_dispatch(Runtime *runtime, int32_t thread_ } #endif - // Phase 4b: early-dispatch onto spare cores, but ONLY when this thread is - // otherwise idle — nothing was dispatched this iteration AND no ready work is - // queued for any shape. Early-dispatch competes with normal dispatch for - // pending slots, so gating on "no ready work" keeps it from delaying a real - // ready task; skipping the producer-fanout scan when busy also removes its - // per-iteration cost (the discovery walk only runs on genuinely idle passes). - bool any_ready_work = try_pushed; - for (int s = 0; !any_ready_work && s < PTO2_NUM_RESOURCE_SHAPES; s++) { - if (sched_->ready_queues[s].size() > 0 || local_bufs[s].count > 0) any_ready_work = true; - } + // 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 - // Skip speculative early-dispatch under PMU: dispatch_ready_tasks already - // withholds PENDING dispatch when pmu_active to preserve single-issue PMU - // windows, and staging gated work into idle/pending slots would perturb the - // same windows. [[maybe_unused]] int32_t staged_count = - (pmu_active || any_ready_work) ? 0 : try_speculative_early_dispatch(thread_idx); + 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. @@ -1230,71 +1224,14 @@ int32_t SchedulerContext::resolve_and_dispatch(Runtime *runtime, int32_t thread_ // 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 following second-poll's Complete bar - // starts at the EarlyDispatch end, not before it (otherwise their - // spans overlap and the outer-phase mutual-exclusion breaks). + // 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 - // Second completion poll. dispatch_ready_tasks + try_speculative_early_dispatch - // above can take several us in a busy window; a producer block that FINs - // during them would otherwise wait for the NEXT iteration's top-of-loop - // Phase-1 poll (the ~7us detection latency that delays a flagged - // producer's doorbell). Re-polling here observes those FINs immediately, - // so the doorbell fires this iteration. Idempotent (the poll is a poll); - // we drain deferred releases eagerly to keep the buffer from growing. #if SIMPLER_DFX - uint64_t complete2_t0 = (l2_swimlane_level_ >= L2SwimlaneLevel::SCHED_PHASES) ? get_sys_cnt_aicpu() : 0; -#endif - if (tracker.has_any_running_cores()) { - int32_t completed_2nd = 0; - check_running_cores_for_completion( - thread_idx, hank, completed_2nd, cur_thread_completed, made_progress, deferred_release_slot_states, - deferred_release_count, local_bufs - ); - if (completed_2nd > 0) { -#if SIMPLER_SCHED_PROFILING - sched_->tasks_completed.fetch_add(completed_2nd, std::memory_order_relaxed); -#endif - completed_tasks_.fetch_add(completed_2nd, std::memory_order_relaxed); - last_progress_count = completed_tasks_.load(std::memory_order_relaxed); - } - // Eager drain so the second poll can't push deferred_release toward - // its cap between idle iterations. - while (deferred_release_count >= PTO2_DEFERRED_RELEASE_CAP - 96) { -#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 - // Complete2 outer phase: covers second-poll FIN observation. Without - // this emit, FIN counts from the second poll would carry over into the - // next iter's first-Complete bar and be displayed with a span that - // doesn't actually include those FINs' timestamps (visible mismatch - // between Complete bar span and per-task finish_time in Worker / - // Scheduler View). - if (complete2_t0 != 0 && (l2_swimlane.phase_complete_count > 0 || l2_swimlane.phase_subretire_count > 0)) { - uint64_t complete2_t1 = get_sys_cnt_aicpu(); - int16_t phase_end_local[L2SWIMLANE_NUM_QUEUE_SHAPES]; - capture_local_snapshot(phase_end_local); - l2_swimlane_aicpu_record_sched_phase( - thread_idx, L2SwimlaneSchedPhaseKind::Complete, complete2_t0, complete2_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_start_shared - ); - for (int s = 0; s < L2SWIMLANE_NUM_QUEUE_SHAPES; s++) { - phase_start_local[s] = phase_end_local[s]; - } - _t0_phase = complete2_t1; - l2_swimlane.phase_complete_count = 0; - l2_swimlane.phase_subretire_count = 0; - } - // 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) { @@ -1415,8 +1352,8 @@ int32_t SchedulerContext::resolve_and_dispatch(Runtime *runtime, int32_t thread_ #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 + // dispatch emit (typically the trailing idle loops while waiting for the + // last in-flight tasks to complete) 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. @@ -1427,9 +1364,8 @@ int32_t SchedulerContext::resolve_and_dispatch(Runtime *runtime, int32_t thread_ 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_local[L2SWIMLANE_NUM_QUEUE_SHAPES]; int16_t phase_end_shared[L2SWIMLANE_NUM_QUEUE_SHAPES]; - capture_phase_end(phase_end_local, phase_end_shared); + 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), diff --git a/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_types.h b/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_types.h index 33f94cb09..757ea3ad5 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_types.h +++ b/src/a2a3/runtime/host_build_graph/runtime/scheduler/scheduler_types.h @@ -74,13 +74,11 @@ constexpr int32_t FATAL_ERROR_CHECK_INTERVAL = 1024; // Check orchestrator erro // 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) -// because the safe value differs per variant: onboard trims it to 2 s so the -// AICPU detects a hang and flushes its diagnostics (tensor dump, in-flight -// partial output) before STARS reaps the op and poisons the context (chain: -// this < op-exec < host stream-sync, platform_config.h); sim has no STARS to -// race and keeps the full 5 s #897 headroom. See spin_hint.h for the per-variant -// rationale. +// 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); @@ -221,6 +219,25 @@ class alignas(64) CoreTracker { bool has_any_running_cores() const { return ((~core_states_) & (aic_mask_ | aiv_mask_)).has_value(); } + // True if any core on this thread still has a free slot to stage onto — an + // idle core (running slot) or a running core with a free pending slot. A + // core is unavailable only when running AND its pending slot is occupied; + // idle cores keep pending_occupied_ clear by invariant, so + // ~pending_occupied_ over all cores is exactly "has a free slot". Purely + // local (no shared/atomic access) — used to skip early dispatch, and its + // shared-queue pop, when this thread has no capacity at all. + // BitStates of every core on this thread with a free slot to stage onto: a + // core is unavailable only when running AND its pending slot is occupied. + // Idle cores keep pending_occupied_ clear by invariant, so ~pending_occupied_ + // over aic|aiv is exactly "has a free slot". Spans AIC+AIV, so its .count() is + // an upper bound on the early-dispatch drain's per-shape pop (never exceeds the + // thread's total free cores), and .has_value() is the has_any_free_slot() + // predicate that gates the Phase-4b early-dispatch pass. Purely local (no + // shared/atomic access). + BitStates get_free_slot_states() const { return (~pending_occupied_) & (aic_mask_ | aiv_mask_); } + + bool has_any_free_slot() const { return get_free_slot_states().has_value(); } + template int32_t get_running_count() const { if constexpr (CT == CoreType::AIC) { @@ -318,9 +335,14 @@ class alignas(64) CoreTracker { // Runtime MIX dispatch uses classify_mix_cluster() so the decision follows the task's active_mask. enum class MixPlacement : uint8_t { RUNNING, PENDING, REJECT }; - // A MIX block must place all cores named by active_mask the same way: - // all idle means running placement, all running means pending placement, - // and any mixed state is retried later. + // Placement for the cores named by active_mask, ignoring cores this task does + // not use. All used cores idle -> RUNNING placement (each to its running slot). + // Otherwise -> PENDING placement: at dispatch each used core is filled per its + // own state -- an idle core takes its running slot (and is marked running, so + // the completion poller, which scans only running cores, tracks its FIN), an + // already-running core takes its pending slot and executes after its in-flight + // task. REJECT only when a used core's pending slot is already occupied (no free + // slot) or the mask is empty. MixPlacement classify_mix_cluster(int32_t cluster_offset, uint8_t core_mask) const { BitStates used(0ULL); if (core_mask & PTO2_SUBTASK_MASK_AIC) { @@ -340,10 +362,7 @@ class alignas(64) CoreTracker { if (idle.count() == used.count()) { return MixPlacement::RUNNING; } - if (!idle.has_value()) { - return MixPlacement::PENDING; - } - return MixPlacement::REJECT; + return MixPlacement::PENDING; } BitStates get_mix_running_cluster_offset_states(uint8_t core_mask) const { @@ -432,7 +451,6 @@ struct alignas(64) SchedL2SwimlaneCounters { uint64_t sched_start_ts{0}; uint64_t sched_complete_cycle{0}; uint64_t sched_dispatch_cycle{0}; - uint64_t sched_wiring_cycle{0}; uint64_t sched_idle_cycle{0}; uint64_t sched_loop_count{0}; uint32_t phase_complete_count{0}; @@ -450,7 +468,6 @@ struct alignas(64) SchedL2SwimlaneCounters { uint64_t pop_hit_at_last_emit{0}; uint64_t pop_miss_at_last_emit{0}; #if SIMPLER_SCHED_PROFILING - uint32_t phase_wiring_count{0}; uint64_t complete_probe_count{0}; uint64_t complete_hit_count{0}; uint64_t sched_complete_perf_cycle{0}; diff --git a/src/a2a3/runtime/host_build_graph/runtime/shared/pto_runtime2_init.cpp b/src/a2a3/runtime/host_build_graph/runtime/shared/pto_runtime2_init.cpp index 8626516f2..fdddade21 100644 --- a/src/a2a3/runtime/host_build_graph/runtime/shared/pto_runtime2_init.cpp +++ b/src/a2a3/runtime/host_build_graph/runtime/shared/pto_runtime2_init.cpp @@ -108,20 +108,23 @@ void PTO2SchedulerState::RingSchedState::destroy() { ring = nullptr; } PTO2SchedulerLayout PTO2SchedulerState::reserve_layout(DeviceArena &arena, int32_t dep_pool_capacity) { PTO2SchedulerLayout layout{}; layout.ready_queue_capacity = PTO2_READY_QUEUE_SIZE; - layout.spsc_capacity = PTO2_WRIRING_QUEUE_SIZE; layout.dep_pool_capacity = dep_pool_capacity; for (int i = 0; i < PTO2_NUM_RESOURCE_SHAPES; i++) { layout.off_ready_queue_slots[i] = ready_queue_reserve_layout(arena, PTO2_READY_QUEUE_SIZE); } + 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); - layout.off_early_dispatch_queue_slots = ready_queue_reserve_layout(arena, PTO2_EARLY_DISPATCH_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); + } // Force a cache-line base so writes from scheduler thread 0 (sole writer of // the dep_pool) do not invalidate adjacent multi-threaded regions like // ready_queue.slots. layout.off_dep_pool_entries = arena.reserve(static_cast(dep_pool_capacity) * sizeof(PTO2DepListEntry), PTO2_ALIGN_SIZE); - layout.off_wiring_spsc_buffer = PTO2SpscQueue::reserve_layout(arena, PTO2_WRIRING_QUEUE_SIZE); return layout; } @@ -146,15 +149,25 @@ bool PTO2SchedulerState::init_data_from_layout( 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; } - if (!ready_queue_init_data_from_layout( - &sched->early_dispatch_queue, arena, layout.off_early_dispatch_queue_slots, PTO2_EARLY_DISPATCH_QUEUE_SIZE - )) { - 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; + } } auto *orch_err = pto2_sm_layout::orch_error_code_addr(sm_dev_base); @@ -162,13 +175,6 @@ bool PTO2SchedulerState::init_data_from_layout( memset(dep_entries, 0, static_cast(layout.dep_pool_capacity) * sizeof(PTO2DepListEntry)); sched->ring_sched_state.dep_pool.init(dep_entries, layout.dep_pool_capacity, orch_err); - if (!sched->wiring.queue.init_data_from_layout(arena, layout.off_wiring_spsc_buffer, layout.spsc_capacity)) { - return false; - } - sched->wiring.batch_count = 0; - sched->wiring.batch_index = 0; - sched->wiring.backoff_counter = 0; - return true; } @@ -177,23 +183,33 @@ void PTO2SchedulerState::wire_arena_pointers(const PTO2SchedulerLayout &layout, 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); - ready_queue_wire_arena_pointers(&sched->early_dispatch_queue, arena, layout.off_early_dispatch_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] + ); + } sched->ring_sched_state.dep_pool.base = static_cast(arena.region_ptr(layout.off_dep_pool_entries)); - sched->wiring.queue.wire_arena_pointers(arena, layout.off_wiring_spsc_buffer); } void PTO2SchedulerState::destroy() { PTO2SchedulerState *sched = this; sched->ring_sched_state.destroy(); sched->ring_sched_state.dep_pool.base = nullptr; - 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); - ready_queue_destroy(&sched->early_dispatch_queue); + for (int i = 0; i < PTO2_NUM_RESOURCE_SHAPES; i++) { + ready_queue_destroy(&sched->early_dispatch_queues[i]); + } } // =============================================================================