diff --git a/crates/onnx-runtime-ep-cpu-plugin/tests/plugin_ort_e2e.rs b/crates/onnx-runtime-ep-cpu-plugin/tests/plugin_ort_e2e.rs index 25f9e7a4d9..b82e119fec 100644 --- a/crates/onnx-runtime-ep-cpu-plugin/tests/plugin_ort_e2e.rs +++ b/crates/onnx-runtime-ep-cpu-plugin/tests/plugin_ort_e2e.rs @@ -5743,6 +5743,10 @@ fn unary_case( " attribute: [{ name: \"approximate\" type: STRING s: \"tanh\" }]" } "Gelu" => " attribute: [{ name: \"approximate\" type: STRING s: \"none\" }]", + // Celu's alpha divides the input *and* scales the output, so the + // default and a non-default value exercise different arithmetic. The + // name carries which one the case means. + "Celu" if name.contains("_a2") => " attribute: [{ name: \"alpha\" type: FLOAT f: 2.0 }]", _ => "", }; let model = format!( @@ -5870,6 +5874,54 @@ fn unary_bench_cases() -> Vec { false, 0.0, ), + // `Celu` and `Mish` had no kernel here at all before, so ORT ran them + // and there is no "before" arm to compare against: ours-vs-ORT *is* + // the A/B. `Mish` is opset 18, `Celu` is 12. + unary_case( + "bench_celu_f32_4k", + "Celu", + "", + ELEM_F32, + K4, + 17, + false, + 1e-4, + ), + unary_case( + "bench_celu_f32_1m", + "Celu", + "", + ELEM_F32, + M1, + 17, + false, + 1e-4, + ), + unary_case( + "bench_mish_f32_4k", + "Mish", + "", + ELEM_F32, + K4, + 18, + false, + 1e-4, + ), + unary_case( + "bench_mish_f32_1m", + "Mish", + "", + ELEM_F32, + M1, + 18, + false, + 1e-4, + ), + // `Log` was the last elementwise op still evaluated one `libm` call at + // a time, so it gets both sizes: 4k to see the per-call overhead with + // the data L1-resident, 1M to see the steady-state throughput. + unary_case("bench_log_f32_4k", "Log", "", ELEM_F32, K4, 17, true, 1e-4), + unary_case("bench_log_f32_1m", "Log", "", ELEM_F32, M1, 17, true, 1e-4), unary_case( "bench_sqrt_f32_4k", "Sqrt", @@ -6863,3 +6915,199 @@ fn every_assigned_node_is_also_executed_by_this_ep() { skipped.len() ); } + +// ─── Celu, Mish and Log: new native kernels, checked against ORT's ─────────── + +/// The activation buffer a unary parity check has to survive. +/// +/// A sweep of well-behaved values in `[-1, 1]` is the one input set on which +/// every plausible implementation agrees, so it proves almost nothing. The +/// interesting disagreements live at the edges each kernel has to special-case +/// by hand: both signed zeros, both infinities, NaN, the ends of the normal +/// range, the denormal band (which `log` handles by pre-scaling and fixing the +/// exponent afterwards), and the saturation tails where `exp` overflows and the +/// kernels are expected to clamp rather than produce Inf/NaN. +/// +/// The negative half is kept even for `Log`, where it is undefined: "both +/// return NaN" is a real contract, and a kernel that returned a huge negative +/// number instead would pass a positives-only sweep. +fn unary_special_values() -> Vec { + let mut v = vec![ + f32::NEG_INFINITY, + -f32::MAX, + -1e30, + -100.0, + -88.72284, + -88.0, + -20.0, + -10.0, + -2.0, + -1.5, + -1.0, + -0.5, + -std::f32::consts::LN_2, + -f32::MIN_POSITIVE, + -1e-40, + -1e-45, + -0.0, + 0.0, + 1e-45, + 1e-40, + 5.877472e-39, + f32::MIN_POSITIVE, + 0.5, + std::f32::consts::LN_2, + std::f32::consts::FRAC_1_SQRT_2, + 0.99999994, + 1.0, + 1.0000001, + 1.5, + 2.0, + 10.0, + 20.0, + 88.0, + 88.72284, + 100.0, + 1e30, + f32::MAX, + f32::INFINITY, + f32::NAN, + ]; + // A dense two-sided sweep on top of the corner cases, so the comparison is + // not made only of exceptional lanes, and so the length crosses several + // SIMD blocks with a ragged tail. + let n = 421; + for i in 0..n { + v.push(-30.0 + 60.0 * (i as f32) / ((n - 1) as f32)); + } + v +} + +/// Whether two results are compatible, given that neither side is the oracle. +/// +/// Exact equality is required wherever the value is not a rounding of a real +/// number -- NaN, the infinities, and the sign of zero -- because those are +/// contract, not accuracy. Everything else is compared relative to the larger +/// magnitude, which is the right scale for `log`: `ln(1e-38)` is `-87.3`, where +/// one ulp is `7.6e-6`, so an absolute bound would either be vacuous there or +/// unmeetable near `ln(1) = 0`. +fn unary_result_matches(ours: f32, ort: f32, rel: f32) -> bool { + if ours.is_nan() || ort.is_nan() { + return ours.is_nan() && ort.is_nan(); + } + if ours.is_infinite() || ort.is_infinite() { + return ours == ort; + } + if ours == 0.0 && ort == 0.0 { + return true; + } + (ours - ort).abs() <= rel * ours.abs().max(ort.abs()).max(1.0) +} + +/// `Celu`, `Mish` and `Log` run on this EP and agree with ORT's own kernels. +/// +/// These three are the ops this EP was quietly not doing properly. `Celu` and +/// `Mish` had no kernel at all, so the plugin's fail-closed capability filter +/// dropped them and ORT ran them -- which the architectural rule forbids, and +/// which no existing test could see because a missing op cannot appear in a +/// coverage list written from the ops that exist. `Log` was registered but fell +/// into the scalar catch-all in `unary_math`, one `libm` call per element. +/// +/// Every case runs with `session.disable_cpu_ep_fallback=1`, so a node this EP +/// declines does not quietly move to ORT -- session creation fails instead. +/// Assignment is then checked against ORT's own record, the case is run so the +/// claim is backed by an execution, and the output is compared elementwise +/// against a second session with no EP appended at all. +#[test] +fn celu_mish_and_log_execute_locally_and_match_ort_numerics() { + let _lock = lock_ort_ep(); + let x = unary_special_values(); + let len = x.len(); + let bytes: Vec = x.iter().flat_map(|v| v.to_le_bytes()).collect(); + + // `Celu` is opset 12, `Mish` is opset 18; ask for each op's own opset + // rather than one number that happens to satisfy both. + let specs: [(&'static str, &'static str, u64, f32); 4] = [ + ("parity_celu_f32", "Celu", 12, 2e-6), + ("parity_celu_a2_f32", "Celu", 12, 2e-6), + ("parity_mish_f32", "Mish", 18, 2e-6), + ("parity_log_f32", "Log", 17, 2e-6), + ]; + + let mut checked = 0usize; + for (name, op, opset, rel) in specs { + let mut case = unary_case(name, op, "", ELEM_F32, len, opset, false, rel); + case.inputs[0].data = bytes.clone(); + let path = write_generated_model(name, &case.model); + let reg = format!("cpu_ep_parity_{name}"); + let Some((_lib, api, env, _opts, session)) = + (unsafe { conformance_setup(®, &path, true) }) + else { + eprintln!("*** SKIPPED: {name} — ORT or EP cdylib not found ***"); + return; + }; + + unsafe { + let info = query_ep_assignment(api, session); + let ours_ops = info.ops_on_our_ep(); + let theirs = info.ops_not_on_our_ep(); + assert!( + ours_ops.contains(&op), + "{name}: '{op}' must run on this EP, got {:?}", + info.assignments + ); + assert!( + theirs.is_empty(), + "{name}: nodes {theirs:?} were left to ORT's CPU EP. This EP does not \ + defer -- a missing kernel is a kernel to write, not a node to give away", + ); + + let ours_out = run_generated_case(api, session, &case, &format!("Run({name})")); + let (ort_opts, ort_session) = match try_plain_ort_session(api, env, &path) { + Ok(pair) => pair, + Err(e) => panic!("{name}: ORT could not build its own kernel: {e}"), + }; + let ort_out = run_generated_case( + api, + ort_session, + &case, + &format!("Run(ORT baseline {name})"), + ); + assert_eq!(ours_out.len(), ort_out.len(), "{name}: output length"); + + let mut worst = 0.0f32; + let mut worst_at = usize::MAX; + let mut bad: Vec = Vec::new(); + for (i, (&a, &b)) in ours_out.iter().zip(ort_out.iter()).enumerate() { + if !unary_result_matches(a, b, rel) { + if bad.len() < 12 { + bad.push(format!(" x={:e}: ours={a:e} ort={b:e}", x[i])); + } + let d = (a - b).abs() / a.abs().max(b.abs()).max(1.0); + if d > worst { + worst = d; + worst_at = i; + } + } + } + assert!( + bad.is_empty(), + "{name}: {} of {len} elements disagree with ORT (worst {worst:e} at \ + x={:e}, tolerance {rel:e}):\n{}", + bad.len(), + if worst_at == usize::MAX { + f32::NAN + } else { + x[worst_at] + }, + bad.join("\n"), + ); + eprintln!(" [{name}] {len} elements match ORT within {rel:e}"); + ((*api).ReleaseSession.unwrap())(ort_session); + ((*api).ReleaseSessionOptions.unwrap())(ort_opts); + } + checked += 1; + } + assert_eq!(checked, 4, "every parity case must have been checked"); + eprintln!("\n✅ celu_mish_and_log_execute_locally_and_match_ort_numerics: PASSED"); +} diff --git a/crates/onnx-runtime-ep-cpu-plugin/tests/shape_inference_coverage.rs b/crates/onnx-runtime-ep-cpu-plugin/tests/shape_inference_coverage.rs index 9f20cda1b4..716f772f04 100644 --- a/crates/onnx-runtime-ep-cpu-plugin/tests/shape_inference_coverage.rs +++ b/crates/onnx-runtime-ep-cpu-plugin/tests/shape_inference_coverage.rs @@ -264,6 +264,8 @@ fn no_activation_or_norm_op_is_left_to_ort() { ("", "Softsign"), ("", "PRelu"), ("", "ThresholdedRelu"), + ("", "Celu"), + ("", "Mish"), ("", "Softmax"), ("", "LogSoftmax"), ("", "Swish"), @@ -363,12 +365,13 @@ fn activation_and_norm_ops_clear_every_capability_filter() { /// Ops whose whole purpose is an activation or a normalisation. If one of /// these reaches ORT's CPU EP, the EP is not doing its job. /// - /// `Celu` and `Mish` are deliberately absent: this EP has no kernel for - /// either, so ORT running them is a missing feature rather than a decline. - /// They are tracked in `docs/performance/CPU_ACTIVATION_GAPS.md`; add them - /// here the moment a kernel lands, and this test will then hold them to the - /// same standard as the rest. + /// `Celu` and `Mish` used to be listed here as deliberate absences, + /// with a note to add them the moment a kernel landed. It has landed, so + /// they are held to the same standard as the rest: both the shape filter + /// and the dtype filter, which is the pair `PRelu` slipped between. const OWNED: &[(&str, &str)] = &[ + ("", "Celu"), + ("", "Mish"), ("", "Tanh"), ("", "Sigmoid"), ("", "Erf"), diff --git a/crates/onnx-runtime-ep-cpu/src/kernels/activations.rs b/crates/onnx-runtime-ep-cpu/src/kernels/activations.rs index 818b961396..204977ce57 100644 --- a/crates/onnx-runtime-ep-cpu/src/kernels/activations.rs +++ b/crates/onnx-runtime-ep-cpu/src/kernels/activations.rs @@ -18,6 +18,8 @@ enum Activation { ThresholdedRelu { alpha: f32 }, Swish { alpha: f32 }, Silu, + Celu { alpha: f32 }, + Mish, } impl Activation { @@ -30,6 +32,8 @@ impl Activation { Self::ThresholdedRelu { .. } => "ThresholdedRelu", Self::Swish { .. } => "Swish", Self::Silu => "Silu", + Self::Celu { .. } => "Celu", + Self::Mish => "Mish", } } @@ -71,6 +75,13 @@ impl Activation { x * s } Self::Silu => silu(x), + // Shared with the slice path rather than transcribed again: the + // literal ONNX formula loses NaN through `f32::max`/`min`, and one + // copy of that guard is easier to keep true than two. + Self::Celu { alpha } => crate::kernels::simd_activations::celu_scalar(x, alpha), + // Mish(x) = x * tanh(softplus(x)), with softplus in its stable + // form so large `x` neither overflows nor loses the identity. + Self::Mish => x * (x.max(0.0) + (-x.abs()).exp().ln_1p()).tanh(), } } @@ -119,6 +130,18 @@ impl Activation { x * s } Self::Silu => silu_f64(x), + Self::Celu { alpha } => { + // See `simd_activations::celu_scalar`: `max`/`min` return the + // other operand for NaN, so without this the f64 path -- which + // every `Float64` tensor takes -- would answer `0.0` where ORT + // and the f32 paths answer NaN. + if x.is_nan() { + return x; + } + let a = f64::from(alpha); + x.max(0.0) + (a * ((x / a).exp() - 1.0)).min(0.0) + } + Self::Mish => x * (x.max(0.0) + (-x.abs()).exp().ln_1p()).tanh(), } } } @@ -169,6 +192,8 @@ pub struct SeluFactory; pub struct ThresholdedReluFactory; pub struct SwishFactory; pub struct SiluFactory; +pub struct CeluFactory; +pub struct MishFactory; impl KernelFactory for EluFactory { fn create(&self, node: &Node, _shapes: &[Vec]) -> Result> { @@ -245,6 +270,27 @@ impl KernelFactory for SwishFactory { } } +impl KernelFactory for CeluFactory { + fn create(&self, node: &Node, _shapes: &[Vec]) -> Result> { + let alpha = node.attr("alpha").and_then(|a| a.as_float()).unwrap_or(1.0); + // ONNX requires alpha != 0 (the definition divides by it). A model that + // ships 0 would otherwise reach a division by zero in the kernel, so + // fall back to the documented default rather than produce Inf/NaN. + let alpha = if alpha == 0.0 { 1.0 } else { alpha }; + Ok(Box::new(ActivationKernel { + activation: Activation::Celu { alpha }, + })) + } +} + +impl KernelFactory for MishFactory { + fn create(&self, _node: &Node, _shapes: &[Vec]) -> Result> { + Ok(Box::new(ActivationKernel { + activation: Activation::Mish, + })) + } +} + impl KernelFactory for SiluFactory { fn create(&self, _node: &Node, _shapes: &[Vec]) -> Result> { Ok(Box::new(ActivationKernel { @@ -273,6 +319,14 @@ impl Kernel for ActivationKernel { let mut output = vec![0.0; input.len()]; silu_f32_slice(&input, &mut output); output + } else if let Activation::Celu { alpha } = self.activation { + let mut output = vec![0.0; input.len()]; + crate::kernels::simd_activations::celu_f32_slice(&input, &mut output, alpha); + output + } else if matches!(self.activation, Activation::Mish) { + let mut output = vec![0.0; input.len()]; + crate::kernels::simd_activations::mish_f32_slice(&input, &mut output); + output } else { input .iter() @@ -474,6 +528,27 @@ mod tests { use crate::kernels::testutil::Owned; use onnx_runtime_ir::{Attribute, NodeId}; + /// `Float64` tensors never touch the SIMD kernels, so `apply_f64` is the + /// only implementation a `double` model sees. It has to agree with the + /// other two about NaN -- `f64::max`/`min` drop it exactly as the f32 ones + /// do, and ORT propagates it. + #[test] + fn celu_and_mish_propagate_nan_on_every_path() { + for alpha in [0.5f32, 1.0, 3.0] { + let act = Activation::Celu { alpha }; + assert!(act.apply(f32::NAN).is_nan(), "f32 Celu(NaN), alpha={alpha}"); + assert!( + act.apply_f64(f64::NAN).is_nan(), + "f64 Celu(NaN), alpha={alpha}" + ); + } + assert!(Activation::Mish.apply(f32::NAN).is_nan(), "f32 Mish(NaN)"); + assert!( + Activation::Mish.apply_f64(f64::NAN).is_nan(), + "f64 Mish(NaN)" + ); + } + #[test] fn activation_formulas_and_defaults() { let x = Owned::f32(&[3], &[-1.0, 0.0, 1.0]); diff --git a/crates/onnx-runtime-ep-cpu/src/kernels/mod.rs b/crates/onnx-runtime-ep-cpu/src/kernels/mod.rs index b527a57bc1..d0882c36db 100644 --- a/crates/onnx-runtime-ep-cpu/src/kernels/mod.rs +++ b/crates/onnx-runtime-ep-cpu/src/kernels/mod.rs @@ -345,7 +345,9 @@ pub fn supported_dtypes_for_op(op_type: &str, domain: &str) -> &'static [DataTyp | ("LeakyRelu", "") | ("HardSigmoid", "") | ("Selu", "") - | ("ThresholdedRelu", "") => FLOAT_DTYPES, + | ("ThresholdedRelu", "") + | ("Celu", "") + | ("Mish", "") => FLOAT_DTYPES, // Softmax, LogSoftmax, ReduceMean, LayerNorm, etc.: float-only. ("Softmax", "") @@ -662,6 +664,8 @@ pub const PHASE1_OPS: &[&str] = &[ "Elu", "LeakyRelu", "HardSigmoid", + "Celu", + "Mish", // Logical / selection. "And", "Or", @@ -1341,6 +1345,16 @@ fn build_cpu_registry_recorded_inner( rec.register(OpKey::new("Sinh", "", 1), Box::new(unary_math::SinhFactory)); rec.register(OpKey::new("Tan", "", 1), Box::new(unary_math::TanFactory)); rec.register(OpKey::new("Elu", "", 1), Box::new(activations::EluFactory)); + // Celu is opset 12, Mish opset 18 -- registering at their introducing + // opset keeps older models routing to whatever handled them before. + rec.register( + OpKey::new("Celu", "", 12), + Box::new(activations::CeluFactory), + ); + rec.register( + OpKey::new("Mish", "", 18), + Box::new(activations::MishFactory), + ); rec.register( OpKey::new("LeakyRelu", "", 1), Box::new(activations::LeakyReluFactory), diff --git a/crates/onnx-runtime-ep-cpu/src/kernels/simd_activations.rs b/crates/onnx-runtime-ep-cpu/src/kernels/simd_activations.rs index 0a8a454db4..732b60d696 100644 --- a/crates/onnx-runtime-ep-cpu/src/kernels/simd_activations.rs +++ b/crates/onnx-runtime-ep-cpu/src/kernels/simd_activations.rs @@ -792,6 +792,37 @@ mod erf_c { pub(super) const EXP_C: f32 = 1.25829120e7; } +#[cfg(target_arch = "x86_64")] +/// Cephes `logf` constants. +/// +/// These are the same coefficients Eigen's `plog` carries, which is what ORT's +/// CPU `Log` ultimately evaluates, so matching them is what buys us parity +/// rather than merely accuracy. +#[cfg(target_arch = "x86_64")] +mod log_c { + /// `sqrt(0.5)`. The mantissa is folded into `[sqrt(0.5), sqrt(2))` around + /// this point so the polynomial argument `f = m - 1` stays small and + /// symmetric about zero. + pub(super) const SQRT_HALF: f32 = core::f32::consts::FRAC_1_SQRT_2; + pub(super) const P0: f32 = 7.037_683_6E-2; + pub(super) const P1: f32 = -1.151_461_0E-1; + pub(super) const P2: f32 = 1.167_699_9E-1; + pub(super) const P3: f32 = -1.242_014_1E-1; + pub(super) const P4: f32 = 1.424_932_3E-1; + pub(super) const P5: f32 = -1.666_805_8E-1; + pub(super) const P6: f32 = 2.000_071_5E-1; + pub(super) const P7: f32 = -2.499_999_4E-1; + pub(super) const P8: f32 = 3.333_333_1E-1; + /// `ln 2` split so that `e * LN2_HI` is exact for every `e` a `f32` + /// exponent can hold; `LN2_LO` carries the remainder. + pub(super) const LN2_HI: f32 = 0.693_359_38; + pub(super) const LN2_LO: f32 = -2.121_944_4E-4; + /// Denormal inputs have no usable exponent field, so they are scaled by + /// `2^25` first and `25` is subtracted from the recovered exponent. + pub(super) const DENORM_SCALE: f32 = 33_554_432.0; + pub(super) const DENORM_EXP: f32 = 25.0; +} + /// `MlasExpConstants` (`onnxruntime/core/mlas/lib/compute.cpp`), the parameter /// set behind `MlasComputeExp`. /// @@ -1121,6 +1152,41 @@ pub(crate) fn tanh_gelu_f32_slice(input: &[f32], output: &mut [f32]) { dispatch!(input, output, tanh_gelu_scalar, tanh_gelu_avx2); } +/// `y = ln(x)`. +/// +/// Vectorised rather than left to `f32::ln`: `Log` was the last unary op in +/// this module still taking a `libm` call per element. +pub(crate) fn log_f32_slice(input: &[f32], output: &mut [f32]) { + dispatch!(input, output, log_scalar, log_avx2); +} + +/// `y = x·tanh(softplus(x))`, `Mish`. +pub(crate) fn mish_f32_slice(input: &[f32], output: &mut [f32]) { + dispatch!(input, output, mish_scalar, mish_avx2); +} + +/// `y = max(0,x) + min(0, alpha·(exp(x/alpha) − 1))`, `Celu`. +/// +/// `alpha` must be non-zero; ONNX requires it and the reciprocal below would +/// otherwise be infinite. +pub(crate) fn celu_f32_slice(input: &[f32], output: &mut [f32], alpha: f32) { + debug_assert_eq!(input.len(), output.len()); + debug_assert!(alpha != 0.0); + #[cfg(target_arch = "x86_64")] + { + if input.len() >= SIMD_MIN_LEN && vector_path_available() { + // SAFETY: guarded by the runtime AVX2+FMA detection above. + run_chunked(input, output, |i, o| unsafe { celu_avx2(i, o, alpha) }); + return; + } + } + run_chunked(input, output, |i, o| { + for (o, &i) in o.iter_mut().zip(i) { + *o = celu_scalar(i, alpha); + } + }); +} + /// `y = x·sigmoid(alpha·x)`, the `QuickGelu` / Swish form. pub(crate) fn quick_gelu_f32_slice(input: &[f32], output: &mut [f32], alpha: f32) { debug_assert_eq!(input.len(), output.len()); @@ -1334,6 +1400,40 @@ pub(crate) fn softsign_scalar(x: f32) -> f32 { // --------------------------------------------------------------------------- #[inline] +/// `ln(x)`, the scalar reference for [`log_f32_slice`]. +fn log_scalar(x: f32) -> f32 { + x.ln() +} + +/// `Celu(x) = max(0, x) + min(0, alpha * (exp(x / alpha) - 1))`. +/// +/// This is ONNX's definition evaluated literally, which is also what ORT does. +/// `alpha` is required to be non-zero by the spec. +pub(crate) fn celu_scalar(x: f32, alpha: f32) -> f32 { + // Transcribing the formula literally is not enough. Rust's `f32::max` and + // `f32::min` are IEEE-754 `maxNum`/`minNum`: they return the *other* + // operand when one is NaN, so both terms collapse to zero and `Celu(NaN)` + // comes out `0.0`. ORT propagates NaN, and so does `celu_ps` (whose + // `maxps` keeps it) -- without this guard the answer would depend on + // whether the slice was long enough for the vector path and on whether the + // host has AVX2, which is exactly what `dispatch!` exists to prevent. + if x.is_nan() { + return x; + } + x.max(0.0) + (alpha * ((x / alpha).exp() - 1.0)).min(0.0) +} + +/// `Softplus(x) = ln(1 + e^x)`, in the stable form +/// `max(x, 0) + ln(1 + e^-|x|)`. +fn softplus_scalar(x: f32) -> f32 { + x.max(0.0) + (-x.abs()).exp().ln_1p() +} + +/// `Mish(x) = x * tanh(softplus(x))`. +fn mish_scalar(x: f32) -> f32 { + x * softplus_scalar(x).tanh() +} + fn tanh_scalar(x: f32) -> f32 { x.tanh() } @@ -1392,7 +1492,7 @@ fn quick_gelu_scalar(x: f32, alpha: f32) -> f32 { #[cfg(target_arch = "x86_64")] mod avx2 { - use super::{GELU_B, GELU_C, erf_c, exp_c, logistic_c, tanh_c}; + use super::{GELU_B, GELU_C, erf_c, exp_c, log_c, logistic_c, tanh_c}; use core::arch::x86_64::*; /// `[-1; 7] ++ [0; 8]`. Loading 8 lanes at offset `7 - rem` yields a mask @@ -1737,6 +1837,153 @@ mod avx2 { p } + /// `ln(x)` for eight lanes, the Cephes decomposition. + /// + /// `x = m * 2^e` with `m` folded into `[sqrt(0.5), sqrt(2))`, then + /// `ln x = ln m + e ln 2`, with `ln m` from a degree-8 polynomial in + /// `f = m - 1`. `ln 2` is split `HI + LO` so `e * HI` is exact. + /// + /// Denormals are handled rather than flushed: they carry no usable exponent + /// field, so they are pre-scaled by `2^25` and `25` is taken back off the + /// recovered exponent. That keeps the whole subnormal range accurate + /// instead of returning `-Inf` for it. + /// + /// Every IEEE special is restored by a final blend, because the bit + /// surgery in the middle destroys them: + /// + /// | input | result | + /// |---|---| + /// | `+Inf` | `+Inf` | + /// | `±0` | `-Inf` | + /// | `x < 0` | `NaN` | + /// | `NaN` | that same `NaN`, payload intact | + #[inline] + #[target_feature(enable = "avx2,fma")] + pub(super) unsafe fn log_ps(x: __m256) -> __m256 { + let zero = _mm256_setzero_ps(); + + // Denormals (and only denormals: the compare is ordered, so NaN + // and negatives answer false, and they are overwritten at the end + // anyway) get scaled into the normal range first. + let is_denorm = _mm256_and_ps( + _mm256_cmp_ps(x, _mm256_set1_ps(f32::MIN_POSITIVE), _CMP_LT_OQ), + _mm256_cmp_ps(x, zero, _CMP_GT_OQ), + ); + let scale = _mm256_blendv_ps( + _mm256_set1_ps(1.0), + _mm256_set1_ps(log_c::DENORM_SCALE), + is_denorm, + ); + let scaled = _mm256_mul_ps(x, scale); + let exp_fix = _mm256_and_ps(is_denorm, _mm256_set1_ps(log_c::DENORM_EXP)); + + let bits = _mm256_castps_si256(scaled); + // 126, not 127: the mantissa below is forced into [0.5, 1) rather than + // [1, 2), so `x == m * 2^(field - 126)`. Using the usual 127 bias here + // costs exactly one `ln 2` and is invisible near `x = 1`, which is why + // it has to be caught at the top of the range (`ln(f32::MAX)`). + let raw_exp = _mm256_sub_epi32(_mm256_srli_epi32::<23>(bits), _mm256_set1_epi32(126)); + let mut e = _mm256_sub_ps(_mm256_cvtepi32_ps(raw_exp), exp_fix); + + // Mantissa into [0.5, 1): keep the fraction field, force the + // exponent field to 126. + let m = _mm256_or_ps( + _mm256_and_ps( + scaled, + _mm256_castsi256_ps(_mm256_set1_epi32(0x807f_ffffu32 as i32)), + ), + _mm256_set1_ps(0.5), + ); + + // Fold to [sqrt(0.5), sqrt(2)): below the split point, double the + // mantissa and borrow one from the exponent. + let lt = _mm256_cmp_ps(m, _mm256_set1_ps(log_c::SQRT_HALF), _CMP_LT_OQ); + e = _mm256_sub_ps(e, _mm256_and_ps(lt, _mm256_set1_ps(1.0))); + let f = _mm256_sub_ps(_mm256_add_ps(m, _mm256_and_ps(lt, m)), _mm256_set1_ps(1.0)); + + let z = _mm256_mul_ps(f, f); + let mut poly = _mm256_set1_ps(log_c::P0); + poly = _mm256_fmadd_ps(poly, f, _mm256_set1_ps(log_c::P1)); + poly = _mm256_fmadd_ps(poly, f, _mm256_set1_ps(log_c::P2)); + poly = _mm256_fmadd_ps(poly, f, _mm256_set1_ps(log_c::P3)); + poly = _mm256_fmadd_ps(poly, f, _mm256_set1_ps(log_c::P4)); + poly = _mm256_fmadd_ps(poly, f, _mm256_set1_ps(log_c::P5)); + poly = _mm256_fmadd_ps(poly, f, _mm256_set1_ps(log_c::P6)); + poly = _mm256_fmadd_ps(poly, f, _mm256_set1_ps(log_c::P7)); + poly = _mm256_fmadd_ps(poly, f, _mm256_set1_ps(log_c::P8)); + poly = _mm256_mul_ps(poly, _mm256_mul_ps(f, z)); + + poly = _mm256_fmadd_ps(e, _mm256_set1_ps(log_c::LN2_LO), poly); + poly = _mm256_fnmadd_ps(z, _mm256_set1_ps(0.5), poly); + let mut r = _mm256_add_ps(f, poly); + r = _mm256_fmadd_ps(e, _mm256_set1_ps(log_c::LN2_HI), r); + + // Specials, innermost first. + let nan = _mm256_set1_ps(f32::NAN); + r = _mm256_blendv_ps( + r, + _mm256_set1_ps(f32::INFINITY), + _mm256_cmp_ps(x, _mm256_set1_ps(f32::INFINITY), _CMP_EQ_OQ), + ); + // `-0.0 < 0.0` is false in IEEE, so both zeros land here and both + // give `-Inf`, which is what `f32::ln` does. + r = _mm256_blendv_ps( + r, + _mm256_set1_ps(f32::NEG_INFINITY), + _mm256_cmp_ps(x, zero, _CMP_EQ_OQ), + ); + r = _mm256_blendv_ps(r, nan, _mm256_cmp_ps(x, zero, _CMP_LT_OQ)); + // Pass the original through so the payload survives. + _mm256_blendv_ps(r, x, _mm256_cmp_ps(x, x, _CMP_UNORD_Q)) + } + + /// `Celu(x) = max(0, x) + min(0, alpha * (exp(x / alpha) - 1))`. + /// + /// Evaluated exactly as ONNX writes it. `alpha` is reciprocated once by the + /// caller so the per-element cost is a multiply rather than a divide. + #[inline] + #[target_feature(enable = "avx2,fma")] + pub(super) unsafe fn celu_ps(x: __m256, alpha: __m256, inv_alpha: __m256) -> __m256 { + unsafe { + let zero = _mm256_setzero_ps(); + let e = exp_full_ps(_mm256_mul_ps(x, inv_alpha)); + let neg = _mm256_mul_ps(alpha, _mm256_sub_ps(e, _mm256_set1_ps(1.0))); + // Operand order is load-bearing: `maxps a, b` returns `b` when + // either operand is NaN, so `max(x, zero)` would quietly turn NaN + // into 0. Passing `x` second propagates it, matching the scalar + // reference. Same argument for `minps` on the negative half. + _mm256_add_ps(_mm256_max_ps(zero, x), _mm256_min_ps(zero, neg)) + } + } + + /// `Softplus(x) = ln(1 + e^x)` in the stable form `max(x,0) + ln(1+e^-|x|)`. + /// + /// The `-|x|` argument means the exponential never overflows, so the only + /// rounding that matters is `1 + e^-|x|` losing the tail for large `|x|` — + /// where the result is dominated by `max(x,0)` anyway. + #[inline] + #[target_feature(enable = "avx2,fma")] + pub(super) unsafe fn softplus_ps(x: __m256) -> __m256 { + unsafe { + let zero = _mm256_setzero_ps(); + let abs = _mm256_andnot_ps(_mm256_set1_ps(-0.0), x); + let e = exp_full_ps(_mm256_sub_ps(zero, abs)); + // `max(zero, x)`, not `max(x, zero)`: see `celu_ps`. Here it is + // NaN propagation that depends on it. + _mm256_add_ps( + _mm256_max_ps(zero, x), + log_ps(_mm256_add_ps(_mm256_set1_ps(1.0), e)), + ) + } + } + + /// `Mish(x) = x * tanh(softplus(x))`. + #[inline] + #[target_feature(enable = "avx2,fma")] + pub(super) unsafe fn mish_ps(x: __m256) -> __m256 { + unsafe { _mm256_mul_ps(x, tanh_ps(softplus_ps(x))) } + } + /// `2^k` for integer-valued `k`, built by biasing and shifting into the /// exponent field (`MlasPowerOf2Float32x4`). #[inline] @@ -1988,6 +2235,28 @@ unsafe fn quick_gelu_avx2(input: &[f32], output: &mut [f32], alpha: f32) { } } +#[cfg(target_arch = "x86_64")] +#[target_feature(enable = "avx2,fma")] +unsafe fn log_avx2(input: &[f32], output: &mut [f32]) { + unsafe { avx2::map_ps(input, output, |v| avx2::log_ps(v)) } +} + +#[cfg(target_arch = "x86_64")] +#[target_feature(enable = "avx2,fma")] +unsafe fn mish_avx2(input: &[f32], output: &mut [f32]) { + unsafe { avx2::map_ps(input, output, |v| avx2::mish_ps(v)) } +} + +#[cfg(target_arch = "x86_64")] +#[target_feature(enable = "avx2,fma")] +unsafe fn celu_avx2(input: &[f32], output: &mut [f32], alpha: f32) { + unsafe { + let a = core::arch::x86_64::_mm256_set1_ps(alpha); + let inv = core::arch::x86_64::_mm256_set1_ps(1.0 / alpha); + avx2::map_ps(input, output, |v| avx2::celu_ps(v, a, inv)) + } +} + #[cfg(test)] mod tests { use super::*; @@ -2334,6 +2603,228 @@ mod tests { } } + /// `log`'s error is relative to the *result*, not the input: `ln(1e-38)` + /// is `-87.3`, where one ulp is `7.6e-6` and the shared `scaled_err` + /// (which divides by `max(|x|,1)`) would call that a huge error. Divide by + /// the reference magnitude instead, flooring at 1 so the region around + /// `x = 1` — where `ln` legitimately passes through zero — stays an + /// absolute test. + fn log_err(got: f32, want: f32) -> f64 { + if (got.is_nan() && want.is_nan()) || got == want { + return 0.0; + } + (f64::from(got) - f64::from(want)).abs() / f64::from(want).abs().max(1.0) + } + + const LOG_BOUND: f64 = 2e-7; + + #[test] + fn log_special_values() { + let mut x = special_inputs(); + x.extend_from_slice(&[1.0, std::f32::consts::E, 0.5]); + let mut o = vec![0.0f32; x.len()]; + log_f32_slice(&x, &mut o); + assert!(o[PAD].is_nan(), "log(-Inf) is NaN"); + assert_eq!(o[PAD + 1], f32::NEG_INFINITY, "log(-0) is -Inf"); + assert_eq!(o[PAD + 2], f32::NEG_INFINITY, "log(+0) is -Inf"); + assert_eq!(o[PAD + 3], f32::INFINITY, "log(+Inf) is +Inf"); + assert!(o[PAD + 4].is_nan(), "log(NaN)"); + assert!(o[PAD + 5].is_nan(), "log(-MAX) is NaN"); + assert!(log_err(o[PAD + 6], f32::MAX.ln()) <= LOG_BOUND); + assert!(o[PAD + 7].is_nan(), "log(-1e30) is NaN"); + assert!(log_err(o[PAD + 8], 1e30f32.ln()) <= LOG_BOUND); + assert_eq!(o[x.len() - 3], 0.0, "log(1) is exactly 0"); + assert!(log_err(o[x.len() - 2], 1.0) <= LOG_BOUND, "log(e) is 1"); + assert!(log_err(o[x.len() - 1], 0.5f32.ln()) <= LOG_BOUND); + } + + /// The vector path must not flush subnormals: it scales them into the + /// normal range instead, so the whole subnormal decade stays accurate. + #[test] + fn log_handles_subnormals() { + let mut x = vec![0.0f32; PAD]; + x.extend_from_slice(&[ + f32::from_bits(1), + f32::from_bits(2), + f32::from_bits(0x0000_0f00), + f32::from_bits(0x007f_ffff), + f32::MIN_POSITIVE, + ]); + let mut o = vec![0.0f32; x.len()]; + log_f32_slice(&x, &mut o); + for (&xi, &got) in x[PAD..].iter().zip(&o[PAD..]) { + let want = xi.ln(); + assert!( + got.is_finite(), + "log({xi:e}) returned {got}, so a subnormal was flushed" + ); + let e = log_err(got, want); + assert!(e <= LOG_BOUND, "log({xi:e}): {got} vs {want}, err {e:e}"); + } + } + + #[test] + fn log_dense_sweep_matches_scalar() { + // Spread over the whole positive exponent range, not just [0,1]: the + // exponent reconstruction is the part most likely to be wrong. + let mut x: Vec = (0..200_003) + .map(|i| { + let t = i as f32 / 200_002.0; + (t * 176.0 - 88.0).exp2() + }) + .collect(); + x.extend(grid(1e-6, 20.0, 100_003, &[])); + x.extend(grid(0.9, 1.1, 50_003, &[])); + let mut o = vec![0.0f32; x.len()]; + log_f32_slice(&x, &mut o); + let mut worst = 0.0f64; + for (&xi, &got) in x.iter().zip(&o) { + let e = log_err(got, xi.ln()); + if e > worst { + worst = e; + } + assert!(e <= LOG_BOUND, "log({xi:e}) = {got}, want {}", xi.ln()); + } + assert!(worst > 0.0, "sweep did not exercise anything"); + } + + #[test] + fn celu_special_values() { + let mut x = special_inputs(); + x.extend_from_slice(&[1.0, -1.0, 0.25]); + let mut o = vec![0.0f32; x.len()]; + celu_f32_slice(&x, &mut o, 1.0); + // Celu(-Inf) = 0 + min(0, 1*(exp(-Inf)-1)) = -1. + assert_eq!(o[PAD], -1.0, "celu(-Inf)"); + // Celu(x) = max(0, x) + min(0, alpha * (exp(x / alpha) - 1)). At x = -0 + // both terms are zero, and IEEE-754 round-to-nearest addition of two + // zeros yields -0 only when *both* are -0. Here the right-hand term is + // exp(-0) - 1 = +0, so the sign is destroyed by the formula itself -- + // for the scalar reference and for ORT exactly as for us. Pin both. + assert_eq!(o[PAD + 1].to_bits(), 0, "celu(-0) is +0"); + assert_eq!(celu_scalar(-0.0, 1.0).to_bits(), 0, "scalar celu(-0) is +0"); + assert_eq!(o[PAD + 2], 0.0, "celu(+0)"); + assert_eq!(o[PAD + 3], f32::INFINITY, "celu(+Inf)"); + assert!(o[PAD + 4].is_nan(), "celu(NaN)"); + assert_eq!(o[PAD + 5], -1.0, "celu(-MAX) saturates to -alpha"); + assert_eq!(o[PAD + 6], f32::MAX, "celu(MAX)"); + for (&xi, &got) in x[x.len() - 3..].iter().zip(&o[x.len() - 3..]) { + let want = celu_scalar(xi, 1.0); + assert!( + (f64::from(got) - f64::from(want)).abs() <= 4e-7, + "celu({xi}) = {got}, want {want}" + ); + } + } + + /// The scalar fallback must answer exactly what the vector path answers. + /// + /// `special_inputs()` pads to `SIMD_MIN_LEN`, so every other special-value + /// test in this file takes the vector path -- which is how `celu_scalar` + /// silently turning NaN into `0.0` survived being written. A slice shorter + /// than `SIMD_MIN_LEN` is the only way to reach the scalar branch from + /// outside, and it is reached in production by small tensors and by any + /// host without AVX2. Which one runs must not be observable. + #[test] + fn celu_scalar_path_agrees_with_the_vector_path() { + let specials = [ + f32::NAN, + f32::NEG_INFINITY, + f32::INFINITY, + -0.0, + 0.0, + -1.0, + 1.0, + -30.0, + 0.25, + ]; + assert!( + specials.len() < SIMD_MIN_LEN, + "the short slice must be short enough to take the scalar path" + ); + for alpha in [0.25f32, 1.0, 2.0, 7.5] { + let mut short_out = vec![0.0f32; specials.len()]; + celu_f32_slice(&specials, &mut short_out, alpha); + + let mut long = vec![0.0f32; PAD]; + long.extend_from_slice(&specials); + let mut long_out = vec![0.0f32; long.len()]; + celu_f32_slice(&long, &mut long_out, alpha); + + for (i, &x) in specials.iter().enumerate() { + let scalar = short_out[i]; + let vector = long_out[PAD + i]; + assert_eq!( + scalar.is_nan(), + vector.is_nan(), + "celu({x}, alpha={alpha}): scalar {scalar}, vector {vector}" + ); + if !scalar.is_nan() { + assert_eq!( + scalar.to_bits(), + vector.to_bits(), + "celu({x}, alpha={alpha}): scalar {scalar}, vector {vector}" + ); + } + } + } + } + + #[test] + fn celu_matches_scalar_across_alphas() { + let x = grid(-30.0, 30.0, 20_003, &[]); + let mut o = vec![0.0f32; x.len()]; + for alpha in [0.25f32, 0.5, 1.0, 2.0, 7.5] { + celu_f32_slice(&x, &mut o, alpha); + for (&xi, &got) in x.iter().zip(&o) { + let want = celu_scalar(xi, alpha); + let e = (f64::from(got) - f64::from(want)).abs() / f64::from(want).abs().max(1.0); + // The negative half is alpha * (exp(u) - 1), so alpha directly + // multiplies the rounding error of exp: at alpha = 7.5 a + // half-ulp slip in exp shows up as ~7.5 half-ulps here. The + // bound therefore has to scale with alpha rather than pretend + // the amplification is not real. + let bound = 4e-7 * f64::from(alpha).max(1.0); + assert!( + e <= bound, + "celu({xi}, alpha={alpha}) = {got}, want {want}, err {e:e}" + ); + } + } + } + + #[test] + fn mish_special_values() { + let mut x = special_inputs(); + x.extend_from_slice(&[1.0, -1.0, 0.0]); + let mut o = vec![0.0f32; x.len()]; + mish_f32_slice(&x, &mut o); + // softplus(-Inf) = 0, tanh(0) = 0, so x*0 with x = -Inf is NaN -- + // which is what the scalar reference produces too. + assert_eq!( + o[PAD].is_nan(), + mish_scalar(f32::NEG_INFINITY).is_nan(), + "mish(-Inf) must agree with the scalar reference" + ); + assert_eq!(o[PAD + 1].to_bits(), (-0.0f32).to_bits(), "mish(-0) is -0"); + assert_eq!(o[PAD + 2], 0.0, "mish(+0)"); + assert_eq!(o[PAD + 3], f32::INFINITY, "mish(+Inf)"); + assert!(o[PAD + 4].is_nan(), "mish(NaN)"); + assert_eq!(o[PAD + 6], f32::MAX, "mish(MAX) = MAX"); + } + + #[test] + fn mish_dense_sweep_matches_scalar() { + let x = grid(-30.0, 30.0, 200_003, &[]); + let mut o = vec![0.0f32; x.len()]; + mish_f32_slice(&x, &mut o); + for (&xi, &got) in x.iter().zip(&o) { + let want = mish_scalar(xi); + let e = (f64::from(got) - f64::from(want)).abs() / f64::from(xi).abs().max(1.0); + assert!(e <= 6e-7, "mish({xi}) = {got}, want {want}, err {e:e}"); + } + } + #[test] fn sqrt_special_values() { let mut x = special_inputs(); diff --git a/crates/onnx-runtime-ep-cpu/src/kernels/unary_math.rs b/crates/onnx-runtime-ep-cpu/src/kernels/unary_math.rs index 635ca2bb8a..82e00c4a9e 100644 --- a/crates/onnx-runtime-ep-cpu/src/kernels/unary_math.rs +++ b/crates/onnx-runtime-ep-cpu/src/kernels/unary_math.rs @@ -9,7 +9,8 @@ //! and narrows back to the requested output float type: //! //! * `Abs`, `Neg`, `Reciprocal` — trivial arithmetic. -//! * `Exp`, `Log`, `Sin`, `Cos` — `std` intrinsics. +//! * `Exp`, `Log` — vectorised on AVX2+FMA (see [`simd_activations`]). +//! * `Sin`, `Cos` — `std` intrinsics. //! * `Floor`, `Ceil` — `std` intrinsics. //! * `Round` — ONNX "round half to even" (banker's rounding), not Rust's //! round-half-away-from-zero `f32::round`; uses [`f32::round_ties_even`]. @@ -197,6 +198,7 @@ impl UnaryMathKernel { MathOp::Round => simd_activations::round_ties_even_f32_slice(x, y), MathOp::Sign => simd_activations::sign_f32_slice(x, y), MathOp::Softsign => simd_activations::softsign_f32_slice(x, y), + MathOp::Log => simd_activations::log_f32_slice(x, y), op => { for (o, &v) in y.iter_mut().zip(x.iter()) { *o = op.apply(v); diff --git a/crates/onnx-runtime-ep-cpu/src/provider.rs b/crates/onnx-runtime-ep-cpu/src/provider.rs index 7f087d8ee4..be2db049a0 100644 --- a/crates/onnx-runtime-ep-cpu/src/provider.rs +++ b/crates/onnx-runtime-ep-cpu/src/provider.rs @@ -1141,7 +1141,7 @@ mod tests { node.attributes .insert("direction".into(), Attribute::String(b"RIGHT".to_vec())); } - assert!(ep.get_kernel(&node, &[], 17).is_ok(), "no kernel for {op}"); + assert!(ep.get_kernel(&node, &[], 18).is_ok(), "no kernel for {op}"); } let bad = Node::new(onnx_runtime_ir::NodeId(99), "UnknownOp", vec![], vec![]); assert!(ep.get_kernel(&bad, &[], 17).is_err()); diff --git a/crates/onnx-runtime-ep-plugin/src/compute.rs b/crates/onnx-runtime-ep-plugin/src/compute.rs index 348a03eeed..de4711e9c3 100644 --- a/crates/onnx-runtime-ep-plugin/src/compute.rs +++ b/crates/onnx-runtime-ep-plugin/src/compute.rs @@ -229,7 +229,9 @@ impl ShapeInference { | "HardSwish" | "LeakyRelu" | "Elu" + | "Celu" | "Selu" + | "Mish" | "Softplus" | "Softsign" // Trigonometric and hyperbolic ops. Unary and shape-preserving, @@ -405,7 +407,9 @@ impl ShapeInference { | "HardSwish" | "LeakyRelu" | "Elu" + | "Celu" | "Selu" + | "Mish" | "Softplus" | "Softsign" // Trigonometric and hyperbolic unaries, plus the remaining diff --git a/docs/performance/CPU_ACTIVATION_GAPS.md b/docs/performance/CPU_ACTIVATION_GAPS.md index 0a0c929466..2fc73481c2 100644 --- a/docs/performance/CPU_ACTIVATION_GAPS.md +++ b/docs/performance/CPU_ACTIVATION_GAPS.md @@ -149,18 +149,23 @@ Two effects explain most of the table, and neither is the transcendental math: | `Exp` f32 | 0.14 | 0.82–0.92 | ported `MlasComputeExpVector` to AVX2 (#1093) | | bf16 ⇄ f32 conversion | scalar | 1.6–2.8× | bulk AVX2 conversion (#1041) | | `com.microsoft` activations unreachable | never ran | reachable | shape-inference table entries (#1082) | +| `Celu`, `Mish` — no kernel at all | ORT ran them | 0.29–0.64 | native AVX2 kernels (#1235) | +| `Log` f32 | 1.76 (4k), 1.66 (1M) | 0.77, 0.55 | AVX2 `log_ps`, was scalar `libm` per element (#1235) | ## Activations with no kernel at all -Distinct from the sections below: these are not declines, they are missing -features. ORT runs them because this EP has nothing to run. - -- **`Celu`** — ONNX opset 12. `max(0,x) + min(0, alpha*(exp(x/alpha)-1))`. -- **`Mish`** — ONNX opset 18. `x * tanh(softplus(x))`. - -Both are cheap to add on top of the existing `exp`/`tanh` AVX2 primitives, and -`activation_and_norm_ops_clear_every_capability_filter` has a comment naming -them so they are added to that test the moment a kernel lands. +None, as of #1235. `Celu` (opset 12) and `Mish` (opset 18) were the last two: +distinct from a decline, they were missing features, and ORT ran them because +this EP had nothing to run. Both now have native AVX2 kernels built on the +existing `exp`/`tanh` primitives, and both are in +`activation_and_norm_ops_clear_every_capability_filter`'s `OWNED` list, so a +future regression through either the shape filter or the dtype filter fails +the suite. + +This category is worth keeping as a heading even while it is empty, because +the failure mode it describes is invisible to the obvious test: a coverage +list written from the ops that exist cannot report an op nobody implemented. +Both of these cleared `supports_op` by never reaching it. ## Ops that still reach ORT despite being supported