From 651d0f4205b8fc5558305b029c2ad5d40523594f Mon Sep 17 00:00:00 2001 From: Chao Wang <26245345+ChaoWao@users.noreply.github.com> Date: Sun, 31 May 2026 12:31:40 +0800 Subject: [PATCH] Fix: 4 deferred L2 swimlane correctness bugs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Bundled fixes for issues CodeRabbit flagged during #916 review and agreed to defer to a follow-up. All four are tight per-site fixes without behavior change beyond the affected corner case. 1. scheduler_cold_path::init — phase pool count: pass `aicpu_thread_num_` (not `sched_thread_num_`) when `orch_to_sched_=true`. When orch merges into sched, phase records flow through aicpu-thread-indexed pools; the old code under-allocated when `aicpu_thread_num_ > sched_thread_num_`, silently dropping the tail. Matches the existing branch in `scheduler_dispatch.cpp:489` (dump_tensor_init). 2. scheduler_cold_path::init — reset cached level on disabled runs: if `is_l2_swimlane_enabled()` is false this launch, also clear `l2_swimlane_level_` to DISABLED. Without this, a stale value from a prior enabled launch on the same scheduler instance leaks into the `>= SCHED_PHASES` gates in scheduler_dispatch, which then writes into a non-initialized phase pool. 3. onboard/aicore/kernel.cpp::KERNEL_ENTRY — always publish the rotation slot (nullptr when disabled this launch). `[[block_local]]` storage persists across launches on the same loaded binary, so without an explicit nullptr publication the stale slot from a prior enabled launch survives, and `get_l2_swimlane_aicore_rotation()` returns a freed device pointer. Mirrors the sim wrapper's always-publish pattern. 4. l2_swimlane_collector.h::resolve_entry — validate `entry.kind` before branching. Previously any out-of-range value silently fell into the AicoreTask branch and read a wrong-typed pool. The new guard logs and returns nullopt for unknown kinds. Test coverage: existing `test_l2_swimlane` / `test_l2_swimlane_mixed` sim tests pass unchanged. Bug 1 trigger requires `PTO2_ORCH_TO_SCHED=1` + `SCHED_PHASES` level + phase-pool-count assertion, which the current swimlane test infra doesn't model; bugs 2/3 trigger requires enabled→disabled launch cycling within one process. These test gaps are deferred to a separate test-infra PR rather than expanded inline to keep this fix focused. --- .../include/host/l2_swimlane_collector.h | 8 ++++++++ src/a2a3/platform/onboard/aicore/kernel.cpp | 20 +++++++++++++++++-- .../runtime/scheduler/scheduler_cold_path.cpp | 12 +++++++++-- 3 files changed, 36 insertions(+), 4 deletions(-) diff --git a/src/a2a3/platform/include/host/l2_swimlane_collector.h b/src/a2a3/platform/include/host/l2_swimlane_collector.h index a86dfb7a59..9056aedeed 100644 --- a/src/a2a3/platform/include/host/l2_swimlane_collector.h +++ b/src/a2a3/platform/include/host/l2_swimlane_collector.h @@ -111,6 +111,14 @@ struct L2SwimlaneModule { const int num_cores = static_cast(header->num_cores); const L2SwimlaneBufferKind kind = entry.kind; + // Validate kind first — without this, an out-of-range value silently + // falls into the AicoreTask branch below and reads a wrong-typed pool. + if (kind != L2SwimlaneBufferKind::AicpuTask && kind != L2SwimlaneBufferKind::AicpuPhase && + kind != L2SwimlaneBufferKind::AicoreTask) { + LOG_ERROR("L2SwimlaneModule: invalid entry kind=%u", static_cast(kind)); + return std::nullopt; + } + if (kind == L2SwimlaneBufferKind::AicpuPhase) { if (entry.core_index >= static_cast(PLATFORM_MAX_AICPU_THREADS)) { LOG_ERROR("L2SwimlaneModule: invalid phase entry: thread=%u", entry.core_index); diff --git a/src/a2a3/platform/onboard/aicore/kernel.cpp b/src/a2a3/platform/onboard/aicore/kernel.cpp index 5f504f54db..7f0a639734 100644 --- a/src/a2a3/platform/onboard/aicore/kernel.cpp +++ b/src/a2a3/platform/onboard/aicore/kernel.cpp @@ -104,7 +104,21 @@ extern "C" __global__ __aicore__ void KERNEL_ENTRY(aicore_kernel)(__gm__ KernelA // executor runs. AICore reads via get_aicore_profiling_flag() / // get_l2_swimlane_aicore_rotation() — never touches Handshake for profiling. set_aicore_profiling_flag(k_args->enable_profiling_flag); - if (GET_PROFILING_FLAG(k_args->enable_profiling_flag, PROFILING_FLAG_L2_SWIMLANE)) { + // Always publish the rotation slot (nullptr when this launch is disabled + // or has no rotation table). [[block_local]] storage persists across + // launches on the same loaded kernel binary, so without an explicit + // nullptr publication a sequence like enabled(valid)→enabled(NULL table) + // or enabled→disabled would leave `get_l2_swimlane_aicore_rotation()` + // returning the prior launch's freed pointer. AICore call sites are + // additionally flag-gated on this-launch `l2_swimlane_enabled`, so the + // disabled-launch read path is currently unreachable through the + // executors; the unconditional reset is defensive against future call + // sites that don't carry that gate. Mirrors the publish-nullptr branch + // in sim/aicore/kernel.cpp (sim keys only on the table pointer; onboard + // additionally AND-gates on PROFILING_FLAG_L2_SWIMLANE — intentional, + // since the onboard table is shared across collectors). + if (GET_PROFILING_FLAG(k_args->enable_profiling_flag, PROFILING_FLAG_L2_SWIMLANE) && + k_args->l2_swimlane_aicore_rotation_table != 0) { // Stash only the slot pointer. The slot CONTENTS are written by // AICPU's `l2_swimlane_aicpu_init` which runs concurrently with this // entry; dereferencing here would race with AICPU's write. The @@ -113,7 +127,9 @@ extern "C" __global__ __aicore__ void KERNEL_ENTRY(aicore_kernel)(__gm__ KernelA // done and the slot is populated. __gm__ uint64_t *rotation_table = reinterpret_cast<__gm__ uint64_t *>(k_args->l2_swimlane_aicore_rotation_table); - set_l2_swimlane_aicore_rotation_slot(rotation_table != nullptr ? &rotation_table[block_idx] : nullptr); + set_l2_swimlane_aicore_rotation_slot(&rotation_table[block_idx]); + } else { + set_l2_swimlane_aicore_rotation_slot(nullptr); } aicore_execute(k_args->runtime_args, block_idx, core_type); diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_cold_path.cpp b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_cold_path.cpp index 88c8a749be..b5efc63710 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_cold_path.cpp +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_cold_path.cpp @@ -841,13 +841,21 @@ int32_t SchedulerContext::init( // l2_swimlane_aicpu_init promotes g_l2_swimlane_level from the shared-memory // header — must be called BEFORE caching the level, otherwise the cached // value would still be 0 (only the binary enable bit has been seeded by - // kernel.cpp at this point). + // kernel.cpp at this point). Reset the cached level on disabled runs so a + // prior enabled launch's level can't leak into the phase-record gates in + // scheduler_dispatch. if (is_l2_swimlane_enabled()) { l2_swimlane_aicpu_init(runtime->worker_count); l2_swimlane_level_ = get_l2_swimlane_level(); if (l2_swimlane_level_ >= L2SwimlaneLevel::SCHED_PHASES) { - l2_swimlane_aicpu_init_phase(runtime->worker_count, sched_thread_num_); + // When orchestrator phases merge into scheduler threads, phase + // records flow through aicpu_thread_num_ pools — matches the same + // branch in dump_tensor_init (scheduler_dispatch.cpp). + const int phase_threads = orch_to_sched_ ? aicpu_thread_num_ : sched_thread_num_; + l2_swimlane_aicpu_init_phase(runtime->worker_count, phase_threads); } + } else { + l2_swimlane_level_ = L2SwimlaneLevel::DISABLED; } #endif