diff --git a/ext/crates/algebra/Cargo.toml b/ext/crates/algebra/Cargo.toml index 4dbf1fc02a..701fb8d55b 100644 --- a/ext/crates/algebra/Cargo.toml +++ b/ext/crates/algebra/Cargo.toml @@ -17,6 +17,7 @@ maybe-rayon = { path = "../maybe-rayon" } once = { path = "../once" } anyhow = "1.0.98" +arc-swap = "1.7.1" auto_impl = "1.3.0" hashbrown = "0.15.4" itertools = { version = "0.14.0", default-features = false, features = [ @@ -59,3 +60,7 @@ harness = false [[bench]] name = "motivic" harness = false + +[[bench]] +name = "seqno" +harness = false diff --git a/ext/crates/algebra/EXPERIMENTS.md b/ext/crates/algebra/EXPERIMENTS.md new file mode 100644 index 0000000000..bb0d8f860b --- /dev/null +++ b/ext/crates/algebra/EXPERIMENTS.md @@ -0,0 +1,41 @@ +# algebra experiment log + +What we tried, what it gained, and why the rejected alternatives were rejected. This is the place +for that record — the code comments are not, and neither is the README, which describes the crate +as it is rather than how it got there. + +## Hash-free Milnor basis indexing ("seqno") + +`MilnorAlgebra::basis_element_to_index` is a `HashMap` lookup. The +table-based alternative, `MilnorAlgebra::seqno`, ranks a `p_part` by summing differences of a +precomputed `g` array, with no hash. + +Which one wins depends on the degree, and the crossover is a cache effect. The hashmap is +per-degree, so its working set grows with the dimension of that degree and eventually falls out of +cache. The `g` table is shared across degrees and grows only linearly in the degree, so it stays +resident and its cost is flat. The hashmap wins while it is cache-resident and loses once it is +not; `benches/seqno.rs` sweeps a range of degrees that spans the crossover. + +Because of that, `compute_basis` deliberately does not build the seqno tables: a resolution that +never reaches the crossover should not pay for them. Callers that want the hash-free index — a GPU +backend, or a high-degree CPU run — call `compute_seqno_tables` themselves. + +### Rejected: `OnceVec>` storage + +**Rejected.** The first version stored the tables in a `OnceVec>`. That paid two atomics +*per table access*, which was enough to make the table lose to the hashmap at every degree measured. +Storing one flat, row-major `Vec` behind an `arc_swap::ArcSwapOption` reduced a read to a single +guard load followed by direct indexing. + +### Rejected: re-deriving the degree inside `rank` + +**Rejected.** `rank` could recover the degree as `Σ rᵢ·ξᵢ` instead of taking it as an argument, but +every caller already knows it, and the hashmap it competes with reads it straight off the basis +element. Re-deriving it would have put a loop in the measurement that the competing path does not +pay. It survives as a `debug_assert!`. + +### Hoisting the `arc_swap` guard + +**Kept.** `seqno` acquires the guard on every call, which is one atomic per lookup and pure overhead +in a loop that ranks many elements. `seqno_ranker` hoists the acquisition out of the loop; the +`seqno` vs `seqno_naive` gap in `benches/seqno.rs` is what that is worth. diff --git a/ext/crates/algebra/benches/seqno.rs b/ext/crates/algebra/benches/seqno.rs new file mode 100644 index 0000000000..dbd536892e --- /dev/null +++ b/ext/crates/algebra/benches/seqno.rs @@ -0,0 +1,71 @@ +//! A/B benchmark: the hash-free table index against the basis hashmap. + +use algebra::{Algebra, MilnorAlgebra}; +use criterion::{Criterion, Throughput, black_box, criterion_group, criterion_main}; +use fp::prime::TWO; + +/// Degrees to sample, spanning the point where the table index overtakes the hashmap (see +/// `EXPERIMENTS.md`). +/// +/// `compute_basis` builds every degree below the maximum, so raising the top of this range costs +/// memory as well as time. +const DEGREES: &[i32] = &[32, 64, 128, 192, 256, 300, 340]; + +/// Time both indices over every basis element of each degree in [`DEGREES`]. +fn seqno(c: &mut Criterion) { + let algebra = MilnorAlgebra::new(TWO, false); + let max_degree = *DEGREES.iter().max().unwrap(); + algebra.compute_basis(max_degree); + algebra.compute_seqno_tables(max_degree); + + let mut g = c.benchmark_group("seqno"); + + for °ree in DEGREES { + let dim = algebra.dimension(degree); + if dim == 0 { + continue; + } + // Snapshot the basis so neither index pays to walk the algebra's storage while timed. + let basis: Vec<_> = (0..dim) + .map(|i| algebra.basis_element_from_index(degree, i)) + .collect(); + + g.throughput(Throughput::Elements(dim as u64)); + + g.bench_function(format!("hashmap/deg{degree}"), |b| { + b.iter(|| { + for elt in &basis { + black_box(algebra.basis_element_to_index(elt)); + } + }); + }); + + g.bench_function(format!("seqno/deg{degree}"), |b| { + let ranker = algebra.seqno_ranker(); + b.iter(|| { + for elt in &basis { + black_box(ranker.rank(elt.p_part, degree)); + } + }); + }); + + // `seqno` hoists the table guard out of the loop as a hot caller would; the gap against + // `seqno_naive`, which re-acquires it per call, is what that hoisting is worth. + g.bench_function(format!("seqno_naive/deg{degree}"), |b| { + b.iter(|| { + for elt in &basis { + black_box(algebra.seqno(elt.p_part, degree)); + } + }); + }); + } + + g.finish(); +} + +criterion_group! { + name = benches; + config = Criterion::default(); + targets = seqno +} +criterion_main!(benches); diff --git a/ext/crates/algebra/src/algebra/milnor_algebra.rs b/ext/crates/algebra/src/algebra/milnor_algebra.rs index 38d54960df..187f03fa5f 100644 --- a/ext/crates/algebra/src/algebra/milnor_algebra.rs +++ b/ext/crates/algebra/src/algebra/milnor_algebra.rs @@ -1,4 +1,4 @@ -use std::cell::Cell; +use std::{cell::Cell, sync::Arc}; use fp::{ prime::{Binomial, Prime, ValidPrime, factor_pk, iter::BitflagIterator}, @@ -411,6 +411,71 @@ impl std::fmt::Display for MilnorBasisElement { } } +/// Flat, contiguous storage for the "seqno" (hash-free index) computation. See +/// [`MilnorAlgebra::compute_seqno_tables`] for how `g` is derived and [`MilnorAlgebra::seqno`] for +/// how it is read. Row-major with a fixed `width` (the number of ξ-degrees), so entry `(e, h)` lives +/// at `g[e * width + h]`; degrees `0..=max_degree` are populated. +struct SeqnoTables { + max_degree: i32, + width: usize, + g: Vec, +} + +/// A borrowed view of the seqno tables, acquired once for a batch of lookups. +/// +/// See [`MilnorAlgebra::seqno_ranker`]. Holding this pins the revision of the tables that was +/// current at acquisition, so the per-lookup cost is the rank itself with no atomic and no table +/// re-acquisition. It is therefore valid only for the degrees that revision covered: a ranker held +/// across a concurrent [`MilnorAlgebra::compute_seqno_tables`] that grew the tables will not see +/// the new degrees, and ranking one panics. +pub struct SeqnoRanker { + tables: Arc, + xi: &'static [i32], +} + +impl SeqnoRanker { + /// The index of `P(p_part)`, which must have degree `degree`, in the Milnor basis of that + /// degree. + /// + /// The basis is enumerated by increasing highest ξ-index, so the rank of `P` accumulates, for + /// each populated position `h`, the number of basis elements whose highest index is `< h` + /// together with `h` — which is exactly the `g` difference across the degree consumed at that + /// position. + #[inline] + pub fn rank(&self, p_part: PPart, degree: i32) -> usize { + let t = &*self.tables; + let w = t.width; + debug_assert_eq!( + degree, + p_part + .iter() + .zip(self.xi) + .map(|(r, &x)| r as i32 * x) + .sum::(), + "degree {degree} does not match the p-part {p_part:?}" + ); + // `cur_d` only decreases below, so this one check bounds every `t.g` index. + debug_assert!( + degree <= t.max_degree, + "degree {degree} exceeds seqno tables built to {}; call compute_seqno_tables first", + t.max_degree + ); + let mut cur_d = degree; + let mut rank = 0; + // Position 0 contributes nothing. + for h in (1..p_part.len()).rev() { + let r = p_part.get(h) as i32; + if r == 0 { + continue; + } + let below = cur_d - r * self.xi[h]; + rank += t.g[cur_d as usize * w + h] - t.g[below as usize * w + h]; + cur_d = below; + } + rank + } +} + pub struct MilnorAlgebra { profile: MilnorProfile, p: ValidPrime, @@ -435,6 +500,14 @@ pub struct MilnorAlgebra { /// degree -> MilnorBasisElement -> index basis_element_to_index_map: OnceVec>, + /// Table backing the "seqno" (hash-free index) computation, populated only when + /// [`Self::seqno_applicable`] holds (p = 2, trivial profile, stable). It holds the flat, + /// row-major `g` array described in [`Self::compute_seqno_tables`]; [`Self::seqno`] ranks a + /// `p_part` from it with plain array indexing and no hash lookup. Stored behind an + /// [`arc_swap::ArcSwapOption`] rather than a [`OnceVec`] so that reads on the hot path are a + /// single guard load followed by direct indexing. + seqno_tables: arc_swap::ArcSwapOption, + #[cfg(feature = "cache-multiplication")] /// source_deg -> target_deg -> source_op -> target_op multiplication_table: OnceVec>>>, @@ -463,6 +536,7 @@ impl MilnorAlgebra { basis_table: OnceVec::new(), excess_table: OnceVec::new(), basis_element_to_index_map: OnceVec::new(), + seqno_tables: arc_swap::ArcSwapOption::empty(), #[cfg(feature = "cache-multiplication")] multiplication_table: OnceVec::new(), } @@ -511,6 +585,10 @@ impl MilnorAlgebra { } pub fn try_basis_element_to_index(&self, elt: &MilnorBasisElement) -> Option { + // NB: [`Self::seqno`] computes this same index without a hash, and wins above the degree + // where this map stops being cache-resident (see `EXPERIMENTS.md`). This path stays on the + // hashmap regardless, because `compute_basis` does not build the seqno tables, so they are + // not guaranteed to exist here. self.basis_element_to_index_map[elt.degree as usize] .get(elt) .copied() @@ -1085,6 +1163,123 @@ impl MilnorAlgebra { }); } + /// Whether the fast table-based [`Self::seqno`] can be used instead of the hashmap. It requires + /// `p = 2` (single-generator Milnor basis), a trivial profile (so *every* `P(R)` of a degree is + /// a basis element, matching the partition counts), and the stable ordering (unstable sorts the + /// basis by excess, breaking the enumeration-order = index correspondence). + fn seqno_applicable(&self) -> bool { + !self.generic() && !self.unstable_enabled && self.profile.is_trivial() + } + + /// Build the flat `SeqnoTables` up to `max_degree`, so that [`Self::seqno`] can be used. + /// Requires `seqno_applicable`. Idempotent: if the stored tables already reach + /// `max_degree` this returns immediately; otherwise it rebuilds the whole (cheap, + /// `O(max_degree · width)`) table from scratch and atomically swaps it in, so readers always see + /// either the old complete table or the new one. + /// + /// The `n[e][m]` intermediate — the number of `P(R)` of degree `e` using only `ξ₁ … ξ_{m+1}` — + /// is built locally and discarded; only the `g` row-progression it feeds is stored, since that + /// is all [`Self::seqno`] reads. `g[e][h]` sums `n[·][h−1]` along the arithmetic progression of + /// step `ξ_{h+1}`, letting `seqno` rank a `p_part` without a hash lookup. + pub fn compute_seqno_tables(&self, max_degree: i32) { + assert!(self.seqno_applicable()); + // As in `compute_basis`: a negative degree would wrap `rows`, and one past the bound would + // build rows for elements the packing cannot hold. + assert!( + (0..=PPart::MAX_DEGREE).contains(&max_degree), + "seqno tables are only supported for degrees 0..={}, got {max_degree}", + PPart::MAX_DEGREE, + ); + if let Some(t) = &*self.seqno_tables.load() + && t.max_degree >= max_degree + { + return; + } + + let xi = combinatorics::xi_degrees(self.prime()); + let width = xi.len(); + let rows = max_degree as usize + 1; + + // n[e * width + m] = #{ P(R) of degree e using only ξ₁ … ξ_{m+1} } + // = n[e][m-1] + [ξ_{m+1} ≤ e] · n[e − ξ_{m+1}][m] + let mut n = vec![0usize; rows * width]; + for e in 0..=max_degree { + let base = e as usize * width; + for m in 0..width { + // m = 0: partitions into {1} — always exactly one, P(e), for e ≥ 0. + let without = if m == 0 { + (e == 0) as usize + } else { + n[base + m - 1] + }; + let with = if xi[m] <= e { + n[(e - xi[m]) as usize * width + m] + } else { + 0 + }; + n[base + m] = without + with; + } + } + + // g[e * width + h] = Σ_{j ≥ 0} n[e − j·ξ_{h+1}][h−1] (h ≥ 1; g[·][0] unused) + // = n[e][h−1] + [ξ_{h+1} ≤ e] · g[e − ξ_{h+1}][h] + let mut g = vec![0usize; rows * width]; + for e in 0..=max_degree { + let base = e as usize * width; + for h in 1..width { + let head = n[base + h - 1]; + let tail = if xi[h] <= e { + g[(e - xi[h]) as usize * width + h] + } else { + 0 + }; + g[base + h] = head + tail; + } + } + + // Parallel `get_partial_matrix` builds race here under `concurrent`, and `seqno` would + // panic on a table that shrank under it, so only ever replace with one that reaches as far. + let new_tables = Arc::new(SeqnoTables { + max_degree, + width, + g, + }); + self.seqno_tables.rcu(|current| match current.as_deref() { + Some(t) if t.max_degree >= max_degree => current.clone(), + _ => Some(new_tables.clone()), + }); + } + + /// A handle that borrows the seqno tables once for a batch of lookups. + /// + /// [`Self::seqno`] acquires the [`arc_swap`] guard on every call, which is one atomic per + /// lookup. A hot loop should hoist the acquisition with this instead, subject to the staleness + /// bound documented on [`SeqnoRanker`]. + /// + /// # Panics + /// + /// If the tables have not been built; call [`Self::compute_seqno_tables`] first. + pub fn seqno_ranker(&self) -> SeqnoRanker { + debug_assert!(self.seqno_applicable()); + SeqnoRanker { + tables: self + .seqno_tables + .load_full() + .expect("seqno tables not built; call compute_seqno_tables first"), + xi: combinatorics::xi_degrees(self.prime()), + } + } + + /// The index ("sequence number") of `P(p_part)`, of degree `degree`, in the Milnor basis of + /// that degree — computed in O(number of `p_part` entries) from the precomputed tables, with + /// no hash lookup. Assumes `seqno_applicable` and that `p_part` is a genuine basis element + /// (trimmed, in range) of `degree`. + /// + /// Ranking many elements? Use [`Self::seqno_ranker`] to acquire the tables once. + pub fn seqno(&self, p_part: PPart, degree: i32) -> usize { + self.seqno_ranker().rank(p_part, degree) + } + fn generate_basis_generic(&self, max_degree: i32) { let q = 2 * self.prime() - 2; let tau_degrees = combinatorics::tau_degrees(self.prime()); @@ -2003,6 +2198,68 @@ mod tests { use super::*; + /// The table-based [`MilnorAlgebra::seqno`] must return the position of every basis element in + /// its degree — i.e. agree with the enumeration order that defines the index — for the stable + /// `p = 2` full algebra. + #[test] + fn seqno_matches_enumeration_order() { + let algebra = MilnorAlgebra::new(ValidPrime::new(2), false); + assert!(algebra.seqno_applicable()); + let max_degree = 100; + algebra.compute_basis(max_degree); + + let check = |upto: i32| { + for d in 0..=upto { + let dim = algebra.dimension(d); + for i in 0..dim { + let elt = algebra.basis_element_from_index(d, i); + assert_eq!( + algebra.seqno(elt.p_part, d), + i, + "seqno mismatch at degree {d}, index {i}: {elt:?}" + ); + assert_eq!(algebra.basis_element_to_index(&elt), i); + } + } + }; + + // Exercise the idempotent, monotonic (non-shrinking) publish documented on + // `compute_seqno_tables`: build partially, rebuild identically (no-op), grow, then request + // a smaller degree (must not shrink the cached table). + algebra.compute_seqno_tables(50); + check(50); + algebra.compute_seqno_tables(50); + check(50); + algebra.compute_seqno_tables(max_degree); + check(max_degree); + algebra.compute_seqno_tables(50); + check(max_degree); + } + + /// The tables are per-algebra and built on demand, so ranking through an algebra that never + /// called `compute_seqno_tables` must say so rather than read uninitialised state. + #[test] + #[should_panic(expected = "seqno tables not built")] + fn seqno_without_tables_panics() { + let algebra = MilnorAlgebra::new(ValidPrime::new(2), false); + algebra.compute_basis(10); + algebra.seqno(algebra.basis_element_from_index(1, 0).p_part, 1); + } + + /// A ranker pins the revision current at acquisition, so it covers only the degrees that + /// revision reached even after the tables have grown. + #[test] + #[should_panic(expected = "exceeds seqno tables built to")] + fn seqno_ranker_is_stale_past_its_degree() { + let algebra = MilnorAlgebra::new(ValidPrime::new(2), false); + algebra.compute_basis(60); + algebra.compute_seqno_tables(20); + let ranker = algebra.seqno_ranker(); + algebra.compute_seqno_tables(60); + let elt = algebra.basis_element_from_index(60, 0); + ranker.rank(elt.p_part, 60); + } + #[rstest] #[trace] #[case(2, 32, None)]