1use ndarray::{Array1, Array2, ArrayView2};
21
22use crate::arrow_schur::{ArrowPcgDiagnostics, ArrowSchurSystem, DeviceSaePcgData};
23#[cfg(target_os = "linux")]
27use crate::arrow_schur::{
28 ArrowBetaGaugeQuotient, ArrowSolveOptions, solve_dense_reduced_system,
29};
30use gam_linalg::triangular::{CholeskyGuard, cholesky_factor_in_place, cholesky_solve_vector};
31
32pub struct ArrowSchurGpuSolution {
34 pub delta_t: Array1<f64>,
35 pub delta_beta: Array1<f64>,
36 pub log_det_hessian: f64,
39}
40
41#[derive(Debug, Clone)]
45pub enum ArrowSchurGpuFailure {
46 Unavailable,
48 RidgeBumpRequired { row: usize, bump: f64 },
51 SchurFactorFailed { reason: String },
54 GpuRequiresDenseSystem {
65 had_hbb_matvec: bool,
66 had_htbeta_matvec: bool,
67 },
68}
69
70#[cfg(target_os = "linux")]
81fn resolve_runtime_for_device_path(
82) -> Result<Option<&'static gam_gpu::GpuRuntime>, ArrowSchurGpuFailure> {
83 gam_gpu::device_runtime::GpuRuntime::resolve(gam_gpu::global_policy()).map_err(|error| {
84 ArrowSchurGpuFailure::SchurFactorFailed {
85 reason: format!("GPU runtime resolution failed: {error}"),
86 }
87 })
88}
89
90const RIDGE_BUMP_EPS_MARGIN: f64 = 1024.0;
101
102#[must_use]
138fn ridge_bump_to_make_pd(htt: ArrayView2<'_, f64>, ridge_t: f64) -> f64 {
139 let d = htt.nrows();
140 let mut scale = 1.0_f64;
143 let mut min_gershgorin_edge = f64::INFINITY;
144 for i in 0..d {
145 let diag = htt[[i, i]];
146 scale = scale.max(diag.abs());
147 let mut off_sum = 0.0_f64;
148 for j in 0..d {
149 if j != i {
150 off_sum += htt[[i, j]].abs();
151 }
152 }
153 min_gershgorin_edge = min_gershgorin_edge.min(diag - off_sum);
154 }
155 if !min_gershgorin_edge.is_finite() {
156 return scale * f64::EPSILON.sqrt() * RIDGE_BUMP_EPS_MARGIN;
159 }
160 let deficit = -(min_gershgorin_edge + ridge_t);
163 let margin = scale * f64::EPSILON.sqrt() * RIDGE_BUMP_EPS_MARGIN;
164 deficit.max(0.0) + margin
167}
168
169#[cfg(target_os = "linux")]
182#[must_use]
183fn ridge_bump_to_make_pd_colmajor(block: &[f64], d: usize) -> f64 {
184 if d == 0 || block.len() < d * d {
185 return f64::EPSILON.sqrt() * RIDGE_BUMP_EPS_MARGIN;
186 }
187 let mut scale = 1.0_f64;
190 let mut min_gershgorin_edge = f64::INFINITY;
191 for i in 0..d {
192 let diag = block[i * d + i];
193 scale = scale.max(diag.abs());
194 let mut off_sum = 0.0_f64;
195 for j in 0..d {
196 if j != i {
197 off_sum += block[j * d + i].abs();
198 }
199 }
200 min_gershgorin_edge = min_gershgorin_edge.min(diag - off_sum);
201 }
202 let margin = scale * f64::EPSILON.sqrt() * RIDGE_BUMP_EPS_MARGIN;
203 if !min_gershgorin_edge.is_finite() {
204 return margin;
205 }
206 (-min_gershgorin_edge).max(0.0) + margin
207}
208
209#[cfg(target_os = "linux")]
242pub(crate) fn canonicalize_device_beta_factor(
243 quotient: Option<&ArrowBetaGaugeQuotient>,
244 newton_schur_tikhonov_rel_floor: Option<f64>,
245 k: usize,
246 schur_col_major: &mut [f64],
247) -> Result<bool, ArrowSchurGpuFailure> {
248 if quotient.is_none() && newton_schur_tikhonov_rel_floor.is_none() {
249 return Ok(false);
250 }
251 if schur_col_major.len() != k * k {
252 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
253 reason: format!(
254 "device canonical reduced factor received Schur length {} for k={k}",
255 schur_col_major.len(),
256 ),
257 });
258 }
259
260 let mut schur = Array2::<f64>::zeros((k, k));
261 for col in 0..k {
262 for row in 0..k {
263 schur[[row, col]] = schur_col_major[col * k + row];
264 }
265 }
266 let conditioned_input = match quotient {
267 Some(quotient) => quotient.pin_reduced_schur(schur.view()),
268 None => schur,
269 };
270 let mut canonical_options = ArrowSolveOptions::direct();
271 canonical_options.newton_schur_tikhonov_rel_floor =
272 newton_schur_tikhonov_rel_floor;
273 let zero_rhs = Array1::<f64>::zeros(k);
274 let (_, factor, _) =
275 solve_dense_reduced_system(&conditioned_input, &zero_rhs, &canonical_options, None)
276 .map_err(|error| ArrowSchurGpuFailure::SchurFactorFailed {
277 reason: format!("canonical reduced-border factorization failed: {error}"),
278 })?;
279 let factor = factor.ok_or_else(|| ArrowSchurGpuFailure::SchurFactorFailed {
280 reason: "canonical Direct reduced-border factorization returned no factor".to_string(),
281 })?;
282 for col in 0..k {
283 for row in 0..k {
284 schur_col_major[col * k + row] = factor[[row, col]];
285 }
286 }
287 Ok(true)
288}
289
290#[cfg(target_os = "linux")]
297fn project_device_beta_vector(
298 quotient: Option<&ArrowBetaGaugeQuotient>,
299 k: usize,
300 vector: Vec<f64>,
301) -> Result<Vec<f64>, ArrowSchurGpuFailure> {
302 if vector.len() != k {
303 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
304 reason: format!(
305 "device beta vector length {} does not match border {}",
306 vector.len(),
307 k,
308 ),
309 });
310 }
311 Ok(match quotient {
312 Some(quotient) => quotient
313 .project_complement(Array1::from_vec(vector).view())
314 .to_vec(),
315 None => vector,
316 })
317}
318
319pub fn solve_arrow_newton_step(
323 sys: &ArrowSchurSystem,
324 ridge_t: f64,
325 ridge_beta: f64,
326 newton_schur_tikhonov_rel_floor: Option<f64>,
327) -> Result<ArrowSchurGpuSolution, ArrowSchurGpuFailure> {
328 let n = sys.rows.len();
329 let d = sys.d;
330 let k = sys.k;
331
332 let had_hbb_matvec = sys.hbb_matvec.is_some();
337 let had_htbeta_matvec = sys.htbeta_matvec.is_some();
338 if had_hbb_matvec || had_htbeta_matvec {
339 return Err(ArrowSchurGpuFailure::GpuRequiresDenseSystem {
340 had_hbb_matvec,
341 had_htbeta_matvec,
342 });
343 }
344
345 if sys.penalty_op.is_some() {
362 return Err(ArrowSchurGpuFailure::GpuRequiresDenseSystem {
363 had_hbb_matvec: false,
364 had_htbeta_matvec: false,
365 });
366 }
367
368 if sys.hbb.dim() != (k, k) {
369 return Err(ArrowSchurGpuFailure::GpuRequiresDenseSystem {
383 had_hbb_matvec: false,
384 had_htbeta_matvec: false,
385 });
386 }
387 if n == 0 || d == 0 {
388 return Err(ArrowSchurGpuFailure::Unavailable);
389 }
390 if sys
391 .rows
392 .iter()
393 .any(|row| row.htt.dim() != (d, d) || row.htbeta.dim() != (d, k) || row.gt.len() != d)
394 {
395 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
396 reason: "row block dimension mismatch".to_string(),
397 });
398 }
399
400 #[cfg(not(target_os = "linux"))]
401 {
402 if let Some(floor) = newton_schur_tikhonov_rel_floor {
410 if !floor.is_finite() {
411 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
412 reason: format!(
413 "newton_schur_tikhonov_rel_floor must be finite, got {floor}"
414 ),
415 });
416 }
417 }
418 if ridge_t.is_nan() || ridge_beta.is_nan() {
419 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
420 reason: "ridge is NaN".to_string(),
421 });
422 }
423 Err(ArrowSchurGpuFailure::Unavailable)
424 }
425
426 #[cfg(target_os = "linux")]
427 {
428 if resolve_runtime_for_device_path()?
437 .map(gam_gpu::device_runtime::GpuRuntime::device_count)
438 .unwrap_or(0)
439 > 1
440 {
441 match cuda::solve_multi_gpu(
442 sys,
443 ridge_t,
444 ridge_beta,
445 newton_schur_tikhonov_rel_floor,
446 ) {
447 Ok(sol) => return Ok(sol),
448 Err(ArrowSchurGpuFailure::RidgeBumpRequired { row, bump }) => {
449 return Err(ArrowSchurGpuFailure::RidgeBumpRequired { row, bump });
450 }
451 Err(ArrowSchurGpuFailure::SchurFactorFailed { reason }) => {
452 return Err(ArrowSchurGpuFailure::SchurFactorFailed { reason });
453 }
454 Err(
459 ArrowSchurGpuFailure::Unavailable
460 | ArrowSchurGpuFailure::GpuRequiresDenseSystem { .. },
461 ) => {}
462 }
463 }
464 if crate::gpu_kernels::arrow_schur_nvrtc::system_admits_fused_path(sys) {
470 match cuda::solve_fused(
471 sys,
472 ridge_t,
473 ridge_beta,
474 newton_schur_tikhonov_rel_floor,
475 ) {
476 Ok(sol) => return Ok(sol),
477 Err(ArrowSchurGpuFailure::RidgeBumpRequired { row, bump }) => {
481 return Err(ArrowSchurGpuFailure::RidgeBumpRequired { row, bump });
482 }
483 Err(
489 ArrowSchurGpuFailure::Unavailable
490 | ArrowSchurGpuFailure::SchurFactorFailed { .. }
491 | ArrowSchurGpuFailure::GpuRequiresDenseSystem { .. },
492 ) => {}
493 }
494 }
495 cuda::solve(
496 sys,
497 ridge_t,
498 ridge_beta,
499 newton_schur_tikhonov_rel_floor,
500 )
501 }
502}
503
504#[cfg(target_os = "linux")]
510fn pack_host(sys: &ArrowSchurSystem, ridge_t: f64) -> (Vec<f64>, Vec<f64>, Vec<f64>) {
511 let n = sys.rows.len();
512 let d = sys.d;
513 let k = sys.k;
514 let mut d_buf = Vec::with_capacity(n * d * d);
515 let mut b_buf = Vec::with_capacity(n * d * k);
516 let mut g_buf = Vec::with_capacity(n * d);
517 for row in &sys.rows {
518 pack_block(row, ridge_t, d, k, &mut d_buf, &mut b_buf, &mut g_buf);
519 }
520 (d_buf, b_buf, g_buf)
521}
522
523#[cfg(target_os = "linux")]
533fn pack_host_d_and_stacked_b(sys: &ArrowSchurSystem, ridge_t: f64) -> (Vec<f64>, Vec<f64>) {
534 let n = sys.rows.len();
535 let d = sys.d;
536 let k = sys.k;
537 let rows = n * d;
538 let mut d_buf = Vec::with_capacity(n * d * d);
539 let mut b_buf = vec![0.0_f64; rows * k];
540 for (i, row) in sys.rows.iter().enumerate() {
541 for col in 0..d {
542 for r in 0..d {
543 let mut value = row.htt[[r, col]];
544 if r == col {
545 value += ridge_t;
546 }
547 d_buf.push(value);
548 }
549 }
550 for col in 0..k {
551 let base = col * rows + i * d;
552 for r in 0..d {
553 b_buf[base + r] = row.htbeta[[r, col]];
554 }
555 }
556 }
557 (d_buf, b_buf)
558}
559
560#[cfg(target_os = "linux")]
561#[inline]
562fn pack_block(
563 row: &crate::arrow_schur::ArrowRowBlock,
564 ridge_t: f64,
565 d: usize,
566 k: usize,
567 d_buf: &mut Vec<f64>,
568 b_buf: &mut Vec<f64>,
569 g_buf: &mut Vec<f64>,
570) {
571 for col in 0..d {
572 for r in 0..d {
573 let mut value = row.htt[[r, col]];
574 if r == col {
575 value += ridge_t;
576 }
577 d_buf.push(value);
578 }
579 }
580 for col in 0..k {
581 for r in 0..d {
582 b_buf.push(row.htbeta[[r, col]]);
583 }
584 }
585 for r in 0..d {
586 g_buf.push(row.gt[r]);
587 }
588}
589
590#[cfg(target_os = "linux")]
596pub fn solve_arrow_newton_step_fused_force(
597 sys: &ArrowSchurSystem,
598 ridge_t: f64,
599 ridge_beta: f64,
600 newton_schur_tikhonov_rel_floor: Option<f64>,
601) -> Result<ArrowSchurGpuSolution, ArrowSchurGpuFailure> {
602 if ridge_t.is_nan() || ridge_beta.is_nan() {
603 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
604 reason: "ridge is NaN".to_string(),
605 });
606 }
607 if crate::gpu_kernels::arrow_schur_nvrtc::plan_fused_launch(sys.rows.len(), sys.d, sys.k)
608 .is_none()
609 {
610 return Err(ArrowSchurGpuFailure::Unavailable);
611 }
612 cuda::solve_fused(
613 sys,
614 ridge_t,
615 ridge_beta,
616 newton_schur_tikhonov_rel_floor,
617 )
618}
619
620#[cfg(target_os = "linux")]
630pub struct ResidentArrowFrameHandle {
631 inner: cuda::ResidentArrowFrame,
632}
633
634#[cfg(not(target_os = "linux"))]
638pub enum ResidentArrowFrameHandle {}
639
640impl ResidentArrowFrameHandle {
641 pub fn new(
643 sys: &ArrowSchurSystem,
644 ridge_t: f64,
645 ridge_beta: f64,
646 newton_schur_tikhonov_rel_floor: Option<f64>,
647 ) -> Result<Self, ArrowSchurGpuFailure> {
648 if sys.hbb_matvec.is_some() || sys.htbeta_matvec.is_some() {
651 return Err(ArrowSchurGpuFailure::GpuRequiresDenseSystem {
652 had_hbb_matvec: sys.hbb_matvec.is_some(),
653 had_htbeta_matvec: sys.htbeta_matvec.is_some(),
654 });
655 }
656 #[cfg(not(target_os = "linux"))]
657 {
658 if let Some(floor) = newton_schur_tikhonov_rel_floor {
661 if !floor.is_finite() {
662 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
663 reason: format!(
664 "newton_schur_tikhonov_rel_floor must be finite, got {floor}"
665 ),
666 });
667 }
668 }
669 if ridge_t.is_nan() || ridge_beta.is_nan() {
670 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
671 reason: "ridge is NaN".to_string(),
672 });
673 }
674 Err(ArrowSchurGpuFailure::Unavailable)
675 }
676 #[cfg(target_os = "linux")]
677 {
678 Ok(Self {
679 inner: cuda::ResidentArrowFrame::new(
680 sys,
681 ridge_t,
682 ridge_beta,
683 newton_schur_tikhonov_rel_floor,
684 )?,
685 })
686 }
687 }
688
689 pub fn solve_gradient(
691 &self,
692 g_t: &[f64],
693 g_beta: &[f64],
694 ) -> Result<ArrowSchurGpuSolution, ArrowSchurGpuFailure> {
695 #[cfg(not(target_os = "linux"))]
696 {
697 if g_t.iter().chain(g_beta).any(|v| !v.is_finite()) {
698 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
699 reason: "non-finite gradient entry".to_string(),
700 });
701 }
702 Err(ArrowSchurGpuFailure::Unavailable)
703 }
704 #[cfg(target_os = "linux")]
705 {
706 self.inner.solve_gradient(g_t, g_beta)
707 }
708 }
709
710 #[must_use]
712 pub fn log_det_hessian(&self) -> f64 {
713 #[cfg(not(target_os = "linux"))]
714 {
715 match *self {}
716 }
717 #[cfg(target_os = "linux")]
718 {
719 self.inner.log_det_hessian()
720 }
721 }
722}
723
724#[cfg(target_os = "linux")]
739pub struct ResidentBaseArrowFrameHandle {
740 inner: cuda::ResidentBaseArrowFrame,
741}
742
743#[cfg(not(target_os = "linux"))]
746pub enum ResidentBaseArrowFrameHandle {}
747
748impl ResidentBaseArrowFrameHandle {
749 pub fn new(
755 sys: &ArrowSchurSystem,
756 newton_schur_tikhonov_rel_floor: Option<f64>,
757 ) -> Result<Self, ArrowSchurGpuFailure> {
758 if sys.hbb_matvec.is_some() || sys.htbeta_matvec.is_some() {
759 return Err(ArrowSchurGpuFailure::GpuRequiresDenseSystem {
760 had_hbb_matvec: sys.hbb_matvec.is_some(),
761 had_htbeta_matvec: sys.htbeta_matvec.is_some(),
762 });
763 }
764 #[cfg(not(target_os = "linux"))]
765 {
766 if let Some(floor) = newton_schur_tikhonov_rel_floor {
769 if !floor.is_finite() {
770 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
771 reason: format!(
772 "newton_schur_tikhonov_rel_floor must be finite, got {floor}"
773 ),
774 });
775 }
776 }
777 Err(ArrowSchurGpuFailure::Unavailable)
778 }
779 #[cfg(target_os = "linux")]
780 {
781 Ok(Self {
782 inner: cuda::ResidentBaseArrowFrame::new(
783 sys,
784 newton_schur_tikhonov_rel_floor,
785 )?,
786 })
787 }
788 }
789
790 pub fn refactor_and_solve(
797 &self,
798 ridge_t: f64,
799 ridge_beta: f64,
800 ) -> Result<ArrowSchurGpuSolution, ArrowSchurGpuFailure> {
801 #[cfg(not(target_os = "linux"))]
802 {
803 if ridge_t.is_nan() || ridge_beta.is_nan() {
804 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
805 reason: "ridge is NaN".to_string(),
806 });
807 }
808 Err(ArrowSchurGpuFailure::Unavailable)
809 }
810 #[cfg(target_os = "linux")]
811 {
812 self.inner.refactor_and_solve(ridge_t, ridge_beta)
813 }
814 }
815
816 pub fn factor_at(
827 &self,
828 ridge_t: f64,
829 ridge_beta: f64,
830 ) -> Result<ResidentBaseRidgeFactorsHandle, ArrowSchurGpuFailure> {
831 #[cfg(not(target_os = "linux"))]
832 {
833 if ridge_t.is_nan() || ridge_beta.is_nan() {
834 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
835 reason: "ridge is NaN".to_string(),
836 });
837 }
838 Err(ArrowSchurGpuFailure::Unavailable)
839 }
840 #[cfg(target_os = "linux")]
841 {
842 Ok(ResidentBaseRidgeFactorsHandle {
843 inner: self.inner.factor_at(ridge_t, ridge_beta)?,
844 })
845 }
846 }
847
848 pub fn solve_with_factors(
852 &self,
853 factors: &ResidentBaseRidgeFactorsHandle,
854 g_t: &[f64],
855 g_beta: &[f64],
856 ) -> Result<ArrowSchurGpuSolution, ArrowSchurGpuFailure> {
857 #[cfg(not(target_os = "linux"))]
858 {
859 if g_t.iter().chain(g_beta).any(|v| !v.is_finite()) {
860 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
861 reason: "gradient entry is not finite".to_string(),
862 });
863 }
864 match *factors {}
865 }
866 #[cfg(target_os = "linux")]
867 {
868 self.inner
869 .solve_with_gradient_checked(&factors.inner, g_t, g_beta)
870 }
871 }
872}
873
874#[cfg(target_os = "linux")]
883pub struct ResidentBaseRidgeFactorsHandle {
884 inner: cuda::BaseRidgeFactors,
885}
886
887#[cfg(not(target_os = "linux"))]
890pub enum ResidentBaseRidgeFactorsHandle {}
891
892impl ResidentBaseRidgeFactorsHandle {
893 #[must_use]
898 pub fn matches_ridge(&self, ridge_t: f64, ridge_beta: f64) -> bool {
899 #[cfg(not(target_os = "linux"))]
900 {
901 if ridge_t.is_nan() || ridge_beta.is_nan() {
903 return false;
904 }
905 match *self {}
906 }
907 #[cfg(target_os = "linux")]
908 {
909 self.inner.matches_ridge(ridge_t, ridge_beta)
910 }
911 }
912}
913
914pub fn gpu_schur_matvec_backend(
957 sys: &ArrowSchurSystem,
958 ridge_t: f64,
959 ridge_beta: f64,
960) -> Result<crate::arrow_schur::GpuSchurMatvec, ArrowSchurGpuFailure> {
961 if sys.htbeta_matvec.is_some() {
964 return build_row_procedural_matvec(sys, ridge_t, ridge_beta);
965 }
966
967 #[cfg(not(target_os = "linux"))]
968 {
969 if ridge_t.is_nan() || ridge_beta.is_nan() {
972 return Err(ArrowSchurGpuFailure::Unavailable);
973 }
974 Err(ArrowSchurGpuFailure::Unavailable)
975 }
976
977 #[cfg(target_os = "linux")]
978 {
979 cuda::build_schur_matvec_backend(sys, ridge_t, ridge_beta)
980 }
981}
982
983pub fn build_framed_resident_evidence_matvec(
991 sys: &ArrowSchurSystem,
992 ridge_t: f64,
993 ridge_beta: f64,
994 apply_budget: usize,
995) -> Result<Option<crate::arrow_schur::GpuSchurMatvec>, ArrowSchurGpuFailure> {
996 #[cfg(not(target_os = "linux"))]
997 {
998 if sys.k == 0 || !ridge_t.is_finite() || !ridge_beta.is_finite() || apply_budget == 0 {
1002 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
1003 reason: "resident evidence matvec received degenerate dimensions or ridge"
1004 .to_string(),
1005 });
1006 }
1007 Ok(None)
1008 }
1009 #[cfg(target_os = "linux")]
1010 {
1011 cuda::build_framed_resident_evidence_matvec(sys, ridge_t, ridge_beta, apply_budget)
1012 }
1013}
1014
1015fn build_row_procedural_matvec(
1032 sys: &ArrowSchurSystem,
1033 ridge_t: f64,
1034 ridge_beta: f64,
1035) -> Result<crate::arrow_schur::GpuSchurMatvec, ArrowSchurGpuFailure> {
1036 use std::sync::Arc;
1037 let n = sys.rows.len();
1038 let k = sys.k;
1039 let forward = sys
1040 .htbeta_matvec
1041 .clone()
1042 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
1043 let transpose = sys.htbeta_transpose_matvec.clone().ok_or_else(|| {
1044 ArrowSchurGpuFailure::SchurFactorFailed {
1049 reason: "row-procedural Schur matvec requires htbeta_transpose_matvec; \
1050 forward operator installed without its sparse adjoint"
1051 .to_string(),
1052 }
1053 })?;
1054
1055 let mut factors: Vec<Array2<f64>> = Vec::with_capacity(n);
1060 for (i, row) in sys.rows.iter().enumerate() {
1061 let di = row.htt.nrows();
1062 if row.htt.ncols() != di || row.gt.len() != di {
1063 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
1064 reason: format!("row {i}: malformed H_tt block {:?}", row.htt.dim()),
1065 });
1066 }
1067 let mut block = row.htt.clone();
1068 for r in 0..di {
1069 block[[r, r]] += ridge_t;
1070 }
1071 let factor = cholesky_factor_in_place(block.view(), CholeskyGuard::NonnegativePivot)
1072 .ok_or_else(|| {
1073 ArrowSchurGpuFailure::RidgeBumpRequired {
1077 row: i,
1078 bump: ridge_bump_to_make_pd(row.htt.view(), ridge_t),
1079 }
1080 })?;
1081 factors.push(factor);
1082 }
1083
1084 let penalty_op = sys.effective_penalty_op();
1091 let row_dims: Vec<usize> = sys.rows.iter().map(|row| row.htt.nrows()).collect();
1092
1093 let closure: crate::arrow_schur::GpuSchurMatvec =
1094 Arc::new(move |x: &Array1<f64>, out: &mut Array1<f64>| {
1095 assert_eq!(x.len(), k, "row-procedural matvec: x.len() != k");
1096 assert_eq!(out.len(), k, "row-procedural matvec: out.len() != k");
1097
1098 {
1101 let x_slice = x.as_slice().expect("x must be contiguous");
1102 let out_slice = out.as_slice_mut().expect("out must be contiguous");
1103 for a in 0..k {
1104 out_slice[a] = ridge_beta * x_slice[a];
1105 }
1106 penalty_op.matvec(x_slice, out_slice);
1107 }
1108
1109 let parallel = n >= crate::arrow_schur::SCHUR_MATVEC_PARALLEL_ROW_MIN
1136 && rayon::current_thread_index().is_none();
1137 if parallel {
1138 use rayon::prelude::*;
1139 const CHUNK: usize = 64;
1140 let partials: Vec<Array1<f64>> = (0..n)
1141 .into_par_iter()
1142 .chunks(CHUNK)
1143 .map(|idxs| {
1144 let mut neg = Array1::<f64>::zeros(k);
1148 for i in idxs {
1149 let di = row_dims[i];
1150 let mut v_i = Array1::<f64>::zeros(di);
1152 forward(i, x.view(), &mut v_i);
1153 let w_i = cholesky_solve_vector(factors[i].view(), v_i.view());
1155 transpose(i, w_i.view(), &mut neg);
1157 }
1158 neg
1159 })
1160 .collect();
1161 let mut neg = Array1::<f64>::zeros(k);
1173 for part in &partials {
1174 for a in 0..k {
1175 neg[a] += part[a];
1176 }
1177 }
1178 for a in 0..k {
1179 out[a] -= neg[a];
1180 }
1181 } else {
1182 let mut neg = Array1::<f64>::zeros(k);
1184 for i in 0..n {
1185 let di = row_dims[i];
1186 let mut v_i = Array1::<f64>::zeros(di);
1188 forward(i, x.view(), &mut v_i);
1189 let w_i = cholesky_solve_vector(factors[i].view(), v_i.view());
1191 transpose(i, w_i.view(), &mut neg);
1193 }
1194 for a in 0..k {
1195 out[a] -= neg[a];
1196 }
1197 }
1198 });
1199
1200 Ok(closure)
1201}
1202
1203pub fn solve_reduced_beta_pcg_with_diagnostics(
1204 s_acc: &Array2<f64>,
1205 rhs_beta: &Array1<f64>,
1206 max_iterations: usize,
1207 relative_tolerance: f64,
1208) -> Result<(Array1<f64>, ArrowPcgDiagnostics), ArrowSchurGpuFailure> {
1209 let k = rhs_beta.len();
1210 if s_acc.dim() != (k, k) {
1211 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
1212 reason: format!(
1213 "reduced-β GPU PCG requires a square (k×k) Schur block; got {:?} for k={k}",
1214 s_acc.dim()
1215 ),
1216 });
1217 }
1218 if k == 0 {
1219 return Err(ArrowSchurGpuFailure::Unavailable);
1220 }
1221
1222 #[cfg(not(target_os = "linux"))]
1223 {
1224 if relative_tolerance.is_nan() || max_iterations == 0 {
1225 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
1226 reason: "reduced-β GPU PCG: invalid CG controls".to_string(),
1227 });
1228 }
1229 Err(ArrowSchurGpuFailure::Unavailable)
1230 }
1231
1232 #[cfg(target_os = "linux")]
1233 {
1234 cuda::solve_reduced_beta_pcg_with_diagnostics(
1235 s_acc,
1236 rhs_beta,
1237 max_iterations,
1238 relative_tolerance,
1239 )
1240 }
1241}
1242
1243pub fn solve_sae_matrix_free_pcg(
1244 sys: &ArrowSchurSystem,
1245 data: &DeviceSaePcgData,
1246 ridge_t: f64,
1247 ridge_beta: f64,
1248 rhs_beta: &Array1<f64>,
1249 max_iterations: usize,
1250 relative_tolerance: f64,
1251) -> Result<(Array1<f64>, ArrowPcgDiagnostics), ArrowSchurGpuFailure> {
1252 if sys.k != data.beta_dim || rhs_beta.len() != data.beta_dim || data.p == 0 {
1253 return Err(ArrowSchurGpuFailure::Unavailable);
1254 }
1255 #[cfg(not(target_os = "linux"))]
1256 {
1257 if ridge_t.is_nan()
1258 || ridge_beta.is_nan()
1259 || relative_tolerance.is_nan()
1260 || max_iterations == 0
1261 {
1262 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
1263 reason: "SAE matrix-free GPU PCG: invalid controls".to_string(),
1264 });
1265 }
1266 Err(ArrowSchurGpuFailure::Unavailable)
1267 }
1268 #[cfg(target_os = "linux")]
1269 {
1270 if data.frame.is_some() {
1277 cuda::solve_sae_matrix_free_pcg_framed(
1278 sys,
1279 data,
1280 ridge_t,
1281 ridge_beta,
1282 rhs_beta,
1283 max_iterations,
1284 relative_tolerance,
1285 )
1286 } else {
1287 cuda::solve_sae_matrix_free_pcg(
1288 sys,
1289 data,
1290 ridge_t,
1291 ridge_beta,
1292 rhs_beta,
1293 max_iterations,
1294 relative_tolerance,
1295 )
1296 }
1297 }
1298}
1299
1300pub trait SaeResidentFrame {
1320 fn refresh(&self, sys: &ArrowSchurSystem) -> Result<(), ArrowSchurGpuFailure>;
1326
1327 fn resolve(
1335 &self,
1336 sys: &ArrowSchurSystem,
1337 ridge_t: f64,
1338 ridge_beta: f64,
1339 rhs_beta: &Array1<f64>,
1340 max_iterations: usize,
1341 relative_tolerance: f64,
1342 ) -> Result<(Array1<f64>, ArrowPcgDiagnostics), ArrowSchurGpuFailure>;
1343}
1344
1345pub fn build_sae_resident_frame(
1353 sys: &ArrowSchurSystem,
1354 cg_iters: usize,
1355) -> Result<
1356 Option<std::sync::Arc<dyn SaeResidentFrame + Send + Sync>>,
1357 ArrowSchurGpuFailure,
1358> {
1359 if sys.k == 0 || cg_iters == 0 {
1364 return Ok(None);
1365 }
1366 #[cfg(target_os = "linux")]
1367 {
1368 cuda::ResidentSaeFrameHandle::build(sys, cg_iters).map(|handle| {
1369 handle.map(|frame| {
1370 std::sync::Arc::new(frame)
1371 as std::sync::Arc<dyn SaeResidentFrame + Send + Sync>
1372 })
1373 })
1374 }
1375 #[cfg(not(target_os = "linux"))]
1377 {
1378 Ok(None)
1379 }
1380}
1381
1382#[cfg(target_os = "linux")]
1390pub struct FrameHostOperands {
1391 pub s_off: Vec<i32>,
1392 pub s_m: Vec<i32>,
1393 pub s_r: Vec<i32>,
1394 pub s_ptr: Vec<i32>,
1395 pub s_data: Vec<f64>,
1396 pub s_blocks: usize,
1397 pub g_off_i: Vec<i32>,
1398 pub g_off_j: Vec<i32>,
1399 pub g_ri: Vec<i32>,
1400 pub g_rj: Vec<i32>,
1401 pub g_mj: Vec<i32>,
1402 pub g_ptr: Vec<i32>,
1403 pub g_data: Vec<f64>,
1404 pub w_ptr: Vec<i32>,
1405 pub w_data: Vec<f64>,
1406 pub g_group_ptr: Vec<i32>,
1410 pub g_group_blocks: Vec<i32>,
1411 pub g_group_off: Vec<i32>,
1412 pub g_group_width: Vec<i32>,
1413 pub g_groups: usize,
1414 pub g_max_work: usize,
1415 pub htb_ptr: Vec<i32>,
1416 pub htb: Vec<f64>,
1417 pub q_of: Vec<i32>,
1418 pub n_rows: usize,
1419 pub k: usize,
1420 pub max_q: usize,
1421}
1422
1423#[cfg(target_os = "linux")]
1424fn frame_checked_i32(value: usize) -> Result<i32, ArrowSchurGpuFailure> {
1425 i32::try_from(value).map_err(|_| ArrowSchurGpuFailure::Unavailable)
1426}
1427
1428#[cfg(target_os = "linux")]
1438fn deterministic_output_groups(
1439 spans: &[(usize, usize)],
1440 output_len: usize,
1441 label: &str,
1442) -> Result<DeterministicOutputGroups, ArrowSchurGpuFailure> {
1443 use std::collections::HashMap;
1444
1445 let malformed = |reason: String| ArrowSchurGpuFailure::SchurFactorFailed {
1446 reason: format!("{label} deterministic output ownership: {reason}"),
1447 };
1448 let mut lookup = HashMap::<(usize, usize), usize>::new();
1449 let mut group_spans = Vec::<(usize, usize)>::new();
1450 let mut group_blocks = Vec::<Vec<usize>>::new();
1451 for (block, &(offset, width)) in spans.iter().enumerate() {
1452 let end = offset.checked_add(width).ok_or_else(|| {
1453 malformed(format!(
1454 "block {block} output span overflows ({offset} + {width})"
1455 ))
1456 })?;
1457 if width == 0 || end > output_len {
1458 return Err(malformed(format!(
1459 "block {block} has output span [{offset}, {end}) outside [0, {output_len})"
1460 )));
1461 }
1462 let group = if let Some(&group) = lookup.get(&(offset, width)) {
1463 group
1464 } else {
1465 let group = group_spans.len();
1466 lookup.insert((offset, width), group);
1467 group_spans.push((offset, width));
1468 group_blocks.push(Vec::new());
1469 group
1470 };
1471 group_blocks[group].push(block);
1472 }
1473
1474 let mut by_offset: Vec<usize> = (0..group_spans.len()).collect();
1475 by_offset.sort_unstable_by_key(|&group| group_spans[group].0);
1476 for adjacent in by_offset.windows(2) {
1477 let left = group_spans[adjacent[0]];
1478 let right = group_spans[adjacent[1]];
1479 let left_end = left.0 + left.1;
1480 if left_end > right.0 {
1481 return Err(malformed(format!(
1482 "distinct output spans [{}, {}) and [{}, {}) overlap",
1483 left.0,
1484 left_end,
1485 right.0,
1486 right.0 + right.1
1487 )));
1488 }
1489 }
1490
1491 let mut ptr = Vec::with_capacity(group_spans.len() + 1);
1492 let mut blocks = Vec::with_capacity(spans.len());
1493 ptr.push(0_i32);
1494 for members in &group_blocks {
1495 for &block in members {
1496 blocks.push(frame_checked_i32(block)?);
1497 }
1498 ptr.push(frame_checked_i32(blocks.len())?);
1499 }
1500 let mut off = Vec::with_capacity(group_spans.len());
1501 let mut width = Vec::with_capacity(group_spans.len());
1502 let mut max_width = 0usize;
1503 for &(group_off, group_width) in &group_spans {
1504 off.push(frame_checked_i32(group_off)?);
1505 width.push(frame_checked_i32(group_width)?);
1506 max_width = max_width.max(group_width);
1507 }
1508 Ok(DeterministicOutputGroups {
1509 ptr,
1510 blocks,
1511 off,
1512 width,
1513 max_width,
1514 })
1515}
1516
1517#[cfg(target_os = "linux")]
1518struct DeterministicOutputGroups {
1519 ptr: Vec<i32>,
1520 blocks: Vec<i32>,
1521 off: Vec<i32>,
1522 width: Vec<i32>,
1523 max_width: usize,
1524}
1525
1526#[cfg(target_os = "linux")]
1531pub fn flatten_frame_host_operands(
1532 sys: &ArrowSchurSystem,
1533 data: &DeviceSaePcgData,
1534 frame: &crate::arrow_schur::DeviceSaeFrameData,
1535) -> Result<FrameHostOperands, ArrowSchurGpuFailure> {
1536 let n_rows = sys.rows.len();
1537 let k = data.beta_dim;
1538 if frame.row_htbeta.len() != n_rows
1539 || frame.ranks.len() != frame.basis_sizes.len()
1540 || frame.border_offsets.len() != frame.ranks.len()
1541 || data.smooth_blocks.len() != frame.smooth_ranks.len()
1542 {
1543 return Err(ArrowSchurGpuFailure::Unavailable);
1544 }
1545
1546 let mut s_off = Vec::new();
1548 let mut s_m = Vec::new();
1549 let mut s_r = Vec::new();
1550 let mut s_ptr = vec![0_i32];
1551 let mut s_data = Vec::<f64>::new();
1552 for (blk, &r) in data.smooth_blocks.iter().zip(frame.smooth_ranks.iter()) {
1553 let (m, mc) = blk.factor_a.dim();
1554 if m != mc {
1555 return Err(ArrowSchurGpuFailure::Unavailable);
1556 }
1557 s_off.push(frame_checked_i32(blk.global_offset)?);
1558 s_m.push(frame_checked_i32(m)?);
1559 s_r.push(frame_checked_i32(r)?);
1560 for ri in 0..m {
1561 for ci in 0..m {
1562 s_data.push(blk.factor_a[[ri, ci]]);
1563 }
1564 }
1565 s_ptr.push(frame_checked_i32(s_data.len())?);
1566 }
1567
1568 let mut g_off_i = Vec::new();
1570 let mut g_off_j = Vec::new();
1571 let mut g_ri = Vec::new();
1572 let mut g_rj = Vec::new();
1573 let mut g_mj = Vec::new();
1574 let mut g_ptr = vec![0_i32];
1575 let mut g_data = Vec::<f64>::new();
1576 let mut w_ptr = vec![0_i32];
1577 let mut w_data = Vec::<f64>::new();
1578 let mut g_output_spans = Vec::with_capacity(frame.frame_blocks.len());
1579 for blk in &frame.frame_blocks {
1580 let ri = frame.ranks[blk.atom_i];
1581 let rj = frame.ranks[blk.atom_j];
1582 let (mi, mj) = blk.g.dim();
1583 if blk.w.dim() != (ri, rj) {
1584 return Err(ArrowSchurGpuFailure::Unavailable);
1585 }
1586 g_off_i.push(frame_checked_i32(frame.border_offsets[blk.atom_i])?);
1587 g_off_j.push(frame_checked_i32(frame.border_offsets[blk.atom_j])?);
1588 g_ri.push(frame_checked_i32(ri)?);
1589 g_rj.push(frame_checked_i32(rj)?);
1590 g_mj.push(frame_checked_i32(mj)?);
1591 for r in 0..mi {
1592 for c in 0..mj {
1593 g_data.push(blk.g[[r, c]]);
1594 }
1595 }
1596 g_ptr.push(frame_checked_i32(g_data.len())?);
1597 for a in 0..ri {
1598 for b in 0..rj {
1599 w_data.push(blk.w[[a, b]]);
1600 }
1601 }
1602 w_ptr.push(frame_checked_i32(w_data.len())?);
1603 let width = mi
1604 .checked_mul(ri)
1605 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
1606 g_output_spans.push((frame.border_offsets[blk.atom_i], width));
1607 }
1608 let g_groups =
1609 deterministic_output_groups(&g_output_spans, k, "framed G matvec")?;
1610 let g_group_count = g_groups.off.len();
1611
1612 let mut htb_ptr = vec![0_i32];
1615 let mut htb = Vec::<f64>::new();
1616 let mut q_of = Vec::<i32>::with_capacity(n_rows);
1617 let mut max_q = 0usize;
1618 for (i, slab) in frame.row_htbeta.iter().enumerate() {
1619 let qi = sys.row_dims[i];
1620 let q_eff = if !slab.is_empty() && slab.len() == qi * k {
1621 qi
1622 } else {
1623 0
1624 };
1625 q_of.push(frame_checked_i32(q_eff)?);
1626 max_q = max_q.max(q_eff);
1627 if q_eff > 0 {
1628 htb.extend_from_slice(slab);
1629 }
1630 htb_ptr.push(frame_checked_i32(htb.len())?);
1631 }
1632 if max_q == 0 {
1633 max_q = 1;
1636 }
1637
1638 Ok(FrameHostOperands {
1639 s_off,
1640 s_m,
1641 s_r,
1642 s_ptr,
1643 s_data,
1644 s_blocks: data.smooth_blocks.len(),
1645 g_off_i,
1646 g_off_j,
1647 g_ri,
1648 g_rj,
1649 g_mj,
1650 g_ptr,
1651 g_data,
1652 w_ptr,
1653 w_data,
1654 g_group_ptr: g_groups.ptr,
1655 g_group_blocks: g_groups.blocks,
1656 g_group_off: g_groups.off,
1657 g_group_width: g_groups.width,
1658 g_groups: g_group_count,
1659 g_max_work: g_groups.max_width,
1660 htb_ptr,
1661 htb,
1662 q_of,
1663 n_rows,
1664 k,
1665 max_q,
1666 })
1667}
1668
1669#[cfg(target_os = "linux")]
1676pub fn compute_ainv_host(
1677 sys: &ArrowSchurSystem,
1678 q_of: &[i32],
1679 max_q: usize,
1680 n_rows: usize,
1681 ridge_t: f64,
1682) -> Result<Vec<f64>, ArrowSchurGpuFailure> {
1683 let mut ainv = vec![0.0_f64; n_rows * max_q * max_q];
1684 for (i, row) in sys.rows.iter().enumerate() {
1685 let q = q_of[i] as usize;
1686 if q == 0 {
1687 continue;
1688 }
1689 if row.htt.dim() != (q, q) {
1690 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
1691 reason: format!(
1692 "framed SAE device PCG row {i}: H_tt shape {:?} != ({q}, {q})",
1693 row.htt.dim()
1694 ),
1695 });
1696 }
1697 let mut block = row.htt.clone();
1698 for d in 0..q {
1699 block[[d, d]] += ridge_t;
1700 }
1701 let factor = cholesky_factor_in_place(block.view(), CholeskyGuard::NonnegativePivot)
1702 .ok_or_else(|| ArrowSchurGpuFailure::RidgeBumpRequired {
1703 row: i,
1704 bump: ridge_bump_to_make_pd(row.htt.view(), ridge_t),
1705 })?;
1706 for col in 0..q {
1707 let mut e = Array1::<f64>::zeros(q);
1708 e[col] = 1.0;
1709 let solved = cholesky_solve_vector(factor.view(), e.view());
1710 for r in 0..q {
1711 ainv[i * max_q * max_q + r * max_q + col] = solved[r];
1712 }
1713 }
1714 }
1715 Ok(ainv)
1716}
1717
1718pub fn solve_arrow_newton_step_dense_reference(
1722 sys: &ArrowSchurSystem,
1723 ridge_t: f64,
1724 ridge_beta: f64,
1725) -> Result<ArrowSchurGpuSolution, String> {
1726 let n = sys.rows.len();
1727 let d = sys.d;
1728 let k = sys.k;
1729 let total = n.checked_mul(d).ok_or("dimension overflow")? + k;
1730 let mut h = Array2::<f64>::zeros((total, total));
1731 let mut rhs = Array1::<f64>::zeros(total);
1732 for (i, row) in sys.rows.iter().enumerate() {
1733 let base = i * d;
1734 for c in 0..d {
1735 for r in 0..d {
1736 h[[base + r, base + c]] = row.htt[[r, c]];
1737 }
1738 h[[base + c, base + c]] += ridge_t;
1739 }
1740 for c in 0..k {
1741 for r in 0..d {
1742 let value = row.htbeta[[r, c]];
1743 h[[base + r, n * d + c]] = value;
1744 h[[n * d + c, base + r]] = value;
1745 }
1746 }
1747 for r in 0..d {
1748 rhs[base + r] = -row.gt[r];
1749 }
1750 }
1751 for c in 0..k {
1752 for r in 0..k {
1753 h[[n * d + r, n * d + c]] += sys.hbb[[r, c]];
1754 }
1755 h[[n * d + c, n * d + c]] += ridge_beta;
1756 rhs[n * d + c] = -sys.gb[c];
1757 }
1758 let factor = cholesky_factor_in_place(h.view(), CholeskyGuard::NonnegativePivot)
1767 .ok_or_else(|| "dense reference Cholesky failed".to_string())?;
1768 let mut log_det = 0.0_f64;
1769 for i in 0..total {
1770 log_det += factor[[i, i]].ln();
1771 }
1772 log_det *= 2.0;
1773 let solved = cholesky_solve_vector(factor.view(), rhs.view());
1774 let delta_t = solved.slice(ndarray::s![..n * d]).to_owned();
1775 let delta_beta = solved.slice(ndarray::s![n * d..]).to_owned();
1776 Ok(ArrowSchurGpuSolution {
1777 delta_t,
1778 delta_beta,
1779 log_det_hessian: log_det,
1780 })
1781}
1782
1783pub fn sae_framed_penalty_matvec_cpu(
1794 data: &DeviceSaePcgData,
1795 ridge_beta: f64,
1796 x: &[f64],
1797 out: &mut [f64],
1798) {
1799 let frame = data
1800 .frame
1801 .as_ref()
1802 .expect("sae_framed_penalty_matvec_cpu requires frame metadata");
1803 let k = data.beta_dim;
1804 for a in 0..k {
1805 out[a] = ridge_beta * x[a];
1806 }
1807 for (blk, &r) in data.smooth_blocks.iter().zip(frame.smooth_ranks.iter()) {
1809 let off = blk.global_offset;
1810 let m = blk.factor_a.nrows();
1811 for i_a in 0..m {
1812 for i_b in 0..r {
1813 let mut acc = 0.0_f64;
1814 for j_a in 0..m {
1815 let s = blk.factor_a[[i_a, j_a]];
1816 if s == 0.0 {
1817 continue;
1818 }
1819 acc += s * x[off + j_a * r + i_b];
1820 }
1821 out[off + i_a * r + i_b] += acc;
1822 }
1823 }
1824 }
1825 for blk in &frame.frame_blocks {
1827 let r_i = frame.ranks[blk.atom_i];
1828 let r_j = frame.ranks[blk.atom_j];
1829 let off_i = frame.border_offsets[blk.atom_i];
1830 let off_j = frame.border_offsets[blk.atom_j];
1831 let (m_i, m_j) = blk.g.dim();
1832 for li in 0..m_i {
1833 let yi_base = off_i + li * r_i;
1834 for lj in 0..m_j {
1835 let g = blk.g[[li, lj]];
1836 if g == 0.0 {
1837 continue;
1838 }
1839 let xj_base = off_j + lj * r_j;
1840 for a in 0..r_i {
1841 let mut acc = 0.0_f64;
1842 for b in 0..r_j {
1843 acc += blk.w[[a, b]] * x[xj_base + b];
1844 }
1845 out[yi_base + a] += g * acc;
1846 }
1847 }
1848 }
1849 }
1850}
1851
1852#[cfg(target_os = "linux")]
1853mod cuda {
1854 use super::{
1855 canonicalize_device_beta_factor,
1856 ArrowSchurGpuFailure, ArrowSchurGpuSolution, pack_block, pack_host,
1857 pack_host_d_and_stacked_b, project_device_beta_vector,
1858 };
1859 use crate::arrow_schur::{
1860 ArrowBetaGaugeQuotient, ArrowPcgDiagnostics, ArrowSchurSystem, DeviceSaeFrameData,
1861 DeviceSaePcgData, PcgStopReason,
1862 };
1863 use cudarc::cublas::sys::{
1864 cublasDiagType_t, cublasFillMode_t, cublasOperation_t, cublasSideMode_t, cublasStatus_t,
1865 };
1866 use cudarc::cublas::{CudaBlas, Gemm, GemmConfig, Gemv, GemvConfig};
1867 use cudarc::cusolver::{DnHandle, sys as cusolver_sys};
1868 use cudarc::driver::{
1869 CudaContext, CudaModule, CudaSlice, CudaStream, DevicePtr, DevicePtrMut, LaunchConfig,
1870 PushKernelArg,
1871 };
1872 use gam_gpu::driver::to_i32;
1873 use gam_gpu::linalg_dispatch::{DispatchOp, route_through_gpu};
1874 use ndarray::Array1;
1875 use std::sync::{Arc, OnceLock};
1876
1877 struct RowSlot {
1882 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,
1891 bump: Option<f64>,
1894 tile_partial_schur: Option<Vec<f64>>, tile_partial_rhs: Option<Vec<f64>>, delta_t_block: Vec<f64>, }
1900
1901 pub(super) fn solve_multi_gpu(
1922 sys: &ArrowSchurSystem,
1923 ridge_t: f64,
1924 ridge_beta: f64,
1925 newton_schur_tikhonov_rel_floor: Option<f64>,
1926 ) -> Result<ArrowSchurGpuSolution, ArrowSchurGpuFailure> {
1927 let n = sys.rows.len();
1928 let d = sys.d;
1929 let k = sys.k;
1930 if n == 0 || d == 0 || k == 0 {
1931 return Err(ArrowSchurGpuFailure::Unavailable);
1932 }
1933 if sys.hbb_matvec.is_some() || sys.htbeta_matvec.is_some() || sys.hbb.dim() != (k, k) {
1937 return Err(ArrowSchurGpuFailure::Unavailable);
1938 }
1939
1940 let runtime = super::resolve_runtime_for_device_path()?
1941 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
1942 if runtime.device_count() < 2 {
1943 return Err(ArrowSchurGpuFailure::Unavailable);
1944 }
1945
1946 let mut slots: Vec<RowSlot> = Vec::with_capacity(n);
1948 for row in &sys.rows {
1949 if row.htt.dim() != (d, d) || row.htbeta.dim() != (d, k) || row.gt.len() != d {
1950 return Err(ArrowSchurGpuFailure::Unavailable);
1951 }
1952 let mut d_block = Vec::with_capacity(d * d);
1953 let mut b_block = Vec::with_capacity(d * k);
1954 let mut g_vec = Vec::with_capacity(d);
1955 pack_block(row, ridge_t, d, k, &mut d_block, &mut b_block, &mut g_vec);
1956 slots.push(RowSlot {
1957 d_block,
1958 b_block,
1959 g_vec,
1960 l_block: Vec::new(),
1961 u_vec: Vec::new(),
1962 y_block: Vec::new(),
1963 log_det_local: 0.0,
1964 bump: None,
1965 tile_partial_schur: None,
1966 tile_partial_rhs: None,
1967 delta_t_block: vec![0.0; d],
1968 });
1969 }
1970
1971 let forward_ok = gam_gpu::pool::scatter_batched(runtime, &mut slots, |ordinal, tile| {
1973 forward_tile(ordinal, d, k, tile)
1974 });
1975 if forward_ok.is_none() {
1976 return Err(ArrowSchurGpuFailure::Unavailable);
1977 }
1978
1979 let row_base_of_tile = gam_gpu::pool::balanced_partition(runtime, n);
1981 if let Some((row, bump)) = slots
1982 .iter()
1983 .enumerate()
1984 .find_map(|(i, slot)| slot.bump.map(|b| (i, b)))
1985 {
1986 return Err(ArrowSchurGpuFailure::RidgeBumpRequired { row, bump });
1987 }
1988
1989 let mut schur_host = vec![0.0_f64; k * k];
1994 for col in 0..k {
1995 for row in 0..k {
1996 let mut v = sys.hbb[[row, col]];
1997 if row == col {
1998 v += ridge_beta;
1999 }
2000 schur_host[col * k + row] = v;
2001 }
2002 }
2003 let mut rhs_host: Vec<f64> = sys.gb.iter().map(|v| -v).collect();
2004 let mut log_det = 0.0_f64;
2005 for start in tile_starts(&row_base_of_tile) {
2006 let slot = &slots[start];
2007 let partial_schur = slot
2008 .tile_partial_schur
2009 .as_ref()
2010 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
2011 let partial_rhs = slot
2012 .tile_partial_rhs
2013 .as_ref()
2014 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
2015 for idx in 0..k * k {
2020 schur_host[idx] += partial_schur[idx];
2021 }
2022 for a in 0..k {
2023 rhs_host[a] += partial_rhs[a];
2024 }
2025 }
2026 for slot in &slots {
2027 log_det += slot.log_det_local;
2028 }
2029 let canonical_factor = canonicalize_device_beta_factor(
2030 sys.beta_gauge_quotient.as_ref(),
2031 newton_schur_tikhonov_rel_floor,
2032 sys.k,
2033 &mut schur_host,
2034 )?;
2035 rhs_host = project_device_beta_vector(sys.beta_gauge_quotient.as_ref(), sys.k, rhs_host)?;
2036
2037 let primary = runtime.selected_device().ordinal;
2041 let stream = gam_gpu::device_runtime::cuda_context_for(primary)
2042 .and_then(|ctx| ctx.new_stream().ok())
2043 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
2044 let solver =
2045 DnHandle::new(stream.clone()).map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2046 let blas = CudaBlas::new(stream.clone()).map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2047 let mut schur_dev = stream
2048 .clone_htod(&schur_host)
2049 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2050 let mut rhs_dev = stream
2051 .clone_htod(&rhs_host)
2052 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2053 if !canonical_factor {
2054 let info = potrf_single(&solver, &stream, k, &mut schur_dev)?;
2055 if info != 0 {
2056 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
2057 reason: format!("multi-GPU Schur Cholesky failed at pivot {info}"),
2058 });
2059 }
2060 }
2061 trsm_single(&blas, &stream, k, &schur_dev, &mut rhs_dev, false, false)?;
2062 trsm_single(&blas, &stream, k, &schur_dev, &mut rhs_dev, false, true)?;
2063 let delta_beta_host = project_device_beta_vector(
2064 sys.beta_gauge_quotient.as_ref(),
2065 sys.k,
2066 stream
2067 .clone_dtoh(&rhs_dev)
2068 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
2069 )?;
2070 let delta_beta = Array1::from_vec(delta_beta_host.clone());
2071 let l_schur_host = stream
2072 .clone_dtoh(&schur_dev)
2073 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2074 for j in 0..k {
2075 log_det += l_schur_host[j * k + j].ln();
2076 }
2077 log_det *= 2.0;
2078
2079 let delta_beta_ref = &delta_beta_host;
2081 let back_ok = gam_gpu::pool::scatter_batched(runtime, &mut slots, |ordinal, tile| {
2082 back_sub_tile(ordinal, d, k, delta_beta_ref, tile)
2083 });
2084 if back_ok.is_none() {
2085 return Err(ArrowSchurGpuFailure::Unavailable);
2086 }
2087
2088 let mut delta_t = Array1::<f64>::zeros(n * d);
2090 for (i, slot) in slots.iter().enumerate() {
2091 let base = i * d;
2092 for r in 0..d {
2093 delta_t[base + r] = slot.delta_t_block[r];
2094 }
2095 }
2096
2097 Ok(ArrowSchurGpuSolution {
2098 delta_t,
2099 delta_beta,
2100 log_det_hessian: log_det,
2101 })
2102 }
2103
2104 fn tile_starts(tiles: &[(usize, std::ops::Range<usize>)]) -> impl Iterator<Item = usize> + '_ {
2107 tiles.iter().map(|(_, range)| range.start)
2108 }
2109
2110 fn forward_tile(ordinal: usize, d: usize, k: usize, tile: &mut [RowSlot]) -> Option<()> {
2118 if tile.is_empty() {
2119 return Some(());
2120 }
2121 let stream = gam_gpu::device_runtime::cuda_context_for(ordinal)
2124 .and_then(|ctx| ctx.new_stream().ok())?;
2125 let solver = DnHandle::new(stream.clone()).ok()?;
2126 let blas = CudaBlas::new(stream.clone()).ok()?;
2127 let m = tile.len();
2128
2129 let mut d_host = Vec::with_capacity(m * d * d);
2132 let mut b_host = Vec::with_capacity(m * d * k);
2133 let mut g_host = Vec::with_capacity(m * d);
2134 for slot in tile.iter() {
2135 d_host.extend_from_slice(&slot.d_block);
2136 b_host.extend_from_slice(&slot.b_block);
2137 g_host.extend_from_slice(&slot.g_vec);
2138 }
2139 let mut d_dev = stream.clone_htod(&d_host).ok()?;
2140 let mut b_dev = stream.clone_htod(&b_host).ok()?;
2141 let mut g_dev = stream.clone_htod(&g_host).ok()?;
2142
2143 let info_host = potrf_batched(&solver, &stream, d, m, &mut d_dev).ok()?;
2149 if let Some(local) = info_host.iter().position(|info| *info != 0) {
2150 tile[local].bump = Some(super::ridge_bump_to_make_pd_colmajor(
2151 &tile[local].d_block,
2152 d,
2153 ));
2154 return Some(());
2155 }
2156
2157 trsm_batched_lower_inplace(&blas, &stream, d, m, 1, &d_dev, &mut g_dev).ok()?;
2159 trsm_batched_lower_inplace(&blas, &stream, d, m, k, &d_dev, &mut b_dev).ok()?;
2160
2161 let mut schur_dev = stream.alloc_zeros::<f64>(k * k).ok()?;
2163 let mut rhs_dev = stream.alloc_zeros::<f64>(k).ok()?;
2164 accumulate_schur(&blas, d, k, m, &b_dev, &g_dev, &mut schur_dev, &mut rhs_dev).ok()?;
2165
2166 let l_host = stream.clone_dtoh(&d_dev).ok()?;
2168 let u_host = stream.clone_dtoh(&g_dev).ok()?;
2169 let y_host = stream.clone_dtoh(&b_dev).ok()?;
2170 let partial_schur = stream.clone_dtoh(&schur_dev).ok()?;
2171 let partial_rhs = stream.clone_dtoh(&rhs_dev).ok()?;
2172
2173 for (local, slot) in tile.iter_mut().enumerate() {
2174 let l_base = local * d * d;
2175 let u_base = local * d;
2176 let y_base = local * d * k;
2177 slot.l_block = l_host[l_base..l_base + d * d].to_vec();
2178 slot.u_vec = u_host[u_base..u_base + d].to_vec();
2179 slot.y_block = y_host[y_base..y_base + d * k].to_vec();
2180 let mut log_det_local = 0.0_f64;
2181 for j in 0..d {
2182 log_det_local += l_host[l_base + j * d + j].ln();
2183 }
2184 slot.log_det_local = log_det_local;
2185 }
2186 tile[0].tile_partial_schur = Some(partial_schur);
2187 tile[0].tile_partial_rhs = Some(partial_rhs);
2188 Some(())
2189 }
2190
2191 fn back_sub_tile(
2195 ordinal: usize,
2196 d: usize,
2197 k: usize,
2198 delta_beta: &[f64],
2199 tile: &mut [RowSlot],
2200 ) -> Option<()> {
2201 if tile.is_empty() {
2202 return Some(());
2203 }
2204 let stream = gam_gpu::device_runtime::cuda_context_for(ordinal)
2207 .and_then(|ctx| ctx.new_stream().ok())?;
2208 let blas = CudaBlas::new(stream.clone()).ok()?;
2209 let m = tile.len();
2210
2211 let mut l_host = Vec::with_capacity(m * d * d);
2212 let mut u_host = Vec::with_capacity(m * d);
2213 let mut y_host = Vec::with_capacity(m * d * k);
2214 for slot in tile.iter() {
2215 l_host.extend_from_slice(&slot.l_block);
2216 u_host.extend_from_slice(&slot.u_vec);
2217 y_host.extend_from_slice(&slot.y_block);
2218 }
2219 let d_dev = stream.clone_htod(&l_host).ok()?;
2220 let mut g_dev = stream.clone_htod(&u_host).ok()?;
2221 let b_dev = stream.clone_htod(&y_host).ok()?;
2222 let rhs_dev = stream.clone_htod(&delta_beta.to_vec()).ok()?;
2223
2224 accumulate_back_sub_rhs(&blas, d, k, m, &b_dev, &rhs_dev, &mut g_dev).ok()?;
2226 trsm_batched_lower_inplace_transposed(&blas, &stream, d, m, 1, &d_dev, &mut g_dev).ok()?;
2227 let x_host = stream.clone_dtoh(&g_dev).ok()?;
2228 for (local, slot) in tile.iter_mut().enumerate() {
2229 let base = local * d;
2230 for r in 0..d {
2231 slot.delta_t_block[r] = -x_host[base + r];
2232 }
2233 }
2234 Some(())
2235 }
2236
2237 pub(super) fn solve(
2238 sys: &ArrowSchurSystem,
2239 ridge_t: f64,
2240 ridge_beta: f64,
2241 newton_schur_tikhonov_rel_floor: Option<f64>,
2242 ) -> Result<ArrowSchurGpuSolution, ArrowSchurGpuFailure> {
2243 let n = sys.rows.len();
2244 let d = sys.d;
2245 let k = sys.k;
2246 let runtime = route_through_gpu(DispatchOp::SmallDenseBatchedPotrf { p: d, batch: n })
2247 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
2248
2249 let stream = gam_gpu::device_runtime::cuda_context_for(runtime.device.ordinal)
2250 .and_then(|ctx| ctx.new_stream().ok())
2251 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
2252 let solver =
2253 DnHandle::new(stream.clone()).map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2254 let blas = CudaBlas::new(stream.clone()).map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2255
2256 let (d_host, b_host, g_host) = pack_host(sys, ridge_t);
2258 let mut d_dev = stream
2259 .clone_htod(&d_host)
2260 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2261 let mut b_dev = stream
2262 .clone_htod(&b_host)
2263 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2264 let mut g_dev = stream
2265 .clone_htod(&g_host)
2266 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2267
2268 let info_host = potrf_batched(&solver, &stream, d, n, &mut d_dev)?;
2276 if let Some(idx) = info_host.iter().position(|info| *info != 0) {
2277 return Err(ArrowSchurGpuFailure::RidgeBumpRequired {
2281 row: idx,
2282 bump: super::ridge_bump_to_make_pd(sys.rows[idx].htt.view(), ridge_t),
2283 });
2284 }
2285
2286 trsm_batched_lower_inplace(&blas, &stream, d, n, 1, &d_dev, &mut g_dev)?;
2289 trsm_batched_lower_inplace(&blas, &stream, d, n, k, &d_dev, &mut b_dev)?;
2292
2293 let schur_init: Vec<f64> = {
2312 let mut tmp = Vec::with_capacity(k * k);
2313 for col in 0..k {
2314 for row in 0..k {
2315 let mut v = sys.hbb[[row, col]];
2316 if row == col {
2317 v += ridge_beta;
2318 }
2319 tmp.push(v);
2320 }
2321 }
2322 tmp
2323 };
2324 let mut schur_dev = stream
2325 .clone_htod(&schur_init)
2326 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2327 let rhs_init: Vec<f64> = sys.gb.iter().map(|v| -v).collect();
2328 let mut rhs_dev = stream
2329 .clone_htod(&rhs_init)
2330 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2331
2332 accumulate_schur(&blas, d, k, n, &b_dev, &g_dev, &mut schur_dev, &mut rhs_dev)?;
2333 let canonical_factor = if sys.beta_gauge_quotient.is_some()
2334 || newton_schur_tikhonov_rel_floor.is_some()
2335 {
2336 let mut schur_host = stream
2337 .clone_dtoh(&schur_dev)
2338 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2339 let factorized = canonicalize_device_beta_factor(
2340 sys.beta_gauge_quotient.as_ref(),
2341 newton_schur_tikhonov_rel_floor,
2342 sys.k,
2343 &mut schur_host,
2344 )?;
2345 schur_dev = stream
2346 .clone_htod(&schur_host)
2347 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2348 if sys.beta_gauge_quotient.is_some() {
2349 let rhs_host = project_device_beta_vector(
2350 sys.beta_gauge_quotient.as_ref(),
2351 sys.k,
2352 stream
2353 .clone_dtoh(&rhs_dev)
2354 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
2355 )?;
2356 rhs_dev = stream
2357 .clone_htod(&rhs_host)
2358 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2359 }
2360 factorized
2361 } else {
2362 false
2363 };
2364
2365 if !canonical_factor {
2367 let info = potrf_single(&solver, &stream, k, &mut schur_dev)?;
2368 if info != 0 {
2369 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
2370 reason: format!("Schur Cholesky failed at pivot {info}"),
2371 });
2372 }
2373 }
2374 trsm_single(&blas, &stream, k, &schur_dev, &mut rhs_dev, false, false)?;
2376 trsm_single(&blas, &stream, k, &schur_dev, &mut rhs_dev, false, true)?;
2377 let delta_beta_host = project_device_beta_vector(
2378 sys.beta_gauge_quotient.as_ref(),
2379 sys.k,
2380 stream
2381 .clone_dtoh(&rhs_dev)
2382 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
2383 )?;
2384 if sys.beta_gauge_quotient.is_some() {
2385 rhs_dev = stream
2386 .clone_htod(&delta_beta_host)
2387 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2388 }
2389 let delta_beta = Array1::from_vec(delta_beta_host);
2390
2391 accumulate_back_sub_rhs(&blas, d, k, n, &b_dev, &rhs_dev, &mut g_dev)?;
2399 trsm_batched_lower_inplace_transposed(&blas, &stream, d, n, 1, &d_dev, &mut g_dev)?;
2400
2401 let x_host = stream
2402 .clone_dtoh(&g_dev)
2403 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2404 let mut delta_t = Array1::<f64>::zeros(n * d);
2405 for (i, v) in x_host.iter().enumerate() {
2406 delta_t[i] = -*v;
2407 }
2408
2409 let l_local_host = stream
2411 .clone_dtoh(&d_dev)
2412 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2413 let l_schur_host = stream
2414 .clone_dtoh(&schur_dev)
2415 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2416 let mut log_det = 0.0_f64;
2417 for i in 0..n {
2418 let base = i * d * d;
2419 for j in 0..d {
2420 log_det += l_local_host[base + j * d + j].ln();
2421 }
2422 }
2423 for j in 0..k {
2424 log_det += l_schur_host[j * k + j].ln();
2425 }
2426 log_det *= 2.0;
2427
2428 Ok(ArrowSchurGpuSolution {
2429 delta_t,
2430 delta_beta,
2431 log_det_hessian: log_det,
2432 })
2433 }
2434
2435 fn potrf_batched(
2436 solver: &DnHandle,
2437 stream: &Arc<CudaStream>,
2438 p: usize,
2439 batch: usize,
2440 matrices: &mut CudaSlice<f64>,
2441 ) -> Result<Vec<i32>, ArrowSchurGpuFailure> {
2442 let p_i = to_i32(p).ok_or(ArrowSchurGpuFailure::Unavailable)?;
2443 let batch_i = to_i32(batch).ok_or(ArrowSchurGpuFailure::Unavailable)?;
2444 let matrix_len = p * p;
2445 let bytes_per = (matrix_len * std::mem::size_of::<f64>()) as u64;
2446 let (base_ptr, _record) = matrices.device_ptr_mut(stream);
2447 let mut ptrs = Vec::with_capacity(batch);
2448 for idx in 0..batch {
2449 ptrs.push(base_ptr + (idx as u64) * bytes_per);
2450 }
2451 let mut ptrs_dev = stream
2452 .clone_htod(&ptrs)
2453 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2454 let mut info_dev = stream
2455 .alloc_zeros::<i32>(batch)
2456 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2457 let status = {
2458 let (ptrs_ptr, _ptrs_record) = ptrs_dev.device_ptr_mut(stream);
2459 let (info_ptr, _info_record) = info_dev.device_ptr_mut(stream);
2460 unsafe {
2463 cusolver_sys::cusolverDnDpotrfBatched(
2464 solver.cu(),
2465 cusolver_sys::cublasFillMode_t::CUBLAS_FILL_MODE_LOWER,
2466 p_i,
2467 ptrs_ptr as *mut *mut f64,
2468 p_i,
2469 info_ptr as *mut i32,
2470 batch_i,
2471 )
2472 }
2473 };
2474 if status != cusolver_sys::cusolverStatus_t::CUSOLVER_STATUS_SUCCESS {
2475 return Err(ArrowSchurGpuFailure::Unavailable);
2476 }
2477 stream
2478 .clone_dtoh(&info_dev)
2479 .map_err(|_| ArrowSchurGpuFailure::Unavailable)
2480 }
2481
2482 fn potrf_single(
2483 solver: &DnHandle,
2484 stream: &Arc<CudaStream>,
2485 p: usize,
2486 matrix: &mut CudaSlice<f64>,
2487 ) -> Result<i32, ArrowSchurGpuFailure> {
2488 let p_i = to_i32(p).ok_or(ArrowSchurGpuFailure::Unavailable)?;
2489 let uplo = cusolver_sys::cublasFillMode_t::CUBLAS_FILL_MODE_LOWER;
2490 let mut lwork = 0_i32;
2491 {
2492 let (mat_ptr, _rec) = matrix.device_ptr_mut(stream);
2493 let status = unsafe {
2495 cusolver_sys::cusolverDnDpotrf_bufferSize(
2496 solver.cu(),
2497 uplo,
2498 p_i,
2499 mat_ptr as *mut f64,
2500 p_i,
2501 &mut lwork,
2502 )
2503 };
2504 if status != cusolver_sys::cusolverStatus_t::CUSOLVER_STATUS_SUCCESS {
2505 return Err(ArrowSchurGpuFailure::Unavailable);
2506 }
2507 }
2508 let lwork_usize = usize::try_from(lwork).map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2509 let mut workspace = stream
2510 .alloc_zeros::<f64>(lwork_usize.max(1))
2511 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2512 let mut info_dev = stream
2513 .alloc_zeros::<i32>(1)
2514 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2515 {
2516 let (mat_ptr, _rec) = matrix.device_ptr_mut(stream);
2517 let (work_ptr, _wrec) = workspace.device_ptr_mut(stream);
2518 let (info_ptr, _irec) = info_dev.device_ptr_mut(stream);
2519 let status = unsafe {
2521 cusolver_sys::cusolverDnDpotrf(
2522 solver.cu(),
2523 uplo,
2524 p_i,
2525 mat_ptr as *mut f64,
2526 p_i,
2527 work_ptr as *mut f64,
2528 lwork,
2529 info_ptr as *mut i32,
2530 )
2531 };
2532 if status != cusolver_sys::cusolverStatus_t::CUSOLVER_STATUS_SUCCESS {
2533 return Err(ArrowSchurGpuFailure::Unavailable);
2534 }
2535 }
2536 let info_host = stream
2537 .clone_dtoh(&info_dev)
2538 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2539 Ok(info_host[0])
2540 }
2541
2542 #[derive(Clone, Copy)]
2556 struct BatchedTrsmSpec {
2557 d: usize,
2558 n: usize,
2559 nrhs: usize,
2560 rhs_leading_dim: usize,
2561 rhs_block_stride: usize,
2562 transposed: bool,
2563 }
2564
2565 impl BatchedTrsmSpec {
2566 fn block_tiles(d: usize, n: usize, nrhs: usize) -> Self {
2568 Self {
2569 d,
2570 n,
2571 nrhs,
2572 rhs_leading_dim: d,
2573 rhs_block_stride: d * nrhs,
2574 transposed: false,
2575 }
2576 }
2577
2578 fn stacked_matrix(d: usize, n: usize, nrhs: usize) -> Self {
2581 Self {
2582 d,
2583 n,
2584 nrhs,
2585 rhs_leading_dim: n * d,
2586 rhs_block_stride: d,
2587 transposed: false,
2588 }
2589 }
2590
2591 fn transposed(self) -> Self {
2592 Self {
2593 transposed: true,
2594 ..self
2595 }
2596 }
2597 }
2598
2599 fn trsm_batched_lower_inplace(
2600 blas: &CudaBlas,
2601 stream: &Arc<CudaStream>,
2602 d: usize,
2603 n: usize,
2604 nrhs: usize,
2605 l_stack: &CudaSlice<f64>,
2606 rhs_stack: &mut CudaSlice<f64>,
2607 ) -> Result<(), ArrowSchurGpuFailure> {
2608 trsm_batched_inplace_inner(
2609 blas,
2610 stream,
2611 BatchedTrsmSpec::block_tiles(d, n, nrhs),
2612 l_stack,
2613 rhs_stack,
2614 )
2615 }
2616
2617 fn trsm_batched_lower_inplace_transposed(
2619 blas: &CudaBlas,
2620 stream: &Arc<CudaStream>,
2621 d: usize,
2622 n: usize,
2623 nrhs: usize,
2624 l_stack: &CudaSlice<f64>,
2625 rhs_stack: &mut CudaSlice<f64>,
2626 ) -> Result<(), ArrowSchurGpuFailure> {
2627 trsm_batched_inplace_inner(
2628 blas,
2629 stream,
2630 BatchedTrsmSpec::block_tiles(d, n, nrhs).transposed(),
2631 l_stack,
2632 rhs_stack,
2633 )
2634 }
2635
2636 fn trsm_batched_lower_inplace_stacked(
2639 blas: &CudaBlas,
2640 stream: &Arc<CudaStream>,
2641 d: usize,
2642 n: usize,
2643 nrhs: usize,
2644 l_stack: &CudaSlice<f64>,
2645 rhs_stacked: &mut CudaSlice<f64>,
2646 ) -> Result<(), ArrowSchurGpuFailure> {
2647 trsm_batched_inplace_inner(
2648 blas,
2649 stream,
2650 BatchedTrsmSpec::stacked_matrix(d, n, nrhs),
2651 l_stack,
2652 rhs_stacked,
2653 )
2654 }
2655
2656 fn trsm_batched_inplace_inner(
2657 blas: &CudaBlas,
2658 stream: &Arc<CudaStream>,
2659 spec: BatchedTrsmSpec,
2660 l_stack: &CudaSlice<f64>,
2661 rhs_stack: &mut CudaSlice<f64>,
2662 ) -> Result<(), ArrowSchurGpuFailure> {
2663 let BatchedTrsmSpec {
2664 d,
2665 n,
2666 nrhs,
2667 rhs_leading_dim,
2668 rhs_block_stride,
2669 transposed,
2670 } = spec;
2671 let alpha = 1.0_f64;
2672 let d_i = to_i32(d).ok_or(ArrowSchurGpuFailure::Unavailable)?;
2673 let nrhs_i = to_i32(nrhs).ok_or(ArrowSchurGpuFailure::Unavailable)?;
2674 let rhs_ld_i = to_i32(rhs_leading_dim).ok_or(ArrowSchurGpuFailure::Unavailable)?;
2675 let batch_i = to_i32(n).ok_or(ArrowSchurGpuFailure::Unavailable)?;
2676 let l_bytes_per = (d * d * std::mem::size_of::<f64>()) as u64;
2677 let rhs_bytes_per = (rhs_block_stride * std::mem::size_of::<f64>()) as u64;
2678 let (l_base, _l_record) = l_stack.device_ptr(stream);
2679 let (rhs_base, _rhs_record) = rhs_stack.device_ptr_mut(stream);
2680 let mut l_ptrs = Vec::with_capacity(n);
2681 let mut rhs_ptrs = Vec::with_capacity(n);
2682 for i in 0..n {
2683 l_ptrs.push(l_base + (i as u64) * l_bytes_per);
2684 rhs_ptrs.push(rhs_base + (i as u64) * rhs_bytes_per);
2685 }
2686 let mut l_ptrs_dev = stream
2687 .clone_htod(&l_ptrs)
2688 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2689 let mut rhs_ptrs_dev = stream
2690 .clone_htod(&rhs_ptrs)
2691 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2692 let (l_ptrs_ptr, _l_ptrs_rec) = l_ptrs_dev.device_ptr_mut(stream);
2693 let (rhs_ptrs_ptr, _rhs_ptrs_rec) = rhs_ptrs_dev.device_ptr_mut(stream);
2694 let op = if transposed {
2695 cublasOperation_t::CUBLAS_OP_T
2696 } else {
2697 cublasOperation_t::CUBLAS_OP_N
2698 };
2699 let handle = *blas.handle();
2700 let status = unsafe {
2703 cudarc::cublas::sys::cublasDtrsmBatched(
2704 handle,
2705 cublasSideMode_t::CUBLAS_SIDE_LEFT,
2706 cublasFillMode_t::CUBLAS_FILL_MODE_LOWER,
2707 op,
2708 cublasDiagType_t::CUBLAS_DIAG_NON_UNIT,
2709 d_i,
2710 nrhs_i,
2711 &alpha,
2712 l_ptrs_ptr as *const *const f64,
2713 d_i,
2714 rhs_ptrs_ptr as *const *mut f64,
2715 rhs_ld_i,
2716 batch_i,
2717 )
2718 };
2719 if status != cublasStatus_t::CUBLAS_STATUS_SUCCESS {
2720 return Err(ArrowSchurGpuFailure::Unavailable);
2721 }
2722 Ok(())
2723 }
2724
2725 fn trsm_single(
2728 blas: &CudaBlas,
2729 stream: &Arc<CudaStream>,
2730 n: usize,
2731 l: &CudaSlice<f64>,
2732 rhs: &mut CudaSlice<f64>,
2733 upper: bool,
2734 transposed: bool,
2735 ) -> Result<(), ArrowSchurGpuFailure> {
2736 let alpha = 1.0_f64;
2737 let n_i = to_i32(n).ok_or(ArrowSchurGpuFailure::Unavailable)?;
2738 let handle = *blas.handle();
2739 let (l_ptr, _l_rec) = l.device_ptr(stream);
2740 let (rhs_ptr, _rhs_rec) = rhs.device_ptr_mut(stream);
2741 let status = unsafe {
2743 cudarc::cublas::sys::cublasDtrsm_v2(
2744 handle,
2745 cublasSideMode_t::CUBLAS_SIDE_LEFT,
2746 if upper {
2747 cublasFillMode_t::CUBLAS_FILL_MODE_UPPER
2748 } else {
2749 cublasFillMode_t::CUBLAS_FILL_MODE_LOWER
2750 },
2751 if transposed {
2752 cublasOperation_t::CUBLAS_OP_T
2753 } else {
2754 cublasOperation_t::CUBLAS_OP_N
2755 },
2756 cublasDiagType_t::CUBLAS_DIAG_NON_UNIT,
2757 n_i,
2758 1,
2759 &alpha,
2760 l_ptr as *const f64,
2761 n_i,
2762 rhs_ptr as *mut f64,
2763 n_i,
2764 )
2765 };
2766 if status != cublasStatus_t::CUBLAS_STATUS_SUCCESS {
2767 return Err(ArrowSchurGpuFailure::Unavailable);
2768 }
2769 Ok(())
2770 }
2771
2772 fn accumulate_schur_reduce_only(
2782 blas: &CudaBlas,
2783 d: usize,
2784 k: usize,
2785 n: usize,
2786 y_stack: &CudaSlice<f64>,
2787 schur: &mut CudaSlice<f64>,
2788 ) -> Result<(), ArrowSchurGpuFailure> {
2789 let y_block_elems = d * k;
2790 for i in 0..n {
2791 let y_slice = y_stack.slice(i * y_block_elems..(i + 1) * y_block_elems);
2792 let gemm_cfg = GemmConfig::<f64> {
2793 transa: cublasOperation_t::CUBLAS_OP_T,
2794 transb: cublasOperation_t::CUBLAS_OP_N,
2795 m: to_i32(k).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2796 n: to_i32(k).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2797 k: to_i32(d).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2798 alpha: -1.0,
2799 lda: to_i32(d).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2800 ldb: to_i32(d).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2801 beta: 1.0,
2802 ldc: to_i32(k).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2803 };
2804 unsafe { blas.gemm(gemm_cfg, &y_slice, &y_slice, schur) }
2806 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2807 }
2808 Ok(())
2809 }
2810
2811 fn accumulate_schur_rhs_only(
2815 blas: &CudaBlas,
2816 d: usize,
2817 k: usize,
2818 n: usize,
2819 y_stack: &CudaSlice<f64>,
2820 u_stack: &CudaSlice<f64>,
2821 rhs: &mut CudaSlice<f64>,
2822 ) -> Result<(), ArrowSchurGpuFailure> {
2823 let y_block_elems = d * k;
2824 let u_block_elems = d;
2825 for i in 0..n {
2826 let y_slice = y_stack.slice(i * y_block_elems..(i + 1) * y_block_elems);
2827 let u_slice = u_stack.slice(i * u_block_elems..(i + 1) * u_block_elems);
2828 let gemv_cfg = GemvConfig::<f64> {
2829 trans: cublasOperation_t::CUBLAS_OP_T,
2830 m: to_i32(d).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2831 n: to_i32(k).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2832 alpha: 1.0,
2833 lda: to_i32(d).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2834 incx: 1,
2835 beta: 1.0,
2836 incy: 1,
2837 };
2838 unsafe { blas.gemv(gemv_cfg, &y_slice, &u_slice, rhs) }
2841 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2842 }
2843 Ok(())
2844 }
2845
2846 fn accumulate_schur(
2850 blas: &CudaBlas,
2851 d: usize,
2852 k: usize,
2853 n: usize,
2854 y_stack: &CudaSlice<f64>,
2855 u_stack: &CudaSlice<f64>,
2856 schur: &mut CudaSlice<f64>,
2857 rhs: &mut CudaSlice<f64>,
2858 ) -> Result<(), ArrowSchurGpuFailure> {
2859 let y_block_elems = d * k;
2860 let u_block_elems = d;
2861 for i in 0..n {
2862 let y_slice = y_stack.slice(i * y_block_elems..(i + 1) * y_block_elems);
2863 let u_slice = u_stack.slice(i * u_block_elems..(i + 1) * u_block_elems);
2864 let gemm_cfg = GemmConfig::<f64> {
2866 transa: cublasOperation_t::CUBLAS_OP_T,
2867 transb: cublasOperation_t::CUBLAS_OP_N,
2868 m: to_i32(k).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2869 n: to_i32(k).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2870 k: to_i32(d).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2871 alpha: -1.0,
2872 lda: to_i32(d).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2873 ldb: to_i32(d).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2874 beta: 1.0,
2875 ldc: to_i32(k).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2876 };
2877 unsafe { blas.gemm(gemm_cfg, &y_slice, &y_slice, schur) }
2879 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2880 let gemv_cfg = GemvConfig::<f64> {
2882 trans: cublasOperation_t::CUBLAS_OP_T,
2883 m: to_i32(d).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2884 n: to_i32(k).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2885 alpha: 1.0,
2886 lda: to_i32(d).ok_or(ArrowSchurGpuFailure::Unavailable)?,
2887 incx: 1,
2888 beta: 1.0,
2889 incy: 1,
2890 };
2891 unsafe { blas.gemv(gemv_cfg, &y_slice, &u_slice, rhs) }
2894 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
2895 }
2896 Ok(())
2897 }
2898
2899 fn schur_gemm_stacked(
2915 blas: &CudaBlas,
2916 rows: usize,
2917 k: usize,
2918 y_stacked: &CudaSlice<f64>,
2919 schur: &mut CudaSlice<f64>,
2920 ) -> Result<(), ArrowSchurGpuFailure> {
2921 let k_i = to_i32(k).ok_or(ArrowSchurGpuFailure::Unavailable)?;
2922 let rows_i = to_i32(rows).ok_or(ArrowSchurGpuFailure::Unavailable)?;
2923 let cfg = GemmConfig::<f64> {
2924 transa: cublasOperation_t::CUBLAS_OP_T,
2925 transb: cublasOperation_t::CUBLAS_OP_N,
2926 m: k_i,
2927 n: k_i,
2928 k: rows_i,
2929 alpha: -1.0,
2930 lda: rows_i,
2931 ldb: rows_i,
2932 beta: 1.0,
2933 ldc: k_i,
2934 };
2935 unsafe { blas.gemm(cfg, y_stacked, y_stacked, schur) }
2939 .map_err(|_| ArrowSchurGpuFailure::Unavailable)
2940 }
2941
2942 fn schur_rhs_stacked(
2947 blas: &CudaBlas,
2948 rows: usize,
2949 k: usize,
2950 y_stacked: &CudaSlice<f64>,
2951 u_stacked: &CudaSlice<f64>,
2952 rhs: &mut CudaSlice<f64>,
2953 ) -> Result<(), ArrowSchurGpuFailure> {
2954 let k_i = to_i32(k).ok_or(ArrowSchurGpuFailure::Unavailable)?;
2955 let rows_i = to_i32(rows).ok_or(ArrowSchurGpuFailure::Unavailable)?;
2956 let cfg = GemvConfig::<f64> {
2957 trans: cublasOperation_t::CUBLAS_OP_T,
2958 m: rows_i,
2959 n: k_i,
2960 alpha: 1.0,
2961 lda: rows_i,
2962 incx: 1,
2963 beta: 1.0,
2964 incy: 1,
2965 };
2966 unsafe { blas.gemv(cfg, y_stacked, u_stacked, rhs) }
2970 .map_err(|_| ArrowSchurGpuFailure::Unavailable)
2971 }
2972
2973 fn back_sub_rhs_stacked(
2977 blas: &CudaBlas,
2978 rows: usize,
2979 k: usize,
2980 y_stacked: &CudaSlice<f64>,
2981 delta_beta: &CudaSlice<f64>,
2982 u_stacked: &mut CudaSlice<f64>,
2983 ) -> Result<(), ArrowSchurGpuFailure> {
2984 let k_i = to_i32(k).ok_or(ArrowSchurGpuFailure::Unavailable)?;
2985 let rows_i = to_i32(rows).ok_or(ArrowSchurGpuFailure::Unavailable)?;
2986 let cfg = GemvConfig::<f64> {
2987 trans: cublasOperation_t::CUBLAS_OP_N,
2988 m: rows_i,
2989 n: k_i,
2990 alpha: 1.0,
2991 lda: rows_i,
2992 incx: 1,
2993 beta: 1.0,
2994 incy: 1,
2995 };
2996 unsafe { blas.gemv(cfg, y_stacked, delta_beta, u_stacked) }
3000 .map_err(|_| ArrowSchurGpuFailure::Unavailable)
3001 }
3002
3003 fn accumulate_back_sub_rhs(
3006 blas: &CudaBlas,
3007 d: usize,
3008 k: usize,
3009 n: usize,
3010 y_stack: &CudaSlice<f64>,
3011 delta_beta: &CudaSlice<f64>,
3012 u_stack: &mut CudaSlice<f64>,
3013 ) -> Result<(), ArrowSchurGpuFailure> {
3014 let y_block_elems = d * k;
3015 let u_block_elems = d;
3016 for i in 0..n {
3017 let y_slice = y_stack.slice(i * y_block_elems..(i + 1) * y_block_elems);
3018 let mut u_slice = u_stack.slice_mut(i * u_block_elems..(i + 1) * u_block_elems);
3019 let gemv_cfg = GemvConfig::<f64> {
3020 trans: cublasOperation_t::CUBLAS_OP_N,
3021 m: to_i32(d).ok_or(ArrowSchurGpuFailure::Unavailable)?,
3022 n: to_i32(k).ok_or(ArrowSchurGpuFailure::Unavailable)?,
3023 alpha: 1.0,
3024 lda: to_i32(d).ok_or(ArrowSchurGpuFailure::Unavailable)?,
3025 incx: 1,
3026 beta: 1.0,
3027 incy: 1,
3028 };
3029 unsafe { blas.gemv(gemv_cfg, &y_slice, delta_beta, &mut u_slice) }
3032 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
3033 }
3034 Ok(())
3035 }
3036
3037 use std::collections::HashMap;
3053 use std::sync::Mutex;
3054
3055 struct FusedModuleCache {
3060 modules: Mutex<
3061 HashMap<crate::gpu_kernels::arrow_schur_nvrtc::FusedModuleCacheKey, Arc<CudaModule>>,
3062 >,
3063 }
3064
3065 fn fused_module_cache() -> &'static FusedModuleCache {
3066 static CACHE: OnceLock<FusedModuleCache> = OnceLock::new();
3067 CACHE.get_or_init(|| FusedModuleCache {
3068 modules: Mutex::new(HashMap::new()),
3069 })
3070 }
3071
3072 fn fused_module_for(
3073 ctx: &Arc<CudaContext>,
3074 key: crate::gpu_kernels::arrow_schur_nvrtc::FusedModuleCacheKey,
3075 ) -> Result<Arc<CudaModule>, ArrowSchurGpuFailure> {
3076 let cache = fused_module_cache();
3077 if let Ok(guard) = cache.modules.lock() {
3078 if let Some(existing) = guard.get(&key) {
3079 return Ok(existing.clone());
3080 }
3081 }
3082 let src = crate::gpu_kernels::arrow_schur_nvrtc::forward_kernel_source(
3083 key.p_max as usize,
3084 key.r_template as usize,
3085 );
3086 let ptx = gam_gpu::device_cache::compile_ptx_arch(&src).map_err(|err| {
3087 ArrowSchurGpuFailure::SchurFactorFailed {
3088 reason: format!(
3089 "arrow-schur fused NVRTC compile (p_max={}, r={}): {err}",
3090 key.p_max, key.r_template
3091 ),
3092 }
3093 })?;
3094 let module = ctx
3095 .load_module(ptx)
3096 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
3097 if let Ok(mut guard) = cache.modules.lock() {
3098 guard.entry(key).or_insert_with(|| module.clone());
3099 }
3100 Ok(module)
3101 }
3102
3103 const PCG_VECTOR_KERNEL_SOURCE: &str = r#"
3104extern "C" __global__ void arrow_pcg_jacobi_mul(
3105 const double* __restrict__ inv_diag,
3106 const double* __restrict__ r,
3107 double* __restrict__ z,
3108 int n
3109) {
3110 int idx = blockIdx.x * blockDim.x + threadIdx.x;
3111 if (idx < n) {
3112 z[idx] = inv_diag[idx] * r[idx];
3113 }
3114}
3115
3116extern "C" __global__ void arrow_pcg_update_p(
3117 const double* __restrict__ z,
3118 double beta,
3119 double* __restrict__ p,
3120 int n
3121) {
3122 int idx = blockIdx.x * blockDim.x + threadIdx.x;
3123 if (idx < n) {
3124 p[idx] = z[idx] + beta * p[idx];
3125 }
3126}
3127
3128extern "C" __global__ void arrow_sae_init(
3129 double* __restrict__ out,
3130 const double* __restrict__ x,
3131 double ridge,
3132 int n
3133) {
3134 int idx = blockIdx.x * blockDim.x + threadIdx.x;
3135 if (idx < n) {
3136 out[idx] = ridge * x[idx];
3137 }
3138}
3139
3140extern "C" __global__ void arrow_sae_smooth_matvec(
3141 const double* __restrict__ x,
3142 double* __restrict__ out,
3143 const int* __restrict__ block_offsets,
3144 const int* __restrict__ block_m,
3145 const int* __restrict__ factor_ptr,
3146 const double* __restrict__ factors,
3147 int p,
3148 int n_blocks
3149) {
3150 int block_id = blockIdx.y;
3151 int linear = blockIdx.x * blockDim.x + threadIdx.x;
3152 if (block_id >= n_blocks) {
3153 return;
3154 }
3155 int m = block_m[block_id];
3156 int total = m * p;
3157 if (linear >= total) {
3158 return;
3159 }
3160 int li = linear / p;
3161 int oc = linear - li * p;
3162 int off = block_offsets[block_id];
3163 int fbase = factor_ptr[block_id];
3164 double acc = 0.0;
3165 for (int lj = 0; lj < m; ++lj) {
3166 double a = factors[fbase + li * m + lj];
3167 acc += a * x[off + lj * p + oc];
3168 }
3169 out[off + li * p + oc] += acc;
3170}
3171
3172extern "C" __global__ void arrow_sae_sparse_g_matvec(
3173 const double* __restrict__ x,
3174 double* __restrict__ out,
3175 const int* __restrict__ col_off,
3176 const int* __restrict__ cols,
3177 const int* __restrict__ data_ptr,
3178 const double* __restrict__ data,
3179 const int* __restrict__ group_ptr,
3180 const int* __restrict__ group_blocks,
3181 const int* __restrict__ group_off,
3182 const int* __restrict__ group_width,
3183 int p,
3184 int n_groups
3185) {
3186 int group = blockIdx.y;
3187 int linear = blockIdx.x * blockDim.x + threadIdx.x;
3188 if (group >= n_groups || linear >= group_width[group]) {
3189 return;
3190 }
3191 int out_index = group_off[group] + linear;
3192 int li = linear / p;
3193 int oc = linear - li * p;
3194 double acc = out[out_index];
3195 // #2535 — one thread OWNS this scalar output and walks every contributing
3196 // G block in original host order. Equal output spans were grouped into the
3197 // CSR above; distinct spans are disjoint. The scheduler therefore chooses
3198 // only when an independent output runs, never the order of its additions.
3199 for (int position = group_ptr[group]; position < group_ptr[group + 1]; ++position) {
3200 int block_id = group_blocks[position];
3201 int m_j = cols[block_id];
3202 int cbase = col_off[block_id];
3203 int dbase = data_ptr[block_id];
3204 for (int lj = 0; lj < m_j; ++lj) {
3205 acc += data[dbase + li * m_j + lj] * x[(cbase + lj) * p + oc];
3206 }
3207 }
3208 out[out_index] = acc;
3209}
3210
3211extern "C" __global__ void arrow_sae_gather_u(
3212 const double* __restrict__ x,
3213 const int* __restrict__ row_ptr,
3214 const int* __restrict__ beta_base,
3215 const double* __restrict__ phi,
3216 double* __restrict__ u,
3217 int p,
3218 int n_rows
3219) {
3220 int row = blockIdx.y;
3221 int oc = blockIdx.x * blockDim.x + threadIdx.x;
3222 if (row >= n_rows || oc >= p) {
3223 return;
3224 }
3225 double acc = 0.0;
3226 int start = row_ptr[row];
3227 int end = row_ptr[row + 1];
3228 for (int e = start; e < end; ++e) {
3229 acc += phi[e] * x[beta_base[e] + oc];
3230 }
3231 u[row * p + oc] = acc;
3232}
3233
3234extern "C" __global__ void arrow_sae_apply_l(
3235 const double* __restrict__ u,
3236 const int* __restrict__ jac_ptr,
3237 const double* __restrict__ jac,
3238 double* __restrict__ w,
3239 int p,
3240 int max_q,
3241 int n_rows
3242) {
3243 int row = blockIdx.y;
3244 int c = blockIdx.x * blockDim.x + threadIdx.x;
3245 if (row >= n_rows) {
3246 return;
3247 }
3248 int jstart = jac_ptr[row];
3249 int q = (jac_ptr[row + 1] - jstart) / p;
3250 if (c >= q) {
3251 return;
3252 }
3253 double acc = 0.0;
3254 for (int oc = 0; oc < p; ++oc) {
3255 acc += jac[jstart + c * p + oc] * u[row * p + oc];
3256 }
3257 w[row * max_q + c] = acc;
3258}
3259
3260extern "C" __global__ void arrow_sae_apply_ainv(
3261 const double* __restrict__ ainv,
3262 const double* __restrict__ w,
3263 double* __restrict__ v,
3264 int max_q,
3265 int n_rows
3266) {
3267 int row = blockIdx.y;
3268 int c = blockIdx.x * blockDim.x + threadIdx.x;
3269 if (row >= n_rows || c >= max_q) {
3270 return;
3271 }
3272 double acc = 0.0;
3273 int base = row * max_q * max_q;
3274 for (int j = 0; j < max_q; ++j) {
3275 acc += ainv[base + c * max_q + j] * w[row * max_q + j];
3276 }
3277 v[row * max_q + c] = acc;
3278}
3279
3280/* #2512 DETERMINISTIC legacy row-Schur scatter, STAGE 1 (partials).
3281
3282 `arrow_sae_scatter_sub` above ran one block per row and folded the rows into
3283 `out` with `atomicAdd`, so the summation order was the block scheduler's and
3284 the matvec returned a different value on every launch. Here thread `(oc,
3285 chunk)` owns output channel `oc` over the contiguous row range
3286 [chunk·rows_per_chunk, …) and walks its rows in fixed index order, so the
3287 reassociation is a property of the shape rather than of the schedule.
3288
3289 The thread may write `partial[cbase + beta_base[e] + oc]` for several `e`
3290 without colliding with any other thread because every `beta_base` is a
3291 multiple of `p`: the host builds the support as
3292 `a_phi.push((atom_beta_off + basis_col * p, w))` and each atom's border block
3293 is `M_k · p` wide, so the entries a thread touches are exactly the indices
3294 congruent to `oc` modulo `p`. That is also why the thread can zero its own
3295 progression first and needs no separate clearing pass or barrier. */
3296extern "C" __global__ void arrow_sae_scatter_sub_det_partial(
3297 const double* __restrict__ v,
3298 const int* __restrict__ jac_ptr,
3299 const double* __restrict__ jac,
3300 const int* __restrict__ row_ptr,
3301 const int* __restrict__ beta_base,
3302 const double* __restrict__ phi,
3303 double* __restrict__ partial,
3304 int p,
3305 int k,
3306 int max_q,
3307 int n_rows,
3308 int rows_per_chunk
3309) {
3310 int oc = blockIdx.x * blockDim.x + threadIdx.x;
3311 if (oc >= p) { return; }
3312 int chunk = blockIdx.y;
3313 long long cbase = (long long)chunk * k;
3314 for (int i = oc; i < k; i += p) {
3315 partial[cbase + i] = 0.0;
3316 }
3317 int row0 = chunk * rows_per_chunk;
3318 int row1 = row0 + rows_per_chunk;
3319 if (row1 > n_rows) { row1 = n_rows; }
3320 for (int row = row0; row < row1; ++row) {
3321 int jstart = jac_ptr[row];
3322 int q = (jac_ptr[row + 1] - jstart) / p;
3323 double lt_v = 0.0;
3324 for (int c = 0; c < q; ++c) {
3325 lt_v += jac[jstart + c * p + oc] * v[row * max_q + c];
3326 }
3327 int start = row_ptr[row];
3328 int end = row_ptr[row + 1];
3329 for (int e = start; e < end; ++e) {
3330 partial[cbase + beta_base[e] + oc] -= phi[e] * lt_v;
3331 }
3332 }
3333}
3334
3335/* #2512 STAGE 2 (reduce): out[a] += Σ_chunk partial[chunk][a], chunks summed in
3336 fixed order 0..n_chunks. One thread per output coord `a`, so no two threads
3337 touch the same `out[a]`. */
3338extern "C" __global__ void arrow_sae_scatter_sub_det_reduce(
3339 const double* __restrict__ partial,
3340 double* __restrict__ out,
3341 int k,
3342 int n_chunks
3343) {
3344 int a = blockIdx.x * blockDim.x + threadIdx.x;
3345 if (a >= k) { return; }
3346 double acc = 0.0;
3347 for (int chunk = 0; chunk < n_chunks; ++chunk) {
3348 acc += partial[(long long)chunk * k + a];
3349 }
3350 out[a] += acc;
3351}
3352
3353/* #2512 DETERMINISTIC Jacobi-diagonal subtraction, STAGE 1 (partials).
3354
3355 The former kernel ran one block per row and folded the rows into `diag` with
3356 `atomicAdd`, so the PRECONDITIONER was a fresh draw per call — and a moved
3357 preconditioner moves the whole CG trajectory, not merely the last bits of one
3358 matvec. Thread `(oc, chunk)` owns output channel `oc` over a contiguous row
3359 range and walks it in fixed index order; as in the scatter, every `beta_base`
3360 is a multiple of `p`, so the indices a thread touches are exactly those
3361 congruent to `oc` modulo `p` and it can clear its own progression first.
3362 Stage 2 is the shared `arrow_sae_scatter_sub_det_reduce`. */
3363extern "C" __global__ void arrow_sae_diag_sub_det_partial(
3364 const double* __restrict__ ainv,
3365 const int* __restrict__ jac_ptr,
3366 const double* __restrict__ jac,
3367 const int* __restrict__ row_ptr,
3368 const int* __restrict__ beta_base,
3369 const double* __restrict__ phi,
3370 double* __restrict__ partial,
3371 int p,
3372 int k,
3373 int max_q,
3374 int n_rows,
3375 int rows_per_chunk
3376) {
3377 int oc = blockIdx.x * blockDim.x + threadIdx.x;
3378 if (oc >= p) { return; }
3379 int chunk = blockIdx.y;
3380 long long cbase = (long long)chunk * k;
3381 for (int i = oc; i < k; i += p) {
3382 partial[cbase + i] = 0.0;
3383 }
3384 int row0 = chunk * rows_per_chunk;
3385 int row1 = row0 + rows_per_chunk;
3386 if (row1 > n_rows) { row1 = n_rows; }
3387 for (int row = row0; row < row1; ++row) {
3388 int jstart = jac_ptr[row];
3389 int q = (jac_ptr[row + 1] - jstart) / p;
3390 int abase = row * max_q * max_q;
3391 double quad = 0.0;
3392 for (int c = 0; c < q; ++c) {
3393 double lc = jac[jstart + c * p + oc];
3394 for (int d = 0; d < q; ++d) {
3395 quad += lc * ainv[abase + c * max_q + d] * jac[jstart + d * p + oc];
3396 }
3397 }
3398 int start = row_ptr[row];
3399 int end = row_ptr[row + 1];
3400 for (int e = start; e < end; ++e) {
3401 double pe = phi[e];
3402 partial[cbase + beta_base[e] + oc] -= (pe * pe) * quad;
3403 }
3404 }
3405}
3406
3407/* ── #1017/#1026 frames-engaged device kernels ─────────────────────────────
3408 * The factored β border is C-space (width Σ M_k·r_k). The penalty side is the
3409 * smooth `λ S_k ⊗ I_{r_k}` (per-block right-width r_k) plus the data-fit
3410 * `G_{ij} ⊗ W_{ij}` (W = U_iᵀU_j, dense r_i×r_j). The reduced-Schur term uses
3411 * the per-row DENSE cross-block H_tβ^(i) (q_i × border_dim, row-major). */
3412
3413extern "C" __global__ void arrow_sae_frame_smooth_matvec(
3414 const double* __restrict__ x,
3415 double* __restrict__ out,
3416 const int* __restrict__ block_offsets,
3417 const int* __restrict__ block_m,
3418 const int* __restrict__ block_r,
3419 const int* __restrict__ factor_ptr,
3420 const double* __restrict__ factors,
3421 int n_blocks
3422) {
3423 int block_id = blockIdx.y;
3424 int linear = blockIdx.x * blockDim.x + threadIdx.x;
3425 if (block_id >= n_blocks) {
3426 return;
3427 }
3428 int m = block_m[block_id];
3429 int r = block_r[block_id];
3430 int total = m * r;
3431 if (linear >= total) {
3432 return;
3433 }
3434 int li = linear / r;
3435 int ib = linear - li * r;
3436 int off = block_offsets[block_id];
3437 int fbase = factor_ptr[block_id];
3438 double acc = 0.0;
3439 for (int lj = 0; lj < m; ++lj) {
3440 double a = factors[fbase + li * m + lj];
3441 acc += a * x[off + lj * r + ib];
3442 }
3443 out[off + li * r + ib] += acc;
3444}
3445
3446extern "C" __global__ void arrow_sae_frame_g_matvec(
3447 const double* __restrict__ x,
3448 double* __restrict__ out,
3449 const int* __restrict__ off_i,
3450 const int* __restrict__ off_j,
3451 const int* __restrict__ r_i,
3452 const int* __restrict__ r_j,
3453 const int* __restrict__ m_j,
3454 const int* __restrict__ g_ptr,
3455 const double* __restrict__ g_data,
3456 const int* __restrict__ w_ptr,
3457 const double* __restrict__ w_data,
3458 const int* __restrict__ group_ptr,
3459 const int* __restrict__ group_blocks,
3460 const int* __restrict__ group_off,
3461 const int* __restrict__ group_width,
3462 int n_groups
3463) {
3464 int group = blockIdx.y;
3465 int linear = blockIdx.x * blockDim.x + threadIdx.x;
3466 if (group >= n_groups || linear >= group_width[group]) {
3467 return;
3468 }
3469 int out_index = group_off[group] + linear;
3470 double acc = out[out_index];
3471 // #2535 — fixed-order G⊗W accumulation. The owning output thread walks the
3472 // original block subsequence and, within each block, the original lj/b
3473 // order used by the CPU oracle. No partial buffer and no host reduction.
3474 for (int position = group_ptr[group]; position < group_ptr[group + 1]; ++position) {
3475 int block_id = group_blocks[position];
3476 int ri = r_i[block_id];
3477 int rj = r_j[block_id];
3478 int mj = m_j[block_id];
3479 int local = out_index - off_i[block_id];
3480 int li = local / ri;
3481 int a = local - li * ri;
3482 int oj = off_j[block_id];
3483 int gbase = g_ptr[block_id];
3484 int wbase = w_ptr[block_id];
3485 for (int lj = 0; lj < mj; ++lj) {
3486 double g = g_data[gbase + li * mj + lj];
3487 if (g == 0.0) { continue; }
3488 int xj_base = oj + lj * rj;
3489 double inner = 0.0;
3490 for (int b = 0; b < rj; ++b) {
3491 inner += w_data[wbase + a * rj + b] * x[xj_base + b];
3492 }
3493 acc += g * inner;
3494 }
3495 }
3496 out[out_index] = acc;
3497}
3498
3499/* Per-row reduced-Schur subtraction with a DENSE cross-block H_tβ^(i).
3500 * h_i = H_tβ^(i) · x (length q_i)
3501 * s_i = (H_tt^(i)+ρ_t I)⁻¹ h_i (apply cached ainv, length q_i)
3502 * out -= (H_tβ^(i))ᵀ · s_i (scatter into border_dim)
3503 * `htb` is row-major (q_i × k) flattened, `htb_ptr` gives each row's base and
3504 * (htb_ptr[row+1]-htb_ptr[row])/k == q_i. `q_of` carries q_i directly. */
3505extern "C" __global__ void arrow_sae_frame_apply_h(
3506 const double* __restrict__ x,
3507 const int* __restrict__ htb_ptr,
3508 const double* __restrict__ htb,
3509 const int* __restrict__ q_of,
3510 double* __restrict__ hvec,
3511 int k,
3512 int max_q,
3513 int n_rows
3514) {
3515 int row = blockIdx.y;
3516 int c = blockIdx.x * blockDim.x + threadIdx.x;
3517 if (row >= n_rows) { return; }
3518 int q = q_of[row];
3519 if (c >= q) { return; }
3520 int base = htb_ptr[row] + c * k;
3521 double acc = 0.0;
3522 for (int a = 0; a < k; ++a) {
3523 acc += htb[base + a] * x[a];
3524 }
3525 hvec[row * max_q + c] = acc;
3526}
3527
3528extern "C" __global__ void arrow_sae_frame_apply_ainv(
3529 const double* __restrict__ ainv,
3530 const double* __restrict__ hvec,
3531 const int* __restrict__ q_of,
3532 double* __restrict__ svec,
3533 int max_q,
3534 int n_rows
3535) {
3536 int row = blockIdx.y;
3537 int c = blockIdx.x * blockDim.x + threadIdx.x;
3538 if (row >= n_rows || c >= max_q) { return; }
3539 int q = q_of[row];
3540 double acc = 0.0;
3541 int abase = row * max_q * max_q;
3542 for (int j = 0; j < q; ++j) {
3543 acc += ainv[abase + c * max_q + j] * hvec[row * max_q + j];
3544 }
3545 svec[row * max_q + c] = acc;
3546}
3547
3548/* #1017 evidence-lane DETERMINISTIC reduced-Schur scatter:
3549 out[a] = -Σ_i Σ_c H_tβ[i][c,a]·svec[i,c]. One thread owns output coord `a` and
3550 sums the rows in fixed index order 0..n_rows — NO atomics, so the result is
3551 run-to-run bit-stable (the SLQ log|S| determinism contract): the reduction
3552 order over rows is pinned rather than decided by the block scheduler.
3553 `out` is fully assigned (no init needed). */
3554extern "C" __global__ void arrow_sae_frame_scatter_h_det(
3555 const double* __restrict__ svec,
3556 const int* __restrict__ htb_ptr,
3557 const double* __restrict__ htb,
3558 const int* __restrict__ q_of,
3559 double* __restrict__ out,
3560 int k,
3561 int max_q,
3562 int n_rows
3563) {
3564 int a = blockIdx.x * blockDim.x + threadIdx.x;
3565 if (a >= k) { return; }
3566 double acc = 0.0;
3567 for (int row = 0; row < n_rows; ++row) {
3568 int q = q_of[row];
3569 int hbase = htb_ptr[row];
3570 int sbase = row * max_q;
3571 for (int c = 0; c < q; ++c) {
3572 acc += htb[hbase + c * k + a] * svec[sbase + c];
3573 }
3574 }
3575 out[a] = -acc;
3576}
3577
3578/* #1017 evidence-lane 2-STAGE deterministic scatter, STAGE 1 (partials):
3579 partial[chunk][a] = Σ_{row∈chunk} Σ_c H_tβ[row][c,a]·svec[row,c], for the
3580 contiguous row range [chunk·rows_per_chunk, …). grid = (⌈k/256⌉, n_chunks);
3581 thread owns (a, chunk). Replaces the single-strip `arrow_sae_frame_scatter_h_det`
3582 (⌈k/256⌉ CTAs — only 4 at k=911, one thread serial over ALL n_rows, ~94% of a
3583 72-SM A10 idle) with ⌈k/256⌉·n_chunks CTAs. Rows are summed in fixed order
3584 within the chunk and the chunks are reduced in order by stage 2, so the result
3585 is a FIXED reassociation of the same ordered row sum — run-to-run bit-stable
3586 (the SLQ log|S| determinism contract) and within the ≤1e-9 CPU-oracle gate. */
3587extern "C" __global__ void arrow_sae_frame_scatter_h_det_partial(
3588 const double* __restrict__ svec,
3589 const int* __restrict__ htb_ptr,
3590 const double* __restrict__ htb,
3591 const int* __restrict__ q_of,
3592 double* __restrict__ partial,
3593 int k,
3594 int max_q,
3595 int n_rows,
3596 int rows_per_chunk
3597) {
3598 int a = blockIdx.x * blockDim.x + threadIdx.x;
3599 if (a >= k) { return; }
3600 int chunk = blockIdx.y;
3601 int row0 = chunk * rows_per_chunk;
3602 int row1 = row0 + rows_per_chunk;
3603 if (row1 > n_rows) { row1 = n_rows; }
3604 double acc = 0.0;
3605 for (int row = row0; row < row1; ++row) {
3606 int q = q_of[row];
3607 int hbase = htb_ptr[row];
3608 int sbase = row * max_q;
3609 for (int c = 0; c < q; ++c) {
3610 acc += htb[hbase + c * k + a] * svec[sbase + c];
3611 }
3612 }
3613 partial[(long long)chunk * k + a] = acc;
3614}
3615
3616/* #2512 STAGE 2, ACCUMULATING form: out[a] -= Σ_chunk partial[chunk][a], chunks
3617 summed in fixed order 0..n_chunks. Identical reduction to
3618 `arrow_sae_frame_scatter_h_det_reduce`, but subtracts into an `out` that the
3619 penalty and data-Gram matvecs have already written, which is what the step-PCG
3620 matvec needs. One thread per output coord `a`, so no two threads touch the same
3621 `out[a]` and the result is run-to-run bit-stable. */
3622extern "C" __global__ void arrow_sae_frame_scatter_h_det_reduce_sub(
3623 const double* __restrict__ partial,
3624 double* __restrict__ out,
3625 int k,
3626 int n_chunks
3627) {
3628 int a = blockIdx.x * blockDim.x + threadIdx.x;
3629 if (a >= k) { return; }
3630 double acc = 0.0;
3631 for (int chunk = 0; chunk < n_chunks; ++chunk) {
3632 acc += partial[(long long)chunk * k + a];
3633 }
3634 out[a] -= acc;
3635}
3636
3637/* #1017 STAGE 2 (reduce): out[a] = -Σ_chunk partial[chunk][a], chunks summed in
3638 fixed order 0..n_chunks. One thread per output coord `a`; ⌈k/256⌉ CTAs. */
3639extern "C" __global__ void arrow_sae_frame_scatter_h_det_reduce(
3640 const double* __restrict__ partial,
3641 double* __restrict__ out,
3642 int k,
3643 int n_chunks
3644) {
3645 int a = blockIdx.x * blockDim.x + threadIdx.x;
3646 if (a >= k) { return; }
3647 double acc = 0.0;
3648 for (int chunk = 0; chunk < n_chunks; ++chunk) {
3649 acc += partial[(long long)chunk * k + a];
3650 }
3651 out[a] = -acc;
3652}
3653
3654/* #1017 evidence-lane WARP-COOPERATIVE apply_h: hvec[i][c] = Σ_a H_tβ[i][c,a]·x[a].
3655 One WARP owns (row, c): lane `l` strides `a = l, l+32, …` over the contiguous
3656 `H_tβ[i][c,·]` slab (fully coalesced across the warp) and a fixed-order
3657 __shfl_down tree reduces to lane 0. Replaces the shared `arrow_sae_frame_apply_h`
3658 (256-thread block, only `q_i` threads active, stride-`k` uncoalesced reads) on
3659 the evidence path ONLY — the shared kernel is untouched (step-PCG relies on it).
3660 Reduction order is fixed (lane stride + tree), so the result is run-to-run
3661 bit-stable; it differs from the scalar kernel only by ULP reassociation, within
3662 the ≤1e-9 parity gate. Launch: block = max_q·32 (≤1024), grid.x = n_rows;
3663 warp `w = threadIdx.x/32` handles `c = w` for `w < q_i`. */
3664extern "C" __global__ void arrow_sae_frame_apply_h_warp(
3665 const double* __restrict__ x,
3666 const int* __restrict__ htb_ptr,
3667 const double* __restrict__ htb,
3668 const int* __restrict__ q_of,
3669 double* __restrict__ hvec,
3670 int k,
3671 int max_q,
3672 int n_rows
3673) {
3674 int row = blockIdx.x;
3675 if (row >= n_rows) { return; }
3676 int warp = threadIdx.x >> 5;
3677 int lane = threadIdx.x & 31;
3678 int q = q_of[row];
3679 if (warp >= q) { return; }
3680 int base = htb_ptr[row] + warp * k;
3681 double acc = 0.0;
3682 for (int a = lane; a < k; a += 32) {
3683 acc += htb[base + a] * x[a];
3684 }
3685 #pragma unroll
3686 for (int off = 16; off > 0; off >>= 1) {
3687 acc += __shfl_down_sync(0xffffffffu, acc, off);
3688 }
3689 if (lane == 0) {
3690 hvec[row * max_q + warp] = acc;
3691 }
3692}
3693
3694/* Frame Jacobi diagonal subtraction:
3695 * diag[a] -= Σ_row Σ_c Σ_d H_tβ[row,c,a]·ainv[row,c,d]·H_tβ[row,d,a].
3696 *
3697 * #2535 — one thread owns each diagonal scalar and walks rows in host order.
3698 * The former (row,a) grid raced through atomicAdd(&diag[a], -quad), making the
3699 * preconditioner depend on CTA scheduling even after the two G writers were
3700 * made deterministic. */
3701extern "C" __global__ void arrow_sae_frame_diag_sub(
3702 double* __restrict__ diag,
3703 const double* __restrict__ ainv,
3704 const int* __restrict__ htb_ptr,
3705 const double* __restrict__ htb,
3706 const int* __restrict__ q_of,
3707 int k,
3708 int max_q,
3709 int n_rows
3710) {
3711 int a = blockIdx.x * blockDim.x + threadIdx.x;
3712 if (a >= k) { return; }
3713 double acc = diag[a];
3714 for (int row = 0; row < n_rows; ++row) {
3715 int q = q_of[row];
3716 int hbase = htb_ptr[row];
3717 int abase = row * max_q * max_q;
3718 double quad = 0.0;
3719 for (int c = 0; c < q; ++c) {
3720 double hc = htb[hbase + c * k + a];
3721 for (int d = 0; d < q; ++d) {
3722 quad += hc * ainv[abase + c * max_q + d] * htb[hbase + d * k + a];
3723 }
3724 }
3725 acc -= quad;
3726 }
3727 diag[a] = acc;
3728}
3729"#;
3730
3731 fn pcg_vector_module(
3732 ctx: &Arc<CudaContext>,
3733 ) -> Result<&'static Arc<CudaModule>, ArrowSchurGpuFailure> {
3734 static CACHE: gam_gpu::device_cache::PtxModuleCache =
3735 gam_gpu::device_cache::PtxModuleCache::new();
3736 CACHE
3737 .get_or_compile(ctx, "arrow_pcg_vector", PCG_VECTOR_KERNEL_SOURCE)
3738 .map_err(|err| {
3739 log::warn!("[#1551] pcg_vector_module get_or_compile failed: {err}");
3745 ArrowSchurGpuFailure::Unavailable
3746 })
3747 }
3748
3749 fn pcg_launch_config(n: usize) -> Result<LaunchConfig, ArrowSchurGpuFailure> {
3750 let threads = 256u32;
3751 let blocks = ((n as u32).saturating_add(threads - 1) / threads).max(1);
3752 Ok(LaunchConfig {
3753 grid_dim: (blocks, 1, 1),
3754 block_dim: (threads, 1, 1),
3755 shared_mem_bytes: 0,
3756 })
3757 }
3758
3759 fn launch_jacobi_mul(
3760 stream: &Arc<CudaStream>,
3761 module: &Arc<CudaModule>,
3762 inv_diag: &CudaSlice<f64>,
3763 r: &CudaSlice<f64>,
3764 z: &mut CudaSlice<f64>,
3765 n: usize,
3766 ) -> Result<(), ArrowSchurGpuFailure> {
3767 let kernel = module
3768 .load_function("arrow_pcg_jacobi_mul")
3769 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
3770 let n_i32 = to_i32(n).ok_or(ArrowSchurGpuFailure::Unavailable)?;
3771 let mut builder = stream.launch_builder(&kernel);
3772 builder.arg(inv_diag).arg(r).arg(z).arg(&n_i32);
3773 unsafe { builder.launch(pcg_launch_config(n)?) }
3776 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
3777 Ok(())
3778 }
3779
3780 fn launch_update_p(
3781 stream: &Arc<CudaStream>,
3782 module: &Arc<CudaModule>,
3783 z: &CudaSlice<f64>,
3784 beta: f64,
3785 p: &mut CudaSlice<f64>,
3786 n: usize,
3787 ) -> Result<(), ArrowSchurGpuFailure> {
3788 let kernel = module
3789 .load_function("arrow_pcg_update_p")
3790 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
3791 let n_i32 = to_i32(n).ok_or(ArrowSchurGpuFailure::Unavailable)?;
3792 let mut builder = stream.launch_builder(&kernel);
3793 builder.arg(z).arg(&beta).arg(p).arg(&n_i32);
3794 unsafe { builder.launch(pcg_launch_config(n)?) }
3797 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
3798 Ok(())
3799 }
3800
3801 struct DeviceSaePcgBuffers {
3802 row_ptr: CudaSlice<i32>,
3803 beta_base: CudaSlice<i32>,
3804 phi: CudaSlice<f64>,
3805 jac_ptr: CudaSlice<i32>,
3806 jac: CudaSlice<f64>,
3807 smooth_offsets: CudaSlice<i32>,
3808 smooth_m: CudaSlice<i32>,
3809 smooth_ptr: CudaSlice<i32>,
3810 smooth_data: CudaSlice<f64>,
3811 g_col_off: CudaSlice<i32>,
3812 g_cols: CudaSlice<i32>,
3813 g_ptr: CudaSlice<i32>,
3814 g_data: CudaSlice<f64>,
3815 g_group_ptr: CudaSlice<i32>,
3816 g_group_blocks: CudaSlice<i32>,
3817 g_group_off: CudaSlice<i32>,
3818 g_group_width: CudaSlice<i32>,
3819 ainv: CudaSlice<f64>,
3820 u: CudaSlice<f64>,
3821 w: CudaSlice<f64>,
3822 v: CudaSlice<f64>,
3823 scatter_partial: CudaSlice<f64>,
3825 scatter_n_chunks: usize,
3826 scatter_rows_per_chunk: usize,
3827 n_rows: usize,
3828 p: usize,
3829 k: usize,
3830 max_q: usize,
3831 smooth_blocks: usize,
3832 g_groups: usize,
3833 g_max_work: usize,
3834 }
3835
3836 fn checked_i32(value: usize) -> Result<i32, ArrowSchurGpuFailure> {
3837 to_i32(value).ok_or(ArrowSchurGpuFailure::Unavailable)
3838 }
3839
3840 fn sae_penalty_diag_host(
3841 data: &DeviceSaePcgData,
3842 ridge_beta: f64,
3843 ) -> Result<Vec<f64>, ArrowSchurGpuFailure> {
3844 let mut diag = vec![ridge_beta; data.beta_dim];
3845 for block in &data.smooth_blocks {
3846 let (rows, cols) = block.factor_a.dim();
3847 if rows != cols {
3848 return Err(ArrowSchurGpuFailure::Unavailable);
3849 }
3850 for row in 0..rows {
3851 let coeff = block.factor_a[[row, row]];
3852 let base = block
3853 .global_offset
3854 .checked_add(
3855 row.checked_mul(data.p)
3856 .ok_or(ArrowSchurGpuFailure::Unavailable)?,
3857 )
3858 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
3859 let end = base
3860 .checked_add(data.p)
3861 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
3862 if end > diag.len() {
3863 return Err(ArrowSchurGpuFailure::Unavailable);
3864 }
3865 for channel in 0..data.p {
3866 diag[base + channel] += coeff;
3867 }
3868 }
3869 }
3870 for block in &data.sparse_g_blocks {
3871 if block.row_off != block.col_off {
3872 continue;
3873 }
3874 let (rows, cols) = block.data.dim();
3875 for row in 0..rows.min(cols) {
3876 let coeff = block.data[[row, row]];
3877 let beta_row = block
3878 .row_off
3879 .checked_add(row)
3880 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
3881 let base = beta_row
3882 .checked_mul(data.p)
3883 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
3884 let end = base
3885 .checked_add(data.p)
3886 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
3887 if end > diag.len() {
3888 return Err(ArrowSchurGpuFailure::Unavailable);
3889 }
3890 for channel in 0..data.p {
3891 diag[base + channel] += coeff;
3892 }
3893 }
3894 }
3895 Ok(diag)
3896 }
3897
3898 fn flatten_device_sae_data(
3899 sys: &ArrowSchurSystem,
3900 data: &DeviceSaePcgData,
3901 ridge_t: f64,
3902 stream: &Arc<CudaStream>,
3903 ) -> Result<DeviceSaePcgBuffers, ArrowSchurGpuFailure> {
3904 let n_rows = sys.rows.len();
3905 let p = data.p;
3906 let k = data.beta_dim;
3907 if data.a_phi.len() != n_rows || data.local_jac.len() != n_rows {
3908 return Err(ArrowSchurGpuFailure::Unavailable);
3909 }
3910
3911 let (scatter_rows_per_chunk, scatter_n_chunks) = scatter_chunking(n_rows);
3912 let mut row_ptr_host = Vec::with_capacity(n_rows + 1);
3913 let mut beta_base_host = Vec::<i32>::new();
3914 let mut phi_host = Vec::<f64>::new();
3915 row_ptr_host.push(0_i32);
3916 for row in data.a_phi.iter() {
3917 for &(base, phi) in row {
3918 beta_base_host.push(checked_i32(base)?);
3919 phi_host.push(phi);
3920 }
3921 row_ptr_host.push(checked_i32(beta_base_host.len())?);
3922 }
3923
3924 let mut jac_ptr_host = Vec::with_capacity(n_rows + 1);
3925 let mut jac_host = Vec::<f64>::new();
3926 let mut max_q = 0usize;
3927 jac_ptr_host.push(0_i32);
3928 for row_jac in data.local_jac.iter() {
3929 if row_jac.len() % p != 0 {
3930 return Err(ArrowSchurGpuFailure::Unavailable);
3931 }
3932 max_q = max_q.max(row_jac.len() / p);
3933 jac_host.extend_from_slice(row_jac);
3934 jac_ptr_host.push(checked_i32(jac_host.len())?);
3935 }
3936 if max_q == 0 {
3937 return Err(ArrowSchurGpuFailure::Unavailable);
3938 }
3939
3940 let mut smooth_offsets_host = Vec::with_capacity(data.smooth_blocks.len());
3941 let mut smooth_m_host = Vec::with_capacity(data.smooth_blocks.len());
3942 let mut smooth_ptr_host = Vec::with_capacity(data.smooth_blocks.len() + 1);
3943 let mut smooth_data_host = Vec::<f64>::new();
3944 smooth_ptr_host.push(0_i32);
3945 for block in &data.smooth_blocks {
3946 let (rows, cols) = block.factor_a.dim();
3947 if rows != cols {
3948 return Err(ArrowSchurGpuFailure::Unavailable);
3949 }
3950 smooth_offsets_host.push(checked_i32(block.global_offset)?);
3951 smooth_m_host.push(checked_i32(rows)?);
3952 for r in 0..rows {
3953 for c in 0..cols {
3954 smooth_data_host.push(block.factor_a[[r, c]]);
3955 }
3956 }
3957 smooth_ptr_host.push(checked_i32(smooth_data_host.len())?);
3958 }
3959
3960 let mut g_col_off_host = Vec::with_capacity(data.sparse_g_blocks.len());
3961 let mut g_cols_host = Vec::with_capacity(data.sparse_g_blocks.len());
3962 let mut g_ptr_host = Vec::with_capacity(data.sparse_g_blocks.len() + 1);
3963 let mut g_data_host = Vec::<f64>::new();
3964 let mut g_output_spans = Vec::with_capacity(data.sparse_g_blocks.len());
3965 g_ptr_host.push(0_i32);
3966 for block in &data.sparse_g_blocks {
3967 let (rows, cols) = block.data.dim();
3968 g_col_off_host.push(checked_i32(block.col_off)?);
3969 g_cols_host.push(checked_i32(cols)?);
3970 for r in 0..rows {
3971 for c in 0..cols {
3972 g_data_host.push(block.data[[r, c]]);
3973 }
3974 }
3975 g_ptr_host.push(checked_i32(g_data_host.len())?);
3976 let offset = block
3977 .row_off
3978 .checked_mul(p)
3979 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
3980 let width = rows
3981 .checked_mul(p)
3982 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
3983 g_output_spans.push((offset, width));
3984 }
3985 let g_groups =
3986 super::deterministic_output_groups(&g_output_spans, k, "sparse G matvec")?;
3987 let g_group_count = g_groups.off.len();
3988
3989 let mut ainv_host = vec![0.0_f64; n_rows * max_q * max_q];
3990 for (row_idx, row) in sys.rows.iter().enumerate() {
3991 let q = data.local_jac[row_idx].len() / p;
3992 if row.htt.dim() != (q, q) {
3993 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
3994 reason: format!(
3995 "SAE device PCG row {row_idx}: H_tt shape {:?} != ({q}, {q})",
3996 row.htt.dim()
3997 ),
3998 });
3999 }
4000 let mut block = row.htt.clone();
4001 for d in 0..q {
4002 block[[d, d]] += ridge_t;
4003 }
4004 let factor = gam_linalg::triangular::cholesky_factor_in_place(
4005 block.view(),
4006 gam_linalg::triangular::CholeskyGuard::NonnegativePivot,
4007 )
4008 .ok_or_else(|| {
4009 ArrowSchurGpuFailure::RidgeBumpRequired {
4012 row: row_idx,
4013 bump: super::ridge_bump_to_make_pd(row.htt.view(), ridge_t),
4014 }
4015 })?;
4016 for col in 0..q {
4017 let mut e = Array1::<f64>::zeros(q);
4018 e[col] = 1.0;
4019 let solved = gam_linalg::triangular::cholesky_solve_vector(factor.view(), e.view());
4020 for r in 0..q {
4021 ainv_host[row_idx * max_q * max_q + r * max_q + col] = solved[r];
4022 }
4023 }
4024 }
4025
4026 Ok(DeviceSaePcgBuffers {
4027 row_ptr: stream
4028 .clone_htod(&row_ptr_host)
4029 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
4030 beta_base: stream
4031 .clone_htod(&beta_base_host)
4032 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
4033 phi: stream
4034 .clone_htod(&phi_host)
4035 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
4036 jac_ptr: stream
4037 .clone_htod(&jac_ptr_host)
4038 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
4039 jac: stream
4040 .clone_htod(&jac_host)
4041 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
4042 smooth_offsets: stream
4043 .clone_htod(&smooth_offsets_host)
4044 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
4045 smooth_m: stream
4046 .clone_htod(&smooth_m_host)
4047 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
4048 smooth_ptr: stream
4049 .clone_htod(&smooth_ptr_host)
4050 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
4051 smooth_data: stream
4052 .clone_htod(&smooth_data_host)
4053 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
4054 g_col_off: stream
4055 .clone_htod(&g_col_off_host)
4056 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
4057 g_cols: stream
4058 .clone_htod(&g_cols_host)
4059 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
4060 g_ptr: stream
4061 .clone_htod(&g_ptr_host)
4062 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
4063 g_data: stream
4064 .clone_htod(&g_data_host)
4065 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
4066 g_group_ptr: stream
4067 .clone_htod(&g_groups.ptr)
4068 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
4069 g_group_blocks: stream
4070 .clone_htod(&g_groups.blocks)
4071 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
4072 g_group_off: stream
4073 .clone_htod(&g_groups.off)
4074 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
4075 g_group_width: stream
4076 .clone_htod(&g_groups.width)
4077 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
4078 ainv: stream
4079 .clone_htod(&ainv_host)
4080 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
4081 u: stream
4082 .alloc_zeros::<f64>(n_rows * p)
4083 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
4084 w: stream
4085 .alloc_zeros::<f64>(n_rows * max_q)
4086 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
4087 v: stream
4088 .alloc_zeros::<f64>(n_rows * max_q)
4089 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
4090 scatter_partial: stream
4091 .alloc_zeros::<f64>(scatter_n_chunks * k)
4092 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
4093 scatter_n_chunks,
4094 scatter_rows_per_chunk,
4095 n_rows,
4096 p,
4097 k,
4098 max_q,
4099 smooth_blocks: data.smooth_blocks.len(),
4100 g_groups: g_group_count,
4101 g_max_work: g_groups.max_width,
4102 })
4103 }
4104
4105 fn launch_sae_init(
4106 stream: &Arc<CudaStream>,
4107 module: &Arc<CudaModule>,
4108 out: &mut CudaSlice<f64>,
4109 x: &CudaSlice<f64>,
4110 ridge: f64,
4111 n: usize,
4112 ) -> Result<(), ArrowSchurGpuFailure> {
4113 let kernel = module
4114 .load_function("arrow_sae_init")
4115 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4116 let n_i32 = checked_i32(n)?;
4117 let mut builder = stream.launch_builder(&kernel);
4118 builder.arg(out).arg(x).arg(&ridge).arg(&n_i32);
4119 unsafe { builder.launch(pcg_launch_config(n)?) }
4123 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4124 Ok(())
4125 }
4126
4127 fn launch_sae_penalty_matvec(
4128 stream: &Arc<CudaStream>,
4129 module: &Arc<CudaModule>,
4130 buffers: &mut DeviceSaePcgBuffers,
4131 x: &CudaSlice<f64>,
4132 out: &mut CudaSlice<f64>,
4133 ridge_beta: f64,
4134 ) -> Result<(), ArrowSchurGpuFailure> {
4135 launch_sae_init(stream, module, out, x, ridge_beta, buffers.k)?;
4136 if buffers.smooth_blocks > 0 {
4137 let kernel = module
4138 .load_function("arrow_sae_smooth_matvec")
4139 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4140 let max_m = buffers.k;
4141 let p_i32 = checked_i32(buffers.p)?;
4142 let blocks_i32 = checked_i32(buffers.smooth_blocks)?;
4143 let cfg = LaunchConfig {
4144 grid_dim: (
4145 ((max_m as u32).saturating_add(255) / 256).max(1),
4146 checked_i32(buffers.smooth_blocks)? as u32,
4147 1,
4148 ),
4149 block_dim: (256, 1, 1),
4150 shared_mem_bytes: 0,
4151 };
4152 let mut builder = stream.launch_builder(&kernel);
4153 builder
4154 .arg(x)
4155 .arg(&mut *out)
4156 .arg(&buffers.smooth_offsets)
4157 .arg(&buffers.smooth_m)
4158 .arg(&buffers.smooth_ptr)
4159 .arg(&buffers.smooth_data)
4160 .arg(&p_i32)
4161 .arg(&blocks_i32);
4162 unsafe { builder.launch(cfg) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4167 }
4168 if buffers.g_groups > 0 {
4169 let kernel = module
4170 .load_function("arrow_sae_sparse_g_matvec")
4171 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4172 let p_i32 = checked_i32(buffers.p)?;
4173 let groups_i32 = checked_i32(buffers.g_groups)?;
4174 let cfg = LaunchConfig {
4175 grid_dim: (
4176 ((buffers.g_max_work as u32).saturating_add(255) / 256).max(1),
4177 groups_i32 as u32,
4178 1,
4179 ),
4180 block_dim: (256, 1, 1),
4181 shared_mem_bytes: 0,
4182 };
4183 let mut builder = stream.launch_builder(&kernel);
4184 builder
4185 .arg(x)
4186 .arg(&mut *out)
4187 .arg(&buffers.g_col_off)
4188 .arg(&buffers.g_cols)
4189 .arg(&buffers.g_ptr)
4190 .arg(&buffers.g_data)
4191 .arg(&buffers.g_group_ptr)
4192 .arg(&buffers.g_group_blocks)
4193 .arg(&buffers.g_group_off)
4194 .arg(&buffers.g_group_width)
4195 .arg(&p_i32)
4196 .arg(&groups_i32);
4197 unsafe { builder.launch(cfg) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4202 }
4203 Ok(())
4204 }
4205
4206 fn launch_sae_row_schur_sub(
4207 stream: &Arc<CudaStream>,
4208 module: &Arc<CudaModule>,
4209 buffers: &mut DeviceSaePcgBuffers,
4210 x: &CudaSlice<f64>,
4211 out: &mut CudaSlice<f64>,
4212 ) -> Result<(), ArrowSchurGpuFailure> {
4213 let p_i32 = checked_i32(buffers.p)?;
4214 let max_q_i32 = checked_i32(buffers.max_q)?;
4215 let n_rows_i32 = checked_i32(buffers.n_rows)?;
4216 let cfg_p_rows = LaunchConfig {
4217 grid_dim: (
4218 ((buffers.p as u32).saturating_add(255) / 256).max(1),
4219 checked_i32(buffers.n_rows)? as u32,
4220 1,
4221 ),
4222 block_dim: (256, 1, 1),
4223 shared_mem_bytes: 0,
4224 };
4225 let gather = module
4226 .load_function("arrow_sae_gather_u")
4227 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4228 {
4229 let mut builder = stream.launch_builder(&gather);
4230 builder
4231 .arg(x)
4232 .arg(&buffers.row_ptr)
4233 .arg(&buffers.beta_base)
4234 .arg(&buffers.phi)
4235 .arg(&mut buffers.u)
4236 .arg(&p_i32)
4237 .arg(&n_rows_i32);
4238 unsafe { builder.launch(cfg_p_rows) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4242 }
4243
4244 let cfg_q_rows = LaunchConfig {
4245 grid_dim: (
4246 ((buffers.max_q as u32).saturating_add(255) / 256).max(1),
4247 checked_i32(buffers.n_rows)? as u32,
4248 1,
4249 ),
4250 block_dim: (256, 1, 1),
4251 shared_mem_bytes: 0,
4252 };
4253 let apply_l = module
4254 .load_function("arrow_sae_apply_l")
4255 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4256 {
4257 let mut builder = stream.launch_builder(&apply_l);
4258 builder
4259 .arg(&buffers.u)
4260 .arg(&buffers.jac_ptr)
4261 .arg(&buffers.jac)
4262 .arg(&mut buffers.w)
4263 .arg(&p_i32)
4264 .arg(&max_q_i32)
4265 .arg(&n_rows_i32);
4266 unsafe { builder.launch(cfg_q_rows) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4270 }
4271
4272 let apply_ainv = module
4273 .load_function("arrow_sae_apply_ainv")
4274 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4275 {
4276 let mut builder = stream.launch_builder(&apply_ainv);
4277 builder
4278 .arg(&buffers.ainv)
4279 .arg(&buffers.w)
4280 .arg(&mut buffers.v)
4281 .arg(&max_q_i32)
4282 .arg(&n_rows_i32);
4283 unsafe { builder.launch(cfg_q_rows) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4287 }
4288
4289 let k_i32 = checked_i32(buffers.k)?;
4296 let rows_per_chunk_i32 = checked_i32(buffers.scatter_rows_per_chunk)?;
4297 let n_chunks_i32 = checked_i32(buffers.scatter_n_chunks)?;
4298 let k_blocks = ((buffers.k as u32).saturating_add(255) / 256).max(1);
4299 {
4300 let kernel = module
4301 .load_function("arrow_sae_scatter_sub_det_partial")
4302 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4303 let cfg = LaunchConfig {
4304 grid_dim: (
4305 ((buffers.p as u32).saturating_add(255) / 256).max(1),
4306 n_chunks_i32 as u32,
4307 1,
4308 ),
4309 block_dim: (256, 1, 1),
4310 shared_mem_bytes: 0,
4311 };
4312 let mut builder = stream.launch_builder(&kernel);
4313 builder
4314 .arg(&buffers.v)
4315 .arg(&buffers.jac_ptr)
4316 .arg(&buffers.jac)
4317 .arg(&buffers.row_ptr)
4318 .arg(&buffers.beta_base)
4319 .arg(&buffers.phi)
4320 .arg(&mut buffers.scatter_partial)
4321 .arg(&p_i32)
4322 .arg(&k_i32)
4323 .arg(&max_q_i32)
4324 .arg(&n_rows_i32)
4325 .arg(&rows_per_chunk_i32);
4326 unsafe { builder.launch(cfg) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4331 }
4332 {
4333 let kernel = module
4334 .load_function("arrow_sae_scatter_sub_det_reduce")
4335 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4336 let cfg = LaunchConfig {
4337 grid_dim: (k_blocks, 1, 1),
4338 block_dim: (256, 1, 1),
4339 shared_mem_bytes: 0,
4340 };
4341 let mut builder = stream.launch_builder(&kernel);
4342 builder
4343 .arg(&buffers.scatter_partial)
4344 .arg(&mut *out)
4345 .arg(&k_i32)
4346 .arg(&n_chunks_i32);
4347 unsafe { builder.launch(cfg) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4350 }
4351 Ok(())
4352 }
4353
4354 fn launch_sae_diag_sub(
4355 stream: &Arc<CudaStream>,
4356 module: &Arc<CudaModule>,
4357 buffers: &mut DeviceSaePcgBuffers,
4358 diag: &mut CudaSlice<f64>,
4359 ) -> Result<(), ArrowSchurGpuFailure> {
4360 let p_i32 = checked_i32(buffers.p)?;
4361 let k_i32 = checked_i32(buffers.k)?;
4362 let max_q_i32 = checked_i32(buffers.max_q)?;
4363 let n_rows_i32 = checked_i32(buffers.n_rows)?;
4364 let rows_per_chunk_i32 = checked_i32(buffers.scatter_rows_per_chunk)?;
4365 let n_chunks_i32 = checked_i32(buffers.scatter_n_chunks)?;
4366 {
4367 let kernel = module
4368 .load_function("arrow_sae_diag_sub_det_partial")
4369 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4370 let cfg = LaunchConfig {
4371 grid_dim: (
4372 ((buffers.p as u32).saturating_add(255) / 256).max(1),
4373 n_chunks_i32 as u32,
4374 1,
4375 ),
4376 block_dim: (256, 1, 1),
4377 shared_mem_bytes: 0,
4378 };
4379 let mut builder = stream.launch_builder(&kernel);
4380 builder
4381 .arg(&buffers.ainv)
4382 .arg(&buffers.jac_ptr)
4383 .arg(&buffers.jac)
4384 .arg(&buffers.row_ptr)
4385 .arg(&buffers.beta_base)
4386 .arg(&buffers.phi)
4387 .arg(&mut buffers.scatter_partial)
4388 .arg(&p_i32)
4389 .arg(&k_i32)
4390 .arg(&max_q_i32)
4391 .arg(&n_rows_i32)
4392 .arg(&rows_per_chunk_i32);
4393 unsafe { builder.launch(cfg) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4398 }
4399 {
4400 let kernel = module
4401 .load_function("arrow_sae_scatter_sub_det_reduce")
4402 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4403 let cfg = LaunchConfig {
4404 grid_dim: (
4405 ((buffers.k as u32).saturating_add(255) / 256).max(1),
4406 1,
4407 1,
4408 ),
4409 block_dim: (256, 1, 1),
4410 shared_mem_bytes: 0,
4411 };
4412 let mut builder = stream.launch_builder(&kernel);
4413 builder
4414 .arg(&buffers.scatter_partial)
4415 .arg(diag)
4416 .arg(&k_i32)
4417 .arg(&n_chunks_i32);
4418 unsafe { builder.launch(cfg) }
4421 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4422 Ok(())
4423 }
4424 }
4425
4426 fn launch_sae_matvec(
4427 stream: &Arc<CudaStream>,
4428 module: &Arc<CudaModule>,
4429 buffers: &mut DeviceSaePcgBuffers,
4430 x: &CudaSlice<f64>,
4431 out: &mut CudaSlice<f64>,
4432 ridge_beta: f64,
4433 ) -> Result<(), ArrowSchurGpuFailure> {
4434 launch_sae_penalty_matvec(stream, module, buffers, x, out, ridge_beta)?;
4435 launch_sae_row_schur_sub(stream, module, buffers, x, out)
4436 }
4437
4438 fn pack_fused_host(
4443 sys: &ArrowSchurSystem,
4444 ridge_t: f64,
4445 p_max: usize,
4446 r_template: usize,
4447 ) -> (Vec<f64>, Vec<f64>, Vec<f64>) {
4448 let n = sys.rows.len();
4449 let d = sys.d;
4450 let k = sys.k;
4451 let mut d_buf = vec![0.0_f64; n * p_max * p_max];
4452 let mut b_buf = vec![0.0_f64; n * p_max * r_template];
4453 let mut g_buf = vec![0.0_f64; n * p_max];
4454 for (i, row) in sys.rows.iter().enumerate() {
4455 for col in 0..d {
4457 let base = (i * p_max + col) * p_max;
4458 for r in 0..d {
4459 let mut value = row.htt[[r, col]];
4460 if r == col {
4461 value += ridge_t;
4462 }
4463 d_buf[base + r] = value;
4464 }
4465 }
4466 for col in 0..k {
4474 let base = (i * r_template + col) * p_max;
4475 for r in 0..d {
4476 b_buf[base + r] = row.htbeta[[r, col]];
4477 }
4478 }
4479 let g_base = i * p_max;
4481 for r in 0..d {
4482 g_buf[g_base + r] = row.gt[r];
4483 }
4484 }
4485 (d_buf, b_buf, g_buf)
4486 }
4487
4488 pub(super) struct ResidentArrowFrame {
4515 n: usize,
4516 d: usize,
4517 k: usize,
4518 beta_gauge_quotient: Option<ArrowBetaGaugeQuotient>,
4519 stream: Arc<CudaStream>,
4520 blas: CudaBlas,
4521 l_dev: CudaSlice<f64>,
4524 y_dev: CudaSlice<f64>,
4527 schur_dev: CudaSlice<f64>,
4530 log_det_hessian: f64,
4533 }
4534
4535 impl ResidentArrowFrame {
4536 pub(super) fn new(
4540 sys: &ArrowSchurSystem,
4541 ridge_t: f64,
4542 ridge_beta: f64,
4543 newton_schur_tikhonov_rel_floor: Option<f64>,
4544 ) -> Result<Self, ArrowSchurGpuFailure> {
4545 if ridge_t.is_nan() || ridge_beta.is_nan() {
4546 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
4547 reason: "ridge is NaN".to_string(),
4548 });
4549 }
4550 let n = sys.rows.len();
4551 let d = sys.d;
4552 let k = sys.k;
4553 let runtime = route_through_gpu(DispatchOp::SmallDenseBatchedPotrf { p: d, batch: n })
4554 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
4555 let stream = gam_gpu::device_runtime::cuda_context_for(runtime.device.ordinal)
4556 .and_then(|ctx| ctx.new_stream().ok())
4557 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
4558 let solver =
4559 DnHandle::new(stream.clone()).map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4560 let blas =
4561 CudaBlas::new(stream.clone()).map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4562
4563 let (d_host, b_host_stacked) = pack_host_d_and_stacked_b(sys, ridge_t);
4569 let mut l_dev = stream
4570 .clone_htod(&d_host)
4571 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4572 let mut y_dev = stream
4573 .clone_htod(&b_host_stacked)
4574 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4575
4576 let info_host = potrf_batched(&solver, &stream, d, n, &mut l_dev)?;
4578 if let Some(idx) = info_host.iter().position(|info| *info != 0) {
4579 return Err(ArrowSchurGpuFailure::RidgeBumpRequired {
4583 row: idx,
4584 bump: super::ridge_bump_to_make_pd(sys.rows[idx].htt.view(), ridge_t),
4585 });
4586 }
4587
4588 trsm_batched_lower_inplace_stacked(&blas, &stream, d, n, k, &l_dev, &mut y_dev)?;
4590
4591 let schur_init: Vec<f64> = {
4595 let mut tmp = Vec::with_capacity(k * k);
4596 for col in 0..k {
4597 for row in 0..k {
4598 let mut v = sys.hbb[[row, col]];
4599 if row == col {
4600 v += ridge_beta;
4601 }
4602 tmp.push(v);
4603 }
4604 }
4605 tmp
4606 };
4607 let mut schur_dev = stream
4608 .clone_htod(&schur_init)
4609 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4610 schur_gemm_stacked(&blas, n * d, k, &y_dev, &mut schur_dev)?;
4611 let canonical_factor = if sys.beta_gauge_quotient.is_some()
4612 || newton_schur_tikhonov_rel_floor.is_some()
4613 {
4614 let mut schur_host = stream
4615 .clone_dtoh(&schur_dev)
4616 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4617 let factorized = canonicalize_device_beta_factor(
4618 sys.beta_gauge_quotient.as_ref(),
4619 newton_schur_tikhonov_rel_floor,
4620 k,
4621 &mut schur_host,
4622 )?;
4623 schur_dev = stream
4624 .clone_htod(&schur_host)
4625 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4626 factorized
4627 } else {
4628 false
4629 };
4630 if !canonical_factor {
4631 let info = potrf_single(&solver, &stream, k, &mut schur_dev)?;
4632 if info != 0 {
4633 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
4634 reason: format!("Schur Cholesky failed at pivot {info}"),
4635 });
4636 }
4637 }
4638
4639 let l_local_host = stream
4641 .clone_dtoh(&l_dev)
4642 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4643 let l_schur_host = stream
4644 .clone_dtoh(&schur_dev)
4645 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4646 let mut log_det = 0.0_f64;
4647 for i in 0..n {
4648 let base = i * d * d;
4649 for j in 0..d {
4650 log_det += l_local_host[base + j * d + j].ln();
4651 }
4652 }
4653 for j in 0..k {
4654 log_det += l_schur_host[j * k + j].ln();
4655 }
4656 log_det *= 2.0;
4657
4658 Ok(Self {
4659 n,
4660 d,
4661 k,
4662 beta_gauge_quotient: sys.beta_gauge_quotient.clone(),
4663 stream,
4664 blas,
4665 l_dev,
4666 y_dev,
4667 schur_dev,
4668 log_det_hessian: log_det,
4669 })
4670 }
4671
4672 #[inline]
4673 pub(super) fn log_det_hessian(&self) -> f64 {
4674 self.log_det_hessian
4675 }
4676
4677 pub(super) fn solve_gradient(
4681 &self,
4682 g_t: &[f64],
4683 g_beta: &[f64],
4684 ) -> Result<ArrowSchurGpuSolution, ArrowSchurGpuFailure> {
4685 let n = self.n;
4686 let d = self.d;
4687 let k = self.k;
4688 if g_t.len() != n * d || g_beta.len() != k {
4689 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
4690 reason: format!(
4691 "resident gradient shape mismatch: g_t={} (want {}), g_beta={} (want {})",
4692 g_t.len(),
4693 n * d,
4694 g_beta.len(),
4695 k
4696 ),
4697 });
4698 }
4699 let mut u_dev = self
4701 .stream
4702 .clone_htod(&g_t.to_vec())
4703 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4704 trsm_batched_lower_inplace(&self.blas, &self.stream, d, n, 1, &self.l_dev, &mut u_dev)?;
4705
4706 let rhs_init: Vec<f64> = g_beta.iter().map(|v| -v).collect();
4709 let mut rhs_dev = self
4710 .stream
4711 .clone_htod(&rhs_init)
4712 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4713 schur_rhs_stacked(&self.blas, n * d, k, &self.y_dev, &u_dev, &mut rhs_dev)?;
4714 if self.beta_gauge_quotient.is_some() {
4715 let projected_rhs = project_device_beta_vector(
4716 self.beta_gauge_quotient.as_ref(),
4717 k,
4718 self.stream
4719 .clone_dtoh(&rhs_dev)
4720 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
4721 )?;
4722 rhs_dev = self
4723 .stream
4724 .clone_htod(&projected_rhs)
4725 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4726 }
4727
4728 trsm_single(
4730 &self.blas,
4731 &self.stream,
4732 k,
4733 &self.schur_dev,
4734 &mut rhs_dev,
4735 false,
4736 false,
4737 )?;
4738 trsm_single(
4739 &self.blas,
4740 &self.stream,
4741 k,
4742 &self.schur_dev,
4743 &mut rhs_dev,
4744 false,
4745 true,
4746 )?;
4747 let delta_beta_host = project_device_beta_vector(
4748 self.beta_gauge_quotient.as_ref(),
4749 k,
4750 self.stream
4751 .clone_dtoh(&rhs_dev)
4752 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
4753 )?;
4754 if self.beta_gauge_quotient.is_some() {
4755 rhs_dev = self
4756 .stream
4757 .clone_htod(&delta_beta_host)
4758 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4759 }
4760 let delta_beta = Array1::from_vec(delta_beta_host);
4761
4762 back_sub_rhs_stacked(&self.blas, n * d, k, &self.y_dev, &rhs_dev, &mut u_dev)?;
4764 trsm_batched_lower_inplace_transposed(
4765 &self.blas,
4766 &self.stream,
4767 d,
4768 n,
4769 1,
4770 &self.l_dev,
4771 &mut u_dev,
4772 )?;
4773 let x_host = self
4774 .stream
4775 .clone_dtoh(&u_dev)
4776 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4777 let mut delta_t = Array1::<f64>::zeros(n * d);
4778 for (i, v) in x_host.iter().enumerate() {
4779 delta_t[i] = -*v;
4780 }
4781
4782 Ok(ArrowSchurGpuSolution {
4783 delta_t,
4784 delta_beta,
4785 log_det_hessian: self.log_det_hessian,
4786 })
4787 }
4788 }
4789
4790 pub(super) struct ResidentBaseArrowFrame {
4796 n: usize,
4797 d: usize,
4798 k: usize,
4799 beta_gauge_quotient: Option<ArrowBetaGaugeQuotient>,
4800 newton_schur_tikhonov_rel_floor: Option<f64>,
4801 stream: Arc<CudaStream>,
4802 solver: DnHandle,
4803 blas: CudaBlas,
4804 d_base_host: Vec<f64>,
4810 base_b_dev: CudaSlice<f64>,
4813 base_hbb_dev: CudaSlice<f64>,
4816 g_t_dev: CudaSlice<f64>,
4818 gb_host: Vec<f64>,
4821 ones_k_dev: CudaSlice<f64>,
4824 }
4825
4826 impl ResidentBaseArrowFrame {
4827 pub(super) fn new(
4829 sys: &ArrowSchurSystem,
4830 newton_schur_tikhonov_rel_floor: Option<f64>,
4831 ) -> Result<Self, ArrowSchurGpuFailure> {
4832 let n = sys.rows.len();
4833 let d = sys.d;
4834 let k = sys.k;
4835 let runtime = route_through_gpu(DispatchOp::SmallDenseBatchedPotrf { p: d, batch: n })
4836 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
4837 let stream = gam_gpu::device_runtime::cuda_context_for(runtime.device.ordinal)
4838 .and_then(|ctx| ctx.new_stream().ok())
4839 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
4840 let solver =
4841 DnHandle::new(stream.clone()).map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4842 let blas =
4843 CudaBlas::new(stream.clone()).map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4844
4845 let (d_base_host, b_host, g_host) = pack_host(sys, 0.0);
4847 let base_b_dev = stream
4848 .clone_htod(&b_host)
4849 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4850 let g_t_dev = stream
4851 .clone_htod(&g_host)
4852 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4853
4854 let hbb_base: Vec<f64> = {
4856 let mut tmp = Vec::with_capacity(k * k);
4857 for col in 0..k {
4858 for row in 0..k {
4859 tmp.push(sys.hbb[[row, col]]);
4860 }
4861 }
4862 tmp
4863 };
4864 let base_hbb_dev = stream
4865 .clone_htod(&hbb_base)
4866 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4867
4868 let gb_host: Vec<f64> = sys.gb.iter().copied().collect();
4869 let ones_k_dev = stream
4870 .clone_htod(&vec![1.0_f64; k])
4871 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
4872
4873 Ok(Self {
4874 n,
4875 d,
4876 k,
4877 beta_gauge_quotient: sys.beta_gauge_quotient.clone(),
4878 newton_schur_tikhonov_rel_floor,
4879 stream,
4880 solver,
4881 blas,
4882 d_base_host,
4883 base_b_dev,
4884 base_hbb_dev,
4885 g_t_dev,
4886 gb_host,
4887 ones_k_dev,
4888 })
4889 }
4890
4891 pub(super) fn refactor_and_solve(
4897 &self,
4898 ridge_t: f64,
4899 ridge_beta: f64,
4900 ) -> Result<ArrowSchurGpuSolution, ArrowSchurGpuFailure> {
4901 self.refactor_and_solve_from(ridge_t, ridge_beta, None)
4902 }
4903
4904 pub(super) fn solve_with_gradient_checked(
4911 &self,
4912 factors: &BaseRidgeFactors,
4913 g_t: &[f64],
4914 g_beta: &[f64],
4915 ) -> Result<ArrowSchurGpuSolution, ArrowSchurGpuFailure> {
4916 self.check_gradient_lens(g_t, g_beta)?;
4917 self.solve_with_factors(factors, Some((g_t, g_beta)))
4918 }
4919
4920 fn check_gradient_lens(
4921 &self,
4922 g_t: &[f64],
4923 g_beta: &[f64],
4924 ) -> Result<(), ArrowSchurGpuFailure> {
4925 if g_t.len() != self.n * self.d {
4926 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
4927 reason: format!(
4928 "base-resident frame row gradient length mismatch: got {}, expected {}",
4929 g_t.len(),
4930 self.n * self.d
4931 ),
4932 });
4933 }
4934 if g_beta.len() != self.k {
4935 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
4936 reason: format!(
4937 "base-resident frame border gradient length mismatch: got {}, expected {}",
4938 g_beta.len(),
4939 self.k
4940 ),
4941 });
4942 }
4943 Ok(())
4944 }
4945
4946 fn refactor_and_solve_from(
4958 &self,
4959 ridge_t: f64,
4960 ridge_beta: f64,
4961 gradient: Option<(&[f64], &[f64])>,
4962 ) -> Result<ArrowSchurGpuSolution, ArrowSchurGpuFailure> {
4963 let factors = self.factor_at(ridge_t, ridge_beta)?;
4964 self.solve_with_factors(&factors, gradient)
4965 }
4966
4967 pub(super) fn factor_at(
4977 &self,
4978 ridge_t: f64,
4979 ridge_beta: f64,
4980 ) -> Result<BaseRidgeFactors, ArrowSchurGpuFailure> {
4981 if ridge_t.is_nan() || ridge_beta.is_nan() {
4982 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
4983 reason: "ridge is NaN".to_string(),
4984 });
4985 }
4986 let (n, d, k) = (self.n, self.d, self.k);
4987
4988 let mut d_ridged = self.d_base_host.clone();
4992 for i in 0..n {
4993 for j in 0..d {
4994 d_ridged[i * d * d + j * (d + 1)] += ridge_t;
4995 }
4996 }
4997 let mut l_dev = self
4998 .stream
4999 .clone_htod(&d_ridged)
5000 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5001
5002 let info_host = potrf_batched(&self.solver, &self.stream, d, n, &mut l_dev)?;
5004 if let Some(idx) = info_host.iter().position(|info| *info != 0) {
5005 let base = idx * d * d;
5006 return Err(ArrowSchurGpuFailure::RidgeBumpRequired {
5007 row: idx,
5008 bump: super::ridge_bump_to_make_pd_colmajor(&d_ridged[base..base + d * d], d),
5012 });
5013 }
5014
5015 let mut y_dev = self
5017 .stream
5018 .alloc_zeros::<f64>(n * d * k)
5019 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5020 self.stream
5021 .memcpy_dtod(&self.base_b_dev, &mut y_dev)
5022 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5023 trsm_batched_lower_inplace(&self.blas, &self.stream, d, n, k, &l_dev, &mut y_dev)?;
5024
5025 let mut schur_dev = self
5027 .stream
5028 .alloc_zeros::<f64>(k * k)
5029 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5030 self.stream
5031 .memcpy_dtod(&self.base_hbb_dev, &mut schur_dev)
5032 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5033 device_axpy_strided(
5035 &self.blas,
5036 &self.stream,
5037 k,
5038 ridge_beta,
5039 &self.ones_k_dev,
5040 1,
5041 &mut schur_dev,
5042 k + 1,
5043 )?;
5044 accumulate_schur_reduce_only(&self.blas, d, k, n, &y_dev, &mut schur_dev)?;
5045 let canonical_factor = if self.beta_gauge_quotient.is_some()
5046 || self.newton_schur_tikhonov_rel_floor.is_some()
5047 {
5048 let mut schur_host = self
5049 .stream
5050 .clone_dtoh(&schur_dev)
5051 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5052 let factorized = canonicalize_device_beta_factor(
5053 self.beta_gauge_quotient.as_ref(),
5054 self.newton_schur_tikhonov_rel_floor,
5055 k,
5056 &mut schur_host,
5057 )?;
5058 schur_dev = self
5059 .stream
5060 .clone_htod(&schur_host)
5061 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5062 factorized
5063 } else {
5064 false
5065 };
5066
5067 if !canonical_factor {
5069 let info = potrf_single(&self.solver, &self.stream, k, &mut schur_dev)?;
5070 if info != 0 {
5071 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
5072 reason: format!("Schur Cholesky failed at pivot {info}"),
5073 });
5074 }
5075 }
5076
5077 let l_local_host = self
5080 .stream
5081 .clone_dtoh(&l_dev)
5082 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5083 let l_schur_host = self
5084 .stream
5085 .clone_dtoh(&schur_dev)
5086 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5087 let mut log_det = 0.0_f64;
5088 for i in 0..n {
5089 let base = i * d * d;
5090 for j in 0..d {
5091 log_det += l_local_host[base + j * d + j].ln();
5092 }
5093 }
5094 for j in 0..k {
5095 log_det += l_schur_host[j * k + j].ln();
5096 }
5097 log_det *= 2.0;
5098
5099 Ok(BaseRidgeFactors {
5100 ridge_t,
5101 ridge_beta,
5102 l_dev,
5103 y_dev,
5104 schur_dev,
5105 log_det_hessian: log_det,
5106 })
5107 }
5108
5109 pub(super) fn solve_with_factors(
5118 &self,
5119 factors: &BaseRidgeFactors,
5120 gradient: Option<(&[f64], &[f64])>,
5121 ) -> Result<ArrowSchurGpuSolution, ArrowSchurGpuFailure> {
5122 let (n, d, k) = (self.n, self.d, self.k);
5123 let l_dev = &factors.l_dev;
5124 let y_dev = &factors.y_dev;
5125 let schur_dev = &factors.schur_dev;
5126
5127 let mut u_dev = match gradient {
5131 Some((g_t, _)) => self
5132 .stream
5133 .clone_htod(g_t)
5134 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
5135 None => {
5136 let mut resident = self
5137 .stream
5138 .alloc_zeros::<f64>(n * d)
5139 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5140 self.stream
5141 .memcpy_dtod(&self.g_t_dev, &mut resident)
5142 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5143 resident
5144 }
5145 };
5146 trsm_batched_lower_inplace(&self.blas, &self.stream, d, n, 1, l_dev, &mut u_dev)?;
5147
5148 let rhs_init: Vec<f64> = match gradient {
5150 Some((_, g_beta)) => g_beta.iter().map(|v| -v).collect(),
5151 None => self.gb_host.iter().map(|v| -v).collect(),
5152 };
5153 let mut rhs_dev = self
5154 .stream
5155 .clone_htod(&rhs_init)
5156 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5157 accumulate_schur_rhs_only(&self.blas, d, k, n, y_dev, &u_dev, &mut rhs_dev)?;
5158 if self.beta_gauge_quotient.is_some() {
5159 let projected_rhs = project_device_beta_vector(
5160 self.beta_gauge_quotient.as_ref(),
5161 k,
5162 self.stream
5163 .clone_dtoh(&rhs_dev)
5164 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
5165 )?;
5166 rhs_dev = self
5167 .stream
5168 .clone_htod(&projected_rhs)
5169 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5170 }
5171
5172 trsm_single(
5174 &self.blas,
5175 &self.stream,
5176 k,
5177 schur_dev,
5178 &mut rhs_dev,
5179 false,
5180 false,
5181 )?;
5182 trsm_single(
5183 &self.blas,
5184 &self.stream,
5185 k,
5186 schur_dev,
5187 &mut rhs_dev,
5188 false,
5189 true,
5190 )?;
5191 let delta_beta_host = project_device_beta_vector(
5192 self.beta_gauge_quotient.as_ref(),
5193 k,
5194 self.stream
5195 .clone_dtoh(&rhs_dev)
5196 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
5197 )?;
5198 if self.beta_gauge_quotient.is_some() {
5199 rhs_dev = self
5200 .stream
5201 .clone_htod(&delta_beta_host)
5202 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5203 }
5204 let delta_beta = Array1::from_vec(delta_beta_host);
5205
5206 accumulate_back_sub_rhs(&self.blas, d, k, n, y_dev, &rhs_dev, &mut u_dev)?;
5208 trsm_batched_lower_inplace_transposed(
5209 &self.blas,
5210 &self.stream,
5211 d,
5212 n,
5213 1,
5214 l_dev,
5215 &mut u_dev,
5216 )?;
5217 let x_host = self
5218 .stream
5219 .clone_dtoh(&u_dev)
5220 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5221 let mut delta_t = Array1::<f64>::zeros(n * d);
5222 for (i, v) in x_host.iter().enumerate() {
5223 delta_t[i] = -*v;
5224 }
5225
5226 Ok(ArrowSchurGpuSolution {
5227 delta_t,
5228 delta_beta,
5229 log_det_hessian: factors.log_det_hessian,
5230 })
5231 }
5232 }
5233
5234 pub(super) struct BaseRidgeFactors {
5245 ridge_t: f64,
5246 ridge_beta: f64,
5247 l_dev: CudaSlice<f64>,
5250 y_dev: CudaSlice<f64>,
5252 schur_dev: CudaSlice<f64>,
5254 log_det_hessian: f64,
5256 }
5257
5258 impl BaseRidgeFactors {
5259 #[inline]
5263 pub(super) fn matches_ridge(&self, ridge_t: f64, ridge_beta: f64) -> bool {
5264 self.ridge_t == ridge_t && self.ridge_beta == ridge_beta
5265 }
5266 }
5267
5268 pub(super) fn solve_fused(
5269 sys: &ArrowSchurSystem,
5270 ridge_t: f64,
5271 ridge_beta: f64,
5272 newton_schur_tikhonov_rel_floor: Option<f64>,
5273 ) -> Result<ArrowSchurGpuSolution, ArrowSchurGpuFailure> {
5274 let n = sys.rows.len();
5275 let d = sys.d;
5276 let k = sys.k;
5277 let plan = crate::gpu_kernels::arrow_schur_nvrtc::plan_fused_launch(n, d, k)
5278 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
5279 let p_max = plan.p_max;
5280 let r_template = plan.r_template;
5281
5282 let runtime = gam_gpu::linalg_dispatch::route_through_gpu(
5283 gam_gpu::linalg_dispatch::DispatchOp::SmallDenseBatchedPotrf { p: d, batch: n },
5284 )
5285 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
5286 let ctx = gam_gpu::device_runtime::cuda_context_for(runtime.device.ordinal)
5287 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
5288 let stream = ctx
5289 .new_stream()
5290 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5291 let cap = &runtime.device.capability;
5292 let key = crate::gpu_kernels::arrow_schur_nvrtc::FusedModuleCacheKey {
5293 cc_major: cap.compute_major,
5294 cc_minor: cap.compute_minor,
5295 p_max: p_max as u32,
5296 r_template: r_template as u32,
5297 };
5298 let module = fused_module_for(&ctx, key)?;
5299 let forward = module
5300 .load_function("arrow_schur_forward_pgroup")
5301 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5302 let back_sub = module
5303 .load_function("arrow_schur_back_sub_pgroup")
5304 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5305
5306 let (d_host, b_host, g_host) = pack_fused_host(sys, ridge_t, p_max, r_template);
5308 let d_dev = stream
5309 .clone_htod(&d_host)
5310 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5311 let b_dev = stream
5312 .clone_htod(&b_host)
5313 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5314 let g_dev = stream
5315 .clone_htod(&g_host)
5316 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5317 let mut l_out = stream
5318 .alloc_zeros::<f64>(n * p_max * p_max)
5319 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5320 let mut u_out = stream
5321 .alloc_zeros::<f64>(n * p_max)
5322 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5323 let mut y_out = stream
5324 .alloc_zeros::<f64>(n * p_max * r_template)
5325 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5326 let mut partial_s = stream
5327 .alloc_zeros::<f64>(plan.partial_s_doubles)
5328 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5329 let mut partial_r = stream
5330 .alloc_zeros::<f64>(plan.partial_r_doubles)
5331 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5332 let mut status_dev = stream
5333 .alloc_zeros::<i32>(n)
5334 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5335
5336 let cfg = LaunchConfig {
5338 grid_dim: (plan.blocks, 1, 1),
5339 block_dim: (plan.threads_per_block, 1, 1),
5340 shared_mem_bytes: 0,
5341 };
5342 let n_i32 = to_i32(n).ok_or(ArrowSchurGpuFailure::Unavailable)?;
5343 let p_i32 = to_i32(d).ok_or(ArrowSchurGpuFailure::Unavailable)?;
5344 let r_i32 = to_i32(k).ok_or(ArrowSchurGpuFailure::Unavailable)?;
5345 let ridge_arg = ridge_t;
5346 {
5347 let mut builder = stream.launch_builder(&forward);
5348 builder
5349 .arg(&d_dev)
5350 .arg(&b_dev)
5351 .arg(&g_dev)
5352 .arg(&n_i32)
5353 .arg(&p_i32)
5354 .arg(&r_i32)
5355 .arg(&ridge_arg)
5356 .arg(&mut l_out)
5357 .arg(&mut u_out)
5358 .arg(&mut y_out)
5359 .arg(&mut partial_s)
5360 .arg(&mut partial_r)
5361 .arg(&mut status_dev);
5362 unsafe { builder.launch(cfg) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5366 }
5367 stream
5368 .synchronize()
5369 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5370
5371 let status_host = stream
5373 .clone_dtoh(&status_dev)
5374 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5375 if let Some(row) = status_host.iter().position(|s| *s != 0) {
5376 return Err(ArrowSchurGpuFailure::RidgeBumpRequired {
5380 row,
5381 bump: super::ridge_bump_to_make_pd(sys.rows[row].htt.view(), ridge_t),
5382 });
5383 }
5384
5385 let partial_s_host = stream
5387 .clone_dtoh(&partial_s)
5388 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5389 let partial_r_host = stream
5390 .clone_dtoh(&partial_r)
5391 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5392 let mut schur_host = vec![0.0_f64; k * k];
5393 for col in 0..k {
5394 for row in 0..k {
5395 let mut v = sys.hbb[[row, col]];
5396 if row == col {
5397 v += ridge_beta;
5398 }
5399 schur_host[col * k + row] = v;
5400 }
5401 }
5402 let mut rhs_host: Vec<f64> = sys.gb.iter().map(|v| -v).collect();
5403 for i in 0..n {
5404 let s_base = i * r_template * r_template;
5407 for col in 0..k {
5408 let col_base = s_base + col * r_template;
5409 let dst_col_base = col * k;
5410 for row in 0..k {
5411 schur_host[dst_col_base + row] -= partial_s_host[col_base + row];
5412 }
5413 }
5414 let r_base = i * r_template;
5415 for a in 0..k {
5416 rhs_host[a] += partial_r_host[r_base + a];
5417 }
5418 }
5419 let canonical_factor = canonicalize_device_beta_factor(
5420 sys.beta_gauge_quotient.as_ref(),
5421 newton_schur_tikhonov_rel_floor,
5422 sys.k,
5423 &mut schur_host,
5424 )?;
5425 rhs_host = project_device_beta_vector(sys.beta_gauge_quotient.as_ref(), sys.k, rhs_host)?;
5426
5427 let mut schur_dev = stream
5429 .clone_htod(&schur_host)
5430 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5431 let mut rhs_dev = stream
5432 .clone_htod(&rhs_host)
5433 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5434 let solver =
5435 DnHandle::new(stream.clone()).map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5436 let blas = CudaBlas::new(stream.clone()).map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5437 if !canonical_factor {
5438 let info = potrf_single(&solver, &stream, k, &mut schur_dev)?;
5439 if info != 0 {
5440 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
5441 reason: format!("fused Schur Cholesky failed at pivot {info}"),
5442 });
5443 }
5444 }
5445 trsm_single(&blas, &stream, k, &schur_dev, &mut rhs_dev, false, false)?;
5446 trsm_single(&blas, &stream, k, &schur_dev, &mut rhs_dev, false, true)?;
5447 let delta_beta_host = project_device_beta_vector(
5448 sys.beta_gauge_quotient.as_ref(),
5449 sys.k,
5450 stream
5451 .clone_dtoh(&rhs_dev)
5452 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
5453 )?;
5454 if sys.beta_gauge_quotient.is_some() {
5455 rhs_dev = stream
5456 .clone_htod(&delta_beta_host)
5457 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5458 }
5459 let delta_beta = Array1::from_vec(delta_beta_host);
5460
5461 let mut delta_t_dev = stream
5463 .alloc_zeros::<f64>(n * p_max)
5464 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5465 let back_cfg = LaunchConfig {
5466 grid_dim: (plan.blocks, 1, 1),
5467 block_dim: (plan.threads_per_block, 1, 1),
5468 shared_mem_bytes: 0,
5469 };
5470 {
5471 let mut builder = stream.launch_builder(&back_sub);
5472 builder
5473 .arg(&l_out)
5474 .arg(&u_out)
5475 .arg(&y_out)
5476 .arg(&rhs_dev)
5477 .arg(&n_i32)
5478 .arg(&p_i32)
5479 .arg(&r_i32)
5480 .arg(&mut delta_t_dev);
5481 unsafe { builder.launch(back_cfg) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5485 }
5486 stream
5487 .synchronize()
5488 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5489
5490 let delta_t_host = stream
5491 .clone_dtoh(&delta_t_dev)
5492 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5493 let mut delta_t = Array1::<f64>::zeros(n * d);
5494 for i in 0..n {
5495 let src_base = i * p_max;
5496 let dst_base = i * d;
5497 for r in 0..d {
5498 delta_t[dst_base + r] = delta_t_host[src_base + r];
5499 }
5500 }
5501
5502 let l_local_host = stream
5504 .clone_dtoh(&l_out)
5505 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5506 let l_schur_host = stream
5507 .clone_dtoh(&schur_dev)
5508 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5509 let mut log_det = 0.0_f64;
5510 for i in 0..n {
5511 let base = i * p_max * p_max;
5512 for j in 0..d {
5513 log_det += l_local_host[base + j * p_max + j].ln();
5514 }
5515 }
5516 for j in 0..k {
5517 log_det += l_schur_host[j * k + j].ln();
5518 }
5519 log_det *= 2.0;
5520
5521 Ok(ArrowSchurGpuSolution {
5522 delta_t,
5523 delta_beta,
5524 log_det_hessian: log_det,
5525 })
5526 }
5527
5528 pub(super) fn build_schur_matvec_backend(
5538 sys: &ArrowSchurSystem,
5539 ridge_t: f64,
5540 ridge_beta: f64,
5541 ) -> Result<crate::arrow_schur::GpuSchurMatvec, super::ArrowSchurGpuFailure> {
5542 let n = sys.rows.len();
5543 let d = sys.d;
5544 let k = sys.k;
5545 let plan = crate::gpu_kernels::arrow_schur_nvrtc::plan_fused_launch(n, d, k)
5546 .ok_or(super::ArrowSchurGpuFailure::Unavailable)?;
5547 let p_max = plan.p_max;
5548 let r_template = plan.r_template;
5549
5550 let runtime = gam_gpu::linalg_dispatch::route_through_gpu(
5551 gam_gpu::linalg_dispatch::DispatchOp::SmallDenseBatchedPotrf { p: d, batch: n },
5552 )
5553 .ok_or(super::ArrowSchurGpuFailure::Unavailable)?;
5554 let ctx = gam_gpu::device_runtime::cuda_context_for(runtime.device.ordinal)
5555 .ok_or(super::ArrowSchurGpuFailure::Unavailable)?;
5556 let stream = ctx
5557 .new_stream()
5558 .map_err(|_| super::ArrowSchurGpuFailure::Unavailable)?;
5559 let cap = &runtime.device.capability;
5560 let key = crate::gpu_kernels::arrow_schur_nvrtc::FusedModuleCacheKey {
5561 cc_major: cap.compute_major,
5562 cc_minor: cap.compute_minor,
5563 p_max: p_max as u32,
5564 r_template: r_template as u32,
5565 };
5566 let module = fused_module_for(&ctx, key)?;
5567 let forward = module
5568 .load_function("arrow_schur_forward_pgroup")
5569 .map_err(|_| super::ArrowSchurGpuFailure::Unavailable)?;
5570
5571 let (d_host, b_host, g_host) = pack_fused_host(sys, ridge_t, p_max, r_template);
5572 let d_dev = stream
5573 .clone_htod(&d_host)
5574 .map_err(|_| super::ArrowSchurGpuFailure::Unavailable)?;
5575 let b_dev = stream
5576 .clone_htod(&b_host)
5577 .map_err(|_| super::ArrowSchurGpuFailure::Unavailable)?;
5578 let g_dev = stream
5579 .clone_htod(&g_host)
5580 .map_err(|_| super::ArrowSchurGpuFailure::Unavailable)?;
5581 let mut l_out = stream
5582 .alloc_zeros::<f64>(n * p_max * p_max)
5583 .map_err(|_| super::ArrowSchurGpuFailure::Unavailable)?;
5584 let mut u_out = stream
5585 .alloc_zeros::<f64>(n * p_max)
5586 .map_err(|_| super::ArrowSchurGpuFailure::Unavailable)?;
5587 let mut y_out = stream
5588 .alloc_zeros::<f64>(n * p_max * r_template)
5589 .map_err(|_| super::ArrowSchurGpuFailure::Unavailable)?;
5590 let mut partial_s = stream
5591 .alloc_zeros::<f64>(plan.partial_s_doubles)
5592 .map_err(|_| super::ArrowSchurGpuFailure::Unavailable)?;
5593 let mut partial_r = stream
5594 .alloc_zeros::<f64>(plan.partial_r_doubles)
5595 .map_err(|_| super::ArrowSchurGpuFailure::Unavailable)?;
5596 let mut status_dev = stream
5597 .alloc_zeros::<i32>(n)
5598 .map_err(|_| super::ArrowSchurGpuFailure::Unavailable)?;
5599
5600 let cfg = LaunchConfig {
5601 grid_dim: (plan.blocks, 1, 1),
5602 block_dim: (plan.threads_per_block, 1, 1),
5603 shared_mem_bytes: 0,
5604 };
5605 let n_i32 = to_i32(n).ok_or(super::ArrowSchurGpuFailure::Unavailable)?;
5606 let p_i32 = to_i32(d).ok_or(super::ArrowSchurGpuFailure::Unavailable)?;
5607 let r_i32 = to_i32(k).ok_or(super::ArrowSchurGpuFailure::Unavailable)?;
5608 let ridge_arg = ridge_t;
5609 {
5610 let mut builder = stream.launch_builder(&forward);
5611 builder
5612 .arg(&d_dev)
5613 .arg(&b_dev)
5614 .arg(&g_dev)
5615 .arg(&n_i32)
5616 .arg(&p_i32)
5617 .arg(&r_i32)
5618 .arg(&ridge_arg)
5619 .arg(&mut l_out)
5620 .arg(&mut u_out)
5621 .arg(&mut y_out)
5622 .arg(&mut partial_s)
5623 .arg(&mut partial_r)
5624 .arg(&mut status_dev);
5625 unsafe { builder.launch(cfg) }.map_err(|_| super::ArrowSchurGpuFailure::Unavailable)?;
5628 }
5629 stream
5630 .synchronize()
5631 .map_err(|_| super::ArrowSchurGpuFailure::Unavailable)?;
5632
5633 let status_host = stream
5634 .clone_dtoh(&status_dev)
5635 .map_err(|_| super::ArrowSchurGpuFailure::Unavailable)?;
5636 if let Some(row) = status_host.iter().position(|s| *s != 0) {
5637 return Err(super::ArrowSchurGpuFailure::RidgeBumpRequired {
5641 row,
5642 bump: super::ridge_bump_to_make_pd(sys.rows[row].htt.view(), ridge_t),
5643 });
5644 }
5645
5646 let y_host = stream
5648 .clone_dtoh(&y_out)
5649 .map_err(|_| super::ArrowSchurGpuFailure::Unavailable)?;
5650
5651 let hbb_host: Vec<f64> = sys.hbb.iter().copied().collect();
5654 let hbb_is_kk = sys.hbb.dim() == (k, k);
5655 let hbb_matvec_opt = sys.hbb_matvec.clone();
5656
5657 let closure: crate::arrow_schur::GpuSchurMatvec =
5658 Arc::new(move |x: &Array1<f64>, out: &mut Array1<f64>| {
5659 assert_eq!(x.len(), k, "gpu_schur_matvec: x.len() != k");
5660 assert_eq!(out.len(), k, "gpu_schur_matvec: out.len() != k");
5661
5662 if let Some(ref mv) = hbb_matvec_opt {
5664 mv(x.view(), out);
5665 for a in 0..k {
5666 out[a] += ridge_beta * x[a];
5667 }
5668 } else if hbb_is_kk {
5669 for a in 0..k {
5671 let mut acc = ridge_beta * x[a];
5672 for b in 0..k {
5673 acc += hbb_host[a * k + b] * x[b];
5674 }
5675 out[a] = acc;
5676 }
5677 } else {
5678 for a in 0..k {
5679 out[a] = ridge_beta * x[a];
5680 }
5681 }
5682
5683 let mut z = vec![0.0_f64; d];
5686 for i in 0..n {
5687 let y_base = i * p_max * r_template;
5688 for r in 0..d {
5689 let mut acc = 0.0;
5690 for c in 0..k {
5691 acc += y_host[y_base + c * p_max + r] * x[c];
5692 }
5693 z[r] = acc;
5694 }
5695 for c in 0..k {
5696 let mut acc = 0.0;
5697 for r in 0..d {
5698 acc += y_host[y_base + c * p_max + r] * z[r];
5699 }
5700 out[c] -= acc;
5701 }
5702 }
5703 });
5704
5705 Ok(closure)
5706 }
5707
5708 struct DeviceSaeFrameBuffers {
5711 s_off: CudaSlice<i32>,
5713 s_m: CudaSlice<i32>,
5714 s_r: CudaSlice<i32>,
5715 s_ptr: CudaSlice<i32>,
5716 s_data: CudaSlice<f64>,
5717 s_blocks: usize,
5718 g_off_i: CudaSlice<i32>,
5720 g_off_j: CudaSlice<i32>,
5721 g_ri: CudaSlice<i32>,
5722 g_rj: CudaSlice<i32>,
5723 g_mj: CudaSlice<i32>,
5724 g_ptr: CudaSlice<i32>,
5725 g_data: CudaSlice<f64>,
5726 w_ptr: CudaSlice<i32>,
5727 w_data: CudaSlice<f64>,
5728 g_group_ptr: CudaSlice<i32>,
5729 g_group_blocks: CudaSlice<i32>,
5730 g_group_off: CudaSlice<i32>,
5731 g_group_width: CudaSlice<i32>,
5732 g_groups: usize,
5733 g_max_work: usize,
5734 htb_ptr: CudaSlice<i32>,
5736 htb: CudaSlice<f64>,
5737 q_of: CudaSlice<i32>,
5738 ainv: CudaSlice<f64>,
5739 hvec: CudaSlice<f64>,
5740 svec: CudaSlice<f64>,
5741 scatter_partial: CudaSlice<f64>,
5746 n_chunks: usize,
5747 rows_per_chunk: usize,
5748 n_rows: usize,
5749 k: usize,
5750 max_q: usize,
5751 }
5752
5753 fn scatter_chunking(n_rows: usize) -> (usize, usize) {
5759 let target_chunks = 128usize;
5760 let rows_per_chunk = n_rows.div_ceil(target_chunks).max(1);
5761 let n_chunks = n_rows.div_ceil(rows_per_chunk).max(1);
5762 (rows_per_chunk, n_chunks)
5763 }
5764
5765 fn flatten_device_sae_frame_data(
5766 sys: &ArrowSchurSystem,
5767 data: &DeviceSaePcgData,
5768 frame: &DeviceSaeFrameData,
5769 ridge_t: f64,
5770 stream: &Arc<CudaStream>,
5771 ) -> Result<DeviceSaeFrameBuffers, ArrowSchurGpuFailure> {
5772 let host = super::flatten_frame_host_operands(sys, data, frame)?;
5779 let ainv = super::compute_ainv_host(sys, &host.q_of, host.max_q, host.n_rows, ridge_t)?;
5780 upload_frame_buffers(&host, &ainv, stream)
5781 }
5782
5783 fn upload_frame_buffers(
5789 host: &super::FrameHostOperands,
5790 ainv: &[f64],
5791 stream: &Arc<CudaStream>,
5792 ) -> Result<DeviceSaeFrameBuffers, ArrowSchurGpuFailure> {
5793 let htod_i = |v: &[i32]| {
5794 stream
5795 .clone_htod(v)
5796 .map_err(|_| ArrowSchurGpuFailure::Unavailable)
5797 };
5798 let htod_f = |v: &[f64]| {
5799 stream
5800 .clone_htod(v)
5801 .map_err(|_| ArrowSchurGpuFailure::Unavailable)
5802 };
5803 let (rows_per_chunk, n_chunks) = scatter_chunking(host.n_rows);
5804 Ok(DeviceSaeFrameBuffers {
5805 s_off: htod_i(&host.s_off)?,
5806 s_m: htod_i(&host.s_m)?,
5807 s_r: htod_i(&host.s_r)?,
5808 s_ptr: htod_i(&host.s_ptr)?,
5809 s_data: htod_f(&host.s_data)?,
5810 s_blocks: host.s_blocks,
5811 g_off_i: htod_i(&host.g_off_i)?,
5812 g_off_j: htod_i(&host.g_off_j)?,
5813 g_ri: htod_i(&host.g_ri)?,
5814 g_rj: htod_i(&host.g_rj)?,
5815 g_mj: htod_i(&host.g_mj)?,
5816 g_ptr: htod_i(&host.g_ptr)?,
5817 g_data: htod_f(&host.g_data)?,
5818 w_ptr: htod_i(&host.w_ptr)?,
5819 w_data: htod_f(&host.w_data)?,
5820 g_group_ptr: htod_i(&host.g_group_ptr)?,
5821 g_group_blocks: htod_i(&host.g_group_blocks)?,
5822 g_group_off: htod_i(&host.g_group_off)?,
5823 g_group_width: htod_i(&host.g_group_width)?,
5824 g_groups: host.g_groups,
5825 g_max_work: host.g_max_work,
5826 htb_ptr: htod_i(&host.htb_ptr)?,
5827 htb: htod_f(&host.htb)?,
5828 q_of: htod_i(&host.q_of)?,
5829 ainv: htod_f(ainv)?,
5830 hvec: stream
5831 .alloc_zeros::<f64>(host.n_rows * host.max_q)
5832 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
5833 svec: stream
5834 .alloc_zeros::<f64>(host.n_rows * host.max_q)
5835 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
5836 scatter_partial: stream
5837 .alloc_zeros::<f64>(n_chunks * host.k)
5838 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?,
5839 n_chunks,
5840 rows_per_chunk,
5841 n_rows: host.n_rows,
5842 k: host.k,
5843 max_q: host.max_q,
5844 })
5845 }
5846
5847 fn refresh_frame_buffers(
5852 host: &super::FrameHostOperands,
5853 buffers: &mut DeviceSaeFrameBuffers,
5854 stream: &Arc<CudaStream>,
5855 ) -> Result<(), ArrowSchurGpuFailure> {
5856 macro_rules! refresh {
5857 ($host:expr, $device:expr) => {{
5858 if $host.len() != $device.len() {
5859 return Err(ArrowSchurGpuFailure::Unavailable);
5860 }
5861 stream
5862 .memcpy_htod($host, &mut $device)
5863 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5864 }};
5865 }
5866 if host.s_blocks != buffers.s_blocks
5867 || host.g_groups != buffers.g_groups
5868 || host.g_max_work != buffers.g_max_work
5869 || host.n_rows != buffers.n_rows
5870 || host.k != buffers.k
5871 || host.max_q != buffers.max_q
5872 {
5873 return Err(ArrowSchurGpuFailure::Unavailable);
5874 }
5875 refresh!(&host.s_off, buffers.s_off);
5876 refresh!(&host.s_m, buffers.s_m);
5877 refresh!(&host.s_r, buffers.s_r);
5878 refresh!(&host.s_ptr, buffers.s_ptr);
5879 refresh!(&host.s_data, buffers.s_data);
5880 refresh!(&host.g_off_i, buffers.g_off_i);
5881 refresh!(&host.g_off_j, buffers.g_off_j);
5882 refresh!(&host.g_ri, buffers.g_ri);
5883 refresh!(&host.g_rj, buffers.g_rj);
5884 refresh!(&host.g_mj, buffers.g_mj);
5885 refresh!(&host.g_ptr, buffers.g_ptr);
5886 refresh!(&host.g_data, buffers.g_data);
5887 refresh!(&host.w_ptr, buffers.w_ptr);
5888 refresh!(&host.w_data, buffers.w_data);
5889 refresh!(&host.g_group_ptr, buffers.g_group_ptr);
5890 refresh!(&host.g_group_blocks, buffers.g_group_blocks);
5891 refresh!(&host.g_group_off, buffers.g_group_off);
5892 refresh!(&host.g_group_width, buffers.g_group_width);
5893 refresh!(&host.htb_ptr, buffers.htb_ptr);
5894 refresh!(&host.htb, buffers.htb);
5895 refresh!(&host.q_of, buffers.q_of);
5896 Ok(())
5897 }
5898
5899 fn sae_frame_penalty_diag_host(
5900 data: &DeviceSaePcgData,
5901 frame: &DeviceSaeFrameData,
5902 ridge_beta: f64,
5903 ) -> Result<Vec<f64>, ArrowSchurGpuFailure> {
5904 let mut diag = vec![ridge_beta; data.beta_dim];
5905 for (blk, &r) in data.smooth_blocks.iter().zip(frame.smooth_ranks.iter()) {
5907 let m = blk.factor_a.nrows();
5908 for ia in 0..m {
5909 let coeff = blk.factor_a[[ia, ia]];
5910 let base = blk.global_offset + ia * r;
5911 for ib in 0..r {
5912 if base + ib >= diag.len() {
5913 return Err(ArrowSchurGpuFailure::Unavailable);
5914 }
5915 diag[base + ib] += coeff;
5916 }
5917 }
5918 }
5919 for blk in &frame.frame_blocks {
5921 if blk.atom_i != blk.atom_j {
5922 continue;
5923 }
5924 let r = frame.ranks[blk.atom_i];
5925 let off = frame.border_offsets[blk.atom_i];
5926 let (mi, mj) = blk.g.dim();
5927 for li in 0..mi.min(mj) {
5928 let gii = blk.g[[li, li]];
5929 let base = off + li * r;
5930 for a in 0..r {
5931 if base + a >= diag.len() {
5932 return Err(ArrowSchurGpuFailure::Unavailable);
5933 }
5934 diag[base + a] += gii * blk.w[[a, a]];
5935 }
5936 }
5937 }
5938 Ok(diag)
5939 }
5940
5941 fn frame_grid(work: usize, n_rows: usize) -> Result<LaunchConfig, ArrowSchurGpuFailure> {
5942 Ok(LaunchConfig {
5943 grid_dim: (
5944 ((work as u32).saturating_add(255) / 256).max(1),
5945 checked_i32(n_rows)? as u32,
5946 1,
5947 ),
5948 block_dim: (256, 1, 1),
5949 shared_mem_bytes: 0,
5950 })
5951 }
5952
5953 fn launch_sae_frame_matvec(
5954 stream: &Arc<CudaStream>,
5955 module: &Arc<CudaModule>,
5956 buffers: &mut DeviceSaeFrameBuffers,
5957 x: &CudaSlice<f64>,
5958 out: &mut CudaSlice<f64>,
5959 ridge_beta: f64,
5960 ) -> Result<(), ArrowSchurGpuFailure> {
5961 launch_sae_init(stream, module, out, x, ridge_beta, buffers.k)?;
5962 if buffers.s_blocks > 0 {
5964 let kernel = module
5965 .load_function("arrow_sae_frame_smooth_matvec")
5966 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5967 let blocks_i32 = checked_i32(buffers.s_blocks)?;
5968 let cfg = frame_grid(buffers.k, buffers.s_blocks)?;
5969 let mut b = stream.launch_builder(&kernel);
5970 b.arg(x)
5971 .arg(&mut *out)
5972 .arg(&buffers.s_off)
5973 .arg(&buffers.s_m)
5974 .arg(&buffers.s_r)
5975 .arg(&buffers.s_ptr)
5976 .arg(&buffers.s_data)
5977 .arg(&blocks_i32);
5978 unsafe { b.launch(cfg) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5981 }
5982 if buffers.g_groups > 0 {
5984 let kernel = module
5985 .load_function("arrow_sae_frame_g_matvec")
5986 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
5987 let groups_i32 = checked_i32(buffers.g_groups)?;
5988 let cfg = frame_grid(buffers.g_max_work.max(1), buffers.g_groups)?;
5989 let mut b = stream.launch_builder(&kernel);
5990 b.arg(x)
5991 .arg(&mut *out)
5992 .arg(&buffers.g_off_i)
5993 .arg(&buffers.g_off_j)
5994 .arg(&buffers.g_ri)
5995 .arg(&buffers.g_rj)
5996 .arg(&buffers.g_mj)
5997 .arg(&buffers.g_ptr)
5998 .arg(&buffers.g_data)
5999 .arg(&buffers.w_ptr)
6000 .arg(&buffers.w_data)
6001 .arg(&buffers.g_group_ptr)
6002 .arg(&buffers.g_group_blocks)
6003 .arg(&buffers.g_group_off)
6004 .arg(&buffers.g_group_width)
6005 .arg(&groups_i32);
6006 unsafe { b.launch(cfg) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6010 }
6011 let k_i32 = checked_i32(buffers.k)?;
6013 let max_q_i32 = checked_i32(buffers.max_q)?;
6014 let n_rows_i32 = checked_i32(buffers.n_rows)?;
6015 {
6016 let kernel = module
6017 .load_function("arrow_sae_frame_apply_h")
6018 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6019 let cfg = frame_grid(buffers.max_q, buffers.n_rows)?;
6020 let mut b = stream.launch_builder(&kernel);
6021 b.arg(x)
6022 .arg(&buffers.htb_ptr)
6023 .arg(&buffers.htb)
6024 .arg(&buffers.q_of)
6025 .arg(&mut buffers.hvec)
6026 .arg(&k_i32)
6027 .arg(&max_q_i32)
6028 .arg(&n_rows_i32);
6029 unsafe { b.launch(cfg) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6032 }
6033 {
6034 let kernel = module
6035 .load_function("arrow_sae_frame_apply_ainv")
6036 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6037 let cfg = frame_grid(buffers.max_q, buffers.n_rows)?;
6038 let mut b = stream.launch_builder(&kernel);
6039 b.arg(&buffers.ainv)
6040 .arg(&buffers.hvec)
6041 .arg(&buffers.q_of)
6042 .arg(&mut buffers.svec)
6043 .arg(&max_q_i32)
6044 .arg(&n_rows_i32);
6045 unsafe { b.launch(cfg) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6048 }
6049 let k_blocks = ((buffers.k as u32).saturating_add(255) / 256).max(1);
6061 {
6062 let kernel = module
6063 .load_function("arrow_sae_frame_scatter_h_det_partial")
6064 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6065 let rows_per_chunk_i32 = checked_i32(buffers.rows_per_chunk)?;
6066 let cfg = LaunchConfig {
6067 grid_dim: (k_blocks, checked_i32(buffers.n_chunks)? as u32, 1),
6068 block_dim: (256, 1, 1),
6069 shared_mem_bytes: 0,
6070 };
6071 let mut b = stream.launch_builder(&kernel);
6072 b.arg(&buffers.svec)
6073 .arg(&buffers.htb_ptr)
6074 .arg(&buffers.htb)
6075 .arg(&buffers.q_of)
6076 .arg(&mut buffers.scatter_partial)
6077 .arg(&k_i32)
6078 .arg(&max_q_i32)
6079 .arg(&n_rows_i32)
6080 .arg(&rows_per_chunk_i32);
6081 unsafe { b.launch(cfg) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6084 }
6085 {
6086 let kernel = module
6087 .load_function("arrow_sae_frame_scatter_h_det_reduce_sub")
6088 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6089 let n_chunks_i32 = checked_i32(buffers.n_chunks)?;
6090 let cfg = LaunchConfig {
6091 grid_dim: (k_blocks, 1, 1),
6092 block_dim: (256, 1, 1),
6093 shared_mem_bytes: 0,
6094 };
6095 let mut b = stream.launch_builder(&kernel);
6096 b.arg(&buffers.scatter_partial)
6097 .arg(&mut *out)
6098 .arg(&k_i32)
6099 .arg(&n_chunks_i32);
6100 unsafe { b.launch(cfg) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6103 }
6104 Ok(())
6105 }
6106
6107 fn launch_sae_frame_reduced_schur_det(
6117 stream: &Arc<CudaStream>,
6118 module: &Arc<CudaModule>,
6119 buffers: &mut DeviceSaeFrameBuffers,
6120 x: &CudaSlice<f64>,
6121 out: &mut CudaSlice<f64>,
6122 ) -> Result<(), ArrowSchurGpuFailure> {
6123 let k_i32 = checked_i32(buffers.k)?;
6124 let max_q_i32 = checked_i32(buffers.max_q)?;
6125 let n_rows_i32 = checked_i32(buffers.n_rows)?;
6126 {
6130 let kernel = module
6131 .load_function("arrow_sae_frame_apply_h_warp")
6132 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6133 let block = checked_i32(buffers.max_q)?
6134 .checked_mul(32)
6135 .ok_or(ArrowSchurGpuFailure::Unavailable)? as u32;
6136 let cfg = LaunchConfig {
6137 grid_dim: (checked_i32(buffers.n_rows)? as u32, 1, 1),
6138 block_dim: (block, 1, 1),
6139 shared_mem_bytes: 0,
6140 };
6141 let mut b = stream.launch_builder(&kernel);
6142 b.arg(x)
6143 .arg(&buffers.htb_ptr)
6144 .arg(&buffers.htb)
6145 .arg(&buffers.q_of)
6146 .arg(&mut buffers.hvec)
6147 .arg(&k_i32)
6148 .arg(&max_q_i32)
6149 .arg(&n_rows_i32);
6150 unsafe { b.launch(cfg) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6154 }
6155 {
6157 let kernel = module
6158 .load_function("arrow_sae_frame_apply_ainv")
6159 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6160 let cfg = frame_grid(buffers.max_q, buffers.n_rows)?;
6161 let mut b = stream.launch_builder(&kernel);
6162 b.arg(&buffers.ainv)
6163 .arg(&buffers.hvec)
6164 .arg(&buffers.q_of)
6165 .arg(&mut buffers.svec)
6166 .arg(&max_q_i32)
6167 .arg(&n_rows_i32);
6168 unsafe { b.launch(cfg) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6171 }
6172 let k_blocks = ((buffers.k as u32).saturating_add(255) / 256).max(1);
6177 {
6178 let kernel = module
6179 .load_function("arrow_sae_frame_scatter_h_det_partial")
6180 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6181 let rows_per_chunk_i32 = checked_i32(buffers.rows_per_chunk)?;
6182 let cfg = LaunchConfig {
6183 grid_dim: (k_blocks, checked_i32(buffers.n_chunks)? as u32, 1),
6184 block_dim: (256, 1, 1),
6185 shared_mem_bytes: 0,
6186 };
6187 let mut b = stream.launch_builder(&kernel);
6188 b.arg(&buffers.svec)
6189 .arg(&buffers.htb_ptr)
6190 .arg(&buffers.htb)
6191 .arg(&buffers.q_of)
6192 .arg(&mut buffers.scatter_partial)
6193 .arg(&k_i32)
6194 .arg(&max_q_i32)
6195 .arg(&n_rows_i32)
6196 .arg(&rows_per_chunk_i32);
6197 unsafe { b.launch(cfg) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6200 }
6201 {
6204 let kernel = module
6205 .load_function("arrow_sae_frame_scatter_h_det_reduce")
6206 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6207 let n_chunks_i32 = checked_i32(buffers.n_chunks)?;
6208 let cfg = LaunchConfig {
6209 grid_dim: (k_blocks, 1, 1),
6210 block_dim: (256, 1, 1),
6211 shared_mem_bytes: 0,
6212 };
6213 let mut b = stream.launch_builder(&kernel);
6214 b.arg(&buffers.scatter_partial)
6215 .arg(&mut *out)
6216 .arg(&k_i32)
6217 .arg(&n_chunks_i32);
6218 unsafe { b.launch(cfg) }.map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6221 }
6222 Ok(())
6223 }
6224
6225 fn launch_sae_frame_diag_sub(
6226 stream: &Arc<CudaStream>,
6227 module: &Arc<CudaModule>,
6228 buffers: &DeviceSaeFrameBuffers,
6229 diag: &mut CudaSlice<f64>,
6230 ) -> Result<(), ArrowSchurGpuFailure> {
6231 let kernel = module
6232 .load_function("arrow_sae_frame_diag_sub")
6233 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6234 let k_i32 = checked_i32(buffers.k)?;
6235 let max_q_i32 = checked_i32(buffers.max_q)?;
6236 let n_rows_i32 = checked_i32(buffers.n_rows)?;
6237 let cfg = frame_grid(buffers.k, 1)?;
6238 let mut b = stream.launch_builder(&kernel);
6239 b.arg(diag)
6240 .arg(&buffers.ainv)
6241 .arg(&buffers.htb_ptr)
6242 .arg(&buffers.htb)
6243 .arg(&buffers.q_of)
6244 .arg(&k_i32)
6245 .arg(&max_q_i32)
6246 .arg(&n_rows_i32);
6247 unsafe { b.launch(cfg) }
6250 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6251 Ok(())
6252 }
6253
6254 pub(super) fn solve_sae_matrix_free_pcg_framed(
6255 sys: &ArrowSchurSystem,
6256 data: &DeviceSaePcgData,
6257 ridge_t: f64,
6258 ridge_beta: f64,
6259 rhs_beta: &Array1<f64>,
6260 max_iterations: usize,
6261 relative_tolerance: f64,
6262 ) -> Result<(Array1<f64>, ArrowPcgDiagnostics), ArrowSchurGpuFailure> {
6263 let k = rhs_beta.len();
6264 if k == 0 || data.beta_dim != k || sys.k != k {
6265 return Err(ArrowSchurGpuFailure::Unavailable);
6266 }
6267 let frame = data
6268 .frame
6269 .as_ref()
6270 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
6271 let runtime = super::resolve_runtime_for_device_path()?
6272 .filter(|rt| {
6273 rt.policy().reduced_schur_matvec_should_offload(
6274 sys.rows.len(),
6275 sys.k,
6276 sys.d,
6277 max_iterations,
6278 )
6279 })
6280 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
6281 let ctx = gam_gpu::device_runtime::cuda_context_for(runtime.selected_device().ordinal)
6282 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
6283 let stream = ctx
6284 .new_stream()
6285 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6286 let blas = CudaBlas::new(stream.clone()).map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6287 let vector_module = pcg_vector_module(&ctx)?;
6288 log::info!(
6293 "#1017/#2230 {} ridge_t={ridge_t:e} ridge_beta={ridge_beta:e}",
6294 data.operand_byte_report()
6295 );
6296 let mut buffers = flatten_device_sae_frame_data(sys, data, frame, ridge_t, &stream)?;
6297 let dctx = FramedPcgCtx {
6298 stream: &stream,
6299 blas: &blas,
6300 module: vector_module,
6301 max_iterations,
6302 relative_tolerance,
6303 };
6304 pcg_solve_framed_body(data, frame, &mut buffers, ridge_beta, rhs_beta, &dctx)
6305 }
6306
6307 struct FramedPcgCtx<'a> {
6310 stream: &'a Arc<CudaStream>,
6311 blas: &'a CudaBlas,
6312 module: &'a Arc<CudaModule>,
6313 max_iterations: usize,
6314 relative_tolerance: f64,
6315 }
6316
6317 fn pcg_solve_framed_body(
6325 data: &DeviceSaePcgData,
6326 frame: &DeviceSaeFrameData,
6327 buffers: &mut DeviceSaeFrameBuffers,
6328 ridge_beta: f64,
6329 rhs_beta: &Array1<f64>,
6330 dctx: &FramedPcgCtx<'_>,
6331 ) -> Result<(Array1<f64>, ArrowPcgDiagnostics), ArrowSchurGpuFailure> {
6332 let stream = dctx.stream;
6333 let blas = dctx.blas;
6334 let vector_module = dctx.module;
6335 let max_iterations = dctx.max_iterations;
6336 let relative_tolerance = dctx.relative_tolerance;
6337 let k = rhs_beta.len();
6338 let rhs_norm = rhs_beta.iter().map(|v| v * v).sum::<f64>().sqrt();
6339 if rhs_norm == 0.0 {
6340 return Ok((Array1::<f64>::zeros(k), ArrowPcgDiagnostics::default()));
6341 }
6342 let tol = (relative_tolerance.max(0.0) * rhs_norm).max(1e-12);
6343 let rhs_dev = stream
6344 .clone_htod(
6345 rhs_beta
6346 .as_slice()
6347 .ok_or(ArrowSchurGpuFailure::Unavailable)?,
6348 )
6349 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6350 let diag_host = sae_frame_penalty_diag_host(data, frame, ridge_beta)?;
6351 let mut diag_dev = stream
6352 .clone_htod(&diag_host)
6353 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6354 launch_sae_frame_diag_sub(stream, vector_module, buffers, &mut diag_dev)?;
6355 let diag_host = stream
6356 .clone_dtoh(&diag_dev)
6357 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6358 let mut inv_diag = Vec::with_capacity(k);
6359 for (idx, &d) in diag_host.iter().enumerate() {
6360 if !d.is_finite() || d <= 1.0e-18 {
6361 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
6362 reason: format!(
6363 "framed SAE GPU PCG: non-positive Jacobi diagonal at {idx}: {d:e}"
6364 ),
6365 });
6366 }
6367 inv_diag.push(1.0 / d);
6368 }
6369 let inv_diag_dev = stream
6370 .clone_htod(&inv_diag)
6371 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6372
6373 let mut x_dev = stream
6374 .alloc_zeros::<f64>(k)
6375 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6376 let mut r_dev = stream
6377 .alloc_zeros::<f64>(k)
6378 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6379 device_copy(blas, stream, k, &rhs_dev, &mut r_dev)?;
6380 let mut z_dev = stream
6381 .alloc_zeros::<f64>(k)
6382 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6383 launch_jacobi_mul(stream, vector_module, &inv_diag_dev, &r_dev, &mut z_dev, k)?;
6384 let mut p_dev = stream
6385 .alloc_zeros::<f64>(k)
6386 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6387 device_copy(blas, stream, k, &z_dev, &mut p_dev)?;
6388 let mut ap_dev = stream
6389 .alloc_zeros::<f64>(k)
6390 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6391
6392 let mut rz = device_dot(blas, stream, k, &r_dev, &z_dev)?;
6393 if rz <= 0.0 || !rz.is_finite() {
6394 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
6395 reason: format!("framed SAE GPU PCG: non-positive initial rᵀM⁻¹r={rz:e}"),
6396 });
6397 }
6398 let mut diag = ArrowPcgDiagnostics {
6399 precond_apply_calls: 1,
6400 stopping_reason: PcgStopReason::MaxIter,
6401 ..ArrowPcgDiagnostics::default()
6402 };
6403 for _ in 0..max_iterations.max(1) {
6404 launch_sae_frame_matvec(
6405 stream,
6406 vector_module,
6407 buffers,
6408 &p_dev,
6409 &mut ap_dev,
6410 ridge_beta,
6411 )?;
6412 diag.matvec_calls += 1;
6413 diag.iterations += 1;
6414 let pap = device_dot(blas, stream, k, &p_dev, &ap_dev)?;
6415 if pap <= 0.0 || !pap.is_finite() {
6416 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
6417 reason: format!("framed SAE GPU PCG: non-positive curvature pᵀAp={pap:e}"),
6418 });
6419 }
6420 let alpha = rz / pap;
6421 device_axpy(blas, stream, k, alpha, &p_dev, &mut x_dev)?;
6422 device_axpy(blas, stream, k, -alpha, &ap_dev, &mut r_dev)?;
6423 let r_norm = device_nrm2(blas, stream, k, &r_dev)?;
6424 if r_norm <= tol {
6425 diag.final_relative_residual = r_norm / rhs_norm;
6426 diag.stopping_reason = PcgStopReason::Converged;
6427 break;
6428 }
6429 launch_jacobi_mul(stream, vector_module, &inv_diag_dev, &r_dev, &mut z_dev, k)?;
6430 diag.precond_apply_calls += 1;
6431 let rz_new = device_dot(blas, stream, k, &r_dev, &z_dev)?;
6432 if rz_new <= 0.0 || !rz_new.is_finite() {
6433 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
6434 reason: format!("framed SAE GPU PCG: non-positive rᵀM⁻¹r={rz_new:e}"),
6435 });
6436 }
6437 let beta = rz_new / rz;
6438 launch_update_p(stream, vector_module, &z_dev, beta, &mut p_dev, k)?;
6439 rz = rz_new;
6440 }
6441 if diag.stopping_reason != PcgStopReason::Converged {
6442 let r_norm = device_nrm2(blas, stream, k, &r_dev)?;
6443 diag.final_relative_residual = r_norm / rhs_norm;
6444 diag.stopping_reason = PcgStopReason::MaxIter;
6445 }
6446 let x = stream
6447 .clone_dtoh(&x_dev)
6448 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6449 Ok((Array1::from_vec(x), diag))
6450 }
6451
6452 pub(crate) struct ResidentSaeFrameHandle {
6461 ctx: Arc<CudaContext>,
6462 stream: Arc<CudaStream>,
6463 buffers: std::sync::Mutex<DeviceSaeFrameBuffers>,
6464 q_of: Vec<i32>,
6465 max_q: usize,
6466 n_rows: usize,
6467 k: usize,
6468 }
6469
6470 impl ResidentSaeFrameHandle {
6471 pub(crate) fn build(
6479 sys: &ArrowSchurSystem,
6480 cg_iters: usize,
6481 ) -> Result<Option<Self>, ArrowSchurGpuFailure> {
6482 let Some(data) = sys.device_sae_pcg.as_ref() else {
6483 return Ok(None);
6484 };
6485 let Some(frame) = data.frame.as_ref() else {
6486 return Ok(None);
6487 };
6488 if sys.k == 0 || data.beta_dim != sys.k {
6489 return Ok(None);
6490 }
6491 let Some(runtime) = super::resolve_runtime_for_device_path()? else {
6492 return Ok(None);
6493 };
6494 if !runtime.policy().reduced_schur_matvec_should_offload(
6495 sys.rows.len(),
6496 sys.k,
6497 sys.d,
6498 cg_iters,
6499 ) {
6500 return Ok(None);
6501 }
6502 let ctx = gam_gpu::device_runtime::cuda_context_for(runtime.selected_device().ordinal)
6503 .ok_or_else(|| ArrowSchurGpuFailure::SchurFactorFailed {
6504 reason: "resident SAE frame could not bind the admitted CUDA context"
6505 .to_string(),
6506 })?;
6507 let stream = ctx.new_stream().map_err(|error| {
6508 ArrowSchurGpuFailure::SchurFactorFailed {
6509 reason: format!("resident SAE frame stream creation failed: {error}"),
6510 }
6511 })?;
6512 let host = super::flatten_frame_host_operands(sys, data, frame)?;
6513 log::info!(
6520 "#1017 SAE resident frame ENGAGED: {} uploaded ONCE for the ladder; \
6521 per-trial re-upload now only ainv ({}rows × {}²·8B)",
6522 data.operand_byte_report(),
6523 host.n_rows,
6524 host.max_q
6525 );
6526 let zero_ainv = vec![0.0_f64; host.n_rows * host.max_q * host.max_q];
6527 let buffers = upload_frame_buffers(&host, &zero_ainv, &stream)?;
6528 Ok(Some(Self {
6529 ctx,
6530 stream,
6531 buffers: std::sync::Mutex::new(buffers),
6532 q_of: host.q_of,
6533 max_q: host.max_q,
6534 n_rows: host.n_rows,
6535 k: host.k,
6536 }))
6537 }
6538
6539 pub(super) fn prime_ainv(
6546 &self,
6547 sys: &ArrowSchurSystem,
6548 ridge_t: f64,
6549 ) -> Result<(), ArrowSchurGpuFailure> {
6550 if sys.k != self.k || sys.rows.len() != self.n_rows {
6551 return Err(ArrowSchurGpuFailure::Unavailable);
6552 }
6553 let ainv =
6554 super::compute_ainv_host(sys, &self.q_of, self.max_q, self.n_rows, ridge_t)?;
6555 let mut buffers = self
6556 .buffers
6557 .lock()
6558 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6559 if ainv.len() != buffers.ainv.len() {
6560 return Err(ArrowSchurGpuFailure::Unavailable);
6561 }
6562 self.stream
6563 .memcpy_htod(&ainv, &mut buffers.ainv)
6564 .map_err(|_| ArrowSchurGpuFailure::Unavailable)
6565 }
6566 }
6567
6568 impl super::SaeResidentFrame for ResidentSaeFrameHandle {
6569 fn refresh(&self, sys: &ArrowSchurSystem) -> Result<(), ArrowSchurGpuFailure> {
6570 if sys.k != self.k || sys.rows.len() != self.n_rows {
6571 return Err(ArrowSchurGpuFailure::Unavailable);
6572 }
6573 let data = sys
6574 .device_sae_pcg
6575 .as_ref()
6576 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
6577 let frame = data
6578 .frame
6579 .as_ref()
6580 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
6581 let host = super::flatten_frame_host_operands(sys, data, frame)?;
6582 if host.q_of != self.q_of || host.max_q != self.max_q {
6583 return Err(ArrowSchurGpuFailure::Unavailable);
6584 }
6585 let mut buffers = self
6586 .buffers
6587 .lock()
6588 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6589 refresh_frame_buffers(&host, &mut buffers, &self.stream)
6590 }
6591
6592 fn resolve(
6593 &self,
6594 sys: &ArrowSchurSystem,
6595 ridge_t: f64,
6596 ridge_beta: f64,
6597 rhs_beta: &Array1<f64>,
6598 max_iterations: usize,
6599 relative_tolerance: f64,
6600 ) -> Result<(Array1<f64>, ArrowPcgDiagnostics), ArrowSchurGpuFailure> {
6601 if sys.k != self.k || sys.rows.len() != self.n_rows || rhs_beta.len() != self.k {
6605 return Err(ArrowSchurGpuFailure::Unavailable);
6606 }
6607 let data = sys
6608 .device_sae_pcg
6609 .as_ref()
6610 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
6611 let frame = data
6612 .frame
6613 .as_ref()
6614 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
6615 let ainv = super::compute_ainv_host(sys, &self.q_of, self.max_q, self.n_rows, ridge_t)?;
6619 let mut buffers = self
6620 .buffers
6621 .lock()
6622 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6623 if ainv.len() != buffers.ainv.len() {
6624 return Err(ArrowSchurGpuFailure::Unavailable);
6625 }
6626 self.stream
6627 .memcpy_htod(&ainv, &mut buffers.ainv)
6628 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6629 let blas = CudaBlas::new(self.stream.clone())
6630 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6631 let vector_module = pcg_vector_module(&self.ctx)?;
6632 let dctx = FramedPcgCtx {
6633 stream: &self.stream,
6634 blas: &blas,
6635 module: vector_module,
6636 max_iterations,
6637 relative_tolerance,
6638 };
6639 pcg_solve_framed_body(data, frame, &mut *buffers, ridge_beta, rhs_beta, &dctx)
6640 }
6641 }
6642
6643 pub(super) fn build_framed_resident_evidence_matvec(
6655 sys: &ArrowSchurSystem,
6656 ridge_t: f64,
6657 ridge_beta: f64,
6658 apply_budget: usize,
6659 ) -> Result<Option<crate::arrow_schur::GpuSchurMatvec>, ArrowSchurGpuFailure> {
6660 let Some(handle) = ResidentSaeFrameHandle::build(sys, apply_budget)? else {
6661 return Ok(None);
6662 };
6663 handle.prime_ainv(sys, ridge_t)?;
6664 let Some(data) = sys.device_sae_pcg.as_ref().cloned() else {
6665 return Ok(None);
6666 };
6667 let handle = Arc::new(handle);
6668 let k = handle.k;
6669 let closure: crate::arrow_schur::GpuSchurMatvec =
6670 Arc::new(move |x: &Array1<f64>, out: &mut Array1<f64>| {
6671 assert_eq!(x.len(), k, "#1017 framed evidence matvec: x.len() != k");
6672 assert_eq!(out.len(), k, "#1017 framed evidence matvec: out.len() != k");
6673 let module = pcg_vector_module(&handle.ctx)
6674 .expect("#1017 framed evidence matvec: pcg_vector_module unavailable");
6675 let x_slice = x
6676 .as_slice()
6677 .expect("#1017 framed evidence matvec: x not contiguous");
6678 let x_dev = handle
6679 .stream
6680 .clone_htod(x_slice)
6681 .expect("#1017 framed evidence matvec: htod(x) failed");
6682 let mut reduced_dev = handle
6683 .stream
6684 .alloc_zeros::<f64>(k)
6685 .expect("#1017 framed evidence matvec: device alloc failed");
6686 {
6687 let mut buffers = handle
6688 .buffers
6689 .lock()
6690 .expect("#1017 framed evidence matvec: resident frame poisoned");
6691 launch_sae_frame_reduced_schur_det(
6692 &handle.stream,
6693 module,
6694 &mut buffers,
6695 &x_dev,
6696 &mut reduced_dev,
6697 )
6698 .expect("#1017 framed evidence matvec: device reduced-Schur launch failed");
6699 }
6700 let reduced = handle
6701 .stream
6702 .clone_dtoh(&reduced_dev)
6703 .expect("#1017 framed evidence matvec: dtoh(out) failed");
6704 let out_slice = out
6705 .as_slice_mut()
6706 .expect("#1017 framed evidence matvec: out not contiguous");
6707 super::sae_framed_penalty_matvec_cpu(&data, ridge_beta, x_slice, out_slice);
6708 for a in 0..k {
6709 out_slice[a] += reduced[a];
6710 }
6711 });
6712 Ok(Some(closure))
6713 }
6714
6715 pub(super) fn solve_sae_matrix_free_pcg(
6722 sys: &ArrowSchurSystem,
6723 data: &DeviceSaePcgData,
6724 ridge_t: f64,
6725 ridge_beta: f64,
6726 rhs_beta: &Array1<f64>,
6727 max_iterations: usize,
6728 relative_tolerance: f64,
6729 ) -> Result<(Array1<f64>, ArrowPcgDiagnostics), ArrowSchurGpuFailure> {
6730 let k = rhs_beta.len();
6731 if k == 0 || data.beta_dim != k || sys.k != k {
6732 return Err(ArrowSchurGpuFailure::Unavailable);
6733 }
6734 if data.sparse_g_blocks.is_empty() && data.smooth_blocks.is_empty() && data.frame.is_none() {
6740 return Err(ArrowSchurGpuFailure::Unavailable);
6741 }
6742 if data.frame.is_some() {
6746 return Err(ArrowSchurGpuFailure::Unavailable);
6747 }
6748 let runtime = super::resolve_runtime_for_device_path()?
6762 .filter(|rt| {
6763 rt.policy().reduced_schur_matvec_should_offload(
6764 sys.rows.len(),
6765 sys.k,
6766 sys.d,
6767 max_iterations,
6768 )
6769 })
6770 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
6771 let ctx = gam_gpu::device_runtime::cuda_context_for(runtime.selected_device().ordinal)
6772 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
6773 let stream = ctx
6774 .new_stream()
6775 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6776 let blas = CudaBlas::new(stream.clone()).map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6777 let vector_module = pcg_vector_module(&ctx)?;
6778 log::info!(
6782 "#1017/#2230 {} ridge_t={ridge_t:e} ridge_beta={ridge_beta:e}",
6783 data.operand_byte_report()
6784 );
6785 let mut buffers = flatten_device_sae_data(sys, data, ridge_t, &stream)?;
6786
6787 let rhs_norm = rhs_beta.iter().map(|v| v * v).sum::<f64>().sqrt();
6788 if rhs_norm == 0.0 {
6789 return Ok((Array1::<f64>::zeros(k), ArrowPcgDiagnostics::default()));
6790 }
6791 let tol = (relative_tolerance.max(0.0) * rhs_norm).max(1e-12);
6792 let rhs_dev = stream
6793 .clone_htod(
6794 rhs_beta
6795 .as_slice()
6796 .ok_or(ArrowSchurGpuFailure::Unavailable)?,
6797 )
6798 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6799 let diag_host = sae_penalty_diag_host(data, ridge_beta)?;
6800 let mut diag_dev = stream
6801 .clone_htod(&diag_host)
6802 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6803 launch_sae_diag_sub(&stream, vector_module, &mut buffers, &mut diag_dev)?;
6804 let diag_host = stream
6805 .clone_dtoh(&diag_dev)
6806 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6807 let mut inv_diag = Vec::with_capacity(k);
6808 for (idx, &d) in diag_host.iter().enumerate() {
6809 if !d.is_finite() || d <= 1.0e-18 {
6810 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
6811 reason: format!(
6812 "SAE matrix-free GPU PCG: non-positive Schur Jacobi diagonal at {idx}: {d:e}"
6813 ),
6814 });
6815 }
6816 inv_diag.push(1.0 / d);
6817 }
6818 let inv_diag_dev = stream
6819 .clone_htod(&inv_diag)
6820 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6821
6822 let mut x_dev = stream
6823 .alloc_zeros::<f64>(k)
6824 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6825 let mut r_dev = stream
6826 .alloc_zeros::<f64>(k)
6827 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6828 device_copy(&blas, &stream, k, &rhs_dev, &mut r_dev)?;
6829 let mut z_dev = stream
6830 .alloc_zeros::<f64>(k)
6831 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6832 launch_jacobi_mul(&stream, vector_module, &inv_diag_dev, &r_dev, &mut z_dev, k)?;
6833 let mut p_dev = stream
6834 .alloc_zeros::<f64>(k)
6835 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6836 device_copy(&blas, &stream, k, &z_dev, &mut p_dev)?;
6837 let mut ap_dev = stream
6838 .alloc_zeros::<f64>(k)
6839 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6840
6841 let mut rz = device_dot(&blas, &stream, k, &r_dev, &z_dev)?;
6842 if rz <= 0.0 || !rz.is_finite() {
6843 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
6844 reason: format!("SAE matrix-free GPU PCG: non-positive initial rᵀM⁻¹r={rz:e}"),
6845 });
6846 }
6847 let mut diag = ArrowPcgDiagnostics {
6848 precond_apply_calls: 1,
6849 stopping_reason: PcgStopReason::MaxIter,
6850 ..ArrowPcgDiagnostics::default()
6851 };
6852
6853 for _ in 0..max_iterations.max(1) {
6854 launch_sae_matvec(
6855 &stream,
6856 vector_module,
6857 &mut buffers,
6858 &p_dev,
6859 &mut ap_dev,
6860 ridge_beta,
6861 )?;
6862 diag.matvec_calls += 1;
6863 diag.iterations += 1;
6864 let pap = device_dot(&blas, &stream, k, &p_dev, &ap_dev)?;
6865 if pap <= 0.0 || !pap.is_finite() {
6866 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
6867 reason: format!("SAE matrix-free GPU PCG: non-positive curvature pᵀAp={pap:e}"),
6868 });
6869 }
6870 let alpha = rz / pap;
6871 device_axpy(&blas, &stream, k, alpha, &p_dev, &mut x_dev)?;
6872 device_axpy(&blas, &stream, k, -alpha, &ap_dev, &mut r_dev)?;
6873 let r_norm = device_nrm2(&blas, &stream, k, &r_dev)?;
6874 if r_norm <= tol {
6875 diag.final_relative_residual = r_norm / rhs_norm;
6876 diag.stopping_reason = PcgStopReason::Converged;
6877 break;
6878 }
6879 launch_jacobi_mul(&stream, vector_module, &inv_diag_dev, &r_dev, &mut z_dev, k)?;
6880 diag.precond_apply_calls += 1;
6881 let rz_new = device_dot(&blas, &stream, k, &r_dev, &z_dev)?;
6882 if rz_new <= 0.0 || !rz_new.is_finite() {
6883 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
6884 reason: format!("SAE matrix-free GPU PCG: non-positive rᵀM⁻¹r={rz_new:e}"),
6885 });
6886 }
6887 let beta = rz_new / rz;
6888 launch_update_p(&stream, vector_module, &z_dev, beta, &mut p_dev, k)?;
6889 rz = rz_new;
6890 }
6891 if diag.stopping_reason != PcgStopReason::Converged {
6892 let r_norm = device_nrm2(&blas, &stream, k, &r_dev)?;
6893 diag.final_relative_residual = r_norm / rhs_norm;
6894 diag.stopping_reason = PcgStopReason::MaxIter;
6895 }
6896 let x = stream
6897 .clone_dtoh(&x_dev)
6898 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6899 Ok((Array1::from_vec(x), diag))
6900 }
6901
6902 pub(super) fn solve_reduced_beta_pcg_with_diagnostics(
6903 s_acc: &ndarray::Array2<f64>,
6904 rhs_beta: &Array1<f64>,
6905 max_iterations: usize,
6906 relative_tolerance: f64,
6907 ) -> Result<(Array1<f64>, ArrowPcgDiagnostics), ArrowSchurGpuFailure> {
6908 let k = rhs_beta.len();
6909 let cg_iters = max_iterations.max(1);
6921 let runtime = gam_gpu::linalg_dispatch::route_through_gpu(
6922 gam_gpu::linalg_dispatch::DispatchOp::Gemm {
6923 m: k,
6924 n: k,
6925 k: cg_iters,
6926 },
6927 )
6928 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
6929 let stream = gam_gpu::device_runtime::cuda_context_for(runtime.device.ordinal)
6930 .and_then(|ctx| ctx.new_stream().ok())
6931 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
6932 let blas = CudaBlas::new(stream.clone()).map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6933 let ctx = gam_gpu::device_runtime::cuda_context_for(runtime.device.ordinal)
6934 .ok_or(ArrowSchurGpuFailure::Unavailable)?;
6935 let vector_module = pcg_vector_module(&ctx)?;
6936
6937 let mut inv_diag = vec![0.0_f64; k];
6939 for j in 0..k {
6940 let djj = s_acc[[j, j]];
6941 if !(djj > 0.0) {
6942 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
6943 reason: format!(
6944 "reduced-β GPU PCG: Jacobi diagonal S[{j},{j}]={djj:e} not positive"
6945 ),
6946 });
6947 }
6948 inv_diag[j] = 1.0 / djj;
6949 }
6950
6951 let mut s_host = vec![0.0_f64; k * k];
6953 for col in 0..k {
6954 for row in 0..k {
6955 s_host[col * k + row] = s_acc[[row, col]];
6956 }
6957 }
6958 let s_dev = stream
6959 .clone_htod(&s_host)
6960 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6961
6962 let rhs_norm = rhs_beta.iter().map(|v| v * v).sum::<f64>().sqrt();
6966 if rhs_norm == 0.0 {
6967 return Ok((Array1::<f64>::zeros(k), ArrowPcgDiagnostics::default()));
6968 }
6969 let tol = (relative_tolerance.max(0.0) * rhs_norm).max(1e-12);
6970
6971 let mut x_dev = stream
6974 .alloc_zeros::<f64>(k)
6975 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6976 let mut r_dev = stream
6977 .clone_htod(
6978 rhs_beta
6979 .as_slice()
6980 .ok_or(ArrowSchurGpuFailure::Unavailable)?,
6981 )
6982 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6983 let inv_diag_dev = stream
6984 .clone_htod(&inv_diag)
6985 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6986 let mut z_dev = stream
6987 .alloc_zeros::<f64>(k)
6988 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6989 launch_jacobi_mul(&stream, vector_module, &inv_diag_dev, &r_dev, &mut z_dev, k)?;
6990 let mut p_dev = stream
6991 .alloc_zeros::<f64>(k)
6992 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6993 device_copy(&blas, &stream, k, &z_dev, &mut p_dev)?;
6994 let mut sp_dev = stream
6995 .alloc_zeros::<f64>(k)
6996 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
6997 let mut rz = device_dot(&blas, &stream, k, &r_dev, &z_dev)?;
6998 let mut diag = ArrowPcgDiagnostics {
6999 precond_apply_calls: 1,
7000 stopping_reason: PcgStopReason::MaxIter,
7001 ..ArrowPcgDiagnostics::default()
7002 };
7003 if rz <= 0.0 || !rz.is_finite() {
7004 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
7005 reason: format!("reduced-β GPU PCG: non-positive initial rᵀM⁻¹r={rz:e}"),
7006 });
7007 }
7008
7009 let max_iters = max_iterations.max(1);
7010 for _ in 0..max_iters {
7011 let gemv_cfg = GemvConfig::<f64> {
7013 trans: cublasOperation_t::CUBLAS_OP_N,
7014 m: to_i32(k).ok_or(ArrowSchurGpuFailure::Unavailable)?,
7015 n: to_i32(k).ok_or(ArrowSchurGpuFailure::Unavailable)?,
7016 alpha: 1.0,
7017 lda: to_i32(k).ok_or(ArrowSchurGpuFailure::Unavailable)?,
7018 incx: 1,
7019 beta: 0.0,
7020 incy: 1,
7021 };
7022 unsafe { blas.gemv(gemv_cfg, &s_dev, &p_dev, &mut sp_dev) }
7024 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
7025 diag.matvec_calls += 1;
7026 diag.iterations += 1;
7027
7028 let p_sp = device_dot(&blas, &stream, k, &p_dev, &sp_dev)?;
7029 if !(p_sp > 0.0) {
7030 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
7033 reason: format!("reduced-β GPU PCG: non-positive curvature pᵀSp={p_sp:e}"),
7034 });
7035 }
7036 let alpha = rz / p_sp;
7037 device_axpy(&blas, &stream, k, alpha, &p_dev, &mut x_dev)?;
7038 device_axpy(&blas, &stream, k, -alpha, &sp_dev, &mut r_dev)?;
7039 let r_norm = device_nrm2(&blas, &stream, k, &r_dev)?;
7040 if r_norm <= tol {
7041 diag.final_relative_residual = r_norm / rhs_norm;
7042 diag.stopping_reason = PcgStopReason::Converged;
7043 break;
7044 }
7045 launch_jacobi_mul(&stream, vector_module, &inv_diag_dev, &r_dev, &mut z_dev, k)?;
7046 diag.precond_apply_calls += 1;
7047 let rz_new = device_dot(&blas, &stream, k, &r_dev, &z_dev)?;
7048 if rz_new <= 0.0 || !rz_new.is_finite() {
7049 return Err(ArrowSchurGpuFailure::SchurFactorFailed {
7050 reason: format!("reduced-β GPU PCG: non-positive rᵀM⁻¹r={rz_new:e}"),
7051 });
7052 }
7053 let beta = rz_new / rz;
7054 launch_update_p(&stream, vector_module, &z_dev, beta, &mut p_dev, k)?;
7055 rz = rz_new;
7056 }
7057 if diag.stopping_reason != PcgStopReason::Converged {
7058 let r_norm = device_nrm2(&blas, &stream, k, &r_dev)?;
7059 diag.final_relative_residual = r_norm / rhs_norm;
7060 diag.stopping_reason = PcgStopReason::MaxIter;
7061 }
7062
7063 let x = stream
7064 .clone_dtoh(&x_dev)
7065 .map_err(|_| ArrowSchurGpuFailure::Unavailable)?;
7066 Ok((Array1::from_vec(x), diag))
7067 }
7068
7069 fn device_copy(
7070 blas: &CudaBlas,
7071 stream: &Arc<CudaStream>,
7072 n: usize,
7073 src: &CudaSlice<f64>,
7074 dst: &mut CudaSlice<f64>,
7075 ) -> Result<(), ArrowSchurGpuFailure> {
7076 let n_i = to_i32(n).ok_or(ArrowSchurGpuFailure::Unavailable)?;
7077 let (src_ptr, _src_rec) = src.device_ptr(stream);
7078 let (dst_ptr, _dst_rec) = dst.device_ptr_mut(stream);
7079 let status = unsafe {
7082 cudarc::cublas::sys::cublasDcopy_v2(
7083 *blas.handle(),
7084 n_i,
7085 src_ptr as *const f64,
7086 1,
7087 dst_ptr as *mut f64,
7088 1,
7089 )
7090 };
7091 if status == cublasStatus_t::CUBLAS_STATUS_SUCCESS {
7092 Ok(())
7093 } else {
7094 Err(ArrowSchurGpuFailure::Unavailable)
7095 }
7096 }
7097
7098 fn device_axpy(
7099 blas: &CudaBlas,
7100 stream: &Arc<CudaStream>,
7101 n: usize,
7102 alpha: f64,
7103 x: &CudaSlice<f64>,
7104 y: &mut CudaSlice<f64>,
7105 ) -> Result<(), ArrowSchurGpuFailure> {
7106 let n_i = to_i32(n).ok_or(ArrowSchurGpuFailure::Unavailable)?;
7107 let (x_ptr, _x_rec) = x.device_ptr(stream);
7108 let (y_ptr, _y_rec) = y.device_ptr_mut(stream);
7109 let status = unsafe {
7112 cudarc::cublas::sys::cublasDaxpy_v2(
7113 *blas.handle(),
7114 n_i,
7115 &alpha,
7116 x_ptr as *const f64,
7117 1,
7118 y_ptr as *mut f64,
7119 1,
7120 )
7121 };
7122 if status == cublasStatus_t::CUBLAS_STATUS_SUCCESS {
7123 Ok(())
7124 } else {
7125 Err(ArrowSchurGpuFailure::Unavailable)
7126 }
7127 }
7128
7129 fn device_axpy_strided(
7134 blas: &CudaBlas,
7135 stream: &Arc<CudaStream>,
7136 n: usize,
7137 alpha: f64,
7138 x: &CudaSlice<f64>,
7139 incx: usize,
7140 y: &mut CudaSlice<f64>,
7141 incy: usize,
7142 ) -> Result<(), ArrowSchurGpuFailure> {
7143 let n_i = to_i32(n).ok_or(ArrowSchurGpuFailure::Unavailable)?;
7144 let incx_i = to_i32(incx).ok_or(ArrowSchurGpuFailure::Unavailable)?;
7145 let incy_i = to_i32(incy).ok_or(ArrowSchurGpuFailure::Unavailable)?;
7146 let (x_ptr, _x_rec) = x.device_ptr(stream);
7147 let (y_ptr, _y_rec) = y.device_ptr_mut(stream);
7148 let status = unsafe {
7151 cudarc::cublas::sys::cublasDaxpy_v2(
7152 *blas.handle(),
7153 n_i,
7154 &alpha,
7155 x_ptr as *const f64,
7156 incx_i,
7157 y_ptr as *mut f64,
7158 incy_i,
7159 )
7160 };
7161 if status == cublasStatus_t::CUBLAS_STATUS_SUCCESS {
7162 Ok(())
7163 } else {
7164 Err(ArrowSchurGpuFailure::Unavailable)
7165 }
7166 }
7167
7168 fn device_dot(
7169 blas: &CudaBlas,
7170 stream: &Arc<CudaStream>,
7171 n: usize,
7172 x: &CudaSlice<f64>,
7173 y: &CudaSlice<f64>,
7174 ) -> Result<f64, ArrowSchurGpuFailure> {
7175 let n_i = to_i32(n).ok_or(ArrowSchurGpuFailure::Unavailable)?;
7176 let (x_ptr, _x_rec) = x.device_ptr(stream);
7177 let (y_ptr, _y_rec) = y.device_ptr(stream);
7178 let mut result = 0.0_f64;
7179 let status = unsafe {
7183 cudarc::cublas::sys::cublasDdot_v2(
7184 *blas.handle(),
7185 n_i,
7186 x_ptr as *const f64,
7187 1,
7188 y_ptr as *const f64,
7189 1,
7190 &mut result,
7191 )
7192 };
7193 if status == cublasStatus_t::CUBLAS_STATUS_SUCCESS {
7194 Ok(result)
7195 } else {
7196 Err(ArrowSchurGpuFailure::Unavailable)
7197 }
7198 }
7199
7200 fn device_nrm2(
7201 blas: &CudaBlas,
7202 stream: &Arc<CudaStream>,
7203 n: usize,
7204 x: &CudaSlice<f64>,
7205 ) -> Result<f64, ArrowSchurGpuFailure> {
7206 let n_i = to_i32(n).ok_or(ArrowSchurGpuFailure::Unavailable)?;
7207 let (x_ptr, _x_rec) = x.device_ptr(stream);
7208 let mut result = 0.0_f64;
7209 let status = unsafe {
7213 cudarc::cublas::sys::cublasDnrm2_v2(
7214 *blas.handle(),
7215 n_i,
7216 x_ptr as *const f64,
7217 1,
7218 &mut result,
7219 )
7220 };
7221 if status == cublasStatus_t::CUBLAS_STATUS_SUCCESS {
7222 Ok(result)
7223 } else {
7224 Err(ArrowSchurGpuFailure::Unavailable)
7225 }
7226 }
7227
7228 #[cfg(test)]
7229 mod tests {
7230 use crate::arrow_schur::{ArrowSchurSystem, DeviceSaeFrameData, DeviceSaePcgData, DeviceSaeSmoothBlock, FactoredFrameGBlock};
7235 use ndarray::Array2;
7236
7237 fn tiny_framed_fixture() -> (ArrowSchurSystem, DeviceSaePcgData) {
7241 let p = 3usize;
7242 let ranks = vec![2usize, 3usize];
7243 let basis_sizes = vec![2usize, 2usize];
7244 let mut border_offsets = Vec::new();
7245 let mut acc = 0usize;
7246 for k in 0..2 {
7247 border_offsets.push(acc);
7248 acc += basis_sizes[k] * ranks[k];
7249 }
7250 let border_dim = acc;
7251 let frame_of = |k: usize| -> Array2<f64> {
7252 Array2::from_shape_fn((p, ranks[k]), |(i, j)| {
7253 0.1 + 0.2 * ((i + 1) as f64) * ((j + 1 + 2 * k) as f64)
7254 })
7255 };
7256 let frames: Vec<Array2<f64>> = (0..2).map(frame_of).collect();
7257 let w_of = |i: usize, j: usize| -> Array2<f64> {
7258 let (ui, uj) = (&frames[i], &frames[j]);
7259 Array2::from_shape_fn((ranks[i], ranks[j]), |(a, b)| {
7260 (0..p).map(|c| ui[[c, a]] * uj[[c, b]]).sum()
7261 })
7262 };
7263 let mut frame_blocks = Vec::new();
7264 for &(i, j) in &[(0usize, 0usize), (1usize, 1usize), (0, 1), (1, 0)] {
7265 let (mi, mj) = (basis_sizes[i], basis_sizes[j]);
7266 let mut g =
7267 Array2::<f64>::from_shape_fn((mi, mj), |(r, c)| 0.1 * (r + 2 * c + 1) as f64);
7268 if i == j {
7269 for r in 0..mi.min(mj) {
7270 g[[r, r]] += mi as f64 + 2.0;
7271 }
7272 }
7273 frame_blocks.push(FactoredFrameGBlock {
7274 atom_i: i,
7275 atom_j: j,
7276 g,
7277 w: w_of(i, j),
7278 });
7279 }
7280 let mut smooth_blocks = Vec::new();
7281 for k in 0..2 {
7282 let m = basis_sizes[k];
7283 let mut s =
7284 Array2::<f64>::from_shape_fn((m, m), |(r, c)| 0.05 * (r + c + 1) as f64);
7285 for r in 0..m {
7286 s[[r, r]] += 1.0;
7287 }
7288 smooth_blocks.push(DeviceSaeSmoothBlock {
7289 global_offset: border_offsets[k],
7290 factor_a: s,
7291 });
7292 }
7293 let smooth_ranks = ranks.clone();
7294 let n = 2usize;
7295 let q = 2usize;
7296 let mut sys = ArrowSchurSystem::new(n, q, border_dim);
7297 let mut row_htbeta = Vec::new();
7298 for i in 0..n {
7299 let mut htt =
7300 Array2::<f64>::from_shape_fn((q, q), |(r, c)| 0.3 * (r + c + 1) as f64);
7301 for r in 0..q {
7302 htt[[r, r]] += q as f64 + 2.0;
7303 }
7304 sys.rows[i].htt = htt;
7305 let mut slab = vec![0.0_f64; q * border_dim];
7306 for c in 0..q {
7307 for col in 0..border_dim {
7308 let v = 0.01 * ((c + 1) * (col + 1) + i) as f64;
7309 slab[c * border_dim + col] = v;
7310 sys.rows[i].htbeta[[c, col]] = v;
7311 }
7312 }
7313 row_htbeta.push(slab);
7314 }
7315 let data = DeviceSaePcgData {
7316 p,
7317 beta_dim: border_dim,
7318 a_phi: std::sync::Arc::from(Vec::new().into_boxed_slice()),
7319 local_jac: std::sync::Arc::from(Vec::new().into_boxed_slice()),
7320 smooth_blocks,
7321 sparse_g_blocks: Vec::new(),
7322 frame: Some(DeviceSaeFrameData {
7323 ranks,
7324 basis_sizes,
7325 border_offsets,
7326 frame_blocks,
7327 smooth_ranks,
7328 row_htbeta,
7329 }),
7330 };
7331 (sys, data)
7332 }
7333
7334 #[test]
7346 fn sae_resident_frame_only_ainv_is_ridge_dependent_1017() {
7347 let (sys, data) = tiny_framed_fixture();
7348 let frame = data.frame.as_ref().expect("framed fixture");
7349
7350 let host_a = super::super::flatten_frame_host_operands(&sys, &data, frame)
7353 .expect("host operands a");
7354 let host_b = super::super::flatten_frame_host_operands(&sys, &data, frame)
7355 .expect("host operands b");
7356 assert_eq!(
7357 host_a.s_data, host_b.s_data,
7358 "smooth λS must be ridge-independent"
7359 );
7360 assert_eq!(
7361 host_a.g_data, host_b.g_data,
7362 "frame G must be ridge-independent"
7363 );
7364 assert_eq!(
7365 host_a.w_data, host_b.w_data,
7366 "frame W must be ridge-independent"
7367 );
7368 assert_eq!(host_a.htb, host_b.htb, "row H_tβ must be ridge-independent");
7369 assert_eq!(host_a.q_of, host_b.q_of);
7370
7371 let report = data.operand_byte_report();
7372 assert!(report.framed, "fixture must exercise the framed lane");
7373 assert!(
7374 report.total_bytes > 0,
7375 "framed operands must have nonzero bytes"
7376 );
7377
7378 let ainv_lo = super::super::compute_ainv_host(
7382 &sys,
7383 &host_a.q_of,
7384 host_a.max_q,
7385 host_a.n_rows,
7386 1e-3,
7387 )
7388 .expect("ainv lo");
7389 let ainv_lo2 = super::super::compute_ainv_host(
7390 &sys,
7391 &host_a.q_of,
7392 host_a.max_q,
7393 host_a.n_rows,
7394 1e-3,
7395 )
7396 .expect("ainv lo repeat");
7397 let ainv_hi = super::super::compute_ainv_host(
7398 &sys,
7399 &host_a.q_of,
7400 host_a.max_q,
7401 host_a.n_rows,
7402 1e3,
7403 )
7404 .expect("ainv hi");
7405 assert_eq!(
7406 ainv_lo, ainv_lo2,
7407 "ainv must be deterministic at a fixed ridge"
7408 );
7409 let max_diff = ainv_lo
7410 .iter()
7411 .zip(&ainv_hi)
7412 .map(|(a, b)| (a - b).abs())
7413 .fold(0.0_f64, f64::max);
7414 assert!(
7415 max_diff > 1e-6,
7416 "ainv MUST change with ridge_t (else the split is wrong); max_diff={max_diff:e}"
7417 );
7418
7419 let trials = crate::arrow_schur::DEFAULT_PROXIMAL_MAX_ATTEMPTS + 1;
7423 let saved = report.total_bytes * (trials - 1);
7424 assert!(saved > 0);
7425 eprintln!(
7426 "#1017 resident-frame ladder saving: {trials} trials × {}B ridge-independent \
7427 operand upload → resident removes {saved}B, re-uploading only ainv \
7428 ({}rows × {}² × 8B) per trial",
7429 report.total_bytes, host_a.n_rows, host_a.max_q
7430 );
7431 }
7432
7433 }
7434}
7435
7436#[cfg(test)]
7437mod tests {
7438 use super::*;
7439 use crate::arrow_schur::ArrowSchurSystem;
7440 use ndarray::{Array2, ArrayView1};
7441
7442 fn build_fixture(n: usize, d: usize, k: usize, seed: u64) -> ArrowSchurSystem {
7443 let mut sys = ArrowSchurSystem::new(n, d, k);
7444 let mut state = seed.wrapping_mul(0x9E37_79B9_7F4A_7C15);
7445 let mut sample = || -> f64 {
7446 state = state
7447 .wrapping_mul(6364136223846793005)
7448 .wrapping_add(1442695040888963407);
7449 ((state >> 33) as f64) / ((1u64 << 31) as f64) - 1.0
7450 };
7451 for row in &mut sys.rows {
7452 let mut a = Array2::<f64>::zeros((d, d));
7453 for r in 0..d {
7454 for c in 0..d {
7455 a[[r, c]] = sample();
7456 }
7457 }
7458 let mut htt = a.t().dot(&a);
7459 for r in 0..d {
7460 htt[[r, r]] += d as f64 + 1.0;
7461 }
7462 row.htt = htt;
7463 for r in 0..d {
7464 for c in 0..k {
7465 row.htbeta[[r, c]] = 0.1 * sample();
7466 }
7467 row.gt[r] = sample();
7468 }
7469 }
7470 let mut hbb_a = Array2::<f64>::zeros((k, k));
7471 for r in 0..k {
7472 for c in 0..k {
7473 hbb_a[[r, c]] = sample();
7474 }
7475 }
7476 let mut hbb = hbb_a.t().dot(&hbb_a);
7477 for r in 0..k {
7478 hbb[[r, r]] += k as f64 + 1.0;
7479 }
7480 sys.hbb = hbb;
7481 for r in 0..k {
7482 sys.gb[r] = sample();
7483 }
7484 sys
7485 }
7486
7487 #[test]
7492 fn ridge_bump_makes_known_indefinite_blocks_pd() {
7493 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));
7501 indef2[[0, 0]] = 1.0;
7502 indef2[[1, 1]] = 1.0;
7503 indef2[[0, 1]] = 2.0;
7504 indef2[[1, 0]] = 2.0;
7505 let pd = Array2::<f64>::from_diag(&Array1::from_elem(3, 5.0));
7508
7509 for (label, block) in [
7510 ("-I (λ_min=-1)", neg_identity),
7511 ("-250·I (λ_min=-250)", scaled_neg),
7512 ("[[1,2],[2,1]] (λ_min=-1)", indef2),
7513 ("5·I (PD)", pd),
7514 ] {
7515 let ridge_t = 0.0;
7516 let bump = ridge_bump_to_make_pd(block.view(), ridge_t);
7517 assert!(
7518 bump > 0.0 && bump.is_finite(),
7519 "[{label}] bump must be strictly positive and finite, got {bump:e}"
7520 );
7521 let d = block.nrows();
7522 let mut shifted = block.clone();
7523 for i in 0..d {
7524 shifted[[i, i]] += ridge_t + bump;
7525 }
7526 assert!(
7527 cholesky_factor_in_place(shifted.view(), CholeskyGuard::NonnegativePivot).is_some(),
7528 "[{label}] H_tt + (ridge_t + bump={bump:e})·I must be PD after the \
7529 Gershgorin bump, but the Cholesky still rejected it"
7530 );
7531 }
7532 }
7533
7534 #[cfg(target_os = "linux")]
7544 #[test]
7545 fn ridge_bump_colmajor_matches_rowmajor_for_symmetric_block() {
7546 let mut a = Array2::<f64>::zeros((3, 3));
7548 a[[0, 0]] = -2.0;
7549 a[[1, 1]] = 0.5;
7550 a[[2, 2]] = 1.0;
7551 a[[0, 1]] = 0.3;
7552 a[[1, 0]] = 0.3;
7553 a[[1, 2]] = -0.4;
7554 a[[2, 1]] = -0.4;
7555 a[[0, 2]] = 0.1;
7556 a[[2, 0]] = 0.1;
7557
7558 let row_major_bump = ridge_bump_to_make_pd(a.view(), 0.0);
7559
7560 let d = 3;
7562 let mut col_major = vec![0.0_f64; d * d];
7563 for c in 0..d {
7564 for r in 0..d {
7565 col_major[c * d + r] = a[[r, c]];
7566 }
7567 }
7568 let col_major_bump = ridge_bump_to_make_pd_colmajor(&col_major, d);
7569
7570 assert!(
7571 (row_major_bump - col_major_bump).abs() <= 1e-12 * row_major_bump.max(1.0),
7572 "colmajor bump {col_major_bump:e} must match rowmajor bump \
7573 {row_major_bump:e} for a symmetric block"
7574 );
7575
7576 let mut shifted = a.clone();
7578 for i in 0..d {
7579 shifted[[i, i]] += col_major_bump;
7580 }
7581 assert!(
7582 cholesky_factor_in_place(shifted.view(), CholeskyGuard::NonnegativePivot).is_some(),
7583 "colmajor Gershgorin bump must make the symmetric block PD"
7584 );
7585 }
7586
7587 fn device_pcg_fixture(k: usize) -> (Array2<f64>, Array1<f64>) {
7588 let mut s = Array2::<f64>::zeros((k, k));
7589 for row in 0..k {
7590 s[[row, row]] = 2.5 + 0.001 * ((row % 17) as f64);
7591 if row + 1 < k {
7592 s[[row, row + 1]] = -0.05;
7593 s[[row + 1, row]] = -0.05;
7594 }
7595 if row + 7 < k {
7596 s[[row, row + 7]] = 0.01;
7597 s[[row + 7, row]] = 0.01;
7598 }
7599 }
7600 let rhs = Array1::from_shape_fn(k, |idx| ((idx as f64 + 1.0) * 0.013).sin());
7601 (s, rhs)
7602 }
7603
7604 fn dense_pcg_cpu_reference(
7605 s: &Array2<f64>,
7606 rhs: &Array1<f64>,
7607 max_iterations: usize,
7608 relative_tolerance: f64,
7609 ) -> Array1<f64> {
7610 let k = rhs.len();
7611 let rhs_norm = rhs.iter().map(|v| v * v).sum::<f64>().sqrt();
7612 if rhs_norm == 0.0 {
7613 return Array1::<f64>::zeros(k);
7614 }
7615 let tol = (relative_tolerance.max(0.0) * rhs_norm).max(1e-12);
7616 let inv_diag: Vec<f64> = (0..k).map(|idx| 1.0 / s[[idx, idx]]).collect();
7617 let mut x = Array1::<f64>::zeros(k);
7618 let mut r = rhs.clone();
7619 let mut z = Array1::from_shape_fn(k, |idx| inv_diag[idx] * r[idx]);
7620 let mut p = z.clone();
7621 let mut sp = Array1::<f64>::zeros(k);
7622 let mut rz = r.iter().zip(z.iter()).map(|(a, b)| a * b).sum::<f64>();
7623 for _ in 0..max_iterations.max(1) {
7624 for row in 0..k {
7625 let mut acc = 0.0;
7626 for col in 0..k {
7627 acc += s[[row, col]] * p[col];
7628 }
7629 sp[row] = acc;
7630 }
7631 let p_sp = p.iter().zip(sp.iter()).map(|(a, b)| a * b).sum::<f64>();
7632 let alpha = rz / p_sp;
7633 for idx in 0..k {
7634 x[idx] += alpha * p[idx];
7635 r[idx] -= alpha * sp[idx];
7636 }
7637 let r_norm = r.iter().map(|v| v * v).sum::<f64>().sqrt();
7638 if r_norm <= tol {
7639 break;
7640 }
7641 for idx in 0..k {
7642 z[idx] = inv_diag[idx] * r[idx];
7643 }
7644 let rz_next = r.iter().zip(z.iter()).map(|(a, b)| a * b).sum::<f64>();
7645 let beta = rz_next / rz;
7646 for idx in 0..k {
7647 p[idx] = z[idx] + beta * p[idx];
7648 }
7649 rz = rz_next;
7650 }
7651 x
7652 }
7653
7654 #[test]
7655 fn device_resident_pcg_matches_cpu_reference_when_cuda_admits() {
7656 let (s, rhs) = device_pcg_fixture(512);
7657 let max_iterations = 200usize;
7658 let relative_tolerance = 1.0e-12;
7659 let cpu = dense_pcg_cpu_reference(&s, &rhs, max_iterations, relative_tolerance);
7660 let (device, diag) = match solve_reduced_beta_pcg_with_diagnostics(
7661 &s,
7662 &rhs,
7663 max_iterations,
7664 relative_tolerance,
7665 ) {
7666 Ok(result) => result,
7667 Err(failure) => {
7674 assert!(
7675 gam_gpu::device_runtime::GpuRuntime::resolve(gam_gpu::GpuPolicy::Auto)
7676 .unwrap_or_else(|error| {
7677 panic!("GPU probe fault in reduced-beta PCG test: {error}")
7678 })
7679 .is_none(),
7680 "#1017: CUDA device present but the device reduced-beta PCG \
7681 declined/faulted instead of returning a result (tag: {failure:?}) — \
7682 the kernel does not run correctly on GPU"
7683 );
7684 return;
7685 }
7686 };
7687 let max_err = cpu
7688 .iter()
7689 .zip(device.iter())
7690 .map(|(a, b)| (a - b).abs())
7691 .fold(0.0_f64, f64::max);
7692 assert!(
7693 max_err <= 1.0e-10,
7694 "device resident PCG parity failed: max_err={max_err:e}, diag={diag:?}"
7695 );
7696 assert!(diag.matvec_calls > 0);
7697 assert_eq!(diag.matvec_calls, diag.iterations);
7698 }
7699
7700 #[test]
7701 fn dense_reference_matches_independent_solve() {
7702 let sys = build_fixture(4, 5, 3, 7);
7703 let solution = solve_arrow_newton_step_dense_reference(&sys, 0.0, 0.0).unwrap();
7704 let n = sys.rows.len();
7708 let d = sys.d;
7709 let k = sys.k;
7710 let total = n * d + k;
7711 let mut h = Array2::<f64>::zeros((total, total));
7712 let mut g = ndarray::Array1::<f64>::zeros(total);
7713 for (i, row) in sys.rows.iter().enumerate() {
7714 let base = i * d;
7715 for c in 0..d {
7716 for r in 0..d {
7717 h[[base + r, base + c]] = row.htt[[r, c]];
7718 }
7719 }
7720 for c in 0..k {
7721 for r in 0..d {
7722 h[[base + r, n * d + c]] = row.htbeta[[r, c]];
7723 h[[n * d + c, base + r]] = row.htbeta[[r, c]];
7724 }
7725 }
7726 for r in 0..d {
7727 g[base + r] = row.gt[r];
7728 }
7729 }
7730 for c in 0..k {
7731 for r in 0..k {
7732 h[[n * d + r, n * d + c]] += sys.hbb[[r, c]];
7733 }
7734 g[n * d + c] = sys.gb[c];
7735 }
7736 let l = cholesky_factor_in_place(h.view(), CholeskyGuard::NonnegativePivot).unwrap();
7737 let rhs = g.mapv(|v| -v);
7738 let expected = cholesky_solve_vector(l.view(), rhs.view());
7739 for i in 0..n * d {
7740 assert!(
7741 (solution.delta_t[i] - expected[i]).abs() < 1e-10 * (1.0 + expected[i].abs()),
7742 "delta_t[{i}] mismatch: got {} expected {}",
7743 solution.delta_t[i],
7744 expected[i]
7745 );
7746 }
7747 for a in 0..k {
7748 assert!(
7749 (solution.delta_beta[a] - expected[n * d + a]).abs()
7750 < 1e-10 * (1.0 + expected[n * d + a].abs()),
7751 "delta_beta[{a}] mismatch"
7752 );
7753 }
7754 }
7755
7756 #[test]
7770 fn row_procedural_matvec_parallel_deterministic_and_matches_serial() {
7771 use crate::arrow_schur::SCHUR_MATVEC_PARALLEL_ROW_MIN;
7772 let n = SCHUR_MATVEC_PARALLEL_ROW_MIN + 96; let d = 3usize;
7774 let k = 24usize;
7775 let mut sys = build_fixture(n, d, k, 0xA17C_0FFE);
7776 let slabs: Vec<Array2<f64>> = sys.rows.iter().map(|row| row.htbeta.clone()).collect();
7781 let forward_slabs = slabs.clone();
7782 let transpose_slabs = slabs;
7783 sys.set_row_htbeta_operator(
7784 move |row: usize, x: ArrayView1<'_, f64>, out: &mut Array1<f64>| {
7785 let h = &forward_slabs[row];
7786 for r in 0..h.nrows() {
7787 let mut acc = 0.0_f64;
7788 for c in 0..h.ncols() {
7789 acc += h[[r, c]] * x[c];
7790 }
7791 out[r] = acc;
7792 }
7793 },
7794 move |row: usize, v: ArrayView1<'_, f64>, out: &mut Array1<f64>| {
7795 let h = &transpose_slabs[row];
7796 for r in 0..h.nrows() {
7797 for c in 0..h.ncols() {
7798 out[c] += h[[r, c]] * v[r];
7799 }
7800 }
7801 },
7802 );
7803
7804 let matvec = gpu_schur_matvec_backend(&sys, 0.0, 0.0)
7805 .expect("row-procedural matvec backend builds for matrix-free system");
7806 let x = Array1::from_shape_fn(k, |i| ((i as f64 + 1.0) * 0.37).sin());
7807
7808 let mut out_parallel_a = Array1::<f64>::zeros(k);
7812 matvec(&x, &mut out_parallel_a);
7813 let mut out_parallel_b = Array1::<f64>::zeros(k);
7814 matvec(&x, &mut out_parallel_b);
7815 for a in 0..k {
7816 assert_eq!(
7817 out_parallel_a[a].to_bits(),
7818 out_parallel_b[a].to_bits(),
7819 "row-procedural matvec parallel reduction is non-deterministic at index {a}"
7820 );
7821 }
7822
7823 let mut out_serial = Array1::<f64>::zeros(k);
7828 rayon::ThreadPoolBuilder::new()
7829 .num_threads(2)
7830 .build()
7831 .expect("build rayon pool")
7832 .install(|| matvec(&x, &mut out_serial));
7833
7834 let max_abs = out_serial.iter().fold(0.0_f64, |m, v| m.max(v.abs()));
7835 for a in 0..k {
7836 let diff = (out_parallel_a[a] - out_serial[a]).abs();
7837 assert!(
7838 diff <= 1e-12 * (1.0 + max_abs),
7839 "row-procedural matvec parallel vs serial diverged beyond reassociation \
7840 at index {a}: {} vs {} (diff={diff:e})",
7841 out_parallel_a[a],
7842 out_serial[a]
7843 );
7844 }
7845 }
7846
7847}