tract_linalg/generic/exp.rs
1#![allow(clippy::excessive_precision)]
2
3/// Coefficients of the degree-6 minimax fit of `e^r` over `[-ln 2 / 2, ln 2 / 2]`, highest
4/// power first.
5///
6/// Shared by every exp kernel whatever its width, like [`crate::generic::ln::POLY`]. The
7/// softmax kernels run the same reduction and fit over a domain of their own, where the
8/// argument is never positive and everything under `-103` is zero; nothing here may assume
9/// either.
10pub const POLY: [f32; 7] = [
11 1.383684405e-03,
12 8.374815793e-03,
13 4.166822560e-02,
14 1.666642017e-01,
15 4.999999208e-01,
16 1.000000036e+00,
17 1.000000001e+00,
18];
19
20pub const LOG2E: f32 = 1.442_695_04;
21
22/// `ln 2`, split so that `k * LN2_HI` is exact for every `k` the reduction produces and
23/// `LN2_LO` holds what it drops, leaving `|r|` under half an ulp of `ln 2 / 2`.
24pub const LN2_HI: f32 = 0.693_145_75;
25pub const LN2_LO: f32 = 1.428_606_8e-6;
26
27/// The input clamp. Above `HIGH` every result overflows f32 and below `LOW` every result
28/// rounds to zero, so clamping there is what the answer already is -- and it holds `k`
29/// inside the range [`SCALE_BIAS`] can build from a pair of valid exponent fields.
30pub const LOW: f32 = -104.0;
31pub const HIGH: f32 = 89.0;
32
33/// What each half of `2^k` adds to its exponent field. `k` is halved and rebuilt as two
34/// factors because a single `2^k` has no representation once `k` leaves `[-126, 127]`,
35/// which the subnormal and overflowing tails both need.
36pub const SCALE_BIAS: i32 = 127;
37
38/// The Cody-Waite rounding constant: `x * LOG2E + MAGIC - MAGIC` is that product rounded
39/// to an integer, without an integer conversion the polynomial's lane would have to leave.
40pub const MAGIC: f32 = 12_582_912.0;
41
42/// f32 exponential: a Cody-Waite reduction to `e^r * 2^k`, the [`POLY`] fit for `e^r`, and
43/// `2^k` rebuilt from a pair of exponent fields.
44///
45/// Within two ulp of a correctly rounded `exp` over `[LOW, HIGH]`, subnormal results
46/// included. Outside the clamp the answer is what the clamp gives: `+inf` above, zero
47/// below, so `±inf` answer `+inf` and zero. NaN propagates.
48pub fn sexp(x: f32) -> f32 {
49 let x = x.clamp(LOW, HIGH);
50 let kf = (x * LOG2E + MAGIC) - MAGIC;
51 let r = kf.mul_add(-LN2_LO, kf.mul_add(-LN2_HI, x));
52 let mut q = POLY[0];
53 for c in &POLY[1..] {
54 q = q.mul_add(r, *c);
55 }
56 let k = kf as i32;
57 let low = k >> 1;
58 let high = k - low;
59 let scale_low = f32::from_bits(((low + SCALE_BIAS) as u32) << 23);
60 let scale_high = f32::from_bits(((high + SCALE_BIAS) as u32) << 23);
61 q * scale_low * scale_high
62}
63
64routine_ew_rust!(generic;
65 f32,
66 generic_exp_f32_4n,
67 4,
68 4,
69 fn run(x: &mut [f32], _: ()) {
70 debug_assert!(x.len() % Self::nr() == 0);
71 debug_assert!(x.as_ptr() as usize % Self::alignment_bytes() == 0);
72 x.iter_mut().for_each(|px| *px = sexp(*px))
73 },
74 func(Exp)
75);