gam_gpu/
numerics_device.rs

1//! Shared device-side probit numerics for NVRTC kernels.
2//!
3//! [`PROBIT_NUMERICS_CU`] is prepended to every NVRTC kernel source that needs
4//! stable probit/normal-CDF arithmetic.  Keeping one copy here means a
5//! numerics fix is a one-line change instead of a coordination problem across
6//! multiple kernel source strings.
7//!
8//! Covered device functions (all `__device__ __forceinline__`, double precision):
9//!   - `erfcx_nonnegative(x)`   — scaled complementary error function for x ≥ 0
10//!   - `log_ndtr(x)`            — log Φ(x), numerically stable in the deep left tail
11//!   - `log_ndtr_and_mills(x, *log_cdf, *lambda)` — joint (log Φ(x), φ(x)/Φ(x))
12//!   - `log_ndtr_mills_curvature(...)` — also returns `−d² log Φ(x)/dx²`
13
14/// Device-side probit numerics injected at the top of every NVRTC kernel that
15/// needs them.  Prepend this string to a kernel-specific body before passing to
16/// `cudarc::nvrtc::compile_ptx` or `PtxModuleCache::get_or_compile`.
17pub const PROBIT_NUMERICS_CU: &str = r#"
18// -------- shared probit numerics -----------------------------------------
19// All math in double precision; fast-math is disabled at compile time
20// (see `device_cache`'s `--fmad=false`) and the source is kept free of any
21// fast-math / single-precision intrinsic, guarded by the numerics_host tests.
22//
23// `log_ndtr(x)` = log Φ(x).  For x < 0 uses the erfcx representation
24//   log Φ(x) = -u² + log(½ · erfcx(u)),   u = -x / √2
25// which preserves digits all the way into the deep left tail (matches
26// the CPU `normal_logcdf`).  For x ≥ 0 falls back to log1p(-½·erfc(x/√2)).
27//
28// `log_ndtr_and_mills(x, *log_cdf, *lambda)` returns both log Φ(x) and the
29// Mills ratio φ(x)/Φ(x) in a single pass.  For x < 0 the erfcx path keeps
30// the ratio stable even when Φ(x) underflows to zero.
31
32#ifndef PROBIT_NUMERICS_INCLUDED
33#define PROBIT_NUMERICS_INCLUDED
34
35#define INV_SQRT_2PI 0.3989422804014327
36#define SQRT_2       1.4142135623730951
37#define LN_2         0.6931471805599453
38
39extern "C" __device__ __forceinline__ double erfcx_nonnegative(double x) {
40    if (isnan(x) || x < 0.0) return nan("");
41    if (isinf(x)) return 0.0;
42    if (x < 26.0) {
43        // Carry x*x EXACTLY into exp. Rounding the square perturbs it by a
44        // relative eps/2, and exp converts a relative perturbation of its
45        // ARGUMENT into x^2 times that in its RESULT -- 5.7e-14 at the top of
46        // this branch, against the 3e-16 the asymptotic branch below already
47        // delivers, so the crossover at 26 was a 190x step DOWN in error into
48        // the interval every probit consumer lives in. `lo` is the whole of the
49        // discarded term and is itself exactly representable, so
50        // exp(x^2) = exp(hi)*exp(lo) with exp(lo) = 1 + lo to within 1e-27.
51        //
52        // `fma` is the IEEE fused operation (`fma.rn.f64`, one instruction).
53        // `--fmad=false` disables CONTRACTION of a separate `a*b+c`, not an
54        // explicit `fma` call, so this stays operation-for-operation identical
55        // to the host oracle's `f64::mul_add`.
56        double hi   = x * x;
57        double lo   = fma(x, x, -hi);
58        double head = exp(hi) * erfc(x);
59        return fma(head, lo, head);
60    }
61    // Six-correction asymptotic expansion of erfcx for large x. At x=26,
62    // the first omitted term is below 2e-17 relative to the leading term.
63    double inv  = 1.0 / x;
64    double inv2 = inv * inv;
65    double poly = 1.0
66                + inv2 * (-0.5
67                + inv2 * (0.75
68                + inv2 * (-1.875
69                + inv2 * (6.5625
70                + inv2 * (-29.53125
71                + inv2 * 162.421875)))));
72    const double inv_sqrt_pi = 0.5641895835477563; // 1/√π
73    return inv * poly * inv_sqrt_pi;
74}
75
76extern "C" __device__ __forceinline__ double log_ndtr(double x) {
77    if (isnan(x)) return x;
78    if (isinf(x)) return (x > 0.0) ? 0.0 : x;
79    if (x < 0.0) {
80        double u   = -x / SQRT_2;
81        double ex  = erfcx_nonnegative(u);
82        return -u * u + log(ex) - LN_2;
83    } else {
84        double upper_tail = 0.5 * erfc(x / SQRT_2);
85        return log1p(-upper_tail);
86    }
87}
88
89// Returns (log Φ(x), φ(x)/Φ(x)).
90extern "C" __device__ __forceinline__ void
91log_ndtr_and_mills(double x, double *log_cdf, double *lambda) {
92    if (isnan(x))          { *log_cdf = x;              *lambda = x;              return; }
93    if (isinf(x)) {
94        if (x > 0.0) { *log_cdf = 0.0; *lambda = 0.0; }
95        else         { *log_cdf = x;   *lambda = -x;  }
96        return;
97    }
98    if (x < 0.0) {
99        double u   = -x / SQRT_2;
100        double ex  = erfcx_nonnegative(u);
101        *log_cdf = -u * u + log(ex) - LN_2;
102        const double sqrt_2_over_pi = 0.7978845608028654; // √(2/π)
103        *lambda  = sqrt_2_over_pi / ex;
104    } else {
105        double upper_tail = 0.5 * erfc(x / SQRT_2);
106        double cdf = 1.0 - upper_tail;
107        // Same exact-square correction as `erfcx_nonnegative`: `exp(-0.5*x*x)`
108        // otherwise carries x^2*eps/2 relative error. `x` is finite and
109        // non-negative here (the isnan/isinf guards above returned), so the
110        // residual is always a finite number and needs no further defence.
111        double xx  = x * x;
112        double pdf = INV_SQRT_2PI * exp(-0.5 * xx);
113        pdf = fma(pdf, -0.5 * fma(x, x, -xx), pdf);
114        *log_cdf = log1p(-upper_tail);
115        *lambda  = pdf / cdf;
116    }
117}
118
119// Joint log Φ(x), Mills ratio, and positive negated log-CDF curvature
120// `-d² log Φ(x)/dx²`. The direct `lambda * (x + lambda)` spelling loses the
121// unit left-tail limit when x and lambda cancel, so the deep tail differentiates
122// the Laplace continued fraction used by the CPU kernel.
123extern "C" __device__ __forceinline__ void
124log_ndtr_mills_curvature(double x, double *log_cdf, double *lambda, double *curvature) {
125    log_ndtr_and_mills(x, log_cdf, lambda);
126    if (isnan(x)) { *curvature = x; return; }
127    if (isinf(x)) { *curvature = (x > 0.0) ? 0.0 : 1.0; return; }
128    if (x <= -4.0) {
129        double t = -x;
130        double q = 0.0;
131        double q_first = 0.0;
132        for (int n = 32; n >= 1; --n) {
133            double denominator = t + q;
134            double value = ((double)n) / denominator;
135            q_first = -value * (1.0 + q_first) / denominator;
136            q = value;
137        }
138        *curvature = 1.0 + q_first;
139    } else {
140        *curvature = *lambda * (x + *lambda);
141    }
142}
143
144#endif // PROBIT_NUMERICS_INCLUDED
145"#;
gam_gpu/numerics_device.rs

gam_gpu/
numerics_device.rs
