diff --git a/src/qdldl/qdldl.rs b/src/qdldl/qdldl.rs index b36b82e6..cae71ceb 100644 --- a/src/qdldl/qdldl.rs +++ b/src/qdldl/qdldl.rs @@ -325,8 +325,44 @@ struct QDLDLWorkspace { // number of regularized entries in D regularize_count: usize, + + // wavefront schedule for the parallel numeric factorization: nodes 1..n + // in ascending elimination-tree height order, with per-height offsets + wf_order: Vec, + wf_offsets: Vec, + + // lazily allocated per-thread scratch for the parallel factorization + par_scratch: Vec>, +} + +// per-thread scratch mirroring the y_markers/y_idx/elim_buffer/y_vals split +// of the serial factorization workspace +#[derive(Debug)] +struct ParScratch { + y_markers: Vec, + y_idx: Vec, + elim_buffer: Vec, + y_vals: Vec, +} + +impl ParScratch { + fn new(n: usize) -> Self { + Self { + y_markers: vec![QDLDL_UNUSED; n], + y_idx: vec![0; n], + elim_buffer: vec![0; n], + y_vals: vec![T::zero(); n], + } + } } +// raw-pointer wrapper for the shared factorization outputs; safety of the +// shared mutation is argued at the use site in _factor_inner_parallel +#[derive(Clone, Copy)] +struct SendPtr

(P); +unsafe impl

Send for SendPtr

{} +unsafe impl

Sync for SendPtr

{} + impl QDLDLWorkspace where T: FloatT, @@ -355,6 +391,35 @@ where &mut etree, )?; + // wavefront schedule: group nodes by elimination-tree height. Parents + // always have larger indices than their children, so a single ascending + // pass computes heights. + let n = triuA.ncols(); + let mut height = vec![0usize; n]; + for k in 0..n { + let p = etree[k]; + if p != QDLDL_UNKNOWN { + height[p] = height[p].max(height[k] + 1); + } + } + let maxh = height.iter().copied().max().unwrap_or(0); + // counting sort of nodes 1..n by height (node 0 is handled separately + // by the factorization prologue) + let mut wf_offsets = vec![0usize; maxh + 2]; + for k in 1..n { + wf_offsets[height[k] + 1] += 1; + } + for h in 0..(maxh + 1) { + wf_offsets[h + 1] += wf_offsets[h]; + } + let mut wf_order = vec![0usize; n.saturating_sub(1)]; + let mut pos = wf_offsets.clone(); + for k in 1..n { + let h = height[k]; + wf_order[pos[h]] = k; + pos[h] += 1; + } + // positive inertia count. let positive_inertia = 0; @@ -375,6 +440,9 @@ where regularize_eps, regularize_delta, regularize_count, + wf_order, + wf_offsets, + par_scratch: Vec::new(), }) } } @@ -392,6 +460,16 @@ fn _factor( Dinv.fill(T::one()); } + // Use the wavefront-parallel numeric factorization for large problems. + // The serial path is kept bit-identical for small problems and for + // logical (symbolic-only) factorization. + const PARALLEL_MIN_DIM: usize = 20_000; + if !logical && workspace.triuA.ncols() >= PARALLEL_MIN_DIM && rayon::current_num_threads() > 1 { + let pos_d_count = _factor_inner_parallel(L, D, Dinv, workspace)?; + workspace.positive_inertia = pos_d_count; + return Ok(()); + } + // factor using QDLDL C style converted code let A = &workspace.triuA; @@ -668,6 +746,274 @@ fn _factor_inner( Ok(positiveValuesInD) } +// Wavefront-parallel variant of _factor_inner. +// +// Rows are processed in ascending elimination-tree height order; rows of equal +// height are processed in parallel. Safety of the shared mutation relies on +// standard elimination-tree structure properties: +// +// * the sparsity pattern of row k of L (the columns its triangular solve +// reads, and the columns it appends one entry to) consists only of +// DESCENDANTS of k in the elimination tree; +// * every entry of column j of L is written by a row that is an ANCESTOR of +// j; ancestors of j form a chain with strictly increasing heights (and +// strictly increasing row indices), so within one wavefront at most one +// row touches any given column, appends across wavefronts occur in +// ascending row order (preserving CSC ordering), and every column read by +// a row was completed at strictly lower heights; +// * two rows of equal height have disjoint descendant sets (a common +// descendant would make one an ancestor of the other, contradicting equal +// heights). +// +// D[k]/Dinv[k] are written only by row k and read only for descendants +// (strictly lower heights). Each worker uses its own y workspace. Results are +// bitwise identical to the serial factorization: every entry is computed from +// the same inputs with the same inner loop order, independent of scheduling. +#[allow(non_snake_case)] +fn _factor_inner_parallel( + L: &mut CscMatrix, + D: &mut [T], + Dinv: &mut [T], + workspace: &mut QDLDLWorkspace, +) -> Result { + use core::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; + + let A = &workspace.triuA; + let n = A.n; + let (Ap, Ai, Ax) = (&A.colptr, &A.rowval, &A.nzval); + let Lnz = &workspace.Lnz; + let etree = &workspace.etree; + let Dsigns = &workspace.Dsigns; + let regularize_enable = workspace.regularize_enable; + let regularize_eps = workspace.regularize_eps; + let regularize_delta = workspace.regularize_delta; + + let nthreads = rayon::current_num_threads(); + if workspace.par_scratch.len() < nthreads || workspace.par_scratch.first().is_some_and(|w| w.y_vals.len() != n) { + workspace.par_scratch = (0..nthreads).map(|_| ParScratch::new(n)).collect(); + } + + // prologue, identical to _factor_inner: Lp = cumsum(Lnz) and the first + // diagonal element + let Lp = &mut L.colptr; + Lp[0] = 0; + let mut acc = 0; + for (Lp, Lnz) in zip(&mut Lp[1..], Lnz) { + acc += Lnz; + *Lp = acc; + } + + D.fill(T::zero()); + // next available space in each column of L; stored in the same iwork + // segment the serial code uses, though only this function touches it here + let (_, iwork) = workspace.iwork.split_at_mut(n); + let (_, next_colspace) = iwork.split_at_mut(n); + next_colspace.copy_from_slice(&Lp[0..Lp.len() - 1]); + + let mut regularize_count_total = 0usize; + let mut positive_total = 0usize; + + D[0] = Ax[0]; + if regularize_enable { + let sign = T::from_i8(Dsigns[0]).unwrap(); + if D[0] * sign < regularize_eps { + D[0] = regularize_delta * sign; + regularize_count_total += 1; + } + } + if D[0].is_zero() { + return Err(QDLDLError::ZeroPivot); + } + if D[0] > T::zero() { + positive_total += 1; + } + Dinv[0] = T::recip(D[0]); + + // shared outputs, mutated disjointly per the wavefront argument above + let Lp_ptr = SendPtr(Lp.as_ptr()); + let Li_ptr = SendPtr(L.rowval.as_mut_ptr()); + let Lx_ptr = SendPtr(L.nzval.as_mut_ptr()); + let D_ptr = SendPtr(D.as_mut_ptr()); + let Dinv_ptr = SendPtr(Dinv.as_mut_ptr()); + let colspace_ptr = SendPtr(next_colspace.as_mut_ptr()); + + let zero_pivot = AtomicBool::new(false); + let positive_count = AtomicUsize::new(0); + let regularize_count_atomic = AtomicUsize::new(0); + + // the per-row kernel; caller guarantees rows within a call batch are of + // equal elimination-tree height and scratch is exclusive to this call + let factor_row = |k: usize, scratch: &mut ParScratch| -> bool { + let y_markers = &mut scratch.y_markers; + let y_idx = &mut scratch.y_idx; + let elim_buffer = &mut scratch.elim_buffer; + let y_vals = &mut scratch.y_vals; + + let (Lp, Li, Lx) = (Lp_ptr, Li_ptr, Lx_ptr); + let (D, Dinv, next_colspace) = (D_ptr, Dinv_ptr, colspace_ptr); + + let mut nnz_y = 0usize; + let mut dk = T::zero(); + + for i in Ap[k]..Ap[k + 1] { + let bidx = Ai[i]; + if bidx == k { + dk = Ax[i]; + continue; + } + + y_vals[bidx] = Ax[i]; + + let next_idx = bidx; + if y_markers[next_idx] == QDLDL_UNUSED { + y_markers[next_idx] = QDLDL_USED; + elim_buffer[0] = next_idx; + let mut nnz_e = 1; + + let mut next_idx = etree[bidx]; + while next_idx != QDLDL_UNKNOWN && next_idx < k { + if y_markers[next_idx] == QDLDL_USED { + break; + } + y_markers[next_idx] = QDLDL_USED; + elim_buffer[nnz_e] = next_idx; + next_idx = etree[next_idx]; + nnz_e += 1; + } + + while nnz_e != 0 { + nnz_e -= 1; + y_idx[nnz_y] = elim_buffer[nnz_e]; + nnz_y += 1; + } + } + } + + for i in (0..nnz_y).rev() { + let cidx = y_idx[i]; + + unsafe { + // Safety: cidx is a descendant of k, so no other row in this + // wavefront reads or writes column cidx, its L entries, D, or + // its next_colspace slot; all entries read were completed at + // strictly lower wavefronts. + let tmp_idx = *next_colspace.0.add(cidx); + let y_vals_cidx = y_vals[cidx]; + + let (f, l) = (*Lp_ptr.0.add(cidx), tmp_idx); + for idx in f..l { + let lij = *Li.0.add(idx); + *(y_vals.get_unchecked_mut(lij)) -= *Lx.0.add(idx) * y_vals_cidx; + } + + let lx_tmp = y_vals_cidx * *Dinv.0.add(cidx); + *Lx.0.add(tmp_idx) = lx_tmp; + dk -= y_vals_cidx * lx_tmp; + + *Li.0.add(tmp_idx) = k; + *next_colspace.0.add(cidx) = tmp_idx + 1; + } + + y_vals[cidx] = T::zero(); + y_markers[cidx] = QDLDL_UNUSED; + } + + if regularize_enable { + let sign = T::from_i8(Dsigns[k]).unwrap(); + if dk * sign < regularize_eps { + dk = regularize_delta * sign; + regularize_count_atomic.fetch_add(1, Ordering::Relaxed); + } + } + + if dk.is_zero() { + zero_pivot.store(true, Ordering::Relaxed); + return false; + } + if dk > T::zero() { + positive_count.fetch_add(1, Ordering::Relaxed); + } + + unsafe { + *D.0.add(k) = dk; + *Dinv.0.add(k) = T::recip(dk); + } + let _ = Lp; + true + }; + + // process the wavefronts; small levels run on the current thread to avoid + // scheduling overhead, larger ones fan out across a scope + const PARALLEL_MIN_LEVEL: usize = 32; + let wf_offsets = &workspace.wf_offsets; + let wf_order = &workspace.wf_order; + + let mut scratches: Vec<&mut ParScratch> = workspace.par_scratch.iter_mut().collect(); + + if std::env::var("CLARABEL_DEBUG_WF").is_ok() { + let nlv = wf_offsets.len() - 1; + let mut wide_rows = 0usize; + let mut wide_work = 0usize; + let mut total_work = 0usize; + for h in 0..nlv { + let level = &wf_order[wf_offsets[h]..wf_offsets[h + 1]]; + let work: usize = level.iter().map(|&k| Lnz[k]).sum(); + total_work += work; + if level.len() >= PARALLEL_MIN_LEVEL { + wide_rows += level.len(); + wide_work += work; + } + } + eprintln!( + "CLARABEL_DEBUG_WF: n={} levels={} wide_rows={} ({:.1}%) wide_Lnz={} of {} ({:.1}%)", + n, nlv, wide_rows, + 100.0 * wide_rows as f64 / (n - 1) as f64, + wide_work, total_work, + 100.0 * wide_work as f64 / total_work.max(1) as f64, + ); + } + + for h in 0..wf_offsets.len() - 1 { + if zero_pivot.load(Ordering::Relaxed) { + break; + } + let level = &wf_order[wf_offsets[h]..wf_offsets[h + 1]]; + if level.is_empty() { + continue; + } + + if level.len() < PARALLEL_MIN_LEVEL { + let scratch = &mut scratches[0]; + for &k in level { + if !factor_row(k, scratch) { + break; + } + } + } else { + let chunk = level.len().div_ceil(scratches.len()); + rayon::scope(|s| { + for (rows, scratch) in zip(level.chunks(chunk), scratches.iter_mut()) { + let factor_row = &factor_row; + s.spawn(move |_| { + for &k in rows { + if !factor_row(k, scratch) { + break; + } + } + }); + } + }); + } + } + + if zero_pivot.load(Ordering::Relaxed) { + return Err(QDLDLError::ZeroPivot); + } + + workspace.regularize_count = regularize_count_total + regularize_count_atomic.load(Ordering::Relaxed); + Ok(positive_total + positive_count.load(Ordering::Relaxed)) +} + // Solves (L+I)x = b, with x replacing b (with standard bounds checks) fn _lsolve_safe(Lp: &[usize], Li: &[usize], Lx: &[T], x: &mut [T]) { for i in 0..x.len() {