1use std::sync::OnceLock;
46
47use gam_gpu::gpu_error::GpuError;
48#[cfg(target_os = "linux")]
49use gam_gpu::gpu_error::GpuResultExt;
50use gam_problem::EstimationError;
51
52#[cfg(target_os = "linux")]
53use std::sync::{Arc, Mutex};
54
55#[cfg(target_os = "linux")]
56use cudarc::driver::{CudaContext, CudaModule};
57
58#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
68pub enum PirlsRowFamily {
69 BernoulliLogit,
70 BernoulliProbit,
71 BernoulliCLogLog,
72 PoissonLog,
73 GaussianIdentity,
74 GammaLog,
75}
76
77impl PirlsRowFamily {
78 pub const ALL: [Self; 6] = [
79 Self::BernoulliLogit,
80 Self::BernoulliProbit,
81 Self::BernoulliCLogLog,
82 Self::PoissonLog,
83 Self::GaussianIdentity,
84 Self::GammaLog,
85 ];
86
87 pub const fn as_str(self) -> &'static str {
88 match self {
89 Self::BernoulliLogit => "bernoulli-logit",
90 Self::BernoulliProbit => "bernoulli-probit",
91 Self::BernoulliCLogLog => "bernoulli-cloglog",
92 Self::PoissonLog => "poisson-log",
93 Self::GaussianIdentity => "gaussian-identity",
94 Self::GammaLog => "gamma-log",
95 }
96 }
97
98 pub const fn kernel_name(self) -> &'static str {
100 match self {
101 Self::BernoulliLogit => "pirls_row_bernoulli_logit",
102 Self::BernoulliProbit => "pirls_row_bernoulli_probit",
103 Self::BernoulliCLogLog => "pirls_row_bernoulli_cloglog",
104 Self::PoissonLog => "pirls_row_poisson_log",
105 Self::GaussianIdentity => "pirls_row_gaussian_identity",
106 Self::GammaLog => "pirls_row_gamma_log",
107 }
108 }
109
110 pub const fn solve_kernel_name(self) -> &'static str {
113 match self {
114 Self::BernoulliLogit => "pirls_solve_bernoulli_logit",
115 Self::BernoulliProbit => "pirls_solve_bernoulli_probit",
116 Self::BernoulliCLogLog => "pirls_solve_bernoulli_cloglog",
117 Self::PoissonLog => "pirls_solve_poisson_log",
118 Self::GaussianIdentity => "pirls_solve_gaussian_identity",
119 Self::GammaLog => "pirls_solve_gamma_log",
120 }
121 }
122
123 pub const fn ladder_kernel_name(self) -> &'static str {
127 match self {
128 Self::BernoulliLogit => "pirls_ladder_bernoulli_logit",
129 Self::BernoulliProbit => "pirls_ladder_bernoulli_probit",
130 Self::BernoulliCLogLog => "pirls_ladder_bernoulli_cloglog",
131 Self::PoissonLog => "pirls_ladder_poisson_log",
132 Self::GaussianIdentity => "pirls_ladder_gaussian_identity",
133 Self::GammaLog => "pirls_ladder_gamma_log",
134 }
135 }
136}
137
138#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
145pub enum CurvatureMode {
146 Fisher,
147 Observed,
148}
149
150impl CurvatureMode {
151 pub const fn as_str(self) -> &'static str {
152 match self {
153 Self::Fisher => "fisher",
154 Self::Observed => "observed",
155 }
156 }
157}
158
159pub mod status_codes {
165 pub const OK: u32 = 0;
166 pub const ETA_DOMAIN: u32 = 1;
167 pub const PRIOR_WEIGHT: u32 = 2;
168 pub const RESPONSE: u32 = 3;
169 pub const GAMMA_SHAPE: u32 = 4;
170 pub const INVERSE_LINK: u32 = 5;
171 pub const FISHER_WEIGHT: u32 = 6;
172 pub const OBSERVED_WEIGHT: u32 = 7;
173 pub const GRADIENT: u32 = 8;
174 pub const DEVIANCE: u32 = 9;
175 pub const FINAL_OUTPUT: u32 = 10;
176
177 pub const fn quantity(code: u32) -> &'static str {
178 match code {
179 ETA_DOMAIN => "inverse-link eta domain",
180 PRIOR_WEIGHT => "prior weight",
181 RESPONSE => "response",
182 GAMMA_SHAPE => "Gamma shape",
183 INVERSE_LINK => "inverse-link jet",
184 FISHER_WEIGHT => "Fisher weight",
185 OBSERVED_WEIGHT => "observed Hessian weight",
186 GRADIENT => "eta gradient",
187 DEVIANCE => "deviance contribution",
188 FINAL_OUTPUT => "final row output",
189 _ => "unknown GPU PIRLS refusal",
190 }
191 }
192}
193
194#[derive(Clone, Copy, Debug)]
207pub struct RowInput {
208 pub eta: f64,
209 pub y: f64,
210 pub prior_weight: f64,
211}
212
213#[derive(Clone, Copy, Debug, Default)]
215pub struct RowOutput {
216 pub mu: f64,
217 pub grad_eta: f64,
218 pub w_fisher: f64,
219 pub w_hessian: f64,
220 pub w_solver: f64,
221 pub deviance: f64,
222}
223
224pub fn row_reweight_cpu(
230 family: PirlsRowFamily,
231 mode: CurvatureMode,
232 input: RowInput,
233 gamma_shape: f64,
234) -> Result<RowOutput, EstimationError> {
235 row_reweight_cpu_at(0, family, mode, input, gamma_shape)
236}
237
238pub fn row_reweight_cpu_at(
241 row: usize,
242 family: PirlsRowFamily,
243 mode: CurvatureMode,
244 input: RowInput,
245 gamma_shape: f64,
246) -> Result<RowOutput, EstimationError> {
247 match family {
248 PirlsRowFamily::GaussianIdentity => row_gaussian_identity(row, input, mode),
249 PirlsRowFamily::PoissonLog => row_poisson_log(row, input, mode),
250 PirlsRowFamily::GammaLog => row_gamma_log(row, input, mode, gamma_shape),
251 PirlsRowFamily::BernoulliLogit => row_bernoulli_logit(row, input, mode),
252 PirlsRowFamily::BernoulliProbit => row_bernoulli_probit(row, input, mode),
253 PirlsRowFamily::BernoulliCLogLog => row_bernoulli_cloglog(row, input, mode),
254 }
255}
256
257pub fn replay_first_refusal(
262 family: PirlsRowFamily,
263 mode: CurvatureMode,
264 gamma_shape: f64,
265 eta: &[f64],
266 y: &[f64],
267 prior_weight: &[f64],
268 status: &[u32],
269) -> Result<(), EstimationError> {
270 let n = eta.len();
271 if y.len() != n || prior_weight.len() != n || status.len() != n {
272 return Err(EstimationError::InvalidInput(format!(
273 "GPU PIRLS refusal replay length mismatch: eta={n}, y={}, prior_weight={}, status={}",
274 y.len(),
275 prior_weight.len(),
276 status.len(),
277 )));
278 }
279 let Some((row, &code)) = status
280 .iter()
281 .enumerate()
282 .find(|(_, code)| **code != status_codes::OK)
283 else {
284 return Ok(());
285 };
286 let input = RowInput {
287 eta: eta[row],
288 y: y[row],
289 prior_weight: prior_weight[row],
290 };
291 match row_reweight_cpu_at(row, family, mode, input, gamma_shape) {
292 Err(error) => Err(error),
293 Ok(_) => Err(row_error(
294 row,
295 status_codes::quantity(code),
296 input.eta,
297 f64::from(code),
298 )),
299 }
300}
301
302#[inline]
308fn select_w_hessian(mode: CurvatureMode, w_fisher: f64, observed_correction: f64) -> f64 {
309 match mode {
310 CurvatureMode::Fisher => w_fisher,
311 CurvatureMode::Observed => w_fisher + observed_correction,
312 }
313}
314
315#[inline]
316fn row_error(row: usize, quantity: &'static str, eta: f64, value: f64) -> EstimationError {
317 EstimationError::PirlsRowGeometryUnrepresentable {
318 row,
319 quantity,
320 eta,
321 value,
322 }
323}
324
325#[inline]
326fn finite_eta(link: &'static str, eta: f64) -> Result<(), EstimationError> {
327 if eta.is_finite() {
328 Ok(())
329 } else {
330 Err(EstimationError::InverseLinkDomainViolation {
331 link,
332 eta,
333 lower: -f64::MAX,
334 upper: f64::MAX,
335 })
336 }
337}
338
339#[inline]
340fn prior_weight(row: usize, input: RowInput) -> Result<f64, EstimationError> {
341 if input.prior_weight.is_finite() && input.prior_weight >= 0.0 {
342 Ok(input.prior_weight)
343 } else {
344 Err(row_error(
345 row,
346 "prior weight",
347 input.eta,
348 input.prior_weight,
349 ))
350 }
351}
352
353#[inline]
354fn certify_output(row: usize, eta: f64, output: RowOutput) -> Result<RowOutput, EstimationError> {
355 for (quantity, value) in [
356 ("mean", output.mu),
357 ("eta gradient", output.grad_eta),
358 ("Fisher weight", output.w_fisher),
359 ("observed Hessian weight", output.w_hessian),
360 ("solver Hessian weight", output.w_solver),
361 ("deviance contribution", output.deviance),
362 ] {
363 if !value.is_finite() {
364 return Err(row_error(row, quantity, eta, value));
365 }
366 }
367 Ok(output)
368}
369
370#[inline]
375fn positive_mul_div(a: f64, b: f64, c: f64) -> f64 {
376 let product = a * b;
377 if product.is_finite() && product > 0.0 {
378 let value = product / c;
379 if value.is_finite() && value > 0.0 {
380 return value;
381 }
382 }
383 let quotient_a = a / c;
384 if quotient_a.is_finite() && quotient_a > 0.0 {
385 let value = quotient_a * b;
386 if value.is_finite() && value > 0.0 {
387 return value;
388 }
389 }
390 let quotient_b = b / c;
391 if quotient_b.is_finite() && quotient_b > 0.0 {
392 let value = quotient_b * a;
393 if value.is_finite() && value > 0.0 {
394 return value;
395 }
396 }
397 product / c
398}
399
400#[inline]
402fn gamma_unit_deviance_near_one(u: f64) -> f64 {
403 if u.abs() > 0.125 {
404 return u - u.ln_1p();
405 }
406 let mut power = u * u;
407 let mut sum = 0.5 * power;
408 for degree in 3..=32 {
409 power *= u;
410 let term = power / f64::from(degree);
411 let next = if degree % 2 == 0 {
412 sum + term
413 } else {
414 sum - term
415 };
416 if next == sum {
417 break;
418 }
419 sum = next;
420 }
421 sum
422}
423
424#[inline]
426fn poisson_unit_deviance_near_one(u: f64) -> f64 {
427 if u.abs() > 0.125 {
428 return (1.0 + u) * u.ln_1p() - u;
429 }
430 let mut power = u * u;
431 let mut sum = 0.5 * power;
432 for degree in 3..=32 {
433 power *= u;
434 let coefficient =
435 if degree % 2 == 0 { 1.0 } else { -1.0 } / (f64::from(degree) * f64::from(degree - 1));
436 let next = sum + coefficient * power;
437 if next == sum {
438 break;
439 }
440 sum = next;
441 }
442 sum
443}
444
445#[inline]
446fn row_gaussian_identity(
447 row: usize,
448 input: RowInput,
449 mode: CurvatureMode,
450) -> Result<RowOutput, EstimationError> {
451 finite_eta("standard identity inverse link", input.eta)?;
452 let w = prior_weight(row, input)?;
453 let mu = input.eta;
454 if w > 0.0 && !input.y.is_finite() {
455 return Err(row_error(row, "Gaussian response", input.eta, input.y));
456 }
457 let resid = input.y - mu;
458 let (grad_eta, dev) = if w == 0.0 {
459 (0.0, 0.0)
460 } else {
461 (w * resid, w * resid * resid)
462 };
463 let w_hessian = select_w_hessian(mode, w, 0.0);
464 certify_output(
465 row,
466 input.eta,
467 RowOutput {
468 mu,
469 grad_eta,
470 w_fisher: w,
471 w_hessian,
472 w_solver: w_hessian,
473 deviance: dev,
474 },
475 )
476}
477
478#[inline]
479fn row_poisson_log(
480 row: usize,
481 input: RowInput,
482 mode: CurvatureMode,
483) -> Result<RowOutput, EstimationError> {
484 let mu = crate::mixture_link::log_link_solver_exp(input.eta)?;
485 let w_prior = prior_weight(row, input)?;
486 if w_prior > 0.0 && !(input.y.is_finite() && input.y >= 0.0) {
487 return Err(row_error(row, "Poisson response", input.eta, input.y));
488 }
489 if w_prior == 0.0 {
490 return certify_output(
491 row,
492 input.eta,
493 RowOutput {
494 mu,
495 ..RowOutput::default()
496 },
497 );
498 }
499 let w_fisher = w_prior * mu;
500 if !(w_fisher.is_finite() && w_fisher > 0.0) {
501 return Err(row_error(row, "Poisson Fisher weight", input.eta, w_fisher));
502 }
503 let grad_eta = w_prior * (input.y - mu);
504 let u = (input.y - mu) / mu;
505 let dev_base = if input.y == 0.0 {
506 w_fisher
507 } else {
508 let scaled_unit = w_fisher * poisson_unit_deviance_near_one(u);
513 if scaled_unit.is_finite() && scaled_unit >= 0.0 {
514 scaled_unit
515 } else {
516 let weighted_y = positive_mul_div(w_fisher, input.y, mu);
517 weighted_y * (input.y.ln() - input.eta - 1.0) + w_fisher
518 }
519 };
520 let dev = 2.0 * dev_base;
521 let w_hessian = select_w_hessian(mode, w_fisher, 0.0);
522 certify_output(
523 row,
524 input.eta,
525 RowOutput {
526 mu,
527 grad_eta,
528 w_fisher,
529 w_hessian,
530 w_solver: w_hessian,
531 deviance: dev,
532 },
533 )
534}
535
536#[inline]
537fn row_gamma_log(
538 row: usize,
539 input: RowInput,
540 mode: CurvatureMode,
541 shape: f64,
542) -> Result<RowOutput, EstimationError> {
543 let mu = crate::mixture_link::log_link_solver_exp(input.eta)?;
544 if !(shape.is_finite() && shape > 0.0) {
545 return Err(row_error(row, "Gamma shape", input.eta, shape));
546 }
547 let w_prior = prior_weight(row, input)?;
548 if w_prior > 0.0 && !(input.y.is_finite() && input.y > 0.0) {
549 return Err(row_error(row, "Gamma response", input.eta, input.y));
550 }
551 if w_prior == 0.0 {
552 return certify_output(
553 row,
554 input.eta,
555 RowOutput {
556 mu,
557 ..RowOutput::default()
558 },
559 );
560 }
561 let w_fisher = w_prior * shape;
562 if !(w_fisher.is_finite() && w_fisher > 0.0) {
563 return Err(row_error(row, "Gamma Fisher weight", input.eta, w_fisher));
564 }
565 let observed_ratio = match mode {
566 CurvatureMode::Fisher => None,
567 CurvatureMode::Observed => {
568 let direct = w_fisher * (input.y / mu);
574 let weighted_ratio = if direct.is_finite() && direct > 0.0 {
575 direct
576 } else {
577 positive_mul_div(w_fisher, input.y, mu)
578 };
579 if !(weighted_ratio.is_finite() && weighted_ratio > 0.0) {
580 return Err(row_error(
581 row,
582 "Gamma observed Hessian weight",
583 input.eta,
584 weighted_ratio,
585 ));
586 }
587 Some(weighted_ratio)
588 }
589 };
590 let w_hessian = observed_ratio.unwrap_or(w_fisher);
591 if !w_hessian.is_finite() {
592 return Err(row_error(
593 row,
594 "Gamma observed Hessian weight",
595 input.eta,
596 w_hessian,
597 ));
598 }
599 let u = (input.y - mu) / mu;
600 let scaled_unit = w_fisher * gamma_unit_deviance_near_one(u);
605 let need_weighted_ratio = !u.is_finite() || !(scaled_unit.is_finite() && scaled_unit >= 0.0);
606 let weighted_ratio = if need_weighted_ratio {
607 observed_ratio.unwrap_or_else(|| positive_mul_div(w_fisher, input.y, mu))
608 } else {
609 0.0
610 };
611 let grad_eta = if u.is_finite() {
612 w_fisher * u
613 } else {
614 weighted_ratio - w_fisher
615 };
616 let dev_base = if scaled_unit.is_finite() && scaled_unit >= 0.0 {
617 scaled_unit
618 } else {
619 weighted_ratio - w_fisher * (1.0 + input.y.ln() - input.eta)
620 };
621 let dev = 2.0 * dev_base;
622 certify_output(
623 row,
624 input.eta,
625 RowOutput {
626 mu,
627 grad_eta,
628 w_fisher,
629 w_hessian,
630 w_solver: w_hessian,
631 deviance: dev,
632 },
633 )
634}
635
636#[inline]
637fn bernoulli_response(row: usize, input: RowInput, w: f64) -> Result<(), EstimationError> {
638 if w == 0.0 || (input.y.is_finite() && (0.0..=1.0).contains(&input.y)) {
639 Ok(())
640 } else {
641 Err(row_error(row, "binomial response", input.eta, input.y))
642 }
643}
644
645#[inline]
646fn row_bernoulli_logit(
647 row: usize,
648 input: RowInput,
649 mode: CurvatureMode,
650) -> Result<RowOutput, EstimationError> {
651 finite_eta("standard logit inverse link", input.eta)?;
652 let w_prior = prior_weight(row, input)?;
653 bernoulli_response(row, input, w_prior)?;
654 let tail = (-input.eta.abs()).exp();
655 let denom = 1.0 + tail;
656 let (mu, residual) = if input.eta >= 0.0 {
657 let one_minus_mu = tail / denom;
658 let residual = if input.y == 1.0 {
659 one_minus_mu
660 } else {
661 (input.y - 1.0) + one_minus_mu
662 };
663 (1.0 / denom, residual)
664 } else {
665 let mu = tail / denom;
666 (mu, input.y - mu)
667 };
668 let dmu_deta = tail / (denom * denom);
669 if !(dmu_deta.is_finite() && dmu_deta > 0.0) {
670 return Err(row_error(
671 row,
672 "canonical-logit inverse-link jet",
673 input.eta,
674 dmu_deta,
675 ));
676 }
677 if w_prior == 0.0 {
678 return certify_output(
679 row,
680 input.eta,
681 RowOutput {
682 mu,
683 ..RowOutput::default()
684 },
685 );
686 }
687 let w_fisher = w_prior * dmu_deta;
688 if !(w_fisher.is_finite() && w_fisher > 0.0) {
689 return Err(row_error(row, "logit Fisher weight", input.eta, w_fisher));
690 }
691 let grad_eta = w_prior * residual;
692 let dev = bernoulli_logit_deviance(input.y, input.eta, w_prior);
693 let w_hessian = select_w_hessian(mode, w_fisher, 0.0);
694 certify_output(
695 row,
696 input.eta,
697 RowOutput {
698 mu,
699 grad_eta,
700 w_fisher,
701 w_hessian,
702 w_solver: w_hessian,
703 deviance: dev,
704 },
705 )
706}
707
708#[inline]
709fn row_bernoulli_probit(
710 row: usize,
711 input: RowInput,
712 mode: CurvatureMode,
713) -> Result<RowOutput, EstimationError> {
714 finite_eta("standard probit inverse link", input.eta)?;
715 let d1 = standard_normal_pdf(input.eta);
716 row_bernoulli_noncanonical(
717 row,
718 input,
719 mode,
720 standard_normal_cdf(input.eta),
721 d1,
722 -input.eta * d1,
723 )
724}
725
726#[inline]
727fn row_bernoulli_cloglog(
728 row: usize,
729 input: RowInput,
730 mode: CurvatureMode,
731) -> Result<RowOutput, EstimationError> {
732 finite_eta("standard complementary-log-log inverse link", input.eta)?;
733 let inner = input.eta.exp();
734 let mu = -(-inner).exp_m1();
735 let complement = (-inner).exp();
736 let d1 = inner * complement;
737 row_bernoulli_noncanonical(row, input, mode, mu, d1, d1 * (1.0 - inner))
738}
739
740#[inline]
741fn row_bernoulli_noncanonical(
742 row: usize,
743 input: RowInput,
744 mode: CurvatureMode,
745 mu: f64,
746 d1: f64,
747 d2: f64,
748) -> Result<RowOutput, EstimationError> {
749 let w_prior = prior_weight(row, input)?;
750 bernoulli_response(row, input, w_prior)?;
751 if !(mu.is_finite() && mu > 0.0 && mu < 1.0 && d1.is_finite() && d1 > 0.0 && d2.is_finite()) {
752 return Err(row_error(row, "inverse-link jet", input.eta, mu));
753 }
754 if w_prior == 0.0 {
755 return certify_output(
756 row,
757 input.eta,
758 RowOutput {
759 mu,
760 ..RowOutput::default()
761 },
762 );
763 }
764 let v = mu * (1.0 - mu);
765 let fisher_per_prior = d1 * d1 / v;
766 let w_fisher = w_prior * fisher_per_prior;
767 if !(v.is_finite()
768 && v > 0.0
769 && fisher_per_prior.is_finite()
770 && fisher_per_prior > 0.0
771 && w_fisher.is_finite()
772 && w_fisher > 0.0)
773 {
774 return Err(row_error(
775 row,
776 "Bernoulli Fisher weight",
777 input.eta,
778 w_fisher,
779 ));
780 }
781 let resid = input.y - mu;
782 let grad_eta = w_prior * resid * d1 / v;
783 let bracket = d2 / v - d1 * d1 * (1.0 - 2.0 * mu) / (v * v);
784 let observed_correction = -w_prior * resid * bracket;
786 let w_hessian = select_w_hessian(mode, w_fisher, observed_correction);
787 if !w_hessian.is_finite() {
788 return Err(row_error(
789 row,
790 "Bernoulli observed Hessian weight",
791 input.eta,
792 w_hessian,
793 ));
794 }
795 let dev = bernoulli_deviance(input.y, mu, w_prior);
796 certify_output(
797 row,
798 input.eta,
799 RowOutput {
800 mu,
801 grad_eta,
802 w_fisher,
803 w_hessian,
804 w_solver: w_hessian,
805 deviance: dev,
806 },
807 )
808}
809
810#[inline]
811fn softplus(x: f64) -> f64 {
812 x.max(0.0) + (-x.abs()).exp().ln_1p()
813}
814
815#[inline]
816fn expm1_minus_x(x: f64) -> f64 {
817 if x.abs() > 0.5 {
818 return x.exp_m1() - x;
819 }
820 let mut term = 0.5 * x * x;
821 let mut sum = term;
822 let mut degree = 2.0;
823 loop {
824 degree += 1.0;
825 term *= x / degree;
826 let next = sum + term;
827 if next == sum {
828 return next;
829 }
830 sum = next;
831 }
832}
833
834#[inline]
835fn log1p_minus_x(x: f64) -> f64 {
836 if x.abs() > 0.5 {
837 return x.ln_1p() - x;
838 }
839 let mut power = x * x;
840 let mut sign = -1.0;
841 let mut degree = 2.0;
842 let mut sum = sign * power / degree;
843 loop {
844 power *= x;
845 sign = -sign;
846 degree += 1.0;
847 let next = sum + sign * power / degree;
848 if next == sum {
849 return next;
850 }
851 sum = next;
852 }
853}
854
855#[inline]
856fn logistic(x: f64) -> f64 {
857 if x >= 0.0 {
858 1.0 / (1.0 + (-x).exp())
859 } else {
860 let e = x.exp();
861 e / (1.0 + e)
862 }
863}
864
865#[inline]
868fn bernoulli_kl_from_logits(a: f64, b: f64) -> f64 {
869 if a == b {
870 return 0.0;
871 }
872 let h = b - a;
873 if h.abs() <= 0.5 {
874 let (p, local_h) = if a <= 0.0 {
875 (logistic(a), h)
876 } else {
877 (logistic(-a), -h)
878 };
879 let em1 = local_h.exp_m1();
880 let x = p * em1;
881 return log1p_minus_x(x) + p * expm1_minus_x(local_h);
882 }
883 if a <= 0.0 {
884 let p = logistic(a);
885 p * (a - b) + softplus(b) - softplus(a)
886 } else {
887 let q = logistic(-a);
888 q * (b - a) + softplus(-b) - softplus(-a)
889 }
890}
891
892#[inline]
896fn bd0(x: f64, m: f64) -> f64 {
897 if x == 0.0 {
898 return m;
899 }
900 if x == m {
901 return 0.0;
902 }
903 let hi = x.max(m);
904 let lo = x.min(m);
905 if (x - m).abs() / hi < 0.2 {
906 let v = ((x - m) / hi) / (1.0 + lo / hi);
907 let mut sum = (x - m) * v;
908 let mut term = 2.0 * x * v;
909 let v2 = v * v;
910 let mut denominator = 3.0;
911 loop {
912 term *= v2;
913 let next = sum + term / denominator;
914 if next == sum {
915 return next;
916 }
917 sum = next;
918 denominator += 2.0;
919 }
920 }
921 x * (x.ln() - m.ln()) + (m - x)
922}
923
924#[inline]
925fn bernoulli_logit_deviance(y: f64, eta: f64, w: f64) -> f64 {
926 let unit = if y == 0.0 {
927 softplus(eta)
928 } else if y == 1.0 {
929 softplus(-eta)
930 } else {
931 let response_logit = y.ln() - (-y).ln_1p();
932 bernoulli_kl_from_logits(response_logit, eta)
933 };
934 2.0 * w * unit
935}
936
937#[inline]
938fn bernoulli_deviance(y: f64, mu: f64, w: f64) -> f64 {
939 2.0 * w * (bd0(y, mu) + bd0(1.0 - y, 1.0 - mu))
940}
941
942#[inline]
945fn standard_normal_cdf(x: f64) -> f64 {
946 0.5 * gam_gpu::numerics_host::erfc(-x * std::f64::consts::FRAC_1_SQRT_2)
947}
948
949#[inline]
950fn standard_normal_pdf(x: f64) -> f64 {
951 const COEFF: f64 = 0.398_942_280_401_432_7; COEFF * (-0.5 * x * x).exp()
953}
954
955#[must_use]
961pub struct PirlsRowBackend {
962 #[cfg(target_os = "linux")]
963 inner: PirlsRowBackendLinux,
964}
965
966#[cfg(target_os = "linux")]
967struct PirlsRowBackendLinux {
968 ctx: Arc<CudaContext>,
969 modules: Mutex<std::collections::HashMap<ModuleKey, Arc<CudaModule>>>,
970 jit_modules: Mutex<std::collections::HashMap<JitKey, Arc<CudaModule>>>,
974}
975
976#[cfg(target_os = "linux")]
978#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
979enum KernelMode {
980 FinalRow,
983 SolveRow,
985 AlphaLadder,
987}
988
989#[cfg(target_os = "linux")]
990#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
991struct ModuleKey {
992 family: PirlsRowFamily,
993 curvature: CurvatureMode,
994 mode: KernelMode,
995}
996
997impl PirlsRowBackend {
998 pub const fn compiled() -> bool {
999 cfg!(target_os = "linux")
1000 }
1001
1002 pub fn probe() -> Result<&'static Self, GpuError> {
1003 static BACKEND: OnceLock<Result<PirlsRowBackend, GpuError>> = OnceLock::new();
1004 BACKEND
1005 .get_or_init(|| {
1006 #[cfg(target_os = "linux")]
1007 {
1008 Self::probe_linux()
1009 }
1010 #[cfg(not(target_os = "linux"))]
1011 {
1012 Err(GpuError::DriverLibraryUnavailable {
1013 reason: "pirls_row GPU backend is Linux-only".to_string(),
1014 })
1015 }
1016 })
1017 .as_ref()
1018 .map_err(GpuError::clone)
1019 }
1020
1021 #[cfg(target_os = "linux")]
1022 fn probe_linux() -> Result<Self, GpuError> {
1023 let parts = gam_gpu::backend_probe::probe_cuda_backend("pirls_row")?;
1024 Ok(Self {
1025 inner: PirlsRowBackendLinux {
1026 ctx: parts.ctx,
1027 modules: Mutex::new(std::collections::HashMap::new()),
1028 jit_modules: Mutex::new(std::collections::HashMap::new()),
1029 },
1030 })
1031 }
1032
1033 #[cfg(target_os = "linux")]
1039 fn module_for_kind(
1040 &self,
1041 family: PirlsRowFamily,
1042 curvature: CurvatureMode,
1043 mode: KernelMode,
1044 label: &str,
1045 ) -> Result<Arc<CudaModule>, GpuError> {
1046 let key = ModuleKey {
1047 family,
1048 curvature,
1049 mode,
1050 };
1051 if let Some(existing) = self
1052 .inner
1053 .modules
1054 .lock()
1055 .gpu_ctx_with(|err| format!("pirls_row {label}module cache mutex poisoned: {err}"))?
1056 .get(&key)
1057 {
1058 return Ok(existing.clone());
1059 }
1060 let source = match mode {
1061 KernelMode::FinalRow => cuda_source_for(family, curvature),
1062 KernelMode::SolveRow => solve_row_source_for(family, curvature),
1063 KernelMode::AlphaLadder => ladder_source_for(family, curvature),
1064 };
1065 let ptx = gam_gpu::device_cache::compile_ptx_arch(&source).gpu_ctx_with(|err| {
1069 format!(
1070 "pirls_row {label}NVRTC compile failed for {family}/{curv}: {err}",
1071 family = family.as_str(),
1072 curv = curvature.as_str(),
1073 )
1074 })?;
1075 let module = self
1076 .inner
1077 .ctx
1078 .load_module(ptx)
1079 .gpu_ctx_with(|err| format!("pirls_row {label}module load failed: {err}"))?;
1080 self.inner
1081 .modules
1082 .lock()
1083 .gpu_ctx_with(|err| format!("pirls_row {label}module cache mutex poisoned: {err}"))?
1084 .insert(key, module.clone());
1085 Ok(module)
1086 }
1087
1088 #[cfg(target_os = "linux")]
1091 pub fn module_for(
1092 &self,
1093 family: PirlsRowFamily,
1094 curvature: CurvatureMode,
1095 ) -> Result<Arc<CudaModule>, GpuError> {
1096 self.module_for_kind(family, curvature, KernelMode::FinalRow, "")
1097 }
1098
1099 #[cfg(target_os = "linux")]
1103 pub fn module_for_solve(
1104 &self,
1105 family: PirlsRowFamily,
1106 curvature: CurvatureMode,
1107 ) -> Result<Arc<CudaModule>, GpuError> {
1108 self.module_for_kind(family, curvature, KernelMode::SolveRow, "solve ")
1109 }
1110
1111 #[cfg(target_os = "linux")]
1115 pub fn module_for_ladder(
1116 &self,
1117 family: PirlsRowFamily,
1118 curvature: CurvatureMode,
1119 ) -> Result<Arc<CudaModule>, GpuError> {
1120 self.module_for_kind(family, curvature, KernelMode::AlphaLadder, "ladder ")
1121 }
1122
1123 #[cfg(target_os = "linux")]
1137 pub fn module_for_jit(
1138 &self,
1139 spec: &JitFamilySpec,
1140 curvature: CurvatureMode,
1141 ) -> Result<Arc<CudaModule>, GpuError> {
1142 let key = JitKey {
1147 spec_id: spec.spec_id,
1148 curvature,
1149 };
1150 if let Some(existing) = self
1151 .inner
1152 .jit_modules
1153 .lock()
1154 .gpu_ctx("pirls_row jit cache poisoned")?
1155 .get(&key)
1156 {
1157 return Ok(existing.clone());
1158 }
1159 let source = spec.cuda_source(curvature);
1160 let ptx = gam_gpu::device_cache::compile_ptx_arch(&source).gpu_ctx_with(|err| {
1162 format!(
1163 "pirls_row JIT NVRTC compile failed for spec_id={} curvature={}: {err}",
1164 spec.spec_id,
1165 curvature.as_str(),
1166 )
1167 })?;
1168 let module = self
1169 .inner
1170 .ctx
1171 .load_module(ptx)
1172 .gpu_ctx("pirls_row JIT module load failed")?;
1173 self.inner
1174 .jit_modules
1175 .lock()
1176 .gpu_ctx("pirls_row jit cache poisoned (insert)")?
1177 .insert(key, module.clone());
1178 Ok(module)
1179 }
1180}
1181
1182#[cfg(target_os = "linux")]
1184#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
1185struct JitKey {
1186 spec_id: u64,
1187 curvature: CurvatureMode,
1188}
1189
1190#[derive(Clone, Debug)]
1206pub struct JitFamilySpec {
1207 pub spec_id: u64,
1211 pub body: String,
1216}
1217
1218impl JitFamilySpec {
1219 #[cfg(target_os = "linux")]
1224 pub fn glm(
1225 spec_id: u64,
1226 family: PirlsRowFamily,
1227 curvature: CurvatureMode,
1228 gamma_shape: f64,
1229 ) -> Self {
1230 let mut body = match family {
1231 PirlsRowFamily::GaussianIdentity => gaussian_identity_body(curvature),
1232 PirlsRowFamily::PoissonLog => poisson_log_body(curvature),
1233 PirlsRowFamily::GammaLog => gamma_log_body(curvature),
1234 PirlsRowFamily::BernoulliLogit => bernoulli_logit_body(curvature),
1235 PirlsRowFamily::BernoulliProbit => bernoulli_probit_body(curvature),
1236 PirlsRowFamily::BernoulliCLogLog => bernoulli_cloglog_body(curvature),
1237 };
1238 if matches!(family, PirlsRowFamily::GammaLog) {
1239 body.insert_str(0, &format!(" const double shape = {gamma_shape:?};\n"));
1240 }
1241 Self { spec_id, body }
1242 }
1243
1244 pub fn raw(spec_id: u64, body: impl Into<String>) -> Self {
1248 Self {
1249 spec_id,
1250 body: body.into(),
1251 }
1252 }
1253
1254 pub fn kernel_name(&self) -> String {
1256 format!("pirls_row_jit_{}", self.spec_id)
1257 }
1258
1259 #[cfg(target_os = "linux")]
1264 pub fn cuda_source(&self, curvature: CurvatureMode) -> String {
1265 let curvature_define = match curvature {
1266 CurvatureMode::Fisher => "#define PIRLS_CURVATURE_FISHER 1",
1267 CurvatureMode::Observed => "#define PIRLS_CURVATURE_OBSERVED 1",
1268 };
1269 let kernel_name = self.kernel_name();
1270 let body = &self.body;
1271 format!(
1272 r#"
1273{curvature_define}
1274{prolog}
1275
1276extern "C" __global__ void {kernel_name}(
1277 int n,
1278 const double* __restrict__ eta,
1279 const double* __restrict__ y,
1280 const double* __restrict__ prior_w,
1281 double* __restrict__ mu_out,
1282 double* __restrict__ grad_eta_out,
1283 double* __restrict__ w_hessian_out,
1284 double* __restrict__ w_solver_out,
1285 double* __restrict__ deviance_out,
1286 unsigned int* __restrict__ status_out
1287) {{
1288 int i = blockIdx.x * blockDim.x + threadIdx.x;
1289 if (i >= n) return;
1290 unsigned int status = PIRLS_OK;
1291 double eta_i = eta[i];
1292 double y_i = y[i];
1293 double wp = prior_w[i];
1294{body}
1295 if (status == PIRLS_OK) {{
1296 mu_out[i] = mu;
1297 grad_eta_out[i] = grad_eta;
1298 w_hessian_out[i] = w_hessian;
1299 w_solver_out[i] = w_solver;
1300 deviance_out[i] = dev;
1301 }}
1302 status_out[i] = status;
1303}}
1304"#,
1305 prolog = common_device_prolog(),
1306 )
1307 }
1308}
1309
1310#[cfg(target_os = "linux")]
1318pub struct RowOutputDevBuffers {
1319 pub mu: cudarc::driver::CudaSlice<f64>,
1320 pub grad_eta: cudarc::driver::CudaSlice<f64>,
1321 pub w_hessian: cudarc::driver::CudaSlice<f64>,
1322 pub w_solver: cudarc::driver::CudaSlice<f64>,
1323 pub deviance: cudarc::driver::CudaSlice<f64>,
1324 pub status: cudarc::driver::CudaSlice<u32>,
1325 pub n: usize,
1326}
1327
1328#[cfg(target_os = "linux")]
1329impl RowOutputDevBuffers {
1330 pub fn allocate(stream: &Arc<cudarc::driver::CudaStream>, n: usize) -> Result<Self, GpuError> {
1332 let alloc_f64 = |label: &'static str| {
1333 stream
1334 .alloc_zeros::<f64>(n)
1335 .gpu_ctx_with(|err| format!("pirls_row alloc {label}: {err}"))
1336 };
1337 let alloc_u32 = |label: &'static str| {
1338 stream
1339 .alloc_zeros::<u32>(n)
1340 .gpu_ctx_with(|err| format!("pirls_row alloc {label}: {err}"))
1341 };
1342 Ok(Self {
1343 mu: alloc_f64("mu")?,
1344 grad_eta: alloc_f64("grad_eta")?,
1345 w_hessian: alloc_f64("w_hessian")?,
1346 w_solver: alloc_f64("w_solver")?,
1347 deviance: alloc_f64("deviance")?,
1348 status: alloc_u32("status")?,
1349 n,
1350 })
1351 }
1352}
1353
1354#[cfg(target_os = "linux")]
1363pub struct SolveRowBuffers {
1364 pub grad_eta: cudarc::driver::CudaSlice<f64>,
1366 pub w_solver: cudarc::driver::CudaSlice<f64>,
1368 pub deviance: cudarc::driver::CudaSlice<f64>,
1370 pub status: cudarc::driver::CudaSlice<u32>,
1372 pub n: usize,
1373}
1374
1375#[cfg(target_os = "linux")]
1376impl SolveRowBuffers {
1377 pub fn allocate(stream: &Arc<cudarc::driver::CudaStream>, n: usize) -> Result<Self, GpuError> {
1379 let alloc_f64 = |label: &'static str| {
1380 stream
1381 .alloc_zeros::<f64>(n)
1382 .gpu_ctx_with(|err| format!("pirls_row solve alloc {label}: {err}"))
1383 };
1384 let alloc_u32 = |label: &'static str| {
1385 stream
1386 .alloc_zeros::<u32>(n)
1387 .gpu_ctx_with(|err| format!("pirls_row solve alloc {label}: {err}"))
1388 };
1389 Ok(Self {
1390 grad_eta: alloc_f64("grad_eta")?,
1391 w_solver: alloc_f64("w_solver")?,
1392 deviance: alloc_f64("deviance")?,
1393 status: alloc_u32("status")?,
1394 n,
1395 })
1396 }
1397}
1398
1399pub const ALPHA_LADDER_LEN: usize = 7;
1401
1402pub const ALPHA_LADDER: [f64; ALPHA_LADDER_LEN] =
1404 [1.0, 0.5, 0.25, 0.125, 0.0625, 0.03125, 0.015625];
1405
1406#[cfg(target_os = "linux")]
1416pub struct AlphaLadderDevBuffers {
1417 pub objective_dev: cudarc::driver::CudaSlice<f64>,
1419 pub status_dev: cudarc::driver::CudaSlice<u32>,
1422 pub n: usize,
1423}
1424
1425#[cfg(target_os = "linux")]
1426impl AlphaLadderDevBuffers {
1427 pub fn allocate(stream: &Arc<cudarc::driver::CudaStream>, n: usize) -> Result<Self, GpuError> {
1429 let status_len = ALPHA_LADDER_LEN.checked_mul(n).ok_or_else(|| {
1430 gam_gpu::gpu_err!("pirls_row ladder status length overflows: {ALPHA_LADDER_LEN} * {n}")
1431 })?;
1432 Ok(Self {
1433 objective_dev: stream
1434 .alloc_zeros::<f64>(ALPHA_LADDER_LEN)
1435 .gpu_ctx_with(|err| format!("pirls_row ladder alloc objective: {err}"))?,
1436 status_dev: stream
1437 .alloc_zeros::<u32>(status_len)
1438 .gpu_ctx_with(|err| format!("pirls_row ladder alloc status: {err}"))?,
1439 n,
1440 })
1441 }
1442
1443 pub fn zero(&mut self, stream: &Arc<cudarc::driver::CudaStream>) -> Result<(), GpuError> {
1445 stream
1446 .memset_zeros(&mut self.objective_dev)
1447 .gpu_ctx_with(|err| format!("pirls_row ladder zero objective: {err}"))?;
1448 stream
1449 .memset_zeros(&mut self.status_dev)
1450 .gpu_ctx_with(|err| format!("pirls_row ladder zero status: {err}"))
1451 }
1452}
1453
1454#[cfg(target_os = "linux")]
1468pub fn launch_row_reweight_on_stream(
1469 backend: &PirlsRowBackend,
1470 family: PirlsRowFamily,
1471 curvature: CurvatureMode,
1472 gamma_shape: f64,
1473 stream: &Arc<cudarc::driver::CudaStream>,
1474 n: usize,
1475 eta_dev: &cudarc::driver::CudaSlice<f64>,
1476 y_dev: &cudarc::driver::CudaSlice<f64>,
1477 prior_w_dev: &cudarc::driver::CudaSlice<f64>,
1478 out: &mut RowOutputDevBuffers,
1479) -> Result<(), GpuError> {
1480 use cudarc::driver::{LaunchConfig, PushKernelArg};
1481 if out.n != n {
1482 gam_gpu::gpu_bail!("row reweight buffers shape {} mismatches n={n}", out.n);
1483 }
1484 let module = backend.module_for(family, curvature)?;
1485 let func = module
1486 .load_function(family.kernel_name())
1487 .gpu_ctx_with(|err| {
1488 format!(
1489 "row reweight load_function({}): {err}",
1490 family.kernel_name()
1491 )
1492 })?;
1493 const THREADS_PER_BLOCK: u32 = 256;
1494 let n_u32 = u32::try_from(n)
1495 .map_err(|_| gam_gpu::gpu_err!("n={n} exceeds u32 for row reweight grid sizing"))?;
1496 let grid_x = n_u32.div_ceil(THREADS_PER_BLOCK).max(1);
1497 let n_i32 = i32::try_from(n)
1498 .map_err(|_| gam_gpu::gpu_err!("n={n} exceeds i32 for row reweight kernel argument"))?;
1499 let cfg = LaunchConfig {
1500 grid_dim: (grid_x, 1, 1),
1501 block_dim: (THREADS_PER_BLOCK, 1, 1),
1502 shared_mem_bytes: 0,
1503 };
1504 let mut builder = stream.launch_builder(&func);
1505 builder.arg(&n_i32);
1506 builder.arg(eta_dev);
1507 builder.arg(y_dev);
1508 builder.arg(prior_w_dev);
1509 if matches!(family, PirlsRowFamily::GammaLog) {
1511 builder.arg(&gamma_shape);
1512 }
1513 builder.arg(&mut out.mu);
1514 builder.arg(&mut out.grad_eta);
1515 builder.arg(&mut out.w_hessian);
1516 builder.arg(&mut out.w_solver);
1517 builder.arg(&mut out.deviance);
1518 builder.arg(&mut out.status);
1519 unsafe { builder.launch(cfg) }
1526 .map(|_event_pair| ())
1527 .gpu_ctx_with(|err| format!("row reweight launch({}): {err}", family.kernel_name()))
1528}
1529
1530#[cfg(target_os = "linux")]
1536pub fn launch_row_reweight_jit_on_stream(
1537 backend: &PirlsRowBackend,
1538 spec: &JitFamilySpec,
1539 curvature: CurvatureMode,
1540 stream: &Arc<cudarc::driver::CudaStream>,
1541 n: usize,
1542 eta_dev: &cudarc::driver::CudaSlice<f64>,
1543 y_dev: &cudarc::driver::CudaSlice<f64>,
1544 prior_w_dev: &cudarc::driver::CudaSlice<f64>,
1545 out: &mut RowOutputDevBuffers,
1546) -> Result<(), GpuError> {
1547 use cudarc::driver::{LaunchConfig, PushKernelArg};
1548 if out.n != n {
1549 gam_gpu::gpu_bail!("JIT row reweight buffers shape {} mismatches n={n}", out.n);
1550 }
1551 let module = backend.module_for_jit(spec, curvature)?;
1552 let kernel_name = spec.kernel_name();
1553 let func = module
1554 .load_function(&kernel_name)
1555 .gpu_ctx_with(|err| format!("JIT row reweight load_function({kernel_name}): {err}"))?;
1556 const THREADS_PER_BLOCK: u32 = 256;
1557 let n_u32 = u32::try_from(n)
1558 .map_err(|_| gam_gpu::gpu_err!("n={n} exceeds u32 for JIT row reweight grid sizing"))?;
1559 let grid_x = n_u32.div_ceil(THREADS_PER_BLOCK).max(1);
1560 let n_i32 = i32::try_from(n)
1561 .map_err(|_| gam_gpu::gpu_err!("n={n} exceeds i32 for JIT row reweight kernel argument"))?;
1562 let cfg = LaunchConfig {
1563 grid_dim: (grid_x, 1, 1),
1564 block_dim: (THREADS_PER_BLOCK, 1, 1),
1565 shared_mem_bytes: 0,
1566 };
1567 let mut builder = stream.launch_builder(&func);
1568 builder.arg(&n_i32);
1569 builder.arg(eta_dev);
1570 builder.arg(y_dev);
1571 builder.arg(prior_w_dev);
1572 builder.arg(&mut out.mu);
1573 builder.arg(&mut out.grad_eta);
1574 builder.arg(&mut out.w_hessian);
1575 builder.arg(&mut out.w_solver);
1576 builder.arg(&mut out.deviance);
1577 builder.arg(&mut out.status);
1578 unsafe { builder.launch(cfg) }
1581 .map(|_event_pair| ())
1582 .gpu_ctx_with(|err| format!("JIT row reweight launch({kernel_name}): {err}"))
1583}
1584
1585#[cfg(target_os = "linux")]
1601pub fn launch_solve_row_on_stream(
1602 backend: &PirlsRowBackend,
1603 family: PirlsRowFamily,
1604 curvature: CurvatureMode,
1605 gamma_shape: f64,
1606 stream: &Arc<cudarc::driver::CudaStream>,
1607 n: usize,
1608 eta_dev: &cudarc::driver::CudaSlice<f64>,
1609 y_dev: &cudarc::driver::CudaSlice<f64>,
1610 prior_w_dev: &cudarc::driver::CudaSlice<f64>,
1611 out: &mut SolveRowBuffers,
1612) -> Result<(), GpuError> {
1613 use cudarc::driver::{LaunchConfig, PushKernelArg};
1614 if out.n != n {
1615 gam_gpu::gpu_bail!("solve-row buffers shape {} mismatches n={n}", out.n);
1616 }
1617 let module = backend.module_for_solve(family, curvature)?;
1618 let kernel_name = family.solve_kernel_name();
1619 let func = module
1620 .load_function(kernel_name)
1621 .gpu_ctx_with(|err| format!("solve-row load_function({kernel_name}): {err}"))?;
1622 const THREADS_PER_BLOCK: u32 = 256;
1623 let n_u32 = u32::try_from(n)
1624 .map_err(|_| gam_gpu::gpu_err!("n={n} exceeds u32 for solve-row grid sizing"))?;
1625 let grid_x = n_u32.div_ceil(THREADS_PER_BLOCK).max(1);
1626 let n_i32 = i32::try_from(n)
1627 .map_err(|_| gam_gpu::gpu_err!("n={n} exceeds i32 for solve-row kernel argument"))?;
1628 let cfg = LaunchConfig {
1629 grid_dim: (grid_x, 1, 1),
1630 block_dim: (THREADS_PER_BLOCK, 1, 1),
1631 shared_mem_bytes: 0,
1632 };
1633 let mut builder = stream.launch_builder(&func);
1634 builder.arg(&n_i32);
1635 builder.arg(eta_dev);
1636 builder.arg(y_dev);
1637 builder.arg(prior_w_dev);
1638 if matches!(family, PirlsRowFamily::GammaLog) {
1640 builder.arg(&gamma_shape);
1641 }
1642 builder.arg(&mut out.grad_eta);
1643 builder.arg(&mut out.w_solver);
1644 builder.arg(&mut out.deviance);
1645 builder.arg(&mut out.status);
1646 unsafe { builder.launch(cfg) }
1653 .map(|_event_pair| ())
1654 .gpu_ctx_with(|err| format!("solve-row launch({kernel_name}): {err}"))
1655}
1656
1657#[cfg(target_os = "linux")]
1671pub fn launch_alpha_ladder_on_stream(
1672 backend: &PirlsRowBackend,
1673 family: PirlsRowFamily,
1674 curvature: CurvatureMode,
1675 gamma_shape: f64,
1676 stream: &Arc<cudarc::driver::CudaStream>,
1677 n: usize,
1678 eta_dev: &cudarc::driver::CudaSlice<f64>,
1679 xd_dev: &cudarc::driver::CudaSlice<f64>,
1680 y_dev: &cudarc::driver::CudaSlice<f64>,
1681 prior_w_dev: &cudarc::driver::CudaSlice<f64>,
1682 out: &mut AlphaLadderDevBuffers,
1683) -> Result<(), GpuError> {
1684 use cudarc::driver::{LaunchConfig, PushKernelArg};
1685 if out.n != n {
1686 gam_gpu::gpu_bail!("alpha-ladder buffers shape {} mismatches n={n}", out.n);
1687 }
1688 let module = backend.module_for_ladder(family, curvature)?;
1689 let kernel_name = family.ladder_kernel_name();
1690 let func = module
1691 .load_function(kernel_name)
1692 .gpu_ctx_with(|err| format!("alpha-ladder load_function({kernel_name}): {err}"))?;
1693 const THREADS_PER_BLOCK: u32 = 256;
1694 let n_u32 = u32::try_from(n)
1695 .map_err(|_| gam_gpu::gpu_err!("n={n} exceeds u32 for alpha-ladder grid sizing"))?;
1696 let row_blocks = n_u32.div_ceil(THREADS_PER_BLOCK).max(1);
1697 let n_i32 = i32::try_from(n)
1698 .map_err(|_| gam_gpu::gpu_err!("n={n} exceeds i32 for alpha-ladder kernel argument"))?;
1699 let cfg = LaunchConfig {
1701 grid_dim: (row_blocks, ALPHA_LADDER_LEN as u32, 1),
1702 block_dim: (THREADS_PER_BLOCK, 1, 1),
1703 shared_mem_bytes: 0,
1704 };
1705 let mut builder = stream.launch_builder(&func);
1706 builder.arg(&n_i32);
1707 builder.arg(eta_dev);
1708 builder.arg(xd_dev);
1709 builder.arg(y_dev);
1710 builder.arg(prior_w_dev);
1711 if matches!(family, PirlsRowFamily::GammaLog) {
1713 builder.arg(&gamma_shape);
1714 }
1715 builder.arg(&mut out.objective_dev);
1716 builder.arg(&mut out.status_dev);
1717 unsafe { builder.launch(cfg) }
1726 .map(|_event_pair| ())
1727 .gpu_ctx_with(|err| format!("alpha-ladder launch({kernel_name}): {err}"))
1728}
1729
1730#[cfg(target_os = "linux")]
1738fn common_device_prolog() -> String {
1739 r#"
1743// NVRTC math builtins: prototypes must carry an execution space. Newer
1744// NVRTC (CUDA 12.x JIT semantics) rejects unannotated declarations outright
1745// ("host functions are not allowed in JIT mode"), which failed every
1746// pirls_row kernel compile on real hardware while CPU-only CI stayed green
1747// (#2313 hardware sweep). `__device__` matches how the CUDA math library
1748// declares them; the definitions come from libdevice as before.
1749extern "C" {
1750 __device__ double exp(double);
1751 __device__ double log(double);
1752 __device__ double log1p(double);
1753 __device__ double expm1(double);
1754 __device__ double fabs(double);
1755 __device__ double erfc(double);
1756}
1757
1758static constexpr double PIRLS_LOG_ETA_MIN = __PIRLS_LOG_ETA_MIN__;
1759static constexpr double PIRLS_LOG_ETA_MAX = __PIRLS_LOG_ETA_MAX__;
1760
1761static constexpr unsigned int PIRLS_OK = 0u;
1762static constexpr unsigned int PIRLS_ETA_DOMAIN = 1u;
1763static constexpr unsigned int PIRLS_PRIOR_WEIGHT = 2u;
1764static constexpr unsigned int PIRLS_RESPONSE = 3u;
1765static constexpr unsigned int PIRLS_GAMMA_SHAPE = 4u;
1766static constexpr unsigned int PIRLS_INVERSE_LINK = 5u;
1767static constexpr unsigned int PIRLS_FISHER_WEIGHT = 6u;
1768static constexpr unsigned int PIRLS_OBSERVED_WEIGHT = 7u;
1769static constexpr unsigned int PIRLS_GRADIENT = 8u;
1770static constexpr unsigned int PIRLS_DEVIANCE = 9u;
1771static constexpr unsigned int PIRLS_FINAL_OUTPUT = 10u;
1772
1773__device__ __forceinline__ void pirls_refuse(unsigned int* status, unsigned int code) {
1774 if (*status == PIRLS_OK) *status = code;
1775}
1776
1777__device__ __forceinline__ bool pirls_log_eta_valid(double eta) {
1778 return eta >= PIRLS_LOG_ETA_MIN && eta <= PIRLS_LOG_ETA_MAX;
1779}
1780
1781__device__ __forceinline__ double softplus(double x) {
1782 return (x > 0.0 ? x : 0.0) + log1p(exp(-fabs(x)));
1783}
1784
1785__device__ __forceinline__ double expm1_minus_x(double x) {
1786 if (fabs(x) > 0.5) return expm1(x) - x;
1787 double term = 0.5 * x * x;
1788 double sum = term;
1789 double degree = 2.0;
1790 for (;;) {
1791 degree += 1.0;
1792 term *= x / degree;
1793 double next = sum + term;
1794 if (next == sum) return next;
1795 sum = next;
1796 }
1797}
1798
1799__device__ __forceinline__ double log1p_minus_x(double x) {
1800 if (fabs(x) > 0.5) return log1p(x) - x;
1801 double power = x * x;
1802 double sign = -1.0;
1803 double degree = 2.0;
1804 double sum = sign * power / degree;
1805 for (;;) {
1806 power *= x;
1807 sign = -sign;
1808 degree += 1.0;
1809 double next = sum + sign * power / degree;
1810 if (next == sum) return next;
1811 sum = next;
1812 }
1813}
1814
1815__device__ __forceinline__ double logistic(double x) {
1816 if (x >= 0.0) return 1.0 / (1.0 + exp(-x));
1817 double e = exp(x);
1818 return e / (1.0 + e);
1819}
1820
1821__device__ __forceinline__ double bernoulli_kl_from_logits(double a, double b) {
1822 if (a == b) return 0.0;
1823 double h = b - a;
1824 if (fabs(h) <= 0.5) {
1825 double p = a <= 0.0 ? logistic(a) : logistic(-a);
1826 double local_h = a <= 0.0 ? h : -h;
1827 double em1 = expm1(local_h);
1828 double x = p * em1;
1829 return log1p_minus_x(x) + p * expm1_minus_x(local_h);
1830 }
1831 if (a <= 0.0) {
1832 double p = logistic(a);
1833 return p * (a - b) + softplus(b) - softplus(a);
1834 }
1835 double q = logistic(-a);
1836 return q * (b - a) + softplus(-b) - softplus(-a);
1837}
1838
1839__device__ __forceinline__ double bd0(double x, double m) {
1840 if (x == 0.0) return m;
1841 if (x == m) return 0.0;
1842 double hi = x > m ? x : m;
1843 double lo = x < m ? x : m;
1844 if (fabs(x - m) / hi < 0.2) {
1845 double v = ((x - m) / hi) / (1.0 + lo / hi);
1846 double sum = (x - m) * v;
1847 double term = 2.0 * x * v;
1848 double v2 = v * v;
1849 double denominator = 3.0;
1850 for (;;) {
1851 term *= v2;
1852 double next = sum + term / denominator;
1853 if (next == sum) return next;
1854 sum = next;
1855 denominator += 2.0;
1856 }
1857 }
1858 return x * (log(x) - log(m)) + (m - x);
1859}
1860
1861__device__ __forceinline__ double bernoulli_deviance(double y, double mu, double w) {
1862 return 2.0 * w * (bd0(y, mu) + bd0(1.0 - y, 1.0 - mu));
1863}
1864
1865__device__ __forceinline__ double logit_deviance(double y, double eta, double w) {
1866 double unit;
1867 if (y == 0.0) unit = softplus(eta);
1868 else if (y == 1.0) unit = softplus(-eta);
1869 else {
1870 double response_logit = log(y) - log1p(-y);
1871 unit = bernoulli_kl_from_logits(response_logit, eta);
1872 }
1873 return 2.0 * w * unit;
1874}
1875
1876__device__ __forceinline__ double std_norm_cdf(double x) {
1877 return 0.5 * erfc(-x * 0.7071067811865475);
1878}
1879
1880__device__ __forceinline__ double std_norm_pdf(double x) {
1881 return 0.3989422804014327 * exp(-0.5 * x * x);
1882}
1883
1884__device__ __forceinline__ double positive_mul_div(double a, double b, double c) {
1885 double product = a * b;
1886 if (isfinite(product) && product > 0.0) {
1887 double value = product / c;
1888 if (isfinite(value) && value > 0.0) return value;
1889 }
1890 double quotient_a = a / c;
1891 if (isfinite(quotient_a) && quotient_a > 0.0) {
1892 double value = quotient_a * b;
1893 if (isfinite(value) && value > 0.0) return value;
1894 }
1895 double quotient_b = b / c;
1896 if (isfinite(quotient_b) && quotient_b > 0.0) {
1897 double value = quotient_b * a;
1898 if (isfinite(value) && value > 0.0) return value;
1899 }
1900 return product / c;
1901}
1902
1903__device__ __forceinline__ double gamma_unit_deviance_near_one(double u) {
1904 if (fabs(u) > 0.125) return u - log1p(u);
1905 double power = u * u;
1906 double sum = 0.5 * power;
1907 for (int degree = 3; degree <= 32; ++degree) {
1908 power *= u;
1909 double term = power / (double)degree;
1910 double next = sum + ((degree & 1) ? -term : term);
1911 if (next == sum) break;
1912 sum = next;
1913 }
1914 return sum;
1915}
1916
1917__device__ __forceinline__ double poisson_unit_deviance_near_one(double u) {
1918 if (fabs(u) > 0.125) return (1.0 + u) * log1p(u) - u;
1919 double power = u * u;
1920 double sum = 0.5 * power;
1921 for (int degree = 3; degree <= 32; ++degree) {
1922 power *= u;
1923 double coefficient = ((degree & 1) ? -1.0 : 1.0)
1924 / ((double)degree * (double)(degree - 1));
1925 double next = sum + coefficient * power;
1926 if (next == sum) break;
1927 sum = next;
1928 }
1929 return sum;
1930}
1931
1932__device__ __forceinline__ bool pirls_outputs_finite(
1933 double mu, double grad_eta, double w_fisher, double w_hessian,
1934 double w_solver, double dev
1935) {
1936 return isfinite(mu) && isfinite(grad_eta) && isfinite(w_fisher)
1937 && isfinite(w_hessian) && isfinite(w_solver) && isfinite(dev);
1938}
1939"#
1940 .replace(
1941 "__PIRLS_LOG_ETA_MIN__",
1942 &format!("{:?}", crate::mixture_link::LOG_LINK_SOLVER_ETA_MIN),
1943 )
1944 .replace(
1945 "__PIRLS_LOG_ETA_MAX__",
1946 &format!("{:?}", crate::mixture_link::LOG_LINK_SOLVER_ETA_MAX),
1947 )
1948}
1949
1950#[cfg(target_os = "linux")]
1958fn cuda_source_for(family: PirlsRowFamily, curvature: CurvatureMode) -> String {
1959 let body = match family {
1960 PirlsRowFamily::GaussianIdentity => gaussian_identity_body(curvature),
1961 PirlsRowFamily::PoissonLog => poisson_log_body(curvature),
1962 PirlsRowFamily::GammaLog => gamma_log_body(curvature),
1963 PirlsRowFamily::BernoulliLogit => bernoulli_logit_body(curvature),
1964 PirlsRowFamily::BernoulliProbit => bernoulli_probit_body(curvature),
1965 PirlsRowFamily::BernoulliCLogLog => bernoulli_cloglog_body(curvature),
1966 };
1967 let kernel_name = family.kernel_name();
1968 let curvature_define = match curvature {
1973 CurvatureMode::Fisher => "#define PIRLS_CURVATURE_FISHER 1",
1974 CurvatureMode::Observed => "#define PIRLS_CURVATURE_OBSERVED 1",
1975 };
1976 let shape_param = if matches!(family, PirlsRowFamily::GammaLog) {
1979 " double shape,\n"
1980 } else {
1981 ""
1982 };
1983 format!(
1984 r#"
1985{curvature_define}
1986{prolog}
1987
1988extern "C" __global__ void {kernel_name}(
1989 int n,
1990 const double* __restrict__ eta,
1991 const double* __restrict__ y,
1992 const double* __restrict__ prior_w,
1993{shape_param} double* __restrict__ mu_out,
1994 double* __restrict__ grad_eta_out,
1995 double* __restrict__ w_hessian_out,
1996 double* __restrict__ w_solver_out,
1997 double* __restrict__ deviance_out,
1998 unsigned int* __restrict__ status_out
1999) {{
2000 int i = blockIdx.x * blockDim.x + threadIdx.x;
2001 if (i >= n) return;
2002 unsigned int status = PIRLS_OK;
2003 double eta_i = eta[i];
2004 double y_i = y[i];
2005 double wp = prior_w[i];
2006{body}
2007 if (status == PIRLS_OK) {{
2008 mu_out[i] = mu;
2009 grad_eta_out[i] = grad_eta;
2010 w_hessian_out[i] = w_hessian;
2011 w_solver_out[i] = w_solver;
2012 deviance_out[i] = dev;
2013 }}
2014 status_out[i] = status;
2015}}
2016"#,
2017 prolog = common_device_prolog(),
2018 )
2019}
2020
2021#[cfg(target_os = "linux")]
2026#[inline]
2027fn curvature_tag(curvature: CurvatureMode) -> &'static str {
2028 match curvature {
2029 CurvatureMode::Fisher => " // curvature: fisher\n",
2030 CurvatureMode::Observed => " // curvature: observed\n",
2031 }
2032}
2033
2034#[cfg(target_os = "linux")]
2035fn gaussian_identity_body(curvature: CurvatureMode) -> String {
2036 let tag = curvature_tag(curvature);
2037 format!(
2038 r#"{tag} double mu = 0.0, grad_eta = 0.0, w_fisher = 0.0;
2039 double w_hessian = 0.0, w_solver = 0.0, dev = 0.0;
2040 if (!isfinite(eta_i)) pirls_refuse(&status, PIRLS_ETA_DOMAIN);
2041 if (status == PIRLS_OK && !(isfinite(wp) && wp >= 0.0))
2042 pirls_refuse(&status, PIRLS_PRIOR_WEIGHT);
2043 if (status == PIRLS_OK && wp > 0.0 && !isfinite(y_i))
2044 pirls_refuse(&status, PIRLS_RESPONSE);
2045 if (status == PIRLS_OK) {{
2046 mu = eta_i;
2047 w_fisher = wp;
2048 w_hessian = wp;
2049 w_solver = w_hessian;
2050 if (wp > 0.0) {{
2051 double resid = y_i - mu;
2052 grad_eta = wp * resid;
2053 dev = wp * resid * resid;
2054 }}
2055 }}
2056 if (status == PIRLS_OK && !pirls_outputs_finite(
2057 mu, grad_eta, w_fisher, w_hessian, w_solver, dev))
2058 pirls_refuse(&status, PIRLS_FINAL_OUTPUT);
2059"#
2060 )
2061}
2062
2063#[cfg(target_os = "linux")]
2064fn poisson_log_body(curvature: CurvatureMode) -> String {
2065 let tag = curvature_tag(curvature);
2066 format!(
2067 r#"{tag} double mu = 0.0, grad_eta = 0.0, w_fisher = 0.0;
2068 double w_hessian = 0.0, w_solver = 0.0, dev = 0.0;
2069 if (!pirls_log_eta_valid(eta_i)) pirls_refuse(&status, PIRLS_ETA_DOMAIN);
2070 if (status == PIRLS_OK && !(isfinite(wp) && wp >= 0.0))
2071 pirls_refuse(&status, PIRLS_PRIOR_WEIGHT);
2072 if (status == PIRLS_OK && wp > 0.0 && !(isfinite(y_i) && y_i >= 0.0))
2073 pirls_refuse(&status, PIRLS_RESPONSE);
2074 if (status == PIRLS_OK) {{
2075 mu = exp(eta_i);
2076 if (!(isfinite(mu) && mu > 0.0)) pirls_refuse(&status, PIRLS_INVERSE_LINK);
2077 }}
2078 if (status == PIRLS_OK && wp > 0.0) {{
2079 w_fisher = wp * mu;
2080 if (!(isfinite(w_fisher) && w_fisher > 0.0))
2081 pirls_refuse(&status, PIRLS_FISHER_WEIGHT);
2082 if (status == PIRLS_OK) {{
2083 w_hessian = w_fisher;
2084 w_solver = w_hessian;
2085 grad_eta = wp * (y_i - mu);
2086 double u = (y_i - mu) / mu;
2087 double dev_base;
2088 if (y_i == 0.0) {{
2089 dev_base = w_fisher;
2090 }} else {{
2091 double scaled_unit = w_fisher * poisson_unit_deviance_near_one(u);
2092 if (isfinite(scaled_unit) && scaled_unit >= 0.0) {{
2093 dev_base = scaled_unit;
2094 }} else {{
2095 double weighted_y = positive_mul_div(w_fisher, y_i, mu);
2096 dev_base = weighted_y * (log(y_i) - eta_i - 1.0) + w_fisher;
2097 }}
2098 }}
2099 if (!isfinite(grad_eta)) pirls_refuse(&status, PIRLS_GRADIENT);
2100 dev = 2.0 * dev_base;
2101 if (!isfinite(dev)) pirls_refuse(&status, PIRLS_DEVIANCE);
2102 }}
2103 }}
2104 if (status == PIRLS_OK && !pirls_outputs_finite(
2105 mu, grad_eta, w_fisher, w_hessian, w_solver, dev))
2106 pirls_refuse(&status, PIRLS_FINAL_OUTPUT);
2107"#
2108 )
2109}
2110
2111#[cfg(target_os = "linux")]
2112fn gamma_log_body(curvature: CurvatureMode) -> String {
2113 let tag = curvature_tag(curvature);
2116 format!(
2117 r#"{tag} double mu = 0.0, grad_eta = 0.0, w_fisher = 0.0;
2118 double w_hessian = 0.0, w_solver = 0.0, dev = 0.0;
2119 if (!pirls_log_eta_valid(eta_i)) pirls_refuse(&status, PIRLS_ETA_DOMAIN);
2120 if (status == PIRLS_OK && !(isfinite(shape) && shape > 0.0))
2121 pirls_refuse(&status, PIRLS_GAMMA_SHAPE);
2122 if (status == PIRLS_OK && !(isfinite(wp) && wp >= 0.0))
2123 pirls_refuse(&status, PIRLS_PRIOR_WEIGHT);
2124 if (status == PIRLS_OK && wp > 0.0 && !(isfinite(y_i) && y_i > 0.0))
2125 pirls_refuse(&status, PIRLS_RESPONSE);
2126 if (status == PIRLS_OK) {{
2127 mu = exp(eta_i);
2128 if (!(isfinite(mu) && mu > 0.0)) pirls_refuse(&status, PIRLS_INVERSE_LINK);
2129 }}
2130 if (status == PIRLS_OK && wp > 0.0) {{
2131 w_fisher = wp * shape;
2132 if (!(isfinite(w_fisher) && w_fisher > 0.0))
2133 pirls_refuse(&status, PIRLS_FISHER_WEIGHT);
2134#ifdef PIRLS_CURVATURE_OBSERVED
2135 double weighted_ratio_observed = positive_mul_div(w_fisher, y_i, mu);
2136 if (!(isfinite(weighted_ratio_observed) && weighted_ratio_observed > 0.0))
2137 pirls_refuse(&status, PIRLS_OBSERVED_WEIGHT);
2138 w_hessian = weighted_ratio_observed;
2139#else
2140 w_hessian = w_fisher;
2141#endif
2142 if (!isfinite(w_hessian)) pirls_refuse(&status, PIRLS_OBSERVED_WEIGHT);
2143 w_solver = w_hessian;
2144 double u = (y_i - mu) / mu;
2145 double scaled_unit = w_fisher * gamma_unit_deviance_near_one(u);
2146 bool need_weighted_ratio = !isfinite(u)
2147 || !(isfinite(scaled_unit) && scaled_unit >= 0.0);
2148 double weighted_ratio = 0.0;
2149#ifdef PIRLS_CURVATURE_OBSERVED
2150 weighted_ratio = weighted_ratio_observed;
2151#else
2152 if (need_weighted_ratio)
2153 weighted_ratio = positive_mul_div(w_fisher, y_i, mu);
2154#endif
2155 grad_eta = isfinite(u) ? w_fisher * u : weighted_ratio - w_fisher;
2156 double dev_base;
2157 if (isfinite(scaled_unit) && scaled_unit >= 0.0) {{
2158 dev_base = scaled_unit;
2159 }} else {{
2160 dev_base = weighted_ratio - w_fisher * (1.0 + log(y_i) - eta_i);
2161 }}
2162 if (!isfinite(grad_eta)) pirls_refuse(&status, PIRLS_GRADIENT);
2163 dev = 2.0 * dev_base;
2164 if (!isfinite(dev)) pirls_refuse(&status, PIRLS_DEVIANCE);
2165 }}
2166 if (status == PIRLS_OK && !pirls_outputs_finite(
2167 mu, grad_eta, w_fisher, w_hessian, w_solver, dev))
2168 pirls_refuse(&status, PIRLS_FINAL_OUTPUT);
2169"#
2170 )
2171}
2172
2173#[cfg(target_os = "linux")]
2174fn bernoulli_logit_body(curvature: CurvatureMode) -> String {
2175 let tag = curvature_tag(curvature);
2176 format!(
2177 r#"{tag} double mu = 0.0, grad_eta = 0.0, w_fisher = 0.0;
2178 double w_hessian = 0.0, w_solver = 0.0, dev = 0.0;
2179 if (!isfinite(eta_i)) pirls_refuse(&status, PIRLS_ETA_DOMAIN);
2180 if (status == PIRLS_OK && !(isfinite(wp) && wp >= 0.0))
2181 pirls_refuse(&status, PIRLS_PRIOR_WEIGHT);
2182 if (status == PIRLS_OK && wp > 0.0
2183 && !(isfinite(y_i) && y_i >= 0.0 && y_i <= 1.0))
2184 pirls_refuse(&status, PIRLS_RESPONSE);
2185 double tail = exp(-fabs(eta_i));
2186 double denom = 1.0 + tail;
2187 double dmu_deta = tail / (denom * denom);
2188 if (status == PIRLS_OK) {{
2189 mu = eta_i >= 0.0 ? 1.0 / denom : tail / denom;
2190 if (!(isfinite(mu) && mu >= 0.0 && mu <= 1.0
2191 && isfinite(dmu_deta) && dmu_deta > 0.0))
2192 pirls_refuse(&status, PIRLS_INVERSE_LINK);
2193 }}
2194 if (status == PIRLS_OK && wp > 0.0) {{
2195 double residual;
2196 if (eta_i >= 0.0) {{
2197 double one_minus_mu = tail / denom;
2198 residual = y_i == 1.0 ? one_minus_mu : (y_i - 1.0) + one_minus_mu;
2199 }} else {{
2200 residual = y_i - mu;
2201 }}
2202 w_fisher = wp * dmu_deta;
2203 if (!(isfinite(w_fisher) && w_fisher > 0.0))
2204 pirls_refuse(&status, PIRLS_FISHER_WEIGHT);
2205 w_hessian = w_fisher;
2206 w_solver = w_hessian;
2207 grad_eta = wp * residual;
2208 if (!isfinite(grad_eta)) pirls_refuse(&status, PIRLS_GRADIENT);
2209 dev = logit_deviance(y_i, eta_i, wp);
2210 if (!isfinite(dev)) pirls_refuse(&status, PIRLS_DEVIANCE);
2211 }}
2212 if (status == PIRLS_OK && !pirls_outputs_finite(
2213 mu, grad_eta, w_fisher, w_hessian, w_solver, dev))
2214 pirls_refuse(&status, PIRLS_FINAL_OUTPUT);
2215"#
2216 )
2217}
2218
2219#[cfg(target_os = "linux")]
2220fn bernoulli_probit_body(curvature: CurvatureMode) -> String {
2221 let tag = curvature_tag(curvature);
2222 format!(
2223 r#"{tag} double mu = 0.0, grad_eta = 0.0, w_fisher = 0.0;
2224 double w_hessian = 0.0, w_solver = 0.0, dev = 0.0;
2225 if (!isfinite(eta_i)) pirls_refuse(&status, PIRLS_ETA_DOMAIN);
2226 if (status == PIRLS_OK && !(isfinite(wp) && wp >= 0.0))
2227 pirls_refuse(&status, PIRLS_PRIOR_WEIGHT);
2228 if (status == PIRLS_OK && wp > 0.0
2229 && !(isfinite(y_i) && y_i >= 0.0 && y_i <= 1.0))
2230 pirls_refuse(&status, PIRLS_RESPONSE);
2231 double dmu_deta = 0.0, d2mu_deta2 = 0.0, v = 0.0;
2232 if (status == PIRLS_OK) {{
2233 mu = std_norm_cdf(eta_i);
2234 dmu_deta = std_norm_pdf(eta_i);
2235 d2mu_deta2 = -eta_i * dmu_deta;
2236 if (!(isfinite(mu) && mu > 0.0 && mu < 1.0
2237 && isfinite(dmu_deta) && dmu_deta > 0.0
2238 && isfinite(d2mu_deta2)))
2239 pirls_refuse(&status, PIRLS_INVERSE_LINK);
2240 }}
2241 if (status == PIRLS_OK && wp > 0.0) {{
2242 v = mu * (1.0 - mu);
2243 double fisher_per_prior = dmu_deta * dmu_deta / v;
2244 w_fisher = wp * fisher_per_prior;
2245 if (!(isfinite(v) && v > 0.0 && isfinite(fisher_per_prior)
2246 && fisher_per_prior > 0.0 && isfinite(w_fisher) && w_fisher > 0.0))
2247 pirls_refuse(&status, PIRLS_FISHER_WEIGHT);
2248 double resid = y_i - mu;
2249#ifdef PIRLS_CURVATURE_OBSERVED
2250 double bracket = d2mu_deta2 / v
2251 - (dmu_deta * dmu_deta) * (1.0 - 2.0 * mu) / (v * v);
2252 w_hessian = w_fisher - wp * resid * bracket;
2253#else
2254 w_hessian = w_fisher;
2255#endif
2256 if (!isfinite(w_hessian)) pirls_refuse(&status, PIRLS_OBSERVED_WEIGHT);
2257 w_solver = w_hessian;
2258 grad_eta = wp * resid * dmu_deta / v;
2259 if (!isfinite(grad_eta)) pirls_refuse(&status, PIRLS_GRADIENT);
2260 dev = bernoulli_deviance(y_i, mu, wp);
2261 if (!isfinite(dev)) pirls_refuse(&status, PIRLS_DEVIANCE);
2262 }}
2263 if (status == PIRLS_OK && !pirls_outputs_finite(
2264 mu, grad_eta, w_fisher, w_hessian, w_solver, dev))
2265 pirls_refuse(&status, PIRLS_FINAL_OUTPUT);
2266"#
2267 )
2268}
2269
2270#[cfg(target_os = "linux")]
2271fn bernoulli_cloglog_body(curvature: CurvatureMode) -> String {
2272 let tag = curvature_tag(curvature);
2273 format!(
2274 r#"{tag} double mu = 0.0, grad_eta = 0.0, w_fisher = 0.0;
2275 double w_hessian = 0.0, w_solver = 0.0, dev = 0.0;
2276 if (!isfinite(eta_i)) pirls_refuse(&status, PIRLS_ETA_DOMAIN);
2277 if (status == PIRLS_OK && !(isfinite(wp) && wp >= 0.0))
2278 pirls_refuse(&status, PIRLS_PRIOR_WEIGHT);
2279 if (status == PIRLS_OK && wp > 0.0
2280 && !(isfinite(y_i) && y_i >= 0.0 && y_i <= 1.0))
2281 pirls_refuse(&status, PIRLS_RESPONSE);
2282 double inner = 0.0, dmu_deta = 0.0, d2mu_deta2 = 0.0, v = 0.0;
2283 if (status == PIRLS_OK) {{
2284 inner = exp(eta_i);
2285 double complement = exp(-inner);
2286 mu = -expm1(-inner);
2287 dmu_deta = inner * complement;
2288 d2mu_deta2 = dmu_deta * (1.0 - inner);
2289 if (!(isfinite(mu) && mu > 0.0 && mu < 1.0
2290 && isfinite(dmu_deta) && dmu_deta > 0.0
2291 && isfinite(d2mu_deta2)))
2292 pirls_refuse(&status, PIRLS_INVERSE_LINK);
2293 }}
2294 if (status == PIRLS_OK && wp > 0.0) {{
2295 v = mu * (1.0 - mu);
2296 double fisher_per_prior = dmu_deta * dmu_deta / v;
2297 w_fisher = wp * fisher_per_prior;
2298 if (!(isfinite(v) && v > 0.0 && isfinite(fisher_per_prior)
2299 && fisher_per_prior > 0.0 && isfinite(w_fisher) && w_fisher > 0.0))
2300 pirls_refuse(&status, PIRLS_FISHER_WEIGHT);
2301 double resid = y_i - mu;
2302#ifdef PIRLS_CURVATURE_OBSERVED
2303 double bracket = d2mu_deta2 / v
2304 - (dmu_deta * dmu_deta) * (1.0 - 2.0 * mu) / (v * v);
2305 w_hessian = w_fisher - wp * resid * bracket;
2306#else
2307 w_hessian = w_fisher;
2308#endif
2309 if (!isfinite(w_hessian)) pirls_refuse(&status, PIRLS_OBSERVED_WEIGHT);
2310 w_solver = w_hessian;
2311 grad_eta = wp * resid * dmu_deta / v;
2312 if (!isfinite(grad_eta)) pirls_refuse(&status, PIRLS_GRADIENT);
2313 dev = bernoulli_deviance(y_i, mu, wp);
2314 if (!isfinite(dev)) pirls_refuse(&status, PIRLS_DEVIANCE);
2315 }}
2316 if (status == PIRLS_OK && !pirls_outputs_finite(
2317 mu, grad_eta, w_fisher, w_hessian, w_solver, dev))
2318 pirls_refuse(&status, PIRLS_FINAL_OUTPUT);
2319"#
2320 )
2321}
2322
2323#[cfg(target_os = "linux")]
2337fn solve_row_source_for(family: PirlsRowFamily, curvature: CurvatureMode) -> String {
2338 let body = match family {
2339 PirlsRowFamily::GaussianIdentity => gaussian_identity_body(curvature),
2340 PirlsRowFamily::PoissonLog => poisson_log_body(curvature),
2341 PirlsRowFamily::GammaLog => gamma_log_body(curvature),
2342 PirlsRowFamily::BernoulliLogit => bernoulli_logit_body(curvature),
2343 PirlsRowFamily::BernoulliProbit => bernoulli_probit_body(curvature),
2344 PirlsRowFamily::BernoulliCLogLog => bernoulli_cloglog_body(curvature),
2345 };
2346 let kernel_name = family.solve_kernel_name();
2347 let curvature_define = match curvature {
2348 CurvatureMode::Fisher => "#define PIRLS_CURVATURE_FISHER 1",
2349 CurvatureMode::Observed => "#define PIRLS_CURVATURE_OBSERVED 1",
2350 };
2351 let shape_param = if matches!(family, PirlsRowFamily::GammaLog) {
2353 " double shape,\n"
2354 } else {
2355 ""
2356 };
2357 format!(
2358 r#"
2359{curvature_define}
2360{prolog}
2361
2362extern "C" __global__ void {kernel_name}(
2363 int n,
2364 const double* __restrict__ eta,
2365 const double* __restrict__ y,
2366 const double* __restrict__ prior_w,
2367{shape_param} double* __restrict__ grad_eta_out,
2368 double* __restrict__ w_solver_out,
2369 double* __restrict__ deviance_out,
2370 unsigned int* __restrict__ status_out
2371) {{
2372 int i = blockIdx.x * blockDim.x + threadIdx.x;
2373 if (i >= n) return;
2374 unsigned int status = PIRLS_OK;
2375 double eta_i = eta[i];
2376 double y_i = y[i];
2377 double wp = prior_w[i];
2378{body}
2379 if (status == PIRLS_OK) {{
2380 grad_eta_out[i] = grad_eta;
2381 w_solver_out[i] = w_solver;
2382 deviance_out[i] = dev;
2383 }}
2384 status_out[i] = status;
2385}}
2386"#,
2387 prolog = common_device_prolog(),
2388 )
2389}
2390
2391#[cfg(target_os = "linux")]
2398const ALPHA_LADDER_CUDA_ARRAY: &str =
2399 "__constant__ double PIRLS_ALPHAS[7] = {1.0, 0.5, 0.25, 0.125, 0.0625, 0.03125, 0.015625};";
2400
2401#[cfg(target_os = "linux")]
2416fn ladder_source_for(family: PirlsRowFamily, curvature: CurvatureMode) -> String {
2417 let body = match family {
2418 PirlsRowFamily::GaussianIdentity => gaussian_identity_body(curvature),
2419 PirlsRowFamily::PoissonLog => poisson_log_body(curvature),
2420 PirlsRowFamily::GammaLog => gamma_log_body(curvature),
2421 PirlsRowFamily::BernoulliLogit => bernoulli_logit_body(curvature),
2422 PirlsRowFamily::BernoulliProbit => bernoulli_probit_body(curvature),
2423 PirlsRowFamily::BernoulliCLogLog => bernoulli_cloglog_body(curvature),
2424 };
2425 let kernel_name = family.ladder_kernel_name();
2426 let curvature_define = match curvature {
2427 CurvatureMode::Fisher => "#define PIRLS_CURVATURE_FISHER 1",
2428 CurvatureMode::Observed => "#define PIRLS_CURVATURE_OBSERVED 1",
2429 };
2430 let shape_param = if matches!(family, PirlsRowFamily::GammaLog) {
2437 " double shape,\n"
2438 } else {
2439 ""
2440 };
2441 format!(
2442 r#"
2443{curvature_define}
2444{prolog}
2445{alphas}
2446
2447extern "C" __global__ void {kernel_name}(
2448 int n,
2449 const double* __restrict__ eta,
2450 const double* __restrict__ xd,
2451 const double* __restrict__ y,
2452 const double* __restrict__ prior_w,
2453{shape_param} double* __restrict__ objective_out,
2454 unsigned int* __restrict__ status_out
2455) {{
2456 int i = blockIdx.x * blockDim.x + threadIdx.x;
2457 int k = (int)blockIdx.y;
2458 if (i >= n) return;
2459 unsigned int status = PIRLS_OK;
2460 double alpha = PIRLS_ALPHAS[k];
2461 double eta_i = eta[i] + alpha * xd[i];
2462 double y_i = y[i];
2463 double wp = prior_w[i];
2464{body}
2465 if (status == PIRLS_OK) atomicAdd(&objective_out[k], dev);
2466 status_out[k * n + i] = status;
2467}}
2468"#,
2469 prolog = common_device_prolog(),
2470 alphas = ALPHA_LADDER_CUDA_ARRAY,
2471 )
2472}
2473
2474#[cfg(test)]
2479#[path = "pirls_row_tests.rs"]
2480mod pirls_row_tests;