1use ndarray::{Array1, Array2, ArrayView2};
21
22use crate::arrow_schur::{ArrowPcgDiagnostics, ArrowSchurSystem, DeviceSaePcgData};
23use gam_linalg::triangular::{CholeskyGuard, cholesky_factor_in_place, cholesky_solve_vector};
24
25pub struct ArrowSchurGpuSolution {
27 pub delta_t: Array1<f64>,
28 pub delta_beta: Array1<f64>,
29 pub log_det_hessian: f64,
32}
33
34#[derive(Debug, Clone)]
38pub enum ArrowSchurGpuFailure {
39 Unavailable,
41 RidgeBumpRequired { row: usize, bump: f64 },
44 SchurFactorFailed { reason: String },
47 GpuRequiresDenseSystem {
58 had_hbb_matvec: bool,
59 had_htbeta_matvec: bool,
60 },
61}
62
63#[cfg(target_os = "linux")]
74fn resolve_runtime_for_device_path(
75) -> Result<Option<&'static gam_gpu::GpuRuntime>, ArrowSchurGpuFailure> {
76 gam_gpu::device_runtime::GpuRuntime::resolve(gam_gpu::global_policy()).map_err(|error| {
77 ArrowSchurGpuFailure::SchurFactorFailed {
78 reason: format!("GPU runtime resolution failed: {error}"),
79 }
80 })
81}
82
83const RIDGE_BUMP_EPS_MARGIN: f64 = 1024.0;
94
95#[must_use]
131fn ridge_bump_to_make_pd(htt: ArrayView2<'_, f64>, ridge_t: f64) -> f64 {
132 let d = htt.nrows();
133 let mut scale = 1.0_f64;
136 let mut min_gershgorin_edge = f64::INFINITY;
137 for i in 0..d {
138 let diag = htt[[i, i]];
139 scale = scale.max(diag.abs());
140 let mut off_sum = 0.0_f64;
141 for j in 0..d {
142 if j != i {
143 off_sum += htt[[i, j]].abs();
144 }
145 }
146 min_gershgorin_edge = min_gershgorin_edge.min(diag - off_sum);
147 }
148 if !min_gershgorin_edge.is_finite() {
149 return scale * f64::EPSILON.sqrt() * RIDGE_BUMP_EPS_MARGIN;
152 }
153 let deficit = -(min_gershgorin_edge + ridge_t);
156 let margin = scale * f64::EPSILON.sqrt() * RIDGE_BUMP_EPS_MARGIN;
157 deficit.max(0.0) + margin
160}
161
162#[cfg(target_os = "linux")]
175#[must_use]
176fn ridge_bump_to_make_pd_colmajor(block: &[f64], d: usize) -> f64 {
177 if d == 0 || block.len() < d * d {
178 return f64::EPSILON.sqrt() * RIDGE_BUMP_EPS_MARGIN;
179 }
180 let mut scale = 1.0_f64;
183 let mut min_gershgorin_edge = f64::INFINITY;
184 for i in 0..d {
185 let diag = block[i * d + i];
186 scale = scale.max(diag.abs());
187 let mut off_sum = 0.0_f64;
188 for j in 0..d {
189 if j != i {
190 off_sum += block[j * d + i].abs();
191 }
192 }
193 min_gershgorin_edge = min_gershgorin_edge.min(diag - off_sum);
194 }
195 let margin = scale * f64::EPSILON.sqrt() * RIDGE_BUMP_EPS_MARGIN;
196 if !min_gershgorin_edge.is_finite() {
197 return margin;
198 }
199 (-min_gershgorin_edge).max(0.0) + margin
200}
201
202pub fn solve_arrow_newton_step(
206 sys: &ArrowSchurSystem,
207 ridge_t: f64,
208 ridge_beta: f64,
209) -> Result<ArrowSchurGpuSolution, ArrowSchurGpuFailure> {
210 let n = sys.rows.len();
211 let d = sys.d;
212 let k = sys.k;
213
214 let had_hbb_matvec = sys.hbb_matvec.is_some();
219 let had_htbeta_matvec = sys.htbeta_matvec.is_some();
220 if had_hbb_matvec || had_htbeta_matvec {
221 return Err(ArrowSchurGpuFailure::GpuRequiresDenseSystem {
222 had_hbb_matvec,
223 had_htbeta_matvec,
224 });
225 }
226
227 if sys.penalty_op.is_some() {
244 return Err(ArrowSchurGpuFailure::GpuRequiresDenseSystem {
245 had_hbb_matvec: false,
246 had_htbeta_matvec: false,
247 });
248 }
249
250 if sys.hbb.dim() != (k, k) {
251 return Err(ArrowSchurGpuFailure::GpuRequiresDenseSystem {
265 had_hbb_matvec: false,
266 had_htbeta_matvec: false,
267 });
268 }
269 if n == 0 || d == 0 {
270 return Err(ArrowSchurGpuFailure::Unavailable);
271 }
272 if sys
273 .rows
274 .iter()
275 .any(|row| row.htt.dim() != (d, d) || row.htbeta.dim() != (d, k) || row.gt.len() != d)
276 {
277 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
278 reason: "row block dimension mismatch".to_string(),
279 });
280 }
281
282 #[cfg(not(target_os = "linux"))]
283 {
284 if ridge_t.is_nan() || ridge_beta.is_nan() {
285 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
286 reason: "ridge is NaN".to_string(),
287 });
288 }
289 Err(ArrowSchurGpuFailure::Unavailable)
290 }
291
292 #[cfg(target_os = "linux")]
293 {
294 if resolve_runtime_for_device_path()?
303 .map(gam_gpu::device_runtime::GpuRuntime::device_count)
304 .unwrap_or(0)
305 > 1
306 {
307 match cuda::solve_multi_gpu(sys, ridge_t, ridge_beta) {
308 Ok(sol) => return Ok(sol),
309 Err(ArrowSchurGpuFailure::RidgeBumpRequired { row, bump }) => {
310 return Err(ArrowSchurGpuFailure::RidgeBumpRequired { row, bump });
311 }
312 Err(ArrowSchurGpuFailure::SchurFactorFailed { reason }) => {
313 return Err(ArrowSchurGpuFailure::SchurFactorFailed { reason });
314 }
315 Err(_) => {}
318 }
319 }
320 if crate::gpu_kernels::arrow_schur_nvrtc::system_admits_fused_path(sys) {
326 match cuda::solve_fused(sys, ridge_t, ridge_beta) {
327 Ok(sol) => return Ok(sol),
328 Err(ArrowSchurGpuFailure::RidgeBumpRequired { row, bump }) => {
332 return Err(ArrowSchurGpuFailure::RidgeBumpRequired { row, bump });
333 }
334 Err(_) => {}
338 }
339 }
340 cuda::solve(sys, ridge_t, ridge_beta)
341 }
342}
343
344#[cfg(target_os = "linux")]
350fn pack_host(sys: &ArrowSchurSystem, ridge_t: f64) -> (Vec<f64>, Vec<f64>, Vec<f64>) {
351 let n = sys.rows.len();
352 let d = sys.d;
353 let k = sys.k;
354 let mut d_buf = Vec::with_capacity(n * d * d);
355 let mut b_buf = Vec::with_capacity(n * d * k);
356 let mut g_buf = Vec::with_capacity(n * d);
357 for row in &sys.rows {
358 pack_block(row, ridge_t, d, k, &mut d_buf, &mut b_buf, &mut g_buf);
359 }
360 (d_buf, b_buf, g_buf)
361}
362
363#[cfg(target_os = "linux")]
364#[inline]
365fn pack_block(
366 row: &crate::arrow_schur::ArrowRowBlock,
367 ridge_t: f64,
368 d: usize,
369 k: usize,
370 d_buf: &mut Vec<f64>,
371 b_buf: &mut Vec<f64>,
372 g_buf: &mut Vec<f64>,
373) {
374 for col in 0..d {
375 for r in 0..d {
376 let mut value = row.htt[[r, col]];
377 if r == col {
378 value += ridge_t;
379 }
380 d_buf.push(value);
381 }
382 }
383 for col in 0..k {
384 for r in 0..d {
385 b_buf.push(row.htbeta[[r, col]]);
386 }
387 }
388 for r in 0..d {
389 g_buf.push(row.gt[r]);
390 }
391}
392
393#[doc(hidden)]
399#[cfg(target_os = "linux")]
400pub fn solve_arrow_newton_step_fused_force(
401 sys: &ArrowSchurSystem,
402 ridge_t: f64,
403 ridge_beta: f64,
404) -> Result<ArrowSchurGpuSolution, ArrowSchurGpuFailure> {
405 if ridge_t.is_nan() || ridge_beta.is_nan() {
406 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
407 reason: "ridge is NaN".to_string(),
408 });
409 }
410 if crate::gpu_kernels::arrow_schur_nvrtc::plan_fused_launch(sys.rows.len(), sys.d, sys.k)
411 .is_none()
412 {
413 return Err(ArrowSchurGpuFailure::Unavailable);
414 }
415 cuda::solve_fused(sys, ridge_t, ridge_beta)
416}
417
418#[cfg(target_os = "linux")]
428pub struct ResidentArrowFrameHandle {
429 inner: cuda::ResidentArrowFrame,
430}
431
432#[cfg(not(target_os = "linux"))]
436pub enum ResidentArrowFrameHandle {}
437
438impl ResidentArrowFrameHandle {
439 pub fn new(
441 sys: &ArrowSchurSystem,
442 ridge_t: f64,
443 ridge_beta: f64,
444 ) -> Result<Self, ArrowSchurGpuFailure> {
445 if sys.hbb_matvec.is_some() || sys.htbeta_matvec.is_some() {
448 return Err(ArrowSchurGpuFailure::GpuRequiresDenseSystem {
449 had_hbb_matvec: sys.hbb_matvec.is_some(),
450 had_htbeta_matvec: sys.htbeta_matvec.is_some(),
451 });
452 }
453 #[cfg(not(target_os = "linux"))]
454 {
455 if ridge_t.is_nan() || ridge_beta.is_nan() {
456 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
457 reason: "ridge is NaN".to_string(),
458 });
459 }
460 Err(ArrowSchurGpuFailure::Unavailable)
461 }
462 #[cfg(target_os = "linux")]
463 {
464 Ok(Self {
465 inner: cuda::ResidentArrowFrame::new(sys, ridge_t, ridge_beta)?,
466 })
467 }
468 }
469
470 pub fn solve_gradient(
472 &self,
473 g_t: &[f64],
474 g_beta: &[f64],
475 ) -> Result<ArrowSchurGpuSolution, ArrowSchurGpuFailure> {
476 #[cfg(not(target_os = "linux"))]
477 {
478 if g_t.iter().chain(g_beta).any(|v| !v.is_finite()) {
479 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
480 reason: "non-finite gradient entry".to_string(),
481 });
482 }
483 Err(ArrowSchurGpuFailure::Unavailable)
484 }
485 #[cfg(target_os = "linux")]
486 {
487 self.inner.solve_gradient(g_t, g_beta)
488 }
489 }
490
491 #[must_use]
493 pub fn log_det_hessian(&self) -> f64 {
494 #[cfg(not(target_os = "linux"))]
495 {
496 match *self {}
497 }
498 #[cfg(target_os = "linux")]
499 {
500 self.inner.log_det_hessian()
501 }
502 }
503}
504
505#[cfg(target_os = "linux")]
520pub struct ResidentBaseArrowFrameHandle {
521 inner: cuda::ResidentBaseArrowFrame,
522}
523
524#[cfg(not(target_os = "linux"))]
527pub enum ResidentBaseArrowFrameHandle {}
528
529impl ResidentBaseArrowFrameHandle {
530 pub fn new(sys: &ArrowSchurSystem) -> Result<Self, ArrowSchurGpuFailure> {
536 if sys.hbb_matvec.is_some() || sys.htbeta_matvec.is_some() {
537 return Err(ArrowSchurGpuFailure::GpuRequiresDenseSystem {
538 had_hbb_matvec: sys.hbb_matvec.is_some(),
539 had_htbeta_matvec: sys.htbeta_matvec.is_some(),
540 });
541 }
542 #[cfg(not(target_os = "linux"))]
543 {
544 Err(ArrowSchurGpuFailure::Unavailable)
545 }
546 #[cfg(target_os = "linux")]
547 {
548 Ok(Self {
549 inner: cuda::ResidentBaseArrowFrame::new(sys)?,
550 })
551 }
552 }
553
554 pub fn refactor_and_solve(
561 &self,
562 ridge_t: f64,
563 ridge_beta: f64,
564 ) -> Result<ArrowSchurGpuSolution, ArrowSchurGpuFailure> {
565 #[cfg(not(target_os = "linux"))]
566 {
567 if ridge_t.is_nan() || ridge_beta.is_nan() {
568 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
569 reason: "ridge is NaN".to_string(),
570 });
571 }
572 Err(ArrowSchurGpuFailure::Unavailable)
573 }
574 #[cfg(target_os = "linux")]
575 {
576 self.inner.refactor_and_solve(ridge_t, ridge_beta)
577 }
578 }
579}
580
581pub fn gpu_schur_matvec_backend(
624 sys: &ArrowSchurSystem,
625 ridge_t: f64,
626 ridge_beta: f64,
627) -> Result<crate::arrow_schur::GpuSchurMatvec, ArrowSchurGpuFailure> {
628 if sys.htbeta_matvec.is_some() {
631 return build_row_procedural_matvec(sys, ridge_t, ridge_beta);
632 }
633
634 #[cfg(not(target_os = "linux"))]
635 {
636 if ridge_t.is_nan() || ridge_beta.is_nan() {
639 return Err(ArrowSchurGpuFailure::Unavailable);
640 }
641 Err(ArrowSchurGpuFailure::Unavailable)
642 }
643
644 #[cfg(target_os = "linux")]
645 {
646 cuda::build_schur_matvec_backend(sys, ridge_t, ridge_beta)
647 }
648}
649
650pub fn build_framed_resident_evidence_matvec(
658 sys: &ArrowSchurSystem,
659 ridge_t: f64,
660 ridge_beta: f64,
661 apply_budget: usize,
662) -> Result<Option<crate::arrow_schur::GpuSchurMatvec>, ArrowSchurGpuFailure> {
663 #[cfg(not(target_os = "linux"))]
664 {
665 if sys.k == 0 || !ridge_t.is_finite() || !ridge_beta.is_finite() || apply_budget == 0 {
669 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
670 reason: "resident evidence matvec received degenerate dimensions or ridge"
671 .to_string(),
672 });
673 }
674 Ok(None)
675 }
676 #[cfg(target_os = "linux")]
677 {
678 cuda::build_framed_resident_evidence_matvec(sys, ridge_t, ridge_beta, apply_budget)
679 }
680}
681
682fn build_row_procedural_matvec(
699 sys: &ArrowSchurSystem,
700 ridge_t: f64,
701 ridge_beta: f64,
702) -> Result<crate::arrow_schur::GpuSchurMatvec, ArrowSchurGpuFailure> {
703 use std::sync::Arc;
704 let n = sys.rows.len();
705 let k = sys.k;
706 let forward = sys
707 .htbeta_matvec
708 .clone()
709 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
710 let transpose = sys.htbeta_transpose_matvec.clone().ok_or_else(|| {
711 ArrowSchurGpuFailure::SchurFactorFailed {
716 reason: "row-procedural Schur matvec requires htbeta_transpose_matvec; \
717 forward operator installed without its sparse adjoint"
718 .to_string(),
719 }
720 })?;
721
722 let mut factors: Vec<Array2<f64>> = Vec::with_capacity(n);
727 for (i, row) in sys.rows.iter().enumerate() {
728 let di = row.htt.nrows();
729 if row.htt.ncols() != di || row.gt.len() != di {
730 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
731 reason: format!("row {i}: malformed H_tt block {:?}", row.htt.dim()),
732 });
733 }
734 let mut block = row.htt.clone();
735 for r in 0..di {
736 block[[r, r]] += ridge_t;
737 }
738 let factor = cholesky_factor_in_place(block.view(), CholeskyGuard::NonnegativePivot)
739 .ok_or_else(|| {
740 ArrowSchurGpuFailure::RidgeBumpRequired {
744 row: i,
745 bump: ridge_bump_to_make_pd(row.htt.view(), ridge_t),
746 }
747 })?;
748 factors.push(factor);
749 }
750
751 let penalty_op = sys.effective_penalty_op();
758 let row_dims: Vec<usize> = sys.rows.iter().map(|row| row.htt.nrows()).collect();
759
760 let closure: crate::arrow_schur::GpuSchurMatvec =
761 Arc::new(move |x: &Array1<f64>, out: &mut Array1<f64>| {
762 assert_eq!(x.len(), k, "row-procedural matvec: x.len() != k");
763 assert_eq!(out.len(), k, "row-procedural matvec: out.len() != k");
764
765 {
768 let x_slice = x.as_slice().expect("x must be contiguous");
769 let out_slice = out.as_slice_mut().expect("out must be contiguous");
770 for a in 0..k {
771 out_slice[a] = ridge_beta * x_slice[a];
772 }
773 penalty_op.matvec(x_slice, out_slice);
774 }
775
776 let parallel = n >= crate::arrow_schur::SCHUR_MATVEC_PARALLEL_ROW_MIN
803 && rayon::current_thread_index().is_none();
804 if parallel {
805 use rayon::prelude::*;
806 const CHUNK: usize = 64;
807 let partials: Vec<Array1<f64>> = (0..n)
808 .into_par_iter()
809 .chunks(CHUNK)
810 .map(|idxs| {
811 let mut neg = Array1::<f64>::zeros(k);
815 for i in idxs {
816 let di = row_dims[i];
817 let mut v_i = Array1::<f64>::zeros(di);
819 forward(i, x.view(), &mut v_i);
820 let w_i = cholesky_solve_vector(factors[i].view(), v_i.view());
822 transpose(i, w_i.view(), &mut neg);
824 }
825 neg
826 })
827 .collect();
828 let mut neg = Array1::<f64>::zeros(k);
840 for part in &partials {
841 for a in 0..k {
842 neg[a] += part[a];
843 }
844 }
845 for a in 0..k {
846 out[a] -= neg[a];
847 }
848 } else {
849 let mut neg = Array1::<f64>::zeros(k);
851 for i in 0..n {
852 let di = row_dims[i];
853 let mut v_i = Array1::<f64>::zeros(di);
855 forward(i, x.view(), &mut v_i);
856 let w_i = cholesky_solve_vector(factors[i].view(), v_i.view());
858 transpose(i, w_i.view(), &mut neg);
860 }
861 for a in 0..k {
862 out[a] -= neg[a];
863 }
864 }
865 });
866
867 Ok(closure)
868}
869
870pub fn solve_reduced_beta_pcg(
892 s_acc: &Array2<f64>,
893 rhs_beta: &Array1<f64>,
894 max_iterations: usize,
895 relative_tolerance: f64,
896) -> Result<Array1<f64>, ArrowSchurGpuFailure> {
897 solve_reduced_beta_pcg_with_diagnostics(s_acc, rhs_beta, max_iterations, relative_tolerance)
898 .map(|(x, _)| x)
899}
900
901#[doc(hidden)]
902pub fn solve_reduced_beta_pcg_with_diagnostics(
903 s_acc: &Array2<f64>,
904 rhs_beta: &Array1<f64>,
905 max_iterations: usize,
906 relative_tolerance: f64,
907) -> Result<(Array1<f64>, ArrowPcgDiagnostics), ArrowSchurGpuFailure> {
908 let k = rhs_beta.len();
909 if s_acc.dim() != (k, k) {
910 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
911 reason: format!(
912 "reduced-β GPU PCG requires a square (k×k) Schur block; got {:?} for k={k}",
913 s_acc.dim()
914 ),
915 });
916 }
917 if k == 0 {
918 return Err(ArrowSchurGpuFailure::Unavailable);
919 }
920
921 #[cfg(not(target_os = "linux"))]
922 {
923 if relative_tolerance.is_nan() || max_iterations == 0 {
924 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
925 reason: "reduced-β GPU PCG: invalid CG controls".to_string(),
926 });
927 }
928 Err(ArrowSchurGpuFailure::Unavailable)
929 }
930
931 #[cfg(target_os = "linux")]
932 {
933 cuda::solve_reduced_beta_pcg_with_diagnostics(
934 s_acc,
935 rhs_beta,
936 max_iterations,
937 relative_tolerance,
938 )
939 }
940}
941
942pub fn solve_sae_matrix_free_pcg(
943 sys: &ArrowSchurSystem,
944 data: &DeviceSaePcgData,
945 ridge_t: f64,
946 ridge_beta: f64,
947 rhs_beta: &Array1<f64>,
948 max_iterations: usize,
949 relative_tolerance: f64,
950) -> Result<(Array1<f64>, ArrowPcgDiagnostics), ArrowSchurGpuFailure> {
951 if sys.k != data.beta_dim || rhs_beta.len() != data.beta_dim || data.p == 0 {
952 return Err(ArrowSchurGpuFailure::Unavailable);
953 }
954 #[cfg(not(target_os = "linux"))]
955 {
956 if ridge_t.is_nan()
957 || ridge_beta.is_nan()
958 || relative_tolerance.is_nan()
959 || max_iterations == 0
960 {
961 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
962 reason: "SAE matrix-free GPU PCG: invalid controls".to_string(),
963 });
964 }
965 Err(ArrowSchurGpuFailure::Unavailable)
966 }
967 #[cfg(target_os = "linux")]
968 {
969 if data.frame.is_some() {
976 cuda::solve_sae_matrix_free_pcg_framed(
977 sys,
978 data,
979 ridge_t,
980 ridge_beta,
981 rhs_beta,
982 max_iterations,
983 relative_tolerance,
984 )
985 } else {
986 cuda::solve_sae_matrix_free_pcg(
987 sys,
988 data,
989 ridge_t,
990 ridge_beta,
991 rhs_beta,
992 max_iterations,
993 relative_tolerance,
994 )
995 }
996 }
997}
998
999pub trait SaeResidentFrame {
1019 fn refresh(&self, sys: &ArrowSchurSystem) -> Result<(), ArrowSchurGpuFailure>;
1025
1026 fn resolve(
1034 &self,
1035 sys: &ArrowSchurSystem,
1036 ridge_t: f64,
1037 ridge_beta: f64,
1038 rhs_beta: &Array1<f64>,
1039 max_iterations: usize,
1040 relative_tolerance: f64,
1041 ) -> Result<(Array1<f64>, ArrowPcgDiagnostics), ArrowSchurGpuFailure>;
1042}
1043
1044pub fn build_sae_resident_frame(
1052 sys: &ArrowSchurSystem,
1053 cg_iters: usize,
1054) -> Result<
1055 Option<std::sync::Arc<dyn SaeResidentFrame + Send + Sync>>,
1056 ArrowSchurGpuFailure,
1057> {
1058 if sys.k == 0 || cg_iters == 0 {
1063 return Ok(None);
1064 }
1065 #[cfg(target_os = "linux")]
1066 {
1067 cuda::ResidentSaeFrameHandle::build(sys, cg_iters).map(|handle| {
1068 handle.map(|frame| {
1069 std::sync::Arc::new(frame)
1070 as std::sync::Arc<dyn SaeResidentFrame + Send + Sync>
1071 })
1072 })
1073 }
1074 #[cfg(not(target_os = "linux"))]
1076 {
1077 Ok(None)
1078 }
1079}
1080
1081#[cfg(target_os = "linux")]
1089pub struct FrameHostOperands {
1090 pub s_off: Vec<i32>,
1091 pub s_m: Vec<i32>,
1092 pub s_r: Vec<i32>,
1093 pub s_ptr: Vec<i32>,
1094 pub s_data: Vec<f64>,
1095 pub s_blocks: usize,
1096 pub g_off_i: Vec<i32>,
1097 pub g_off_j: Vec<i32>,
1098 pub g_ri: Vec<i32>,
1099 pub g_rj: Vec<i32>,
1100 pub g_mi: Vec<i32>,
1101 pub g_mj: Vec<i32>,
1102 pub g_ptr: Vec<i32>,
1103 pub g_data: Vec<f64>,
1104 pub w_ptr: Vec<i32>,
1105 pub w_data: Vec<f64>,
1106 pub g_blocks: usize,
1107 pub g_max_work: usize,
1108 pub htb_ptr: Vec<i32>,
1109 pub htb: Vec<f64>,
1110 pub q_of: Vec<i32>,
1111 pub n_rows: usize,
1112 pub k: usize,
1113 pub max_q: usize,
1114}
1115
1116#[cfg(target_os = "linux")]
1117fn frame_checked_i32(value: usize) -> Result<i32, ArrowSchurGpuFailure> {
1118 i32::try_from(value).map_err(|_| ArrowSchurGpuFailure::Unavailable)
1119}
1120
1121#[cfg(target_os = "linux")]
1126pub fn flatten_frame_host_operands(
1127 sys: &ArrowSchurSystem,
1128 data: &DeviceSaePcgData,
1129 frame: &crate::arrow_schur::DeviceSaeFrameData,
1130) -> Result<FrameHostOperands, ArrowSchurGpuFailure> {
1131 let n_rows = sys.rows.len();
1132 let k = data.beta_dim;
1133 if frame.row_htbeta.len() != n_rows
1134 || frame.ranks.len() != frame.basis_sizes.len()
1135 || frame.border_offsets.len() != frame.ranks.len()
1136 || data.smooth_blocks.len() != frame.smooth_ranks.len()
1137 {
1138 return Err(ArrowSchurGpuFailure::Unavailable);
1139 }
1140
1141 let mut s_off = Vec::new();
1143 let mut s_m = Vec::new();
1144 let mut s_r = Vec::new();
1145 let mut s_ptr = vec![0_i32];
1146 let mut s_data = Vec::<f64>::new();
1147 for (blk, &r) in data.smooth_blocks.iter().zip(frame.smooth_ranks.iter()) {
1148 let (m, mc) = blk.factor_a.dim();
1149 if m != mc {
1150 return Err(ArrowSchurGpuFailure::Unavailable);
1151 }
1152 s_off.push(frame_checked_i32(blk.global_offset)?);
1153 s_m.push(frame_checked_i32(m)?);
1154 s_r.push(frame_checked_i32(r)?);
1155 for ri in 0..m {
1156 for ci in 0..m {
1157 s_data.push(blk.factor_a[[ri, ci]]);
1158 }
1159 }
1160 s_ptr.push(frame_checked_i32(s_data.len())?);
1161 }
1162
1163 let mut g_off_i = Vec::new();
1165 let mut g_off_j = Vec::new();
1166 let mut g_ri = Vec::new();
1167 let mut g_rj = Vec::new();
1168 let mut g_mi = Vec::new();
1169 let mut g_mj = Vec::new();
1170 let mut g_ptr = vec![0_i32];
1171 let mut g_data = Vec::<f64>::new();
1172 let mut w_ptr = vec![0_i32];
1173 let mut w_data = Vec::<f64>::new();
1174 let mut g_max_work = 0usize;
1175 for blk in &frame.frame_blocks {
1176 let ri = frame.ranks[blk.atom_i];
1177 let rj = frame.ranks[blk.atom_j];
1178 let (mi, mj) = blk.g.dim();
1179 if blk.w.dim() != (ri, rj) {
1180 return Err(ArrowSchurGpuFailure::Unavailable);
1181 }
1182 g_off_i.push(frame_checked_i32(frame.border_offsets[blk.atom_i])?);
1183 g_off_j.push(frame_checked_i32(frame.border_offsets[blk.atom_j])?);
1184 g_ri.push(frame_checked_i32(ri)?);
1185 g_rj.push(frame_checked_i32(rj)?);
1186 g_mi.push(frame_checked_i32(mi)?);
1187 g_mj.push(frame_checked_i32(mj)?);
1188 for r in 0..mi {
1189 for c in 0..mj {
1190 g_data.push(blk.g[[r, c]]);
1191 }
1192 }
1193 g_ptr.push(frame_checked_i32(g_data.len())?);
1194 for a in 0..ri {
1195 for b in 0..rj {
1196 w_data.push(blk.w[[a, b]]);
1197 }
1198 }
1199 w_ptr.push(frame_checked_i32(w_data.len())?);
1200 g_max_work = g_max_work.max(mi * ri);
1201 }
1202
1203 let mut htb_ptr = vec![0_i32];
1206 let mut htb = Vec::<f64>::new();
1207 let mut q_of = Vec::<i32>::with_capacity(n_rows);
1208 let mut max_q = 0usize;
1209 for (i, slab) in frame.row_htbeta.iter().enumerate() {
1210 let qi = sys.row_dims[i];
1211 let q_eff = if !slab.is_empty() && slab.len() == qi * k {
1212 qi
1213 } else {
1214 0
1215 };
1216 q_of.push(frame_checked_i32(q_eff)?);
1217 max_q = max_q.max(q_eff);
1218 if q_eff > 0 {
1219 htb.extend_from_slice(slab);
1220 }
1221 htb_ptr.push(frame_checked_i32(htb.len())?);
1222 }
1223 if max_q == 0 {
1224 max_q = 1;
1227 }
1228
1229 Ok(FrameHostOperands {
1230 s_off,
1231 s_m,
1232 s_r,
1233 s_ptr,
1234 s_data,
1235 s_blocks: data.smooth_blocks.len(),
1236 g_off_i,
1237 g_off_j,
1238 g_ri,
1239 g_rj,
1240 g_mi,
1241 g_mj,
1242 g_ptr,
1243 g_data,
1244 w_ptr,
1245 w_data,
1246 g_blocks: frame.frame_blocks.len(),
1247 g_max_work,
1248 htb_ptr,
1249 htb,
1250 q_of,
1251 n_rows,
1252 k,
1253 max_q,
1254 })
1255}
1256
1257#[cfg(target_os = "linux")]
1264pub fn compute_ainv_host(
1265 sys: &ArrowSchurSystem,
1266 q_of: &[i32],
1267 max_q: usize,
1268 n_rows: usize,
1269 ridge_t: f64,
1270) -> Result<Vec<f64>, ArrowSchurGpuFailure> {
1271 let mut ainv = vec![0.0_f64; n_rows * max_q * max_q];
1272 for (i, row) in sys.rows.iter().enumerate() {
1273 let q = q_of[i] as usize;
1274 if q == 0 {
1275 continue;
1276 }
1277 if row.htt.dim() != (q, q) {
1278 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
1279 reason: format!(
1280 "framed SAE device PCG row {i}: H_tt shape {:?} != ({q}, {q})",
1281 row.htt.dim()
1282 ),
1283 });
1284 }
1285 let mut block = row.htt.clone();
1286 for d in 0..q {
1287 block[[d, d]] += ridge_t;
1288 }
1289 let factor = cholesky_factor_in_place(block.view(), CholeskyGuard::NonnegativePivot)
1290 .ok_or_else(|| ArrowSchurGpuFailure::RidgeBumpRequired {
1291 row: i,
1292 bump: ridge_bump_to_make_pd(row.htt.view(), ridge_t),
1293 })?;
1294 for col in 0..q {
1295 let mut e = Array1::<f64>::zeros(q);
1296 e[col] = 1.0;
1297 let solved = cholesky_solve_vector(factor.view(), e.view());
1298 for r in 0..q {
1299 ainv[i * max_q * max_q + r * max_q + col] = solved[r];
1300 }
1301 }
1302 }
1303 Ok(ainv)
1304}
1305
1306#[doc(hidden)]
1314#[cfg(target_os = "linux")]
1315pub fn framed_schur_matvec_once_on_device(
1316 sys: &ArrowSchurSystem,
1317 data: &DeviceSaePcgData,
1318 ridge_t: f64,
1319 ridge_beta: f64,
1320 x: &Array1<f64>,
1321) -> Result<Array1<f64>, ArrowSchurGpuFailure> {
1322 if sys.k != data.beta_dim || x.len() != data.beta_dim || data.p == 0 {
1323 return Err(ArrowSchurGpuFailure::Unavailable);
1324 }
1325 if data.frame.is_none() {
1326 return Err(ArrowSchurGpuFailure::Unavailable);
1327 }
1328 cuda::framed_schur_matvec_once_on_device(sys, data, ridge_t, ridge_beta, x)
1329}
1330
1331#[doc(hidden)]
1338#[cfg(target_os = "linux")]
1339pub fn framed_reduced_schur_det_once_on_device(
1340 sys: &ArrowSchurSystem,
1341 data: &DeviceSaePcgData,
1342 ridge_t: f64,
1343 ridge_beta: f64,
1344 x: &Array1<f64>,
1345) -> Result<Array1<f64>, ArrowSchurGpuFailure> {
1346 if sys.k != data.beta_dim || x.len() != data.beta_dim || data.p == 0 {
1347 return Err(ArrowSchurGpuFailure::Unavailable);
1348 }
1349 if data.frame.is_none() {
1350 return Err(ArrowSchurGpuFailure::Unavailable);
1351 }
1352 cuda::framed_reduced_schur_det_once_on_device(sys, data, ridge_t, ridge_beta, x)
1353}
1354
1355#[doc(hidden)]
1359pub fn solve_arrow_newton_step_dense_reference(
1360 sys: &ArrowSchurSystem,
1361 ridge_t: f64,
1362 ridge_beta: f64,
1363) -> Result<ArrowSchurGpuSolution, String> {
1364 let n = sys.rows.len();
1365 let d = sys.d;
1366 let k = sys.k;
1367 let total = n.checked_mul(d).ok_or("dimension overflow")? + k;
1368 let mut h = Array2::<f64>::zeros((total, total));
1369 let mut rhs = Array1::<f64>::zeros(total);
1370 for (i, row) in sys.rows.iter().enumerate() {
1371 let base = i * d;
1372 for c in 0..d {
1373 for r in 0..d {
1374 h[[base + r, base + c]] = row.htt[[r, c]];
1375 }
1376 h[[base + c, base + c]] += ridge_t;
1377 }
1378 for c in 0..k {
1379 for r in 0..d {
1380 let value = row.htbeta[[r, c]];
1381 h[[base + r, n * d + c]] = value;
1382 h[[n * d + c, base + r]] = value;
1383 }
1384 }
1385 for r in 0..d {
1386 rhs[base + r] = -row.gt[r];
1387 }
1388 }
1389 for c in 0..k {
1390 for r in 0..k {
1391 h[[n * d + r, n * d + c]] += sys.hbb[[r, c]];
1392 }
1393 h[[n * d + c, n * d + c]] += ridge_beta;
1394 rhs[n * d + c] = -sys.gb[c];
1395 }
1396 let factor = cholesky_factor_in_place(h.view(), CholeskyGuard::NonnegativePivot)
1405 .ok_or_else(|| "dense reference Cholesky failed".to_string())?;
1406 let mut log_det = 0.0_f64;
1407 for i in 0..total {
1408 log_det += factor[[i, i]].ln();
1409 }
1410 log_det *= 2.0;
1411 let solved = cholesky_solve_vector(factor.view(), rhs.view());
1412 let delta_t = solved.slice(ndarray::s![..n * d]).to_owned();
1413 let delta_beta = solved.slice(ndarray::s![n * d..]).to_owned();
1414 Ok(ArrowSchurGpuSolution {
1415 delta_t,
1416 delta_beta,
1417 log_det_hessian: log_det,
1418 })
1419}
1420
1421#[doc(hidden)]
1432pub fn sae_framed_penalty_matvec_cpu(
1433 data: &DeviceSaePcgData,
1434 ridge_beta: f64,
1435 x: &[f64],
1436 out: &mut [f64],
1437) {
1438 let frame = data
1439 .frame
1440 .as_ref()
1441 .expect("sae_framed_penalty_matvec_cpu requires frame metadata");
1442 let k = data.beta_dim;
1443 for a in 0..k {
1444 out[a] = ridge_beta * x[a];
1445 }
1446 for (blk, &r) in data.smooth_blocks.iter().zip(frame.smooth_ranks.iter()) {
1448 let off = blk.global_offset;
1449 let m = blk.factor_a.nrows();
1450 for i_a in 0..m {
1451 for i_b in 0..r {
1452 let mut acc = 0.0_f64;
1453 for j_a in 0..m {
1454 let s = blk.factor_a[[i_a, j_a]];
1455 if s == 0.0 {
1456 continue;
1457 }
1458 acc += s * x[off + j_a * r + i_b];
1459 }
1460 out[off + i_a * r + i_b] += acc;
1461 }
1462 }
1463 }
1464 for blk in &frame.frame_blocks {
1466 let r_i = frame.ranks[blk.atom_i];
1467 let r_j = frame.ranks[blk.atom_j];
1468 let off_i = frame.border_offsets[blk.atom_i];
1469 let off_j = frame.border_offsets[blk.atom_j];
1470 let (m_i, m_j) = blk.g.dim();
1471 for li in 0..m_i {
1472 let yi_base = off_i + li * r_i;
1473 for lj in 0..m_j {
1474 let g = blk.g[[li, lj]];
1475 if g == 0.0 {
1476 continue;
1477 }
1478 let xj_base = off_j + lj * r_j;
1479 for a in 0..r_i {
1480 let mut acc = 0.0_f64;
1481 for b in 0..r_j {
1482 acc += blk.w[[a, b]] * x[xj_base + b];
1483 }
1484 out[yi_base + a] += g * acc;
1485 }
1486 }
1487 }
1488 }
1489}
1490
1491#[doc(hidden)]
1500pub fn sae_framed_schur_matvec_cpu(
1501 sys: &ArrowSchurSystem,
1502 data: &DeviceSaePcgData,
1503 ridge_t: f64,
1504 ridge_beta: f64,
1505 x: &[f64],
1506 out: &mut [f64],
1507) -> Result<(), String> {
1508 let frame = data
1509 .frame
1510 .as_ref()
1511 .ok_or("sae_framed_schur_matvec_cpu requires frame metadata")?;
1512 let k = data.beta_dim;
1513 sae_framed_penalty_matvec_cpu(data, ridge_beta, x, out);
1514 if frame.row_htbeta.len() != sys.rows.len() {
1515 return Err(format!(
1516 "sae_framed_schur_matvec_cpu: {} row_htbeta slabs but {} rows",
1517 frame.row_htbeta.len(),
1518 sys.rows.len()
1519 ));
1520 }
1521 for (i, row) in sys.rows.iter().enumerate() {
1522 let slab = &frame.row_htbeta[i];
1523 if slab.is_empty() {
1524 continue;
1525 }
1526 let qi = sys.row_dims[i];
1527 if qi == 0 || slab.len() != qi * k {
1528 continue;
1529 }
1530 let mut h = vec![0.0_f64; qi];
1532 for c in 0..qi {
1533 let base = c * k;
1534 let mut acc = 0.0_f64;
1535 for a in 0..k {
1536 acc += slab[base + a] * x[a];
1537 }
1538 h[c] = acc;
1539 }
1540 let mut block = row.htt.clone();
1542 for d in 0..qi {
1543 block[[d, d]] += ridge_t;
1544 }
1545 let factor = cholesky_factor_in_place(block.view(), CholeskyGuard::NonnegativePivot)
1546 .ok_or_else(|| format!("sae_framed_schur_matvec_cpu: row {i} H_tt not PD"))?;
1547 let s = cholesky_solve_vector(factor.view(), Array1::from_vec(h).view());
1548 for c in 0..qi {
1550 let sc = s[c];
1551 if sc == 0.0 {
1552 continue;
1553 }
1554 let base = c * k;
1555 for a in 0..k {
1556 out[a] -= slab[base + a] * sc;
1557 }
1558 }
1559 }
1560 Ok(())
1561}
1562
1563#[cfg(target_os = "linux")]
1564mod cuda {
1565 use super::{ArrowSchurGpuFailure, ArrowSchurGpuSolution, pack_block, pack_host};
1566 use crate::arrow_schur::{
1567 ArrowPcgDiagnostics, ArrowSchurSystem, DeviceSaeFrameData, DeviceSaePcgData, PcgStopReason,
1568 };
1569 use cudarc::cublas::sys::{
1570 cublasDiagType_t, cublasFillMode_t, cublasOperation_t, cublasSideMode_t, cublasStatus_t,
1571 };
1572 use cudarc::cublas::{CudaBlas, Gemm, GemmConfig, Gemv, GemvConfig};
1573 use cudarc::cusolver::{DnHandle, sys as cusolver_sys};
1574 use cudarc::driver::{
1575 CudaContext, CudaModule, CudaSlice, CudaStream, DevicePtr, DevicePtrMut, LaunchConfig,
1576 PushKernelArg,
1577 };
1578 use gam_gpu::driver::to_i32;
1579 use gam_gpu::linalg_dispatch::{DispatchOp, route_through_gpu};
1580 use ndarray::Array1;
1581 use std::sync::{Arc, OnceLock};
1582
1583 struct RowSlot {
1588 d_block: Vec<f64>, b_block: Vec<f64>, g_vec: Vec<f64>, l_block: Vec<f64>, u_vec: Vec<f64>, y_block: Vec<f64>, log_det_local: f64,
1597 bump: Option<f64>,
1600 tile_partial_schur: Option<Vec<f64>>, tile_partial_rhs: Option<Vec<f64>>, delta_t_block: Vec<f64>, }
1606
1607 pub(super) fn solve_multi_gpu(
1628 sys: &ArrowSchurSystem,
1629 ridge_t: f64,
1630 ridge_beta: f64,
1631 ) -> Result<ArrowSchurGpuSolution, ArrowSchurGpuFailure> {
1632 let n = sys.rows.len();
1633 let d = sys.d;
1634 let k = sys.k;
1635 if n == 0 || d == 0 || k == 0 {
1636 return Err(ArrowSchurGpuFailure::Unavailable);
1637 }
1638 if sys.hbb_matvec.is_some() || sys.htbeta_matvec.is_some() || sys.hbb.dim() != (k, k) {
1642 return Err(ArrowSchurGpuFailure::Unavailable);
1643 }
1644
1645 let runtime = super::resolve_runtime_for_device_path()?
1646 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
1647 if runtime.device_count() < 2 {
1648 return Err(ArrowSchurGpuFailure::Unavailable);
1649 }
1650
1651 let mut slots: Vec<RowSlot> = Vec::with_capacity(n);
1653 for row in &sys.rows {
1654 if row.htt.dim() != (d, d) || row.htbeta.dim() != (d, k) || row.gt.len() != d {
1655 return Err(ArrowSchurGpuFailure::Unavailable);
1656 }
1657 let mut d_block = Vec::with_capacity(d * d);
1658 let mut b_block = Vec::with_capacity(d * k);
1659 let mut g_vec = Vec::with_capacity(d);
1660 pack_block(row, ridge_t, d, k, &mut d_block, &mut b_block, &mut g_vec);
1661 slots.push(RowSlot {
1662 d_block,
1663 b_block,
1664 g_vec,
1665 l_block: Vec::new(),
1666 u_vec: Vec::new(),
1667 y_block: Vec::new(),
1668 log_det_local: 0.0,
1669 bump: None,
1670 tile_partial_schur: None,
1671 tile_partial_rhs: None,
1672 delta_t_block: vec![0.0; d],
1673 });
1674 }
1675
1676 let forward_ok = gam_gpu::pool::scatter_batched(runtime, &mut slots, |ordinal, tile| {
1678 forward_tile(ordinal, d, k, tile)
1679 });
1680 if forward_ok.is_none() {
1681 return Err(ArrowSchurGpuFailure::Unavailable);
1682 }
1683
1684 let row_base_of_tile = gam_gpu::pool::balanced_partition(runtime, n);
1686 if let Some((row, bump)) = slots
1687 .iter()
1688 .enumerate()
1689 .find_map(|(i, slot)| slot.bump.map(|b| (i, b)))
1690 {
1691 return Err(ArrowSchurGpuFailure::RidgeBumpRequired { row, bump });
1692 }
1693
1694 let mut schur_host = vec![0.0_f64; k * k];
1699 for col in 0..k {
1700 for row in 0..k {
1701 let mut v = sys.hbb[[row, col]];
1702 if row == col {
1703 v += ridge_beta;
1704 }
1705 schur_host[col * k + row] = v;
1706 }
1707 }
1708 let mut rhs_host: Vec<f64> = sys.gb.iter().map(|v| -v).collect();
1709 let mut log_det = 0.0_f64;
1710 for start in tile_starts(&row_base_of_tile) {
1711 let slot = &slots[start];
1712 let partial_schur = slot
1713 .tile_partial_schur
1714 .as_ref()
1715 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
1716 let partial_rhs = slot
1717 .tile_partial_rhs
1718 .as_ref()
1719 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
1720 for idx in 0..k * k {
1725 schur_host[idx] += partial_schur[idx];
1726 }
1727 for a in 0..k {
1728 rhs_host[a] += partial_rhs[a];
1729 }
1730 }
1731 for slot in &slots {
1732 log_det += slot.log_det_local;
1733 }
1734
1735 let primary = runtime.selected_device().ordinal;
1739 let stream = gam_gpu::device_runtime::cuda_context_for(primary)
1740 .and_then(|ctx| ctx.new_stream().ok())
1741 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
1742 let solver =
1743 DnHandle::new(stream.clone()).map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
1744 let blas = CudaBlas::new(stream.clone()).map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
1745 let mut schur_dev = stream
1746 .clone_htod(&schur_host)
1747 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
1748 let mut rhs_dev = stream
1749 .clone_htod(&rhs_host)
1750 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
1751 let info = potrf_single(&solver, &stream, k, &mut schur_dev)?;
1752 if info != 0 {
1753 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
1754 reason: format!("multi-GPU Schur Cholesky failed at pivot {info}"),
1755 });
1756 }
1757 trsm_single(&blas, &stream, k, &schur_dev, &mut rhs_dev, false, false)?;
1758 trsm_single(&blas, &stream, k, &schur_dev, &mut rhs_dev, false, true)?;
1759 let delta_beta_host = stream
1760 .clone_dtoh(&rhs_dev)
1761 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
1762 let delta_beta = Array1::from_vec(delta_beta_host.clone());
1763 let l_schur_host = stream
1764 .clone_dtoh(&schur_dev)
1765 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
1766 for j in 0..k {
1767 log_det += l_schur_host[j * k + j].ln();
1768 }
1769 log_det *= 2.0;
1770
1771 let delta_beta_ref = &delta_beta_host;
1773 let back_ok = gam_gpu::pool::scatter_batched(runtime, &mut slots, |ordinal, tile| {
1774 back_sub_tile(ordinal, d, k, delta_beta_ref, tile)
1775 });
1776 if back_ok.is_none() {
1777 return Err(ArrowSchurGpuFailure::Unavailable);
1778 }
1779
1780 let mut delta_t = Array1::<f64>::zeros(n * d);
1782 for (i, slot) in slots.iter().enumerate() {
1783 let base = i * d;
1784 for r in 0..d {
1785 delta_t[base + r] = slot.delta_t_block[r];
1786 }
1787 }
1788
1789 Ok(ArrowSchurGpuSolution {
1790 delta_t,
1791 delta_beta,
1792 log_det_hessian: log_det,
1793 })
1794 }
1795
1796 fn tile_starts(tiles: &[(usize, std::ops::Range<usize>)]) -> impl Iterator<Item = usize> + '_ {
1799 tiles.iter().map(|(_, range)| range.start)
1800 }
1801
1802 fn forward_tile(ordinal: usize, d: usize, k: usize, tile: &mut [RowSlot]) -> Option<()> {
1810 if tile.is_empty() {
1811 return Some(());
1812 }
1813 let stream = gam_gpu::device_runtime::cuda_context_for(ordinal)
1816 .and_then(|ctx| ctx.new_stream().ok())?;
1817 let solver = DnHandle::new(stream.clone()).ok()?;
1818 let blas = CudaBlas::new(stream.clone()).ok()?;
1819 let m = tile.len();
1820
1821 let mut d_host = Vec::with_capacity(m * d * d);
1824 let mut b_host = Vec::with_capacity(m * d * k);
1825 let mut g_host = Vec::with_capacity(m * d);
1826 for slot in tile.iter() {
1827 d_host.extend_from_slice(&slot.d_block);
1828 b_host.extend_from_slice(&slot.b_block);
1829 g_host.extend_from_slice(&slot.g_vec);
1830 }
1831 let mut d_dev = stream.clone_htod(&d_host).ok()?;
1832 let mut b_dev = stream.clone_htod(&b_host).ok()?;
1833 let mut g_dev = stream.clone_htod(&g_host).ok()?;
1834
1835 let info_host = potrf_batched(&solver, &stream, d, m, &mut d_dev).ok()?;
1841 if let Some(local) = info_host.iter().position(|info| *info != 0) {
1842 tile[local].bump = Some(super::ridge_bump_to_make_pd_colmajor(
1843 &tile[local].d_block,
1844 d,
1845 ));
1846 return Some(());
1847 }
1848
1849 trsm_batched_lower_inplace(&blas, &stream, d, m, 1, &d_dev, &mut g_dev).ok()?;
1851 trsm_batched_lower_inplace(&blas, &stream, d, m, k, &d_dev, &mut b_dev).ok()?;
1852
1853 let mut schur_dev = stream.alloc_zeros::<f64>(k * k).ok()?;
1855 let mut rhs_dev = stream.alloc_zeros::<f64>(k).ok()?;
1856 accumulate_schur(&blas, d, k, m, &b_dev, &g_dev, &mut schur_dev, &mut rhs_dev).ok()?;
1857
1858 let l_host = stream.clone_dtoh(&d_dev).ok()?;
1860 let u_host = stream.clone_dtoh(&g_dev).ok()?;
1861 let y_host = stream.clone_dtoh(&b_dev).ok()?;
1862 let partial_schur = stream.clone_dtoh(&schur_dev).ok()?;
1863 let partial_rhs = stream.clone_dtoh(&rhs_dev).ok()?;
1864
1865 for (local, slot) in tile.iter_mut().enumerate() {
1866 let l_base = local * d * d;
1867 let u_base = local * d;
1868 let y_base = local * d * k;
1869 slot.l_block = l_host[l_base..l_base + d * d].to_vec();
1870 slot.u_vec = u_host[u_base..u_base + d].to_vec();
1871 slot.y_block = y_host[y_base..y_base + d * k].to_vec();
1872 let mut log_det_local = 0.0_f64;
1873 for j in 0..d {
1874 log_det_local += l_host[l_base + j * d + j].ln();
1875 }
1876 slot.log_det_local = log_det_local;
1877 }
1878 tile[0].tile_partial_schur = Some(partial_schur);
1879 tile[0].tile_partial_rhs = Some(partial_rhs);
1880 Some(())
1881 }
1882
1883 fn back_sub_tile(
1887 ordinal: usize,
1888 d: usize,
1889 k: usize,
1890 delta_beta: &[f64],
1891 tile: &mut [RowSlot],
1892 ) -> Option<()> {
1893 if tile.is_empty() {
1894 return Some(());
1895 }
1896 let stream = gam_gpu::device_runtime::cuda_context_for(ordinal)
1899 .and_then(|ctx| ctx.new_stream().ok())?;
1900 let blas = CudaBlas::new(stream.clone()).ok()?;
1901 let m = tile.len();
1902
1903 let mut l_host = Vec::with_capacity(m * d * d);
1904 let mut u_host = Vec::with_capacity(m * d);
1905 let mut y_host = Vec::with_capacity(m * d * k);
1906 for slot in tile.iter() {
1907 l_host.extend_from_slice(&slot.l_block);
1908 u_host.extend_from_slice(&slot.u_vec);
1909 y_host.extend_from_slice(&slot.y_block);
1910 }
1911 let d_dev = stream.clone_htod(&l_host).ok()?;
1912 let mut g_dev = stream.clone_htod(&u_host).ok()?;
1913 let b_dev = stream.clone_htod(&y_host).ok()?;
1914 let rhs_dev = stream.clone_htod(&delta_beta.to_vec()).ok()?;
1915
1916 accumulate_back_sub_rhs(&blas, d, k, m, &b_dev, &rhs_dev, &mut g_dev).ok()?;
1918 trsm_batched_lower_inplace_transposed(&blas, &stream, d, m, 1, &d_dev, &mut g_dev).ok()?;
1919 let x_host = stream.clone_dtoh(&g_dev).ok()?;
1920 for (local, slot) in tile.iter_mut().enumerate() {
1921 let base = local * d;
1922 for r in 0..d {
1923 slot.delta_t_block[r] = -x_host[base + r];
1924 }
1925 }
1926 Some(())
1927 }
1928
1929 pub(super) fn solve(
1930 sys: &ArrowSchurSystem,
1931 ridge_t: f64,
1932 ridge_beta: f64,
1933 ) -> Result<ArrowSchurGpuSolution, ArrowSchurGpuFailure> {
1934 let n = sys.rows.len();
1935 let d = sys.d;
1936 let k = sys.k;
1937 let runtime = route_through_gpu(DispatchOp::SmallDenseBatchedPotrf { p: d, batch: n })
1938 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
1939
1940 let stream = gam_gpu::device_runtime::cuda_context_for(runtime.device.ordinal)
1941 .and_then(|ctx| ctx.new_stream().ok())
1942 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
1943 let solver =
1944 DnHandle::new(stream.clone()).map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
1945 let blas = CudaBlas::new(stream.clone()).map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
1946
1947 let (d_host, b_host, g_host) = pack_host(sys, ridge_t);
1949 let mut d_dev = stream
1950 .clone_htod(&d_host)
1951 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
1952 let mut b_dev = stream
1953 .clone_htod(&b_host)
1954 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
1955 let mut g_dev = stream
1956 .clone_htod(&g_host)
1957 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
1958
1959 let info_host = potrf_batched(&solver, &stream, d, n, &mut d_dev)?;
1967 if let Some(idx) = info_host.iter().position(|info| *info != 0) {
1968 return Err(ArrowSchurGpuFailure::RidgeBumpRequired {
1972 row: idx,
1973 bump: super::ridge_bump_to_make_pd(sys.rows[idx].htt.view(), ridge_t),
1974 });
1975 }
1976
1977 trsm_batched_lower_inplace(&blas, &stream, d, n, 1, &d_dev, &mut g_dev)?;
1980 trsm_batched_lower_inplace(&blas, &stream, d, n, k, &d_dev, &mut b_dev)?;
1983
1984 let schur_init: Vec<f64> = {
2003 let mut tmp = Vec::with_capacity(k * k);
2004 for col in 0..k {
2005 for row in 0..k {
2006 let mut v = sys.hbb[[row, col]];
2007 if row == col {
2008 v += ridge_beta;
2009 }
2010 tmp.push(v);
2011 }
2012 }
2013 tmp
2014 };
2015 let mut schur_dev = stream
2016 .clone_htod(&schur_init)
2017 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2018 let rhs_init: Vec<f64> = sys.gb.iter().map(|v| -v).collect();
2019 let mut rhs_dev = stream
2020 .clone_htod(&rhs_init)
2021 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2022
2023 accumulate_schur(&blas, d, k, n, &b_dev, &g_dev, &mut schur_dev, &mut rhs_dev)?;
2024
2025 let info = potrf_single(&solver, &stream, k, &mut schur_dev)?;
2027 if info != 0 {
2028 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
2029 reason: format!("Schur Cholesky failed at pivot {info}"),
2030 });
2031 }
2032 trsm_single(&blas, &stream, k, &schur_dev, &mut rhs_dev, false, false)?;
2034 trsm_single(&blas, &stream, k, &schur_dev, &mut rhs_dev, false, true)?;
2035 let delta_beta_host = stream
2036 .clone_dtoh(&rhs_dev)
2037 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2038 let delta_beta = Array1::from_vec(delta_beta_host.clone());
2039
2040 accumulate_back_sub_rhs(&blas, d, k, n, &b_dev, &rhs_dev, &mut g_dev)?;
2048 trsm_batched_lower_inplace_transposed(&blas, &stream, d, n, 1, &d_dev, &mut g_dev)?;
2049
2050 let x_host = stream
2051 .clone_dtoh(&g_dev)
2052 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2053 let mut delta_t = Array1::<f64>::zeros(n * d);
2054 for (i, v) in x_host.iter().enumerate() {
2055 delta_t[i] = -*v;
2056 }
2057
2058 let l_local_host = stream
2060 .clone_dtoh(&d_dev)
2061 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2062 let l_schur_host = stream
2063 .clone_dtoh(&schur_dev)
2064 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2065 let mut log_det = 0.0_f64;
2066 for i in 0..n {
2067 let base = i * d * d;
2068 for j in 0..d {
2069 log_det += l_local_host[base + j * d + j].ln();
2070 }
2071 }
2072 for j in 0..k {
2073 log_det += l_schur_host[j * k + j].ln();
2074 }
2075 log_det *= 2.0;
2076
2077 Ok(ArrowSchurGpuSolution {
2078 delta_t,
2079 delta_beta,
2080 log_det_hessian: log_det,
2081 })
2082 }
2083
2084 fn potrf_batched(
2085 solver: &DnHandle,
2086 stream: &Arc<CudaStream>,
2087 p: usize,
2088 batch: usize,
2089 matrices: &mut CudaSlice<f64>,
2090 ) -> Result<Vec<i32>, ArrowSchurGpuFailure> {
2091 let p_i = to_i32(p).ok_or(ArrowSchurGpuFailure::Unavailable)?;
2092 let batch_i = to_i32(batch).ok_or(ArrowSchurGpuFailure::Unavailable)?;
2093 let matrix_len = p * p;
2094 let bytes_per = (matrix_len * std::mem::size_of::<f64>()) as u64;
2095 let (base_ptr, _record) = matrices.device_ptr_mut(stream);
2096 let mut ptrs = Vec::with_capacity(batch);
2097 for idx in 0..batch {
2098 ptrs.push(base_ptr + (idx as u64) * bytes_per);
2099 }
2100 let mut ptrs_dev = stream
2101 .clone_htod(&ptrs)
2102 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2103 let mut info_dev = stream
2104 .alloc_zeros::<i32>(batch)
2105 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2106 let status = {
2107 let (ptrs_ptr, _ptrs_record) = ptrs_dev.device_ptr_mut(stream);
2108 let (info_ptr, _info_record) = info_dev.device_ptr_mut(stream);
2109 unsafe {
2112 cusolver_sys::cusolverDnDpotrfBatched(
2113 solver.cu(),
2114 cusolver_sys::cublasFillMode_t::CUBLAS_FILL_MODE_LOWER,
2115 p_i,
2116 ptrs_ptr as *mut *mut f64,
2117 p_i,
2118 info_ptr as *mut i32,
2119 batch_i,
2120 )
2121 }
2122 };
2123 if status != cusolver_sys::cusolverStatus_t::CUSOLVER_STATUS_SUCCESS {
2124 return Err(ArrowSchurGpuFailure::Unavailable);
2125 }
2126 stream
2127 .clone_dtoh(&info_dev)
2128 .map_err(|_| ArrowSchurGpuFailure::Unavailable)
2129 }
2130
2131 fn potrf_single(
2132 solver: &DnHandle,
2133 stream: &Arc<CudaStream>,
2134 p: usize,
2135 matrix: &mut CudaSlice<f64>,
2136 ) -> Result<i32, ArrowSchurGpuFailure> {
2137 let p_i = to_i32(p).ok_or(ArrowSchurGpuFailure::Unavailable)?;
2138 let uplo = cusolver_sys::cublasFillMode_t::CUBLAS_FILL_MODE_LOWER;
2139 let mut lwork = 0_i32;
2140 {
2141 let (mat_ptr, _rec) = matrix.device_ptr_mut(stream);
2142 let status = unsafe {
2144 cusolver_sys::cusolverDnDpotrf_bufferSize(
2145 solver.cu(),
2146 uplo,
2147 p_i,
2148 mat_ptr as *mut f64,
2149 p_i,
2150 &mut lwork,
2151 )
2152 };
2153 if status != cusolver_sys::cusolverStatus_t::CUSOLVER_STATUS_SUCCESS {
2154 return Err(ArrowSchurGpuFailure::Unavailable);
2155 }
2156 }
2157 let lwork_usize = usize::try_from(lwork).map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2158 let mut workspace = stream
2159 .alloc_zeros::<f64>(lwork_usize.max(1))
2160 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2161 let mut info_dev = stream
2162 .alloc_zeros::<i32>(1)
2163 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2164 {
2165 let (mat_ptr, _rec) = matrix.device_ptr_mut(stream);
2166 let (work_ptr, _wrec) = workspace.device_ptr_mut(stream);
2167 let (info_ptr, _irec) = info_dev.device_ptr_mut(stream);
2168 let status = unsafe {
2170 cusolver_sys::cusolverDnDpotrf(
2171 solver.cu(),
2172 uplo,
2173 p_i,
2174 mat_ptr as *mut f64,
2175 p_i,
2176 work_ptr as *mut f64,
2177 lwork,
2178 info_ptr as *mut i32,
2179 )
2180 };
2181 if status != cusolver_sys::cusolverStatus_t::CUSOLVER_STATUS_SUCCESS {
2182 return Err(ArrowSchurGpuFailure::Unavailable);
2183 }
2184 }
2185 let info_host = stream
2186 .clone_dtoh(&info_dev)
2187 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2188 Ok(info_host[0])
2189 }
2190
2191 fn trsm_batched_lower_inplace(
2195 blas: &CudaBlas,
2196 stream: &Arc<CudaStream>,
2197 d: usize,
2198 n: usize,
2199 nrhs: usize,
2200 l_stack: &CudaSlice<f64>,
2201 rhs_stack: &mut CudaSlice<f64>,
2202 ) -> Result<(), ArrowSchurGpuFailure> {
2203 trsm_batched_inplace_inner(blas, stream, d, n, nrhs, l_stack, rhs_stack, false)
2204 }
2205
2206 fn trsm_batched_lower_inplace_transposed(
2208 blas: &CudaBlas,
2209 stream: &Arc<CudaStream>,
2210 d: usize,
2211 n: usize,
2212 nrhs: usize,
2213 l_stack: &CudaSlice<f64>,
2214 rhs_stack: &mut CudaSlice<f64>,
2215 ) -> Result<(), ArrowSchurGpuFailure> {
2216 trsm_batched_inplace_inner(blas, stream, d, n, nrhs, l_stack, rhs_stack, true)
2217 }
2218
2219 fn trsm_batched_inplace_inner(
2220 blas: &CudaBlas,
2221 stream: &Arc<CudaStream>,
2222 d: usize,
2223 n: usize,
2224 nrhs: usize,
2225 l_stack: &CudaSlice<f64>,
2226 rhs_stack: &mut CudaSlice<f64>,
2227 transposed: bool,
2228 ) -> Result<(), ArrowSchurGpuFailure> {
2229 let alpha = 1.0_f64;
2230 let d_i = to_i32(d).ok_or(ArrowSchurGpuFailure::Unavailable)?;
2231 let nrhs_i = to_i32(nrhs).ok_or(ArrowSchurGpuFailure::Unavailable)?;
2232 let batch_i = to_i32(n).ok_or(ArrowSchurGpuFailure::Unavailable)?;
2233 let l_bytes_per = (d * d * std::mem::size_of::<f64>()) as u64;
2234 let rhs_bytes_per = (d * nrhs * std::mem::size_of::<f64>()) as u64;
2235 let (l_base, _l_record) = l_stack.device_ptr(stream);
2236 let (rhs_base, _rhs_record) = rhs_stack.device_ptr_mut(stream);
2237 let mut l_ptrs = Vec::with_capacity(n);
2238 let mut rhs_ptrs = Vec::with_capacity(n);
2239 for i in 0..n {
2240 l_ptrs.push(l_base + (i as u64) * l_bytes_per);
2241 rhs_ptrs.push(rhs_base + (i as u64) * rhs_bytes_per);
2242 }
2243 let mut l_ptrs_dev = stream
2244 .clone_htod(&l_ptrs)
2245 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2246 let mut rhs_ptrs_dev = stream
2247 .clone_htod(&rhs_ptrs)
2248 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2249 let (l_ptrs_ptr, _l_ptrs_rec) = l_ptrs_dev.device_ptr_mut(stream);
2250 let (rhs_ptrs_ptr, _rhs_ptrs_rec) = rhs_ptrs_dev.device_ptr_mut(stream);
2251 let op = if transposed {
2252 cublasOperation_t::CUBLAS_OP_T
2253 } else {
2254 cublasOperation_t::CUBLAS_OP_N
2255 };
2256 let handle = *blas.handle();
2257 let status = unsafe {
2260 cudarc::cublas::sys::cublasDtrsmBatched(
2261 handle,
2262 cublasSideMode_t::CUBLAS_SIDE_LEFT,
2263 cublasFillMode_t::CUBLAS_FILL_MODE_LOWER,
2264 op,
2265 cublasDiagType_t::CUBLAS_DIAG_NON_UNIT,
2266 d_i,
2267 nrhs_i,
2268 &alpha,
2269 l_ptrs_ptr as *const *const f64,
2270 d_i,
2271 rhs_ptrs_ptr as *const *mut f64,
2272 d_i,
2273 batch_i,
2274 )
2275 };
2276 if status != cublasStatus_t::CUBLAS_STATUS_SUCCESS {
2277 return Err(ArrowSchurGpuFailure::Unavailable);
2278 }
2279 Ok(())
2280 }
2281
2282 fn trsm_single(
2285 blas: &CudaBlas,
2286 stream: &Arc<CudaStream>,
2287 n: usize,
2288 l: &CudaSlice<f64>,
2289 rhs: &mut CudaSlice<f64>,
2290 upper: bool,
2291 transposed: bool,
2292 ) -> Result<(), ArrowSchurGpuFailure> {
2293 let alpha = 1.0_f64;
2294 let n_i = to_i32(n).ok_or(ArrowSchurGpuFailure::Unavailable)?;
2295 let handle = *blas.handle();
2296 let (l_ptr, _l_rec) = l.device_ptr(stream);
2297 let (rhs_ptr, _rhs_rec) = rhs.device_ptr_mut(stream);
2298 let status = unsafe {
2300 cudarc::cublas::sys::cublasDtrsm_v2(
2301 handle,
2302 cublasSideMode_t::CUBLAS_SIDE_LEFT,
2303 if upper {
2304 cublasFillMode_t::CUBLAS_FILL_MODE_UPPER
2305 } else {
2306 cublasFillMode_t::CUBLAS_FILL_MODE_LOWER
2307 },
2308 if transposed {
2309 cublasOperation_t::CUBLAS_OP_T
2310 } else {
2311 cublasOperation_t::CUBLAS_OP_N
2312 },
2313 cublasDiagType_t::CUBLAS_DIAG_NON_UNIT,
2314 n_i,
2315 1,
2316 &alpha,
2317 l_ptr as *const f64,
2318 n_i,
2319 rhs_ptr as *mut f64,
2320 n_i,
2321 )
2322 };
2323 if status != cublasStatus_t::CUBLAS_STATUS_SUCCESS {
2324 return Err(ArrowSchurGpuFailure::Unavailable);
2325 }
2326 Ok(())
2327 }
2328
2329 fn accumulate_schur(
2333 blas: &CudaBlas,
2334 d: usize,
2335 k: usize,
2336 n: usize,
2337 y_stack: &CudaSlice<f64>,
2338 u_stack: &CudaSlice<f64>,
2339 schur: &mut CudaSlice<f64>,
2340 rhs: &mut CudaSlice<f64>,
2341 ) -> Result<(), ArrowSchurGpuFailure> {
2342 let y_block_elems = d * k;
2343 let u_block_elems = d;
2344 for i in 0..n {
2345 let y_slice = y_stack.slice(i * y_block_elems..(i + 1) * y_block_elems);
2346 let u_slice = u_stack.slice(i * u_block_elems..(i + 1) * u_block_elems);
2347 let gemm_cfg = GemmConfig::<f64> {
2349 transa: cublasOperation_t::CUBLAS_OP_T,
2350 transb: cublasOperation_t::CUBLAS_OP_N,
2351 m: to_i32(k).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2352 n: to_i32(k).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2353 k: to_i32(d).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2354 alpha: -1.0,
2355 lda: to_i32(d).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2356 ldb: to_i32(d).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2357 beta: 1.0,
2358 ldc: to_i32(k).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2359 };
2360 unsafe { blas.gemm(gemm_cfg, &y_slice, &y_slice, schur) }
2362 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2363 let gemv_cfg = GemvConfig::<f64> {
2365 trans: cublasOperation_t::CUBLAS_OP_T,
2366 m: to_i32(d).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2367 n: to_i32(k).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2368 alpha: 1.0,
2369 lda: to_i32(d).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2370 incx: 1,
2371 beta: 1.0,
2372 incy: 1,
2373 };
2374 unsafe { blas.gemv(gemv_cfg, &y_slice, &u_slice, rhs) }
2377 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2378 }
2379 Ok(())
2380 }
2381
2382 fn accumulate_schur_rhs_only(
2390 blas: &CudaBlas,
2391 d: usize,
2392 k: usize,
2393 n: usize,
2394 y_stack: &CudaSlice<f64>,
2395 u_stack: &CudaSlice<f64>,
2396 rhs: &mut CudaSlice<f64>,
2397 ) -> Result<(), ArrowSchurGpuFailure> {
2398 let y_block_elems = d * k;
2399 let u_block_elems = d;
2400 for i in 0..n {
2401 let y_slice = y_stack.slice(i * y_block_elems..(i + 1) * y_block_elems);
2402 let u_slice = u_stack.slice(i * u_block_elems..(i + 1) * u_block_elems);
2403 let gemv_cfg = GemvConfig::<f64> {
2404 trans: cublasOperation_t::CUBLAS_OP_T,
2405 m: to_i32(d).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2406 n: to_i32(k).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2407 alpha: 1.0,
2408 lda: to_i32(d).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2409 incx: 1,
2410 beta: 1.0,
2411 incy: 1,
2412 };
2413 unsafe { blas.gemv(gemv_cfg, &y_slice, &u_slice, rhs) }
2416 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2417 }
2418 Ok(())
2419 }
2420
2421 fn accumulate_back_sub_rhs(
2424 blas: &CudaBlas,
2425 d: usize,
2426 k: usize,
2427 n: usize,
2428 y_stack: &CudaSlice<f64>,
2429 delta_beta: &CudaSlice<f64>,
2430 u_stack: &mut CudaSlice<f64>,
2431 ) -> Result<(), ArrowSchurGpuFailure> {
2432 let y_block_elems = d * k;
2433 let u_block_elems = d;
2434 for i in 0..n {
2435 let y_slice = y_stack.slice(i * y_block_elems..(i + 1) * y_block_elems);
2436 let mut u_slice = u_stack.slice_mut(i * u_block_elems..(i + 1) * u_block_elems);
2437 let gemv_cfg = GemvConfig::<f64> {
2438 trans: cublasOperation_t::CUBLAS_OP_N,
2439 m: to_i32(d).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2440 n: to_i32(k).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2441 alpha: 1.0,
2442 lda: to_i32(d).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2443 incx: 1,
2444 beta: 1.0,
2445 incy: 1,
2446 };
2447 unsafe { blas.gemv(gemv_cfg, &y_slice, delta_beta, &mut u_slice) }
2450 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2451 }
2452 Ok(())
2453 }
2454
2455 use std::collections::HashMap;
2471 use std::sync::Mutex;
2472
2473 struct FusedModuleCache {
2478 modules: Mutex<
2479 HashMap<crate::gpu_kernels::arrow_schur_nvrtc::FusedModuleCacheKey, Arc<CudaModule>>,
2480 >,
2481 }
2482
2483 fn fused_module_cache() -> &'static FusedModuleCache {
2484 static CACHE: OnceLock<FusedModuleCache> = OnceLock::new();
2485 CACHE.get_or_init(|| FusedModuleCache {
2486 modules: Mutex::new(HashMap::new()),
2487 })
2488 }
2489
2490 fn fused_module_for(
2491 ctx: &Arc<CudaContext>,
2492 key: crate::gpu_kernels::arrow_schur_nvrtc::FusedModuleCacheKey,
2493 ) -> Result<Arc<CudaModule>, ArrowSchurGpuFailure> {
2494 let cache = fused_module_cache();
2495 if let Ok(guard) = cache.modules.lock() {
2496 if let Some(existing) = guard.get(&key) {
2497 return Ok(existing.clone());
2498 }
2499 }
2500 let src = crate::gpu_kernels::arrow_schur_nvrtc::forward_kernel_source(
2501 key.p_max as usize,
2502 key.r_template as usize,
2503 );
2504 let ptx = gam_gpu::device_cache::compile_ptx_arch(&src).map_err(|err| {
2505 ArrowSchurGpuFailure::SchurFactorFailed {
2506 reason: format!(
2507 "arrow-schur fused NVRTC compile (p_max={}, r={}): {err}",
2508 key.p_max, key.r_template
2509 ),
2510 }
2511 })?;
2512 let module = ctx
2513 .load_module(ptx)
2514 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2515 if let Ok(mut guard) = cache.modules.lock() {
2516 guard.entry(key).or_insert_with(|| module.clone());
2517 }
2518 Ok(module)
2519 }
2520
2521 const PCG_VECTOR_KERNEL_SOURCE: &str = r#"
2522extern "C" __global__ void arrow_pcg_jacobi_mul(
2523 const double* __restrict__ inv_diag,
2524 const double* __restrict__ r,
2525 double* __restrict__ z,
2526 int n
2527) {
2528 int idx = blockIdx.x * blockDim.x + threadIdx.x;
2529 if (idx < n) {
2530 z[idx] = inv_diag[idx] * r[idx];
2531 }
2532}
2533
2534extern "C" __global__ void arrow_pcg_update_p(
2535 const double* __restrict__ z,
2536 double beta,
2537 double* __restrict__ p,
2538 int n
2539) {
2540 int idx = blockIdx.x * blockDim.x + threadIdx.x;
2541 if (idx < n) {
2542 p[idx] = z[idx] + beta * p[idx];
2543 }
2544}
2545
2546extern "C" __global__ void arrow_sae_init(
2547 double* __restrict__ out,
2548 const double* __restrict__ x,
2549 double ridge,
2550 int n
2551) {
2552 int idx = blockIdx.x * blockDim.x + threadIdx.x;
2553 if (idx < n) {
2554 out[idx] = ridge * x[idx];
2555 }
2556}
2557
2558extern "C" __global__ void arrow_sae_smooth_matvec(
2559 const double* __restrict__ x,
2560 double* __restrict__ out,
2561 const int* __restrict__ block_offsets,
2562 const int* __restrict__ block_m,
2563 const int* __restrict__ factor_ptr,
2564 const double* __restrict__ factors,
2565 int p,
2566 int n_blocks
2567) {
2568 int block_id = blockIdx.y;
2569 int linear = blockIdx.x * blockDim.x + threadIdx.x;
2570 if (block_id >= n_blocks) {
2571 return;
2572 }
2573 int m = block_m[block_id];
2574 int total = m * p;
2575 if (linear >= total) {
2576 return;
2577 }
2578 int li = linear / p;
2579 int oc = linear - li * p;
2580 int off = block_offsets[block_id];
2581 int fbase = factor_ptr[block_id];
2582 double acc = 0.0;
2583 for (int lj = 0; lj < m; ++lj) {
2584 double a = factors[fbase + li * m + lj];
2585 acc += a * x[off + lj * p + oc];
2586 }
2587 out[off + li * p + oc] += acc;
2588}
2589
2590extern "C" __global__ void arrow_sae_sparse_g_matvec(
2591 const double* __restrict__ x,
2592 double* __restrict__ out,
2593 const int* __restrict__ row_off,
2594 const int* __restrict__ col_off,
2595 const int* __restrict__ rows,
2596 const int* __restrict__ cols,
2597 const int* __restrict__ data_ptr,
2598 const double* __restrict__ data,
2599 int p,
2600 int n_blocks
2601) {
2602 int block_id = blockIdx.y;
2603 int linear = blockIdx.x * blockDim.x + threadIdx.x;
2604 if (block_id >= n_blocks) {
2605 return;
2606 }
2607 int m_i = rows[block_id];
2608 int m_j = cols[block_id];
2609 int total = m_i * p;
2610 if (linear >= total) {
2611 return;
2612 }
2613 int li = linear / p;
2614 int oc = linear - li * p;
2615 int rbase = row_off[block_id];
2616 int cbase = col_off[block_id];
2617 int dbase = data_ptr[block_id];
2618 double acc = 0.0;
2619 for (int lj = 0; lj < m_j; ++lj) {
2620 acc += data[dbase + li * m_j + lj] * x[(cbase + lj) * p + oc];
2621 }
2622 // #1017 — a row atom co-occurs with multiple column atoms, so several
2623 // concurrent (atom_i, atom_j) blocks (blockIdx.y) write the SAME output
2624 // element `out[(rbase+li)*p+oc]`. A plain `+=` races and loses updates
2625 // (silently-wrong Schur matvec); accumulate atomically. `double` atomicAdd
2626 // needs sm_60+, guaranteed by the NVRTC arch pin (#1551).
2627 atomicAdd(&out[(rbase + li) * p + oc], acc);
2628}
2629
2630extern "C" __global__ void arrow_sae_gather_u(
2631 const double* __restrict__ x,
2632 const int* __restrict__ row_ptr,
2633 const int* __restrict__ beta_base,
2634 const double* __restrict__ phi,
2635 double* __restrict__ u,
2636 int p,
2637 int n_rows
2638) {
2639 int row = blockIdx.y;
2640 int oc = blockIdx.x * blockDim.x + threadIdx.x;
2641 if (row >= n_rows || oc >= p) {
2642 return;
2643 }
2644 double acc = 0.0;
2645 int start = row_ptr[row];
2646 int end = row_ptr[row + 1];
2647 for (int e = start; e < end; ++e) {
2648 acc += phi[e] * x[beta_base[e] + oc];
2649 }
2650 u[row * p + oc] = acc;
2651}
2652
2653extern "C" __global__ void arrow_sae_apply_l(
2654 const double* __restrict__ u,
2655 const int* __restrict__ jac_ptr,
2656 const double* __restrict__ jac,
2657 double* __restrict__ w,
2658 int p,
2659 int max_q,
2660 int n_rows
2661) {
2662 int row = blockIdx.y;
2663 int c = blockIdx.x * blockDim.x + threadIdx.x;
2664 if (row >= n_rows) {
2665 return;
2666 }
2667 int jstart = jac_ptr[row];
2668 int q = (jac_ptr[row + 1] - jstart) / p;
2669 if (c >= q) {
2670 return;
2671 }
2672 double acc = 0.0;
2673 for (int oc = 0; oc < p; ++oc) {
2674 acc += jac[jstart + c * p + oc] * u[row * p + oc];
2675 }
2676 w[row * max_q + c] = acc;
2677}
2678
2679extern "C" __global__ void arrow_sae_apply_ainv(
2680 const double* __restrict__ ainv,
2681 const double* __restrict__ w,
2682 double* __restrict__ v,
2683 int max_q,
2684 int n_rows
2685) {
2686 int row = blockIdx.y;
2687 int c = blockIdx.x * blockDim.x + threadIdx.x;
2688 if (row >= n_rows || c >= max_q) {
2689 return;
2690 }
2691 double acc = 0.0;
2692 int base = row * max_q * max_q;
2693 for (int j = 0; j < max_q; ++j) {
2694 acc += ainv[base + c * max_q + j] * w[row * max_q + j];
2695 }
2696 v[row * max_q + c] = acc;
2697}
2698
2699extern "C" __global__ void arrow_sae_scatter_sub(
2700 const double* __restrict__ v,
2701 const int* __restrict__ jac_ptr,
2702 const double* __restrict__ jac,
2703 const int* __restrict__ row_ptr,
2704 const int* __restrict__ beta_base,
2705 const double* __restrict__ phi,
2706 double* __restrict__ out,
2707 int p,
2708 int max_q,
2709 int n_rows
2710) {
2711 int row = blockIdx.y;
2712 int oc = blockIdx.x * blockDim.x + threadIdx.x;
2713 if (row >= n_rows || oc >= p) {
2714 return;
2715 }
2716 int jstart = jac_ptr[row];
2717 int q = (jac_ptr[row + 1] - jstart) / p;
2718 double lt_v = 0.0;
2719 for (int c = 0; c < q; ++c) {
2720 lt_v += jac[jstart + c * p + oc] * v[row * max_q + c];
2721 }
2722 int start = row_ptr[row];
2723 int end = row_ptr[row + 1];
2724 for (int e = start; e < end; ++e) {
2725 atomicAdd(&out[beta_base[e] + oc], -phi[e] * lt_v);
2726 }
2727}
2728
2729extern "C" __global__ void arrow_sae_diag_sub(
2730 double* __restrict__ diag,
2731 const double* __restrict__ ainv,
2732 const int* __restrict__ jac_ptr,
2733 const double* __restrict__ jac,
2734 const int* __restrict__ row_ptr,
2735 const int* __restrict__ beta_base,
2736 const double* __restrict__ phi,
2737 int p,
2738 int max_q,
2739 int n_rows
2740) {
2741 int row = blockIdx.y;
2742 int oc = blockIdx.x * blockDim.x + threadIdx.x;
2743 if (row >= n_rows || oc >= p) {
2744 return;
2745 }
2746 int jstart = jac_ptr[row];
2747 int q = (jac_ptr[row + 1] - jstart) / p;
2748 int abase = row * max_q * max_q;
2749 double quad = 0.0;
2750 for (int c = 0; c < q; ++c) {
2751 double lc = jac[jstart + c * p + oc];
2752 for (int d = 0; d < q; ++d) {
2753 quad += lc * ainv[abase + c * max_q + d] * jac[jstart + d * p + oc];
2754 }
2755 }
2756 int start = row_ptr[row];
2757 int end = row_ptr[row + 1];
2758 for (int e = start; e < end; ++e) {
2759 double pe = phi[e];
2760 atomicAdd(&diag[beta_base[e] + oc], -(pe * pe) * quad);
2761 }
2762}
2763
2764/* ── #1017/#1026 frames-engaged device kernels ─────────────────────────────
2765 * The factored β border is C-space (width Σ M_k·r_k). The penalty side is the
2766 * smooth `λ S_k ⊗ I_{r_k}` (per-block right-width r_k) plus the data-fit
2767 * `G_{ij} ⊗ W_{ij}` (W = U_iᵀU_j, dense r_i×r_j). The reduced-Schur term uses
2768 * the per-row DENSE cross-block H_tβ^(i) (q_i × border_dim, row-major). */
2769
2770extern "C" __global__ void arrow_sae_frame_smooth_matvec(
2771 const double* __restrict__ x,
2772 double* __restrict__ out,
2773 const int* __restrict__ block_offsets,
2774 const int* __restrict__ block_m,
2775 const int* __restrict__ block_r,
2776 const int* __restrict__ factor_ptr,
2777 const double* __restrict__ factors,
2778 int n_blocks
2779) {
2780 int block_id = blockIdx.y;
2781 int linear = blockIdx.x * blockDim.x + threadIdx.x;
2782 if (block_id >= n_blocks) {
2783 return;
2784 }
2785 int m = block_m[block_id];
2786 int r = block_r[block_id];
2787 int total = m * r;
2788 if (linear >= total) {
2789 return;
2790 }
2791 int li = linear / r;
2792 int ib = linear - li * r;
2793 int off = block_offsets[block_id];
2794 int fbase = factor_ptr[block_id];
2795 double acc = 0.0;
2796 for (int lj = 0; lj < m; ++lj) {
2797 double a = factors[fbase + li * m + lj];
2798 acc += a * x[off + lj * r + ib];
2799 }
2800 out[off + li * r + ib] += acc;
2801}
2802
2803extern "C" __global__ void arrow_sae_frame_g_matvec(
2804 const double* __restrict__ x,
2805 double* __restrict__ out,
2806 const int* __restrict__ off_i,
2807 const int* __restrict__ off_j,
2808 const int* __restrict__ r_i,
2809 const int* __restrict__ r_j,
2810 const int* __restrict__ m_i,
2811 const int* __restrict__ m_j,
2812 const int* __restrict__ g_ptr,
2813 const double* __restrict__ g_data,
2814 const int* __restrict__ w_ptr,
2815 const double* __restrict__ w_data,
2816 int n_blocks
2817) {
2818 int block_id = blockIdx.y;
2819 int linear = blockIdx.x * blockDim.x + threadIdx.x;
2820 if (block_id >= n_blocks) {
2821 return;
2822 }
2823 int ri = r_i[block_id];
2824 int rj = r_j[block_id];
2825 int mi = m_i[block_id];
2826 int mj = m_j[block_id];
2827 int total = mi * ri;
2828 if (linear >= total) {
2829 return;
2830 }
2831 int li = linear / ri; // basis row in atom i
2832 int a = linear - li * ri; // frame coord in atom i
2833 int oi = off_i[block_id];
2834 int oj = off_j[block_id];
2835 int gbase = g_ptr[block_id];
2836 int wbase = w_ptr[block_id];
2837 double acc = 0.0;
2838 for (int lj = 0; lj < mj; ++lj) {
2839 double g = g_data[gbase + li * mj + lj];
2840 if (g == 0.0) { continue; }
2841 int xj_base = oj + lj * rj;
2842 double inner = 0.0;
2843 for (int b = 0; b < rj; ++b) {
2844 inner += w_data[wbase + a * rj + b] * x[xj_base + b];
2845 }
2846 acc += g * inner;
2847 }
2848 // #1017 — same race as `arrow_sae_sparse_g_matvec`: atom i is the row atom of
2849 // multiple co-occurring (i,j) frame blocks running concurrently on
2850 // blockIdx.y, all targeting `out[oi+li*ri+a]`. Accumulate atomically so the
2851 // framed G⊗W matvec is correct (the CPU oracle sums these sequentially).
2852 atomicAdd(&out[oi + li * ri + a], acc);
2853}
2854
2855/* Per-row reduced-Schur subtraction with a DENSE cross-block H_tβ^(i).
2856 * h_i = H_tβ^(i) · x (length q_i)
2857 * s_i = (H_tt^(i)+ρ_t I)⁻¹ h_i (apply cached ainv, length q_i)
2858 * out -= (H_tβ^(i))ᵀ · s_i (scatter into border_dim)
2859 * `htb` is row-major (q_i × k) flattened, `htb_ptr` gives each row's base and
2860 * (htb_ptr[row+1]-htb_ptr[row])/k == q_i. `q_of` carries q_i directly. */
2861extern "C" __global__ void arrow_sae_frame_apply_h(
2862 const double* __restrict__ x,
2863 const int* __restrict__ htb_ptr,
2864 const double* __restrict__ htb,
2865 const int* __restrict__ q_of,
2866 double* __restrict__ hvec,
2867 int k,
2868 int max_q,
2869 int n_rows
2870) {
2871 int row = blockIdx.y;
2872 int c = blockIdx.x * blockDim.x + threadIdx.x;
2873 if (row >= n_rows) { return; }
2874 int q = q_of[row];
2875 if (c >= q) { return; }
2876 int base = htb_ptr[row] + c * k;
2877 double acc = 0.0;
2878 for (int a = 0; a < k; ++a) {
2879 acc += htb[base + a] * x[a];
2880 }
2881 hvec[row * max_q + c] = acc;
2882}
2883
2884extern "C" __global__ void arrow_sae_frame_apply_ainv(
2885 const double* __restrict__ ainv,
2886 const double* __restrict__ hvec,
2887 const int* __restrict__ q_of,
2888 double* __restrict__ svec,
2889 int max_q,
2890 int n_rows
2891) {
2892 int row = blockIdx.y;
2893 int c = blockIdx.x * blockDim.x + threadIdx.x;
2894 if (row >= n_rows || c >= max_q) { return; }
2895 int q = q_of[row];
2896 double acc = 0.0;
2897 int abase = row * max_q * max_q;
2898 for (int j = 0; j < q; ++j) {
2899 acc += ainv[abase + c * max_q + j] * hvec[row * max_q + j];
2900 }
2901 svec[row * max_q + c] = acc;
2902}
2903
2904extern "C" __global__ void arrow_sae_frame_scatter_h(
2905 const double* __restrict__ svec,
2906 const int* __restrict__ htb_ptr,
2907 const double* __restrict__ htb,
2908 const int* __restrict__ q_of,
2909 double* __restrict__ out,
2910 int k,
2911 int max_q,
2912 int n_rows
2913) {
2914 int row = blockIdx.y;
2915 int a = blockIdx.x * blockDim.x + threadIdx.x;
2916 if (row >= n_rows || a >= k) { return; }
2917 int q = q_of[row];
2918 int hbase = htb_ptr[row];
2919 double acc = 0.0;
2920 for (int c = 0; c < q; ++c) {
2921 acc += htb[hbase + c * k + a] * svec[row * max_q + c];
2922 }
2923 atomicAdd(&out[a], -acc);
2924}
2925
2926/* #1017 evidence-lane DETERMINISTIC reduced-Schur scatter:
2927 out[a] = -Σ_i Σ_c H_tβ[i][c,a]·svec[i,c]. One thread owns output coord `a` and
2928 sums the rows in fixed index order 0..n_rows — NO atomics, so the result is
2929 run-to-run bit-stable (the SLQ log|S| determinism contract). Same arithmetic as
2930 arrow_sae_frame_scatter_h; only the reduction order is pinned (there the sum
2931 over rows is an atomicAdd race). The shared atomic kernel is left untouched —
2932 the step-PCG relies on it. `out` is fully assigned (no init needed). */
2933extern "C" __global__ void arrow_sae_frame_scatter_h_det(
2934 const double* __restrict__ svec,
2935 const int* __restrict__ htb_ptr,
2936 const double* __restrict__ htb,
2937 const int* __restrict__ q_of,
2938 double* __restrict__ out,
2939 int k,
2940 int max_q,
2941 int n_rows
2942) {
2943 int a = blockIdx.x * blockDim.x + threadIdx.x;
2944 if (a >= k) { return; }
2945 double acc = 0.0;
2946 for (int row = 0; row < n_rows; ++row) {
2947 int q = q_of[row];
2948 int hbase = htb_ptr[row];
2949 int sbase = row * max_q;
2950 for (int c = 0; c < q; ++c) {
2951 acc += htb[hbase + c * k + a] * svec[sbase + c];
2952 }
2953 }
2954 out[a] = -acc;
2955}
2956
2957/* #1017 evidence-lane 2-STAGE deterministic scatter, STAGE 1 (partials):
2958 partial[chunk][a] = Σ_{row∈chunk} Σ_c H_tβ[row][c,a]·svec[row,c], for the
2959 contiguous row range [chunk·rows_per_chunk, …). grid = (⌈k/256⌉, n_chunks);
2960 thread owns (a, chunk). Replaces the single-strip `arrow_sae_frame_scatter_h_det`
2961 (⌈k/256⌉ CTAs — only 4 at k=911, one thread serial over ALL n_rows, ~94% of a
2962 72-SM A10 idle) with ⌈k/256⌉·n_chunks CTAs. Rows are summed in fixed order
2963 within the chunk and the chunks are reduced in order by stage 2, so the result
2964 is a FIXED reassociation of the same ordered row sum — run-to-run bit-stable
2965 (the SLQ log|S| determinism contract) and within the ≤1e-9 CPU-oracle gate. */
2966extern "C" __global__ void arrow_sae_frame_scatter_h_det_partial(
2967 const double* __restrict__ svec,
2968 const int* __restrict__ htb_ptr,
2969 const double* __restrict__ htb,
2970 const int* __restrict__ q_of,
2971 double* __restrict__ partial,
2972 int k,
2973 int max_q,
2974 int n_rows,
2975 int rows_per_chunk
2976) {
2977 int a = blockIdx.x * blockDim.x + threadIdx.x;
2978 if (a >= k) { return; }
2979 int chunk = blockIdx.y;
2980 int row0 = chunk * rows_per_chunk;
2981 int row1 = row0 + rows_per_chunk;
2982 if (row1 > n_rows) { row1 = n_rows; }
2983 double acc = 0.0;
2984 for (int row = row0; row < row1; ++row) {
2985 int q = q_of[row];
2986 int hbase = htb_ptr[row];
2987 int sbase = row * max_q;
2988 for (int c = 0; c < q; ++c) {
2989 acc += htb[hbase + c * k + a] * svec[sbase + c];
2990 }
2991 }
2992 partial[(long long)chunk * k + a] = acc;
2993}
2994
2995/* #1017 STAGE 2 (reduce): out[a] = -Σ_chunk partial[chunk][a], chunks summed in
2996 fixed order 0..n_chunks. One thread per output coord `a`; ⌈k/256⌉ CTAs. */
2997extern "C" __global__ void arrow_sae_frame_scatter_h_det_reduce(
2998 const double* __restrict__ partial,
2999 double* __restrict__ out,
3000 int k,
3001 int n_chunks
3002) {
3003 int a = blockIdx.x * blockDim.x + threadIdx.x;
3004 if (a >= k) { return; }
3005 double acc = 0.0;
3006 for (int chunk = 0; chunk < n_chunks; ++chunk) {
3007 acc += partial[(long long)chunk * k + a];
3008 }
3009 out[a] = -acc;
3010}
3011
3012/* #1017 evidence-lane WARP-COOPERATIVE apply_h: hvec[i][c] = Σ_a H_tβ[i][c,a]·x[a].
3013 One WARP owns (row, c): lane `l` strides `a = l, l+32, …` over the contiguous
3014 `H_tβ[i][c,·]` slab (fully coalesced across the warp) and a fixed-order
3015 __shfl_down tree reduces to lane 0. Replaces the shared `arrow_sae_frame_apply_h`
3016 (256-thread block, only `q_i` threads active, stride-`k` uncoalesced reads) on
3017 the evidence path ONLY — the shared kernel is untouched (step-PCG relies on it).
3018 Reduction order is fixed (lane stride + tree), so the result is run-to-run
3019 bit-stable; it differs from the scalar kernel only by ULP reassociation, within
3020 the ≤1e-9 parity gate. Launch: block = max_q·32 (≤1024), grid.x = n_rows;
3021 warp `w = threadIdx.x/32` handles `c = w` for `w < q_i`. */
3022extern "C" __global__ void arrow_sae_frame_apply_h_warp(
3023 const double* __restrict__ x,
3024 const int* __restrict__ htb_ptr,
3025 const double* __restrict__ htb,
3026 const int* __restrict__ q_of,
3027 double* __restrict__ hvec,
3028 int k,
3029 int max_q,
3030 int n_rows
3031) {
3032 int row = blockIdx.x;
3033 if (row >= n_rows) { return; }
3034 int warp = threadIdx.x >> 5;
3035 int lane = threadIdx.x & 31;
3036 int q = q_of[row];
3037 if (warp >= q) { return; }
3038 int base = htb_ptr[row] + warp * k;
3039 double acc = 0.0;
3040 for (int a = lane; a < k; a += 32) {
3041 acc += htb[base + a] * x[a];
3042 }
3043 #pragma unroll
3044 for (int off = 16; off > 0; off >>= 1) {
3045 acc += __shfl_down_sync(0xffffffffu, acc, off);
3046 }
3047 if (lane == 0) {
3048 hvec[row * max_q + warp] = acc;
3049 }
3050}
3051
3052/* Frame Jacobi diagonal subtraction: diag[a] -= Σ_c Σ_d H_tβ[c,a]·ainv[c,d]·H_tβ[d,a]. */
3053extern "C" __global__ void arrow_sae_frame_diag_sub(
3054 double* __restrict__ diag,
3055 const double* __restrict__ ainv,
3056 const int* __restrict__ htb_ptr,
3057 const double* __restrict__ htb,
3058 const int* __restrict__ q_of,
3059 int k,
3060 int max_q,
3061 int n_rows
3062) {
3063 int row = blockIdx.y;
3064 int a = blockIdx.x * blockDim.x + threadIdx.x;
3065 if (row >= n_rows || a >= k) { return; }
3066 int q = q_of[row];
3067 int hbase = htb_ptr[row];
3068 int abase = row * max_q * max_q;
3069 double quad = 0.0;
3070 for (int c = 0; c < q; ++c) {
3071 double hc = htb[hbase + c * k + a];
3072 for (int d = 0; d < q; ++d) {
3073 quad += hc * ainv[abase + c * max_q + d] * htb[hbase + d * k + a];
3074 }
3075 }
3076 atomicAdd(&diag[a], -quad);
3077}
3078"#;
3079
3080 fn pcg_vector_module(
3081 ctx: &Arc<CudaContext>,
3082 ) -> Result<&'static Arc<CudaModule>, ArrowSchurGpuFailure> {
3083 static CACHE: gam_gpu::device_cache::PtxModuleCache =
3084 gam_gpu::device_cache::PtxModuleCache::new();
3085 CACHE
3086 .get_or_compile(ctx, "arrow_pcg_vector", PCG_VECTOR_KERNEL_SOURCE)
3087 .map_err(|err| {
3088 log::warn!("[#1551] pcg_vector_module get_or_compile failed: {err}");
3094 ArrowSchurGpuFailure::Unavailable
3095 })
3096 }
3097
3098 fn pcg_launch_config(n: usize) -> Result<LaunchConfig, ArrowSchurGpuFailure> {
3099 let threads = 256u32;
3100 let blocks = ((n as u32).saturating_add(threads - 1) / threads).max(1);
3101 Ok(LaunchConfig {
3102 grid_dim: (blocks, 1, 1),
3103 block_dim: (threads, 1, 1),
3104 shared_mem_bytes: 0,
3105 })
3106 }
3107
3108 fn launch_jacobi_mul(
3109 stream: &Arc<CudaStream>,
3110 module: &Arc<CudaModule>,
3111 inv_diag: &CudaSlice<f64>,
3112 r: &CudaSlice<f64>,
3113 z: &mut CudaSlice<f64>,
3114 n: usize,
3115 ) -> Result<(), ArrowSchurGpuFailure> {
3116 let kernel = module
3117 .load_function("arrow_pcg_jacobi_mul")
3118 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
3119 let n_i32 = to_i32(n).ok_or(ArrowSchurGpuFailure::Unavailable)?;
3120 let mut builder = stream.launch_builder(&kernel);
3121 builder.arg(inv_diag).arg(r).arg(z).arg(&n_i32);
3122 unsafe { builder.launch(pcg_launch_config(n)?) }
3125 .map(drop)
3126 .map_err(|_| ArrowSchurGpuFailure::Unavailable)
3127 }
3128
3129 fn launch_update_p(
3130 stream: &Arc<CudaStream>,
3131 module: &Arc<CudaModule>,
3132 z: &CudaSlice<f64>,
3133 beta: f64,
3134 p: &mut CudaSlice<f64>,
3135 n: usize,
3136 ) -> Result<(), ArrowSchurGpuFailure> {
3137 let kernel = module
3138 .load_function("arrow_pcg_update_p")
3139 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
3140 let n_i32 = to_i32(n).ok_or(ArrowSchurGpuFailure::Unavailable)?;
3141 let mut builder = stream.launch_builder(&kernel);
3142 builder.arg(z).arg(&beta).arg(p).arg(&n_i32);
3143 unsafe { builder.launch(pcg_launch_config(n)?) }
3146 .map(drop)
3147 .map_err(|_| ArrowSchurGpuFailure::Unavailable)
3148 }
3149
3150 struct DeviceSaePcgBuffers {
3151 row_ptr: CudaSlice<i32>,
3152 beta_base: CudaSlice<i32>,
3153 phi: CudaSlice<f64>,
3154 jac_ptr: CudaSlice<i32>,
3155 jac: CudaSlice<f64>,
3156 smooth_offsets: CudaSlice<i32>,
3157 smooth_m: CudaSlice<i32>,
3158 smooth_ptr: CudaSlice<i32>,
3159 smooth_data: CudaSlice<f64>,
3160 g_row_off: CudaSlice<i32>,
3161 g_col_off: CudaSlice<i32>,
3162 g_rows: CudaSlice<i32>,
3163 g_cols: CudaSlice<i32>,
3164 g_ptr: CudaSlice<i32>,
3165 g_data: CudaSlice<f64>,
3166 ainv: CudaSlice<f64>,
3167 u: CudaSlice<f64>,
3168 w: CudaSlice<f64>,
3169 v: CudaSlice<f64>,
3170 n_rows: usize,
3171 p: usize,
3172 k: usize,
3173 max_q: usize,
3174 smooth_blocks: usize,
3175 g_blocks: usize,
3176 }
3177
3178 fn checked_i32(value: usize) -> Result<i32, ArrowSchurGpuFailure> {
3179 to_i32(value).ok_or(ArrowSchurGpuFailure::Unavailable)
3180 }
3181
3182 fn sae_penalty_diag_host(
3183 data: &DeviceSaePcgData,
3184 ridge_beta: f64,
3185 ) -> Result<Vec<f64>, ArrowSchurGpuFailure> {
3186 let mut diag = vec![ridge_beta; data.beta_dim];
3187 for block in &data.smooth_blocks {
3188 let (rows, cols) = block.factor_a.dim();
3189 if rows != cols {
3190 return Err(ArrowSchurGpuFailure::Unavailable);
3191 }
3192 for row in 0..rows {
3193 let coeff = block.factor_a[[row, row]];
3194 let base = block
3195 .global_offset
3196 .checked_add(
3197 row.checked_mul(data.p)
3198 .ok_or(ArrowSchurGpuFailure::Unavailable)?,
3199 )
3200 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
3201 let end = base
3202 .checked_add(data.p)
3203 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
3204 if end > diag.len() {
3205 return Err(ArrowSchurGpuFailure::Unavailable);
3206 }
3207 for channel in 0..data.p {
3208 diag[base + channel] += coeff;
3209 }
3210 }
3211 }
3212 for block in &data.sparse_g_blocks {
3213 if block.row_off != block.col_off {
3214 continue;
3215 }
3216 let (rows, cols) = block.data.dim();
3217 for row in 0..rows.min(cols) {
3218 let coeff = block.data[[row, row]];
3219 let beta_row = block
3220 .row_off
3221 .checked_add(row)
3222 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
3223 let base = beta_row
3224 .checked_mul(data.p)
3225 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
3226 let end = base
3227 .checked_add(data.p)
3228 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
3229 if end > diag.len() {
3230 return Err(ArrowSchurGpuFailure::Unavailable);
3231 }
3232 for channel in 0..data.p {
3233 diag[base + channel] += coeff;
3234 }
3235 }
3236 }
3237 Ok(diag)
3238 }
3239
3240 fn flatten_device_sae_data(
3241 sys: &ArrowSchurSystem,
3242 data: &DeviceSaePcgData,
3243 ridge_t: f64,
3244 stream: &Arc<CudaStream>,
3245 ) -> Result<DeviceSaePcgBuffers, ArrowSchurGpuFailure> {
3246 let n_rows = sys.rows.len();
3247 let p = data.p;
3248 let k = data.beta_dim;
3249 if data.a_phi.len() != n_rows || data.local_jac.len() != n_rows {
3250 return Err(ArrowSchurGpuFailure::Unavailable);
3251 }
3252
3253 let mut row_ptr_host = Vec::with_capacity(n_rows + 1);
3254 let mut beta_base_host = Vec::<i32>::new();
3255 let mut phi_host = Vec::<f64>::new();
3256 row_ptr_host.push(0_i32);
3257 for row in data.a_phi.iter() {
3258 for &(base, phi) in row {
3259 beta_base_host.push(checked_i32(base)?);
3260 phi_host.push(phi);
3261 }
3262 row_ptr_host.push(checked_i32(beta_base_host.len())?);
3263 }
3264
3265 let mut jac_ptr_host = Vec::with_capacity(n_rows + 1);
3266 let mut jac_host = Vec::<f64>::new();
3267 let mut max_q = 0usize;
3268 jac_ptr_host.push(0_i32);
3269 for row_jac in data.local_jac.iter() {
3270 if row_jac.len() % p != 0 {
3271 return Err(ArrowSchurGpuFailure::Unavailable);
3272 }
3273 max_q = max_q.max(row_jac.len() / p);
3274 jac_host.extend_from_slice(row_jac);
3275 jac_ptr_host.push(checked_i32(jac_host.len())?);
3276 }
3277 if max_q == 0 {
3278 return Err(ArrowSchurGpuFailure::Unavailable);
3279 }
3280
3281 let mut smooth_offsets_host = Vec::with_capacity(data.smooth_blocks.len());
3282 let mut smooth_m_host = Vec::with_capacity(data.smooth_blocks.len());
3283 let mut smooth_ptr_host = Vec::with_capacity(data.smooth_blocks.len() + 1);
3284 let mut smooth_data_host = Vec::<f64>::new();
3285 smooth_ptr_host.push(0_i32);
3286 for block in &data.smooth_blocks {
3287 let (rows, cols) = block.factor_a.dim();
3288 if rows != cols {
3289 return Err(ArrowSchurGpuFailure::Unavailable);
3290 }
3291 smooth_offsets_host.push(checked_i32(block.global_offset)?);
3292 smooth_m_host.push(checked_i32(rows)?);
3293 for r in 0..rows {
3294 for c in 0..cols {
3295 smooth_data_host.push(block.factor_a[[r, c]]);
3296 }
3297 }
3298 smooth_ptr_host.push(checked_i32(smooth_data_host.len())?);
3299 }
3300
3301 let mut g_row_off_host = Vec::with_capacity(data.sparse_g_blocks.len());
3302 let mut g_col_off_host = Vec::with_capacity(data.sparse_g_blocks.len());
3303 let mut g_rows_host = Vec::with_capacity(data.sparse_g_blocks.len());
3304 let mut g_cols_host = Vec::with_capacity(data.sparse_g_blocks.len());
3305 let mut g_ptr_host = Vec::with_capacity(data.sparse_g_blocks.len() + 1);
3306 let mut g_data_host = Vec::<f64>::new();
3307 g_ptr_host.push(0_i32);
3308 for block in &data.sparse_g_blocks {
3309 let (rows, cols) = block.data.dim();
3310 g_row_off_host.push(checked_i32(block.row_off)?);
3311 g_col_off_host.push(checked_i32(block.col_off)?);
3312 g_rows_host.push(checked_i32(rows)?);
3313 g_cols_host.push(checked_i32(cols)?);
3314 for r in 0..rows {
3315 for c in 0..cols {
3316 g_data_host.push(block.data[[r, c]]);
3317 }
3318 }
3319 g_ptr_host.push(checked_i32(g_data_host.len())?);
3320 }
3321
3322 let mut ainv_host = vec![0.0_f64; n_rows * max_q * max_q];
3323 for (row_idx, row) in sys.rows.iter().enumerate() {
3324 let q = data.local_jac[row_idx].len() / p;
3325 if row.htt.dim() != (q, q) {
3326 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
3327 reason: format!(
3328 "SAE device PCG row {row_idx}: H_tt shape {:?} != ({q}, {q})",
3329 row.htt.dim()
3330 ),
3331 });
3332 }
3333 let mut block = row.htt.clone();
3334 for d in 0..q {
3335 block[[d, d]] += ridge_t;
3336 }
3337 let factor = gam_linalg::triangular::cholesky_factor_in_place(
3338 block.view(),
3339 gam_linalg::triangular::CholeskyGuard::NonnegativePivot,
3340 )
3341 .ok_or_else(|| {
3342 ArrowSchurGpuFailure::RidgeBumpRequired {
3345 row: row_idx,
3346 bump: super::ridge_bump_to_make_pd(row.htt.view(), ridge_t),
3347 }
3348 })?;
3349 for col in 0..q {
3350 let mut e = Array1::<f64>::zeros(q);
3351 e[col] = 1.0;
3352 let solved = gam_linalg::triangular::cholesky_solve_vector(factor.view(), e.view());
3353 for r in 0..q {
3354 ainv_host[row_idx * max_q * max_q + r * max_q + col] = solved[r];
3355 }
3356 }
3357 }
3358
3359 Ok(DeviceSaePcgBuffers {
3360 row_ptr: stream
3361 .clone_htod(&row_ptr_host)
3362 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
3363 beta_base: stream
3364 .clone_htod(&beta_base_host)
3365 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
3366 phi: stream
3367 .clone_htod(&phi_host)
3368 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
3369 jac_ptr: stream
3370 .clone_htod(&jac_ptr_host)
3371 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
3372 jac: stream
3373 .clone_htod(&jac_host)
3374 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
3375 smooth_offsets: stream
3376 .clone_htod(&smooth_offsets_host)
3377 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
3378 smooth_m: stream
3379 .clone_htod(&smooth_m_host)
3380 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
3381 smooth_ptr: stream
3382 .clone_htod(&smooth_ptr_host)
3383 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
3384 smooth_data: stream
3385 .clone_htod(&smooth_data_host)
3386 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
3387 g_row_off: stream
3388 .clone_htod(&g_row_off_host)
3389 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
3390 g_col_off: stream
3391 .clone_htod(&g_col_off_host)
3392 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
3393 g_rows: stream
3394 .clone_htod(&g_rows_host)
3395 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
3396 g_cols: stream
3397 .clone_htod(&g_cols_host)
3398 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
3399 g_ptr: stream
3400 .clone_htod(&g_ptr_host)
3401 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
3402 g_data: stream
3403 .clone_htod(&g_data_host)
3404 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
3405 ainv: stream
3406 .clone_htod(&ainv_host)
3407 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
3408 u: stream
3409 .alloc_zeros::<f64>(n_rows * p)
3410 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
3411 w: stream
3412 .alloc_zeros::<f64>(n_rows * max_q)
3413 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
3414 v: stream
3415 .alloc_zeros::<f64>(n_rows * max_q)
3416 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
3417 n_rows,
3418 p,
3419 k,
3420 max_q,
3421 smooth_blocks: data.smooth_blocks.len(),
3422 g_blocks: data.sparse_g_blocks.len(),
3423 })
3424 }
3425
3426 fn launch_sae_init(
3427 stream: &Arc<CudaStream>,
3428 module: &Arc<CudaModule>,
3429 out: &mut CudaSlice<f64>,
3430 x: &CudaSlice<f64>,
3431 ridge: f64,
3432 n: usize,
3433 ) -> Result<(), ArrowSchurGpuFailure> {
3434 let kernel = module
3435 .load_function("arrow_sae_init")
3436 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
3437 let n_i32 = checked_i32(n)?;
3438 let mut builder = stream.launch_builder(&kernel);
3439 builder.arg(out).arg(x).arg(&ridge).arg(&n_i32);
3440 unsafe { builder.launch(pcg_launch_config(n)?) }
3444 .map(drop)
3445 .map_err(|_| ArrowSchurGpuFailure::Unavailable)
3446 }
3447
3448 fn launch_sae_penalty_matvec(
3449 stream: &Arc<CudaStream>,
3450 module: &Arc<CudaModule>,
3451 buffers: &mut DeviceSaePcgBuffers,
3452 x: &CudaSlice<f64>,
3453 out: &mut CudaSlice<f64>,
3454 ridge_beta: f64,
3455 ) -> Result<(), ArrowSchurGpuFailure> {
3456 launch_sae_init(stream, module, out, x, ridge_beta, buffers.k)?;
3457 if buffers.smooth_blocks > 0 {
3458 let kernel = module
3459 .load_function("arrow_sae_smooth_matvec")
3460 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
3461 let max_m = buffers.k;
3462 let p_i32 = checked_i32(buffers.p)?;
3463 let blocks_i32 = checked_i32(buffers.smooth_blocks)?;
3464 let cfg = LaunchConfig {
3465 grid_dim: (
3466 ((max_m as u32).saturating_add(255) / 256).max(1),
3467 checked_i32(buffers.smooth_blocks)? as u32,
3468 1,
3469 ),
3470 block_dim: (256, 1, 1),
3471 shared_mem_bytes: 0,
3472 };
3473 let mut builder = stream.launch_builder(&kernel);
3474 builder
3475 .arg(x)
3476 .arg(&mut *out)
3477 .arg(&buffers.smooth_offsets)
3478 .arg(&buffers.smooth_m)
3479 .arg(&buffers.smooth_ptr)
3480 .arg(&buffers.smooth_data)
3481 .arg(&p_i32)
3482 .arg(&blocks_i32);
3483 unsafe { builder.launch(cfg) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
3488 }
3489 if buffers.g_blocks > 0 {
3490 let kernel = module
3491 .load_function("arrow_sae_sparse_g_matvec")
3492 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
3493 let max_work = buffers
3494 .k
3495 .checked_div(buffers.p)
3496 .unwrap_or(0)
3497 .saturating_mul(buffers.p);
3498 let p_i32 = checked_i32(buffers.p)?;
3499 let blocks_i32 = checked_i32(buffers.g_blocks)?;
3500 let cfg = LaunchConfig {
3501 grid_dim: (
3502 ((max_work as u32).saturating_add(255) / 256).max(1),
3503 checked_i32(buffers.g_blocks)? as u32,
3504 1,
3505 ),
3506 block_dim: (256, 1, 1),
3507 shared_mem_bytes: 0,
3508 };
3509 let mut builder = stream.launch_builder(&kernel);
3510 builder
3511 .arg(x)
3512 .arg(&mut *out)
3513 .arg(&buffers.g_row_off)
3514 .arg(&buffers.g_col_off)
3515 .arg(&buffers.g_rows)
3516 .arg(&buffers.g_cols)
3517 .arg(&buffers.g_ptr)
3518 .arg(&buffers.g_data)
3519 .arg(&p_i32)
3520 .arg(&blocks_i32);
3521 unsafe { builder.launch(cfg) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
3526 }
3527 Ok(())
3528 }
3529
3530 fn launch_sae_row_schur_sub(
3531 stream: &Arc<CudaStream>,
3532 module: &Arc<CudaModule>,
3533 buffers: &mut DeviceSaePcgBuffers,
3534 x: &CudaSlice<f64>,
3535 out: &mut CudaSlice<f64>,
3536 ) -> Result<(), ArrowSchurGpuFailure> {
3537 let p_i32 = checked_i32(buffers.p)?;
3538 let max_q_i32 = checked_i32(buffers.max_q)?;
3539 let n_rows_i32 = checked_i32(buffers.n_rows)?;
3540 let cfg_p_rows = LaunchConfig {
3541 grid_dim: (
3542 ((buffers.p as u32).saturating_add(255) / 256).max(1),
3543 checked_i32(buffers.n_rows)? as u32,
3544 1,
3545 ),
3546 block_dim: (256, 1, 1),
3547 shared_mem_bytes: 0,
3548 };
3549 let gather = module
3550 .load_function("arrow_sae_gather_u")
3551 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
3552 {
3553 let mut builder = stream.launch_builder(&gather);
3554 builder
3555 .arg(x)
3556 .arg(&buffers.row_ptr)
3557 .arg(&buffers.beta_base)
3558 .arg(&buffers.phi)
3559 .arg(&mut buffers.u)
3560 .arg(&p_i32)
3561 .arg(&n_rows_i32);
3562 unsafe { builder.launch(cfg_p_rows) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
3566 }
3567
3568 let cfg_q_rows = LaunchConfig {
3569 grid_dim: (
3570 ((buffers.max_q as u32).saturating_add(255) / 256).max(1),
3571 checked_i32(buffers.n_rows)? as u32,
3572 1,
3573 ),
3574 block_dim: (256, 1, 1),
3575 shared_mem_bytes: 0,
3576 };
3577 let apply_l = module
3578 .load_function("arrow_sae_apply_l")
3579 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
3580 {
3581 let mut builder = stream.launch_builder(&apply_l);
3582 builder
3583 .arg(&buffers.u)
3584 .arg(&buffers.jac_ptr)
3585 .arg(&buffers.jac)
3586 .arg(&mut buffers.w)
3587 .arg(&p_i32)
3588 .arg(&max_q_i32)
3589 .arg(&n_rows_i32);
3590 unsafe { builder.launch(cfg_q_rows) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
3594 }
3595
3596 let apply_ainv = module
3597 .load_function("arrow_sae_apply_ainv")
3598 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
3599 {
3600 let mut builder = stream.launch_builder(&apply_ainv);
3601 builder
3602 .arg(&buffers.ainv)
3603 .arg(&buffers.w)
3604 .arg(&mut buffers.v)
3605 .arg(&max_q_i32)
3606 .arg(&n_rows_i32);
3607 unsafe { builder.launch(cfg_q_rows) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
3611 }
3612
3613 let scatter = module
3614 .load_function("arrow_sae_scatter_sub")
3615 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
3616 {
3617 let mut builder = stream.launch_builder(&scatter);
3618 builder
3619 .arg(&buffers.v)
3620 .arg(&buffers.jac_ptr)
3621 .arg(&buffers.jac)
3622 .arg(&buffers.row_ptr)
3623 .arg(&buffers.beta_base)
3624 .arg(&buffers.phi)
3625 .arg(out)
3626 .arg(&p_i32)
3627 .arg(&max_q_i32)
3628 .arg(&n_rows_i32);
3629 unsafe { builder.launch(cfg_p_rows) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
3633 }
3634 Ok(())
3635 }
3636
3637 fn launch_sae_diag_sub(
3638 stream: &Arc<CudaStream>,
3639 module: &Arc<CudaModule>,
3640 buffers: &DeviceSaePcgBuffers,
3641 diag: &mut CudaSlice<f64>,
3642 ) -> Result<(), ArrowSchurGpuFailure> {
3643 let kernel = module
3644 .load_function("arrow_sae_diag_sub")
3645 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
3646 let p_i32 = checked_i32(buffers.p)?;
3647 let max_q_i32 = checked_i32(buffers.max_q)?;
3648 let n_rows_i32 = checked_i32(buffers.n_rows)?;
3649 let cfg = LaunchConfig {
3650 grid_dim: (
3651 ((buffers.p as u32).saturating_add(255) / 256).max(1),
3652 checked_i32(buffers.n_rows)? as u32,
3653 1,
3654 ),
3655 block_dim: (256, 1, 1),
3656 shared_mem_bytes: 0,
3657 };
3658 let mut builder = stream.launch_builder(&kernel);
3659 builder
3660 .arg(diag)
3661 .arg(&buffers.ainv)
3662 .arg(&buffers.jac_ptr)
3663 .arg(&buffers.jac)
3664 .arg(&buffers.row_ptr)
3665 .arg(&buffers.beta_base)
3666 .arg(&buffers.phi)
3667 .arg(&p_i32)
3668 .arg(&max_q_i32)
3669 .arg(&n_rows_i32);
3670 unsafe { builder.launch(cfg) }
3674 .map(drop)
3675 .map_err(|_| ArrowSchurGpuFailure::Unavailable)
3676 }
3677
3678 fn launch_sae_matvec(
3679 stream: &Arc<CudaStream>,
3680 module: &Arc<CudaModule>,
3681 buffers: &mut DeviceSaePcgBuffers,
3682 x: &CudaSlice<f64>,
3683 out: &mut CudaSlice<f64>,
3684 ridge_beta: f64,
3685 ) -> Result<(), ArrowSchurGpuFailure> {
3686 launch_sae_penalty_matvec(stream, module, buffers, x, out, ridge_beta)?;
3687 launch_sae_row_schur_sub(stream, module, buffers, x, out)
3688 }
3689
3690 fn pack_fused_host(
3695 sys: &ArrowSchurSystem,
3696 ridge_t: f64,
3697 p_max: usize,
3698 r_template: usize,
3699 ) -> (Vec<f64>, Vec<f64>, Vec<f64>) {
3700 let n = sys.rows.len();
3701 let d = sys.d;
3702 let k = sys.k;
3703 let mut d_buf = vec![0.0_f64; n * p_max * p_max];
3704 let mut b_buf = vec![0.0_f64; n * p_max * r_template];
3705 let mut g_buf = vec![0.0_f64; n * p_max];
3706 for (i, row) in sys.rows.iter().enumerate() {
3707 for col in 0..d {
3709 let base = (i * p_max + col) * p_max;
3710 for r in 0..d {
3711 let mut value = row.htt[[r, col]];
3712 if r == col {
3713 value += ridge_t;
3714 }
3715 d_buf[base + r] = value;
3716 }
3717 }
3718 for col in 0..k {
3726 let base = (i * r_template + col) * p_max;
3727 for r in 0..d {
3728 b_buf[base + r] = row.htbeta[[r, col]];
3729 }
3730 }
3731 let g_base = i * p_max;
3733 for r in 0..d {
3734 g_buf[g_base + r] = row.gt[r];
3735 }
3736 }
3737 (d_buf, b_buf, g_buf)
3738 }
3739
3740 pub(super) struct ResidentArrowFrame {
3767 n: usize,
3768 d: usize,
3769 k: usize,
3770 stream: Arc<CudaStream>,
3771 blas: CudaBlas,
3772 l_dev: CudaSlice<f64>,
3775 y_dev: CudaSlice<f64>,
3778 schur_dev: CudaSlice<f64>,
3781 log_det_hessian: f64,
3784 }
3785
3786 impl ResidentArrowFrame {
3787 pub(super) fn new(
3791 sys: &ArrowSchurSystem,
3792 ridge_t: f64,
3793 ridge_beta: f64,
3794 ) -> Result<Self, ArrowSchurGpuFailure> {
3795 if ridge_t.is_nan() || ridge_beta.is_nan() {
3796 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
3797 reason: "ridge is NaN".to_string(),
3798 });
3799 }
3800 let n = sys.rows.len();
3801 let d = sys.d;
3802 let k = sys.k;
3803 let runtime = route_through_gpu(DispatchOp::SmallDenseBatchedPotrf { p: d, batch: n })
3804 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
3805 let stream = gam_gpu::device_runtime::cuda_context_for(runtime.device.ordinal)
3806 .and_then(|ctx| ctx.new_stream().ok())
3807 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
3808 let solver =
3809 DnHandle::new(stream.clone()).map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
3810 let blas =
3811 CudaBlas::new(stream.clone()).map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
3812
3813 let (d_host, b_host, _g_host) = pack_host(sys, ridge_t);
3815 let mut l_dev = stream
3816 .clone_htod(&d_host)
3817 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
3818 let mut y_dev = stream
3819 .clone_htod(&b_host)
3820 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
3821
3822 let info_host = potrf_batched(&solver, &stream, d, n, &mut l_dev)?;
3824 if let Some(idx) = info_host.iter().position(|info| *info != 0) {
3825 return Err(ArrowSchurGpuFailure::RidgeBumpRequired {
3829 row: idx,
3830 bump: super::ridge_bump_to_make_pd(sys.rows[idx].htt.view(), ridge_t),
3831 });
3832 }
3833
3834 trsm_batched_lower_inplace(&blas, &stream, d, n, k, &l_dev, &mut y_dev)?;
3836
3837 let schur_init: Vec<f64> = {
3842 let mut tmp = Vec::with_capacity(k * k);
3843 for col in 0..k {
3844 for row in 0..k {
3845 let mut v = sys.hbb[[row, col]];
3846 if row == col {
3847 v += ridge_beta;
3848 }
3849 tmp.push(v);
3850 }
3851 }
3852 tmp
3853 };
3854 let mut schur_dev = stream
3855 .clone_htod(&schur_init)
3856 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
3857 let zero_u = stream
3860 .clone_htod(&vec![0.0_f64; n * d])
3861 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
3862 let mut throwaway_rhs = stream
3863 .clone_htod(&vec![0.0_f64; k])
3864 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
3865 accumulate_schur(
3866 &blas,
3867 d,
3868 k,
3869 n,
3870 &y_dev,
3871 &zero_u,
3872 &mut schur_dev,
3873 &mut throwaway_rhs,
3874 )?;
3875 let info = potrf_single(&solver, &stream, k, &mut schur_dev)?;
3876 if info != 0 {
3877 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
3878 reason: format!("Schur Cholesky failed at pivot {info}"),
3879 });
3880 }
3881
3882 let l_local_host = stream
3884 .clone_dtoh(&l_dev)
3885 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
3886 let l_schur_host = stream
3887 .clone_dtoh(&schur_dev)
3888 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
3889 let mut log_det = 0.0_f64;
3890 for i in 0..n {
3891 let base = i * d * d;
3892 for j in 0..d {
3893 log_det += l_local_host[base + j * d + j].ln();
3894 }
3895 }
3896 for j in 0..k {
3897 log_det += l_schur_host[j * k + j].ln();
3898 }
3899 log_det *= 2.0;
3900
3901 Ok(Self {
3902 n,
3903 d,
3904 k,
3905 stream,
3906 blas,
3907 l_dev,
3908 y_dev,
3909 schur_dev,
3910 log_det_hessian: log_det,
3911 })
3912 }
3913
3914 #[inline]
3915 pub(super) fn log_det_hessian(&self) -> f64 {
3916 self.log_det_hessian
3917 }
3918
3919 pub(super) fn solve_gradient(
3923 &self,
3924 g_t: &[f64],
3925 g_beta: &[f64],
3926 ) -> Result<ArrowSchurGpuSolution, ArrowSchurGpuFailure> {
3927 let n = self.n;
3928 let d = self.d;
3929 let k = self.k;
3930 if g_t.len() != n * d || g_beta.len() != k {
3931 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
3932 reason: format!(
3933 "resident gradient shape mismatch: g_t={} (want {}), g_beta={} (want {})",
3934 g_t.len(),
3935 n * d,
3936 g_beta.len(),
3937 k
3938 ),
3939 });
3940 }
3941 let mut u_dev = self
3943 .stream
3944 .clone_htod(&g_t.to_vec())
3945 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
3946 trsm_batched_lower_inplace(&self.blas, &self.stream, d, n, 1, &self.l_dev, &mut u_dev)?;
3947
3948 let rhs_init: Vec<f64> = g_beta.iter().map(|v| -v).collect();
3951 let mut rhs_dev = self
3952 .stream
3953 .clone_htod(&rhs_init)
3954 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
3955 accumulate_schur_rhs_only(&self.blas, d, k, n, &self.y_dev, &u_dev, &mut rhs_dev)?;
3956
3957 trsm_single(
3959 &self.blas,
3960 &self.stream,
3961 k,
3962 &self.schur_dev,
3963 &mut rhs_dev,
3964 false,
3965 false,
3966 )?;
3967 trsm_single(
3968 &self.blas,
3969 &self.stream,
3970 k,
3971 &self.schur_dev,
3972 &mut rhs_dev,
3973 false,
3974 true,
3975 )?;
3976 let delta_beta_host = self
3977 .stream
3978 .clone_dtoh(&rhs_dev)
3979 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
3980 let delta_beta = Array1::from_vec(delta_beta_host);
3981
3982 accumulate_back_sub_rhs(&self.blas, d, k, n, &self.y_dev, &rhs_dev, &mut u_dev)?;
3984 trsm_batched_lower_inplace_transposed(
3985 &self.blas,
3986 &self.stream,
3987 d,
3988 n,
3989 1,
3990 &self.l_dev,
3991 &mut u_dev,
3992 )?;
3993 let x_host = self
3994 .stream
3995 .clone_dtoh(&u_dev)
3996 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
3997 let mut delta_t = Array1::<f64>::zeros(n * d);
3998 for (i, v) in x_host.iter().enumerate() {
3999 delta_t[i] = -*v;
4000 }
4001
4002 Ok(ArrowSchurGpuSolution {
4003 delta_t,
4004 delta_beta,
4005 log_det_hessian: self.log_det_hessian,
4006 })
4007 }
4008 }
4009
4010 pub(super) struct ResidentBaseArrowFrame {
4016 n: usize,
4017 d: usize,
4018 k: usize,
4019 stream: Arc<CudaStream>,
4020 solver: DnHandle,
4021 blas: CudaBlas,
4022 d_base_host: Vec<f64>,
4028 base_b_dev: CudaSlice<f64>,
4031 base_hbb_dev: CudaSlice<f64>,
4034 g_t_dev: CudaSlice<f64>,
4036 gb_host: Vec<f64>,
4039 ones_k_dev: CudaSlice<f64>,
4042 }
4043
4044 impl ResidentBaseArrowFrame {
4045 pub(super) fn new(sys: &ArrowSchurSystem) -> Result<Self, ArrowSchurGpuFailure> {
4047 let n = sys.rows.len();
4048 let d = sys.d;
4049 let k = sys.k;
4050 let runtime = route_through_gpu(DispatchOp::SmallDenseBatchedPotrf { p: d, batch: n })
4051 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
4052 let stream = gam_gpu::device_runtime::cuda_context_for(runtime.device.ordinal)
4053 .and_then(|ctx| ctx.new_stream().ok())
4054 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
4055 let solver =
4056 DnHandle::new(stream.clone()).map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4057 let blas =
4058 CudaBlas::new(stream.clone()).map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4059
4060 let (d_base_host, b_host, g_host) = pack_host(sys, 0.0);
4062 let base_b_dev = stream
4063 .clone_htod(&b_host)
4064 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4065 let g_t_dev = stream
4066 .clone_htod(&g_host)
4067 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4068
4069 let hbb_base: Vec<f64> = {
4071 let mut tmp = Vec::with_capacity(k * k);
4072 for col in 0..k {
4073 for row in 0..k {
4074 tmp.push(sys.hbb[[row, col]]);
4075 }
4076 }
4077 tmp
4078 };
4079 let base_hbb_dev = stream
4080 .clone_htod(&hbb_base)
4081 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4082
4083 let gb_host: Vec<f64> = sys.gb.iter().copied().collect();
4084 let ones_k_dev = stream
4085 .clone_htod(&vec![1.0_f64; k])
4086 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4087
4088 Ok(Self {
4089 n,
4090 d,
4091 k,
4092 stream,
4093 solver,
4094 blas,
4095 d_base_host,
4096 base_b_dev,
4097 base_hbb_dev,
4098 g_t_dev,
4099 gb_host,
4100 ones_k_dev,
4101 })
4102 }
4103
4104 pub(super) fn refactor_and_solve(
4110 &self,
4111 ridge_t: f64,
4112 ridge_beta: f64,
4113 ) -> Result<ArrowSchurGpuSolution, ArrowSchurGpuFailure> {
4114 if ridge_t.is_nan() || ridge_beta.is_nan() {
4115 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
4116 reason: "ridge is NaN".to_string(),
4117 });
4118 }
4119 let (n, d, k) = (self.n, self.d, self.k);
4120
4121 let mut d_ridged = self.d_base_host.clone();
4125 for i in 0..n {
4126 for j in 0..d {
4127 d_ridged[i * d * d + j * (d + 1)] += ridge_t;
4128 }
4129 }
4130 let mut l_dev = self
4131 .stream
4132 .clone_htod(&d_ridged)
4133 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4134
4135 let info_host = potrf_batched(&self.solver, &self.stream, d, n, &mut l_dev)?;
4137 if let Some(idx) = info_host.iter().position(|info| *info != 0) {
4138 let base = idx * d * d;
4139 return Err(ArrowSchurGpuFailure::RidgeBumpRequired {
4140 row: idx,
4141 bump: super::ridge_bump_to_make_pd_colmajor(&d_ridged[base..base + d * d], d),
4145 });
4146 }
4147
4148 let mut y_dev = self
4150 .stream
4151 .alloc_zeros::<f64>(n * d * k)
4152 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4153 self.stream
4154 .memcpy_dtod(&self.base_b_dev, &mut y_dev)
4155 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4156 trsm_batched_lower_inplace(&self.blas, &self.stream, d, n, k, &l_dev, &mut y_dev)?;
4157
4158 let mut u_dev = self
4160 .stream
4161 .alloc_zeros::<f64>(n * d)
4162 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4163 self.stream
4164 .memcpy_dtod(&self.g_t_dev, &mut u_dev)
4165 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4166 trsm_batched_lower_inplace(&self.blas, &self.stream, d, n, 1, &l_dev, &mut u_dev)?;
4167
4168 let mut schur_dev = self
4170 .stream
4171 .alloc_zeros::<f64>(k * k)
4172 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4173 self.stream
4174 .memcpy_dtod(&self.base_hbb_dev, &mut schur_dev)
4175 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4176 device_axpy_strided(
4178 &self.blas,
4179 &self.stream,
4180 k,
4181 ridge_beta,
4182 &self.ones_k_dev,
4183 1,
4184 &mut schur_dev,
4185 k + 1,
4186 )?;
4187 let rhs_init: Vec<f64> = self.gb_host.iter().map(|v| -v).collect();
4188 let mut rhs_dev = self
4189 .stream
4190 .clone_htod(&rhs_init)
4191 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4192
4193 accumulate_schur(
4194 &self.blas,
4195 d,
4196 k,
4197 n,
4198 &y_dev,
4199 &u_dev,
4200 &mut schur_dev,
4201 &mut rhs_dev,
4202 )?;
4203
4204 let info = potrf_single(&self.solver, &self.stream, k, &mut schur_dev)?;
4206 if info != 0 {
4207 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
4208 reason: format!("Schur Cholesky failed at pivot {info}"),
4209 });
4210 }
4211 trsm_single(
4212 &self.blas,
4213 &self.stream,
4214 k,
4215 &schur_dev,
4216 &mut rhs_dev,
4217 false,
4218 false,
4219 )?;
4220 trsm_single(
4221 &self.blas,
4222 &self.stream,
4223 k,
4224 &schur_dev,
4225 &mut rhs_dev,
4226 false,
4227 true,
4228 )?;
4229 let delta_beta_host = self
4230 .stream
4231 .clone_dtoh(&rhs_dev)
4232 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4233 let delta_beta = Array1::from_vec(delta_beta_host);
4234
4235 accumulate_back_sub_rhs(&self.blas, d, k, n, &y_dev, &rhs_dev, &mut u_dev)?;
4237 trsm_batched_lower_inplace_transposed(
4238 &self.blas,
4239 &self.stream,
4240 d,
4241 n,
4242 1,
4243 &l_dev,
4244 &mut u_dev,
4245 )?;
4246 let x_host = self
4247 .stream
4248 .clone_dtoh(&u_dev)
4249 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4250 let mut delta_t = Array1::<f64>::zeros(n * d);
4251 for (i, v) in x_host.iter().enumerate() {
4252 delta_t[i] = -*v;
4253 }
4254
4255 let l_local_host = self
4257 .stream
4258 .clone_dtoh(&l_dev)
4259 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4260 let l_schur_host = self
4261 .stream
4262 .clone_dtoh(&schur_dev)
4263 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4264 let mut log_det = 0.0_f64;
4265 for i in 0..n {
4266 let base = i * d * d;
4267 for j in 0..d {
4268 log_det += l_local_host[base + j * d + j].ln();
4269 }
4270 }
4271 for j in 0..k {
4272 log_det += l_schur_host[j * k + j].ln();
4273 }
4274 log_det *= 2.0;
4275
4276 Ok(ArrowSchurGpuSolution {
4277 delta_t,
4278 delta_beta,
4279 log_det_hessian: log_det,
4280 })
4281 }
4282 }
4283
4284 pub(super) fn solve_fused(
4285 sys: &ArrowSchurSystem,
4286 ridge_t: f64,
4287 ridge_beta: f64,
4288 ) -> Result<ArrowSchurGpuSolution, ArrowSchurGpuFailure> {
4289 let n = sys.rows.len();
4290 let d = sys.d;
4291 let k = sys.k;
4292 let plan = crate::gpu_kernels::arrow_schur_nvrtc::plan_fused_launch(n, d, k)
4293 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
4294 let p_max = plan.p_max;
4295 let r_template = plan.r_template;
4296
4297 let runtime = gam_gpu::linalg_dispatch::route_through_gpu(
4298 gam_gpu::linalg_dispatch::DispatchOp::SmallDenseBatchedPotrf { p: d, batch: n },
4299 )
4300 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
4301 let ctx = gam_gpu::device_runtime::cuda_context_for(runtime.device.ordinal)
4302 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
4303 let stream = ctx
4304 .new_stream()
4305 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4306 let cap = &runtime.device.capability;
4307 let key = crate::gpu_kernels::arrow_schur_nvrtc::FusedModuleCacheKey {
4308 cc_major: cap.compute_major,
4309 cc_minor: cap.compute_minor,
4310 p_max: p_max as u32,
4311 r_template: r_template as u32,
4312 };
4313 let module = fused_module_for(&ctx, key)?;
4314 let forward = module
4315 .load_function("arrow_schur_forward_pgroup")
4316 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4317 let back_sub = module
4318 .load_function("arrow_schur_back_sub_pgroup")
4319 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4320
4321 let (d_host, b_host, g_host) = pack_fused_host(sys, ridge_t, p_max, r_template);
4323 let d_dev = stream
4324 .clone_htod(&d_host)
4325 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4326 let b_dev = stream
4327 .clone_htod(&b_host)
4328 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4329 let g_dev = stream
4330 .clone_htod(&g_host)
4331 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4332 let mut l_out = stream
4333 .alloc_zeros::<f64>(n * p_max * p_max)
4334 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4335 let mut u_out = stream
4336 .alloc_zeros::<f64>(n * p_max)
4337 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4338 let mut y_out = stream
4339 .alloc_zeros::<f64>(n * p_max * r_template)
4340 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4341 let mut partial_s = stream
4342 .alloc_zeros::<f64>(plan.partial_s_doubles)
4343 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4344 let mut partial_r = stream
4345 .alloc_zeros::<f64>(plan.partial_r_doubles)
4346 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4347 let mut status_dev = stream
4348 .alloc_zeros::<i32>(n)
4349 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4350
4351 let cfg = LaunchConfig {
4353 grid_dim: (plan.blocks, 1, 1),
4354 block_dim: (plan.threads_per_block, 1, 1),
4355 shared_mem_bytes: 0,
4356 };
4357 let n_i32 = to_i32(n).ok_or(ArrowSchurGpuFailure::Unavailable)?;
4358 let p_i32 = to_i32(d).ok_or(ArrowSchurGpuFailure::Unavailable)?;
4359 let r_i32 = to_i32(k).ok_or(ArrowSchurGpuFailure::Unavailable)?;
4360 let ridge_arg = ridge_t;
4361 {
4362 let mut builder = stream.launch_builder(&forward);
4363 builder
4364 .arg(&d_dev)
4365 .arg(&b_dev)
4366 .arg(&g_dev)
4367 .arg(&n_i32)
4368 .arg(&p_i32)
4369 .arg(&r_i32)
4370 .arg(&ridge_arg)
4371 .arg(&mut l_out)
4372 .arg(&mut u_out)
4373 .arg(&mut y_out)
4374 .arg(&mut partial_s)
4375 .arg(&mut partial_r)
4376 .arg(&mut status_dev);
4377 unsafe { builder.launch(cfg) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4381 }
4382 stream
4383 .synchronize()
4384 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4385
4386 let status_host = stream
4388 .clone_dtoh(&status_dev)
4389 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4390 if let Some(row) = status_host.iter().position(|s| *s != 0) {
4391 return Err(ArrowSchurGpuFailure::RidgeBumpRequired {
4395 row,
4396 bump: super::ridge_bump_to_make_pd(sys.rows[row].htt.view(), ridge_t),
4397 });
4398 }
4399
4400 let partial_s_host = stream
4402 .clone_dtoh(&partial_s)
4403 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4404 let partial_r_host = stream
4405 .clone_dtoh(&partial_r)
4406 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4407 let mut schur_host = vec![0.0_f64; k * k];
4408 for col in 0..k {
4409 for row in 0..k {
4410 let mut v = sys.hbb[[row, col]];
4411 if row == col {
4412 v += ridge_beta;
4413 }
4414 schur_host[col * k + row] = v;
4415 }
4416 }
4417 let mut rhs_host: Vec<f64> = sys.gb.iter().map(|v| -v).collect();
4418 for i in 0..n {
4419 let s_base = i * r_template * r_template;
4422 for col in 0..k {
4423 let col_base = s_base + col * r_template;
4424 let dst_col_base = col * k;
4425 for row in 0..k {
4426 schur_host[dst_col_base + row] -= partial_s_host[col_base + row];
4427 }
4428 }
4429 let r_base = i * r_template;
4430 for a in 0..k {
4431 rhs_host[a] += partial_r_host[r_base + a];
4432 }
4433 }
4434
4435 let mut schur_dev = stream
4437 .clone_htod(&schur_host)
4438 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4439 let mut rhs_dev = stream
4440 .clone_htod(&rhs_host)
4441 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4442 let solver =
4443 DnHandle::new(stream.clone()).map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4444 let blas = CudaBlas::new(stream.clone()).map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4445 let info = potrf_single(&solver, &stream, k, &mut schur_dev)?;
4446 if info != 0 {
4447 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
4448 reason: format!("fused Schur Cholesky failed at pivot {info}"),
4449 });
4450 }
4451 trsm_single(&blas, &stream, k, &schur_dev, &mut rhs_dev, false, false)?;
4452 trsm_single(&blas, &stream, k, &schur_dev, &mut rhs_dev, false, true)?;
4453 let delta_beta_host = stream
4454 .clone_dtoh(&rhs_dev)
4455 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4456 let delta_beta = Array1::from_vec(delta_beta_host.clone());
4457
4458 let mut delta_t_dev = stream
4460 .alloc_zeros::<f64>(n * p_max)
4461 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4462 let back_cfg = LaunchConfig {
4463 grid_dim: (plan.blocks, 1, 1),
4464 block_dim: (plan.threads_per_block, 1, 1),
4465 shared_mem_bytes: 0,
4466 };
4467 {
4468 let mut builder = stream.launch_builder(&back_sub);
4469 builder
4470 .arg(&l_out)
4471 .arg(&u_out)
4472 .arg(&y_out)
4473 .arg(&rhs_dev)
4474 .arg(&n_i32)
4475 .arg(&p_i32)
4476 .arg(&r_i32)
4477 .arg(&mut delta_t_dev);
4478 unsafe { builder.launch(back_cfg) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4482 }
4483 stream
4484 .synchronize()
4485 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4486
4487 let delta_t_host = stream
4488 .clone_dtoh(&delta_t_dev)
4489 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4490 let mut delta_t = Array1::<f64>::zeros(n * d);
4491 for i in 0..n {
4492 let src_base = i * p_max;
4493 let dst_base = i * d;
4494 for r in 0..d {
4495 delta_t[dst_base + r] = delta_t_host[src_base + r];
4496 }
4497 }
4498
4499 let l_local_host = stream
4501 .clone_dtoh(&l_out)
4502 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4503 let l_schur_host = stream
4504 .clone_dtoh(&schur_dev)
4505 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4506 let mut log_det = 0.0_f64;
4507 for i in 0..n {
4508 let base = i * p_max * p_max;
4509 for j in 0..d {
4510 log_det += l_local_host[base + j * p_max + j].ln();
4511 }
4512 }
4513 for j in 0..k {
4514 log_det += l_schur_host[j * k + j].ln();
4515 }
4516 log_det *= 2.0;
4517
4518 Ok(ArrowSchurGpuSolution {
4519 delta_t,
4520 delta_beta,
4521 log_det_hessian: log_det,
4522 })
4523 }
4524
4525 pub(super) fn build_schur_matvec_backend(
4535 sys: &ArrowSchurSystem,
4536 ridge_t: f64,
4537 ridge_beta: f64,
4538 ) -> Result<crate::arrow_schur::GpuSchurMatvec, super::ArrowSchurGpuFailure> {
4539 let n = sys.rows.len();
4540 let d = sys.d;
4541 let k = sys.k;
4542 let plan = crate::gpu_kernels::arrow_schur_nvrtc::plan_fused_launch(n, d, k)
4543 .ok_or(super::ArrowSchurGpuFailure::Unavailable)?;
4544 let p_max = plan.p_max;
4545 let r_template = plan.r_template;
4546
4547 let runtime = gam_gpu::linalg_dispatch::route_through_gpu(
4548 gam_gpu::linalg_dispatch::DispatchOp::SmallDenseBatchedPotrf { p: d, batch: n },
4549 )
4550 .ok_or(super::ArrowSchurGpuFailure::Unavailable)?;
4551 let ctx = gam_gpu::device_runtime::cuda_context_for(runtime.device.ordinal)
4552 .ok_or(super::ArrowSchurGpuFailure::Unavailable)?;
4553 let stream = ctx
4554 .new_stream()
4555 .map_err(|_| super::ArrowSchurGpuFailure::Unavailable)?;
4556 let cap = &runtime.device.capability;
4557 let key = crate::gpu_kernels::arrow_schur_nvrtc::FusedModuleCacheKey {
4558 cc_major: cap.compute_major,
4559 cc_minor: cap.compute_minor,
4560 p_max: p_max as u32,
4561 r_template: r_template as u32,
4562 };
4563 let module = fused_module_for(&ctx, key)?;
4564 let forward = module
4565 .load_function("arrow_schur_forward_pgroup")
4566 .map_err(|_| super::ArrowSchurGpuFailure::Unavailable)?;
4567
4568 let (d_host, b_host, g_host) = pack_fused_host(sys, ridge_t, p_max, r_template);
4569 let d_dev = stream
4570 .clone_htod(&d_host)
4571 .map_err(|_| super::ArrowSchurGpuFailure::Unavailable)?;
4572 let b_dev = stream
4573 .clone_htod(&b_host)
4574 .map_err(|_| super::ArrowSchurGpuFailure::Unavailable)?;
4575 let g_dev = stream
4576 .clone_htod(&g_host)
4577 .map_err(|_| super::ArrowSchurGpuFailure::Unavailable)?;
4578 let mut l_out = stream
4579 .alloc_zeros::<f64>(n * p_max * p_max)
4580 .map_err(|_| super::ArrowSchurGpuFailure::Unavailable)?;
4581 let mut u_out = stream
4582 .alloc_zeros::<f64>(n * p_max)
4583 .map_err(|_| super::ArrowSchurGpuFailure::Unavailable)?;
4584 let mut y_out = stream
4585 .alloc_zeros::<f64>(n * p_max * r_template)
4586 .map_err(|_| super::ArrowSchurGpuFailure::Unavailable)?;
4587 let mut partial_s = stream
4588 .alloc_zeros::<f64>(plan.partial_s_doubles)
4589 .map_err(|_| super::ArrowSchurGpuFailure::Unavailable)?;
4590 let mut partial_r = stream
4591 .alloc_zeros::<f64>(plan.partial_r_doubles)
4592 .map_err(|_| super::ArrowSchurGpuFailure::Unavailable)?;
4593 let mut status_dev = stream
4594 .alloc_zeros::<i32>(n)
4595 .map_err(|_| super::ArrowSchurGpuFailure::Unavailable)?;
4596
4597 let cfg = LaunchConfig {
4598 grid_dim: (plan.blocks, 1, 1),
4599 block_dim: (plan.threads_per_block, 1, 1),
4600 shared_mem_bytes: 0,
4601 };
4602 let n_i32 = to_i32(n).ok_or(super::ArrowSchurGpuFailure::Unavailable)?;
4603 let p_i32 = to_i32(d).ok_or(super::ArrowSchurGpuFailure::Unavailable)?;
4604 let r_i32 = to_i32(k).ok_or(super::ArrowSchurGpuFailure::Unavailable)?;
4605 let ridge_arg = ridge_t;
4606 {
4607 let mut builder = stream.launch_builder(&forward);
4608 builder
4609 .arg(&d_dev)
4610 .arg(&b_dev)
4611 .arg(&g_dev)
4612 .arg(&n_i32)
4613 .arg(&p_i32)
4614 .arg(&r_i32)
4615 .arg(&ridge_arg)
4616 .arg(&mut l_out)
4617 .arg(&mut u_out)
4618 .arg(&mut y_out)
4619 .arg(&mut partial_s)
4620 .arg(&mut partial_r)
4621 .arg(&mut status_dev);
4622 unsafe { builder.launch(cfg) }.map_err(|_| super::ArrowSchurGpuFailure::Unavailable)?;
4625 }
4626 stream
4627 .synchronize()
4628 .map_err(|_| super::ArrowSchurGpuFailure::Unavailable)?;
4629
4630 let status_host = stream
4631 .clone_dtoh(&status_dev)
4632 .map_err(|_| super::ArrowSchurGpuFailure::Unavailable)?;
4633 if let Some(row) = status_host.iter().position(|s| *s != 0) {
4634 return Err(super::ArrowSchurGpuFailure::RidgeBumpRequired {
4638 row,
4639 bump: super::ridge_bump_to_make_pd(sys.rows[row].htt.view(), ridge_t),
4640 });
4641 }
4642
4643 let y_host = stream
4645 .clone_dtoh(&y_out)
4646 .map_err(|_| super::ArrowSchurGpuFailure::Unavailable)?;
4647
4648 let hbb_host: Vec<f64> = sys.hbb.iter().copied().collect();
4651 let hbb_is_kk = sys.hbb.dim() == (k, k);
4652 let hbb_matvec_opt = sys.hbb_matvec.clone();
4653
4654 let closure: crate::arrow_schur::GpuSchurMatvec =
4655 Arc::new(move |x: &Array1<f64>, out: &mut Array1<f64>| {
4656 assert_eq!(x.len(), k, "gpu_schur_matvec: x.len() != k");
4657 assert_eq!(out.len(), k, "gpu_schur_matvec: out.len() != k");
4658
4659 if let Some(ref mv) = hbb_matvec_opt {
4661 mv(x.view(), out);
4662 for a in 0..k {
4663 out[a] += ridge_beta * x[a];
4664 }
4665 } else if hbb_is_kk {
4666 for a in 0..k {
4668 let mut acc = ridge_beta * x[a];
4669 for b in 0..k {
4670 acc += hbb_host[a * k + b] * x[b];
4671 }
4672 out[a] = acc;
4673 }
4674 } else {
4675 for a in 0..k {
4676 out[a] = ridge_beta * x[a];
4677 }
4678 }
4679
4680 let mut z = vec![0.0_f64; d];
4683 for i in 0..n {
4684 let y_base = i * p_max * r_template;
4685 for r in 0..d {
4686 let mut acc = 0.0;
4687 for c in 0..k {
4688 acc += y_host[y_base + c * p_max + r] * x[c];
4689 }
4690 z[r] = acc;
4691 }
4692 for c in 0..k {
4693 let mut acc = 0.0;
4694 for r in 0..d {
4695 acc += y_host[y_base + c * p_max + r] * z[r];
4696 }
4697 out[c] -= acc;
4698 }
4699 }
4700 });
4701
4702 Ok(closure)
4703 }
4704
4705 struct DeviceSaeFrameBuffers {
4708 s_off: CudaSlice<i32>,
4710 s_m: CudaSlice<i32>,
4711 s_r: CudaSlice<i32>,
4712 s_ptr: CudaSlice<i32>,
4713 s_data: CudaSlice<f64>,
4714 s_blocks: usize,
4715 g_off_i: CudaSlice<i32>,
4717 g_off_j: CudaSlice<i32>,
4718 g_ri: CudaSlice<i32>,
4719 g_rj: CudaSlice<i32>,
4720 g_mi: CudaSlice<i32>,
4721 g_mj: CudaSlice<i32>,
4722 g_ptr: CudaSlice<i32>,
4723 g_data: CudaSlice<f64>,
4724 w_ptr: CudaSlice<i32>,
4725 w_data: CudaSlice<f64>,
4726 g_blocks: usize,
4727 g_max_work: usize,
4728 htb_ptr: CudaSlice<i32>,
4730 htb: CudaSlice<f64>,
4731 q_of: CudaSlice<i32>,
4732 ainv: CudaSlice<f64>,
4733 hvec: CudaSlice<f64>,
4734 svec: CudaSlice<f64>,
4735 scatter_partial: CudaSlice<f64>,
4740 n_chunks: usize,
4741 rows_per_chunk: usize,
4742 n_rows: usize,
4743 k: usize,
4744 max_q: usize,
4745 }
4746
4747 fn scatter_chunking(n_rows: usize) -> (usize, usize) {
4753 let target_chunks = 128usize;
4754 let rows_per_chunk = n_rows.div_ceil(target_chunks).max(1);
4755 let n_chunks = n_rows.div_ceil(rows_per_chunk).max(1);
4756 (rows_per_chunk, n_chunks)
4757 }
4758
4759 fn flatten_device_sae_frame_data(
4760 sys: &ArrowSchurSystem,
4761 data: &DeviceSaePcgData,
4762 frame: &DeviceSaeFrameData,
4763 ridge_t: f64,
4764 stream: &Arc<CudaStream>,
4765 ) -> Result<DeviceSaeFrameBuffers, ArrowSchurGpuFailure> {
4766 let host = super::flatten_frame_host_operands(sys, data, frame)?;
4773 let ainv = super::compute_ainv_host(sys, &host.q_of, host.max_q, host.n_rows, ridge_t)?;
4774 upload_frame_buffers(&host, &ainv, stream)
4775 }
4776
4777 fn upload_frame_buffers(
4783 host: &super::FrameHostOperands,
4784 ainv: &[f64],
4785 stream: &Arc<CudaStream>,
4786 ) -> Result<DeviceSaeFrameBuffers, ArrowSchurGpuFailure> {
4787 let htod_i = |v: &[i32]| {
4788 stream
4789 .clone_htod(v)
4790 .map_err(|_| ArrowSchurGpuFailure::Unavailable)
4791 };
4792 let htod_f = |v: &[f64]| {
4793 stream
4794 .clone_htod(v)
4795 .map_err(|_| ArrowSchurGpuFailure::Unavailable)
4796 };
4797 let (rows_per_chunk, n_chunks) = scatter_chunking(host.n_rows);
4798 Ok(DeviceSaeFrameBuffers {
4799 s_off: htod_i(&host.s_off)?,
4800 s_m: htod_i(&host.s_m)?,
4801 s_r: htod_i(&host.s_r)?,
4802 s_ptr: htod_i(&host.s_ptr)?,
4803 s_data: htod_f(&host.s_data)?,
4804 s_blocks: host.s_blocks,
4805 g_off_i: htod_i(&host.g_off_i)?,
4806 g_off_j: htod_i(&host.g_off_j)?,
4807 g_ri: htod_i(&host.g_ri)?,
4808 g_rj: htod_i(&host.g_rj)?,
4809 g_mi: htod_i(&host.g_mi)?,
4810 g_mj: htod_i(&host.g_mj)?,
4811 g_ptr: htod_i(&host.g_ptr)?,
4812 g_data: htod_f(&host.g_data)?,
4813 w_ptr: htod_i(&host.w_ptr)?,
4814 w_data: htod_f(&host.w_data)?,
4815 g_blocks: host.g_blocks,
4816 g_max_work: host.g_max_work,
4817 htb_ptr: htod_i(&host.htb_ptr)?,
4818 htb: htod_f(&host.htb)?,
4819 q_of: htod_i(&host.q_of)?,
4820 ainv: htod_f(ainv)?,
4821 hvec: stream
4822 .alloc_zeros::<f64>(host.n_rows * host.max_q)
4823 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
4824 svec: stream
4825 .alloc_zeros::<f64>(host.n_rows * host.max_q)
4826 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
4827 scatter_partial: stream
4828 .alloc_zeros::<f64>(n_chunks * host.k)
4829 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
4830 n_chunks,
4831 rows_per_chunk,
4832 n_rows: host.n_rows,
4833 k: host.k,
4834 max_q: host.max_q,
4835 })
4836 }
4837
4838 fn refresh_frame_buffers(
4843 host: &super::FrameHostOperands,
4844 buffers: &mut DeviceSaeFrameBuffers,
4845 stream: &Arc<CudaStream>,
4846 ) -> Result<(), ArrowSchurGpuFailure> {
4847 macro_rules! refresh {
4848 ($host:expr, $device:expr) => {{
4849 if $host.len() != $device.len() {
4850 return Err(ArrowSchurGpuFailure::Unavailable);
4851 }
4852 stream
4853 .memcpy_htod($host, &mut $device)
4854 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4855 }};
4856 }
4857 if host.s_blocks != buffers.s_blocks
4858 || host.g_blocks != buffers.g_blocks
4859 || host.g_max_work != buffers.g_max_work
4860 || host.n_rows != buffers.n_rows
4861 || host.k != buffers.k
4862 || host.max_q != buffers.max_q
4863 {
4864 return Err(ArrowSchurGpuFailure::Unavailable);
4865 }
4866 refresh!(&host.s_off, buffers.s_off);
4867 refresh!(&host.s_m, buffers.s_m);
4868 refresh!(&host.s_r, buffers.s_r);
4869 refresh!(&host.s_ptr, buffers.s_ptr);
4870 refresh!(&host.s_data, buffers.s_data);
4871 refresh!(&host.g_off_i, buffers.g_off_i);
4872 refresh!(&host.g_off_j, buffers.g_off_j);
4873 refresh!(&host.g_ri, buffers.g_ri);
4874 refresh!(&host.g_rj, buffers.g_rj);
4875 refresh!(&host.g_mi, buffers.g_mi);
4876 refresh!(&host.g_mj, buffers.g_mj);
4877 refresh!(&host.g_ptr, buffers.g_ptr);
4878 refresh!(&host.g_data, buffers.g_data);
4879 refresh!(&host.w_ptr, buffers.w_ptr);
4880 refresh!(&host.w_data, buffers.w_data);
4881 refresh!(&host.htb_ptr, buffers.htb_ptr);
4882 refresh!(&host.htb, buffers.htb);
4883 refresh!(&host.q_of, buffers.q_of);
4884 Ok(())
4885 }
4886
4887 fn sae_frame_penalty_diag_host(
4888 data: &DeviceSaePcgData,
4889 frame: &DeviceSaeFrameData,
4890 ridge_beta: f64,
4891 ) -> Result<Vec<f64>, ArrowSchurGpuFailure> {
4892 let mut diag = vec![ridge_beta; data.beta_dim];
4893 for (blk, &r) in data.smooth_blocks.iter().zip(frame.smooth_ranks.iter()) {
4895 let m = blk.factor_a.nrows();
4896 for ia in 0..m {
4897 let coeff = blk.factor_a[[ia, ia]];
4898 let base = blk.global_offset + ia * r;
4899 for ib in 0..r {
4900 if base + ib >= diag.len() {
4901 return Err(ArrowSchurGpuFailure::Unavailable);
4902 }
4903 diag[base + ib] += coeff;
4904 }
4905 }
4906 }
4907 for blk in &frame.frame_blocks {
4909 if blk.atom_i != blk.atom_j {
4910 continue;
4911 }
4912 let r = frame.ranks[blk.atom_i];
4913 let off = frame.border_offsets[blk.atom_i];
4914 let (mi, mj) = blk.g.dim();
4915 for li in 0..mi.min(mj) {
4916 let gii = blk.g[[li, li]];
4917 let base = off + li * r;
4918 for a in 0..r {
4919 if base + a >= diag.len() {
4920 return Err(ArrowSchurGpuFailure::Unavailable);
4921 }
4922 diag[base + a] += gii * blk.w[[a, a]];
4923 }
4924 }
4925 }
4926 Ok(diag)
4927 }
4928
4929 fn frame_grid(work: usize, n_rows: usize) -> Result<LaunchConfig, ArrowSchurGpuFailure> {
4930 Ok(LaunchConfig {
4931 grid_dim: (
4932 ((work as u32).saturating_add(255) / 256).max(1),
4933 checked_i32(n_rows)? as u32,
4934 1,
4935 ),
4936 block_dim: (256, 1, 1),
4937 shared_mem_bytes: 0,
4938 })
4939 }
4940
4941 fn launch_sae_frame_matvec(
4942 stream: &Arc<CudaStream>,
4943 module: &Arc<CudaModule>,
4944 buffers: &mut DeviceSaeFrameBuffers,
4945 x: &CudaSlice<f64>,
4946 out: &mut CudaSlice<f64>,
4947 ridge_beta: f64,
4948 ) -> Result<(), ArrowSchurGpuFailure> {
4949 launch_sae_init(stream, module, out, x, ridge_beta, buffers.k)?;
4950 if buffers.s_blocks > 0 {
4952 let kernel = module
4953 .load_function("arrow_sae_frame_smooth_matvec")
4954 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4955 let blocks_i32 = checked_i32(buffers.s_blocks)?;
4956 let cfg = frame_grid(buffers.k, buffers.s_blocks)?;
4957 let mut b = stream.launch_builder(&kernel);
4958 b.arg(x)
4959 .arg(&mut *out)
4960 .arg(&buffers.s_off)
4961 .arg(&buffers.s_m)
4962 .arg(&buffers.s_r)
4963 .arg(&buffers.s_ptr)
4964 .arg(&buffers.s_data)
4965 .arg(&blocks_i32);
4966 unsafe { b.launch(cfg) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4969 }
4970 if buffers.g_blocks > 0 {
4972 let kernel = module
4973 .load_function("arrow_sae_frame_g_matvec")
4974 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4975 let blocks_i32 = checked_i32(buffers.g_blocks)?;
4976 let cfg = frame_grid(buffers.g_max_work.max(1), buffers.g_blocks)?;
4977 let mut b = stream.launch_builder(&kernel);
4978 b.arg(x)
4979 .arg(&mut *out)
4980 .arg(&buffers.g_off_i)
4981 .arg(&buffers.g_off_j)
4982 .arg(&buffers.g_ri)
4983 .arg(&buffers.g_rj)
4984 .arg(&buffers.g_mi)
4985 .arg(&buffers.g_mj)
4986 .arg(&buffers.g_ptr)
4987 .arg(&buffers.g_data)
4988 .arg(&buffers.w_ptr)
4989 .arg(&buffers.w_data)
4990 .arg(&blocks_i32);
4991 unsafe { b.launch(cfg) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4994 }
4995 let k_i32 = checked_i32(buffers.k)?;
4997 let max_q_i32 = checked_i32(buffers.max_q)?;
4998 let n_rows_i32 = checked_i32(buffers.n_rows)?;
4999 {
5000 let kernel = module
5001 .load_function("arrow_sae_frame_apply_h")
5002 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5003 let cfg = frame_grid(buffers.max_q, buffers.n_rows)?;
5004 let mut b = stream.launch_builder(&kernel);
5005 b.arg(x)
5006 .arg(&buffers.htb_ptr)
5007 .arg(&buffers.htb)
5008 .arg(&buffers.q_of)
5009 .arg(&mut buffers.hvec)
5010 .arg(&k_i32)
5011 .arg(&max_q_i32)
5012 .arg(&n_rows_i32);
5013 unsafe { b.launch(cfg) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5016 }
5017 {
5018 let kernel = module
5019 .load_function("arrow_sae_frame_apply_ainv")
5020 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5021 let cfg = frame_grid(buffers.max_q, buffers.n_rows)?;
5022 let mut b = stream.launch_builder(&kernel);
5023 b.arg(&buffers.ainv)
5024 .arg(&buffers.hvec)
5025 .arg(&buffers.q_of)
5026 .arg(&mut buffers.svec)
5027 .arg(&max_q_i32)
5028 .arg(&n_rows_i32);
5029 unsafe { b.launch(cfg) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5032 }
5033 {
5034 let kernel = module
5035 .load_function("arrow_sae_frame_scatter_h")
5036 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5037 let cfg = frame_grid(buffers.k, buffers.n_rows)?;
5038 let mut b = stream.launch_builder(&kernel);
5039 b.arg(&buffers.svec)
5040 .arg(&buffers.htb_ptr)
5041 .arg(&buffers.htb)
5042 .arg(&buffers.q_of)
5043 .arg(out)
5044 .arg(&k_i32)
5045 .arg(&max_q_i32)
5046 .arg(&n_rows_i32);
5047 unsafe { b.launch(cfg) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5050 }
5051 Ok(())
5052 }
5053
5054 fn launch_sae_frame_reduced_schur_det(
5064 stream: &Arc<CudaStream>,
5065 module: &Arc<CudaModule>,
5066 buffers: &mut DeviceSaeFrameBuffers,
5067 x: &CudaSlice<f64>,
5068 out: &mut CudaSlice<f64>,
5069 ) -> Result<(), ArrowSchurGpuFailure> {
5070 let k_i32 = checked_i32(buffers.k)?;
5071 let max_q_i32 = checked_i32(buffers.max_q)?;
5072 let n_rows_i32 = checked_i32(buffers.n_rows)?;
5073 {
5077 let kernel = module
5078 .load_function("arrow_sae_frame_apply_h_warp")
5079 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5080 let block = checked_i32(buffers.max_q)?
5081 .checked_mul(32)
5082 .ok_or(ArrowSchurGpuFailure::Unavailable)? as u32;
5083 let cfg = LaunchConfig {
5084 grid_dim: (checked_i32(buffers.n_rows)? as u32, 1, 1),
5085 block_dim: (block, 1, 1),
5086 shared_mem_bytes: 0,
5087 };
5088 let mut b = stream.launch_builder(&kernel);
5089 b.arg(x)
5090 .arg(&buffers.htb_ptr)
5091 .arg(&buffers.htb)
5092 .arg(&buffers.q_of)
5093 .arg(&mut buffers.hvec)
5094 .arg(&k_i32)
5095 .arg(&max_q_i32)
5096 .arg(&n_rows_i32);
5097 unsafe { b.launch(cfg) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5101 }
5102 {
5104 let kernel = module
5105 .load_function("arrow_sae_frame_apply_ainv")
5106 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5107 let cfg = frame_grid(buffers.max_q, buffers.n_rows)?;
5108 let mut b = stream.launch_builder(&kernel);
5109 b.arg(&buffers.ainv)
5110 .arg(&buffers.hvec)
5111 .arg(&buffers.q_of)
5112 .arg(&mut buffers.svec)
5113 .arg(&max_q_i32)
5114 .arg(&n_rows_i32);
5115 unsafe { b.launch(cfg) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5118 }
5119 let k_blocks = ((buffers.k as u32).saturating_add(255) / 256).max(1);
5124 {
5125 let kernel = module
5126 .load_function("arrow_sae_frame_scatter_h_det_partial")
5127 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5128 let rows_per_chunk_i32 = checked_i32(buffers.rows_per_chunk)?;
5129 let cfg = LaunchConfig {
5130 grid_dim: (k_blocks, checked_i32(buffers.n_chunks)? as u32, 1),
5131 block_dim: (256, 1, 1),
5132 shared_mem_bytes: 0,
5133 };
5134 let mut b = stream.launch_builder(&kernel);
5135 b.arg(&buffers.svec)
5136 .arg(&buffers.htb_ptr)
5137 .arg(&buffers.htb)
5138 .arg(&buffers.q_of)
5139 .arg(&mut buffers.scatter_partial)
5140 .arg(&k_i32)
5141 .arg(&max_q_i32)
5142 .arg(&n_rows_i32)
5143 .arg(&rows_per_chunk_i32);
5144 unsafe { b.launch(cfg) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5147 }
5148 {
5151 let kernel = module
5152 .load_function("arrow_sae_frame_scatter_h_det_reduce")
5153 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5154 let n_chunks_i32 = checked_i32(buffers.n_chunks)?;
5155 let cfg = LaunchConfig {
5156 grid_dim: (k_blocks, 1, 1),
5157 block_dim: (256, 1, 1),
5158 shared_mem_bytes: 0,
5159 };
5160 let mut b = stream.launch_builder(&kernel);
5161 b.arg(&buffers.scatter_partial)
5162 .arg(&mut *out)
5163 .arg(&k_i32)
5164 .arg(&n_chunks_i32);
5165 unsafe { b.launch(cfg) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5168 }
5169 Ok(())
5170 }
5171
5172 fn launch_sae_frame_diag_sub(
5173 stream: &Arc<CudaStream>,
5174 module: &Arc<CudaModule>,
5175 buffers: &DeviceSaeFrameBuffers,
5176 diag: &mut CudaSlice<f64>,
5177 ) -> Result<(), ArrowSchurGpuFailure> {
5178 let kernel = module
5179 .load_function("arrow_sae_frame_diag_sub")
5180 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5181 let k_i32 = checked_i32(buffers.k)?;
5182 let max_q_i32 = checked_i32(buffers.max_q)?;
5183 let n_rows_i32 = checked_i32(buffers.n_rows)?;
5184 let cfg = frame_grid(buffers.k, buffers.n_rows)?;
5185 let mut b = stream.launch_builder(&kernel);
5186 b.arg(diag)
5187 .arg(&buffers.ainv)
5188 .arg(&buffers.htb_ptr)
5189 .arg(&buffers.htb)
5190 .arg(&buffers.q_of)
5191 .arg(&k_i32)
5192 .arg(&max_q_i32)
5193 .arg(&n_rows_i32);
5194 unsafe { b.launch(cfg) }
5196 .map(drop)
5197 .map_err(|_| ArrowSchurGpuFailure::Unavailable)
5198 }
5199
5200 pub(super) fn framed_schur_matvec_once_on_device(
5211 sys: &ArrowSchurSystem,
5212 data: &DeviceSaePcgData,
5213 ridge_t: f64,
5214 ridge_beta: f64,
5215 x: &Array1<f64>,
5216 ) -> Result<Array1<f64>, ArrowSchurGpuFailure> {
5217 let k = x.len();
5218 if k == 0 || data.beta_dim != k || sys.k != k {
5219 return Err(ArrowSchurGpuFailure::Unavailable);
5220 }
5221 let frame = data
5222 .frame
5223 .as_ref()
5224 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
5225 let runtime = super::resolve_runtime_for_device_path()?
5228 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
5229 let ctx = gam_gpu::device_runtime::cuda_context_for(runtime.selected_device().ordinal)
5230 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
5231 let stream = ctx
5232 .new_stream()
5233 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5234 let vector_module = pcg_vector_module(&ctx)?;
5235 let mut buffers = flatten_device_sae_frame_data(sys, data, frame, ridge_t, &stream)?;
5236 let x_dev = stream
5237 .clone_htod(x.as_slice().ok_or(ArrowSchurGpuFailure::Unavailable)?)
5238 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5239 let mut out_dev = stream
5240 .alloc_zeros::<f64>(k)
5241 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5242 launch_sae_frame_matvec(
5243 &stream,
5244 vector_module,
5245 &mut buffers,
5246 &x_dev,
5247 &mut out_dev,
5248 ridge_beta,
5249 )?;
5250 let out = stream
5251 .clone_dtoh(&out_dev)
5252 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5253 Ok(Array1::from_vec(out))
5254 }
5255
5256 pub(super) fn framed_reduced_schur_det_once_on_device(
5264 sys: &ArrowSchurSystem,
5265 data: &DeviceSaePcgData,
5266 ridge_t: f64,
5267 ridge_beta: f64,
5268 x: &Array1<f64>,
5269 ) -> Result<Array1<f64>, ArrowSchurGpuFailure> {
5270 let k = x.len();
5271 if k == 0 || data.beta_dim != k || sys.k != k {
5272 return Err(ArrowSchurGpuFailure::Unavailable);
5273 }
5274 let frame = data
5275 .frame
5276 .as_ref()
5277 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
5278 let runtime = super::resolve_runtime_for_device_path()?
5279 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
5280 let ctx = gam_gpu::device_runtime::cuda_context_for(runtime.selected_device().ordinal)
5281 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
5282 let stream = ctx
5283 .new_stream()
5284 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5285 let vector_module = pcg_vector_module(&ctx)?;
5286 let mut buffers = flatten_device_sae_frame_data(sys, data, frame, ridge_t, &stream)?;
5287 let x_slice = x.as_slice().ok_or(ArrowSchurGpuFailure::Unavailable)?;
5288 let x_dev = stream
5289 .clone_htod(x_slice)
5290 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5291 let mut reduced_dev = stream
5292 .alloc_zeros::<f64>(k)
5293 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5294 launch_sae_frame_reduced_schur_det(
5295 &stream,
5296 vector_module,
5297 &mut buffers,
5298 &x_dev,
5299 &mut reduced_dev,
5300 )?;
5301 let reduced = stream
5302 .clone_dtoh(&reduced_dev)
5303 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5304 let mut out = vec![0.0_f64; k];
5306 super::sae_framed_penalty_matvec_cpu(data, ridge_beta, x_slice, &mut out);
5307 for a in 0..k {
5308 out[a] += reduced[a];
5309 }
5310 Ok(Array1::from_vec(out))
5311 }
5312
5313 pub(super) fn solve_sae_matrix_free_pcg_framed(
5314 sys: &ArrowSchurSystem,
5315 data: &DeviceSaePcgData,
5316 ridge_t: f64,
5317 ridge_beta: f64,
5318 rhs_beta: &Array1<f64>,
5319 max_iterations: usize,
5320 relative_tolerance: f64,
5321 ) -> Result<(Array1<f64>, ArrowPcgDiagnostics), ArrowSchurGpuFailure> {
5322 let k = rhs_beta.len();
5323 if k == 0 || data.beta_dim != k || sys.k != k {
5324 return Err(ArrowSchurGpuFailure::Unavailable);
5325 }
5326 let frame = data
5327 .frame
5328 .as_ref()
5329 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
5330 let runtime = super::resolve_runtime_for_device_path()?
5331 .filter(|rt| {
5332 rt.policy().reduced_schur_matvec_should_offload(
5333 sys.rows.len(),
5334 sys.k,
5335 sys.d,
5336 max_iterations,
5337 )
5338 })
5339 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
5340 let ctx = gam_gpu::device_runtime::cuda_context_for(runtime.selected_device().ordinal)
5341 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
5342 let stream = ctx
5343 .new_stream()
5344 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5345 let blas = CudaBlas::new(stream.clone()).map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5346 let vector_module = pcg_vector_module(&ctx)?;
5347 log::info!(
5352 "#1017/#2230 {} ridge_t={ridge_t:e} ridge_beta={ridge_beta:e}",
5353 data.operand_byte_report()
5354 );
5355 let mut buffers = flatten_device_sae_frame_data(sys, data, frame, ridge_t, &stream)?;
5356 let dctx = FramedPcgCtx {
5357 stream: &stream,
5358 blas: &blas,
5359 module: vector_module,
5360 max_iterations,
5361 relative_tolerance,
5362 };
5363 pcg_solve_framed_body(data, frame, &mut buffers, ridge_beta, rhs_beta, &dctx)
5364 }
5365
5366 struct FramedPcgCtx<'a> {
5369 stream: &'a Arc<CudaStream>,
5370 blas: &'a CudaBlas,
5371 module: &'a Arc<CudaModule>,
5372 max_iterations: usize,
5373 relative_tolerance: f64,
5374 }
5375
5376 fn pcg_solve_framed_body(
5384 data: &DeviceSaePcgData,
5385 frame: &DeviceSaeFrameData,
5386 buffers: &mut DeviceSaeFrameBuffers,
5387 ridge_beta: f64,
5388 rhs_beta: &Array1<f64>,
5389 dctx: &FramedPcgCtx<'_>,
5390 ) -> Result<(Array1<f64>, ArrowPcgDiagnostics), ArrowSchurGpuFailure> {
5391 let stream = dctx.stream;
5392 let blas = dctx.blas;
5393 let vector_module = dctx.module;
5394 let max_iterations = dctx.max_iterations;
5395 let relative_tolerance = dctx.relative_tolerance;
5396 let k = rhs_beta.len();
5397 let rhs_norm = rhs_beta.iter().map(|v| v * v).sum::<f64>().sqrt();
5398 if rhs_norm == 0.0 {
5399 return Ok((Array1::<f64>::zeros(k), ArrowPcgDiagnostics::default()));
5400 }
5401 let tol = (relative_tolerance.max(0.0) * rhs_norm).max(1e-12);
5402 let rhs_dev = stream
5403 .clone_htod(
5404 rhs_beta
5405 .as_slice()
5406 .ok_or(ArrowSchurGpuFailure::Unavailable)?,
5407 )
5408 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5409 let diag_host = sae_frame_penalty_diag_host(data, frame, ridge_beta)?;
5410 let mut diag_dev = stream
5411 .clone_htod(&diag_host)
5412 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5413 launch_sae_frame_diag_sub(stream, vector_module, buffers, &mut diag_dev)?;
5414 let diag_host = stream
5415 .clone_dtoh(&diag_dev)
5416 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5417 let mut inv_diag = Vec::with_capacity(k);
5418 for (idx, &d) in diag_host.iter().enumerate() {
5419 if !d.is_finite() || d <= 1.0e-18 {
5420 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
5421 reason: format!(
5422 "framed SAE GPU PCG: non-positive Jacobi diagonal at {idx}: {d:e}"
5423 ),
5424 });
5425 }
5426 inv_diag.push(1.0 / d);
5427 }
5428 let inv_diag_dev = stream
5429 .clone_htod(&inv_diag)
5430 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5431
5432 let mut x_dev = stream
5433 .alloc_zeros::<f64>(k)
5434 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5435 let mut r_dev = stream
5436 .alloc_zeros::<f64>(k)
5437 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5438 device_copy(blas, stream, k, &rhs_dev, &mut r_dev)?;
5439 let mut z_dev = stream
5440 .alloc_zeros::<f64>(k)
5441 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5442 launch_jacobi_mul(stream, vector_module, &inv_diag_dev, &r_dev, &mut z_dev, k)?;
5443 let mut p_dev = stream
5444 .alloc_zeros::<f64>(k)
5445 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5446 device_copy(blas, stream, k, &z_dev, &mut p_dev)?;
5447 let mut ap_dev = stream
5448 .alloc_zeros::<f64>(k)
5449 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5450
5451 let mut rz = device_dot(blas, stream, k, &r_dev, &z_dev)?;
5452 if rz <= 0.0 || !rz.is_finite() {
5453 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
5454 reason: format!("framed SAE GPU PCG: non-positive initial rᵀM⁻¹r={rz:e}"),
5455 });
5456 }
5457 let mut diag = ArrowPcgDiagnostics {
5458 precond_apply_calls: 1,
5459 stopping_reason: PcgStopReason::MaxIter,
5460 ..ArrowPcgDiagnostics::default()
5461 };
5462 for _ in 0..max_iterations.max(1) {
5463 launch_sae_frame_matvec(
5464 stream,
5465 vector_module,
5466 buffers,
5467 &p_dev,
5468 &mut ap_dev,
5469 ridge_beta,
5470 )?;
5471 diag.matvec_calls += 1;
5472 diag.iterations += 1;
5473 let pap = device_dot(blas, stream, k, &p_dev, &ap_dev)?;
5474 if pap <= 0.0 || !pap.is_finite() {
5475 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
5476 reason: format!("framed SAE GPU PCG: non-positive curvature pᵀAp={pap:e}"),
5477 });
5478 }
5479 let alpha = rz / pap;
5480 device_axpy(blas, stream, k, alpha, &p_dev, &mut x_dev)?;
5481 device_axpy(blas, stream, k, -alpha, &ap_dev, &mut r_dev)?;
5482 let r_norm = device_nrm2(blas, stream, k, &r_dev)?;
5483 if r_norm <= tol {
5484 diag.final_relative_residual = r_norm / rhs_norm;
5485 diag.stopping_reason = PcgStopReason::Converged;
5486 break;
5487 }
5488 launch_jacobi_mul(stream, vector_module, &inv_diag_dev, &r_dev, &mut z_dev, k)?;
5489 diag.precond_apply_calls += 1;
5490 let rz_new = device_dot(blas, stream, k, &r_dev, &z_dev)?;
5491 if rz_new <= 0.0 || !rz_new.is_finite() {
5492 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
5493 reason: format!("framed SAE GPU PCG: non-positive rᵀM⁻¹r={rz_new:e}"),
5494 });
5495 }
5496 let beta = rz_new / rz;
5497 launch_update_p(stream, vector_module, &z_dev, beta, &mut p_dev, k)?;
5498 rz = rz_new;
5499 }
5500 if diag.stopping_reason != PcgStopReason::Converged {
5501 let r_norm = device_nrm2(blas, stream, k, &r_dev)?;
5502 diag.final_relative_residual = r_norm / rhs_norm;
5503 diag.stopping_reason = PcgStopReason::MaxIter;
5504 }
5505 let x = stream
5506 .clone_dtoh(&x_dev)
5507 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5508 Ok((Array1::from_vec(x), diag))
5509 }
5510
5511 pub(crate) struct ResidentSaeFrameHandle {
5520 ctx: Arc<CudaContext>,
5521 stream: Arc<CudaStream>,
5522 buffers: std::sync::Mutex<DeviceSaeFrameBuffers>,
5523 q_of: Vec<i32>,
5524 max_q: usize,
5525 n_rows: usize,
5526 k: usize,
5527 }
5528
5529 impl ResidentSaeFrameHandle {
5530 pub(crate) fn build(
5538 sys: &ArrowSchurSystem,
5539 cg_iters: usize,
5540 ) -> Result<Option<Self>, ArrowSchurGpuFailure> {
5541 let Some(data) = sys.device_sae_pcg.as_ref() else {
5542 return Ok(None);
5543 };
5544 let Some(frame) = data.frame.as_ref() else {
5545 return Ok(None);
5546 };
5547 if sys.k == 0 || data.beta_dim != sys.k {
5548 return Ok(None);
5549 }
5550 let Some(runtime) = super::resolve_runtime_for_device_path()? else {
5551 return Ok(None);
5552 };
5553 if !runtime.policy().reduced_schur_matvec_should_offload(
5554 sys.rows.len(),
5555 sys.k,
5556 sys.d,
5557 cg_iters,
5558 ) {
5559 return Ok(None);
5560 }
5561 let ctx = gam_gpu::device_runtime::cuda_context_for(runtime.selected_device().ordinal)
5562 .ok_or_else(|| ArrowSchurGpuFailure::SchurFactorFailed {
5563 reason: "resident SAE frame could not bind the admitted CUDA context"
5564 .to_string(),
5565 })?;
5566 let stream = ctx.new_stream().map_err(|error| {
5567 ArrowSchurGpuFailure::SchurFactorFailed {
5568 reason: format!("resident SAE frame stream creation failed: {error}"),
5569 }
5570 })?;
5571 let host = super::flatten_frame_host_operands(sys, data, frame)?;
5572 log::info!(
5579 "#1017 SAE resident frame ENGAGED: {} uploaded ONCE for the ladder; \
5580 per-trial re-upload now only ainv ({}rows × {}²·8B)",
5581 data.operand_byte_report(),
5582 host.n_rows,
5583 host.max_q
5584 );
5585 let zero_ainv = vec![0.0_f64; host.n_rows * host.max_q * host.max_q];
5586 let buffers = upload_frame_buffers(&host, &zero_ainv, &stream)?;
5587 Ok(Some(Self {
5588 ctx,
5589 stream,
5590 buffers: std::sync::Mutex::new(buffers),
5591 q_of: host.q_of,
5592 max_q: host.max_q,
5593 n_rows: host.n_rows,
5594 k: host.k,
5595 }))
5596 }
5597
5598 pub(super) fn prime_ainv(
5605 &self,
5606 sys: &ArrowSchurSystem,
5607 ridge_t: f64,
5608 ) -> Result<(), ArrowSchurGpuFailure> {
5609 if sys.k != self.k || sys.rows.len() != self.n_rows {
5610 return Err(ArrowSchurGpuFailure::Unavailable);
5611 }
5612 let ainv =
5613 super::compute_ainv_host(sys, &self.q_of, self.max_q, self.n_rows, ridge_t)?;
5614 let mut buffers = self
5615 .buffers
5616 .lock()
5617 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5618 if ainv.len() != buffers.ainv.len() {
5619 return Err(ArrowSchurGpuFailure::Unavailable);
5620 }
5621 self.stream
5622 .memcpy_htod(&ainv, &mut buffers.ainv)
5623 .map_err(|_| ArrowSchurGpuFailure::Unavailable)
5624 }
5625 }
5626
5627 impl super::SaeResidentFrame for ResidentSaeFrameHandle {
5628 fn refresh(&self, sys: &ArrowSchurSystem) -> Result<(), ArrowSchurGpuFailure> {
5629 if sys.k != self.k || sys.rows.len() != self.n_rows {
5630 return Err(ArrowSchurGpuFailure::Unavailable);
5631 }
5632 let data = sys
5633 .device_sae_pcg
5634 .as_ref()
5635 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
5636 let frame = data
5637 .frame
5638 .as_ref()
5639 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
5640 let host = super::flatten_frame_host_operands(sys, data, frame)?;
5641 if host.q_of != self.q_of || host.max_q != self.max_q {
5642 return Err(ArrowSchurGpuFailure::Unavailable);
5643 }
5644 let mut buffers = self
5645 .buffers
5646 .lock()
5647 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5648 refresh_frame_buffers(&host, &mut buffers, &self.stream)
5649 }
5650
5651 fn resolve(
5652 &self,
5653 sys: &ArrowSchurSystem,
5654 ridge_t: f64,
5655 ridge_beta: f64,
5656 rhs_beta: &Array1<f64>,
5657 max_iterations: usize,
5658 relative_tolerance: f64,
5659 ) -> Result<(Array1<f64>, ArrowPcgDiagnostics), ArrowSchurGpuFailure> {
5660 if sys.k != self.k || sys.rows.len() != self.n_rows || rhs_beta.len() != self.k {
5664 return Err(ArrowSchurGpuFailure::Unavailable);
5665 }
5666 let data = sys
5667 .device_sae_pcg
5668 .as_ref()
5669 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
5670 let frame = data
5671 .frame
5672 .as_ref()
5673 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
5674 let ainv = super::compute_ainv_host(sys, &self.q_of, self.max_q, self.n_rows, ridge_t)?;
5678 let mut buffers = self
5679 .buffers
5680 .lock()
5681 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5682 if ainv.len() != buffers.ainv.len() {
5683 return Err(ArrowSchurGpuFailure::Unavailable);
5684 }
5685 self.stream
5686 .memcpy_htod(&ainv, &mut buffers.ainv)
5687 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5688 let blas = CudaBlas::new(self.stream.clone())
5689 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5690 let vector_module = pcg_vector_module(&self.ctx)?;
5691 let dctx = FramedPcgCtx {
5692 stream: &self.stream,
5693 blas: &blas,
5694 module: vector_module,
5695 max_iterations,
5696 relative_tolerance,
5697 };
5698 pcg_solve_framed_body(data, frame, &mut *buffers, ridge_beta, rhs_beta, &dctx)
5699 }
5700 }
5701
5702 pub(super) fn build_framed_resident_evidence_matvec(
5714 sys: &ArrowSchurSystem,
5715 ridge_t: f64,
5716 ridge_beta: f64,
5717 apply_budget: usize,
5718 ) -> Result<Option<crate::arrow_schur::GpuSchurMatvec>, ArrowSchurGpuFailure> {
5719 let Some(handle) = ResidentSaeFrameHandle::build(sys, apply_budget)? else {
5720 return Ok(None);
5721 };
5722 handle.prime_ainv(sys, ridge_t)?;
5723 let Some(data) = sys.device_sae_pcg.as_ref().cloned() else {
5724 return Ok(None);
5725 };
5726 let handle = Arc::new(handle);
5727 let k = handle.k;
5728 let closure: crate::arrow_schur::GpuSchurMatvec =
5729 Arc::new(move |x: &Array1<f64>, out: &mut Array1<f64>| {
5730 assert_eq!(x.len(), k, "#1017 framed evidence matvec: x.len() != k");
5731 assert_eq!(out.len(), k, "#1017 framed evidence matvec: out.len() != k");
5732 let module = pcg_vector_module(&handle.ctx)
5733 .expect("#1017 framed evidence matvec: pcg_vector_module unavailable");
5734 let x_slice = x
5735 .as_slice()
5736 .expect("#1017 framed evidence matvec: x not contiguous");
5737 let x_dev = handle
5738 .stream
5739 .clone_htod(x_slice)
5740 .expect("#1017 framed evidence matvec: htod(x) failed");
5741 let mut reduced_dev = handle
5742 .stream
5743 .alloc_zeros::<f64>(k)
5744 .expect("#1017 framed evidence matvec: device alloc failed");
5745 {
5746 let mut buffers = handle
5747 .buffers
5748 .lock()
5749 .expect("#1017 framed evidence matvec: resident frame poisoned");
5750 launch_sae_frame_reduced_schur_det(
5751 &handle.stream,
5752 module,
5753 &mut buffers,
5754 &x_dev,
5755 &mut reduced_dev,
5756 )
5757 .expect("#1017 framed evidence matvec: device reduced-Schur launch failed");
5758 }
5759 let reduced = handle
5760 .stream
5761 .clone_dtoh(&reduced_dev)
5762 .expect("#1017 framed evidence matvec: dtoh(out) failed");
5763 let out_slice = out
5764 .as_slice_mut()
5765 .expect("#1017 framed evidence matvec: out not contiguous");
5766 super::sae_framed_penalty_matvec_cpu(&data, ridge_beta, x_slice, out_slice);
5767 for a in 0..k {
5768 out_slice[a] += reduced[a];
5769 }
5770 });
5771 Ok(Some(closure))
5772 }
5773
5774 pub(super) fn solve_sae_matrix_free_pcg(
5781 sys: &ArrowSchurSystem,
5782 data: &DeviceSaePcgData,
5783 ridge_t: f64,
5784 ridge_beta: f64,
5785 rhs_beta: &Array1<f64>,
5786 max_iterations: usize,
5787 relative_tolerance: f64,
5788 ) -> Result<(Array1<f64>, ArrowPcgDiagnostics), ArrowSchurGpuFailure> {
5789 let k = rhs_beta.len();
5790 if k == 0 || data.beta_dim != k || sys.k != k {
5791 return Err(ArrowSchurGpuFailure::Unavailable);
5792 }
5793 if data.frame.is_some() {
5797 return Err(ArrowSchurGpuFailure::Unavailable);
5798 }
5799 let runtime = super::resolve_runtime_for_device_path()?
5813 .filter(|rt| {
5814 rt.policy().reduced_schur_matvec_should_offload(
5815 sys.rows.len(),
5816 sys.k,
5817 sys.d,
5818 max_iterations,
5819 )
5820 })
5821 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
5822 let ctx = gam_gpu::device_runtime::cuda_context_for(runtime.selected_device().ordinal)
5823 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
5824 let stream = ctx
5825 .new_stream()
5826 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5827 let blas = CudaBlas::new(stream.clone()).map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5828 let vector_module = pcg_vector_module(&ctx)?;
5829 log::info!(
5833 "#1017/#2230 {} ridge_t={ridge_t:e} ridge_beta={ridge_beta:e}",
5834 data.operand_byte_report()
5835 );
5836 let mut buffers = flatten_device_sae_data(sys, data, ridge_t, &stream)?;
5837
5838 let rhs_norm = rhs_beta.iter().map(|v| v * v).sum::<f64>().sqrt();
5839 if rhs_norm == 0.0 {
5840 return Ok((Array1::<f64>::zeros(k), ArrowPcgDiagnostics::default()));
5841 }
5842 let tol = (relative_tolerance.max(0.0) * rhs_norm).max(1e-12);
5843 let rhs_dev = stream
5844 .clone_htod(
5845 rhs_beta
5846 .as_slice()
5847 .ok_or(ArrowSchurGpuFailure::Unavailable)?,
5848 )
5849 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5850 let diag_host = sae_penalty_diag_host(data, ridge_beta)?;
5851 let mut diag_dev = stream
5852 .clone_htod(&diag_host)
5853 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5854 launch_sae_diag_sub(&stream, vector_module, &buffers, &mut diag_dev)?;
5855 let diag_host = stream
5856 .clone_dtoh(&diag_dev)
5857 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5858 let mut inv_diag = Vec::with_capacity(k);
5859 for (idx, &d) in diag_host.iter().enumerate() {
5860 if !d.is_finite() || d <= 1.0e-18 {
5861 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
5862 reason: format!(
5863 "SAE matrix-free GPU PCG: non-positive Schur Jacobi diagonal at {idx}: {d:e}"
5864 ),
5865 });
5866 }
5867 inv_diag.push(1.0 / d);
5868 }
5869 let inv_diag_dev = stream
5870 .clone_htod(&inv_diag)
5871 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5872
5873 let mut x_dev = stream
5874 .alloc_zeros::<f64>(k)
5875 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5876 let mut r_dev = stream
5877 .alloc_zeros::<f64>(k)
5878 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5879 device_copy(&blas, &stream, k, &rhs_dev, &mut r_dev)?;
5880 let mut z_dev = stream
5881 .alloc_zeros::<f64>(k)
5882 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5883 launch_jacobi_mul(&stream, vector_module, &inv_diag_dev, &r_dev, &mut z_dev, k)?;
5884 let mut p_dev = stream
5885 .alloc_zeros::<f64>(k)
5886 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5887 device_copy(&blas, &stream, k, &z_dev, &mut p_dev)?;
5888 let mut ap_dev = stream
5889 .alloc_zeros::<f64>(k)
5890 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5891
5892 let mut rz = device_dot(&blas, &stream, k, &r_dev, &z_dev)?;
5893 if rz <= 0.0 || !rz.is_finite() {
5894 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
5895 reason: format!("SAE matrix-free GPU PCG: non-positive initial rᵀM⁻¹r={rz:e}"),
5896 });
5897 }
5898 let mut diag = ArrowPcgDiagnostics {
5899 precond_apply_calls: 1,
5900 stopping_reason: PcgStopReason::MaxIter,
5901 ..ArrowPcgDiagnostics::default()
5902 };
5903
5904 for _ in 0..max_iterations.max(1) {
5905 launch_sae_matvec(
5906 &stream,
5907 vector_module,
5908 &mut buffers,
5909 &p_dev,
5910 &mut ap_dev,
5911 ridge_beta,
5912 )?;
5913 diag.matvec_calls += 1;
5914 diag.iterations += 1;
5915 let pap = device_dot(&blas, &stream, k, &p_dev, &ap_dev)?;
5916 if pap <= 0.0 || !pap.is_finite() {
5917 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
5918 reason: format!("SAE matrix-free GPU PCG: non-positive curvature pᵀAp={pap:e}"),
5919 });
5920 }
5921 let alpha = rz / pap;
5922 device_axpy(&blas, &stream, k, alpha, &p_dev, &mut x_dev)?;
5923 device_axpy(&blas, &stream, k, -alpha, &ap_dev, &mut r_dev)?;
5924 let r_norm = device_nrm2(&blas, &stream, k, &r_dev)?;
5925 if r_norm <= tol {
5926 diag.final_relative_residual = r_norm / rhs_norm;
5927 diag.stopping_reason = PcgStopReason::Converged;
5928 break;
5929 }
5930 launch_jacobi_mul(&stream, vector_module, &inv_diag_dev, &r_dev, &mut z_dev, k)?;
5931 diag.precond_apply_calls += 1;
5932 let rz_new = device_dot(&blas, &stream, k, &r_dev, &z_dev)?;
5933 if rz_new <= 0.0 || !rz_new.is_finite() {
5934 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
5935 reason: format!("SAE matrix-free GPU PCG: non-positive rᵀM⁻¹r={rz_new:e}"),
5936 });
5937 }
5938 let beta = rz_new / rz;
5939 launch_update_p(&stream, vector_module, &z_dev, beta, &mut p_dev, k)?;
5940 rz = rz_new;
5941 }
5942 if diag.stopping_reason != PcgStopReason::Converged {
5943 let r_norm = device_nrm2(&blas, &stream, k, &r_dev)?;
5944 diag.final_relative_residual = r_norm / rhs_norm;
5945 diag.stopping_reason = PcgStopReason::MaxIter;
5946 }
5947 let x = stream
5948 .clone_dtoh(&x_dev)
5949 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5950 Ok((Array1::from_vec(x), diag))
5951 }
5952
5953 pub(super) fn solve_reduced_beta_pcg_with_diagnostics(
5954 s_acc: &ndarray::Array2<f64>,
5955 rhs_beta: &Array1<f64>,
5956 max_iterations: usize,
5957 relative_tolerance: f64,
5958 ) -> Result<(Array1<f64>, ArrowPcgDiagnostics), ArrowSchurGpuFailure> {
5959 let k = rhs_beta.len();
5960 let cg_iters = max_iterations.max(1);
5972 let runtime = gam_gpu::linalg_dispatch::route_through_gpu(
5973 gam_gpu::linalg_dispatch::DispatchOp::Gemm {
5974 m: k,
5975 n: k,
5976 k: cg_iters,
5977 },
5978 )
5979 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
5980 let stream = gam_gpu::device_runtime::cuda_context_for(runtime.device.ordinal)
5981 .and_then(|ctx| ctx.new_stream().ok())
5982 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
5983 let blas = CudaBlas::new(stream.clone()).map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5984 let ctx = gam_gpu::device_runtime::cuda_context_for(runtime.device.ordinal)
5985 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
5986 let vector_module = pcg_vector_module(&ctx)?;
5987
5988 let mut inv_diag = vec![0.0_f64; k];
5990 for j in 0..k {
5991 let djj = s_acc[[j, j]];
5992 if !(djj > 0.0) {
5993 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
5994 reason: format!(
5995 "reduced-β GPU PCG: Jacobi diagonal S[{j},{j}]={djj:e} not positive"
5996 ),
5997 });
5998 }
5999 inv_diag[j] = 1.0 / djj;
6000 }
6001
6002 let mut s_host = vec![0.0_f64; k * k];
6004 for col in 0..k {
6005 for row in 0..k {
6006 s_host[col * k + row] = s_acc[[row, col]];
6007 }
6008 }
6009 let s_dev = stream
6010 .clone_htod(&s_host)
6011 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6012
6013 let rhs_norm = rhs_beta.iter().map(|v| v * v).sum::<f64>().sqrt();
6017 if rhs_norm == 0.0 {
6018 return Ok((Array1::<f64>::zeros(k), ArrowPcgDiagnostics::default()));
6019 }
6020 let tol = (relative_tolerance.max(0.0) * rhs_norm).max(1e-12);
6021
6022 let mut x_dev = stream
6025 .alloc_zeros::<f64>(k)
6026 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6027 let mut r_dev = stream
6028 .clone_htod(
6029 rhs_beta
6030 .as_slice()
6031 .ok_or(ArrowSchurGpuFailure::Unavailable)?,
6032 )
6033 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6034 let inv_diag_dev = stream
6035 .clone_htod(&inv_diag)
6036 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6037 let mut z_dev = stream
6038 .alloc_zeros::<f64>(k)
6039 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6040 launch_jacobi_mul(&stream, vector_module, &inv_diag_dev, &r_dev, &mut z_dev, k)?;
6041 let mut p_dev = stream
6042 .alloc_zeros::<f64>(k)
6043 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6044 device_copy(&blas, &stream, k, &z_dev, &mut p_dev)?;
6045 let mut sp_dev = stream
6046 .alloc_zeros::<f64>(k)
6047 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6048 let mut rz = device_dot(&blas, &stream, k, &r_dev, &z_dev)?;
6049 let mut diag = ArrowPcgDiagnostics {
6050 precond_apply_calls: 1,
6051 stopping_reason: PcgStopReason::MaxIter,
6052 ..ArrowPcgDiagnostics::default()
6053 };
6054 if rz <= 0.0 || !rz.is_finite() {
6055 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
6056 reason: format!("reduced-β GPU PCG: non-positive initial rᵀM⁻¹r={rz:e}"),
6057 });
6058 }
6059
6060 let max_iters = max_iterations.max(1);
6061 for _ in 0..max_iters {
6062 let gemv_cfg = GemvConfig::<f64> {
6064 trans: cublasOperation_t::CUBLAS_OP_N,
6065 m: to_i32(k).ok_or(ArrowSchurGpuFailure::Unavailable)?,
6066 n: to_i32(k).ok_or(ArrowSchurGpuFailure::Unavailable)?,
6067 alpha: 1.0,
6068 lda: to_i32(k).ok_or(ArrowSchurGpuFailure::Unavailable)?,
6069 incx: 1,
6070 beta: 0.0,
6071 incy: 1,
6072 };
6073 unsafe { blas.gemv(gemv_cfg, &s_dev, &p_dev, &mut sp_dev) }
6075 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6076 diag.matvec_calls += 1;
6077 diag.iterations += 1;
6078
6079 let p_sp = device_dot(&blas, &stream, k, &p_dev, &sp_dev)?;
6080 if !(p_sp > 0.0) {
6081 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
6084 reason: format!("reduced-β GPU PCG: non-positive curvature pᵀSp={p_sp:e}"),
6085 });
6086 }
6087 let alpha = rz / p_sp;
6088 device_axpy(&blas, &stream, k, alpha, &p_dev, &mut x_dev)?;
6089 device_axpy(&blas, &stream, k, -alpha, &sp_dev, &mut r_dev)?;
6090 let r_norm = device_nrm2(&blas, &stream, k, &r_dev)?;
6091 if r_norm <= tol {
6092 diag.final_relative_residual = r_norm / rhs_norm;
6093 diag.stopping_reason = PcgStopReason::Converged;
6094 break;
6095 }
6096 launch_jacobi_mul(&stream, vector_module, &inv_diag_dev, &r_dev, &mut z_dev, k)?;
6097 diag.precond_apply_calls += 1;
6098 let rz_new = device_dot(&blas, &stream, k, &r_dev, &z_dev)?;
6099 if rz_new <= 0.0 || !rz_new.is_finite() {
6100 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
6101 reason: format!("reduced-β GPU PCG: non-positive rᵀM⁻¹r={rz_new:e}"),
6102 });
6103 }
6104 let beta = rz_new / rz;
6105 launch_update_p(&stream, vector_module, &z_dev, beta, &mut p_dev, k)?;
6106 rz = rz_new;
6107 }
6108 if diag.stopping_reason != PcgStopReason::Converged {
6109 let r_norm = device_nrm2(&blas, &stream, k, &r_dev)?;
6110 diag.final_relative_residual = r_norm / rhs_norm;
6111 diag.stopping_reason = PcgStopReason::MaxIter;
6112 }
6113
6114 let x = stream
6115 .clone_dtoh(&x_dev)
6116 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6117 Ok((Array1::from_vec(x), diag))
6118 }
6119
6120 fn device_copy(
6121 blas: &CudaBlas,
6122 stream: &Arc<CudaStream>,
6123 n: usize,
6124 src: &CudaSlice<f64>,
6125 dst: &mut CudaSlice<f64>,
6126 ) -> Result<(), ArrowSchurGpuFailure> {
6127 let n_i = to_i32(n).ok_or(ArrowSchurGpuFailure::Unavailable)?;
6128 let (src_ptr, _src_rec) = src.device_ptr(stream);
6129 let (dst_ptr, _dst_rec) = dst.device_ptr_mut(stream);
6130 let status = unsafe {
6133 cudarc::cublas::sys::cublasDcopy_v2(
6134 *blas.handle(),
6135 n_i,
6136 src_ptr as *const f64,
6137 1,
6138 dst_ptr as *mut f64,
6139 1,
6140 )
6141 };
6142 if status == cublasStatus_t::CUBLAS_STATUS_SUCCESS {
6143 Ok(())
6144 } else {
6145 Err(ArrowSchurGpuFailure::Unavailable)
6146 }
6147 }
6148
6149 fn device_axpy(
6150 blas: &CudaBlas,
6151 stream: &Arc<CudaStream>,
6152 n: usize,
6153 alpha: f64,
6154 x: &CudaSlice<f64>,
6155 y: &mut CudaSlice<f64>,
6156 ) -> Result<(), ArrowSchurGpuFailure> {
6157 let n_i = to_i32(n).ok_or(ArrowSchurGpuFailure::Unavailable)?;
6158 let (x_ptr, _x_rec) = x.device_ptr(stream);
6159 let (y_ptr, _y_rec) = y.device_ptr_mut(stream);
6160 let status = unsafe {
6163 cudarc::cublas::sys::cublasDaxpy_v2(
6164 *blas.handle(),
6165 n_i,
6166 &alpha,
6167 x_ptr as *const f64,
6168 1,
6169 y_ptr as *mut f64,
6170 1,
6171 )
6172 };
6173 if status == cublasStatus_t::CUBLAS_STATUS_SUCCESS {
6174 Ok(())
6175 } else {
6176 Err(ArrowSchurGpuFailure::Unavailable)
6177 }
6178 }
6179
6180 fn device_axpy_strided(
6185 blas: &CudaBlas,
6186 stream: &Arc<CudaStream>,
6187 n: usize,
6188 alpha: f64,
6189 x: &CudaSlice<f64>,
6190 incx: usize,
6191 y: &mut CudaSlice<f64>,
6192 incy: usize,
6193 ) -> Result<(), ArrowSchurGpuFailure> {
6194 let n_i = to_i32(n).ok_or(ArrowSchurGpuFailure::Unavailable)?;
6195 let incx_i = to_i32(incx).ok_or(ArrowSchurGpuFailure::Unavailable)?;
6196 let incy_i = to_i32(incy).ok_or(ArrowSchurGpuFailure::Unavailable)?;
6197 let (x_ptr, _x_rec) = x.device_ptr(stream);
6198 let (y_ptr, _y_rec) = y.device_ptr_mut(stream);
6199 let status = unsafe {
6202 cudarc::cublas::sys::cublasDaxpy_v2(
6203 *blas.handle(),
6204 n_i,
6205 &alpha,
6206 x_ptr as *const f64,
6207 incx_i,
6208 y_ptr as *mut f64,
6209 incy_i,
6210 )
6211 };
6212 if status == cublasStatus_t::CUBLAS_STATUS_SUCCESS {
6213 Ok(())
6214 } else {
6215 Err(ArrowSchurGpuFailure::Unavailable)
6216 }
6217 }
6218
6219 fn device_dot(
6220 blas: &CudaBlas,
6221 stream: &Arc<CudaStream>,
6222 n: usize,
6223 x: &CudaSlice<f64>,
6224 y: &CudaSlice<f64>,
6225 ) -> Result<f64, ArrowSchurGpuFailure> {
6226 let n_i = to_i32(n).ok_or(ArrowSchurGpuFailure::Unavailable)?;
6227 let (x_ptr, _x_rec) = x.device_ptr(stream);
6228 let (y_ptr, _y_rec) = y.device_ptr(stream);
6229 let mut result = 0.0_f64;
6230 let status = unsafe {
6234 cudarc::cublas::sys::cublasDdot_v2(
6235 *blas.handle(),
6236 n_i,
6237 x_ptr as *const f64,
6238 1,
6239 y_ptr as *const f64,
6240 1,
6241 &mut result,
6242 )
6243 };
6244 if status == cublasStatus_t::CUBLAS_STATUS_SUCCESS {
6245 Ok(result)
6246 } else {
6247 Err(ArrowSchurGpuFailure::Unavailable)
6248 }
6249 }
6250
6251 fn device_nrm2(
6252 blas: &CudaBlas,
6253 stream: &Arc<CudaStream>,
6254 n: usize,
6255 x: &CudaSlice<f64>,
6256 ) -> Result<f64, ArrowSchurGpuFailure> {
6257 let n_i = to_i32(n).ok_or(ArrowSchurGpuFailure::Unavailable)?;
6258 let (x_ptr, _x_rec) = x.device_ptr(stream);
6259 let mut result = 0.0_f64;
6260 let status = unsafe {
6264 cudarc::cublas::sys::cublasDnrm2_v2(
6265 *blas.handle(),
6266 n_i,
6267 x_ptr as *const f64,
6268 1,
6269 &mut result,
6270 )
6271 };
6272 if status == cublasStatus_t::CUBLAS_STATUS_SUCCESS {
6273 Ok(result)
6274 } else {
6275 Err(ArrowSchurGpuFailure::Unavailable)
6276 }
6277 }
6278
6279 #[cfg(test)]
6280 mod tests {
6281 use super::*;
6286 use crate::arrow_schur::{
6287 ArrowSchurSystem, DeviceSaeFrameData, DeviceSaePcgData, DeviceSaeSmoothBlock,
6288 FactoredFrameGBlock,
6289 };
6290 use ndarray::Array2;
6291
6292 fn tiny_framed_fixture() -> (ArrowSchurSystem, DeviceSaePcgData) {
6296 let p = 3usize;
6297 let ranks = vec![2usize, 3usize];
6298 let basis_sizes = vec![2usize, 2usize];
6299 let mut border_offsets = Vec::new();
6300 let mut acc = 0usize;
6301 for k in 0..2 {
6302 border_offsets.push(acc);
6303 acc += basis_sizes[k] * ranks[k];
6304 }
6305 let border_dim = acc;
6306 let frame_of = |k: usize| -> Array2<f64> {
6307 Array2::from_shape_fn((p, ranks[k]), |(i, j)| {
6308 0.1 + 0.2 * ((i + 1) as f64) * ((j + 1 + 2 * k) as f64)
6309 })
6310 };
6311 let frames: Vec<Array2<f64>> = (0..2).map(frame_of).collect();
6312 let w_of = |i: usize, j: usize| -> Array2<f64> {
6313 let (ui, uj) = (&frames[i], &frames[j]);
6314 Array2::from_shape_fn((ranks[i], ranks[j]), |(a, b)| {
6315 (0..p).map(|c| ui[[c, a]] * uj[[c, b]]).sum()
6316 })
6317 };
6318 let mut frame_blocks = Vec::new();
6319 for &(i, j) in &[(0usize, 0usize), (1usize, 1usize), (0, 1), (1, 0)] {
6320 let (mi, mj) = (basis_sizes[i], basis_sizes[j]);
6321 let mut g =
6322 Array2::<f64>::from_shape_fn((mi, mj), |(r, c)| 0.1 * (r + 2 * c + 1) as f64);
6323 if i == j {
6324 for r in 0..mi.min(mj) {
6325 g[[r, r]] += mi as f64 + 2.0;
6326 }
6327 }
6328 frame_blocks.push(FactoredFrameGBlock {
6329 atom_i: i,
6330 atom_j: j,
6331 g,
6332 w: w_of(i, j),
6333 });
6334 }
6335 let mut smooth_blocks = Vec::new();
6336 for k in 0..2 {
6337 let m = basis_sizes[k];
6338 let mut s =
6339 Array2::<f64>::from_shape_fn((m, m), |(r, c)| 0.05 * (r + c + 1) as f64);
6340 for r in 0..m {
6341 s[[r, r]] += 1.0;
6342 }
6343 smooth_blocks.push(DeviceSaeSmoothBlock {
6344 global_offset: border_offsets[k],
6345 factor_a: s,
6346 });
6347 }
6348 let smooth_ranks = ranks.clone();
6349 let n = 2usize;
6350 let q = 2usize;
6351 let mut sys = ArrowSchurSystem::new(n, q, border_dim);
6352 let mut row_htbeta = Vec::new();
6353 for i in 0..n {
6354 let mut htt =
6355 Array2::<f64>::from_shape_fn((q, q), |(r, c)| 0.3 * (r + c + 1) as f64);
6356 for r in 0..q {
6357 htt[[r, r]] += q as f64 + 2.0;
6358 }
6359 sys.rows[i].htt = htt;
6360 let mut slab = vec![0.0_f64; q * border_dim];
6361 for c in 0..q {
6362 for col in 0..border_dim {
6363 let v = 0.01 * ((c + 1) * (col + 1) + i) as f64;
6364 slab[c * border_dim + col] = v;
6365 sys.rows[i].htbeta[[c, col]] = v;
6366 }
6367 }
6368 row_htbeta.push(slab);
6369 }
6370 let data = DeviceSaePcgData {
6371 p,
6372 beta_dim: border_dim,
6373 a_phi: std::sync::Arc::from(Vec::new().into_boxed_slice()),
6374 local_jac: std::sync::Arc::from(Vec::new().into_boxed_slice()),
6375 smooth_blocks,
6376 sparse_g_blocks: Vec::new(),
6377 frame: Some(DeviceSaeFrameData {
6378 ranks,
6379 basis_sizes,
6380 border_offsets,
6381 frame_blocks,
6382 smooth_ranks,
6383 row_htbeta,
6384 }),
6385 };
6386 (sys, data)
6387 }
6388
6389 #[test]
6401 fn sae_resident_frame_only_ainv_is_ridge_dependent_1017() {
6402 let (sys, data) = tiny_framed_fixture();
6403 let frame = data.frame.as_ref().expect("framed fixture");
6404
6405 let host_a = super::super::flatten_frame_host_operands(&sys, &data, frame)
6408 .expect("host operands a");
6409 let host_b = super::super::flatten_frame_host_operands(&sys, &data, frame)
6410 .expect("host operands b");
6411 assert_eq!(
6412 host_a.s_data, host_b.s_data,
6413 "smooth λS must be ridge-independent"
6414 );
6415 assert_eq!(
6416 host_a.g_data, host_b.g_data,
6417 "frame G must be ridge-independent"
6418 );
6419 assert_eq!(
6420 host_a.w_data, host_b.w_data,
6421 "frame W must be ridge-independent"
6422 );
6423 assert_eq!(host_a.htb, host_b.htb, "row H_tβ must be ridge-independent");
6424 assert_eq!(host_a.q_of, host_b.q_of);
6425
6426 let report = data.operand_byte_report();
6427 assert!(report.framed, "fixture must exercise the framed lane");
6428 assert!(
6429 report.total_bytes > 0,
6430 "framed operands must have nonzero bytes"
6431 );
6432
6433 let ainv_lo = super::super::compute_ainv_host(
6437 &sys,
6438 &host_a.q_of,
6439 host_a.max_q,
6440 host_a.n_rows,
6441 1e-3,
6442 )
6443 .expect("ainv lo");
6444 let ainv_lo2 = super::super::compute_ainv_host(
6445 &sys,
6446 &host_a.q_of,
6447 host_a.max_q,
6448 host_a.n_rows,
6449 1e-3,
6450 )
6451 .expect("ainv lo repeat");
6452 let ainv_hi = super::super::compute_ainv_host(
6453 &sys,
6454 &host_a.q_of,
6455 host_a.max_q,
6456 host_a.n_rows,
6457 1e3,
6458 )
6459 .expect("ainv hi");
6460 assert_eq!(
6461 ainv_lo, ainv_lo2,
6462 "ainv must be deterministic at a fixed ridge"
6463 );
6464 let max_diff = ainv_lo
6465 .iter()
6466 .zip(&ainv_hi)
6467 .map(|(a, b)| (a - b).abs())
6468 .fold(0.0_f64, f64::max);
6469 assert!(
6470 max_diff > 1e-6,
6471 "ainv MUST change with ridge_t (else the split is wrong); max_diff={max_diff:e}"
6472 );
6473
6474 let trials = crate::arrow_schur::DEFAULT_PROXIMAL_MAX_ATTEMPTS + 1;
6478 let saved = report.total_bytes * (trials - 1);
6479 assert!(saved > 0);
6480 eprintln!(
6481 "#1017 resident-frame ladder saving: {trials} trials × {}B ridge-independent \
6482 operand upload → resident removes {saved}B, re-uploading only ainv \
6483 ({}rows × {}² × 8B) per trial",
6484 report.total_bytes, host_a.n_rows, host_a.max_q
6485 );
6486 }
6487
6488 fn device_matvec_once(
6491 sys: &ArrowSchurSystem,
6492 data: &DeviceSaePcgData,
6493 ridge_t: f64,
6494 ridge_beta: f64,
6495 x_host: &[f64],
6496 ) -> Result<Vec<f64>, ArrowSchurGpuFailure> {
6497 let k = x_host.len();
6498 let frame = data
6499 .frame
6500 .as_ref()
6501 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
6502 let runtime = super::super::resolve_runtime_for_device_path()?
6503 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
6504 let ctx = gam_gpu::device_runtime::cuda_context_for(runtime.selected_device().ordinal)
6505 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
6506 let stream = ctx
6507 .new_stream()
6508 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6509 let vector_module = pcg_vector_module(&ctx)?;
6510 let mut buffers = flatten_device_sae_frame_data(sys, data, frame, ridge_t, &stream)?;
6511 let x_dev = stream
6512 .clone_htod(x_host)
6513 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6514 let mut out_dev = stream
6515 .alloc_zeros::<f64>(k)
6516 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6517 launch_sae_frame_matvec(
6518 &stream,
6519 vector_module,
6520 &mut buffers,
6521 &x_dev,
6522 &mut out_dev,
6523 ridge_beta,
6524 )?;
6525 stream
6526 .clone_dtoh(&out_dev)
6527 .map_err(|_| ArrowSchurGpuFailure::Unavailable)
6528 }
6529
6530 #[test]
6536 fn framed_sae_device_matvec_stage_diff_tiny_1551() {
6537 if gam_gpu::device_runtime::GpuRuntime::resolve(gam_gpu::GpuPolicy::Auto)
6538 .unwrap_or_else(|error| panic!("GPU probe fault in framed matvec test: {error}"))
6539 .is_none()
6540 {
6541 return;
6542 }
6543 let p = 3usize;
6544 let ranks = vec![2usize, 3usize];
6545 let basis_sizes = vec![2usize, 2usize];
6546 let mut border_offsets = Vec::new();
6547 let mut acc = 0usize;
6548 for k in 0..2 {
6549 border_offsets.push(acc);
6550 acc += basis_sizes[k] * ranks[k];
6551 }
6552 let border_dim = acc; let frame_of = |k: usize| -> Array2<f64> {
6554 Array2::from_shape_fn((p, ranks[k]), |(i, j)| {
6555 0.1 + 0.2 * ((i + 1) as f64) * ((j + 1 + 2 * k) as f64)
6556 })
6557 };
6558 let frames: Vec<Array2<f64>> = (0..2).map(frame_of).collect();
6559 let w_of = |i: usize, j: usize| -> Array2<f64> {
6560 let (ui, uj) = (&frames[i], &frames[j]);
6561 Array2::from_shape_fn((ranks[i], ranks[j]), |(a, b)| {
6562 (0..p).map(|c| ui[[c, a]] * uj[[c, b]]).sum()
6563 })
6564 };
6565 let mut frame_blocks = Vec::new();
6566 for &(i, j) in &[(0usize, 0usize), (1usize, 1usize), (0, 1), (1, 0)] {
6567 let (mi, mj) = (basis_sizes[i], basis_sizes[j]);
6568 let mut g =
6569 Array2::<f64>::from_shape_fn((mi, mj), |(r, c)| 0.1 * (r + 2 * c + 1) as f64);
6570 if i == j {
6571 for r in 0..mi.min(mj) {
6572 g[[r, r]] += mi as f64 + 2.0;
6573 }
6574 }
6575 frame_blocks.push(FactoredFrameGBlock {
6576 atom_i: i,
6577 atom_j: j,
6578 g,
6579 w: w_of(i, j),
6580 });
6581 }
6582 let mut smooth_blocks = Vec::new();
6583 for k in 0..2 {
6584 let m = basis_sizes[k];
6585 let mut s =
6586 Array2::<f64>::from_shape_fn((m, m), |(r, c)| 0.05 * (r + c + 1) as f64);
6587 for r in 0..m {
6588 s[[r, r]] += 1.0;
6589 }
6590 smooth_blocks.push(DeviceSaeSmoothBlock {
6591 global_offset: border_offsets[k],
6592 factor_a: s,
6593 });
6594 }
6595 let smooth_ranks = ranks.clone();
6596 let n = 2usize;
6597 let q = 2usize;
6598 let mut sys = ArrowSchurSystem::new(n, q, border_dim);
6599 let mut row_htbeta = Vec::new();
6600 for i in 0..n {
6601 let mut htt =
6602 Array2::<f64>::from_shape_fn((q, q), |(r, c)| 0.3 * (r + c + 1) as f64);
6603 for r in 0..q {
6604 htt[[r, r]] += q as f64 + 2.0;
6605 }
6606 sys.rows[i].htt = htt;
6607 let mut slab = vec![0.0_f64; q * border_dim];
6608 for c in 0..q {
6609 for col in 0..border_dim {
6610 let v = 0.01 * ((c + 1) * (col + 1) + i) as f64;
6611 slab[c * border_dim + col] = v;
6612 sys.rows[i].htbeta[[c, col]] = v;
6613 }
6614 }
6615 row_htbeta.push(slab);
6616 }
6617 let data = DeviceSaePcgData {
6618 p,
6619 beta_dim: border_dim,
6620 a_phi: std::sync::Arc::from(Vec::new().into_boxed_slice()),
6621 local_jac: std::sync::Arc::from(Vec::new().into_boxed_slice()),
6622 smooth_blocks,
6623 sparse_g_blocks: Vec::new(),
6624 frame: Some(DeviceSaeFrameData {
6625 ranks,
6626 basis_sizes,
6627 border_offsets,
6628 frame_blocks,
6629 smooth_ranks,
6630 row_htbeta,
6631 }),
6632 };
6633 let ridge_t = 1e-7;
6634 let ridge_beta = 1e-6;
6635 let mut first_bad: Option<usize> = None;
6636 let mut worst = 0.0_f64;
6637 let mut worst_at = 0usize;
6638 let mut worst_dev = 0.0_f64;
6639 let mut worst_cpu = 0.0_f64;
6640 for col in 0..border_dim {
6641 let mut x = vec![0.0_f64; border_dim];
6642 x[col] = 1.0;
6643 let dev = match device_matvec_once(&sys, &data, ridge_t, ridge_beta, &x) {
6644 Ok(v) => v,
6645 Err(_) => return,
6646 };
6647 let mut cpu = vec![0.0_f64; border_dim];
6648 super::super::sae_framed_schur_matvec_cpu(
6649 &sys, &data, ridge_t, ridge_beta, &x, &mut cpu,
6650 )
6651 .expect("cpu matvec");
6652 for r in 0..border_dim {
6653 let d = (dev[r] - cpu[r]).abs();
6654 if d > 1e-9 && first_bad.is_none() {
6655 first_bad = Some(r * border_dim + col);
6656 }
6657 if d > worst {
6658 worst = d;
6659 worst_at = r * border_dim + col;
6660 worst_dev = dev[r];
6661 worst_cpu = cpu[r];
6662 }
6663 }
6664 }
6665 assert!(
6666 worst <= 1e-9,
6667 "[#1551 stage-diff] device framed matvec != CPU oracle: worst abs={worst:e} at \
6668 (row*K+col)={worst_at} (dev={worst_dev:e} cpu={worst_cpu:e}), \
6669 first_bad_idx={first_bad:?}; border layout: atom0 [0..4) rank2, atom1 [4..10) \
6670 rank3 — which atom-range the bad row/col falls in pins the stage (smooth=diag, \
6671 G⊗W=cross, reduced-Schur=dense per-row)",
6672 );
6673 }
6674 }
6675}
6676
6677#[cfg(test)]
6678mod tests {
6679 use super::*;
6680 use crate::arrow_schur::ArrowSchurSystem;
6681 use ndarray::{Array2, ArrayView1};
6682
6683 fn build_fixture(n: usize, d: usize, k: usize, seed: u64) -> ArrowSchurSystem {
6684 let mut sys = ArrowSchurSystem::new(n, d, k);
6685 let mut state = seed.wrapping_mul(0x9E37_79B9_7F4A_7C15);
6686 let mut sample = || -> f64 {
6687 state = state
6688 .wrapping_mul(6364136223846793005)
6689 .wrapping_add(1442695040888963407);
6690 ((state >> 33) as f64) / ((1u64 << 31) as f64) - 1.0
6691 };
6692 for row in &mut sys.rows {
6693 let mut a = Array2::<f64>::zeros((d, d));
6694 for r in 0..d {
6695 for c in 0..d {
6696 a[[r, c]] = sample();
6697 }
6698 }
6699 let mut htt = a.t().dot(&a);
6700 for r in 0..d {
6701 htt[[r, r]] += d as f64 + 1.0;
6702 }
6703 row.htt = htt;
6704 for r in 0..d {
6705 for c in 0..k {
6706 row.htbeta[[r, c]] = 0.1 * sample();
6707 }
6708 row.gt[r] = sample();
6709 }
6710 }
6711 let mut hbb_a = Array2::<f64>::zeros((k, k));
6712 for r in 0..k {
6713 for c in 0..k {
6714 hbb_a[[r, c]] = sample();
6715 }
6716 }
6717 let mut hbb = hbb_a.t().dot(&hbb_a);
6718 for r in 0..k {
6719 hbb[[r, r]] += k as f64 + 1.0;
6720 }
6721 sys.hbb = hbb;
6722 for r in 0..k {
6723 sys.gb[r] = sample();
6724 }
6725 sys
6726 }
6727
6728 #[test]
6733 fn ridge_bump_makes_known_indefinite_blocks_pd() {
6734 let neg_identity = Array2::<f64>::from_diag(&Array1::from_elem(8, -1.0)); let scaled_neg = Array2::<f64>::from_diag(&Array1::from_elem(4, -250.0)); let mut indef2 = Array2::<f64>::zeros((2, 2));
6742 indef2[[0, 0]] = 1.0;
6743 indef2[[1, 1]] = 1.0;
6744 indef2[[0, 1]] = 2.0;
6745 indef2[[1, 0]] = 2.0;
6746 let pd = Array2::<f64>::from_diag(&Array1::from_elem(3, 5.0));
6749
6750 for (label, block) in [
6751 ("-I (λ_min=-1)", neg_identity),
6752 ("-250·I (λ_min=-250)", scaled_neg),
6753 ("[[1,2],[2,1]] (λ_min=-1)", indef2),
6754 ("5·I (PD)", pd),
6755 ] {
6756 let ridge_t = 0.0;
6757 let bump = ridge_bump_to_make_pd(block.view(), ridge_t);
6758 assert!(
6759 bump > 0.0 && bump.is_finite(),
6760 "[{label}] bump must be strictly positive and finite, got {bump:e}"
6761 );
6762 let d = block.nrows();
6763 let mut shifted = block.clone();
6764 for i in 0..d {
6765 shifted[[i, i]] += ridge_t + bump;
6766 }
6767 assert!(
6768 cholesky_factor_in_place(shifted.view(), CholeskyGuard::NonnegativePivot).is_some(),
6769 "[{label}] H_tt + (ridge_t + bump={bump:e})·I must be PD after the \
6770 Gershgorin bump, but the Cholesky still rejected it"
6771 );
6772 }
6773 }
6774
6775 #[cfg(target_os = "linux")]
6785 #[test]
6786 fn ridge_bump_colmajor_matches_rowmajor_for_symmetric_block() {
6787 let mut a = Array2::<f64>::zeros((3, 3));
6789 a[[0, 0]] = -2.0;
6790 a[[1, 1]] = 0.5;
6791 a[[2, 2]] = 1.0;
6792 a[[0, 1]] = 0.3;
6793 a[[1, 0]] = 0.3;
6794 a[[1, 2]] = -0.4;
6795 a[[2, 1]] = -0.4;
6796 a[[0, 2]] = 0.1;
6797 a[[2, 0]] = 0.1;
6798
6799 let row_major_bump = ridge_bump_to_make_pd(a.view(), 0.0);
6800
6801 let d = 3;
6803 let mut col_major = vec![0.0_f64; d * d];
6804 for c in 0..d {
6805 for r in 0..d {
6806 col_major[c * d + r] = a[[r, c]];
6807 }
6808 }
6809 let col_major_bump = ridge_bump_to_make_pd_colmajor(&col_major, d);
6810
6811 assert!(
6812 (row_major_bump - col_major_bump).abs() <= 1e-12 * row_major_bump.max(1.0),
6813 "colmajor bump {col_major_bump:e} must match rowmajor bump \
6814 {row_major_bump:e} for a symmetric block"
6815 );
6816
6817 let mut shifted = a.clone();
6819 for i in 0..d {
6820 shifted[[i, i]] += col_major_bump;
6821 }
6822 assert!(
6823 cholesky_factor_in_place(shifted.view(), CholeskyGuard::NonnegativePivot).is_some(),
6824 "colmajor Gershgorin bump must make the symmetric block PD"
6825 );
6826 }
6827
6828 fn device_pcg_fixture(k: usize) -> (Array2<f64>, Array1<f64>) {
6829 let mut s = Array2::<f64>::zeros((k, k));
6830 for row in 0..k {
6831 s[[row, row]] = 2.5 + 0.001 * ((row % 17) as f64);
6832 if row + 1 < k {
6833 s[[row, row + 1]] = -0.05;
6834 s[[row + 1, row]] = -0.05;
6835 }
6836 if row + 7 < k {
6837 s[[row, row + 7]] = 0.01;
6838 s[[row + 7, row]] = 0.01;
6839 }
6840 }
6841 let rhs = Array1::from_shape_fn(k, |idx| ((idx as f64 + 1.0) * 0.013).sin());
6842 (s, rhs)
6843 }
6844
6845 fn dense_pcg_cpu_reference(
6846 s: &Array2<f64>,
6847 rhs: &Array1<f64>,
6848 max_iterations: usize,
6849 relative_tolerance: f64,
6850 ) -> Array1<f64> {
6851 let k = rhs.len();
6852 let rhs_norm = rhs.iter().map(|v| v * v).sum::<f64>().sqrt();
6853 if rhs_norm == 0.0 {
6854 return Array1::<f64>::zeros(k);
6855 }
6856 let tol = (relative_tolerance.max(0.0) * rhs_norm).max(1e-12);
6857 let inv_diag: Vec<f64> = (0..k).map(|idx| 1.0 / s[[idx, idx]]).collect();
6858 let mut x = Array1::<f64>::zeros(k);
6859 let mut r = rhs.clone();
6860 let mut z = Array1::from_shape_fn(k, |idx| inv_diag[idx] * r[idx]);
6861 let mut p = z.clone();
6862 let mut sp = Array1::<f64>::zeros(k);
6863 let mut rz = r.iter().zip(z.iter()).map(|(a, b)| a * b).sum::<f64>();
6864 for _ in 0..max_iterations.max(1) {
6865 for row in 0..k {
6866 let mut acc = 0.0;
6867 for col in 0..k {
6868 acc += s[[row, col]] * p[col];
6869 }
6870 sp[row] = acc;
6871 }
6872 let p_sp = p.iter().zip(sp.iter()).map(|(a, b)| a * b).sum::<f64>();
6873 let alpha = rz / p_sp;
6874 for idx in 0..k {
6875 x[idx] += alpha * p[idx];
6876 r[idx] -= alpha * sp[idx];
6877 }
6878 let r_norm = r.iter().map(|v| v * v).sum::<f64>().sqrt();
6879 if r_norm <= tol {
6880 break;
6881 }
6882 for idx in 0..k {
6883 z[idx] = inv_diag[idx] * r[idx];
6884 }
6885 let rz_next = r.iter().zip(z.iter()).map(|(a, b)| a * b).sum::<f64>();
6886 let beta = rz_next / rz;
6887 for idx in 0..k {
6888 p[idx] = z[idx] + beta * p[idx];
6889 }
6890 rz = rz_next;
6891 }
6892 x
6893 }
6894
6895 #[test]
6896 fn device_resident_pcg_matches_cpu_reference_when_cuda_admits() {
6897 let (s, rhs) = device_pcg_fixture(512);
6898 let max_iterations = 200usize;
6899 let relative_tolerance = 1.0e-12;
6900 let cpu = dense_pcg_cpu_reference(&s, &rhs, max_iterations, relative_tolerance);
6901 let (device, diag) = match solve_reduced_beta_pcg_with_diagnostics(
6902 &s,
6903 &rhs,
6904 max_iterations,
6905 relative_tolerance,
6906 ) {
6907 Ok(result) => result,
6908 Err(failure) => {
6915 assert!(
6916 gam_gpu::device_runtime::GpuRuntime::resolve(gam_gpu::GpuPolicy::Auto)
6917 .unwrap_or_else(|error| {
6918 panic!("GPU probe fault in reduced-beta PCG test: {error}")
6919 })
6920 .is_none(),
6921 "#1017: CUDA device present but the device reduced-beta PCG \
6922 declined/faulted instead of returning a result (tag: {failure:?}) — \
6923 the kernel does not run correctly on GPU"
6924 );
6925 return;
6926 }
6927 };
6928 let max_err = cpu
6929 .iter()
6930 .zip(device.iter())
6931 .map(|(a, b)| (a - b).abs())
6932 .fold(0.0_f64, f64::max);
6933 assert!(
6934 max_err <= 1.0e-10,
6935 "device resident PCG parity failed: max_err={max_err:e}, diag={diag:?}"
6936 );
6937 assert!(diag.matvec_calls > 0);
6938 assert_eq!(diag.matvec_calls, diag.iterations);
6939 }
6940
6941 #[test]
6942 fn dense_reference_matches_independent_solve() {
6943 let sys = build_fixture(4, 5, 3, 7);
6944 let solution = solve_arrow_newton_step_dense_reference(&sys, 0.0, 0.0).unwrap();
6945 let n = sys.rows.len();
6949 let d = sys.d;
6950 let k = sys.k;
6951 let total = n * d + k;
6952 let mut h = Array2::<f64>::zeros((total, total));
6953 let mut g = ndarray::Array1::<f64>::zeros(total);
6954 for (i, row) in sys.rows.iter().enumerate() {
6955 let base = i * d;
6956 for c in 0..d {
6957 for r in 0..d {
6958 h[[base + r, base + c]] = row.htt[[r, c]];
6959 }
6960 }
6961 for c in 0..k {
6962 for r in 0..d {
6963 h[[base + r, n * d + c]] = row.htbeta[[r, c]];
6964 h[[n * d + c, base + r]] = row.htbeta[[r, c]];
6965 }
6966 }
6967 for r in 0..d {
6968 g[base + r] = row.gt[r];
6969 }
6970 }
6971 for c in 0..k {
6972 for r in 0..k {
6973 h[[n * d + r, n * d + c]] += sys.hbb[[r, c]];
6974 }
6975 g[n * d + c] = sys.gb[c];
6976 }
6977 let l = cholesky_factor_in_place(h.view(), CholeskyGuard::NonnegativePivot).unwrap();
6978 let rhs = g.mapv(|v| -v);
6979 let expected = cholesky_solve_vector(l.view(), rhs.view());
6980 for i in 0..n * d {
6981 assert!(
6982 (solution.delta_t[i] - expected[i]).abs() < 1e-10 * (1.0 + expected[i].abs()),
6983 "delta_t[{i}] mismatch: got {} expected {}",
6984 solution.delta_t[i],
6985 expected[i]
6986 );
6987 }
6988 for a in 0..k {
6989 assert!(
6990 (solution.delta_beta[a] - expected[n * d + a]).abs()
6991 < 1e-10 * (1.0 + expected[n * d + a].abs()),
6992 "delta_beta[{a}] mismatch"
6993 );
6994 }
6995 }
6996
6997 #[test]
7011 fn row_procedural_matvec_parallel_deterministic_and_matches_serial() {
7012 use crate::arrow_schur::SCHUR_MATVEC_PARALLEL_ROW_MIN;
7013 let n = SCHUR_MATVEC_PARALLEL_ROW_MIN + 96; let d = 3usize;
7015 let k = 24usize;
7016 let mut sys = build_fixture(n, d, k, 0xA17C_0FFE);
7017 let slabs: Vec<Array2<f64>> = sys.rows.iter().map(|row| row.htbeta.clone()).collect();
7022 let forward_slabs = slabs.clone();
7023 let transpose_slabs = slabs;
7024 sys.set_row_htbeta_operator(
7025 move |row: usize, x: ArrayView1<'_, f64>, out: &mut Array1<f64>| {
7026 let h = &forward_slabs[row];
7027 for r in 0..h.nrows() {
7028 let mut acc = 0.0_f64;
7029 for c in 0..h.ncols() {
7030 acc += h[[r, c]] * x[c];
7031 }
7032 out[r] = acc;
7033 }
7034 },
7035 move |row: usize, v: ArrayView1<'_, f64>, out: &mut Array1<f64>| {
7036 let h = &transpose_slabs[row];
7037 for r in 0..h.nrows() {
7038 for c in 0..h.ncols() {
7039 out[c] += h[[r, c]] * v[r];
7040 }
7041 }
7042 },
7043 );
7044
7045 let matvec = gpu_schur_matvec_backend(&sys, 0.0, 0.0)
7046 .expect("row-procedural matvec backend builds for matrix-free system");
7047 let x = Array1::from_shape_fn(k, |i| ((i as f64 + 1.0) * 0.37).sin());
7048
7049 let mut out_parallel_a = Array1::<f64>::zeros(k);
7053 matvec(&x, &mut out_parallel_a);
7054 let mut out_parallel_b = Array1::<f64>::zeros(k);
7055 matvec(&x, &mut out_parallel_b);
7056 for a in 0..k {
7057 assert_eq!(
7058 out_parallel_a[a].to_bits(),
7059 out_parallel_b[a].to_bits(),
7060 "row-procedural matvec parallel reduction is non-deterministic at index {a}"
7061 );
7062 }
7063
7064 let mut out_serial = Array1::<f64>::zeros(k);
7069 rayon::ThreadPoolBuilder::new()
7070 .num_threads(2)
7071 .build()
7072 .expect("build rayon pool")
7073 .install(|| matvec(&x, &mut out_serial));
7074
7075 let max_abs = out_serial.iter().fold(0.0_f64, |m, v| m.max(v.abs()));
7076 for a in 0..k {
7077 let diff = (out_parallel_a[a] - out_serial[a]).abs();
7078 assert!(
7079 diff <= 1e-12 * (1.0 + max_abs),
7080 "row-procedural matvec parallel vs serial diverged beyond reassociation \
7081 at index {a}: {} vs {} (diff={diff:e})",
7082 out_parallel_a[a],
7083 out_serial[a]
7084 );
7085 }
7086 }
7087
7088 #[test]
7095 fn framed_sae_schur_matvec_matches_dense_reference() {
7096 use crate::arrow_schur::{
7097 BetaPenaltyOp, DeviceSaeFrameData, DeviceSaePcgData, DeviceSaeSmoothBlock,
7098 FactoredFrameGBlock, FactoredFrameKroneckerOp, IdentityRightKroneckerPenaltyOp,
7099 };
7100
7101 let p = 4usize;
7102 let ranks = vec![2usize, 4usize, 3usize];
7104 let basis_sizes = vec![2usize, 1usize, 2usize];
7105 let n_atoms = ranks.len();
7106 let mut border_offsets = Vec::with_capacity(n_atoms);
7107 let mut acc = 0usize;
7108 for k in 0..n_atoms {
7109 border_offsets.push(acc);
7110 acc += basis_sizes[k] * ranks[k];
7111 }
7112 let border_dim = acc; let mut state = 0x1234_5678_9abc_def0u64;
7115 let mut sample = || -> f64 {
7116 state = state
7117 .wrapping_mul(6364136223846793005)
7118 .wrapping_add(1442695040888963407);
7119 ((state >> 33) as f64) / ((1u64 << 31) as f64) - 1.0
7120 };
7121
7122 let mut frames: Vec<Array2<f64>> = Vec::with_capacity(n_atoms);
7125 for k in 0..n_atoms {
7126 let r = ranks[k];
7127 let mut u = Array2::<f64>::zeros((p, r));
7128 for i in 0..p {
7129 for j in 0..r {
7130 u[[i, j]] = if r == p && i == j {
7131 1.0
7132 } else if r == p {
7133 0.0
7134 } else {
7135 sample()
7136 };
7137 }
7138 }
7139 frames.push(u);
7140 }
7141 let w_of = |i: usize, j: usize| -> Array2<f64> {
7142 let (ui, uj) = (&frames[i], &frames[j]);
7143 let (ri, rj) = (ranks[i], ranks[j]);
7144 let mut w = Array2::<f64>::zeros((ri, rj));
7145 for a in 0..ri {
7146 for b in 0..rj {
7147 let mut s = 0.0;
7148 for c in 0..p {
7149 s += ui[[c, a]] * uj[[c, b]];
7150 }
7151 w[[a, b]] = s;
7152 }
7153 }
7154 w
7155 };
7156
7157 let mut frame_blocks: Vec<FactoredFrameGBlock> = Vec::new();
7159 let mut pairs = vec![(0usize, 0usize), (1, 1), (2, 2), (0, 2), (2, 0)];
7160 pairs.sort();
7161 for &(i, j) in &pairs {
7162 let (mi, mj) = (basis_sizes[i], basis_sizes[j]);
7163 let mut g = Array2::<f64>::zeros((mi, mj));
7164 for r in 0..mi {
7165 for c in 0..mj {
7166 g[[r, c]] = 0.3 * sample();
7167 }
7168 }
7169 if i == j {
7171 for r in 0..mi.min(mj) {
7172 g[[r, r]] += mi as f64 + 2.0;
7173 }
7174 }
7175 frame_blocks.push(FactoredFrameGBlock {
7176 atom_i: i,
7177 atom_j: j,
7178 g,
7179 w: w_of(i, j),
7180 });
7181 }
7182
7183 let mut smooth_blocks: Vec<DeviceSaeSmoothBlock> = Vec::with_capacity(n_atoms);
7185 let mut smooth_ranks: Vec<usize> = Vec::with_capacity(n_atoms);
7186 for k in 0..n_atoms {
7187 let m = basis_sizes[k];
7188 let mut a = Array2::<f64>::zeros((m, m));
7189 for r in 0..m {
7190 for c in 0..m {
7191 a[[r, c]] = 0.2 * sample();
7192 }
7193 }
7194 let mut s = a.t().dot(&a);
7195 for r in 0..m {
7196 s[[r, r]] += 1.0;
7197 }
7198 smooth_blocks.push(DeviceSaeSmoothBlock {
7199 global_offset: border_offsets[k],
7200 factor_a: s,
7201 });
7202 smooth_ranks.push(ranks[k]);
7203 }
7204
7205 let n = 6usize;
7207 let q = 3usize;
7208 let mut sys = ArrowSchurSystem::new(n, q, border_dim);
7209 let mut row_htbeta: Vec<Vec<f64>> = Vec::with_capacity(n);
7210 for i in 0..n {
7211 let mut a = Array2::<f64>::zeros((q, q));
7213 for r in 0..q {
7214 for c in 0..q {
7215 a[[r, c]] = sample();
7216 }
7217 }
7218 let mut htt = a.t().dot(&a);
7219 for r in 0..q {
7220 htt[[r, r]] += q as f64 + 1.0;
7221 }
7222 sys.rows[i].htt = htt;
7223 let mut slab = vec![0.0_f64; q * border_dim];
7224 for c in 0..q {
7225 for col in 0..border_dim {
7226 let v = 0.15 * sample();
7227 slab[c * border_dim + col] = v;
7228 sys.rows[i].htbeta[[c, col]] = v;
7229 }
7230 }
7231 row_htbeta.push(slab);
7232 }
7233
7234 let data_op =
7237 FactoredFrameKroneckerOp::new(ranks.clone(), basis_sizes.clone(), frame_blocks.clone())
7238 .expect("frame op");
7239 let mut hbb = data_op.to_dense();
7240 for k in 0..n_atoms {
7241 let op = IdentityRightKroneckerPenaltyOp {
7242 factor_a: smooth_blocks[k].factor_a.clone(),
7243 p: ranks[k],
7244 global_offset: border_offsets[k],
7245 k: border_dim,
7246 };
7247 let d = op.to_dense();
7248 for r in 0..border_dim {
7249 for c in 0..border_dim {
7250 hbb[[r, c]] += d[[r, c]];
7251 }
7252 }
7253 }
7254 sys.hbb = hbb;
7255
7256 let data = DeviceSaePcgData {
7257 p,
7258 beta_dim: border_dim,
7259 a_phi: std::sync::Arc::from(Vec::new().into_boxed_slice()),
7260 local_jac: std::sync::Arc::from(Vec::new().into_boxed_slice()),
7261 smooth_blocks,
7262 sparse_g_blocks: Vec::new(),
7263 frame: Some(DeviceSaeFrameData {
7264 ranks: ranks.clone(),
7265 basis_sizes: basis_sizes.clone(),
7266 border_offsets: border_offsets.clone(),
7267 frame_blocks,
7268 smooth_ranks,
7269 row_htbeta,
7270 }),
7271 };
7272
7273 let ridge_t = 1e-7;
7274 let ridge_beta = 1e-6;
7275
7276 let mut s_dense = Array2::<f64>::zeros((border_dim, border_dim));
7280 for r in 0..border_dim {
7281 for c in 0..border_dim {
7282 s_dense[[r, c]] = sys.hbb[[r, c]];
7283 }
7284 s_dense[[r, r]] += ridge_beta;
7285 }
7286 for row in &sys.rows {
7287 let mut htt = row.htt.clone();
7288 for d in 0..q {
7289 htt[[d, d]] += ridge_t;
7290 }
7291 let factor = cholesky_factor_in_place(htt.view(), CholeskyGuard::NonnegativePivot)
7292 .expect("htt PD");
7293 let mut y = Array2::<f64>::zeros((q, border_dim));
7295 for col in 0..border_dim {
7296 let mut e = Array1::<f64>::zeros(q);
7297 for r in 0..q {
7298 e[r] = row.htbeta[[r, col]];
7299 }
7300 let solved = cholesky_solve_vector(factor.view(), e.view());
7301 for r in 0..q {
7302 y[[r, col]] = solved[r];
7303 }
7304 }
7305 for r in 0..border_dim {
7306 for c in 0..border_dim {
7307 let mut acc = 0.0;
7308 for d in 0..q {
7309 acc += row.htbeta[[d, r]] * y[[d, c]];
7310 }
7311 s_dense[[r, c]] -= acc;
7312 }
7313 }
7314 }
7315
7316 let mut max_rel = 0.0_f64;
7318 for trial in 0..4 {
7319 let x: Vec<f64> = (0..border_dim)
7320 .map(|a| 0.3 * ((a as f64 + trial as f64) * 0.21).cos() - 0.1)
7321 .collect();
7322 let mut got = vec![0.0_f64; border_dim];
7323 sae_framed_schur_matvec_cpu(&sys, &data, ridge_t, ridge_beta, &x, &mut got)
7324 .expect("framed matvec");
7325 let mut want = vec![0.0_f64; border_dim];
7326 for r in 0..border_dim {
7327 let mut acc = 0.0;
7328 for c in 0..border_dim {
7329 acc += s_dense[[r, c]] * x[c];
7330 }
7331 want[r] = acc;
7332 }
7333 let scale = want.iter().fold(0.0_f64, |m, v| m.max(v.abs())).max(1.0);
7334 for a in 0..border_dim {
7335 let rel = (got[a] - want[a]).abs() / scale;
7336 max_rel = max_rel.max(rel);
7337 }
7338 }
7339 assert!(
7340 max_rel <= 1e-10,
7341 "framed SAE Schur matvec vs dense reference diverged: max_rel={max_rel:e}"
7342 );
7343 }
7344
7345 #[test]
7355 fn framed_sae_schur_matvec_matches_dense_reference_large_k_1026() {
7356 use crate::arrow_schur::{
7357 BetaPenaltyOp, DeviceSaeFrameData, DeviceSaePcgData, DeviceSaeSmoothBlock,
7358 FactoredFrameGBlock, FactoredFrameKroneckerOp, IdentityRightKroneckerPenaltyOp,
7359 };
7360
7361 let p = 12usize;
7362 let n_atoms = 40usize;
7363 let ranks: Vec<usize> = (0..n_atoms)
7366 .map(|k| if k % 5 == 0 { p } else { 2 + (k % 3) })
7367 .collect();
7368 let basis_sizes: Vec<usize> = (0..n_atoms).map(|k| 1 + (k % 3)).collect();
7369 let mut border_offsets = Vec::with_capacity(n_atoms);
7370 let mut acc = 0usize;
7371 for k in 0..n_atoms {
7372 border_offsets.push(acc);
7373 acc += basis_sizes[k] * ranks[k];
7374 }
7375 let border_dim = acc;
7376
7377 let mut state = 0x0bad_c0de_dead_beefu64;
7378 let mut sample = || -> f64 {
7379 state = state
7380 .wrapping_mul(6364136223846793005)
7381 .wrapping_add(1442695040888963407);
7382 ((state >> 33) as f64) / ((1u64 << 31) as f64) - 1.0
7383 };
7384
7385 let mut frames: Vec<Array2<f64>> = Vec::with_capacity(n_atoms);
7387 for k in 0..n_atoms {
7388 let r = ranks[k];
7389 let mut u = Array2::<f64>::zeros((p, r));
7390 for i in 0..p {
7391 for j in 0..r {
7392 u[[i, j]] = if r == p {
7393 if i == j { 1.0 } else { 0.0 }
7394 } else {
7395 sample()
7396 };
7397 }
7398 }
7399 frames.push(u);
7400 }
7401 let w_of = |i: usize, j: usize| -> Array2<f64> {
7402 let (ui, uj) = (&frames[i], &frames[j]);
7403 let (ri, rj) = (ranks[i], ranks[j]);
7404 let mut w = Array2::<f64>::zeros((ri, rj));
7405 for a in 0..ri {
7406 for b in 0..rj {
7407 let mut s = 0.0;
7408 for c in 0..p {
7409 s += ui[[c, a]] * uj[[c, b]];
7410 }
7411 w[[a, b]] = s;
7412 }
7413 }
7414 w
7415 };
7416
7417 let mut pairs: Vec<(usize, usize)> = Vec::new();
7420 for k in 0..n_atoms {
7421 pairs.push((k, k));
7422 }
7423 for k in 0..n_atoms - 1 {
7424 pairs.push((k, k + 1));
7425 pairs.push((k + 1, k));
7426 }
7427 pairs.sort_unstable();
7428 let mut frame_blocks: Vec<FactoredFrameGBlock> = Vec::new();
7429 for &(i, j) in &pairs {
7430 let (mi, mj) = (basis_sizes[i], basis_sizes[j]);
7431 let mut g = Array2::<f64>::zeros((mi, mj));
7432 for r in 0..mi {
7433 for c in 0..mj {
7434 g[[r, c]] = 0.3 * sample();
7435 }
7436 }
7437 if i == j {
7438 for r in 0..mi.min(mj) {
7439 g[[r, r]] += mi as f64 + 2.0;
7440 }
7441 }
7442 frame_blocks.push(FactoredFrameGBlock {
7443 atom_i: i,
7444 atom_j: j,
7445 g,
7446 w: w_of(i, j),
7447 });
7448 }
7449
7450 let mut smooth_blocks: Vec<DeviceSaeSmoothBlock> = Vec::with_capacity(n_atoms);
7452 let mut smooth_ranks: Vec<usize> = Vec::with_capacity(n_atoms);
7453 for k in 0..n_atoms {
7454 let m = basis_sizes[k];
7455 let mut a = Array2::<f64>::zeros((m, m));
7456 for r in 0..m {
7457 for c in 0..m {
7458 a[[r, c]] = 0.2 * sample();
7459 }
7460 }
7461 let mut s = a.t().dot(&a);
7462 for r in 0..m {
7463 s[[r, r]] += 1.0;
7464 }
7465 smooth_blocks.push(DeviceSaeSmoothBlock {
7466 global_offset: border_offsets[k],
7467 factor_a: s,
7468 });
7469 smooth_ranks.push(ranks[k]);
7470 }
7471
7472 let n = 8usize;
7474 let q = 3usize;
7475 let mut sys = ArrowSchurSystem::new(n, q, border_dim);
7476 let mut row_htbeta: Vec<Vec<f64>> = Vec::with_capacity(n);
7477 for i in 0..n {
7478 let mut a = Array2::<f64>::zeros((q, q));
7479 for r in 0..q {
7480 for c in 0..q {
7481 a[[r, c]] = sample();
7482 }
7483 }
7484 let mut htt = a.t().dot(&a);
7485 for r in 0..q {
7486 htt[[r, r]] += q as f64 + 1.0;
7487 }
7488 sys.rows[i].htt = htt;
7489 let mut slab = vec![0.0_f64; q * border_dim];
7490 for c in 0..q {
7491 for col in 0..border_dim {
7492 let v = 0.15 * sample();
7493 slab[c * border_dim + col] = v;
7494 sys.rows[i].htbeta[[c, col]] = v;
7495 }
7496 }
7497 row_htbeta.push(slab);
7498 }
7499
7500 let data_op =
7503 FactoredFrameKroneckerOp::new(ranks.clone(), basis_sizes.clone(), frame_blocks.clone())
7504 .expect("frame op");
7505 let mut hbb = data_op.to_dense();
7506 for k in 0..n_atoms {
7507 let op = IdentityRightKroneckerPenaltyOp {
7508 factor_a: smooth_blocks[k].factor_a.clone(),
7509 p: ranks[k],
7510 global_offset: border_offsets[k],
7511 k: border_dim,
7512 };
7513 let d = op.to_dense();
7514 for r in 0..border_dim {
7515 for c in 0..border_dim {
7516 hbb[[r, c]] += d[[r, c]];
7517 }
7518 }
7519 }
7520 sys.hbb = hbb;
7521
7522 let data = DeviceSaePcgData {
7523 p,
7524 beta_dim: border_dim,
7525 a_phi: std::sync::Arc::from(Vec::new().into_boxed_slice()),
7526 local_jac: std::sync::Arc::from(Vec::new().into_boxed_slice()),
7527 smooth_blocks,
7528 sparse_g_blocks: Vec::new(),
7529 frame: Some(DeviceSaeFrameData {
7530 ranks: ranks.clone(),
7531 basis_sizes: basis_sizes.clone(),
7532 border_offsets: border_offsets.clone(),
7533 frame_blocks,
7534 smooth_ranks,
7535 row_htbeta,
7536 }),
7537 };
7538
7539 let ridge_t = 1e-7;
7540 let ridge_beta = 1e-6;
7541
7542 let mut s_dense = sys.hbb.clone();
7544 for r in 0..border_dim {
7545 s_dense[[r, r]] += ridge_beta;
7546 }
7547 for row in &sys.rows {
7548 let mut htt = row.htt.clone();
7549 for d in 0..q {
7550 htt[[d, d]] += ridge_t;
7551 }
7552 let factor = cholesky_factor_in_place(htt.view(), CholeskyGuard::NonnegativePivot)
7553 .expect("htt PD");
7554 let mut y = Array2::<f64>::zeros((q, border_dim));
7555 for col in 0..border_dim {
7556 let mut e = Array1::<f64>::zeros(q);
7557 for r in 0..q {
7558 e[r] = row.htbeta[[r, col]];
7559 }
7560 let solved = cholesky_solve_vector(factor.view(), e.view());
7561 for r in 0..q {
7562 y[[r, col]] = solved[r];
7563 }
7564 }
7565 for r in 0..border_dim {
7566 for c in 0..border_dim {
7567 let mut acc = 0.0;
7568 for d in 0..q {
7569 acc += row.htbeta[[d, r]] * y[[d, c]];
7570 }
7571 s_dense[[r, c]] -= acc;
7572 }
7573 }
7574 }
7575
7576 let mut max_rel = 0.0_f64;
7577 for trial in 0..4 {
7578 let x: Vec<f64> = (0..border_dim)
7579 .map(|a| 0.3 * ((a as f64 + trial as f64) * 0.21).cos() - 0.1)
7580 .collect();
7581 let mut got = vec![0.0_f64; border_dim];
7582 sae_framed_schur_matvec_cpu(&sys, &data, ridge_t, ridge_beta, &x, &mut got)
7583 .expect("framed matvec");
7584 let mut want = vec![0.0_f64; border_dim];
7585 for r in 0..border_dim {
7586 let mut acc = 0.0;
7587 for c in 0..border_dim {
7588 acc += s_dense[[r, c]] * x[c];
7589 }
7590 want[r] = acc;
7591 }
7592 let scale = want.iter().fold(0.0_f64, |m, v| m.max(v.abs())).max(1.0);
7593 for a in 0..border_dim {
7594 let rel = (got[a] - want[a]).abs() / scale;
7595 max_rel = max_rel.max(rel);
7596 }
7597 }
7598 assert!(
7599 max_rel <= 1e-10,
7600 "large-K framed SAE Schur matvec vs dense reference diverged: \
7601 max_rel={max_rel:e} (n_atoms={n_atoms}, border_dim={border_dim})"
7602 );
7603 }
7604
7605 #[test]
7611 fn framed_sae_device_pcg_matches_cpu_when_cuda_admits() {
7612 use crate::arrow_schur::{
7613 BetaPenaltyOp, DeviceSaeFrameData, DeviceSaePcgData, DeviceSaeSmoothBlock,
7614 FactoredFrameGBlock, FactoredFrameKroneckerOp, IdentityRightKroneckerPenaltyOp,
7615 };
7616
7617 let p = 6usize;
7621 let n_atoms = 8usize;
7622 let ranks: Vec<usize> = (0..n_atoms)
7623 .map(|k| if k % 2 == 0 { 3usize } else { p })
7624 .collect();
7625 let basis_sizes: Vec<usize> = (0..n_atoms).map(|_| 3usize).collect();
7626 let mut border_offsets = Vec::with_capacity(n_atoms);
7627 let mut acc = 0usize;
7628 for k in 0..n_atoms {
7629 border_offsets.push(acc);
7630 acc += basis_sizes[k] * ranks[k];
7631 }
7632 let border_dim = acc; let mut state = 0xfeed_face_dead_beefu64;
7635 let mut sample = || -> f64 {
7636 state = state
7637 .wrapping_mul(6364136223846793005)
7638 .wrapping_add(1442695040888963407);
7639 ((state >> 33) as f64) / ((1u64 << 31) as f64) - 1.0
7640 };
7641 let mut frames: Vec<Array2<f64>> = Vec::new();
7642 for k in 0..n_atoms {
7643 let r = ranks[k];
7644 let mut u = Array2::<f64>::zeros((p, r));
7645 for i in 0..p {
7646 for j in 0..r {
7647 u[[i, j]] = if r == p && i == j {
7648 1.0
7649 } else if r == p {
7650 0.0
7651 } else {
7652 sample()
7653 };
7654 }
7655 }
7656 frames.push(u);
7657 }
7658 let w_of = |i: usize, j: usize| {
7659 let (ui, uj) = (&frames[i], &frames[j]);
7660 let (ri, rj) = (ranks[i], ranks[j]);
7661 let mut w = Array2::<f64>::zeros((ri, rj));
7662 for a in 0..ri {
7663 for b in 0..rj {
7664 let mut s = 0.0;
7665 for c in 0..p {
7666 s += ui[[c, a]] * uj[[c, b]];
7667 }
7668 w[[a, b]] = s;
7669 }
7670 }
7671 w
7672 };
7673 let mut pairs: Vec<(usize, usize)> = (0..n_atoms).map(|k| (k, k)).collect();
7674 for &(i, j) in &[(0usize, 1usize), (2, 4), (3, 6)] {
7676 pairs.push((i, j));
7677 pairs.push((j, i));
7678 }
7679 let mut frame_blocks = Vec::new();
7680 for &(i, j) in &pairs {
7681 let (mi, mj) = (basis_sizes[i], basis_sizes[j]);
7682 let mut g = Array2::<f64>::zeros((mi, mj));
7683 for r in 0..mi {
7684 for c in 0..mj {
7685 g[[r, c]] = 0.25 * sample();
7686 }
7687 }
7688 if i == j {
7689 for r in 0..mi.min(mj) {
7690 g[[r, r]] += mi as f64 + 2.0;
7691 }
7692 }
7693 frame_blocks.push(FactoredFrameGBlock {
7694 atom_i: i,
7695 atom_j: j,
7696 g,
7697 w: w_of(i, j),
7698 });
7699 }
7700 let mut smooth_blocks = Vec::new();
7701 let mut smooth_ranks = Vec::new();
7702 for k in 0..n_atoms {
7703 let m = basis_sizes[k];
7704 let mut a = Array2::<f64>::zeros((m, m));
7705 for r in 0..m {
7706 for c in 0..m {
7707 a[[r, c]] = 0.2 * sample();
7708 }
7709 }
7710 let mut s = a.t().dot(&a);
7711 for r in 0..m {
7712 s[[r, r]] += 1.0;
7713 }
7714 smooth_blocks.push(DeviceSaeSmoothBlock {
7715 global_offset: border_offsets[k],
7716 factor_a: s,
7717 });
7718 smooth_ranks.push(ranks[k]);
7719 }
7720 let n = 400usize;
7721 let q = 4usize;
7722 let mut sys = ArrowSchurSystem::new(n, q, border_dim);
7723 let mut row_htbeta = Vec::new();
7724 for i in 0..n {
7725 let mut a = Array2::<f64>::zeros((q, q));
7726 for r in 0..q {
7727 for c in 0..q {
7728 a[[r, c]] = sample();
7729 }
7730 }
7731 let mut htt = a.t().dot(&a);
7732 for r in 0..q {
7733 htt[[r, r]] += q as f64 + 1.0;
7734 }
7735 sys.rows[i].htt = htt;
7736 let mut slab = vec![0.0_f64; q * border_dim];
7737 for c in 0..q {
7738 for col in 0..border_dim {
7739 let v = 0.02 * sample();
7742 slab[c * border_dim + col] = v;
7743 sys.rows[i].htbeta[[c, col]] = v;
7744 }
7745 }
7746 row_htbeta.push(slab);
7747 }
7748 let data_op =
7749 FactoredFrameKroneckerOp::new(ranks.clone(), basis_sizes.clone(), frame_blocks.clone())
7750 .expect("frame op");
7751 let mut hbb = data_op.to_dense();
7752 for k in 0..n_atoms {
7753 let op = IdentityRightKroneckerPenaltyOp {
7754 factor_a: smooth_blocks[k].factor_a.clone(),
7755 p: ranks[k],
7756 global_offset: border_offsets[k],
7757 k: border_dim,
7758 };
7759 let d = op.to_dense();
7760 for r in 0..border_dim {
7761 for c in 0..border_dim {
7762 hbb[[r, c]] += d[[r, c]];
7763 }
7764 }
7765 }
7766 sys.hbb = hbb;
7767 let data = DeviceSaePcgData {
7768 p,
7769 beta_dim: border_dim,
7770 a_phi: std::sync::Arc::from(Vec::new().into_boxed_slice()),
7771 local_jac: std::sync::Arc::from(Vec::new().into_boxed_slice()),
7772 smooth_blocks,
7773 sparse_g_blocks: Vec::new(),
7774 frame: Some(DeviceSaeFrameData {
7775 ranks: ranks.clone(),
7776 basis_sizes: basis_sizes.clone(),
7777 border_offsets: border_offsets.clone(),
7778 frame_blocks,
7779 smooth_ranks,
7780 row_htbeta,
7781 }),
7782 };
7783 let ridge_t = 1e-7;
7784 let ridge_beta = 1e-6;
7785 let rhs: Array1<f64> =
7786 Array1::from_shape_fn(border_dim, |a| ((a as f64 + 1.0) * 0.17).sin());
7787
7788 let (device, diag) =
7789 match solve_sae_matrix_free_pcg(&sys, &data, ridge_t, ridge_beta, &rhs, 400, 1e-12) {
7790 Ok(result) => result,
7791 Err(failure) => {
7797 assert!(
7798 gam_gpu::device_runtime::GpuRuntime::resolve(gam_gpu::GpuPolicy::Auto)
7799 .unwrap_or_else(|error| {
7800 panic!("GPU probe fault in framed PCG test: {error}")
7801 })
7802 .is_none(),
7803 "#1017: CUDA device present but the framed device SAE PCG \
7804 declined/faulted instead of returning a result (tag: {failure:?}) — \
7805 the kernel does not run correctly on GPU"
7806 );
7807 return;
7808 }
7809 };
7810
7811 let rhs_norm = rhs.iter().map(|v| v * v).sum::<f64>().sqrt();
7829 let oracle_resid = |x: &Array1<f64>| -> f64 {
7830 let mut sx = vec![0.0_f64; border_dim];
7831 sae_framed_schur_matvec_cpu(
7832 &sys,
7833 &data,
7834 ridge_t,
7835 ridge_beta,
7836 x.as_slice().unwrap(),
7837 &mut sx,
7838 )
7839 .expect("cpu oracle matvec");
7840 let mut acc = 0.0_f64;
7841 for a in 0..border_dim {
7842 let e = sx[a] - rhs[a];
7843 acc += e * e;
7844 }
7845 acc.sqrt()
7846 };
7847 let s_dev_resid = oracle_resid(&device);
7848 let dev_rel_resid = s_dev_resid / rhs_norm.max(1e-300);
7849
7850 let precond = {
7855 let d = sae_frame_penalty_diag_host_for_test(&data, ridge_beta);
7856 let mut diag = d;
7859 for (i, row) in sys.rows.iter().enumerate() {
7860 let slab = &data.frame.as_ref().unwrap().row_htbeta[i];
7861 let qi = sys.row_dims[i];
7862 if slab.is_empty() || qi == 0 || slab.len() != qi * border_dim {
7863 continue;
7864 }
7865 let mut block = row.htt.clone();
7866 for dd in 0..qi {
7867 block[[dd, dd]] += ridge_t;
7868 }
7869 let factor =
7870 cholesky_factor_in_place(block.view(), CholeskyGuard::NonnegativePivot)
7871 .expect("row htt PD");
7872 let mut ainv = Array2::<f64>::zeros((qi, qi));
7874 for col in 0..qi {
7875 let mut e = Array1::<f64>::zeros(qi);
7876 e[col] = 1.0;
7877 let s = cholesky_solve_vector(factor.view(), e.view());
7878 for r in 0..qi {
7879 ainv[[r, col]] = s[r];
7880 }
7881 }
7882 for a in 0..border_dim {
7883 let mut quad = 0.0_f64;
7884 for c in 0..qi {
7885 let hc = slab[c * border_dim + a];
7886 for dd in 0..qi {
7887 quad += hc * ainv[[c, dd]] * slab[dd * border_dim + a];
7888 }
7889 }
7890 diag[a] -= quad;
7891 }
7892 }
7893 Array1::from_vec(diag)
7894 };
7895 let mut cpu = Array1::<f64>::zeros(border_dim);
7896 let cpu_result = {
7897 let mut apply = |v: &Array1<f64>, out: &mut Array1<f64>| {
7898 let mut tmp = vec![0.0_f64; border_dim];
7899 sae_framed_schur_matvec_cpu(
7900 &sys,
7901 &data,
7902 ridge_t,
7903 ridge_beta,
7904 v.as_slice().unwrap(),
7905 &mut tmp,
7906 )
7907 .expect("cpu oracle matvec");
7908 out.assign(&Array1::from_vec(tmp));
7909 };
7910 gam_linalg::pcg::pcg_core(
7911 &mut apply,
7912 &rhs.view(),
7913 &precond.view(),
7914 1e-12,
7915 800,
7916 32,
7917 false,
7918 gam_linalg::pcg::DotReduction::Serial,
7919 &mut cpu.view_mut(),
7920 )
7921 };
7922 let s_cpu_resid = oracle_resid(&cpu);
7923 let cpu_rel_resid = s_cpu_resid / rhs_norm.max(1e-300);
7924
7925 assert!(
7928 dev_rel_resid <= 1e-7,
7929 "[#1551] device δβ does not solve the CPU-oracle system: \
7930 ‖S_cpu·device−rhs‖/‖rhs‖={dev_rel_resid:e} (>1e-7) | abs={s_dev_resid:e} | \
7931 device PCG stop={:?} iters={} final_rel_resid={:e} — a large operator residual \
7932 means the device matvec is a DIFFERENT operator (kernel bug)",
7933 diag.stopping_reason,
7934 diag.iterations,
7935 diag.final_relative_residual,
7936 );
7937 assert!(
7940 cpu_rel_resid <= 1e-6,
7941 "[#1551] CPU pcg_core failed to solve the oracle system: \
7942 ‖S_cpu·cpu−rhs‖/‖rhs‖={cpu_rel_resid:e} (stop={:?}, iters={}) — fixture/oracle issue",
7943 cpu_result.stop,
7944 cpu_result.iterations,
7945 );
7946 }
7947
7948 fn sae_frame_penalty_diag_host_for_test(data: &DeviceSaePcgData, ridge_beta: f64) -> Vec<f64> {
7952 let frame = data.frame.as_ref().expect("frame");
7953 let mut diag = vec![ridge_beta; data.beta_dim];
7954 for (blk, &r) in data.smooth_blocks.iter().zip(frame.smooth_ranks.iter()) {
7955 let m = blk.factor_a.nrows();
7956 for ia in 0..m {
7957 let coeff = blk.factor_a[[ia, ia]];
7958 let base = blk.global_offset + ia * r;
7959 for ib in 0..r {
7960 diag[base + ib] += coeff;
7961 }
7962 }
7963 }
7964 for blk in &frame.frame_blocks {
7965 if blk.atom_i != blk.atom_j {
7966 continue;
7967 }
7968 let r = frame.ranks[blk.atom_i];
7969 let off = frame.border_offsets[blk.atom_i];
7970 let (mi, mj) = blk.g.dim();
7971 for li in 0..mi.min(mj) {
7972 let gii = blk.g[[li, li]];
7973 let base = off + li * r;
7974 for a in 0..r {
7975 diag[base + a] += gii * blk.w[[a, a]];
7976 }
7977 }
7978 }
7979 diag
7980 }
7981
7982 #[test]
7993 #[cfg(target_os = "linux")]
7994 fn framed_sae_device_matvec_matches_cpu_oracle_when_cuda_admits() {
7995 use crate::arrow_schur::{
7996 DeviceSaeFrameData, DeviceSaePcgData, DeviceSaeSmoothBlock, FactoredFrameGBlock,
7997 };
7998
7999 let p = 6usize;
8002 let n_atoms = 8usize;
8003 let ranks: Vec<usize> = (0..n_atoms)
8004 .map(|k| if k % 2 == 0 { 3usize } else { p })
8005 .collect();
8006 let basis_sizes: Vec<usize> = (0..n_atoms).map(|_| 3usize).collect();
8007 let mut border_offsets = Vec::with_capacity(n_atoms);
8008 let mut acc = 0usize;
8009 for k in 0..n_atoms {
8010 border_offsets.push(acc);
8011 acc += basis_sizes[k] * ranks[k];
8012 }
8013 let border_dim = acc;
8014
8015 let mut state = 0x1551_0017_1026_0922u64;
8016 let mut sample = || -> f64 {
8017 state = state
8018 .wrapping_mul(6364136223846793005)
8019 .wrapping_add(1442695040888963407);
8020 ((state >> 33) as f64) / ((1u64 << 31) as f64) - 1.0
8021 };
8022 let mut frames: Vec<Array2<f64>> = Vec::new();
8023 for k in 0..n_atoms {
8024 let r = ranks[k];
8025 let mut u = Array2::<f64>::zeros((p, r));
8026 for i in 0..p {
8027 for j in 0..r {
8028 u[[i, j]] = if r == p && i == j {
8029 1.0
8030 } else if r == p {
8031 0.0
8032 } else {
8033 sample()
8034 };
8035 }
8036 }
8037 frames.push(u);
8038 }
8039 let w_of = |i: usize, j: usize| {
8040 let (ui, uj) = (&frames[i], &frames[j]);
8041 let (ri, rj) = (ranks[i], ranks[j]);
8042 let mut w = Array2::<f64>::zeros((ri, rj));
8043 for a in 0..ri {
8044 for b in 0..rj {
8045 let mut s = 0.0;
8046 for c in 0..p {
8047 s += ui[[c, a]] * uj[[c, b]];
8048 }
8049 w[[a, b]] = s;
8050 }
8051 }
8052 w
8053 };
8054 let mut pairs: Vec<(usize, usize)> = (0..n_atoms).map(|k| (k, k)).collect();
8055 for &(i, j) in &[(0usize, 1usize), (2, 4), (3, 6)] {
8056 pairs.push((i, j));
8057 pairs.push((j, i));
8058 }
8059 let mut frame_blocks = Vec::new();
8060 for &(i, j) in &pairs {
8061 let (mi, mj) = (basis_sizes[i], basis_sizes[j]);
8062 let mut g = Array2::<f64>::zeros((mi, mj));
8063 for r in 0..mi {
8064 for c in 0..mj {
8065 g[[r, c]] = 0.25 * sample();
8066 }
8067 }
8068 if i == j {
8069 for r in 0..mi.min(mj) {
8070 g[[r, r]] += mi as f64 + 2.0;
8071 }
8072 }
8073 frame_blocks.push(FactoredFrameGBlock {
8074 atom_i: i,
8075 atom_j: j,
8076 g,
8077 w: w_of(i, j),
8078 });
8079 }
8080 let mut smooth_blocks = Vec::new();
8081 let mut smooth_ranks = Vec::new();
8082 for k in 0..n_atoms {
8083 let m = basis_sizes[k];
8084 let mut a = Array2::<f64>::zeros((m, m));
8085 for r in 0..m {
8086 for c in 0..m {
8087 a[[r, c]] = 0.2 * sample();
8088 }
8089 }
8090 let mut s = a.t().dot(&a);
8091 for r in 0..m {
8092 s[[r, r]] += 1.0;
8093 }
8094 smooth_blocks.push(DeviceSaeSmoothBlock {
8095 global_offset: border_offsets[k],
8096 factor_a: s,
8097 });
8098 smooth_ranks.push(ranks[k]);
8099 }
8100 let n = 32usize;
8104 let q = 4usize;
8105 let mut sys = ArrowSchurSystem::new(n, q, border_dim);
8106 let mut row_htbeta = Vec::new();
8107 for i in 0..n {
8108 let mut a = Array2::<f64>::zeros((q, q));
8109 for r in 0..q {
8110 for c in 0..q {
8111 a[[r, c]] = sample();
8112 }
8113 }
8114 let mut htt = a.t().dot(&a);
8115 for r in 0..q {
8116 htt[[r, r]] += q as f64 + 1.0;
8117 }
8118 sys.rows[i].htt = htt;
8119 let mut slab = vec![0.0_f64; q * border_dim];
8120 for c in 0..q {
8121 for col in 0..border_dim {
8122 let v = 0.3 * sample();
8123 slab[c * border_dim + col] = v;
8124 sys.rows[i].htbeta[[c, col]] = v;
8125 }
8126 }
8127 row_htbeta.push(slab);
8128 }
8129 let ridge_t = 1e-7;
8130 let ridge_beta = 1e-6;
8131 let data = DeviceSaePcgData {
8132 p,
8133 beta_dim: border_dim,
8134 a_phi: std::sync::Arc::from(Vec::new().into_boxed_slice()),
8135 local_jac: std::sync::Arc::from(Vec::new().into_boxed_slice()),
8136 smooth_blocks,
8137 sparse_g_blocks: Vec::new(),
8138 frame: Some(DeviceSaeFrameData {
8139 ranks: ranks.clone(),
8140 basis_sizes: basis_sizes.clone(),
8141 border_offsets: border_offsets.clone(),
8142 frame_blocks,
8143 smooth_ranks,
8144 row_htbeta,
8145 }),
8146 };
8147
8148 let mut probes: Vec<Array1<f64>> = Vec::new();
8152 probes.push(Array1::from_shape_fn(border_dim, |a| {
8153 ((a as f64 + 1.0) * 0.37).sin()
8154 }));
8155 probes.push(Array1::from_shape_fn(border_dim, |_| sample()));
8156 for axis in [0usize, border_dim / 3, border_dim - 1] {
8157 let mut e = Array1::<f64>::zeros(border_dim);
8158 e[axis] = 1.0;
8159 probes.push(e);
8160 }
8161
8162 let mut any_ran = false;
8163 let mut worst = 0.0_f64;
8164 for (pi, x) in probes.iter().enumerate() {
8165 let device = match super::framed_schur_matvec_once_on_device(
8166 &sys, &data, ridge_t, ridge_beta, x,
8167 ) {
8168 Ok(out) => out,
8169 Err(failure) => {
8170 assert!(
8174 gam_gpu::device_runtime::GpuRuntime::resolve(gam_gpu::GpuPolicy::Auto)
8175 .unwrap_or_else(|error| {
8176 panic!("GPU probe fault in framed matvec parity test: {error}")
8177 })
8178 .is_none(),
8179 "#1551: CUDA device present but the framed device matvec \
8180 declined/faulted (probe {pi}, tag: {failure:?}) — the kernel \
8181 does not run on GPU"
8182 );
8183 return;
8184 }
8185 };
8186 any_ran = true;
8187 let mut cpu = vec![0.0_f64; border_dim];
8188 sae_framed_schur_matvec_cpu(
8189 &sys,
8190 &data,
8191 ridge_t,
8192 ridge_beta,
8193 x.as_slice().unwrap(),
8194 &mut cpu,
8195 )
8196 .expect("cpu oracle matvec");
8197 let scale = cpu.iter().fold(0.0_f64, |m, v| m.max(v.abs())).max(1.0);
8198 for a in 0..border_dim {
8199 let rel = (device[a] - cpu[a]).abs() / scale;
8200 worst = worst.max(rel);
8201 assert!(
8202 rel <= 1e-9,
8203 "[#1551 matvec-parity] probe {pi} component {a}: device={:e} cpu={:e} \
8204 rel={rel:e} (>1e-9) — framed S·x kernel diverges from the CPU oracle",
8205 device[a],
8206 cpu[a],
8207 );
8208 }
8209 }
8210 if any_ran {
8211 assert!(
8216 gam_gpu::device_runtime::GpuRuntime::resolve(gam_gpu::GpuPolicy::Auto)
8217 .unwrap_or_else(|error| {
8218 panic!("GPU probe fault after framed matvec execution: {error}")
8219 })
8220 .is_some(),
8221 "#1551: matvec ran but no GPU runtime — unexpected"
8222 );
8223 assert!(worst <= 1e-9, "framed matvec parity worst rel = {worst:e}");
8224 }
8225 }
8226}