From 2d229b466bf1877ba00c84ee7e47c92155b5657f Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Sat, 11 Jul 2026 02:38:55 -0400 Subject: [PATCH 01/10] Add a criterion benchmark suite for the algebra crate Criterion benches for Milnor multiplication, module actions, and the Nassau-regime operation-by-element multiply, sharing a small harness in benches/common. Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_01Fsxqsgjoa5RWYhAvRG1bT2 --- ext/crates/algebra/Cargo.toml | 15 +- ext/crates/algebra/benches/common/mod.rs | 325 +++++++++++++++++++ ext/crates/algebra/benches/milnor.rs | 121 ++++--- ext/crates/algebra/benches/module_action.rs | 58 ++++ ext/crates/algebra/benches/multiplication.rs | 43 +++ ext/crates/algebra/benches/nassau_milnor.rs | 125 +++++++ 6 files changed, 649 insertions(+), 38 deletions(-) create mode 100644 ext/crates/algebra/benches/common/mod.rs create mode 100644 ext/crates/algebra/benches/module_action.rs create mode 100644 ext/crates/algebra/benches/multiplication.rs create mode 100644 ext/crates/algebra/benches/nassau_milnor.rs diff --git a/ext/crates/algebra/Cargo.toml b/ext/crates/algebra/Cargo.toml index 9a45a269fb..3b98825cb5 100644 --- a/ext/crates/algebra/Cargo.toml +++ b/ext/crates/algebra/Cargo.toml @@ -29,8 +29,9 @@ serde_json = "1.0.141" enum_dispatch = "0.3.13" [dev-dependencies] -bencher = "0.1.5" +criterion = { version = "0.5", features = ["html_reports"] } expect-test = "1.5.1" +pprof = { version = "0.15.0", features = ["criterion", "flamegraph"] } rstest = "0.25.0" [features] @@ -42,3 +43,15 @@ odd-primes = ["fp/odd-primes"] [[bench]] name = "milnor" harness = false + +[[bench]] +name = "multiplication" +harness = false + +[[bench]] +name = "module_action" +harness = false + +[[bench]] +name = "nassau_milnor" +harness = false diff --git a/ext/crates/algebra/benches/common/mod.rs b/ext/crates/algebra/benches/common/mod.rs new file mode 100644 index 0000000000..01b7c359e7 --- /dev/null +++ b/ext/crates/algebra/benches/common/mod.rs @@ -0,0 +1,325 @@ +//! Shared helpers for the algebra crate's criterion benchmarks. +//! +//! These are generic over *any* [`Algebra`] or [`Module`], so the same routines can benchmark the +//! Adem and Milnor bases (see `multiplication.rs`) or any of the predefined Steenrod modules in +//! `ext/steenrod_modules` (see `module_action.rs`). +//! +//! Every bench target that needs these includes the file with `mod common;`; each target compiles +//! its own copy, which is the standard criterion pattern for sharing code between benches. + +// Not every bench target uses every helper, and each compiles this module independently, so silence +// the unavoidable dead-code warnings. +#![allow(dead_code)] + +use std::{hint::black_box, path::PathBuf, sync::Arc}; + +use algebra::{ + Algebra, AlgebraType, SteenrodAlgebra, + module::{Module, SteenrodModule, steenrod_module}, +}; +use criterion::{BenchmarkGroup, Throughput, measurement::WallTime}; +use fp::{ + prime::{Prime, ValidPrime}, + vector::FpVector, +}; +use serde_json::Value; + +/// The directory shipping the predefined Steenrod modules. We read the JSON directly rather than +/// going through `ext::utils`, because the `ext` crate depends on `algebra` (a dependency cycle). +pub const STEENROD_MODULES_DIR: &str = + concat!(env!("CARGO_MANIFEST_DIR"), "/../../steenrod_modules"); + +/// Build a dense element of the given length with every entry non-zero and reduced mod `p`. Used to +/// exercise the general-element multiplication and action code paths. +pub fn dense_vector(p: ValidPrime, len: usize) -> FpVector { + let mut v = FpVector::new(p, len); + let modulus = p.as_u32() - 1; + for i in 0..len { + // Entries in `1..p`, so the vector is genuinely dense. + v.set_entry(i, (i as u32 % modulus) + 1); + } + v +} + +// ------------------------------------------------------------------------------------------------- +// Algebra multiplication +// ------------------------------------------------------------------------------------------------- + +/// Benchmark multiplication in `algebra` for each `(r_degree, s_degree)` pair in `pairs`. +/// +/// For each pair we register two benchmarks: +/// - `basis_products`: multiply every pair of basis elements via +/// [`Algebra::multiply_basis_elements`], the required multiplication primitive. +/// - `element_product`: a single [`Algebra::multiply_element_by_element`] of two dense elements, +/// exercising the general-element path. +pub fn bench_algebra_multiplication( + group: &mut BenchmarkGroup, + label: &str, + algebra: &A, + pairs: &[(i32, i32)], +) { + let p = algebra.prime(); + for &(r_degree, s_degree) in pairs { + let out_degree = r_degree + s_degree; + // `compute_basis` is cumulative, so this readies every degree up to `out_degree`. + algebra.compute_basis(out_degree); + let r_dim = algebra.dimension(r_degree); + let s_dim = algebra.dimension(s_degree); + let out_dim = algebra.dimension(out_degree); + if r_dim == 0 || s_dim == 0 { + continue; + } + + group.throughput(Throughput::Elements((r_dim * s_dim) as u64)); + + { + let mut scratch = FpVector::new(p, out_dim); + group.bench_function( + format!("{label}/basis_products/{r_degree}x{s_degree}"), + |b| { + b.iter(|| { + scratch.set_to_zero(); + for i in 0..r_dim { + for j in 0..s_dim { + algebra.multiply_basis_elements( + scratch.as_slice_mut(), + 1, + r_degree, + i, + s_degree, + j, + ); + } + } + black_box(&scratch); + }); + }, + ); + } + + { + let r = dense_vector(p, r_dim); + let s = dense_vector(p, s_dim); + let mut scratch = FpVector::new(p, out_dim); + group.bench_function( + format!("{label}/element_product/{r_degree}x{s_degree}"), + |b| { + b.iter(|| { + scratch.set_to_zero(); + algebra.multiply_element_by_element( + scratch.as_slice_mut(), + 1, + r_degree, + r.as_slice(), + s_degree, + s.as_slice(), + ); + black_box(&scratch); + }); + }, + ); + } + } +} + +// ------------------------------------------------------------------------------------------------- +// Module action +// ------------------------------------------------------------------------------------------------- + +/// The lowest few positive operation degrees in which `algebra` is non-trivial, up to `max` of them +/// and bounded by `op_cap`. Keeps the number of benchmark cases per module small and bounded. +fn representative_op_degrees(algebra: &SteenrodAlgebra, op_cap: i32, max: usize) -> Vec { + let mut degrees = Vec::new(); + let mut d = 1; + while d <= op_cap && degrees.len() < max { + algebra.compute_basis(d); + if algebra.dimension(d) > 0 { + degrees.push(d); + } + d += 1; + } + degrees +} + +/// Benchmark the module action of `module` up to total degree `cap`. +/// +/// For a handful of representative operation degrees, we register two benchmarks that sweep over all +/// module degrees (bounded by `cap`, which keeps infinite modules such as `RP_inf` in check): +/// - `act_on_basis`: apply every operation basis element to every module basis element via the +/// required primitive [`Module::act_on_basis`]. +/// - `act`: apply every operation basis element to a dense element via [`Module::act`]. +pub fn bench_module_action( + group: &mut BenchmarkGroup, + label: &str, + module: &dyn Module, + cap: i32, +) { + let algebra = module.algebra(); + let p = module.prime(); + + let min_degree = module.min_degree(); + let max_degree = module.max_degree().map_or(cap, |m| m.min(cap)); + if max_degree < min_degree { + return; + } + + let op_cap = cap.min(12); + let op_degrees = representative_op_degrees(&algebra, op_cap, 4); + + for op_degree in op_degrees { + // Ready every module (and algebra) degree we might touch. + algebra.compute_basis(op_degree); + module.compute_basis(max_degree + op_degree); + let op_dim = algebra.dimension(op_degree); + if op_dim == 0 { + continue; + } + + // The module degrees that actually contribute work at this operation degree. + let sweep: Vec<(i32, usize, usize)> = (min_degree..=max_degree) + .filter_map(|mod_degree| { + let mod_dim = module.dimension(mod_degree); + let out_dim = module.dimension(mod_degree + op_degree); + (mod_dim > 0 && out_dim > 0).then_some((mod_degree, mod_dim, out_dim)) + }) + .collect(); + if sweep.is_empty() { + continue; + } + + let total_pairs: u64 = sweep + .iter() + .map(|&(_, mod_dim, _)| (op_dim * mod_dim) as u64) + .sum(); + group.throughput(Throughput::Elements(total_pairs)); + + { + let mut scratch = FpVector::new(p, 0); + group.bench_function(format!("{label}/act_on_basis/op{op_degree}"), |b| { + b.iter(|| { + for &(mod_degree, mod_dim, out_dim) in &sweep { + scratch.set_scratch_vector_size(out_dim); + scratch.set_to_zero(); + for op_idx in 0..op_dim { + for mod_idx in 0..mod_dim { + module.act_on_basis( + scratch.as_slice_mut(), + 1, + op_degree, + op_idx, + mod_degree, + mod_idx, + ); + } + } + black_box(&scratch); + } + }); + }); + } + + { + let inputs: Vec<(i32, FpVector, usize)> = sweep + .iter() + .map(|&(mod_degree, mod_dim, out_dim)| { + (mod_degree, dense_vector(p, mod_dim), out_dim) + }) + .collect(); + let mut scratch = FpVector::new(p, 0); + group.bench_function(format!("{label}/act/op{op_degree}"), |b| { + b.iter(|| { + for op_idx in 0..op_dim { + for (mod_degree, input, out_dim) in &inputs { + scratch.set_scratch_vector_size(*out_dim); + scratch.set_to_zero(); + module.act( + scratch.as_slice_mut(), + 1, + op_degree, + op_idx, + *mod_degree, + input.as_slice(), + ); + black_box(&scratch); + } + } + }); + }); + } + } +} + +// ------------------------------------------------------------------------------------------------- +// Loading the predefined Steenrod modules +// ------------------------------------------------------------------------------------------------- + +/// All `*.json` module files shipped in [`STEENROD_MODULES_DIR`], sorted by file name for a stable +/// benchmark ordering. +pub fn module_json_files() -> Vec { + let mut files: Vec = std::fs::read_dir(STEENROD_MODULES_DIR) + .unwrap_or_else(|e| panic!("failed to read {STEENROD_MODULES_DIR}: {e}")) + .filter_map(Result::ok) + .map(|entry| entry.path()) + .filter(|path| path.extension().is_some_and(|ext| ext == "json")) + .collect(); + files.sort(); + files +} + +/// A parsed module JSON together with its file stem (used as the benchmark label). +pub struct ModuleSpec { + pub name: String, + pub json: Value, +} + +/// Parse every module file, returning `(name, json)` for each. Files that fail to read or parse are +/// reported to stderr and skipped. +pub fn load_module_specs() -> Vec { + let mut specs = Vec::new(); + for path in module_json_files() { + let name = path + .file_stem() + .and_then(|s| s.to_str()) + .unwrap_or("") + .to_owned(); + match std::fs::read_to_string(&path) + .map_err(anyhow::Error::from) + .and_then(|s| serde_json::from_str::(&s).map_err(anyhow::Error::from)) + { + Ok(json) => specs.push(ModuleSpec { name, json }), + Err(e) => eprintln!("skipping module {name}: could not parse JSON ({e})"), + } + } + specs +} + +/// Whether the module spec requires machinery only available in the `ext` crate. The one such case +/// is a `cofiber`, which is realized as a Yoneda representative via `Resolution` / +/// `yoneda_representative` — neither of which lives in the `algebra` crate. Such modules cannot be +/// constructed from raw JSON here and are skipped. +pub fn requires_ext_machinery(json: &Value) -> bool { + !json["cofiber"].is_null() +} + +/// Whether the module's optional `"algebra"` restriction permits the given basis. Modules without an +/// `"algebra"` field support both bases. +pub fn module_supports(json: &Value, ty: AlgebraType) -> bool { + match json["algebra"].as_array() { + None => true, + Some(list) => { + let wanted = ty.to_string(); + list.iter().any(|x| x.as_str() == Some(&wanted)) + } + } +} + +/// Build the algebra and module for a given spec and basis, using only the `algebra` crate's JSON +/// constructors. +pub fn load_named_module( + json: &Value, + ty: AlgebraType, +) -> anyhow::Result<(Arc, SteenrodModule)> { + let algebra = Arc::new(SteenrodAlgebra::from_json(json, ty, false)?); + let module = steenrod_module::from_json(Arc::clone(&algebra), json)?; + Ok((algebra, module)) +} diff --git a/ext/crates/algebra/benches/milnor.rs b/ext/crates/algebra/benches/milnor.rs index 86e1cd711a..9dcac82dcd 100644 --- a/ext/crates/algebra/benches/milnor.rs +++ b/ext/crates/algebra/benches/milnor.rs @@ -1,54 +1,101 @@ +//! Benchmarks for the low-level Milnor `PPartMultiplier` kernel. + use algebra::milnor_algebra::{PPartAllocation, PPartEntry, PPartMultiplier}; -use bencher::{Bencher, benchmark_group, benchmark_main}; +use criterion::{ + BenchmarkGroup, Criterion, criterion_group, criterion_main, measurement::WallTime, +}; use fp::prime::{Prime, ValidPrime}; +use pprof::criterion::{Output, PProfProfiler}; -fn ppart_inner( - bench: &mut Bencher, +fn bench_ppart( + g: &mut BenchmarkGroup, + name: &str, p: u32, r: Vec, s: Vec, ) { let p = ValidPrime::new(p); + g.bench_function(name, |bench| { + bench.iter(|| { + let m = PPartMultiplier::::new_from_allocation( + p, + &r, + &s, + PPartAllocation::default(), + 0, + 0, + ); - bench.iter(move || { - let m = PPartMultiplier::::new_from_allocation( - p, - &r, - &s, - PPartAllocation::default(), - 0, - 0, - ); - - for c in m { - if MOD4 { - assert!(c < 4); - } else { - assert!(c < p.as_u32()); + for c in m { + if MOD4 { + assert!(c < 4); + } else { + assert!(c < p.as_u32()); + } } - } + }); }); } -fn ppart_2(bench: &mut Bencher) { - ppart_inner::(bench, 2, vec![60, 30, 8, 2, 1], vec![20, 30, 20, 4, 1, 2]); - ppart_inner::(bench, 2, vec![35, 12, 20, 14, 1, 3], vec![60, 30, 0, 2, 1]); -} +fn ppart(c: &mut Criterion) { + let mut g = c.benchmark_group("milnor_ppart"); -fn ppart_4(bench: &mut Bencher) { - ppart_inner::(bench, 2, vec![60, 30, 8, 2, 1], vec![20, 30, 20, 4, 1, 2]); - ppart_inner::(bench, 2, vec![35, 12, 20, 14, 1, 3], vec![60, 30, 0, 2, 1]); -} + bench_ppart::( + &mut g, + "ppart_2/a", + 2, + vec![60, 30, 8, 2, 1], + vec![20, 30, 20, 4, 1, 2], + ); + bench_ppart::( + &mut g, + "ppart_2/b", + 2, + vec![35, 12, 20, 14, 1, 3], + vec![60, 30, 0, 2, 1], + ); -#[cfg(feature = "odd-primes")] -fn ppart_3(bench: &mut Bencher) { - ppart_inner::(bench, 3, vec![120, 70, 40, 2], vec![60, 35, 21, 6]); - ppart_inner::(bench, 3, vec![30, 12, 35, 24], vec![100, 80, 16, 2, 3]); -} + bench_ppart::( + &mut g, + "ppart_4/a", + 2, + vec![60, 30, 8, 2, 1], + vec![20, 30, 20, 4, 1, 2], + ); + bench_ppart::( + &mut g, + "ppart_4/b", + 2, + vec![35, 12, 20, 14, 1, 3], + vec![60, 30, 0, 2, 1], + ); -#[cfg(feature = "odd-primes")] -benchmark_group!(benches, ppart_2, ppart_4, ppart_3); + #[cfg(feature = "odd-primes")] + { + bench_ppart::( + &mut g, + "ppart_3/a", + 3, + vec![120, 70, 40, 2], + vec![60, 35, 21, 6], + ); + bench_ppart::( + &mut g, + "ppart_3/b", + 3, + vec![30, 12, 35, 24], + vec![100, 80, 16, 2, 3], + ); + } -#[cfg(not(feature = "odd-primes"))] -benchmark_group!(benches, ppart_2, ppart_4); -benchmark_main!(benches); + 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 = ppart +} +criterion_main!(benches); diff --git a/ext/crates/algebra/benches/module_action.rs b/ext/crates/algebra/benches/module_action.rs new file mode 100644 index 0000000000..19ab32174d --- /dev/null +++ b/ext/crates/algebra/benches/module_action.rs @@ -0,0 +1,58 @@ +//! Benchmarks for the module action, run against every predefined module in +//! `ext/steenrod_modules`. +//! +//! Each module is loaded directly from its JSON specification (via the `algebra` crate's own +//! constructors) and benchmarked under both the Milnor and Adem bases it supports. Modules that +//! require `ext`-level machinery (i.e. those defined via a `cofiber`) are skipped. +//! +//! To benchmark a single module, filter by name, e.g. +//! `cargo bench --bench module_action -- Joker`. + +mod common; + +use algebra::AlgebraType; +use criterion::{Criterion, criterion_group, criterion_main}; +use pprof::criterion::{Output, PProfProfiler}; + +/// Cap on total degree, so infinite or high-dimensional modules (`RP_inf`, finitely presented +/// modules, …) stay bounded. +const DEGREE_CAP: i32 = 20; + +fn module_action(c: &mut Criterion) { + for spec in common::load_module_specs() { + if common::requires_ext_machinery(&spec.json) { + eprintln!( + "skipping module {}: cofiber modules require ext-level resolution machinery", + spec.name + ); + continue; + } + + for ty in [AlgebraType::Milnor, AlgebraType::Adem] { + if !common::module_supports(&spec.json, ty) { + continue; + } + let basis = ty.to_string(); + let (_algebra, module) = match common::load_named_module(&spec.json, ty) { + Ok(pair) => pair, + Err(e) => { + eprintln!("skipping module {} ({basis}): {e}", spec.name); + continue; + } + }; + + let mut g = c.benchmark_group(format!("module_action/{}", spec.name)); + common::bench_module_action(&mut g, &basis, &*module, DEGREE_CAP); + 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 = module_action +} +criterion_main!(benches); diff --git a/ext/crates/algebra/benches/multiplication.rs b/ext/crates/algebra/benches/multiplication.rs new file mode 100644 index 0000000000..cd30e46a43 --- /dev/null +++ b/ext/crates/algebra/benches/multiplication.rs @@ -0,0 +1,43 @@ +//! Benchmarks for algebra multiplication. +//! +//! Exercises both the Adem and Milnor bases at primes 2, 3 and 5 (odd primes require the +//! `odd-primes` feature). See `common::bench_algebra_multiplication` for what is measured. + +mod common; + +use algebra::{AdemAlgebra, MilnorAlgebra}; +use criterion::{Criterion, criterion_group, criterion_main}; +use fp::prime::ValidPrime; +use pprof::criterion::{Output, PProfProfiler}; + +/// Benchmark both bases at prime `p` over the given `(r_degree, s_degree)` pairs. Degrees that turn +/// out to be empty for a basis are skipped by the helper. +fn bench_prime(c: &mut Criterion, p: u32, pairs: &[(i32, i32)]) { + let prime = ValidPrime::new(p); + + let adem = AdemAlgebra::new(prime, false); + let mut g = c.benchmark_group(format!("multiplication/p{p}")); + common::bench_algebra_multiplication(&mut g, "adem", &adem, pairs); + let milnor = MilnorAlgebra::new(prime, false); + common::bench_algebra_multiplication(&mut g, "milnor", &milnor, pairs); + g.finish(); +} + +fn multiplication(c: &mut Criterion) { + bench_prime(c, 2, &[(8, 8), (16, 16), (24, 24), (16, 8)]); + #[cfg(feature = "odd-primes")] + { + // Odd-prime degrees are spaced by q = 2(p - 1), so we pick multiples of q. + bench_prime(c, 3, &[(8, 8), (12, 12), (16, 16), (12, 8)]); + bench_prime(c, 5, &[(8, 8), (16, 16), (24, 24), (16, 8)]); + } +} + +criterion_group! { + name = benches; + config = Criterion::default() + .measurement_time(std::time::Duration::from_secs(3)) + .with_profiler(PProfProfiler::new(100, Output::Flamegraph(None))); + targets = multiplication +} +criterion_main!(benches); diff --git a/ext/crates/algebra/benches/nassau_milnor.rs b/ext/crates/algebra/benches/nassau_milnor.rs new file mode 100644 index 0000000000..41021a792d --- /dev/null +++ b/ext/crates/algebra/benches/nassau_milnor.rs @@ -0,0 +1,125 @@ +//! Benchmarks the Milnor multiplication kernel in the exact regime Nassau's algorithm +//! exercises when resolving the sphere `S` at the prime 2. +//! +//! ## What Nassau actually calls +//! +//! Nassau (`ext/src/nassau.rs`) resolves over `FreeModule`, which is a *stable* +//! (`U = false`) free module built on `MilnorAlgebra::new(TWO, false)`. Building the masked +//! differential matrices (`get_partial_matrix` → `FreeModuleHomomorphism::apply_to_basis_element` +//! → `FreeModule::act`) issues, per generator block, one +//! +//! ```text +//! Algebra::multiply_basis_element_by_element(result, 1, op_degree, op_index, elt_degree, elt) +//! ``` +//! +//! — a fixed *operation* basis element times a general algebra *element*. Because the free module +//! is `U = false`, the `MuAlgebra` shim discards `excess` and routes to the plain **stable** +//! multiply; every basis-by-basis product funnels through `MilnorAlgebra::multiply_with_allocation` +//! (`excess = i32::MAX`, no instability truncation) into the `PPartMultiplier::` sweep — the +//! innermost kernel. The `*_unstable` entry points are never taken. +//! +//! ## Where the regime comes from +//! +//! The `(operation_degree, element_degree)` pairs in [`REGIME`] were captured by instrumenting +//! `multiply_with_allocation` during a real `construct_nassau(("S_2", "milnor"), None)` resolved +//! through bidegree (70, 70): ~3.77M basis-by-basis products over ~2200 distinct degree pairs. The +//! cost mass sits in target degrees `op + elt ≈ 40..52`; within that, products land heavily on +//! degree-24 and degree-32 element components, alongside a frequent "large operation × small +//! element" band (e.g. a degree-40+ operation applied to a degree-1 element). The table samples +//! those regions plus a couple of smaller pairs for the cheap end of the cost curve. To regenerate +//! after algorithm changes, re-instrument `multiply_with_allocation` to tally `(m1.degree, +//! m2.degree)` and rerun the resolution (see the plan/commit history). + +use std::hint::black_box; + +use algebra::{Algebra, MilnorAlgebra}; +use criterion::{Criterion, Throughput, criterion_group, criterion_main}; +use fp::{prime::TWO, vector::FpVector}; +use pprof::criterion::{Output, PProfProfiler}; + +/// `(operation_degree, element_degree)` pairs sampled from Nassau's `S_2` @ p=2 regime. +const REGIME: &[(i32, i32)] = &[ + // cheap end of the curve + (8, 8), + (16, 16), + // operation × degree-24 element component (a hot band) + (8, 24), + (20, 24), + (24, 24), + (26, 24), + // operation × degree-32 element component + (16, 32), + (24, 32), + // operation applied to a smaller element + (24, 16), + (32, 8), + // large operation × small element (the frequent Sq^1-style band) + (40, 1), + (46, 1), +]; + +/// A dense element of the given length. At p=2 every present coordinate is 1, so this is the +/// all-ones vector — the densest element, giving a stable upper bound on the inner-loop cost. +fn dense_element(len: usize) -> FpVector { + let mut v = FpVector::new(TWO, len); + for i in 0..len { + v.set_entry(i, 1); + } + v +} + +fn nassau_milnor(c: &mut Criterion) { + // Exactly the algebra Nassau uses: the full Milnor algebra at p=2, stable (not unstable). + let algebra = MilnorAlgebra::new(TWO, false); + let mut g = c.benchmark_group("nassau_milnor"); + + for &(op_degree, elt_degree) in REGIME { + let out_degree = op_degree + elt_degree; + // `compute_basis` is cumulative, so this readies every degree up to `out_degree`. + algebra.compute_basis(out_degree); + let op_dim = algebra.dimension(op_degree); + let elt_dim = algebra.dimension(elt_degree); + let out_dim = algebra.dimension(out_degree); + if op_dim == 0 || elt_dim == 0 { + continue; + } + + // A dense element sends every basis product to the kernel: op_dim * elt_dim of them. + g.throughput(Throughput::Elements((op_dim * elt_dim) as u64)); + + let elt = dense_element(elt_dim); + let mut scratch = FpVector::new(TWO, out_dim); + g.bench_function( + format!("basis_by_element/op{op_degree}xel{elt_degree}"), + |b| { + b.iter(|| { + scratch.set_to_zero(); + // Sweep every operation basis element against the dense element — exactly the + // `multiply_basis_element_by_element` call `FreeModule::act` makes in Nassau. + for i in 0..op_dim { + algebra.multiply_basis_element_by_element( + scratch.as_slice_mut(), + 1, + op_degree, + i, + elt_degree, + elt.as_slice(), + ); + } + black_box(&scratch); + }); + }, + ); + } + + 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 = nassau_milnor +} +criterion_main!(benches); From 2f71106ad025dfe5df73117bbebf17b4dbc8b0b3 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Sat, 11 Jul 2026 02:38:55 -0400 Subject: [PATCH 02/10] Add an admissible-matrix Milnor multiply at p = 2 An operation-by-element multiply that precomputes admissible matrices. It is later moved off the CPU hot path and kept as the model for the GPU port. Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_01Fsxqsgjoa5RWYhAvRG1bT2 --- .../algebra/src/algebra/milnor_algebra.rs | 328 ++++++++++++++++++ 1 file changed, 328 insertions(+) diff --git a/ext/crates/algebra/src/algebra/milnor_algebra.rs b/ext/crates/algebra/src/algebra/milnor_algebra.rs index 07b3bd35f8..95fb45020c 100644 --- a/ext/crates/algebra/src/algebra/milnor_algebra.rs +++ b/ext/crates/algebra/src/algebra/milnor_algebra.rs @@ -521,6 +521,13 @@ impl Algebra for MilnorAlgebra { s_degree: i32, s: FpSlice, ) { + // At p = 2 we use the admissible-matrix algorithm, which enumerates the admissible matrices + // for the fixed operation `r` *once* and tests every term of `s` against them, instead of + // re-running the `PPartMultiplier` sweep once per term of `s`. + if !self.generic() { + self.multiply_basis_element_by_element_2(result, coeff, r_degree, r_idx, s_degree, s); + return; + } let p = self.prime(); let r = self.basis_element_from_index(r_degree, r_idx); PPartAllocation::with_local(|mut allocation| { @@ -1059,6 +1066,145 @@ 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`. + // 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)] + 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, @@ -1146,6 +1292,110 @@ 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 { + 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 +2043,84 @@ mod tests { use super::*; + /// The `p = 2` admissible-matrix multiply (`multiply_basis_element_by_element_2`, reached via + /// the `multiply_basis_element_by_element` override) 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); + + // New path: 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( + 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( + 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)] From 02ab61623f361fa63af9f452fabfcca50a89cf68 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Sat, 11 Jul 2026 02:38:55 -0400 Subject: [PATCH 03/10] Add a Nassau end-to-end timing example and a Milnor multiply profiler An S_2-at-p=2 end-to-end timing harness with an opt-in flamegraph, a compile-time-gated multiply profiler reporting per-operation and per-launch GPU-occupancy metrics, and a trim of wasted work in the p=2 multiply hot path. Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_01Fsxqsgjoa5RWYhAvRG1bT2 --- ext/Cargo.toml | 3 + ext/crates/algebra/build.rs | 13 + .../algebra/src/algebra/milnor_algebra.rs | 271 +++++++++++++++++- ext/crates/algebra/src/module/free_module.rs | 5 + ext/examples/nassau_e2e.rs | 76 +++++ 5 files changed, 366 insertions(+), 2 deletions(-) create mode 100644 ext/crates/algebra/build.rs create mode 100644 ext/examples/nassau_e2e.rs diff --git a/ext/Cargo.toml b/ext/Cargo.toml index 8b3919aba8..9e9b23fed9 100644 --- a/ext/Cargo.toml +++ b/ext/Cargo.toml @@ -33,6 +33,7 @@ serde_json = { version = "1.0.141", features = ["preserve_order"] } tracing = "0.1.41" tracing-subscriber = { version = "0.3.19", features = ["env-filter"] } +pprof = { version = "0.15.0", features = ["flamegraph"], optional = true } zstd = { version = "0.13.3", optional = true } [target.'cfg(unix)'.dependencies] @@ -59,6 +60,8 @@ concurrent = [ odd-primes = ["fp/odd-primes", "algebra/odd-primes", "sseq/odd-primes"] logging = [] nassau = [] +# Enables the optional pprof-based flamegraph in the `nassau_e2e` example (set `FLAME=`). +flamegraph = ["dep:pprof"] [workspace] members = [ diff --git a/ext/crates/algebra/build.rs b/ext/crates/algebra/build.rs new file mode 100644 index 0000000000..cf792dc271 --- /dev/null +++ b/ext/crates/algebra/build.rs @@ -0,0 +1,13 @@ +//! Turns the build-time `MILNOR_PROFILE` environment variable into `cfg(milnor_profile)`, which +//! gates the (otherwise zero-cost) Milnor-multiplication profiling counters in +//! `src/algebra/milnor_algebra.rs`. Build with e.g. `MILNOR_PROFILE=1 cargo build` to enable them. + +fn main() { + // Declare the cfg so `cfg(milnor_profile)` does not trip the `unexpected_cfgs` lint. + println!("cargo::rustc-check-cfg=cfg(milnor_profile)"); + // Rebuild when the toggle changes. + println!("cargo::rerun-if-env-changed=MILNOR_PROFILE"); + if std::env::var_os("MILNOR_PROFILE").is_some() { + println!("cargo::rustc-cfg=milnor_profile"); + } +} diff --git a/ext/crates/algebra/src/algebra/milnor_algebra.rs b/ext/crates/algebra/src/algebra/milnor_algebra.rs index 95fb45020c..cb29198fe9 100644 --- a/ext/crates/algebra/src/algebra/milnor_algebra.rs +++ b/ext/crates/algebra/src/algebra/milnor_algebra.rs @@ -11,6 +11,251 @@ use serde::{Deserialize, Serialize}; use crate::algebra::{Algebra, Bialgebra, GeneratedAlgebra, UnstableAlgebra, combinatorics}; +/// Wrap profiling-only statements. Expands to nothing unless the crate is built with the +/// `MILNOR_PROFILE` environment variable set (which turns on `cfg(milnor_profile)`; see `build.rs`), +/// so the hot-path hooks below cost *nothing* — not even argument evaluation — in normal builds. +macro_rules! profile { + ($($body:tt)*) => { + #[cfg(milnor_profile)] + { + $($body)* + } + }; +} + +/// Compile-time-gated counters for the Milnor multiplication hot path, to pinpoint what to +/// optimize (how sparse the operands are, how much matrix enumeration is wasted, how many index +/// lookups we pay, which path is taken, …). +/// +/// Enable by building with `MILNOR_PROFILE=1` (e.g. +/// `MILNOR_PROFILE=1 cargo run --release --example nassau_e2e -- 80 42 1`); otherwise every counter +/// and hook is removed by `#[cfg]`. Call [`report`] to print a summary to stderr. +pub mod profile { + #[cfg(milnor_profile)] + pub use enabled::*; + + /// Print the collected multiply statistics to stderr (or a note if profiling is disabled). + #[cfg(not(milnor_profile))] + pub fn report() { + eprintln!( + "[milnor profile] disabled — rebuild with `MILNOR_PROFILE=1` to collect multiply stats" + ); + } + + #[cfg(milnor_profile)] + mod enabled { + use std::sync::{ + LazyLock, Mutex, + atomic::{AtomicU64, Ordering::Relaxed}, + }; + + use rustc_hash::FxHashMap; + + macro_rules! counters { + ($($(#[$m:meta])* $name:ident),* $(,)?) => { + $($(#[$m])* pub static $name: AtomicU64 = AtomicU64::new(0);)* + }; + } + counters! { + /// `multiply_basis_element_by_element` (p=2) calls that do work. + MBE_CALLS, + /// … that took the admissible-matrix sweep. + MBE_ADMISSIBLE, + /// … that fell back to the per-term `PPartMultiplier` path. + MBE_PERTERM, + /// … where the operation was the identity `Sq(∅)`. + MBE_IDENTITY, + /// Basis×basis `multiply_with_allocation` invocations. + KERNEL_CALLS, + /// `PPartMultiplier` candidate terms yielded (before the index/bound checks). + PPART_TERMS, + /// `basis_element_to_index` lookups on the multiply hot path. + INDEX_LOOKUPS, + /// `add_basis_element` calls that actually contribute an output term. + OUTPUT_ADDS, + /// Admissible matrices enumerated (admissible path only). + ADM_MATRICES, + /// (matrix, term) compatibility tests (admissible path only). + ADM_TESTS, + } + + /// Histogram of `nnz(s)` — the number of non-zero terms of the element being acted on. + pub static NNZ_HIST: LazyLock>> = + LazyLock::new(|| Mutex::new(FxHashMap::default())); + /// Histogram of `(operation_degree, element_degree)` per call. + pub static DEG_HIST: LazyLock>> = + LazyLock::new(|| Mutex::new(FxHashMap::default())); + /// Total element terms multiplied by each distinct operation `R = (degree, index)`. This is + /// the GPU-occupancy metric for Nassau's `Sq(R) · Σ Sq(Sⱼ)` kernel: batching all work sharing + /// one `R` (uniform matrix enumeration) gives this many threads' worth of uniform work, so + /// the distribution says how well workgroups would fill. + pub static R_TERMS: LazyLock>> = + LazyLock::new(|| Mutex::new(FxHashMap::default())); + + pub fn record_call(nnz: usize, r_degree: i32, r_index: usize, s_degree: i32) { + MBE_CALLS.fetch_add(1, Relaxed); + *NNZ_HIST.lock().unwrap().entry(nnz).or_default() += 1; + *DEG_HIST + .lock() + .unwrap() + .entry((r_degree, s_degree)) + .or_default() += 1; + *R_TERMS + .lock() + .unwrap() + .entry((r_degree, r_index)) + .or_default() += nnz as u64; + } + + pub fn admissible() { + MBE_ADMISSIBLE.fetch_add(1, Relaxed); + } + pub fn perterm() { + MBE_PERTERM.fetch_add(1, Relaxed); + } + pub fn identity() { + MBE_IDENTITY.fetch_add(1, Relaxed); + } + pub fn kernel_call() { + KERNEL_CALLS.fetch_add(1, Relaxed); + } + pub fn ppart_term() { + PPART_TERMS.fetch_add(1, Relaxed); + } + pub fn index_lookup() { + INDEX_LOOKUPS.fetch_add(1, Relaxed); + } + pub fn output_add() { + OUTPUT_ADDS.fetch_add(1, Relaxed); + } + pub fn adm_matrix() { + ADM_MATRICES.fetch_add(1, Relaxed); + } + pub fn adm_test() { + ADM_TESTS.fetch_add(1, Relaxed); + } + + pub fn report() { + let calls = MBE_CALLS.load(Relaxed); + if calls == 0 { + eprintln!("[milnor profile] no multiply_basis_element_by_element (p=2) calls seen"); + return; + } + let (adm, per, ident) = ( + MBE_ADMISSIBLE.load(Relaxed), + MBE_PERTERM.load(Relaxed), + MBE_IDENTITY.load(Relaxed), + ); + let zero = calls - adm - per - ident; + let adds = OUTPUT_ADDS.load(Relaxed); + let pct = |x: u64| 100.0 * x as f64 / calls as f64; + let per_add = |x: u64| { + if adds > 0 { + x as f64 / adds as f64 + } else { + 0.0 + } + }; + + eprintln!("================= Milnor multiply profile ================="); + eprintln!("multiply_basis_element_by_element (p=2) calls : {calls}"); + eprintln!(" admissible-matrix path : {adm:>12} ({:5.1}%)", pct(adm)); + eprintln!(" per-term path : {per:>12} ({:5.1}%)", pct(per)); + eprintln!( + " identity Sq(∅) : {ident:>12} ({:5.1}%)", + pct(ident) + ); + eprintln!(" zero element (nnz=0) : {zero:>12} ({:5.1}%)", pct(zero)); + + let hist = NNZ_HIST.lock().unwrap(); + let n: u64 = hist.values().sum(); + let weighted: u64 = hist.iter().map(|(k, v)| *k as u64 * v).sum(); + eprintln!( + "element term-count nnz : mean {:.2} over {n} calls", + weighted as f64 / n.max(1) as f64 + ); + let mut keys: Vec = hist.keys().copied().collect(); + keys.sort_unstable(); + let mut cum = 0u64; + for k in keys { + let v = hist[&k]; + cum += v; + if k <= 6 || k.is_multiple_of(10) { + eprintln!( + " nnz={k:<3} {:>6.2}% cum {:>6.2}%", + 100.0 * v as f64 / n as f64, + 100.0 * cum as f64 / n as f64 + ); + } + } + + eprintln!( + "kernel basis×basis multiply_with_allocation : {}", + KERNEL_CALLS.load(Relaxed) + ); + eprintln!("output basis-element adds : {adds}"); + eprintln!( + "PPartMultiplier candidate terms : {} ({:.2} per output add)", + PPART_TERMS.load(Relaxed), + per_add(PPART_TERMS.load(Relaxed)) + ); + eprintln!( + "basis_element_to_index lookups : {} ({:.2} per output add)", + INDEX_LOOKUPS.load(Relaxed), + per_add(INDEX_LOOKUPS.load(Relaxed)) + ); + let matrices = ADM_MATRICES.load(Relaxed); + if matrices > 0 { + eprintln!( + "admissible matrices enumerated : {matrices} ({:.1} term-tests each)", + ADM_TESTS.load(Relaxed) as f64 / matrices as f64 + ); + } + + let deg = DEG_HIST.lock().unwrap(); + let mut pairs: Vec<((i32, i32), u64)> = deg.iter().map(|(k, v)| (*k, *v)).collect(); + pairs.sort_by_key(|(_, v)| std::cmp::Reverse(*v)); + eprintln!("top (operation_degree, element_degree) call sites:"); + for ((r, s), v) in pairs.iter().take(10) { + eprintln!( + " op={r:<3} elt={s:<3} : {v} ({:.1}%)", + 100.0 * *v as f64 / calls as f64 + ); + } + + // GPU-occupancy view: group all element terms by the operation `R` they are multiplied + // by. Nassau's kernel (`Sq(R) · Σ Sq(Sⱼ)`) parallelizes over terms sharing one `R` with + // uniform matrix enumeration, so a workgroup of size `W` is well-filled only by `R`s that + // accumulate ≥ `W` terms. We report what fraction of all term-work lives in such `R`s. + let r_terms = R_TERMS.lock().unwrap(); + let distinct = r_terms.len() as u64; + let total_terms: u64 = r_terms.values().sum(); + let max_terms = r_terms.values().copied().max().unwrap_or(0); + eprintln!( + "GPU occupancy — distinct operations R: {distinct}, total element terms: \ + {total_terms}, mean terms/R: {:.1}, max: {max_terms}", + total_terms as f64 / distinct.max(1) as f64 + ); + eprintln!(" fraction of term-work in R's with ≥ W terms (W = workgroup size):"); + for w in [32u64, 64, 128, 256, 512, 1024] { + let (n_r, work): (u64, u64) = + r_terms.values().fold( + (0, 0), + |(n, s), &t| { + if t >= w { (n + 1, s + t) } else { (n, s) } + }, + ); + eprintln!( + " W={w:<4} : {n_r:>6} R's ({:>5.1}% of R's) cover {:>5.1}% of term-work", + 100.0 * n_r as f64 / distinct.max(1) as f64, + 100.0 * work as f64 / total_terms.max(1) as f64 + ); + } + eprintln!("==========================================================="); + } + } +} + fn q_part_default() -> u32 { !0 } @@ -1096,11 +1341,18 @@ impl MilnorAlgebra { if coeff.is_multiple_of(2) { return; } + profile!(profile::record_call( + s.iter_nonzero().count(), + r_degree, + r_idx, + s_degree + )); 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() { + profile!(profile::identity()); for (i, _) in s.iter_nonzero() { result.add_basis_element(i, 1); } @@ -1118,6 +1370,7 @@ impl MilnorAlgebra { return; // s = 0 }; let Some((i1, _)) = second else { + profile!(profile::perterm()); PPartAllocation::with_local(|allocation| { self.multiply_with_allocation( result, @@ -1137,6 +1390,7 @@ impl MilnorAlgebra { 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))); + profile!(profile::admissible()); let out_degree = r_degree + s_degree; let mut matrix = AdmissibleMatrix::new(&r.p_part); @@ -1147,7 +1401,9 @@ impl MilnorAlgebra { }; loop { + profile!(profile::adm_matrix()); 'outer: for term in &terms { + profile!(profile::adm_test()); let basis = &term.p_part; working.p_part.clear(); @@ -1197,7 +1453,9 @@ impl MilnorAlgebra { } let idx = self.basis_element_to_index(&working); + profile!(profile::index_lookup()); result.add_basis_element(idx, 1); + profile!(profile::output_add()); } if !matrix.next() { break; @@ -1215,6 +1473,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 { @@ -1229,13 +1492,14 @@ 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); } } allocation = multiplier.into_allocation() } } else { + profile!(profile::kernel_call()); let mut multiplier = PPartMultiplier::::new_from_allocation( self.prime(), &m1.p_part, @@ -1246,9 +1510,12 @@ impl MilnorAlgebra { ); while let Some(c) = multiplier.next() { + profile!(profile::ppart_term()); let idx = self.basis_element_to_index(&multiplier.ans); - if idx < self.dimension_unstable(target_deg, excess) { + profile!(profile::index_lookup()); + if idx < dim { res.add_basis_element(idx, c * coef); + profile!(profile::output_add()); } } allocation = multiplier.into_allocation() 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..388636374f --- /dev/null +++ b/ext/examples/nassau_e2e.rs @@ -0,0 +1,76 @@ +//! End-to-end timing harness for Nassau's algorithm on the sphere `S_2` at p = 2. +//! +//! Resolves `S_2` via [`construct_nassau`] and [`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). +//! +//! To see *where* the time goes (e.g. multiply vs. linear algebra), build with the `flamegraph` +//! feature and set `FLAME` to an output path — a sampling flamegraph of every run is written there: +//! ```text +//! FLAME=/tmp/nassau.svg cargo run --release --features flamegraph --example nassau_e2e -- 60 32 1 +//! ``` + +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); + + // When built `--features flamegraph` and `FLAME=` is set, sample the whole run and write a + // flamegraph there. Kept alive until after the timing loop so it captures every resolution. + #[cfg(feature = "flamegraph")] + let flame = std::env::var("FLAME").ok().map(|path| { + ( + path, + pprof::ProfilerGuard::new(1000).expect("failed to start profiler"), + ) + }); + + 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"); + + // Prints Milnor-multiply counters when built with `MILNOR_PROFILE=1`, else a one-line note. + algebra::milnor_algebra::profile::report(); + + #[cfg(feature = "flamegraph")] + if let Some((path, guard)) = flame { + let report = guard + .report() + .build() + .expect("failed to build profiler report"); + let file = std::fs::File::create(&path).expect("failed to create flamegraph file"); + report.flamegraph(file).expect("failed to write flamegraph"); + eprintln!("wrote flamegraph to {path}"); + } +} From ea2de1bcbf2c2a9f121b612b6d4cf4c8bfd9a0e1 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Sat, 11 Jul 2026 02:38:55 -0400 Subject: [PATCH 04/10] Add hash-free seqno index tables for the Milnor multiply Flat, arc-swapped seqno tables (the output-indexing primitive the GPU multiply needs) with a hashmap A/B bench. Moves the admissible-matrix multiply off the CPU path, keeping it as the GPU model. Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_01Fsxqsgjoa5RWYhAvRG1bT2 --- ext/crates/algebra/Cargo.toml | 5 + ext/crates/algebra/benches/seqno.rs | 70 ++++ .../algebra/src/algebra/milnor_algebra.rs | 378 +++++++++++++++++- .../algebra/src/module/homomorphism/mod.rs | 20 + 4 files changed, 458 insertions(+), 15 deletions(-) create mode 100644 ext/crates/algebra/benches/seqno.rs diff --git a/ext/crates/algebra/Cargo.toml b/ext/crates/algebra/Cargo.toml index 3b98825cb5..aded31fc16 100644 --- a/ext/crates/algebra/Cargo.toml +++ b/ext/crates/algebra/Cargo.toml @@ -17,6 +17,7 @@ maybe-rayon = { path = "../maybe-rayon" } once = { path = "../once" } anyhow = "1.0.98" +arc-swap = "1.7.1" auto_impl = "1.3.0" hashbrown = "0.15.4" itertools = { version = "0.14.0", default-features = false, features = [ @@ -55,3 +56,7 @@ harness = false [[bench]] name = "nassau_milnor" harness = false + +[[bench]] +name = "seqno" +harness = false diff --git a/ext/crates/algebra/benches/seqno.rs b/ext/crates/algebra/benches/seqno.rs new file mode 100644 index 0000000000..22ad312edf --- /dev/null +++ b/ext/crates/algebra/benches/seqno.rs @@ -0,0 +1,70 @@ +//! A/B benchmark: the hash-free table-based Milnor index ([`MilnorAlgebra::seqno`]) versus the +//! `FxHashMap`-backed [`MilnorAlgebra::basis_element_to_index`], over the exact operation of +//! turning a basis element back into its index. +//! +//! This is the lookup that a resolution performs on *every* output term of a multiply — Nassau's +//! `S_2` @ p=2 run issues ~1.3B of them — so a lookup that is even a little cheaper is worth having, +//! and a GPU kernel (which cannot carry a hashmap) *needs* the table-based form regardless. +//! +//! The `seqno` group here reads the flat [`arc_swap`]-backed table; compare its `basis_to_index/*` +//! ids against the `hashmap/*` ids at the same degree to see which wins on the CPU. + +use std::hint::black_box; + +use algebra::{Algebra, MilnorAlgebra}; +use criterion::{Criterion, Throughput, criterion_group, criterion_main}; +use fp::prime::TWO; +use pprof::criterion::{Output, PProfProfiler}; + +/// Degrees to sweep. Chosen to bracket the range a real `S_2` resolution reaches, from the cheap +/// low-dimensional end up to degrees whose basis has thousands of elements. +const DEGREES: &[i32] = &[16, 24, 32, 40, 48, 56, 64]; + +fn seqno(c: &mut Criterion) { + let algebra = MilnorAlgebra::new(TWO, false); + let max_degree = *DEGREES.iter().max().unwrap(); + algebra.compute_basis(max_degree); + algebra.compute_seqno_tables(max_degree); + + let mut g = c.benchmark_group("seqno"); + + for °ree in DEGREES { + let dim = algebra.dimension(degree); + if dim == 0 { + continue; + } + // Snapshot the basis so neither method pays to walk the algebra's storage during timing. + let basis: Vec<_> = (0..dim) + .map(|i| algebra.basis_element_from_index(degree, i).clone()) + .collect(); + + g.throughput(Throughput::Elements(dim as u64)); + + g.bench_function(format!("hashmap/deg{degree}"), |b| { + b.iter(|| { + for elt in &basis { + black_box(algebra.basis_element_to_index(elt)); + } + }); + }); + + g.bench_function(format!("basis_to_index/deg{degree}"), |b| { + b.iter(|| { + for elt in &basis { + black_box(algebra.seqno(&elt.p_part)); + } + }); + }); + } + + g.finish(); +} + +criterion_group! { + name = benches; + config = Criterion::default() + .measurement_time(std::time::Duration::from_secs(3)) + .with_profiler(PProfProfiler::new(100, Output::Flamegraph(None))); + targets = seqno +} +criterion_main!(benches); diff --git a/ext/crates/algebra/src/algebra/milnor_algebra.rs b/ext/crates/algebra/src/algebra/milnor_algebra.rs index cb29198fe9..d788f85e8b 100644 --- a/ext/crates/algebra/src/algebra/milnor_algebra.rs +++ b/ext/crates/algebra/src/algebra/milnor_algebra.rs @@ -107,6 +107,115 @@ pub mod profile { .or_default() += nnz as u64; } + /// Per–GPU-launch accumulation of `R → element-terms`, where one launch = one + /// `get_partial_matrix` (matrix build). Unlike [`R_TERMS`], which sums an operation `R` + /// across the *whole* resolution, this asks how much same-`R` work is co-located within a + /// single matrix build — the largest unit a kernel could batch without buffering across the + /// streaming algorithm. `depth` keeps [`scope_begin`]/[`scope_end`] reentrancy-safe: only the + /// outermost pair clears and folds, so any nested matrix build is attributed to the outer + /// launch rather than corrupting it. (Measurement must be run serially — the `concurrent` + /// feature would let two launches interleave into one scope and inflate the batch sizes.) + struct ScopeState { + depth: u32, + /// `(homomorphism id, degree)` of the currently open launch — its merged-scope key. + key: (usize, i32), + map: FxHashMap<(i32, usize), u64>, + } + static SCOPE: LazyLock> = LazyLock::new(|| { + Mutex::new(ScopeState { + depth: 0, + key: (0, 0), + map: FxHashMap::default(), + }) + }); + + /// Coarser accumulation: `R → element-terms` merged across every launch sharing a + /// `(homomorphism id, degree)` key — i.e. all the per-signature `get_partial_matrix` builds + /// of one differential at one bidegree, which all read the same matrix and so *could* be + /// fused into a single kernel launch (computing the full matrix once and slicing it). Held to + /// report time and folded there, to measure how much the realizable occupancy improves when + /// the launch scope is widened from one masked build to one whole bidegree. + #[allow(clippy::type_complexity)] + static MERGED: LazyLock>>> = + LazyLock::new(|| Mutex::new(FxHashMap::default())); + + /// Workgroup sizes for the realizable-coverage report (mirrors the global one). + const SCOPE_WS: [u64; 6] = [32, 64, 128, 256, 512, 1024]; + /// Number of launch scopes (matrix builds) that did any work. + static SCOPE_COUNT: AtomicU64 = AtomicU64::new(0); + /// Sum over scopes of each scope's total element-terms (equals the global term-work). + static SCOPE_TERMS: AtomicU64 = AtomicU64::new(0); + /// Largest single-scope total element-terms. + static SCOPE_TERMS_MAX: AtomicU64 = AtomicU64::new(0); + /// Per `W`, term-work in `R`s that reach ≥ `W` terms *within their own launch*. + static SCOPE_COVER: [AtomicU64; 6] = [ + AtomicU64::new(0), + AtomicU64::new(0), + AtomicU64::new(0), + AtomicU64::new(0), + AtomicU64::new(0), + AtomicU64::new(0), + ]; + + /// Open a launch scope around a matrix build of homomorphism `hom_id` at `degree`. + /// Reentrancy-safe (see [`ScopeState`]); the key is set by the outermost open. + pub fn scope_begin(hom_id: usize, degree: i32) { + let mut s = SCOPE.lock().unwrap(); + if s.depth == 0 { + s.map.clear(); + s.key = (hom_id, degree); + } + s.depth += 1; + } + + /// Attribute `nnz` element-terms of operation `R = (r_degree, r_index)` to the open launch, + /// both to its per-launch histogram and to its `(hom, degree)` merged bucket. + pub fn scope_record(r_degree: i32, r_index: usize, nnz: usize) { + if nnz == 0 { + return; + } + let key = { + let mut s = SCOPE.lock().unwrap(); + if s.depth == 0 { + return; + } + *s.map.entry((r_degree, r_index)).or_default() += nnz as u64; + s.key + }; + *MERGED + .lock() + .unwrap() + .entry(key) + .or_default() + .entry((r_degree, r_index)) + .or_default() += nnz as u64; + } + + /// Close a launch scope; the outermost close folds the scope's `R`-histogram into the + /// realizable-coverage accumulators. + pub fn scope_end() { + let mut s = SCOPE.lock().unwrap(); + if s.depth == 0 { + return; + } + s.depth -= 1; + if s.depth > 0 { + return; + } + let total: u64 = s.map.values().sum(); + if total == 0 { + return; + } + SCOPE_COUNT.fetch_add(1, Relaxed); + SCOPE_TERMS.fetch_add(total, Relaxed); + SCOPE_TERMS_MAX.fetch_max(total, Relaxed); + for (i, &w) in SCOPE_WS.iter().enumerate() { + let work: u64 = s.map.values().filter(|&&t| t >= w).sum(); + SCOPE_COVER[i].fetch_add(work, Relaxed); + } + s.map.clear(); + } + pub fn admissible() { MBE_ADMISSIBLE.fetch_add(1, Relaxed); } @@ -251,6 +360,64 @@ pub mod profile { 100.0 * work as f64 / total_terms.max(1) as f64 ); } + + // The number above aggregates each R across the *whole* resolution. A kernel can only + // batch work that is co-located in one launch, so this is the realizable version: R-work + // is re-counted per `get_partial_matrix` (matrix build), the largest unit batchable + // without buffering across the streaming algorithm. If these percentages collapse + // relative to the global ones, the amortization-by-R shape does not survive at launch + // granularity, and a GPU kernel must lean on raw parallel width (many independent + // products per launch) rather than on many terms sharing one R. + let scopes = SCOPE_COUNT.load(Relaxed); + if scopes > 0 { + let scope_terms = SCOPE_TERMS.load(Relaxed); + eprintln!( + "GPU occupancy — REALIZABLE per matrix-build launch ({scopes} launches, mean \ + {:.1} terms/launch, max {}):", + scope_terms as f64 / scopes as f64, + SCOPE_TERMS_MAX.load(Relaxed) + ); + eprintln!(" fraction of term-work in R's reaching ≥ W terms *within one launch*:"); + for (i, w) in SCOPE_WS.iter().enumerate() { + let work = SCOPE_COVER[i].load(Relaxed); + eprintln!( + " W={w:<4} : {:>5.1}% of term-work", + 100.0 * work as f64 / scope_terms.max(1) as f64 + ); + } + } + + // Coarser scope: merge the per-signature builds of one differential at one bidegree into + // a single launch (they read the same matrix, so fusing them is free of extra multiply + // work). This upper-bounds the realizable occupancy for a kernel that computes the full + // bidegree matrix once instead of one masked slice per signature. + let merged = MERGED.lock().unwrap(); + if !merged.is_empty() { + let n_launches = merged.len() as u64; + let mut total = 0u64; + let mut max_terms = 0u64; + let mut cover = [0u64; 6]; + for hist in merged.values() { + let t: u64 = hist.values().sum(); + total += t; + max_terms = max_terms.max(t); + for (i, &w) in SCOPE_WS.iter().enumerate() { + cover[i] += hist.values().filter(|&&x| x >= w).sum::(); + } + } + eprintln!( + "GPU occupancy — MERGED per (differential, bidegree) launch ({n_launches} \ + launches, mean {:.1} terms/launch, max {max_terms}):", + total as f64 / n_launches.max(1) as f64 + ); + eprintln!(" fraction of term-work in R's reaching ≥ W terms *within one launch*:"); + for (i, w) in SCOPE_WS.iter().enumerate() { + eprintln!( + " W={w:<4} : {:>5.1}% of term-work", + 100.0 * cover[i] as f64 / total.max(1) as f64 + ); + } + } eprintln!("==========================================================="); } } @@ -490,6 +657,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, @@ -510,6 +687,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>>>, @@ -538,6 +724,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(), } @@ -573,6 +760,14 @@ 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 it loses + // to this hashmap on the CPU. Even after moving its tables to flat, contiguous + // `arc_swap`-backed storage (removing the earlier `OnceVec` per-access atomics), the + // `benches/seqno` A/B still measures raw lookups at ~50 Melem/s for `seqno` vs ~115 Melem/s + // for this hashmap — a ~2.3× gap that is flat across degree: computing the rank (a degree + // sum plus two indexed table reads per populated ξ-position) is simply more work than one + // hash and probe. `seqno` is therefore 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() @@ -670,6 +865,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| { @@ -766,13 +965,25 @@ impl Algebra for MilnorAlgebra { s_degree: i32, s: FpSlice, ) { - // At p = 2 we use the admissible-matrix algorithm, which enumerates the admissible matrices - // for the fixed operation `r` *once* and tests every term of `s` against them, instead of - // re-running the `PPartMultiplier` sweep once per term of `s`. - if !self.generic() { - self.multiply_basis_element_by_element_2(result, coeff, r_degree, r_idx, s_degree, s); - return; - } + // 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 end-to-end + // A/Bs of Nassau's `S_2` regime measured it a consistent net regression on the CPU (~8% at + // stem 80, ~3% at stem 100): the regime is dominated by sparse elements (≈31% single-term), + // for which the up-front enumeration cannot be amortized. It is retained as the reference + // model for a future GPU kernel (where the enumerate-once/test-all-terms shape is ideal), + // not wired here. See the commit history for the measurements. + profile!({ + let nnz = s.iter_nonzero().count(); + profile::record_call(nnz, r_degree, r_idx, s_degree); + if r_degree == 0 { + profile::identity(); + } else if nnz > 0 { + profile::perterm(); + } + profile::scope_record(r_degree, r_idx, nnz); + }); let p = self.prime(); let r = self.basis_element_from_index(r_degree, r_idx); PPartAllocation::with_local(|mut allocation| { @@ -1136,6 +1347,111 @@ 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 [`Self::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() { + if 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; + } + } + + self.seqno_tables + .store(Some(std::sync::Arc::new(SeqnoTables { + max_degree, + width, + g, + }))); + } + + /// 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 + /// [`Self::seqno_applicable`] and that `p_part` is a genuine basis element (trimmed, in range). + /// + /// The basis is enumerated by increasing highest ξ-index, so the rank of `P` accumulates, for + /// each populated position `h`, the number of basis elements whose highest index is `< h` + /// together with `h` — which is exactly the `g_table` difference across the degree consumed at + /// that position. + pub fn seqno(&self, p_part: &[PPartEntry]) -> usize { + let xi = combinatorics::xi_degrees(self.prime()); + let guard = self.seqno_tables.load(); + let t = guard + .as_ref() + .expect("seqno tables not built; call compute_seqno_tables first"); + let w = t.width; + let mut cur_d: i32 = p_part.iter().zip(xi).map(|(&r, &x)| r as i32 * x).sum(); + let mut rank = 0; + // Consume positions from the highest down; position 0 contributes nothing. + for h in (1..p_part.len()).rev() { + let r = p_part[h] as i32; + if r == 0 { + continue; + } + let below = cur_d - r * xi[h]; + rank += t.g[cur_d as usize * w + h] - t.g[below as usize * w + h]; + cur_d = below; + } + rank + } + fn generate_basis_generic(&self, max_degree: i32) { let q = 2 * self.prime() - 2; let tau_degrees = combinatorics::tau_degrees(self.prime()); @@ -1320,10 +1636,18 @@ impl MilnorAlgebra { /// 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.** End-to-end A/Bs of Nassau's `S_2` regime measured this a net + /// regression versus the per-term sweep in [`Self::multiply_basis_element_by_element`] (the + /// regime is too sparse for the up-front enumeration to pay off). It is kept as the reference + /// model for a future GPU kernel: enumerate `Sq(R)`'s matrices once per operation and test every + /// element term against them in parallel — a shape that batches extremely well on a GPU (a real + /// resolution presents tens of thousands of terms per distinct `R`). 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)] - fn multiply_basis_element_by_element_2( + pub fn multiply_basis_element_by_element_2( &self, mut result: FpSliceMut, coeff: u32, @@ -2310,10 +2634,34 @@ mod tests { use super::*; - /// The `p = 2` admissible-matrix multiply (`multiply_basis_element_by_element_2`, reached via - /// the `multiply_basis_element_by_element` override) 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. + /// The table-based [`MilnorAlgebra::seqno`] must return the position of every basis element in + /// its degree — i.e. agree with the enumeration order that defines the index — for the stable + /// `p = 2` full algebra, and reject non-basis elements via `try_`. + #[test] + fn seqno_matches_enumeration_order() { + let algebra = MilnorAlgebra::new(ValidPrime::new(2), false); + assert!(algebra.seqno_applicable()); + let max_degree = 100; + algebra.compute_basis(max_degree); + algebra.compute_seqno_tables(max_degree); + for d in 0..=max_degree { + let dim = algebra.dimension(d); + for i in 0..dim { + let elt = algebra.basis_element_from_index(d, i); + assert_eq!( + algebra.seqno(&elt.p_part), + i, + "seqno mismatch at degree {d}, index {i}: {elt:?}" + ); + } + } + } + + /// 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); @@ -2344,11 +2692,11 @@ mod tests { ); expected_dense.add(&expected, 1); - // New path: element `s = e_j`. + // 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( + algebra.multiply_basis_element_by_element_2( got.as_slice_mut(), 1, r_degree, @@ -2370,7 +2718,7 @@ mod tests { s.set_entry(j, 1); } let mut got = FpVector::new(p, out_dim); - algebra.multiply_basis_element_by_element( + algebra.multiply_basis_element_by_element_2( got.as_slice_mut(), 1, r_degree, diff --git a/ext/crates/algebra/src/module/homomorphism/mod.rs b/ext/crates/algebra/src/module/homomorphism/mod.rs index 540b34ac4b..b6d3b08d89 100644 --- a/ext/crates/algebra/src/module/homomorphism/mod.rs +++ b/ext/crates/algebra/src/module/homomorphism/mod.rs @@ -130,10 +130,20 @@ pub trait ModuleHomomorphism: Send + Sync { return; } + // One matrix build is the largest unit a GPU kernel could batch without buffering across + // the streaming algorithm; mark it as a launch scope for the realizable-occupancy metric. + // The `(self pointer, degree)` key groups builds of the same differential at one bidegree. + #[cfg(milnor_profile)] + crate::algebra::milnor_algebra::profile::scope_begin( + std::ptr::from_ref(self) as *const () as usize, + degree, + ); matrix .maybe_par_iter_mut() .enumerate() .for_each(|(i, row)| self.apply_to_basis_element(row, 1, degree, i)); + #[cfg(milnor_profile)] + crate::algebra::milnor_algebra::profile::scope_end(); } /// Get the values of the homomorphism on the specified inputs to `matrix`. @@ -144,10 +154,20 @@ pub trait ModuleHomomorphism: Send + Sync { return matrix; } + // One matrix build is the largest unit a GPU kernel could batch without buffering across + // the streaming algorithm; mark it as a launch scope for the realizable-occupancy metric. + // The `(self pointer, degree)` key groups builds of the same differential at one bidegree. + #[cfg(milnor_profile)] + crate::algebra::milnor_algebra::profile::scope_begin( + std::ptr::from_ref(self) as *const () as usize, + degree, + ); matrix .maybe_par_iter_mut() .enumerate() .for_each(|(i, row)| self.apply_to_basis_element(row, 1, degree, inputs[i])); + #[cfg(milnor_profile)] + crate::algebra::milnor_algebra::profile::scope_end(); matrix } From bdc2223c7f9af84eccdfd8182dc63e96d94781b6 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Sat, 11 Jul 2026 02:38:55 -0400 Subject: [PATCH 05/10] Add CubeCL GPU offload for Nassau's Milnor multiply Staged CubeCL kernels (F2 xor, on-device seqno, single-R and batched multiply) wired into Nassau's get_partial_matrix, with u16-packed admissible/term transfers and parallel host marshalling. Gated behind the gpu feature; cargo check/build need no CUDA toolkit. Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_01Fsxqsgjoa5RWYhAvRG1bT2 --- ext/Cargo.toml | 8 + ext/benches/milnor_gpu_ab.rs | 212 ++++ ext/crates/algebra/Cargo.toml | 9 + .../algebra/src/algebra/milnor_algebra.rs | 55 + ext/crates/algebra/src/algebra/milnor_gpu.rs | 1028 +++++++++++++++++ ext/crates/algebra/src/algebra/mod.rs | 3 + ext/flake.nix | 27 + ext/src/lib.rs | 2 + ext/src/nassau.rs | 41 +- ext/src/nassau_gpu.rs | 179 +++ ext/tests/nassau_gpu.rs | 50 + ext/tests/nassau_gpu_timing.rs | 49 + 12 files changed, 1660 insertions(+), 3 deletions(-) create mode 100644 ext/benches/milnor_gpu_ab.rs create mode 100644 ext/crates/algebra/src/algebra/milnor_gpu.rs create mode 100644 ext/src/nassau_gpu.rs create mode 100644 ext/tests/nassau_gpu.rs create mode 100644 ext/tests/nassau_gpu_timing.rs diff --git a/ext/Cargo.toml b/ext/Cargo.toml index 9e9b23fed9..6aa52ef8ec 100644 --- a/ext/Cargo.toml +++ b/ext/Cargo.toml @@ -62,6 +62,8 @@ logging = [] nassau = [] # Enables the optional pprof-based flamegraph in the `nassau_e2e` example (set `FLAME=`). flamegraph = ["dep:pprof"] +# Enables the GPU Milnor-multiply backend and the `milnor_gpu_ab` A/B microbenchmark. +gpu = ["algebra/gpu"] [workspace] members = [ @@ -89,3 +91,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..97faa2678e --- /dev/null +++ b/ext/benches/milnor_gpu_ab.rs @@ -0,0 +1,212 @@ +//! A/B microbenchmark: GPU batch Milnor-multiply vs the CPU `get_partial_matrix`. +//! +//! Answers the Stage 5 go/no-go question from `GPU_KERNEL_HANDOFF.md` — does offloading +//! the Milnor multiply to the GPU beat the CPU per-term path *including* host↔device +//! transfer — before committing to the invasive wire-in. +//! +//! 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 aded31fc16..ccb6eda417 100644 --- a/ext/crates/algebra/Cargo.toml +++ b/ext/crates/algebra/Cargo.toml @@ -29,6 +29,14 @@ 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", +] } + [dev-dependencies] criterion = { version = "0.5", features = ["html_reports"] } expect-test = "1.5.1" @@ -39,6 +47,7 @@ rstest = "0.25.0" default = ["odd-primes"] cache-multiplication = [] concurrent = ["fp/concurrent", "maybe-rayon/concurrent"] +gpu = ["dep:cubecl"] odd-primes = ["fp/odd-primes"] [[bench]] diff --git a/ext/crates/algebra/src/algebra/milnor_algebra.rs b/ext/crates/algebra/src/algebra/milnor_algebra.rs index d788f85e8b..dfcbe2460b 100644 --- a/ext/crates/algebra/src/algebra/milnor_algebra.rs +++ b/ext/crates/algebra/src/algebra/milnor_algebra.rs @@ -1452,6 +1452,61 @@ impl MilnorAlgebra { 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 [`Self::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()); 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..a4db426e06 --- /dev/null +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -0,0 +1,1028 @@ +//! GPU offload for the Milnor multiply at `p = 2`, built on [CubeCL]. +//! +//! This is the entry point for the staged port described in +//! `crates/algebra/GPU_KERNEL_HANDOFF.md`: the admissible-matrix multiply +//! ([`super::milnor_algebra::MilnorAlgebra::multiply_basis_element_by_element_2`]) +//! and the hash-free `seqno` index run as CubeCL kernels, batched per +//! `get_partial_matrix` launch. +//! +//! Stages landed so far: +//! - **Stage 1 — toolchain slice.** A trivial F₂ kernel ([`xor_f2`]) proving the +//! CubeCL `cuda` runtime works end to end. F₂ addition is XOR of the bit-packed +//! limbs, the exact accumulation every multiply kernel performs. +//! - **Stage 2 — `seqno` on device.** [`seqno_core`]/[`seqno_kernel`] port the +//! hash-free index [`MilnorAlgebra::seqno`] as pure integer arithmetic over the +//! flat `g` table — the output-indexing primitive the multiply kernel needs. +//! - **Stage 3 — single-`R` multiply.** [`multiply_pair`] ports the per-term test +//! + output assembly of `multiply_basis_element_by_element_2`; +//! [`multiply_single_r_kernel`] runs one `Sq(R)·s` product per launch. +//! - **Stage 4 — batched launch.** [`multiply_batch_kernel`] fuses all `(R, s)` +//! products of one `get_partial_matrix` into a single launch (one thread per +//! `(product, matrix, term)` pair). The pair is decoded on-device from a prefix-sum +//! over per-product pair counts, with admissible-matrix data deduplicated by distinct +//! `R` — so a launch uploads compact per-`R`/per-product tables, not a per-pair table +//! (which at scale is gigabytes of almost-entirely-redundant data). On an RTX 3050 Ti +//! this made the batch ~2× faster than the CPU `get_partial_matrix` at stem 100 +//! (`benches/milnor_gpu_ab.rs`); the naive per-pair table was ~7× slower. +//! +//! 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 crate::algebra::{ + Algebra, MilnorAlgebra, + combinatorics::{MAX_XI_TAU, 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; + +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), + ) +} + +/// 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. +#[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. +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(n * size_of::()); + + // 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). +#[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`. +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`]. +#[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`. +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| v as u16).collect(); + let masks: Vec = mk32.iter().map(|&v| v as u16).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 = v as u16; + } + } + + 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 +/// Stage 4 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); + + // Optional host-marshal vs device timing split (set MILNOR_GPU_TIMING). + let timing = std::env::var("MILNOR_GPU_TIMING").is_ok(); + 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); + } + + // Parallel: each distinct `R`'s admissible matrices. Values ship as u16 (p-part + // magnitudes ≤ the output degree, far below 65535) to halve the transfer. + let r_mats: Vec<(usize, usize, Vec, Vec)> = (0..distinct_r.len()) + .into_maybe_par_iter() + .map(|i| { + let (rd, ridx) = distinct_r[i]; + let r = algebra.basis_element_from_index(rd, ridx); + assert!(!r.p_part.is_empty(), "each R must be non-empty"); + let (cs_len, mk_len, cs, mk) = algebra.admissible_matrices(&r.p_part); + let cs16 = cs.iter().map(|&v| v as u16).collect(); + let mk16 = mk.iter().map(|&v| v as u16).collect(); + (cs_len, mk_len, cs16, mk16) + }) + .collect(); + + // Lay out per-`R` matrix data (sequential concat + offsets). + let mut col_sums: Vec = Vec::new(); + let mut masks: Vec = Vec::new(); + let mut r_cs_offset: Vec = Vec::with_capacity(r_mats.len()); + let mut r_mk_offset: Vec = Vec::with_capacity(r_mats.len()); + let mut r_cs_len: Vec = Vec::with_capacity(r_mats.len()); + let mut r_mk_len: Vec = Vec::with_capacity(r_mats.len()); + let mut r_num_matrices: Vec = Vec::with_capacity(r_mats.len()); + for (cs_len, mk_len, cs, mk) in &r_mats { + r_cs_offset.push(col_sums.len() as u32); + r_mk_offset.push(masks.len() as u32); + r_cs_len.push(*cs_len as u32); + r_mk_len.push(*mk_len as u32); + r_num_matrices.push(mk.len() / mk_len); + col_sums.extend_from_slice(cs); + masks.extend_from_slice(mk); + } + + // 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 = v as u16; + } + } + (tp, tl) + }) + .collect(); + + // 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]; + } + + // Device buffers must be non-empty even when a dimension is zero for every product. + if col_sums.is_empty() { + col_sums.push(0); + } + if masks.is_empty() { + masks.push(0); + } + if term_pparts.is_empty() { + term_pparts.push(0); + } + + let marshal_ms = t_marshal.elapsed().as_secs_f64() * 1e3; + + if timing { + let mb = |n: usize| (n * 4) as f64 / (1024.0 * 1024.0); + let mb16 = |n: usize| (n * 2) as f64 / (1024.0 * 1024.0); + eprintln!( + " [sizes] {} products, {} distinct R | upload(u16): col_sums {:.1}MB masks \ + {:.1}MB term_pparts {:.1}MB per-product {:.1}MB (was ~{:.0}MB pair table)", + products.len(), + r_cs_len.len(), + mb16(col_sums.len()), + mb16(masks.len()), + mb16(term_pparts.len()), + mb(term_lens.len() + prod_r_index.len() * 4 + prod_pair_start.len()), + mb(total_pairs * 7), + ); + } + + let t_device = std::time::Instant::now(); + + let client = CudaRuntime::client(&CudaDevice::default()); + let t_upload = std::time::Instant::now(); + 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 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)); + // Barrier so the upload timer excludes kernel/readback (timing runs only). + if timing { + let _ = cubecl::future::block_on(client.sync()); + } + let upload_ms = t_upload.elapsed().as_secs_f64() * 1e3; + + let t_kernel = std::time::Instant::now(); + 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, 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(), 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, + ); + } + + // Barrier so the kernel timer excludes readback (timing runs only). + if timing { + let _ = cubecl::future::block_on(client.sync()); + } + let kernel_ms = t_kernel.elapsed().as_secs_f64() * 1e3; + + let t_read = std::time::Instant::now(); + let bytes = client.read_one(out_h).unwrap(); + let read_ms = t_read.elapsed().as_secs_f64() * 1e3; + let flat = u32::from_bytes(&bytes); + let result = (0..num_rows) + .map(|r| flat[r * num_limbs..(r + 1) * num_limbs].to_vec()) + .collect(); + + // 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); + + if timing { + eprintln!( + " [gpu] {total_pairs} pairs: marshal {marshal_ms:.1} upload {upload_ms:.1} kernel \ + {kernel_ms:.1} readback {read_ms:.1} device {device_ms:.1} (ms)" + ); + } + 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/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..e929b9fd13 100644 --- a/ext/src/nassau.rs +++ b/ext/src/nassau.rs @@ -388,6 +388,42 @@ 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) +} + /// A resolution of `S_2` using Nassau's algorithm. /// /// This aims to have an API similar to that of @@ -618,7 +654,7 @@ impl> Resolution { let full_matrix = { let _guard = ParallelGuard::new(); - self.differentials[b.s() - 1].get_partial_matrix(b.t(), &target_mask) + 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]); @@ -687,8 +723,7 @@ impl> Resolution { let full_matrix = { let _guard = ParallelGuard::new(); - self.differential(b.s() - 1) - .get_partial_matrix(b.t(), &target_mask) + 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..c63e0462dd --- /dev/null +++ b/ext/src/nassau_gpu.rs @@ -0,0 +1,179 @@ +//! 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`], which on an RTX 3050 Ti beats the +//! parallel CPU per-term sweep ~2× at stem 100 (see `benches/milnor_gpu_ab.rs`). +//! +//! 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::{ + 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) +} + +/// Same extraction as [`get_partial_matrix`] but the products are multiplied on the CPU +/// via [`MilnorAlgebra::multiply_basis_element_by_element_2`] (the kernel's reference +/// algorithm). Isolates the extraction from the GPU batch: if this matches the CPU +/// `get_partial_matrix` but [`get_partial_matrix`] does not, the fault is the GPU batch, +/// not the extraction. +pub fn get_partial_matrix_cpu2(hom: &NassauDifferential, degree: i32, inputs: &[usize]) -> Matrix { + use fp::vector::FpVector; + let (mut matrix, products) = extract(hom, degree, inputs); + let algebra = hom.target().algebra(); + for prod in &products { + let out_deg = prod.r_degree + prod.s_degree; + let s_dim = algebra.dimension(prod.s_degree); + let mut s = FpVector::new(hom.prime(), s_dim); + for &ti in &prod.term_indices { + s.set_entry(ti, 1); + } + let mut tmp = FpVector::new(hom.prime(), algebra.dimension(out_deg)); + algebra.multiply_basis_element_by_element_2( + tmp.as_slice_mut(), + 1, + prod.r_degree, + prod.r_idx, + prod.s_degree, + s.as_slice(), + ); + let mut target_row = matrix.row_mut(prod.row); + for (i, _) in tmp.iter_nonzero() { + target_row.add_basis_element(prod.out_offset + i, 1); + } + } + matrix +} + +/// 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_timing.rs b/ext/tests/nassau_gpu_timing.rs new file mode 100644 index 0000000000..d64c1564e8 --- /dev/null +++ b/ext/tests/nassau_gpu_timing.rs @@ -0,0 +1,49 @@ +//! 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`. +#![cfg(feature = "gpu")] + +use std::time::Instant; + +use ext::{chain_complex::ChainComplex, 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() +} + +#[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); +} From a3e5ed231d5c57ed870924ef8b3d7393edebe790 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Sat, 11 Jul 2026 02:38:55 -0400 Subject: [PATCH 06/10] Optimize the GPU Milnor multiply: reuse, caching, and memory Compute the Nassau differential matrix once per bidegree, cache admissible-matrix enumeration across launches, keep the admissible buffers resident on-device, and release per-launch memory on a single pinned CUDA stream so cleanup actually reclaims it. Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_01Fsxqsgjoa5RWYhAvRG1bT2 --- ext/crates/algebra/Cargo.toml | 6 +- ext/crates/algebra/src/algebra/milnor_gpu.rs | 320 ++++++++++++------- ext/src/nassau.rs | 62 +++- ext/tests/nassau_gpu_reuse.rs | 55 ++++ ext/tests/nassau_gpu_timing.rs | 27 ++ 5 files changed, 354 insertions(+), 116 deletions(-) create mode 100644 ext/tests/nassau_gpu_reuse.rs diff --git a/ext/crates/algebra/Cargo.toml b/ext/crates/algebra/Cargo.toml index ccb6eda417..09031e4c4b 100644 --- a/ext/crates/algebra/Cargo.toml +++ b/ext/crates/algebra/Cargo.toml @@ -36,6 +36,10 @@ enum_dispatch = "0.3.13" 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"] } @@ -47,7 +51,7 @@ rstest = "0.25.0" default = ["odd-primes"] cache-multiplication = [] concurrent = ["fp/concurrent", "maybe-rayon/concurrent"] -gpu = ["dep:cubecl"] +gpu = ["dep:cubecl", "dep:cubecl-common"] odd-primes = ["fp/odd-primes"] [[bench]] diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index a4db426e06..317db01a02 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -35,6 +35,17 @@ 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 }; use crate::algebra::{ Algebra, MilnorAlgebra, @@ -65,6 +76,75 @@ pub fn take_batch_stats() -> (u64, u64, u64, u64) { ) } +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| v as u16)); + self.masks.extend(mk.iter().map(|&v| v as u16)); + 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 @@ -602,40 +682,11 @@ pub fn multiply_batch_on_gpu( prod_r_index.push(ri); } - // Parallel: each distinct `R`'s admissible matrices. Values ship as u16 (p-part - // magnitudes ≤ the output degree, far below 65535) to halve the transfer. - let r_mats: Vec<(usize, usize, Vec, Vec)> = (0..distinct_r.len()) - .into_maybe_par_iter() - .map(|i| { - let (rd, ridx) = distinct_r[i]; - let r = algebra.basis_element_from_index(rd, ridx); - assert!(!r.p_part.is_empty(), "each R must be non-empty"); - let (cs_len, mk_len, cs, mk) = algebra.admissible_matrices(&r.p_part); - let cs16 = cs.iter().map(|&v| v as u16).collect(); - let mk16 = mk.iter().map(|&v| v as u16).collect(); - (cs_len, mk_len, cs16, mk16) - }) - .collect(); - - // Lay out per-`R` matrix data (sequential concat + offsets). - let mut col_sums: Vec = Vec::new(); - let mut masks: Vec = Vec::new(); - let mut r_cs_offset: Vec = Vec::with_capacity(r_mats.len()); - let mut r_mk_offset: Vec = Vec::with_capacity(r_mats.len()); - let mut r_cs_len: Vec = Vec::with_capacity(r_mats.len()); - let mut r_mk_len: Vec = Vec::with_capacity(r_mats.len()); - let mut r_num_matrices: Vec = Vec::with_capacity(r_mats.len()); - for (cs_len, mk_len, cs, mk) in &r_mats { - r_cs_offset.push(col_sums.len() as u32); - r_mk_offset.push(masks.len() as u32); - r_cs_len.push(*cs_len as u32); - r_mk_len.push(*mk_len as u32); - r_num_matrices.push(mk.len() / mk_len); - col_sums.extend_from_slice(cs); - masks.extend_from_slice(mk); - } + // 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 t_terms = std::time::Instant::now(); let per_prod: Vec<(Vec, Vec)> = (0..products.len()) .into_maybe_par_iter() .map(|pi| { @@ -653,8 +704,32 @@ pub fn multiply_batch_on_gpu( (tp, tl) }) .collect(); + let terms_ms = t_terms.elapsed().as_secs_f64() * 1e3; + + // 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 t_resident = std::time::Instant::now(); + 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); + } + let resident_ms = t_resident.elapsed().as_secs_f64() * 1e3; // Lay out per-product term data + records + the pair-count prefix sum (sequential). + let t_glue = std::time::Instant::now(); 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()); @@ -683,105 +758,130 @@ pub fn multiply_batch_on_gpu( return vec![vec![0u32; num_limbs]; num_rows]; } - // Device buffers must be non-empty even when a dimension is zero for every product. - if col_sums.is_empty() { - col_sums.push(0); - } - if masks.is_empty() { - masks.push(0); - } + // 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 glue_ms = t_glue.elapsed().as_secs_f64() * 1e3; let marshal_ms = t_marshal.elapsed().as_secs_f64() * 1e3; if timing { + eprintln!( + " [marshal] {marshal_ms:.1}ms = terms {terms_ms:.1} + resident {resident_ms:.1} + \ + glue {glue_ms:.1}", + ); let mb = |n: usize| (n * 4) as f64 / (1024.0 * 1024.0); let mb16 = |n: usize| (n * 2) as f64 / (1024.0 * 1024.0); eprintln!( - " [sizes] {} products, {} distinct R | upload(u16): col_sums {:.1}MB masks \ - {:.1}MB term_pparts {:.1}MB per-product {:.1}MB (was ~{:.0}MB pair table)", + " [sizes] {} products, {} distinct R | resident master(u16): col_sums {:.1}MB \ + masks {:.1}MB | per-launch upload: term_pparts {:.1}MB per-product {:.1}MB", products.len(), - r_cs_len.len(), - mb16(col_sums.len()), - mb16(masks.len()), + distinct_r.len(), + mb16(resident.col_sums.len()), + mb16(resident.masks.len()), mb16(term_pparts.len()), mb(term_lens.len() + prod_r_index.len() * 4 + prod_pair_start.len()), - mb(total_pairs * 7), ); } let t_device = std::time::Instant::now(); - let client = CudaRuntime::client(&CudaDevice::default()); - let t_upload = std::time::Instant::now(); - 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 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)); - // Barrier so the upload timer excludes kernel/readback (timing runs only). - if timing { - let _ = cubecl::future::block_on(client.sync()); - } - let upload_ms = t_upload.elapsed().as_secs_f64() * 1e3; + // 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, upload_ms, kernel_ms, read_ms) = GPU_STREAM.executes(|| { + let client = CudaRuntime::client(&CudaDevice::default()); + let t_upload = std::time::Instant::now(); + // 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)); + // Barrier so the upload timer excludes kernel/readback (timing runs only). + if timing { + let _ = cubecl::future::block_on(client.sync()); + } + let upload_ms = t_upload.elapsed().as_secs_f64() * 1e3; + + let t_kernel = std::time::Instant::now(); + 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 t_kernel = std::time::Instant::now(); - 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, 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(), 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, - ); - } + // Barrier so the kernel timer excludes readback (timing runs only). + if timing { + let _ = cubecl::future::block_on(client.sync()); + } + let kernel_ms = t_kernel.elapsed().as_secs_f64() * 1e3; + + let t_read = std::time::Instant::now(); + let bytes = client.read_one(out_h).unwrap(); + let read_ms = t_read.elapsed().as_secs_f64() * 1e3; + 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(); - // Barrier so the kernel timer excludes readback (timing runs only). - if timing { - let _ = cubecl::future::block_on(client.sync()); - } - let kernel_ms = t_kernel.elapsed().as_secs_f64() * 1e3; + // `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(); - let t_read = std::time::Instant::now(); - let bytes = client.read_one(out_h).unwrap(); - let read_ms = t_read.elapsed().as_secs_f64() * 1e3; - let flat = u32::from_bytes(&bytes); - let result = (0..num_rows) - .map(|r| flat[r * num_limbs..(r + 1) * num_limbs].to_vec()) - .collect(); + (result, upload_ms, kernel_ms, read_ms) + }); // 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`]. diff --git a/ext/src/nassau.rs b/ext/src/nassau.rs index e929b9fd13..709cb427ea 100644 --- a/ext/src/nassau.rs +++ b/ext/src/nassau.rs @@ -424,6 +424,38 @@ fn build_partial_matrix( 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 @@ -652,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(); - build_partial_matrix(&self.differentials[b.s() - 1], 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]); @@ -721,9 +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(); - build_partial_matrix(&self.differential(b.s() - 1), 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/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 index d64c1564e8..176163ad87 100644 --- a/ext/tests/nassau_gpu_timing.rs +++ b/ext/tests/nassau_gpu_timing.rs @@ -16,6 +16,33 @@ fn resolve_time(stem: i32, filt: i32) -> f64 { 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() { From 97c53818e236fdb5bcddacf30ba7ffa71c59a7a8 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Sat, 11 Jul 2026 02:39:13 -0400 Subject: [PATCH 07/10] Fix lint and rustdoc warnings in the GPU code Collapse a nested if into a let-chain, drop an unused import, and repair broken/ambiguous intra-doc links (fully-qualify a cross-crate link, and demote links to private items and cube-macro-shadowed kernels to code spans). Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_01Fsxqsgjoa5RWYhAvRG1bT2 --- .../algebra/src/algebra/milnor_algebra.rs | 18 +++++++------- ext/crates/algebra/src/algebra/milnor_gpu.rs | 24 +++++++++---------- ext/examples/nassau_e2e.rs | 2 +- ext/tests/nassau_gpu_timing.rs | 2 +- 4 files changed, 23 insertions(+), 23 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_algebra.rs b/ext/crates/algebra/src/algebra/milnor_algebra.rs index dfcbe2460b..6ce97addb7 100644 --- a/ext/crates/algebra/src/algebra/milnor_algebra.rs +++ b/ext/crates/algebra/src/algebra/milnor_algebra.rs @@ -29,7 +29,7 @@ macro_rules! profile { /// /// Enable by building with `MILNOR_PROFILE=1` (e.g. /// `MILNOR_PROFILE=1 cargo run --release --example nassau_e2e -- 80 42 1`); otherwise every counter -/// and hook is removed by `#[cfg]`. Call [`report`] to print a summary to stderr. +/// and hook is removed by `#[cfg]`. Call [`profile::report`] to print a summary to stderr. pub mod profile { #[cfg(milnor_profile)] pub use enabled::*; @@ -1355,8 +1355,8 @@ impl MilnorAlgebra { !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 [`Self::seqno_applicable`]. Idempotent: if the stored tables already reach + /// 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. @@ -1367,10 +1367,10 @@ impl MilnorAlgebra { /// 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() { - if t.max_degree >= max_degree { - return; - } + if let Some(t) = &*self.seqno_tables.load() + && t.max_degree >= max_degree + { + return; } let xi = combinatorics::xi_degrees(self.prime()); @@ -1424,7 +1424,7 @@ impl MilnorAlgebra { /// 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 - /// [`Self::seqno_applicable`] and that `p_part` is a genuine basis element (trimmed, in range). + /// `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` @@ -1459,7 +1459,7 @@ impl MilnorAlgebra { /// 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 [`Self::seqno_applicable`] regime (`p = 2`, trivial profile, stable), since + /// 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")] diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 317db01a02..58ab008a5e 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -7,16 +7,16 @@ //! `get_partial_matrix` launch. //! //! Stages landed so far: -//! - **Stage 1 — toolchain slice.** A trivial F₂ kernel ([`xor_f2`]) proving the +//! - **Stage 1 — toolchain slice.** A trivial F₂ kernel (`xor_f2`) proving the //! CubeCL `cuda` runtime works end to end. F₂ addition is XOR of the bit-packed //! limbs, the exact accumulation every multiply kernel performs. -//! - **Stage 2 — `seqno` on device.** [`seqno_core`]/[`seqno_kernel`] port the +//! - **Stage 2 — `seqno` on device.** `seqno_core`/`seqno_kernel` port the //! hash-free index [`MilnorAlgebra::seqno`] as pure integer arithmetic over the //! flat `g` table — the output-indexing primitive the multiply kernel needs. -//! - **Stage 3 — single-`R` multiply.** [`multiply_pair`] ports the per-term test +//! - **Stage 3 — single-`R` multiply.** `multiply_pair` ports the per-term test //! + output assembly of `multiply_basis_element_by_element_2`; -//! [`multiply_single_r_kernel`] runs one `Sq(R)·s` product per launch. -//! - **Stage 4 — batched launch.** [`multiply_batch_kernel`] fuses all `(R, s)` +//! `multiply_single_r_kernel` runs one `Sq(R)·s` product per launch. +//! - **Stage 4 — batched launch.** `multiply_batch_kernel` fuses all `(R, s)` //! products of one `get_partial_matrix` into a single launch (one thread per //! `(product, matrix, term)` pair). The pair is decoded on-device from a prefix-sum //! over per-product pair counts, with admissible-matrix data deduplicated by distinct @@ -158,7 +158,7 @@ fn xor_f2(a: &Array, b: &Array, out: &mut Array) { /// Compute `a ^ b` limb-wise on the default CUDA device. /// -/// Host-side driver for [`xor_f2`]: uploads both operands, launches one thread per +/// 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. pub fn xor_f2_on_gpu(a: &[u32], b: &[u32]) -> Vec { assert_eq!(a.len(), b.len(), "operands must have equal limb counts"); @@ -194,8 +194,8 @@ pub fn xor_f2_on_gpu(a: &[u32], b: &[u32]) -> Vec { /// /// 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. +/// `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, @@ -250,9 +250,9 @@ fn seqno_kernel( out[idx] = seqno_core(g, xi, &working, width, width); } -/// Run [`seqno_kernel`] over `n` padded p_parts and return their seqno indices. +/// Run `seqno_kernel` over `n` padded p_parts and return their seqno indices. /// -/// `g`/`xi` come from [`MilnorAlgebra::seqno_table_u32`] and +/// `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`. pub fn seqno_batch_on_gpu( @@ -302,7 +302,7 @@ pub fn seqno_batch_on_gpu( /// - `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 +/// 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 @@ -385,7 +385,7 @@ fn multiply_pair( } /// 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`]. +/// One thread per `(matrix, term)` pair; delegates the assembly to `multiply_pair`. #[cube(launch)] #[allow(clippy::too_many_arguments)] fn multiply_single_r_kernel( diff --git a/ext/examples/nassau_e2e.rs b/ext/examples/nassau_e2e.rs index 388636374f..352c813abc 100644 --- a/ext/examples/nassau_e2e.rs +++ b/ext/examples/nassau_e2e.rs @@ -1,6 +1,6 @@ //! End-to-end timing harness for Nassau's algorithm on the sphere `S_2` at p = 2. //! -//! Resolves `S_2` via [`construct_nassau`] and [`Resolution::compute_through_stem`] up to a given +//! 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 diff --git a/ext/tests/nassau_gpu_timing.rs b/ext/tests/nassau_gpu_timing.rs index 176163ad87..756c3d8ac9 100644 --- a/ext/tests/nassau_gpu_timing.rs +++ b/ext/tests/nassau_gpu_timing.rs @@ -6,7 +6,7 @@ use std::time::Instant; -use ext::{chain_complex::ChainComplex, utils::construct_nassau}; +use ext::utils::construct_nassau; use sseq::coordinates::Bidegree; fn resolve_time(stem: i32, filt: i32) -> f64 { From 1476426361595149725e777bafc4312c606fdcae Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Sat, 11 Jul 2026 02:58:35 -0400 Subject: [PATCH 08/10] Remove the compile-time Milnor multiply profiler Drops the dev-only MILNOR_PROFILE profiling harness that guided the GPU work but has no place in the shipped code: the `profile!` macro and `profile` module (counters, occupancy metrics, report), their hot-path hooks in the multiply and in homomorphism matrix builds, the build.rs cfg wiring, and the nassau_e2e example's opt-in pprof flamegraph. The algebra's mathematical `profile` field (the Milnor sub-Hopf-algebra truncation) and the criterion benchmark suite are untouched. Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_01Fsxqsgjoa5RWYhAvRG1bT2 --- ext/Cargo.toml | 3 - ext/crates/algebra/build.rs | 13 - .../algebra/src/algebra/milnor_algebra.rs | 439 ------------------ .../algebra/src/module/homomorphism/mod.rs | 20 - ext/examples/nassau_e2e.rs | 30 -- 5 files changed, 505 deletions(-) delete mode 100644 ext/crates/algebra/build.rs diff --git a/ext/Cargo.toml b/ext/Cargo.toml index 6aa52ef8ec..017a0f859d 100644 --- a/ext/Cargo.toml +++ b/ext/Cargo.toml @@ -33,7 +33,6 @@ serde_json = { version = "1.0.141", features = ["preserve_order"] } tracing = "0.1.41" tracing-subscriber = { version = "0.3.19", features = ["env-filter"] } -pprof = { version = "0.15.0", features = ["flamegraph"], optional = true } zstd = { version = "0.13.3", optional = true } [target.'cfg(unix)'.dependencies] @@ -60,8 +59,6 @@ concurrent = [ odd-primes = ["fp/odd-primes", "algebra/odd-primes", "sseq/odd-primes"] logging = [] nassau = [] -# Enables the optional pprof-based flamegraph in the `nassau_e2e` example (set `FLAME=`). -flamegraph = ["dep:pprof"] # Enables the GPU Milnor-multiply backend and the `milnor_gpu_ab` A/B microbenchmark. gpu = ["algebra/gpu"] diff --git a/ext/crates/algebra/build.rs b/ext/crates/algebra/build.rs deleted file mode 100644 index cf792dc271..0000000000 --- a/ext/crates/algebra/build.rs +++ /dev/null @@ -1,13 +0,0 @@ -//! Turns the build-time `MILNOR_PROFILE` environment variable into `cfg(milnor_profile)`, which -//! gates the (otherwise zero-cost) Milnor-multiplication profiling counters in -//! `src/algebra/milnor_algebra.rs`. Build with e.g. `MILNOR_PROFILE=1 cargo build` to enable them. - -fn main() { - // Declare the cfg so `cfg(milnor_profile)` does not trip the `unexpected_cfgs` lint. - println!("cargo::rustc-check-cfg=cfg(milnor_profile)"); - // Rebuild when the toggle changes. - println!("cargo::rerun-if-env-changed=MILNOR_PROFILE"); - if std::env::var_os("MILNOR_PROFILE").is_some() { - println!("cargo::rustc-cfg=milnor_profile"); - } -} diff --git a/ext/crates/algebra/src/algebra/milnor_algebra.rs b/ext/crates/algebra/src/algebra/milnor_algebra.rs index 6ce97addb7..6a5de7bc20 100644 --- a/ext/crates/algebra/src/algebra/milnor_algebra.rs +++ b/ext/crates/algebra/src/algebra/milnor_algebra.rs @@ -11,418 +11,6 @@ use serde::{Deserialize, Serialize}; use crate::algebra::{Algebra, Bialgebra, GeneratedAlgebra, UnstableAlgebra, combinatorics}; -/// Wrap profiling-only statements. Expands to nothing unless the crate is built with the -/// `MILNOR_PROFILE` environment variable set (which turns on `cfg(milnor_profile)`; see `build.rs`), -/// so the hot-path hooks below cost *nothing* — not even argument evaluation — in normal builds. -macro_rules! profile { - ($($body:tt)*) => { - #[cfg(milnor_profile)] - { - $($body)* - } - }; -} - -/// Compile-time-gated counters for the Milnor multiplication hot path, to pinpoint what to -/// optimize (how sparse the operands are, how much matrix enumeration is wasted, how many index -/// lookups we pay, which path is taken, …). -/// -/// Enable by building with `MILNOR_PROFILE=1` (e.g. -/// `MILNOR_PROFILE=1 cargo run --release --example nassau_e2e -- 80 42 1`); otherwise every counter -/// and hook is removed by `#[cfg]`. Call [`profile::report`] to print a summary to stderr. -pub mod profile { - #[cfg(milnor_profile)] - pub use enabled::*; - - /// Print the collected multiply statistics to stderr (or a note if profiling is disabled). - #[cfg(not(milnor_profile))] - pub fn report() { - eprintln!( - "[milnor profile] disabled — rebuild with `MILNOR_PROFILE=1` to collect multiply stats" - ); - } - - #[cfg(milnor_profile)] - mod enabled { - use std::sync::{ - LazyLock, Mutex, - atomic::{AtomicU64, Ordering::Relaxed}, - }; - - use rustc_hash::FxHashMap; - - macro_rules! counters { - ($($(#[$m:meta])* $name:ident),* $(,)?) => { - $($(#[$m])* pub static $name: AtomicU64 = AtomicU64::new(0);)* - }; - } - counters! { - /// `multiply_basis_element_by_element` (p=2) calls that do work. - MBE_CALLS, - /// … that took the admissible-matrix sweep. - MBE_ADMISSIBLE, - /// … that fell back to the per-term `PPartMultiplier` path. - MBE_PERTERM, - /// … where the operation was the identity `Sq(∅)`. - MBE_IDENTITY, - /// Basis×basis `multiply_with_allocation` invocations. - KERNEL_CALLS, - /// `PPartMultiplier` candidate terms yielded (before the index/bound checks). - PPART_TERMS, - /// `basis_element_to_index` lookups on the multiply hot path. - INDEX_LOOKUPS, - /// `add_basis_element` calls that actually contribute an output term. - OUTPUT_ADDS, - /// Admissible matrices enumerated (admissible path only). - ADM_MATRICES, - /// (matrix, term) compatibility tests (admissible path only). - ADM_TESTS, - } - - /// Histogram of `nnz(s)` — the number of non-zero terms of the element being acted on. - pub static NNZ_HIST: LazyLock>> = - LazyLock::new(|| Mutex::new(FxHashMap::default())); - /// Histogram of `(operation_degree, element_degree)` per call. - pub static DEG_HIST: LazyLock>> = - LazyLock::new(|| Mutex::new(FxHashMap::default())); - /// Total element terms multiplied by each distinct operation `R = (degree, index)`. This is - /// the GPU-occupancy metric for Nassau's `Sq(R) · Σ Sq(Sⱼ)` kernel: batching all work sharing - /// one `R` (uniform matrix enumeration) gives this many threads' worth of uniform work, so - /// the distribution says how well workgroups would fill. - pub static R_TERMS: LazyLock>> = - LazyLock::new(|| Mutex::new(FxHashMap::default())); - - pub fn record_call(nnz: usize, r_degree: i32, r_index: usize, s_degree: i32) { - MBE_CALLS.fetch_add(1, Relaxed); - *NNZ_HIST.lock().unwrap().entry(nnz).or_default() += 1; - *DEG_HIST - .lock() - .unwrap() - .entry((r_degree, s_degree)) - .or_default() += 1; - *R_TERMS - .lock() - .unwrap() - .entry((r_degree, r_index)) - .or_default() += nnz as u64; - } - - /// Per–GPU-launch accumulation of `R → element-terms`, where one launch = one - /// `get_partial_matrix` (matrix build). Unlike [`R_TERMS`], which sums an operation `R` - /// across the *whole* resolution, this asks how much same-`R` work is co-located within a - /// single matrix build — the largest unit a kernel could batch without buffering across the - /// streaming algorithm. `depth` keeps [`scope_begin`]/[`scope_end`] reentrancy-safe: only the - /// outermost pair clears and folds, so any nested matrix build is attributed to the outer - /// launch rather than corrupting it. (Measurement must be run serially — the `concurrent` - /// feature would let two launches interleave into one scope and inflate the batch sizes.) - struct ScopeState { - depth: u32, - /// `(homomorphism id, degree)` of the currently open launch — its merged-scope key. - key: (usize, i32), - map: FxHashMap<(i32, usize), u64>, - } - static SCOPE: LazyLock> = LazyLock::new(|| { - Mutex::new(ScopeState { - depth: 0, - key: (0, 0), - map: FxHashMap::default(), - }) - }); - - /// Coarser accumulation: `R → element-terms` merged across every launch sharing a - /// `(homomorphism id, degree)` key — i.e. all the per-signature `get_partial_matrix` builds - /// of one differential at one bidegree, which all read the same matrix and so *could* be - /// fused into a single kernel launch (computing the full matrix once and slicing it). Held to - /// report time and folded there, to measure how much the realizable occupancy improves when - /// the launch scope is widened from one masked build to one whole bidegree. - #[allow(clippy::type_complexity)] - static MERGED: LazyLock>>> = - LazyLock::new(|| Mutex::new(FxHashMap::default())); - - /// Workgroup sizes for the realizable-coverage report (mirrors the global one). - const SCOPE_WS: [u64; 6] = [32, 64, 128, 256, 512, 1024]; - /// Number of launch scopes (matrix builds) that did any work. - static SCOPE_COUNT: AtomicU64 = AtomicU64::new(0); - /// Sum over scopes of each scope's total element-terms (equals the global term-work). - static SCOPE_TERMS: AtomicU64 = AtomicU64::new(0); - /// Largest single-scope total element-terms. - static SCOPE_TERMS_MAX: AtomicU64 = AtomicU64::new(0); - /// Per `W`, term-work in `R`s that reach ≥ `W` terms *within their own launch*. - static SCOPE_COVER: [AtomicU64; 6] = [ - AtomicU64::new(0), - AtomicU64::new(0), - AtomicU64::new(0), - AtomicU64::new(0), - AtomicU64::new(0), - AtomicU64::new(0), - ]; - - /// Open a launch scope around a matrix build of homomorphism `hom_id` at `degree`. - /// Reentrancy-safe (see [`ScopeState`]); the key is set by the outermost open. - pub fn scope_begin(hom_id: usize, degree: i32) { - let mut s = SCOPE.lock().unwrap(); - if s.depth == 0 { - s.map.clear(); - s.key = (hom_id, degree); - } - s.depth += 1; - } - - /// Attribute `nnz` element-terms of operation `R = (r_degree, r_index)` to the open launch, - /// both to its per-launch histogram and to its `(hom, degree)` merged bucket. - pub fn scope_record(r_degree: i32, r_index: usize, nnz: usize) { - if nnz == 0 { - return; - } - let key = { - let mut s = SCOPE.lock().unwrap(); - if s.depth == 0 { - return; - } - *s.map.entry((r_degree, r_index)).or_default() += nnz as u64; - s.key - }; - *MERGED - .lock() - .unwrap() - .entry(key) - .or_default() - .entry((r_degree, r_index)) - .or_default() += nnz as u64; - } - - /// Close a launch scope; the outermost close folds the scope's `R`-histogram into the - /// realizable-coverage accumulators. - pub fn scope_end() { - let mut s = SCOPE.lock().unwrap(); - if s.depth == 0 { - return; - } - s.depth -= 1; - if s.depth > 0 { - return; - } - let total: u64 = s.map.values().sum(); - if total == 0 { - return; - } - SCOPE_COUNT.fetch_add(1, Relaxed); - SCOPE_TERMS.fetch_add(total, Relaxed); - SCOPE_TERMS_MAX.fetch_max(total, Relaxed); - for (i, &w) in SCOPE_WS.iter().enumerate() { - let work: u64 = s.map.values().filter(|&&t| t >= w).sum(); - SCOPE_COVER[i].fetch_add(work, Relaxed); - } - s.map.clear(); - } - - pub fn admissible() { - MBE_ADMISSIBLE.fetch_add(1, Relaxed); - } - pub fn perterm() { - MBE_PERTERM.fetch_add(1, Relaxed); - } - pub fn identity() { - MBE_IDENTITY.fetch_add(1, Relaxed); - } - pub fn kernel_call() { - KERNEL_CALLS.fetch_add(1, Relaxed); - } - pub fn ppart_term() { - PPART_TERMS.fetch_add(1, Relaxed); - } - pub fn index_lookup() { - INDEX_LOOKUPS.fetch_add(1, Relaxed); - } - pub fn output_add() { - OUTPUT_ADDS.fetch_add(1, Relaxed); - } - pub fn adm_matrix() { - ADM_MATRICES.fetch_add(1, Relaxed); - } - pub fn adm_test() { - ADM_TESTS.fetch_add(1, Relaxed); - } - - pub fn report() { - let calls = MBE_CALLS.load(Relaxed); - if calls == 0 { - eprintln!("[milnor profile] no multiply_basis_element_by_element (p=2) calls seen"); - return; - } - let (adm, per, ident) = ( - MBE_ADMISSIBLE.load(Relaxed), - MBE_PERTERM.load(Relaxed), - MBE_IDENTITY.load(Relaxed), - ); - let zero = calls - adm - per - ident; - let adds = OUTPUT_ADDS.load(Relaxed); - let pct = |x: u64| 100.0 * x as f64 / calls as f64; - let per_add = |x: u64| { - if adds > 0 { - x as f64 / adds as f64 - } else { - 0.0 - } - }; - - eprintln!("================= Milnor multiply profile ================="); - eprintln!("multiply_basis_element_by_element (p=2) calls : {calls}"); - eprintln!(" admissible-matrix path : {adm:>12} ({:5.1}%)", pct(adm)); - eprintln!(" per-term path : {per:>12} ({:5.1}%)", pct(per)); - eprintln!( - " identity Sq(∅) : {ident:>12} ({:5.1}%)", - pct(ident) - ); - eprintln!(" zero element (nnz=0) : {zero:>12} ({:5.1}%)", pct(zero)); - - let hist = NNZ_HIST.lock().unwrap(); - let n: u64 = hist.values().sum(); - let weighted: u64 = hist.iter().map(|(k, v)| *k as u64 * v).sum(); - eprintln!( - "element term-count nnz : mean {:.2} over {n} calls", - weighted as f64 / n.max(1) as f64 - ); - let mut keys: Vec = hist.keys().copied().collect(); - keys.sort_unstable(); - let mut cum = 0u64; - for k in keys { - let v = hist[&k]; - cum += v; - if k <= 6 || k.is_multiple_of(10) { - eprintln!( - " nnz={k:<3} {:>6.2}% cum {:>6.2}%", - 100.0 * v as f64 / n as f64, - 100.0 * cum as f64 / n as f64 - ); - } - } - - eprintln!( - "kernel basis×basis multiply_with_allocation : {}", - KERNEL_CALLS.load(Relaxed) - ); - eprintln!("output basis-element adds : {adds}"); - eprintln!( - "PPartMultiplier candidate terms : {} ({:.2} per output add)", - PPART_TERMS.load(Relaxed), - per_add(PPART_TERMS.load(Relaxed)) - ); - eprintln!( - "basis_element_to_index lookups : {} ({:.2} per output add)", - INDEX_LOOKUPS.load(Relaxed), - per_add(INDEX_LOOKUPS.load(Relaxed)) - ); - let matrices = ADM_MATRICES.load(Relaxed); - if matrices > 0 { - eprintln!( - "admissible matrices enumerated : {matrices} ({:.1} term-tests each)", - ADM_TESTS.load(Relaxed) as f64 / matrices as f64 - ); - } - - let deg = DEG_HIST.lock().unwrap(); - let mut pairs: Vec<((i32, i32), u64)> = deg.iter().map(|(k, v)| (*k, *v)).collect(); - pairs.sort_by_key(|(_, v)| std::cmp::Reverse(*v)); - eprintln!("top (operation_degree, element_degree) call sites:"); - for ((r, s), v) in pairs.iter().take(10) { - eprintln!( - " op={r:<3} elt={s:<3} : {v} ({:.1}%)", - 100.0 * *v as f64 / calls as f64 - ); - } - - // GPU-occupancy view: group all element terms by the operation `R` they are multiplied - // by. Nassau's kernel (`Sq(R) · Σ Sq(Sⱼ)`) parallelizes over terms sharing one `R` with - // uniform matrix enumeration, so a workgroup of size `W` is well-filled only by `R`s that - // accumulate ≥ `W` terms. We report what fraction of all term-work lives in such `R`s. - let r_terms = R_TERMS.lock().unwrap(); - let distinct = r_terms.len() as u64; - let total_terms: u64 = r_terms.values().sum(); - let max_terms = r_terms.values().copied().max().unwrap_or(0); - eprintln!( - "GPU occupancy — distinct operations R: {distinct}, total element terms: \ - {total_terms}, mean terms/R: {:.1}, max: {max_terms}", - total_terms as f64 / distinct.max(1) as f64 - ); - eprintln!(" fraction of term-work in R's with ≥ W terms (W = workgroup size):"); - for w in [32u64, 64, 128, 256, 512, 1024] { - let (n_r, work): (u64, u64) = - r_terms.values().fold( - (0, 0), - |(n, s), &t| { - if t >= w { (n + 1, s + t) } else { (n, s) } - }, - ); - eprintln!( - " W={w:<4} : {n_r:>6} R's ({:>5.1}% of R's) cover {:>5.1}% of term-work", - 100.0 * n_r as f64 / distinct.max(1) as f64, - 100.0 * work as f64 / total_terms.max(1) as f64 - ); - } - - // The number above aggregates each R across the *whole* resolution. A kernel can only - // batch work that is co-located in one launch, so this is the realizable version: R-work - // is re-counted per `get_partial_matrix` (matrix build), the largest unit batchable - // without buffering across the streaming algorithm. If these percentages collapse - // relative to the global ones, the amortization-by-R shape does not survive at launch - // granularity, and a GPU kernel must lean on raw parallel width (many independent - // products per launch) rather than on many terms sharing one R. - let scopes = SCOPE_COUNT.load(Relaxed); - if scopes > 0 { - let scope_terms = SCOPE_TERMS.load(Relaxed); - eprintln!( - "GPU occupancy — REALIZABLE per matrix-build launch ({scopes} launches, mean \ - {:.1} terms/launch, max {}):", - scope_terms as f64 / scopes as f64, - SCOPE_TERMS_MAX.load(Relaxed) - ); - eprintln!(" fraction of term-work in R's reaching ≥ W terms *within one launch*:"); - for (i, w) in SCOPE_WS.iter().enumerate() { - let work = SCOPE_COVER[i].load(Relaxed); - eprintln!( - " W={w:<4} : {:>5.1}% of term-work", - 100.0 * work as f64 / scope_terms.max(1) as f64 - ); - } - } - - // Coarser scope: merge the per-signature builds of one differential at one bidegree into - // a single launch (they read the same matrix, so fusing them is free of extra multiply - // work). This upper-bounds the realizable occupancy for a kernel that computes the full - // bidegree matrix once instead of one masked slice per signature. - let merged = MERGED.lock().unwrap(); - if !merged.is_empty() { - let n_launches = merged.len() as u64; - let mut total = 0u64; - let mut max_terms = 0u64; - let mut cover = [0u64; 6]; - for hist in merged.values() { - let t: u64 = hist.values().sum(); - total += t; - max_terms = max_terms.max(t); - for (i, &w) in SCOPE_WS.iter().enumerate() { - cover[i] += hist.values().filter(|&&x| x >= w).sum::(); - } - } - eprintln!( - "GPU occupancy — MERGED per (differential, bidegree) launch ({n_launches} \ - launches, mean {:.1} terms/launch, max {max_terms}):", - total as f64 / n_launches.max(1) as f64 - ); - eprintln!(" fraction of term-work in R's reaching ≥ W terms *within one launch*:"); - for (i, w) in SCOPE_WS.iter().enumerate() { - eprintln!( - " W={w:<4} : {:>5.1}% of term-work", - 100.0 * cover[i] as f64 / total.max(1) as f64 - ); - } - } - eprintln!("==========================================================="); - } - } -} - fn q_part_default() -> u32 { !0 } @@ -974,16 +562,6 @@ impl Algebra for MilnorAlgebra { // for which the up-front enumeration cannot be amortized. It is retained as the reference // model for a future GPU kernel (where the enumerate-once/test-all-terms shape is ideal), // not wired here. See the commit history for the measurements. - profile!({ - let nnz = s.iter_nonzero().count(); - profile::record_call(nnz, r_degree, r_idx, s_degree); - if r_degree == 0 { - profile::identity(); - } else if nnz > 0 { - profile::perterm(); - } - profile::scope_record(r_degree, r_idx, nnz); - }); let p = self.prime(); let r = self.basis_element_from_index(r_degree, r_idx); PPartAllocation::with_local(|mut allocation| { @@ -1720,18 +1298,11 @@ impl MilnorAlgebra { if coeff.is_multiple_of(2) { return; } - profile!(profile::record_call( - s.iter_nonzero().count(), - r_degree, - r_idx, - s_degree - )); 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() { - profile!(profile::identity()); for (i, _) in s.iter_nonzero() { result.add_basis_element(i, 1); } @@ -1749,7 +1320,6 @@ impl MilnorAlgebra { return; // s = 0 }; let Some((i1, _)) = second else { - profile!(profile::perterm()); PPartAllocation::with_local(|allocation| { self.multiply_with_allocation( result, @@ -1769,7 +1339,6 @@ impl MilnorAlgebra { 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))); - profile!(profile::admissible()); let out_degree = r_degree + s_degree; let mut matrix = AdmissibleMatrix::new(&r.p_part); @@ -1780,9 +1349,7 @@ impl MilnorAlgebra { }; loop { - profile!(profile::adm_matrix()); 'outer: for term in &terms { - profile!(profile::adm_test()); let basis = &term.p_part; working.p_part.clear(); @@ -1832,9 +1399,7 @@ impl MilnorAlgebra { } let idx = self.basis_element_to_index(&working); - profile!(profile::index_lookup()); result.add_basis_element(idx, 1); - profile!(profile::output_add()); } if !matrix.next() { break; @@ -1878,7 +1443,6 @@ impl MilnorAlgebra { allocation = multiplier.into_allocation() } } else { - profile!(profile::kernel_call()); let mut multiplier = PPartMultiplier::::new_from_allocation( self.prime(), &m1.p_part, @@ -1889,12 +1453,9 @@ impl MilnorAlgebra { ); while let Some(c) = multiplier.next() { - profile!(profile::ppart_term()); let idx = self.basis_element_to_index(&multiplier.ans); - profile!(profile::index_lookup()); if idx < dim { res.add_basis_element(idx, c * coef); - profile!(profile::output_add()); } } allocation = multiplier.into_allocation() diff --git a/ext/crates/algebra/src/module/homomorphism/mod.rs b/ext/crates/algebra/src/module/homomorphism/mod.rs index b6d3b08d89..540b34ac4b 100644 --- a/ext/crates/algebra/src/module/homomorphism/mod.rs +++ b/ext/crates/algebra/src/module/homomorphism/mod.rs @@ -130,20 +130,10 @@ pub trait ModuleHomomorphism: Send + Sync { return; } - // One matrix build is the largest unit a GPU kernel could batch without buffering across - // the streaming algorithm; mark it as a launch scope for the realizable-occupancy metric. - // The `(self pointer, degree)` key groups builds of the same differential at one bidegree. - #[cfg(milnor_profile)] - crate::algebra::milnor_algebra::profile::scope_begin( - std::ptr::from_ref(self) as *const () as usize, - degree, - ); matrix .maybe_par_iter_mut() .enumerate() .for_each(|(i, row)| self.apply_to_basis_element(row, 1, degree, i)); - #[cfg(milnor_profile)] - crate::algebra::milnor_algebra::profile::scope_end(); } /// Get the values of the homomorphism on the specified inputs to `matrix`. @@ -154,20 +144,10 @@ pub trait ModuleHomomorphism: Send + Sync { return matrix; } - // One matrix build is the largest unit a GPU kernel could batch without buffering across - // the streaming algorithm; mark it as a launch scope for the realizable-occupancy metric. - // The `(self pointer, degree)` key groups builds of the same differential at one bidegree. - #[cfg(milnor_profile)] - crate::algebra::milnor_algebra::profile::scope_begin( - std::ptr::from_ref(self) as *const () as usize, - degree, - ); matrix .maybe_par_iter_mut() .enumerate() .for_each(|(i, row)| self.apply_to_basis_element(row, 1, degree, inputs[i])); - #[cfg(milnor_profile)] - crate::algebra::milnor_algebra::profile::scope_end(); matrix } diff --git a/ext/examples/nassau_e2e.rs b/ext/examples/nassau_e2e.rs index 352c813abc..04220e2525 100644 --- a/ext/examples/nassau_e2e.rs +++ b/ext/examples/nassau_e2e.rs @@ -15,12 +15,6 @@ //! //! 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). -//! -//! To see *where* the time goes (e.g. multiply vs. linear algebra), build with the `flamegraph` -//! feature and set `FLAME` to an output path — a sampling flamegraph of every run is written there: -//! ```text -//! FLAME=/tmp/nassau.svg cargo run --release --features flamegraph --example nassau_e2e -- 60 32 1 -//! ``` use std::time::Instant; @@ -33,16 +27,6 @@ fn main() { 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); - // When built `--features flamegraph` and `FLAME=` is set, sample the whole run and write a - // flamegraph there. Kept alive until after the timing loop so it captures every resolution. - #[cfg(feature = "flamegraph")] - let flame = std::env::var("FLAME").ok().map(|path| { - ( - path, - pprof::ProfilerGuard::new(1000).expect("failed to start profiler"), - ) - }); - let target = Bidegree::n_s(n, s); let mut best = f64::INFINITY; for run in 1..=runs { @@ -59,18 +43,4 @@ fn main() { ); } println!("S_2 @ p=2 stem={n} filtration={s} runs={runs} best={best:.2}s"); - - // Prints Milnor-multiply counters when built with `MILNOR_PROFILE=1`, else a one-line note. - algebra::milnor_algebra::profile::report(); - - #[cfg(feature = "flamegraph")] - if let Some((path, guard)) = flame { - let report = guard - .report() - .build() - .expect("failed to build profiler report"); - let file = std::fs::File::create(&path).expect("failed to create flamegraph file"); - report.flamegraph(file).expect("failed to write flamegraph"); - eprintln!("wrote flamegraph to {path}"); - } } From c22f1c894a15ae093f3c6937942644cdc359e9f6 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Sat, 11 Jul 2026 03:05:55 -0400 Subject: [PATCH 09/10] Trim development-history comments from the GPU code Cut the process archaeology that documented how the port was built rather than what the code does: the stage-by-stage narrative in the milnor_gpu module doc, the A/B benchmark figures and "measured a regression / see the commit history" notes around the admissible multiply and seqno index, and the now-dangling references to the removed GPU_KERNEL_HANDOFF.md. The design rationale that explains the current code (why seqno is GPU-only, why the admissible multiply is the GPU reference model, the stream-pinning / OOM reasoning) is kept. Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_01Fsxqsgjoa5RWYhAvRG1bT2 --- ext/benches/milnor_gpu_ab.rs | 5 +-- .../algebra/src/algebra/milnor_algebra.rs | 37 ++++++++----------- ext/crates/algebra/src/algebra/milnor_gpu.rs | 36 ++++++------------ ext/src/nassau_gpu.rs | 4 +- 4 files changed, 31 insertions(+), 51 deletions(-) diff --git a/ext/benches/milnor_gpu_ab.rs b/ext/benches/milnor_gpu_ab.rs index 97faa2678e..1fb354a47f 100644 --- a/ext/benches/milnor_gpu_ab.rs +++ b/ext/benches/milnor_gpu_ab.rs @@ -1,8 +1,7 @@ //! A/B microbenchmark: GPU batch Milnor-multiply vs the CPU `get_partial_matrix`. //! -//! Answers the Stage 5 go/no-go question from `GPU_KERNEL_HANDOFF.md` — does offloading -//! the Milnor multiply to the GPU beat the CPU per-term path *including* host↔device -//! transfer — before committing to the invasive wire-in. +//! 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: diff --git a/ext/crates/algebra/src/algebra/milnor_algebra.rs b/ext/crates/algebra/src/algebra/milnor_algebra.rs index 6a5de7bc20..492229d09b 100644 --- a/ext/crates/algebra/src/algebra/milnor_algebra.rs +++ b/ext/crates/algebra/src/algebra/milnor_algebra.rs @@ -348,14 +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 it loses - // to this hashmap on the CPU. Even after moving its tables to flat, contiguous - // `arc_swap`-backed storage (removing the earlier `OnceVec` per-access atomics), the - // `benches/seqno` A/B still measures raw lookups at ~50 Melem/s for `seqno` vs ~115 Melem/s - // for this hashmap — a ~2.3× gap that is flat across degree: computing the rank (a degree - // sum plus two indexed table reads per populated ξ-position) is simply more work than one - // hash and probe. `seqno` is therefore 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. + // 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() @@ -555,13 +552,11 @@ impl Algebra for MilnorAlgebra { ) { // 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 end-to-end - // A/Bs of Nassau's `S_2` regime measured it a consistent net regression on the CPU (~8% at - // stem 80, ~3% at stem 100): the regime is dominated by sparse elements (≈31% single-term), - // for which the up-front enumeration cannot be amortized. It is retained as the reference - // model for a future GPU kernel (where the enumerate-once/test-all-terms shape is ideal), - // not wired here. See the commit history for the measurements. + // (`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| { @@ -1270,13 +1265,11 @@ impl MilnorAlgebra { /// 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.** End-to-end A/Bs of Nassau's `S_2` regime measured this a net - /// regression versus the per-term sweep in [`Self::multiply_basis_element_by_element`] (the - /// regime is too sparse for the up-front enumeration to pay off). It is kept as the reference - /// model for a future GPU kernel: enumerate `Sq(R)`'s matrices once per operation and test every - /// element term against them in parallel — a shape that batches extremely well on a GPU (a real - /// resolution presents tens of thousands of terms per distinct `R`). Exercised by the - /// `admissible_multiply_agrees_with_reference` test. + /// **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)] diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 58ab008a5e..40e3ec9aef 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -1,29 +1,17 @@ //! GPU offload for the Milnor multiply at `p = 2`, built on [CubeCL]. //! -//! This is the entry point for the staged port described in -//! `crates/algebra/GPU_KERNEL_HANDOFF.md`: the admissible-matrix multiply -//! ([`super::milnor_algebra::MilnorAlgebra::multiply_basis_element_by_element_2`]) -//! and the hash-free `seqno` index run as CubeCL kernels, batched per -//! `get_partial_matrix` launch. +//! 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. //! -//! Stages landed so far: -//! - **Stage 1 — toolchain slice.** A trivial F₂ kernel (`xor_f2`) proving the -//! CubeCL `cuda` runtime works end to end. F₂ addition is XOR of the bit-packed -//! limbs, the exact accumulation every multiply kernel performs. -//! - **Stage 2 — `seqno` on device.** `seqno_core`/`seqno_kernel` port the -//! hash-free index [`MilnorAlgebra::seqno`] as pure integer arithmetic over the -//! flat `g` table — the output-indexing primitive the multiply kernel needs. -//! - **Stage 3 — single-`R` multiply.** `multiply_pair` ports the per-term test -//! + output assembly of `multiply_basis_element_by_element_2`; -//! `multiply_single_r_kernel` runs one `Sq(R)·s` product per launch. -//! - **Stage 4 — batched launch.** `multiply_batch_kernel` fuses all `(R, s)` -//! products of one `get_partial_matrix` into a single launch (one thread per -//! `(product, matrix, term)` pair). The pair is decoded on-device from a prefix-sum -//! over per-product pair counts, with admissible-matrix data deduplicated by distinct -//! `R` — so a launch uploads compact per-`R`/per-product tables, not a per-pair table -//! (which at scale is gigabytes of almost-entirely-redundant data). On an RTX 3050 Ti -//! this made the batch ~2× faster than the CPU `get_partial_matrix` at stem 100 -//! (`benches/milnor_gpu_ab.rs`); the naive per-pair table was ~7× slower. +//! 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 @@ -630,7 +618,7 @@ pub struct GpuProduct { } /// Compute a whole batch of `Sq(R) · s` products in a single GPU launch — the -/// Stage 4 unit of one `get_partial_matrix` call. `R`s may differ (each contributes its +/// 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. /// diff --git a/ext/src/nassau_gpu.rs b/ext/src/nassau_gpu.rs index c63e0462dd..9bf3dfff43 100644 --- a/ext/src/nassau_gpu.rs +++ b/ext/src/nassau_gpu.rs @@ -4,8 +4,8 @@ //! 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`], which on an RTX 3050 Ti beats the -//! parallel CPU per-term sweep ~2× at stem 100 (see `benches/milnor_gpu_ab.rs`). +//! [`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 From 758901358a76890a9581d368a957e2f81bf7cc37 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Sat, 11 Jul 2026 03:16:20 -0400 Subject: [PATCH 10/10] Address CodeRabbit review: seqno-table race, u16 checks, invariants - compute_seqno_tables: guard the publish with an rcu that only replaces the cache when the new table reaches at least as far, so parallel `get_partial_matrix` builds can't shrink it out from under `seqno` (which would then index out of bounds and panic). - GPU marshalling: narrow admissible-matrix / p-part entries to u16 through a checked `narrow_u16` helper that panics loudly instead of silently wrapping if a future larger-degree run exceeds u16. - AdmissibleMatrix::new: state the non-empty-R precondition with a debug_assert instead of relying on a bare `.max().unwrap()`. - nassau_gpu_timing doc: document `--test-threads=1`, required because both ignored tests toggle the process-wide NASSAU_GPU env var. - xor_f2_on_gpu: use size_of_val for the output buffer size (clippy). Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_01Fsxqsgjoa5RWYhAvRG1bT2 --- .../algebra/src/algebra/milnor_algebra.rs | 23 ++++++++++++++----- ext/crates/algebra/src/algebra/milnor_gpu.rs | 22 ++++++++++++------ ext/tests/nassau_gpu_timing.rs | 5 +++- 3 files changed, 36 insertions(+), 14 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_algebra.rs b/ext/crates/algebra/src/algebra/milnor_algebra.rs index 492229d09b..f031d786ab 100644 --- a/ext/crates/algebra/src/algebra/milnor_algebra.rs +++ b/ext/crates/algebra/src/algebra/milnor_algebra.rs @@ -987,12 +987,19 @@ impl MilnorAlgebra { } } - self.seqno_tables - .store(Some(std::sync::Arc::new(SeqnoTables { - max_degree, - width, - g, - }))); + // 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 @@ -1507,6 +1514,10 @@ struct AdmissibleMatrix { 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() diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 40e3ec9aef..8139b1a6be 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -45,6 +45,14 @@ use crate::algebra::{ /// `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 @@ -124,8 +132,8 @@ impl Resident { mk_len: mk_len as u32, num_mats: (mk.len() / mk_len) as u32, }; - self.col_sums.extend(cs.iter().map(|&v| v as u16)); - self.masks.extend(mk.iter().map(|&v| v as u16)); + 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 } @@ -155,7 +163,7 @@ pub fn xor_f2_on_gpu(a: &[u32], b: &[u32]) -> Vec { 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(n * size_of::()); + 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; @@ -533,8 +541,8 @@ pub fn multiply_single_r_on_gpu( ); 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| v as u16).collect(); - let masks: Vec = mk32.iter().map(|&v| v as u16).collect(); + 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. @@ -548,7 +556,7 @@ pub fn multiply_single_r_on_gpu( .iter_mut() .zip(&elt.p_part) { - *slot = v as u16; + *slot = narrow_u16(v); } } @@ -686,7 +694,7 @@ pub fn multiply_batch_on_gpu( 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 = v as u16; + *slot = narrow_u16(v); } } (tp, tl) diff --git a/ext/tests/nassau_gpu_timing.rs b/ext/tests/nassau_gpu_timing.rs index 756c3d8ac9..04b83d0c92 100644 --- a/ext/tests/nassau_gpu_timing.rs +++ b/ext/tests/nassau_gpu_timing.rs @@ -1,7 +1,10 @@ //! 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`. +//! `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;