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