diff --git a/ext/Cargo.toml b/ext/Cargo.toml index 8b3919aba8..017a0f859d 100644 --- a/ext/Cargo.toml +++ b/ext/Cargo.toml @@ -59,6 +59,8 @@ concurrent = [ odd-primes = ["fp/odd-primes", "algebra/odd-primes", "sseq/odd-primes"] logging = [] nassau = [] +# Enables the GPU Milnor-multiply backend and the `milnor_gpu_ab` A/B microbenchmark. +gpu = ["algebra/gpu"] [workspace] members = [ @@ -86,3 +88,9 @@ harness = false [[bench]] name = "load_resolution" harness = false + +# A/B: GPU batch Milnor-multiply vs CPU get_partial_matrix on a real Nassau launch. +# Requires the `gpu` feature and a live CUDA device (run under the `gpu` dev shell). +[[bench]] +name = "milnor_gpu_ab" +harness = false diff --git a/ext/benches/milnor_gpu_ab.rs b/ext/benches/milnor_gpu_ab.rs new file mode 100644 index 0000000000..1fb354a47f --- /dev/null +++ b/ext/benches/milnor_gpu_ab.rs @@ -0,0 +1,211 @@ +//! A/B microbenchmark: GPU batch Milnor-multiply vs the CPU `get_partial_matrix`. +//! +//! Measures whether offloading the Milnor multiply to the GPU beats the CPU per-term path +//! *including* host↔device transfer. +//! +//! It resolves `S_2` with Nassau's algorithm, finds the largest `get_partial_matrix` +//! launch (by total element-term work), then times, best-of-N: +//! - **CPU**: the real `get_partial_matrix(degree, inputs)` (parallel, the default +//! per-term path) — this is the baseline to beat. +//! - **GPU**: extract the `(R, s)` products of that launch via the public API, then +//! `multiply_batch_on_gpu` (marshalling + upload + kernel + readback). +//! +//! Caveats (all *favour* the GPU, so a GPU loss here is conclusive): the GPU timing +//! excludes identity (`Sq(∅)`) copies and F₂ result *placement* (bit offsets) — both +//! cheap and not the multiply work being offloaded. +//! +//! Run under the `gpu` dev shell (unsandboxed): +//! `nix develop .#gpu --command cargo bench --bench milnor_gpu_ab --features gpu -- ` + +#[cfg(not(feature = "gpu"))] +fn main() { + eprintln!("milnor_gpu_ab requires --features gpu (and a live CUDA device)."); +} + +#[cfg(feature = "gpu")] +fn main() { + use std::time::Instant; + + use algebra::{ + MilnorAlgebra, + milnor_gpu::{GpuProduct, multiply_batch_on_gpu}, + module::{Module, homomorphism::ModuleHomomorphism}, + }; + use ext::{ + chain_complex::ChainComplex, + utils::{construct_nassau, init_logging}, + }; + use sseq::coordinates::Bidegree; + + // With `--features logging` this installs the tracing subscriber so the + // resolution's per-bidegree spans stream to stderr (RUST_LOG=info); without + // it, it's a no-op NoSubscriber. + let _ = init_logging(); + + let args: Vec = std::env::args().collect(); + let stem: i32 = args.get(1).and_then(|x| x.parse().ok()).unwrap_or(80); + let filt: i32 = args.get(2).and_then(|x| x.parse().ok()).unwrap_or(42); + let n_runs: u32 = args.get(3).and_then(|x| x.parse().ok()).unwrap_or(3); + let max_t = stem + filt; + + eprintln!("Resolving S_2 (Nassau) through stem {stem}, filtration {filt} ..."); + let res = construct_nassau(("S_2", "milnor"), None).unwrap(); + let t0 = Instant::now(); + res.compute_through_stem(Bidegree::n_s(stem, filt)); + eprintln!(" resolved in {:.1}s", t0.elapsed().as_secs_f64()); + + // The algebra (shared) — build the seqno tables the GPU path indexes with. + let algebra: std::sync::Arc = res.differential(1).target().algebra(); + algebra.compute_seqno_tables(max_t); + + // Cheap traversal that only *counts* the multiply work of one full launch (all + // source basis elements as inputs) at bidegree (s, t) — used to find the biggest + // launch without allocating products for every bidegree. + let count_work = |s: i32, t: i32| -> (usize, usize) { + let diff = res.differential(s); + let source = diff.source(); + let target = diff.target(); + let shift = diff.degree_shift(); + let (mut n_products, mut n_terms) = (0usize, 0usize); + for input_index in 0..source.dimension(t) { + let ogp = source.index_to_op_gen(t, input_index); + if ogp.generator_degree < diff.min_degree() || ogp.operation_degree == 0 { + continue; + } + let out_on_gen = diff.output(ogp.generator_degree, ogp.generator_index); + let act_input_degree = ogp.generator_degree - shift; + for gd in target + .iter_gen_offsets::<2>([act_input_degree, act_input_degree + ogp.operation_degree]) + { + let (input_start, input_end) = (gd.start[0], gd.end[0]); + if input_start >= out_on_gen.len() { + break; + } + let s_slice = out_on_gen.as_slice().restrict(input_start, input_end); + if !s_slice.is_zero() { + n_products += 1; + n_terms += s_slice.iter_nonzero().count(); + } + } + } + (n_products, n_terms) + }; + + // Extract the (R, s) products of one full launch at bidegree (s, t), mirroring + // apply_to_basis_element → FreeModule::act via the public API. Identity operations + // (Sq(∅)) carry no multiply work and are skipped. + let extract = |s: i32, t: i32| -> Vec { + let diff = res.differential(s); + let source = diff.source(); + let target = diff.target(); + let shift = diff.degree_shift(); + let mut products = Vec::new(); + for input_index in 0..source.dimension(t) { + let ogp = source.index_to_op_gen(t, input_index); + if ogp.generator_degree < diff.min_degree() || ogp.operation_degree == 0 { + continue; + } + let out_on_gen = diff.output(ogp.generator_degree, ogp.generator_index); + let act_input_degree = ogp.generator_degree - shift; + for gd in target + .iter_gen_offsets::<2>([act_input_degree, act_input_degree + ogp.operation_degree]) + { + let (input_start, input_end) = (gd.start[0], gd.end[0]); + if input_start >= out_on_gen.len() { + break; + } + let s_slice = out_on_gen.as_slice().restrict(input_start, input_end); + if s_slice.is_zero() { + continue; + } + products.push(GpuProduct { + r_degree: ogp.operation_degree, + r_idx: ogp.operation_index, + s_degree: act_input_degree - gd.gen_deg, + term_indices: s_slice.iter_nonzero().map(|(i, _)| i).collect(), + row: input_index, + out_offset: gd.start[1], + }); + } + } + products + }; + + // Find the launch with the most element-term work across the resolved region. + let mut best: Option<(i32, i32, usize, usize)> = None; // (s, t, n_products, n_terms) + eprintln!("Scanning {filt} filtrations for the largest launch ..."); + for s in 1..=filt { + eprintln!(" scan s={s}/{filt} (best so far: {best:?})"); + for t in s..=(s + stem) { + let diff = res.differential(s); + if diff.source().dimension(t) == 0 || diff.target().dimension(t) == 0 { + continue; + } + let (n_products, n_terms) = count_work(s, t); + if n_products == 0 { + continue; + } + if best.is_none_or(|(_, _, _, bt)| n_terms > bt) { + best = Some((s, t, n_products, n_terms)); + } + } + } + + let Some((s, t, n_products, n_terms)) = best else { + eprintln!("No non-trivial launch found; resolve to a larger bidegree."); + return; + }; + + let source_dim = res.differential(s).source().dimension(t); + let target_dim = res.differential(s).target().dimension(t); + eprintln!( + "\nLargest launch: (s={s}, t={t}) rows={source_dim} cols={target_dim}\n {n_products} \ + products, {n_terms} element-terms" + ); + + let inputs: Vec = (0..source_dim).collect(); + let diff = res.differential(s); + + // Best-of-N wall time helper (borrows locals via &mut dyn). + let best_of = |n: u32, f: &mut dyn FnMut()| -> f64 { + let mut best = f64::INFINITY; + for _ in 0..n { + let t0 = Instant::now(); + f(); + best = best.min(t0.elapsed().as_secs_f64()); + } + best + }; + + // CPU baseline: the real get_partial_matrix (parallel, per-term path). + let cpu = best_of(n_runs, &mut || { + let m = diff.get_partial_matrix(t, &inputs); + std::hint::black_box(&m); + }); + + // GPU candidate, split into host extraction and the batch multiply itself. + let num_rows = source_dim.max(1); + let extract_time = best_of(n_runs, &mut || { + let p = extract(s, t); + std::hint::black_box(&p); + }); + let products = extract(s, t); + let batch_time = best_of(n_runs, &mut || { + let out = multiply_batch_on_gpu(&algebra, target_dim, num_rows, &products); + std::hint::black_box(&out); + }); + let gpu = extract_time + batch_time; + + eprintln!("\n=== A/B (best of {n_runs}) ==="); + eprintln!(" CPU get_partial_matrix : {:8.3} ms", cpu * 1e3); + eprintln!(" GPU extract (host) : {:8.3} ms", extract_time * 1e3); + eprintln!(" GPU batch multiply : {:8.3} ms", batch_time * 1e3); + eprintln!(" GPU total : {:8.3} ms", gpu * 1e3); + eprintln!(" speedup (CPU / GPU) : {:8.2}×", cpu / gpu); + if gpu < cpu { + eprintln!(" → GPU wins: the invasive wire-in is justified."); + } else { + eprintln!(" → CPU wins at this scale on this device."); + } + eprintln!("(set MILNOR_GPU_TIMING=1 for the batch's marshal/device split)"); +} diff --git a/ext/crates/algebra/Cargo.toml b/ext/crates/algebra/Cargo.toml index 3b98825cb5..09031e4c4b 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 = [ @@ -28,6 +29,18 @@ serde = { version = "1.0.219", features = ["derive"] } serde_json = "1.0.141" enum_dispatch = "0.3.13" +# Optional GPU backend for the Milnor multiply kernel (see src/algebra/milnor_gpu.rs). +# `cuda` targets the local NVIDIA card via NVRTC (needs the CUDA toolkit — wired into +# the ext dev shell in flake.nix). Kernels are runtime-agnostic, so `wgpu` (Vulkan) +# remains a drop-in portable fallback. +cubecl = { version = "0.10.0", optional = true, default-features = false, features = [ + "cuda", +] } +# For pinning all GPU work to one CUDA stream (`StreamId`), so a single memory pool is +# reclaimed by `memory_cleanup` — CubeCL's pools are per-stream, and rayon spreads launches +# across threads/streams, which otherwise accumulates buffers until the card OOMs. +cubecl-common = { version = "0.10.0", optional = true } + [dev-dependencies] criterion = { version = "0.5", features = ["html_reports"] } expect-test = "1.5.1" @@ -38,6 +51,7 @@ rstest = "0.25.0" default = ["odd-primes"] cache-multiplication = [] concurrent = ["fp/concurrent", "maybe-rayon/concurrent"] +gpu = ["dep:cubecl", "dep:cubecl-common"] odd-primes = ["fp/odd-primes"] [[bench]] @@ -55,3 +69,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..d07ba273e4 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| { @@ -521,6 +550,13 @@ impl Algebra for MilnorAlgebra { s_degree: i32, s: FpSlice, ) { + // Per-term reference sweep: run the `PPartMultiplier` multiply once for each term of `s`, + // reusing one `PPartAllocation`. At p = 2 the admissible-matrix algorithm + // (`Self::multiply_basis_element_by_element_2`) computes the same product by enumerating + // `Sq(R)`'s admissible matrices once and amortizing over the terms of `s`, but Nassau's + // `S_2` regime is too sparse (dominated by single-term elements) for that up-front + // enumeration to pay off on the CPU. It is retained as the reference model for the GPU + // kernel, not wired here. let p = self.prime(); let r = self.basis_element_from_index(r_degree, r_idx); PPartAllocation::with_local(|mut allocation| { @@ -884,6 +920,181 @@ 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 + } + + /// Snapshot of the flat seqno `g` table as `u32`, for GPU upload: returns + /// `(width, g)` with `g[e * width + h]` row-major (see [`Self::seqno`]). + /// + /// Requires [`Self::compute_seqno_tables`] to have been built to at least the + /// degrees that will be indexed; panics if any entry exceeds `u32` (the device + /// representation). + /// Whether the GPU multiply path ([`crate::algebra::milnor_gpu`]) applies: exactly + /// the `seqno_applicable` regime (`p = 2`, trivial profile, stable), since + /// the kernel indexes its output with the table-based `seqno`. Public so the + /// resolution can gate its GPU dispatch without reaching into private state. + #[cfg(feature = "gpu")] + pub fn gpu_multiply_applicable(&self) -> bool { + self.seqno_applicable() + } + + #[cfg(feature = "gpu")] + pub(crate) fn seqno_table_u32(&self) -> (usize, Vec) { + let guard = self.seqno_tables.load(); + let t = guard + .as_ref() + .expect("seqno tables not built; call compute_seqno_tables first"); + let g = + t.g.iter() + .map(|&x| u32::try_from(x).expect("seqno g table entry exceeds u32")) + .collect(); + (t.width, g) + } + + /// Enumerate every admissible matrix of `Sq(R)` (for non-empty `r_p_part`) for + /// GPU upload. Returns `(cs_len, mk_len, col_sums, masks)`: `col_sums` and + /// `masks` are row-major flattenings (`num_matrices × cs_len` and + /// `num_matrices × mk_len`) in enumeration order — exactly the per-matrix data + /// the multiply kernel's term test consumes (see + /// [`Self::multiply_basis_element_by_element_2`]). Every matrix of a fixed `R` + /// shares the same `cs_len`/`mk_len`, so the flattening is rectangular. + #[cfg(feature = "gpu")] + pub(crate) fn admissible_matrices( + &self, + r_p_part: &[PPartEntry], + ) -> (usize, usize, Vec, Vec) { + let mut matrix = AdmissibleMatrix::new(r_p_part); + let cs_len = matrix.col_sums.len(); + let mk_len = matrix.masks.len(); + let mut col_sums = Vec::new(); + let mut masks = Vec::new(); + loop { + col_sums.extend_from_slice(&matrix.col_sums); + masks.extend_from_slice(&matrix.masks); + if !matrix.next() { + break; + } + } + (cs_len, mk_len, col_sums, masks) + } + fn generate_basis_generic(&self, max_degree: i32) { let q = 2 * self.prime() - 2; let tau_degrees = combinatorics::tau_degrees(self.prime()); @@ -1059,6 +1270,151 @@ impl MilnorAlgebra { }); } + /// Compute `Sq(R) * s` for a fixed operation `Sq(R)` and a general element `s`, adding the + /// result to `result`. Only valid at `p = 2`. + /// + /// Algorithm due to Christian Nassau (ported from the previously disabled + /// `FreeModule::custom_milnor_act`). To compute `Sq(R) * (Sq(S₁) + Sq(S₂) + ⋯)` we build the + /// admissible matrices for `Sq(R)` once and, for each matrix, test every `Sq(Sₖ)` against it: + /// a matrix contributes iff each column sum is at most the corresponding entry of `Sₖ` and the + /// relevant bits are disjoint. This amortizes the (expensive) matrix enumeration over the whole + /// element, whereas [`Self::multiply_with_allocation`] re-runs it per term of `s`. + /// + /// **Not on the CPU hot path** — Nassau's `S_2` regime is too sparse for the up-front + /// enumeration to beat the per-term sweep in [`Self::multiply_basis_element_by_element`]. It is + /// kept as the reference model for the GPU kernel: enumerate `Sq(R)`'s matrices once per + /// operation and test every element term against them in parallel — a shape that batches well + /// on a GPU. Exercised by the `admissible_multiply_agrees_with_reference` test. + // The `working`-building loops below legitimately index `basis`, `col_sums`, and `masks` by the + // same `j`, so a range loop is clearer than zipping three slices. + #[allow(clippy::needless_range_loop)] + pub fn multiply_basis_element_by_element_2( + &self, + mut result: FpSliceMut, + coeff: u32, + r_degree: i32, + r_idx: usize, + s_degree: i32, + s: FpSlice, + ) { + debug_assert!( + !self.generic(), + "multiply_basis_element_by_element_2 is p = 2 only" + ); + // Coefficients live in F₂, so an even coefficient kills the whole product, and every + // non-zero term of `s` has coefficient 1. + if coeff.is_multiple_of(2) { + return; + } + + let r = self.basis_element_from_index(r_degree, r_idx); + // `Sq(∅) = 1`, so `Sq(R) * s = s`. (Also avoids an empty `AdmissibleMatrix`.) The output + // degree equals `s_degree`, so basis indices are unchanged. + if r.p_part.is_empty() { + for (i, _) in s.iter_nonzero() { + result.add_basis_element(i, 1); + } + return; + } + + // The admissible-matrix sweep enumerates *all* matrices of `Sq(R)` up front and amortizes + // that over the terms of `s`. With a single term there is nothing to amortize, and for a + // large operation the wasted enumeration makes it several times slower than the + // `PPartMultiplier` path (which is constrained by `S` and so enumerates far fewer matrices). + // So peek the first two terms in one pass: with fewer than two, fall back to the per-term + // path — byte-for-byte the generic multiply, so that case never regresses. + let mut nonzero = s.iter_nonzero(); + let (Some((i0, _)), second) = (nonzero.next(), nonzero.next()) else { + return; // s = 0 + }; + let Some((i1, _)) = second else { + PPartAllocation::with_local(|allocation| { + self.multiply_with_allocation( + result, + 1, + r, + self.basis_element_from_index(s_degree, i0), + i32::MAX, + allocation, + ) + }); + return; + }; + + // Two or more terms: use the admissible-matrix sweep. Cache the (already-peeked) input + // basis elements once; they are reused across every admissible matrix. + let mut terms: Vec<&MilnorBasisElement> = Vec::with_capacity(s.len()); + terms.push(self.basis_element_from_index(s_degree, i0)); + terms.push(self.basis_element_from_index(s_degree, i1)); + terms.extend(nonzero.map(|(i, _)| self.basis_element_from_index(s_degree, i))); + + let out_degree = r_degree + s_degree; + let mut matrix = AdmissibleMatrix::new(&r.p_part); + let mut working = MilnorBasisElement { + q_part: 0, + p_part: PPart::new(), + degree: out_degree, + }; + + loop { + 'outer: for term in &terms { + let basis = &term.p_part; + working.p_part.clear(); + + for j in 0..std::cmp::min(basis.len(), matrix.col_sums.len()) { + if matrix.col_sums[j] > basis[j] { + continue 'outer; + } + if (basis[j] - matrix.col_sums[j]) & matrix.masks[j] != 0 { + continue 'outer; + } + // We should add the diagonal sum, but that equals the mask, and there are no + // bit conflicts, so a bitwise-or is the same thing. + working + .p_part + .push((basis[j] - matrix.col_sums[j]) | matrix.masks[j]); + } + + if basis.len() < matrix.col_sums.len() { + for &col_sum in &matrix.col_sums[basis.len()..] { + if col_sum > 0 { + continue 'outer; + } + } + for &mask in &matrix.masks[basis.len()..] { + working.p_part.push(mask); + } + } else { + for j in matrix.col_sums.len()..std::cmp::min(basis.len(), matrix.masks.len()) { + if basis[j] & matrix.masks[j] != 0 { + continue 'outer; + } + working.p_part.push(basis[j] | matrix.masks[j]); + } + if basis.len() < matrix.masks.len() { + for &mask in &matrix.masks[basis.len()..] { + working.p_part.push(mask); + } + } else { + for &entry in &basis[matrix.masks.len()..] { + working.p_part.push(entry); + } + } + } + + while let Some(0) = working.p_part.last() { + working.p_part.pop(); + } + + let idx = self.basis_element_to_index(&working); + result.add_basis_element(idx, 1); + } + if !matrix.next() { + break; + } + } + } + pub fn multiply_with_allocation( &self, mut res: FpSliceMut, @@ -1069,6 +1425,11 @@ impl MilnorAlgebra { mut allocation: PPartAllocation, ) -> PPartAllocation { let target_deg = m1.degree + m2.degree; + // The unstable dimension only depends on `target_deg` and `excess`, both loop-invariant, so + // compute the truncation bound once instead of per output term. In Nassau's stable path + // (`excess = i32::MAX`) this is the full dimension and the check never fires, but keeping it + // hoisted is correct for the unstable callers too. + let dim = self.dimension_unstable(target_deg, excess); if self.generic() { let m1f = self.multiply_qpart(m1, m2.q_part); for (cc, basis) in m1f { @@ -1083,7 +1444,7 @@ impl MilnorAlgebra { while let Some(c) = multiplier.next() { let idx = self.basis_element_to_index(&multiplier.ans); - if idx < self.dimension_unstable(target_deg, excess) { + if idx < dim { res.add_basis_element(idx, c * cc * coef); } } @@ -1101,7 +1462,7 @@ impl MilnorAlgebra { while let Some(c) = multiplier.next() { let idx = self.basis_element_to_index(&multiplier.ans); - if idx < self.dimension_unstable(target_deg, excess) { + if idx < dim { res.add_basis_element(idx, c * coef); } } @@ -1146,6 +1507,114 @@ impl MilnorAlgebra { } } +/// The state for enumerating the admissible matrices of a fixed operation `Sq(R)` at `p = 2`, used +/// by [`MilnorAlgebra::multiply_basis_element_by_element_2`]. See that method (and the original +/// `FreeModule::custom_milnor_act`) for the algorithm. Rows are indexed by the entries of `R`; the +/// stored `matrix` is row-major with `cols` columns. +struct AdmissibleMatrix { + cols: usize, + rows: usize, + matrix: Vec, + totals: Vec, + col_sums: Vec, + masks: Vec, +} + +impl AdmissibleMatrix { + fn new(ps: &[PPartEntry]) -> Self { + debug_assert!( + !ps.is_empty(), + "AdmissibleMatrix::new requires a non-empty R; Sq(∅) = 1 is handled by the caller" + ); + let rows = ps.len(); + let cols = ps + .iter() + .map(|x| (PPartEntry::BITS - x.leading_zeros()) as usize) + .max() + .unwrap(); + let mut matrix = vec![0; rows * cols]; + for (i, &x) in ps.iter().enumerate() { + matrix[i * cols] = x; + } + + let mut masks = Vec::with_capacity(rows + cols - 1); + masks.extend_from_slice(ps); + masks.resize(rows + cols - 1, 0); + + Self { + rows, + cols, + totals: vec![0; rows], // only used by `next`; no need to initialize + col_sums: vec![0; cols - 1], + matrix, + masks, + } + } + + #[inline] + fn get(&self, row: usize, col: usize) -> PPartEntry { + self.matrix[row * self.cols + col] + } + + #[inline] + fn set(&mut self, row: usize, col: usize, val: PPartEntry) { + self.matrix[row * self.cols + col] = val; + } + + /// Advance to the next admissible matrix, returning `false` when the enumeration is exhausted. + fn next(&mut self) -> bool { + for row in 0..self.rows { + let mut p_to_the_j: PPartEntry = 1; + self.totals[row] = self.get(row, 0); + 'mid: for col in 1..self.cols { + p_to_the_j *= 2; + // Quick check before computing the bitsums. + if p_to_the_j <= self.totals[row] { + // Compute the bitsum along the anti-diagonal to the bottom-left. + let mut d = 0; + for c in (row + col + 1).saturating_sub(self.rows)..col { + d |= self.get(row + col - c, c); + } + // Magic: the next number greater than `self[row][col]` whose bitwise-and with + // `d` is 0. + let new_entry = ((self.get(row, col) | d) + 1) & !d; + let inc = new_entry - self.get(row, col); + let sub = inc * p_to_the_j; + if self.totals[row] < sub { + self.totals[row] += p_to_the_j * self.get(row, col); + continue 'mid; + } + self.set(row, 0, self.totals[row] - sub); + self.masks[row] = self.get(row, 0); + self.col_sums[col - 1] += inc; + for j in 1..col { + self.masks[row + j] &= !self.get(row, j); + self.col_sums[j - 1] -= self.get(row, j); + self.set(row, j, 0); + } + self.set(row, col, new_entry); + + for i in 0..row { + self.set(i, 0, self.totals[i]); + self.masks[i] = self.totals[i]; + for j in 1..self.cols { + if i + j > row { + self.masks[i + j] &= !self.get(i, j); + } + self.col_sums[j - 1] -= self.get(i, j); + self.set(i, j, 0); + } + } + self.masks[row + col] = d | new_entry; + return true; + } + self.totals[row] += p_to_the_j * self.get(row, col); + } + } + false + } +} + #[derive(Debug, Default)] struct Matrix2D { cols: usize, @@ -1793,6 +2262,123 @@ 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); + } + + /// The `p = 2` admissible-matrix multiply ([`MilnorAlgebra::multiply_basis_element_by_element_2`], + /// retained as the GPU reference model — no longer wired into the CPU path) must agree + /// bit-for-bit with the reference `PPartMultiplier` path (`multiply_basis_elements`), both for + /// single basis elements and for dense (multi-term) elements — the latter also exercising mod-2 + /// cancellation. + #[test] + fn admissible_multiply_agrees_with_reference() { + let p = ValidPrime::new(2); + let algebra = MilnorAlgebra::new(p, false); + let max_degree = 32; + algebra.compute_basis(max_degree); + + for r_degree in 0..=max_degree { + let r_dim = algebra.dimension(r_degree); + for s_degree in 0..=(max_degree - r_degree) { + let s_dim = algebra.dimension(s_degree); + let out_degree = r_degree + s_degree; + let out_dim = algebra.dimension(out_degree); + + for i in 0..r_dim { + let mut expected_dense = FpVector::new(p, out_dim); + + for j in 0..s_dim { + // Reference: R_i * S_j via the PPartMultiplier path. + let mut expected = FpVector::new(p, out_dim); + algebra.multiply_basis_elements( + expected.as_slice_mut(), + 1, + r_degree, + i, + s_degree, + j, + ); + expected_dense.add(&expected, 1); + + // Admissible model: element `s = e_j`. + let mut s = FpVector::new(p, s_dim); + s.set_entry(j, 1); + let mut got = FpVector::new(p, out_dim); + algebra.multiply_basis_element_by_element_2( + got.as_slice_mut(), + 1, + r_degree, + i, + s_degree, + s.as_slice(), + ); + assert_eq!( + expected, got, + "single-term mismatch: R(deg {r_degree}, idx {i}) * S(deg {s_degree}, \ + idx {j})", + ); + } + + // Dense element (all ones): multi-term handling and mod-2 cancellation. + if s_dim > 0 { + let mut s = FpVector::new(p, s_dim); + for j in 0..s_dim { + s.set_entry(j, 1); + } + let mut got = FpVector::new(p, out_dim); + algebra.multiply_basis_element_by_element_2( + got.as_slice_mut(), + 1, + r_degree, + i, + s_degree, + s.as_slice(), + ); + assert_eq!( + expected_dense, got, + "dense mismatch: R(deg {r_degree}, idx {i}) * (all of deg {s_degree})", + ); + } + } + } + } + } + #[rstest] #[trace] #[case(2, 32, None)] diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs new file mode 100644 index 0000000000..1b77d5f85a --- /dev/null +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -0,0 +1,1082 @@ +//! GPU offload for the Milnor multiply at `p = 2`, built on [CubeCL]. +//! +//! Runs the admissible-matrix multiply +//! ([`super::milnor_algebra::MilnorAlgebra::multiply_basis_element_by_element_2`]) and the +//! hash-free `seqno` index as CubeCL kernels, batched per `get_partial_matrix` launch. +//! +//! The batched kernel `multiply_batch_kernel` fuses all `(R, s)` products of one +//! `get_partial_matrix` into a single launch — one thread per `(product, matrix, term)` pair, +//! decoded on-device from a prefix-sum over per-product pair counts. Admissible-matrix data is +//! deduplicated by distinct `R`, so a launch uploads compact per-`R`/per-product tables rather +//! than a per-pair table (which at scale would be gigabytes of almost-entirely-redundant data). +//! Its building blocks are the F₂ XOR accumulation (`xor_f2`), the on-device `seqno` index +//! (`seqno_core`/`seqno_kernel`, porting [`MilnorAlgebra::seqno`] as integer arithmetic over the +//! flat `g` table), and the single-`R` product (`multiply_pair`). +//! +//! Gated behind the `gpu` feature. Running needs the CUDA toolkit on `CUDA_PATH` / +//! `LD_LIBRARY_PATH` (the `gpu` dev shell in `ext/flake.nix` sets both) and a live +//! device; `cargo check`/`build` need neither (cudarc dlopens at runtime). +//! +//! [CubeCL]: https://github.com/tracel-ai/cubecl + +use cubecl::{ + cuda::{CudaDevice, CudaRuntime}, + prelude::*, +}; +use cubecl_common::stream_id::StreamId; + +/// The single CUDA stream all GPU work is pinned to (via [`StreamId::executes`]). +/// +/// CubeCL's memory pools are per-stream, and the resolution issues launches from many +/// rayon worker threads (each its own stream). Left alone, each stream's pool retains its +/// freed per-launch buffers (chiefly the hundreds-of-MB `out_h`), and across ~16 streams +/// they accumulate until the 4 GB card OOMs — `memory_cleanup` only trims the *calling* +/// stream's pool. Pinning every launch to one stream gives one pool that each launch's +/// `memory_cleanup` fully reclaims. Value 0 is a valid stream id (the first thread's). +const GPU_STREAM: StreamId = StreamId { value: 0 }; + +// Only the `#[cfg(test)]` standalone `seqno_kernel` sizes its working array by this bound; the +// production kernels use `WORKING_CAP`. +#[cfg(test)] +use crate::algebra::combinatorics::MAX_XI_TAU; +use crate::algebra::{Algebra, MilnorAlgebra, combinatorics::xi_degrees}; + +/// Comptime capacity for the per-thread `working` p_part in the multiply kernel. +/// The assembled p_part has length `max(term_len, mk_len)` before trimming, where +/// `mk_len = rows + cols − 1 ≤ MAX_XI_TAU + ⌈log2⌉`; 32 covers every in-range case. +const WORKING_CAP: usize = 32; + +/// Narrow an admissible-matrix / p-part entry to the `u16` the GPU buffers use, failing loudly +/// instead of silently wrapping. Every entry is well within `u16` for the stem ranges this path +/// targets; a panic here means that assumption was pushed past its limit, which must not ship +/// truncated data to the device. +fn narrow_u16(v: u32) -> u16 { + u16::try_from(v).expect("admissible/term entry exceeds u16") +} + +use std::sync::atomic::{AtomicU64, Ordering}; + +/// Aggregate [`multiply_batch_on_gpu`] counters across all launches (call count, host +/// marshal µs, device µs, total pairs), for splitting a whole resolution's GPU overhead. +static BATCH_CALLS: AtomicU64 = AtomicU64::new(0); +static BATCH_MARSHAL_US: AtomicU64 = AtomicU64::new(0); +static BATCH_DEVICE_US: AtomicU64 = AtomicU64::new(0); +static BATCH_PAIRS: AtomicU64 = AtomicU64::new(0); + +/// Read and reset the aggregate batch counters: `(calls, marshal_us, device_us, pairs)`. +pub fn take_batch_stats() -> (u64, u64, u64, u64) { + ( + BATCH_CALLS.swap(0, Ordering::Relaxed), + BATCH_MARSHAL_US.swap(0, Ordering::Relaxed), + BATCH_DEVICE_US.swap(0, Ordering::Relaxed), + BATCH_PAIRS.swap(0, Ordering::Relaxed), + ) +} + +use std::{ + collections::HashMap, + sync::{LazyLock, Mutex}, +}; + +use cubecl::server::Handle; + +use crate::algebra::milnor_algebra::PPartEntry; + +/// Where one `R`'s admissible-matrix data lives inside the resident master buffers. +#[derive(Clone, Copy)] +struct RInfo { + cs_off: u32, + mk_off: u32, + cs_len: u32, + mk_len: u32, + num_mats: u32, +} + +/// Process-global resident store of admissible-matrix data, both host- and device-side. +/// +/// Admissible-matrix enumeration is a pure function of `R`'s p-part and the same +/// low-degree `R`s recur in essentially every bidegree, so the host master (`col_sums` / +/// `masks`, append-only, keyed by p-part in `index`) is enumerated once per distinct `R` +/// and never recomputed. The device copies (`cs_handle` / `mk_handle`) mirror the master +/// and are re-uploaded *only when it grows* — after the `R`s saturate (early in a +/// resolution) launches upload no admissible data at all, cutting the dominant transfer. +/// +/// Guarded by a `Mutex` so the device section serializes across rayon worker threads: each +/// launch runs to its blocking readback before releasing, giving a happens-before edge and +/// no concurrent access — which is what CubeCL's single-device-thread managed-memory model +/// (its `unsafe impl Sync`) requires for a handle created on one thread to be reused on +/// another. Safe as a global because in the GPU path's regime (`p = 2`, trivial profile) +/// `admissible_matrices` depends only on the p-part, not on the algebra instance. +#[derive(Default)] +struct Resident { + col_sums: Vec, + masks: Vec, + index: HashMap, RInfo>, + cs_handle: Option, + mk_handle: Option, + cs_uploaded: usize, + mk_uploaded: usize, +} + +impl Resident { + /// Global offsets/lengths of `R`'s admissible matrices in the master, enumerating and + /// appending them on first sight (the append order fixes the offsets forever). + fn ensure(&mut self, algebra: &MilnorAlgebra, p_part: &[PPartEntry]) -> RInfo { + if let Some(info) = self.index.get(p_part) { + return *info; + } + let (cs_len, mk_len, cs, mk) = algebra.admissible_matrices(p_part); + let info = RInfo { + cs_off: self.col_sums.len() as u32, + mk_off: self.masks.len() as u32, + cs_len: cs_len as u32, + mk_len: mk_len as u32, + num_mats: (mk.len() / mk_len) as u32, + }; + self.col_sums.extend(cs.iter().map(|&v| narrow_u16(v))); + self.masks.extend(mk.iter().map(|&v| narrow_u16(v))); + self.index.insert(p_part.to_vec(), info); + info + } +} + +static RESIDENT: LazyLock> = LazyLock::new(|| Mutex::new(Resident::default())); + +/// Elementwise F₂ addition of two bit-packed vectors: `out[i] = a[i] ^ b[i]`. +/// +/// One thread per `u32` limb. F₂ addition is XOR of the packed limbs, so this is +/// the output primitive the multiply kernels accumulate with. +#[cfg(test)] +#[cube(launch)] +fn xor_f2(a: &Array, b: &Array, out: &mut Array) { + if ABSOLUTE_POS < out.len() { + out[ABSOLUTE_POS] = a[ABSOLUTE_POS] ^ b[ABSOLUTE_POS]; + } +} + +/// Compute `a ^ b` limb-wise on the default CUDA device. +/// +/// Host-side driver for `xor_f2`: uploads both operands, launches one thread per +/// limb, and reads the result back. Panics if the operands differ in length. +#[cfg(test)] +pub fn xor_f2_on_gpu(a: &[u32], b: &[u32]) -> Vec { + assert_eq!(a.len(), b.len(), "operands must have equal limb counts"); + let n = a.len(); + let client = CudaRuntime::client(&CudaDevice::default()); + + let a_handle = client.create_from_slice(u32::as_bytes(a)); + let b_handle = client.create_from_slice(u32::as_bytes(b)); + let out_handle = client.empty(std::mem::size_of_val(a)); + + // One 1-D block of `THREADS` units, enough blocks to cover every limb. + const THREADS: u32 = 256; + let cubes = (n as u32).div_ceil(THREADS); + unsafe { + xor_f2::launch::( + &client, + CubeCount::Static(cubes, 1, 1), + CubeDim::new_1d(THREADS), + ArrayArg::from_raw_parts(a_handle, n), + ArrayArg::from_raw_parts(b_handle, n), + ArrayArg::from_raw_parts(out_handle.clone(), n), + ); + } + + let bytes = client.read_one(out_handle).unwrap(); + u32::from_bytes(&bytes).to_vec() +} + +/// Device port of [`MilnorAlgebra::seqno`]: the index of `P(working)` in the Milnor +/// basis of its degree, from the flat `g` table with no hashing. `working` holds the +/// (trimmed) p_part in its first `wlen` entries; `g` has row width `width`, entry +/// `(e, h)` at `g[e*width + h]`; `xi` are the ξ-degrees. +/// +/// Thread/array indices are `usize`; p_part and table *values* are `u32`. A degree +/// (`cur_d`) is a value computed from `u32`s but also indexes `g`, so it is cast to +/// `usize` at the index sites. Shared by `seqno_kernel` and +/// `multiply_single_r_kernel` so both index outputs identically. +#[cube] +fn seqno_core( + g: &Array, + xi: &Array, + working: &Array, + wlen: usize, + width: usize, +) -> u32 { + // cur_d = Σ working[h] · xi[h]. + let mut cur_d = 0u32; + for h in 0..wlen { + cur_d += working[h] * xi[h]; + } + + // Rank by consuming positions from high to low; position 0 contributes nothing. + let mut rank = 0u32; + for hh in 1..wlen { + let h = wlen - hh; // wlen-1 down to 1 + let r = working[h]; + if r != 0 { + let below = cur_d - r * xi[h]; + let cur_row = usize::cast_from(cur_d) * width + h; + let below_row = usize::cast_from(below) * width + h; + rank += g[cur_row] - g[below_row]; + cur_d = below; + } + } + rank +} + +/// One thread per padded p_part: `out[i] = seqno(p_parts[i])`. `p_parts` is +/// `n × width` row-major, each row a p_part zero-padded to `width` (padding entries +/// are zero and skipped, so `wlen == width` matches the CPU's trimmed loop). +#[cfg(test)] +#[cube(launch)] +fn seqno_kernel( + g: &Array, + xi: &Array, + p_parts: &Array, + out: &mut Array, + width: usize, +) { + let idx = ABSOLUTE_POS; + if idx >= out.len() { + terminate!(); + } + let base = idx * width; + + let mut working = Array::::new(MAX_XI_TAU); + for h in 0..width { + working[h] = p_parts[base + h]; + } + out[idx] = seqno_core(g, xi, &working, width, width); +} + +/// Run `seqno_kernel` over `n` padded p_parts and return their seqno indices. +/// +/// `g`/`xi` come from `MilnorAlgebra::seqno_table_u32` and +/// [`crate::algebra::combinatorics::xi_degrees`]; `p_parts` is `n × width` row-major, +/// each row a p_part zero-padded to `width`. +#[cfg(test)] +pub fn seqno_batch_on_gpu( + width: usize, + xi: &[u32], + g: &[u32], + p_parts: &[u32], + n: usize, +) -> Vec { + assert_eq!(xi.len(), width, "xi must have `width` entries"); + assert_eq!(p_parts.len(), n * width, "p_parts must be n × width"); + let client = CudaRuntime::client(&CudaDevice::default()); + + let g_h = client.create_from_slice(u32::as_bytes(g)); + let xi_h = client.create_from_slice(u32::as_bytes(xi)); + let pp_h = client.create_from_slice(u32::as_bytes(p_parts)); + let out_h = client.empty(n * size_of::()); + + const THREADS: u32 = 256; + let cubes = (n as u32).div_ceil(THREADS); + unsafe { + seqno_kernel::launch::( + &client, + CubeCount::Static(cubes, 1, 1), + CubeDim::new_1d(THREADS), + ArrayArg::from_raw_parts(g_h, g.len()), + ArrayArg::from_raw_parts(xi_h, xi.len()), + ArrayArg::from_raw_parts(pp_h, p_parts.len()), + ArrayArg::from_raw_parts(out_h.clone(), n), + width, + ); + } + + let bytes = client.read_one(out_h).unwrap(); + u32::from_bytes(&bytes).to_vec() +} + +/// Assemble one `(admissible matrix, term)` product and XOR its F₂ output bit into +/// `out` at `row_base + idx`. The whole per-term test + output assembly of +/// [`MilnorAlgebra::multiply_basis_element_by_element_2`] lives here; both the +/// single-`R` and batch kernels call it with per-pair offsets. +/// +/// The reference's three tail branches collapse into one uniform per-position rule: +/// for column `j`, with `b`, `cs`, `mk` the term / `col_sums` / `masks` entries (zero +/// outside their lengths) and `low = min(term_len, cs_len)` — +/// - `j < low`: reject if `cs > b` or `(b−cs) & mk`; else `working[j] = (b−cs) | mk`. +/// - `j ≥ low`: reject if `cs > 0` or `b & mk`; else `working[j] = b | mk`. +/// +/// (For `j ≥ low` at most one of `b`, `cs` is in range, so this reproduces every +/// branch.) `seqno_core` gives the output index; the F₂ bit is XORed atomically +/// (collisions cancel mod 2). No explicit trailing-zero trim is needed — `seqno_core` +/// skips zero entries and `working` beyond the assembled length is zero, so the full +/// `WORKING_CAP` length is equivalent to the CPU's trimmed p_part (`xi` is host-padded +/// to `WORKING_CAP` so the `cur_d` sum stays in bounds; the extra terms are `0 · xi`). +#[cube] +#[allow(clippy::too_many_arguments)] +fn multiply_pair( + col_sums: &Array, + masks: &Array, + term_pparts: &Array, + g: &Array, + xi: &Array, + out: &mut Array>, + cs_base: usize, + mk_base: usize, + b_base: usize, + term_len: usize, + cs_len: usize, + mk_len: usize, + row_base: usize, + out_offset: usize, + width: usize, +) { + let mut low = cs_len; + if term_len < cs_len { + low = term_len; + } + + let mut working = Array::::new(WORKING_CAP); + let mut rejected = false; + + for j in 0..WORKING_CAP { + let mut b = 0u32; + if j < term_len { + b = u32::cast_from(term_pparts[b_base + j]); + } + let mut cs = 0u32; + if j < cs_len { + cs = u32::cast_from(col_sums[cs_base + j]); + } + let mut mk = 0u32; + if j < mk_len { + mk = u32::cast_from(masks[mk_base + j]); + } + + let mut val = 0u32; + if j < low { + if cs > b { + rejected = true; + } else { + let diff = b - cs; + if (diff & mk) != 0u32 { + rejected = true; + } else { + val = diff | mk; + } + } + } else { + if cs > 0u32 { + rejected = true; + } + if (b & mk) != 0u32 { + rejected = true; + } + val = b | mk; + } + working[j] = val; + } + + if !rejected { + // `seqno` indexes the algebra basis of the output degree; `out_offset` shifts it + // to this product's target-generator block within the row (0 for a single-block + // output). Both are bit offsets, added before splitting into (limb, bit). + let idx = seqno_core(g, xi, &working, WORKING_CAP, width); + let global_bit = out_offset + usize::cast_from(idx); + let word = row_base + global_bit / 32; + let bit = u32::cast_from(global_bit % 32); + out[word].fetch_xor(1u32 << bit); + } +} + +/// Multiply `Sq(R) · s` for a single fixed operation `R` into one F₂ output vector. +/// One thread per `(matrix, term)` pair; delegates the assembly to `multiply_pair`. +#[cfg(test)] +#[cube(launch)] +#[allow(clippy::too_many_arguments)] +fn multiply_single_r_kernel( + col_sums: &Array, + masks: &Array, + term_pparts: &Array, + term_lens: &Array, + g: &Array, + xi: &Array, + out: &mut Array>, + num_terms: usize, + num_matrices: usize, + cs_len: usize, + mk_len: usize, + width: usize, +) { + let pair = ABSOLUTE_POS; + if pair >= num_matrices * num_terms { + terminate!(); + } + let m = pair / num_terms; + let t = pair % num_terms; + let term_len = usize::cast_from(term_lens[t]); + multiply_pair( + col_sums, + masks, + term_pparts, + g, + xi, + out, + m * cs_len, + m * mk_len, + t * width, + term_len, + cs_len, + mk_len, + 0, + 0, + width, + ); +} + +/// Batched multiply: one launch covering all `(R, s)` products of (e.g.) a +/// `get_partial_matrix` call. One thread per `(product, matrix, term)` pair. +/// +/// Rather than a per-pair table (7 arrays × total-pairs — up to gigabytes at scale, +/// almost all redundant), the pair a thread handles is *decoded* from compact data: +/// - `prod_pair_start` is a prefix-sum of each product's pair count (`num_matrices × +/// num_terms`), length `num_products + 1`. A binary search finds the product `p` +/// owning thread `k`, then `local = k − prod_pair_start[p]` splits into matrix +/// `m = local / num_terms` and term `t = local % num_terms`. +/// - Admissible-matrix data (`col_sums`/`masks`) is deduplicated by distinct `R`: +/// `prod_r_index[p]` indexes the per-`R` `r_*` tables, so an `R` shared across many +/// rows is stored (and uploaded) once. +/// +/// Output is `num_rows` F₂ vectors of `num_limbs` `u32` limbs, row `r` at +/// `out[r*num_limbs ..]`. +#[cube(launch)] +#[allow(clippy::too_many_arguments)] +fn multiply_batch_kernel( + col_sums: &Array, + masks: &Array, + term_pparts: &Array, + term_lens: &Array, + g: &Array, + xi: &Array, + out: &mut Array>, + r_cs_offset: &Array, + r_mk_offset: &Array, + r_cs_len: &Array, + r_mk_len: &Array, + prod_r_index: &Array, + prod_term_start: &Array, + prod_num_terms: &Array, + prod_row_base: &Array, + prod_out_offset: &Array, + prod_pair_start: &Array, + width: usize, +) { + let k = ABSOLUTE_POS; + let num_products = prod_pair_start.len() - 1; + if k >= usize::cast_from(prod_pair_start[num_products]) { + terminate!(); + } + + // Largest product `p` with `prod_pair_start[p] <= k` (every product owns ≥ 1 pair, + // so `prod_pair_start` is strictly increasing and `p` is unique). 32 iterations + // cover any realistic product count; once `hi = lo + 1` the update is idempotent. + let mut lo = 0usize; + let mut hi = num_products; + for _ in 0..32 { + if hi - lo > 1 { + let mid = (lo + hi) / 2; + if usize::cast_from(prod_pair_start[mid]) <= k { + lo = mid; + } else { + hi = mid; + } + } + } + let p = lo; + + let ri = usize::cast_from(prod_r_index[p]); + let nt = usize::cast_from(prod_num_terms[p]); + let local = k - usize::cast_from(prod_pair_start[p]); + let m = local / nt; + let t = local % nt; + + let cs_len = usize::cast_from(r_cs_len[ri]); + let mk_len = usize::cast_from(r_mk_len[ri]); + let term_slot = usize::cast_from(prod_term_start[p]) + t; + multiply_pair( + col_sums, + masks, + term_pparts, + g, + xi, + out, + usize::cast_from(r_cs_offset[ri]) + m * cs_len, + usize::cast_from(r_mk_offset[ri]) + m * mk_len, + term_slot * width, + usize::cast_from(term_lens[term_slot]), + cs_len, + mk_len, + usize::cast_from(prod_row_base[p]), + usize::cast_from(prod_out_offset[p]), + width, + ); +} + +/// Compute `Sq(R) · s` on the GPU for a single operation `R = (r_degree, r_idx)`, +/// returning the F₂ result as bit-packed `u32` limbs (bit `i` = basis index `i`). +/// +/// `term_indices` are the nonzero indices of `s` in the degree-`s_degree` basis. +/// `R` must be non-empty (`Sq(∅) = 1` is the trivial identity the caller handles). +/// Requires the algebra's basis and seqno tables built through `r_degree + s_degree`. +#[cfg(test)] +pub fn multiply_single_r_on_gpu( + algebra: &MilnorAlgebra, + r_degree: i32, + r_idx: usize, + s_degree: i32, + term_indices: &[usize], +) -> Vec { + let (width, g) = algebra.seqno_table_u32(); + // Pad `xi` to `WORKING_CAP` so the kernel's `cur_d` sum (which runs to the full + // working capacity) never reads out of bounds; padding entries multiply zero. + let mut xi: Vec = xi_degrees(algebra.prime()) + .iter() + .map(|&x| x as u32) + .collect(); + xi.resize(WORKING_CAP, 0); + + let r = algebra.basis_element_from_index(r_degree, r_idx); + assert!( + !r.p_part.is_empty(), + "R must be non-empty (Sq(∅) = 1 is the identity)" + ); + let (cs_len, mk_len, cs32, mk32) = algebra.admissible_matrices(&r.p_part); + // Ship admissible-matrix / term data as u16 (see `multiply_batch_on_gpu`). + let mut col_sums: Vec = cs32.iter().map(|&v| narrow_u16(v)).collect(); + let masks: Vec = mk32.iter().map(|&v| narrow_u16(v)).collect(); + let num_matrices = masks.len() / mk_len; + + // Terms of s, each p_part padded to `width`, with their true (trimmed) lengths. + let num_terms = term_indices.len(); + let mut term_pparts = vec![0u16; num_terms * width]; + let mut term_lens = vec![0u32; num_terms]; + for (t, &ti) in term_indices.iter().enumerate() { + let elt = algebra.basis_element_from_index(s_degree, ti); + term_lens[t] = elt.p_part.len() as u32; + for (slot, &v) in term_pparts[t * width..(t + 1) * width] + .iter_mut() + .zip(&elt.p_part) + { + *slot = narrow_u16(v); + } + } + + let out_degree = r_degree + s_degree; + let dim = algebra.dimension(out_degree); + let num_limbs = dim.div_ceil(32).max(1); + + // Device buffers must be non-empty; `cs_len == 0` (R's max entry is 1) leaves + // `col_sums` empty. The kernel never reads past the real lengths. + if col_sums.is_empty() { + col_sums.push(0); + } + + let client = CudaRuntime::client(&CudaDevice::default()); + let cs_h = client.create_from_slice(u16::as_bytes(&col_sums)); + let mk_h = client.create_from_slice(u16::as_bytes(&masks)); + let tp_h = client.create_from_slice(u16::as_bytes(&term_pparts)); + let tl_h = client.create_from_slice(u32::as_bytes(&term_lens)); + let g_h = client.create_from_slice(u32::as_bytes(&g)); + let xi_h = client.create_from_slice(u32::as_bytes(&xi)); + let zeros = vec![0u32; num_limbs]; + let out_h = client.create_from_slice(u32::as_bytes(&zeros)); + + let total_pairs = num_matrices * num_terms; + const THREADS: u32 = 256; + let cubes = (total_pairs as u32).div_ceil(THREADS).max(1); + unsafe { + multiply_single_r_kernel::launch::( + &client, + CubeCount::Static(cubes, 1, 1), + CubeDim::new_1d(THREADS), + ArrayArg::from_raw_parts(cs_h, col_sums.len()), + ArrayArg::from_raw_parts(mk_h, masks.len()), + ArrayArg::from_raw_parts(tp_h, term_pparts.len()), + ArrayArg::from_raw_parts(tl_h, term_lens.len()), + ArrayArg::from_raw_parts(g_h, g.len()), + ArrayArg::from_raw_parts(xi_h, xi.len()), + ArrayArg::from_raw_parts(out_h.clone(), num_limbs), + num_terms, + num_matrices, + cs_len, + mk_len, + width, + ); + } + + let bytes = client.read_one(out_h).unwrap(); + u32::from_bytes(&bytes).to_vec() +} + +/// One `Sq(R) · s` product of a batched launch, written into output row `row` at bit +/// offset `out_offset`. +/// +/// `term_indices` are the nonzero indices of `s` in the degree-`s_degree` basis. +/// Multiple products may target the same `row` (their F₂ contributions XOR together), +/// mirroring how `get_partial_matrix` accumulates a row over generator blocks. The +/// product's `seqno` output indexes the algebra basis of the output degree; `out_offset` +/// is the start of the target-generator block that basis maps into within the row (0 when +/// the whole row is a single algebra element, as in the single-generator tests). +pub struct GpuProduct { + pub r_degree: i32, + pub r_idx: usize, + pub s_degree: i32, + pub term_indices: Vec, + pub row: usize, + pub out_offset: usize, +} + +/// Compute a whole batch of `Sq(R) · s` products in a single GPU launch — the +/// The batched unit of one `get_partial_matrix` call. `R`s may differ (each contributes its +/// own admissible matrices). Returns `num_rows` F₂ vectors, each `⌈num_cols/32⌉` +/// bit-packed `u32` limbs. +/// +/// `num_cols` is the *row* width — for a module row that is the module dimension (a sum +/// over generator blocks, generally larger than any single algebra degree's dimension), +/// with each product's `out_offset` selecting its block. Every product's +/// `out_offset + index` must be `< num_cols`. Every `R` must be non-empty; the algebra's +/// basis and seqno tables must reach each product's output degree (`r_degree + s_degree`). +pub fn multiply_batch_on_gpu( + algebra: &MilnorAlgebra, + num_cols: usize, + num_rows: usize, + products: &[GpuProduct], +) -> Vec> { + let (width, g) = algebra.seqno_table_u32(); + let mut xi: Vec = xi_degrees(algebra.prime()) + .iter() + .map(|&x| x as u32) + .collect(); + xi.resize(WORKING_CAP, 0); + + let num_limbs = num_cols.div_ceil(32).max(1); + + let t_marshal = std::time::Instant::now(); + + // The two heavy parts of marshalling — enumerating each distinct `R`'s admissible + // matrices, and looking up + padding every term's p-part — are independent per item, + // so they run in parallel (rayon via `concurrent`; serial otherwise). The cheap + // sequential glue (interning `R`s, concatenation, prefix sums) stays on one thread. + use maybe_rayon::prelude::*; + + // Intern distinct `R`s in first-seen order (cheap, sequential); record each product's + // `R` index. Admissible-matrix data is thus deduplicated: an `R` shared across many + // rows is enumerated and uploaded once. + let mut r_index: std::collections::HashMap<(i32, usize), u32> = + std::collections::HashMap::new(); + let mut distinct_r: Vec<(i32, usize)> = Vec::new(); + let mut prod_r_index: Vec = Vec::with_capacity(products.len()); + for prod in products { + let ri = *r_index + .entry((prod.r_degree, prod.r_idx)) + .or_insert_with(|| { + let i = distinct_r.len() as u32; + distinct_r.push((prod.r_degree, prod.r_idx)); + i + }); + prod_r_index.push(ri); + } + + // Admissible-matrix data (`col_sums`/`masks` + per-`R` offsets) is resident (built + // under the `RESIDENT` lock below), so nothing to enumerate or lay out here. + + // Parallel: each product's term p-parts (padded to `width`) and lengths. + let per_prod: Vec<(Vec, Vec)> = (0..products.len()) + .into_maybe_par_iter() + .map(|pi| { + let prod = &products[pi]; + let nt = prod.term_indices.len(); + let mut tp = vec![0u16; nt * width]; + let mut tl = Vec::with_capacity(nt); + for (k, &ti) in prod.term_indices.iter().enumerate() { + let elt = algebra.basis_element_from_index(prod.s_degree, ti); + tl.push(elt.p_part.len() as u32); + for (slot, &v) in tp[k * width..(k + 1) * width].iter_mut().zip(&elt.p_part) { + *slot = narrow_u16(v); + } + } + (tp, tl) + }) + .collect(); + + // Resident admissible-matrix store: enumerate each new `R` once and reuse forever; + // the per-`R` offsets are global (into the master `col_sums`/`masks`). Taking the lock + // here also serializes the device section across rayon workers (see [`Resident`]). + let mut resident = RESIDENT.lock().unwrap(); + let mut r_cs_offset: Vec = Vec::with_capacity(distinct_r.len()); + let mut r_mk_offset: Vec = Vec::with_capacity(distinct_r.len()); + let mut r_cs_len: Vec = Vec::with_capacity(distinct_r.len()); + let mut r_mk_len: Vec = Vec::with_capacity(distinct_r.len()); + let mut r_num_matrices: Vec = Vec::with_capacity(distinct_r.len()); + for &(rd, ridx) in &distinct_r { + let r = algebra.basis_element_from_index(rd, ridx); + assert!(!r.p_part.is_empty(), "each R must be non-empty"); + let info = resident.ensure(algebra, &r.p_part); + r_cs_offset.push(info.cs_off); + r_mk_offset.push(info.mk_off); + r_cs_len.push(info.cs_len); + r_mk_len.push(info.mk_len); + r_num_matrices.push(info.num_mats as usize); + } + + // Lay out per-product term data + records + the pair-count prefix sum (sequential). + let mut term_pparts: Vec = Vec::new(); + let mut term_lens: Vec = Vec::new(); + let mut prod_term_start: Vec = Vec::with_capacity(products.len()); + let mut prod_num_terms: Vec = Vec::with_capacity(products.len()); + let mut prod_row_base: Vec = Vec::with_capacity(products.len()); + let mut prod_out_offset: Vec = Vec::with_capacity(products.len()); + let mut prod_pair_start: Vec = Vec::with_capacity(products.len() + 1); + let mut pair_acc: usize = 0; + for (pi, (tp, tl)) in per_prod.iter().enumerate() { + let prod = &products[pi]; + let ri = prod_r_index[pi]; + prod_term_start.push(term_lens.len() as u32); + term_lens.extend_from_slice(tl); + term_pparts.extend_from_slice(tp); + prod_pair_start.push(pair_acc as u32); + pair_acc += r_num_matrices[ri as usize] * prod.term_indices.len(); + prod_num_terms.push(prod.term_indices.len() as u32); + prod_row_base.push((prod.row * num_limbs) as u32); + prod_out_offset.push(prod.out_offset as u32); + } + prod_pair_start.push(pair_acc as u32); // sentinel: total pair count at index num_products + + let total_pairs = pair_acc; + let out_len = num_rows * num_limbs; + if total_pairs == 0 { + return vec![vec![0u32; num_limbs]; num_rows]; + } + + // The resident `col_sums`/`masks` are non-empty once any `R` is present (guaranteed + // here, since `total_pairs > 0`); only `term_pparts` needs the non-empty guard. + if term_pparts.is_empty() { + term_pparts.push(0); + } + + let marshal_ms = t_marshal.elapsed().as_secs_f64() * 1e3; + + let t_device = std::time::Instant::now(); + + // Pin the whole device section to one CUDA stream (see [`GPU_STREAM`]) so a single + // memory pool is reclaimed by `memory_cleanup`. Held under the `resident` lock, so this + // stream is used by at most one thread at a time. + let result = GPU_STREAM.executes(|| { + let client = CudaRuntime::client(&CudaDevice::default()); + // Resident admissible buffers: (re-)upload the master only when it grew this + // launch; otherwise reuse the handle from a previous launch and upload nothing. + if resident.cs_handle.is_none() || resident.cs_uploaded != resident.col_sums.len() { + resident.cs_handle = Some(client.create_from_slice(u16::as_bytes(&resident.col_sums))); + resident.cs_uploaded = resident.col_sums.len(); + } + if resident.mk_handle.is_none() || resident.mk_uploaded != resident.masks.len() { + resident.mk_handle = Some(client.create_from_slice(u16::as_bytes(&resident.masks))); + resident.mk_uploaded = resident.masks.len(); + } + let cs_len_master = resident.col_sums.len(); + let mk_len_master = resident.masks.len(); + let cs_h = resident.cs_handle.clone().unwrap(); + let mk_h = resident.mk_handle.clone().unwrap(); + let tp_h = client.create_from_slice(u16::as_bytes(&term_pparts)); + let tl_h = client.create_from_slice(u32::as_bytes(&term_lens)); + let g_h = client.create_from_slice(u32::as_bytes(&g)); + let xi_h = client.create_from_slice(u32::as_bytes(&xi)); + let rco_h = client.create_from_slice(u32::as_bytes(&r_cs_offset)); + let rmo_h = client.create_from_slice(u32::as_bytes(&r_mk_offset)); + let rcl_h = client.create_from_slice(u32::as_bytes(&r_cs_len)); + let rml_h = client.create_from_slice(u32::as_bytes(&r_mk_len)); + let pri_h = client.create_from_slice(u32::as_bytes(&prod_r_index)); + let pts_h = client.create_from_slice(u32::as_bytes(&prod_term_start)); + let pnt_h = client.create_from_slice(u32::as_bytes(&prod_num_terms)); + let prb_h = client.create_from_slice(u32::as_bytes(&prod_row_base)); + let poo_h = client.create_from_slice(u32::as_bytes(&prod_out_offset)); + let pps_h = client.create_from_slice(u32::as_bytes(&prod_pair_start)); + let zeros = vec![0u32; out_len]; + let out_h = client.create_from_slice(u32::as_bytes(&zeros)); + const THREADS: u32 = 256; + let cubes = (total_pairs as u32).div_ceil(THREADS).max(1); + unsafe { + multiply_batch_kernel::launch::( + &client, + CubeCount::Static(cubes, 1, 1), + CubeDim::new_1d(THREADS), + ArrayArg::from_raw_parts(cs_h, cs_len_master), + ArrayArg::from_raw_parts(mk_h, mk_len_master), + ArrayArg::from_raw_parts(tp_h, term_pparts.len()), + ArrayArg::from_raw_parts(tl_h, term_lens.len()), + ArrayArg::from_raw_parts(g_h, g.len()), + ArrayArg::from_raw_parts(xi_h, xi.len()), + ArrayArg::from_raw_parts(out_h.clone(), out_len), + ArrayArg::from_raw_parts(rco_h, r_cs_offset.len()), + ArrayArg::from_raw_parts(rmo_h, r_mk_offset.len()), + ArrayArg::from_raw_parts(rcl_h, r_cs_len.len()), + ArrayArg::from_raw_parts(rml_h, r_mk_len.len()), + ArrayArg::from_raw_parts(pri_h, prod_r_index.len()), + ArrayArg::from_raw_parts(pts_h, prod_term_start.len()), + ArrayArg::from_raw_parts(pnt_h, prod_num_terms.len()), + ArrayArg::from_raw_parts(prb_h, prod_row_base.len()), + ArrayArg::from_raw_parts(poo_h, prod_out_offset.len()), + ArrayArg::from_raw_parts(pps_h, prod_pair_start.len()), + width, + ); + } + + let bytes = client.read_one(out_h).unwrap(); + let flat = u32::from_bytes(&bytes); + let result: Vec> = (0..num_rows) + .map(|r| flat[r * num_limbs..(r + 1) * num_limbs].to_vec()) + .collect(); + + // `out_h` alone is `num_rows × num_limbs` u32 — hundreds of MB at record degrees. + // It (and the small per-launch buffers, now dropped) varies in size launch to + // launch, so CubeCL's pool cannot reuse the slab and would accumulate them until + // the 4 GB card OOMs. Return the freed memory to the driver each launch; the + // resident admissible handles stay alive (refcount > 0) so cleanup skips them. + client.memory_cleanup(); + + result + }); + + // Aggregate marshal/device totals across every launch (cheap, always on) so a whole + // resolution's GPU overhead can be split host-vs-device via [`take_batch_stats`]. + let device_ms = t_device.elapsed().as_secs_f64() * 1e3; + BATCH_CALLS.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + BATCH_MARSHAL_US.fetch_add( + (marshal_ms * 1e3) as u64, + std::sync::atomic::Ordering::Relaxed, + ); + BATCH_DEVICE_US.fetch_add( + (device_ms * 1e3) as u64, + std::sync::atomic::Ordering::Relaxed, + ); + BATCH_PAIRS.fetch_add(total_pairs as u64, std::sync::atomic::Ordering::Relaxed); + + result +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Smoke test proving the CubeCL `cuda` runtime launches and returns correct + /// results. Requires a live GPU + the CUDA toolkit env (run under the `gpu` + /// dev shell, unsandboxed). + #[test] + fn xor_f2_matches_host() { + let a: Vec = (0..1000u32).map(|i| i.wrapping_mul(2654435761)).collect(); + let b: Vec = (0..1000u32).map(|i| i.wrapping_mul(40503)).collect(); + let expected: Vec = a.iter().zip(&b).map(|(x, y)| x ^ y).collect(); + assert_eq!(xor_f2_on_gpu(&a, &b), expected); + } + + /// The device `seqno` must reproduce the CPU basis order exactly: for every + /// basis element of every degree, `seqno(elt.p_part) == index`. Mirrors the CPU + /// `seqno_matches_enumeration_order` test on-device. Requires a live GPU + the + /// CUDA toolkit env (run under the `gpu` dev shell, unsandboxed). + #[test] + fn seqno_matches_index_on_gpu() { + use fp::prime::ValidPrime; + + let p = ValidPrime::new(2); + let algebra = MilnorAlgebra::new(p, false); + let max_degree = 60; + algebra.compute_basis(max_degree); + algebra.compute_seqno_tables(max_degree); + + let (width, g) = algebra.seqno_table_u32(); + assert_eq!(width, MAX_XI_TAU); + let xi: Vec = xi_degrees(p).iter().map(|&x| x as u32).collect(); + + // Marshal every basis element, padded to `width`; the expected seqno is the + // element's own index (the identity permutation the CPU proves). + let mut p_parts = Vec::new(); + let mut expected = Vec::new(); + 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); + let mut row = vec![0u32; width]; + for (slot, &v) in row.iter_mut().zip(&elt.p_part) { + *slot = v; + } + p_parts.extend_from_slice(&row); + expected.push(i as u32); + } + } + + let n = expected.len(); + let got = seqno_batch_on_gpu(width, &xi, &g, &p_parts, n); + assert_eq!(got, expected, "device seqno diverged from CPU basis order"); + } + + /// The single-`R` multiply kernel must match the CPU reference + /// `multiply_basis_element_by_element_2` bit-for-bit. For many `(R, s)` with `R` + /// non-empty and `s` the dense (all-ones) element — exercising the admissible + /// path and mod-2 cancellation — compare the GPU's packed F₂ output to the CPU's. + /// Requires a live GPU + the CUDA toolkit env (run under the `gpu` dev shell). + #[test] + fn multiply_single_r_matches_reference() { + use fp::{prime::ValidPrime, vector::FpVector}; + + let p = ValidPrime::new(2); + let algebra = MilnorAlgebra::new(p, false); + let max_degree = 40; + algebra.compute_basis(max_degree); + algebra.compute_seqno_tables(max_degree); + + let mut checked = 0usize; + for r_degree in 1..=12 { + let r_dim = algebra.dimension(r_degree); + for r_idx in 0..r_dim { + if algebra + .basis_element_from_index(r_degree, r_idx) + .p_part + .is_empty() + { + continue; // Sq(∅) = 1 is handled separately + } + for s_degree in 1..=(max_degree - r_degree) { + let s_dim = algebra.dimension(s_degree); + if s_dim == 0 { + continue; + } + let out_dim = algebra.dimension(r_degree + s_degree); + + // s = dense (all basis elements): multi-term, mod-2 cancellation. + let mut s = FpVector::new(p, s_dim); + for j in 0..s_dim { + s.set_entry(j, 1); + } + let mut cpu = FpVector::new(p, out_dim); + algebra.multiply_basis_element_by_element_2( + cpu.as_slice_mut(), + 1, + r_degree, + r_idx, + s_degree, + s.as_slice(), + ); + let num_limbs = out_dim.div_ceil(32).max(1); + let mut golden = vec![0u32; num_limbs]; + for (i, _) in cpu.iter_nonzero() { + golden[i / 32] ^= 1u32 << (i % 32); + } + + let term_indices: Vec = (0..s_dim).collect(); + let got = multiply_single_r_on_gpu( + &algebra, + r_degree, + r_idx, + s_degree, + &term_indices, + ); + assert_eq!( + got, golden, + "GPU multiply diverged from reference: R(deg {r_degree}, idx {r_idx}) * \ + dense s(deg {s_degree})", + ); + checked += 1; + } + } + } + assert!(checked > 0, "no (R, s) cases exercised"); + eprintln!("multiply_single_r: {checked} (R, s) cases matched reference"); + } + + /// The batched kernel must reproduce a whole output matrix: many heterogeneous + /// `(R, s)` products, several accumulating into the same row (XOR), computed in a + /// single launch, must equal the CPU reference matrix built product-by-product. + /// Requires a live GPU + the CUDA toolkit env (run under the `gpu` dev shell). + #[test] + fn multiply_batch_matches_reference() { + use fp::{prime::ValidPrime, vector::FpVector}; + + let p = ValidPrime::new(2); + let algebra = MilnorAlgebra::new(p, false); + let max_degree = 40; + algebra.compute_basis(max_degree); + algebra.compute_seqno_tables(max_degree); + + let out_degree = 24; + let out_dim = algebra.dimension(out_degree); + let num_rows = 8; + + // Products: every non-empty R of degree 1..out_degree, s dense at the + // complementary degree, assigned round-robin to rows so rows accumulate. + let mut products = Vec::new(); + for r_degree in 1..out_degree { + let s_degree = out_degree - r_degree; + let s_dim = algebra.dimension(s_degree); + if s_dim == 0 { + continue; + } + let r_dim = algebra.dimension(r_degree); + for r_idx in 0..r_dim { + if algebra + .basis_element_from_index(r_degree, r_idx) + .p_part + .is_empty() + { + continue; + } + let row = products.len() % num_rows; + products.push(GpuProduct { + r_degree, + r_idx, + s_degree, + term_indices: (0..s_dim).collect(), + row, + out_offset: 0, + }); + } + } + + // CPU golden matrix: accumulate each product into its row. + let mut cpu_rows: Vec = + (0..num_rows).map(|_| FpVector::new(p, out_dim)).collect(); + for prod in &products { + let s_dim = algebra.dimension(prod.s_degree); + let mut s = FpVector::new(p, s_dim); + for &ti in &prod.term_indices { + s.set_entry(ti, 1); + } + let mut tmp = FpVector::new(p, out_dim); + algebra.multiply_basis_element_by_element_2( + tmp.as_slice_mut(), + 1, + prod.r_degree, + prod.r_idx, + prod.s_degree, + s.as_slice(), + ); + cpu_rows[prod.row].add(&tmp, 1); + } + let num_limbs = out_dim.div_ceil(32).max(1); + let golden: Vec> = cpu_rows + .iter() + .map(|row| { + let mut packed = vec![0u32; num_limbs]; + for (i, _) in row.iter_nonzero() { + packed[i / 32] ^= 1u32 << (i % 32); + } + packed + }) + .collect(); + + let got = multiply_batch_on_gpu(&algebra, out_dim, num_rows, &products); + assert_eq!( + got, golden, + "batched GPU multiply diverged from reference matrix" + ); + eprintln!( + "multiply_batch: {} products across {num_rows} rows matched reference", + products.len() + ); + } +} diff --git a/ext/crates/algebra/src/algebra/mod.rs b/ext/crates/algebra/src/algebra/mod.rs index 67baed2b04..4513884e2a 100644 --- a/ext/crates/algebra/src/algebra/mod.rs +++ b/ext/crates/algebra/src/algebra/mod.rs @@ -18,6 +18,9 @@ pub use field::Field; pub mod milnor_algebra; pub use milnor_algebra::MilnorAlgebra; +#[cfg(feature = "gpu")] +pub mod milnor_gpu; + mod steenrod_algebra; pub use steenrod_algebra::{AlgebraType, SteenrodAlgebra}; diff --git a/ext/crates/algebra/src/module/free_module.rs b/ext/crates/algebra/src/module/free_module.rs index 45af27b2c8..9d878f6ce0 100644 --- a/ext/crates/algebra/src/module/free_module.rs +++ b/ext/crates/algebra/src/module/free_module.rs @@ -196,6 +196,11 @@ impl> Module for MuFreeModule { break; } let input_slice = input.restrict(input_start, input_end); + // A zero generator block contributes nothing; skip it rather than paying a multiply + // call that does no work (empirically ~a quarter of calls during a sphere resolution). + if input_slice.is_zero() { + continue; + } self.algebra.multiply_basis_element_by_element_unstable( result.slice_mut(output_start, output_end), coeff, diff --git a/ext/examples/nassau_e2e.rs b/ext/examples/nassau_e2e.rs new file mode 100644 index 0000000000..04220e2525 --- /dev/null +++ b/ext/examples/nassau_e2e.rs @@ -0,0 +1,46 @@ +//! End-to-end timing harness for Nassau's algorithm on the sphere `S_2` at p = 2. +//! +//! Resolves `S_2` via [`construct_nassau`] and [`ext::nassau::Resolution::compute_through_stem`] up to a given +//! `(stem, filtration)`, reporting the fastest wall time over several fresh runs. This exercises the +//! whole resolution — the Milnor multiplication kernel *and* the F₂ linear algebra — so it is the +//! right tool for judging whether a change to `MilnorAlgebra` multiplication actually moves the +//! needle in practice (unlike a micro-benchmark of the multiply alone). +//! +//! Usage: +//! ```text +//! cargo run --release --example nassau_e2e -- [stem=80] [filtration=42] [runs=3] +//! # for the parallel driver on a big case, enable rayon: +//! cargo run --release --features concurrent --example nassau_e2e -- 100 55 3 +//! ``` +//! +//! To A/B a multiplication change: build/run on the change, then `git stash`/checkout the baseline +//! and run again; compare the `best=` figures (min over runs is the most throttling-robust). + +use std::time::Instant; + +use ext::{chain_complex::ChainComplex, utils::construct_nassau}; +use sseq::coordinates::Bidegree; + +fn main() { + let mut args = std::env::args().skip(1); + let n: i32 = args.next().and_then(|s| s.parse().ok()).unwrap_or(80); + let s: i32 = args.next().and_then(|s| s.parse().ok()).unwrap_or(42); + let runs: usize = args.next().and_then(|s| s.parse().ok()).unwrap_or(3); + + let target = Bidegree::n_s(n, s); + let mut best = f64::INFINITY; + for run in 1..=runs { + // Each run builds a fresh resolution: the resolution caches results, so it cannot be reused. + let start = Instant::now(); + let res = construct_nassau(("S_2", "milnor"), None).unwrap(); + res.compute_through_stem(target); + let secs = start.elapsed().as_secs_f64(); + best = best.min(secs); + // Touch the result so nothing is optimized away (the compute has side effects regardless). + eprintln!( + " run {run}/{runs}: {secs:.2} s (hom degrees resolved: {})", + res.next_homological_degree() + ); + } + println!("S_2 @ p=2 stem={n} filtration={s} runs={runs} best={best:.2}s"); +} diff --git a/ext/flake.nix b/ext/flake.nix index 35e0d02fc8..b43a0e0d15 100644 --- a/ext/flake.nix +++ b/ext/flake.nix @@ -9,6 +9,13 @@ super.flake-utils.lib.eachDefaultSystem (system: let pkgs = import super.nixpkgs {inherit system;}; + # CUDA is unfree, so it needs its own nixpkgs instance. Only the `gpu` dev + # shell pulls it in — the default shell and `nix run .#test` stay CUDA-free. + cudaPkgs = import super.nixpkgs { + inherit system; + config.allowUnfree = true; + }; + pythonEnv = pkgs.python3.withPackages (ps: [ ps.black ps.pytest @@ -27,6 +34,15 @@ pkgs.perf ] ++ super.defaultPackages.devTools.${system}; + + # CUDA toolkit for the CubeCL `cuda` backend (algebra `gpu` feature). + # `cubecl-cuda` JIT-compiles kernels with NVRTC — it needs `CUDA_PATH` to + # point at a tree with `include/` (NVRTC `--include-path`) plus libnvrtc, and + # drives them via the CUDA driver API (`libcuda`, supplied by the host NVIDIA + # driver at /run/opengl-driver/lib, not nixpkgs). The monolithic `cudatoolkit` + # gives one prefix with both headers and libs. cudarc dlopens the libs at + # runtime (no build-time link), so only running — not building — needs this. + cudatoolkit = cudaPkgs.cudaPackages.cudatoolkit; in { devShells.default = pkgs.mkShell { packages = commonPackages; @@ -35,6 +51,17 @@ ''; }; + # GPU dev shell: `nix develop .#gpu`. Adds the CUDA toolkit and points the + # loader at both it and the host driver's libcuda. + devShells.gpu = pkgs.mkShell { + packages = commonPackages ++ [cudatoolkit]; + shellHook = '' + export RUST_LOG=info + export CUDA_PATH="${cudatoolkit}" + export LD_LIBRARY_PATH="${cudatoolkit}/lib:/run/opengl-driver/lib''${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}" + ''; + }; + apps.test = { type = "app"; packages = commonPackages; diff --git a/ext/src/lib.rs b/ext/src/lib.rs index a0b9b6fd7b..f261900616 100644 --- a/ext/src/lib.rs +++ b/ext/src/lib.rs @@ -181,6 +181,8 @@ use crate::chain_complex::FiniteChainComplex; pub type CCC = FiniteChainComplex; pub mod nassau; +#[cfg(feature = "gpu")] +pub mod nassau_gpu; pub mod secondary; pub mod utils; diff --git a/ext/src/nassau.rs b/ext/src/nassau.rs index 8d1919136a..709cb427ea 100644 --- a/ext/src/nassau.rs +++ b/ext/src/nassau.rs @@ -388,6 +388,74 @@ enum Magic { Fix = -3, } +/// Build the partial matrix of a differential, dispatching to the GPU Milnor-multiply +/// path when it is compiled in, opted into (`NASSAU_GPU`), applicable, and the launch is +/// large enough to amortise the fixed per-launch GPU cost. +/// +/// Defaults to the CPU per-term sweep, so behaviour is unchanged unless a caller sets +/// `NASSAU_GPU`. A resolution issues thousands of small signature-masked launches (avg +/// ~10³ term-pairs), for which the GPU's per-launch overhead (kernel launch + readback +/// sync, ~0.7 ms) dwarfs the multiply; only launches whose `rows × cols` exceeds +/// `NASSAU_GPU_MIN_WORK` (default 4M) are offloaded. `NASSAU_GPU_VERIFY` builds the CPU +/// matrix too and asserts they agree. Without the `gpu` feature this is exactly +/// `diff.get_partial_matrix(t, mask)`. +fn build_partial_matrix( + diff: &FreeModuleHomomorphism>, + t: i32, + mask: &[usize], +) -> Matrix { + #[cfg(feature = "gpu")] + { + if std::env::var_os("NASSAU_GPU").is_some() && crate::nassau_gpu::applicable(diff) { + let min_work: u64 = std::env::var("NASSAU_GPU_MIN_WORK") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(4_000_000); + let work = mask.len() as u64 * diff.target().dimension(t) as u64; + if work >= min_work { + return if std::env::var_os("NASSAU_GPU_VERIFY").is_some() { + crate::nassau_gpu::get_partial_matrix_verified(diff, t, mask) + } else { + crate::nassau_gpu::get_partial_matrix(diff, t, mask) + }; + } + } + } + diff.get_partial_matrix(t, mask) +} + +/// Whether to compute the full differential matrix once per bidegree and reuse row slices +/// across the signature passes, instead of relaunching the multiply once per signature. +/// +/// Each signature's [`build_partial_matrix`] is a *row subset* of one full matrix — the +/// masks partition the source basis, so the per-signature builds together compute every +/// row exactly once, the same total multiply work as one all-rows build. On the CPU that +/// restructuring is roughly neutral, but for the GPU it turns thousands of small +/// (often sub-threshold, CPU-fallback) launches into one big launch per bidegree that +/// amortises all fixed per-launch overhead. So it is gated on the same opt-in as the GPU +/// path; without it the per-signature build is unchanged. +fn reuse_full_matrix(_diff: &FreeModuleHomomorphism>) -> bool { + #[cfg(feature = "gpu")] + { + std::env::var_os("NASSAU_GPU").is_some() && crate::nassau_gpu::applicable(_diff) + } + #[cfg(not(feature = "gpu"))] + { + false + } +} + +/// Extract `rows` of `full` into a fresh matrix (`out.row(i) = full.row(rows[i])`), +/// preserving the column layout. Slices a precomputed full differential matrix into one +/// signature's partial matrix (see [`reuse_full_matrix`]). +fn select_rows(full: &Matrix, rows: &[usize]) -> Matrix { + let mut out = Matrix::new(full.prime(), rows.len(), full.columns()); + for (dst, &src) in rows.iter().enumerate() { + out.row_mut(dst).assign(full.row(src)); + } + out +} + /// A resolution of `S_2` using Nassau's algorithm. /// /// This aims to have an API similar to that of @@ -616,9 +684,26 @@ impl> Resolution { .collect(); let next_masked_dim = next_mask.len(); - let full_matrix = { + // Compute the full differential matrix once when reuse is active, then slice each + // signature's rows out of it instead of relaunching the multiply per signature. + let full_reuse: Option = if reuse_full_matrix(&self.differentials[b.s() - 1]) { + let all_rows: Vec = (0..target_dim).collect(); let _guard = ParallelGuard::new(); - self.differentials[b.s() - 1].get_partial_matrix(b.t(), &target_mask) + Some(build_partial_matrix( + &self.differentials[b.s() - 1], + b.t(), + &all_rows, + )) + } else { + None + }; + + let full_matrix = match &full_reuse { + Some(full) => select_rows(full, &target_mask), + None => { + let _guard = ParallelGuard::new(); + build_partial_matrix(&self.differentials[b.s() - 1], b.t(), &target_mask) + } }; let mut masked_matrix = AugmentedMatrix::new(p, target_masked_dim, [next_masked_dim, target_masked_dim]); @@ -685,10 +770,12 @@ impl> Resolution { target_mask.extend(subalgebra.signature_mask(&algebra, target, b.t(), &signature)); next_mask.extend(subalgebra.signature_mask(&algebra, next, b.t(), &signature)); - let full_matrix = { - let _guard = ParallelGuard::new(); - self.differential(b.s() - 1) - .get_partial_matrix(b.t(), &target_mask) + let full_matrix = match &full_reuse { + Some(full) => select_rows(full, &target_mask), + None => { + let _guard = ParallelGuard::new(); + build_partial_matrix(&self.differential(b.s() - 1), b.t(), &target_mask) + } }; let mut masked_matrix = diff --git a/ext/src/nassau_gpu.rs b/ext/src/nassau_gpu.rs new file mode 100644 index 0000000000..0f3f3ecbc2 --- /dev/null +++ b/ext/src/nassau_gpu.rs @@ -0,0 +1,146 @@ +//! GPU-accelerated `get_partial_matrix` for Nassau's Milnor differentials. +//! +//! A Nassau differential is a `FreeModuleHomomorphism>`; +//! building its (partial) matrix applies the differential to each input basis element, +//! whose dominant cost is the Milnor multiply `Sq(R) · s` in [`FreeModule::act`]. This +//! module batches every such multiply of one matrix build into a single GPU launch via +//! [`algebra::milnor_gpu::multiply_batch_on_gpu`]. See `benches/milnor_gpu_ab.rs` for a +//! CPU/GPU A/B of that launch. +//! +//! Only the *multiply* work is offloaded. Identity operations (`Sq(∅) = 1`, i.e. +//! `operation_degree == 0`) are plain copies with no admissible-matrix work, so they +//! are left to the CPU `apply_to_basis_element` per row. The output F₂ bits the kernel +//! returns are XORed into the matrix rows (bit `i` → `add_basis_element(i, 1)`), the +//! same layout the CPU path produces. +//! +//! Gated behind the `gpu` feature. Callers must ensure +//! [`MilnorAlgebra::gpu_multiply_applicable`] (`p = 2`, trivial profile, stable) — the +//! Nassau `S_2` regime — and that the algebra's basis and seqno tables reach `degree`. + +use algebra::{ + MilnorAlgebra, + milnor_gpu::{GpuProduct, multiply_batch_on_gpu}, + module::{ + FreeModule, Module, + homomorphism::{FreeModuleHomomorphism, ModuleHomomorphism}, + }, +}; +use fp::matrix::Matrix; + +type NassauDifferential = FreeModuleHomomorphism>; + +/// Whether the GPU `get_partial_matrix` path applies to this differential — the +/// seqno-table regime (`p = 2`, trivial profile, stable), i.e. Nassau `S_2`. The +/// (cheap, idempotent) seqno tables are built on demand in [`get_partial_matrix`]. +pub fn applicable(hom: &NassauDifferential) -> bool { + hom.target().algebra().gpu_multiply_applicable() +} + +/// GPU analogue of [`ModuleHomomorphism::get_partial_matrix`] for a Nassau differential. +/// +/// Produces the same matrix as the CPU path: rows are the differential applied to +/// `inputs[i]`, columns the target basis at `degree`. The multiply work of every input +/// is fused into one GPU launch; identity operations are filled in on the CPU. +/// +/// Callers must ensure [`applicable`]; this mirrors the extraction the CPU +/// [`FreeModuleHomomorphism::apply_to_basis_element`] performs. +pub fn get_partial_matrix(hom: &NassauDifferential, degree: i32, inputs: &[usize]) -> Matrix { + let (mut matrix, products) = extract(hom, degree, inputs); + if !products.is_empty() { + let target = hom.target(); + let algebra = target.algebra(); + // Idempotent + cheap (O(degree · width)); returns immediately once built. + algebra.compute_seqno_tables(degree); + let num_cols = target.dimension(degree); + let rows = multiply_batch_on_gpu(&algebra, num_cols, inputs.len(), &products); + for (row, limbs) in rows.iter().enumerate() { + let mut target_row = matrix.row_mut(row); + for (limb_idx, &limb) in limbs.iter().enumerate() { + let mut bits = limb; + while bits != 0 { + let b = bits.trailing_zeros() as usize; + target_row.add_basis_element(limb_idx * 32 + b, 1); + bits &= bits - 1; + } + } + } + } + matrix +} + +/// Extract the non-identity Milnor multiplies of one matrix build as [`GpuProduct`]s, +/// mirroring `apply_to_basis_element` → [`FreeModule::act`]. Returns the matrix with its +/// identity (`Sq(∅)`) and out-of-range rows already filled, plus the products whose GPU +/// (or CPU-reference) output completes it. +fn extract(hom: &NassauDifferential, degree: i32, inputs: &[usize]) -> (Matrix, Vec) { + let p = hom.prime(); + let source = hom.source(); + let target = hom.target(); + let shift = hom.degree_shift(); + let out_dim = target.dimension(degree); + let mut matrix = Matrix::new(p, inputs.len(), out_dim); + let mut products: Vec = Vec::new(); + + if out_dim == 0 { + return (matrix, products); + } + + for (row, &input_index) in inputs.iter().enumerate() { + let ogp = source.index_to_op_gen(degree, input_index); + if ogp.generator_degree < hom.min_degree() { + continue; + } + if ogp.operation_degree == 0 { + // Sq(∅) · s = s: let the CPU copy it directly into this row. + hom.apply_to_basis_element(matrix.row_mut(row), 1, degree, input_index); + continue; + } + let out_on_gen = hom.output(ogp.generator_degree, ogp.generator_index); + let act_input_degree = ogp.generator_degree - shift; + for gd in target + .iter_gen_offsets::<2>([act_input_degree, act_input_degree + ogp.operation_degree]) + { + let (input_start, input_end) = (gd.start[0], gd.end[0]); + if input_start >= out_on_gen.len() { + break; + } + let s_slice = out_on_gen.as_slice().restrict(input_start, input_end); + if s_slice.is_zero() { + continue; + } + products.push(GpuProduct { + r_degree: ogp.operation_degree, + r_idx: ogp.operation_index, + s_degree: act_input_degree - gd.gen_deg, + term_indices: s_slice.iter_nonzero().map(|(i, _)| i).collect(), + row, + // `Sq(R) · s` lands in this target generator's block of the output row, + // which starts at gd.start[1] (the offsets at the *output* degree). + out_offset: gd.start[1], + }); + } + } + (matrix, products) +} + +/// Build the matrix both ways and assert they agree; returns the CPU matrix. For the +/// `NASSAU_GPU_VERIFY` env gate — validates the GPU path over a real resolution before +/// it is trusted unconditionally. +pub fn get_partial_matrix_verified( + hom: &NassauDifferential, + degree: i32, + inputs: &[usize], +) -> Matrix { + let gpu = get_partial_matrix(hom, degree, inputs); + let cpu = hom.get_partial_matrix(degree, inputs); + for row in 0..inputs.len() { + let g: Vec = gpu.row(row).iter_nonzero().map(|(i, _)| i).collect(); + let c: Vec = cpu.row(row).iter_nonzero().map(|(i, _)| i).collect(); + assert_eq!( + g, c, + "GPU/CPU get_partial_matrix mismatch at degree {degree}, row {row} (input {})", + inputs[row], + ); + } + cpu +} diff --git a/ext/tests/nassau_gpu.rs b/ext/tests/nassau_gpu.rs new file mode 100644 index 0000000000..1adcfd37dd --- /dev/null +++ b/ext/tests/nassau_gpu.rs @@ -0,0 +1,50 @@ +//! End-to-end validation of the GPU `get_partial_matrix` wire-in against the CPU path. +//! +//! Resolves `S_2` (Nassau) over a modest region and, for every non-trivial differential +//! bidegree, asserts [`ext::nassau_gpu::get_partial_matrix`] reproduces the CPU +//! [`ModuleHomomorphism::get_partial_matrix`] bit-for-bit. Requires a live CUDA device + +//! the toolkit env (run under the `gpu` dev shell, unsandboxed); the whole test is +//! `gpu`-gated so it is skipped in ordinary CI. +#![cfg(feature = "gpu")] + +use algebra::module::{Module, homomorphism::ModuleHomomorphism}; +use ext::{chain_complex::ChainComplex, nassau_gpu, utils::construct_nassau}; +use sseq::coordinates::Bidegree; + +#[test] +fn gpu_partial_matrix_matches_cpu() { + let stem = 30; + let filt = 17; + let res = construct_nassau(("S_2", "milnor"), None).unwrap(); + res.compute_through_stem(Bidegree::n_s(stem, filt)); + + let mut checked_bidegrees = 0usize; + let mut checked_rows = 0usize; + for s in 1..=filt { + let diff = res.differential(s); + for t in s..=(s + stem) { + let source_dim = diff.source().dimension(t); + let target_dim = diff.target().dimension(t); + if source_dim == 0 || target_dim == 0 { + continue; + } + // Full launch: every source basis element as an input, like the resolution's + // `get_partial_matrix` call over a signature mask (here the trivial mask). + let inputs: Vec = (0..source_dim).collect(); + let gpu = nassau_gpu::get_partial_matrix(&diff, t, &inputs); + let cpu = diff.get_partial_matrix(t, &inputs); + for row in 0..source_dim { + let g: Vec = gpu.row(row).iter_nonzero().map(|(i, _)| i).collect(); + let c: Vec = cpu.row(row).iter_nonzero().map(|(i, _)| i).collect(); + assert_eq!(g, c, "mismatch at (s={s}, t={t}), row {row}"); + } + checked_bidegrees += 1; + checked_rows += source_dim; + } + } + assert!( + checked_bidegrees > 0, + "no non-trivial bidegrees were checked" + ); + eprintln!("verified {checked_rows} rows across {checked_bidegrees} bidegrees"); +} diff --git a/ext/tests/nassau_gpu_reuse.rs b/ext/tests/nassau_gpu_reuse.rs new file mode 100644 index 0000000000..7b81e1c429 --- /dev/null +++ b/ext/tests/nassau_gpu_reuse.rs @@ -0,0 +1,55 @@ +//! End-to-end validation of the GPU compute-once-reuse path in `step_resolution`. +//! +//! Resolves `S_2` (Nassau) twice — once with the CPU per-signature `build_partial_matrix` +//! and once with `NASSAU_GPU` set, which computes the full differential matrix once per +//! bidegree on the GPU and slices each signature's rows out of it — and asserts the two +//! resolutions agree bidegree-by-bidegree. This exercises the `reuse_full_matrix` / +//! `select_rows` restructuring (and the `dx == 0` correction invariant inside the +//! signature loop), which the per-matrix `nassau_gpu` test does not touch. +//! +//! Requires a live CUDA device + toolkit env (run under the `gpu` dev shell, unsandboxed); +//! `gpu`-gated so it is skipped in ordinary CI. Isolated in its own test binary because it +//! toggles the `NASSAU_GPU` process env var, which would race parallel tests. +#![cfg(feature = "gpu")] + +use ext::{chain_complex::FreeChainComplex, utils::construct_nassau}; +use sseq::coordinates::Bidegree; + +fn betti(stem: i32, filt: i32) -> Vec<(i32, i32, usize)> { + let res = construct_nassau(("S_2", "milnor"), None).unwrap(); + res.compute_through_stem(Bidegree::n_s(stem, filt)); + let mut out = Vec::new(); + for s in 0..=filt { + for n in 0..=stem { + let b = Bidegree::n_s(n, s); + out.push((n, s, res.number_of_gens_in_bidegree(b))); + } + } + out +} + +#[test] +fn gpu_reuse_resolution_matches_cpu() { + let stem: i32 = std::env::var("STEM") + .ok() + .and_then(|s| s.parse().ok()) + .unwrap_or(30); + let filt: i32 = std::env::var("FILT") + .ok() + .and_then(|s| s.parse().ok()) + .unwrap_or(17); + + // SAFETY: single-threaded test (own binary); the env var only gates build dispatch. + unsafe { std::env::remove_var("NASSAU_GPU") }; + let cpu = betti(stem, filt); + + unsafe { std::env::set_var("NASSAU_GPU", "1") }; + let gpu = betti(stem, filt); + + assert_eq!( + cpu, gpu, + "GPU compute-once-reuse resolution disagrees with the CPU resolution" + ); + let total: usize = cpu.iter().map(|&(_, _, g)| g).sum(); + eprintln!("matched {total} generators across {} bidegrees", cpu.len()); +} diff --git a/ext/tests/nassau_gpu_timing.rs b/ext/tests/nassau_gpu_timing.rs new file mode 100644 index 0000000000..04b83d0c92 --- /dev/null +++ b/ext/tests/nassau_gpu_timing.rs @@ -0,0 +1,79 @@ +//! End-to-end timing of the GPU wire-in: resolve `S_2` (Nassau) with and without the +//! GPU `get_partial_matrix` path and print wall times. `#[ignore]` by default (it is a +//! benchmark, not a correctness check); run explicitly under the `gpu` dev shell: +//! `cargo test --test nassau_gpu_timing --features gpu -- --ignored --nocapture --test-threads=1`. +//! `--test-threads=1` is required: both tests toggle the process-wide `NASSAU_GPU` env var, so +//! running them concurrently would race (and `set_var`/`remove_var` are themselves unsound under +//! concurrency). +#![cfg(feature = "gpu")] + +use std::time::Instant; + +use ext::utils::construct_nassau; +use sseq::coordinates::Bidegree; + +fn resolve_time(stem: i32, filt: i32) -> f64 { + let res = construct_nassau(("S_2", "milnor"), None).unwrap(); + let t0 = Instant::now(); + res.compute_through_stem(Bidegree::n_s(stem, filt)); + t0.elapsed().as_secs_f64() +} + +/// GPU-only resolution (no CPU baseline) — for iterating on the GPU path at large stems +/// without paying the multi-minute CPU run each time. `STEM`/`FILT` env, default 100/55. +#[test] +#[ignore = "GPU-only benchmark; run explicitly with --ignored"] +fn gpu_only_resolution_time() { + let stem: i32 = std::env::var("STEM") + .ok() + .and_then(|s| s.parse().ok()) + .unwrap_or(100); + let filt: i32 = std::env::var("FILT") + .ok() + .and_then(|s| s.parse().ok()) + .unwrap_or(55); + // SAFETY: single-threaded test setup; the env var only gates build_partial_matrix. + unsafe { std::env::set_var("NASSAU_GPU", "1") }; + let _ = algebra::milnor_gpu::take_batch_stats(); + let gpu = resolve_time(stem, filt); + let (calls, marshal_us, device_us, pairs) = algebra::milnor_gpu::take_batch_stats(); + eprintln!("GPU-only (stem {stem}, filt {filt}): {gpu:.1}s"); + eprintln!( + "GPU launches: {calls} | host marshal {:.2}s device {:.2}s | avg {:.0} pairs/launch", + marshal_us as f64 / 1e6, + device_us as f64 / 1e6, + pairs as f64 / calls.max(1) as f64, + ); +} + +#[test] +#[ignore = "timing benchmark; run explicitly with --ignored"] +fn gpu_vs_cpu_resolution_time() { + let stem: i32 = std::env::var("STEM") + .ok() + .and_then(|s| s.parse().ok()) + .unwrap_or(50); + let filt: i32 = std::env::var("FILT") + .ok() + .and_then(|s| s.parse().ok()) + .unwrap_or(27); + + // SAFETY: single-threaded test setup; the env var only gates build_partial_matrix. + unsafe { std::env::remove_var("NASSAU_GPU") }; + let cpu = resolve_time(stem, filt); + eprintln!("CPU (stem {stem}, filt {filt}): {cpu:.1}s"); + + let _ = algebra::milnor_gpu::take_batch_stats(); // reset before the GPU run + unsafe { std::env::set_var("NASSAU_GPU", "1") }; + let gpu = resolve_time(stem, filt); + eprintln!("GPU (stem {stem}, filt {filt}): {gpu:.1}s"); + + let (calls, marshal_us, device_us, pairs) = algebra::milnor_gpu::take_batch_stats(); + eprintln!( + "GPU launches: {calls} | host marshal {:.2}s device {:.2}s | avg {:.0} pairs/launch", + marshal_us as f64 / 1e6, + device_us as f64 / 1e6, + pairs as f64 / calls.max(1) as f64, + ); + eprintln!("speedup (CPU / GPU): {:.2}×", cpu / gpu); +}