diff --git a/core/src/ops/nn/gelu_approximate.rs b/core/src/ops/nn/gelu_approximate.rs index 81cc8997d5..7bb804daf1 100644 --- a/core/src/ops/nn/gelu_approximate.rs +++ b/core/src/ops/nn/gelu_approximate.rs @@ -12,11 +12,15 @@ fn gelu_approx_f32(x: f32, pow: i32) -> f32 { element_wise!(gelu_approximate, GeluApproximate { fast_impl: bool }, [f16] => |op, xs| { - let pow = if op.fast_impl { 2 } else { 3 }; - xs.iter_mut().for_each(|x| { - *x = f16::from_f32(gelu_approx_f32(x.to_f32(), pow)); - }); - Ok(()) + if op.fast_impl { + // pow=2 fast path: no linalg kernel yet, scalar fallback. + xs.iter_mut().for_each(|x| { + *x = f16::from_f32(gelu_approx_f32(x.to_f32(), 2)); + }); + Ok(()) + } else { + (tract_linalg::ops().gelu_f16)().run(xs) + } }, [f32] => |op, xs| { if op.fast_impl { diff --git a/linalg/benches/gelu.rs b/linalg/benches/gelu.rs index 31b1df8e05..54d425c111 100644 --- a/linalg/benches/gelu.rs +++ b/linalg/benches/gelu.rs @@ -23,6 +23,24 @@ fn gelu_f32(c: &mut Criterion) { }); } +fn gelu_f16(c: &mut Criterion) { + for n in [1024usize, 65536, 1 << 20] { + let mut group = c.benchmark_group(format!("gelu_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::HGelu8::run(input, ())) + }); + group + .bench_function("lut", |b| b.iter(|| tract_linalg::generic::HGeluLut8::run(input, ()))); + group.finish(); + } +} + #[inline(never)] fn rust_scalar(input: &mut [f32]) { // Match tract's GeluApproximate scalar formula (pow=3). @@ -40,5 +58,5 @@ fn linalg(input: &mut [f32]) { (tract_linalg::ops().gelu_f32)().run(input).unwrap(); } -criterion_group!(benches, gelu_f32); +criterion_group!(benches, gelu_f32, gelu_f16); criterion_main!(benches); diff --git a/linalg/src/generic.rs b/linalg/src/generic.rs index c9e62014f2..845d32d125 100644 --- a/linalg/src/generic.rs +++ b/linalg/src/generic.rs @@ -21,7 +21,7 @@ use crate::{BinOp, LinalgRegistry}; pub use self::by_scalar::{HMulByScalar8, SMulByScalar4}; pub use self::erf::SErf4; -pub use self::gelu::{HGelu8, SGelu4}; +pub use self::gelu::{HGelu8, HGeluLut8, SGelu4}; pub use self::hardswish::{HHardSwish8, SHardSwish4}; pub use self::leaky_relu::{HLeakyRelu8, SLeakyRelu4}; pub use self::lut::GenericLut8; diff --git a/linalg/src/generic/gelu.rs b/linalg/src/generic/gelu.rs index d62d09dd46..32b5c51f42 100644 --- a/linalg/src/generic/gelu.rs +++ b/linalg/src/generic/gelu.rs @@ -13,6 +13,15 @@ use tract_data::internal::*; const SQRT_2_OVER_PI: f32 = 0.7978845608028654; const COEF: f32 = 0.044715; +/// The scalar reference every GELU kernel in this module is defined against. +/// `HGeluLut8`'s table is built from this function, so the table is bit-identical +/// to `HGelu8` by construction rather than by approximation. +#[inline] +fn gelu(v: f32) -> f32 { + let inner = SQRT_2_OVER_PI * (v + COEF * v * v * v); + 0.5 * v * (1.0 + inner.tanh()) +} + #[derive(Clone, Debug)] pub struct SGelu4; @@ -36,11 +45,7 @@ impl ElementWiseKer for SGelu4 { fn run(x: &mut [f32], _: ()) { debug_assert!(x.len() % Self::nr() == 0); debug_assert!(x.as_ptr() as usize % Self::alignment_bytes() == 0); - x.iter_mut().for_each(|px| { - let v = *px; - let inner = SQRT_2_OVER_PI * (v + COEF * v * v * v); - *px = 0.5 * v * (1.0 + inner.tanh()); - }); + x.iter_mut().for_each(|px| *px = gelu(*px)); } } @@ -67,11 +72,49 @@ impl ElementWiseKer for HGelu8 { fn run(x: &mut [f16], _: ()) { debug_assert!(x.len() % Self::nr() == 0); debug_assert!(x.as_ptr() as usize % Self::alignment_bytes() == 0); - x.iter_mut().for_each(|px| { - let v = px.to_f32(); - let inner = SQRT_2_OVER_PI * (v + COEF * v * v * v); - *px = f16::from_f32(0.5 * v * (1.0 + inner.tanh())); - }); + x.iter_mut().for_each(|px| *px = f16::from_f32(gelu(px.to_f32()))); + } +} + +/// Every f16 bit pattern mapped through `gelu`, so the whole activation is one +/// load per element. 128 KiB, built on first use: a model with no f16 GELU never +/// pays for it, and the build costs 65536 scalar evaluations. +fn gelu_lut() -> &'static [u16; 1 << 16] { + static LUT: std::sync::OnceLock> = std::sync::OnceLock::new(); + LUT.get_or_init(|| { + let mut lut = Box::new([0u16; 1 << 16]); + for (bits, slot) in lut.iter_mut().enumerate() { + *slot = f16::from_f32(gelu(f16::from_bits(bits as u16).to_f32())).to_bits(); + } + lut + }) +} + +#[derive(Clone, Debug)] +pub struct HGeluLut8; + +impl ElementWiseKer for HGeluLut8 { + fn name() -> &'static str { + "lut" + } + + fn alignment_bytes() -> usize { + 16 + } + + fn alignment_items() -> usize { + 4 + } + + fn nr() -> usize { + 8 + } + + fn run(x: &mut [f16], _: ()) { + debug_assert!(x.len() % Self::nr() == 0); + debug_assert!(x.as_ptr() as usize % Self::alignment_bytes() == 0); + let lut = gelu_lut(); + x.iter_mut().for_each(|px| *px = f16::from_bits(lut[px.to_bits() as usize])); } } @@ -86,3 +129,45 @@ pub mod s { pub mod h { gelu_frame_tests!(true, tract_data::internal::f16, crate::generic::gelu::HGelu8); } + +#[cfg(test)] +mod lut { + use super::*; + + #[test] + fn lut_matches_scalar_kernel_on_every_f16() { + let all: Vec = (0..=u16::MAX).map(f16::from_bits).collect(); + let mut reference = all.clone(); + let mut lut = all; + HGelu8::ew().run(&mut reference).unwrap(); + HGeluLut8::ew().run(&mut lut).unwrap(); + let mismatch = reference + .iter() + .zip(&lut) + .position(|(a, b)| a.to_bits() != b.to_bits()) + .map(|i| (f16::from_bits(i as u16), reference[i], lut[i])); + assert_eq!(mismatch, None); + } + + fn ordered(x: f16) -> i32 { + let b = x.to_bits(); + if b & 0x8000 != 0 { !b as i32 } else { (b | 0x8000) as i32 } + } + + #[test] + fn registered_kernel_tracks_the_scalar_kernel_on_every_f16() { + let all: Vec = (0..=u16::MAX).map(f16::from_bits).collect(); + let mut reference = all.clone(); + let mut registered = all; + HGelu8::ew().run(&mut reference).unwrap(); + (crate::ops().gelu_f16)().run(&mut registered).unwrap(); + let worst = reference + .iter() + .zip(®istered) + .filter(|(a, b)| !(a.is_nan() && b.is_nan())) + .map(|(a, b)| (ordered(*a) - ordered(*b)).abs()) + .max() + .unwrap(); + assert!(worst <= 1, "registered gelu_f16 drifts {worst} ulp from the scalar kernel"); + } +} diff --git a/linalg/src/lib.rs b/linalg/src/lib.rs index feba41da18..31c62c9d8d 100644 --- a/linalg/src/lib.rs +++ b/linalg/src/lib.rs @@ -240,7 +240,7 @@ pub fn generic() -> Ops { hardswish_f32: Box::new(|| generic::SHardSwish4::ew()), silu_f16: Box::new(|| generic::HSiLU8::ew()), silu_f32: Box::new(|| generic::SSiLU4::ew()), - gelu_f16: Box::new(|| generic::HGelu8::ew()), + gelu_f16: Box::new(|| generic::HGeluLut8::ew()), gelu_f32: Box::new(|| generic::SGelu4::ew()), lut_u8: Box::new(|table: &[u8]| Box::new(lut::LutImpl::::new(table))), max_f16: Box::new(|| generic::reduce::max::HMax8::red()),