gam_solve/gpu/
arrow_schur_gpu.rs

1//! Caller-facing thin wrapper around `crate::gpu_kernels::arrow_schur`.
2//!
3//! The entire dense per-row factor + Schur reduce + back-sub pipeline lives
4//! device-side; this module only translates the device failure enum into the
5//! `ArrowSchurError` variant the PIRLS outer loop already understands, so
6//! call-sites do not need to learn the device-specific reason codes.
7//!
8//! ## Dispatch logic for matrix-free systems
9//!
10//! When `solve_arrow_newton_step` returns `GpuRequiresDenseSystem`, the GPU
11//! dense-Schur path is structurally incompatible with the supplied operators.
12//! This wrapper routes such systems to CPU `InexactPCG` — the mode that was
13//! designed precisely for SAE-manifold scale callers that cannot materialise
14//! a dense `K × K` block. No information is lost: `GpuRequiresDenseSystem`
15//! is not a numerical failure, just a capability mismatch, so the CPU solver
16//! receives the full system without escalating any ridge.
17
18use crate::arrow_schur::{ArrowSchurError, ArrowSchurSystem, ArrowSolveOptions, ArrowSolverMode};
19use crate::gpu_kernels::arrow_schur::{
20    ArrowSchurGpuFailure, gpu_schur_matvec_backend, solve_arrow_newton_step,
21};
22use ndarray::Array1;
23
24pub fn solve_arrow_newton_step_gpu(
25    sys: &ArrowSchurSystem,
26    ridge_t: f64,
27    ridge_beta: f64,
28) -> Result<(Array1<f64>, Array1<f64>), ArrowSchurError> {
29    match solve_arrow_newton_step(sys, ridge_t, ridge_beta) {
30        Ok(solution) => Ok((solution.delta_t, solution.delta_beta)),
31        Err(ArrowSchurGpuFailure::Unavailable) => {
32            // Mirror the CPU path's failure variant so the outer loop falls
33            // through to its existing recovery logic.
34            sys.solve_with_options(ridge_t, ridge_beta, &ArrowSolveOptions::automatic(sys.k))
35                .map(|(dt, db, _diag)| (dt, db))
36        }
37        Err(ArrowSchurGpuFailure::GpuRequiresDenseSystem { .. }) => {
38            // Matrix-free H_ββ or H_tβ operators present — the dense GPU Schur
39            // path cannot consume them, but the reduced K-system PCG can.
40            // Build the GPU-backed reduced Schur matvec (row-procedural sparse
41            // Kronecker apply over active atoms; per-row latent eliminated via
42            // cached factors) and run `InexactPCG` against it. Only when the
43            // device matvec is genuinely `Unavailable` do we fall back to the
44            // pure-CPU `InexactPCG` matvec.
45            let mut opts = ArrowSolveOptions::automatic(sys.k);
46            opts.mode = ArrowSolverMode::InexactPCG;
47            match gpu_schur_matvec_backend(sys, ridge_t, ridge_beta) {
48                Ok(gpu_matvec) => {
49                    opts.gpu_matvec = Some(gpu_matvec);
50                }
51                Err(ArrowSchurGpuFailure::Unavailable) => {
52                    // No device matvec available; CPU InexactPCG owns the solve.
53                }
54                Err(ArrowSchurGpuFailure::RidgeBumpRequired { row, bump }) => {
55                    return Err(ArrowSchurError::PerRowFactorFailed {
56                        row,
57                        reason: format!(
58                            "GPU row-procedural factor failed; suggested ridge bump {bump:.3e}"
59                        ),
60                    });
61                }
62                Err(ArrowSchurGpuFailure::GpuRequiresDenseSystem { .. }) => {
63                    // The matvec builder cannot lift this system either; CPU
64                    // InexactPCG matvec handles the reduction.
65                }
66                Err(ArrowSchurGpuFailure::SchurFactorFailed { reason }) => {
67                    return Err(ArrowSchurError::SchurFactorFailed { reason });
68                }
69            }
70            sys.solve_with_options(ridge_t, ridge_beta, &opts)
71                .map(|(dt, db, _diag)| (dt, db))
72        }
73        Err(ArrowSchurGpuFailure::RidgeBumpRequired { row, bump }) => {
74            Err(ArrowSchurError::PerRowFactorFailed {
75                row,
76                reason: format!("GPU Cholesky factor failed; suggested ridge bump {bump:.3e}"),
77            })
78        }
79        Err(ArrowSchurGpuFailure::SchurFactorFailed { reason }) => {
80            Err(ArrowSchurError::SchurFactorFailed { reason })
81        }
82    }
83}
gam_solve/gpu/arrow_schur_gpu.rs

gam_solve/gpu/
arrow_schur_gpu.rs
