Skip to main content

gam_solve/gpu_kernels/
arrow_schur.rs

1//! Fully GPU-resident batched Arrow-Schur dense Cholesky solver.
2//!
3//! Implements the square-root Schur form: each local block `D_i = L_i L_i^T`
4//! is factored on device, `u_i = L_i^{-1} g_i` and `Y_i = L_i^{-1} B_i` are
5//! formed by triangular solves, the reduced shared system
6//!     `S_β = C + ρ_β I − Σ_i Y_i^T Y_i,  r_β = -g_β + Σ_i Y_i^T u_i`
7//! is assembled on device, factored once, and the back-substitution
8//!     `w_i = u_i + Y_i · δβ,  L_i^T x_i = w_i,  δt_i = -x_i`
9//! is run on device. Only the final `(δt, δβ, log|H|)` triple is downloaded.
10//!
11//! The current caller (Arrow-Schur Newton step inside PIRLS) feeds uniform
12//! local block size `d` and uniform shared width `k`, so the entire pipeline
13//! is dispatched as a single p-group; per-p grouping for heterogenous blocks
14//! is Layer D's NVRTC fused-kernel concern and lives in this module's
15//! follow-up implementation rather than in policy plumbing.
16//!
17//! CUDA-only probes are exported only on Linux. Platform-neutral dispatch
18//! entries remain available so their callers can report a typed device decline.
19
20use ndarray::{Array1, Array2, ArrayView2};
21
22use crate::arrow_schur::{ArrowPcgDiagnostics, ArrowSchurSystem, DeviceSaePcgData};
23// Consumed only by the Linux device path (`canonicalize_device_beta_factor` and
24// the `cuda` module). Importing them unconditionally makes them dead on every
25// other target, which `-D warnings` rejects on the windows-gnu cross-check.
26#[cfg(target_os = "linux")]
27use crate::arrow_schur::{
28    ArrowBetaGaugeQuotient, ArrowSolveOptions, solve_dense_reduced_system,
29};
30use gam_linalg::triangular::{CholeskyGuard, cholesky_factor_in_place, cholesky_solve_vector};
31
32/// Outcome of a single Arrow-Schur Newton solve.
33pub struct ArrowSchurGpuSolution {
34    pub delta_t: Array1<f64>,
35    pub delta_beta: Array1<f64>,
36    /// Natural log of the determinant of the full bordered Hessian, computed
37    /// from the local Cholesky factors and the Schur factor on device.
38    pub log_det_hessian: f64,
39}
40
41/// Reason a device path declined to run; lets the host caller decide between
42/// CPU fallback and per-row escalation. `RidgeBumpRequired` carries the
43/// estimated diagonal bump needed to clear the failed pivot.
44#[derive(Debug, Clone)]
45pub enum ArrowSchurGpuFailure {
46    /// CUDA runtime unavailable, allocation failed, or workload below policy.
47    Unavailable,
48    /// A row block was not positive definite even after the requested ridge.
49    /// Caller may retry with `ridge_t + bump`.
50    RidgeBumpRequired { row: usize, bump: f64 },
51    /// Shared Schur factor failed; bordered system is rank-deficient at the
52    /// requested ridges and the CPU path should handle escalation.
53    SchurFactorFailed { reason: String },
54    /// The dense GPU Schur path cannot consume this system's β-block. Either
55    /// the system carries matrix-free `H_ββ` / per-row `H_tβ` operators
56    /// (`had_*_matvec` set), OR the dense `(K×K)` `H_ββ` block is simply absent
57    /// (both flags false) — e.g. an SAE-manifold system whose β-curvature lives
58    /// in a `penalty_op` / factored-frame representation with `hbb` reclaimed to
59    /// a `0×0` workspace. In BOTH cases this is a capability mismatch, NOT a
60    /// numerical failure: the caller should route to CPU `InexactPCG` (or supply
61    /// dense buffers) rather than escalating a ridge. See `gpu/arrow_schur.rs`
62    /// Part B for the planned GPU PCG path that will lift this restriction at
63    /// K ≥ 5000.
64    GpuRequiresDenseSystem {
65        had_hbb_matvec: bool,
66        had_htbeta_matvec: bool,
67    },
68}
69
70/// Resolve the configured runtime without conflating probe faults with an
71/// ordinary device decline. The existing GPU failure surface has a diagnostic
72/// string variant, so faults flow through it; only typed Auto/Off absence
73/// remains `Ok(None)` and may become `Unavailable` at a caller-specific gate.
74///
75/// Every call site sits inside a `cfg(target_os = "linux")` block (the CUDA
76/// backend compiles only there), so off-Linux this resolver has no callers and
77/// `-D dead-code` rejects it — which is what silently killed the Windows and
78/// macOS wheel jobs. Gate it to the platform that owns its callers rather than
79/// suppressing the lint.
80#[cfg(target_os = "linux")]
81fn resolve_runtime_for_device_path(
82) -> Result<Option<&'static gam_gpu::GpuRuntime>, ArrowSchurGpuFailure> {
83    gam_gpu::device_runtime::GpuRuntime::resolve(gam_gpu::global_policy()).map_err(|error| {
84        ArrowSchurGpuFailure::SchurFactorFailed {
85            reason: format!("GPU runtime resolution failed: {error}"),
86        }
87    })
88}
89
90/// Relative rounding margin (multiplier on `diag_scale · √ε`) added on top of
91/// the deficit-clearing shift in [`ridge_bump_to_make_pd`].
92///
93/// The exact shift `-(λ_min)` makes a block PD in exact arithmetic, but a
94/// single retry at precisely that magnitude is routinely re-rejected by the
95/// next POTRF because the rounding error of forming `D + ridge·I` and
96/// re-factoring is itself O(√ε). The 1024× headroom (≈ 2¹⁰, ten extra bits
97/// below the f64 mantissa's 52) clears the pivot on the first retry without
98/// materially perturbing the curvature the Newton step sees. Shared by every
99/// per-row / batched / fused producer so they suggest a consistent bump.
100const RIDGE_BUMP_EPS_MARGIN: f64 = 1024.0;
101
102/// Diagonal ridge bump that is GUARANTEED to make `H_tt + (ridge_t + bump)·I`
103/// positive definite for a *symmetric* per-row block, sized from the block's
104/// own entries rather than from the factorization's pivot index.
105///
106/// # Why the old `scale · |pivot| · √ε · 1024` estimate is wrong
107///
108/// The batched/fused device paths derive the suggested bump from the
109/// factorization "pivot" — but cuSOLVER's `potrf` (and the NVRTC kernel's
110/// status code) report the failing pivot as a **1-based row index**, NOT the
111/// magnitude of the negative pivot. A block that is indefinite by `O(1)`
112/// (e.g. `H_tt = -I`, whose smallest eigenvalue is `-1`) then yields the same
113/// `bump ≈ √ε · 1024 ≈ 1.5e-5` as a block that is indefinite by `O(√ε)`. The
114/// outer LM escalation, which retries at `ridge_t + bump` and grows
115/// geometrically with a bounded step count, can never lift a strongly
116/// indefinite block out of the negative regime, so the solve fails to recover
117/// even though the block is trivially regularizable. (Surfaced by the V100
118/// `ridge_bump_required_on_non_pd_row_recovers_after_bump` validation test.)
119///
120/// # The bound
121///
122/// By the Gershgorin circle theorem every eigenvalue `λ` of the symmetric
123/// matrix `A = H_tt` satisfies, for some row `i`,
124///   `λ ≥ A[i,i] − Σ_{j≠i} |A[i,j]|`,
125/// so `λ_min(A) ≥ min_i ( A[i,i] − Σ_{j≠i} |A[i,j]| ) =: g` (the most negative
126/// Gershgorin left edge). Adding `t·I` shifts every eigenvalue up by `t`, so
127/// `A + t·I` is PD as soon as `t > -g`. We are already sitting at `ridge_t`, so
128/// the ADDITIONAL bump needed is `-(g + ridge_t)` when that is positive. We add
129/// a relative safety margin (`√ε · scale · 1024`, the same headroom the legacy
130/// estimate used) so the re-factored, rounding-perturbed block clears the pivot
131/// on the first retry, and a `max(1)`-scaled floor so a marginally-indefinite
132/// block still gets a strictly positive, non-vanishing bump.
133///
134/// The returned value is the bump to ADD to the current `ridge_t`. It is always
135/// strictly positive (the caller only constructs `RidgeBumpRequired` on an
136/// actual non-PD failure, but the bound is defensive regardless).
137#[must_use]
138fn ridge_bump_to_make_pd(htt: ArrayView2<'_, f64>, ridge_t: f64) -> f64 {
139    let d = htt.nrows();
140    // Diagonal magnitude scale (also the legacy `scale`), and the most-negative
141    // Gershgorin left edge `g = min_i (A_ii − Σ_{j≠i} |A_ij|)`.
142    let mut scale = 1.0_f64;
143    let mut min_gershgorin_edge = f64::INFINITY;
144    for i in 0..d {
145        let diag = htt[[i, i]];
146        scale = scale.max(diag.abs());
147        let mut off_sum = 0.0_f64;
148        for j in 0..d {
149            if j != i {
150                off_sum += htt[[i, j]].abs();
151            }
152        }
153        min_gershgorin_edge = min_gershgorin_edge.min(diag - off_sum);
154    }
155    if !min_gershgorin_edge.is_finite() {
156        // d == 0 (no rows) or non-finite entries: fall back to the scale-only
157        // floor so the caller still gets a strictly positive bump.
158        return scale * f64::EPSILON.sqrt() * RIDGE_BUMP_EPS_MARGIN;
159    }
160    // Additional shift needed so `λ_min(A) + ridge_t + bump > 0`, i.e.
161    // `bump > -(min_gershgorin_edge + ridge_t)`.
162    let deficit = -(min_gershgorin_edge + ridge_t);
163    let margin = scale * f64::EPSILON.sqrt() * RIDGE_BUMP_EPS_MARGIN;
164    // Lift past the deficit (when positive) plus a rounding margin; never below
165    // the scale-relative floor so a marginal block still moves.
166    deficit.max(0.0) + margin
167}
168
169/// [`ridge_bump_to_make_pd`] for a `d × d` symmetric block stored column-major
170/// in a flat slice with the current ridge ALREADY baked into the diagonal
171/// (the device packers emit `D = H_tt + ridge_t·I` this way). Because the shift
172/// is already present, the Gershgorin bound is taken at `ridge_t = 0` and the
173/// returned value is still the ADDITIONAL bump to add on top of the current
174/// ridge. Returns the scale-only floor when `block` is mis-sized.
175// The column-major bump is a helper for the device tile packers, which only
176// exist on the linux CUDA path (`mod cuda`, `#[cfg(target_os = "linux")]`). It
177// therefore lives where it is used: gate it to linux. Its parity unit test is
178// gated to linux to match (a `test` token in the cfg would trip the build.rs
179// `#[cfg(test)]`-on-a-src-item ban; including `test` to dodge the non-linux
180// dead_code lint is exactly the escape hatch that ban forbids).
181#[cfg(target_os = "linux")]
182#[must_use]
183fn ridge_bump_to_make_pd_colmajor(block: &[f64], d: usize) -> f64 {
184    if d == 0 || block.len() < d * d {
185        return f64::EPSILON.sqrt() * RIDGE_BUMP_EPS_MARGIN;
186    }
187    // Column-major: element (row r, col c) at block[c*d + r]. The matrix is
188    // symmetric, so reading by column gives the same Gershgorin edges as by row.
189    let mut scale = 1.0_f64;
190    let mut min_gershgorin_edge = f64::INFINITY;
191    for i in 0..d {
192        let diag = block[i * d + i];
193        scale = scale.max(diag.abs());
194        let mut off_sum = 0.0_f64;
195        for j in 0..d {
196            if j != i {
197                off_sum += block[j * d + i].abs();
198            }
199        }
200        min_gershgorin_edge = min_gershgorin_edge.min(diag - off_sum);
201    }
202    let margin = scale * f64::EPSILON.sqrt() * RIDGE_BUMP_EPS_MARGIN;
203    if !min_gershgorin_edge.is_finite() {
204        return margin;
205    }
206    (-min_gershgorin_edge).max(0.0) + margin
207}
208
209/// Canonically condition and factor a device Direct solve's small reduced
210/// border without downloading any per-row Arrow slabs.
211///
212/// Device Direct forms the same reduced equation as the host path,
213/// `S_beta delta_beta = r_beta`, but stores `S_beta` column-major for
214/// cuSOLVER. When the caller declares an
215/// [`ArrowBetaGaugeQuotient`](crate::arrow_schur::ArrowBetaGaugeQuotient), the
216/// mathematical system is instead
217///
218/// `P S_beta P + Q Q^T`, with right-hand side `P r_beta`.
219///
220/// Production SAE Direct solves additionally request a reduced-Schur spectral
221/// Tikhonov floor. The host Direct path owns both canonical contracts:
222///
223/// * the quotient operations in
224/// [`ArrowBetaGaugeQuotient::pin_reduced_schur`](crate::arrow_schur::ArrowBetaGaugeQuotient::pin_reduced_schur)
225/// and
226/// [`ArrowBetaGaugeQuotient::project_complement`](crate::arrow_schur::ArrowBetaGaugeQuotient::project_complement);
227/// * Jacobi/Van-der-Sluis equilibration, the PD/condition-number gate, and the
228/// optional spectral floor in
229/// [`solve_dense_reduced_system`](crate::arrow_schur::solve_dense_reduced_system).
230///
231/// Reuse those exact implementations rather than maintaining numerically
232/// distinct CUDA policies. The returned `bool` says that `schur_col_major` now
233/// contains the lower Cholesky factor, so the caller must skip device POTRF and
234/// may proceed directly to its device triangular solves. When neither contract
235/// is active the buffer is untouched and the caller retains its fast raw-POTRF
236/// path.
237///
238/// Only the small `k x k` reduced block crosses this seam. Per-row factors,
239/// whitened cross blocks, Schur assembly, and back-substitution remain device
240/// work.
241#[cfg(target_os = "linux")]
242pub(crate) fn canonicalize_device_beta_factor(
243    quotient: Option<&ArrowBetaGaugeQuotient>,
244    newton_schur_tikhonov_rel_floor: Option<f64>,
245    k: usize,
246    schur_col_major: &mut [f64],
247) -> Result<bool, ArrowSchurGpuFailure> {
248    if quotient.is_none() && newton_schur_tikhonov_rel_floor.is_none() {
249        return Ok(false);
250    }
251    if schur_col_major.len() != k * k {
252        return Err(ArrowSchurGpuFailure::SchurFactorFailed {
253            reason: format!(
254                "device canonical reduced factor received Schur length {} for k={k}",
255                schur_col_major.len(),
256            ),
257        });
258    }
259
260    let mut schur = Array2::<f64>::zeros((k, k));
261    for col in 0..k {
262        for row in 0..k {
263            schur[[row, col]] = schur_col_major[col * k + row];
264        }
265    }
266    let conditioned_input = match quotient {
267        Some(quotient) => quotient.pin_reduced_schur(schur.view()),
268        None => schur,
269    };
270    let mut canonical_options = ArrowSolveOptions::direct();
271    canonical_options.newton_schur_tikhonov_rel_floor =
272        newton_schur_tikhonov_rel_floor;
273    let zero_rhs = Array1::<f64>::zeros(k);
274    let (_, factor, _) =
275        solve_dense_reduced_system(&conditioned_input, &zero_rhs, &canonical_options, None)
276            .map_err(|error| ArrowSchurGpuFailure::SchurFactorFailed {
277                reason: format!("canonical reduced-border factorization failed: {error}"),
278            })?;
279    let factor = factor.ok_or_else(|| ArrowSchurGpuFailure::SchurFactorFailed {
280        reason: "canonical Direct reduced-border factorization returned no factor".to_string(),
281    })?;
282    for col in 0..k {
283        for row in 0..k {
284            schur_col_major[col * k + row] = factor[[row, col]];
285        }
286    }
287    Ok(true)
288}
289
290/// Project a device reduced vector onto the identifiable beta complement.
291///
292/// Both the reduced right-hand side and the solved step use this exact helper:
293/// `P r_beta` is the Faddeev--Popov equation's right-hand side, and projecting
294/// the solution removes the last finite-precision orbit component before
295/// row-local back-substitution.
296#[cfg(target_os = "linux")]
297fn project_device_beta_vector(
298    quotient: Option<&ArrowBetaGaugeQuotient>,
299    k: usize,
300    vector: Vec<f64>,
301) -> Result<Vec<f64>, ArrowSchurGpuFailure> {
302    if vector.len() != k {
303        return Err(ArrowSchurGpuFailure::SchurFactorFailed {
304            reason: format!(
305                "device beta vector length {} does not match border {}",
306                vector.len(),
307                k,
308            ),
309        });
310    }
311    Ok(match quotient {
312        Some(quotient) => quotient
313            .project_complement(Array1::from_vec(vector).view())
314            .to_vec(),
315        None => vector,
316    })
317}
318
319/// Entry point: attempt the fully device-resident Arrow-Schur Newton solve.
320/// Returns `Err(ArrowSchurGpuFailure::Unavailable)` to indicate "device path
321/// declined, fall back to CPU" — never panics.
322pub fn solve_arrow_newton_step(
323    sys: &ArrowSchurSystem,
324    ridge_t: f64,
325    ridge_beta: f64,
326    newton_schur_tikhonov_rel_floor: Option<f64>,
327) -> Result<ArrowSchurGpuSolution, ArrowSchurGpuFailure> {
328    let n = sys.rows.len();
329    let d = sys.d;
330    let k = sys.k;
331
332    // Detect matrix-free operators before any dim() checks so callers get a
333    // clear, actionable error instead of a generic SchurFactorFailed. The GPU
334    // dense-Schur path requires materialised H_ββ and per-row H_tβ slabs;
335    // CPU InexactPCG is the correct fallback when either operator is abstract.
336    let had_hbb_matvec = sys.hbb_matvec.is_some();
337    let had_htbeta_matvec = sys.htbeta_matvec.is_some();
338    if had_hbb_matvec || had_htbeta_matvec {
339        return Err(ArrowSchurGpuFailure::GpuRequiresDenseSystem {
340            had_hbb_matvec,
341            had_htbeta_matvec,
342        });
343    }
344
345    // A `penalty_op` is the AUTHORITATIVE β-curvature source whenever it is
346    // installed: assembly bypasses the dense `hbb` accumulator (and, for
347    // frames-engaged SAE systems, reclaims it to a 0×0 workspace), so whatever
348    // `hbb` survives is STALE relative to the operator. The dense device Schur
349    // path reads ONLY `hbb`, so it must decline here — BEFORE the shape gate
350    // below — even when a stale `(k, k)` `hbb` would pass that gate: proceeding
351    // would silently compute the WRONG Newton step from stale curvature instead
352    // of routing to the CPU matrix-free lane (which reads `penalty_op`) that
353    // returns the correct one. The production caller `try_device_arrow_direct`
354    // already short-circuits this shape, but enforcing it at the entry keeps
355    // EVERY caller covered — the resident / reupload harnesses and any future
356    // direct caller — so no path can reach the device solve with a stale dense
357    // block. Both matvec flags are false here: control only reaches this point
358    // when neither matrix-free operator is installed (returned above), and the
359    // frames-engaged path installs `penalty_op` with no `hbb_matvec` /
360    // `htbeta_matvec`.
361    if sys.penalty_op.is_some() {
362        return Err(ArrowSchurGpuFailure::GpuRequiresDenseSystem {
363            had_hbb_matvec: false,
364            had_htbeta_matvec: false,
365        });
366    }
367
368    if sys.hbb.dim() != (k, k) {
369        // The dense (K×K) H_ββ block is absent (e.g. an SAE-manifold system
370        // whose β-curvature is carried by a matrix-free `penalty_op` /
371        // factored-frame representation, with `hbb` reclaimed to a 0×0
372        // workspace at the end of assembly). This is a CAPABILITY decline, not
373        // a numerical failure: the dense device Schur path simply cannot
374        // consume this system, so the host must route it to the CPU lane
375        // (which reads the matrix-free operators) exactly as it does for the
376        // `hbb_matvec` / `htbeta_matvec` case above. Returning `SchurFactorFailed`
377        // here would masquerade as a non-PD/rank-deficient factorization and be
378        // escalated (and ultimately surfaced as a FATAL RemlConvergenceError)
379        // by the outer LM loop instead of falling back. Decline instead. Both
380        // matvec flags are false because the absence is structural (no dense
381        // block was materialized), not caused by an installed matrix-free op.
382        return Err(ArrowSchurGpuFailure::GpuRequiresDenseSystem {
383            had_hbb_matvec: false,
384            had_htbeta_matvec: false,
385        });
386    }
387    if n == 0 || d == 0 {
388        return Err(ArrowSchurGpuFailure::Unavailable);
389    }
390    if sys
391        .rows
392        .iter()
393        .any(|row| row.htt.dim() != (d, d) || row.htbeta.dim() != (d, k) || row.gt.len() != d)
394    {
395        return Err(ArrowSchurGpuFailure::SchurFactorFailed {
396            reason: "row block dimension mismatch".to_string(),
397        });
398    }
399
400    #[cfg(not(target_os = "linux"))]
401    {
402        // The floor is consumed only by the Linux device path, but it is still
403        // the caller's input on this target, so validate it rather than discard
404        // it: a host that merely lacks the device should reject a malformed
405        // request the same way, not accept it silently. (`drop` cannot serve as
406        // the "deliberately unused" marker here -- the floor is `Copy`, so
407        // `drop` is a no-op that `-D warnings` rejects, and the workspace
408        // ban-scanner forbids every underscore form.)
409        if let Some(floor) = newton_schur_tikhonov_rel_floor {
410            if !floor.is_finite() {
411                return Err(ArrowSchurGpuFailure::SchurFactorFailed {
412                    reason: format!(
413                        "newton_schur_tikhonov_rel_floor must be finite, got {floor}"
414                    ),
415                });
416            }
417        }
418        if ridge_t.is_nan() || ridge_beta.is_nan() {
419            return Err(ArrowSchurGpuFailure::SchurFactorFailed {
420                reason: "ridge is NaN".to_string(),
421            });
422        }
423        Err(ArrowSchurGpuFailure::Unavailable)
424    }
425
426    #[cfg(target_os = "linux")]
427    {
428        // Multi-GPU: the arrow-Schur solve is row-block separable in its forward
429        // (per-row factor / whiten / partial-Schur) and backward (per-row
430        // back-sub) phases — only the small shared K×K reduce+factor+δβ is
431        // central. When more than one device is usable, split the WHOLE solve at
432        // row-block granularity across all GPUs. The POTRF stays fused with its
433        // dependent TRSM+GEMM on each tile's own stream, so no on-stream solve is
434        // orphaned. On `Unavailable` (one device, shape below policy, transient)
435        // fall through to the single-device fused / Layer-A paths below.
436        if resolve_runtime_for_device_path()?
437            .map(gam_gpu::device_runtime::GpuRuntime::device_count)
438            .unwrap_or(0)
439            > 1
440        {
441            match cuda::solve_multi_gpu(
442                sys,
443                ridge_t,
444                ridge_beta,
445                newton_schur_tikhonov_rel_floor,
446            ) {
447                Ok(sol) => return Ok(sol),
448                Err(ArrowSchurGpuFailure::RidgeBumpRequired { row, bump }) => {
449                    return Err(ArrowSchurGpuFailure::RidgeBumpRequired { row, bump });
450                }
451                Err(ArrowSchurGpuFailure::SchurFactorFailed { reason }) => {
452                    return Err(ArrowSchurGpuFailure::SchurFactorFailed { reason });
453                }
454                // Unavailable / GpuRequiresDenseSystem: fall through to the
455                // single-device paths (already shape-validated above).
456                // Named rather than wildcarded so a new failure variant has to
457                // declare whether it is fall-through or must surface.
458                Err(
459                    ArrowSchurGpuFailure::Unavailable
460                    | ArrowSchurGpuFailure::GpuRequiresDenseSystem { .. },
461                ) => {}
462            }
463        }
464        // Layer D admission: when the system shape passes the
465        // (Σ p³ ≥ 1e5 OR R ≥ 16) heuristic and `p ≤ MAX_FUSED_P`, the fused
466        // NVRTC kernel replaces the cuSOLVER/cuBLAS Layer A+B+C path with a
467        // single per-row block. Layer C↔D parity (math block 3 §16 test 6)
468        // requires both paths to agree to 1e-10 on identical inputs.
469        if crate::gpu_kernels::arrow_schur_nvrtc::system_admits_fused_path(sys) {
470            match cuda::solve_fused(
471                sys,
472                ridge_t,
473                ridge_beta,
474                newton_schur_tikhonov_rel_floor,
475            ) {
476                Ok(sol) => return Ok(sol),
477                // RidgeBumpRequired must surface to the outer escalation loop —
478                // the fused path's pivot diagnostic is identical in semantics
479                // to the cuSOLVER batched POTRF info code.
480                Err(ArrowSchurGpuFailure::RidgeBumpRequired { row, bump }) => {
481                    return Err(ArrowSchurGpuFailure::RidgeBumpRequired { row, bump });
482                }
483                // Unavailable / SchurFactorFailed / GpuRequiresDenseSystem fall
484                // through to the unfused path so a flaky NVRTC compile or
485                // shared-mem allocation does not abort the outer Newton step.
486                // Named rather than wildcarded so a new failure variant has to
487                // declare whether it is fall-through or must surface.
488                Err(
489                    ArrowSchurGpuFailure::Unavailable
490                    | ArrowSchurGpuFailure::SchurFactorFailed { .. }
491                    | ArrowSchurGpuFailure::GpuRequiresDenseSystem { .. },
492                ) => {}
493            }
494        }
495        cuda::solve(
496            sys,
497            ridge_t,
498            ridge_beta,
499            newton_schur_tikhonov_rel_floor,
500        )
501    }
502}
503
504/// Build the stacked column-major D buffer (n local d×d blocks), the stacked
505/// stacked B buffer (n local d×k blocks), and the stacked g buffer
506/// (n local d-vectors) consumed by the device pipeline. Each block is laid
507/// out column-major so a single allocation + `cuMemcpyHtoD` reaches the
508/// device without per-row dispatch overhead.
509#[cfg(target_os = "linux")]
510fn pack_host(sys: &ArrowSchurSystem, ridge_t: f64) -> (Vec<f64>, Vec<f64>, Vec<f64>) {
511    let n = sys.rows.len();
512    let d = sys.d;
513    let k = sys.k;
514    let mut d_buf = Vec::with_capacity(n * d * d);
515    let mut b_buf = Vec::with_capacity(n * d * k);
516    let mut g_buf = Vec::with_capacity(n * d);
517    for row in &sys.rows {
518        pack_block(row, ridge_t, d, k, &mut d_buf, &mut b_buf, &mut g_buf);
519    }
520    (d_buf, b_buf, g_buf)
521}
522
523/// Pack the per-row `D = H_tt + ρ_t I` blocks (block-contiguous `d×d`
524/// column-major, as the batched POTRF/TRSM require) together with the cross
525/// blocks `B` as ONE stacked `(n·d) × k` column-major matrix with leading
526/// dimension `n·d`: row `i·d + r` of the stack is row `r` of block `i`.
527///
528/// That single-matrix view is what makes the whole Schur reduction one GEMM
529/// (`schur_gemm_stacked`) instead of n rank-`d` updates, and both per-iterate
530/// accumulations one GEMV each — the per-op granularity #2393 identifies as the
531/// reason the LLM shape trips no dispatch gate despite being large in aggregate.
532#[cfg(target_os = "linux")]
533fn pack_host_d_and_stacked_b(sys: &ArrowSchurSystem, ridge_t: f64) -> (Vec<f64>, Vec<f64>) {
534    let n = sys.rows.len();
535    let d = sys.d;
536    let k = sys.k;
537    let rows = n * d;
538    let mut d_buf = Vec::with_capacity(n * d * d);
539    let mut b_buf = vec![0.0_f64; rows * k];
540    for (i, row) in sys.rows.iter().enumerate() {
541        for col in 0..d {
542            for r in 0..d {
543                let mut value = row.htt[[r, col]];
544                if r == col {
545                    value += ridge_t;
546                }
547                d_buf.push(value);
548            }
549        }
550        for col in 0..k {
551            let base = col * rows + i * d;
552            for r in 0..d {
553                b_buf[base + r] = row.htbeta[[r, col]];
554            }
555        }
556    }
557    (d_buf, b_buf)
558}
559
560#[cfg(target_os = "linux")]
561#[inline]
562fn pack_block(
563    row: &crate::arrow_schur::ArrowRowBlock,
564    ridge_t: f64,
565    d: usize,
566    k: usize,
567    d_buf: &mut Vec<f64>,
568    b_buf: &mut Vec<f64>,
569    g_buf: &mut Vec<f64>,
570) {
571    for col in 0..d {
572        for r in 0..d {
573            let mut value = row.htt[[r, col]];
574            if r == col {
575                value += ridge_t;
576            }
577            d_buf.push(value);
578        }
579    }
580    for col in 0..k {
581        for r in 0..d {
582            b_buf.push(row.htbeta[[r, col]]);
583        }
584    }
585    for r in 0..d {
586        g_buf.push(row.gt[r]);
587    }
588}
589
590/// Entry that forces the Layer D + E fused NVRTC path regardless of the
591/// admission heuristic. Used by the V100 Layer C↔D parity harness to drive
592/// the fused kernel at small shapes the heuristic would otherwise route through
593/// the cuSOLVER/cuBLAS Layer A+B+C path. The symbol only exists on the target
594/// that provides the NVRTC implementation.
595#[cfg(target_os = "linux")]
596pub fn solve_arrow_newton_step_fused_force(
597    sys: &ArrowSchurSystem,
598    ridge_t: f64,
599    ridge_beta: f64,
600    newton_schur_tikhonov_rel_floor: Option<f64>,
601) -> Result<ArrowSchurGpuSolution, ArrowSchurGpuFailure> {
602    if ridge_t.is_nan() || ridge_beta.is_nan() {
603        return Err(ArrowSchurGpuFailure::SchurFactorFailed {
604            reason: "ridge is NaN".to_string(),
605        });
606    }
607    if crate::gpu_kernels::arrow_schur_nvrtc::plan_fused_launch(sys.rows.len(), sys.d, sys.k)
608        .is_none()
609    {
610        return Err(ArrowSchurGpuFailure::Unavailable);
611    }
612    cuda::solve_fused(
613        sys,
614        ridge_t,
615        ridge_beta,
616        newton_schur_tikhonov_rel_floor,
617    )
618}
619
620/// #1017 Phase 3: a device-resident Arrow-Schur frame whose constant Hessian
621/// blocks (`D = H_tt`, `B = H_tβ`, border `H_ββ`) and their factors stay on the
622/// device across the inner Newton loop. Construct once per frozen gate/basis
623/// frame, then call [`ResidentArrowFrameHandle::solve_gradient`] once per
624/// iterate with the fresh residual gradient — only the `O(n·d + p)` gradient
625/// crosses to the device and only `δ` crosses back, in contrast to
626/// [`solve_arrow_newton_step`] which re-uploads and re-factors the full system
627/// every call. On a non-CUDA host construction returns
628/// `ArrowSchurGpuFailure::Unavailable`.
629#[cfg(target_os = "linux")]
630pub struct ResidentArrowFrameHandle {
631    inner: cuda::ResidentArrowFrame,
632}
633
634/// The resident CUDA frame has no value on hosts that cannot construct it.
635/// Keeping this type uninhabited preserves the fail-loud platform contract
636/// without exposing a fake non-CUDA implementation.
637#[cfg(not(target_os = "linux"))]
638pub enum ResidentArrowFrameHandle {}
639
640impl ResidentArrowFrameHandle {
641    /// Upload the constant Hessian blocks and perform the one-time factor work.
642    pub fn new(
643        sys: &ArrowSchurSystem,
644        ridge_t: f64,
645        ridge_beta: f64,
646        newton_schur_tikhonov_rel_floor: Option<f64>,
647    ) -> Result<Self, ArrowSchurGpuFailure> {
648        // The dense device path requires materialised blocks, same admission as
649        // `solve_arrow_newton_step`.
650        if sys.hbb_matvec.is_some() || sys.htbeta_matvec.is_some() {
651            return Err(ArrowSchurGpuFailure::GpuRequiresDenseSystem {
652                had_hbb_matvec: sys.hbb_matvec.is_some(),
653                had_htbeta_matvec: sys.htbeta_matvec.is_some(),
654            });
655        }
656        #[cfg(not(target_os = "linux"))]
657        {
658            // Validated rather than discarded on this target -- see the note in
659            // `solve_arrow_newton_step`.
660            if let Some(floor) = newton_schur_tikhonov_rel_floor {
661                if !floor.is_finite() {
662                    return Err(ArrowSchurGpuFailure::SchurFactorFailed {
663                        reason: format!(
664                            "newton_schur_tikhonov_rel_floor must be finite, got {floor}"
665                        ),
666                    });
667                }
668            }
669            if ridge_t.is_nan() || ridge_beta.is_nan() {
670                return Err(ArrowSchurGpuFailure::SchurFactorFailed {
671                    reason: "ridge is NaN".to_string(),
672                });
673            }
674            Err(ArrowSchurGpuFailure::Unavailable)
675        }
676        #[cfg(target_os = "linux")]
677        {
678            Ok(Self {
679                inner: cuda::ResidentArrowFrame::new(
680                    sys,
681                    ridge_t,
682                    ridge_beta,
683                    newton_schur_tikhonov_rel_floor,
684                )?,
685            })
686        }
687    }
688
689    /// Solve `H δ = −gradient` for a fresh gradient reusing the resident factors.
690    pub fn solve_gradient(
691        &self,
692        g_t: &[f64],
693        g_beta: &[f64],
694    ) -> Result<ArrowSchurGpuSolution, ArrowSchurGpuFailure> {
695        #[cfg(not(target_os = "linux"))]
696        {
697            if g_t.iter().chain(g_beta).any(|v| !v.is_finite()) {
698                return Err(ArrowSchurGpuFailure::SchurFactorFailed {
699                    reason: "non-finite gradient entry".to_string(),
700                });
701            }
702            Err(ArrowSchurGpuFailure::Unavailable)
703        }
704        #[cfg(target_os = "linux")]
705        {
706            self.inner.solve_gradient(g_t, g_beta)
707        }
708    }
709
710    /// `log|H|` for the frame (constant; depends only on the factored Hessian).
711    #[must_use]
712    pub fn log_det_hessian(&self) -> f64 {
713        #[cfg(not(target_os = "linux"))]
714        {
715            match *self {}
716        }
717        #[cfg(target_os = "linux")]
718        {
719            self.inner.log_det_hessian()
720        }
721    }
722}
723
724/// #1017: a BASE-block-resident Arrow-Schur frame for the LM ridge ladder.
725///
726/// Unlike [`ResidentArrowFrameHandle`] — which BAKES one ridge into its factors
727/// and then serves cheap re-solves for a NEW GRADIENT at that SAME ridge — this
728/// frame holds the ridge-INDEPENDENT base blocks (`D = H_tt`, `B = H_tβ`, border
729/// `H_ββ`, gradient) resident and RE-FACTORS on-device at each requested
730/// `(ridge_t, ridge_beta)`. That is the regime `solve_with_lm_escalation_inner`
731/// actually runs: its trials re-solve the SAME system (same gradient) at
732/// ESCALATING ridges, so the factor changes every trial but the base blocks do
733/// not. The base blocks upload ONCE; each trial pays only a device-to-device
734/// copy of the base blocks into scratch, an on-device diagonal ridge add, and the
735/// factor/solve — in place of the full `O(n·d·k)` host→device re-upload that
736/// [`solve_arrow_newton_step`] performs every trial. The per-trial numerics are
737/// bit-identical to that re-upload path (same POTRF/TRSM/Schur/back-sub order).
738#[cfg(target_os = "linux")]
739pub struct ResidentBaseArrowFrameHandle {
740    inner: cuda::ResidentBaseArrowFrame,
741}
742
743/// The base-resident CUDA frame is unavailable, rather than emulated, on a
744/// non-CUDA host.
745#[cfg(not(target_os = "linux"))]
746pub enum ResidentBaseArrowFrameHandle {}
747
748impl ResidentBaseArrowFrameHandle {
749    /// Upload the ridge-independent base blocks once. No factorization runs here;
750    /// each [`Self::refactor_and_solve`] performs the ridge-dependent factor+solve.
751    /// The dense device path requires materialised blocks, so a matrix-free
752    /// `H_ββ` / `H_tβ` operator is rejected (same admission as
753    /// [`solve_arrow_newton_step`]).
754    pub fn new(
755        sys: &ArrowSchurSystem,
756        newton_schur_tikhonov_rel_floor: Option<f64>,
757    ) -> Result<Self, ArrowSchurGpuFailure> {
758        if sys.hbb_matvec.is_some() || sys.htbeta_matvec.is_some() {
759            return Err(ArrowSchurGpuFailure::GpuRequiresDenseSystem {
760                had_hbb_matvec: sys.hbb_matvec.is_some(),
761                had_htbeta_matvec: sys.htbeta_matvec.is_some(),
762            });
763        }
764        #[cfg(not(target_os = "linux"))]
765        {
766            // Validated rather than discarded on this target -- see the note in
767            // `solve_arrow_newton_step`.
768            if let Some(floor) = newton_schur_tikhonov_rel_floor {
769                if !floor.is_finite() {
770                    return Err(ArrowSchurGpuFailure::SchurFactorFailed {
771                        reason: format!(
772                            "newton_schur_tikhonov_rel_floor must be finite, got {floor}"
773                        ),
774                    });
775                }
776            }
777            Err(ArrowSchurGpuFailure::Unavailable)
778        }
779        #[cfg(target_os = "linux")]
780        {
781            Ok(Self {
782                inner: cuda::ResidentBaseArrowFrame::new(
783                    sys,
784                    newton_schur_tikhonov_rel_floor,
785                )?,
786            })
787        }
788    }
789
790    /// Factor the resident base blocks at `(ridge_t, ridge_beta)` and solve
791    /// `(H + ridge)·δ = −gradient`. Only the two ridge scalars and the tiny
792    /// re-diagonalised `D` cross to the device; only `δ` crosses back. A non-PD
793    /// per-row block surfaces as [`ArrowSchurGpuFailure::RidgeBumpRequired`] so
794    /// the LM escalation bumps and retries at the larger ridge exactly as the
795    /// re-upload path does.
796    pub fn refactor_and_solve(
797        &self,
798        ridge_t: f64,
799        ridge_beta: f64,
800    ) -> Result<ArrowSchurGpuSolution, ArrowSchurGpuFailure> {
801        #[cfg(not(target_os = "linux"))]
802        {
803            if ridge_t.is_nan() || ridge_beta.is_nan() {
804                return Err(ArrowSchurGpuFailure::SchurFactorFailed {
805                    reason: "ridge is NaN".to_string(),
806                });
807            }
808            Err(ArrowSchurGpuFailure::Unavailable)
809        }
810        #[cfg(target_os = "linux")]
811        {
812            self.inner.refactor_and_solve(ridge_t, ridge_beta)
813        }
814    }
815
816    /// Run only the ridge-dependent FACTOR work and hand the factors back, so a
817    /// caller whose ridge does not move on the next iterate can re-solve without
818    /// re-factoring (#2539).
819    ///
820    /// `Self::refactor_and_solve_with_gradient` is exactly this followed by
821    /// [`Self::solve_with_factors`]. An inner Newton needs them separately: its
822    /// ridge changes only on an LM accept/reject, while its gradient changes
823    /// every iterate, so factoring per iterate would pay a POTRF/TRSM/Schur
824    /// chain for a factor that did not move — the accepted-step cost the
825    /// ridge-keyed frame used to avoid.
826    pub fn factor_at(
827        &self,
828        ridge_t: f64,
829        ridge_beta: f64,
830    ) -> Result<ResidentBaseRidgeFactorsHandle, ArrowSchurGpuFailure> {
831        #[cfg(not(target_os = "linux"))]
832        {
833            if ridge_t.is_nan() || ridge_beta.is_nan() {
834                return Err(ArrowSchurGpuFailure::SchurFactorFailed {
835                    reason: "ridge is NaN".to_string(),
836                });
837            }
838            Err(ArrowSchurGpuFailure::Unavailable)
839        }
840        #[cfg(target_os = "linux")]
841        {
842            Ok(ResidentBaseRidgeFactorsHandle {
843                inner: self.inner.factor_at(ridge_t, ridge_beta)?,
844            })
845        }
846    }
847
848    /// Solve `(H + ridge)·δ = −gradient` against factors from
849    /// [`Self::factor_at`]. No POTRF, no `B` whitening: only the gradient
850    /// (`n·d + k` doubles) crosses to the device and only `δ` crosses back.
851    pub fn solve_with_factors(
852        &self,
853        factors: &ResidentBaseRidgeFactorsHandle,
854        g_t: &[f64],
855        g_beta: &[f64],
856    ) -> Result<ArrowSchurGpuSolution, ArrowSchurGpuFailure> {
857        #[cfg(not(target_os = "linux"))]
858        {
859            if g_t.iter().chain(g_beta).any(|v| !v.is_finite()) {
860                return Err(ArrowSchurGpuFailure::SchurFactorFailed {
861                    reason: "gradient entry is not finite".to_string(),
862                });
863            }
864            match *factors {}
865        }
866        #[cfg(target_os = "linux")]
867        {
868            self.inner
869                .solve_with_gradient_checked(&factors.inner, g_t, g_beta)
870        }
871    }
872}
873
874/// Ridge-derived factors of a [`ResidentBaseArrowFrameHandle`], held on the
875/// device between solves (#2539).
876///
877/// The base frame owns the ridge-INDEPENDENT blocks; this owns everything the
878/// ridge moves. A caller keeps one of these alongside the frame and asks
879/// [`Self::matches_ridge`] before re-solving: a hit is a gradient-only solve, a
880/// miss is a [`ResidentBaseArrowFrameHandle::factor_at`] whose inputs never
881/// leave the device.
882#[cfg(target_os = "linux")]
883pub struct ResidentBaseRidgeFactorsHandle {
884    inner: cuda::BaseRidgeFactors,
885}
886
887/// The base-resident CUDA factors are unavailable, rather than emulated, on a
888/// non-CUDA host.
889#[cfg(not(target_os = "linux"))]
890pub enum ResidentBaseRidgeFactorsHandle {}
891
892impl ResidentBaseRidgeFactorsHandle {
893    /// Whether these factors were built at exactly `(ridge_t, ridge_beta)`.
894    ///
895    /// The key lives with the factors rather than beside them so a cache cannot
896    /// hold a key that has drifted from the factor state it names.
897    #[must_use]
898    pub fn matches_ridge(&self, ridge_t: f64, ridge_beta: f64) -> bool {
899        #[cfg(not(target_os = "linux"))]
900        {
901            // A NaN ridge never keys anything, on either host.
902            if ridge_t.is_nan() || ridge_beta.is_nan() {
903                return false;
904            }
905            match *self {}
906        }
907        #[cfg(target_os = "linux")]
908        {
909            self.inner.matches_ridge(ridge_t, ridge_beta)
910        }
911    }
912}
913
914/// Build a GPU-backed Schur matvec closure for CPU-driven PCG at K ≥ 5000.
915///
916/// Runs the fused NVRTC forward kernel once on the dense per-row `H_tβ` slabs
917/// to compute `Y_i = L_i^{-1} H_tβ^(i)` for all rows, persists the `Y_i`
918/// factors in a host-side buffer, and returns an `Arc<dyn Fn(...)>` closure
919/// that computes the full Schur matvec
920///
921/// ```text
922/// S·x = (H_ββ + ridge_beta·I)·x  −  Σ_i Y_i^T (Y_i·x)
923/// ```
924///
925/// each time it is called. At K ≥ 5000 the `Σ_i Y_i^T (Y_i·x)` term
926/// dominates over the host↔device transfer of the K-vector `x`, so the GPU
927/// path is a clear win even with per-iteration transfer.
928///
929/// `H_ββ·x` is evaluated on the CPU using `sys.hbb_matvec` when present (the
930/// matrix-free hook for SAE-manifold scale callers) or the dense `sys.hbb`
931/// block otherwise. The `Y_i` term uses cuBLAS batched GEMV device-side; only
932/// `x` (K doubles) and `out` (K doubles) cross the host↔device boundary per
933/// PCG iteration.
934///
935/// Returns `Err(ArrowSchurGpuFailure::Unavailable)` if CUDA is unavailable or
936/// the system shape is outside the fused kernel's admission range (e.g.
937/// `d > MAX_FUSED_P = 32` or no CUDA context). Callers should fall back to CPU
938/// `InexactPCG` on `Unavailable`.
939///
940/// Returns `Err(ArrowSchurGpuFailure::RidgeBumpRequired)` if a per-row Cholesky
941/// factor failed at the requested `ridge_t`; the outer LM escalation should
942/// bump `ridge_t` and retry.
943///
944/// # Composition with the matrix-free SAE Kronecker operator
945///
946/// When `sys.htbeta_matvec` is set (matrix-free `H_tβ` Kronecker operator),
947/// the dense `H_tβ` slabs are absent — the dense forward kernel above cannot
948/// run, and at `K = 100K` the dense `Y_i = L_i^{-1} H_tβ^(i)` (`d × K` per row)
949/// could not be materialised anyway. Instead, `build_row_procedural_matvec`
950/// returns a row-procedural Schur matvec: per row it gathers
951/// `v_i = H_tβ^(i)·x` through the forward operator (sparse `O(m_i · p)`),
952/// solves `(H_tt^(i) + ρ_t·I)^{-1} v_i` through a pre-computed per-row Cholesky
953/// factor, and scatters `H_βt^(i)·w_i` through the sparse transpose operator
954/// (`O(m_i · p)`, replacing the old `O(K)` column-probe). This is the
955/// row-procedural `a_ik · Φ_k[i,m]` Kronecker apply over the active atoms only.
956pub fn gpu_schur_matvec_backend(
957    sys: &ArrowSchurSystem,
958    ridge_t: f64,
959    ridge_beta: f64,
960) -> Result<crate::arrow_schur::GpuSchurMatvec, ArrowSchurGpuFailure> {
961    // Matrix-free H_tβ operator present: drive the row-procedural sparse
962    // Kronecker apply (active atoms only) instead of the dense forward kernel.
963    if sys.htbeta_matvec.is_some() {
964        return build_row_procedural_matvec(sys, ridge_t, ridge_beta);
965    }
966
967    #[cfg(not(target_os = "linux"))]
968    {
969        // No CUDA runtime on non-Linux. NaN ridges are validated to ensure the
970        // same contract as the Linux path.
971        if ridge_t.is_nan() || ridge_beta.is_nan() {
972            return Err(ArrowSchurGpuFailure::Unavailable);
973        }
974        Err(ArrowSchurGpuFailure::Unavailable)
975    }
976
977    #[cfg(target_os = "linux")]
978    {
979        cuda::build_schur_matvec_backend(sys, ridge_t, ridge_beta)
980    }
981}
982
983/// #1017 evidence lane: a device-resident, RUN-TO-RUN DETERMINISTIC framed
984/// reduced-Schur `S·v` for the SLQ/surrogate `log|S|` matvec, or `None` when the
985/// device/shape declines. CPU and non-Linux always return `None`, so the evidence
986/// matvec is byte-identical there. Unlike `gpu_schur_matvec_backend` (whose
987/// matrix-free branch returns the CPU row-procedural closure), this engages the
988/// resident device apply — but only via an atomics-free reduction, so it upholds
989/// `slq_reduced_schur_log_det`'s determinism contract.
990pub fn build_framed_resident_evidence_matvec(
991    sys: &ArrowSchurSystem,
992    ridge_t: f64,
993    ridge_beta: f64,
994    apply_budget: usize,
995) -> Result<Option<crate::arrow_schur::GpuSchurMatvec>, ArrowSchurGpuFailure> {
996    #[cfg(not(target_os = "linux"))]
997    {
998        // No CUDA runtime exists off Linux. Refuse the same degenerate
999        // requests the Linux path would (mirroring the stubs above), then
1000        // report that no resident evidence matvec is available.
1001        if sys.k == 0 || !ridge_t.is_finite() || !ridge_beta.is_finite() || apply_budget == 0 {
1002            return Err(ArrowSchurGpuFailure::SchurFactorFailed {
1003                reason: "resident evidence matvec received degenerate dimensions or ridge"
1004                    .to_string(),
1005            });
1006        }
1007        Ok(None)
1008    }
1009    #[cfg(target_os = "linux")]
1010    {
1011        cuda::build_framed_resident_evidence_matvec(sys, ridge_t, ridge_beta, apply_budget)
1012    }
1013}
1014
1015/// Build a row-procedural reduced-Schur matvec for matrix-free SAE Kronecker
1016/// systems, eliminating the per-row latent block via cached per-row Cholesky
1017/// factors and applying the cross-block through the sparse forward/transpose
1018/// Kronecker operators (active atoms only).
1019///
1020/// The returned closure evaluates
1021/// `S·x = (H_ββ + ρ_β·I)·x − Σ_i H_βt^(i) (H_tt^(i) + ρ_t·I)^{-1} H_tβ^(i)·x`,
1022/// the same reduced Schur complement the dense path forms, but never
1023/// materialises the `d × K` cross-block `H_tβ^(i)`: the forward operator
1024/// (`out = H_tβ^(i)·x`) and transpose operator (`out += H_βt^(i)·v`) are the
1025/// sparse Kronecker gather/scatter from `SaeKroneckerRows`. The per-row factor
1026/// of `H_tt^(i) + ρ_t·I` is computed once when the closure is built and reused
1027/// across every CG iteration.
1028///
1029/// Returns `RidgeBumpRequired` if a per-row block is not positive definite at
1030/// the requested `ridge_t`; the outer LM escalation bumps `ridge_t` and retries.
1031fn build_row_procedural_matvec(
1032    sys: &ArrowSchurSystem,
1033    ridge_t: f64,
1034    ridge_beta: f64,
1035) -> Result<crate::arrow_schur::GpuSchurMatvec, ArrowSchurGpuFailure> {
1036    use std::sync::Arc;
1037    let n = sys.rows.len();
1038    let k = sys.k;
1039    let forward = sys
1040        .htbeta_matvec
1041        .clone()
1042        .ok_or(ArrowSchurGpuFailure::Unavailable)?;
1043    let transpose = sys.htbeta_transpose_matvec.clone().ok_or_else(|| {
1044        // A forward operator without its sparse adjoint cannot be applied
1045        // row-procedurally; this is a wiring error, surfaced as a Schur failure
1046        // so the caller routes to the dense CPU path rather than misreporting a
1047        // numerical bump.
1048        ArrowSchurGpuFailure::SchurFactorFailed {
1049            reason: "row-procedural Schur matvec requires htbeta_transpose_matvec; \
1050                     forward operator installed without its sparse adjoint"
1051                .to_string(),
1052        }
1053    })?;
1054
1055    // Pre-factor each per-row block H_tt^(i) + ρ_t·I = L_i L_iᵀ on the host.
1056    // The blocks are tiny (d_i ≲ 32) and the dense cross-block slabs are
1057    // absent, so there is no device forward-kernel work to amortise here; the
1058    // GPU win is the reduced K-system solve in `solve_reduced_beta_pcg`.
1059    let mut factors: Vec<Array2<f64>> = Vec::with_capacity(n);
1060    for (i, row) in sys.rows.iter().enumerate() {
1061        let di = row.htt.nrows();
1062        if row.htt.ncols() != di || row.gt.len() != di {
1063            return Err(ArrowSchurGpuFailure::SchurFactorFailed {
1064                reason: format!("row {i}: malformed H_tt block {:?}", row.htt.dim()),
1065            });
1066        }
1067        let mut block = row.htt.clone();
1068        for r in 0..di {
1069            block[[r, r]] += ridge_t;
1070        }
1071        let factor = cholesky_factor_in_place(block.view(), CholeskyGuard::NonnegativePivot)
1072            .ok_or_else(|| {
1073                // Deficit-aware bump from the block's own entries (Gershgorin),
1074                // so the outer LM escalation lifts a strongly-indefinite block
1075                // out of the negative regime in one retry.
1076                ArrowSchurGpuFailure::RidgeBumpRequired {
1077                    row: i,
1078                    bump: ridge_bump_to_make_pd(row.htt.view(), ridge_t),
1079                }
1080            })?;
1081        factors.push(factor);
1082    }
1083
1084    // The SAE-manifold β-Hessian lives in the structured penalty operator
1085    // (data-fit Gauss-Newton `G ⊗ I_p` + smoothness Kronecker blocks + any
1086    // dense analytic-β residual), NOT in the dense `hbb` accumulator — for
1087    // matrix-free systems `hbb` is zero/absent. Capture the effective penalty
1088    // operator so `H_ββ·x` matches the CPU `schur_matvec` path exactly. The
1089    // operator's `matvec` adds (`y += P x`), so seed `out` from the ridge term.
1090    let penalty_op = sys.effective_penalty_op();
1091    let row_dims: Vec<usize> = sys.rows.iter().map(|row| row.htt.nrows()).collect();
1092
1093    let closure: crate::arrow_schur::GpuSchurMatvec =
1094        Arc::new(move |x: &Array1<f64>, out: &mut Array1<f64>| {
1095            assert_eq!(x.len(), k, "row-procedural matvec: x.len() != k");
1096            assert_eq!(out.len(), k, "row-procedural matvec: out.len() != k");
1097
1098            // (H_ββ + ρ_β·I)·x into out. Seed with the ridge term, then add the
1099            // structured penalty-side product (penalty_op.matvec is additive).
1100            {
1101                let x_slice = x.as_slice().expect("x must be contiguous");
1102                let out_slice = out.as_slice_mut().expect("out must be contiguous");
1103                for a in 0..k {
1104                    out_slice[a] = ridge_beta * x_slice[a];
1105                }
1106                penalty_op.matvec(x_slice, out_slice);
1107            }
1108
1109            // out -= Σ_i H_βt^(i) (H_tt^(i) + ρ_t·I)^{-1} H_tβ^(i)·x.
1110            //
1111            // #1017: this row-procedural reduced-Schur term is the matrix-free
1112            // SAE path's matvec hot loop (`build_row_procedural_matvec` is the
1113            // host backend `gpu_schur_matvec_backend` returns when the dense
1114            // `H_tβ` slabs are absent — the production Qwen shape). At
1115            // (n≈2000 rows) it ran SERIALLY on one core and allocated a fresh
1116            // length-`K` `neg` plus per-row `v_i`/`w_i` on EVERY CG iteration —
1117            // tens of thousands of tiny heap allocations across a solve. Each
1118            // row contributes an independent length-`K` scatter, so the sum is
1119            // embarrassingly parallel; fan it across rayon over fixed row chunks
1120            // and fold the per-chunk length-`K` partials in chunk order so the
1121            // f64 reduction is deterministic (bit-identical run-to-run)
1122            // regardless of thread scheduling — it agrees with the serial sum up
1123            // to ULP-scale chunk reassociation (the #1017 verification gate).
1124            // Because that reassociation is a real (if tiny) departure from
1125            // serial, the criterion ranking across topology candidates is stable
1126            // except for candidates separated by less than the reassociation
1127            // margin, where the near-tie winner can flip — not an exact no-move
1128            // guarantee (#1211). Stay
1129            // sequential below
1130            // `SCHUR_MATVEC_PARALLEL_ROW_MIN` rows and when already inside a
1131            // rayon worker (the topology race fans candidates with
1132            // `run_topology_race_parallel`) — the same nested-rayon guard the
1133            // CPU `schur_matvec` uses. Buffers (`v_i`, `neg`) are reused across
1134            // rows within a chunk, so the per-row allocation churn is gone.
1135            let parallel = n >= crate::arrow_schur::SCHUR_MATVEC_PARALLEL_ROW_MIN
1136                && rayon::current_thread_index().is_none();
1137            if parallel {
1138                use rayon::prelude::*;
1139                const CHUNK: usize = 64;
1140                let partials: Vec<Array1<f64>> = (0..n)
1141                    .into_par_iter()
1142                    .chunks(CHUNK)
1143                    .map(|idxs| {
1144                        // One length-`K` scatter accumulator per chunk; the
1145                        // per-row latent vector `v_i` (length `d_i ≲ 32`) is the
1146                        // only per-row buffer, sized to the row's own `d_i`.
1147                        let mut neg = Array1::<f64>::zeros(k);
1148                        for i in idxs {
1149                            let di = row_dims[i];
1150                            // v_i = H_tβ^(i)·x (sparse Kronecker gather).
1151                            let mut v_i = Array1::<f64>::zeros(di);
1152                            forward(i, x.view(), &mut v_i);
1153                            // w_i = (H_tt^(i) + ρ_t·I)^{-1} v_i via L_i L_iᵀ.
1154                            let w_i = cholesky_solve_vector(factors[i].view(), v_i.view());
1155                            // neg += H_βt^(i)·w_i (sparse scatter).
1156                            transpose(i, w_i.view(), &mut neg);
1157                        }
1158                        neg
1159                    })
1160                    .collect();
1161                // #1017/#1175 floating-point parity contract: Rayon may
1162                // schedule chunks on any worker, but `.chunks(CHUNK).collect()`
1163                // returns partials in chunk-index order. Each chunk's row sum
1164                // is formed locally in increasing row order, then chunk
1165                // partials are folded left-to-right below. That makes the
1166                // parallel row-procedural Schur term deterministic for a fixed
1167                // input and chunking (no scheduling-dependent gather/scatter
1168                // reordering), but it is not required to be bit-identical to
1169                // the serial path because additions are reassociated at chunk
1170                // boundaries. CPU/GPU validation should therefore allow
1171                // ULP-scale drift while expecting stable run-to-run results.
1172                let mut neg = Array1::<f64>::zeros(k);
1173                for part in &partials {
1174                    for a in 0..k {
1175                        neg[a] += part[a];
1176                    }
1177                }
1178                for a in 0..k {
1179                    out[a] -= neg[a];
1180                }
1181            } else {
1182                // Serial path: reuse one `neg` and one `v_i` across rows.
1183                let mut neg = Array1::<f64>::zeros(k);
1184                for i in 0..n {
1185                    let di = row_dims[i];
1186                    // v_i = H_tβ^(i)·x (sparse Kronecker gather, length d_i).
1187                    let mut v_i = Array1::<f64>::zeros(di);
1188                    forward(i, x.view(), &mut v_i);
1189                    // w_i = (H_tt^(i) + ρ_t·I)^{-1} v_i via L_i L_iᵀ.
1190                    let w_i = cholesky_solve_vector(factors[i].view(), v_i.view());
1191                    // neg += H_βt^(i)·w_i (sparse scatter); subtract once at end.
1192                    transpose(i, w_i.view(), &mut neg);
1193                }
1194                for a in 0..k {
1195                    out[a] -= neg[a];
1196                }
1197            }
1198        });
1199
1200    Ok(closure)
1201}
1202
1203pub fn solve_reduced_beta_pcg_with_diagnostics(
1204    s_acc: &Array2<f64>,
1205    rhs_beta: &Array1<f64>,
1206    max_iterations: usize,
1207    relative_tolerance: f64,
1208) -> Result<(Array1<f64>, ArrowPcgDiagnostics), ArrowSchurGpuFailure> {
1209    let k = rhs_beta.len();
1210    if s_acc.dim() != (k, k) {
1211        return Err(ArrowSchurGpuFailure::SchurFactorFailed {
1212            reason: format!(
1213                "reduced-β GPU PCG requires a square (k×k) Schur block; got {:?} for k={k}",
1214                s_acc.dim()
1215            ),
1216        });
1217    }
1218    if k == 0 {
1219        return Err(ArrowSchurGpuFailure::Unavailable);
1220    }
1221
1222    #[cfg(not(target_os = "linux"))]
1223    {
1224        if relative_tolerance.is_nan() || max_iterations == 0 {
1225            return Err(ArrowSchurGpuFailure::SchurFactorFailed {
1226                reason: "reduced-β GPU PCG: invalid CG controls".to_string(),
1227            });
1228        }
1229        Err(ArrowSchurGpuFailure::Unavailable)
1230    }
1231
1232    #[cfg(target_os = "linux")]
1233    {
1234        cuda::solve_reduced_beta_pcg_with_diagnostics(
1235            s_acc,
1236            rhs_beta,
1237            max_iterations,
1238            relative_tolerance,
1239        )
1240    }
1241}
1242
1243pub fn solve_sae_matrix_free_pcg(
1244    sys: &ArrowSchurSystem,
1245    data: &DeviceSaePcgData,
1246    ridge_t: f64,
1247    ridge_beta: f64,
1248    rhs_beta: &Array1<f64>,
1249    max_iterations: usize,
1250    relative_tolerance: f64,
1251) -> Result<(Array1<f64>, ArrowPcgDiagnostics), ArrowSchurGpuFailure> {
1252    if sys.k != data.beta_dim || rhs_beta.len() != data.beta_dim || data.p == 0 {
1253        return Err(ArrowSchurGpuFailure::Unavailable);
1254    }
1255    #[cfg(not(target_os = "linux"))]
1256    {
1257        if ridge_t.is_nan()
1258            || ridge_beta.is_nan()
1259            || relative_tolerance.is_nan()
1260            || max_iterations == 0
1261        {
1262            return Err(ArrowSchurGpuFailure::SchurFactorFailed {
1263                reason: "SAE matrix-free GPU PCG: invalid controls".to_string(),
1264            });
1265        }
1266        Err(ArrowSchurGpuFailure::Unavailable)
1267    }
1268    #[cfg(target_os = "linux")]
1269    {
1270        // #1017/#1026 dispatch GUARD: framed data (frame metadata present) carries
1271        // a factored β border `G ⊗ W_{ij}` data Hessian and dense per-row cross
1272        // blocks the legacy `⊗ I_p` kernel CANNOT represent — feeding it framed
1273        // data would silently return a WRONG Newton step (it returns Ok with no
1274        // fallback). Route framed systems to the dedicated framed kernel and
1275        // legacy full-`B` systems to the legacy kernel; the two never cross.
1276        if data.frame.is_some() {
1277            cuda::solve_sae_matrix_free_pcg_framed(
1278                sys,
1279                data,
1280                ridge_t,
1281                ridge_beta,
1282                rhs_beta,
1283                max_iterations,
1284                relative_tolerance,
1285            )
1286        } else {
1287            cuda::solve_sae_matrix_free_pcg(
1288                sys,
1289                data,
1290                ridge_t,
1291                ridge_beta,
1292                rhs_beta,
1293                max_iterations,
1294                relative_tolerance,
1295            )
1296        }
1297    }
1298}
1299
1300/// #1017 device-resident SAE frame across the LM ridge ladder.
1301///
1302/// A single inner Newton step drives the proximal ridge ladder (up to
1303/// [`crate::arrow_schur::DEFAULT_PROXIMAL_MAX_ATTEMPTS`] trials) at a FIXED
1304/// system: only `ridge_t`/`ridge_beta` change per trial. In the per-trial
1305/// [`solve_sae_matrix_free_pcg`] path, `flatten_device_sae_frame_data` re-marshals
1306/// AND re-uploads every device operand each trial — yet the ONLY ridge-dependent
1307/// buffer is the per-row factored inverse `ainv = (H_tt + ridge_t·I)⁻¹` (the
1308/// smooth `λ S_k`, the framed `G ⊗ W`, and the dense per-row cross `H_tβ` are all
1309/// ridge-independent and constant across the ladder). This handle uploads the
1310/// ridge-independent buffers ONCE (at [`build_sae_resident_frame`]) and, per trial,
1311/// recomputes only `ainv` before running the identical framed PCG loop — so the
1312/// numbers are bit-identical to the per-trial re-flatten path while the
1313/// `(trials − 1) × (ridge-independent operand bytes)` re-upload is eliminated.
1314///
1315/// It is a trait object (mirroring [`crate::arrow_schur::GpuSchurMatvec`]) so the
1316/// concrete CUDA implementation — which owns `mod cuda`-only device buffers — can
1317/// be carried through the cfg-independent [`crate::arrow_schur::ArrowSolveOptions`]
1318/// without leaking a CUDA-only type into the shared solve options.
1319pub trait SaeResidentFrame {
1320    /// Refresh every ridge-independent numerical operand from a newly
1321    /// assembled nonlinear iterate while retaining the device allocations.
1322    /// Implementations must return `Unavailable` when any shape changes; the
1323    /// caller then builds a new frame. No factor or old numerical value may be
1324    /// reused across accepted iterates.
1325    fn refresh(&self, sys: &ArrowSchurSystem) -> Result<(), ArrowSchurGpuFailure>;
1326
1327    /// Recompute only the ridge-dependent per-row `ainv` at this trial's ridge,
1328    /// then run the framed reduced-Schur PCG against the resident buffers.
1329    /// Returns the reduced-β step `Δβ` and the PCG diagnostics, exactly as
1330    /// [`solve_sae_matrix_free_pcg`] would on the framed path. `Unavailable`
1331    /// signals a resident-path decline the caller should retry via the per-trial
1332    /// flatten; `RidgeBumpRequired`/`SchurFactorFailed` are genuine numerical
1333    /// signals the LM escalation must respond to (propagated unchanged).
1334    fn resolve(
1335        &self,
1336        sys: &ArrowSchurSystem,
1337        ridge_t: f64,
1338        ridge_beta: f64,
1339        rhs_beta: &Array1<f64>,
1340        max_iterations: usize,
1341        relative_tolerance: f64,
1342    ) -> Result<(Array1<f64>, ArrowPcgDiagnostics), ArrowSchurGpuFailure>;
1343}
1344
1345/// Build the device-resident SAE frame for the LM ridge ladder.
1346/// `Err(Unavailable)` is the decline signal — non-CUDA host, no framed device
1347/// data, or the offload predicate rejects the shape — exactly the contract of
1348/// the sibling device entry points ([`gpu_schur_matvec_backend`]): the caller
1349/// keeps the established per-trial re-flatten path completely unchanged.
1350/// `cg_iters` is the CG budget the offload gate scores (same value the
1351/// per-trial framed solve uses).
1352pub fn build_sae_resident_frame(
1353    sys: &ArrowSchurSystem,
1354    cg_iters: usize,
1355) -> Result<
1356    Option<std::sync::Arc<dyn SaeResidentFrame + Send + Sync>>,
1357    ArrowSchurGpuFailure,
1358> {
1359    // Target-independent admission: a zero-K system has no reduced-Schur block
1360    // to keep resident, and a zero CG budget can never consume the frame — both
1361    // decline on every host, keeping the per-trial flatten the single fallback
1362    // (on CUDA hosts this also spares a doomed device build attempt).
1363    if sys.k == 0 || cg_iters == 0 {
1364        return Ok(None);
1365    }
1366    #[cfg(target_os = "linux")]
1367    {
1368        cuda::ResidentSaeFrameHandle::build(sys, cg_iters).map(|handle| {
1369            handle.map(|frame| {
1370                std::sync::Arc::new(frame)
1371                    as std::sync::Arc<dyn SaeResidentFrame + Send + Sync>
1372            })
1373        })
1374    }
1375    // Non-CUDA host: there is no device to build a frame on.
1376    #[cfg(not(target_os = "linux"))]
1377    {
1378        Ok(None)
1379    }
1380}
1381
1382/// The ridge-INDEPENDENT host operands of the framed SAE reduced-Schur system,
1383/// marshalled into the contiguous upload layout `flatten_device_sae_frame_data`
1384/// consumes. Split out (with [`compute_ainv_host`], the sole ridge-DEPENDENT
1385/// buffer) so a single source builds both the per-trial flatten and the resident
1386/// frame, and so the host-marshalling cost is measurable off-device. Every field
1387/// here is a pure function of `(sys, data, frame)` — invariant across the ridge
1388/// ladder — which is exactly why the resident frame can upload them once.
1389#[cfg(target_os = "linux")]
1390pub struct FrameHostOperands {
1391    pub s_off: Vec<i32>,
1392    pub s_m: Vec<i32>,
1393    pub s_r: Vec<i32>,
1394    pub s_ptr: Vec<i32>,
1395    pub s_data: Vec<f64>,
1396    pub s_blocks: usize,
1397    pub g_off_i: Vec<i32>,
1398    pub g_off_j: Vec<i32>,
1399    pub g_ri: Vec<i32>,
1400    pub g_rj: Vec<i32>,
1401    pub g_mj: Vec<i32>,
1402    pub g_ptr: Vec<i32>,
1403    pub g_data: Vec<f64>,
1404    pub w_ptr: Vec<i32>,
1405    pub w_data: Vec<f64>,
1406    /// CSR over disjoint output spans. `g_group_blocks[g_group_ptr[g]..]`
1407    /// lists, in original block order, every G block contributing to the scalar
1408    /// output range `[g_group_off[g], g_group_off[g] + g_group_width[g])`.
1409    pub g_group_ptr: Vec<i32>,
1410    pub g_group_blocks: Vec<i32>,
1411    pub g_group_off: Vec<i32>,
1412    pub g_group_width: Vec<i32>,
1413    pub g_groups: usize,
1414    pub g_max_work: usize,
1415    pub htb_ptr: Vec<i32>,
1416    pub htb: Vec<f64>,
1417    pub q_of: Vec<i32>,
1418    pub n_rows: usize,
1419    pub k: usize,
1420    pub max_q: usize,
1421}
1422
1423#[cfg(target_os = "linux")]
1424fn frame_checked_i32(value: usize) -> Result<i32, ArrowSchurGpuFailure> {
1425    i32::try_from(value).map_err(|_| ArrowSchurGpuFailure::Unavailable)
1426}
1427
1428/// Fixed-order ownership plan for additive device blocks.
1429///
1430/// Each input span describes the complete scalar output range written by one
1431/// block. Equal spans are one ownership group; their block indices are retained
1432/// in input order. Distinct spans must be disjoint, because partially
1433/// overlapping ranges cannot be assigned to one owning thread per output
1434/// without expanding to a per-coordinate incidence matrix. SAE atom blocks
1435/// satisfy the stronger equal-or-disjoint invariant by construction, so a
1436/// partial overlap is malformed geometry and fails loudly.
1437#[cfg(target_os = "linux")]
1438fn deterministic_output_groups(
1439    spans: &[(usize, usize)],
1440    output_len: usize,
1441    label: &str,
1442) -> Result<DeterministicOutputGroups, ArrowSchurGpuFailure> {
1443    use std::collections::HashMap;
1444
1445    let malformed = |reason: String| ArrowSchurGpuFailure::SchurFactorFailed {
1446        reason: format!("{label} deterministic output ownership: {reason}"),
1447    };
1448    let mut lookup = HashMap::<(usize, usize), usize>::new();
1449    let mut group_spans = Vec::<(usize, usize)>::new();
1450    let mut group_blocks = Vec::<Vec<usize>>::new();
1451    for (block, &(offset, width)) in spans.iter().enumerate() {
1452        let end = offset.checked_add(width).ok_or_else(|| {
1453            malformed(format!(
1454                "block {block} output span overflows ({offset} + {width})"
1455            ))
1456        })?;
1457        if width == 0 || end > output_len {
1458            return Err(malformed(format!(
1459                "block {block} has output span [{offset}, {end}) outside [0, {output_len})"
1460            )));
1461        }
1462        let group = if let Some(&group) = lookup.get(&(offset, width)) {
1463            group
1464        } else {
1465            let group = group_spans.len();
1466            lookup.insert((offset, width), group);
1467            group_spans.push((offset, width));
1468            group_blocks.push(Vec::new());
1469            group
1470        };
1471        group_blocks[group].push(block);
1472    }
1473
1474    let mut by_offset: Vec<usize> = (0..group_spans.len()).collect();
1475    by_offset.sort_unstable_by_key(|&group| group_spans[group].0);
1476    for adjacent in by_offset.windows(2) {
1477        let left = group_spans[adjacent[0]];
1478        let right = group_spans[adjacent[1]];
1479        let left_end = left.0 + left.1;
1480        if left_end > right.0 {
1481            return Err(malformed(format!(
1482                "distinct output spans [{}, {}) and [{}, {}) overlap",
1483                left.0,
1484                left_end,
1485                right.0,
1486                right.0 + right.1
1487            )));
1488        }
1489    }
1490
1491    let mut ptr = Vec::with_capacity(group_spans.len() + 1);
1492    let mut blocks = Vec::with_capacity(spans.len());
1493    ptr.push(0_i32);
1494    for members in &group_blocks {
1495        for &block in members {
1496            blocks.push(frame_checked_i32(block)?);
1497        }
1498        ptr.push(frame_checked_i32(blocks.len())?);
1499    }
1500    let mut off = Vec::with_capacity(group_spans.len());
1501    let mut width = Vec::with_capacity(group_spans.len());
1502    let mut max_width = 0usize;
1503    for &(group_off, group_width) in &group_spans {
1504        off.push(frame_checked_i32(group_off)?);
1505        width.push(frame_checked_i32(group_width)?);
1506        max_width = max_width.max(group_width);
1507    }
1508    Ok(DeterministicOutputGroups {
1509        ptr,
1510        blocks,
1511        off,
1512        width,
1513        max_width,
1514    })
1515}
1516
1517#[cfg(target_os = "linux")]
1518struct DeterministicOutputGroups {
1519    ptr: Vec<i32>,
1520    blocks: Vec<i32>,
1521    off: Vec<i32>,
1522    width: Vec<i32>,
1523    max_width: usize,
1524}
1525
1526/// Marshal the ridge-INDEPENDENT framed operands into contiguous host buffers.
1527/// Bit-for-bit the same layout `flatten_device_sae_frame_data` produced inline;
1528/// factored out so the per-trial flatten and the resident frame share one source
1529/// and so the marshalling is measurable without a device.
1530#[cfg(target_os = "linux")]
1531pub fn flatten_frame_host_operands(
1532    sys: &ArrowSchurSystem,
1533    data: &DeviceSaePcgData,
1534    frame: &crate::arrow_schur::DeviceSaeFrameData,
1535) -> Result<FrameHostOperands, ArrowSchurGpuFailure> {
1536    let n_rows = sys.rows.len();
1537    let k = data.beta_dim;
1538    if frame.row_htbeta.len() != n_rows
1539        || frame.ranks.len() != frame.basis_sizes.len()
1540        || frame.border_offsets.len() != frame.ranks.len()
1541        || data.smooth_blocks.len() != frame.smooth_ranks.len()
1542    {
1543        return Err(ArrowSchurGpuFailure::Unavailable);
1544    }
1545
1546    // Smooth blocks.
1547    let mut s_off = Vec::new();
1548    let mut s_m = Vec::new();
1549    let mut s_r = Vec::new();
1550    let mut s_ptr = vec![0_i32];
1551    let mut s_data = Vec::<f64>::new();
1552    for (blk, &r) in data.smooth_blocks.iter().zip(frame.smooth_ranks.iter()) {
1553        let (m, mc) = blk.factor_a.dim();
1554        if m != mc {
1555            return Err(ArrowSchurGpuFailure::Unavailable);
1556        }
1557        s_off.push(frame_checked_i32(blk.global_offset)?);
1558        s_m.push(frame_checked_i32(m)?);
1559        s_r.push(frame_checked_i32(r)?);
1560        for ri in 0..m {
1561            for ci in 0..m {
1562                s_data.push(blk.factor_a[[ri, ci]]);
1563            }
1564        }
1565        s_ptr.push(frame_checked_i32(s_data.len())?);
1566    }
1567
1568    // Data blocks (g + w).
1569    let mut g_off_i = Vec::new();
1570    let mut g_off_j = Vec::new();
1571    let mut g_ri = Vec::new();
1572    let mut g_rj = Vec::new();
1573    let mut g_mj = Vec::new();
1574    let mut g_ptr = vec![0_i32];
1575    let mut g_data = Vec::<f64>::new();
1576    let mut w_ptr = vec![0_i32];
1577    let mut w_data = Vec::<f64>::new();
1578    let mut g_output_spans = Vec::with_capacity(frame.frame_blocks.len());
1579    for blk in &frame.frame_blocks {
1580        let ri = frame.ranks[blk.atom_i];
1581        let rj = frame.ranks[blk.atom_j];
1582        let (mi, mj) = blk.g.dim();
1583        if blk.w.dim() != (ri, rj) {
1584            return Err(ArrowSchurGpuFailure::Unavailable);
1585        }
1586        g_off_i.push(frame_checked_i32(frame.border_offsets[blk.atom_i])?);
1587        g_off_j.push(frame_checked_i32(frame.border_offsets[blk.atom_j])?);
1588        g_ri.push(frame_checked_i32(ri)?);
1589        g_rj.push(frame_checked_i32(rj)?);
1590        g_mj.push(frame_checked_i32(mj)?);
1591        for r in 0..mi {
1592            for c in 0..mj {
1593                g_data.push(blk.g[[r, c]]);
1594            }
1595        }
1596        g_ptr.push(frame_checked_i32(g_data.len())?);
1597        for a in 0..ri {
1598            for b in 0..rj {
1599                w_data.push(blk.w[[a, b]]);
1600            }
1601        }
1602        w_ptr.push(frame_checked_i32(w_data.len())?);
1603        let width = mi
1604            .checked_mul(ri)
1605            .ok_or(ArrowSchurGpuFailure::Unavailable)?;
1606        g_output_spans.push((frame.border_offsets[blk.atom_i], width));
1607    }
1608    let g_groups =
1609        deterministic_output_groups(&g_output_spans, k, "framed G matvec")?;
1610    let g_group_count = g_groups.off.len();
1611
1612    // Per-row dense cross-block + q (the factored ainv is ridge-dependent and
1613    // lives in `compute_ainv_host`, not here).
1614    let mut htb_ptr = vec![0_i32];
1615    let mut htb = Vec::<f64>::new();
1616    let mut q_of = Vec::<i32>::with_capacity(n_rows);
1617    let mut max_q = 0usize;
1618    for (i, slab) in frame.row_htbeta.iter().enumerate() {
1619        let qi = sys.row_dims[i];
1620        let q_eff = if !slab.is_empty() && slab.len() == qi * k {
1621            qi
1622        } else {
1623            0
1624        };
1625        q_of.push(frame_checked_i32(q_eff)?);
1626        max_q = max_q.max(q_eff);
1627        if q_eff > 0 {
1628            htb.extend_from_slice(slab);
1629        }
1630        htb_ptr.push(frame_checked_i32(htb.len())?);
1631    }
1632    if max_q == 0 {
1633        // No row contributes a reduced term — pure-penalty system. Give max_q=1
1634        // so the ainv buffer is non-empty.
1635        max_q = 1;
1636    }
1637
1638    Ok(FrameHostOperands {
1639        s_off,
1640        s_m,
1641        s_r,
1642        s_ptr,
1643        s_data,
1644        s_blocks: data.smooth_blocks.len(),
1645        g_off_i,
1646        g_off_j,
1647        g_ri,
1648        g_rj,
1649        g_mj,
1650        g_ptr,
1651        g_data,
1652        w_ptr,
1653        w_data,
1654        g_group_ptr: g_groups.ptr,
1655        g_group_blocks: g_groups.blocks,
1656        g_group_off: g_groups.off,
1657        g_group_width: g_groups.width,
1658        g_groups: g_group_count,
1659        g_max_work: g_groups.max_width,
1660        htb_ptr,
1661        htb,
1662        q_of,
1663        n_rows,
1664        k,
1665        max_q,
1666    })
1667}
1668
1669/// Recompute the ridge-DEPENDENT per-row factored inverse `ainv[i] = (H_tt^(i) +
1670/// ridge_t·I)⁻¹` as a row-major `n_rows × max_q × max_q` host buffer — the ONLY
1671/// buffer that changes across the ridge ladder. Bit-for-bit the computation
1672/// `flatten_device_sae_frame_data` did inline (per-row Cholesky with the
1673/// nonnegative-pivot guard, dense inverse via unit-column back-substitution, and
1674/// the Gershgorin `RidgeBumpRequired` deficit on a non-PD block).
1675#[cfg(target_os = "linux")]
1676pub fn compute_ainv_host(
1677    sys: &ArrowSchurSystem,
1678    q_of: &[i32],
1679    max_q: usize,
1680    n_rows: usize,
1681    ridge_t: f64,
1682) -> Result<Vec<f64>, ArrowSchurGpuFailure> {
1683    let mut ainv = vec![0.0_f64; n_rows * max_q * max_q];
1684    for (i, row) in sys.rows.iter().enumerate() {
1685        let q = q_of[i] as usize;
1686        if q == 0 {
1687            continue;
1688        }
1689        if row.htt.dim() != (q, q) {
1690            return Err(ArrowSchurGpuFailure::SchurFactorFailed {
1691                reason: format!(
1692                    "framed SAE device PCG row {i}: H_tt shape {:?} != ({q}, {q})",
1693                    row.htt.dim()
1694                ),
1695            });
1696        }
1697        let mut block = row.htt.clone();
1698        for d in 0..q {
1699            block[[d, d]] += ridge_t;
1700        }
1701        let factor = cholesky_factor_in_place(block.view(), CholeskyGuard::NonnegativePivot)
1702            .ok_or_else(|| ArrowSchurGpuFailure::RidgeBumpRequired {
1703                row: i,
1704                bump: ridge_bump_to_make_pd(row.htt.view(), ridge_t),
1705            })?;
1706        for col in 0..q {
1707            let mut e = Array1::<f64>::zeros(q);
1708            e[col] = 1.0;
1709            let solved = cholesky_solve_vector(factor.view(), e.view());
1710            for r in 0..q {
1711                ainv[i * max_q * max_q + r * max_q + col] = solved[r];
1712            }
1713        }
1714    }
1715    Ok(ainv)
1716}
1717
1718/// Reference dense back-end used by tests and as the fallback when the
1719/// GPU declines. Kept here (not in `arrow_schur_gpu.rs`) so the validation
1720/// suite has one canonical baseline.
1721pub fn solve_arrow_newton_step_dense_reference(
1722    sys: &ArrowSchurSystem,
1723    ridge_t: f64,
1724    ridge_beta: f64,
1725) -> Result<ArrowSchurGpuSolution, String> {
1726    let n = sys.rows.len();
1727    let d = sys.d;
1728    let k = sys.k;
1729    let total = n.checked_mul(d).ok_or("dimension overflow")? + k;
1730    let mut h = Array2::<f64>::zeros((total, total));
1731    let mut rhs = Array1::<f64>::zeros(total);
1732    for (i, row) in sys.rows.iter().enumerate() {
1733        let base = i * d;
1734        for c in 0..d {
1735            for r in 0..d {
1736                h[[base + r, base + c]] = row.htt[[r, c]];
1737            }
1738            h[[base + c, base + c]] += ridge_t;
1739        }
1740        for c in 0..k {
1741            for r in 0..d {
1742                let value = row.htbeta[[r, c]];
1743                h[[base + r, n * d + c]] = value;
1744                h[[n * d + c, base + r]] = value;
1745            }
1746        }
1747        for r in 0..d {
1748            rhs[base + r] = -row.gt[r];
1749        }
1750    }
1751    for c in 0..k {
1752        for r in 0..k {
1753            h[[n * d + r, n * d + c]] += sys.hbb[[r, c]];
1754        }
1755        h[[n * d + c, n * d + c]] += ridge_beta;
1756        rhs[n * d + c] = -sys.gb[c];
1757    }
1758    // #2015 — this GPU/device-reference dense factorization is INDEPENDENT of
1759    // the CPU dense reduced-Schur path's Jacobi/Van der Sluis diagonal
1760    // equilibration fix (`gam_solve::arrow_schur::reduced_solve::factor_dense_reduced_schur`).
1761    // Both paths are exact (a correctly-formed Cholesky of the true joint/Schur
1762    // matrix), so no correctness gap exists between them, but this path does
1763    // not yet get the CPU path's improved conditioning on an ill-scaled `H`;
1764    // porting the same equilibrate-then-reconstruct technique here is a
1765    // deliberate follow-up, not done in this change.
1766    let factor = cholesky_factor_in_place(h.view(), CholeskyGuard::NonnegativePivot)
1767        .ok_or_else(|| "dense reference Cholesky failed".to_string())?;
1768    let mut log_det = 0.0_f64;
1769    for i in 0..total {
1770        log_det += factor[[i, i]].ln();
1771    }
1772    log_det *= 2.0;
1773    let solved = cholesky_solve_vector(factor.view(), rhs.view());
1774    let delta_t = solved.slice(ndarray::s![..n * d]).to_owned();
1775    let delta_beta = solved.slice(ndarray::s![n * d..]).to_owned();
1776    Ok(ArrowSchurGpuSolution {
1777        delta_t,
1778        delta_beta,
1779        log_det_hessian: log_det,
1780    })
1781}
1782
1783/// Frames-engaged reduced-Schur penalty-side matvec `out = (P_ββ + ρ_β I)·x`,
1784/// computed purely from the factored device data (issue #1017/#1026). This is
1785/// the CPU bit-parity ORACLE for the GPU `arrow_sae_*` penalty kernels on the
1786/// frames path: smooth `λ S_k ⊗ I_{r_k}` (each `smooth_blocks[i]` at its
1787/// `global_offset` with right-width `frame.smooth_ranks[i]`) plus data-fit
1788/// `G_{ij} ⊗ W_{ij}` (each `frame.frame_blocks` entry, with the `μ`-major /
1789/// frame-minor index `border_offset[atom] + basis·r + frame_coord`). The
1790/// accumulation order matches the device kernels exactly.
1791///
1792/// `out` is OVERWRITTEN: first set to `ρ_β·x`, then the penalty blocks add in.
1793pub fn sae_framed_penalty_matvec_cpu(
1794    data: &DeviceSaePcgData,
1795    ridge_beta: f64,
1796    x: &[f64],
1797    out: &mut [f64],
1798) {
1799    let frame = data
1800        .frame
1801        .as_ref()
1802        .expect("sae_framed_penalty_matvec_cpu requires frame metadata");
1803    let k = data.beta_dim;
1804    for a in 0..k {
1805        out[a] = ridge_beta * x[a];
1806    }
1807    // Smooth penalty `λ S_k ⊗ I_{r_k}`: y[off + ia·r + ib] += Σ_ja S[ia,ja]·x[off + ja·r + ib].
1808    for (blk, &r) in data.smooth_blocks.iter().zip(frame.smooth_ranks.iter()) {
1809        let off = blk.global_offset;
1810        let m = blk.factor_a.nrows();
1811        for i_a in 0..m {
1812            for i_b in 0..r {
1813                let mut acc = 0.0_f64;
1814                for j_a in 0..m {
1815                    let s = blk.factor_a[[i_a, j_a]];
1816                    if s == 0.0 {
1817                        continue;
1818                    }
1819                    acc += s * x[off + j_a * r + i_b];
1820                }
1821                out[off + i_a * r + i_b] += acc;
1822            }
1823        }
1824    }
1825    // Data-fit penalty `G_{ij} ⊗ W_{ij}`.
1826    for blk in &frame.frame_blocks {
1827        let r_i = frame.ranks[blk.atom_i];
1828        let r_j = frame.ranks[blk.atom_j];
1829        let off_i = frame.border_offsets[blk.atom_i];
1830        let off_j = frame.border_offsets[blk.atom_j];
1831        let (m_i, m_j) = blk.g.dim();
1832        for li in 0..m_i {
1833            let yi_base = off_i + li * r_i;
1834            for lj in 0..m_j {
1835                let g = blk.g[[li, lj]];
1836                if g == 0.0 {
1837                    continue;
1838                }
1839                let xj_base = off_j + lj * r_j;
1840                for a in 0..r_i {
1841                    let mut acc = 0.0_f64;
1842                    for b in 0..r_j {
1843                        acc += blk.w[[a, b]] * x[xj_base + b];
1844                    }
1845                    out[yi_base + a] += g * acc;
1846                }
1847            }
1848        }
1849    }
1850}
1851
1852#[cfg(target_os = "linux")]
1853mod cuda {
1854    use super::{
1855        canonicalize_device_beta_factor,
1856        ArrowSchurGpuFailure, ArrowSchurGpuSolution, pack_block, pack_host,
1857        pack_host_d_and_stacked_b, project_device_beta_vector,
1858    };
1859    use crate::arrow_schur::{
1860        ArrowBetaGaugeQuotient, ArrowPcgDiagnostics, ArrowSchurSystem, DeviceSaeFrameData,
1861        DeviceSaePcgData, PcgStopReason,
1862    };
1863    use cudarc::cublas::sys::{
1864        cublasDiagType_t, cublasFillMode_t, cublasOperation_t, cublasSideMode_t, cublasStatus_t,
1865    };
1866    use cudarc::cublas::{CudaBlas, Gemm, GemmConfig, Gemv, GemvConfig};
1867    use cudarc::cusolver::{DnHandle, sys as cusolver_sys};
1868    use cudarc::driver::{
1869        CudaContext, CudaModule, CudaSlice, CudaStream, DevicePtr, DevicePtrMut, LaunchConfig,
1870        PushKernelArg,
1871    };
1872    use gam_gpu::driver::to_i32;
1873    use gam_gpu::linalg_dispatch::{DispatchOp, route_through_gpu};
1874    use ndarray::Array1;
1875    use std::sync::{Arc, OnceLock};
1876
1877    /// Per-row work slot for the row-block-granular multi-GPU solve. Inputs are
1878    /// the packed single-row buffers (`d×d` D block + ρ_t ridge, `d×k` B block,
1879    /// `d` g vector); the forward pass fills the whitened factors `l/u/y` and the
1880    /// per-tile reduction lands in the tile's leading slot.
1881    struct RowSlot {
1882        // Inputs (packed once on the host, column-major).
1883        d_block: Vec<f64>, // d*d
1884        b_block: Vec<f64>, // d*k
1885        g_vec: Vec<f64>,   // d
1886        // Forward outputs, kept on the host for the back-sub pass.
1887        l_block: Vec<f64>, // d*d lower factor, column-major
1888        u_vec: Vec<f64>,   // d   (= L^{-1} g)
1889        y_block: Vec<f64>, // d*k (= L^{-1} B), column-major
1890        log_det_local: f64,
1891        // Set on a non-PD pivot so the orchestrator can raise RidgeBumpRequired
1892        // for the offending global row instead of silently falling back.
1893        bump: Option<f64>,
1894        // Tile-level reduction, written into the tile's first slot only.
1895        tile_partial_schur: Option<Vec<f64>>, // k*k col-major, = Σ Y_iᵀY_i
1896        tile_partial_rhs: Option<Vec<f64>>,   // k, = Σ Y_iᵀu_i
1897        // Back-sub output for this row.
1898        delta_t_block: Vec<f64>, // d
1899    }
1900
1901    /// Row-block-granular multi-GPU Arrow-Schur Newton solve.
1902    ///
1903    /// The solve is separable across row blocks in both phases:
1904    ///   * forward — each row's local Cholesky `L_i`, whitening
1905    ///     `u_i = L_i⁻¹g_i`, `Y_i = L_i⁻¹B_i`, and partial Schur
1906    ///     `(Σ Y_iᵀY_i, Σ Y_iᵀu_i)` are independent;
1907    ///   * backward — `δt_i = -L_iᵀ⁻¹(u_i + Y_iδβ)` is independent.
1908    /// Only the small shared `K×K` reduce + factor + `δβ` solve is central.
1909    ///
1910    /// `gam_gpu::pool::scatter_batched` hands each device a contiguous row
1911    /// tile on its own bound context/stream; the per-tile forward keeps the
1912    /// POTRF fused with its dependent TRSM + Schur GEMM on that one stream, so no
1913    /// on-stream solve is orphaned. Tile partials and per-tile `log|L|` are
1914    /// reduced on the host (in tile/row order), `S_β` is factored on the primary
1915    /// device, and the back-sub is scattered back across the same tiles.
1916    ///
1917    /// Returns `Unavailable` (caller uses a single-device path) when the system
1918    /// carries matrix-free operators, the shared block is not dense `K×K`, the
1919    /// pool is single-device, or any tile's device work declines. A non-PD tip
1920    /// block surfaces as `RidgeBumpRequired` for the precise global row.
1921    pub(super) fn solve_multi_gpu(
1922        sys: &ArrowSchurSystem,
1923        ridge_t: f64,
1924        ridge_beta: f64,
1925        newton_schur_tikhonov_rel_floor: Option<f64>,
1926    ) -> Result<ArrowSchurGpuSolution, ArrowSchurGpuFailure> {
1927        let n = sys.rows.len();
1928        let d = sys.d;
1929        let k = sys.k;
1930        if n == 0 || d == 0 || k == 0 {
1931            return Err(ArrowSchurGpuFailure::Unavailable);
1932        }
1933        // Dense shared block + materialised per-row slabs are required; the
1934        // public entry already rejected matrix-free operators, but re-check so
1935        // this routine is safe in isolation.
1936        if sys.hbb_matvec.is_some() || sys.htbeta_matvec.is_some() || sys.hbb.dim() != (k, k) {
1937            return Err(ArrowSchurGpuFailure::Unavailable);
1938        }
1939
1940        let runtime = super::resolve_runtime_for_device_path()?
1941            .ok_or(ArrowSchurGpuFailure::Unavailable)?;
1942        if runtime.device_count() < 2 {
1943            return Err(ArrowSchurGpuFailure::Unavailable);
1944        }
1945
1946        // Pack one slot per row (column-major), folding ρ_t into each D block.
1947        let mut slots: Vec<RowSlot> = Vec::with_capacity(n);
1948        for row in &sys.rows {
1949            if row.htt.dim() != (d, d) || row.htbeta.dim() != (d, k) || row.gt.len() != d {
1950                return Err(ArrowSchurGpuFailure::Unavailable);
1951            }
1952            let mut d_block = Vec::with_capacity(d * d);
1953            let mut b_block = Vec::with_capacity(d * k);
1954            let mut g_vec = Vec::with_capacity(d);
1955            pack_block(row, ridge_t, d, k, &mut d_block, &mut b_block, &mut g_vec);
1956            slots.push(RowSlot {
1957                d_block,
1958                b_block,
1959                g_vec,
1960                l_block: Vec::new(),
1961                u_vec: Vec::new(),
1962                y_block: Vec::new(),
1963                log_det_local: 0.0,
1964                bump: None,
1965                tile_partial_schur: None,
1966                tile_partial_rhs: None,
1967                delta_t_block: vec![0.0; d],
1968            });
1969        }
1970
1971        // ---- Forward pass: per-device row tile, fused on its own stream ----
1972        let forward_ok = gam_gpu::pool::scatter_batched(runtime, &mut slots, |ordinal, tile| {
1973            forward_tile(ordinal, d, k, tile)
1974        });
1975        if forward_ok.is_none() {
1976            return Err(ArrowSchurGpuFailure::Unavailable);
1977        }
1978
1979        // Surface a non-PD tip block as a precise per-row ridge bump.
1980        let row_base_of_tile = gam_gpu::pool::balanced_partition(runtime, n);
1981        if let Some((row, bump)) = slots
1982            .iter()
1983            .enumerate()
1984            .find_map(|(i, slot)| slot.bump.map(|b| (i, b)))
1985        {
1986            return Err(ArrowSchurGpuFailure::RidgeBumpRequired { row, bump });
1987        }
1988
1989        // ---- Central: reduce tile partials → S_β, r_β; factor; solve δβ ----
1990        // Seed S_β with H_ββ + ρ_β I (column-major) and r_β with -g_β, then fold
1991        // in the per-tile partials in tile order so the reduction order tracks
1992        // the single-device accumulation (up to inter-tile reassociation).
1993        let mut schur_host = vec![0.0_f64; k * k];
1994        for col in 0..k {
1995            for row in 0..k {
1996                let mut v = sys.hbb[[row, col]];
1997                if row == col {
1998                    v += ridge_beta;
1999                }
2000                schur_host[col * k + row] = v;
2001            }
2002        }
2003        let mut rhs_host: Vec<f64> = sys.gb.iter().map(|v| -v).collect();
2004        let mut log_det = 0.0_f64;
2005        for start in tile_starts(&row_base_of_tile) {
2006            let slot = &slots[start];
2007            let partial_schur = slot
2008                .tile_partial_schur
2009                .as_ref()
2010                .ok_or(ArrowSchurGpuFailure::Unavailable)?;
2011            let partial_rhs = slot
2012                .tile_partial_rhs
2013                .as_ref()
2014                .ok_or(ArrowSchurGpuFailure::Unavailable)?;
2015            // `accumulate_schur` writes `partial_schur = -Σ_tile Y_iᵀY_i` (GEMM
2016            // α=-1, β=1 into a zero seed) and `partial_rhs = +Σ_tile Y_iᵀu_i`.
2017            // The reduced Schur is `S = (H_ββ+ρI) − Σ_all Y_iᵀY_i`, so adding the
2018            // (already-negated) partials reproduces the single-device sign.
2019            for idx in 0..k * k {
2020                schur_host[idx] += partial_schur[idx];
2021            }
2022            for a in 0..k {
2023                rhs_host[a] += partial_rhs[a];
2024            }
2025        }
2026        for slot in &slots {
2027            log_det += slot.log_det_local;
2028        }
2029        let canonical_factor = canonicalize_device_beta_factor(
2030            sys.beta_gauge_quotient.as_ref(),
2031            newton_schur_tikhonov_rel_floor,
2032            sys.k,
2033            &mut schur_host,
2034        )?;
2035        rhs_host = project_device_beta_vector(sys.beta_gauge_quotient.as_ref(), sys.k, rhs_host)?;
2036
2037        // Factor S_β and solve δβ on the primary device (small K×K leaf). The
2038        // stream carries the primary context (same pattern as `solve()`); no
2039        // thread bind is needed for the cuSOLVER/cuBLAS handles created from it.
2040        let primary = runtime.selected_device().ordinal;
2041        let stream = gam_gpu::device_runtime::cuda_context_for(primary)
2042            .and_then(|ctx| ctx.new_stream().ok())
2043            .ok_or(ArrowSchurGpuFailure::Unavailable)?;
2044        let solver =
2045            DnHandle::new(stream.clone()).map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2046        let blas = CudaBlas::new(stream.clone()).map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2047        let mut schur_dev = stream
2048            .clone_htod(&schur_host)
2049            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2050        let mut rhs_dev = stream
2051            .clone_htod(&rhs_host)
2052            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2053        if !canonical_factor {
2054            let info = potrf_single(&solver, &stream, k, &mut schur_dev)?;
2055            if info != 0 {
2056                return Err(ArrowSchurGpuFailure::SchurFactorFailed {
2057                    reason: format!("multi-GPU Schur Cholesky failed at pivot {info}"),
2058                });
2059            }
2060        }
2061        trsm_single(&blas, &stream, k, &schur_dev, &mut rhs_dev, false, false)?;
2062        trsm_single(&blas, &stream, k, &schur_dev, &mut rhs_dev, false, true)?;
2063        let delta_beta_host = project_device_beta_vector(
2064            sys.beta_gauge_quotient.as_ref(),
2065            sys.k,
2066            stream
2067                .clone_dtoh(&rhs_dev)
2068                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
2069        )?;
2070        let delta_beta = Array1::from_vec(delta_beta_host.clone());
2071        let l_schur_host = stream
2072            .clone_dtoh(&schur_dev)
2073            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2074        for j in 0..k {
2075            log_det += l_schur_host[j * k + j].ln();
2076        }
2077        log_det *= 2.0;
2078
2079        // ---- Backward pass: δt_i = -L_iᵀ⁻¹(u_i + Y_iδβ), per-device tile ----
2080        let delta_beta_ref = &delta_beta_host;
2081        let back_ok = gam_gpu::pool::scatter_batched(runtime, &mut slots, |ordinal, tile| {
2082            back_sub_tile(ordinal, d, k, delta_beta_ref, tile)
2083        });
2084        if back_ok.is_none() {
2085            return Err(ArrowSchurGpuFailure::Unavailable);
2086        }
2087
2088        // Stitch per-row δt into the stacked (n*d) result.
2089        let mut delta_t = Array1::<f64>::zeros(n * d);
2090        for (i, slot) in slots.iter().enumerate() {
2091            let base = i * d;
2092            for r in 0..d {
2093                delta_t[base + r] = slot.delta_t_block[r];
2094            }
2095        }
2096
2097        Ok(ArrowSchurGpuSolution {
2098            delta_t,
2099            delta_beta,
2100            log_det_hessian: log_det,
2101        })
2102    }
2103
2104    /// Tile starts: the leading global row index of each device tile (where the
2105    /// tile-level partial reduction was written by the forward pass).
2106    fn tile_starts(tiles: &[(usize, std::ops::Range<usize>)]) -> impl Iterator<Item = usize> + '_ {
2107        tiles.iter().map(|(_, range)| range.start)
2108    }
2109
2110    /// Forward pass for one device row tile, running on `ordinal`'s bound stream.
2111    /// Factors each row block, whitens `u`/`Y`, accumulates the tile's partial
2112    /// Schur `(Σ Y_iᵀY_i, Σ Y_iᵀu_i)` into the tile's leading slot, keeps the
2113    /// per-row `L`/`u`/`Y` on the host for back-sub, and records the per-row
2114    /// `Σ_j log L_jj`. A non-PD pivot is recorded in `slot.bump` (the tile still
2115    /// returns `Some(())` so the orchestrator raises a precise `RidgeBumpRequired`
2116    /// rather than collapsing the whole batch to CPU).
2117    fn forward_tile(ordinal: usize, d: usize, k: usize, tile: &mut [RowSlot]) -> Option<()> {
2118        if tile.is_empty() {
2119            return Some(());
2120        }
2121        // `scatter_batched` has already bound this ordinal's context on this
2122        // worker thread; the stream below targets that same device.
2123        let stream = gam_gpu::device_runtime::cuda_context_for(ordinal)
2124            .and_then(|ctx| ctx.new_stream().ok())?;
2125        let solver = DnHandle::new(stream.clone()).ok()?;
2126        let blas = CudaBlas::new(stream.clone()).ok()?;
2127        let m = tile.len();
2128
2129        // Stack the tile's D, B, g into contiguous device buffers (same layout
2130        // the single-device path packs for `m` rows).
2131        let mut d_host = Vec::with_capacity(m * d * d);
2132        let mut b_host = Vec::with_capacity(m * d * k);
2133        let mut g_host = Vec::with_capacity(m * d);
2134        for slot in tile.iter() {
2135            d_host.extend_from_slice(&slot.d_block);
2136            b_host.extend_from_slice(&slot.b_block);
2137            g_host.extend_from_slice(&slot.g_vec);
2138        }
2139        let mut d_dev = stream.clone_htod(&d_host).ok()?;
2140        let mut b_dev = stream.clone_htod(&b_host).ok()?;
2141        let mut g_dev = stream.clone_htod(&g_host).ok()?;
2142
2143        // Batched POTRF; a non-PD block records its bump and stops the tile.
2144        // The bump is deficit-aware (Gershgorin lower bound on λ_min of the
2145        // already-ridged `d_block`), NOT derived from the cuSOLVER `info` —
2146        // which is a 1-based pivot ROW INDEX, not a pivot magnitude — so a
2147        // strongly-indefinite block recovers in one outer-loop retry.
2148        let info_host = potrf_batched(&solver, &stream, d, m, &mut d_dev).ok()?;
2149        if let Some(local) = info_host.iter().position(|info| *info != 0) {
2150            tile[local].bump = Some(super::ridge_bump_to_make_pd_colmajor(
2151                &tile[local].d_block,
2152                d,
2153            ));
2154            return Some(());
2155        }
2156
2157        // Whiten: u = L⁻¹ g, Y = L⁻¹ B.
2158        trsm_batched_lower_inplace(&blas, &stream, d, m, 1, &d_dev, &mut g_dev).ok()?;
2159        trsm_batched_lower_inplace(&blas, &stream, d, m, k, &d_dev, &mut b_dev).ok()?;
2160
2161        // Tile partial Schur: zero-seeded so the host adds the H_ββ seed once.
2162        let mut schur_dev = stream.alloc_zeros::<f64>(k * k).ok()?;
2163        let mut rhs_dev = stream.alloc_zeros::<f64>(k).ok()?;
2164        accumulate_schur(&blas, d, k, m, &b_dev, &g_dev, &mut schur_dev, &mut rhs_dev).ok()?;
2165
2166        // Download L, u, Y, and the tile partials.
2167        let l_host = stream.clone_dtoh(&d_dev).ok()?;
2168        let u_host = stream.clone_dtoh(&g_dev).ok()?;
2169        let y_host = stream.clone_dtoh(&b_dev).ok()?;
2170        let partial_schur = stream.clone_dtoh(&schur_dev).ok()?;
2171        let partial_rhs = stream.clone_dtoh(&rhs_dev).ok()?;
2172
2173        for (local, slot) in tile.iter_mut().enumerate() {
2174            let l_base = local * d * d;
2175            let u_base = local * d;
2176            let y_base = local * d * k;
2177            slot.l_block = l_host[l_base..l_base + d * d].to_vec();
2178            slot.u_vec = u_host[u_base..u_base + d].to_vec();
2179            slot.y_block = y_host[y_base..y_base + d * k].to_vec();
2180            let mut log_det_local = 0.0_f64;
2181            for j in 0..d {
2182                log_det_local += l_host[l_base + j * d + j].ln();
2183            }
2184            slot.log_det_local = log_det_local;
2185        }
2186        tile[0].tile_partial_schur = Some(partial_schur);
2187        tile[0].tile_partial_rhs = Some(partial_rhs);
2188        Some(())
2189    }
2190
2191    /// Back-substitution for one device row tile: `δt_i = -L_iᵀ⁻¹(u_i + Y_iδβ)`.
2192    /// Re-uploads the tile's kept `L`/`u`/`Y` to `ordinal`, applies the GEMV
2193    /// accumulate + transposed TRSM, and writes each row's `δt` into its slot.
2194    fn back_sub_tile(
2195        ordinal: usize,
2196        d: usize,
2197        k: usize,
2198        delta_beta: &[f64],
2199        tile: &mut [RowSlot],
2200    ) -> Option<()> {
2201        if tile.is_empty() {
2202            return Some(());
2203        }
2204        // `scatter_batched` has already bound this ordinal's context on this
2205        // worker thread; the stream below targets that same device.
2206        let stream = gam_gpu::device_runtime::cuda_context_for(ordinal)
2207            .and_then(|ctx| ctx.new_stream().ok())?;
2208        let blas = CudaBlas::new(stream.clone()).ok()?;
2209        let m = tile.len();
2210
2211        let mut l_host = Vec::with_capacity(m * d * d);
2212        let mut u_host = Vec::with_capacity(m * d);
2213        let mut y_host = Vec::with_capacity(m * d * k);
2214        for slot in tile.iter() {
2215            l_host.extend_from_slice(&slot.l_block);
2216            u_host.extend_from_slice(&slot.u_vec);
2217            y_host.extend_from_slice(&slot.y_block);
2218        }
2219        let d_dev = stream.clone_htod(&l_host).ok()?;
2220        let mut g_dev = stream.clone_htod(&u_host).ok()?;
2221        let b_dev = stream.clone_htod(&y_host).ok()?;
2222        let rhs_dev = stream.clone_htod(&delta_beta.to_vec()).ok()?;
2223
2224        // g ← u + Y·δβ, then x = L⁻ᵀ g; δt = -x.
2225        accumulate_back_sub_rhs(&blas, d, k, m, &b_dev, &rhs_dev, &mut g_dev).ok()?;
2226        trsm_batched_lower_inplace_transposed(&blas, &stream, d, m, 1, &d_dev, &mut g_dev).ok()?;
2227        let x_host = stream.clone_dtoh(&g_dev).ok()?;
2228        for (local, slot) in tile.iter_mut().enumerate() {
2229            let base = local * d;
2230            for r in 0..d {
2231                slot.delta_t_block[r] = -x_host[base + r];
2232            }
2233        }
2234        Some(())
2235    }
2236
2237    pub(super) fn solve(
2238        sys: &ArrowSchurSystem,
2239        ridge_t: f64,
2240        ridge_beta: f64,
2241        newton_schur_tikhonov_rel_floor: Option<f64>,
2242    ) -> Result<ArrowSchurGpuSolution, ArrowSchurGpuFailure> {
2243        let n = sys.rows.len();
2244        let d = sys.d;
2245        let k = sys.k;
2246        let runtime = route_through_gpu(DispatchOp::SmallDenseBatchedPotrf { p: d, batch: n })
2247            .ok_or(ArrowSchurGpuFailure::Unavailable)?;
2248
2249        let stream = gam_gpu::device_runtime::cuda_context_for(runtime.device.ordinal)
2250            .and_then(|ctx| ctx.new_stream().ok())
2251            .ok_or(ArrowSchurGpuFailure::Unavailable)?;
2252        let solver =
2253            DnHandle::new(stream.clone()).map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2254        let blas = CudaBlas::new(stream.clone()).map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2255
2256        // ----- Pack + upload D, B, g -----
2257        let (d_host, b_host, g_host) = pack_host(sys, ridge_t);
2258        let mut d_dev = stream
2259            .clone_htod(&d_host)
2260            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2261        let mut b_dev = stream
2262            .clone_htod(&b_host)
2263            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2264        let mut g_dev = stream
2265            .clone_htod(&g_host)
2266            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2267
2268        // ----- Layer A: batched lower Cholesky of D in place -----
2269        // This POTRF is fused with the downstream TRSM + Schur GEMM + back-sub
2270        // on this one stream, so splitting only the POTRF across devices would
2271        // orphan the dependent on-stream solves. Multi-GPU here is the
2272        // whole-solve row-block split in `solve_arrow_newton_step` (see
2273        // `solve_multi_gpu`), not a per-layer split — this device-resident path
2274        // is the single-device leaf the split dispatches per tile.
2275        let info_host = potrf_batched(&solver, &stream, d, n, &mut d_dev)?;
2276        if let Some(idx) = info_host.iter().position(|info| *info != 0) {
2277            // `info` is cuSOLVER's 1-based pivot ROW INDEX, not a magnitude;
2278            // size the bump from the block's own entries (Gershgorin λ_min
2279            // bound) so a strongly-indefinite block recovers in one retry.
2280            return Err(ArrowSchurGpuFailure::RidgeBumpRequired {
2281                row: idx,
2282                bump: super::ridge_bump_to_make_pd(sys.rows[idx].htt.view(), ridge_t),
2283            });
2284        }
2285
2286        // ----- Layer B (1/2): in-place triangular solves -----
2287        // u_i = L_i^{-1} g_i, packed as a stacked (n*d) column-vector.
2288        trsm_batched_lower_inplace(&blas, &stream, d, n, 1, &d_dev, &mut g_dev)?;
2289        // Y_i = L_i^{-1} B_i, in place over the (n*d) × k buffer (laid out as
2290        // n stacked column-major d×k tiles).
2291        trsm_batched_lower_inplace(&blas, &stream, d, n, k, &d_dev, &mut b_dev)?;
2292
2293        // ----- Layer B (2/2): Schur reduction via single big GEMM / GEMV -----
2294        // Y_all is (n*d) × k column-major: viewing all n stacked d×k tiles as
2295        // one big matrix is bit-exact because each tile is column-major with
2296        // leading dim d and the tiles are contiguous in memory, so the
2297        // combined leading dim is n*d only for the *outer* matrix view. To
2298        // make the single-GEMM equivalence hold we must treat the stacked
2299        // buffer as (n*d) × k column-major with leading dim = n*d, which
2300        // means columns of Y_all are interleaved by row across blocks.
2301        // That is NOT what we packed. So we use the cuBLAS stride pattern
2302        // instead: stride-by-block, transpose-A, and *accumulate* into one
2303        // S_β buffer via beta=1 across batches. Equivalent flop count, no
2304        // extra reduction kernel, and correct layout.
2305        //
2306        // Concretely: schur ← C + ρ_β I; rhs ← -g_β; then for each block
2307        //   schur -= Y_i^T Y_i      (k×k)
2308        //   rhs   += Y_i^T u_i      (k)
2309        // We launch this as `n` sequential GEMMs/GEMVs with beta=1 on the
2310        // accumulator. Layer D fuses these into one NVRTC launch.
2311        let schur_init: Vec<f64> = {
2312            let mut tmp = Vec::with_capacity(k * k);
2313            for col in 0..k {
2314                for row in 0..k {
2315                    let mut v = sys.hbb[[row, col]];
2316                    if row == col {
2317                        v += ridge_beta;
2318                    }
2319                    tmp.push(v);
2320                }
2321            }
2322            tmp
2323        };
2324        let mut schur_dev = stream
2325            .clone_htod(&schur_init)
2326            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2327        let rhs_init: Vec<f64> = sys.gb.iter().map(|v| -v).collect();
2328        let mut rhs_dev = stream
2329            .clone_htod(&rhs_init)
2330            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2331
2332        accumulate_schur(&blas, d, k, n, &b_dev, &g_dev, &mut schur_dev, &mut rhs_dev)?;
2333        let canonical_factor = if sys.beta_gauge_quotient.is_some()
2334            || newton_schur_tikhonov_rel_floor.is_some()
2335        {
2336            let mut schur_host = stream
2337                .clone_dtoh(&schur_dev)
2338                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2339            let factorized = canonicalize_device_beta_factor(
2340                sys.beta_gauge_quotient.as_ref(),
2341                newton_schur_tikhonov_rel_floor,
2342                sys.k,
2343                &mut schur_host,
2344            )?;
2345            schur_dev = stream
2346                .clone_htod(&schur_host)
2347                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2348            if sys.beta_gauge_quotient.is_some() {
2349                let rhs_host = project_device_beta_vector(
2350                    sys.beta_gauge_quotient.as_ref(),
2351                    sys.k,
2352                    stream
2353                        .clone_dtoh(&rhs_dev)
2354                        .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
2355                )?;
2356                rhs_dev = stream
2357                    .clone_htod(&rhs_host)
2358                    .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2359            }
2360            factorized
2361        } else {
2362            false
2363        };
2364
2365        // ----- Layer C (1/2): factor S_β and solve for δβ -----
2366        if !canonical_factor {
2367            let info = potrf_single(&solver, &stream, k, &mut schur_dev)?;
2368            if info != 0 {
2369                return Err(ArrowSchurGpuFailure::SchurFactorFailed {
2370                    reason: format!("Schur Cholesky failed at pivot {info}"),
2371                });
2372            }
2373        }
2374        // δβ ← L_S^{-T} L_S^{-1} rhs
2375        trsm_single(&blas, &stream, k, &schur_dev, &mut rhs_dev, false, false)?;
2376        trsm_single(&blas, &stream, k, &schur_dev, &mut rhs_dev, false, true)?;
2377        let delta_beta_host = project_device_beta_vector(
2378            sys.beta_gauge_quotient.as_ref(),
2379            sys.k,
2380            stream
2381                .clone_dtoh(&rhs_dev)
2382                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
2383        )?;
2384        if sys.beta_gauge_quotient.is_some() {
2385            rhs_dev = stream
2386                .clone_htod(&delta_beta_host)
2387                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2388        }
2389        let delta_beta = Array1::from_vec(delta_beta_host);
2390
2391        // ----- Layer C (2/2): back-sub δt_i = -L_i^{-T} (u_i + Y_i δβ) -----
2392        // Already on device:
2393        //   g_dev holds u_i stacked (n*d).
2394        //   b_dev holds Y_i stacked column-major n×(d×k) tiles.
2395        // Compute g_dev ← g_dev + Y_block · δβ per block (cuBLAS gemv with beta=1),
2396        // then in-place trsm with L_i^T (CUBLAS_OP_T) to obtain x_i, and finally
2397        // δt_i = -x_i on host after download.
2398        accumulate_back_sub_rhs(&blas, d, k, n, &b_dev, &rhs_dev, &mut g_dev)?;
2399        trsm_batched_lower_inplace_transposed(&blas, &stream, d, n, 1, &d_dev, &mut g_dev)?;
2400
2401        let x_host = stream
2402            .clone_dtoh(&g_dev)
2403            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2404        let mut delta_t = Array1::<f64>::zeros(n * d);
2405        for (i, v) in x_host.iter().enumerate() {
2406            delta_t[i] = -*v;
2407        }
2408
2409        // ----- log|H| = 2 Σ log L_{i,jj} + 2 Σ log R_{β,aa} -----
2410        let l_local_host = stream
2411            .clone_dtoh(&d_dev)
2412            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2413        let l_schur_host = stream
2414            .clone_dtoh(&schur_dev)
2415            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2416        let mut log_det = 0.0_f64;
2417        for i in 0..n {
2418            let base = i * d * d;
2419            for j in 0..d {
2420                log_det += l_local_host[base + j * d + j].ln();
2421            }
2422        }
2423        for j in 0..k {
2424            log_det += l_schur_host[j * k + j].ln();
2425        }
2426        log_det *= 2.0;
2427
2428        Ok(ArrowSchurGpuSolution {
2429            delta_t,
2430            delta_beta,
2431            log_det_hessian: log_det,
2432        })
2433    }
2434
2435    fn potrf_batched(
2436        solver: &DnHandle,
2437        stream: &Arc<CudaStream>,
2438        p: usize,
2439        batch: usize,
2440        matrices: &mut CudaSlice<f64>,
2441    ) -> Result<Vec<i32>, ArrowSchurGpuFailure> {
2442        let p_i = to_i32(p).ok_or(ArrowSchurGpuFailure::Unavailable)?;
2443        let batch_i = to_i32(batch).ok_or(ArrowSchurGpuFailure::Unavailable)?;
2444        let matrix_len = p * p;
2445        let bytes_per = (matrix_len * std::mem::size_of::<f64>()) as u64;
2446        let (base_ptr, _record) = matrices.device_ptr_mut(stream);
2447        let mut ptrs = Vec::with_capacity(batch);
2448        for idx in 0..batch {
2449            ptrs.push(base_ptr + (idx as u64) * bytes_per);
2450        }
2451        let mut ptrs_dev = stream
2452            .clone_htod(&ptrs)
2453            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2454        let mut info_dev = stream
2455            .alloc_zeros::<i32>(batch)
2456            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2457        let status = {
2458            let (ptrs_ptr, _ptrs_record) = ptrs_dev.device_ptr_mut(stream);
2459            let (info_ptr, _info_record) = info_dev.device_ptr_mut(stream);
2460            // SAFETY: pointer array and info buffer live on the device,
2461            // matrices_dev holds `batch` contiguous p×p column-major blocks.
2462            unsafe {
2463                cusolver_sys::cusolverDnDpotrfBatched(
2464                    solver.cu(),
2465                    cusolver_sys::cublasFillMode_t::CUBLAS_FILL_MODE_LOWER,
2466                    p_i,
2467                    ptrs_ptr as *mut *mut f64,
2468                    p_i,
2469                    info_ptr as *mut i32,
2470                    batch_i,
2471                )
2472            }
2473        };
2474        if status != cusolver_sys::cusolverStatus_t::CUSOLVER_STATUS_SUCCESS {
2475            return Err(ArrowSchurGpuFailure::Unavailable);
2476        }
2477        stream
2478            .clone_dtoh(&info_dev)
2479            .map_err(|_| ArrowSchurGpuFailure::Unavailable)
2480    }
2481
2482    fn potrf_single(
2483        solver: &DnHandle,
2484        stream: &Arc<CudaStream>,
2485        p: usize,
2486        matrix: &mut CudaSlice<f64>,
2487    ) -> Result<i32, ArrowSchurGpuFailure> {
2488        let p_i = to_i32(p).ok_or(ArrowSchurGpuFailure::Unavailable)?;
2489        let uplo = cusolver_sys::cublasFillMode_t::CUBLAS_FILL_MODE_LOWER;
2490        let mut lwork = 0_i32;
2491        {
2492            let (mat_ptr, _rec) = matrix.device_ptr_mut(stream);
2493            // SAFETY: buffer query against a live p-by-p column-major device matrix.
2494            let status = unsafe {
2495                cusolver_sys::cusolverDnDpotrf_bufferSize(
2496                    solver.cu(),
2497                    uplo,
2498                    p_i,
2499                    mat_ptr as *mut f64,
2500                    p_i,
2501                    &mut lwork,
2502                )
2503            };
2504            if status != cusolver_sys::cusolverStatus_t::CUSOLVER_STATUS_SUCCESS {
2505                return Err(ArrowSchurGpuFailure::Unavailable);
2506            }
2507        }
2508        let lwork_usize = usize::try_from(lwork).map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2509        let mut workspace = stream
2510            .alloc_zeros::<f64>(lwork_usize.max(1))
2511            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2512        let mut info_dev = stream
2513            .alloc_zeros::<i32>(1)
2514            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2515        {
2516            let (mat_ptr, _rec) = matrix.device_ptr_mut(stream);
2517            let (work_ptr, _wrec) = workspace.device_ptr_mut(stream);
2518            let (info_ptr, _irec) = info_dev.device_ptr_mut(stream);
2519            // SAFETY: all three pointers refer to live, correctly sized device buffers.
2520            let status = unsafe {
2521                cusolver_sys::cusolverDnDpotrf(
2522                    solver.cu(),
2523                    uplo,
2524                    p_i,
2525                    mat_ptr as *mut f64,
2526                    p_i,
2527                    work_ptr as *mut f64,
2528                    lwork,
2529                    info_ptr as *mut i32,
2530                )
2531            };
2532            if status != cusolver_sys::cusolverStatus_t::CUSOLVER_STATUS_SUCCESS {
2533                return Err(ArrowSchurGpuFailure::Unavailable);
2534            }
2535        }
2536        let info_host = stream
2537            .clone_dtoh(&info_dev)
2538            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2539        Ok(info_host[0])
2540    }
2541
2542    /// In-place lower-triangular solves `X_i ← L_i^{-1} X_i` over the n stacked
2543    /// d×nrhs RHS tiles in `rhs`. Uses `cublasDtrsmBatched` so all n solves
2544    /// hit the device in one launch.
2545    /// One batched triangular solve `X_i ← L_i^{-∗} X_i`: the block geometry
2546    /// plus the RHS memory layout the solve walks.
2547    ///
2548    /// The layout is a parameter because the same whitening feeds two different
2549    /// downstream reductions. Block-contiguous tiles ([`Self::block_tiles`])
2550    /// suit per-row consumers; the stacked layout ([`Self::stacked_matrix`])
2551    /// makes the n whitened blocks ONE `(n·d) × nrhs` matrix, which is what lets
2552    /// the Schur reduction be a single GEMM ([`schur_gemm_stacked`]) instead of
2553    /// n rank-`d` updates. `cublasDtrsmBatched` takes an explicit `ldb`, so the
2554    /// whitening costs the same single launch either way.
2555    #[derive(Clone, Copy)]
2556    struct BatchedTrsmSpec {
2557        d: usize,
2558        n: usize,
2559        nrhs: usize,
2560        rhs_leading_dim: usize,
2561        rhs_block_stride: usize,
2562        transposed: bool,
2563    }
2564
2565    impl BatchedTrsmSpec {
2566        /// Block `i` is a `d × nrhs` column-major tile at offset `i·d·nrhs`.
2567        fn block_tiles(d: usize, n: usize, nrhs: usize) -> Self {
2568            Self {
2569                d,
2570                n,
2571                nrhs,
2572                rhs_leading_dim: d,
2573                rhs_block_stride: d * nrhs,
2574                transposed: false,
2575            }
2576        }
2577
2578        /// Block `i` occupies rows `i·d .. i·d+d` of ONE `(n·d) × nrhs`
2579        /// column-major matrix with leading dimension `n·d`.
2580        fn stacked_matrix(d: usize, n: usize, nrhs: usize) -> Self {
2581            Self {
2582                d,
2583                n,
2584                nrhs,
2585                rhs_leading_dim: n * d,
2586                rhs_block_stride: d,
2587                transposed: false,
2588            }
2589        }
2590
2591        fn transposed(self) -> Self {
2592            Self {
2593                transposed: true,
2594                ..self
2595            }
2596        }
2597    }
2598
2599    fn trsm_batched_lower_inplace(
2600        blas: &CudaBlas,
2601        stream: &Arc<CudaStream>,
2602        d: usize,
2603        n: usize,
2604        nrhs: usize,
2605        l_stack: &CudaSlice<f64>,
2606        rhs_stack: &mut CudaSlice<f64>,
2607    ) -> Result<(), ArrowSchurGpuFailure> {
2608        trsm_batched_inplace_inner(
2609            blas,
2610            stream,
2611            BatchedTrsmSpec::block_tiles(d, n, nrhs),
2612            l_stack,
2613            rhs_stack,
2614        )
2615    }
2616
2617    /// As above but with `L_i^T` instead of `L_i`.
2618    fn trsm_batched_lower_inplace_transposed(
2619        blas: &CudaBlas,
2620        stream: &Arc<CudaStream>,
2621        d: usize,
2622        n: usize,
2623        nrhs: usize,
2624        l_stack: &CudaSlice<f64>,
2625        rhs_stack: &mut CudaSlice<f64>,
2626    ) -> Result<(), ArrowSchurGpuFailure> {
2627        trsm_batched_inplace_inner(
2628            blas,
2629            stream,
2630            BatchedTrsmSpec::block_tiles(d, n, nrhs).transposed(),
2631            l_stack,
2632            rhs_stack,
2633        )
2634    }
2635
2636    /// `X_i ← L_i^{-1} X_i` over the n row blocks of a SINGLE stacked
2637    /// `(n·d) × nrhs` column-major matrix.
2638    fn trsm_batched_lower_inplace_stacked(
2639        blas: &CudaBlas,
2640        stream: &Arc<CudaStream>,
2641        d: usize,
2642        n: usize,
2643        nrhs: usize,
2644        l_stack: &CudaSlice<f64>,
2645        rhs_stacked: &mut CudaSlice<f64>,
2646    ) -> Result<(), ArrowSchurGpuFailure> {
2647        trsm_batched_inplace_inner(
2648            blas,
2649            stream,
2650            BatchedTrsmSpec::stacked_matrix(d, n, nrhs),
2651            l_stack,
2652            rhs_stacked,
2653        )
2654    }
2655
2656    fn trsm_batched_inplace_inner(
2657        blas: &CudaBlas,
2658        stream: &Arc<CudaStream>,
2659        spec: BatchedTrsmSpec,
2660        l_stack: &CudaSlice<f64>,
2661        rhs_stack: &mut CudaSlice<f64>,
2662    ) -> Result<(), ArrowSchurGpuFailure> {
2663        let BatchedTrsmSpec {
2664            d,
2665            n,
2666            nrhs,
2667            rhs_leading_dim,
2668            rhs_block_stride,
2669            transposed,
2670        } = spec;
2671        let alpha = 1.0_f64;
2672        let d_i = to_i32(d).ok_or(ArrowSchurGpuFailure::Unavailable)?;
2673        let nrhs_i = to_i32(nrhs).ok_or(ArrowSchurGpuFailure::Unavailable)?;
2674        let rhs_ld_i = to_i32(rhs_leading_dim).ok_or(ArrowSchurGpuFailure::Unavailable)?;
2675        let batch_i = to_i32(n).ok_or(ArrowSchurGpuFailure::Unavailable)?;
2676        let l_bytes_per = (d * d * std::mem::size_of::<f64>()) as u64;
2677        let rhs_bytes_per = (rhs_block_stride * std::mem::size_of::<f64>()) as u64;
2678        let (l_base, _l_record) = l_stack.device_ptr(stream);
2679        let (rhs_base, _rhs_record) = rhs_stack.device_ptr_mut(stream);
2680        let mut l_ptrs = Vec::with_capacity(n);
2681        let mut rhs_ptrs = Vec::with_capacity(n);
2682        for i in 0..n {
2683            l_ptrs.push(l_base + (i as u64) * l_bytes_per);
2684            rhs_ptrs.push(rhs_base + (i as u64) * rhs_bytes_per);
2685        }
2686        let mut l_ptrs_dev = stream
2687            .clone_htod(&l_ptrs)
2688            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2689        let mut rhs_ptrs_dev = stream
2690            .clone_htod(&rhs_ptrs)
2691            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2692        let (l_ptrs_ptr, _l_ptrs_rec) = l_ptrs_dev.device_ptr_mut(stream);
2693        let (rhs_ptrs_ptr, _rhs_ptrs_rec) = rhs_ptrs_dev.device_ptr_mut(stream);
2694        let op = if transposed {
2695            cublasOperation_t::CUBLAS_OP_T
2696        } else {
2697            cublasOperation_t::CUBLAS_OP_N
2698        };
2699        let handle = *blas.handle();
2700        // SAFETY: pointer arrays and base buffers were just constructed from
2701        // live device allocations covering the entire batch.
2702        let status = unsafe {
2703            cudarc::cublas::sys::cublasDtrsmBatched(
2704                handle,
2705                cublasSideMode_t::CUBLAS_SIDE_LEFT,
2706                cublasFillMode_t::CUBLAS_FILL_MODE_LOWER,
2707                op,
2708                cublasDiagType_t::CUBLAS_DIAG_NON_UNIT,
2709                d_i,
2710                nrhs_i,
2711                &alpha,
2712                l_ptrs_ptr as *const *const f64,
2713                d_i,
2714                rhs_ptrs_ptr as *const *mut f64,
2715                rhs_ld_i,
2716                batch_i,
2717            )
2718        };
2719        if status != cublasStatus_t::CUBLAS_STATUS_SUCCESS {
2720            return Err(ArrowSchurGpuFailure::Unavailable);
2721        }
2722        Ok(())
2723    }
2724
2725    /// Single-matrix lower-triangular solve: `rhs ← L^{-1} rhs` (or
2726    /// `L^{-T} rhs` if `transposed`). For the Schur Cholesky back-sub.
2727    fn trsm_single(
2728        blas: &CudaBlas,
2729        stream: &Arc<CudaStream>,
2730        n: usize,
2731        l: &CudaSlice<f64>,
2732        rhs: &mut CudaSlice<f64>,
2733        upper: bool,
2734        transposed: bool,
2735    ) -> Result<(), ArrowSchurGpuFailure> {
2736        let alpha = 1.0_f64;
2737        let n_i = to_i32(n).ok_or(ArrowSchurGpuFailure::Unavailable)?;
2738        let handle = *blas.handle();
2739        let (l_ptr, _l_rec) = l.device_ptr(stream);
2740        let (rhs_ptr, _rhs_rec) = rhs.device_ptr_mut(stream);
2741        // SAFETY: single n×n lower factor and n-vector RHS on device.
2742        let status = unsafe {
2743            cudarc::cublas::sys::cublasDtrsm_v2(
2744                handle,
2745                cublasSideMode_t::CUBLAS_SIDE_LEFT,
2746                if upper {
2747                    cublasFillMode_t::CUBLAS_FILL_MODE_UPPER
2748                } else {
2749                    cublasFillMode_t::CUBLAS_FILL_MODE_LOWER
2750                },
2751                if transposed {
2752                    cublasOperation_t::CUBLAS_OP_T
2753                } else {
2754                    cublasOperation_t::CUBLAS_OP_N
2755                },
2756                cublasDiagType_t::CUBLAS_DIAG_NON_UNIT,
2757                n_i,
2758                1,
2759                &alpha,
2760                l_ptr as *const f64,
2761                n_i,
2762                rhs_ptr as *mut f64,
2763                n_i,
2764            )
2765        };
2766        if status != cublasStatus_t::CUBLAS_STATUS_SUCCESS {
2767            return Err(ArrowSchurGpuFailure::Unavailable);
2768        }
2769        Ok(())
2770    }
2771
2772    /// The GRADIENT-INDEPENDENT half of [`accumulate_schur`]: `schur ← schur −
2773    /// Σ_i Y_i^T Y_i`, one GEMM per block in ascending block order.
2774    ///
2775    /// Splitting the fused loop is a pure regrouping: the GEMM accumulates into
2776    /// `schur` and the GEMV of [`accumulate_schur_rhs_only`] accumulates into a
2777    /// disjoint `rhs`, both in ascending `i` on the same stream, so running the
2778    /// two loops separately produces bit-identical buffers to running them
2779    /// interleaved. That is what lets a ridge-derived factor be cached and
2780    /// re-solved against a fresh gradient without changing any number.
2781    fn accumulate_schur_reduce_only(
2782        blas: &CudaBlas,
2783        d: usize,
2784        k: usize,
2785        n: usize,
2786        y_stack: &CudaSlice<f64>,
2787        schur: &mut CudaSlice<f64>,
2788    ) -> Result<(), ArrowSchurGpuFailure> {
2789        let y_block_elems = d * k;
2790        for i in 0..n {
2791            let y_slice = y_stack.slice(i * y_block_elems..(i + 1) * y_block_elems);
2792            let gemm_cfg = GemmConfig::<f64> {
2793                transa: cublasOperation_t::CUBLAS_OP_T,
2794                transb: cublasOperation_t::CUBLAS_OP_N,
2795                m: to_i32(k).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2796                n: to_i32(k).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2797                k: to_i32(d).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2798                alpha: -1.0,
2799                lda: to_i32(d).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2800                ldb: to_i32(d).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2801                beta: 1.0,
2802                ldc: to_i32(k).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2803            };
2804            // SAFETY: y_slice is d×k col-major, schur is k×k col-major; alpha/beta scalars set above.
2805            unsafe { blas.gemm(gemm_cfg, &y_slice, &y_slice, schur) }
2806                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2807        }
2808        Ok(())
2809    }
2810
2811    /// The GRADIENT-DEPENDENT half of [`accumulate_schur`]: `rhs ← rhs + Σ_i
2812    /// Y_i^T u_i`, one GEMV per block in ascending block order. See
2813    /// [`accumulate_schur_reduce_only`] for why the split is exact.
2814    fn accumulate_schur_rhs_only(
2815        blas: &CudaBlas,
2816        d: usize,
2817        k: usize,
2818        n: usize,
2819        y_stack: &CudaSlice<f64>,
2820        u_stack: &CudaSlice<f64>,
2821        rhs: &mut CudaSlice<f64>,
2822    ) -> Result<(), ArrowSchurGpuFailure> {
2823        let y_block_elems = d * k;
2824        let u_block_elems = d;
2825        for i in 0..n {
2826            let y_slice = y_stack.slice(i * y_block_elems..(i + 1) * y_block_elems);
2827            let u_slice = u_stack.slice(i * u_block_elems..(i + 1) * u_block_elems);
2828            let gemv_cfg = GemvConfig::<f64> {
2829                trans: cublasOperation_t::CUBLAS_OP_T,
2830                m: to_i32(d).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2831                n: to_i32(k).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2832                alpha: 1.0,
2833                lda: to_i32(d).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2834                incx: 1,
2835                beta: 1.0,
2836                incy: 1,
2837            };
2838            // SAFETY: y_slice (d×k col-major) and u_slice (length d) are live
2839            // device buffers; `rhs` is the length-k accumulator.
2840            unsafe { blas.gemv(gemv_cfg, &y_slice, &u_slice, rhs) }
2841                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2842        }
2843        Ok(())
2844    }
2845
2846    /// Accumulate `schur ← schur − Σ_i Y_i^T Y_i` and `rhs ← rhs + Σ_i Y_i^T u_i`
2847    /// using one GEMM and one GEMV per block. Each call uses beta=1 to chain
2848    /// the accumulation device-side.
2849    fn accumulate_schur(
2850        blas: &CudaBlas,
2851        d: usize,
2852        k: usize,
2853        n: usize,
2854        y_stack: &CudaSlice<f64>,
2855        u_stack: &CudaSlice<f64>,
2856        schur: &mut CudaSlice<f64>,
2857        rhs: &mut CudaSlice<f64>,
2858    ) -> Result<(), ArrowSchurGpuFailure> {
2859        let y_block_elems = d * k;
2860        let u_block_elems = d;
2861        for i in 0..n {
2862            let y_slice = y_stack.slice(i * y_block_elems..(i + 1) * y_block_elems);
2863            let u_slice = u_stack.slice(i * u_block_elems..(i + 1) * u_block_elems);
2864            // GEMM: schur += (-1) · Y_i^T · Y_i  (Y_i is d×k col-major; out is k×k)
2865            let gemm_cfg = GemmConfig::<f64> {
2866                transa: cublasOperation_t::CUBLAS_OP_T,
2867                transb: cublasOperation_t::CUBLAS_OP_N,
2868                m: to_i32(k).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2869                n: to_i32(k).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2870                k: to_i32(d).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2871                alpha: -1.0,
2872                lda: to_i32(d).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2873                ldb: to_i32(d).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2874                beta: 1.0,
2875                ldc: to_i32(k).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2876            };
2877            // SAFETY: y_slice is d×k col-major, schur is k×k col-major; alpha/beta scalars set above.
2878            unsafe { blas.gemm(gemm_cfg, &y_slice, &y_slice, schur) }
2879                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2880            // GEMV: rhs += 1 · Y_i^T · u_i
2881            let gemv_cfg = GemvConfig::<f64> {
2882                trans: cublasOperation_t::CUBLAS_OP_T,
2883                m: to_i32(d).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2884                n: to_i32(k).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2885                alpha: 1.0,
2886                lda: to_i32(d).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2887                incx: 1,
2888                beta: 1.0,
2889                incy: 1,
2890            };
2891            // SAFETY: y_slice (d×k col-major) and u_slice (length d) are live
2892            // device buffers; `rhs` is the length-k accumulator.
2893            unsafe { blas.gemv(gemv_cfg, &y_slice, &u_slice, rhs) }
2894                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2895        }
2896        Ok(())
2897    }
2898
2899    /// `#2393` single-launch Schur reduction `S ← S − Yᵀ Y` for a STACKED
2900    /// whitened `Y` (`(n·d) × k` column-major, leading dimension `n·d`).
2901    ///
2902    /// [`accumulate_schur`] performs the SAME arithmetic as n rank-`d` updates,
2903    /// one cuBLAS launch per row block. At the SAE LLM shape (n≈2000, d=2,
2904    /// k≈2048) that is 2000 launches whose individual arithmetic (≈17 MFLOP)
2905    /// never fills the device, and each one reads AND writes the whole `k×k`
2906    /// accumulator — ≈67 GB of device traffic for 33 GFLOP of work. Stacking the
2907    /// row blocks into one matrix turns the whole reduction into a single
2908    /// `k × k × (n·d)` GEMM: same flops, the accumulator touched once, and the
2909    /// device saturated for the duration instead of drip-fed 2000 times.
2910    ///
2911    /// The reduction ORDER changes (cuBLAS's internal split rather than
2912    /// ascending row blocks), so results differ from [`accumulate_schur`] by
2913    /// float reassociation only.
2914    fn schur_gemm_stacked(
2915        blas: &CudaBlas,
2916        rows: usize,
2917        k: usize,
2918        y_stacked: &CudaSlice<f64>,
2919        schur: &mut CudaSlice<f64>,
2920    ) -> Result<(), ArrowSchurGpuFailure> {
2921        let k_i = to_i32(k).ok_or(ArrowSchurGpuFailure::Unavailable)?;
2922        let rows_i = to_i32(rows).ok_or(ArrowSchurGpuFailure::Unavailable)?;
2923        let cfg = GemmConfig::<f64> {
2924            transa: cublasOperation_t::CUBLAS_OP_T,
2925            transb: cublasOperation_t::CUBLAS_OP_N,
2926            m: k_i,
2927            n: k_i,
2928            k: rows_i,
2929            alpha: -1.0,
2930            lda: rows_i,
2931            ldb: rows_i,
2932            beta: 1.0,
2933            ldc: k_i,
2934        };
2935        // SAFETY: `y_stacked` is the live `(n·d)×k` column-major whitened block
2936        // stack with leading dimension `n·d`; `schur` is the live `k×k`
2937        // column-major accumulator with leading dimension `k`.
2938        unsafe { blas.gemm(cfg, y_stacked, y_stacked, schur) }
2939            .map_err(|_| ArrowSchurGpuFailure::Unavailable)
2940    }
2941
2942    /// `#2393` single-launch Schur RHS accumulation `rhs += Yᵀ u` for the same
2943    /// stacked `Y` — the counterpart of [`accumulate_schur_rhs_only`]'s n GEMV
2944    /// launches, and the per-iterate half of the residency win (it runs on
2945    /// EVERY resident solve, not just at frame build).
2946    fn schur_rhs_stacked(
2947        blas: &CudaBlas,
2948        rows: usize,
2949        k: usize,
2950        y_stacked: &CudaSlice<f64>,
2951        u_stacked: &CudaSlice<f64>,
2952        rhs: &mut CudaSlice<f64>,
2953    ) -> Result<(), ArrowSchurGpuFailure> {
2954        let k_i = to_i32(k).ok_or(ArrowSchurGpuFailure::Unavailable)?;
2955        let rows_i = to_i32(rows).ok_or(ArrowSchurGpuFailure::Unavailable)?;
2956        let cfg = GemvConfig::<f64> {
2957            trans: cublasOperation_t::CUBLAS_OP_T,
2958            m: rows_i,
2959            n: k_i,
2960            alpha: 1.0,
2961            lda: rows_i,
2962            incx: 1,
2963            beta: 1.0,
2964            incy: 1,
2965        };
2966        // SAFETY: `y_stacked` is `(n·d)×k` column-major with leading dimension
2967        // `n·d`; under OP_T the operand `u_stacked` has `n·d` entries and the
2968        // accumulator `rhs` has `k`, both unit-stride and live.
2969        unsafe { blas.gemv(cfg, y_stacked, u_stacked, rhs) }
2970            .map_err(|_| ArrowSchurGpuFailure::Unavailable)
2971    }
2972
2973    /// `#2393` single-launch back-substitution accumulation `u += Y δβ` for the
2974    /// stacked `Y` — the counterpart of [`accumulate_back_sub_rhs`]'s n GEMV
2975    /// launches.
2976    fn back_sub_rhs_stacked(
2977        blas: &CudaBlas,
2978        rows: usize,
2979        k: usize,
2980        y_stacked: &CudaSlice<f64>,
2981        delta_beta: &CudaSlice<f64>,
2982        u_stacked: &mut CudaSlice<f64>,
2983    ) -> Result<(), ArrowSchurGpuFailure> {
2984        let k_i = to_i32(k).ok_or(ArrowSchurGpuFailure::Unavailable)?;
2985        let rows_i = to_i32(rows).ok_or(ArrowSchurGpuFailure::Unavailable)?;
2986        let cfg = GemvConfig::<f64> {
2987            trans: cublasOperation_t::CUBLAS_OP_N,
2988            m: rows_i,
2989            n: k_i,
2990            alpha: 1.0,
2991            lda: rows_i,
2992            incx: 1,
2993            beta: 1.0,
2994            incy: 1,
2995        };
2996        // SAFETY: `y_stacked` is `(n·d)×k` column-major with leading dimension
2997        // `n·d`; under OP_N the operand `delta_beta` has `k` entries and the
2998        // accumulator `u_stacked` has `n·d`, both unit-stride and live.
2999        unsafe { blas.gemv(cfg, y_stacked, delta_beta, u_stacked) }
3000            .map_err(|_| ArrowSchurGpuFailure::Unavailable)
3001    }
3002
3003    /// Accumulate `g_dev[i] ← u_i + Y_i · δβ` per block. This is the
3004    /// pre-trsm RHS for the back-substitution `L_i^T x_i = w_i`.
3005    fn accumulate_back_sub_rhs(
3006        blas: &CudaBlas,
3007        d: usize,
3008        k: usize,
3009        n: usize,
3010        y_stack: &CudaSlice<f64>,
3011        delta_beta: &CudaSlice<f64>,
3012        u_stack: &mut CudaSlice<f64>,
3013    ) -> Result<(), ArrowSchurGpuFailure> {
3014        let y_block_elems = d * k;
3015        let u_block_elems = d;
3016        for i in 0..n {
3017            let y_slice = y_stack.slice(i * y_block_elems..(i + 1) * y_block_elems);
3018            let mut u_slice = u_stack.slice_mut(i * u_block_elems..(i + 1) * u_block_elems);
3019            let gemv_cfg = GemvConfig::<f64> {
3020                trans: cublasOperation_t::CUBLAS_OP_N,
3021                m: to_i32(d).ok_or(ArrowSchurGpuFailure::Unavailable)?,
3022                n: to_i32(k).ok_or(ArrowSchurGpuFailure::Unavailable)?,
3023                alpha: 1.0,
3024                lda: to_i32(d).ok_or(ArrowSchurGpuFailure::Unavailable)?,
3025                incx: 1,
3026                beta: 1.0,
3027                incy: 1,
3028            };
3029            // SAFETY: y_slice / delta_beta / u_slice are live device buffers
3030            // of the expected sizes (d×k, k, d).
3031            unsafe { blas.gemv(gemv_cfg, &y_slice, delta_beta, &mut u_slice) }
3032                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
3033        }
3034        Ok(())
3035    }
3036
3037    // ────────────────────────────────────────────────────────────────────
3038    // Layer D + E — fused NVRTC dispatch.
3039    //
3040    // The forward kernel (`arrow_schur_forward_pgroup`) is a single launch
3041    // that, per row block, factors `D_i + ρI = L_i L_iᵀ` in shared memory,
3042    // forward-solves `u_i = L_i⁻¹ g_i` and `Y_i = L_i⁻¹ B_i`, and emits the
3043    // per-block Schur partials `partial_S[i] = Yᵀ Y` (R×R) and
3044    // `partial_r[i] = Yᵀ u` (R). The host reduces partials on the CPU after
3045    // dtoh (one fused sum across `n` blocks of R²+R doubles; cheap because
3046    // n·R² ≲ 5M doubles at large scale), assembles `S_β`, factors it via
3047    // cuSOLVER, and launches the back-substitution kernel
3048    // `arrow_schur_back_sub_pgroup` to recover `δt_i = -L_i⁻ᵀ(u_i + Y_i δβ)`
3049    // without re-uploading the local factors.
3050    // ────────────────────────────────────────────────────────────────────
3051
3052    use std::collections::HashMap;
3053    use std::sync::Mutex;
3054
3055    /// One compiled NVRTC module per `(cc_major, cc_minor, p_max, r_template)`.
3056    /// `cc_*` lets one process drive multiple device generations; the
3057    /// `(p_max, r_template)` pair selects the shared-memory layout baked into
3058    /// the kernel source.
3059    struct FusedModuleCache {
3060        modules: Mutex<
3061            HashMap<crate::gpu_kernels::arrow_schur_nvrtc::FusedModuleCacheKey, Arc<CudaModule>>,
3062        >,
3063    }
3064
3065    fn fused_module_cache() -> &'static FusedModuleCache {
3066        static CACHE: OnceLock<FusedModuleCache> = OnceLock::new();
3067        CACHE.get_or_init(|| FusedModuleCache {
3068            modules: Mutex::new(HashMap::new()),
3069        })
3070    }
3071
3072    fn fused_module_for(
3073        ctx: &Arc<CudaContext>,
3074        key: crate::gpu_kernels::arrow_schur_nvrtc::FusedModuleCacheKey,
3075    ) -> Result<Arc<CudaModule>, ArrowSchurGpuFailure> {
3076        let cache = fused_module_cache();
3077        if let Ok(guard) = cache.modules.lock() {
3078            if let Some(existing) = guard.get(&key) {
3079                return Ok(existing.clone());
3080            }
3081        }
3082        let src = crate::gpu_kernels::arrow_schur_nvrtc::forward_kernel_source(
3083            key.p_max as usize,
3084            key.r_template as usize,
3085        );
3086        let ptx = gam_gpu::device_cache::compile_ptx_arch(&src).map_err(|err| {
3087            ArrowSchurGpuFailure::SchurFactorFailed {
3088                reason: format!(
3089                    "arrow-schur fused NVRTC compile (p_max={}, r={}): {err}",
3090                    key.p_max, key.r_template
3091                ),
3092            }
3093        })?;
3094        let module = ctx
3095            .load_module(ptx)
3096            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
3097        if let Ok(mut guard) = cache.modules.lock() {
3098            guard.entry(key).or_insert_with(|| module.clone());
3099        }
3100        Ok(module)
3101    }
3102
3103    const PCG_VECTOR_KERNEL_SOURCE: &str = r#"
3104extern "C" __global__ void arrow_pcg_jacobi_mul(
3105    const double* __restrict__ inv_diag,
3106    const double* __restrict__ r,
3107    double* __restrict__ z,
3108    int n
3109) {
3110    int idx = blockIdx.x * blockDim.x + threadIdx.x;
3111    if (idx < n) {
3112        z[idx] = inv_diag[idx] * r[idx];
3113    }
3114}
3115
3116extern "C" __global__ void arrow_pcg_update_p(
3117    const double* __restrict__ z,
3118    double beta,
3119    double* __restrict__ p,
3120    int n
3121) {
3122    int idx = blockIdx.x * blockDim.x + threadIdx.x;
3123    if (idx < n) {
3124        p[idx] = z[idx] + beta * p[idx];
3125    }
3126}
3127
3128extern "C" __global__ void arrow_sae_init(
3129    double* __restrict__ out,
3130    const double* __restrict__ x,
3131    double ridge,
3132    int n
3133) {
3134    int idx = blockIdx.x * blockDim.x + threadIdx.x;
3135    if (idx < n) {
3136        out[idx] = ridge * x[idx];
3137    }
3138}
3139
3140extern "C" __global__ void arrow_sae_smooth_matvec(
3141    const double* __restrict__ x,
3142    double* __restrict__ out,
3143    const int* __restrict__ block_offsets,
3144    const int* __restrict__ block_m,
3145    const int* __restrict__ factor_ptr,
3146    const double* __restrict__ factors,
3147    int p,
3148    int n_blocks
3149) {
3150    int block_id = blockIdx.y;
3151    int linear = blockIdx.x * blockDim.x + threadIdx.x;
3152    if (block_id >= n_blocks) {
3153        return;
3154    }
3155    int m = block_m[block_id];
3156    int total = m * p;
3157    if (linear >= total) {
3158        return;
3159    }
3160    int li = linear / p;
3161    int oc = linear - li * p;
3162    int off = block_offsets[block_id];
3163    int fbase = factor_ptr[block_id];
3164    double acc = 0.0;
3165    for (int lj = 0; lj < m; ++lj) {
3166        double a = factors[fbase + li * m + lj];
3167        acc += a * x[off + lj * p + oc];
3168    }
3169    out[off + li * p + oc] += acc;
3170}
3171
3172extern "C" __global__ void arrow_sae_sparse_g_matvec(
3173    const double* __restrict__ x,
3174    double* __restrict__ out,
3175    const int* __restrict__ col_off,
3176    const int* __restrict__ cols,
3177    const int* __restrict__ data_ptr,
3178    const double* __restrict__ data,
3179    const int* __restrict__ group_ptr,
3180    const int* __restrict__ group_blocks,
3181    const int* __restrict__ group_off,
3182    const int* __restrict__ group_width,
3183    int p,
3184    int n_groups
3185) {
3186    int group = blockIdx.y;
3187    int linear = blockIdx.x * blockDim.x + threadIdx.x;
3188    if (group >= n_groups || linear >= group_width[group]) {
3189        return;
3190    }
3191    int out_index = group_off[group] + linear;
3192    int li = linear / p;
3193    int oc = linear - li * p;
3194    double acc = out[out_index];
3195    // #2535 — one thread OWNS this scalar output and walks every contributing
3196    // G block in original host order. Equal output spans were grouped into the
3197    // CSR above; distinct spans are disjoint. The scheduler therefore chooses
3198    // only when an independent output runs, never the order of its additions.
3199    for (int position = group_ptr[group]; position < group_ptr[group + 1]; ++position) {
3200        int block_id = group_blocks[position];
3201        int m_j = cols[block_id];
3202        int cbase = col_off[block_id];
3203        int dbase = data_ptr[block_id];
3204        for (int lj = 0; lj < m_j; ++lj) {
3205            acc += data[dbase + li * m_j + lj] * x[(cbase + lj) * p + oc];
3206        }
3207    }
3208    out[out_index] = acc;
3209}
3210
3211extern "C" __global__ void arrow_sae_gather_u(
3212    const double* __restrict__ x,
3213    const int* __restrict__ row_ptr,
3214    const int* __restrict__ beta_base,
3215    const double* __restrict__ phi,
3216    double* __restrict__ u,
3217    int p,
3218    int n_rows
3219) {
3220    int row = blockIdx.y;
3221    int oc = blockIdx.x * blockDim.x + threadIdx.x;
3222    if (row >= n_rows || oc >= p) {
3223        return;
3224    }
3225    double acc = 0.0;
3226    int start = row_ptr[row];
3227    int end = row_ptr[row + 1];
3228    for (int e = start; e < end; ++e) {
3229        acc += phi[e] * x[beta_base[e] + oc];
3230    }
3231    u[row * p + oc] = acc;
3232}
3233
3234extern "C" __global__ void arrow_sae_apply_l(
3235    const double* __restrict__ u,
3236    const int* __restrict__ jac_ptr,
3237    const double* __restrict__ jac,
3238    double* __restrict__ w,
3239    int p,
3240    int max_q,
3241    int n_rows
3242) {
3243    int row = blockIdx.y;
3244    int c = blockIdx.x * blockDim.x + threadIdx.x;
3245    if (row >= n_rows) {
3246        return;
3247    }
3248    int jstart = jac_ptr[row];
3249    int q = (jac_ptr[row + 1] - jstart) / p;
3250    if (c >= q) {
3251        return;
3252    }
3253    double acc = 0.0;
3254    for (int oc = 0; oc < p; ++oc) {
3255        acc += jac[jstart + c * p + oc] * u[row * p + oc];
3256    }
3257    w[row * max_q + c] = acc;
3258}
3259
3260extern "C" __global__ void arrow_sae_apply_ainv(
3261    const double* __restrict__ ainv,
3262    const double* __restrict__ w,
3263    double* __restrict__ v,
3264    int max_q,
3265    int n_rows
3266) {
3267    int row = blockIdx.y;
3268    int c = blockIdx.x * blockDim.x + threadIdx.x;
3269    if (row >= n_rows || c >= max_q) {
3270        return;
3271    }
3272    double acc = 0.0;
3273    int base = row * max_q * max_q;
3274    for (int j = 0; j < max_q; ++j) {
3275        acc += ainv[base + c * max_q + j] * w[row * max_q + j];
3276    }
3277    v[row * max_q + c] = acc;
3278}
3279
3280/* #2512 DETERMINISTIC legacy row-Schur scatter, STAGE 1 (partials).
3281
3282   `arrow_sae_scatter_sub` above ran one block per row and folded the rows into
3283   `out` with `atomicAdd`, so the summation order was the block scheduler's and
3284   the matvec returned a different value on every launch. Here thread `(oc,
3285   chunk)` owns output channel `oc` over the contiguous row range
3286   [chunk·rows_per_chunk, …) and walks its rows in fixed index order, so the
3287   reassociation is a property of the shape rather than of the schedule.
3288
3289   The thread may write `partial[cbase + beta_base[e] + oc]` for several `e`
3290   without colliding with any other thread because every `beta_base` is a
3291   multiple of `p`: the host builds the support as
3292   `a_phi.push((atom_beta_off + basis_col * p, w))` and each atom's border block
3293   is `M_k · p` wide, so the entries a thread touches are exactly the indices
3294   congruent to `oc` modulo `p`. That is also why the thread can zero its own
3295   progression first and needs no separate clearing pass or barrier. */
3296extern "C" __global__ void arrow_sae_scatter_sub_det_partial(
3297    const double* __restrict__ v,
3298    const int* __restrict__ jac_ptr,
3299    const double* __restrict__ jac,
3300    const int* __restrict__ row_ptr,
3301    const int* __restrict__ beta_base,
3302    const double* __restrict__ phi,
3303    double* __restrict__ partial,
3304    int p,
3305    int k,
3306    int max_q,
3307    int n_rows,
3308    int rows_per_chunk
3309) {
3310    int oc = blockIdx.x * blockDim.x + threadIdx.x;
3311    if (oc >= p) { return; }
3312    int chunk = blockIdx.y;
3313    long long cbase = (long long)chunk * k;
3314    for (int i = oc; i < k; i += p) {
3315        partial[cbase + i] = 0.0;
3316    }
3317    int row0 = chunk * rows_per_chunk;
3318    int row1 = row0 + rows_per_chunk;
3319    if (row1 > n_rows) { row1 = n_rows; }
3320    for (int row = row0; row < row1; ++row) {
3321        int jstart = jac_ptr[row];
3322        int q = (jac_ptr[row + 1] - jstart) / p;
3323        double lt_v = 0.0;
3324        for (int c = 0; c < q; ++c) {
3325            lt_v += jac[jstart + c * p + oc] * v[row * max_q + c];
3326        }
3327        int start = row_ptr[row];
3328        int end = row_ptr[row + 1];
3329        for (int e = start; e < end; ++e) {
3330            partial[cbase + beta_base[e] + oc] -= phi[e] * lt_v;
3331        }
3332    }
3333}
3334
3335/* #2512 STAGE 2 (reduce): out[a] += Σ_chunk partial[chunk][a], chunks summed in
3336   fixed order 0..n_chunks. One thread per output coord `a`, so no two threads
3337   touch the same `out[a]`. */
3338extern "C" __global__ void arrow_sae_scatter_sub_det_reduce(
3339    const double* __restrict__ partial,
3340    double* __restrict__ out,
3341    int k,
3342    int n_chunks
3343) {
3344    int a = blockIdx.x * blockDim.x + threadIdx.x;
3345    if (a >= k) { return; }
3346    double acc = 0.0;
3347    for (int chunk = 0; chunk < n_chunks; ++chunk) {
3348        acc += partial[(long long)chunk * k + a];
3349    }
3350    out[a] += acc;
3351}
3352
3353/* #2512 DETERMINISTIC Jacobi-diagonal subtraction, STAGE 1 (partials).
3354
3355   The former kernel ran one block per row and folded the rows into `diag` with
3356   `atomicAdd`, so the PRECONDITIONER was a fresh draw per call — and a moved
3357   preconditioner moves the whole CG trajectory, not merely the last bits of one
3358   matvec. Thread `(oc, chunk)` owns output channel `oc` over a contiguous row
3359   range and walks it in fixed index order; as in the scatter, every `beta_base`
3360   is a multiple of `p`, so the indices a thread touches are exactly those
3361   congruent to `oc` modulo `p` and it can clear its own progression first.
3362   Stage 2 is the shared `arrow_sae_scatter_sub_det_reduce`. */
3363extern "C" __global__ void arrow_sae_diag_sub_det_partial(
3364    const double* __restrict__ ainv,
3365    const int* __restrict__ jac_ptr,
3366    const double* __restrict__ jac,
3367    const int* __restrict__ row_ptr,
3368    const int* __restrict__ beta_base,
3369    const double* __restrict__ phi,
3370    double* __restrict__ partial,
3371    int p,
3372    int k,
3373    int max_q,
3374    int n_rows,
3375    int rows_per_chunk
3376) {
3377    int oc = blockIdx.x * blockDim.x + threadIdx.x;
3378    if (oc >= p) { return; }
3379    int chunk = blockIdx.y;
3380    long long cbase = (long long)chunk * k;
3381    for (int i = oc; i < k; i += p) {
3382        partial[cbase + i] = 0.0;
3383    }
3384    int row0 = chunk * rows_per_chunk;
3385    int row1 = row0 + rows_per_chunk;
3386    if (row1 > n_rows) { row1 = n_rows; }
3387    for (int row = row0; row < row1; ++row) {
3388        int jstart = jac_ptr[row];
3389        int q = (jac_ptr[row + 1] - jstart) / p;
3390        int abase = row * max_q * max_q;
3391        double quad = 0.0;
3392        for (int c = 0; c < q; ++c) {
3393            double lc = jac[jstart + c * p + oc];
3394            for (int d = 0; d < q; ++d) {
3395                quad += lc * ainv[abase + c * max_q + d] * jac[jstart + d * p + oc];
3396            }
3397        }
3398        int start = row_ptr[row];
3399        int end = row_ptr[row + 1];
3400        for (int e = start; e < end; ++e) {
3401            double pe = phi[e];
3402            partial[cbase + beta_base[e] + oc] -= (pe * pe) * quad;
3403        }
3404    }
3405}
3406
3407/* ── #1017/#1026 frames-engaged device kernels ─────────────────────────────
3408 * The factored β border is C-space (width Σ M_k·r_k). The penalty side is the
3409 * smooth `λ S_k ⊗ I_{r_k}` (per-block right-width r_k) plus the data-fit
3410 * `G_{ij} ⊗ W_{ij}` (W = U_iᵀU_j, dense r_i×r_j). The reduced-Schur term uses
3411 * the per-row DENSE cross-block H_tβ^(i) (q_i × border_dim, row-major). */
3412
3413extern "C" __global__ void arrow_sae_frame_smooth_matvec(
3414    const double* __restrict__ x,
3415    double* __restrict__ out,
3416    const int* __restrict__ block_offsets,
3417    const int* __restrict__ block_m,
3418    const int* __restrict__ block_r,
3419    const int* __restrict__ factor_ptr,
3420    const double* __restrict__ factors,
3421    int n_blocks
3422) {
3423    int block_id = blockIdx.y;
3424    int linear = blockIdx.x * blockDim.x + threadIdx.x;
3425    if (block_id >= n_blocks) {
3426        return;
3427    }
3428    int m = block_m[block_id];
3429    int r = block_r[block_id];
3430    int total = m * r;
3431    if (linear >= total) {
3432        return;
3433    }
3434    int li = linear / r;
3435    int ib = linear - li * r;
3436    int off = block_offsets[block_id];
3437    int fbase = factor_ptr[block_id];
3438    double acc = 0.0;
3439    for (int lj = 0; lj < m; ++lj) {
3440        double a = factors[fbase + li * m + lj];
3441        acc += a * x[off + lj * r + ib];
3442    }
3443    out[off + li * r + ib] += acc;
3444}
3445
3446extern "C" __global__ void arrow_sae_frame_g_matvec(
3447    const double* __restrict__ x,
3448    double* __restrict__ out,
3449    const int* __restrict__ off_i,
3450    const int* __restrict__ off_j,
3451    const int* __restrict__ r_i,
3452    const int* __restrict__ r_j,
3453    const int* __restrict__ m_j,
3454    const int* __restrict__ g_ptr,
3455    const double* __restrict__ g_data,
3456    const int* __restrict__ w_ptr,
3457    const double* __restrict__ w_data,
3458    const int* __restrict__ group_ptr,
3459    const int* __restrict__ group_blocks,
3460    const int* __restrict__ group_off,
3461    const int* __restrict__ group_width,
3462    int n_groups
3463) {
3464    int group = blockIdx.y;
3465    int linear = blockIdx.x * blockDim.x + threadIdx.x;
3466    if (group >= n_groups || linear >= group_width[group]) {
3467        return;
3468    }
3469    int out_index = group_off[group] + linear;
3470    double acc = out[out_index];
3471    // #2535 — fixed-order G⊗W accumulation. The owning output thread walks the
3472    // original block subsequence and, within each block, the original lj/b
3473    // order used by the CPU oracle. No partial buffer and no host reduction.
3474    for (int position = group_ptr[group]; position < group_ptr[group + 1]; ++position) {
3475        int block_id = group_blocks[position];
3476        int ri = r_i[block_id];
3477        int rj = r_j[block_id];
3478        int mj = m_j[block_id];
3479        int local = out_index - off_i[block_id];
3480        int li = local / ri;
3481        int a = local - li * ri;
3482        int oj = off_j[block_id];
3483        int gbase = g_ptr[block_id];
3484        int wbase = w_ptr[block_id];
3485        for (int lj = 0; lj < mj; ++lj) {
3486            double g = g_data[gbase + li * mj + lj];
3487            if (g == 0.0) { continue; }
3488            int xj_base = oj + lj * rj;
3489            double inner = 0.0;
3490            for (int b = 0; b < rj; ++b) {
3491                inner += w_data[wbase + a * rj + b] * x[xj_base + b];
3492            }
3493            acc += g * inner;
3494        }
3495    }
3496    out[out_index] = acc;
3497}
3498
3499/* Per-row reduced-Schur subtraction with a DENSE cross-block H_tβ^(i).
3500 *   h_i   = H_tβ^(i) · x                (length q_i)
3501 *   s_i   = (H_tt^(i)+ρ_t I)⁻¹ h_i      (apply cached ainv, length q_i)
3502 *   out  -= (H_tβ^(i))ᵀ · s_i           (scatter into border_dim)
3503 * `htb` is row-major (q_i × k) flattened, `htb_ptr` gives each row's base and
3504 * (htb_ptr[row+1]-htb_ptr[row])/k == q_i. `q_of` carries q_i directly. */
3505extern "C" __global__ void arrow_sae_frame_apply_h(
3506    const double* __restrict__ x,
3507    const int* __restrict__ htb_ptr,
3508    const double* __restrict__ htb,
3509    const int* __restrict__ q_of,
3510    double* __restrict__ hvec,
3511    int k,
3512    int max_q,
3513    int n_rows
3514) {
3515    int row = blockIdx.y;
3516    int c = blockIdx.x * blockDim.x + threadIdx.x;
3517    if (row >= n_rows) { return; }
3518    int q = q_of[row];
3519    if (c >= q) { return; }
3520    int base = htb_ptr[row] + c * k;
3521    double acc = 0.0;
3522    for (int a = 0; a < k; ++a) {
3523        acc += htb[base + a] * x[a];
3524    }
3525    hvec[row * max_q + c] = acc;
3526}
3527
3528extern "C" __global__ void arrow_sae_frame_apply_ainv(
3529    const double* __restrict__ ainv,
3530    const double* __restrict__ hvec,
3531    const int* __restrict__ q_of,
3532    double* __restrict__ svec,
3533    int max_q,
3534    int n_rows
3535) {
3536    int row = blockIdx.y;
3537    int c = blockIdx.x * blockDim.x + threadIdx.x;
3538    if (row >= n_rows || c >= max_q) { return; }
3539    int q = q_of[row];
3540    double acc = 0.0;
3541    int abase = row * max_q * max_q;
3542    for (int j = 0; j < q; ++j) {
3543        acc += ainv[abase + c * max_q + j] * hvec[row * max_q + j];
3544    }
3545    svec[row * max_q + c] = acc;
3546}
3547
3548/* #1017 evidence-lane DETERMINISTIC reduced-Schur scatter:
3549   out[a] = -Σ_i Σ_c H_tβ[i][c,a]·svec[i,c]. One thread owns output coord `a` and
3550   sums the rows in fixed index order 0..n_rows — NO atomics, so the result is
3551   run-to-run bit-stable (the SLQ log|S| determinism contract): the reduction
3552   order over rows is pinned rather than decided by the block scheduler.
3553   `out` is fully assigned (no init needed). */
3554extern "C" __global__ void arrow_sae_frame_scatter_h_det(
3555    const double* __restrict__ svec,
3556    const int* __restrict__ htb_ptr,
3557    const double* __restrict__ htb,
3558    const int* __restrict__ q_of,
3559    double* __restrict__ out,
3560    int k,
3561    int max_q,
3562    int n_rows
3563) {
3564    int a = blockIdx.x * blockDim.x + threadIdx.x;
3565    if (a >= k) { return; }
3566    double acc = 0.0;
3567    for (int row = 0; row < n_rows; ++row) {
3568        int q = q_of[row];
3569        int hbase = htb_ptr[row];
3570        int sbase = row * max_q;
3571        for (int c = 0; c < q; ++c) {
3572            acc += htb[hbase + c * k + a] * svec[sbase + c];
3573        }
3574    }
3575    out[a] = -acc;
3576}
3577
3578/* #1017 evidence-lane 2-STAGE deterministic scatter, STAGE 1 (partials):
3579   partial[chunk][a] = Σ_{row∈chunk} Σ_c H_tβ[row][c,a]·svec[row,c], for the
3580   contiguous row range [chunk·rows_per_chunk, …). grid = (⌈k/256⌉, n_chunks);
3581   thread owns (a, chunk). Replaces the single-strip `arrow_sae_frame_scatter_h_det`
3582   (⌈k/256⌉ CTAs — only 4 at k=911, one thread serial over ALL n_rows, ~94% of a
3583   72-SM A10 idle) with ⌈k/256⌉·n_chunks CTAs. Rows are summed in fixed order
3584   within the chunk and the chunks are reduced in order by stage 2, so the result
3585   is a FIXED reassociation of the same ordered row sum — run-to-run bit-stable
3586   (the SLQ log|S| determinism contract) and within the ≤1e-9 CPU-oracle gate. */
3587extern "C" __global__ void arrow_sae_frame_scatter_h_det_partial(
3588    const double* __restrict__ svec,
3589    const int* __restrict__ htb_ptr,
3590    const double* __restrict__ htb,
3591    const int* __restrict__ q_of,
3592    double* __restrict__ partial,
3593    int k,
3594    int max_q,
3595    int n_rows,
3596    int rows_per_chunk
3597) {
3598    int a = blockIdx.x * blockDim.x + threadIdx.x;
3599    if (a >= k) { return; }
3600    int chunk = blockIdx.y;
3601    int row0 = chunk * rows_per_chunk;
3602    int row1 = row0 + rows_per_chunk;
3603    if (row1 > n_rows) { row1 = n_rows; }
3604    double acc = 0.0;
3605    for (int row = row0; row < row1; ++row) {
3606        int q = q_of[row];
3607        int hbase = htb_ptr[row];
3608        int sbase = row * max_q;
3609        for (int c = 0; c < q; ++c) {
3610            acc += htb[hbase + c * k + a] * svec[sbase + c];
3611        }
3612    }
3613    partial[(long long)chunk * k + a] = acc;
3614}
3615
3616/* #2512 STAGE 2, ACCUMULATING form: out[a] -= Σ_chunk partial[chunk][a], chunks
3617   summed in fixed order 0..n_chunks. Identical reduction to
3618   `arrow_sae_frame_scatter_h_det_reduce`, but subtracts into an `out` that the
3619   penalty and data-Gram matvecs have already written, which is what the step-PCG
3620   matvec needs. One thread per output coord `a`, so no two threads touch the same
3621   `out[a]` and the result is run-to-run bit-stable. */
3622extern "C" __global__ void arrow_sae_frame_scatter_h_det_reduce_sub(
3623    const double* __restrict__ partial,
3624    double* __restrict__ out,
3625    int k,
3626    int n_chunks
3627) {
3628    int a = blockIdx.x * blockDim.x + threadIdx.x;
3629    if (a >= k) { return; }
3630    double acc = 0.0;
3631    for (int chunk = 0; chunk < n_chunks; ++chunk) {
3632        acc += partial[(long long)chunk * k + a];
3633    }
3634    out[a] -= acc;
3635}
3636
3637/* #1017 STAGE 2 (reduce): out[a] = -Σ_chunk partial[chunk][a], chunks summed in
3638   fixed order 0..n_chunks. One thread per output coord `a`; ⌈k/256⌉ CTAs. */
3639extern "C" __global__ void arrow_sae_frame_scatter_h_det_reduce(
3640    const double* __restrict__ partial,
3641    double* __restrict__ out,
3642    int k,
3643    int n_chunks
3644) {
3645    int a = blockIdx.x * blockDim.x + threadIdx.x;
3646    if (a >= k) { return; }
3647    double acc = 0.0;
3648    for (int chunk = 0; chunk < n_chunks; ++chunk) {
3649        acc += partial[(long long)chunk * k + a];
3650    }
3651    out[a] = -acc;
3652}
3653
3654/* #1017 evidence-lane WARP-COOPERATIVE apply_h: hvec[i][c] = Σ_a H_tβ[i][c,a]·x[a].
3655   One WARP owns (row, c): lane `l` strides `a = l, l+32, …` over the contiguous
3656   `H_tβ[i][c,·]` slab (fully coalesced across the warp) and a fixed-order
3657   __shfl_down tree reduces to lane 0. Replaces the shared `arrow_sae_frame_apply_h`
3658   (256-thread block, only `q_i` threads active, stride-`k` uncoalesced reads) on
3659   the evidence path ONLY — the shared kernel is untouched (step-PCG relies on it).
3660   Reduction order is fixed (lane stride + tree), so the result is run-to-run
3661   bit-stable; it differs from the scalar kernel only by ULP reassociation, within
3662   the ≤1e-9 parity gate. Launch: block = max_q·32 (≤1024), grid.x = n_rows;
3663   warp `w = threadIdx.x/32` handles `c = w` for `w < q_i`. */
3664extern "C" __global__ void arrow_sae_frame_apply_h_warp(
3665    const double* __restrict__ x,
3666    const int* __restrict__ htb_ptr,
3667    const double* __restrict__ htb,
3668    const int* __restrict__ q_of,
3669    double* __restrict__ hvec,
3670    int k,
3671    int max_q,
3672    int n_rows
3673) {
3674    int row = blockIdx.x;
3675    if (row >= n_rows) { return; }
3676    int warp = threadIdx.x >> 5;
3677    int lane = threadIdx.x & 31;
3678    int q = q_of[row];
3679    if (warp >= q) { return; }
3680    int base = htb_ptr[row] + warp * k;
3681    double acc = 0.0;
3682    for (int a = lane; a < k; a += 32) {
3683        acc += htb[base + a] * x[a];
3684    }
3685    #pragma unroll
3686    for (int off = 16; off > 0; off >>= 1) {
3687        acc += __shfl_down_sync(0xffffffffu, acc, off);
3688    }
3689    if (lane == 0) {
3690        hvec[row * max_q + warp] = acc;
3691    }
3692}
3693
3694/* Frame Jacobi diagonal subtraction:
3695 *   diag[a] -= Σ_row Σ_c Σ_d H_tβ[row,c,a]·ainv[row,c,d]·H_tβ[row,d,a].
3696 *
3697 * #2535 — one thread owns each diagonal scalar and walks rows in host order.
3698 * The former (row,a) grid raced through atomicAdd(&diag[a], -quad), making the
3699 * preconditioner depend on CTA scheduling even after the two G writers were
3700 * made deterministic. */
3701extern "C" __global__ void arrow_sae_frame_diag_sub(
3702    double* __restrict__ diag,
3703    const double* __restrict__ ainv,
3704    const int* __restrict__ htb_ptr,
3705    const double* __restrict__ htb,
3706    const int* __restrict__ q_of,
3707    int k,
3708    int max_q,
3709    int n_rows
3710) {
3711    int a = blockIdx.x * blockDim.x + threadIdx.x;
3712    if (a >= k) { return; }
3713    double acc = diag[a];
3714    for (int row = 0; row < n_rows; ++row) {
3715        int q = q_of[row];
3716        int hbase = htb_ptr[row];
3717        int abase = row * max_q * max_q;
3718        double quad = 0.0;
3719        for (int c = 0; c < q; ++c) {
3720            double hc = htb[hbase + c * k + a];
3721            for (int d = 0; d < q; ++d) {
3722                quad += hc * ainv[abase + c * max_q + d] * htb[hbase + d * k + a];
3723            }
3724        }
3725        acc -= quad;
3726    }
3727    diag[a] = acc;
3728}
3729"#;
3730
3731    fn pcg_vector_module(
3732        ctx: &Arc<CudaContext>,
3733    ) -> Result<&'static Arc<CudaModule>, ArrowSchurGpuFailure> {
3734        static CACHE: gam_gpu::device_cache::PtxModuleCache =
3735            gam_gpu::device_cache::PtxModuleCache::new();
3736        CACHE
3737            .get_or_compile(ctx, "arrow_pcg_vector", PCG_VECTOR_KERNEL_SOURCE)
3738            .map_err(|err| {
3739                // #1551: an NVRTC compile / module-load failure of
3740                // PCG_VECTOR_KERNEL_SOURCE means the device SAE PCG cannot run;
3741                // log it (the historical silent collapse to `Unavailable` is what
3742                // masked the missing `--gpu-architecture` for so long) and fall
3743                // back to the CPU.
3744                log::warn!("[#1551] pcg_vector_module get_or_compile failed: {err}");
3745                ArrowSchurGpuFailure::Unavailable
3746            })
3747    }
3748
3749    fn pcg_launch_config(n: usize) -> Result<LaunchConfig, ArrowSchurGpuFailure> {
3750        let threads = 256u32;
3751        let blocks = ((n as u32).saturating_add(threads - 1) / threads).max(1);
3752        Ok(LaunchConfig {
3753            grid_dim: (blocks, 1, 1),
3754            block_dim: (threads, 1, 1),
3755            shared_mem_bytes: 0,
3756        })
3757    }
3758
3759    fn launch_jacobi_mul(
3760        stream: &Arc<CudaStream>,
3761        module: &Arc<CudaModule>,
3762        inv_diag: &CudaSlice<f64>,
3763        r: &CudaSlice<f64>,
3764        z: &mut CudaSlice<f64>,
3765        n: usize,
3766    ) -> Result<(), ArrowSchurGpuFailure> {
3767        let kernel = module
3768            .load_function("arrow_pcg_jacobi_mul")
3769            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
3770        let n_i32 = to_i32(n).ok_or(ArrowSchurGpuFailure::Unavailable)?;
3771        let mut builder = stream.launch_builder(&kernel);
3772        builder.arg(inv_diag).arg(r).arg(z).arg(&n_i32);
3773        // SAFETY: all buffers have length n and belong to `stream`; the kernel only
3774        // reads/writes indices `< n`.
3775        unsafe { builder.launch(pcg_launch_config(n)?) }
3776            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
3777        Ok(())
3778    }
3779
3780    fn launch_update_p(
3781        stream: &Arc<CudaStream>,
3782        module: &Arc<CudaModule>,
3783        z: &CudaSlice<f64>,
3784        beta: f64,
3785        p: &mut CudaSlice<f64>,
3786        n: usize,
3787    ) -> Result<(), ArrowSchurGpuFailure> {
3788        let kernel = module
3789            .load_function("arrow_pcg_update_p")
3790            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
3791        let n_i32 = to_i32(n).ok_or(ArrowSchurGpuFailure::Unavailable)?;
3792        let mut builder = stream.launch_builder(&kernel);
3793        builder.arg(z).arg(&beta).arg(p).arg(&n_i32);
3794        // SAFETY: z/p both have length n and belong to `stream`; the kernel only
3795        // reads/writes indices `< n`.
3796        unsafe { builder.launch(pcg_launch_config(n)?) }
3797            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
3798        Ok(())
3799    }
3800
3801    struct DeviceSaePcgBuffers {
3802        row_ptr: CudaSlice<i32>,
3803        beta_base: CudaSlice<i32>,
3804        phi: CudaSlice<f64>,
3805        jac_ptr: CudaSlice<i32>,
3806        jac: CudaSlice<f64>,
3807        smooth_offsets: CudaSlice<i32>,
3808        smooth_m: CudaSlice<i32>,
3809        smooth_ptr: CudaSlice<i32>,
3810        smooth_data: CudaSlice<f64>,
3811        g_col_off: CudaSlice<i32>,
3812        g_cols: CudaSlice<i32>,
3813        g_ptr: CudaSlice<i32>,
3814        g_data: CudaSlice<f64>,
3815        g_group_ptr: CudaSlice<i32>,
3816        g_group_blocks: CudaSlice<i32>,
3817        g_group_off: CudaSlice<i32>,
3818        g_group_width: CudaSlice<i32>,
3819        ainv: CudaSlice<f64>,
3820        u: CudaSlice<f64>,
3821        w: CudaSlice<f64>,
3822        v: CudaSlice<f64>,
3823        /// #2512 two-stage deterministic scatter scratch: `partial[n_chunks × k]`.
3824        scatter_partial: CudaSlice<f64>,
3825        scatter_n_chunks: usize,
3826        scatter_rows_per_chunk: usize,
3827        n_rows: usize,
3828        p: usize,
3829        k: usize,
3830        max_q: usize,
3831        smooth_blocks: usize,
3832        g_groups: usize,
3833        g_max_work: usize,
3834    }
3835
3836    fn checked_i32(value: usize) -> Result<i32, ArrowSchurGpuFailure> {
3837        to_i32(value).ok_or(ArrowSchurGpuFailure::Unavailable)
3838    }
3839
3840    fn sae_penalty_diag_host(
3841        data: &DeviceSaePcgData,
3842        ridge_beta: f64,
3843    ) -> Result<Vec<f64>, ArrowSchurGpuFailure> {
3844        let mut diag = vec![ridge_beta; data.beta_dim];
3845        for block in &data.smooth_blocks {
3846            let (rows, cols) = block.factor_a.dim();
3847            if rows != cols {
3848                return Err(ArrowSchurGpuFailure::Unavailable);
3849            }
3850            for row in 0..rows {
3851                let coeff = block.factor_a[[row, row]];
3852                let base = block
3853                    .global_offset
3854                    .checked_add(
3855                        row.checked_mul(data.p)
3856                            .ok_or(ArrowSchurGpuFailure::Unavailable)?,
3857                    )
3858                    .ok_or(ArrowSchurGpuFailure::Unavailable)?;
3859                let end = base
3860                    .checked_add(data.p)
3861                    .ok_or(ArrowSchurGpuFailure::Unavailable)?;
3862                if end > diag.len() {
3863                    return Err(ArrowSchurGpuFailure::Unavailable);
3864                }
3865                for channel in 0..data.p {
3866                    diag[base + channel] += coeff;
3867                }
3868            }
3869        }
3870        for block in &data.sparse_g_blocks {
3871            if block.row_off != block.col_off {
3872                continue;
3873            }
3874            let (rows, cols) = block.data.dim();
3875            for row in 0..rows.min(cols) {
3876                let coeff = block.data[[row, row]];
3877                let beta_row = block
3878                    .row_off
3879                    .checked_add(row)
3880                    .ok_or(ArrowSchurGpuFailure::Unavailable)?;
3881                let base = beta_row
3882                    .checked_mul(data.p)
3883                    .ok_or(ArrowSchurGpuFailure::Unavailable)?;
3884                let end = base
3885                    .checked_add(data.p)
3886                    .ok_or(ArrowSchurGpuFailure::Unavailable)?;
3887                if end > diag.len() {
3888                    return Err(ArrowSchurGpuFailure::Unavailable);
3889                }
3890                for channel in 0..data.p {
3891                    diag[base + channel] += coeff;
3892                }
3893            }
3894        }
3895        Ok(diag)
3896    }
3897
3898    fn flatten_device_sae_data(
3899        sys: &ArrowSchurSystem,
3900        data: &DeviceSaePcgData,
3901        ridge_t: f64,
3902        stream: &Arc<CudaStream>,
3903    ) -> Result<DeviceSaePcgBuffers, ArrowSchurGpuFailure> {
3904        let n_rows = sys.rows.len();
3905        let p = data.p;
3906        let k = data.beta_dim;
3907        if data.a_phi.len() != n_rows || data.local_jac.len() != n_rows {
3908            return Err(ArrowSchurGpuFailure::Unavailable);
3909        }
3910
3911        let (scatter_rows_per_chunk, scatter_n_chunks) = scatter_chunking(n_rows);
3912        let mut row_ptr_host = Vec::with_capacity(n_rows + 1);
3913        let mut beta_base_host = Vec::<i32>::new();
3914        let mut phi_host = Vec::<f64>::new();
3915        row_ptr_host.push(0_i32);
3916        for row in data.a_phi.iter() {
3917            for &(base, phi) in row {
3918                beta_base_host.push(checked_i32(base)?);
3919                phi_host.push(phi);
3920            }
3921            row_ptr_host.push(checked_i32(beta_base_host.len())?);
3922        }
3923
3924        let mut jac_ptr_host = Vec::with_capacity(n_rows + 1);
3925        let mut jac_host = Vec::<f64>::new();
3926        let mut max_q = 0usize;
3927        jac_ptr_host.push(0_i32);
3928        for row_jac in data.local_jac.iter() {
3929            if row_jac.len() % p != 0 {
3930                return Err(ArrowSchurGpuFailure::Unavailable);
3931            }
3932            max_q = max_q.max(row_jac.len() / p);
3933            jac_host.extend_from_slice(row_jac);
3934            jac_ptr_host.push(checked_i32(jac_host.len())?);
3935        }
3936        if max_q == 0 {
3937            return Err(ArrowSchurGpuFailure::Unavailable);
3938        }
3939
3940        let mut smooth_offsets_host = Vec::with_capacity(data.smooth_blocks.len());
3941        let mut smooth_m_host = Vec::with_capacity(data.smooth_blocks.len());
3942        let mut smooth_ptr_host = Vec::with_capacity(data.smooth_blocks.len() + 1);
3943        let mut smooth_data_host = Vec::<f64>::new();
3944        smooth_ptr_host.push(0_i32);
3945        for block in &data.smooth_blocks {
3946            let (rows, cols) = block.factor_a.dim();
3947            if rows != cols {
3948                return Err(ArrowSchurGpuFailure::Unavailable);
3949            }
3950            smooth_offsets_host.push(checked_i32(block.global_offset)?);
3951            smooth_m_host.push(checked_i32(rows)?);
3952            for r in 0..rows {
3953                for c in 0..cols {
3954                    smooth_data_host.push(block.factor_a[[r, c]]);
3955                }
3956            }
3957            smooth_ptr_host.push(checked_i32(smooth_data_host.len())?);
3958        }
3959
3960        let mut g_col_off_host = Vec::with_capacity(data.sparse_g_blocks.len());
3961        let mut g_cols_host = Vec::with_capacity(data.sparse_g_blocks.len());
3962        let mut g_ptr_host = Vec::with_capacity(data.sparse_g_blocks.len() + 1);
3963        let mut g_data_host = Vec::<f64>::new();
3964        let mut g_output_spans = Vec::with_capacity(data.sparse_g_blocks.len());
3965        g_ptr_host.push(0_i32);
3966        for block in &data.sparse_g_blocks {
3967            let (rows, cols) = block.data.dim();
3968            g_col_off_host.push(checked_i32(block.col_off)?);
3969            g_cols_host.push(checked_i32(cols)?);
3970            for r in 0..rows {
3971                for c in 0..cols {
3972                    g_data_host.push(block.data[[r, c]]);
3973                }
3974            }
3975            g_ptr_host.push(checked_i32(g_data_host.len())?);
3976            let offset = block
3977                .row_off
3978                .checked_mul(p)
3979                .ok_or(ArrowSchurGpuFailure::Unavailable)?;
3980            let width = rows
3981                .checked_mul(p)
3982                .ok_or(ArrowSchurGpuFailure::Unavailable)?;
3983            g_output_spans.push((offset, width));
3984        }
3985        let g_groups =
3986            super::deterministic_output_groups(&g_output_spans, k, "sparse G matvec")?;
3987        let g_group_count = g_groups.off.len();
3988
3989        let mut ainv_host = vec![0.0_f64; n_rows * max_q * max_q];
3990        for (row_idx, row) in sys.rows.iter().enumerate() {
3991            let q = data.local_jac[row_idx].len() / p;
3992            if row.htt.dim() != (q, q) {
3993                return Err(ArrowSchurGpuFailure::SchurFactorFailed {
3994                    reason: format!(
3995                        "SAE device PCG row {row_idx}: H_tt shape {:?} != ({q}, {q})",
3996                        row.htt.dim()
3997                    ),
3998                });
3999            }
4000            let mut block = row.htt.clone();
4001            for d in 0..q {
4002                block[[d, d]] += ridge_t;
4003            }
4004            let factor = gam_linalg::triangular::cholesky_factor_in_place(
4005                block.view(),
4006                gam_linalg::triangular::CholeskyGuard::NonnegativePivot,
4007            )
4008            .ok_or_else(|| {
4009                // Deficit-aware bump (Gershgorin λ_min bound) so a strongly
4010                // indefinite per-row block recovers in one outer-loop retry.
4011                ArrowSchurGpuFailure::RidgeBumpRequired {
4012                    row: row_idx,
4013                    bump: super::ridge_bump_to_make_pd(row.htt.view(), ridge_t),
4014                }
4015            })?;
4016            for col in 0..q {
4017                let mut e = Array1::<f64>::zeros(q);
4018                e[col] = 1.0;
4019                let solved = gam_linalg::triangular::cholesky_solve_vector(factor.view(), e.view());
4020                for r in 0..q {
4021                    ainv_host[row_idx * max_q * max_q + r * max_q + col] = solved[r];
4022                }
4023            }
4024        }
4025
4026        Ok(DeviceSaePcgBuffers {
4027            row_ptr: stream
4028                .clone_htod(&row_ptr_host)
4029                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
4030            beta_base: stream
4031                .clone_htod(&beta_base_host)
4032                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
4033            phi: stream
4034                .clone_htod(&phi_host)
4035                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
4036            jac_ptr: stream
4037                .clone_htod(&jac_ptr_host)
4038                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
4039            jac: stream
4040                .clone_htod(&jac_host)
4041                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
4042            smooth_offsets: stream
4043                .clone_htod(&smooth_offsets_host)
4044                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
4045            smooth_m: stream
4046                .clone_htod(&smooth_m_host)
4047                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
4048            smooth_ptr: stream
4049                .clone_htod(&smooth_ptr_host)
4050                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
4051            smooth_data: stream
4052                .clone_htod(&smooth_data_host)
4053                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
4054            g_col_off: stream
4055                .clone_htod(&g_col_off_host)
4056                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
4057            g_cols: stream
4058                .clone_htod(&g_cols_host)
4059                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
4060            g_ptr: stream
4061                .clone_htod(&g_ptr_host)
4062                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
4063            g_data: stream
4064                .clone_htod(&g_data_host)
4065                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
4066            g_group_ptr: stream
4067                .clone_htod(&g_groups.ptr)
4068                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
4069            g_group_blocks: stream
4070                .clone_htod(&g_groups.blocks)
4071                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
4072            g_group_off: stream
4073                .clone_htod(&g_groups.off)
4074                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
4075            g_group_width: stream
4076                .clone_htod(&g_groups.width)
4077                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
4078            ainv: stream
4079                .clone_htod(&ainv_host)
4080                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
4081            u: stream
4082                .alloc_zeros::<f64>(n_rows * p)
4083                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
4084            w: stream
4085                .alloc_zeros::<f64>(n_rows * max_q)
4086                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
4087            v: stream
4088                .alloc_zeros::<f64>(n_rows * max_q)
4089                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
4090            scatter_partial: stream
4091                .alloc_zeros::<f64>(scatter_n_chunks * k)
4092                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
4093            scatter_n_chunks,
4094            scatter_rows_per_chunk,
4095            n_rows,
4096            p,
4097            k,
4098            max_q,
4099            smooth_blocks: data.smooth_blocks.len(),
4100            g_groups: g_group_count,
4101            g_max_work: g_groups.max_width,
4102        })
4103    }
4104
4105    fn launch_sae_init(
4106        stream: &Arc<CudaStream>,
4107        module: &Arc<CudaModule>,
4108        out: &mut CudaSlice<f64>,
4109        x: &CudaSlice<f64>,
4110        ridge: f64,
4111        n: usize,
4112    ) -> Result<(), ArrowSchurGpuFailure> {
4113        let kernel = module
4114            .load_function("arrow_sae_init")
4115            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4116        let n_i32 = checked_i32(n)?;
4117        let mut builder = stream.launch_builder(&kernel);
4118        builder.arg(out).arg(x).arg(&ridge).arg(&n_i32);
4119        // SAFETY: `out` and `x` are live device buffers with at least `n`
4120        // entries on `stream`; the kernel writes one in-bounds element per
4121        // launched index below `n`.
4122        unsafe { builder.launch(pcg_launch_config(n)?) }
4123            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4124        Ok(())
4125    }
4126
4127    fn launch_sae_penalty_matvec(
4128        stream: &Arc<CudaStream>,
4129        module: &Arc<CudaModule>,
4130        buffers: &mut DeviceSaePcgBuffers,
4131        x: &CudaSlice<f64>,
4132        out: &mut CudaSlice<f64>,
4133        ridge_beta: f64,
4134    ) -> Result<(), ArrowSchurGpuFailure> {
4135        launch_sae_init(stream, module, out, x, ridge_beta, buffers.k)?;
4136        if buffers.smooth_blocks > 0 {
4137            let kernel = module
4138                .load_function("arrow_sae_smooth_matvec")
4139                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4140            let max_m = buffers.k;
4141            let p_i32 = checked_i32(buffers.p)?;
4142            let blocks_i32 = checked_i32(buffers.smooth_blocks)?;
4143            let cfg = LaunchConfig {
4144                grid_dim: (
4145                    ((max_m as u32).saturating_add(255) / 256).max(1),
4146                    checked_i32(buffers.smooth_blocks)? as u32,
4147                    1,
4148                ),
4149                block_dim: (256, 1, 1),
4150                shared_mem_bytes: 0,
4151            };
4152            let mut builder = stream.launch_builder(&kernel);
4153            builder
4154                .arg(x)
4155                .arg(&mut *out)
4156                .arg(&buffers.smooth_offsets)
4157                .arg(&buffers.smooth_m)
4158                .arg(&buffers.smooth_ptr)
4159                .arg(&buffers.smooth_data)
4160                .arg(&p_i32)
4161                .arg(&blocks_i32);
4162            // SAFETY: smooth block metadata and dense smooth data were flattened
4163            // into live device buffers; the 2D grid covers only declared block
4164            // and coefficient-channel work items, and the kernel bounds-checks
4165            // against each block's stored size.
4166            unsafe { builder.launch(cfg) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4167        }
4168        if buffers.g_groups > 0 {
4169            let kernel = module
4170                .load_function("arrow_sae_sparse_g_matvec")
4171                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4172            let p_i32 = checked_i32(buffers.p)?;
4173            let groups_i32 = checked_i32(buffers.g_groups)?;
4174            let cfg = LaunchConfig {
4175                grid_dim: (
4176                    ((buffers.g_max_work as u32).saturating_add(255) / 256).max(1),
4177                    groups_i32 as u32,
4178                    1,
4179                ),
4180                block_dim: (256, 1, 1),
4181                shared_mem_bytes: 0,
4182            };
4183            let mut builder = stream.launch_builder(&kernel);
4184            builder
4185                .arg(x)
4186                .arg(&mut *out)
4187                .arg(&buffers.g_col_off)
4188                .arg(&buffers.g_cols)
4189                .arg(&buffers.g_ptr)
4190                .arg(&buffers.g_data)
4191                .arg(&buffers.g_group_ptr)
4192                .arg(&buffers.g_group_blocks)
4193                .arg(&buffers.g_group_off)
4194                .arg(&buffers.g_group_width)
4195                .arg(&p_i32)
4196                .arg(&groups_i32);
4197            // SAFETY: each CSR ownership group covers one in-bounds scalar output
4198            // span, distinct group spans are disjoint, and every listed block's
4199            // column/data metadata is live. One thread writes each output after
4200            // walking its contributor list in fixed order.
4201            unsafe { builder.launch(cfg) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4202        }
4203        Ok(())
4204    }
4205
4206    fn launch_sae_row_schur_sub(
4207        stream: &Arc<CudaStream>,
4208        module: &Arc<CudaModule>,
4209        buffers: &mut DeviceSaePcgBuffers,
4210        x: &CudaSlice<f64>,
4211        out: &mut CudaSlice<f64>,
4212    ) -> Result<(), ArrowSchurGpuFailure> {
4213        let p_i32 = checked_i32(buffers.p)?;
4214        let max_q_i32 = checked_i32(buffers.max_q)?;
4215        let n_rows_i32 = checked_i32(buffers.n_rows)?;
4216        let cfg_p_rows = LaunchConfig {
4217            grid_dim: (
4218                ((buffers.p as u32).saturating_add(255) / 256).max(1),
4219                checked_i32(buffers.n_rows)? as u32,
4220                1,
4221            ),
4222            block_dim: (256, 1, 1),
4223            shared_mem_bytes: 0,
4224        };
4225        let gather = module
4226            .load_function("arrow_sae_gather_u")
4227            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4228        {
4229            let mut builder = stream.launch_builder(&gather);
4230            builder
4231                .arg(x)
4232                .arg(&buffers.row_ptr)
4233                .arg(&buffers.beta_base)
4234                .arg(&buffers.phi)
4235                .arg(&mut buffers.u)
4236                .arg(&p_i32)
4237                .arg(&n_rows_i32);
4238            // SAFETY: `x`, row pointers, beta offsets, basis rows, and `u` are
4239            // live device buffers sized for `n_rows` by `p`; the kernel guards
4240            // row/channel indices before gathering.
4241            unsafe { builder.launch(cfg_p_rows) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4242        }
4243
4244        let cfg_q_rows = LaunchConfig {
4245            grid_dim: (
4246                ((buffers.max_q as u32).saturating_add(255) / 256).max(1),
4247                checked_i32(buffers.n_rows)? as u32,
4248                1,
4249            ),
4250            block_dim: (256, 1, 1),
4251            shared_mem_bytes: 0,
4252        };
4253        let apply_l = module
4254            .load_function("arrow_sae_apply_l")
4255            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4256        {
4257            let mut builder = stream.launch_builder(&apply_l);
4258            builder
4259                .arg(&buffers.u)
4260                .arg(&buffers.jac_ptr)
4261                .arg(&buffers.jac)
4262                .arg(&mut buffers.w)
4263                .arg(&p_i32)
4264                .arg(&max_q_i32)
4265                .arg(&n_rows_i32);
4266            // SAFETY: `u`, Jacobian row pointers/data, and `w` are live buffers
4267            // sized for the `(n_rows, p)` to `(n_rows, max_q)` multiply; the
4268            // kernel checks row and local-coordinate bounds.
4269            unsafe { builder.launch(cfg_q_rows) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4270        }
4271
4272        let apply_ainv = module
4273            .load_function("arrow_sae_apply_ainv")
4274            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4275        {
4276            let mut builder = stream.launch_builder(&apply_ainv);
4277            builder
4278                .arg(&buffers.ainv)
4279                .arg(&buffers.w)
4280                .arg(&mut buffers.v)
4281                .arg(&max_q_i32)
4282                .arg(&n_rows_i32);
4283            // SAFETY: `ainv`, `w`, and `v` are live device buffers sized for
4284            // `n_rows * max_q`; the kernel guards all row/local-coordinate
4285            // indices before reading or writing.
4286            unsafe { builder.launch(cfg_q_rows) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4287        }
4288
4289        // #2512 — fold the per-row contributions into `out` through a fixed
4290        // reassociation instead of `atomicAdd`. Stage 1 sums each contiguous row
4291        // chunk in fixed index order into `scatter_partial`; stage 2 reduces the
4292        // chunks in fixed order. The former single kernel let the block scheduler
4293        // pick the row order, which made this matvec — and therefore the Newton
4294        // step and the fitted decoder — a different value on every call.
4295        let k_i32 = checked_i32(buffers.k)?;
4296        let rows_per_chunk_i32 = checked_i32(buffers.scatter_rows_per_chunk)?;
4297        let n_chunks_i32 = checked_i32(buffers.scatter_n_chunks)?;
4298        let k_blocks = ((buffers.k as u32).saturating_add(255) / 256).max(1);
4299        {
4300            let kernel = module
4301                .load_function("arrow_sae_scatter_sub_det_partial")
4302                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4303            let cfg = LaunchConfig {
4304                grid_dim: (
4305                    ((buffers.p as u32).saturating_add(255) / 256).max(1),
4306                    n_chunks_i32 as u32,
4307                    1,
4308                ),
4309                block_dim: (256, 1, 1),
4310                shared_mem_bytes: 0,
4311            };
4312            let mut builder = stream.launch_builder(&kernel);
4313            builder
4314                .arg(&buffers.v)
4315                .arg(&buffers.jac_ptr)
4316                .arg(&buffers.jac)
4317                .arg(&buffers.row_ptr)
4318                .arg(&buffers.beta_base)
4319                .arg(&buffers.phi)
4320                .arg(&mut buffers.scatter_partial)
4321                .arg(&p_i32)
4322                .arg(&k_i32)
4323                .arg(&max_q_i32)
4324                .arg(&n_rows_i32)
4325                .arg(&rows_per_chunk_i32);
4326            // SAFETY: `v`, Jacobian metadata, row pointers, beta offsets and basis
4327            // rows are live buffers for `n_rows` by `p`; `scatter_partial` is sized
4328            // `n_chunks · k` and each thread clears then writes only the indices of
4329            // its own chunk congruent to its channel modulo `p`.
4330            unsafe { builder.launch(cfg) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4331        }
4332        {
4333            let kernel = module
4334                .load_function("arrow_sae_scatter_sub_det_reduce")
4335                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4336            let cfg = LaunchConfig {
4337                grid_dim: (k_blocks, 1, 1),
4338                block_dim: (256, 1, 1),
4339                shared_mem_bytes: 0,
4340            };
4341            let mut builder = stream.launch_builder(&kernel);
4342            builder
4343                .arg(&buffers.scatter_partial)
4344                .arg(&mut *out)
4345                .arg(&k_i32)
4346                .arg(&n_chunks_i32);
4347            // SAFETY: `scatter_partial` is sized `n_chunks · k` and `out` is sized
4348            // `k`; exactly one thread updates each in-bounds `out[a]`.
4349            unsafe { builder.launch(cfg) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4350        }
4351        Ok(())
4352    }
4353
4354    fn launch_sae_diag_sub(
4355        stream: &Arc<CudaStream>,
4356        module: &Arc<CudaModule>,
4357        buffers: &mut DeviceSaePcgBuffers,
4358        diag: &mut CudaSlice<f64>,
4359    ) -> Result<(), ArrowSchurGpuFailure> {
4360        let p_i32 = checked_i32(buffers.p)?;
4361        let k_i32 = checked_i32(buffers.k)?;
4362        let max_q_i32 = checked_i32(buffers.max_q)?;
4363        let n_rows_i32 = checked_i32(buffers.n_rows)?;
4364        let rows_per_chunk_i32 = checked_i32(buffers.scatter_rows_per_chunk)?;
4365        let n_chunks_i32 = checked_i32(buffers.scatter_n_chunks)?;
4366        {
4367            let kernel = module
4368                .load_function("arrow_sae_diag_sub_det_partial")
4369                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4370            let cfg = LaunchConfig {
4371                grid_dim: (
4372                    ((buffers.p as u32).saturating_add(255) / 256).max(1),
4373                    n_chunks_i32 as u32,
4374                    1,
4375                ),
4376                block_dim: (256, 1, 1),
4377                shared_mem_bytes: 0,
4378            };
4379            let mut builder = stream.launch_builder(&kernel);
4380            builder
4381                .arg(&buffers.ainv)
4382                .arg(&buffers.jac_ptr)
4383                .arg(&buffers.jac)
4384                .arg(&buffers.row_ptr)
4385                .arg(&buffers.beta_base)
4386                .arg(&buffers.phi)
4387                .arg(&mut buffers.scatter_partial)
4388                .arg(&p_i32)
4389                .arg(&k_i32)
4390                .arg(&max_q_i32)
4391                .arg(&n_rows_i32)
4392                .arg(&rows_per_chunk_i32);
4393            // SAFETY: every read-only SAE row metadata buffer is live on `stream`
4394            // with sizes matching `n_rows`, `p` and `max_q`; `scatter_partial` is
4395            // sized `n_chunks · k` and each thread clears then writes only the
4396            // indices of its own chunk congruent to its channel modulo `p`.
4397            unsafe { builder.launch(cfg) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4398        }
4399        {
4400            let kernel = module
4401                .load_function("arrow_sae_scatter_sub_det_reduce")
4402                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4403            let cfg = LaunchConfig {
4404                grid_dim: (
4405                    ((buffers.k as u32).saturating_add(255) / 256).max(1),
4406                    1,
4407                    1,
4408                ),
4409                block_dim: (256, 1, 1),
4410                shared_mem_bytes: 0,
4411            };
4412            let mut builder = stream.launch_builder(&kernel);
4413            builder
4414                .arg(&buffers.scatter_partial)
4415                .arg(diag)
4416                .arg(&k_i32)
4417                .arg(&n_chunks_i32);
4418            // SAFETY: `scatter_partial` is sized `n_chunks · k` and `diag` is sized
4419            // `k`; exactly one thread updates each in-bounds `diag[a]`.
4420            unsafe { builder.launch(cfg) }
4421                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4422            Ok(())
4423        }
4424    }
4425
4426    fn launch_sae_matvec(
4427        stream: &Arc<CudaStream>,
4428        module: &Arc<CudaModule>,
4429        buffers: &mut DeviceSaePcgBuffers,
4430        x: &CudaSlice<f64>,
4431        out: &mut CudaSlice<f64>,
4432        ridge_beta: f64,
4433    ) -> Result<(), ArrowSchurGpuFailure> {
4434        launch_sae_penalty_matvec(stream, module, buffers, x, out, ridge_beta)?;
4435        launch_sae_row_schur_sub(stream, module, buffers, x, out)
4436    }
4437
4438    /// Pack `D + ρ_t I`, `B`, and `g` into the strided `(n × P_MAX × P_MAX)`
4439    /// / `(n × P_MAX × R_TEMPLATE)` / `(n × P_MAX)` layout the fused kernel
4440    /// expects. Entries outside the runtime `(p, r)` window stay at zero so
4441    /// the kernel's per-element loops are safe to no-op there.
4442    fn pack_fused_host(
4443        sys: &ArrowSchurSystem,
4444        ridge_t: f64,
4445        p_max: usize,
4446        r_template: usize,
4447    ) -> (Vec<f64>, Vec<f64>, Vec<f64>) {
4448        let n = sys.rows.len();
4449        let d = sys.d;
4450        let k = sys.k;
4451        let mut d_buf = vec![0.0_f64; n * p_max * p_max];
4452        let mut b_buf = vec![0.0_f64; n * p_max * r_template];
4453        let mut g_buf = vec![0.0_f64; n * p_max];
4454        for (i, row) in sys.rows.iter().enumerate() {
4455            // D_i + ρI, column-major in P_MAX×P_MAX strided block.
4456            for col in 0..d {
4457                let base = (i * p_max + col) * p_max;
4458                for r in 0..d {
4459                    let mut value = row.htt[[r, col]];
4460                    if r == col {
4461                        value += ridge_t;
4462                    }
4463                    d_buf[base + r] = value;
4464                }
4465            }
4466            // B_i in P_MAX×R_TEMPLATE strided block. The per-row (per-i) block
4467            // stride is `p_max · r_template` (matching the `b_buf` allocation
4468            // above and the kernel's `b_stack`/`y_out` layout), NOT
4469            // `p_max · p_max`: using the D-block multiplier here overflows the
4470            // buffer whenever `p_max > r_template` (e.g. d=30→p_max=32,
4471            // k=5→r_template=5). The within-block element offset stays
4472            // column-major `col·p_max + r` (P_MAX rows per column).
4473            for col in 0..k {
4474                let base = (i * r_template + col) * p_max;
4475                for r in 0..d {
4476                    b_buf[base + r] = row.htbeta[[r, col]];
4477                }
4478            }
4479            // g_i in P_MAX strided vector.
4480            let g_base = i * p_max;
4481            for r in 0..d {
4482                g_buf[g_base + r] = row.gt[r];
4483            }
4484        }
4485        (d_buf, b_buf, g_buf)
4486    }
4487
4488    // -----------------------------------------------------------------------
4489    // #1017 Phase 3: across-iteration device residency.
4490    //
4491    // `solve()` re-packs and re-uploads `D` (`H_tt`), `B` (`H_tβ`) and `g`,
4492    // then re-runs the per-row POTRF and the border Schur factorization on
4493    // EVERY call. For the SAE joint inner Newton at a frozen gate/basis frame
4494    // the Hessian blocks `D`, `B`, `H_ββ` are CONSTANT across the inner loop —
4495    // only the gradient `g = r(z) = H z − g₀` changes per iterate. So the
4496    // factor work (`O(n·d³ + p³)`) and the dominant `O(n·d·p)` cross-block
4497    // upload are pure waste when repeated per iterate.
4498    //
4499    // `ResidentArrowFrame` performs that constant work ONCE at construction:
4500    // upload+ridge+POTRF of `D` (keeping `L_i` resident in `l_dev`), the
4501    // forward solve `Y_i = L_i^{-1} B_i` (kept resident in `y_dev`), and the
4502    // Schur assembly + border POTRF (keeping `L_S` resident in `schur_dev`).
4503    // Each subsequent `solve_gradient(g)` uploads only the `n·d` row gradient,
4504    // runs the cheap residual path — `u_i = L_i^{-1} g_i` (one batched TRSM),
4505    // Schur RHS `−g_β + Σ Y_iᵀ u_i`, `δβ = L_S^{-T} L_S^{-1} rhs` (two TRSM,
4506    // NO POTRF), back-sub `δt_i = −L_i^{-T}(u_i + Y_i δβ)` — and reads back only
4507    // `δ` and the cached log|H|. The heavy buffers never leave the device
4508    // across iterations; the per-iterate host transfer is `O(n·d + p)`, not
4509    // `O(n·d·p)`. Numerics are bit-identical to a `solve()` at the same
4510    // `(D, B, H_ββ, g, ridge_t, ridge_beta)` because the factor buffers and the
4511    // helper kernels are the same; the resident path merely SKIPS re-deriving
4512    // the parts that do not depend on `g`. The CPU dense reference
4513    // (`solve_arrow_newton_step_dense_reference`) is the parity oracle.
4514    pub(super) struct ResidentArrowFrame {
4515        n: usize,
4516        d: usize,
4517        k: usize,
4518        beta_gauge_quotient: Option<ArrowBetaGaugeQuotient>,
4519        stream: Arc<CudaStream>,
4520        blas: CudaBlas,
4521        /// Per-row lower Cholesky factors `L_i` of `H_tt + ρ_t I`, stacked
4522        /// column-major (`n` tiles of `d×d`). Resident across iterations.
4523        l_dev: CudaSlice<f64>,
4524        /// Whitened cross blocks `Y_i = L_i^{-1} H_tβ^(i)`, stacked column-major
4525        /// (`n` tiles of `d×k`). Resident across iterations.
4526        y_dev: CudaSlice<f64>,
4527        /// Lower Cholesky factor `L_S` of the reduced Schur complement
4528        /// `S_β = H_ββ + ρ_β I − Σ_i Y_iᵀ Y_i`. Resident across iterations.
4529        schur_dev: CudaSlice<f64>,
4530        /// `log|H| = 2 Σ log L_{i,jj} + 2 Σ log L_{S,aa}`, constant for the
4531        /// frame (depends only on the factored Hessian, not on `g`).
4532        log_det_hessian: f64,
4533    }
4534
4535    impl ResidentArrowFrame {
4536        /// Upload the constant Hessian blocks and perform the one-time factor
4537        /// work (`POTRF(D)`, `Y_i = L_i^{-1} B_i`, Schur assembly + border
4538        /// `POTRF`). The frame then serves cheap per-gradient solves.
4539        pub(super) fn new(
4540            sys: &ArrowSchurSystem,
4541            ridge_t: f64,
4542            ridge_beta: f64,
4543            newton_schur_tikhonov_rel_floor: Option<f64>,
4544        ) -> Result<Self, ArrowSchurGpuFailure> {
4545            if ridge_t.is_nan() || ridge_beta.is_nan() {
4546                return Err(ArrowSchurGpuFailure::SchurFactorFailed {
4547                    reason: "ridge is NaN".to_string(),
4548                });
4549            }
4550            let n = sys.rows.len();
4551            let d = sys.d;
4552            let k = sys.k;
4553            let runtime = route_through_gpu(DispatchOp::SmallDenseBatchedPotrf { p: d, batch: n })
4554                .ok_or(ArrowSchurGpuFailure::Unavailable)?;
4555            let stream = gam_gpu::device_runtime::cuda_context_for(runtime.device.ordinal)
4556                .and_then(|ctx| ctx.new_stream().ok())
4557                .ok_or(ArrowSchurGpuFailure::Unavailable)?;
4558            let solver =
4559                DnHandle::new(stream.clone()).map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4560            let blas =
4561                CudaBlas::new(stream.clone()).map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4562
4563            // Upload the constant blocks. `g` is uploaded per-gradient, not here.
4564            // `B` goes up in the STACKED `(n·d)×k` column-major layout so the
4565            // Schur reduction and both per-iterate accumulations are each ONE
4566            // launch (see `schur_gemm_stacked`); the `d×d` `D` blocks keep the
4567            // block-contiguous layout the batched POTRF/TRSM expect.
4568            let (d_host, b_host_stacked) = pack_host_d_and_stacked_b(sys, ridge_t);
4569            let mut l_dev = stream
4570                .clone_htod(&d_host)
4571                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4572            let mut y_dev = stream
4573                .clone_htod(&b_host_stacked)
4574                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4575
4576            // POTRF(D) → L_i, kept resident in l_dev.
4577            let info_host = potrf_batched(&solver, &stream, d, n, &mut l_dev)?;
4578            if let Some(idx) = info_host.iter().position(|info| *info != 0) {
4579                // cuSOLVER `info` is a 1-based pivot row index; size the bump
4580                // from the block (Gershgorin λ_min bound) so a strongly
4581                // indefinite block recovers in one retry.
4582                return Err(ArrowSchurGpuFailure::RidgeBumpRequired {
4583                    row: idx,
4584                    bump: super::ridge_bump_to_make_pd(sys.rows[idx].htt.view(), ridge_t),
4585                });
4586            }
4587
4588            // Y_i = L_i^{-1} B_i, in place over the stacked y_dev. Kept resident.
4589            trsm_batched_lower_inplace_stacked(&blas, &stream, d, n, k, &l_dev, &mut y_dev)?;
4590
4591            // Schur assembly S_β = (H_ββ + ρ_β I) − Σ Y_iᵀ Y_i, then POTRF → L_S.
4592            // The RHS accumulation is folded into the gradient path; here we
4593            // only need the (gradient-independent) Schur factor.
4594            let schur_init: Vec<f64> = {
4595                let mut tmp = Vec::with_capacity(k * k);
4596                for col in 0..k {
4597                    for row in 0..k {
4598                        let mut v = sys.hbb[[row, col]];
4599                        if row == col {
4600                            v += ridge_beta;
4601                        }
4602                        tmp.push(v);
4603                    }
4604                }
4605                tmp
4606            };
4607            let mut schur_dev = stream
4608                .clone_htod(&schur_init)
4609                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4610            schur_gemm_stacked(&blas, n * d, k, &y_dev, &mut schur_dev)?;
4611            let canonical_factor = if sys.beta_gauge_quotient.is_some()
4612                || newton_schur_tikhonov_rel_floor.is_some()
4613            {
4614                let mut schur_host = stream
4615                    .clone_dtoh(&schur_dev)
4616                    .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4617                let factorized = canonicalize_device_beta_factor(
4618                    sys.beta_gauge_quotient.as_ref(),
4619                    newton_schur_tikhonov_rel_floor,
4620                    k,
4621                    &mut schur_host,
4622                )?;
4623                schur_dev = stream
4624                    .clone_htod(&schur_host)
4625                    .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4626                factorized
4627            } else {
4628                false
4629            };
4630            if !canonical_factor {
4631                let info = potrf_single(&solver, &stream, k, &mut schur_dev)?;
4632                if info != 0 {
4633                    return Err(ArrowSchurGpuFailure::SchurFactorFailed {
4634                        reason: format!("Schur Cholesky failed at pivot {info}"),
4635                    });
4636                }
4637            }
4638
4639            // log|H| from the resident factors (constant for the frame).
4640            let l_local_host = stream
4641                .clone_dtoh(&l_dev)
4642                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4643            let l_schur_host = stream
4644                .clone_dtoh(&schur_dev)
4645                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4646            let mut log_det = 0.0_f64;
4647            for i in 0..n {
4648                let base = i * d * d;
4649                for j in 0..d {
4650                    log_det += l_local_host[base + j * d + j].ln();
4651                }
4652            }
4653            for j in 0..k {
4654                log_det += l_schur_host[j * k + j].ln();
4655            }
4656            log_det *= 2.0;
4657
4658            Ok(Self {
4659                n,
4660                d,
4661                k,
4662                beta_gauge_quotient: sys.beta_gauge_quotient.clone(),
4663                stream,
4664                blas,
4665                l_dev,
4666                y_dev,
4667                schur_dev,
4668                log_det_hessian: log_det,
4669            })
4670        }
4671
4672        #[inline]
4673        pub(super) fn log_det_hessian(&self) -> f64 {
4674            self.log_det_hessian
4675        }
4676
4677        /// Solve `H δ = −gradient` for a fresh gradient `(g_t, g_β)` reusing the
4678        /// resident factors. Uploads only `g_t` (`n·d` scalars); reads back only
4679        /// `δ`. No POTRF runs here — all factorization is amortized into `new`.
4680        pub(super) fn solve_gradient(
4681            &self,
4682            g_t: &[f64],
4683            g_beta: &[f64],
4684        ) -> Result<ArrowSchurGpuSolution, ArrowSchurGpuFailure> {
4685            let n = self.n;
4686            let d = self.d;
4687            let k = self.k;
4688            if g_t.len() != n * d || g_beta.len() != k {
4689                return Err(ArrowSchurGpuFailure::SchurFactorFailed {
4690                    reason: format!(
4691                        "resident gradient shape mismatch: g_t={} (want {}), g_beta={} (want {})",
4692                        g_t.len(),
4693                        n * d,
4694                        g_beta.len(),
4695                        k
4696                    ),
4697                });
4698            }
4699            // Upload the per-iterate row gradient → u_i = L_i^{-1} g_i in place.
4700            let mut u_dev = self
4701                .stream
4702                .clone_htod(&g_t.to_vec())
4703                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4704            trsm_batched_lower_inplace(&self.blas, &self.stream, d, n, 1, &self.l_dev, &mut u_dev)?;
4705
4706            // Schur RHS = −g_β + Σ_i Y_iᵀ u_i. Reuse the resident Schur factor
4707            // (no POTRF, and skip the −Σ Y_iᵀ Y_i GEMM already baked into L_S).
4708            let rhs_init: Vec<f64> = g_beta.iter().map(|v| -v).collect();
4709            let mut rhs_dev = self
4710                .stream
4711                .clone_htod(&rhs_init)
4712                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4713            schur_rhs_stacked(&self.blas, n * d, k, &self.y_dev, &u_dev, &mut rhs_dev)?;
4714            if self.beta_gauge_quotient.is_some() {
4715                let projected_rhs = project_device_beta_vector(
4716                    self.beta_gauge_quotient.as_ref(),
4717                    k,
4718                    self.stream
4719                        .clone_dtoh(&rhs_dev)
4720                        .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
4721                )?;
4722                rhs_dev = self
4723                    .stream
4724                    .clone_htod(&projected_rhs)
4725                    .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4726            }
4727
4728            // δβ ← L_S^{-T} L_S^{-1} rhs using the resident border factor.
4729            trsm_single(
4730                &self.blas,
4731                &self.stream,
4732                k,
4733                &self.schur_dev,
4734                &mut rhs_dev,
4735                false,
4736                false,
4737            )?;
4738            trsm_single(
4739                &self.blas,
4740                &self.stream,
4741                k,
4742                &self.schur_dev,
4743                &mut rhs_dev,
4744                false,
4745                true,
4746            )?;
4747            let delta_beta_host = project_device_beta_vector(
4748                self.beta_gauge_quotient.as_ref(),
4749                k,
4750                self.stream
4751                    .clone_dtoh(&rhs_dev)
4752                    .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
4753            )?;
4754            if self.beta_gauge_quotient.is_some() {
4755                rhs_dev = self
4756                    .stream
4757                    .clone_htod(&delta_beta_host)
4758                    .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4759            }
4760            let delta_beta = Array1::from_vec(delta_beta_host);
4761
4762            // Back-sub δt_i = −L_i^{-T}(u_i + Y_i δβ).
4763            back_sub_rhs_stacked(&self.blas, n * d, k, &self.y_dev, &rhs_dev, &mut u_dev)?;
4764            trsm_batched_lower_inplace_transposed(
4765                &self.blas,
4766                &self.stream,
4767                d,
4768                n,
4769                1,
4770                &self.l_dev,
4771                &mut u_dev,
4772            )?;
4773            let x_host = self
4774                .stream
4775                .clone_dtoh(&u_dev)
4776                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4777            let mut delta_t = Array1::<f64>::zeros(n * d);
4778            for (i, v) in x_host.iter().enumerate() {
4779                delta_t[i] = -*v;
4780            }
4781
4782            Ok(ArrowSchurGpuSolution {
4783                delta_t,
4784                delta_beta,
4785                log_det_hessian: self.log_det_hessian,
4786            })
4787        }
4788    }
4789
4790    /// #1017 base-resident frame for the LM ridge ladder: holds the ridge-
4791    /// INDEPENDENT base blocks resident and re-factors at each requested ridge.
4792    /// This is the residency counterpart to [`ResidentArrowFrame`], which bakes
4793    /// the ridge into its factors (wrong invariant for the ladder, whose trials
4794    /// vary the ridge while the gradient stays fixed).
4795    pub(super) struct ResidentBaseArrowFrame {
4796        n: usize,
4797        d: usize,
4798        k: usize,
4799        beta_gauge_quotient: Option<ArrowBetaGaugeQuotient>,
4800        newton_schur_tikhonov_rel_floor: Option<f64>,
4801        stream: Arc<CudaStream>,
4802        solver: DnHandle,
4803        blas: CudaBlas,
4804        /// `D = H_tt` at ridge 0, host-side (`n` stacked column-major `d×d`
4805        /// tiles). Per trial a copy gets `ridge_t` added to every tile diagonal
4806        /// and is uploaded (the `O(n·d·d)` re-diagonalised `D` is tiny relative to
4807        /// the resident `B`); the same copy feeds the `RidgeBumpRequired`
4808        /// Gershgorin bound on a non-PD pivot.
4809        d_base_host: Vec<f64>,
4810        /// `B = H_tβ` resident (`n` stacked column-major `d×k` tiles). Ridge-free;
4811        /// this is the bulk of the per-trial transfer the residency eliminates.
4812        base_b_dev: CudaSlice<f64>,
4813        /// Border `H_ββ` resident (column-major `k×k`). Ridge-free; `ridge_beta`
4814        /// is added to a per-trial device copy's diagonal.
4815        base_hbb_dev: CudaSlice<f64>,
4816        /// Row gradient `g_t` resident (`n·d`). Ridge-free.
4817        g_t_dev: CudaSlice<f64>,
4818        /// Border gradient `g_β` host-side (`k`); the tiny `−g_β` RHS is rebuilt
4819        /// per trial.
4820        gb_host: Vec<f64>,
4821        /// Resident all-ones vector (length `k`) whose strided daxpy adds
4822        /// `ridge_beta` to the `k×k` Schur base diagonal on-device.
4823        ones_k_dev: CudaSlice<f64>,
4824    }
4825
4826    impl ResidentBaseArrowFrame {
4827        /// Upload the ridge-independent base blocks once. No POTRF runs here.
4828        pub(super) fn new(
4829            sys: &ArrowSchurSystem,
4830            newton_schur_tikhonov_rel_floor: Option<f64>,
4831        ) -> Result<Self, ArrowSchurGpuFailure> {
4832            let n = sys.rows.len();
4833            let d = sys.d;
4834            let k = sys.k;
4835            let runtime = route_through_gpu(DispatchOp::SmallDenseBatchedPotrf { p: d, batch: n })
4836                .ok_or(ArrowSchurGpuFailure::Unavailable)?;
4837            let stream = gam_gpu::device_runtime::cuda_context_for(runtime.device.ordinal)
4838                .and_then(|ctx| ctx.new_stream().ok())
4839                .ok_or(ArrowSchurGpuFailure::Unavailable)?;
4840            let solver =
4841                DnHandle::new(stream.clone()).map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4842            let blas =
4843                CudaBlas::new(stream.clone()).map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4844
4845            // Base blocks at ridge 0 (ridge-independent); g_t stacked (n·d).
4846            let (d_base_host, b_host, g_host) = pack_host(sys, 0.0);
4847            let base_b_dev = stream
4848                .clone_htod(&b_host)
4849                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4850            let g_t_dev = stream
4851                .clone_htod(&g_host)
4852                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4853
4854            // Border H_ββ, column-major, NO ridge (ridge_beta added per trial).
4855            let hbb_base: Vec<f64> = {
4856                let mut tmp = Vec::with_capacity(k * k);
4857                for col in 0..k {
4858                    for row in 0..k {
4859                        tmp.push(sys.hbb[[row, col]]);
4860                    }
4861                }
4862                tmp
4863            };
4864            let base_hbb_dev = stream
4865                .clone_htod(&hbb_base)
4866                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4867
4868            let gb_host: Vec<f64> = sys.gb.iter().copied().collect();
4869            let ones_k_dev = stream
4870                .clone_htod(&vec![1.0_f64; k])
4871                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4872
4873            Ok(Self {
4874                n,
4875                d,
4876                k,
4877                beta_gauge_quotient: sys.beta_gauge_quotient.clone(),
4878                newton_schur_tikhonov_rel_floor,
4879                stream,
4880                solver,
4881                blas,
4882                d_base_host,
4883                base_b_dev,
4884                base_hbb_dev,
4885                g_t_dev,
4886                gb_host,
4887                ones_k_dev,
4888            })
4889        }
4890
4891        /// Factor the resident base blocks at `(ridge_t, ridge_beta)` and solve.
4892        /// Mirrors the full [`solve`] sequence, but sources `D`/`B`/`H_ββ`/`g`
4893        /// from the resident buffers (device-to-device copy into scratch) instead
4894        /// of re-uploading them, so it is bit-identical to [`solve`] at the same
4895        /// ridge while moving only the ridge scalars + re-diagonalised `D`.
4896        pub(super) fn refactor_and_solve(
4897            &self,
4898            ridge_t: f64,
4899            ridge_beta: f64,
4900        ) -> Result<ArrowSchurGpuSolution, ArrowSchurGpuFailure> {
4901            self.refactor_and_solve_from(ridge_t, ridge_beta, None)
4902        }
4903
4904        /// #2539: solve against ALREADY-BUILT factors for a fresh gradient. The
4905        /// checked entry point behind
4906        /// [`super::ResidentBaseArrowFrameHandle::solve_with_factors`]; the
4907        /// length contract is the same one
4908        /// [`Self::refactor_and_solve_with_gradient`] enforces, so a caller that
4909        /// caches factors cannot smuggle a differently shaped gradient past it.
4910        pub(super) fn solve_with_gradient_checked(
4911            &self,
4912            factors: &BaseRidgeFactors,
4913            g_t: &[f64],
4914            g_beta: &[f64],
4915        ) -> Result<ArrowSchurGpuSolution, ArrowSchurGpuFailure> {
4916            self.check_gradient_lens(g_t, g_beta)?;
4917            self.solve_with_factors(factors, Some((g_t, g_beta)))
4918        }
4919
4920        fn check_gradient_lens(
4921            &self,
4922            g_t: &[f64],
4923            g_beta: &[f64],
4924        ) -> Result<(), ArrowSchurGpuFailure> {
4925            if g_t.len() != self.n * self.d {
4926                return Err(ArrowSchurGpuFailure::SchurFactorFailed {
4927                    reason: format!(
4928                        "base-resident frame row gradient length mismatch: got {}, expected {}",
4929                        g_t.len(),
4930                        self.n * self.d
4931                    ),
4932                });
4933            }
4934            if g_beta.len() != self.k {
4935                return Err(ArrowSchurGpuFailure::SchurFactorFailed {
4936                    reason: format!(
4937                        "base-resident frame border gradient length mismatch: got {}, expected {}",
4938                        g_beta.len(),
4939                        self.k
4940                    ),
4941                });
4942            }
4943            Ok(())
4944        }
4945
4946        /// The shared ridge-dependent factor+solve. `gradient` selects where the
4947        /// right-hand side comes from: `None` reads the resident `g_t_dev` /
4948        /// `gb_host` captured at construction (the LM ridge ladder), `Some`
4949        /// uploads a fresh one (an inner Newton iterate). Every step after the
4950        /// gradient sourcing is common, so the two callers are bit-identical at
4951        /// equal ridge and equal gradient.
4952        ///
4953        /// It is the composition `factor_at` ∘ `solve_with_factors` and holds no
4954        /// arithmetic of its own, so a caller that keeps the factors (an inner
4955        /// Newton at an unchanged ridge) gets numbers identical to a caller that
4956        /// re-derives them here.
4957        fn refactor_and_solve_from(
4958            &self,
4959            ridge_t: f64,
4960            ridge_beta: f64,
4961            gradient: Option<(&[f64], &[f64])>,
4962        ) -> Result<ArrowSchurGpuSolution, ArrowSchurGpuFailure> {
4963            let factors = self.factor_at(ridge_t, ridge_beta)?;
4964            self.solve_with_factors(&factors, gradient)
4965        }
4966
4967        /// The GRADIENT-INDEPENDENT ridge-dependent factor work: `POTRF(D +
4968        /// ρ_t I)`, `Y_i = L_i^{-1} B_i`, the Schur reduction `S = (H_ββ + ρ_β I)
4969        /// − Σ Y_iᵀ Y_i` and its `POTRF`, plus `log|H|`. Nothing here reads a
4970        /// gradient, so the result is a complete, reusable factor state for
4971        /// `(ridge_t, ridge_beta)` (#2539): an inner Newton whose ridge did not
4972        /// move re-solves against it and pays no POTRF at all.
4973        ///
4974        /// Sources `D`/`B`/`H_ββ` from the resident base buffers, so the only
4975        /// host→device traffic is the `n·d·d` re-diagonalised `D`.
4976        pub(super) fn factor_at(
4977            &self,
4978            ridge_t: f64,
4979            ridge_beta: f64,
4980        ) -> Result<BaseRidgeFactors, ArrowSchurGpuFailure> {
4981            if ridge_t.is_nan() || ridge_beta.is_nan() {
4982                return Err(ArrowSchurGpuFailure::SchurFactorFailed {
4983                    reason: "ridge is NaN".to_string(),
4984                });
4985            }
4986            let (n, d, k) = (self.n, self.d, self.k);
4987
4988            // ----- D + ridge_t·I: add on a host copy (tiny) and upload as work L.
4989            // Tile i is column-major d×d, so its diagonal entries are at
4990            // i·d·d + j·(d+1) — matching pack_block's `value += ridge_t` on r==col.
4991            let mut d_ridged = self.d_base_host.clone();
4992            for i in 0..n {
4993                for j in 0..d {
4994                    d_ridged[i * d * d + j * (d + 1)] += ridge_t;
4995                }
4996            }
4997            let mut l_dev = self
4998                .stream
4999                .clone_htod(&d_ridged)
5000                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5001
5002            // POTRF(D) → L_i. A non-PD pivot is the LM escalation's signal.
5003            let info_host = potrf_batched(&self.solver, &self.stream, d, n, &mut l_dev)?;
5004            if let Some(idx) = info_host.iter().position(|info| *info != 0) {
5005                let base = idx * d * d;
5006                return Err(ArrowSchurGpuFailure::RidgeBumpRequired {
5007                    row: idx,
5008                    // The tile already carries ridge_t on its diagonal, so the
5009                    // Gershgorin bound is taken at ridge 0 (see
5010                    // `ridge_bump_to_make_pd_colmajor`).
5011                    bump: super::ridge_bump_to_make_pd_colmajor(&d_ridged[base..base + d * d], d),
5012                });
5013            }
5014
5015            // ----- Y_i = L_i^{-1} B_i on a device copy of the resident base B.
5016            let mut y_dev = self
5017                .stream
5018                .alloc_zeros::<f64>(n * d * k)
5019                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5020            self.stream
5021                .memcpy_dtod(&self.base_b_dev, &mut y_dev)
5022                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5023            trsm_batched_lower_inplace(&self.blas, &self.stream, d, n, k, &l_dev, &mut y_dev)?;
5024
5025            // ----- Schur S = (H_ββ + ridge_β I) − Σ Y_iᵀ Y_i.
5026            let mut schur_dev = self
5027                .stream
5028                .alloc_zeros::<f64>(k * k)
5029                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5030            self.stream
5031                .memcpy_dtod(&self.base_hbb_dev, &mut schur_dev)
5032                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5033            // schur diag += ridge_β (column-major k×k → stride k+1 over the ones).
5034            device_axpy_strided(
5035                &self.blas,
5036                &self.stream,
5037                k,
5038                ridge_beta,
5039                &self.ones_k_dev,
5040                1,
5041                &mut schur_dev,
5042                k + 1,
5043            )?;
5044            accumulate_schur_reduce_only(&self.blas, d, k, n, &y_dev, &mut schur_dev)?;
5045            let canonical_factor = if self.beta_gauge_quotient.is_some()
5046                || self.newton_schur_tikhonov_rel_floor.is_some()
5047            {
5048                let mut schur_host = self
5049                    .stream
5050                    .clone_dtoh(&schur_dev)
5051                    .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5052                let factorized = canonicalize_device_beta_factor(
5053                    self.beta_gauge_quotient.as_ref(),
5054                    self.newton_schur_tikhonov_rel_floor,
5055                    k,
5056                    &mut schur_host,
5057                )?;
5058                schur_dev = self
5059                    .stream
5060                    .clone_htod(&schur_host)
5061                    .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5062                factorized
5063            } else {
5064                false
5065            };
5066
5067            // ----- Factor S_β.
5068            if !canonical_factor {
5069                let info = potrf_single(&self.solver, &self.stream, k, &mut schur_dev)?;
5070                if info != 0 {
5071                    return Err(ArrowSchurGpuFailure::SchurFactorFailed {
5072                        reason: format!("Schur Cholesky failed at pivot {info}"),
5073                    });
5074                }
5075            }
5076
5077            // log|H| = 2 Σ log L_{i,jj} + 2 Σ log L_{S,aa}: a property of the
5078            // factored Hessian, so it belongs to the factors, not the solve.
5079            let l_local_host = self
5080                .stream
5081                .clone_dtoh(&l_dev)
5082                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5083            let l_schur_host = self
5084                .stream
5085                .clone_dtoh(&schur_dev)
5086                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5087            let mut log_det = 0.0_f64;
5088            for i in 0..n {
5089                let base = i * d * d;
5090                for j in 0..d {
5091                    log_det += l_local_host[base + j * d + j].ln();
5092                }
5093            }
5094            for j in 0..k {
5095                log_det += l_schur_host[j * k + j].ln();
5096            }
5097            log_det *= 2.0;
5098
5099            Ok(BaseRidgeFactors {
5100                ridge_t,
5101                ridge_beta,
5102                l_dev,
5103                y_dev,
5104                schur_dev,
5105                log_det_hessian: log_det,
5106            })
5107        }
5108
5109        /// Solve `(H + ridge)·δ = −gradient` against factors already produced by
5110        /// [`Self::factor_at`]. No POTRF and no `B` whitening run here — the only
5111        /// host→device traffic is the gradient (`n·d + k` doubles) and the only
5112        /// readback is `δ`.
5113        ///
5114        /// `gradient` sources the right-hand side exactly as
5115        /// [`Self::refactor_and_solve_from`] documents: `None` is the resident
5116        /// gradient captured at construction, `Some` a fresh per-iterate one.
5117        pub(super) fn solve_with_factors(
5118            &self,
5119            factors: &BaseRidgeFactors,
5120            gradient: Option<(&[f64], &[f64])>,
5121        ) -> Result<ArrowSchurGpuSolution, ArrowSchurGpuFailure> {
5122            let (n, d, k) = (self.n, self.d, self.k);
5123            let l_dev = &factors.l_dev;
5124            let y_dev = &factors.y_dev;
5125            let schur_dev = &factors.schur_dev;
5126
5127            // ----- u_i = L_i^{-1} g_i. The resident base g_t is a device-to-device
5128            // copy; a caller-supplied gradient is the one extra H2D (`n·d`
5129            // doubles, the same order as the ridged `D` the factor pass uploads).
5130            let mut u_dev = match gradient {
5131                Some((g_t, _)) => self
5132                    .stream
5133                    .clone_htod(g_t)
5134                    .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
5135                None => {
5136                    let mut resident = self
5137                        .stream
5138                        .alloc_zeros::<f64>(n * d)
5139                        .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5140                    self.stream
5141                        .memcpy_dtod(&self.g_t_dev, &mut resident)
5142                        .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5143                    resident
5144                }
5145            };
5146            trsm_batched_lower_inplace(&self.blas, &self.stream, d, n, 1, l_dev, &mut u_dev)?;
5147
5148            // ----- Schur RHS = −g_β + Σ_i Y_iᵀ u_i.
5149            let rhs_init: Vec<f64> = match gradient {
5150                Some((_, g_beta)) => g_beta.iter().map(|v| -v).collect(),
5151                None => self.gb_host.iter().map(|v| -v).collect(),
5152            };
5153            let mut rhs_dev = self
5154                .stream
5155                .clone_htod(&rhs_init)
5156                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5157            accumulate_schur_rhs_only(&self.blas, d, k, n, y_dev, &u_dev, &mut rhs_dev)?;
5158            if self.beta_gauge_quotient.is_some() {
5159                let projected_rhs = project_device_beta_vector(
5160                    self.beta_gauge_quotient.as_ref(),
5161                    k,
5162                    self.stream
5163                        .clone_dtoh(&rhs_dev)
5164                        .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
5165                )?;
5166                rhs_dev = self
5167                    .stream
5168                    .clone_htod(&projected_rhs)
5169                    .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5170            }
5171
5172            // ----- Solve δβ = L_S^{-T} L_S^{-1} rhs.
5173            trsm_single(
5174                &self.blas,
5175                &self.stream,
5176                k,
5177                schur_dev,
5178                &mut rhs_dev,
5179                false,
5180                false,
5181            )?;
5182            trsm_single(
5183                &self.blas,
5184                &self.stream,
5185                k,
5186                schur_dev,
5187                &mut rhs_dev,
5188                false,
5189                true,
5190            )?;
5191            let delta_beta_host = project_device_beta_vector(
5192                self.beta_gauge_quotient.as_ref(),
5193                k,
5194                self.stream
5195                    .clone_dtoh(&rhs_dev)
5196                    .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
5197            )?;
5198            if self.beta_gauge_quotient.is_some() {
5199                rhs_dev = self
5200                    .stream
5201                    .clone_htod(&delta_beta_host)
5202                    .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5203            }
5204            let delta_beta = Array1::from_vec(delta_beta_host);
5205
5206            // ----- Back-sub δt_i = −L_i^{-T}(u_i + Y_i δβ).
5207            accumulate_back_sub_rhs(&self.blas, d, k, n, y_dev, &rhs_dev, &mut u_dev)?;
5208            trsm_batched_lower_inplace_transposed(
5209                &self.blas,
5210                &self.stream,
5211                d,
5212                n,
5213                1,
5214                l_dev,
5215                &mut u_dev,
5216            )?;
5217            let x_host = self
5218                .stream
5219                .clone_dtoh(&u_dev)
5220                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5221            let mut delta_t = Array1::<f64>::zeros(n * d);
5222            for (i, v) in x_host.iter().enumerate() {
5223                delta_t[i] = -*v;
5224            }
5225
5226            Ok(ArrowSchurGpuSolution {
5227                delta_t,
5228                delta_beta,
5229                log_det_hessian: factors.log_det_hessian,
5230            })
5231        }
5232    }
5233
5234    /// The ridge-derived factor state of a [`ResidentBaseArrowFrame`], resident
5235    /// on the device and keyed by the `(ridge_t, ridge_beta)` it was built at
5236    /// (#2539).
5237    ///
5238    /// The base frame holds the ridge-INDEPENDENT blocks; this holds everything
5239    /// the ridge does move. Separating them is what lets an inner Newton keep
5240    /// the accepted-step behaviour it had before the base frame was adopted: a
5241    /// same-ridge iterate re-solves against these factors with no POTRF, and a
5242    /// ridge change rebuilds only this — from blocks that never leave the
5243    /// device.
5244    pub(super) struct BaseRidgeFactors {
5245        ridge_t: f64,
5246        ridge_beta: f64,
5247        /// Per-row lower Cholesky factors `L_i` of `H_tt + ρ_t I` (`n` tiles of
5248        /// `d×d`, column-major).
5249        l_dev: CudaSlice<f64>,
5250        /// Whitened cross blocks `Y_i = L_i^{-1} H_tβ^(i)` (`n` tiles of `d×k`).
5251        y_dev: CudaSlice<f64>,
5252        /// Lower Cholesky factor `L_S` of `S_β = H_ββ + ρ_β I − Σ_i Y_iᵀ Y_i`.
5253        schur_dev: CudaSlice<f64>,
5254        /// `log|H|`, a property of the factored Hessian alone.
5255        log_det_hessian: f64,
5256    }
5257
5258    impl BaseRidgeFactors {
5259        /// The ridge pair these factors were built at. A caller caching them
5260        /// compares against this rather than tracking the key separately, so the
5261        /// key can never drift from the factors it names.
5262        #[inline]
5263        pub(super) fn matches_ridge(&self, ridge_t: f64, ridge_beta: f64) -> bool {
5264            self.ridge_t == ridge_t && self.ridge_beta == ridge_beta
5265        }
5266    }
5267
5268    pub(super) fn solve_fused(
5269        sys: &ArrowSchurSystem,
5270        ridge_t: f64,
5271        ridge_beta: f64,
5272        newton_schur_tikhonov_rel_floor: Option<f64>,
5273    ) -> Result<ArrowSchurGpuSolution, ArrowSchurGpuFailure> {
5274        let n = sys.rows.len();
5275        let d = sys.d;
5276        let k = sys.k;
5277        let plan = crate::gpu_kernels::arrow_schur_nvrtc::plan_fused_launch(n, d, k)
5278            .ok_or(ArrowSchurGpuFailure::Unavailable)?;
5279        let p_max = plan.p_max;
5280        let r_template = plan.r_template;
5281
5282        let runtime = gam_gpu::linalg_dispatch::route_through_gpu(
5283            gam_gpu::linalg_dispatch::DispatchOp::SmallDenseBatchedPotrf { p: d, batch: n },
5284        )
5285        .ok_or(ArrowSchurGpuFailure::Unavailable)?;
5286        let ctx = gam_gpu::device_runtime::cuda_context_for(runtime.device.ordinal)
5287            .ok_or(ArrowSchurGpuFailure::Unavailable)?;
5288        let stream = ctx
5289            .new_stream()
5290            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5291        let cap = &runtime.device.capability;
5292        let key = crate::gpu_kernels::arrow_schur_nvrtc::FusedModuleCacheKey {
5293            cc_major: cap.compute_major,
5294            cc_minor: cap.compute_minor,
5295            p_max: p_max as u32,
5296            r_template: r_template as u32,
5297        };
5298        let module = fused_module_for(&ctx, key)?;
5299        let forward = module
5300            .load_function("arrow_schur_forward_pgroup")
5301            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5302        let back_sub = module
5303            .load_function("arrow_schur_back_sub_pgroup")
5304            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5305
5306        // ----- Upload packed D, B, g -----
5307        let (d_host, b_host, g_host) = pack_fused_host(sys, ridge_t, p_max, r_template);
5308        let d_dev = stream
5309            .clone_htod(&d_host)
5310            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5311        let b_dev = stream
5312            .clone_htod(&b_host)
5313            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5314        let g_dev = stream
5315            .clone_htod(&g_host)
5316            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5317        let mut l_out = stream
5318            .alloc_zeros::<f64>(n * p_max * p_max)
5319            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5320        let mut u_out = stream
5321            .alloc_zeros::<f64>(n * p_max)
5322            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5323        let mut y_out = stream
5324            .alloc_zeros::<f64>(n * p_max * r_template)
5325            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5326        let mut partial_s = stream
5327            .alloc_zeros::<f64>(plan.partial_s_doubles)
5328            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5329        let mut partial_r = stream
5330            .alloc_zeros::<f64>(plan.partial_r_doubles)
5331            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5332        let mut status_dev = stream
5333            .alloc_zeros::<i32>(n)
5334            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5335
5336        // ----- Launch forward kernel: 1 block per row, P_MAX threads -----
5337        let cfg = LaunchConfig {
5338            grid_dim: (plan.blocks, 1, 1),
5339            block_dim: (plan.threads_per_block, 1, 1),
5340            shared_mem_bytes: 0,
5341        };
5342        let n_i32 = to_i32(n).ok_or(ArrowSchurGpuFailure::Unavailable)?;
5343        let p_i32 = to_i32(d).ok_or(ArrowSchurGpuFailure::Unavailable)?;
5344        let r_i32 = to_i32(k).ok_or(ArrowSchurGpuFailure::Unavailable)?;
5345        let ridge_arg = ridge_t;
5346        {
5347            let mut builder = stream.launch_builder(&forward);
5348            builder
5349                .arg(&d_dev)
5350                .arg(&b_dev)
5351                .arg(&g_dev)
5352                .arg(&n_i32)
5353                .arg(&p_i32)
5354                .arg(&r_i32)
5355                .arg(&ridge_arg)
5356                .arg(&mut l_out)
5357                .arg(&mut u_out)
5358                .arg(&mut y_out)
5359                .arg(&mut partial_s)
5360                .arg(&mut partial_r)
5361                .arg(&mut status_dev);
5362            // SAFETY: all buffers were just allocated on `stream` with sizes
5363            // derived from `plan`; kernel parameter list matches the
5364            // FORWARD_KERNEL_SOURCE signature.
5365            unsafe { builder.launch(cfg) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5366        }
5367        stream
5368            .synchronize()
5369            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5370
5371        // ----- Check per-block pivot status -----
5372        let status_host = stream
5373            .clone_dtoh(&status_dev)
5374            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5375        if let Some(row) = status_host.iter().position(|s| *s != 0) {
5376            // The NVRTC kernel's status code is a 1-based pivot row index, not
5377            // a magnitude; size the bump from the block (Gershgorin λ_min
5378            // bound) so a strongly indefinite block recovers in one retry.
5379            return Err(ArrowSchurGpuFailure::RidgeBumpRequired {
5380                row,
5381                bump: super::ridge_bump_to_make_pd(sys.rows[row].htt.view(), ridge_t),
5382            });
5383        }
5384
5385        // ----- Reduce partials on host into S_β and r_β -----
5386        let partial_s_host = stream
5387            .clone_dtoh(&partial_s)
5388            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5389        let partial_r_host = stream
5390            .clone_dtoh(&partial_r)
5391            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5392        let mut schur_host = vec![0.0_f64; k * k];
5393        for col in 0..k {
5394            for row in 0..k {
5395                let mut v = sys.hbb[[row, col]];
5396                if row == col {
5397                    v += ridge_beta;
5398                }
5399                schur_host[col * k + row] = v;
5400            }
5401        }
5402        let mut rhs_host: Vec<f64> = sys.gb.iter().map(|v| -v).collect();
5403        for i in 0..n {
5404            // partial_S[i] stride is R_TEMPLATE × R_TEMPLATE column-major; we
5405            // only read the leading (k × k) sub-block.
5406            let s_base = i * r_template * r_template;
5407            for col in 0..k {
5408                let col_base = s_base + col * r_template;
5409                let dst_col_base = col * k;
5410                for row in 0..k {
5411                    schur_host[dst_col_base + row] -= partial_s_host[col_base + row];
5412                }
5413            }
5414            let r_base = i * r_template;
5415            for a in 0..k {
5416                rhs_host[a] += partial_r_host[r_base + a];
5417            }
5418        }
5419        let canonical_factor = canonicalize_device_beta_factor(
5420            sys.beta_gauge_quotient.as_ref(),
5421            newton_schur_tikhonov_rel_floor,
5422            sys.k,
5423            &mut schur_host,
5424        )?;
5425        rhs_host = project_device_beta_vector(sys.beta_gauge_quotient.as_ref(), sys.k, rhs_host)?;
5426
5427        // ----- Factor S_β on device (cuSOLVER), solve for δβ -----
5428        let mut schur_dev = stream
5429            .clone_htod(&schur_host)
5430            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5431        let mut rhs_dev = stream
5432            .clone_htod(&rhs_host)
5433            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5434        let solver =
5435            DnHandle::new(stream.clone()).map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5436        let blas = CudaBlas::new(stream.clone()).map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5437        if !canonical_factor {
5438            let info = potrf_single(&solver, &stream, k, &mut schur_dev)?;
5439            if info != 0 {
5440                return Err(ArrowSchurGpuFailure::SchurFactorFailed {
5441                    reason: format!("fused Schur Cholesky failed at pivot {info}"),
5442                });
5443            }
5444        }
5445        trsm_single(&blas, &stream, k, &schur_dev, &mut rhs_dev, false, false)?;
5446        trsm_single(&blas, &stream, k, &schur_dev, &mut rhs_dev, false, true)?;
5447        let delta_beta_host = project_device_beta_vector(
5448            sys.beta_gauge_quotient.as_ref(),
5449            sys.k,
5450            stream
5451                .clone_dtoh(&rhs_dev)
5452                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
5453        )?;
5454        if sys.beta_gauge_quotient.is_some() {
5455            rhs_dev = stream
5456                .clone_htod(&delta_beta_host)
5457                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5458        }
5459        let delta_beta = Array1::from_vec(delta_beta_host);
5460
5461        // ----- Layer E: launch back-sub kernel using persisted L, u, Y -----
5462        let mut delta_t_dev = stream
5463            .alloc_zeros::<f64>(n * p_max)
5464            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5465        let back_cfg = LaunchConfig {
5466            grid_dim: (plan.blocks, 1, 1),
5467            block_dim: (plan.threads_per_block, 1, 1),
5468            shared_mem_bytes: 0,
5469        };
5470        {
5471            let mut builder = stream.launch_builder(&back_sub);
5472            builder
5473                .arg(&l_out)
5474                .arg(&u_out)
5475                .arg(&y_out)
5476                .arg(&rhs_dev)
5477                .arg(&n_i32)
5478                .arg(&p_i32)
5479                .arg(&r_i32)
5480                .arg(&mut delta_t_dev);
5481            // SAFETY: kernel parameter list matches FORWARD_KERNEL_SOURCE
5482            // back-sub signature; `rhs_dev` holds δβ in the leading k entries
5483            // (R_TEMPLATE-strided indexing is column 0..k of the R-vector).
5484            unsafe { builder.launch(back_cfg) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5485        }
5486        stream
5487            .synchronize()
5488            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5489
5490        let delta_t_host = stream
5491            .clone_dtoh(&delta_t_dev)
5492            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5493        let mut delta_t = Array1::<f64>::zeros(n * d);
5494        for i in 0..n {
5495            let src_base = i * p_max;
5496            let dst_base = i * d;
5497            for r in 0..d {
5498                delta_t[dst_base + r] = delta_t_host[src_base + r];
5499            }
5500        }
5501
5502        // ----- log|H| = 2·Σ log L_{i,jj} + 2·Σ log R_{β,aa} -----
5503        let l_local_host = stream
5504            .clone_dtoh(&l_out)
5505            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5506        let l_schur_host = stream
5507            .clone_dtoh(&schur_dev)
5508            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5509        let mut log_det = 0.0_f64;
5510        for i in 0..n {
5511            let base = i * p_max * p_max;
5512            for j in 0..d {
5513                log_det += l_local_host[base + j * p_max + j].ln();
5514            }
5515        }
5516        for j in 0..k {
5517            log_det += l_schur_host[j * k + j].ln();
5518        }
5519        log_det *= 2.0;
5520
5521        Ok(ArrowSchurGpuSolution {
5522            delta_t,
5523            delta_beta,
5524            log_det_hessian: log_det,
5525        })
5526    }
5527
5528    /// Pre-compute `Y_i = L_i^{-1} H_tβ^(i)` via the fused forward kernel and
5529    /// return a closure that evaluates the full Schur matvec
5530    /// `S·x = (H_ββ + ρ·I)·x − Σ_i Y_i^T (Y_i·x)` for each PCG iteration.
5531    ///
5532    /// The `Y_i` factors are kept in a host-side buffer after one GPU forward
5533    /// pass. Each matvec call runs O(N·d·K) host loops over the pre-computed
5534    /// buffer plus an optional `H_ββ·x` call (matrix-free or dense). This is
5535    /// the first landing of the GPU matvec; a future iteration can move the
5536    /// `Y_i·x` / `Y_i^T z_i` steps to cuBLAS batched GEMV.
5537    pub(super) fn build_schur_matvec_backend(
5538        sys: &ArrowSchurSystem,
5539        ridge_t: f64,
5540        ridge_beta: f64,
5541    ) -> Result<crate::arrow_schur::GpuSchurMatvec, super::ArrowSchurGpuFailure> {
5542        let n = sys.rows.len();
5543        let d = sys.d;
5544        let k = sys.k;
5545        let plan = crate::gpu_kernels::arrow_schur_nvrtc::plan_fused_launch(n, d, k)
5546            .ok_or(super::ArrowSchurGpuFailure::Unavailable)?;
5547        let p_max = plan.p_max;
5548        let r_template = plan.r_template;
5549
5550        let runtime = gam_gpu::linalg_dispatch::route_through_gpu(
5551            gam_gpu::linalg_dispatch::DispatchOp::SmallDenseBatchedPotrf { p: d, batch: n },
5552        )
5553        .ok_or(super::ArrowSchurGpuFailure::Unavailable)?;
5554        let ctx = gam_gpu::device_runtime::cuda_context_for(runtime.device.ordinal)
5555            .ok_or(super::ArrowSchurGpuFailure::Unavailable)?;
5556        let stream = ctx
5557            .new_stream()
5558            .map_err(|_| super::ArrowSchurGpuFailure::Unavailable)?;
5559        let cap = &runtime.device.capability;
5560        let key = crate::gpu_kernels::arrow_schur_nvrtc::FusedModuleCacheKey {
5561            cc_major: cap.compute_major,
5562            cc_minor: cap.compute_minor,
5563            p_max: p_max as u32,
5564            r_template: r_template as u32,
5565        };
5566        let module = fused_module_for(&ctx, key)?;
5567        let forward = module
5568            .load_function("arrow_schur_forward_pgroup")
5569            .map_err(|_| super::ArrowSchurGpuFailure::Unavailable)?;
5570
5571        let (d_host, b_host, g_host) = pack_fused_host(sys, ridge_t, p_max, r_template);
5572        let d_dev = stream
5573            .clone_htod(&d_host)
5574            .map_err(|_| super::ArrowSchurGpuFailure::Unavailable)?;
5575        let b_dev = stream
5576            .clone_htod(&b_host)
5577            .map_err(|_| super::ArrowSchurGpuFailure::Unavailable)?;
5578        let g_dev = stream
5579            .clone_htod(&g_host)
5580            .map_err(|_| super::ArrowSchurGpuFailure::Unavailable)?;
5581        let mut l_out = stream
5582            .alloc_zeros::<f64>(n * p_max * p_max)
5583            .map_err(|_| super::ArrowSchurGpuFailure::Unavailable)?;
5584        let mut u_out = stream
5585            .alloc_zeros::<f64>(n * p_max)
5586            .map_err(|_| super::ArrowSchurGpuFailure::Unavailable)?;
5587        let mut y_out = stream
5588            .alloc_zeros::<f64>(n * p_max * r_template)
5589            .map_err(|_| super::ArrowSchurGpuFailure::Unavailable)?;
5590        let mut partial_s = stream
5591            .alloc_zeros::<f64>(plan.partial_s_doubles)
5592            .map_err(|_| super::ArrowSchurGpuFailure::Unavailable)?;
5593        let mut partial_r = stream
5594            .alloc_zeros::<f64>(plan.partial_r_doubles)
5595            .map_err(|_| super::ArrowSchurGpuFailure::Unavailable)?;
5596        let mut status_dev = stream
5597            .alloc_zeros::<i32>(n)
5598            .map_err(|_| super::ArrowSchurGpuFailure::Unavailable)?;
5599
5600        let cfg = LaunchConfig {
5601            grid_dim: (plan.blocks, 1, 1),
5602            block_dim: (plan.threads_per_block, 1, 1),
5603            shared_mem_bytes: 0,
5604        };
5605        let n_i32 = to_i32(n).ok_or(super::ArrowSchurGpuFailure::Unavailable)?;
5606        let p_i32 = to_i32(d).ok_or(super::ArrowSchurGpuFailure::Unavailable)?;
5607        let r_i32 = to_i32(k).ok_or(super::ArrowSchurGpuFailure::Unavailable)?;
5608        let ridge_arg = ridge_t;
5609        {
5610            let mut builder = stream.launch_builder(&forward);
5611            builder
5612                .arg(&d_dev)
5613                .arg(&b_dev)
5614                .arg(&g_dev)
5615                .arg(&n_i32)
5616                .arg(&p_i32)
5617                .arg(&r_i32)
5618                .arg(&ridge_arg)
5619                .arg(&mut l_out)
5620                .arg(&mut u_out)
5621                .arg(&mut y_out)
5622                .arg(&mut partial_s)
5623                .arg(&mut partial_r)
5624                .arg(&mut status_dev);
5625            // SAFETY: all buffers were allocated on `stream` with sizes
5626            // derived from `plan`; parameter list matches FORWARD_KERNEL_SOURCE.
5627            unsafe { builder.launch(cfg) }.map_err(|_| super::ArrowSchurGpuFailure::Unavailable)?;
5628        }
5629        stream
5630            .synchronize()
5631            .map_err(|_| super::ArrowSchurGpuFailure::Unavailable)?;
5632
5633        let status_host = stream
5634            .clone_dtoh(&status_dev)
5635            .map_err(|_| super::ArrowSchurGpuFailure::Unavailable)?;
5636        if let Some(row) = status_host.iter().position(|s| *s != 0) {
5637            // Status code is a 1-based pivot row index, not a magnitude; size
5638            // the bump from the block (Gershgorin λ_min bound) so a strongly
5639            // indefinite block recovers in one retry.
5640            return Err(super::ArrowSchurGpuFailure::RidgeBumpRequired {
5641                row,
5642                bump: super::ridge_bump_to_make_pd(sys.rows[row].htt.view(), ridge_t),
5643            });
5644        }
5645
5646        // Download Y_i factors: n × p_max × r_template column-major per block.
5647        let y_host = stream
5648            .clone_dtoh(&y_out)
5649            .map_err(|_| super::ArrowSchurGpuFailure::Unavailable)?;
5650
5651        // Capture H_ββ data for the closure. Use the matrix-free hook if present
5652        // (SAE-manifold callers), otherwise fall back to the dense matrix rows.
5653        let hbb_host: Vec<f64> = sys.hbb.iter().copied().collect();
5654        let hbb_is_kk = sys.hbb.dim() == (k, k);
5655        let hbb_matvec_opt = sys.hbb_matvec.clone();
5656
5657        let closure: crate::arrow_schur::GpuSchurMatvec =
5658            Arc::new(move |x: &Array1<f64>, out: &mut Array1<f64>| {
5659                assert_eq!(x.len(), k, "gpu_schur_matvec: x.len() != k");
5660                assert_eq!(out.len(), k, "gpu_schur_matvec: out.len() != k");
5661
5662                // (H_ββ + ρ·I)·x into out.
5663                if let Some(ref mv) = hbb_matvec_opt {
5664                    mv(x.view(), out);
5665                    for a in 0..k {
5666                        out[a] += ridge_beta * x[a];
5667                    }
5668                } else if hbb_is_kk {
5669                    // hbb_host row-major: hbb[a, b] = hbb_host[a * k + b].
5670                    for a in 0..k {
5671                        let mut acc = ridge_beta * x[a];
5672                        for b in 0..k {
5673                            acc += hbb_host[a * k + b] * x[b];
5674                        }
5675                        out[a] = acc;
5676                    }
5677                } else {
5678                    for a in 0..k {
5679                        out[a] = ridge_beta * x[a];
5680                    }
5681                }
5682
5683                // out[c] -= Σ_i (Y_i^T (Y_i·x))[c].
5684                // Y_i column-major at y_host[i·p_max·r_template + col·p_max + row].
5685                let mut z = vec![0.0_f64; d];
5686                for i in 0..n {
5687                    let y_base = i * p_max * r_template;
5688                    for r in 0..d {
5689                        let mut acc = 0.0;
5690                        for c in 0..k {
5691                            acc += y_host[y_base + c * p_max + r] * x[c];
5692                        }
5693                        z[r] = acc;
5694                    }
5695                    for c in 0..k {
5696                        let mut acc = 0.0;
5697                        for r in 0..d {
5698                            acc += y_host[y_base + c * p_max + r] * z[r];
5699                        }
5700                        out[c] -= acc;
5701                    }
5702                }
5703            });
5704
5705        Ok(closure)
5706    }
5707
5708    // ── #1017/#1026 frames-engaged device PCG ──────────────────────────────
5709
5710    struct DeviceSaeFrameBuffers {
5711        // Smooth `λ S_k ⊗ I_{r_k}`.
5712        s_off: CudaSlice<i32>,
5713        s_m: CudaSlice<i32>,
5714        s_r: CudaSlice<i32>,
5715        s_ptr: CudaSlice<i32>,
5716        s_data: CudaSlice<f64>,
5717        s_blocks: usize,
5718        // Data `G_{ij} ⊗ W_{ij}`.
5719        g_off_i: CudaSlice<i32>,
5720        g_off_j: CudaSlice<i32>,
5721        g_ri: CudaSlice<i32>,
5722        g_rj: CudaSlice<i32>,
5723        g_mj: CudaSlice<i32>,
5724        g_ptr: CudaSlice<i32>,
5725        g_data: CudaSlice<f64>,
5726        w_ptr: CudaSlice<i32>,
5727        w_data: CudaSlice<f64>,
5728        g_group_ptr: CudaSlice<i32>,
5729        g_group_blocks: CudaSlice<i32>,
5730        g_group_off: CudaSlice<i32>,
5731        g_group_width: CudaSlice<i32>,
5732        g_groups: usize,
5733        g_max_work: usize,
5734        // Per-row dense cross-block H_tβ^(i) + row q + factored ainv.
5735        htb_ptr: CudaSlice<i32>,
5736        htb: CudaSlice<f64>,
5737        q_of: CudaSlice<i32>,
5738        ainv: CudaSlice<f64>,
5739        hvec: CudaSlice<f64>,
5740        svec: CudaSlice<f64>,
5741        // #1017 2-stage deterministic scatter scratch: partial[n_chunks × k]
5742        // holds each row-chunk's reduced-Schur contribution before the fixed-order
5743        // reduce. Ridge-independent shape (derived from n_rows), so it is allocated
5744        // once with the resident frame and reused across every apply.
5745        scatter_partial: CudaSlice<f64>,
5746        n_chunks: usize,
5747        rows_per_chunk: usize,
5748        n_rows: usize,
5749        k: usize,
5750        max_q: usize,
5751    }
5752
5753    /// #1017 chunking for the 2-stage deterministic reduced-Schur scatter: split
5754    /// the `n_rows` row reduction into contiguous chunks so stage 1 launches
5755    /// `⌈k/256⌉·n_chunks` CTAs (vs the single-strip `⌈k/256⌉`). ~128 chunks fills a
5756    /// 72-SM A10 with several CTAs even at small `k`, while keeping the partial
5757    /// buffer (`n_chunks·k`) small and the stage-2 reduce (`n_chunks` adds) cheap.
5758    fn scatter_chunking(n_rows: usize) -> (usize, usize) {
5759        let target_chunks = 128usize;
5760        let rows_per_chunk = n_rows.div_ceil(target_chunks).max(1);
5761        let n_chunks = n_rows.div_ceil(rows_per_chunk).max(1);
5762        (rows_per_chunk, n_chunks)
5763    }
5764
5765    fn flatten_device_sae_frame_data(
5766        sys: &ArrowSchurSystem,
5767        data: &DeviceSaePcgData,
5768        frame: &DeviceSaeFrameData,
5769        ridge_t: f64,
5770        stream: &Arc<CudaStream>,
5771    ) -> Result<DeviceSaeFrameBuffers, ArrowSchurGpuFailure> {
5772        // #1017: single-source the marshalling. The ridge-INDEPENDENT operands
5773        // (smooth `λ S_k`, framed `G ⊗ W`, dense per-row cross `H_tβ`) come from
5774        // `flatten_frame_host_operands`; only the ridge-DEPENDENT per-row `ainv`
5775        // is recomputed here. `upload_frame_buffers` performs the identical
5776        // host→device transfer the inline body used, so the resulting buffers are
5777        // byte-for-byte what this function produced before the split.
5778        let host = super::flatten_frame_host_operands(sys, data, frame)?;
5779        let ainv = super::compute_ainv_host(sys, &host.q_of, host.max_q, host.n_rows, ridge_t)?;
5780        upload_frame_buffers(&host, &ainv, stream)
5781    }
5782
5783    /// Upload the marshalled host operands (ridge-independent) plus the supplied
5784    /// per-row `ainv` (ridge-dependent) into device buffers. Shared by the
5785    /// per-trial [`flatten_device_sae_frame_data`] and the resident-frame build
5786    /// (which uploads a zero `ainv` placeholder once and overwrites only `ainv`
5787    /// per ladder trial), so both paths marshal through one code path.
5788    fn upload_frame_buffers(
5789        host: &super::FrameHostOperands,
5790        ainv: &[f64],
5791        stream: &Arc<CudaStream>,
5792    ) -> Result<DeviceSaeFrameBuffers, ArrowSchurGpuFailure> {
5793        let htod_i = |v: &[i32]| {
5794            stream
5795                .clone_htod(v)
5796                .map_err(|_| ArrowSchurGpuFailure::Unavailable)
5797        };
5798        let htod_f = |v: &[f64]| {
5799            stream
5800                .clone_htod(v)
5801                .map_err(|_| ArrowSchurGpuFailure::Unavailable)
5802        };
5803        let (rows_per_chunk, n_chunks) = scatter_chunking(host.n_rows);
5804        Ok(DeviceSaeFrameBuffers {
5805            s_off: htod_i(&host.s_off)?,
5806            s_m: htod_i(&host.s_m)?,
5807            s_r: htod_i(&host.s_r)?,
5808            s_ptr: htod_i(&host.s_ptr)?,
5809            s_data: htod_f(&host.s_data)?,
5810            s_blocks: host.s_blocks,
5811            g_off_i: htod_i(&host.g_off_i)?,
5812            g_off_j: htod_i(&host.g_off_j)?,
5813            g_ri: htod_i(&host.g_ri)?,
5814            g_rj: htod_i(&host.g_rj)?,
5815            g_mj: htod_i(&host.g_mj)?,
5816            g_ptr: htod_i(&host.g_ptr)?,
5817            g_data: htod_f(&host.g_data)?,
5818            w_ptr: htod_i(&host.w_ptr)?,
5819            w_data: htod_f(&host.w_data)?,
5820            g_group_ptr: htod_i(&host.g_group_ptr)?,
5821            g_group_blocks: htod_i(&host.g_group_blocks)?,
5822            g_group_off: htod_i(&host.g_group_off)?,
5823            g_group_width: htod_i(&host.g_group_width)?,
5824            g_groups: host.g_groups,
5825            g_max_work: host.g_max_work,
5826            htb_ptr: htod_i(&host.htb_ptr)?,
5827            htb: htod_f(&host.htb)?,
5828            q_of: htod_i(&host.q_of)?,
5829            ainv: htod_f(ainv)?,
5830            hvec: stream
5831                .alloc_zeros::<f64>(host.n_rows * host.max_q)
5832                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
5833            svec: stream
5834                .alloc_zeros::<f64>(host.n_rows * host.max_q)
5835                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
5836            scatter_partial: stream
5837                .alloc_zeros::<f64>(n_chunks * host.k)
5838                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
5839            n_chunks,
5840            rows_per_chunk,
5841            n_rows: host.n_rows,
5842            k: host.k,
5843            max_q: host.max_q,
5844        })
5845    }
5846
5847    /// Overwrite every ridge-independent resident operand in place. This is the
5848    /// accepted-nonlinear-iterate boundary: allocations persist, numerical
5849    /// content does not. A shape change declines so the caller replaces the
5850    /// complete frame rather than partially refreshing incompatible buffers.
5851    fn refresh_frame_buffers(
5852        host: &super::FrameHostOperands,
5853        buffers: &mut DeviceSaeFrameBuffers,
5854        stream: &Arc<CudaStream>,
5855    ) -> Result<(), ArrowSchurGpuFailure> {
5856        macro_rules! refresh {
5857            ($host:expr, $device:expr) => {{
5858                if $host.len() != $device.len() {
5859                    return Err(ArrowSchurGpuFailure::Unavailable);
5860                }
5861                stream
5862                    .memcpy_htod($host, &mut $device)
5863                    .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5864            }};
5865        }
5866        if host.s_blocks != buffers.s_blocks
5867            || host.g_groups != buffers.g_groups
5868            || host.g_max_work != buffers.g_max_work
5869            || host.n_rows != buffers.n_rows
5870            || host.k != buffers.k
5871            || host.max_q != buffers.max_q
5872        {
5873            return Err(ArrowSchurGpuFailure::Unavailable);
5874        }
5875        refresh!(&host.s_off, buffers.s_off);
5876        refresh!(&host.s_m, buffers.s_m);
5877        refresh!(&host.s_r, buffers.s_r);
5878        refresh!(&host.s_ptr, buffers.s_ptr);
5879        refresh!(&host.s_data, buffers.s_data);
5880        refresh!(&host.g_off_i, buffers.g_off_i);
5881        refresh!(&host.g_off_j, buffers.g_off_j);
5882        refresh!(&host.g_ri, buffers.g_ri);
5883        refresh!(&host.g_rj, buffers.g_rj);
5884        refresh!(&host.g_mj, buffers.g_mj);
5885        refresh!(&host.g_ptr, buffers.g_ptr);
5886        refresh!(&host.g_data, buffers.g_data);
5887        refresh!(&host.w_ptr, buffers.w_ptr);
5888        refresh!(&host.w_data, buffers.w_data);
5889        refresh!(&host.g_group_ptr, buffers.g_group_ptr);
5890        refresh!(&host.g_group_blocks, buffers.g_group_blocks);
5891        refresh!(&host.g_group_off, buffers.g_group_off);
5892        refresh!(&host.g_group_width, buffers.g_group_width);
5893        refresh!(&host.htb_ptr, buffers.htb_ptr);
5894        refresh!(&host.htb, buffers.htb);
5895        refresh!(&host.q_of, buffers.q_of);
5896        Ok(())
5897    }
5898
5899    fn sae_frame_penalty_diag_host(
5900        data: &DeviceSaePcgData,
5901        frame: &DeviceSaeFrameData,
5902        ridge_beta: f64,
5903    ) -> Result<Vec<f64>, ArrowSchurGpuFailure> {
5904        let mut diag = vec![ridge_beta; data.beta_dim];
5905        // Smooth: diag[off + ia·r + ib] += S[ia,ia].
5906        for (blk, &r) in data.smooth_blocks.iter().zip(frame.smooth_ranks.iter()) {
5907            let m = blk.factor_a.nrows();
5908            for ia in 0..m {
5909                let coeff = blk.factor_a[[ia, ia]];
5910                let base = blk.global_offset + ia * r;
5911                for ib in 0..r {
5912                    if base + ib >= diag.len() {
5913                        return Err(ArrowSchurGpuFailure::Unavailable);
5914                    }
5915                    diag[base + ib] += coeff;
5916                }
5917            }
5918        }
5919        // Data: on-diagonal atom blocks contribute g[li,li]·w[a,a].
5920        for blk in &frame.frame_blocks {
5921            if blk.atom_i != blk.atom_j {
5922                continue;
5923            }
5924            let r = frame.ranks[blk.atom_i];
5925            let off = frame.border_offsets[blk.atom_i];
5926            let (mi, mj) = blk.g.dim();
5927            for li in 0..mi.min(mj) {
5928                let gii = blk.g[[li, li]];
5929                let base = off + li * r;
5930                for a in 0..r {
5931                    if base + a >= diag.len() {
5932                        return Err(ArrowSchurGpuFailure::Unavailable);
5933                    }
5934                    diag[base + a] += gii * blk.w[[a, a]];
5935                }
5936            }
5937        }
5938        Ok(diag)
5939    }
5940
5941    fn frame_grid(work: usize, n_rows: usize) -> Result<LaunchConfig, ArrowSchurGpuFailure> {
5942        Ok(LaunchConfig {
5943            grid_dim: (
5944                ((work as u32).saturating_add(255) / 256).max(1),
5945                checked_i32(n_rows)? as u32,
5946                1,
5947            ),
5948            block_dim: (256, 1, 1),
5949            shared_mem_bytes: 0,
5950        })
5951    }
5952
5953    fn launch_sae_frame_matvec(
5954        stream: &Arc<CudaStream>,
5955        module: &Arc<CudaModule>,
5956        buffers: &mut DeviceSaeFrameBuffers,
5957        x: &CudaSlice<f64>,
5958        out: &mut CudaSlice<f64>,
5959        ridge_beta: f64,
5960    ) -> Result<(), ArrowSchurGpuFailure> {
5961        launch_sae_init(stream, module, out, x, ridge_beta, buffers.k)?;
5962        // Smooth penalty.
5963        if buffers.s_blocks > 0 {
5964            let kernel = module
5965                .load_function("arrow_sae_frame_smooth_matvec")
5966                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5967            let blocks_i32 = checked_i32(buffers.s_blocks)?;
5968            let cfg = frame_grid(buffers.k, buffers.s_blocks)?;
5969            let mut b = stream.launch_builder(&kernel);
5970            b.arg(x)
5971                .arg(&mut *out)
5972                .arg(&buffers.s_off)
5973                .arg(&buffers.s_m)
5974                .arg(&buffers.s_r)
5975                .arg(&buffers.s_ptr)
5976                .arg(&buffers.s_data)
5977                .arg(&blocks_i32);
5978            // SAFETY: smooth block metadata/data are live device buffers; the grid
5979            // covers (k channels × n_blocks) and the kernel bounds-checks m·r.
5980            unsafe { b.launch(cfg) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5981        }
5982        // Data penalty.
5983        if buffers.g_groups > 0 {
5984            let kernel = module
5985                .load_function("arrow_sae_frame_g_matvec")
5986                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5987            let groups_i32 = checked_i32(buffers.g_groups)?;
5988            let cfg = frame_grid(buffers.g_max_work.max(1), buffers.g_groups)?;
5989            let mut b = stream.launch_builder(&kernel);
5990            b.arg(x)
5991                .arg(&mut *out)
5992                .arg(&buffers.g_off_i)
5993                .arg(&buffers.g_off_j)
5994                .arg(&buffers.g_ri)
5995                .arg(&buffers.g_rj)
5996                .arg(&buffers.g_mj)
5997                .arg(&buffers.g_ptr)
5998                .arg(&buffers.g_data)
5999                .arg(&buffers.w_ptr)
6000                .arg(&buffers.w_data)
6001                .arg(&buffers.g_group_ptr)
6002                .arg(&buffers.g_group_blocks)
6003                .arg(&buffers.g_group_off)
6004                .arg(&buffers.g_group_width)
6005                .arg(&groups_i32);
6006            // SAFETY: each CSR ownership group covers one in-bounds scalar output
6007            // span, distinct spans are disjoint, and its block list indexes live
6008            // G/W metadata. Exactly one thread writes each scalar output.
6009            unsafe { b.launch(cfg) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6010        }
6011        // Reduced-Schur subtraction via dense per-row cross-blocks.
6012        let k_i32 = checked_i32(buffers.k)?;
6013        let max_q_i32 = checked_i32(buffers.max_q)?;
6014        let n_rows_i32 = checked_i32(buffers.n_rows)?;
6015        {
6016            let kernel = module
6017                .load_function("arrow_sae_frame_apply_h")
6018                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6019            let cfg = frame_grid(buffers.max_q, buffers.n_rows)?;
6020            let mut b = stream.launch_builder(&kernel);
6021            b.arg(x)
6022                .arg(&buffers.htb_ptr)
6023                .arg(&buffers.htb)
6024                .arg(&buffers.q_of)
6025                .arg(&mut buffers.hvec)
6026                .arg(&k_i32)
6027                .arg(&max_q_i32)
6028                .arg(&n_rows_i32);
6029            // SAFETY: dense cross-block + pointers + hvec are live buffers sized
6030            // for (n_rows × max_q) / (n_rows × k); kernel guards q_i and k.
6031            unsafe { b.launch(cfg) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6032        }
6033        {
6034            let kernel = module
6035                .load_function("arrow_sae_frame_apply_ainv")
6036                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6037            let cfg = frame_grid(buffers.max_q, buffers.n_rows)?;
6038            let mut b = stream.launch_builder(&kernel);
6039            b.arg(&buffers.ainv)
6040                .arg(&buffers.hvec)
6041                .arg(&buffers.q_of)
6042                .arg(&mut buffers.svec)
6043                .arg(&max_q_i32)
6044                .arg(&n_rows_i32);
6045            // SAFETY: ainv/hvec/svec are live buffers sized for n_rows·max_q²
6046            // and n_rows·max_q; the kernel guards row/coord bounds.
6047            unsafe { b.launch(cfg) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6048        }
6049        // #2512 — subtract the reduced-Schur term through the SAME two-stage
6050        // atomics-free scatter the evidence lane uses. The former single-kernel
6051        // form accumulated with `atomicAdd(double*, double)`, whose summation
6052        // order is chosen by the block scheduler; floating-point addition is not
6053        // associative, so the matvec — and with it the Newton step and the fitted
6054        // decoder — took a different value on every call. That made every fit
6055        // whose shape clears `reduced_schur_matvec_should_offload` irreproducible
6056        // on any host with a CUDA device, and reproducible on any host without
6057        // one. Stage 1 sums each contiguous row chunk in fixed index order and
6058        // stage 2 reduces the chunks in fixed order, so the reassociation is a
6059        // property of the shape rather than of the schedule.
6060        let k_blocks = ((buffers.k as u32).saturating_add(255) / 256).max(1);
6061        {
6062            let kernel = module
6063                .load_function("arrow_sae_frame_scatter_h_det_partial")
6064                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6065            let rows_per_chunk_i32 = checked_i32(buffers.rows_per_chunk)?;
6066            let cfg = LaunchConfig {
6067                grid_dim: (k_blocks, checked_i32(buffers.n_chunks)? as u32, 1),
6068                block_dim: (256, 1, 1),
6069                shared_mem_bytes: 0,
6070            };
6071            let mut b = stream.launch_builder(&kernel);
6072            b.arg(&buffers.svec)
6073                .arg(&buffers.htb_ptr)
6074                .arg(&buffers.htb)
6075                .arg(&buffers.q_of)
6076                .arg(&mut buffers.scatter_partial)
6077                .arg(&k_i32)
6078                .arg(&max_q_i32)
6079                .arg(&n_rows_i32)
6080                .arg(&rows_per_chunk_i32);
6081            // SAFETY: svec/cross-block are live buffers; scatter_partial is sized
6082            // n_chunks·k; each thread writes one in-bounds partial[chunk·k+a], a<k.
6083            unsafe { b.launch(cfg) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6084        }
6085        {
6086            let kernel = module
6087                .load_function("arrow_sae_frame_scatter_h_det_reduce_sub")
6088                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6089            let n_chunks_i32 = checked_i32(buffers.n_chunks)?;
6090            let cfg = LaunchConfig {
6091                grid_dim: (k_blocks, 1, 1),
6092                block_dim: (256, 1, 1),
6093                shared_mem_bytes: 0,
6094            };
6095            let mut b = stream.launch_builder(&kernel);
6096            b.arg(&buffers.scatter_partial)
6097                .arg(&mut *out)
6098                .arg(&k_i32)
6099                .arg(&n_chunks_i32);
6100            // SAFETY: scatter_partial sized n_chunks·k, out sized k; one in-bounds
6101            // out[a] updated per a<k, by exactly one thread.
6102            unsafe { b.launch(cfg) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6103        }
6104        Ok(())
6105    }
6106
6107    /// #1017 evidence lane: the DETERMINISTIC device reduced-Schur term
6108    /// `out[a] = -Σ_i H_βt^(i)(H_tt^(i)+ρ_t I)⁻¹H_tβ^(i)x`. Reuses the per-row
6109    /// `apply_h`→`apply_ainv` chain (each output written by a single thread — no
6110    /// cross-thread race) and the atomics-free `arrow_sae_frame_scatter_h_det`
6111    /// (one thread per output coord, fixed row order), so the value is run-to-run
6112    /// bit-stable. This is the reduced-Schur half of `S·x` ONLY; the penalty side
6113    /// `(P_ββ + ρ_β I)x` is added by the caller on the host via the already
6114    /// deterministic `sae_framed_penalty_matvec_cpu`. `out` is fully assigned.
6115    /// `ainv` must already be primed for the target `ρ_t` in `buffers`.
6116    fn launch_sae_frame_reduced_schur_det(
6117        stream: &Arc<CudaStream>,
6118        module: &Arc<CudaModule>,
6119        buffers: &mut DeviceSaeFrameBuffers,
6120        x: &CudaSlice<f64>,
6121        out: &mut CudaSlice<f64>,
6122    ) -> Result<(), ArrowSchurGpuFailure> {
6123        let k_i32 = checked_i32(buffers.k)?;
6124        let max_q_i32 = checked_i32(buffers.max_q)?;
6125        let n_rows_i32 = checked_i32(buffers.n_rows)?;
6126        // hvec[i][c] = Σ_a H_tβ[i][c,a]·x[a] — warp-cooperative (one warp per
6127        // (row, c), coalesced reads). Block = max_q·32 threads (one warp per
6128        // possible c), grid.x = n_rows.
6129        {
6130            let kernel = module
6131                .load_function("arrow_sae_frame_apply_h_warp")
6132                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6133            let block = checked_i32(buffers.max_q)?
6134                .checked_mul(32)
6135                .ok_or(ArrowSchurGpuFailure::Unavailable)? as u32;
6136            let cfg = LaunchConfig {
6137                grid_dim: (checked_i32(buffers.n_rows)? as u32, 1, 1),
6138                block_dim: (block, 1, 1),
6139                shared_mem_bytes: 0,
6140            };
6141            let mut b = stream.launch_builder(&kernel);
6142            b.arg(x)
6143                .arg(&buffers.htb_ptr)
6144                .arg(&buffers.htb)
6145                .arg(&buffers.q_of)
6146                .arg(&mut buffers.hvec)
6147                .arg(&k_i32)
6148                .arg(&max_q_i32)
6149                .arg(&n_rows_i32);
6150            // SAFETY: dense cross-block + pointers + hvec are live buffers sized
6151            // for (n_rows × max_q) / (n_rows × k); block = max_q·32 ≤ 1024, each
6152            // warp guards `warp < q_i` and strides `a < k`.
6153            unsafe { b.launch(cfg) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6154        }
6155        // svec[i] = ainv_i · hvec_i.
6156        {
6157            let kernel = module
6158                .load_function("arrow_sae_frame_apply_ainv")
6159                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6160            let cfg = frame_grid(buffers.max_q, buffers.n_rows)?;
6161            let mut b = stream.launch_builder(&kernel);
6162            b.arg(&buffers.ainv)
6163                .arg(&buffers.hvec)
6164                .arg(&buffers.q_of)
6165                .arg(&mut buffers.svec)
6166                .arg(&max_q_i32)
6167                .arg(&n_rows_i32);
6168            // SAFETY: ainv/hvec/svec are live buffers sized for n_rows·max_q²
6169            // and n_rows·max_q; the kernel guards row/coord bounds.
6170            unsafe { b.launch(cfg) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6171        }
6172        // out[a] = -Σ_i Σ_c H_tβ[i][c,a]·svec[i,c] — 2-stage DETERMINISTIC scatter.
6173        // Stage 1: partial[chunk][a] over contiguous row chunks, launching
6174        // ⌈k/256⌉·n_chunks CTAs (vs the single-strip ⌈k/256⌉ = 4 at k=911 that left
6175        // ~94% of the SMs idle). Rows summed in fixed order within each chunk.
6176        let k_blocks = ((buffers.k as u32).saturating_add(255) / 256).max(1);
6177        {
6178            let kernel = module
6179                .load_function("arrow_sae_frame_scatter_h_det_partial")
6180                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6181            let rows_per_chunk_i32 = checked_i32(buffers.rows_per_chunk)?;
6182            let cfg = LaunchConfig {
6183                grid_dim: (k_blocks, checked_i32(buffers.n_chunks)? as u32, 1),
6184                block_dim: (256, 1, 1),
6185                shared_mem_bytes: 0,
6186            };
6187            let mut b = stream.launch_builder(&kernel);
6188            b.arg(&buffers.svec)
6189                .arg(&buffers.htb_ptr)
6190                .arg(&buffers.htb)
6191                .arg(&buffers.q_of)
6192                .arg(&mut buffers.scatter_partial)
6193                .arg(&k_i32)
6194                .arg(&max_q_i32)
6195                .arg(&n_rows_i32)
6196                .arg(&rows_per_chunk_i32);
6197            // SAFETY: svec/cross-block are live buffers; scatter_partial is sized
6198            // n_chunks·k; each thread writes one in-bounds partial[chunk·k+a], a<k.
6199            unsafe { b.launch(cfg) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6200        }
6201        // Stage 2: out[a] = -Σ_chunk partial[chunk][a], chunks reduced in fixed
6202        // order — the fixed reassociation that keeps the scatter deterministic.
6203        {
6204            let kernel = module
6205                .load_function("arrow_sae_frame_scatter_h_det_reduce")
6206                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6207            let n_chunks_i32 = checked_i32(buffers.n_chunks)?;
6208            let cfg = LaunchConfig {
6209                grid_dim: (k_blocks, 1, 1),
6210                block_dim: (256, 1, 1),
6211                shared_mem_bytes: 0,
6212            };
6213            let mut b = stream.launch_builder(&kernel);
6214            b.arg(&buffers.scatter_partial)
6215                .arg(&mut *out)
6216                .arg(&k_i32)
6217                .arg(&n_chunks_i32);
6218            // SAFETY: scatter_partial sized n_chunks·k, out sized k; one in-bounds
6219            // out[a] written per a<k.
6220            unsafe { b.launch(cfg) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6221        }
6222        Ok(())
6223    }
6224
6225    fn launch_sae_frame_diag_sub(
6226        stream: &Arc<CudaStream>,
6227        module: &Arc<CudaModule>,
6228        buffers: &DeviceSaeFrameBuffers,
6229        diag: &mut CudaSlice<f64>,
6230    ) -> Result<(), ArrowSchurGpuFailure> {
6231        let kernel = module
6232            .load_function("arrow_sae_frame_diag_sub")
6233            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6234        let k_i32 = checked_i32(buffers.k)?;
6235        let max_q_i32 = checked_i32(buffers.max_q)?;
6236        let n_rows_i32 = checked_i32(buffers.n_rows)?;
6237        let cfg = frame_grid(buffers.k, 1)?;
6238        let mut b = stream.launch_builder(&kernel);
6239        b.arg(diag)
6240            .arg(&buffers.ainv)
6241            .arg(&buffers.htb_ptr)
6242            .arg(&buffers.htb)
6243            .arg(&buffers.q_of)
6244            .arg(&k_i32)
6245            .arg(&max_q_i32)
6246            .arg(&n_rows_i32);
6247        // SAFETY: diag + cross-block + ainv are live buffers; one thread owns
6248        // each a<k and walks every row, with c/d bounded by q<=max_q.
6249        unsafe { b.launch(cfg) }
6250            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6251        Ok(())
6252    }
6253
6254    pub(super) fn solve_sae_matrix_free_pcg_framed(
6255        sys: &ArrowSchurSystem,
6256        data: &DeviceSaePcgData,
6257        ridge_t: f64,
6258        ridge_beta: f64,
6259        rhs_beta: &Array1<f64>,
6260        max_iterations: usize,
6261        relative_tolerance: f64,
6262    ) -> Result<(Array1<f64>, ArrowPcgDiagnostics), ArrowSchurGpuFailure> {
6263        let k = rhs_beta.len();
6264        if k == 0 || data.beta_dim != k || sys.k != k {
6265            return Err(ArrowSchurGpuFailure::Unavailable);
6266        }
6267        let frame = data
6268            .frame
6269            .as_ref()
6270            .ok_or(ArrowSchurGpuFailure::Unavailable)?;
6271        let runtime = super::resolve_runtime_for_device_path()?
6272            .filter(|rt| {
6273                rt.policy().reduced_schur_matvec_should_offload(
6274                    sys.rows.len(),
6275                    sys.k,
6276                    sys.d,
6277                    max_iterations,
6278                )
6279            })
6280            .ok_or(ArrowSchurGpuFailure::Unavailable)?;
6281        let ctx = gam_gpu::device_runtime::cuda_context_for(runtime.selected_device().ordinal)
6282            .ok_or(ArrowSchurGpuFailure::Unavailable)?;
6283        let stream = ctx
6284            .new_stream()
6285            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6286        let blas = CudaBlas::new(stream.clone()).map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6287        let vector_module = pcg_vector_module(&ctx)?;
6288        // #1017/#2230 residency measurement: one line per matrix-free PCG solve
6289        // (hence per LM ridge-ladder trial) — operand bytes by category + the
6290        // ridge pair, so the a100 job (RUST_LOG=info) confirms the sub-lane and
6291        // sizes the per-trial re-upload a base-resident frame would remove.
6292        log::info!(
6293            "#1017/#2230 {} ridge_t={ridge_t:e} ridge_beta={ridge_beta:e}",
6294            data.operand_byte_report()
6295        );
6296        let mut buffers = flatten_device_sae_frame_data(sys, data, frame, ridge_t, &stream)?;
6297        let dctx = FramedPcgCtx {
6298            stream: &stream,
6299            blas: &blas,
6300            module: vector_module,
6301            max_iterations,
6302            relative_tolerance,
6303        };
6304        pcg_solve_framed_body(data, frame, &mut buffers, ridge_beta, rhs_beta, &dctx)
6305    }
6306
6307    /// Device handles + CG controls for the framed PCG loop, bundled so the
6308    /// shared body stays under the argument-count lint without an `#[allow]`.
6309    struct FramedPcgCtx<'a> {
6310        stream: &'a Arc<CudaStream>,
6311        blas: &'a CudaBlas,
6312        module: &'a Arc<CudaModule>,
6313        max_iterations: usize,
6314        relative_tolerance: f64,
6315    }
6316
6317    /// The framed reduced-Schur Jacobi-PCG loop over already-built device
6318    /// buffers. Factored out of [`solve_sae_matrix_free_pcg_framed`] so the
6319    /// resident-frame path ([`ResidentSaeFrameHandle::resolve`]) — which reuses
6320    /// the ridge-independent buffers and re-uploads only `ainv` — runs the
6321    /// byte-for-byte identical solve. Everything after the buffer build is
6322    /// unchanged from the inline body: penalty-diagonal Jacobi preconditioner,
6323    /// the CG recurrence, and the `Δβ` readback.
6324    fn pcg_solve_framed_body(
6325        data: &DeviceSaePcgData,
6326        frame: &DeviceSaeFrameData,
6327        buffers: &mut DeviceSaeFrameBuffers,
6328        ridge_beta: f64,
6329        rhs_beta: &Array1<f64>,
6330        dctx: &FramedPcgCtx<'_>,
6331    ) -> Result<(Array1<f64>, ArrowPcgDiagnostics), ArrowSchurGpuFailure> {
6332        let stream = dctx.stream;
6333        let blas = dctx.blas;
6334        let vector_module = dctx.module;
6335        let max_iterations = dctx.max_iterations;
6336        let relative_tolerance = dctx.relative_tolerance;
6337        let k = rhs_beta.len();
6338        let rhs_norm = rhs_beta.iter().map(|v| v * v).sum::<f64>().sqrt();
6339        if rhs_norm == 0.0 {
6340            return Ok((Array1::<f64>::zeros(k), ArrowPcgDiagnostics::default()));
6341        }
6342        let tol = (relative_tolerance.max(0.0) * rhs_norm).max(1e-12);
6343        let rhs_dev = stream
6344            .clone_htod(
6345                rhs_beta
6346                    .as_slice()
6347                    .ok_or(ArrowSchurGpuFailure::Unavailable)?,
6348            )
6349            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6350        let diag_host = sae_frame_penalty_diag_host(data, frame, ridge_beta)?;
6351        let mut diag_dev = stream
6352            .clone_htod(&diag_host)
6353            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6354        launch_sae_frame_diag_sub(stream, vector_module, buffers, &mut diag_dev)?;
6355        let diag_host = stream
6356            .clone_dtoh(&diag_dev)
6357            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6358        let mut inv_diag = Vec::with_capacity(k);
6359        for (idx, &d) in diag_host.iter().enumerate() {
6360            if !d.is_finite() || d <= 1.0e-18 {
6361                return Err(ArrowSchurGpuFailure::SchurFactorFailed {
6362                    reason: format!(
6363                        "framed SAE GPU PCG: non-positive Jacobi diagonal at {idx}: {d:e}"
6364                    ),
6365                });
6366            }
6367            inv_diag.push(1.0 / d);
6368        }
6369        let inv_diag_dev = stream
6370            .clone_htod(&inv_diag)
6371            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6372
6373        let mut x_dev = stream
6374            .alloc_zeros::<f64>(k)
6375            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6376        let mut r_dev = stream
6377            .alloc_zeros::<f64>(k)
6378            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6379        device_copy(blas, stream, k, &rhs_dev, &mut r_dev)?;
6380        let mut z_dev = stream
6381            .alloc_zeros::<f64>(k)
6382            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6383        launch_jacobi_mul(stream, vector_module, &inv_diag_dev, &r_dev, &mut z_dev, k)?;
6384        let mut p_dev = stream
6385            .alloc_zeros::<f64>(k)
6386            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6387        device_copy(blas, stream, k, &z_dev, &mut p_dev)?;
6388        let mut ap_dev = stream
6389            .alloc_zeros::<f64>(k)
6390            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6391
6392        let mut rz = device_dot(blas, stream, k, &r_dev, &z_dev)?;
6393        if rz <= 0.0 || !rz.is_finite() {
6394            return Err(ArrowSchurGpuFailure::SchurFactorFailed {
6395                reason: format!("framed SAE GPU PCG: non-positive initial rᵀM⁻¹r={rz:e}"),
6396            });
6397        }
6398        let mut diag = ArrowPcgDiagnostics {
6399            precond_apply_calls: 1,
6400            stopping_reason: PcgStopReason::MaxIter,
6401            ..ArrowPcgDiagnostics::default()
6402        };
6403        for _ in 0..max_iterations.max(1) {
6404            launch_sae_frame_matvec(
6405                stream,
6406                vector_module,
6407                buffers,
6408                &p_dev,
6409                &mut ap_dev,
6410                ridge_beta,
6411            )?;
6412            diag.matvec_calls += 1;
6413            diag.iterations += 1;
6414            let pap = device_dot(blas, stream, k, &p_dev, &ap_dev)?;
6415            if pap <= 0.0 || !pap.is_finite() {
6416                return Err(ArrowSchurGpuFailure::SchurFactorFailed {
6417                    reason: format!("framed SAE GPU PCG: non-positive curvature pᵀAp={pap:e}"),
6418                });
6419            }
6420            let alpha = rz / pap;
6421            device_axpy(blas, stream, k, alpha, &p_dev, &mut x_dev)?;
6422            device_axpy(blas, stream, k, -alpha, &ap_dev, &mut r_dev)?;
6423            let r_norm = device_nrm2(blas, stream, k, &r_dev)?;
6424            if r_norm <= tol {
6425                diag.final_relative_residual = r_norm / rhs_norm;
6426                diag.stopping_reason = PcgStopReason::Converged;
6427                break;
6428            }
6429            launch_jacobi_mul(stream, vector_module, &inv_diag_dev, &r_dev, &mut z_dev, k)?;
6430            diag.precond_apply_calls += 1;
6431            let rz_new = device_dot(blas, stream, k, &r_dev, &z_dev)?;
6432            if rz_new <= 0.0 || !rz_new.is_finite() {
6433                return Err(ArrowSchurGpuFailure::SchurFactorFailed {
6434                    reason: format!("framed SAE GPU PCG: non-positive rᵀM⁻¹r={rz_new:e}"),
6435                });
6436            }
6437            let beta = rz_new / rz;
6438            launch_update_p(stream, vector_module, &z_dev, beta, &mut p_dev, k)?;
6439            rz = rz_new;
6440        }
6441        if diag.stopping_reason != PcgStopReason::Converged {
6442            let r_norm = device_nrm2(blas, stream, k, &r_dev)?;
6443            diag.final_relative_residual = r_norm / rhs_norm;
6444            diag.stopping_reason = PcgStopReason::MaxIter;
6445        }
6446        let x = stream
6447            .clone_dtoh(&x_dev)
6448            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6449        Ok((Array1::from_vec(x), diag))
6450    }
6451
6452    /// #1017 device-resident framed SAE frame across the LM ridge ladder. Holds
6453    /// the ridge-INDEPENDENT device operand buffers (uploaded once) plus the CUDA
6454    /// context/stream they live on; each [`resolve`](Self::resolve) recomputes
6455    /// ONLY the ridge-dependent per-row `ainv` and re-runs the identical framed
6456    /// PCG. Persists only `Send + Sync` cudarc handles (`Arc<CudaContext>`,
6457    /// `Arc<CudaStream>`, `CudaSlice`); the cheap `CudaBlas`/module are re-derived
6458    /// per solve, so the whole handle is safely shareable through
6459    /// [`crate::arrow_schur::ArrowSolveOptions`].
6460    pub(crate) struct ResidentSaeFrameHandle {
6461        ctx: Arc<CudaContext>,
6462        stream: Arc<CudaStream>,
6463        buffers: std::sync::Mutex<DeviceSaeFrameBuffers>,
6464        q_of: Vec<i32>,
6465        max_q: usize,
6466        n_rows: usize,
6467        k: usize,
6468    }
6469
6470    impl ResidentSaeFrameHandle {
6471        /// Build the resident frame: gate exactly as
6472        /// [`solve_sae_matrix_free_pcg_framed`] (framed data present, offload
6473        /// predicate over the CG budget, live runtime), upload the
6474        /// ridge-independent operands once with a zero `ainv` placeholder, and
6475        /// stash the metadata needed to recompute `ainv` per trial. Ordinary
6476        /// shape/policy absence is `Ok(None)`; runtime and upload faults retain
6477        /// their typed identity.
6478        pub(crate) fn build(
6479            sys: &ArrowSchurSystem,
6480            cg_iters: usize,
6481        ) -> Result<Option<Self>, ArrowSchurGpuFailure> {
6482            let Some(data) = sys.device_sae_pcg.as_ref() else {
6483                return Ok(None);
6484            };
6485            let Some(frame) = data.frame.as_ref() else {
6486                return Ok(None);
6487            };
6488            if sys.k == 0 || data.beta_dim != sys.k {
6489                return Ok(None);
6490            }
6491            let Some(runtime) = super::resolve_runtime_for_device_path()? else {
6492                return Ok(None);
6493            };
6494            if !runtime.policy().reduced_schur_matvec_should_offload(
6495                sys.rows.len(),
6496                sys.k,
6497                sys.d,
6498                cg_iters,
6499            ) {
6500                return Ok(None);
6501            }
6502            let ctx = gam_gpu::device_runtime::cuda_context_for(runtime.selected_device().ordinal)
6503                .ok_or_else(|| ArrowSchurGpuFailure::SchurFactorFailed {
6504                    reason: "resident SAE frame could not bind the admitted CUDA context"
6505                        .to_string(),
6506                })?;
6507            let stream = ctx.new_stream().map_err(|error| {
6508                ArrowSchurGpuFailure::SchurFactorFailed {
6509                    reason: format!("resident SAE frame stream creation failed: {error}"),
6510                }
6511            })?;
6512            let host = super::flatten_frame_host_operands(sys, data, frame)?;
6513            // #1017/#2230 residency measurement: the ridge-independent operand
6514            // bytes this frame uploads ONCE for the whole LM ridge ladder. The
6515            // per-trial flatten path re-uploaded this same total on EVERY trial
6516            // (its `#1017/#2230 …` info line fires once per trial); against a
6517            // ladder of `T` trials the resident frame removes `(T − 1) ×` this,
6518            // re-uploading only the per-row `ainv` (n_rows·max_q² f64) per trial.
6519            log::info!(
6520                "#1017 SAE resident frame ENGAGED: {} uploaded ONCE for the ladder; \
6521                 per-trial re-upload now only ainv ({}rows × {}²·8B)",
6522                data.operand_byte_report(),
6523                host.n_rows,
6524                host.max_q
6525            );
6526            let zero_ainv = vec![0.0_f64; host.n_rows * host.max_q * host.max_q];
6527            let buffers = upload_frame_buffers(&host, &zero_ainv, &stream)?;
6528            Ok(Some(Self {
6529                ctx,
6530                stream,
6531                buffers: std::sync::Mutex::new(buffers),
6532                q_of: host.q_of,
6533                max_q: host.max_q,
6534                n_rows: host.n_rows,
6535                k: host.k,
6536            }))
6537        }
6538
6539        /// Overwrite the resident per-row `ainv` for a fixed `ridge_t` (the single
6540        /// evidence ridge), WITHOUT running the PCG — so an evidence matvec closure
6541        /// primes the factors once and every apply reuses them. Mirrors the
6542        /// ridge-dependent refresh in [`SaeResidentFrame::resolve`]. A genuinely
6543        /// non-PD row surfaces `RidgeBumpRequired`/`SchurFactorFailed` (the caller
6544        /// then declines to the CPU matvec, which handles the escalation).
6545        pub(super) fn prime_ainv(
6546            &self,
6547            sys: &ArrowSchurSystem,
6548            ridge_t: f64,
6549        ) -> Result<(), ArrowSchurGpuFailure> {
6550            if sys.k != self.k || sys.rows.len() != self.n_rows {
6551                return Err(ArrowSchurGpuFailure::Unavailable);
6552            }
6553            let ainv =
6554                super::compute_ainv_host(sys, &self.q_of, self.max_q, self.n_rows, ridge_t)?;
6555            let mut buffers = self
6556                .buffers
6557                .lock()
6558                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6559            if ainv.len() != buffers.ainv.len() {
6560                return Err(ArrowSchurGpuFailure::Unavailable);
6561            }
6562            self.stream
6563                .memcpy_htod(&ainv, &mut buffers.ainv)
6564                .map_err(|_| ArrowSchurGpuFailure::Unavailable)
6565        }
6566    }
6567
6568    impl super::SaeResidentFrame for ResidentSaeFrameHandle {
6569        fn refresh(&self, sys: &ArrowSchurSystem) -> Result<(), ArrowSchurGpuFailure> {
6570            if sys.k != self.k || sys.rows.len() != self.n_rows {
6571                return Err(ArrowSchurGpuFailure::Unavailable);
6572            }
6573            let data = sys
6574                .device_sae_pcg
6575                .as_ref()
6576                .ok_or(ArrowSchurGpuFailure::Unavailable)?;
6577            let frame = data
6578                .frame
6579                .as_ref()
6580                .ok_or(ArrowSchurGpuFailure::Unavailable)?;
6581            let host = super::flatten_frame_host_operands(sys, data, frame)?;
6582            if host.q_of != self.q_of || host.max_q != self.max_q {
6583                return Err(ArrowSchurGpuFailure::Unavailable);
6584            }
6585            let mut buffers = self
6586                .buffers
6587                .lock()
6588                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6589            refresh_frame_buffers(&host, &mut buffers, &self.stream)
6590        }
6591
6592        fn resolve(
6593            &self,
6594            sys: &ArrowSchurSystem,
6595            ridge_t: f64,
6596            ridge_beta: f64,
6597            rhs_beta: &Array1<f64>,
6598            max_iterations: usize,
6599            relative_tolerance: f64,
6600        ) -> Result<(Array1<f64>, ArrowPcgDiagnostics), ArrowSchurGpuFailure> {
6601            // The ladder holds `sys` fixed across trials; if any shape drifted
6602            // from build time the resident buffers no longer match — decline so
6603            // the caller retries via the per-trial flatten.
6604            if sys.k != self.k || sys.rows.len() != self.n_rows || rhs_beta.len() != self.k {
6605                return Err(ArrowSchurGpuFailure::Unavailable);
6606            }
6607            let data = sys
6608                .device_sae_pcg
6609                .as_ref()
6610                .ok_or(ArrowSchurGpuFailure::Unavailable)?;
6611            let frame = data
6612                .frame
6613                .as_ref()
6614                .ok_or(ArrowSchurGpuFailure::Unavailable)?;
6615            // Ridge-dependent buffer only: recompute per-row ainv and overwrite
6616            // the resident `ainv` slice in place (a genuine non-PD block still
6617            // surfaces `RidgeBumpRequired` for the LM escalation, unchanged).
6618            let ainv = super::compute_ainv_host(sys, &self.q_of, self.max_q, self.n_rows, ridge_t)?;
6619            let mut buffers = self
6620                .buffers
6621                .lock()
6622                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6623            if ainv.len() != buffers.ainv.len() {
6624                return Err(ArrowSchurGpuFailure::Unavailable);
6625            }
6626            self.stream
6627                .memcpy_htod(&ainv, &mut buffers.ainv)
6628                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6629            let blas = CudaBlas::new(self.stream.clone())
6630                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6631            let vector_module = pcg_vector_module(&self.ctx)?;
6632            let dctx = FramedPcgCtx {
6633                stream: &self.stream,
6634                blas: &blas,
6635                module: vector_module,
6636                max_iterations,
6637                relative_tolerance,
6638            };
6639            pcg_solve_framed_body(data, frame, &mut *buffers, ridge_beta, rhs_beta, &dctx)
6640        }
6641    }
6642
6643    /// #1017 evidence lane: build a device-resident, RUN-TO-RUN DETERMINISTIC
6644    /// framed reduced-Schur `S·v` as the SLQ/surrogate `GpuSchurMatvec`. Uploads
6645    /// the ridge-independent framed operands ONCE ([`ResidentSaeFrameHandle`]),
6646    /// primes `ainv` at the single evidence `ridge_t`, and returns a closure that
6647    /// per apply crosses only `x` (down) and `out` (up): the deterministic host
6648    /// penalty ([`super::sae_framed_penalty_matvec_cpu`]) plus the atomics-free
6649    /// device reduced-Schur term ([`launch_sae_frame_reduced_schur_det`]). `None`
6650    /// on any decline (no device / shape / offload floor / non-PD at this ridge),
6651    /// so the caller keeps the CPU row-procedural matvec. A per-apply device fault
6652    /// after a validated build is a genuine bug/OOM and panics LOUD (the #1551
6653    /// no-silent-CPU discipline) rather than silently degrading the evidence.
6654    pub(super) fn build_framed_resident_evidence_matvec(
6655        sys: &ArrowSchurSystem,
6656        ridge_t: f64,
6657        ridge_beta: f64,
6658        apply_budget: usize,
6659    ) -> Result<Option<crate::arrow_schur::GpuSchurMatvec>, ArrowSchurGpuFailure> {
6660        let Some(handle) = ResidentSaeFrameHandle::build(sys, apply_budget)? else {
6661            return Ok(None);
6662        };
6663        handle.prime_ainv(sys, ridge_t)?;
6664        let Some(data) = sys.device_sae_pcg.as_ref().cloned() else {
6665            return Ok(None);
6666        };
6667        let handle = Arc::new(handle);
6668        let k = handle.k;
6669        let closure: crate::arrow_schur::GpuSchurMatvec =
6670            Arc::new(move |x: &Array1<f64>, out: &mut Array1<f64>| {
6671                assert_eq!(x.len(), k, "#1017 framed evidence matvec: x.len() != k");
6672                assert_eq!(out.len(), k, "#1017 framed evidence matvec: out.len() != k");
6673                let module = pcg_vector_module(&handle.ctx)
6674                    .expect("#1017 framed evidence matvec: pcg_vector_module unavailable");
6675                let x_slice = x
6676                    .as_slice()
6677                    .expect("#1017 framed evidence matvec: x not contiguous");
6678                let x_dev = handle
6679                    .stream
6680                    .clone_htod(x_slice)
6681                    .expect("#1017 framed evidence matvec: htod(x) failed");
6682                let mut reduced_dev = handle
6683                    .stream
6684                    .alloc_zeros::<f64>(k)
6685                    .expect("#1017 framed evidence matvec: device alloc failed");
6686                {
6687                    let mut buffers = handle
6688                        .buffers
6689                        .lock()
6690                        .expect("#1017 framed evidence matvec: resident frame poisoned");
6691                    launch_sae_frame_reduced_schur_det(
6692                        &handle.stream,
6693                        module,
6694                        &mut buffers,
6695                        &x_dev,
6696                        &mut reduced_dev,
6697                    )
6698                    .expect("#1017 framed evidence matvec: device reduced-Schur launch failed");
6699                }
6700                let reduced = handle
6701                    .stream
6702                    .clone_dtoh(&reduced_dev)
6703                    .expect("#1017 framed evidence matvec: dtoh(out) failed");
6704                let out_slice = out
6705                    .as_slice_mut()
6706                    .expect("#1017 framed evidence matvec: out not contiguous");
6707                super::sae_framed_penalty_matvec_cpu(&data, ridge_beta, x_slice, out_slice);
6708                for a in 0..k {
6709                    out_slice[a] += reduced[a];
6710                }
6711            });
6712        Ok(Some(closure))
6713    }
6714
6715    /// #1551 stage-isolating triage seam: run the framed reduced-Schur matvec
6716    /// `out = S·x` ONCE on the device (no PCG, no offload-floor gate) and return
6717    /// `out`, so a tiny hand-verifiable fixture can diff it against the CPU oracle
6718    /// `sae_framed_schur_matvec_cpu` element-by-element to localize the structural
6719    /// divergence to a single kernel stage. Returns `Unavailable` only when CUDA
6720    /// is genuinely absent (so the test skips cleanly off-device).
6721    pub(super) fn solve_sae_matrix_free_pcg(
6722        sys: &ArrowSchurSystem,
6723        data: &DeviceSaePcgData,
6724        ridge_t: f64,
6725        ridge_beta: f64,
6726        rhs_beta: &Array1<f64>,
6727        max_iterations: usize,
6728        relative_tolerance: f64,
6729    ) -> Result<(Array1<f64>, ArrowPcgDiagnostics), ArrowSchurGpuFailure> {
6730        let k = rhs_beta.len();
6731        if k == 0 || data.beta_dim != k || sys.k != k {
6732            return Err(ArrowSchurGpuFailure::Unavailable);
6733        }
6734        // #2502 GUARD: a residency-only payload (a_phi/local_jac staged for the
6735        // CPU resident reduced-Schur, no H_bb blocks) carries NO data-fit Gram
6736        // and NO smooth penalty; the legacy kernel would solve against a
6737        // ridge-only H_bb. Decline so the host falls back to the CPU operator,
6738        // which composes H_bb through the system closures.
6739        if data.sparse_g_blocks.is_empty() && data.smooth_blocks.is_empty() && data.frame.is_none() {
6740            return Err(ArrowSchurGpuFailure::Unavailable);
6741        }
6742        // #1017/#1026 GUARD: the legacy `⊗ I_p` kernel must NEVER receive framed
6743        // data (factored `G ⊗ W_{ij}` + dense per-row cross blocks); decline so a
6744        // mis-route falls back to the CPU rather than returning a wrong step.
6745        if data.frame.is_some() {
6746            return Err(ArrowSchurGpuFailure::Unavailable);
6747        }
6748        // #1017 Phase-1 dispatch re-key: this is the matrix-free SAE reduced-Schur
6749        // PCG — the production hot path, not a single dense factorization. The
6750        // dense-Direct floor `dense_hessian_work_target_is_gpu(n, k)` keys on
6751        // `2·n·k²` and is the WRONG gate here: it ignores the per-row frame depth
6752        // `d` (the M dimension that multiplies the per-apply work) and the
6753        // `1/cg_iters` staging amortisation, so it both undercounts the SAE batched
6754        // work `n·k·d` and applies a cold single-launch breakeven to an apply that
6755        // reuses device-resident frames `max_iterations` times. Key instead on the
6756        // CG-amortised total batched work — the same predicate the host injection
6757        // gate (`maybe_inject_gpu_schur_matvec`) consults — so few-row/wide-`k`/
6758        // modest-`d` LLM shapes register the real `n × k × d × cg_iters` arithmetic.
6759        // Kernels and numerics are untouched; only where the matvec runs changes,
6760        // and the host falls back to the bit-identical CPU matvec when this declines.
6761        let runtime = super::resolve_runtime_for_device_path()?
6762            .filter(|rt| {
6763                rt.policy().reduced_schur_matvec_should_offload(
6764                    sys.rows.len(),
6765                    sys.k,
6766                    sys.d,
6767                    max_iterations,
6768                )
6769            })
6770            .ok_or(ArrowSchurGpuFailure::Unavailable)?;
6771        let ctx = gam_gpu::device_runtime::cuda_context_for(runtime.selected_device().ordinal)
6772            .ok_or(ArrowSchurGpuFailure::Unavailable)?;
6773        let stream = ctx
6774            .new_stream()
6775            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6776        let blas = CudaBlas::new(stream.clone()).map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6777        let vector_module = pcg_vector_module(&ctx)?;
6778        // #1017/#2230 residency measurement (legacy sparse ⊗I_p lane): see the
6779        // framed twin — one line per solve/ladder-trial for the a100 job to size
6780        // the per-trial operand re-upload.
6781        log::info!(
6782            "#1017/#2230 {} ridge_t={ridge_t:e} ridge_beta={ridge_beta:e}",
6783            data.operand_byte_report()
6784        );
6785        let mut buffers = flatten_device_sae_data(sys, data, ridge_t, &stream)?;
6786
6787        let rhs_norm = rhs_beta.iter().map(|v| v * v).sum::<f64>().sqrt();
6788        if rhs_norm == 0.0 {
6789            return Ok((Array1::<f64>::zeros(k), ArrowPcgDiagnostics::default()));
6790        }
6791        let tol = (relative_tolerance.max(0.0) * rhs_norm).max(1e-12);
6792        let rhs_dev = stream
6793            .clone_htod(
6794                rhs_beta
6795                    .as_slice()
6796                    .ok_or(ArrowSchurGpuFailure::Unavailable)?,
6797            )
6798            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6799        let diag_host = sae_penalty_diag_host(data, ridge_beta)?;
6800        let mut diag_dev = stream
6801            .clone_htod(&diag_host)
6802            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6803        launch_sae_diag_sub(&stream, vector_module, &mut buffers, &mut diag_dev)?;
6804        let diag_host = stream
6805            .clone_dtoh(&diag_dev)
6806            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6807        let mut inv_diag = Vec::with_capacity(k);
6808        for (idx, &d) in diag_host.iter().enumerate() {
6809            if !d.is_finite() || d <= 1.0e-18 {
6810                return Err(ArrowSchurGpuFailure::SchurFactorFailed {
6811                    reason: format!(
6812                        "SAE matrix-free GPU PCG: non-positive Schur Jacobi diagonal at {idx}: {d:e}"
6813                    ),
6814                });
6815            }
6816            inv_diag.push(1.0 / d);
6817        }
6818        let inv_diag_dev = stream
6819            .clone_htod(&inv_diag)
6820            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6821
6822        let mut x_dev = stream
6823            .alloc_zeros::<f64>(k)
6824            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6825        let mut r_dev = stream
6826            .alloc_zeros::<f64>(k)
6827            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6828        device_copy(&blas, &stream, k, &rhs_dev, &mut r_dev)?;
6829        let mut z_dev = stream
6830            .alloc_zeros::<f64>(k)
6831            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6832        launch_jacobi_mul(&stream, vector_module, &inv_diag_dev, &r_dev, &mut z_dev, k)?;
6833        let mut p_dev = stream
6834            .alloc_zeros::<f64>(k)
6835            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6836        device_copy(&blas, &stream, k, &z_dev, &mut p_dev)?;
6837        let mut ap_dev = stream
6838            .alloc_zeros::<f64>(k)
6839            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6840
6841        let mut rz = device_dot(&blas, &stream, k, &r_dev, &z_dev)?;
6842        if rz <= 0.0 || !rz.is_finite() {
6843            return Err(ArrowSchurGpuFailure::SchurFactorFailed {
6844                reason: format!("SAE matrix-free GPU PCG: non-positive initial rᵀM⁻¹r={rz:e}"),
6845            });
6846        }
6847        let mut diag = ArrowPcgDiagnostics {
6848            precond_apply_calls: 1,
6849            stopping_reason: PcgStopReason::MaxIter,
6850            ..ArrowPcgDiagnostics::default()
6851        };
6852
6853        for _ in 0..max_iterations.max(1) {
6854            launch_sae_matvec(
6855                &stream,
6856                vector_module,
6857                &mut buffers,
6858                &p_dev,
6859                &mut ap_dev,
6860                ridge_beta,
6861            )?;
6862            diag.matvec_calls += 1;
6863            diag.iterations += 1;
6864            let pap = device_dot(&blas, &stream, k, &p_dev, &ap_dev)?;
6865            if pap <= 0.0 || !pap.is_finite() {
6866                return Err(ArrowSchurGpuFailure::SchurFactorFailed {
6867                    reason: format!("SAE matrix-free GPU PCG: non-positive curvature pᵀAp={pap:e}"),
6868                });
6869            }
6870            let alpha = rz / pap;
6871            device_axpy(&blas, &stream, k, alpha, &p_dev, &mut x_dev)?;
6872            device_axpy(&blas, &stream, k, -alpha, &ap_dev, &mut r_dev)?;
6873            let r_norm = device_nrm2(&blas, &stream, k, &r_dev)?;
6874            if r_norm <= tol {
6875                diag.final_relative_residual = r_norm / rhs_norm;
6876                diag.stopping_reason = PcgStopReason::Converged;
6877                break;
6878            }
6879            launch_jacobi_mul(&stream, vector_module, &inv_diag_dev, &r_dev, &mut z_dev, k)?;
6880            diag.precond_apply_calls += 1;
6881            let rz_new = device_dot(&blas, &stream, k, &r_dev, &z_dev)?;
6882            if rz_new <= 0.0 || !rz_new.is_finite() {
6883                return Err(ArrowSchurGpuFailure::SchurFactorFailed {
6884                    reason: format!("SAE matrix-free GPU PCG: non-positive rᵀM⁻¹r={rz_new:e}"),
6885                });
6886            }
6887            let beta = rz_new / rz;
6888            launch_update_p(&stream, vector_module, &z_dev, beta, &mut p_dev, k)?;
6889            rz = rz_new;
6890        }
6891        if diag.stopping_reason != PcgStopReason::Converged {
6892            let r_norm = device_nrm2(&blas, &stream, k, &r_dev)?;
6893            diag.final_relative_residual = r_norm / rhs_norm;
6894            diag.stopping_reason = PcgStopReason::MaxIter;
6895        }
6896        let x = stream
6897            .clone_dtoh(&x_dev)
6898            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6899        Ok((Array1::from_vec(x), diag))
6900    }
6901
6902    pub(super) fn solve_reduced_beta_pcg_with_diagnostics(
6903        s_acc: &ndarray::Array2<f64>,
6904        rhs_beta: &Array1<f64>,
6905        max_iterations: usize,
6906        relative_tolerance: f64,
6907    ) -> Result<(Array1<f64>, ArrowPcgDiagnostics), ArrowSchurGpuFailure> {
6908        let k = rhs_beta.len();
6909        // #1017 dispatch re-key: this is an ITERATIVE device-resident PCG, not a
6910        // single GEMV. `S` (k×k) is uploaded once and reused for `max_iterations`
6911        // `S·p` GEMVs while only convergence scalars cross PCIe, so the staging
6912        // cost is amortised over the whole CG solve. Gating on the flops of ONE
6913        // `Gemv{k,k}` (`2·k²`) understates the work by the iteration count and
6914        // declines shapes (e.g. k≈512) whose total iterated arithmetic
6915        // `2·k²·iters` clears the device floor by orders of magnitude — the same
6916        // single-launch-breakeven miskey #1017 fixed for the framed reduced-Schur
6917        // matvec. Key on the CG-amortised total work via a `Gemm{k,k,iters}` whose
6918        // `flops()` is exactly `2·k²·iters`; numerics and kernels are untouched,
6919        // and the host falls back to the bit-identical CPU PCG when this declines.
6920        let cg_iters = max_iterations.max(1);
6921        let runtime = gam_gpu::linalg_dispatch::route_through_gpu(
6922            gam_gpu::linalg_dispatch::DispatchOp::Gemm {
6923                m: k,
6924                n: k,
6925                k: cg_iters,
6926            },
6927        )
6928        .ok_or(ArrowSchurGpuFailure::Unavailable)?;
6929        let stream = gam_gpu::device_runtime::cuda_context_for(runtime.device.ordinal)
6930            .and_then(|ctx| ctx.new_stream().ok())
6931            .ok_or(ArrowSchurGpuFailure::Unavailable)?;
6932        let blas = CudaBlas::new(stream.clone()).map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6933        let ctx = gam_gpu::device_runtime::cuda_context_for(runtime.device.ordinal)
6934            .ok_or(ArrowSchurGpuFailure::Unavailable)?;
6935        let vector_module = pcg_vector_module(&ctx)?;
6936
6937        // Jacobi diagonal from S; must be strictly positive for SPD.
6938        let mut inv_diag = vec![0.0_f64; k];
6939        for j in 0..k {
6940            let djj = s_acc[[j, j]];
6941            if !(djj > 0.0) {
6942                return Err(ArrowSchurGpuFailure::SchurFactorFailed {
6943                    reason: format!(
6944                        "reduced-β GPU PCG: Jacobi diagonal S[{j},{j}]={djj:e} not positive"
6945                    ),
6946                });
6947            }
6948            inv_diag[j] = 1.0 / djj;
6949        }
6950
6951        // Upload S column-major (S[row,col] at col*k + row).
6952        let mut s_host = vec![0.0_f64; k * k];
6953        for col in 0..k {
6954            for row in 0..k {
6955                s_host[col * k + row] = s_acc[[row, col]];
6956            }
6957        }
6958        let s_dev = stream
6959            .clone_htod(&s_host)
6960            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6961
6962        // Steihaug truncated-CG with Jacobi preconditioner, host scalar
6963        // recurrences and a device `S·p` matvec. The streaming reduced solve
6964        // uses an unbounded trust region (pure CG to tolerance).
6965        let rhs_norm = rhs_beta.iter().map(|v| v * v).sum::<f64>().sqrt();
6966        if rhs_norm == 0.0 {
6967            return Ok((Array1::<f64>::zeros(k), ArrowPcgDiagnostics::default()));
6968        }
6969        let tol = (relative_tolerance.max(0.0) * rhs_norm).max(1e-12);
6970
6971        // Device-resident PCG state. Only convergence scalars cross back during
6972        // the loop; x/r/z/p/Sp stay on CUDA until the final solution download.
6973        let mut x_dev = stream
6974            .alloc_zeros::<f64>(k)
6975            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6976        let mut r_dev = stream
6977            .clone_htod(
6978                rhs_beta
6979                    .as_slice()
6980                    .ok_or(ArrowSchurGpuFailure::Unavailable)?,
6981            )
6982            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6983        let inv_diag_dev = stream
6984            .clone_htod(&inv_diag)
6985            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6986        let mut z_dev = stream
6987            .alloc_zeros::<f64>(k)
6988            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6989        launch_jacobi_mul(&stream, vector_module, &inv_diag_dev, &r_dev, &mut z_dev, k)?;
6990        let mut p_dev = stream
6991            .alloc_zeros::<f64>(k)
6992            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6993        device_copy(&blas, &stream, k, &z_dev, &mut p_dev)?;
6994        let mut sp_dev = stream
6995            .alloc_zeros::<f64>(k)
6996            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6997        let mut rz = device_dot(&blas, &stream, k, &r_dev, &z_dev)?;
6998        let mut diag = ArrowPcgDiagnostics {
6999            precond_apply_calls: 1,
7000            stopping_reason: PcgStopReason::MaxIter,
7001            ..ArrowPcgDiagnostics::default()
7002        };
7003        if rz <= 0.0 || !rz.is_finite() {
7004            return Err(ArrowSchurGpuFailure::SchurFactorFailed {
7005                reason: format!("reduced-β GPU PCG: non-positive initial rᵀM⁻¹r={rz:e}"),
7006            });
7007        }
7008
7009        let max_iters = max_iterations.max(1);
7010        for _ in 0..max_iters {
7011            // sp = S · p (device GEMV, S column-major k×k, op = N).
7012            let gemv_cfg = GemvConfig::<f64> {
7013                trans: cublasOperation_t::CUBLAS_OP_N,
7014                m: to_i32(k).ok_or(ArrowSchurGpuFailure::Unavailable)?,
7015                n: to_i32(k).ok_or(ArrowSchurGpuFailure::Unavailable)?,
7016                alpha: 1.0,
7017                lda: to_i32(k).ok_or(ArrowSchurGpuFailure::Unavailable)?,
7018                incx: 1,
7019                beta: 0.0,
7020                incy: 1,
7021            };
7022            // SAFETY: s_dev is k×k column-major, p_dev / sp_dev length k.
7023            unsafe { blas.gemv(gemv_cfg, &s_dev, &p_dev, &mut sp_dev) }
7024                .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
7025            diag.matvec_calls += 1;
7026            diag.iterations += 1;
7027
7028            let p_sp = device_dot(&blas, &stream, k, &p_dev, &sp_dev)?;
7029            if !(p_sp > 0.0) {
7030                // Non-positive curvature on a (proximal-ridged) SPD system means
7031                // numerical breakdown; surface so the caller escalates.
7032                return Err(ArrowSchurGpuFailure::SchurFactorFailed {
7033                    reason: format!("reduced-β GPU PCG: non-positive curvature pᵀSp={p_sp:e}"),
7034                });
7035            }
7036            let alpha = rz / p_sp;
7037            device_axpy(&blas, &stream, k, alpha, &p_dev, &mut x_dev)?;
7038            device_axpy(&blas, &stream, k, -alpha, &sp_dev, &mut r_dev)?;
7039            let r_norm = device_nrm2(&blas, &stream, k, &r_dev)?;
7040            if r_norm <= tol {
7041                diag.final_relative_residual = r_norm / rhs_norm;
7042                diag.stopping_reason = PcgStopReason::Converged;
7043                break;
7044            }
7045            launch_jacobi_mul(&stream, vector_module, &inv_diag_dev, &r_dev, &mut z_dev, k)?;
7046            diag.precond_apply_calls += 1;
7047            let rz_new = device_dot(&blas, &stream, k, &r_dev, &z_dev)?;
7048            if rz_new <= 0.0 || !rz_new.is_finite() {
7049                return Err(ArrowSchurGpuFailure::SchurFactorFailed {
7050                    reason: format!("reduced-β GPU PCG: non-positive rᵀM⁻¹r={rz_new:e}"),
7051                });
7052            }
7053            let beta = rz_new / rz;
7054            launch_update_p(&stream, vector_module, &z_dev, beta, &mut p_dev, k)?;
7055            rz = rz_new;
7056        }
7057        if diag.stopping_reason != PcgStopReason::Converged {
7058            let r_norm = device_nrm2(&blas, &stream, k, &r_dev)?;
7059            diag.final_relative_residual = r_norm / rhs_norm;
7060            diag.stopping_reason = PcgStopReason::MaxIter;
7061        }
7062
7063        let x = stream
7064            .clone_dtoh(&x_dev)
7065            .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
7066        Ok((Array1::from_vec(x), diag))
7067    }
7068
7069    fn device_copy(
7070        blas: &CudaBlas,
7071        stream: &Arc<CudaStream>,
7072        n: usize,
7073        src: &CudaSlice<f64>,
7074        dst: &mut CudaSlice<f64>,
7075    ) -> Result<(), ArrowSchurGpuFailure> {
7076        let n_i = to_i32(n).ok_or(ArrowSchurGpuFailure::Unavailable)?;
7077        let (src_ptr, _src_rec) = src.device_ptr(stream);
7078        let (dst_ptr, _dst_rec) = dst.device_ptr_mut(stream);
7079        // SAFETY: src and dst are live device allocations on this stream with at
7080        // least n contiguous f64 entries and unit stride.
7081        let status = unsafe {
7082            cudarc::cublas::sys::cublasDcopy_v2(
7083                *blas.handle(),
7084                n_i,
7085                src_ptr as *const f64,
7086                1,
7087                dst_ptr as *mut f64,
7088                1,
7089            )
7090        };
7091        if status == cublasStatus_t::CUBLAS_STATUS_SUCCESS {
7092            Ok(())
7093        } else {
7094            Err(ArrowSchurGpuFailure::Unavailable)
7095        }
7096    }
7097
7098    fn device_axpy(
7099        blas: &CudaBlas,
7100        stream: &Arc<CudaStream>,
7101        n: usize,
7102        alpha: f64,
7103        x: &CudaSlice<f64>,
7104        y: &mut CudaSlice<f64>,
7105    ) -> Result<(), ArrowSchurGpuFailure> {
7106        let n_i = to_i32(n).ok_or(ArrowSchurGpuFailure::Unavailable)?;
7107        let (x_ptr, _x_rec) = x.device_ptr(stream);
7108        let (y_ptr, _y_rec) = y.device_ptr_mut(stream);
7109        // SAFETY: x and y are live device allocations on this stream with at
7110        // least n contiguous f64 entries and unit stride; cuBLAS only reads alpha.
7111        let status = unsafe {
7112            cudarc::cublas::sys::cublasDaxpy_v2(
7113                *blas.handle(),
7114                n_i,
7115                &alpha,
7116                x_ptr as *const f64,
7117                1,
7118                y_ptr as *mut f64,
7119                1,
7120            )
7121        };
7122        if status == cublasStatus_t::CUBLAS_STATUS_SUCCESS {
7123            Ok(())
7124        } else {
7125            Err(ArrowSchurGpuFailure::Unavailable)
7126        }
7127    }
7128
7129    /// As [`device_axpy`] but with explicit strides, so a unit-stride source
7130    /// (e.g. an all-ones vector) can target a matrix diagonal with `incy = dim+1`.
7131    /// Used by [`ResidentBaseArrowFrame`] to add `ridge_beta` to the resident
7132    /// `k×k` Schur base on-device without re-uploading the border block.
7133    fn device_axpy_strided(
7134        blas: &CudaBlas,
7135        stream: &Arc<CudaStream>,
7136        n: usize,
7137        alpha: f64,
7138        x: &CudaSlice<f64>,
7139        incx: usize,
7140        y: &mut CudaSlice<f64>,
7141        incy: usize,
7142    ) -> Result<(), ArrowSchurGpuFailure> {
7143        let n_i = to_i32(n).ok_or(ArrowSchurGpuFailure::Unavailable)?;
7144        let incx_i = to_i32(incx).ok_or(ArrowSchurGpuFailure::Unavailable)?;
7145        let incy_i = to_i32(incy).ok_or(ArrowSchurGpuFailure::Unavailable)?;
7146        let (x_ptr, _x_rec) = x.device_ptr(stream);
7147        let (y_ptr, _y_rec) = y.device_ptr_mut(stream);
7148        // SAFETY: x spans ≥ 1+(n−1)·incx entries and y spans ≥ 1+(n−1)·incy
7149        // entries, both live on this stream; cuBLAS only reads alpha by pointer.
7150        let status = unsafe {
7151            cudarc::cublas::sys::cublasDaxpy_v2(
7152                *blas.handle(),
7153                n_i,
7154                &alpha,
7155                x_ptr as *const f64,
7156                incx_i,
7157                y_ptr as *mut f64,
7158                incy_i,
7159            )
7160        };
7161        if status == cublasStatus_t::CUBLAS_STATUS_SUCCESS {
7162            Ok(())
7163        } else {
7164            Err(ArrowSchurGpuFailure::Unavailable)
7165        }
7166    }
7167
7168    fn device_dot(
7169        blas: &CudaBlas,
7170        stream: &Arc<CudaStream>,
7171        n: usize,
7172        x: &CudaSlice<f64>,
7173        y: &CudaSlice<f64>,
7174    ) -> Result<f64, ArrowSchurGpuFailure> {
7175        let n_i = to_i32(n).ok_or(ArrowSchurGpuFailure::Unavailable)?;
7176        let (x_ptr, _x_rec) = x.device_ptr(stream);
7177        let (y_ptr, _y_rec) = y.device_ptr(stream);
7178        let mut result = 0.0_f64;
7179        // SAFETY: x and y are live device allocations on this stream with at
7180        // least n contiguous f64 entries and unit stride; result is a valid host
7181        // out-pointer for the cuBLAS scalar.
7182        let status = unsafe {
7183            cudarc::cublas::sys::cublasDdot_v2(
7184                *blas.handle(),
7185                n_i,
7186                x_ptr as *const f64,
7187                1,
7188                y_ptr as *const f64,
7189                1,
7190                &mut result,
7191            )
7192        };
7193        if status == cublasStatus_t::CUBLAS_STATUS_SUCCESS {
7194            Ok(result)
7195        } else {
7196            Err(ArrowSchurGpuFailure::Unavailable)
7197        }
7198    }
7199
7200    fn device_nrm2(
7201        blas: &CudaBlas,
7202        stream: &Arc<CudaStream>,
7203        n: usize,
7204        x: &CudaSlice<f64>,
7205    ) -> Result<f64, ArrowSchurGpuFailure> {
7206        let n_i = to_i32(n).ok_or(ArrowSchurGpuFailure::Unavailable)?;
7207        let (x_ptr, _x_rec) = x.device_ptr(stream);
7208        let mut result = 0.0_f64;
7209        // SAFETY: x is a live device allocation on this stream with at least n
7210        // contiguous f64 entries and unit stride; result is a valid host
7211        // out-pointer for the cuBLAS scalar.
7212        let status = unsafe {
7213            cudarc::cublas::sys::cublasDnrm2_v2(
7214                *blas.handle(),
7215                n_i,
7216                x_ptr as *const f64,
7217                1,
7218                &mut result,
7219            )
7220        };
7221        if status == cublasStatus_t::CUBLAS_STATUS_SUCCESS {
7222            Ok(result)
7223        } else {
7224            Err(ArrowSchurGpuFailure::Unavailable)
7225        }
7226    }
7227
7228    #[cfg(test)]
7229    mod tests {
7230        //! #1551 device-side framed-matvec triage. Lives inside `mod cuda` so it
7231        //! can call the private kernel launchers directly (no test-only public
7232        //! seam, which the ban-scanner forbids). A bare `#[cfg(test)] mod tests`
7233        //! is the one form the scanner permits.
7234        use crate::arrow_schur::{ArrowSchurSystem, DeviceSaeFrameData, DeviceSaePcgData, DeviceSaeSmoothBlock, FactoredFrameGBlock};
7235        use ndarray::Array2;
7236
7237        /// Build the tiny hand-verifiable framed SAE fixture (2 atoms, 2 rows).
7238        /// Shared by the resident-frame residency test below and mirrors the
7239        /// `#1551` matvec-triage fixture so both exercise the same operand shapes.
7240        fn tiny_framed_fixture() -> (ArrowSchurSystem, DeviceSaePcgData) {
7241            let p = 3usize;
7242            let ranks = vec![2usize, 3usize];
7243            let basis_sizes = vec![2usize, 2usize];
7244            let mut border_offsets = Vec::new();
7245            let mut acc = 0usize;
7246            for k in 0..2 {
7247                border_offsets.push(acc);
7248                acc += basis_sizes[k] * ranks[k];
7249            }
7250            let border_dim = acc;
7251            let frame_of = |k: usize| -> Array2<f64> {
7252                Array2::from_shape_fn((p, ranks[k]), |(i, j)| {
7253                    0.1 + 0.2 * ((i + 1) as f64) * ((j + 1 + 2 * k) as f64)
7254                })
7255            };
7256            let frames: Vec<Array2<f64>> = (0..2).map(frame_of).collect();
7257            let w_of = |i: usize, j: usize| -> Array2<f64> {
7258                let (ui, uj) = (&frames[i], &frames[j]);
7259                Array2::from_shape_fn((ranks[i], ranks[j]), |(a, b)| {
7260                    (0..p).map(|c| ui[[c, a]] * uj[[c, b]]).sum()
7261                })
7262            };
7263            let mut frame_blocks = Vec::new();
7264            for &(i, j) in &[(0usize, 0usize), (1usize, 1usize), (0, 1), (1, 0)] {
7265                let (mi, mj) = (basis_sizes[i], basis_sizes[j]);
7266                let mut g =
7267                    Array2::<f64>::from_shape_fn((mi, mj), |(r, c)| 0.1 * (r + 2 * c + 1) as f64);
7268                if i == j {
7269                    for r in 0..mi.min(mj) {
7270                        g[[r, r]] += mi as f64 + 2.0;
7271                    }
7272                }
7273                frame_blocks.push(FactoredFrameGBlock {
7274                    atom_i: i,
7275                    atom_j: j,
7276                    g,
7277                    w: w_of(i, j),
7278                });
7279            }
7280            let mut smooth_blocks = Vec::new();
7281            for k in 0..2 {
7282                let m = basis_sizes[k];
7283                let mut s =
7284                    Array2::<f64>::from_shape_fn((m, m), |(r, c)| 0.05 * (r + c + 1) as f64);
7285                for r in 0..m {
7286                    s[[r, r]] += 1.0;
7287                }
7288                smooth_blocks.push(DeviceSaeSmoothBlock {
7289                    global_offset: border_offsets[k],
7290                    factor_a: s,
7291                });
7292            }
7293            let smooth_ranks = ranks.clone();
7294            let n = 2usize;
7295            let q = 2usize;
7296            let mut sys = ArrowSchurSystem::new(n, q, border_dim);
7297            let mut row_htbeta = Vec::new();
7298            for i in 0..n {
7299                let mut htt =
7300                    Array2::<f64>::from_shape_fn((q, q), |(r, c)| 0.3 * (r + c + 1) as f64);
7301                for r in 0..q {
7302                    htt[[r, r]] += q as f64 + 2.0;
7303                }
7304                sys.rows[i].htt = htt;
7305                let mut slab = vec![0.0_f64; q * border_dim];
7306                for c in 0..q {
7307                    for col in 0..border_dim {
7308                        let v = 0.01 * ((c + 1) * (col + 1) + i) as f64;
7309                        slab[c * border_dim + col] = v;
7310                        sys.rows[i].htbeta[[c, col]] = v;
7311                    }
7312                }
7313                row_htbeta.push(slab);
7314            }
7315            let data = DeviceSaePcgData {
7316                p,
7317                beta_dim: border_dim,
7318                a_phi: std::sync::Arc::from(Vec::new().into_boxed_slice()),
7319                local_jac: std::sync::Arc::from(Vec::new().into_boxed_slice()),
7320                smooth_blocks,
7321                sparse_g_blocks: Vec::new(),
7322                frame: Some(DeviceSaeFrameData {
7323                    ranks,
7324                    basis_sizes,
7325                    border_offsets,
7326                    frame_blocks,
7327                    smooth_ranks,
7328                    row_htbeta,
7329                }),
7330            };
7331            (sys, data)
7332        }
7333
7334        /// #1017 residency invariant (no GPU needed — pure host marshalling).
7335        ///
7336        /// Verifies the property the whole resident-frame optimization rests on:
7337        /// across the LM ridge ladder the framed SAE operands split into a
7338        /// ridge-INDEPENDENT part (`flatten_frame_host_operands` — which takes no
7339        /// `ridge_t` at all, a compile-time proof, and is deterministic build to
7340        /// build) and a single ridge-DEPENDENT buffer, the per-row factored
7341        /// `ainv` (`compute_ainv_host`). So the resident frame can upload the
7342        /// ridge-independent operands once and recompute only `ainv` per trial,
7343        /// removing `(trials − 1) × operand_bytes` of re-upload with a
7344        /// bit-identical solve.
7345        #[test]
7346        fn sae_resident_frame_only_ainv_is_ridge_dependent_1017() {
7347            let (sys, data) = tiny_framed_fixture();
7348            let frame = data.frame.as_ref().expect("framed fixture");
7349
7350            // Ridge-INDEPENDENT operands: identical across builds (deterministic;
7351            // no `ridge_t` in the signature).
7352            let host_a = super::super::flatten_frame_host_operands(&sys, &data, frame)
7353                .expect("host operands a");
7354            let host_b = super::super::flatten_frame_host_operands(&sys, &data, frame)
7355                .expect("host operands b");
7356            assert_eq!(
7357                host_a.s_data, host_b.s_data,
7358                "smooth λS must be ridge-independent"
7359            );
7360            assert_eq!(
7361                host_a.g_data, host_b.g_data,
7362                "frame G must be ridge-independent"
7363            );
7364            assert_eq!(
7365                host_a.w_data, host_b.w_data,
7366                "frame W must be ridge-independent"
7367            );
7368            assert_eq!(host_a.htb, host_b.htb, "row H_tβ must be ridge-independent");
7369            assert_eq!(host_a.q_of, host_b.q_of);
7370
7371            let report = data.operand_byte_report();
7372            assert!(report.framed, "fixture must exercise the framed lane");
7373            assert!(
7374                report.total_bytes > 0,
7375                "framed operands must have nonzero bytes"
7376            );
7377
7378            // The ONLY ridge-dependent buffer: ainv. Deterministic at a fixed
7379            // ridge (safe to reuse across the ladder), changing with ridge_t
7380            // (must be recomputed each trial).
7381            let ainv_lo = super::super::compute_ainv_host(
7382                &sys,
7383                &host_a.q_of,
7384                host_a.max_q,
7385                host_a.n_rows,
7386                1e-3,
7387            )
7388            .expect("ainv lo");
7389            let ainv_lo2 = super::super::compute_ainv_host(
7390                &sys,
7391                &host_a.q_of,
7392                host_a.max_q,
7393                host_a.n_rows,
7394                1e-3,
7395            )
7396            .expect("ainv lo repeat");
7397            let ainv_hi = super::super::compute_ainv_host(
7398                &sys,
7399                &host_a.q_of,
7400                host_a.max_q,
7401                host_a.n_rows,
7402                1e3,
7403            )
7404            .expect("ainv hi");
7405            assert_eq!(
7406                ainv_lo, ainv_lo2,
7407                "ainv must be deterministic at a fixed ridge"
7408            );
7409            let max_diff = ainv_lo
7410                .iter()
7411                .zip(&ainv_hi)
7412                .map(|(a, b)| (a - b).abs())
7413                .fold(0.0_f64, f64::max);
7414            assert!(
7415                max_diff > 1e-6,
7416                "ainv MUST change with ridge_t (else the split is wrong); max_diff={max_diff:e}"
7417            );
7418
7419            // Ladder projection: the per-trial flatten re-uploaded `total_bytes`
7420            // on every trial; the resident frame uploads it once, so a ladder of
7421            // `trials` removes `(trials − 1) × total_bytes`.
7422            let trials = crate::arrow_schur::DEFAULT_PROXIMAL_MAX_ATTEMPTS + 1;
7423            let saved = report.total_bytes * (trials - 1);
7424            assert!(saved > 0);
7425            eprintln!(
7426                "#1017 resident-frame ladder saving: {trials} trials × {}B ridge-independent \
7427                 operand upload → resident removes {saved}B, re-uploading only ainv \
7428                 ({}rows × {}² × 8B) per trial",
7429                report.total_bytes, host_a.n_rows, host_a.max_q
7430            );
7431        }
7432
7433    }
7434}
7435
7436#[cfg(test)]
7437mod tests {
7438    use super::*;
7439    use crate::arrow_schur::ArrowSchurSystem;
7440    use ndarray::{Array2, ArrayView1};
7441
7442    fn build_fixture(n: usize, d: usize, k: usize, seed: u64) -> ArrowSchurSystem {
7443        let mut sys = ArrowSchurSystem::new(n, d, k);
7444        let mut state = seed.wrapping_mul(0x9E37_79B9_7F4A_7C15);
7445        let mut sample = || -> f64 {
7446            state = state
7447                .wrapping_mul(6364136223846793005)
7448                .wrapping_add(1442695040888963407);
7449            ((state >> 33) as f64) / ((1u64 << 31) as f64) - 1.0
7450        };
7451        for row in &mut sys.rows {
7452            let mut a = Array2::<f64>::zeros((d, d));
7453            for r in 0..d {
7454                for c in 0..d {
7455                    a[[r, c]] = sample();
7456                }
7457            }
7458            let mut htt = a.t().dot(&a);
7459            for r in 0..d {
7460                htt[[r, r]] += d as f64 + 1.0;
7461            }
7462            row.htt = htt;
7463            for r in 0..d {
7464                for c in 0..k {
7465                    row.htbeta[[r, c]] = 0.1 * sample();
7466                }
7467                row.gt[r] = sample();
7468            }
7469        }
7470        let mut hbb_a = Array2::<f64>::zeros((k, k));
7471        for r in 0..k {
7472            for c in 0..k {
7473                hbb_a[[r, c]] = sample();
7474            }
7475        }
7476        let mut hbb = hbb_a.t().dot(&hbb_a);
7477        for r in 0..k {
7478            hbb[[r, r]] += k as f64 + 1.0;
7479        }
7480        sys.hbb = hbb;
7481        for r in 0..k {
7482            sys.gb[r] = sample();
7483        }
7484        sys
7485    }
7486
7487    /// The Gershgorin ridge bump must actually make a known-indefinite block PD
7488    /// on the first retry — the whole point of #1711. Verified directly by
7489    /// re-factoring `H_tt + (ridge_t + bump)·I` with the same Cholesky guard the
7490    /// device readback uses, for blocks whose `λ_min` is known in closed form.
7491    #[test]
7492    fn ridge_bump_makes_known_indefinite_blocks_pd() {
7493        // `cholesky_factor_in_place` / `CholeskyGuard` are already in scope via
7494        // `super::*` (imported at the top of the module).
7495        // A few blocks with a CLOSED-FORM smallest eigenvalue, all at ridge_t=0.
7496        // (label, matrix, λ_min) — the bump must clear each one.
7497        let neg_identity = Array2::<f64>::from_diag(&Array1::from_elem(8, -1.0)); // λ_min = -1
7498        let scaled_neg = Array2::<f64>::from_diag(&Array1::from_elem(4, -250.0)); // λ_min = -250
7499        // Symmetric 2×2 [[1, 2], [2, 1]] has eigenvalues 3 and -1 → indefinite.
7500        let mut indef2 = Array2::<f64>::zeros((2, 2));
7501        indef2[[0, 0]] = 1.0;
7502        indef2[[1, 1]] = 1.0;
7503        indef2[[0, 1]] = 2.0;
7504        indef2[[1, 0]] = 2.0;
7505        // A genuinely PD block must get a bump that is the bare rounding margin
7506        // only (deficit 0), and must still factor — the helper is defensive.
7507        let pd = Array2::<f64>::from_diag(&Array1::from_elem(3, 5.0));
7508
7509        for (label, block) in [
7510            ("-I (λ_min=-1)", neg_identity),
7511            ("-250·I (λ_min=-250)", scaled_neg),
7512            ("[[1,2],[2,1]] (λ_min=-1)", indef2),
7513            ("5·I (PD)", pd),
7514        ] {
7515            let ridge_t = 0.0;
7516            let bump = ridge_bump_to_make_pd(block.view(), ridge_t);
7517            assert!(
7518                bump > 0.0 && bump.is_finite(),
7519                "[{label}] bump must be strictly positive and finite, got {bump:e}"
7520            );
7521            let d = block.nrows();
7522            let mut shifted = block.clone();
7523            for i in 0..d {
7524                shifted[[i, i]] += ridge_t + bump;
7525            }
7526            assert!(
7527                cholesky_factor_in_place(shifted.view(), CholeskyGuard::NonnegativePivot).is_some(),
7528                "[{label}] H_tt + (ridge_t + bump={bump:e})·I must be PD after the \
7529                 Gershgorin bump, but the Cholesky still rejected it"
7530            );
7531        }
7532    }
7533
7534    /// The column-major variant (multi-GPU tile path) must agree with the
7535    /// row-major helper for a symmetric block, since Gershgorin edges are
7536    /// invariant under reading the symmetric matrix by row vs by column. The
7537    /// colmajor variant takes the bound at ridge_t=0 (the ridge is already baked
7538    /// into the diagonal it reads), so compare against `ridge_bump_to_make_pd`
7539    /// with `ridge_t = 0`.
7540    ///
7541    /// Gated to linux: `ridge_bump_to_make_pd_colmajor` only exists on the
7542    /// linux CUDA tile path, so the parity test runs where the function does.
7543    #[cfg(target_os = "linux")]
7544    #[test]
7545    fn ridge_bump_colmajor_matches_rowmajor_for_symmetric_block() {
7546        // Symmetric 3×3 with a negative-definite-ish diagonal and off-diagonals.
7547        let mut a = Array2::<f64>::zeros((3, 3));
7548        a[[0, 0]] = -2.0;
7549        a[[1, 1]] = 0.5;
7550        a[[2, 2]] = 1.0;
7551        a[[0, 1]] = 0.3;
7552        a[[1, 0]] = 0.3;
7553        a[[1, 2]] = -0.4;
7554        a[[2, 1]] = -0.4;
7555        a[[0, 2]] = 0.1;
7556        a[[2, 0]] = 0.1;
7557
7558        let row_major_bump = ridge_bump_to_make_pd(a.view(), 0.0);
7559
7560        // Flatten column-major: block[c*d + r] = a[[r, c]].
7561        let d = 3;
7562        let mut col_major = vec![0.0_f64; d * d];
7563        for c in 0..d {
7564            for r in 0..d {
7565                col_major[c * d + r] = a[[r, c]];
7566            }
7567        }
7568        let col_major_bump = ridge_bump_to_make_pd_colmajor(&col_major, d);
7569
7570        assert!(
7571            (row_major_bump - col_major_bump).abs() <= 1e-12 * row_major_bump.max(1.0),
7572            "colmajor bump {col_major_bump:e} must match rowmajor bump \
7573             {row_major_bump:e} for a symmetric block"
7574        );
7575
7576        // And the bump must actually make it PD (sanity, same as the row-major test).
7577        let mut shifted = a.clone();
7578        for i in 0..d {
7579            shifted[[i, i]] += col_major_bump;
7580        }
7581        assert!(
7582            cholesky_factor_in_place(shifted.view(), CholeskyGuard::NonnegativePivot).is_some(),
7583            "colmajor Gershgorin bump must make the symmetric block PD"
7584        );
7585    }
7586
7587    fn device_pcg_fixture(k: usize) -> (Array2<f64>, Array1<f64>) {
7588        let mut s = Array2::<f64>::zeros((k, k));
7589        for row in 0..k {
7590            s[[row, row]] = 2.5 + 0.001 * ((row % 17) as f64);
7591            if row + 1 < k {
7592                s[[row, row + 1]] = -0.05;
7593                s[[row + 1, row]] = -0.05;
7594            }
7595            if row + 7 < k {
7596                s[[row, row + 7]] = 0.01;
7597                s[[row + 7, row]] = 0.01;
7598            }
7599        }
7600        let rhs = Array1::from_shape_fn(k, |idx| ((idx as f64 + 1.0) * 0.013).sin());
7601        (s, rhs)
7602    }
7603
7604    fn dense_pcg_cpu_reference(
7605        s: &Array2<f64>,
7606        rhs: &Array1<f64>,
7607        max_iterations: usize,
7608        relative_tolerance: f64,
7609    ) -> Array1<f64> {
7610        let k = rhs.len();
7611        let rhs_norm = rhs.iter().map(|v| v * v).sum::<f64>().sqrt();
7612        if rhs_norm == 0.0 {
7613            return Array1::<f64>::zeros(k);
7614        }
7615        let tol = (relative_tolerance.max(0.0) * rhs_norm).max(1e-12);
7616        let inv_diag: Vec<f64> = (0..k).map(|idx| 1.0 / s[[idx, idx]]).collect();
7617        let mut x = Array1::<f64>::zeros(k);
7618        let mut r = rhs.clone();
7619        let mut z = Array1::from_shape_fn(k, |idx| inv_diag[idx] * r[idx]);
7620        let mut p = z.clone();
7621        let mut sp = Array1::<f64>::zeros(k);
7622        let mut rz = r.iter().zip(z.iter()).map(|(a, b)| a * b).sum::<f64>();
7623        for _ in 0..max_iterations.max(1) {
7624            for row in 0..k {
7625                let mut acc = 0.0;
7626                for col in 0..k {
7627                    acc += s[[row, col]] * p[col];
7628                }
7629                sp[row] = acc;
7630            }
7631            let p_sp = p.iter().zip(sp.iter()).map(|(a, b)| a * b).sum::<f64>();
7632            let alpha = rz / p_sp;
7633            for idx in 0..k {
7634                x[idx] += alpha * p[idx];
7635                r[idx] -= alpha * sp[idx];
7636            }
7637            let r_norm = r.iter().map(|v| v * v).sum::<f64>().sqrt();
7638            if r_norm <= tol {
7639                break;
7640            }
7641            for idx in 0..k {
7642                z[idx] = inv_diag[idx] * r[idx];
7643            }
7644            let rz_next = r.iter().zip(z.iter()).map(|(a, b)| a * b).sum::<f64>();
7645            let beta = rz_next / rz;
7646            for idx in 0..k {
7647                p[idx] = z[idx] + beta * p[idx];
7648            }
7649            rz = rz_next;
7650        }
7651        x
7652    }
7653
7654    #[test]
7655    fn device_resident_pcg_matches_cpu_reference_when_cuda_admits() {
7656        let (s, rhs) = device_pcg_fixture(512);
7657        let max_iterations = 200usize;
7658        let relative_tolerance = 1.0e-12;
7659        let cpu = dense_pcg_cpu_reference(&s, &rhs, max_iterations, relative_tolerance);
7660        let (device, diag) = match solve_reduced_beta_pcg_with_diagnostics(
7661            &s,
7662            &rhs,
7663            max_iterations,
7664            relative_tolerance,
7665        ) {
7666            Ok(result) => result,
7667            // #1017 — fail loud, never skip-pass: this fixture clears the device
7668            // offload floor, so a CUDA device that is PRESENT yet declines/returns
7669            // Err means the device PCG kernel does not run on GPU (a real fault that
7670            // must not masquerade as a pass via this skip). Legit skip ONLY when no
7671            // usable CUDA device exists (CPU CI). The exact `ArrowSchurGpuFailure`
7672            // variant is folded into the assert message as the diagnostic.
7673            Err(failure) => {
7674                assert!(
7675                    gam_gpu::device_runtime::GpuRuntime::resolve(gam_gpu::GpuPolicy::Auto)
7676                        .unwrap_or_else(|error| {
7677                            panic!("GPU probe fault in reduced-beta PCG test: {error}")
7678                        })
7679                        .is_none(),
7680                    "#1017: CUDA device present but the device reduced-beta PCG \
7681                     declined/faulted instead of returning a result (tag: {failure:?}) — \
7682                     the kernel does not run correctly on GPU"
7683                );
7684                return;
7685            }
7686        };
7687        let max_err = cpu
7688            .iter()
7689            .zip(device.iter())
7690            .map(|(a, b)| (a - b).abs())
7691            .fold(0.0_f64, f64::max);
7692        assert!(
7693            max_err <= 1.0e-10,
7694            "device resident PCG parity failed: max_err={max_err:e}, diag={diag:?}"
7695        );
7696        assert!(diag.matvec_calls > 0);
7697        assert_eq!(diag.matvec_calls, diag.iterations);
7698    }
7699
7700    #[test]
7701    fn dense_reference_matches_independent_solve() {
7702        let sys = build_fixture(4, 5, 3, 7);
7703        let solution = solve_arrow_newton_step_dense_reference(&sys, 0.0, 0.0).unwrap();
7704        // Re-solve by an independent matrix build and a textbook
7705        // Gaussian-elimination Cholesky to guard against typos in the
7706        // reference implementation itself.
7707        let n = sys.rows.len();
7708        let d = sys.d;
7709        let k = sys.k;
7710        let total = n * d + k;
7711        let mut h = Array2::<f64>::zeros((total, total));
7712        let mut g = ndarray::Array1::<f64>::zeros(total);
7713        for (i, row) in sys.rows.iter().enumerate() {
7714            let base = i * d;
7715            for c in 0..d {
7716                for r in 0..d {
7717                    h[[base + r, base + c]] = row.htt[[r, c]];
7718                }
7719            }
7720            for c in 0..k {
7721                for r in 0..d {
7722                    h[[base + r, n * d + c]] = row.htbeta[[r, c]];
7723                    h[[n * d + c, base + r]] = row.htbeta[[r, c]];
7724                }
7725            }
7726            for r in 0..d {
7727                g[base + r] = row.gt[r];
7728            }
7729        }
7730        for c in 0..k {
7731            for r in 0..k {
7732                h[[n * d + r, n * d + c]] += sys.hbb[[r, c]];
7733            }
7734            g[n * d + c] = sys.gb[c];
7735        }
7736        let l = cholesky_factor_in_place(h.view(), CholeskyGuard::NonnegativePivot).unwrap();
7737        let rhs = g.mapv(|v| -v);
7738        let expected = cholesky_solve_vector(l.view(), rhs.view());
7739        for i in 0..n * d {
7740            assert!(
7741                (solution.delta_t[i] - expected[i]).abs() < 1e-10 * (1.0 + expected[i].abs()),
7742                "delta_t[{i}] mismatch: got {} expected {}",
7743                solution.delta_t[i],
7744                expected[i]
7745            );
7746        }
7747        for a in 0..k {
7748            assert!(
7749                (solution.delta_beta[a] - expected[n * d + a]).abs()
7750                    < 1e-10 * (1.0 + expected[n * d + a].abs()),
7751                "delta_beta[{a}] mismatch"
7752            );
7753        }
7754    }
7755
7756    /// #1017: the row-procedural reduced-Schur matvec (the matrix-free SAE
7757    /// host backend) auto-fans its per-row point-elimination sum across rayon
7758    /// over fixed row chunks when at the top level (`n ≥
7759    /// SCHUR_MATVEC_PARALLEL_ROW_MIN`), and stays serial when already inside a
7760    /// rayon worker. The chunk-ordered fold makes the parallel result
7761    /// **deterministic** (two parallel calls are bit-identical — scheduling
7762    /// cannot change the numbers) and it agrees with the serial accumulation up
7763    /// to ULP-scale chunk reassociation (the #1017 verification gate). That
7764    /// reassociation is a genuine f64 departure from serial, so the criterion
7765    /// ranking across topology candidates is stable only up to the reassociation
7766    /// margin: a near-tie winner inside that margin can flip. This is NOT an
7767    /// exact no-move guarantee (#1211); for that, the ranking path must use the
7768    /// fixed-order serial accumulation.
7769    #[test]
7770    fn row_procedural_matvec_parallel_deterministic_and_matches_serial() {
7771        use crate::arrow_schur::SCHUR_MATVEC_PARALLEL_ROW_MIN;
7772        let n = SCHUR_MATVEC_PARALLEL_ROW_MIN + 96; // trips the parallel path
7773        let d = 3usize;
7774        let k = 24usize;
7775        let mut sys = build_fixture(n, d, k, 0xA17C_0FFE);
7776        // Install a matrix-free forward/transpose pair that reads the dense
7777        // `htbeta` slabs the fixture already populated, so the procedural
7778        // backend has a well-defined operator to apply (and exercises exactly
7779        // the sparse gather/scatter the SAE Kronecker path drives).
7780        let slabs: Vec<Array2<f64>> = sys.rows.iter().map(|row| row.htbeta.clone()).collect();
7781        let forward_slabs = slabs.clone();
7782        let transpose_slabs = slabs;
7783        sys.set_row_htbeta_operator(
7784            move |row: usize, x: ArrayView1<'_, f64>, out: &mut Array1<f64>| {
7785                let h = &forward_slabs[row];
7786                for r in 0..h.nrows() {
7787                    let mut acc = 0.0_f64;
7788                    for c in 0..h.ncols() {
7789                        acc += h[[r, c]] * x[c];
7790                    }
7791                    out[r] = acc;
7792                }
7793            },
7794            move |row: usize, v: ArrayView1<'_, f64>, out: &mut Array1<f64>| {
7795                let h = &transpose_slabs[row];
7796                for r in 0..h.nrows() {
7797                    for c in 0..h.ncols() {
7798                        out[c] += h[[r, c]] * v[r];
7799                    }
7800                }
7801            },
7802        );
7803
7804        let matvec = gpu_schur_matvec_backend(&sys, 0.0, 0.0)
7805            .expect("row-procedural matvec backend builds for matrix-free system");
7806        let x = Array1::from_shape_fn(k, |i| ((i as f64 + 1.0) * 0.37).sin());
7807
7808        // Top-level call: auto-selects the parallel chunk-fold. Run twice and
7809        // assert bit-identity — the chunk-ordered reduction must not depend on
7810        // thread scheduling.
7811        let mut out_parallel_a = Array1::<f64>::zeros(k);
7812        matvec(&x, &mut out_parallel_a);
7813        let mut out_parallel_b = Array1::<f64>::zeros(k);
7814        matvec(&x, &mut out_parallel_b);
7815        for a in 0..k {
7816            assert_eq!(
7817                out_parallel_a[a].to_bits(),
7818                out_parallel_b[a].to_bits(),
7819                "row-procedural matvec parallel reduction is non-deterministic at index {a}"
7820            );
7821        }
7822
7823        // Inside a rayon worker: auto-selects the serial path (nested-rayon
7824        // guard). `install` runs the closure on a pool thread, so
7825        // `current_thread_index()` is `Some`. The serial running sum and the
7826        // chunk-ordered parallel fold differ only by f64 reassociation.
7827        let mut out_serial = Array1::<f64>::zeros(k);
7828        rayon::ThreadPoolBuilder::new()
7829            .num_threads(2)
7830            .build()
7831            .expect("build rayon pool")
7832            .install(|| matvec(&x, &mut out_serial));
7833
7834        let max_abs = out_serial.iter().fold(0.0_f64, |m, v| m.max(v.abs()));
7835        for a in 0..k {
7836            let diff = (out_parallel_a[a] - out_serial[a]).abs();
7837            assert!(
7838                diff <= 1e-12 * (1.0 + max_abs),
7839                "row-procedural matvec parallel vs serial diverged beyond reassociation \
7840                 at index {a}: {} vs {} (diff={diff:e})",
7841                out_parallel_a[a],
7842                out_serial[a]
7843            );
7844        }
7845    }
7846
7847}