diff --git a/src/a5/platform/onboard/aicore/inner_kernel.h b/src/a5/platform/onboard/aicore/inner_kernel.h index 2f4296030f..96d8c3d6e6 100644 --- a/src/a5/platform/onboard/aicore/inner_kernel.h +++ b/src/a5/platform/onboard/aicore/inner_kernel.h @@ -68,6 +68,20 @@ __aicore__ inline uint64_t read_reg(RegId reg) { } } +/** + * Read the high 32 bits of DATA_MAIN_BASE. + * + * AICore reads the full 64-bit SPR via MOV; the high half is the + * early-dispatch doorbell written by AICPU (low half stays the dispatch + * token). Read-only on the AICore side, so this is always valid (unlike writes + * to DATA_MAIN_BASE, which the SPR-write port rejects). + */ +__aicore__ inline uint32_t read_dmb_high32() { + uint64_t v; + __asm__ volatile("MOV %0, DATA_MAIN_BASE\n" : "=l"(v)); + return static_cast(v >> 32); +} + /** * Write to an AICore register * diff --git a/src/a5/platform/sim/aicore/inner_kernel.h b/src/a5/platform/sim/aicore/inner_kernel.h index 58dd448871..1ffa776928 100644 --- a/src/a5/platform/sim/aicore/inner_kernel.h +++ b/src/a5/platform/sim/aicore/inner_kernel.h @@ -176,6 +176,19 @@ inline uint64_t read_reg(RegId reg) { return static_cast(__atomic_load_n(ptr, __ATOMIC_ACQUIRE)); } +/** + * Read the high 32 bits of DATA_MAIN_BASE (early-dispatch doorbell). + * The high word lives one 32-bit slot above the dispatch token. + * + * Same atomic-acquire rule as read_reg(): in sim the cell is plain host memory + * shared with the AICPU thread that rings via a 64-bit STR of (token<<32)|token. + */ +inline uint32_t read_dmb_high32() { + uint32_t offset = reg_offset(RegId::DATA_MAIN_BASE); + volatile uint32_t *ptr = reinterpret_cast(sparse_reg_ptr(sim_get_reg_base(), offset + 4)); + return __atomic_load_n(ptr, __ATOMIC_ACQUIRE); +} + /** * Write to an AICore register in simulated register memory * diff --git a/src/a5/runtime/tensormap_and_ringbuffer/aicore/aicore_executor.cpp b/src/a5/runtime/tensormap_and_ringbuffer/aicore/aicore_executor.cpp index 5fb76f4149..0cd99085e3 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/aicore/aicore_executor.cpp +++ b/src/a5/runtime/tensormap_and_ringbuffer/aicore/aicore_executor.cpp @@ -119,6 +119,7 @@ __aicore__ __attribute__((weak)) void aicore_execute(__gm__ Runtime *runtime, in // Register encoding: AICPU_IDLE_TASK_ID=idle, task_id=task, AICORE_EXIT_SIGNAL=exit uint32_t reg_val = AICPU_IDLE_TASK_ID; uint32_t last_reg_val = AICPU_IDLE_TASK_ID; + bool exiting = false; while (true) { reg_val = static_cast(read_reg(RegId::DATA_MAIN_BASE)); @@ -142,6 +143,9 @@ __aicore__ __attribute__((weak)) void aicore_execute(__gm__ Runtime *runtime, in // hardware-bound) and the AICore-local dcci+ack cost // (receive_time → start_time, software-tunable). Stored in the // record as a 32-bit delta `start_time - receive_time`. + // + // Early-dispatch (src_payload != 0): receive_time stays at pickup — + // before the doorbell wait — so it may precede the producer's end. uint64_t receive_time = get_sys_cnt_aicore(); uint32_t task_id = reg_val; // Decode: register holds task_id directly @@ -157,6 +161,48 @@ __aicore__ __attribute__((weak)) void aicore_execute(__gm__ Runtime *runtime, in // Invalidate payload buffer (AICPU updates its content each dispatch) dcci(exec_payload, ENTIRE_DATA_CACHE); + // Early-dispatch gate. A gated task was staged on this core before its + // dependencies resolved; wait until AICPU rings the doorbell + // (DATA_MAIN_BASE high 32 == task_id) before executing. The ACK is + // deferred until AFTER the gate so the scheduler keeps the core + // off-limits (pending_occupied stays set) while the task is gated. + // src_payload == 0 (ready path) skips this; a non-zero src_payload is + // both the gate flag and the source PTO2TaskPayload. + if (exec_payload->src_payload != 0) { + __gm__ char *src = reinterpret_cast<__gm__ char *>(exec_payload->src_payload); + int32_t tensor_count = *reinterpret_cast<__gm__ int32_t *>(src + PTO2_TASKPAYLOAD_TENSOR_COUNT_OFFSET); + int32_t scalar_count = *reinterpret_cast<__gm__ int32_t *>(src + PTO2_TASKPAYLOAD_SCALAR_COUNT_OFFSET); + __gm__ uint64_t *src_scalars = + reinterpret_cast<__gm__ uint64_t *>(src + PTO2_TASKPAYLOAD_SCALARS_OFFSET); + int n = 0; + for (int32_t i = 0; i < tensor_count; i++) { + exec_payload->args[n++] = reinterpret_cast( + src + PTO2_TASKPAYLOAD_TENSORS_OFFSET + i * PTO2_TASKPAYLOAD_TENSOR_STRIDE + ); + } + for (int32_t i = 0; i < scalar_count; i++) { + exec_payload->args[n++] = src_scalars[i]; + } + OUT_OF_ORDER_STORE_BARRIER(); + while (true) { + if (read_dmb_high32() == task_id) { + if (static_cast(read_reg(RegId::DATA_MAIN_BASE)) == AICORE_EXIT_SIGNAL) { + exiting = true; + } + break; + } + if (static_cast(read_reg(RegId::DATA_MAIN_BASE)) == AICORE_EXIT_SIGNAL) { + exiting = true; + break; + } + SPIN_WAIT_HINT(); + } + if (exiting) { + write_reg(RegId::COND, AICORE_EXITED_VALUE); + break; + } + } + write_reg(RegId::COND, MAKE_ACK_VALUE(task_id)); // Performance profiling: record start time diff --git a/src/a5/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md b/src/a5/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md index 01c4b09b09..2d4be2711d 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md +++ b/src/a5/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md @@ -546,11 +546,12 @@ Each scheduler thread runs a tight loop with two main phases: - Poll register `COND` on each managed core - When `TASK_FIN_STATE` detected: record completion timestamps, call `on_subtask_complete(task_id, subslot)` to increment the completion counter; when `completed_subtasks == total_required_subtasks`, trigger `on_task_complete(task_id)` which marks `task_state[slot] = COMPLETED`, acquires fanout lock, traverses fanout list (incrementing consumers' `fanin_refcount`), marks `task_state[slot] = CONSUMED`, and advances `last_task_alive` watermark -**Phase 2 — Dispatch**: +**Phase 2 — Dispatch** (full model in §8.6): - For each idle core: pop a task from the matching shape-based ready queue (lock-free MPMC Vyukov queue, one per resource shape) - Build `PTO2DispatchPayload` from `TaskDescriptor` with `task_id`, `subslot`, `kernel_id`, and `core_type` -- Write task pointer to `Handshake.task`, signal AICore via register `DATA_MAIN_BASE` +- Write task pointer to `Handshake.task`, signal AICore via register `DATA_MAIN_BASE` (DMB offset `0xD0` on a5) +- After normal ready queues are empty, Phase **4b** may stage speculative early-dispatch candidates onto spare slots (`early_dispatch_queues[]` / `early_sync_start_queue`) After these phases, the scheduler updates profiling headers and checks for termination (all tasks completed and orchestrator done). @@ -558,12 +559,14 @@ After these phases, the scheduler updates profiling headers and checks for termi Ready queues use a lock-free bounded MPMC (Vyukov) design: -- One `PTO2ReadyQueue` per resource shape (5 shapes: `AIC_ONLY`, `AIV_X1`, `AIV_X2`, `AIC_AIV_X1`, `AIC_AIV_X2`) -- **Push**: any thread (orchestrator via `init_task`, or scheduler on completion) pushes newly-ready tasks to the queue matching `task->active_mask.to_shape()` +- One `PTO2ReadyQueue` per resource shape (`MIX` / `AIC` / `AIV` in the production tensormap path) +- **Push**: any thread (orchestrator via wiring, or scheduler on completion) pushes newly-ready tasks to the queue matching `task->active_mask.to_shape()` - **Pop**: scheduler threads pop from the queue matching the idle core's resource shape - Per-slot sequence counters prevent ABA problems - `enqueue_pos` and `dequeue_pos` are on separate cache lines to avoid false sharing +Unlike a2a3, a5 does **not** keep a separate `ready_sync_queues[]` tier: ready `require_sync_start` cohorts share `ready_queues[]`. Speculative sync_start early candidates still use the dedicated `early_sync_start_queue` (see §8.6). + ### 8.4 Watermark Advancement (last_task_alive) After a task reaches state CONSUMED (4), the scheduler tries to advance `last_task_alive`: @@ -605,6 +608,84 @@ Private internals are split across three .cpp files by responsibility: `AicpuExecutor` calls neither `handshake_*`, `assign_*`, `reassign_*`, nor `emergency_shutdown` directly — they are private, invoked only by `init` and `on_orchestration_done`. +### 8.6 Dispatch model — two sources, sync tiers, occupancy order + +`resolve_and_dispatch` places ready and speculative work onto AICore cores under one +occupancy model (ported from a2a3 early-dispatch; a5 specifics called out below). Two +orthogonal axes decide *what* runs and *where*: + +- **Source** — `NORMAL` (all producers done; the task sits in a ready queue and launches on + pickup) vs `EARLY` (a *speculative* pre-stage of a not-yet-released task; its dispatch + payload carries a non-zero `src_payload` gate and launches later by a high-32 doorbell on + `DATA_MAIN_BASE`). Normal strictly precedes early. +- **Cohort** — `SYNC_START` (an SPMD cohort that must launch atomically) vs `REGULAR` (each + block launches independently). "is it ready" (source) and "does it need a rendezvous" + (cohort) are orthogonal. + +Within each source the occupancy order is **`sync_start` ▸ MIX ▸ AIC/AIV`** (shape), and per +shape **idle ▸ pending** (an idle core takes its running slot; a busy core takes its gated +pending slot, promoted on completion). a5 implements this order inline in +`dispatch_ready_tasks` / `try_early_dispatch` (no separate `run_staging_order` helper). + +#### Queues + +| Source | Regular lanes | sync_start lane | +| ------ | ------------- | --------------- | +| NORMAL (ready) | `ready_queues[MIX\|AIC\|AIV]` | *(same `ready_queues[]` — a5 has no `ready_sync_queues[]`)* | +| EARLY (speculative) | `early_dispatch_queues[MIX\|AIC\|AIV]` | `early_sync_start_queue` (single) | + +A task routes to the early sync lane iff `active_mask.requires_sync_start()`. Early +dispatch runs only once normal `ready_queues[]` are empty **and** the local +`CoreTracker` has a spare slot (`has_any_free_slot`, a2a3 #1288). + +**Direct-only eligibility (a2a3 #1285/#1292):** a consumer is an early candidate only when +every *direct* producer is flagged `allow_early_resolve` (slot-state hint from Arg, or +unconditional true for hidden alloc creators). There is no auto-chain inheritance. + +#### sync_start drain + rendezvous + +A sync_start cohort of `block_num` cores must occupy all its cores before any of them run. +When it cannot fit inline, `enter_drain_mode` arms a stop-the-world drain: + +1. **Single election** — a CAS on `sync_start_pending` makes drains mutually exclusive. +2. **All-or-nothing** — the elected thread checks global available capacity ≥ `block_num` + before staging; if short it aborts and retries after completions free cores. +3. **Parallel stage** — threads barrier, then each CAS-claims a block range and stages its + own cores with a non-zero `src_payload` gate (`drain_stage_cores`): idle → running, + busy → pending (`pending_gated` when still waiting for the doorbell). +4. **Rendezvous launch** — `running_slot_count` counts staged running-slot cores; when it + reaches `popcount(staged_core_mask)` **and** the producer has released, + `maybe_rendezvous_ring` rings every gated core's doorbell together — the cohort starts as one. + +Doorbell ownership is exclusive (`claim_all_staged_doorbell_bits` / +`claim_late_staged_doorbell_bits`); launch is latched via `early_dispatch_launch_state`. + +#### Early-candidate gate: producer must publish every block + +Producer-side `propagate_dispatch_fanin` no-ops until the producer is **fully published**: +`published_block_count == logical_block_num` (a2a3 #1326). Publication is recorded after +prepare/publish makes payload + low-32 dispatch tokens visible. Early queue entries carry a +`task_id_snapshot` generation tag (#1336) so recycled slots cannot revive stale candidates. + +Wiring and propagation serialize under `fanout_lock` with `PTO2_DISPATCH_PROPAGATED` so +late-wired consumers still receive exactly one early-candidate contribution per eligible +edge (#1405). + +#### a5 platform notes + +- `RUNTIME_MAX_WORKER = 108` — doorbell table / `staged_core_mask` words must cover 108 cores. +- AICore index fields are `s_block_idx` / `s_block_num` (not a2a3 `block_idx` / `block_num`). +- DMB MMIO offset is **`0xD0`** (a2a3 uses `0xA0`). Ready dispatch writes the low 32 bits; + early release rings high 32 via 64-bit STR `(token<<32)|token`. AICore gated path spins on + `read_dmb_high32() == task_id` before ACK. +- Ready path keeps `src_payload == 0` and ACK-then-execute behavior unchanged. + +#### MIX per-core placement + +A MIX task spans a cluster (1 AIC + 2 AIV). Gated MIX may place **per core** +(`to_pending && !is_core_idle`): idle cores → running, busy cores → pending. Cross-core start +skew within a block is tolerated by AICore incore synchronization. + --- ## 9. AICore Worker Interaction @@ -632,11 +713,13 @@ Instead of polling a shared-memory status flag, the production protocol uses har **AICore execution loop**: -1. Poll `DATA_MAIN_BASE` for value != AICPU_IDLE_TASK_ID +1. Poll `DATA_MAIN_BASE` (low 32) for value != AICPU_IDLE_TASK_ID 2. Read payload from `Handshake.task` -3. Write ACK to `COND` -4. Execute kernel function via `func_id_to_addr` lookup -5. Write FIN to `COND` +3. If `src_payload != 0` (early/gated): materialize args from `src_payload`, spin until + `read_dmb_high32() == task_id`, **then** ACK +4. Else (ready path): ACK immediately +5. Execute kernel function via `func_id_to_addr` lookup +6. Write FIN to `COND` ### 9.3 PTO2DispatchPayload @@ -649,8 +732,9 @@ Built by the scheduler from `PTO2TaskDescriptor`: | `kernel_id` | Function ID for this subtask slot | | `core_type` | AIC or AIV | | `function_bin_addr` | GM address of compiled kernel binary | +| `src_payload` | Non-zero ⇒ gated early path (AICore materializes args + waits high-32 doorbell) | | `num_args` | Number of arguments | -| `args[]` | Tensor addresses and scalar values | +| `args[]` | Tensor addresses and scalar values (filled by AICPU on ready path; by AICore when gated) | --- diff --git a/src/a5/runtime/tensormap_and_ringbuffer/orchestration/pto_arg_with_deps.h b/src/a5/runtime/tensormap_and_ringbuffer/orchestration/pto_arg_with_deps.h index 20144d3a1b..00c36211c5 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/orchestration/pto_arg_with_deps.h +++ b/src/a5/runtime/tensormap_and_ringbuffer/orchestration/pto_arg_with_deps.h @@ -54,15 +54,17 @@ class L0TaskArgsWithDeps : private L0TaskArgs { using L0TaskArgs::add_scalar; using L0TaskArgs::add_scalars; using L0TaskArgs::add_scalars_i32; + using L0TaskArgs::allow_early_resolve; // early-dispatch hint (getter) using L0TaskArgs::copy_scalars_from; + using L0TaskArgs::set_allow_early_resolve; // early-dispatch hint (setter) + using L0TaskArgs::set_task_timing_slot; // selective task-timing slot (setter) + using L0TaskArgs::task_timing_slot; // selective task-timing slot (getter) // Error / status — forward to Arg using L0TaskArgs::error_msg; using L0TaskArgs::has_error; using L0TaskArgs::launch_spec; using L0TaskArgs::set_error; - using L0TaskArgs::set_task_timing_slot; // selective task-timing slot (setter) - using L0TaskArgs::task_timing_slot; // selective task-timing slot (getter) // NOT exposed: set_dependencies, explicit_dep_count, explicit_dep, // explicit_deps_data — these are the primitive-layer dep API. Users of diff --git a/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto2_dispatch_payload.h b/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto2_dispatch_payload.h index cae2756253..c47f980e96 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto2_dispatch_payload.h +++ b/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto2_dispatch_payload.h @@ -20,7 +20,7 @@ * GlobalContext (sub_block_id) is initialized once at runtime startup via * init_global_context() and never modified afterwards. * - * LocalContext (block_idx, block_num) and args[] are rebuilt by build_payload() + * LocalContext (s_block_idx, s_block_num) and args[] are rebuilt by build_payload() * before each dispatch. Both context struct pointers are written into the * args[] suffix on every dispatch (since args[] is rebuilt entirely each time). * @@ -34,6 +34,7 @@ #pragma once +#include #include #include "arg_direction.h" @@ -57,6 +58,19 @@ static_assert( "GLOBAL_CONTEXT_INDEX out of sync with intrinsic.h" ); +// Byte offsets into PTO2TaskPayload used by AICore to materialize args[] on the +// not_ready (early-dispatch) path. The AICore .o does not include +// pto_runtime2_types.h, so it reads tensor_count / scalar_count / tensors[] / +// scalars[] through these constants. The AICPU side (scheduler_dispatch.cpp) +// static_asserts them against offsetof where the full struct is visible, so any +// PTO2TaskPayload layout drift fails the build rather than corrupting args[]. +constexpr uint32_t PTO2_TASKPAYLOAD_TENSOR_COUNT_OFFSET = 0; +constexpr uint32_t PTO2_TASKPAYLOAD_SCALAR_COUNT_OFFSET = 4; +// Cache line 9 (byte 576) holds the AICPU-only DispatchPredicate; tensors follow it. +constexpr uint32_t PTO2_TASKPAYLOAD_TENSORS_OFFSET = 640; +constexpr uint32_t PTO2_TASKPAYLOAD_SCALARS_OFFSET = 4736; +constexpr uint32_t PTO2_TASKPAYLOAD_TENSOR_STRIDE = 128; // sizeof(Tensor) + /** * Per-core dispatch payload: function address + args[] + SPMD context. * @@ -68,20 +82,40 @@ static_assert( * concurrently dispatched cores. */ struct alignas(64) PTO2DispatchPayload { - uint64_t function_bin_addr; /**< Kernel entry address in GM (set by Scheduler) */ - uint64_t args[PTO2_DISPATCH_MAX_ARGS]; /**< Kernel arguments (GM pointers + scalars + ext params) */ + // === Cache line 0 (64B): control block, the only line written per dispatch === + // function_bin_addr, local_context.{s_block_idx,s_block_num,async_ctx.task_token} + // and src_payload are the per-dispatch writes; async_ctx's slab pointers + + // capacity are cold (prefilled once at init / rebuilt per a5 dispatch) but + // ride this hot line for free. + // Sized to exactly 64B so both dispatch paths write one control line: the + // ready path (src_payload = 0) then also fills args[0..num_args); the gated + // path (src_payload = &PTO2TaskPayload) leaves args[] to the idle AICore. + uint64_t function_bin_addr; /**< Kernel entry address in GM (set by Scheduler). */ - /** Per-dispatch context: block_idx and block_num. - * Written by build_payload() before each dispatch. - * args[SPMD_LOCAL_CONTEXT_INDEX] points here. */ + /** Per-dispatch context: s_block_idx/s_block_num (hot) + async_ctx (task_token hot, + * slab pointers + capacity). args[SPMD_LOCAL_CONTEXT_INDEX] points here. */ LocalContext local_context; - /** Per-core global context: sub_block_id (AIV lane identity). - * Initialized once by init_global_context() at runtime startup. - * args[SPMD_GLOBAL_CONTEXT_INDEX] points here. */ - GlobalContext global_context; + /** Early-dispatch gate AND source pointer, folded into one field. 0 = ready: + * AICore executes on pickup (args[] already filled by the AICPU). Non-zero = + * gated: the value is the source PTO2TaskPayload address; the AICore fills + * args[0..num_args) from it, then waits for the doorbell (DATA_MAIN_BASE + * high 32 == this dispatch's reg_task_id) before executing. A payload address + * is never 0, so the gate flag is lossless. */ + volatile uint64_t src_payload; - uint8_t reserved_payload_abi_pad[8]; + // === Cache lines 1..7: kernel argument vector === + /** [0..num_args) = GM tensor pointers + scalar values; [SPMD_LOCAL_CONTEXT_INDEX] + * = &local_context, [SPMD_GLOBAL_CONTEXT_INDEX] = &global_context (both written + * each dispatch on a5). On the gated path the AICPU leaves [0..num_args) + * unwritten and the idle AICore fills them from src_payload during its gate wait. */ + uint64_t args[PTO2_DISPATCH_MAX_ARGS]; + + /** Per-core global context: sub_block_id (AIV lane identity). Cold: written once + * at init, never per dispatch, so it lives in the tail (not the CL0 control + * block). args[SPMD_GLOBAL_CONTEXT_INDEX] points here. */ + GlobalContext global_context; + // No explicit tail padding: alignas(64) rounds sizeof up to 512 (8 cache lines). static_assert(sizeof(args[0]) == 8); static_assert( @@ -91,3 +125,5 @@ struct alignas(64) PTO2DispatchPayload { }; static_assert(sizeof(PTO2DispatchPayload) == 512, "PTO2DispatchPayload hardware ABI size drift"); +static_assert(offsetof(PTO2DispatchPayload, args) == 64, "args[] must start at cache line 1 (control block = CL0)"); +static_assert(offsetof(PTO2DispatchPayload, src_payload) < 64, "src_payload (gate) must live on the CL0 control block"); diff --git a/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.cpp b/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.cpp index 9cfcb49234..625289819a 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.cpp +++ b/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_orchestrator.cpp @@ -359,6 +359,13 @@ static bool all_claimed_fanin_completed(const PTO2FaninBuilder &fanin_builder) { }); } +static bool all_claimed_fanin_allow_early_resolve(const PTO2FaninBuilder &fanin_builder) { + if (fanin_builder.count == 0) return true; + return fanin_builder.for_each([](PTO2TaskSlotState *producer) -> bool { + return producer != nullptr && producer->allow_early_resolve; + }); +} + void PTO2OrchestratorState::mark_dep_pool_position(PTO2TaskSlotState &slot_state) { PTO2SchedulerState *sched = scheduler; auto &rss = sched->ring_sched_states[slot_state.ring_id]; @@ -377,22 +384,45 @@ void PTO2OrchestratorState::wire_fanin_task(PTO2TaskSlotState &slot_state, int32 slot_state.fanin_count = wfanin + 1; int32_t completed_fanin = 0; + int32_t early_propagated = 0; + // Direct-only (#1285): one unflagged producer => consumer never early-dispatches. + bool early_disqualified = false; for_each_fanin_slot_state(*payload, [&](PTO2TaskSlotState *producer) { producer->lock_fanout(); int32_t pstate = producer->task_state.load(std::memory_order_acquire); + if (!early_disqualified && !producer->allow_early_resolve) { + early_disqualified = true; + } if (pstate >= PTO2_TASK_COMPLETED) { completed_fanin++; } else { producer->fanout_head = rss.dep_pool.prepend(producer->fanout_head, &slot_state); + // Marker shares fanout_lock with propagation's snapshot. A set marker + // means this edge is outside that snapshot and needs a seed. + if (!early_disqualified && producer->has_dispatch_propagated()) { + early_propagated++; + } } producer->unlock_fanout(); }); + // Seed only when every direct producer is codegen-flagged; one unflagged + // producer disqualifies the task (no auto-chain inheritance). + int32_t early_seed = completed_fanin + early_propagated; + if (!early_disqualified && early_seed != 0) { + int32_t dispatch_fanin = payload->dispatch_fanin.fetch_add(early_seed, std::memory_order_acq_rel) + early_seed; + // Fully pre-completed fanin routes normally. If any producer was live, + // the exact-full increment must enqueue the early candidate. + if (completed_fanin != payload->fanin_actual_count && dispatch_fanin == payload->fanin_actual_count) { + sched->try_enqueue_early_dispatch_candidate(slot_state); + } + } + int32_t init_rc = completed_fanin + 1; int32_t new_rc = slot_state.fanin_refcount.fetch_add(init_rc, std::memory_order_acq_rel) + init_rc; mark_dep_pool_position(slot_state); if (new_rc >= slot_state.fanin_count) { - sched->push_ready_routed(&slot_state); + sched->route_ready_once(slot_state); } } @@ -816,7 +846,7 @@ static TaskOutputTensors submit_task_common( } const int32_t kernel_ids_capture[3] = {aic_kernel_id, aiv0_kernel_id, aiv1_kernel_id}; dep_gen_aicpu_record_submit( - task_id.raw, orch->in_manual_scope(), /*early_dispatch=*/false, tc, tensor_ptrs, arg_types_u8, + task_id.raw, orch->in_manual_scope(), args.allow_early_resolve(), tc, tensor_ptrs, arg_types_u8, static_cast(args.explicit_dep_count()), reinterpret_cast(args.explicit_deps_data()), args.launch_spec.core_num(), kernel_ids_capture ); @@ -936,6 +966,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()); // Dispatch predicate: resolve the (tensor, indices) to an absolute GM address // now so the scheduler can read it at the dispatch point with a single load, @@ -986,6 +1017,9 @@ static TaskOutputTensors submit_task_common( } else if (all_claimed_fanin_completed(fanin_builder)) { int32_t ready_seed = fanin_builder.count + 1; cur_slot_state.fanin_count = ready_seed; + if (all_claimed_fanin_allow_early_resolve(fanin_builder)) { + payload.dispatch_fanin.store(fanin_builder.count, std::memory_order_release); + } cur_slot_state.fanin_refcount.store(ready_seed, std::memory_order_release); orch->mark_dep_pool_position(cur_slot_state); sched->push_ready_routed(&cur_slot_state); @@ -1185,7 +1219,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. (a2a3 #1285/#1292) + prepared.slot_state->allow_early_resolve = true; + prepared.slot_state->mark_completed(); } orch->inline_completed_tasks++; diff --git a/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2_types.h b/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2_types.h index 920b39566d..25d433c841 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2_types.h +++ b/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_runtime2_types.h @@ -85,7 +85,8 @@ #define PTO2_SCOPE_TASKS_CAP (PTO2_TASK_WINDOW_SIZE * PTO2_MAX_RING_DEPTH) // Ready queue -#define PTO2_READY_QUEUE_SIZE 65536 // Per-shape queue size +#define PTO2_READY_QUEUE_SIZE 65536 // Per-shape queue size +#define PTO2_EARLY_DISPATCH_QUEUE_SIZE 64 // Per-shape early-dispatch candidate queue // Fanin storage #define PTO2_FANIN_INLINE_CAP 64 @@ -222,6 +223,37 @@ static_assert(offsetof(PTO2TaskDescriptor, packed_buffer_base) == 24, "packed_bu // Per-Slot Scheduling State // ============================================================================= +// 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, // reserved + PTO2_EARLY_DISPATCH_DISPATCHED = 3 // producers released; staged blocks may still be gated +}; + +enum PTO2EarlyDispatchLaunchState : uint8_t { + PTO2_EARLY_DISPATCH_LAUNCH_NONE = 0, + PTO2_EARLY_DISPATCH_LAUNCH_RINGING = 1, + PTO2_EARLY_DISPATCH_LAUNCH_COMPLETE = 2, +}; + +enum PTO2EarlySyncDrainState : uint8_t { + PTO2_EARLY_SYNC_DRAIN_NONE = 0, + PTO2_EARLY_SYNC_DRAIN_OWNER = 1 << 0, + PTO2_EARLY_SYNC_DRAIN_ARMED = 1 << 1, + PTO2_EARLY_SYNC_DRAIN_READY = 1 << 2, + PTO2_EARLY_SYNC_DRAIN_COMPLETE = 1 << 3, +}; + +// A pre-staged consumer occupies one core per gated subtask block. WHICH cores +// it occupies is recorded as a bitmask (staged_core_mask, 1 bit per global +// core_id); the completion-path release iterates the set bits and rings each +// core's doorbell from the scheduler's per-core doorbell table. Bounded by the +// chip's core count (RUNTIME_MAX_WORKER = 108 on a5; 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 >= 108. +inline constexpr int PTO2_EARLY_DISPATCH_CORE_MASK_WORDS = 2; + /** * Task payload data (cold path - only accessed during orchestration and dispatch) * @@ -237,6 +269,15 @@ 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]; + // Early-dispatch metadata (AICPU-side only). Fits in the 40B between the fanin + // array (offset 536) and the 64B-aligned predicate (offset 576). + std::atomic staged_core_mask[PTO2_EARLY_DISPATCH_CORE_MASK_WORDS]{}; + std::atomic dispatch_fanin{0}; // CONSUMER side: fully-published + pre-completed producers + std::atomic published_block_count{0}; + std::atomic early_dispatch_state{0}; + std::atomic early_dispatch_launch_state{PTO2_EARLY_DISPATCH_LAUNCH_NONE}; + std::atomic running_slot_count{0}; + std::atomic early_sync_drain_state{PTO2_EARLY_SYNC_DRAIN_NONE}; // === Cache line 9 (byte 576) — dispatch predicate (AICPU-only) === // Offset is a fixed 576, independent of MAX_TENSOR_ARGS / MAX_SCALAR_ARGS. // AICore never reads it — args are materialized from the tensor_count / tensors @@ -252,6 +293,20 @@ struct PTO2TaskPayload { static_assert(sizeof(Tensor) == 128, "Tensor must be 2 cache lines"); static_assert(MAX_SCALAR_ARGS * sizeof(uint64_t) == 128, "scalar region must be 128B (2 cache lines)"); + void prefetch(int32_t tensor_count, int32_t scalar_count) const { + for (int32_t i = 0; i < tensor_count; i++) { + __builtin_prefetch(&tensors[i], 1, 3); + __builtin_prefetch(reinterpret_cast(&tensors[i]) + 64, 1, 3); + } + for (int32_t i = 0; i < scalar_count; i += 8) { + __builtin_prefetch(&scalars[i], 1, 3); + } + __builtin_prefetch(this, 1, 3); + __builtin_prefetch(reinterpret_cast(this) + 64, 1, 3); + __builtin_prefetch(reinterpret_cast(this) + 128, 1, 3); + __builtin_prefetch(reinterpret_cast(this) + 512, 1, 3); // early-dispatch fields (cache line 8) + } + /** * Initialize payload: copy tensors, store scalars. * @@ -285,6 +340,16 @@ struct PTO2TaskPayload { // Round up to cache line boundary. Both arrays are 128B so no overrun. // Eliminates branches; extra bytes within the same CL have zero additional cost. memcpy(scalars, args.scalars(), PTO2_ALIGN_UP(args.scalar_count() * sizeof(uint64_t), 64)); + + // Early-dispatch metadata — reset on every submit (reset_for_reuse skips payload). + 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); + published_block_count.store(0, std::memory_order_relaxed); + early_dispatch_launch_state.store(PTO2_EARLY_DISPATCH_LAUNCH_NONE, std::memory_order_relaxed); + running_slot_count.store(0, std::memory_order_relaxed); + early_sync_drain_state.store(PTO2_EARLY_SYNC_DRAIN_NONE, std::memory_order_relaxed); } }; @@ -338,6 +403,20 @@ static_assert( // never reach -> provable deadlock. static constexpr uint32_t PTO2_FANOUT_SCOPE_BIT = 0x80000000u; +enum PTO2TaskLifecycleFlag : uint8_t { + PTO2_LIFECYCLE_FLAGS_NONE = 0, + PTO2_READY_CLAIMED = 1U << 0, + PTO2_COMPLETION_DONE = 1U << 1, + // a5 deferred-completion discriminator (formerly std::atomic + // any_subtask_deferred). Packed into lifecycle_flags so sizeof stays 64 + // while READY_CLAIMED / DISPATCH_PROPAGATED also fit. Call sites use + // mark_any_subtask_deferred() / has_any_subtask_deferred(). + PTO2_SUBTASK_DEFERRED = 1U << 2, + PTO2_DISPATCH_PROPAGATED = 1U << 3, +}; + +static_assert((PTO2_DISPATCH_PROPAGATED & (PTO2_READY_CLAIMED | PTO2_COMPLETION_DONE | PTO2_SUBTASK_DEFERRED)) == 0); + struct alignas(64) PTO2TaskSlotState { // Fanout lock + list (accessed together under lock in on_task_complete) std::atomic fanout_lock; // Per-task spinlock (0=unlocked, 1=locked) @@ -366,19 +445,36 @@ struct alignas(64) PTO2TaskSlotState { // --- Set per-submit (depend on task inputs) --- ActiveMask active_mask; // Bitmask of active subtask slots (set once) uint8_t ring_id; // Ring layer (immutable after init) - // Set by any subtask FIN that pushed deferred-completion CONDITIONs to - // the runtime mailbox; read by the last subtask FIN to decide MPSC vs - // inline completion. Mirrors a2a3; see that mirror for the full - // memory-order argument. Carved out of the padding byte between ring_id - // and dep_pool_mark to keep PTO2TaskSlotState at 64 bytes. - std::atomic any_subtask_deferred{false}; - uint8_t _async_pad{0}; + // These one-byte flags live in the padding before dep_pool_mark to keep + // PTO2TaskSlotState at 64 bytes. + bool allow_early_resolve{false}; + std::atomic lifecycle_flags{PTO2_LIFECYCLE_FLAGS_NONE}; int32_t dep_pool_mark{0}; // Dep pool top after Orch-side wiring 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) - int16_t next_block_idx{0}; // Next block to dispatch (scheduler state) + // Next block to dispatch. Normal dispatch and late early-dispatch stagers + // can run concurrently after a partial staged release. All paths claim + // ranges through claim_block_range(). + std::atomic next_block_idx{0}; + + int32_t claim_block_range(int32_t block_limit, int32_t max_count, int32_t &start) { + int16_t current = next_block_idx.load(std::memory_order_relaxed); + while (current < block_limit && max_count > 0) { + int32_t count = block_limit - current; + if (count > max_count) count = max_count; + int16_t desired = static_cast(current + count); + if (next_block_idx.compare_exchange_weak( + current, desired, std::memory_order_seq_cst, std::memory_order_relaxed + )) { + start = current; + return count; + } + } + start = current; + return 0; + } /** * Bind the slot-invariant ring id. Called once per slot during @@ -398,6 +494,35 @@ struct alignas(64) PTO2TaskSlotState { task = t; } + bool has_dispatch_propagated() const { + return (lifecycle_flags.load(std::memory_order_acquire) & PTO2_DISPATCH_PROPAGATED) != 0; + } + + bool try_mark_dispatch_propagated() { + return (lifecycle_flags.fetch_or(PTO2_DISPATCH_PROPAGATED, std::memory_order_acq_rel) & + PTO2_DISPATCH_PROPAGATED) == 0; + } + + void mark_completed() { + task_state.store(PTO2_TASK_COMPLETED, std::memory_order_release); + lifecycle_flags.fetch_or(PTO2_COMPLETION_DONE, std::memory_order_release); + } + + bool is_completion_flag_set() const { + return (lifecycle_flags.load(std::memory_order_acquire) & PTO2_COMPLETION_DONE) != 0; + } + + // Set by any subtask FIN that pushed deferred-completion CONDITIONs to the + // runtime mailbox; read by the last subtask FIN to decide whether the task + // needs MPSC-deferred completion or can complete inline on this thread. + void mark_any_subtask_deferred() { lifecycle_flags.fetch_or(PTO2_SUBTASK_DEFERRED, std::memory_order_release); } + + bool has_any_subtask_deferred() const { + return (lifecycle_flags.load(std::memory_order_acquire) & PTO2_SUBTASK_DEFERRED) != 0; + } + + void set_allow_early_resolve(bool v) { allow_early_resolve = v; } + /** * Reset dynamic scheduling fields for slot reuse. * Called by advance_ring_pointers() after a slot transitions to CONSUMED @@ -416,8 +541,9 @@ struct alignas(64) PTO2TaskSlotState { fanin_refcount.store(0, std::memory_order_relaxed); fanout_refcount.store(0, std::memory_order_relaxed); completed_subtasks.store(0, std::memory_order_relaxed); - next_block_idx = 0; - any_subtask_deferred.store(false, std::memory_order_relaxed); + next_block_idx.store(0, std::memory_order_relaxed); + lifecycle_flags.store(PTO2_LIFECYCLE_FLAGS_NONE, std::memory_order_relaxed); + allow_early_resolve = false; } // === Per-task fanout spinlock === diff --git a/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_types.h b/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_types.h index bd0cd66be1..27829ffb68 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_types.h +++ b/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_types.h @@ -242,6 +242,15 @@ struct Arg : TaskArgsTpl { const char *error_msg{nullptr}; PTO2LaunchSpec launch_spec; // SPMD launch parameters (block_num, etc.) + // Early-dispatch hint (codegen-author set, off by default). When + // true, the scheduler may stage this task on an idle core before its producer + // finishes, gating execution on the DATA_MAIN_BASE doorbell — only safe when + // the author knows the task's data dependencies allow it. Read in-process by + // the runtime; never crosses the wire format. + bool allow_early_resolve_{false}; + void set_allow_early_resolve(bool v = true) { allow_early_resolve_ = v; } + bool allow_early_resolve() const { return allow_early_resolve_; } + // Dispatch predicate (codegen-author set; default op == NONE = always // dispatch). A FALSE result at the dispatch point retires the task inline // through the dep-only path — never dispatched to an AICore — while still @@ -273,6 +282,7 @@ struct Arg : TaskArgsTpl { #endif explicit_deps_ = nullptr; explicit_dep_count_ = 0; + allow_early_resolve_ = false; predicate_ = L0TaskPredicate{}; task_timing_slot_ = TASK_TIMING_SLOT_NONE; } diff --git a/src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.h b/src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.h index 4b69e5c6fc..61e59215dd 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.h +++ b/src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/pto_scheduler.h @@ -37,6 +37,8 @@ #include "pto_ring_buffer.h" #include "pto_runtime2_types.h" #include "pto_shared_memory.h" +#include "aicpu/platform_regs.h" +#include "common/memory_barrier.h" #if SIMPLER_SCHED_PROFILING #include "aicpu/device_time.h" @@ -59,6 +61,7 @@ struct PTO2ReadyQueueSlot { std::atomic sequence; PTO2TaskSlotState *slot_state; + uint64_t task_id_snapshot; // generation tag for early-dispatch queue entries }; /** @@ -91,7 +94,9 @@ struct alignas(64) PTO2ReadyQueue { void reset_for_reuse() {} - bool push(PTO2TaskSlotState *slot_state) { + bool push(PTO2TaskSlotState *slot_state) { return push_tagged(slot_state, 0); } + + bool push_tagged(PTO2TaskSlotState *slot_state, uint64_t task_id_snapshot) { uint64_t pos; PTO2ReadyQueueSlot *slot; while (true) { @@ -111,13 +116,16 @@ struct alignas(64) PTO2ReadyQueue { } slot->slot_state = slot_state; + slot->task_id_snapshot = task_id_snapshot; slot->sequence.store(static_cast(pos + 1), std::memory_order_release); return true; } // Batch push: reserve count slots with a single CAS after confirming // every target slot is available under the usual Vyukov sequence check. - void push_batch(PTO2TaskSlotState **items, int count) { + void push_batch(PTO2TaskSlotState **items, int count) { push_batch_tagged(items, nullptr, count); } + + void push_batch_tagged(PTO2TaskSlotState **items, const uint64_t *task_id_snapshots, int count) { if (count == 0) return; uint64_t pos; @@ -146,6 +154,7 @@ struct alignas(64) PTO2ReadyQueue { for (int i = 0; i < count; i++) { PTO2ReadyQueueSlot *slot = &slots[(pos + i) & mask]; slot->slot_state = items[i]; + slot->task_id_snapshot = task_id_snapshots == nullptr ? 0 : task_id_snapshots[i]; slot->sequence.store(static_cast(pos + i + 1), std::memory_order_release); } } @@ -185,12 +194,15 @@ struct alignas(64) PTO2ReadyQueue { } slot->slot_state = slot_state; + slot->task_id_snapshot = 0; slot->sequence.store(static_cast(pos + 1), std::memory_order_release); return true; } #endif - PTO2TaskSlotState *pop() { + PTO2TaskSlotState *pop() { return pop_tagged(nullptr); } + + PTO2TaskSlotState *pop_tagged(uint64_t *task_id_snapshot) { // Fast-path: skip slot load when queue is clearly empty uint64_t d = dequeue_pos.load(std::memory_order_relaxed); uint64_t e = enqueue_pos.load(std::memory_order_relaxed); @@ -216,6 +228,7 @@ struct alignas(64) PTO2ReadyQueue { } PTO2TaskSlotState *result = slot->slot_state; + if (task_id_snapshot != nullptr) *task_id_snapshot = slot->task_id_snapshot; slot->sequence.store(static_cast(pos + mask + 1), std::memory_order_release); return result; } @@ -271,7 +284,9 @@ struct alignas(64) PTO2ReadyQueue { // Batch pop: reserve a contiguous run of ready slots with a single CAS. // Returns actual number of items popped (may be less than max_count). - int pop_batch(PTO2TaskSlotState **out, int max_count) { + int pop_batch(PTO2TaskSlotState **out, int max_count) { return pop_batch_tagged(out, nullptr, max_count); } + + int pop_batch_tagged(PTO2TaskSlotState **out, uint64_t *task_id_snapshots, int max_count) { uint64_t pos; int count; while (true) { @@ -303,6 +318,7 @@ struct alignas(64) PTO2ReadyQueue { for (int i = 0; i < count; i++) { PTO2ReadyQueueSlot *slot = &slots[(pos + i) & mask]; out[i] = slot->slot_state; + if (task_id_snapshots != nullptr) task_id_snapshots[i] = slot->task_id_snapshot; slot->sequence.store(static_cast(pos + i + mask + 1), std::memory_order_release); } return count; @@ -403,6 +419,8 @@ struct CompletionStats { struct PTO2SchedulerLayout { size_t off_ready_queue_slots[PTO2_NUM_RESOURCE_SHAPES]; size_t off_dummy_ready_queue_slots; + size_t off_early_dispatch_queue_slots[PTO2_NUM_RESOURCE_SHAPES]; + size_t off_early_sync_start_queue_slots; size_t off_dep_pool_entries[PTO2_MAX_RING_DEPTH]; uint64_t ready_queue_capacity; int32_t dep_pool_capacities[PTO2_MAX_RING_DEPTH]; @@ -623,36 +641,334 @@ struct PTO2SchedulerState { } #endif - bool release_fanin_and_check_ready(PTO2TaskSlotState &slot_state) { + // ========================================================================= + // Early-dispatch (#1079 foundation + #1304 sync_start early queue) + // ========================================================================= + + struct EarlyDispatchDoorbell { + uint64_t addr{0}; + uint32_t token{0}; + }; + EarlyDispatchDoorbell early_dispatch_doorbell_table[PTO2_EARLY_DISPATCH_CORE_MASK_WORDS * 64]{}; + PTO2ReadyQueue early_dispatch_queues[PTO2_NUM_RESOURCE_SHAPES]; + PTO2ReadyQueue early_sync_start_queue; + + static inline void ring_one_doorbell(uint64_t reg_addr, uint32_t token) { + volatile uint64_t *dmb = reinterpret_cast(get_reg_ptr(reg_addr, RegId::DATA_MAIN_BASE)); + uint64_t tk = static_cast(token); + *dmb = (tk << 32) | tk; // 64-bit STR: high=low=token releases the gated AICore + } + + inline void ring_staged_doorbell_bits(int word, uint64_t bits) { + while (bits != 0) { + int core_id = word * 64 + __builtin_ctzll(bits); + bits &= bits - 1; + ring_one_doorbell( + early_dispatch_doorbell_table[core_id].addr, early_dispatch_doorbell_table[core_id].token + ); + } + } + + static inline uint64_t claim_all_staged_doorbell_bits(std::atomic &mask) { + return mask.exchange(0, std::memory_order_seq_cst); + } + + static inline uint64_t claim_late_staged_doorbell_bits(std::atomic &mask, uint64_t candidates) { + return mask.fetch_and(~candidates, std::memory_order_seq_cst) & candidates; + } + + static inline bool should_gate_early_dispatch(bool force_gate, uint8_t early_dispatch_state) { + return force_gate || early_dispatch_state == PTO2_EARLY_DISPATCH_STAGING; + } + + static inline bool + ring_claimed_local_doorbell(uint64_t claimed_word, int core_id, uint64_t reg_addr, uint32_t token) { + if ((claimed_word & (1ULL << (core_id & 63))) == 0) return false; + ring_one_doorbell(reg_addr, token); + return true; + } + + static inline bool try_claim_early_dispatch_launch(PTO2TaskPayload &payload) { + uint8_t expected = PTO2_EARLY_DISPATCH_LAUNCH_NONE; + return payload.early_dispatch_launch_state.compare_exchange_strong( + expected, PTO2_EARLY_DISPATCH_LAUNCH_RINGING, std::memory_order_seq_cst, std::memory_order_seq_cst + ); + } + + inline void record_published_blocks(PTO2TaskSlotState &slot_state, int32_t count) { + if (count <= 0 || !slot_state.allow_early_resolve) return; + slot_state.payload->published_block_count.fetch_add(static_cast(count), std::memory_order_seq_cst); + } + + inline void ring_all_staged_doorbells(PTO2TaskSlotState &slot_state) { + for (int w = 0; w < PTO2_EARLY_DISPATCH_CORE_MASK_WORDS; w++) { + uint64_t bits = slot_state.payload->staged_core_mask[w].load(std::memory_order_seq_cst); + while (bits != 0) { + int core_id = w * 64 + __builtin_ctzll(bits); + bits &= bits - 1; + ring_one_doorbell( + early_dispatch_doorbell_table[core_id].addr, early_dispatch_doorbell_table[core_id].token + ); + } + } + } + + static inline bool try_claim_early_sync_drain(PTO2TaskPayload &payload) { + uint8_t expected = PTO2_EARLY_SYNC_DRAIN_NONE; + return payload.early_sync_drain_state.compare_exchange_strong( + expected, PTO2_EARLY_SYNC_DRAIN_OWNER, std::memory_order_seq_cst, std::memory_order_seq_cst + ); + } + + static inline bool owns_early_sync_drain(const PTO2TaskPayload &payload) { + return (payload.early_sync_drain_state.load(std::memory_order_acquire) & PTO2_EARLY_SYNC_DRAIN_OWNER) != 0; + } + + static inline void mark_early_sync_drain_armed(PTO2TaskPayload &payload) { + payload.early_sync_drain_state.fetch_or(PTO2_EARLY_SYNC_DRAIN_ARMED, std::memory_order_seq_cst); + } + + static inline bool publish_ready_to_early_sync_drain(PTO2TaskPayload &payload) { + uint8_t previous = + payload.early_sync_drain_state.fetch_or(PTO2_EARLY_SYNC_DRAIN_READY, std::memory_order_seq_cst); + return (previous & PTO2_EARLY_SYNC_DRAIN_OWNER) != 0; + } + + inline void cancel_early_sync_drain(PTO2TaskSlotState &slot_state) { + uint8_t previous = + slot_state.payload->early_sync_drain_state.exchange(PTO2_EARLY_SYNC_DRAIN_NONE, std::memory_order_seq_cst); + if ((previous & PTO2_EARLY_SYNC_DRAIN_OWNER) == 0) return; + if ((previous & PTO2_EARLY_SYNC_DRAIN_READY) != 0) { + push_ready_routed(&slot_state); + return; + } + if (slot_state.payload->early_dispatch_state.load(std::memory_order_seq_cst) == PTO2_EARLY_DISPATCH_STAGING) { + early_sync_start_queue.push_tagged(&slot_state, static_cast(slot_state.task->task_id.raw)); + } + } + + static inline void finish_early_sync_drain(PTO2TaskPayload &payload) { + uint8_t state = payload.early_sync_drain_state.load(std::memory_order_seq_cst); + while ((state & PTO2_EARLY_SYNC_DRAIN_OWNER) != 0 && (state & PTO2_EARLY_SYNC_DRAIN_COMPLETE) == 0) { + uint8_t desired = state | PTO2_EARLY_SYNC_DRAIN_COMPLETE; + if (payload.early_sync_drain_state.compare_exchange_weak( + state, desired, std::memory_order_seq_cst, std::memory_order_seq_cst + )) { + return; + } + } + } + + inline bool maybe_rendezvous_ring(PTO2TaskSlotState &slot_state) { + int32_t staged_cores = 0; + for (int w = 0; w < PTO2_EARLY_DISPATCH_CORE_MASK_WORDS; w++) + staged_cores += + __builtin_popcountll(slot_state.payload->staged_core_mask[w].load(std::memory_order_seq_cst)); + if (staged_cores == 0) return false; + if (slot_state.payload->running_slot_count.load(std::memory_order_seq_cst) != staged_cores) return false; + if (slot_state.payload->early_dispatch_state.load(std::memory_order_seq_cst) != PTO2_EARLY_DISPATCH_DISPATCHED) + return false; + if (!try_claim_early_dispatch_launch(*slot_state.payload)) return false; + ring_all_staged_doorbells(slot_state); + wmb(); + slot_state.payload->early_dispatch_launch_state.store( + PTO2_EARLY_DISPATCH_LAUNCH_COMPLETE, std::memory_order_release + ); + return true; + } + + inline bool retry_sync_start_rendezvous_after_drain(PTO2TaskSlotState &slot_state) { + if (!maybe_rendezvous_ring(slot_state)) return false; + propagate_dispatch_fanin(slot_state); + return true; + } + + inline void try_enqueue_early_dispatch_candidate(PTO2TaskSlotState &consumer) { + if (consumer.active_mask.has_predicate()) return; + PTO2ResourceShape shape = consumer.active_mask.to_shape(); + if (shape != PTO2ResourceShape::AIC && shape != PTO2ResourceShape::AIV && shape != PTO2ResourceShape::MIX) { + return; + } + + uint8_t expected = PTO2_EARLY_DISPATCH_NONE; + if (!consumer.payload->early_dispatch_state.compare_exchange_strong( + expected, PTO2_EARLY_DISPATCH_STAGING, std::memory_order_seq_cst, std::memory_order_seq_cst + )) { + return; + } + + uint64_t task_id = static_cast(consumer.task->task_id.raw); + if (consumer.active_mask.requires_sync_start()) { + early_sync_start_queue.push_tagged(&consumer, task_id); + } else { + early_dispatch_queues[static_cast(shape)].push_tagged(&consumer, task_id); + } + } + + // Only codegen-flagged (direct) producers propagate: a task's successors + // early-dispatch off its DIRECT producers' marks, never an inherited chain + // (a2a3 #1285 — a5 never had auto-chain / spec_chain_*). + void propagate_dispatch_fanin(PTO2TaskSlotState &p) { + if (!p.allow_early_resolve) return; + if (p.payload->published_block_count.load(std::memory_order_seq_cst) < p.logical_block_num) return; + if (p.payload->early_dispatch_state.load(std::memory_order_acquire) == PTO2_EARLY_DISPATCH_STAGING) return; + uint8_t launch_state = p.payload->early_dispatch_launch_state.load(std::memory_order_acquire); + if (launch_state == PTO2_EARLY_DISPATCH_LAUNCH_RINGING) return; + if (p.active_mask.requires_sync_start()) { + bool was_pre_staged = false; + for (int w = 0; w < PTO2_EARLY_DISPATCH_CORE_MASK_WORDS; w++) { + was_pre_staged |= p.payload->staged_core_mask[w].load(std::memory_order_acquire) != 0; + } + if (was_pre_staged && launch_state != PTO2_EARLY_DISPATCH_LAUNCH_COMPLETE) return; + } + if (p.has_dispatch_propagated()) return; + p.lock_fanout(); + if (!p.try_mark_dispatch_propagated()) { + p.unlock_fanout(); + return; + } + PTO2DepListEntry *edge = p.fanout_head; + p.unlock_fanout(); + for (; edge != nullptr; edge = edge->next) { + PTO2TaskSlotState *c = edge->slot_state; + if (c->active_mask.has_predicate()) continue; + int32_t nf = c->payload->dispatch_fanin.fetch_add(1, std::memory_order_acq_rel) + 1; + if (nf != c->payload->fanin_actual_count) continue; + try_enqueue_early_dispatch_candidate(*c); + } + } + + struct EarlyDispatchReleaseSink { + static constexpr int CAP = 32; + PTO2TaskSlotState *items[CAP]; + int n = 0; + inline bool push(PTO2TaskSlotState *s) { + if (n >= CAP) return false; + items[n++] = s; + return true; + } + }; + + inline bool try_early_dispatch_release(PTO2TaskSlotState &slot_state, EarlyDispatchReleaseSink *sink = nullptr) { + if (slot_state.active_mask.has_predicate()) return false; + 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; + } + if (expect != PTO2_EARLY_DISPATCH_STAGING) return true; + bool sync_start = slot_state.active_mask.requires_sync_start(); + if (!sync_start && !try_claim_early_dispatch_launch(*slot_state.payload)) return true; + expect = PTO2_EARLY_DISPATCH_STAGING; + slot_state.payload->early_dispatch_state.compare_exchange_strong( + expect, PTO2_EARLY_DISPATCH_DISPATCHED, std::memory_order_seq_cst, std::memory_order_seq_cst + ); + bool launched = true; + if (sync_start) { + launched = maybe_rendezvous_ring(slot_state); + } else { + for (int w = 0; w < PTO2_EARLY_DISPATCH_CORE_MASK_WORDS; w++) { + uint64_t owned = claim_all_staged_doorbell_bits(slot_state.payload->staged_core_mask[w]); + ring_staged_doorbell_bits(w, owned); + } + wmb(); + slot_state.payload->early_dispatch_launch_state.store( + PTO2_EARLY_DISPATCH_LAUNCH_COMPLETE, std::memory_order_seq_cst + ); + } + if (launched) { + if (sink == nullptr || !sink->push(&slot_state)) { + propagate_dispatch_fanin(slot_state); + } + } + return slot_state.next_block_idx.load(std::memory_order_seq_cst) >= slot_state.logical_block_num; + } + + bool try_claim_ready_once(PTO2TaskSlotState &slot_state) { + uint8_t flags = slot_state.lifecycle_flags.load(std::memory_order_acquire); + for (;;) { + if ((flags & PTO2_READY_CLAIMED) != 0) return false; + uint8_t desired_flags = flags | PTO2_READY_CLAIMED; + if (slot_state.lifecycle_flags.compare_exchange_weak( + flags, desired_flags, std::memory_order_acq_rel, std::memory_order_acquire + )) { + return true; + } + } + } + + bool route_ready_once(PTO2TaskSlotState &slot_state, EarlyDispatchReleaseSink *sink = nullptr) { + if (!try_claim_ready_once(slot_state)) return false; + bool early_handled = try_early_dispatch_release(slot_state, sink); + if (slot_state.active_mask.requires_sync_start()) { + bool drain_owned = publish_ready_to_early_sync_drain(*slot_state.payload); + if (early_handled || drain_owned) return true; + } else if (early_handled) { + return true; + } + push_ready_routed(&slot_state); + return true; + } + +#if SIMPLER_ORCH_PROFILING || SIMPLER_SCHED_PROFILING + bool route_ready_once( + PTO2TaskSlotState &slot_state, uint64_t &atomic_count, uint64_t &push_wait, + EarlyDispatchReleaseSink *sink = nullptr + ) { + uint8_t flags = slot_state.lifecycle_flags.load(std::memory_order_acquire); + atomic_count += 1; + for (;;) { + if ((flags & PTO2_READY_CLAIMED) != 0) return false; + uint8_t desired_flags = flags | PTO2_READY_CLAIMED; + if (slot_state.lifecycle_flags.compare_exchange_weak( + flags, desired_flags, std::memory_order_acq_rel, std::memory_order_acquire + )) { + atomic_count += 1; + break; + } + atomic_count += 1; + } + bool early_handled = try_early_dispatch_release(slot_state, sink); + if (slot_state.active_mask.requires_sync_start()) { + bool drain_owned = publish_ready_to_early_sync_drain(*slot_state.payload); + if (early_handled || drain_owned) return true; + } else if (early_handled) { + return true; + } + PTO2ResourceShape shape = slot_state.active_mask.to_shape(); + if (shape == PTO2ResourceShape::DUMMY || + (slot_state.active_mask.has_predicate() && !slot_state.payload->predicate.pass())) { + dummy_ready_queue.push(&slot_state, atomic_count, push_wait); + } else { + 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) { - push_ready_routed(&slot_state); - return true; + return route_ready_once(slot_state, sink); } return false; } #if SIMPLER_ORCH_PROFILING || SIMPLER_SCHED_PROFILING - bool release_fanin_and_check_ready(PTO2TaskSlotState &slot_state, uint64_t &atomic_count, uint64_t &push_wait) { + bool release_fanin_and_check_ready( + PTO2TaskSlotState &slot_state, uint64_t &atomic_count, uint64_t &push_wait, + EarlyDispatchReleaseSink *sink = nullptr + ) { int32_t new_refcount = slot_state.fanin_refcount.fetch_add(1, std::memory_order_acq_rel) + 1; atomic_count += 1; // fanin_refcount.fetch_add if (new_refcount == slot_state.fanin_count) { - // Dummy slots go to dummy_ready_queue; everything else to the per-shape - // ready_queues[]. 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 || - (slot_state.active_mask.has_predicate() && !slot_state.payload->predicate.pass())) { - dummy_ready_queue.push(&slot_state, atomic_count, push_wait); - } else { - 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; } @@ -734,7 +1050,7 @@ 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(); @@ -749,18 +1065,22 @@ struct PTO2SchedulerState { #if SIMPLER_SCHED_PROFILING uint64_t fanout_atomics = 0, push_wait = 0; #endif + 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)) { + if (release_fanin_and_check_ready(consumer_slot, fanout_atomics, push_wait, &rel_sink)) { stats.tasks_enqueued++; } #else - release_fanin_and_check_ready(consumer_slot); + release_fanin_and_check_ready(consumer_slot, &rel_sink); #endif current = current->next; } + for (int i = 0; i < rel_sink.n; i++) { + propagate_dispatch_fanin(*rel_sink.items[i]); + } #if SIMPLER_SCHED_PROFILING g_sched_fanout_atomic_count[thread_idx] += fanout_atomics; diff --git a/src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_cold_path.cpp b/src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_cold_path.cpp index 567bb5dab7..3ac164a774 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_cold_path.cpp +++ b/src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_cold_path.cpp @@ -1259,6 +1259,9 @@ void SchedulerContext::deinit() { drain_state_.drain_worker_elected.store(0, std::memory_order_release); drain_state_.drain_ack_mask.store(0, std::memory_order_release); drain_state_.pending_task.store(nullptr, std::memory_order_release); + drain_state_.drain_stage_go.store(0, std::memory_order_release); + drain_state_.drain_stage_done_mask.store(0, std::memory_order_release); + drain_state_.drain_running_staged.store(0, std::memory_order_release); // Reset task counters and orchestrator state completed_tasks_.store(0, std::memory_order_release); diff --git a/src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_completion.cpp b/src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_completion.cpp index e15862272f..7a65524aad 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_completion.cpp +++ b/src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_completion.cpp @@ -10,6 +10,8 @@ */ #include "scheduler_context.h" +#include + #include "common/unified_log.h" #include "aicpu/device_time.h" #include "aicpu/device_phase_aicpu.h" @@ -36,7 +38,7 @@ inline constexpr int32_t PTO2_DEFERRED_RELEASE_CAP = 256; // Pure function: read register result -> SlotTransition (no side effects). SlotTransition SchedulerContext::decide_slot_transition( - int32_t reg_task_id, int32_t reg_state, int32_t running_id, int32_t pending_id + int32_t reg_task_id, int32_t reg_state, int32_t running_id, int32_t pending_id, bool pending_gated ) { SlotTransition t; if (pending_id != AICPU_TASK_INVALID && reg_task_id == pending_id) { @@ -55,9 +57,14 @@ SlotTransition SchedulerContext::decide_slot_transition( t.matched = true; t.running_done = true; t.running_freed = true; + } else if (pending_gated) { + // Case 3.3: running FIN, pending is early-dispatch gated. Complete + // running now and promote; waiting for a gated ACK would deadlock. + t.matched = true; + t.running_done = true; + t.running_freed = true; } - // Case 3.1: running FIN, pending exists -> skip (transient state). - // Case 1/2 (pending ACK/FIN) will complete running implicitly via running_done=true. + // Case 3.1: running FIN, NON-gated pending exists -> skip (transient). } else { // Case 4: running ACK -- only pending_freed (slot now hardware-latched) t.matched = true; @@ -112,7 +119,7 @@ void SchedulerContext::complete_slot_task( } if (cond_count > 0) { - 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) { @@ -129,8 +136,7 @@ void SchedulerContext::complete_slot_task( bool task_complete = sched_->on_subtask_complete(slot_state); - 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()) { while (!mailbox->try_push_normal_done(slot_state.task->task_id, reinterpret_cast(&slot_state))) { sched_->async_wait_list.mpsc_skipped_count.fetch_add(1, std::memory_order_relaxed); SPIN_WAIT_HINT(); @@ -253,6 +259,15 @@ void SchedulerContext::check_running_cores_for_completion( int32_t core_id = tracker.get_core_id_by_offset(bit_pos); CoreExecState &core = core_exec_states_[core_id]; + // Skip gated early-dispatch cores still waiting for their doorbell. + { + PTO2TaskSlotState *rs = core.running_slot_state; + if (rs != nullptr && rs->payload != nullptr && + rs->payload->early_dispatch_state.load(std::memory_order_relaxed) == PTO2_EARLY_DISPATCH_STAGING) { + continue; + } + } + // --- Judgment phase: read register, derive transition --- // Use the precomputed cond_ptr (resolved once in handshake) to skip // the reg_offset switch and reg_addr addition on every poll. @@ -275,8 +290,19 @@ void SchedulerContext::check_running_cores_for_completion( } #endif - SlotTransition t = - decide_slot_transition(reg_task_id, reg_state, core.running_reg_task_id, core.pending_reg_task_id); + // Phase 1+: gated == not yet launched (STAGING, or sync_start DISPATCHED rendezvous). + uint8_t pending_ss = + (core.pending_slot_state != nullptr && core.pending_slot_state->payload != nullptr) ? + core.pending_slot_state->payload->early_dispatch_state.load(std::memory_order_relaxed) : + static_cast(PTO2_EARLY_DISPATCH_NONE); + bool pending_gated = + (core.pending_slot_state != nullptr && core.pending_slot_state->payload != nullptr && + (pending_ss == PTO2_EARLY_DISPATCH_STAGING || + (pending_ss == PTO2_EARLY_DISPATCH_DISPATCHED && + core.pending_slot_state->active_mask.requires_sync_start()))); + SlotTransition t = decide_slot_transition( + reg_task_id, reg_state, core.running_reg_task_id, core.pending_reg_task_id, pending_gated + ); if (!t.matched) continue; #if SIMPLER_SCHED_PROFILING @@ -336,7 +362,15 @@ void SchedulerContext::check_running_cores_for_completion( // 2. Update slot data if (t.running_freed) { if (core.pending_slot_state != nullptr && !t.pending_done) { - promote_pending_to_running(core); // Case 2 or Case 3 (with pending) + PTO2TaskSlotState *promoted = core.pending_slot_state; + bool sync_start_promote = pending_gated && promoted->active_mask.requires_sync_start(); + promote_pending_to_running(core); + if (sync_start_promote) { + promoted->payload->running_slot_count.fetch_add(1, std::memory_order_seq_cst); + if (sched_->maybe_rendezvous_ring(*promoted)) { + sched_->propagate_dispatch_fanin(*promoted); + } + } } else { clear_running_slot(core); // Case 1 or Case 3 (no pending) if (t.pending_done) { @@ -390,6 +424,9 @@ bool SchedulerContext::enter_drain_mode(PTO2TaskSlotState *slot_state, int32_t b drain_state_.pending_task.store(slot_state, std::memory_order_release); drain_state_.drain_ack_mask.store(0, std::memory_order_relaxed); drain_state_.drain_worker_elected.store(0, std::memory_order_relaxed); + drain_state_.drain_stage_go.store(0, std::memory_order_relaxed); + drain_state_.drain_stage_done_mask.store(0, std::memory_order_relaxed); + drain_state_.drain_running_staged.store(0, std::memory_order_relaxed); // Release store: all stores above are now visible to any thread that // acquire-loads sync_start_pending and sees block_num > 0. drain_state_.sync_start_pending.store(block_num, std::memory_order_release); @@ -397,74 +434,97 @@ bool SchedulerContext::enter_drain_mode(PTO2TaskSlotState *slot_state, int32_t b } // Count total available resources across all scheduler threads for a given shape. -int32_t SchedulerContext::count_global_available(PTO2ResourceShape shape, uint8_t core_mask) { +int32_t SchedulerContext::count_global_available(PTO2ResourceShape shape, uint8_t core_mask, bool include_pending) { int32_t total = 0; for (int32_t t = 0; t < active_sched_threads_; t++) { if (shape == PTO2ResourceShape::MIX) { - total += core_trackers_[t].count_mix_running_clusters(core_mask); + total += include_pending ? core_trackers_[t].count_mix_split_clusters(core_mask) : + core_trackers_[t].count_mix_running_clusters(core_mask); } else { total += core_trackers_[t].get_idle_core_offset_states(shape).count(); + if (include_pending) { + total += core_trackers_[t].get_pending_core_offset_states(shape).count(); + } } } return total; } -// Drain worker: dispatch all blocks in one pass across all threads' trackers. -// Called only when global resources >= block_num, so one pass always suffices. -// All other threads are spinning -- the drain worker has exclusive tracker access. -void SchedulerContext::drain_worker_dispatch(Runtime *runtime, int32_t block_num) { - PTO2TaskSlotState *slot_state = drain_state_.pending_task.load(std::memory_order_acquire); - if (!slot_state) { - drain_state_.sync_start_pending.store(0, std::memory_order_release); - return; - } +int32_t +SchedulerContext::drain_stage_cores(PTO2TaskSlotState *slot_state, int32_t block_num, int32_t thread_idx, bool gated) { + CoreTracker &tracker = core_trackers_[thread_idx]; PTO2ResourceShape shape = slot_state->active_mask.to_shape(); uint8_t core_mask = slot_state->active_mask.core_mask(); - - for (int32_t t = 0; t < active_sched_threads_ && slot_state->next_block_idx < block_num; t++) { - auto valid = (shape == PTO2ResourceShape::MIX) ? - core_trackers_[t].get_mix_running_cluster_offset_states(core_mask) : - core_trackers_[t].get_idle_core_offset_states(shape); - while (valid.has_value() && slot_state->next_block_idx < block_num) { - dispatch_block(runtime, t, valid.pop_first(), *slot_state, shape, false, slot_state->next_block_idx); - slot_state->next_block_idx++; + bool mix_split = gated && shape == PTO2ResourceShape::MIX; + int32_t running_staged = 0; + + auto stage = [&](CoreTracker::BitStates valid, bool to_pending) { + while (valid.has_value()) { + int32_t avail = valid.count(); + int32_t start = 0; + int32_t claim = slot_state->claim_block_range(block_num, avail, start); + if (claim == 0) return; + PublishHandle handles[CoreTracker::MAX_CLUSTERS * 3]; + int handle_count = 0; + int32_t claimed[CoreTracker::MAX_CLUSTERS * 3]; + for (int32_t b = 0; b < claim; b++) + claimed[b] = valid.pop_first(); + for (int32_t b = 0; b < claim; b++) { + auto core_offset = claimed[b]; + if (shape == PTO2ResourceShape::MIX) { + running_staged += tracker.mix_cluster_idle_core_count(core_offset, core_mask); + } + handle_count += prepare_block_for_dispatch( + thread_idx, core_offset, *slot_state, shape, to_pending, start + b, &handles[handle_count], gated + ); + } + wmb(); + uint64_t dispatch_ts = 0; +#if SIMPLER_DFX + if (l2_swimlane_level_ >= L2SwimlaneLevel::AICPU_TIMING) { + dispatch_ts = get_sys_cnt_aicpu(); + } +#endif + uint64_t my_mask[PTO2_EARLY_DISPATCH_CORE_MASK_WORDS] = {0}; + for (int i = 0; i < handle_count; i++) { + publish_subtask_to_core(handles[i], dispatch_ts, thread_idx); + if (gated) { + int32_t cid = tracker.get_core_id_by_offset(handles[i].core_offset); + sched_->early_dispatch_doorbell_table[cid].addr = handles[i].reg_addr; + sched_->early_dispatch_doorbell_table[cid].token = handles[i].reg_task_id; + my_mask[cid >> 6] |= 1ULL << (cid & 63); + } + } + if (gated) { + for (int w = 0; w < PTO2_EARLY_DISPATCH_CORE_MASK_WORDS; w++) { + if (my_mask[w] != 0) { + slot_state->payload->staged_core_mask[w].fetch_or(my_mask[w], std::memory_order_seq_cst); + } + } + } + sched_->record_published_blocks(*slot_state, claim); + if (gated && shape != PTO2ResourceShape::MIX && !to_pending) running_staged += handle_count; + } + }; + + if (mix_split) { + stage(tracker.get_mix_split_cluster_offset_states(core_mask), /*to_pending=*/true); + } else { + auto idle = (shape == PTO2ResourceShape::MIX) ? tracker.get_mix_running_cluster_offset_states(core_mask) : + tracker.get_idle_core_offset_states(shape); + stage(idle, /*to_pending=*/false); + if (gated) { + stage(tracker.get_pending_core_offset_states(shape), /*to_pending=*/true); } } - - // 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. - std::atomic_thread_fence(std::memory_order_release); - drain_state_.pending_task.store(nullptr, std::memory_order_release); - drain_state_.drain_ack_mask.store(0, std::memory_order_relaxed); - drain_state_.drain_worker_elected.store(0, std::memory_order_relaxed); - drain_state_.sync_start_pending.store(0, std::memory_order_release); + return running_staged; } -// Called by each scheduler thread when drain_state_.sync_start_pending != 0. -// -// Protocol (single-stage ack barrier): -// 1. Ack barrier: all threads signal they've stopped dispatch, then spin -// until all ack bits are set. -// If this thread's bit gets cleared while waiting, a reset occurred -- return. -// 2. Election: one thread wins the CAS and becomes the drain worker. -// If resources are insufficient, reset ack/election fields and return -- -// all threads resume completion polling to free running cores, then retry. -// 3. Dispatch: elected thread dispatches all blocks (one pass, resources guaranteed). -// Non-elected threads spin-wait until sync_start_pending == 0. -// During dispatch the elected thread has exclusive tracker access. -void SchedulerContext::handle_drain_mode(Runtime *runtime, int32_t thread_idx) { - // Every spin in this function honors is_completed(): once the run latches - // completed_ (all tasks done, or a fatal error raised elsewhere), peers leave - // the dispatch loop and stop participating in the drain. A thread parked in a - // drain spin would then wait forever for acks / a gate-open that can no longer - // arrive -- the AICPU watchdog never fires here because these spins live - // outside the dispatch loop's wall-clock budget, so the hang escalates straight - // to the 3 s STARS op-exec timeout (507018) and poisons the device. Bailing on - // completed_ is always safe: any pending sync_start task is either already - // dispatched (a stale re-popped slot) or moot under teardown, and deinit() - // resets drain_state_ before the next run, so leaving it dirty is harmless. - // Spin until drain is fully initialized (sentinel -1 -> block_num > 0). +void SchedulerContext::handle_drain_mode(int32_t thread_idx, [[maybe_unused]] uint64_t *out_stage_wall_cycles) { +#if SIMPLER_DFX + bool drain_prof = (l2_swimlane_level_ >= L2SwimlaneLevel::SCHED_PHASES && out_stage_wall_cycles != nullptr); + uint64_t drain_acked_ts = 0; +#endif int32_t block_num; do { if (is_completed()) return; @@ -474,11 +534,8 @@ void SchedulerContext::handle_drain_mode(Runtime *runtime, int32_t thread_idx) { uint32_t all_acked = (1u << active_sched_threads_) - 1; - // Ack barrier -- signal this thread has stopped dispatch. drain_state_.drain_ack_mask.fetch_or(1u << thread_idx, std::memory_order_release); - // Spin until all threads have acked. - // If our bit is cleared while waiting, elected reset due to insufficient resources. while (true) { if (is_completed()) return; uint32_t ack = drain_state_.drain_ack_mask.load(std::memory_order_acquire); @@ -487,45 +544,84 @@ void SchedulerContext::handle_drain_mode(Runtime *runtime, int32_t thread_idx) { SPIN_WAIT_HINT(); } - // Election -- exactly one thread wins the CAS. int32_t expected = 0; drain_state_.drain_worker_elected.compare_exchange_strong( expected, thread_idx + 1, std::memory_order_acquire, std::memory_order_relaxed ); + bool elected = drain_state_.drain_worker_elected.load(std::memory_order_relaxed) == thread_idx + 1; - if (drain_state_.drain_worker_elected.load(std::memory_order_relaxed) != thread_idx + 1) { - // Non-elected: spin-wait for drain completion or resource-insufficient reset. - while (drain_state_.sync_start_pending.load(std::memory_order_acquire) != 0) { + PTO2TaskSlotState *slot_state = drain_state_.pending_task.load(std::memory_order_acquire); + bool gated = slot_state != nullptr && slot_state->payload != nullptr && + PTO2SchedulerState::owns_early_sync_drain(*slot_state->payload); + + if (elected) { + if (slot_state == nullptr) { + drain_state_.drain_worker_elected.store(0, std::memory_order_release); + return; + } + PTO2ResourceShape shape = slot_state->active_mask.to_shape(); + int32_t available = + count_global_available(shape, slot_state->active_mask.core_mask(), /*include_pending=*/gated); + if (available < block_num) { + drain_state_.drain_ack_mask.store(0, std::memory_order_release); + drain_state_.drain_worker_elected.store(0, std::memory_order_release); + return; + } + drain_state_.drain_running_staged.store(0, std::memory_order_relaxed); + drain_state_.drain_stage_done_mask.store(0, std::memory_order_relaxed); + drain_state_.drain_stage_go.store(1, std::memory_order_release); + } else { + while (drain_state_.drain_stage_go.load(std::memory_order_acquire) == 0) { if (is_completed()) return; if (drain_state_.drain_worker_elected.load(std::memory_order_acquire) == 0) return; SPIN_WAIT_HINT(); } - return; + slot_state = drain_state_.pending_task.load(std::memory_order_acquire); + if (slot_state == nullptr) return; + gated = slot_state->payload != nullptr && PTO2SchedulerState::owns_early_sync_drain(*slot_state->payload); } - // Elected: check if global resources are sufficient. - PTO2TaskSlotState *slot_state = drain_state_.pending_task.load(std::memory_order_acquire); - if (slot_state == nullptr) { - // pending_task is observed null only when a concurrent drain completion - // already cleared it (drain_worker_dispatch nulls it before reopening the - // gate). That drain is done and this is a stale-elected thread, so just - // release the election lock and return. Do NOT clear drain_ack_mask or - // sync_start_pending: a *new* drain run may already be active and - // accumulating acks, and zeroing them would corrupt it into a hang. - drain_state_.drain_worker_elected.store(0, std::memory_order_release); +#if SIMPLER_DFX + if (drain_prof) drain_acked_ts = get_sys_cnt_aicpu(); +#endif + int32_t my_running = drain_stage_cores(slot_state, block_num, thread_idx, gated); +#if SIMPLER_DFX + if (drain_prof && drain_acked_ts != 0) *out_stage_wall_cycles = get_sys_cnt_aicpu() - drain_acked_ts; +#endif + drain_state_.drain_running_staged.fetch_add(my_running, std::memory_order_acq_rel); + drain_state_.drain_stage_done_mask.fetch_or(1u << thread_idx, std::memory_order_release); + + if (!elected) { + while (drain_state_.sync_start_pending.load(std::memory_order_acquire) != 0) { + if (is_completed()) return; + if (drain_state_.drain_worker_elected.load(std::memory_order_acquire) == 0) return; + SPIN_WAIT_HINT(); + } return; } - PTO2ResourceShape shape = slot_state->active_mask.to_shape(); - int32_t available = count_global_available(shape, slot_state->active_mask.core_mask()); - if (available < block_num) { - // Insufficient resources -- reset drain fields so threads can resume - // completion polling to free running cores, then retry. - drain_state_.drain_ack_mask.store(0, std::memory_order_release); - drain_state_.drain_worker_elected.store(0, std::memory_order_release); - return; + while ((drain_state_.drain_stage_done_mask.load(std::memory_order_acquire) & all_acked) != all_acked) { + if (is_completed()) return; + SPIN_WAIT_HINT(); } + if (gated) { + slot_state->payload->running_slot_count.store( + static_cast(drain_state_.drain_running_staged.load(std::memory_order_acquire)), + std::memory_order_seq_cst + ); + } + std::atomic_thread_fence(std::memory_order_release); + drain_state_.pending_task.store(nullptr, std::memory_order_release); + drain_state_.drain_stage_go.store(0, std::memory_order_relaxed); + drain_state_.drain_stage_done_mask.store(0, std::memory_order_relaxed); + drain_state_.drain_ack_mask.store(0, std::memory_order_relaxed); + drain_state_.drain_worker_elected.store(0, std::memory_order_relaxed); + drain_state_.sync_start_pending.store(0, std::memory_order_release); - // Dispatch -- all other threads are spinning, elected thread has exclusive tracker access. - drain_worker_dispatch(runtime, block_num); + if (gated) { + sched_->retry_sync_start_rendezvous_after_drain(*slot_state); + } else { + sched_->propagate_dispatch_fanin(*slot_state); + } + PTO2SchedulerState::finish_early_sync_drain(*slot_state->payload); } diff --git a/src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_context.h b/src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_context.h index c60340ef33..1f13d02aa9 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_context.h +++ b/src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_context.h @@ -11,6 +11,8 @@ #ifndef SCHEDULER_CONTEXT_H #define SCHEDULER_CONTEXT_H +#include "aicpu/device_phase_aicpu.h" +#include "aicpu/platform_regs.h" #include "common/l2_swimlane_profiling.h" #include "common/unified_log.h" #include "scheduler_types.h" @@ -247,22 +249,51 @@ class SchedulerContext { void build_payload( PTO2DispatchPayload &dispatch_payload, PTO2TaskSlotState &slot_state, PTO2SubtaskSlot subslot, - const AsyncCtx &async_ctx, int32_t block_idx + const AsyncCtx &async_ctx, int32_t block_idx, bool force_gate = false ); - void dispatch_subtask_to_core( - Runtime *runtime, int32_t thread_idx, int32_t core_offset, PTO2TaskSlotState &slot_state, - PTO2SubtaskSlot subslot, bool to_pending, int32_t block_idx - ); + // Batched-dispatch primitives. prepare_* builds the payload and per-core + // state; publish_* issues the MMIO register write. Callers must wmb() + // between the prepare batch and the publish batch, then sample + // get_sys_cnt_aicpu() once and pass it to publish_* for every handle. + // + // dispatch_timestamp_slot points to the CoreExecState slot + // (pending_dispatch_timestamp / running_dispatch_timestamp) selected at + // prepare time, or nullptr when L2 swimlane is below AICPU_TIMING and no + // dispatch timestamp is being recorded. + struct PublishHandle { + uint64_t reg_addr; + uint32_t reg_task_id; + int32_t core_offset; + uint64_t *dispatch_timestamp_slot; + int32_t task_timing_slot; // TASK_TIMING_SLOT_NONE unless the task is tagged + }; - void dispatch_mix_block_to_cluster( - Runtime *runtime, int32_t thread_idx, int32_t cluster_offset, PTO2TaskSlotState &slot_state, bool to_pending, - int32_t block_idx + PublishHandle prepare_subtask_to_core( + int32_t thread_idx, int32_t core_offset, PTO2TaskSlotState &slot_state, PTO2SubtaskSlot subslot, + bool to_pending, int32_t block_idx, bool force_gate = false ); - void dispatch_block( - Runtime *runtime, int32_t thread_idx, int32_t core_offset, PTO2TaskSlotState &slot_state, - PTO2ResourceShape shape, bool to_pending, int32_t block_idx + // `thread_idx` is the publishing Scheduler thread's index, used to select the + // per-thread task-timing record; every call site already has it in scope. + inline void publish_subtask_to_core(const PublishHandle &h, uint64_t dispatch_ts, int32_t thread_idx) { + if (h.dispatch_timestamp_slot != nullptr) { + *h.dispatch_timestamp_slot = dispatch_ts; + } + // Task-timing dispatch: earliest DATA_MAIN_BASE publication for a tagged + // task, folded as min. Untagged tasks pay only this cache-hot compare and + // never read the sys counter. Independent of L2 swimlane level. + if (h.task_timing_slot != TASK_TIMING_SLOT_NONE) { + aicpu_task_timing_dispatch(h.task_timing_slot, thread_idx); + } + write_reg(h.reg_addr, RegId::DATA_MAIN_BASE, static_cast(h.reg_task_id)); + } + + // Fan out one block's subtasks (1 for AIC/AIV, 1-3 for MIX) into the + // caller-supplied handles buffer. Returns the number of handles written. + int prepare_block_for_dispatch( + int32_t thread_idx, int32_t core_offset, PTO2TaskSlotState &slot_state, PTO2ResourceShape shape, + bool to_pending, int32_t block_idx, PublishHandle *out_handles, bool force_gate = false ); void dispatch_shape( @@ -287,6 +318,16 @@ class SchedulerContext { bool &try_pushed ); + // Phase 4b: early-dispatch onto spare cores after normal dispatch. + int32_t try_early_dispatch( + int32_t thread_idx, CoreTracker &tracker, bool pmu_active, bool &made_progress, bool &try_pushed + ); + int32_t stage_consumer_blocks( + int32_t thread_idx, PTO2TaskSlotState *c, PTO2ResourceShape shape, int32_t start, int32_t count, + CoreTracker::BitStates &idle, CoreTracker::BitStates &pend + ); + int32_t early_dispatch_shape(int32_t thread_idx, PTO2ResourceShape shape, CoreTracker::DispatchPhase phase); + // 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 @@ -302,12 +343,17 @@ class SchedulerContext { return sched_->ready_queues[static_cast(PTO2ResourceShape::MIX)].size() > 0; } + bool has_residual_early_mix() const { + return sched_->early_dispatch_queues[static_cast(PTO2ResourceShape::MIX)].size() > 0; + } + // ========================================================================= // Completion & drain (scheduler_completion.cpp) // ========================================================================= - static SlotTransition - decide_slot_transition(int32_t reg_task_id, int32_t reg_state, int32_t running_id, int32_t pending_id); + static SlotTransition decide_slot_transition( + int32_t reg_task_id, int32_t reg_state, int32_t running_id, int32_t pending_id, bool pending_gated = false + ); void complete_slot_task( PTO2TaskSlotState &slot_state, int32_t expected_reg_task_id, PTO2SubtaskSlot subslot, int32_t thread_idx, @@ -328,9 +374,9 @@ class SchedulerContext { ); bool enter_drain_mode(PTO2TaskSlotState *slot_state, int32_t block_num); - int32_t count_global_available(PTO2ResourceShape shape, uint8_t core_mask); - void drain_worker_dispatch(Runtime *runtime, int32_t block_num); - void handle_drain_mode(Runtime *runtime, int32_t thread_idx); + int32_t count_global_available(PTO2ResourceShape shape, uint8_t core_mask, bool include_pending = false); + int32_t drain_stage_cores(PTO2TaskSlotState *slot_state, int32_t block_num, int32_t thread_idx, bool gated); + void handle_drain_mode(int32_t thread_idx, uint64_t *out_stage_wall_cycles = nullptr); // ========================================================================= // Cold path: exit checks, stall diagnostics, profiling (scheduler_cold_path.cpp) diff --git a/src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_dispatch.cpp b/src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_dispatch.cpp index 8007169b05..864f8581bb 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_dispatch.cpp +++ b/src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_dispatch.cpp @@ -41,6 +41,15 @@ namespace { inline constexpr int32_t PTO2_DEFERRED_RELEASE_CAP = 256; } +// AICore materializes args[] from src_payload on the gated path using these +// offsets; pin them against the live PTO2TaskPayload layout. +static_assert(offsetof(PTO2TaskPayload, tensor_count) == PTO2_TASKPAYLOAD_TENSOR_COUNT_OFFSET); +static_assert(offsetof(PTO2TaskPayload, scalar_count) == PTO2_TASKPAYLOAD_SCALAR_COUNT_OFFSET); +static_assert(offsetof(PTO2TaskPayload, tensors) == PTO2_TASKPAYLOAD_TENSORS_OFFSET); +static_assert(offsetof(PTO2TaskPayload, scalars) == PTO2_TASKPAYLOAD_SCALARS_OFFSET); +static_assert(sizeof(Tensor) == PTO2_TASKPAYLOAD_TENSOR_STRIDE); +static_assert(RUNTIME_MAX_WORKER <= PTO2_EARLY_DISPATCH_CORE_MASK_WORDS * 64); + const char *SchedulerContext::shape_name(PTO2ResourceShape shape) { switch (shape) { case PTO2ResourceShape::AIC: @@ -104,19 +113,26 @@ int SchedulerContext::pop_ready_tasks_batch( void SchedulerContext::build_payload( PTO2DispatchPayload &dispatch_payload, PTO2TaskSlotState &slot_state, PTO2SubtaskSlot subslot, - const AsyncCtx &async_ctx, int32_t block_idx + const AsyncCtx &async_ctx, int32_t block_idx, bool force_gate ) { int32_t slot_idx = static_cast(subslot); uint64_t callable_addr = get_function_bin_addr(slot_state.task->kernel_id[slot_idx]); const CoreCallable *callable = reinterpret_cast(callable_addr); dispatch_payload.function_bin_addr = callable->resolved_addr(); auto &payload = *slot_state.payload; - int n = 0; - for (int32_t i = 0; i < payload.tensor_count; i++) { - dispatch_payload.args[n++] = reinterpret_cast(&payload.tensors[i]); - } - for (int32_t i = 0; i < payload.scalar_count; i++) { - dispatch_payload.args[n++] = payload.scalars[i]; + if (PTO2SchedulerState::should_gate_early_dispatch( + force_gate, payload.early_dispatch_state.load(std::memory_order_relaxed) + )) { + dispatch_payload.src_payload = reinterpret_cast(&payload); + } else { + dispatch_payload.src_payload = 0; + int n = 0; + for (int32_t i = 0; i < payload.tensor_count; i++) { + dispatch_payload.args[n++] = reinterpret_cast(&payload.tensors[i]); + } + for (int32_t i = 0; i < payload.scalar_count; i++) { + dispatch_payload.args[n++] = payload.scalars[i]; + } } dispatch_payload.local_context.s_block_idx = block_idx; dispatch_payload.local_context.s_block_num = slot_state.logical_block_num; @@ -125,14 +141,14 @@ void SchedulerContext::build_payload( dispatch_payload.args[PAYLOAD_GLOBAL_CONTEXT_INDEX] = reinterpret_cast(&dispatch_payload.global_context); } -void SchedulerContext::dispatch_subtask_to_core( - Runtime *runtime, int32_t thread_idx, int32_t core_offset, PTO2TaskSlotState &slot_state, PTO2SubtaskSlot subslot, - bool to_pending, int32_t block_idx +SchedulerContext::PublishHandle SchedulerContext::prepare_subtask_to_core( + int32_t thread_idx, int32_t core_offset, PTO2TaskSlotState &slot_state, PTO2SubtaskSlot subslot, bool to_pending, + int32_t block_idx, bool force_gate ) { CoreTracker &tracker = core_trackers_[thread_idx]; auto core_id = tracker.get_core_id_by_offset(core_offset); - (void)runtime; CoreExecState &core_exec_state = core_exec_states_[core_id]; + core_exec_state.dispatch_seq++; uint32_t reg_task_id = core_exec_state.dispatch_seq & TASK_ID_MASK; static_assert( @@ -145,11 +161,15 @@ void SchedulerContext::dispatch_subtask_to_core( uint32_t buf_idx = reg_task_id & 1u; PTO2DispatchPayload &payload = payload_per_core_[core_id][buf_idx]; + // a5 clears the deferred slab per dispatch (unlike a2a3's init-once path): + // AsyncCtx::make wires the slab into the payload, and a stale count from a + // prior deferred completion on this (core, buf) would be observed as a live + // wait entry. DeferredCompletionSlab *deferred_slab = &deferred_slab_per_core_[core_id][buf_idx]; deferred_slab->count = 0; deferred_slab->error_code = PTO2_ERROR_NONE; AsyncCtx async_ctx = AsyncCtx::make(slot_state.task->task_id, deferred_slab); - build_payload(payload, slot_state, subslot, async_ctx, block_idx); + build_payload(payload, slot_state, subslot, async_ctx, block_idx, force_gate); if (to_pending) { core_exec_state.pending_subslot = subslot; @@ -161,6 +181,7 @@ void SchedulerContext::dispatch_subtask_to_core( core_exec_state.running_reg_task_id = static_cast(reg_task_id); tracker.change_core_state(core_offset); } + tracker.set_pending_occupied(core_offset); LOG_DEBUG( "Thread %d: Dispatched %s %s task %" PRId64 " kernel_id=[%d,%d,%d] block_idx=%d/total_blocks=%d to" @@ -182,68 +203,22 @@ void SchedulerContext::dispatch_subtask_to_core( } #endif - // Publish task data (slot_state / args writes done above) before AICore - // can observe the dispatched task_id. ARM64 needs an explicit store-store - // fence across Normal-cacheable -> Device-nGnRnE; the old write_reg() - // helper provided this implicitly via __sync_synchronize. - wmb(); - - // Capture dispatch timestamp at the latest possible moment — after wmb, - // immediately before the DATA_MAIN_BASE write. + uint64_t *dispatch_timestamp_slot = nullptr; #if SIMPLER_DFX if (l2_swimlane_level_ >= L2SwimlaneLevel::AICPU_TIMING) { - uint64_t dispatch_ts = get_sys_cnt_aicpu(); - if (to_pending) { - core_exec_state.pending_dispatch_timestamp = dispatch_ts; - } else { - core_exec_state.running_dispatch_timestamp = dispatch_ts; - } + dispatch_timestamp_slot = + to_pending ? &core_exec_state.pending_dispatch_timestamp : &core_exec_state.running_dispatch_timestamp; } #endif - // Task-timing dispatch: earliest DATA_MAIN_BASE publication for a tagged - // task, folded as min. Untagged tasks pay only this cache-hot compare and - // never read the sys counter. Independent of L2 swimlane level. - if (slot_state.task->task_timing_slot != TASK_TIMING_SLOT_NONE) { - aicpu_task_timing_dispatch(slot_state.task->task_timing_slot, thread_idx); - } - - write_reg(core_exec_state.reg_addr, RegId::DATA_MAIN_BASE, static_cast(reg_task_id)); - tracker.set_pending_occupied(core_offset); -} - -void SchedulerContext::dispatch_mix_block_to_cluster( - Runtime *runtime, int32_t thread_idx, int32_t cluster_offset, PTO2TaskSlotState &slot_state, bool to_pending, - int32_t block_idx -) { - CoreTracker &tracker = core_trackers_[thread_idx]; - uint8_t cmask = slot_state.active_mask.core_mask(); - if (cmask & PTO2_SUBTASK_MASK_AIC) { - bool aic_to_pending = to_pending && !tracker.is_aic_core_idle(cluster_offset); - dispatch_subtask_to_core( - runtime, thread_idx, tracker.get_aic_core_offset(cluster_offset), slot_state, PTO2SubtaskSlot::AIC, - aic_to_pending, block_idx - ); - } - if (cmask & PTO2_SUBTASK_MASK_AIV0) { - bool aiv0_to_pending = to_pending && !tracker.is_aiv0_core_idle(cluster_offset); - dispatch_subtask_to_core( - runtime, thread_idx, tracker.get_aiv0_core_offset(cluster_offset), slot_state, PTO2SubtaskSlot::AIV0, - aiv0_to_pending, block_idx - ); - } - if (cmask & PTO2_SUBTASK_MASK_AIV1) { - bool aiv1_to_pending = to_pending && !tracker.is_aiv1_core_idle(cluster_offset); - dispatch_subtask_to_core( - runtime, thread_idx, tracker.get_aiv1_core_offset(cluster_offset), slot_state, PTO2SubtaskSlot::AIV1, - aiv1_to_pending, block_idx - ); - } + return PublishHandle{ + core_exec_state.reg_addr, reg_task_id, core_offset, dispatch_timestamp_slot, slot_state.task->task_timing_slot + }; } -void SchedulerContext::dispatch_block( - Runtime *runtime, int32_t thread_idx, int32_t core_offset, PTO2TaskSlotState &slot_state, PTO2ResourceShape shape, - bool to_pending, int32_t block_idx +int SchedulerContext::prepare_block_for_dispatch( + int32_t thread_idx, int32_t core_offset, PTO2TaskSlotState &slot_state, PTO2ResourceShape shape, bool to_pending, + int32_t block_idx, PublishHandle *out_handles, bool force_gate ) { #if SIMPLER_DFX if (is_dump_args_enabled()) { @@ -258,26 +233,63 @@ void SchedulerContext::dispatch_block( ); } #endif + CoreTracker &tracker = core_trackers_[thread_idx]; if (shape == PTO2ResourceShape::MIX) { - dispatch_mix_block_to_cluster(runtime, thread_idx, core_offset, slot_state, to_pending, block_idx); + uint8_t cmask = slot_state.active_mask.core_mask(); + int n = 0; + // Preserve a5 MIX per-core slot placement: idle used cores take the + // running slot; busy used cores take pending when to_pending. Cluster + // selection remains classify_mix_cluster (uniform RUNNING/PENDING), not + // a2a3 gated MIX split. + if (cmask & PTO2_SUBTASK_MASK_AIC) { + bool p = to_pending && !tracker.is_aic_core_idle(core_offset); + out_handles[n++] = prepare_subtask_to_core( + thread_idx, tracker.get_aic_core_offset(core_offset), slot_state, PTO2SubtaskSlot::AIC, p, block_idx, + force_gate + ); + } + if (cmask & PTO2_SUBTASK_MASK_AIV0) { + bool p = to_pending && !tracker.is_aiv0_core_idle(core_offset); + out_handles[n++] = prepare_subtask_to_core( + thread_idx, tracker.get_aiv0_core_offset(core_offset), slot_state, PTO2SubtaskSlot::AIV0, p, block_idx, + force_gate + ); + } + if (cmask & PTO2_SUBTASK_MASK_AIV1) { + bool p = to_pending && !tracker.is_aiv1_core_idle(core_offset); + out_handles[n++] = prepare_subtask_to_core( + thread_idx, tracker.get_aiv1_core_offset(core_offset), slot_state, PTO2SubtaskSlot::AIV1, p, block_idx, + force_gate + ); + } +#if SIMPLER_DFX + sched_l2_swimlane_[thread_idx].phase_dispatch_count += __builtin_popcount(cmask); +#endif + return n; } else if (shape == PTO2ResourceShape::AIC) { - dispatch_subtask_to_core( - runtime, thread_idx, core_offset, slot_state, PTO2SubtaskSlot::AIC, to_pending, block_idx + out_handles[0] = prepare_subtask_to_core( + thread_idx, core_offset, slot_state, PTO2SubtaskSlot::AIC, to_pending, block_idx, force_gate ); +#if SIMPLER_DFX + sched_l2_swimlane_[thread_idx].phase_dispatch_count += 1; +#endif + return 1; } else { - dispatch_subtask_to_core( - runtime, thread_idx, core_offset, slot_state, PTO2SubtaskSlot::AIV0, to_pending, block_idx + out_handles[0] = prepare_subtask_to_core( + thread_idx, core_offset, slot_state, PTO2SubtaskSlot::AIV0, to_pending, block_idx, force_gate ); - } #if SIMPLER_DFX - sched_l2_swimlane_[thread_idx].phase_dispatch_count += __builtin_popcount(slot_state.active_mask.core_mask()); + sched_l2_swimlane_[thread_idx].phase_dispatch_count += 1; #endif + return 1; + } } void SchedulerContext::dispatch_shape( Runtime *runtime, int32_t thread_idx, PTO2ResourceShape shape, CoreTracker::DispatchPhase phase, CoreTracker &tracker, bool &entered_drain, bool &made_progress, bool &try_pushed ) { + (void)runtime; #if SIMPLER_SCHED_PROFILING auto &l2_swimlane = sched_l2_swimlane_[thread_idx]; #endif @@ -294,7 +306,67 @@ void SchedulerContext::dispatch_shape( int got = pop_ready_tasks_batch(shape, thread_idx, batch, want); if (got == 0) break; + // sync_start exclusion gate. + // + // When the popped batch contains a sync_start task we MUST publish each + // prior task with its own wmb so AICore receives them with time + // separation. The drain coordinator's `count_global_available()` check + // reads the per-thread CoreTracker, and although `prepare_block_for_dispatch` + // marks cores occupied synchronously, the head-start between successive + // tasks is what lets the surrounding completion loop catch up on FINs in + // the retry window when the sync_start task hits insufficient resources. + // Bursting all prior tasks at the end of the pop (cross-task batching) + // collapses that head-start and causes spmd_sync_start_stress to time + // out via 507018 on ~40% of runs — see + // docs/investigations/2026-06-cross-task-batched-publish.md. + // + // When the batch carries no sync_start task, no drain entry can happen + // in this pop, so we hoist `handles[]`, `wmb()`, and the publish loop + // out of the per-task body. One wmb amortizes across all tasks and one + // dispatch_ts is shared, which restores ~60 ns first-to-last AICore + // start span for single-block decode kernels (out_proj, q_proj, ...). + // Detection is a single mask check per task — cheap relative to even + // one register write. + bool any_sync_start = false; + for (int bi = 0; bi < got; bi++) { + if (batch[bi]->active_mask.requires_sync_start()) { + any_sync_start = true; + break; + } + } + + // handles[] is sized for the MIX worst case: total claims across the + // pop bounded by `cores.count() ≤ MAX_CLUSTERS`, and each block + // contributes ≤ 3 subtasks for MIX. + PublishHandle handles[CoreTracker::MAX_CLUSTERS * 3]; + int handle_count = 0; + PTO2TaskSlotState *published_list[CoreTracker::MAX_CLUSTERS * 3]; + int16_t published_counts[CoreTracker::MAX_CLUSTERS * 3]; + int published_n = 0; bool dispatched_any = false; +#if SIMPLER_SCHED_PROFILING + uint64_t t_setup_start = get_sys_cnt_aicpu(); +#endif + + // Flush prepared-but-unpublished handles. Required before + // `enter_drain_mode` so the drain coordinator sees cores as occupied, + // and at the per-task boundary when `any_sync_start` is true. + auto flush_publish = [&]() { + if (handle_count == 0) return; + wmb(); + uint64_t dispatch_ts = 0; +#if SIMPLER_DFX + if (l2_swimlane_level_ >= L2SwimlaneLevel::AICPU_TIMING) { + dispatch_ts = get_sys_cnt_aicpu(); + } +#endif + for (int i = 0; i < handle_count; i++) { + publish_subtask_to_core(handles[i], dispatch_ts, thread_idx); + } + handle_count = 0; + made_progress = true; + }; + for (int bi = 0; bi < got; bi++) { PTO2TaskSlotState *slot_state = batch[bi]; CoreTracker::BitStates selected_mix_clusters(0ULL); @@ -322,6 +394,7 @@ void SchedulerContext::dispatch_shape( } 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); } @@ -334,29 +407,25 @@ void SchedulerContext::dispatch_shape( } if (!cores.has_value()) { + flush_publish(); sched_->ready_queues[static_cast(shape)].push_batch(&batch[bi], got - bi); break; } - dispatched_any = true; - try_pushed = true; -#if SIMPLER_SCHED_PROFILING - uint64_t t_setup_start = get_sys_cnt_aicpu(); -#endif // Claim a contiguous range of blocks, hand the slot back to the // ready queue immediately, then perform the expensive dispatches. - // This lets other schedulers concurrently claim and dispatch the - // remaining blocks of the same SPMD task instead of spinning while - // this thread fills all its own cores. Only local `start + b` is - // read after the push -- `next_block_idx` may already be advanced - // by another scheduler that popped the slot. - int32_t remaining = slot_state->logical_block_num - slot_state->next_block_idx; int32_t available = is_mix ? selected_mix_clusters.count() : cores.count(); - int32_t claim = std::min(available, remaining); - int32_t start = slot_state->next_block_idx; - slot_state->next_block_idx += claim; + int32_t start = 0; + int32_t claim = slot_state->claim_block_range(slot_state->logical_block_num, available, start); + if (claim == 0) continue; + dispatched_any = true; + try_pushed = true; + + published_list[published_n] = slot_state; + published_counts[published_n] = static_cast(claim); + published_n++; - if (slot_state->next_block_idx < slot_state->logical_block_num) { + if (start + claim < slot_state->logical_block_num) { sched_->ready_queues[static_cast(shape)].push(slot_state); } @@ -365,13 +434,28 @@ void SchedulerContext::dispatch_shape( if (is_mix) { cores.clear_bit(core_offset); } - dispatch_block(runtime, thread_idx, core_offset, *slot_state, shape, is_pending, start + b); + handle_count += prepare_block_for_dispatch( + thread_idx, core_offset, *slot_state, shape, is_pending, start + b, &handles[handle_count] + ); } - made_progress = true; + + // Sync_start exclusion: flush per task so prior tasks have head- + // start time before any sync_start drain check. Normal batches + // fall through and accumulate for one cross-task flush at the + // end of the pop. + if (any_sync_start) { + flush_publish(); + } + } + + flush_publish(); + for (int i = 0; i < published_n; i++) { + sched_->record_published_blocks(*published_list[i], published_counts[i]); + sched_->propagate_dispatch_fanin(*published_list[i]); + } #if SIMPLER_SCHED_PROFILING - l2_swimlane.sched_dispatch_setup_cycle += (get_sys_cnt_aicpu() - t_setup_start); + l2_swimlane.sched_dispatch_setup_cycle += (get_sys_cnt_aicpu() - t_setup_start); #endif - } if (!dispatched_any) break; @@ -381,6 +465,188 @@ void SchedulerContext::dispatch_shape( } } +int32_t SchedulerContext::stage_consumer_blocks( + int32_t thread_idx, PTO2TaskSlotState *c, PTO2ResourceShape shape, int32_t start, int32_t count, + CoreTracker::BitStates &idle, CoreTracker::BitStates &pend +) { + CoreTracker &tracker = core_trackers_[thread_idx]; + uint64_t early_dispatch_ts = get_sys_cnt_aicpu(); + uint64_t my_cores[PTO2_EARLY_DISPATCH_CORE_MASK_WORDS] = {0}; + int32_t staged = 0; + int32_t block = start; + PublishHandle handles[CoreTracker::MAX_CLUSTERS * 3]; + int n = 0; + auto prepare_from = [&](CoreTracker::BitStates &avail, bool to_pending) { + while (count > 0 && avail.has_value()) { + int32_t core_offset = avail.pop_first(); + n += prepare_block_for_dispatch( + thread_idx, core_offset, *c, shape, to_pending, block, &handles[n], /*force_gate=*/true + ); + block++; + count--; + staged++; + } + }; + if (idle.has_value()) prepare_from(idle, /*to_pending=*/false); + if (pend.has_value()) prepare_from(pend, /*to_pending=*/true); + if (n > 0) { + wmb(); + for (int i = 0; i < n; i++) { + publish_subtask_to_core(handles[i], early_dispatch_ts, thread_idx); + int32_t cid = tracker.get_core_id_by_offset(handles[i].core_offset); + sched_->early_dispatch_doorbell_table[cid].addr = handles[i].reg_addr; + sched_->early_dispatch_doorbell_table[cid].token = handles[i].reg_task_id; + my_cores[cid >> 6] |= (1ULL << (cid & 63)); + } + } + 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); + + bool released = staged > 0 && + c->payload->early_dispatch_state.load(std::memory_order_seq_cst) == PTO2_EARLY_DISPATCH_DISPATCHED; + if (released) { + uint64_t owned[PTO2_EARLY_DISPATCH_CORE_MASK_WORDS] = {0}; + for (int w = 0; w < PTO2_EARLY_DISPATCH_CORE_MASK_WORDS; w++) { + if (my_cores[w] != 0) { + owned[w] = + PTO2SchedulerState::claim_late_staged_doorbell_bits(c->payload->staged_core_mask[w], my_cores[w]); + } + } + for (int i = 0; i < n; i++) { + int32_t cid = tracker.get_core_id_by_offset(handles[i].core_offset); + PTO2SchedulerState::ring_claimed_local_doorbell( + owned[cid >> 6], cid, handles[i].reg_addr, handles[i].reg_task_id + ); + } + wmb(); + } + sched_->record_published_blocks(*c, staged); + sched_->propagate_dispatch_fanin(*c); + return staged; +} + +int32_t +SchedulerContext::early_dispatch_shape(int32_t thread_idx, PTO2ResourceShape shape, CoreTracker::DispatchPhase phase) { + CoreTracker &tracker = core_trackers_[thread_idx]; + int32_t s = static_cast(shape); + bool is_mix = (shape == PTO2ResourceShape::MIX); + bool is_idle = (phase == CoreTracker::DispatchPhase::IDLE); + + CoreTracker::BitStates cores = + is_mix ? tracker.get_cluster_offset_states() : tracker.get_dispatchable_cores(shape, phase); + if (!cores.has_value()) return 0; + + int32_t total_staged = 0; + PTO2TaskSlotState *batch[CoreTracker::MAX_CLUSTERS * 3]; + uint64_t task_id_snapshots[CoreTracker::MAX_CLUSTERS * 3]; + int got = sched_->early_dispatch_queues[s].pop_batch_tagged(batch, task_id_snapshots, cores.count()); + for (int bi = 0; bi < got; bi++) { + PTO2TaskSlotState *c = batch[bi]; + if (static_cast(c->task->task_id.raw) != task_id_snapshots[bi]) continue; + if (c->payload->early_dispatch_state.load(std::memory_order_acquire) != PTO2_EARLY_DISPATCH_STAGING) continue; + + 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) { + sched_->early_dispatch_queues[s].push_batch_tagged(&batch[bi], &task_id_snapshots[bi], got - bi); + break; + } + int32_t start = 0; + int32_t claim = c->claim_block_range(c->logical_block_num, freecores, start); + if (claim == 0) continue; + if (start + claim < c->logical_block_num) { + if (!sched_->early_dispatch_queues[s].push_tagged(c, task_id_snapshots[bi])) + LOG_INFO_V9( + "[EARLY_DISPATCH] queue full on re-push, consumer=%" PRId64, + static_cast(c->task->task_id.raw) + ); + } + 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; +} + +int32_t SchedulerContext::try_early_dispatch( + int32_t thread_idx, CoreTracker &tracker, bool pmu_active, bool &made_progress, bool &try_pushed +) { + // Gate (a2a3 #1288): owned here rather than by the caller. + // - pmu_active: staging gated work perturbs single-issue PMU windows. + // - has_any_free_slot: spare capacity (local read; fully-occupied bails + // before touching shared queues) — not the old fully-idle pass. + // - ready queues empty: normal dispatch strictly precedes early. + // a5 has no ready_sync_queues[]; ready_queues[] cover the normal lane. + if (pmu_active || !tracker.has_any_free_slot()) return 0; + for (int s = 0; s < PTO2_NUM_RESOURCE_SHAPES; s++) { + if (sched_->ready_queues[s].size() > 0) return 0; + } + + // Tier 0: sync_start early cohorts (shape-agnostic, all-or-nothing drain). + uint64_t sync_task_id_snapshot = 0; + if (PTO2TaskSlotState *c = sched_->early_sync_start_queue.pop_tagged(&sync_task_id_snapshot)) { + bool current_sync_task = static_cast(c->task->task_id.raw) == sync_task_id_snapshot && + c->active_mask.requires_sync_start(); + if (current_sync_task && PTO2SchedulerState::try_claim_early_sync_drain(*c->payload)) { + if (c->payload->early_dispatch_state.load(std::memory_order_seq_cst) != PTO2_EARLY_DISPATCH_STAGING) { + sched_->cancel_early_sync_drain(*c); + } else if (enter_drain_mode(c, c->logical_block_num)) { + PTO2SchedulerState::mark_early_sync_drain_armed(*c->payload); + } else { + sched_->cancel_early_sync_drain(*c); + } + } + } + + using Phase = CoreTracker::DispatchPhase; + static constexpr PTO2ResourceShape kAicAivOrder[2][2] = { + {PTO2ResourceShape::AIC, PTO2ResourceShape::AIV}, + {PTO2ResourceShape::AIV, PTO2ResourceShape::AIC}, + }; + const PTO2ResourceShape *aic_aiv = kAicAivOrder[thread_idx & 1]; + + int32_t total_staged = 0; + total_staged += early_dispatch_shape(thread_idx, PTO2ResourceShape::MIX, Phase::IDLE); + bool skip_aic_aiv = has_residual_early_mix(); + if (!skip_aic_aiv) { + for (int i = 0; i < 2; i++) { + total_staged += early_dispatch_shape(thread_idx, aic_aiv[i], Phase::IDLE); + } + } + if (!pmu_active) { + if (!has_idle_in_other_threads(thread_idx, PTO2ResourceShape::MIX)) { + total_staged += early_dispatch_shape(thread_idx, PTO2ResourceShape::MIX, Phase::PENDING); + } + if (!skip_aic_aiv && has_residual_early_mix()) skip_aic_aiv = true; + if (!skip_aic_aiv) { + for (int i = 0; i < 2; i++) { + PTO2ResourceShape s = aic_aiv[i]; + if (has_idle_in_other_threads(thread_idx, s)) continue; + total_staged += early_dispatch_shape(thread_idx, s, Phase::PENDING); + } + } + } + + if (total_staged > 0) { + made_progress = true; + try_pushed = true; + } + return total_staged; +} + void SchedulerContext::dispatch_ready_tasks( Runtime *runtime, int32_t thread_idx, CoreTracker &tracker, bool pmu_active, bool &made_progress, bool &try_pushed ) { @@ -730,7 +996,7 @@ int32_t SchedulerContext::resolve_and_dispatch(Runtime *runtime, int32_t thread_ // Phase 2 drain check if (drain_state_.sync_start_pending.load(std::memory_order_acquire) != 0) { - handle_drain_mode(runtime, thread_idx); + handle_drain_mode(thread_idx); continue; } @@ -812,6 +1078,9 @@ int32_t SchedulerContext::resolve_and_dispatch(Runtime *runtime, int32_t thread_ // pmu_active is cached at function scope above (loop-invariant). dispatch_ready_tasks(runtime, thread_idx, tracker, pmu_active, made_progress, try_pushed); + // Phase 4b: early-dispatch onto spare cores after normal dispatch. + (void)try_early_dispatch(thread_idx, tracker, pmu_active, made_progress, try_pushed); + #if SIMPLER_DFX if (!try_pushed) { CYCLE_COUNT_LAP(l2_swimlane.sched_idle_cycle); diff --git a/src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_types.h b/src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_types.h index 887569f22c..7f2f10f1bb 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_types.h +++ b/src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_types.h @@ -232,6 +232,12 @@ class alignas(64) CoreTracker { BitStates get_all_running_cores() const { return (~core_states_) & (aic_mask_ | aiv_mask_); } BitStates get_cluster_offset_states() const { return aic_mask_; } + // Free capacity for early-dispatch staging: any core whose pending slot is + // not occupied (idle RUNNING or dual-issue PENDING). Matches a2a3 #1288-ish + // spare-slot gate (not full-idle-only). + 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(); } + // --- Cluster matching --- BitStates get_valid_cluster_offset_states(PTO2ResourceShape shape) const { @@ -352,6 +358,43 @@ class alignas(64) CoreTracker { return get_mix_running_cluster_offset_states(core_mask).count(); } + // Gated MIX split placement (#1304): each used core independently takes + // running-if-idle / pending-if-busy while every used core waits on the doorbell. + BitStates mix_used_cores(int32_t cluster_offset, uint8_t core_mask) const { + BitStates used(0ULL); + if (core_mask & PTO2_SUBTASK_MASK_AIC) used |= BitStates(1ULL << cluster_offset); + if (core_mask & PTO2_SUBTASK_MASK_AIV0) used |= BitStates(1ULL << (cluster_offset + 1)); + if (core_mask & PTO2_SUBTASK_MASK_AIV1) used |= BitStates(1ULL << (cluster_offset + 2)); + return used; + } + + bool mix_cluster_all_slots(int32_t cluster_offset, uint8_t core_mask) const { + BitStates used = mix_used_cores(cluster_offset, core_mask); + if (!used.has_value()) return false; + BitStates no_slot = (~core_states_) & pending_occupied_; + return !(used & no_slot).has_value(); + } + + int32_t mix_cluster_idle_core_count(int32_t cluster_offset, uint8_t core_mask) const { + return (mix_used_cores(cluster_offset, core_mask) & core_states_).count(); + } + + BitStates get_mix_split_cluster_offset_states(uint8_t core_mask) const { + BitStates result(0ULL); + BitStates candidates = get_cluster_offset_states(); + while (candidates.has_value()) { + int32_t off = candidates.pop_first(); + if (mix_cluster_all_slots(off, core_mask)) { + result |= BitStates(1ULL << off); + } + } + return result; + } + + int32_t count_mix_split_clusters(uint8_t core_mask) const { + return get_mix_split_cluster_offset_states(core_mask).count(); + } + BitStates get_pending_core_offset_states(PTO2ResourceShape shape) const { if (shape == PTO2ResourceShape::MIX) { // Shape-level query kept conservative for legacy callers/tests. @@ -456,7 +499,10 @@ struct alignas(64) SyncStartDrainState { std::atomic drain_worker_elected{0}; // 0=none; >0: elected thread's (thread_idx+1) std::atomic drain_ack_mask{0}; // bit per thread; all-set = all threads reached ack barrier std::atomic pending_task{nullptr}; // held task (not re-queued) - int32_t _pad[10]; + std::atomic drain_stage_go{0}; + std::atomic drain_stage_done_mask{0}; + std::atomic drain_running_staged{0}; + int32_t _pad[7]; }; static_assert(sizeof(SyncStartDrainState) == 64); diff --git a/src/a5/runtime/tensormap_and_ringbuffer/runtime/shared/pto_runtime2_init.cpp b/src/a5/runtime/tensormap_and_ringbuffer/runtime/shared/pto_runtime2_init.cpp index b99960ca47..e9a3d8a2e8 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/runtime/shared/pto_runtime2_init.cpp +++ b/src/a5/runtime/tensormap_and_ringbuffer/runtime/shared/pto_runtime2_init.cpp @@ -139,6 +139,10 @@ PTO2SchedulerState::reserve_layout(DeviceArena &arena, const int32_t dep_pool_ca layout.off_ready_queue_slots[i] = ready_queue_reserve_layout(arena, PTO2_READY_QUEUE_SIZE); } layout.off_dummy_ready_queue_slots = ready_queue_reserve_layout(arena, PTO2_READY_QUEUE_SIZE); + for (int i = 0; i < PTO2_NUM_RESOURCE_SHAPES; i++) { + layout.off_early_dispatch_queue_slots[i] = ready_queue_reserve_layout(arena, PTO2_EARLY_DISPATCH_QUEUE_SIZE); + } + layout.off_early_sync_start_queue_slots = ready_queue_reserve_layout(arena, PTO2_EARLY_DISPATCH_QUEUE_SIZE); for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { // Force a cache-line base so Orch-side dep_pool writes do not invalidate // adjacent multi-threaded regions like ready_queue.slots. @@ -176,6 +180,20 @@ 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->early_dispatch_queues[i], arena, layout.off_early_dispatch_queue_slots[i], + PTO2_EARLY_DISPATCH_QUEUE_SIZE + )) { + return false; + } + } + if (!ready_queue_init_data_from_layout( + &sched->early_sync_start_queue, arena, layout.off_early_sync_start_queue_slots, + PTO2_EARLY_DISPATCH_QUEUE_SIZE + )) { + return false; + } auto *orch_err = pto2_sm_layout::orch_error_code_addr(sm_dev_base); for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { @@ -204,6 +222,10 @@ void PTO2SchedulerState::reset_for_reuse(const PTO2SchedulerLayout &layout, void sched->ready_queues[i].reset_for_reuse(); } sched->dummy_ready_queue.reset_for_reuse(); + for (int i = 0; i < PTO2_NUM_RESOURCE_SHAPES; i++) { + sched->early_dispatch_queues[i].reset_for_reuse(); + } + sched->early_sync_start_queue.reset_for_reuse(); sched->async_wait_list.reset_for_reuse(); (void)layout; @@ -215,6 +237,12 @@ void PTO2SchedulerState::wire_arena_pointers(const PTO2SchedulerLayout &layout, ready_queue_wire_arena_pointers(&sched->ready_queues[i], arena, layout.off_ready_queue_slots[i]); } ready_queue_wire_arena_pointers(&sched->dummy_ready_queue, arena, layout.off_dummy_ready_queue_slots); + for (int i = 0; i < PTO2_NUM_RESOURCE_SHAPES; i++) { + ready_queue_wire_arena_pointers( + &sched->early_dispatch_queues[i], arena, layout.off_early_dispatch_queue_slots[i] + ); + } + ready_queue_wire_arena_pointers(&sched->early_sync_start_queue, arena, layout.off_early_sync_start_queue_slots); for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { sched->ring_sched_states[r].dep_pool.base = static_cast(arena.region_ptr(layout.off_dep_pool_entries[r])); @@ -231,6 +259,10 @@ void PTO2SchedulerState::destroy() { ready_queue_destroy(&sched->ready_queues[i]); } ready_queue_destroy(&sched->dummy_ready_queue); + for (int i = 0; i < PTO2_NUM_RESOURCE_SHAPES; i++) { + ready_queue_destroy(&sched->early_dispatch_queues[i]); + } + ready_queue_destroy(&sched->early_sync_start_queue); } // ============================================================================= diff --git a/tests/st/a5/tensormap_and_ringbuffer/spmd_sync_start_early_dispatch/kernels/aiv/kernel_spmd_write_slow.cpp b/tests/st/a5/tensormap_and_ringbuffer/spmd_sync_start_early_dispatch/kernels/aiv/kernel_spmd_write_slow.cpp new file mode 100644 index 0000000000..9b0cae98ea --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/spmd_sync_start_early_dispatch/kernels/aiv/kernel_spmd_write_slow.cpp @@ -0,0 +1,76 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ + +/** + * SPMD Multi-Block Write Kernel with a spin delay (AIV) + * + * Same result as the plain write kernel, but spins for `spin_iters` before the + * store. A slow producer keeps its cores occupied long enough that a dependent + * consumer is processed as an early-dispatch candidate WHILE the producer is still + * running — the only way a trivial-kernel scene exercises the speculative + * gated-dispatch path (a fast producer finishes before the consumer is ever staged). + * + * out[(base_cl + block_idx) * FLOATS_PER_CACHE_LINE] = float(block_idx) + * + * Args: + * args[0] = output Tensor* (INOUT) + * args[1] = scalar: base_cl (starting cache line index for this task) + * args[2] = scalar: spin_iters (0 = no delay) + */ + +#include +#include + +#include "tensor.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] // NOLINT(whitespace/braces) +#endif + +#include "intrinsic.h" + +static constexpr int32_t FLOATS_PER_CACHE_LINE = 16; + +#ifdef PTO_CPUSTUB_HPP +#define dcci(...) \ + do { \ + } while (0) +#endif +#ifndef SINGLE_CACHE_LINE +#define SINGLE_CACHE_LINE 0 +#endif +#ifndef CACHELINE_OUT +#define CACHELINE_OUT 0 +#endif + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *out_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ float *out = reinterpret_cast<__gm__ float *>(out_tensor->buffer.addr) + out_tensor->start_offset; + + int32_t base_cl = static_cast(args[1]); + int32_t spin_iters = static_cast(args[2]); + int32_t block_idx = get_block_idx(args); + int32_t offset = (base_cl + block_idx) * FLOATS_PER_CACHE_LINE; + + volatile int32_t acc = 0; + for (int32_t i = 0; i < spin_iters; i++) { + acc++; // ++ not += i: += i overflows int32 for large spin_iters (UB); volatile keeps the loop + } + (void)acc; + + out[offset] = static_cast(block_idx); + + dcci(&out[offset], SINGLE_CACHE_LINE, CACHELINE_OUT); +} diff --git a/tests/st/a5/tensormap_and_ringbuffer/spmd_sync_start_early_dispatch/kernels/orchestration/spmd_sync_start_early_dispatch_orch.cpp b/tests/st/a5/tensormap_and_ringbuffer/spmd_sync_start_early_dispatch/kernels/orchestration/spmd_sync_start_early_dispatch_orch.cpp new file mode 100644 index 0000000000..32509981a9 --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/spmd_sync_start_early_dispatch/kernels/orchestration/spmd_sync_start_early_dispatch_orch.cpp @@ -0,0 +1,88 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ + +/** + * SPMD Sync-Start Early-Dispatch Orchestration (a5) + * + * Port of a2a3 spmd_sync_start_early_dispatch. Host scalar[0] selects whether + * the producer is flagged allow_early_resolve (1) or not (0), so the same scene + * can compare early vs ready-only swimlanes. + * + * P: AIC core_num=50, base_cl=0, allow_early_resolve = scalar early_on + * C: MIX core_num=24, base_cl=50, require_sync_start=true, dep=[P] + * + * Args layout: [output], scalar: early_on + */ + +#include +#include + +#include "pto_orchestration_api.h" // NOLINT(build/include_subdir) +#include "pto_arg_with_deps.h" // NOLINT(build/include_subdir) + +#define FUNC_SPMD_WRITE_AIC 0 +#define FUNC_SPMD_MIX_AIC 1 +#define FUNC_SPMD_MIX_AIV0 2 +#define FUNC_SPMD_MIX_AIV1 3 + +extern "C" { + +__attribute__((visibility("default"))) PTO2OrchestrationConfig aicpu_orchestration_config(const L2TaskArgs &orch_args) { + (void)orch_args; // NOLINT(readability/casting) + return PTO2OrchestrationConfig{ + .expected_arg_count = 1, + }; +} + +// Match a2a3 ST (1e7): board is fast enough that 2e6 often finishes before the +// scheduler can stage the sync_start consumer on the early path. +static constexpr int64_t PRODUCER_SPIN_ITERS = 10000000; + +static PTO2TaskId submit_producer(const Tensor &out, int16_t core_num, int64_t base_cl, bool early_on) { + L0TaskArgs args; + args.add_inout(out); + args.add_scalar(base_cl); + args.add_scalar(PRODUCER_SPIN_ITERS); + args.launch_spec.set_core_num(core_num); + args.set_allow_early_resolve(early_on); + return rt_submit_aic_task(FUNC_SPMD_WRITE_AIC, args).task_id(); +} + +static void submit_sync_consumer(const Tensor &out, int16_t core_num, int64_t base_cl, PTO2TaskId dep) { + MixedKernels kernels; + kernels.aic_kernel_id = FUNC_SPMD_MIX_AIC; + kernels.aiv0_kernel_id = FUNC_SPMD_MIX_AIV0; + kernels.aiv1_kernel_id = FUNC_SPMD_MIX_AIV1; + L0TaskArgsWithDeps<4> args; + args.add_inout(out); + args.add_scalar(base_cl); + args.launch_spec.set_core_num(core_num); + args.launch_spec.set_require_sync_start(true); + args.add_dep(dep); + rt_submit_task(kernels, args); +} + +__attribute__((visibility("default"))) void aicpu_orchestration_entry(const L2TaskArgs &orch_args) { + const Tensor &ext_output = orch_args.tensor(0).ref(); + const bool early_on = orch_args.scalar(0) != 0; + + rt_scope_begin(PTO2ScopeMode::MANUAL); + PTO2TaskId prod = submit_producer(ext_output, 50, 0, early_on); + submit_sync_consumer(ext_output, 24, 50, prod); + rt_scope_end(); + + LOG_INFO_V9( + "[spmd_sync_start_early_dispatch] early_on=%d wide AIC producer + MIX sync_start consumer submitted", + early_on ? 1 : 0 + ); +} + +} // extern "C" diff --git a/tests/st/a5/tensormap_and_ringbuffer/spmd_sync_start_early_dispatch/test_spmd_sync_start_early_dispatch.py b/tests/st/a5/tensormap_and_ringbuffer/spmd_sync_start_early_dispatch/test_spmd_sync_start_early_dispatch.py new file mode 100644 index 0000000000..205439397d --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/spmd_sync_start_early_dispatch/test_spmd_sync_start_early_dispatch.py @@ -0,0 +1,102 @@ +#!/usr/bin/env python3 +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +"""A wide flagged producer feeds a MIX sync_start early-dispatch consumer (a5). + +Cases EarlyOn / EarlyOff toggle producer ``allow_early_resolve`` via orch scalar +so swimlanes can be compared under identical topology. +""" + +import torch +from simpler.task_interface import ArgDirection as D + +from simpler_setup import Scalar, SceneTestCase, TaskArgsBuilder, Tensor, scene_test + +FLOATS_PER_CACHE_LINE = 16 +PRODUCER_BLOCKS = 50 +SYNC_BLOCKS = 24 +SYNC_BASE_CL = PRODUCER_BLOCKS +TOTAL_CL = SYNC_BASE_CL + SYNC_BLOCKS * 3 + + +@scene_test(level=2, runtime="tensormap_and_ringbuffer") +class TestSpmdSyncStartEarlyDispatch(SceneTestCase): + RTOL = 0 + ATOL = 0 + + CALLABLE = { + "orchestration": { + "source": "kernels/orchestration/spmd_sync_start_early_dispatch_orch.cpp", + "function_name": "aicpu_orchestration_entry", + "signature": [D.INOUT], + }, + "incores": [ + { + "func_id": 0, + "name": "SPMD_WRITE_AIC", + "source": "kernels/aiv/kernel_spmd_write_slow.cpp", + "core_type": "aic", + "signature": [D.INOUT], + }, + { + "func_id": 1, + "name": "SPMD_MIX_AIC", + "source": "../spmd_multiblock_mix/kernels/aic/kernel_spmd_mix.cpp", + "core_type": "aic", + "signature": [D.INOUT], + }, + { + "func_id": 2, + "name": "SPMD_MIX_AIV0", + "source": "../spmd_multiblock_mix/kernels/aiv/kernel_spmd_mix.cpp", + "core_type": "aiv", + "signature": [D.INOUT], + }, + { + "func_id": 3, + "name": "SPMD_MIX_AIV1", + "source": "../spmd_multiblock_mix/kernels/aiv/kernel_spmd_mix.cpp", + "core_type": "aiv", + "signature": [D.INOUT], + }, + ], + } + + CASES = [ + { + "name": "EarlyOn", + "platforms": ["a5sim", "a5"], + "config": {"aicpu_thread_num": 4, "block_dim": 24}, + "params": {"early_on": 1}, + }, + { + "name": "EarlyOff", + "platforms": ["a5sim", "a5"], + "config": {"aicpu_thread_num": 4, "block_dim": 24}, + "params": {"early_on": 0}, + }, + ] + + def generate_args(self, params): + return TaskArgsBuilder( + Tensor("output", torch.zeros(TOTAL_CL * FLOATS_PER_CACHE_LINE, dtype=torch.float32)), + Scalar("early_on", int(params.get("early_on", 1))), + ) + + def compute_golden(self, args, params): + out = args.output + for block_idx in range(PRODUCER_BLOCKS): + out[block_idx * FLOATS_PER_CACHE_LINE] = float(block_idx) + for block_idx in range(SYNC_BLOCKS): + for slot in range(3): + out[(SYNC_BASE_CL + block_idx * 3 + slot) * FLOATS_PER_CACHE_LINE] = float(block_idx) + + +if __name__ == "__main__": + SceneTestCase.run_module(__name__) diff --git a/tests/st/a5/tensormap_and_ringbuffer/spmd_sync_start_mix_spill/kernels/aic/kernel_spmd_mix_slow.cpp b/tests/st/a5/tensormap_and_ringbuffer/spmd_sync_start_mix_spill/kernels/aic/kernel_spmd_mix_slow.cpp new file mode 100644 index 0000000000..0485d14b09 --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/spmd_sync_start_mix_spill/kernels/aic/kernel_spmd_mix_slow.cpp @@ -0,0 +1,68 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ + +/** + * SPMD MIX slow kernel (AIC): writes float(block_idx) at cache line + * (base_cl + block_idx * 3 + 0), then optionally spins so the task holds its + * cluster long enough for a dependent sync_start cohort to pre-stage while it runs. + * + * Args: args[0] = output Tensor* (INOUT), args[1] = base_cl, args[2] = spin_iters. + */ + +#include +#include + +#include "tensor.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] // NOLINT(whitespace/braces) +#endif + +#include "intrinsic.h" + +static constexpr int32_t FLOATS_PER_CACHE_LINE = 16; +static constexpr int32_t SLOTS_PER_BLOCK = 3; + +#ifdef PTO_CPUSTUB_HPP +#define dcci(...) \ + do { \ + } while (0) +#endif +#ifndef SINGLE_CACHE_LINE +#define SINGLE_CACHE_LINE 0 +#endif +#ifndef CACHELINE_OUT +#define CACHELINE_OUT 0 +#endif + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *out_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ float *out = reinterpret_cast<__gm__ float *>(out_tensor->buffer.addr) + out_tensor->start_offset; + + int32_t base_cl = static_cast(args[1]); + int32_t spin_iters = static_cast(args[2]); + int32_t block_idx = get_block_idx(args); + int32_t offset = (base_cl + block_idx * SLOTS_PER_BLOCK + 0) * FLOATS_PER_CACHE_LINE; + + volatile int32_t acc = 0; + for (int32_t i = 0; i < spin_iters; i++) { + acc++; // ++ not += i: += i overflows int32 for large spin_iters (UB); volatile keeps the loop + } + (void)acc; + + out[offset] = static_cast(block_idx); + + dcci(&out[offset], SINGLE_CACHE_LINE, CACHELINE_OUT); +} diff --git a/tests/st/a5/tensormap_and_ringbuffer/spmd_sync_start_mix_spill/kernels/aiv/kernel_spmd_mix_slow.cpp b/tests/st/a5/tensormap_and_ringbuffer/spmd_sync_start_mix_spill/kernels/aiv/kernel_spmd_mix_slow.cpp new file mode 100644 index 0000000000..6d7fcb73dc --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/spmd_sync_start_mix_spill/kernels/aiv/kernel_spmd_mix_slow.cpp @@ -0,0 +1,69 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ + +/** + * SPMD MIX slow kernel (AIV): writes float(block_idx) at cache line + * (base_cl + block_idx * 3 + 1 + sub_block_id), then optionally spins so the task + * holds its cluster while a dependent sync_start cohort pre-stages. + * + * Args: args[0] = output Tensor* (INOUT), args[1] = base_cl, args[2] = spin_iters. + */ + +#include +#include + +#include "tensor.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] // NOLINT(whitespace/braces) +#endif + +#include "intrinsic.h" + +static constexpr int32_t FLOATS_PER_CACHE_LINE = 16; +static constexpr int32_t SLOTS_PER_BLOCK = 3; + +#ifdef PTO_CPUSTUB_HPP +#define dcci(...) \ + do { \ + } while (0) +#endif +#ifndef SINGLE_CACHE_LINE +#define SINGLE_CACHE_LINE 0 +#endif +#ifndef CACHELINE_OUT +#define CACHELINE_OUT 0 +#endif + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *out_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ float *out = reinterpret_cast<__gm__ float *>(out_tensor->buffer.addr) + out_tensor->start_offset; + + int32_t base_cl = static_cast(args[1]); + int32_t spin_iters = static_cast(args[2]); + int32_t block_idx = get_block_idx(args); + int32_t sub_block_id = get_sub_block_id(args); + int32_t offset = (base_cl + block_idx * SLOTS_PER_BLOCK + 1 + sub_block_id) * FLOATS_PER_CACHE_LINE; + + volatile int32_t acc = 0; + for (int32_t i = 0; i < spin_iters; i++) { + acc++; // ++ not += i: += i overflows int32 for large spin_iters (UB); volatile keeps the loop + } + (void)acc; + + out[offset] = static_cast(block_idx); + + dcci(&out[offset], SINGLE_CACHE_LINE, CACHELINE_OUT); +} diff --git a/tests/st/a5/tensormap_and_ringbuffer/spmd_sync_start_mix_spill/kernels/aiv/kernel_spmd_write_slow.cpp b/tests/st/a5/tensormap_and_ringbuffer/spmd_sync_start_mix_spill/kernels/aiv/kernel_spmd_write_slow.cpp new file mode 100644 index 0000000000..9b0cae98ea --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/spmd_sync_start_mix_spill/kernels/aiv/kernel_spmd_write_slow.cpp @@ -0,0 +1,76 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ + +/** + * SPMD Multi-Block Write Kernel with a spin delay (AIV) + * + * Same result as the plain write kernel, but spins for `spin_iters` before the + * store. A slow producer keeps its cores occupied long enough that a dependent + * consumer is processed as an early-dispatch candidate WHILE the producer is still + * running — the only way a trivial-kernel scene exercises the speculative + * gated-dispatch path (a fast producer finishes before the consumer is ever staged). + * + * out[(base_cl + block_idx) * FLOATS_PER_CACHE_LINE] = float(block_idx) + * + * Args: + * args[0] = output Tensor* (INOUT) + * args[1] = scalar: base_cl (starting cache line index for this task) + * args[2] = scalar: spin_iters (0 = no delay) + */ + +#include +#include + +#include "tensor.h" + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] // NOLINT(whitespace/braces) +#endif + +#include "intrinsic.h" + +static constexpr int32_t FLOATS_PER_CACHE_LINE = 16; + +#ifdef PTO_CPUSTUB_HPP +#define dcci(...) \ + do { \ + } while (0) +#endif +#ifndef SINGLE_CACHE_LINE +#define SINGLE_CACHE_LINE 0 +#endif +#ifndef CACHELINE_OUT +#define CACHELINE_OUT 0 +#endif + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *out_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ float *out = reinterpret_cast<__gm__ float *>(out_tensor->buffer.addr) + out_tensor->start_offset; + + int32_t base_cl = static_cast(args[1]); + int32_t spin_iters = static_cast(args[2]); + int32_t block_idx = get_block_idx(args); + int32_t offset = (base_cl + block_idx) * FLOATS_PER_CACHE_LINE; + + volatile int32_t acc = 0; + for (int32_t i = 0; i < spin_iters; i++) { + acc++; // ++ not += i: += i overflows int32 for large spin_iters (UB); volatile keeps the loop + } + (void)acc; + + out[offset] = static_cast(block_idx); + + dcci(&out[offset], SINGLE_CACHE_LINE, CACHELINE_OUT); +} diff --git a/tests/st/a5/tensormap_and_ringbuffer/spmd_sync_start_mix_spill/kernels/orchestration/spmd_sync_start_mix_spill_orch.cpp b/tests/st/a5/tensormap_and_ringbuffer/spmd_sync_start_mix_spill/kernels/orchestration/spmd_sync_start_mix_spill_orch.cpp new file mode 100644 index 0000000000..7e2a213086 --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/spmd_sync_start_mix_spill/kernels/orchestration/spmd_sync_start_mix_spill_orch.cpp @@ -0,0 +1,87 @@ +/* + * Copyright (c) PyPTO Contributors. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + * ----------------------------------------------------------------------------------------------------------- + */ + +/** + * SPMD sync_start MIX pending-spill Orchestration (a5) + * + * Port of a2a3 spmd_sync_start_mix_spill, sized for a5 (36 AIC + 72 AIV): + * P: AIV core_num=72, base_cl=0, allow_early_resolve = scalar early_on (fills all AIV) + * C: MIX core_num=24, base_cl=72, require_sync_start=true, dep=[P] + * → 24 AIC idle→running, 48 AIV busy→pending (per-core spill + rendezvous) + * + * Args layout: [output], scalar: early_on + */ + +#include +#include + +#include "pto_orchestration_api.h" // NOLINT(build/include_subdir) +#include "pto_arg_with_deps.h" // NOLINT(build/include_subdir) + +#define FUNC_SPMD_MIX_AIC 0 +#define FUNC_SPMD_MIX_AIV0 1 +#define FUNC_SPMD_MIX_AIV1 2 +#define FUNC_SPMD_WRITE_AIV 3 + +extern "C" { + +__attribute__((visibility("default"))) PTO2OrchestrationConfig aicpu_orchestration_config(const L2TaskArgs &orch_args) { + (void)orch_args; // NOLINT(readability/casting) + return PTO2OrchestrationConfig{ + .expected_arg_count = 1, + }; +} + +static constexpr int64_t PRODUCER_SPIN_ITERS = 2000000; + +static MixedKernels mix_kernels() { + MixedKernels mk; + mk.aic_kernel_id = FUNC_SPMD_MIX_AIC; + mk.aiv0_kernel_id = FUNC_SPMD_MIX_AIV0; + mk.aiv1_kernel_id = FUNC_SPMD_MIX_AIV1; + return mk; +} + +static PTO2TaskId submit_aiv_producer(const Tensor &out, int16_t core_num, int64_t base_cl, bool early_on) { + L0TaskArgs args; + args.add_inout(out); + args.add_scalar(base_cl); + args.add_scalar(PRODUCER_SPIN_ITERS); + args.launch_spec.set_core_num(core_num); + args.set_allow_early_resolve(early_on); + return rt_submit_aiv_task(FUNC_SPMD_WRITE_AIV, args).task_id(); +} + +static void submit_mix_sync_consumer(const Tensor &out, int16_t core_num, int64_t base_cl, PTO2TaskId dep) { + L0TaskArgsWithDeps<4> args; + args.add_inout(out); + args.add_scalar(base_cl); + args.add_scalar(0); // consumer does not spin + args.launch_spec.set_core_num(core_num); + args.launch_spec.set_require_sync_start(true); + args.add_dep(dep); + rt_submit_task(mix_kernels(), args); +} + +__attribute__((visibility("default"))) void aicpu_orchestration_entry(const L2TaskArgs &orch_args) { + const Tensor &ext_output = orch_args.tensor(0).ref(); + const bool early_on = orch_args.scalar(0) != 0; + + PTO2TaskId prod = submit_aiv_producer(ext_output, 72, 0, early_on); + submit_mix_sync_consumer(ext_output, 24, 72, prod); + + LOG_INFO_V9( + "[spmd_sync_start_mix_spill] early_on=%d AIV producer(72) + MIX sync_start consumer(24) submitted", + early_on ? 1 : 0 + ); +} + +} // extern "C" diff --git a/tests/st/a5/tensormap_and_ringbuffer/spmd_sync_start_mix_spill/test_spmd_sync_start_mix_spill.py b/tests/st/a5/tensormap_and_ringbuffer/spmd_sync_start_mix_spill/test_spmd_sync_start_mix_spill.py new file mode 100644 index 0000000000..c2aa65e618 --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/spmd_sync_start_mix_spill/test_spmd_sync_start_mix_spill.py @@ -0,0 +1,105 @@ +#!/usr/bin/env python3 +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +"""sync_start MIX per-core pending-spill on a5 (36 AIC + 72 AIV). + +A flagged AIV producer occupies all 72 AIV cores and spins; the require_sync_start +MIX consumer (24 clusters) pre-stages with AIC on idle running slots and AIVs on +busy pending slots. EarlyOn / EarlyOff toggle producer ``allow_early_resolve``. +""" + +import torch +from simpler.task_interface import ArgDirection as D + +from simpler_setup import Scalar, SceneTestCase, TaskArgsBuilder, Tensor, scene_test + +FLOATS_PER_CACHE_LINE = 16 +SLOTS_PER_BLOCK = 3 +PRODUCER_BLOCKS = 72 # fill all a5 AIV cores +PRODUCER_BASE_CL = 0 +CONSUMER_BLOCKS = 24 +CONSUMER_BASE_CL = PRODUCER_BLOCKS +TOTAL_CL = PRODUCER_BLOCKS + CONSUMER_BLOCKS * SLOTS_PER_BLOCK # 144 + + +@scene_test(level=2, runtime="tensormap_and_ringbuffer") +class TestSpmdSyncStartMixSpill(SceneTestCase): + RTOL = 0 + ATOL = 0 + + CALLABLE = { + "orchestration": { + "source": "kernels/orchestration/spmd_sync_start_mix_spill_orch.cpp", + "function_name": "aicpu_orchestration_entry", + "signature": [D.INOUT], + }, + "incores": [ + { + "func_id": 0, + "name": "SPMD_MIX_AIC", + "source": "kernels/aic/kernel_spmd_mix_slow.cpp", + "core_type": "aic", + "signature": [D.INOUT], + }, + { + "func_id": 1, + "name": "SPMD_MIX_AIV0", + "source": "kernels/aiv/kernel_spmd_mix_slow.cpp", + "core_type": "aiv", + "signature": [D.INOUT], + }, + { + "func_id": 2, + "name": "SPMD_MIX_AIV1", + "source": "kernels/aiv/kernel_spmd_mix_slow.cpp", + "core_type": "aiv", + "signature": [D.INOUT], + }, + { + "func_id": 3, + "name": "SPMD_WRITE_AIV", + "source": "kernels/aiv/kernel_spmd_write_slow.cpp", + "core_type": "aiv", + "signature": [D.INOUT], + }, + ], + } + + CASES = [ + { + "name": "EarlyOn", + "platforms": ["a5sim", "a5"], + "config": {"aicpu_thread_num": 4, "block_dim": 24}, + "params": {"early_on": 1}, + }, + { + "name": "EarlyOff", + "platforms": ["a5sim", "a5"], + "config": {"aicpu_thread_num": 4, "block_dim": 24}, + "params": {"early_on": 0}, + }, + ] + + def generate_args(self, params): + return TaskArgsBuilder( + Tensor("output", torch.zeros(TOTAL_CL * FLOATS_PER_CACHE_LINE, dtype=torch.float32)), + Scalar("early_on", int(params.get("early_on", 1))), + ) + + def compute_golden(self, args, params): + out = args.output + for block_idx in range(PRODUCER_BLOCKS): + out[(PRODUCER_BASE_CL + block_idx) * FLOATS_PER_CACHE_LINE] = float(block_idx) + for block_idx in range(CONSUMER_BLOCKS): + for slot in range(SLOTS_PER_BLOCK): + out[(CONSUMER_BASE_CL + block_idx * SLOTS_PER_BLOCK + slot) * FLOATS_PER_CACHE_LINE] = float(block_idx) + + +if __name__ == "__main__": + SceneTestCase.run_module(__name__) diff --git a/tests/ut/cpp/a5/test_scheduler_state.cpp b/tests/ut/cpp/a5/test_scheduler_state.cpp index 8d808799aa..6db6def8ca 100644 --- a/tests/ut/cpp/a5/test_scheduler_state.cpp +++ b/tests/ut/cpp/a5/test_scheduler_state.cpp @@ -30,6 +30,13 @@ class SchedulerStateTest : public ::testing::Test { DeviceArena sm_arena; DeviceArena sched_arena; + // Each init_slot()'d slot gets a distinct zeroed payload from this pool, + // mirroring orch::prepare_task's bind_buffers: every production slot has a + // payload, and the scheduler's release/propagate paths dereference it. + static constexpr int kSlotPayloadPoolSize = 16; + PTO2TaskPayload slot_payload_pool_[kSlotPayloadPoolSize]; + int slot_payload_pool_idx_ = 0; + void SetUp() override { sm_handle = PTO2SharedMemoryHandle::create_and_init_default(sm_arena); ASSERT_NE(sm_handle, nullptr); @@ -61,6 +68,9 @@ class SchedulerStateTest : public ::testing::Test { slot.completed_subtasks.store(0); slot.total_required_subtasks = 1; slot.logical_block_num = 1; + PTO2TaskPayload &slot_pl = slot_payload_pool_[slot_payload_pool_idx_++ % kSlotPayloadPoolSize]; + memset(&slot_pl, 0, sizeof(slot_pl)); + slot.payload = &slot_pl; } }; diff --git a/tests/ut/cpp/a5/test_task_state.cpp b/tests/ut/cpp/a5/test_task_state.cpp index c0773ec222..916d9144f1 100644 --- a/tests/ut/cpp/a5/test_task_state.cpp +++ b/tests/ut/cpp/a5/test_task_state.cpp @@ -38,6 +38,13 @@ class TaskStateTest : public ::testing::Test { DeviceArena sm_arena; DeviceArena sched_arena; + // Each init_slot()'d slot gets a distinct zeroed payload from this pool, + // mirroring orch::prepare_task's bind_buffers: every production slot has a + // payload, and the scheduler's release/propagate paths dereference it. + static constexpr int kSlotPayloadPoolSize = 16; + PTO2TaskPayload slot_payload_pool_[kSlotPayloadPoolSize]; + int slot_payload_pool_idx_ = 0; + void SetUp() override { sm_handle = PTO2SharedMemoryHandle::create_and_init_default(sm_arena); ASSERT_NE(sm_handle, nullptr); @@ -67,6 +74,9 @@ class TaskStateTest : public ::testing::Test { slot.completed_subtasks.store(0); slot.total_required_subtasks = 1; slot.logical_block_num = 1; + PTO2TaskPayload &slot_pl = slot_payload_pool_[slot_payload_pool_idx_++ % kSlotPayloadPoolSize]; + memset(&slot_pl, 0, sizeof(slot_pl)); + slot.payload = &slot_pl; } }; diff --git a/tests/ut/cpp/a5/test_wiring.cpp b/tests/ut/cpp/a5/test_wiring.cpp index 62723ac39f..24b25a663c 100644 --- a/tests/ut/cpp/a5/test_wiring.cpp +++ b/tests/ut/cpp/a5/test_wiring.cpp @@ -45,6 +45,14 @@ class WiringTest : public ::testing::Test { DeviceArena sm_arena; DeviceArena sched_arena; + // Each init_slot()'d slot gets a distinct zeroed payload from this pool, + // mirroring orch::prepare_task's bind_buffers: every production slot has a + // payload, and the scheduler's release/propagate paths dereference it. + static constexpr int kSlotPayloadPoolSize = 16; + PTO2TaskPayload slot_payload_pool_[kSlotPayloadPoolSize]; + PTO2TaskDescriptor slot_task_pool_[kSlotPayloadPoolSize]; + int slot_payload_pool_idx_ = 0; + void SetUp() override { sm_handle = PTO2SharedMemoryHandle::create_and_init_default(sm_arena); ASSERT_NE(sm_handle, nullptr); @@ -79,6 +87,12 @@ class WiringTest : public ::testing::Test { slot.total_required_subtasks = 1; slot.logical_block_num = 1; slot.dep_pool_mark = 0; + PTO2TaskPayload &slot_pl = slot_payload_pool_[slot_payload_pool_idx_++ % kSlotPayloadPoolSize]; + memset(&slot_pl, 0, sizeof(slot_pl)); + slot.payload = &slot_pl; + PTO2TaskDescriptor &slot_task = slot_task_pool_[(slot_payload_pool_idx_ - 1) % kSlotPayloadPoolSize]; + memset(&slot_task, 0, sizeof(slot_task)); + slot.task = &slot_task; } void publish_no_fanin(PTO2TaskSlotState &slot) {