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); + } +}