perf(cpu): route Transpose and the elementwise fallback through the task runtime

crates/onnx-runtime-ep-cpu/src/kernels/simd_activations.rs

@@ -570,24 +570,49 @@ where
         body(input, output);
         return;
     }
-    // Nesting a `par_chunks` inside an outer parallel region only adds
-    // scheduling overhead: the outer region is already keeping the pool busy.
-    if rayon::current_thread_index().is_some() {
+    // Nesting a split inside an outer parallel region only adds scheduling
+    // overhead: the outer region is already keeping the pool busy. Rayon is
+    // still checked because plenty of this crate fans out through it directly.
+    if rayon::current_thread_index().is_some() || crate::task_runtime::in_task() {
         body(input, output);
         return;
     }
-    let Some(chunk) = par_chunk_len(len, rayon::current_num_threads()) else {
+    let Some(chunk) = par_chunk_len(len, crate::task_runtime::width()) else {
         body(input, output);
         return;
     };
 
     note_parallel_dispatch();
+    run_on_task_runtime(input, output, chunk, &body);
+}
 
-    use rayon::prelude::*;
-    output
-        .par_chunks_mut(chunk)
-        .zip(input.par_chunks(chunk))
-        .for_each(|(o, i)| body(i, o));
+/// Splits `input`/`output` into `chunk`-sized pieces across the native task
+/// runtime and runs `body` on each.
+///
+/// Bit-identical to calling `body` on the whole slice, for the same reason
+/// [`run_on_host`] is: the kernels are elementwise, every chunk is a multiple of
+/// eight lanes, and none is shorter than [`PAR_MIN_CHUNK`], so no chunk takes a
+/// different code path from its neighbours.
+///
+/// The runtime's pool replaces rayon here because rayon parks its workers
+/// between regions and an elementwise op is *all* region boundary: an isolated
+/// region measured 67 us back-to-back and 226 us when the previous one ended
+/// 20 us earlier, which is the spacing between two activations in a decode step.
+///
+/// `#[inline(never)]` for the same reason as [`try_host`]: `run_chunked`'s
+/// callers are the hottest elementwise kernels in the crate, and growing it has
+/// repartitioned codegen units before -- 34% on `Relu` at 1 Mi with no path
+/// change at all.
+#[inline(never)]
+fn run_on_task_runtime<F>(input: &[f32], output: &mut [f32], chunk: usize, body: &F)
+where
+    F: Fn(&[f32], &mut [f32]) + Sync,
+{
+    debug_assert_eq!(input.len(), output.len());
+    crate::task_runtime::chunk_runs_mut(output, chunk, 1, |first, out| {
+        let start = first * chunk;
+        body(&input[start..start + out.len()], out);
+    });
 }
 
 #[cfg(test)]
@@ -675,20 +700,16 @@ where
         body(input, output);
         return;
     }
