[Optimization] Replace wiring with polling-based task readiness test (~17% median device speedup)

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

@@ -257,74 +257,44 @@ int32_t AicpuExecutor::init(Runtime *runtime) {
         if (init_failed_.load(std::memory_order_acquire)) return -1;
     }
 
-    // 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) {}
+    // Multi-thread runs take the barrier-free init path: the orchestrator (last
+    // thread) skips the handshake and returns to build the graph immediately, and
+    // each scheduler thread handshakes + self-assigns only the clusters it will
+    // dispatch to (ci % sched_thread_num_ == tidx) — no all-thread barrier and no
+    // leader-serialized post_handshake_init. DFX-safe: polling inits PMU / dep-gen
+    // per-thread inside the dispatch loop (no global pmu_aicpu_init over all cores),
+    // so unlike #1345 upstream there is no leader-only profiling init needing a
+    // barrier. Single-thread runs keep the original path (no orch/scheduler split).
+    const bool barrier_free = (nthreads >= 2);
+    if (barrier_free) {
+        const bool is_orchestrator = (tidx == sched_thread_num_);
+        if (!is_orchestrator) {
+            sched_ctx_.handshake_owned_clusters(runtime, tidx, sched_thread_num_);
+            sched_ctx_.assign_own_clusters(tidx);
             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;
+        // No thread blocks: schedulers dispatch their own cores as soon as those
+        // cores are up; the orchestrator overlaps graph-build with peer handshakes.
+        if (is_leader) {
+            finished_count_.store(0, std::memory_order_release);
+            init_done_.store(true, std::memory_order_release);
+            LOG_INFO_V0("AicpuExecutor: Init complete (barrier-free)");
         }
-#endif
         return 0;
     }
 
-    // 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);
+    // 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) < hs_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);
@@ -660,11 +630,7 @@ 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).
-                // 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);
+                runtime_finalize_after_wire(rt, sched_ctx_.aic_count(), sched_ctx_.aiv_count());
 #if SIMPLER_DFX
                 rt->orchestrator.l2_swimlane_level = get_l2_swimlane_level();
                 {

src/a2a3/runtime/tensormap_and_ringbuffer/common/intrinsic.h

@@ -63,7 +63,7 @@
  *     compiled, ran without error, and produced wrong output. Use
  *     `get_sub_block_id(args)` instead, which reads from the runtime's
  *     `GlobalContext.sub_block_id` that the scheduler initializes per
- *     AIV core in `scheduler_cold_path.cpp::SchedulerContext::init`.
+ *     AIV core in `scheduler_context.h::SchedulerContext::init`.
  *
  *   - `get_block_idx()` and `get_block_num()` are not redirected to
  *     simpler's LocalContext either — use the `(args)` variants below
@@ -97,7 +97,7 @@ static constexpr int32_t PTO2_EXT_PARAMS_COUNT = 2;
 
 /**
  * Args[] suffix indices for context pointers.
- * Derived from MAX_TENSOR_ARGS(32) + MAX_SCALAR_ARGS(16).
+ * Derived from MAX_TENSOR_ARGS(16) + MAX_SCALAR_ARGS(32).
  * Users should not depend on these values; use the Get* functions below.
  */
 static constexpr int32_t SPMD_LOCAL_CONTEXT_INDEX = 48;

src/a2a3/runtime/tensormap_and_ringbuffer/docs/MULTI_RING.md

@@ -178,8 +178,9 @@ Each ring's `last_task_alive` advances independently:
 
 ```text
 advance_ring_pointers(ring_id):  // protected by per-ring advance_lock
-    la = ring->fc.last_task_alive
-    while ring->get_slot_state_by_task_id(la).task_state >= CONSUMED:
+    watermark = ring->completed_watermark
+    la = last_task_alive
+    while la <= watermark and watermark >= slot[la].last_consumer_local_id:
         reset slot for reuse
         la++
     sync_to_sm()  // release-store last_task_alive
@@ -234,91 +235,25 @@ AICore uses `last_reg_val` to detect new dispatches — identical values cause s
 | `PTO2_HEAP_SIZE` | 256 MB | 1 GB |
 | `PTO2_DEP_LIST_POOL_SIZE` | 16384 | 65536 |
 
-### 7.2 Runtime Overrides
-
-Each ring resource (`ring_task_window` / `ring_heap` / `ring_dep_pool`) is a
-single `CallConfig.runtime_env` field that accepts **either** a scalar (broadcast
-to every ring) **or** a list of four per-ring values. Precedence is resolved
-independently for each resource and ring:
-
-```text
-per-ring CallConfig entry (a scalar is broadcast to every entry)
-  > per-ring PTO2_RING_* env value
-  > scalar PTO2_RING_* env value
-  > compile-time default
-```
-
-`ring_id` is the scope-depth ring selected by the runtime:
-
-```text
-scope depth 0 -> ring 0
-scope depth 1 -> ring 1
-scope depth 2 -> ring 2
-scope depth >=3 -> ring 3
-```
+### 7.2 Runtime Environment Overrides
 
-Per-task via `CallConfig.runtime_env` — different L2 tasks in one launch can
-each carry their own sizes. Invalid values raise at submit time (`validate()`).
-Assign a scalar to size every ring the same:
-
-```python
-cfg = CallConfig()
-cfg.runtime_env.ring_task_window = 128   # power of 2, >= 4
-cfg.runtime_env.ring_heap = 262144       # bytes/ring, >= 1024
-cfg.runtime_env.ring_dep_pool = 256      # 4 .. INT32_MAX
-orchestrator.submit_next_level(handle, args, cfg)
-```
-
-Assign a four-entry list to tune the scope-depth rings independently. The list
-must contain exactly four entries; use `0` for an entry that should fall through
-to the next precedence tier. All `CallConfig` values are integer byte/count
-values, and each field always reads back as a four-entry list.
-
-```python
-cfg = CallConfig()
-cfg.runtime_env.ring_task_window = [8192, 16384, 131072, 524288]
-cfg.runtime_env.ring_heap = [
-    128 * 1024 * 1024,
-    256 * 1024 * 1024,
-    384 * 1024 * 1024,
-    512 * 1024 * 1024,
-]
-cfg.runtime_env.ring_dep_pool = [4096, 8192, 16384, 32768]
-orchestrator.submit_next_level(handle, args, cfg)
-```
-
-Scene tests set the same keys under a nested `runtime_env` block in the
-per-case `config` dict — each value is a scalar or a four-entry list:
-
-```python
-"config": {
-    "runtime_env": {
-        "ring_task_window": [8192, 16384, 131072, 524288],
-        "ring_heap": [134217728, 268435456, 402653184, 536870912],
-        "ring_dep_pool": 256,  # scalar broadcasts to every ring
-    }
-}
-```
-
-Process-wide env fallback accepts either one scalar value or exactly four
-comma-separated per-ring values. Invalid env values are logged and ignored, then
-fall through to defaults. `PTO2_RING_HEAP` values are integer bytes:
+Uniform (applies to all rings):
 
 ```bash
-# Uniform, old behavior:
 PTO2_RING_TASK_WINDOW=1024
 PTO2_RING_HEAP=1048576
 PTO2_RING_DEP_POOL=1024
-
-# Per-ring, indexed by ring_id 0..3:
-PTO2_RING_TASK_WINDOW=8192,16384,131072,524288
-PTO2_RING_HEAP=134217728,268435456,402653184,536870912
-PTO2_RING_DEP_POOL=4096,8192,16384,32768
 ```
 
-Use `--enable-scope-stats` to confirm the effective values for a real run. The
-first line of `scope_stats/scope_stats.jsonl` includes `task_window_max`,
-`heap_max`, and `dep_pool_max`, indexed by `ring`.
+In `kernel_config.py`:
+
+```python
+RUNTIME_ENV = {
+    "PTO2_RING_TASK_WINDOW": "128",
+    "PTO2_RING_HEAP": "262144",
+    "PTO2_RING_DEP_POOL": "256",
+}
+```
 
 ### 7.3 Sizing Guidelines
 

src/a2a3/runtime/tensormap_and_ringbuffer/docs/RUNTIME_LOGIC.md

@@ -446,7 +446,7 @@ The orchestrator completes fanout wiring before publishing a task to the ready q
    - Checks `task_state >= COMPLETED` (early-finished optimization)
    - If not completed: prepends consumer to producer's `fanout_head` via `dep_pool.prepend`
    - **Releases** `fanout_lock`
-3. Atomically releases the +1 redundance + completed-fanin count via `fanin_refcount.fetch_add`
+3. Atomically releases the +1 redundance + early_finished count via `fanin_refcount.fetch_add`
 4. If all deps satisfied: pushes task to the routed ready queue
 
 Zero-fanin tasks and tasks whose claimed producers are already completed skip dep_pool entry allocation and publish directly to the routed ready queue.
@@ -586,7 +586,7 @@ This is protected by a per-ring try-lock (`advance_lock`) in `RingSchedState`, e
 
 ### 8.5 SchedulerContext
 
-All scheduler-side state and methods live in `SchedulerContext` (`runtime/scheduler/scheduler_context.h`). It is held as a `sched_ctx_` member of `AicpuExecutor`; `AicpuExecutor` is a thin wrapper that owns the lifecycle atomics and the orchestration SO handle, and delegates everything else to `SchedulerContext`.
+All scheduler-side state and methods live in `SchedulerContext` (`runtime/scheduler_context.h`). It is held as a `sched_ctx_` member of `AicpuExecutor`; `AicpuExecutor` is a thin wrapper that owns the lifecycle atomics and the orchestration SO handle, and delegates everything else to `SchedulerContext`.
 
 Public surface (called from `AicpuExecutor::init/run/deinit`):
 
@@ -600,11 +600,7 @@ Public surface (called from `AicpuExecutor::init/run/deinit`):
 | `deinit()` | once per run | Reset every scheduler-owned field to its post-construction default |
 | Read-only accessors | various | `aic_count()` / `aiv_count()` / `is_completed()` / `completed_tasks_count()` |
 
-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 (`pre_handshake_init` / `handshake_partition` / `post_handshake_init` / `deinit`), core management (`assign_cores_to_threads` / `emergency_shutdown`), and `on_orchestration_done`
+Private internals all live inline in `scheduler_context.h`, covering completion polling, drain protocol, task dispatch loop and helpers, exit checks, stall diagnostics, profiling, lifecycle (`init/deinit`), core management (`handshake_all_cores` / `assign_cores_to_threads` / `reassign_cores_for_all_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`.
 
@@ -683,13 +679,11 @@ reserved core slot and a launch-visible payload before a consumer can pre-occupy
 `next_block_idx` records reservation progress; `published_block_count` independently establishes
 publication and early-candidate readiness.
 
-The producer's slot-local dispatch-propagated bit in `lifecycle_flags` and fanout snapshot are
-serialized under `fanout_lock` with consumer wiring. Wiring already locks and reads the
-producer's 64-byte `PTO2TaskSlotState`, so testing this bit does not read a producer payload
-cache line. An edge already in the snapshot is counted by scheduler propagation; wiring seeds
-an edge added after the claim and enqueues the consumer if that seed completes
-`dispatch_fanin`. This gives each eligible producer-consumer edge exactly one early-candidate
-contribution regardless of which side acquires the lock first.
+The producer's `dispatch_propagated` claim and fanout snapshot are serialized under
+`fanout_lock` with consumer wiring. An edge already in the snapshot is counted by scheduler
+propagation; wiring seeds an edge added after the claim and enqueues the consumer if that seed
+completes `dispatch_fanin`. This gives each eligible producer-consumer edge exactly one
+early-candidate contribution regardless of which side acquires the lock first.
 
 #### MIX per-core placement
 

src/a2a3/runtime/tensormap_and_ringbuffer/docs/SCALAR_DATA_ACCESS.md

@@ -32,7 +32,7 @@ addr null-check → TensorMap lookup → spin-wait producer COMPLETED → comput
 
 - **addr null-check**: `buffer.addr == 0` means unallocated — log error, return 0
 - **TensorMap lookup**: find producer task by `buffer.addr`
-- **spin-wait**: wait until producer `task_state >= PTO2_TASK_COMPLETED`
+- **spin-wait**: wait until producer's `completion_flags[local_id & mask] == 1`
 - **No producer** (lookup callback never fires): skip waiting, read immediately
 
 ### 3.2 set_tensor_data Flow

src/a2a3/runtime/tensormap_and_ringbuffer/docs/device_log_profiling.md

@@ -52,8 +52,8 @@ Thread 3: PTO2 total submitted tasks = 16704
 
 ### Field Reference
 
-| Field | Source (`pto_orchestrator.cpp`) | Description |
-| ----- | ------------------------------- | ----------- |
+| Field | Source (`pto_orchestrator.h`) | Description |
+| ----- | ----------------------------- | ----------- |
 | **cost** | Wall-clock around `orch_func()` call | Total time including orchestration logic + scope overhead |
 | **total** | Sum of all sub-steps below | Accumulated time inside `submit_task` across all tasks |
 | **sync_tensormap** | `g_orch_sync_cycle` | TensorMap validity sync and optional cleanup before each submission |

src/a2a3/runtime/tensormap_and_ringbuffer/docs/profiling_levels.md

@@ -48,7 +48,7 @@ Each sub-level macro requires `SIMPLER_DFX=1`:
 
 - Debug/diagnostic logs (always present)
 - Progress tracking (`PTO2 progress: completed=...`)
-- Stall detection and dump (triggered after the `SCHEDULER_TIMEOUT_MS` wall-clock no-progress budget)
+- Stall detection and dump (triggered only after `MAX_IDLE_ITERATIONS` idle loops)
 - Deadlock/livelock detection (`diagnose_stuck_state`, called on stall)
 
 **What's NOT compiled:**
@@ -255,7 +255,7 @@ Identity fields the AICPU side used to write at level 1 (`func_id`,
   collector (`L2SwimlaneCollector::set_core_types`).
 
 AICore buffer rotation no longer piggy-backs on `complete_task`. AICPU
-counts dispatches per core in the dispatch path (scheduler_dispatch in
+counts dispatches per core in the dispatch path (scheduler_context in
 tensormap_and_ringbuffer; aicpu_executor in host_build_graph) and rotates
 the AICore buffer when the count is about to cross a
 `PLATFORM_AICORE_BUFFER_SIZE` boundary — strictly before
@@ -428,7 +428,7 @@ definitions to runtime headers.
 ### Code Locations
 
 - Macro defaults and validation: `src/common/task_interface/profiling_config.h`
-- Scheduler profiling: `src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_dispatch.cpp` and `scheduler_cold_path.cpp`
+- Scheduler profiling: `src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler_context.h`
 - Orchestrator profiling: `src/a2a3/runtime/tensormap_and_ringbuffer/aicpu/aicpu_executor.cpp`
 - TensorMap profiling: `src/a2a3/runtime/tensormap_and_ringbuffer/runtime/pto_tensormap.h`
 

src/a2a3/runtime/tensormap_and_ringbuffer/host/dep_gen_replay.cpp

@@ -558,7 +558,7 @@ dep_gen_replay_emit_deps_json(const DepGenRecord *records, size_t num_records, c
         // `explicit_dep_count` / `over->dep_count` originate from device
         // shared memory and are bounded by the writer to the array sizes, but
         // we clamp on read too so a corrupted record never drives an OOB read
-        // off the end of rec.explicit_deps[64] / over->deps[582].
+        // off the end of rec.explicit_deps[64] / over->deps[326].
         const uint64_t *deps_data;
         int32_t dc;
         if (rec.flags & DEP_GEN_FLAG_HAS_OVERFLOW) {

src/a2a3/runtime/tensormap_and_ringbuffer/orchestration/common.cpp

@@ -10,6 +10,13 @@
  */
 #include "common.h"
 
+// LOG_ERROR can't be pulled from common/unified_log.h here because that header
+// would re-#define LOG_INFO_V0..V9 already provided by pto_orchestration_api.h
+// (orchestration routes them through the runtime ops table). For the limited
+// use inside this file, write directly to stderr.
+#include <cstdio>
+#define LOG_ERROR(fmt, ...) std::fprintf(stderr, "[ERROR] " fmt "\n", ##__VA_ARGS__)
+
 #ifdef __linux__
 #include <cxxabi.h>
 #include <dlfcn.h>