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        return exp(x * x) * erfc(x);
44    }
45    // Six-correction asymptotic expansion of erfcx for large x. At x=26,
46    // the first omitted term is below 2e-17 relative to the leading term.
47    double inv  = 1.0 / x;
48    double inv2 = inv * inv;
49    double poly = 1.0
50                + inv2 * (-0.5
51                + inv2 * (0.75
52                + inv2 * (-1.875
53                + inv2 * (6.5625
54                + inv2 * (-29.53125
55                + inv2 * 162.421875)))));
56    const double inv_sqrt_pi = 0.5641895835477563; // 1/√π
57    return inv * poly * inv_sqrt_pi;
58}
59
60extern "C" __device__ __forceinline__ double log_ndtr(double x) {
61    if (isnan(x)) return x;
62    if (isinf(x)) return (x > 0.0) ? 0.0 : x;
63    if (x < 0.0) {
64        double u   = -x / SQRT_2;
65        double ex  = erfcx_nonnegative(u);
66        return -u * u + log(ex) - LN_2;
67    } else {
68        double upper_tail = 0.5 * erfc(x / SQRT_2);
69        return log1p(-upper_tail);
70    }
71}
72
73// Returns (log Φ(x), φ(x)/Φ(x)).
74extern "C" __device__ __forceinline__ void
75log_ndtr_and_mills(double x, double *log_cdf, double *lambda) {
76    if (isnan(x))          { *log_cdf = x;              *lambda = x;              return; }
77    if (isinf(x)) {
78        if (x > 0.0) { *log_cdf = 0.0; *lambda = 0.0; }
79        else         { *log_cdf = x;   *lambda = -x;  }
80        return;
81    }
82    if (x < 0.0) {
83        double u   = -x / SQRT_2;
84        double ex  = erfcx_nonnegative(u);
85        *log_cdf = -u * u + log(ex) - LN_2;
86        const double sqrt_2_over_pi = 0.7978845608028654; // √(2/π)
87        *lambda  = sqrt_2_over_pi / ex;
88    } else {
89        double upper_tail = 0.5 * erfc(x / SQRT_2);
90        double cdf = 1.0 - upper_tail;
91        double pdf = INV_SQRT_2PI * exp(-0.5 * x * x);
92        *log_cdf = log1p(-upper_tail);
93        *lambda  = pdf / cdf;
94    }
95}
96
97// Joint log Φ(x), Mills ratio, and positive negated log-CDF curvature
98// `-d² log Φ(x)/dx²`. The direct `lambda * (x + lambda)` spelling loses the
99// unit left-tail limit when x and lambda cancel, so the deep tail differentiates
100// the Laplace continued fraction used by the CPU kernel.
101extern "C" __device__ __forceinline__ void
102log_ndtr_mills_curvature(double x, double *log_cdf, double *lambda, double *curvature) {
103    log_ndtr_and_mills(x, log_cdf, lambda);
104    if (isnan(x)) { *curvature = x; return; }
105    if (isinf(x)) { *curvature = (x > 0.0) ? 0.0 : 1.0; return; }
106    if (x <= -4.0) {
107        double t = -x;
108        double q = 0.0;
109        double q_first = 0.0;
110        for (int n = 32; n >= 1; --n) {
111            double denominator = t + q;
112            double value = ((double)n) / denominator;
113            q_first = -value * (1.0 + q_first) / denominator;
114            q = value;
115        }
116        *curvature = 1.0 + q_first;
117    } else {
118        *curvature = *lambda * (x + *lambda);
119    }
120}
121
122#endif // PROBIT_NUMERICS_INCLUDED
123"#;
gam_gpu/numerics_device.rs

gam_gpu/
numerics_device.rs
