[Optimization] Parallelize AICore handshake across AICPU threads -- 3~22% Device Speedup

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;

src/a2a3/runtime/tensormap_and_ringbuffer/aicpu/aicpu_executor.cpp

@@ -119,11 +119,20 @@ struct AicpuExecutor {
 
     // ===== Thread management state =====
     std::atomic<int32_t> thread_idx_{0};
-    std::atomic<bool> initialized_{false};
     std::atomic<bool> init_done_{false};
     std::atomic<bool> init_failed_{false};
     std::atomic<bool> finished_{false};
 
+    // Parallel-handshake coordination (see AicpuExecutor::init). hs_setup_done_
+    // is published by the leader once the shared pre-handshake setup is visible;
+    // hs_arrived_ is the barrier counting threads that finished their core slice.
+    // hs_thread_seq_ hands out a distinct [0, nthreads) index when the platform
+    // exposes no affinity idx (sim, where platform_aicpu_affinity_thread_idx()
+    // is -1 during init) so the threads don't all collapse to leader 0.
+    std::atomic<bool> hs_setup_done_{false};
+    std::atomic<int32_t> hs_arrived_{0};
+    std::atomic<int32_t> hs_thread_seq_{0};
+
     int32_t aicpu_thread_num_{0};
 
     // ===== Task queue state (managed by scheduler ready queues) =====
@@ -187,42 +196,88 @@ 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 (tidx 0) does the shared setup, every
+    // thread handshakes a disjoint slice of cores, then the leader finishes init
+    // after a barrier. Non-leaders spin on init_done_.
+    int32_t nthreads = runtime->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);
+    // Each thread needs a distinct index in [0, nthreads) to pick the leader and
+    // partition the cores. Onboard the gate filter assigns it (exec_idx); sim's
+    // gate does not, so platform_aicpu_affinity_thread_idx() is -1 here for every
+    // thread — hand those a distinct index from a counter (mirrors run()'s
+    // thread_idx_++ fallback) instead of collapsing them all to leader 0, which
+    // would run pre_/post_handshake_init on every thread and race the shared
+    // scheduler state. Exactly nthreads threads reach init (the gate drops the
+    // rest), so the counter yields a gap-free [0, nthreads).
+    int32_t tidx = platform_aicpu_affinity_thread_idx();
+    if (tidx < 0) tidx = hs_thread_seq_.fetch_add(1, std::memory_order_acq_rel);
+    // A thread whose index still falls outside [0, nthreads) owns no core slice:
+    // handshake_partition would compute lo/hi past cores_total_num_ and index
+    // all_handshakes[]/core_exec_states_ out of bounds. Reject it here (mirrors
+    // the bounds guard already in run()). Fail only this thread and do NOT set
+    // init_failed_ — that would make the valid peers abort before their
+    // hs_arrived_ increment and hang the leader at the barrier below.
+    if (tidx >= nthreads) {
+        LOG_ERROR("AICPU affinity thread idx %d out of range [0,%d) in init", tidx, nthreads);
         return -1;
     }
+    const bool is_leader = (tidx == 0);
 
-    finished_count_.store(0, std::memory_order_release);
+    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;
 
-    init_done_.store(true, std::memory_order_release);
-    LOG_INFO_V0("AicpuExecutor: Init complete");
+        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;
+    }
+
+    // All threads: handshake this thread's slice of cores in parallel.
+    sched_ctx_.handshake_partition(runtime, tidx, nthreads);
+
+    // Barrier: leader waits for every slice to finish, then completes init.
+    hs_arrived_.fetch_add(1, std::memory_order_acq_rel);
+    if (is_leader) {
+        while (hs_arrived_.load(std::memory_order_acquire) < nthreads) {}
+        finished_count_.store(0, std::memory_order_release);
+        if (sched_ctx_.post_handshake_init(runtime) != 0) {
+            init_failed_.store(true, std::memory_order_release);
+            init_done_.store(true, std::memory_order_release);
+            return -1;
+        }
+        init_done_.store(true, std::memory_order_release);
+        LOG_INFO_V0("AicpuExecutor: Init complete");
+    } else {
+        while (!init_done_.load(std::memory_order_acquire)) {
+            if (init_failed_.load(std::memory_order_acquire)) return -1;
+        }
+        if (init_failed_.load(std::memory_order_acquire)) return -1;
+    }
     return 0;
 }
 
@@ -814,9 +869,11 @@ void AicpuExecutor::deinit(Runtime *runtime) {
 
     LOG_INFO_V0("DeInit: Runtime execution state reset");
 
-    initialized_.store(false, std::memory_order_release);
     init_done_.store(false, std::memory_order_release);
     init_failed_.store(false, std::memory_order_release);
+    hs_setup_done_.store(false, std::memory_order_release);
+    hs_arrived_.store(0, std::memory_order_release);
+    hs_thread_seq_.store(0, std::memory_order_release);
     thread_idx_.store(0, std::memory_order_release);
     finished_.store(false, std::memory_order_release);
 
@@ -855,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
@@ -876,12 +934,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;
         }
     }
 

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`.
 

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<Handshake *>(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<int32_t>((static_cast<int64_t>(tidx) * total) / nthreads);
+    const int32_t hi = static_cast<int32_t>((static_cast<int64_t>(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<uint64_t>(&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<Handshake *>(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).