silu: serve f16 from a table instead of widening every chunk to f32

linalg/benches/silu.rs

@@ -36,5 +36,31 @@ fn linalg(input: &mut [f32]) {
     (tract_linalg::ops().silu_f32)().run(input).unwrap();
 }
 
+#[cfg(target_arch = "aarch64")]
+fn silu_f16(c: &mut Criterion) {
+    for n in [1024usize, 65536, 1 << 20] {
+        let mut group = c.benchmark_group(format!("silu_f16/{n}"));
+        group.throughput(Throughput::Elements(n as u64));
+        let mut input = unsafe { Tensor::uninitialized_aligned::<f16>(&[n], 16).unwrap() };
+        let input = unsafe { input.as_slice_mut_unchecked::<f16>() };
+        for (i, x) in input.iter_mut().enumerate() {
+            *x = f16::from_f32((i as f32 / 10.0).sin() * 5.0);
+        }
+        group.bench_function("generic", |b| {
+            b.iter(|| tract_linalg::generic::HSiLU8::run(input, ()))
+        });
+        group.bench_function("f32-roundtrip", |b| {
+            b.iter(|| tract_linalg::arm64::arm64simd_silu_f16_4n::run(input, ()))
+        });
+        group.bench_function("lut", |b| {
+            b.iter(|| tract_linalg::arm64::arm64simd_silu_f16_lut_8n::run(input, ()))
+        });
+        group.finish();
+    }
+}
+
+#[cfg(target_arch = "aarch64")]
+criterion_group!(benches, silu_f32, silu_f16);
+#[cfg(not(target_arch = "aarch64"))]
 criterion_group!(benches, silu_f32);
 criterion_main!(benches);

linalg/src/arm64.rs

@@ -489,6 +489,7 @@ pub fn plug(ops: &mut Ops) {
     ops.leaky_relu_f32 = Box::new(|| arm64simd_leaky_relu_f32_8n::ew());
     ops.hardswish_f32 = Box::new(|| arm64simd_hardswish_f32_8n::ew());
     ops.silu_f32 = Box::new(|| arm64simd_silu_f32_4n_fused::ew());
+    ops.silu_f16 = Box::new(|| arm64simd_silu_f16_lut_8n::ew());
     ops.gelu_f32 = Box::new(|| arm64simd_gelu_f32_4n_fused::ew());
     ops.sigmoid_f32 = Box::new(|| arm64simd_sigmoid_f32_4n::ew());
     ops.tanh_f32 = Box::new(|| arm64simd_tanh_f32_4n::ew());
@@ -507,12 +508,9 @@ pub fn plug(ops: &mut Ops) {
         ops.max_f16 = Box::new(|| arm64fp16_max_f16_32n::red());
         ops.sum_f16 = Box::new(|| arm64fp16_sum_f16_32n::red());
         ops.mul_by_scalar_f16 = Box::new(|| arm64fp16_mul_by_scalar_f16_32n::ew());
-        // TODO: Change this SiLU kernel once we have a native-FP16 one
-        ops.silu_f16 = Box::new(|| arm64simd_silu_f16_4n::ew());
     } else {
-        log::info!("No native fp16 support; f32-roundtrip NEON sigmoid_f16 and silu_f16 activated");
+        log::info!("No native fp16 support; f32-roundtrip NEON sigmoid_f16 activated");
         ops.sigmoid_f16 = Box::new(|| arm64simd_sigmoid_f16_4n::ew());
-        ops.silu_f16 = Box::new(|| arm64simd_silu_f16_4n::ew());
     }
     #[cfg(any(target_os = "macos", all(target_os = "ios", feature = "apple-amx-ios")))]
     {

linalg/src/arm64/arm64simd.rs

@@ -16,6 +16,7 @@ mod unicast;
 
 pub use act_f16::arm64simd_sigmoid_f16_4n;
 pub use act_f16::arm64simd_silu_f16_4n;
+pub use act_f16::arm64simd_silu_f16_lut_8n;
 pub use by_scalar::*;
 pub use gelu::arm64simd_gelu_f32_4n;
 pub use gelu_fused::arm64simd_gelu_f32_4n_fused;

linalg/src/arm64/arm64simd/act_f16.rs

@@ -173,3 +173,69 @@ pub mod test_arm64simd_silu_f16_4n {
     use super::*;
     silu_frame_tests!(true, f16, arm64simd_silu_f16_4n);
 }
+
+/// Every f16 bit pattern mapped through the NEON f32 SiLU kernel and rounded
+/// back, so the activation is one load per element.
+///
+/// The table is filled by the same kernel the f32-roundtrip path runs, over an
+/// aligned whole number of `nr`-blocks — which is the case the element-wise
+/// frame hands to that kernel directly — so the two agree bit for bit. Calling
+/// the kernel raw rather than through the frame also keeps this off the frame's
+/// thread-local scratch, which is already borrowed whenever a kernel runs.
+/// 128 KiB, built on first use.
+fn silu_lut() -> &'static [u16; 1 << 16] {
+    use crate::frame::element_wise::ElementWiseKer;
+    use tract_data::prelude::Tensor;
+    static LUT: std::sync::OnceLock<Box<[u16; 1 << 16]>> = std::sync::OnceLock::new();
+    LUT.get_or_init(|| {
+        let mut values = unsafe {
+            Tensor::uninitialized_aligned::<f32>(&[1 << 16], 16)
+                .expect("silu lookup table allocation")
+        };
+        let widened = unsafe { values.as_slice_mut_unchecked::<f32>() };
+        widened
+            .iter_mut()
+            .enumerate()
+            .for_each(|(bits, v)| *v = f16::from_bits(bits as u16).to_f32());
+        super::arm64simd_silu_f32_4n_fused::run(widened, ());
+        let mut lut = Box::new([0u16; 1 << 16]);
+        lut.iter_mut()
+            .zip(widened.iter())
+            .for_each(|(slot, v)| *slot = f16::from_f32(*v).to_bits());
+        lut
+    })
+}
+
+ew_impl_wrap!(
+    f16,
+    arm64simd_silu_f16_lut_8n,
+    8,
+    4,
+    (),
+    #[inline(never)]
+    fn run(buf: &mut [f16], _params: ()) {
+        let lut = silu_lut();
+        buf.iter_mut().for_each(|x| *x = f16::from_bits(lut[x.to_bits() as usize]));
+    }
+);
+
+#[cfg(test)]
+mod silu_f16_agreement {
+    use super::*;
+    use crate::frame::element_wise::ElementWiseKer;
+
+    #[test]
+    fn lut_matches_the_f32_roundtrip_on_every_f16() {
+        let all: Vec<f16> = (0..=u16::MAX).map(f16::from_bits).collect();
+        let mut roundtrip = all.clone();
+        let mut lut = all;
+        arm64simd_silu_f16_4n::ew().run(&mut roundtrip).unwrap();
+        arm64simd_silu_f16_lut_8n::ew().run(&mut lut).unwrap();
+        let mismatch = roundtrip
+            .iter()
+            .zip(&lut)
+            .position(|(a, b)| a.to_bits() != b.to_bits())
+            .map(|i| (f16::from_bits(i as u16), roundtrip[i], lut[i]));
+        assert_eq!(mismatch, None);
+    }
+}