use std::sync::OnceLock;
use gam_gpu::gpu_error::GpuError;
#[cfg(target_os = "linux")]
use gam_gpu::gpu_error::GpuResultExt;
use gam_problem::EstimationError;
#[cfg(target_os = "linux")]
use std::sync::{Arc, Mutex};
#[cfg(target_os = "linux")]
use cudarc::driver::{CudaContext, CudaModule};
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub enum PirlsRowFamily {
BernoulliLogit,
BernoulliProbit,
BernoulliCLogLog,
PoissonLog,
GaussianIdentity,
GammaLog,
}
impl PirlsRowFamily {
pub const ALL: [Self; 6] = [
Self::BernoulliLogit,
Self::BernoulliProbit,
Self::BernoulliCLogLog,
Self::PoissonLog,
Self::GaussianIdentity,
Self::GammaLog,
];
pub const fn as_str(self) -> &'static str {
match self {
Self::BernoulliLogit => "bernoulli-logit",
Self::BernoulliProbit => "bernoulli-probit",
Self::BernoulliCLogLog => "bernoulli-cloglog",
Self::PoissonLog => "poisson-log",
Self::GaussianIdentity => "gaussian-identity",
Self::GammaLog => "gamma-log",
}
}
pub const fn kernel_name(self) -> &'static str {
match self {
Self::BernoulliLogit => "pirls_row_bernoulli_logit",
Self::BernoulliProbit => "pirls_row_bernoulli_probit",
Self::BernoulliCLogLog => "pirls_row_bernoulli_cloglog",
Self::PoissonLog => "pirls_row_poisson_log",
Self::GaussianIdentity => "pirls_row_gaussian_identity",
Self::GammaLog => "pirls_row_gamma_log",
}
}
pub const fn solve_kernel_name(self) -> &'static str {
match self {
Self::BernoulliLogit => "pirls_solve_bernoulli_logit",
Self::BernoulliProbit => "pirls_solve_bernoulli_probit",
Self::BernoulliCLogLog => "pirls_solve_bernoulli_cloglog",
Self::PoissonLog => "pirls_solve_poisson_log",
Self::GaussianIdentity => "pirls_solve_gaussian_identity",
Self::GammaLog => "pirls_solve_gamma_log",
}
}
pub const fn ladder_kernel_name(self) -> &'static str {
match self {
Self::BernoulliLogit => "pirls_ladder_bernoulli_logit",
Self::BernoulliProbit => "pirls_ladder_bernoulli_probit",
Self::BernoulliCLogLog => "pirls_ladder_bernoulli_cloglog",
Self::PoissonLog => "pirls_ladder_poisson_log",
Self::GaussianIdentity => "pirls_ladder_gaussian_identity",
Self::GammaLog => "pirls_ladder_gamma_log",
}
}
}
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub enum CurvatureMode {
Fisher,
Observed,
}
impl CurvatureMode {
pub const fn as_str(self) -> &'static str {
match self {
Self::Fisher => "fisher",
Self::Observed => "observed",
}
}
}
pub mod status_codes {
pub const OK: u32 = 0;
pub const ETA_DOMAIN: u32 = 1;
pub const PRIOR_WEIGHT: u32 = 2;
pub const RESPONSE: u32 = 3;
pub const GAMMA_SHAPE: u32 = 4;
pub const INVERSE_LINK: u32 = 5;
pub const FISHER_WEIGHT: u32 = 6;
pub const OBSERVED_WEIGHT: u32 = 7;
pub const GRADIENT: u32 = 8;
pub const DEVIANCE: u32 = 9;
pub const FINAL_OUTPUT: u32 = 10;
pub const fn quantity(code: u32) -> &'static str {
match code {
ETA_DOMAIN => "inverse-link eta domain",
PRIOR_WEIGHT => "prior weight",
RESPONSE => "response",
GAMMA_SHAPE => "Gamma shape",
INVERSE_LINK => "inverse-link jet",
FISHER_WEIGHT => "Fisher weight",
OBSERVED_WEIGHT => "observed Hessian weight",
GRADIENT => "eta gradient",
DEVIANCE => "deviance contribution",
FINAL_OUTPUT => "final row output",
_ => "unknown GPU PIRLS refusal",
}
}
}
#[derive(Clone, Copy, Debug)]
pub struct RowInput {
pub eta: f64,
pub y: f64,
pub prior_weight: f64,
}
#[derive(Clone, Copy, Debug, Default)]
pub struct RowOutput {
pub mu: f64,
pub grad_eta: f64,
pub w_fisher: f64,
pub w_hessian: f64,
pub w_solver: f64,
pub deviance: f64,
}
pub fn row_reweight_cpu(
family: PirlsRowFamily,
mode: CurvatureMode,
input: RowInput,
gamma_shape: f64,
) -> Result<RowOutput, EstimationError> {
row_reweight_cpu_at(0, family, mode, input, gamma_shape)
}
pub fn row_reweight_cpu_at(
row: usize,
family: PirlsRowFamily,
mode: CurvatureMode,
input: RowInput,
gamma_shape: f64,
) -> Result<RowOutput, EstimationError> {
match family {
PirlsRowFamily::GaussianIdentity => row_gaussian_identity(row, input, mode),
PirlsRowFamily::PoissonLog => row_poisson_log(row, input, mode),
PirlsRowFamily::GammaLog => row_gamma_log(row, input, mode, gamma_shape),
PirlsRowFamily::BernoulliLogit => row_bernoulli_logit(row, input, mode),
PirlsRowFamily::BernoulliProbit => row_bernoulli_probit(row, input, mode),
PirlsRowFamily::BernoulliCLogLog => row_bernoulli_cloglog(row, input, mode),
}
}
pub fn replay_first_refusal(
family: PirlsRowFamily,
mode: CurvatureMode,
gamma_shape: f64,
eta: &[f64],
y: &[f64],
prior_weight: &[f64],
status: &[u32],
) -> Result<(), EstimationError> {
let n = eta.len();
if y.len() != n || prior_weight.len() != n || status.len() != n {
return Err(EstimationError::InvalidInput(format!(
"GPU PIRLS refusal replay length mismatch: eta={n}, y={}, prior_weight={}, status={}",
y.len(),
prior_weight.len(),
status.len(),
)));
}
let Some((row, &code)) = status
.iter()
.enumerate()
.find(|(_, code)| **code != status_codes::OK)
else {
return Ok(());
};
let input = RowInput {
eta: eta[row],
y: y[row],
prior_weight: prior_weight[row],
};
match row_reweight_cpu_at(row, family, mode, input, gamma_shape) {
Err(error) => Err(error),
Ok(_) => Err(row_error(
row,
status_codes::quantity(code),
input.eta,
f64::from(code),
)),
}
}
#[inline]
fn select_w_hessian(mode: CurvatureMode, w_fisher: f64, observed_correction: f64) -> f64 {
match mode {
CurvatureMode::Fisher => w_fisher,
CurvatureMode::Observed => w_fisher + observed_correction,
}
}
#[inline]
fn row_error(row: usize, quantity: &'static str, eta: f64, value: f64) -> EstimationError {
EstimationError::PirlsRowGeometryUnrepresentable {
row,
quantity,
eta,
value,
}
}
#[inline]
fn finite_eta(link: &'static str, eta: f64) -> Result<(), EstimationError> {
if eta.is_finite() {
Ok(())
} else {
Err(EstimationError::InverseLinkDomainViolation {
link,
eta,
lower: -f64::MAX,
upper: f64::MAX,
})
}
}
#[inline]
fn prior_weight(row: usize, input: RowInput) -> Result<f64, EstimationError> {
if input.prior_weight.is_finite() && input.prior_weight >= 0.0 {
Ok(input.prior_weight)
} else {
Err(row_error(
row,
"prior weight",
input.eta,
input.prior_weight,
))
}
}
#[inline]
fn certify_output(row: usize, eta: f64, output: RowOutput) -> Result<RowOutput, EstimationError> {
for (quantity, value) in [
("mean", output.mu),
("eta gradient", output.grad_eta),
("Fisher weight", output.w_fisher),
("observed Hessian weight", output.w_hessian),
("solver Hessian weight", output.w_solver),
("deviance contribution", output.deviance),
] {
if !value.is_finite() {
return Err(row_error(row, quantity, eta, value));
}
}
Ok(output)
}
#[inline]
fn positive_mul_div(a: f64, b: f64, c: f64) -> f64 {
let product = a * b;
if product.is_finite() && product > 0.0 {
let value = product / c;
if value.is_finite() && value > 0.0 {
return value;
}
}
let quotient_a = a / c;
if quotient_a.is_finite() && quotient_a > 0.0 {
let value = quotient_a * b;
if value.is_finite() && value > 0.0 {
return value;
}
}
let quotient_b = b / c;
if quotient_b.is_finite() && quotient_b > 0.0 {
let value = quotient_b * a;
if value.is_finite() && value > 0.0 {
return value;
}
}
product / c
}
#[inline]
fn gamma_unit_deviance_near_one(u: f64) -> f64 {
if u.abs() > 0.125 {
return u - u.ln_1p();
}
let mut power = u * u;
let mut sum = 0.5 * power;
for degree in 3..=32 {
power *= u;
let term = power / f64::from(degree);
let next = if degree % 2 == 0 {
sum + term
} else {
sum - term
};
if next == sum {
break;
}
sum = next;
}
sum
}
#[inline]
fn poisson_unit_deviance_near_one(u: f64) -> f64 {
if u.abs() > 0.125 {
return (1.0 + u) * u.ln_1p() - u;
}
let mut power = u * u;
let mut sum = 0.5 * power;
for degree in 3..=32 {
power *= u;
let coefficient =
if degree % 2 == 0 { 1.0 } else { -1.0 } / (f64::from(degree) * f64::from(degree - 1));
let next = sum + coefficient * power;
if next == sum {
break;
}
sum = next;
}
sum
}
#[inline]
fn row_gaussian_identity(
row: usize,
input: RowInput,
mode: CurvatureMode,
) -> Result<RowOutput, EstimationError> {
finite_eta("standard identity inverse link", input.eta)?;
let w = prior_weight(row, input)?;
let mu = input.eta;
if w > 0.0 && !input.y.is_finite() {
return Err(row_error(row, "Gaussian response", input.eta, input.y));
}
let resid = input.y - mu;
let (grad_eta, dev) = if w == 0.0 {
(0.0, 0.0)
} else {
(w * resid, w * resid * resid)
};
let w_hessian = select_w_hessian(mode, w, 0.0);
certify_output(
row,
input.eta,
RowOutput {
mu,
grad_eta,
w_fisher: w,
w_hessian,
w_solver: w_hessian,
deviance: dev,
},
)
}
#[inline]
fn row_poisson_log(
row: usize,
input: RowInput,
mode: CurvatureMode,
) -> Result<RowOutput, EstimationError> {
let mu = crate::mixture_link::log_link_solver_exp(input.eta)?;
let w_prior = prior_weight(row, input)?;
if w_prior > 0.0 && !(input.y.is_finite() && input.y >= 0.0) {
return Err(row_error(row, "Poisson response", input.eta, input.y));
}
if w_prior == 0.0 {
return certify_output(
row,
input.eta,
RowOutput {
mu,
..RowOutput::default()
},
);
}
let w_fisher = w_prior * mu;
if !(w_fisher.is_finite() && w_fisher > 0.0) {
return Err(row_error(row, "Poisson Fisher weight", input.eta, w_fisher));
}
let grad_eta = w_prior * (input.y - mu);
let u = (input.y - mu) / mu;
let dev_base = if input.y == 0.0 {
w_fisher
} else {
let scaled_unit = w_fisher * poisson_unit_deviance_near_one(u);
if scaled_unit.is_finite() && scaled_unit >= 0.0 {
scaled_unit
} else {
let weighted_y = positive_mul_div(w_fisher, input.y, mu);
weighted_y * (input.y.ln() - input.eta - 1.0) + w_fisher
}
};
let dev = 2.0 * dev_base;
let w_hessian = select_w_hessian(mode, w_fisher, 0.0);
certify_output(
row,
input.eta,
RowOutput {
mu,
grad_eta,
w_fisher,
w_hessian,
w_solver: w_hessian,
deviance: dev,
},
)
}
#[inline]
fn row_gamma_log(
row: usize,
input: RowInput,
mode: CurvatureMode,
shape: f64,
) -> Result<RowOutput, EstimationError> {
let mu = crate::mixture_link::log_link_solver_exp(input.eta)?;
if !(shape.is_finite() && shape > 0.0) {
return Err(row_error(row, "Gamma shape", input.eta, shape));
}
let w_prior = prior_weight(row, input)?;
if w_prior > 0.0 && !(input.y.is_finite() && input.y > 0.0) {
return Err(row_error(row, "Gamma response", input.eta, input.y));
}
if w_prior == 0.0 {
return certify_output(
row,
input.eta,
RowOutput {
mu,
..RowOutput::default()
},
);
}
let w_fisher = w_prior * shape;
if !(w_fisher.is_finite() && w_fisher > 0.0) {
return Err(row_error(row, "Gamma Fisher weight", input.eta, w_fisher));
}
let observed_ratio = match mode {
CurvatureMode::Fisher => None,
CurvatureMode::Observed => {
let direct = w_fisher * (input.y / mu);
let weighted_ratio = if direct.is_finite() && direct > 0.0 {
direct
} else {
positive_mul_div(w_fisher, input.y, mu)
};
if !(weighted_ratio.is_finite() && weighted_ratio > 0.0) {
return Err(row_error(
row,
"Gamma observed Hessian weight",
input.eta,
weighted_ratio,
));
}
Some(weighted_ratio)
}
};
let w_hessian = observed_ratio.unwrap_or(w_fisher);
if !w_hessian.is_finite() {
return Err(row_error(
row,
"Gamma observed Hessian weight",
input.eta,
w_hessian,
));
}
let u = (input.y - mu) / mu;
let scaled_unit = w_fisher * gamma_unit_deviance_near_one(u);
let need_weighted_ratio = !u.is_finite() || !(scaled_unit.is_finite() && scaled_unit >= 0.0);
let weighted_ratio = if need_weighted_ratio {
observed_ratio.unwrap_or_else(|| positive_mul_div(w_fisher, input.y, mu))
} else {
0.0
};
let grad_eta = if u.is_finite() {
w_fisher * u
} else {
weighted_ratio - w_fisher
};
let dev_base = if scaled_unit.is_finite() && scaled_unit >= 0.0 {
scaled_unit
} else {
weighted_ratio - w_fisher * (1.0 + input.y.ln() - input.eta)
};
let dev = 2.0 * dev_base;
certify_output(
row,
input.eta,
RowOutput {
mu,
grad_eta,
w_fisher,
w_hessian,
w_solver: w_hessian,
deviance: dev,
},
)
}
#[inline]
fn bernoulli_response(row: usize, input: RowInput, w: f64) -> Result<(), EstimationError> {
if w == 0.0 || (input.y.is_finite() && (0.0..=1.0).contains(&input.y)) {
Ok(())
} else {
Err(row_error(row, "binomial response", input.eta, input.y))
}
}
#[inline]
fn row_bernoulli_logit(
row: usize,
input: RowInput,
mode: CurvatureMode,
) -> Result<RowOutput, EstimationError> {
finite_eta("standard logit inverse link", input.eta)?;
let w_prior = prior_weight(row, input)?;
bernoulli_response(row, input, w_prior)?;
let tail = (-input.eta.abs()).exp();
let denom = 1.0 + tail;
let (mu, residual) = if input.eta >= 0.0 {
let one_minus_mu = tail / denom;
let residual = if input.y == 1.0 {
one_minus_mu
} else {
(input.y - 1.0) + one_minus_mu
};
(1.0 / denom, residual)
} else {
let mu = tail / denom;
(mu, input.y - mu)
};
let dmu_deta = tail / (denom * denom);
if !(dmu_deta.is_finite() && dmu_deta > 0.0) {
return Err(row_error(
row,
"canonical-logit inverse-link jet",
input.eta,
dmu_deta,
));
}
if w_prior == 0.0 {
return certify_output(
row,
input.eta,
RowOutput {
mu,
..RowOutput::default()
},
);
}
let w_fisher = w_prior * dmu_deta;
if !(w_fisher.is_finite() && w_fisher > 0.0) {
return Err(row_error(row, "logit Fisher weight", input.eta, w_fisher));
}
let grad_eta = w_prior * residual;
let dev = bernoulli_logit_deviance(input.y, input.eta, w_prior);
let w_hessian = select_w_hessian(mode, w_fisher, 0.0);
certify_output(
row,
input.eta,
RowOutput {
mu,
grad_eta,
w_fisher,
w_hessian,
w_solver: w_hessian,
deviance: dev,
},
)
}
#[inline]
fn row_bernoulli_probit(
row: usize,
input: RowInput,
mode: CurvatureMode,
) -> Result<RowOutput, EstimationError> {
finite_eta("standard probit inverse link", input.eta)?;
let d1 = standard_normal_pdf(input.eta);
row_bernoulli_noncanonical(
row,
input,
mode,
standard_normal_cdf(input.eta),
d1,
-input.eta * d1,
)
}
#[inline]
fn row_bernoulli_cloglog(
row: usize,
input: RowInput,
mode: CurvatureMode,
) -> Result<RowOutput, EstimationError> {
finite_eta("standard complementary-log-log inverse link", input.eta)?;
let inner = input.eta.exp();
let mu = -(-inner).exp_m1();
let complement = (-inner).exp();
let d1 = inner * complement;
row_bernoulli_noncanonical(row, input, mode, mu, d1, d1 * (1.0 - inner))
}
#[inline]
fn row_bernoulli_noncanonical(
row: usize,
input: RowInput,
mode: CurvatureMode,
mu: f64,
d1: f64,
d2: f64,
) -> Result<RowOutput, EstimationError> {
let w_prior = prior_weight(row, input)?;
bernoulli_response(row, input, w_prior)?;
if !(mu.is_finite() && mu > 0.0 && mu < 1.0 && d1.is_finite() && d1 > 0.0 && d2.is_finite()) {
return Err(row_error(row, "inverse-link jet", input.eta, mu));
}
if w_prior == 0.0 {
return certify_output(
row,
input.eta,
RowOutput {
mu,
..RowOutput::default()
},
);
}
let v = mu * (1.0 - mu);
let fisher_per_prior = d1 * d1 / v;
let w_fisher = w_prior * fisher_per_prior;
if !(v.is_finite()
&& v > 0.0
&& fisher_per_prior.is_finite()
&& fisher_per_prior > 0.0
&& w_fisher.is_finite()
&& w_fisher > 0.0)
{
return Err(row_error(
row,
"Bernoulli Fisher weight",
input.eta,
w_fisher,
));
}
let resid = input.y - mu;
let grad_eta = w_prior * resid * d1 / v;
let bracket = d2 / v - d1 * d1 * (1.0 - 2.0 * mu) / (v * v);
let observed_correction = -w_prior * resid * bracket;
let w_hessian = select_w_hessian(mode, w_fisher, observed_correction);
if !w_hessian.is_finite() {
return Err(row_error(
row,
"Bernoulli observed Hessian weight",
input.eta,
w_hessian,
));
}
let dev = bernoulli_deviance(input.y, mu, w_prior);
certify_output(
row,
input.eta,
RowOutput {
mu,
grad_eta,
w_fisher,
w_hessian,
w_solver: w_hessian,
deviance: dev,
},
)
}
#[inline]
fn softplus(x: f64) -> f64 {
x.max(0.0) + (-x.abs()).exp().ln_1p()
}
#[inline]
fn expm1_minus_x(x: f64) -> f64 {
if x.abs() > 0.5 {
return x.exp_m1() - x;
}
let mut term = 0.5 * x * x;
let mut sum = term;
let mut degree = 2.0;
loop {
degree += 1.0;
term *= x / degree;
let next = sum + term;
if next == sum {
return next;
}
sum = next;
}
}
#[inline]
fn log1p_minus_x(x: f64) -> f64 {
if x.abs() > 0.5 {
return x.ln_1p() - x;
}
let mut power = x * x;
let mut sign = -1.0;
let mut degree = 2.0;
let mut sum = sign * power / degree;
loop {
power *= x;
sign = -sign;
degree += 1.0;
let next = sum + sign * power / degree;
if next == sum {
return next;
}
sum = next;
}
}
#[inline]
fn logistic(x: f64) -> f64 {
if x >= 0.0 {
1.0 / (1.0 + (-x).exp())
} else {
let e = x.exp();
e / (1.0 + e)
}
}
#[inline]
fn bernoulli_kl_from_logits(a: f64, b: f64) -> f64 {
if a == b {
return 0.0;
}
let h = b - a;
if h.abs() <= 0.5 {
let (p, local_h) = if a <= 0.0 {
(logistic(a), h)
} else {
(logistic(-a), -h)
};
let em1 = local_h.exp_m1();
let x = p * em1;
return log1p_minus_x(x) + p * expm1_minus_x(local_h);
}
if a <= 0.0 {
let p = logistic(a);
p * (a - b) + softplus(b) - softplus(a)
} else {
let q = logistic(-a);
q * (b - a) + softplus(-b) - softplus(-a)
}
}
#[inline]
fn bd0(x: f64, m: f64) -> f64 {
if x == 0.0 {
return m;
}
if x == m {
return 0.0;
}
let hi = x.max(m);
let lo = x.min(m);
if (x - m).abs() / hi < 0.2 {
let v = ((x - m) / hi) / (1.0 + lo / hi);
let mut sum = (x - m) * v;
let mut term = 2.0 * x * v;
let v2 = v * v;
let mut denominator = 3.0;
loop {
term *= v2;
let next = sum + term / denominator;
if next == sum {
return next;
}
sum = next;
denominator += 2.0;
}
}
x * (x.ln() - m.ln()) + (m - x)
}
#[inline]
fn bernoulli_logit_deviance(y: f64, eta: f64, w: f64) -> f64 {
let unit = if y == 0.0 {
softplus(eta)
} else if y == 1.0 {
softplus(-eta)
} else {
let response_logit = y.ln() - (-y).ln_1p();
bernoulli_kl_from_logits(response_logit, eta)
};
2.0 * w * unit
}
#[inline]
fn bernoulli_deviance(y: f64, mu: f64, w: f64) -> f64 {
2.0 * w * (bd0(y, mu) + bd0(1.0 - y, 1.0 - mu))
}
#[inline]
fn standard_normal_cdf(x: f64) -> f64 {
0.5 * gam_gpu::numerics_host::erfc(-x * std::f64::consts::FRAC_1_SQRT_2)
}
#[inline]
fn standard_normal_pdf(x: f64) -> f64 {
const COEFF: f64 = 0.398_942_280_401_432_7; COEFF * (-0.5 * x * x).exp()
}
#[must_use]
pub struct PirlsRowBackend {
#[cfg(target_os = "linux")]
inner: PirlsRowBackendLinux,
}
#[cfg(target_os = "linux")]
struct PirlsRowBackendLinux {
ctx: Arc<CudaContext>,
modules: Mutex<std::collections::HashMap<ModuleKey, Arc<CudaModule>>>,
jit_modules: Mutex<std::collections::HashMap<JitKey, Arc<CudaModule>>>,
}
#[cfg(target_os = "linux")]
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
enum KernelMode {
FinalRow,
SolveRow,
AlphaLadder,
}
#[cfg(target_os = "linux")]
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
struct ModuleKey {
family: PirlsRowFamily,
curvature: CurvatureMode,
mode: KernelMode,
}
impl PirlsRowBackend {
pub const fn compiled() -> bool {
cfg!(target_os = "linux")
}
pub fn probe() -> Result<&'static Self, GpuError> {
static BACKEND: OnceLock<Result<PirlsRowBackend, GpuError>> = OnceLock::new();
BACKEND
.get_or_init(|| {
#[cfg(target_os = "linux")]
{
Self::probe_linux()
}
#[cfg(not(target_os = "linux"))]
{
Err(GpuError::DriverLibraryUnavailable {
reason: "pirls_row GPU backend is Linux-only".to_string(),
})
}
})
.as_ref()
.map_err(GpuError::clone)
}
#[cfg(target_os = "linux")]
fn probe_linux() -> Result<Self, GpuError> {
let parts = gam_gpu::backend_probe::probe_cuda_backend("pirls_row")?;
Ok(Self {
inner: PirlsRowBackendLinux {
ctx: parts.ctx,
modules: Mutex::new(std::collections::HashMap::new()),
jit_modules: Mutex::new(std::collections::HashMap::new()),
},
})
}
#[cfg(target_os = "linux")]
fn module_for_kind(
&self,
family: PirlsRowFamily,
curvature: CurvatureMode,
mode: KernelMode,
label: &str,
) -> Result<Arc<CudaModule>, GpuError> {
let key = ModuleKey {
family,
curvature,
mode,
};
if let Some(existing) = self
.inner
.modules
.lock()
.gpu_ctx_with(|err| format!("pirls_row {label}module cache mutex poisoned: {err}"))?
.get(&key)
{
return Ok(existing.clone());
}
let source = match mode {
KernelMode::FinalRow => cuda_source_for(family, curvature),
KernelMode::SolveRow => solve_row_source_for(family, curvature),
KernelMode::AlphaLadder => ladder_source_for(family, curvature),
};
let ptx = gam_gpu::device_cache::compile_ptx_arch(&source).gpu_ctx_with(|err| {
format!(
"pirls_row {label}NVRTC compile failed for {family}/{curv}: {err}",
family = family.as_str(),
curv = curvature.as_str(),
)
})?;
let module = self
.inner
.ctx
.load_module(ptx)
.gpu_ctx_with(|err| format!("pirls_row {label}module load failed: {err}"))?;
self.inner
.modules
.lock()
.gpu_ctx_with(|err| format!("pirls_row {label}module cache mutex poisoned: {err}"))?
.insert(key, module.clone());
Ok(module)
}
#[cfg(target_os = "linux")]
pub fn module_for(
&self,
family: PirlsRowFamily,
curvature: CurvatureMode,
) -> Result<Arc<CudaModule>, GpuError> {
self.module_for_kind(family, curvature, KernelMode::FinalRow, "")
}
#[cfg(target_os = "linux")]
pub fn module_for_solve(
&self,
family: PirlsRowFamily,
curvature: CurvatureMode,
) -> Result<Arc<CudaModule>, GpuError> {
self.module_for_kind(family, curvature, KernelMode::SolveRow, "solve ")
}
#[cfg(target_os = "linux")]
pub fn module_for_ladder(
&self,
family: PirlsRowFamily,
curvature: CurvatureMode,
) -> Result<Arc<CudaModule>, GpuError> {
self.module_for_kind(family, curvature, KernelMode::AlphaLadder, "ladder ")
}
#[cfg(target_os = "linux")]
pub fn module_for_jit(
&self,
spec: &JitFamilySpec,
curvature: CurvatureMode,
) -> Result<Arc<CudaModule>, GpuError> {
let key = JitKey {
spec_id: spec.spec_id,
curvature,
};
if let Some(existing) = self
.inner
.jit_modules
.lock()
.gpu_ctx("pirls_row jit cache poisoned")?
.get(&key)
{
return Ok(existing.clone());
}
let source = spec.cuda_source(curvature);
let ptx = gam_gpu::device_cache::compile_ptx_arch(&source).gpu_ctx_with(|err| {
format!(
"pirls_row JIT NVRTC compile failed for spec_id={} curvature={}: {err}",
spec.spec_id,
curvature.as_str(),
)
})?;
let module = self
.inner
.ctx
.load_module(ptx)
.gpu_ctx("pirls_row JIT module load failed")?;
self.inner
.jit_modules
.lock()
.gpu_ctx("pirls_row jit cache poisoned (insert)")?
.insert(key, module.clone());
Ok(module)
}
}
#[cfg(target_os = "linux")]
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
struct JitKey {
spec_id: u64,
curvature: CurvatureMode,
}
#[derive(Clone, Debug)]
pub struct JitFamilySpec {
pub spec_id: u64,
pub body: String,
}
impl JitFamilySpec {
#[cfg(target_os = "linux")]
pub fn glm(
spec_id: u64,
family: PirlsRowFamily,
curvature: CurvatureMode,
gamma_shape: f64,
) -> Self {
let mut body = match family {
PirlsRowFamily::GaussianIdentity => gaussian_identity_body(curvature),
PirlsRowFamily::PoissonLog => poisson_log_body(curvature),
PirlsRowFamily::GammaLog => gamma_log_body(curvature),
PirlsRowFamily::BernoulliLogit => bernoulli_logit_body(curvature),
PirlsRowFamily::BernoulliProbit => bernoulli_probit_body(curvature),
PirlsRowFamily::BernoulliCLogLog => bernoulli_cloglog_body(curvature),
};
if matches!(family, PirlsRowFamily::GammaLog) {
body.insert_str(0, &format!(" const double shape = {gamma_shape:?};\n"));
}
Self { spec_id, body }
}
pub fn raw(spec_id: u64, body: impl Into<String>) -> Self {
Self {
spec_id,
body: body.into(),
}
}
pub fn kernel_name(&self) -> String {
format!("pirls_row_jit_{}", self.spec_id)
}
#[cfg(target_os = "linux")]
pub fn cuda_source(&self, curvature: CurvatureMode) -> String {
let curvature_define = match curvature {
CurvatureMode::Fisher => "#define PIRLS_CURVATURE_FISHER 1",
CurvatureMode::Observed => "#define PIRLS_CURVATURE_OBSERVED 1",
};
let kernel_name = self.kernel_name();
let body = &self.body;
format!(
r#"
{curvature_define}
{prolog}
extern "C" __global__ void {kernel_name}(
int n,
const double* __restrict__ eta,
const double* __restrict__ y,
const double* __restrict__ prior_w,
double* __restrict__ mu_out,
double* __restrict__ grad_eta_out,
double* __restrict__ w_hessian_out,
double* __restrict__ w_solver_out,
double* __restrict__ deviance_out,
unsigned int* __restrict__ status_out
) {{
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= n) return;
unsigned int status = PIRLS_OK;
double eta_i = eta[i];
double y_i = y[i];
double wp = prior_w[i];
{body}
if (status == PIRLS_OK) {{
mu_out[i] = mu;
grad_eta_out[i] = grad_eta;
w_hessian_out[i] = w_hessian;
w_solver_out[i] = w_solver;
deviance_out[i] = dev;
}}
status_out[i] = status;
}}
"#,
prolog = common_device_prolog(),
)
}
}
#[cfg(target_os = "linux")]
pub struct RowOutputDevBuffers {
pub mu: cudarc::driver::CudaSlice<f64>,
pub grad_eta: cudarc::driver::CudaSlice<f64>,
pub w_hessian: cudarc::driver::CudaSlice<f64>,
pub w_solver: cudarc::driver::CudaSlice<f64>,
pub deviance: cudarc::driver::CudaSlice<f64>,
pub status: cudarc::driver::CudaSlice<u32>,
pub n: usize,
}
#[cfg(target_os = "linux")]
impl RowOutputDevBuffers {
pub fn allocate(stream: &Arc<cudarc::driver::CudaStream>, n: usize) -> Result<Self, GpuError> {
let alloc_f64 = |label: &'static str| {
stream
.alloc_zeros::<f64>(n)
.gpu_ctx_with(|err| format!("pirls_row alloc {label}: {err}"))
};
let alloc_u32 = |label: &'static str| {
stream
.alloc_zeros::<u32>(n)
.gpu_ctx_with(|err| format!("pirls_row alloc {label}: {err}"))
};
Ok(Self {
mu: alloc_f64("mu")?,
grad_eta: alloc_f64("grad_eta")?,
w_hessian: alloc_f64("w_hessian")?,
w_solver: alloc_f64("w_solver")?,
deviance: alloc_f64("deviance")?,
status: alloc_u32("status")?,
n,
})
}
}
#[cfg(target_os = "linux")]
pub struct SolveRowBuffers {
pub grad_eta: cudarc::driver::CudaSlice<f64>,
pub w_solver: cudarc::driver::CudaSlice<f64>,
pub deviance: cudarc::driver::CudaSlice<f64>,
pub status: cudarc::driver::CudaSlice<u32>,
pub n: usize,
}
#[cfg(target_os = "linux")]
impl SolveRowBuffers {
pub fn allocate(stream: &Arc<cudarc::driver::CudaStream>, n: usize) -> Result<Self, GpuError> {
let alloc_f64 = |label: &'static str| {
stream
.alloc_zeros::<f64>(n)
.gpu_ctx_with(|err| format!("pirls_row solve alloc {label}: {err}"))
};
let alloc_u32 = |label: &'static str| {
stream
.alloc_zeros::<u32>(n)
.gpu_ctx_with(|err| format!("pirls_row solve alloc {label}: {err}"))
};
Ok(Self {
grad_eta: alloc_f64("grad_eta")?,
w_solver: alloc_f64("w_solver")?,
deviance: alloc_f64("deviance")?,
status: alloc_u32("status")?,
n,
})
}
}
pub const ALPHA_LADDER_LEN: usize = 7;
pub const ALPHA_LADDER: [f64; ALPHA_LADDER_LEN] =
[1.0, 0.5, 0.25, 0.125, 0.0625, 0.03125, 0.015625];
#[cfg(target_os = "linux")]
pub struct AlphaLadderDevBuffers {
pub objective_dev: cudarc::driver::CudaSlice<f64>,
pub status_dev: cudarc::driver::CudaSlice<u32>,
pub n: usize,
}
#[cfg(target_os = "linux")]
impl AlphaLadderDevBuffers {
pub fn allocate(stream: &Arc<cudarc::driver::CudaStream>, n: usize) -> Result<Self, GpuError> {
let status_len = ALPHA_LADDER_LEN.checked_mul(n).ok_or_else(|| {
gam_gpu::gpu_err!("pirls_row ladder status length overflows: {ALPHA_LADDER_LEN} * {n}")
})?;
Ok(Self {
objective_dev: stream
.alloc_zeros::<f64>(ALPHA_LADDER_LEN)
.gpu_ctx_with(|err| format!("pirls_row ladder alloc objective: {err}"))?,
status_dev: stream
.alloc_zeros::<u32>(status_len)
.gpu_ctx_with(|err| format!("pirls_row ladder alloc status: {err}"))?,
n,
})
}
pub fn zero(&mut self, stream: &Arc<cudarc::driver::CudaStream>) -> Result<(), GpuError> {
stream
.memset_zeros(&mut self.objective_dev)
.gpu_ctx_with(|err| format!("pirls_row ladder zero objective: {err}"))?;
stream
.memset_zeros(&mut self.status_dev)
.gpu_ctx_with(|err| format!("pirls_row ladder zero status: {err}"))
}
}
#[cfg(target_os = "linux")]
pub fn launch_row_reweight_on_stream(
backend: &PirlsRowBackend,
family: PirlsRowFamily,
curvature: CurvatureMode,
gamma_shape: f64,
stream: &Arc<cudarc::driver::CudaStream>,
n: usize,
eta_dev: &cudarc::driver::CudaSlice<f64>,
y_dev: &cudarc::driver::CudaSlice<f64>,
prior_w_dev: &cudarc::driver::CudaSlice<f64>,
out: &mut RowOutputDevBuffers,
) -> Result<(), GpuError> {
use cudarc::driver::{LaunchConfig, PushKernelArg};
if out.n != n {
gam_gpu::gpu_bail!("row reweight buffers shape {} mismatches n={n}", out.n);
}
let module = backend.module_for(family, curvature)?;
let func = module
.load_function(family.kernel_name())
.gpu_ctx_with(|err| {
format!(
"row reweight load_function({}): {err}",
family.kernel_name()
)
})?;
const THREADS_PER_BLOCK: u32 = 256;
let n_u32 = u32::try_from(n)
.map_err(|_| gam_gpu::gpu_err!("n={n} exceeds u32 for row reweight grid sizing"))?;
let grid_x = n_u32.div_ceil(THREADS_PER_BLOCK).max(1);
let n_i32 = i32::try_from(n)
.map_err(|_| gam_gpu::gpu_err!("n={n} exceeds i32 for row reweight kernel argument"))?;
let cfg = LaunchConfig {
grid_dim: (grid_x, 1, 1),
block_dim: (THREADS_PER_BLOCK, 1, 1),
shared_mem_bytes: 0,
};
let mut builder = stream.launch_builder(&func);
builder.arg(&n_i32);
builder.arg(eta_dev);
builder.arg(y_dev);
builder.arg(prior_w_dev);
if matches!(family, PirlsRowFamily::GammaLog) {
builder.arg(&gamma_shape);
}
builder.arg(&mut out.mu);
builder.arg(&mut out.grad_eta);
builder.arg(&mut out.w_hessian);
builder.arg(&mut out.w_solver);
builder.arg(&mut out.deviance);
builder.arg(&mut out.status);
unsafe { builder.launch(cfg) }
.map(|_event_pair| ())
.gpu_ctx_with(|err| format!("row reweight launch({}): {err}", family.kernel_name()))
}
#[cfg(target_os = "linux")]
pub fn launch_row_reweight_jit_on_stream(
backend: &PirlsRowBackend,
spec: &JitFamilySpec,
curvature: CurvatureMode,
stream: &Arc<cudarc::driver::CudaStream>,
n: usize,
eta_dev: &cudarc::driver::CudaSlice<f64>,
y_dev: &cudarc::driver::CudaSlice<f64>,
prior_w_dev: &cudarc::driver::CudaSlice<f64>,
out: &mut RowOutputDevBuffers,
) -> Result<(), GpuError> {
use cudarc::driver::{LaunchConfig, PushKernelArg};
if out.n != n {
gam_gpu::gpu_bail!("JIT row reweight buffers shape {} mismatches n={n}", out.n);
}
let module = backend.module_for_jit(spec, curvature)?;
let kernel_name = spec.kernel_name();
let func = module
.load_function(&kernel_name)
.gpu_ctx_with(|err| format!("JIT row reweight load_function({kernel_name}): {err}"))?;
const THREADS_PER_BLOCK: u32 = 256;
let n_u32 = u32::try_from(n)
.map_err(|_| gam_gpu::gpu_err!("n={n} exceeds u32 for JIT row reweight grid sizing"))?;
let grid_x = n_u32.div_ceil(THREADS_PER_BLOCK).max(1);
let n_i32 = i32::try_from(n)
.map_err(|_| gam_gpu::gpu_err!("n={n} exceeds i32 for JIT row reweight kernel argument"))?;
let cfg = LaunchConfig {
grid_dim: (grid_x, 1, 1),
block_dim: (THREADS_PER_BLOCK, 1, 1),
shared_mem_bytes: 0,
};
let mut builder = stream.launch_builder(&func);
builder.arg(&n_i32);
builder.arg(eta_dev);
builder.arg(y_dev);
builder.arg(prior_w_dev);
builder.arg(&mut out.mu);
builder.arg(&mut out.grad_eta);
builder.arg(&mut out.w_hessian);
builder.arg(&mut out.w_solver);
builder.arg(&mut out.deviance);
builder.arg(&mut out.status);
unsafe { builder.launch(cfg) }
.map(|_event_pair| ())
.gpu_ctx_with(|err| format!("JIT row reweight launch({kernel_name}): {err}"))
}
#[cfg(target_os = "linux")]
pub fn launch_solve_row_on_stream(
backend: &PirlsRowBackend,
family: PirlsRowFamily,
curvature: CurvatureMode,
gamma_shape: f64,
stream: &Arc<cudarc::driver::CudaStream>,
n: usize,
eta_dev: &cudarc::driver::CudaSlice<f64>,
y_dev: &cudarc::driver::CudaSlice<f64>,
prior_w_dev: &cudarc::driver::CudaSlice<f64>,
out: &mut SolveRowBuffers,
) -> Result<(), GpuError> {
use cudarc::driver::{LaunchConfig, PushKernelArg};
if out.n != n {
gam_gpu::gpu_bail!("solve-row buffers shape {} mismatches n={n}", out.n);
}
let module = backend.module_for_solve(family, curvature)?;
let kernel_name = family.solve_kernel_name();
let func = module
.load_function(kernel_name)
.gpu_ctx_with(|err| format!("solve-row load_function({kernel_name}): {err}"))?;
const THREADS_PER_BLOCK: u32 = 256;
let n_u32 = u32::try_from(n)
.map_err(|_| gam_gpu::gpu_err!("n={n} exceeds u32 for solve-row grid sizing"))?;
let grid_x = n_u32.div_ceil(THREADS_PER_BLOCK).max(1);
let n_i32 = i32::try_from(n)
.map_err(|_| gam_gpu::gpu_err!("n={n} exceeds i32 for solve-row kernel argument"))?;
let cfg = LaunchConfig {
grid_dim: (grid_x, 1, 1),
block_dim: (THREADS_PER_BLOCK, 1, 1),
shared_mem_bytes: 0,
};
let mut builder = stream.launch_builder(&func);
builder.arg(&n_i32);
builder.arg(eta_dev);
builder.arg(y_dev);
builder.arg(prior_w_dev);
if matches!(family, PirlsRowFamily::GammaLog) {
builder.arg(&gamma_shape);
}
builder.arg(&mut out.grad_eta);
builder.arg(&mut out.w_solver);
builder.arg(&mut out.deviance);
builder.arg(&mut out.status);
unsafe { builder.launch(cfg) }
.map(|_event_pair| ())
.gpu_ctx_with(|err| format!("solve-row launch({kernel_name}): {err}"))
}
#[cfg(target_os = "linux")]
pub fn launch_alpha_ladder_on_stream(
backend: &PirlsRowBackend,
family: PirlsRowFamily,
curvature: CurvatureMode,
gamma_shape: f64,
stream: &Arc<cudarc::driver::CudaStream>,
n: usize,
eta_dev: &cudarc::driver::CudaSlice<f64>,
xd_dev: &cudarc::driver::CudaSlice<f64>,
y_dev: &cudarc::driver::CudaSlice<f64>,
prior_w_dev: &cudarc::driver::CudaSlice<f64>,
out: &mut AlphaLadderDevBuffers,
) -> Result<(), GpuError> {
use cudarc::driver::{LaunchConfig, PushKernelArg};
if out.n != n {
gam_gpu::gpu_bail!("alpha-ladder buffers shape {} mismatches n={n}", out.n);
}
let module = backend.module_for_ladder(family, curvature)?;
let kernel_name = family.ladder_kernel_name();
let func = module
.load_function(kernel_name)
.gpu_ctx_with(|err| format!("alpha-ladder load_function({kernel_name}): {err}"))?;
const THREADS_PER_BLOCK: u32 = 256;
let n_u32 = u32::try_from(n)
.map_err(|_| gam_gpu::gpu_err!("n={n} exceeds u32 for alpha-ladder grid sizing"))?;
let row_blocks = n_u32.div_ceil(THREADS_PER_BLOCK).max(1);
let n_i32 = i32::try_from(n)
.map_err(|_| gam_gpu::gpu_err!("n={n} exceeds i32 for alpha-ladder kernel argument"))?;
let cfg = LaunchConfig {
grid_dim: (row_blocks, ALPHA_LADDER_LEN as u32, 1),
block_dim: (THREADS_PER_BLOCK, 1, 1),
shared_mem_bytes: 0,
};
let mut builder = stream.launch_builder(&func);
builder.arg(&n_i32);
builder.arg(eta_dev);
builder.arg(xd_dev);
builder.arg(y_dev);
builder.arg(prior_w_dev);
if matches!(family, PirlsRowFamily::GammaLog) {
builder.arg(&gamma_shape);
}
builder.arg(&mut out.objective_dev);
builder.arg(&mut out.status_dev);
unsafe { builder.launch(cfg) }
.map(|_event_pair| ())
.gpu_ctx_with(|err| format!("alpha-ladder launch({kernel_name}): {err}"))
}
#[cfg(target_os = "linux")]
fn common_device_prolog() -> String {
r#"
// NVRTC math builtins: prototypes must carry an execution space. Newer
// NVRTC (CUDA 12.x JIT semantics) rejects unannotated declarations outright
// ("host functions are not allowed in JIT mode"), which failed every
// pirls_row kernel compile on real hardware while CPU-only CI stayed green
// (#2313 hardware sweep). `__device__` matches how the CUDA math library
// declares them; the definitions come from libdevice as before.
extern "C" {
__device__ double exp(double);
__device__ double log(double);
__device__ double log1p(double);
__device__ double expm1(double);
__device__ double fabs(double);
__device__ double erfc(double);
}
static constexpr double PIRLS_LOG_ETA_MIN = __PIRLS_LOG_ETA_MIN__;
static constexpr double PIRLS_LOG_ETA_MAX = __PIRLS_LOG_ETA_MAX__;
static constexpr unsigned int PIRLS_OK = 0u;
static constexpr unsigned int PIRLS_ETA_DOMAIN = 1u;
static constexpr unsigned int PIRLS_PRIOR_WEIGHT = 2u;
static constexpr unsigned int PIRLS_RESPONSE = 3u;
static constexpr unsigned int PIRLS_GAMMA_SHAPE = 4u;
static constexpr unsigned int PIRLS_INVERSE_LINK = 5u;
static constexpr unsigned int PIRLS_FISHER_WEIGHT = 6u;
static constexpr unsigned int PIRLS_OBSERVED_WEIGHT = 7u;
static constexpr unsigned int PIRLS_GRADIENT = 8u;
static constexpr unsigned int PIRLS_DEVIANCE = 9u;
static constexpr unsigned int PIRLS_FINAL_OUTPUT = 10u;
__device__ __forceinline__ void pirls_refuse(unsigned int* status, unsigned int code) {
if (*status == PIRLS_OK) *status = code;
}
__device__ __forceinline__ bool pirls_log_eta_valid(double eta) {
return eta >= PIRLS_LOG_ETA_MIN && eta <= PIRLS_LOG_ETA_MAX;
}
__device__ __forceinline__ double softplus(double x) {
return (x > 0.0 ? x : 0.0) + log1p(exp(-fabs(x)));
}
__device__ __forceinline__ double expm1_minus_x(double x) {
if (fabs(x) > 0.5) return expm1(x) - x;
double term = 0.5 * x * x;
double sum = term;
double degree = 2.0;
for (;;) {
degree += 1.0;
term *= x / degree;
double next = sum + term;
if (next == sum) return next;
sum = next;
}
}
__device__ __forceinline__ double log1p_minus_x(double x) {
if (fabs(x) > 0.5) return log1p(x) - x;
double power = x * x;
double sign = -1.0;
double degree = 2.0;
double sum = sign * power / degree;
for (;;) {
power *= x;
sign = -sign;
degree += 1.0;
double next = sum + sign * power / degree;
if (next == sum) return next;
sum = next;
}
}
__device__ __forceinline__ double logistic(double x) {
if (x >= 0.0) return 1.0 / (1.0 + exp(-x));
double e = exp(x);
return e / (1.0 + e);
}
__device__ __forceinline__ double bernoulli_kl_from_logits(double a, double b) {
if (a == b) return 0.0;
double h = b - a;
if (fabs(h) <= 0.5) {
double p = a <= 0.0 ? logistic(a) : logistic(-a);
double local_h = a <= 0.0 ? h : -h;
double em1 = expm1(local_h);
double x = p * em1;
return log1p_minus_x(x) + p * expm1_minus_x(local_h);
}
if (a <= 0.0) {
double p = logistic(a);
return p * (a - b) + softplus(b) - softplus(a);
}
double q = logistic(-a);
return q * (b - a) + softplus(-b) - softplus(-a);
}
__device__ __forceinline__ double bd0(double x, double m) {
if (x == 0.0) return m;
if (x == m) return 0.0;
double hi = x > m ? x : m;
double lo = x < m ? x : m;
if (fabs(x - m) / hi < 0.2) {
double v = ((x - m) / hi) / (1.0 + lo / hi);
double sum = (x - m) * v;
double term = 2.0 * x * v;
double v2 = v * v;
double denominator = 3.0;
for (;;) {
term *= v2;
double next = sum + term / denominator;
if (next == sum) return next;
sum = next;
denominator += 2.0;
}
}
return x * (log(x) - log(m)) + (m - x);
}
__device__ __forceinline__ double bernoulli_deviance(double y, double mu, double w) {
return 2.0 * w * (bd0(y, mu) + bd0(1.0 - y, 1.0 - mu));
}
__device__ __forceinline__ double logit_deviance(double y, double eta, double w) {
double unit;
if (y == 0.0) unit = softplus(eta);
else if (y == 1.0) unit = softplus(-eta);
else {
double response_logit = log(y) - log1p(-y);
unit = bernoulli_kl_from_logits(response_logit, eta);
}
return 2.0 * w * unit;
}
__device__ __forceinline__ double std_norm_cdf(double x) {
return 0.5 * erfc(-x * 0.7071067811865475);
}
__device__ __forceinline__ double std_norm_pdf(double x) {
return 0.3989422804014327 * exp(-0.5 * x * x);
}
__device__ __forceinline__ double positive_mul_div(double a, double b, double c) {
double product = a * b;
if (isfinite(product) && product > 0.0) {
double value = product / c;
if (isfinite(value) && value > 0.0) return value;
}
double quotient_a = a / c;
if (isfinite(quotient_a) && quotient_a > 0.0) {
double value = quotient_a * b;
if (isfinite(value) && value > 0.0) return value;
}
double quotient_b = b / c;
if (isfinite(quotient_b) && quotient_b > 0.0) {
double value = quotient_b * a;
if (isfinite(value) && value > 0.0) return value;
}
return product / c;
}
__device__ __forceinline__ double gamma_unit_deviance_near_one(double u) {
if (fabs(u) > 0.125) return u - log1p(u);
double power = u * u;
double sum = 0.5 * power;
for (int degree = 3; degree <= 32; ++degree) {
power *= u;
double term = power / (double)degree;
double next = sum + ((degree & 1) ? -term : term);
if (next == sum) break;
sum = next;
}
return sum;
}
__device__ __forceinline__ double poisson_unit_deviance_near_one(double u) {
if (fabs(u) > 0.125) return (1.0 + u) * log1p(u) - u;
double power = u * u;
double sum = 0.5 * power;
for (int degree = 3; degree <= 32; ++degree) {
power *= u;
double coefficient = ((degree & 1) ? -1.0 : 1.0)
/ ((double)degree * (double)(degree - 1));
double next = sum + coefficient * power;
if (next == sum) break;
sum = next;
}
return sum;
}
__device__ __forceinline__ bool pirls_outputs_finite(
double mu, double grad_eta, double w_fisher, double w_hessian,
double w_solver, double dev
) {
return isfinite(mu) && isfinite(grad_eta) && isfinite(w_fisher)
&& isfinite(w_hessian) && isfinite(w_solver) && isfinite(dev);
}
"#
.replace(
"__PIRLS_LOG_ETA_MIN__",
&format!("{:?}", crate::mixture_link::LOG_LINK_SOLVER_ETA_MIN),
)
.replace(
"__PIRLS_LOG_ETA_MAX__",
&format!("{:?}", crate::mixture_link::LOG_LINK_SOLVER_ETA_MAX),
)
}
#[cfg(target_os = "linux")]
fn cuda_source_for(family: PirlsRowFamily, curvature: CurvatureMode) -> String {
let body = match family {
PirlsRowFamily::GaussianIdentity => gaussian_identity_body(curvature),
PirlsRowFamily::PoissonLog => poisson_log_body(curvature),
PirlsRowFamily::GammaLog => gamma_log_body(curvature),
PirlsRowFamily::BernoulliLogit => bernoulli_logit_body(curvature),
PirlsRowFamily::BernoulliProbit => bernoulli_probit_body(curvature),
PirlsRowFamily::BernoulliCLogLog => bernoulli_cloglog_body(curvature),
};
let kernel_name = family.kernel_name();
let curvature_define = match curvature {
CurvatureMode::Fisher => "#define PIRLS_CURVATURE_FISHER 1",
CurvatureMode::Observed => "#define PIRLS_CURVATURE_OBSERVED 1",
};
let shape_param = if matches!(family, PirlsRowFamily::GammaLog) {
" double shape,\n"
} else {
""
};
format!(
r#"
{curvature_define}
{prolog}
extern "C" __global__ void {kernel_name}(
int n,
const double* __restrict__ eta,
const double* __restrict__ y,
const double* __restrict__ prior_w,
{shape_param} double* __restrict__ mu_out,
double* __restrict__ grad_eta_out,
double* __restrict__ w_hessian_out,
double* __restrict__ w_solver_out,
double* __restrict__ deviance_out,
unsigned int* __restrict__ status_out
) {{
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= n) return;
unsigned int status = PIRLS_OK;
double eta_i = eta[i];
double y_i = y[i];
double wp = prior_w[i];
{body}
if (status == PIRLS_OK) {{
mu_out[i] = mu;
grad_eta_out[i] = grad_eta;
w_hessian_out[i] = w_hessian;
w_solver_out[i] = w_solver;
deviance_out[i] = dev;
}}
status_out[i] = status;
}}
"#,
prolog = common_device_prolog(),
)
}
#[cfg(target_os = "linux")]
#[inline]
fn curvature_tag(curvature: CurvatureMode) -> &'static str {
match curvature {
CurvatureMode::Fisher => " // curvature: fisher\n",
CurvatureMode::Observed => " // curvature: observed\n",
}
}
#[cfg(target_os = "linux")]
fn gaussian_identity_body(curvature: CurvatureMode) -> String {
let tag = curvature_tag(curvature);
format!(
r#"{tag} double mu = 0.0, grad_eta = 0.0, w_fisher = 0.0;
double w_hessian = 0.0, w_solver = 0.0, dev = 0.0;
if (!isfinite(eta_i)) pirls_refuse(&status, PIRLS_ETA_DOMAIN);
if (status == PIRLS_OK && !(isfinite(wp) && wp >= 0.0))
pirls_refuse(&status, PIRLS_PRIOR_WEIGHT);
if (status == PIRLS_OK && wp > 0.0 && !isfinite(y_i))
pirls_refuse(&status, PIRLS_RESPONSE);
if (status == PIRLS_OK) {{
mu = eta_i;
w_fisher = wp;
w_hessian = wp;
w_solver = w_hessian;
if (wp > 0.0) {{
double resid = y_i - mu;
grad_eta = wp * resid;
dev = wp * resid * resid;
}}
}}
if (status == PIRLS_OK && !pirls_outputs_finite(
mu, grad_eta, w_fisher, w_hessian, w_solver, dev))
pirls_refuse(&status, PIRLS_FINAL_OUTPUT);
"#
)
}
#[cfg(target_os = "linux")]
fn poisson_log_body(curvature: CurvatureMode) -> String {
let tag = curvature_tag(curvature);
format!(
r#"{tag} double mu = 0.0, grad_eta = 0.0, w_fisher = 0.0;
double w_hessian = 0.0, w_solver = 0.0, dev = 0.0;
if (!pirls_log_eta_valid(eta_i)) pirls_refuse(&status, PIRLS_ETA_DOMAIN);
if (status == PIRLS_OK && !(isfinite(wp) && wp >= 0.0))
pirls_refuse(&status, PIRLS_PRIOR_WEIGHT);
if (status == PIRLS_OK && wp > 0.0 && !(isfinite(y_i) && y_i >= 0.0))
pirls_refuse(&status, PIRLS_RESPONSE);
if (status == PIRLS_OK) {{
mu = exp(eta_i);
if (!(isfinite(mu) && mu > 0.0)) pirls_refuse(&status, PIRLS_INVERSE_LINK);
}}
if (status == PIRLS_OK && wp > 0.0) {{
w_fisher = wp * mu;
if (!(isfinite(w_fisher) && w_fisher > 0.0))
pirls_refuse(&status, PIRLS_FISHER_WEIGHT);
if (status == PIRLS_OK) {{
w_hessian = w_fisher;
w_solver = w_hessian;
grad_eta = wp * (y_i - mu);
double u = (y_i - mu) / mu;
double dev_base;
if (y_i == 0.0) {{
dev_base = w_fisher;
}} else {{
double scaled_unit = w_fisher * poisson_unit_deviance_near_one(u);
if (isfinite(scaled_unit) && scaled_unit >= 0.0) {{
dev_base = scaled_unit;
}} else {{
double weighted_y = positive_mul_div(w_fisher, y_i, mu);
dev_base = weighted_y * (log(y_i) - eta_i - 1.0) + w_fisher;
}}
}}
if (!isfinite(grad_eta)) pirls_refuse(&status, PIRLS_GRADIENT);
dev = 2.0 * dev_base;
if (!isfinite(dev)) pirls_refuse(&status, PIRLS_DEVIANCE);
}}
}}
if (status == PIRLS_OK && !pirls_outputs_finite(
mu, grad_eta, w_fisher, w_hessian, w_solver, dev))
pirls_refuse(&status, PIRLS_FINAL_OUTPUT);
"#
)
}
#[cfg(target_os = "linux")]
fn gamma_log_body(curvature: CurvatureMode) -> String {
let tag = curvature_tag(curvature);
format!(
r#"{tag} double mu = 0.0, grad_eta = 0.0, w_fisher = 0.0;
double w_hessian = 0.0, w_solver = 0.0, dev = 0.0;
if (!pirls_log_eta_valid(eta_i)) pirls_refuse(&status, PIRLS_ETA_DOMAIN);
if (status == PIRLS_OK && !(isfinite(shape) && shape > 0.0))
pirls_refuse(&status, PIRLS_GAMMA_SHAPE);
if (status == PIRLS_OK && !(isfinite(wp) && wp >= 0.0))
pirls_refuse(&status, PIRLS_PRIOR_WEIGHT);
if (status == PIRLS_OK && wp > 0.0 && !(isfinite(y_i) && y_i > 0.0))
pirls_refuse(&status, PIRLS_RESPONSE);
if (status == PIRLS_OK) {{
mu = exp(eta_i);
if (!(isfinite(mu) && mu > 0.0)) pirls_refuse(&status, PIRLS_INVERSE_LINK);
}}
if (status == PIRLS_OK && wp > 0.0) {{
w_fisher = wp * shape;
if (!(isfinite(w_fisher) && w_fisher > 0.0))
pirls_refuse(&status, PIRLS_FISHER_WEIGHT);
#ifdef PIRLS_CURVATURE_OBSERVED
double weighted_ratio_observed = positive_mul_div(w_fisher, y_i, mu);
if (!(isfinite(weighted_ratio_observed) && weighted_ratio_observed > 0.0))
pirls_refuse(&status, PIRLS_OBSERVED_WEIGHT);
w_hessian = weighted_ratio_observed;
#else
w_hessian = w_fisher;
#endif
if (!isfinite(w_hessian)) pirls_refuse(&status, PIRLS_OBSERVED_WEIGHT);
w_solver = w_hessian;
double u = (y_i - mu) / mu;
double scaled_unit = w_fisher * gamma_unit_deviance_near_one(u);
bool need_weighted_ratio = !isfinite(u)
|| !(isfinite(scaled_unit) && scaled_unit >= 0.0);
double weighted_ratio = 0.0;
#ifdef PIRLS_CURVATURE_OBSERVED
weighted_ratio = weighted_ratio_observed;
#else
if (need_weighted_ratio)
weighted_ratio = positive_mul_div(w_fisher, y_i, mu);
#endif
grad_eta = isfinite(u) ? w_fisher * u : weighted_ratio - w_fisher;
double dev_base;
if (isfinite(scaled_unit) && scaled_unit >= 0.0) {{
dev_base = scaled_unit;
}} else {{
dev_base = weighted_ratio - w_fisher * (1.0 + log(y_i) - eta_i);
}}
if (!isfinite(grad_eta)) pirls_refuse(&status, PIRLS_GRADIENT);
dev = 2.0 * dev_base;
if (!isfinite(dev)) pirls_refuse(&status, PIRLS_DEVIANCE);
}}
if (status == PIRLS_OK && !pirls_outputs_finite(
mu, grad_eta, w_fisher, w_hessian, w_solver, dev))
pirls_refuse(&status, PIRLS_FINAL_OUTPUT);
"#
)
}
#[cfg(target_os = "linux")]
fn bernoulli_logit_body(curvature: CurvatureMode) -> String {
let tag = curvature_tag(curvature);
format!(
r#"{tag} double mu = 0.0, grad_eta = 0.0, w_fisher = 0.0;
double w_hessian = 0.0, w_solver = 0.0, dev = 0.0;
if (!isfinite(eta_i)) pirls_refuse(&status, PIRLS_ETA_DOMAIN);
if (status == PIRLS_OK && !(isfinite(wp) && wp >= 0.0))
pirls_refuse(&status, PIRLS_PRIOR_WEIGHT);
if (status == PIRLS_OK && wp > 0.0
&& !(isfinite(y_i) && y_i >= 0.0 && y_i <= 1.0))
pirls_refuse(&status, PIRLS_RESPONSE);
double tail = exp(-fabs(eta_i));
double denom = 1.0 + tail;
double dmu_deta = tail / (denom * denom);
if (status == PIRLS_OK) {{
mu = eta_i >= 0.0 ? 1.0 / denom : tail / denom;
if (!(isfinite(mu) && mu >= 0.0 && mu <= 1.0
&& isfinite(dmu_deta) && dmu_deta > 0.0))
pirls_refuse(&status, PIRLS_INVERSE_LINK);
}}
if (status == PIRLS_OK && wp > 0.0) {{
double residual;
if (eta_i >= 0.0) {{
double one_minus_mu = tail / denom;
residual = y_i == 1.0 ? one_minus_mu : (y_i - 1.0) + one_minus_mu;
}} else {{
residual = y_i - mu;
}}
w_fisher = wp * dmu_deta;
if (!(isfinite(w_fisher) && w_fisher > 0.0))
pirls_refuse(&status, PIRLS_FISHER_WEIGHT);
w_hessian = w_fisher;
w_solver = w_hessian;
grad_eta = wp * residual;
if (!isfinite(grad_eta)) pirls_refuse(&status, PIRLS_GRADIENT);
dev = logit_deviance(y_i, eta_i, wp);
if (!isfinite(dev)) pirls_refuse(&status, PIRLS_DEVIANCE);
}}
if (status == PIRLS_OK && !pirls_outputs_finite(
mu, grad_eta, w_fisher, w_hessian, w_solver, dev))
pirls_refuse(&status, PIRLS_FINAL_OUTPUT);
"#
)
}
#[cfg(target_os = "linux")]
fn bernoulli_probit_body(curvature: CurvatureMode) -> String {
let tag = curvature_tag(curvature);
format!(
r#"{tag} double mu = 0.0, grad_eta = 0.0, w_fisher = 0.0;
double w_hessian = 0.0, w_solver = 0.0, dev = 0.0;
if (!isfinite(eta_i)) pirls_refuse(&status, PIRLS_ETA_DOMAIN);
if (status == PIRLS_OK && !(isfinite(wp) && wp >= 0.0))
pirls_refuse(&status, PIRLS_PRIOR_WEIGHT);
if (status == PIRLS_OK && wp > 0.0
&& !(isfinite(y_i) && y_i >= 0.0 && y_i <= 1.0))
pirls_refuse(&status, PIRLS_RESPONSE);
double dmu_deta = 0.0, d2mu_deta2 = 0.0, v = 0.0;
if (status == PIRLS_OK) {{
mu = std_norm_cdf(eta_i);
dmu_deta = std_norm_pdf(eta_i);
d2mu_deta2 = -eta_i * dmu_deta;
if (!(isfinite(mu) && mu > 0.0 && mu < 1.0
&& isfinite(dmu_deta) && dmu_deta > 0.0
&& isfinite(d2mu_deta2)))
pirls_refuse(&status, PIRLS_INVERSE_LINK);
}}
if (status == PIRLS_OK && wp > 0.0) {{
v = mu * (1.0 - mu);
double fisher_per_prior = dmu_deta * dmu_deta / v;
w_fisher = wp * fisher_per_prior;
if (!(isfinite(v) && v > 0.0 && isfinite(fisher_per_prior)
&& fisher_per_prior > 0.0 && isfinite(w_fisher) && w_fisher > 0.0))
pirls_refuse(&status, PIRLS_FISHER_WEIGHT);
double resid = y_i - mu;
#ifdef PIRLS_CURVATURE_OBSERVED
double bracket = d2mu_deta2 / v
- (dmu_deta * dmu_deta) * (1.0 - 2.0 * mu) / (v * v);
w_hessian = w_fisher - wp * resid * bracket;
#else
w_hessian = w_fisher;
#endif
if (!isfinite(w_hessian)) pirls_refuse(&status, PIRLS_OBSERVED_WEIGHT);
w_solver = w_hessian;
grad_eta = wp * resid * dmu_deta / v;
if (!isfinite(grad_eta)) pirls_refuse(&status, PIRLS_GRADIENT);
dev = bernoulli_deviance(y_i, mu, wp);
if (!isfinite(dev)) pirls_refuse(&status, PIRLS_DEVIANCE);
}}
if (status == PIRLS_OK && !pirls_outputs_finite(
mu, grad_eta, w_fisher, w_hessian, w_solver, dev))
pirls_refuse(&status, PIRLS_FINAL_OUTPUT);
"#
)
}
#[cfg(target_os = "linux")]
fn bernoulli_cloglog_body(curvature: CurvatureMode) -> String {
let tag = curvature_tag(curvature);
format!(
r#"{tag} double mu = 0.0, grad_eta = 0.0, w_fisher = 0.0;
double w_hessian = 0.0, w_solver = 0.0, dev = 0.0;
if (!isfinite(eta_i)) pirls_refuse(&status, PIRLS_ETA_DOMAIN);
if (status == PIRLS_OK && !(isfinite(wp) && wp >= 0.0))
pirls_refuse(&status, PIRLS_PRIOR_WEIGHT);
if (status == PIRLS_OK && wp > 0.0
&& !(isfinite(y_i) && y_i >= 0.0 && y_i <= 1.0))
pirls_refuse(&status, PIRLS_RESPONSE);
double inner = 0.0, dmu_deta = 0.0, d2mu_deta2 = 0.0, v = 0.0;
if (status == PIRLS_OK) {{
inner = exp(eta_i);
double complement = exp(-inner);
mu = -expm1(-inner);
dmu_deta = inner * complement;
d2mu_deta2 = dmu_deta * (1.0 - inner);
if (!(isfinite(mu) && mu > 0.0 && mu < 1.0
&& isfinite(dmu_deta) && dmu_deta > 0.0
&& isfinite(d2mu_deta2)))
pirls_refuse(&status, PIRLS_INVERSE_LINK);
}}
if (status == PIRLS_OK && wp > 0.0) {{
v = mu * (1.0 - mu);
double fisher_per_prior = dmu_deta * dmu_deta / v;
w_fisher = wp * fisher_per_prior;
if (!(isfinite(v) && v > 0.0 && isfinite(fisher_per_prior)
&& fisher_per_prior > 0.0 && isfinite(w_fisher) && w_fisher > 0.0))
pirls_refuse(&status, PIRLS_FISHER_WEIGHT);
double resid = y_i - mu;
#ifdef PIRLS_CURVATURE_OBSERVED
double bracket = d2mu_deta2 / v
- (dmu_deta * dmu_deta) * (1.0 - 2.0 * mu) / (v * v);
w_hessian = w_fisher - wp * resid * bracket;
#else
w_hessian = w_fisher;
#endif
if (!isfinite(w_hessian)) pirls_refuse(&status, PIRLS_OBSERVED_WEIGHT);
w_solver = w_hessian;
grad_eta = wp * resid * dmu_deta / v;
if (!isfinite(grad_eta)) pirls_refuse(&status, PIRLS_GRADIENT);
dev = bernoulli_deviance(y_i, mu, wp);
if (!isfinite(dev)) pirls_refuse(&status, PIRLS_DEVIANCE);
}}
if (status == PIRLS_OK && !pirls_outputs_finite(
mu, grad_eta, w_fisher, w_hessian, w_solver, dev))
pirls_refuse(&status, PIRLS_FINAL_OUTPUT);
"#
)
}
#[cfg(target_os = "linux")]
fn solve_row_source_for(family: PirlsRowFamily, curvature: CurvatureMode) -> String {
let body = match family {
PirlsRowFamily::GaussianIdentity => gaussian_identity_body(curvature),
PirlsRowFamily::PoissonLog => poisson_log_body(curvature),
PirlsRowFamily::GammaLog => gamma_log_body(curvature),
PirlsRowFamily::BernoulliLogit => bernoulli_logit_body(curvature),
PirlsRowFamily::BernoulliProbit => bernoulli_probit_body(curvature),
PirlsRowFamily::BernoulliCLogLog => bernoulli_cloglog_body(curvature),
};
let kernel_name = family.solve_kernel_name();
let curvature_define = match curvature {
CurvatureMode::Fisher => "#define PIRLS_CURVATURE_FISHER 1",
CurvatureMode::Observed => "#define PIRLS_CURVATURE_OBSERVED 1",
};
let shape_param = if matches!(family, PirlsRowFamily::GammaLog) {
" double shape,\n"
} else {
""
};
format!(
r#"
{curvature_define}
{prolog}
extern "C" __global__ void {kernel_name}(
int n,
const double* __restrict__ eta,
const double* __restrict__ y,
const double* __restrict__ prior_w,
{shape_param} double* __restrict__ grad_eta_out,
double* __restrict__ w_solver_out,
double* __restrict__ deviance_out,
unsigned int* __restrict__ status_out
) {{
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= n) return;
unsigned int status = PIRLS_OK;
double eta_i = eta[i];
double y_i = y[i];
double wp = prior_w[i];
{body}
if (status == PIRLS_OK) {{
grad_eta_out[i] = grad_eta;
w_solver_out[i] = w_solver;
deviance_out[i] = dev;
}}
status_out[i] = status;
}}
"#,
prolog = common_device_prolog(),
)
}
#[cfg(target_os = "linux")]
const ALPHA_LADDER_CUDA_ARRAY: &str =
"__constant__ double PIRLS_ALPHAS[7] = {1.0, 0.5, 0.25, 0.125, 0.0625, 0.03125, 0.015625};";
#[cfg(target_os = "linux")]
fn ladder_source_for(family: PirlsRowFamily, curvature: CurvatureMode) -> String {
let body = match family {
PirlsRowFamily::GaussianIdentity => gaussian_identity_body(curvature),
PirlsRowFamily::PoissonLog => poisson_log_body(curvature),
PirlsRowFamily::GammaLog => gamma_log_body(curvature),
PirlsRowFamily::BernoulliLogit => bernoulli_logit_body(curvature),
PirlsRowFamily::BernoulliProbit => bernoulli_probit_body(curvature),
PirlsRowFamily::BernoulliCLogLog => bernoulli_cloglog_body(curvature),
};
let kernel_name = family.ladder_kernel_name();
let curvature_define = match curvature {
CurvatureMode::Fisher => "#define PIRLS_CURVATURE_FISHER 1",
CurvatureMode::Observed => "#define PIRLS_CURVATURE_OBSERVED 1",
};
let shape_param = if matches!(family, PirlsRowFamily::GammaLog) {
" double shape,\n"
} else {
""
};
format!(
r#"
{curvature_define}
{prolog}
{alphas}
extern "C" __global__ void {kernel_name}(
int n,
const double* __restrict__ eta,
const double* __restrict__ xd,
const double* __restrict__ y,
const double* __restrict__ prior_w,
{shape_param} double* __restrict__ objective_out,
unsigned int* __restrict__ status_out
) {{
int i = blockIdx.x * blockDim.x + threadIdx.x;
int k = (int)blockIdx.y;
if (i >= n) return;
unsigned int status = PIRLS_OK;
double alpha = PIRLS_ALPHAS[k];
double eta_i = eta[i] + alpha * xd[i];
double y_i = y[i];
double wp = prior_w[i];
{body}
if (status == PIRLS_OK) atomicAdd(&objective_out[k], dev);
status_out[k * n + i] = status;
}}
"#,
prolog = common_device_prolog(),
alphas = ALPHA_LADDER_CUDA_ARRAY,
)
}
#[cfg(test)]
#[path = "pirls_row_tests.rs"]
mod pirls_row_tests;