Skip to main content

tract_linalg/generic/
sigmoid.rs

1#![allow(clippy::excessive_precision)]
2use tract_data::internal::*;
3
4/// The input clamp of the f32 sigmoid fit, shared by [`ssigmoid`] and by the SiLU kernels
5/// built on the same coefficients.
6///
7/// It alone holds the result inside `[0, 1]`, so a SiLU kernel needs it as the floor of
8/// the factor its sigmoid multiplies. See [`ssigmoid`] for what fixes its value.
9pub const LOW: f32 = -14.5;
10
11/// f32 sigmoid, as a rational minimax fit of `sigmoid(x) - 0.5` over `[LOW, HIGH]`.
12///
13/// The clamp keeps the `[0, 1]` range consumers rely on without an output clamp: on either
14/// tail `p / q` approaches ±0.5 and the `+ 0.5` cancels against it, so the cancellation
15/// has to stay above the ~6e-8 f32 rounding error of `p / q` rather than run down to zero.
16/// At `±14.5` the sum keeps six f32 steps of margin at both ends for every f32 input. It
17/// costs the tails their last few ulps — a saturated input returns `1 - 4.8e-7` or
18/// `4.8e-7`, not exactly 1 or 0. NaN propagates.
19///
20/// Raising the clamp voids this: `1 - sigmoid(16.29)` is already under the rounding error,
21/// so the sum crosses 1 at scattered inputs from there up. `14.5` rather than something
22/// nearer 16 because the value is shared, and `armv7neon_sigmoid_f32_4n` crosses from
23/// 15.92 — it refines a `vrecpe` estimate instead of dividing.
24pub fn ssigmoid(x: f32) -> f32 {
25    const HIGH: f32 = -LOW;
26
27    const ALPHA_13: f32 = -4.433153405e-18;
28    const ALPHA_11: f32 = 1.169974371e-14;
29    const ALPHA_9: f32 = -1.875289645e-11;
30    const ALPHA_7: f32 = 4.257889523e-8;
31    const ALPHA_5: f32 = 0.00004811817576;
32    const ALPHA_3: f32 = 0.008163842030;
33    const ALPHA_1: f32 = 0.2499999971;
34    const BETA_6: f32 = 3.922935744e-6;
35    const BETA_4: f32 = 0.001524872358;
36    const BETA_2: f32 = 0.1159886749;
37    const BETA_0: f32 = 1.0;
38
39    let x = x.clamp(LOW, HIGH);
40
41    let x2 = x * x;
42
43    let p = ALPHA_13;
44    let p = x2 * p + ALPHA_11;
45    let p = x2 * p + ALPHA_9;
46    let p = x2 * p + ALPHA_7;
47    let p = x2 * p + ALPHA_5;
48    let p = x2 * p + ALPHA_3;
49    let p = x2 * p + ALPHA_1;
50    let p = p * x;
51
52    let q = BETA_6;
53    let q = x2 * q + BETA_4;
54    let q = x2 * q + BETA_2;
55    let q = x2 * q + BETA_0;
56
57    p / q + 0.5
58}
59
60/// f16 sigmoid, as a rational minimax fit of `sigmoid(x) - 0.5` over `[LOW, HIGH]`.
61///
62/// Needs no output clamp, like [`ssigmoid`]: the `+ 0.5` cancellation stays inside
63/// (0, 1) for every non-NaN f16 input, whatever the clamp.
64pub fn hsigmoid(x: f16) -> f16 {
65    /*
66     * (x (0.249895 + x^2 (0.00400222 - 0.0000124702 x^2)))
67     * /
68     * (1. + 0.098734 x^2)
69     */
70
71    const LOW: f16 = f16::from_f32_const(-6.92);
72    const HIGH: f16 = f16::from_f32_const(6.92);
73
74    const ALPHA_5: f16 = f16::from_f32_const(-0.0000124702);
75    const ALPHA_3: f16 = f16::from_f32_const(0.00400222);
76    const ALPHA_1: f16 = f16::from_f32_const(0.249895);
77
78    const BETA_2: f16 = f16::from_f32_const(0.098734);
79    const BETA_0: f16 = f16::from_f32_const(1.0);
80
81    let x = x.clamp(LOW, HIGH);
82
83    let x2 = x * x;
84
85    let p = ALPHA_5;
86    let p = x2 * p + ALPHA_3;
87    let p = x2 * p + ALPHA_1;
88    let p = p * x;
89
90    let q = BETA_2;
91    let q = x2 * q + BETA_0;
92
93    p / q + f16::from_f32_const(0.5)
94}
95
96routine_ew_rust!(generic;
97    f32,
98    generic_sigmoid_f32_4n,
99    4,
100    4,
101    fn run(x: &mut [f32], _: ()) {
102        debug_assert!(x.len() % Self::nr() == 0);
103        debug_assert!(x.as_ptr() as usize % Self::alignment_bytes() == 0);
104        x.iter_mut().for_each(|px| *px = ssigmoid(*px))
105    },
106    func(Sigmoid)
107);
108
109routine_ew_rust!(generic;
110    f16,
111    generic_sigmoid_f16_8n,
112    8,
113    8,
114    fn run(x: &mut [f16], _: ()) {
115        debug_assert!(x.len() % Self::nr() == 0);
116        debug_assert!(x.as_ptr() as usize % Self::alignment_bytes() == 0);
117        x.iter_mut().for_each(|px| *px = hsigmoid(*px))
118    },
119    func(Sigmoid)
120);