From b66a556a1cf1e0e0293dddfa71ea0ed7031f24f8 Mon Sep 17 00:00:00 2001 From: yanghaoran29 Date: Tue, 21 Jul 2026 20:03:51 +0800 Subject: [PATCH 1/3] feat(a5/tmr): port allow_early_resolve / early-dispatch from a2a3 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Bring a5 tensormap_and_ringbuffer in line with a2a3 early-dispatch (#989/#1079/#1285/#1288/#1297/#1304/#1326/#1329/#1336/#1405/#1340): prepare/publish, src_payload gate + DMB high32 doorbell (sim uses __atomic_load_n ACQUIRE), direct-only eligibility, spare-slot Phase 4b, sync_start early rendezvous, and dep_gen early_dispatch truth. Document the a5 model in RUNTIME_LOGIC §8.6. Add a5 STs: plain early_dispatch plus sync_start early_dispatch and mix_spill ports (EarlyOn/Off) for board validation. --- src/a5/platform/onboard/aicore/inner_kernel.h | 14 + src/a5/platform/sim/aicore/inner_kernel.h | 13 + .../aicore/aicore_executor.cpp | 46 ++ .../docs/RUNTIME_LOGIC.md | 102 +++- .../orchestration/pto_arg_with_deps.h | 6 +- .../runtime/pto2_dispatch_payload.h | 58 ++- .../runtime/pto_orchestrator.cpp | 49 +- .../runtime/pto_runtime2_types.h | 148 +++++- .../runtime/pto_types.h | 10 + .../runtime/scheduler/pto_scheduler.h | 364 +++++++++++++- .../runtime/scheduler/scheduler_cold_path.cpp | 3 + .../scheduler/scheduler_completion.cpp | 274 +++++++---- .../runtime/scheduler/scheduler_context.h | 78 ++- .../runtime/scheduler/scheduler_dispatch.cpp | 457 ++++++++++++++---- .../runtime/scheduler/scheduler_types.h | 48 +- .../runtime/shared/pto_runtime2_init.cpp | 32 ++ .../spmd_early_dispatch_orch.cpp | 87 ++++ .../test_spmd_early_dispatch.py | 83 ++++ .../kernels/aiv/kernel_spmd_write_slow.cpp | 76 +++ .../spmd_sync_start_early_dispatch_orch.cpp | 88 ++++ .../test_spmd_sync_start_early_dispatch.py | 102 ++++ .../kernels/aic/kernel_spmd_mix_slow.cpp | 68 +++ .../kernels/aiv/kernel_spmd_mix_slow.cpp | 69 +++ .../kernels/aiv/kernel_spmd_write_slow.cpp | 76 +++ .../spmd_sync_start_mix_spill_orch.cpp | 87 ++++ .../test_spmd_sync_start_mix_spill.py | 105 ++++ 26 files changed, 2285 insertions(+), 258 deletions(-) create mode 100644 tests/st/a5/tensormap_and_ringbuffer/spmd_early_dispatch/kernels/orchestration/spmd_early_dispatch_orch.cpp create mode 100644 tests/st/a5/tensormap_and_ringbuffer/spmd_early_dispatch/test_spmd_early_dispatch.py create mode 100644 tests/st/a5/tensormap_and_ringbuffer/spmd_sync_start_early_dispatch/kernels/aiv/kernel_spmd_write_slow.cpp create mode 100644 tests/st/a5/tensormap_and_ringbuffer/spmd_sync_start_early_dispatch/kernels/orchestration/spmd_sync_start_early_dispatch_orch.cpp create mode 100644 tests/st/a5/tensormap_and_ringbuffer/spmd_sync_start_early_dispatch/test_spmd_sync_start_early_dispatch.py create mode 100644 tests/st/a5/tensormap_and_ringbuffer/spmd_sync_start_mix_spill/kernels/aic/kernel_spmd_mix_slow.cpp create mode 100644 tests/st/a5/tensormap_and_ringbuffer/spmd_sync_start_mix_spill/kernels/aiv/kernel_spmd_mix_slow.cpp create mode 100644 tests/st/a5/tensormap_and_ringbuffer/spmd_sync_start_mix_spill/kernels/aiv/kernel_spmd_write_slow.cpp create mode 100644 tests/st/a5/tensormap_and_ringbuffer/spmd_sync_start_mix_spill/kernels/orchestration/spmd_sync_start_mix_spill_orch.cpp create mode 100644 tests/st/a5/tensormap_and_ringbuffer/spmd_sync_start_mix_spill/test_spmd_sync_start_mix_spill.py 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_early_dispatch/kernels/orchestration/spmd_early_dispatch_orch.cpp b/tests/st/a5/tensormap_and_ringbuffer/spmd_early_dispatch/kernels/orchestration/spmd_early_dispatch_orch.cpp new file mode 100644 index 0000000000..52e3f05003 --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/spmd_early_dispatch/kernels/orchestration/spmd_early_dispatch_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 Early-Dispatch Orchestration (a5) — plain early_resolve, NO sync_start + * + * Exercises the shape-queued early path (`early_dispatch_queues[]`), not the + * sync_start cohort lane (`early_sync_start_queue`). + * + * Topology is sized so the producer leaves spare AIC cores: while P spins on + * 8 AICs, C can gated-stage onto the remaining idle AICs. + * + * P: AIC core_num=8, base_cl=0, allow_early_resolve = scalar early_on, spin + * C: AIC core_num=8, base_cl=8, dep=[P], require_sync_start=false + * + * 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 + +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 #1079-era spin: long enough that C can stage while P is still on-core. +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_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); // no spin + args.launch_spec.set_core_num(core_num); + // deliberately NOT require_sync_start — plain early_dispatch_queues path + args.add_dep(dep); + rt_submit_aic_task(FUNC_SPMD_WRITE_AIC, args); +} + +__attribute__((visibility("default"))) void aicpu_orchestration_entry(const L2TaskArgs &orch_args) { + const Tensor &ext_output = orch_args.tensor(0).ref(); + // Host Scalar("early_on") is orch scalar(0). Also stamp a sentinel into the + // output so board runs can prove the scalar reached the orch (out[1] = early_on). + const bool early_on = orch_args.scalar(0) != 0; + { + // out layout: CL0..7 producer, CL8..15 consumer; float index 1 is unused by kernels. + float *out = reinterpret_cast(ext_output.buffer.addr) + ext_output.start_offset; + out[1] = early_on ? 1.0f : 0.0f; + } + + rt_scope_begin(PTO2ScopeMode::MANUAL); + PTO2TaskId prod = submit_producer(ext_output, 8, 0, early_on); + submit_consumer(ext_output, 8, 8, prod); + rt_scope_end(); + + LOG_INFO_V9("[spmd_early_dispatch] early_on=%d AIC producer + AIC consumer (no sync_start)", early_on ? 1 : 0); +} + +} // extern "C" diff --git a/tests/st/a5/tensormap_and_ringbuffer/spmd_early_dispatch/test_spmd_early_dispatch.py b/tests/st/a5/tensormap_and_ringbuffer/spmd_early_dispatch/test_spmd_early_dispatch.py new file mode 100644 index 0000000000..2cbb473873 --- /dev/null +++ b/tests/st/a5/tensormap_and_ringbuffer/spmd_early_dispatch/test_spmd_early_dispatch.py @@ -0,0 +1,83 @@ +#!/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. +# ----------------------------------------------------------------------------------------------------------- +"""Plain allow_early_resolve ST (a5) — NO require_sync_start. + +Producer leaves spare AIC cores so the consumer can stage via +``early_dispatch_queues`` while the producer is still spinning. EarlyOn / +EarlyOff toggle the producer flag for swimlane comparison. +""" + +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 = 8 +CONSUMER_BLOCKS = 8 +CONSUMER_BASE_CL = PRODUCER_BLOCKS +TOTAL_CL = CONSUMER_BASE_CL + CONSUMER_BLOCKS + + +@scene_test(level=2, runtime="tensormap_and_ringbuffer") +class TestSpmdEarlyDispatch(SceneTestCase): + RTOL = 0 + ATOL = 0 + + CALLABLE = { + "orchestration": { + "source": "kernels/orchestration/spmd_early_dispatch_orch.cpp", + "function_name": "aicpu_orchestration_entry", + "signature": [D.INOUT], + }, + "incores": [ + { + "func_id": 0, + "name": "SPMD_WRITE_AIC", + "source": "../spmd_sync_start_early_dispatch/kernels/aiv/kernel_spmd_write_slow.cpp", + "core_type": "aic", + "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 + # orch stamps early_on into out[1] as a host-side probe (kernels leave it alone). + out[1] = float(int(params.get("early_on", 1))) + for block_idx in range(PRODUCER_BLOCKS): + out[block_idx * FLOATS_PER_CACHE_LINE] = float(block_idx) + for block_idx in range(CONSUMER_BLOCKS): + out[(CONSUMER_BASE_CL + block_idx) * 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_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__) From b8900bdbd395e3f21f01e24cd5cee5d379efe4eb Mon Sep 17 00:00:00 2001 From: yanghaoran29 Date: Tue, 21 Jul 2026 20:31:46 +0800 Subject: [PATCH 2/3] test(a5): drop local early-dispatch ST ports from this PR Remove the a5-only early_dispatch / sync_start early ST scenes added for board experiments; keep the runtime port focused for CI. --- _env_dump.sh | 2 + _migrate_patches/early_phase0b.diff | 552 ++++++++++ _migrate_patches/pto_scheduler_partial.h | 993 ++++++++++++++++++ _run_early_st_board.sh | 22 + _run_st_dev2.sh | 9 + .../acc_c2v_strip_validshape/README.md | 62 ++ .../kernels/aic/kernel_matmul_tstore.cpp | 69 ++ .../kernels/mix/kernel_acc_c2v_compare.cpp | 194 ++++ .../kernels/mix/kernel_acc_c2v_full.cpp | 192 ++++ .../kernels/mix/kernel_acc_c2v_strip.cpp | 192 ++++ .../orchestration/acc_c2v_compare_orch.cpp | 53 + .../orchestration/acc_c2v_strip_orch.cpp | 71 ++ .../test_acc_c2v_full_control.py | 80 ++ .../test_acc_c2v_strip_validshape.py | 99 ++ .../test_acc_c2v_strip_vs_full.py | 124 +++ .../test_acc_matmul_tstore.py | 53 + .../spmd_early_dispatch_orch.cpp | 87 -- .../test_spmd_early_dispatch.py | 83 -- tests/ut/cpp/a5/test_scheduler_state.cpp | 10 + tests/ut/cpp/a5/test_task_state.cpp | 10 + tests/ut/cpp/a5/test_wiring.cpp | 14 + 21 files changed, 2801 insertions(+), 170 deletions(-) create mode 100755 _env_dump.sh create mode 100644 _migrate_patches/early_phase0b.diff create mode 100644 _migrate_patches/pto_scheduler_partial.h create mode 100755 _run_early_st_board.sh create mode 100755 _run_st_dev2.sh create mode 100644 examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/README.md create mode 100644 examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/kernels/aic/kernel_matmul_tstore.cpp create mode 100644 examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/kernels/mix/kernel_acc_c2v_compare.cpp create mode 100644 examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/kernels/mix/kernel_acc_c2v_full.cpp create mode 100644 examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/kernels/mix/kernel_acc_c2v_strip.cpp create mode 100644 examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/kernels/orchestration/acc_c2v_compare_orch.cpp create mode 100644 examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/kernels/orchestration/acc_c2v_strip_orch.cpp create mode 100644 examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/test_acc_c2v_full_control.py create mode 100644 examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/test_acc_c2v_strip_validshape.py create mode 100644 examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/test_acc_c2v_strip_vs_full.py create mode 100644 examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/test_acc_matmul_tstore.py delete mode 100644 tests/st/a5/tensormap_and_ringbuffer/spmd_early_dispatch/kernels/orchestration/spmd_early_dispatch_orch.cpp delete mode 100644 tests/st/a5/tensormap_and_ringbuffer/spmd_early_dispatch/test_spmd_early_dispatch.py diff --git a/_env_dump.sh b/_env_dump.sh new file mode 100755 index 0000000000..6519a196f2 --- /dev/null +++ b/_env_dump.sh @@ -0,0 +1,2 @@ +#!/bin/bash +env | sort diff --git a/_migrate_patches/early_phase0b.diff b/_migrate_patches/early_phase0b.diff new file mode 100644 index 0000000000..e1bd555f26 --- /dev/null +++ b/_migrate_patches/early_phase0b.diff @@ -0,0 +1,552 @@ +diff --git a/src/a5/platform/onboard/aicore/inner_kernel.h b/src/a5/platform/onboard/aicore/inner_kernel.h +index 2f429603..96d8c3d6 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 58dd4488..8b997b28 100644 +--- a/src/a5/platform/sim/aicore/inner_kernel.h ++++ b/src/a5/platform/sim/aicore/inner_kernel.h +@@ -176,6 +176,17 @@ 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. ++ */ ++inline uint32_t read_dmb_high32() { ++ uint32_t offset = reg_offset(RegId::DATA_MAIN_BASE); ++ uint32_t hi = *reinterpret_cast(sparse_reg_ptr(sim_get_reg_base(), offset + 4)); ++ OUT_OF_ORDER_LOAD_BARRIER(); ++ return hi; ++} ++ + /** + * Write to an AICore register in simulated register memory + * +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 20144d3a..00c36211 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 @@ public: + 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/pto_types.h b/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_types.h +index bd0cd66b..27829ffb 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/scheduler_completion.cpp b/src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_completion.cpp +index e1586227..4a5c13f9 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" +@@ -413,6 +415,7 @@ int32_t SchedulerContext::count_global_available(PTO2ResourceShape shape, uint8_ + // 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) { ++ (void)runtime; + 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); +@@ -425,9 +428,27 @@ void SchedulerContext::drain_worker_dispatch(Runtime *runtime, int32_t block_num + 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++; ++ int32_t remaining = slot_state->logical_block_num - slot_state->next_block_idx; ++ int32_t claim = std::min(valid.count(), remaining); ++ int32_t start = slot_state->next_block_idx; ++ slot_state->next_block_idx += claim; ++ PublishHandle handles[CoreTracker::MAX_CLUSTERS * 3]; ++ int handle_count = 0; ++ for (int32_t b = 0; b < claim; b++) { ++ auto core_offset = valid.pop_first(); ++ handle_count += prepare_block_for_dispatch( ++ t, core_offset, *slot_state, shape, false, start + b, &handles[handle_count] ++ ); ++ } ++ wmb(); ++ uint64_t dispatch_ts = 0; ++#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, t); + } + } + +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 c60340ef..9aa25a0d 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" +@@ -250,19 +252,62 @@ private: + const AsyncCtx &async_ctx, int32_t block_idx + ); + +- 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 + ); + +- 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); ++ } ++ // PTO2_DMB64_DISPATCH_PROBE: seed both halves of DATA_MAIN_BASE so AICore ++ // read_dmb_high32()/MOV can observe the doorbell half on silicon. Matches ++ // a2a3 ring_one_doorbell encoding `(tok<<32)|tok`. Remove after Phase-0 ++ // onboard probe (tests/st/a5/.../dmb_high32_probe) lands a verdict. ++#if defined(PTO2_DMB64_DISPATCH_PROBE) ++ { ++ volatile uint64_t *dmb = reinterpret_cast( ++ get_reg_ptr(h.reg_addr, RegId::DATA_MAIN_BASE) ++ ); ++ uint64_t tk = static_cast(h.reg_task_id); ++ *dmb = (tk << 32) | tk; ++ } ++#else ++ write_reg(h.reg_addr, RegId::DATA_MAIN_BASE, static_cast(h.reg_task_id)); ++#endif ++ } ++ ++ // 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 + ); + + void dispatch_shape( +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 8007169b..30c00e4e 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 +@@ -125,14 +125,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 + ) { + 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,6 +145,10 @@ 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; +@@ -161,6 +165,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 +187,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 + ) { + #if SIMPLER_DFX + if (is_dump_args_enabled()) { +@@ -258,26 +217,58 @@ 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 ++ ); ++ } ++ 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 ++ ); ++ } ++ 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 ++ ); ++ } ++#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); ++#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); + #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 +285,64 @@ 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; + 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 +370,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,21 +383,19 @@ 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 ++ // 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(); +@@ -365,13 +412,24 @@ 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(); + #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; + diff --git a/_migrate_patches/pto_scheduler_partial.h b/_migrate_patches/pto_scheduler_partial.h new file mode 100644 index 0000000000..eca3496924 --- /dev/null +++ b/_migrate_patches/pto_scheduler_partial.h @@ -0,0 +1,993 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +/** + * PTO Runtime2 - Scheduler Interface + * + * The Scheduler is responsible for: + * 1. Maintaining per-resource-shape ready queues + * 2. Tracking task state (PENDING -> COMPLETED -> CONSUMED) + * 3. Managing fanin/fanout refcounts for dependency resolution + * 4. Advancing last_task_alive for heap reclamation + * 5. Two-stage mixed-task completion (subtask done bits → mixed-task complete) + * + * The Scheduler runs on Device AI_CPU and processes: + * - Task state transitions based on fanin_refcount + * - Buffer lifecycle based on fanout_refcount + * - Ring pointer advancement for flow control + * + * Based on: docs/RUNTIME_LOGIC.md + */ + +#pragma once + +#include + +#include "common/core_type.h" +#include "utils/device_arena.h" +#include "pto_async_wait.h" +#include "pto_ring_buffer.h" +#include "pto_runtime2_types.h" +#include "pto_shared_memory.h" + +#if SIMPLER_SCHED_PROFILING +#include "aicpu/device_time.h" +#define PTO2_SCHED_CYCLE_START() uint64_t _st0 = get_sys_cnt_aicpu(), _st1 +#define PTO2_SCHED_CYCLE_LAP(acc) \ + do { \ + _st1 = get_sys_cnt_aicpu(); \ + acc += (_st1 - _st0); \ + _st0 = _st1; \ + } while (0) +#endif + +// ============================================================================= +// Ready Queue (Lock-free bounded MPMC — Vyukov design) +// ============================================================================= + +/** + * Per-slot entry: sequence counter for ABA safety + task payload + */ +struct PTO2ReadyQueueSlot { + std::atomic sequence; + PTO2TaskSlotState *slot_state; + uint64_t task_id_snapshot; // generation tag for early-dispatch queue entries +}; + +/** + * Lock-free bounded MPMC queue (Dmitry Vyukov design) + * + * Key properties: + * - enqueue_pos and dequeue_pos on separate cache lines (no false sharing) + * - Per-slot sequence counter prevents ABA problem + * - Empty queue pop returns immediately (single atomic load, no lock) + * - CAS contention is split: producers only touch enqueue_pos, + * consumers only touch dequeue_pos + */ +struct alignas(64) PTO2ReadyQueue { + PTO2ReadyQueueSlot *slots; + uint64_t capacity; + uint64_t mask; // capacity - 1 + char _pad0[64 - 24]; // Pad to own cache line + + std::atomic enqueue_pos; + char _pad1[64 - sizeof(std::atomic)]; // Own cache line + + std::atomic dequeue_pos; + char _pad2[64 - sizeof(std::atomic)]; // Own cache line + + uint64_t size() { + uint64_t e = enqueue_pos.load(std::memory_order_relaxed); + uint64_t d = dequeue_pos.load(std::memory_order_relaxed); + return (e >= d) ? (e - d) : 0; + } + + void reset_for_reuse() {} + + bool push(PTO2TaskSlotState *slot_state) { return push_tagged(slot_state, 0); } + + bool push_tagged(PTO2TaskSlotState *slot_state, uint64_t task_id_snapshot) { + uint64_t pos; + PTO2ReadyQueueSlot *slot; + while (true) { + pos = enqueue_pos.load(std::memory_order_relaxed); + slot = &slots[pos & mask]; + int64_t seq = slot->sequence.load(std::memory_order_acquire); + int64_t diff = seq - static_cast(pos); + if (diff == 0) { + if (enqueue_pos.compare_exchange_weak( + pos, pos + 1, std::memory_order_relaxed, std::memory_order_relaxed + )) { + break; + } + } else if (diff < 0) { + return false; // Queue full + } + } + + slot->slot_state = slot_state; + slot->task_id_snapshot = task_id_snapshot; + slot->sequence.store(static_cast(pos + 1), std::memory_order_release); + return true; + } + + // Batch push: reserve count slots with a single CAS after confirming + // every target slot is available under the usual Vyukov sequence check. + void push_batch(PTO2TaskSlotState **items, int count) { push_batch_tagged(items, nullptr, count); } + + void push_batch_tagged(PTO2TaskSlotState **items, const uint64_t *task_id_snapshots, int count) { + if (count == 0) return; + + uint64_t pos; + while (true) { + pos = enqueue_pos.load(std::memory_order_relaxed); + bool ready = true; + for (int i = 0; i < count; i++) { + PTO2ReadyQueueSlot *slot = &slots[(pos + i) & mask]; + int64_t seq = slot->sequence.load(std::memory_order_acquire); + int64_t diff = seq - static_cast(pos + i); + if (diff != 0) { + ready = false; + break; + } + } + if (!ready) { + continue; + } + if (enqueue_pos.compare_exchange_weak( + pos, pos + count, std::memory_order_relaxed, std::memory_order_relaxed + )) { + break; + } + } + + for (int i = 0; i < count; i++) { + PTO2ReadyQueueSlot *slot = &slots[(pos + i) & mask]; + slot->slot_state = items[i]; + slot->task_id_snapshot = task_id_snapshots == nullptr ? 0 : task_id_snapshots[i]; + slot->sequence.store(static_cast(pos + i + 1), std::memory_order_release); + } + } + +#if SIMPLER_ORCH_PROFILING || SIMPLER_SCHED_PROFILING + bool push(PTO2TaskSlotState *slot_state, uint64_t &atomic_count, uint64_t &wait_cycle) { + uint64_t pos; + PTO2ReadyQueueSlot *slot; + uint64_t t0 = get_sys_cnt_aicpu(); + bool contended = false; + uint32_t atomic_ops = 0; + while (true) { + pos = enqueue_pos.load(std::memory_order_relaxed); + slot = &slots[pos & mask]; + int64_t seq = slot->sequence.load(std::memory_order_acquire); + int64_t diff = seq - static_cast(pos); + atomic_ops += 2; // enqueue_pos.load + sequence.load + if (diff == 0) { + if (enqueue_pos.compare_exchange_weak( + pos, pos + 1, std::memory_order_relaxed, std::memory_order_relaxed + )) { + atomic_ops++; // successful CAS + break; + } + contended = true; + atomic_ops++; // failed CAS + } else if (diff < 0) { + return false; // Queue full + } else { + contended = true; // diff > 0: slot not yet released, spin + } + } + atomic_ops++; // final sequence.store + atomic_count += atomic_ops; + if (contended) { + wait_cycle += (get_sys_cnt_aicpu() - t0); + } + + slot->slot_state = slot_state; + slot->task_id_snapshot = 0; + slot->sequence.store(static_cast(pos + 1), std::memory_order_release); + return true; + } +#endif + + PTO2TaskSlotState *pop() { return pop_tagged(nullptr); } + + PTO2TaskSlotState *pop_tagged(uint64_t *task_id_snapshot) { + // Fast-path: skip slot load when queue is clearly empty + uint64_t d = dequeue_pos.load(std::memory_order_relaxed); + uint64_t e = enqueue_pos.load(std::memory_order_relaxed); + if (d >= e) { + return nullptr; + } + + uint64_t pos; + PTO2ReadyQueueSlot *slot; + while (true) { + pos = dequeue_pos.load(std::memory_order_relaxed); + slot = &slots[pos & mask]; + int64_t seq = slot->sequence.load(std::memory_order_acquire); + int64_t diff = seq - static_cast(pos + 1); + if (diff == 0) { + if (dequeue_pos.compare_exchange_weak( + pos, pos + 1, std::memory_order_relaxed, std::memory_order_relaxed + )) + break; + } else if (diff < 0) { + return nullptr; // Queue empty + } + } + + PTO2TaskSlotState *result = slot->slot_state; + slot->sequence.store(static_cast(pos + mask + 1), std::memory_order_release); + return result; + } + +#if SIMPLER_SCHED_PROFILING + PTO2TaskSlotState *pop(uint64_t &atomic_count, uint64_t &wait_cycle) { + // Fast-path: skip slot load when queue is clearly empty + uint64_t d = dequeue_pos.load(std::memory_order_relaxed); + uint64_t e = enqueue_pos.load(std::memory_order_relaxed); + atomic_count += 2; // dequeue_pos.load + enqueue_pos.load + if (d >= e) { + return nullptr; + } + + uint64_t pos; + PTO2ReadyQueueSlot *slot; + uint64_t t0 = get_sys_cnt_aicpu(); + bool contended = false; + uint32_t atomic_ops = 0; + while (true) { + pos = dequeue_pos.load(std::memory_order_relaxed); + slot = &slots[pos & mask]; + int64_t seq = slot->sequence.load(std::memory_order_acquire); + int64_t diff = seq - static_cast(pos + 1); + atomic_ops += 2; // dequeue_pos.load + sequence.load + if (diff == 0) { + if (dequeue_pos.compare_exchange_weak( + pos, pos + 1, std::memory_order_relaxed, std::memory_order_relaxed + )) { + atomic_ops++; // successful CAS + break; + } + contended = true; + atomic_ops++; // failed CAS + } else if (diff < 0) { + atomic_count += atomic_ops; + return nullptr; // Queue empty + } else { + contended = true; + } + } + atomic_ops++; // final sequence.store + atomic_count += atomic_ops; + if (contended) { + wait_cycle += (get_sys_cnt_aicpu() - t0); + } + + PTO2TaskSlotState *result = slot->slot_state; + slot->sequence.store(static_cast(pos + mask + 1), std::memory_order_release); + return result; + } +#endif + + // Batch pop: reserve a contiguous run of ready slots with a single CAS. + // Returns actual number of items popped (may be less than max_count). + int pop_batch(PTO2TaskSlotState **out, int max_count) { + uint64_t pos; + int count; + while (true) { + pos = dequeue_pos.load(std::memory_order_relaxed); + count = 0; + while (count < max_count) { + PTO2ReadyQueueSlot *slot = &slots[(pos + count) & mask]; + int64_t seq = slot->sequence.load(std::memory_order_acquire); + int64_t diff = seq - static_cast(pos + count + 1); + if (diff == 0) { + count++; + continue; + } + if (diff < 0) { + break; + } + count = -1; + break; + } + if (count == 0) return 0; + if (count < 0) continue; + if (dequeue_pos.compare_exchange_weak( + pos, pos + count, std::memory_order_relaxed, std::memory_order_relaxed + )) { + break; + } + } + + for (int i = 0; i < count; i++) { + PTO2ReadyQueueSlot *slot = &slots[(pos + i) & mask]; + out[i] = slot->slot_state; + slot->sequence.store(static_cast(pos + i + mask + 1), std::memory_order_release); + } + return count; + } + +#if SIMPLER_SCHED_PROFILING + int pop_batch(PTO2TaskSlotState **out, int max_count, uint64_t &atomic_count, uint64_t &wait_cycle) { + uint64_t pos; + int count; + uint64_t t0 = get_sys_cnt_aicpu(); + bool contended = false; + uint32_t atomic_ops = 0; + while (true) { + pos = dequeue_pos.load(std::memory_order_relaxed); + atomic_ops++; // dequeue_pos.load + count = 0; + while (count < max_count) { + PTO2ReadyQueueSlot *slot = &slots[(pos + count) & mask]; + int64_t seq = slot->sequence.load(std::memory_order_acquire); + int64_t diff = seq - static_cast(pos + count + 1); + atomic_ops++; // sequence.load + if (diff == 0) { + count++; + continue; + } + if (diff < 0) { + break; + } + contended = true; + count = -1; + break; + } + if (count == 0) { + atomic_count += atomic_ops; + return 0; + } + if (count < 0) { + continue; + } + if (dequeue_pos.compare_exchange_weak( + pos, pos + count, std::memory_order_relaxed, std::memory_order_relaxed + )) { + atomic_ops++; // successful CAS + break; + } + contended = true; + atomic_ops++; // failed CAS + } + + for (int i = 0; i < count; i++) { + PTO2ReadyQueueSlot *slot = &slots[(pos + i) & mask]; + out[i] = slot->slot_state; + slot->sequence.store(static_cast(pos + i + mask + 1), std::memory_order_release); + atomic_ops++; // sequence.store + } + atomic_count += atomic_ops; + if (contended) { + wait_cycle += (get_sys_cnt_aicpu() - t0); + } + return count; + } +#endif +}; + +// Cold-path ready queue operations (defined in pto_scheduler.cpp). Declared +// as non-member so PTO2ReadyQueue stays a POD-like struct with cache-line +// alignment. Storage is owned by the caller-supplied arena. +// reserve_layout: declare the slots[] region on the arena (must precede commit) +// init_from_layout: bind slots pointer from arena.region_ptr(off) and +// initialize sequence counters +// destroy: forget the slots pointer (arena owns the buffer) +size_t ready_queue_reserve_layout(DeviceArena &arena, uint64_t capacity); +// Writes everything *except* the arena-internal `slots` pointer field +// (sequences/positions on the slot array, capacity, mask). Uses +// arena.region_ptr(slots_off) only to address the slot array for writes; +// does NOT store the pointer in `queue->slots`. Call +// `ready_queue_wire_arena_pointers` afterwards to set the field itself. +bool ready_queue_init_data_from_layout(PTO2ReadyQueue *queue, DeviceArena &arena, size_t slots_off, uint64_t capacity); +// Stores queue->slots = arena.region_ptr(slots_off). Idempotent. +void ready_queue_wire_arena_pointers(PTO2ReadyQueue *queue, DeviceArena &arena, size_t slots_off); +void ready_queue_destroy(PTO2ReadyQueue *queue); + +/** + * Statistics returned by mixed-task completion processing + */ +struct CompletionStats { + int32_t fanout_edges; // Number of fanout edges traversed (notify consumers) + int32_t tasks_enqueued; // Number of consumers that became READY + int32_t fanin_edges; // Number of fanin edges traversed (release producers) + bool mixed_task_completed; // True only when this callback completed a mixed task +}; + +/** + * Layout descriptor produced by PTO2SchedulerState::reserve_layout(). Holds + * the arena offsets of every sub-region the scheduler needs plus the + * capacities used at layout time (init_from_layout reuses them). + */ +struct PTO2SchedulerLayout { + size_t off_ready_queue_slots[PTO2_NUM_RESOURCE_SHAPES]; + size_t off_dummy_ready_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]; +}; + +/** + * Scheduler state structure + * + * Contains dynamic state updated during task execution. + * Separated from shared memory for cache efficiency. + * Hot-path methods are defined inline (implicitly inline as member functions). + */ +struct PTO2SchedulerState { + // Shared memory access + PTO2SharedMemoryHeader *sm_header; + + // Per-ring state + struct alignas(64) RingSchedState { + // --- Cache Line 0: ring pointer (read-only) + hot path (read-write) --- + PTO2SharedMemoryRingHeader *ring; + int32_t last_task_alive; + std::atomic advance_lock; // multi-thread CAS + + // --- Cache Line 1+: Orch-side wiring dep_pool --- + alignas(64) PTO2DepListPool dep_pool; +#if SIMPLER_DFX + // Published only for scope_stats; orchestrator must not read dep_pool's non-atomic counters directly. + alignas(64) std::atomic dep_pool_snapshot_tail; + std::atomic dep_pool_snapshot_top; +#endif + + // Initialize arena-internal data + arena-external pointers; does NOT + // store dep_pool.base (that lives in the runtime arena and is wired + // by SchedulerState::wire_arena_pointers). The `ring` field stores + // the device address of the SM ring header — computed via offset + // arithmetic, no SM dereference. + bool init_data_from_layout(void *sm_dev_base, int32_t ring_id); + void reset_for_reuse(void *sm_dev_base, int32_t ring_id, std::atomic *orch_err); + void destroy(); + + void sync_to_sm() { ring->fc.last_task_alive.store(last_task_alive, std::memory_order_release); } + +#if SIMPLER_DFX + void publish_dep_pool_snapshot() { + dep_pool_snapshot_tail.store(dep_pool.tail, std::memory_order_release); + dep_pool_snapshot_top.store(dep_pool.top, std::memory_order_release); + } + + void read_dep_pool_snapshot(int32_t &tail, int32_t &top) const { + top = dep_pool_snapshot_top.load(std::memory_order_acquire); + tail = dep_pool_snapshot_tail.load(std::memory_order_acquire); + if (tail > top) tail = top; + } +#endif + + void advance_ring_pointers() { + int32_t current_task_index = ring->fc.current_task_index.load(std::memory_order_acquire); + int32_t old_last_task_alive = last_task_alive; + + while (last_task_alive < current_task_index) { + PTO2TaskSlotState &slot_state = ring->get_slot_state_by_task_id(last_task_alive); + if (slot_state.task_state.load(std::memory_order_acquire) != PTO2_TASK_CONSUMED) { + break; + } + last_task_alive++; + } + + // Eager reset: prepare reclaimed slots for reuse while still hot in cache. + // Safe because last_task_alive has advanced past these slots but + // sync_to_sm has not yet published — the orchestrator cannot reuse + // them until the release store below. + // Skips payload, task, ring_id — immutable after RingSchedState::init(). + for (int32_t id = old_last_task_alive; id < last_task_alive; id++) { + ring->get_slot_state_by_task_id(id).reset_for_reuse(); + } + + sync_to_sm(); + } + } ring_sched_states[PTO2_MAX_RING_DEPTH]; + + // Ready queues remain global (scheduling is ring-agnostic) + PTO2ReadyQueue ready_queues[PTO2_NUM_RESOURCE_SHAPES]; + + // Dependency-only tasks (active_mask is empty, shape == DUMMY). Drained by + // the dispatch loop and completed inline -- never goes to AICore. + PTO2ReadyQueue dummy_ready_queue; + + alignas(64) AsyncWaitList async_wait_list; + + // Statistics (cold path, isolated from hot-path fields) +#if SIMPLER_SCHED_PROFILING + alignas(64) std::atomic tasks_completed; + std::atomic tasks_consumed; +#endif + // ========================================================================= + // Inline hot-path methods + // ========================================================================= + + // Route a ready slot to the right global queue. Dummy tasks (empty + // active_mask) live in dummy_ready_queue; everything else goes to the + // per-shape ready_queues[]. + void push_ready_routed(PTO2TaskSlotState *slot_state) { + PTO2ResourceShape shape = slot_state->active_mask.to_shape(); + if (shape == PTO2ResourceShape::DUMMY || + (slot_state->active_mask.has_predicate() && !slot_state->payload->predicate.pass())) { + dummy_ready_queue.push(slot_state); + } else { + ready_queues[static_cast(shape)].push(slot_state); + } + } + + void check_and_handle_consumed(PTO2TaskSlotState &slot_state) { + // Read fanout_refcount/fanout_count and flip COMPLETED->CONSUMED under + // fanout_lock. The orchestrator claims producers (fanout_count++) under the + // same lock, so the consume decision is serialized against a concurrent + // claim: either the ++ lands first (count then exceeds refcount, so we do + // not consume and the producer stays pinned until released) or the consume + // lands first (the orchestrator then observes CONSUMED and skips the + // claim). Without this lock a claim racing the consume desyncs the slot's + // refcount and wedges in-order reclaim. + bool became_consumed = false; + slot_state.lock_fanout(); + if (slot_state.fanout_refcount.load(std::memory_order_acquire) == slot_state.fanout_count) { + PTO2TaskState expected = PTO2_TASK_COMPLETED; + became_consumed = slot_state.task_state.compare_exchange_strong( + expected, PTO2_TASK_CONSUMED, std::memory_order_acq_rel, std::memory_order_acquire + ); + } + slot_state.unlock_fanout(); + if (!became_consumed) return; + +#if SIMPLER_SCHED_PROFILING + tasks_consumed.fetch_add(1, std::memory_order_relaxed); +#endif + + int32_t ring_id = slot_state.ring_id; + // advance_ring_pointers (and the reset_for_reuse it triggers) MUST run + // outside fanout_lock: reset_for_reuse stores fanout_lock=0 and would + // clobber a held lock. Safe here — the slot is CONSUMED and quiescent. + // Try-lock — if another thread is advancing this ring, it will scan our CONSUMED task + int32_t expected_lock = 0; + if (ring_sched_states[ring_id].advance_lock.compare_exchange_strong( + expected_lock, 1, std::memory_order_acquire, std::memory_order_relaxed + )) { + ring_sched_states[ring_id].advance_ring_pointers(); + ring_sched_states[ring_id].advance_lock.store(0, std::memory_order_release); + } + } + +#if SIMPLER_ORCH_PROFILING || SIMPLER_SCHED_PROFILING + void check_and_handle_consumed(PTO2TaskSlotState &slot_state, uint64_t &atomic_count) { + // See the non-profiling overload for why the read + COMPLETED->CONSUMED + // flip is serialized against the orchestrator's claim under fanout_lock. + bool became_consumed = false; + slot_state.lock_fanout(); + atomic_count += 1; // lock CAS + uint32_t fc = slot_state.fanout_count; + uint32_t rc = slot_state.fanout_refcount.load(std::memory_order_acquire); + atomic_count += 1; // fanout_refcount.load (fanout_count is a plain read under lock) + if (rc == fc) { + PTO2TaskState expected = PTO2_TASK_COMPLETED; + became_consumed = slot_state.task_state.compare_exchange_strong( + expected, PTO2_TASK_CONSUMED, std::memory_order_acq_rel, std::memory_order_acquire + ); + atomic_count += 1; // CAS + } + slot_state.unlock_fanout(); + atomic_count += 1; // unlock store + if (!became_consumed) return; + +#if SIMPLER_SCHED_PROFILING + tasks_consumed.fetch_add(1, std::memory_order_relaxed); +#endif + + int32_t ring_id = slot_state.ring_id; + // advance_ring_pointers + reset_for_reuse run outside fanout_lock (reset + // stores fanout_lock=0). Safe — the slot is CONSUMED and quiescent. + // Try-lock — if another thread is advancing this ring, it will scan our CONSUMED task + int32_t expected_lock = 0; + if (ring_sched_states[ring_id].advance_lock.compare_exchange_strong( + expected_lock, 1, std::memory_order_acquire, std::memory_order_relaxed + )) { + ring_sched_states[ring_id].advance_ring_pointers(); + ring_sched_states[ring_id].advance_lock.store(0, std::memory_order_release); + atomic_count += 2; // try-lock CAS + unlock store + } else { + atomic_count += 1; // failed try-lock CAS + } + } +#endif + + void release_producer(PTO2TaskSlotState &slot_state) { + slot_state.fanout_refcount.fetch_add(1, std::memory_order_acq_rel); + check_and_handle_consumed(slot_state); + } + + // Scope-end release: sets bit31 (PTO2_FANOUT_SCOPE_BIT) instead of bumping a + // consumer ref. Called exactly once per task from on_scope_end. Keeping it a + // distinct add lets a consumer release leave the scope bit unset, so "all + // consumers done but scope still open" stays distinguishable from "fully + // consumed". + void release_producer_scope(PTO2TaskSlotState &slot_state) { + slot_state.fanout_refcount.fetch_add(PTO2_FANOUT_SCOPE_BIT, std::memory_order_acq_rel); + check_and_handle_consumed(slot_state); + } + +#if SIMPLER_ORCH_PROFILING || SIMPLER_SCHED_PROFILING + void release_producer(PTO2TaskSlotState &slot_state, uint64_t &atomic_count) { + slot_state.fanout_refcount.fetch_add(1, std::memory_order_acq_rel); + atomic_count += 1; // fanout_refcount.fetch_add + check_and_handle_consumed(slot_state, atomic_count); + } + + void release_producer_scope(PTO2TaskSlotState &slot_state, uint64_t &atomic_count) { + slot_state.fanout_refcount.fetch_add(PTO2_FANOUT_SCOPE_BIT, std::memory_order_acq_rel); + atomic_count += 1; // fanout_refcount.fetch_add + check_and_handle_consumed(slot_state, atomic_count); + } +#endif + + bool release_fanin_and_check_ready(PTO2TaskSlotState &slot_state) { + // 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 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) { + 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 false; + } +#endif + + int get_ready_tasks_batch(PTO2ResourceShape shape, PTO2TaskSlotState **out, int max_count) { + return ready_queues[static_cast(shape)].pop_batch(out, max_count); + } + +#if SIMPLER_SCHED_PROFILING + int get_ready_tasks_batch( + PTO2ResourceShape shape, PTO2TaskSlotState **out, int max_count, uint64_t &atomic_count, uint64_t &wait_cycle + ) { + return ready_queues[static_cast(shape)].pop_batch(out, max_count, atomic_count, wait_cycle); + } +#endif + + void on_scope_end(PTO2TaskSlotState **task_slot_states, int32_t count) { +#if SIMPLER_ORCH_PROFILING + extern uint64_t g_orch_scope_end_atomic_count; + if (count > 0) __builtin_prefetch(task_slot_states[0], 1, 0); + for (int32_t i = 0; i < count; i++) { + if (i + 1 < count) __builtin_prefetch(task_slot_states[i + 1], 1, 0); + release_producer_scope(*task_slot_states[i], g_orch_scope_end_atomic_count); + } +#else + if (count > 0) __builtin_prefetch(task_slot_states[0], 1, 0); + for (int32_t i = 0; i < count; i++) { + if (i + 1 < count) __builtin_prefetch(task_slot_states[i + 1], 1, 0); + release_producer_scope(*task_slot_states[i]); + } +#endif + } + + /** + * Subtask completion: atomic counter model. + * Called when a single subtask (AIC, AIV0, or AIV1) finishes on any block. + * Atomically increments completed_subtasks and checks whether all subtasks + * across all blocks are done. + * + * @return true if this was the last subtask, completing the entire task. + */ + bool on_subtask_complete(PTO2TaskSlotState &slot_state) { + int16_t prev = slot_state.completed_subtasks.fetch_add(1, std::memory_order_acq_rel); + return (prev + 1) == slot_state.total_required_subtasks; + } + + /** + * Two-stage completion: second stage. + * Called exactly once when all subtasks of a mixed task are done + * (i.e., on_subtask_complete returned true). + * Handles fanout notification, fanin release, and self-consumption check. + */ +#if SIMPLER_SCHED_PROFILING + CompletionStats +#else + void +#endif + on_task_complete( + PTO2TaskSlotState &slot_state +#if SIMPLER_SCHED_PROFILING + , + int thread_idx +#endif + ) { +#if SIMPLER_SCHED_PROFILING + CompletionStats stats = {0, 0, 0, true}; +#endif +#if SIMPLER_SCHED_PROFILING + extern uint64_t g_sched_lock_cycle[], g_sched_fanout_cycle[]; + extern uint64_t g_sched_lock_atomic_count[], g_sched_lock_wait_cycle[]; + extern uint64_t g_sched_fanout_atomic_count[], g_sched_push_wait_cycle[]; + uint64_t lock_atomics = 0, lock_wait = 0; + PTO2_SCHED_CYCLE_START(); +#endif + +#if SIMPLER_SCHED_PROFILING + slot_state.lock_fanout(lock_atomics, lock_wait); +#else + slot_state.lock_fanout(); +#endif + slot_state.task_state.store(PTO2_TASK_COMPLETED, std::memory_order_release); + PTO2DepListEntry *current = slot_state.fanout_head; // Protected by fanout_lock + slot_state.unlock_fanout(); + +#if SIMPLER_SCHED_PROFILING + lock_atomics += 2; // state.store + unlock.store + g_sched_lock_atomic_count[thread_idx] += lock_atomics; + g_sched_lock_wait_cycle[thread_idx] += lock_wait; + PTO2_SCHED_CYCLE_LAP(g_sched_lock_cycle[thread_idx]); +#endif + + // Fanout: notify consumers +#if SIMPLER_SCHED_PROFILING + uint64_t fanout_atomics = 0, push_wait = 0; +#endif + 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)) { + stats.tasks_enqueued++; + } +#else + release_fanin_and_check_ready(consumer_slot); +#endif + current = current->next; + } + +#if SIMPLER_SCHED_PROFILING + g_sched_fanout_atomic_count[thread_idx] += fanout_atomics; + g_sched_push_wait_cycle[thread_idx] += push_wait; + PTO2_SCHED_CYCLE_LAP(g_sched_fanout_cycle[thread_idx]); + return stats; +#endif + } + + /** + * Cold path: release producers (fanin traversal) + check self for CONSUMED. + * Returns fanin edge count for profiling. + */ + +#if SIMPLER_SCHED_PROFILING + int32_t on_task_release(PTO2TaskSlotState &slot_state, int32_t thread_idx) { + PTO2_SCHED_CYCLE_START(); + extern uint64_t g_sched_fanin_cycle[], g_sched_fanin_atomic_count[]; + extern uint64_t g_sched_self_atomic_count[]; + extern uint64_t g_sched_self_consumed_cycle[]; + extern uint64_t g_sched_complete_count[]; + uint64_t fanin_atomics = 0; +#else + int32_t on_task_release(PTO2TaskSlotState &slot_state) { +#endif + PTO2TaskPayload *payload = slot_state.payload; + for_each_fanin_slot_state(*payload, [&](PTO2TaskSlotState *producer_slot_state) { +#if SIMPLER_SCHED_PROFILING + release_producer(*producer_slot_state, fanin_atomics); +#else + release_producer(*producer_slot_state); +#endif + }); +#if SIMPLER_SCHED_PROFILING + g_sched_fanin_atomic_count[thread_idx] += fanin_atomics; + PTO2_SCHED_CYCLE_LAP(g_sched_fanin_cycle[thread_idx]); +#endif + + // Self consumed check +#if SIMPLER_SCHED_PROFILING + uint64_t self_atomics = 0; + check_and_handle_consumed(slot_state, self_atomics); + g_sched_self_atomic_count[thread_idx] += self_atomics; + PTO2_SCHED_CYCLE_LAP(g_sched_self_consumed_cycle[thread_idx]); + g_sched_complete_count[thread_idx]++; +#else + check_and_handle_consumed(slot_state); +#endif + return payload->fanin_actual_count; + } + + // === Cold-path API (defined in pto_scheduler.cpp) === + + // Phase 1: declare every sub-region (ready_queue slots, dummy queue slots, + // per-ring dep_pool entries) on the supplied arena. + // Capacities are baked into the returned layout; init_data_from_layout uses + // the same values. + static PTO2SchedulerLayout reserve_layout(DeviceArena &arena, int32_t dep_pool_capacity = PTO2_DEP_LIST_POOL_SIZE); + static PTO2SchedulerLayout + reserve_layout(DeviceArena &arena, const int32_t dep_pool_capacities[PTO2_MAX_RING_DEPTH]); + + // Phase 3a: write everything *except* arena-internal pointer fields. + // `sm_dev_base` is the device address of the SM (only stored, never + // dereferenced here). Safe to call on a host arena that holds the + // prebuilt image buffer. (The orchestrator counterpart takes + // task_window_size for ring task_descriptors address arithmetic; the + // scheduler only needs the SM header / ring header base addresses, + // both window-size-independent.) + bool init_data_from_layout(const PTO2SchedulerLayout &layout, DeviceArena &arena, void *sm_dev_base); + void reset_for_reuse(const PTO2SchedulerLayout &layout, void *sm_dev_base); + + // Phase 3b: write the arena-internal pointer fields + // (ready_queues[].slots, dummy_ready_queue.slots, dep_pool.base for each + // ring). Called on both host and device sides. + void wire_arena_pointers(const PTO2SchedulerLayout &layout, DeviceArena &arena); + + // Forget per-region pointers; arena owns the backing memory. + void destroy(); + void print_stats(); + void print_queues(); +}; + +// Scheduler cold-path API is declared as PTO2SchedulerState member functions. +// See init()/destroy()/print_stats()/print_queues() below the struct definition. + +// Short-circuit NotDeferred completions seen during drain so they don't grow +// entries[]. Mirrors the a2a3 impl; see that mirror for the rationale. +inline bool +AsyncWaitList::try_inline_complete_locked(AsyncWaitList::DrainCompletionSink &sink, PTO2TaskSlotState &slot_state) { +#if SIMPLER_SCHED_PROFILING + sink.sched->on_task_complete(slot_state, sink.thread_idx); +#else + sink.sched->on_task_complete(slot_state); +#endif + if (*sink.deferred_release_count >= sink.deferred_release_capacity) { + while (*sink.deferred_release_count > 0) { +#if SIMPLER_SCHED_PROFILING + (void)sink.sched->on_task_release( + *sink.deferred_release_slot_states[--(*sink.deferred_release_count)], sink.thread_idx + ); +#else + sink.sched->on_task_release(*sink.deferred_release_slot_states[--(*sink.deferred_release_count)]); +#endif + } + } + sink.deferred_release_slot_states[(*sink.deferred_release_count)++] = &slot_state; + sink.inline_completed++; + return true; +} + +template +inline AsyncPollResult AsyncWaitList::poll_and_complete( + AICoreCompletionMailbox *aicore_mailbox, PTO2SchedulerState *sched, + PTO2TaskSlotState **deferred_release_slot_states, int32_t &deferred_release_count, int32_t deferred_release_capacity +#if SIMPLER_SCHED_PROFILING + , + int thread_idx +#endif +) { + AsyncPollResult result; + if (!try_lock()) return result; + + AsyncWaitList::DrainCompletionSink sink{}; + sink.sched = sched; + sink.deferred_release_slot_states = deferred_release_slot_states; + sink.deferred_release_count = &deferred_release_count; + sink.deferred_release_capacity = deferred_release_capacity; +#if SIMPLER_SCHED_PROFILING + sink.thread_idx = thread_idx; +#endif + + int32_t drain_err = PTO2_ERROR_NONE; + drain_aicore_completion_mailbox_locked(aicore_mailbox, sink, drain_err); + if (drain_err != PTO2_ERROR_NONE) { + result.error_code = drain_err; + unlock(); + return result; + } + result.completed += sink.inline_completed; + + for (int32_t i = count - 1; i >= 0; --i) { + AsyncWaitEntry &entry = entries[i]; + for (int32_t c = 0; c < entry.condition_count; c++) { + CompletionCondition &cond = entry.conditions[c]; + if (cond.satisfied) continue; + CompletionPollResult poll = cond.test(); + if (poll.state == CompletionPollState::FAILED) { + result.error_code = poll.error_code; + result.failed_slot_state = entry.slot_state; + unlock(); + return result; + } + if (poll.state == CompletionPollState::READY) { + cond.satisfied = true; + cond.retire(); + entry.waiting_completion_count--; + } + } + + if (entry.normal_done && entry.waiting_completion_count <= 0) { +#if SIMPLER_SCHED_PROFILING + sched->on_task_complete(*entry.slot_state, thread_idx); +#else + sched->on_task_complete(*entry.slot_state); +#endif + if (deferred_release_count >= deferred_release_capacity) { + while (deferred_release_count > 0) { +#if SIMPLER_SCHED_PROFILING + (void)sched->on_task_release(*deferred_release_slot_states[--deferred_release_count], thread_idx); +#else + sched->on_task_release(*deferred_release_slot_states[--deferred_release_count]); +#endif + } + } + deferred_release_slot_states[deferred_release_count++] = entry.slot_state; + result.completed++; + + int32_t last = count - 1; + if (i != last) entries[i] = entries[last]; + count = last; + } + } + + unlock(); + return result; +} + +// ============================================================================= +// Scheduler Profiling Data +// ============================================================================= + +#if SIMPLER_SCHED_PROFILING +struct PTO2SchedProfilingData { + // Sub-phase cycle breakdown within on_task_complete + uint64_t lock_cycle; // lock_fanout + state store + unlock + uint64_t fanout_cycle; // fanout traversal + uint64_t fanin_cycle; // fanin traversal + uint64_t self_consumed_cycle; // self check_and_handle_consumed + + // Wait times + uint64_t lock_wait_cycle; // spin-wait in fanout_lock + uint64_t push_wait_cycle; // CAS contention in push() + uint64_t pop_wait_cycle; // CAS contention in pop() + + // Atomic counts per sub-phase + uint64_t lock_atomic_count; + uint64_t fanout_atomic_count; + uint64_t fanin_atomic_count; + uint64_t self_atomic_count; + uint64_t pop_atomic_count; + + int64_t complete_count; +}; + +/** + * Get and reset scheduler profiling data for a specific thread. + * Returns accumulated profiling data and resets counters. + */ +PTO2SchedProfilingData scheduler_get_profiling(int thread_idx); +#endif diff --git a/_run_early_st_board.sh b/_run_early_st_board.sh new file mode 100755 index 0000000000..1b3df0e9f6 --- /dev/null +++ b/_run_early_st_board.sh @@ -0,0 +1,22 @@ +#!/bin/bash +set -euo pipefail +cd /home/pyptouser/yanghaoran/Desktop/simpler2 +source .venv/bin/activate +DEV="${ASCEND_RT_VISIBLE_DEVICES:-${DEVICE:-}}" +# Prefer lock-assigned device: task-submit often sets ASCEND_RT_VISIBLE_DEVICES +if [[ -z "${DEV}" ]]; then + echo "No device env; failing" + env | sort | head -40 + exit 2 +fi +# If ASCEND is set to a single id by task-submit, use it +echo "ASCEND_RT_VISIBLE_DEVICES=${ASCEND_RT_VISIBLE_DEVICES:-unset} DEVICE=${DEVICE:-unset}" +# Use first id for --device +DID="${ASCEND_RT_VISIBLE_DEVICES%%,*}" +DID="${DID:-$DEV}" +echo "Using --device $DID" +TEST="${1:?test path}" +NAME="${2:?label}" +python -m pytest "$TEST" --platform a5 --device "$DID" -v --tb=short \ + --enable-l2-swimlane 4 --enable-dep-gen \ + 2>&1 | tee "_early_st_out/${NAME}_board.log" diff --git a/_run_st_dev2.sh b/_run_st_dev2.sh new file mode 100755 index 0000000000..9e9e031551 --- /dev/null +++ b/_run_st_dev2.sh @@ -0,0 +1,9 @@ +#!/bin/bash +set -euo pipefail +cd /home/pyptouser/yanghaoran/Desktop/simpler2 +source .venv/bin/activate +# task-submit --device N may not export ASCEND; hardcode from argv +DEV="${1:?device}" +SHIFT=1 +shift +python -m pytest "$@" --platform a5 --device "$DEV" -v --tb=line diff --git a/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/README.md b/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/README.md new file mode 100644 index 0000000000..16a3d9e7de --- /dev/null +++ b/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/README.md @@ -0,0 +1,62 @@ +# Acc ValidShape strip C2V (`srcStride` repro) + +Minimal Ascend950 **simpler** sample for +`ISSUE_A5_ACC_VALIDSHAPE_C2V_SRCSTRIDE`. + +## Primary test + +`test_acc_c2v_strip_vs_full.py` — same Acc drained two ways: + +| Output | Path | +|--------|------| +| `C_full` | one Acc TPUSH, `Valid≡Rows` | +| `C_strip` | Acc+`ValidShape(H=16)` windows @ `addr=row*64` | + +Host: `C_strip ≈ C_full`. Shapes (UB-safe): `M=32, N=128, H=16, K=32` +(still `validRow < Rows`). Seed `torch.manual_seed(0)`. + +### Board results (2026-07-21, device dump) + +#### 修复前 (`srcStride=align(validRow)`) + +典型现象：strip0（行 0…15）与 full 一致；strip1（行 16…）系统性偏离。 +`col=0` 在本 seed 下仍可能碰巧相等，看同行列最大偏差处： + +```text +# strip0 — 接近/相等 +row=0 C_full[0,0]=10.4928 C_strip[0,0]=10.4928 diff=0 +row=15 C_full[15,0]=-2.02309 C_strip[15,0]=-2.02309 diff=0 + +# strip1 — 明显不等（本 seed 最大差在下列） +row=16 C_full[16,71]=-12.8381 C_strip[16,71]=13.6886 |diff|=26.5268 +row=17 C_full[17,44]=-11.6004 C_strip[17,44]=11.7646 |diff|=23.365 +# 汇总: strip0_max=0.0 strip1_max=31.9782 → 主测 FAILED +``` + +#### 修复后 (`srcStride=align(Rows)`, pin `0ebbd03d`) + +```text +row=0 C_full[0,0]=10.4928 C_strip[0,0]=10.4928 diff=0 +row=15 C_full[15,0]=-2.02309 C_strip[15,0]=-2.02309 diff=0 +row=16 C_full[16,0]=2.29007 C_strip[16,0]=2.29007 diff=0 +row=17 C_full[17,0]=5.26025 C_strip[17,0]=5.26025 diff=0 +# 汇总: strip0_max=0 strip1_max=0 → PASSED +``` + +```bash +# 在已配置 A5 的 simpler 仓库根目录： +python examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/test_acc_c2v_strip_vs_full.py \ + -p a5 -d + +# 或用包根的 board 包装（改 SIMPLER 后）： +task-submit --device auto --max-time 900 --run \ + "bash /path/to/run_acc_c2v_strip_vs_full_board.sh" +``` + +## Files + +| Path | Role | +|------|------| +| `kernels/mix/kernel_acc_c2v_compare.cpp` | matmul + full + strip TPUSH | +| `kernels/orchestration/acc_c2v_compare_orch.cpp` | `[A,B,C_full,C_strip]` | +| `test_acc_c2v_strip_vs_full.py` | self-compare | diff --git a/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/kernels/aic/kernel_matmul_tstore.cpp b/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/kernels/aic/kernel_matmul_tstore.cpp new file mode 100644 index 0000000000..360f4f6273 --- /dev/null +++ b/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/kernels/aic/kernel_matmul_tstore.cpp @@ -0,0 +1,69 @@ +/* + * CONTROL: Acc float matmul → TSTORE Acc to GM (no C2V). Isolates matmul golden match. + */ +#include +#include +#include "tensor.h" + +using namespace pto; + +#ifndef __gm__ +#define __gm__ +#endif +#ifndef __aicore__ +#define __aicore__ [aicore] +#endif +#include "intrinsic.h" + +constexpr int M = 128; +constexpr int N = 256; +constexpr int K = 32; + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *a_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ Tensor *b_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ Tensor *c_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); + + __gm__ float *a_ptr = reinterpret_cast<__gm__ float *>(a_tensor->buffer.addr) + a_tensor->start_offset; + __gm__ float *b_ptr = reinterpret_cast<__gm__ float *>(b_tensor->buffer.addr) + b_tensor->start_offset; + __gm__ float *c_ptr = reinterpret_cast<__gm__ float *>(c_tensor->buffer.addr) + c_tensor->start_offset; + + using GlobalA = GlobalTensor, pto::Stride>; + using GlobalB = GlobalTensor, pto::Stride>; + using GlobalC = GlobalTensor, pto::Stride>; + GlobalA aGlobal(a_ptr); + GlobalB bGlobal(b_ptr); + GlobalC cGlobal(c_ptr); + + using TileMatA = Tile; + using TileMatB = Tile; + using LeftT = TileLeft; + using RightT = TileRight; + using AccT = TileAcc; + + TileMatA aMat; + TileMatB bMat; + TASSIGN(aMat, 0x0); + TASSIGN(bMat, 0x20000); + LeftT aL0; + RightT bL0; + AccT acc; + TASSIGN(aL0, 0x0); + TASSIGN(bL0, 0x0); + TASSIGN(acc, 0x0); + + TLOAD(aMat, aGlobal); + TLOAD(bMat, bGlobal); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + TMOV(aL0, aMat); + TMOV(bL0, bMat); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + TMATMUL(acc, aL0, bL0); + set_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + wait_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + TSTORE(cGlobal, acc); + set_flag(PIPE_FIX, PIPE_S, EVENT_ID7); + wait_flag(PIPE_FIX, PIPE_S, EVENT_ID7); +} diff --git a/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/kernels/mix/kernel_acc_c2v_compare.cpp b/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/kernels/mix/kernel_acc_c2v_compare.cpp new file mode 100644 index 0000000000..eb093ea56b --- /dev/null +++ b/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/kernels/mix/kernel_acc_c2v_compare.cpp @@ -0,0 +1,194 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +/** + * A5 Acc ValidShape strip C2V repro (self-compare). + * + * AIC: Acc = A@B once, then + * 1) full Acc TPUSH (Valid≡Rows) → AIV stores C_full + * 2) 8× Acc+Valid(H) TPUSH @ addr=row*64 → AIV stores C_strip + * + * Host golden: C_strip == C_full. + * - unfixed TMovCcToUb (srcStride=validRow): FAIL + * - fixed (srcStride=Rows): PASS + * + * args[0]=A, args[1]=B, args[2]=C_full, args[3]=C_strip + */ + +#include +#include +#include +#include "tensor.h" + +using pto::BLayout; +using pto::Direction; +using pto::GlobalTensor; +using pto::Shape; +using pto::SLayout; +using pto::Tile; +using pto::TileAcc; +using pto::TileLeft; +using pto::TileRight; +using pto::TileSplitAxis; +using pto::TileType; +using pto::TPipe; + +#ifndef __gm__ +#define __gm__ +#endif +#ifndef __aicore__ +#define __aicore__ [aicore] +#endif +#include "intrinsic.h" + +#ifdef __DAV_CUBE__ +constexpr bool DAV_CUBE = true; +#else +constexpr bool DAV_CUBE = false; +#endif +#ifdef __DAV_VEC__ +constexpr bool DAV_VEC = true; +#else +constexpr bool DAV_VEC = false; +#endif + +constexpr int M = 32; +constexpr int N = 128; +constexpr int H = 16; +constexpr int K = 32; +constexpr uint64_t kAccRowByteStride = 64; +constexpr uint16_t PP_FLAG_ID = 0; +constexpr uint8_t PP_FIFO_DEPTH = 1; + +constexpr uint32_t kFullSlot = static_cast(M * N * sizeof(float)); +constexpr uint32_t kStripSlot = static_cast(H * N * sizeof(float)); + +using AccFullT = TileAcc; +using AccWinFullT = Tile; +using AccWinStripT = Tile; +using VecFullT = Tile; +using VecStripT = Tile; + +using PipeFullT = TPipe; +using PipeStripT = TPipe; + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *a_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ Tensor *b_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ Tensor *c_full_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); + __gm__ Tensor *c_strip_tensor = reinterpret_cast<__gm__ Tensor *>(args[3]); + + PipeFullT pipeFull(nullptr, 0U, 0U); + PipeStripT pipeStrip(nullptr, 0U, 0U); + + if constexpr (DAV_CUBE) { + __gm__ float *a_ptr = + reinterpret_cast<__gm__ float *>(a_tensor->buffer.addr) + a_tensor->start_offset; + __gm__ float *b_ptr = + reinterpret_cast<__gm__ float *>(b_tensor->buffer.addr) + b_tensor->start_offset; + + using GlobalA = GlobalTensor, pto::Stride>; + using GlobalB = GlobalTensor, pto::Stride>; + GlobalA aGlobal(a_ptr); + GlobalB bGlobal(b_ptr); + + using TileMatA = Tile; + using TileMatB = Tile; + using LeftT = TileLeft; + using RightT = TileRight; + + TileMatA aMat; + TileMatB bMat; + TASSIGN(aMat, 0x0); + TASSIGN(bMat, 0x20000); + LeftT aL0; + RightT bL0; + AccFullT acc; + TASSIGN(aL0, 0x0); + TASSIGN(bL0, 0x0); + TASSIGN(acc, 0x0); + + TLOAD(aMat, aGlobal); + TLOAD(bMat, bGlobal); + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + TMOV(aL0, aMat); + TMOV(bL0, bMat); + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + TMATMUL(acc, aL0, bL0); + set_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + wait_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + + // (1) full Acc C2V — Valid≡Rows, always correct for srcStride either formula + { + AccWinFullT full; + TASSIGN(full, 0x0); + TPUSH(pipeFull, full); + } + + // (2) strip Acc C2V — Valid(H) windows (the buggy pattern) + for (int row = 0; row < M; row += H) { + AccWinStripT strip; + TASSIGN(strip, static_cast(row) * kAccRowByteStride); + TPUSH(pipeStrip, strip); + } + + set_flag(PIPE_FIX, PIPE_S, EVENT_ID7); + wait_flag(PIPE_FIX, PIPE_S, EVENT_ID7); + } + + if constexpr (DAV_VEC) { + if (get_sub_block_id(args) != 0) { + return; + } + __gm__ float *c_full = + reinterpret_cast<__gm__ float *>(c_full_tensor->buffer.addr) + c_full_tensor->start_offset; + __gm__ float *c_strip = + reinterpret_cast<__gm__ float *>(c_strip_tensor->buffer.addr) + c_strip_tensor->start_offset; + + // UB layout: full fifo + strip fifo + scratch + constexpr uint64_t kFullFifo = static_cast(PP_FIFO_DEPTH) * kFullSlot; + constexpr uint64_t kStripFifo = static_cast(PP_FIFO_DEPTH) * kStripSlot; + VecFullT vFull; + VecStripT vStrip; + TASSIGN(vFull, kFullFifo + kStripFifo); + TASSIGN(vStrip, kFullFifo + kStripFifo + kFullSlot); + + // Pop full + TPOP(pipeFull, vFull); + set_flag(PIPE_S, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_S, PIPE_MTE3, EVENT_ID0); + using GlobalFull = GlobalTensor, pto::Stride>; + GlobalFull gFull(c_full); + TSTORE(gFull, vFull); + TFREE(pipeFull); + set_flag(PIPE_MTE3, PIPE_S, EVENT_ID1); + wait_flag(PIPE_MTE3, PIPE_S, EVENT_ID1); + + // Pop strips + for (int row = 0; row < M; row += H) { + TPOP(pipeStrip, vStrip); + set_flag(PIPE_S, PIPE_MTE3, EVENT_ID2); + wait_flag(PIPE_S, PIPE_MTE3, EVENT_ID2); + using GlobalStrip = + GlobalTensor, pto::Stride>; + GlobalStrip gStrip(c_strip + static_cast(row) * N); + TSTORE(gStrip, vStrip); + TFREE(pipeStrip); + set_flag(PIPE_MTE3, PIPE_S, EVENT_ID3); + wait_flag(PIPE_MTE3, PIPE_S, EVENT_ID3); + } + } +} diff --git a/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/kernels/mix/kernel_acc_c2v_full.cpp b/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/kernels/mix/kernel_acc_c2v_full.cpp new file mode 100644 index 0000000000..e1bda584a0 --- /dev/null +++ b/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/kernels/mix/kernel_acc_c2v_full.cpp @@ -0,0 +1,192 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +/** + * A5 Acc ValidShape strip C2V repro — matches ISSUE_A5_ACC_VALIDSHAPE_C2V_SRCSTRIDE §样例. + * + * AIC: Acc[M,N] = A@B; then M/H × TPUSH Acc[M,N]+Valid(H,N) @ addr=row*64 (TILE_NO_SPLIT). + * AIV0: TPOP Vec[H,N] ND strip → TSTORE GM (AIV1 idle for C2V). + * + * USE_STRIP_C2V=0: one full Acc TPUSH (Valid≡Rows) control path. + */ + +#include +#include +#include + +#include "tensor.h" + +using pto::BLayout; +using pto::Direction; +using pto::GlobalTensor; +using pto::Shape; +using pto::SLayout; +using pto::Tile; +using pto::TileAcc; +using pto::TileLeft; +using pto::TileRight; +using pto::TileSplitAxis; +using pto::TileType; +using pto::TPipe; + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] +#endif + +#include "intrinsic.h" + +#ifdef __DAV_CUBE__ +constexpr bool DAV_CUBE = true; +#else +constexpr bool DAV_CUBE = false; +#endif + +#ifdef __DAV_VEC__ +constexpr bool DAV_VEC = true; +#else +constexpr bool DAV_VEC = false; +#endif + +#ifndef USE_STRIP_C2V +#define USE_STRIP_C2V 0 +#endif + +constexpr int M = 128; +constexpr int N = 256; +constexpr int H = 16; +constexpr int K = 32; + +constexpr uint64_t kAccRowByteStride = 64; + +constexpr uint16_t PP_FLAG_ID = 0; +constexpr uint8_t PP_FIFO_DEPTH = 2; + +#if USE_STRIP_C2V +constexpr int PUSH_ROWS = H; +constexpr int kNumPush = M / H; +#else +constexpr int PUSH_ROWS = M; +constexpr int kNumPush = 1; +#endif + +constexpr uint32_t kSlotBytes = static_cast(PUSH_ROWS * N * sizeof(float)); + +using AccFullT = TileAcc; +using AccPushT = Tile; +using VecPushT = Tile; +// IsNoSplit=true matches TILE_NO_SPLIT (AIV0 only), same as qr_proj full-Acc hand patch. +using PipeT = TPipe; + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *a_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ Tensor *b_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ Tensor *c_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); + + PipeT pipe(nullptr, 0U, 0U); + + if constexpr (DAV_CUBE) { + __gm__ float *a_ptr = + reinterpret_cast<__gm__ float *>(a_tensor->buffer.addr) + a_tensor->start_offset; + __gm__ float *b_ptr = + reinterpret_cast<__gm__ float *>(b_tensor->buffer.addr) + b_tensor->start_offset; + + using GlobalA = GlobalTensor, pto::Stride>; + using GlobalB = GlobalTensor, pto::Stride>; + GlobalA aGlobal(a_ptr); + GlobalB bGlobal(b_ptr); + + using TileMatA = Tile; + using TileMatB = Tile; + using LeftT = TileLeft; + using RightT = TileRight; + + TileMatA aMat; + TileMatB bMat; + TASSIGN(aMat, 0x0); + TASSIGN(bMat, 0x20000); + + LeftT aL0; + RightT bL0; + AccFullT acc; + TASSIGN(aL0, 0x0); + TASSIGN(bL0, 0x0); + TASSIGN(acc, 0x0); + + TLOAD(aMat, aGlobal); + TLOAD(bMat, bGlobal); + + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + + TMOV(aL0, aMat); + TMOV(bL0, bMat); + + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + + TMATMUL(acc, aL0, bL0); + + set_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + wait_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + +#if USE_STRIP_C2V + for (int row = 0; row < M; row += H) { + AccPushT strip; + TASSIGN(strip, static_cast(row) * kAccRowByteStride); + TPUSH(pipe, strip); + } +#else + { + AccPushT full; + TASSIGN(full, 0x0); + TPUSH(pipe, full); + } +#endif + + set_flag(PIPE_FIX, PIPE_S, EVENT_ID7); + wait_flag(PIPE_FIX, PIPE_S, EVENT_ID7); + } + + if constexpr (DAV_VEC) { + if (get_sub_block_id(args) != 0) { + return; + } + + __gm__ float *c_ptr = + reinterpret_cast<__gm__ float *>(c_tensor->buffer.addr) + c_tensor->start_offset; + + constexpr uint64_t kScratch = static_cast(PP_FIFO_DEPTH) * kSlotBytes; + VecPushT vec; + TASSIGN(vec, kScratch); + + for (int s = 0; s < kNumPush; ++s) { + TPOP(pipe, vec); + + int row0 = s * PUSH_ROWS; + using GlobalC = GlobalTensor, + pto::Stride>; + GlobalC cGlobal(c_ptr + static_cast(row0) * N); + + set_flag(PIPE_S, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_S, PIPE_MTE3, EVENT_ID0); + TSTORE(cGlobal, vec); + TFREE(pipe); + + set_flag(PIPE_MTE3, PIPE_S, EVENT_ID1); + wait_flag(PIPE_MTE3, PIPE_S, EVENT_ID1); + } + } +} diff --git a/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/kernels/mix/kernel_acc_c2v_strip.cpp b/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/kernels/mix/kernel_acc_c2v_strip.cpp new file mode 100644 index 0000000000..e6be66ec92 --- /dev/null +++ b/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/kernels/mix/kernel_acc_c2v_strip.cpp @@ -0,0 +1,192 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +/** + * A5 Acc ValidShape strip C2V repro — matches ISSUE_A5_ACC_VALIDSHAPE_C2V_SRCSTRIDE §样例. + * + * AIC: Acc[M,N] = A@B; then M/H × TPUSH Acc[M,N]+Valid(H,N) @ addr=row*64 (TILE_NO_SPLIT). + * AIV0: TPOP Vec[H,N] ND strip → TSTORE GM (AIV1 idle for C2V). + * + * USE_STRIP_C2V=0: one full Acc TPUSH (Valid≡Rows) control path. + */ + +#include +#include +#include + +#include "tensor.h" + +using pto::BLayout; +using pto::Direction; +using pto::GlobalTensor; +using pto::Shape; +using pto::SLayout; +using pto::Tile; +using pto::TileAcc; +using pto::TileLeft; +using pto::TileRight; +using pto::TileSplitAxis; +using pto::TileType; +using pto::TPipe; + +#ifndef __gm__ +#define __gm__ +#endif + +#ifndef __aicore__ +#define __aicore__ [aicore] +#endif + +#include "intrinsic.h" + +#ifdef __DAV_CUBE__ +constexpr bool DAV_CUBE = true; +#else +constexpr bool DAV_CUBE = false; +#endif + +#ifdef __DAV_VEC__ +constexpr bool DAV_VEC = true; +#else +constexpr bool DAV_VEC = false; +#endif + +#ifndef USE_STRIP_C2V +#define USE_STRIP_C2V 1 +#endif + +constexpr int M = 128; +constexpr int N = 256; +constexpr int H = 16; +constexpr int K = 32; + +constexpr uint64_t kAccRowByteStride = 64; + +constexpr uint16_t PP_FLAG_ID = 0; +constexpr uint8_t PP_FIFO_DEPTH = 2; + +#if USE_STRIP_C2V +constexpr int PUSH_ROWS = H; +constexpr int kNumPush = M / H; +#else +constexpr int PUSH_ROWS = M; +constexpr int kNumPush = 1; +#endif + +constexpr uint32_t kSlotBytes = static_cast(PUSH_ROWS * N * sizeof(float)); + +using AccFullT = TileAcc; +using AccPushT = Tile; +using VecPushT = Tile; +// IsNoSplit=true matches TILE_NO_SPLIT (AIV0 only), same as qr_proj full-Acc hand patch. +using PipeT = TPipe; + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + __gm__ Tensor *a_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); + __gm__ Tensor *b_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); + __gm__ Tensor *c_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); + + PipeT pipe(nullptr, 0U, 0U); + + if constexpr (DAV_CUBE) { + __gm__ float *a_ptr = + reinterpret_cast<__gm__ float *>(a_tensor->buffer.addr) + a_tensor->start_offset; + __gm__ float *b_ptr = + reinterpret_cast<__gm__ float *>(b_tensor->buffer.addr) + b_tensor->start_offset; + + using GlobalA = GlobalTensor, pto::Stride>; + using GlobalB = GlobalTensor, pto::Stride>; + GlobalA aGlobal(a_ptr); + GlobalB bGlobal(b_ptr); + + using TileMatA = Tile; + using TileMatB = Tile; + using LeftT = TileLeft; + using RightT = TileRight; + + TileMatA aMat; + TileMatB bMat; + TASSIGN(aMat, 0x0); + TASSIGN(bMat, 0x20000); + + LeftT aL0; + RightT bL0; + AccFullT acc; + TASSIGN(aL0, 0x0); + TASSIGN(bL0, 0x0); + TASSIGN(acc, 0x0); + + TLOAD(aMat, aGlobal); + TLOAD(bMat, bGlobal); + + set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); + + TMOV(aL0, aMat); + TMOV(bL0, bMat); + + set_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); + + TMATMUL(acc, aL0, bL0); + + set_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + wait_flag(PIPE_M, PIPE_FIX, EVENT_ID0); + +#if USE_STRIP_C2V + for (int row = 0; row < M; row += H) { + AccPushT strip; + TASSIGN(strip, static_cast(row) * kAccRowByteStride); + TPUSH(pipe, strip); + } +#else + { + AccPushT full; + TASSIGN(full, 0x0); + TPUSH(pipe, full); + } +#endif + + set_flag(PIPE_FIX, PIPE_S, EVENT_ID7); + wait_flag(PIPE_FIX, PIPE_S, EVENT_ID7); + } + + if constexpr (DAV_VEC) { + if (get_sub_block_id(args) != 0) { + return; + } + + __gm__ float *c_ptr = + reinterpret_cast<__gm__ float *>(c_tensor->buffer.addr) + c_tensor->start_offset; + + constexpr uint64_t kScratch = static_cast(PP_FIFO_DEPTH) * kSlotBytes; + VecPushT vec; + TASSIGN(vec, kScratch); + + for (int s = 0; s < kNumPush; ++s) { + TPOP(pipe, vec); + + int row0 = s * PUSH_ROWS; + using GlobalC = GlobalTensor, + pto::Stride>; + GlobalC cGlobal(c_ptr + static_cast(row0) * N); + + set_flag(PIPE_S, PIPE_MTE3, EVENT_ID0); + wait_flag(PIPE_S, PIPE_MTE3, EVENT_ID0); + TSTORE(cGlobal, vec); + TFREE(pipe); + + set_flag(PIPE_MTE3, PIPE_S, EVENT_ID1); + wait_flag(PIPE_MTE3, PIPE_S, EVENT_ID1); + } + } +} diff --git a/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/kernels/orchestration/acc_c2v_compare_orch.cpp b/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/kernels/orchestration/acc_c2v_compare_orch.cpp new file mode 100644 index 0000000000..54cfc64a33 --- /dev/null +++ b/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/kernels/orchestration/acc_c2v_compare_orch.cpp @@ -0,0 +1,53 @@ +/* + * Orchestration: A, B, C_full, C_strip — one MixedKernels. + */ +#include +#include "pto_orchestration_api.h" // NOLINT(build/include_subdir) + +#define FUNC_AIC 0 +#define FUNC_AIV 1 + +static constexpr int M = 32; +static constexpr int N = 128; +static constexpr int K = 32; +static constexpr uint32_t A_ELEMS = static_cast(M * K); +static constexpr uint32_t B_ELEMS = static_cast(K * N); +static constexpr uint32_t C_ELEMS = static_cast(M * N); + +extern "C" { + +__attribute__((visibility("default"))) PTO2OrchestrationConfig aicpu_orchestration_config(const L2TaskArgs &orch_args) { + (void)orch_args; + return PTO2OrchestrationConfig{.expected_arg_count = 4}; +} + +__attribute__((visibility("default"))) void aicpu_orchestration_entry(const L2TaskArgs &orch_args) { + const Tensor &ext_A = orch_args.tensor(0).ref(); + const Tensor &ext_B = orch_args.tensor(1).ref(); + const Tensor &ext_Cf = orch_args.tensor(2).ref(); + const Tensor &ext_Cs = orch_args.tensor(3).ref(); + + uint32_t a_shapes[1] = {A_ELEMS}; + uint32_t b_shapes[1] = {B_ELEMS}; + uint32_t c_shapes[1] = {C_ELEMS}; + uint32_t zero[1] = {0}; + + Tensor A_view = ext_A.view(a_shapes, zero); + Tensor B_view = ext_B.view(b_shapes, zero); + Tensor Cf_view = ext_Cf.view(c_shapes, zero); + Tensor Cs_view = ext_Cs.view(c_shapes, zero); + + L0TaskArgs args; + args.add_input(A_view); + args.add_input(B_view); + args.add_output(Cf_view); + args.add_output(Cs_view); + + MixedKernels mk; + mk.aic_kernel_id = FUNC_AIC; + mk.aiv0_kernel_id = FUNC_AIV; + mk.aiv1_kernel_id = FUNC_AIV; + rt_submit_task(mk, args); +} + +} // extern "C" diff --git a/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/kernels/orchestration/acc_c2v_strip_orch.cpp b/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/kernels/orchestration/acc_c2v_strip_orch.cpp new file mode 100644 index 0000000000..e6670d24a7 --- /dev/null +++ b/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/kernels/orchestration/acc_c2v_strip_orch.cpp @@ -0,0 +1,71 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ +/** + * Orchestration for Acc ValidShape strip C2V repro. + * One MixedKernels task: C = A @ B via Acc strip TPUSH/TPOP. + * Arg layout: [A, B, C] + */ + +#include + +#include "pto_orchestration_api.h" // NOLINT(build/include_subdir) + +#define FUNC_AIC 0 +#define FUNC_AIV 1 + +static constexpr int M = 128; +static constexpr int N = 256; +static constexpr int K = 32; + +static constexpr uint32_t A_ELEMS = static_cast(M * K); +static constexpr uint32_t B_ELEMS = static_cast(K * N); +static constexpr uint32_t C_ELEMS = static_cast(M * N); + +extern "C" { + +__attribute__((visibility("default"))) PTO2OrchestrationConfig aicpu_orchestration_config(const L2TaskArgs &orch_args) { + (void)orch_args; + return PTO2OrchestrationConfig{ + .expected_arg_count = 3, + }; +} + +__attribute__((visibility("default"))) void aicpu_orchestration_entry(const L2TaskArgs &orch_args) { + const Tensor &ext_A = orch_args.tensor(0).ref(); + const Tensor &ext_B = orch_args.tensor(1).ref(); + const Tensor &ext_C = orch_args.tensor(2).ref(); + + LOG_INFO_V0("[acc_c2v_strip] M=%d N=%d K=%d Acc ValidShape strip C2V repro", M, N, K); + + uint32_t a_shapes[1] = {A_ELEMS}; + uint32_t b_shapes[1] = {B_ELEMS}; + uint32_t c_shapes[1] = {C_ELEMS}; + uint32_t zero[1] = {0}; + + Tensor A_view = ext_A.view(a_shapes, zero); + Tensor B_view = ext_B.view(b_shapes, zero); + Tensor C_view = ext_C.view(c_shapes, zero); + + L0TaskArgs args; + args.add_input(A_view); + args.add_input(B_view); + args.add_output(C_view); + + MixedKernels mk; + mk.aic_kernel_id = FUNC_AIC; + mk.aiv0_kernel_id = FUNC_AIV; + mk.aiv1_kernel_id = FUNC_AIV; + rt_submit_task(mk, args); + + LOG_INFO_V0("[acc_c2v_strip] submitted 1 MixedKernels (AIC+AIV strip C2V)"); +} + +} // extern "C" diff --git a/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/test_acc_c2v_full_control.py b/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/test_acc_c2v_full_control.py new file mode 100644 index 0000000000..8ce3c826ec --- /dev/null +++ b/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/test_acc_c2v_full_control.py @@ -0,0 +1,80 @@ +#!/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. +# ----------------------------------------------------------------------------------------------------------- +"""CONTROL: full Acc C2V (Valid≡Rows) — should PASS with or without the ISA fix. + +Same matmul shapes as the strip repro; only the Acc→Vec transfer differs +(one full TPUSH vs ValidShape strip windows). +""" + +import torch +from simpler.task_interface import ArgDirection as D + +from simpler_setup import SceneTestCase, TaskArgsBuilder, Tensor, scene_test + +M, N, K = 128, 256, 32 + + +@scene_test(level=2, runtime="tensormap_and_ringbuffer") +class TestAccC2VFullControl(SceneTestCase): + RTOL = 5e-2 + ATOL = 0.5 + + CALLABLE = { + "orchestration": { + "source": "kernels/orchestration/acc_c2v_strip_orch.cpp", + "function_name": "aicpu_orchestration_entry", + "signature": [D.IN, D.IN, D.OUT], + }, + "incores": [ + { + "func_id": 0, + "name": "ACC_C2V_FULL_AIC", + "source": "kernels/mix/kernel_acc_c2v_full.cpp", + "core_type": "aic", + "signature": [D.IN, D.IN, D.OUT], + }, + { + "func_id": 1, + "name": "ACC_C2V_FULL_AIV", + "source": "kernels/mix/kernel_acc_c2v_full.cpp", + "core_type": "aiv", + "signature": [D.IN, D.IN, D.OUT], + }, + ], + } + + CASES = [ + { + "name": "default", + "platforms": ["a5"], + "config": {"aicpu_thread_num": 4, "block_dim": 1}, + "params": {}, + } + ] + + def generate_args(self, params): + torch.manual_seed(0) + A = torch.randn(M, K, dtype=torch.float32) * 1.0 + B = torch.randn(K, N, dtype=torch.float32) * 1.0 + C = torch.zeros(M, N, dtype=torch.float32) + return TaskArgsBuilder( + Tensor("A", A.reshape(-1)), + Tensor("B", B.reshape(-1)), + Tensor("C", C.reshape(-1)), + ) + + def compute_golden(self, args, params): + A = args.A.reshape(M, K) + B = args.B.reshape(K, N) + args.C[:] = torch.matmul(A, B).reshape(-1) + + +if __name__ == "__main__": + SceneTestCase.run_module(__name__) diff --git a/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/test_acc_c2v_strip_validshape.py b/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/test_acc_c2v_strip_validshape.py new file mode 100644 index 0000000000..b98e7289b3 --- /dev/null +++ b/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/test_acc_c2v_strip_validshape.py @@ -0,0 +1,99 @@ +#!/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. +# ----------------------------------------------------------------------------------------------------------- +"""A5 Acc ValidShape strip C2V repro (pto-isa TMovCcToUb srcStride). + +Shapes: Acc[M=128,N=256] drained as 8× ValidShape(H=16,N) TPUSH with +addr = row * 64. Golden is float32 C = A @ B (K=64). + +Expect: + - unfixed pto-isa (srcStride=align(validRow)): FAIL (~90%+ mistmatch after strip0) + - fixed pto-isa (srcStride=align(Rows)): PASS + +Run (from simpler repo root, with installed simpler package):: + + # Fixed ISA (e.g. pypto/build/pto-isa @ 0ebbd03d): + PTO_ISA_ROOT=/path/to/fixed/pto-isa \\ + python examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/test_acc_c2v_strip_validshape.py -p a5 -d 0 + + # Unfixed ISA (checkout before srcStride fix): + PTO_ISA_ROOT=/path/to/unfixed/pto-isa \\ + python .../test_acc_c2v_strip_validshape.py -p a5 -d 0 +""" + +import torch +from simpler.task_interface import ArgDirection as D + +from simpler_setup import SceneTestCase, TaskArgsBuilder, Tensor, scene_test + +M, N, K = 128, 256, 32 +H = 16 + + +@scene_test(level=2, runtime="tensormap_and_ringbuffer") +class TestAccC2VStripValidShape(SceneTestCase): + """Strip Acc C2V vs matmul golden — fails before ISA srcStride=Rows fix.""" + + # Match issue / board thresholds used for Acc C2V strip diagnosis. + RTOL = 5e-2 + ATOL = 0.5 + + CALLABLE = { + "orchestration": { + "source": "kernels/orchestration/acc_c2v_strip_orch.cpp", + "function_name": "aicpu_orchestration_entry", + "signature": [D.IN, D.IN, D.OUT], + }, + "incores": [ + { + "func_id": 0, + "name": "ACC_C2V_STRIP_AIC", + "source": "kernels/mix/kernel_acc_c2v_strip.cpp", + "core_type": "aic", + "signature": [D.IN, D.IN, D.OUT], + }, + { + "func_id": 1, + "name": "ACC_C2V_STRIP_AIV", + "source": "kernels/mix/kernel_acc_c2v_strip.cpp", + "core_type": "aiv", + "signature": [D.IN, D.IN, D.OUT], + }, + ], + } + + CASES = [ + { + "name": "default", + "platforms": ["a5"], + "config": {"aicpu_thread_num": 4, "block_dim": 1}, + "params": {}, + } + ] + + def generate_args(self, params): + # Scale down so Acc values stay moderate for atol/rtol. + torch.manual_seed(0) + A = torch.randn(M, K, dtype=torch.float32) * 1.0 + B = torch.randn(K, N, dtype=torch.float32) * 1.0 + C = torch.zeros(M, N, dtype=torch.float32) + return TaskArgsBuilder( + Tensor("A", A.reshape(-1)), + Tensor("B", B.reshape(-1)), + Tensor("C", C.reshape(-1)), + ) + + def compute_golden(self, args, params): + A = args.A.reshape(M, K) + B = args.B.reshape(K, N) + args.C[:] = torch.matmul(A, B).reshape(-1) + + +if __name__ == "__main__": + SceneTestCase.run_module(__name__) diff --git a/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/test_acc_c2v_strip_vs_full.py b/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/test_acc_c2v_strip_vs_full.py new file mode 100644 index 0000000000..228236679d --- /dev/null +++ b/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/test_acc_c2v_strip_vs_full.py @@ -0,0 +1,124 @@ +#!/usr/bin/env python3 +# Copyright (c) PyPTO Contributors. +"""A5 Acc ValidShape strip C2V repro — device C_strip vs device C_full. + +Drains the same Acc twice (full Valid≡Rows vs ValidShape strip windows). +After the run we require C_strip ≈ C_full. + + unfixed ISA (srcStride=validRow): FAIL (often strip0 OK, strip1+ bad) + fixed ISA (srcStride=Rows): PASS +""" + +import torch +from simpler.task_interface import ArgDirection as D + +from simpler_setup import SceneTestCase, TaskArgsBuilder, Tensor, scene_test +from simpler_setup.scene_test import _build_chip_task_args, _temporary_env + +M, N, K = 32, 128, 32 +H = 16 + + +@scene_test(level=2, runtime="tensormap_and_ringbuffer") +class TestAccC2VStripVsFull(SceneTestCase): + RTOL = 1e-4 + ATOL = 1e-4 + + CALLABLE = { + "orchestration": { + "source": "kernels/orchestration/acc_c2v_compare_orch.cpp", + "function_name": "aicpu_orchestration_entry", + "signature": [D.IN, D.IN, D.OUT, D.OUT], + }, + "incores": [ + { + "func_id": 0, + "name": "ACC_C2V_COMPARE_AIC", + "source": "kernels/mix/kernel_acc_c2v_compare.cpp", + "core_type": "aic", + "signature": [D.IN, D.IN, D.OUT, D.OUT], + }, + { + "func_id": 1, + "name": "ACC_C2V_COMPARE_AIV", + "source": "kernels/mix/kernel_acc_c2v_compare.cpp", + "core_type": "aiv", + "signature": [D.IN, D.IN, D.OUT, D.OUT], + }, + ], + } + + CASES = [ + { + "name": "default", + "platforms": ["a5"], + "config": {"aicpu_thread_num": 4, "block_dim": 1}, + "params": {}, + } + ] + + def generate_args(self, params): + torch.manual_seed(0) + A = torch.randn(M, K, dtype=torch.float32) + B = torch.randn(K, N, dtype=torch.float32) + return TaskArgsBuilder( + Tensor("A", A.reshape(-1)), + Tensor("B", B.reshape(-1)), + Tensor("C_full", torch.zeros(M * N, dtype=torch.float32)), + Tensor("C_strip", torch.zeros(M * N, dtype=torch.float32)), + ) + + def compute_golden(self, args, params): + pass + + def _run_and_validate_l2( # noqa: PLR0913 + self, + worker, + callable_obj, + case, + rounds=1, + skip_golden=False, + enable_l2_swimlane=0, + enable_dump_args=False, + enable_pmu=0, + enable_dep_gen=False, + enable_scope_stats=False, + output_prefix="", + ): + del rounds, skip_golden # single-shot self-compare + params = case.get("params", {}) + config_dict = case.get("config", {}) + orch_sig = self.CALLABLE.get("orchestration", {}).get("signature", []) + handle = getattr(type(self), "_st_l2_handle", None) + if handle is None: + handle = worker.register(callable_obj) + type(self)._st_l2_handle = handle + + test_args = self.generate_args(params) + chip_args, _ = _build_chip_task_args(test_args, orch_sig) + config = self._build_config( + config_dict, + enable_l2_swimlane=enable_l2_swimlane, + enable_dump_args=enable_dump_args, + enable_pmu=enable_pmu, + enable_dep_gen=enable_dep_gen, + enable_scope_stats=enable_scope_stats, + output_prefix=output_prefix, + ) + with _temporary_env(self._resolve_env()): + worker.run(handle, chip_args, config=config) + + full = test_args.C_full.reshape(M, N) + strip = test_args.C_strip.reshape(M, N) + if not torch.allclose(strip, full, rtol=self.RTOL, atol=self.ATOL): + diff = (strip - full).abs().max().item() + s0 = (strip[:H] - full[:H]).abs().max().item() + s1 = (strip[H : 2 * H] - full[H : 2 * H]).abs().max().item() + raise AssertionError( + f"C_strip vs C_full mismatch: max_diff={diff}, " + f"strip0_max={s0}, strip1_max={s1}, rtol={self.RTOL}, atol={self.ATOL}" + ) + + +if __name__ == "__main__": + SceneTestCase.run_module(__name__) diff --git a/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/test_acc_matmul_tstore.py b/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/test_acc_matmul_tstore.py new file mode 100644 index 0000000000..7e19a421be --- /dev/null +++ b/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/test_acc_matmul_tstore.py @@ -0,0 +1,53 @@ +#!/usr/bin/env python3 +"""CONTROL: AIC-only Acc matmul + TSTORE (no C2V).""" +import torch +from simpler.task_interface import ArgDirection as D +from simpler_setup import SceneTestCase, TaskArgsBuilder, Tensor, scene_test + +M, N, K = 128, 256, 32 + + +@scene_test(level=2, runtime="tensormap_and_ringbuffer") +class TestAccMatmulTstore(SceneTestCase): + RTOL = 5e-2 + ATOL = 0.5 + + CALLABLE = { + "orchestration": { + "source": "kernels/orchestration/acc_c2v_strip_orch.cpp", + "function_name": "aicpu_orchestration_entry", + "signature": [D.IN, D.IN, D.OUT], + }, + "incores": [ + { + "func_id": 0, + "name": "MATMUL_TSTORE", + "source": "kernels/aic/kernel_matmul_tstore.cpp", + "core_type": "aic", + "signature": [D.IN, D.IN, D.OUT], + }, + ], + } + + CASES = [ + { + "name": "default", + "platforms": ["a5"], + "config": {"aicpu_thread_num": 4, "block_dim": 1}, + "params": {}, + } + ] + + def generate_args(self, params): + torch.manual_seed(0) + A = torch.randn(M, K, dtype=torch.float32) * 1.0 + B = torch.randn(K, N, dtype=torch.float32) * 1.0 + C = torch.zeros(M, N, dtype=torch.float32) + return TaskArgsBuilder(Tensor("A", A.reshape(-1)), Tensor("B", B.reshape(-1)), Tensor("C", C.reshape(-1))) + + def compute_golden(self, args, params): + args.C[:] = torch.matmul(args.A.reshape(M, K), args.B.reshape(K, N)).reshape(-1) + + +if __name__ == "__main__": + SceneTestCase.run_module(__name__) diff --git a/tests/st/a5/tensormap_and_ringbuffer/spmd_early_dispatch/kernels/orchestration/spmd_early_dispatch_orch.cpp b/tests/st/a5/tensormap_and_ringbuffer/spmd_early_dispatch/kernels/orchestration/spmd_early_dispatch_orch.cpp deleted file mode 100644 index 52e3f05003..0000000000 --- a/tests/st/a5/tensormap_and_ringbuffer/spmd_early_dispatch/kernels/orchestration/spmd_early_dispatch_orch.cpp +++ /dev/null @@ -1,87 +0,0 @@ -/* - * 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 Early-Dispatch Orchestration (a5) — plain early_resolve, NO sync_start - * - * Exercises the shape-queued early path (`early_dispatch_queues[]`), not the - * sync_start cohort lane (`early_sync_start_queue`). - * - * Topology is sized so the producer leaves spare AIC cores: while P spins on - * 8 AICs, C can gated-stage onto the remaining idle AICs. - * - * P: AIC core_num=8, base_cl=0, allow_early_resolve = scalar early_on, spin - * C: AIC core_num=8, base_cl=8, dep=[P], require_sync_start=false - * - * 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 - -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 #1079-era spin: long enough that C can stage while P is still on-core. -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_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); // no spin - args.launch_spec.set_core_num(core_num); - // deliberately NOT require_sync_start — plain early_dispatch_queues path - args.add_dep(dep); - rt_submit_aic_task(FUNC_SPMD_WRITE_AIC, args); -} - -__attribute__((visibility("default"))) void aicpu_orchestration_entry(const L2TaskArgs &orch_args) { - const Tensor &ext_output = orch_args.tensor(0).ref(); - // Host Scalar("early_on") is orch scalar(0). Also stamp a sentinel into the - // output so board runs can prove the scalar reached the orch (out[1] = early_on). - const bool early_on = orch_args.scalar(0) != 0; - { - // out layout: CL0..7 producer, CL8..15 consumer; float index 1 is unused by kernels. - float *out = reinterpret_cast(ext_output.buffer.addr) + ext_output.start_offset; - out[1] = early_on ? 1.0f : 0.0f; - } - - rt_scope_begin(PTO2ScopeMode::MANUAL); - PTO2TaskId prod = submit_producer(ext_output, 8, 0, early_on); - submit_consumer(ext_output, 8, 8, prod); - rt_scope_end(); - - LOG_INFO_V9("[spmd_early_dispatch] early_on=%d AIC producer + AIC consumer (no sync_start)", early_on ? 1 : 0); -} - -} // extern "C" diff --git a/tests/st/a5/tensormap_and_ringbuffer/spmd_early_dispatch/test_spmd_early_dispatch.py b/tests/st/a5/tensormap_and_ringbuffer/spmd_early_dispatch/test_spmd_early_dispatch.py deleted file mode 100644 index 2cbb473873..0000000000 --- a/tests/st/a5/tensormap_and_ringbuffer/spmd_early_dispatch/test_spmd_early_dispatch.py +++ /dev/null @@ -1,83 +0,0 @@ -#!/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. -# ----------------------------------------------------------------------------------------------------------- -"""Plain allow_early_resolve ST (a5) — NO require_sync_start. - -Producer leaves spare AIC cores so the consumer can stage via -``early_dispatch_queues`` while the producer is still spinning. EarlyOn / -EarlyOff toggle the producer flag for swimlane comparison. -""" - -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 = 8 -CONSUMER_BLOCKS = 8 -CONSUMER_BASE_CL = PRODUCER_BLOCKS -TOTAL_CL = CONSUMER_BASE_CL + CONSUMER_BLOCKS - - -@scene_test(level=2, runtime="tensormap_and_ringbuffer") -class TestSpmdEarlyDispatch(SceneTestCase): - RTOL = 0 - ATOL = 0 - - CALLABLE = { - "orchestration": { - "source": "kernels/orchestration/spmd_early_dispatch_orch.cpp", - "function_name": "aicpu_orchestration_entry", - "signature": [D.INOUT], - }, - "incores": [ - { - "func_id": 0, - "name": "SPMD_WRITE_AIC", - "source": "../spmd_sync_start_early_dispatch/kernels/aiv/kernel_spmd_write_slow.cpp", - "core_type": "aic", - "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 - # orch stamps early_on into out[1] as a host-side probe (kernels leave it alone). - out[1] = float(int(params.get("early_on", 1))) - for block_idx in range(PRODUCER_BLOCKS): - out[block_idx * FLOATS_PER_CACHE_LINE] = float(block_idx) - for block_idx in range(CONSUMER_BLOCKS): - out[(CONSUMER_BASE_CL + block_idx) * 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) { From 08f67b61fd603ef6a2e5b365bdf9d14b9b3e96ce Mon Sep 17 00:00:00 2001 From: yanghaoran29 Date: Tue, 21 Jul 2026 21:09:44 +0800 Subject: [PATCH 3/3] chore(a5): drop unrelated files that broke pre-commit on the PR Remove accidental local helpers, migrate patches, and the acc_c2v example that were committed with the UT payload-pool fix. Keep the early-dispatch runtime port, a2a3-ported sync_start STs, and a5 UT payload/task pool binding. --- _env_dump.sh | 2 - _migrate_patches/early_phase0b.diff | 552 ---------- _migrate_patches/pto_scheduler_partial.h | 993 ------------------ _run_early_st_board.sh | 22 - _run_st_dev2.sh | 9 - .../acc_c2v_strip_validshape/README.md | 62 -- .../kernels/aic/kernel_matmul_tstore.cpp | 69 -- .../kernels/mix/kernel_acc_c2v_compare.cpp | 194 ---- .../kernels/mix/kernel_acc_c2v_full.cpp | 192 ---- .../kernels/mix/kernel_acc_c2v_strip.cpp | 192 ---- .../orchestration/acc_c2v_compare_orch.cpp | 53 - .../orchestration/acc_c2v_strip_orch.cpp | 71 -- .../test_acc_c2v_full_control.py | 80 -- .../test_acc_c2v_strip_validshape.py | 99 -- .../test_acc_c2v_strip_vs_full.py | 124 --- .../test_acc_matmul_tstore.py | 53 - 16 files changed, 2767 deletions(-) delete mode 100755 _env_dump.sh delete mode 100644 _migrate_patches/early_phase0b.diff delete mode 100644 _migrate_patches/pto_scheduler_partial.h delete mode 100755 _run_early_st_board.sh delete mode 100755 _run_st_dev2.sh delete mode 100644 examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/README.md delete mode 100644 examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/kernels/aic/kernel_matmul_tstore.cpp delete mode 100644 examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/kernels/mix/kernel_acc_c2v_compare.cpp delete mode 100644 examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/kernels/mix/kernel_acc_c2v_full.cpp delete mode 100644 examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/kernels/mix/kernel_acc_c2v_strip.cpp delete mode 100644 examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/kernels/orchestration/acc_c2v_compare_orch.cpp delete mode 100644 examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/kernels/orchestration/acc_c2v_strip_orch.cpp delete mode 100644 examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/test_acc_c2v_full_control.py delete mode 100644 examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/test_acc_c2v_strip_validshape.py delete mode 100644 examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/test_acc_c2v_strip_vs_full.py delete mode 100644 examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/test_acc_matmul_tstore.py diff --git a/_env_dump.sh b/_env_dump.sh deleted file mode 100755 index 6519a196f2..0000000000 --- a/_env_dump.sh +++ /dev/null @@ -1,2 +0,0 @@ -#!/bin/bash -env | sort diff --git a/_migrate_patches/early_phase0b.diff b/_migrate_patches/early_phase0b.diff deleted file mode 100644 index e1bd555f26..0000000000 --- a/_migrate_patches/early_phase0b.diff +++ /dev/null @@ -1,552 +0,0 @@ -diff --git a/src/a5/platform/onboard/aicore/inner_kernel.h b/src/a5/platform/onboard/aicore/inner_kernel.h -index 2f429603..96d8c3d6 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 58dd4488..8b997b28 100644 ---- a/src/a5/platform/sim/aicore/inner_kernel.h -+++ b/src/a5/platform/sim/aicore/inner_kernel.h -@@ -176,6 +176,17 @@ 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. -+ */ -+inline uint32_t read_dmb_high32() { -+ uint32_t offset = reg_offset(RegId::DATA_MAIN_BASE); -+ uint32_t hi = *reinterpret_cast(sparse_reg_ptr(sim_get_reg_base(), offset + 4)); -+ OUT_OF_ORDER_LOAD_BARRIER(); -+ return hi; -+} -+ - /** - * Write to an AICore register in simulated register memory - * -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 20144d3a..00c36211 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 @@ public: - 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/pto_types.h b/src/a5/runtime/tensormap_and_ringbuffer/runtime/pto_types.h -index bd0cd66b..27829ffb 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/scheduler_completion.cpp b/src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_completion.cpp -index e1586227..4a5c13f9 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" -@@ -413,6 +415,7 @@ int32_t SchedulerContext::count_global_available(PTO2ResourceShape shape, uint8_ - // 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) { -+ (void)runtime; - 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); -@@ -425,9 +428,27 @@ void SchedulerContext::drain_worker_dispatch(Runtime *runtime, int32_t block_num - 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++; -+ int32_t remaining = slot_state->logical_block_num - slot_state->next_block_idx; -+ int32_t claim = std::min(valid.count(), remaining); -+ int32_t start = slot_state->next_block_idx; -+ slot_state->next_block_idx += claim; -+ PublishHandle handles[CoreTracker::MAX_CLUSTERS * 3]; -+ int handle_count = 0; -+ for (int32_t b = 0; b < claim; b++) { -+ auto core_offset = valid.pop_first(); -+ handle_count += prepare_block_for_dispatch( -+ t, core_offset, *slot_state, shape, false, start + b, &handles[handle_count] -+ ); -+ } -+ wmb(); -+ uint64_t dispatch_ts = 0; -+#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, t); - } - } - -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 c60340ef..9aa25a0d 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" -@@ -250,19 +252,62 @@ private: - const AsyncCtx &async_ctx, int32_t block_idx - ); - -- 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 - ); - -- 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); -+ } -+ // PTO2_DMB64_DISPATCH_PROBE: seed both halves of DATA_MAIN_BASE so AICore -+ // read_dmb_high32()/MOV can observe the doorbell half on silicon. Matches -+ // a2a3 ring_one_doorbell encoding `(tok<<32)|tok`. Remove after Phase-0 -+ // onboard probe (tests/st/a5/.../dmb_high32_probe) lands a verdict. -+#if defined(PTO2_DMB64_DISPATCH_PROBE) -+ { -+ volatile uint64_t *dmb = reinterpret_cast( -+ get_reg_ptr(h.reg_addr, RegId::DATA_MAIN_BASE) -+ ); -+ uint64_t tk = static_cast(h.reg_task_id); -+ *dmb = (tk << 32) | tk; -+ } -+#else -+ write_reg(h.reg_addr, RegId::DATA_MAIN_BASE, static_cast(h.reg_task_id)); -+#endif -+ } -+ -+ // 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 - ); - - void dispatch_shape( -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 8007169b..30c00e4e 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 -@@ -125,14 +125,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 - ) { - 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,6 +145,10 @@ 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; -@@ -161,6 +165,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 +187,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 - ) { - #if SIMPLER_DFX - if (is_dump_args_enabled()) { -@@ -258,26 +217,58 @@ 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 -+ ); -+ } -+ 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 -+ ); -+ } -+ 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 -+ ); -+ } -+#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); -+#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); - #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 +285,64 @@ 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; - 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 +370,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,21 +383,19 @@ 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 -+ // 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(); -@@ -365,13 +412,24 @@ 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(); - #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; - diff --git a/_migrate_patches/pto_scheduler_partial.h b/_migrate_patches/pto_scheduler_partial.h deleted file mode 100644 index eca3496924..0000000000 --- a/_migrate_patches/pto_scheduler_partial.h +++ /dev/null @@ -1,993 +0,0 @@ -/* - * 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. - * ----------------------------------------------------------------------------------------------------------- - */ - -/** - * PTO Runtime2 - Scheduler Interface - * - * The Scheduler is responsible for: - * 1. Maintaining per-resource-shape ready queues - * 2. Tracking task state (PENDING -> COMPLETED -> CONSUMED) - * 3. Managing fanin/fanout refcounts for dependency resolution - * 4. Advancing last_task_alive for heap reclamation - * 5. Two-stage mixed-task completion (subtask done bits → mixed-task complete) - * - * The Scheduler runs on Device AI_CPU and processes: - * - Task state transitions based on fanin_refcount - * - Buffer lifecycle based on fanout_refcount - * - Ring pointer advancement for flow control - * - * Based on: docs/RUNTIME_LOGIC.md - */ - -#pragma once - -#include - -#include "common/core_type.h" -#include "utils/device_arena.h" -#include "pto_async_wait.h" -#include "pto_ring_buffer.h" -#include "pto_runtime2_types.h" -#include "pto_shared_memory.h" - -#if SIMPLER_SCHED_PROFILING -#include "aicpu/device_time.h" -#define PTO2_SCHED_CYCLE_START() uint64_t _st0 = get_sys_cnt_aicpu(), _st1 -#define PTO2_SCHED_CYCLE_LAP(acc) \ - do { \ - _st1 = get_sys_cnt_aicpu(); \ - acc += (_st1 - _st0); \ - _st0 = _st1; \ - } while (0) -#endif - -// ============================================================================= -// Ready Queue (Lock-free bounded MPMC — Vyukov design) -// ============================================================================= - -/** - * Per-slot entry: sequence counter for ABA safety + task payload - */ -struct PTO2ReadyQueueSlot { - std::atomic sequence; - PTO2TaskSlotState *slot_state; - uint64_t task_id_snapshot; // generation tag for early-dispatch queue entries -}; - -/** - * Lock-free bounded MPMC queue (Dmitry Vyukov design) - * - * Key properties: - * - enqueue_pos and dequeue_pos on separate cache lines (no false sharing) - * - Per-slot sequence counter prevents ABA problem - * - Empty queue pop returns immediately (single atomic load, no lock) - * - CAS contention is split: producers only touch enqueue_pos, - * consumers only touch dequeue_pos - */ -struct alignas(64) PTO2ReadyQueue { - PTO2ReadyQueueSlot *slots; - uint64_t capacity; - uint64_t mask; // capacity - 1 - char _pad0[64 - 24]; // Pad to own cache line - - std::atomic enqueue_pos; - char _pad1[64 - sizeof(std::atomic)]; // Own cache line - - std::atomic dequeue_pos; - char _pad2[64 - sizeof(std::atomic)]; // Own cache line - - uint64_t size() { - uint64_t e = enqueue_pos.load(std::memory_order_relaxed); - uint64_t d = dequeue_pos.load(std::memory_order_relaxed); - return (e >= d) ? (e - d) : 0; - } - - void reset_for_reuse() {} - - bool push(PTO2TaskSlotState *slot_state) { return push_tagged(slot_state, 0); } - - bool push_tagged(PTO2TaskSlotState *slot_state, uint64_t task_id_snapshot) { - uint64_t pos; - PTO2ReadyQueueSlot *slot; - while (true) { - pos = enqueue_pos.load(std::memory_order_relaxed); - slot = &slots[pos & mask]; - int64_t seq = slot->sequence.load(std::memory_order_acquire); - int64_t diff = seq - static_cast(pos); - if (diff == 0) { - if (enqueue_pos.compare_exchange_weak( - pos, pos + 1, std::memory_order_relaxed, std::memory_order_relaxed - )) { - break; - } - } else if (diff < 0) { - return false; // Queue full - } - } - - slot->slot_state = slot_state; - slot->task_id_snapshot = task_id_snapshot; - slot->sequence.store(static_cast(pos + 1), std::memory_order_release); - return true; - } - - // Batch push: reserve count slots with a single CAS after confirming - // every target slot is available under the usual Vyukov sequence check. - void push_batch(PTO2TaskSlotState **items, int count) { push_batch_tagged(items, nullptr, count); } - - void push_batch_tagged(PTO2TaskSlotState **items, const uint64_t *task_id_snapshots, int count) { - if (count == 0) return; - - uint64_t pos; - while (true) { - pos = enqueue_pos.load(std::memory_order_relaxed); - bool ready = true; - for (int i = 0; i < count; i++) { - PTO2ReadyQueueSlot *slot = &slots[(pos + i) & mask]; - int64_t seq = slot->sequence.load(std::memory_order_acquire); - int64_t diff = seq - static_cast(pos + i); - if (diff != 0) { - ready = false; - break; - } - } - if (!ready) { - continue; - } - if (enqueue_pos.compare_exchange_weak( - pos, pos + count, std::memory_order_relaxed, std::memory_order_relaxed - )) { - break; - } - } - - for (int i = 0; i < count; i++) { - PTO2ReadyQueueSlot *slot = &slots[(pos + i) & mask]; - slot->slot_state = items[i]; - slot->task_id_snapshot = task_id_snapshots == nullptr ? 0 : task_id_snapshots[i]; - slot->sequence.store(static_cast(pos + i + 1), std::memory_order_release); - } - } - -#if SIMPLER_ORCH_PROFILING || SIMPLER_SCHED_PROFILING - bool push(PTO2TaskSlotState *slot_state, uint64_t &atomic_count, uint64_t &wait_cycle) { - uint64_t pos; - PTO2ReadyQueueSlot *slot; - uint64_t t0 = get_sys_cnt_aicpu(); - bool contended = false; - uint32_t atomic_ops = 0; - while (true) { - pos = enqueue_pos.load(std::memory_order_relaxed); - slot = &slots[pos & mask]; - int64_t seq = slot->sequence.load(std::memory_order_acquire); - int64_t diff = seq - static_cast(pos); - atomic_ops += 2; // enqueue_pos.load + sequence.load - if (diff == 0) { - if (enqueue_pos.compare_exchange_weak( - pos, pos + 1, std::memory_order_relaxed, std::memory_order_relaxed - )) { - atomic_ops++; // successful CAS - break; - } - contended = true; - atomic_ops++; // failed CAS - } else if (diff < 0) { - return false; // Queue full - } else { - contended = true; // diff > 0: slot not yet released, spin - } - } - atomic_ops++; // final sequence.store - atomic_count += atomic_ops; - if (contended) { - wait_cycle += (get_sys_cnt_aicpu() - t0); - } - - slot->slot_state = slot_state; - slot->task_id_snapshot = 0; - slot->sequence.store(static_cast(pos + 1), std::memory_order_release); - return true; - } -#endif - - PTO2TaskSlotState *pop() { return pop_tagged(nullptr); } - - PTO2TaskSlotState *pop_tagged(uint64_t *task_id_snapshot) { - // Fast-path: skip slot load when queue is clearly empty - uint64_t d = dequeue_pos.load(std::memory_order_relaxed); - uint64_t e = enqueue_pos.load(std::memory_order_relaxed); - if (d >= e) { - return nullptr; - } - - uint64_t pos; - PTO2ReadyQueueSlot *slot; - while (true) { - pos = dequeue_pos.load(std::memory_order_relaxed); - slot = &slots[pos & mask]; - int64_t seq = slot->sequence.load(std::memory_order_acquire); - int64_t diff = seq - static_cast(pos + 1); - if (diff == 0) { - if (dequeue_pos.compare_exchange_weak( - pos, pos + 1, std::memory_order_relaxed, std::memory_order_relaxed - )) - break; - } else if (diff < 0) { - return nullptr; // Queue empty - } - } - - PTO2TaskSlotState *result = slot->slot_state; - slot->sequence.store(static_cast(pos + mask + 1), std::memory_order_release); - return result; - } - -#if SIMPLER_SCHED_PROFILING - PTO2TaskSlotState *pop(uint64_t &atomic_count, uint64_t &wait_cycle) { - // Fast-path: skip slot load when queue is clearly empty - uint64_t d = dequeue_pos.load(std::memory_order_relaxed); - uint64_t e = enqueue_pos.load(std::memory_order_relaxed); - atomic_count += 2; // dequeue_pos.load + enqueue_pos.load - if (d >= e) { - return nullptr; - } - - uint64_t pos; - PTO2ReadyQueueSlot *slot; - uint64_t t0 = get_sys_cnt_aicpu(); - bool contended = false; - uint32_t atomic_ops = 0; - while (true) { - pos = dequeue_pos.load(std::memory_order_relaxed); - slot = &slots[pos & mask]; - int64_t seq = slot->sequence.load(std::memory_order_acquire); - int64_t diff = seq - static_cast(pos + 1); - atomic_ops += 2; // dequeue_pos.load + sequence.load - if (diff == 0) { - if (dequeue_pos.compare_exchange_weak( - pos, pos + 1, std::memory_order_relaxed, std::memory_order_relaxed - )) { - atomic_ops++; // successful CAS - break; - } - contended = true; - atomic_ops++; // failed CAS - } else if (diff < 0) { - atomic_count += atomic_ops; - return nullptr; // Queue empty - } else { - contended = true; - } - } - atomic_ops++; // final sequence.store - atomic_count += atomic_ops; - if (contended) { - wait_cycle += (get_sys_cnt_aicpu() - t0); - } - - PTO2TaskSlotState *result = slot->slot_state; - slot->sequence.store(static_cast(pos + mask + 1), std::memory_order_release); - return result; - } -#endif - - // Batch pop: reserve a contiguous run of ready slots with a single CAS. - // Returns actual number of items popped (may be less than max_count). - int pop_batch(PTO2TaskSlotState **out, int max_count) { - uint64_t pos; - int count; - while (true) { - pos = dequeue_pos.load(std::memory_order_relaxed); - count = 0; - while (count < max_count) { - PTO2ReadyQueueSlot *slot = &slots[(pos + count) & mask]; - int64_t seq = slot->sequence.load(std::memory_order_acquire); - int64_t diff = seq - static_cast(pos + count + 1); - if (diff == 0) { - count++; - continue; - } - if (diff < 0) { - break; - } - count = -1; - break; - } - if (count == 0) return 0; - if (count < 0) continue; - if (dequeue_pos.compare_exchange_weak( - pos, pos + count, std::memory_order_relaxed, std::memory_order_relaxed - )) { - break; - } - } - - for (int i = 0; i < count; i++) { - PTO2ReadyQueueSlot *slot = &slots[(pos + i) & mask]; - out[i] = slot->slot_state; - slot->sequence.store(static_cast(pos + i + mask + 1), std::memory_order_release); - } - return count; - } - -#if SIMPLER_SCHED_PROFILING - int pop_batch(PTO2TaskSlotState **out, int max_count, uint64_t &atomic_count, uint64_t &wait_cycle) { - uint64_t pos; - int count; - uint64_t t0 = get_sys_cnt_aicpu(); - bool contended = false; - uint32_t atomic_ops = 0; - while (true) { - pos = dequeue_pos.load(std::memory_order_relaxed); - atomic_ops++; // dequeue_pos.load - count = 0; - while (count < max_count) { - PTO2ReadyQueueSlot *slot = &slots[(pos + count) & mask]; - int64_t seq = slot->sequence.load(std::memory_order_acquire); - int64_t diff = seq - static_cast(pos + count + 1); - atomic_ops++; // sequence.load - if (diff == 0) { - count++; - continue; - } - if (diff < 0) { - break; - } - contended = true; - count = -1; - break; - } - if (count == 0) { - atomic_count += atomic_ops; - return 0; - } - if (count < 0) { - continue; - } - if (dequeue_pos.compare_exchange_weak( - pos, pos + count, std::memory_order_relaxed, std::memory_order_relaxed - )) { - atomic_ops++; // successful CAS - break; - } - contended = true; - atomic_ops++; // failed CAS - } - - for (int i = 0; i < count; i++) { - PTO2ReadyQueueSlot *slot = &slots[(pos + i) & mask]; - out[i] = slot->slot_state; - slot->sequence.store(static_cast(pos + i + mask + 1), std::memory_order_release); - atomic_ops++; // sequence.store - } - atomic_count += atomic_ops; - if (contended) { - wait_cycle += (get_sys_cnt_aicpu() - t0); - } - return count; - } -#endif -}; - -// Cold-path ready queue operations (defined in pto_scheduler.cpp). Declared -// as non-member so PTO2ReadyQueue stays a POD-like struct with cache-line -// alignment. Storage is owned by the caller-supplied arena. -// reserve_layout: declare the slots[] region on the arena (must precede commit) -// init_from_layout: bind slots pointer from arena.region_ptr(off) and -// initialize sequence counters -// destroy: forget the slots pointer (arena owns the buffer) -size_t ready_queue_reserve_layout(DeviceArena &arena, uint64_t capacity); -// Writes everything *except* the arena-internal `slots` pointer field -// (sequences/positions on the slot array, capacity, mask). Uses -// arena.region_ptr(slots_off) only to address the slot array for writes; -// does NOT store the pointer in `queue->slots`. Call -// `ready_queue_wire_arena_pointers` afterwards to set the field itself. -bool ready_queue_init_data_from_layout(PTO2ReadyQueue *queue, DeviceArena &arena, size_t slots_off, uint64_t capacity); -// Stores queue->slots = arena.region_ptr(slots_off). Idempotent. -void ready_queue_wire_arena_pointers(PTO2ReadyQueue *queue, DeviceArena &arena, size_t slots_off); -void ready_queue_destroy(PTO2ReadyQueue *queue); - -/** - * Statistics returned by mixed-task completion processing - */ -struct CompletionStats { - int32_t fanout_edges; // Number of fanout edges traversed (notify consumers) - int32_t tasks_enqueued; // Number of consumers that became READY - int32_t fanin_edges; // Number of fanin edges traversed (release producers) - bool mixed_task_completed; // True only when this callback completed a mixed task -}; - -/** - * Layout descriptor produced by PTO2SchedulerState::reserve_layout(). Holds - * the arena offsets of every sub-region the scheduler needs plus the - * capacities used at layout time (init_from_layout reuses them). - */ -struct PTO2SchedulerLayout { - size_t off_ready_queue_slots[PTO2_NUM_RESOURCE_SHAPES]; - size_t off_dummy_ready_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]; -}; - -/** - * Scheduler state structure - * - * Contains dynamic state updated during task execution. - * Separated from shared memory for cache efficiency. - * Hot-path methods are defined inline (implicitly inline as member functions). - */ -struct PTO2SchedulerState { - // Shared memory access - PTO2SharedMemoryHeader *sm_header; - - // Per-ring state - struct alignas(64) RingSchedState { - // --- Cache Line 0: ring pointer (read-only) + hot path (read-write) --- - PTO2SharedMemoryRingHeader *ring; - int32_t last_task_alive; - std::atomic advance_lock; // multi-thread CAS - - // --- Cache Line 1+: Orch-side wiring dep_pool --- - alignas(64) PTO2DepListPool dep_pool; -#if SIMPLER_DFX - // Published only for scope_stats; orchestrator must not read dep_pool's non-atomic counters directly. - alignas(64) std::atomic dep_pool_snapshot_tail; - std::atomic dep_pool_snapshot_top; -#endif - - // Initialize arena-internal data + arena-external pointers; does NOT - // store dep_pool.base (that lives in the runtime arena and is wired - // by SchedulerState::wire_arena_pointers). The `ring` field stores - // the device address of the SM ring header — computed via offset - // arithmetic, no SM dereference. - bool init_data_from_layout(void *sm_dev_base, int32_t ring_id); - void reset_for_reuse(void *sm_dev_base, int32_t ring_id, std::atomic *orch_err); - void destroy(); - - void sync_to_sm() { ring->fc.last_task_alive.store(last_task_alive, std::memory_order_release); } - -#if SIMPLER_DFX - void publish_dep_pool_snapshot() { - dep_pool_snapshot_tail.store(dep_pool.tail, std::memory_order_release); - dep_pool_snapshot_top.store(dep_pool.top, std::memory_order_release); - } - - void read_dep_pool_snapshot(int32_t &tail, int32_t &top) const { - top = dep_pool_snapshot_top.load(std::memory_order_acquire); - tail = dep_pool_snapshot_tail.load(std::memory_order_acquire); - if (tail > top) tail = top; - } -#endif - - void advance_ring_pointers() { - int32_t current_task_index = ring->fc.current_task_index.load(std::memory_order_acquire); - int32_t old_last_task_alive = last_task_alive; - - while (last_task_alive < current_task_index) { - PTO2TaskSlotState &slot_state = ring->get_slot_state_by_task_id(last_task_alive); - if (slot_state.task_state.load(std::memory_order_acquire) != PTO2_TASK_CONSUMED) { - break; - } - last_task_alive++; - } - - // Eager reset: prepare reclaimed slots for reuse while still hot in cache. - // Safe because last_task_alive has advanced past these slots but - // sync_to_sm has not yet published — the orchestrator cannot reuse - // them until the release store below. - // Skips payload, task, ring_id — immutable after RingSchedState::init(). - for (int32_t id = old_last_task_alive; id < last_task_alive; id++) { - ring->get_slot_state_by_task_id(id).reset_for_reuse(); - } - - sync_to_sm(); - } - } ring_sched_states[PTO2_MAX_RING_DEPTH]; - - // Ready queues remain global (scheduling is ring-agnostic) - PTO2ReadyQueue ready_queues[PTO2_NUM_RESOURCE_SHAPES]; - - // Dependency-only tasks (active_mask is empty, shape == DUMMY). Drained by - // the dispatch loop and completed inline -- never goes to AICore. - PTO2ReadyQueue dummy_ready_queue; - - alignas(64) AsyncWaitList async_wait_list; - - // Statistics (cold path, isolated from hot-path fields) -#if SIMPLER_SCHED_PROFILING - alignas(64) std::atomic tasks_completed; - std::atomic tasks_consumed; -#endif - // ========================================================================= - // Inline hot-path methods - // ========================================================================= - - // Route a ready slot to the right global queue. Dummy tasks (empty - // active_mask) live in dummy_ready_queue; everything else goes to the - // per-shape ready_queues[]. - void push_ready_routed(PTO2TaskSlotState *slot_state) { - PTO2ResourceShape shape = slot_state->active_mask.to_shape(); - if (shape == PTO2ResourceShape::DUMMY || - (slot_state->active_mask.has_predicate() && !slot_state->payload->predicate.pass())) { - dummy_ready_queue.push(slot_state); - } else { - ready_queues[static_cast(shape)].push(slot_state); - } - } - - void check_and_handle_consumed(PTO2TaskSlotState &slot_state) { - // Read fanout_refcount/fanout_count and flip COMPLETED->CONSUMED under - // fanout_lock. The orchestrator claims producers (fanout_count++) under the - // same lock, so the consume decision is serialized against a concurrent - // claim: either the ++ lands first (count then exceeds refcount, so we do - // not consume and the producer stays pinned until released) or the consume - // lands first (the orchestrator then observes CONSUMED and skips the - // claim). Without this lock a claim racing the consume desyncs the slot's - // refcount and wedges in-order reclaim. - bool became_consumed = false; - slot_state.lock_fanout(); - if (slot_state.fanout_refcount.load(std::memory_order_acquire) == slot_state.fanout_count) { - PTO2TaskState expected = PTO2_TASK_COMPLETED; - became_consumed = slot_state.task_state.compare_exchange_strong( - expected, PTO2_TASK_CONSUMED, std::memory_order_acq_rel, std::memory_order_acquire - ); - } - slot_state.unlock_fanout(); - if (!became_consumed) return; - -#if SIMPLER_SCHED_PROFILING - tasks_consumed.fetch_add(1, std::memory_order_relaxed); -#endif - - int32_t ring_id = slot_state.ring_id; - // advance_ring_pointers (and the reset_for_reuse it triggers) MUST run - // outside fanout_lock: reset_for_reuse stores fanout_lock=0 and would - // clobber a held lock. Safe here — the slot is CONSUMED and quiescent. - // Try-lock — if another thread is advancing this ring, it will scan our CONSUMED task - int32_t expected_lock = 0; - if (ring_sched_states[ring_id].advance_lock.compare_exchange_strong( - expected_lock, 1, std::memory_order_acquire, std::memory_order_relaxed - )) { - ring_sched_states[ring_id].advance_ring_pointers(); - ring_sched_states[ring_id].advance_lock.store(0, std::memory_order_release); - } - } - -#if SIMPLER_ORCH_PROFILING || SIMPLER_SCHED_PROFILING - void check_and_handle_consumed(PTO2TaskSlotState &slot_state, uint64_t &atomic_count) { - // See the non-profiling overload for why the read + COMPLETED->CONSUMED - // flip is serialized against the orchestrator's claim under fanout_lock. - bool became_consumed = false; - slot_state.lock_fanout(); - atomic_count += 1; // lock CAS - uint32_t fc = slot_state.fanout_count; - uint32_t rc = slot_state.fanout_refcount.load(std::memory_order_acquire); - atomic_count += 1; // fanout_refcount.load (fanout_count is a plain read under lock) - if (rc == fc) { - PTO2TaskState expected = PTO2_TASK_COMPLETED; - became_consumed = slot_state.task_state.compare_exchange_strong( - expected, PTO2_TASK_CONSUMED, std::memory_order_acq_rel, std::memory_order_acquire - ); - atomic_count += 1; // CAS - } - slot_state.unlock_fanout(); - atomic_count += 1; // unlock store - if (!became_consumed) return; - -#if SIMPLER_SCHED_PROFILING - tasks_consumed.fetch_add(1, std::memory_order_relaxed); -#endif - - int32_t ring_id = slot_state.ring_id; - // advance_ring_pointers + reset_for_reuse run outside fanout_lock (reset - // stores fanout_lock=0). Safe — the slot is CONSUMED and quiescent. - // Try-lock — if another thread is advancing this ring, it will scan our CONSUMED task - int32_t expected_lock = 0; - if (ring_sched_states[ring_id].advance_lock.compare_exchange_strong( - expected_lock, 1, std::memory_order_acquire, std::memory_order_relaxed - )) { - ring_sched_states[ring_id].advance_ring_pointers(); - ring_sched_states[ring_id].advance_lock.store(0, std::memory_order_release); - atomic_count += 2; // try-lock CAS + unlock store - } else { - atomic_count += 1; // failed try-lock CAS - } - } -#endif - - void release_producer(PTO2TaskSlotState &slot_state) { - slot_state.fanout_refcount.fetch_add(1, std::memory_order_acq_rel); - check_and_handle_consumed(slot_state); - } - - // Scope-end release: sets bit31 (PTO2_FANOUT_SCOPE_BIT) instead of bumping a - // consumer ref. Called exactly once per task from on_scope_end. Keeping it a - // distinct add lets a consumer release leave the scope bit unset, so "all - // consumers done but scope still open" stays distinguishable from "fully - // consumed". - void release_producer_scope(PTO2TaskSlotState &slot_state) { - slot_state.fanout_refcount.fetch_add(PTO2_FANOUT_SCOPE_BIT, std::memory_order_acq_rel); - check_and_handle_consumed(slot_state); - } - -#if SIMPLER_ORCH_PROFILING || SIMPLER_SCHED_PROFILING - void release_producer(PTO2TaskSlotState &slot_state, uint64_t &atomic_count) { - slot_state.fanout_refcount.fetch_add(1, std::memory_order_acq_rel); - atomic_count += 1; // fanout_refcount.fetch_add - check_and_handle_consumed(slot_state, atomic_count); - } - - void release_producer_scope(PTO2TaskSlotState &slot_state, uint64_t &atomic_count) { - slot_state.fanout_refcount.fetch_add(PTO2_FANOUT_SCOPE_BIT, std::memory_order_acq_rel); - atomic_count += 1; // fanout_refcount.fetch_add - check_and_handle_consumed(slot_state, atomic_count); - } -#endif - - bool release_fanin_and_check_ready(PTO2TaskSlotState &slot_state) { - // 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 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) { - 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 false; - } -#endif - - int get_ready_tasks_batch(PTO2ResourceShape shape, PTO2TaskSlotState **out, int max_count) { - return ready_queues[static_cast(shape)].pop_batch(out, max_count); - } - -#if SIMPLER_SCHED_PROFILING - int get_ready_tasks_batch( - PTO2ResourceShape shape, PTO2TaskSlotState **out, int max_count, uint64_t &atomic_count, uint64_t &wait_cycle - ) { - return ready_queues[static_cast(shape)].pop_batch(out, max_count, atomic_count, wait_cycle); - } -#endif - - void on_scope_end(PTO2TaskSlotState **task_slot_states, int32_t count) { -#if SIMPLER_ORCH_PROFILING - extern uint64_t g_orch_scope_end_atomic_count; - if (count > 0) __builtin_prefetch(task_slot_states[0], 1, 0); - for (int32_t i = 0; i < count; i++) { - if (i + 1 < count) __builtin_prefetch(task_slot_states[i + 1], 1, 0); - release_producer_scope(*task_slot_states[i], g_orch_scope_end_atomic_count); - } -#else - if (count > 0) __builtin_prefetch(task_slot_states[0], 1, 0); - for (int32_t i = 0; i < count; i++) { - if (i + 1 < count) __builtin_prefetch(task_slot_states[i + 1], 1, 0); - release_producer_scope(*task_slot_states[i]); - } -#endif - } - - /** - * Subtask completion: atomic counter model. - * Called when a single subtask (AIC, AIV0, or AIV1) finishes on any block. - * Atomically increments completed_subtasks and checks whether all subtasks - * across all blocks are done. - * - * @return true if this was the last subtask, completing the entire task. - */ - bool on_subtask_complete(PTO2TaskSlotState &slot_state) { - int16_t prev = slot_state.completed_subtasks.fetch_add(1, std::memory_order_acq_rel); - return (prev + 1) == slot_state.total_required_subtasks; - } - - /** - * Two-stage completion: second stage. - * Called exactly once when all subtasks of a mixed task are done - * (i.e., on_subtask_complete returned true). - * Handles fanout notification, fanin release, and self-consumption check. - */ -#if SIMPLER_SCHED_PROFILING - CompletionStats -#else - void -#endif - on_task_complete( - PTO2TaskSlotState &slot_state -#if SIMPLER_SCHED_PROFILING - , - int thread_idx -#endif - ) { -#if SIMPLER_SCHED_PROFILING - CompletionStats stats = {0, 0, 0, true}; -#endif -#if SIMPLER_SCHED_PROFILING - extern uint64_t g_sched_lock_cycle[], g_sched_fanout_cycle[]; - extern uint64_t g_sched_lock_atomic_count[], g_sched_lock_wait_cycle[]; - extern uint64_t g_sched_fanout_atomic_count[], g_sched_push_wait_cycle[]; - uint64_t lock_atomics = 0, lock_wait = 0; - PTO2_SCHED_CYCLE_START(); -#endif - -#if SIMPLER_SCHED_PROFILING - slot_state.lock_fanout(lock_atomics, lock_wait); -#else - slot_state.lock_fanout(); -#endif - slot_state.task_state.store(PTO2_TASK_COMPLETED, std::memory_order_release); - PTO2DepListEntry *current = slot_state.fanout_head; // Protected by fanout_lock - slot_state.unlock_fanout(); - -#if SIMPLER_SCHED_PROFILING - lock_atomics += 2; // state.store + unlock.store - g_sched_lock_atomic_count[thread_idx] += lock_atomics; - g_sched_lock_wait_cycle[thread_idx] += lock_wait; - PTO2_SCHED_CYCLE_LAP(g_sched_lock_cycle[thread_idx]); -#endif - - // Fanout: notify consumers -#if SIMPLER_SCHED_PROFILING - uint64_t fanout_atomics = 0, push_wait = 0; -#endif - 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)) { - stats.tasks_enqueued++; - } -#else - release_fanin_and_check_ready(consumer_slot); -#endif - current = current->next; - } - -#if SIMPLER_SCHED_PROFILING - g_sched_fanout_atomic_count[thread_idx] += fanout_atomics; - g_sched_push_wait_cycle[thread_idx] += push_wait; - PTO2_SCHED_CYCLE_LAP(g_sched_fanout_cycle[thread_idx]); - return stats; -#endif - } - - /** - * Cold path: release producers (fanin traversal) + check self for CONSUMED. - * Returns fanin edge count for profiling. - */ - -#if SIMPLER_SCHED_PROFILING - int32_t on_task_release(PTO2TaskSlotState &slot_state, int32_t thread_idx) { - PTO2_SCHED_CYCLE_START(); - extern uint64_t g_sched_fanin_cycle[], g_sched_fanin_atomic_count[]; - extern uint64_t g_sched_self_atomic_count[]; - extern uint64_t g_sched_self_consumed_cycle[]; - extern uint64_t g_sched_complete_count[]; - uint64_t fanin_atomics = 0; -#else - int32_t on_task_release(PTO2TaskSlotState &slot_state) { -#endif - PTO2TaskPayload *payload = slot_state.payload; - for_each_fanin_slot_state(*payload, [&](PTO2TaskSlotState *producer_slot_state) { -#if SIMPLER_SCHED_PROFILING - release_producer(*producer_slot_state, fanin_atomics); -#else - release_producer(*producer_slot_state); -#endif - }); -#if SIMPLER_SCHED_PROFILING - g_sched_fanin_atomic_count[thread_idx] += fanin_atomics; - PTO2_SCHED_CYCLE_LAP(g_sched_fanin_cycle[thread_idx]); -#endif - - // Self consumed check -#if SIMPLER_SCHED_PROFILING - uint64_t self_atomics = 0; - check_and_handle_consumed(slot_state, self_atomics); - g_sched_self_atomic_count[thread_idx] += self_atomics; - PTO2_SCHED_CYCLE_LAP(g_sched_self_consumed_cycle[thread_idx]); - g_sched_complete_count[thread_idx]++; -#else - check_and_handle_consumed(slot_state); -#endif - return payload->fanin_actual_count; - } - - // === Cold-path API (defined in pto_scheduler.cpp) === - - // Phase 1: declare every sub-region (ready_queue slots, dummy queue slots, - // per-ring dep_pool entries) on the supplied arena. - // Capacities are baked into the returned layout; init_data_from_layout uses - // the same values. - static PTO2SchedulerLayout reserve_layout(DeviceArena &arena, int32_t dep_pool_capacity = PTO2_DEP_LIST_POOL_SIZE); - static PTO2SchedulerLayout - reserve_layout(DeviceArena &arena, const int32_t dep_pool_capacities[PTO2_MAX_RING_DEPTH]); - - // Phase 3a: write everything *except* arena-internal pointer fields. - // `sm_dev_base` is the device address of the SM (only stored, never - // dereferenced here). Safe to call on a host arena that holds the - // prebuilt image buffer. (The orchestrator counterpart takes - // task_window_size for ring task_descriptors address arithmetic; the - // scheduler only needs the SM header / ring header base addresses, - // both window-size-independent.) - bool init_data_from_layout(const PTO2SchedulerLayout &layout, DeviceArena &arena, void *sm_dev_base); - void reset_for_reuse(const PTO2SchedulerLayout &layout, void *sm_dev_base); - - // Phase 3b: write the arena-internal pointer fields - // (ready_queues[].slots, dummy_ready_queue.slots, dep_pool.base for each - // ring). Called on both host and device sides. - void wire_arena_pointers(const PTO2SchedulerLayout &layout, DeviceArena &arena); - - // Forget per-region pointers; arena owns the backing memory. - void destroy(); - void print_stats(); - void print_queues(); -}; - -// Scheduler cold-path API is declared as PTO2SchedulerState member functions. -// See init()/destroy()/print_stats()/print_queues() below the struct definition. - -// Short-circuit NotDeferred completions seen during drain so they don't grow -// entries[]. Mirrors the a2a3 impl; see that mirror for the rationale. -inline bool -AsyncWaitList::try_inline_complete_locked(AsyncWaitList::DrainCompletionSink &sink, PTO2TaskSlotState &slot_state) { -#if SIMPLER_SCHED_PROFILING - sink.sched->on_task_complete(slot_state, sink.thread_idx); -#else - sink.sched->on_task_complete(slot_state); -#endif - if (*sink.deferred_release_count >= sink.deferred_release_capacity) { - while (*sink.deferred_release_count > 0) { -#if SIMPLER_SCHED_PROFILING - (void)sink.sched->on_task_release( - *sink.deferred_release_slot_states[--(*sink.deferred_release_count)], sink.thread_idx - ); -#else - sink.sched->on_task_release(*sink.deferred_release_slot_states[--(*sink.deferred_release_count)]); -#endif - } - } - sink.deferred_release_slot_states[(*sink.deferred_release_count)++] = &slot_state; - sink.inline_completed++; - return true; -} - -template -inline AsyncPollResult AsyncWaitList::poll_and_complete( - AICoreCompletionMailbox *aicore_mailbox, PTO2SchedulerState *sched, - PTO2TaskSlotState **deferred_release_slot_states, int32_t &deferred_release_count, int32_t deferred_release_capacity -#if SIMPLER_SCHED_PROFILING - , - int thread_idx -#endif -) { - AsyncPollResult result; - if (!try_lock()) return result; - - AsyncWaitList::DrainCompletionSink sink{}; - sink.sched = sched; - sink.deferred_release_slot_states = deferred_release_slot_states; - sink.deferred_release_count = &deferred_release_count; - sink.deferred_release_capacity = deferred_release_capacity; -#if SIMPLER_SCHED_PROFILING - sink.thread_idx = thread_idx; -#endif - - int32_t drain_err = PTO2_ERROR_NONE; - drain_aicore_completion_mailbox_locked(aicore_mailbox, sink, drain_err); - if (drain_err != PTO2_ERROR_NONE) { - result.error_code = drain_err; - unlock(); - return result; - } - result.completed += sink.inline_completed; - - for (int32_t i = count - 1; i >= 0; --i) { - AsyncWaitEntry &entry = entries[i]; - for (int32_t c = 0; c < entry.condition_count; c++) { - CompletionCondition &cond = entry.conditions[c]; - if (cond.satisfied) continue; - CompletionPollResult poll = cond.test(); - if (poll.state == CompletionPollState::FAILED) { - result.error_code = poll.error_code; - result.failed_slot_state = entry.slot_state; - unlock(); - return result; - } - if (poll.state == CompletionPollState::READY) { - cond.satisfied = true; - cond.retire(); - entry.waiting_completion_count--; - } - } - - if (entry.normal_done && entry.waiting_completion_count <= 0) { -#if SIMPLER_SCHED_PROFILING - sched->on_task_complete(*entry.slot_state, thread_idx); -#else - sched->on_task_complete(*entry.slot_state); -#endif - if (deferred_release_count >= deferred_release_capacity) { - while (deferred_release_count > 0) { -#if SIMPLER_SCHED_PROFILING - (void)sched->on_task_release(*deferred_release_slot_states[--deferred_release_count], thread_idx); -#else - sched->on_task_release(*deferred_release_slot_states[--deferred_release_count]); -#endif - } - } - deferred_release_slot_states[deferred_release_count++] = entry.slot_state; - result.completed++; - - int32_t last = count - 1; - if (i != last) entries[i] = entries[last]; - count = last; - } - } - - unlock(); - return result; -} - -// ============================================================================= -// Scheduler Profiling Data -// ============================================================================= - -#if SIMPLER_SCHED_PROFILING -struct PTO2SchedProfilingData { - // Sub-phase cycle breakdown within on_task_complete - uint64_t lock_cycle; // lock_fanout + state store + unlock - uint64_t fanout_cycle; // fanout traversal - uint64_t fanin_cycle; // fanin traversal - uint64_t self_consumed_cycle; // self check_and_handle_consumed - - // Wait times - uint64_t lock_wait_cycle; // spin-wait in fanout_lock - uint64_t push_wait_cycle; // CAS contention in push() - uint64_t pop_wait_cycle; // CAS contention in pop() - - // Atomic counts per sub-phase - uint64_t lock_atomic_count; - uint64_t fanout_atomic_count; - uint64_t fanin_atomic_count; - uint64_t self_atomic_count; - uint64_t pop_atomic_count; - - int64_t complete_count; -}; - -/** - * Get and reset scheduler profiling data for a specific thread. - * Returns accumulated profiling data and resets counters. - */ -PTO2SchedProfilingData scheduler_get_profiling(int thread_idx); -#endif diff --git a/_run_early_st_board.sh b/_run_early_st_board.sh deleted file mode 100755 index 1b3df0e9f6..0000000000 --- a/_run_early_st_board.sh +++ /dev/null @@ -1,22 +0,0 @@ -#!/bin/bash -set -euo pipefail -cd /home/pyptouser/yanghaoran/Desktop/simpler2 -source .venv/bin/activate -DEV="${ASCEND_RT_VISIBLE_DEVICES:-${DEVICE:-}}" -# Prefer lock-assigned device: task-submit often sets ASCEND_RT_VISIBLE_DEVICES -if [[ -z "${DEV}" ]]; then - echo "No device env; failing" - env | sort | head -40 - exit 2 -fi -# If ASCEND is set to a single id by task-submit, use it -echo "ASCEND_RT_VISIBLE_DEVICES=${ASCEND_RT_VISIBLE_DEVICES:-unset} DEVICE=${DEVICE:-unset}" -# Use first id for --device -DID="${ASCEND_RT_VISIBLE_DEVICES%%,*}" -DID="${DID:-$DEV}" -echo "Using --device $DID" -TEST="${1:?test path}" -NAME="${2:?label}" -python -m pytest "$TEST" --platform a5 --device "$DID" -v --tb=short \ - --enable-l2-swimlane 4 --enable-dep-gen \ - 2>&1 | tee "_early_st_out/${NAME}_board.log" diff --git a/_run_st_dev2.sh b/_run_st_dev2.sh deleted file mode 100755 index 9e9e031551..0000000000 --- a/_run_st_dev2.sh +++ /dev/null @@ -1,9 +0,0 @@ -#!/bin/bash -set -euo pipefail -cd /home/pyptouser/yanghaoran/Desktop/simpler2 -source .venv/bin/activate -# task-submit --device N may not export ASCEND; hardcode from argv -DEV="${1:?device}" -SHIFT=1 -shift -python -m pytest "$@" --platform a5 --device "$DEV" -v --tb=line diff --git a/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/README.md b/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/README.md deleted file mode 100644 index 16a3d9e7de..0000000000 --- a/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/README.md +++ /dev/null @@ -1,62 +0,0 @@ -# Acc ValidShape strip C2V (`srcStride` repro) - -Minimal Ascend950 **simpler** sample for -`ISSUE_A5_ACC_VALIDSHAPE_C2V_SRCSTRIDE`. - -## Primary test - -`test_acc_c2v_strip_vs_full.py` — same Acc drained two ways: - -| Output | Path | -|--------|------| -| `C_full` | one Acc TPUSH, `Valid≡Rows` | -| `C_strip` | Acc+`ValidShape(H=16)` windows @ `addr=row*64` | - -Host: `C_strip ≈ C_full`. Shapes (UB-safe): `M=32, N=128, H=16, K=32` -(still `validRow < Rows`). Seed `torch.manual_seed(0)`. - -### Board results (2026-07-21, device dump) - -#### 修复前 (`srcStride=align(validRow)`) - -典型现象：strip0（行 0…15）与 full 一致；strip1（行 16…）系统性偏离。 -`col=0` 在本 seed 下仍可能碰巧相等，看同行列最大偏差处： - -```text -# strip0 — 接近/相等 -row=0 C_full[0,0]=10.4928 C_strip[0,0]=10.4928 diff=0 -row=15 C_full[15,0]=-2.02309 C_strip[15,0]=-2.02309 diff=0 - -# strip1 — 明显不等（本 seed 最大差在下列） -row=16 C_full[16,71]=-12.8381 C_strip[16,71]=13.6886 |diff|=26.5268 -row=17 C_full[17,44]=-11.6004 C_strip[17,44]=11.7646 |diff|=23.365 -# 汇总: strip0_max=0.0 strip1_max=31.9782 → 主测 FAILED -``` - -#### 修复后 (`srcStride=align(Rows)`, pin `0ebbd03d`) - -```text -row=0 C_full[0,0]=10.4928 C_strip[0,0]=10.4928 diff=0 -row=15 C_full[15,0]=-2.02309 C_strip[15,0]=-2.02309 diff=0 -row=16 C_full[16,0]=2.29007 C_strip[16,0]=2.29007 diff=0 -row=17 C_full[17,0]=5.26025 C_strip[17,0]=5.26025 diff=0 -# 汇总: strip0_max=0 strip1_max=0 → PASSED -``` - -```bash -# 在已配置 A5 的 simpler 仓库根目录： -python examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/test_acc_c2v_strip_vs_full.py \ - -p a5 -d - -# 或用包根的 board 包装（改 SIMPLER 后）： -task-submit --device auto --max-time 900 --run \ - "bash /path/to/run_acc_c2v_strip_vs_full_board.sh" -``` - -## Files - -| Path | Role | -|------|------| -| `kernels/mix/kernel_acc_c2v_compare.cpp` | matmul + full + strip TPUSH | -| `kernels/orchestration/acc_c2v_compare_orch.cpp` | `[A,B,C_full,C_strip]` | -| `test_acc_c2v_strip_vs_full.py` | self-compare | diff --git a/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/kernels/aic/kernel_matmul_tstore.cpp b/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/kernels/aic/kernel_matmul_tstore.cpp deleted file mode 100644 index 360f4f6273..0000000000 --- a/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/kernels/aic/kernel_matmul_tstore.cpp +++ /dev/null @@ -1,69 +0,0 @@ -/* - * CONTROL: Acc float matmul → TSTORE Acc to GM (no C2V). Isolates matmul golden match. - */ -#include -#include -#include "tensor.h" - -using namespace pto; - -#ifndef __gm__ -#define __gm__ -#endif -#ifndef __aicore__ -#define __aicore__ [aicore] -#endif -#include "intrinsic.h" - -constexpr int M = 128; -constexpr int N = 256; -constexpr int K = 32; - -extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { - __gm__ Tensor *a_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); - __gm__ Tensor *b_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); - __gm__ Tensor *c_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); - - __gm__ float *a_ptr = reinterpret_cast<__gm__ float *>(a_tensor->buffer.addr) + a_tensor->start_offset; - __gm__ float *b_ptr = reinterpret_cast<__gm__ float *>(b_tensor->buffer.addr) + b_tensor->start_offset; - __gm__ float *c_ptr = reinterpret_cast<__gm__ float *>(c_tensor->buffer.addr) + c_tensor->start_offset; - - using GlobalA = GlobalTensor, pto::Stride>; - using GlobalB = GlobalTensor, pto::Stride>; - using GlobalC = GlobalTensor, pto::Stride>; - GlobalA aGlobal(a_ptr); - GlobalB bGlobal(b_ptr); - GlobalC cGlobal(c_ptr); - - using TileMatA = Tile; - using TileMatB = Tile; - using LeftT = TileLeft; - using RightT = TileRight; - using AccT = TileAcc; - - TileMatA aMat; - TileMatB bMat; - TASSIGN(aMat, 0x0); - TASSIGN(bMat, 0x20000); - LeftT aL0; - RightT bL0; - AccT acc; - TASSIGN(aL0, 0x0); - TASSIGN(bL0, 0x0); - TASSIGN(acc, 0x0); - - TLOAD(aMat, aGlobal); - TLOAD(bMat, bGlobal); - set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); - wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); - TMOV(aL0, aMat); - TMOV(bL0, bMat); - set_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); - wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); - TMATMUL(acc, aL0, bL0); - set_flag(PIPE_M, PIPE_FIX, EVENT_ID0); - wait_flag(PIPE_M, PIPE_FIX, EVENT_ID0); - TSTORE(cGlobal, acc); - set_flag(PIPE_FIX, PIPE_S, EVENT_ID7); - wait_flag(PIPE_FIX, PIPE_S, EVENT_ID7); -} diff --git a/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/kernels/mix/kernel_acc_c2v_compare.cpp b/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/kernels/mix/kernel_acc_c2v_compare.cpp deleted file mode 100644 index eb093ea56b..0000000000 --- a/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/kernels/mix/kernel_acc_c2v_compare.cpp +++ /dev/null @@ -1,194 +0,0 @@ -/* - * 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. - * ----------------------------------------------------------------------------------------------------------- - */ -/** - * A5 Acc ValidShape strip C2V repro (self-compare). - * - * AIC: Acc = A@B once, then - * 1) full Acc TPUSH (Valid≡Rows) → AIV stores C_full - * 2) 8× Acc+Valid(H) TPUSH @ addr=row*64 → AIV stores C_strip - * - * Host golden: C_strip == C_full. - * - unfixed TMovCcToUb (srcStride=validRow): FAIL - * - fixed (srcStride=Rows): PASS - * - * args[0]=A, args[1]=B, args[2]=C_full, args[3]=C_strip - */ - -#include -#include -#include -#include "tensor.h" - -using pto::BLayout; -using pto::Direction; -using pto::GlobalTensor; -using pto::Shape; -using pto::SLayout; -using pto::Tile; -using pto::TileAcc; -using pto::TileLeft; -using pto::TileRight; -using pto::TileSplitAxis; -using pto::TileType; -using pto::TPipe; - -#ifndef __gm__ -#define __gm__ -#endif -#ifndef __aicore__ -#define __aicore__ [aicore] -#endif -#include "intrinsic.h" - -#ifdef __DAV_CUBE__ -constexpr bool DAV_CUBE = true; -#else -constexpr bool DAV_CUBE = false; -#endif -#ifdef __DAV_VEC__ -constexpr bool DAV_VEC = true; -#else -constexpr bool DAV_VEC = false; -#endif - -constexpr int M = 32; -constexpr int N = 128; -constexpr int H = 16; -constexpr int K = 32; -constexpr uint64_t kAccRowByteStride = 64; -constexpr uint16_t PP_FLAG_ID = 0; -constexpr uint8_t PP_FIFO_DEPTH = 1; - -constexpr uint32_t kFullSlot = static_cast(M * N * sizeof(float)); -constexpr uint32_t kStripSlot = static_cast(H * N * sizeof(float)); - -using AccFullT = TileAcc; -using AccWinFullT = Tile; -using AccWinStripT = Tile; -using VecFullT = Tile; -using VecStripT = Tile; - -using PipeFullT = TPipe; -using PipeStripT = TPipe; - -extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { - __gm__ Tensor *a_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); - __gm__ Tensor *b_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); - __gm__ Tensor *c_full_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); - __gm__ Tensor *c_strip_tensor = reinterpret_cast<__gm__ Tensor *>(args[3]); - - PipeFullT pipeFull(nullptr, 0U, 0U); - PipeStripT pipeStrip(nullptr, 0U, 0U); - - if constexpr (DAV_CUBE) { - __gm__ float *a_ptr = - reinterpret_cast<__gm__ float *>(a_tensor->buffer.addr) + a_tensor->start_offset; - __gm__ float *b_ptr = - reinterpret_cast<__gm__ float *>(b_tensor->buffer.addr) + b_tensor->start_offset; - - using GlobalA = GlobalTensor, pto::Stride>; - using GlobalB = GlobalTensor, pto::Stride>; - GlobalA aGlobal(a_ptr); - GlobalB bGlobal(b_ptr); - - using TileMatA = Tile; - using TileMatB = Tile; - using LeftT = TileLeft; - using RightT = TileRight; - - TileMatA aMat; - TileMatB bMat; - TASSIGN(aMat, 0x0); - TASSIGN(bMat, 0x20000); - LeftT aL0; - RightT bL0; - AccFullT acc; - TASSIGN(aL0, 0x0); - TASSIGN(bL0, 0x0); - TASSIGN(acc, 0x0); - - TLOAD(aMat, aGlobal); - TLOAD(bMat, bGlobal); - set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); - wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); - TMOV(aL0, aMat); - TMOV(bL0, bMat); - set_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); - wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); - TMATMUL(acc, aL0, bL0); - set_flag(PIPE_M, PIPE_FIX, EVENT_ID0); - wait_flag(PIPE_M, PIPE_FIX, EVENT_ID0); - - // (1) full Acc C2V — Valid≡Rows, always correct for srcStride either formula - { - AccWinFullT full; - TASSIGN(full, 0x0); - TPUSH(pipeFull, full); - } - - // (2) strip Acc C2V — Valid(H) windows (the buggy pattern) - for (int row = 0; row < M; row += H) { - AccWinStripT strip; - TASSIGN(strip, static_cast(row) * kAccRowByteStride); - TPUSH(pipeStrip, strip); - } - - set_flag(PIPE_FIX, PIPE_S, EVENT_ID7); - wait_flag(PIPE_FIX, PIPE_S, EVENT_ID7); - } - - if constexpr (DAV_VEC) { - if (get_sub_block_id(args) != 0) { - return; - } - __gm__ float *c_full = - reinterpret_cast<__gm__ float *>(c_full_tensor->buffer.addr) + c_full_tensor->start_offset; - __gm__ float *c_strip = - reinterpret_cast<__gm__ float *>(c_strip_tensor->buffer.addr) + c_strip_tensor->start_offset; - - // UB layout: full fifo + strip fifo + scratch - constexpr uint64_t kFullFifo = static_cast(PP_FIFO_DEPTH) * kFullSlot; - constexpr uint64_t kStripFifo = static_cast(PP_FIFO_DEPTH) * kStripSlot; - VecFullT vFull; - VecStripT vStrip; - TASSIGN(vFull, kFullFifo + kStripFifo); - TASSIGN(vStrip, kFullFifo + kStripFifo + kFullSlot); - - // Pop full - TPOP(pipeFull, vFull); - set_flag(PIPE_S, PIPE_MTE3, EVENT_ID0); - wait_flag(PIPE_S, PIPE_MTE3, EVENT_ID0); - using GlobalFull = GlobalTensor, pto::Stride>; - GlobalFull gFull(c_full); - TSTORE(gFull, vFull); - TFREE(pipeFull); - set_flag(PIPE_MTE3, PIPE_S, EVENT_ID1); - wait_flag(PIPE_MTE3, PIPE_S, EVENT_ID1); - - // Pop strips - for (int row = 0; row < M; row += H) { - TPOP(pipeStrip, vStrip); - set_flag(PIPE_S, PIPE_MTE3, EVENT_ID2); - wait_flag(PIPE_S, PIPE_MTE3, EVENT_ID2); - using GlobalStrip = - GlobalTensor, pto::Stride>; - GlobalStrip gStrip(c_strip + static_cast(row) * N); - TSTORE(gStrip, vStrip); - TFREE(pipeStrip); - set_flag(PIPE_MTE3, PIPE_S, EVENT_ID3); - wait_flag(PIPE_MTE3, PIPE_S, EVENT_ID3); - } - } -} diff --git a/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/kernels/mix/kernel_acc_c2v_full.cpp b/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/kernels/mix/kernel_acc_c2v_full.cpp deleted file mode 100644 index e1bda584a0..0000000000 --- a/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/kernels/mix/kernel_acc_c2v_full.cpp +++ /dev/null @@ -1,192 +0,0 @@ -/* - * 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. - * ----------------------------------------------------------------------------------------------------------- - */ -/** - * A5 Acc ValidShape strip C2V repro — matches ISSUE_A5_ACC_VALIDSHAPE_C2V_SRCSTRIDE §样例. - * - * AIC: Acc[M,N] = A@B; then M/H × TPUSH Acc[M,N]+Valid(H,N) @ addr=row*64 (TILE_NO_SPLIT). - * AIV0: TPOP Vec[H,N] ND strip → TSTORE GM (AIV1 idle for C2V). - * - * USE_STRIP_C2V=0: one full Acc TPUSH (Valid≡Rows) control path. - */ - -#include -#include -#include - -#include "tensor.h" - -using pto::BLayout; -using pto::Direction; -using pto::GlobalTensor; -using pto::Shape; -using pto::SLayout; -using pto::Tile; -using pto::TileAcc; -using pto::TileLeft; -using pto::TileRight; -using pto::TileSplitAxis; -using pto::TileType; -using pto::TPipe; - -#ifndef __gm__ -#define __gm__ -#endif - -#ifndef __aicore__ -#define __aicore__ [aicore] -#endif - -#include "intrinsic.h" - -#ifdef __DAV_CUBE__ -constexpr bool DAV_CUBE = true; -#else -constexpr bool DAV_CUBE = false; -#endif - -#ifdef __DAV_VEC__ -constexpr bool DAV_VEC = true; -#else -constexpr bool DAV_VEC = false; -#endif - -#ifndef USE_STRIP_C2V -#define USE_STRIP_C2V 0 -#endif - -constexpr int M = 128; -constexpr int N = 256; -constexpr int H = 16; -constexpr int K = 32; - -constexpr uint64_t kAccRowByteStride = 64; - -constexpr uint16_t PP_FLAG_ID = 0; -constexpr uint8_t PP_FIFO_DEPTH = 2; - -#if USE_STRIP_C2V -constexpr int PUSH_ROWS = H; -constexpr int kNumPush = M / H; -#else -constexpr int PUSH_ROWS = M; -constexpr int kNumPush = 1; -#endif - -constexpr uint32_t kSlotBytes = static_cast(PUSH_ROWS * N * sizeof(float)); - -using AccFullT = TileAcc; -using AccPushT = Tile; -using VecPushT = Tile; -// IsNoSplit=true matches TILE_NO_SPLIT (AIV0 only), same as qr_proj full-Acc hand patch. -using PipeT = TPipe; - -extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { - __gm__ Tensor *a_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); - __gm__ Tensor *b_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); - __gm__ Tensor *c_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); - - PipeT pipe(nullptr, 0U, 0U); - - if constexpr (DAV_CUBE) { - __gm__ float *a_ptr = - reinterpret_cast<__gm__ float *>(a_tensor->buffer.addr) + a_tensor->start_offset; - __gm__ float *b_ptr = - reinterpret_cast<__gm__ float *>(b_tensor->buffer.addr) + b_tensor->start_offset; - - using GlobalA = GlobalTensor, pto::Stride>; - using GlobalB = GlobalTensor, pto::Stride>; - GlobalA aGlobal(a_ptr); - GlobalB bGlobal(b_ptr); - - using TileMatA = Tile; - using TileMatB = Tile; - using LeftT = TileLeft; - using RightT = TileRight; - - TileMatA aMat; - TileMatB bMat; - TASSIGN(aMat, 0x0); - TASSIGN(bMat, 0x20000); - - LeftT aL0; - RightT bL0; - AccFullT acc; - TASSIGN(aL0, 0x0); - TASSIGN(bL0, 0x0); - TASSIGN(acc, 0x0); - - TLOAD(aMat, aGlobal); - TLOAD(bMat, bGlobal); - - set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); - wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); - - TMOV(aL0, aMat); - TMOV(bL0, bMat); - - set_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); - wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); - - TMATMUL(acc, aL0, bL0); - - set_flag(PIPE_M, PIPE_FIX, EVENT_ID0); - wait_flag(PIPE_M, PIPE_FIX, EVENT_ID0); - -#if USE_STRIP_C2V - for (int row = 0; row < M; row += H) { - AccPushT strip; - TASSIGN(strip, static_cast(row) * kAccRowByteStride); - TPUSH(pipe, strip); - } -#else - { - AccPushT full; - TASSIGN(full, 0x0); - TPUSH(pipe, full); - } -#endif - - set_flag(PIPE_FIX, PIPE_S, EVENT_ID7); - wait_flag(PIPE_FIX, PIPE_S, EVENT_ID7); - } - - if constexpr (DAV_VEC) { - if (get_sub_block_id(args) != 0) { - return; - } - - __gm__ float *c_ptr = - reinterpret_cast<__gm__ float *>(c_tensor->buffer.addr) + c_tensor->start_offset; - - constexpr uint64_t kScratch = static_cast(PP_FIFO_DEPTH) * kSlotBytes; - VecPushT vec; - TASSIGN(vec, kScratch); - - for (int s = 0; s < kNumPush; ++s) { - TPOP(pipe, vec); - - int row0 = s * PUSH_ROWS; - using GlobalC = GlobalTensor, - pto::Stride>; - GlobalC cGlobal(c_ptr + static_cast(row0) * N); - - set_flag(PIPE_S, PIPE_MTE3, EVENT_ID0); - wait_flag(PIPE_S, PIPE_MTE3, EVENT_ID0); - TSTORE(cGlobal, vec); - TFREE(pipe); - - set_flag(PIPE_MTE3, PIPE_S, EVENT_ID1); - wait_flag(PIPE_MTE3, PIPE_S, EVENT_ID1); - } - } -} diff --git a/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/kernels/mix/kernel_acc_c2v_strip.cpp b/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/kernels/mix/kernel_acc_c2v_strip.cpp deleted file mode 100644 index e6be66ec92..0000000000 --- a/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/kernels/mix/kernel_acc_c2v_strip.cpp +++ /dev/null @@ -1,192 +0,0 @@ -/* - * 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. - * ----------------------------------------------------------------------------------------------------------- - */ -/** - * A5 Acc ValidShape strip C2V repro — matches ISSUE_A5_ACC_VALIDSHAPE_C2V_SRCSTRIDE §样例. - * - * AIC: Acc[M,N] = A@B; then M/H × TPUSH Acc[M,N]+Valid(H,N) @ addr=row*64 (TILE_NO_SPLIT). - * AIV0: TPOP Vec[H,N] ND strip → TSTORE GM (AIV1 idle for C2V). - * - * USE_STRIP_C2V=0: one full Acc TPUSH (Valid≡Rows) control path. - */ - -#include -#include -#include - -#include "tensor.h" - -using pto::BLayout; -using pto::Direction; -using pto::GlobalTensor; -using pto::Shape; -using pto::SLayout; -using pto::Tile; -using pto::TileAcc; -using pto::TileLeft; -using pto::TileRight; -using pto::TileSplitAxis; -using pto::TileType; -using pto::TPipe; - -#ifndef __gm__ -#define __gm__ -#endif - -#ifndef __aicore__ -#define __aicore__ [aicore] -#endif - -#include "intrinsic.h" - -#ifdef __DAV_CUBE__ -constexpr bool DAV_CUBE = true; -#else -constexpr bool DAV_CUBE = false; -#endif - -#ifdef __DAV_VEC__ -constexpr bool DAV_VEC = true; -#else -constexpr bool DAV_VEC = false; -#endif - -#ifndef USE_STRIP_C2V -#define USE_STRIP_C2V 1 -#endif - -constexpr int M = 128; -constexpr int N = 256; -constexpr int H = 16; -constexpr int K = 32; - -constexpr uint64_t kAccRowByteStride = 64; - -constexpr uint16_t PP_FLAG_ID = 0; -constexpr uint8_t PP_FIFO_DEPTH = 2; - -#if USE_STRIP_C2V -constexpr int PUSH_ROWS = H; -constexpr int kNumPush = M / H; -#else -constexpr int PUSH_ROWS = M; -constexpr int kNumPush = 1; -#endif - -constexpr uint32_t kSlotBytes = static_cast(PUSH_ROWS * N * sizeof(float)); - -using AccFullT = TileAcc; -using AccPushT = Tile; -using VecPushT = Tile; -// IsNoSplit=true matches TILE_NO_SPLIT (AIV0 only), same as qr_proj full-Acc hand patch. -using PipeT = TPipe; - -extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { - __gm__ Tensor *a_tensor = reinterpret_cast<__gm__ Tensor *>(args[0]); - __gm__ Tensor *b_tensor = reinterpret_cast<__gm__ Tensor *>(args[1]); - __gm__ Tensor *c_tensor = reinterpret_cast<__gm__ Tensor *>(args[2]); - - PipeT pipe(nullptr, 0U, 0U); - - if constexpr (DAV_CUBE) { - __gm__ float *a_ptr = - reinterpret_cast<__gm__ float *>(a_tensor->buffer.addr) + a_tensor->start_offset; - __gm__ float *b_ptr = - reinterpret_cast<__gm__ float *>(b_tensor->buffer.addr) + b_tensor->start_offset; - - using GlobalA = GlobalTensor, pto::Stride>; - using GlobalB = GlobalTensor, pto::Stride>; - GlobalA aGlobal(a_ptr); - GlobalB bGlobal(b_ptr); - - using TileMatA = Tile; - using TileMatB = Tile; - using LeftT = TileLeft; - using RightT = TileRight; - - TileMatA aMat; - TileMatB bMat; - TASSIGN(aMat, 0x0); - TASSIGN(bMat, 0x20000); - - LeftT aL0; - RightT bL0; - AccFullT acc; - TASSIGN(aL0, 0x0); - TASSIGN(bL0, 0x0); - TASSIGN(acc, 0x0); - - TLOAD(aMat, aGlobal); - TLOAD(bMat, bGlobal); - - set_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); - wait_flag(PIPE_MTE2, PIPE_MTE1, EVENT_ID0); - - TMOV(aL0, aMat); - TMOV(bL0, bMat); - - set_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); - wait_flag(PIPE_MTE1, PIPE_M, EVENT_ID0); - - TMATMUL(acc, aL0, bL0); - - set_flag(PIPE_M, PIPE_FIX, EVENT_ID0); - wait_flag(PIPE_M, PIPE_FIX, EVENT_ID0); - -#if USE_STRIP_C2V - for (int row = 0; row < M; row += H) { - AccPushT strip; - TASSIGN(strip, static_cast(row) * kAccRowByteStride); - TPUSH(pipe, strip); - } -#else - { - AccPushT full; - TASSIGN(full, 0x0); - TPUSH(pipe, full); - } -#endif - - set_flag(PIPE_FIX, PIPE_S, EVENT_ID7); - wait_flag(PIPE_FIX, PIPE_S, EVENT_ID7); - } - - if constexpr (DAV_VEC) { - if (get_sub_block_id(args) != 0) { - return; - } - - __gm__ float *c_ptr = - reinterpret_cast<__gm__ float *>(c_tensor->buffer.addr) + c_tensor->start_offset; - - constexpr uint64_t kScratch = static_cast(PP_FIFO_DEPTH) * kSlotBytes; - VecPushT vec; - TASSIGN(vec, kScratch); - - for (int s = 0; s < kNumPush; ++s) { - TPOP(pipe, vec); - - int row0 = s * PUSH_ROWS; - using GlobalC = GlobalTensor, - pto::Stride>; - GlobalC cGlobal(c_ptr + static_cast(row0) * N); - - set_flag(PIPE_S, PIPE_MTE3, EVENT_ID0); - wait_flag(PIPE_S, PIPE_MTE3, EVENT_ID0); - TSTORE(cGlobal, vec); - TFREE(pipe); - - set_flag(PIPE_MTE3, PIPE_S, EVENT_ID1); - wait_flag(PIPE_MTE3, PIPE_S, EVENT_ID1); - } - } -} diff --git a/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/kernels/orchestration/acc_c2v_compare_orch.cpp b/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/kernels/orchestration/acc_c2v_compare_orch.cpp deleted file mode 100644 index 54cfc64a33..0000000000 --- a/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/kernels/orchestration/acc_c2v_compare_orch.cpp +++ /dev/null @@ -1,53 +0,0 @@ -/* - * Orchestration: A, B, C_full, C_strip — one MixedKernels. - */ -#include -#include "pto_orchestration_api.h" // NOLINT(build/include_subdir) - -#define FUNC_AIC 0 -#define FUNC_AIV 1 - -static constexpr int M = 32; -static constexpr int N = 128; -static constexpr int K = 32; -static constexpr uint32_t A_ELEMS = static_cast(M * K); -static constexpr uint32_t B_ELEMS = static_cast(K * N); -static constexpr uint32_t C_ELEMS = static_cast(M * N); - -extern "C" { - -__attribute__((visibility("default"))) PTO2OrchestrationConfig aicpu_orchestration_config(const L2TaskArgs &orch_args) { - (void)orch_args; - return PTO2OrchestrationConfig{.expected_arg_count = 4}; -} - -__attribute__((visibility("default"))) void aicpu_orchestration_entry(const L2TaskArgs &orch_args) { - const Tensor &ext_A = orch_args.tensor(0).ref(); - const Tensor &ext_B = orch_args.tensor(1).ref(); - const Tensor &ext_Cf = orch_args.tensor(2).ref(); - const Tensor &ext_Cs = orch_args.tensor(3).ref(); - - uint32_t a_shapes[1] = {A_ELEMS}; - uint32_t b_shapes[1] = {B_ELEMS}; - uint32_t c_shapes[1] = {C_ELEMS}; - uint32_t zero[1] = {0}; - - Tensor A_view = ext_A.view(a_shapes, zero); - Tensor B_view = ext_B.view(b_shapes, zero); - Tensor Cf_view = ext_Cf.view(c_shapes, zero); - Tensor Cs_view = ext_Cs.view(c_shapes, zero); - - L0TaskArgs args; - args.add_input(A_view); - args.add_input(B_view); - args.add_output(Cf_view); - args.add_output(Cs_view); - - MixedKernels mk; - mk.aic_kernel_id = FUNC_AIC; - mk.aiv0_kernel_id = FUNC_AIV; - mk.aiv1_kernel_id = FUNC_AIV; - rt_submit_task(mk, args); -} - -} // extern "C" diff --git a/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/kernels/orchestration/acc_c2v_strip_orch.cpp b/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/kernels/orchestration/acc_c2v_strip_orch.cpp deleted file mode 100644 index e6670d24a7..0000000000 --- a/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/kernels/orchestration/acc_c2v_strip_orch.cpp +++ /dev/null @@ -1,71 +0,0 @@ -/* - * 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. - * ----------------------------------------------------------------------------------------------------------- - */ -/** - * Orchestration for Acc ValidShape strip C2V repro. - * One MixedKernels task: C = A @ B via Acc strip TPUSH/TPOP. - * Arg layout: [A, B, C] - */ - -#include - -#include "pto_orchestration_api.h" // NOLINT(build/include_subdir) - -#define FUNC_AIC 0 -#define FUNC_AIV 1 - -static constexpr int M = 128; -static constexpr int N = 256; -static constexpr int K = 32; - -static constexpr uint32_t A_ELEMS = static_cast(M * K); -static constexpr uint32_t B_ELEMS = static_cast(K * N); -static constexpr uint32_t C_ELEMS = static_cast(M * N); - -extern "C" { - -__attribute__((visibility("default"))) PTO2OrchestrationConfig aicpu_orchestration_config(const L2TaskArgs &orch_args) { - (void)orch_args; - return PTO2OrchestrationConfig{ - .expected_arg_count = 3, - }; -} - -__attribute__((visibility("default"))) void aicpu_orchestration_entry(const L2TaskArgs &orch_args) { - const Tensor &ext_A = orch_args.tensor(0).ref(); - const Tensor &ext_B = orch_args.tensor(1).ref(); - const Tensor &ext_C = orch_args.tensor(2).ref(); - - LOG_INFO_V0("[acc_c2v_strip] M=%d N=%d K=%d Acc ValidShape strip C2V repro", M, N, K); - - uint32_t a_shapes[1] = {A_ELEMS}; - uint32_t b_shapes[1] = {B_ELEMS}; - uint32_t c_shapes[1] = {C_ELEMS}; - uint32_t zero[1] = {0}; - - Tensor A_view = ext_A.view(a_shapes, zero); - Tensor B_view = ext_B.view(b_shapes, zero); - Tensor C_view = ext_C.view(c_shapes, zero); - - L0TaskArgs args; - args.add_input(A_view); - args.add_input(B_view); - args.add_output(C_view); - - MixedKernels mk; - mk.aic_kernel_id = FUNC_AIC; - mk.aiv0_kernel_id = FUNC_AIV; - mk.aiv1_kernel_id = FUNC_AIV; - rt_submit_task(mk, args); - - LOG_INFO_V0("[acc_c2v_strip] submitted 1 MixedKernels (AIC+AIV strip C2V)"); -} - -} // extern "C" diff --git a/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/test_acc_c2v_full_control.py b/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/test_acc_c2v_full_control.py deleted file mode 100644 index 8ce3c826ec..0000000000 --- a/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/test_acc_c2v_full_control.py +++ /dev/null @@ -1,80 +0,0 @@ -#!/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. -# ----------------------------------------------------------------------------------------------------------- -"""CONTROL: full Acc C2V (Valid≡Rows) — should PASS with or without the ISA fix. - -Same matmul shapes as the strip repro; only the Acc→Vec transfer differs -(one full TPUSH vs ValidShape strip windows). -""" - -import torch -from simpler.task_interface import ArgDirection as D - -from simpler_setup import SceneTestCase, TaskArgsBuilder, Tensor, scene_test - -M, N, K = 128, 256, 32 - - -@scene_test(level=2, runtime="tensormap_and_ringbuffer") -class TestAccC2VFullControl(SceneTestCase): - RTOL = 5e-2 - ATOL = 0.5 - - CALLABLE = { - "orchestration": { - "source": "kernels/orchestration/acc_c2v_strip_orch.cpp", - "function_name": "aicpu_orchestration_entry", - "signature": [D.IN, D.IN, D.OUT], - }, - "incores": [ - { - "func_id": 0, - "name": "ACC_C2V_FULL_AIC", - "source": "kernels/mix/kernel_acc_c2v_full.cpp", - "core_type": "aic", - "signature": [D.IN, D.IN, D.OUT], - }, - { - "func_id": 1, - "name": "ACC_C2V_FULL_AIV", - "source": "kernels/mix/kernel_acc_c2v_full.cpp", - "core_type": "aiv", - "signature": [D.IN, D.IN, D.OUT], - }, - ], - } - - CASES = [ - { - "name": "default", - "platforms": ["a5"], - "config": {"aicpu_thread_num": 4, "block_dim": 1}, - "params": {}, - } - ] - - def generate_args(self, params): - torch.manual_seed(0) - A = torch.randn(M, K, dtype=torch.float32) * 1.0 - B = torch.randn(K, N, dtype=torch.float32) * 1.0 - C = torch.zeros(M, N, dtype=torch.float32) - return TaskArgsBuilder( - Tensor("A", A.reshape(-1)), - Tensor("B", B.reshape(-1)), - Tensor("C", C.reshape(-1)), - ) - - def compute_golden(self, args, params): - A = args.A.reshape(M, K) - B = args.B.reshape(K, N) - args.C[:] = torch.matmul(A, B).reshape(-1) - - -if __name__ == "__main__": - SceneTestCase.run_module(__name__) diff --git a/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/test_acc_c2v_strip_validshape.py b/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/test_acc_c2v_strip_validshape.py deleted file mode 100644 index b98e7289b3..0000000000 --- a/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/test_acc_c2v_strip_validshape.py +++ /dev/null @@ -1,99 +0,0 @@ -#!/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. -# ----------------------------------------------------------------------------------------------------------- -"""A5 Acc ValidShape strip C2V repro (pto-isa TMovCcToUb srcStride). - -Shapes: Acc[M=128,N=256] drained as 8× ValidShape(H=16,N) TPUSH with -addr = row * 64. Golden is float32 C = A @ B (K=64). - -Expect: - - unfixed pto-isa (srcStride=align(validRow)): FAIL (~90%+ mistmatch after strip0) - - fixed pto-isa (srcStride=align(Rows)): PASS - -Run (from simpler repo root, with installed simpler package):: - - # Fixed ISA (e.g. pypto/build/pto-isa @ 0ebbd03d): - PTO_ISA_ROOT=/path/to/fixed/pto-isa \\ - python examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/test_acc_c2v_strip_validshape.py -p a5 -d 0 - - # Unfixed ISA (checkout before srcStride fix): - PTO_ISA_ROOT=/path/to/unfixed/pto-isa \\ - python .../test_acc_c2v_strip_validshape.py -p a5 -d 0 -""" - -import torch -from simpler.task_interface import ArgDirection as D - -from simpler_setup import SceneTestCase, TaskArgsBuilder, Tensor, scene_test - -M, N, K = 128, 256, 32 -H = 16 - - -@scene_test(level=2, runtime="tensormap_and_ringbuffer") -class TestAccC2VStripValidShape(SceneTestCase): - """Strip Acc C2V vs matmul golden — fails before ISA srcStride=Rows fix.""" - - # Match issue / board thresholds used for Acc C2V strip diagnosis. - RTOL = 5e-2 - ATOL = 0.5 - - CALLABLE = { - "orchestration": { - "source": "kernels/orchestration/acc_c2v_strip_orch.cpp", - "function_name": "aicpu_orchestration_entry", - "signature": [D.IN, D.IN, D.OUT], - }, - "incores": [ - { - "func_id": 0, - "name": "ACC_C2V_STRIP_AIC", - "source": "kernels/mix/kernel_acc_c2v_strip.cpp", - "core_type": "aic", - "signature": [D.IN, D.IN, D.OUT], - }, - { - "func_id": 1, - "name": "ACC_C2V_STRIP_AIV", - "source": "kernels/mix/kernel_acc_c2v_strip.cpp", - "core_type": "aiv", - "signature": [D.IN, D.IN, D.OUT], - }, - ], - } - - CASES = [ - { - "name": "default", - "platforms": ["a5"], - "config": {"aicpu_thread_num": 4, "block_dim": 1}, - "params": {}, - } - ] - - def generate_args(self, params): - # Scale down so Acc values stay moderate for atol/rtol. - torch.manual_seed(0) - A = torch.randn(M, K, dtype=torch.float32) * 1.0 - B = torch.randn(K, N, dtype=torch.float32) * 1.0 - C = torch.zeros(M, N, dtype=torch.float32) - return TaskArgsBuilder( - Tensor("A", A.reshape(-1)), - Tensor("B", B.reshape(-1)), - Tensor("C", C.reshape(-1)), - ) - - def compute_golden(self, args, params): - A = args.A.reshape(M, K) - B = args.B.reshape(K, N) - args.C[:] = torch.matmul(A, B).reshape(-1) - - -if __name__ == "__main__": - SceneTestCase.run_module(__name__) diff --git a/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/test_acc_c2v_strip_vs_full.py b/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/test_acc_c2v_strip_vs_full.py deleted file mode 100644 index 228236679d..0000000000 --- a/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/test_acc_c2v_strip_vs_full.py +++ /dev/null @@ -1,124 +0,0 @@ -#!/usr/bin/env python3 -# Copyright (c) PyPTO Contributors. -"""A5 Acc ValidShape strip C2V repro — device C_strip vs device C_full. - -Drains the same Acc twice (full Valid≡Rows vs ValidShape strip windows). -After the run we require C_strip ≈ C_full. - - unfixed ISA (srcStride=validRow): FAIL (often strip0 OK, strip1+ bad) - fixed ISA (srcStride=Rows): PASS -""" - -import torch -from simpler.task_interface import ArgDirection as D - -from simpler_setup import SceneTestCase, TaskArgsBuilder, Tensor, scene_test -from simpler_setup.scene_test import _build_chip_task_args, _temporary_env - -M, N, K = 32, 128, 32 -H = 16 - - -@scene_test(level=2, runtime="tensormap_and_ringbuffer") -class TestAccC2VStripVsFull(SceneTestCase): - RTOL = 1e-4 - ATOL = 1e-4 - - CALLABLE = { - "orchestration": { - "source": "kernels/orchestration/acc_c2v_compare_orch.cpp", - "function_name": "aicpu_orchestration_entry", - "signature": [D.IN, D.IN, D.OUT, D.OUT], - }, - "incores": [ - { - "func_id": 0, - "name": "ACC_C2V_COMPARE_AIC", - "source": "kernels/mix/kernel_acc_c2v_compare.cpp", - "core_type": "aic", - "signature": [D.IN, D.IN, D.OUT, D.OUT], - }, - { - "func_id": 1, - "name": "ACC_C2V_COMPARE_AIV", - "source": "kernels/mix/kernel_acc_c2v_compare.cpp", - "core_type": "aiv", - "signature": [D.IN, D.IN, D.OUT, D.OUT], - }, - ], - } - - CASES = [ - { - "name": "default", - "platforms": ["a5"], - "config": {"aicpu_thread_num": 4, "block_dim": 1}, - "params": {}, - } - ] - - def generate_args(self, params): - torch.manual_seed(0) - A = torch.randn(M, K, dtype=torch.float32) - B = torch.randn(K, N, dtype=torch.float32) - return TaskArgsBuilder( - Tensor("A", A.reshape(-1)), - Tensor("B", B.reshape(-1)), - Tensor("C_full", torch.zeros(M * N, dtype=torch.float32)), - Tensor("C_strip", torch.zeros(M * N, dtype=torch.float32)), - ) - - def compute_golden(self, args, params): - pass - - def _run_and_validate_l2( # noqa: PLR0913 - self, - worker, - callable_obj, - case, - rounds=1, - skip_golden=False, - enable_l2_swimlane=0, - enable_dump_args=False, - enable_pmu=0, - enable_dep_gen=False, - enable_scope_stats=False, - output_prefix="", - ): - del rounds, skip_golden # single-shot self-compare - params = case.get("params", {}) - config_dict = case.get("config", {}) - orch_sig = self.CALLABLE.get("orchestration", {}).get("signature", []) - handle = getattr(type(self), "_st_l2_handle", None) - if handle is None: - handle = worker.register(callable_obj) - type(self)._st_l2_handle = handle - - test_args = self.generate_args(params) - chip_args, _ = _build_chip_task_args(test_args, orch_sig) - config = self._build_config( - config_dict, - enable_l2_swimlane=enable_l2_swimlane, - enable_dump_args=enable_dump_args, - enable_pmu=enable_pmu, - enable_dep_gen=enable_dep_gen, - enable_scope_stats=enable_scope_stats, - output_prefix=output_prefix, - ) - with _temporary_env(self._resolve_env()): - worker.run(handle, chip_args, config=config) - - full = test_args.C_full.reshape(M, N) - strip = test_args.C_strip.reshape(M, N) - if not torch.allclose(strip, full, rtol=self.RTOL, atol=self.ATOL): - diff = (strip - full).abs().max().item() - s0 = (strip[:H] - full[:H]).abs().max().item() - s1 = (strip[H : 2 * H] - full[H : 2 * H]).abs().max().item() - raise AssertionError( - f"C_strip vs C_full mismatch: max_diff={diff}, " - f"strip0_max={s0}, strip1_max={s1}, rtol={self.RTOL}, atol={self.ATOL}" - ) - - -if __name__ == "__main__": - SceneTestCase.run_module(__name__) diff --git a/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/test_acc_matmul_tstore.py b/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/test_acc_matmul_tstore.py deleted file mode 100644 index 7e19a421be..0000000000 --- a/examples/a5/tensormap_and_ringbuffer/acc_c2v_strip_validshape/test_acc_matmul_tstore.py +++ /dev/null @@ -1,53 +0,0 @@ -#!/usr/bin/env python3 -"""CONTROL: AIC-only Acc matmul + TSTORE (no C2V).""" -import torch -from simpler.task_interface import ArgDirection as D -from simpler_setup import SceneTestCase, TaskArgsBuilder, Tensor, scene_test - -M, N, K = 128, 256, 32 - - -@scene_test(level=2, runtime="tensormap_and_ringbuffer") -class TestAccMatmulTstore(SceneTestCase): - RTOL = 5e-2 - ATOL = 0.5 - - CALLABLE = { - "orchestration": { - "source": "kernels/orchestration/acc_c2v_strip_orch.cpp", - "function_name": "aicpu_orchestration_entry", - "signature": [D.IN, D.IN, D.OUT], - }, - "incores": [ - { - "func_id": 0, - "name": "MATMUL_TSTORE", - "source": "kernels/aic/kernel_matmul_tstore.cpp", - "core_type": "aic", - "signature": [D.IN, D.IN, D.OUT], - }, - ], - } - - CASES = [ - { - "name": "default", - "platforms": ["a5"], - "config": {"aicpu_thread_num": 4, "block_dim": 1}, - "params": {}, - } - ] - - def generate_args(self, params): - torch.manual_seed(0) - A = torch.randn(M, K, dtype=torch.float32) * 1.0 - B = torch.randn(K, N, dtype=torch.float32) * 1.0 - C = torch.zeros(M, N, dtype=torch.float32) - return TaskArgsBuilder(Tensor("A", A.reshape(-1)), Tensor("B", B.reshape(-1)), Tensor("C", C.reshape(-1))) - - def compute_golden(self, args, params): - args.C[:] = torch.matmul(args.A.reshape(M, K), args.B.reshape(K, N)).reshape(-1) - - -if __name__ == "__main__": - SceneTestCase.run_module(__name__)