Skip to content

Repository navigation

perf(cpu): stop routing Tanh/Sigmoid through MLAS #1124

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
justinchuby merged 2 commits into from
Aug 17, 2026
Merged

perf(cpu): stop routing Tanh/Sigmoid through MLAS #1124

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
4c2f661
perf(cpu): stop routing Tanh/Sigmoid through MLAS
justinchuby Aug 17, 2026
4b47896
docs(cpu): fix review findings on the MLAS route removal
justinchuby Aug 17, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
134 changes: 72 additions & 62 deletions crates/onnx-runtime-ep-cpu/src/kernels/simd_activations.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -588,6 +588,9 @@ macro_rules! dispatch {

/// Route an elementwise f32 kernel to MLAS when the `mlas` feature is on.
///
/// Only `Erf` and exact `Gelu` still take this route. `Tanh` and `Sigmoid`
/// used to and no longer do — see the note below this macro for why.
///
/// MLAS is the library ONNX Runtime's own CPU activation kernels call, so this
/// is not "a faster approximation" — it is *the same* polynomial ORT uses.
///
Expand All @@ -597,7 +600,7 @@ macro_rules! dispatch {
/// CI (AVX2+FMA, AVX-512 masked) the two routes are bit-identical, pinned by
/// `mlas_ab::mlas_matches_rust_simd_on_special_values`; elsewhere they may
/// differ within the documented tolerance. Before this change an `mlas`-on and
/// an `mlas`-off build agreed bitwise for these four ops on every host, and
/// an `mlas`-off build agreed bitwise for these ops on every host, and
/// that is what is being traded for the speed. Measured on this
/// host (AMD EPYC 9V74, AVX2+FMA, AVX-512 masked off), the pure-Rust AVX2 path
/// is compute-bound at ~0.78 ns/elem while MLAS reaches ~0.35 ns/elem, which is
Expand Down Expand Up @@ -627,24 +630,31 @@ macro_rules! dispatch_mlas {
}};
}

/// MLAS's `tanh`, clamped back into `[-1, 1]`.
///
/// `MlasComputeTanh` is a rational approximation that is not range-preserving:
/// over a dense sweep of [-20, 20] it lands outside `[-1, 1]` for 2934 of
/// 1048576 points, by up to 2 ulp (`tanh(-8.990999) = -1.0000001`). ORT's
/// `Tanh` calls it directly and inherits that. Our pure-Rust kernel has always
/// guaranteed the range, `monotonicity_within_documented_slack` asserts it, and
/// downstream code is entitled to rely on it — `sqrt(1 - tanh²)` turns a 1-ulp
/// overshoot into a NaN.
///
/// Clamping is never *less* accurate: the true `tanh` is inside `(-1, 1)`, so
/// pulling an out-of-range value to the endpoint moves it toward the exact
/// answer. NaN is preserved because all of `clamp`'s comparisons are false for
/// NaN.
#[cfg(feature = "mlas")]
fn tanh_mlas(input: &[f32], output: &mut [f32]) {
mlas_clamped(input, output, mlas_sys::compute_tanh, -1.0, 1.0);
}
// Why `Tanh` and `Sigmoid` do *not* have an MLAS route.
//
// They used to. MLAS's `MlasComputeTanh`/`MlasComputeLogistic` are not
// range-preserving — over a dense sweep of [-20, 20] tanh lands outside
// `[-1, 1]` for 2934 of 1048576 points (by up to 2 ulp) and logistic outside
// `[0, 1]` for 2. Our kernels guarantee the range and
// `monotonicity_within_documented_slack` asserts it, so the MLAS route had to
// re-read every block it had just written and clamp it.
//
// That fix-up pass is now more expensive than the polynomial it was buying.
// The two kernels used the *same* Eigen rational as MLAS even before this,
// and once the redundant saturation blends came out of them (#1121) the
// pure-Rust AVX2 path became outright faster than MLAS-plus-clamp. Measured on
// this host, `mlas`-off vs `mlas`-on p50, 1 thread:
//
// | op | 64 Ki | 256 Ki | 1 Mi | 4 Mi |
// |---------|-------:|--------:|--------:|--------:|
// | Tanh | 32.5 / 36.2 | 105.6 / 146.7 | 438.4 / 469.4 | 1896 / 2226 |
// | Sigmoid | 31.2 / 38.3 | 113.7 / 129.3 | 408.9 / 525.8 | 1887 / 2268 |
//
// The `mlas` route lost at every size — up to 1.39x on `Tanh` at 256 Ki — so
// it is gone and both builds now run the same kernel. `Erf` and exact `Gelu`
// keep theirs: `Erf` needs no fix-up at all and `Gelu`'s is a non-writing
// compare scan, and both still win (1.4-1.6x and 1.17x at 4 Mi).
//

