From aaca7801f84168ef88b65e7dd2147c3bca9cff32 Mon Sep 17 00:00:00 2001 From: s00831018 Date: Mon, 6 Jul 2026 15:19:32 +0200 Subject: [PATCH 1/2] Update: parallelize AICore handshake across AICPU threads The AICPU->AICore handshake in scheduler init discovered all cores serially on a single thread while the other AICPU threads spun idle waiting on init_done_. On a2a3 (72 cores) this serial per-core MMIO (reg init + two spin round-trips per core) measured ~217 us of the ~283 us preamble -- a fixed per-round cost that dominates device time for small workloads and is a flat tax on every run. Parallelize it: split SchedulerContext::init into three parts so every AICPU thread handshakes a disjoint slice of cores. - pre_handshake_init (leader): zero per-core state, wire config, init l2-swimlane buffers, set core count -- published before any thread handshakes so the swimlane-before-aicpu_ready ordering still holds. - handshake_partition (all threads): each thread handshakes its contiguous [lo, hi) slice; core writes are race-free (disjoint), and an invalid physical_core_id sets an atomic flag. - post_handshake_init (leader, after a barrier): build the AIC/AIV worker-id lists in core-index order (identical clusters to before, read from the Handshake struct so it stays correct under PTO2_PROFILING), assign cores, init profiling subsystems, payloads. AicpuExecutor::init orchestrates it: the leader (affinity exec_idx 0) publishes hs_setup_done_, all threads handshake their slice, then an hs_arrived_ barrier gates the leader's finish. init_failed_ is stored before hs_setup_done_ (both release) so a pre-handshake failure is always visible before any thread proceeds. Drops the now-dead initialized_ CAS and replaces the preamble spin with a direct init-return check. Measured on a2a3: preamble ~283 us -> ~85 us, a ~187 us/round device reduction (-7.5% on qwen3-14b decode, up to -22% on small workloads where the fixed handshake cost dominated). Compute window unchanged. Golden-verified across vector_example, benchmark_bgemm, paged_attention (65k tasks), and qwen3_14b_decode. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01MqeALZTPEnDXTnbYfcnPfq --- .../aicore/aicore_executor.cpp | 4 +- .../aicpu/aicpu_executor.cpp | 101 +++++++++++----- .../docs/RUNTIME_LOGIC.md | 2 +- .../runtime/scheduler/scheduler_cold_path.cpp | 114 ++++++++++-------- .../runtime/scheduler/scheduler_context.h | 35 ++++-- 5 files changed, 160 insertions(+), 96 deletions(-) diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/aicore/aicore_executor.cpp b/src/a2a3/runtime/tensormap_and_ringbuffer/aicore/aicore_executor.cpp index f9ec9c97c6..38d7a9d1ea 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/aicore/aicore_executor.cpp +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/aicore/aicore_executor.cpp @@ -103,8 +103,8 @@ __aicore__ __attribute__((weak)) void aicore_execute(__gm__ Runtime *runtime, in bool pmu_enabled = GET_PROFILING_FLAG(enable_profiling_flag, PROFILING_FLAG_PMU); // Per-core L2SwimlaneActiveHead channel. AICPU completes - // `l2_swimlane_aicpu_init` before writing `aicpu_ready = 1` in - // `handshake_all_cores`, and Phase 1 above has already observed + // `l2_swimlane_aicpu_init` (in pre_handshake_init) before any thread writes + // `aicpu_ready = 1` in `handshake_partition`, and Phase 1 above has already observed // `aicpu_ready == 1`, so the rotation-table slot is populated and the // first deref is safe here — off the dispatch→start critical path. __gm__ L2SwimlaneActiveHead *l2_swimlane_head = l2_swimlane_enabled ? get_l2_swimlane_aicore_head() : nullptr; diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/aicpu/aicpu_executor.cpp b/src/a2a3/runtime/tensormap_and_ringbuffer/aicpu/aicpu_executor.cpp index 152f82fc44..8ede034da2 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/aicpu/aicpu_executor.cpp +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/aicpu/aicpu_executor.cpp @@ -119,11 +119,16 @@ struct AicpuExecutor { // ===== Thread management state ===== std::atomic thread_idx_{0}; - std::atomic initialized_{false}; std::atomic init_done_{false}; std::atomic init_failed_{false}; std::atomic finished_{false}; + // Parallel-handshake coordination (see AicpuExecutor::init). hs_setup_done_ + // is published by the leader once the shared pre-handshake setup is visible; + // hs_arrived_ is the barrier counting threads that finished their core slice. + std::atomic hs_setup_done_{false}; + std::atomic hs_arrived_{0}; + int32_t aicpu_thread_num_{0}; // ===== Task queue state (managed by scheduler ready queues) ===== @@ -187,42 +192,73 @@ static_assert( // ===== AicpuExecutor Method Implementations ===== int32_t AicpuExecutor::init(Runtime *runtime) { - bool expected = false; - if (!initialized_.compare_exchange_strong(expected, true, std::memory_order_acq_rel, std::memory_order_acquire)) { - return 0; - } - - LOG_INFO_V0("AicpuExecutor: Initializing"); - if (runtime == nullptr) { LOG_ERROR("runtime is nullptr"); init_failed_.store(true, std::memory_order_release); return -1; } - // Read execution parameters from runtime. The 0 → 1 fixup runs before the - // sched_thread_num_ derivation so a zero input doesn't leave the scheduler - // count at -1. - aicpu_thread_num_ = runtime->dev.aicpu_thread_num; - if (aicpu_thread_num_ == 0) aicpu_thread_num_ = 1; - sched_thread_num_ = aicpu_thread_num_ - 1; - serial_orch_sched_ = runtime->dev.serial_orch_sched; - - if (aicpu_thread_num_ < 1 || aicpu_thread_num_ > MAX_AICPU_THREADS) { - LOG_ERROR("Invalid aicpu_thread_num: %d", aicpu_thread_num_); + // All AICPU threads enter init. The per-core AICore handshake is the + // dominant preamble cost (serial MMIO, ~217 µs of ~283 µs for 72 cores), so + // it is parallelized: the leader (exec_idx 0) does the shared setup, every + // thread handshakes a disjoint slice of cores, then the leader finishes init + // after a barrier. Non-leaders spin on init_done_. The affinity thread idx + // is already valid here (set by the gate filter before aicpu_execute) and is + // used both to pick the leader and to partition cores. + int32_t nthreads = runtime->dev.aicpu_thread_num; + if (nthreads == 0) nthreads = 1; + if (nthreads < 1 || nthreads > MAX_AICPU_THREADS) { + LOG_ERROR("Invalid aicpu_thread_num: %d", nthreads); init_failed_.store(true, std::memory_order_release); return -1; } - - if (sched_ctx_.init(runtime, aicpu_thread_num_, sched_thread_num_, get_platform_regs()) != 0) { - init_failed_.store(true, std::memory_order_release); - return -1; + int32_t tidx = platform_aicpu_affinity_thread_idx(); + if (tidx < 0) tidx = 0; + const bool is_leader = (tidx == 0); + + if (is_leader) { + LOG_INFO_V0("AicpuExecutor: Initializing"); + // The 0 → 1 fixup already applied above; derive scheduler count from it. + aicpu_thread_num_ = nthreads; + sched_thread_num_ = nthreads - 1; + serial_orch_sched_ = runtime->dev.serial_orch_sched; + + hs_arrived_.store(0, std::memory_order_relaxed); + if (sched_ctx_.pre_handshake_init(runtime, aicpu_thread_num_, sched_thread_num_, get_platform_regs()) != 0) { + init_failed_.store(true, std::memory_order_release); + hs_setup_done_.store(true, std::memory_order_release); + return -1; + } + hs_setup_done_.store(true, std::memory_order_release); + } else { + while (!hs_setup_done_.load(std::memory_order_acquire)) { + if (init_failed_.load(std::memory_order_acquire)) return -1; + } + if (init_failed_.load(std::memory_order_acquire)) return -1; } - finished_count_.store(0, std::memory_order_release); + // All threads: handshake this thread's slice of cores in parallel. + sched_ctx_.handshake_partition(runtime, tidx, nthreads); - init_done_.store(true, std::memory_order_release); - LOG_INFO_V0("AicpuExecutor: Init complete"); + // Barrier: leader waits for every slice to finish, then completes init. + hs_arrived_.fetch_add(1, std::memory_order_acq_rel); + if (is_leader) { + while (hs_arrived_.load(std::memory_order_acquire) < nthreads) { + } + finished_count_.store(0, std::memory_order_release); + if (sched_ctx_.post_handshake_init(runtime) != 0) { + init_failed_.store(true, std::memory_order_release); + init_done_.store(true, std::memory_order_release); + return -1; + } + init_done_.store(true, std::memory_order_release); + LOG_INFO_V0("AicpuExecutor: Init complete"); + } else { + while (!init_done_.load(std::memory_order_acquire)) { + if (init_failed_.load(std::memory_order_acquire)) return -1; + } + if (init_failed_.load(std::memory_order_acquire)) return -1; + } return 0; } @@ -814,9 +850,10 @@ void AicpuExecutor::deinit(Runtime *runtime) { LOG_INFO_V0("DeInit: Runtime execution state reset"); - initialized_.store(false, std::memory_order_release); init_done_.store(false, std::memory_order_release); init_failed_.store(false, std::memory_order_release); + hs_setup_done_.store(false, std::memory_order_release); + hs_arrived_.store(0, std::memory_order_release); thread_idx_.store(0, std::memory_order_release); finished_.store(false, std::memory_order_release); @@ -876,12 +913,12 @@ extern "C" int32_t aicpu_execute(Runtime *runtime) { // rc / runtime_rc are declared out here because they outlive their phase. { AicpuPhaseScope preamble(AicpuPhase::Preamble); - g_aicpu_executor.init(runtime); - while (!g_aicpu_executor.init_done_.load(std::memory_order_acquire)) { - if (g_aicpu_executor.init_failed_.load(std::memory_order_acquire)) { - LOG_ERROR("%s", "aicpu_execute: Initialization failed, aborting execution"); - return -1; - } + // init() barriers every thread internally until init is complete on the + // leader (or a thread failed), then returns the status — so a non-zero + // return is authoritative on all threads and no extra spin is needed. + if (g_aicpu_executor.init(runtime) != 0) { + LOG_ERROR("%s", "aicpu_execute: Initialization failed, aborting execution"); + return -1; } } diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md b/src/a2a3/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md index bcc50f25d7..9bbd9f8bfc 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md @@ -599,7 +599,7 @@ Private internals are split across three .cpp files by responsibility: - `scheduler_completion.cpp` — completion polling, drain protocol - `scheduler_dispatch.cpp` — task dispatch loop and helpers -- `scheduler_cold_path.cpp` — exit checks, stall diagnostics, profiling, lifecycle (`init/deinit`), core management (`handshake_all_cores` / `assign_cores_to_threads` / `emergency_shutdown`), and `on_orchestration_done` +- `scheduler_cold_path.cpp` — exit checks, stall diagnostics, profiling, lifecycle (`pre_handshake_init` / `handshake_partition` / `post_handshake_init` / `deinit`), core management (`assign_cores_to_threads` / `emergency_shutdown`), and `on_orchestration_done` `AicpuExecutor` calls neither `handshake_*`, `assign_*`, `reassign_*`, nor `emergency_shutdown` directly — they are private, invoked only by `init` and `on_orchestration_done`. 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 8e1813367f..8891b90752 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 @@ -669,27 +669,23 @@ int32_t SchedulerContext::shutdown(int32_t thread_idx) { } // ============================================================================= -// Handshake with all AICore workers; discover core type and reg address. +// Handshake a contiguous slice of AICore workers. Runs on every AICPU thread in +// parallel (partitioned by tidx/nthreads); the leader's pre_handshake_init has +// already zeroed state, set cores_total_num_, and reset the counts/flag. The +// serial per-core MMIO here is what dominates preamble, so splitting the slice +// across threads is the whole point. Worker-id lists are built serially in +// post_handshake_init (core-index order) once every slice has landed. // ============================================================================= -int32_t SchedulerContext::handshake_all_cores(Runtime *runtime) { +void SchedulerContext::handshake_partition(Runtime *runtime, int32_t tidx, int32_t nthreads) { Handshake *all_handshakes = reinterpret_cast(runtime->dev.workers); - cores_total_num_ = runtime->dev.worker_count; - - // Validate cores_total_num_ before using as array index - if (cores_total_num_ == 0 || cores_total_num_ > RUNTIME_MAX_WORKER) { - LOG_ERROR("Invalid cores_total_num %d (expected 1-%d)", cores_total_num_, RUNTIME_MAX_WORKER); - return -1; - } - - aic_count_ = 0; - aiv_count_ = 0; - - LOG_INFO_V0("Handshaking with %d cores", cores_total_num_); + const int32_t total = cores_total_num_; + const int32_t lo = static_cast((static_cast(tidx) * total) / nthreads); + const int32_t hi = static_cast((static_cast(tidx + 1) * total) / nthreads); // Step 1: Write per-core payload addresses and send handshake signal. // OUT_OF_ORDER_STORE_BARRIER() ensures task is globally visible before // aicpu_ready=1, so AICore reads the correct payload pointer after waking up. - for (int32_t i = 0; i < cores_total_num_; i++) { + for (int32_t i = lo; i < hi; i++) { all_handshakes[i].task = reinterpret_cast(&payload_per_core_[i][0]); OUT_OF_ORDER_STORE_BARRIER(); all_handshakes[i].aicpu_ready = 1; @@ -699,9 +695,8 @@ int32_t SchedulerContext::handshake_all_cores(Runtime *runtime) { // Get platform physical cores count for validation uint32_t max_physical_cores_count = platform_get_physical_cores_count(); - // Step 2: Wait for all cores to respond, collect core type and register addresses - bool handshake_failed = false; - for (int32_t i = 0; i < cores_total_num_; i++) { + // Step 2: Wait for this slice's cores, init registers, collect state. + for (int32_t i = lo; i < hi; i++) { Handshake *hank = &all_handshakes[i]; while (hank->aicore_regs_ready == 0) { @@ -715,7 +710,7 @@ int32_t SchedulerContext::handshake_all_cores(Runtime *runtime) { "Core %d reported invalid physical_core_id=%u (platform max=%u)", i, physical_core_id, max_physical_cores_count ); - handshake_failed = true; + handshake_failed_.store(true, std::memory_order_release); continue; } @@ -733,8 +728,6 @@ int32_t SchedulerContext::handshake_all_cores(Runtime *runtime) { SPIN_WAIT_HINT(); } - CoreType type = hank->core_type; - core_exec_states_[i].reg_addr = reg_addr; core_exec_states_[i].cond_ptr = get_reg_ptr(reg_addr, RegId::COND); @@ -746,25 +739,9 @@ int32_t SchedulerContext::handshake_all_cores(Runtime *runtime) { #if !PTO2_PROFILING core_exec_states_[i].worker_id = i; core_exec_states_[i].physical_core_id = physical_core_id; - core_exec_states_[i].core_type = type; + core_exec_states_[i].core_type = hank->core_type; #endif - - if (type == CoreType::AIC) { - aic_worker_ids_[aic_count_++] = i; - LOG_INFO_V0("Core %d: AIC, physical_id=%u, reg_addr=0x%lx", i, physical_core_id, reg_addr); - } else { - aiv_worker_ids_[aiv_count_++] = i; - LOG_INFO_V0("Core %d: AIV, physical_id=%u, reg_addr=0x%lx", i, physical_core_id, reg_addr); - } } - - if (handshake_failed) { - emergency_shutdown(runtime); - return -1; - } - - LOG_INFO_V0("Core discovery complete: %d AIC, %d AIV", aic_count_, aiv_count_); - return 0; } // ============================================================================= @@ -858,8 +835,9 @@ void SchedulerContext::emergency_shutdown(Runtime *runtime) { // ============================================================================= // Lifecycle: init / deinit // ============================================================================= -int32_t -SchedulerContext::init(Runtime *runtime, int32_t aicpu_thread_num, int32_t sched_thread_num, uint64_t regs_base) { +int32_t SchedulerContext::pre_handshake_init( + Runtime *runtime, int32_t aicpu_thread_num, int32_t sched_thread_num, uint64_t regs_base +) { always_assert(runtime != nullptr); // Zero all per-core execution state before handshake @@ -876,7 +854,9 @@ SchedulerContext::init(Runtime *runtime, int32_t aicpu_thread_num, int32_t sched // value would still be 0 (only the binary enable bit has been seeded by // kernel.cpp at this point). Reset the cached level on disabled runs so a // prior enabled launch's level can't leak into the phase-record gates in - // scheduler_dispatch. + // scheduler_dispatch. This runs on the leader before it publishes + // hs_setup_done_, so it happens-before every thread's handshake_partition + // (and therefore before any aicpu_ready=1 write). if (is_l2_swimlane_enabled()) { l2_swimlane_aicpu_init(runtime->dev.worker_count); l2_swimlane_level_ = get_l2_swimlane_level(); @@ -899,19 +879,55 @@ SchedulerContext::init(Runtime *runtime, int32_t aicpu_thread_num, int32_t sched } #endif - // Discover cores and assign to scheduler threads. - int32_t rc = handshake_all_cores(runtime); - if (rc != 0) { - LOG_ERROR("handshake_all_cores failed"); - return rc; + // Core count is needed by every thread to compute its handshake slice. + cores_total_num_ = runtime->dev.worker_count; + if (cores_total_num_ == 0 || cores_total_num_ > RUNTIME_MAX_WORKER) { + LOG_ERROR("Invalid cores_total_num %d (expected 1-%d)", cores_total_num_, RUNTIME_MAX_WORKER); + return -1; + } + aic_count_ = 0; + aiv_count_ = 0; + handshake_failed_.store(false, std::memory_order_release); + + LOG_INFO_V0("Handshaking with %d cores", cores_total_num_); + return 0; +} + +int32_t SchedulerContext::post_handshake_init(Runtime *runtime) { + if (handshake_failed_.load(std::memory_order_acquire)) { + emergency_shutdown(runtime); + return -1; + } + + // Build cluster-ordered AIC/AIV worker-id lists from the discovered cores. + // Serial and MMIO-free — the expensive per-core handshake already ran in + // parallel across threads. Core-index order matches the original + // single-thread handshake so assign_cores_to_threads forms identical + // clusters. core_type / physical_core_id are read from the Handshake struct + // (stable after the handshake) so this stays correct under PTO2_PROFILING, + // where they are not mirrored into CoreExecState. + Handshake *all_handshakes = reinterpret_cast(runtime->dev.workers); + for (int32_t i = 0; i < cores_total_num_; i++) { + CoreType type = all_handshakes[i].core_type; + uint32_t physical_core_id = all_handshakes[i].physical_core_id; + uint64_t reg_addr = core_exec_states_[i].reg_addr; + if (type == CoreType::AIC) { + aic_worker_ids_[aic_count_++] = i; + LOG_INFO_V0("Core %d: AIC, physical_id=%u, reg_addr=0x%lx", i, physical_core_id, reg_addr); + } else { + aiv_worker_ids_[aiv_count_++] = i; + LOG_INFO_V0("Core %d: AIV, physical_id=%u, reg_addr=0x%lx", i, physical_core_id, reg_addr); + } } + LOG_INFO_V0("Core discovery complete: %d AIC, %d AIV", aic_count_, aiv_count_); + if (!assign_cores_to_threads()) { return -1; } - // Profiling-subsystem buffer/state init: single-threaded cold path, so the - // "do it once" guarantee is structural (no CAS needed). Runs after - // handshake_all_cores / assign_cores_to_threads because pmu_aicpu_init needs + // Profiling-subsystem buffer/state init: single-threaded cold path (leader + // only), so the "do it once" guarantee is structural (no CAS needed). Runs + // after the handshake / assign_cores_to_threads because pmu_aicpu_init needs // physical_core_ids_ / cores_total_num_. Mirrors the l2_swimlane_aicpu_init // convention above; the per-thread *_set_orch_thread_idx setters stay on the // orchestrator thread (see aicpu_executor.cpp). diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_context.h b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_context.h index 9ce98ac5fa..67f57302fc 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_context.h +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_context.h @@ -54,14 +54,24 @@ class SchedulerContext { // Lifecycle // ========================================================================= - // Initialize scheduler state from the given runtime and thread layout. - // - Discovers cores via handshake_all_cores() - // - Assigns cores to scheduler threads - // - Resets task counters, payloads, per-core GlobalContext - // - Binds func_id_to_addr_ / initial sched_ (if rt is already known) - // - Captures AICore-register base (consumed by handshake_all_cores()) - // Returns 0 on success, negative on failure (handshake / assignment error). - int32_t init(Runtime *runtime, int32_t aicpu_thread_num, int32_t sched_thread_num, uint64_t regs_base); + // Initialize scheduler state from the given runtime and thread layout. Split + // into three parts so the per-core AICore handshake — a serial, MMIO-bound + // loop that dominates preamble (~217 µs of ~283 µs for 72 cores) — can run in + // parallel across all AICPU threads. Orchestrated by AicpuExecutor::init: + // the leader runs pre_handshake_init, every thread handshakes a disjoint + // slice of cores via handshake_partition, then the leader runs + // post_handshake_init after a barrier. + // + // Leader-only: per-core state + config + swimlane buffers + core count. Must + // be published before any thread enters handshake_partition. Returns 0 on + // success, negative on failure. + int32_t pre_handshake_init(Runtime *runtime, int32_t aicpu_thread_num, int32_t sched_thread_num, uint64_t regs_base); + // All threads: handshake this thread's contiguous slice [lo, hi) of cores + // (partitioned by tidx/nthreads). Each core is touched by exactly one thread. + void handshake_partition(Runtime *runtime, int32_t tidx, int32_t nthreads); + // Leader-only, after the handshake barrier: build worker-id lists, assign + // cores, init profiling subsystems, read task counts, init payloads. + int32_t post_handshake_init(Runtime *runtime); // Reset all SchedulerContext-owned state to its post-construction defaults. // Called by AicpuExecutor::deinit() during per-run teardown. @@ -157,12 +167,16 @@ class SchedulerContext { int32_t aicpu_thread_num_{0}; int32_t cores_total_num_{0}; - // Cluster-ordered worker_id lists, populated by handshake_all_cores(). + // Cluster-ordered worker_id lists, populated by post_handshake_init(). int32_t aic_worker_ids_[RUNTIME_MAX_WORKER]{}; int32_t aiv_worker_ids_[RUNTIME_MAX_WORKER]{}; int32_t aic_count_{0}; int32_t aiv_count_{0}; + // Set by any thread whose slice hits an invalid physical_core_id in + // handshake_partition; checked by the leader in post_handshake_init. + std::atomic handshake_failed_{false}; + // Platform AICore-register base array (set by AicpuExecutor before init()). uint64_t regs_{0}; @@ -177,9 +191,6 @@ class SchedulerContext { // Core management (scheduler_cold_path.cpp) // ========================================================================= - // Handshake with all AICore workers; populates core_exec_states_, worker id lists. - int32_t handshake_all_cores(Runtime *runtime); - // Assign discovered cores (cluster = 1 AIC + 2 AIV) round-robin across scheduler threads. bool assign_cores_to_threads(); From 18fee9cf2014d9cd75e4a28da097b50a788f3b2a Mon Sep 17 00:00:00 2001 From: Chao Wang <26245345+ChaoWao@users.noreply.github.com> Date: Wed, 8 Jul 2026 23:19:55 -0700 Subject: [PATCH 2/2] fix(pr): harden AICPU parallel-init thread indexing for #1279 - Sim regression (segfault): in a2a3sim, platform_aicpu_affinity_thread_idx() returns -1 during init() (sim gates with platform_aicpu_affinity_gate, which never sets the exec idx; set_thread_idx runs later in run()). The prior `if (tidx<0) tidx=0` made every sim AICPU thread the leader, so all of them ran pre_/post_handshake_init concurrently and raced the shared scheduler state (aic_count_++, worker-list writes, memset) -> memory corruption -> SIGSEGV on the simplest tensormap sim runs (vector_example, per_task_runtime_env). Fix: hand a distinct [0,nthreads) index from a counter (hs_thread_seq_, reset in deinit) when affinity is unavailable, mirroring run()'s thread_idx_++ fallback, so sim gets one leader + disjoint slices like onboard. - Add the upper-bound guard init() was missing vs run(): a tidx >= nthreads would index the partition arrays out of bounds. Fail only that thread (not init_failed_, which would hang the leader at the barrier). - Refresh the now-stale aicpu_execute docstring (init is no longer "first thread only" with an external wait). Co-Authored-By: Claude Opus 4.8 (1M context) --- .../aicpu/aicpu_executor.cpp | 43 ++++++++++++++----- .../runtime/scheduler/scheduler_context.h | 3 +- 2 files changed, 34 insertions(+), 12 deletions(-) diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/aicpu/aicpu_executor.cpp b/src/a2a3/runtime/tensormap_and_ringbuffer/aicpu/aicpu_executor.cpp index 8ede034da2..b9cc75f6bb 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/aicpu/aicpu_executor.cpp +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/aicpu/aicpu_executor.cpp @@ -126,8 +126,12 @@ struct AicpuExecutor { // Parallel-handshake coordination (see AicpuExecutor::init). hs_setup_done_ // is published by the leader once the shared pre-handshake setup is visible; // hs_arrived_ is the barrier counting threads that finished their core slice. + // hs_thread_seq_ hands out a distinct [0, nthreads) index when the platform + // exposes no affinity idx (sim, where platform_aicpu_affinity_thread_idx() + // is -1 during init) so the threads don't all collapse to leader 0. std::atomic hs_setup_done_{false}; std::atomic hs_arrived_{0}; + std::atomic hs_thread_seq_{0}; int32_t aicpu_thread_num_{0}; @@ -200,11 +204,9 @@ int32_t AicpuExecutor::init(Runtime *runtime) { // All AICPU threads enter init. The per-core AICore handshake is the // dominant preamble cost (serial MMIO, ~217 µs of ~283 µs for 72 cores), so - // it is parallelized: the leader (exec_idx 0) does the shared setup, every + // it is parallelized: the leader (tidx 0) does the shared setup, every // thread handshakes a disjoint slice of cores, then the leader finishes init - // after a barrier. Non-leaders spin on init_done_. The affinity thread idx - // is already valid here (set by the gate filter before aicpu_execute) and is - // used both to pick the leader and to partition cores. + // after a barrier. Non-leaders spin on init_done_. int32_t nthreads = runtime->dev.aicpu_thread_num; if (nthreads == 0) nthreads = 1; if (nthreads < 1 || nthreads > MAX_AICPU_THREADS) { @@ -212,8 +214,26 @@ int32_t AicpuExecutor::init(Runtime *runtime) { init_failed_.store(true, std::memory_order_release); return -1; } + // Each thread needs a distinct index in [0, nthreads) to pick the leader and + // partition the cores. Onboard the gate filter assigns it (exec_idx); sim's + // gate does not, so platform_aicpu_affinity_thread_idx() is -1 here for every + // thread — hand those a distinct index from a counter (mirrors run()'s + // thread_idx_++ fallback) instead of collapsing them all to leader 0, which + // would run pre_/post_handshake_init on every thread and race the shared + // scheduler state. Exactly nthreads threads reach init (the gate drops the + // rest), so the counter yields a gap-free [0, nthreads). int32_t tidx = platform_aicpu_affinity_thread_idx(); - if (tidx < 0) tidx = 0; + if (tidx < 0) tidx = hs_thread_seq_.fetch_add(1, std::memory_order_acq_rel); + // A thread whose index still falls outside [0, nthreads) owns no core slice: + // handshake_partition would compute lo/hi past cores_total_num_ and index + // all_handshakes[]/core_exec_states_ out of bounds. Reject it here (mirrors + // the bounds guard already in run()). Fail only this thread and do NOT set + // init_failed_ — that would make the valid peers abort before their + // hs_arrived_ increment and hang the leader at the barrier below. + if (tidx >= nthreads) { + LOG_ERROR("AICPU affinity thread idx %d out of range [0,%d) in init", tidx, nthreads); + return -1; + } const bool is_leader = (tidx == 0); if (is_leader) { @@ -243,8 +263,7 @@ int32_t AicpuExecutor::init(Runtime *runtime) { // Barrier: leader waits for every slice to finish, then completes init. hs_arrived_.fetch_add(1, std::memory_order_acq_rel); if (is_leader) { - while (hs_arrived_.load(std::memory_order_acquire) < nthreads) { - } + while (hs_arrived_.load(std::memory_order_acquire) < nthreads) {} finished_count_.store(0, std::memory_order_release); if (sched_ctx_.post_handshake_init(runtime) != 0) { init_failed_.store(true, std::memory_order_release); @@ -854,6 +873,7 @@ void AicpuExecutor::deinit(Runtime *runtime) { init_failed_.store(false, std::memory_order_release); hs_setup_done_.store(false, std::memory_order_release); hs_arrived_.store(0, std::memory_order_release); + hs_thread_seq_.store(0, std::memory_order_release); thread_idx_.store(0, std::memory_order_release); finished_.store(false, std::memory_order_release); @@ -892,10 +912,11 @@ extern "C" __attribute__((visibility("default"))) int simpler_aicpu_register_cal * * This is called by DynTileFwkBackendKernelServer in kernel.cpp. * Orchestrates the complete task runtime execution: - * 1. Initialize executor (thread-safe, first thread only) - * 2. Wait for initialization to complete - * 3. Execute tasks on managed cores - * 4. Cleanup when last thread finishes + * 1. Initialize executor: all threads enter init(), which handshakes the cores + * in parallel and barriers internally until init is complete (or a thread + * failed); its return value is authoritative on every thread. + * 2. Execute tasks on managed cores + * 3. Cleanup when last thread finishes * * @param runtime Pointer to Runtime structure * @return 0 on success, non-zero on error diff --git a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_context.h b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_context.h index 67f57302fc..b35b93cbaf 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_context.h +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_context.h @@ -65,7 +65,8 @@ class SchedulerContext { // Leader-only: per-core state + config + swimlane buffers + core count. Must // be published before any thread enters handshake_partition. Returns 0 on // success, negative on failure. - int32_t pre_handshake_init(Runtime *runtime, int32_t aicpu_thread_num, int32_t sched_thread_num, uint64_t regs_base); + int32_t + pre_handshake_init(Runtime *runtime, int32_t aicpu_thread_num, int32_t sched_thread_num, uint64_t regs_base); // All threads: handshake this thread's contiguous slice [lo, hi) of cores // (partitioned by tidx/nthreads). Each core is touched by exactly one thread. void handshake_partition(Runtime *runtime, int32_t tidx, int32_t nthreads);