Skip to main content

gam_solve/gpu_kernels/
pirls_row.rs

1//! Generic GPU PIRLS row-reweight primitives.
2//!
3//! Stage 1 of the device-resident PIRLS port: for every row `i` and a fixed
4//! exponential-family `(family, link)` pair, evaluate the working IRLS state
5//! on the GPU and write it back into device-resident output buffers. The
6//! Hessian/gradient assembly that consumes those buffers (Stage 2) and the
7//! full device-resident PIRLS iteration loop (Stage 3) plug in on top of this
8//! contract without touching the per-family math.
9//!
10//! ## Output contract (per row `i`)
11//!
12//! | Field        | Meaning                                                    |
13//! |--------------|------------------------------------------------------------|
14//! | `mu`         | Inverse-link mean μ_i = g⁻¹(η_i).                          |
15//! | `grad_eta`   | Score wrt η: ∂ℓ/∂η_i = wᵢ·(yᵢ − μᵢ)·dη/dV for canonical    |
16//! |              | links; equals priorweight·(yᵢ − μᵢ)·h'(ηᵢ)/V(μᵢ) for non-  |
17//! |              | canonical Bernoulli; equals priorweight·(yᵢ − μᵢ) for      |
18//! |              | Gaussian/Poisson and priorweight·shape·(yᵢ/μᵢ−1) for Gamma.|
19//! | `w_fisher`   | Fisher expected weight (priorweight · h'(η)² / V(μ)).      |
20//! |              | Used for inference (Var(β̂)).                              |
21//! | `w_hessian`  | Curvature weight for the Newton/Laplace Hessian.            |
22//! |              | == w_fisher on canonical links; observed correction on     |
23//! |              | non-canonical Bernoulli + Gamma-log (Stage 5 populates).   |
24//! | `w_solver`   | Exact `w_hessian` consumed by matrix assembly. Statistical  |
25//! |              | rows are never floored or sign-projected; conditioning is  |
26//! |              | exclusively an assembled-matrix ridge operation.           |
27//! | `deviance`   | Per-row deviance contribution dᵢ; sum aggregated host-side  |
28//! |              | for line search and convergence checks.                    |
29//! | `status`     | Zero on success; otherwise a refusal code. A failed thread  |
30//! |              | writes no numerical outputs.                               |
31//!
32//! Two-weight discipline is structural in the contract: the gradient is
33//! emitted **directly** and never reconstructed from an artificial working
34//! response, because that product is both unnecessary and unstable in tails.
35//!
36//! ## Per-family source layout
37//!
38//! Each `(family, link, curvature_mode)` triple has its own specialised CUDA
39//! source compiled to a dedicated module and cached. The kernel never branches
40//! on a runtime `family` enum — that pattern collapses ILP and forces the
41//! compiler to keep dead paths warm. The module cache is keyed by
42//! `(family_id, curvature_mode, precision)` so a single process compiles each
43//! kernel exactly once across all fits.
44
45use std::sync::OnceLock;
46
47use gam_gpu::gpu_error::GpuError;
48#[cfg(target_os = "linux")]
49use gam_gpu::gpu_error::GpuResultExt;
50use gam_math::special::{bd0, bernoulli_kl_from_logits, softplus};
51use gam_problem::EstimationError;
52
53#[cfg(target_os = "linux")]
54use std::sync::{Arc, Mutex};
55
56#[cfg(target_os = "linux")]
57use cudarc::driver::{CudaContext, CudaModule};
58
59// ────────────────────────────────────────────────────────────────────────
60// Public selectors
61// ────────────────────────────────────────────────────────────────────────
62
63/// Which built-in `(response, link)` PIRLS family the row kernel evaluates.
64///
65/// One enum value ↔ one specialised CUDA source ↔ one cached module. Custom
66/// families come in Stage 6 via NVRTC JIT (Level A / Level B) and reuse the
67/// same host harness.
68#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
69pub enum PirlsRowFamily {
70    BernoulliLogit,
71    BernoulliProbit,
72    BernoulliCLogLog,
73    PoissonLog,
74    GaussianIdentity,
75    GammaLog,
76}
77
78impl PirlsRowFamily {
79    pub const ALL: [Self; 6] = [
80        Self::BernoulliLogit,
81        Self::BernoulliProbit,
82        Self::BernoulliCLogLog,
83        Self::PoissonLog,
84        Self::GaussianIdentity,
85        Self::GammaLog,
86    ];
87
88    pub const fn as_str(self) -> &'static str {
89        match self {
90            Self::BernoulliLogit => "bernoulli-logit",
91            Self::BernoulliProbit => "bernoulli-probit",
92            Self::BernoulliCLogLog => "bernoulli-cloglog",
93            Self::PoissonLog => "poisson-log",
94            Self::GaussianIdentity => "gaussian-identity",
95            Self::GammaLog => "gamma-log",
96        }
97    }
98
99    /// CUDA `extern "C"` entry symbol for this family's full (final-row) kernel.
100    pub const fn kernel_name(self) -> &'static str {
101        match self {
102            Self::BernoulliLogit => "pirls_row_bernoulli_logit",
103            Self::BernoulliProbit => "pirls_row_bernoulli_probit",
104            Self::BernoulliCLogLog => "pirls_row_bernoulli_cloglog",
105            Self::PoissonLog => "pirls_row_poisson_log",
106            Self::GaussianIdentity => "pirls_row_gaussian_identity",
107            Self::GammaLog => "pirls_row_gamma_log",
108        }
109    }
110
111    /// CUDA `extern "C"` entry symbol for this family's solve-row kernel
112    /// (writes only `grad_eta`, `w_solver`, `deviance`, `status`).
113    pub const fn solve_kernel_name(self) -> &'static str {
114        match self {
115            Self::BernoulliLogit => "pirls_solve_bernoulli_logit",
116            Self::BernoulliProbit => "pirls_solve_bernoulli_probit",
117            Self::BernoulliCLogLog => "pirls_solve_bernoulli_cloglog",
118            Self::PoissonLog => "pirls_solve_poisson_log",
119            Self::GaussianIdentity => "pirls_solve_gaussian_identity",
120            Self::GammaLog => "pirls_solve_gamma_log",
121        }
122    }
123
124    /// CUDA `extern "C"` entry symbol for this family's alpha-ladder kernel
125    /// (evaluates all step sizes in a single launch, outputs `objective[]` and
126    /// `status[]` per alpha slot).
127    pub const fn ladder_kernel_name(self) -> &'static str {
128        match self {
129            Self::BernoulliLogit => "pirls_ladder_bernoulli_logit",
130            Self::BernoulliProbit => "pirls_ladder_bernoulli_probit",
131            Self::BernoulliCLogLog => "pirls_ladder_bernoulli_cloglog",
132            Self::PoissonLog => "pirls_ladder_poisson_log",
133            Self::GaussianIdentity => "pirls_ladder_gaussian_identity",
134            Self::GammaLog => "pirls_ladder_gamma_log",
135        }
136    }
137}
138
139/// Curvature surface used to populate `w_hessian` / `w_solver`.
140///
141/// `Fisher` is the default and matches the CPU Stage-1 path bit-for-bit. The
142/// `Observed` mode is populated by Stage 5 for non-canonical Bernoulli and
143/// Gamma-log fits where the negative-log-likelihood Hessian uses the observed
144/// information surface instead of the Fisher expected information.
145#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
146pub enum CurvatureMode {
147    Fisher,
148    Observed,
149}
150
151impl CurvatureMode {
152    pub const fn as_str(self) -> &'static str {
153        match self {
154            Self::Fisher => "fisher",
155            Self::Observed => "observed",
156        }
157    }
158}
159
160/// Device refusal codes.  Zero is the only success value; a non-zero code
161/// means that the row kernel wrote no numerical outputs for that row.  The
162/// host deterministically selects the smallest failing row and replays the
163/// CPU oracle to recover the full typed [`EstimationError`] (including the
164/// exact offending value).
165pub mod status_codes {
166    pub const OK: u32 = 0;
167    pub const ETA_DOMAIN: u32 = 1;
168    pub const PRIOR_WEIGHT: u32 = 2;
169    pub const RESPONSE: u32 = 3;
170    pub const GAMMA_SHAPE: u32 = 4;
171    pub const INVERSE_LINK: u32 = 5;
172    pub const FISHER_WEIGHT: u32 = 6;
173    pub const OBSERVED_WEIGHT: u32 = 7;
174    pub const GRADIENT: u32 = 8;
175    pub const DEVIANCE: u32 = 9;
176    pub const FINAL_OUTPUT: u32 = 10;
177
178    pub const fn quantity(code: u32) -> &'static str {
179        match code {
180            ETA_DOMAIN => "inverse-link eta domain",
181            PRIOR_WEIGHT => "prior weight",
182            RESPONSE => "response",
183            GAMMA_SHAPE => "Gamma shape",
184            INVERSE_LINK => "inverse-link jet",
185            FISHER_WEIGHT => "Fisher weight",
186            OBSERVED_WEIGHT => "observed Hessian weight",
187            GRADIENT => "eta gradient",
188            DEVIANCE => "deviance contribution",
189            FINAL_OUTPUT => "final row output",
190            _ => "unknown GPU PIRLS refusal",
191        }
192    }
193}
194
195// ────────────────────────────────────────────────────────────────────────
196// Reference CPU evaluator (parity gate against the GPU kernel).
197//
198// These functions reproduce, byte-for-byte in f64, the formulas in
199// `src/solver/pirls.rs`'s `update_glmvectors` / `write_poisson_log_working_state`
200// / `write_gamma_log_working_state` / `write_identityworking_state`. Stage 1
201// parity tests compare the GPU buffers to these on the V100; mac builds
202// exercise only the CPU reference (the GPU launcher returns
203// `DriverLibraryUnavailable` without a CUDA runtime).
204// ────────────────────────────────────────────────────────────────────────
205
206/// Per-row inputs in scalar form.
207#[derive(Clone, Copy, Debug)]
208pub struct RowInput {
209    pub eta: f64,
210    pub y: f64,
211    pub prior_weight: f64,
212}
213
214/// Per-row outputs matching the GPU kernel contract.
215#[derive(Clone, Copy, Debug, Default)]
216pub struct RowOutput {
217    pub mu: f64,
218    pub grad_eta: f64,
219    pub w_fisher: f64,
220    pub w_hessian: f64,
221    pub w_solver: f64,
222    pub deviance: f64,
223}
224
225/// Indexed form of `row_reweight_cpu`, used to reproduce a device refusal
226/// with the correct row in its typed error.
227pub fn row_reweight_cpu_at(
228    row: usize,
229    family: PirlsRowFamily,
230    mode: CurvatureMode,
231    input: RowInput,
232    gamma_shape: f64,
233) -> Result<RowOutput, EstimationError> {
234    match family {
235        PirlsRowFamily::GaussianIdentity => row_gaussian_identity(row, input, mode),
236        PirlsRowFamily::PoissonLog => row_poisson_log(row, input, mode),
237        PirlsRowFamily::GammaLog => row_gamma_log(row, input, mode, gamma_shape),
238        PirlsRowFamily::BernoulliLogit => row_bernoulli_logit(row, input, mode),
239        PirlsRowFamily::BernoulliProbit => row_bernoulli_probit(row, input, mode),
240        PirlsRowFamily::BernoulliCLogLog => row_bernoulli_cloglog(row, input, mode),
241    }
242}
243
244/// Recover the typed error represented by a device status vector.  Device
245/// threads write one code per row, so scanning in index order makes concurrent
246/// failures deterministic.  The scalar CPU replay supplies the exact
247/// quantity/value payload without expanding the hot GPU ABI.
248pub fn replay_first_refusal(
249    family: PirlsRowFamily,
250    mode: CurvatureMode,
251    gamma_shape: f64,
252    eta: &[f64],
253    y: &[f64],
254    prior_weight: &[f64],
255    status: &[u32],
256) -> Result<(), EstimationError> {
257    let n = eta.len();
258    if y.len() != n || prior_weight.len() != n || status.len() != n {
259        return Err(EstimationError::InvalidInput(format!(
260            "GPU PIRLS refusal replay length mismatch: eta={n}, y={}, prior_weight={}, status={}",
261            y.len(),
262            prior_weight.len(),
263            status.len(),
264        )));
265    }
266    let Some((row, &code)) = status
267        .iter()
268        .enumerate()
269        .find(|(_, code)| **code != status_codes::OK)
270    else {
271        return Ok(());
272    };
273    let input = RowInput {
274        eta: eta[row],
275        y: y[row],
276        prior_weight: prior_weight[row],
277    };
278    match row_reweight_cpu_at(row, family, mode, input, gamma_shape) {
279        Err(error) => Err(error),
280        Ok(_) => Err(EstimationError::pirls_row_geometry_unrepresentable(
281            row,
282            status_codes::quantity(code),
283            input.eta,
284            f64::from(code),
285        )),
286    }
287}
288
289/// Resolve `(w_fisher, observed_correction)` into the `w_hessian` value that
290/// matches the selected curvature surface. Stage 1 returns `w_fisher` for both
291/// modes (parity with the CPU PIRLS path that, today, uses Fisher weights
292/// even for non-canonical links); Stage 5 will switch the `Observed` arm to
293/// `w_fisher + observed_correction` and the call sites stay unchanged.
294#[inline]
295fn select_w_hessian(mode: CurvatureMode, w_fisher: f64, observed_correction: f64) -> f64 {
296    match mode {
297        CurvatureMode::Fisher => w_fisher,
298        CurvatureMode::Observed => w_fisher + observed_correction,
299    }
300}
301
302
303#[inline]
304fn finite_eta(link: &'static str, eta: f64) -> Result<(), EstimationError> {
305    if eta.is_finite() {
306        Ok(())
307    } else {
308        Err(EstimationError::InverseLinkDomainViolation {
309            link,
310            eta,
311            lower: -f64::MAX,
312            upper: f64::MAX,
313        })
314    }
315}
316
317#[inline]
318fn prior_weight(row: usize, input: RowInput) -> Result<f64, EstimationError> {
319    if input.prior_weight.is_finite() && input.prior_weight >= 0.0 {
320        Ok(input.prior_weight)
321    } else {
322        Err(EstimationError::pirls_row_geometry_unrepresentable(
323            row,
324            "prior weight",
325            input.eta,
326            input.prior_weight,
327        ))
328    }
329}
330
331#[inline]
332fn certify_output(row: usize, eta: f64, output: RowOutput) -> Result<RowOutput, EstimationError> {
333    for (quantity, value) in [
334        ("mean", output.mu),
335        ("eta gradient", output.grad_eta),
336        ("Fisher weight", output.w_fisher),
337        ("observed Hessian weight", output.w_hessian),
338        ("solver Hessian weight", output.w_solver),
339        ("deviance contribution", output.deviance),
340    ] {
341        if !value.is_finite() {
342            return Err(EstimationError::pirls_row_geometry_unrepresentable(row, quantity, eta, value));
343        }
344    }
345    Ok(output)
346}
347
348/// Evaluate `a*b/c` while avoiding a false overflow/underflow caused solely by
349/// operation order.  At least one of the product-first or quotient-first forms
350/// is normally representable whenever the final positive f64 is; all three are
351/// tried in a fixed order so CPU/device refusal and rounding stay deterministic.
352#[inline]
353fn positive_mul_div(a: f64, b: f64, c: f64) -> f64 {
354    let product = a * b;
355    if product.is_finite() && product > 0.0 {
356        let value = product / c;
357        if value.is_finite() && value > 0.0 {
358            return value;
359        }
360    }
361    let quotient_a = a / c;
362    if quotient_a.is_finite() && quotient_a > 0.0 {
363        let value = quotient_a * b;
364        if value.is_finite() && value > 0.0 {
365            return value;
366        }
367    }
368    let quotient_b = b / c;
369    if quotient_b.is_finite() && quotient_b > 0.0 {
370        let value = quotient_b * a;
371        if value.is_finite() && value > 0.0 {
372            return value;
373        }
374    }
375    product / c
376}
377
378/// `u - log1p(u)` without cancellation around zero.
379#[inline]
380fn gamma_unit_deviance_near_one(u: f64) -> f64 {
381    if u.abs() > 0.125 {
382        return u - u.ln_1p();
383    }
384    let mut power = u * u;
385    let mut sum = 0.5 * power;
386    for degree in 3..=32 {
387        power *= u;
388        let term = power / f64::from(degree);
389        let next = if degree % 2 == 0 {
390            sum + term
391        } else {
392            sum - term
393        };
394        if next == sum {
395            break;
396        }
397        sum = next;
398    }
399    sum
400}
401
402/// `(1+u)log1p(u)-u` without cancellation around zero.
403#[inline]
404fn poisson_unit_deviance_near_one(u: f64) -> f64 {
405    if u.abs() > 0.125 {
406        return (1.0 + u) * u.ln_1p() - u;
407    }
408    let mut power = u * u;
409    let mut sum = 0.5 * power;
410    for degree in 3..=32 {
411        power *= u;
412        let coefficient =
413            if degree % 2 == 0 { 1.0 } else { -1.0 } / (f64::from(degree) * f64::from(degree - 1));
414        let next = sum + coefficient * power;
415        if next == sum {
416            break;
417        }
418        sum = next;
419    }
420    sum
421}
422
423#[inline]
424fn row_gaussian_identity(
425    row: usize,
426    input: RowInput,
427    mode: CurvatureMode,
428) -> Result<RowOutput, EstimationError> {
429    finite_eta("standard identity inverse link", input.eta)?;
430    let w = prior_weight(row, input)?;
431    let mu = input.eta;
432    if w > 0.0 && !input.y.is_finite() {
433        return Err(EstimationError::pirls_row_geometry_unrepresentable(row, "Gaussian response", input.eta, input.y));
434    }
435    let resid = input.y - mu;
436    let (grad_eta, dev) = if w == 0.0 {
437        (0.0, 0.0)
438    } else {
439        (w * resid, w * resid * resid)
440    };
441    let w_hessian = select_w_hessian(mode, w, 0.0);
442    certify_output(
443        row,
444        input.eta,
445        RowOutput {
446            mu,
447            grad_eta,
448            w_fisher: w,
449            w_hessian,
450            w_solver: w_hessian,
451            deviance: dev,
452        },
453    )
454}
455
456#[inline]
457fn row_poisson_log(
458    row: usize,
459    input: RowInput,
460    mode: CurvatureMode,
461) -> Result<RowOutput, EstimationError> {
462    let mu = crate::mixture_link::log_link_solver_exp(input.eta)?;
463    let w_prior = prior_weight(row, input)?;
464    if w_prior > 0.0 && !(input.y.is_finite() && input.y >= 0.0) {
465        return Err(EstimationError::pirls_row_geometry_unrepresentable(row, "Poisson response", input.eta, input.y));
466    }
467    if w_prior == 0.0 {
468        return certify_output(
469            row,
470            input.eta,
471            RowOutput {
472                mu,
473                ..RowOutput::default()
474            },
475        );
476    }
477    let w_fisher = w_prior * mu;
478    if !(w_fisher.is_finite() && w_fisher > 0.0) {
479        return Err(EstimationError::pirls_row_geometry_unrepresentable(row, "Poisson Fisher weight", input.eta, w_fisher));
480    }
481    let grad_eta = w_prior * (input.y - mu);
482    let u = (input.y - mu) / mu;
483    let dev_base = if input.y == 0.0 {
484        w_fisher
485    } else {
486        // Accurate around saturation. In either far tail this dimensionless
487        // ratio can become non-finite before multiplication by a tiny weight;
488        // only then switch to the algebraically identical absolute-coordinate
489        // expression, whose products are balanced independently.
490        let scaled_unit = w_fisher * poisson_unit_deviance_near_one(u);
491        if scaled_unit.is_finite() && scaled_unit >= 0.0 {
492            scaled_unit
493        } else {
494            let weighted_y = positive_mul_div(w_fisher, input.y, mu);
495            weighted_y * (input.y.ln() - input.eta - 1.0) + w_fisher
496        }
497    };
498    let dev = 2.0 * dev_base;
499    let w_hessian = select_w_hessian(mode, w_fisher, 0.0);
500    certify_output(
501        row,
502        input.eta,
503        RowOutput {
504            mu,
505            grad_eta,
506            w_fisher,
507            w_hessian,
508            w_solver: w_hessian,
509            deviance: dev,
510        },
511    )
512}
513
514#[inline]
515fn row_gamma_log(
516    row: usize,
517    input: RowInput,
518    mode: CurvatureMode,
519    shape: f64,
520) -> Result<RowOutput, EstimationError> {
521    let mu = crate::mixture_link::log_link_solver_exp(input.eta)?;
522    if !(shape.is_finite() && shape > 0.0) {
523        return Err(EstimationError::pirls_row_geometry_unrepresentable(row, "Gamma shape", input.eta, shape));
524    }
525    let w_prior = prior_weight(row, input)?;
526    if w_prior > 0.0 && !(input.y.is_finite() && input.y > 0.0) {
527        return Err(EstimationError::pirls_row_geometry_unrepresentable(row, "Gamma response", input.eta, input.y));
528    }
529    if w_prior == 0.0 {
530        return certify_output(
531            row,
532            input.eta,
533            RowOutput {
534                mu,
535                ..RowOutput::default()
536            },
537        );
538    }
539    let w_fisher = w_prior * shape;
540    if !(w_fisher.is_finite() && w_fisher > 0.0) {
541        return Err(EstimationError::pirls_row_geometry_unrepresentable(row, "Gamma Fisher weight", input.eta, w_fisher));
542    }
543    let observed_ratio = match mode {
544        CurvatureMode::Fisher => None,
545        CurvatureMode::Observed => {
546            // Ratio-first: at `y == mu` the observed weight must reproduce the
547            // Fisher weight EXACTLY (`y/mu` is bit-for-bit 1.0), which the
548            // product-first reordering destroys by one ulp. The reordering
549            // fallback remains for tails where the direct form is not
550            // representable.
551            let direct = w_fisher * (input.y / mu);
552            let weighted_ratio = if direct.is_finite() && direct > 0.0 {
553                direct
554            } else {
555                positive_mul_div(w_fisher, input.y, mu)
556            };
557            if !(weighted_ratio.is_finite() && weighted_ratio > 0.0) {
558                return Err(EstimationError::pirls_row_geometry_unrepresentable(
559                    row,
560                    "Gamma observed Hessian weight",
561                    input.eta,
562                    weighted_ratio,
563                ));
564            }
565            Some(weighted_ratio)
566        }
567    };
568    let w_hessian = observed_ratio.unwrap_or(w_fisher);
569    if !w_hessian.is_finite() {
570        return Err(EstimationError::pirls_row_geometry_unrepresentable(
571            row,
572            "Gamma observed Hessian weight",
573            input.eta,
574            w_hessian,
575        ));
576    }
577    let u = (input.y - mu) / mu;
578    // `u` rounds to exactly -1 when y/mu is a representable but very small
579    // ratio, and the dimensionless expression can overflow for the opposite
580    // tail. Preserve the local-series path whenever it succeeds, then use the
581    // weighted absolute-coordinate identity for those two tail cases.
582    let scaled_unit = w_fisher * gamma_unit_deviance_near_one(u);
583    let need_weighted_ratio = !u.is_finite() || !(scaled_unit.is_finite() && scaled_unit >= 0.0);
584    let weighted_ratio = if need_weighted_ratio {
585        observed_ratio.unwrap_or_else(|| positive_mul_div(w_fisher, input.y, mu))
586    } else {
587        0.0
588    };
589    let grad_eta = if u.is_finite() {
590        w_fisher * u
591    } else {
592        weighted_ratio - w_fisher
593    };
594    let dev_base = if scaled_unit.is_finite() && scaled_unit >= 0.0 {
595        scaled_unit
596    } else {
597        weighted_ratio - w_fisher * (1.0 + input.y.ln() - input.eta)
598    };
599    let dev = 2.0 * dev_base;
600    certify_output(
601        row,
602        input.eta,
603        RowOutput {
604            mu,
605            grad_eta,
606            w_fisher,
607            w_hessian,
608            w_solver: w_hessian,
609            deviance: dev,
610        },
611    )
612}
613
614#[inline]
615fn bernoulli_response(row: usize, input: RowInput, w: f64) -> Result<(), EstimationError> {
616    if w == 0.0 || (input.y.is_finite() && (0.0..=1.0).contains(&input.y)) {
617        Ok(())
618    } else {
619        Err(EstimationError::pirls_row_geometry_unrepresentable(row, "binomial response", input.eta, input.y))
620    }
621}
622
623#[inline]
624fn row_bernoulli_logit(
625    row: usize,
626    input: RowInput,
627    mode: CurvatureMode,
628) -> Result<RowOutput, EstimationError> {
629    finite_eta("standard logit inverse link", input.eta)?;
630    let w_prior = prior_weight(row, input)?;
631    bernoulli_response(row, input, w_prior)?;
632    let tail = (-input.eta.abs()).exp();
633    let denom = 1.0 + tail;
634    let (mu, residual) = if input.eta >= 0.0 {
635        let one_minus_mu = tail / denom;
636        let residual = if input.y == 1.0 {
637            one_minus_mu
638        } else {
639            (input.y - 1.0) + one_minus_mu
640        };
641        (1.0 / denom, residual)
642    } else {
643        let mu = tail / denom;
644        (mu, input.y - mu)
645    };
646    let dmu_deta = tail / (denom * denom);
647    if !(dmu_deta.is_finite() && dmu_deta > 0.0) {
648        return Err(EstimationError::pirls_row_geometry_unrepresentable(
649            row,
650            "canonical-logit inverse-link jet",
651            input.eta,
652            dmu_deta,
653        ));
654    }
655    if w_prior == 0.0 {
656        return certify_output(
657            row,
658            input.eta,
659            RowOutput {
660                mu,
661                ..RowOutput::default()
662            },
663        );
664    }
665    let w_fisher = w_prior * dmu_deta;
666    if !(w_fisher.is_finite() && w_fisher > 0.0) {
667        return Err(EstimationError::pirls_row_geometry_unrepresentable(row, "logit Fisher weight", input.eta, w_fisher));
668    }
669    let grad_eta = w_prior * residual;
670    let dev = bernoulli_logit_deviance(input.y, input.eta, w_prior);
671    let w_hessian = select_w_hessian(mode, w_fisher, 0.0);
672    certify_output(
673        row,
674        input.eta,
675        RowOutput {
676            mu,
677            grad_eta,
678            w_fisher,
679            w_hessian,
680            w_solver: w_hessian,
681            deviance: dev,
682        },
683    )
684}
685
686#[inline]
687fn row_bernoulli_probit(
688    row: usize,
689    input: RowInput,
690    mode: CurvatureMode,
691) -> Result<RowOutput, EstimationError> {
692    finite_eta("standard probit inverse link", input.eta)?;
693    let d1 = standard_normal_pdf(input.eta);
694    row_bernoulli_noncanonical(
695        row,
696        input,
697        mode,
698        standard_normal_cdf(input.eta),
699        d1,
700        -input.eta * d1,
701    )
702}
703
704#[inline]
705fn row_bernoulli_cloglog(
706    row: usize,
707    input: RowInput,
708    mode: CurvatureMode,
709) -> Result<RowOutput, EstimationError> {
710    finite_eta("standard complementary-log-log inverse link", input.eta)?;
711    let inner = input.eta.exp();
712    let mu = -(-inner).exp_m1();
713    let complement = (-inner).exp();
714    let d1 = inner * complement;
715    row_bernoulli_noncanonical(row, input, mode, mu, d1, d1 * (1.0 - inner))
716}
717
718#[inline]
719fn row_bernoulli_noncanonical(
720    row: usize,
721    input: RowInput,
722    mode: CurvatureMode,
723    mu: f64,
724    d1: f64,
725    d2: f64,
726) -> Result<RowOutput, EstimationError> {
727    let w_prior = prior_weight(row, input)?;
728    bernoulli_response(row, input, w_prior)?;
729    if !(mu.is_finite() && mu > 0.0 && mu < 1.0 && d1.is_finite() && d1 > 0.0 && d2.is_finite()) {
730        return Err(EstimationError::pirls_row_geometry_unrepresentable(row, "inverse-link jet", input.eta, mu));
731    }
732    if w_prior == 0.0 {
733        return certify_output(
734            row,
735            input.eta,
736            RowOutput {
737                mu,
738                ..RowOutput::default()
739            },
740        );
741    }
742    let v = mu * (1.0 - mu);
743    let fisher_per_prior = d1 * d1 / v;
744    let w_fisher = w_prior * fisher_per_prior;
745    if !(v.is_finite()
746        && v > 0.0
747        && fisher_per_prior.is_finite()
748        && fisher_per_prior > 0.0
749        && w_fisher.is_finite()
750        && w_fisher > 0.0)
751    {
752        return Err(EstimationError::pirls_row_geometry_unrepresentable(
753            row,
754            "Bernoulli Fisher weight",
755            input.eta,
756            w_fisher,
757        ));
758    }
759    let resid = input.y - mu;
760    let grad_eta = w_prior * resid * d1 / v;
761    let bracket = d2 / v - d1 * d1 * (1.0 - 2.0 * mu) / (v * v);
762    // -d²ℓ/dη² = W_F - prior·(y-μ)·d(h'/V)/dη.
763    let observed_correction = -w_prior * resid * bracket;
764    let w_hessian = select_w_hessian(mode, w_fisher, observed_correction);
765    if !w_hessian.is_finite() {
766        return Err(EstimationError::pirls_row_geometry_unrepresentable(
767            row,
768            "Bernoulli observed Hessian weight",
769            input.eta,
770            w_hessian,
771        ));
772    }
773    let dev = bernoulli_deviance(input.y, mu, w_prior);
774    certify_output(
775        row,
776        input.eta,
777        RowOutput {
778            mu,
779            grad_eta,
780            w_fisher,
781            w_hessian,
782            w_solver: w_hessian,
783            deviance: dev,
784        },
785    )
786}
787
788#[inline]
789fn bernoulli_logit_deviance(y: f64, eta: f64, w: f64) -> f64 {
790    let unit = if y == 0.0 {
791        softplus(eta)
792    } else if y == 1.0 {
793        softplus(-eta)
794    } else {
795        let response_logit = y.ln() - (-y).ln_1p();
796        bernoulli_kl_from_logits(response_logit, eta)
797    };
798    2.0 * w * unit
799}
800
801#[inline]
802fn bernoulli_deviance(y: f64, mu: f64, w: f64) -> f64 {
803    2.0 * w * (bd0(y, mu) + bd0(1.0 - y, 1.0 - mu))
804}
805
806/// Stable Φ(x) using the complementary error function with the same identity
807/// `erfc(-x/√2)/2 = Φ(x)` used by libstd. Keeps mass at the tails accurate.
808#[inline]
809fn standard_normal_cdf(x: f64) -> f64 {
810    0.5 * gam_gpu::numerics_host::erfc(-x * std::f64::consts::FRAC_1_SQRT_2)
811}
812
813#[inline]
814fn standard_normal_pdf(x: f64) -> f64 {
815    const COEFF: f64 = 0.398_942_280_401_432_7; // 1 / sqrt(2π)
816    COEFF * (-0.5 * x * x).exp()
817}
818
819// ────────────────────────────────────────────────────────────────────────
820// CUDA host harness
821// ────────────────────────────────────────────────────────────────────────
822
823/// Process-wide cache of compiled per-family modules.
824#[must_use]
825pub struct PirlsRowBackend {
826    #[cfg(target_os = "linux")]
827    inner: PirlsRowBackendLinux,
828}
829
830#[cfg(target_os = "linux")]
831struct PirlsRowBackendLinux {
832    ctx: Arc<CudaContext>,
833    modules: Mutex<std::collections::HashMap<ModuleKey, Arc<CudaModule>>>,
834}
835
836/// Distinguishes the three kernel modes in the per-process module cache.
837#[cfg(target_os = "linux")]
838#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
839enum KernelMode {
840    /// Full final-row kernel (mu, grad_eta, w_hessian, w_solver, deviance,
841    /// status).
842    FinalRow,
843    /// Solve-row kernel (4 outputs: grad_eta, w_solver, deviance, status).
844    SolveRow,
845    /// Alpha-ladder kernel (2 per-alpha outputs: objective[], status[]).
846    AlphaLadder,
847}
848
849#[cfg(target_os = "linux")]
850#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
851struct ModuleKey {
852    family: PirlsRowFamily,
853    curvature: CurvatureMode,
854    mode: KernelMode,
855}
856
857impl PirlsRowBackend {
858    pub const fn compiled() -> bool {
859        cfg!(target_os = "linux")
860    }
861
862    pub fn probe() -> Result<&'static Self, GpuError> {
863        static BACKEND: OnceLock<Result<PirlsRowBackend, GpuError>> = OnceLock::new();
864        BACKEND
865            .get_or_init(|| {
866                #[cfg(target_os = "linux")]
867                {
868                    Self::probe_linux()
869                }
870                #[cfg(not(target_os = "linux"))]
871                {
872                    Err(GpuError::DriverLibraryUnavailable {
873                        reason: "pirls_row GPU backend is Linux-only".to_string(),
874                    })
875                }
876            })
877            .as_ref()
878            .map_err(GpuError::clone)
879    }
880
881    #[cfg(target_os = "linux")]
882    fn probe_linux() -> Result<Self, GpuError> {
883        let parts = gam_gpu::backend_probe::probe_cuda_backend("pirls_row")?;
884        Ok(Self {
885            inner: PirlsRowBackendLinux {
886                ctx: parts.ctx,
887                modules: Mutex::new(std::collections::HashMap::new()),
888            },
889        })
890    }
891
892    /// Compile (or fetch from cache) the kernel module for `(family, curvature)`
893    /// in the given [`KernelMode`]. This is the single source of truth behind
894    /// [`module_for`], [`module_for_solve`], and [`module_for_ladder`]; the only
895    /// per-mode variation is which CUDA source generator is used (selected by
896    /// `mode`) and the error label `label` woven into compile/load diagnostics.
897    #[cfg(target_os = "linux")]
898    fn module_for_kind(
899        &self,
900        family: PirlsRowFamily,
901        curvature: CurvatureMode,
902        mode: KernelMode,
903        label: &str,
904    ) -> Result<Arc<CudaModule>, GpuError> {
905        let key = ModuleKey {
906            family,
907            curvature,
908            mode,
909        };
910        if let Some(existing) = self
911            .inner
912            .modules
913            .lock()
914            .gpu_ctx_with(|err| format!("pirls_row {label}module cache mutex poisoned: {err}"))?
915            .get(&key)
916        {
917            return Ok(existing.clone());
918        }
919        let source = match mode {
920            KernelMode::FinalRow => cuda_source_for(family, curvature),
921            KernelMode::SolveRow => solve_row_source_for(family, curvature),
922            KernelMode::AlphaLadder => ladder_source_for(family, curvature),
923        };
924        // #1551: route through the device-arch-pinned compile — this kernel uses
925        // `atomicAdd(double*, double)` (objective_out), which NVRTC rejects under
926        // its default sub-sm_60 arch, silently disabling the device PIRLS path.
927        let ptx = gam_gpu::device_cache::compile_ptx_arch(&source).gpu_ctx_with(|err| {
928            format!(
929                "pirls_row {label}NVRTC compile failed for {family}/{curv}: {err}",
930                family = family.as_str(),
931                curv = curvature.as_str(),
932            )
933        })?;
934        let module = self
935            .inner
936            .ctx
937            .load_module(ptx)
938            .gpu_ctx_with(|err| format!("pirls_row {label}module load failed: {err}"))?;
939        self.inner
940            .modules
941            .lock()
942            .gpu_ctx_with(|err| format!("pirls_row {label}module cache mutex poisoned: {err}"))?
943            .insert(key, module.clone());
944        Ok(module)
945    }
946
947    /// Compile (or fetch from cache) the **final-row** kernel module for
948    /// `(family, curvature)`. Writes the full production row surface.
949    #[cfg(target_os = "linux")]
950    pub fn module_for(
951        &self,
952        family: PirlsRowFamily,
953        curvature: CurvatureMode,
954    ) -> Result<Arc<CudaModule>, GpuError> {
955        self.module_for_kind(family, curvature, KernelMode::FinalRow, "")
956    }
957
958    /// Compile (or fetch from cache) the **solve-row** kernel module for
959    /// `(family, curvature)`. Writes only `grad_eta`, `w_solver`, `deviance`,
960    /// `status` — used on every hot Newton iteration.
961    #[cfg(target_os = "linux")]
962    pub fn module_for_solve(
963        &self,
964        family: PirlsRowFamily,
965        curvature: CurvatureMode,
966    ) -> Result<Arc<CudaModule>, GpuError> {
967        self.module_for_kind(family, curvature, KernelMode::SolveRow, "solve ")
968    }
969
970    /// Compile (or fetch from cache) the **alpha-ladder** kernel module for
971    /// `(family, curvature)`. Evaluates all [`ALPHA_LADDER_LEN`] step sizes in
972    /// a single launch, accumulating `objective[]` and `status[]` per alpha slot.
973    #[cfg(target_os = "linux")]
974    pub fn module_for_ladder(
975        &self,
976        family: PirlsRowFamily,
977        curvature: CurvatureMode,
978    ) -> Result<Arc<CudaModule>, GpuError> {
979        self.module_for_kind(family, curvature, KernelMode::AlphaLadder, "ladder ")
980    }
981
982}
983
984/// Stage 6 cache key for JIT-compiled family modules.
985#[cfg(target_os = "linux")]
986/// Stage 6 custom-family JIT specification.
987///
988/// Two levels per the charter:
989/// - **Level A** (`JitFamilySpec::glm`): provide a `(family, link)` enum
990///   value plus the Gamma shape constant. The generator emits the
991///   matching built-in row body, identical to the cached built-in
992///   kernel — useful for end-to-end JIT path validation against the
993///   built-in cache.
994/// - **Level B** (`JitFamilySpec::raw`): provide raw CUDA source for
995///   the row body. The body must define the same per-row locals the
996///   kernel shell expects: `mu`, `grad_eta`, `w_hessian`, `w_solver`,
997///   `dev`, and update `status`. The shell
998///   wraps it in the canonical
999///   `extern "C" __global__ void pirls_row_jit_{spec_id}(...)`
1000///   signature that [`launch_row_reweight_on_stream`] expects.
1001#[derive(Clone, Debug)]
1002pub struct JitFamilySpec {
1003    /// Process-unique identifier for this spec; the module cache uses
1004    /// it as a key so callers must reuse the same `spec_id` for the
1005    /// same body and pick a new one whenever the body changes.
1006    pub spec_id: u64,
1007    /// CUDA body source. Must read from `eta_i`, `y_i`, `wp`, set
1008    /// `status`, and assign to `mu`, `grad_eta`, `w_hessian`, `w_solver`,
1009    /// and `dev`. See `common_device_prolog` for
1010    /// the available helpers.
1011    pub body: String,
1012}
1013
1014#[cfg(target_os = "linux")]
1015impl JitFamilySpec {
1016    /// Level A: build a spec from a built-in `(family, curvature)`
1017    /// pair. The generator reuses the same per-family body as the
1018    /// built-in cached kernel — useful to validate the JIT pipeline
1019    /// end-to-end against the built-in numerical reference.
1020    #[cfg(target_os = "linux")]
1021    pub fn glm(
1022        spec_id: u64,
1023        family: PirlsRowFamily,
1024        curvature: CurvatureMode,
1025        gamma_shape: f64,
1026    ) -> Self {
1027        let mut body = match family {
1028            PirlsRowFamily::GaussianIdentity => gaussian_identity_body(curvature),
1029            PirlsRowFamily::PoissonLog => poisson_log_body(curvature),
1030            PirlsRowFamily::GammaLog => gamma_log_body(curvature),
1031            PirlsRowFamily::BernoulliLogit => bernoulli_logit_body(curvature),
1032            PirlsRowFamily::BernoulliProbit => bernoulli_probit_body(curvature),
1033            PirlsRowFamily::BernoulliCLogLog => bernoulli_cloglog_body(curvature),
1034        };
1035        if matches!(family, PirlsRowFamily::GammaLog) {
1036            body.insert_str(0, &format!("    const double shape = {gamma_shape:?};\n"));
1037        }
1038        Self { spec_id, body }
1039    }
1040
1041    /// Level B: build a spec from caller-supplied body source. The
1042    /// kernel shell wraps it; the body must define the required locals
1043    /// listed on [`JitFamilySpec`].
1044    pub fn raw(spec_id: u64, body: impl Into<String>) -> Self {
1045        Self {
1046            spec_id,
1047            body: body.into(),
1048        }
1049    }
1050
1051}
1052
1053/// Device-resident per-row output buffers for the GPU row-reweight kernel.
1054///
1055/// **final-row mode**: the five production numerical fields plus status,
1056/// length `n`. Written
1057/// once at convergence by [`launch_row_reweight_on_stream`]. For the hot
1058/// inner-loop use [`SolveRowBuffers`]; for line-search use
1059/// [`AlphaLadderDevBuffers`].
1060#[cfg(target_os = "linux")]
1061pub struct RowOutputDevBuffers {
1062    pub mu: cudarc::driver::CudaSlice<f64>,
1063    pub grad_eta: cudarc::driver::CudaSlice<f64>,
1064    pub w_hessian: cudarc::driver::CudaSlice<f64>,
1065    pub w_solver: cudarc::driver::CudaSlice<f64>,
1066    pub deviance: cudarc::driver::CudaSlice<f64>,
1067    pub status: cudarc::driver::CudaSlice<u32>,
1068    pub n: usize,
1069}
1070
1071#[cfg(target_os = "linux")]
1072impl RowOutputDevBuffers {
1073    /// Allocate all production per-row output buffers (length `n`) on `stream`.
1074    pub fn allocate(stream: &Arc<cudarc::driver::CudaStream>, n: usize) -> Result<Self, GpuError> {
1075        let alloc_f64 = |label: &'static str| {
1076            stream
1077                .alloc_zeros::<f64>(n)
1078                .gpu_ctx_with(|err| format!("pirls_row alloc {label}: {err}"))
1079        };
1080        let alloc_u32 = |label: &'static str| {
1081            stream
1082                .alloc_zeros::<u32>(n)
1083                .gpu_ctx_with(|err| format!("pirls_row alloc {label}: {err}"))
1084        };
1085        Ok(Self {
1086            mu: alloc_f64("mu")?,
1087            grad_eta: alloc_f64("grad_eta")?,
1088            w_hessian: alloc_f64("w_hessian")?,
1089            w_solver: alloc_f64("w_solver")?,
1090            deviance: alloc_f64("deviance")?,
1091            status: alloc_u32("status")?,
1092            n,
1093        })
1094    }
1095}
1096
1097/// Device-resident per-row output buffers for the **solve-row** mode.
1098///
1099/// Allocates only the four fields the PIRLS solver reads on every Newton
1100/// iteration: `grad_eta` (score for Xᵀg RHS), `w_solver` (working weight
1101/// for XᵀWX assembly), `deviance` (per-row deviance for convergence check),
1102/// and `status` (one exact refusal code per row). Written
1103/// by [`launch_solve_row_on_stream`]; used instead of [`RowOutputDevBuffers`]
1104/// during the hot inner loop to reduce device memory and kernel store traffic.
1105#[cfg(target_os = "linux")]
1106pub struct SolveRowBuffers {
1107    /// ∂ℓ/∂η_i — score for Xᵀg RHS formation.
1108    pub grad_eta: cudarc::driver::CudaSlice<f64>,
1109    /// Stabilised Hessian weight — fed to XᵀWX assembly.
1110    pub w_solver: cudarc::driver::CudaSlice<f64>,
1111    /// Per-row deviance contribution — summed for convergence check.
1112    pub deviance: cudarc::driver::CudaSlice<f64>,
1113    /// Zero on success, otherwise the row's deterministic refusal code.
1114    pub status: cudarc::driver::CudaSlice<u32>,
1115    pub n: usize,
1116}
1117
1118#[cfg(target_os = "linux")]
1119impl SolveRowBuffers {
1120    /// Allocate the four solve-row output buffers (length `n`) on `stream`.
1121    pub fn allocate(stream: &Arc<cudarc::driver::CudaStream>, n: usize) -> Result<Self, GpuError> {
1122        let alloc_f64 = |label: &'static str| {
1123            stream
1124                .alloc_zeros::<f64>(n)
1125                .gpu_ctx_with(|err| format!("pirls_row solve alloc {label}: {err}"))
1126        };
1127        let alloc_u32 = |label: &'static str| {
1128            stream
1129                .alloc_zeros::<u32>(n)
1130                .gpu_ctx_with(|err| format!("pirls_row solve alloc {label}: {err}"))
1131        };
1132        Ok(Self {
1133            grad_eta: alloc_f64("grad_eta")?,
1134            w_solver: alloc_f64("w_solver")?,
1135            deviance: alloc_f64("deviance")?,
1136            status: alloc_u32("status")?,
1137            n,
1138        })
1139    }
1140}
1141
1142/// Number of alpha step sizes in the fused alpha ladder.
1143pub const ALPHA_LADDER_LEN: usize = 7;
1144
1145/// The fixed alpha step-size ladder: `[1, 0.5, 0.25, 0.125, 0.0625, 0.03125, 0.015625]`.
1146pub const ALPHA_LADDER: [f64; ALPHA_LADDER_LEN] =
1147    [1.0, 0.5, 0.25, 0.125, 0.0625, 0.03125, 0.015625];
1148
1149/// Device buffers for the fused alpha-ladder candidate-objective kernel.
1150///
1151/// **candidate-objective mode**: for each of the [`ALPHA_LADDER_LEN`] step
1152/// sizes α_k the kernel evaluates `η_trial_i = η_i + α_k · xδ_i`, computes
1153/// the per-row deviance, and atomically accumulates the sum into
1154/// `objective_dev[k]`. Each row writes its refusal code to
1155/// `status_dev[k * n + i]`, preserving deterministic row identity.
1156/// After a single `memcpy_dtoh` the host picks the first α that achieves
1157/// deviance descent — no per-α kernel launch, no full row-output write.
1158#[cfg(target_os = "linux")]
1159pub struct AlphaLadderDevBuffers {
1160    /// Device: summed deviance for each alpha step, length [`ALPHA_LADDER_LEN`].
1161    pub objective_dev: cudarc::driver::CudaSlice<f64>,
1162    /// Device: row refusal codes in alpha-major order, length
1163    /// `ALPHA_LADDER_LEN * n`.
1164    pub status_dev: cudarc::driver::CudaSlice<u32>,
1165    pub n: usize,
1166}
1167
1168#[cfg(target_os = "linux")]
1169impl AlphaLadderDevBuffers {
1170    /// Allocate the ladder device buffers on `stream`.
1171    pub fn allocate(stream: &Arc<cudarc::driver::CudaStream>, n: usize) -> Result<Self, GpuError> {
1172        let status_len = ALPHA_LADDER_LEN.checked_mul(n).ok_or_else(|| {
1173            gam_gpu::gpu_err!("pirls_row ladder status length overflows: {ALPHA_LADDER_LEN} * {n}")
1174        })?;
1175        Ok(Self {
1176            objective_dev: stream
1177                .alloc_zeros::<f64>(ALPHA_LADDER_LEN)
1178                .gpu_ctx_with(|err| format!("pirls_row ladder alloc objective: {err}"))?,
1179            status_dev: stream
1180                .alloc_zeros::<u32>(status_len)
1181                .gpu_ctx_with(|err| format!("pirls_row ladder alloc status: {err}"))?,
1182            n,
1183        })
1184    }
1185
1186    /// Zero all per-alpha accumulators in-place (call before each ladder launch).
1187    pub fn zero(&mut self, stream: &Arc<cudarc::driver::CudaStream>) -> Result<(), GpuError> {
1188        stream
1189            .memset_zeros(&mut self.objective_dev)
1190            .gpu_ctx_with(|err| format!("pirls_row ladder zero objective: {err}"))?;
1191        stream
1192            .memset_zeros(&mut self.status_dev)
1193            .gpu_ctx_with(|err| format!("pirls_row ladder zero status: {err}"))
1194    }
1195}
1196
1197/// Device-side row reweight launcher.
1198///
1199/// Resolves the cached per-family kernel from [`PirlsRowBackend::module_for`],
1200/// dispatches a 1D grid of `THREADS_PER_BLOCK = 256` threads across `n`
1201/// rows, and returns once the launch is enqueued on `stream`. The kernel
1202/// writes the per-row IRLS state into `out` in place; no host transfers.
1203///
1204/// The kernel's `extern "C"` signature is fixed at the top of `cuda_source_for`
1205/// (see `extern "C" __global__ void {kernel_name}(int n, …)` in this file).
1206/// `gamma_shape`: active Gamma dispersion shape (α > 0). Forwarded as a
1207/// scalar kernel argument only for `PirlsRowFamily::GammaLog`; all other
1208/// families compile a ten-argument kernel and ignore this value. Pass `1.0`
1209/// for non-Gamma fits.
1210#[cfg(target_os = "linux")]
1211pub fn launch_row_reweight_on_stream(
1212    backend: &PirlsRowBackend,
1213    family: PirlsRowFamily,
1214    curvature: CurvatureMode,
1215    gamma_shape: f64,
1216    stream: &Arc<cudarc::driver::CudaStream>,
1217    n: usize,
1218    eta_dev: &cudarc::driver::CudaSlice<f64>,
1219    y_dev: &cudarc::driver::CudaSlice<f64>,
1220    prior_w_dev: &cudarc::driver::CudaSlice<f64>,
1221    out: &mut RowOutputDevBuffers,
1222) -> Result<(), GpuError> {
1223    use cudarc::driver::{LaunchConfig, PushKernelArg};
1224    if out.n != n {
1225        gam_gpu::gpu_bail!("row reweight buffers shape {} mismatches n={n}", out.n);
1226    }
1227    let module = backend.module_for(family, curvature)?;
1228    let func = module
1229        .load_function(family.kernel_name())
1230        .gpu_ctx_with(|err| {
1231            format!(
1232                "row reweight load_function({}): {err}",
1233                family.kernel_name()
1234            )
1235        })?;
1236    const THREADS_PER_BLOCK: u32 = 256;
1237    let n_u32 = u32::try_from(n)
1238        .map_err(|_| gam_gpu::gpu_err!("n={n} exceeds u32 for row reweight grid sizing"))?;
1239    let grid_x = n_u32.div_ceil(THREADS_PER_BLOCK).max(1);
1240    let n_i32 = i32::try_from(n)
1241        .map_err(|_| gam_gpu::gpu_err!("n={n} exceeds i32 for row reweight kernel argument"))?;
1242    let cfg = LaunchConfig {
1243        grid_dim: (grid_x, 1, 1),
1244        block_dim: (THREADS_PER_BLOCK, 1, 1),
1245        shared_mem_bytes: 0,
1246    };
1247    let mut builder = stream.launch_builder(&func);
1248    builder.arg(&n_i32);
1249    builder.arg(eta_dev);
1250    builder.arg(y_dev);
1251    builder.arg(prior_w_dev);
1252    // GammaLog kernel has `double shape` before the output buffers.
1253    if matches!(family, PirlsRowFamily::GammaLog) {
1254        builder.arg(&gamma_shape);
1255    }
1256    builder.arg(&mut out.mu);
1257    builder.arg(&mut out.grad_eta);
1258    builder.arg(&mut out.w_hessian);
1259    builder.arg(&mut out.w_solver);
1260    builder.arg(&mut out.deviance);
1261    builder.arg(&mut out.status);
1262    // SAFETY: kernel signature for non-GammaLog is (n:i32, 3×const f64*,
1263    // 5×mut f64*, 1×mut u32*). GammaLog extends this with `double shape`
1264    // after `prior_w` and before the output buffers — see `cuda_source_for`.
1265    // Arg order/types match one-for-one. Output buffers were allocated with
1266    // `n` elements each (validated above); input buffers are caller-supplied
1267    // with length n. Grid covers all n rows; threads guard `if (i >= n) return`.
1268    unsafe { builder.launch(cfg) }
1269        .gpu_ctx_with(|err| format!("row reweight launch({}): {err}", family.kernel_name()))?;
1270    Ok(())
1271}
1272
1273/// **solve-row** mode launcher.
1274///
1275/// Runs the per-family row math and writes only the four fields needed by the
1276/// PIRLS solver on each Newton iteration: `grad_eta`, `w_solver`, `deviance`,
1277/// `status`. The CUDA kernel is compiled from a specialised source
1278/// (`solve_row_source_for`) that skips the `mu` and `w_hessian` stores,
1279/// reducing both bandwidth and register
1280/// pressure relative to [`launch_row_reweight_on_stream`].
1281///
1282/// Call once per Newton step on the accepted η. At convergence, call
1283/// [`launch_row_reweight_on_stream`] (final-row mode) to populate the full
1284/// output surface before downloading.
1285///
1286/// `gamma_shape`: active Gamma dispersion shape (α > 0). Forwarded as a kernel
1287/// argument only for `PirlsRowFamily::GammaLog`. Pass `1.0` for non-Gamma fits.
1288#[cfg(target_os = "linux")]
1289pub fn launch_solve_row_on_stream(
1290    backend: &PirlsRowBackend,
1291    family: PirlsRowFamily,
1292    curvature: CurvatureMode,
1293    gamma_shape: f64,
1294    stream: &Arc<cudarc::driver::CudaStream>,
1295    n: usize,
1296    eta_dev: &cudarc::driver::CudaSlice<f64>,
1297    y_dev: &cudarc::driver::CudaSlice<f64>,
1298    prior_w_dev: &cudarc::driver::CudaSlice<f64>,
1299    out: &mut SolveRowBuffers,
1300) -> Result<(), GpuError> {
1301    use cudarc::driver::{LaunchConfig, PushKernelArg};
1302    if out.n != n {
1303        gam_gpu::gpu_bail!("solve-row buffers shape {} mismatches n={n}", out.n);
1304    }
1305    let module = backend.module_for_solve(family, curvature)?;
1306    let kernel_name = family.solve_kernel_name();
1307    let func = module
1308        .load_function(kernel_name)
1309        .gpu_ctx_with(|err| format!("solve-row load_function({kernel_name}): {err}"))?;
1310    const THREADS_PER_BLOCK: u32 = 256;
1311    let n_u32 = u32::try_from(n)
1312        .map_err(|_| gam_gpu::gpu_err!("n={n} exceeds u32 for solve-row grid sizing"))?;
1313    let grid_x = n_u32.div_ceil(THREADS_PER_BLOCK).max(1);
1314    let n_i32 = i32::try_from(n)
1315        .map_err(|_| gam_gpu::gpu_err!("n={n} exceeds i32 for solve-row kernel argument"))?;
1316    let cfg = LaunchConfig {
1317        grid_dim: (grid_x, 1, 1),
1318        block_dim: (THREADS_PER_BLOCK, 1, 1),
1319        shared_mem_bytes: 0,
1320    };
1321    let mut builder = stream.launch_builder(&func);
1322    builder.arg(&n_i32);
1323    builder.arg(eta_dev);
1324    builder.arg(y_dev);
1325    builder.arg(prior_w_dev);
1326    // GammaLog solve kernel has `double shape` before the output buffers.
1327    if matches!(family, PirlsRowFamily::GammaLog) {
1328        builder.arg(&gamma_shape);
1329    }
1330    builder.arg(&mut out.grad_eta);
1331    builder.arg(&mut out.w_solver);
1332    builder.arg(&mut out.deviance);
1333    builder.arg(&mut out.status);
1334    // SAFETY: solve_row_source_for emits for non-GammaLog:
1335    //   (int n, const f64* eta, const f64* y, const f64* prior_w,
1336    //    f64* grad_eta_out, f64* w_solver_out, f64* deviance_out, u32* status_out)
1337    // For GammaLog the signature inserts `double shape` after `prior_w`.
1338    // 4 outputs match the 4 SolveRowBuffers fields; grid covers all n rows with
1339    // per-thread guard `if (i >= n) return`.
1340    unsafe { builder.launch(cfg) }
1341        .gpu_ctx_with(|err| format!("solve-row launch({kernel_name}): {err}"))?;
1342    Ok(())
1343}
1344
1345/// **candidate-objective / fused alpha-ladder** launcher.
1346///
1347/// Evaluates `η_trial_i = η_i + α_k · xδ_i` for all `i ∈ [0,n)` and all
1348/// `k ∈ [0, ALPHA_LADDER_LEN)` simultaneously.  Each thread atomically
1349/// accumulates the per-row deviance into `out.objective_dev[k]` and
1350/// writes each refusal code to `out.status_dev[k * n + i]`.
1351///
1352/// The grid is `(row_blocks × ALPHA_LADDER_LEN)`: block index `bx / n_blocks`
1353/// selects the alpha slot, `bx % n_blocks` selects the row tile.
1354///
1355/// Caller must call [`AlphaLadderDevBuffers::zero`] before each launch, then
1356/// reduce the alpha-major status matrix with the deterministic smallest-row
1357/// reducer, then read the seven objectives plus seven row/code summaries.
1358#[cfg(target_os = "linux")]
1359pub fn launch_alpha_ladder_on_stream(
1360    backend: &PirlsRowBackend,
1361    family: PirlsRowFamily,
1362    curvature: CurvatureMode,
1363    gamma_shape: f64,
1364    stream: &Arc<cudarc::driver::CudaStream>,
1365    n: usize,
1366    eta_dev: &cudarc::driver::CudaSlice<f64>,
1367    xd_dev: &cudarc::driver::CudaSlice<f64>,
1368    y_dev: &cudarc::driver::CudaSlice<f64>,
1369    prior_w_dev: &cudarc::driver::CudaSlice<f64>,
1370    out: &mut AlphaLadderDevBuffers,
1371) -> Result<(), GpuError> {
1372    use cudarc::driver::{LaunchConfig, PushKernelArg};
1373    if out.n != n {
1374        gam_gpu::gpu_bail!("alpha-ladder buffers shape {} mismatches n={n}", out.n);
1375    }
1376    let module = backend.module_for_ladder(family, curvature)?;
1377    let kernel_name = family.ladder_kernel_name();
1378    let func = module
1379        .load_function(kernel_name)
1380        .gpu_ctx_with(|err| format!("alpha-ladder load_function({kernel_name}): {err}"))?;
1381    const THREADS_PER_BLOCK: u32 = 256;
1382    let n_u32 = u32::try_from(n)
1383        .map_err(|_| gam_gpu::gpu_err!("n={n} exceeds u32 for alpha-ladder grid sizing"))?;
1384    let row_blocks = n_u32.div_ceil(THREADS_PER_BLOCK).max(1);
1385    let n_i32 = i32::try_from(n)
1386        .map_err(|_| gam_gpu::gpu_err!("n={n} exceeds i32 for alpha-ladder kernel argument"))?;
1387    // Grid: x = row tile index (0..row_blocks), y = alpha index (0..ALPHA_LADDER_LEN).
1388    let cfg = LaunchConfig {
1389        grid_dim: (row_blocks, ALPHA_LADDER_LEN as u32, 1),
1390        block_dim: (THREADS_PER_BLOCK, 1, 1),
1391        shared_mem_bytes: 0,
1392    };
1393    let mut builder = stream.launch_builder(&func);
1394    builder.arg(&n_i32);
1395    builder.arg(eta_dev);
1396    builder.arg(xd_dev);
1397    builder.arg(y_dev);
1398    builder.arg(prior_w_dev);
1399    // GammaLog ladder kernel has `double shape` before the output buffers.
1400    if matches!(family, PirlsRowFamily::GammaLog) {
1401        builder.arg(&gamma_shape);
1402    }
1403    builder.arg(&mut out.objective_dev);
1404    builder.arg(&mut out.status_dev);
1405    // SAFETY: for non-GammaLog, ladder_source_for emits:
1406    //   (int n, const f64* eta, const f64* xd, const f64* y, const f64* prior_w,
1407    //    f64* objective_out, u32* status_out)
1408    // For GammaLog `double shape` is inserted after `prior_w`.
1409    // Grid is (row_blocks × ALPHA_LADDER_LEN); each thread reads alphas[] via
1410    // blockIdx.y, rows via blockIdx.x * blockDim.x + threadIdx.x (guarded by n).
1411    // Atomic double-precision add to objective_out[blockIdx.y]; each thread
1412    // owns status_out[blockIdx.y * n + i].
1413    unsafe { builder.launch(cfg) }
1414        .gpu_ctx_with(|err| format!("alpha-ladder launch({kernel_name}): {err}"))?;
1415    Ok(())
1416}
1417
1418// ────────────────────────────────────────────────────────────────────────
1419// CUDA sources (one per family / curvature pair)
1420// ────────────────────────────────────────────────────────────────────────
1421
1422/// Common device-side helpers shared across every family kernel.  The log-link
1423/// interval is injected from the same Rust constants used by the CPU oracle,
1424/// so the generated CUDA source cannot drift to a second domain policy.
1425#[cfg(target_os = "linux")]
1426fn common_device_prolog() -> String {
1427    // Brace-heavy CUDA source: keep it a plain raw string literal so every C
1428    // brace stays literal (no `format!` grammar to escape), and substitute the
1429    // two host solver bounds by unique sentinel token.
1430    r#"
1431// NVRTC math builtins: prototypes must carry an execution space. Newer
1432// NVRTC (CUDA 12.x JIT semantics) rejects unannotated declarations outright
1433// ("host functions are not allowed in JIT mode"), which failed every
1434// pirls_row kernel compile on real hardware while CPU-only CI stayed green
1435// (#2313 hardware sweep). `__device__` matches how the CUDA math library
1436// declares them; the definitions come from libdevice as before.
1437extern "C" {
1438    __device__ double exp(double);
1439    __device__ double log(double);
1440    __device__ double log1p(double);
1441    __device__ double expm1(double);
1442    __device__ double fabs(double);
1443    __device__ double erfc(double);
1444}
1445
1446static constexpr double PIRLS_LOG_ETA_MIN = __PIRLS_LOG_ETA_MIN__;
1447static constexpr double PIRLS_LOG_ETA_MAX = __PIRLS_LOG_ETA_MAX__;
1448
1449static constexpr unsigned int PIRLS_OK = 0u;
1450static constexpr unsigned int PIRLS_ETA_DOMAIN = 1u;
1451static constexpr unsigned int PIRLS_PRIOR_WEIGHT = 2u;
1452static constexpr unsigned int PIRLS_RESPONSE = 3u;
1453static constexpr unsigned int PIRLS_GAMMA_SHAPE = 4u;
1454static constexpr unsigned int PIRLS_INVERSE_LINK = 5u;
1455static constexpr unsigned int PIRLS_FISHER_WEIGHT = 6u;
1456static constexpr unsigned int PIRLS_OBSERVED_WEIGHT = 7u;
1457static constexpr unsigned int PIRLS_GRADIENT = 8u;
1458static constexpr unsigned int PIRLS_DEVIANCE = 9u;
1459static constexpr unsigned int PIRLS_FINAL_OUTPUT = 10u;
1460
1461__device__ __forceinline__ void pirls_refuse(unsigned int* status, unsigned int code) {
1462    if (*status == PIRLS_OK) *status = code;
1463}
1464
1465__device__ __forceinline__ bool pirls_log_eta_valid(double eta) {
1466    return eta >= PIRLS_LOG_ETA_MIN && eta <= PIRLS_LOG_ETA_MAX;
1467}
1468
1469__device__ __forceinline__ double softplus(double x) {
1470    return (x > 0.0 ? x : 0.0) + log1p(exp(-fabs(x)));
1471}
1472
1473__device__ __forceinline__ double expm1_minus_x(double x) {
1474    if (fabs(x) > 0.5) return expm1(x) - x;
1475    double term = 0.5 * x * x;
1476    double sum = term;
1477    double degree = 2.0;
1478    for (;;) {
1479        degree += 1.0;
1480        term *= x / degree;
1481        double next = sum + term;
1482        if (next == sum) return next;
1483        sum = next;
1484    }
1485}
1486
1487__device__ __forceinline__ double log1p_minus_x(double x) {
1488    if (fabs(x) > 0.5) return log1p(x) - x;
1489    double power = x * x;
1490    double sign = -1.0;
1491    double degree = 2.0;
1492    double sum = sign * power / degree;
1493    for (;;) {
1494        power *= x;
1495        sign = -sign;
1496        degree += 1.0;
1497        double next = sum + sign * power / degree;
1498        if (next == sum) return next;
1499        sum = next;
1500    }
1501}
1502
1503__device__ __forceinline__ double logistic(double x) {
1504    if (x >= 0.0) return 1.0 / (1.0 + exp(-x));
1505    double e = exp(x);
1506    return e / (1.0 + e);
1507}
1508
1509__device__ __forceinline__ double bernoulli_kl_from_logits(double a, double b) {
1510    if (a == b) return 0.0;
1511    double h = b - a;
1512    if (fabs(h) <= 0.5) {
1513        double p = a <= 0.0 ? logistic(a) : logistic(-a);
1514        double local_h = a <= 0.0 ? h : -h;
1515        double em1 = expm1(local_h);
1516        double x = p * em1;
1517        return log1p_minus_x(x) + p * expm1_minus_x(local_h);
1518    }
1519    if (a <= 0.0) {
1520        double p = logistic(a);
1521        return p * (a - b) + softplus(b) - softplus(a);
1522    }
1523    double q = logistic(-a);
1524    return q * (b - a) + softplus(-b) - softplus(-a);
1525}
1526
1527__device__ __forceinline__ double bd0(double x, double m) {
1528    if (x == 0.0) return m;
1529    if (x == m) return 0.0;
1530    double hi = x > m ? x : m;
1531    double lo = x < m ? x : m;
1532    if (fabs(x - m) / hi < 0.2) {
1533        double v = ((x - m) / hi) / (1.0 + lo / hi);
1534        double sum = (x - m) * v;
1535        double term = 2.0 * x * v;
1536        double v2 = v * v;
1537        double denominator = 3.0;
1538        for (;;) {
1539            term *= v2;
1540            double next = sum + term / denominator;
1541            if (next == sum) return next;
1542            sum = next;
1543            denominator += 2.0;
1544        }
1545    }
1546    return x * (log(x) - log(m)) + (m - x);
1547}
1548
1549__device__ __forceinline__ double bernoulli_deviance(double y, double mu, double w) {
1550    return 2.0 * w * (bd0(y, mu) + bd0(1.0 - y, 1.0 - mu));
1551}
1552
1553__device__ __forceinline__ double logit_deviance(double y, double eta, double w) {
1554    double unit;
1555    if (y == 0.0) unit = softplus(eta);
1556    else if (y == 1.0) unit = softplus(-eta);
1557    else {
1558        double response_logit = log(y) - log1p(-y);
1559        unit = bernoulli_kl_from_logits(response_logit, eta);
1560    }
1561    return 2.0 * w * unit;
1562}
1563
1564__device__ __forceinline__ double std_norm_cdf(double x) {
1565    return 0.5 * erfc(-x * 0.7071067811865475);
1566}
1567
1568__device__ __forceinline__ double std_norm_pdf(double x) {
1569    return 0.3989422804014327 * exp(-0.5 * x * x);
1570}
1571
1572__device__ __forceinline__ double positive_mul_div(double a, double b, double c) {
1573    double product = a * b;
1574    if (isfinite(product) && product > 0.0) {
1575        double value = product / c;
1576        if (isfinite(value) && value > 0.0) return value;
1577    }
1578    double quotient_a = a / c;
1579    if (isfinite(quotient_a) && quotient_a > 0.0) {
1580        double value = quotient_a * b;
1581        if (isfinite(value) && value > 0.0) return value;
1582    }
1583    double quotient_b = b / c;
1584    if (isfinite(quotient_b) && quotient_b > 0.0) {
1585        double value = quotient_b * a;
1586        if (isfinite(value) && value > 0.0) return value;
1587    }
1588    return product / c;
1589}
1590
1591__device__ __forceinline__ double gamma_unit_deviance_near_one(double u) {
1592    if (fabs(u) > 0.125) return u - log1p(u);
1593    double power = u * u;
1594    double sum = 0.5 * power;
1595    for (int degree = 3; degree <= 32; ++degree) {
1596        power *= u;
1597        double term = power / (double)degree;
1598        double next = sum + ((degree & 1) ? -term : term);
1599        if (next == sum) break;
1600        sum = next;
1601    }
1602    return sum;
1603}
1604
1605__device__ __forceinline__ double poisson_unit_deviance_near_one(double u) {
1606    if (fabs(u) > 0.125) return (1.0 + u) * log1p(u) - u;
1607    double power = u * u;
1608    double sum = 0.5 * power;
1609    for (int degree = 3; degree <= 32; ++degree) {
1610        power *= u;
1611        double coefficient = ((degree & 1) ? -1.0 : 1.0)
1612            / ((double)degree * (double)(degree - 1));
1613        double next = sum + coefficient * power;
1614        if (next == sum) break;
1615        sum = next;
1616    }
1617    return sum;
1618}
1619
1620__device__ __forceinline__ bool pirls_outputs_finite(
1621    double mu, double grad_eta, double w_fisher, double w_hessian,
1622    double w_solver, double dev
1623) {
1624    return isfinite(mu) && isfinite(grad_eta) && isfinite(w_fisher)
1625        && isfinite(w_hessian) && isfinite(w_solver) && isfinite(dev);
1626}
1627"#
1628    .replace(
1629        "__PIRLS_LOG_ETA_MIN__",
1630        &format!("{:?}", crate::mixture_link::LOG_LINK_SOLVER_ETA_MIN),
1631    )
1632    .replace(
1633        "__PIRLS_LOG_ETA_MAX__",
1634        &format!("{:?}", crate::mixture_link::LOG_LINK_SOLVER_ETA_MAX),
1635    )
1636}
1637
1638/// Build the per-family CUDA source. Each source defines exactly one entry
1639/// kernel (`family.kernel_name()`) reading the input arrays and writing the
1640/// output arrays defined by the [`RowOutput`] contract above.
1641///
1642/// The GammaLog kernel has an extra `double shape` parameter after `prior_w`
1643/// so the host can forward the active dispersion shape. All other families
1644/// use the standard ten-argument signature.
1645#[cfg(target_os = "linux")]
1646fn cuda_source_for(family: PirlsRowFamily, curvature: CurvatureMode) -> String {
1647    let body = match family {
1648        PirlsRowFamily::GaussianIdentity => gaussian_identity_body(curvature),
1649        PirlsRowFamily::PoissonLog => poisson_log_body(curvature),
1650        PirlsRowFamily::GammaLog => gamma_log_body(curvature),
1651        PirlsRowFamily::BernoulliLogit => bernoulli_logit_body(curvature),
1652        PirlsRowFamily::BernoulliProbit => bernoulli_probit_body(curvature),
1653        PirlsRowFamily::BernoulliCLogLog => bernoulli_cloglog_body(curvature),
1654    };
1655    let kernel_name = family.kernel_name();
1656    // Inject a curvature marker into the source so each `(family, curvature)`
1657    // pair compiles a distinct PTX module (cache key) and Stage 5 can branch
1658    // on `PIRLS_CURVATURE_OBSERVED` from inside each per-family body without
1659    // changing the host harness.
1660    let curvature_define = match curvature {
1661        CurvatureMode::Fisher => "#define PIRLS_CURVATURE_FISHER 1",
1662        CurvatureMode::Observed => "#define PIRLS_CURVATURE_OBSERVED 1",
1663    };
1664    // GammaLog receives the active shape as a scalar kernel argument so the
1665    // host can forward any positive shape without recompiling the PTX.
1666    let shape_param = if matches!(family, PirlsRowFamily::GammaLog) {
1667        "    double         shape,\n"
1668    } else {
1669        ""
1670    };
1671    format!(
1672        r#"
1673{curvature_define}
1674{prolog}
1675
1676extern "C" __global__ void {kernel_name}(
1677    int            n,
1678    const double* __restrict__ eta,
1679    const double* __restrict__ y,
1680    const double* __restrict__ prior_w,
1681{shape_param}    double* __restrict__ mu_out,
1682    double* __restrict__ grad_eta_out,
1683    double* __restrict__ w_hessian_out,
1684    double* __restrict__ w_solver_out,
1685    double* __restrict__ deviance_out,
1686    unsigned int* __restrict__ status_out
1687) {{
1688    int i = blockIdx.x * blockDim.x + threadIdx.x;
1689    if (i >= n) return;
1690    unsigned int status = PIRLS_OK;
1691    double eta_i = eta[i];
1692    double y_i = y[i];
1693    double wp = prior_w[i];
1694{body}
1695    if (status == PIRLS_OK) {{
1696        mu_out[i] = mu;
1697        grad_eta_out[i] = grad_eta;
1698        w_hessian_out[i] = w_hessian;
1699        w_solver_out[i] = w_solver;
1700        deviance_out[i] = dev;
1701    }}
1702    status_out[i] = status;
1703}}
1704"#,
1705        prolog = common_device_prolog(),
1706    )
1707}
1708
1709/// Emits a CUDA comment tag identifying the curvature mode the kernel was
1710/// compiled for. Each body builder prepends this so the body source actually
1711/// consumes the `curvature` argument and Stage 5 can keyed-extend the bodies
1712/// behind `#ifdef PIRLS_CURVATURE_OBSERVED`.
1713#[cfg(target_os = "linux")]
1714#[inline]
1715fn curvature_tag(curvature: CurvatureMode) -> &'static str {
1716    match curvature {
1717        CurvatureMode::Fisher => "    // curvature: fisher\n",
1718        CurvatureMode::Observed => "    // curvature: observed\n",
1719    }
1720}
1721
1722#[cfg(target_os = "linux")]
1723fn gaussian_identity_body(curvature: CurvatureMode) -> String {
1724    let tag = curvature_tag(curvature);
1725    format!(
1726        r#"{tag}    double mu = 0.0, grad_eta = 0.0, w_fisher = 0.0;
1727    double w_hessian = 0.0, w_solver = 0.0, dev = 0.0;
1728    if (!isfinite(eta_i)) pirls_refuse(&status, PIRLS_ETA_DOMAIN);
1729    if (status == PIRLS_OK && !(isfinite(wp) && wp >= 0.0))
1730        pirls_refuse(&status, PIRLS_PRIOR_WEIGHT);
1731    if (status == PIRLS_OK && wp > 0.0 && !isfinite(y_i))
1732        pirls_refuse(&status, PIRLS_RESPONSE);
1733    if (status == PIRLS_OK) {{
1734        mu = eta_i;
1735        w_fisher = wp;
1736        w_hessian = wp;
1737        w_solver = w_hessian;
1738        if (wp > 0.0) {{
1739            double resid = y_i - mu;
1740            grad_eta = wp * resid;
1741            dev = wp * resid * resid;
1742        }}
1743    }}
1744    if (status == PIRLS_OK && !pirls_outputs_finite(
1745            mu, grad_eta, w_fisher, w_hessian, w_solver, dev))
1746        pirls_refuse(&status, PIRLS_FINAL_OUTPUT);
1747"#
1748    )
1749}
1750
1751#[cfg(target_os = "linux")]
1752fn poisson_log_body(curvature: CurvatureMode) -> String {
1753    let tag = curvature_tag(curvature);
1754    format!(
1755        r#"{tag}    double mu = 0.0, grad_eta = 0.0, w_fisher = 0.0;
1756    double w_hessian = 0.0, w_solver = 0.0, dev = 0.0;
1757    if (!pirls_log_eta_valid(eta_i)) pirls_refuse(&status, PIRLS_ETA_DOMAIN);
1758    if (status == PIRLS_OK && !(isfinite(wp) && wp >= 0.0))
1759        pirls_refuse(&status, PIRLS_PRIOR_WEIGHT);
1760    if (status == PIRLS_OK && wp > 0.0 && !(isfinite(y_i) && y_i >= 0.0))
1761        pirls_refuse(&status, PIRLS_RESPONSE);
1762    if (status == PIRLS_OK) {{
1763        mu = exp(eta_i);
1764        if (!(isfinite(mu) && mu > 0.0)) pirls_refuse(&status, PIRLS_INVERSE_LINK);
1765    }}
1766    if (status == PIRLS_OK && wp > 0.0) {{
1767        w_fisher = wp * mu;
1768        if (!(isfinite(w_fisher) && w_fisher > 0.0))
1769            pirls_refuse(&status, PIRLS_FISHER_WEIGHT);
1770        if (status == PIRLS_OK) {{
1771            w_hessian = w_fisher;
1772            w_solver = w_hessian;
1773            grad_eta = wp * (y_i - mu);
1774            double u = (y_i - mu) / mu;
1775            double dev_base;
1776            if (y_i == 0.0) {{
1777                dev_base = w_fisher;
1778            }} else {{
1779                double scaled_unit = w_fisher * poisson_unit_deviance_near_one(u);
1780                if (isfinite(scaled_unit) && scaled_unit >= 0.0) {{
1781                    dev_base = scaled_unit;
1782                }} else {{
1783                    double weighted_y = positive_mul_div(w_fisher, y_i, mu);
1784                    dev_base = weighted_y * (log(y_i) - eta_i - 1.0) + w_fisher;
1785                }}
1786            }}
1787            if (!isfinite(grad_eta)) pirls_refuse(&status, PIRLS_GRADIENT);
1788            dev = 2.0 * dev_base;
1789            if (!isfinite(dev)) pirls_refuse(&status, PIRLS_DEVIANCE);
1790        }}
1791    }}
1792    if (status == PIRLS_OK && !pirls_outputs_finite(
1793            mu, grad_eta, w_fisher, w_hessian, w_solver, dev))
1794        pirls_refuse(&status, PIRLS_FINAL_OUTPUT);
1795"#
1796    )
1797}
1798
1799#[cfg(target_os = "linux")]
1800fn gamma_log_body(curvature: CurvatureMode) -> String {
1801    // `shape` is a kernel parameter (see `cuda_source_for`); the body reads it
1802    // directly. No local shadowing needed.
1803    let tag = curvature_tag(curvature);
1804    format!(
1805        r#"{tag}    double mu = 0.0, grad_eta = 0.0, w_fisher = 0.0;
1806    double w_hessian = 0.0, w_solver = 0.0, dev = 0.0;
1807    if (!pirls_log_eta_valid(eta_i)) pirls_refuse(&status, PIRLS_ETA_DOMAIN);
1808    if (status == PIRLS_OK && !(isfinite(shape) && shape > 0.0))
1809        pirls_refuse(&status, PIRLS_GAMMA_SHAPE);
1810    if (status == PIRLS_OK && !(isfinite(wp) && wp >= 0.0))
1811        pirls_refuse(&status, PIRLS_PRIOR_WEIGHT);
1812    if (status == PIRLS_OK && wp > 0.0 && !(isfinite(y_i) && y_i > 0.0))
1813        pirls_refuse(&status, PIRLS_RESPONSE);
1814    if (status == PIRLS_OK) {{
1815        mu = exp(eta_i);
1816        if (!(isfinite(mu) && mu > 0.0)) pirls_refuse(&status, PIRLS_INVERSE_LINK);
1817    }}
1818    if (status == PIRLS_OK && wp > 0.0) {{
1819        w_fisher = wp * shape;
1820        if (!(isfinite(w_fisher) && w_fisher > 0.0))
1821            pirls_refuse(&status, PIRLS_FISHER_WEIGHT);
1822#ifdef PIRLS_CURVATURE_OBSERVED
1823        double weighted_ratio_observed = positive_mul_div(w_fisher, y_i, mu);
1824        if (!(isfinite(weighted_ratio_observed) && weighted_ratio_observed > 0.0))
1825            pirls_refuse(&status, PIRLS_OBSERVED_WEIGHT);
1826        w_hessian = weighted_ratio_observed;
1827#else
1828        w_hessian = w_fisher;
1829#endif
1830        if (!isfinite(w_hessian)) pirls_refuse(&status, PIRLS_OBSERVED_WEIGHT);
1831        w_solver = w_hessian;
1832        double u = (y_i - mu) / mu;
1833        double scaled_unit = w_fisher * gamma_unit_deviance_near_one(u);
1834        bool need_weighted_ratio = !isfinite(u)
1835            || !(isfinite(scaled_unit) && scaled_unit >= 0.0);
1836        double weighted_ratio = 0.0;
1837#ifdef PIRLS_CURVATURE_OBSERVED
1838        weighted_ratio = weighted_ratio_observed;
1839#else
1840        if (need_weighted_ratio)
1841            weighted_ratio = positive_mul_div(w_fisher, y_i, mu);
1842#endif
1843        grad_eta = isfinite(u) ? w_fisher * u : weighted_ratio - w_fisher;
1844        double dev_base;
1845        if (isfinite(scaled_unit) && scaled_unit >= 0.0) {{
1846            dev_base = scaled_unit;
1847        }} else {{
1848            dev_base = weighted_ratio - w_fisher * (1.0 + log(y_i) - eta_i);
1849        }}
1850        if (!isfinite(grad_eta)) pirls_refuse(&status, PIRLS_GRADIENT);
1851        dev = 2.0 * dev_base;
1852        if (!isfinite(dev)) pirls_refuse(&status, PIRLS_DEVIANCE);
1853    }}
1854    if (status == PIRLS_OK && !pirls_outputs_finite(
1855            mu, grad_eta, w_fisher, w_hessian, w_solver, dev))
1856        pirls_refuse(&status, PIRLS_FINAL_OUTPUT);
1857"#
1858    )
1859}
1860
1861#[cfg(target_os = "linux")]
1862fn bernoulli_logit_body(curvature: CurvatureMode) -> String {
1863    let tag = curvature_tag(curvature);
1864    format!(
1865        r#"{tag}    double mu = 0.0, grad_eta = 0.0, w_fisher = 0.0;
1866    double w_hessian = 0.0, w_solver = 0.0, dev = 0.0;
1867    if (!isfinite(eta_i)) pirls_refuse(&status, PIRLS_ETA_DOMAIN);
1868    if (status == PIRLS_OK && !(isfinite(wp) && wp >= 0.0))
1869        pirls_refuse(&status, PIRLS_PRIOR_WEIGHT);
1870    if (status == PIRLS_OK && wp > 0.0
1871            && !(isfinite(y_i) && y_i >= 0.0 && y_i <= 1.0))
1872        pirls_refuse(&status, PIRLS_RESPONSE);
1873    double tail = exp(-fabs(eta_i));
1874    double denom = 1.0 + tail;
1875    double dmu_deta = tail / (denom * denom);
1876    if (status == PIRLS_OK) {{
1877        mu = eta_i >= 0.0 ? 1.0 / denom : tail / denom;
1878        if (!(isfinite(mu) && mu >= 0.0 && mu <= 1.0
1879                && isfinite(dmu_deta) && dmu_deta > 0.0))
1880            pirls_refuse(&status, PIRLS_INVERSE_LINK);
1881    }}
1882    if (status == PIRLS_OK && wp > 0.0) {{
1883        double residual;
1884        if (eta_i >= 0.0) {{
1885            double one_minus_mu = tail / denom;
1886            residual = y_i == 1.0 ? one_minus_mu : (y_i - 1.0) + one_minus_mu;
1887        }} else {{
1888            residual = y_i - mu;
1889        }}
1890        w_fisher = wp * dmu_deta;
1891        if (!(isfinite(w_fisher) && w_fisher > 0.0))
1892            pirls_refuse(&status, PIRLS_FISHER_WEIGHT);
1893        w_hessian = w_fisher;
1894        w_solver = w_hessian;
1895        grad_eta = wp * residual;
1896        if (!isfinite(grad_eta)) pirls_refuse(&status, PIRLS_GRADIENT);
1897        dev = logit_deviance(y_i, eta_i, wp);
1898        if (!isfinite(dev)) pirls_refuse(&status, PIRLS_DEVIANCE);
1899    }}
1900    if (status == PIRLS_OK && !pirls_outputs_finite(
1901            mu, grad_eta, w_fisher, w_hessian, w_solver, dev))
1902        pirls_refuse(&status, PIRLS_FINAL_OUTPUT);
1903"#
1904    )
1905}
1906
1907#[cfg(target_os = "linux")]
1908fn bernoulli_probit_body(curvature: CurvatureMode) -> String {
1909    let tag = curvature_tag(curvature);
1910    format!(
1911        r#"{tag}    double mu = 0.0, grad_eta = 0.0, w_fisher = 0.0;
1912    double w_hessian = 0.0, w_solver = 0.0, dev = 0.0;
1913    if (!isfinite(eta_i)) pirls_refuse(&status, PIRLS_ETA_DOMAIN);
1914    if (status == PIRLS_OK && !(isfinite(wp) && wp >= 0.0))
1915        pirls_refuse(&status, PIRLS_PRIOR_WEIGHT);
1916    if (status == PIRLS_OK && wp > 0.0
1917            && !(isfinite(y_i) && y_i >= 0.0 && y_i <= 1.0))
1918        pirls_refuse(&status, PIRLS_RESPONSE);
1919    double dmu_deta = 0.0, d2mu_deta2 = 0.0, v = 0.0;
1920    if (status == PIRLS_OK) {{
1921        mu = std_norm_cdf(eta_i);
1922        dmu_deta = std_norm_pdf(eta_i);
1923        d2mu_deta2 = -eta_i * dmu_deta;
1924        if (!(isfinite(mu) && mu > 0.0 && mu < 1.0
1925                && isfinite(dmu_deta) && dmu_deta > 0.0
1926                && isfinite(d2mu_deta2)))
1927            pirls_refuse(&status, PIRLS_INVERSE_LINK);
1928    }}
1929    if (status == PIRLS_OK && wp > 0.0) {{
1930        v = mu * (1.0 - mu);
1931        double fisher_per_prior = dmu_deta * dmu_deta / v;
1932        w_fisher = wp * fisher_per_prior;
1933        if (!(isfinite(v) && v > 0.0 && isfinite(fisher_per_prior)
1934                && fisher_per_prior > 0.0 && isfinite(w_fisher) && w_fisher > 0.0))
1935            pirls_refuse(&status, PIRLS_FISHER_WEIGHT);
1936        double resid = y_i - mu;
1937#ifdef PIRLS_CURVATURE_OBSERVED
1938        double bracket = d2mu_deta2 / v
1939            - (dmu_deta * dmu_deta) * (1.0 - 2.0 * mu) / (v * v);
1940        w_hessian = w_fisher - wp * resid * bracket;
1941#else
1942        w_hessian = w_fisher;
1943#endif
1944        if (!isfinite(w_hessian)) pirls_refuse(&status, PIRLS_OBSERVED_WEIGHT);
1945        w_solver = w_hessian;
1946        grad_eta = wp * resid * dmu_deta / v;
1947        if (!isfinite(grad_eta)) pirls_refuse(&status, PIRLS_GRADIENT);
1948        dev = bernoulli_deviance(y_i, mu, wp);
1949        if (!isfinite(dev)) pirls_refuse(&status, PIRLS_DEVIANCE);
1950    }}
1951    if (status == PIRLS_OK && !pirls_outputs_finite(
1952            mu, grad_eta, w_fisher, w_hessian, w_solver, dev))
1953        pirls_refuse(&status, PIRLS_FINAL_OUTPUT);
1954"#
1955    )
1956}
1957
1958#[cfg(target_os = "linux")]
1959fn bernoulli_cloglog_body(curvature: CurvatureMode) -> String {
1960    let tag = curvature_tag(curvature);
1961    format!(
1962        r#"{tag}    double mu = 0.0, grad_eta = 0.0, w_fisher = 0.0;
1963    double w_hessian = 0.0, w_solver = 0.0, dev = 0.0;
1964    if (!isfinite(eta_i)) pirls_refuse(&status, PIRLS_ETA_DOMAIN);
1965    if (status == PIRLS_OK && !(isfinite(wp) && wp >= 0.0))
1966        pirls_refuse(&status, PIRLS_PRIOR_WEIGHT);
1967    if (status == PIRLS_OK && wp > 0.0
1968            && !(isfinite(y_i) && y_i >= 0.0 && y_i <= 1.0))
1969        pirls_refuse(&status, PIRLS_RESPONSE);
1970    double inner = 0.0, dmu_deta = 0.0, d2mu_deta2 = 0.0, v = 0.0;
1971    if (status == PIRLS_OK) {{
1972        inner = exp(eta_i);
1973        double complement = exp(-inner);
1974        mu = -expm1(-inner);
1975        dmu_deta = inner * complement;
1976        d2mu_deta2 = dmu_deta * (1.0 - inner);
1977        if (!(isfinite(mu) && mu > 0.0 && mu < 1.0
1978                && isfinite(dmu_deta) && dmu_deta > 0.0
1979                && isfinite(d2mu_deta2)))
1980            pirls_refuse(&status, PIRLS_INVERSE_LINK);
1981    }}
1982    if (status == PIRLS_OK && wp > 0.0) {{
1983        v = mu * (1.0 - mu);
1984        double fisher_per_prior = dmu_deta * dmu_deta / v;
1985        w_fisher = wp * fisher_per_prior;
1986        if (!(isfinite(v) && v > 0.0 && isfinite(fisher_per_prior)
1987                && fisher_per_prior > 0.0 && isfinite(w_fisher) && w_fisher > 0.0))
1988            pirls_refuse(&status, PIRLS_FISHER_WEIGHT);
1989        double resid = y_i - mu;
1990#ifdef PIRLS_CURVATURE_OBSERVED
1991        double bracket = d2mu_deta2 / v
1992            - (dmu_deta * dmu_deta) * (1.0 - 2.0 * mu) / (v * v);
1993        w_hessian = w_fisher - wp * resid * bracket;
1994#else
1995        w_hessian = w_fisher;
1996#endif
1997        if (!isfinite(w_hessian)) pirls_refuse(&status, PIRLS_OBSERVED_WEIGHT);
1998        w_solver = w_hessian;
1999        grad_eta = wp * resid * dmu_deta / v;
2000        if (!isfinite(grad_eta)) pirls_refuse(&status, PIRLS_GRADIENT);
2001        dev = bernoulli_deviance(y_i, mu, wp);
2002        if (!isfinite(dev)) pirls_refuse(&status, PIRLS_DEVIANCE);
2003    }}
2004    if (status == PIRLS_OK && !pirls_outputs_finite(
2005            mu, grad_eta, w_fisher, w_hessian, w_solver, dev))
2006        pirls_refuse(&status, PIRLS_FINAL_OUTPUT);
2007"#
2008    )
2009}
2010
2011// ────────────────────────────────────────────────────────────────────────
2012// solve-row CUDA source (4-output variant)
2013// ────────────────────────────────────────────────────────────────────────
2014
2015/// Build the solve-row CUDA source for `(family, curvature)`.
2016///
2017/// The kernel has a reduced signature:
2018///   `(int n, const f64* eta, const f64* y, const f64* prior_w,
2019///     f64* grad_eta_out, f64* w_solver_out, f64* deviance_out, u32* status_out)`
2020///
2021/// It executes the same per-family math as `cuda_source_for` but skips the
2022/// `mu` and `w_hessian` stores, reducing
2023/// both bandwidth and L1/register pressure in the hot Newton iteration.
2024#[cfg(target_os = "linux")]
2025fn solve_row_source_for(family: PirlsRowFamily, curvature: CurvatureMode) -> String {
2026    let body = match family {
2027        PirlsRowFamily::GaussianIdentity => gaussian_identity_body(curvature),
2028        PirlsRowFamily::PoissonLog => poisson_log_body(curvature),
2029        PirlsRowFamily::GammaLog => gamma_log_body(curvature),
2030        PirlsRowFamily::BernoulliLogit => bernoulli_logit_body(curvature),
2031        PirlsRowFamily::BernoulliProbit => bernoulli_probit_body(curvature),
2032        PirlsRowFamily::BernoulliCLogLog => bernoulli_cloglog_body(curvature),
2033    };
2034    let kernel_name = family.solve_kernel_name();
2035    let curvature_define = match curvature {
2036        CurvatureMode::Fisher => "#define PIRLS_CURVATURE_FISHER 1",
2037        CurvatureMode::Observed => "#define PIRLS_CURVATURE_OBSERVED 1",
2038    };
2039    // GammaLog solve kernel also takes `double shape` after `prior_w`.
2040    let shape_param = if matches!(family, PirlsRowFamily::GammaLog) {
2041        "    double         shape,\n"
2042    } else {
2043        ""
2044    };
2045    format!(
2046        r#"
2047{curvature_define}
2048{prolog}
2049
2050extern "C" __global__ void {kernel_name}(
2051    int            n,
2052    const double* __restrict__ eta,
2053    const double* __restrict__ y,
2054    const double* __restrict__ prior_w,
2055{shape_param}    double* __restrict__ grad_eta_out,
2056    double* __restrict__ w_solver_out,
2057    double* __restrict__ deviance_out,
2058    unsigned int* __restrict__ status_out
2059) {{
2060    int i = blockIdx.x * blockDim.x + threadIdx.x;
2061    if (i >= n) return;
2062    unsigned int status = PIRLS_OK;
2063    double eta_i = eta[i];
2064    double y_i = y[i];
2065    double wp = prior_w[i];
2066{body}
2067    if (status == PIRLS_OK) {{
2068        grad_eta_out[i] = grad_eta;
2069        w_solver_out[i] = w_solver;
2070        deviance_out[i] = dev;
2071    }}
2072    status_out[i] = status;
2073}}
2074"#,
2075        prolog = common_device_prolog(),
2076    )
2077}
2078
2079// ────────────────────────────────────────────────────────────────────────
2080// alpha-ladder CUDA source (fused all-alpha candidate-objective kernel)
2081// ────────────────────────────────────────────────────────────────────────
2082
2083/// The alpha constants embedded into the ladder kernel source as a
2084/// `__constant__` array. Must stay in sync with [`ALPHA_LADDER`].
2085#[cfg(target_os = "linux")]
2086const ALPHA_LADDER_CUDA_ARRAY: &str =
2087    "__constant__ double PIRLS_ALPHAS[7] = {1.0, 0.5, 0.25, 0.125, 0.0625, 0.03125, 0.015625};";
2088
2089/// Build the fused alpha-ladder CUDA source for `(family, curvature)`.
2090///
2091/// Grid layout: `grid = (row_blocks, ALPHA_LADDER_LEN, 1)`.
2092///   - `blockIdx.y` selects the alpha slot `k`.
2093///   - `blockIdx.x * blockDim.x + threadIdx.x` selects row `i`.
2094///
2095/// Each thread evaluates `eta_trial = eta[i] + PIRLS_ALPHAS[k] * xd[i]`,
2096/// runs the per-family deviance math, and atomically adds to
2097/// `objective_out[k]` (double-precision atomic add) and writes the exact row
2098/// code to `status_out[k*n+i]`.
2099///
2100/// The kernel signature is:
2101///   `(int n, const f64* eta, const f64* xd, const f64* y, const f64* prior_w,
2102///     f64* objective_out, u32* status_out)`
2103#[cfg(target_os = "linux")]
2104fn ladder_source_for(family: PirlsRowFamily, curvature: CurvatureMode) -> String {
2105    let body = match family {
2106        PirlsRowFamily::GaussianIdentity => gaussian_identity_body(curvature),
2107        PirlsRowFamily::PoissonLog => poisson_log_body(curvature),
2108        PirlsRowFamily::GammaLog => gamma_log_body(curvature),
2109        PirlsRowFamily::BernoulliLogit => bernoulli_logit_body(curvature),
2110        PirlsRowFamily::BernoulliProbit => bernoulli_probit_body(curvature),
2111        PirlsRowFamily::BernoulliCLogLog => bernoulli_cloglog_body(curvature),
2112    };
2113    let kernel_name = family.ladder_kernel_name();
2114    let curvature_define = match curvature {
2115        CurvatureMode::Fisher => "#define PIRLS_CURVATURE_FISHER 1",
2116        CurvatureMode::Observed => "#define PIRLS_CURVATURE_OBSERVED 1",
2117    };
2118    // The body uses the name `eta_i` for the exact trial linear predictor.
2119    // For the ladder we substitute `eta[i] + alpha * xd[i]` as the trial eta,
2120    // so we define `eta_i` before the body runs. The body's own local variable
2121    // For GaussianIdentity the body reads `eta_i` directly as `mu`; log-link
2122    // bodies certify the shared solver interval and never modify the value.
2123    // GammaLog ladder kernel also takes `double shape` after `prior_w`.
2124    let shape_param = if matches!(family, PirlsRowFamily::GammaLog) {
2125        "    double         shape,\n"
2126    } else {
2127        ""
2128    };
2129    format!(
2130        r#"
2131{curvature_define}
2132{prolog}
2133{alphas}
2134
2135extern "C" __global__ void {kernel_name}(
2136    int            n,
2137    const double* __restrict__ eta,
2138    const double* __restrict__ xd,
2139    const double* __restrict__ y,
2140    const double* __restrict__ prior_w,
2141{shape_param}    double* __restrict__ objective_out,
2142    unsigned int* __restrict__ status_out
2143) {{
2144    int i = blockIdx.x * blockDim.x + threadIdx.x;
2145    int k = (int)blockIdx.y;
2146    if (i >= n) return;
2147    unsigned int status = PIRLS_OK;
2148    double alpha = PIRLS_ALPHAS[k];
2149    double eta_i = eta[i] + alpha * xd[i];
2150    double y_i = y[i];
2151    double wp = prior_w[i];
2152{body}
2153    if (status == PIRLS_OK) atomicAdd(&objective_out[k], dev);
2154    status_out[k * n + i] = status;
2155}}
2156"#,
2157        prolog = common_device_prolog(),
2158        alphas = ALPHA_LADDER_CUDA_ARRAY,
2159    )
2160}
2161
2162// ────────────────────────────────────────────────────────────────────────
2163// Tests
2164// ────────────────────────────────────────────────────────────────────────
2165
2166#[cfg(test)]
2167#[path = "pirls_row_tests.rs"]
2168mod pirls_row_tests;