perf(ep-cpu): route the activation transcendentals through MLAS

crates/mlas-sys/Cargo.toml

@@ -16,6 +16,9 @@ cc = "1"
 
 [dev-dependencies]
 rayon = "1"
+# Correctly-rounded `erf` for the transcendental smoke tests, so they pin the
+# mathematical function rather than MLAS's own polynomial.
+libm = "0.2"
 
 [[bench]]
 # Keep the executable name free of Windows Installer Detection keywords. In

crates/mlas-sys/src/lib.rs

@@ -136,6 +136,15 @@ unsafe extern "C" {
     /// Vectorized fused SiLU over `n` contiguous f32s. MLAS runtime-dispatches
     /// to its one-pass AVX-512F kernel when supported.
     fn mlas_compute_silu(input: *const f32, output: *mut f32, n: usize);
+    /// Vectorized `tanh` over `n` contiguous f32s — the same polynomial ONNX
+    /// Runtime's own `Tanh` CPU kernel calls.
+    fn mlas_compute_tanh(input: *const f32, output: *mut f32, n: usize);
+    /// Vectorized `erf` over `n` contiguous f32s — the same polynomial ONNX
+    /// Runtime's own `Erf` CPU kernel calls.
+    fn mlas_compute_erf(input: *const f32, output: *mut f32, n: usize);
+    /// Vectorized exact (erf-based) GELU over `n` contiguous f32s. Input and
+    /// output must not overlap.
+    fn mlas_compute_gelu_erf(input: *const f32, output: *mut f32, n: usize);
     /// Row-wise softmax over `n` rows of `d` contiguous f32s, single-threaded,
     /// using MLAS's SIMD max reduction and polynomial exp.
     fn mlas_compute_softmax_in_place(data: *mut f32, n: usize, d: usize);
@@ -734,6 +743,65 @@ pub fn compute_silu(input: &[f32], output: &mut [f32]) {
     unsafe { mlas_compute_silu(input.as_ptr(), output.as_mut_ptr(), input.len()) };
 }
 
+/// Compute elementwise `tanh` over equal-length contiguous f32 slices using
+/// MLAS's SIMD polynomial — the same one ONNX Runtime's `Tanh` CPU kernel
+/// calls. MLAS dispatches by ISA at runtime, so which kernel runs (and hence
+/// the exact bits) depends on the host. Single threaded; callers shard
+/// across threads themselves.
+pub fn compute_tanh(input: &[f32], output: &mut [f32]) {
+    assert_eq!(
+        input.len(),
+        output.len(),
+        "compute_tanh input and output must have equal length"
+    );
+    if input.is_empty() {
+        return;
+    }
+    // SAFETY: both slices are valid for `n` contiguous f32s; MLAS reads `input`
+    // and writes `output`, and Rust's borrow rules prove they do not alias.
+    unsafe { mlas_compute_tanh(input.as_ptr(), output.as_mut_ptr(), input.len()) };
+}
+
+/// Compute elementwise `erf` over equal-length contiguous f32 slices using
+/// MLAS's SIMD polynomial — the same one ONNX Runtime's `Erf` CPU kernel
+/// calls. Which kernel runs depends on MLAS's runtime ISA dispatch.
+/// Single threaded; callers shard across threads themselves.
+pub fn compute_erf(input: &[f32], output: &mut [f32]) {
+    assert_eq!(
+        input.len(),
+        output.len(),
+        "compute_erf input and output must have equal length"
+    );
+    if input.is_empty() {
+        return;
+    }
+    // SAFETY: both slices are valid for `n` contiguous f32s; MLAS reads `input`
+    // and writes `output`, and Rust's borrow rules prove they do not alias.
+    unsafe { mlas_compute_erf(input.as_ptr(), output.as_mut_ptr(), input.len()) };
+}
+
+/// Compute elementwise exact GELU `x * 0.5 * (1 + erf(x / sqrt(2)))` over
+/// equal-length contiguous f32 slices, fused in one MLAS pass.
+///
+/// MLAS requires that input and output not overlap (`mlas.h:1166`); the
+/// `&[f32]` / `&mut [f32]` signature makes that unrepresentable.
+///
+/// Single threaded; callers shard across threads themselves.
+pub fn compute_gelu_erf(input: &[f32], output: &mut [f32]) {
+    assert_eq!(
+        input.len(),
+        output.len(),
+        "compute_gelu_erf input and output must have equal length"
+    );
+    if input.is_empty() {
+        return;
+    }
+    // SAFETY: both slices are valid for `n` contiguous f32s; MLAS reads `input`
+    // and writes `output`, and Rust's borrow rules prove they do not alias —
+    // which is also what discharges MLAS's no-overlap precondition.
+    unsafe { mlas_compute_gelu_erf(input.as_ptr(), output.as_mut_ptr(), input.len()) };
+}
+
 /// Row-wise in-place softmax, replacing a scalar `expf` loop: normalizes `n` rows of
 /// `d` contiguous f32s with MLAS's SIMD max reduction and polynomial exp — the
 /// same primitive ONNX Runtime's own Softmax and attention kernels use.

crates/mlas-sys/tests/transcendentals.rs

@@ -0,0 +1,72 @@
+//! Smoke tests pinning the `tanh`/`erf`/`gelu_erf` FFI at the crate boundary.
+//!
+//! The callers in `onnx-runtime-ep-cpu` compare these against a pure-Rust
+//! reference, but that only runs when the `mlas` feature is on there. These
+//! tests fix the shim signatures and the basic contract here, where the
+//! bindings live, so a bad `extern "C"` declaration or a template
+//! instantiation that silently stops linking fails in this crate.
+
+/// Reference values are the correctly-rounded `f64` result of the same
+/// function, so this pins the *function*, not a particular polynomial.
+fn assert_close(actual: &[f32], expect: impl Fn(f64) -> f64, input: &[f32], tol: f64, what: &str) {
+    for (&y, &x) in actual.iter().zip(input) {
+        let want = expect(f64::from(x));
+        let got = f64::from(y);
+        assert!(
+            (got - want).abs() <= tol,
+            "{what}({x}) = {got}, expected ~{want} (tol {tol})"
+        );
+    }
+}
+
+/// Long enough that the vector body runs and leaves a ragged tail.
+fn inputs() -> Vec<f32> {
+    (0..1000).map(|i| (i as f32 - 500.0) / 64.0).collect()
+}
+
+#[test]
+fn compute_tanh_matches_a_f64_reference() {
+    let x = inputs();
+    let mut y = vec![0.0f32; x.len()];
+    mlas_sys::compute_tanh(&x, &mut y);
+    assert_close(&y, f64::tanh, &x, 1e-6, "tanh");
+}
+
+#[test]
+fn compute_erf_matches_a_f64_reference() {
+    let x = inputs();
+    let mut y = vec![0.0f32; x.len()];
+    mlas_sys::compute_erf(&x, &mut y);
+    assert_close(&y, libm::erf, &x, 1e-6, "erf");
+}
+
+#[test]
+fn compute_gelu_erf_matches_a_f64_reference() {
+    let x = inputs();
+    let mut y = vec![0.0f32; x.len()];
+    mlas_sys::compute_gelu_erf(&x, &mut y);
+    assert_close(
+        &y,
+        |v| v * 0.5 * (1.0 + libm::erf(v / std::f64::consts::SQRT_2)),
+        &x,
+        1e-6,
+        "gelu_erf",
+    );
+}
+
+/// Empty input must not call into MLAS with a dangling pointer.
+#[test]
+fn empty_slices_are_a_no_op() {
+    let mut empty: [f32; 0] = [];
+    mlas_sys::compute_tanh(&[], &mut empty);
+    mlas_sys::compute_erf(&[], &mut empty);
+    mlas_sys::compute_gelu_erf(&[], &mut empty);
+}
+
+#[test]
+#[should_panic(expected = "equal length")]
+fn mismatched_lengths_panic_rather_than_read_out_of_bounds() {
+    let x = vec![0.0f32; 8];
+    let mut y = vec![0.0f32; 4];
+    mlas_sys::compute_tanh(&x, &mut y);
+}

crates/mlas-sys/vendor/shim.cpp

@@ -477,6 +477,41 @@ extern "C" void mlas_compute_logistic(
     MlasComputeLogistic(input, output, n);
 }
 
+// ---- Vectorized tanh, erf and exact GELU -----------------------------------
+// The remaining MLAS transcendentals ONNX Runtime's own CPU activation kernels
+// call. Exposing them lets our elementwise kernels share the exact polynomial
+// ORT ships rather than carry a second, slower approximation of the same
+// function -- which also makes our output bit-identical to ORT's for these ops.
+//
+// All three are single-threaded; the caller drives any outer sharding, exactly
+// as `mlas_compute_logistic` above.
+extern "C" void mlas_compute_tanh(
+    const float* input,
+    float* output,
+    size_t n)
+{
+    MlasComputeTanh<float>(input, output, n);
+}
+
+extern "C" void mlas_compute_erf(
+    const float* input,
+    float* output,
+    size_t n)
+{
+    MlasComputeErf(input, output, n);
+}
+
+// NOTE: MLAS documents that Input and Output must not overlap for this one
+// (mlas.h:1166). The Rust wrapper takes `&[f32]` and `&mut [f32]`, so the
+// borrow checker already rules aliasing out.
+extern "C" void mlas_compute_gelu_erf(
+    const float* input,
+    float* output,
+    size_t n)
+{
+    MlasComputeGeluErf(input, output, n);
+}
+
 // ---- Vectorized softmax ----------------------------------------------------
 // `n` independent rows of `d` contiguous floats, normalized in the same way
 // ONNX Runtime's own Softmax/attention kernels do: MLAS finds each row max with