thermite_special/specialized/generic/
trigamma.rs

1use thermite::{
2    element::FloatElementWithBits,
3    mask::GenericMask,
4    math::{CoreMathWithPolicy as _, TranscendentalMathWithPolicy as _, policy::Policy},
5    prelude::*,
6};
7
8use crate::specialized::SpecializedSpecialMath;
9use crate::tables::gamma::Trigamma;
10
11/// Shared trigamma (`psi_1`) implementation for all real element types.
12///
13/// `psi_1(x) = d/dx psi(x)`, the second derivative of `ln Gamma`. Structurally much
14/// cheaper than `digamma`: the reduction to the fitted range is a *single* step
15/// rather than a masked walk, because the `[1, 2]` rational and the `x > 4` rational
16/// between them already cover everything from 1 upward.
17///
18/// The poles at zero and the negative integers evaluate to `+inf`, which is the
19/// correct two-sided limit (`psi_1` has a double pole there, so unlike `psi` the two
20/// one-sided limits agree) and falls out of the reflection term for free.
21#[inline(always)]
22pub fn trigamma_impl<P, E, V>(x_in: V, t: &Trigamma<E>) -> V
23where
24    P: Policy,
25    E: FloatElementWithBits,
26    V: FloatVectorWithBits<Element = E> + SpecializedSpecialMath<E>,
27{
28    let x0 = x_in.flush_denormals_p::<P>();
29
30    // --- Reflection for x <= 0: psi_1(x) = -psi_1(1 - x) + pi^2 / sin^2(pi*x) ---
31    let reflect = x0.cmp_le(V::ZERO);
32    let mut refl = V::ZERO;
33    let mut x = x0;
34
35    if const { P::POLICY.avoid_branching } || reflect.any() {
36        // Boost takes sin_pi of whichever of {x, 1 - x} is smaller in magnitude; on
37        // this path that is always x, since |x| < |1 - x| for every x <= 0.
38        let s = x0.sin_pi_p::<P>();
39
40        // sin_pi is exactly zero at the negative integers, so those lanes divide by
41        // zero and yield +inf - the intended pole value, no special case needed.
42        refl = (V::PI_SQUARED / (s * s)).zz(reflect);
43        x = reflect.select(V::ONE - x0, x);
44    }
45
46    // --- One recurrence step for 0 < x < 1: psi_1(x) = 1/x^2 + psi_1(x + 1) ---
47    // One step suffices: the rational below is fit from 1 up, and the reflected lanes
48    // already satisfy x >= 1.
49    let mut acc = V::ZERO;
50    let below_one = x.cmp_lt(V::ONE);
51
52    if const { P::POLICY.avoid_branching } || below_one.any() {
53        acc = (x * x).approx_reciprocal_p::<P>().zz(below_one);
54        x = x.add_c(below_one, V::ONE);
55    }
56
57    // --- Three minimax rational regions ---
58    let small = x.cmp_le(V::TWO);
59    let mid = x.cmp_le(V::splat(E::from_int(4))) & !small;
60    let large = !(small | mid);
61
62    // A single reciprocal covers both uses: it is the trailing 1/x^2 (small) or 1/x
63    // (mid, large) scale factor, and on the x > 2 lanes it is *also* the polynomial
64    // argument y = 1/x. Note the small lanes take 1/(x*x) directly rather than
65    // squaring a reciprocal, so they keep the accuracy of a single rounding.
66    let y = small.select(x * x, x).approx_reciprocal_p::<P>();
67
68    let mut num = V::EMPTY;
69    let mut den = V::EMPTY;
70    let mut base = V::ONE;
71
72    if const { P::POLICY.avoid_branching } || small.any() {
73        num = x.poly_n_p::<P, _>(&t.p_1_2);
74        den = x.poly_n_p::<P, _>(&t.q_1_2);
75        base = small.select(V::splat(t.offset), V::ONE);
76    }
77
78    if const { P::POLICY.avoid_branching } || mid.any() {
79        num = mid.select(y.poly_n_p::<P, _>(&t.p_2_4), num);
80        den = mid.select(y.poly_n_p::<P, _>(&t.q_2_4), den);
81    }
82
83    if const { P::POLICY.avoid_branching } || large.any() {
84        num = large.select(y.poly_n_p::<P, _>(&t.p_4_inf), num);
85        den = large.select(y.poly_n_p::<P, _>(&t.q_4_inf), den);
86    }
87
88    // Selecting the numerator and denominator before dividing keeps this to one
89    // division for all three regions - and means an unselected region's overflow can
90    // never turn into an inf/inf NaN in a lane that survives.
91    let main = (base + num / den) * y;
92
93    reflect.select(refl - main, acc + main)
94}
thermite_special/specialized/generic/trigamma.rs

thermite_special/specialized/generic/
trigamma.rs
