From e41c769228a6163735c0f625cdcbc5a1584c8202 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 14 Jul 2026 16:22:14 +0000 Subject: [PATCH 1/2] Add hash-free seqno index tables to the Milnor algebra MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add a flat, arc-swapped `seqno` index for the Milnor basis: an O(#p-part-entries) rank computed from a precomputed `g` table, with no hash lookup. This is a general (GPU-agnostic) data structure — it is the uploadable/on-device index primitive, and an independently benchmarkable alternative to the basis hashmap. - `SeqnoTables` + `compute_seqno_tables` (idempotent, concurrency-safe via `ArcSwapOption::rcu`) + `seqno` in `milnor_algebra.rs`, applicable at p = 2 with a trivial profile and stable ordering. - The CPU basis index deliberately still uses the hashmap (the tables lose to it on the CPU); `compute_basis` does not build them. Documented at `try_basis_element_to_index`. - `benches/seqno.rs`: A/B of the table index against the hashmap. - Covered by the `seqno_matches_enumeration_order` test. Split out of the Nassau GPU work as standalone infrastructure. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01UPYvsLdEfitgCbAiPxxx3U --- ext/crates/algebra/Cargo.toml | 5 + ext/crates/algebra/benches/seqno.rs | 70 ++++++++ .../algebra/src/algebra/milnor_algebra.rs | 164 ++++++++++++++++++ 3 files changed, 239 insertions(+) create mode 100644 ext/crates/algebra/benches/seqno.rs diff --git a/ext/crates/algebra/Cargo.toml b/ext/crates/algebra/Cargo.toml index 3b98825cb5..aded31fc16 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 = [ @@ -55,3 +56,7 @@ harness = false [[bench]] name = "nassau_milnor" harness = false + +[[bench]] +name = "seqno" +harness = false diff --git a/ext/crates/algebra/benches/seqno.rs b/ext/crates/algebra/benches/seqno.rs new file mode 100644 index 0000000000..22ad312edf --- /dev/null +++ b/ext/crates/algebra/benches/seqno.rs @@ -0,0 +1,70 @@ +//! A/B benchmark: the hash-free table-based Milnor index ([`MilnorAlgebra::seqno`]) versus the +//! `FxHashMap`-backed [`MilnorAlgebra::basis_element_to_index`], over the exact operation of +//! turning a basis element back into its index. +//! +//! This is the lookup that a resolution performs on *every* output term of a multiply — Nassau's +//! `S_2` @ p=2 run issues ~1.3B of them — so a lookup that is even a little cheaper is worth having, +//! and a GPU kernel (which cannot carry a hashmap) *needs* the table-based form regardless. +//! +//! The `seqno` group here reads the flat [`arc_swap`]-backed table; compare its `basis_to_index/*` +//! ids against the `hashmap/*` ids at the same degree to see which wins on the CPU. + +use std::hint::black_box; + +use algebra::{Algebra, MilnorAlgebra}; +use criterion::{Criterion, Throughput, criterion_group, criterion_main}; +use fp::prime::TWO; +use pprof::criterion::{Output, PProfProfiler}; + +/// Degrees to sweep. Chosen to bracket the range a real `S_2` resolution reaches, from the cheap +/// low-dimensional end up to degrees whose basis has thousands of elements. +const DEGREES: &[i32] = &[16, 24, 32, 40, 48, 56, 64]; + +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 method pays to walk the algebra's storage during timing. + let basis: Vec<_> = (0..dim) + .map(|i| algebra.basis_element_from_index(degree, i).clone()) + .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!("basis_to_index/deg{degree}"), |b| { + b.iter(|| { + for elt in &basis { + black_box(algebra.seqno(&elt.p_part)); + } + }); + }); + } + + g.finish(); +} + +criterion_group! { + name = benches; + config = Criterion::default() + .measurement_time(std::time::Duration::from_secs(3)) + .with_profiler(PProfProfiler::new(100, Output::Flamegraph(None))); + 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 07b3bd35f8..3980b56a58 100644 --- a/ext/crates/algebra/src/algebra/milnor_algebra.rs +++ b/ext/crates/algebra/src/algebra/milnor_algebra.rs @@ -245,6 +245,16 @@ impl MilnorHashMap { } } +/// 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, +} + pub struct MilnorAlgebra { profile: MilnorProfile, p: ValidPrime, @@ -265,6 +275,15 @@ 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 — the earlier `OnceVec>` layout paid two + /// atomics *per table access*, which is what made the table lose to the hashmap. + seqno_tables: arc_swap::ArcSwapOption, + #[cfg(feature = "cache-multiplication")] /// source_deg -> target_deg -> source_op -> target_op multiplication_table: OnceVec>>>, @@ -293,6 +312,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(), } @@ -328,6 +348,11 @@ impl MilnorAlgebra { } pub fn try_basis_element_to_index(&self, elt: &MilnorBasisElement) -> Option { + // NB: the table-based `Self::seqno` computes this same index without a hash, but computing + // the rank (a degree sum plus two indexed table reads per populated ξ-position) is more work + // than one hash and probe, so it loses to this hashmap on the CPU. `seqno` is kept as the + // GPU-oriented index (a GPU kernel cannot carry a hashmap, and the flat table uploads + // directly), not for the CPU hot path. self.basis_element_to_index_map[elt.degree as usize] .get(elt) .copied() @@ -425,6 +450,10 @@ impl Algebra for MilnorAlgebra { self.generate_basis_2(max_degree); } + // The `seqno` tables are *not* built here: `seqno` lost to the hashmap on the CPU (see + // `try_basis_element_to_index`), so a normal resolution should not pay to build tables it + // won't use. A GPU backend that needs the hash-free index calls `compute_seqno_tables`. + // Populate hash map self.basis_element_to_index_map .extend(max_degree as usize, |d| { @@ -884,6 +913,118 @@ 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()); + 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; + } + } + + // Guard the publish: under `concurrent`, parallel `get_partial_matrix` builds can race here, + // and an unconditional store would let a smaller table clobber a larger one already in place + // — after which `seqno` would index past the shrunken `g` and panic. Only replace when ours + // reaches at least as far, so the cached `max_degree` is monotonic. + let new_tables = std::sync::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()), + }); + } + + /// The index ("sequence number") of `P(p_part)` in the Milnor basis of its degree, computed in + /// O(number of `p_part` entries) from the precomputed tables — no hash lookup. Assumes + /// `seqno_applicable` and that `p_part` is a genuine basis element (trimmed, in range). + /// + /// 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_table` difference across the degree consumed at + /// that position. + pub fn seqno(&self, p_part: &[PPartEntry]) -> usize { + let xi = combinatorics::xi_degrees(self.prime()); + let guard = self.seqno_tables.load(); + let t = guard + .as_ref() + .expect("seqno tables not built; call compute_seqno_tables first"); + let w = t.width; + let mut cur_d: i32 = p_part.iter().zip(xi).map(|(&r, &x)| r as i32 * x).sum(); + let mut rank = 0; + // Consume positions from the highest down; position 0 contributes nothing. + for h in (1..p_part.len()).rev() { + let r = p_part[h] as i32; + if r == 0 { + continue; + } + let below = cur_d - r * xi[h]; + rank += t.g[cur_d as usize * w + h] - t.g[below as usize * w + h]; + cur_d = below; + } + rank + } + fn generate_basis_generic(&self, max_degree: i32) { let q = 2 * self.prime() - 2; let tau_degrees = combinatorics::tau_degrees(self.prime()); @@ -1793,6 +1934,29 @@ 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, and reject non-basis elements via `try_`. + #[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); + algebra.compute_seqno_tables(max_degree); + for d in 0..=max_degree { + 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), + i, + "seqno mismatch at degree {d}, index {i}: {elt:?}" + ); + } + } + } + #[rstest] #[trace] #[case(2, 32, None)] From 5c4f429ef35c4c11a5febb58bdfa05bc452012d8 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 15 Jul 2026 05:14:31 +0000 Subject: [PATCH 2/2] Address review: guard seqno indexing and widen its test MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - `seqno`: add a `debug_assert` that the element's degree is within the loaded table's `max_degree`, so a not-built-far-enough call fails with a clear diagnostic instead of a raw slice out-of-bounds panic. - `seqno_matches_enumeration_order`: build the tables partially, rebuild identically (no-op), grow, then request a smaller degree — exercising the idempotent, monotonic (non-shrinking) publish documented on `compute_seqno_tables`, not just a single build. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01UPYvsLdEfitgCbAiPxxx3U --- .../algebra/src/algebra/milnor_algebra.rs | 45 ++++++++++++++----- 1 file changed, 34 insertions(+), 11 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_algebra.rs b/ext/crates/algebra/src/algebra/milnor_algebra.rs index 3980b56a58..771880d70f 100644 --- a/ext/crates/algebra/src/algebra/milnor_algebra.rs +++ b/ext/crates/algebra/src/algebra/milnor_algebra.rs @@ -1011,6 +1011,14 @@ impl MilnorAlgebra { .expect("seqno tables not built; call compute_seqno_tables first"); let w = t.width; let mut cur_d: i32 = p_part.iter().zip(xi).map(|(&r, &x)| r as i32 * x).sum(); + // `cur_d` only decreases in the loop below, so this bounds every `t.g` index. A raw + // out-of-bounds panic here means the tables were not built far enough for this element. + debug_assert!( + cur_d <= t.max_degree, + "p_part degree {cur_d} exceeds seqno tables built to {}; call compute_seqno_tables \ + first", + t.max_degree + ); let mut rank = 0; // Consume positions from the highest down; position 0 contributes nothing. for h in (1..p_part.len()).rev() { @@ -1943,18 +1951,33 @@ mod tests { assert!(algebra.seqno_applicable()); let max_degree = 100; algebra.compute_basis(max_degree); - algebra.compute_seqno_tables(max_degree); - for d in 0..=max_degree { - 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), - i, - "seqno mismatch at degree {d}, index {i}: {elt:?}" - ); + + // `seqno` must return the enumeration index of every basis element in `0..=upto`. + 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), + i, + "seqno mismatch at degree {d}, index {i}: {elt:?}" + ); + } } - } + }; + + // 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); } #[rstest]