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")]
373fn pack_host_d_and_stacked_b(sys: &ArrowSchurSystem, ridge_t: f64) -> (Vec<f64>, Vec<f64>) {
374 let n = sys.rows.len();
375 let d = sys.d;
376 let k = sys.k;
377 let rows = n * d;
378 let mut d_buf = Vec::with_capacity(n * d * d);
379 let mut b_buf = vec![0.0_f64; rows * k];
380 for (i, row) in sys.rows.iter().enumerate() {
381 for col in 0..d {
382 for r in 0..d {
383 let mut value = row.htt[[r, col]];
384 if r == col {
385 value += ridge_t;
386 }
387 d_buf.push(value);
388 }
389 }
390 for col in 0..k {
391 let base = col * rows + i * d;
392 for r in 0..d {
393 b_buf[base + r] = row.htbeta[[r, col]];
394 }
395 }
396 }
397 (d_buf, b_buf)
398}
399
400#[cfg(target_os = "linux")]
401#[inline]
402fn pack_block(
403 row: &crate::arrow_schur::ArrowRowBlock,
404 ridge_t: f64,
405 d: usize,
406 k: usize,
407 d_buf: &mut Vec<f64>,
408 b_buf: &mut Vec<f64>,
409 g_buf: &mut Vec<f64>,
410) {
411 for col in 0..d {
412 for r in 0..d {
413 let mut value = row.htt[[r, col]];
414 if r == col {
415 value += ridge_t;
416 }
417 d_buf.push(value);
418 }
419 }
420 for col in 0..k {
421 for r in 0..d {
422 b_buf.push(row.htbeta[[r, col]]);
423 }
424 }
425 for r in 0..d {
426 g_buf.push(row.gt[r]);
427 }
428}
429
430#[doc(hidden)]
436#[cfg(target_os = "linux")]
437pub fn solve_arrow_newton_step_fused_force(
438 sys: &ArrowSchurSystem,
439 ridge_t: f64,
440 ridge_beta: f64,
441) -> Result<ArrowSchurGpuSolution, ArrowSchurGpuFailure> {
442 if ridge_t.is_nan() || ridge_beta.is_nan() {
443 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
444 reason: "ridge is NaN".to_string(),
445 });
446 }
447 if crate::gpu_kernels::arrow_schur_nvrtc::plan_fused_launch(sys.rows.len(), sys.d, sys.k)
448 .is_none()
449 {
450 return Err(ArrowSchurGpuFailure::Unavailable);
451 }
452 cuda::solve_fused(sys, ridge_t, ridge_beta)
453}
454
455#[cfg(target_os = "linux")]
465pub struct ResidentArrowFrameHandle {
466 inner: cuda::ResidentArrowFrame,
467}
468
469#[cfg(not(target_os = "linux"))]
473pub enum ResidentArrowFrameHandle {}
474
475impl ResidentArrowFrameHandle {
476 pub fn new(
478 sys: &ArrowSchurSystem,
479 ridge_t: f64,
480 ridge_beta: f64,
481 ) -> Result<Self, ArrowSchurGpuFailure> {
482 if sys.hbb_matvec.is_some() || sys.htbeta_matvec.is_some() {
485 return Err(ArrowSchurGpuFailure::GpuRequiresDenseSystem {
486 had_hbb_matvec: sys.hbb_matvec.is_some(),
487 had_htbeta_matvec: sys.htbeta_matvec.is_some(),
488 });
489 }
490 #[cfg(not(target_os = "linux"))]
491 {
492 if ridge_t.is_nan() || ridge_beta.is_nan() {
493 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
494 reason: "ridge is NaN".to_string(),
495 });
496 }
497 Err(ArrowSchurGpuFailure::Unavailable)
498 }
499 #[cfg(target_os = "linux")]
500 {
501 Ok(Self {
502 inner: cuda::ResidentArrowFrame::new(sys, ridge_t, ridge_beta)?,
503 })
504 }
505 }
506
507 pub fn solve_gradient(
509 &self,
510 g_t: &[f64],
511 g_beta: &[f64],
512 ) -> Result<ArrowSchurGpuSolution, ArrowSchurGpuFailure> {
513 #[cfg(not(target_os = "linux"))]
514 {
515 if g_t.iter().chain(g_beta).any(|v| !v.is_finite()) {
516 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
517 reason: "non-finite gradient entry".to_string(),
518 });
519 }
520 Err(ArrowSchurGpuFailure::Unavailable)
521 }
522 #[cfg(target_os = "linux")]
523 {
524 self.inner.solve_gradient(g_t, g_beta)
525 }
526 }
527
528 #[must_use]
530 pub fn log_det_hessian(&self) -> f64 {
531 #[cfg(not(target_os = "linux"))]
532 {
533 match *self {}
534 }
535 #[cfg(target_os = "linux")]
536 {
537 self.inner.log_det_hessian()
538 }
539 }
540}
541
542#[cfg(target_os = "linux")]
557pub struct ResidentBaseArrowFrameHandle {
558 inner: cuda::ResidentBaseArrowFrame,
559}
560
561#[cfg(not(target_os = "linux"))]
564pub enum ResidentBaseArrowFrameHandle {}
565
566impl ResidentBaseArrowFrameHandle {
567 pub fn new(sys: &ArrowSchurSystem) -> Result<Self, ArrowSchurGpuFailure> {
573 if sys.hbb_matvec.is_some() || sys.htbeta_matvec.is_some() {
574 return Err(ArrowSchurGpuFailure::GpuRequiresDenseSystem {
575 had_hbb_matvec: sys.hbb_matvec.is_some(),
576 had_htbeta_matvec: sys.htbeta_matvec.is_some(),
577 });
578 }
579 #[cfg(not(target_os = "linux"))]
580 {
581 Err(ArrowSchurGpuFailure::Unavailable)
582 }
583 #[cfg(target_os = "linux")]
584 {
585 Ok(Self {
586 inner: cuda::ResidentBaseArrowFrame::new(sys)?,
587 })
588 }
589 }
590
591 pub fn refactor_and_solve(
598 &self,
599 ridge_t: f64,
600 ridge_beta: f64,
601 ) -> Result<ArrowSchurGpuSolution, ArrowSchurGpuFailure> {
602 #[cfg(not(target_os = "linux"))]
603 {
604 if ridge_t.is_nan() || ridge_beta.is_nan() {
605 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
606 reason: "ridge is NaN".to_string(),
607 });
608 }
609 Err(ArrowSchurGpuFailure::Unavailable)
610 }
611 #[cfg(target_os = "linux")]
612 {
613 self.inner.refactor_and_solve(ridge_t, ridge_beta)
614 }
615 }
616}
617
618pub fn gpu_schur_matvec_backend(
661 sys: &ArrowSchurSystem,
662 ridge_t: f64,
663 ridge_beta: f64,
664) -> Result<crate::arrow_schur::GpuSchurMatvec, ArrowSchurGpuFailure> {
665 if sys.htbeta_matvec.is_some() {
668 return build_row_procedural_matvec(sys, ridge_t, ridge_beta);
669 }
670
671 #[cfg(not(target_os = "linux"))]
672 {
673 if ridge_t.is_nan() || ridge_beta.is_nan() {
676 return Err(ArrowSchurGpuFailure::Unavailable);
677 }
678 Err(ArrowSchurGpuFailure::Unavailable)
679 }
680
681 #[cfg(target_os = "linux")]
682 {
683 cuda::build_schur_matvec_backend(sys, ridge_t, ridge_beta)
684 }
685}
686
687pub fn build_framed_resident_evidence_matvec(
695 sys: &ArrowSchurSystem,
696 ridge_t: f64,
697 ridge_beta: f64,
698 apply_budget: usize,
699) -> Result<Option<crate::arrow_schur::GpuSchurMatvec>, ArrowSchurGpuFailure> {
700 #[cfg(not(target_os = "linux"))]
701 {
702 if sys.k == 0 || !ridge_t.is_finite() || !ridge_beta.is_finite() || apply_budget == 0 {
706 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
707 reason: "resident evidence matvec received degenerate dimensions or ridge"
708 .to_string(),
709 });
710 }
711 Ok(None)
712 }
713 #[cfg(target_os = "linux")]
714 {
715 cuda::build_framed_resident_evidence_matvec(sys, ridge_t, ridge_beta, apply_budget)
716 }
717}
718
719fn build_row_procedural_matvec(
736 sys: &ArrowSchurSystem,
737 ridge_t: f64,
738 ridge_beta: f64,
739) -> Result<crate::arrow_schur::GpuSchurMatvec, ArrowSchurGpuFailure> {
740 use std::sync::Arc;
741 let n = sys.rows.len();
742 let k = sys.k;
743 let forward = sys
744 .htbeta_matvec
745 .clone()
746 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
747 let transpose = sys.htbeta_transpose_matvec.clone().ok_or_else(|| {
748 ArrowSchurGpuFailure::SchurFactorFailed {
753 reason: "row-procedural Schur matvec requires htbeta_transpose_matvec; \
754 forward operator installed without its sparse adjoint"
755 .to_string(),
756 }
757 })?;
758
759 let mut factors: Vec<Array2<f64>> = Vec::with_capacity(n);
764 for (i, row) in sys.rows.iter().enumerate() {
765 let di = row.htt.nrows();
766 if row.htt.ncols() != di || row.gt.len() != di {
767 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
768 reason: format!("row {i}: malformed H_tt block {:?}", row.htt.dim()),
769 });
770 }
771 let mut block = row.htt.clone();
772 for r in 0..di {
773 block[[r, r]] += ridge_t;
774 }
775 let factor = cholesky_factor_in_place(block.view(), CholeskyGuard::NonnegativePivot)
776 .ok_or_else(|| {
777 ArrowSchurGpuFailure::RidgeBumpRequired {
781 row: i,
782 bump: ridge_bump_to_make_pd(row.htt.view(), ridge_t),
783 }
784 })?;
785 factors.push(factor);
786 }
787
788 let penalty_op = sys.effective_penalty_op();
795 let row_dims: Vec<usize> = sys.rows.iter().map(|row| row.htt.nrows()).collect();
796
797 let closure: crate::arrow_schur::GpuSchurMatvec =
798 Arc::new(move |x: &Array1<f64>, out: &mut Array1<f64>| {
799 assert_eq!(x.len(), k, "row-procedural matvec: x.len() != k");
800 assert_eq!(out.len(), k, "row-procedural matvec: out.len() != k");
801
802 {
805 let x_slice = x.as_slice().expect("x must be contiguous");
806 let out_slice = out.as_slice_mut().expect("out must be contiguous");
807 for a in 0..k {
808 out_slice[a] = ridge_beta * x_slice[a];
809 }
810 penalty_op.matvec(x_slice, out_slice);
811 }
812
813 let parallel = n >= crate::arrow_schur::SCHUR_MATVEC_PARALLEL_ROW_MIN
840 && rayon::current_thread_index().is_none();
841 if parallel {
842 use rayon::prelude::*;
843 const CHUNK: usize = 64;
844 let partials: Vec<Array1<f64>> = (0..n)
845 .into_par_iter()
846 .chunks(CHUNK)
847 .map(|idxs| {
848 let mut neg = Array1::<f64>::zeros(k);
852 for i in idxs {
853 let di = row_dims[i];
854 let mut v_i = Array1::<f64>::zeros(di);
856 forward(i, x.view(), &mut v_i);
857 let w_i = cholesky_solve_vector(factors[i].view(), v_i.view());
859 transpose(i, w_i.view(), &mut neg);
861 }
862 neg
863 })
864 .collect();
865 let mut neg = Array1::<f64>::zeros(k);
877 for part in &partials {
878 for a in 0..k {
879 neg[a] += part[a];
880 }
881 }
882 for a in 0..k {
883 out[a] -= neg[a];
884 }
885 } else {
886 let mut neg = Array1::<f64>::zeros(k);
888 for i in 0..n {
889 let di = row_dims[i];
890 let mut v_i = Array1::<f64>::zeros(di);
892 forward(i, x.view(), &mut v_i);
893 let w_i = cholesky_solve_vector(factors[i].view(), v_i.view());
895 transpose(i, w_i.view(), &mut neg);
897 }
898 for a in 0..k {
899 out[a] -= neg[a];
900 }
901 }
902 });
903
904 Ok(closure)
905}
906
907pub fn solve_reduced_beta_pcg(
929 s_acc: &Array2<f64>,
930 rhs_beta: &Array1<f64>,
931 max_iterations: usize,
932 relative_tolerance: f64,
933) -> Result<Array1<f64>, ArrowSchurGpuFailure> {
934 solve_reduced_beta_pcg_with_diagnostics(s_acc, rhs_beta, max_iterations, relative_tolerance)
935 .map(|(x, _)| x)
936}
937
938#[doc(hidden)]
939pub fn solve_reduced_beta_pcg_with_diagnostics(
940 s_acc: &Array2<f64>,
941 rhs_beta: &Array1<f64>,
942 max_iterations: usize,
943 relative_tolerance: f64,
944) -> Result<(Array1<f64>, ArrowPcgDiagnostics), ArrowSchurGpuFailure> {
945 let k = rhs_beta.len();
946 if s_acc.dim() != (k, k) {
947 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
948 reason: format!(
949 "reduced-β GPU PCG requires a square (k×k) Schur block; got {:?} for k={k}",
950 s_acc.dim()
951 ),
952 });
953 }
954 if k == 0 {
955 return Err(ArrowSchurGpuFailure::Unavailable);
956 }
957
958 #[cfg(not(target_os = "linux"))]
959 {
960 if relative_tolerance.is_nan() || max_iterations == 0 {
961 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
962 reason: "reduced-β GPU PCG: invalid CG controls".to_string(),
963 });
964 }
965 Err(ArrowSchurGpuFailure::Unavailable)
966 }
967
968 #[cfg(target_os = "linux")]
969 {
970 cuda::solve_reduced_beta_pcg_with_diagnostics(
971 s_acc,
972 rhs_beta,
973 max_iterations,
974 relative_tolerance,
975 )
976 }
977}
978
979pub fn solve_sae_matrix_free_pcg(
980 sys: &ArrowSchurSystem,
981 data: &DeviceSaePcgData,
982 ridge_t: f64,
983 ridge_beta: f64,
984 rhs_beta: &Array1<f64>,
985 max_iterations: usize,
986 relative_tolerance: f64,
987) -> Result<(Array1<f64>, ArrowPcgDiagnostics), ArrowSchurGpuFailure> {
988 if sys.k != data.beta_dim || rhs_beta.len() != data.beta_dim || data.p == 0 {
989 return Err(ArrowSchurGpuFailure::Unavailable);
990 }
991 #[cfg(not(target_os = "linux"))]
992 {
993 if ridge_t.is_nan()
994 || ridge_beta.is_nan()
995 || relative_tolerance.is_nan()
996 || max_iterations == 0
997 {
998 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
999 reason: "SAE matrix-free GPU PCG: invalid controls".to_string(),
1000 });
1001 }
1002 Err(ArrowSchurGpuFailure::Unavailable)
1003 }
1004 #[cfg(target_os = "linux")]
1005 {
1006 if data.frame.is_some() {
1013 cuda::solve_sae_matrix_free_pcg_framed(
1014 sys,
1015 data,
1016 ridge_t,
1017 ridge_beta,
1018 rhs_beta,
1019 max_iterations,
1020 relative_tolerance,
1021 )
1022 } else {
1023 cuda::solve_sae_matrix_free_pcg(
1024 sys,
1025 data,
1026 ridge_t,
1027 ridge_beta,
1028 rhs_beta,
1029 max_iterations,
1030 relative_tolerance,
1031 )
1032 }
1033 }
1034}
1035
1036pub trait SaeResidentFrame {
1056 fn refresh(&self, sys: &ArrowSchurSystem) -> Result<(), ArrowSchurGpuFailure>;
1062
1063 fn resolve(
1071 &self,
1072 sys: &ArrowSchurSystem,
1073 ridge_t: f64,
1074 ridge_beta: f64,
1075 rhs_beta: &Array1<f64>,
1076 max_iterations: usize,
1077 relative_tolerance: f64,
1078 ) -> Result<(Array1<f64>, ArrowPcgDiagnostics), ArrowSchurGpuFailure>;
1079}
1080
1081pub fn build_sae_resident_frame(
1089 sys: &ArrowSchurSystem,
1090 cg_iters: usize,
1091) -> Result<
1092 Option<std::sync::Arc<dyn SaeResidentFrame + Send + Sync>>,
1093 ArrowSchurGpuFailure,
1094> {
1095 if sys.k == 0 || cg_iters == 0 {
1100 return Ok(None);
1101 }
1102 #[cfg(target_os = "linux")]
1103 {
1104 cuda::ResidentSaeFrameHandle::build(sys, cg_iters).map(|handle| {
1105 handle.map(|frame| {
1106 std::sync::Arc::new(frame)
1107 as std::sync::Arc<dyn SaeResidentFrame + Send + Sync>
1108 })
1109 })
1110 }
1111 #[cfg(not(target_os = "linux"))]
1113 {
1114 Ok(None)
1115 }
1116}
1117
1118#[cfg(target_os = "linux")]
1126pub struct FrameHostOperands {
1127 pub s_off: Vec<i32>,
1128 pub s_m: Vec<i32>,
1129 pub s_r: Vec<i32>,
1130 pub s_ptr: Vec<i32>,
1131 pub s_data: Vec<f64>,
1132 pub s_blocks: usize,
1133 pub g_off_i: Vec<i32>,
1134 pub g_off_j: Vec<i32>,
1135 pub g_ri: Vec<i32>,
1136 pub g_rj: Vec<i32>,
1137 pub g_mi: Vec<i32>,
1138 pub g_mj: Vec<i32>,
1139 pub g_ptr: Vec<i32>,
1140 pub g_data: Vec<f64>,
1141 pub w_ptr: Vec<i32>,
1142 pub w_data: Vec<f64>,
1143 pub g_blocks: usize,
1144 pub g_max_work: usize,
1145 pub htb_ptr: Vec<i32>,
1146 pub htb: Vec<f64>,
1147 pub q_of: Vec<i32>,
1148 pub n_rows: usize,
1149 pub k: usize,
1150 pub max_q: usize,
1151}
1152
1153#[cfg(target_os = "linux")]
1154fn frame_checked_i32(value: usize) -> Result<i32, ArrowSchurGpuFailure> {
1155 i32::try_from(value).map_err(|_| ArrowSchurGpuFailure::Unavailable)
1156}
1157
1158#[cfg(target_os = "linux")]
1163pub fn flatten_frame_host_operands(
1164 sys: &ArrowSchurSystem,
1165 data: &DeviceSaePcgData,
1166 frame: &crate::arrow_schur::DeviceSaeFrameData,
1167) -> Result<FrameHostOperands, ArrowSchurGpuFailure> {
1168 let n_rows = sys.rows.len();
1169 let k = data.beta_dim;
1170 if frame.row_htbeta.len() != n_rows
1171 || frame.ranks.len() != frame.basis_sizes.len()
1172 || frame.border_offsets.len() != frame.ranks.len()
1173 || data.smooth_blocks.len() != frame.smooth_ranks.len()
1174 {
1175 return Err(ArrowSchurGpuFailure::Unavailable);
1176 }
1177
1178 let mut s_off = Vec::new();
1180 let mut s_m = Vec::new();
1181 let mut s_r = Vec::new();
1182 let mut s_ptr = vec![0_i32];
1183 let mut s_data = Vec::<f64>::new();
1184 for (blk, &r) in data.smooth_blocks.iter().zip(frame.smooth_ranks.iter()) {
1185 let (m, mc) = blk.factor_a.dim();
1186 if m != mc {
1187 return Err(ArrowSchurGpuFailure::Unavailable);
1188 }
1189 s_off.push(frame_checked_i32(blk.global_offset)?);
1190 s_m.push(frame_checked_i32(m)?);
1191 s_r.push(frame_checked_i32(r)?);
1192 for ri in 0..m {
1193 for ci in 0..m {
1194 s_data.push(blk.factor_a[[ri, ci]]);
1195 }
1196 }
1197 s_ptr.push(frame_checked_i32(s_data.len())?);
1198 }
1199
1200 let mut g_off_i = Vec::new();
1202 let mut g_off_j = Vec::new();
1203 let mut g_ri = Vec::new();
1204 let mut g_rj = Vec::new();
1205 let mut g_mi = Vec::new();
1206 let mut g_mj = Vec::new();
1207 let mut g_ptr = vec![0_i32];
1208 let mut g_data = Vec::<f64>::new();
1209 let mut w_ptr = vec![0_i32];
1210 let mut w_data = Vec::<f64>::new();
1211 let mut g_max_work = 0usize;
1212 for blk in &frame.frame_blocks {
1213 let ri = frame.ranks[blk.atom_i];
1214 let rj = frame.ranks[blk.atom_j];
1215 let (mi, mj) = blk.g.dim();
1216 if blk.w.dim() != (ri, rj) {
1217 return Err(ArrowSchurGpuFailure::Unavailable);
1218 }
1219 g_off_i.push(frame_checked_i32(frame.border_offsets[blk.atom_i])?);
1220 g_off_j.push(frame_checked_i32(frame.border_offsets[blk.atom_j])?);
1221 g_ri.push(frame_checked_i32(ri)?);
1222 g_rj.push(frame_checked_i32(rj)?);
1223 g_mi.push(frame_checked_i32(mi)?);
1224 g_mj.push(frame_checked_i32(mj)?);
1225 for r in 0..mi {
1226 for c in 0..mj {
1227 g_data.push(blk.g[[r, c]]);
1228 }
1229 }
1230 g_ptr.push(frame_checked_i32(g_data.len())?);
1231 for a in 0..ri {
1232 for b in 0..rj {
1233 w_data.push(blk.w[[a, b]]);
1234 }
1235 }
1236 w_ptr.push(frame_checked_i32(w_data.len())?);
1237 g_max_work = g_max_work.max(mi * ri);
1238 }
1239
1240 let mut htb_ptr = vec![0_i32];
1243 let mut htb = Vec::<f64>::new();
1244 let mut q_of = Vec::<i32>::with_capacity(n_rows);
1245 let mut max_q = 0usize;
1246 for (i, slab) in frame.row_htbeta.iter().enumerate() {
1247 let qi = sys.row_dims[i];
1248 let q_eff = if !slab.is_empty() && slab.len() == qi * k {
1249 qi
1250 } else {
1251 0
1252 };
1253 q_of.push(frame_checked_i32(q_eff)?);
1254 max_q = max_q.max(q_eff);
1255 if q_eff > 0 {
1256 htb.extend_from_slice(slab);
1257 }
1258 htb_ptr.push(frame_checked_i32(htb.len())?);
1259 }
1260 if max_q == 0 {
1261 max_q = 1;
1264 }
1265
1266 Ok(FrameHostOperands {
1267 s_off,
1268 s_m,
1269 s_r,
1270 s_ptr,
1271 s_data,
1272 s_blocks: data.smooth_blocks.len(),
1273 g_off_i,
1274 g_off_j,
1275 g_ri,
1276 g_rj,
1277 g_mi,
1278 g_mj,
1279 g_ptr,
1280 g_data,
1281 w_ptr,
1282 w_data,
1283 g_blocks: frame.frame_blocks.len(),
1284 g_max_work,
1285 htb_ptr,
1286 htb,
1287 q_of,
1288 n_rows,
1289 k,
1290 max_q,
1291 })
1292}
1293
1294#[cfg(target_os = "linux")]
1301pub fn compute_ainv_host(
1302 sys: &ArrowSchurSystem,
1303 q_of: &[i32],
1304 max_q: usize,
1305 n_rows: usize,
1306 ridge_t: f64,
1307) -> Result<Vec<f64>, ArrowSchurGpuFailure> {
1308 let mut ainv = vec![0.0_f64; n_rows * max_q * max_q];
1309 for (i, row) in sys.rows.iter().enumerate() {
1310 let q = q_of[i] as usize;
1311 if q == 0 {
1312 continue;
1313 }
1314 if row.htt.dim() != (q, q) {
1315 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
1316 reason: format!(
1317 "framed SAE device PCG row {i}: H_tt shape {:?} != ({q}, {q})",
1318 row.htt.dim()
1319 ),
1320 });
1321 }
1322 let mut block = row.htt.clone();
1323 for d in 0..q {
1324 block[[d, d]] += ridge_t;
1325 }
1326 let factor = cholesky_factor_in_place(block.view(), CholeskyGuard::NonnegativePivot)
1327 .ok_or_else(|| ArrowSchurGpuFailure::RidgeBumpRequired {
1328 row: i,
1329 bump: ridge_bump_to_make_pd(row.htt.view(), ridge_t),
1330 })?;
1331 for col in 0..q {
1332 let mut e = Array1::<f64>::zeros(q);
1333 e[col] = 1.0;
1334 let solved = cholesky_solve_vector(factor.view(), e.view());
1335 for r in 0..q {
1336 ainv[i * max_q * max_q + r * max_q + col] = solved[r];
1337 }
1338 }
1339 }
1340 Ok(ainv)
1341}
1342
1343#[doc(hidden)]
1351#[cfg(target_os = "linux")]
1352pub fn framed_schur_matvec_once_on_device(
1353 sys: &ArrowSchurSystem,
1354 data: &DeviceSaePcgData,
1355 ridge_t: f64,
1356 ridge_beta: f64,
1357 x: &Array1<f64>,
1358) -> Result<Array1<f64>, ArrowSchurGpuFailure> {
1359 if sys.k != data.beta_dim || x.len() != data.beta_dim || data.p == 0 {
1360 return Err(ArrowSchurGpuFailure::Unavailable);
1361 }
1362 if data.frame.is_none() {
1363 return Err(ArrowSchurGpuFailure::Unavailable);
1364 }
1365 cuda::framed_schur_matvec_once_on_device(sys, data, ridge_t, ridge_beta, x)
1366}
1367
1368#[doc(hidden)]
1375#[cfg(target_os = "linux")]
1376pub fn framed_reduced_schur_det_once_on_device(
1377 sys: &ArrowSchurSystem,
1378 data: &DeviceSaePcgData,
1379 ridge_t: f64,
1380 ridge_beta: f64,
1381 x: &Array1<f64>,
1382) -> Result<Array1<f64>, ArrowSchurGpuFailure> {
1383 if sys.k != data.beta_dim || x.len() != data.beta_dim || data.p == 0 {
1384 return Err(ArrowSchurGpuFailure::Unavailable);
1385 }
1386 if data.frame.is_none() {
1387 return Err(ArrowSchurGpuFailure::Unavailable);
1388 }
1389 cuda::framed_reduced_schur_det_once_on_device(sys, data, ridge_t, ridge_beta, x)
1390}
1391
1392#[doc(hidden)]
1396pub fn solve_arrow_newton_step_dense_reference(
1397 sys: &ArrowSchurSystem,
1398 ridge_t: f64,
1399 ridge_beta: f64,
1400) -> Result<ArrowSchurGpuSolution, String> {
1401 let n = sys.rows.len();
1402 let d = sys.d;
1403 let k = sys.k;
1404 let total = n.checked_mul(d).ok_or("dimension overflow")? + k;
1405 let mut h = Array2::<f64>::zeros((total, total));
1406 let mut rhs = Array1::<f64>::zeros(total);
1407 for (i, row) in sys.rows.iter().enumerate() {
1408 let base = i * d;
1409 for c in 0..d {
1410 for r in 0..d {
1411 h[[base + r, base + c]] = row.htt[[r, c]];
1412 }
1413 h[[base + c, base + c]] += ridge_t;
1414 }
1415 for c in 0..k {
1416 for r in 0..d {
1417 let value = row.htbeta[[r, c]];
1418 h[[base + r, n * d + c]] = value;
1419 h[[n * d + c, base + r]] = value;
1420 }
1421 }
1422 for r in 0..d {
1423 rhs[base + r] = -row.gt[r];
1424 }
1425 }
1426 for c in 0..k {
1427 for r in 0..k {
1428 h[[n * d + r, n * d + c]] += sys.hbb[[r, c]];
1429 }
1430 h[[n * d + c, n * d + c]] += ridge_beta;
1431 rhs[n * d + c] = -sys.gb[c];
1432 }
1433 let factor = cholesky_factor_in_place(h.view(), CholeskyGuard::NonnegativePivot)
1442 .ok_or_else(|| "dense reference Cholesky failed".to_string())?;
1443 let mut log_det = 0.0_f64;
1444 for i in 0..total {
1445 log_det += factor[[i, i]].ln();
1446 }
1447 log_det *= 2.0;
1448 let solved = cholesky_solve_vector(factor.view(), rhs.view());
1449 let delta_t = solved.slice(ndarray::s![..n * d]).to_owned();
1450 let delta_beta = solved.slice(ndarray::s![n * d..]).to_owned();
1451 Ok(ArrowSchurGpuSolution {
1452 delta_t,
1453 delta_beta,
1454 log_det_hessian: log_det,
1455 })
1456}
1457
1458#[doc(hidden)]
1469pub fn sae_framed_penalty_matvec_cpu(
1470 data: &DeviceSaePcgData,
1471 ridge_beta: f64,
1472 x: &[f64],
1473 out: &mut [f64],
1474) {
1475 let frame = data
1476 .frame
1477 .as_ref()
1478 .expect("sae_framed_penalty_matvec_cpu requires frame metadata");
1479 let k = data.beta_dim;
1480 for a in 0..k {
1481 out[a] = ridge_beta * x[a];
1482 }
1483 for (blk, &r) in data.smooth_blocks.iter().zip(frame.smooth_ranks.iter()) {
1485 let off = blk.global_offset;
1486 let m = blk.factor_a.nrows();
1487 for i_a in 0..m {
1488 for i_b in 0..r {
1489 let mut acc = 0.0_f64;
1490 for j_a in 0..m {
1491 let s = blk.factor_a[[i_a, j_a]];
1492 if s == 0.0 {
1493 continue;
1494 }
1495 acc += s * x[off + j_a * r + i_b];
1496 }
1497 out[off + i_a * r + i_b] += acc;
1498 }
1499 }
1500 }
1501 for blk in &frame.frame_blocks {
1503 let r_i = frame.ranks[blk.atom_i];
1504 let r_j = frame.ranks[blk.atom_j];
1505 let off_i = frame.border_offsets[blk.atom_i];
1506 let off_j = frame.border_offsets[blk.atom_j];
1507 let (m_i, m_j) = blk.g.dim();
1508 for li in 0..m_i {
1509 let yi_base = off_i + li * r_i;
1510 for lj in 0..m_j {
1511 let g = blk.g[[li, lj]];
1512 if g == 0.0 {
1513 continue;
1514 }
1515 let xj_base = off_j + lj * r_j;
1516 for a in 0..r_i {
1517 let mut acc = 0.0_f64;
1518 for b in 0..r_j {
1519 acc += blk.w[[a, b]] * x[xj_base + b];
1520 }
1521 out[yi_base + a] += g * acc;
1522 }
1523 }
1524 }
1525 }
1526}
1527
1528#[doc(hidden)]
1537pub fn sae_framed_schur_matvec_cpu(
1538 sys: &ArrowSchurSystem,
1539 data: &DeviceSaePcgData,
1540 ridge_t: f64,
1541 ridge_beta: f64,
1542 x: &[f64],
1543 out: &mut [f64],
1544) -> Result<(), String> {
1545 let frame = data
1546 .frame
1547 .as_ref()
1548 .ok_or("sae_framed_schur_matvec_cpu requires frame metadata")?;
1549 let k = data.beta_dim;
1550 sae_framed_penalty_matvec_cpu(data, ridge_beta, x, out);
1551 if frame.row_htbeta.len() != sys.rows.len() {
1552 return Err(format!(
1553 "sae_framed_schur_matvec_cpu: {} row_htbeta slabs but {} rows",
1554 frame.row_htbeta.len(),
1555 sys.rows.len()
1556 ));
1557 }
1558 for (i, row) in sys.rows.iter().enumerate() {
1559 let slab = &frame.row_htbeta[i];
1560 if slab.is_empty() {
1561 continue;
1562 }
1563 let qi = sys.row_dims[i];
1564 if qi == 0 || slab.len() != qi * k {
1565 continue;
1566 }
1567 let mut h = vec![0.0_f64; qi];
1569 for c in 0..qi {
1570 let base = c * k;
1571 let mut acc = 0.0_f64;
1572 for a in 0..k {
1573 acc += slab[base + a] * x[a];
1574 }
1575 h[c] = acc;
1576 }
1577 let mut block = row.htt.clone();
1579 for d in 0..qi {
1580 block[[d, d]] += ridge_t;
1581 }
1582 let factor = cholesky_factor_in_place(block.view(), CholeskyGuard::NonnegativePivot)
1583 .ok_or_else(|| format!("sae_framed_schur_matvec_cpu: row {i} H_tt not PD"))?;
1584 let s = cholesky_solve_vector(factor.view(), Array1::from_vec(h).view());
1585 for c in 0..qi {
1587 let sc = s[c];
1588 if sc == 0.0 {
1589 continue;
1590 }
1591 let base = c * k;
1592 for a in 0..k {
1593 out[a] -= slab[base + a] * sc;
1594 }
1595 }
1596 }
1597 Ok(())
1598}
1599
1600#[cfg(target_os = "linux")]
1601mod cuda {
1602 use super::{
1603 ArrowSchurGpuFailure, ArrowSchurGpuSolution, pack_block, pack_host,
1604 pack_host_d_and_stacked_b,
1605 };
1606 use crate::arrow_schur::{
1607 ArrowPcgDiagnostics, ArrowSchurSystem, DeviceSaeFrameData, DeviceSaePcgData, PcgStopReason,
1608 };
1609 use cudarc::cublas::sys::{
1610 cublasDiagType_t, cublasFillMode_t, cublasOperation_t, cublasSideMode_t, cublasStatus_t,
1611 };
1612 use cudarc::cublas::{CudaBlas, Gemm, GemmConfig, Gemv, GemvConfig};
1613 use cudarc::cusolver::{DnHandle, sys as cusolver_sys};
1614 use cudarc::driver::{
1615 CudaContext, CudaModule, CudaSlice, CudaStream, DevicePtr, DevicePtrMut, LaunchConfig,
1616 PushKernelArg,
1617 };
1618 use gam_gpu::driver::to_i32;
1619 use gam_gpu::linalg_dispatch::{DispatchOp, route_through_gpu};
1620 use ndarray::Array1;
1621 use std::sync::{Arc, OnceLock};
1622
1623 struct RowSlot {
1628 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,
1637 bump: Option<f64>,
1640 tile_partial_schur: Option<Vec<f64>>, tile_partial_rhs: Option<Vec<f64>>, delta_t_block: Vec<f64>, }
1646
1647 pub(super) fn solve_multi_gpu(
1668 sys: &ArrowSchurSystem,
1669 ridge_t: f64,
1670 ridge_beta: f64,
1671 ) -> Result<ArrowSchurGpuSolution, ArrowSchurGpuFailure> {
1672 let n = sys.rows.len();
1673 let d = sys.d;
1674 let k = sys.k;
1675 if n == 0 || d == 0 || k == 0 {
1676 return Err(ArrowSchurGpuFailure::Unavailable);
1677 }
1678 if sys.hbb_matvec.is_some() || sys.htbeta_matvec.is_some() || sys.hbb.dim() != (k, k) {
1682 return Err(ArrowSchurGpuFailure::Unavailable);
1683 }
1684
1685 let runtime = super::resolve_runtime_for_device_path()?
1686 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
1687 if runtime.device_count() < 2 {
1688 return Err(ArrowSchurGpuFailure::Unavailable);
1689 }
1690
1691 let mut slots: Vec<RowSlot> = Vec::with_capacity(n);
1693 for row in &sys.rows {
1694 if row.htt.dim() != (d, d) || row.htbeta.dim() != (d, k) || row.gt.len() != d {
1695 return Err(ArrowSchurGpuFailure::Unavailable);
1696 }
1697 let mut d_block = Vec::with_capacity(d * d);
1698 let mut b_block = Vec::with_capacity(d * k);
1699 let mut g_vec = Vec::with_capacity(d);
1700 pack_block(row, ridge_t, d, k, &mut d_block, &mut b_block, &mut g_vec);
1701 slots.push(RowSlot {
1702 d_block,
1703 b_block,
1704 g_vec,
1705 l_block: Vec::new(),
1706 u_vec: Vec::new(),
1707 y_block: Vec::new(),
1708 log_det_local: 0.0,
1709 bump: None,
1710 tile_partial_schur: None,
1711 tile_partial_rhs: None,
1712 delta_t_block: vec![0.0; d],
1713 });
1714 }
1715
1716 let forward_ok = gam_gpu::pool::scatter_batched(runtime, &mut slots, |ordinal, tile| {
1718 forward_tile(ordinal, d, k, tile)
1719 });
1720 if forward_ok.is_none() {
1721 return Err(ArrowSchurGpuFailure::Unavailable);
1722 }
1723
1724 let row_base_of_tile = gam_gpu::pool::balanced_partition(runtime, n);
1726 if let Some((row, bump)) = slots
1727 .iter()
1728 .enumerate()
1729 .find_map(|(i, slot)| slot.bump.map(|b| (i, b)))
1730 {
1731 return Err(ArrowSchurGpuFailure::RidgeBumpRequired { row, bump });
1732 }
1733
1734 let mut schur_host = vec![0.0_f64; k * k];
1739 for col in 0..k {
1740 for row in 0..k {
1741 let mut v = sys.hbb[[row, col]];
1742 if row == col {
1743 v += ridge_beta;
1744 }
1745 schur_host[col * k + row] = v;
1746 }
1747 }
1748 let mut rhs_host: Vec<f64> = sys.gb.iter().map(|v| -v).collect();
1749 let mut log_det = 0.0_f64;
1750 for start in tile_starts(&row_base_of_tile) {
1751 let slot = &slots[start];
1752 let partial_schur = slot
1753 .tile_partial_schur
1754 .as_ref()
1755 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
1756 let partial_rhs = slot
1757 .tile_partial_rhs
1758 .as_ref()
1759 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
1760 for idx in 0..k * k {
1765 schur_host[idx] += partial_schur[idx];
1766 }
1767 for a in 0..k {
1768 rhs_host[a] += partial_rhs[a];
1769 }
1770 }
1771 for slot in &slots {
1772 log_det += slot.log_det_local;
1773 }
1774
1775 let primary = runtime.selected_device().ordinal;
1779 let stream = gam_gpu::device_runtime::cuda_context_for(primary)
1780 .and_then(|ctx| ctx.new_stream().ok())
1781 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
1782 let solver =
1783 DnHandle::new(stream.clone()).map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
1784 let blas = CudaBlas::new(stream.clone()).map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
1785 let mut schur_dev = stream
1786 .clone_htod(&schur_host)
1787 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
1788 let mut rhs_dev = stream
1789 .clone_htod(&rhs_host)
1790 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
1791 let info = potrf_single(&solver, &stream, k, &mut schur_dev)?;
1792 if info != 0 {
1793 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
1794 reason: format!("multi-GPU Schur Cholesky failed at pivot {info}"),
1795 });
1796 }
1797 trsm_single(&blas, &stream, k, &schur_dev, &mut rhs_dev, false, false)?;
1798 trsm_single(&blas, &stream, k, &schur_dev, &mut rhs_dev, false, true)?;
1799 let delta_beta_host = stream
1800 .clone_dtoh(&rhs_dev)
1801 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
1802 let delta_beta = Array1::from_vec(delta_beta_host.clone());
1803 let l_schur_host = stream
1804 .clone_dtoh(&schur_dev)
1805 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
1806 for j in 0..k {
1807 log_det += l_schur_host[j * k + j].ln();
1808 }
1809 log_det *= 2.0;
1810
1811 let delta_beta_ref = &delta_beta_host;
1813 let back_ok = gam_gpu::pool::scatter_batched(runtime, &mut slots, |ordinal, tile| {
1814 back_sub_tile(ordinal, d, k, delta_beta_ref, tile)
1815 });
1816 if back_ok.is_none() {
1817 return Err(ArrowSchurGpuFailure::Unavailable);
1818 }
1819
1820 let mut delta_t = Array1::<f64>::zeros(n * d);
1822 for (i, slot) in slots.iter().enumerate() {
1823 let base = i * d;
1824 for r in 0..d {
1825 delta_t[base + r] = slot.delta_t_block[r];
1826 }
1827 }
1828
1829 Ok(ArrowSchurGpuSolution {
1830 delta_t,
1831 delta_beta,
1832 log_det_hessian: log_det,
1833 })
1834 }
1835
1836 fn tile_starts(tiles: &[(usize, std::ops::Range<usize>)]) -> impl Iterator<Item = usize> + '_ {
1839 tiles.iter().map(|(_, range)| range.start)
1840 }
1841
1842 fn forward_tile(ordinal: usize, d: usize, k: usize, tile: &mut [RowSlot]) -> Option<()> {
1850 if tile.is_empty() {
1851 return Some(());
1852 }
1853 let stream = gam_gpu::device_runtime::cuda_context_for(ordinal)
1856 .and_then(|ctx| ctx.new_stream().ok())?;
1857 let solver = DnHandle::new(stream.clone()).ok()?;
1858 let blas = CudaBlas::new(stream.clone()).ok()?;
1859 let m = tile.len();
1860
1861 let mut d_host = Vec::with_capacity(m * d * d);
1864 let mut b_host = Vec::with_capacity(m * d * k);
1865 let mut g_host = Vec::with_capacity(m * d);
1866 for slot in tile.iter() {
1867 d_host.extend_from_slice(&slot.d_block);
1868 b_host.extend_from_slice(&slot.b_block);
1869 g_host.extend_from_slice(&slot.g_vec);
1870 }
1871 let mut d_dev = stream.clone_htod(&d_host).ok()?;
1872 let mut b_dev = stream.clone_htod(&b_host).ok()?;
1873 let mut g_dev = stream.clone_htod(&g_host).ok()?;
1874
1875 let info_host = potrf_batched(&solver, &stream, d, m, &mut d_dev).ok()?;
1881 if let Some(local) = info_host.iter().position(|info| *info != 0) {
1882 tile[local].bump = Some(super::ridge_bump_to_make_pd_colmajor(
1883 &tile[local].d_block,
1884 d,
1885 ));
1886 return Some(());
1887 }
1888
1889 trsm_batched_lower_inplace(&blas, &stream, d, m, 1, &d_dev, &mut g_dev).ok()?;
1891 trsm_batched_lower_inplace(&blas, &stream, d, m, k, &d_dev, &mut b_dev).ok()?;
1892
1893 let mut schur_dev = stream.alloc_zeros::<f64>(k * k).ok()?;
1895 let mut rhs_dev = stream.alloc_zeros::<f64>(k).ok()?;
1896 accumulate_schur(&blas, d, k, m, &b_dev, &g_dev, &mut schur_dev, &mut rhs_dev).ok()?;
1897
1898 let l_host = stream.clone_dtoh(&d_dev).ok()?;
1900 let u_host = stream.clone_dtoh(&g_dev).ok()?;
1901 let y_host = stream.clone_dtoh(&b_dev).ok()?;
1902 let partial_schur = stream.clone_dtoh(&schur_dev).ok()?;
1903 let partial_rhs = stream.clone_dtoh(&rhs_dev).ok()?;
1904
1905 for (local, slot) in tile.iter_mut().enumerate() {
1906 let l_base = local * d * d;
1907 let u_base = local * d;
1908 let y_base = local * d * k;
1909 slot.l_block = l_host[l_base..l_base + d * d].to_vec();
1910 slot.u_vec = u_host[u_base..u_base + d].to_vec();
1911 slot.y_block = y_host[y_base..y_base + d * k].to_vec();
1912 let mut log_det_local = 0.0_f64;
1913 for j in 0..d {
1914 log_det_local += l_host[l_base + j * d + j].ln();
1915 }
1916 slot.log_det_local = log_det_local;
1917 }
1918 tile[0].tile_partial_schur = Some(partial_schur);
1919 tile[0].tile_partial_rhs = Some(partial_rhs);
1920 Some(())
1921 }
1922
1923 fn back_sub_tile(
1927 ordinal: usize,
1928 d: usize,
1929 k: usize,
1930 delta_beta: &[f64],
1931 tile: &mut [RowSlot],
1932 ) -> Option<()> {
1933 if tile.is_empty() {
1934 return Some(());
1935 }
1936 let stream = gam_gpu::device_runtime::cuda_context_for(ordinal)
1939 .and_then(|ctx| ctx.new_stream().ok())?;
1940 let blas = CudaBlas::new(stream.clone()).ok()?;
1941 let m = tile.len();
1942
1943 let mut l_host = Vec::with_capacity(m * d * d);
1944 let mut u_host = Vec::with_capacity(m * d);
1945 let mut y_host = Vec::with_capacity(m * d * k);
1946 for slot in tile.iter() {
1947 l_host.extend_from_slice(&slot.l_block);
1948 u_host.extend_from_slice(&slot.u_vec);
1949 y_host.extend_from_slice(&slot.y_block);
1950 }
1951 let d_dev = stream.clone_htod(&l_host).ok()?;
1952 let mut g_dev = stream.clone_htod(&u_host).ok()?;
1953 let b_dev = stream.clone_htod(&y_host).ok()?;
1954 let rhs_dev = stream.clone_htod(&delta_beta.to_vec()).ok()?;
1955
1956 accumulate_back_sub_rhs(&blas, d, k, m, &b_dev, &rhs_dev, &mut g_dev).ok()?;
1958 trsm_batched_lower_inplace_transposed(&blas, &stream, d, m, 1, &d_dev, &mut g_dev).ok()?;
1959 let x_host = stream.clone_dtoh(&g_dev).ok()?;
1960 for (local, slot) in tile.iter_mut().enumerate() {
1961 let base = local * d;
1962 for r in 0..d {
1963 slot.delta_t_block[r] = -x_host[base + r];
1964 }
1965 }
1966 Some(())
1967 }
1968
1969 pub(super) fn solve(
1970 sys: &ArrowSchurSystem,
1971 ridge_t: f64,
1972 ridge_beta: f64,
1973 ) -> Result<ArrowSchurGpuSolution, ArrowSchurGpuFailure> {
1974 let n = sys.rows.len();
1975 let d = sys.d;
1976 let k = sys.k;
1977 let runtime = route_through_gpu(DispatchOp::SmallDenseBatchedPotrf { p: d, batch: n })
1978 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
1979
1980 let stream = gam_gpu::device_runtime::cuda_context_for(runtime.device.ordinal)
1981 .and_then(|ctx| ctx.new_stream().ok())
1982 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
1983 let solver =
1984 DnHandle::new(stream.clone()).map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
1985 let blas = CudaBlas::new(stream.clone()).map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
1986
1987 let (d_host, b_host, g_host) = pack_host(sys, ridge_t);
1989 let mut d_dev = stream
1990 .clone_htod(&d_host)
1991 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
1992 let mut b_dev = stream
1993 .clone_htod(&b_host)
1994 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
1995 let mut g_dev = stream
1996 .clone_htod(&g_host)
1997 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
1998
1999 let info_host = potrf_batched(&solver, &stream, d, n, &mut d_dev)?;
2007 if let Some(idx) = info_host.iter().position(|info| *info != 0) {
2008 return Err(ArrowSchurGpuFailure::RidgeBumpRequired {
2012 row: idx,
2013 bump: super::ridge_bump_to_make_pd(sys.rows[idx].htt.view(), ridge_t),
2014 });
2015 }
2016
2017 trsm_batched_lower_inplace(&blas, &stream, d, n, 1, &d_dev, &mut g_dev)?;
2020 trsm_batched_lower_inplace(&blas, &stream, d, n, k, &d_dev, &mut b_dev)?;
2023
2024 let schur_init: Vec<f64> = {
2043 let mut tmp = Vec::with_capacity(k * k);
2044 for col in 0..k {
2045 for row in 0..k {
2046 let mut v = sys.hbb[[row, col]];
2047 if row == col {
2048 v += ridge_beta;
2049 }
2050 tmp.push(v);
2051 }
2052 }
2053 tmp
2054 };
2055 let mut schur_dev = stream
2056 .clone_htod(&schur_init)
2057 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2058 let rhs_init: Vec<f64> = sys.gb.iter().map(|v| -v).collect();
2059 let mut rhs_dev = stream
2060 .clone_htod(&rhs_init)
2061 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2062
2063 accumulate_schur(&blas, d, k, n, &b_dev, &g_dev, &mut schur_dev, &mut rhs_dev)?;
2064
2065 let info = potrf_single(&solver, &stream, k, &mut schur_dev)?;
2067 if info != 0 {
2068 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
2069 reason: format!("Schur Cholesky failed at pivot {info}"),
2070 });
2071 }
2072 trsm_single(&blas, &stream, k, &schur_dev, &mut rhs_dev, false, false)?;
2074 trsm_single(&blas, &stream, k, &schur_dev, &mut rhs_dev, false, true)?;
2075 let delta_beta_host = stream
2076 .clone_dtoh(&rhs_dev)
2077 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2078 let delta_beta = Array1::from_vec(delta_beta_host.clone());
2079
2080 accumulate_back_sub_rhs(&blas, d, k, n, &b_dev, &rhs_dev, &mut g_dev)?;
2088 trsm_batched_lower_inplace_transposed(&blas, &stream, d, n, 1, &d_dev, &mut g_dev)?;
2089
2090 let x_host = stream
2091 .clone_dtoh(&g_dev)
2092 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2093 let mut delta_t = Array1::<f64>::zeros(n * d);
2094 for (i, v) in x_host.iter().enumerate() {
2095 delta_t[i] = -*v;
2096 }
2097
2098 let l_local_host = stream
2100 .clone_dtoh(&d_dev)
2101 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2102 let l_schur_host = stream
2103 .clone_dtoh(&schur_dev)
2104 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2105 let mut log_det = 0.0_f64;
2106 for i in 0..n {
2107 let base = i * d * d;
2108 for j in 0..d {
2109 log_det += l_local_host[base + j * d + j].ln();
2110 }
2111 }
2112 for j in 0..k {
2113 log_det += l_schur_host[j * k + j].ln();
2114 }
2115 log_det *= 2.0;
2116
2117 Ok(ArrowSchurGpuSolution {
2118 delta_t,
2119 delta_beta,
2120 log_det_hessian: log_det,
2121 })
2122 }
2123
2124 fn potrf_batched(
2125 solver: &DnHandle,
2126 stream: &Arc<CudaStream>,
2127 p: usize,
2128 batch: usize,
2129 matrices: &mut CudaSlice<f64>,
2130 ) -> Result<Vec<i32>, ArrowSchurGpuFailure> {
2131 let p_i = to_i32(p).ok_or(ArrowSchurGpuFailure::Unavailable)?;
2132 let batch_i = to_i32(batch).ok_or(ArrowSchurGpuFailure::Unavailable)?;
2133 let matrix_len = p * p;
2134 let bytes_per = (matrix_len * std::mem::size_of::<f64>()) as u64;
2135 let (base_ptr, _record) = matrices.device_ptr_mut(stream);
2136 let mut ptrs = Vec::with_capacity(batch);
2137 for idx in 0..batch {
2138 ptrs.push(base_ptr + (idx as u64) * bytes_per);
2139 }
2140 let mut ptrs_dev = stream
2141 .clone_htod(&ptrs)
2142 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2143 let mut info_dev = stream
2144 .alloc_zeros::<i32>(batch)
2145 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2146 let status = {
2147 let (ptrs_ptr, _ptrs_record) = ptrs_dev.device_ptr_mut(stream);
2148 let (info_ptr, _info_record) = info_dev.device_ptr_mut(stream);
2149 unsafe {
2152 cusolver_sys::cusolverDnDpotrfBatched(
2153 solver.cu(),
2154 cusolver_sys::cublasFillMode_t::CUBLAS_FILL_MODE_LOWER,
2155 p_i,
2156 ptrs_ptr as *mut *mut f64,
2157 p_i,
2158 info_ptr as *mut i32,
2159 batch_i,
2160 )
2161 }
2162 };
2163 if status != cusolver_sys::cusolverStatus_t::CUSOLVER_STATUS_SUCCESS {
2164 return Err(ArrowSchurGpuFailure::Unavailable);
2165 }
2166 stream
2167 .clone_dtoh(&info_dev)
2168 .map_err(|_| ArrowSchurGpuFailure::Unavailable)
2169 }
2170
2171 fn potrf_single(
2172 solver: &DnHandle,
2173 stream: &Arc<CudaStream>,
2174 p: usize,
2175 matrix: &mut CudaSlice<f64>,
2176 ) -> Result<i32, ArrowSchurGpuFailure> {
2177 let p_i = to_i32(p).ok_or(ArrowSchurGpuFailure::Unavailable)?;
2178 let uplo = cusolver_sys::cublasFillMode_t::CUBLAS_FILL_MODE_LOWER;
2179 let mut lwork = 0_i32;
2180 {
2181 let (mat_ptr, _rec) = matrix.device_ptr_mut(stream);
2182 let status = unsafe {
2184 cusolver_sys::cusolverDnDpotrf_bufferSize(
2185 solver.cu(),
2186 uplo,
2187 p_i,
2188 mat_ptr as *mut f64,
2189 p_i,
2190 &mut lwork,
2191 )
2192 };
2193 if status != cusolver_sys::cusolverStatus_t::CUSOLVER_STATUS_SUCCESS {
2194 return Err(ArrowSchurGpuFailure::Unavailable);
2195 }
2196 }
2197 let lwork_usize = usize::try_from(lwork).map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2198 let mut workspace = stream
2199 .alloc_zeros::<f64>(lwork_usize.max(1))
2200 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2201 let mut info_dev = stream
2202 .alloc_zeros::<i32>(1)
2203 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2204 {
2205 let (mat_ptr, _rec) = matrix.device_ptr_mut(stream);
2206 let (work_ptr, _wrec) = workspace.device_ptr_mut(stream);
2207 let (info_ptr, _irec) = info_dev.device_ptr_mut(stream);
2208 let status = unsafe {
2210 cusolver_sys::cusolverDnDpotrf(
2211 solver.cu(),
2212 uplo,
2213 p_i,
2214 mat_ptr as *mut f64,
2215 p_i,
2216 work_ptr as *mut f64,
2217 lwork,
2218 info_ptr as *mut i32,
2219 )
2220 };
2221 if status != cusolver_sys::cusolverStatus_t::CUSOLVER_STATUS_SUCCESS {
2222 return Err(ArrowSchurGpuFailure::Unavailable);
2223 }
2224 }
2225 let info_host = stream
2226 .clone_dtoh(&info_dev)
2227 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2228 Ok(info_host[0])
2229 }
2230
2231 #[derive(Clone, Copy)]
2245 struct BatchedTrsmSpec {
2246 d: usize,
2247 n: usize,
2248 nrhs: usize,
2249 rhs_leading_dim: usize,
2250 rhs_block_stride: usize,
2251 transposed: bool,
2252 }
2253
2254 impl BatchedTrsmSpec {
2255 fn block_tiles(d: usize, n: usize, nrhs: usize) -> Self {
2257 Self {
2258 d,
2259 n,
2260 nrhs,
2261 rhs_leading_dim: d,
2262 rhs_block_stride: d * nrhs,
2263 transposed: false,
2264 }
2265 }
2266
2267 fn stacked_matrix(d: usize, n: usize, nrhs: usize) -> Self {
2270 Self {
2271 d,
2272 n,
2273 nrhs,
2274 rhs_leading_dim: n * d,
2275 rhs_block_stride: d,
2276 transposed: false,
2277 }
2278 }
2279
2280 fn transposed(self) -> Self {
2281 Self {
2282 transposed: true,
2283 ..self
2284 }
2285 }
2286 }
2287
2288 fn trsm_batched_lower_inplace(
2289 blas: &CudaBlas,
2290 stream: &Arc<CudaStream>,
2291 d: usize,
2292 n: usize,
2293 nrhs: usize,
2294 l_stack: &CudaSlice<f64>,
2295 rhs_stack: &mut CudaSlice<f64>,
2296 ) -> Result<(), ArrowSchurGpuFailure> {
2297 trsm_batched_inplace_inner(
2298 blas,
2299 stream,
2300 BatchedTrsmSpec::block_tiles(d, n, nrhs),
2301 l_stack,
2302 rhs_stack,
2303 )
2304 }
2305
2306 fn trsm_batched_lower_inplace_transposed(
2308 blas: &CudaBlas,
2309 stream: &Arc<CudaStream>,
2310 d: usize,
2311 n: usize,
2312 nrhs: usize,
2313 l_stack: &CudaSlice<f64>,
2314 rhs_stack: &mut CudaSlice<f64>,
2315 ) -> Result<(), ArrowSchurGpuFailure> {
2316 trsm_batched_inplace_inner(
2317 blas,
2318 stream,
2319 BatchedTrsmSpec::block_tiles(d, n, nrhs).transposed(),
2320 l_stack,
2321 rhs_stack,
2322 )
2323 }
2324
2325 fn trsm_batched_lower_inplace_stacked(
2328 blas: &CudaBlas,
2329 stream: &Arc<CudaStream>,
2330 d: usize,
2331 n: usize,
2332 nrhs: usize,
2333 l_stack: &CudaSlice<f64>,
2334 rhs_stacked: &mut CudaSlice<f64>,
2335 ) -> Result<(), ArrowSchurGpuFailure> {
2336 trsm_batched_inplace_inner(
2337 blas,
2338 stream,
2339 BatchedTrsmSpec::stacked_matrix(d, n, nrhs),
2340 l_stack,
2341 rhs_stacked,
2342 )
2343 }
2344
2345 fn trsm_batched_inplace_inner(
2346 blas: &CudaBlas,
2347 stream: &Arc<CudaStream>,
2348 spec: BatchedTrsmSpec,
2349 l_stack: &CudaSlice<f64>,
2350 rhs_stack: &mut CudaSlice<f64>,
2351 ) -> Result<(), ArrowSchurGpuFailure> {
2352 let BatchedTrsmSpec {
2353 d,
2354 n,
2355 nrhs,
2356 rhs_leading_dim,
2357 rhs_block_stride,
2358 transposed,
2359 } = spec;
2360 let alpha = 1.0_f64;
2361 let d_i = to_i32(d).ok_or(ArrowSchurGpuFailure::Unavailable)?;
2362 let nrhs_i = to_i32(nrhs).ok_or(ArrowSchurGpuFailure::Unavailable)?;
2363 let rhs_ld_i = to_i32(rhs_leading_dim).ok_or(ArrowSchurGpuFailure::Unavailable)?;
2364 let batch_i = to_i32(n).ok_or(ArrowSchurGpuFailure::Unavailable)?;
2365 let l_bytes_per = (d * d * std::mem::size_of::<f64>()) as u64;
2366 let rhs_bytes_per = (rhs_block_stride * std::mem::size_of::<f64>()) as u64;
2367 let (l_base, _l_record) = l_stack.device_ptr(stream);
2368 let (rhs_base, _rhs_record) = rhs_stack.device_ptr_mut(stream);
2369 let mut l_ptrs = Vec::with_capacity(n);
2370 let mut rhs_ptrs = Vec::with_capacity(n);
2371 for i in 0..n {
2372 l_ptrs.push(l_base + (i as u64) * l_bytes_per);
2373 rhs_ptrs.push(rhs_base + (i as u64) * rhs_bytes_per);
2374 }
2375 let mut l_ptrs_dev = stream
2376 .clone_htod(&l_ptrs)
2377 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2378 let mut rhs_ptrs_dev = stream
2379 .clone_htod(&rhs_ptrs)
2380 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2381 let (l_ptrs_ptr, _l_ptrs_rec) = l_ptrs_dev.device_ptr_mut(stream);
2382 let (rhs_ptrs_ptr, _rhs_ptrs_rec) = rhs_ptrs_dev.device_ptr_mut(stream);
2383 let op = if transposed {
2384 cublasOperation_t::CUBLAS_OP_T
2385 } else {
2386 cublasOperation_t::CUBLAS_OP_N
2387 };
2388 let handle = *blas.handle();
2389 let status = unsafe {
2392 cudarc::cublas::sys::cublasDtrsmBatched(
2393 handle,
2394 cublasSideMode_t::CUBLAS_SIDE_LEFT,
2395 cublasFillMode_t::CUBLAS_FILL_MODE_LOWER,
2396 op,
2397 cublasDiagType_t::CUBLAS_DIAG_NON_UNIT,
2398 d_i,
2399 nrhs_i,
2400 &alpha,
2401 l_ptrs_ptr as *const *const f64,
2402 d_i,
2403 rhs_ptrs_ptr as *const *mut f64,
2404 rhs_ld_i,
2405 batch_i,
2406 )
2407 };
2408 if status != cublasStatus_t::CUBLAS_STATUS_SUCCESS {
2409 return Err(ArrowSchurGpuFailure::Unavailable);
2410 }
2411 Ok(())
2412 }
2413
2414 fn trsm_single(
2417 blas: &CudaBlas,
2418 stream: &Arc<CudaStream>,
2419 n: usize,
2420 l: &CudaSlice<f64>,
2421 rhs: &mut CudaSlice<f64>,
2422 upper: bool,
2423 transposed: bool,
2424 ) -> Result<(), ArrowSchurGpuFailure> {
2425 let alpha = 1.0_f64;
2426 let n_i = to_i32(n).ok_or(ArrowSchurGpuFailure::Unavailable)?;
2427 let handle = *blas.handle();
2428 let (l_ptr, _l_rec) = l.device_ptr(stream);
2429 let (rhs_ptr, _rhs_rec) = rhs.device_ptr_mut(stream);
2430 let status = unsafe {
2432 cudarc::cublas::sys::cublasDtrsm_v2(
2433 handle,
2434 cublasSideMode_t::CUBLAS_SIDE_LEFT,
2435 if upper {
2436 cublasFillMode_t::CUBLAS_FILL_MODE_UPPER
2437 } else {
2438 cublasFillMode_t::CUBLAS_FILL_MODE_LOWER
2439 },
2440 if transposed {
2441 cublasOperation_t::CUBLAS_OP_T
2442 } else {
2443 cublasOperation_t::CUBLAS_OP_N
2444 },
2445 cublasDiagType_t::CUBLAS_DIAG_NON_UNIT,
2446 n_i,
2447 1,
2448 &alpha,
2449 l_ptr as *const f64,
2450 n_i,
2451 rhs_ptr as *mut f64,
2452 n_i,
2453 )
2454 };
2455 if status != cublasStatus_t::CUBLAS_STATUS_SUCCESS {
2456 return Err(ArrowSchurGpuFailure::Unavailable);
2457 }
2458 Ok(())
2459 }
2460
2461 fn accumulate_schur(
2465 blas: &CudaBlas,
2466 d: usize,
2467 k: usize,
2468 n: usize,
2469 y_stack: &CudaSlice<f64>,
2470 u_stack: &CudaSlice<f64>,
2471 schur: &mut CudaSlice<f64>,
2472 rhs: &mut CudaSlice<f64>,
2473 ) -> Result<(), ArrowSchurGpuFailure> {
2474 let y_block_elems = d * k;
2475 let u_block_elems = d;
2476 for i in 0..n {
2477 let y_slice = y_stack.slice(i * y_block_elems..(i + 1) * y_block_elems);
2478 let u_slice = u_stack.slice(i * u_block_elems..(i + 1) * u_block_elems);
2479 let gemm_cfg = GemmConfig::<f64> {
2481 transa: cublasOperation_t::CUBLAS_OP_T,
2482 transb: cublasOperation_t::CUBLAS_OP_N,
2483 m: to_i32(k).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2484 n: to_i32(k).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2485 k: to_i32(d).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2486 alpha: -1.0,
2487 lda: to_i32(d).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2488 ldb: to_i32(d).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2489 beta: 1.0,
2490 ldc: to_i32(k).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2491 };
2492 unsafe { blas.gemm(gemm_cfg, &y_slice, &y_slice, schur) }
2494 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2495 let gemv_cfg = GemvConfig::<f64> {
2497 trans: cublasOperation_t::CUBLAS_OP_T,
2498 m: to_i32(d).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2499 n: to_i32(k).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2500 alpha: 1.0,
2501 lda: to_i32(d).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2502 incx: 1,
2503 beta: 1.0,
2504 incy: 1,
2505 };
2506 unsafe { blas.gemv(gemv_cfg, &y_slice, &u_slice, rhs) }
2509 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2510 }
2511 Ok(())
2512 }
2513
2514 fn schur_gemm_stacked(
2530 blas: &CudaBlas,
2531 rows: usize,
2532 k: usize,
2533 y_stacked: &CudaSlice<f64>,
2534 schur: &mut CudaSlice<f64>,
2535 ) -> Result<(), ArrowSchurGpuFailure> {
2536 let k_i = to_i32(k).ok_or(ArrowSchurGpuFailure::Unavailable)?;
2537 let rows_i = to_i32(rows).ok_or(ArrowSchurGpuFailure::Unavailable)?;
2538 let cfg = GemmConfig::<f64> {
2539 transa: cublasOperation_t::CUBLAS_OP_T,
2540 transb: cublasOperation_t::CUBLAS_OP_N,
2541 m: k_i,
2542 n: k_i,
2543 k: rows_i,
2544 alpha: -1.0,
2545 lda: rows_i,
2546 ldb: rows_i,
2547 beta: 1.0,
2548 ldc: k_i,
2549 };
2550 unsafe { blas.gemm(cfg, y_stacked, y_stacked, schur) }
2554 .map_err(|_| ArrowSchurGpuFailure::Unavailable)
2555 }
2556
2557 fn schur_rhs_stacked(
2562 blas: &CudaBlas,
2563 rows: usize,
2564 k: usize,
2565 y_stacked: &CudaSlice<f64>,
2566 u_stacked: &CudaSlice<f64>,
2567 rhs: &mut CudaSlice<f64>,
2568 ) -> Result<(), ArrowSchurGpuFailure> {
2569 let k_i = to_i32(k).ok_or(ArrowSchurGpuFailure::Unavailable)?;
2570 let rows_i = to_i32(rows).ok_or(ArrowSchurGpuFailure::Unavailable)?;
2571 let cfg = GemvConfig::<f64> {
2572 trans: cublasOperation_t::CUBLAS_OP_T,
2573 m: rows_i,
2574 n: k_i,
2575 alpha: 1.0,
2576 lda: rows_i,
2577 incx: 1,
2578 beta: 1.0,
2579 incy: 1,
2580 };
2581 unsafe { blas.gemv(cfg, y_stacked, u_stacked, rhs) }
2585 .map_err(|_| ArrowSchurGpuFailure::Unavailable)
2586 }
2587
2588 fn back_sub_rhs_stacked(
2592 blas: &CudaBlas,
2593 rows: usize,
2594 k: usize,
2595 y_stacked: &CudaSlice<f64>,
2596 delta_beta: &CudaSlice<f64>,
2597 u_stacked: &mut CudaSlice<f64>,
2598 ) -> Result<(), ArrowSchurGpuFailure> {
2599 let k_i = to_i32(k).ok_or(ArrowSchurGpuFailure::Unavailable)?;
2600 let rows_i = to_i32(rows).ok_or(ArrowSchurGpuFailure::Unavailable)?;
2601 let cfg = GemvConfig::<f64> {
2602 trans: cublasOperation_t::CUBLAS_OP_N,
2603 m: rows_i,
2604 n: k_i,
2605 alpha: 1.0,
2606 lda: rows_i,
2607 incx: 1,
2608 beta: 1.0,
2609 incy: 1,
2610 };
2611 unsafe { blas.gemv(cfg, y_stacked, delta_beta, u_stacked) }
2615 .map_err(|_| ArrowSchurGpuFailure::Unavailable)
2616 }
2617
2618 fn accumulate_back_sub_rhs(
2621 blas: &CudaBlas,
2622 d: usize,
2623 k: usize,
2624 n: usize,
2625 y_stack: &CudaSlice<f64>,
2626 delta_beta: &CudaSlice<f64>,
2627 u_stack: &mut CudaSlice<f64>,
2628 ) -> Result<(), ArrowSchurGpuFailure> {
2629 let y_block_elems = d * k;
2630 let u_block_elems = d;
2631 for i in 0..n {
2632 let y_slice = y_stack.slice(i * y_block_elems..(i + 1) * y_block_elems);
2633 let mut u_slice = u_stack.slice_mut(i * u_block_elems..(i + 1) * u_block_elems);
2634 let gemv_cfg = GemvConfig::<f64> {
2635 trans: cublasOperation_t::CUBLAS_OP_N,
2636 m: to_i32(d).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2637 n: to_i32(k).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2638 alpha: 1.0,
2639 lda: to_i32(d).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2640 incx: 1,
2641 beta: 1.0,
2642 incy: 1,
2643 };
2644 unsafe { blas.gemv(gemv_cfg, &y_slice, delta_beta, &mut u_slice) }
2647 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2648 }
2649 Ok(())
2650 }
2651
2652 use std::collections::HashMap;
2668 use std::sync::Mutex;
2669
2670 struct FusedModuleCache {
2675 modules: Mutex<
2676 HashMap<crate::gpu_kernels::arrow_schur_nvrtc::FusedModuleCacheKey, Arc<CudaModule>>,
2677 >,
2678 }
2679
2680 fn fused_module_cache() -> &'static FusedModuleCache {
2681 static CACHE: OnceLock<FusedModuleCache> = OnceLock::new();
2682 CACHE.get_or_init(|| FusedModuleCache {
2683 modules: Mutex::new(HashMap::new()),
2684 })
2685 }
2686
2687 fn fused_module_for(
2688 ctx: &Arc<CudaContext>,
2689 key: crate::gpu_kernels::arrow_schur_nvrtc::FusedModuleCacheKey,
2690 ) -> Result<Arc<CudaModule>, ArrowSchurGpuFailure> {
2691 let cache = fused_module_cache();
2692 if let Ok(guard) = cache.modules.lock() {
2693 if let Some(existing) = guard.get(&key) {
2694 return Ok(existing.clone());
2695 }
2696 }
2697 let src = crate::gpu_kernels::arrow_schur_nvrtc::forward_kernel_source(
2698 key.p_max as usize,
2699 key.r_template as usize,
2700 );
2701 let ptx = gam_gpu::device_cache::compile_ptx_arch(&src).map_err(|err| {
2702 ArrowSchurGpuFailure::SchurFactorFailed {
2703 reason: format!(
2704 "arrow-schur fused NVRTC compile (p_max={}, r={}): {err}",
2705 key.p_max, key.r_template
2706 ),
2707 }
2708 })?;
2709 let module = ctx
2710 .load_module(ptx)
2711 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2712 if let Ok(mut guard) = cache.modules.lock() {
2713 guard.entry(key).or_insert_with(|| module.clone());
2714 }
2715 Ok(module)
2716 }
2717
2718 const PCG_VECTOR_KERNEL_SOURCE: &str = r#"
2719extern "C" __global__ void arrow_pcg_jacobi_mul(
2720 const double* __restrict__ inv_diag,
2721 const double* __restrict__ r,
2722 double* __restrict__ z,
2723 int n
2724) {
2725 int idx = blockIdx.x * blockDim.x + threadIdx.x;
2726 if (idx < n) {
2727 z[idx] = inv_diag[idx] * r[idx];
2728 }
2729}
2730
2731extern "C" __global__ void arrow_pcg_update_p(
2732 const double* __restrict__ z,
2733 double beta,
2734 double* __restrict__ p,
2735 int n
2736) {
2737 int idx = blockIdx.x * blockDim.x + threadIdx.x;
2738 if (idx < n) {
2739 p[idx] = z[idx] + beta * p[idx];
2740 }
2741}
2742
2743extern "C" __global__ void arrow_sae_init(
2744 double* __restrict__ out,
2745 const double* __restrict__ x,
2746 double ridge,
2747 int n
2748) {
2749 int idx = blockIdx.x * blockDim.x + threadIdx.x;
2750 if (idx < n) {
2751 out[idx] = ridge * x[idx];
2752 }
2753}
2754
2755extern "C" __global__ void arrow_sae_smooth_matvec(
2756 const double* __restrict__ x,
2757 double* __restrict__ out,
2758 const int* __restrict__ block_offsets,
2759 const int* __restrict__ block_m,
2760 const int* __restrict__ factor_ptr,
2761 const double* __restrict__ factors,
2762 int p,
2763 int n_blocks
2764) {
2765 int block_id = blockIdx.y;
2766 int linear = blockIdx.x * blockDim.x + threadIdx.x;
2767 if (block_id >= n_blocks) {
2768 return;
2769 }
2770 int m = block_m[block_id];
2771 int total = m * p;
2772 if (linear >= total) {
2773 return;
2774 }
2775 int li = linear / p;
2776 int oc = linear - li * p;
2777 int off = block_offsets[block_id];
2778 int fbase = factor_ptr[block_id];
2779 double acc = 0.0;
2780 for (int lj = 0; lj < m; ++lj) {
2781 double a = factors[fbase + li * m + lj];
2782 acc += a * x[off + lj * p + oc];
2783 }
2784 out[off + li * p + oc] += acc;
2785}
2786
2787extern "C" __global__ void arrow_sae_sparse_g_matvec(
2788 const double* __restrict__ x,
2789 double* __restrict__ out,
2790 const int* __restrict__ row_off,
2791 const int* __restrict__ col_off,
2792 const int* __restrict__ rows,
2793 const int* __restrict__ cols,
2794 const int* __restrict__ data_ptr,
2795 const double* __restrict__ data,
2796 int p,
2797 int n_blocks
2798) {
2799 int block_id = blockIdx.y;
2800 int linear = blockIdx.x * blockDim.x + threadIdx.x;
2801 if (block_id >= n_blocks) {
2802 return;
2803 }
2804 int m_i = rows[block_id];
2805 int m_j = cols[block_id];
2806 int total = m_i * p;
2807 if (linear >= total) {
2808 return;
2809 }
2810 int li = linear / p;
2811 int oc = linear - li * p;
2812 int rbase = row_off[block_id];
2813 int cbase = col_off[block_id];
2814 int dbase = data_ptr[block_id];
2815 double acc = 0.0;
2816 for (int lj = 0; lj < m_j; ++lj) {
2817 acc += data[dbase + li * m_j + lj] * x[(cbase + lj) * p + oc];
2818 }
2819 // #1017 — a row atom co-occurs with multiple column atoms, so several
2820 // concurrent (atom_i, atom_j) blocks (blockIdx.y) write the SAME output
2821 // element `out[(rbase+li)*p+oc]`. A plain `+=` races and loses updates
2822 // (silently-wrong Schur matvec); accumulate atomically. `double` atomicAdd
2823 // needs sm_60+, guaranteed by the NVRTC arch pin (#1551).
2824 atomicAdd(&out[(rbase + li) * p + oc], acc);
2825}
2826
2827extern "C" __global__ void arrow_sae_gather_u(
2828 const double* __restrict__ x,
2829 const int* __restrict__ row_ptr,
2830 const int* __restrict__ beta_base,
2831 const double* __restrict__ phi,
2832 double* __restrict__ u,
2833 int p,
2834 int n_rows
2835) {
2836 int row = blockIdx.y;
2837 int oc = blockIdx.x * blockDim.x + threadIdx.x;
2838 if (row >= n_rows || oc >= p) {
2839 return;
2840 }
2841 double acc = 0.0;
2842 int start = row_ptr[row];
2843 int end = row_ptr[row + 1];
2844 for (int e = start; e < end; ++e) {
2845 acc += phi[e] * x[beta_base[e] + oc];
2846 }
2847 u[row * p + oc] = acc;
2848}
2849
2850extern "C" __global__ void arrow_sae_apply_l(
2851 const double* __restrict__ u,
2852 const int* __restrict__ jac_ptr,
2853 const double* __restrict__ jac,
2854 double* __restrict__ w,
2855 int p,
2856 int max_q,
2857 int n_rows
2858) {
2859 int row = blockIdx.y;
2860 int c = blockIdx.x * blockDim.x + threadIdx.x;
2861 if (row >= n_rows) {
2862 return;
2863 }
2864 int jstart = jac_ptr[row];
2865 int q = (jac_ptr[row + 1] - jstart) / p;
2866 if (c >= q) {
2867 return;
2868 }
2869 double acc = 0.0;
2870 for (int oc = 0; oc < p; ++oc) {
2871 acc += jac[jstart + c * p + oc] * u[row * p + oc];
2872 }
2873 w[row * max_q + c] = acc;
2874}
2875
2876extern "C" __global__ void arrow_sae_apply_ainv(
2877 const double* __restrict__ ainv,
2878 const double* __restrict__ w,
2879 double* __restrict__ v,
2880 int max_q,
2881 int n_rows
2882) {
2883 int row = blockIdx.y;
2884 int c = blockIdx.x * blockDim.x + threadIdx.x;
2885 if (row >= n_rows || c >= max_q) {
2886 return;
2887 }
2888 double acc = 0.0;
2889 int base = row * max_q * max_q;
2890 for (int j = 0; j < max_q; ++j) {
2891 acc += ainv[base + c * max_q + j] * w[row * max_q + j];
2892 }
2893 v[row * max_q + c] = acc;
2894}
2895
2896extern "C" __global__ void arrow_sae_scatter_sub(
2897 const double* __restrict__ v,
2898 const int* __restrict__ jac_ptr,
2899 const double* __restrict__ jac,
2900 const int* __restrict__ row_ptr,
2901 const int* __restrict__ beta_base,
2902 const double* __restrict__ phi,
2903 double* __restrict__ out,
2904 int p,
2905 int max_q,
2906 int n_rows
2907) {
2908 int row = blockIdx.y;
2909 int oc = blockIdx.x * blockDim.x + threadIdx.x;
2910 if (row >= n_rows || oc >= p) {
2911 return;
2912 }
2913 int jstart = jac_ptr[row];
2914 int q = (jac_ptr[row + 1] - jstart) / p;
2915 double lt_v = 0.0;
2916 for (int c = 0; c < q; ++c) {
2917 lt_v += jac[jstart + c * p + oc] * v[row * max_q + c];
2918 }
2919 int start = row_ptr[row];
2920 int end = row_ptr[row + 1];
2921 for (int e = start; e < end; ++e) {
2922 atomicAdd(&out[beta_base[e] + oc], -phi[e] * lt_v);
2923 }
2924}
2925
2926extern "C" __global__ void arrow_sae_diag_sub(
2927 double* __restrict__ diag,
2928 const double* __restrict__ ainv,
2929 const int* __restrict__ jac_ptr,
2930 const double* __restrict__ jac,
2931 const int* __restrict__ row_ptr,
2932 const int* __restrict__ beta_base,
2933 const double* __restrict__ phi,
2934 int p,
2935 int max_q,
2936 int n_rows
2937) {
2938 int row = blockIdx.y;
2939 int oc = blockIdx.x * blockDim.x + threadIdx.x;
2940 if (row >= n_rows || oc >= p) {
2941 return;
2942 }
2943 int jstart = jac_ptr[row];
2944 int q = (jac_ptr[row + 1] - jstart) / p;
2945 int abase = row * max_q * max_q;
2946 double quad = 0.0;
2947 for (int c = 0; c < q; ++c) {
2948 double lc = jac[jstart + c * p + oc];
2949 for (int d = 0; d < q; ++d) {
2950 quad += lc * ainv[abase + c * max_q + d] * jac[jstart + d * p + oc];
2951 }
2952 }
2953 int start = row_ptr[row];
2954 int end = row_ptr[row + 1];
2955 for (int e = start; e < end; ++e) {
2956 double pe = phi[e];
2957 atomicAdd(&diag[beta_base[e] + oc], -(pe * pe) * quad);
2958 }
2959}
2960
2961/* ── #1017/#1026 frames-engaged device kernels ─────────────────────────────
2962 * The factored β border is C-space (width Σ M_k·r_k). The penalty side is the
2963 * smooth `λ S_k ⊗ I_{r_k}` (per-block right-width r_k) plus the data-fit
2964 * `G_{ij} ⊗ W_{ij}` (W = U_iᵀU_j, dense r_i×r_j). The reduced-Schur term uses
2965 * the per-row DENSE cross-block H_tβ^(i) (q_i × border_dim, row-major). */
2966
2967extern "C" __global__ void arrow_sae_frame_smooth_matvec(
2968 const double* __restrict__ x,
2969 double* __restrict__ out,
2970 const int* __restrict__ block_offsets,
2971 const int* __restrict__ block_m,
2972 const int* __restrict__ block_r,
2973 const int* __restrict__ factor_ptr,
2974 const double* __restrict__ factors,
2975 int n_blocks
2976) {
2977 int block_id = blockIdx.y;
2978 int linear = blockIdx.x * blockDim.x + threadIdx.x;
2979 if (block_id >= n_blocks) {
2980 return;
2981 }
2982 int m = block_m[block_id];
2983 int r = block_r[block_id];
2984 int total = m * r;
2985 if (linear >= total) {
2986 return;
2987 }
2988 int li = linear / r;
2989 int ib = linear - li * r;
2990 int off = block_offsets[block_id];
2991 int fbase = factor_ptr[block_id];
2992 double acc = 0.0;
2993 for (int lj = 0; lj < m; ++lj) {
2994 double a = factors[fbase + li * m + lj];
2995 acc += a * x[off + lj * r + ib];
2996 }
2997 out[off + li * r + ib] += acc;
2998}
2999
3000extern "C" __global__ void arrow_sae_frame_g_matvec(
3001 const double* __restrict__ x,
3002 double* __restrict__ out,
3003 const int* __restrict__ off_i,
3004 const int* __restrict__ off_j,
3005 const int* __restrict__ r_i,
3006 const int* __restrict__ r_j,
3007 const int* __restrict__ m_i,
3008 const int* __restrict__ m_j,
3009 const int* __restrict__ g_ptr,
3010 const double* __restrict__ g_data,
3011 const int* __restrict__ w_ptr,
3012 const double* __restrict__ w_data,
3013 int n_blocks
3014) {
3015 int block_id = blockIdx.y;
3016 int linear = blockIdx.x * blockDim.x + threadIdx.x;
3017 if (block_id >= n_blocks) {
3018 return;
3019 }
3020 int ri = r_i[block_id];
3021 int rj = r_j[block_id];
3022 int mi = m_i[block_id];
3023 int mj = m_j[block_id];
3024 int total = mi * ri;
3025 if (linear >= total) {
3026 return;
3027 }
3028 int li = linear / ri; // basis row in atom i
3029 int a = linear - li * ri; // frame coord in atom i
3030 int oi = off_i[block_id];
3031 int oj = off_j[block_id];
3032 int gbase = g_ptr[block_id];
3033 int wbase = w_ptr[block_id];
3034 double acc = 0.0;
3035 for (int lj = 0; lj < mj; ++lj) {
3036 double g = g_data[gbase + li * mj + lj];
3037 if (g == 0.0) { continue; }
3038 int xj_base = oj + lj * rj;
3039 double inner = 0.0;
3040 for (int b = 0; b < rj; ++b) {
3041 inner += w_data[wbase + a * rj + b] * x[xj_base + b];
3042 }
3043 acc += g * inner;
3044 }
3045 // #1017 — same race as `arrow_sae_sparse_g_matvec`: atom i is the row atom of
3046 // multiple co-occurring (i,j) frame blocks running concurrently on
3047 // blockIdx.y, all targeting `out[oi+li*ri+a]`. Accumulate atomically so the
3048 // framed G⊗W matvec is correct (the CPU oracle sums these sequentially).
3049 atomicAdd(&out[oi + li * ri + a], acc);
3050}
3051
3052/* Per-row reduced-Schur subtraction with a DENSE cross-block H_tβ^(i).
3053 * h_i = H_tβ^(i) · x (length q_i)
3054 * s_i = (H_tt^(i)+ρ_t I)⁻¹ h_i (apply cached ainv, length q_i)
3055 * out -= (H_tβ^(i))ᵀ · s_i (scatter into border_dim)
3056 * `htb` is row-major (q_i × k) flattened, `htb_ptr` gives each row's base and
3057 * (htb_ptr[row+1]-htb_ptr[row])/k == q_i. `q_of` carries q_i directly. */
3058extern "C" __global__ void arrow_sae_frame_apply_h(
3059 const double* __restrict__ x,
3060 const int* __restrict__ htb_ptr,
3061 const double* __restrict__ htb,
3062 const int* __restrict__ q_of,
3063 double* __restrict__ hvec,
3064 int k,
3065 int max_q,
3066 int n_rows
3067) {
3068 int row = blockIdx.y;
3069 int c = blockIdx.x * blockDim.x + threadIdx.x;
3070 if (row >= n_rows) { return; }
3071 int q = q_of[row];
3072 if (c >= q) { return; }
3073 int base = htb_ptr[row] + c * k;
3074 double acc = 0.0;
3075 for (int a = 0; a < k; ++a) {
3076 acc += htb[base + a] * x[a];
3077 }
3078 hvec[row * max_q + c] = acc;
3079}
3080
3081extern "C" __global__ void arrow_sae_frame_apply_ainv(
3082 const double* __restrict__ ainv,
3083 const double* __restrict__ hvec,
3084 const int* __restrict__ q_of,
3085 double* __restrict__ svec,
3086 int max_q,
3087 int n_rows
3088) {
3089 int row = blockIdx.y;
3090 int c = blockIdx.x * blockDim.x + threadIdx.x;
3091 if (row >= n_rows || c >= max_q) { return; }
3092 int q = q_of[row];
3093 double acc = 0.0;
3094 int abase = row * max_q * max_q;
3095 for (int j = 0; j < q; ++j) {
3096 acc += ainv[abase + c * max_q + j] * hvec[row * max_q + j];
3097 }
3098 svec[row * max_q + c] = acc;
3099}
3100
3101extern "C" __global__ void arrow_sae_frame_scatter_h(
3102 const double* __restrict__ svec,
3103 const int* __restrict__ htb_ptr,
3104 const double* __restrict__ htb,
3105 const int* __restrict__ q_of,
3106 double* __restrict__ out,
3107 int k,
3108 int max_q,
3109 int n_rows
3110) {
3111 int row = blockIdx.y;
3112 int a = blockIdx.x * blockDim.x + threadIdx.x;
3113 if (row >= n_rows || a >= k) { return; }
3114 int q = q_of[row];
3115 int hbase = htb_ptr[row];
3116 double acc = 0.0;
3117 for (int c = 0; c < q; ++c) {
3118 acc += htb[hbase + c * k + a] * svec[row * max_q + c];
3119 }
3120 atomicAdd(&out[a], -acc);
3121}
3122
3123/* #1017 evidence-lane DETERMINISTIC reduced-Schur scatter:
3124 out[a] = -Σ_i Σ_c H_tβ[i][c,a]·svec[i,c]. One thread owns output coord `a` and
3125 sums the rows in fixed index order 0..n_rows — NO atomics, so the result is
3126 run-to-run bit-stable (the SLQ log|S| determinism contract). Same arithmetic as
3127 arrow_sae_frame_scatter_h; only the reduction order is pinned (there the sum
3128 over rows is an atomicAdd race). The shared atomic kernel is left untouched —
3129 the step-PCG relies on it. `out` is fully assigned (no init needed). */
3130extern "C" __global__ void arrow_sae_frame_scatter_h_det(
3131 const double* __restrict__ svec,
3132 const int* __restrict__ htb_ptr,
3133 const double* __restrict__ htb,
3134 const int* __restrict__ q_of,
3135 double* __restrict__ out,
3136 int k,
3137 int max_q,
3138 int n_rows
3139) {
3140 int a = blockIdx.x * blockDim.x + threadIdx.x;
3141 if (a >= k) { return; }
3142 double acc = 0.0;
3143 for (int row = 0; row < n_rows; ++row) {
3144 int q = q_of[row];
3145 int hbase = htb_ptr[row];
3146 int sbase = row * max_q;
3147 for (int c = 0; c < q; ++c) {
3148 acc += htb[hbase + c * k + a] * svec[sbase + c];
3149 }
3150 }
3151 out[a] = -acc;
3152}
3153
3154/* #1017 evidence-lane 2-STAGE deterministic scatter, STAGE 1 (partials):
3155 partial[chunk][a] = Σ_{row∈chunk} Σ_c H_tβ[row][c,a]·svec[row,c], for the
3156 contiguous row range [chunk·rows_per_chunk, …). grid = (⌈k/256⌉, n_chunks);
3157 thread owns (a, chunk). Replaces the single-strip `arrow_sae_frame_scatter_h_det`
3158 (⌈k/256⌉ CTAs — only 4 at k=911, one thread serial over ALL n_rows, ~94% of a
3159 72-SM A10 idle) with ⌈k/256⌉·n_chunks CTAs. Rows are summed in fixed order
3160 within the chunk and the chunks are reduced in order by stage 2, so the result
3161 is a FIXED reassociation of the same ordered row sum — run-to-run bit-stable
3162 (the SLQ log|S| determinism contract) and within the ≤1e-9 CPU-oracle gate. */
3163extern "C" __global__ void arrow_sae_frame_scatter_h_det_partial(
3164 const double* __restrict__ svec,
3165 const int* __restrict__ htb_ptr,
3166 const double* __restrict__ htb,
3167 const int* __restrict__ q_of,
3168 double* __restrict__ partial,
3169 int k,
3170 int max_q,
3171 int n_rows,
3172 int rows_per_chunk
3173) {
3174 int a = blockIdx.x * blockDim.x + threadIdx.x;
3175 if (a >= k) { return; }
3176 int chunk = blockIdx.y;
3177 int row0 = chunk * rows_per_chunk;
3178 int row1 = row0 + rows_per_chunk;
3179 if (row1 > n_rows) { row1 = n_rows; }
3180 double acc = 0.0;
3181 for (int row = row0; row < row1; ++row) {
3182 int q = q_of[row];
3183 int hbase = htb_ptr[row];
3184 int sbase = row * max_q;
3185 for (int c = 0; c < q; ++c) {
3186 acc += htb[hbase + c * k + a] * svec[sbase + c];
3187 }
3188 }
3189 partial[(long long)chunk * k + a] = acc;
3190}
3191
3192/* #1017 STAGE 2 (reduce): out[a] = -Σ_chunk partial[chunk][a], chunks summed in
3193 fixed order 0..n_chunks. One thread per output coord `a`; ⌈k/256⌉ CTAs. */
3194extern "C" __global__ void arrow_sae_frame_scatter_h_det_reduce(
3195 const double* __restrict__ partial,
3196 double* __restrict__ out,
3197 int k,
3198 int n_chunks
3199) {
3200 int a = blockIdx.x * blockDim.x + threadIdx.x;
3201 if (a >= k) { return; }
3202 double acc = 0.0;
3203 for (int chunk = 0; chunk < n_chunks; ++chunk) {
3204 acc += partial[(long long)chunk * k + a];
3205 }
3206 out[a] = -acc;
3207}
3208
3209/* #1017 evidence-lane WARP-COOPERATIVE apply_h: hvec[i][c] = Σ_a H_tβ[i][c,a]·x[a].
3210 One WARP owns (row, c): lane `l` strides `a = l, l+32, …` over the contiguous
3211 `H_tβ[i][c,·]` slab (fully coalesced across the warp) and a fixed-order
3212 __shfl_down tree reduces to lane 0. Replaces the shared `arrow_sae_frame_apply_h`
3213 (256-thread block, only `q_i` threads active, stride-`k` uncoalesced reads) on
3214 the evidence path ONLY — the shared kernel is untouched (step-PCG relies on it).
3215 Reduction order is fixed (lane stride + tree), so the result is run-to-run
3216 bit-stable; it differs from the scalar kernel only by ULP reassociation, within
3217 the ≤1e-9 parity gate. Launch: block = max_q·32 (≤1024), grid.x = n_rows;
3218 warp `w = threadIdx.x/32` handles `c = w` for `w < q_i`. */
3219extern "C" __global__ void arrow_sae_frame_apply_h_warp(
3220 const double* __restrict__ x,
3221 const int* __restrict__ htb_ptr,
3222 const double* __restrict__ htb,
3223 const int* __restrict__ q_of,
3224 double* __restrict__ hvec,
3225 int k,
3226 int max_q,
3227 int n_rows
3228) {
3229 int row = blockIdx.x;
3230 if (row >= n_rows) { return; }
3231 int warp = threadIdx.x >> 5;
3232 int lane = threadIdx.x & 31;
3233 int q = q_of[row];
3234 if (warp >= q) { return; }
3235 int base = htb_ptr[row] + warp * k;
3236 double acc = 0.0;
3237 for (int a = lane; a < k; a += 32) {
3238 acc += htb[base + a] * x[a];
3239 }
3240 #pragma unroll
3241 for (int off = 16; off > 0; off >>= 1) {
3242 acc += __shfl_down_sync(0xffffffffu, acc, off);
3243 }
3244 if (lane == 0) {
3245 hvec[row * max_q + warp] = acc;
3246 }
3247}
3248
3249/* Frame Jacobi diagonal subtraction: diag[a] -= Σ_c Σ_d H_tβ[c,a]·ainv[c,d]·H_tβ[d,a]. */
3250extern "C" __global__ void arrow_sae_frame_diag_sub(
3251 double* __restrict__ diag,
3252 const double* __restrict__ ainv,
3253 const int* __restrict__ htb_ptr,
3254 const double* __restrict__ htb,
3255 const int* __restrict__ q_of,
3256 int k,
3257 int max_q,
3258 int n_rows
3259) {
3260 int row = blockIdx.y;
3261 int a = blockIdx.x * blockDim.x + threadIdx.x;
3262 if (row >= n_rows || a >= k) { return; }
3263 int q = q_of[row];
3264 int hbase = htb_ptr[row];
3265 int abase = row * max_q * max_q;
3266 double quad = 0.0;
3267 for (int c = 0; c < q; ++c) {
3268 double hc = htb[hbase + c * k + a];
3269 for (int d = 0; d < q; ++d) {
3270 quad += hc * ainv[abase + c * max_q + d] * htb[hbase + d * k + a];
3271 }
3272 }
3273 atomicAdd(&diag[a], -quad);
3274}
3275"#;
3276
3277 fn pcg_vector_module(
3278 ctx: &Arc<CudaContext>,
3279 ) -> Result<&'static Arc<CudaModule>, ArrowSchurGpuFailure> {
3280 static CACHE: gam_gpu::device_cache::PtxModuleCache =
3281 gam_gpu::device_cache::PtxModuleCache::new();
3282 CACHE
3283 .get_or_compile(ctx, "arrow_pcg_vector", PCG_VECTOR_KERNEL_SOURCE)
3284 .map_err(|err| {
3285 log::warn!("[#1551] pcg_vector_module get_or_compile failed: {err}");
3291 ArrowSchurGpuFailure::Unavailable
3292 })
3293 }
3294
3295 fn pcg_launch_config(n: usize) -> Result<LaunchConfig, ArrowSchurGpuFailure> {
3296 let threads = 256u32;
3297 let blocks = ((n as u32).saturating_add(threads - 1) / threads).max(1);
3298 Ok(LaunchConfig {
3299 grid_dim: (blocks, 1, 1),
3300 block_dim: (threads, 1, 1),
3301 shared_mem_bytes: 0,
3302 })
3303 }
3304
3305 fn launch_jacobi_mul(
3306 stream: &Arc<CudaStream>,
3307 module: &Arc<CudaModule>,
3308 inv_diag: &CudaSlice<f64>,
3309 r: &CudaSlice<f64>,
3310 z: &mut CudaSlice<f64>,
3311 n: usize,
3312 ) -> Result<(), ArrowSchurGpuFailure> {
3313 let kernel = module
3314 .load_function("arrow_pcg_jacobi_mul")
3315 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
3316 let n_i32 = to_i32(n).ok_or(ArrowSchurGpuFailure::Unavailable)?;
3317 let mut builder = stream.launch_builder(&kernel);
3318 builder.arg(inv_diag).arg(r).arg(z).arg(&n_i32);
3319 unsafe { builder.launch(pcg_launch_config(n)?) }
3322 .map(drop)
3323 .map_err(|_| ArrowSchurGpuFailure::Unavailable)
3324 }
3325
3326 fn launch_update_p(
3327 stream: &Arc<CudaStream>,
3328 module: &Arc<CudaModule>,
3329 z: &CudaSlice<f64>,
3330 beta: f64,
3331 p: &mut CudaSlice<f64>,
3332 n: usize,
3333 ) -> Result<(), ArrowSchurGpuFailure> {
3334 let kernel = module
3335 .load_function("arrow_pcg_update_p")
3336 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
3337 let n_i32 = to_i32(n).ok_or(ArrowSchurGpuFailure::Unavailable)?;
3338 let mut builder = stream.launch_builder(&kernel);
3339 builder.arg(z).arg(&beta).arg(p).arg(&n_i32);
3340 unsafe { builder.launch(pcg_launch_config(n)?) }
3343 .map(drop)
3344 .map_err(|_| ArrowSchurGpuFailure::Unavailable)
3345 }
3346
3347 struct DeviceSaePcgBuffers {
3348 row_ptr: CudaSlice<i32>,
3349 beta_base: CudaSlice<i32>,
3350 phi: CudaSlice<f64>,
3351 jac_ptr: CudaSlice<i32>,
3352 jac: CudaSlice<f64>,
3353 smooth_offsets: CudaSlice<i32>,
3354 smooth_m: CudaSlice<i32>,
3355 smooth_ptr: CudaSlice<i32>,
3356 smooth_data: CudaSlice<f64>,
3357 g_row_off: CudaSlice<i32>,
3358 g_col_off: CudaSlice<i32>,
3359 g_rows: CudaSlice<i32>,
3360 g_cols: CudaSlice<i32>,
3361 g_ptr: CudaSlice<i32>,
3362 g_data: CudaSlice<f64>,
3363 ainv: CudaSlice<f64>,
3364 u: CudaSlice<f64>,
3365 w: CudaSlice<f64>,
3366 v: CudaSlice<f64>,
3367 n_rows: usize,
3368 p: usize,
3369 k: usize,
3370 max_q: usize,
3371 smooth_blocks: usize,
3372 g_blocks: usize,
3373 }
3374
3375 fn checked_i32(value: usize) -> Result<i32, ArrowSchurGpuFailure> {
3376 to_i32(value).ok_or(ArrowSchurGpuFailure::Unavailable)
3377 }
3378
3379 fn sae_penalty_diag_host(
3380 data: &DeviceSaePcgData,
3381 ridge_beta: f64,
3382 ) -> Result<Vec<f64>, ArrowSchurGpuFailure> {
3383 let mut diag = vec![ridge_beta; data.beta_dim];
3384 for block in &data.smooth_blocks {
3385 let (rows, cols) = block.factor_a.dim();
3386 if rows != cols {
3387 return Err(ArrowSchurGpuFailure::Unavailable);
3388 }
3389 for row in 0..rows {
3390 let coeff = block.factor_a[[row, row]];
3391 let base = block
3392 .global_offset
3393 .checked_add(
3394 row.checked_mul(data.p)
3395 .ok_or(ArrowSchurGpuFailure::Unavailable)?,
3396 )
3397 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
3398 let end = base
3399 .checked_add(data.p)
3400 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
3401 if end > diag.len() {
3402 return Err(ArrowSchurGpuFailure::Unavailable);
3403 }
3404 for channel in 0..data.p {
3405 diag[base + channel] += coeff;
3406 }
3407 }
3408 }
3409 for block in &data.sparse_g_blocks {
3410 if block.row_off != block.col_off {
3411 continue;
3412 }
3413 let (rows, cols) = block.data.dim();
3414 for row in 0..rows.min(cols) {
3415 let coeff = block.data[[row, row]];
3416 let beta_row = block
3417 .row_off
3418 .checked_add(row)
3419 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
3420 let base = beta_row
3421 .checked_mul(data.p)
3422 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
3423 let end = base
3424 .checked_add(data.p)
3425 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
3426 if end > diag.len() {
3427 return Err(ArrowSchurGpuFailure::Unavailable);
3428 }
3429 for channel in 0..data.p {
3430 diag[base + channel] += coeff;
3431 }
3432 }
3433 }
3434 Ok(diag)
3435 }
3436
3437 fn flatten_device_sae_data(
3438 sys: &ArrowSchurSystem,
3439 data: &DeviceSaePcgData,
3440 ridge_t: f64,
3441 stream: &Arc<CudaStream>,
3442 ) -> Result<DeviceSaePcgBuffers, ArrowSchurGpuFailure> {
3443 let n_rows = sys.rows.len();
3444 let p = data.p;
3445 let k = data.beta_dim;
3446 if data.a_phi.len() != n_rows || data.local_jac.len() != n_rows {
3447 return Err(ArrowSchurGpuFailure::Unavailable);
3448 }
3449
3450 let mut row_ptr_host = Vec::with_capacity(n_rows + 1);
3451 let mut beta_base_host = Vec::<i32>::new();
3452 let mut phi_host = Vec::<f64>::new();
3453 row_ptr_host.push(0_i32);
3454 for row in data.a_phi.iter() {
3455 for &(base, phi) in row {
3456 beta_base_host.push(checked_i32(base)?);
3457 phi_host.push(phi);
3458 }
3459 row_ptr_host.push(checked_i32(beta_base_host.len())?);
3460 }
3461
3462 let mut jac_ptr_host = Vec::with_capacity(n_rows + 1);
3463 let mut jac_host = Vec::<f64>::new();
3464 let mut max_q = 0usize;
3465 jac_ptr_host.push(0_i32);
3466 for row_jac in data.local_jac.iter() {
3467 if row_jac.len() % p != 0 {
3468 return Err(ArrowSchurGpuFailure::Unavailable);
3469 }
3470 max_q = max_q.max(row_jac.len() / p);
3471 jac_host.extend_from_slice(row_jac);
3472 jac_ptr_host.push(checked_i32(jac_host.len())?);
3473 }
3474 if max_q == 0 {
3475 return Err(ArrowSchurGpuFailure::Unavailable);
3476 }
3477
3478 let mut smooth_offsets_host = Vec::with_capacity(data.smooth_blocks.len());
3479 let mut smooth_m_host = Vec::with_capacity(data.smooth_blocks.len());
3480 let mut smooth_ptr_host = Vec::with_capacity(data.smooth_blocks.len() + 1);
3481 let mut smooth_data_host = Vec::<f64>::new();
3482 smooth_ptr_host.push(0_i32);
3483 for block in &data.smooth_blocks {
3484 let (rows, cols) = block.factor_a.dim();
3485 if rows != cols {
3486 return Err(ArrowSchurGpuFailure::Unavailable);
3487 }
3488 smooth_offsets_host.push(checked_i32(block.global_offset)?);
3489 smooth_m_host.push(checked_i32(rows)?);
3490 for r in 0..rows {
3491 for c in 0..cols {
3492 smooth_data_host.push(block.factor_a[[r, c]]);
3493 }
3494 }
3495 smooth_ptr_host.push(checked_i32(smooth_data_host.len())?);
3496 }
3497
3498 let mut g_row_off_host = Vec::with_capacity(data.sparse_g_blocks.len());
3499 let mut g_col_off_host = Vec::with_capacity(data.sparse_g_blocks.len());
3500 let mut g_rows_host = Vec::with_capacity(data.sparse_g_blocks.len());
3501 let mut g_cols_host = Vec::with_capacity(data.sparse_g_blocks.len());
3502 let mut g_ptr_host = Vec::with_capacity(data.sparse_g_blocks.len() + 1);
3503 let mut g_data_host = Vec::<f64>::new();
3504 g_ptr_host.push(0_i32);
3505 for block in &data.sparse_g_blocks {
3506 let (rows, cols) = block.data.dim();
3507 g_row_off_host.push(checked_i32(block.row_off)?);
3508 g_col_off_host.push(checked_i32(block.col_off)?);
3509 g_rows_host.push(checked_i32(rows)?);
3510 g_cols_host.push(checked_i32(cols)?);
3511 for r in 0..rows {
3512 for c in 0..cols {
3513 g_data_host.push(block.data[[r, c]]);
3514 }
3515 }
3516 g_ptr_host.push(checked_i32(g_data_host.len())?);
3517 }
3518
3519 let mut ainv_host = vec![0.0_f64; n_rows * max_q * max_q];
3520 for (row_idx, row) in sys.rows.iter().enumerate() {
3521 let q = data.local_jac[row_idx].len() / p;
3522 if row.htt.dim() != (q, q) {
3523 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
3524 reason: format!(
3525 "SAE device PCG row {row_idx}: H_tt shape {:?} != ({q}, {q})",
3526 row.htt.dim()
3527 ),
3528 });
3529 }
3530 let mut block = row.htt.clone();
3531 for d in 0..q {
3532 block[[d, d]] += ridge_t;
3533 }
3534 let factor = gam_linalg::triangular::cholesky_factor_in_place(
3535 block.view(),
3536 gam_linalg::triangular::CholeskyGuard::NonnegativePivot,
3537 )
3538 .ok_or_else(|| {
3539 ArrowSchurGpuFailure::RidgeBumpRequired {
3542 row: row_idx,
3543 bump: super::ridge_bump_to_make_pd(row.htt.view(), ridge_t),
3544 }
3545 })?;
3546 for col in 0..q {
3547 let mut e = Array1::<f64>::zeros(q);
3548 e[col] = 1.0;
3549 let solved = gam_linalg::triangular::cholesky_solve_vector(factor.view(), e.view());
3550 for r in 0..q {
3551 ainv_host[row_idx * max_q * max_q + r * max_q + col] = solved[r];
3552 }
3553 }
3554 }
3555
3556 Ok(DeviceSaePcgBuffers {
3557 row_ptr: stream
3558 .clone_htod(&row_ptr_host)
3559 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
3560 beta_base: stream
3561 .clone_htod(&beta_base_host)
3562 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
3563 phi: stream
3564 .clone_htod(&phi_host)
3565 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
3566 jac_ptr: stream
3567 .clone_htod(&jac_ptr_host)
3568 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
3569 jac: stream
3570 .clone_htod(&jac_host)
3571 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
3572 smooth_offsets: stream
3573 .clone_htod(&smooth_offsets_host)
3574 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
3575 smooth_m: stream
3576 .clone_htod(&smooth_m_host)
3577 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
3578 smooth_ptr: stream
3579 .clone_htod(&smooth_ptr_host)
3580 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
3581 smooth_data: stream
3582 .clone_htod(&smooth_data_host)
3583 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
3584 g_row_off: stream
3585 .clone_htod(&g_row_off_host)
3586 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
3587 g_col_off: stream
3588 .clone_htod(&g_col_off_host)
3589 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
3590 g_rows: stream
3591 .clone_htod(&g_rows_host)
3592 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
3593 g_cols: stream
3594 .clone_htod(&g_cols_host)
3595 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
3596 g_ptr: stream
3597 .clone_htod(&g_ptr_host)
3598 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
3599 g_data: stream
3600 .clone_htod(&g_data_host)
3601 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
3602 ainv: stream
3603 .clone_htod(&ainv_host)
3604 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
3605 u: stream
3606 .alloc_zeros::<f64>(n_rows * p)
3607 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
3608 w: stream
3609 .alloc_zeros::<f64>(n_rows * max_q)
3610 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
3611 v: stream
3612 .alloc_zeros::<f64>(n_rows * max_q)
3613 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
3614 n_rows,
3615 p,
3616 k,
3617 max_q,
3618 smooth_blocks: data.smooth_blocks.len(),
3619 g_blocks: data.sparse_g_blocks.len(),
3620 })
3621 }
3622
3623 fn launch_sae_init(
3624 stream: &Arc<CudaStream>,
3625 module: &Arc<CudaModule>,
3626 out: &mut CudaSlice<f64>,
3627 x: &CudaSlice<f64>,
3628 ridge: f64,
3629 n: usize,
3630 ) -> Result<(), ArrowSchurGpuFailure> {
3631 let kernel = module
3632 .load_function("arrow_sae_init")
3633 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
3634 let n_i32 = checked_i32(n)?;
3635 let mut builder = stream.launch_builder(&kernel);
3636 builder.arg(out).arg(x).arg(&ridge).arg(&n_i32);
3637 unsafe { builder.launch(pcg_launch_config(n)?) }
3641 .map(drop)
3642 .map_err(|_| ArrowSchurGpuFailure::Unavailable)
3643 }
3644
3645 fn launch_sae_penalty_matvec(
3646 stream: &Arc<CudaStream>,
3647 module: &Arc<CudaModule>,
3648 buffers: &mut DeviceSaePcgBuffers,
3649 x: &CudaSlice<f64>,
3650 out: &mut CudaSlice<f64>,
3651 ridge_beta: f64,
3652 ) -> Result<(), ArrowSchurGpuFailure> {
3653 launch_sae_init(stream, module, out, x, ridge_beta, buffers.k)?;
3654 if buffers.smooth_blocks > 0 {
3655 let kernel = module
3656 .load_function("arrow_sae_smooth_matvec")
3657 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
3658 let max_m = buffers.k;
3659 let p_i32 = checked_i32(buffers.p)?;
3660 let blocks_i32 = checked_i32(buffers.smooth_blocks)?;
3661 let cfg = LaunchConfig {
3662 grid_dim: (
3663 ((max_m as u32).saturating_add(255) / 256).max(1),
3664 checked_i32(buffers.smooth_blocks)? as u32,
3665 1,
3666 ),
3667 block_dim: (256, 1, 1),
3668 shared_mem_bytes: 0,
3669 };
3670 let mut builder = stream.launch_builder(&kernel);
3671 builder
3672 .arg(x)
3673 .arg(&mut *out)
3674 .arg(&buffers.smooth_offsets)
3675 .arg(&buffers.smooth_m)
3676 .arg(&buffers.smooth_ptr)
3677 .arg(&buffers.smooth_data)
3678 .arg(&p_i32)
3679 .arg(&blocks_i32);
3680 unsafe { builder.launch(cfg) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
3685 }
3686 if buffers.g_blocks > 0 {
3687 let kernel = module
3688 .load_function("arrow_sae_sparse_g_matvec")
3689 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
3690 let max_work = buffers
3691 .k
3692 .checked_div(buffers.p)
3693 .unwrap_or(0)
3694 .saturating_mul(buffers.p);
3695 let p_i32 = checked_i32(buffers.p)?;
3696 let blocks_i32 = checked_i32(buffers.g_blocks)?;
3697 let cfg = LaunchConfig {
3698 grid_dim: (
3699 ((max_work as u32).saturating_add(255) / 256).max(1),
3700 checked_i32(buffers.g_blocks)? as u32,
3701 1,
3702 ),
3703 block_dim: (256, 1, 1),
3704 shared_mem_bytes: 0,
3705 };
3706 let mut builder = stream.launch_builder(&kernel);
3707 builder
3708 .arg(x)
3709 .arg(&mut *out)
3710 .arg(&buffers.g_row_off)
3711 .arg(&buffers.g_col_off)
3712 .arg(&buffers.g_rows)
3713 .arg(&buffers.g_cols)
3714 .arg(&buffers.g_ptr)
3715 .arg(&buffers.g_data)
3716 .arg(&p_i32)
3717 .arg(&blocks_i32);
3718 unsafe { builder.launch(cfg) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
3723 }
3724 Ok(())
3725 }
3726
3727 fn launch_sae_row_schur_sub(
3728 stream: &Arc<CudaStream>,
3729 module: &Arc<CudaModule>,
3730 buffers: &mut DeviceSaePcgBuffers,
3731 x: &CudaSlice<f64>,
3732 out: &mut CudaSlice<f64>,
3733 ) -> Result<(), ArrowSchurGpuFailure> {
3734 let p_i32 = checked_i32(buffers.p)?;
3735 let max_q_i32 = checked_i32(buffers.max_q)?;
3736 let n_rows_i32 = checked_i32(buffers.n_rows)?;
3737 let cfg_p_rows = LaunchConfig {
3738 grid_dim: (
3739 ((buffers.p as u32).saturating_add(255) / 256).max(1),
3740 checked_i32(buffers.n_rows)? as u32,
3741 1,
3742 ),
3743 block_dim: (256, 1, 1),
3744 shared_mem_bytes: 0,
3745 };
3746 let gather = module
3747 .load_function("arrow_sae_gather_u")
3748 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
3749 {
3750 let mut builder = stream.launch_builder(&gather);
3751 builder
3752 .arg(x)
3753 .arg(&buffers.row_ptr)
3754 .arg(&buffers.beta_base)
3755 .arg(&buffers.phi)
3756 .arg(&mut buffers.u)
3757 .arg(&p_i32)
3758 .arg(&n_rows_i32);
3759 unsafe { builder.launch(cfg_p_rows) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
3763 }
3764
3765 let cfg_q_rows = LaunchConfig {
3766 grid_dim: (
3767 ((buffers.max_q as u32).saturating_add(255) / 256).max(1),
3768 checked_i32(buffers.n_rows)? as u32,
3769 1,
3770 ),
3771 block_dim: (256, 1, 1),
3772 shared_mem_bytes: 0,
3773 };
3774 let apply_l = module
3775 .load_function("arrow_sae_apply_l")
3776 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
3777 {
3778 let mut builder = stream.launch_builder(&apply_l);
3779 builder
3780 .arg(&buffers.u)
3781 .arg(&buffers.jac_ptr)
3782 .arg(&buffers.jac)
3783 .arg(&mut buffers.w)
3784 .arg(&p_i32)
3785 .arg(&max_q_i32)
3786 .arg(&n_rows_i32);
3787 unsafe { builder.launch(cfg_q_rows) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
3791 }
3792
3793 let apply_ainv = module
3794 .load_function("arrow_sae_apply_ainv")
3795 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
3796 {
3797 let mut builder = stream.launch_builder(&apply_ainv);
3798 builder
3799 .arg(&buffers.ainv)
3800 .arg(&buffers.w)
3801 .arg(&mut buffers.v)
3802 .arg(&max_q_i32)
3803 .arg(&n_rows_i32);
3804 unsafe { builder.launch(cfg_q_rows) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
3808 }
3809
3810 let scatter = module
3811 .load_function("arrow_sae_scatter_sub")
3812 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
3813 {
3814 let mut builder = stream.launch_builder(&scatter);
3815 builder
3816 .arg(&buffers.v)
3817 .arg(&buffers.jac_ptr)
3818 .arg(&buffers.jac)
3819 .arg(&buffers.row_ptr)
3820 .arg(&buffers.beta_base)
3821 .arg(&buffers.phi)
3822 .arg(out)
3823 .arg(&p_i32)
3824 .arg(&max_q_i32)
3825 .arg(&n_rows_i32);
3826 unsafe { builder.launch(cfg_p_rows) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
3830 }
3831 Ok(())
3832 }
3833
3834 fn launch_sae_diag_sub(
3835 stream: &Arc<CudaStream>,
3836 module: &Arc<CudaModule>,
3837 buffers: &DeviceSaePcgBuffers,
3838 diag: &mut CudaSlice<f64>,
3839 ) -> Result<(), ArrowSchurGpuFailure> {
3840 let kernel = module
3841 .load_function("arrow_sae_diag_sub")
3842 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
3843 let p_i32 = checked_i32(buffers.p)?;
3844 let max_q_i32 = checked_i32(buffers.max_q)?;
3845 let n_rows_i32 = checked_i32(buffers.n_rows)?;
3846 let cfg = LaunchConfig {
3847 grid_dim: (
3848 ((buffers.p as u32).saturating_add(255) / 256).max(1),
3849 checked_i32(buffers.n_rows)? as u32,
3850 1,
3851 ),
3852 block_dim: (256, 1, 1),
3853 shared_mem_bytes: 0,
3854 };
3855 let mut builder = stream.launch_builder(&kernel);
3856 builder
3857 .arg(diag)
3858 .arg(&buffers.ainv)
3859 .arg(&buffers.jac_ptr)
3860 .arg(&buffers.jac)
3861 .arg(&buffers.row_ptr)
3862 .arg(&buffers.beta_base)
3863 .arg(&buffers.phi)
3864 .arg(&p_i32)
3865 .arg(&max_q_i32)
3866 .arg(&n_rows_i32);
3867 unsafe { builder.launch(cfg) }
3871 .map(drop)
3872 .map_err(|_| ArrowSchurGpuFailure::Unavailable)
3873 }
3874
3875 fn launch_sae_matvec(
3876 stream: &Arc<CudaStream>,
3877 module: &Arc<CudaModule>,
3878 buffers: &mut DeviceSaePcgBuffers,
3879 x: &CudaSlice<f64>,
3880 out: &mut CudaSlice<f64>,
3881 ridge_beta: f64,
3882 ) -> Result<(), ArrowSchurGpuFailure> {
3883 launch_sae_penalty_matvec(stream, module, buffers, x, out, ridge_beta)?;
3884 launch_sae_row_schur_sub(stream, module, buffers, x, out)
3885 }
3886
3887 fn pack_fused_host(
3892 sys: &ArrowSchurSystem,
3893 ridge_t: f64,
3894 p_max: usize,
3895 r_template: usize,
3896 ) -> (Vec<f64>, Vec<f64>, Vec<f64>) {
3897 let n = sys.rows.len();
3898 let d = sys.d;
3899 let k = sys.k;
3900 let mut d_buf = vec![0.0_f64; n * p_max * p_max];
3901 let mut b_buf = vec![0.0_f64; n * p_max * r_template];
3902 let mut g_buf = vec![0.0_f64; n * p_max];
3903 for (i, row) in sys.rows.iter().enumerate() {
3904 for col in 0..d {
3906 let base = (i * p_max + col) * p_max;
3907 for r in 0..d {
3908 let mut value = row.htt[[r, col]];
3909 if r == col {
3910 value += ridge_t;
3911 }
3912 d_buf[base + r] = value;
3913 }
3914 }
3915 for col in 0..k {
3923 let base = (i * r_template + col) * p_max;
3924 for r in 0..d {
3925 b_buf[base + r] = row.htbeta[[r, col]];
3926 }
3927 }
3928 let g_base = i * p_max;
3930 for r in 0..d {
3931 g_buf[g_base + r] = row.gt[r];
3932 }
3933 }
3934 (d_buf, b_buf, g_buf)
3935 }
3936
3937 pub(super) struct ResidentArrowFrame {
3964 n: usize,
3965 d: usize,
3966 k: usize,
3967 stream: Arc<CudaStream>,
3968 blas: CudaBlas,
3969 l_dev: CudaSlice<f64>,
3972 y_dev: CudaSlice<f64>,
3975 schur_dev: CudaSlice<f64>,
3978 log_det_hessian: f64,
3981 }
3982
3983 impl ResidentArrowFrame {
3984 pub(super) fn new(
3988 sys: &ArrowSchurSystem,
3989 ridge_t: f64,
3990 ridge_beta: f64,
3991 ) -> Result<Self, ArrowSchurGpuFailure> {
3992 if ridge_t.is_nan() || ridge_beta.is_nan() {
3993 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
3994 reason: "ridge is NaN".to_string(),
3995 });
3996 }
3997 let n = sys.rows.len();
3998 let d = sys.d;
3999 let k = sys.k;
4000 let runtime = route_through_gpu(DispatchOp::SmallDenseBatchedPotrf { p: d, batch: n })
4001 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
4002 let stream = gam_gpu::device_runtime::cuda_context_for(runtime.device.ordinal)
4003 .and_then(|ctx| ctx.new_stream().ok())
4004 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
4005 let solver =
4006 DnHandle::new(stream.clone()).map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4007 let blas =
4008 CudaBlas::new(stream.clone()).map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4009
4010 let (d_host, b_host_stacked) = pack_host_d_and_stacked_b(sys, ridge_t);
4016 let mut l_dev = stream
4017 .clone_htod(&d_host)
4018 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4019 let mut y_dev = stream
4020 .clone_htod(&b_host_stacked)
4021 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4022
4023 let info_host = potrf_batched(&solver, &stream, d, n, &mut l_dev)?;
4025 if let Some(idx) = info_host.iter().position(|info| *info != 0) {
4026 return Err(ArrowSchurGpuFailure::RidgeBumpRequired {
4030 row: idx,
4031 bump: super::ridge_bump_to_make_pd(sys.rows[idx].htt.view(), ridge_t),
4032 });
4033 }
4034
4035 trsm_batched_lower_inplace_stacked(&blas, &stream, d, n, k, &l_dev, &mut y_dev)?;
4037
4038 let schur_init: Vec<f64> = {
4042 let mut tmp = Vec::with_capacity(k * k);
4043 for col in 0..k {
4044 for row in 0..k {
4045 let mut v = sys.hbb[[row, col]];
4046 if row == col {
4047 v += ridge_beta;
4048 }
4049 tmp.push(v);
4050 }
4051 }
4052 tmp
4053 };
4054 let mut schur_dev = stream
4055 .clone_htod(&schur_init)
4056 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4057 schur_gemm_stacked(&blas, n * d, k, &y_dev, &mut schur_dev)?;
4058 let info = potrf_single(&solver, &stream, k, &mut schur_dev)?;
4059 if info != 0 {
4060 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
4061 reason: format!("Schur Cholesky failed at pivot {info}"),
4062 });
4063 }
4064
4065 let l_local_host = stream
4067 .clone_dtoh(&l_dev)
4068 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4069 let l_schur_host = stream
4070 .clone_dtoh(&schur_dev)
4071 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4072 let mut log_det = 0.0_f64;
4073 for i in 0..n {
4074 let base = i * d * d;
4075 for j in 0..d {
4076 log_det += l_local_host[base + j * d + j].ln();
4077 }
4078 }
4079 for j in 0..k {
4080 log_det += l_schur_host[j * k + j].ln();
4081 }
4082 log_det *= 2.0;
4083
4084 Ok(Self {
4085 n,
4086 d,
4087 k,
4088 stream,
4089 blas,
4090 l_dev,
4091 y_dev,
4092 schur_dev,
4093 log_det_hessian: log_det,
4094 })
4095 }
4096
4097 #[inline]
4098 pub(super) fn log_det_hessian(&self) -> f64 {
4099 self.log_det_hessian
4100 }
4101
4102 pub(super) fn solve_gradient(
4106 &self,
4107 g_t: &[f64],
4108 g_beta: &[f64],
4109 ) -> Result<ArrowSchurGpuSolution, ArrowSchurGpuFailure> {
4110 let n = self.n;
4111 let d = self.d;
4112 let k = self.k;
4113 if g_t.len() != n * d || g_beta.len() != k {
4114 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
4115 reason: format!(
4116 "resident gradient shape mismatch: g_t={} (want {}), g_beta={} (want {})",
4117 g_t.len(),
4118 n * d,
4119 g_beta.len(),
4120 k
4121 ),
4122 });
4123 }
4124 let mut u_dev = self
4126 .stream
4127 .clone_htod(&g_t.to_vec())
4128 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4129 trsm_batched_lower_inplace(&self.blas, &self.stream, d, n, 1, &self.l_dev, &mut u_dev)?;
4130
4131 let rhs_init: Vec<f64> = g_beta.iter().map(|v| -v).collect();
4134 let mut rhs_dev = self
4135 .stream
4136 .clone_htod(&rhs_init)
4137 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4138 schur_rhs_stacked(&self.blas, n * d, k, &self.y_dev, &u_dev, &mut rhs_dev)?;
4139
4140 trsm_single(
4142 &self.blas,
4143 &self.stream,
4144 k,
4145 &self.schur_dev,
4146 &mut rhs_dev,
4147 false,
4148 false,
4149 )?;
4150 trsm_single(
4151 &self.blas,
4152 &self.stream,
4153 k,
4154 &self.schur_dev,
4155 &mut rhs_dev,
4156 false,
4157 true,
4158 )?;
4159 let delta_beta_host = self
4160 .stream
4161 .clone_dtoh(&rhs_dev)
4162 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4163 let delta_beta = Array1::from_vec(delta_beta_host);
4164
4165 back_sub_rhs_stacked(&self.blas, n * d, k, &self.y_dev, &rhs_dev, &mut u_dev)?;
4167 trsm_batched_lower_inplace_transposed(
4168 &self.blas,
4169 &self.stream,
4170 d,
4171 n,
4172 1,
4173 &self.l_dev,
4174 &mut u_dev,
4175 )?;
4176 let x_host = self
4177 .stream
4178 .clone_dtoh(&u_dev)
4179 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4180 let mut delta_t = Array1::<f64>::zeros(n * d);
4181 for (i, v) in x_host.iter().enumerate() {
4182 delta_t[i] = -*v;
4183 }
4184
4185 Ok(ArrowSchurGpuSolution {
4186 delta_t,
4187 delta_beta,
4188 log_det_hessian: self.log_det_hessian,
4189 })
4190 }
4191 }
4192
4193 pub(super) struct ResidentBaseArrowFrame {
4199 n: usize,
4200 d: usize,
4201 k: usize,
4202 stream: Arc<CudaStream>,
4203 solver: DnHandle,
4204 blas: CudaBlas,
4205 d_base_host: Vec<f64>,
4211 base_b_dev: CudaSlice<f64>,
4214 base_hbb_dev: CudaSlice<f64>,
4217 g_t_dev: CudaSlice<f64>,
4219 gb_host: Vec<f64>,
4222 ones_k_dev: CudaSlice<f64>,
4225 }
4226
4227 impl ResidentBaseArrowFrame {
4228 pub(super) fn new(sys: &ArrowSchurSystem) -> Result<Self, ArrowSchurGpuFailure> {
4230 let n = sys.rows.len();
4231 let d = sys.d;
4232 let k = sys.k;
4233 let runtime = route_through_gpu(DispatchOp::SmallDenseBatchedPotrf { p: d, batch: n })
4234 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
4235 let stream = gam_gpu::device_runtime::cuda_context_for(runtime.device.ordinal)
4236 .and_then(|ctx| ctx.new_stream().ok())
4237 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
4238 let solver =
4239 DnHandle::new(stream.clone()).map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4240 let blas =
4241 CudaBlas::new(stream.clone()).map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4242
4243 let (d_base_host, b_host, g_host) = pack_host(sys, 0.0);
4245 let base_b_dev = stream
4246 .clone_htod(&b_host)
4247 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4248 let g_t_dev = stream
4249 .clone_htod(&g_host)
4250 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4251
4252 let hbb_base: Vec<f64> = {
4254 let mut tmp = Vec::with_capacity(k * k);
4255 for col in 0..k {
4256 for row in 0..k {
4257 tmp.push(sys.hbb[[row, col]]);
4258 }
4259 }
4260 tmp
4261 };
4262 let base_hbb_dev = stream
4263 .clone_htod(&hbb_base)
4264 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4265
4266 let gb_host: Vec<f64> = sys.gb.iter().copied().collect();
4267 let ones_k_dev = stream
4268 .clone_htod(&vec![1.0_f64; k])
4269 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4270
4271 Ok(Self {
4272 n,
4273 d,
4274 k,
4275 stream,
4276 solver,
4277 blas,
4278 d_base_host,
4279 base_b_dev,
4280 base_hbb_dev,
4281 g_t_dev,
4282 gb_host,
4283 ones_k_dev,
4284 })
4285 }
4286
4287 pub(super) fn refactor_and_solve(
4293 &self,
4294 ridge_t: f64,
4295 ridge_beta: f64,
4296 ) -> Result<ArrowSchurGpuSolution, ArrowSchurGpuFailure> {
4297 if ridge_t.is_nan() || ridge_beta.is_nan() {
4298 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
4299 reason: "ridge is NaN".to_string(),
4300 });
4301 }
4302 let (n, d, k) = (self.n, self.d, self.k);
4303
4304 let mut d_ridged = self.d_base_host.clone();
4308 for i in 0..n {
4309 for j in 0..d {
4310 d_ridged[i * d * d + j * (d + 1)] += ridge_t;
4311 }
4312 }
4313 let mut l_dev = self
4314 .stream
4315 .clone_htod(&d_ridged)
4316 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4317
4318 let info_host = potrf_batched(&self.solver, &self.stream, d, n, &mut l_dev)?;
4320 if let Some(idx) = info_host.iter().position(|info| *info != 0) {
4321 let base = idx * d * d;
4322 return Err(ArrowSchurGpuFailure::RidgeBumpRequired {
4323 row: idx,
4324 bump: super::ridge_bump_to_make_pd_colmajor(&d_ridged[base..base + d * d], d),
4328 });
4329 }
4330
4331 let mut y_dev = self
4333 .stream
4334 .alloc_zeros::<f64>(n * d * k)
4335 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4336 self.stream
4337 .memcpy_dtod(&self.base_b_dev, &mut y_dev)
4338 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4339 trsm_batched_lower_inplace(&self.blas, &self.stream, d, n, k, &l_dev, &mut y_dev)?;
4340
4341 let mut u_dev = self
4343 .stream
4344 .alloc_zeros::<f64>(n * d)
4345 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4346 self.stream
4347 .memcpy_dtod(&self.g_t_dev, &mut u_dev)
4348 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4349 trsm_batched_lower_inplace(&self.blas, &self.stream, d, n, 1, &l_dev, &mut u_dev)?;
4350
4351 let mut schur_dev = self
4353 .stream
4354 .alloc_zeros::<f64>(k * k)
4355 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4356 self.stream
4357 .memcpy_dtod(&self.base_hbb_dev, &mut schur_dev)
4358 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4359 device_axpy_strided(
4361 &self.blas,
4362 &self.stream,
4363 k,
4364 ridge_beta,
4365 &self.ones_k_dev,
4366 1,
4367 &mut schur_dev,
4368 k + 1,
4369 )?;
4370 let rhs_init: Vec<f64> = self.gb_host.iter().map(|v| -v).collect();
4371 let mut rhs_dev = self
4372 .stream
4373 .clone_htod(&rhs_init)
4374 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4375
4376 accumulate_schur(
4377 &self.blas,
4378 d,
4379 k,
4380 n,
4381 &y_dev,
4382 &u_dev,
4383 &mut schur_dev,
4384 &mut rhs_dev,
4385 )?;
4386
4387 let info = potrf_single(&self.solver, &self.stream, k, &mut schur_dev)?;
4389 if info != 0 {
4390 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
4391 reason: format!("Schur Cholesky failed at pivot {info}"),
4392 });
4393 }
4394 trsm_single(
4395 &self.blas,
4396 &self.stream,
4397 k,
4398 &schur_dev,
4399 &mut rhs_dev,
4400 false,
4401 false,
4402 )?;
4403 trsm_single(
4404 &self.blas,
4405 &self.stream,
4406 k,
4407 &schur_dev,
4408 &mut rhs_dev,
4409 false,
4410 true,
4411 )?;
4412 let delta_beta_host = self
4413 .stream
4414 .clone_dtoh(&rhs_dev)
4415 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4416 let delta_beta = Array1::from_vec(delta_beta_host);
4417
4418 accumulate_back_sub_rhs(&self.blas, d, k, n, &y_dev, &rhs_dev, &mut u_dev)?;
4420 trsm_batched_lower_inplace_transposed(
4421 &self.blas,
4422 &self.stream,
4423 d,
4424 n,
4425 1,
4426 &l_dev,
4427 &mut u_dev,
4428 )?;
4429 let x_host = self
4430 .stream
4431 .clone_dtoh(&u_dev)
4432 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4433 let mut delta_t = Array1::<f64>::zeros(n * d);
4434 for (i, v) in x_host.iter().enumerate() {
4435 delta_t[i] = -*v;
4436 }
4437
4438 let l_local_host = self
4440 .stream
4441 .clone_dtoh(&l_dev)
4442 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4443 let l_schur_host = self
4444 .stream
4445 .clone_dtoh(&schur_dev)
4446 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4447 let mut log_det = 0.0_f64;
4448 for i in 0..n {
4449 let base = i * d * d;
4450 for j in 0..d {
4451 log_det += l_local_host[base + j * d + j].ln();
4452 }
4453 }
4454 for j in 0..k {
4455 log_det += l_schur_host[j * k + j].ln();
4456 }
4457 log_det *= 2.0;
4458
4459 Ok(ArrowSchurGpuSolution {
4460 delta_t,
4461 delta_beta,
4462 log_det_hessian: log_det,
4463 })
4464 }
4465 }
4466
4467 pub(super) fn solve_fused(
4468 sys: &ArrowSchurSystem,
4469 ridge_t: f64,
4470 ridge_beta: f64,
4471 ) -> Result<ArrowSchurGpuSolution, ArrowSchurGpuFailure> {
4472 let n = sys.rows.len();
4473 let d = sys.d;
4474 let k = sys.k;
4475 let plan = crate::gpu_kernels::arrow_schur_nvrtc::plan_fused_launch(n, d, k)
4476 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
4477 let p_max = plan.p_max;
4478 let r_template = plan.r_template;
4479
4480 let runtime = gam_gpu::linalg_dispatch::route_through_gpu(
4481 gam_gpu::linalg_dispatch::DispatchOp::SmallDenseBatchedPotrf { p: d, batch: n },
4482 )
4483 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
4484 let ctx = gam_gpu::device_runtime::cuda_context_for(runtime.device.ordinal)
4485 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
4486 let stream = ctx
4487 .new_stream()
4488 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4489 let cap = &runtime.device.capability;
4490 let key = crate::gpu_kernels::arrow_schur_nvrtc::FusedModuleCacheKey {
4491 cc_major: cap.compute_major,
4492 cc_minor: cap.compute_minor,
4493 p_max: p_max as u32,
4494 r_template: r_template as u32,
4495 };
4496 let module = fused_module_for(&ctx, key)?;
4497 let forward = module
4498 .load_function("arrow_schur_forward_pgroup")
4499 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4500 let back_sub = module
4501 .load_function("arrow_schur_back_sub_pgroup")
4502 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4503
4504 let (d_host, b_host, g_host) = pack_fused_host(sys, ridge_t, p_max, r_template);
4506 let d_dev = stream
4507 .clone_htod(&d_host)
4508 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4509 let b_dev = stream
4510 .clone_htod(&b_host)
4511 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4512 let g_dev = stream
4513 .clone_htod(&g_host)
4514 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4515 let mut l_out = stream
4516 .alloc_zeros::<f64>(n * p_max * p_max)
4517 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4518 let mut u_out = stream
4519 .alloc_zeros::<f64>(n * p_max)
4520 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4521 let mut y_out = stream
4522 .alloc_zeros::<f64>(n * p_max * r_template)
4523 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4524 let mut partial_s = stream
4525 .alloc_zeros::<f64>(plan.partial_s_doubles)
4526 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4527 let mut partial_r = stream
4528 .alloc_zeros::<f64>(plan.partial_r_doubles)
4529 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4530 let mut status_dev = stream
4531 .alloc_zeros::<i32>(n)
4532 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4533
4534 let cfg = LaunchConfig {
4536 grid_dim: (plan.blocks, 1, 1),
4537 block_dim: (plan.threads_per_block, 1, 1),
4538 shared_mem_bytes: 0,
4539 };
4540 let n_i32 = to_i32(n).ok_or(ArrowSchurGpuFailure::Unavailable)?;
4541 let p_i32 = to_i32(d).ok_or(ArrowSchurGpuFailure::Unavailable)?;
4542 let r_i32 = to_i32(k).ok_or(ArrowSchurGpuFailure::Unavailable)?;
4543 let ridge_arg = ridge_t;
4544 {
4545 let mut builder = stream.launch_builder(&forward);
4546 builder
4547 .arg(&d_dev)
4548 .arg(&b_dev)
4549 .arg(&g_dev)
4550 .arg(&n_i32)
4551 .arg(&p_i32)
4552 .arg(&r_i32)
4553 .arg(&ridge_arg)
4554 .arg(&mut l_out)
4555 .arg(&mut u_out)
4556 .arg(&mut y_out)
4557 .arg(&mut partial_s)
4558 .arg(&mut partial_r)
4559 .arg(&mut status_dev);
4560 unsafe { builder.launch(cfg) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4564 }
4565 stream
4566 .synchronize()
4567 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4568
4569 let status_host = stream
4571 .clone_dtoh(&status_dev)
4572 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4573 if let Some(row) = status_host.iter().position(|s| *s != 0) {
4574 return Err(ArrowSchurGpuFailure::RidgeBumpRequired {
4578 row,
4579 bump: super::ridge_bump_to_make_pd(sys.rows[row].htt.view(), ridge_t),
4580 });
4581 }
4582
4583 let partial_s_host = stream
4585 .clone_dtoh(&partial_s)
4586 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4587 let partial_r_host = stream
4588 .clone_dtoh(&partial_r)
4589 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4590 let mut schur_host = vec![0.0_f64; k * k];
4591 for col in 0..k {
4592 for row in 0..k {
4593 let mut v = sys.hbb[[row, col]];
4594 if row == col {
4595 v += ridge_beta;
4596 }
4597 schur_host[col * k + row] = v;
4598 }
4599 }
4600 let mut rhs_host: Vec<f64> = sys.gb.iter().map(|v| -v).collect();
4601 for i in 0..n {
4602 let s_base = i * r_template * r_template;
4605 for col in 0..k {
4606 let col_base = s_base + col * r_template;
4607 let dst_col_base = col * k;
4608 for row in 0..k {
4609 schur_host[dst_col_base + row] -= partial_s_host[col_base + row];
4610 }
4611 }
4612 let r_base = i * r_template;
4613 for a in 0..k {
4614 rhs_host[a] += partial_r_host[r_base + a];
4615 }
4616 }
4617
4618 let mut schur_dev = stream
4620 .clone_htod(&schur_host)
4621 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4622 let mut rhs_dev = stream
4623 .clone_htod(&rhs_host)
4624 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4625 let solver =
4626 DnHandle::new(stream.clone()).map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4627 let blas = CudaBlas::new(stream.clone()).map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4628 let info = potrf_single(&solver, &stream, k, &mut schur_dev)?;
4629 if info != 0 {
4630 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
4631 reason: format!("fused Schur Cholesky failed at pivot {info}"),
4632 });
4633 }
4634 trsm_single(&blas, &stream, k, &schur_dev, &mut rhs_dev, false, false)?;
4635 trsm_single(&blas, &stream, k, &schur_dev, &mut rhs_dev, false, true)?;
4636 let delta_beta_host = stream
4637 .clone_dtoh(&rhs_dev)
4638 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4639 let delta_beta = Array1::from_vec(delta_beta_host.clone());
4640
4641 let mut delta_t_dev = stream
4643 .alloc_zeros::<f64>(n * p_max)
4644 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4645 let back_cfg = LaunchConfig {
4646 grid_dim: (plan.blocks, 1, 1),
4647 block_dim: (plan.threads_per_block, 1, 1),
4648 shared_mem_bytes: 0,
4649 };
4650 {
4651 let mut builder = stream.launch_builder(&back_sub);
4652 builder
4653 .arg(&l_out)
4654 .arg(&u_out)
4655 .arg(&y_out)
4656 .arg(&rhs_dev)
4657 .arg(&n_i32)
4658 .arg(&p_i32)
4659 .arg(&r_i32)
4660 .arg(&mut delta_t_dev);
4661 unsafe { builder.launch(back_cfg) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4665 }
4666 stream
4667 .synchronize()
4668 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4669
4670 let delta_t_host = stream
4671 .clone_dtoh(&delta_t_dev)
4672 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4673 let mut delta_t = Array1::<f64>::zeros(n * d);
4674 for i in 0..n {
4675 let src_base = i * p_max;
4676 let dst_base = i * d;
4677 for r in 0..d {
4678 delta_t[dst_base + r] = delta_t_host[src_base + r];
4679 }
4680 }
4681
4682 let l_local_host = stream
4684 .clone_dtoh(&l_out)
4685 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4686 let l_schur_host = stream
4687 .clone_dtoh(&schur_dev)
4688 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4689 let mut log_det = 0.0_f64;
4690 for i in 0..n {
4691 let base = i * p_max * p_max;
4692 for j in 0..d {
4693 log_det += l_local_host[base + j * p_max + j].ln();
4694 }
4695 }
4696 for j in 0..k {
4697 log_det += l_schur_host[j * k + j].ln();
4698 }
4699 log_det *= 2.0;
4700
4701 Ok(ArrowSchurGpuSolution {
4702 delta_t,
4703 delta_beta,
4704 log_det_hessian: log_det,
4705 })
4706 }
4707
4708 pub(super) fn build_schur_matvec_backend(
4718 sys: &ArrowSchurSystem,
4719 ridge_t: f64,
4720 ridge_beta: f64,
4721 ) -> Result<crate::arrow_schur::GpuSchurMatvec, super::ArrowSchurGpuFailure> {
4722 let n = sys.rows.len();
4723 let d = sys.d;
4724 let k = sys.k;
4725 let plan = crate::gpu_kernels::arrow_schur_nvrtc::plan_fused_launch(n, d, k)
4726 .ok_or(super::ArrowSchurGpuFailure::Unavailable)?;
4727 let p_max = plan.p_max;
4728 let r_template = plan.r_template;
4729
4730 let runtime = gam_gpu::linalg_dispatch::route_through_gpu(
4731 gam_gpu::linalg_dispatch::DispatchOp::SmallDenseBatchedPotrf { p: d, batch: n },
4732 )
4733 .ok_or(super::ArrowSchurGpuFailure::Unavailable)?;
4734 let ctx = gam_gpu::device_runtime::cuda_context_for(runtime.device.ordinal)
4735 .ok_or(super::ArrowSchurGpuFailure::Unavailable)?;
4736 let stream = ctx
4737 .new_stream()
4738 .map_err(|_| super::ArrowSchurGpuFailure::Unavailable)?;
4739 let cap = &runtime.device.capability;
4740 let key = crate::gpu_kernels::arrow_schur_nvrtc::FusedModuleCacheKey {
4741 cc_major: cap.compute_major,
4742 cc_minor: cap.compute_minor,
4743 p_max: p_max as u32,
4744 r_template: r_template as u32,
4745 };
4746 let module = fused_module_for(&ctx, key)?;
4747 let forward = module
4748 .load_function("arrow_schur_forward_pgroup")
4749 .map_err(|_| super::ArrowSchurGpuFailure::Unavailable)?;
4750
4751 let (d_host, b_host, g_host) = pack_fused_host(sys, ridge_t, p_max, r_template);
4752 let d_dev = stream
4753 .clone_htod(&d_host)
4754 .map_err(|_| super::ArrowSchurGpuFailure::Unavailable)?;
4755 let b_dev = stream
4756 .clone_htod(&b_host)
4757 .map_err(|_| super::ArrowSchurGpuFailure::Unavailable)?;
4758 let g_dev = stream
4759 .clone_htod(&g_host)
4760 .map_err(|_| super::ArrowSchurGpuFailure::Unavailable)?;
4761 let mut l_out = stream
4762 .alloc_zeros::<f64>(n * p_max * p_max)
4763 .map_err(|_| super::ArrowSchurGpuFailure::Unavailable)?;
4764 let mut u_out = stream
4765 .alloc_zeros::<f64>(n * p_max)
4766 .map_err(|_| super::ArrowSchurGpuFailure::Unavailable)?;
4767 let mut y_out = stream
4768 .alloc_zeros::<f64>(n * p_max * r_template)
4769 .map_err(|_| super::ArrowSchurGpuFailure::Unavailable)?;
4770 let mut partial_s = stream
4771 .alloc_zeros::<f64>(plan.partial_s_doubles)
4772 .map_err(|_| super::ArrowSchurGpuFailure::Unavailable)?;
4773 let mut partial_r = stream
4774 .alloc_zeros::<f64>(plan.partial_r_doubles)
4775 .map_err(|_| super::ArrowSchurGpuFailure::Unavailable)?;
4776 let mut status_dev = stream
4777 .alloc_zeros::<i32>(n)
4778 .map_err(|_| super::ArrowSchurGpuFailure::Unavailable)?;
4779
4780 let cfg = LaunchConfig {
4781 grid_dim: (plan.blocks, 1, 1),
4782 block_dim: (plan.threads_per_block, 1, 1),
4783 shared_mem_bytes: 0,
4784 };
4785 let n_i32 = to_i32(n).ok_or(super::ArrowSchurGpuFailure::Unavailable)?;
4786 let p_i32 = to_i32(d).ok_or(super::ArrowSchurGpuFailure::Unavailable)?;
4787 let r_i32 = to_i32(k).ok_or(super::ArrowSchurGpuFailure::Unavailable)?;
4788 let ridge_arg = ridge_t;
4789 {
4790 let mut builder = stream.launch_builder(&forward);
4791 builder
4792 .arg(&d_dev)
4793 .arg(&b_dev)
4794 .arg(&g_dev)
4795 .arg(&n_i32)
4796 .arg(&p_i32)
4797 .arg(&r_i32)
4798 .arg(&ridge_arg)
4799 .arg(&mut l_out)
4800 .arg(&mut u_out)
4801 .arg(&mut y_out)
4802 .arg(&mut partial_s)
4803 .arg(&mut partial_r)
4804 .arg(&mut status_dev);
4805 unsafe { builder.launch(cfg) }.map_err(|_| super::ArrowSchurGpuFailure::Unavailable)?;
4808 }
4809 stream
4810 .synchronize()
4811 .map_err(|_| super::ArrowSchurGpuFailure::Unavailable)?;
4812
4813 let status_host = stream
4814 .clone_dtoh(&status_dev)
4815 .map_err(|_| super::ArrowSchurGpuFailure::Unavailable)?;
4816 if let Some(row) = status_host.iter().position(|s| *s != 0) {
4817 return Err(super::ArrowSchurGpuFailure::RidgeBumpRequired {
4821 row,
4822 bump: super::ridge_bump_to_make_pd(sys.rows[row].htt.view(), ridge_t),
4823 });
4824 }
4825
4826 let y_host = stream
4828 .clone_dtoh(&y_out)
4829 .map_err(|_| super::ArrowSchurGpuFailure::Unavailable)?;
4830
4831 let hbb_host: Vec<f64> = sys.hbb.iter().copied().collect();
4834 let hbb_is_kk = sys.hbb.dim() == (k, k);
4835 let hbb_matvec_opt = sys.hbb_matvec.clone();
4836
4837 let closure: crate::arrow_schur::GpuSchurMatvec =
4838 Arc::new(move |x: &Array1<f64>, out: &mut Array1<f64>| {
4839 assert_eq!(x.len(), k, "gpu_schur_matvec: x.len() != k");
4840 assert_eq!(out.len(), k, "gpu_schur_matvec: out.len() != k");
4841
4842 if let Some(ref mv) = hbb_matvec_opt {
4844 mv(x.view(), out);
4845 for a in 0..k {
4846 out[a] += ridge_beta * x[a];
4847 }
4848 } else if hbb_is_kk {
4849 for a in 0..k {
4851 let mut acc = ridge_beta * x[a];
4852 for b in 0..k {
4853 acc += hbb_host[a * k + b] * x[b];
4854 }
4855 out[a] = acc;
4856 }
4857 } else {
4858 for a in 0..k {
4859 out[a] = ridge_beta * x[a];
4860 }
4861 }
4862
4863 let mut z = vec![0.0_f64; d];
4866 for i in 0..n {
4867 let y_base = i * p_max * r_template;
4868 for r in 0..d {
4869 let mut acc = 0.0;
4870 for c in 0..k {
4871 acc += y_host[y_base + c * p_max + r] * x[c];
4872 }
4873 z[r] = acc;
4874 }
4875 for c in 0..k {
4876 let mut acc = 0.0;
4877 for r in 0..d {
4878 acc += y_host[y_base + c * p_max + r] * z[r];
4879 }
4880 out[c] -= acc;
4881 }
4882 }
4883 });
4884
4885 Ok(closure)
4886 }
4887
4888 struct DeviceSaeFrameBuffers {
4891 s_off: CudaSlice<i32>,
4893 s_m: CudaSlice<i32>,
4894 s_r: CudaSlice<i32>,
4895 s_ptr: CudaSlice<i32>,
4896 s_data: CudaSlice<f64>,
4897 s_blocks: usize,
4898 g_off_i: CudaSlice<i32>,
4900 g_off_j: CudaSlice<i32>,
4901 g_ri: CudaSlice<i32>,
4902 g_rj: CudaSlice<i32>,
4903 g_mi: CudaSlice<i32>,
4904 g_mj: CudaSlice<i32>,
4905 g_ptr: CudaSlice<i32>,
4906 g_data: CudaSlice<f64>,
4907 w_ptr: CudaSlice<i32>,
4908 w_data: CudaSlice<f64>,
4909 g_blocks: usize,
4910 g_max_work: usize,
4911 htb_ptr: CudaSlice<i32>,
4913 htb: CudaSlice<f64>,
4914 q_of: CudaSlice<i32>,
4915 ainv: CudaSlice<f64>,
4916 hvec: CudaSlice<f64>,
4917 svec: CudaSlice<f64>,
4918 scatter_partial: CudaSlice<f64>,
4923 n_chunks: usize,
4924 rows_per_chunk: usize,
4925 n_rows: usize,
4926 k: usize,
4927 max_q: usize,
4928 }
4929
4930 fn scatter_chunking(n_rows: usize) -> (usize, usize) {
4936 let target_chunks = 128usize;
4937 let rows_per_chunk = n_rows.div_ceil(target_chunks).max(1);
4938 let n_chunks = n_rows.div_ceil(rows_per_chunk).max(1);
4939 (rows_per_chunk, n_chunks)
4940 }
4941
4942 fn flatten_device_sae_frame_data(
4943 sys: &ArrowSchurSystem,
4944 data: &DeviceSaePcgData,
4945 frame: &DeviceSaeFrameData,
4946 ridge_t: f64,
4947 stream: &Arc<CudaStream>,
4948 ) -> Result<DeviceSaeFrameBuffers, ArrowSchurGpuFailure> {
4949 let host = super::flatten_frame_host_operands(sys, data, frame)?;
4956 let ainv = super::compute_ainv_host(sys, &host.q_of, host.max_q, host.n_rows, ridge_t)?;
4957 upload_frame_buffers(&host, &ainv, stream)
4958 }
4959
4960 fn upload_frame_buffers(
4966 host: &super::FrameHostOperands,
4967 ainv: &[f64],
4968 stream: &Arc<CudaStream>,
4969 ) -> Result<DeviceSaeFrameBuffers, ArrowSchurGpuFailure> {
4970 let htod_i = |v: &[i32]| {
4971 stream
4972 .clone_htod(v)
4973 .map_err(|_| ArrowSchurGpuFailure::Unavailable)
4974 };
4975 let htod_f = |v: &[f64]| {
4976 stream
4977 .clone_htod(v)
4978 .map_err(|_| ArrowSchurGpuFailure::Unavailable)
4979 };
4980 let (rows_per_chunk, n_chunks) = scatter_chunking(host.n_rows);
4981 Ok(DeviceSaeFrameBuffers {
4982 s_off: htod_i(&host.s_off)?,
4983 s_m: htod_i(&host.s_m)?,
4984 s_r: htod_i(&host.s_r)?,
4985 s_ptr: htod_i(&host.s_ptr)?,
4986 s_data: htod_f(&host.s_data)?,
4987 s_blocks: host.s_blocks,
4988 g_off_i: htod_i(&host.g_off_i)?,
4989 g_off_j: htod_i(&host.g_off_j)?,
4990 g_ri: htod_i(&host.g_ri)?,
4991 g_rj: htod_i(&host.g_rj)?,
4992 g_mi: htod_i(&host.g_mi)?,
4993 g_mj: htod_i(&host.g_mj)?,
4994 g_ptr: htod_i(&host.g_ptr)?,
4995 g_data: htod_f(&host.g_data)?,
4996 w_ptr: htod_i(&host.w_ptr)?,
4997 w_data: htod_f(&host.w_data)?,
4998 g_blocks: host.g_blocks,
4999 g_max_work: host.g_max_work,
5000 htb_ptr: htod_i(&host.htb_ptr)?,
5001 htb: htod_f(&host.htb)?,
5002 q_of: htod_i(&host.q_of)?,
5003 ainv: htod_f(ainv)?,
5004 hvec: stream
5005 .alloc_zeros::<f64>(host.n_rows * host.max_q)
5006 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
5007 svec: stream
5008 .alloc_zeros::<f64>(host.n_rows * host.max_q)
5009 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
5010 scatter_partial: stream
5011 .alloc_zeros::<f64>(n_chunks * host.k)
5012 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
5013 n_chunks,
5014 rows_per_chunk,
5015 n_rows: host.n_rows,
5016 k: host.k,
5017 max_q: host.max_q,
5018 })
5019 }
5020
5021 fn refresh_frame_buffers(
5026 host: &super::FrameHostOperands,
5027 buffers: &mut DeviceSaeFrameBuffers,
5028 stream: &Arc<CudaStream>,
5029 ) -> Result<(), ArrowSchurGpuFailure> {
5030 macro_rules! refresh {
5031 ($host:expr, $device:expr) => {{
5032 if $host.len() != $device.len() {
5033 return Err(ArrowSchurGpuFailure::Unavailable);
5034 }
5035 stream
5036 .memcpy_htod($host, &mut $device)
5037 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5038 }};
5039 }
5040 if host.s_blocks != buffers.s_blocks
5041 || host.g_blocks != buffers.g_blocks
5042 || host.g_max_work != buffers.g_max_work
5043 || host.n_rows != buffers.n_rows
5044 || host.k != buffers.k
5045 || host.max_q != buffers.max_q
5046 {
5047 return Err(ArrowSchurGpuFailure::Unavailable);
5048 }
5049 refresh!(&host.s_off, buffers.s_off);
5050 refresh!(&host.s_m, buffers.s_m);
5051 refresh!(&host.s_r, buffers.s_r);
5052 refresh!(&host.s_ptr, buffers.s_ptr);
5053 refresh!(&host.s_data, buffers.s_data);
5054 refresh!(&host.g_off_i, buffers.g_off_i);
5055 refresh!(&host.g_off_j, buffers.g_off_j);
5056 refresh!(&host.g_ri, buffers.g_ri);
5057 refresh!(&host.g_rj, buffers.g_rj);
5058 refresh!(&host.g_mi, buffers.g_mi);
5059 refresh!(&host.g_mj, buffers.g_mj);
5060 refresh!(&host.g_ptr, buffers.g_ptr);
5061 refresh!(&host.g_data, buffers.g_data);
5062 refresh!(&host.w_ptr, buffers.w_ptr);
5063 refresh!(&host.w_data, buffers.w_data);
5064 refresh!(&host.htb_ptr, buffers.htb_ptr);
5065 refresh!(&host.htb, buffers.htb);
5066 refresh!(&host.q_of, buffers.q_of);
5067 Ok(())
5068 }
5069
5070 fn sae_frame_penalty_diag_host(
5071 data: &DeviceSaePcgData,
5072 frame: &DeviceSaeFrameData,
5073 ridge_beta: f64,
5074 ) -> Result<Vec<f64>, ArrowSchurGpuFailure> {
5075 let mut diag = vec![ridge_beta; data.beta_dim];
5076 for (blk, &r) in data.smooth_blocks.iter().zip(frame.smooth_ranks.iter()) {
5078 let m = blk.factor_a.nrows();
5079 for ia in 0..m {
5080 let coeff = blk.factor_a[[ia, ia]];
5081 let base = blk.global_offset + ia * r;
5082 for ib in 0..r {
5083 if base + ib >= diag.len() {
5084 return Err(ArrowSchurGpuFailure::Unavailable);
5085 }
5086 diag[base + ib] += coeff;
5087 }
5088 }
5089 }
5090 for blk in &frame.frame_blocks {
5092 if blk.atom_i != blk.atom_j {
5093 continue;
5094 }
5095 let r = frame.ranks[blk.atom_i];
5096 let off = frame.border_offsets[blk.atom_i];
5097 let (mi, mj) = blk.g.dim();
5098 for li in 0..mi.min(mj) {
5099 let gii = blk.g[[li, li]];
5100 let base = off + li * r;
5101 for a in 0..r {
5102 if base + a >= diag.len() {
5103 return Err(ArrowSchurGpuFailure::Unavailable);
5104 }
5105 diag[base + a] += gii * blk.w[[a, a]];
5106 }
5107 }
5108 }
5109 Ok(diag)
5110 }
5111
5112 fn frame_grid(work: usize, n_rows: usize) -> Result<LaunchConfig, ArrowSchurGpuFailure> {
5113 Ok(LaunchConfig {
5114 grid_dim: (
5115 ((work as u32).saturating_add(255) / 256).max(1),
5116 checked_i32(n_rows)? as u32,
5117 1,
5118 ),
5119 block_dim: (256, 1, 1),
5120 shared_mem_bytes: 0,
5121 })
5122 }
5123
5124 fn launch_sae_frame_matvec(
5125 stream: &Arc<CudaStream>,
5126 module: &Arc<CudaModule>,
5127 buffers: &mut DeviceSaeFrameBuffers,
5128 x: &CudaSlice<f64>,
5129 out: &mut CudaSlice<f64>,
5130 ridge_beta: f64,
5131 ) -> Result<(), ArrowSchurGpuFailure> {
5132 launch_sae_init(stream, module, out, x, ridge_beta, buffers.k)?;
5133 if buffers.s_blocks > 0 {
5135 let kernel = module
5136 .load_function("arrow_sae_frame_smooth_matvec")
5137 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5138 let blocks_i32 = checked_i32(buffers.s_blocks)?;
5139 let cfg = frame_grid(buffers.k, buffers.s_blocks)?;
5140 let mut b = stream.launch_builder(&kernel);
5141 b.arg(x)
5142 .arg(&mut *out)
5143 .arg(&buffers.s_off)
5144 .arg(&buffers.s_m)
5145 .arg(&buffers.s_r)
5146 .arg(&buffers.s_ptr)
5147 .arg(&buffers.s_data)
5148 .arg(&blocks_i32);
5149 unsafe { b.launch(cfg) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5152 }
5153 if buffers.g_blocks > 0 {
5155 let kernel = module
5156 .load_function("arrow_sae_frame_g_matvec")
5157 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5158 let blocks_i32 = checked_i32(buffers.g_blocks)?;
5159 let cfg = frame_grid(buffers.g_max_work.max(1), buffers.g_blocks)?;
5160 let mut b = stream.launch_builder(&kernel);
5161 b.arg(x)
5162 .arg(&mut *out)
5163 .arg(&buffers.g_off_i)
5164 .arg(&buffers.g_off_j)
5165 .arg(&buffers.g_ri)
5166 .arg(&buffers.g_rj)
5167 .arg(&buffers.g_mi)
5168 .arg(&buffers.g_mj)
5169 .arg(&buffers.g_ptr)
5170 .arg(&buffers.g_data)
5171 .arg(&buffers.w_ptr)
5172 .arg(&buffers.w_data)
5173 .arg(&blocks_i32);
5174 unsafe { b.launch(cfg) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5177 }
5178 let k_i32 = checked_i32(buffers.k)?;
5180 let max_q_i32 = checked_i32(buffers.max_q)?;
5181 let n_rows_i32 = checked_i32(buffers.n_rows)?;
5182 {
5183 let kernel = module
5184 .load_function("arrow_sae_frame_apply_h")
5185 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5186 let cfg = frame_grid(buffers.max_q, buffers.n_rows)?;
5187 let mut b = stream.launch_builder(&kernel);
5188 b.arg(x)
5189 .arg(&buffers.htb_ptr)
5190 .arg(&buffers.htb)
5191 .arg(&buffers.q_of)
5192 .arg(&mut buffers.hvec)
5193 .arg(&k_i32)
5194 .arg(&max_q_i32)
5195 .arg(&n_rows_i32);
5196 unsafe { b.launch(cfg) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5199 }
5200 {
5201 let kernel = module
5202 .load_function("arrow_sae_frame_apply_ainv")
5203 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5204 let cfg = frame_grid(buffers.max_q, buffers.n_rows)?;
5205 let mut b = stream.launch_builder(&kernel);
5206 b.arg(&buffers.ainv)
5207 .arg(&buffers.hvec)
5208 .arg(&buffers.q_of)
5209 .arg(&mut buffers.svec)
5210 .arg(&max_q_i32)
5211 .arg(&n_rows_i32);
5212 unsafe { b.launch(cfg) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5215 }
5216 {
5217 let kernel = module
5218 .load_function("arrow_sae_frame_scatter_h")
5219 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5220 let cfg = frame_grid(buffers.k, buffers.n_rows)?;
5221 let mut b = stream.launch_builder(&kernel);
5222 b.arg(&buffers.svec)
5223 .arg(&buffers.htb_ptr)
5224 .arg(&buffers.htb)
5225 .arg(&buffers.q_of)
5226 .arg(out)
5227 .arg(&k_i32)
5228 .arg(&max_q_i32)
5229 .arg(&n_rows_i32);
5230 unsafe { b.launch(cfg) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5233 }
5234 Ok(())
5235 }
5236
5237 fn launch_sae_frame_reduced_schur_det(
5247 stream: &Arc<CudaStream>,
5248 module: &Arc<CudaModule>,
5249 buffers: &mut DeviceSaeFrameBuffers,
5250 x: &CudaSlice<f64>,
5251 out: &mut CudaSlice<f64>,
5252 ) -> Result<(), ArrowSchurGpuFailure> {
5253 let k_i32 = checked_i32(buffers.k)?;
5254 let max_q_i32 = checked_i32(buffers.max_q)?;
5255 let n_rows_i32 = checked_i32(buffers.n_rows)?;
5256 {
5260 let kernel = module
5261 .load_function("arrow_sae_frame_apply_h_warp")
5262 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5263 let block = checked_i32(buffers.max_q)?
5264 .checked_mul(32)
5265 .ok_or(ArrowSchurGpuFailure::Unavailable)? as u32;
5266 let cfg = LaunchConfig {
5267 grid_dim: (checked_i32(buffers.n_rows)? as u32, 1, 1),
5268 block_dim: (block, 1, 1),
5269 shared_mem_bytes: 0,
5270 };
5271 let mut b = stream.launch_builder(&kernel);
5272 b.arg(x)
5273 .arg(&buffers.htb_ptr)
5274 .arg(&buffers.htb)
5275 .arg(&buffers.q_of)
5276 .arg(&mut buffers.hvec)
5277 .arg(&k_i32)
5278 .arg(&max_q_i32)
5279 .arg(&n_rows_i32);
5280 unsafe { b.launch(cfg) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5284 }
5285 {
5287 let kernel = module
5288 .load_function("arrow_sae_frame_apply_ainv")
5289 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5290 let cfg = frame_grid(buffers.max_q, buffers.n_rows)?;
5291 let mut b = stream.launch_builder(&kernel);
5292 b.arg(&buffers.ainv)
5293 .arg(&buffers.hvec)
5294 .arg(&buffers.q_of)
5295 .arg(&mut buffers.svec)
5296 .arg(&max_q_i32)
5297 .arg(&n_rows_i32);
5298 unsafe { b.launch(cfg) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5301 }
5302 let k_blocks = ((buffers.k as u32).saturating_add(255) / 256).max(1);
5307 {
5308 let kernel = module
5309 .load_function("arrow_sae_frame_scatter_h_det_partial")
5310 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5311 let rows_per_chunk_i32 = checked_i32(buffers.rows_per_chunk)?;
5312 let cfg = LaunchConfig {
5313 grid_dim: (k_blocks, checked_i32(buffers.n_chunks)? as u32, 1),
5314 block_dim: (256, 1, 1),
5315 shared_mem_bytes: 0,
5316 };
5317 let mut b = stream.launch_builder(&kernel);
5318 b.arg(&buffers.svec)
5319 .arg(&buffers.htb_ptr)
5320 .arg(&buffers.htb)
5321 .arg(&buffers.q_of)
5322 .arg(&mut buffers.scatter_partial)
5323 .arg(&k_i32)
5324 .arg(&max_q_i32)
5325 .arg(&n_rows_i32)
5326 .arg(&rows_per_chunk_i32);
5327 unsafe { b.launch(cfg) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5330 }
5331 {
5334 let kernel = module
5335 .load_function("arrow_sae_frame_scatter_h_det_reduce")
5336 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5337 let n_chunks_i32 = checked_i32(buffers.n_chunks)?;
5338 let cfg = LaunchConfig {
5339 grid_dim: (k_blocks, 1, 1),
5340 block_dim: (256, 1, 1),
5341 shared_mem_bytes: 0,
5342 };
5343 let mut b = stream.launch_builder(&kernel);
5344 b.arg(&buffers.scatter_partial)
5345 .arg(&mut *out)
5346 .arg(&k_i32)
5347 .arg(&n_chunks_i32);
5348 unsafe { b.launch(cfg) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5351 }
5352 Ok(())
5353 }
5354
5355 fn launch_sae_frame_diag_sub(
5356 stream: &Arc<CudaStream>,
5357 module: &Arc<CudaModule>,
5358 buffers: &DeviceSaeFrameBuffers,
5359 diag: &mut CudaSlice<f64>,
5360 ) -> Result<(), ArrowSchurGpuFailure> {
5361 let kernel = module
5362 .load_function("arrow_sae_frame_diag_sub")
5363 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5364 let k_i32 = checked_i32(buffers.k)?;
5365 let max_q_i32 = checked_i32(buffers.max_q)?;
5366 let n_rows_i32 = checked_i32(buffers.n_rows)?;
5367 let cfg = frame_grid(buffers.k, buffers.n_rows)?;
5368 let mut b = stream.launch_builder(&kernel);
5369 b.arg(diag)
5370 .arg(&buffers.ainv)
5371 .arg(&buffers.htb_ptr)
5372 .arg(&buffers.htb)
5373 .arg(&buffers.q_of)
5374 .arg(&k_i32)
5375 .arg(&max_q_i32)
5376 .arg(&n_rows_i32);
5377 unsafe { b.launch(cfg) }
5379 .map(drop)
5380 .map_err(|_| ArrowSchurGpuFailure::Unavailable)
5381 }
5382
5383 pub(super) fn framed_schur_matvec_once_on_device(
5394 sys: &ArrowSchurSystem,
5395 data: &DeviceSaePcgData,
5396 ridge_t: f64,
5397 ridge_beta: f64,
5398 x: &Array1<f64>,
5399 ) -> Result<Array1<f64>, ArrowSchurGpuFailure> {
5400 let k = x.len();
5401 if k == 0 || data.beta_dim != k || sys.k != k {
5402 return Err(ArrowSchurGpuFailure::Unavailable);
5403 }
5404 let frame = data
5405 .frame
5406 .as_ref()
5407 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
5408 let runtime = super::resolve_runtime_for_device_path()?
5411 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
5412 let ctx = gam_gpu::device_runtime::cuda_context_for(runtime.selected_device().ordinal)
5413 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
5414 let stream = ctx
5415 .new_stream()
5416 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5417 let vector_module = pcg_vector_module(&ctx)?;
5418 let mut buffers = flatten_device_sae_frame_data(sys, data, frame, ridge_t, &stream)?;
5419 let x_dev = stream
5420 .clone_htod(x.as_slice().ok_or(ArrowSchurGpuFailure::Unavailable)?)
5421 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5422 let mut out_dev = stream
5423 .alloc_zeros::<f64>(k)
5424 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5425 launch_sae_frame_matvec(
5426 &stream,
5427 vector_module,
5428 &mut buffers,
5429 &x_dev,
5430 &mut out_dev,
5431 ridge_beta,
5432 )?;
5433 let out = stream
5434 .clone_dtoh(&out_dev)
5435 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5436 Ok(Array1::from_vec(out))
5437 }
5438
5439 pub(super) fn framed_reduced_schur_det_once_on_device(
5447 sys: &ArrowSchurSystem,
5448 data: &DeviceSaePcgData,
5449 ridge_t: f64,
5450 ridge_beta: f64,
5451 x: &Array1<f64>,
5452 ) -> Result<Array1<f64>, ArrowSchurGpuFailure> {
5453 let k = x.len();
5454 if k == 0 || data.beta_dim != k || sys.k != k {
5455 return Err(ArrowSchurGpuFailure::Unavailable);
5456 }
5457 let frame = data
5458 .frame
5459 .as_ref()
5460 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
5461 let runtime = super::resolve_runtime_for_device_path()?
5462 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
5463 let ctx = gam_gpu::device_runtime::cuda_context_for(runtime.selected_device().ordinal)
5464 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
5465 let stream = ctx
5466 .new_stream()
5467 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5468 let vector_module = pcg_vector_module(&ctx)?;
5469 let mut buffers = flatten_device_sae_frame_data(sys, data, frame, ridge_t, &stream)?;
5470 let x_slice = x.as_slice().ok_or(ArrowSchurGpuFailure::Unavailable)?;
5471 let x_dev = stream
5472 .clone_htod(x_slice)
5473 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5474 let mut reduced_dev = stream
5475 .alloc_zeros::<f64>(k)
5476 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5477 launch_sae_frame_reduced_schur_det(
5478 &stream,
5479 vector_module,
5480 &mut buffers,
5481 &x_dev,
5482 &mut reduced_dev,
5483 )?;
5484 let reduced = stream
5485 .clone_dtoh(&reduced_dev)
5486 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5487 let mut out = vec![0.0_f64; k];
5489 super::sae_framed_penalty_matvec_cpu(data, ridge_beta, x_slice, &mut out);
5490 for a in 0..k {
5491 out[a] += reduced[a];
5492 }
5493 Ok(Array1::from_vec(out))
5494 }
5495
5496 pub(super) fn solve_sae_matrix_free_pcg_framed(
5497 sys: &ArrowSchurSystem,
5498 data: &DeviceSaePcgData,
5499 ridge_t: f64,
5500 ridge_beta: f64,
5501 rhs_beta: &Array1<f64>,
5502 max_iterations: usize,
5503 relative_tolerance: f64,
5504 ) -> Result<(Array1<f64>, ArrowPcgDiagnostics), ArrowSchurGpuFailure> {
5505 let k = rhs_beta.len();
5506 if k == 0 || data.beta_dim != k || sys.k != k {
5507 return Err(ArrowSchurGpuFailure::Unavailable);
5508 }
5509 let frame = data
5510 .frame
5511 .as_ref()
5512 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
5513 let runtime = super::resolve_runtime_for_device_path()?
5514 .filter(|rt| {
5515 rt.policy().reduced_schur_matvec_should_offload(
5516 sys.rows.len(),
5517 sys.k,
5518 sys.d,
5519 max_iterations,
5520 )
5521 })
5522 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
5523 let ctx = gam_gpu::device_runtime::cuda_context_for(runtime.selected_device().ordinal)
5524 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
5525 let stream = ctx
5526 .new_stream()
5527 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5528 let blas = CudaBlas::new(stream.clone()).map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5529 let vector_module = pcg_vector_module(&ctx)?;
5530 log::info!(
5535 "#1017/#2230 {} ridge_t={ridge_t:e} ridge_beta={ridge_beta:e}",
5536 data.operand_byte_report()
5537 );
5538 let mut buffers = flatten_device_sae_frame_data(sys, data, frame, ridge_t, &stream)?;
5539 let dctx = FramedPcgCtx {
5540 stream: &stream,
5541 blas: &blas,
5542 module: vector_module,
5543 max_iterations,
5544 relative_tolerance,
5545 };
5546 pcg_solve_framed_body(data, frame, &mut buffers, ridge_beta, rhs_beta, &dctx)
5547 }
5548
5549 struct FramedPcgCtx<'a> {
5552 stream: &'a Arc<CudaStream>,
5553 blas: &'a CudaBlas,
5554 module: &'a Arc<CudaModule>,
5555 max_iterations: usize,
5556 relative_tolerance: f64,
5557 }
5558
5559 fn pcg_solve_framed_body(
5567 data: &DeviceSaePcgData,
5568 frame: &DeviceSaeFrameData,
5569 buffers: &mut DeviceSaeFrameBuffers,
5570 ridge_beta: f64,
5571 rhs_beta: &Array1<f64>,
5572 dctx: &FramedPcgCtx<'_>,
5573 ) -> Result<(Array1<f64>, ArrowPcgDiagnostics), ArrowSchurGpuFailure> {
5574 let stream = dctx.stream;
5575 let blas = dctx.blas;
5576 let vector_module = dctx.module;
5577 let max_iterations = dctx.max_iterations;
5578 let relative_tolerance = dctx.relative_tolerance;
5579 let k = rhs_beta.len();
5580 let rhs_norm = rhs_beta.iter().map(|v| v * v).sum::<f64>().sqrt();
5581 if rhs_norm == 0.0 {
5582 return Ok((Array1::<f64>::zeros(k), ArrowPcgDiagnostics::default()));
5583 }
5584 let tol = (relative_tolerance.max(0.0) * rhs_norm).max(1e-12);
5585 let rhs_dev = stream
5586 .clone_htod(
5587 rhs_beta
5588 .as_slice()
5589 .ok_or(ArrowSchurGpuFailure::Unavailable)?,
5590 )
5591 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5592 let diag_host = sae_frame_penalty_diag_host(data, frame, ridge_beta)?;
5593 let mut diag_dev = stream
5594 .clone_htod(&diag_host)
5595 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5596 launch_sae_frame_diag_sub(stream, vector_module, buffers, &mut diag_dev)?;
5597 let diag_host = stream
5598 .clone_dtoh(&diag_dev)
5599 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5600 let mut inv_diag = Vec::with_capacity(k);
5601 for (idx, &d) in diag_host.iter().enumerate() {
5602 if !d.is_finite() || d <= 1.0e-18 {
5603 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
5604 reason: format!(
5605 "framed SAE GPU PCG: non-positive Jacobi diagonal at {idx}: {d:e}"
5606 ),
5607 });
5608 }
5609 inv_diag.push(1.0 / d);
5610 }
5611 let inv_diag_dev = stream
5612 .clone_htod(&inv_diag)
5613 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5614
5615 let mut x_dev = stream
5616 .alloc_zeros::<f64>(k)
5617 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5618 let mut r_dev = stream
5619 .alloc_zeros::<f64>(k)
5620 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5621 device_copy(blas, stream, k, &rhs_dev, &mut r_dev)?;
5622 let mut z_dev = stream
5623 .alloc_zeros::<f64>(k)
5624 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5625 launch_jacobi_mul(stream, vector_module, &inv_diag_dev, &r_dev, &mut z_dev, k)?;
5626 let mut p_dev = stream
5627 .alloc_zeros::<f64>(k)
5628 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5629 device_copy(blas, stream, k, &z_dev, &mut p_dev)?;
5630 let mut ap_dev = stream
5631 .alloc_zeros::<f64>(k)
5632 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5633
5634 let mut rz = device_dot(blas, stream, k, &r_dev, &z_dev)?;
5635 if rz <= 0.0 || !rz.is_finite() {
5636 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
5637 reason: format!("framed SAE GPU PCG: non-positive initial rᵀM⁻¹r={rz:e}"),
5638 });
5639 }
5640 let mut diag = ArrowPcgDiagnostics {
5641 precond_apply_calls: 1,
5642 stopping_reason: PcgStopReason::MaxIter,
5643 ..ArrowPcgDiagnostics::default()
5644 };
5645 for _ in 0..max_iterations.max(1) {
5646 launch_sae_frame_matvec(
5647 stream,
5648 vector_module,
5649 buffers,
5650 &p_dev,
5651 &mut ap_dev,
5652 ridge_beta,
5653 )?;
5654 diag.matvec_calls += 1;
5655 diag.iterations += 1;
5656 let pap = device_dot(blas, stream, k, &p_dev, &ap_dev)?;
5657 if pap <= 0.0 || !pap.is_finite() {
5658 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
5659 reason: format!("framed SAE GPU PCG: non-positive curvature pᵀAp={pap:e}"),
5660 });
5661 }
5662 let alpha = rz / pap;
5663 device_axpy(blas, stream, k, alpha, &p_dev, &mut x_dev)?;
5664 device_axpy(blas, stream, k, -alpha, &ap_dev, &mut r_dev)?;
5665 let r_norm = device_nrm2(blas, stream, k, &r_dev)?;
5666 if r_norm <= tol {
5667 diag.final_relative_residual = r_norm / rhs_norm;
5668 diag.stopping_reason = PcgStopReason::Converged;
5669 break;
5670 }
5671 launch_jacobi_mul(stream, vector_module, &inv_diag_dev, &r_dev, &mut z_dev, k)?;
5672 diag.precond_apply_calls += 1;
5673 let rz_new = device_dot(blas, stream, k, &r_dev, &z_dev)?;
5674 if rz_new <= 0.0 || !rz_new.is_finite() {
5675 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
5676 reason: format!("framed SAE GPU PCG: non-positive rᵀM⁻¹r={rz_new:e}"),
5677 });
5678 }
5679 let beta = rz_new / rz;
5680 launch_update_p(stream, vector_module, &z_dev, beta, &mut p_dev, k)?;
5681 rz = rz_new;
5682 }
5683 if diag.stopping_reason != PcgStopReason::Converged {
5684 let r_norm = device_nrm2(blas, stream, k, &r_dev)?;
5685 diag.final_relative_residual = r_norm / rhs_norm;
5686 diag.stopping_reason = PcgStopReason::MaxIter;
5687 }
5688 let x = stream
5689 .clone_dtoh(&x_dev)
5690 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5691 Ok((Array1::from_vec(x), diag))
5692 }
5693
5694 pub(crate) struct ResidentSaeFrameHandle {
5703 ctx: Arc<CudaContext>,
5704 stream: Arc<CudaStream>,
5705 buffers: std::sync::Mutex<DeviceSaeFrameBuffers>,
5706 q_of: Vec<i32>,
5707 max_q: usize,
5708 n_rows: usize,
5709 k: usize,
5710 }
5711
5712 impl ResidentSaeFrameHandle {
5713 pub(crate) fn build(
5721 sys: &ArrowSchurSystem,
5722 cg_iters: usize,
5723 ) -> Result<Option<Self>, ArrowSchurGpuFailure> {
5724 let Some(data) = sys.device_sae_pcg.as_ref() else {
5725 return Ok(None);
5726 };
5727 let Some(frame) = data.frame.as_ref() else {
5728 return Ok(None);
5729 };
5730 if sys.k == 0 || data.beta_dim != sys.k {
5731 return Ok(None);
5732 }
5733 let Some(runtime) = super::resolve_runtime_for_device_path()? else {
5734 return Ok(None);
5735 };
5736 if !runtime.policy().reduced_schur_matvec_should_offload(
5737 sys.rows.len(),
5738 sys.k,
5739 sys.d,
5740 cg_iters,
5741 ) {
5742 return Ok(None);
5743 }
5744 let ctx = gam_gpu::device_runtime::cuda_context_for(runtime.selected_device().ordinal)
5745 .ok_or_else(|| ArrowSchurGpuFailure::SchurFactorFailed {
5746 reason: "resident SAE frame could not bind the admitted CUDA context"
5747 .to_string(),
5748 })?;
5749 let stream = ctx.new_stream().map_err(|error| {
5750 ArrowSchurGpuFailure::SchurFactorFailed {
5751 reason: format!("resident SAE frame stream creation failed: {error}"),
5752 }
5753 })?;
5754 let host = super::flatten_frame_host_operands(sys, data, frame)?;
5755 log::info!(
5762 "#1017 SAE resident frame ENGAGED: {} uploaded ONCE for the ladder; \
5763 per-trial re-upload now only ainv ({}rows × {}²·8B)",
5764 data.operand_byte_report(),
5765 host.n_rows,
5766 host.max_q
5767 );
5768 let zero_ainv = vec![0.0_f64; host.n_rows * host.max_q * host.max_q];
5769 let buffers = upload_frame_buffers(&host, &zero_ainv, &stream)?;
5770 Ok(Some(Self {
5771 ctx,
5772 stream,
5773 buffers: std::sync::Mutex::new(buffers),
5774 q_of: host.q_of,
5775 max_q: host.max_q,
5776 n_rows: host.n_rows,
5777 k: host.k,
5778 }))
5779 }
5780
5781 pub(super) fn prime_ainv(
5788 &self,
5789 sys: &ArrowSchurSystem,
5790 ridge_t: f64,
5791 ) -> Result<(), ArrowSchurGpuFailure> {
5792 if sys.k != self.k || sys.rows.len() != self.n_rows {
5793 return Err(ArrowSchurGpuFailure::Unavailable);
5794 }
5795 let ainv =
5796 super::compute_ainv_host(sys, &self.q_of, self.max_q, self.n_rows, ridge_t)?;
5797 let mut buffers = self
5798 .buffers
5799 .lock()
5800 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5801 if ainv.len() != buffers.ainv.len() {
5802 return Err(ArrowSchurGpuFailure::Unavailable);
5803 }
5804 self.stream
5805 .memcpy_htod(&ainv, &mut buffers.ainv)
5806 .map_err(|_| ArrowSchurGpuFailure::Unavailable)
5807 }
5808 }
5809
5810 impl super::SaeResidentFrame for ResidentSaeFrameHandle {
5811 fn refresh(&self, sys: &ArrowSchurSystem) -> Result<(), ArrowSchurGpuFailure> {
5812 if sys.k != self.k || sys.rows.len() != self.n_rows {
5813 return Err(ArrowSchurGpuFailure::Unavailable);
5814 }
5815 let data = sys
5816 .device_sae_pcg
5817 .as_ref()
5818 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
5819 let frame = data
5820 .frame
5821 .as_ref()
5822 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
5823 let host = super::flatten_frame_host_operands(sys, data, frame)?;
5824 if host.q_of != self.q_of || host.max_q != self.max_q {
5825 return Err(ArrowSchurGpuFailure::Unavailable);
5826 }
5827 let mut buffers = self
5828 .buffers
5829 .lock()
5830 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5831 refresh_frame_buffers(&host, &mut buffers, &self.stream)
5832 }
5833
5834 fn resolve(
5835 &self,
5836 sys: &ArrowSchurSystem,
5837 ridge_t: f64,
5838 ridge_beta: f64,
5839 rhs_beta: &Array1<f64>,
5840 max_iterations: usize,
5841 relative_tolerance: f64,
5842 ) -> Result<(Array1<f64>, ArrowPcgDiagnostics), ArrowSchurGpuFailure> {
5843 if sys.k != self.k || sys.rows.len() != self.n_rows || rhs_beta.len() != self.k {
5847 return Err(ArrowSchurGpuFailure::Unavailable);
5848 }
5849 let data = sys
5850 .device_sae_pcg
5851 .as_ref()
5852 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
5853 let frame = data
5854 .frame
5855 .as_ref()
5856 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
5857 let ainv = super::compute_ainv_host(sys, &self.q_of, self.max_q, self.n_rows, ridge_t)?;
5861 let mut buffers = self
5862 .buffers
5863 .lock()
5864 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5865 if ainv.len() != buffers.ainv.len() {
5866 return Err(ArrowSchurGpuFailure::Unavailable);
5867 }
5868 self.stream
5869 .memcpy_htod(&ainv, &mut buffers.ainv)
5870 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5871 let blas = CudaBlas::new(self.stream.clone())
5872 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5873 let vector_module = pcg_vector_module(&self.ctx)?;
5874 let dctx = FramedPcgCtx {
5875 stream: &self.stream,
5876 blas: &blas,
5877 module: vector_module,
5878 max_iterations,
5879 relative_tolerance,
5880 };
5881 pcg_solve_framed_body(data, frame, &mut *buffers, ridge_beta, rhs_beta, &dctx)
5882 }
5883 }
5884
5885 pub(super) fn build_framed_resident_evidence_matvec(
5897 sys: &ArrowSchurSystem,
5898 ridge_t: f64,
5899 ridge_beta: f64,
5900 apply_budget: usize,
5901 ) -> Result<Option<crate::arrow_schur::GpuSchurMatvec>, ArrowSchurGpuFailure> {
5902 let Some(handle) = ResidentSaeFrameHandle::build(sys, apply_budget)? else {
5903 return Ok(None);
5904 };
5905 handle.prime_ainv(sys, ridge_t)?;
5906 let Some(data) = sys.device_sae_pcg.as_ref().cloned() else {
5907 return Ok(None);
5908 };
5909 let handle = Arc::new(handle);
5910 let k = handle.k;
5911 let closure: crate::arrow_schur::GpuSchurMatvec =
5912 Arc::new(move |x: &Array1<f64>, out: &mut Array1<f64>| {
5913 assert_eq!(x.len(), k, "#1017 framed evidence matvec: x.len() != k");
5914 assert_eq!(out.len(), k, "#1017 framed evidence matvec: out.len() != k");
5915 let module = pcg_vector_module(&handle.ctx)
5916 .expect("#1017 framed evidence matvec: pcg_vector_module unavailable");
5917 let x_slice = x
5918 .as_slice()
5919 .expect("#1017 framed evidence matvec: x not contiguous");
5920 let x_dev = handle
5921 .stream
5922 .clone_htod(x_slice)
5923 .expect("#1017 framed evidence matvec: htod(x) failed");
5924 let mut reduced_dev = handle
5925 .stream
5926 .alloc_zeros::<f64>(k)
5927 .expect("#1017 framed evidence matvec: device alloc failed");
5928 {
5929 let mut buffers = handle
5930 .buffers
5931 .lock()
5932 .expect("#1017 framed evidence matvec: resident frame poisoned");
5933 launch_sae_frame_reduced_schur_det(
5934 &handle.stream,
5935 module,
5936 &mut buffers,
5937 &x_dev,
5938 &mut reduced_dev,
5939 )
5940 .expect("#1017 framed evidence matvec: device reduced-Schur launch failed");
5941 }
5942 let reduced = handle
5943 .stream
5944 .clone_dtoh(&reduced_dev)
5945 .expect("#1017 framed evidence matvec: dtoh(out) failed");
5946 let out_slice = out
5947 .as_slice_mut()
5948 .expect("#1017 framed evidence matvec: out not contiguous");
5949 super::sae_framed_penalty_matvec_cpu(&data, ridge_beta, x_slice, out_slice);
5950 for a in 0..k {
5951 out_slice[a] += reduced[a];
5952 }
5953 });
5954 Ok(Some(closure))
5955 }
5956
5957 pub(super) fn solve_sae_matrix_free_pcg(
5964 sys: &ArrowSchurSystem,
5965 data: &DeviceSaePcgData,
5966 ridge_t: f64,
5967 ridge_beta: f64,
5968 rhs_beta: &Array1<f64>,
5969 max_iterations: usize,
5970 relative_tolerance: f64,
5971 ) -> Result<(Array1<f64>, ArrowPcgDiagnostics), ArrowSchurGpuFailure> {
5972 let k = rhs_beta.len();
5973 if k == 0 || data.beta_dim != k || sys.k != k {
5974 return Err(ArrowSchurGpuFailure::Unavailable);
5975 }
5976 if data.frame.is_some() {
5980 return Err(ArrowSchurGpuFailure::Unavailable);
5981 }
5982 let runtime = super::resolve_runtime_for_device_path()?
5996 .filter(|rt| {
5997 rt.policy().reduced_schur_matvec_should_offload(
5998 sys.rows.len(),
5999 sys.k,
6000 sys.d,
6001 max_iterations,
6002 )
6003 })
6004 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
6005 let ctx = gam_gpu::device_runtime::cuda_context_for(runtime.selected_device().ordinal)
6006 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
6007 let stream = ctx
6008 .new_stream()
6009 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6010 let blas = CudaBlas::new(stream.clone()).map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6011 let vector_module = pcg_vector_module(&ctx)?;
6012 log::info!(
6016 "#1017/#2230 {} ridge_t={ridge_t:e} ridge_beta={ridge_beta:e}",
6017 data.operand_byte_report()
6018 );
6019 let mut buffers = flatten_device_sae_data(sys, data, ridge_t, &stream)?;
6020
6021 let rhs_norm = rhs_beta.iter().map(|v| v * v).sum::<f64>().sqrt();
6022 if rhs_norm == 0.0 {
6023 return Ok((Array1::<f64>::zeros(k), ArrowPcgDiagnostics::default()));
6024 }
6025 let tol = (relative_tolerance.max(0.0) * rhs_norm).max(1e-12);
6026 let rhs_dev = stream
6027 .clone_htod(
6028 rhs_beta
6029 .as_slice()
6030 .ok_or(ArrowSchurGpuFailure::Unavailable)?,
6031 )
6032 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6033 let diag_host = sae_penalty_diag_host(data, ridge_beta)?;
6034 let mut diag_dev = stream
6035 .clone_htod(&diag_host)
6036 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6037 launch_sae_diag_sub(&stream, vector_module, &buffers, &mut diag_dev)?;
6038 let diag_host = stream
6039 .clone_dtoh(&diag_dev)
6040 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6041 let mut inv_diag = Vec::with_capacity(k);
6042 for (idx, &d) in diag_host.iter().enumerate() {
6043 if !d.is_finite() || d <= 1.0e-18 {
6044 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
6045 reason: format!(
6046 "SAE matrix-free GPU PCG: non-positive Schur Jacobi diagonal at {idx}: {d:e}"
6047 ),
6048 });
6049 }
6050 inv_diag.push(1.0 / d);
6051 }
6052 let inv_diag_dev = stream
6053 .clone_htod(&inv_diag)
6054 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6055
6056 let mut x_dev = stream
6057 .alloc_zeros::<f64>(k)
6058 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6059 let mut r_dev = stream
6060 .alloc_zeros::<f64>(k)
6061 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6062 device_copy(&blas, &stream, k, &rhs_dev, &mut r_dev)?;
6063 let mut z_dev = stream
6064 .alloc_zeros::<f64>(k)
6065 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6066 launch_jacobi_mul(&stream, vector_module, &inv_diag_dev, &r_dev, &mut z_dev, k)?;
6067 let mut p_dev = stream
6068 .alloc_zeros::<f64>(k)
6069 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6070 device_copy(&blas, &stream, k, &z_dev, &mut p_dev)?;
6071 let mut ap_dev = stream
6072 .alloc_zeros::<f64>(k)
6073 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6074
6075 let mut rz = device_dot(&blas, &stream, k, &r_dev, &z_dev)?;
6076 if rz <= 0.0 || !rz.is_finite() {
6077 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
6078 reason: format!("SAE matrix-free GPU PCG: non-positive initial rᵀM⁻¹r={rz:e}"),
6079 });
6080 }
6081 let mut diag = ArrowPcgDiagnostics {
6082 precond_apply_calls: 1,
6083 stopping_reason: PcgStopReason::MaxIter,
6084 ..ArrowPcgDiagnostics::default()
6085 };
6086
6087 for _ in 0..max_iterations.max(1) {
6088 launch_sae_matvec(
6089 &stream,
6090 vector_module,
6091 &mut buffers,
6092 &p_dev,
6093 &mut ap_dev,
6094 ridge_beta,
6095 )?;
6096 diag.matvec_calls += 1;
6097 diag.iterations += 1;
6098 let pap = device_dot(&blas, &stream, k, &p_dev, &ap_dev)?;
6099 if pap <= 0.0 || !pap.is_finite() {
6100 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
6101 reason: format!("SAE matrix-free GPU PCG: non-positive curvature pᵀAp={pap:e}"),
6102 });
6103 }
6104 let alpha = rz / pap;
6105 device_axpy(&blas, &stream, k, alpha, &p_dev, &mut x_dev)?;
6106 device_axpy(&blas, &stream, k, -alpha, &ap_dev, &mut r_dev)?;
6107 let r_norm = device_nrm2(&blas, &stream, k, &r_dev)?;
6108 if r_norm <= tol {
6109 diag.final_relative_residual = r_norm / rhs_norm;
6110 diag.stopping_reason = PcgStopReason::Converged;
6111 break;
6112 }
6113 launch_jacobi_mul(&stream, vector_module, &inv_diag_dev, &r_dev, &mut z_dev, k)?;
6114 diag.precond_apply_calls += 1;
6115 let rz_new = device_dot(&blas, &stream, k, &r_dev, &z_dev)?;
6116 if rz_new <= 0.0 || !rz_new.is_finite() {
6117 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
6118 reason: format!("SAE matrix-free GPU PCG: non-positive rᵀM⁻¹r={rz_new:e}"),
6119 });
6120 }
6121 let beta = rz_new / rz;
6122 launch_update_p(&stream, vector_module, &z_dev, beta, &mut p_dev, k)?;
6123 rz = rz_new;
6124 }
6125 if diag.stopping_reason != PcgStopReason::Converged {
6126 let r_norm = device_nrm2(&blas, &stream, k, &r_dev)?;
6127 diag.final_relative_residual = r_norm / rhs_norm;
6128 diag.stopping_reason = PcgStopReason::MaxIter;
6129 }
6130 let x = stream
6131 .clone_dtoh(&x_dev)
6132 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6133 Ok((Array1::from_vec(x), diag))
6134 }
6135
6136 pub(super) fn solve_reduced_beta_pcg_with_diagnostics(
6137 s_acc: &ndarray::Array2<f64>,
6138 rhs_beta: &Array1<f64>,
6139 max_iterations: usize,
6140 relative_tolerance: f64,
6141 ) -> Result<(Array1<f64>, ArrowPcgDiagnostics), ArrowSchurGpuFailure> {
6142 let k = rhs_beta.len();
6143 let cg_iters = max_iterations.max(1);
6155 let runtime = gam_gpu::linalg_dispatch::route_through_gpu(
6156 gam_gpu::linalg_dispatch::DispatchOp::Gemm {
6157 m: k,
6158 n: k,
6159 k: cg_iters,
6160 },
6161 )
6162 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
6163 let stream = gam_gpu::device_runtime::cuda_context_for(runtime.device.ordinal)
6164 .and_then(|ctx| ctx.new_stream().ok())
6165 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
6166 let blas = CudaBlas::new(stream.clone()).map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6167 let ctx = gam_gpu::device_runtime::cuda_context_for(runtime.device.ordinal)
6168 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
6169 let vector_module = pcg_vector_module(&ctx)?;
6170
6171 let mut inv_diag = vec![0.0_f64; k];
6173 for j in 0..k {
6174 let djj = s_acc[[j, j]];
6175 if !(djj > 0.0) {
6176 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
6177 reason: format!(
6178 "reduced-β GPU PCG: Jacobi diagonal S[{j},{j}]={djj:e} not positive"
6179 ),
6180 });
6181 }
6182 inv_diag[j] = 1.0 / djj;
6183 }
6184
6185 let mut s_host = vec![0.0_f64; k * k];
6187 for col in 0..k {
6188 for row in 0..k {
6189 s_host[col * k + row] = s_acc[[row, col]];
6190 }
6191 }
6192 let s_dev = stream
6193 .clone_htod(&s_host)
6194 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6195
6196 let rhs_norm = rhs_beta.iter().map(|v| v * v).sum::<f64>().sqrt();
6200 if rhs_norm == 0.0 {
6201 return Ok((Array1::<f64>::zeros(k), ArrowPcgDiagnostics::default()));
6202 }
6203 let tol = (relative_tolerance.max(0.0) * rhs_norm).max(1e-12);
6204
6205 let mut x_dev = stream
6208 .alloc_zeros::<f64>(k)
6209 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6210 let mut r_dev = stream
6211 .clone_htod(
6212 rhs_beta
6213 .as_slice()
6214 .ok_or(ArrowSchurGpuFailure::Unavailable)?,
6215 )
6216 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6217 let inv_diag_dev = stream
6218 .clone_htod(&inv_diag)
6219 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6220 let mut z_dev = stream
6221 .alloc_zeros::<f64>(k)
6222 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6223 launch_jacobi_mul(&stream, vector_module, &inv_diag_dev, &r_dev, &mut z_dev, k)?;
6224 let mut p_dev = stream
6225 .alloc_zeros::<f64>(k)
6226 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6227 device_copy(&blas, &stream, k, &z_dev, &mut p_dev)?;
6228 let mut sp_dev = stream
6229 .alloc_zeros::<f64>(k)
6230 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6231 let mut rz = device_dot(&blas, &stream, k, &r_dev, &z_dev)?;
6232 let mut diag = ArrowPcgDiagnostics {
6233 precond_apply_calls: 1,
6234 stopping_reason: PcgStopReason::MaxIter,
6235 ..ArrowPcgDiagnostics::default()
6236 };
6237 if rz <= 0.0 || !rz.is_finite() {
6238 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
6239 reason: format!("reduced-β GPU PCG: non-positive initial rᵀM⁻¹r={rz:e}"),
6240 });
6241 }
6242
6243 let max_iters = max_iterations.max(1);
6244 for _ in 0..max_iters {
6245 let gemv_cfg = GemvConfig::<f64> {
6247 trans: cublasOperation_t::CUBLAS_OP_N,
6248 m: to_i32(k).ok_or(ArrowSchurGpuFailure::Unavailable)?,
6249 n: to_i32(k).ok_or(ArrowSchurGpuFailure::Unavailable)?,
6250 alpha: 1.0,
6251 lda: to_i32(k).ok_or(ArrowSchurGpuFailure::Unavailable)?,
6252 incx: 1,
6253 beta: 0.0,
6254 incy: 1,
6255 };
6256 unsafe { blas.gemv(gemv_cfg, &s_dev, &p_dev, &mut sp_dev) }
6258 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6259 diag.matvec_calls += 1;
6260 diag.iterations += 1;
6261
6262 let p_sp = device_dot(&blas, &stream, k, &p_dev, &sp_dev)?;
6263 if !(p_sp > 0.0) {
6264 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
6267 reason: format!("reduced-β GPU PCG: non-positive curvature pᵀSp={p_sp:e}"),
6268 });
6269 }
6270 let alpha = rz / p_sp;
6271 device_axpy(&blas, &stream, k, alpha, &p_dev, &mut x_dev)?;
6272 device_axpy(&blas, &stream, k, -alpha, &sp_dev, &mut r_dev)?;
6273 let r_norm = device_nrm2(&blas, &stream, k, &r_dev)?;
6274 if r_norm <= tol {
6275 diag.final_relative_residual = r_norm / rhs_norm;
6276 diag.stopping_reason = PcgStopReason::Converged;
6277 break;
6278 }
6279 launch_jacobi_mul(&stream, vector_module, &inv_diag_dev, &r_dev, &mut z_dev, k)?;
6280 diag.precond_apply_calls += 1;
6281 let rz_new = device_dot(&blas, &stream, k, &r_dev, &z_dev)?;
6282 if rz_new <= 0.0 || !rz_new.is_finite() {
6283 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
6284 reason: format!("reduced-β GPU PCG: non-positive rᵀM⁻¹r={rz_new:e}"),
6285 });
6286 }
6287 let beta = rz_new / rz;
6288 launch_update_p(&stream, vector_module, &z_dev, beta, &mut p_dev, k)?;
6289 rz = rz_new;
6290 }
6291 if diag.stopping_reason != PcgStopReason::Converged {
6292 let r_norm = device_nrm2(&blas, &stream, k, &r_dev)?;
6293 diag.final_relative_residual = r_norm / rhs_norm;
6294 diag.stopping_reason = PcgStopReason::MaxIter;
6295 }
6296
6297 let x = stream
6298 .clone_dtoh(&x_dev)
6299 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6300 Ok((Array1::from_vec(x), diag))
6301 }
6302
6303 fn device_copy(
6304 blas: &CudaBlas,
6305 stream: &Arc<CudaStream>,
6306 n: usize,
6307 src: &CudaSlice<f64>,
6308 dst: &mut CudaSlice<f64>,
6309 ) -> Result<(), ArrowSchurGpuFailure> {
6310 let n_i = to_i32(n).ok_or(ArrowSchurGpuFailure::Unavailable)?;
6311 let (src_ptr, _src_rec) = src.device_ptr(stream);
6312 let (dst_ptr, _dst_rec) = dst.device_ptr_mut(stream);
6313 let status = unsafe {
6316 cudarc::cublas::sys::cublasDcopy_v2(
6317 *blas.handle(),
6318 n_i,
6319 src_ptr as *const f64,
6320 1,
6321 dst_ptr as *mut f64,
6322 1,
6323 )
6324 };
6325 if status == cublasStatus_t::CUBLAS_STATUS_SUCCESS {
6326 Ok(())
6327 } else {
6328 Err(ArrowSchurGpuFailure::Unavailable)
6329 }
6330 }
6331
6332 fn device_axpy(
6333 blas: &CudaBlas,
6334 stream: &Arc<CudaStream>,
6335 n: usize,
6336 alpha: f64,
6337 x: &CudaSlice<f64>,
6338 y: &mut CudaSlice<f64>,
6339 ) -> Result<(), ArrowSchurGpuFailure> {
6340 let n_i = to_i32(n).ok_or(ArrowSchurGpuFailure::Unavailable)?;
6341 let (x_ptr, _x_rec) = x.device_ptr(stream);
6342 let (y_ptr, _y_rec) = y.device_ptr_mut(stream);
6343 let status = unsafe {
6346 cudarc::cublas::sys::cublasDaxpy_v2(
6347 *blas.handle(),
6348 n_i,
6349 &alpha,
6350 x_ptr as *const f64,
6351 1,
6352 y_ptr as *mut f64,
6353 1,
6354 )
6355 };
6356 if status == cublasStatus_t::CUBLAS_STATUS_SUCCESS {
6357 Ok(())
6358 } else {
6359 Err(ArrowSchurGpuFailure::Unavailable)
6360 }
6361 }
6362
6363 fn device_axpy_strided(
6368 blas: &CudaBlas,
6369 stream: &Arc<CudaStream>,
6370 n: usize,
6371 alpha: f64,
6372 x: &CudaSlice<f64>,
6373 incx: usize,
6374 y: &mut CudaSlice<f64>,
6375 incy: usize,
6376 ) -> Result<(), ArrowSchurGpuFailure> {
6377 let n_i = to_i32(n).ok_or(ArrowSchurGpuFailure::Unavailable)?;
6378 let incx_i = to_i32(incx).ok_or(ArrowSchurGpuFailure::Unavailable)?;
6379 let incy_i = to_i32(incy).ok_or(ArrowSchurGpuFailure::Unavailable)?;
6380 let (x_ptr, _x_rec) = x.device_ptr(stream);
6381 let (y_ptr, _y_rec) = y.device_ptr_mut(stream);
6382 let status = unsafe {
6385 cudarc::cublas::sys::cublasDaxpy_v2(
6386 *blas.handle(),
6387 n_i,
6388 &alpha,
6389 x_ptr as *const f64,
6390 incx_i,
6391 y_ptr as *mut f64,
6392 incy_i,
6393 )
6394 };
6395 if status == cublasStatus_t::CUBLAS_STATUS_SUCCESS {
6396 Ok(())
6397 } else {
6398 Err(ArrowSchurGpuFailure::Unavailable)
6399 }
6400 }
6401
6402 fn device_dot(
6403 blas: &CudaBlas,
6404 stream: &Arc<CudaStream>,
6405 n: usize,
6406 x: &CudaSlice<f64>,
6407 y: &CudaSlice<f64>,
6408 ) -> Result<f64, ArrowSchurGpuFailure> {
6409 let n_i = to_i32(n).ok_or(ArrowSchurGpuFailure::Unavailable)?;
6410 let (x_ptr, _x_rec) = x.device_ptr(stream);
6411 let (y_ptr, _y_rec) = y.device_ptr(stream);
6412 let mut result = 0.0_f64;
6413 let status = unsafe {
6417 cudarc::cublas::sys::cublasDdot_v2(
6418 *blas.handle(),
6419 n_i,
6420 x_ptr as *const f64,
6421 1,
6422 y_ptr as *const f64,
6423 1,
6424 &mut result,
6425 )
6426 };
6427 if status == cublasStatus_t::CUBLAS_STATUS_SUCCESS {
6428 Ok(result)
6429 } else {
6430 Err(ArrowSchurGpuFailure::Unavailable)
6431 }
6432 }
6433
6434 fn device_nrm2(
6435 blas: &CudaBlas,
6436 stream: &Arc<CudaStream>,
6437 n: usize,
6438 x: &CudaSlice<f64>,
6439 ) -> Result<f64, ArrowSchurGpuFailure> {
6440 let n_i = to_i32(n).ok_or(ArrowSchurGpuFailure::Unavailable)?;
6441 let (x_ptr, _x_rec) = x.device_ptr(stream);
6442 let mut result = 0.0_f64;
6443 let status = unsafe {
6447 cudarc::cublas::sys::cublasDnrm2_v2(
6448 *blas.handle(),
6449 n_i,
6450 x_ptr as *const f64,
6451 1,
6452 &mut result,
6453 )
6454 };
6455 if status == cublasStatus_t::CUBLAS_STATUS_SUCCESS {
6456 Ok(result)
6457 } else {
6458 Err(ArrowSchurGpuFailure::Unavailable)
6459 }
6460 }
6461
6462 #[cfg(test)]
6463 mod tests {
6464 use super::*;
6469 use crate::arrow_schur::{
6470 ArrowSchurSystem, DeviceSaeFrameData, DeviceSaePcgData, DeviceSaeSmoothBlock,
6471 FactoredFrameGBlock,
6472 };
6473 use ndarray::Array2;
6474
6475 fn tiny_framed_fixture() -> (ArrowSchurSystem, DeviceSaePcgData) {
6479 let p = 3usize;
6480 let ranks = vec![2usize, 3usize];
6481 let basis_sizes = vec![2usize, 2usize];
6482 let mut border_offsets = Vec::new();
6483 let mut acc = 0usize;
6484 for k in 0..2 {
6485 border_offsets.push(acc);
6486 acc += basis_sizes[k] * ranks[k];
6487 }
6488 let border_dim = acc;
6489 let frame_of = |k: usize| -> Array2<f64> {
6490 Array2::from_shape_fn((p, ranks[k]), |(i, j)| {
6491 0.1 + 0.2 * ((i + 1) as f64) * ((j + 1 + 2 * k) as f64)
6492 })
6493 };
6494 let frames: Vec<Array2<f64>> = (0..2).map(frame_of).collect();
6495 let w_of = |i: usize, j: usize| -> Array2<f64> {
6496 let (ui, uj) = (&frames[i], &frames[j]);
6497 Array2::from_shape_fn((ranks[i], ranks[j]), |(a, b)| {
6498 (0..p).map(|c| ui[[c, a]] * uj[[c, b]]).sum()
6499 })
6500 };
6501 let mut frame_blocks = Vec::new();
6502 for &(i, j) in &[(0usize, 0usize), (1usize, 1usize), (0, 1), (1, 0)] {
6503 let (mi, mj) = (basis_sizes[i], basis_sizes[j]);
6504 let mut g =
6505 Array2::<f64>::from_shape_fn((mi, mj), |(r, c)| 0.1 * (r + 2 * c + 1) as f64);
6506 if i == j {
6507 for r in 0..mi.min(mj) {
6508 g[[r, r]] += mi as f64 + 2.0;
6509 }
6510 }
6511 frame_blocks.push(FactoredFrameGBlock {
6512 atom_i: i,
6513 atom_j: j,
6514 g,
6515 w: w_of(i, j),
6516 });
6517 }
6518 let mut smooth_blocks = Vec::new();
6519 for k in 0..2 {
6520 let m = basis_sizes[k];
6521 let mut s =
6522 Array2::<f64>::from_shape_fn((m, m), |(r, c)| 0.05 * (r + c + 1) as f64);
6523 for r in 0..m {
6524 s[[r, r]] += 1.0;
6525 }
6526 smooth_blocks.push(DeviceSaeSmoothBlock {
6527 global_offset: border_offsets[k],
6528 factor_a: s,
6529 });
6530 }
6531 let smooth_ranks = ranks.clone();
6532 let n = 2usize;
6533 let q = 2usize;
6534 let mut sys = ArrowSchurSystem::new(n, q, border_dim);
6535 let mut row_htbeta = Vec::new();
6536 for i in 0..n {
6537 let mut htt =
6538 Array2::<f64>::from_shape_fn((q, q), |(r, c)| 0.3 * (r + c + 1) as f64);
6539 for r in 0..q {
6540 htt[[r, r]] += q as f64 + 2.0;
6541 }
6542 sys.rows[i].htt = htt;
6543 let mut slab = vec![0.0_f64; q * border_dim];
6544 for c in 0..q {
6545 for col in 0..border_dim {
6546 let v = 0.01 * ((c + 1) * (col + 1) + i) as f64;
6547 slab[c * border_dim + col] = v;
6548 sys.rows[i].htbeta[[c, col]] = v;
6549 }
6550 }
6551 row_htbeta.push(slab);
6552 }
6553 let data = DeviceSaePcgData {
6554 p,
6555 beta_dim: border_dim,
6556 a_phi: std::sync::Arc::from(Vec::new().into_boxed_slice()),
6557 local_jac: std::sync::Arc::from(Vec::new().into_boxed_slice()),
6558 smooth_blocks,
6559 sparse_g_blocks: Vec::new(),
6560 frame: Some(DeviceSaeFrameData {
6561 ranks,
6562 basis_sizes,
6563 border_offsets,
6564 frame_blocks,
6565 smooth_ranks,
6566 row_htbeta,
6567 }),
6568 };
6569 (sys, data)
6570 }
6571
6572 #[test]
6584 fn sae_resident_frame_only_ainv_is_ridge_dependent_1017() {
6585 let (sys, data) = tiny_framed_fixture();
6586 let frame = data.frame.as_ref().expect("framed fixture");
6587
6588 let host_a = super::super::flatten_frame_host_operands(&sys, &data, frame)
6591 .expect("host operands a");
6592 let host_b = super::super::flatten_frame_host_operands(&sys, &data, frame)
6593 .expect("host operands b");
6594 assert_eq!(
6595 host_a.s_data, host_b.s_data,
6596 "smooth λS must be ridge-independent"
6597 );
6598 assert_eq!(
6599 host_a.g_data, host_b.g_data,
6600 "frame G must be ridge-independent"
6601 );
6602 assert_eq!(
6603 host_a.w_data, host_b.w_data,
6604 "frame W must be ridge-independent"
6605 );
6606 assert_eq!(host_a.htb, host_b.htb, "row H_tβ must be ridge-independent");
6607 assert_eq!(host_a.q_of, host_b.q_of);
6608
6609 let report = data.operand_byte_report();
6610 assert!(report.framed, "fixture must exercise the framed lane");
6611 assert!(
6612 report.total_bytes > 0,
6613 "framed operands must have nonzero bytes"
6614 );
6615
6616 let ainv_lo = super::super::compute_ainv_host(
6620 &sys,
6621 &host_a.q_of,
6622 host_a.max_q,
6623 host_a.n_rows,
6624 1e-3,
6625 )
6626 .expect("ainv lo");
6627 let ainv_lo2 = super::super::compute_ainv_host(
6628 &sys,
6629 &host_a.q_of,
6630 host_a.max_q,
6631 host_a.n_rows,
6632 1e-3,
6633 )
6634 .expect("ainv lo repeat");
6635 let ainv_hi = super::super::compute_ainv_host(
6636 &sys,
6637 &host_a.q_of,
6638 host_a.max_q,
6639 host_a.n_rows,
6640 1e3,
6641 )
6642 .expect("ainv hi");
6643 assert_eq!(
6644 ainv_lo, ainv_lo2,
6645 "ainv must be deterministic at a fixed ridge"
6646 );
6647 let max_diff = ainv_lo
6648 .iter()
6649 .zip(&ainv_hi)
6650 .map(|(a, b)| (a - b).abs())
6651 .fold(0.0_f64, f64::max);
6652 assert!(
6653 max_diff > 1e-6,
6654 "ainv MUST change with ridge_t (else the split is wrong); max_diff={max_diff:e}"
6655 );
6656
6657 let trials = crate::arrow_schur::DEFAULT_PROXIMAL_MAX_ATTEMPTS + 1;
6661 let saved = report.total_bytes * (trials - 1);
6662 assert!(saved > 0);
6663 eprintln!(
6664 "#1017 resident-frame ladder saving: {trials} trials × {}B ridge-independent \
6665 operand upload → resident removes {saved}B, re-uploading only ainv \
6666 ({}rows × {}² × 8B) per trial",
6667 report.total_bytes, host_a.n_rows, host_a.max_q
6668 );
6669 }
6670
6671 fn device_matvec_once(
6674 sys: &ArrowSchurSystem,
6675 data: &DeviceSaePcgData,
6676 ridge_t: f64,
6677 ridge_beta: f64,
6678 x_host: &[f64],
6679 ) -> Result<Vec<f64>, ArrowSchurGpuFailure> {
6680 let k = x_host.len();
6681 let frame = data
6682 .frame
6683 .as_ref()
6684 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
6685 let runtime = super::super::resolve_runtime_for_device_path()?
6686 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
6687 let ctx = gam_gpu::device_runtime::cuda_context_for(runtime.selected_device().ordinal)
6688 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
6689 let stream = ctx
6690 .new_stream()
6691 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6692 let vector_module = pcg_vector_module(&ctx)?;
6693 let mut buffers = flatten_device_sae_frame_data(sys, data, frame, ridge_t, &stream)?;
6694 let x_dev = stream
6695 .clone_htod(x_host)
6696 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6697 let mut out_dev = stream
6698 .alloc_zeros::<f64>(k)
6699 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6700 launch_sae_frame_matvec(
6701 &stream,
6702 vector_module,
6703 &mut buffers,
6704 &x_dev,
6705 &mut out_dev,
6706 ridge_beta,
6707 )?;
6708 stream
6709 .clone_dtoh(&out_dev)
6710 .map_err(|_| ArrowSchurGpuFailure::Unavailable)
6711 }
6712
6713 #[test]
6719 fn framed_sae_device_matvec_stage_diff_tiny_1551() {
6720 if gam_gpu::device_runtime::GpuRuntime::resolve(gam_gpu::GpuPolicy::Auto)
6721 .unwrap_or_else(|error| panic!("GPU probe fault in framed matvec test: {error}"))
6722 .is_none()
6723 {
6724 return;
6725 }
6726 let p = 3usize;
6727 let ranks = vec![2usize, 3usize];
6728 let basis_sizes = vec![2usize, 2usize];
6729 let mut border_offsets = Vec::new();
6730 let mut acc = 0usize;
6731 for k in 0..2 {
6732 border_offsets.push(acc);
6733 acc += basis_sizes[k] * ranks[k];
6734 }
6735 let border_dim = acc; let frame_of = |k: usize| -> Array2<f64> {
6737 Array2::from_shape_fn((p, ranks[k]), |(i, j)| {
6738 0.1 + 0.2 * ((i + 1) as f64) * ((j + 1 + 2 * k) as f64)
6739 })
6740 };
6741 let frames: Vec<Array2<f64>> = (0..2).map(frame_of).collect();
6742 let w_of = |i: usize, j: usize| -> Array2<f64> {
6743 let (ui, uj) = (&frames[i], &frames[j]);
6744 Array2::from_shape_fn((ranks[i], ranks[j]), |(a, b)| {
6745 (0..p).map(|c| ui[[c, a]] * uj[[c, b]]).sum()
6746 })
6747 };
6748 let mut frame_blocks = Vec::new();
6749 for &(i, j) in &[(0usize, 0usize), (1usize, 1usize), (0, 1), (1, 0)] {
6750 let (mi, mj) = (basis_sizes[i], basis_sizes[j]);
6751 let mut g =
6752 Array2::<f64>::from_shape_fn((mi, mj), |(r, c)| 0.1 * (r + 2 * c + 1) as f64);
6753 if i == j {
6754 for r in 0..mi.min(mj) {
6755 g[[r, r]] += mi as f64 + 2.0;
6756 }
6757 }
6758 frame_blocks.push(FactoredFrameGBlock {
6759 atom_i: i,
6760 atom_j: j,
6761 g,
6762 w: w_of(i, j),
6763 });
6764 }
6765 let mut smooth_blocks = Vec::new();
6766 for k in 0..2 {
6767 let m = basis_sizes[k];
6768 let mut s =
6769 Array2::<f64>::from_shape_fn((m, m), |(r, c)| 0.05 * (r + c + 1) as f64);
6770 for r in 0..m {
6771 s[[r, r]] += 1.0;
6772 }
6773 smooth_blocks.push(DeviceSaeSmoothBlock {
6774 global_offset: border_offsets[k],
6775 factor_a: s,
6776 });
6777 }
6778 let smooth_ranks = ranks.clone();
6779 let n = 2usize;
6780 let q = 2usize;
6781 let mut sys = ArrowSchurSystem::new(n, q, border_dim);
6782 let mut row_htbeta = Vec::new();
6783 for i in 0..n {
6784 let mut htt =
6785 Array2::<f64>::from_shape_fn((q, q), |(r, c)| 0.3 * (r + c + 1) as f64);
6786 for r in 0..q {
6787 htt[[r, r]] += q as f64 + 2.0;
6788 }
6789 sys.rows[i].htt = htt;
6790 let mut slab = vec![0.0_f64; q * border_dim];
6791 for c in 0..q {
6792 for col in 0..border_dim {
6793 let v = 0.01 * ((c + 1) * (col + 1) + i) as f64;
6794 slab[c * border_dim + col] = v;
6795 sys.rows[i].htbeta[[c, col]] = v;
6796 }
6797 }
6798 row_htbeta.push(slab);
6799 }
6800 let data = DeviceSaePcgData {
6801 p,
6802 beta_dim: border_dim,
6803 a_phi: std::sync::Arc::from(Vec::new().into_boxed_slice()),
6804 local_jac: std::sync::Arc::from(Vec::new().into_boxed_slice()),
6805 smooth_blocks,
6806 sparse_g_blocks: Vec::new(),
6807 frame: Some(DeviceSaeFrameData {
6808 ranks,
6809 basis_sizes,
6810 border_offsets,
6811 frame_blocks,
6812 smooth_ranks,
6813 row_htbeta,
6814 }),
6815 };
6816 let ridge_t = 1e-7;
6817 let ridge_beta = 1e-6;
6818 let mut first_bad: Option<usize> = None;
6819 let mut worst = 0.0_f64;
6820 let mut worst_at = 0usize;
6821 let mut worst_dev = 0.0_f64;
6822 let mut worst_cpu = 0.0_f64;
6823 for col in 0..border_dim {
6824 let mut x = vec![0.0_f64; border_dim];
6825 x[col] = 1.0;
6826 let dev = match device_matvec_once(&sys, &data, ridge_t, ridge_beta, &x) {
6827 Ok(v) => v,
6828 Err(_) => return,
6829 };
6830 let mut cpu = vec![0.0_f64; border_dim];
6831 super::super::sae_framed_schur_matvec_cpu(
6832 &sys, &data, ridge_t, ridge_beta, &x, &mut cpu,
6833 )
6834 .expect("cpu matvec");
6835 for r in 0..border_dim {
6836 let d = (dev[r] - cpu[r]).abs();
6837 if d > 1e-9 && first_bad.is_none() {
6838 first_bad = Some(r * border_dim + col);
6839 }
6840 if d > worst {
6841 worst = d;
6842 worst_at = r * border_dim + col;
6843 worst_dev = dev[r];
6844 worst_cpu = cpu[r];
6845 }
6846 }
6847 }
6848 assert!(
6849 worst <= 1e-9,
6850 "[#1551 stage-diff] device framed matvec != CPU oracle: worst abs={worst:e} at \
6851 (row*K+col)={worst_at} (dev={worst_dev:e} cpu={worst_cpu:e}), \
6852 first_bad_idx={first_bad:?}; border layout: atom0 [0..4) rank2, atom1 [4..10) \
6853 rank3 — which atom-range the bad row/col falls in pins the stage (smooth=diag, \
6854 G⊗W=cross, reduced-Schur=dense per-row)",
6855 );
6856 }
6857 }
6858}
6859
6860#[cfg(test)]
6861mod tests {
6862 use super::*;
6863 use crate::arrow_schur::ArrowSchurSystem;
6864 use ndarray::{Array2, ArrayView1};
6865
6866 fn build_fixture(n: usize, d: usize, k: usize, seed: u64) -> ArrowSchurSystem {
6867 let mut sys = ArrowSchurSystem::new(n, d, k);
6868 let mut state = seed.wrapping_mul(0x9E37_79B9_7F4A_7C15);
6869 let mut sample = || -> f64 {
6870 state = state
6871 .wrapping_mul(6364136223846793005)
6872 .wrapping_add(1442695040888963407);
6873 ((state >> 33) as f64) / ((1u64 << 31) as f64) - 1.0
6874 };
6875 for row in &mut sys.rows {
6876 let mut a = Array2::<f64>::zeros((d, d));
6877 for r in 0..d {
6878 for c in 0..d {
6879 a[[r, c]] = sample();
6880 }
6881 }
6882 let mut htt = a.t().dot(&a);
6883 for r in 0..d {
6884 htt[[r, r]] += d as f64 + 1.0;
6885 }
6886 row.htt = htt;
6887 for r in 0..d {
6888 for c in 0..k {
6889 row.htbeta[[r, c]] = 0.1 * sample();
6890 }
6891 row.gt[r] = sample();
6892 }
6893 }
6894 let mut hbb_a = Array2::<f64>::zeros((k, k));
6895 for r in 0..k {
6896 for c in 0..k {
6897 hbb_a[[r, c]] = sample();
6898 }
6899 }
6900 let mut hbb = hbb_a.t().dot(&hbb_a);
6901 for r in 0..k {
6902 hbb[[r, r]] += k as f64 + 1.0;
6903 }
6904 sys.hbb = hbb;
6905 for r in 0..k {
6906 sys.gb[r] = sample();
6907 }
6908 sys
6909 }
6910
6911 #[test]
6916 fn ridge_bump_makes_known_indefinite_blocks_pd() {
6917 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));
6925 indef2[[0, 0]] = 1.0;
6926 indef2[[1, 1]] = 1.0;
6927 indef2[[0, 1]] = 2.0;
6928 indef2[[1, 0]] = 2.0;
6929 let pd = Array2::<f64>::from_diag(&Array1::from_elem(3, 5.0));
6932
6933 for (label, block) in [
6934 ("-I (λ_min=-1)", neg_identity),
6935 ("-250·I (λ_min=-250)", scaled_neg),
6936 ("[[1,2],[2,1]] (λ_min=-1)", indef2),
6937 ("5·I (PD)", pd),
6938 ] {
6939 let ridge_t = 0.0;
6940 let bump = ridge_bump_to_make_pd(block.view(), ridge_t);
6941 assert!(
6942 bump > 0.0 && bump.is_finite(),
6943 "[{label}] bump must be strictly positive and finite, got {bump:e}"
6944 );
6945 let d = block.nrows();
6946 let mut shifted = block.clone();
6947 for i in 0..d {
6948 shifted[[i, i]] += ridge_t + bump;
6949 }
6950 assert!(
6951 cholesky_factor_in_place(shifted.view(), CholeskyGuard::NonnegativePivot).is_some(),
6952 "[{label}] H_tt + (ridge_t + bump={bump:e})·I must be PD after the \
6953 Gershgorin bump, but the Cholesky still rejected it"
6954 );
6955 }
6956 }
6957
6958 #[cfg(target_os = "linux")]
6968 #[test]
6969 fn ridge_bump_colmajor_matches_rowmajor_for_symmetric_block() {
6970 let mut a = Array2::<f64>::zeros((3, 3));
6972 a[[0, 0]] = -2.0;
6973 a[[1, 1]] = 0.5;
6974 a[[2, 2]] = 1.0;
6975 a[[0, 1]] = 0.3;
6976 a[[1, 0]] = 0.3;
6977 a[[1, 2]] = -0.4;
6978 a[[2, 1]] = -0.4;
6979 a[[0, 2]] = 0.1;
6980 a[[2, 0]] = 0.1;
6981
6982 let row_major_bump = ridge_bump_to_make_pd(a.view(), 0.0);
6983
6984 let d = 3;
6986 let mut col_major = vec![0.0_f64; d * d];
6987 for c in 0..d {
6988 for r in 0..d {
6989 col_major[c * d + r] = a[[r, c]];
6990 }
6991 }
6992 let col_major_bump = ridge_bump_to_make_pd_colmajor(&col_major, d);
6993
6994 assert!(
6995 (row_major_bump - col_major_bump).abs() <= 1e-12 * row_major_bump.max(1.0),
6996 "colmajor bump {col_major_bump:e} must match rowmajor bump \
6997 {row_major_bump:e} for a symmetric block"
6998 );
6999
7000 let mut shifted = a.clone();
7002 for i in 0..d {
7003 shifted[[i, i]] += col_major_bump;
7004 }
7005 assert!(
7006 cholesky_factor_in_place(shifted.view(), CholeskyGuard::NonnegativePivot).is_some(),
7007 "colmajor Gershgorin bump must make the symmetric block PD"
7008 );
7009 }
7010
7011 fn device_pcg_fixture(k: usize) -> (Array2<f64>, Array1<f64>) {
7012 let mut s = Array2::<f64>::zeros((k, k));
7013 for row in 0..k {
7014 s[[row, row]] = 2.5 + 0.001 * ((row % 17) as f64);
7015 if row + 1 < k {
7016 s[[row, row + 1]] = -0.05;
7017 s[[row + 1, row]] = -0.05;
7018 }
7019 if row + 7 < k {
7020 s[[row, row + 7]] = 0.01;
7021 s[[row + 7, row]] = 0.01;
7022 }
7023 }
7024 let rhs = Array1::from_shape_fn(k, |idx| ((idx as f64 + 1.0) * 0.013).sin());
7025 (s, rhs)
7026 }
7027
7028 fn dense_pcg_cpu_reference(
7029 s: &Array2<f64>,
7030 rhs: &Array1<f64>,
7031 max_iterations: usize,
7032 relative_tolerance: f64,
7033 ) -> Array1<f64> {
7034 let k = rhs.len();
7035 let rhs_norm = rhs.iter().map(|v| v * v).sum::<f64>().sqrt();
7036 if rhs_norm == 0.0 {
7037 return Array1::<f64>::zeros(k);
7038 }
7039 let tol = (relative_tolerance.max(0.0) * rhs_norm).max(1e-12);
7040 let inv_diag: Vec<f64> = (0..k).map(|idx| 1.0 / s[[idx, idx]]).collect();
7041 let mut x = Array1::<f64>::zeros(k);
7042 let mut r = rhs.clone();
7043 let mut z = Array1::from_shape_fn(k, |idx| inv_diag[idx] * r[idx]);
7044 let mut p = z.clone();
7045 let mut sp = Array1::<f64>::zeros(k);
7046 let mut rz = r.iter().zip(z.iter()).map(|(a, b)| a * b).sum::<f64>();
7047 for _ in 0..max_iterations.max(1) {
7048 for row in 0..k {
7049 let mut acc = 0.0;
7050 for col in 0..k {
7051 acc += s[[row, col]] * p[col];
7052 }
7053 sp[row] = acc;
7054 }
7055 let p_sp = p.iter().zip(sp.iter()).map(|(a, b)| a * b).sum::<f64>();
7056 let alpha = rz / p_sp;
7057 for idx in 0..k {
7058 x[idx] += alpha * p[idx];
7059 r[idx] -= alpha * sp[idx];
7060 }
7061 let r_norm = r.iter().map(|v| v * v).sum::<f64>().sqrt();
7062 if r_norm <= tol {
7063 break;
7064 }
7065 for idx in 0..k {
7066 z[idx] = inv_diag[idx] * r[idx];
7067 }
7068 let rz_next = r.iter().zip(z.iter()).map(|(a, b)| a * b).sum::<f64>();
7069 let beta = rz_next / rz;
7070 for idx in 0..k {
7071 p[idx] = z[idx] + beta * p[idx];
7072 }
7073 rz = rz_next;
7074 }
7075 x
7076 }
7077
7078 #[test]
7079 fn device_resident_pcg_matches_cpu_reference_when_cuda_admits() {
7080 let (s, rhs) = device_pcg_fixture(512);
7081 let max_iterations = 200usize;
7082 let relative_tolerance = 1.0e-12;
7083 let cpu = dense_pcg_cpu_reference(&s, &rhs, max_iterations, relative_tolerance);
7084 let (device, diag) = match solve_reduced_beta_pcg_with_diagnostics(
7085 &s,
7086 &rhs,
7087 max_iterations,
7088 relative_tolerance,
7089 ) {
7090 Ok(result) => result,
7091 Err(failure) => {
7098 assert!(
7099 gam_gpu::device_runtime::GpuRuntime::resolve(gam_gpu::GpuPolicy::Auto)
7100 .unwrap_or_else(|error| {
7101 panic!("GPU probe fault in reduced-beta PCG test: {error}")
7102 })
7103 .is_none(),
7104 "#1017: CUDA device present but the device reduced-beta PCG \
7105 declined/faulted instead of returning a result (tag: {failure:?}) — \
7106 the kernel does not run correctly on GPU"
7107 );
7108 return;
7109 }
7110 };
7111 let max_err = cpu
7112 .iter()
7113 .zip(device.iter())
7114 .map(|(a, b)| (a - b).abs())
7115 .fold(0.0_f64, f64::max);
7116 assert!(
7117 max_err <= 1.0e-10,
7118 "device resident PCG parity failed: max_err={max_err:e}, diag={diag:?}"
7119 );
7120 assert!(diag.matvec_calls > 0);
7121 assert_eq!(diag.matvec_calls, diag.iterations);
7122 }
7123
7124 #[test]
7125 fn dense_reference_matches_independent_solve() {
7126 let sys = build_fixture(4, 5, 3, 7);
7127 let solution = solve_arrow_newton_step_dense_reference(&sys, 0.0, 0.0).unwrap();
7128 let n = sys.rows.len();
7132 let d = sys.d;
7133 let k = sys.k;
7134 let total = n * d + k;
7135 let mut h = Array2::<f64>::zeros((total, total));
7136 let mut g = ndarray::Array1::<f64>::zeros(total);
7137 for (i, row) in sys.rows.iter().enumerate() {
7138 let base = i * d;
7139 for c in 0..d {
7140 for r in 0..d {
7141 h[[base + r, base + c]] = row.htt[[r, c]];
7142 }
7143 }
7144 for c in 0..k {
7145 for r in 0..d {
7146 h[[base + r, n * d + c]] = row.htbeta[[r, c]];
7147 h[[n * d + c, base + r]] = row.htbeta[[r, c]];
7148 }
7149 }
7150 for r in 0..d {
7151 g[base + r] = row.gt[r];
7152 }
7153 }
7154 for c in 0..k {
7155 for r in 0..k {
7156 h[[n * d + r, n * d + c]] += sys.hbb[[r, c]];
7157 }
7158 g[n * d + c] = sys.gb[c];
7159 }
7160 let l = cholesky_factor_in_place(h.view(), CholeskyGuard::NonnegativePivot).unwrap();
7161 let rhs = g.mapv(|v| -v);
7162 let expected = cholesky_solve_vector(l.view(), rhs.view());
7163 for i in 0..n * d {
7164 assert!(
7165 (solution.delta_t[i] - expected[i]).abs() < 1e-10 * (1.0 + expected[i].abs()),
7166 "delta_t[{i}] mismatch: got {} expected {}",
7167 solution.delta_t[i],
7168 expected[i]
7169 );
7170 }
7171 for a in 0..k {
7172 assert!(
7173 (solution.delta_beta[a] - expected[n * d + a]).abs()
7174 < 1e-10 * (1.0 + expected[n * d + a].abs()),
7175 "delta_beta[{a}] mismatch"
7176 );
7177 }
7178 }
7179
7180 #[test]
7194 fn row_procedural_matvec_parallel_deterministic_and_matches_serial() {
7195 use crate::arrow_schur::SCHUR_MATVEC_PARALLEL_ROW_MIN;
7196 let n = SCHUR_MATVEC_PARALLEL_ROW_MIN + 96; let d = 3usize;
7198 let k = 24usize;
7199 let mut sys = build_fixture(n, d, k, 0xA17C_0FFE);
7200 let slabs: Vec<Array2<f64>> = sys.rows.iter().map(|row| row.htbeta.clone()).collect();
7205 let forward_slabs = slabs.clone();
7206 let transpose_slabs = slabs;
7207 sys.set_row_htbeta_operator(
7208 move |row: usize, x: ArrayView1<'_, f64>, out: &mut Array1<f64>| {
7209 let h = &forward_slabs[row];
7210 for r in 0..h.nrows() {
7211 let mut acc = 0.0_f64;
7212 for c in 0..h.ncols() {
7213 acc += h[[r, c]] * x[c];
7214 }
7215 out[r] = acc;
7216 }
7217 },
7218 move |row: usize, v: ArrayView1<'_, f64>, out: &mut Array1<f64>| {
7219 let h = &transpose_slabs[row];
7220 for r in 0..h.nrows() {
7221 for c in 0..h.ncols() {
7222 out[c] += h[[r, c]] * v[r];
7223 }
7224 }
7225 },
7226 );
7227
7228 let matvec = gpu_schur_matvec_backend(&sys, 0.0, 0.0)
7229 .expect("row-procedural matvec backend builds for matrix-free system");
7230 let x = Array1::from_shape_fn(k, |i| ((i as f64 + 1.0) * 0.37).sin());
7231
7232 let mut out_parallel_a = Array1::<f64>::zeros(k);
7236 matvec(&x, &mut out_parallel_a);
7237 let mut out_parallel_b = Array1::<f64>::zeros(k);
7238 matvec(&x, &mut out_parallel_b);
7239 for a in 0..k {
7240 assert_eq!(
7241 out_parallel_a[a].to_bits(),
7242 out_parallel_b[a].to_bits(),
7243 "row-procedural matvec parallel reduction is non-deterministic at index {a}"
7244 );
7245 }
7246
7247 let mut out_serial = Array1::<f64>::zeros(k);
7252 rayon::ThreadPoolBuilder::new()
7253 .num_threads(2)
7254 .build()
7255 .expect("build rayon pool")
7256 .install(|| matvec(&x, &mut out_serial));
7257
7258 let max_abs = out_serial.iter().fold(0.0_f64, |m, v| m.max(v.abs()));
7259 for a in 0..k {
7260 let diff = (out_parallel_a[a] - out_serial[a]).abs();
7261 assert!(
7262 diff <= 1e-12 * (1.0 + max_abs),
7263 "row-procedural matvec parallel vs serial diverged beyond reassociation \
7264 at index {a}: {} vs {} (diff={diff:e})",
7265 out_parallel_a[a],
7266 out_serial[a]
7267 );
7268 }
7269 }
7270
7271 #[test]
7278 fn framed_sae_schur_matvec_matches_dense_reference() {
7279 use crate::arrow_schur::{
7280 BetaPenaltyOp, DeviceSaeFrameData, DeviceSaePcgData, DeviceSaeSmoothBlock,
7281 FactoredFrameGBlock, FactoredFrameKroneckerOp, IdentityRightKroneckerPenaltyOp,
7282 };
7283
7284 let p = 4usize;
7285 let ranks = vec![2usize, 4usize, 3usize];
7287 let basis_sizes = vec![2usize, 1usize, 2usize];
7288 let n_atoms = ranks.len();
7289 let mut border_offsets = Vec::with_capacity(n_atoms);
7290 let mut acc = 0usize;
7291 for k in 0..n_atoms {
7292 border_offsets.push(acc);
7293 acc += basis_sizes[k] * ranks[k];
7294 }
7295 let border_dim = acc; let mut state = 0x1234_5678_9abc_def0u64;
7298 let mut sample = || -> f64 {
7299 state = state
7300 .wrapping_mul(6364136223846793005)
7301 .wrapping_add(1442695040888963407);
7302 ((state >> 33) as f64) / ((1u64 << 31) as f64) - 1.0
7303 };
7304
7305 let mut frames: Vec<Array2<f64>> = Vec::with_capacity(n_atoms);
7308 for k in 0..n_atoms {
7309 let r = ranks[k];
7310 let mut u = Array2::<f64>::zeros((p, r));
7311 for i in 0..p {
7312 for j in 0..r {
7313 u[[i, j]] = if r == p && i == j {
7314 1.0
7315 } else if r == p {
7316 0.0
7317 } else {
7318 sample()
7319 };
7320 }
7321 }
7322 frames.push(u);
7323 }
7324 let w_of = |i: usize, j: usize| -> Array2<f64> {
7325 let (ui, uj) = (&frames[i], &frames[j]);
7326 let (ri, rj) = (ranks[i], ranks[j]);
7327 let mut w = Array2::<f64>::zeros((ri, rj));
7328 for a in 0..ri {
7329 for b in 0..rj {
7330 let mut s = 0.0;
7331 for c in 0..p {
7332 s += ui[[c, a]] * uj[[c, b]];
7333 }
7334 w[[a, b]] = s;
7335 }
7336 }
7337 w
7338 };
7339
7340 let mut frame_blocks: Vec<FactoredFrameGBlock> = Vec::new();
7342 let mut pairs = vec![(0usize, 0usize), (1, 1), (2, 2), (0, 2), (2, 0)];
7343 pairs.sort();
7344 for &(i, j) in &pairs {
7345 let (mi, mj) = (basis_sizes[i], basis_sizes[j]);
7346 let mut g = Array2::<f64>::zeros((mi, mj));
7347 for r in 0..mi {
7348 for c in 0..mj {
7349 g[[r, c]] = 0.3 * sample();
7350 }
7351 }
7352 if i == j {
7354 for r in 0..mi.min(mj) {
7355 g[[r, r]] += mi as f64 + 2.0;
7356 }
7357 }
7358 frame_blocks.push(FactoredFrameGBlock {
7359 atom_i: i,
7360 atom_j: j,
7361 g,
7362 w: w_of(i, j),
7363 });
7364 }
7365
7366 let mut smooth_blocks: Vec<DeviceSaeSmoothBlock> = Vec::with_capacity(n_atoms);
7368 let mut smooth_ranks: Vec<usize> = Vec::with_capacity(n_atoms);
7369 for k in 0..n_atoms {
7370 let m = basis_sizes[k];
7371 let mut a = Array2::<f64>::zeros((m, m));
7372 for r in 0..m {
7373 for c in 0..m {
7374 a[[r, c]] = 0.2 * sample();
7375 }
7376 }
7377 let mut s = a.t().dot(&a);
7378 for r in 0..m {
7379 s[[r, r]] += 1.0;
7380 }
7381 smooth_blocks.push(DeviceSaeSmoothBlock {
7382 global_offset: border_offsets[k],
7383 factor_a: s,
7384 });
7385 smooth_ranks.push(ranks[k]);
7386 }
7387
7388 let n = 6usize;
7390 let q = 3usize;
7391 let mut sys = ArrowSchurSystem::new(n, q, border_dim);
7392 let mut row_htbeta: Vec<Vec<f64>> = Vec::with_capacity(n);
7393 for i in 0..n {
7394 let mut a = Array2::<f64>::zeros((q, q));
7396 for r in 0..q {
7397 for c in 0..q {
7398 a[[r, c]] = sample();
7399 }
7400 }
7401 let mut htt = a.t().dot(&a);
7402 for r in 0..q {
7403 htt[[r, r]] += q as f64 + 1.0;
7404 }
7405 sys.rows[i].htt = htt;
7406 let mut slab = vec![0.0_f64; q * border_dim];
7407 for c in 0..q {
7408 for col in 0..border_dim {
7409 let v = 0.15 * sample();
7410 slab[c * border_dim + col] = v;
7411 sys.rows[i].htbeta[[c, col]] = v;
7412 }
7413 }
7414 row_htbeta.push(slab);
7415 }
7416
7417 let data_op =
7420 FactoredFrameKroneckerOp::new(ranks.clone(), basis_sizes.clone(), frame_blocks.clone())
7421 .expect("frame op");
7422 let mut hbb = data_op.to_dense();
7423 for k in 0..n_atoms {
7424 let op = IdentityRightKroneckerPenaltyOp {
7425 factor_a: smooth_blocks[k].factor_a.clone(),
7426 p: ranks[k],
7427 global_offset: border_offsets[k],
7428 k: border_dim,
7429 };
7430 let d = op.to_dense();
7431 for r in 0..border_dim {
7432 for c in 0..border_dim {
7433 hbb[[r, c]] += d[[r, c]];
7434 }
7435 }
7436 }
7437 sys.hbb = hbb;
7438
7439 let data = DeviceSaePcgData {
7440 p,
7441 beta_dim: border_dim,
7442 a_phi: std::sync::Arc::from(Vec::new().into_boxed_slice()),
7443 local_jac: std::sync::Arc::from(Vec::new().into_boxed_slice()),
7444 smooth_blocks,
7445 sparse_g_blocks: Vec::new(),
7446 frame: Some(DeviceSaeFrameData {
7447 ranks: ranks.clone(),
7448 basis_sizes: basis_sizes.clone(),
7449 border_offsets: border_offsets.clone(),
7450 frame_blocks,
7451 smooth_ranks,
7452 row_htbeta,
7453 }),
7454 };
7455
7456 let ridge_t = 1e-7;
7457 let ridge_beta = 1e-6;
7458
7459 let mut s_dense = Array2::<f64>::zeros((border_dim, border_dim));
7463 for r in 0..border_dim {
7464 for c in 0..border_dim {
7465 s_dense[[r, c]] = sys.hbb[[r, c]];
7466 }
7467 s_dense[[r, r]] += ridge_beta;
7468 }
7469 for row in &sys.rows {
7470 let mut htt = row.htt.clone();
7471 for d in 0..q {
7472 htt[[d, d]] += ridge_t;
7473 }
7474 let factor = cholesky_factor_in_place(htt.view(), CholeskyGuard::NonnegativePivot)
7475 .expect("htt PD");
7476 let mut y = Array2::<f64>::zeros((q, border_dim));
7478 for col in 0..border_dim {
7479 let mut e = Array1::<f64>::zeros(q);
7480 for r in 0..q {
7481 e[r] = row.htbeta[[r, col]];
7482 }
7483 let solved = cholesky_solve_vector(factor.view(), e.view());
7484 for r in 0..q {
7485 y[[r, col]] = solved[r];
7486 }
7487 }
7488 for r in 0..border_dim {
7489 for c in 0..border_dim {
7490 let mut acc = 0.0;
7491 for d in 0..q {
7492 acc += row.htbeta[[d, r]] * y[[d, c]];
7493 }
7494 s_dense[[r, c]] -= acc;
7495 }
7496 }
7497 }
7498
7499 let mut max_rel = 0.0_f64;
7501 for trial in 0..4 {
7502 let x: Vec<f64> = (0..border_dim)
7503 .map(|a| 0.3 * ((a as f64 + trial as f64) * 0.21).cos() - 0.1)
7504 .collect();
7505 let mut got = vec![0.0_f64; border_dim];
7506 sae_framed_schur_matvec_cpu(&sys, &data, ridge_t, ridge_beta, &x, &mut got)
7507 .expect("framed matvec");
7508 let mut want = vec![0.0_f64; border_dim];
7509 for r in 0..border_dim {
7510 let mut acc = 0.0;
7511 for c in 0..border_dim {
7512 acc += s_dense[[r, c]] * x[c];
7513 }
7514 want[r] = acc;
7515 }
7516 let scale = want.iter().fold(0.0_f64, |m, v| m.max(v.abs())).max(1.0);
7517 for a in 0..border_dim {
7518 let rel = (got[a] - want[a]).abs() / scale;
7519 max_rel = max_rel.max(rel);
7520 }
7521 }
7522 assert!(
7523 max_rel <= 1e-10,
7524 "framed SAE Schur matvec vs dense reference diverged: max_rel={max_rel:e}"
7525 );
7526 }
7527
7528 #[test]
7538 fn framed_sae_schur_matvec_matches_dense_reference_large_k_1026() {
7539 use crate::arrow_schur::{
7540 BetaPenaltyOp, DeviceSaeFrameData, DeviceSaePcgData, DeviceSaeSmoothBlock,
7541 FactoredFrameGBlock, FactoredFrameKroneckerOp, IdentityRightKroneckerPenaltyOp,
7542 };
7543
7544 let p = 12usize;
7545 let n_atoms = 40usize;
7546 let ranks: Vec<usize> = (0..n_atoms)
7549 .map(|k| if k % 5 == 0 { p } else { 2 + (k % 3) })
7550 .collect();
7551 let basis_sizes: Vec<usize> = (0..n_atoms).map(|k| 1 + (k % 3)).collect();
7552 let mut border_offsets = Vec::with_capacity(n_atoms);
7553 let mut acc = 0usize;
7554 for k in 0..n_atoms {
7555 border_offsets.push(acc);
7556 acc += basis_sizes[k] * ranks[k];
7557 }
7558 let border_dim = acc;
7559
7560 let mut state = 0x0bad_c0de_dead_beefu64;
7561 let mut sample = || -> f64 {
7562 state = state
7563 .wrapping_mul(6364136223846793005)
7564 .wrapping_add(1442695040888963407);
7565 ((state >> 33) as f64) / ((1u64 << 31) as f64) - 1.0
7566 };
7567
7568 let mut frames: Vec<Array2<f64>> = Vec::with_capacity(n_atoms);
7570 for k in 0..n_atoms {
7571 let r = ranks[k];
7572 let mut u = Array2::<f64>::zeros((p, r));
7573 for i in 0..p {
7574 for j in 0..r {
7575 u[[i, j]] = if r == p {
7576 if i == j { 1.0 } else { 0.0 }
7577 } else {
7578 sample()
7579 };
7580 }
7581 }
7582 frames.push(u);
7583 }
7584 let w_of = |i: usize, j: usize| -> Array2<f64> {
7585 let (ui, uj) = (&frames[i], &frames[j]);
7586 let (ri, rj) = (ranks[i], ranks[j]);
7587 let mut w = Array2::<f64>::zeros((ri, rj));
7588 for a in 0..ri {
7589 for b in 0..rj {
7590 let mut s = 0.0;
7591 for c in 0..p {
7592 s += ui[[c, a]] * uj[[c, b]];
7593 }
7594 w[[a, b]] = s;
7595 }
7596 }
7597 w
7598 };
7599
7600 let mut pairs: Vec<(usize, usize)> = Vec::new();
7603 for k in 0..n_atoms {
7604 pairs.push((k, k));
7605 }
7606 for k in 0..n_atoms - 1 {
7607 pairs.push((k, k + 1));
7608 pairs.push((k + 1, k));
7609 }
7610 pairs.sort_unstable();
7611 let mut frame_blocks: Vec<FactoredFrameGBlock> = Vec::new();
7612 for &(i, j) in &pairs {
7613 let (mi, mj) = (basis_sizes[i], basis_sizes[j]);
7614 let mut g = Array2::<f64>::zeros((mi, mj));
7615 for r in 0..mi {
7616 for c in 0..mj {
7617 g[[r, c]] = 0.3 * sample();
7618 }
7619 }
7620 if i == j {
7621 for r in 0..mi.min(mj) {
7622 g[[r, r]] += mi as f64 + 2.0;
7623 }
7624 }
7625 frame_blocks.push(FactoredFrameGBlock {
7626 atom_i: i,
7627 atom_j: j,
7628 g,
7629 w: w_of(i, j),
7630 });
7631 }
7632
7633 let mut smooth_blocks: Vec<DeviceSaeSmoothBlock> = Vec::with_capacity(n_atoms);
7635 let mut smooth_ranks: Vec<usize> = Vec::with_capacity(n_atoms);
7636 for k in 0..n_atoms {
7637 let m = basis_sizes[k];
7638 let mut a = Array2::<f64>::zeros((m, m));
7639 for r in 0..m {
7640 for c in 0..m {
7641 a[[r, c]] = 0.2 * sample();
7642 }
7643 }
7644 let mut s = a.t().dot(&a);
7645 for r in 0..m {
7646 s[[r, r]] += 1.0;
7647 }
7648 smooth_blocks.push(DeviceSaeSmoothBlock {
7649 global_offset: border_offsets[k],
7650 factor_a: s,
7651 });
7652 smooth_ranks.push(ranks[k]);
7653 }
7654
7655 let n = 8usize;
7657 let q = 3usize;
7658 let mut sys = ArrowSchurSystem::new(n, q, border_dim);
7659 let mut row_htbeta: Vec<Vec<f64>> = Vec::with_capacity(n);
7660 for i in 0..n {
7661 let mut a = Array2::<f64>::zeros((q, q));
7662 for r in 0..q {
7663 for c in 0..q {
7664 a[[r, c]] = sample();
7665 }
7666 }
7667 let mut htt = a.t().dot(&a);
7668 for r in 0..q {
7669 htt[[r, r]] += q as f64 + 1.0;
7670 }
7671 sys.rows[i].htt = htt;
7672 let mut slab = vec![0.0_f64; q * border_dim];
7673 for c in 0..q {
7674 for col in 0..border_dim {
7675 let v = 0.15 * sample();
7676 slab[c * border_dim + col] = v;
7677 sys.rows[i].htbeta[[c, col]] = v;
7678 }
7679 }
7680 row_htbeta.push(slab);
7681 }
7682
7683 let data_op =
7686 FactoredFrameKroneckerOp::new(ranks.clone(), basis_sizes.clone(), frame_blocks.clone())
7687 .expect("frame op");
7688 let mut hbb = data_op.to_dense();
7689 for k in 0..n_atoms {
7690 let op = IdentityRightKroneckerPenaltyOp {
7691 factor_a: smooth_blocks[k].factor_a.clone(),
7692 p: ranks[k],
7693 global_offset: border_offsets[k],
7694 k: border_dim,
7695 };
7696 let d = op.to_dense();
7697 for r in 0..border_dim {
7698 for c in 0..border_dim {
7699 hbb[[r, c]] += d[[r, c]];
7700 }
7701 }
7702 }
7703 sys.hbb = hbb;
7704
7705 let data = DeviceSaePcgData {
7706 p,
7707 beta_dim: border_dim,
7708 a_phi: std::sync::Arc::from(Vec::new().into_boxed_slice()),
7709 local_jac: std::sync::Arc::from(Vec::new().into_boxed_slice()),
7710 smooth_blocks,
7711 sparse_g_blocks: Vec::new(),
7712 frame: Some(DeviceSaeFrameData {
7713 ranks: ranks.clone(),
7714 basis_sizes: basis_sizes.clone(),
7715 border_offsets: border_offsets.clone(),
7716 frame_blocks,
7717 smooth_ranks,
7718 row_htbeta,
7719 }),
7720 };
7721
7722 let ridge_t = 1e-7;
7723 let ridge_beta = 1e-6;
7724
7725 let mut s_dense = sys.hbb.clone();
7727 for r in 0..border_dim {
7728 s_dense[[r, r]] += ridge_beta;
7729 }
7730 for row in &sys.rows {
7731 let mut htt = row.htt.clone();
7732 for d in 0..q {
7733 htt[[d, d]] += ridge_t;
7734 }
7735 let factor = cholesky_factor_in_place(htt.view(), CholeskyGuard::NonnegativePivot)
7736 .expect("htt PD");
7737 let mut y = Array2::<f64>::zeros((q, border_dim));
7738 for col in 0..border_dim {
7739 let mut e = Array1::<f64>::zeros(q);
7740 for r in 0..q {
7741 e[r] = row.htbeta[[r, col]];
7742 }
7743 let solved = cholesky_solve_vector(factor.view(), e.view());
7744 for r in 0..q {
7745 y[[r, col]] = solved[r];
7746 }
7747 }
7748 for r in 0..border_dim {
7749 for c in 0..border_dim {
7750 let mut acc = 0.0;
7751 for d in 0..q {
7752 acc += row.htbeta[[d, r]] * y[[d, c]];
7753 }
7754 s_dense[[r, c]] -= acc;
7755 }
7756 }
7757 }
7758
7759 let mut max_rel = 0.0_f64;
7760 for trial in 0..4 {
7761 let x: Vec<f64> = (0..border_dim)
7762 .map(|a| 0.3 * ((a as f64 + trial as f64) * 0.21).cos() - 0.1)
7763 .collect();
7764 let mut got = vec![0.0_f64; border_dim];
7765 sae_framed_schur_matvec_cpu(&sys, &data, ridge_t, ridge_beta, &x, &mut got)
7766 .expect("framed matvec");
7767 let mut want = vec![0.0_f64; border_dim];
7768 for r in 0..border_dim {
7769 let mut acc = 0.0;
7770 for c in 0..border_dim {
7771 acc += s_dense[[r, c]] * x[c];
7772 }
7773 want[r] = acc;
7774 }
7775 let scale = want.iter().fold(0.0_f64, |m, v| m.max(v.abs())).max(1.0);
7776 for a in 0..border_dim {
7777 let rel = (got[a] - want[a]).abs() / scale;
7778 max_rel = max_rel.max(rel);
7779 }
7780 }
7781 assert!(
7782 max_rel <= 1e-10,
7783 "large-K framed SAE Schur matvec vs dense reference diverged: \
7784 max_rel={max_rel:e} (n_atoms={n_atoms}, border_dim={border_dim})"
7785 );
7786 }
7787
7788 #[test]
7794 fn framed_sae_device_pcg_matches_cpu_when_cuda_admits() {
7795 use crate::arrow_schur::{
7796 BetaPenaltyOp, DeviceSaeFrameData, DeviceSaePcgData, DeviceSaeSmoothBlock,
7797 FactoredFrameGBlock, FactoredFrameKroneckerOp, IdentityRightKroneckerPenaltyOp,
7798 };
7799
7800 let p = 6usize;
7804 let n_atoms = 8usize;
7805 let ranks: Vec<usize> = (0..n_atoms)
7806 .map(|k| if k % 2 == 0 { 3usize } else { p })
7807 .collect();
7808 let basis_sizes: Vec<usize> = (0..n_atoms).map(|_| 3usize).collect();
7809 let mut border_offsets = Vec::with_capacity(n_atoms);
7810 let mut acc = 0usize;
7811 for k in 0..n_atoms {
7812 border_offsets.push(acc);
7813 acc += basis_sizes[k] * ranks[k];
7814 }
7815 let border_dim = acc; let mut state = 0xfeed_face_dead_beefu64;
7818 let mut sample = || -> f64 {
7819 state = state
7820 .wrapping_mul(6364136223846793005)
7821 .wrapping_add(1442695040888963407);
7822 ((state >> 33) as f64) / ((1u64 << 31) as f64) - 1.0
7823 };
7824 let mut frames: Vec<Array2<f64>> = Vec::new();
7825 for k in 0..n_atoms {
7826 let r = ranks[k];
7827 let mut u = Array2::<f64>::zeros((p, r));
7828 for i in 0..p {
7829 for j in 0..r {
7830 u[[i, j]] = if r == p && i == j {
7831 1.0
7832 } else if r == p {
7833 0.0
7834 } else {
7835 sample()
7836 };
7837 }
7838 }
7839 frames.push(u);
7840 }
7841 let w_of = |i: usize, j: usize| {
7842 let (ui, uj) = (&frames[i], &frames[j]);
7843 let (ri, rj) = (ranks[i], ranks[j]);
7844 let mut w = Array2::<f64>::zeros((ri, rj));
7845 for a in 0..ri {
7846 for b in 0..rj {
7847 let mut s = 0.0;
7848 for c in 0..p {
7849 s += ui[[c, a]] * uj[[c, b]];
7850 }
7851 w[[a, b]] = s;
7852 }
7853 }
7854 w
7855 };
7856 let mut pairs: Vec<(usize, usize)> = (0..n_atoms).map(|k| (k, k)).collect();
7857 for &(i, j) in &[(0usize, 1usize), (2, 4), (3, 6)] {
7859 pairs.push((i, j));
7860 pairs.push((j, i));
7861 }
7862 let mut frame_blocks = Vec::new();
7863 for &(i, j) in &pairs {
7864 let (mi, mj) = (basis_sizes[i], basis_sizes[j]);
7865 let mut g = Array2::<f64>::zeros((mi, mj));
7866 for r in 0..mi {
7867 for c in 0..mj {
7868 g[[r, c]] = 0.25 * sample();
7869 }
7870 }
7871 if i == j {
7872 for r in 0..mi.min(mj) {
7873 g[[r, r]] += mi as f64 + 2.0;
7874 }
7875 }
7876 frame_blocks.push(FactoredFrameGBlock {
7877 atom_i: i,
7878 atom_j: j,
7879 g,
7880 w: w_of(i, j),
7881 });
7882 }
7883 let mut smooth_blocks = Vec::new();
7884 let mut smooth_ranks = Vec::new();
7885 for k in 0..n_atoms {
7886 let m = basis_sizes[k];
7887 let mut a = Array2::<f64>::zeros((m, m));
7888 for r in 0..m {
7889 for c in 0..m {
7890 a[[r, c]] = 0.2 * sample();
7891 }
7892 }
7893 let mut s = a.t().dot(&a);
7894 for r in 0..m {
7895 s[[r, r]] += 1.0;
7896 }
7897 smooth_blocks.push(DeviceSaeSmoothBlock {
7898 global_offset: border_offsets[k],
7899 factor_a: s,
7900 });
7901 smooth_ranks.push(ranks[k]);
7902 }
7903 let n = 400usize;
7904 let q = 4usize;
7905 let mut sys = ArrowSchurSystem::new(n, q, border_dim);
7906 let mut row_htbeta = Vec::new();
7907 for i in 0..n {
7908 let mut a = Array2::<f64>::zeros((q, q));
7909 for r in 0..q {
7910 for c in 0..q {
7911 a[[r, c]] = sample();
7912 }
7913 }
7914 let mut htt = a.t().dot(&a);
7915 for r in 0..q {
7916 htt[[r, r]] += q as f64 + 1.0;
7917 }
7918 sys.rows[i].htt = htt;
7919 let mut slab = vec![0.0_f64; q * border_dim];
7920 for c in 0..q {
7921 for col in 0..border_dim {
7922 let v = 0.02 * sample();
7925 slab[c * border_dim + col] = v;
7926 sys.rows[i].htbeta[[c, col]] = v;
7927 }
7928 }
7929 row_htbeta.push(slab);
7930 }
7931 let data_op =
7932 FactoredFrameKroneckerOp::new(ranks.clone(), basis_sizes.clone(), frame_blocks.clone())
7933 .expect("frame op");
7934 let mut hbb = data_op.to_dense();
7935 for k in 0..n_atoms {
7936 let op = IdentityRightKroneckerPenaltyOp {
7937 factor_a: smooth_blocks[k].factor_a.clone(),
7938 p: ranks[k],
7939 global_offset: border_offsets[k],
7940 k: border_dim,
7941 };
7942 let d = op.to_dense();
7943 for r in 0..border_dim {
7944 for c in 0..border_dim {
7945 hbb[[r, c]] += d[[r, c]];
7946 }
7947 }
7948 }
7949 sys.hbb = hbb;
7950 let data = DeviceSaePcgData {
7951 p,
7952 beta_dim: border_dim,
7953 a_phi: std::sync::Arc::from(Vec::new().into_boxed_slice()),
7954 local_jac: std::sync::Arc::from(Vec::new().into_boxed_slice()),
7955 smooth_blocks,
7956 sparse_g_blocks: Vec::new(),
7957 frame: Some(DeviceSaeFrameData {
7958 ranks: ranks.clone(),
7959 basis_sizes: basis_sizes.clone(),
7960 border_offsets: border_offsets.clone(),
7961 frame_blocks,
7962 smooth_ranks,
7963 row_htbeta,
7964 }),
7965 };
7966 let ridge_t = 1e-7;
7967 let ridge_beta = 1e-6;
7968 let rhs: Array1<f64> =
7969 Array1::from_shape_fn(border_dim, |a| ((a as f64 + 1.0) * 0.17).sin());
7970
7971 let (device, diag) =
7972 match solve_sae_matrix_free_pcg(&sys, &data, ridge_t, ridge_beta, &rhs, 400, 1e-12) {
7973 Ok(result) => result,
7974 Err(failure) => {
7980 assert!(
7981 gam_gpu::device_runtime::GpuRuntime::resolve(gam_gpu::GpuPolicy::Auto)
7982 .unwrap_or_else(|error| {
7983 panic!("GPU probe fault in framed PCG test: {error}")
7984 })
7985 .is_none(),
7986 "#1017: CUDA device present but the framed device SAE PCG \
7987 declined/faulted instead of returning a result (tag: {failure:?}) — \
7988 the kernel does not run correctly on GPU"
7989 );
7990 return;
7991 }
7992 };
7993
7994 let rhs_norm = rhs.iter().map(|v| v * v).sum::<f64>().sqrt();
8012 let oracle_resid = |x: &Array1<f64>| -> f64 {
8013 let mut sx = vec![0.0_f64; border_dim];
8014 sae_framed_schur_matvec_cpu(
8015 &sys,
8016 &data,
8017 ridge_t,
8018 ridge_beta,
8019 x.as_slice().unwrap(),
8020 &mut sx,
8021 )
8022 .expect("cpu oracle matvec");
8023 let mut acc = 0.0_f64;
8024 for a in 0..border_dim {
8025 let e = sx[a] - rhs[a];
8026 acc += e * e;
8027 }
8028 acc.sqrt()
8029 };
8030 let s_dev_resid = oracle_resid(&device);
8031 let dev_rel_resid = s_dev_resid / rhs_norm.max(1e-300);
8032
8033 let precond = {
8038 let d = sae_frame_penalty_diag_host_for_test(&data, ridge_beta);
8039 let mut diag = d;
8042 for (i, row) in sys.rows.iter().enumerate() {
8043 let slab = &data.frame.as_ref().unwrap().row_htbeta[i];
8044 let qi = sys.row_dims[i];
8045 if slab.is_empty() || qi == 0 || slab.len() != qi * border_dim {
8046 continue;
8047 }
8048 let mut block = row.htt.clone();
8049 for dd in 0..qi {
8050 block[[dd, dd]] += ridge_t;
8051 }
8052 let factor =
8053 cholesky_factor_in_place(block.view(), CholeskyGuard::NonnegativePivot)
8054 .expect("row htt PD");
8055 let mut ainv = Array2::<f64>::zeros((qi, qi));
8057 for col in 0..qi {
8058 let mut e = Array1::<f64>::zeros(qi);
8059 e[col] = 1.0;
8060 let s = cholesky_solve_vector(factor.view(), e.view());
8061 for r in 0..qi {
8062 ainv[[r, col]] = s[r];
8063 }
8064 }
8065 for a in 0..border_dim {
8066 let mut quad = 0.0_f64;
8067 for c in 0..qi {
8068 let hc = slab[c * border_dim + a];
8069 for dd in 0..qi {
8070 quad += hc * ainv[[c, dd]] * slab[dd * border_dim + a];
8071 }
8072 }
8073 diag[a] -= quad;
8074 }
8075 }
8076 Array1::from_vec(diag)
8077 };
8078 let mut cpu = Array1::<f64>::zeros(border_dim);
8079 let cpu_result = {
8080 let mut apply = |v: &Array1<f64>, out: &mut Array1<f64>| {
8081 let mut tmp = vec![0.0_f64; border_dim];
8082 sae_framed_schur_matvec_cpu(
8083 &sys,
8084 &data,
8085 ridge_t,
8086 ridge_beta,
8087 v.as_slice().unwrap(),
8088 &mut tmp,
8089 )
8090 .expect("cpu oracle matvec");
8091 out.assign(&Array1::from_vec(tmp));
8092 };
8093 gam_linalg::pcg::pcg_core(
8094 &mut apply,
8095 &rhs.view(),
8096 &precond.view(),
8097 1e-12,
8098 800,
8099 32,
8100 false,
8101 gam_linalg::pcg::DotReduction::Serial,
8102 &mut cpu.view_mut(),
8103 )
8104 };
8105 let s_cpu_resid = oracle_resid(&cpu);
8106 let cpu_rel_resid = s_cpu_resid / rhs_norm.max(1e-300);
8107
8108 assert!(
8111 dev_rel_resid <= 1e-7,
8112 "[#1551] device δβ does not solve the CPU-oracle system: \
8113 ‖S_cpu·device−rhs‖/‖rhs‖={dev_rel_resid:e} (>1e-7) | abs={s_dev_resid:e} | \
8114 device PCG stop={:?} iters={} final_rel_resid={:e} — a large operator residual \
8115 means the device matvec is a DIFFERENT operator (kernel bug)",
8116 diag.stopping_reason,
8117 diag.iterations,
8118 diag.final_relative_residual,
8119 );
8120 assert!(
8123 cpu_rel_resid <= 1e-6,
8124 "[#1551] CPU pcg_core failed to solve the oracle system: \
8125 ‖S_cpu·cpu−rhs‖/‖rhs‖={cpu_rel_resid:e} (stop={:?}, iters={}) — fixture/oracle issue",
8126 cpu_result.stop,
8127 cpu_result.iterations,
8128 );
8129 }
8130
8131 fn sae_frame_penalty_diag_host_for_test(data: &DeviceSaePcgData, ridge_beta: f64) -> Vec<f64> {
8135 let frame = data.frame.as_ref().expect("frame");
8136 let mut diag = vec![ridge_beta; data.beta_dim];
8137 for (blk, &r) in data.smooth_blocks.iter().zip(frame.smooth_ranks.iter()) {
8138 let m = blk.factor_a.nrows();
8139 for ia in 0..m {
8140 let coeff = blk.factor_a[[ia, ia]];
8141 let base = blk.global_offset + ia * r;
8142 for ib in 0..r {
8143 diag[base + ib] += coeff;
8144 }
8145 }
8146 }
8147 for blk in &frame.frame_blocks {
8148 if blk.atom_i != blk.atom_j {
8149 continue;
8150 }
8151 let r = frame.ranks[blk.atom_i];
8152 let off = frame.border_offsets[blk.atom_i];
8153 let (mi, mj) = blk.g.dim();
8154 for li in 0..mi.min(mj) {
8155 let gii = blk.g[[li, li]];
8156 let base = off + li * r;
8157 for a in 0..r {
8158 diag[base + a] += gii * blk.w[[a, a]];
8159 }
8160 }
8161 }
8162 diag
8163 }
8164
8165 #[test]
8176 #[cfg(target_os = "linux")]
8177 fn framed_sae_device_matvec_matches_cpu_oracle_when_cuda_admits() {
8178 use crate::arrow_schur::{
8179 DeviceSaeFrameData, DeviceSaePcgData, DeviceSaeSmoothBlock, FactoredFrameGBlock,
8180 };
8181
8182 let p = 6usize;
8185 let n_atoms = 8usize;
8186 let ranks: Vec<usize> = (0..n_atoms)
8187 .map(|k| if k % 2 == 0 { 3usize } else { p })
8188 .collect();
8189 let basis_sizes: Vec<usize> = (0..n_atoms).map(|_| 3usize).collect();
8190 let mut border_offsets = Vec::with_capacity(n_atoms);
8191 let mut acc = 0usize;
8192 for k in 0..n_atoms {
8193 border_offsets.push(acc);
8194 acc += basis_sizes[k] * ranks[k];
8195 }
8196 let border_dim = acc;
8197
8198 let mut state = 0x1551_0017_1026_0922u64;
8199 let mut sample = || -> f64 {
8200 state = state
8201 .wrapping_mul(6364136223846793005)
8202 .wrapping_add(1442695040888963407);
8203 ((state >> 33) as f64) / ((1u64 << 31) as f64) - 1.0
8204 };
8205 let mut frames: Vec<Array2<f64>> = Vec::new();
8206 for k in 0..n_atoms {
8207 let r = ranks[k];
8208 let mut u = Array2::<f64>::zeros((p, r));
8209 for i in 0..p {
8210 for j in 0..r {
8211 u[[i, j]] = if r == p && i == j {
8212 1.0
8213 } else if r == p {
8214 0.0
8215 } else {
8216 sample()
8217 };
8218 }
8219 }
8220 frames.push(u);
8221 }
8222 let w_of = |i: usize, j: usize| {
8223 let (ui, uj) = (&frames[i], &frames[j]);
8224 let (ri, rj) = (ranks[i], ranks[j]);
8225 let mut w = Array2::<f64>::zeros((ri, rj));
8226 for a in 0..ri {
8227 for b in 0..rj {
8228 let mut s = 0.0;
8229 for c in 0..p {
8230 s += ui[[c, a]] * uj[[c, b]];
8231 }
8232 w[[a, b]] = s;
8233 }
8234 }
8235 w
8236 };
8237 let mut pairs: Vec<(usize, usize)> = (0..n_atoms).map(|k| (k, k)).collect();
8238 for &(i, j) in &[(0usize, 1usize), (2, 4), (3, 6)] {
8239 pairs.push((i, j));
8240 pairs.push((j, i));
8241 }
8242 let mut frame_blocks = Vec::new();
8243 for &(i, j) in &pairs {
8244 let (mi, mj) = (basis_sizes[i], basis_sizes[j]);
8245 let mut g = Array2::<f64>::zeros((mi, mj));
8246 for r in 0..mi {
8247 for c in 0..mj {
8248 g[[r, c]] = 0.25 * sample();
8249 }
8250 }
8251 if i == j {
8252 for r in 0..mi.min(mj) {
8253 g[[r, r]] += mi as f64 + 2.0;
8254 }
8255 }
8256 frame_blocks.push(FactoredFrameGBlock {
8257 atom_i: i,
8258 atom_j: j,
8259 g,
8260 w: w_of(i, j),
8261 });
8262 }
8263 let mut smooth_blocks = Vec::new();
8264 let mut smooth_ranks = Vec::new();
8265 for k in 0..n_atoms {
8266 let m = basis_sizes[k];
8267 let mut a = Array2::<f64>::zeros((m, m));
8268 for r in 0..m {
8269 for c in 0..m {
8270 a[[r, c]] = 0.2 * sample();
8271 }
8272 }
8273 let mut s = a.t().dot(&a);
8274 for r in 0..m {
8275 s[[r, r]] += 1.0;
8276 }
8277 smooth_blocks.push(DeviceSaeSmoothBlock {
8278 global_offset: border_offsets[k],
8279 factor_a: s,
8280 });
8281 smooth_ranks.push(ranks[k]);
8282 }
8283 let n = 32usize;
8287 let q = 4usize;
8288 let mut sys = ArrowSchurSystem::new(n, q, border_dim);
8289 let mut row_htbeta = Vec::new();
8290 for i in 0..n {
8291 let mut a = Array2::<f64>::zeros((q, q));
8292 for r in 0..q {
8293 for c in 0..q {
8294 a[[r, c]] = sample();
8295 }
8296 }
8297 let mut htt = a.t().dot(&a);
8298 for r in 0..q {
8299 htt[[r, r]] += q as f64 + 1.0;
8300 }
8301 sys.rows[i].htt = htt;
8302 let mut slab = vec![0.0_f64; q * border_dim];
8303 for c in 0..q {
8304 for col in 0..border_dim {
8305 let v = 0.3 * sample();
8306 slab[c * border_dim + col] = v;
8307 sys.rows[i].htbeta[[c, col]] = v;
8308 }
8309 }
8310 row_htbeta.push(slab);
8311 }
8312 let ridge_t = 1e-7;
8313 let ridge_beta = 1e-6;
8314 let data = DeviceSaePcgData {
8315 p,
8316 beta_dim: border_dim,
8317 a_phi: std::sync::Arc::from(Vec::new().into_boxed_slice()),
8318 local_jac: std::sync::Arc::from(Vec::new().into_boxed_slice()),
8319 smooth_blocks,
8320 sparse_g_blocks: Vec::new(),
8321 frame: Some(DeviceSaeFrameData {
8322 ranks: ranks.clone(),
8323 basis_sizes: basis_sizes.clone(),
8324 border_offsets: border_offsets.clone(),
8325 frame_blocks,
8326 smooth_ranks,
8327 row_htbeta,
8328 }),
8329 };
8330
8331 let mut probes: Vec<Array1<f64>> = Vec::new();
8335 probes.push(Array1::from_shape_fn(border_dim, |a| {
8336 ((a as f64 + 1.0) * 0.37).sin()
8337 }));
8338 probes.push(Array1::from_shape_fn(border_dim, |_| sample()));
8339 for axis in [0usize, border_dim / 3, border_dim - 1] {
8340 let mut e = Array1::<f64>::zeros(border_dim);
8341 e[axis] = 1.0;
8342 probes.push(e);
8343 }
8344
8345 let mut any_ran = false;
8346 let mut worst = 0.0_f64;
8347 for (pi, x) in probes.iter().enumerate() {
8348 let device = match super::framed_schur_matvec_once_on_device(
8349 &sys, &data, ridge_t, ridge_beta, x,
8350 ) {
8351 Ok(out) => out,
8352 Err(failure) => {
8353 assert!(
8357 gam_gpu::device_runtime::GpuRuntime::resolve(gam_gpu::GpuPolicy::Auto)
8358 .unwrap_or_else(|error| {
8359 panic!("GPU probe fault in framed matvec parity test: {error}")
8360 })
8361 .is_none(),
8362 "#1551: CUDA device present but the framed device matvec \
8363 declined/faulted (probe {pi}, tag: {failure:?}) — the kernel \
8364 does not run on GPU"
8365 );
8366 return;
8367 }
8368 };
8369 any_ran = true;
8370 let mut cpu = vec![0.0_f64; border_dim];
8371 sae_framed_schur_matvec_cpu(
8372 &sys,
8373 &data,
8374 ridge_t,
8375 ridge_beta,
8376 x.as_slice().unwrap(),
8377 &mut cpu,
8378 )
8379 .expect("cpu oracle matvec");
8380 let scale = cpu.iter().fold(0.0_f64, |m, v| m.max(v.abs())).max(1.0);
8381 for a in 0..border_dim {
8382 let rel = (device[a] - cpu[a]).abs() / scale;
8383 worst = worst.max(rel);
8384 assert!(
8385 rel <= 1e-9,
8386 "[#1551 matvec-parity] probe {pi} component {a}: device={:e} cpu={:e} \
8387 rel={rel:e} (>1e-9) — framed S·x kernel diverges from the CPU oracle",
8388 device[a],
8389 cpu[a],
8390 );
8391 }
8392 }
8393 if any_ran {
8394 assert!(
8399 gam_gpu::device_runtime::GpuRuntime::resolve(gam_gpu::GpuPolicy::Auto)
8400 .unwrap_or_else(|error| {
8401 panic!("GPU probe fault after framed matvec execution: {error}")
8402 })
8403 .is_some(),
8404 "#1551: matvec ran but no GPU runtime — unexpected"
8405 );
8406 assert!(worst <= 1e-9, "framed matvec parity worst rel = {worst:e}");
8407 }
8408 }
8409}