-    if rayon::current_thread_index().is_some() {
+    if rayon::current_thread_index().is_some() || crate::task_runtime::in_task() {
         body(input, output);
         return;
     }
-    let Some(chunk) = par_chunk_len_rows(len, width, rayon::current_num_threads()) else {
+    let Some(chunk) = par_chunk_len_rows(len, width, crate::task_runtime::width()) else {
         body(input, output);
         return;
     };
 
-    use rayon::prelude::*;
-    output
-        .par_chunks_mut(chunk)
-        .zip(input.par_chunks(chunk))
-        .for_each(|(o, i)| body(i, o));
+    run_on_task_runtime(input, output, chunk, &body);
 }
 
 // ---------------------------------------------------------------------------

crates/onnx-runtime-ep-cpu/src/kernels/transpose.rs

@@ -76,6 +76,17 @@ const TILE: usize = 64;
 /// more cores start helping.
 const MIN_PARALLEL_BYTES: usize = 256 * 1024;
 
+/// Minimum output bytes in one task's chunk.
+///
+/// A quarter of [`MIN_PARALLEL_BYTES`]: the whole-operation bar exists because
+/// below it the copy finishes faster than a fan-out can be scheduled, and a
+/// single chunk faces the same arithmetic against the task runtime's much
+/// cheaper per-task cost (a `fetch_sub` and an indirect call, not a fork/join).
+/// 64 KiB is still comfortably past the point where a copy is bandwidth-bound,
+/// so a chunk this size scales, while being small enough that a 1 MiB transpose
+/// has 16 chunks to balance across rather than one per thread.
+const MIN_PARALLEL_TASK_BYTES: usize = 64 * 1024;
+
 /// Dtype-generic Transpose kernel carrying the resolved `perm`.
 pub struct TransposeKernel {
     /// Axis permutation; `None` means reverse all axes.
@@ -389,11 +400,10 @@ fn block_move(src: &[u8], dst: &mut [u8], shape: &[usize], perm: &[usize], block
     };
 
     match parallel_blocks_per_task(blocks, block) {
-        Some(per_task) => {
-            use rayon::prelude::*;
-            dst.par_chunks_mut(per_task * block)
-                .enumerate()
-                .for_each(|(task, chunk)| body(chunk, task * per_task));
+        Some(per_chunk) => {
+            crate::task_runtime::chunk_runs_mut(dst, per_chunk * block, 1, |first, run| {
+                body(run, first * per_chunk);
+            });
         }
         None => body(dst, 0),
     }
@@ -407,13 +417,10 @@ fn blocked_2d(src: &[u8], dst: &mut [u8], rows: usize, cols: usize, esize: usize
     // worker a `&mut [u8]` without any unsafe pointer splitting.
     let dst_row = rows * esize;
     match parallel_blocks_per_task(cols, dst_row) {
-        Some(per_task) => {
-            use rayon::prelude::*;
-            dst.par_chunks_mut(per_task * dst_row)
-                .enumerate()
-                .for_each(|(task, chunk)| {
-                    blocked_2d_band(src, chunk, rows, cols, esize, task * per_task);
-                });
+        Some(per_chunk) => {
+            crate::task_runtime::chunk_runs_mut(dst, per_chunk * dst_row, 1, |first, run| {
+                blocked_2d_band(src, run, rows, cols, esize, first * per_chunk);
+            });
         }
         None => blocked_2d_band(src, dst, rows, cols, esize, 0),
     }
@@ -469,16 +476,13 @@ fn batched_blocked_2d(
     plane: usize,
 ) {
     match parallel_blocks_per_task(batch, plane) {
-        Some(per_task) => {
-            use rayon::prelude::*;
-            dst.par_chunks_mut(per_task * plane)
-                .enumerate()
-                .for_each(|(task, chunk)| {
-                    for (n, dst_plane) in chunk.chunks_mut(plane).enumerate() {
-                        let base = (task * per_task + n) * plane;
-                        blocked_2d_band(&src[base..base + plane], dst_plane, rows, cols, esize, 0);
-                    }
-                });
+        Some(per_chunk) => {
+            crate::task_runtime::chunk_runs_mut(dst, per_chunk * plane, 1, |first, run| {
+                for (n, dst_plane) in run.chunks_mut(plane).enumerate() {
+                    let base = (first * per_chunk + n) * plane;
+                    blocked_2d_band(&src[base..base + plane], dst_plane, rows, cols, esize, 0);
+                }
+            });
         }
         None => {
             // Too few planes to split, so let each plane decide for itself
@@ -524,16 +528,31 @@ fn unflatten(mut flat: usize, shape: &[usize]) -> Vec<usize> {
     idx
 }
 
-/// `Some(units_per_task)` when splitting `units` of `unit_bytes` each pays off.
+/// `Some(units_per_chunk)` when splitting `units` of `unit_bytes` each pays off.
+///
+/// The chunk is a *balance* unit, not a per-thread share. Sizing it as
+/// `units / workers` gives every worker one static band and makes the whole
+/// transpose wait for whichever band was unlucky -- a descheduled thread, an SMT
+/// sibling, a page fault on first touch. The task runtime claims chunks
+/// dynamically and caps the task count itself, so handing it several chunks per
+/// thread costs nothing on an idle machine and absorbs a straggler on a busy
+/// one.
+///
+/// [`MIN_PARALLEL_TASK_BYTES`] keeps a chunk from falling below the size where
+/// a copy is bandwidth-bound rather than latency-bound; below that the
+/// per-chunk bookkeeping starts to show.
 fn parallel_blocks_per_task(units: usize, unit_bytes: usize) -> Option<usize> {
     if units.saturating_mul(unit_bytes) < MIN_PARALLEL_BYTES {
         return None;
     }
-    let workers = rayon::current_num_threads().max(1);
-    if workers < 2 || units < 2 {
+    if crate::task_runtime::width() < 2 || units < 2 {
         return None;
     }
-    Some(units.div_ceil(workers).max(1))
+    Some(
+        MIN_PARALLEL_TASK_BYTES
+            .div_ceil(unit_bytes.max(1))
+            .clamp(1, units),
+    )
 }
 
 #[cfg(test)]

crates/onnx-runtime-ep-cpu/src/task_runtime/pool.rs

@@ -736,21 +736,47 @@ mod tests {
         let pool = Arc::new(TaskPool::new(4));
         let hold = Arc::new(AtomicBool::new(true));
         let declined = Arc::new(AtomicUsize::new(0));
+        // Count how many holder dispatches have actually entered their body, so
+        // the probe below runs only once every slot is provably occupied.
+        // Sleeping a fixed 50 ms instead assumes the OS scheduled all
+        // `SLOT_COUNT` holder threads inside that window -- which fails on a
+        // contended box (e.g. the rest of the test suite running in parallel),
+        // leaving slots free, the probe succeeding, and the assertion flaking.
+        // Constructing the precondition removes the race rather than papering
+        // it over.
+        let held = Arc::new(AtomicUsize::new(0));
         thread::scope(|scope| {
             for _ in 0..SLOT_COUNT {
                 let pool = Arc::clone(&pool);
                 let hold = Arc::clone(&hold);
+                let held = Arc::clone(&held);
                 scope.spawn(move || {
+                    // One increment per holder even though the fan-out has two
+                    // tasks and a pool worker may run the second.
+                    let counted = AtomicBool::new(false);
                     let _ = pool.dispatch(2, &|_| {
+                        if !counted.swap(true, Ordering::Relaxed) {
+                            held.fetch_add(1, Ordering::Release);
+                        }
                         while hold.load(Ordering::Relaxed) {
                             std::hint::spin_loop();
                         }
                     });
                 });
             }
-            // Give the holders time to take every slot, then prove the next
-            // dispatcher is declined rather than blocked.
-            thread::sleep(Duration::from_millis(50));
+            // Wait until every holder is provably inside its body -- i.e. every
+            // slot is claimed -- then prove the next dispatcher is declined
+            // rather than blocked. A generous deadline keeps a starved runner
+            // from hanging the suite forever.
+            let deadline = Instant::now() + Duration::from_secs(10);
+            while held.load(Ordering::Acquire) < SLOT_COUNT && Instant::now() < deadline {
+                std::hint::spin_loop();
+            }
+            assert_eq!(
+                held.load(Ordering::Acquire),
+                SLOT_COUNT,
+                "not every holder took a slot within the deadline"
+            );
             for _ in 0..4 {
                 if !pool.dispatch(2, &|_| {}) {
                     declined.fetch_add(1, Ordering::Relaxed);

docs/benchmarks/2026-08-15-cpu-ep-vs-ort-attention-moe.md

@@ -2973,7 +2973,7 @@ path reads an environment variable for test purposes, and the tests do not race
 on a global by construction: `isolated_pool` builds a private pool rather than
 reconfiguring the shared one.
 
-### 31.7 RoPE and Softmax, moved onto it
+### 33.7 RoPE and Softmax, moved onto it
 
 Real ORT A/B, one driver invocation per pair, 7 trials × 7 runs. Ratio is
 native ÷ ORT, lower is better, `base>new`:
@@ -3009,7 +3009,7 @@ The decode-shaped cells are deliberately unmoved: `rope_*_s1` and
 arms, flat to within noise at every width. A one-token RoPE should not be
 fanning out and still does not.
 
-### 31.7.1 The scar: an unrelated edit de-vectorised the activations
+### 33.7.1 The scar: an unrelated edit de-vectorised the activations
 
 Rewiring RoPE and Softmax made **`Tanh` and `Sigmoid` 2× slower at one thread**,
 in a file the change did not touch, on inputs that never reach a parallel path.
@@ -3037,7 +3037,7 @@ Two things about this are worth keeping:
 neither helper carries the latter; every caller is already inside an `avx2,fma`
 function, so the features arrive with the inline.
 
-### 31.7.2 Re-verified against §29's Softmax, not the one it replaced
+### 33.7.2 Re-verified against §29's Softmax, not the one it replaced
 
 The Softmax numbers above were taken before #1219 landed, so their baseline was
 a `softmax_rows` that still copied `src` into `dst` before reducing. Merging
@@ -3061,3 +3061,61 @@ because what it removes is a per-region park, not a per-element cost. The
 decode-shaped `sm_decode_h32_kv*` cells stay serial and stay flat (18–182 µs,
 within ±1 µs of base at every width), and `sm_whisper_cross` is neutral at every
 width — both are the intended outcome, not a miss.
+
+### 33.8 Transpose and the activations, and what the controls say
+
+Activations, our own native ms, `base>new`, 5 trials × 7 runs:
+
+| cell | t=1 | t=2 | t=4 | t=8 | t=16 | t=32 |
+|---|---|---|---|---|---|---|
+| act_relu_hidden_t512 | 0.838>0.822 | 0.546>0.488 | 0.295>0.251 | 0.276>0.183 | 0.381>0.130 | 1.201>0.303 |
+| act_tanh_hidden_t512 | 0.804>0.803 | 0.806>0.794 | 0.457>0.417 | 0.419>0.317 | 0.485>0.204 | 1.441>0.386 |
+| act_sigmoid_hidden_t512 | 0.826>0.837 | 0.832>0.808 | 0.640>0.559 | 0.462>0.296 | 0.513>0.214 | 1.176>0.437 |
+| act_tanh_hidden_t4096 | 6.398>6.542 | 6.331>6.436 | 4.300>3.412 | 3.609>3.494 | 3.691>3.490 | 5.991>2.933 |
+| act_sigmoid_hidden_t4096 | 6.532>6.680 | 6.488>6.605 | 3.547>3.465 | 3.560>3.476 | 3.672>3.539 | 6.070>2.356 |
+| act_fastgelu_mlp_t512 | 5.466>5.448 | 5.169>5.294 | 2.815>2.765 | 2.156>1.864 | 1.949>1.553 | 4.461>1.463 |
+
+27 improved cells against 5 regressed at a 5% deadband, and the five are 1–6 µs
+absolute or bandwidth-bound cells whose distributions overlap outright. Up to
+3.6× at t=16 and 4× at t=32.
+
+Transpose needed 15 trials rather than 7 before it said anything stable — at 7
+it produced a `tr_bert_b8_s128` "regression" that was entirely dispersion:
+
+| cell | t=1 | t=2 | t=4 | t=8 | t=16 | t=32 |
+|---|---|---|---|---|---|---|
+| tr_bert_b8_s128 | 0.117>0.119 | 0.152>0.153 | 0.111>0.112 | 0.134>0.140 | 0.261>0.248 | 2.658>2.670 |
+| tr_llama3_s512 | 0.718>0.705 | 0.390>0.389 | 0.227>0.228 | 0.253>0.239 | 0.430>0.352 | 3.412>3.101 |
+| tr_whisper_s1500 | 0.566>0.593 | 0.448>0.442 | 0.353>0.358 | 0.307>0.296 | 0.397>0.380 | 3.292>3.017 |
+
+4 improved, 0 regressed. Transpose is a straight copy, so it is bandwidth-bound
+long before it is scheduler-bound; the gain is small and lives at the wide end.
+
+**The controls.** Two families this work does not touch were run at the same six
+widths purely to establish the noise floor: 24 GQA attention cells (32 improved
+/ 27 regressed) and 20 GEMM cells (19 / 12). Both are symmetric, which is what
+licenses reading the 3–15× numbers as real rather than as luck. Anyone repeating
+this on this host should budget for ±25% on a single 7-trial cell and treat any
+claim smaller than that as unmeasured.
+
+### 33.9 The biggest remaining loss is GEMM, and it is the same disease
+
+The GEMM control was meant to be a null result and instead found the largest
+single loss in the CPU EP. Decode-shaped `MatMulNBits`, native ms:
+
+| cell | t=8 | t=16 | t=32 | ratio vs ORT at t=32 |
+|---|---|---|---|---|
+| gemm_nbits_qwen3_0p6b_qkv_t8 | 0.67 | 0.64 | **2.98** | 25.99 |
+| gemm_nbits_qwen3_0p6b_mlp_t8 | 1.17 | 1.03 | **3.81** | 18.73 |
+| gemm_nbits_llama3_8b_qkv_t8 | 4.15 | 3.05 | **6.07** | 8.64 |
+
+Going from 16 to 32 threads makes these **4.6× slower**, on the operator that
+carries most of a decode step's arithmetic. That is the §26 signature exactly,
+and `matmul_nbits.rs` is still the largest raw-Rayon fan-out in the crate.
+
+Prefill-shaped GEMM (`t128`, `t512`) does not show it: those regions are long
+enough to amortise a park. It is specifically the small, frequent, decode-shaped
+region that the global pool cannot serve, which is the whole thesis of §26.
+
+Nothing in this phase touches it. It is the next target, and the runtime it
+needs already exists.