From ad719457dd2a5eb4b244bc7cbbea765111b46919 Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 27 Sep 2026 05:56:37 +0000 Subject: [PATCH] perf(lindblad): store the expm generator cache as blocked flat CSC build_mf_cols / build_orbit_rep_cols returned one Vec per column, each reserved at the full L* output count although only in-basis entries are kept. On large bases this over-reserves the cache severalfold and leaves one small allocation per basis string per expm call for the allocator to retain. The cache is now a BlockCsc: blocks of up to 4096 columns, each an exactly sized offsets/rows/vals triple (structure of arrays, 12 B per real entry instead of 16 B). CscOp::dot keeps the same per-thread column partition, so results are bit-identical. Co-Authored-By: Claude Opus 5.5 --- crates/ppvm-lindblad/src/mf_expm.rs | 243 +++++++++++++++++----------- 1 file changed, 150 insertions(+), 93 deletions(-) diff --git a/crates/ppvm-lindblad/src/mf_expm.rs b/crates/ppvm-lindblad/src/mf_expm.rs index 9437396f8..3eb8c7ab7 100644 --- a/crates/ppvm-lindblad/src/mf_expm.rs +++ b/crates/ppvm-lindblad/src/mf_expm.rs @@ -35,12 +35,101 @@ use rayon::prelude::*; use std::iter::Sum; use std::ops::{AddAssign, Div, Mul, Sub}; -/// CSC columns of a cached in-basis action: `cols[c]` = `(row, coeff)`. -type Cols = Vec>; /// Per-column `(raw, diag)` for the `μ`/1-norm selection: `raw` bounds /// `Σ_r |M[r,c]|` from above and `diag = M[c,c]`. type PerCol = Vec<(f64, T)>; +/// Scratch buffers for [`LindbladSpec::compute_action_terms`]. +type ActionScratch = (Vec, Vec, FxHashMap>); + +/// Consecutive CSC columns stored flat: local column `j` holds +/// `rows[offsets[j]..offsets[j + 1]]` and the matching `vals`. +struct CscBlock { + offsets: Vec, + rows: Vec, + vals: Vec, +} + +/// Cached in-basis action in CSC form, stored as blocks of `block` columns. +/// +/// One exactly-sized allocation triple per block replaces one `Vec` per +/// column: at `|basis| ~ 10^6` the per-column layout reserved every `L*` +/// output (in- and out-of-basis) and left ~10^6 small allocations for the +/// system allocator to retain after the expm call. +pub(crate) struct BlockCsc { + blocks: Vec>, + block: usize, + dim: usize, +} + +impl BlockCsc { + /// Visit the columns `range` in order as `(col, rows, vals)`. + fn for_each_col(&self, range: std::ops::Range, mut f: impl FnMut(usize, &[u32], &[T])) { + let mut c = range.start; + while c < range.end { + let b = &self.blocks[c / self.block]; + let base = (c / self.block) * self.block; + let stop = range.end.min(base + b.offsets.len() - 1); + for j in (c - base)..(stop - base) { + let (lo, hi) = (b.offsets[j] as usize, b.offsets[j + 1] as usize); + f(base + j, &b.rows[lo..hi], &b.vals[lo..hi]); + } + c = stop; + } + } +} + +/// Build the [`BlockCsc`] cache and the per-column `(raw, diag)` data for +/// a `dim`-column generator in one parallel pass. `col(c, scratch, rows, +/// vals)` appends the in-basis entries of column `c` to `rows`/`vals` and +/// returns its `(raw, diag)`. +fn build_block_csc(spec: &LindbladSpec, dim: usize, col: F) -> (BlockCsc, PerCol) +where + T: Copy + Send + Sync, + F: Fn(usize, &mut ActionScratch, &mut Vec, &mut Vec) -> (f64, T) + Sync, +{ + // ~16 blocks per thread for load balance, but never so small that the + // per-block allocations matter. + let block = dim + .div_ceil(16 * rayon::current_num_threads().max(1)) + .clamp(64, 4096); + let (blocks, per_col): (Vec>, Vec>) = (0..dim.div_ceil(block)) + .into_par_iter() + .map_init( + || { + let scratch: ActionScratch = ( + Vec::with_capacity(spec.n_qubits()), + Vec::with_capacity(128), + FxHashMap::with_capacity_and_hasher(128, FxBuildHasher::default()), + ); + (scratch, Vec::::new(), Vec::::new()) + }, + |(scratch, rows, vals), b| { + let cols = (b * block)..dim.min((b + 1) * block); + rows.clear(); + vals.clear(); + let mut offsets = Vec::with_capacity(cols.len() + 1); + let mut per_col = Vec::with_capacity(cols.len()); + offsets.push(0); + for c in cols { + per_col.push(col(c, scratch, rows, vals)); + offsets.push(u32::try_from(rows.len()).expect("CSC block exceeds u32 entries")); + } + // `to_vec` sizes the stored block exactly; the staging + // buffers are reused for the next block on this thread. + let blk = CscBlock { + offsets, + rows: rows.to_vec(), + vals: vals.to_vec(), + }; + (blk, per_col) + }, + ) + .unzip(); + let per_col = per_col.into_iter().flatten().collect(); + (BlockCsc { blocks, block, dim }, per_col) +} + /// Per-column in-basis action of the real generator `M`, plus the data the /// `(m, s)`/`μ` selection needs — all from ONE action pass over the basis. /// @@ -54,38 +143,24 @@ fn build_mf_cols( spec: &LindbladSpec, basis: &[Word], index: &FxHashMap, -) -> (Cols, PerCol) { - basis - .par_iter() - .map_init( - || { - ( - Vec::::with_capacity(spec.n_qubits()), - Vec::::with_capacity(128), - FxHashMap::>::with_capacity_and_hasher( - 128, - FxBuildHasher::default(), - ), - ) - }, - |(s1, s2, lm), p| { - let terms = spec.compute_action_terms(p, s1, s2, lm); - let mut out = Vec::with_capacity(terms.len()); - let mut raw = 0.0; - let mut diag = 0.0; - for (w, c) in terms.iter() { - raw += c.abs(); - if w == p { - diag = *c; - } - if let Some(&row) = index.get(w) { - out.push((row, *c)); - } - } - (out, (raw, diag)) - }, - ) - .unzip() +) -> (BlockCsc, PerCol) { + build_block_csc(spec, basis.len(), |c, (s1, s2, lm), rows, vals| { + let p = &basis[c]; + let terms = spec.compute_action_terms(p, s1, s2, lm); + let mut raw = 0.0; + let mut diag = 0.0; + for (w, v) in terms.iter() { + raw += v.abs(); + if w == p { + diag = *v; + } + if let Some(&row) = index.get(w) { + rows.push(row); + vals.push(*v); + } + } + (raw, diag) + }) } /// Per-column **phase-aware** action of the in-basis-restricted orbit-rep @@ -118,48 +193,33 @@ fn build_orbit_rep_cols( basis: &[Word], index: &FxHashMap, sector: Sector<'_>, -) -> (Cols>, PerCol>) { - basis - .par_iter() - .enumerate() - .map_init( - || { - ( - Vec::::with_capacity(spec.n_qubits()), - Vec::::with_capacity(128), - FxHashMap::>::with_capacity_and_hasher( - 128, - FxBuildHasher::default(), - ), - ) - }, - |(s1, s2, lm), (c, r)| { - // A rep that cannot carry the sector has coefficient zero - // identically, so its column is empty. - let Some(orbit_in) = sector.orbit_size(r) else { - return (Vec::new(), (0.0, Complex::new(0.0, 0.0))); - }; - let terms = spec.compute_action_terms(r, s1, s2, lm); - let mut out = Vec::with_capacity(terms.len()); - let mut raw = 0.0; - let mut diag = Complex::new(0.0, 0.0); - for (q, v) in terms.iter() { - let Some((r_q, phase, orbit_out)) = sector.canonicalize_phase(q) else { - continue; - }; - if let Some(&row) = index.get(&r_q) { - let val = phase * *v * (orbit_in as f64 / orbit_out as f64); - raw += val.norm(); - if row as usize == c { - diag += val; - } - out.push((row, val)); - } +) -> (BlockCsc>, PerCol>) { + build_block_csc(spec, basis.len(), |c, (s1, s2, lm), rows, vals| { + let r = &basis[c]; + // A rep that cannot carry the sector has coefficient zero + // identically, so its column is empty. + let Some(orbit_in) = sector.orbit_size(r) else { + return (0.0, Complex::new(0.0, 0.0)); + }; + let terms = spec.compute_action_terms(r, s1, s2, lm); + let mut raw = 0.0; + let mut diag = Complex::new(0.0, 0.0); + for (q, v) in terms.iter() { + let Some((r_q, phase, orbit_out)) = sector.canonicalize_phase(q) else { + continue; + }; + if let Some(&row) = index.get(&r_q) { + let val = phase * *v * (orbit_in as f64 / orbit_out as f64); + raw += val.norm(); + if row as usize == c { + diag += val; } - (out, (raw, diag)) - }, - ) - .unzip() + rows.push(row); + vals.push(val); + } + } + (raw, diag) + }) } /// Borrowed CSC-style view of an in-basis-restricted generator `M`, backed @@ -172,8 +232,7 @@ fn build_orbit_rep_cols( /// impl for `&T`, so `ExpmOp::from_parts(op, ...)` accepts a `CscOp` by /// value while it keeps borrowing `cols`. pub(crate) struct CscOp<'a, T> { - pub(crate) cols: &'a [Vec<(u32, T)>], - pub(crate) dim: usize, + pub(crate) cols: &'a BlockCsc, } impl LinearOperator for CscOp<'_, T> @@ -188,7 +247,7 @@ where + Sync, { fn dim(&self) -> usize { - self.dim + self.cols.dim } fn parallel_hint(&self) -> bool { @@ -199,7 +258,7 @@ where } fn dot(&self, overwrite: bool, input: &[T], output: &mut [T]) -> Result<(), QuSpinError> { - let n = self.dim; + let n = self.cols.dim; if n == 0 { return Ok(()); } @@ -209,22 +268,20 @@ where // Parallelise over column chunks; each thread accumulates into a dense // local `y` of length `dim`, reading the cached action; the partials // are reduced into `output` sequentially at the end. - let partial_ys: Vec> = self - .cols - .par_chunks(chunk_size) - .enumerate() - .map(|(chunk_idx, chunk)| { - let c_offset = chunk_idx * chunk_size; + let partial_ys: Vec> = (0..n.div_ceil(chunk_size)) + .into_par_iter() + .map(|chunk_idx| { + let cols = (chunk_idx * chunk_size)..n.min((chunk_idx + 1) * chunk_size); let mut y_local = vec![T::zero(); n]; - for (c_local, col) in chunk.iter().enumerate() { - let xc = input[c_offset + c_local]; + self.cols.for_each_col(cols, |c, rows, vals| { + let xc = input[c]; if xc == T::zero() { - continue; + return; } - for &(row, val) in col.iter() { + for (&row, &val) in rows.iter().zip(vals) { y_local[row as usize] += val * xc; } - } + }); y_local }) .collect(); @@ -305,7 +362,7 @@ where /// `select` maps `‖dt·(M−μI)‖₁` to `(m*, s, backward-error tol)`; the two /// call sites differ only in that choice. fn expm_apply_cached( - cols: &Cols, + cols: &BlockCsc, per_col: &PerCol, dt: f64, coeffs: &[T], @@ -323,7 +380,7 @@ where + From + Sum, { - let n = cols.len(); + let n = cols.dim; let trace: T = per_col.iter().map(|(_, d)| *d).sum(); let mu = trace / n as f64; let onenorm = per_col @@ -333,7 +390,7 @@ where let (m_star, s, expm_tol) = select(dt.abs() * onenorm); let mut v = coeffs.to_vec(); - let op = CscOp { cols, dim: n }; + let op = CscOp { cols }; let expm = quspin_expm::ExpmOp::from_parts(op, T::from(dt), mu, s as usize, m_star as usize, expm_tol); expm.apply(ndarray::ArrayViewMut1::from(v.as_mut_slice()))