perf(linalg): run SiLU on a dedicated AVX-512 kernel

linalg/benches/activations_avx512.rs

@@ -1,11 +1,9 @@
 // Microbenchmark: AVX-512 (zmm, 16-wide) element-wise activation kernels vs
 // their x86 predecessor.
 //
-//   sigmoid, tanh              : predecessor = FMA (256-bit, 8-wide) kernel
+//   sigmoid, tanh, silu        : 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)
 //
 // 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
@@ -71,11 +69,13 @@ fn benches(c: &mut Criterion) {
         enable_ftz_daz();
         use tract_linalg::x86_64_fma::act::*;
         use tract_linalg::x86_64_fma::{
-            avx512_sigmoid_f32, avx512_tanh_f32, fma_sigmoid_f32, fma_tanh_f32,
+            avx512_sigmoid_f32, avx512_silu_f32, avx512_tanh_f32, fma_sigmoid_f32, fma_silu_f32,
+            fma_tanh_f32,
         };
 
         bench_pair!(c, "sigmoid_f32", "fma", fma_sigmoid_f32, avx512_sigmoid_f32);
         bench_pair!(c, "tanh_f32", "fma", fma_tanh_f32, avx512_tanh_f32);
+        bench_pair!(c, "silu_f32", "fma", fma_silu_f32, avx512_silu_f32);
         bench_pair!(
             c,
             "hardswish_f32",
@@ -91,13 +91,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",

linalg/src/x86_64_fma.rs

@@ -76,6 +76,7 @@ silu_impl!(f32, fma_silu_f32, 8, 8, is_x86_feature_detected!("fma"));
 // nr()=16 (any multiple of 16 is safe).
 tanh_impl!(f32, avx512_tanh_f32, 16, 16, is_x86_feature_detected!("avx512f"));
 sigmoid_impl!(f32, avx512_sigmoid_f32, 16, 16, is_x86_feature_detected!("avx512f"));
+silu_impl!(f32, avx512_silu_f32, 16, 16, is_x86_feature_detected!("avx512f"));
 
 fn plug_avx2(_ops: &mut Ops) {}
 
@@ -117,7 +118,7 @@ fn plug_avx512f(ops: &mut Ops) {
     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.silu_f32 = Box::new(|| avx512_silu_f32::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());

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

linalg/src/x86_64_fma/act_f16.rs

@@ -3,19 +3,30 @@
 //! 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
+//! back.
+//! 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.
 
 use tract_data::internal::f16;
 
 const CHUNK: usize = 256;
 
+/// SiLU's f32 scratch length, wider than `CHUNK`.
+///
+/// Every call into an f32 kernel pays a fixed cost that does not shrink with the
+/// length passed to it: MXCSR is saved, overwritten and restored, resynchronising
+/// the FP pipeline at both ends. `avx512_silu_f32` needs a longer call than
+/// `CHUNK` to amortise that and to get its four 64-lane groups in flight across
+/// the divide; at `CHUNK` the ymm `fma_silu_f32`, whose groups are 16 lanes wide,
+/// wins instead. Must stay a multiple of `nr`.
+const SILU_CHUNK: usize = 1024;
+
 // Vectorized f16 <-> f32 helpers using vcvtph2ps / vcvtps2ph. Rustc + LLVM
 // do NOT autovectorize the scalar `.to_f32()` loop (the half crate's method
 // has branches / function-call overhead), so we drive the conversion with
 // intrinsics directly. Both helpers process 16 lanes per iteration; the tail
-// (which only fires for the 1-15 leftover lanes inside a CHUNK = 256 batch)
+// (which only fires for the 1-15 leftover lanes inside a CHUNK-sized batch)
 // falls back to scalar.
 #[target_feature(enable = "avx512f")]
 unsafe fn cvt_f16_to_f32(src: &[f16], dst: &mut [f32]) {
@@ -137,11 +148,11 @@ ew_impl_f16_via_f32!(
     x86_64_avx512_silu_f16_16n,
     16,
     16,
-    CHUNK,
+    SILU_CHUNK,
     64,
     cvt_f16_to_f32,
     cvt_f32_to_f16,
-    super::act::x86_64_avx512_silu_f32_16n
+    super::avx512_silu_f32
 );
 
 #[cfg(test)]