From 11b623eb318eaa5e5f0ae0f242bdb546a1fb4487 Mon Sep 17 00:00:00 2001 From: Clement Verrier <17490658+cverrier@users.noreply.github.com> Date: Mon, 3 Aug 2026 12:53:50 +0000 Subject: [PATCH] perf(linalg): run SiLU through the fused FMA kernel on AVX-512 hosts AVX-512 hosts ran silu_f32 and silu_f16 through a kernel composed over the zmm sigmoid, which pays a scratch copy and an extra traversal that the fused FMA kernel avoids. They now use the fused kernel at both widths. The composed kernel has no remaining users, and is removed. --- linalg/benches/activations_avx512.rs | 12 ++----- linalg/src/x86_64_fma.rs | 6 ++-- linalg/src/x86_64_fma/act.rs | 50 ---------------------------- linalg/src/x86_64_fma/act_f16.rs | 12 ++++--- 4 files changed, 14 insertions(+), 66 deletions(-) diff --git a/linalg/benches/activations_avx512.rs b/linalg/benches/activations_avx512.rs index 662cbe14ad..d3199dd499 100644 --- a/linalg/benches/activations_avx512.rs +++ b/linalg/benches/activations_avx512.rs @@ -4,8 +4,9 @@ // sigmoid, tanh : predecessor = FMA (256-bit, 8-wide) kernel // hardswish, leaky_relu, gelu: predecessor = generic scalar kernel // (no FMA kernel exists on x86) -// silu : predecessor = generic scalar kernel -// (the FMA kernel is benched in silu.rs) +// +// silu is absent: the fused FMA kernel wins on AVX-512 hosts too, so it is the +// plugged choice at every x86 tier and is benched against generic in silu.rs. // // All buffers are 64-byte aligned (AVX-512 alignment_bytes) and a multiple of // 64 elements so every kernel's nr() divides the length. Criterion reports the @@ -91,13 +92,6 @@ fn benches(c: &mut Criterion) { x86_64_avx512_leaky_relu_f32_64n, 0.1f32 ); - bench_pair!( - c, - "silu_f32", - "generic", - tract_linalg::generic::SSiLU4, - x86_64_avx512_silu_f32_16n - ); bench_pair!( c, "gelu_f32", diff --git a/linalg/src/x86_64_fma.rs b/linalg/src/x86_64_fma.rs index b8c941e449..56c4cb75dc 100644 --- a/linalg/src/x86_64_fma.rs +++ b/linalg/src/x86_64_fma.rs @@ -112,12 +112,14 @@ fn plug_avx512fp16(ops: &mut Ops) { log::info!("hardswish_f16: x86_64/avx512fp16 native activated"); } +/// `silu_f32` is deliberately not overridden here: the fused `fma_silu_f32` plugged by +/// `plug_fma` also beats a zmm SiLU composed over `avx512_sigmoid_f32`, because the +/// composition costs a scratch copy and a second traversal that the fused kernel does not. fn plug_avx512f(ops: &mut Ops) { ops.sigmoid_f32 = Box::new(|| avx512_sigmoid_f32::ew()); ops.tanh_f32 = Box::new(|| avx512_tanh_f32::ew()); ops.hardswish_f32 = Box::new(|| act::x86_64_avx512_hardswish_f32_64n::ew()); ops.leaky_relu_f32 = Box::new(|| act::x86_64_avx512_leaky_relu_f32_64n::ew()); - ops.silu_f32 = Box::new(|| act::x86_64_avx512_silu_f32_16n::ew()); ops.gelu_f32 = Box::new(|| act::x86_64_avx512_gelu_f32_16n::ew()); ops.sigmoid_f16 = Box::new(|| act_f16::x86_64_avx512_sigmoid_f16_16n::ew()); @@ -138,7 +140,7 @@ fn plug_avx512f(ops: &mut Ops) { log::info!( "sigmoid_f32, tanh_f32, hardswish_f32, leaky_relu_f32, \ - silu_f32, gelu_f32, \ + gelu_f32, \ sigmoid_f16, tanh_f16, hardswish_f16, leaky_relu_f16, \ silu_f16, gelu_f16, \ max_f32, softmax2_fastcompact_f32, softmax2_fastcompact_f16, erf_f32, \ diff --git a/linalg/src/x86_64_fma/act.rs b/linalg/src/x86_64_fma/act.rs index 27fcca353c..77675a31bf 100644 --- a/linalg/src/x86_64_fma/act.rs +++ b/linalg/src/x86_64_fma/act.rs @@ -175,56 +175,6 @@ pub mod test_x86_64_avx512_leaky_relu_f32_64n { ); } -// SiLU(x) = f * sigmoid(z), with z = clamp(x, -18.0, 18.0) and f = max(x, -18.0). -// The kernel-level composition mirrors the arm64 fused SiLU clamping strategy, not its -// exact bounds: save the input chunk, run the AVX-512 sigmoid kernel in place (it clamps -// its argument to z internally), then multiply back by the factor f. The factor floor -// must equal the sigmoid kernel's own clamp, which here is -18.0 (the AVX-512 sigmoid's -// range). -// -// The factor f is clamped below at -18.0 but left unbounded above: SiLU(x) ~ x as -// x -> +inf, so the factor must grow, whereas the upper clamp on z is only there to -// keep the sigmoid polynomial in range. Clamping f below keeps the negative tail -// bounded: since sigmoid is floored at sigmoid(-18.0), an unclamped factor would let -// x * sigmoid(-18.0) diverge toward -inf, while the clamped factor saturates at the -// constant -18.0 * sigmoid(-18.0) ~= -2.74e-7, which the true SiLU approaches from -// below as x -> -inf. -// -// nr() and CHUNK (256) are multiples of 16 so the sigmoid kernel always receives a -// 64-byte-aligned slice whose length is a multiple of 16. -ew_impl_wrap!( - f32, - x86_64_avx512_silu_f32_16n, - 16, - 16, - (), - #[inline(never)] - fn run(buf: &mut [f32], _: ()) { - debug_assert!(buf.len() % Self::nr() == 0); - debug_assert!(buf.as_ptr() as usize % Self::alignment_bytes() == 0); - const CHUNK: usize = 256; - let mut scratch = [0f32; CHUNK]; - let mut start = 0; - while start < buf.len() { - let end = (start + CHUNK).min(buf.len()); - let chunk = &mut buf[start..end]; - let n = chunk.len(); - scratch[..n].copy_from_slice(chunk); - super::avx512_sigmoid_f32::run(chunk, ()); - for i in 0..n { - chunk[i] *= scratch[i].max(-18.0); - } - start = end; - } - } -); - -#[cfg(test)] -pub mod test_x86_64_avx512_silu_f32_16n { - use super::*; - silu_frame_tests!(is_x86_feature_detected!("avx512f"), f32, x86_64_avx512_silu_f32_16n); -} - // Tanh-form GELU (pow=3) matching tract's GeluApproximate: // gelu(x) = 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3))) // Composed at the kernel level (mirrors arm64): save the original x, compute diff --git a/linalg/src/x86_64_fma/act_f16.rs b/linalg/src/x86_64_fma/act_f16.rs index c47d0055e2..751b6dfcc2 100644 --- a/linalg/src/x86_64_fma/act_f16.rs +++ b/linalg/src/x86_64_fma/act_f16.rs @@ -1,9 +1,11 @@ //! AVX-512 f16 element-wise activations for cores without native f16 arithmetic. //! -//! Each kernel round-trips through the matching f32 AVX-512 kernel via -//! `ew_impl_f16_via_f32!`: convert an f16 chunk into a 64-byte-aligned f32 scratch -//! (the f32 kernels assume 64-byte-aligned input), run the f32 kernel, convert -//! back. Conversion is driven through `std::arch` intrinsics directly (see the +//! Each kernel round-trips through an f32 kernel via `ew_impl_f16_via_f32!`: convert +//! an f16 chunk into a 64-byte-aligned f32 scratch (the widest input-alignment +//! contract among the f32 kernels reused here), run the f32 kernel, convert back. +//! SiLU reuses the fused `fma_silu_f32` rather than a zmm composition over +//! `avx512_sigmoid_f32`, which would pay a scratch copy and a second traversal. +//! Conversion is driven through `std::arch` intrinsics directly (see the //! helpers below) because rustc + LLVM do not autovectorize the scalar //! `f16::to_f32` / `f16::from_f32` loops. @@ -141,7 +143,7 @@ ew_impl_f16_via_f32!( 64, cvt_f16_to_f32, cvt_f32_to_f16, - super::act::x86_64_avx512_silu_f32_16n + super::fma_silu_f32 ); #[cfg(test)]