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 unsafe { cudarc::driver::result::free_host(self.ptr as *mut std::ffi::c_void) }.ok();
342 }
343}
344
345#[cfg(target_os = "linux")]
351unsafe impl Send for PinnedF64 {}
352
353#[cfg(target_os = "linux")]
362const PINNED_POOL_MAX_BUFFERS: usize = 4;
363
364#[cfg(target_os = "linux")]
365static PINNED_POOL: OnceLock<Mutex<Vec<PinnedF64>>> = OnceLock::new();
366
367#[cfg(target_os = "linux")]
371struct PinnedLease {
372 buf: Option<PinnedF64>,
373}
374
375#[cfg(target_os = "linux")]
376impl PinnedLease {
377 fn acquire(ctx: &Arc<CudaContext>, len: usize) -> Result<Self, GpuError> {
380 let pool = PINNED_POOL.get_or_init(|| Mutex::new(Vec::new()));
381 if let Ok(mut guard) = pool.lock() {
382 if let Some(pos) = guard.iter().position(|b| b.len == len) {
383 return Ok(Self {
384 buf: Some(guard.swap_remove(pos)),
385 });
386 }
387 }
388 Ok(Self {
389 buf: Some(PinnedF64::alloc(ctx, len)?),
390 })
391 }
392
393 fn as_mut_slice(&mut self) -> &mut [f64] {
394 self.buf
395 .as_mut()
396 .expect("PinnedLease buffer present until drop")
397 .as_mut_slice()
398 }
399
400 fn as_slice(&self) -> &[f64] {
401 self.buf
402 .as_ref()
403 .expect("PinnedLease buffer present until drop")
404 .as_slice()
405 }
406}
407
408#[cfg(target_os = "linux")]
409impl Drop for PinnedLease {
410 fn drop(&mut self) {
411 let Some(buf) = self.buf.take() else {
412 return;
413 };
414 if let Some(pool) = PINNED_POOL.get() {
415 if let Ok(mut guard) = pool.lock() {
416 if guard.len() < PINNED_POOL_MAX_BUFFERS {
417 guard.push(buf);
418 return;
419 }
420 guard.remove(0);
424 guard.push(buf);
425 return;
426 }
427 }
428 drop(buf);
430 }
431}
432
433#[derive(Clone, Debug)]
444pub struct S2KernelBuildInputs<'a> {
445 pub n: usize,
446 pub m: usize,
447 pub lmax: usize,
448 pub data_xyz: &'a [f64],
449 pub centers_xyz: &'a [f64],
450 pub coeffs: &'a [f64],
451 pub kind: SphereSpectralKernelKind,
452 pub layout: DeviceMatrixLayout,
453}
454
455impl<'a> S2KernelBuildInputs<'a> {
456 fn validate(&self) -> Result<(), GpuError> {
457 if self.lmax == 0 {
458 return Err(GpuError::DriverCallFailed {
459 reason: "S2KernelBuildInputs: lmax must be >= 1".into(),
460 });
461 }
462 if self.data_xyz.len() != 3 * self.n {
463 gam_gpu::gpu_bail!(
464 "S2KernelBuildInputs: data_xyz.len()={} != 3*n={}",
465 self.data_xyz.len(),
466 3 * self.n
467 );
468 }
469 if self.centers_xyz.len() != 3 * self.m {
470 gam_gpu::gpu_bail!(
471 "S2KernelBuildInputs: centers_xyz.len()={} != 3*m={}",
472 self.centers_xyz.len(),
473 3 * self.m
474 );
475 }
476 if self.coeffs.len() != self.lmax + 1 {
477 gam_gpu::gpu_bail!(
478 "S2KernelBuildInputs: coeffs.len()={} != lmax+1={}",
479 self.coeffs.len(),
480 self.lmax + 1
481 );
482 }
483 if self.coeffs[0] != 0.0 {
484 return Err(GpuError::DriverCallFailed {
485 reason: "S2KernelBuildInputs: coeffs[0] must be 0 (mean-zero kernel)".into(),
486 });
487 }
488 Ok(())
489 }
490}
491
492#[cfg(target_os = "linux")]
502const KERNEL_TEMPLATE: &str = r#"
503// LMAX is supplied by the host via a `#define LMAX ...` prepended to
504// this source before NVRTC compilation (see `SphereGpuBackend::module_for`).
505extern "C" __global__
506__launch_bounds__(256)
507void s2_wahba_legendre_colmajor(
508 const double* __restrict__ data_xyz, // n × 3 (row-major flat)
509 const double* __restrict__ centers_xyz, // m × 3 (row-major flat)
510 const double* __restrict__ coeffs, // length LMAX + 1, coeffs[0] = 0
511 int n,
512 int m,
513 long long ld,
514 double* __restrict__ out // ld × m column-major
515) {
516 const int i = blockIdx.y * blockDim.y + threadIdx.y;
517 const int j = blockIdx.x * blockDim.x + threadIdx.x;
518 if (i >= n || j >= m) return;
519
520 // Load (x_i, y_i, z_i) and (cx_j, cy_j, cz_j) into registers.
521 const double xi = data_xyz[3 * i + 0];
522 const double yi = data_xyz[3 * i + 1];
523 const double zi = data_xyz[3 * i + 2];
524 const double cxj = centers_xyz[3 * j + 0];
525 const double cyj = centers_xyz[3 * j + 1];
526 const double czj = centers_xyz[3 * j + 2];
527
528 // t = clamp(x_i · z_j, -1, +1).
529 double t = fma(xi, cxj, fma(yi, cyj, zi * czj));
530 if (t > 1.0) t = 1.0;
531 if (t < -1.0) t = -1.0;
532
533 // Legendre 3-term recurrence in registers.
534 // P_0(t) = 1, P_1(t) = t.
535 double p_prev = 1.0;
536 double p_curr = t;
537 double acc = coeffs[0] * p_prev + coeffs[1] * p_curr;
538
539 #pragma unroll 8
540 for (int ell = 1; ell < LMAX; ++ell) {
541 const double lf = (double) ell;
542 const double inv = 1.0 / (lf + 1.0);
543 // p_{ell+1} = ((2ell+1) * t * p_curr - ell * p_prev) / (ell+1)
544 const double p_next =
545 fma((2.0 * lf + 1.0) * t, p_curr, -lf * p_prev) * inv;
546 acc = fma(coeffs[ell + 1], p_next, acc);
547 p_prev = p_curr;
548 p_curr = p_next;
549 }
550
551 out[(long long) j * ld + (long long) i] = acc;
552}
553
554// Fused Householder-constrained kernel (Phase 3). Z = I - beta · v · v^T,
555// the constrained design is X_s = B[:, 1..m] - beta * (B · v) · v[1..m]^T,
556// i.e. drop the first column after applying Z. Each thread computes one
557// row of B in registers (m kernel evaluations), forms d_i = B_row · v,
558// then emits X_s[i, j_out] = B_row[j_out + 1] - beta * d_i * v[j_out + 1]
559// for j_out in 0..m-1.
560//
561// Grid: 1D over rows (block_dim.x rows per block). Each thread iterates
562// over centers in an inner loop — register-bound by the per-row state
563// (xyz_i, p_prev, p_curr, acc, and a small per-center scratch).
564extern "C" __global__
565__launch_bounds__(128)
566void s2_wahba_householder_constrained_colmajor(
567 const double* __restrict__ data_xyz, // n × 3
568 const double* __restrict__ centers_xyz, // m × 3
569 const double* __restrict__ coeffs, // length LMAX + 1
570 const double* __restrict__ v, // length m, Householder vector
571 double beta,
572 int n,
573 int m,
574 long long ld_out,
575 double* __restrict__ out // ld_out × (m-1) column-major
576) {
577 const int i = blockIdx.x * blockDim.x + threadIdx.x;
578 if (i >= n) return;
579
580 const double xi = data_xyz[3 * i + 0];
581 const double yi = data_xyz[3 * i + 1];
582 const double zi = data_xyz[3 * i + 2];
583
584 // Pass 1: compute d_i = sum_j v[j] * B[i, j].
585 double d_i = 0.0;
586 for (int j = 0; j < m; ++j) {
587 const double cxj = centers_xyz[3 * j + 0];
588 const double cyj = centers_xyz[3 * j + 1];
589 const double czj = centers_xyz[3 * j + 2];
590 double t = fma(xi, cxj, fma(yi, cyj, zi * czj));
591 if (t > 1.0) t = 1.0;
592 if (t < -1.0) t = -1.0;
593
594 double p_prev = 1.0;
595 double p_curr = t;
596 double acc = coeffs[0] * p_prev + coeffs[1] * p_curr;
597 #pragma unroll 8
598 for (int ell = 1; ell < LMAX; ++ell) {
599 const double lf = (double) ell;
600 const double inv = 1.0 / (lf + 1.0);
601 const double p_next =
602 fma((2.0 * lf + 1.0) * t, p_curr, -lf * p_prev) * inv;
603 acc = fma(coeffs[ell + 1], p_next, acc);
604 p_prev = p_curr;
605 p_curr = p_next;
606 }
607 d_i = fma(v[j], acc, d_i);
608 }
609
610 // Pass 2: emit X_s[i, j_out] = B[i, j_out+1] - beta * d_i * v[j_out+1].
611 const double bd = beta * d_i;
612 for (int j_out = 0; j_out < m - 1; ++j_out) {
613 const int j = j_out + 1;
614 const double cxj = centers_xyz[3 * j + 0];
615 const double cyj = centers_xyz[3 * j + 1];
616 const double czj = centers_xyz[3 * j + 2];
617 double t = fma(xi, cxj, fma(yi, cyj, zi * czj));
618 if (t > 1.0) t = 1.0;
619 if (t < -1.0) t = -1.0;
620
621 double p_prev = 1.0;
622 double p_curr = t;
623 double acc = coeffs[0] * p_prev + coeffs[1] * p_curr;
624 #pragma unroll 8
625 for (int ell = 1; ell < LMAX; ++ell) {
626 const double lf = (double) ell;
627 const double inv = 1.0 / (lf + 1.0);
628 const double p_next =
629 fma((2.0 * lf + 1.0) * t, p_curr, -lf * p_prev) * inv;
630 acc = fma(coeffs[ell + 1], p_next, acc);
631 p_prev = p_curr;
632 p_curr = p_next;
633 }
634 const double xs = acc - bd * v[j];
635 out[(long long) j_out * ld_out + (long long) i] = xs;
636 }
637}
638"#;
639
640#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
650pub struct S2ModuleCacheKey {
651 pub cc_major: i32,
652 pub cc_minor: i32,
653 pub lmax: u32,
654 pub kind: SphereSpectralKernelKind,
655 pub layout: DeviceMatrixLayout,
656}
657
658pub const fn sphere_gpu_compiled() -> bool {
661 cfg!(target_os = "linux")
662}
663
664#[must_use]
671pub fn sphere_kernel_decision(n: usize, m: usize, lmax: usize) -> Result<GpuDecision, GpuError> {
672 let large_enough = match gam_gpu::device_runtime::GpuRuntime::resolve(gam_gpu::global_policy())?
673 {
674 Some(runtime) => {
675 let ld = ((n + 31) / 32) * 32;
676 let needed_bytes = ld
677 .saturating_mul(m)
678 .saturating_mul(std::mem::size_of::<f64>());
679 let budget = runtime.memory_budget_bytes;
680 n.saturating_mul(m) >= 1_000_000 && lmax <= 200 && needed_bytes <= budget
681 }
682 None => false,
683 };
684 decide(
685 GpuKernel::SpatialKernelOperator,
686 gam_gpu::GpuEligibility::from_flags(sphere_gpu_compiled(), large_enough),
687 )
688}
689
690#[must_use]
696pub fn truncated_device_kind(
697 kernel: crate::basis::SphereWahbaKernel,
698) -> Option<(SphereSpectralKernelKind, u16)> {
699 use crate::basis::SphereWahbaKernel;
700 match kernel {
701 SphereWahbaKernel::SobolevTruncated { lmax } => {
702 Some((SphereSpectralKernelKind::Sobolev, lmax))
703 }
704 SphereWahbaKernel::PseudoTruncated { lmax } => {
705 Some((SphereSpectralKernelKind::Pseudo, lmax))
706 }
707 SphereWahbaKernel::Sobolev | SphereWahbaKernel::Pseudo => None,
708 }
709}
710
711pub fn try_build_truncated_kernel_matrix_gpu(
733 data: ArrayView2<'_, f64>,
734 centers: ArrayView2<'_, f64>,
735 penalty_order: usize,
736 radians: bool,
737 kernel: crate::basis::SphereWahbaKernel,
738) -> Option<Result<Array2<f64>, GpuError>> {
739 let (kind, lmax) = truncated_device_kind(kernel)?;
740 let n = data.nrows();
741 let m = centers.nrows();
742 if n == 0 || m == 0 || lmax == 0 {
743 return None;
744 }
745 let decision = match sphere_kernel_decision(n, m, lmax as usize) {
746 Ok(decision) => decision,
747 Err(error) => return Some(Err(error)),
748 };
749 if !decision.use_gpu {
750 return None;
753 }
754 Some(build_truncated_kernel_matrix_gpu_admitted(
756 data,
757 centers,
758 penalty_order,
759 radians,
760 kind,
761 lmax,
762 ))
763}
764
765fn build_truncated_kernel_matrix_gpu_admitted(
769 data: ArrayView2<'_, f64>,
770 centers: ArrayView2<'_, f64>,
771 penalty_order: usize,
772 radians: bool,
773 kind: SphereSpectralKernelKind,
774 lmax: u16,
775) -> Result<Array2<f64>, GpuError> {
776 let n = data.nrows();
777 let m = centers.nrows();
778 let data_xyz = latlon_to_xyz_host(data, radians)
779 .map_err(|reason| GpuError::DriverCallFailed { reason })?;
780 let centers_xyz = latlon_to_xyz_host(centers, radians)
781 .map_err(|reason| GpuError::DriverCallFailed { reason })?;
782 let coeffs = kind.coefficients(lmax as usize, penalty_order);
786 let inputs = S2KernelBuildInputs {
787 n,
788 m,
789 lmax: lmax as usize,
790 data_xyz: &data_xyz,
791 centers_xyz: ¢ers_xyz,
792 coeffs: &coeffs,
793 kind,
794 layout: DeviceMatrixLayout::ColumnMajor,
795 };
796 let device_matrix = build_kernel_matrix_device(inputs)?;
797 let out = device_matrix.to_host_array()?;
798 if !out.sum().is_finite() {
809 return Err(GpuError::DriverCallFailed {
810 reason: "sphere GPU truncated kernel produced a non-finite value".to_string(),
811 });
812 }
813 Ok(out)
814}
815
816#[cfg(target_os = "linux")]
817struct SphereGpuContext {
818 ctx: Arc<CudaContext>,
819 stream: Arc<CudaStream>,
820 modules: Mutex<HashMap<S2ModuleCacheKey, Arc<CudaModule>>>,
821 cc_major: i32,
822 cc_minor: i32,
823}
824
825pub struct SphereGpuBackend {
828 #[cfg(target_os = "linux")]
829 inner: SphereGpuContext,
830}
831
832impl SphereGpuBackend {
833 pub fn probe() -> Result<&'static Self, GpuError> {
835 static BACKEND: OnceLock<Result<SphereGpuBackend, GpuError>> = OnceLock::new();
836 BACKEND
837 .get_or_init(|| {
838 #[cfg(target_os = "linux")]
839 {
840 Self::probe_linux()
841 }
842 #[cfg(not(target_os = "linux"))]
843 {
844 Err(GpuError::DriverLibraryUnavailable {
845 reason: "sphere GPU backend is Linux-only".to_string(),
846 })
847 }
848 })
849 .as_ref()
850 .map_err(GpuError::clone)
851 }
852
853 #[cfg(target_os = "linux")]
854 fn probe_linux() -> Result<Self, GpuError> {
855 let parts = gam_gpu::backend_probe::probe_cuda_backend("sphere")?;
856 Ok(SphereGpuBackend {
857 inner: SphereGpuContext {
858 ctx: parts.ctx,
859 stream: parts.stream,
860 modules: Mutex::new(HashMap::new()),
861 cc_major: parts.capability.compute_major,
862 cc_minor: parts.capability.compute_minor,
863 },
864 })
865 }
866
867 #[cfg(target_os = "linux")]
870 fn module_for(&self, key: S2ModuleCacheKey) -> Result<Arc<CudaModule>, GpuError> {
871 if let Ok(guard) = self.inner.modules.lock() {
872 if let Some(existing) = guard.get(&key) {
873 return Ok(existing.clone());
874 }
875 }
876 let src = format!("#define LMAX {}\n{}", key.lmax, KERNEL_TEMPLATE);
886 let ptx = gam_gpu::device_cache::compile_ptx_arch(&src).gpu_ctx_with(|err| {
887 format!(
888 "sphere NVRTC compile (kind={}, lmax={}): {err}",
889 key.kind.tag(),
890 key.lmax
891 )
892 })?;
893 let module = self
894 .inner
895 .ctx
896 .load_module(ptx)
897 .gpu_ctx("sphere module load")?;
898 if let Ok(mut guard) = self.inner.modules.lock() {
899 guard.entry(key).or_insert_with(|| module.clone());
900 }
901 Ok(module)
902 }
903
904 #[cfg(target_os = "linux")]
905 fn cc(&self) -> (i32, i32) {
906 (self.inner.cc_major, self.inner.cc_minor)
907 }
908}
909
910pub fn build_kernel_matrix_device(
917 inputs: S2KernelBuildInputs<'_>,
918) -> Result<DeviceS2KernelMatrix, GpuError> {
919 inputs.validate()?;
920
921 #[cfg(target_os = "linux")]
922 {
923 use cudarc::driver::{LaunchConfig, PushKernelArg};
924 let backend = SphereGpuBackend::probe()?;
925 let (cc_major, cc_minor) = backend.cc();
926 let key = S2ModuleCacheKey {
927 cc_major,
928 cc_minor,
929 lmax: inputs.lmax as u32,
930 kind: inputs.kind,
931 layout: inputs.layout,
932 };
933 let module = backend.module_for(key)?;
934 let func = module
935 .load_function("s2_wahba_legendre_colmajor")
936 .gpu_ctx("sphere load_function raw")?;
937 let stream = backend.inner.stream.clone();
938
939 let data_dev = stream
940 .clone_htod(inputs.data_xyz)
941 .gpu_ctx("sphere htod data_xyz")?;
942 let centers_dev = stream
943 .clone_htod(inputs.centers_xyz)
944 .gpu_ctx("sphere htod centers_xyz")?;
945 let coeffs_dev = stream
946 .clone_htod(inputs.coeffs)
947 .gpu_ctx("sphere htod coeffs")?;
948
949 let n = inputs.n;
950 let m = inputs.m;
951 let ld = ((n + 31) / 32) * 32;
952 let mut out_dev = stream
953 .alloc_zeros::<f64>(ld * m)
954 .gpu_ctx_with(|err| format!("sphere alloc out (ld={ld}, m={m}): {err}"))?;
955
956 let block_x: u32 = 32;
958 let block_y: u32 = 8;
959 let grid_x: u32 = ((m as u32) + block_x - 1) / block_x;
960 let grid_y: u32 = ((n as u32) + block_y - 1) / block_y;
961 let cfg = LaunchConfig {
962 grid_dim: (grid_x, grid_y, 1),
963 block_dim: (block_x, block_y, 1),
964 shared_mem_bytes: 0,
965 };
966 let n_i32: i32 =
967 i32::try_from(n).map_err(|_| gam_gpu::gpu_err!("sphere n={n} overflows i32"))?;
968 let m_i32: i32 =
969 i32::try_from(m).map_err(|_| gam_gpu::gpu_err!("sphere m={m} overflows i32"))?;
970 let ld_i64: i64 = ld as i64;
971
972 let mut builder = stream.launch_builder(&func);
973 builder
974 .arg(&data_dev)
975 .arg(¢ers_dev)
976 .arg(&coeffs_dev)
977 .arg(&n_i32)
978 .arg(&m_i32)
979 .arg(&ld_i64)
980 .arg(&mut out_dev);
981 unsafe { builder.launch(cfg) }.gpu_ctx("sphere raw kernel launch")?;
986 stream
987 .synchronize()
988 .gpu_ctx("sphere raw kernel synchronize")?;
989
990 Ok(DeviceS2KernelMatrix {
991 rows: n,
992 cols: m,
993 ld,
994 col_major_dev: out_dev,
995 stream,
996 })
997 }
998
999 #[cfg(not(target_os = "linux"))]
1000 {
1001 Err(GpuError::DriverLibraryUnavailable {
1002 reason: "sphere GPU backend is Linux-only".to_string(),
1003 })
1004 }
1005}
1006
1007pub fn build_householder_constrained_design_device(
1011 inputs: S2KernelBuildInputs<'_>,
1012 v: &[f64],
1013 beta: f64,
1014) -> Result<DeviceS2KernelMatrix, GpuError> {
1015 inputs.validate()?;
1016 if v.len() != inputs.m {
1017 gam_gpu::gpu_bail!(
1018 "build_householder_constrained_design_device: v.len()={} != m={}",
1019 v.len(),
1020 inputs.m
1021 );
1022 }
1023 if inputs.m < 2 {
1024 gam_gpu::gpu_bail!(
1025 "build_householder_constrained_design_device: m must be >= 2 (got {})",
1026 inputs.m
1027 );
1028 }
1029 if !beta.is_finite() {
1030 gam_gpu::gpu_bail!(
1031 "build_householder_constrained_design_device: beta must be finite (got {beta})"
1032 );
1033 }
1034
1035 #[cfg(target_os = "linux")]
1036 {
1037 use cudarc::driver::{LaunchConfig, PushKernelArg};
1038 let backend = SphereGpuBackend::probe()?;
1039 let (cc_major, cc_minor) = backend.cc();
1040 let key = S2ModuleCacheKey {
1041 cc_major,
1042 cc_minor,
1043 lmax: inputs.lmax as u32,
1044 kind: inputs.kind,
1045 layout: inputs.layout,
1046 };
1047 let module = backend.module_for(key)?;
1048 let func = module
1049 .load_function("s2_wahba_householder_constrained_colmajor")
1050 .gpu_ctx("sphere load_function householder")?;
1051 let stream = backend.inner.stream.clone();
1052
1053 let data_dev = stream
1054 .clone_htod(inputs.data_xyz)
1055 .gpu_ctx("sphere-hh htod data_xyz")?;
1056 let centers_dev = stream
1057 .clone_htod(inputs.centers_xyz)
1058 .gpu_ctx("sphere-hh htod centers_xyz")?;
1059 let coeffs_dev = stream
1060 .clone_htod(inputs.coeffs)
1061 .gpu_ctx("sphere-hh htod coeffs")?;
1062 let v_dev = stream.clone_htod(v).gpu_ctx("sphere-hh htod v")?;
1063
1064 let n = inputs.n;
1065 let m = inputs.m;
1066 let cols_out = m - 1;
1067 let ld_out = ((n + 31) / 32) * 32;
1068 let mut out_dev = stream
1069 .alloc_zeros::<f64>(ld_out * cols_out)
1070 .gpu_ctx_with(|err| {
1071 format!("sphere-hh alloc out (ld={ld_out}, cols={cols_out}): {err}")
1072 })?;
1073
1074 let block_x: u32 = 128;
1075 let grid_x: u32 = ((n as u32) + block_x - 1) / block_x;
1076 let cfg = LaunchConfig {
1077 grid_dim: (grid_x, 1, 1),
1078 block_dim: (block_x, 1, 1),
1079 shared_mem_bytes: 0,
1080 };
1081 let n_i32: i32 =
1082 i32::try_from(n).map_err(|_| gam_gpu::gpu_err!("sphere-hh n={n} overflows i32"))?;
1083 let m_i32: i32 =
1084 i32::try_from(m).map_err(|_| gam_gpu::gpu_err!("sphere-hh m={m} overflows i32"))?;
1085 let ld_out_i64: i64 = ld_out as i64;
1086
1087 let mut builder = stream.launch_builder(&func);
1088 builder
1089 .arg(&data_dev)
1090 .arg(¢ers_dev)
1091 .arg(&coeffs_dev)
1092 .arg(&v_dev)
1093 .arg(&beta)
1094 .arg(&n_i32)
1095 .arg(&m_i32)
1096 .arg(&ld_out_i64)
1097 .arg(&mut out_dev);
1098 unsafe { builder.launch(cfg) }.gpu_ctx("sphere-hh kernel launch")?;
1101 stream
1102 .synchronize()
1103 .gpu_ctx("sphere-hh kernel synchronize")?;
1104
1105 Ok(DeviceS2KernelMatrix {
1106 rows: n,
1107 cols: cols_out,
1108 ld: ld_out,
1109 col_major_dev: out_dev,
1110 stream,
1111 })
1112 }
1113
1114 #[cfg(not(target_os = "linux"))]
1115 {
1116 Err(GpuError::DriverLibraryUnavailable {
1117 reason: "sphere GPU backend is Linux-only".to_string(),
1118 })
1119 }
1120}
1121
1122pub fn householder_reflector_from_weights(w: &[f64]) -> (Vec<f64>, f64) {
1135 let m = w.len();
1136 if m == 0 {
1137 return (Vec::new(), 0.0);
1138 }
1139 let norm = w.iter().map(|x| x * x).sum::<f64>().sqrt();
1140 if norm == 0.0 {
1141 return (vec![0.0; m], 0.0);
1142 }
1143 let sigma = if w[0] >= 0.0 { norm } else { -norm };
1144 let mut v = w.to_vec();
1145 v[0] += sigma;
1146 let v0 = v[0];
1147 if v0 == 0.0 {
1148 return (vec![0.0; m], 0.0);
1149 }
1150 for entry in v.iter_mut() {
1152 *entry /= v0;
1153 }
1154 let vv: f64 = v.iter().map(|x| x * x).sum();
1156 let beta = 2.0 / vv;
1157 (v, beta)
1158}
1159
1160pub fn build_center_kernel_device(
1179 centers_xyz: &[f64],
1180 lmax: usize,
1181 coeffs: &[f64],
1182 kind: SphereSpectralKernelKind,
1183) -> Result<DeviceS2KernelMatrix, GpuError> {
1184 let m = centers_xyz.len() / 3;
1185 if centers_xyz.len() != 3 * m {
1186 return Err(GpuError::DriverCallFailed {
1187 reason: "build_center_kernel_device: centers_xyz length not divisible by 3".into(),
1188 });
1189 }
1190 let inputs = S2KernelBuildInputs {
1191 n: m,
1192 m,
1193 lmax,
1194 data_xyz: centers_xyz,
1195 centers_xyz,
1196 coeffs,
1197 kind,
1198 layout: DeviceMatrixLayout::ColumnMajor,
1199 };
1200 build_kernel_matrix_device(inputs)
1201}
1202
1203pub fn constrained_penalty_host(
1208 c: ArrayView2<'_, f64>,
1209 w: &[f64],
1210) -> Result<Array2<f64>, GpuError> {
1211 let (m1, m2) = c.dim();
1212 if m1 != m2 {
1213 gam_gpu::gpu_bail!("constrained_penalty_host: C must be square, got {m1}x{m2}");
1214 }
1215 let m = m1;
1216 if w.len() != m {
1217 gam_gpu::gpu_bail!("constrained_penalty_host: w.len()={} != m={}", w.len(), m);
1218 }
1219 if m < 2 {
1220 gam_gpu::gpu_bail!("constrained_penalty_host: m must be >= 2 (got {m})");
1221 }
1222 let (v, beta) = householder_reflector_from_weights(w);
1223
1224 let mut u = vec![0.0_f64; m];
1227 for i in 0..m {
1228 let mut acc = 0.0_f64;
1229 for j in 0..m {
1230 acc += c[(i, j)] * v[j];
1231 }
1232 u[i] = acc;
1233 }
1234 let vtcv: f64 = v.iter().zip(&u).map(|(vi, ui)| vi * ui).sum();
1235 let mut hch = Array2::<f64>::zeros((m, m));
1236 for i in 0..m {
1237 for j in 0..m {
1238 hch[(i, j)] =
1239 c[(i, j)] - beta * (v[i] * u[j] + u[i] * v[j]) + beta * beta * vtcv * v[i] * v[j];
1240 }
1241 }
1242 let mut s = Array2::<f64>::zeros((m - 1, m - 1));
1244 for i in 0..(m - 1) {
1245 for j in 0..(m - 1) {
1246 s[(i, j)] = hch[(i + 1, j + 1)];
1247 }
1248 }
1249 Ok(s)
1250}
1251
1252#[derive(Clone, Debug)]
1283pub struct PenalisedLsSolution {
1284 pub beta: Vec<f64>,
1286 pub weighted_residual_ssq: f64,
1288 pub log_det_hessian: f64,
1290}
1291
1292#[cfg(target_os = "linux")]
1302pub fn solve_penalised_ls_device(
1303 x_s_device: &DeviceS2KernelMatrix,
1304 wy: &[f64],
1305 r_s: ArrayView2<'_, f64>,
1306) -> Result<PenalisedLsSolution, GpuError> {
1307 use cudarc::cusolver::{DnHandle, sys as cusolver_sys};
1308 use cudarc::driver::DevicePtrMut;
1309
1310 let n = x_s_device.rows;
1311 let p = x_s_device.cols;
1312 if wy.len() != n {
1313 gam_gpu::gpu_bail!("solve_penalised_ls_device: wy.len()={} != n={n}", wy.len());
1314 }
1315 if r_s.dim() != (p, p) {
1316 gam_gpu::gpu_bail!(
1317 "solve_penalised_ls_device: r_s.dim()={:?} != ({p}, {p})",
1318 r_s.dim()
1319 );
1320 }
1321 if p == 0 {
1322 return Ok(PenalisedLsSolution {
1323 beta: Vec::new(),
1324 weighted_residual_ssq: wy.iter().map(|v| v * v).sum(),
1325 log_det_hessian: 0.0,
1326 });
1327 }
1328
1329 let stream = x_s_device.stream.clone();
1330 let n_aug = n + p;
1331
1332 let mut a_aug_host = vec![0.0_f64; n_aug * p];
1337 let mut x_host_colmajor = vec![0.0_f64; x_s_device.ld * p];
1339 x_s_device.copy_to_host_col_major(&mut x_host_colmajor)?;
1340 for j in 0..p {
1341 let src_off = j * x_s_device.ld;
1342 let dst_off = j * n_aug;
1343 a_aug_host[dst_off..dst_off + n].copy_from_slice(&x_host_colmajor[src_off..src_off + n]);
1344 for i in 0..p {
1345 a_aug_host[dst_off + n + i] = r_s[(i, j)];
1348 }
1349 }
1350 let mut a_dev = stream
1351 .clone_htod(&a_aug_host)
1352 .gpu_ctx("solve_penalised_ls_device htod A_aug")?;
1353
1354 let mut b_host = vec![0.0_f64; n_aug];
1356 b_host[..n].copy_from_slice(wy);
1357 let mut b_dev = stream
1358 .clone_htod(&b_host)
1359 .gpu_ctx("solve_penalised_ls_device htod b_aug")?;
1360
1361 let solver = DnHandle::new(stream.clone()).gpu_ctx("solve_penalised_ls_device DnHandle")?;
1362 let n_aug_i: i32 = i32::try_from(n_aug)
1363 .map_err(|_| gam_gpu::gpu_err!("solve_penalised_ls_device: n_aug={n_aug} overflows i32"))?;
1364 let p_i: i32 = i32::try_from(p)
1365 .map_err(|_| gam_gpu::gpu_err!("solve_penalised_ls_device: p={p} overflows i32"))?;
1366
1367 let mut lwork: i32 = 0;
1369 {
1370 let (a_ptr, _rec) = a_dev.device_ptr_mut(&stream);
1371 let status = unsafe {
1374 cusolver_sys::cusolverDnDgeqrf_bufferSize(
1375 solver.cu(),
1376 n_aug_i,
1377 p_i,
1378 a_ptr as *mut f64,
1379 n_aug_i,
1380 &mut lwork,
1381 )
1382 };
1383 if status != cusolver_sys::cusolverStatus_t::CUSOLVER_STATUS_SUCCESS {
1384 gam_gpu::gpu_bail!("cusolverDnDgeqrf_bufferSize status={status:?}");
1385 }
1386 }
1387 let lwork_us = usize::try_from(lwork)
1388 .map_err(|_| gam_gpu::gpu_err!("solve_penalised_ls_device: negative lwork={lwork}"))?;
1389 let mut workspace = stream
1390 .alloc_zeros::<f64>(lwork_us.max(1))
1391 .gpu_ctx("solve_penalised_ls_device alloc workspace")?;
1392 let mut tau = stream
1393 .alloc_zeros::<f64>(p)
1394 .gpu_ctx("solve_penalised_ls_device alloc tau")?;
1395 let mut info = stream
1396 .alloc_zeros::<i32>(1)
1397 .gpu_ctx("solve_penalised_ls_device alloc info")?;
1398
1399 {
1401 let (a_ptr, _rec_a) = a_dev.device_ptr_mut(&stream);
1402 let (tau_ptr, _rec_t) = tau.device_ptr_mut(&stream);
1403 let (work_ptr, _rec_w) = workspace.device_ptr_mut(&stream);
1404 let (info_ptr, _rec_i) = info.device_ptr_mut(&stream);
1405 let status = unsafe {
1408 cusolver_sys::cusolverDnDgeqrf(
1409 solver.cu(),
1410 n_aug_i,
1411 p_i,
1412 a_ptr as *mut f64,
1413 n_aug_i,
1414 tau_ptr as *mut f64,
1415 work_ptr as *mut f64,
1416 lwork,
1417 info_ptr as *mut i32,
1418 )
1419 };
1420 if status != cusolver_sys::cusolverStatus_t::CUSOLVER_STATUS_SUCCESS {
1421 gam_gpu::gpu_bail!("cusolverDnDgeqrf status={status:?}");
1422 }
1423 }
1424
1425 let mut ormqr_lwork: i32 = 0;
1427 {
1428 let (a_ptr, _rec_a) = a_dev.device_ptr_mut(&stream);
1429 let (tau_ptr, _rec_t) = tau.device_ptr_mut(&stream);
1430 let (b_ptr, _rec_b) = b_dev.device_ptr_mut(&stream);
1431 let status = unsafe {
1434 cusolver_sys::cusolverDnDormqr_bufferSize(
1435 solver.cu(),
1436 cusolver_sys::cublasSideMode_t::CUBLAS_SIDE_LEFT,
1437 cusolver_sys::cublasOperation_t::CUBLAS_OP_T,
1438 n_aug_i,
1439 1,
1440 p_i,
1441 a_ptr as *const f64,
1442 n_aug_i,
1443 tau_ptr as *const f64,
1444 b_ptr as *mut f64,
1445 n_aug_i,
1446 &mut ormqr_lwork,
1447 )
1448 };
1449 if status != cusolver_sys::cusolverStatus_t::CUSOLVER_STATUS_SUCCESS {
1450 gam_gpu::gpu_bail!("cusolverDnDormqr_bufferSize status={status:?}");
1451 }
1452 }
1453 if ormqr_lwork > lwork {
1454 workspace = stream
1455 .alloc_zeros::<f64>(usize::try_from(ormqr_lwork).unwrap_or(1))
1456 .gpu_ctx("solve_penalised_ls_device realloc workspace ormqr")?;
1457 }
1458 {
1459 let (a_ptr, _rec_a) = a_dev.device_ptr_mut(&stream);
1460 let (tau_ptr, _rec_t) = tau.device_ptr_mut(&stream);
1461 let (b_ptr, _rec_b) = b_dev.device_ptr_mut(&stream);
1462 let (work_ptr, _rec_w) = workspace.device_ptr_mut(&stream);
1463 let (info_ptr, _rec_i) = info.device_ptr_mut(&stream);
1464 let status = unsafe {
1468 cusolver_sys::cusolverDnDormqr(
1469 solver.cu(),
1470 cusolver_sys::cublasSideMode_t::CUBLAS_SIDE_LEFT,
1471 cusolver_sys::cublasOperation_t::CUBLAS_OP_T,
1472 n_aug_i,
1473 1,
1474 p_i,
1475 a_ptr as *const f64,
1476 n_aug_i,
1477 tau_ptr as *const f64,
1478 b_ptr as *mut f64,
1479 n_aug_i,
1480 work_ptr as *mut f64,
1481 ormqr_lwork.max(lwork),
1482 info_ptr as *mut i32,
1483 )
1484 };
1485 if status != cusolver_sys::cusolverStatus_t::CUSOLVER_STATUS_SUCCESS {
1486 gam_gpu::gpu_bail!("cusolverDnDormqr status={status:?}");
1487 }
1488 }
1489
1490 {
1493 use cudarc::cublas::CudaBlas;
1494 let blas = CudaBlas::new(stream.clone()).gpu_ctx("solve_penalised_ls_device CudaBlas")?;
1495 let alpha = 1.0_f64;
1496 let (a_ptr, _rec_a) = a_dev.device_ptr_mut(&stream);
1497 let (b_ptr, _rec_b) = b_dev.device_ptr_mut(&stream);
1498 let handle = *blas.handle();
1503 let status = unsafe {
1504 cudarc::cublas::sys::cublasDtrsm_v2(
1505 handle,
1506 cudarc::cublas::sys::cublasSideMode_t::CUBLAS_SIDE_LEFT,
1507 cudarc::cublas::sys::cublasFillMode_t::CUBLAS_FILL_MODE_UPPER,
1508 cudarc::cublas::sys::cublasOperation_t::CUBLAS_OP_N,
1509 cudarc::cublas::sys::cublasDiagType_t::CUBLAS_DIAG_NON_UNIT,
1510 p_i,
1511 1,
1512 &alpha,
1513 a_ptr as *const f64,
1514 n_aug_i,
1515 b_ptr as *mut f64,
1516 n_aug_i,
1517 )
1518 };
1519 if status != cudarc::cublas::sys::cublasStatus_t::CUBLAS_STATUS_SUCCESS {
1520 gam_gpu::gpu_bail!("cublasDtrsm_v2 status={status:?}");
1521 }
1522 }
1523
1524 let mut b_out = vec![0.0_f64; n_aug];
1526 stream
1527 .memcpy_dtoh(&b_dev, &mut b_out)
1528 .gpu_ctx("solve_penalised_ls_device dtoh b_out")?;
1529 let mut a_back = vec![0.0_f64; n_aug * p];
1530 stream
1531 .memcpy_dtoh(&a_dev, &mut a_back)
1532 .gpu_ctx("solve_penalised_ls_device dtoh A_back")?;
1533 stream
1534 .synchronize()
1535 .gpu_ctx("solve_penalised_ls_device synchronize")?;
1536
1537 let beta: Vec<f64> = b_out[..p].to_vec();
1538 let augmented_residual_ssq: f64 = b_out[p..].iter().map(|v| v * v).sum();
1547
1548 let mut log_abs_r = 0.0_f64;
1550 for k in 0..p {
1551 let r_kk = a_back[k * n_aug + k];
1552 log_abs_r += r_kk.abs().ln();
1553 }
1554 let log_det_hessian = 2.0 * log_abs_r;
1555
1556 Ok(PenalisedLsSolution {
1557 beta,
1558 weighted_residual_ssq: augmented_residual_ssq,
1559 log_det_hessian,
1560 })
1561}
1562
1563#[cfg(not(target_os = "linux"))]
1564pub fn solve_penalised_ls_device(
1565 x_s_device: &DeviceS2KernelMatrix,
1566 wy: &[f64],
1567 r_s: ArrayView2<'_, f64>,
1568) -> Result<PenalisedLsSolution, GpuError> {
1569 Err(GpuError::DriverLibraryUnavailable {
1570 reason: format!(
1571 "sphere GPU cuSOLVER QR path is Linux-only (n={}, p={}, wy.len()={}, r_s={:?})",
1572 x_s_device.rows,
1573 x_s_device.cols,
1574 wy.len(),
1575 r_s.dim()
1576 ),
1577 })
1578}
1579
1580#[cfg(test)]
1585mod sphere_gpu_tests {
1586 use super::*;
1587 use crate::basis::{
1588 SphereWahbaKernel, sobolev_s2_truncated_coefficients, sphere_truncated_spectral_eval,
1589 spherical_wahba_kernel_matrix_with_kind,
1590 };
1591 use ndarray::Array2;
1592
1593 fn small_latlon_grid(n_lat: usize, n_lon: usize) -> Array2<f64> {
1594 let mut rows = Vec::with_capacity(n_lat * n_lon);
1596 for i in 0..n_lat {
1597 let lat = -85.0 + (170.0 * i as f64) / (n_lat.saturating_sub(1).max(1) as f64);
1598 for j in 0..n_lon {
1599 let lon = -180.0 + (360.0 * j as f64) / (n_lon.saturating_sub(1).max(1) as f64);
1600 rows.push(lat);
1601 rows.push(lon);
1602 }
1603 }
1604 Array2::from_shape_vec((n_lat * n_lon, 2), rows).unwrap()
1605 }
1606
1607 fn cuda_available_for_test(label: &str) -> bool {
1608 match gam_gpu::device_runtime::GpuRuntime::resolve(gam_gpu::GpuPolicy::Auto) {
1609 Ok(Some(_)) => true,
1610 Ok(None) => {
1611 eprintln!("[sphere_gpu test] no CUDA device — skipping {label}");
1612 false
1613 }
1614 Err(error) => panic!("[sphere_gpu test] CUDA resolution failed for {label}: {error}"),
1615 }
1616 }
1617
1618 fn assert_sphere_decision_declines_without_device(n: usize, m: usize, lmax: usize) {
1626 let decision = sphere_kernel_decision(n, m, lmax)
1627 .expect("the sphere GPU decision must not fault on a device-free host");
1628 assert!(
1629 !decision.use_gpu,
1630 "no CUDA runtime on this host, yet the sphere dispatch decision admitted the \
1631 device for (n={n}, m={m}, lmax={lmax}) — reason={}",
1632 decision.reason
1633 );
1634 }
1635
1636 fn assert_device_kernel_entry_refuses(inputs: S2KernelBuildInputs<'_>) {
1641 assert!(
1642 build_kernel_matrix_device(inputs).is_err(),
1643 "no CUDA runtime on this host, yet the device kernel entry returned a matrix \
1644 — the admitted-only device path fabricated a host answer (#1551 class)"
1645 );
1646 }
1647
1648 fn assert_cpu_kernel_matches_spectral_definition(
1654 kernel_matrix: &Array2<f64>,
1655 data_xyz: &[f64],
1656 centers_xyz: &[f64],
1657 coeffs: &[f64],
1658 ) {
1659 let (n, m) = kernel_matrix.dim();
1660 let mut max_abs = 0.0_f64;
1661 for i in 0..n {
1662 for j in 0..m {
1663 let dot = data_xyz[3 * i] * centers_xyz[3 * j]
1664 + data_xyz[3 * i + 1] * centers_xyz[3 * j + 1]
1665 + data_xyz[3 * i + 2] * centers_xyz[3 * j + 2];
1666 let expected = sphere_truncated_spectral_eval(dot.clamp(-1.0, 1.0), coeffs);
1667 max_abs = max_abs.max((kernel_matrix[(i, j)] - expected).abs());
1668 }
1669 }
1670 assert!(
1671 max_abs < 1e-12,
1672 "CPU truncated-spectral kernel matrix departs from its elementwise definition \
1673 Σ_ℓ c_ℓ P_ℓ(x·c): max |Δ| = {max_abs:.3e}"
1674 );
1675 }
1676
1677 #[test]
1678 fn sum_finite_guard_accepts_finite_rejects_nonfinite() {
1679 let finite = Array2::<f64>::from_shape_fn((5, 7), |(i, j)| (i as f64 - 2.0) * (j as f64));
1684 assert!(finite.sum().is_finite());
1685
1686 let mut with_nan = finite.clone();
1687 with_nan[[3, 4]] = f64::NAN;
1688 assert!(!with_nan.sum().is_finite());
1689
1690 let mut with_pos_inf = finite.clone();
1691 with_pos_inf[[0, 0]] = f64::INFINITY;
1692 assert!(!with_pos_inf.sum().is_finite());
1693
1694 let mut with_neg_inf = finite.clone();
1695 with_neg_inf[[4, 6]] = f64::NEG_INFINITY;
1696 assert!(!with_neg_inf.sum().is_finite());
1697 }
1698
1699 #[test]
1700 fn xyz_preprocessing_matches_unit_sphere() {
1701 let latlon = ndarray::array![
1702 [0.0, 0.0],
1703 [90.0, 0.0],
1704 [0.0, 90.0],
1705 [-90.0, 17.5],
1706 [45.0, -120.0],
1707 ];
1708 let xyz = latlon_to_xyz_host(latlon.view(), false).expect("xyz");
1709 assert_eq!(xyz.len(), 3 * 5);
1710 for i in 0..5 {
1711 let nrm2 = xyz[3 * i] * xyz[3 * i]
1712 + xyz[3 * i + 1] * xyz[3 * i + 1]
1713 + xyz[3 * i + 2] * xyz[3 * i + 2];
1714 assert!((nrm2 - 1.0).abs() < 1e-15, "row {i} not unit norm: {nrm2}");
1715 }
1716 assert!((xyz[0] - 1.0).abs() < 1e-15);
1718 assert!(xyz[1].abs() < 1e-15);
1719 assert!(xyz[2].abs() < 1e-15);
1720 assert!(xyz[3].abs() < 1e-15);
1722 assert!(xyz[4].abs() < 1e-15);
1723 assert!((xyz[5] - 1.0).abs() < 1e-15);
1724 assert!(xyz[6].abs() < 1e-15);
1726 assert!((xyz[7] - 1.0).abs() < 1e-15);
1727 assert!(xyz[8].abs() < 1e-15);
1728 }
1729
1730 #[test]
1731 fn truncated_spectral_at_same_point_matches_sum_of_coefficients() {
1732 for m_penalty in 1..=4 {
1736 for &lmax in &[5_usize, 20, 50] {
1737 let coeffs = sobolev_s2_truncated_coefficients(lmax, m_penalty);
1738 let expected: f64 = coeffs.iter().sum();
1739 let got = sphere_truncated_spectral_eval(1.0, &coeffs);
1740 assert!(
1741 (got - expected).abs() < 1e-13,
1742 "K(x,x) identity broken at m={m_penalty}, L={lmax}: got {got:.6e}, expected {expected:.6e}"
1743 );
1744 }
1745 }
1746 }
1747
1748 #[test]
1749 fn truncated_spectral_at_antipode_matches_alternating_sum() {
1750 for m_penalty in 1..=4 {
1753 for &lmax in &[5_usize, 20, 50] {
1754 let coeffs = sobolev_s2_truncated_coefficients(lmax, m_penalty);
1755 let expected: f64 = coeffs
1756 .iter()
1757 .enumerate()
1758 .map(|(ell, c)| if ell % 2 == 0 { *c } else { -*c })
1759 .sum();
1760 let got = sphere_truncated_spectral_eval(-1.0, &coeffs);
1761 assert!(
1762 (got - expected).abs() < 1e-13,
1763 "K(x,-x) identity broken at m={m_penalty}, L={lmax}: got {got:.6e}, expected {expected:.6e}"
1764 );
1765 }
1766 }
1767 }
1768
1769 #[test]
1770 fn truncated_spectral_matrix_is_symmetric() {
1771 let centers = ndarray::array![
1775 [10.0_f64, 20.0],
1776 [-30.0, 100.0],
1777 [45.0, -60.0],
1778 [-89.0, 0.0],
1779 [0.0, 180.0],
1780 [60.0, -179.9],
1781 ];
1782 for m_penalty in [1usize, 2, 4] {
1783 for &lmax in &[10_usize, 30] {
1784 let mat = spherical_wahba_kernel_matrix_with_kind(
1785 centers.view(),
1786 centers.view(),
1787 m_penalty,
1788 false,
1789 SphereWahbaKernel::SobolevTruncated { lmax: lmax as u16 },
1790 )
1791 .expect("kernel matrix");
1792 let n = centers.nrows();
1793 let mut max_asym = 0.0_f64;
1794 for i in 0..n {
1795 for j in 0..n {
1796 let d = (mat[(i, j)] - mat[(j, i)]).abs();
1797 if d > max_asym {
1798 max_asym = d;
1799 }
1800 }
1801 }
1802 assert!(
1803 max_asym < 1e-13,
1804 "K not symmetric at m={m_penalty}, L={lmax}: max |K - Kᵀ| = {max_asym:.3e}"
1805 );
1806 }
1807 }
1808 }
1809
1810 #[test]
1811 fn truncated_coefficients_have_zero_constant_mode() {
1812 for m in 1..=4 {
1813 let c = sobolev_s2_truncated_coefficients(50, m);
1814 assert_eq!(c.len(), 51);
1815 assert_eq!(c[0], 0.0);
1816 assert!(c[1] > 0.0);
1817 for ell in 2..=50 {
1819 assert!(
1820 c[ell] < c[ell - 1] + 1e-15,
1821 "Sobolev coefficient not non-increasing at m={m}, ell={ell}: {} vs {}",
1822 c[ell],
1823 c[ell - 1]
1824 );
1825 }
1826 }
1827 }
1828
1829 #[test]
1830 fn truncated_spectral_matches_matrix_helper() {
1831 let m_penalty = 2;
1835 let lmax = 20;
1836 let coeffs = sobolev_s2_truncated_coefficients(lmax, m_penalty);
1837 let data = ndarray::array![[12.5, -34.0]];
1838 let centers = ndarray::array![[40.0, 10.0]];
1839 let mat = spherical_wahba_kernel_matrix_with_kind(
1840 data.view(),
1841 centers.view(),
1842 m_penalty,
1843 false,
1844 SphereWahbaKernel::SobolevTruncated { lmax: lmax as u16 },
1845 )
1846 .expect("kernel matrix");
1847 let xyz_d = latlon_to_xyz_host(data.view(), false).unwrap();
1849 let xyz_c = latlon_to_xyz_host(centers.view(), false).unwrap();
1850 let cos_g = xyz_d[0] * xyz_c[0] + xyz_d[1] * xyz_c[1] + xyz_d[2] * xyz_c[2];
1851 let expected = sphere_truncated_spectral_eval(cos_g, &coeffs);
1852 assert!(
1853 (mat[(0, 0)] - expected).abs() < 1e-13,
1854 "matrix helper differs from scalar evaluator: {} vs {}",
1855 mat[(0, 0)],
1856 expected
1857 );
1858 }
1859
1860 #[test]
1861 fn constrained_penalty_is_symmetric_and_drops_constraint_direction() {
1862 let m = 6;
1867 let mut c = Array2::<f64>::zeros((m, m));
1868 for i in 0..m {
1869 for j in 0..m {
1870 let d = (i as f64 - j as f64).abs();
1871 c[(i, j)] = (-0.5 * d).exp();
1872 }
1873 }
1874 let w = vec![1.0_f64; m];
1875 let s = constrained_penalty_host(c.view(), &w).expect("constrained S");
1876 assert_eq!(s.dim(), (m - 1, m - 1));
1877 let mut max_asym = 0.0_f64;
1879 for i in 0..(m - 1) {
1880 for j in 0..(m - 1) {
1881 let d = (s[(i, j)] - s[(j, i)]).abs();
1882 if d > max_asym {
1883 max_asym = d;
1884 }
1885 }
1886 }
1887 assert!(
1888 max_asym < 1e-13,
1889 "S not symmetric: max |S - Sᵀ| = {max_asym:.3e}"
1890 );
1891
1892 let ones = ndarray::Array1::<f64>::ones(m - 1);
1900 let sx = s.dot(&ones);
1901 assert!(sx.iter().all(|v| v.is_finite()));
1902 }
1903
1904 #[test]
1905 fn householder_reflector_zeroes_target_vector() {
1906 let w = vec![3.0, 4.0, 0.0, -1.0];
1907 let (v, beta) = householder_reflector_from_weights(&w);
1908 let dot: f64 = v.iter().zip(&w).map(|(a, b)| a * b).sum();
1911 let hw: Vec<f64> = w
1912 .iter()
1913 .zip(&v)
1914 .map(|(wj, vj)| wj - beta * dot * vj)
1915 .collect();
1916 for entry in hw.iter().skip(1) {
1917 assert!(entry.abs() < 1e-12, "H · w not e_1 multiple: {hw:?}");
1918 }
1919 assert!(hw[0].abs() > 0.0);
1920 }
1921
1922 #[test]
1928 fn sphere_gpu_raw_kernel_parity_vs_cpu_truncated() {
1929 let data_ll = small_latlon_grid(7, 9);
1930 let centers_ll = small_latlon_grid(5, 7);
1931 let data_xyz = latlon_to_xyz_host(data_ll.view(), false).unwrap();
1932 let centers_xyz = latlon_to_xyz_host(centers_ll.view(), false).unwrap();
1933 let n = data_ll.nrows();
1934 let m = centers_ll.nrows();
1935 let penalty = 2usize;
1936 let lmax = 20usize;
1937 let coeffs = sobolev_s2_truncated_coefficients(lmax, penalty);
1938
1939 let inputs = S2KernelBuildInputs {
1940 n,
1941 m,
1942 lmax,
1943 data_xyz: &data_xyz,
1944 centers_xyz: ¢ers_xyz,
1945 coeffs: &coeffs,
1946 kind: SphereSpectralKernelKind::Sobolev,
1947 layout: DeviceMatrixLayout::ColumnMajor,
1948 };
1949
1950 let cpu = spherical_wahba_kernel_matrix_with_kind(
1951 data_ll.view(),
1952 centers_ll.view(),
1953 penalty,
1954 false,
1955 SphereWahbaKernel::SobolevTruncated { lmax: lmax as u16 },
1956 )
1957 .expect("cpu kernel matrix");
1958
1959 assert_cpu_kernel_matches_spectral_definition(&cpu, &data_xyz, ¢ers_xyz, &coeffs);
1962
1963 if !cuda_available_for_test("raw-kernel parity") {
1964 assert_sphere_decision_declines_without_device(n, m, lmax);
1965 assert_device_kernel_entry_refuses(inputs);
1966 return;
1967 }
1968 SphereGpuBackend::probe()
1971 .expect("[sphere_gpu test] backend probe must succeed on a CUDA host");
1972 let dev_mat = build_kernel_matrix_device(inputs).expect("device kernel matrix");
1973 let gpu = dev_mat.to_host_array().expect("dtoh kernel matrix");
1974
1975 let mut max_abs = 0.0_f64;
1976 for i in 0..n {
1977 for j in 0..m {
1978 let d = (gpu[(i, j)] - cpu[(i, j)]).abs();
1979 if d > max_abs {
1980 max_abs = d;
1981 }
1982 }
1983 }
1984 assert!(
1985 max_abs < 1e-11,
1986 "GPU vs CPU truncated parity max |Δ| = {max_abs:.3e} >= 1e-11"
1987 );
1988 }
1989
1990 #[test]
2004 fn sphere_gpu_end_to_end_dispatch_parity_vs_cpu_truncated() {
2005 use crate::basis::{
2006 CenterStrategy, SphereMethod, SphericalSplineBasisSpec, SphericalSplineIdentifiability,
2007 build_spherical_spline_basis, spherical_wahba_kernel_matrix_cpu,
2008 spherical_wahba_kernel_matrix_with_kind,
2009 };
2010 let on_cuda = cuda_available_for_test("end-to-end dispatch parity");
2011 if on_cuda {
2012 SphereGpuBackend::probe()
2016 .expect("[sphere_gpu test] backend probe must succeed on a CUDA host");
2017 }
2018
2019 let data = small_latlon_grid(100, 100);
2021 let lmax: u16 = 30;
2022 let penalty_order = 2usize;
2023 let centers =
2024 crate::basis::select_spherical_farthest_point_centers(data.view(), 200, false)
2025 .expect("centers");
2026 let n = data.nrows();
2027 let m = centers.nrows();
2028
2029 if on_cuda {
2035 let decision = sphere_kernel_decision(n, m, lmax as usize)
2036 .expect("GPU decision must preserve CUDA resolution faults");
2037 assert!(
2038 decision.use_gpu,
2039 "expected GPU dispatch for (n={n}, m={m}, lmax={lmax}); decision said CPU \
2040 (reason={}); the engagement gate regressed",
2041 decision.reason
2042 );
2043 } else {
2044 assert_sphere_decision_declines_without_device(n, m, lmax as usize);
2045 assert!(
2046 try_build_truncated_kernel_matrix_gpu(
2047 data.view(),
2048 centers.view(),
2049 penalty_order,
2050 false,
2051 SphereWahbaKernel::SobolevTruncated { lmax },
2052 )
2053 .is_none(),
2054 "no CUDA runtime on this host, yet the production sphere seam did not take \
2055 the quiet CPU route at the device-eligible shape (n={n}, m={m}, lmax={lmax})"
2056 );
2057 }
2058
2059 let dispatched_kernel = spherical_wahba_kernel_matrix_with_kind(
2063 data.view(),
2064 centers.view(),
2065 penalty_order,
2066 false,
2067 SphereWahbaKernel::SobolevTruncated { lmax },
2068 )
2069 .expect("GPU-eligible production kernel build succeeds");
2070
2071 let cpu_kernel = spherical_wahba_kernel_matrix_cpu(
2073 data.view(),
2074 centers.view(),
2075 penalty_order,
2076 false,
2077 SphereWahbaKernel::SobolevTruncated { lmax },
2078 )
2079 .expect("cpu oracle kernel build succeeds");
2080
2081 assert_eq!(dispatched_kernel.dim(), cpu_kernel.dim());
2082 let mut max_abs = 0.0_f64;
2083 let mut max_rel = 0.0_f64;
2084 for (g, c) in dispatched_kernel.iter().zip(cpu_kernel.iter()) {
2085 let d = (g - c).abs();
2086 if d > max_abs {
2087 max_abs = d;
2088 }
2089 let denom = g.abs().max(c.abs()).max(1e-300);
2090 let r = d / denom;
2091 if r > max_rel {
2092 max_rel = r;
2093 }
2094 }
2095 assert!(
2096 max_rel < 1e-9,
2097 "GPU-dispatch vs CPU-oracle kernel parity max relative |Δ| = {max_rel:.3e} \
2098 >= 1e-9 (abs {max_abs:.3e})"
2099 );
2100 if !on_cuda {
2101 for (a, b) in dispatched_kernel.iter().zip(cpu_kernel.iter()) {
2105 assert_eq!(
2106 a.to_bits(),
2107 b.to_bits(),
2108 "device-free dispatcher must equal the CPU oracle bit-for-bit"
2109 );
2110 }
2111 }
2112
2113 let spec_gpu = SphericalSplineBasisSpec {
2117 center_strategy: CenterStrategy::FarthestPoint { num_centers: 200 },
2118 penalty_order,
2119 double_penalty: false,
2120 radians: false,
2121 method: SphereMethod::Wahba,
2122 max_degree: None,
2123 wahba_kernel: SphereWahbaKernel::SobolevTruncated { lmax },
2124 identifiability: SphericalSplineIdentifiability::CenterSumToZero,
2125 };
2126 let result_gpu = build_spherical_spline_basis(data.view(), &spec_gpu)
2127 .expect("GPU-eligible build_spherical_spline_basis succeeds");
2128 let design = result_gpu.design.as_dense().expect("dense design");
2129 assert_eq!(design.nrows(), n, "design row count must match data rows");
2130 assert!(
2131 design.iter().all(|v| v.is_finite()),
2132 "engaged-device spherical design must be finite"
2133 );
2134 }
2135
2136 fn householder_apply_host(b: &Array2<f64>, v: &[f64], beta: f64) -> Array2<f64> {
2140 let (n, m) = b.dim();
2141 let mut xs = Array2::<f64>::zeros((n, m - 1));
2142 for i in 0..n {
2143 let d_i: f64 = (0..m).map(|j| v[j] * b[(i, j)]).sum();
2144 for j_out in 0..(m - 1) {
2145 xs[(i, j_out)] = b[(i, j_out + 1)] - beta * d_i * v[j_out + 1];
2146 }
2147 }
2148 xs
2149 }
2150
2151 fn assert_householder_fused_matches_explicit_product(b: &Array2<f64>, v: &[f64], beta: f64) {
2156 let (n, m) = b.dim();
2157 let mut reflector = Array2::<f64>::eye(m);
2158 for i in 0..m {
2159 for j in 0..m {
2160 reflector[(i, j)] -= beta * v[i] * v[j];
2161 }
2162 }
2163 let full = b.dot(&reflector);
2164 let fused = householder_apply_host(b, v, beta);
2165 let mut max_abs = 0.0_f64;
2166 for i in 0..n {
2167 for j in 0..(m - 1) {
2168 max_abs = max_abs.max((fused[(i, j)] - full[(i, j + 1)]).abs());
2169 }
2170 }
2171 assert!(
2172 max_abs < 1e-13,
2173 "fused Householder host expression departs from B·(I − β·v·vᵀ): \
2174 max |Δ| = {max_abs:.3e}"
2175 );
2176 }
2177
2178 #[test]
2183 fn sphere_gpu_householder_parity_vs_raw_dot_z() {
2184 let data_ll = small_latlon_grid(6, 8);
2185 let centers_ll = small_latlon_grid(4, 5);
2186 let data_xyz = latlon_to_xyz_host(data_ll.view(), false).unwrap();
2187 let centers_xyz = latlon_to_xyz_host(centers_ll.view(), false).unwrap();
2188 let n = data_ll.nrows();
2189 let m = centers_ll.nrows();
2190 let penalty = 2usize;
2191 let lmax = 15usize;
2192 let coeffs = sobolev_s2_truncated_coefficients(lmax, penalty);
2193
2194 let inputs_raw = S2KernelBuildInputs {
2196 n,
2197 m,
2198 lmax,
2199 data_xyz: &data_xyz,
2200 centers_xyz: ¢ers_xyz,
2201 coeffs: &coeffs,
2202 kind: SphereSpectralKernelKind::Sobolev,
2203 layout: DeviceMatrixLayout::ColumnMajor,
2204 };
2205 let w = vec![1.0_f64; m];
2208 let (v, beta) = householder_reflector_from_weights(&w);
2209
2210 let b_cpu = spherical_wahba_kernel_matrix_with_kind(
2214 data_ll.view(),
2215 centers_ll.view(),
2216 penalty,
2217 false,
2218 SphereWahbaKernel::SobolevTruncated { lmax: lmax as u16 },
2219 )
2220 .expect("cpu kernel matrix");
2221 assert_cpu_kernel_matches_spectral_definition(&b_cpu, &data_xyz, ¢ers_xyz, &coeffs);
2222 assert_householder_fused_matches_explicit_product(&b_cpu, &v, beta);
2223
2224 if !cuda_available_for_test("householder parity") {
2225 assert_sphere_decision_declines_without_device(n, m, lmax);
2226 assert!(
2227 build_householder_constrained_design_device(inputs_raw, &v, beta).is_err(),
2228 "no CUDA runtime on this host, yet the fused Householder device entry \
2229 returned a design — the admitted-only device path fabricated a host \
2230 answer (#1551 class)"
2231 );
2232 return;
2233 }
2234 SphereGpuBackend::probe()
2237 .expect("[sphere_gpu test] backend probe must succeed on a CUDA host");
2238 let b_dev = build_kernel_matrix_device(inputs_raw.clone()).expect("raw kernel");
2239 let b = b_dev.to_host_array().expect("dtoh raw");
2240
2241 let xs_host = householder_apply_host(&b, &v, beta);
2243
2244 let xs_dev =
2245 build_householder_constrained_design_device(inputs_raw, &v, beta).expect("hh design");
2246 let xs_gpu = xs_dev.to_host_array().expect("dtoh hh");
2247
2248 let mut max_abs = 0.0_f64;
2249 for i in 0..n {
2250 for j in 0..(m - 1) {
2251 let d = (xs_host[(i, j)] - xs_gpu[(i, j)]).abs();
2252 if d > max_abs {
2253 max_abs = d;
2254 }
2255 }
2256 }
2257 assert!(
2258 max_abs < 1e-12,
2259 "Householder fused parity max |Δ| = {max_abs:.3e} >= 1e-12"
2260 );
2261 }
2262
2263 #[test]
2275 fn sphere_gpu_kernel_matrix_hill_climb_declines_without_device_else_20x_vs_cpu() {
2276 let n_lat = 500usize;
2278 let n_lon = 400usize;
2279 assert_eq!(n_lat * n_lon, 200_000);
2280 let m = 200usize;
2281 let lmax = 50usize;
2282
2283 if !cuda_available_for_test("kernel-matrix hill climb") {
2284 assert_sphere_decision_declines_without_device(n_lat * n_lon, m, lmax);
2285 return;
2286 }
2287 SphereGpuBackend::probe()
2290 .expect("[sphere_gpu hill-climb] backend probe must succeed on a CUDA host");
2291
2292 let data_ll = small_latlon_grid(n_lat, n_lon);
2294 let centers_ll =
2295 crate::basis::select_spherical_farthest_point_centers(data_ll.view(), m, false)
2296 .expect("centers");
2297 let n = data_ll.nrows();
2298 let data_xyz = latlon_to_xyz_host(data_ll.view(), false).unwrap();
2299 let centers_xyz = latlon_to_xyz_host(centers_ll.view(), false).unwrap();
2300 let penalty_order = 2usize;
2301 let coeffs = sobolev_s2_truncated_coefficients(lmax, penalty_order);
2302
2303 let inputs_warm = S2KernelBuildInputs {
2305 n,
2306 m,
2307 lmax,
2308 data_xyz: &data_xyz,
2309 centers_xyz: ¢ers_xyz,
2310 coeffs: &coeffs,
2311 kind: SphereSpectralKernelKind::Sobolev,
2312 layout: DeviceMatrixLayout::ColumnMajor,
2313 };
2314 {
2319 let warm = build_kernel_matrix_device(inputs_warm.clone()).expect("warmup");
2320 drop(warm.to_host_array().expect("warmup to_host"));
2321 }
2322
2323 let t0 = std::time::Instant::now();
2325 let dev = build_kernel_matrix_device(inputs_warm.clone()).expect("gpu kernel matrix");
2326 dev.to_host_array().expect("dtoh");
2327 let gpu_secs = t0.elapsed().as_secs_f64();
2328
2329 let t1 = std::time::Instant::now();
2336 crate::basis::spherical_wahba_kernel_matrix_cpu(
2337 data_ll.view(),
2338 centers_ll.view(),
2339 penalty_order,
2340 false,
2341 SphereWahbaKernel::SobolevTruncated { lmax: lmax as u16 },
2342 )
2343 .expect("cpu kernel matrix");
2344 let cpu_secs = t1.elapsed().as_secs_f64();
2345
2346 let ratio = cpu_secs / gpu_secs.max(1e-9);
2347 eprintln!(
2348 "[sphere_gpu hill-climb] n={n} m={m} L={lmax} cpu={cpu_secs:.3}s gpu={gpu_secs:.3}s ratio={ratio:.2}x"
2349 );
2350 assert!(
2357 ratio >= 3.0,
2358 "GPU kernel matrix only {ratio:.2}× faster than CPU (dispatch-worthiness ≥ 3×) at \
2359 n={n} m={m} L={lmax}: cpu={cpu_secs:.3}s gpu={gpu_secs:.3}s"
2360 );
2361 }
2362
2363 #[test]
2377 fn sphere_gpu_end_to_end_fit_hill_climb_declines_without_device_else_10x_vs_cpu() {
2378 use crate::basis::{
2379 CenterStrategy, SphereMethod, SphericalSplineBasisSpec, SphericalSplineIdentifiability,
2380 build_spherical_spline_basis,
2381 };
2382
2383 let n_lat = 500usize;
2384 let n_lon = 400usize;
2385 let m: usize = 200;
2386 let lmax: u16 = 50;
2387
2388 if !cuda_available_for_test("end-to-end fit hill climb") {
2389 assert_sphere_decision_declines_without_device(n_lat * n_lon, m, lmax as usize);
2390 return;
2391 }
2392 SphereGpuBackend::probe()
2395 .expect("[sphere_gpu hill-climb fit] backend probe must succeed on a CUDA host");
2396
2397 let data_ll = small_latlon_grid(n_lat, n_lon);
2398 let spec_gpu = SphericalSplineBasisSpec {
2399 center_strategy: CenterStrategy::FarthestPoint { num_centers: m },
2400 penalty_order: 2,
2401 double_penalty: false,
2402 radians: false,
2403 method: SphereMethod::Wahba,
2404 max_degree: None,
2405 wahba_kernel: SphereWahbaKernel::SobolevTruncated { lmax },
2406 identifiability: SphericalSplineIdentifiability::CenterSumToZero,
2407 };
2408
2409 drop(build_spherical_spline_basis(data_ll.view(), &spec_gpu).expect("warmup build"));
2411
2412 let t0 = std::time::Instant::now();
2413 drop(build_spherical_spline_basis(data_ll.view(), &spec_gpu).expect("gpu build"));
2414 let gpu_secs = t0.elapsed().as_secs_f64();
2415
2416 let centers =
2423 crate::basis::select_spherical_farthest_point_centers(data_ll.view(), m, false)
2424 .expect("centers");
2425 let z = Array2::<f64>::eye(centers.nrows());
2426 let t1 = std::time::Instant::now();
2427 let raw_cpu = crate::basis::spherical_wahba_kernel_matrix_cpu(
2432 data_ll.view(),
2433 centers.view(),
2434 2,
2435 false,
2436 SphereWahbaKernel::SobolevTruncated { lmax },
2437 )
2438 .expect("cpu raw");
2439 raw_cpu.dot(&z);
2440 let cpu_secs = t1.elapsed().as_secs_f64();
2441
2442 let ratio = cpu_secs / gpu_secs.max(1e-9);
2443 eprintln!(
2444 "[sphere_gpu hill-climb fit] n={} m={m} L={lmax} cpu={cpu_secs:.3}s gpu={gpu_secs:.3}s ratio={ratio:.2}x",
2445 data_ll.nrows()
2446 );
2447 assert!(
2448 ratio >= 10.0,
2449 "End-to-end sphere fit only {ratio:.2}× faster on GPU (target ≥ 10×): \
2450 cpu={cpu_secs:.3}s gpu={gpu_secs:.3}s"
2451 );
2452 }
2453
2454 #[test]
2476 fn sphere_gpu_end_to_end_fit_parity_vs_cpu_truncated() {
2477 use crate::basis::{
2478 select_spherical_farthest_point_centers, spherical_wahba_kernel_matrix_with_kind,
2479 };
2480 use faer::Side;
2481 use gam_linalg::faer_ndarray::FaerCholesky;
2482
2483 let data_ll = small_latlon_grid(25, 40);
2485 assert_eq!(data_ll.nrows(), 1000);
2486 let n = data_ll.nrows();
2487 let m: usize = 80;
2488 let lmax_u16: u16 = 15;
2489 let lmax: usize = lmax_u16 as usize;
2490 let penalty_order: usize = 2;
2491 let kernel = SphereWahbaKernel::SobolevTruncated { lmax: lmax_u16 };
2492 let lambda: f64 = 1.0e-3;
2493
2494 let centers_ll = select_spherical_farthest_point_centers(data_ll.view(), m, false)
2496 .expect("farthest-point centers");
2497 assert_eq!(centers_ll.nrows(), m);
2498
2499 let z = Array2::<f64>::eye(centers_ll.nrows());
2502 let p = z.ncols();
2503 assert_eq!(p, m);
2504
2505 let k_cc = spherical_wahba_kernel_matrix_with_kind(
2510 centers_ll.view(),
2511 centers_ll.view(),
2512 penalty_order,
2513 false,
2514 kernel,
2515 )
2516 .expect("centers×centers kernel");
2517 let s_full = z.t().dot(&k_cc).dot(&z);
2518
2519 let raw_design_cpu = spherical_wahba_kernel_matrix_with_kind(
2521 data_ll.view(),
2522 centers_ll.view(),
2523 penalty_order,
2524 false,
2525 kernel,
2526 )
2527 .expect("CPU raw design");
2528 let x_s_cpu = raw_design_cpu.dot(&z);
2529
2530 let data_xyz = latlon_to_xyz_host(data_ll.view(), false).expect("data xyz");
2532 let centers_xyz = latlon_to_xyz_host(centers_ll.view(), false).expect("centers xyz");
2533 let coeffs = crate::basis::sobolev_s2_truncated_coefficients(lmax, penalty_order);
2534 let inputs = S2KernelBuildInputs {
2535 n,
2536 m,
2537 lmax,
2538 data_xyz: &data_xyz,
2539 centers_xyz: ¢ers_xyz,
2540 coeffs: &coeffs,
2541 kind: SphereSpectralKernelKind::Sobolev,
2542 layout: DeviceMatrixLayout::ColumnMajor,
2543 };
2544 let mut y = ndarray::Array1::<f64>::zeros(n);
2550 for i in 0..n {
2551 let lat_rad = data_ll[(i, 0)].to_radians();
2552 let lon_rad = data_ll[(i, 1)].to_radians();
2553 y[i] = (2.0 * lat_rad).sin() * (3.0 * lon_rad).cos()
2555 + 0.25 * lat_rad.cos() * (5.0 * lon_rad).sin();
2556 }
2557
2558 let solve_penalised = |x_s: &ndarray::Array2<f64>| -> ndarray::Array1<f64> {
2563 let xtx = x_s.t().dot(x_s);
2564 let mut a = xtx;
2565 for i in 0..p {
2566 for j in 0..p {
2567 a[(i, j)] += lambda * s_full[(i, j)];
2568 }
2569 }
2570 let rhs = x_s.t().dot(&y);
2571 let factor = a
2572 .cholesky(Side::Lower)
2573 .expect("penalised normal equations are SPD under λ > 0");
2574 factor.solvevec(&rhs)
2575 };
2576
2577 let beta_cpu = solve_penalised(&x_s_cpu);
2578 assert_eq!(beta_cpu.len(), p);
2579 let yhat_cpu = x_s_cpu.dot(&beta_cpu);
2580 assert_eq!(x_s_cpu.dim(), (n, p));
2581
2582 assert_cpu_kernel_matches_spectral_definition(
2587 &raw_design_cpu,
2588 &data_xyz,
2589 ¢ers_xyz,
2590 &coeffs,
2591 );
2592 {
2593 let mut a = x_s_cpu.t().dot(&x_s_cpu);
2594 for i in 0..p {
2595 for j in 0..p {
2596 a[(i, j)] += lambda * s_full[(i, j)];
2597 }
2598 }
2599 let residual = a.dot(&beta_cpu) - x_s_cpu.t().dot(&y);
2600 let rhs_scale = x_s_cpu
2601 .t()
2602 .dot(&y)
2603 .iter()
2604 .fold(0.0_f64, |acc, v| acc.max(v.abs()))
2605 .max(1.0);
2606 let max_residual = residual.iter().fold(0.0_f64, |acc, v| acc.max(v.abs()));
2607 assert!(
2608 max_residual <= 1e-9 * rhs_scale,
2609 "CPU penalised normal equations not solved: ‖(XᵀX + λS)β − Xᵀy‖∞ = \
2610 {max_residual:.3e} (rhs scale {rhs_scale:.3e})"
2611 );
2612 }
2613
2614 if !cuda_available_for_test("end-to-end fit parity") {
2615 assert_sphere_decision_declines_without_device(n, m, lmax);
2616 assert_device_kernel_entry_refuses(inputs);
2617 return;
2618 }
2619 SphereGpuBackend::probe()
2622 .expect("[sphere gpu parity] sphere GPU backend probe must succeed on a CUDA host");
2623 let raw_dev = build_kernel_matrix_device(inputs).expect("GPU raw design");
2624 let raw_design_gpu = raw_dev.to_host_array().expect("dtoh GPU raw design");
2625 let x_s_gpu = raw_design_gpu.dot(&z);
2626
2627 assert_eq!(x_s_gpu.dim(), (n, p));
2628
2629 let mut raw_xs_delta = 0.0_f64;
2639 let mut xs_scale = 0.0_f64;
2640 for (a, b) in x_s_cpu.iter().zip(x_s_gpu.iter()) {
2641 raw_xs_delta = raw_xs_delta.max((a - b).abs());
2642 xs_scale = xs_scale.max(a.abs());
2643 }
2644 let cond = {
2647 use gam_linalg::faer_ndarray::FaerEigh;
2648 let xtx = x_s_cpu.t().dot(&x_s_cpu);
2649 let mut a = xtx;
2650 for i in 0..p {
2651 for j in 0..p {
2652 a[(i, j)] += lambda * s_full[(i, j)];
2653 }
2654 }
2655 let (mut lo, mut hi) = (f64::INFINITY, 0.0_f64);
2656 if let Ok((vals, _)) = a.eigh(faer::Side::Lower) {
2657 for &v in vals.iter() {
2658 lo = lo.min(v);
2659 hi = hi.max(v);
2660 }
2661 }
2662 hi / lo.max(1e-300)
2663 };
2664 assert!(
2669 raw_xs_delta <= 1e-12 * xs_scale.max(1.0),
2670 "GPU vs CPU sphere design matrix max |Δ| = {raw_xs_delta:.3e} > {:.3e} \
2671 (scale {xs_scale:.3e}) — the kernel itself drifted (this is the genuine \
2672 GPU output, NOT a conditioning artifact)",
2673 1e-12 * xs_scale.max(1.0)
2674 );
2675
2676 let beta_gpu = solve_penalised(&x_s_gpu);
2677 assert_eq!(beta_gpu.len(), p);
2678
2679 let yhat_gpu = x_s_gpu.dot(&beta_gpu);
2683
2684 let mut max_beta_delta = 0.0_f64;
2685 for k in 0..p {
2686 let d = (beta_cpu[k] - beta_gpu[k]).abs();
2687 if d > max_beta_delta {
2688 max_beta_delta = d;
2689 }
2690 }
2691 let mut max_fit_delta = 0.0_f64;
2692 for i in 0..n {
2693 let d = (yhat_cpu[i] - yhat_gpu[i]).abs();
2694 if d > max_fit_delta {
2695 max_fit_delta = d;
2696 }
2697 }
2698
2699 eprintln!(
2700 "[sphere_gpu fit parity] n={n} m={m} p={p} lmax={lmax} λ={lambda:.1e} \
2701 raw_xs|Δ|={raw_xs_delta:.3e} cond={cond:.3e} \
2702 max|Δβ|={max_beta_delta:.3e} max|Δŷ|={max_fit_delta:.3e}"
2703 );
2704
2705 assert!(
2712 max_fit_delta <= 1.0e-9,
2713 "GPU vs CPU truncated-spectral fitted-value max |Δ| = {max_fit_delta:.3e} > 1e-9"
2714 );
2715
2716 let beta_tol = (1e-15 * cond * (1.0 + xs_scale)).max(1e-9) * 16.0;
2728 assert!(
2729 max_beta_delta <= beta_tol,
2730 "GPU vs CPU truncated-spectral coefficient max |Δ| = {max_beta_delta:.3e} > \
2731 condition-aware tol {beta_tol:.3e} (cond={cond:.3e}). Raw design parity is \
2732 {raw_xs_delta:.3e}; a drift THIS much larger than cond·ULP is a real solve/kernel \
2733 mismatch, not conditioning."
2734 );
2735 }
2736}