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