diff --git a/CHANGELOG.md b/CHANGELOG.md index f23e27b1..791e7cd0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -23,6 +23,7 @@ and this project follows [Semantic Versioning](https://semver.org/). - **BREAKING:** `mlsmr(op, b, m, x, start, &options, stall, par, scratch)` writes the solution into `x` and returns `LsmrInfo`; `start: InitialGuess` says whether `x` holds a warm start. Mismatched lengths or shapes panic instead of returning `SolveError::InvalidInput`. `MlsmrOptions`, `LsmrResult`, `EscalationPolicy`, `EscalationHandler` and `Progress` are removed; `stall: Option` replaces `escalation`, and `LsmrStopReason::Escalated` is `Stalled` (#365). - **BREAKING:** `within::LsmrOptions` is `schwarz_precond::LsmrOptions`, which has a hidden field, so a struct literal needs `..Default::default()`. - **BREAKING:** `within::Preconditioner::apply` is renamed `apply_slices`; `Preconditioner` implements faer `LinOp`, and `within` re-exports `LinOp`, `MatMut`, `MatRef`, `MemBuffer`, `MemStack`, `Par` and `StackReq` for it. +- `solve` and `solve_batch` read the thread count from the ambient rayon pool once; a batch divides it across right-hand sides (#365). - **BREAKING:** The serialized `Preconditioner` wire format moved v19 → v20; 0.4.0 bytes no longer decode. ### Fixed diff --git a/crates/within/src/operator/schwarz.rs b/crates/within/src/operator/schwarz.rs index 5a7fc681..8d143784 100644 --- a/crates/within/src/operator/schwarz.rs +++ b/crates/within/src/operator/schwarz.rs @@ -86,14 +86,6 @@ impl Preconditioner { } } - /// The ladder a reusing solver resumes, while this map has not escalated. - pub(crate) fn ladder(&self) -> Option<&AdaptiveLadder> { - match &self.inner { - Variant::Adaptive(ladder) => Some(ladder), - Variant::Additive { .. } | Variant::Diagonal(_) => None, - } - } - /// Drop the ladder for a design with no factor pair to escalate to, keeping its diagonal. pub(crate) fn settle(self) -> Self { match self.inner { @@ -260,16 +252,24 @@ fn diagonal_map(prepared: &PreparedDesign<'_>) -> Result, BuildError> Ok(Arc::from(diag)) } -/// Build the additive Schwarz map plus its warnings; `None` when the design has no factor-pair -/// subdomain to build on, where plain LSMR is the fallback. +/// What one preconditioner build yields: the map (`None` adds nothing below it) and its warnings. +pub(crate) struct PreconditionerBuild { + pub(crate) map: Option, + pub(crate) warnings: Vec, +} + +/// The additive Schwarz build; no map when the design has no factor-pair subdomain. pub(crate) fn build_schwarz( prepared: &PreparedDesign<'_>, local_solver: &LocalSolverConfig, -) -> Result<(Option, Vec), BuildError> { +) -> Result { let build_started = Instant::now(); let (domains, warnings) = crate::domain::build_local_domains(prepared, local_solver)?; if domains.is_empty() { - return Ok((None, warnings)); + return Ok(PreconditionerBuild { + map: None, + warnings, + }); } let schwarz = build_additive(domains, prepared.design.n_dofs)?; let preconditioner = Preconditioner { @@ -279,5 +279,8 @@ pub(crate) fn build_schwarz( }, build_duration: build_started.elapsed(), }; - Ok((Some(preconditioner), warnings)) + Ok(PreconditionerBuild { + map: Some(preconditioner), + warnings, + }) } diff --git a/crates/within/src/solver.rs b/crates/within/src/solver.rs index 7b2435a6..35c840bd 100644 --- a/crates/within/src/solver.rs +++ b/crates/within/src/solver.rs @@ -6,16 +6,12 @@ use std::borrow::Cow; use std::time::Instant; use ndarray::ArrayView2; -use rayon::prelude::*; -use schwarz_precond::{ - mlsmr, IdentityPrecond, InitialGuess, LinOp, LsmrInfo, LsmrStopReason, MemBuffer, Par, - StackReq, Staleness, -}; +use schwarz_precond::{LsmrInfo, Par}; use crate::channel::CoefficientAddress; use crate::config::{LsmrOptions, PreconditionerConfig}; use crate::domain::{Design, Effect, PreparedDesign}; -use crate::operator::schwarz::Preconditioner; +use crate::operator::schwarz::{Preconditioner, PreconditionerBuild}; use crate::{BuildError, BuildWarning, SolveError, WithinError}; mod ladder; @@ -23,7 +19,7 @@ mod layout; #[cfg(test)] mod tests; -use ladder::PrecondSlot; +use ladder::Rung; pub use layout::CoefficientLayout; /// Fallible conversion into a [`Design`] for [`Solver::new`]: a categories @@ -211,7 +207,10 @@ impl BatchSolveResult { /// Python boundary — uses owned columns. pub struct Solver<'a> { prepared: PreparedDesign<'a>, - slot: PrecondSlot, + /// The map built with the solver; `None` inside it runs unpreconditioned. + base: PreconditionerBuild, + /// Only an `Adaptive` config on two or more terms has one. + rungs: Vec, } impl std::fmt::Debug for Solver<'_> { @@ -256,51 +255,14 @@ impl PreparedRhs<'_> { } } -/// One finished LSMR run and the seconds it alone took, without any build between rungs. +/// One RHS's runs so far: its iterate, the last report with all iterations, and their seconds. struct Run { x: Vec, info: LsmrInfo, + /// LSMR alone, without any build between rungs. solve_secs: f64, } -impl Run { - /// `mlsmr` on the ambient pool from `x`, which is zero unless `start` says otherwise. - fn timed( - prepared: &PreparedDesign<'_>, - b: &[f64], - m: &impl LinOp, - mut x: Vec, - start: InitialGuess, - options: &LsmrOptions, - stall: Option, - ) -> Result { - let t_start = Instant::now(); - let mut scratch = MemBuffer::new(StackReq::EMPTY); - let info = mlsmr( - prepared, - b, - m, - &mut x, - start, - options, - stall, - Par::rayon(0), - &mut scratch, - )?; - Ok(Self { - x, - info, - solve_secs: t_start.elapsed().as_secs_f64(), - }) - } -} - -/// A first pass over one RHS: finished, or stalled on the diagonal and kept for the resume. -enum Pass<'s> { - Done(RhsSolution), - Stalled(PreparedRhs<'s>, Run), -} - impl<'a> Solver<'a> { /// Construct a solver. /// @@ -326,44 +288,21 @@ impl<'a> Solver<'a> { ) -> Result { // Whiten the slope columns (if any) before the preconditioner reads them. let prepared = PreparedDesign::new(design.into_design()?, weights)?; - let n_dofs = prepared.design.n_dofs; - - let slot = match preconditioner.into() { - PreconditionerInput::Default => { - PrecondSlot::build(&prepared, PreconditionerConfig::default())? - } - PreconditionerInput::Config(c) => PrecondSlot::build(&prepared, c)?, - PreconditionerInput::Prebuilt(p) => { - if p.nrows() != n_dofs || p.ncols() != n_dofs { - return Err(BuildError::PreconditionerDimensionMismatch { - expected: n_dofs, - actual_rows: p.nrows(), - actual_cols: p.ncols(), - }); - } - PrecondSlot::reuse(&prepared, p)? - } - }; - - Ok(Self { prepared, slot }) + match preconditioner.into() { + PreconditionerInput::Default => Self::build(prepared, PreconditionerConfig::default()), + PreconditionerInput::Config(c) => Self::build(prepared, c), + PreconditionerInput::Prebuilt(p) => Self::reuse(prepared, p), + } } - /// Non-fatal events from the preconditioner build; a reused pre-built preconditioner - /// contributes none (its own were reported when built). An Adaptive solver reports its - /// deferred build's once a solve escalates. + /// Non-fatal events from the latest build: the base's, or an escalated rung's. pub fn warnings(&self) -> &[BuildWarning] { - match &self.slot { - PrecondSlot::Static { warnings, .. } => warnings, - PrecondSlot::Adaptive(a) => a.build().map_or(&[], |b| &b.warnings), - } + &self.builds().last().expect("the base is built").warnings } /// Whether an [`Adaptive`](crate::PreconditionerConfig::Adaptive) solve has built Schwarz. pub fn has_escalated(&self) -> bool { - match &self.slot { - PrecondSlot::Adaptive(a) => a.schwarz().is_some(), - PrecondSlot::Static { .. } => false, - } + self.builds().skip(1).any(|b| b.map.is_some()) } /// Validate `y` and move it into the internal observation frame. @@ -401,9 +340,9 @@ impl<'a> Solver<'a> { /// Demean and back-transform one finished run. Excludes the design-level `layout` / /// `warnings` / `unidentified`, which the public entry points attach once. - fn finish(&self, rhs: PreparedRhs<'_>, run: Run) -> RhsSolution { + fn finish(&self, rhs: PreparedRhs<'_>, run: Run, par: Par) -> RhsSolution { let Run { mut x, info, .. } = run; - let demeaned = self.prepared.demeaned(&x, &rhs.y, Par::rayon(0)); + let demeaned = self.prepared.demeaned(&x, &rhs.y, par); self.prepared.back_transform(&mut x); RhsSolution { @@ -419,82 +358,6 @@ impl<'a> Solver<'a> { } } - /// Every RHS in one pass; an unsettled ladder then builds Schwarz once and resumes the stalled. - fn solve_all( - &self, - ys: &[&[f64]], - lsmr: &LsmrOptions, - ) -> Result<(Vec, f64), WithinError> { - let (map, ladder) = match &self.slot { - PrecondSlot::Static { map, .. } => (map.as_ref(), None), - // A settled ladder never re-probes: its map, or its kept build error, is final. - PrecondSlot::Adaptive(a) => match a.built.get() { - Some(Ok(_)) => (Some(a.rung()), None), - Some(Err(e)) => return Err(e.clone().into()), - None => (Some(&a.base), Some(a.as_ref())), - }, - }; - // A stall on the last permitted iteration leaves rung 2 nothing to spend, so no build. - let stalled = - |r: &LsmrInfo| r.stop_reason == LsmrStopReason::Stalled && r.iterations < lsmr.maxiter; - // Collecting into `Result` fails fast on the first per-RHS error, not during the fold. - let passes = ys - .par_iter() - .map(|y| { - let rhs = self.prepare(y)?; - let stall = ladder.map(|a| a.ladder().stall); - let (prepared, b) = (&self.prepared, rhs.b()); - let x = vec![0.0; prepared.design.n_dofs]; - let zero = InitialGuess::Zero; - let run = match map { - Some(m) => Run::timed(prepared, b, m, x, zero, lsmr, stall)?, - None => { - let identity = IdentityPrecond { - dim: prepared.design.n_dofs, - }; - Run::timed(prepared, b, &identity, x, zero, lsmr, stall)? - } - }; - Ok(if stalled(&run.info) { - Pass::Stalled(rhs, run) - } else { - Pass::Done(self.finish(rhs, run)) - }) - }) - .collect::, SolveError>>()?; - - // Built outside the fan-out, so sibling RHS never race for it. - let build_secs = match ladder { - Some(a) if passes.iter().any(|p| matches!(p, Pass::Stalled(..))) => { - a.escalate(&self.prepared)? - } - _ => 0.0, - }; - let solutions = passes - .into_par_iter() - .map(|pass| { - let (rhs, probe) = match pass { - Pass::Done(solution) => return Ok(solution), - Pass::Stalled(rhs, probe) => (rhs, probe), - }; - let map = ladder.expect("only a ladder stalls").rung(); - let spent = probe.info.iterations; - // The whole ladder shares the caller's budget; rung 2 gets what rung 1 left. - let options = LsmrOptions { - maxiter: lsmr.maxiter - spent, - ..*lsmr - }; - let start = InitialGuess::Given; - let (prepared, b) = (&self.prepared, rhs.b()); - let mut resumed = Run::timed(prepared, b, map, probe.x, start, &options, None)?; - resumed.info.iterations += spent; - resumed.solve_secs += probe.solve_secs; - Ok(self.finish(rhs, resumed)) - }) - .collect::, SolveError>>()?; - Ok((solutions, build_secs)) - } - /// Per-level directions the data cannot identify, shared across all RHS: /// identification depends only on the design and weights, never on `y`. fn unidentified(&self) -> Vec { @@ -514,7 +377,8 @@ impl<'a> Solver<'a> { let lsmr = lsmr.into().unwrap_or(&default); let t_start = Instant::now(); - let (solutions, build_secs) = self.solve_all(&[y], lsmr)?; + // The ambient pool's width is read here, at the boundary, and nowhere inside. + let (solutions, build_secs) = self.solve_all(&[y], lsmr, Par::rayon(0))?; let solution = solutions .into_iter() .next() @@ -537,7 +401,7 @@ impl<'a> Solver<'a> { }) } - /// Solve for multiple RHS vectors in parallel. + /// Solve for multiple RHS vectors; the ambient pool's width is split across them. pub fn solve_batch<'o>( &self, ys: &[&[f64]], @@ -548,7 +412,8 @@ impl<'a> Solver<'a> { let lsmr = lsmr.into().unwrap_or(&default); let n_rhs = ys.len(); - let (solutions, build_secs) = self.solve_all(ys, lsmr)?; + // The ambient pool's width is read here, at the boundary, and nowhere inside. + let (solutions, build_secs) = self.solve_all(ys, lsmr, Par::rayon(0))?; let mut x = Vec::with_capacity(self.prepared.design.n_dofs * n_rhs); let mut demeaned = Vec::with_capacity(self.prepared.design.n_obs * n_rhs); @@ -588,10 +453,7 @@ impl<'a> Solver<'a> { /// Access the preconditioner (for serialization or reuse across solvers). /// Under Adaptive: the Schwarz map once built, otherwise the diagonal base carrying the ladder. pub fn preconditioner(&self) -> Option<&Preconditioner> { - match &self.slot { - PrecondSlot::Static { map, .. } => map.as_ref(), - PrecondSlot::Adaptive(a) => Some(a.rung()), - } + self.builds().filter_map(|b| b.map.as_ref()).last() } /// Number of DOFs (coefficients). diff --git a/crates/within/src/solver/ladder.rs b/crates/within/src/solver/ladder.rs index a7a69f88..d31aa7d4 100644 --- a/crates/within/src/solver/ladder.rs +++ b/crates/within/src/solver/ladder.rs @@ -1,46 +1,48 @@ -//! The solver's preconditioner slot: a fixed map, or the diagonal→Schwarz ladder that -//! [`PreconditionerConfig::Adaptive`] builds on a stalled solve. +//! The maps a solve runs on, each `Adaptive` rung built on a stall, and the lanes over the RHS. use std::time::Instant; use once_cell::sync::OnceCell; +use schwarz_precond::par::join; +use schwarz_precond::Par; +use schwarz_precond::{ + mlsmr, IdentityPrecond, InitialGuess, LsmrOptions, LsmrStopReason, MemBuffer, StackReq, + Staleness, +}; -use crate::config::PreconditionerConfig; +use super::{PreparedRhs, RhsSolution, Run, Solver}; +use crate::config::{LocalSolverConfig, PreconditionerConfig}; use crate::domain::PreparedDesign; use crate::operator::schwarz::{ - build_adaptive, build_diagonal, build_schwarz, AdaptiveLadder, Preconditioner, + build_adaptive, build_diagonal, build_schwarz, Preconditioner, PreconditionerBuild, }; -use crate::{BuildError, BuildWarning}; - -/// The solver's preconditioner: a fixed map, or an adaptive diagonal→Schwarz ladder. -pub(super) enum PrecondSlot { - /// A single map (`None` = unpreconditioned) built at construction, with its build's warnings. - Static { - map: Option, - warnings: Vec, - }, - /// Diagonal now, Schwarz built lazily on a stalled contraction. - Adaptive(Box), +use crate::{BuildError, SolveError, WithinError}; + +/// One escalation: the stall that triggers it and the map it builds, at most once per solver. +pub(super) struct Rung { + stall: Staleness, + local_solver: LocalSolverConfig, + /// A design's own build failure settles here; a failed isolation pool retries instead. + pub(super) built: OnceCell>, } -impl PrecondSlot { - /// Resolve a strategy: a fixed map built now, or a diagonal base with Schwarz deferred to a stall. +impl<'a> Solver<'a> { + /// `Adaptive` goes through `reuse`, the one place a rung is derived. pub(super) fn build( - prepared: &PreparedDesign<'_>, + prepared: PreparedDesign<'a>, config: PreconditionerConfig, ) -> Result { - Ok(match config { - PreconditionerConfig::Off => Self::Static { + let base = match config { + PreconditionerConfig::Off => PreconditionerBuild { map: None, warnings: Vec::new(), }, - PreconditionerConfig::Diagonal => Self::Static { - map: Some(build_diagonal(prepared)?), + PreconditionerConfig::Diagonal => PreconditionerBuild { + map: Some(build_diagonal(&prepared)?), warnings: Vec::new(), }, PreconditionerConfig::Additive { local_solver } => { - let (map, warnings) = build_schwarz(prepared, &local_solver)?; - Self::Static { map, warnings } + build_schwarz(&prepared, &local_solver)? } PreconditionerConfig::Adaptive { local_solver, @@ -48,93 +50,232 @@ impl PrecondSlot { } => { // Fail before the diagonal pass; `reuse` repeats this for a deserialized ladder. local_solver.validate()?; - let base = build_adaptive(prepared, stall, local_solver)?; - Self::reuse(prepared, base)? + let map = build_adaptive(&prepared, stall, local_solver)?; + return Self::reuse(prepared, map); } + }; + Ok(Self { + prepared, + base, + rungs: Vec::new(), }) } - /// An unescalated ladder resumes, or settles on one term; any other map stays fixed. + /// The map as base plus its `Adaptive` rung; on one term it settles to its diagonal, no rung. pub(super) fn reuse( - prepared: &PreparedDesign<'_>, + prepared: PreparedDesign<'a>, preconditioner: Preconditioner, ) -> Result { - let Some(ladder) = preconditioner.ladder() else { - return Ok(Self::Static { - map: Some(preconditioner), - warnings: Vec::new(), + let n_dofs = prepared.design.n_dofs; + let (actual_rows, actual_cols) = (preconditioner.nrows(), preconditioner.ncols()); + if actual_rows != n_dofs || actual_cols != n_dofs { + return Err(BuildError::PreconditionerDimensionMismatch { + expected: n_dofs, + actual_rows, + actual_cols, }); + } + let rung = match preconditioner.config() { + PreconditionerConfig::Adaptive { + local_solver, + stall, + } => { + // A deserialized ladder skipped `Adaptive`'s build-time check. + local_solver.validate()?; + (prepared.design.n_factors() >= 2).then(|| Rung { + stall, + local_solver, + built: OnceCell::new(), + }) + } + _ => None, + }; + let map = match rung { + Some(_) => preconditioner, + None => preconditioner.settle(), }; - // A deserialized ladder skipped `Adaptive`'s build-time check, and a one-term design settles. - ladder.escalated.validate()?; - Ok(if prepared.design.n_factors() < 2 { - Self::Static { - map: Some(preconditioner.settle()), + Ok(Self { + prepared, + base: PreconditionerBuild { + map: Some(map), warnings: Vec::new(), - } - } else { - Self::Adaptive(Box::new(AdaptivePrecond { - base: preconditioner, - built: OnceCell::new(), - })) + }, + rungs: rung.into_iter().collect(), }) } -} -/// Diagonal-first strategy holding everything needed to build the Schwarz rung on demand. -pub(super) struct AdaptivePrecond { - /// Applies the diagonal and carries the ladder, so handing it out keeps the strategy. - pub(super) base: Preconditioner, - /// A design's own build failure settles here too; a failed isolation pool retries instead. - pub(super) built: OnceCell>, -} + /// `base`, then each rung's build up to the first rung not built. + pub(super) fn builds(&self) -> impl Iterator { + std::iter::once(&self.base).chain( + self.rungs + .iter() + .map_while(|r| r.built.get()?.as_ref().ok()), + ) + } -/// Outcome of the deferred build: the Schwarz map, or `None` when no factor-pair target exists. -pub(super) struct AdaptiveBuild { - pub(super) schwarz: Option, - pub(super) warnings: Vec, + /// Every RHS in lanes; a stall builds the next rung once under all of `par`, then resumes. + pub(super) fn solve_all<'s>( + &'s self, + ys: &'s [&'s [f64]], + lsmr: &LsmrOptions, + par: Par, + ) -> Result<(Vec, f64), WithinError> { + let prepared = &self.prepared; + // One pass over the cells: a concurrent build must not move `map` and `next` apart. + let (mut map, mut built) = (None, 0); + for build in self.builds() { + (map, built) = (build.map.as_ref().or(map), built + 1); + } + let mut next = self.rungs[built - 1..].iter(); + // A settled rung never re-probes: its kept build error is final. + if let Some(Some(Err(e))) = self.rungs.get(built - 1).map(|r| r.built.get()) { + return Err(e.clone().into()); + } + let (mut done, mut build_secs) = (Vec::with_capacity(ys.len()), 0.0); + let mut pending: Vec<_> = (0..ys.len()).map(|i| (i, None)).collect(); + loop { + let rung = next.next(); + let stall = rung.map(|r| r.stall); + let lanes = Lanes { + solver: self, + map, + stall, + ys, + lsmr, + }; + let (finished, stalled) = lanes.run(pending, par)?; + done.extend(finished); + let Some(rung) = rung.filter(|_| !stalled.is_empty()) else { + break; + }; + let t_build = Instant::now(); + let built = rung.built.get_or_try_init(|| { + isolated(par, || build_schwarz(prepared, &rung.local_solver)) + })?; + map = built.as_ref().map_err(Clone::clone)?.map.as_ref().or(map); + build_secs += t_build.elapsed().as_secs_f64(); + pending = stalled + .into_iter() + .map(|(i, rhs, run)| (i, Some((rhs, run)))) + .collect(); + } + done.sort_unstable_by_key(|&(i, _)| i); + Ok((done.into_iter().map(|(_, s)| s).collect(), build_secs)) + } } -impl AdaptivePrecond { - pub(super) fn ladder(&self) -> &AdaptiveLadder { - self.base - .ladder() - .expect("an adaptive slot is built only from a ladder") - } +type Solved = (usize, RhsSolution); +type Stalled<'s> = (usize, PreparedRhs<'s>, Run); +type Pending<'s> = (usize, Option<(PreparedRhs<'s>, Run)>); - pub(super) fn build(&self) -> Option<&AdaptiveBuild> { - self.built.get().and_then(|b| b.as_ref().ok()) - } +/// One pass over pending RHS on one map, under `stall` when a rung waits above it. +struct Lanes<'s, 'a> { + solver: &'s Solver<'a>, + map: Option<&'s Preconditioner>, + stall: Option, + ys: &'s [&'s [f64]], + lsmr: &'s LsmrOptions, +} - pub(super) fn schwarz(&self) -> Option<&Preconditioner> { - self.build().and_then(|b| b.schwarz.as_ref()) +impl<'s> Lanes<'s, '_> { + /// Halves `pending` and `par` together; each lane runs its RHS in order on one scratch. + fn run( + &self, + mut pending: Vec>, + par: Par, + ) -> Result<(Vec, Vec>), SolveError> { + if pending.len() <= 1 || par.degree() <= 1 { + return self.lane(pending, par); + } + let right = pending.split_off(pending.len().div_ceil(2)); + let (l, r) = join(par, |p| self.run(pending, p), |p| self.run(right, p)); + let ((mut done, mut stalled), (more_done, more_stalled)) = (l?, r?); + done.extend(more_done); + stalled.extend(more_stalled); + Ok((done, stalled)) } - /// The rung a solve should run on: the escalated Schwarz map once built, else the diagonal. - pub(super) fn rung(&self) -> &Preconditioner { - self.schwarz().unwrap_or(&self.base) + fn lane( + &self, + pending: Vec>, + par: Par, + ) -> Result<(Vec, Vec>), SolveError> { + let mut scratch = MemBuffer::new(StackReq::EMPTY); + let (mut done, mut stalled) = (Vec::new(), Vec::new()); + for (i, prior) in pending { + let (rhs, prior) = match prior { + Some((rhs, run)) => (rhs, Some(run)), + None => (self.solver.prepare(self.ys[i])?, None), + }; + let run = self.resume(&rhs, prior, par, &mut scratch)?; + // A stall on the last permitted iteration leaves the rung nothing to spend. + if run.info.stop_reason == LsmrStopReason::Stalled + && run.info.iterations < self.lsmr.maxiter + { + stalled.push((i, rhs, run)); + } else { + done.push((i, self.solver.finish(rhs, run, par))); + } + } + Ok((done, stalled)) } - /// Build once and return this call's build seconds; every other caller waits for it. - pub(super) fn escalate(&self, prepared: &PreparedDesign<'_>) -> Result { - let mut build_secs = 0.0; - let built = self.built.get_or_try_init(|| { - let t_build = Instant::now(); - let outcome = isolated(|| build_schwarz(prepared, &self.ladder().escalated))? - .map(|(schwarz, warnings)| AdaptiveBuild { schwarz, warnings }); - build_secs = t_build.elapsed().as_secs_f64(); - Ok(outcome) - })?; - built.as_ref().map(|_| build_secs).map_err(Clone::clone) + /// One run on `rhs`, warm-started from `prior`, which spent part of the iteration budget. + fn resume( + &self, + rhs: &PreparedRhs<'_>, + prior: Option, + par: Par, + scratch: &mut MemBuffer, + ) -> Result { + let prepared = &self.solver.prepared; + let (mut x, start, spent, spent_secs) = match prior { + Some(run) => ( + run.x, + InitialGuess::Given, + run.info.iterations, + run.solve_secs, + ), + None => ( + vec![0.0; prepared.design.n_dofs], + InitialGuess::Zero, + 0, + 0.0, + ), + }; + let options = LsmrOptions { + maxiter: self.lsmr.maxiter - spent, + ..*self.lsmr + }; + let b = rhs.b(); + let t_start = Instant::now(); + let (xs, stall) = (&mut x[..], self.stall); + let mut info = match self.map { + Some(m) => mlsmr(prepared, b, m, xs, start, &options, stall, par, scratch)?, + None => { + let identity = IdentityPrecond { + dim: prepared.design.n_dofs, + }; + mlsmr( + prepared, b, &identity, xs, start, &options, stall, par, scratch, + )? + } + }; + info.iterations += spent; + Ok(Run { + x, + info, + solve_secs: t_start.elapsed().as_secs_f64() + spent_secs, + }) } } /// Run `f` on a pool of its own, entered from a thread outside every pool, so no wait inside it can /// steal a solve off the shared pool; a stolen solve blocking on this build is the #371 deadlock. /// Nothing run here may wait on the shared pool. -fn isolated(f: impl FnOnce() -> R + Send) -> Result { +pub(super) fn isolated(par: Par, f: impl FnOnce() -> R + Send) -> Result { let pool = rayon::ThreadPoolBuilder::new() - .num_threads(rayon::current_num_threads()) + .num_threads(par.degree()) .build() .map_err(|e| BuildError::ThreadPool(e.to_string()))?; std::thread::scope(|s| { @@ -156,7 +297,7 @@ mod tests { use std::sync::Arc; use std::time::Duration; - use super::isolated; + use super::{isolated, Par}; thread_local! { static ON_CALLER_POOL: Cell = const { Cell::new(false) }; @@ -176,7 +317,7 @@ mod tests { ); assert!( !caller - .install(|| isolated(|| ON_CALLER_POOL.with(Cell::get))) + .install(|| isolated(Par::rayon(0), || ON_CALLER_POOL.with(Cell::get))) .expect("isolated build"), "the build ran on the caller's pool" ); @@ -194,7 +335,10 @@ mod tests { let queued = Arc::clone(&ran); caller.install(move || { rayon::spawn(move || queued.store(true, Ordering::SeqCst)); - isolated(|| std::thread::sleep(Duration::from_millis(50))).expect("isolated build"); + isolated(Par::rayon(0), || { + std::thread::sleep(Duration::from_millis(50)) + }) + .expect("isolated build"); assert!( !ran.load(Ordering::SeqCst), "the waiting worker stole a queued job while the build ran" diff --git a/crates/within/src/solver/tests.rs b/crates/within/src/solver/tests.rs index a99e2333..ca3363e8 100644 --- a/crates/within/src/solver/tests.rs +++ b/crates/within/src/solver/tests.rs @@ -340,7 +340,6 @@ impl SharedCovariatePair { /// same error without the ladder rebuilding, and re-failing, the same map. #[test] fn a_failed_deferred_build_is_kept_and_reported_again() { - use super::ladder::PrecondSlot; use crate::config::{ScalingConfig, ScalingFailure}; use crate::{BuildError, WithinError}; @@ -367,11 +366,11 @@ fn a_failed_deferred_build_is_kept_and_reported_again() { }; let first = solver.solve(&panel.y, None); assert!(unscalable(&first), "{first:?}"); - let PrecondSlot::Adaptive(a) = &solver.slot else { + let [rung] = &solver.rungs[..] else { panic!("the ladder") }; assert!(matches!( - a.built.get(), + rung.built.get(), Some(Err(BuildError::UnscalableComponent { .. })) )); assert!(unscalable(&solver.solve(&panel.y, None)));