gelu: dispatch f16 through linalg and back it with a lookup table

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 {

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::<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::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);

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;

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<f32> 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<f16> 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<Box<[u16; 1 << 16]>> = 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<f16> 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<f16> = (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<f16> = (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(&registered)
+            .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");
+    }
+}

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::<generic::GenericLut8>::new(table))),
         max_f16: Box::new(|| generic::reduce::max::HMax8::red()),