oxfuse_runtime/
math.rs

1//! `exp` and `sigmoid` with the same bits on every platform.
2//!
3//! The standard library's `exp` comes from the platform's math library, so x86
4//! and ARM, or two versions of one libc, can differ in the last bit. This
5//! module uses only IEEE multiplications, additions and fused multiply-adds,
6//! integer operations and a clamp, all of which are exact or rounded the same
7//! way everywhere, so the result depends on nothing but the input.
8//!
9//! The algorithm is the specification. It is within 1 ulp of the correctly
10//! rounded result, and equal to it for all but 6,412,479 of the 2^32 `f32` bit
11//! patterns (0.15 percent), checked against CORE-MATH over every one of them,
12//! see `scripts/exp-oracle`. The output bits of every input are pinned by hash in
13//! the tests, so changing a constant here is a visible, deliberate change.
14//!
15//! This file imports nothing from the crate, so `scripts/exp-oracle` can
16//! include it as it is.
17
18pub(crate) const LOG2E: f32 = std::f32::consts::LOG2_E;
19/// `ln 2` split into a leading `f32` and the rest, so `n * LN2_HI` loses
20/// nothing in the fused step that subtracts it.
21pub(crate) const LN2_HI: f32 = std::f32::consts::LN_2;
22pub(crate) const LN2_LO: f32 = -1.904_654_3e-9;
23/// `1.5 * 2^23`: adding it rounds to the nearest integer in the low mantissa
24/// bits, and subtracting it gives that integer back.
25pub(crate) const MAGIC: f32 = 12_582_912.0;
26/// Inputs beyond these give infinity and zero already; clamping keeps the
27/// power-of-two scaling below in range.
28pub(crate) const X_MAX: f32 = 89.0;
29pub(crate) const X_MIN: f32 = -104.0;
30/// `q(r) = (e^r - 1 - r) / r^2` on `[-ln2/2, ln2/2]`, degree 5, least squares
31/// fit in `f64` and rounded to `f32`. Lowest order first.
32pub(crate) const Q: [f32; 6] = [
33    0.5,
34    0.166_666_67,
35    0.041_666_463,
36    0.008_333_298,
37    0.001_393_368,
38    0.000_198_993_18,
39];
40
41/// `2^n` for `n` in `[-126, 127]`.
42#[inline(always)]
43fn pow2(n: i32) -> f32 {
44    f32::from_bits(((n + 127) as u32) << 23)
45}
46
47/// `e^x`. NaN gives NaN, `+inf` gives `+inf`, `-inf` gives 0, and a result
48/// below the smallest subnormal rounds to 0.
49#[inline(always)]
50pub(crate) fn exp(x: f32) -> f32 {
51    // `clamp` passes NaN through.
52    let x = x.clamp(X_MIN, X_MAX);
53    let n_f = x.mul_add(LOG2E, MAGIC) - MAGIC;
54    let n = n_f as i32;
55    let r = (-n_f).mul_add(LN2_HI, x);
56    let r = (-n_f).mul_add(LN2_LO, r);
57    let q = Q[5]
58        .mul_add(r, Q[4])
59        .mul_add(r, Q[3])
60        .mul_add(r, Q[2])
61        .mul_add(r, Q[1])
62        .mul_add(r, Q[0]);
63    let tail = (r * r) * q;
64    // 1 + r kept as a sum and its rounding error, so the single rounding that
65    // matters is the last addition.
66    let hi = 1.0 + r;
67    let lo = r - (hi - 1.0);
68    let e = hi + (lo + tail);
69    // Two scalings, so a result in the subnormal range is rounded once.
70    let half = n >> 1;
71    e * pow2(half) * pow2(n - half)
72}
73
74/// `1 / (1 + e^-x)`, the definition of the `Sigmoid` op: the same bits on every
75/// platform. Large negative `x` makes `exp` overflow to infinity and the
76/// result 0, without NaN.
77#[inline(always)]
78pub fn sigmoid(x: f32) -> f32 {
79    1.0 / (1.0 + exp(-x))
80}
81
82#[cfg(test)]
83mod tests {
84    use super::*;
85
86    /// Distance in representable values between two finite `f32`.
87    fn ulps(a: f32, b: f32) -> u32 {
88        let key = |x: f32| {
89            let bits = x.to_bits() as i64;
90            if bits & 0x8000_0000 != 0 {
91                -(bits & 0x7fff_ffff)
92            } else {
93                bits
94            }
95        };
96        (key(a) - key(b)).unsigned_abs() as u32
97    }
98
99    /// The correctly rounded result, up to the rare double rounding of `f64::exp`.
100    fn nearly_exact(x: f32) -> f32 {
101        f64::from(x).exp() as f32
102    }
103
104    #[test]
105    fn exp_of_the_special_values() {
106        assert_eq!(exp(0.0), 1.0);
107        assert_eq!(exp(-0.0), 1.0);
108        assert_eq!(exp(f32::INFINITY), f32::INFINITY);
109        assert_eq!(exp(f32::NEG_INFINITY).to_bits(), 0);
110        assert!(exp(f32::NAN).is_nan());
111        assert!(exp(-f32::NAN).is_nan());
112    }
113
114    #[test]
115    fn exp_of_one_is_e() {
116        assert_eq!(exp(1.0), std::f32::consts::E);
117        assert_eq!(exp(-1.0), 1.0 / std::f32::consts::E);
118    }
119
120    #[test]
121    fn exp_overflows_exactly_where_the_result_leaves_f32() {
122        // A result rounds to infinity once it reaches f32::MAX plus half an ulp.
123        let threshold = f64::from(f32::MAX) * (1.0 + 2f64.powi(-24));
124        let edge = 88.722_84_f32.to_bits();
125        for bits in edge - 40..edge + 40 {
126            let x = f32::from_bits(bits);
127            let finite = f64::from(x).exp() < threshold;
128            assert_eq!(exp(x).is_finite(), finite, "x={x}");
129            if finite {
130                assert!(ulps(exp(x), nearly_exact(x)) <= 1, "x={x}");
131            }
132        }
133        assert_eq!(exp(100.0), f32::INFINITY);
134        assert_eq!(exp(1.0e30), f32::INFINITY);
135    }
136
137    #[test]
138    fn exp_underflows_through_the_subnormals_to_zero() {
139        let smallest = f32::from_bits(1);
140        // ln(2^-149) = -103.2789, ln(2^-150) = -103.9720.
141        assert_eq!(exp(-103.0), nearly_exact(-103.0));
142        assert_eq!(exp(-103.27), smallest);
143        assert_eq!(exp(-103.99).to_bits(), 0);
144        assert_eq!(exp(-104.0).to_bits(), 0);
145        assert_eq!(exp(-1.0e30).to_bits(), 0);
146        for k in 1..=40 {
147            let x = -87.0 - 0.4 * k as f32;
148            assert!(ulps(exp(x), nearly_exact(x)) <= 1, "x={x}");
149        }
150    }
151
152    #[test]
153    fn exp_is_within_one_ulp_of_the_exact_result() {
154        let mut state = 7u64;
155        for _ in 0..200_000 {
156            state = state
157                .wrapping_mul(6364136223846793005)
158                .wrapping_add(1442695040888963407);
159            let x = f32::from_bits((state >> 32) as u32);
160            if !x.is_finite() || !(-103.0..=88.0).contains(&x) {
161                continue;
162            }
163            assert!(ulps(exp(x), nearly_exact(x)) <= 1, "x={x}");
164        }
165    }
166
167    #[test]
168    fn exp_is_exact_for_tiny_arguments() {
169        // e^x = 1 + x to within half an ulp for |x| below 2^-25.
170        for e in 26..60 {
171            let x = 2f32.powi(-e);
172            assert_eq!(exp(x), 1.0, "x=2^-{e}");
173            assert_eq!(exp(-x), 1.0, "x=-2^-{e}");
174        }
175    }
176
177    #[test]
178    fn sigmoid_of_the_special_values_and_the_edges() {
179        assert_eq!(sigmoid(0.0), 0.5);
180        assert_eq!(sigmoid(f32::INFINITY), 1.0);
181        assert_eq!(sigmoid(f32::NEG_INFINITY), 0.0);
182        assert!(sigmoid(f32::NAN).is_nan());
183        assert_eq!(sigmoid(20.0), 1.0);
184        // exp(-x) overflows to infinity here, so the result is 0 although
185        // the true value is subnormal (docs/semantics.md).
186        assert_eq!(sigmoid(-88.8).to_bits(), 0);
187        assert_eq!(sigmoid(-100.0).to_bits(), 0);
188        assert_eq!(sigmoid(-1.0e30).to_bits(), 0);
189    }
190
191    #[test]
192    fn sigmoid_matches_the_f64_formula_within_two_ulp() {
193        for i in -2000..=2000 {
194            let x = i as f32 * 0.01;
195            let want = (1.0 / (1.0 + (-f64::from(x)).exp())) as f32;
196            assert!(ulps(sigmoid(x), want) <= 2, "x={x}");
197        }
198    }
199}
oxfuse_runtime/math.rs

oxfuse_runtime/
math.rs
