thermite_special/specialized/generic/
digamma.rs

1use thermite::{
2    element::FloatElementWithBits,
3    mask::GenericMask,
4    math::{CoreMathWithPolicy as _, TranscendentalMathWithPolicy as _, policy::Policy, specialized::FlushDenormals},
5    prelude::*,
6};
7
8use crate::specialized::SpecializedSpecialMath;
9use crate::tables::gamma::Digamma;
10
11/// Shared digamma (`psi`) implementation for all real element types.
12///
13/// `psi(x) = d/dx ln(Gamma(x))`. The element-specific rational/asymptotic
14/// coefficients are passed in so the f32 and f64 specializations can share this
15/// body:
16///
17/// * `y` / `roots` / `p_12` / `q_12`: the `[1, 2]` rational `psi(x) = (x - root)(Y + R(x-1))`,
18///   where `root` is summed from `roots` via staged subtraction to preserve bits.
19/// * `p_large`: the `x >= 10` asymptotic expansion in `1/(x-1)^2`.
20#[inline(always)]
21pub fn digamma_impl<P, E, V, const NR: usize, const NL: usize, const NP: usize, const NQ: usize>(
22    x_in: V,
23    t: &Digamma<E, NR, NL, NP, NQ>,
24) -> V
25where
26    P: Policy,
27    E: FloatElementWithBits,
28    V: FloatVectorWithBits<Element = E> + SpecializedSpecialMath<E>,
29{
30    let mut x0 = x_in;
31
32    #[cfg(not(target_arch = "spirv"))]
33    if let Some(new_x) = FlushDenormals::<P>::flush_denormals([x0]) {
34        x0 = new_x[0];
35    }
36
37    let mut result = V::ZERO;
38    let mut x = x0;
39
40    // --- Reflection for x <= -1: psi(x) = psi(1-x) + pi*cot(pi*(1-x)) ---
41    let reflect = x0.cmp_le(V::NEG_ONE);
42    let mut refl_pole = GenericMask::FALSY;
43
44    if const { P::POLICY.avoid_branching } || reflect.any() {
45        let xr = V::ONE - x0; // 1 - x, >= 2 for reflected lanes
46        // fractional part shifted to (-1/2, 1/2] for tan argument reduction
47        let mut rem = xr - xr.floor();
48        rem = rem.sub_c(rem.cmp_gt(V::HALF), V::ONE);
49        // pi * cot(pi*rem) = pi / tan(pi*rem); tan_pi is accurate near the poles
50        let refl_term = V::PI / rem.tan_pi_p::<P>();
51        result = refl_term.zz(reflect); // result is still zero here
52        x = reflect.select(xr, x);
53        refl_pole = reflect & rem.is_zero(); // reflected negative integer is a pole
54    }
55
56    // Large lanes (x >= 10) use the asymptotic expansion directly; smaller lanes are
57    // reduced into [1, 2] via the recurrence psi(x) = psi(x+1) - 1/x.
58    let large = x.cmp_ge(V::splat(E::from_int(10)));
59
60    // Reduce into [1, 2]: lanes above 2 walk down (x -= 1, result += 1/x), lanes
61    // below 1 walk up (result -= 1/x, x += 1). The two directions are disjoint per
62    // lane, so one loop handles both with a single division per iteration.
63    let mut active = (x.cmp_gt(V::TWO) | x.cmp_lt(V::ONE)) & !large;
64    while active.any() {
65        V::_loop_hint();
66
67        // sign(x-1) is +1 above the interval (walk down: x -= 1, add +1/(x-1)) and
68        // -1 below it (walk up: x += 1, add -1/x). The reciprocal point is the smaller
69        // of {x, x-sign}: x-1 when walking down, x when walking up.
70        let sign = (x - V::ONE).signum();
71        let xs = x.sub_c(active, sign); // step toward [1, 2]; inactive lanes keep x
72        let term = sign * x.min(xs).approx_reciprocal_p::<P>();
73        result = result.add_c(active, term);
74        x = xs;
75        active = (x.cmp_gt(V::TWO) | x.cmp_lt(V::ONE)) & !large;
76    }
77
78    // x - 1 is shared by both the [1, 2] rational and the asymptotic expansion.
79    let xm1 = x - V::ONE;
80
81    // --- Rational approximation on [1, 2] (small lanes) ---
82    // staged subtraction preserves bits: root = sum(roots)
83    let mut g = x;
84    let mut i = 0;
85    while i < NR {
86        g -= V::splat(t.roots[i]);
87        i += 1;
88    }
89    let r = xm1.poly_n_p::<P, _>(&t.p_12) / xm1.poly_n_p::<P, _>(&t.q_12);
90    let rational = g * (V::splat(t.y) + r);
91
92    // --- Asymptotic expansion for x >= 10 (large lanes) ---
93    // ln(x-1) + 1/(2(x-1)) - z*P(z), with the trailing product fused into an FMA.
94    let z = (xm1 * xm1).approx_reciprocal_p::<P>();
95    let asymptotic = z.nmul_adde(
96        z.poly_n_p::<P, _>(&t.p_large),
97        xm1.ln_p::<P>() + (xm1 + xm1).approx_reciprocal_p::<P>(),
98    );
99
100    // both paths share the accumulated recurrence term
101    let mut res = result + large.select(asymptotic, rational);
102
103    // --- Poles: x == 0 and the negative integers -> NaN ---
104    if const { P::POLICY.check_overflow } {
105        let pole = x0.is_zero() | refl_pole;
106        res = pole.select(V::NAN, res);
107        res = x0.is_nan().select(V::NAN, res);
108    }
109
110    res
111}
thermite_special/specialized/generic/digamma.rs

thermite_special/specialized/generic/
digamma.rs
