Add hash-free seqno index tables to the Milnor algebra

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

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 &degree 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);

ext/crates/algebra/src/algebra/milnor_algebra.rs

@@ -245,6 +245,16 @@ impl<V> MilnorHashMap<V> {
     }
 }
 
+/// 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<usize>,
+}
+
 pub struct MilnorAlgebra {
     profile: MilnorProfile,
     p: ValidPrime,
@@ -265,6 +275,15 @@ pub struct MilnorAlgebra {
     /// degree -> MilnorBasisElement -> index
     basis_element_to_index_map: OnceVec<MilnorHashMap<usize>>,
 
+    /// 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<Vec<_>>` layout paid two
+    /// atomics *per table access*, which is what made the table lose to the hashmap.
+    seqno_tables: arc_swap::ArcSwapOption<SeqnoTables>,
+
     #[cfg(feature = "cache-multiplication")]
     /// source_deg -> target_deg -> source_op -> target_op
     multiplication_table: OnceVec<OnceVec<Vec<Vec<FpVector>>>>,
@@ -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<usize> {
+        // 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,126 @@ 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();
+        // `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() {
+            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 +1942,44 @@ 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);
+
+        // `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]
     #[trace]
     #[case(2, 32, None)]