Skip to content

Repository navigation

perf(cpu-ep): AVX2 kernels for LeakyRelu, HardSigmoid, ThresholdedRelu and Selu #1243

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 1 commit into from
Aug 18, 2026
Merged

perf(cpu-ep): AVX2 kernels for LeakyRelu, HardSigmoid, ThresholdedRelu and Selu #1243

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
113 changes: 104 additions & 9 deletions crates/onnx-runtime-ep-cpu-plugin/tests/plugin_ort_e2e.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5878,11 +5878,10 @@ fn unary_bench_cases() -> Vec<MatmulFamilyCase> {
false,
0.0,
),
// `Elu` exercises the AVX2 kernel and the zero-copy tensor path.
// `LeakyRelu` deliberately does not: it still has no SIMD kernel, so it
// measures what the generic scalar activation path costs -- two extra
// passes over the tensor and a per-element `match` -- and is the
// control this file needs to keep that cost visible.
// The activation family that moved off the generic scalar path: Elu,
// LeakyRelu, HardSigmoid, ThresholdedRelu and Selu all have AVX2
// kernels and the zero-copy tensor path now. `Swish(alpha != 1)` is
// the only one left on the generic path.
unary_case("bench_elu_f32_4k", "Elu", "", ELEM_F32, K4, 17, false, 1e-4),
unary_case("bench_elu_f32_1m", "Elu", "", ELEM_F32, M1, 17, false, 1e-4),
unary_case(
Expand All @@ -5895,6 +5894,76 @@ fn unary_bench_cases() -> Vec<MatmulFamilyCase> {
false,
1e-4,
),
unary_case(
"bench_leakyrelu_f32_4k",
"LeakyRelu",
"",
ELEM_F32,
K4,
17,
false,
1e-4,
),
unary_case(
"bench_hardsigmoid_f32_4k",
"HardSigmoid",
"",
ELEM_F32,
K4,
17,
false,
1e-4,
),
unary_case(
"bench_hardsigmoid_f32_1m",
"HardSigmoid",
"",
ELEM_F32,
M1,
17,
false,
1e-4,
),
unary_case(
"bench_thresholdedrelu_f32_4k",
"ThresholdedRelu",
"",
ELEM_F32,
K4,
17,
false,
1e-4,
),
unary_case(
"bench_thresholdedrelu_f32_1m",
"ThresholdedRelu",
"",
ELEM_F32,
M1,
17,
false,
1e-4,
),
unary_case(
"bench_selu_f32_4k",
"Selu",
"",
ELEM_F32,
K4,
17,
false,
1e-4,
),
unary_case(
"bench_selu_f32_1m",
"Selu",
"",
ELEM_F32,
M1,
17,
false,
1e-4,
),
// `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.
Expand Down Expand Up @@ -7040,21 +7109,25 @@ fn unary_result_matches(ours: f32, ort: f32, rel: f32) -> bool {
/// 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_elu_and_log_execute_locally_and_match_ort_numerics() {
fn the_native_activation_family_executes_locally_and_matches_ort_numerics() {
let _lock = lock_ort_ep();
let x = unary_special_values();
let len = x.len();
let bytes: Vec<u8> = 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); 6] = [
let specs: [(&'static str, &'static str, u64, f32); 10] = [
("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),
("parity_elu_f32", "Elu", 17, 2e-6),
("parity_elu_a2_f32", "Elu", 17, 2e-6),
("parity_leakyrelu_f32", "LeakyRelu", 17, 2e-6),
("parity_hardsigmoid_f32", "HardSigmoid", 17, 2e-6),
("parity_thresholdedrelu_f32", "ThresholdedRelu", 17, 2e-6),
("parity_selu_f32", "Selu", 17, 2e-6),
];

let mut checked = 0usize;
Expand Down Expand Up @@ -7098,6 +7171,26 @@ fn celu_mish_elu_and_log_execute_locally_and_match_ort_numerics() {
);
assert_eq!(ours_out.len(), ort_out.len(), "{name}: output length");

// `unary_result_matches` calls +0 and -0 equal, which is right for
// a tolerance comparison and wrong for the question these kernels
// actually had to answer: several of them are implemented as a
// *select* precisely so a signed zero survives, and one (`Selu`)
// has its branch written on `x < 0` rather than ONNX's `x > 0` for
// no other reason. Comparing the sign bit against ORT is the only
// thing that makes those choices evidence rather than assertion.
for (i, &xi) in x.iter().enumerate() {
if xi == 0.0 && ours_out[i] == 0.0 && ort_out[i] == 0.0 {
assert_eq!(
ours_out[i].to_bits(),
ort_out[i].to_bits(),
"{name}: sign of zero disagrees with ORT at x={xi:e} \
(ours {:e}, ORT {:e})",
ours_out[i],
ort_out[i]
);
}
}

let mut worst = 0.0f32;
let mut worst_at = usize::MAX;
let mut bad: Vec<String> = Vec::new();
Expand Down Expand Up @@ -7131,6 +7224,8 @@ fn celu_mish_elu_and_log_execute_locally_and_match_ort_numerics() {
}
checked += 1;
}
assert_eq!(checked, 6, "every parity case must have been checked");
eprintln!("\n✅ celu_mish_elu_and_log_execute_locally_and_match_ort_numerics: PASSED");
assert_eq!(checked, 10, "every parity case must have been checked");
eprintln!(
"\n✅ the_native_activation_family_executes_locally_and_matches_ort_numerics: PASSED"
);
}
92 changes: 79 additions & 13 deletions crates/onnx-runtime-ep-cpu/src/kernels/activations.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,21 +43,22 @@ impl Activation {
// has no `expm1` to match `exp_m1` with, and with ORT, which spells
// it the same way. See `elu_scalar`.
Self::Elu { alpha } => crate::kernels::simd_activations::elu_scalar(x, alpha),
// All four delegate to the slice module's scalars rather than
// spelling the formula twice: that is the only way the scalar
// fallback and the AVX2 kernel are guaranteed to answer the same
// thing for NaN and for signed zero, which is where transcribing
// an ONNX formula literally goes wrong.
Self::LeakyRelu { alpha } => {
if x >= 0.0 {
x
} else {
alpha * x
}
crate::kernels::simd_activations::leaky_relu_scalar(x, alpha)
}
Self::HardSigmoid { alpha, beta } => {
crate::kernels::simd_activations::hard_sigmoid_scalar(x, alpha, beta)
}
Self::Selu { alpha, gamma } => {
crate::kernels::simd_activations::selu_scalar(x, alpha, gamma)
}
Self::HardSigmoid { alpha, beta } => (alpha * x + beta).clamp(0.0, 1.0),
Self::Selu { alpha, gamma } => gamma * if x > 0.0 { x } else { alpha * x.exp_m1() },
Self::ThresholdedRelu { alpha } => {
if x > alpha {
x
} else {
0.0
}
crate::kernels::simd_activations::thresholded_relu_scalar(x, alpha)
}
// Swish/SiLU: x·sigmoid(alpha·x), evaluated via the numerically
// stable logistic to avoid overflow at large-magnitude inputs.
Expand Down Expand Up @@ -101,12 +102,16 @@ impl Activation {
Self::HardSigmoid { alpha, beta } => {
(f64::from(alpha) * x + f64::from(beta)).clamp(0.0, 1.0)
}
// `exp() - 1`, not `exp_m1`, for the same reason as the f32 path:
// at `x = -0` the two disagree on the sign of the result, and ORT
// answers `+0`. Keeping the f64 path on `exp_m1` would make the
// sign depend on the tensor's dtype.
Self::Selu { alpha, gamma } => {
f64::from(gamma)
* if x > 0.0 {
x
} else {
f64::from(alpha) * x.exp_m1()
f64::from(alpha) * (x.exp() - 1.0)
}
}
Self::ThresholdedRelu { alpha } => {
Expand Down Expand Up @@ -317,6 +322,28 @@ impl Kernel for ActivationKernel {
Activation::Elu { alpha } => contiguous_f32(&inputs[0], &mut outputs[0], |i, o| {
crate::kernels::simd_activations::elu_f32_slice(i, o, alpha);
}),
Activation::LeakyRelu { alpha } => {
contiguous_f32(&inputs[0], &mut outputs[0], |i, o| {
crate::kernels::simd_activations::leaky_relu_f32_slice(i, o, alpha);
})
}
Activation::HardSigmoid { alpha, beta } => {
contiguous_f32(&inputs[0], &mut outputs[0], |i, o| {
crate::kernels::simd_activations::hard_sigmoid_f32_slice(i, o, alpha, beta);
})
}
Activation::ThresholdedRelu { alpha } => {
contiguous_f32(&inputs[0], &mut outputs[0], |i, o| {
crate::kernels::simd_activations::thresholded_relu_f32_slice(i, o, alpha);
})
}
Activation::Selu { alpha, gamma } => {
contiguous_f32(&inputs[0], &mut outputs[0], |i, o| {
crate::kernels::simd_activations::selu_f32_slice(i, o, alpha, gamma);
})
}
// `Swish(alpha)` with alpha != 1 is the only activation left on the
// generic path; `SwishFactory` canonicalises alpha == 1 to SiLU.
_ => false,
};
if ran {
Expand Down Expand Up @@ -559,6 +586,45 @@ mod tests {
/// 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.
/// `Selu(-0) = +0` in both dtypes, because ORT answers `+0`.
///
/// The f32 path is pinned against ORT itself in the plugin's parity test;
/// this is the f64 arm, which no ORT comparison reaches (the parity models
/// are f32) and which used `exp_m1` -- the spelling that returns `-0` --
/// until this was fixed. Without it the sign of the answer would depend on
/// the tensor's dtype, which is not something a model author can see.
#[test]
fn selu_zero_sign_does_not_depend_on_dtype() {
let act = Activation::Selu {
alpha: 1.673_263_2,
gamma: 1.050_701,
};
assert_eq!(act.apply(-0.0f32).to_bits(), 0, "f32 Selu(-0) is +0");
assert_eq!(act.apply(0.0f32).to_bits(), 0, "f32 Selu(+0) is +0");
assert_eq!(act.apply_f64(-0.0f64).to_bits(), 0, "f64 Selu(-0) is +0");
assert_eq!(act.apply_f64(0.0f64).to_bits(), 0, "f64 Selu(+0) is +0");
// `Elu` and `LeakyRelu` are selects on `x < 0`, so `-0` takes the
// identity branch and keeps its sign in both dtypes. Pinned here so the
// contrast with `Selu` is deliberate rather than accidental.
for act in [
Activation::Elu { alpha: 1.0 },
Activation::LeakyRelu { alpha: 0.125 },
] {
assert_eq!(
act.apply(-0.0f32).to_bits(),
0x8000_0000,
"f32 {}(-0) stays -0",
act.name()
);
assert_eq!(
act.apply_f64(-0.0f64).to_bits(),
0x8000_0000_0000_0000,
"f64 {}(-0) stays -0",
act.name()
);
}
}

#[test]
fn celu_and_mish_propagate_nan_on_every_path() {
for alpha in [0.5f32, 1.0, 3.0] {
Expand Down
Loading
Loading