From 4573163f66e02a8619643be7e5e074d172fd49e2 Mon Sep 17 00:00:00 2001 From: czoli1976 Date: Mon, 3 Aug 2026 22:54:07 +0100 Subject: [PATCH] silu: serve f16 from a table instead of widening every chunk to f32 The aarch64 f16 SiLU converted each chunk into an f32 scratch, ran the NEON f32 kernel over it and converted back, on every call. Map the 65536 f16 bit patterns through that same kernel once and keep the result, so the activation is a single load per element. The table is filled by calling the f32 kernel over an aligned whole number of nr-blocks, which is what the element-wise frame hands it directly, so the two agree bit for bit. Co-Authored-By: Claude Opus 5 --- linalg/benches/silu.rs | 26 +++++++++++ linalg/src/arm64.rs | 6 +-- linalg/src/arm64/arm64simd.rs | 1 + linalg/src/arm64/arm64simd/act_f16.rs | 66 +++++++++++++++++++++++++++ 4 files changed, 95 insertions(+), 4 deletions(-) diff --git a/linalg/benches/silu.rs b/linalg/benches/silu.rs index acc44ebe20..b07fc19384 100644 --- a/linalg/benches/silu.rs +++ b/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::(&[n], 16).unwrap() }; + let input = unsafe { input.as_slice_mut_unchecked::() }; + 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); diff --git a/linalg/src/arm64.rs b/linalg/src/arm64.rs index 211b7be997..24e6e6348e 100644 --- a/linalg/src/arm64.rs +++ b/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")))] { diff --git a/linalg/src/arm64/arm64simd.rs b/linalg/src/arm64/arm64simd.rs index 0f6fc17e4e..ecc2edd568 100644 --- a/linalg/src/arm64/arm64simd.rs +++ b/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; diff --git a/linalg/src/arm64/arm64simd/act_f16.rs b/linalg/src/arm64/arm64simd/act_f16.rs index 6d5cdfbe19..983a6ddb70 100644 --- a/linalg/src/arm64/arm64simd/act_f16.rs +++ b/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> = std::sync::OnceLock::new(); + LUT.get_or_init(|| { + let mut values = unsafe { + Tensor::uninitialized_aligned::(&[1 << 16], 16) + .expect("silu lookup table allocation") + }; + let widened = unsafe { values.as_slice_mut_unchecked::() }; + 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 = (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); + } +}