Skip to main content

tract_linalg/generic/
silu.rs

1#![allow(clippy::excessive_precision)]
2use crate::generic::sigmoid::{LOW, ssigmoid};
3use tract_data::internal::*;
4
5// f32 SiLU, as `max(x, LOW) * ssigmoid(x)`.
6//
7// The factor is floored at [`LOW`] rather than left as `x`: past the clamp [`ssigmoid`]
8// returns the constant `4.8e-7` instead of exactly 0, so an unfloored factor would take
9// the negative tail to `-inf` instead of decaying to 0. Floored, `x < LOW` saturates at
10// `LOW * ssigmoid(LOW)` ~= `-6.9e-6`, which the true SiLU approaches from below.
11routine_ew_rust!(generic;
12    f32,
13    generic_silu_f32_4n,
14    4,
15    4,
16    fn run(x: &mut [f32], _: ()) {
17        debug_assert!(x.len() % Self::nr() == 0);
18        debug_assert!(x.as_ptr() as usize % Self::alignment_bytes() == 0);
19        x.iter_mut().for_each(|px| *px = px.max(LOW) * ssigmoid(*px));
20    },
21    func(Silu)
22);
23
24// f16 SiLU, evaluated on the f32 fit and narrowed, and floored the same way.
25routine_ew_rust!(generic;
26    f16,
27    generic_silu_f16_8n,
28    8,
29    8,
30    fn run(x: &mut [f16], _: ()) {
31        debug_assert!(x.len() % Self::nr() == 0);
32        debug_assert!(x.as_ptr() as usize % Self::alignment_bytes() == 0);
33        x.iter_mut().for_each(|px| {
34            let x_f32 = px.to_f32();
35            *px = f16::from_f32(x_f32.max(LOW) * ssigmoid(x_f32));
36        });
37    },
38    func(Silu)
39);