1use 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#[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 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 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 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#[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
160pub 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#[derive(Clone, Copy, Debug)]
208pub struct RowInput {
209 pub eta: f64,
210 pub y: f64,
211 pub prior_weight: f64,
212}
213
214#[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
225pub 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
244pub 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#[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#[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#[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#[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 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 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 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 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#[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; COEFF * (-0.5 * x * x).exp()
817}
818
819#[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#[cfg(target_os = "linux")]
838#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
839enum KernelMode {
840 FinalRow,
843 SolveRow,
845 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 #[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 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 #[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 #[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 #[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#[cfg(target_os = "linux")]
986#[derive(Clone, Debug)]
1002pub struct JitFamilySpec {
1003 pub spec_id: u64,
1007 pub body: String,
1012}
1013
1014#[cfg(target_os = "linux")]
1015impl JitFamilySpec {
1016 #[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 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#[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 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#[cfg(target_os = "linux")]
1106pub struct SolveRowBuffers {
1107 pub grad_eta: cudarc::driver::CudaSlice<f64>,
1109 pub w_solver: cudarc::driver::CudaSlice<f64>,
1111 pub deviance: cudarc::driver::CudaSlice<f64>,
1113 pub status: cudarc::driver::CudaSlice<u32>,
1115 pub n: usize,
1116}
1117
1118#[cfg(target_os = "linux")]
1119impl SolveRowBuffers {
1120 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
1142pub const ALPHA_LADDER_LEN: usize = 7;
1144
1145pub const ALPHA_LADDER: [f64; ALPHA_LADDER_LEN] =
1147 [1.0, 0.5, 0.25, 0.125, 0.0625, 0.03125, 0.015625];
1148
1149#[cfg(target_os = "linux")]
1159pub struct AlphaLadderDevBuffers {
1160 pub objective_dev: cudarc::driver::CudaSlice<f64>,
1162 pub status_dev: cudarc::driver::CudaSlice<u32>,
1165 pub n: usize,
1166}
1167
1168#[cfg(target_os = "linux")]
1169impl AlphaLadderDevBuffers {
1170 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 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#[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 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 unsafe { builder.launch(cfg) }
1269 .gpu_ctx_with(|err| format!("row reweight launch({}): {err}", family.kernel_name()))?;
1270 Ok(())
1271}
1272
1273#[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 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 unsafe { builder.launch(cfg) }
1341 .gpu_ctx_with(|err| format!("solve-row launch({kernel_name}): {err}"))?;
1342 Ok(())
1343}
1344
1345#[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 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 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 unsafe { builder.launch(cfg) }
1414 .gpu_ctx_with(|err| format!("alpha-ladder launch({kernel_name}): {err}"))?;
1415 Ok(())
1416}
1417
1418#[cfg(target_os = "linux")]
1426fn common_device_prolog() -> String {
1427 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#[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 let curvature_define = match curvature {
1661 CurvatureMode::Fisher => "#define PIRLS_CURVATURE_FISHER 1",
1662 CurvatureMode::Observed => "#define PIRLS_CURVATURE_OBSERVED 1",
1663 };
1664 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#[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 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#[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 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#[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#[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 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#[cfg(test)]
2167#[path = "pirls_row_tests.rs"]
2168mod pirls_row_tests;