Skip to main content

gam_terms/basis/
sphere_gpu.rs

1//! GPU NVRTC Wahba intrinsic-S2 kernel matrix construction.
2//!
3//! This module owns the device-side construction of the Wahba reproducing
4//! kernel basis matrix on the 2-sphere using the **finite truncated
5//! spectral Legendre series**
6//!
7//! `K_L(γ) = Σ_{ℓ=1..L} c_ℓ · P_ℓ(cos γ)`,
8//!
9//! evaluated entry-by-entry against the 3-term Legendre recurrence kept
10//! in registers. The host CPU parity target is the matching
11//! `SphereWahbaKernel::SobolevTruncated { lmax }` /
12//! `SphereWahbaKernel::PseudoTruncated { lmax }` variant added to
13//! `src/terms/basis.rs` (single source: same recurrence, same c_ℓ).
14//!
15//! The device path evaluates the raw column-major kernel matrix with `f64`
16//! Legendre recurrence math. Host code owns centering, constraints, and solver
17//! assembly in `basis.rs`.
18
19use std::sync::OnceLock;
20
21use ndarray::{Array2, ArrayView2};
22
23use gam_gpu::gpu_error::GpuError;
24#[cfg(target_os = "linux")]
25use gam_gpu::gpu_error::GpuResultExt;
26use gam_gpu::{GpuDecision, GpuKernel, decide};
27
28#[cfg(target_os = "linux")]
29use std::collections::HashMap;
30#[cfg(target_os = "linux")]
31use std::sync::{Arc, Mutex};
32
33#[cfg(target_os = "linux")]
34use cudarc::driver::{CudaContext, CudaModule, CudaSlice, CudaStream};
35
36/// Which truncated-spectral Wahba kernel to evaluate on device. Matches
37/// the CPU `SphereWahbaKernel::{SobolevTruncated, PseudoTruncated}` so
38/// parity tests are well-defined.
39#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
40pub enum SphereSpectralKernelKind {
41    /// `c_ℓ = (2ℓ+1) / (4π · [ℓ(ℓ+1)]^m)` — true `H^m(S²)` Sobolev RKHS.
42    Sobolev,
43    /// `c_ℓ = 2 / (4π · Π_{k=1..m+1}(ℓ + k))` — Wahba 1981 pseudo-spline.
44    Pseudo,
45}
46
47impl SphereSpectralKernelKind {
48    /// `c_0 = 0`, `c_ℓ = c_ℓ(m)` for `ℓ = 1..=lmax`. Returned vector has
49    /// length `lmax + 1` and is uploaded verbatim to constant/global
50    /// memory before kernel launch.
51    pub fn coefficients(self, lmax: usize, m: usize) -> Vec<f64> {
52        match self {
53            SphereSpectralKernelKind::Sobolev => {
54                crate::basis::sobolev_s2_truncated_coefficients(lmax, m)
55            }
56            SphereSpectralKernelKind::Pseudo => {
57                crate::basis::pseudo_s2_truncated_coefficients(lmax, m)
58            }
59        }
60    }
61
62    /// Stable string tag used in the NVRTC module cache key + logs.
63    pub const fn tag(self) -> &'static str {
64        match self {
65            SphereSpectralKernelKind::Sobolev => "sobolev",
66            SphereSpectralKernelKind::Pseudo => "pseudo",
67        }
68    }
69}
70
71/// Layout of the (n,m) kernel design matrix on device. The Wahba
72/// pipeline downstream of this kernel (cuBLAS GEMM, cuSOLVER GEQRF)
73/// requires column-major.
74#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
75pub enum DeviceMatrixLayout {
76    ColumnMajor,
77}
78
79/// Lat/lon (degrees or radians) → unit vector `(x, y, z)` on S² ⊂ ℝ³.
80/// Returns a flat `Vec<f64>` of length `3 * n` in the row-major layout
81/// `[x_0, y_0, z_0, x_1, y_1, z_1, …]`, ready for one `htod` upload.
82///
83/// `radians = false` interprets inputs as degrees (the codebase default
84/// for `SphericalSplineBasisSpec`).
85pub fn latlon_to_xyz_host(latlon: ArrayView2<'_, f64>, radians: bool) -> Result<Vec<f64>, String> {
86    if latlon.ncols() != 2 {
87        return Err(format!(
88            "latlon_to_xyz_host: expected (_, 2) lat/lon matrix, got shape {:?}",
89            latlon.shape()
90        ));
91    }
92    let deg = if radians {
93        1.0
94    } else {
95        std::f64::consts::PI / 180.0
96    };
97    let n = latlon.nrows();
98    let mut out = Vec::with_capacity(3 * n);
99    for row in latlon.outer_iter() {
100        let lat = row[0] * deg;
101        let lon = row[1] * deg;
102        let (s_lat, c_lat) = lat.sin_cos();
103        let (s_lon, c_lon) = lon.sin_cos();
104        // Standard geographic→cartesian: pole on +z.
105        out.push(c_lat * c_lon);
106        out.push(c_lat * s_lon);
107        out.push(s_lat);
108    }
109    Ok(out)
110}
111
112/// Device-resident `(rows × cols)` matrix in column-major layout with
113/// leading dimension `ld ≥ rows`. The slice holds `ld * cols` `f64`
114/// elements; entry `(i, j)` lives at `col_major_dev[j * ld + i]`.
115///
116/// On non-Linux builds the type is intentionally a host shadow so the
117/// surrounding orchestration compiles without cudarc.
118#[cfg(target_os = "linux")]
119pub struct DeviceS2KernelMatrix {
120    pub rows: usize,
121    pub cols: usize,
122    pub ld: usize,
123    pub col_major_dev: CudaSlice<f64>,
124    pub stream: Arc<CudaStream>,
125}
126
127#[cfg(not(target_os = "linux"))]
128pub struct DeviceS2KernelMatrix {
129    pub rows: usize,
130    pub cols: usize,
131    pub ld: usize,
132    /// Host shadow for CPU-only builds.
133    pub col_major_dev: Vec<f64>,
134}
135
136impl DeviceS2KernelMatrix {
137    /// Copy the device matrix back to the host as a regular ndarray
138    /// `(rows × cols)` row-major view. Convenience for tests + parity
139    /// comparisons; production paths should keep the matrix resident.
140    ///
141    /// The device matrix is `(ld × cols)` column-major; the host wants
142    /// `(rows × cols)` row-major. Two costs dominate this round-trip on the
143    /// real V100:
144    ///   1. the device→host copy of the full `ld·cols·8 B` payload, and
145    ///   2. the column-major→row-major transpose.
146    /// On Linux the dtoh is staged through a *cacheable* pinned host buffer
147    /// (see `PinnedF64`) so the DMA runs at full PCIe bandwidth (~10 GB/s)
148    /// instead of the ~1.3 GB/s the driver achieves staging a pageable
149    /// destination, and the subsequent host reads during the transpose hit
150    /// L1/L2 normally (unlike write-combined pinned memory). The transpose
151    /// itself is the parallel cache-blocked `col_major_to_row_major_parallel`.
152    #[cfg(target_os = "linux")]
153    pub fn to_host_array(&self) -> Result<Array2<f64>, GpuError> {
154        let needed = self.ld * self.cols;
155        let mut staging = PinnedLease::acquire(self.stream.context(), needed)?;
156        self.stream
157            .memcpy_dtoh(&self.col_major_dev, staging.as_mut_slice())
158            .gpu_ctx("DeviceS2KernelMatrix dtoh (pinned)")?;
159        self.stream
160            .synchronize()
161            .gpu_ctx("DeviceS2KernelMatrix synchronize (pinned)")?;
162        Ok(col_major_to_row_major_parallel(
163            staging.as_slice(),
164            self.rows,
165            self.cols,
166            self.ld,
167        ))
168    }
169
170    #[cfg(not(target_os = "linux"))]
171    pub fn to_host_array(&self) -> Result<Array2<f64>, GpuError> {
172        // Mirror the linux `to_host_array` exactly so both platforms return the
173        // identical row-major layout: pull the padded `(ld × cols)` column-major
174        // payload, then run the cache-blocked parallel transpose.
175        let mut col_major = vec![0.0_f64; self.ld * self.cols];
176        self.copy_to_host_col_major(&mut col_major)?;
177        Ok(col_major_to_row_major_parallel(
178            &col_major, self.rows, self.cols, self.ld,
179        ))
180    }
181
182    // Both variants below are reached only from `to_host_array` on the
183    // platform each is compiled for, so a Linux-built call graph saw the
184    // non-Linux one as unreachable and the sweep (d484a091a) removed both;
185    // the `x86_64-pc-windows-gnu` cross-check then failed at the call site
186    // above. A `cfg`-gated item is linked by the platforms the sweep does
187    // not build.
188    /// Copy the underlying `(ld × cols)` column-major payload to a
189    /// caller-provided buffer. Used by `to_host_array` and by the
190    /// device-resident cuSOLVER consumer when it needs to extract the
191    /// coefficient vector.
192    #[cfg(target_os = "linux")]
193    pub fn copy_to_host_col_major(&self, dst: &mut [f64]) -> Result<(), GpuError> {
194        let needed = self.ld * self.cols;
195        if dst.len() != needed {
196            gam_gpu::gpu_bail!(
197                "DeviceS2KernelMatrix::copy_to_host_col_major: dst.len()={} expected {}",
198                dst.len(),
199                needed
200            );
201        }
202        self.stream
203            .memcpy_dtoh(&self.col_major_dev, dst)
204            .gpu_ctx("DeviceS2KernelMatrix dtoh")?;
205        self.stream
206            .synchronize()
207            .gpu_ctx("DeviceS2KernelMatrix synchronize")?;
208        Ok(())
209    }
210
211    #[cfg(not(target_os = "linux"))]
212    pub fn copy_to_host_col_major(&self, dst: &mut [f64]) -> Result<(), GpuError> {
213        let needed = self.ld * self.cols;
214        if dst.len() != needed {
215            gam_gpu::gpu_bail!(
216                "DeviceS2KernelMatrix::copy_to_host_col_major: dst.len()={} expected {}",
217                dst.len(),
218                needed
219            );
220        }
221        dst.copy_from_slice(&self.col_major_dev);
222        Ok(())
223    }
224
225}
226
227/// Convert a `(ld × cols)` column-major device payload into a row-major
228/// `(rows × cols)` host `Array2`, in parallel with a cache-blocked tiled
229/// transpose.
230///
231/// Entry `(i, j)` lives at `col_major[j * ld + i]` and must land at
232/// `out[i * cols + j]`. A naive scalar `out[(i, j)] = col_major[j*ld+i]`
233/// loop over an `n·m` design (e.g. 200_000 × 200 ⇒ 320 MB) is utterly
234/// cache-hostile — the read stride is `ld` doubles — and measured at ~9 s,
235/// which alone made the GPU path lose to CPU. Here we:
236///   * tile the output rows into blocks small enough that one block's
237///     output stays L2-resident (`BLOCK_ROWS` rows × `cols` doubles),
238///   * read each source column slice contiguously (`col_major[j*ld+r0..]`),
239///   * run the row-blocks across the rayon pool.
240/// Reads are fully sequential per column; writes are bounded to the hot
241/// block. This drops the transpose from seconds to tens of milliseconds.
242fn col_major_to_row_major_parallel(
243    col_major: &[f64],
244    rows: usize,
245    cols: usize,
246    ld: usize,
247) -> Array2<f64> {
248    use rayon::prelude::*;
249
250    assert!(ld >= rows, "ld {ld} must be >= rows {rows}");
251    assert!(
252        col_major.len() >= ld * cols,
253        "col_major len {} < ld*cols {}",
254        col_major.len(),
255        ld * cols
256    );
257
258    // Block size chosen so one output block (BLOCK_ROWS × cols × 8 B) plus the
259    // source column slices stay roughly within L2 for the common `cols ≲ 200`.
260    const BLOCK_ROWS: usize = 128;
261
262    let mut out_flat = vec![0.0_f64; rows * cols];
263    out_flat
264        .par_chunks_mut(BLOCK_ROWS * cols)
265        .enumerate()
266        .for_each(|(block_idx, out_block)| {
267            let r0 = block_idx * BLOCK_ROWS;
268            let block_rows = out_block.len() / cols;
269            for j in 0..cols {
270                let base = j * ld + r0;
271                let src_col = &col_major[base..base + block_rows];
272                // Strided write within the hot block; contiguous column read.
273                for (local_i, &v) in src_col.iter().enumerate() {
274                    out_block[local_i * cols + j] = v;
275                }
276            }
277        });
278
279    Array2::from_shape_vec((rows, cols), out_flat).expect("row-major buffer has rows*cols elements")
280}
281
282/// RAII handle for a *cacheable* page-locked (pinned) host `f64` buffer.
283///
284/// cudarc's `CudaContext::alloc_pinned` always passes
285/// `CU_MEMHOSTALLOC_WRITECOMBINED`, which is excellent for host→device
286/// uploads but pathological for the host *reads* the transpose performs
287/// (write-combined memory is uncached on the CPU side). For the device→host
288/// return path we instead allocate plain pinned memory (`flags = 0`) directly
289/// via the driver: pinned so the dtoh DMA runs at full PCIe bandwidth, and
290/// cacheable so the parallel transpose can read it through the normal cache
291/// hierarchy. The buffer is freed with `cuMemFreeHost` on drop.
292#[cfg(target_os = "linux")]
293struct PinnedF64 {
294    ptr: *mut f64,
295    len: usize,
296    freed: bool,
297}
298
299#[cfg(target_os = "linux")]
300impl PinnedF64 {
301    /// Allocate `len` cacheable pinned `f64`s. Binds the context to the
302    /// calling thread first (required before any driver allocation call).
303    fn alloc(ctx: &Arc<CudaContext>, len: usize) -> Result<Self, GpuError> {
304        ctx.bind_to_thread().gpu_ctx("PinnedF64 bind_to_thread")?;
305        let bytes = len
306            .checked_mul(std::mem::size_of::<f64>())
307            .ok_or_else(|| gam_gpu::gpu_err!("PinnedF64: len={len} byte size overflows usize"))?;
308        // flags = 0 ⇒ cacheable pinned (NOT write-combined): fast DMA *and*
309        // fast host reads for the subsequent transpose.
310        // SAFETY: `bytes` is a valid non-overflowing size; the returned host
311        // pointer is owned by this struct and freed exactly once in `drop`.
312        let raw = unsafe { cudarc::driver::result::malloc_host(bytes, 0) }
313            .gpu_ctx("PinnedF64 cuMemHostAlloc")?;
314        let ptr = raw as *mut f64;
315        if ptr.is_null() {
316            gam_gpu::gpu_bail!("PinnedF64: cuMemHostAlloc returned null for {bytes} bytes");
317        }
318        Ok(Self {
319            ptr,
320            len,
321            freed: false,
322        })
323    }
324
325    fn as_mut_slice(&mut self) -> &mut [f64] {
326        // SAFETY: `ptr` points to `len` f64s of live pinned memory owned by
327        // self; the borrow is bounded by `&mut self`.
328        unsafe { std::slice::from_raw_parts_mut(self.ptr, self.len) }
329    }
330
331    fn as_slice(&self) -> &[f64] {
332        // SAFETY: as above; shared borrow bounded by `&self`.
333        unsafe { std::slice::from_raw_parts(self.ptr, self.len) }
334    }
335}
336
337#[cfg(target_os = "linux")]
338impl Drop for PinnedF64 {
339    fn drop(&mut self) {
340        if self.freed {
341            return;
342        }
343        self.freed = true;
344        // SAFETY: `ptr` was returned by `cuMemHostAlloc` in `alloc` and is
345        // freed exactly once (guarded by `freed`). A free failure during Drop
346        // is unrecoverable here; absorb it (the host process is tearing the
347        // allocation down regardless) without unwinding out of Drop.
348        if let Err(err) =
349            unsafe { cudarc::driver::result::free_host(self.ptr as *mut std::ffi::c_void) }
350        {
351            log::debug!(
352                "PinnedF64::drop: cuMemFreeHost failed ({err}); the pinned host allocation \
353                 is leaked for the remaining process lifetime"
354            );
355        }
356    }
357}
358
359// SAFETY: `PinnedF64` owns a single raw host allocation. The pointer is only
360// dereferenced by the thread holding the (mutable or shared) borrow; the pool
361// below moves the *handle* between threads while no borrow is outstanding, and
362// the rayon transpose only ever sees a `&[f64]` (already `Send + Sync`). The
363// raw pointer itself is never shared concurrently.
364#[cfg(target_os = "linux")]
365unsafe impl Send for PinnedF64 {}
366
367/// Bounded free-list of cacheable pinned host buffers, keyed by length.
368///
369/// Page-locking 320 MB via `cuMemHostAlloc` costs ~140 ms on the V100 — far
370/// more than the dtoh (~25 ms) it accelerates. During a REML fit the sphere
371/// design matrix is rebuilt and copied back at the *same* `(ld·cols)` size on
372/// every outer iteration, so caching the page-locked buffer turns that 140 ms
373/// into a one-time cost. The pool keeps at most [`PINNED_POOL_MAX_BUFFERS`]
374/// buffers (LRU-ish: oldest dropped first) to bound resident pinned memory.
375#[cfg(target_os = "linux")]
376const PINNED_POOL_MAX_BUFFERS: usize = 4;
377
378#[cfg(target_os = "linux")]
379static PINNED_POOL: OnceLock<Mutex<Vec<PinnedF64>>> = OnceLock::new();
380
381/// RAII lease of a pooled pinned buffer. Returns the buffer to [`PINNED_POOL`]
382/// on drop instead of freeing it, so the next same-size request reuses the
383/// page-locked allocation.
384#[cfg(target_os = "linux")]
385struct PinnedLease {
386    buf: Option<PinnedF64>,
387}
388
389#[cfg(target_os = "linux")]
390impl PinnedLease {
391    /// Acquire a pinned buffer of at least `len` f64s, reusing a pooled one of
392    /// exactly `len` when available, else allocating fresh.
393    fn acquire(ctx: &Arc<CudaContext>, len: usize) -> Result<Self, GpuError> {
394        let pool = PINNED_POOL.get_or_init(|| Mutex::new(Vec::new()));
395        if let Ok(mut guard) = pool.lock() {
396            if let Some(pos) = guard.iter().position(|b| b.len == len) {
397                return Ok(Self {
398                    buf: Some(guard.swap_remove(pos)),
399                });
400            }
401        }
402        Ok(Self {
403            buf: Some(PinnedF64::alloc(ctx, len)?),
404        })
405    }
406
407    fn as_mut_slice(&mut self) -> &mut [f64] {
408        self.buf
409            .as_mut()
410            .expect("PinnedLease buffer present until drop")
411            .as_mut_slice()
412    }
413
414    fn as_slice(&self) -> &[f64] {
415        self.buf
416            .as_ref()
417            .expect("PinnedLease buffer present until drop")
418            .as_slice()
419    }
420}
421
422#[cfg(target_os = "linux")]
423impl Drop for PinnedLease {
424    fn drop(&mut self) {
425        let Some(buf) = self.buf.take() else {
426            return;
427        };
428        if let Some(pool) = PINNED_POOL.get() {
429            if let Ok(mut guard) = pool.lock() {
430                if guard.len() < PINNED_POOL_MAX_BUFFERS {
431                    guard.push(buf);
432                    return;
433                }
434                // Pool full: evict the oldest cached buffer to make room for
435                // this (most-recently-used) one, keeping resident pinned memory
436                // bounded while favouring the hot size.
437                guard.remove(0);
438                guard.push(buf);
439                return;
440            }
441        }
442        // No pool / poisoned lock: fall back to freeing via PinnedF64::drop.
443        drop(buf);
444    }
445}
446
447// ────────────────────────────────────────────────────────────────────────
448// Inputs
449// ────────────────────────────────────────────────────────────────────────
450
451/// Host-side inputs needed to launch `s2_wahba_legendre_colmajor`.
452///
453/// `data_xyz` and `centers_xyz` are flat row-major
454/// `[x_0, y_0, z_0, …]` length `3 * n` and `3 * m` respectively, pre-
455/// computed via [`latlon_to_xyz_host`]. `coeffs` has length `lmax + 1`,
456/// indexed as `coeffs[ℓ] = c_ℓ` with `c_0 = 0`.
457#[derive(Clone, Debug)]
458pub struct S2KernelBuildInputs<'a> {
459    pub n: usize,
460    pub m: usize,
461    pub lmax: usize,
462    pub data_xyz: &'a [f64],
463    pub centers_xyz: &'a [f64],
464    pub coeffs: &'a [f64],
465    pub kind: SphereSpectralKernelKind,
466    pub layout: DeviceMatrixLayout,
467}
468
469impl<'a> S2KernelBuildInputs<'a> {
470    fn validate(&self) -> Result<(), GpuError> {
471        if self.lmax == 0 {
472            return Err(GpuError::DriverCallFailed {
473                reason: "S2KernelBuildInputs: lmax must be >= 1".into(),
474            });
475        }
476        if self.data_xyz.len() != 3 * self.n {
477            gam_gpu::gpu_bail!(
478                "S2KernelBuildInputs: data_xyz.len()={} != 3*n={}",
479                self.data_xyz.len(),
480                3 * self.n
481            );
482        }
483        if self.centers_xyz.len() != 3 * self.m {
484            gam_gpu::gpu_bail!(
485                "S2KernelBuildInputs: centers_xyz.len()={} != 3*m={}",
486                self.centers_xyz.len(),
487                3 * self.m
488            );
489        }
490        if self.coeffs.len() != self.lmax + 1 {
491            gam_gpu::gpu_bail!(
492                "S2KernelBuildInputs: coeffs.len()={} != lmax+1={}",
493                self.coeffs.len(),
494                self.lmax + 1
495            );
496        }
497        if self.coeffs[0] != 0.0 {
498            return Err(GpuError::DriverCallFailed {
499                reason: "S2KernelBuildInputs: coeffs[0] must be 0 (mean-zero kernel)".into(),
500            });
501        }
502        Ok(())
503    }
504}
505
506// ────────────────────────────────────────────────────────────────────────
507// NVRTC kernel source — raw and Householder-fused variants.
508//
509// Both compile with `--std=c++17 --gpu-architecture=compute_${cc}` and
510// take LMAX as a compile-time `#define`. Block (32, 8, 1), shared-mem
511// tiles for one data row × 3 doubles per warp and one center × 3
512// doubles per warp.
513// ────────────────────────────────────────────────────────────────────────
514
515#[cfg(target_os = "linux")]
516const KERNEL_TEMPLATE: &str = r#"
517// LMAX is supplied by the host via a `#define LMAX ...` prepended to
518// this source before NVRTC compilation (see `SphereGpuBackend::module_for`).
519// Recover cos(gamma) from the two half-angle chord lengths instead of
520// x dot c. The dot product rounds 1 - O(gamma^2) to 1 near coincidence,
521// permanently destroying the separation before the spectral evaluator sees it.
522// Here u = |x-c|^2 / (|x-c|^2 + |x+c|^2) and
523// v = |x+c|^2 / (|x-c|^2 + |x+c|^2), so both singular ends are carried
524// without cancellation and exact coincidence gives u=0, v=1 by construction.
525__device__ __forceinline__
526double s2_chord_cos_gamma(
527    double xi,
528    double yi,
529    double zi,
530    double cxj,
531    double cyj,
532    double czj
533) {
534    const double dx = xi - cxj;
535    const double dy = yi - cyj;
536    const double dz = zi - czj;
537    const double sx = xi + cxj;
538    const double sy = yi + cyj;
539    const double sz = zi + czj;
540    const double chord_sq = fma(dx, dx, fma(dy, dy, dz * dz));
541    const double anti_chord_sq = fma(sx, sx, fma(sy, sy, sz * sz));
542    const double scale = chord_sq + anti_chord_sq;
543
544    double u = chord_sq / scale;
545    double v = anti_chord_sq / scale;
546    if (u > 1.0) u = 1.0;
547    if (u < 0.0) u = 0.0;
548    if (v > 1.0) v = 1.0;
549    if (v < 0.0) v = 0.0;
550
551    double cos_gamma = v - u;
552    if (cos_gamma >  1.0) cos_gamma =  1.0;
553    if (cos_gamma < -1.0) cos_gamma = -1.0;
554    return cos_gamma;
555}
556
557extern "C" __global__
558__launch_bounds__(256)
559void s2_wahba_legendre_colmajor(
560    const double* __restrict__ data_xyz,    // n × 3 (row-major flat)
561    const double* __restrict__ centers_xyz, // m × 3 (row-major flat)
562    const double* __restrict__ coeffs,      // length LMAX + 1, coeffs[0] = 0
563    int n,
564    int m,
565    long long ld,
566    double* __restrict__ out                // ld × m column-major
567) {
568    const int i = blockIdx.y * blockDim.y + threadIdx.y;
569    const int j = blockIdx.x * blockDim.x + threadIdx.x;
570    if (i >= n || j >= m) return;
571
572    // Load (x_i, y_i, z_i) and (cx_j, cy_j, cz_j) into registers.
573    const double xi = data_xyz[3 * i + 0];
574    const double yi = data_xyz[3 * i + 1];
575    const double zi = data_xyz[3 * i + 2];
576    const double cxj = centers_xyz[3 * j + 0];
577    const double cyj = centers_xyz[3 * j + 1];
578    const double czj = centers_xyz[3 * j + 2];
579
580    // Stable half-angle chord geometry; no near-coincident dot-product loss.
581    const double t = s2_chord_cos_gamma(xi, yi, zi, cxj, cyj, czj);
582
583    // Legendre 3-term recurrence in registers.
584    // P_0(t) = 1, P_1(t) = t.
585    double p_prev = 1.0;
586    double p_curr = t;
587    double acc    = coeffs[0] * p_prev + coeffs[1] * p_curr;
588
589    #pragma unroll 8
590    for (int ell = 1; ell < LMAX; ++ell) {
591        const double lf  = (double) ell;
592        const double inv = 1.0 / (lf + 1.0);
593        // p_{ell+1} = ((2ell+1) * t * p_curr - ell * p_prev) / (ell+1)
594        const double p_next =
595            fma((2.0 * lf + 1.0) * t, p_curr, -lf * p_prev) * inv;
596        acc = fma(coeffs[ell + 1], p_next, acc);
597        p_prev = p_curr;
598        p_curr = p_next;
599    }
600
601    out[(long long) j * ld + (long long) i] = acc;
602}
603
604// Fused Householder-constrained kernel (Phase 3). Z = I - beta · v · v^T,
605// the constrained design is X_s = B[:, 1..m] - beta * (B · v) · v[1..m]^T,
606// i.e. drop the first column after applying Z. Each thread computes one
607// row of B in registers (m kernel evaluations), forms d_i = B_row · v,
608// then emits X_s[i, j_out] = B_row[j_out + 1] - beta * d_i * v[j_out + 1]
609// for j_out in 0..m-1.
610//
611// Grid: 1D over rows (block_dim.x rows per block). Each thread iterates
612// over centers in an inner loop — register-bound by the per-row state
613// (xyz_i, p_prev, p_curr, acc, and a small per-center scratch).
614extern "C" __global__
615__launch_bounds__(128)
616void s2_wahba_householder_constrained_colmajor(
617    const double* __restrict__ data_xyz,    // n × 3
618    const double* __restrict__ centers_xyz, // m × 3
619    const double* __restrict__ coeffs,      // length LMAX + 1
620    const double* __restrict__ v,           // length m, Householder vector
621    double beta,
622    int n,
623    int m,
624    long long ld_out,
625    double* __restrict__ out                // ld_out × (m-1) column-major
626) {
627    const int i = blockIdx.x * blockDim.x + threadIdx.x;
628    if (i >= n) return;
629
630    const double xi = data_xyz[3 * i + 0];
631    const double yi = data_xyz[3 * i + 1];
632    const double zi = data_xyz[3 * i + 2];
633
634    // Pass 1: compute d_i = sum_j v[j] * B[i, j].
635    double d_i = 0.0;
636    for (int j = 0; j < m; ++j) {
637        const double cxj = centers_xyz[3 * j + 0];
638        const double cyj = centers_xyz[3 * j + 1];
639        const double czj = centers_xyz[3 * j + 2];
640        const double t = s2_chord_cos_gamma(xi, yi, zi, cxj, cyj, czj);
641
642        double p_prev = 1.0;
643        double p_curr = t;
644        double acc    = coeffs[0] * p_prev + coeffs[1] * p_curr;
645        #pragma unroll 8
646        for (int ell = 1; ell < LMAX; ++ell) {
647            const double lf  = (double) ell;
648            const double inv = 1.0 / (lf + 1.0);
649            const double p_next =
650                fma((2.0 * lf + 1.0) * t, p_curr, -lf * p_prev) * inv;
651            acc = fma(coeffs[ell + 1], p_next, acc);
652            p_prev = p_curr;
653            p_curr = p_next;
654        }
655        d_i = fma(v[j], acc, d_i);
656    }
657
658    // Pass 2: emit X_s[i, j_out] = B[i, j_out+1] - beta * d_i * v[j_out+1].
659    const double bd = beta * d_i;
660    for (int j_out = 0; j_out < m - 1; ++j_out) {
661        const int j = j_out + 1;
662        const double cxj = centers_xyz[3 * j + 0];
663        const double cyj = centers_xyz[3 * j + 1];
664        const double czj = centers_xyz[3 * j + 2];
665        const double t = s2_chord_cos_gamma(xi, yi, zi, cxj, cyj, czj);
666
667        double p_prev = 1.0;
668        double p_curr = t;
669        double acc    = coeffs[0] * p_prev + coeffs[1] * p_curr;
670        #pragma unroll 8
671        for (int ell = 1; ell < LMAX; ++ell) {
672            const double lf  = (double) ell;
673            const double inv = 1.0 / (lf + 1.0);
674            const double p_next =
675                fma((2.0 * lf + 1.0) * t, p_curr, -lf * p_prev) * inv;
676            acc = fma(coeffs[ell + 1], p_next, acc);
677            p_prev = p_curr;
678            p_curr = p_next;
679        }
680        const double xs = acc - bd * v[j];
681        out[(long long) j_out * ld_out + (long long) i] = xs;
682    }
683}
684"#;
685
686// ────────────────────────────────────────────────────────────────────────
687// Module cache key + per-process backend.
688// ────────────────────────────────────────────────────────────────────────
689
690/// Module cache key: every distinct `(CC, LMAX, kind, layout, kernel
691/// flavor)` compiles to a different PTX. `precision = f64` and the
692/// (32, 8, 1) raw-kernel block / (128, 1, 1) Householder-kernel block
693/// shapes are baked into the kernel source so they are implicit in the
694/// flavor tag and don't appear here.
695#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
696pub struct S2ModuleCacheKey {
697    pub cc_major: i32,
698    pub cc_minor: i32,
699    pub lmax: u32,
700    pub kind: SphereSpectralKernelKind,
701    pub layout: DeviceMatrixLayout,
702}
703
704/// Returns `true` if this build was compiled with the Linux + cudarc GPU
705/// backend that runs the S² Wahba kernels.
706pub const fn sphere_gpu_compiled() -> bool {
707    cfg!(target_os = "linux")
708}
709
710/// Decide whether the GPU sphere kernel matrix path is eligible for
711/// `(n, m, lmax)`. Heuristic per the math spec:
712///   * `n * m >= 1_000_000`
713///   * `lmax <= 200`
714///   * device memory budget admits at least one `(ld × m)` design at
715///     `ld = ((n + 31) / 32) * 32`.
716#[must_use]
717pub fn sphere_kernel_decision(n: usize, m: usize, lmax: usize) -> Result<GpuDecision, GpuError> {
718    let large_enough = match gam_gpu::device_runtime::GpuRuntime::resolve(gam_gpu::global_policy())?
719    {
720        Some(runtime) => {
721            let ld = ((n + 31) / 32) * 32;
722            let needed_bytes = ld
723                .saturating_mul(m)
724                .saturating_mul(std::mem::size_of::<f64>());
725            let budget = runtime.memory_budget_bytes;
726            n.saturating_mul(m) >= 1_000_000 && lmax <= 200 && needed_bytes <= budget
727        }
728        None => false,
729    };
730    decide(
731        GpuKernel::SpatialKernelOperator,
732        gam_gpu::GpuEligibility::from_flags(sphere_gpu_compiled(), large_enough),
733    )
734}
735
736/// Map a truncated `SphereWahbaKernel` variant onto the device kernel kind +
737/// truncation degree. Only the two *truncated* spectral variants have an exact
738/// device counterpart (the closed-form `Sobolev`/`Pseudo` variants use
739/// polylogarithms / deep-`L` series the device kernel does not evaluate), so
740/// `Sobolev`/`Pseudo` return `None` and stay on the CPU closed-form path.
741#[must_use]
742pub fn truncated_device_kind(
743    kernel: crate::basis::SphereWahbaKernel,
744) -> Option<(SphereSpectralKernelKind, u16)> {
745    use crate::basis::SphereWahbaKernel;
746    match kernel {
747        SphereWahbaKernel::SobolevTruncated { lmax } => {
748            Some((SphereSpectralKernelKind::Sobolev, lmax))
749        }
750        SphereWahbaKernel::PseudoTruncated { lmax } => {
751            Some((SphereSpectralKernelKind::Pseudo, lmax))
752        }
753        SphereWahbaKernel::Sobolev | SphereWahbaKernel::Pseudo => None,
754    }
755}
756
757/// Production entry: build the raw `(n × m)` truncated-spectral Wahba kernel
758/// design matrix on the GPU when [`sphere_kernel_decision`] admits the device,
759/// returning `None` to signal the caller to use its CPU oracle.
760///
761/// Contract:
762///   * Returns `None` when the kernel is a non-truncated closed-form variant
763///     (no exact device counterpart), or when the dispatch decision keeps the
764///     work on the CPU (`!use_gpu`). The caller then runs the bit-defining CPU
765///     path. This is the **only** quiet-CPU route and it is taken *before* any
766///     device call — never as a silent fallback after a device failure.
767///   * Returns `Some(Ok(matrix))` with the device-computed host array when the
768///     device path ran and matches the CPU truncated recurrence to roundoff
769///     (proven by the parity tests). `gam_gpu::policy` keeps the same `c_ℓ`
770///     array and the same Legendre 3-term recurrence on both sides.
771///   * Returns `Some(Err(_))` when the device was *admitted* but the launch /
772///     NVRTC compile / copy-back failed — a hard error the caller must surface,
773///     NOT degrade to CPU. Fail-loud once admitted (the recurring silent-CPU
774///     fallback is the bug this path exists to kill).
775///
776/// `data` / `centers` are `(_, 2)` lat/lon matrices (degrees unless
777/// `radians`), matching `spherical_wahba_kernel_matrix_with_kind`.
778pub fn try_build_truncated_kernel_matrix_gpu(
779    data: ArrayView2<'_, f64>,
780    centers: ArrayView2<'_, f64>,
781    penalty_order: usize,
782    radians: bool,
783    kernel: crate::basis::SphereWahbaKernel,
784) -> Option<Result<Array2<f64>, GpuError>> {
785    let (kind, lmax) = truncated_device_kind(kernel)?;
786    let n = data.nrows();
787    let m = centers.nrows();
788    if n == 0 || m == 0 || lmax == 0 {
789        return None;
790    }
791    let decision = match sphere_kernel_decision(n, m, lmax as usize) {
792        Ok(decision) => decision,
793        Err(error) => return Some(Err(error)),
794    };
795    if !decision.use_gpu {
796        // Either backend-not-compiled, runtime-unavailable, or below the
797        // device-work threshold. Quiet CPU route, taken before any device call.
798        return None;
799    }
800    // Admitted: from here a failure is a hard error, never a silent CPU degrade.
801    Some(build_truncated_kernel_matrix_gpu_admitted(
802        data,
803        centers,
804        penalty_order,
805        radians,
806        kind,
807        lmax,
808    ))
809}
810
811/// Run the admitted device build for `try_build_truncated_kernel_matrix_gpu`.
812/// Separated so the admission decision (which returns `None` for the CPU route)
813/// stays distinct from the fail-loud device execution (which returns `Err`).
814fn build_truncated_kernel_matrix_gpu_admitted(
815    data: ArrayView2<'_, f64>,
816    centers: ArrayView2<'_, f64>,
817    penalty_order: usize,
818    radians: bool,
819    kind: SphereSpectralKernelKind,
820    lmax: u16,
821) -> Result<Array2<f64>, GpuError> {
822    let n = data.nrows();
823    let m = centers.nrows();
824    let data_xyz = latlon_to_xyz_host(data, radians)
825        .map_err(|reason| GpuError::DriverCallFailed { reason })?;
826    let centers_xyz = latlon_to_xyz_host(centers, radians)
827        .map_err(|reason| GpuError::DriverCallFailed { reason })?;
828    // Single-source the coefficients: the same `c_ℓ` array the CPU truncated
829    // recurrence consumes (`wahba_sphere_kernel_from_cos_kind`) is uploaded to
830    // the device, so CPU and GPU evaluate an identical zonal series.
831    let coeffs = kind.coefficients(lmax as usize, penalty_order);
832    let inputs = S2KernelBuildInputs {
833        n,
834        m,
835        lmax: lmax as usize,
836        data_xyz: &data_xyz,
837        centers_xyz: &centers_xyz,
838        coeffs: &coeffs,
839        kind,
840        layout: DeviceMatrixLayout::ColumnMajor,
841    };
842    let device_matrix = build_kernel_matrix_device(inputs)?;
843    let out = device_matrix.to_host_array()?;
844    // Guard against a device kernel that emitted NaN/Inf. A whole-matrix sum is
845    // poisoned by any non-finite element (`NaN + x = NaN`, `±Inf + finite =
846    // ±Inf`) and folds the `(n × m)` matrix in a single auto-vectorisable pass,
847    // ~7× faster than a per-element `any(!is_finite)` in the unoptimised
848    // profile (at n=200000, m=200 that scan alone was ~1.8 s — far more than
849    // the entire on-device build). The Wahba zonal kernel is a truncated
850    // Legendre series `Σ c_ℓ P_ℓ(t)` with `|P_ℓ| ≤ 1` and absolutely-summable
851    // coefficients, so every entry is O(1) and the sum of `n·m ≲ 10^8` of them
852    // cannot overflow f64 — a non-finite sum therefore means a genuinely
853    // non-finite entry, never a spurious overflow.
854    if !out.sum().is_finite() {
855        return Err(GpuError::DriverCallFailed {
856            reason: "sphere GPU truncated kernel produced a non-finite value".to_string(),
857        });
858    }
859    Ok(out)
860}
861
862#[cfg(target_os = "linux")]
863struct SphereGpuContext {
864    ctx: Arc<CudaContext>,
865    stream: Arc<CudaStream>,
866    modules: Mutex<HashMap<S2ModuleCacheKey, Arc<CudaModule>>>,
867    cc_major: i32,
868    cc_minor: i32,
869}
870
871/// Process-wide sphere GPU backend. Lazy-initialised on first call to
872/// [`SphereGpuBackend::probe`].
873pub struct SphereGpuBackend {
874    #[cfg(target_os = "linux")]
875    inner: SphereGpuContext,
876}
877
878impl SphereGpuBackend {
879    /// Lazily initialise the process-wide sphere backend.
880    pub fn probe() -> Result<&'static Self, GpuError> {
881        static BACKEND: OnceLock<Result<SphereGpuBackend, GpuError>> = OnceLock::new();
882        BACKEND
883            .get_or_init(|| {
884                #[cfg(target_os = "linux")]
885                {
886                    Self::probe_linux()
887                }
888                #[cfg(not(target_os = "linux"))]
889                {
890                    Err(GpuError::DriverLibraryUnavailable {
891                        reason: "sphere GPU backend is Linux-only".to_string(),
892                    })
893                }
894            })
895            .as_ref()
896            .map_err(GpuError::clone)
897    }
898
899    #[cfg(target_os = "linux")]
900    fn probe_linux() -> Result<Self, GpuError> {
901        let parts = gam_gpu::backend_probe::probe_cuda_backend("sphere")?;
902        Ok(SphereGpuBackend {
903            inner: SphereGpuContext {
904                ctx: parts.ctx,
905                stream: parts.stream,
906                modules: Mutex::new(HashMap::new()),
907                cc_major: parts.capability.compute_major,
908                cc_minor: parts.capability.compute_minor,
909            },
910        })
911    }
912
913    /// NVRTC-compile (or fetch from cache) the module for `key`. The
914    /// returned module exposes both raw and Householder-fused kernels.
915    #[cfg(target_os = "linux")]
916    fn module_for(&self, key: S2ModuleCacheKey) -> Result<Arc<CudaModule>, GpuError> {
917        if let Ok(guard) = self.inner.modules.lock() {
918            if let Some(existing) = guard.get(&key) {
919                return Ok(existing.clone());
920            }
921        }
922        // Prepend the `LMAX` macro directly to the source, then compile through
923        // the shared arch+fmad options (`compile_ptx_arch`). #1686's
924        // `--fmad=false` keeps the spherical-harmonic evaluation bit-comparable
925        // to the separately-rounded CPU reference; the #1551 arch pin keys the
926        // kernel to the device's real compute capability. (The arch is resolved
927        // internally via `nvrtc_arch()` from a `&'static str` table, so the old
928        // "cannot satisfy arch with a runtime string" limitation no longer
929        // applies — the LMAX specialization rides in the source, the arch in
930        // the options.)
931        let src = format!("#define LMAX {}\n{}", key.lmax, KERNEL_TEMPLATE);
932        let ptx = gam_gpu::device_cache::compile_ptx_arch(&src).gpu_ctx_with(|err| {
933            format!(
934                "sphere NVRTC compile (kind={}, lmax={}): {err}",
935                key.kind.tag(),
936                key.lmax
937            )
938        })?;
939        let module = self
940            .inner
941            .ctx
942            .load_module(ptx)
943            .gpu_ctx("sphere module load")?;
944        if let Ok(mut guard) = self.inner.modules.lock() {
945            guard.entry(key).or_insert_with(|| module.clone());
946        }
947        Ok(module)
948    }
949
950    #[cfg(target_os = "linux")]
951    fn cc(&self) -> (i32, i32) {
952        (self.inner.cc_major, self.inner.cc_minor)
953    }
954}
955
956// ────────────────────────────────────────────────────────────────────────
957// Entry points
958// ────────────────────────────────────────────────────────────────────────
959
960/// Build the raw `(n × m)` Wahba kernel matrix on device using
961/// `s2_wahba_legendre_colmajor`. Phase 1 entry point.
962pub fn build_kernel_matrix_device(
963    inputs: S2KernelBuildInputs<'_>,
964) -> Result<DeviceS2KernelMatrix, GpuError> {
965    inputs.validate()?;
966
967    #[cfg(target_os = "linux")]
968    {
969        use cudarc::driver::{LaunchConfig, PushKernelArg};
970        let backend = SphereGpuBackend::probe()?;
971        let (cc_major, cc_minor) = backend.cc();
972        let key = S2ModuleCacheKey {
973            cc_major,
974            cc_minor,
975            lmax: inputs.lmax as u32,
976            kind: inputs.kind,
977            layout: inputs.layout,
978        };
979        let module = backend.module_for(key)?;
980        let func = module
981            .load_function("s2_wahba_legendre_colmajor")
982            .gpu_ctx("sphere load_function raw")?;
983        let stream = backend.inner.stream.clone();
984
985        let data_dev = stream
986            .clone_htod(inputs.data_xyz)
987            .gpu_ctx("sphere htod data_xyz")?;
988        let centers_dev = stream
989            .clone_htod(inputs.centers_xyz)
990            .gpu_ctx("sphere htod centers_xyz")?;
991        let coeffs_dev = stream
992            .clone_htod(inputs.coeffs)
993            .gpu_ctx("sphere htod coeffs")?;
994
995        let n = inputs.n;
996        let m = inputs.m;
997        let ld = ((n + 31) / 32) * 32;
998        let mut out_dev = stream
999            .alloc_zeros::<f64>(ld * m)
1000            .gpu_ctx_with(|err| format!("sphere alloc out (ld={ld}, m={m}): {err}"))?;
1001
1002        // Block (32, 8, 1) — x over centers, y over rows.
1003        let block_x: u32 = 32;
1004        let block_y: u32 = 8;
1005        let grid_x: u32 = ((m as u32) + block_x - 1) / block_x;
1006        let grid_y: u32 = ((n as u32) + block_y - 1) / block_y;
1007        let cfg = LaunchConfig {
1008            grid_dim: (grid_x, grid_y, 1),
1009            block_dim: (block_x, block_y, 1),
1010            shared_mem_bytes: 0,
1011        };
1012        let n_i32: i32 =
1013            i32::try_from(n).map_err(|_| gam_gpu::gpu_err!("sphere n={n} overflows i32"))?;
1014        let m_i32: i32 =
1015            i32::try_from(m).map_err(|_| gam_gpu::gpu_err!("sphere m={m} overflows i32"))?;
1016        let ld_i64: i64 = ld as i64;
1017
1018        let mut builder = stream.launch_builder(&func);
1019        builder
1020            .arg(&data_dev)
1021            .arg(&centers_dev)
1022            .arg(&coeffs_dev)
1023            .arg(&n_i32)
1024            .arg(&m_i32)
1025            .arg(&ld_i64)
1026            .arg(&mut out_dev);
1027        // SAFETY: launch parameters are validated above; all device
1028        // pointers come from cudarc-checked allocations on the same
1029        // stream; the kernel only reads inputs and writes within
1030        // out[0 .. ld*m].
1031        unsafe { builder.launch(cfg) }.gpu_ctx("sphere raw kernel launch")?;
1032        stream
1033            .synchronize()
1034            .gpu_ctx("sphere raw kernel synchronize")?;
1035
1036        Ok(DeviceS2KernelMatrix {
1037            rows: n,
1038            cols: m,
1039            ld,
1040            col_major_dev: out_dev,
1041            stream,
1042        })
1043    }
1044
1045    #[cfg(not(target_os = "linux"))]
1046    {
1047        Err(GpuError::DriverLibraryUnavailable {
1048            reason: "sphere GPU backend is Linux-only".to_string(),
1049        })
1050    }
1051}
1052
1053// ────────────────────────────────────────────────────────────────────────
1054// Householder reflector helpers (host-side; Phase 3 prep).
1055//
1056// Given a non-zero weight vector w ∈ ℝ^m, construct (v, beta) such that
1057// H = I − beta · v · v^T satisfies H · w = ±‖w‖ · e_1 and drops the
1058// weighted-sum constraint into the first column.
1059// ────────────────────────────────────────────────────────────────────────
1060
1061// ────────────────────────────────────────────────────────────────────────
1062// Phase 2 — center-center penalty C + constraint S = Zᵀ C Z.
1063//
1064// `C` is the (m × m) Wahba kernel of centers against themselves and is
1065// computed by reusing the raw GPU kernel with `n = m`. The constraint
1066// transform is the same Householder reflector used by the Phase-3 fused
1067// kernel: Z = (I − β · v · vᵀ) with the first column dropped, so the
1068// constrained penalty is the trailing (m−1)×(m−1) block of HᵀCH.
1069//
1070// At m ≤ 200 the Householder product is cheap on host and the result is
1071// returned as an `ndarray::Array2`. Future calls into cuSOLVER QR can
1072// upload it (or its Cholesky factor) once and keep it device-resident.
1073// ────────────────────────────────────────────────────────────────────────
1074
1075// ────────────────────────────────────────────────────────────────────────
1076// Phase 4 — device-resident cuSOLVER QR penalised solve.
1077//
1078// Solve  min_β  ‖ [√W · X_s] β − [√W · y] ‖² + λ ‖R_S · β‖²
1079//
1080// by stacking the augmented matrix
1081//
1082//     A_aug = [ √W · X_s ;   √λ · R_S ]    shape (n + p) × p,
1083//     b_aug = [ √W · y    ;   0       ]    length n + p,
1084//
1085// where p = m − 1, R_S is the upper-triangular Cholesky factor of the
1086// constrained penalty S = Zᵀ C Z, and (√W·X_s) is the design built by
1087// the fused Householder kernel scaled by sqrt-weights row-by-row on
1088// device. The pipeline is:
1089//
1090//     1. cusolverDnDgeqrf_bufferSize → workspace size.
1091//     2. cusolverDnDgeqrf(A_aug)     → A := [R upper-tri / V Householder]
1092//                                        plus tau vector.
1093//     3. cusolverDnDormqr(side=L, trans=T)
1094//                                  → applies Qᵀ to b_aug.
1095//     4. cublasDtrsm(L = upper) → β := R⁻¹ · (Qᵀ b_aug)[0..p].
1096//
1097// Coefficients (β) come back to host; log|H| can be returned via Σ
1098// log(R_ii²) from the diagonal of the in-place factored R.
1099//
1100// All intermediate state — A_aug, b_aug, tau, workspace, info — stays
1101// device-resident. The host learns only (β, log|H|, residual ssq).
1102// ────────────────────────────────────────────────────────────────────────
1103
1104/// Result returned by `solve_penalised_ls_device`.
1105#[derive(Clone, Debug)]
1106pub struct PenalisedLsSolution {
1107    /// Coefficient vector, length `p = m − 1` (after Householder drop).
1108    pub beta: Vec<f64>,
1109    /// Sum of squared residuals on the unaugmented rows: ‖√W (Xβ − y)‖².
1110    pub weighted_residual_ssq: f64,
1111    /// log|H| = 2 · Σ log |R_ii| of the QR-factored augmented design.
1112    pub log_det_hessian: f64,
1113}
1114
1115// ────────────────────────────────────────────────────────────────────────
1116// Tests
1117// ────────────────────────────────────────────────────────────────────────
1118
1119#[cfg(test)]
1120mod sphere_gpu_tests {
1121    use super::*;
1122    use crate::basis::sphere_half_angle::{SphereTrig, half_angle_separation_scalar};
1123    use crate::basis::{
1124        SphereWahbaKernel, sobolev_s2_truncated_coefficients, sphere_truncated_spectral_eval,
1125        spherical_wahba_kernel_matrix_with_kind,
1126    };
1127    use ndarray::Array2;
1128
1129    fn small_latlon_grid(n_lat: usize, n_lon: usize) -> Array2<f64> {
1130        // Latitude in (-85, 85), longitude in [-180, 180), degrees.
1131        let mut rows = Vec::with_capacity(n_lat * n_lon);
1132        for i in 0..n_lat {
1133            let lat = -85.0 + (170.0 * i as f64) / (n_lat.saturating_sub(1).max(1) as f64);
1134            for j in 0..n_lon {
1135                let lon = -180.0 + (360.0 * j as f64) / (n_lon.saturating_sub(1).max(1) as f64);
1136                rows.push(lat);
1137                rows.push(lon);
1138            }
1139        }
1140        Array2::from_shape_vec((n_lat * n_lon, 2), rows).unwrap()
1141    }
1142
1143    fn cuda_available_for_test(label: &str) -> bool {
1144        match gam_gpu::device_runtime::GpuRuntime::resolve(gam_gpu::GpuPolicy::Auto) {
1145            Ok(Some(_)) => true,
1146            Ok(None) => {
1147                eprintln!("[sphere_gpu test] no CUDA device — skipping {label}");
1148                false
1149            }
1150            Err(error) => panic!("[sphere_gpu test] CUDA resolution failed for {label}: {error}"),
1151        }
1152    }
1153
1154    /// #2424 device-free half: with no CUDA runtime the dispatch decision must
1155    /// DECLINE at `(n, m, lmax)`. Where `n·m` clears the device-work threshold
1156    /// this is a strictly device-dependent claim — only the missing runtime can
1157    /// hold the dispatch back — and below the threshold it additionally pins
1158    /// the size gate. Either way, admitting a device this host does not have is
1159    /// the #1551 silent-device class, and it is exactly what a
1160    /// `return`-before-the-first-assertion skip could never see.
1161    fn assert_sphere_decision_declines_without_device(n: usize, m: usize, lmax: usize) {
1162        let decision = sphere_kernel_decision(n, m, lmax)
1163            .expect("the sphere GPU decision must not fault on a device-free host");
1164        assert!(
1165            !decision.use_gpu,
1166            "no CUDA runtime on this host, yet the sphere dispatch decision admitted the \
1167             device for (n={n}, m={m}, lmax={lmax}) — reason={}",
1168            decision.reason
1169        );
1170    }
1171
1172    /// #2424 device-free half: the admitted-only device entries must REFUSE
1173    /// with an `Err` rather than fabricate a host-side answer. `build_*_device`
1174    /// is reached only after the decision admits the device, so on a host with
1175    /// no runtime every call owes an error — never `Ok`, never a panic.
1176    fn assert_device_kernel_entry_refuses(inputs: S2KernelBuildInputs<'_>) {
1177        assert!(
1178            build_kernel_matrix_device(inputs).is_err(),
1179            "no CUDA runtime on this host, yet the device kernel entry returned a matrix \
1180             — the admitted-only device path fabricated a host answer (#1551 class)"
1181        );
1182    }
1183
1184    /// #2424: the truncated-spectral kernel is defined elementwise as
1185    /// `K(x, c) = Σ_ℓ c_ℓ · P_ℓ(x·c)`. This grades the production CPU matrix
1186    /// against that definition evaluated point-by-point through the Legendre
1187    /// recurrence — the same definition the device kernel implements, so it
1188    /// pins the ORACLE the GPU is compared against, on every host.
1189    fn assert_cpu_kernel_matches_stable_spectral_definition(
1190        kernel_matrix: &Array2<f64>,
1191        data_latlon: &Array2<f64>,
1192        centers_latlon: &Array2<f64>,
1193        coeffs: &[f64],
1194    ) {
1195        let (n, m) = kernel_matrix.dim();
1196        let to_radians = std::f64::consts::PI / 180.0;
1197        let mut max_abs = 0.0_f64;
1198        for i in 0..n {
1199            let point = SphereTrig::from_radians(
1200                data_latlon[(i, 0)] * to_radians,
1201                data_latlon[(i, 1)] * to_radians,
1202            );
1203            for j in 0..m {
1204                let center = SphereTrig::from_radians(
1205                    centers_latlon[(j, 0)] * to_radians,
1206                    centers_latlon[(j, 1)] * to_radians,
1207                );
1208                let separation = half_angle_separation_scalar(point, center);
1209                let expected = sphere_truncated_spectral_eval(separation.cos_gamma(), coeffs);
1210                max_abs = max_abs.max((kernel_matrix[(i, j)] - expected).abs());
1211            }
1212        }
1213        assert!(
1214            max_abs < 1e-12,
1215            "CPU truncated-spectral kernel matrix departs from the stable half-angle \
1216             elementwise definition: max |delta| = {max_abs:.3e}"
1217        );
1218    }
1219
1220    #[test]
1221    fn sum_finite_guard_accepts_finite_rejects_nonfinite() {
1222        // The admitted device path guards its output with `!out.sum().is_finite()`
1223        // instead of a per-element `any(!is_finite)`. This pins the equivalence
1224        // that justifies the swap: a finite matrix has a finite sum, and a single
1225        // NaN or ±Inf entry poisons the sum.
1226        let finite = Array2::<f64>::from_shape_fn((5, 7), |(i, j)| (i as f64 - 2.0) * (j as f64));
1227        assert!(finite.sum().is_finite());
1228
1229        let mut with_nan = finite.clone();
1230        with_nan[[3, 4]] = f64::NAN;
1231        assert!(!with_nan.sum().is_finite());
1232
1233        let mut with_pos_inf = finite.clone();
1234        with_pos_inf[[0, 0]] = f64::INFINITY;
1235        assert!(!with_pos_inf.sum().is_finite());
1236
1237        let mut with_neg_inf = finite.clone();
1238        with_neg_inf[[4, 6]] = f64::NEG_INFINITY;
1239        assert!(!with_neg_inf.sum().is_finite());
1240    }
1241
1242    #[test]
1243    fn xyz_preprocessing_matches_unit_sphere() {
1244        let latlon = ndarray::array![
1245            [0.0, 0.0],
1246            [90.0, 0.0],
1247            [0.0, 90.0],
1248            [-90.0, 17.5],
1249            [45.0, -120.0],
1250        ];
1251        let xyz = latlon_to_xyz_host(latlon.view(), false).expect("xyz");
1252        assert_eq!(xyz.len(), 3 * 5);
1253        for i in 0..5 {
1254            let nrm2 = xyz[3 * i] * xyz[3 * i]
1255                + xyz[3 * i + 1] * xyz[3 * i + 1]
1256                + xyz[3 * i + 2] * xyz[3 * i + 2];
1257            assert!((nrm2 - 1.0).abs() < 1e-15, "row {i} not unit norm: {nrm2}");
1258        }
1259        // Row 0 = equator @ lon=0 → (1, 0, 0).
1260        assert!((xyz[0] - 1.0).abs() < 1e-15);
1261        assert!(xyz[1].abs() < 1e-15);
1262        assert!(xyz[2].abs() < 1e-15);
1263        // Row 1 = north pole (lat=90, lon=0) → (0, 0, 1).
1264        assert!(xyz[3].abs() < 1e-15);
1265        assert!(xyz[4].abs() < 1e-15);
1266        assert!((xyz[5] - 1.0).abs() < 1e-15);
1267        // Row 2 = equator @ lon=90 → (0, 1, 0).
1268        assert!(xyz[6].abs() < 1e-15);
1269        assert!((xyz[7] - 1.0).abs() < 1e-15);
1270        assert!(xyz[8].abs() < 1e-15);
1271    }
1272
1273    #[test]
1274    fn truncated_spectral_at_same_point_matches_sum_of_coefficients() {
1275        // P_ℓ(1) = 1 for all ℓ, so K(x, x) = Σ_{ℓ=0..L} c_ℓ. The Legendre
1276        // recurrence in `sphere_truncated_spectral_eval` must reproduce
1277        // this exact identity to roundoff.
1278        for m_penalty in 1..=4 {
1279            for &lmax in &[5_usize, 20, 50] {
1280                let coeffs = sobolev_s2_truncated_coefficients(lmax, m_penalty);
1281                let expected: f64 = coeffs.iter().sum();
1282                let got = sphere_truncated_spectral_eval(1.0, &coeffs);
1283                assert!(
1284                    (got - expected).abs() < 1e-13,
1285                    "K(x,x) identity broken at m={m_penalty}, L={lmax}: got {got:.6e}, expected {expected:.6e}"
1286                );
1287            }
1288        }
1289    }
1290
1291    #[test]
1292    fn truncated_spectral_at_antipode_matches_alternating_sum() {
1293        // P_ℓ(-1) = (-1)^ℓ, so K(x, -x) = Σ_{ℓ=0..L} c_ℓ · (-1)^ℓ. Same
1294        // exact identity for the recurrence at t = -1.
1295        for m_penalty in 1..=4 {
1296            for &lmax in &[5_usize, 20, 50] {
1297                let coeffs = sobolev_s2_truncated_coefficients(lmax, m_penalty);
1298                let expected: f64 = coeffs
1299                    .iter()
1300                    .enumerate()
1301                    .map(|(ell, c)| if ell % 2 == 0 { *c } else { -*c })
1302                    .sum();
1303                let got = sphere_truncated_spectral_eval(-1.0, &coeffs);
1304                assert!(
1305                    (got - expected).abs() < 1e-13,
1306                    "K(x,-x) identity broken at m={m_penalty}, L={lmax}: got {got:.6e}, expected {expected:.6e}"
1307                );
1308            }
1309        }
1310    }
1311
1312    #[test]
1313    fn truncated_spectral_matrix_is_symmetric() {
1314        // K(γ) depends only on cos γ = x · y = y · x, so the Gram
1315        // matrix B B^T-style kernel evaluation on the same point set
1316        // must be symmetric to roundoff.
1317        let centers = ndarray::array![
1318            [10.0_f64, 20.0],
1319            [-30.0, 100.0],
1320            [45.0, -60.0],
1321            [-89.0, 0.0],
1322            [0.0, 180.0],
1323            [60.0, -179.9],
1324        ];
1325        for m_penalty in [1usize, 2, 4] {
1326            for &lmax in &[10_usize, 30] {
1327                let mat = spherical_wahba_kernel_matrix_with_kind(
1328                    centers.view(),
1329                    centers.view(),
1330                    m_penalty,
1331                    false,
1332                    SphereWahbaKernel::SobolevTruncated { lmax: lmax as u16 },
1333                )
1334                .expect("kernel matrix");
1335                let n = centers.nrows();
1336                let mut max_asym = 0.0_f64;
1337                for i in 0..n {
1338                    for j in 0..n {
1339                        let d = (mat[(i, j)] - mat[(j, i)]).abs();
1340                        if d > max_asym {
1341                            max_asym = d;
1342                        }
1343                    }
1344                }
1345                assert!(
1346                    max_asym < 1e-13,
1347                    "K not symmetric at m={m_penalty}, L={lmax}: max |K - Kᵀ| = {max_asym:.3e}"
1348                );
1349            }
1350        }
1351    }
1352
1353    #[test]
1354    fn truncated_coefficients_have_zero_constant_mode() {
1355        for m in 1..=4 {
1356            let c = sobolev_s2_truncated_coefficients(50, m);
1357            assert_eq!(c.len(), 51);
1358            assert_eq!(c[0], 0.0);
1359            assert!(c[1] > 0.0);
1360            // Spectral decay c_ℓ ~ 1/ℓ^{2m-1}: monotone for ℓ ≥ 1.
1361            for ell in 2..=50 {
1362                assert!(
1363                    c[ell] < c[ell - 1] + 1e-15,
1364                    "Sobolev coefficient not non-increasing at m={m}, ell={ell}: {} vs {}",
1365                    c[ell],
1366                    c[ell - 1]
1367                );
1368            }
1369        }
1370    }
1371
1372    #[test]
1373    fn truncated_spectral_matches_matrix_helper() {
1374        // The Wahba kernel matrix helper, invoked with the truncated
1375        // variant, must produce the same value as the bare scalar
1376        // evaluator.
1377        let m_penalty = 2;
1378        let lmax = 20;
1379        let coeffs = sobolev_s2_truncated_coefficients(lmax, m_penalty);
1380        let data = ndarray::array![[12.5, -34.0]];
1381        let centers = ndarray::array![[40.0, 10.0]];
1382        let mat = spherical_wahba_kernel_matrix_with_kind(
1383            data.view(),
1384            centers.view(),
1385            m_penalty,
1386            false,
1387            SphereWahbaKernel::SobolevTruncated { lmax: lmax as u16 },
1388        )
1389        .expect("kernel matrix");
1390        // Recompute cos gamma through the stable half-angle geometry, not the
1391        // dot-product route whose near-coincident loss this contract must catch.
1392        let to_radians = std::f64::consts::PI / 180.0;
1393        let point = SphereTrig::from_radians(data[(0, 0)] * to_radians, data[(0, 1)] * to_radians);
1394        let center =
1395            SphereTrig::from_radians(centers[(0, 0)] * to_radians, centers[(0, 1)] * to_radians);
1396        let expected = sphere_truncated_spectral_eval(
1397            half_angle_separation_scalar(point, center).cos_gamma(),
1398            &coeffs,
1399        );
1400        assert!(
1401            (mat[(0, 0)] - expected).abs() < 1e-13,
1402            "matrix helper differs from scalar evaluator: {} vs {}",
1403            mat[(0, 0)],
1404            expected
1405        );
1406    }
1407
1408    /// Raw kernel parity vs the CPU truncated-spectral path. The device build
1409    /// is device-only, but the CPU oracle it is graded against owes its own
1410    /// elementwise definition on every host, and a device-free host owes the
1411    /// decline contract (#2424 — this test used to `return` before its first
1412    /// assertion and report a pass on every CI runner).
1413    #[test]
1414    fn sphere_gpu_raw_kernel_parity_vs_cpu_truncated() {
1415        let mut data_ll = small_latlon_grid(7, 9);
1416        let mut centers_ll = small_latlon_grid(5, 7);
1417        // Exercise the region where x dot c rounds away the separation. The
1418        // chord form still carries this distinct pair monotonically.
1419        centers_ll[(0, 0)] = 12.5;
1420        centers_ll[(0, 1)] = -34.0;
1421        data_ll[(0, 0)] = 12.5 + 1.0e-8;
1422        data_ll[(0, 1)] = -34.0;
1423        let data_xyz = latlon_to_xyz_host(data_ll.view(), false).unwrap();
1424        let centers_xyz = latlon_to_xyz_host(centers_ll.view(), false).unwrap();
1425        let n = data_ll.nrows();
1426        let m = centers_ll.nrows();
1427        let penalty = 2usize;
1428        let lmax = 20usize;
1429        let coeffs = sobolev_s2_truncated_coefficients(lmax, penalty);
1430
1431        let inputs = S2KernelBuildInputs {
1432            n,
1433            m,
1434            lmax,
1435            data_xyz: &data_xyz,
1436            centers_xyz: &centers_xyz,
1437            coeffs: &coeffs,
1438            kind: SphereSpectralKernelKind::Sobolev,
1439            layout: DeviceMatrixLayout::ColumnMajor,
1440        };
1441
1442        let cpu = spherical_wahba_kernel_matrix_with_kind(
1443            data_ll.view(),
1444            centers_ll.view(),
1445            penalty,
1446            false,
1447            SphereWahbaKernel::SobolevTruncated { lmax: lmax as u16 },
1448        )
1449        .expect("cpu kernel matrix");
1450
1451        // EVERY HOST: the oracle the device is graded against must itself
1452        // equal the elementwise truncated-spectral definition.
1453        assert_cpu_kernel_matches_stable_spectral_definition(&cpu, &data_ll, &centers_ll, &coeffs);
1454
1455        if !cuda_available_for_test("raw-kernel parity") {
1456            assert_sphere_decision_declines_without_device(n, m, lmax);
1457            assert_device_kernel_entry_refuses(inputs);
1458            return;
1459        }
1460        // Past the runtime Some-gate: a probe failure is a real device fault on a
1461        // CUDA host — fail loud (device-PCG skip-pass class, eee12f6b2).
1462        SphereGpuBackend::probe()
1463            .expect("[sphere_gpu test] backend probe must succeed on a CUDA host");
1464        let dev_mat = build_kernel_matrix_device(inputs).expect("device kernel matrix");
1465        let gpu = dev_mat.to_host_array().expect("dtoh kernel matrix");
1466
1467        let mut max_abs = 0.0_f64;
1468        for i in 0..n {
1469            for j in 0..m {
1470                let d = (gpu[(i, j)] - cpu[(i, j)]).abs();
1471                if d > max_abs {
1472                    max_abs = d;
1473                }
1474            }
1475        }
1476        assert!(
1477            max_abs < 1e-11,
1478            "GPU vs CPU truncated parity max |Δ| = {max_abs:.3e} >= 1e-11"
1479        );
1480    }
1481
1482    /// The end-to-end sphere build routes its kernel to the device exactly
1483    /// when the dispatch policy admits one, on both kinds of host.
1484    ///
1485    /// #2424: the device-free half asserts the decision declines at the
1486    /// end-to-end shape and stops before building the 200k-row fixture.
1487    ///
1488    /// #2372/#2420: this arm used to divide two `Instant::elapsed()` readings
1489    /// and demand `≥ 10×`, and it could never have reached that number on any
1490    /// hardware. The GPU side timed the whole `build_spherical_spline_basis`
1491    /// while the CPU side timed `spherical_wahba_kernel_matrix_cpu` plus one
1492    /// `dot`; the comment justifying that claimed farthest-point center
1493    /// selection was excluded because it "is identical for both paths", but it
1494    /// was excluded only from the CPU side. Measured on an A10 at this exact
1495    /// shape the fit paid `centers = 15.625 s` against a `0.243 s` device
1496    /// kernel and a `7.778 s` host kernel, so the ratio read `0.51×` while the
1497    /// device path itself was running `37×` faster than the host. Timing the
1498    /// same work on both sides puts the shared center selection in both
1499    /// numerator and denominator, which caps the ratio at
1500    /// `(centers + host_kernel + rest) / (centers + rest)` — **1.99×** even
1501    /// with a free kernel, and #2420's landed `55b4367e2` (centers 15.6 s →
1502    /// 6.0 s) lowers that ceiling rather than raising it.
1503    ///
1504    /// So the ratio is retired here rather than re-derived: a smaller constant
1505    /// would still be a stopwatch on a co-tenanted box (SPEC rule 19, and the
1506    /// #2487 precedent that replaced four such gates). What this workload
1507    /// actually needs asserted is that its shape *belongs* on the device, and
1508    /// the calibrated dispatch policy owns that decision — so the gate is the
1509    /// policy's own crossover, pinned by the pair straddling it. The device
1510    /// kernel's speed keeps its gate in
1511    /// `sphere_gpu_kernel_matrix_hill_climb_declines_without_device_else_20x_vs_cpu`,
1512    /// where both arms do time the same work; device/host agreement keeps its
1513    /// gates in the raw-kernel and end-to-end fit parity tests.
1514    #[test]
1515    fn sphere_gpu_end_to_end_fit_dispatches_to_device_else_declines() {
1516        use crate::basis::{CenterStrategy, SphereMethod, SphericalSplineBasisSpec, SphericalSplineIdentifiability, build_spherical_spline_basis};
1517
1518        let n_lat = 500usize;
1519        let n_lon = 400usize;
1520        let m: usize = 200;
1521        let lmax: u16 = 50;
1522
1523        if !cuda_available_for_test("end-to-end fit dispatch") {
1524            assert_sphere_decision_declines_without_device(n_lat * n_lon, m, lmax as usize);
1525            return;
1526        }
1527        // A CUDA runtime is present, so a probe failure is a real device/
1528        // dispatch fault — fail the gate loudly rather than skip-passing.
1529        SphereGpuBackend::probe()
1530            .expect("[sphere_gpu end-to-end dispatch] backend probe must succeed on a CUDA host");
1531
1532        // The policy's device-work crossover, pinned by the adjacent pair that
1533        // straddles it. Both shapes stage `≈ n·m·8 B = 8 MB`, three orders
1534        // under any plausible `memory_budget_bytes`, so neither side of the
1535        // pair can flip on a co-tenanted device — the only thing separating
1536        // them is the crossover itself. Without the negative arm the positive
1537        // one proves nothing: a predicate that admits everything would pass it.
1538        let admit = sphere_kernel_decision(5_000, m, lmax as usize)
1539            .expect("the sphere GPU decision must not fault on a CUDA host");
1540        assert!(
1541            admit.use_gpu,
1542            "a CUDA device is present and (n=5000, m={m}) is exactly at the sphere device-work \
1543             crossover, yet the dispatch decision kept it on the host — reason={}",
1544            admit.reason
1545        );
1546        let refuse = sphere_kernel_decision(4_999, m, lmax as usize)
1547            .expect("the sphere GPU decision must not fault on a CUDA host");
1548        assert!(
1549            !refuse.use_gpu,
1550            "(n=4999, m={m}) is one row below the sphere device-work crossover, yet the dispatch \
1551             decision admitted the device — reason={}",
1552            refuse.reason
1553        );
1554
1555        let data_ll = small_latlon_grid(n_lat, n_lon);
1556        let spec_gpu = SphericalSplineBasisSpec {
1557            center_strategy: CenterStrategy::FarthestPoint { num_centers: m },
1558            penalty_order: 2,
1559            double_penalty: false,
1560            radians: false,
1561            method: SphereMethod::Wahba,
1562            max_degree: None,
1563            wahba_kernel: SphereWahbaKernel::SobolevTruncated { lmax },
1564            identifiability: SphericalSplineIdentifiability::CenterSumToZero,
1565        };
1566
1567        // The decision above is about a shape; this is the production entry
1568        // actually taking it. `n·m = 4·10⁷` clears the crossover by 40×, so on
1569        // this host the build routes its kernel to the device — a path the
1570        // device-free half can never reach, and the reason this test still
1571        // pays for a 200k-row fixture.
1572        let t0 = std::time::Instant::now();
1573        let built = build_spherical_spline_basis(data_ll.view(), &spec_gpu)
1574            .expect("the end-to-end sphere build must succeed on a CUDA host");
1575        let build_secs = t0.elapsed().as_secs_f64();
1576
1577        assert_eq!(
1578            built.design.nrows(),
1579            data_ll.nrows(),
1580            "the device-dispatched sphere build returned {} design rows for {} data rows",
1581            built.design.nrows(),
1582            data_ll.nrows()
1583        );
1584        assert!(
1585            built.design.ncols() > 0 && built.design.ncols() <= m,
1586            "the device-dispatched sphere build returned {} design columns for m={m} centers",
1587            built.design.ncols()
1588        );
1589        // A device path that faulted into NaN would still return the right
1590        // shape; probing rows spread across the grid costs nothing next to the
1591        // build and is the cheapest thing that distinguishes the two.
1592        let beta = ndarray::Array1::<f64>::ones(built.design.ncols());
1593        for row in (0..built.design.nrows()).step_by(data_ll.nrows() / 64 + 1) {
1594            let value = built.design.dot_row(row, &beta);
1595            assert!(
1596                value.is_finite(),
1597                "the device-dispatched sphere design row {row} sums to {value}, not a finite number"
1598            );
1599        }
1600
1601        // Kept as the hill-climbing record this workload is worth, not as a
1602        // gate: `build_secs` is dominated by single-threaded farthest-point
1603        // center selection (#2420), so it tracks the host's core speed and the
1604        // box's load rather than the device kernel's.
1605        eprintln!(
1606            "[sphere_gpu end-to-end dispatch] n={} m={m} L={lmax} build={build_secs:.3}s reason={}",
1607            data_ll.nrows(),
1608            admit.reason
1609        );
1610    }
1611
1612    /// Task #25: end-to-end fit parity between the GPU truncated-spectral
1613    /// path and the CPU truncated-spectral path on a small synthetic
1614    /// intrinsic-S² fixture.
1615    ///
1616    /// Setup: deterministic lat/lon grid (n = 1000 = 25 × 40), 80 centers
1617    /// chosen by farthest-point selection, lmax = 15, penalty order 2,
1618    /// Wahba weighted-sum-to-zero constraint applied via `Z`. We fit a
1619    /// fixed-λ penalised LS problem
1620    ///   β = argmin ‖X_s β − y‖² + λ · βᵀ S β
1621    /// where `X_s = K(data, centers) · Z` and `S = Zᵀ · K(centers, centers) · Z`,
1622    /// solving `(X_sᵀ X_s + λ S) β = X_sᵀ y` via faer LLT for both paths.
1623    /// The only path-dependent quantity is `K(data, centers)`: built on
1624    /// GPU via `build_kernel_matrix_device` for one β, and on CPU via
1625    /// `spherical_wahba_kernel_matrix_with_kind` for the other. The
1626    /// penalty kernel `K(centers, centers)` is m × m and tiny, so we
1627    /// build it once on CPU and share it across paths (it is not the
1628    /// surface under test).
1629    ///
1630    /// Asserts max-absolute coefficient delta ≤ 1e-9 and max-absolute
1631    /// fitted-value delta ≤ 1e-9. `#[ignore = "requires CUDA"]` so the
1632    /// V100 bench runner unignores in their harness.
1633    #[test]
1634    fn sphere_gpu_end_to_end_fit_parity_vs_cpu_truncated() {
1635        use crate::basis::{
1636            select_spherical_farthest_point_centers, spherical_wahba_kernel_matrix_with_kind,
1637        };
1638        use faer::Side;
1639        use gam_linalg::faer_ndarray::FaerCholesky;
1640
1641        // Fixture: 25 × 40 lat/lon grid → n = 1000.
1642        let data_ll = small_latlon_grid(25, 40);
1643        assert_eq!(data_ll.nrows(), 1000);
1644        let n = data_ll.nrows();
1645        let m: usize = 80;
1646        let lmax_u16: u16 = 15;
1647        let lmax: usize = lmax_u16 as usize;
1648        let penalty_order: usize = 2;
1649        let kernel = SphereWahbaKernel::SobolevTruncated { lmax: lmax_u16 };
1650        let lambda: f64 = 1.0e-3;
1651
1652        // Deterministic centers via farthest-point selection.
1653        let centers_ll = select_spherical_farthest_point_centers(data_ll.view(), m, false)
1654            .expect("farthest-point centers");
1655        assert_eq!(centers_ll.nrows(), m);
1656
1657        // The Wahba sphere basis no longer imposes a finite-center coefficient
1658        // gauge; parity compares the raw center coefficient chart.
1659        let z = Array2::<f64>::eye(centers_ll.nrows());
1660        let p = z.ncols();
1661        assert_eq!(p, m);
1662
1663        // Penalty K(centers, centers), built once on CPU. The penalty
1664        // kernel evaluation is m × m (= 6400 entries), well outside the
1665        // GPU dispatch threshold, and identical for both paths under
1666        // test by construction.
1667        let k_cc = spherical_wahba_kernel_matrix_with_kind(
1668            centers_ll.view(),
1669            centers_ll.view(),
1670            penalty_order,
1671            false,
1672            kernel,
1673        )
1674        .expect("centers×centers kernel");
1675        let s_full = z.t().dot(&k_cc).dot(&z);
1676
1677        // CPU path: K(data, centers) via the public CPU helper.
1678        let raw_design_cpu = spherical_wahba_kernel_matrix_with_kind(
1679            data_ll.view(),
1680            centers_ll.view(),
1681            penalty_order,
1682            false,
1683            kernel,
1684        )
1685        .expect("CPU raw design");
1686        let x_s_cpu = raw_design_cpu.dot(&z);
1687
1688        // GPU path: K(data, centers) via `build_kernel_matrix_device`.
1689        let data_xyz = latlon_to_xyz_host(data_ll.view(), false).expect("data xyz");
1690        let centers_xyz = latlon_to_xyz_host(centers_ll.view(), false).expect("centers xyz");
1691        let coeffs = crate::basis::sobolev_s2_truncated_coefficients(lmax, penalty_order);
1692        let inputs = S2KernelBuildInputs {
1693            n,
1694            m,
1695            lmax,
1696            data_xyz: &data_xyz,
1697            centers_xyz: &centers_xyz,
1698            coeffs: &coeffs,
1699            kind: SphereSpectralKernelKind::Sobolev,
1700            layout: DeviceMatrixLayout::ColumnMajor,
1701        };
1702        // Deterministic synthetic response. The intent is to give the
1703        // penalised LS solve a non-trivial right-hand side; any smooth
1704        // function of the lat/lon is fine. Use a fixed-seed pseudo-
1705        // random walk derived from coordinates so the fixture has no
1706        // RNG dependency.
1707        let mut y = ndarray::Array1::<f64>::zeros(n);
1708        for i in 0..n {
1709            let lat_rad = data_ll[(i, 0)].to_radians();
1710            let lon_rad = data_ll[(i, 1)].to_radians();
1711            // Smooth ground truth + a tiny deterministic high-freq jitter.
1712            y[i] = (2.0 * lat_rad).sin() * (3.0 * lon_rad).cos()
1713                + 0.25 * lat_rad.cos() * (5.0 * lon_rad).sin();
1714        }
1715
1716        // Penalised normal-equation solve via faer LLT for each path:
1717        //   (X_sᵀ X_s + λ S) β = X_sᵀ y
1718        // S is symmetric positive semi-definite; λ S makes the system
1719        // strictly positive definite once added to X_sᵀ X_s.
1720        let solve_penalised = |x_s: &ndarray::Array2<f64>| -> ndarray::Array1<f64> {
1721            let xtx = x_s.t().dot(x_s);
1722            let mut a = xtx;
1723            for i in 0..p {
1724                for j in 0..p {
1725                    a[(i, j)] += lambda * s_full[(i, j)];
1726                }
1727            }
1728            let rhs = x_s.t().dot(&y);
1729            let factor = a
1730                .cholesky(Side::Lower)
1731                .expect("penalised normal equations are SPD under λ > 0");
1732            factor.solvevec(&rhs)
1733        };
1734
1735        let beta_cpu = solve_penalised(&x_s_cpu);
1736        assert_eq!(beta_cpu.len(), p);
1737        let yhat_cpu = x_s_cpu.dot(&beta_cpu);
1738        assert_eq!(x_s_cpu.dim(), (n, p));
1739
1740        // EVERY HOST: the CPU side of this comparison owes its own contracts —
1741        // the kernel matrix equals its elementwise spectral definition, and the
1742        // fitted coefficients solve the penalised normal equations. Both are
1743        // the oracle the device output is graded against.
1744        assert_cpu_kernel_matches_stable_spectral_definition(
1745            &raw_design_cpu,
1746            &data_ll,
1747            &centers_ll,
1748            &coeffs,
1749        );
1750        {
1751            let mut a = x_s_cpu.t().dot(&x_s_cpu);
1752            for i in 0..p {
1753                for j in 0..p {
1754                    a[(i, j)] += lambda * s_full[(i, j)];
1755                }
1756            }
1757            let residual = a.dot(&beta_cpu) - x_s_cpu.t().dot(&y);
1758            let rhs_scale = x_s_cpu
1759                .t()
1760                .dot(&y)
1761                .iter()
1762                .fold(0.0_f64, |acc, v| acc.max(v.abs()))
1763                .max(1.0);
1764            let max_residual = residual.iter().fold(0.0_f64, |acc, v| acc.max(v.abs()));
1765            assert!(
1766                max_residual <= 1e-9 * rhs_scale,
1767                "CPU penalised normal equations not solved: ‖(XᵀX + λS)β − Xᵀy‖∞ = \
1768                 {max_residual:.3e} (rhs scale {rhs_scale:.3e})"
1769            );
1770        }
1771
1772        if !cuda_available_for_test("end-to-end fit parity") {
1773            assert_sphere_decision_declines_without_device(n, m, lmax);
1774            assert_device_kernel_entry_refuses(inputs);
1775            return;
1776        }
1777        // Past the runtime Some-gate: a probe failure is a real device fault on a
1778        // CUDA host — fail loud (device-PCG skip-pass class, eee12f6b2).
1779        SphereGpuBackend::probe()
1780            .expect("[sphere gpu parity] sphere GPU backend probe must succeed on a CUDA host");
1781        let raw_dev = build_kernel_matrix_device(inputs).expect("GPU raw design");
1782        let raw_design_gpu = raw_dev.to_host_array().expect("dtoh GPU raw design");
1783        let x_s_gpu = raw_design_gpu.dot(&z);
1784
1785        assert_eq!(x_s_gpu.dim(), (n, p));
1786
1787        // PRIMARY GPU-OUTPUT PARITY (#1175): the only path-dependent quantity is
1788        // the GPU kernel matrix `K(data, centers)` → `x_s`. THIS is the genuine
1789        // device output and it must match the CPU kernel essentially bit-tight.
1790        // The downstream β is the solution of an ill-conditioned normal-equation
1791        // system that AMPLIFIES this difference by cond(XᵀX+λS) (see below), so
1792        // β is the wrong surface to gate at a flat 1e-9 — it tests the
1793        // conditioning of a SHARED CPU solve, not the GPU. Gate the GPU output
1794        // (x_s) tight; gate β with a condition-aware band; gate ŷ (the
1795        // customer-visible prediction) tight.
1796        let mut raw_xs_delta = 0.0_f64;
1797        let mut xs_scale = 0.0_f64;
1798        for (a, b) in x_s_cpu.iter().zip(x_s_gpu.iter()) {
1799            raw_xs_delta = raw_xs_delta.max((a - b).abs());
1800            xs_scale = xs_scale.max(a.abs());
1801        }
1802        // Condition number of A = XᵀX + λS (CPU path) via symmetric eigvals;
1803        // this is the factor that maps the x_s difference into the β difference.
1804        let cond = {
1805            use gam_linalg::faer_ndarray::FaerEigh;
1806            let xtx = x_s_cpu.t().dot(&x_s_cpu);
1807            let mut a = xtx;
1808            for i in 0..p {
1809                for j in 0..p {
1810                    a[(i, j)] += lambda * s_full[(i, j)];
1811                }
1812            }
1813            let (mut lo, mut hi) = (f64::INFINITY, 0.0_f64);
1814            if let Ok((vals, _)) = a.eigh(faer::Side::Lower) {
1815                for &v in vals.iter() {
1816                    lo = lo.min(v);
1817                    hi = hi.max(v);
1818                }
1819            }
1820            // A non-positive smallest eigenvalue is an unbounded condition
1821            // number, reported as such rather than as `hi / 1e-300`.
1822            if lo > 0.0 { hi / lo } else { f64::INFINITY }
1823        };
1824        // GPU kernel output must be bit-tight to the CPU oracle: measured on a
1825        // V100 the raw design parity is ~1e-16 (one ULP, rel ~1.2e-15). Gate at
1826        // a small ULP-scaled band — a real kernel bug perturbs x_s at O(scale),
1827        // 14+ orders above this floor.
1828        assert!(
1829            raw_xs_delta <= 1e-12 * xs_scale.max(1.0),
1830            "GPU vs CPU sphere design matrix max |Δ| = {raw_xs_delta:.3e} > {:.3e} \
1831             (scale {xs_scale:.3e}) — the kernel itself drifted (this is the genuine \
1832             GPU output, NOT a conditioning artifact)",
1833            1e-12 * xs_scale.max(1.0)
1834        );
1835
1836        let beta_gpu = solve_penalised(&x_s_gpu);
1837        assert_eq!(beta_gpu.len(), p);
1838
1839        // Fitted values for both paths use their own design matrices —
1840        // this is the customer-visible quantity (prediction at training
1841        // points).
1842        let yhat_gpu = x_s_gpu.dot(&beta_gpu);
1843
1844        let mut max_beta_delta = 0.0_f64;
1845        for k in 0..p {
1846            let d = (beta_cpu[k] - beta_gpu[k]).abs();
1847            if d > max_beta_delta {
1848                max_beta_delta = d;
1849            }
1850        }
1851        let mut max_fit_delta = 0.0_f64;
1852        for i in 0..n {
1853            let d = (yhat_cpu[i] - yhat_gpu[i]).abs();
1854            if d > max_fit_delta {
1855                max_fit_delta = d;
1856            }
1857        }
1858
1859        eprintln!(
1860            "[sphere_gpu fit parity] n={n} m={m} p={p} lmax={lmax} λ={lambda:.1e} \
1861             raw_xs|Δ|={raw_xs_delta:.3e} cond={cond:.3e} \
1862             max|Δβ|={max_beta_delta:.3e} max|Δŷ|={max_fit_delta:.3e}"
1863        );
1864
1865        // FITTED VALUES (the customer-visible prediction) must be tight. ŷ is a
1866        // well-conditioned functional of the data even when β is not (the
1867        // ill-conditioned directions of A correspond to β components that x_s
1868        // barely projects onto, so they cancel in ŷ = x_s·β). Measured on a
1869        // V100: max|Δŷ| ~7.6e-11. Gate tight — this is the quantity that
1870        // actually matters and it does NOT inherit the conditioning blow-up.
1871        assert!(
1872            max_fit_delta <= 1.0e-9,
1873            "GPU vs CPU truncated-spectral fitted-value max |Δ| = {max_fit_delta:.3e} > 1e-9"
1874        );
1875
1876        // COEFFICIENTS: β = A⁻¹ Xᵀy with A = XᵀX + λS. Standard perturbation
1877        // theory bounds the relative coefficient error by cond(A) times the
1878        // relative input (x_s) error: ‖Δβ‖/‖β‖ ≲ cond(A)·‖Δx_s‖/‖x_s‖. With the
1879        // GPU/CPU x_s difference at the ULP floor (~1e-16 relative) and
1880        // cond(A) ≈ 5e7 on this fixture, β legitimately differs by ~1e-7 — NOT
1881        // a kernel bug (the raw design parity gate above already proved the GPU
1882        // output is bit-tight). A flat 1e-9 β gate is therefore wrong: it
1883        // measures the conditioning of the SHARED CPU solve, not the GPU. Gate
1884        // β against the condition-aware bound with 16× headroom; a genuine
1885        // kernel defect would already have been caught upstream by the raw x_s
1886        // gate (which has no conditioning amplification).
1887        let beta_tol = (1e-15 * cond * (1.0 + xs_scale)).max(1e-9) * 16.0;
1888        assert!(
1889            max_beta_delta <= beta_tol,
1890            "GPU vs CPU truncated-spectral coefficient max |Δ| = {max_beta_delta:.3e} > \
1891             condition-aware tol {beta_tol:.3e} (cond={cond:.3e}). Raw design parity is \
1892             {raw_xs_delta:.3e}; a drift THIS much larger than cond·ULP is a real solve/kernel \
1893             mismatch, not conditioning."
1894        );
1895    }
1896}