Repository navigation
perf(cpu): stop routing Tanh/Sigmoid through MLAS - #1124
Merged
Merged
Conversation
The MLAS route for these two has been a net loss since #1121. Both kernels already used the same Eigen rational MLAS does, so the route never bought a better polynomial -- it bought MLAS's loop. Once the redundant saturation blends came out of our AVX2 kernels the loop stopped being the bottleneck, and the clamp pass the route needs (MLAS's rational is not range-preserving) now costs more than it saves. Measured on this host, mlas-off vs mlas-on p50 at 1 thread, medians of 3 interleaved rounds with EP assignment asserted: Tanh 64 Ki 32.5 / 36.2 256 Ki 105.6 / 146.7 4 Mi 1896 / 2226 Sigmoid 64 Ki 31.2 / 38.3 256 Ki 113.7 / 129.3 4 Mi 1887 / 2268 The route lost at every size, by up to 1.39x. Since #1115 the wheel ships with `mlas` on, so this was shipping. Erf and exact Gelu keep their routes: Erf needs no fix-up and still wins 1.4-1.6x, Gelu's fix-up is a non-writing compare scan and it still wins 1.17x at 4 Mi. The clamped MLAS reference moves into the test module, so the benchmark, special-value and dense-ULP comparisons all keep running against it. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #1124 +/- ##
==========================================
+ Coverage 79.85% 80.53% +0.67%
==========================================
Files 366 366
Lines 157049 157082 +33
Branches 157049 157082 +33
==========================================
+ Hits 125412 126500 +1088
+ Misses 26935 25877 -1058
- Partials 4702 4705 +3
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
Opus review of #1124 found the rationale comment had merged into MLAS_FIXUP_BLOCK's doc comment, so rustdoc showed "Why Tanh and Sigmoid do not have an MLAS route" as the summary of an 8192-element block-size constant -- and, being cfg-gated, the explanation for the route's absence vanished from a default no-mlas build. It is now a free-standing comment. Also corrects the dispatch_mlas! doc, which still said "these four ops" when the macro now has two callers. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
🔴 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
Make the MLAS activation routes go through `run_chunked`, the shared
chunking/parallel seam every pure-Rust route already uses.
One line:
```rust
- f(input, output);
+ run_chunked(input, output, f);
```
## Root cause
`dispatch_mlas!` called its kernel on the whole tensor and returned:
```rust
if input.len() >= SIMD_MIN_LEN {
let f: fn(&[f32], &mut [f32]) = $mlas;
f(input, output); // <-- straight past run_chunked
return;
}
```
`run_chunked` was the only thing that split work across the pool. So
**every MLAS-routed op ran single-threaded, regardless of how many
threads the session had.**
This was harmless when it was written — `run_chunked` was a plain loop
then. #1105 taught it to parallelise above `PAR_MIN_LEN`, and from that
moment the two builds diverged: an `mlas`-off build scaled with the
pool, an `mlas`-on build did not. Since #1115 the wheel ships `mlas` on
x86_64, so the shipped configuration was the non-scaling one.
Nothing caught it because the results were still correct, and
correctness is all the existing tests checked. `thread_invariance`
asserts a kernel gives the same answer serial and parallel — which a
route that *never* goes parallel satisfies trivially.
## Benchmark
Same machine (AMD EPYC 9V74, 32 vCPU, AVX2+FMA), `ab_ort.py`, `--iters
30 --reps 5`, medians of 3 interleaved rounds alternating the two
builds, EP assignment asserted from ORT's own profiler (zero
`NOT-ASSIGNED`). Both arms built `--features mlas --release` from the
same base. p50 µs.
### 16 threads — the ops that take the MLAS route
| op | n | before | after | speedup |
|---|---:|---:|---:|---:|
| Erf | 1 Mi | 897.3 | **393.9** | **2.28x** |
| Erf | 4 Mi | 3558.7 | **561.6** | **6.34x** |
| Gelu (exact) | 1 Mi | 1304.0 | **552.8** | **2.36x** |
| Gelu (exact) | 4 Mi | 5458.7 | **1134.3** | **4.81x** |
| Tanh | 1 Mi | 480.8 | **300.4** | **1.60x** |
| Tanh | 4 Mi | 2219.6 | **425.9** | **5.21x** |
| Sigmoid | 1 Mi | 503.5 | **298.6** | **1.69x** |
| Sigmoid | 4 Mi | 1991.6 | **430.0** | **4.63x** |
(`Tanh`/`Sigmoid` measured on a base that predates #1124; they no longer
take the MLAS route, but they exercised the same mechanism and are kept
here as evidence of it.)
### 16 threads — ops with no MLAS route, as the control
Unchanged within noise, which is what says the win is the routing change
and not drift:
| op | n | before | after |
|---|---:|---:|---:|
| FastGelu | 1 Mi | 413.9 | 395.5 |
| FastGelu | 4 Mi | 600.2 | 605.4 |
| QuickGelu | 1 Mi | 329.7 | 347.1 |
| QuickGelu | 4 Mi | 529.1 | 545.9 |
| Sqrt | 1 Mi | 255.3 | 232.3 |
| Sqrt | 4 Mi | 352.8 | 375.1 |
### 1 thread — no regression
`run_chunked` returns before touching rayon below `PAR_MIN_LEN`, and
`par_chunk_len` declines to split when there is nothing to gain, so the
single-threaded path is untouched:
| op | n | before | after |
|---|---:|---:|---:|
| Erf | 4 Mi | 3628.0 | 3643.8 |
| Gelu | 4 Mi | 5466.7 | 5479.0 |
| Erf | 65 Ki | 62.2 | 63.5 |
| Gelu | 65 Ki | 86.7 | 86.4 |
Every cell is inside run-to-run noise (<0.5%).
## Correctness
`run_chunked` splits a slice into disjoint sub-slices and calls the same
function on each. The MLAS activation entry points (`MlasComputeErf`,
`MlasComputeGeluErf`) are elementwise and take no threadpool argument,
so splitting cannot change a result: element `i` depends only on input
`i`. The existing
`thread_invariance::unary_kernels_are_thread_count_invariant` already
covers `erf` and `erf_gelu` and asserts serial and parallel agree **bit
for bit** — it now actually exercises the parallel branch for them,
where before it compared serial against serial.
## Regression test
`parallel_reachability` asserts the mechanism rather than the output: a
tensor over `PAR_MIN_LEN`, submitted from outside the pool, must
increment `run_chunked`'s parallel-branch counter.
Proven to falsify — reverting just the one line makes it fail with:
```
erf: a 1052675-element call did not reach run_chunked's parallel branch, so it
runs single-threaded no matter how large the pool is. A kernel that calls its
backend directly and returns will fail here.
```
while `pure_rust_kernels_go_through_run_chunked` keeps passing, so it is
not a blanket assertion that would pass on anything. It also guards the
two preconditions that could make it vacuous (pool must have ≥2 threads;
must not already be inside the pool).
The counter is `#[cfg(test)]` and **thread-local**, not a global atomic,
so concurrently running tests cannot bump each other's count. Non-test
builds get an empty `#[inline(always)]` function.
## Tests
- `cargo test -p onnx-runtime-ep-cpu --features mlas --lib` → **1312
passed, 0 failed**
- `cargo test -p onnx-runtime-ep-cpu --lib` (no `mlas`) → **1294 passed,
0 failed**
- `cargo clippy -p onnx-runtime-ep-cpu --features mlas --lib --tests` →
clean
- `cargo fmt --all -- --check` → clean
### A pre-existing intermittent crash, ruled out as mine
While validating I hit an intermittent `SIGSEGV` in the full
`onnx-runtime-ep-cpu` lib test binary. It is **not** from this change:
- Activation tests alone, 40 consecutive runs on this branch: **0
failures**.
- Full suite, 14 consecutive runs on this branch: **0 segfaults**.
- Full suite, 14 consecutive runs on pristine `origin/main`: **1
segfault**.
- `origin/main` also fails
`kernels::sdpa::tests::sdpa_dispatch_matches_scalar_oracle_across_shapes`
intermittently (~1 in 8 full-suite runs), unrelated to activations.
Flagging for whoever owns `sdpa`/`qmoe`; not chased here as it is
outside this PR's scope and reproduces without it.
## Limitations
- The thresholds are still `PAR_MIN_LEN` = 1 Mi and `PAR_MIN_CHUNK` =
256 Ki. #1105 recorded that one global constant cannot serve kernels
with different per-element costs, and this PR does not change that — it
only makes the MLAS routes obey whatever the constants say. Both are now
on the same footing, which is the precondition for tuning them per
kernel.
- This closes a self-inflicted gap; it does not by itself make these ops
beat ORT at 16 threads. At 4 Mi, `Erf` goes from 0.11x to 0.66x of ORT
and `Gelu` from 0.09x to 0.42x. Real wins, still short of parity, and
the remaining multi-thread gap stays open.
- Measured on one 32-vCPU host.
## Review
Independent Opus review: **GO WITH FINDINGS**.
The reviewer independently confirmed the parts that could have been
wrong: that `run_chunked` splits into disjoint sub-slices at identical
input/output offsets (so `erf_gelu_mlas`'s internal 8192-element
blocking and its `(x, y)`-paired NaN repair still line up inside a
chunk, trailing partial block included); that
`MlasComputeErf`/`MlasComputeGeluErf` are elementwise with no threadpool
argument and no shared mutable state, so concurrent calls from rayon
workers are sound; that the `SIMD_MIN_LEN` (32) / `PAR_MIN_LEN` (1 Mi)
composition has no gap and short tensors still return before touching
rayon; that the counter increments on the calling thread and cannot be
cross-contaminated by concurrent tests; that `note_parallel_dispatch`
compiles to nothing in a non-test build; and that the speedup arithmetic
and the "serial against serial" claim are accurate.
Two findings:
1. *(MAJOR, pre-existing, outside this file)* **The same pathology
exists in three sibling elementwise activations.** `silu_f32_slice`
(`kernels/activations.rs:332`), `relu_contiguous_f32_mlas`
(`kernels/relu.rs:144`) and the two Clip paths
(`kernels/selection.rs:182`, `kernels/conv.rs:732`) all call
`mlas_sys::compute_*` on the whole tensor with no chunking seam — and
`mlas-sys` documents those entry points as "Single threaded; callers
shard across threads themselves". `silu_f32_slice` is worse than the ops
fixed here: it follows the MLAS call with a *second* full-tensor serial
correction loop. SiLU is the SwiGLU activation, so this is a hot path.
Confirmed by grep: `run_chunked`/`run_chunked_rows` exist only in
`simd_activations.rs`, and every entry point *inside* that file does
route through the seam — so this PR is complete for its stated scope.
Being a different subsystem with its own correctness wrinkle (the
correction pass), it gets its own PR rather than being folded in here.
2. *(NIT)* `parallel_reachability` hard-fails rather than skips on a
single-threaded pool. This is deliberate and mirrors the pre-existing
`thread_invariance::assert_same` guard, so it adds no environmental
assumption the suite did not already make — a test that silently skipped
would be worse, since the bug it guards is invisible in the output.
---------
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
1 participant
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.
What
Delete the MLAS route for
TanhandSigmoid. Both ops now run the pure-Rust AVX2 kernel in every build,mlason or off.Root cause
#1111 added MLAS routes for
Tanh,Sigmoid,Erfand exactGelubecause MLAS was measurably faster. That was true at the time. Two things have changed since:mlas/lib/tanh.cppandlogistic.cppshowed both are Eigen rational approximationsP(v²)·v / Q(v²)with the range clamp fused into the same pass. Ouravx2::tanh_ps/sigmoid_psalready used the identical rational and the identical constants. The route was buying MLAS's loop, not a better approximation.cmpps+blendvpspairs) that were proven bit-identically redundant over 2 086 666 240 values. With those gone the pure-Rust kernel is no longer the slower loop.Meanwhile the MLAS route always had a tax the native path does not:
MlasComputeTanh/MlasComputeLogisticare not range-preserving. Over a dense sweep of [-20, 20], tanh lands outside[-1, 1]for 2934 of 1048576 points (up to 2 ulp,tanh(-8.990999) = -1.0000001) and logistic outside[0, 1]for 2. Our kernels guarantee the range,monotonicity_within_documented_slackasserts it, and downstream code relies on it —sqrt(1 - tanh²)turns a 1-ulp overshoot into a NaN. So the route had to re-read every block MLAS had just written and clamp it.Net: the fix-up pass now costs more than the polynomial saves.
Benchmark
Same machine (AMD EPYC 9V74, AVX2+FMA, AVX-512 masked), same harness, 1 thread,
--iters 30 --reps 5, medians of 3 interleaved rounds alternating the two builds, EP assignment asserted from ORT's own profiler (zeroNOT-ASSIGNED). p50 µs, lower is better.mlasoffmlason (before)The route lost at every size measured, for both ops.
The route was reachable for any length
>= SIMD_MIN_LEN, which is 32, so the table above left the band[32, 4096)unmeasured. Opus review flagged that "no crossover to preserve" therefore rested on extrapolation, so it was measured too:mlasoffmlason (before)Still no crossover — MLAS loses across the entire reachable range, 64 to 4 Mi. (Below ~4 Ki both builds sit on a ~5.4 µs per-node floor, so these differences are small in relative terms; the point is only that the sign never flips.) This is therefore a deletion rather than a threshold.
Why this matters now
Since #1115 the wheel ships with
mlasenabled on x86_64 (python/nxrt-ep-cpu/setup.py::_mlas_features). This was not a measurement-only artefact — it was the shipped configuration. That is exactly the failure modedocs/performance/ABSORBING_MLAS.mdnames: "the configuration we measured was not the configuration we shipped".Effect against ORT
Restated as ORT p50 ÷ ours (>1 = we win), 1 thread:
Still short of parity at several sizes — that work continues — but this recovers most of a self-inflicted gap.
What is kept
Erfand exactGelukeep their MLAS routes, because the arithmetic is different:Erfneeds no fix-up at all — MLAS's result is used verbatim — and still wins 1.40x at 64 Ki and 1.59x at 4 Mi (5654 → 3551 µs).Gelu's fix-up is a non-writing compare scan (repairing MLAS's-inf → NaN, which is-inf · 0), not a read-modify-write pass, and it still wins 1.17x at 4 Mi (6485 → 5550 µs).So the rule this establishes is narrower than "MLAS is slower": an MLAS route that needs a writing fix-up pass no longer pays for itself; one that needs no fix-up, or only a scan, still does.
Correctness
Nothing about the numerics of the shipped path changes for
Tanh/Sigmoid— it becomes the pure-Rust path, which is what a default (mlas-off) build has always run and what every dense numeric sweep in the file already tests. This strictly reduces cross-build divergence:mlas-on andmlas-off builds are now bit-identical for these two ops on every host, which they were not while MLAS's ISA-dispatched kernel was in play.The clamped MLAS reference is preserved verbatim inside the
mlas_abtest module astanh_mlas_ref/sigmoid_mlas_ref, so all three comparisons keep running against it:mlas_matches_rust_simd_on_special_values— bit-equality on NaN (both signs), ±Inf, ±0, ±MIN_POSITIVE, smallest subnormals,MAX/MIN, ±88.5, ±1e-30, padded pastSIMD_MIN_LENand repeated so every value lands at every lane offset.mlas_vs_rust_dense_ulp_sweep— dense ULP sweep over [-20, 20].mlas_vs_rust_simd— the benchmark that produced the table above.Had these been deleted along with the route, a future MLAS bump could have diverged silently.
Tests
cargo test -p onnx-runtime-ep-cpu --features mlas --lib→ 1309 passed, 0 failedcargo build -p onnx-runtime-ep-cpu(default, nomlas) → cleancargo fmt --all -- --check→ cleanLimitations
Tanh/Sigmoidbeat ORT at every size. It removes a regression; the remaining gap is real and open.Review
Independent Opus review: GO WITH FINDINGS, three findings, all addressed in
4b47896:MLAS_FIXUP_BLOCK's doc comment, so rustdoc rendered "WhyTanhandSigmoiddo not have an MLAS route" as the summary of an 8192-element block-size constant — and, being#[cfg(feature = "mlas")]-gated, the explanation for the route's absence disappeared from a default build. It is now a free-standing comment.dispatch_mlas!'s doc still said "these four ops" when the macro now has two callers. Reworded, with a pointer to the removal note.[32, 4096)to extrapolation. Measured — see the second table above.The reviewer independently verified the route removal is complete (
mlas_sys::compute_tanh/compute_logisticare unreachable from shipped code), that_mlas_featuresreally does shipmlason linux-x86_64, thaterf_gelu_mlas's scan is genuinely non-writing on the all-finite path, thatmlas_clamped_refis faithful to the deleted code rather than a straw man, and that the 1.39x / 1.59x / 1.17x arithmetic checks out.CI
Fast (Linux x86_64)(the required check),EP conformance (Linux x86_64),Rust quality,Miri unsafe-crate soundness,Rust coverage (Linux x86_64),auditand both codecov lanes are green.The red lanes —
CLI ORT (Linux/Windows x86_64),CUDA compile (Linux/Windows x86_64),Rust coverage (macOS arm64),Rust (Windows ARM64)— are red on pristineorigin/maintoo, at every one of the last five commits. WithinCLI ORTthe activation-relevant step fails onshape_inference_coverage::every_registered_op_has_a_shape_rule_or_is_a_known_gap, which reproduces verbatim on an untouchedorigin/maincheckout and names sixNchwc*ops (NchwcConv,NchwcMaxPool, ...) with no shape rule. Nothing to do with activations, and not introduced here.Worth flagging for whoever owns those ops: that test failing means
GetCapabilityis silently handing six registered ops to ORT's CPU EP, which the project's architecture policy forbids. The two activation-side falsifiers in the same file —no_activation_or_norm_op_is_left_to_ortandactivation_and_norm_ops_clear_every_capability_filter— pass.