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 e617415211..1fef4668b9 100644 --- a/src/a2a3/runtime/tensormap_and_ringbuffer/aicpu/aicpu_executor.cpp +++ b/src/a2a3/runtime/tensormap_and_ringbuffer/aicpu/aicpu_executor.cpp @@ -257,13 +257,74 @@ int32_t AicpuExecutor::init(Runtime *runtime) { 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); + // The orchestrator (top thread, tidx == nthreads-1) does not dispatch to + // cores, so it skips the handshake entirely and returns to + // run() immediately to build the graph — overlapping the schedulers' handshake. + // Core counts it needs are derived from cores_total_num_ (fixed 1:2 ratio) in + // run(). The remaining (nthreads-1) threads re-partition ALL cores among + // themselves, so every core still gets its register window opened. + const bool decouple_orch = (nthreads > 1) && !serial_orch_sched_; + const bool is_orchestrator = (tidx == nthreads - 1); + if (decouple_orch && is_orchestrator) { + return 0; // do NOT touch the handshake or hs_arrived_ barrier + } + const int32_t hs_nthreads = decouple_orch ? (nthreads - 1) : nthreads; + + // Barrier-free scheduler init (the decoupled default). Each scheduler thread + // handshakes exactly the clusters it will dispatch to (blocked-layout + // ownership: cluster ci = {ci, N/3+2ci, N/3+2ci+1}, owned by ci % hs_nthreads) + // and self-assigns them, then returns straight to run(). With no all-thread + // barrier a thread starts dispatching to its own cores as soon as they come + // up, independent of peers still handshaking. hs_nthreads == active_sched_threads_ + // in this branch, so handshake ownership matches assign_own_clusters'. + if (decouple_orch) { + sched_ctx_.handshake_owned_clusters(runtime, tidx, hs_nthreads); + sched_ctx_.assign_own_clusters(tidx); +#if SIMPLER_DFX + // Profiling subsystems (pmu/dump/dep) need every core's physical_core_id, + // so gate their one-time leader init behind a barrier — DFX builds only. + hs_arrived_.fetch_add(1, std::memory_order_acq_rel); + if (is_leader) { + while (hs_arrived_.load(std::memory_order_acquire) < hs_nthreads) {} + if (sched_ctx_.handshake_failed()) { + sched_ctx_.abort_and_shutdown(runtime); + init_failed_.store(true, std::memory_order_release); + init_done_.store(true, std::memory_order_release); + return -1; + } + sched_ctx_.post_handshake_profiling_init(); + init_done_.store(true, std::memory_order_release); + } 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; + } +#else + // Perf path: a scheduler that sees an invalid core report (its own or a + // peer's, observed so far) latches completed_ via abort_and_shutdown, which + // stops any peer still entering dispatch (run()'s is_completed() gate). A + // peer that already passed that gate is not joined here — its own cores are + // valid (it handshaked them), and the failure ends in the host device reset + // that reaps every core, so the residual overlap is bounded and + // non-corrupting. finished_count_ is reset per-run in deinit(), not here. + if (sched_ctx_.handshake_failed()) { + sched_ctx_.abort_and_shutdown(runtime); + init_failed_.store(true, std::memory_order_release); + return -1; + } +#endif + return 0; + } - // Barrier: leader waits for every slice to finish, then completes init. + // Serial / single-thread fallback: contiguous handshake + all-thread barrier + + // leader post_handshake_init (the orchestrator participates as a handshaker, + // and hs_nthreads != active_sched_threads_ makes per-thread self-assignment + // unsafe here). + sched_ctx_.handshake_partition(runtime, tidx, hs_nthreads); 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) < hs_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); @@ -599,7 +660,11 @@ int32_t AicpuExecutor::run(Runtime *runtime) { // Fill ops / core counts (host can't resolve s_runtime_ops's // device address nor know the SchedulerContext's core fan-out). - runtime_finalize_after_wire(rt, sched_ctx_.aic_count(), sched_ctx_.aiv_count()); + // Core counts come from cores_total_num_ (fixed 1 AIC : 2 AIV + // cluster ratio). On the decoupled path aic_count()/aiv_count() are + // not populated until after this SM reset, so they cannot be read here. + int32_t spike_total = sched_ctx_.cores_total_num(); + runtime_finalize_after_wire(rt, spike_total / 3, (spike_total * 2) / 3); #if SIMPLER_DFX rt->orchestrator.l2_swimlane_level = get_l2_swimlane_level(); { 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 7f12699be0..fb55ba0186 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 @@ -783,6 +783,188 @@ void SchedulerContext::handshake_partition(Runtime *runtime, int32_t tidx, int32 OUT_OF_ORDER_STORE_BARRIER(); } +// Handshake exactly the cores this scheduler thread will later manage. Blocked +// core layout ([0,N/3) AIC, [N/3,N) AIV) makes ownership predictable before +// handshake: cluster ci = {ci, N/3+2ci, N/3+2ci+1}, assigned to thread +// ci % active_threads. Same protocol as handshake_partition, but over the owned +// set instead of a contiguous slice. +void SchedulerContext::handshake_owned_clusters(Runtime *runtime, int32_t tidx, int32_t active_threads) { + Handshake *all_handshakes = reinterpret_cast(runtime->dev.workers); + const int32_t aic_n = cores_total_num_ / 3; + + int32_t owned[RUNTIME_MAX_WORKER]; + int32_t own_n = 0; + for (int32_t ci = tidx; ci < aic_n; ci += active_threads) { + owned[own_n++] = ci; // AIC + owned[own_n++] = aic_n + 2 * ci; // AIV0 + owned[own_n++] = aic_n + 2 * ci + 1; // AIV1 + } + + uint32_t max_physical_cores_count = platform_get_physical_cores_count(); + uint64_t *regs = reinterpret_cast(regs_); + bool core_serviced[RUNTIME_MAX_WORKER] = {false}; + + // Batched 4-phase handshake (mirrors handshake_partition over the owned set): + // collect reports, then publish tasks / open windows / store CoreExecStates in + // separate passes so posted MMIO STRs and GM stores don't serialize, and only + // two barriers fire for the whole owned set (not one per core). + struct ReadyCore { + int32_t i; + uint32_t pcid; + uint64_t reg_addr; + CoreType core_type; + }; + ReadyCore ready[RUNTIME_MAX_WORKER]; + int32_t n_ready = 0; + + // Phase 1: collect every reported owned core, prefetch its CoreExecState line. + for (int32_t remaining = own_n; remaining > 0;) { + for (int32_t k = 0; k < own_n; k++) { + int32_t i = owned[k]; + if (core_serviced[i]) continue; + Handshake *hank = &all_handshakes[i]; + if (hank->aicore_done == 0) { + SPIN_WAIT_HINT(); + continue; + } + uint32_t physical_core_id = hank->physical_core_id; + if (physical_core_id >= max_physical_cores_count) { + LOG_ERROR( + "Core %d reported invalid physical_core_id=%u (platform max=%u)", i, physical_core_id, + max_physical_cores_count + ); + handshake_failed_.store(true, std::memory_order_release); + core_serviced[i] = true; + remaining--; + continue; + } + __builtin_prefetch(&core_exec_states_[i], 1, 3); + ready[n_ready++] = {i, physical_core_id, regs[physical_core_id], hank->core_type}; + core_serviced[i] = true; + remaining--; + } + } + + // Phase 2: publish every task pointer, then ONE barrier. The core reads task + // only after its window opens (Phase 3), so a single barrier orders all task + // stores before any window STR. + for (int32_t r = 0; r < n_ready; r++) { + all_handshakes[ready[r].i].task = reinterpret_cast(&payload_per_core_[ready[r].i][0]); + } + OUT_OF_ORDER_STORE_BARRIER(); + + // Phase 3: open every window (the IDLE write is also the core's ack). + for (int32_t r = 0; r < n_ready; r++) { + platform_init_aicore_regs(ready[r].reg_addr); + } + + // Phase 4: publish each CoreExecState (AICPU-private, may follow the windows). + // reg_task_id fields start INVALID (pre_handshake_init memset zeroed them, and + // 0 is a valid task id); core_type_compact_ is filled for parity. + for (int32_t r = 0; r < n_ready; r++) { + int32_t i = ready[r].i; + CoreExecState st{}; + st.reg_addr = ready[r].reg_addr; + st.cond_ptr = get_reg_ptr(ready[r].reg_addr, RegId::COND); + st.running_reg_task_id = AICPU_TASK_INVALID; + st.pending_reg_task_id = AICPU_TASK_INVALID; +#if !SIMPLER_DFX + st.worker_id = i; + st.physical_core_id = ready[r].pcid; + st.core_type = ready[r].core_type; +#endif + core_exec_states_[i] = st; + core_type_compact_[i] = static_cast(ready[r].core_type); +#if SIMPLER_DFX + physical_core_ids_[i] = ready[r].pcid; +#endif + } + OUT_OF_ORDER_STORE_BARRIER(); +} + +// ============================================================================= +// Per-thread self-assignment (barrier-free init). Thread tidx owns the clusters +// ci with ci % active_sched_threads_ == tidx (same round-robin as +// assign_cores_to_threads), and the blocked layout gives their worker ids +// directly, so a thread populates its own CoreTracker + per-core payload state +// right after handshaking its own clusters, with no all-thread barrier. +// ============================================================================= +void SchedulerContext::assign_own_clusters(int32_t tidx) { + const int32_t aic_n = cores_total_num_ / 3; + const int32_t active = active_sched_threads_; + + CoreTracker &tracker = core_trackers_[tidx]; + int32_t own_n = 0; + for (int32_t ci = tidx; ci < aic_n; ci += active) + own_n++; + tracker.init(own_n); + + int32_t local = 0; + for (int32_t ci = tidx; ci < aic_n; ci += active) { + tracker.set_cluster(local++, ci, aic_n + 2 * ci, aic_n + 2 * ci + 1); + } + + // Per-cluster GlobalContext sub_block_id (mirrors post_handshake_init) for + // this thread's owned cores only — a thread only ever dispatches to its own. + for (int32_t c = 0; c < tracker.get_cluster_count(); c++) { + int32_t cluster_offset = c * 3; + int32_t aiv0_id = tracker.get_core_id_by_offset(tracker.get_aiv0_core_offset(cluster_offset)); + int32_t aiv1_id = tracker.get_core_id_by_offset(tracker.get_aiv1_core_offset(cluster_offset)); + payload_per_core_[aiv0_id][0].global_context.sub_block_id = 0; + payload_per_core_[aiv0_id][1].global_context.sub_block_id = 0; + payload_per_core_[aiv1_id][0].global_context.sub_block_id = 1; + payload_per_core_[aiv1_id][1].global_context.sub_block_id = 1; + } + + // Per-dispatch AsyncCtx constant prefill + one-time slab clear for owned cores + // (mirrors post_handshake_init's all-core loop, restricted to this thread's). + for (int32_t c = 0; c < tracker.get_cluster_count(); c++) { + for (int32_t sub = 0; sub < 3; sub++) { + int32_t core_id = tracker.get_core_id_by_offset(c * 3 + sub); + for (int32_t buf = 0; buf < 2; buf++) { + PTO2DispatchPayload &dp = payload_per_core_[core_id][buf]; + AsyncCtx &ac = dp.local_context.async_ctx; + volatile DeferredCompletionSlab *slab = &deferred_slab_per_core_[core_id][buf]; + ac.completion_count = &slab->count; + ac.completion_error_code = &slab->error_code; + ac.completion_entries = &slab->entries[0]; + ac.completion_capacity = MAX_COMPLETIONS_PER_TASK; + slab->count = 0; + slab->error_code = PTO2_ERROR_NONE; + dp.args[PAYLOAD_LOCAL_CONTEXT_INDEX] = reinterpret_cast(&dp.local_context); + dp.args[PAYLOAD_GLOBAL_CONTEXT_INDEX] = reinterpret_cast(&dp.global_context); + } + } + } +} + +// Abort the run on a handshake failure discovered without the all-thread barrier +// (non-DFX path): latch completion so every scheduler thread exits its dispatch +// loop, and broadcast exit to whatever cores did come up. Idempotent. +void SchedulerContext::abort_and_shutdown(Runtime *runtime) { + if (!completed_.exchange(true, std::memory_order_acq_rel)) { + emergency_shutdown(runtime); + } +} + +// Profiling-subsystem init (leader-only). pmu_aicpu_init needs every core's +// physical_core_id, so the barrier-free init path calls this behind an +// all-thread barrier compiled only into DFX builds. No-op otherwise. +void SchedulerContext::post_handshake_profiling_init() { +#if SIMPLER_DFX + if (is_dump_args_enabled()) { + dump_args_init(active_sched_threads_); + } + if (is_pmu_enabled()) { + pmu_aicpu_init(physical_core_ids_, cores_total_num_); + LOG_INFO_V0("PMU profiling started on %d cores", cores_total_num_); + } + if (is_dep_gen_enabled()) { + dep_gen_aicpu_init(); + } +#endif +} + // ============================================================================= // Assign discovered cores to scheduler threads (cluster-aligned round-robin). // ============================================================================= @@ -923,10 +1105,47 @@ int32_t SchedulerContext::pre_handshake_init( LOG_ERROR("Invalid cores_total_num %d (expected 1-%d)", cores_total_num_, RUNTIME_MAX_WORKER); return -1; } - aic_count_ = 0; - aiv_count_ = 0; + // The blocked 1 AIC : 2 AIV layout requires an exact multiple of 3: cluster ci + // owns cores {ci, N/3+2ci, N/3+2ci+1}, so a non-zero remainder leaves the tail + // AIV cores [3*(N/3), N) in no cluster — unhandshaked, their windows never open, + // and the run hangs at the op-execute timeout. assign_cores_to_threads pairs + // aiv_worker_ids_[2*ci]/[2*ci+1] on the serial path too, so this holds for both. + if (cores_total_num_ % 3 != 0) { + LOG_ERROR("cores_total_num %d is not a multiple of 3 (blocked 1 AIC : 2 AIV layout)", cores_total_num_); + return -1; + } + // Blocked core layout ([0,N/3) AIC, [N/3,N) AIV) with a fixed 1:2 AIC:AIV + // ratio makes these exact pre-handshake, so scheduler threads can self-assign + // their owned clusters (assign_own_clusters) without the post-handshake + // discovery pass and its all-thread barrier. + aic_count_ = cores_total_num_ / 3; + aiv_count_ = (cores_total_num_ * 2) / 3; + active_sched_threads_ = (sched_thread_num_ > 0) ? sched_thread_num_ : aicpu_thread_num_; handshake_failed_.store(false, std::memory_order_release); + // State the barrier-free per-thread init path no longer reaches via + // post_handshake_init; reset on the leader before any scheduler thread is + // released to dispatch. + completed_tasks_.store(0, std::memory_order_release); + orchestrator_done_.store(false, std::memory_order_release); + func_id_to_addr_ = runtime->dev.func_id_to_addr_; + + // total_tasks_ must be read before hs_setup_done_ is published: on the + // decoupled path the orchestrator resets the SM as soon as it observes + // hs_setup_done_, which zeroes these ring counters, so the read completes here + // (on the leader, before any thread is released) rather than post-handshake. + if (runtime->get_gm_sm_ptr()) { + auto *header = static_cast(runtime->get_gm_sm_ptr()); + int64_t pto2_count = 0; + for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { + int32_t ring_tasks = header->rings[r].fc.current_task_index.load(std::memory_order_acquire); + if (ring_tasks > 0 && ring_tasks <= PTO2_SCOPE_TASKS_CAP) pto2_count += ring_tasks; + } + total_tasks_ = static_cast(pto2_count); + } else { + total_tasks_ = 0; + } + LOG_INFO_V0("Handshaking with %d cores", cores_total_num_); return 0; } @@ -987,25 +1206,8 @@ int32_t SchedulerContext::post_handshake_init(Runtime *runtime) { } #endif - // Initialize task counters. Task count comes from PTO2 shared memory. - if (runtime->get_gm_sm_ptr()) { - auto *header = static_cast(runtime->get_gm_sm_ptr()); - // Read at one-time boot init, before the SM is reset for the run, so a - // ring not yet written holds uninitialized memory (0xbe... under ASAN's - // malloc-fill). Sum in int64 and only count rings whose value is a - // plausible task count — (0, PTO2_SCOPE_TASKS_CAP]; a ring cannot hold - // more than the scope cap. This rejects any garbage pattern (negative - // or positive), so uninitialized rings contribute 0 (the correct boot - // count) while valid counts still add up, with no signed overflow. - int64_t pto2_count = 0; - for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { - int32_t ring_tasks = header->rings[r].fc.current_task_index.load(std::memory_order_acquire); - if (ring_tasks > 0 && ring_tasks <= PTO2_SCOPE_TASKS_CAP) pto2_count += ring_tasks; - } - total_tasks_ = static_cast(pto2_count); - } else { - total_tasks_ = 0; - } + // total_tasks_ is read in pre_handshake_init (before the orchestrator's early + // SM reset on the decoupled path can zero the ring counters). completed_tasks_.store(0, std::memory_order_release); // Device orchestration: the orchestrator thread flips this when the graph is built. 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 c3071f5f57..a6f316173f 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 @@ -71,6 +71,23 @@ class SchedulerContext { // 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); + // Handshake exactly the cores this scheduler thread will later manage: + // clusters {tidx, tidx+active, ...}, cluster ci = + // {ci, N/3+2ci, N/3+2ci+1} (blocked layout: [0,N/3) AIC, [N/3,N) AIV). Matches + // assign_cores_to_threads' round-robin so handshake warms the same + // core_exec_states_ the thread later dispatches from. + void handshake_owned_clusters(Runtime *runtime, int32_t tidx, int32_t active_threads); + // Barrier-free counterpart of assign_cores_to_threads: thread tidx populates + // its own CoreTracker + per-core payload state for the clusters it owns, right + // after handshaking them — no all-thread barrier or leader post_handshake_init. + void assign_own_clusters(int32_t tidx); + // Latch completion + shutdown cores on a handshake failure seen without the + // barrier (non-DFX path). Idempotent. + void abort_and_shutdown(Runtime *runtime); + // Leader-only profiling-subsystem init (DFX builds); called behind a barrier + // in the barrier-free path since pmu_aicpu_init needs all physical_core_ids_. + void post_handshake_profiling_init(); + bool handshake_failed() const { return handshake_failed_.load(std::memory_order_acquire); } // 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); @@ -115,6 +132,7 @@ class SchedulerContext { int32_t aic_count() const { return aic_count_; } int32_t aiv_count() const { return aiv_count_; } + int32_t cores_total_num() const { return cores_total_num_; } bool is_completed() const { return completed_.load(std::memory_order_acquire); } int32_t completed_tasks_count() const { return completed_tasks_.load(std::memory_order_acquire); } bool orchestration_done() const { return orchestrator_done_.load(std::memory_order_relaxed); } diff --git a/src/a5/runtime/tensormap_and_ringbuffer/aicpu/aicpu_executor.cpp b/src/a5/runtime/tensormap_and_ringbuffer/aicpu/aicpu_executor.cpp index 19167d0325..8539ae0f5b 100644 --- a/src/a5/runtime/tensormap_and_ringbuffer/aicpu/aicpu_executor.cpp +++ b/src/a5/runtime/tensormap_and_ringbuffer/aicpu/aicpu_executor.cpp @@ -259,13 +259,74 @@ int32_t AicpuExecutor::init(Runtime *runtime) { 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); + // The orchestrator (top thread, tidx == nthreads-1) does not dispatch to + // cores, so it skips the handshake entirely and returns to + // run() immediately to build the graph — overlapping the schedulers' handshake. + // Core counts it needs are derived from cores_total_num_ (fixed 1:2 ratio) in + // run(). The remaining (nthreads-1) threads re-partition ALL cores among + // themselves, so every core still gets its register window opened. + const bool decouple_orch = (nthreads > 1) && !serial_orch_sched_; + const bool is_orchestrator = (tidx == nthreads - 1); + if (decouple_orch && is_orchestrator) { + return 0; // do NOT touch the handshake or hs_arrived_ barrier + } + const int32_t hs_nthreads = decouple_orch ? (nthreads - 1) : nthreads; + + // Barrier-free scheduler init (the decoupled default). Each scheduler thread + // handshakes exactly the clusters it will dispatch to (blocked-layout + // ownership: cluster ci = {ci, N/3+2ci, N/3+2ci+1}, owned by ci % hs_nthreads) + // and self-assigns them, then returns straight to run(). With no all-thread + // barrier a thread starts dispatching to its own cores as soon as they come + // up, independent of peers still handshaking. hs_nthreads == active_sched_threads_ + // in this branch, so handshake ownership matches assign_own_clusters'. + if (decouple_orch) { + sched_ctx_.handshake_owned_clusters(runtime, tidx, hs_nthreads); + sched_ctx_.assign_own_clusters(tidx); +#if SIMPLER_DFX + // Profiling subsystems (pmu/dump/dep) need every core's physical_core_id, + // so gate their one-time leader init behind a barrier — DFX builds only. + hs_arrived_.fetch_add(1, std::memory_order_acq_rel); + if (is_leader) { + while (hs_arrived_.load(std::memory_order_acquire) < hs_nthreads) {} + if (sched_ctx_.handshake_failed()) { + sched_ctx_.abort_and_shutdown(runtime); + init_failed_.store(true, std::memory_order_release); + init_done_.store(true, std::memory_order_release); + return -1; + } + sched_ctx_.post_handshake_profiling_init(); + init_done_.store(true, std::memory_order_release); + } 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; + } +#else + // Perf path: a scheduler that sees an invalid core report (its own or a + // peer's, observed so far) latches completed_ via abort_and_shutdown, which + // stops any peer still entering dispatch (run()'s is_completed() gate). A + // peer that already passed that gate is not joined here — its own cores are + // valid (it handshaked them), and the failure ends in the host device reset + // that reaps every core, so the residual overlap is bounded and + // non-corrupting. finished_count_ is reset per-run in deinit(), not here. + if (sched_ctx_.handshake_failed()) { + sched_ctx_.abort_and_shutdown(runtime); + init_failed_.store(true, std::memory_order_release); + return -1; + } +#endif + return 0; + } - // Barrier: leader waits for every slice to finish, then completes init. + // Serial / single-thread fallback: contiguous handshake + all-thread barrier + + // leader post_handshake_init (the orchestrator participates as a handshaker, + // and hs_nthreads != active_sched_threads_ makes per-thread self-assignment + // unsafe here). + sched_ctx_.handshake_partition(runtime, tidx, hs_nthreads); 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) < hs_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); @@ -596,7 +657,11 @@ int32_t AicpuExecutor::run(Runtime *runtime) { // Fill ops / core counts (host can't resolve s_runtime_ops's // device address nor know the SchedulerContext's core fan-out). - runtime_finalize_after_wire(rt, sched_ctx_.aic_count(), sched_ctx_.aiv_count()); + // Core counts come from cores_total_num_ (fixed 1 AIC : 2 AIV + // cluster ratio). On the decoupled path aic_count()/aiv_count() are + // not populated until after this SM reset, so they cannot be read here. + int32_t spike_total = sched_ctx_.cores_total_num(); + runtime_finalize_after_wire(rt, spike_total / 3, (spike_total * 2) / 3); #if SIMPLER_DFX rt->orchestrator.l2_swimlane_level = get_l2_swimlane_level(); 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 4f27e542d9..7e245c4d01 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 @@ -786,6 +786,169 @@ void SchedulerContext::handshake_partition(Runtime *runtime, int32_t tidx, int32 OUT_OF_ORDER_STORE_BARRIER(); } +// Handshake exactly the cores this scheduler thread will later manage. Blocked +// core layout ([0,N/3) AIC, [N/3,N) AIV) makes ownership predictable before +// handshake: cluster ci = {ci, N/3+2ci, N/3+2ci+1}, assigned to thread +// ci % active_threads. Same protocol as handshake_partition, but over the owned +// set instead of a contiguous slice. +void SchedulerContext::handshake_owned_clusters(Runtime *runtime, int32_t tidx, int32_t active_threads) { + Handshake *all_handshakes = reinterpret_cast(runtime->dev.workers); + const int32_t aic_n = cores_total_num_ / 3; + + int32_t owned[RUNTIME_MAX_WORKER]; + int32_t own_n = 0; + for (int32_t ci = tidx; ci < aic_n; ci += active_threads) { + owned[own_n++] = ci; // AIC + owned[own_n++] = aic_n + 2 * ci; // AIV0 + owned[own_n++] = aic_n + 2 * ci + 1; // AIV1 + } + + uint32_t max_physical_cores_count = platform_get_physical_cores_count(); + uint64_t *regs = reinterpret_cast(regs_); + bool core_serviced[RUNTIME_MAX_WORKER] = {false}; + + // Batched 4-phase handshake (mirrors handshake_partition over the owned set): + // collect reports, then publish tasks / open windows / store CoreExecStates in + // separate passes so posted MMIO STRs and GM stores don't serialize, and only + // two barriers fire for the whole owned set (not one per core). + struct ReadyCore { + int32_t i; + uint32_t pcid; + uint64_t reg_addr; + CoreType core_type; + }; + ReadyCore ready[RUNTIME_MAX_WORKER]; + int32_t n_ready = 0; + + // Phase 1: collect every reported owned core, prefetch its CoreExecState line. + for (int32_t remaining = own_n; remaining > 0;) { + for (int32_t k = 0; k < own_n; k++) { + int32_t i = owned[k]; + if (core_serviced[i]) continue; + Handshake *hank = &all_handshakes[i]; + if (hank->aicore_done == 0) { + SPIN_WAIT_HINT(); + continue; + } + uint32_t physical_core_id = hank->physical_core_id; + if (physical_core_id >= max_physical_cores_count) { + LOG_ERROR( + "Core %d reported invalid physical_core_id=%u (platform max=%u)", i, physical_core_id, + max_physical_cores_count + ); + handshake_failed_.store(true, std::memory_order_release); + core_serviced[i] = true; + remaining--; + continue; + } + __builtin_prefetch(&core_exec_states_[i], 1, 3); + ready[n_ready++] = {i, physical_core_id, regs[physical_core_id], hank->core_type}; + core_serviced[i] = true; + remaining--; + } + } + + // Phase 2: publish every task pointer, then ONE barrier. The core reads task + // only after its window opens (Phase 3), so a single barrier orders all task + // stores before any window STR. + for (int32_t r = 0; r < n_ready; r++) { + all_handshakes[ready[r].i].task = reinterpret_cast(&payload_per_core_[ready[r].i][0]); + } + OUT_OF_ORDER_STORE_BARRIER(); + + // Phase 3: open every window (the IDLE write is also the core's ack). + for (int32_t r = 0; r < n_ready; r++) { + platform_init_aicore_regs(ready[r].reg_addr); + } + + // Phase 4: publish each CoreExecState (AICPU-private, may follow the windows). + // reg_task_id fields start INVALID (pre_handshake_init memset zeroed them, and + // 0 is a valid task id); core_type_compact_ is filled for parity. + for (int32_t r = 0; r < n_ready; r++) { + int32_t i = ready[r].i; + CoreExecState st{}; + st.reg_addr = ready[r].reg_addr; + st.cond_ptr = get_reg_ptr(ready[r].reg_addr, RegId::COND); + st.running_reg_task_id = AICPU_TASK_INVALID; + st.pending_reg_task_id = AICPU_TASK_INVALID; +#if !SIMPLER_DFX + st.worker_id = i; + st.physical_core_id = ready[r].pcid; + st.core_type = ready[r].core_type; +#endif + core_exec_states_[i] = st; + core_type_compact_[i] = static_cast(ready[r].core_type); +#if SIMPLER_DFX + physical_core_ids_[i] = ready[r].pcid; +#endif + } + OUT_OF_ORDER_STORE_BARRIER(); +} + +// ============================================================================= +// Per-thread self-assignment (barrier-free init). Thread tidx owns the clusters +// ci with ci % active_sched_threads_ == tidx (same round-robin as +// assign_cores_to_threads), and the blocked layout gives their worker ids +// directly, so a thread populates its own CoreTracker + per-core sub_block_id +// right after handshaking its own clusters, with no all-thread barrier. +// ============================================================================= +void SchedulerContext::assign_own_clusters(int32_t tidx) { + const int32_t aic_n = cores_total_num_ / 3; + const int32_t active = active_sched_threads_; + + CoreTracker &tracker = core_trackers_[tidx]; + int32_t own_n = 0; + for (int32_t ci = tidx; ci < aic_n; ci += active) + own_n++; + tracker.init(own_n); + + int32_t local = 0; + for (int32_t ci = tidx; ci < aic_n; ci += active) { + tracker.set_cluster(local++, ci, aic_n + 2 * ci, aic_n + 2 * ci + 1); + } + + // Per-cluster GlobalContext sub_block_id (mirrors post_handshake_init) for + // this thread's owned cores only — a thread only ever dispatches to its own. + // The per-dispatch AsyncCtx / deferred-slab fields are written by build_payload + // at dispatch time (see deinit()'s note), so there is no per-core prefill here. + for (int32_t c = 0; c < tracker.get_cluster_count(); c++) { + int32_t cluster_offset = c * 3; + int32_t aiv0_id = tracker.get_core_id_by_offset(tracker.get_aiv0_core_offset(cluster_offset)); + int32_t aiv1_id = tracker.get_core_id_by_offset(tracker.get_aiv1_core_offset(cluster_offset)); + payload_per_core_[aiv0_id][0].global_context.sub_block_id = 0; + payload_per_core_[aiv0_id][1].global_context.sub_block_id = 0; + payload_per_core_[aiv1_id][0].global_context.sub_block_id = 1; + payload_per_core_[aiv1_id][1].global_context.sub_block_id = 1; + } +} + +// Abort the run on a handshake failure discovered without the all-thread barrier +// (non-DFX path): latch completion so every scheduler thread exits its dispatch +// loop, and broadcast exit to whatever cores did come up. Idempotent. +void SchedulerContext::abort_and_shutdown(Runtime *runtime) { + if (!completed_.exchange(true, std::memory_order_acq_rel)) { + emergency_shutdown(runtime); + } +} + +// Profiling-subsystem init (leader-only). pmu_aicpu_init needs every core's +// physical_core_id, so the barrier-free init path calls this behind an +// all-thread barrier compiled only into DFX builds. No-op otherwise. +void SchedulerContext::post_handshake_profiling_init() { +#if SIMPLER_DFX + if (is_dump_args_enabled()) { + dump_args_init(active_sched_threads_); + } + if (is_pmu_enabled()) { + pmu_aicpu_init(physical_core_ids_, cores_total_num_); + LOG_INFO_V0("PMU profiling started on %d cores", cores_total_num_); + } + if (is_dep_gen_enabled()) { + dep_gen_aicpu_init(); + } +#endif +} + // ============================================================================= // Assign discovered cores to scheduler threads (cluster-aligned round-robin). // ============================================================================= @@ -930,10 +1093,47 @@ int32_t SchedulerContext::pre_handshake_init( LOG_ERROR("Invalid cores_total_num %d (expected 1-%d)", cores_total_num_, RUNTIME_MAX_WORKER); return -1; } - aic_count_ = 0; - aiv_count_ = 0; + // The blocked 1 AIC : 2 AIV layout requires an exact multiple of 3: cluster ci + // owns cores {ci, N/3+2ci, N/3+2ci+1}, so a non-zero remainder leaves the tail + // AIV cores [3*(N/3), N) in no cluster — unhandshaked, their windows never open, + // and the run hangs at the op-execute timeout. assign_cores_to_threads pairs + // aiv_worker_ids_[2*ci]/[2*ci+1] on the serial path too, so this holds for both. + if (cores_total_num_ % 3 != 0) { + LOG_ERROR("cores_total_num %d is not a multiple of 3 (blocked 1 AIC : 2 AIV layout)", cores_total_num_); + return -1; + } + // Blocked core layout ([0,N/3) AIC, [N/3,N) AIV) with a fixed 1:2 AIC:AIV + // ratio makes these exact pre-handshake, so scheduler threads can self-assign + // their owned clusters (assign_own_clusters) without the post-handshake + // discovery pass and its all-thread barrier. + aic_count_ = cores_total_num_ / 3; + aiv_count_ = (cores_total_num_ * 2) / 3; + active_sched_threads_ = (sched_thread_num_ > 0) ? sched_thread_num_ : aicpu_thread_num_; handshake_failed_.store(false, std::memory_order_release); + // State the barrier-free per-thread init path no longer reaches via + // post_handshake_init; reset on the leader before any scheduler thread is + // released to dispatch. + completed_tasks_.store(0, std::memory_order_release); + orchestrator_done_.store(false, std::memory_order_release); + func_id_to_addr_ = runtime->dev.func_id_to_addr_; + + // total_tasks_ must be read before hs_setup_done_ is published: on the + // decoupled path the orchestrator resets the SM as soon as it observes + // hs_setup_done_, which zeroes these ring counters, so the read completes here + // (on the leader, before any thread is released) rather than post-handshake. + if (runtime->get_gm_sm_ptr()) { + auto *header = static_cast(runtime->get_gm_sm_ptr()); + int64_t task_count = 0; + for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { + int32_t ring_tasks = header->rings[r].fc.current_task_index.load(std::memory_order_acquire); + if (ring_tasks > 0 && ring_tasks <= PTO2_SCOPE_TASKS_CAP) task_count += ring_tasks; + } + total_tasks_ = static_cast(task_count); + } else { + total_tasks_ = 0; + } + LOG_INFO_V0("Handshaking with %d cores", cores_total_num_); return 0; } @@ -994,25 +1194,8 @@ int32_t SchedulerContext::post_handshake_init(Runtime *runtime) { } #endif - // Initialize task counters. Task count comes from PTO2 shared memory. - if (runtime->get_gm_sm_ptr()) { - auto *header = static_cast(runtime->get_gm_sm_ptr()); - // Read at one-time boot init, before the SM is reset for the run, so a - // ring not yet written holds uninitialized memory (0xbe... under ASAN's - // malloc-fill). Sum in int64 and only count rings whose value is a - // plausible task count — (0, PTO2_SCOPE_TASKS_CAP]; a ring cannot hold - // more than the scope cap. This rejects any garbage pattern (negative - // or positive), so uninitialized rings contribute 0 (the correct boot - // count) while valid counts still add up, with no signed overflow. - int64_t task_count = 0; - for (int r = 0; r < PTO2_MAX_RING_DEPTH; r++) { - int32_t ring_tasks = header->rings[r].fc.current_task_index.load(std::memory_order_acquire); - if (ring_tasks > 0 && ring_tasks <= PTO2_SCOPE_TASKS_CAP) task_count += ring_tasks; - } - total_tasks_ = static_cast(task_count); - } else { - total_tasks_ = 0; - } + // total_tasks_ is read in pre_handshake_init (before the orchestrator's early + // SM reset on the decoupled path can zero the ring counters). completed_tasks_.store(0, std::memory_order_release); // Device orchestration: the orchestrator thread flips this when the graph is built. 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 4d92a29efa..c60340ef33 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 @@ -68,6 +68,23 @@ class SchedulerContext { // 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); + // Handshake exactly the cores this scheduler thread will later manage: + // clusters {tidx, tidx+active, ...}, cluster ci = + // {ci, N/3+2ci, N/3+2ci+1} (blocked layout: [0,N/3) AIC, [N/3,N) AIV). Matches + // assign_cores_to_threads' round-robin so handshake warms the same + // core_exec_states_ the thread later dispatches from. + void handshake_owned_clusters(Runtime *runtime, int32_t tidx, int32_t active_threads); + // Barrier-free counterpart of assign_cores_to_threads: thread tidx populates + // its own CoreTracker + per-core sub_block_id for the clusters it owns, right + // after handshaking them — no all-thread barrier or leader post_handshake_init. + void assign_own_clusters(int32_t tidx); + // Latch completion + shutdown cores on a handshake failure seen without the + // barrier (non-DFX path). Idempotent. + void abort_and_shutdown(Runtime *runtime); + // Leader-only profiling-subsystem init (DFX builds); called behind a barrier + // in the barrier-free path since pmu_aicpu_init needs all physical_core_ids_. + void post_handshake_profiling_init(); + bool handshake_failed() const { return handshake_failed_.load(std::memory_order_acquire); } // 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); @@ -112,6 +129,7 @@ class SchedulerContext { int32_t aic_count() const { return aic_count_; } int32_t aiv_count() const { return aiv_count_; } + int32_t cores_total_num() const { return cores_total_num_; } bool is_completed() const { return completed_.load(std::memory_order_acquire); } int32_t completed_tasks_count() const { return completed_tasks_.load(std::memory_order_acquire); } bool orchestration_done() const { return orchestrator_done_.load(std::memory_order_relaxed); }