1use 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#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
40pub enum SphereSpectralKernelKind {
41 Sobolev,
43 Pseudo,
45}
46
47impl SphereSpectralKernelKind {
48 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 pub const fn tag(self) -> &'static str {
64 match self {
65 SphereSpectralKernelKind::Sobolev => "sobolev",
66 SphereSpectralKernelKind::Pseudo => "pseudo",
67 }
68 }
69}
70
71#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
75pub enum DeviceMatrixLayout {
76 ColumnMajor,
77}
78
79pub 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 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#[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 pub col_major_dev: Vec<f64>,
134}
135
136impl DeviceS2KernelMatrix {
137 #[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 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 #[cfg(target_os = "linux")]
187 pub fn copy_to_host_col_major(&self, dst: &mut [f64]) -> Result<(), GpuError> {
188 let needed = self.ld * self.cols;
189 if dst.len() != needed {
190 gam_gpu::gpu_bail!(
191 "DeviceS2KernelMatrix::copy_to_host_col_major: dst.len()={} expected {}",
192 dst.len(),
193 needed
194 );
195 }
196 self.stream
197 .memcpy_dtoh(&self.col_major_dev, dst)
198 .gpu_ctx("DeviceS2KernelMatrix dtoh")?;
199 self.stream
200 .synchronize()
201 .gpu_ctx("DeviceS2KernelMatrix synchronize")?;
202 Ok(())
203 }
204
205 #[cfg(not(target_os = "linux"))]
206 pub fn copy_to_host_col_major(&self, dst: &mut [f64]) -> Result<(), GpuError> {
207 let needed = self.ld * self.cols;
208 if dst.len() != needed {
209 gam_gpu::gpu_bail!(
210 "DeviceS2KernelMatrix::copy_to_host_col_major: dst.len()={} expected {}",
211 dst.len(),
212 needed
213 );
214 }
215 dst.copy_from_slice(&self.col_major_dev);
216 Ok(())
217 }
218}
219
220fn col_major_to_row_major_parallel(
236 col_major: &[f64],
237 rows: usize,
238 cols: usize,
239 ld: usize,
240) -> Array2<f64> {
241 use rayon::prelude::*;
242
243 assert!(ld >= rows, "ld {ld} must be >= rows {rows}");
244 assert!(
245 col_major.len() >= ld * cols,
246 "col_major len {} < ld*cols {}",
247 col_major.len(),
248 ld * cols
249 );
250
251 const BLOCK_ROWS: usize = 128;
254
255 let mut out_flat = vec![0.0_f64; rows * cols];
256 out_flat
257 .par_chunks_mut(BLOCK_ROWS * cols)
258 .enumerate()
259 .for_each(|(block_idx, out_block)| {
260 let r0 = block_idx * BLOCK_ROWS;
261 let block_rows = out_block.len() / cols;
262 for j in 0..cols {
263 let base = j * ld + r0;
264 let src_col = &col_major[base..base + block_rows];
265 for (local_i, &v) in src_col.iter().enumerate() {
267 out_block[local_i * cols + j] = v;
268 }
269 }
270 });
271
272 Array2::from_shape_vec((rows, cols), out_flat).expect("row-major buffer has rows*cols elements")
273}
274
275#[cfg(target_os = "linux")]
286struct PinnedF64 {
287 ptr: *mut f64,
288 len: usize,
289 freed: bool,
290}
291
292#[cfg(target_os = "linux")]
293impl PinnedF64 {
294 fn alloc(ctx: &Arc<CudaContext>, len: usize) -> Result<Self, GpuError> {
297 ctx.bind_to_thread().gpu_ctx("PinnedF64 bind_to_thread")?;
298 let bytes = len
299 .checked_mul(std::mem::size_of::<f64>())
300 .ok_or_else(|| gam_gpu::gpu_err!("PinnedF64: len={len} byte size overflows usize"))?;
301 let raw = unsafe { cudarc::driver::result::malloc_host(bytes, 0) }
306 .gpu_ctx("PinnedF64 cuMemHostAlloc")?;
307 let ptr = raw as *mut f64;
308 if ptr.is_null() {
309 gam_gpu::gpu_bail!("PinnedF64: cuMemHostAlloc returned null for {bytes} bytes");
310 }
311 Ok(Self {
312 ptr,
313 len,
314 freed: false,
315 })
316 }
317
318 fn as_mut_slice(&mut self) -> &mut [f64] {
319 unsafe { std::slice::from_raw_parts_mut(self.ptr, self.len) }
322 }
323
324 fn as_slice(&self) -> &[f64] {
325 unsafe { std::slice::from_raw_parts(self.ptr, self.len) }
327 }
328}
329
330#[cfg(target_os = "linux")]
331impl Drop for PinnedF64 {
332 fn drop(&mut self) {
333 if self.freed {
334 return;
335 }
336 self.freed = true;
337 if let Err(err) =
342 unsafe { cudarc::driver::result::free_host(self.ptr as *mut std::ffi::c_void) }
343 {
344 log::debug!(
345 "PinnedF64::drop: cuMemFreeHost failed ({err}); the pinned host allocation \
346 is leaked for the remaining process lifetime"
347 );
348 }
349 }
350}
351
352#[cfg(target_os = "linux")]
358unsafe impl Send for PinnedF64 {}
359
360#[cfg(target_os = "linux")]
369const PINNED_POOL_MAX_BUFFERS: usize = 4;
370
371#[cfg(target_os = "linux")]
372static PINNED_POOL: OnceLock<Mutex<Vec<PinnedF64>>> = OnceLock::new();
373
374#[cfg(target_os = "linux")]
378struct PinnedLease {
379 buf: Option<PinnedF64>,
380}
381
382#[cfg(target_os = "linux")]
383impl PinnedLease {
384 fn acquire(ctx: &Arc<CudaContext>, len: usize) -> Result<Self, GpuError> {
387 let pool = PINNED_POOL.get_or_init(|| Mutex::new(Vec::new()));
388 if let Ok(mut guard) = pool.lock() {
389 if let Some(pos) = guard.iter().position(|b| b.len == len) {
390 return Ok(Self {
391 buf: Some(guard.swap_remove(pos)),
392 });
393 }
394 }
395 Ok(Self {
396 buf: Some(PinnedF64::alloc(ctx, len)?),
397 })
398 }
399
400 fn as_mut_slice(&mut self) -> &mut [f64] {
401 self.buf
402 .as_mut()
403 .expect("PinnedLease buffer present until drop")
404 .as_mut_slice()
405 }
406
407 fn as_slice(&self) -> &[f64] {
408 self.buf
409 .as_ref()
410 .expect("PinnedLease buffer present until drop")
411 .as_slice()
412 }
413}
414
415#[cfg(target_os = "linux")]
416impl Drop for PinnedLease {
417 fn drop(&mut self) {
418 let Some(buf) = self.buf.take() else {
419 return;
420 };
421 if let Some(pool) = PINNED_POOL.get() {
422 if let Ok(mut guard) = pool.lock() {
423 if guard.len() < PINNED_POOL_MAX_BUFFERS {
424 guard.push(buf);
425 return;
426 }
427 guard.remove(0);
431 guard.push(buf);
432 return;
433 }
434 }
435 drop(buf);
437 }
438}
439
440#[derive(Clone, Debug)]
451pub struct S2KernelBuildInputs<'a> {
452 pub n: usize,
453 pub m: usize,
454 pub lmax: usize,
455 pub data_xyz: &'a [f64],
456 pub centers_xyz: &'a [f64],
457 pub coeffs: &'a [f64],
458 pub kind: SphereSpectralKernelKind,
459 pub layout: DeviceMatrixLayout,
460}
461
462impl<'a> S2KernelBuildInputs<'a> {
463 fn validate(&self) -> Result<(), GpuError> {
464 if self.lmax == 0 {
465 return Err(GpuError::DriverCallFailed {
466 reason: "S2KernelBuildInputs: lmax must be >= 1".into(),
467 });
468 }
469 if self.data_xyz.len() != 3 * self.n {
470 gam_gpu::gpu_bail!(
471 "S2KernelBuildInputs: data_xyz.len()={} != 3*n={}",
472 self.data_xyz.len(),
473 3 * self.n
474 );
475 }
476 if self.centers_xyz.len() != 3 * self.m {
477 gam_gpu::gpu_bail!(
478 "S2KernelBuildInputs: centers_xyz.len()={} != 3*m={}",
479 self.centers_xyz.len(),
480 3 * self.m
481 );
482 }
483 if self.coeffs.len() != self.lmax + 1 {
484 gam_gpu::gpu_bail!(
485 "S2KernelBuildInputs: coeffs.len()={} != lmax+1={}",
486 self.coeffs.len(),
487 self.lmax + 1
488 );
489 }
490 if self.coeffs[0] != 0.0 {
491 return Err(GpuError::DriverCallFailed {
492 reason: "S2KernelBuildInputs: coeffs[0] must be 0 (mean-zero kernel)".into(),
493 });
494 }
495 Ok(())
496 }
497}
498
499#[cfg(target_os = "linux")]
509const KERNEL_TEMPLATE: &str = r#"
510// LMAX is supplied by the host via a `#define LMAX ...` prepended to
511// this source before NVRTC compilation (see `SphereGpuBackend::module_for`).
512// Recover cos(gamma) from the two half-angle chord lengths instead of
513// x dot c. The dot product rounds 1 - O(gamma^2) to 1 near coincidence,
514// permanently destroying the separation before the spectral evaluator sees it.
515// Here u = |x-c|^2 / (|x-c|^2 + |x+c|^2) and
516// v = |x+c|^2 / (|x-c|^2 + |x+c|^2), so both singular ends are carried
517// without cancellation and exact coincidence gives u=0, v=1 by construction.
518__device__ __forceinline__
519double s2_chord_cos_gamma(
520 double xi,
521 double yi,
522 double zi,
523 double cxj,
524 double cyj,
525 double czj
526) {
527 const double dx = xi - cxj;
528 const double dy = yi - cyj;
529 const double dz = zi - czj;
530 const double sx = xi + cxj;
531 const double sy = yi + cyj;
532 const double sz = zi + czj;
533 const double chord_sq = fma(dx, dx, fma(dy, dy, dz * dz));
534 const double anti_chord_sq = fma(sx, sx, fma(sy, sy, sz * sz));
535 const double scale = chord_sq + anti_chord_sq;
536
537 double u = chord_sq / scale;
538 double v = anti_chord_sq / scale;
539 if (u > 1.0) u = 1.0;
540 if (u < 0.0) u = 0.0;
541 if (v > 1.0) v = 1.0;
542 if (v < 0.0) v = 0.0;
543
544 double cos_gamma = v - u;
545 if (cos_gamma > 1.0) cos_gamma = 1.0;
546 if (cos_gamma < -1.0) cos_gamma = -1.0;
547 return cos_gamma;
548}
549
550extern "C" __global__
551__launch_bounds__(256)
552void s2_wahba_legendre_colmajor(
553 const double* __restrict__ data_xyz, // n × 3 (row-major flat)
554 const double* __restrict__ centers_xyz, // m × 3 (row-major flat)
555 const double* __restrict__ coeffs, // length LMAX + 1, coeffs[0] = 0
556 int n,
557 int m,
558 long long ld,
559 double* __restrict__ out // ld × m column-major
560) {
561 const int i = blockIdx.y * blockDim.y + threadIdx.y;
562 const int j = blockIdx.x * blockDim.x + threadIdx.x;
563 if (i >= n || j >= m) return;
564
565 // Load (x_i, y_i, z_i) and (cx_j, cy_j, cz_j) into registers.
566 const double xi = data_xyz[3 * i + 0];
567 const double yi = data_xyz[3 * i + 1];
568 const double zi = data_xyz[3 * i + 2];
569 const double cxj = centers_xyz[3 * j + 0];
570 const double cyj = centers_xyz[3 * j + 1];
571 const double czj = centers_xyz[3 * j + 2];
572
573 // Stable half-angle chord geometry; no near-coincident dot-product loss.
574 const double t = s2_chord_cos_gamma(xi, yi, zi, cxj, cyj, czj);
575
576 // Legendre 3-term recurrence in registers.
577 // P_0(t) = 1, P_1(t) = t.
578 double p_prev = 1.0;
579 double p_curr = t;
580 double acc = coeffs[0] * p_prev + coeffs[1] * p_curr;
581
582 #pragma unroll 8
583 for (int ell = 1; ell < LMAX; ++ell) {
584 const double lf = (double) ell;
585 const double inv = 1.0 / (lf + 1.0);
586 // p_{ell+1} = ((2ell+1) * t * p_curr - ell * p_prev) / (ell+1)
587 const double p_next =
588 fma((2.0 * lf + 1.0) * t, p_curr, -lf * p_prev) * inv;
589 acc = fma(coeffs[ell + 1], p_next, acc);
590 p_prev = p_curr;
591 p_curr = p_next;
592 }
593
594 out[(long long) j * ld + (long long) i] = acc;
595}
596
597// Fused Householder-constrained kernel (Phase 3). Z = I - beta · v · v^T,
598// the constrained design is X_s = B[:, 1..m] - beta * (B · v) · v[1..m]^T,
599// i.e. drop the first column after applying Z. Each thread computes one
600// row of B in registers (m kernel evaluations), forms d_i = B_row · v,
601// then emits X_s[i, j_out] = B_row[j_out + 1] - beta * d_i * v[j_out + 1]
602// for j_out in 0..m-1.
603//
604// Grid: 1D over rows (block_dim.x rows per block). Each thread iterates
605// over centers in an inner loop — register-bound by the per-row state
606// (xyz_i, p_prev, p_curr, acc, and a small per-center scratch).
607extern "C" __global__
608__launch_bounds__(128)
609void s2_wahba_householder_constrained_colmajor(
610 const double* __restrict__ data_xyz, // n × 3
611 const double* __restrict__ centers_xyz, // m × 3
612 const double* __restrict__ coeffs, // length LMAX + 1
613 const double* __restrict__ v, // length m, Householder vector
614 double beta,
615 int n,
616 int m,
617 long long ld_out,
618 double* __restrict__ out // ld_out × (m-1) column-major
619) {
620 const int i = blockIdx.x * blockDim.x + threadIdx.x;
621 if (i >= n) return;
622
623 const double xi = data_xyz[3 * i + 0];
624 const double yi = data_xyz[3 * i + 1];
625 const double zi = data_xyz[3 * i + 2];
626
627 // Pass 1: compute d_i = sum_j v[j] * B[i, j].
628 double d_i = 0.0;
629 for (int j = 0; j < m; ++j) {
630 const double cxj = centers_xyz[3 * j + 0];
631 const double cyj = centers_xyz[3 * j + 1];
632 const double czj = centers_xyz[3 * j + 2];
633 const double t = s2_chord_cos_gamma(xi, yi, zi, cxj, cyj, czj);
634
635 double p_prev = 1.0;
636 double p_curr = t;
637 double acc = coeffs[0] * p_prev + coeffs[1] * p_curr;
638 #pragma unroll 8
639 for (int ell = 1; ell < LMAX; ++ell) {
640 const double lf = (double) ell;
641 const double inv = 1.0 / (lf + 1.0);
642 const double p_next =
643 fma((2.0 * lf + 1.0) * t, p_curr, -lf * p_prev) * inv;
644 acc = fma(coeffs[ell + 1], p_next, acc);
645 p_prev = p_curr;
646 p_curr = p_next;
647 }
648 d_i = fma(v[j], acc, d_i);
649 }
650
651 // Pass 2: emit X_s[i, j_out] = B[i, j_out+1] - beta * d_i * v[j_out+1].
652 const double bd = beta * d_i;
653 for (int j_out = 0; j_out < m - 1; ++j_out) {
654 const int j = j_out + 1;
655 const double cxj = centers_xyz[3 * j + 0];
656 const double cyj = centers_xyz[3 * j + 1];
657 const double czj = centers_xyz[3 * j + 2];
658 const double t = s2_chord_cos_gamma(xi, yi, zi, cxj, cyj, czj);
659
660 double p_prev = 1.0;
661 double p_curr = t;
662 double acc = coeffs[0] * p_prev + coeffs[1] * p_curr;
663 #pragma unroll 8
664 for (int ell = 1; ell < LMAX; ++ell) {
665 const double lf = (double) ell;
666 const double inv = 1.0 / (lf + 1.0);
667 const double p_next =
668 fma((2.0 * lf + 1.0) * t, p_curr, -lf * p_prev) * inv;
669 acc = fma(coeffs[ell + 1], p_next, acc);
670 p_prev = p_curr;
671 p_curr = p_next;
672 }
673 const double xs = acc - bd * v[j];
674 out[(long long) j_out * ld_out + (long long) i] = xs;
675 }
676}
677"#;
678
679#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
689pub struct S2ModuleCacheKey {
690 pub cc_major: i32,
691 pub cc_minor: i32,
692 pub lmax: u32,
693 pub kind: SphereSpectralKernelKind,
694 pub layout: DeviceMatrixLayout,
695}
696
697pub const fn sphere_gpu_compiled() -> bool {
700 cfg!(target_os = "linux")
701}
702
703#[must_use]
710pub fn sphere_kernel_decision(n: usize, m: usize, lmax: usize) -> Result<GpuDecision, GpuError> {
711 let large_enough = match gam_gpu::device_runtime::GpuRuntime::resolve(gam_gpu::global_policy())?
712 {
713 Some(runtime) => {
714 let ld = ((n + 31) / 32) * 32;
715 let needed_bytes = ld
716 .saturating_mul(m)
717 .saturating_mul(std::mem::size_of::<f64>());
718 let budget = runtime.memory_budget_bytes;
719 n.saturating_mul(m) >= 1_000_000 && lmax <= 200 && needed_bytes <= budget
720 }
721 None => false,
722 };
723 decide(
724 GpuKernel::SpatialKernelOperator,
725 gam_gpu::GpuEligibility::from_flags(sphere_gpu_compiled(), large_enough),
726 )
727}
728
729#[must_use]
735pub fn truncated_device_kind(
736 kernel: crate::basis::SphereWahbaKernel,
737) -> Option<(SphereSpectralKernelKind, u16)> {
738 use crate::basis::SphereWahbaKernel;
739 match kernel {
740 SphereWahbaKernel::SobolevTruncated { lmax } => {
741 Some((SphereSpectralKernelKind::Sobolev, lmax))
742 }
743 SphereWahbaKernel::PseudoTruncated { lmax } => {
744 Some((SphereSpectralKernelKind::Pseudo, lmax))
745 }
746 SphereWahbaKernel::Sobolev | SphereWahbaKernel::Pseudo => None,
747 }
748}
749
750pub fn try_build_truncated_kernel_matrix_gpu(
772 data: ArrayView2<'_, f64>,
773 centers: ArrayView2<'_, f64>,
774 penalty_order: usize,
775 radians: bool,
776 kernel: crate::basis::SphereWahbaKernel,
777) -> Option<Result<Array2<f64>, GpuError>> {
778 let (kind, lmax) = truncated_device_kind(kernel)?;
779 let n = data.nrows();
780 let m = centers.nrows();
781 if n == 0 || m == 0 || lmax == 0 {
782 return None;
783 }
784 let decision = match sphere_kernel_decision(n, m, lmax as usize) {
785 Ok(decision) => decision,
786 Err(error) => return Some(Err(error)),
787 };
788 if !decision.use_gpu {
789 return None;
792 }
793 Some(build_truncated_kernel_matrix_gpu_admitted(
795 data,
796 centers,
797 penalty_order,
798 radians,
799 kind,
800 lmax,
801 ))
802}
803
804fn build_truncated_kernel_matrix_gpu_admitted(
808 data: ArrayView2<'_, f64>,
809 centers: ArrayView2<'_, f64>,
810 penalty_order: usize,
811 radians: bool,
812 kind: SphereSpectralKernelKind,
813 lmax: u16,
814) -> Result<Array2<f64>, GpuError> {
815 let n = data.nrows();
816 let m = centers.nrows();
817 let data_xyz = latlon_to_xyz_host(data, radians)
818 .map_err(|reason| GpuError::DriverCallFailed { reason })?;
819 let centers_xyz = latlon_to_xyz_host(centers, radians)
820 .map_err(|reason| GpuError::DriverCallFailed { reason })?;
821 let coeffs = kind.coefficients(lmax as usize, penalty_order);
825 let inputs = S2KernelBuildInputs {
826 n,
827 m,
828 lmax: lmax as usize,
829 data_xyz: &data_xyz,
830 centers_xyz: ¢ers_xyz,
831 coeffs: &coeffs,
832 kind,
833 layout: DeviceMatrixLayout::ColumnMajor,
834 };
835 let device_matrix = build_kernel_matrix_device(inputs)?;
836 let out = device_matrix.to_host_array()?;
837 if !out.sum().is_finite() {
848 return Err(GpuError::DriverCallFailed {
849 reason: "sphere GPU truncated kernel produced a non-finite value".to_string(),
850 });
851 }
852 Ok(out)
853}
854
855#[cfg(target_os = "linux")]
856struct SphereGpuContext {
857 ctx: Arc<CudaContext>,
858 stream: Arc<CudaStream>,
859 modules: Mutex<HashMap<S2ModuleCacheKey, Arc<CudaModule>>>,
860 cc_major: i32,
861 cc_minor: i32,
862}
863
864pub struct SphereGpuBackend {
867 #[cfg(target_os = "linux")]
868 inner: SphereGpuContext,
869}
870
871impl SphereGpuBackend {
872 pub fn probe() -> Result<&'static Self, GpuError> {
874 static BACKEND: OnceLock<Result<SphereGpuBackend, GpuError>> = OnceLock::new();
875 BACKEND
876 .get_or_init(|| {
877 #[cfg(target_os = "linux")]
878 {
879 Self::probe_linux()
880 }
881 #[cfg(not(target_os = "linux"))]
882 {
883 Err(GpuError::DriverLibraryUnavailable {
884 reason: "sphere GPU backend is Linux-only".to_string(),
885 })
886 }
887 })
888 .as_ref()
889 .map_err(GpuError::clone)
890 }
891
892 #[cfg(target_os = "linux")]
893 fn probe_linux() -> Result<Self, GpuError> {
894 let parts = gam_gpu::backend_probe::probe_cuda_backend("sphere")?;
895 Ok(SphereGpuBackend {
896 inner: SphereGpuContext {
897 ctx: parts.ctx,
898 stream: parts.stream,
899 modules: Mutex::new(HashMap::new()),
900 cc_major: parts.capability.compute_major,
901 cc_minor: parts.capability.compute_minor,
902 },
903 })
904 }
905
906 #[cfg(target_os = "linux")]
909 fn module_for(&self, key: S2ModuleCacheKey) -> Result<Arc<CudaModule>, GpuError> {
910 if let Ok(guard) = self.inner.modules.lock() {
911 if let Some(existing) = guard.get(&key) {
912 return Ok(existing.clone());
913 }
914 }
915 let src = format!("#define LMAX {}\n{}", key.lmax, KERNEL_TEMPLATE);
925 let ptx = gam_gpu::device_cache::compile_ptx_arch(&src).gpu_ctx_with(|err| {
926 format!(
927 "sphere NVRTC compile (kind={}, lmax={}): {err}",
928 key.kind.tag(),
929 key.lmax
930 )
931 })?;
932 let module = self
933 .inner
934 .ctx
935 .load_module(ptx)
936 .gpu_ctx("sphere module load")?;
937 if let Ok(mut guard) = self.inner.modules.lock() {
938 guard.entry(key).or_insert_with(|| module.clone());
939 }
940 Ok(module)
941 }
942
943 #[cfg(target_os = "linux")]
944 fn cc(&self) -> (i32, i32) {
945 (self.inner.cc_major, self.inner.cc_minor)
946 }
947}
948
949pub fn build_kernel_matrix_device(
956 inputs: S2KernelBuildInputs<'_>,
957) -> Result<DeviceS2KernelMatrix, GpuError> {
958 inputs.validate()?;
959
960 #[cfg(target_os = "linux")]
961 {
962 use cudarc::driver::{LaunchConfig, PushKernelArg};
963 let backend = SphereGpuBackend::probe()?;
964 let (cc_major, cc_minor) = backend.cc();
965 let key = S2ModuleCacheKey {
966 cc_major,
967 cc_minor,
968 lmax: inputs.lmax as u32,
969 kind: inputs.kind,
970 layout: inputs.layout,
971 };
972 let module = backend.module_for(key)?;
973 let func = module
974 .load_function("s2_wahba_legendre_colmajor")
975 .gpu_ctx("sphere load_function raw")?;
976 let stream = backend.inner.stream.clone();
977
978 let data_dev = stream
979 .clone_htod(inputs.data_xyz)
980 .gpu_ctx("sphere htod data_xyz")?;
981 let centers_dev = stream
982 .clone_htod(inputs.centers_xyz)
983 .gpu_ctx("sphere htod centers_xyz")?;
984 let coeffs_dev = stream
985 .clone_htod(inputs.coeffs)
986 .gpu_ctx("sphere htod coeffs")?;
987
988 let n = inputs.n;
989 let m = inputs.m;
990 let ld = ((n + 31) / 32) * 32;
991 let mut out_dev = stream
992 .alloc_zeros::<f64>(ld * m)
993 .gpu_ctx_with(|err| format!("sphere alloc out (ld={ld}, m={m}): {err}"))?;
994
995 let block_x: u32 = 32;
997 let block_y: u32 = 8;
998 let grid_x: u32 = ((m as u32) + block_x - 1) / block_x;
999 let grid_y: u32 = ((n as u32) + block_y - 1) / block_y;
1000 let cfg = LaunchConfig {
1001 grid_dim: (grid_x, grid_y, 1),
1002 block_dim: (block_x, block_y, 1),
1003 shared_mem_bytes: 0,
1004 };
1005 let n_i32: i32 =
1006 i32::try_from(n).map_err(|_| gam_gpu::gpu_err!("sphere n={n} overflows i32"))?;
1007 let m_i32: i32 =
1008 i32::try_from(m).map_err(|_| gam_gpu::gpu_err!("sphere m={m} overflows i32"))?;
1009 let ld_i64: i64 = ld as i64;
1010
1011 let mut builder = stream.launch_builder(&func);
1012 builder
1013 .arg(&data_dev)
1014 .arg(¢ers_dev)
1015 .arg(&coeffs_dev)
1016 .arg(&n_i32)
1017 .arg(&m_i32)
1018 .arg(&ld_i64)
1019 .arg(&mut out_dev);
1020 unsafe { builder.launch(cfg) }.gpu_ctx("sphere raw kernel launch")?;
1025 stream
1026 .synchronize()
1027 .gpu_ctx("sphere raw kernel synchronize")?;
1028
1029 Ok(DeviceS2KernelMatrix {
1030 rows: n,
1031 cols: m,
1032 ld,
1033 col_major_dev: out_dev,
1034 stream,
1035 })
1036 }
1037
1038 #[cfg(not(target_os = "linux"))]
1039 {
1040 Err(GpuError::DriverLibraryUnavailable {
1041 reason: "sphere GPU backend is Linux-only".to_string(),
1042 })
1043 }
1044}
1045
1046pub fn build_householder_constrained_design_device(
1050 inputs: S2KernelBuildInputs<'_>,
1051 v: &[f64],
1052 beta: f64,
1053) -> Result<DeviceS2KernelMatrix, GpuError> {
1054 inputs.validate()?;
1055 if v.len() != inputs.m {
1056 gam_gpu::gpu_bail!(
1057 "build_householder_constrained_design_device: v.len()={} != m={}",
1058 v.len(),
1059 inputs.m
1060 );
1061 }
1062 if inputs.m < 2 {
1063 gam_gpu::gpu_bail!(
1064 "build_householder_constrained_design_device: m must be >= 2 (got {})",
1065 inputs.m
1066 );
1067 }
1068 if !beta.is_finite() {
1069 gam_gpu::gpu_bail!(
1070 "build_householder_constrained_design_device: beta must be finite (got {beta})"
1071 );
1072 }
1073
1074 #[cfg(target_os = "linux")]
1075 {
1076 use cudarc::driver::{LaunchConfig, PushKernelArg};
1077 let backend = SphereGpuBackend::probe()?;
1078 let (cc_major, cc_minor) = backend.cc();
1079 let key = S2ModuleCacheKey {
1080 cc_major,
1081 cc_minor,
1082 lmax: inputs.lmax as u32,
1083 kind: inputs.kind,
1084 layout: inputs.layout,
1085 };
1086 let module = backend.module_for(key)?;
1087 let func = module
1088 .load_function("s2_wahba_householder_constrained_colmajor")
1089 .gpu_ctx("sphere load_function householder")?;
1090 let stream = backend.inner.stream.clone();
1091
1092 let data_dev = stream
1093 .clone_htod(inputs.data_xyz)
1094 .gpu_ctx("sphere-hh htod data_xyz")?;
1095 let centers_dev = stream
1096 .clone_htod(inputs.centers_xyz)
1097 .gpu_ctx("sphere-hh htod centers_xyz")?;
1098 let coeffs_dev = stream
1099 .clone_htod(inputs.coeffs)
1100 .gpu_ctx("sphere-hh htod coeffs")?;
1101 let v_dev = stream.clone_htod(v).gpu_ctx("sphere-hh htod v")?;
1102
1103 let n = inputs.n;
1104 let m = inputs.m;
1105 let cols_out = m - 1;
1106 let ld_out = ((n + 31) / 32) * 32;
1107 let mut out_dev = stream
1108 .alloc_zeros::<f64>(ld_out * cols_out)
1109 .gpu_ctx_with(|err| {
1110 format!("sphere-hh alloc out (ld={ld_out}, cols={cols_out}): {err}")
1111 })?;
1112
1113 let block_x: u32 = 128;
1114 let grid_x: u32 = ((n as u32) + block_x - 1) / block_x;
1115 let cfg = LaunchConfig {
1116 grid_dim: (grid_x, 1, 1),
1117 block_dim: (block_x, 1, 1),
1118 shared_mem_bytes: 0,
1119 };
1120 let n_i32: i32 =
1121 i32::try_from(n).map_err(|_| gam_gpu::gpu_err!("sphere-hh n={n} overflows i32"))?;
1122 let m_i32: i32 =
1123 i32::try_from(m).map_err(|_| gam_gpu::gpu_err!("sphere-hh m={m} overflows i32"))?;
1124 let ld_out_i64: i64 = ld_out as i64;
1125
1126 let mut builder = stream.launch_builder(&func);
1127 builder
1128 .arg(&data_dev)
1129 .arg(¢ers_dev)
1130 .arg(&coeffs_dev)
1131 .arg(&v_dev)
1132 .arg(&beta)
1133 .arg(&n_i32)
1134 .arg(&m_i32)
1135 .arg(&ld_out_i64)
1136 .arg(&mut out_dev);
1137 unsafe { builder.launch(cfg) }.gpu_ctx("sphere-hh kernel launch")?;
1140 stream
1141 .synchronize()
1142 .gpu_ctx("sphere-hh kernel synchronize")?;
1143
1144 Ok(DeviceS2KernelMatrix {
1145 rows: n,
1146 cols: cols_out,
1147 ld: ld_out,
1148 col_major_dev: out_dev,
1149 stream,
1150 })
1151 }
1152
1153 #[cfg(not(target_os = "linux"))]
1154 {
1155 Err(GpuError::DriverLibraryUnavailable {
1156 reason: "sphere GPU backend is Linux-only".to_string(),
1157 })
1158 }
1159}
1160
1161pub fn householder_reflector_from_weights(w: &[f64]) -> (Vec<f64>, f64) {
1174 let m = w.len();
1175 if m == 0 {
1176 return (Vec::new(), 0.0);
1177 }
1178 let norm = w.iter().map(|x| x * x).sum::<f64>().sqrt();
1179 if norm == 0.0 {
1180 return (vec![0.0; m], 0.0);
1181 }
1182 let sigma = if w[0] >= 0.0 { norm } else { -norm };
1183 let mut v = w.to_vec();
1184 v[0] += sigma;
1185 let v0 = v[0];
1186 if v0 == 0.0 {
1187 return (vec![0.0; m], 0.0);
1188 }
1189 for entry in v.iter_mut() {
1191 *entry /= v0;
1192 }
1193 let vv: f64 = v.iter().map(|x| x * x).sum();
1195 let beta = 2.0 / vv;
1196 (v, beta)
1197}
1198
1199pub fn build_center_kernel_device(
1218 centers_xyz: &[f64],
1219 lmax: usize,
1220 coeffs: &[f64],
1221 kind: SphereSpectralKernelKind,
1222) -> Result<DeviceS2KernelMatrix, GpuError> {
1223 let m = centers_xyz.len() / 3;
1224 if centers_xyz.len() != 3 * m {
1225 return Err(GpuError::DriverCallFailed {
1226 reason: "build_center_kernel_device: centers_xyz length not divisible by 3".into(),
1227 });
1228 }
1229 let inputs = S2KernelBuildInputs {
1230 n: m,
1231 m,
1232 lmax,
1233 data_xyz: centers_xyz,
1234 centers_xyz,
1235 coeffs,
1236 kind,
1237 layout: DeviceMatrixLayout::ColumnMajor,
1238 };
1239 build_kernel_matrix_device(inputs)
1240}
1241
1242pub fn constrained_penalty_host(
1247 c: ArrayView2<'_, f64>,
1248 w: &[f64],
1249) -> Result<Array2<f64>, GpuError> {
1250 let (m1, m2) = c.dim();
1251 if m1 != m2 {
1252 gam_gpu::gpu_bail!("constrained_penalty_host: C must be square, got {m1}x{m2}");
1253 }
1254 let m = m1;
1255 if w.len() != m {
1256 gam_gpu::gpu_bail!("constrained_penalty_host: w.len()={} != m={}", w.len(), m);
1257 }
1258 if m < 2 {
1259 gam_gpu::gpu_bail!("constrained_penalty_host: m must be >= 2 (got {m})");
1260 }
1261 let (v, beta) = householder_reflector_from_weights(w);
1262
1263 let mut u = vec![0.0_f64; m];
1266 for i in 0..m {
1267 let mut acc = 0.0_f64;
1268 for j in 0..m {
1269 acc += c[(i, j)] * v[j];
1270 }
1271 u[i] = acc;
1272 }
1273 let vtcv: f64 = v.iter().zip(&u).map(|(vi, ui)| vi * ui).sum();
1274 let mut hch = Array2::<f64>::zeros((m, m));
1275 for i in 0..m {
1276 for j in 0..m {
1277 hch[(i, j)] =
1278 c[(i, j)] - beta * (v[i] * u[j] + u[i] * v[j]) + beta * beta * vtcv * v[i] * v[j];
1279 }
1280 }
1281 let mut s = Array2::<f64>::zeros((m - 1, m - 1));
1283 for i in 0..(m - 1) {
1284 for j in 0..(m - 1) {
1285 s[(i, j)] = hch[(i + 1, j + 1)];
1286 }
1287 }
1288 Ok(s)
1289}
1290
1291#[derive(Clone, Debug)]
1322pub struct PenalisedLsSolution {
1323 pub beta: Vec<f64>,
1325 pub weighted_residual_ssq: f64,
1327 pub log_det_hessian: f64,
1329}
1330
1331#[cfg(target_os = "linux")]
1341pub fn solve_penalised_ls_device(
1342 x_s_device: &DeviceS2KernelMatrix,
1343 wy: &[f64],
1344 r_s: ArrayView2<'_, f64>,
1345) -> Result<PenalisedLsSolution, GpuError> {
1346 use cudarc::cusolver::{DnHandle, sys as cusolver_sys};
1347 use cudarc::driver::DevicePtrMut;
1348
1349 let n = x_s_device.rows;
1350 let p = x_s_device.cols;
1351 if wy.len() != n {
1352 gam_gpu::gpu_bail!("solve_penalised_ls_device: wy.len()={} != n={n}", wy.len());
1353 }
1354 if r_s.dim() != (p, p) {
1355 gam_gpu::gpu_bail!(
1356 "solve_penalised_ls_device: r_s.dim()={:?} != ({p}, {p})",
1357 r_s.dim()
1358 );
1359 }
1360 if p == 0 {
1361 return Ok(PenalisedLsSolution {
1362 beta: Vec::new(),
1363 weighted_residual_ssq: wy.iter().map(|v| v * v).sum(),
1364 log_det_hessian: 0.0,
1365 });
1366 }
1367
1368 let stream = x_s_device.stream.clone();
1369 let n_aug = n + p;
1370
1371 let mut a_aug_host = vec![0.0_f64; n_aug * p];
1376 let mut x_host_colmajor = vec![0.0_f64; x_s_device.ld * p];
1378 x_s_device.copy_to_host_col_major(&mut x_host_colmajor)?;
1379 for j in 0..p {
1380 let src_off = j * x_s_device.ld;
1381 let dst_off = j * n_aug;
1382 a_aug_host[dst_off..dst_off + n].copy_from_slice(&x_host_colmajor[src_off..src_off + n]);
1383 for i in 0..p {
1384 a_aug_host[dst_off + n + i] = r_s[(i, j)];
1387 }
1388 }
1389 let mut a_dev = stream
1390 .clone_htod(&a_aug_host)
1391 .gpu_ctx("solve_penalised_ls_device htod A_aug")?;
1392
1393 let mut b_host = vec![0.0_f64; n_aug];
1395 b_host[..n].copy_from_slice(wy);
1396 let mut b_dev = stream
1397 .clone_htod(&b_host)
1398 .gpu_ctx("solve_penalised_ls_device htod b_aug")?;
1399
1400 let solver = DnHandle::new(stream.clone()).gpu_ctx("solve_penalised_ls_device DnHandle")?;
1401 let n_aug_i: i32 = i32::try_from(n_aug)
1402 .map_err(|_| gam_gpu::gpu_err!("solve_penalised_ls_device: n_aug={n_aug} overflows i32"))?;
1403 let p_i: i32 = i32::try_from(p)
1404 .map_err(|_| gam_gpu::gpu_err!("solve_penalised_ls_device: p={p} overflows i32"))?;
1405
1406 let mut lwork: i32 = 0;
1408 {
1409 let (a_ptr, _rec) = a_dev.device_ptr_mut(&stream);
1410 let status = unsafe {
1413 cusolver_sys::cusolverDnDgeqrf_bufferSize(
1414 solver.cu(),
1415 n_aug_i,
1416 p_i,
1417 a_ptr as *mut f64,
1418 n_aug_i,
1419 &mut lwork,
1420 )
1421 };
1422 if status != cusolver_sys::cusolverStatus_t::CUSOLVER_STATUS_SUCCESS {
1423 gam_gpu::gpu_bail!("cusolverDnDgeqrf_bufferSize status={status:?}");
1424 }
1425 }
1426 let lwork_us = usize::try_from(lwork)
1427 .map_err(|_| gam_gpu::gpu_err!("solve_penalised_ls_device: negative lwork={lwork}"))?;
1428 let mut workspace = stream
1429 .alloc_zeros::<f64>(lwork_us.max(1))
1430 .gpu_ctx("solve_penalised_ls_device alloc workspace")?;
1431 let mut tau = stream
1432 .alloc_zeros::<f64>(p)
1433 .gpu_ctx("solve_penalised_ls_device alloc tau")?;
1434 let mut info = stream
1435 .alloc_zeros::<i32>(1)
1436 .gpu_ctx("solve_penalised_ls_device alloc info")?;
1437
1438 {
1440 let (a_ptr, _rec_a) = a_dev.device_ptr_mut(&stream);
1441 let (tau_ptr, _rec_t) = tau.device_ptr_mut(&stream);
1442 let (work_ptr, _rec_w) = workspace.device_ptr_mut(&stream);
1443 let (info_ptr, _rec_i) = info.device_ptr_mut(&stream);
1444 let status = unsafe {
1447 cusolver_sys::cusolverDnDgeqrf(
1448 solver.cu(),
1449 n_aug_i,
1450 p_i,
1451 a_ptr as *mut f64,
1452 n_aug_i,
1453 tau_ptr as *mut f64,
1454 work_ptr as *mut f64,
1455 lwork,
1456 info_ptr as *mut i32,
1457 )
1458 };
1459 if status != cusolver_sys::cusolverStatus_t::CUSOLVER_STATUS_SUCCESS {
1460 gam_gpu::gpu_bail!("cusolverDnDgeqrf status={status:?}");
1461 }
1462 }
1463
1464 let mut ormqr_lwork: i32 = 0;
1466 {
1467 let (a_ptr, _rec_a) = a_dev.device_ptr_mut(&stream);
1468 let (tau_ptr, _rec_t) = tau.device_ptr_mut(&stream);
1469 let (b_ptr, _rec_b) = b_dev.device_ptr_mut(&stream);
1470 let status = unsafe {
1473 cusolver_sys::cusolverDnDormqr_bufferSize(
1474 solver.cu(),
1475 cusolver_sys::cublasSideMode_t::CUBLAS_SIDE_LEFT,
1476 cusolver_sys::cublasOperation_t::CUBLAS_OP_T,
1477 n_aug_i,
1478 1,
1479 p_i,
1480 a_ptr as *const f64,
1481 n_aug_i,
1482 tau_ptr as *const f64,
1483 b_ptr as *mut f64,
1484 n_aug_i,
1485 &mut ormqr_lwork,
1486 )
1487 };
1488 if status != cusolver_sys::cusolverStatus_t::CUSOLVER_STATUS_SUCCESS {
1489 gam_gpu::gpu_bail!("cusolverDnDormqr_bufferSize status={status:?}");
1490 }
1491 }
1492 if ormqr_lwork > lwork {
1493 workspace = stream
1494 .alloc_zeros::<f64>(usize::try_from(ormqr_lwork).unwrap_or(1))
1495 .gpu_ctx("solve_penalised_ls_device realloc workspace ormqr")?;
1496 }
1497 {
1498 let (a_ptr, _rec_a) = a_dev.device_ptr_mut(&stream);
1499 let (tau_ptr, _rec_t) = tau.device_ptr_mut(&stream);
1500 let (b_ptr, _rec_b) = b_dev.device_ptr_mut(&stream);
1501 let (work_ptr, _rec_w) = workspace.device_ptr_mut(&stream);
1502 let (info_ptr, _rec_i) = info.device_ptr_mut(&stream);
1503 let status = unsafe {
1507 cusolver_sys::cusolverDnDormqr(
1508 solver.cu(),
1509 cusolver_sys::cublasSideMode_t::CUBLAS_SIDE_LEFT,
1510 cusolver_sys::cublasOperation_t::CUBLAS_OP_T,
1511 n_aug_i,
1512 1,
1513 p_i,
1514 a_ptr as *const f64,
1515 n_aug_i,
1516 tau_ptr as *const f64,
1517 b_ptr as *mut f64,
1518 n_aug_i,
1519 work_ptr as *mut f64,
1520 ormqr_lwork.max(lwork),
1521 info_ptr as *mut i32,
1522 )
1523 };
1524 if status != cusolver_sys::cusolverStatus_t::CUSOLVER_STATUS_SUCCESS {
1525 gam_gpu::gpu_bail!("cusolverDnDormqr status={status:?}");
1526 }
1527 }
1528
1529 {
1532 use cudarc::cublas::CudaBlas;
1533 let blas = CudaBlas::new(stream.clone()).gpu_ctx("solve_penalised_ls_device CudaBlas")?;
1534 let alpha = 1.0_f64;
1535 let (a_ptr, _rec_a) = a_dev.device_ptr_mut(&stream);
1536 let (b_ptr, _rec_b) = b_dev.device_ptr_mut(&stream);
1537 let handle = *blas.handle();
1542 let status = unsafe {
1543 cudarc::cublas::sys::cublasDtrsm_v2(
1544 handle,
1545 cudarc::cublas::sys::cublasSideMode_t::CUBLAS_SIDE_LEFT,
1546 cudarc::cublas::sys::cublasFillMode_t::CUBLAS_FILL_MODE_UPPER,
1547 cudarc::cublas::sys::cublasOperation_t::CUBLAS_OP_N,
1548 cudarc::cublas::sys::cublasDiagType_t::CUBLAS_DIAG_NON_UNIT,
1549 p_i,
1550 1,
1551 &alpha,
1552 a_ptr as *const f64,
1553 n_aug_i,
1554 b_ptr as *mut f64,
1555 n_aug_i,
1556 )
1557 };
1558 if status != cudarc::cublas::sys::cublasStatus_t::CUBLAS_STATUS_SUCCESS {
1559 gam_gpu::gpu_bail!("cublasDtrsm_v2 status={status:?}");
1560 }
1561 }
1562
1563 let mut b_out = vec![0.0_f64; n_aug];
1565 stream
1566 .memcpy_dtoh(&b_dev, &mut b_out)
1567 .gpu_ctx("solve_penalised_ls_device dtoh b_out")?;
1568 let mut a_back = vec![0.0_f64; n_aug * p];
1569 stream
1570 .memcpy_dtoh(&a_dev, &mut a_back)
1571 .gpu_ctx("solve_penalised_ls_device dtoh A_back")?;
1572 stream
1573 .synchronize()
1574 .gpu_ctx("solve_penalised_ls_device synchronize")?;
1575
1576 let beta: Vec<f64> = b_out[..p].to_vec();
1577 let augmented_residual_ssq: f64 = b_out[p..].iter().map(|v| v * v).sum();
1586
1587 let mut log_abs_r = 0.0_f64;
1589 for k in 0..p {
1590 let r_kk = a_back[k * n_aug + k];
1591 log_abs_r += r_kk.abs().ln();
1592 }
1593 let log_det_hessian = 2.0 * log_abs_r;
1594
1595 Ok(PenalisedLsSolution {
1596 beta,
1597 weighted_residual_ssq: augmented_residual_ssq,
1598 log_det_hessian,
1599 })
1600}
1601
1602#[cfg(not(target_os = "linux"))]
1603pub fn solve_penalised_ls_device(
1604 x_s_device: &DeviceS2KernelMatrix,
1605 wy: &[f64],
1606 r_s: ArrayView2<'_, f64>,
1607) -> Result<PenalisedLsSolution, GpuError> {
1608 Err(GpuError::DriverLibraryUnavailable {
1609 reason: format!(
1610 "sphere GPU cuSOLVER QR path is Linux-only (n={}, p={}, wy.len()={}, r_s={:?})",
1611 x_s_device.rows,
1612 x_s_device.cols,
1613 wy.len(),
1614 r_s.dim()
1615 ),
1616 })
1617}
1618
1619#[cfg(test)]
1624mod sphere_gpu_tests {
1625 use super::*;
1626 use crate::basis::sphere_half_angle::{SphereTrig, half_angle_separation_scalar};
1627 use crate::basis::{
1628 SphereWahbaKernel, sobolev_s2_truncated_coefficients, sphere_truncated_spectral_eval,
1629 spherical_wahba_kernel_matrix_with_kind,
1630 };
1631 use ndarray::Array2;
1632
1633 fn small_latlon_grid(n_lat: usize, n_lon: usize) -> Array2<f64> {
1634 let mut rows = Vec::with_capacity(n_lat * n_lon);
1636 for i in 0..n_lat {
1637 let lat = -85.0 + (170.0 * i as f64) / (n_lat.saturating_sub(1).max(1) as f64);
1638 for j in 0..n_lon {
1639 let lon = -180.0 + (360.0 * j as f64) / (n_lon.saturating_sub(1).max(1) as f64);
1640 rows.push(lat);
1641 rows.push(lon);
1642 }
1643 }
1644 Array2::from_shape_vec((n_lat * n_lon, 2), rows).unwrap()
1645 }
1646
1647 fn cuda_available_for_test(label: &str) -> bool {
1648 match gam_gpu::device_runtime::GpuRuntime::resolve(gam_gpu::GpuPolicy::Auto) {
1649 Ok(Some(_)) => true,
1650 Ok(None) => {
1651 eprintln!("[sphere_gpu test] no CUDA device — skipping {label}");
1652 false
1653 }
1654 Err(error) => panic!("[sphere_gpu test] CUDA resolution failed for {label}: {error}"),
1655 }
1656 }
1657
1658 fn assert_sphere_decision_declines_without_device(n: usize, m: usize, lmax: usize) {
1666 let decision = sphere_kernel_decision(n, m, lmax)
1667 .expect("the sphere GPU decision must not fault on a device-free host");
1668 assert!(
1669 !decision.use_gpu,
1670 "no CUDA runtime on this host, yet the sphere dispatch decision admitted the \
1671 device for (n={n}, m={m}, lmax={lmax}) — reason={}",
1672 decision.reason
1673 );
1674 }
1675
1676 fn assert_device_kernel_entry_refuses(inputs: S2KernelBuildInputs<'_>) {
1681 assert!(
1682 build_kernel_matrix_device(inputs).is_err(),
1683 "no CUDA runtime on this host, yet the device kernel entry returned a matrix \
1684 — the admitted-only device path fabricated a host answer (#1551 class)"
1685 );
1686 }
1687
1688 fn assert_cpu_kernel_matches_stable_spectral_definition(
1694 kernel_matrix: &Array2<f64>,
1695 data_latlon: &Array2<f64>,
1696 centers_latlon: &Array2<f64>,
1697 coeffs: &[f64],
1698 ) {
1699 let (n, m) = kernel_matrix.dim();
1700 let to_radians = std::f64::consts::PI / 180.0;
1701 let mut max_abs = 0.0_f64;
1702 for i in 0..n {
1703 let point = SphereTrig::from_radians(
1704 data_latlon[(i, 0)] * to_radians,
1705 data_latlon[(i, 1)] * to_radians,
1706 );
1707 for j in 0..m {
1708 let center = SphereTrig::from_radians(
1709 centers_latlon[(j, 0)] * to_radians,
1710 centers_latlon[(j, 1)] * to_radians,
1711 );
1712 let separation = half_angle_separation_scalar(point, center);
1713 let expected = sphere_truncated_spectral_eval(separation.cos_gamma(), coeffs);
1714 max_abs = max_abs.max((kernel_matrix[(i, j)] - expected).abs());
1715 }
1716 }
1717 assert!(
1718 max_abs < 1e-12,
1719 "CPU truncated-spectral kernel matrix departs from the stable half-angle \
1720 elementwise definition: max |delta| = {max_abs:.3e}"
1721 );
1722 }
1723
1724 #[test]
1725 fn sum_finite_guard_accepts_finite_rejects_nonfinite() {
1726 let finite = Array2::<f64>::from_shape_fn((5, 7), |(i, j)| (i as f64 - 2.0) * (j as f64));
1731 assert!(finite.sum().is_finite());
1732
1733 let mut with_nan = finite.clone();
1734 with_nan[[3, 4]] = f64::NAN;
1735 assert!(!with_nan.sum().is_finite());
1736
1737 let mut with_pos_inf = finite.clone();
1738 with_pos_inf[[0, 0]] = f64::INFINITY;
1739 assert!(!with_pos_inf.sum().is_finite());
1740
1741 let mut with_neg_inf = finite.clone();
1742 with_neg_inf[[4, 6]] = f64::NEG_INFINITY;
1743 assert!(!with_neg_inf.sum().is_finite());
1744 }
1745
1746 #[test]
1747 fn xyz_preprocessing_matches_unit_sphere() {
1748 let latlon = ndarray::array![
1749 [0.0, 0.0],
1750 [90.0, 0.0],
1751 [0.0, 90.0],
1752 [-90.0, 17.5],
1753 [45.0, -120.0],
1754 ];
1755 let xyz = latlon_to_xyz_host(latlon.view(), false).expect("xyz");
1756 assert_eq!(xyz.len(), 3 * 5);
1757 for i in 0..5 {
1758 let nrm2 = xyz[3 * i] * xyz[3 * i]
1759 + xyz[3 * i + 1] * xyz[3 * i + 1]
1760 + xyz[3 * i + 2] * xyz[3 * i + 2];
1761 assert!((nrm2 - 1.0).abs() < 1e-15, "row {i} not unit norm: {nrm2}");
1762 }
1763 assert!((xyz[0] - 1.0).abs() < 1e-15);
1765 assert!(xyz[1].abs() < 1e-15);
1766 assert!(xyz[2].abs() < 1e-15);
1767 assert!(xyz[3].abs() < 1e-15);
1769 assert!(xyz[4].abs() < 1e-15);
1770 assert!((xyz[5] - 1.0).abs() < 1e-15);
1771 assert!(xyz[6].abs() < 1e-15);
1773 assert!((xyz[7] - 1.0).abs() < 1e-15);
1774 assert!(xyz[8].abs() < 1e-15);
1775 }
1776
1777 #[test]
1778 fn truncated_spectral_at_same_point_matches_sum_of_coefficients() {
1779 for m_penalty in 1..=4 {
1783 for &lmax in &[5_usize, 20, 50] {
1784 let coeffs = sobolev_s2_truncated_coefficients(lmax, m_penalty);
1785 let expected: f64 = coeffs.iter().sum();
1786 let got = sphere_truncated_spectral_eval(1.0, &coeffs);
1787 assert!(
1788 (got - expected).abs() < 1e-13,
1789 "K(x,x) identity broken at m={m_penalty}, L={lmax}: got {got:.6e}, expected {expected:.6e}"
1790 );
1791 }
1792 }
1793 }
1794
1795 #[test]
1796 fn truncated_spectral_at_antipode_matches_alternating_sum() {
1797 for m_penalty in 1..=4 {
1800 for &lmax in &[5_usize, 20, 50] {
1801 let coeffs = sobolev_s2_truncated_coefficients(lmax, m_penalty);
1802 let expected: f64 = coeffs
1803 .iter()
1804 .enumerate()
1805 .map(|(ell, c)| if ell % 2 == 0 { *c } else { -*c })
1806 .sum();
1807 let got = sphere_truncated_spectral_eval(-1.0, &coeffs);
1808 assert!(
1809 (got - expected).abs() < 1e-13,
1810 "K(x,-x) identity broken at m={m_penalty}, L={lmax}: got {got:.6e}, expected {expected:.6e}"
1811 );
1812 }
1813 }
1814 }
1815
1816 #[test]
1817 fn truncated_spectral_matrix_is_symmetric() {
1818 let centers = ndarray::array![
1822 [10.0_f64, 20.0],
1823 [-30.0, 100.0],
1824 [45.0, -60.0],
1825 [-89.0, 0.0],
1826 [0.0, 180.0],
1827 [60.0, -179.9],
1828 ];
1829 for m_penalty in [1usize, 2, 4] {
1830 for &lmax in &[10_usize, 30] {
1831 let mat = spherical_wahba_kernel_matrix_with_kind(
1832 centers.view(),
1833 centers.view(),
1834 m_penalty,
1835 false,
1836 SphereWahbaKernel::SobolevTruncated { lmax: lmax as u16 },
1837 )
1838 .expect("kernel matrix");
1839 let n = centers.nrows();
1840 let mut max_asym = 0.0_f64;
1841 for i in 0..n {
1842 for j in 0..n {
1843 let d = (mat[(i, j)] - mat[(j, i)]).abs();
1844 if d > max_asym {
1845 max_asym = d;
1846 }
1847 }
1848 }
1849 assert!(
1850 max_asym < 1e-13,
1851 "K not symmetric at m={m_penalty}, L={lmax}: max |K - Kᵀ| = {max_asym:.3e}"
1852 );
1853 }
1854 }
1855 }
1856
1857 #[test]
1858 fn truncated_coefficients_have_zero_constant_mode() {
1859 for m in 1..=4 {
1860 let c = sobolev_s2_truncated_coefficients(50, m);
1861 assert_eq!(c.len(), 51);
1862 assert_eq!(c[0], 0.0);
1863 assert!(c[1] > 0.0);
1864 for ell in 2..=50 {
1866 assert!(
1867 c[ell] < c[ell - 1] + 1e-15,
1868 "Sobolev coefficient not non-increasing at m={m}, ell={ell}: {} vs {}",
1869 c[ell],
1870 c[ell - 1]
1871 );
1872 }
1873 }
1874 }
1875
1876 #[test]
1877 fn truncated_spectral_matches_matrix_helper() {
1878 let m_penalty = 2;
1882 let lmax = 20;
1883 let coeffs = sobolev_s2_truncated_coefficients(lmax, m_penalty);
1884 let data = ndarray::array![[12.5, -34.0]];
1885 let centers = ndarray::array![[40.0, 10.0]];
1886 let mat = spherical_wahba_kernel_matrix_with_kind(
1887 data.view(),
1888 centers.view(),
1889 m_penalty,
1890 false,
1891 SphereWahbaKernel::SobolevTruncated { lmax: lmax as u16 },
1892 )
1893 .expect("kernel matrix");
1894 let to_radians = std::f64::consts::PI / 180.0;
1897 let point = SphereTrig::from_radians(data[(0, 0)] * to_radians, data[(0, 1)] * to_radians);
1898 let center =
1899 SphereTrig::from_radians(centers[(0, 0)] * to_radians, centers[(0, 1)] * to_radians);
1900 let expected = sphere_truncated_spectral_eval(
1901 half_angle_separation_scalar(point, center).cos_gamma(),
1902 &coeffs,
1903 );
1904 assert!(
1905 (mat[(0, 0)] - expected).abs() < 1e-13,
1906 "matrix helper differs from scalar evaluator: {} vs {}",
1907 mat[(0, 0)],
1908 expected
1909 );
1910 }
1911
1912 #[test]
1913 fn constrained_penalty_is_symmetric_and_drops_constraint_direction() {
1914 let m = 6;
1919 let mut c = Array2::<f64>::zeros((m, m));
1920 for i in 0..m {
1921 for j in 0..m {
1922 let d = (i as f64 - j as f64).abs();
1923 c[(i, j)] = (-0.5 * d).exp();
1924 }
1925 }
1926 let w = vec![1.0_f64; m];
1927 let s = constrained_penalty_host(c.view(), &w).expect("constrained S");
1928 assert_eq!(s.dim(), (m - 1, m - 1));
1929 let mut max_asym = 0.0_f64;
1931 for i in 0..(m - 1) {
1932 for j in 0..(m - 1) {
1933 let d = (s[(i, j)] - s[(j, i)]).abs();
1934 if d > max_asym {
1935 max_asym = d;
1936 }
1937 }
1938 }
1939 assert!(
1940 max_asym < 1e-13,
1941 "S not symmetric: max |S - Sᵀ| = {max_asym:.3e}"
1942 );
1943
1944 let ones = ndarray::Array1::<f64>::ones(m - 1);
1952 let sx = s.dot(&ones);
1953 assert!(sx.iter().all(|v| v.is_finite()));
1954 }
1955
1956 #[test]
1957 fn householder_reflector_zeroes_target_vector() {
1958 let w = vec![3.0, 4.0, 0.0, -1.0];
1959 let (v, beta) = householder_reflector_from_weights(&w);
1960 let dot: f64 = v.iter().zip(&w).map(|(a, b)| a * b).sum();
1963 let hw: Vec<f64> = w
1964 .iter()
1965 .zip(&v)
1966 .map(|(wj, vj)| wj - beta * dot * vj)
1967 .collect();
1968 for entry in hw.iter().skip(1) {
1969 assert!(entry.abs() < 1e-12, "H · w not e_1 multiple: {hw:?}");
1970 }
1971 assert!(hw[0].abs() > 0.0);
1972 }
1973
1974 #[test]
1980 fn sphere_gpu_raw_kernel_parity_vs_cpu_truncated() {
1981 let mut data_ll = small_latlon_grid(7, 9);
1982 let mut centers_ll = small_latlon_grid(5, 7);
1983 centers_ll[(0, 0)] = 12.5;
1986 centers_ll[(0, 1)] = -34.0;
1987 data_ll[(0, 0)] = 12.5 + 1.0e-8;
1988 data_ll[(0, 1)] = -34.0;
1989 let data_xyz = latlon_to_xyz_host(data_ll.view(), false).unwrap();
1990 let centers_xyz = latlon_to_xyz_host(centers_ll.view(), false).unwrap();
1991 let n = data_ll.nrows();
1992 let m = centers_ll.nrows();
1993 let penalty = 2usize;
1994 let lmax = 20usize;
1995 let coeffs = sobolev_s2_truncated_coefficients(lmax, penalty);
1996
1997 let inputs = S2KernelBuildInputs {
1998 n,
1999 m,
2000 lmax,
2001 data_xyz: &data_xyz,
2002 centers_xyz: ¢ers_xyz,
2003 coeffs: &coeffs,
2004 kind: SphereSpectralKernelKind::Sobolev,
2005 layout: DeviceMatrixLayout::ColumnMajor,
2006 };
2007
2008 let cpu = spherical_wahba_kernel_matrix_with_kind(
2009 data_ll.view(),
2010 centers_ll.view(),
2011 penalty,
2012 false,
2013 SphereWahbaKernel::SobolevTruncated { lmax: lmax as u16 },
2014 )
2015 .expect("cpu kernel matrix");
2016
2017 assert_cpu_kernel_matches_stable_spectral_definition(&cpu, &data_ll, ¢ers_ll, &coeffs);
2020
2021 if !cuda_available_for_test("raw-kernel parity") {
2022 assert_sphere_decision_declines_without_device(n, m, lmax);
2023 assert_device_kernel_entry_refuses(inputs);
2024 return;
2025 }
2026 SphereGpuBackend::probe()
2029 .expect("[sphere_gpu test] backend probe must succeed on a CUDA host");
2030 let dev_mat = build_kernel_matrix_device(inputs).expect("device kernel matrix");
2031 let gpu = dev_mat.to_host_array().expect("dtoh kernel matrix");
2032
2033 let mut max_abs = 0.0_f64;
2034 for i in 0..n {
2035 for j in 0..m {
2036 let d = (gpu[(i, j)] - cpu[(i, j)]).abs();
2037 if d > max_abs {
2038 max_abs = d;
2039 }
2040 }
2041 }
2042 assert!(
2043 max_abs < 1e-11,
2044 "GPU vs CPU truncated parity max |Δ| = {max_abs:.3e} >= 1e-11"
2045 );
2046 }
2047
2048 #[test]
2062 fn sphere_gpu_end_to_end_dispatch_parity_vs_cpu_truncated() {
2063 use crate::basis::{
2064 CenterStrategy, SphereMethod, SphericalSplineBasisSpec, SphericalSplineIdentifiability,
2065 build_spherical_spline_basis, spherical_wahba_kernel_matrix_cpu,
2066 spherical_wahba_kernel_matrix_with_kind,
2067 };
2068 let on_cuda = cuda_available_for_test("end-to-end dispatch parity");
2069 if on_cuda {
2070 SphereGpuBackend::probe()
2074 .expect("[sphere_gpu test] backend probe must succeed on a CUDA host");
2075 }
2076
2077 let data = small_latlon_grid(100, 100);
2079 let lmax: u16 = 30;
2080 let penalty_order = 2usize;
2081 let centers =
2082 crate::basis::select_spherical_farthest_point_centers(data.view(), 200, false)
2083 .expect("centers");
2084 let n = data.nrows();
2085 let m = centers.nrows();
2086
2087 if on_cuda {
2093 let decision = sphere_kernel_decision(n, m, lmax as usize)
2094 .expect("GPU decision must preserve CUDA resolution faults");
2095 assert!(
2096 decision.use_gpu,
2097 "expected GPU dispatch for (n={n}, m={m}, lmax={lmax}); decision said CPU \
2098 (reason={}); the engagement gate regressed",
2099 decision.reason
2100 );
2101 } else {
2102 assert_sphere_decision_declines_without_device(n, m, lmax as usize);
2103 assert!(
2104 try_build_truncated_kernel_matrix_gpu(
2105 data.view(),
2106 centers.view(),
2107 penalty_order,
2108 false,
2109 SphereWahbaKernel::SobolevTruncated { lmax },
2110 )
2111 .is_none(),
2112 "no CUDA runtime on this host, yet the production sphere seam did not take \
2113 the quiet CPU route at the device-eligible shape (n={n}, m={m}, lmax={lmax})"
2114 );
2115 }
2116
2117 let dispatched_kernel = spherical_wahba_kernel_matrix_with_kind(
2121 data.view(),
2122 centers.view(),
2123 penalty_order,
2124 false,
2125 SphereWahbaKernel::SobolevTruncated { lmax },
2126 )
2127 .expect("GPU-eligible production kernel build succeeds");
2128
2129 let cpu_kernel = spherical_wahba_kernel_matrix_cpu(
2131 data.view(),
2132 centers.view(),
2133 penalty_order,
2134 false,
2135 SphereWahbaKernel::SobolevTruncated { lmax },
2136 )
2137 .expect("cpu oracle kernel build succeeds");
2138
2139 assert_eq!(dispatched_kernel.dim(), cpu_kernel.dim());
2140 let mut max_abs = 0.0_f64;
2141 let mut max_rel = 0.0_f64;
2142 for (g, c) in dispatched_kernel.iter().zip(cpu_kernel.iter()) {
2143 let d = (g - c).abs();
2144 if d > max_abs {
2145 max_abs = d;
2146 }
2147 let denom = g.abs().max(c.abs()).max(1e-300);
2148 let r = d / denom;
2149 if r > max_rel {
2150 max_rel = r;
2151 }
2152 }
2153 assert!(
2154 max_rel < 1e-9,
2155 "GPU-dispatch vs CPU-oracle kernel parity max relative |Δ| = {max_rel:.3e} \
2156 >= 1e-9 (abs {max_abs:.3e})"
2157 );
2158 if !on_cuda {
2159 for (a, b) in dispatched_kernel.iter().zip(cpu_kernel.iter()) {
2163 assert_eq!(
2164 a.to_bits(),
2165 b.to_bits(),
2166 "device-free dispatcher must equal the CPU oracle bit-for-bit"
2167 );
2168 }
2169 }
2170
2171 let spec_gpu = SphericalSplineBasisSpec {
2175 center_strategy: CenterStrategy::FarthestPoint { num_centers: 200 },
2176 penalty_order,
2177 double_penalty: false,
2178 radians: false,
2179 method: SphereMethod::Wahba,
2180 max_degree: None,
2181 wahba_kernel: SphereWahbaKernel::SobolevTruncated { lmax },
2182 identifiability: SphericalSplineIdentifiability::CenterSumToZero,
2183 };
2184 let result_gpu = build_spherical_spline_basis(data.view(), &spec_gpu)
2185 .expect("GPU-eligible build_spherical_spline_basis succeeds");
2186 let design = result_gpu.design.as_dense().expect("dense design");
2187 assert_eq!(design.nrows(), n, "design row count must match data rows");
2188 assert!(
2189 design.iter().all(|v| v.is_finite()),
2190 "engaged-device spherical design must be finite"
2191 );
2192 }
2193
2194 fn householder_apply_host(b: &Array2<f64>, v: &[f64], beta: f64) -> Array2<f64> {
2198 let (n, m) = b.dim();
2199 let mut xs = Array2::<f64>::zeros((n, m - 1));
2200 for i in 0..n {
2201 let d_i: f64 = (0..m).map(|j| v[j] * b[(i, j)]).sum();
2202 for j_out in 0..(m - 1) {
2203 xs[(i, j_out)] = b[(i, j_out + 1)] - beta * d_i * v[j_out + 1];
2204 }
2205 }
2206 xs
2207 }
2208
2209 fn assert_householder_fused_matches_explicit_product(b: &Array2<f64>, v: &[f64], beta: f64) {
2214 let (n, m) = b.dim();
2215 let mut reflector = Array2::<f64>::eye(m);
2216 for i in 0..m {
2217 for j in 0..m {
2218 reflector[(i, j)] -= beta * v[i] * v[j];
2219 }
2220 }
2221 let full = b.dot(&reflector);
2222 let fused = householder_apply_host(b, v, beta);
2223 let mut max_abs = 0.0_f64;
2224 for i in 0..n {
2225 for j in 0..(m - 1) {
2226 max_abs = max_abs.max((fused[(i, j)] - full[(i, j + 1)]).abs());
2227 }
2228 }
2229 assert!(
2230 max_abs < 1e-13,
2231 "fused Householder host expression departs from B·(I − β·v·vᵀ): \
2232 max |Δ| = {max_abs:.3e}"
2233 );
2234 }
2235
2236 #[test]
2241 fn sphere_gpu_householder_parity_vs_raw_dot_z() {
2242 let data_ll = small_latlon_grid(6, 8);
2243 let centers_ll = small_latlon_grid(4, 5);
2244 let data_xyz = latlon_to_xyz_host(data_ll.view(), false).unwrap();
2245 let centers_xyz = latlon_to_xyz_host(centers_ll.view(), false).unwrap();
2246 let n = data_ll.nrows();
2247 let m = centers_ll.nrows();
2248 let penalty = 2usize;
2249 let lmax = 15usize;
2250 let coeffs = sobolev_s2_truncated_coefficients(lmax, penalty);
2251
2252 let inputs_raw = S2KernelBuildInputs {
2254 n,
2255 m,
2256 lmax,
2257 data_xyz: &data_xyz,
2258 centers_xyz: ¢ers_xyz,
2259 coeffs: &coeffs,
2260 kind: SphereSpectralKernelKind::Sobolev,
2261 layout: DeviceMatrixLayout::ColumnMajor,
2262 };
2263 let w = vec![1.0_f64; m];
2266 let (v, beta) = householder_reflector_from_weights(&w);
2267
2268 let b_cpu = spherical_wahba_kernel_matrix_with_kind(
2272 data_ll.view(),
2273 centers_ll.view(),
2274 penalty,
2275 false,
2276 SphereWahbaKernel::SobolevTruncated { lmax: lmax as u16 },
2277 )
2278 .expect("cpu kernel matrix");
2279 assert_cpu_kernel_matches_stable_spectral_definition(
2280 &b_cpu,
2281 &data_ll,
2282 ¢ers_ll,
2283 &coeffs,
2284 );
2285 assert_householder_fused_matches_explicit_product(&b_cpu, &v, beta);
2286
2287 if !cuda_available_for_test("householder parity") {
2288 assert_sphere_decision_declines_without_device(n, m, lmax);
2289 assert!(
2290 build_householder_constrained_design_device(inputs_raw, &v, beta).is_err(),
2291 "no CUDA runtime on this host, yet the fused Householder device entry \
2292 returned a design — the admitted-only device path fabricated a host \
2293 answer (#1551 class)"
2294 );
2295 return;
2296 }
2297 SphereGpuBackend::probe()
2300 .expect("[sphere_gpu test] backend probe must succeed on a CUDA host");
2301 let b_dev = build_kernel_matrix_device(inputs_raw.clone()).expect("raw kernel");
2302 let b = b_dev.to_host_array().expect("dtoh raw");
2303
2304 let xs_host = householder_apply_host(&b, &v, beta);
2306
2307 let xs_dev =
2308 build_householder_constrained_design_device(inputs_raw, &v, beta).expect("hh design");
2309 let xs_gpu = xs_dev.to_host_array().expect("dtoh hh");
2310
2311 let mut max_abs = 0.0_f64;
2312 for i in 0..n {
2313 for j in 0..(m - 1) {
2314 let d = (xs_host[(i, j)] - xs_gpu[(i, j)]).abs();
2315 if d > max_abs {
2316 max_abs = d;
2317 }
2318 }
2319 }
2320 assert!(
2321 max_abs < 1e-12,
2322 "Householder fused parity max |Δ| = {max_abs:.3e} >= 1e-12"
2323 );
2324 }
2325
2326 #[test]
2338 fn sphere_gpu_kernel_matrix_hill_climb_declines_without_device_else_20x_vs_cpu() {
2339 let n_lat = 500usize;
2341 let n_lon = 400usize;
2342 assert_eq!(n_lat * n_lon, 200_000);
2343 let m = 200usize;
2344 let lmax = 50usize;
2345
2346 if !cuda_available_for_test("kernel-matrix hill climb") {
2347 assert_sphere_decision_declines_without_device(n_lat * n_lon, m, lmax);
2348 return;
2349 }
2350 SphereGpuBackend::probe()
2353 .expect("[sphere_gpu hill-climb] backend probe must succeed on a CUDA host");
2354
2355 let data_ll = small_latlon_grid(n_lat, n_lon);
2357 let centers_ll =
2358 crate::basis::select_spherical_farthest_point_centers(data_ll.view(), m, false)
2359 .expect("centers");
2360 let n = data_ll.nrows();
2361 let data_xyz = latlon_to_xyz_host(data_ll.view(), false).unwrap();
2362 let centers_xyz = latlon_to_xyz_host(centers_ll.view(), false).unwrap();
2363 let penalty_order = 2usize;
2364 let coeffs = sobolev_s2_truncated_coefficients(lmax, penalty_order);
2365
2366 let inputs_warm = S2KernelBuildInputs {
2368 n,
2369 m,
2370 lmax,
2371 data_xyz: &data_xyz,
2372 centers_xyz: ¢ers_xyz,
2373 coeffs: &coeffs,
2374 kind: SphereSpectralKernelKind::Sobolev,
2375 layout: DeviceMatrixLayout::ColumnMajor,
2376 };
2377 {
2382 let warm = build_kernel_matrix_device(inputs_warm.clone()).expect("warmup");
2383 drop(warm.to_host_array().expect("warmup to_host"));
2384 }
2385
2386 let t0 = std::time::Instant::now();
2388 let dev = build_kernel_matrix_device(inputs_warm.clone()).expect("gpu kernel matrix");
2389 dev.to_host_array().expect("dtoh");
2390 let gpu_secs = t0.elapsed().as_secs_f64();
2391
2392 let t1 = std::time::Instant::now();
2399 crate::basis::spherical_wahba_kernel_matrix_cpu(
2400 data_ll.view(),
2401 centers_ll.view(),
2402 penalty_order,
2403 false,
2404 SphereWahbaKernel::SobolevTruncated { lmax: lmax as u16 },
2405 )
2406 .expect("cpu kernel matrix");
2407 let cpu_secs = t1.elapsed().as_secs_f64();
2408
2409 let ratio = cpu_secs / gpu_secs.max(1e-9);
2410 eprintln!(
2411 "[sphere_gpu hill-climb] n={n} m={m} L={lmax} cpu={cpu_secs:.3}s gpu={gpu_secs:.3}s ratio={ratio:.2}x"
2412 );
2413 assert!(
2420 ratio >= 3.0,
2421 "GPU kernel matrix only {ratio:.2}× faster than CPU (dispatch-worthiness ≥ 3×) at \
2422 n={n} m={m} L={lmax}: cpu={cpu_secs:.3}s gpu={gpu_secs:.3}s"
2423 );
2424 }
2425
2426 #[test]
2459 fn sphere_gpu_end_to_end_fit_dispatches_to_device_else_declines() {
2460 use crate::basis::{
2461 CenterStrategy, SphereMethod, SphericalSplineBasisSpec, SphericalSplineIdentifiability,
2462 build_spherical_spline_basis,
2463 };
2464
2465 let n_lat = 500usize;
2466 let n_lon = 400usize;
2467 let m: usize = 200;
2468 let lmax: u16 = 50;
2469
2470 if !cuda_available_for_test("end-to-end fit dispatch") {
2471 assert_sphere_decision_declines_without_device(n_lat * n_lon, m, lmax as usize);
2472 return;
2473 }
2474 SphereGpuBackend::probe()
2477 .expect("[sphere_gpu end-to-end dispatch] backend probe must succeed on a CUDA host");
2478
2479 let admit = sphere_kernel_decision(5_000, m, lmax as usize)
2486 .expect("the sphere GPU decision must not fault on a CUDA host");
2487 assert!(
2488 admit.use_gpu,
2489 "a CUDA device is present and (n=5000, m={m}) is exactly at the sphere device-work \
2490 crossover, yet the dispatch decision kept it on the host — reason={}",
2491 admit.reason
2492 );
2493 let refuse = sphere_kernel_decision(4_999, m, lmax as usize)
2494 .expect("the sphere GPU decision must not fault on a CUDA host");
2495 assert!(
2496 !refuse.use_gpu,
2497 "(n=4999, m={m}) is one row below the sphere device-work crossover, yet the dispatch \
2498 decision admitted the device — reason={}",
2499 refuse.reason
2500 );
2501
2502 let data_ll = small_latlon_grid(n_lat, n_lon);
2503 let spec_gpu = SphericalSplineBasisSpec {
2504 center_strategy: CenterStrategy::FarthestPoint { num_centers: m },
2505 penalty_order: 2,
2506 double_penalty: false,
2507 radians: false,
2508 method: SphereMethod::Wahba,
2509 max_degree: None,
2510 wahba_kernel: SphereWahbaKernel::SobolevTruncated { lmax },
2511 identifiability: SphericalSplineIdentifiability::CenterSumToZero,
2512 };
2513
2514 let t0 = std::time::Instant::now();
2520 let built = build_spherical_spline_basis(data_ll.view(), &spec_gpu)
2521 .expect("the end-to-end sphere build must succeed on a CUDA host");
2522 let build_secs = t0.elapsed().as_secs_f64();
2523
2524 assert_eq!(
2525 built.design.nrows(),
2526 data_ll.nrows(),
2527 "the device-dispatched sphere build returned {} design rows for {} data rows",
2528 built.design.nrows(),
2529 data_ll.nrows()
2530 );
2531 assert!(
2532 built.design.ncols() > 0 && built.design.ncols() <= m,
2533 "the device-dispatched sphere build returned {} design columns for m={m} centers",
2534 built.design.ncols()
2535 );
2536 let beta = ndarray::Array1::<f64>::ones(built.design.ncols());
2540 for row in (0..built.design.nrows()).step_by(data_ll.nrows() / 64 + 1) {
2541 let value = built.design.dot_row(row, &beta);
2542 assert!(
2543 value.is_finite(),
2544 "the device-dispatched sphere design row {row} sums to {value}, not a finite number"
2545 );
2546 }
2547
2548 eprintln!(
2553 "[sphere_gpu end-to-end dispatch] n={} m={m} L={lmax} build={build_secs:.3}s reason={}",
2554 data_ll.nrows(),
2555 admit.reason
2556 );
2557 }
2558
2559 #[test]
2581 fn sphere_gpu_end_to_end_fit_parity_vs_cpu_truncated() {
2582 use crate::basis::{
2583 select_spherical_farthest_point_centers, spherical_wahba_kernel_matrix_with_kind,
2584 };
2585 use faer::Side;
2586 use gam_linalg::faer_ndarray::FaerCholesky;
2587
2588 let data_ll = small_latlon_grid(25, 40);
2590 assert_eq!(data_ll.nrows(), 1000);
2591 let n = data_ll.nrows();
2592 let m: usize = 80;
2593 let lmax_u16: u16 = 15;
2594 let lmax: usize = lmax_u16 as usize;
2595 let penalty_order: usize = 2;
2596 let kernel = SphereWahbaKernel::SobolevTruncated { lmax: lmax_u16 };
2597 let lambda: f64 = 1.0e-3;
2598
2599 let centers_ll = select_spherical_farthest_point_centers(data_ll.view(), m, false)
2601 .expect("farthest-point centers");
2602 assert_eq!(centers_ll.nrows(), m);
2603
2604 let z = Array2::<f64>::eye(centers_ll.nrows());
2607 let p = z.ncols();
2608 assert_eq!(p, m);
2609
2610 let k_cc = spherical_wahba_kernel_matrix_with_kind(
2615 centers_ll.view(),
2616 centers_ll.view(),
2617 penalty_order,
2618 false,
2619 kernel,
2620 )
2621 .expect("centers×centers kernel");
2622 let s_full = z.t().dot(&k_cc).dot(&z);
2623
2624 let raw_design_cpu = spherical_wahba_kernel_matrix_with_kind(
2626 data_ll.view(),
2627 centers_ll.view(),
2628 penalty_order,
2629 false,
2630 kernel,
2631 )
2632 .expect("CPU raw design");
2633 let x_s_cpu = raw_design_cpu.dot(&z);
2634
2635 let data_xyz = latlon_to_xyz_host(data_ll.view(), false).expect("data xyz");
2637 let centers_xyz = latlon_to_xyz_host(centers_ll.view(), false).expect("centers xyz");
2638 let coeffs = crate::basis::sobolev_s2_truncated_coefficients(lmax, penalty_order);
2639 let inputs = S2KernelBuildInputs {
2640 n,
2641 m,
2642 lmax,
2643 data_xyz: &data_xyz,
2644 centers_xyz: ¢ers_xyz,
2645 coeffs: &coeffs,
2646 kind: SphereSpectralKernelKind::Sobolev,
2647 layout: DeviceMatrixLayout::ColumnMajor,
2648 };
2649 let mut y = ndarray::Array1::<f64>::zeros(n);
2655 for i in 0..n {
2656 let lat_rad = data_ll[(i, 0)].to_radians();
2657 let lon_rad = data_ll[(i, 1)].to_radians();
2658 y[i] = (2.0 * lat_rad).sin() * (3.0 * lon_rad).cos()
2660 + 0.25 * lat_rad.cos() * (5.0 * lon_rad).sin();
2661 }
2662
2663 let solve_penalised = |x_s: &ndarray::Array2<f64>| -> ndarray::Array1<f64> {
2668 let xtx = x_s.t().dot(x_s);
2669 let mut a = xtx;
2670 for i in 0..p {
2671 for j in 0..p {
2672 a[(i, j)] += lambda * s_full[(i, j)];
2673 }
2674 }
2675 let rhs = x_s.t().dot(&y);
2676 let factor = a
2677 .cholesky(Side::Lower)
2678 .expect("penalised normal equations are SPD under λ > 0");
2679 factor.solvevec(&rhs)
2680 };
2681
2682 let beta_cpu = solve_penalised(&x_s_cpu);
2683 assert_eq!(beta_cpu.len(), p);
2684 let yhat_cpu = x_s_cpu.dot(&beta_cpu);
2685 assert_eq!(x_s_cpu.dim(), (n, p));
2686
2687 assert_cpu_kernel_matches_stable_spectral_definition(
2692 &raw_design_cpu,
2693 &data_ll,
2694 ¢ers_ll,
2695 &coeffs,
2696 );
2697 {
2698 let mut a = x_s_cpu.t().dot(&x_s_cpu);
2699 for i in 0..p {
2700 for j in 0..p {
2701 a[(i, j)] += lambda * s_full[(i, j)];
2702 }
2703 }
2704 let residual = a.dot(&beta_cpu) - x_s_cpu.t().dot(&y);
2705 let rhs_scale = x_s_cpu
2706 .t()
2707 .dot(&y)
2708 .iter()
2709 .fold(0.0_f64, |acc, v| acc.max(v.abs()))
2710 .max(1.0);
2711 let max_residual = residual.iter().fold(0.0_f64, |acc, v| acc.max(v.abs()));
2712 assert!(
2713 max_residual <= 1e-9 * rhs_scale,
2714 "CPU penalised normal equations not solved: ‖(XᵀX + λS)β − Xᵀy‖∞ = \
2715 {max_residual:.3e} (rhs scale {rhs_scale:.3e})"
2716 );
2717 }
2718
2719 if !cuda_available_for_test("end-to-end fit parity") {
2720 assert_sphere_decision_declines_without_device(n, m, lmax);
2721 assert_device_kernel_entry_refuses(inputs);
2722 return;
2723 }
2724 SphereGpuBackend::probe()
2727 .expect("[sphere gpu parity] sphere GPU backend probe must succeed on a CUDA host");
2728 let raw_dev = build_kernel_matrix_device(inputs).expect("GPU raw design");
2729 let raw_design_gpu = raw_dev.to_host_array().expect("dtoh GPU raw design");
2730 let x_s_gpu = raw_design_gpu.dot(&z);
2731
2732 assert_eq!(x_s_gpu.dim(), (n, p));
2733
2734 let mut raw_xs_delta = 0.0_f64;
2744 let mut xs_scale = 0.0_f64;
2745 for (a, b) in x_s_cpu.iter().zip(x_s_gpu.iter()) {
2746 raw_xs_delta = raw_xs_delta.max((a - b).abs());
2747 xs_scale = xs_scale.max(a.abs());
2748 }
2749 let cond = {
2752 use gam_linalg::faer_ndarray::FaerEigh;
2753 let xtx = x_s_cpu.t().dot(&x_s_cpu);
2754 let mut a = xtx;
2755 for i in 0..p {
2756 for j in 0..p {
2757 a[(i, j)] += lambda * s_full[(i, j)];
2758 }
2759 }
2760 let (mut lo, mut hi) = (f64::INFINITY, 0.0_f64);
2761 if let Ok((vals, _)) = a.eigh(faer::Side::Lower) {
2762 for &v in vals.iter() {
2763 lo = lo.min(v);
2764 hi = hi.max(v);
2765 }
2766 }
2767 hi / lo.max(1e-300)
2768 };
2769 assert!(
2774 raw_xs_delta <= 1e-12 * xs_scale.max(1.0),
2775 "GPU vs CPU sphere design matrix max |Δ| = {raw_xs_delta:.3e} > {:.3e} \
2776 (scale {xs_scale:.3e}) — the kernel itself drifted (this is the genuine \
2777 GPU output, NOT a conditioning artifact)",
2778 1e-12 * xs_scale.max(1.0)
2779 );
2780
2781 let beta_gpu = solve_penalised(&x_s_gpu);
2782 assert_eq!(beta_gpu.len(), p);
2783
2784 let yhat_gpu = x_s_gpu.dot(&beta_gpu);
2788
2789 let mut max_beta_delta = 0.0_f64;
2790 for k in 0..p {
2791 let d = (beta_cpu[k] - beta_gpu[k]).abs();
2792 if d > max_beta_delta {
2793 max_beta_delta = d;
2794 }
2795 }
2796 let mut max_fit_delta = 0.0_f64;
2797 for i in 0..n {
2798 let d = (yhat_cpu[i] - yhat_gpu[i]).abs();
2799 if d > max_fit_delta {
2800 max_fit_delta = d;
2801 }
2802 }
2803
2804 eprintln!(
2805 "[sphere_gpu fit parity] n={n} m={m} p={p} lmax={lmax} λ={lambda:.1e} \
2806 raw_xs|Δ|={raw_xs_delta:.3e} cond={cond:.3e} \
2807 max|Δβ|={max_beta_delta:.3e} max|Δŷ|={max_fit_delta:.3e}"
2808 );
2809
2810 assert!(
2817 max_fit_delta <= 1.0e-9,
2818 "GPU vs CPU truncated-spectral fitted-value max |Δ| = {max_fit_delta:.3e} > 1e-9"
2819 );
2820
2821 let beta_tol = (1e-15 * cond * (1.0 + xs_scale)).max(1e-9) * 16.0;
2833 assert!(
2834 max_beta_delta <= beta_tol,
2835 "GPU vs CPU truncated-spectral coefficient max |Δ| = {max_beta_delta:.3e} > \
2836 condition-aware tol {beta_tol:.3e} (cond={cond:.3e}). Raw design parity is \
2837 {raw_xs_delta:.3e}; a drift THIS much larger than cond·ULP is a real solve/kernel \
2838 mismatch, not conditioning."
2839 );
2840 }
2841}