R8-g: parallel scatter + trailing updates + tree solves; force-only TL kernel for modified Newton

- the tangent scatter runs per CSR entry over a transposed contribution
  map, in the serial scatter's order (same bits), on rayon;
- large fronts' trailing gemms run per column block on rayon;
- forward solve multifrontal, backward top-down, subtrees on rayon;
- total_lagrangian::internal_force: the force of internal_force_and_tangent
  alone (bit-identical, tested), used by modified-Newton iterations that
  reuse the factor; the tangent is evaluated only when a refresh is due.

Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
This commit is contained in:
Omar Sobh
2026-09-25 22:52:23 -05:00
co-authored by Claude Opus 5.5
parent 860f5bb37f
commit d50a0e28d2
5 changed files with 459 additions and 79 deletions
@@ -650,8 +650,9 @@ impl<'a> NonlinearDynamicStepper<'a> {
for &dof in &self.free_dofs { for &dof in &self.free_dofs {
a_new[dof] = inv_beta_dt2 * (u_iter[dof] - u_pred[dof]); a_new[dof] = inv_beta_dt2 * (u_iter[dof] - u_pred[dof]);
} }
let (f_int, tangent) = if self.sparse.is_some() { let (f_int, tangent) = if let Some(sparse) = &self.sparse {
(self.sparse_forces(&u_iter)?, None) let with_tangent = sparse.wants_tangent(line_search, inv_beta_dt2);
(self.sparse_forces(&u_iter, with_tangent)?, None)
} else { } else {
let (f_int, tangent) = self.assemble(&u_iter, true, inv_beta_dt2)?; let (f_int, tangent) = self.assemble(&u_iter, true, inv_beta_dt2)?;
(f_int, Some(tangent)) (f_int, Some(tangent))
@@ -676,9 +677,16 @@ impl<'a> NonlinearDynamicStepper<'a> {
None => { None => {
let rate = residual_norm / previous_norm; let rate = residual_norm / previous_norm;
previous_norm = residual_norm; previous_norm = residual_norm;
let sparse = self.sparse.as_ref().expect("sparse path");
if sparse.needs_refresh(line_search, inv_beta_dt2, rate)
&& !sparse.tangents_current()
{
self.sparse_tangents(&u_iter)?;
}
let masses: Vec<&DMatrix<f64>> = self.caches.iter().map(|c| &c.mass).collect();
let sparse = self.sparse.as_mut().expect("sparse path"); let sparse = self.sparse.as_mut().expect("sparse path");
sparse.solve( sparse.solve(
self.caches.iter().map(|c| &c.mass), &masses,
&self.added_mass, &self.added_mass,
&self.free_index, &self.free_index,
&residual, &residual,
@@ -809,13 +817,26 @@ impl<'a> NonlinearDynamicStepper<'a> {
/// forces summed in element order (the banded path's order, so the /// forces summed in element order (the banded path's order, so the
/// internal force is the same float for float), the tangents kept for /// internal force is the same float for float), the tangents kept for
/// the factorisation. /// the factorisation.
fn sparse_forces(&mut self, solution: &DVector<f64>) -> FeaResult<DVector<f64>> { /// With `with_tangent` false (a modified-Newton iteration that will
/// reuse the factor) only the forces are evaluated — the TL force-only
/// kernel gives the same forces bit for bit.
fn sparse_forces(
&mut self,
solution: &DVector<f64>,
with_tangent: bool,
) -> FeaResult<DVector<f64>> {
let started = std::time::Instant::now(); let started = std::time::Instant::now();
let analysis = self.analysis; let analysis = self.analysis;
let results: Vec<(DVector<f64>, DMatrix<f64>)> = self let results: Vec<(DVector<f64>, Option<DMatrix<f64>>)> = self
.caches .caches
.par_iter() .par_iter()
.map(|cache| element_force_and_tangent(analysis, cache, solution)) .map(|cache| {
if with_tangent {
element_force_and_tangent(analysis, cache, solution).map(|(f, k)| (f, Some(k)))
} else {
element_force(analysis, cache, solution).map(|f| (f, None))
}
})
.collect::<FeaResult<_>>()?; .collect::<FeaResult<_>>()?;
let sparse = self.sparse.as_mut().expect("sparse path"); let sparse = self.sparse.as_mut().expect("sparse path");
Ok(sparse.accumulate( Ok(sparse.accumulate(
@@ -826,6 +847,23 @@ impl<'a> NonlinearDynamicStepper<'a> {
)) ))
} }
/// Sparse path: the element tangents at `solution` alone (a refresh
/// after a force-only assembly).
fn sparse_tangents(&mut self, solution: &DVector<f64>) -> FeaResult<()> {
let started = std::time::Instant::now();
let analysis = self.analysis;
let tangents: Vec<DMatrix<f64>> = self
.caches
.par_iter()
.map(|cache| element_force_and_tangent(analysis, cache, solution).map(|(_, k)| k))
.collect::<FeaResult<_>>()?;
self.sparse
.as_mut()
.expect("sparse path")
.set_tangents(tangents, started);
Ok(())
}
/// Internal force and (optionally) tangent at a full displacement /// Internal force and (optionally) tangent at a full displacement
/// vector, reduced to the free DOFs. The tangent includes the Newmark /// vector, reduced to the free DOFs. The tangent includes the Newmark
/// mass term `M / (β Δt²)` at the caller's `inv_beta_dt2` (the /// mass term `M / (β Δt²)` at the caller's `inv_beta_dt2` (the
@@ -935,3 +973,31 @@ fn element_force_and_tangent(
) )
} }
} }
/// One element's internal force alone (the TL force-only kernel; the
/// small-strain path evaluates the pair and drops the tangent).
fn element_force(
analysis: &NonlinearDynamicAnalysis,
cache: &ElementCache,
solution: &DVector<f64>,
) -> FeaResult<DVector<f64>> {
if !analysis.total_lagrangian {
return element_force_and_tangent(analysis, cache, solution).map(|(f, _)| f);
}
let dim = analysis.mesh.spatial_dimension;
let material = analysis.materials.get_material(cache.material_id).unwrap();
let fe = StandardFiniteElement::new(cache.element_type, cache.coords.clone());
let mut element_displacement = DVector::zeros(cache.dofs.len());
for (local, &dof) in cache.dofs.iter().enumerate() {
element_displacement[local] = solution[dof];
}
let (lambda, mu) = material.properties().lame_parameters();
let constitutive = saint_venant_kirchhoff(lambda, mu, dim);
total_lagrangian::internal_force(
&fe,
&cache.coords,
&element_displacement,
constitutive.as_ref(),
None,
)
}
@@ -22,8 +22,12 @@ use crate::assembly::SparseMatrix;
use crate::error::FeaResult; use crate::error::FeaResult;
use crate::solvers::{BandedLu, LinearSolver, SolverOptions, SparseLdlt}; use crate::solvers::{BandedLu, LinearSolver, SolverOptions, SparseLdlt};
use nalgebra::{DMatrix, DVector}; use nalgebra::{DMatrix, DVector};
use rayon::prelude::*;
use std::time::Instant; use std::time::Instant;
/// CSR entries per parallel scatter task.
const SCATTER_CHUNK: usize = 16_384;
/// Which linear solver carries the Newton tangent. /// Which linear solver carries the Newton tangent.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum TangentSolver { pub enum TangentSolver {
@@ -71,8 +75,12 @@ impl TangentSolver {
/// Counters and accumulated wall time of the sparse path. /// Counters and accumulated wall time of the sparse path.
#[derive(Debug, Clone, Copy, Default, PartialEq)] #[derive(Debug, Clone, Copy, Default, PartialEq)]
pub struct TangentStats { pub struct TangentStats {
/// Element force/tangent evaluations (parallel) + force summation. /// Element evaluations (parallel) + force summation; with the
/// tangents under full Newton, forces only on reused iterations.
pub assemblies: usize, pub assemblies: usize,
/// Extra tangent evaluations for a refresh decided after a
/// force-only assembly (slow contraction).
pub tangent_refreshes: usize,
/// Numeric factorisations. /// Numeric factorisations.
pub factorizations: usize, pub factorizations: usize,
/// Linear solves (Newton corrections). /// Linear solves (Newton corrections).
@@ -83,7 +91,9 @@ pub struct TangentStats {
pub fallbacks: usize, pub fallbacks: usize,
/// Seconds in element evaluation + force summation. /// Seconds in element evaluation + force summation.
pub assembly_seconds: f64, pub assembly_seconds: f64,
/// Seconds in tangent scatter + numeric factorisation. /// Seconds in the tangent scatter (element blocks → CSR).
pub scatter_seconds: f64,
/// Seconds in the numeric factorisation.
pub factor_seconds: f64, pub factor_seconds: f64,
/// Seconds in triangular solves. /// Seconds in triangular solves.
pub solve_seconds: f64, pub solve_seconds: f64,
@@ -97,13 +107,18 @@ pub(super) struct SparseTangent {
row_ptr: Vec<usize>, row_ptr: Vec<usize>,
col_idx: Vec<usize>, col_idx: Vec<usize>,
values: Vec<f64>, values: Vec<f64>,
/// Per element, `(local_row · n_e + local_col)` → CSR index, or /// Per CSR entry, its element contributions in element order:
/// `u32::MAX` when either DOF is constrained. /// `contrib[contrib_ptr[p]..contrib_ptr[p + 1]]` = `(element,
elem_pos: Vec<Vec<u32>>, /// column-major local index)` — the transposed scatter, so entries can
/// be summed in parallel in exactly the serial scatter's order.
contrib_ptr: Vec<usize>,
contrib: Vec<(u32, u32)>,
/// CSR index of each free DOF's diagonal. /// CSR index of each free DOF's diagonal.
diag_pos: Vec<usize>, diag_pos: Vec<usize>,
/// The element tangents of the latest assembly. /// The element tangents of the latest assembly that computed them.
element_tangents: Vec<DMatrix<f64>>, element_tangents: Vec<DMatrix<f64>>,
/// Whether `element_tangents` belong to the latest assembly's state.
tangents_current: bool,
solver: SparseLdlt, solver: SparseLdlt,
/// Separate instance for the rest state's mass solve. /// Separate instance for the rest state's mass solve.
pub(super) mass_solver: SparseLdlt, pub(super) mass_solver: SparseLdlt,
@@ -155,19 +170,32 @@ impl SparseTangent {
.binary_search(&c) .binary_search(&c)
.expect("pattern covers every element pair") .expect("pattern covers every element pair")
}; };
let mut elem_pos = Vec::new(); // Element-major, row-major within the element: the order in
for dofs in elements { // which a serial scatter would add the contributions.
let mut flat: Vec<(usize, u32, u32)> = Vec::new();
for (e, dofs) in elements.enumerate() {
let n_e = dofs.len(); let n_e = dofs.len();
let mut pos = vec![u32::MAX; n_e * n_e];
for (a, &r) in dofs.iter().enumerate() { for (a, &r) in dofs.iter().enumerate() {
let Some(fr) = free_index[r] else { continue }; let Some(fr) = free_index[r] else { continue };
for (b, &c) in dofs.iter().enumerate() { for (b, &c) in dofs.iter().enumerate() {
if let Some(fc) = free_index[c] { if let Some(fc) = free_index[c] {
pos[a * n_e + b] = find(fr, fc) as u32; flat.push((find(fr, fc), e as u32, (b * n_e + a) as u32));
} }
} }
} }
elem_pos.push(pos); }
let mut contrib_ptr = vec![0usize; col_idx.len() + 1];
for &(p, _, _) in &flat {
contrib_ptr[p + 1] += 1;
}
for p in 0..col_idx.len() {
contrib_ptr[p + 1] += contrib_ptr[p];
}
let mut fill = contrib_ptr.clone();
let mut contrib = vec![(0u32, 0u32); flat.len()];
for (p, e, l) in flat {
contrib[fill[p]] = (e, l);
fill[p] += 1;
} }
let diag_pos = (0..num_free).map(|d| find(d, d)).collect(); let diag_pos = (0..num_free).map(|d| find(d, d)).collect();
let mut solver = SparseLdlt::new(); let mut solver = SparseLdlt::new();
@@ -183,9 +211,11 @@ impl SparseTangent {
row_ptr, row_ptr,
col_idx, col_idx,
values, values,
elem_pos, contrib_ptr,
contrib,
diag_pos, diag_pos,
element_tangents: Vec::new(), element_tangents: Vec::new(),
tangents_current: false,
solver, solver,
mass_solver: SparseLdlt::new(), mass_solver: SparseLdlt::new(),
reuse: reuse.max(1), reuse: reuse.max(1),
@@ -201,38 +231,73 @@ impl SparseTangent {
self.stats self.stats
} }
/// Whether the next Newton solve will factorise for sure (so the
/// assembly should produce the element tangents): always under full
/// Newton; under modified Newton when the factor is missing, used up,
/// for another dt, or `force_fresh`.
pub(super) fn wants_tangent(&self, force_fresh: bool, coef: f64) -> bool {
#[allow(clippy::float_cmp)]
let same_coef = coef == self.factor_coef;
self.reuse == 1 || force_fresh || !self.factor_valid || self.uses_left == 0 || !same_coef
}
/// Whether the solve at contraction `rate` will factorise.
pub(super) fn needs_refresh(&self, force_fresh: bool, coef: f64, rate: f64) -> bool {
self.wants_tangent(force_fresh, coef) || rate > 0.5
}
/// Whether the stored element tangents are those of the latest state.
pub(super) fn tangents_current(&self) -> bool {
self.tangents_current
}
/// Sum the element internal forces (element order, local order — the /// Sum the element internal forces (element order, local order — the
/// banded path's order) into the free-DOF vector and keep the element /// banded path's order) into the free-DOF vector, and keep the element
/// tangents for a later factorisation. /// tangents when they were computed.
pub(super) fn accumulate( pub(super) fn accumulate(
&mut self, &mut self,
element_results: Vec<(DVector<f64>, DMatrix<f64>)>, element_results: Vec<(DVector<f64>, Option<DMatrix<f64>>)>,
element_dofs: impl Iterator<Item = impl AsRef<[usize]>>, element_dofs: impl Iterator<Item = impl AsRef<[usize]>>,
free_index: &[Option<usize>], free_index: &[Option<usize>],
started: Instant, started: Instant,
) -> DVector<f64> { ) -> DVector<f64> {
let mut internal = DVector::zeros(self.num_free); let mut internal = DVector::zeros(self.num_free);
self.element_tangents.clear(); let with_tangents = element_results.first().is_some_and(|r| r.1.is_some());
if with_tangents {
self.element_tangents.clear();
}
for ((f_e, k_e), dofs) in element_results.into_iter().zip(element_dofs) { for ((f_e, k_e), dofs) in element_results.into_iter().zip(element_dofs) {
for (local, &dof) in dofs.as_ref().iter().enumerate() { for (local, &dof) in dofs.as_ref().iter().enumerate() {
if let Some(free) = free_index[dof] { if let Some(free) = free_index[dof] {
internal[free] += f_e[local]; internal[free] += f_e[local];
} }
} }
self.element_tangents.push(k_e); if let Some(k_e) = k_e {
self.element_tangents.push(k_e);
}
} }
self.tangents_current = with_tangents;
self.stats.assemblies += 1; self.stats.assemblies += 1;
self.stats.assembly_seconds += started.elapsed().as_secs_f64(); self.stats.assembly_seconds += started.elapsed().as_secs_f64();
internal internal
} }
/// Install element tangents computed separately for the latest state
/// (a modified-Newton refresh the force-only assembly did not foresee).
pub(super) fn set_tangents(&mut self, tangents: Vec<DMatrix<f64>>, started: Instant) {
self.element_tangents = tangents;
self.tangents_current = true;
self.stats.tangent_refreshes += 1;
self.stats.assembly_seconds += started.elapsed().as_secs_f64();
}
/// The Newton correction `(K_T + coef·(M + M_added)) δ = r`, with the /// The Newton correction `(K_T + coef·(M + M_added)) δ = r`, with the
/// tangent of the latest [`Self::accumulate`]. `force_fresh` demands a /// tangent of the latest [`Self::accumulate`]. `force_fresh` demands a
/// new factor; `rate` is `‖r_i‖/‖r_{i−1}‖` within the step (0 at the /// new factor; `rate` is `‖r_i‖/‖r_{i−1}‖` within the step (0 at the
/// first iteration). /// first iteration).
pub(super) fn solve<'m>( pub(super) fn solve(
&mut self, &mut self,
masses: impl Iterator<Item = &'m DMatrix<f64>>, masses: &[&DMatrix<f64>],
added_mass: &DVector<f64>, added_mass: &DVector<f64>,
free_index: &[Option<usize>], free_index: &[Option<usize>],
residual: &DVector<f64>, residual: &DVector<f64>,
@@ -240,24 +305,33 @@ impl SparseTangent {
force_fresh: bool, force_fresh: bool,
rate: f64, rate: f64,
) -> FeaResult<DVector<f64>> { ) -> FeaResult<DVector<f64>> {
#[allow(clippy::float_cmp)] if self.needs_refresh(force_fresh, coef, rate) {
let same_coef = coef == self.factor_coef; if !self.tangents_current {
let refresh = return Err(crate::error::SolverError::FactorizationFailed {
force_fresh || !self.factor_valid || self.uses_left == 0 || !same_coef || rate > 0.5; reason: "sparse tangent: refresh without the current element tangents"
if refresh { .to_string(),
let started = Instant::now();
self.values.fill(0.0);
for ((k_e, mass), pos) in self.element_tangents.iter().zip(masses).zip(&self.elem_pos) {
let n_e = k_e.nrows();
for a in 0..n_e {
for b in 0..n_e {
let p = pos[a * n_e + b];
if p != u32::MAX {
self.values[p as usize] += k_e[(a, b)] + coef * mass[(a, b)];
}
}
} }
.into());
} }
let started = Instant::now();
let tangents = &self.element_tangents;
let contrib_ptr = &self.contrib_ptr;
let contrib = &self.contrib;
self.values
.par_chunks_mut(SCATTER_CHUNK)
.enumerate()
.for_each(|(chunk, values)| {
let base = chunk * SCATTER_CHUNK;
for (offset, value) in values.iter_mut().enumerate() {
let p = base + offset;
let mut sum = 0.0;
for &(e, l) in &contrib[contrib_ptr[p]..contrib_ptr[p + 1]] {
let (e, l) = (e as usize, l as usize);
sum += tangents[e].as_slice()[l] + coef * masses[e].as_slice()[l];
}
*value = sum;
}
});
for (dof, &m) in added_mass.iter().enumerate() { for (dof, &m) in added_mass.iter().enumerate() {
if m != 0.0 { if m != 0.0 {
if let Some(free) = free_index[dof] { if let Some(free) = free_index[dof] {
@@ -265,8 +339,10 @@ impl SparseTangent {
} }
} }
} }
let scattered = Instant::now();
self.stats.scatter_seconds += (scattered - started).as_secs_f64();
let outcome = self.solver.factorize_values(&self.values); let outcome = self.solver.factorize_values(&self.values);
self.stats.factor_seconds += started.elapsed().as_secs_f64(); self.stats.factor_seconds += scattered.elapsed().as_secs_f64();
if let Err(error) = outcome { if let Err(error) = outcome {
// A failed pivot: this solve goes to the pivoting banded // A failed pivot: this solve goes to the pivoting banded
// LU on the same matrix; the next one tries LDLᵀ again. // LU on the same matrix; the next one tries LDLᵀ again.
@@ -194,3 +194,96 @@ pub fn internal_force_and_tangent(
Ok((internal_force, tangent)) Ok((internal_force, tangent))
} }
/// The internal force alone of [`internal_force_and_tangent`] — the same
/// operations in the same order for `f_int` (so the result is identical
/// bit for bit), without the tangent's `B_Lᵀ C B_L` and geometric terms.
/// For Newton iterations that reuse a factorised tangent (modified
/// Newton, R8-g).
pub fn internal_force(
element: &dyn FiniteElement,
node_coords: &[Vector3<f64>],
element_displacement: &DVector<f64>,
constitutive: &dyn Fn(&DVector<f64>) -> FeaResult<(DVector<f64>, DMatrix<f64>)>,
quadrature_order: Option<usize>,
) -> FeaResult<DVector<f64>> {
let quad_rule = element.quadrature_rule(quadrature_order)?;
let dim = element.spatial_dimension();
let num_nodes = element.num_nodes();
let total_dofs = num_nodes * dim;
if element_displacement.len() != total_dofs {
return Err(ElementError::MatrixComputationFailed {
reason: format!(
"element displacement has {} entries, element has {total_dofs} DOFs",
element_displacement.len()
),
}
.into());
}
if dim != 2 && dim != 3 {
return Err(ElementError::MatrixComputationFailed {
reason: format!("total-Lagrangian formulation needs 2-D or 3-D, got {dim}"),
}
.into());
}
let n_strain = if dim == 2 { 3 } else { 6 };
let pairs: Vec<(usize, usize)> = if dim == 2 {
vec![(0, 0), (1, 1), (0, 1)]
} else {
vec![(0, 0), (1, 1), (2, 2), (1, 2), (0, 2), (0, 1)]
};
let mut internal_force = DVector::zeros(total_dofs);
for point in &quad_rule.points {
let shape_eval = element.shape_functions(&point.coords)?;
let jacobian_eval = element.jacobian(&point.coords, node_coords)?;
if !jacobian_eval.is_valid() {
return Err(ElementError::JacobianSingular {
det: jacobian_eval.determinant,
}
.into());
}
let g = jacobian_eval.transform_derivatives(&shape_eval.derivatives)?;
let mut f: DMatrix<f64> = DMatrix::identity(dim, dim);
for a in 0..num_nodes {
for i in 0..dim {
let ui = element_displacement[a * dim + i];
for j in 0..dim {
f[(i, j)] += ui * g[(a, j)];
}
}
}
let c_right = f.transpose() * &f;
let mut e_voigt: DVector<f64> = DVector::zeros(n_strain);
for (k, &(i, j)) in pairs.iter().enumerate() {
let e_ij = 0.5 * (c_right[(i, j)] - if i == j { 1.0 } else { 0.0 });
e_voigt[k] = if i == j { e_ij } else { 2.0 * e_ij };
}
let (s_voigt, _) = constitutive(&e_voigt)?;
if s_voigt.len() != n_strain {
return Err(ElementError::MatrixComputationFailed {
reason: format!(
"constitutive closure returned stress of {} components, expected {n_strain}",
s_voigt.len()
),
}
.into());
}
let mut b_l: DMatrix<f64> = DMatrix::zeros(n_strain, total_dofs);
for (k, &(i, j)) in pairs.iter().enumerate() {
for a in 0..num_nodes {
for m in 0..dim {
let col = a * dim + m;
b_l[(k, col)] = if i == j {
f[(m, i)] * g[(a, i)]
} else {
f[(m, i)] * g[(a, j)] + f[(m, j)] * g[(a, i)]
};
}
}
}
let scale = jacobian_eval.determinant().abs() * point.weight;
internal_force += b_l.transpose() * &s_voigt * scale;
}
Ok(internal_force)
}
@@ -21,6 +21,8 @@ const PANEL: usize = 32;
const TRAIL_BLOCK: usize = 64; const TRAIL_BLOCK: usize = 64;
/// Subtrees below this many multiply–adds are factorised serially. /// Subtrees below this many multiply–adds are factorised serially.
const PARALLEL_WORK: f64 = 2.0e6; const PARALLEL_WORK: f64 = 2.0e6;
/// Trailing updates below this many multiply–adds (×2) run serially.
const PARALLEL_FRONT_WORK: f64 = 4.0e6;
/// A zero, non-finite or relatively tiny pivot. /// A zero, non-finite or relatively tiny pivot.
#[derive(Debug, Clone)] #[derive(Debug, Clone)]
@@ -177,8 +179,13 @@ fn partial_ldlt(front: &mut [f64], m: usize, k: usize, floor: f64) -> Result<(),
} }
} }
let (left, right) = front.split_at_mut(p1 * m); let (left, right) = front.split_at_mut(p1 * m);
let mut c0 = 0; let left: &[f64] = left;
while c0 < r { let w: &[f64] = &w;
// Each column block of the trailing matrix is an independent
// gemm; large fronts run them on rayon (the arithmetic of a
// block does not depend on the schedule).
let block_update = |(t, chunk): (usize, &mut [f64])| {
let c0 = t * TRAIL_BLOCK;
let c1 = (c0 + TRAIL_BLOCK).min(r); let c1 = (c0 + TRAIL_BLOCK).min(r);
let rows = r - c0; let rows = r - c0;
// L_P rows (p1 + c0)..m: a view into `left` (columns p0..p1). // L_P rows (p1 + c0)..m: a view into `left` (columns p0..p1).
@@ -198,14 +205,24 @@ fn partial_ldlt(front: &mut [f64], m: usize, k: usize, floor: f64) -> Result<(),
nalgebra::Dyn(r), nalgebra::Dyn(r),
); );
let mut target = DMatrixViewMut::from_slice_with_strides_generic( let mut target = DMatrixViewMut::from_slice_with_strides_generic(
&mut right[c0 * m + p1 + c0..], &mut chunk[p1 + c0..],
nalgebra::Dyn(rows), nalgebra::Dyn(rows),
nalgebra::Dyn(c1 - c0), nalgebra::Dyn(c1 - c0),
nalgebra::Dyn(1), nalgebra::Dyn(1),
nalgebra::Dyn(m), nalgebra::Dyn(m),
); );
target.gemm(-1.0, &lp, &wb.transpose(), 1.0); target.gemm(-1.0, &lp, &wb.transpose(), 1.0);
c0 = c1; };
if (r * r * nb) as f64 > PARALLEL_FRONT_WORK {
right
.par_chunks_mut(TRAIL_BLOCK * m)
.enumerate()
.for_each(block_update);
} else {
right
.chunks_mut(TRAIL_BLOCK * m)
.enumerate()
.for_each(block_update);
} }
} }
p0 = p1; p0 = p1;
@@ -213,49 +230,144 @@ fn partial_ldlt(front: &mut [f64], m: usize, k: usize, floor: f64) -> Result<(),
Ok(()) Ok(())
} }
/// Solve `L D Lᵀ x = b` in the factor numbering, in place. /// Subtrees below this factorisation work solve serially.
const PARALLEL_SOLVE_WORK: f64 = 2.0e7;
/// Solve `L D Lᵀ x = b` in the factor numbering, in place, over the
/// assembly tree: the forward solve is multifrontal (each supernode
/// solves its dense unit-triangular block and hands its parent an update
/// vector, children in order), the backward solve top-down (each
/// supernode gathers its off-diagonal unknowns from its parent's front
/// vector); independent subtrees run on rayon, and every supernode's
/// arithmetic is independent of the schedule.
pub(super) fn solve_in_place(sym: &Symbolic, factor: &Factor, x: &mut [f64]) { pub(super) fn solve_in_place(sym: &Symbolic, factor: &Factor, x: &mut [f64]) {
let ns = sym.num_supernodes();
let blocks: Vec<std::sync::MutexGuard<'_, Vec<f64>>> = let blocks: Vec<std::sync::MutexGuard<'_, Vec<f64>>> =
factor.blocks.iter().map(|b| b.lock().unwrap()).collect(); factor.blocks.iter().map(|b| b.lock().unwrap()).collect();
for s in 0..ns { let blocks: Vec<&[f64]> = blocks.iter().map(|b| b.as_slice()).collect();
let rows = sym.front_rows(s); // One disjoint output slice per supernode (its own columns).
let m = rows.len(); let mut slots: Vec<Mutex<&mut [f64]>> = Vec::with_capacity(sym.num_supernodes());
let f = sym.start[s]; {
let k = sym.start[s + 1] - f; let mut rest: &mut [f64] = x;
let block = &blocks[s]; for s in 0..sym.num_supernodes() {
let k = sym.start[s + 1] - sym.start[s];
let (own, tail) = rest.split_at_mut(k);
slots.push(Mutex::new(own));
rest = tail;
}
}
let ctx = SolveCtx {
sym,
blocks: &blocks,
slots: &slots,
};
let roots = &sym.roots;
let run_forward = |&r: &usize| {
ctx.forward(r);
};
if roots.len() > 1 {
roots.par_iter().for_each(run_forward);
} else {
roots.iter().for_each(run_forward);
}
let run_backward = |&r: &usize| ctx.backward(r, &[]);
if roots.len() > 1 {
roots.par_iter().for_each(run_backward);
} else {
roots.iter().for_each(run_backward);
}
}
struct SolveCtx<'a, 'x> {
sym: &'a Symbolic,
blocks: &'a [&'a [f64]],
slots: &'a [Mutex<&'x mut [f64]>],
}
impl SolveCtx<'_, '_> {
/// Forward (`L y = b`, then `z = D⁻¹ y` on the own columns) for the
/// subtree at `s`; returns the update vector for the parent's front.
fn forward(&self, s: usize) -> Vec<f64> {
let sym = self.sym;
let kids = sym.node_children(s);
let updates: Vec<Vec<f64>> = if kids.len() > 1 && sym.subtree_work[s] > PARALLEL_SOLVE_WORK
{
kids.par_iter().map(|&c| self.forward(c)).collect()
} else {
kids.iter().map(|&c| self.forward(c)).collect()
};
let m = sym.rows_ptr[s + 1] - sym.rows_ptr[s];
let k = sym.start[s + 1] - sym.start[s];
let block = self.blocks[s];
let mut v = vec![0.0f64; m];
{
let own = self.slots[s].lock().unwrap();
v[..k].copy_from_slice(&own);
}
for (&c, update) in kids.iter().zip(&updates) {
let rel = &sym.relind[sym.relind_ptr[c]..sym.relind_ptr[c + 1]];
for (&r, &u) in rel.iter().zip(update) {
v[r as usize] += u;
}
}
let (head, tail) = v.split_at_mut(k);
for j in 0..k { for j in 0..k {
let xj = x[f + j]; let xj = head[j];
if xj != 0.0 { if xj != 0.0 {
let col = &block[j * m..(j + 1) * m]; let col = &block[j * m..(j + 1) * m];
for i in (j + 1)..m { for (y, &l) in head[j + 1..].iter_mut().zip(&col[j + 1..k]) {
x[rows[i]] -= col[i] * xj; *y -= l * xj;
}
for (y, &l) in tail.iter_mut().zip(&col[k..]) {
*y -= l * xj;
} }
} }
} }
{
let mut own = self.slots[s].lock().unwrap();
for j in 0..k {
own[j] = head[j] / block[j + j * m];
}
}
v.drain(..k);
v
} }
for s in 0..ns {
/// Backward (`Lᵀ x = z`) for the subtree at `s`, given the parent's
/// front vector (the solution at the parent's front rows).
fn backward(&self, s: usize, parent_front: &[f64]) {
let sym = self.sym;
let m = sym.rows_ptr[s + 1] - sym.rows_ptr[s]; let m = sym.rows_ptr[s + 1] - sym.rows_ptr[s];
let f = sym.start[s]; let k = sym.start[s + 1] - sym.start[s];
let k = sym.start[s + 1] - f; let block = self.blocks[s];
let block = &blocks[s]; let mut front = vec![0.0f64; m];
for j in 0..k { let rel = &sym.relind[sym.relind_ptr[s]..sym.relind_ptr[s + 1]];
x[f + j] /= block[j + j * m]; for (slot, &r) in front[k..].iter_mut().zip(rel) {
*slot = parent_front[r as usize];
}
{
let mut own = self.slots[s].lock().unwrap();
let (head, tail) = front.split_at_mut(k);
head.copy_from_slice(&own);
for j in (0..k).rev() {
let col = &block[j * m..(j + 1) * m];
let mut acc = 0.0;
for (&l, &y) in col[k..].iter().zip(tail.iter()) {
acc += l * y;
}
for i in (j + 1)..k {
acc += col[i] * head[i];
}
head[j] -= acc;
}
own.copy_from_slice(head);
} }
} let kids = sym.node_children(s);
for s in (0..ns).rev() { if kids.len() > 1 && sym.subtree_work[s] > PARALLEL_SOLVE_WORK {
let rows = sym.front_rows(s); kids.par_iter().for_each(|&c| self.backward(c, &front));
let m = rows.len(); } else {
let f = sym.start[s]; for &c in kids {
let k = sym.start[s + 1] - f; self.backward(c, &front);
let block = &blocks[s];
for j in (0..k).rev() {
let col = &block[j * m..(j + 1) * m];
let mut acc = x[f + j];
for i in (j + 1)..m {
acc -= col[i] * x[rows[i]];
} }
x[f + j] = acc;
} }
} }
} }
@@ -12,7 +12,10 @@
//! (12×2×2), 40 CSM3 steps. //! (12×2×2), 40 CSM3 steps.
//! 4. `modified_newton_reuses_the_factor` — `reuse = 4`: fewer //! 4. `modified_newton_reuses_the_factor` — `reuse = 4`: fewer
//! factorisations than solves, same march to the Newton tolerance. //! factorisations than solves, same march to the Newton tolerance.
//! 5. `tangent_knob_parses` — the `RTX_FEA_TANGENT` values. //! 5. `force_only_kernel_is_bit_identical` — the TL force-only kernel
//! (modified Newton's reused iterations) = the pair's force, bit for
//! bit, on Hex20 and Quad8.
//! 6. `tangent_knob_parses` — the solver enum's defaults.
//! //!
//! Instruments (`#[ignore]`, env-driven, write under `R8G_OUT`): //! Instruments (`#[ignore]`, env-driven, write under `R8G_OUT`):
//! //!
@@ -306,6 +309,34 @@ fn modified_newton_reuses_the_factor() {
assert!(r < 1e-4, "modified Newton drifts: {r:.3e}"); assert!(r < 1e-4, "modified Newton drifts: {r:.3e}");
} }
#[test]
fn force_only_kernel_is_bit_identical() {
use rtx_fea::elements::total_lagrangian::{
internal_force, internal_force_and_tangent, saint_venant_kirchhoff,
};
let flag = Flag3d::build(Flag3dSpec::turek_hron(0.1, -0.05, 3, 1, 1)).unwrap();
let quad = quad8_flag(3, 1);
for mesh in [&flag.mesh, &quad] {
let dim = mesh.spatial_dimension;
let constitutive = saint_venant_kirchhoff(8.0e5, 5.0e5, dim);
for (e, element) in mesh.elements.values().enumerate() {
let coords: Vec<Vector3<f64>> = element
.nodes
.iter()
.map(|id| mesh.get_node(*id).unwrap().position())
.collect();
let fe = StandardFiniteElement::new(element.element_type, coords.clone());
let n = coords.len() * dim;
let u = DVector::from_fn(n, |i, _| 1e-3 * ((i + 7 * e) as f64 * 0.73).sin());
let (f_pair, _) =
internal_force_and_tangent(&fe, &coords, &u, constitutive.as_ref(), None).unwrap();
let f_only = internal_force(&fe, &coords, &u, constitutive.as_ref(), None).unwrap();
assert!(f_pair.norm() > 0.0);
assert_eq!(f_pair, f_only, "force-only kernel differs");
}
}
}
#[test] #[test]
fn tangent_knob_parses() { fn tangent_knob_parses() {
// Only the parser (the knob is read when a stepper is built). // Only the parser (the knob is read when a stepper is built).
@@ -527,12 +558,14 @@ fn r8g_g2_cost() {
} }
let breakdown = match (before, stepper.tangent_stats()) { let breakdown = match (before, stepper.tangent_stats()) {
(Some(b), Some(a)) => format!( (Some(b), Some(a)) => format!(
" | assemblies {} ({:.4} s each), factorisations {} ({:.4} s each incl. \ " | assemblies {} ({:.4} s each), factorisations {} (scatter {:.4} + \
scatter), solves {} ({:.4} s each), fallbacks {}, nnz(L) {}", numeric {:.4} s each), solves {} ({:.4} s each), fallbacks {}, nnz(L) {}",
a.assemblies - b.assemblies, a.assemblies - b.assemblies,
(a.assembly_seconds - b.assembly_seconds) (a.assembly_seconds - b.assembly_seconds)
/ (a.assemblies - b.assemblies).max(1) as f64, / (a.assemblies - b.assemblies).max(1) as f64,
a.factorizations - b.factorizations, a.factorizations - b.factorizations,
(a.scatter_seconds - b.scatter_seconds)
/ (a.factorizations - b.factorizations).max(1) as f64,
(a.factor_seconds - b.factor_seconds) (a.factor_seconds - b.factor_seconds)
/ (a.factorizations - b.factorizations).max(1) as f64, / (a.factorizations - b.factorizations).max(1) as f64,
a.solves - b.solves, a.solves - b.solves,