Repository navigation
perf(cpu): drop provably-redundant saturation from AVX2 tanh/sigmoid - #1121
Merged
Merged
Conversation
The AVX2 `tanh_ps` and `sigmoid_ps` already carry the Eigen rational that MLAS uses, including the input clamp to `±9` / `±18`. On top of that they clamped the result to `[-1, 1]` / `[0, 1]` and *then* ran a second saturation step: two `vcmpps` against the original input and two `vblendvps` to force the endpoint constant. That second step never changes a bit. Because the input is clamped first, the rational's largest argument is the clamp point itself, and there it already evaluates past the output clamp -- `p/q` is `1.0000001` at `v = 9`, and `p/q + 0.5` is negative at `v = -18`. The output clamp therefore saturates on its own, for every input beyond the range including `±Inf`. Removing it takes four vector ops off the per-8-lane cost of two kernels that most of the activation family is built on: `Sigmoid`, `Tanh`, SiLU / Swish, `QuickGelu` and `FastGelu` all route through one of them. Also replaces the three `vblendvps`-against-zero used to pin `-Inf` in `tanh_gelu_ps`, `quick_gelu_ps` and `erf_gelu_ps` with `vandnps`, which is exactly equivalent for an all-ones/all-zeros compare mask and one uop rather than two on Zen. `saturation_blend_is_redundant_exhaustively` keeps the removed sequence as a reference and checks the kernels against it over every finite `f32` past the clamp on both signs -- 2 086 666 240 values -- plus `±Inf`, both signed `NaN`s and both clamp boundaries. Bit-identical throughout. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #1121 +/- ##
==========================================
+ Coverage 79.90% 79.91% +0.01%
==========================================
Files 368 368
Lines 159337 159659 +322
Branches 159337 159659 +322
==========================================
+ Hits 127312 127587 +275
- Misses 27312 27355 +43
- Partials 4713 4717 +4
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
🔴 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
|
Review found the justifying comment overstated the case. At the clamp point `p` and `q` are bit-equal (`0x3fcd33e9`), so `tanh`'s rational returns *exactly* `±1.0` there, and `logistic`'s returns exactly `1.0` at `+18`. Only `-18` lands strictly outside, at `-5.96e-8`. The removal is therefore safe because the clamp is inclusive and passes equality through — not because the rational overshoots. The overshoot to `1.0000001` that the old comment cited is real but happens strictly inside the range, near `|v| = 8.9999971`; the pre-existing comment fifteen lines above already said so, and the new one contradicted it. Also adds a signalling NaN and a non-canonical quiet payload to the special value test, since `andnot` inspects every mask bit where `blendv` inspected only the sign, and notes that the exhaustive proof assumes the default rounding mode. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
justinchuby
added a commit
that referenced
this pull request
Aug 17, 2026
## What ships, and what does not `docs/performance/ABSORBING_MLAS.md` (direction set 2026-08-16, after this PR was opened) is explicit: **we do not bundle MLAS by default**, we absorb its mechanisms into native kernels, and MLAS stays in the tree as a *measurement reference*. This PR is now scoped to that role, and the framing below is corrected accordingly. `mlas` is not in the default features of `onnx-runtime-ep-cpu`, the CLI or the server. Every route added here sits behind `#[cfg(feature = "mlas")]`, so **none of the speedups in this PR reach a default build.** The tables below are measurements of a reference implementation, not a shipping improvement, and should be read as "here is how far our native kernel is from a good one". What they are good for is picking the next absorption. The gap they expose is: | family | our native kernel vs MLAS | absorbed? | |---|---|---| | `Tanh`, `Sigmoid` | 1.03–1.17x | **yes** — #1121, in the default build | | `Erf`, exact `Gelu` | 1.22–1.57x | not yet; `exp_ps` in the `erf_ps` tail is the suspect | #1121 came directly out of reading `tanh.cpp` and `logistic.cpp` for this PR: our rational was already MLAS's, and the difference turned out to be four redundant vector ops on our side. That absorption ships; this PR is the instrument that found it. ## Root cause Our pure-Rust AVX2 activation kernels are **compute-bound where ORT is memory-bound**. Measured on this host (AMD EPYC 9V74, AVX2+FMA, AVX-512 masked off), single-threaded, 1 Mi f32: | | ns/elem | effective GB/s (4 B in + 4 B out) | |---|---|---| | our `tanh_avx2` | 0.46 | 17.4 | | `MlasComputeTanh` | 0.32 | 25.0 | 25 GB/s is single-core stream bandwidth on this part. MLAS is *at* the memory wall; we are ~30% short of it and spending the difference in the polynomial. That means no amount of loop tuning, unrolling or prefetching closes the gap — **the arithmetic itself is the cost**. MLAS is already vendored in this repo (`crates/mlas-sys/vendor/mlas`) and is the library ONNX Runtime's own CPU activation kernels call. So this is not "adopt a faster approximation"; it is *stop maintaining a second, slower approximation of the function ORT already ships*. ## What changed `mlas-sys` gains `compute_tanh`, `compute_erf` and `compute_gelu_erf` next to the existing `compute_logistic`. `simd_activations` routes `Tanh`, `Sigmoid`, `Erf` and exact `Gelu` to them at or above `SIMD_MIN_LEN`. The pure-Rust path is **not** deleted. It remains the fallback: it is what builds without the `mlas` feature, what runs off x86_64, and what still serves tensors below `SIMD_MIN_LEN`, where the FFI call and MLAS's own dispatch cost more than the arithmetic saved. ## Correctness — two real divergences, both found by falsifiers I did not assume MLAS matched us. Two tests were written to try to break the claim, and both succeeded on the first run. **1. MLAS's `tanh` and `logistic` are not range-preserving.** `monotonicity_within_documented_slack` failed immediately with `tanh(-8.990999) = -1.0000001`. A dense sweep of [-20, 20] at 2^20 points: | op | max ulp vs our path | points outside the mathematical range | |---|---|---| | `Tanh` | 2 | **2934 / 1048576** | | `Sigmoid` | 1 | **2 / 1048576** | | `Erf` | 0 | 0 | | exact `Gelu` | 0 | 0 | ORT calls these directly and inherits the overshoot. Our kernels have always guaranteed the range and downstream code is entitled to rely on it — `sqrt(1 - tanh²)` turns a 1-ulp overshoot into a NaN. So we clamp. Clamping is **never less accurate**: the true value lies strictly inside the open interval, so pulling an out-of-range result to the endpoint moves it toward the exact answer. NaN survives because every comparison against it is false. **2. `MlasComputeGeluErf(-inf)` returns NaN.** It evaluates `x·0.5·(1 + erf(x/√2))` without special-casing the limit, so it computes `-inf · 0`. The limit is `0`, which is what our path already returned. Repaired with a branch-free NaN scan over the block MLAS just wrote — a NaN there can only come from a NaN input (on which both paths agree) or from this case, so repairing only the `-inf` lanes is exact. This makes us **deliberately disagree with ORT** on `Gelu(-inf)`. Returning the limit is the better answer and NaN would propagate through the rest of the graph, so this is not something to "fix" by matching ORT. After both repairs the MLAS route is **bit-identical on this host** (max ulp 0) to the pure-Rust route across the dense sweep *and* across NaN (both signs), ±Inf, ±0, ±`MIN_POSITIVE`, the smallest subnormals, `MAX`/`MIN` and the exp saturation thresholds — pinned by `mlas_ab::mlas_matches_rust_simd_on_special_values`, which fails if a future MLAS bump changes any of it. ## Benchmark Same binary, both routes compiled once, **interleaved one iteration each per round** so drift hits both sides equally, best-of-15. The `activation_bench.rs` header documents a uniform cross-build offset on byte-identical kernels, which an in-binary A/B cannot suffer. `cargo test -p onnx-runtime-ep-cpu --features mlas --release --lib mlas_ab -- --ignored --nocapture` Median of 3 full runs (min-max shown where it matters): | op | 1 Ki | 16 Ki | 256 Ki | 1 Mi | 4 Mi | |---|---|---|---|---|---| | `Erf` | 1.56x | 1.57x | 1.45x | 1.46x | 1.45x | | exact `Gelu` | 1.38x | 1.38x | 1.31x | 1.30x | 1.22x | | `Tanh` | 1.15x | 1.17x | 1.08x | 1.08x | 1.07x (1.03-1.07) | | `Sigmoid` | 1.12x | 1.15x | 1.04x | 1.04x | 1.03x (**0.99**-1.04) | **Honest reading — this is not a uniform win.** * `Erf` (1.45-1.57x) and exact `Gelu` (1.22-1.38x) are **real, repeatable wins** at every size. * `Tanh` and `Sigmoid` win 12-17% on small and medium tensors, but at 256 Ki and above they are **bandwidth-bound ties**: 1.03-1.08x, and `Sigmoid` at 4 Mi measured as low as **0.99x**. At those sizes both routes saturate memory, so the polynomial stops mattering and the fix-up pass is pure overhead. Claiming a win there would be claiming noise. An earlier revision of this description said "no size loses". That was wrong, and review caught it: it came from a single non-interleaved run. The interleaved median above is the corrected measurement. Two negative results, kept because they were measured: * A **hand-written branch-free bit-select clamp** was *slower* than `f32::clamp` (0.55 vs 0.43 ns/elem for `Tanh` at 4 Mi). LLVM already lowers `clamp` to `vmaxps`/`vminps`. The measured winner is the one in the diff. * **Blocking the fix-up passes** to keep them in cache was neutral, not a win (the cost was the extra pass, not a DRAM round trip). Kept anyway because it is free and bounds the working set. ## Scope This **does not enable the `mlas` feature anywhere it was not already enabled.** Turning it on for the plugin cdylib flips 280 `#[cfg(feature = "mlas")]` sites across matmul, gemm, sdpa, softmax, conv, pooling and the quantized paths — a far larger behavioural change than activations, needing its own conformance and benchmark evidence. That is a separate PR. So the direct beneficiaries here are builds that already opt in (`onnx-genai-cli`, `onnx-genai-engine`, `onnx-runtime-session`, `onnx-genai-bench`). ## End-to-end: does this help inside a real session? The table above is an in-binary A/B of two kernels. That is the right way to measure a kernel, but it cannot answer whether the win survives ORT's node dispatch. So the same build was run through a real ORT session both ways -- same single-node graph, same input, matched intra-op threads, p50 of 200 interleaved runs, with our EP's assignment asserted from ORT's own profiler -- with the `mlas` feature as the *only* difference: | op | elements | our us, `mlas` off | our us, `mlas` on | faster by | |---|---|---|---|---| | `Tanh` | 65536 | 38.3 | 32.2 | 1.19x | | `Tanh` | 4194304 | 2285.8 | 1625.7 | 1.41x | | `Sigmoid` | 65536 | 42.9 | 35.2 | 1.22x | | `Sigmoid` | 4194304 | 2240.6 | 1735.5 | 1.29x | | `Erf` | 65536 | 91.2 | 63.2 | 1.44x | | `Erf` | 4194304 | 5141.6 | 3623.1 | 1.42x | | exact `Gelu` | 65536 | 111.0 | 86.4 | 1.28x | | exact `Gelu` | 4194304 | 6701.1 | 5433.2 | 1.23x | The win survives, and it is if anything larger end-to-end than in the microbenchmark. **Where that leaves us against ORT.** It is a real improvement but not yet a win. In the same runs, expressed as ORT's time over ours (higher is better): | op | 4096 | 65536 | 262144 | 4194304 | |---|---|---|---|---| | `Tanh` | 0.73 -> 0.76 | 0.71 -> 0.84 | 0.61 -> 0.62 | 0.96 -> 0.85 | | `Sigmoid` | 0.74 -> 0.78 | 0.67 -> 0.81 | 0.84 -> 0.90 | 0.74 -> 0.86 | | `Erf` | 0.71 -> 0.84 | 0.66 -> 0.95 | 0.80 -> **1.02** | 0.71 -> **1.38** | | exact `Gelu` | 0.70 -> 0.81 | 0.67 -> 0.86 | 0.68 -> 0.88 | 0.75 -> 0.91 | `Erf` crosses over. The rest close most of the gap but do not. Two things are still in the way and neither is the polynomial: 1. **The fix-up pass.** ORT calls `MlasComputeTanh` once; we call it and then clamp, which is a second pass over the output. Blocking it to 2048 elements and doing it with `vmaxps`/`vminps` recovered most but not all of that -- about 12% remains at 65 Ki. Removing it entirely means giving up the range guarantee that `monotonicity_within_documented_slack` asserts, which is not a trade this PR makes. 2. **Per-node plugin overhead**, ~1.7-2.7 us, which is most of the loss at 4096 elements and none of it at 4 Mi. The `4194304` column also moves around between runs because ORT's own time there varied 1388-2199 us across repeats; our absolute times were stable to ~2%, which is why the absolute table above is the one to trust. ## Independent review Reviewed by an independent Opus reviewer: **GO WITH FINDINGS**. It reproduced every number in this description, verified the FFI against `mlas.h` (including that `MlasComputeTanh` is a template and that the shim's `<float>` instantiation links), confirmed the `MlasComputeGeluErf` no-overlap precondition is discharged by `output_direct_write_eligible` even for in-place-capable nodes, and ran its **own** adversarial sweep (|x| out to 3.4e38, the +/-88 exp-saturation band, subnormals, and the `SIMD_MIN_LEN` seam) finding **0 mismatches**. It also confirmed this PR does not widen the pre-existing scalar-vs-vector seam. All four findings are addressed on this branch: the ISA-specific bit-reproducibility caveat, the corrected benchmark claim, the new mlas-sys smoke tests, and the interleaved A/B harness. ## Limitations * Numbers are from one host. The dispatch is unconditional above `SIMD_MIN_LEN` rather than ISA-tuned, which is safe because MLAS does its own runtime ISA dispatch internally. * **Bit-reproducibility is host-specific.** Because MLAS dispatches by ISA, an AVX-512 or non-AVX2 machine may pick a different kernel than the AVX2+FMA one the pure-Rust path mirrors, so the two routes need not agree bitwise there. Before this change, `mlas`-on and `mlas`-off builds agreed bitwise for these four ops on *every* host; that is what is traded for the speed. The special-value test pins whatever ISA it runs on. * f16/bf16 still widen to f32 first; this PR does not change that path. * The clamp and NaN-scan fix-ups cost ~0.1 ns/elem. They are the price of keeping our stronger numerical contract, and are why `Tanh`/`Sigmoid` show 1.0-1.2x rather than the 1.3-1.55x the raw MLAS kernels reach. ## Validation * `cargo fmt --all -- --check` * `cargo clippy -p onnx-runtime-ep-cpu --features mlas --all-targets` — clean (one pre-existing `needless_return` in `matmul_nbits.rs`, untouched here) * `cargo test -p onnx-runtime-ep-cpu --lib` — 1323 passed (feature off) * `cargo test -p onnx-runtime-ep-cpu --features mlas --lib` — 1341 passed * `cargo test -p mlas-sys` — 39 passed, plus 5 new `tests/transcendentals.rs` cases pinning the new FFI at the crate boundary against a correctly-rounded `f64` reference (so they check the *function*, not MLAS's own polynomial), including empty-slice and length-mismatch cases --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Resolves the additive-test-module conflict with #1111's mlas_ab reference module: both modules are kept, and the saturation removal is re-applied on top of main's dispatch_mlas routing. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
justinchuby
marked this pull request as ready for review
August 17, 2026 11:36
justinchuby
pushed a commit
that referenced
this pull request
Aug 17, 2026
Same end-of-file module collision as the previous merge, now also against #1121's saturation_absorption module. Resolved the same way: main's modules kept intact, thread_invariance re-appended whole. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
justinchuby
added a commit
that referenced
this pull request
Aug 17, 2026
## What Delete the MLAS route for `Tanh` and `Sigmoid`. Both ops now run the pure-Rust AVX2 kernel in every build, `mlas` on or off. ## Root cause #1111 added MLAS routes for `Tanh`, `Sigmoid`, `Erf` and exact `Gelu` because MLAS was measurably faster. That was true at the time. Two things have changed since: 1. **The polynomial was never the win.** Reading `mlas/lib/tanh.cpp` and `logistic.cpp` showed both are Eigen rational approximations `P(v²)·v / Q(v²)` with the range clamp fused into the same pass. Our `avx2::tanh_ps`/`sigmoid_ps` already used the *identical* rational and the *identical* constants. The route was buying MLAS's loop, not a better approximation. 2. **#1121 removed the loop deficit.** Our kernels carried ~6 redundant vector ops (a result clamp plus two saturation `cmpps`+`blendvps` pairs) 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`/`MlasComputeLogistic` are **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_slack` asserts 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 (zero `NOT-ASSIGNED`). p50 µs, lower is better. | op | n | `mlas` **off** | `mlas` **on** (before) | on/off | |---|---:|---:|---:|---:| | Tanh | 4 Ki | **7.58** | 7.68 | 1.01x slower | | Tanh | 16 Ki | **11.93** | 12.91 | 1.08x slower | | Tanh | 64 Ki | **32.54** | 36.20 | 1.11x slower | | Tanh | 256 Ki | **105.64** | 146.73 | **1.39x slower** | | Tanh | 1 Mi | **438.43** | 469.44 | 1.07x slower | | Tanh | 4 Mi | **1896.32** | 2226.19 | 1.17x slower | | Sigmoid | 4 Ki | **7.27** | 7.49 | 1.03x slower | | Sigmoid | 16 Ki | **11.99** | 13.35 | 1.11x slower | | Sigmoid | 64 Ki | **31.16** | 38.26 | **1.23x slower** | | Sigmoid | 256 Ki | **113.69** | 129.34 | 1.14x slower | | Sigmoid | 1 Mi | **408.89** | 525.83 | **1.29x slower** | | Sigmoid | 4 Mi | **1887.40** | 2268.40 | 1.20x slower | 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: | op | n | `mlas` **off** | `mlas` **on** (before) | on/off | |---|---:|---:|---:|---:| | Tanh | 64 | **5.61** | 5.84 | 1.04x slower | | Tanh | 256 | **5.46** | 5.73 | 1.05x slower | | Tanh | 1024 | **5.99** | 6.23 | 1.04x slower | | Sigmoid | 64 | **5.37** | 5.64 | 1.05x slower | | Sigmoid | 256 | **5.45** | 5.60 | 1.03x slower | | Sigmoid | 1024 | **6.09** | 6.27 | 1.03x slower | 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 `mlas` enabled 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 mode `docs/performance/ABSORBING_MLAS.md` names: "the configuration we measured was not the configuration we shipped". ### Effect against ORT Restated as ORT p50 ÷ ours (>1 = we win), 1 thread: | op | n | shipped before | shipped after | |---|---:|---:|---:| | Tanh | 256 Ki | 0.71 | **0.99** | | Tanh | 4 Mi | 0.99 | **1.16** | | Sigmoid | 64 Ki | 0.75 | **0.93** | | Sigmoid | 1 Mi | 0.75 | **0.96** | Still short of parity at several sizes — that work continues — but this recovers most of a self-inflicted gap. ## What is kept `Erf` and exact `Gelu` keep their MLAS routes, because the arithmetic is different: - **`Erf`** needs 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). - **exact `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 and `mlas`-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_ab` test module as `tanh_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 past `SIMD_MIN_LEN` and 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 failed** - `cargo build -p onnx-runtime-ep-cpu` (default, no `mlas`) → clean - `cargo fmt --all -- --check` → clean ## Limitations - Measured on one host (AMD EPYC 9V74, Zen 4, AVX2+FMA, AVX-512 masked off). On a machine where MLAS dispatches an AVX-512 kernel and we stay on AVX2, the balance could differ — but the fix-up pass is pure memory traffic and does not get cheaper with a wider ISA, and the clamp is required regardless, so the direction should hold. - Only 1-thread numbers are quoted. The 16-thread picture is dominated by a much larger threading gap that this PR does not touch and that is tracked separately. - This does not make `Tanh`/`Sigmoid` beat 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`: 1. *(MINOR)* The rationale comment had merged into `MLAS_FIXUP_BLOCK`'s doc comment, so rustdoc rendered "Why `Tanh` and `Sigmoid` do 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. 2. *(NIT)* `dispatch_mlas!`'s doc still said "these four ops" when the macro now has two callers. Reworded, with a pointer to the removal note. 3. *(NIT)* "lost at every size" was every size *measured*, leaving `[32, 4096)` to extrapolation. Measured — see the second table above. The reviewer independently verified the route removal is complete (`mlas_sys::compute_tanh`/`compute_logistic` are unreachable from shipped code), that `_mlas_features` really does ship `mlas` on linux-x86_64, that `erf_gelu_mlas`'s scan is genuinely non-writing on the all-finite path, that `mlas_clamped_ref` is 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)`, `audit` and 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 pristine `origin/main` too, at every one of the last five commits. Within `CLI ORT` the activation-relevant step fails on `shape_inference_coverage::every_registered_op_has_a_shape_rule_or_is_a_known_gap`, which reproduces verbatim on an untouched `origin/main` checkout and names six `Nchwc*` 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 `GetCapability` **is** 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_ort` and `activation_and_norm_ops_clear_every_capability_filter` — pass. --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
justinchuby
added a commit
that referenced
this pull request
Aug 18, 2026
…ap_bias_ps (#1227) ## What this is `map_ps` is the loop driver behind every unary activation on AVX2. A merge (the `#1037` / `#1074` layering) inserted `map_bias_ps` *between* `map_ps`'s doc comment and `map_ps` itself, so the comment and **both of its attributes** ended up on the wrong function: ```rust /// Apply an 8-lane kernel across a slice. ... <- map_ps's doc #[inline] #[target_feature(enable = "avx2,fma")] <- map_ps's attributes /// Like [`map_ps`], but adds a bias row ... pub(super) unsafe fn map_bias_ps(...) <- ...on map_bias_ps pub(super) unsafe fn map_ps(...) <- bare ``` `map_bias_ps` needs those attributes too, so nothing was broken outright — but `map_ps` was left with no `#[inline]` and no `avx2,fma`. ## Why the missing `target_feature` is a hazard LLVM will not inline a callee that requires a feature its caller lacks (a *strict* superset; equal feature sets inline fine). Every kernel closure `map_ps` takes (`erf_ps`, `tanh_ps`, `sigmoid_ps`, ...) is `avx2,fma`. With `map_ps` compiled at baseline features, those closures are not inlinable into its loop. It is harmless **today** only by luck of ordering: every caller (`tanh_avx2`, `erf_avx2`, ...) is itself `avx2,fma`, so `map_ps` gets inlined *upwards* into the caller first, and once there the closure folds in. That is an inlining decision, not a guarantee. Grow `map_ps` past the inline threshold — which is exactly what an unroll experiment does — and the closure becomes a real call per 8 elements, and a kernel like `erf_ps` has to re-materialise all of its constants on every one of those calls. I hit this while unrolling `map_ps`, which is how it surfaced. ## Evidence that this changes nothing today Disassembled `tanh_avx2` out of the built rlib before and after: ``` diff <(objdump -d ... base) <(objdump -d ... fixed) -> identical ``` Byte-identical, and `erf_avx2` / `sigmoid_avx2` keep their instruction counts (216 / 117). No `call` to any kernel remains in the lib. So this is a latent-hazard and documentation fix with **zero** codegen delta — not a perf change, and it needs no A/B. ## Also recorded here: a negative result on the divide While looking at these loops I measured one algorithmic change and it lost, so it is written down rather than left for someone to retry. `tanh_ps` and `sigmoid_ps` each end in `_mm256_div_ps`. `vdivps ymm` is the only instruction in either loop that is not fully pipelined (Zen 4 retires one about every 4.5 cycles), so replacing it with `vrcpps` + one Newton-Raphson step — `y1 = y0 + y0*(1 - q*y0)`, four pipelined uops for one blocking one — looked like the obvious win. It is not: | case | ORT p50 (control) | divide | rcp + NR | |---|---|---|---| | `bench_tanh_f32_4k` | 0.0031 ms, all 6 runs | 1.518 / 1.525 / 1.526 | 1.550 / 1.570 / 1.570 | Interleaved rebuild-and-alternate, 3 rounds, 1 thread, `taskset -c 8-15`. The 4k case is the only trustworthy one here — it is L1-resident and ORT's own p50 was identical to four decimal places across all six runs, whereas on the 1M cases ORT's p50 moved 4.6x mid-session and neither arm is quotable. Consistent ~3% **regression**. The `rcp` + 2 FMA dependency chain is about as long as the divide it replaces, so a latency-bound kernel gains nothing and just pays two extra uops. It also breaks four exactness proofs that #1121 landed: `tanh(±Inf)` comes back as `0.99999994` instead of exactly `±1.0`, because the input clamp makes `p` and `q` bit-equal at `|v| = 9` and only an exact divide turns that into exactly `1.0`. Buying that back would mean restoring the saturation blend #1121 proved redundant over 1.05 billion inputs — to fund a change that is already slower. Not pursued. No fallback involved; the kernel keeps its own divide. ## Validation - `cargo test --release -p onnx-runtime-ep-cpu --lib` — 1337 passed, 0 failed - `cargo clippy --release -p onnx-runtime-ep-cpu --all-targets` — clean - `cargo fmt --all --check` — clean - `tanh_avx2` disassembly identical to `main` ## Review Opus review: no blockers, no should-fixes. It independently confirmed the three things the PR rests on: every one of the 16 `map_ps` and 2 `map_bias_ps` call sites is already inside a `#[target_feature(enable = "avx2,fma")]` function (so nothing becomes unsound or fails to compile), `map_bias_ps` keeps exactly the attributes it had, and there are **no bare/function-pointer references** to either. It also did the scan I most wanted: all ~130 `fn` definitions in the file were checked for the same orphaned-doc-comment-after-attribute signature that caused this bug. **Zero other anomalies** — every `*_ps` helper and every `*_avx2` dispatcher carries the attributes it should. So this is the only instance. - **NIT (applied)** — the doc comment said LLVM won't inline a callee whose features are "a superset" of the caller's. Read non-strictly that would forbid inlining at *equal* feature sets, which is the opposite of what the fix relies on. Reworded to "requires a feature its caller lacks — a *strict* superset; equal feature sets inline fine". --------- 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
2 participants
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 this is
tanh_psandsigmoid_psare the two AVX2 primitives most of the CPUactivation family is built on:
Tanh,Sigmoid, SiLU / Swish,QuickGeluandFastGeluall end up in one of them. Both already carry the Eigen rationalthat MLAS uses — the polynomial was never the problem. What they also carried
was a redundant saturation step, and it is not free.
The shape of both kernels is: clamp the input to the rational's valid range
(
±9fortanh,±18forlogistic), evaluateP(v²)·v / Q(v²), clamp theoutput to the function's mathematical range. On top of that, both then ran a
second saturation:
Four vector ops, two of them
vblendvps, which is two uops on Zen.It never changes a bit. The input clamp happens first, so the largest
argument the rational ever sees is the clamp point itself — and there it has
already reached the output bound:
tanh,v = 9pandqare bit-equal (0x3fcd33e9), sop/qis exactly1.01.0tanh,v = -9-1.0-1.0logistic,v = 18p/q + 0.5is exactly1.01.0logistic,v = -18p/q + 0.5is-5.96e-80.0Three of the four are equality, not slack, and the inclusive
minps/maxpsclamp passes equality through unchanged. (The rational's overshoot to
1.0000001happens strictly inside the range, near|v| = 8.9999971— thatis what the output clamp is there for, and it is unrelated to saturation.)
Every input past the range,
±Infincluded, therefore already exits with thesaturated value. The blends were re-deriving a result the clamp had produced two
instructions earlier.
Proof, not argument
saturation_blend_is_redundant_exhaustivelykeeps the deleted sequence as areference implementation and checks the shipped kernels against it over every
finite
f32beyond the clamp, both signs — 1 047 527 424 values fortanhand 1 039 138 816 for
sigmoid, 2 086 666 240 in total — plus±Inf, bothsigned
NaNs,f32::MAX/MIN, and both clamp boundaries at ±2 ULP.Bit-identical everywhere. Not within tolerance — identical.
The sweep runs under the default rounding mode. The property was also checked
by hand under round-down, round-up and round-toward-zero, and holds in all
three; only round-to-nearest is exercised in CI.
It runs in 19 s in release and is
#[ignore]d because an unoptimised test buildtakes minutes. Two non-ignored tests cover the decade above each boundary
(32.5 M and 21.7 M values) and the special values, and run in 12 s in the
default
cargo testprofile, so the property stays guarded on every CI run.Also
The three
vblendvps-against-zero that pin-Inf → 0intanh_gelu_ps,quick_gelu_psanderf_gelu_psbecomevandnps. Blending against zero isan
andnotwhen the mask is all-ones or all-zeros, which is exactly whatvcmppsproduces. One uop instead of two, same bits.Measured
Same-machine, alternating build A/B through a real ORT session — one node,
one input,
intra_op_num_threads=1,RAYON_NUM_THREADS=1, p50 of 30 runs,median of 3 alternating rounds per build, with our EP's assignment asserted from
ORT's own profiler on every row. AMD EPYC 9V74, AVX2+FMA, ORT 1.28.0.
Measured on top of #1097 and #1105, because on
maintoday the assignmentpolicy declines these ops to ORT's CPU EP, so our kernel never runs and the
measurement is vacuous. (The first run of this A/B was vacuous for exactly
that reason — both columns were ORT to within 0.1%. The harness's assignment
check is what caught it.)
TanhTanhTanhTanhTanhSigmoidSigmoidSigmoidSigmoidSigmoidFastGeluFastGeluFastGeluFastGeluQuickGeluQuickGeluQuickGeluGeluGeluErf(control)Erf(control)Sqrt(control)Sqrt(control)ErfandSqrttouch neither primitive and are flat, which is the check thatthe rest is real and not drift. Exact
Gelugoes througherf_ps, so it onlycollects the
andnot, and moves ~2-3%.Against ORT
Same runs, ORT's CPU EP as the control, ORT time over ours — above
1.00wewin:
TanhSigmoidQuickGeluFastGeluGeluThis does not claim we now beat ORT. It moves every one of these families
toward it, takes
Tanhpast it at >=1 Mi andQuickGelupast it from 16 Ki up,and leaves
FastGelu, exactGeluandErfstill behind. The remaining gap isnot in these two primitives, and the next steps are elsewhere:
erf_ps'sexp_pstail, and the ~1.7-2.7 us fixed per-node plugin overhead that dominatesat 4096 elements.
Direction
This is an absorption in the sense of
docs/performance/ABSORBING_MLAS.md:the win lands in the native kernel, in the default build, with no
mlasfeature and no dependency. Reading MLAS's
tanh.cppandlogistic.cppis whatmade the redundancy visible — MLAS does not have this step, because it does not
promise the output range we promise, and comparing the two instruction sequences
is what showed that our extra promise costs nothing to keep and four ops to
re-state.
Correctness
saturation_blend_is_redundant_exhaustively— 2 086 666 240 values, bit-identical.onnx-runtime-ep-cpulib suite: 1326 passed.tanh/sigmoid/geluaccuracy test passes unmodified.Limitations
machine-independent; the exact percentages are not.
1048576and4194304rows sit where ORT's own timing is least stable(its control column moved up to 1.6x between rounds on
Tanh); medians ofthree alternating rounds are reported, and the
Erf/Sqrtcontrols are theevidence that the reported deltas are larger than that drift.
Independent review
Reviewed by an independent Opus reviewer, verdict GO WITH FINDINGS. The
reviewer re-derived the exhaustive proof independently (2 095 054 848
tanhand2 078 277 632
sigmoidvalues, zero disagreements), confirmedandnot ≡ blendvacross quiet, signalling and non-canonical
NaNpayloads,±0,±Infandsubnormals, and confirmed the redundancy holds under all four MXCSR rounding
modes.
All findings are applied:
body said
p/qatv = 9is1.0000001. It is exactly1.0—pandqare the same bit pattern — so the safety argument rests on equality with an
inclusive clamp, not on slack, and the old wording also contradicted a
correct comment fifteen lines above it. Both the comment and the body now
state the exact values, and the sigmoid comment no longer claims the result
lands "outside
[0, 1]on both ends" when at+18it is exactly on theboundary. Corrected in
d32e52b.NaNs added to the special-value test.sufficient: on
mainthe assignment policy declines these ops, so thekernels are not reached until Stop handing activation nodes to ORT's CPU EP #1097 lands. That is why the benchmark is
stacked, and it is stated in the section above.