/// Elements per MLAS + fix-up block.
///
Expand All @@ -655,40 +665,6 @@ fn tanh_mlas(input: &[f32], output: &mut [f32]) {
#[cfg(feature = "mlas")]
const MLAS_FIXUP_BLOCK: usize = 8192;

/// Run an MLAS kernel blockwise and clamp each block back into `[lo, hi]`
/// while it is still hot in cache.
#[cfg(feature = "mlas")]
fn mlas_clamped(
input: &[f32],
output: &mut [f32],
kernel: fn(&[f32], &mut [f32]),
lo: f32,
hi: f32,
) {
for (xs, ys) in input
.chunks(MLAS_FIXUP_BLOCK)
.zip(output.chunks_mut(MLAS_FIXUP_BLOCK))
{
kernel(xs, ys);
// `f32::clamp` measured faster here than a hand-written branch-free
// bit-select (0.43 vs 0.55 ns/elem for tanh at 4 Mi) — LLVM already
// lowers it to `vmaxps`/`vminps`, and NaN falls through untouched
// because every comparison against it is false.
for v in ys.iter_mut() {
*v = v.clamp(lo, hi);
}
}
}

/// MLAS's logistic, clamped back into `[0, 1]`.
///
/// Same story as [`tanh_mlas`] but milder: `MlasComputeLogistic` leaves the
/// range for 2 of 1048576 points on the same sweep, by 1 ulp.
#[cfg(feature = "mlas")]
fn sigmoid_mlas(input: &[f32], output: &mut [f32]) {
mlas_clamped(input, output, mlas_sys::compute_logistic, 0.0, 1.0);
}

/// MLAS's exact GELU, with its one special-value divergence repaired.
///
/// `MlasComputeGeluErf` evaluates `x·0.5·(1 + erf(x/√2))` without special-casing
Expand All @@ -712,7 +688,7 @@ fn sigmoid_mlas(input: &[f32], output: &mut [f32]) {
/// vectorisable compare pass on the common all-finite path.
#[cfg(feature = "mlas")]
fn erf_gelu_mlas(input: &[f32], output: &mut [f32]) {
// Blocked for the same cache reason as `mlas_clamped`.
// Blocked so the repair scan reads each block while it is still in L1/L2.
for (xs, ys) in input
.chunks(MLAS_FIXUP_BLOCK)
.zip(output.chunks_mut(MLAS_FIXUP_BLOCK))
Expand All @@ -739,7 +715,7 @@ fn erf_gelu_mlas(input: &[f32], output: &mut [f32]) {

/// `y = tanh(x)`.
pub(crate) fn tanh_f32_slice(input: &[f32], output: &mut [f32]) {
dispatch_mlas!(input, output, tanh_scalar, tanh_avx2, tanh_mlas);
dispatch!(input, output, tanh_scalar, tanh_avx2);
}

/// `y = √x`.
Expand All @@ -763,7 +739,7 @@ pub(crate) fn sqrt_f32_slice(input: &[f32], output: &mut [f32]) {

/// `y = 1 / (1 + e^-x)`.
pub(crate) fn sigmoid_f32_slice(input: &[f32], output: &mut [f32]) {
dispatch_mlas!(input, output, sigmoid_scalar, sigmoid_avx2, sigmoid_mlas);
dispatch!(input, output, sigmoid_scalar, sigmoid_avx2);
}

/// `y = e^x`.
Expand Down Expand Up @@ -3163,8 +3139,8 @@ mod mlas_ab {

#[allow(clippy::type_complexity)]
let cases: [(&str, fn(&[f32], &mut [f32]), fn(&[f32], &mut [f32])); 4] = [
("Tanh", tanh_avx2_shim, tanh_mlas),
("Sigmoid", sigmoid_avx2_shim, sigmoid_mlas),
("Tanh", tanh_avx2_shim, tanh_mlas_ref),
("Sigmoid", sigmoid_avx2_shim, sigmoid_mlas_ref),
("Erf", erf_avx2_shim, mlas_sys::compute_erf),
("GeluErf", erf_gelu_avx2_shim, erf_gelu_mlas),
];
Expand Down Expand Up @@ -3216,8 +3192,8 @@ mod mlas_ab {

#[allow(clippy::type_complexity)]
let cases: [(&str, fn(&[f32], &mut [f32]), fn(&[f32], &mut [f32])); 4] = [
("Tanh", tanh_avx2_shim, tanh_mlas),
("Sigmoid", sigmoid_avx2_shim, sigmoid_mlas),
("Tanh", tanh_avx2_shim, tanh_mlas_ref),
("Sigmoid", sigmoid_avx2_shim, sigmoid_mlas_ref),
("Erf", erf_avx2_shim, mlas_sys::compute_erf),
("GeluErf", erf_gelu_avx2_shim, erf_gelu_mlas),
];
Expand Down Expand Up @@ -3254,8 +3230,8 @@ mod mlas_ab {

#[allow(clippy::type_complexity)]
let cases: [(&str, fn(&[f32], &mut [f32]), fn(&[f32], &mut [f32])); 4] = [
("Tanh", tanh_avx2_shim, tanh_mlas),
("Sigmoid", sigmoid_avx2_shim, sigmoid_mlas),
("Tanh", tanh_avx2_shim, tanh_mlas_ref),
("Sigmoid", sigmoid_avx2_shim, sigmoid_mlas_ref),
("Erf", erf_avx2_shim, mlas_sys::compute_erf),
("GeluErf", erf_gelu_avx2_shim, erf_gelu_mlas),
];
Expand All @@ -3271,6 +3247,40 @@ mod mlas_ab {
}
}

/// What the removed `Tanh` MLAS route did: MLAS's kernel, then the clamp
/// back into `[-1, 1]` that its non-range-preserving rational requires.
///
/// The route is gone from the dispatcher (the clamp pass costs more than
/// the polynomial saves — see the note above `dispatch_mlas!`), but it
/// stays here as the measurement reference `mlas_vs_rust_simd` reports
/// against, and as the thing the special-value and ULP sweeps pin us to.
fn tanh_mlas_ref(x: &[f32], y: &mut [f32]) {
mlas_clamped_ref(x, y, mlas_sys::compute_tanh, -1.0, 1.0);
}

/// Ditto for `Sigmoid` and `[0, 1]`.
fn sigmoid_mlas_ref(x: &[f32], y: &mut [f32]) {
mlas_clamped_ref(x, y, mlas_sys::compute_logistic, 0.0, 1.0);
}

/// Blocked so the clamp re-reads each block while it is still in L1/L2,
/// exactly as the shipped route used to.
fn mlas_clamped_ref(
input: &[f32],
output: &mut [f32],
kernel: fn(&[f32], &mut [f32]),
lo: f32,
hi: f32,
) {
const BLOCK: usize = 8192;
for (xs, ys) in input.chunks(BLOCK).zip(output.chunks_mut(BLOCK)) {
kernel(xs, ys);
for v in ys.iter_mut() {
*v = v.clamp(lo, hi);
}
}
}

// Thin wrappers pinning the pure-Rust route regardless of feature flags.
fn tanh_avx2_shim(x: &[f32], y: &mut [f32]) {
dispatch!(x, y, tanh_scalar, tanh_avx2)
Expand Down
Loading