use std::rc::Rc;
use std::sync::{Arc, Mutex};
use crate::kkt::aug_system_solver::{AugSysCoeffs, AugSysRhs, AugSysSol, AugSystemSolver};
use crate::kkt::std_aug_system_solver::StdAugSystemSolver;
use pounce_common::diagnostics::DiagnosticsState;
use pounce_common::timing::TimingStatistics;
use pounce_common::types::{Index, Number};
use pounce_feral::{FeralBlockSolver, FeralConfig};
use pounce_linsol::summary::LinearSolverSummary;
use pounce_linsol::{ESymSolverStatus, FactorPattern};
pub struct BlockAugSystemSolver {
inner: StdAugSystemSolver,
blocks: FeralBlockSolver,
labels: Vec<i32>,
detect: bool,
decided_for_dim: Option<Index>,
use_blocks: bool,
have_factor: bool,
negevals: Index,
last_status: ESymSolverStatus,
timing: Option<Rc<TimingStatistics>>,
min_block_size: usize,
}
pub const DEFAULT_MIN_BLOCK_SIZE: usize = 256;
impl BlockAugSystemSolver {
pub fn new(inner: StdAugSystemSolver, labels: Vec<i32>, cfg: FeralConfig) -> Self {
Self {
inner,
blocks: FeralBlockSolver::new(cfg),
detect: labels.is_empty(),
labels,
decided_for_dim: None,
use_blocks: false,
have_factor: false,
negevals: 0,
last_status: ESymSolverStatus::Success,
timing: None,
min_block_size: DEFAULT_MIN_BLOCK_SIZE,
}
}
pub fn with_min_block_size(mut self, n: usize) -> Self {
self.min_block_size = n;
self
}
pub fn with_summary_sink(mut self, sink: Arc<Mutex<LinearSolverSummary>>) -> Self {
self.blocks = self.blocks.with_summary_sink(sink);
self
}
fn decide(&mut self, dim: Index) {
if self.decided_for_dim == Some(dim) {
return;
}
self.decided_for_dim = Some(dim);
self.use_blocks = false;
let (irn, jcn) = {
let (a, b, _v) = self.inner.assembled_triplet();
(a.to_vec(), b.to_vec())
};
if self.detect {
let t = std::time::Instant::now();
let found = detect_blocks(dim as usize, &irn, &jcn);
tracing::info!(
target: "pounce::kkt",
secs = t.elapsed().as_secs_f64(),
blocks = found.as_ref().map(|l: &Vec<i32>| l.iter().copied().max().unwrap_or(-1) + 1),
border = found.as_ref().map(|l: &Vec<i32>| l.iter().filter(|&&x| x < 0).count()),
"block detection"
);
match found {
Some(l) => self.labels = l,
None => {
tracing::warn!(
target: "pounce::kkt",
"no block structure found in the KKT; using the standard solver"
);
return;
}
}
}
if self.labels.len() != dim as usize {
tracing::warn!(
target: "pounce::kkt",
labels = self.labels.len(), dim = dim as usize,
"block structure has the wrong length for this KKT; using the standard solver"
);
return;
}
let border = self.labels.iter().filter(|&&l| l < 0).count();
let max_border = (dim as usize / 200).clamp(64, 4096);
if border > max_border {
tracing::warn!(
target: "pounce::kkt",
border, max_border,
"block structure's border is too large to pay off; using the standard solver"
);
return;
}
if self.min_block_size > 0 {
let mut sizes: Vec<usize> = {
let n_blocks = self.labels.iter().copied().max().unwrap_or(-1) + 1;
let mut v = vec![0usize; n_blocks.max(0) as usize];
for &l in &self.labels {
if l >= 0 {
v[l as usize] += 1;
}
}
v
};
if !sizes.is_empty() {
sizes.sort_unstable();
let median = sizes[sizes.len() / 2];
if median < self.min_block_size {
tracing::warn!(
target: "pounce::kkt",
median_block = median,
min_block_size = self.min_block_size,
n_blocks = sizes.len(),
"block structure's blocks are too small to pay off; using the standard solver"
);
return;
}
}
}
let t = std::time::Instant::now();
let st = self
.blocks
.initialize_structure(dim, &irn, &jcn, &self.labels);
tracing::info!(
target: "pounce::kkt", secs = t.elapsed().as_secs_f64(), "block structure initialized"
);
if st == ESymSolverStatus::Success {
self.use_blocks = true;
tracing::info!(
target: "pounce::kkt",
blocks = self.blocks.n_blocks(), border = self.blocks.border_dim(),
"block-parallel KKT path engaged"
);
} else {
tracing::warn!(
target: "pounce::kkt",
"block structure rejected by the backend; using the standard solver"
);
}
}
fn block_solve_one(
&mut self,
rhs: &AugSysRhs<'_>,
sol: &mut AugSysSol<'_>,
check_neg_evals: bool,
num_neg_evals: Index,
) -> ESymSolverStatus {
let dim = self.inner.assembled_dim() as usize;
let vals = self.inner.assembled_triplet().2.to_vec();
self.blocks.values_array_mut().copy_from_slice(&vals);
let status = {
let _g = self
.timing
.as_deref()
.map(|t| t.linear_system_factorization.guard());
self.blocks.factor(check_neg_evals, num_neg_evals)
};
self.last_status = status;
match status {
ESymSolverStatus::Success => {
self.negevals = self.blocks.number_of_neg_evals();
let mut packed = vec![0.0; dim];
self.inner.pack_rhs(rhs, &mut packed);
let bstat = {
let _g = self
.timing
.as_deref()
.map(|t| t.linear_system_back_solve.guard());
self.blocks.backsolve(1, &mut packed)
};
if bstat != ESymSolverStatus::Success {
self.have_factor = false;
self.last_status = bstat;
return bstat;
}
self.inner.unpack_sol(&packed, sol);
self.have_factor = true;
ESymSolverStatus::Success
}
ESymSolverStatus::WrongInertia => {
self.negevals = self.blocks.number_of_neg_evals();
self.have_factor = false;
status
}
ESymSolverStatus::Singular => {
self.have_factor = false;
status
}
other => {
self.have_factor = false;
other
}
}
}
}
pub fn detect_blocks(dim: usize, irn: &[Index], jcn: &[Index]) -> Option<Vec<i32>> {
const DEGREE_GAP: u32 = 3;
const MIN_BLOCKS: usize = 4;
if dim < 4 * MIN_BLOCKS {
return None;
}
let mut deg = vec![0u32; dim];
for p in 0..irn.len() {
let (r, c) = (irn[p] as usize - 1, jcn[p] as usize - 1);
if r != c {
deg[r] += 1;
deg[c] += 1;
}
}
let median = {
let mut d = deg.clone();
let mid = d.len() / 2;
d.select_nth_unstable(mid);
d[mid].max(1)
};
let mut order: Vec<usize> = (0..dim).collect();
order.sort_unstable_by_key(|&i| std::cmp::Reverse(deg[i]));
let max_border = (dim / 20).max(1);
let mut k = 1usize;
while k <= max_border {
if deg[order[k - 1]] < DEGREE_GAP * median {
return None;
}
let mut is_border = vec![false; dim];
for &i in order.iter().take(k) {
is_border[i] = true;
}
if let Some(mut labels) = components_without(dim, irn, jcn, &is_border, MIN_BLOCKS) {
shrink_border(dim, irn, jcn, &mut labels);
merge_small_blocks(&mut labels);
return Some(labels);
}
k *= 2;
}
None
}
fn components_without(
dim: usize,
irn: &[Index],
jcn: &[Index],
is_border: &[bool],
min_blocks: usize,
) -> Option<Vec<i32>> {
let mut parent: Vec<usize> = (0..dim).collect();
fn find(parent: &mut [usize], mut a: usize) -> usize {
while parent[a] != a {
parent[a] = parent[parent[a]];
a = parent[a];
}
a
}
for p in 0..irn.len() {
let (r, c) = (irn[p] as usize - 1, jcn[p] as usize - 1);
if r == c || is_border[r] || is_border[c] {
continue;
}
let (ra, rb) = (find(&mut parent, r), find(&mut parent, c));
if ra != rb {
parent[ra] = rb;
}
}
let mut label = vec![-1i32; dim];
let mut root_to_block = std::collections::HashMap::new();
let mut sizes: Vec<usize> = Vec::new();
for i in 0..dim {
if is_border[i] {
continue;
}
let r = find(&mut parent, i);
let b = *root_to_block.entry(r).or_insert_with(|| {
sizes.push(0);
sizes.len() - 1
});
label[i] = b as i32;
sizes[b] += 1;
}
let interior: usize = sizes.iter().sum();
if sizes.len() < min_blocks || interior == 0 {
return None;
}
if *sizes.iter().max().unwrap() * 2 >= interior {
return None; }
Some(label)
}
fn merge_small_blocks(labels: &mut [i32]) {
let nblocks = labels.iter().copied().max().unwrap_or(-1) + 1;
if nblocks <= 0 {
return;
}
let mut sizes = vec![0usize; nblocks as usize];
for &l in labels.iter() {
if l >= 0 {
sizes[l as usize] += 1;
}
}
let interior: usize = sizes.iter().sum();
let real: Vec<usize> = (0..sizes.len())
.filter(|&b| sizes[b] * sizes.len() * 4 >= interior)
.collect();
if real.is_empty() || real.len() == sizes.len() {
return;
}
let mut new_id = vec![usize::MAX; sizes.len()];
let mut load: Vec<usize> = Vec::new();
for (k, &b) in real.iter().enumerate() {
new_id[b] = k;
load.push(sizes[b]);
}
for b in 0..sizes.len() {
if new_id[b] == usize::MAX {
let smallest = load
.iter()
.enumerate()
.min_by_key(|&(_, l)| *l)
.map(|(i, _)| i)
.unwrap_or(0);
new_id[b] = smallest;
load[smallest] += sizes[b];
}
}
for l in labels.iter_mut() {
if *l >= 0 {
*l = new_id[*l as usize] as i32;
}
}
}
fn shrink_border(dim: usize, irn: &[Index], jcn: &[Index], labels: &mut [i32]) {
let mut deg = vec![0usize; dim];
for p in 0..irn.len() {
let (r, c) = (irn[p] as usize - 1, jcn[p] as usize - 1);
if r != c && (labels[r] < 0 || labels[c] < 0) {
deg[r] += 1;
deg[c] += 1;
}
}
let mut start = vec![0usize; dim + 1];
for i in 0..dim {
start[i + 1] = start[i] + deg[i];
}
let mut adj = vec![0usize; start[dim]];
let mut fill = start.clone();
for p in 0..irn.len() {
let (r, c) = (irn[p] as usize - 1, jcn[p] as usize - 1);
if r != c && (labels[r] < 0 || labels[c] < 0) {
adj[fill[r]] = c;
fill[r] += 1;
adj[fill[c]] = r;
fill[c] += 1;
}
}
let mut queue: Vec<usize> = (0..dim).filter(|&i| labels[i] < 0).collect();
let mut queued = vec![true; dim];
for i in 0..dim {
queued[i] = labels[i] < 0;
}
let mut head = 0usize;
while head < queue.len() {
let i = queue[head];
head += 1;
queued[i] = false;
if labels[i] >= 0 {
continue;
}
let mut seen: Option<i32> = None;
let mut separates = false;
for &j in &adj[start[i]..start[i + 1]] {
if labels[j] < 0 {
continue; }
match seen {
None => seen = Some(labels[j]),
Some(b) if b == labels[j] => {}
Some(_) => {
separates = true;
break;
}
}
}
if separates {
continue;
}
if let Some(b) = seen {
labels[i] = b;
for &j in &adj[start[i]..start[i + 1]] {
if labels[j] < 0 && !queued[j] {
queued[j] = true;
queue.push(j);
}
}
}
}
}
impl AugSystemSolver for BlockAugSystemSolver {
fn provides_inertia(&self) -> bool {
self.inner.provides_inertia()
}
fn singularity_certified(&self) -> bool {
if self.use_blocks {
false
} else {
self.inner.singularity_certified()
}
}
fn number_of_neg_evals(&self) -> Index {
if self.use_blocks {
self.negevals
} else {
self.inner.number_of_neg_evals()
}
}
fn system_dim(&self) -> Index {
self.inner.system_dim()
}
fn kkt_triplets(&self) -> Option<(Index, Vec<Index>, Vec<Index>, Vec<Number>)> {
self.inner.kkt_triplets()
}
fn l_factor(&self, want_values: bool) -> Option<FactorPattern> {
if self.use_blocks {
None
} else {
self.inner.l_factor(want_values)
}
}
fn increase_quality(&mut self) -> bool {
self.have_factor = false;
if self.use_blocks {
self.blocks.increase_quality()
} else {
self.inner.increase_quality()
}
}
fn last_solve_status(&self) -> ESymSolverStatus {
if self.use_blocks {
self.last_status
} else {
self.inner.last_solve_status()
}
}
fn set_timing_stats(&mut self, timing: Rc<TimingStatistics>) {
self.timing = Some(Rc::clone(&timing));
self.inner.set_timing_stats(timing);
}
fn set_diagnostics(&mut self, diag: Rc<DiagnosticsState>) {
self.inner.set_diagnostics(diag);
}
fn solve(
&mut self,
coeffs: &AugSysCoeffs<'_>,
rhs: &AugSysRhs<'_>,
sol: &mut AugSysSol<'_>,
check_neg_evals: bool,
num_neg_evals: Index,
) -> ESymSolverStatus {
let s = self.inner.assemble(coeffs);
if s != ESymSolverStatus::Success {
self.last_status = s;
return s;
}
let dim = self.inner.assembled_dim();
self.decide(dim);
if self.use_blocks {
let st = self.block_solve_one(rhs, sol, check_neg_evals, num_neg_evals);
match st {
ESymSolverStatus::Success
| ESymSolverStatus::WrongInertia
| ESymSolverStatus::Singular => return st,
_ => {
tracing::warn!(
target: "pounce::kkt",
status = ?st,
"block backend could not factor this KKT; falling back to the standard solver"
);
self.use_blocks = false;
return self
.inner
.solve(coeffs, rhs, sol, check_neg_evals, num_neg_evals);
}
}
}
self.inner
.solve(coeffs, rhs, sol, check_neg_evals, num_neg_evals)
}
fn resolve(
&mut self,
coeffs: &AugSysCoeffs<'_>,
rhs: &AugSysRhs<'_>,
sol: &mut AugSysSol<'_>,
) -> ESymSolverStatus {
if self.use_blocks {
if self.have_factor {
let dim = self.inner.assembled_dim() as usize;
let mut packed = vec![0.0; dim];
self.inner.pack_rhs(rhs, &mut packed);
let bstat = {
let _g = self
.timing
.as_deref()
.map(|t| t.linear_system_back_solve.guard());
self.blocks.backsolve(1, &mut packed)
};
if bstat == ESymSolverStatus::Success {
self.inner.unpack_sol(&packed, sol);
}
return bstat;
}
return self.solve(coeffs, rhs, sol, false, 0);
}
self.inner.resolve(coeffs, rhs, sol)
}
}
#[cfg(test)]
mod tests {
use super::*;
fn arrowhead() -> (usize, Vec<Index>, Vec<Index>) {
let (n, nb, nblocks) = (10usize, 2usize, 8usize);
let dim = nblocks * n + nb;
let (mut ia, mut ja) = (Vec::new(), Vec::new());
for b in 0..nblocks {
let off = b * n;
for i in 0..n {
ia.push((off + i + 1) as Index);
ja.push((off + i + 1) as Index);
if i + 1 < n {
ia.push((off + i + 2) as Index);
ja.push((off + i + 1) as Index);
}
}
for k in 0..nb {
ia.push((nblocks * n + k + 1) as Index);
ja.push((off + k + 1) as Index);
}
}
for k in 0..nb {
ia.push((nblocks * n + k + 1) as Index);
ja.push((nblocks * n + k + 1) as Index);
}
(dim, ia, ja)
}
#[test]
fn detection_finds_the_blocks_and_the_border() {
let (dim, ia, ja) = arrowhead();
let labels = detect_blocks(dim, &ia, &ja).expect("structure found");
assert_eq!(labels.len(), dim);
assert_eq!(labels.iter().filter(|&&l| l < 0).count(), 2, "border");
let mut per = std::collections::HashMap::new();
for &l in labels.iter().filter(|&&l| l >= 0) {
*per.entry(l).or_insert(0) += 1;
}
assert_eq!(per.len(), 8, "eight blocks");
assert!(per.values().all(|&c| c == 10), "blocks of ten: {per:?}");
}
#[test]
fn detection_declines_a_matrix_without_structure() {
let n = 400usize;
let (mut ia, mut ja) = (Vec::new(), Vec::new());
for i in 0..n {
ia.push((i + 1) as Index);
ja.push((i + 1) as Index);
if i + 1 < n {
ia.push((i + 2) as Index);
ja.push((i + 1) as Index);
}
}
assert!(detect_blocks(n, &ia, &ja).is_none());
}
}