Repository navigation
perf(cpu): run the SiLU, Relu and Clip MLAS routes through run_chunked - #1130
Merged
Merged
Conversation
The MLAS routes for SiLU, Relu and Clip called their kernel directly instead of going through `run_chunked`, so they never split across the rayon pool no matter how many threads were configured. This is the same pathology #1127 fixed for the `dispatch_mlas!` ops, reported as a MAJOR finding by the independent review on that PR. `mlas-sys` documents `compute_silu`, `compute_relu` and `compute_clip` as single threaded, with sharding left to the caller. Nobody sharded. SiLU additionally needed its correction scan blocked so the parallel chunks stay in L2, and the scan is now a branch-free OR-reduction over the input alone, so the common all-in-band case skips the write loop. Measured at 16 threads, session level through the plugin, zero nodes left unassigned: Clip 4 Mi 1284.90 -> 516.92 us 2.49x Relu 4 Mi 1095.22 -> 494.47 us 2.22x Swish 4 Mi 5629.53 -> 549.99 us 10.24x Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Independent review findings on #1130. MAJOR: `.work_backup.rs`, a 988 line mid-development snapshot of activations.rs, had been committed at the repo root by accident. Removed. MINOR: the `silu_bench` doc comment claimed Silu/Swish have no single-op equivalent in ORT, which this PR itself disproves — ORT 1.28 implements `Swish` in the default domain at opset 24, and the PR reports a session-level A/B built on it. Comment corrected. NIT: `silu_reaches_run_chunked_parallel_branch` hard-failed on a single-core runner instead of skipping. It now returns early with a note. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
The re-review noted that `assert_same` and `assert_parallelises` in simd_activations.rs carry the same unguarded hard assert that was just fixed for `silu_reaches_run_chunked_parallel_branch`. They predate this PR, but they are its closest analogs, so they get the same early-return skip rather than a spurious failure on a single-core runner. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #1130 +/- ##
==========================================
+ Coverage 79.94% 80.60% +0.66%
==========================================
Files 368 366 -2
Lines 160607 157585 -3022
Branches 160607 157585 -3022
==========================================
- Hits 128396 127028 -1368
+ Misses 27494 25853 -1641
+ Partials 4717 4704 -13
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
… deckard/silu-parallel
… deckard/silu-parallel
🔴 Benchmark Regression DetectedComparison of criterion micro-benchmarks: PR head vs merge-base, measured on the same runner in the same job (base first → PR second).
Visual flags: Host infoWhat this cannot catch
|
justinchuby
added a commit
that referenced
this pull request
Aug 17, 2026
## What happened #1130 (mine) wrapped the `Clip` MLAS call in `run_chunked` so it would use the thread pool. `run_chunked` is generic, so doing that from `selection.rs` created a new instantiation in a module that had never had one. The runtime path of every other unary op was untouched — not one instruction — and all 1332 tests still passed. But the crate's codegen units repartitioned and the AVX2 unary kernels in `simd_activations.rs` stopped being vectorised. This shipped. It was found while collecting numbers for the final report: `Sqrt` had gone from beating ORT by 1.9× to losing at 0.7×, and the ORT-relative position of five ops had collapsed in a way no code change explained. ## Evidence Bisected by rebuilding `main` with one file at a time reverted to `1f1ce4b74` (the commit before #1130). n = 65536, 1 thread, `taskset -c 8-23`, 4 interleaved rounds, µs p50. **Only `selection.rs` matters** — reverting `relu.rs` or `conv.rs` alone restored nothing: | op | main (#1130) | revert relu.rs | revert **selection.rs** | revert conv.rs | pre-#1130 | |---|---:|---:|---:|---:|---:| | Sqrt | 48.5 | 48.4 | **21.3** | 44.6 | 21.5 | | Tanh | 55.2 | 57.0 | **31.6** | 60.0 | 31.8 | | Sigmoid | 56.9 | 55.6 | **31.0** | 57.1 | 31.1 | | QuickGelu | 64.5 | 64.5 | **42.5** | 64.2 | 42.5 | | FastGelu | 79.6 | 79.3 | **55.5** | 79.3 | 55.8 | | Erf | 62.2 | 62.4 | 62.2 | 62.2 | 62.5 | | Relu | 21.3 | 26.5 | 22.6 | 23.0 | 20.4 | Reproduced independently by building the same commit in a second worktree, so it is not a build-directory artefact. Adding `#[inline]` to `run_chunked` did **not** help, which rules out a plain inlining decision and points at codegen-unit partitioning. ## The diagnosis, confirmed mechanically If the cause really is codegen-unit partitioning, then forcing the crate into a single codegen unit should erase the regression with no source change at all. It does. `CARGO_PROFILE_RELEASE_CODEGEN_UNITS=1` on the *unfixed* commit `34095af0f`, same machine, 4 interleaved rounds, µs p50: | op | n | main, default CGUs | main, `codegen-units=1` | this PR, default CGUs | |---|---:|---:|---:|---:| | Sqrt | 64 Ki | 48.3 | **21.6** | 21.2 | | Sqrt | 1 Mi | 667.4 | **261.4** | 274.4 | | Sigmoid | 1 Mi | 813.3 | **504.8** | 441.6 | | QuickGelu | 1 Mi | 929.4 | **575.8** | 573.0 | | FastGelu | 1 Mi | 1171.2 | **788.2** | 808.5 | So the diagnosis is not inferred from a bisect alone — the proposed mechanism, applied directly, reproduces the cure. `codegen-units = 1` (or LTO) in the release profile would remove this whole fragility class permanently, and is the more durable answer. It is deliberately **not** in this PR: it is a workspace-wide build-policy change that affects every crate and every contributor's build time (the plugin alone went 14 s to 63 s here), it would need its own measurement across the whole EP rather than the activation kernels, and it does not belong in a regression fix. Filed as the follow-up this PR's limitation section points at. The source fix costs nothing and is independent of it. ## The fix `run_chunked` is private to `simd_activations.rs` again — the compiler now enforces the rule, not a convention. Callers elsewhere go through one of two entry points that are instantiated *in that module*: - `run_chunked_fn(input, output, body: fn(&[f32], &mut [f32]))` — deliberately a `fn` pointer, not `impl Fn`, so every caller shares one instantiation. Used by `Relu` and `SiLU`. - `clip_chunked(input, output, min, max)` — `Clip` needs captured bounds. It takes the serial decision itself so the short case is a direct call rather than one through a closure the optimiser can no longer see into; without that, `Clip` itself paid 12%. ## Result n = 65536 and 1 Mi, 1 thread, 5 interleaved rounds, µs p50: | op | n | main (#1130) | this PR | pre-#1130 | |---|---:|---:|---:|---:| | Sqrt | 64 Ki | 48.5 | **21.4** | 21.7 | | Sqrt | 1 Mi | 667.4 | **259.2** | 275.1 | | Tanh | 1 Mi | 776.1 | **437.8** | 436.4 | | Sigmoid | 1 Mi | 807.6 | **440.1** | 502.2 | | QuickGelu | 1 Mi | 923.1 | **574.7** | 607.0 | | FastGelu | 1 Mi | 1166.9 | **785.3** | 788.1 | | Clip | 1 Mi | 268.7 | **267.5** | 267.4 | | Swish | 1 Mi | 728.8 | **720.8** | 1385.6 | Everything is back to its pre-#1130 level **and** #1130's own win is kept: `Swish` is still 1.92× faster than before #1130, and `Clip`/`Relu` still reach the pool at ≥ `PAR_MIN_LEN`. ## Regression guard `chunking_instantiation_is_local::no_module_outside_this_one_instantiates_run_chunked` walks the crate source and fails if any module other than `simd_activations.rs` instantiates `run_chunked`, naming the offending file and line. Verified to falsify: re-widening the visibility and pointing `relu.rs:144` back at `run_chunked` makes it fail with `Offending call sites: ["…/kernels/relu.rs:144"]`. This class of bug produces no wrong answers, no test failures and no diff in the file that slows down, so a mechanical guard is the only thing that catches it. ## Tests - `cargo test -p onnx-runtime-ep-cpu --features mlas --lib` → **1332 passed, 0 failed** - `cargo test -p onnx-runtime-ep-cpu --lib` → **1305 passed, 0 failed** - `cargo fmt --all` clean. - No numerical change: `clip_chunked`'s serial branch calls exactly the function the closure called, and `run_chunked_fn` forwards unchanged. ## Limitations - The mechanism is codegen-unit partitioning, which the compiler makes no promises about. The guard encodes the rule that was measured to work on this toolchain; it cannot prove the next refactor is safe. That is why the guard names the symptom and the measured cost in its failure message. - Measured on one machine (AMD EPYC 9V74, AVX2/FMA/F16C, `taskset`-pinned, 5 rounds). The *direction* is unambiguous — up to 2.3× on `Sqrt` and roughly 1.5–1.8× on the other four — but the exact figures are not portable. --------- Co-authored-by: Deckard <deckard@users.noreply.github.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Root cause
run_chunkedis the seam that #1105 taught to split activation work across the rayonpool. #1127 fixed
dispatch_mlas!, which called its kernel directly and returned,bypassing that seam entirely — so every MLAS-routed op ran single threaded no matter how
many threads were configured.
The independent review on #1127 reported as a MAJOR finding that three more callers
have the identical pathology, in other files, so they were out of scope there:
silu_f32_slicekernels/activations.rsrelu_contiguous_f32_mlaskernels/relu.rs:144Clip(selection)kernels/selection.rs:182Clip(conv epilogue)kernels/conv.rs:732mlas-sysdocumentscompute_silu,compute_reluandcompute_clipas "Singlethreaded; callers shard across threads themselves." Nobody sharded. This PR wraps all
four call sites in
run_chunked.SiLU needed one extra step. Its MLAS route is followed by a correction scan over the
whole tensor (MLAS's
compute_siluis inaccurate outside±SILU_MLAS_SAFE_BOUND).Run whole-tensor, that scan streams the buffer a second time from DRAM. It is now
blocked at
SILU_CORRECTION_BLOCK = 8192so each block stays in L2, and the scan is abranch-free OR-reduction over the input only — the predicate
!x.is_finite() || x.abs() > SILU_MLAS_SAFE_BOUNDdepends solely on the input, so thecommon all-in-band case skips the write loop entirely.
Benchmarks
Session level through the plugin
.so, base =origin/main@b5309f799, 16 threads,3 interleaved rounds, randomised order,
# NOT-ASSIGNED: 0on every run (no node wasleft to ORT's CPU EP). µs, p50.
Swish(default domain, opset 24) is the ORT-visible spelling of SiLU and is supportedby ORT 1.28, so SiLU does have a real single-node session-level A/B after all — the
earlier note that it did not was wrong, and it is the op that gains the most here.
Kernel-level SiLU,
serial_scopevs parallel in-process, 32 threads, so the MLAS routeis compared against itself with only the split changed:
Correctness
blocked_correction_matches_the_whole_tensor_loop_bit_for_bit— the blocked,OR-reduced scan is compared bit for bit against the original whole-tensor loop over
in-band values, out-of-band values,
±Inf, NaN,±0, denormals and values sittingexactly on
SILU_MLAS_SAFE_BOUND, at lengths that straddle the block boundary.silu_reaches_run_chunked_parallel_branch— asserts the mechanism, not theoutput, using the
PARALLEL_DISPATCHEScounter added in perf(cpu): run the MLAS activation routes through run_chunked #1127. Verified to falsify:reverting the
run_chunkedwrapper makes it fail.silu_is_thread_count_invariant— identical results across pool sizes.changes who runs the arithmetic, plus a blocking/reduction rewrite that is proven
bit identical.
cargo test -p onnx-runtime-ep-cpu --features mlas --lib→ 1322 passed, 0 failed.Both feature configurations build.
cargo fmtclean.Limitations
2.2–10.2× step toward the architectural requirement that our CPU EP beat ORT on every
op it accepts; it does not finish the job, and no fallback was added. The remaining
gap is the general 16-thread scaling gap tracked in
docs/performance/CPU_ACTIVATION_GAPS.md— ORT scales these ops ~14× from 1→16 threads, we manage ~6×, because we split over
our own rayon pool rather than ORT's intra-op pool. The
host_parallelseam overKernelContext_ParallelForis the next step.PAR_MIN_LEN, so gains there are smaller and noisier than at 4 Mi.36% across 3 rounds. The 4 Mi wins are far outside that band. The 1 Mi Clip/Relu
numbers are closer to it and should be read as directional.
run_chunked,PAR_MIN_LENandparallel_dispatchesare widened topub(crate)because the three other callers live in sibling modules.