1use std::rc::Rc;
23use std::sync::{Arc, Mutex};
24
25use crate::kkt::aug_system_solver::{AugSysCoeffs, AugSysRhs, AugSysSol, AugSystemSolver};
26use crate::kkt::std_aug_system_solver::StdAugSystemSolver;
27use pounce_common::diagnostics::DiagnosticsState;
28use pounce_common::timing::TimingStatistics;
29use pounce_common::types::{Index, Number};
30use pounce_feral::{FeralBlockSolver, FeralConfig};
31use pounce_linsol::summary::LinearSolverSummary;
32use pounce_linsol::{ESymSolverStatus, FactorPattern};
33
34pub struct BlockAugSystemSolver {
35 inner: StdAugSystemSolver,
37 blocks: FeralBlockSolver,
38 labels: Vec<i32>,
41 detect: bool,
42 decided_for_dim: Option<Index>,
45 use_blocks: bool,
46 have_factor: bool,
47 negevals: Index,
48 last_status: ESymSolverStatus,
49 timing: Option<Rc<TimingStatistics>>,
50 min_block_size: usize,
53}
54
55pub const DEFAULT_MIN_BLOCK_SIZE: usize = 256;
59
60impl BlockAugSystemSolver {
61 pub fn new(inner: StdAugSystemSolver, labels: Vec<i32>, cfg: FeralConfig) -> Self {
65 Self {
66 inner,
67 blocks: FeralBlockSolver::new(cfg),
68 detect: labels.is_empty(),
69 labels,
70 decided_for_dim: None,
71 use_blocks: false,
72 have_factor: false,
73 negevals: 0,
74 last_status: ESymSolverStatus::Success,
75 timing: None,
76 min_block_size: DEFAULT_MIN_BLOCK_SIZE,
77 }
78 }
79
80 pub fn with_min_block_size(mut self, n: usize) -> Self {
83 self.min_block_size = n;
84 self
85 }
86
87 pub fn with_summary_sink(mut self, sink: Arc<Mutex<LinearSolverSummary>>) -> Self {
89 self.blocks = self.blocks.with_summary_sink(sink);
90 self
91 }
92
93 fn decide(&mut self, dim: Index) {
95 if self.decided_for_dim == Some(dim) {
96 return;
97 }
98 self.decided_for_dim = Some(dim);
99 self.use_blocks = false;
100 let (irn, jcn) = {
101 let (a, b, _v) = self.inner.assembled_triplet();
102 (a.to_vec(), b.to_vec())
103 };
104 if self.detect {
105 let t = std::time::Instant::now();
106 let found = detect_blocks(dim as usize, &irn, &jcn);
107 tracing::info!(
108 target: "pounce::kkt",
109 secs = t.elapsed().as_secs_f64(),
110 blocks = found.as_ref().map(|l: &Vec<i32>| l.iter().copied().max().unwrap_or(-1) + 1),
111 border = found.as_ref().map(|l: &Vec<i32>| l.iter().filter(|&&x| x < 0).count()),
112 "block detection"
113 );
114 match found {
115 Some(l) => self.labels = l,
116 None => {
117 tracing::warn!(
118 target: "pounce::kkt",
119 "no block structure found in the KKT; using the standard solver"
120 );
121 return;
122 }
123 }
124 }
125 if self.labels.len() != dim as usize {
126 tracing::warn!(
127 target: "pounce::kkt",
128 labels = self.labels.len(), dim = dim as usize,
129 "block structure has the wrong length for this KKT; using the standard solver"
130 );
131 return;
132 }
133 let border = self.labels.iter().filter(|&&l| l < 0).count();
138 let max_border = (dim as usize / 200).clamp(64, 4096);
139 if border > max_border {
140 tracing::warn!(
141 target: "pounce::kkt",
142 border, max_border,
143 "block structure's border is too large to pay off; using the standard solver"
144 );
145 return;
146 }
147 if self.min_block_size > 0 {
165 let mut sizes: Vec<usize> = {
166 let n_blocks = self.labels.iter().copied().max().unwrap_or(-1) + 1;
167 let mut v = vec![0usize; n_blocks.max(0) as usize];
168 for &l in &self.labels {
169 if l >= 0 {
170 v[l as usize] += 1;
171 }
172 }
173 v
174 };
175 if !sizes.is_empty() {
176 sizes.sort_unstable();
177 let median = sizes[sizes.len() / 2];
178 if median < self.min_block_size {
179 tracing::warn!(
180 target: "pounce::kkt",
181 median_block = median,
182 min_block_size = self.min_block_size,
183 n_blocks = sizes.len(),
184 "block structure's blocks are too small to pay off; using the standard solver"
185 );
186 return;
187 }
188 }
189 }
190 let t = std::time::Instant::now();
191 let st = self
192 .blocks
193 .initialize_structure(dim, &irn, &jcn, &self.labels);
194 tracing::info!(
195 target: "pounce::kkt", secs = t.elapsed().as_secs_f64(), "block structure initialized"
196 );
197 if st == ESymSolverStatus::Success {
198 self.use_blocks = true;
199 tracing::info!(
200 target: "pounce::kkt",
201 blocks = self.blocks.n_blocks(), border = self.blocks.border_dim(),
202 "block-parallel KKT path engaged"
203 );
204 } else {
205 tracing::warn!(
206 target: "pounce::kkt",
207 "block structure rejected by the backend; using the standard solver"
208 );
209 }
210 }
211
212 fn block_solve_one(
214 &mut self,
215 rhs: &AugSysRhs<'_>,
216 sol: &mut AugSysSol<'_>,
217 check_neg_evals: bool,
218 num_neg_evals: Index,
219 ) -> ESymSolverStatus {
220 let dim = self.inner.assembled_dim() as usize;
221 let vals = self.inner.assembled_triplet().2.to_vec();
222 self.blocks.values_array_mut().copy_from_slice(&vals);
223 let status = {
224 let _g = self
225 .timing
226 .as_deref()
227 .map(|t| t.linear_system_factorization.guard());
228 self.blocks.factor(check_neg_evals, num_neg_evals)
229 };
230 self.last_status = status;
231 match status {
232 ESymSolverStatus::Success => {
233 self.negevals = self.blocks.number_of_neg_evals();
234 let mut packed = vec![0.0; dim];
235 self.inner.pack_rhs(rhs, &mut packed);
236 let bstat = {
237 let _g = self
238 .timing
239 .as_deref()
240 .map(|t| t.linear_system_back_solve.guard());
241 self.blocks.backsolve(1, &mut packed)
242 };
243 if bstat != ESymSolverStatus::Success {
244 self.have_factor = false;
245 self.last_status = bstat;
246 return bstat;
247 }
248 self.inner.unpack_sol(&packed, sol);
249 self.have_factor = true;
250 ESymSolverStatus::Success
251 }
252 ESymSolverStatus::WrongInertia => {
256 self.negevals = self.blocks.number_of_neg_evals();
257 self.have_factor = false;
258 status
259 }
260 ESymSolverStatus::Singular => {
269 self.have_factor = false;
270 status
271 }
272 other => {
273 self.have_factor = false;
274 other
275 }
276 }
277 }
278}
279
280pub fn detect_blocks(dim: usize, irn: &[Index], jcn: &[Index]) -> Option<Vec<i32>> {
306 const DEGREE_GAP: u32 = 3;
307 const MIN_BLOCKS: usize = 4;
308 if dim < 4 * MIN_BLOCKS {
309 return None;
310 }
311 let mut deg = vec![0u32; dim];
312 for p in 0..irn.len() {
313 let (r, c) = (irn[p] as usize - 1, jcn[p] as usize - 1);
314 if r != c {
315 deg[r] += 1;
316 deg[c] += 1;
317 }
318 }
319 let median = {
320 let mut d = deg.clone();
321 let mid = d.len() / 2;
322 d.select_nth_unstable(mid);
323 d[mid].max(1)
324 };
325 let mut order: Vec<usize> = (0..dim).collect();
326 order.sort_unstable_by_key(|&i| std::cmp::Reverse(deg[i]));
327 let max_border = (dim / 20).max(1);
328 let mut k = 1usize;
331 while k <= max_border {
332 if deg[order[k - 1]] < DEGREE_GAP * median {
334 return None;
335 }
336 let mut is_border = vec![false; dim];
337 for &i in order.iter().take(k) {
338 is_border[i] = true;
339 }
340 if let Some(mut labels) = components_without(dim, irn, jcn, &is_border, MIN_BLOCKS) {
341 shrink_border(dim, irn, jcn, &mut labels);
342 merge_small_blocks(&mut labels);
343 return Some(labels);
344 }
345 k *= 2;
346 }
347 None
348}
349
350fn components_without(
353 dim: usize,
354 irn: &[Index],
355 jcn: &[Index],
356 is_border: &[bool],
357 min_blocks: usize,
358) -> Option<Vec<i32>> {
359 let mut parent: Vec<usize> = (0..dim).collect();
360 fn find(parent: &mut [usize], mut a: usize) -> usize {
361 while parent[a] != a {
362 parent[a] = parent[parent[a]];
363 a = parent[a];
364 }
365 a
366 }
367 for p in 0..irn.len() {
368 let (r, c) = (irn[p] as usize - 1, jcn[p] as usize - 1);
369 if r == c || is_border[r] || is_border[c] {
370 continue;
371 }
372 let (ra, rb) = (find(&mut parent, r), find(&mut parent, c));
373 if ra != rb {
374 parent[ra] = rb;
375 }
376 }
377 let mut label = vec![-1i32; dim];
378 let mut root_to_block = std::collections::HashMap::new();
379 let mut sizes: Vec<usize> = Vec::new();
380 for i in 0..dim {
381 if is_border[i] {
382 continue;
383 }
384 let r = find(&mut parent, i);
385 let b = *root_to_block.entry(r).or_insert_with(|| {
386 sizes.push(0);
387 sizes.len() - 1
388 });
389 label[i] = b as i32;
390 sizes[b] += 1;
391 }
392 let interior: usize = sizes.iter().sum();
393 if sizes.len() < min_blocks || interior == 0 {
394 return None;
395 }
396 if *sizes.iter().max().unwrap() * 2 >= interior {
397 return None; }
399 Some(label)
400}
401
402fn merge_small_blocks(labels: &mut [i32]) {
412 let nblocks = labels.iter().copied().max().unwrap_or(-1) + 1;
413 if nblocks <= 0 {
414 return;
415 }
416 let mut sizes = vec![0usize; nblocks as usize];
417 for &l in labels.iter() {
418 if l >= 0 {
419 sizes[l as usize] += 1;
420 }
421 }
422 let interior: usize = sizes.iter().sum();
423 let real: Vec<usize> = (0..sizes.len())
424 .filter(|&b| sizes[b] * sizes.len() * 4 >= interior)
425 .collect();
426 if real.is_empty() || real.len() == sizes.len() {
427 return;
428 }
429 let mut new_id = vec![usize::MAX; sizes.len()];
431 let mut load: Vec<usize> = Vec::new();
432 for (k, &b) in real.iter().enumerate() {
433 new_id[b] = k;
434 load.push(sizes[b]);
435 }
436 for b in 0..sizes.len() {
437 if new_id[b] == usize::MAX {
438 let smallest = load
439 .iter()
440 .enumerate()
441 .min_by_key(|&(_, l)| *l)
442 .map(|(i, _)| i)
443 .unwrap_or(0);
444 new_id[b] = smallest;
445 load[smallest] += sizes[b];
446 }
447 }
448 for l in labels.iter_mut() {
449 if *l >= 0 {
450 *l = new_id[*l as usize] as i32;
451 }
452 }
453}
454
455fn shrink_border(dim: usize, irn: &[Index], jcn: &[Index], labels: &mut [i32]) {
464 let mut deg = vec![0usize; dim];
466 for p in 0..irn.len() {
467 let (r, c) = (irn[p] as usize - 1, jcn[p] as usize - 1);
468 if r != c && (labels[r] < 0 || labels[c] < 0) {
469 deg[r] += 1;
470 deg[c] += 1;
471 }
472 }
473 let mut start = vec![0usize; dim + 1];
474 for i in 0..dim {
475 start[i + 1] = start[i] + deg[i];
476 }
477 let mut adj = vec![0usize; start[dim]];
478 let mut fill = start.clone();
479 for p in 0..irn.len() {
480 let (r, c) = (irn[p] as usize - 1, jcn[p] as usize - 1);
481 if r != c && (labels[r] < 0 || labels[c] < 0) {
482 adj[fill[r]] = c;
483 fill[r] += 1;
484 adj[fill[c]] = r;
485 fill[c] += 1;
486 }
487 }
488 let mut queue: Vec<usize> = (0..dim).filter(|&i| labels[i] < 0).collect();
489 let mut queued = vec![true; dim];
490 for i in 0..dim {
491 queued[i] = labels[i] < 0;
492 }
493 let mut head = 0usize;
494 while head < queue.len() {
495 let i = queue[head];
496 head += 1;
497 queued[i] = false;
498 if labels[i] >= 0 {
499 continue;
500 }
501 let mut seen: Option<i32> = None;
502 let mut separates = false;
503 for &j in &adj[start[i]..start[i + 1]] {
504 if labels[j] < 0 {
505 continue; }
507 match seen {
508 None => seen = Some(labels[j]),
509 Some(b) if b == labels[j] => {}
510 Some(_) => {
511 separates = true;
512 break;
513 }
514 }
515 }
516 if separates {
517 continue;
518 }
519 if let Some(b) = seen {
520 labels[i] = b;
521 for &j in &adj[start[i]..start[i + 1]] {
522 if labels[j] < 0 && !queued[j] {
523 queued[j] = true;
524 queue.push(j);
525 }
526 }
527 }
528 }
529}
530
531impl AugSystemSolver for BlockAugSystemSolver {
532 fn provides_inertia(&self) -> bool {
533 self.inner.provides_inertia()
534 }
535
536 fn singularity_certified(&self) -> bool {
537 if self.use_blocks {
540 false
541 } else {
542 self.inner.singularity_certified()
543 }
544 }
545
546 fn number_of_neg_evals(&self) -> Index {
547 if self.use_blocks {
548 self.negevals
549 } else {
550 self.inner.number_of_neg_evals()
551 }
552 }
553
554 fn system_dim(&self) -> Index {
555 self.inner.system_dim()
556 }
557
558 fn kkt_triplets(&self) -> Option<(Index, Vec<Index>, Vec<Index>, Vec<Number>)> {
559 self.inner.kkt_triplets()
560 }
561
562 fn l_factor(&self, want_values: bool) -> Option<FactorPattern> {
563 if self.use_blocks {
565 None
566 } else {
567 self.inner.l_factor(want_values)
568 }
569 }
570
571 fn increase_quality(&mut self) -> bool {
572 self.have_factor = false;
573 if self.use_blocks {
574 self.blocks.increase_quality()
575 } else {
576 self.inner.increase_quality()
577 }
578 }
579
580 fn last_solve_status(&self) -> ESymSolverStatus {
581 if self.use_blocks {
582 self.last_status
583 } else {
584 self.inner.last_solve_status()
585 }
586 }
587
588 fn set_timing_stats(&mut self, timing: Rc<TimingStatistics>) {
589 self.timing = Some(Rc::clone(&timing));
590 self.inner.set_timing_stats(timing);
591 }
592
593 fn set_diagnostics(&mut self, diag: Rc<DiagnosticsState>) {
594 self.inner.set_diagnostics(diag);
595 }
596
597 fn solve(
598 &mut self,
599 coeffs: &AugSysCoeffs<'_>,
600 rhs: &AugSysRhs<'_>,
601 sol: &mut AugSysSol<'_>,
602 check_neg_evals: bool,
603 num_neg_evals: Index,
604 ) -> ESymSolverStatus {
605 let s = self.inner.assemble(coeffs);
606 if s != ESymSolverStatus::Success {
607 self.last_status = s;
608 return s;
609 }
610 let dim = self.inner.assembled_dim();
611 self.decide(dim);
612
613 if self.use_blocks {
614 let st = self.block_solve_one(rhs, sol, check_neg_evals, num_neg_evals);
615 match st {
616 ESymSolverStatus::Success
617 | ESymSolverStatus::WrongInertia
618 | ESymSolverStatus::Singular => return st,
619 _ => {
620 tracing::warn!(
621 target: "pounce::kkt",
622 status = ?st,
623 "block backend could not factor this KKT; falling back to the standard solver"
624 );
625 self.use_blocks = false;
626 return self
627 .inner
628 .solve(coeffs, rhs, sol, check_neg_evals, num_neg_evals);
629 }
630 }
631 }
632 self.inner
633 .solve(coeffs, rhs, sol, check_neg_evals, num_neg_evals)
634 }
635
636 fn resolve(
637 &mut self,
638 coeffs: &AugSysCoeffs<'_>,
639 rhs: &AugSysRhs<'_>,
640 sol: &mut AugSysSol<'_>,
641 ) -> ESymSolverStatus {
642 if self.use_blocks {
643 if self.have_factor {
644 let dim = self.inner.assembled_dim() as usize;
645 let mut packed = vec![0.0; dim];
646 self.inner.pack_rhs(rhs, &mut packed);
647 let bstat = {
648 let _g = self
649 .timing
650 .as_deref()
651 .map(|t| t.linear_system_back_solve.guard());
652 self.blocks.backsolve(1, &mut packed)
653 };
654 if bstat == ESymSolverStatus::Success {
655 self.inner.unpack_sol(&packed, sol);
656 }
657 return bstat;
658 }
659 return self.solve(coeffs, rhs, sol, false, 0);
660 }
661 self.inner.resolve(coeffs, rhs, sol)
662 }
663}
664
665#[cfg(test)]
666mod tests {
667 use super::*;
668
669 fn arrowhead() -> (usize, Vec<Index>, Vec<Index>) {
672 let (n, nb, nblocks) = (10usize, 2usize, 8usize);
673 let dim = nblocks * n + nb;
674 let (mut ia, mut ja) = (Vec::new(), Vec::new());
675 for b in 0..nblocks {
676 let off = b * n;
677 for i in 0..n {
678 ia.push((off + i + 1) as Index);
679 ja.push((off + i + 1) as Index);
680 if i + 1 < n {
681 ia.push((off + i + 2) as Index);
682 ja.push((off + i + 1) as Index);
683 }
684 }
685 for k in 0..nb {
686 ia.push((nblocks * n + k + 1) as Index);
687 ja.push((off + k + 1) as Index);
688 }
689 }
690 for k in 0..nb {
691 ia.push((nblocks * n + k + 1) as Index);
692 ja.push((nblocks * n + k + 1) as Index);
693 }
694 (dim, ia, ja)
695 }
696
697 #[test]
698 fn detection_finds_the_blocks_and_the_border() {
699 let (dim, ia, ja) = arrowhead();
700 let labels = detect_blocks(dim, &ia, &ja).expect("structure found");
701 assert_eq!(labels.len(), dim);
702 assert_eq!(labels.iter().filter(|&&l| l < 0).count(), 2, "border");
703 let mut per = std::collections::HashMap::new();
704 for &l in labels.iter().filter(|&&l| l >= 0) {
705 *per.entry(l).or_insert(0) += 1;
706 }
707 assert_eq!(per.len(), 8, "eight blocks");
708 assert!(per.values().all(|&c| c == 10), "blocks of ten: {per:?}");
709 }
710
711 #[test]
716 fn detection_declines_a_matrix_without_structure() {
717 let n = 400usize;
718 let (mut ia, mut ja) = (Vec::new(), Vec::new());
719 for i in 0..n {
720 ia.push((i + 1) as Index);
721 ja.push((i + 1) as Index);
722 if i + 1 < n {
723 ia.push((i + 2) as Index);
724 ja.push((i + 1) as Index);
725 }
726 }
727 assert!(detect_blocks(n, &ia, &ja).is_none());
728 }
729}