diff --git a/crates/within/src/domain.rs b/crates/within/src/domain.rs index f01a3fc9..ae38e3ed 100644 --- a/crates/within/src/domain.rs +++ b/crates/within/src/domain.rs @@ -19,13 +19,24 @@ use std::borrow::Cow; use crate::observation::ObservationFrame; use crate::BuildError; -/// Per-factor metadata: level count and global DOF offset. -#[derive(Debug, Clone, Copy)] -pub(crate) struct FactorMeta { +/// Per-term metadata: level count, coefficient-block offset, and column +/// structure. Columns are ordered `[intercept?, slopes…]`; coefficient `c` of +/// level `level` lives at `offset + c * n_levels + level`. +#[derive(Debug, Clone)] +pub(crate) struct TermMeta { pub n_levels: usize, pub offset: usize, /// Non-decreasing in the design's internal row order (fixed at construction). pub sorted: bool, + pub intercept: bool, + /// This term's slope loadings, as indices into the frame's continuous columns. + pub slopes: Vec, +} + +impl TermMeta { + pub fn n_dofs(&self) -> usize { + (usize::from(self.intercept) + self.slopes.len()) * self.n_levels + } } /// Fixed-effects design: observation columns plus coefficient-space layout. @@ -33,7 +44,7 @@ pub(crate) struct FactorMeta { pub struct Design<'a> { /// Columns in internal row order (caller's, or an owned locality-sorted copy). pub(crate) frame: ObservationFrame<'a>, - pub(crate) factors: Vec, + pub(crate) terms: Vec, pub(crate) n_obs: usize, pub(crate) n_dofs: usize, /// `obs_perm[k]` = caller's original index of the observation at internal position `k`. @@ -41,43 +52,53 @@ pub struct Design<'a> { } impl<'a> Design<'a> { - /// Lower effect terms onto the categories path; slope-bearing effects are - /// rejected until a later slice. + /// Lower effect terms into a design: level columns plus each term's slope + /// loadings, laid out term-major (`offset[t] + c * L_t + level`). pub fn new(effects: impl IntoIterator>) -> Result { let mut categorical: Vec> = Vec::new(); - for (idx, effect) in effects.into_iter().enumerate() { - if !effect.slopes().is_empty() { - return Err(BuildError::SlopesNotYetSupported { effect: idx }); - } - // `Effect::new` rejects the intercept-less slope-free shape. - debug_assert!(effect.intercept()); + let mut continuous: Vec> = Vec::new(); + let mut columns: Vec<(bool, Vec)> = Vec::new(); + for effect in effects { + let slots = (continuous.len()..continuous.len() + effect.slopes().len()).collect(); + continuous.extend(effect.slopes().iter().map(|&z| Cow::Borrowed(z))); + columns.push((effect.intercept(), slots)); categorical.push(Cow::Borrowed(effect.levels())); } - let frame = ObservationFrame::new(categorical, Vec::new())?; - Design::from_frame(frame) + let frame = ObservationFrame::new(categorical, continuous)?; + Self::build(frame, columns, true) } - /// Construct from a frame, inferring each factor's level count (`max + 1`); - /// locality-sorts all columns when the dominant factor is unsorted. + /// Construct from a frame of plain factors (each an intercept-only term), + /// inferring each factor's level count (`max + 1`); locality-sorts all + /// columns when the dominant factor is unsorted. pub fn from_frame(frame: ObservationFrame<'a>) -> Result { - Self::build(frame, true) + let columns = vec![(true, Vec::new()); frame.n_factors()]; + Self::build(frame, columns, true) } /// [`from_frame`](Self::from_frame) without the locality sort — profiling escape hatch. #[doc(hidden)] pub fn from_frame_unsorted(frame: ObservationFrame<'a>) -> Result { - Self::build(frame, false) + let columns = vec![(true, Vec::new()); frame.n_factors()]; + Self::build(frame, columns, false) } - fn build(frame: ObservationFrame<'a>, locality_sort: bool) -> Result { + /// `column_structure[q]` = the term's `(intercept, slope column indices)`, + /// aligned with the frame's categorical columns. + fn build( + frame: ObservationFrame<'a>, + column_structure: Vec<(bool, Vec)>, + locality_sort: bool, + ) -> Result { if frame.n_obs() == 0 { return Err(BuildError::EmptyObservations); } + debug_assert_eq!(column_structure.len(), frame.n_factors()); let n_obs = frame.n_obs(); - let mut factors = Vec::with_capacity(frame.n_factors()); + let mut terms = Vec::with_capacity(frame.n_factors()); let mut offset = 0; - for q in 0..frame.n_factors() { + for (q, (intercept, slopes)) in column_structure.into_iter().enumerate() { let col = frame.level_column(q); let mut max = 0; let mut sorted = true; @@ -87,21 +108,24 @@ impl<'a> Design<'a> { sorted &= v >= prev; prev = v; } - let n_levels = max as usize + 1; - factors.push(FactorMeta { - n_levels, + let meta = TermMeta { + n_levels: max as usize + 1, offset, sorted, - }); - offset += n_levels; + intercept, + slopes, + }; + offset += meta.n_dofs(); + terms.push(meta); } - // Sort by the highest-cardinality factor so its gather/scatter runs - // sequentially. `obs_perm` indexes observations as u32; beyond - // u32::MAX rows skip the optimization — the solve itself has no such limit. - let dominant = (0..factors.len()).max_by_key(|&q| factors[q].n_levels); + // Sort by the term contributing the most DOFs (for plain factors, the + // highest-cardinality one) so its gather/scatter runs sequentially. + // `obs_perm` indexes observations as u32; beyond u32::MAX rows skip + // the optimization — the solve itself has no such limit. + let dominant = (0..terms.len()).max_by_key(|&q| terms[q].n_dofs()); let (frame, obs_perm) = match dominant { - Some(d) if locality_sort && !factors[d].sorted && u32::try_from(n_obs).is_ok() => { + Some(d) if locality_sort && !terms[d].sorted && u32::try_from(n_obs).is_ok() => { // Stable argsort. Must be `sort_by_cached_key`, NOT `sort_by_key`: // the latter re-gathers `key[i]` O(n log n) times and dominated // setup at tens of millions of rows. @@ -111,7 +135,7 @@ impl<'a> Design<'a> { let sorted_frame = frame.permuted(&perm); // Rescan sortedness: factors nested in (or duplicating) the // dominant one come out sorted, keeping their coalesced scatter. - for (q, meta) in factors.iter_mut().enumerate() { + for (q, meta) in terms.iter_mut().enumerate() { meta.sorted = sorted_frame.level_column(q).is_sorted(); } (sorted_frame, Some(perm)) @@ -121,7 +145,7 @@ impl<'a> Design<'a> { Ok(Design { frame, - factors, + terms, n_obs, n_dofs: offset, obs_perm, @@ -132,7 +156,7 @@ impl<'a> Design<'a> { pub fn into_owned(self) -> Design<'static> { Design { frame: self.frame.into_owned(), - factors: self.factors, + terms: self.terms, n_obs: self.n_obs, n_dofs: self.n_dofs, obs_perm: self.obs_perm, @@ -190,7 +214,7 @@ impl<'a> Design<'a> { /// Number of categorical factors in the design. #[inline] pub fn n_factors(&self) -> usize { - self.factors.len() + self.terms.len() } /// Number of observations (rows of D). @@ -255,9 +279,9 @@ mod tests { // Stable argsort of [2,0,1,0] → original indices [1,3,2,0]. assert_eq!(design.obs_perm.as_deref(), Some(&[1u32, 3, 2, 0][..])); - assert!(design.factors[0].sorted); + assert!(design.terms[0].sorted); // Factor 1's permuted column [0,1,1,0] is no longer non-decreasing. - assert!(!design.factors[1].sorted); + assert!(!design.terms[1].sorted); assert_eq!(design.frame.level_column(0), [0, 0, 1, 2]); assert_eq!(design.frame.level_column(1), [0, 1, 1, 0]); @@ -271,8 +295,8 @@ mod tests { let col1: Vec = col0.iter().map(|&v| v / 2).collect(); let design = Design::from_frame(frame(vec![col0, col1], vec![])).unwrap(); assert!(design.obs_perm.is_some()); - assert!(design.factors[0].sorted); - assert!(design.factors[1].sorted); + assert!(design.terms[0].sorted); + assert!(design.terms[1].sorted); } #[test] @@ -280,8 +304,8 @@ mod tests { let design = Design::from_frame(frame(vec![vec![0, 0, 1, 2], vec![1, 0, 1, 0]], vec![])).unwrap(); assert!(design.obs_perm.is_none()); - assert!(design.factors[0].sorted); - assert!(!design.factors[1].sorted); + assert!(design.terms[0].sorted); + assert!(!design.terms[1].sorted); } #[test] @@ -299,17 +323,35 @@ mod tests { } #[test] - fn new_rejects_slope_bearing_effect_naming_its_index() { - let plain = [0u32, 1, 0, 1]; - let slope = [1.0, 2.0, 3.0, 4.0]; + fn new_lays_out_slope_terms_term_major() { + // Term 0 is dominant (most DOFs); sorted levels keep the locality + // sort a no-op so the frame columns stay in caller order. + let f0 = [0u32, 0, 1, 1]; + let f1 = [0u32, 2, 1, 0]; + let z0 = [1.0, 2.0, 3.0, 4.0]; + let z1 = [5.0, 6.0, 7.0, 8.0]; let effects = vec![ - Effect::new(&plain, true, []).unwrap(), - Effect::new(&plain, true, [&slope[..]]).unwrap(), + Effect::new(&f0, true, [&z0[..], &z1[..]]).unwrap(), + Effect::new(&f1, true, []).unwrap(), + Effect::new(&f0, false, [&z1[..]]).unwrap(), ]; - let err = Design::new(effects).unwrap_err(); - assert!(matches!( - err, - BuildError::SlopesNotYetSupported { effect: 1 } - )); + let design = Design::new(effects).unwrap(); + + // term 0: [intercept, z0, z1] over 2 levels; term 1: intercept over 3; + // term 2: slope-only over 2. + assert_eq!(design.terms[0].offset, 0); + assert_eq!(design.terms[0].n_dofs(), 6); + assert_eq!(design.terms[1].offset, 6); + assert_eq!(design.terms[1].n_dofs(), 3); + assert_eq!(design.terms[2].offset, 9); + assert!(!design.terms[2].intercept); + assert_eq!(design.terms[2].n_dofs(), 2); + assert_eq!(design.n_dofs, 11); + + // slope indices resolve to the effects' loading columns in the frame. + assert_eq!(design.terms[0].slopes, vec![0, 1]); + assert_eq!(design.terms[2].slopes, vec![2]); + assert_eq!(design.frame.loading_column(0), &z0[..]); + assert_eq!(design.frame.loading_column(2), &z1[..]); } } diff --git a/crates/within/src/domain/cross_tab.rs b/crates/within/src/domain/cross_tab.rs index c4bf531c..eefe8a22 100644 --- a/crates/within/src/domain/cross_tab.rs +++ b/crates/within/src/domain/cross_tab.rs @@ -6,7 +6,7 @@ //! Levels are stored compactly with a `local_to_global` map for active levels only. use crate::csr_block::CsrBlock; -use crate::domain::{Design, FactorMeta}; +use crate::domain::{Design, TermMeta}; mod accumulate; use accumulate::accumulate_cross_block; @@ -41,7 +41,7 @@ struct ActiveLevels { /// Returns `active[f][level]` = true if any observation uses that level of factor f. pub(crate) fn find_all_active_levels(design: &Design<'_>) -> Vec> { let mut active: Vec> = design - .factors + .terms .iter() .map(|f| vec![false; f.n_levels]) .collect(); @@ -78,8 +78,8 @@ fn compact_map(active: &[bool]) -> (Vec, usize) { fn build_compact_mapping( active_q: &[bool], active_r: &[bool], - fq: &FactorMeta, - fr: &FactorMeta, + fq: &TermMeta, + fr: &TermMeta, ) -> Option { let (q_map, n_q) = compact_map(active_q); let (r_map, n_r) = compact_map(active_r); @@ -194,8 +194,8 @@ impl CrossTab { r: usize, all_active: &[Vec], ) -> Option<(Self, BlockDiagonals, Vec)> { - let fq = &design.factors[q]; - let fr = &design.factors[r]; + let fq = &design.terms[q]; + let fr = &design.terms[r]; let active = build_compact_mapping(&all_active[q], &all_active[r], fq, fr)?; let (c, diag_q, diag_r) = accumulate_cross_block(design, weights, q, r, &active); diff --git a/crates/within/src/operator/design.rs b/crates/within/src/operator/design.rs index a1e337ba..34a2b9c1 100644 --- a/crates/within/src/operator/design.rs +++ b/crates/within/src/operator/design.rs @@ -1,206 +1,41 @@ use std::borrow::Cow; -use std::sync::atomic::Ordering; use portable_atomic::AtomicF64; -use rayon::prelude::*; use schwarz_precond::Operator; -use crate::domain::Design; +use crate::domain::{Design, TermMeta}; -// =========================================================================== -// Iteration kernels — module-private, shared between apply / apply_adjoint -// =========================================================================== +mod gather; +mod scatter; + +pub(crate) use gather::gather_apply; +use scatter::scatter_apply; /// Minimum number of rows before scatter/gather loops are parallelized. const PAR_THRESHOLD: usize = 10_000; -/// Factor-level threshold for choosing between fold and atomic scatter-add. -/// -/// Factors with fewer than this many levels use thread-local fold/reduce -/// (O(n_levels * n_threads) memory). Larger factors use atomic CAS instead, -/// which has low contention when bins vastly outnumber threads. -const SCATTER_LOCAL_THRESHOLD: usize = 100_000; - -/// Strategy for a single factor's scatter-add loop. -enum ScatterStrategy { - /// Plain sequential loop — used when n_rows is below `PAR_THRESHOLD`. - Sequential, - /// Parallel fold/reduce with thread-local accumulators — for small factors. - Fold, - /// Parallel atomic CAS — for large factors with low contention. - Atomic, - /// Atomic path for a large sorted factor: equal-level runs coalesce into - /// one atomic add per distinct level per chunk instead of one per row, - /// avoiding the atomic-CAS storm. - SortedCoalesced, -} - -impl ScatterStrategy { - /// Pick the scatter strategy for one factor; `sorted` is the factor's - /// level-column sortedness (`FactorMeta::sorted`). - fn pick(parallel: bool, n_levels: usize, sorted: bool) -> Self { - match (parallel, n_levels < SCATTER_LOCAL_THRESHOLD, sorted) { - (false, _, _) => ScatterStrategy::Sequential, - (true, true, _) => ScatterStrategy::Fold, - (true, false, true) => ScatterStrategy::SortedCoalesced, - (true, false, false) => ScatterStrategy::Atomic, - } - } -} - -/// Gather-apply: `dst[i] = Σ_q src[off_q + level(i, q)]`, times `scale[i]` if given. -/// -/// One sweep over `dst` per factor, plus a scale sweep only when given. -pub(crate) fn gather_apply( - design: &Design<'_>, - src: &[f64], - dst: &mut [f64], - scale: Option<&[f64]>, -) { - debug_assert!(scale.is_none_or(|s| s.len() == design.n_obs)); - debug_assert_eq!(src.len(), design.n_dofs); - debug_assert_eq!(dst.len(), design.n_obs); - - dst.fill(0.0); - - let factors = &design.factors; - let columns: Vec<&[u32]> = (0..factors.len()) - .map(|q| design.frame.level_column(q)) - .collect(); - - let kernel = |chunk: &mut [f64], row_start: usize| { - // `&levels` copies the slice ref out of `columns`; binding `&&[u32]` - // leaves a non-hoisted double deref in the inner loop (~5% measured). - for (f, &levels) in factors.iter().zip(columns.iter()) { - for (local, dst_val) in chunk.iter_mut().enumerate() { - let i = row_start + local; - *dst_val += src[f.offset + levels[i] as usize]; - } - } - - if let Some(scale) = scale { - for (s, dst_val) in scale[row_start..].iter().zip(chunk.iter_mut()) { - *dst_val *= s; - } - } - }; - - if design.n_obs > PAR_THRESHOLD { - const CHUNK_SIZE: usize = 4096; - - dst.par_chunks_mut(CHUNK_SIZE) - .enumerate() - .for_each(|(chunk_idx, chunk)| { - kernel(chunk, chunk_idx * CHUNK_SIZE); - }); - } else { - kernel(dst, 0); - } -} - -/// Sequential scatter-add: `slice[levels[i]] += value_fn(i)`. -fn scatter_sequential( - slice: &mut [f64], - levels: &[u32], - value_fn: &(impl Fn(usize) -> f64 + Sync), -) { - for (i, &l) in levels.iter().enumerate() { - slice[l as usize] += value_fn(i); - } +/// A [`TermMeta`] with its frame columns resolved into borrows for the kernels. +struct ResolvedTerm<'a> { + meta: &'a TermMeta, + levels: &'a [u32], + zs: Vec<&'a [f64]>, } -/// Parallel scatter-add via thread-local fold/reduce — best when `slice.len()` -/// (the factor's level count) is small relative to thread count. -fn scatter_fold(slice: &mut [f64], levels: &[u32], value_fn: &(impl Fn(usize) -> f64 + Sync)) { - let n_levels = slice.len(); - let min_len = (levels.len() / rayon::current_num_threads().max(1)).max(1024); - let identity = || vec![0.0f64; n_levels]; - let fold = |mut acc: Vec, (i, &l): (usize, &u32)| { - acc[l as usize] += value_fn(i); - acc - }; - let reduction = |mut a: Vec, b: Vec| { - for (ai, bi) in a.iter_mut().zip(b.iter()) { - *ai += *bi; - } - a - }; - let result: Vec = levels - .par_iter() +fn resolve_terms<'a>(design: &'a Design<'_>) -> Vec> { + design + .terms + .iter() .enumerate() - .with_min_len(min_len) - .fold(identity, fold) - .reduce(identity, reduction); - for (d, r) in slice.iter_mut().zip(result.iter()) { - *d += *r; - } -} - -/// Seed the operator's atomic scratch with `slice`'s current contents, -/// returning the trimmed view the scatter accumulates into. -fn seed_scatter_scratch<'b>(atomic_buf: &'b [AtomicF64], slice: &[f64]) -> &'b [AtomicF64] { - debug_assert!(atomic_buf.len() >= slice.len()); - let buf = &atomic_buf[..slice.len()]; - for (a, &v) in buf.iter().zip(slice.iter()) { - a.store(v, Ordering::Relaxed); - } - buf -} - -/// Copy accumulated scratch values back into `slice`. -fn writeback_scatter_scratch(slice: &mut [f64], buf: &[AtomicF64]) { - for (d, a) in slice.iter_mut().zip(buf.iter()) { - *d = a.load(Ordering::Relaxed); - } -} - -/// Parallel scatter-add via atomic CAS — best when `slice.len()` is large -/// relative to thread count (low contention). `atomic_buf` is the operator's -/// reusable scratch (sized to the largest factor); we use its first -/// `slice.len()` slots, re-seeding them via `store` so no allocation occurs. -fn scatter_atomic( - slice: &mut [f64], - levels: &[u32], - value_fn: &(impl Fn(usize) -> f64 + Sync), - atomic_buf: &[AtomicF64], -) { - let buf = seed_scatter_scratch(atomic_buf, slice); - levels.par_iter().enumerate().for_each(|(i, &l)| { - buf[l as usize].fetch_add(value_fn(i), Ordering::Relaxed); - }); - writeback_scatter_scratch(slice, buf); -} - -fn scatter_sorted_coalesced( - slice: &mut [f64], - levels: &[u32], - value_fn: &(impl Fn(usize) -> f64 + Sync), - atomic_buf: &[AtomicF64], -) { - let buf = seed_scatter_scratch(atomic_buf, slice); - // Each row-chunk coalesces its equal-level runs locally and commits one - // atomic add per distinct level. A run split across a chunk boundary is - // committed by both chunks — additive, so still correct — keeping chunks - // independent without a carry/fixup pass. - const CHUNK: usize = 65_536; - levels.par_chunks(CHUNK).enumerate().for_each(|(c, chunk)| { - let start = c * CHUNK; - // Single flat pass, one level load per row: accumulate the current - // run's sum and commit it whenever the level changes. - let mut level = chunk[0] as usize; - let mut sum = value_fn(start); - for (i, &li) in (start + 1..).zip(&chunk[1..]) { - let li = li as usize; - if li != level { - buf[level].fetch_add(sum, Ordering::Relaxed); - level = li; - sum = 0.0; - } - sum += value_fn(i); - } - buf[level].fetch_add(sum, Ordering::Relaxed); - }); - writeback_scatter_scratch(slice, buf); + .map(|(q, t)| ResolvedTerm { + meta: t, + levels: design.frame.level_column(q), + zs: t + .slopes + .iter() + .map(|&c| design.frame.loading_column(c)) + .collect(), + }) + .collect() } // =========================================================================== @@ -220,9 +55,9 @@ fn scatter_sorted_coalesced( pub(crate) struct DesignOperator<'a> { design: &'a Design<'a>, sqrt_weights: Option>, - /// Reusable atomic-scatter scratch, sized once to the largest factor's - /// level count and reused across factors and `apply_adjoint` calls so it is - /// allocated once per operator rather than once per LSMR iteration. + /// Reusable atomic-scatter scratch, sized once to the largest term's + /// coefficient block and reused across terms and `apply_adjoint` calls so + /// it is allocated once per operator rather than once per LSMR iteration. /// `apply_adjoint` takes `&self`, but `AtomicF64`'s load/store/fetch_add are /// `&self` operations, so a plain `Vec` (already `Sync`) needs no /// lock: each call re-seeds the buffer via `store` instead of resizing it. @@ -254,11 +89,11 @@ impl<'a> DesignOperator<'a> { ); w.iter().map(|wi| wi.sqrt()).collect() }); - let max_levels = design.factors.iter().map(|f| f.n_levels).max().unwrap_or(0); + let max_block = design.terms.iter().map(|t| t.n_dofs()).max().unwrap_or(0); Self { design, sqrt_weights, - scatter_scratch: (0..max_levels).map(|_| AtomicF64::new(0.0)).collect(), + scatter_scratch: (0..max_block).map(|_| AtomicF64::new(0.0)).collect(), } } @@ -272,31 +107,6 @@ impl<'a> DesignOperator<'a> { Some(sw) => Cow::Owned(y.iter().zip(sw).map(|(&yi, &swi)| swi * yi).collect()), } } - - fn scatter_apply(&self, dst: &mut [f64], value_fn: F) - where - F: Fn(usize) -> f64 + Sync, - { - let design = self.design; - debug_assert_eq!(dst.len(), design.n_dofs); - let parallel = design.n_obs > PAR_THRESHOLD; - - for (q, f) in design.factors.iter().enumerate() { - let slice = &mut dst[f.offset..f.offset + f.n_levels]; - let levels = design.frame.level_column(q); - - match ScatterStrategy::pick(parallel, f.n_levels, f.sorted) { - ScatterStrategy::Sequential => scatter_sequential(slice, levels, &value_fn), - ScatterStrategy::Fold => scatter_fold(slice, levels, &value_fn), - ScatterStrategy::Atomic => { - scatter_atomic(slice, levels, &value_fn, &self.scatter_scratch) - } - ScatterStrategy::SortedCoalesced => { - scatter_sorted_coalesced(slice, levels, &value_fn, &self.scatter_scratch) - } - } - } - } } impl Operator for DesignOperator<'_> { @@ -322,41 +132,9 @@ impl Operator for DesignOperator<'_> { // operator's `apply_adjoint` calls are sequential (`solve_batch` builds // a distinct operator per RHS), so the shared buffer is never raced. match &self.sqrt_weights { - Some(sw) => self.scatter_apply(y, |i| sw[i] * x[i]), - None => self.scatter_apply(y, |i| x[i]), + Some(sw) => scatter_apply(self.design, &self.scatter_scratch, y, &|i| sw[i] * x[i]), + None => scatter_apply(self.design, &self.scatter_scratch, y, &|i| x[i]), } Ok(()) } } - -#[cfg(test)] -mod tests { - use super::*; - - /// Must match a naive per-row scatter-add, including a run straddling the - /// chunk boundary (committed by two chunks via additive atomics). Gated on - /// large sorted factors, so integration tests never reach it — exercised - /// directly here. - #[test] - fn coalesced_scatter_matches_naive() { - // > 65_536 rows so the level-0 run crosses the chunk boundary, exercising - // the two-chunk additive commit for a single level. - let n_rows = 100_000usize; - let col: Vec = (0..n_rows).map(|i| u32::from(i >= 70_000)).collect(); - let n_levels = 2usize; - let values: Vec = (0..n_rows).map(|i| (i % 7) as f64 - 3.0).collect(); - - let mut expected = vec![0.0f64; n_levels]; - for (i, &c) in col.iter().enumerate() { - expected[c as usize] += values[i]; - } - - let buf: Vec = (0..n_levels).map(|_| AtomicF64::new(0.0)).collect(); - let mut got = vec![0.0f64; n_levels]; - scatter_sorted_coalesced(&mut got, &col, &|i| values[i], &buf); - - for (g, e) in got.iter().zip(&expected) { - assert!((g - e).abs() < 1e-6, "got {g}, expected {e}"); - } - } -} diff --git a/crates/within/src/operator/design/gather.rs b/crates/within/src/operator/design/gather.rs new file mode 100644 index 00000000..b1344c60 --- /dev/null +++ b/crates/within/src/operator/design/gather.rs @@ -0,0 +1,95 @@ +//! Gather kernel: coefficient space → observation space (`D x`). + +use rayon::prelude::*; + +use super::{resolve_terms, ResolvedTerm, PAR_THRESHOLD}; +use crate::domain::Design; + +/// Gather-apply: `dst[i] = Σ_t Σ_c src[off_t + c·L_t + level(i,t)] · loading_c(i)`, +/// times `scale[i]` if given (loading is `1` for intercept columns). +pub(crate) fn gather_apply( + design: &Design<'_>, + src: &[f64], + dst: &mut [f64], + scale: Option<&[f64]>, +) { + debug_assert!(scale.is_none_or(|s| s.len() == design.n_obs)); + debug_assert_eq!(src.len(), design.n_dofs); + debug_assert_eq!(dst.len(), design.n_obs); + + dst.fill(0.0); + + let terms = resolve_terms(design); + for_each_chunk(dst, |chunk, row_start| { + for t in &terms { + apply_term(t, src, chunk, row_start); + } + if let Some(scale) = scale { + for (s, dst_val) in scale[row_start..].iter().zip(chunk.iter_mut()) { + *dst_val *= s; + } + } + }); +} + +/// Sweep `dst` in cache-sized chunks, in parallel above [`PAR_THRESHOLD`] rows; +/// the kernel receives each chunk and the row index it starts at. +fn for_each_chunk(dst: &mut [f64], kernel: impl Fn(&mut [f64], usize) + Sync) { + if dst.len() > PAR_THRESHOLD { + const CHUNK_SIZE: usize = 4096; + + dst.par_chunks_mut(CHUNK_SIZE) + .enumerate() + .for_each(|(chunk_idx, chunk)| kernel(chunk, chunk_idx * CHUNK_SIZE)); + } else { + kernel(dst, 0); + } +} + +/// One term's contribution to a chunk of rows: dispatch the term's shape to a +/// monomorphized sweep. +fn apply_term(t: &ResolvedTerm<'_>, src: &[f64], chunk: &mut [f64], row_start: usize) { + let offset = t.meta.offset; + let n_levels = t.meta.n_levels; + let levels = t.levels; + let col = |c: usize| &src[offset + c * n_levels..offset + (c + 1) * n_levels]; + match (t.meta.intercept, t.zs.as_slice()) { + (true, []) => gather_term(chunk, row_start, levels, [col(0)], |_| [1.0]), + (true, &[z0]) => gather_term(chunk, row_start, levels, [col(0), col(1)], |i| [1.0, z0[i]]), + (true, &[z0, z1]) => gather_term(chunk, row_start, levels, [col(0), col(1), col(2)], |i| { + [1.0, z0[i], z1[i]] + }), + (false, &[z0]) => gather_term(chunk, row_start, levels, [col(0)], |i| [z0[i]]), + (intercept, zs) => { + let zoff = usize::from(intercept); + for (local, dst_val) in chunk.iter_mut().enumerate() { + let i = row_start + local; + let lev = levels[i] as usize; + let mut acc = if intercept { src[offset + lev] } else { 0.0 }; + for (v, z) in zs.iter().enumerate() { + acc += src[offset + (zoff + v) * n_levels + lev] * z[i]; + } + *dst_val += acc; + } + } + } +} + +/// One term sweep with a compile-time column count: `chunk[local] += Σ_c +/// cols[c][level(i)] · weights(i)[c]`. +#[inline(always)] +fn gather_term( + chunk: &mut [f64], + row_start: usize, + levels: &[u32], + cols: [&[f64]; N], + weights: impl Fn(usize) -> [f64; N], +) { + for (local, dst_val) in chunk.iter_mut().enumerate() { + let i = row_start + local; + let lev = levels[i] as usize; + let row = cols.iter().zip(weights(i)).map(|(col, w)| col[lev] * w); + // Fold from -0.0, not 0.0: the true additive identity, folds away. + *dst_val += row.fold(-0.0, |acc, term| acc + term); + } +} diff --git a/crates/within/src/operator/design/scatter.rs b/crates/within/src/operator/design/scatter.rs new file mode 100644 index 00000000..47c05c36 --- /dev/null +++ b/crates/within/src/operator/design/scatter.rs @@ -0,0 +1,275 @@ +//! Scatter kernel: observation space → coefficient space (`Dᵀ x`), one +//! strategy per term picked by block size and level-column sortedness. + +use std::sync::atomic::Ordering; + +use portable_atomic::AtomicF64; +use rayon::prelude::*; + +use super::{resolve_terms, ResolvedTerm, PAR_THRESHOLD}; +use crate::domain::Design; + +/// Adjoint scatter over all terms; `base(i)` is the row value (`x[i]`, or +/// `sw[i]·x[i]` when weighted) that each column scales by its loading. +pub(super) fn scatter_apply( + design: &Design<'_>, + scratch: &[AtomicF64], + dst: &mut [f64], + base: &(impl Fn(usize) -> f64 + Sync), +) { + debug_assert_eq!(dst.len(), design.n_dofs); + let parallel = design.n_obs > PAR_THRESHOLD; + + for t in resolve_terms(design) { + let n_levels = t.meta.n_levels; + let block = &mut dst[t.meta.offset..t.meta.offset + t.meta.n_dofs()]; + match (t.meta.intercept, t.zs.as_slice()) { + (true, []) => scatter_term::<1>(block, &t, parallel, scratch, |i| [base(i)]), + (true, &[z0]) => scatter_term::<2>(block, &t, parallel, scratch, |i| { + let b = base(i); + [b, z0[i] * b] + }), + (true, &[z0, z1]) => scatter_term::<3>(block, &t, parallel, scratch, |i| { + let b = base(i); + [b, z0[i] * b, z1[i] * b] + }), + (intercept, zs) => { + let zoff = usize::from(intercept); + if intercept { + scatter_term::<1>(&mut block[..n_levels], &t, parallel, scratch, |i| [base(i)]); + } + for (v, &z) in zs.iter().enumerate() { + let start = (zoff + v) * n_levels; + scatter_term::<1>( + &mut block[start..start + n_levels], + &t, + parallel, + scratch, + move |i| [z[i] * base(i)], + ); + } + } + } + } +} + +/// Scatter one term's coefficient block: `block[c·L + level(i)] += values(i)[c]`. +fn scatter_term( + block: &mut [f64], + term: &ResolvedTerm<'_>, + parallel: bool, + scratch: &[AtomicF64], + values: impl Fn(usize) -> [f64; C] + Sync, +) { + debug_assert_eq!(block.len(), C * term.meta.n_levels); + match ScatterStrategy::pick(parallel, C * term.meta.n_levels, term.meta.sorted) { + ScatterStrategy::Sequential => { + scatter_sequential::(block, term.meta.n_levels, term.levels, &values) + } + ScatterStrategy::Fold => scatter_fold::(block, term.meta.n_levels, term.levels, &values), + ScatterStrategy::Atomic => { + scatter_atomic::(block, term.meta.n_levels, term.levels, &values, scratch) + } + ScatterStrategy::SortedCoalesced => { + scatter_sorted_coalesced::(block, term.meta.n_levels, term.levels, &values, scratch) + } + } +} + +/// Coefficient-block threshold for choosing between fold and atomic scatter-add. +/// +/// Blocks (a term's `n_columns * n_levels` coefficients) smaller than this use +/// thread-local fold/reduce (O(block * n_threads) memory). Larger blocks use +/// atomic CAS instead, which has low contention when bins vastly outnumber +/// threads. +const SCATTER_LOCAL_THRESHOLD: usize = 100_000; + +/// Strategy for a single term's scatter-add loop. +enum ScatterStrategy { + /// Plain sequential loop — used when n_rows is below `PAR_THRESHOLD`. + Sequential, + /// Parallel fold/reduce with thread-local accumulators — for small blocks. + Fold, + /// Parallel atomic CAS — for large blocks with low contention. + Atomic, + /// Atomic path for a large sorted term: equal-level runs coalesce into + /// one atomic add per distinct level per chunk instead of one per row, + /// avoiding the atomic-CAS storm. + SortedCoalesced, +} + +impl ScatterStrategy { + /// Pick the scatter strategy for one term; `block` is the coefficient + /// count written by the kernel call, `sorted` the term's level-column + /// sortedness (`TermMeta::sorted`). + fn pick(parallel: bool, block: usize, sorted: bool) -> Self { + match (parallel, block < SCATTER_LOCAL_THRESHOLD, sorted) { + (false, _, _) => ScatterStrategy::Sequential, + (true, true, _) => ScatterStrategy::Fold, + (true, false, true) => ScatterStrategy::SortedCoalesced, + (true, false, false) => ScatterStrategy::Atomic, + } + } +} + +/// Sequential scatter-add: `block[c·L + levels[i]] += values(i)[c]`. +fn scatter_sequential( + block: &mut [f64], + n_levels: usize, + levels: &[u32], + values: &(impl Fn(usize) -> [f64; C] + Sync), +) { + for (i, &lev) in levels.iter().enumerate() { + let vals = values(i); + for (c, v) in vals.into_iter().enumerate() { + block[c * n_levels + lev as usize] += v; + } + } +} + +/// Parallel scatter-add via thread-local fold/reduce — best when the block +/// (the term's coefficient count) is small relative to thread count. +fn scatter_fold( + block: &mut [f64], + n_levels: usize, + levels: &[u32], + values: &(impl Fn(usize) -> [f64; C] + Sync), +) { + let min_len = (levels.len() / rayon::current_num_threads().max(1)).max(1024); + let identity = || vec![0.0f64; C * n_levels]; + let fold = |mut acc: Vec, (i, &lev): (usize, &u32)| { + let vals = values(i); + for (c, v) in vals.into_iter().enumerate() { + acc[c * n_levels + lev as usize] += v; + } + acc + }; + let reduction = |mut a: Vec, b: Vec| { + for (ai, bi) in a.iter_mut().zip(b.iter()) { + *ai += *bi; + } + a + }; + let result: Vec = levels + .par_iter() + .enumerate() + .with_min_len(min_len) + .fold(identity, fold) + .reduce(identity, reduction); + for (d, r) in block.iter_mut().zip(result.iter()) { + *d += *r; + } +} + +/// Seed the operator's atomic scratch with `block`'s current contents, +/// returning the trimmed view the scatter accumulates into. +fn seed_scatter_scratch<'b>(atomic_buf: &'b [AtomicF64], block: &[f64]) -> &'b [AtomicF64] { + debug_assert!(atomic_buf.len() >= block.len()); + let buf = &atomic_buf[..block.len()]; + for (a, &v) in buf.iter().zip(block.iter()) { + a.store(v, Ordering::Relaxed); + } + buf +} + +/// Copy accumulated scratch values back into `block`. +fn writeback_scatter_scratch(block: &mut [f64], buf: &[AtomicF64]) { + for (d, a) in block.iter_mut().zip(buf.iter()) { + *d = a.load(Ordering::Relaxed); + } +} + +/// Parallel scatter-add via atomic CAS — best when the block is large +/// relative to thread count (low contention). `atomic_buf` is the operator's +/// reusable scratch (sized to the largest term's block); we use its first +/// `block.len()` slots, re-seeding them via `store` so no allocation occurs. +fn scatter_atomic( + block: &mut [f64], + n_levels: usize, + levels: &[u32], + values: &(impl Fn(usize) -> [f64; C] + Sync), + atomic_buf: &[AtomicF64], +) { + let buf = seed_scatter_scratch(atomic_buf, block); + levels.par_iter().enumerate().for_each(|(i, &lev)| { + let vals = values(i); + for (c, v) in vals.into_iter().enumerate() { + buf[c * n_levels + lev as usize].fetch_add(v, Ordering::Relaxed); + } + }); + writeback_scatter_scratch(block, buf); +} + +fn scatter_sorted_coalesced( + block: &mut [f64], + n_levels: usize, + levels: &[u32], + values: &(impl Fn(usize) -> [f64; C] + Sync), + atomic_buf: &[AtomicF64], +) { + let buf = seed_scatter_scratch(atomic_buf, block); + // Each row-chunk coalesces its equal-level runs locally and commits one + // atomic add per distinct level per column. A run split across a chunk + // boundary is committed by both chunks — additive, so still correct — + // keeping chunks independent without a carry/fixup pass. + const CHUNK: usize = 65_536; + levels + .par_chunks(CHUNK) + .enumerate() + .for_each(|(c_idx, chunk)| { + let start = c_idx * CHUNK; + // Single flat pass, one level load per row: accumulate the current + // run's per-column sums and commit them whenever the level changes. + let mut level = chunk[0] as usize; + let mut sums = values(start); + for (i, &li) in (start + 1..).zip(&chunk[1..]) { + let li = li as usize; + if li != level { + for (c, s) in sums.into_iter().enumerate() { + buf[c * n_levels + level].fetch_add(s, Ordering::Relaxed); + } + level = li; + sums = [0.0; C]; + } + let vals = values(i); + for (c, v) in vals.into_iter().enumerate() { + sums[c] += v; + } + } + for (c, s) in sums.into_iter().enumerate() { + buf[c * n_levels + level].fetch_add(s, Ordering::Relaxed); + } + }); + writeback_scatter_scratch(block, buf); +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Must match a naive per-row scatter-add, including a run straddling the + /// chunk boundary (committed by two chunks via additive atomics) and the + /// multi-column fused commit. Gated on large sorted terms, so integration + /// tests never reach it — exercised directly here. + #[test] + fn coalesced_scatter_matches_naive() { + let n = 70_000usize; + let n_levels = 1_000usize; + let levels: Vec = (0..n).map(|i| (i * n_levels / n) as u32).collect(); + let x: Vec = (0..n).map(|i| (i % 13) as f64 - 6.0).collect(); + let z: Vec = (0..n).map(|i| ((i * 7) % 11) as f64 / 11.0 - 0.5).collect(); + + let buf: Vec = (0..2 * n_levels).map(|_| AtomicF64::new(0.0)).collect(); + let mut got = vec![0.0f64; 2 * n_levels]; + scatter_sorted_coalesced::<2>(&mut got, n_levels, &levels, &|i| [x[i], z[i] * x[i]], &buf); + + let mut expect = vec![0.0f64; 2 * n_levels]; + for (i, &l) in levels.iter().enumerate() { + expect[l as usize] += x[i]; + expect[n_levels + l as usize] += z[i] * x[i]; + } + for (g, e) in got.iter().zip(expect.iter()) { + assert!((g - e).abs() < 1e-9, "{g} vs {e}"); + } + } +} diff --git a/crates/within/src/operator/schwarz.rs b/crates/within/src/operator/schwarz.rs index 80aeb2f1..4b56ad93 100644 --- a/crates/within/src/operator/schwarz.rs +++ b/crates/within/src/operator/schwarz.rs @@ -200,7 +200,7 @@ fn build_diagonal( ) -> Result { let mut diag = vec![0.0; design.n_dofs]; - for (factor_idx, factor) in design.factors.iter().enumerate() { + for (factor_idx, factor) in design.terms.iter().enumerate() { let slice = &mut diag[factor.offset..factor.offset + factor.n_levels]; for (uid, &level) in design.frame.level_column(factor_idx).iter().enumerate() { slice[level as usize] += weights.map_or(1.0, |w| w[uid]); diff --git a/crates/within/src/operator/tests.rs b/crates/within/src/operator/tests.rs index a1fa2015..b972681f 100644 --- a/crates/within/src/operator/tests.rs +++ b/crates/within/src/operator/tests.rs @@ -310,6 +310,138 @@ mod design_tests { } } +mod slope_design_tests { + use crate::domain::{Design, Effect}; + use crate::operator::DesignOperator; + use schwarz_precond::Operator; + + /// Deterministic pseudo-random f64 in [-1, 1). + fn noise(seed: usize) -> f64 { + let mut z = seed as u64 ^ 0x9E37_79B9_7F4A_7C15; + z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9); + z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB); + (z >> 11) as f64 / (1u64 << 52) as f64 - 1.0 + } + + /// Dense design matrix from the design's internal (post-sort) columns — + /// the reference the operator must agree with regardless of the locality + /// permutation. + fn dense_matrix(design: &Design<'_>) -> Vec> { + let mut d = vec![vec![0.0; design.n_dofs]; design.n_obs]; + for (q, t) in design.terms.iter().enumerate() { + let levels = design.frame.level_column(q); + let mut c = 0; + if t.intercept { + for (i, &lev) in levels.iter().enumerate() { + d[i][t.offset + lev as usize] = 1.0; + } + c = 1; + } + for (v, &z_idx) in t.slopes.iter().enumerate() { + let z = design.frame.loading_column(z_idx); + for (i, &lev) in levels.iter().enumerate() { + d[i][t.offset + (c + v) * t.n_levels + lev as usize] = z[i]; + } + } + } + d + } + + fn assert_close(a: &[f64], b: &[f64]) { + for (x, y) in a.iter().zip(b) { + assert!( + (x - y).abs() <= 1e-10 * x.abs().max(y.abs()).max(1.0), + "{x} vs {y}" + ); + } + } + + /// Covers every kernel arm on the sequential path: plain, fused V=1/V=2, + /// slope-only, and the generic V=3 fallback — against the dense reference. + #[test] + fn slope_matvec_and_adjoint_match_dense_reference() { + let n = 6; + let f0 = [0u32, 1, 2, 0, 1, 2]; + let f1 = [0u32, 0, 1, 1, 2, 2]; + let f2 = [0u32, 1, 0, 1, 0, 1]; + let f3 = [0u32, 0, 0, 1, 1, 1]; + let f4 = [1u32, 0, 2, 1, 0, 2]; + let zs: Vec> = (0..6) + .map(|k| (0..n).map(|i| noise(k * 100 + i)).collect()) + .collect(); + let effects = vec![ + Effect::new(&f0, true, [&zs[0][..], &zs[1][..], &zs[2][..]]).unwrap(), + Effect::new(&f1, true, [&zs[3][..]]).unwrap(), + Effect::new(&f2, false, [&zs[4][..]]).unwrap(), + Effect::new(&f3, true, []).unwrap(), + Effect::new(&f4, true, [&zs[5][..], &zs[4][..]]).unwrap(), + ]; + let design = Design::new(effects).unwrap(); + let dense = dense_matrix(&design); + let op = DesignOperator::new(&design, None); + + let x: Vec = (0..design.n_dofs).map(|j| noise(7_000 + j)).collect(); + let mut got = vec![0.0; design.n_obs]; + op.apply(&x, &mut got).unwrap(); + let expect: Vec = dense + .iter() + .map(|row| row.iter().zip(&x).map(|(d, xj)| d * xj).sum()) + .collect(); + assert_close(&got, &expect); + + let r: Vec = (0..design.n_obs).map(|i| noise(9_000 + i)).collect(); + let mut got_t = vec![0.0; design.n_dofs]; + op.apply_adjoint(&r, &mut got_t).unwrap(); + let mut expect_t = vec![0.0; design.n_dofs]; + for (row, &ri) in dense.iter().zip(&r) { + for (e, d) in expect_t.iter_mut().zip(row) { + *e += d * ri; + } + } + assert_close(&got_t, &expect_t); + } + + /// Adjoint identity ⟨Dx, r⟩ = ⟨x, Dᵀr⟩ on a design large enough to take + /// the parallel strategies — sorted-coalesced (C=2), atomic (C=2), and + /// fold (C=3) — with weights in play. Gather and scatter share the layout + /// logic but not the kernels, so a per-strategy addressing bug breaks the + /// identity. + #[test] + fn slope_adjoint_property_parallel_strategies() { + let n = 150_000; + let l_big = 60_000usize; + let sorted: Vec = (0..n).map(|i| (i * l_big / n) as u32).collect(); + let unsorted: Vec = (0..n).map(|i| ((i * 7919) % l_big) as u32).collect(); + let small: Vec = (0..n).map(|i| (i % 10) as u32).collect(); + let z: Vec> = (0..4) + .map(|k| (0..n).map(|i| noise(k * n + i)).collect()) + .collect(); + let effects = vec![ + Effect::new(&sorted, true, [&z[0][..]]).unwrap(), + Effect::new(&unsorted, true, [&z[1][..]]).unwrap(), + Effect::new(&small, true, [&z[2][..], &z[3][..]]).unwrap(), + ]; + let design = Design::new(effects).unwrap(); + let weights: Vec = (0..n).map(|i| 0.5 + noise(i).abs()).collect(); + let op = DesignOperator::new(&design, Some(&weights)); + + let x: Vec = (0..design.n_dofs).map(|j| noise(13 * j + 1)).collect(); + let r: Vec = (0..n).map(|i| noise(29 * i + 5)).collect(); + + let mut dx = vec![0.0; n]; + op.apply(&x, &mut dx).unwrap(); + let mut dtr = vec![0.0; design.n_dofs]; + op.apply_adjoint(&r, &mut dtr).unwrap(); + + let lhs: f64 = dx.iter().zip(&r).map(|(a, b)| a * b).sum(); + let rhs: f64 = x.iter().zip(&dtr).map(|(a, b)| a * b).sum(); + assert!( + (lhs - rhs).abs() <= 1e-9 * lhs.abs().max(rhs.abs()).max(1.0), + "{lhs} vs {rhs}" + ); + } +} + // =========================================================================== // weighted adjoint property test // =========================================================================== diff --git a/crates/within/src/solver.rs b/crates/within/src/solver.rs index ba2cd633..00920ac0 100644 --- a/crates/within/src/solver.rs +++ b/crates/within/src/solver.rs @@ -224,6 +224,11 @@ impl<'a> Solver<'a> { preconditioner: impl Into, ) -> Result { let design = design.into_design()?; + // Slope-bearing designs build and their operator is exercisable, but + // the solve transform and preconditioner land in later slices (#59+). + if let Some(idx) = design.terms.iter().position(|t| !t.slopes.is_empty()) { + return Err(BuildError::SlopesNotYetSupported { effect: idx }); + } design.validate_weights(weights.as_deref())?; // Align weights with the design's internal (possibly locality-sorted) diff --git a/crates/within/tests/error_paths.rs b/crates/within/tests/error_paths.rs index 0b9323b0..4e6ff39e 100644 --- a/crates/within/tests/error_paths.rs +++ b/crates/within/tests/error_paths.rs @@ -272,3 +272,21 @@ fn test_build_error_display_preconditioner_dimension_mismatch() { assert!(s.contains("7")); assert!(s.contains("5")); } + +#[test] +fn test_solver_rejects_slope_bearing_design_naming_its_term() { + let levels = [0u32, 1, 0, 1]; + let slope = [1.0, 2.0, 3.0, 4.0]; + let effects = vec![ + within::Effect::new(&levels, true, []).expect("plain effect"), + within::Effect::new(&levels, true, [&slope[..]]).expect("slope effect"), + ]; + // The design builds — the operator contract is exercisable — but solving + // is deferred to the transform slices (#59+). + let design = Design::new(effects).expect("slope design builds"); + let err = Solver::new(design, None, None).unwrap_err(); + match err { + BuildError::SlopesNotYetSupported { effect: 1 } => {} + other => panic!("Expected SlopesNotYetSupported for term 1, got: {other:?}"), + } +}