diff --git a/.gitignore b/.gitignore index 7477b249e6..d4c52a3d33 100644 --- a/.gitignore +++ b/.gitignore @@ -40,3 +40,6 @@ chart/python/*.egg-info # Claude .claude +# GPU run logs. These are multi-hundred-MB analysis artifacts written next to the crate; +# one was committed by accident and blocked a push (GitHub rejects files over 100 MB). +ext/*.log diff --git a/ext/150/nassau_differential/c/126/1 b/ext/150/nassau_differential/c/126/1 new file mode 100644 index 0000000000..79167dadeb Binary files /dev/null and b/ext/150/nassau_differential/c/126/1 differ diff --git a/ext/150/nassau_differential/c/127/0 b/ext/150/nassau_differential/c/127/0 new file mode 100644 index 0000000000..74f8fb3c8d Binary files /dev/null and b/ext/150/nassau_differential/c/127/0 differ diff --git a/ext/150/nassau_differential/c/127/1 b/ext/150/nassau_differential/c/127/1 new file mode 100644 index 0000000000..bc6869a027 Binary files /dev/null and b/ext/150/nassau_differential/c/127/1 differ diff --git a/ext/150/nassau_differential/c/128/0 b/ext/150/nassau_differential/c/128/0 new file mode 100644 index 0000000000..5c9b3950b2 Binary files /dev/null and b/ext/150/nassau_differential/c/128/0 differ diff --git a/ext/150/nassau_differential/c/128/1 b/ext/150/nassau_differential/c/128/1 new file mode 100644 index 0000000000..26fff2ee58 Binary files /dev/null and b/ext/150/nassau_differential/c/128/1 differ diff --git a/ext/150/nassau_differential/c/129/0 b/ext/150/nassau_differential/c/129/0 new file mode 100644 index 0000000000..8dcfc65f72 Binary files /dev/null and b/ext/150/nassau_differential/c/129/0 differ diff --git a/ext/150/nassau_differential/c/129/1 b/ext/150/nassau_differential/c/129/1 new file mode 100644 index 0000000000..2886988399 Binary files /dev/null and b/ext/150/nassau_differential/c/129/1 differ diff --git a/ext/150/nassau_differential/c/130/0 b/ext/150/nassau_differential/c/130/0 new file mode 100644 index 0000000000..9df80ddb5b Binary files /dev/null and b/ext/150/nassau_differential/c/130/0 differ diff --git a/ext/150/nassau_differential/c/130/1 b/ext/150/nassau_differential/c/130/1 new file mode 100644 index 0000000000..cb379032d7 Binary files /dev/null and b/ext/150/nassau_differential/c/130/1 differ diff --git a/ext/150/nassau_differential/c/131/0 b/ext/150/nassau_differential/c/131/0 new file mode 100644 index 0000000000..e9dd7430f4 Binary files /dev/null and b/ext/150/nassau_differential/c/131/0 differ diff --git a/ext/150/nassau_differential/c/131/1 b/ext/150/nassau_differential/c/131/1 new file mode 100644 index 0000000000..2d41ad01ac Binary files /dev/null and b/ext/150/nassau_differential/c/131/1 differ diff --git a/ext/150/nassau_differential/c/132/0 b/ext/150/nassau_differential/c/132/0 new file mode 100644 index 0000000000..4fe1371722 Binary files /dev/null and b/ext/150/nassau_differential/c/132/0 differ diff --git a/ext/150/nassau_differential/c/132/1 b/ext/150/nassau_differential/c/132/1 new file mode 100644 index 0000000000..383932f51c Binary files /dev/null and b/ext/150/nassau_differential/c/132/1 differ diff --git a/ext/150/nassau_differential/c/133/0 b/ext/150/nassau_differential/c/133/0 new file mode 100644 index 0000000000..f0b3eeb080 Binary files /dev/null and b/ext/150/nassau_differential/c/133/0 differ diff --git a/ext/150/nassau_differential/c/133/1 b/ext/150/nassau_differential/c/133/1 new file mode 100644 index 0000000000..a41bc9e030 Binary files /dev/null and b/ext/150/nassau_differential/c/133/1 differ diff --git a/ext/150/nassau_differential/zarr.json b/ext/150/nassau_differential/zarr.json new file mode 100644 index 0000000000..fc55f3183c --- /dev/null +++ b/ext/150/nassau_differential/zarr.json @@ -0,0 +1,63 @@ +{ + "zarr_format": 3, + "node_type": "array", + "shape": [ + 4096, + 1024 + ], + "data_type": "bytes", + "chunk_grid": { + "name": "regular", + "configuration": { + "chunk_shape": [ + 8, + 8 + ] + } + }, + "chunk_key_encoding": { + "name": "default", + "configuration": { + "separator": "/" + } + }, + "fill_value": [], + "codecs": [ + { + "name": "sharding_indexed", + "configuration": { + "chunk_shape": [ + 1, + 1 + ], + "codecs": [ + { + "name": "vlen-bytes" + }, + { + "name": "crc32c" + } + ], + "index_codecs": [ + { + "name": "bytes", + "configuration": { + "endian": "little" + } + }, + { + "name": "crc32c" + } + ], + "index_location": "end" + } + } + ], + "attributes": { + "_zarrs": { + "description": "This array was created with zarrs", + "repository": "https://github.com/zarrs/zarrs", + "version": "0.23.13" + } + } +} \ No newline at end of file diff --git a/ext/150/zarr.json b/ext/150/zarr.json new file mode 100644 index 0000000000..eaf3d8a4e5 --- /dev/null +++ b/ext/150/zarr.json @@ -0,0 +1,18 @@ +{ + "zarr_format": 3, + "node_type": "group", + "attributes": { + "algebra_magic": 163840, + "prime": 2, + "algebra_prefix": "milnor", + "module_spec": { + "p": 2, + "type": "finite dimensional module", + "gens": { + "x0": 0 + }, + "actions": [] + }, + "complex_name": "S_2" + } +} \ No newline at end of file diff --git a/ext/crates/algebra/Cargo.toml b/ext/crates/algebra/Cargo.toml index 09031e4c4b..5845f309a9 100644 --- a/ext/crates/algebra/Cargo.toml +++ b/ext/crates/algebra/Cargo.toml @@ -33,13 +33,20 @@ enum_dispatch = "0.3.13" # `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 = [ +cubecl = { git = "https://github.com/tracel-ai/cubecl", tag = "v0.11.0-pre.1", 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 } +cubecl-common = { git = "https://github.com/tracel-ai/cubecl", tag = "v0.11.0-pre.1", optional = true } +# Spans around the GPU submission (see `gpu_thread`), so a worker blocked waiting for the device +# is visible in the log instead of silent — the stalls it diagnoses emitted nothing at all. +tracing = { version = "0.1.41", optional = true } +# Multi-consumer job queue for the per-device GPU workers (see `gpu_thread`). std's mpsc is +# single-consumer, so sharing one queue across devices there means a `Mutex`; this is an +# actual MPMC queue, which is what the dispatch wants. +crossbeam-channel = { version = "0.5", optional = true } [dev-dependencies] criterion = { version = "0.5", features = ["html_reports"] } @@ -50,8 +57,12 @@ rstest = "0.25.0" [features] default = ["odd-primes"] cache-multiplication = [] +# An arithmetic replacement for the Milnor basis index map. Off by default and not wired in: +# adopting it would renumber the basis and invalidate saved resolutions. See +# `algebra::milnor_rank`. +milnor-rank = [] concurrent = ["fp/concurrent", "maybe-rayon/concurrent"] -gpu = ["dep:cubecl", "dep:cubecl-common"] +gpu = ["dep:cubecl", "dep:cubecl-common", "dep:tracing", "dep:crossbeam-channel"] odd-primes = ["fp/odd-primes"] [[bench]] @@ -70,6 +81,17 @@ harness = false name = "nassau_milnor" harness = false +# GPU batched-multiply throughput / cubecl-regression bench (see benches/nassau_milnor_gpu.rs). +# Requires `--features gpu`; without it the target compiles to a no-op main. +[[bench]] +name = "nassau_milnor_gpu" +harness = false + [[bench]] name = "seqno" harness = false + +[[bench]] +name = "milnor_rank" +harness = false +required-features = ["milnor-rank"] diff --git a/ext/crates/algebra/benches/milnor.rs b/ext/crates/algebra/benches/milnor.rs index 6a38188e59..696bb7325e 100644 --- a/ext/crates/algebra/benches/milnor.rs +++ b/ext/crates/algebra/benches/milnor.rs @@ -1,6 +1,6 @@ //! Benchmarks for the low-level Milnor `PPartMultiplier` kernel. -use algebra::milnor_algebra::{PPartAllocation, PPartEntry, PPartMultiplier}; +use algebra::milnor_algebra::{PPart, PPartAllocation, PPartMultiplier}; use criterion::{ BenchmarkGroup, Criterion, criterion_group, criterion_main, measurement::WallTime, }; @@ -11,8 +11,8 @@ fn bench_ppart( g: &mut BenchmarkGroup, name: &str, p: u32, - r: Vec, - s: Vec, + r: PPart, + s: PPart, ) { let p = ValidPrime::new(p); g.bench_function(name, |bench| { @@ -21,7 +21,7 @@ fn bench_ppart( bench.iter_batched( PPartAllocation::default, |alloc| { - let m = PPartMultiplier::::new_from_allocation(p, &r, &s, alloc, 0, 0); + let m = PPartMultiplier::::new_from_allocation(p, r, s, alloc, 0, 0); for c in m { std::hint::black_box(c); } @@ -38,30 +38,30 @@ fn ppart(c: &mut Criterion) { &mut g, "ppart_2/a", 2, - vec![60, 30, 8, 2, 1], - vec![20, 30, 20, 4, 1, 2], + PPart::from_slice(&[60, 30, 8, 2, 1]), + PPart::from_slice(&[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], + PPart::from_slice(&[35, 12, 20, 14, 1, 3]), + PPart::from_slice(&[60, 30, 0, 2, 1]), ); bench_ppart::( &mut g, "ppart_4/a", 2, - vec![60, 30, 8, 2, 1], - vec![20, 30, 20, 4, 1, 2], + PPart::from_slice(&[60, 30, 8, 2, 1]), + PPart::from_slice(&[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], + PPart::from_slice(&[35, 12, 20, 14, 1, 3]), + PPart::from_slice(&[60, 30, 0, 2, 1]), ); #[cfg(feature = "odd-primes")] @@ -70,15 +70,15 @@ fn ppart(c: &mut Criterion) { &mut g, "ppart_3/a", 3, - vec![120, 70, 40, 2], - vec![60, 35, 21, 6], + PPart::from_slice(&[120, 70, 40, 2]), + PPart::from_slice(&[60, 35, 21, 6]), ); bench_ppart::( &mut g, "ppart_3/b", 3, - vec![30, 12, 35, 24], - vec![100, 80, 16, 2, 3], + PPart::from_slice(&[30, 12, 35, 24]), + PPart::from_slice(&[100, 80, 16, 2, 3]), ); } diff --git a/ext/crates/algebra/benches/milnor_rank.rs b/ext/crates/algebra/benches/milnor_rank.rs new file mode 100644 index 0000000000..bd063f2017 --- /dev/null +++ b/ext/crates/algebra/benches/milnor_rank.rs @@ -0,0 +1,100 @@ +//! Compares [`PPartRanker`] against the hash map lookup it would replace. +//! +//! `MilnorAlgebra::basis_element_to_index` is called once per term of every product, so it is one +//! of the hottest operations in a resolution. It is currently a hash map from the (packed) basis +//! element to its position in `basis_table`. The ranker computes that position arithmetically +//! instead, from a table that covers every degree at once. +//! +//! The two are compared on the same workload: recover the index of every basis element of a +//! degree. Note that they do not agree on *which* index — see [`PPartRanker`] — so this measures +//! the cost of the two strategies, not a drop-in substitution. +//! +//! Each is measured under two access orders, because the choice flatters the map: +//! +//! - **sequential** — sweep the basis in order. This is the map's insertion order, so every probe +//! walks memory linearly and prefetches perfectly. Flattering, and not what callers do. +//! - **scattered** — the same elements in a fixed pseudo-random permutation. This is closer to +//! real use, where `basis_element_to_index` is called on multiplication *outputs*, which arrive +//! in no particular order. It matters because the two structures scale differently: the map +//! stores an entry per basis element and leaves cache as the basis grows (~60 KiB in degree 120 +//! alone), whereas the ranker's table is a few KiB covering every degree at once. +//! +//! [`PPartRanker`]: algebra::milnor_rank::PPartRanker + +use std::hint::black_box; + +use algebra::{Algebra, MilnorAlgebra, milnor_rank::PPartRanker}; +use criterion::{Criterion, Throughput, criterion_group, criterion_main}; +use fp::prime::TWO; +use pprof::criterion::{Output, PProfProfiler}; + +/// Degrees to sweep. +/// +/// The range matters more than it looks, because the two structures live in different parts of the +/// memory hierarchy and the crossover is inside this range. A lookup probes only its own degree's +/// map, which is ~0.1 MB in degree 120 (L2-resident) but ~3 MB in degree 300 and ~12 MB in degree +/// 400 — well past L3, so every probe is a DRAM miss. The ranker's table is ~35 KB for *all* +/// degrees and stays in L1 throughout. Measuring only the small degrees answers a question nobody +/// is asking; the large ones are where expanding the algebra actually hurts. +const DEGREES: &[i32] = &[120, 300, 400, 500]; + +/// A fixed permutation of `0..n`, from a Fisher-Yates shuffle driven by a small LCG. Deterministic +/// so the two variants see exactly the same access order. +fn scattered(n: usize) -> Vec { + let mut order: Vec = (0..n).collect(); + let mut state = 0x2545_f491_4f6c_dd1d_u64; + for i in (1..n).rev() { + state = state + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + order.swap(i, (state >> 33) as usize % (i + 1)); + } + order +} + +fn milnor_rank(c: &mut Criterion) { + let algebra = MilnorAlgebra::new(TWO, false); + let max_degree = *DEGREES.iter().max().unwrap(); + algebra.compute_basis(max_degree); + let ranker = PPartRanker::new(TWO, max_degree); + + let mut g = c.benchmark_group("milnor_rank"); + for °ree in DEGREES { + let dim = algebra.dimension(degree); + g.throughput(Throughput::Elements(dim as u64)); + + // Collect the elements once so neither variant pays for the table walk itself. + let elements: Vec<_> = (0..dim) + .map(|i| algebra.basis_element_from_index(degree, i)) + .collect(); + let shuffled: Vec<_> = scattered(dim).into_iter().map(|i| elements[i]).collect(); + + for (order, elements) in [("seq", &elements), ("scattered", &shuffled)] { + g.bench_function(format!("hashmap_{order}/deg{degree}"), |b| { + b.iter(|| { + for elt in elements { + black_box(algebra.basis_element_to_index(elt)); + } + }); + }); + + g.bench_function(format!("ranker_{order}/deg{degree}"), |b| { + b.iter(|| { + for elt in elements { + black_box(ranker.rank(elt.p_part, degree)); + } + }); + }); + } + } + 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 = milnor_rank +} +criterion_main!(benches); diff --git a/ext/crates/algebra/benches/nassau_milnor_gpu.rs b/ext/crates/algebra/benches/nassau_milnor_gpu.rs new file mode 100644 index 0000000000..a8d0d3fb47 --- /dev/null +++ b/ext/crates/algebra/benches/nassau_milnor_gpu.rs @@ -0,0 +1,122 @@ +//! GPU counterpart to `nassau_milnor.rs`: benchmarks the **batched Milnor multiply on the GPU** +//! ([`algebra::milnor_gpu::multiply_batch_on_gpu`]) — the kernel + resident-master + cubecl-allocator +//! path Nassau's `S_2` resolution drives, WITHOUT the surrounding row-reduction or resolution +//! bookkeeping. This is the "hammer the GPU with multiplications" harness. +//! +//! Why a bench (not just an example): +//! - **Perf + regression tracking.** `cargo bench --bench nassau_milnor_gpu -- --save-baseline pre` +//! then `--baseline pre` across cubecl commits measures the multiply-throughput delta directly — +//! the "how much does the cubecl backend cost us" number, and a guard against kernel regressions. +//! - **Isolation.** A crash here indicts the multiply / cubecl allocator alone (not RREF, which runs +//! on the separate `fp-cuda` runtime, nor Nassau bookkeeping). +//! +//! The output-degree sweep doubles as the memory axis: larger `out_degree` → larger resident master +//! (the shared segmented device buffer warms once and is reused across iterations, exactly as in a +//! real resolution). Requires `--features gpu` and `CUDA_PATH` (see `ext/gpu_prep`); without `gpu` +//! it compiles to a no-op `main`. +//! +//! Scope note: this is the fixed-scale THROUGHPUT bench. The ~100-stream concurrency + unbounded +//! resident-master GROWTH soak that reproduces the cubecl uninit-handle crash +//! (tracel-ai/cubecl#1401) belongs in a `#[test]`, not here — a criterion measurement loop is the +//! wrong shape for a memory-growth soak with correctness assertions. + +#[cfg(feature = "gpu")] +mod gpu { + use algebra::{ + Algebra, MilnorAlgebra, + milnor_gpu::{GpuProduct, multiply_batch_on_gpu}, + }; + use criterion::{Criterion, Throughput, black_box, criterion_group}; + use fp::prime::TWO; + + /// Output degrees to sweep — the cost/memory axis. Chosen around Nassau's hot band + /// (out ≈ 40–52; see `nassau_milnor.rs`'s `REGIME`), plus a cheap and an expensive anchor. + const OUT_DEGREES: &[i32] = &[24, 32, 40, 48]; + /// Output rows per batch. Products round-robin across rows so each launch fills a real matrix + /// rather than a single-row strip. + const NUM_ROWS: usize = 32; + + /// Build one batched `get_partial_matrix`-shaped build at `out_degree`: every non-empty `R` of + /// degree `1..out_degree` times a dense complementary element, round-robin across `NUM_ROWS` + /// rows, single generator block (`out_offset = 0`, `num_cols = dim(out_degree)`). Mirrors the + /// construction in `multiply_batch_matches_reference`, so it hits the same kernel path Nassau does. + fn build_batch(algebra: &MilnorAlgebra, out_degree: i32) -> (usize, Vec) { + let num_cols = algebra.dimension(out_degree); + 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, + }); + } + } + (num_cols, products) + } + + pub fn nassau_milnor_gpu(c: &mut Criterion) { + // Exactly the algebra Nassau uses: the full Milnor algebra at p=2, stable (not unstable). + use std::sync::Arc; + let algebra = Arc::new(MilnorAlgebra::new(TWO, false)); + let mut g = c.benchmark_group("nassau_milnor_gpu"); + + for &out_degree in OUT_DEGREES { + // `compute_basis` is cumulative; seqno tables are what the GPU path indexes by. + algebra.compute_basis(out_degree); + algebra.compute_seqno_tables(out_degree); + let (num_cols, products) = build_batch(&algebra, out_degree); + if products.is_empty() || num_cols == 0 { + continue; + } + + // One "element" = one `Sq(R)·s` product fused into the launch. + g.throughput(Throughput::Elements(products.len() as u64)); + g.bench_function(format!("multiply_batch/out{out_degree}"), |b| { + b.iter(|| { + black_box(multiply_batch_on_gpu( + &algebra, + num_cols, + NUM_ROWS, + black_box(&products), + )); + }); + }); + } + + g.finish(); + } + + criterion_group! { + name = benches; + config = Criterion::default() + .measurement_time(std::time::Duration::from_secs(5)) + .sample_size(30); + targets = nassau_milnor_gpu + } +} + +#[cfg(feature = "gpu")] +criterion::criterion_main!(gpu::benches); + +#[cfg(not(feature = "gpu"))] +fn main() { + eprintln!("nassau_milnor_gpu bench requires --features gpu (and CUDA_PATH; see ext/gpu_prep)"); +} diff --git a/ext/crates/algebra/benches/seqno.rs b/ext/crates/algebra/benches/seqno.rs index 22ad312edf..285617ff6d 100644 --- a/ext/crates/algebra/benches/seqno.rs +++ b/ext/crates/algebra/benches/seqno.rs @@ -51,7 +51,7 @@ fn seqno(c: &mut Criterion) { g.bench_function(format!("basis_to_index/deg{degree}"), |b| { b.iter(|| { for elt in &basis { - black_box(algebra.seqno(&elt.p_part)); + black_box(algebra.seqno(elt.p_part)); } }); }); diff --git a/ext/crates/algebra/src/algebra/milnor_algebra.rs b/ext/crates/algebra/src/algebra/milnor_algebra.rs index d07ba273e4..a3abf12de4 100644 --- a/ext/crates/algebra/src/algebra/milnor_algebra.rs +++ b/ext/crates/algebra/src/algebra/milnor_algebra.rs @@ -24,8 +24,11 @@ pub struct MilnorProfile { #[serde(default = "q_part_default")] pub q_part: u32, /// The profile function for the Q part. + /// + /// Unlike the exponent sequence of a basis element (see [`PPart`]), these are *exponents* of + /// the profile function and use [`PPartEntry::MAX`] to mean infinity, so they stay unpacked. #[serde(default)] - pub p_part: PPart, + pub p_part: Vec, } impl MilnorProfile { @@ -99,9 +102,235 @@ impl Default for MilnorProfile { } pub type PPartEntry = u32; -pub type PPart = Vec; -#[derive(Debug, Clone, Default)] +/// The exponent sequence $(r_1, r_2, \ldots)$ of a Milnor basis element $P(r_1, r_2, \ldots)$, +/// bit-packed into a single `u64`. +/// +/// Entry $r_{i+1}$ occupies `WIDTHS[i]` bits starting at bit `SHIFTS[i]` (both private). The +/// widths are forced by the degree bound: at $p = 2$ the internal degree of $P(R)$ is +/// $\sum_i r_i (2^i - 1)$ and every term is non-negative, so an element of degree at most +/// [`Self::MAX_DEGREE`] has $r_i \le \mathrm{MAX\\_DEGREE}/(2^i - 1)$. At an odd prime the same +/// argument bounds $r_i$ by that quantity divided by $q = 2(p-1)$, so the $p = 2$ widths are valid +/// for every prime and this type is prime-agnostic. +/// +/// Trailing zeros are not represented: $P(2, 1)$ and $P(2, 1, 0)$ have the same packed value. That +/// is what makes the packed value a canonical key, and it makes [`Self::len`] the position of the +/// highest non-zero entry rather than a stored field. +/// +/// # Invariant +/// +/// Every entry fits in its field. This holds for any element of degree at most +/// [`Self::MAX_DEGREE`], which [`MilnorAlgebra::compute_basis`] enforces up front, so the packing +/// can never silently truncate. [`Self::set`] asserts it anyway, and [`Self::try_from_slice`] +/// reports failure instead of panicking for input that has not been through that gate. +#[derive(Clone, Copy, PartialEq, Eq, Hash, Default)] +pub struct PPart(u64); + +impl PPart { + /// `FIELD_OF_BIT[b]` is the index of the entry owning bit `b`, letting [`Self::len`] turn a + /// `leading_zeros` into an entry index without looping. + const FIELD_OF_BIT: [u8; 64] = { + let mut table = [0; 64]; + let mut i = 0; + while i < Self::MAX_LEN { + let mut b = Self::SHIFTS[i]; + while b < Self::SHIFTS[i + 1] { + table[b as usize] = i as u8; + b += 1; + } + i += 1; + } + table + }; + /// The largest internal degree whose exponent sequences are guaranteed to fit. + /// + /// This is the largest bound for which the field widths sum to at most 64. It is far beyond + /// anything reachable — the Milnor algebra already has over 5 million basis elements below + /// degree 300 — and exceeds the degree 1536 the previous hand-rolled packing assumed. + pub const MAX_DEGREE: i32 = 2045; + /// The number of entries that can be stored. This equals `fp`'s `MAX_MULTINOMIAL_LEN`, which + /// already bounds the length of the $\xi$-degree table, so it is not a new restriction. + pub const MAX_LEN: usize = 10; + /// `SHIFTS[i]` is the bit offset of entry `i`; `SHIFTS[MAX_LEN]` is the total width, 64. + const SHIFTS: [u32; Self::TABLE_LEN] = { + let mut shifts = [0; Self::TABLE_LEN]; + let mut i = 0; + while i < Self::MAX_LEN { + shifts[i + 1] = shifts[i] + Self::WIDTHS[i]; + i += 1; + } + shifts + }; + /// The length of the layout tables. This is [`Self::MAX_LEN`] rounded up to a power of two so + /// that [`Self::entry`] can mask its index instead of bounds-checking it; entries at or past + /// `MAX_LEN` are given width 0, so they read as zero. + const TABLE_LEN: usize = 16; + /// `WIDTHS[i]` is the number of bits holding $r_{i+1}$: the number of bits needed to represent + /// `MAX_DEGREE / (2^(i+1) - 1)`. + const WIDTHS: [u32; Self::TABLE_LEN] = [11, 10, 9, 8, 7, 6, 5, 4, 3, 1, 0, 0, 0, 0, 0, 0]; + + /// The largest value entry `i` can hold. + pub const fn max_entry(i: usize) -> PPartEntry { + ((1u64 << Self::WIDTHS[i]) - 1) as PPartEntry + } + + /// The number of bits holding entry `i`. Together with [`Self::shift`] this lets callers build + /// a mask over [`Self::bits`] directly, e.g. to test many entries in one comparison. + pub const fn width(i: usize) -> u32 { + Self::WIDTHS[i] + } + + /// The bit offset of entry `i` within [`Self::bits`]. + pub const fn shift(i: usize) -> u32 { + Self::SHIFTS[i] + } + + const fn mask(i: usize) -> u64 { + ((1u64 << Self::WIDTHS[i]) - 1) << Self::SHIFTS[i] + } + + pub const fn zero() -> Self { + Self(0) + } + + /// The raw packed value. Two exponent sequences are equal exactly when their bits are, so this + /// is a complete hash key, and it can be compared against a packed mask in one operation (see + /// `MilnorSubalgebra::packed_signature` in `ext`). + /// + /// The layout is not uniform, so this is a complete key but not a balanced one: entry `i` sits + /// at [`Self::SHIFTS`]`[i]`, putting `r_1` in the low bits, and `r_1` correlates strongly with + /// internal degree. Taking this value modulo a small number partitions by `r_1` rather than + /// evenly — mix the bits first (see `milnor_gpu::shard_of`). + pub const fn bits(self) -> u64 { + self.0 + } + + /// Reinterpret a raw packed value. + /// + /// Callers that assemble entries by shifting must uphold the type invariant themselves: each + /// entry must lie within its field, which holds for any element of degree at most + /// [`Self::MAX_DEGREE`]. This exists so hot loops can accumulate into a plain `u64` and store + /// once, rather than read-modify-write through [`Self::set`] per entry. + pub(crate) const fn from_bits(bits: u64) -> Self { + Self(bits) + } + + /// Entry `i`, for `i < TABLE_LEN`, with no bounds check. + /// + /// Masking the index keeps the table lookups in range without a branch. Entries in + /// `MAX_LEN..TABLE_LEN` have width 0 and so read as zero, which is the right answer; an index + /// at or beyond `TABLE_LEN` would silently wrap, which is why this is private and + /// `debug_assert`ed. Callers in the multiplier are all bounded by `MAX_LEN`. + #[inline] + const fn entry(self, i: usize) -> PPartEntry { + debug_assert!(i < Self::TABLE_LEN); + let i = i & (Self::TABLE_LEN - 1); + ((self.0 >> Self::SHIFTS[i]) & ((1 << Self::WIDTHS[i]) - 1)) as PPartEntry + } + + /// Entry `i`, or 0 if `i` is past the end. Accepts any index. + #[inline] + pub const fn get(self, i: usize) -> PPartEntry { + if i >= Self::MAX_LEN { + return 0; + } + self.entry(i) + } + + /// Set entry `i` to `v`. + /// + /// # Panics + /// + /// If `i >= MAX_LEN`, or `v` does not fit in entry `i`. Both are unreachable for elements of + /// degree at most [`Self::MAX_DEGREE`]. + #[inline] + pub fn set(&mut self, i: usize, v: PPartEntry) { + assert!(i < Self::MAX_LEN, "p-part index {i} out of range"); + assert!( + v <= Self::max_entry(i), + "p-part entry {v} does not fit in the {} bits at index {i}", + Self::WIDTHS[i], + ); + self.0 = (self.0 & !Self::mask(i)) | ((v as u64) << Self::SHIFTS[i]); + } + + /// The number of entries up to and including the last non-zero one. + #[inline] + pub const fn len(self) -> usize { + if self.0 == 0 { + 0 + } else { + Self::FIELD_OF_BIT[63 - self.0.leading_zeros() as usize] as usize + 1 + } + } + + #[inline] + pub const fn is_empty(self) -> bool { + self.0 == 0 + } + + /// Zero every entry from `n` onwards, i.e. the packed form of `self[..n]`. + #[inline] + pub const fn truncate(self, n: usize) -> Self { + if n >= Self::MAX_LEN { + self + } else { + Self(self.0 & ((1 << Self::SHIFTS[n]) - 1)) + } + } + + pub fn iter(self) -> impl DoubleEndedIterator + ExactSizeIterator { + (0..self.len()).map(move |i| self.get(i)) + } + + /// Pack `entries`, returning `None` if they do not fit. Use this for anything derived from + /// user input; use [`Self::from_slice`] when the degree bound already guarantees a fit. + pub fn try_from_slice(entries: &[PPartEntry]) -> Option { + let mut result = Self::zero(); + for (i, &entry) in entries.iter().enumerate() { + // A zero past the end is just padding, which the packed form drops anyway. + if entry == 0 { + continue; + } + if i >= Self::MAX_LEN || entry > Self::max_entry(i) { + return None; + } + result.set(i, entry); + } + Some(result) + } + + /// Pack `entries`, panicking if they do not fit. + pub fn from_slice(entries: &[PPartEntry]) -> Self { + Self::try_from_slice(entries).unwrap_or_else(|| { + panic!( + "p-part {entries:?} exceeds the degree {} bound", + Self::MAX_DEGREE + ) + }) + } +} + +impl FromIterator for PPart { + fn from_iter>(iter: I) -> Self { + let mut result = Self::zero(); + for (i, entry) in iter.into_iter().enumerate() { + if entry != 0 { + result.set(i, entry); + } + } + result + } +} + +impl std::fmt::Debug for PPart { + fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + f.debug_list().entries(self.iter()).finish() + } +} + +/// A Milnor basis element. This is `Copy` and entirely inline: 16 bytes, no heap. +#[derive(Debug, Clone, Copy, Default)] pub struct MilnorBasisElement { pub q_part: u32, pub p_part: PPart, @@ -126,10 +355,7 @@ impl MilnorBasisElement { } pub fn clone_into(&self, other: &mut Self) { - other.q_part = self.q_part; - other.degree = self.degree; - other.p_part.clear(); - other.p_part.extend_from_slice(&self.p_part); + *other = *self; } /// Update the degree component to the correct degree @@ -138,8 +364,8 @@ impl MilnorBasisElement { let xi_degrees = combinatorics::xi_degrees(p); let tau_degrees = combinatorics::tau_degrees(p); - self.degree = q * std::iter::zip(xi_degrees, &self.p_part) - .map(|(&a, &b)| a * b as i32) + self.degree = q * std::iter::zip(xi_degrees, self.p_part.iter()) + .map(|(&a, b)| a * b as i32) .sum::() + BitflagIterator::set_bit_iterator(self.q_part as u64) .map(|k| tau_degrees[k]) @@ -161,7 +387,9 @@ impl std::cmp::Eq for MilnorBasisElement {} impl std::hash::Hash for MilnorBasisElement { fn hash(&self, state: &mut H) { - self.p_part.hash(state); + // The p-part is a single `u64`, so this is one hasher round rather than a pointer chase + // plus a variable-length slice hash. + self.p_part.bits().hash(state); #[cfg(feature = "odd-primes")] self.q_part.hash(state); } @@ -189,62 +417,13 @@ impl std::fmt::Display for MilnorBasisElement { } } -/// A version of `HashMap` that is more efficient at the prime 2. -#[cfg(feature = "odd-primes")] +/// A map from the basis elements of a single degree to their indices. +/// +/// [`MilnorBasisElement`] hashes and compares on its p-part (and, at odd primes, its q-part), both +/// of which are now single machine words, so a plain `HashMap` is already the specialised form +/// this used to hand-roll for `p = 2`. type MilnorHashMap = HashMap; -#[cfg(not(feature = "odd-primes"))] -struct MilnorHashMap { - degree: i32, - inner: HashMap, -} - -#[cfg(not(feature = "odd-primes"))] -impl Default for MilnorHashMap { - fn default() -> Self { - Self { - degree: -1, - inner: HashMap::default(), - } - } -} - -#[cfg(not(feature = "odd-primes"))] -impl MilnorHashMap { - /// Encode a [`MilnorBasisElement`] of a known degree into a `u64`. This is achieved by packing - /// the PPart into a single `u64`, where we omit the first entry since it can be derived from - /// the degree. This currently supports elements up to degree 2^9 * 3 = 1536. - fn code(x: &MilnorBasisElement) -> u64 { - let mut counter = 0; - let mut shift = 0; - for (idx, &entry) in x.p_part.iter().skip(1).enumerate() { - counter += (entry as u64) << shift; - shift += 9 - idx; - } - counter - } - - fn reserve(&mut self, additional: usize) { - self.inner.reserve(additional); - } - - fn insert(&mut self, k: MilnorBasisElement, v: V) { - if self.degree == -1 { - self.degree = k.degree; - } - assert_eq!(k.degree, self.degree); - assert!( - self.inner.insert(Self::code(&k), v).is_none(), - "Duplicate entry for {k}" - ); - } - - fn get(&self, k: &MilnorBasisElement) -> Option<&V> { - assert_eq!(k.degree, self.degree); - self.inner.get(&Self::code(k)) - } -} - /// 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 @@ -267,7 +446,14 @@ pub struct MilnorAlgebra { /// degree `q * i`. ppart_table: OnceVec>, - /// A list of all basis elements of each degree, constructed from [`Self::ppart_table`] + /// A list of all basis elements of each degree, constructed from [`Self::ppart_table`]. + /// + /// Only populated when [`Self::stores_basis_table`] holds. At `p = 2` with unstable support + /// off, the basis element at index `i` of degree `t` is exactly + /// `MilnorBasisElement::from_p(ppart_table[t][i], t)`, so storing it repeats the p-part with a + /// known q-part and degree bolted on -- 16 bytes per element, about a quarter of the algebra's + /// footprint. [`Self::basis_element_from_index`] reconstructs it instead, which is free now + /// that the type is `Copy` and fits in registers. basis_table: OnceVec>, excess_table: OnceVec>, @@ -343,8 +529,21 @@ impl MilnorAlgebra { &self.profile } - pub fn basis_element_from_index(&self, degree: i32, idx: usize) -> &MilnorBasisElement { - &self.basis_table[degree as usize][idx] + /// Whether the basis of each degree has to be stored rather than derived. + /// + /// At odd primes the q-part varies within a degree, and with unstable support enabled the + /// basis is re-sorted by excess; in both cases the basis is not a re-wrapping of + /// [`Self::ppart_table`] and must be kept. + fn stores_basis_table(&self) -> bool { + self.generic() || self.unstable_enabled + } + + pub fn basis_element_from_index(&self, degree: i32, idx: usize) -> MilnorBasisElement { + if self.stores_basis_table() { + self.basis_table[degree as usize][idx] + } else { + MilnorBasisElement::from_p(self.ppart_table[degree as usize][idx], degree) + } } pub fn try_basis_element_to_index(&self, elt: &MilnorBasisElement) -> Option { @@ -396,7 +595,7 @@ impl Algebra for MilnorAlgebra { MilnorBasisElement { degree: 1, q_part: 1, - p_part: vec![], + p_part: PPart::zero(), }, )); } @@ -408,7 +607,7 @@ impl Algebra for MilnorAlgebra { MilnorBasisElement { degree: (2 * self.prime() - 2) as i32, q_part: 0, - p_part: vec![1], + p_part: PPart::from_iter([1]), }, )); } @@ -427,7 +626,7 @@ impl Algebra for MilnorAlgebra { MilnorBasisElement { degree, q_part: 0, - p_part: vec![1 << i], + p_part: PPart::from_iter([1 << i]), }, )); } @@ -442,6 +641,13 @@ impl Algebra for MilnorAlgebra { } fn compute_basis(&self, max_degree: i32) { + // This is the single gate that makes [`PPart`]'s packing safe: past this degree an + // exponent could outgrow its field. Everything downstream may then assume entries fit. + assert!( + max_degree <= PPart::MAX_DEGREE, + "Milnor basis elements are only supported up to degree {}, got {max_degree}", + PPart::MAX_DEGREE, + ); self.compute_ppart(max_degree); if self.generic() { @@ -457,11 +663,12 @@ impl Algebra for MilnorAlgebra { // Populate hash map self.basis_element_to_index_map .extend(max_degree as usize, |d| { - let basis = &self.basis_table[d]; let mut map = MilnorHashMap::default(); - map.reserve(basis.len()); - for (i, b) in basis.iter().enumerate() { - map.insert(b.clone(), i); + let dim = self.dimension(d as i32); + map.reserve(dim); + for i in 0..dim { + let b = self.basis_element_from_index(d as i32, i); + assert!(map.insert(b, i).is_none(), "Duplicate entry for {b}"); } map }); @@ -482,8 +689,8 @@ impl Algebra for MilnorAlgebra { self.multiply( res.as_slice_mut(), 1, - &self.basis_table[d][i], - &self.basis_table[e][j], + self.basis_element_from_index(d as i32, i), + self.basis_element_from_index(e as i32, j), ); res }) @@ -503,7 +710,11 @@ impl Algebra for MilnorAlgebra { if degree < 0 { return 0; } - self.basis_table[degree as usize].len() + if self.stores_basis_table() { + self.basis_table[degree as usize].len() + } else { + self.ppart_table[degree as usize].len() + } } #[cfg(not(feature = "cache-multiplication"))] @@ -621,19 +832,34 @@ impl Algebra for MilnorAlgebra { map(char('1'), |_| Some((0, 0))), map(char('b'), |_| Some((1, 0))), map(preceded(p_or_sq, digits), |i| self.try_beps_pn(0, i)), - map((tag("P^"), digits, char('_'), digits), |(_, s, _, t)| { - let entry = p.pow(s); - let degree = entry as i32 * self.q() * combinatorics::xi_degrees(p)[t]; - let mut elt = MilnorBasisElement { - degree, - q_part: 0, - p_part: vec![0; t], - }; - elt.p_part[t - 1] = entry as PPartEntry; - self.compute_basis(degree); - self.try_basis_element_to_index(&elt) - .map(|idx| (degree, idx)) - }), + map( + (tag("P^"), digits, char('_'), digits::), + |(_, s, _, t)| { + if t == 0 || t >= combinatorics::xi_degrees(p).len() { + return None; + } + let entry: PPartEntry = p.as_u32().checked_pow(s)?; + if entry > PPart::max_entry(t - 1) { + return None; + } + let degree = (entry as i32) + .checked_mul(self.q())? + .checked_mul(combinatorics::xi_degrees(p)[t])?; + if degree > PPart::MAX_DEGREE { + return None; + } + let mut p_part = PPart::zero(); + p_part.set(t - 1, entry); + let elt = MilnorBasisElement { + degree, + q_part: 0, + p_part, + }; + self.compute_basis(degree); + self.try_basis_element_to_index(&elt) + .map(|idx| (degree, idx)) + }, + ), map( ( many0(preceded(tag("Q_"), digits::)), @@ -644,12 +870,16 @@ impl Algebra for MilnorAlgebra { ), |(q_list, p_list)| { let q_part = q_list.into_iter().fold(0, |acc, q| acc + (1 << q)); + let p_part = PPart::try_from_slice(&p_list.unwrap_or_default())?; let mut elt = MilnorBasisElement { degree: 0, q_part, - p_part: p_list.unwrap_or_default(), + p_part, }; elt.compute_degree(p); + if elt.degree > PPart::MAX_DEGREE { + return None; + } self.compute_basis(elt.degree); self.try_basis_element_to_index(&elt) @@ -673,7 +903,7 @@ impl UnstableAlgebra for MilnorAlgebra { } else if excess < degree { self.excess_table[degree as usize][excess as usize] } else { - self.basis_table[degree as usize].len() + self.dimension(degree) } } @@ -751,7 +981,7 @@ impl GeneratedAlgebra for MilnorAlgebra { return vec![self.basis_element_to_index(&MilnorBasisElement { degree, q_part, - p_part: vec![], + p_part: PPart::zero(), })]; } } @@ -770,7 +1000,7 @@ impl GeneratedAlgebra for MilnorAlgebra { return vec![self.basis_element_to_index(&MilnorBasisElement { degree, q_part: 0, - p_part: vec![(degree as u32 / q) as PPartEntry], + p_part: PPart::from_iter([(degree as u32 / q) as PPartEntry]), })]; } vec![] @@ -789,8 +1019,8 @@ impl GeneratedAlgebra for MilnorAlgebra { if self.profile.get_p_part(j as usize - 1) <= k as PPartEntry { return vec![]; } - let mut p_part = vec![0; j as usize]; - p_part[j as usize - 1] = p.pow(k) as PPartEntry; + let mut p_part = PPart::zero(); + p_part.set(j as usize - 1, p.pow(k) as PPartEntry); return vec![self.basis_element_to_index(&MilnorBasisElement { degree, q_part: 0, @@ -869,7 +1099,7 @@ impl GeneratedAlgebra for MilnorAlgebra { // Compute basis functions impl MilnorAlgebra { fn compute_ppart(&self, max_degree: i32) { - self.ppart_table.extend(0, |_| vec![Vec::new()]); + self.ppart_table.extend(0, |_| vec![PPart::zero()]); let p = self.prime().as_i32(); let q = if p == 2 { 1 } else { 2 * p - 2 }; @@ -899,20 +1129,19 @@ impl MilnorAlgebra { } let rem = (d - xi_degrees[i]) as usize; - for old in &self.ppart_table[rem] { + for &old in &self.ppart_table[rem] { // ppart_table[rem] is arranged in increasing order of highest // xi_i. If we get something too large, we may abort; if old.len() > i + 1 { break; } - if old.len() == i + 1 && old[i] == profile_list[i] { + // `profile_list[i]` is non-zero here, so `old.get(i) == profile_list[i]` + // already implies `old.len() == i + 1`. + if old.get(i) == profile_list[i] { continue; } - let mut new = old.clone(); - if new.len() < i + 1 { - new.resize(i + 1, 0); - } - new[i] += 1; + let mut new = old; + new.set(i, old.get(i) + 1); new_row.push(new); } } @@ -1010,14 +1239,14 @@ impl MilnorAlgebra { /// 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 { + pub fn seqno(&self, p_part: PPart) -> 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 cur_d: i32 = p_part.iter().zip(xi).map(|(r, &x)| r as i32 * x).sum(); // `cur_d` only decreases in the loop below, so this bounds every `t.g` index. A raw // out-of-bounds panic here means the tables were not built far enough for this element. debug_assert!( @@ -1029,7 +1258,7 @@ impl MilnorAlgebra { 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; + let r = p_part.get(h) as i32; if r == 0 { continue; } @@ -1078,7 +1307,7 @@ impl MilnorAlgebra { #[cfg(feature = "gpu")] pub(crate) fn admissible_matrices( &self, - r_p_part: &[PPartEntry], + r_p_part: PPart, ) -> (usize, usize, Vec, Vec) { let mut matrix = AdmissibleMatrix::new(r_p_part); let cs_len = matrix.col_sums.len(); @@ -1129,8 +1358,8 @@ impl MilnorAlgebra { table.extend( self.ppart_table[(d - q_degree as usize) / q as usize] .iter() - .map(|p_part| MilnorBasisElement { - p_part: p_part.clone(), + .map(|&p_part| MilnorBasisElement { + p_part, q_part, degree: d as i32, }), @@ -1144,14 +1373,16 @@ impl MilnorAlgebra { } fn generate_basis_2(&self, max_degree: i32) { + if !self.stores_basis_table() { + // Derived on demand from `ppart_table`; see the field docs. + return; + } self.basis_table.extend(max_degree as usize, |d| { let mut table: Vec<_> = self.ppart_table[d] .iter() - .map(|p| MilnorBasisElement::from_p(p.clone(), d as i32)) + .map(|&p| MilnorBasisElement::from_p(p, d as i32)) .collect(); - if self.unstable_enabled { - table.sort_by_cached_key(|e| e.excess(fp::prime::TWO)); - } + table.sort_by_cached_key(|e| e.excess(fp::prime::TWO)); table }); } @@ -1182,13 +1413,21 @@ impl MilnorAlgebra { /// Return the degree and index of $Q_1^e P(x)$, or `None` if the element is not present /// (e.g. out of range or excluded by the profile). pub fn try_beps_pn(&self, e: u32, x: PPartEntry) -> Option<(i32, usize)> { + // Bound `x` first: `q * x + e` overflows for a large `x`, so the degree cannot be + // computed before it has been rejected. + if x > PPart::max_entry(0) { + return None; + } let q = self.q() as u32; let degree = (q * x + e) as i32; + if degree > PPart::MAX_DEGREE { + return None; + } self.compute_basis(degree); self.try_basis_element_to_index(&MilnorBasisElement { degree, q_part: e, - p_part: vec![x as PPartEntry], + p_part: PPart::from_iter([x]), }) .map(|index| (degree, index)) } @@ -1198,8 +1437,8 @@ impl MilnorAlgebra { self.try_beps_pn(e, x).unwrap() } - fn multiply_qpart(&self, m1: &MilnorBasisElement, f: u32) -> Vec<(u32, MilnorBasisElement)> { - let mut new_result: Vec<(u32, MilnorBasisElement)> = vec![(1, m1.clone())]; + fn multiply_qpart(&self, m1: MilnorBasisElement, f: u32) -> Vec<(u32, MilnorBasisElement)> { + let mut new_result: Vec<(u32, MilnorBasisElement)> = vec![(1, m1)]; let mut old_result: Vec<(u32, MilnorBasisElement)> = Vec::new(); for k in BitflagIterator::set_bit_iterator(f as u64) { @@ -1222,23 +1461,21 @@ impl MilnorAlgebra { if term.q_part & (1 << (k + i as u32)) != 0 { continue; } - // Check if R - p^k e_i < 0. Only do this from the first term onwards. - if i > 0 && term.p_part[i - 1] < pk { - continue; - } - - let mut new_p = term.p_part.clone(); + let mut new_p = term.p_part; if i > 0 { - new_p[i - 1] -= pk; + // Check if R - p^k e_i < 0. Only do this from the first term onwards. + let entry = new_p.get(i - 1); + if entry < pk { + continue; + } + new_p.set(i - 1, entry - pk); } // Now calculate the number of Q's we are moving past let larger_q = (term.q_part >> (k + i as u32 + 1)).count_ones(); - // If new_p ends with 0, drop them - while let Some(0) = new_p.last() { - new_p.pop(); - } + // Trailing zeros are not represented in a packed p-part, so there is nothing + // to trim here. // Now put everything together let m = MilnorBasisElement { p_part: new_p, @@ -1262,8 +1499,8 @@ impl MilnorAlgebra { &self, res: FpSliceMut, coef: u32, - m1: &MilnorBasisElement, - m2: &MilnorBasisElement, + m1: MilnorBasisElement, + m2: MilnorBasisElement, ) { PPartAllocation::with_local(|allocation| { self.multiply_with_allocation(res, coef, m1, m2, i32::MAX, allocation) @@ -1343,36 +1580,42 @@ impl MilnorAlgebra { // 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()); + let mut terms: Vec = 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 matrix = AdmissibleMatrix::new(r.p_part); let mut working = MilnorBasisElement { q_part: 0, - p_part: PPart::new(), + p_part: PPart::zero(), degree: out_degree, }; loop { 'outer: for term in &terms { - let basis = &term.p_part; - working.p_part.clear(); + let basis = term.p_part; + // Packed: assemble by index instead of `push`. Trailing zeros need no trimming — + // the packed form does not represent them, so `len()` is already the position of + // the highest non-zero entry (that is what makes it a canonical key). + working.p_part = PPart::zero(); + let mut n = 0usize; for j in 0..std::cmp::min(basis.len(), matrix.col_sums.len()) { - if matrix.col_sums[j] > basis[j] { + let b = basis.get(j); + if matrix.col_sums[j] > b { continue 'outer; } - if (basis[j] - matrix.col_sums[j]) & matrix.masks[j] != 0 { + if (b - 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]); + .set(n, (b - matrix.col_sums[j]) | matrix.masks[j]); + n += 1; } if basis.len() < matrix.col_sums.len() { @@ -1382,30 +1625,31 @@ impl MilnorAlgebra { } } for &mask in &matrix.masks[basis.len()..] { - working.p_part.push(mask); + working.p_part.set(n, mask); + n += 1; } } else { for j in matrix.col_sums.len()..std::cmp::min(basis.len(), matrix.masks.len()) { - if basis[j] & matrix.masks[j] != 0 { + let b = basis.get(j); + if b & matrix.masks[j] != 0 { continue 'outer; } - working.p_part.push(basis[j] | matrix.masks[j]); + working.p_part.set(n, b | matrix.masks[j]); + n += 1; } if basis.len() < matrix.masks.len() { for &mask in &matrix.masks[basis.len()..] { - working.p_part.push(mask); + working.p_part.set(n, mask); + n += 1; } } else { - for &entry in &basis[matrix.masks.len()..] { - working.p_part.push(entry); + for j in matrix.masks.len()..basis.len() { + working.p_part.set(n, basis.get(j)); + n += 1; } } } - 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); } @@ -1419,8 +1663,8 @@ impl MilnorAlgebra { &self, mut res: FpSliceMut, coef: u32, - m1: &MilnorBasisElement, - m2: &MilnorBasisElement, + m1: MilnorBasisElement, + m2: MilnorBasisElement, excess: i32, mut allocation: PPartAllocation, ) -> PPartAllocation { @@ -1435,8 +1679,8 @@ impl MilnorAlgebra { for (cc, basis) in m1f { let mut multiplier = PPartMultiplier::::new_from_allocation( self.prime(), - &basis.p_part, - &m2.p_part, + basis.p_part, + m2.p_part, allocation, basis.q_part, target_deg, @@ -1453,8 +1697,8 @@ impl MilnorAlgebra { } else { let mut multiplier = PPartMultiplier::::new_from_allocation( self.prime(), - &m1.p_part, - &m2.p_part, + m1.p_part, + m2.p_part, allocation, 0, target_deg, @@ -1475,7 +1719,7 @@ impl MilnorAlgebra { &self, res: FpSliceMut, coef: u32, - m1: &MilnorBasisElement, + m1: MilnorBasisElement, s_deg: i32, s: FpSlice, ) { @@ -1488,7 +1732,7 @@ impl MilnorAlgebra { &self, mut res: FpSliceMut, coef: u32, - m1: &MilnorBasisElement, + m1: MilnorBasisElement, s_deg: i32, s: FpSlice, mut allocation: PPartAllocation, @@ -1511,6 +1755,17 @@ impl MilnorAlgebra { /// 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. +/// +/// These stay `Vec`s deliberately. Replacing them with fixed-size arrays sized to [`PPart`]'s own +/// structural caps (`MAX_LEN` = 10 rows, `width(0)` = 11 cols, so 110/10/10/20 entries) removes four +/// allocations per `R`, and was measured on `benches/nassau_milnor` to be a NET LOSS: 6.7% slower at +/// `op24xel32`, 2.2% at `op40xel1`, and 0.3-0.6% slower across most of the rest, against ~0.1-0.3% +/// gains on three shapes. The allocations are not the cost -- the zeroing is. A typical `R` is far +/// smaller than the cap (~4 rows x ~6 cols = 24 entries against 110), so `vec![0; rows * cols]` +/// clears a third of what `[0; 110]` would, and the saved `malloc`s do not pay for the extra stores. +/// +/// This is the same effect that made shrinking the GPU kernel's `ENUM_COL_CAP` 32 -> 11 worthwhile, +/// seen from the other side: over-sized fixed state costs more than the allocation it avoids. struct AdmissibleMatrix { cols: usize, rows: usize, @@ -1521,7 +1776,7 @@ struct AdmissibleMatrix { } impl AdmissibleMatrix { - fn new(ps: &[PPartEntry]) -> Self { + fn new(ps: PPart) -> Self { debug_assert!( !ps.is_empty(), "AdmissibleMatrix::new requires a non-empty R; Sq(∅) = 1 is handled by the caller" @@ -1533,12 +1788,12 @@ impl AdmissibleMatrix { .max() .unwrap(); let mut matrix = vec![0; rows * cols]; - for (i, &x) in ps.iter().enumerate() { + 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.extend(ps.iter()); masks.resize(rows + cols - 1, 0); Self { @@ -1618,7 +1873,7 @@ impl AdmissibleMatrix { #[derive(Debug, Default)] struct Matrix2D { cols: usize, - inner: PPart, + inner: Vec, } impl std::fmt::Display for Matrix2D { @@ -1671,8 +1926,7 @@ impl std::ops::IndexMut for Matrix2D { pub struct PPartAllocation { m: Matrix2D, #[cfg(feature = "odd-primes")] - diagonal: PPart, - p_part: PPart, + diagonal: Vec, } thread_local! { @@ -1687,9 +1941,6 @@ impl PPartAllocation { m: Matrix2D::with_capacity(n + 1, n), #[cfg(feature = "odd-primes")] diagonal: Vec::with_capacity(n), - // This size should be the number of diagonals. Even though the answer cannot be that - // long, we still insert zeros then pop them out later. - p_part: Vec::with_capacity(2 * n), } } @@ -1701,21 +1952,21 @@ impl PPartAllocation { } #[allow(non_snake_case)] -pub struct PPartMultiplier<'a, const MOD4: bool> { +pub struct PPartMultiplier { p: ValidPrime, M: Matrix2D, - r: &'a PPart, + r: PPart, rows: usize, cols: usize, diag_num: usize, init: bool, pub ans: MilnorBasisElement, #[cfg(feature = "odd-primes")] - diagonal: PPart, + diagonal: Vec, } #[allow(non_snake_case)] -impl<'a, const MOD4: bool> PPartMultiplier<'a, MOD4> { +impl PPartMultiplier { fn prime(&self) -> ValidPrime { self.p } @@ -1723,8 +1974,8 @@ impl<'a, const MOD4: bool> PPartMultiplier<'a, MOD4> { #[allow(unused_mut)] // Mut is only used with odd primes pub fn new_from_allocation( p: ValidPrime, - r: &'a PPart, - s: &'a PPart, + r: PPart, + s: PPart, mut allocation: PPartAllocation, q_part: u32, degree: i32, @@ -1745,20 +1996,18 @@ impl<'a, const MOD4: bool> PPartMultiplier<'a, MOD4> { M.reset(rows, cols); for i in 1..rows { - M[i][0] = r[i - 1]; + M[i][0] = r.entry(i - 1); } - // This is somehow quite significantly faster than copy_from_slice - #[allow(clippy::manual_memcpy)] for k in 1..cols { - M[0][k] = s[k - 1]; + M[0][k] = s.entry(k - 1); } let ans = MilnorBasisElement { q_part, - p_part: allocation.p_part, + p_part: PPart::zero(), degree, }; - PPartMultiplier { + Self { #[cfg(feature = "odd-primes")] diagonal: allocation.diagonal, p, @@ -1777,7 +2026,6 @@ impl<'a, const MOD4: bool> PPartMultiplier<'a, MOD4> { m: self.M, #[cfg(feature = "odd-primes")] diagonal: self.diagonal, - p_part: self.ans.p_part, } } @@ -1861,7 +2109,7 @@ impl<'a, const MOD4: bool> PPartMultiplier<'a, MOD4> { if inc <= max_inc { // If so, we found our next matrix. for row in 1..i { - self.M[row][0] = self.r[row - 1]; + self.M[row][0] = self.r.entry(row - 1); for col in 1..self.cols { self.M[0][col] += self.M[row][col]; self.M[row][col] = 0; @@ -1885,13 +2133,12 @@ impl<'a, const MOD4: bool> PPartMultiplier<'a, MOD4> { } } -impl Iterator for PPartMultiplier<'_, MOD4> { +impl Iterator for PPartMultiplier { type Item = u32; fn next(&mut self) -> Option { let p = self.prime().as_u32() as PPartEntry; 'outer: loop { - self.ans.p_part.clear(); let mut coef = 1; if self.init { @@ -1912,27 +2159,22 @@ impl Iterator for PPartMultiplier<'_, MOD4> { continue 'outer; } } - self.ans - .p_part - .reserve(std::cmp::max(self.cols, self.rows) - 1); - self.ans.p_part.extend(&self.M[0][1..self.cols]); - - if self.rows > self.cols { - self.ans.p_part.resize(self.r.len(), 0); - } - self.ans - .p_part - .iter_mut() - .zip(self.r.iter()) - .for_each(|(l, r)| *l += r); - - // If new_p ends with 0, drop them - while let Some(0) = self.ans.p_part.last() { - self.ans.p_part.pop(); + // The answer is the top row of the matrix plus `r`, entrywise. Accumulate into + // a plain word and store once; trailing zeros contribute nothing, so there is no + // trimming to do. + let mut ans = 0; + for i in 0..std::cmp::max(self.cols, self.rows) - 1 { + let mut entry = self.r.entry(i); + if i + 1 < self.cols { + entry += self.M[0][i + 1]; + } + debug_assert!(entry <= PPart::max_entry(i)); + ans |= (entry as u64) << PPart::shift(i); } + self.ans.p_part = PPart::from_bits(ans); return Some(coef); } else if self.update() { - self.ans.p_part.reserve(self.diag_num); + let mut ans = 0; for diag_idx in 1..=self.diag_num { let i_min = (diag_idx + 1).saturating_sub(self.cols); let i_max = std::cmp::min(diag_idx + 1, self.rows); @@ -1979,12 +2221,17 @@ impl Iterator for PPartMultiplier<'_, MOD4> { } } } - self.ans.p_part.push(sum); - } - // If new_p ends with 0, drop them - while let Some(0) = self.ans.p_part.last() { - self.ans.p_part.pop(); + // `diag_num` counts diagonals of the working matrix, which can exceed the + // number of entries a p-part of this degree can have; those trailing + // diagonals are necessarily zero and need not be stored. + if diag_idx <= PPart::MAX_LEN { + debug_assert!(sum <= PPart::max_entry(diag_idx - 1)); + ans |= (sum as u64) << PPart::shift(diag_idx - 1); + } else { + debug_assert_eq!(sum, 0); + } } + self.ans.p_part = PPart::from_bits(ans); return Some(coef); } else { @@ -2012,7 +2259,7 @@ impl MilnorAlgebra { let p_idx = self .basis_element_to_index(&MilnorBasisElement::from_p( - vec![ppow as PPartEntry], + PPart::from_iter([ppow as PPartEntry]), p_degree, )) .to_owned(); @@ -2020,7 +2267,7 @@ impl MilnorAlgebra { let q_idx = self .basis_element_to_index(&MilnorBasisElement { q_part: 1 << (i - 1), - p_part: Vec::new(), + p_part: PPart::zero(), degree: q_degree, }) .to_owned(); @@ -2036,13 +2283,13 @@ impl MilnorAlgebra { let first_idx = self.basis_element_to_index(&MilnorBasisElement { q_part: 1 << i, - p_part: Vec::new(), + p_part: PPart::zero(), degree: first_degree, }); let second_idx = self.basis_element_to_index(&MilnorBasisElement { q_part: basis.q_part ^ (1 << i), - p_part: basis.p_part.clone(), + p_part: basis.p_part, degree: second_degree, }); @@ -2076,9 +2323,9 @@ impl MilnorAlgebra { let b = self.basis_element_from_index(degree, idx); let len = b.p_part.len(); - if b.p_part[0..len - 1].iter().all(|&x| x == 0) { + if b.p_part.truncate(len - 1).is_empty() { // There is only one entry - let entry = b.p_part[len - 1]; + let entry = b.p_part.get(len - 1); let (k, m) = factor_pk(p, entry); // This is a power of p @@ -2094,12 +2341,12 @@ impl MilnorAlgebra { let l_degree = l_entry as i32 * self.q(); let l_index = self.basis_element_to_index(&MilnorBasisElement { q_part: 0, - p_part: vec![l_entry], + p_part: PPart::from_iter([l_entry]), degree: l_degree, }); - let mut r_p_part = vec![0; len - 1]; - r_p_part[len - 2] = r_entry; + let mut r_p_part = PPart::zero(); + r_p_part.set(len - 2, r_entry); let r_degree = r_entry as i32 * combinatorics::xi_degrees(p)[len - 2] * self.q(); @@ -2123,15 +2370,15 @@ impl MilnorAlgebra { let mut elt = MilnorBasisElement { q_part: 0, degree: 0, - p_part: vec![0; len], + p_part: PPart::zero(), }; - elt.p_part[len - 1] = pk; - elt.degree = entry_deg * elt.p_part[len - 1] as i32; + elt.p_part.set(len - 1, pk); + elt.degree = entry_deg * pk as i32; let first = (elt.degree, self.basis_element_to_index(&elt)); - elt.p_part[len - 1] = rem_entry; - elt.degree = entry_deg * elt.p_part[len - 1] as i32; + elt.p_part.set(len - 1, rem_entry); + elt.degree = entry_deg * rem_entry as i32; let second = (elt.degree, self.basis_element_to_index(&elt)); let coef = @@ -2140,22 +2387,18 @@ impl MilnorAlgebra { } } else { // There is more than one entry. Just separate out the last entry. - let last_entry = b.p_part[len - 1]; + let last_entry = b.p_part.get(len - 1); let last_deg = combinatorics::xi_degrees(p)[len - 1] * self.q() * last_entry as i32; let mut elt = MilnorBasisElement { q_part: 0, - p_part: vec![0; len], + p_part: PPart::zero(), degree: last_deg, }; - elt.p_part[len - 1] = last_entry; + elt.p_part.set(len - 1, last_entry); let first = (elt.degree, self.basis_element_to_index(&elt)); elt.degree = degree - last_deg; - elt.p_part.clear(); - elt.p_part.extend_from_slice(&b.p_part[0..len - 1]); - while let Some(0) = elt.p_part.last() { - elt.p_part.pop(); - } + elt.p_part = b.p_part.truncate(len - 1); let second = (elt.degree, self.basis_element_to_index(&elt)); buffer.extend([(p - c, first, second)]); }; @@ -2177,16 +2420,23 @@ impl MilnorAlgebra { } impl MilnorAlgebra { - /// Returns `true` if the new element is not within the bounds - fn increment_p_part(element: &mut PPart, max: &[PPartEntry]) -> bool { - element[0] += 1; - for i in 0..element.len() - 1 { - if element[i] > max[i] { - element[i] = 0; - element[i + 1] += 1; + /// Advance `element` to the next p-part bounded entrywise by `max`, in odometer order. + /// + /// Returns `true` once the odometer wraps, i.e. when `element` was already `max`. + /// + /// This carries *before* incrementing rather than after. The two orders enumerate the same + /// sequence, but incrementing first would transiently store `max[i] + 1`, which need not fit + /// in a packed field whose width is exactly saturated by `max[i]`. + fn increment_p_part(element: &mut PPart, max: PPart) -> bool { + for i in 0..max.len() { + let entry = element.get(i); + if entry < max.get(i) { + element.set(i, entry + 1); + return false; } + element.set(i, 0); } - element.last().unwrap() > max.last().unwrap() + true } } @@ -2199,7 +2449,7 @@ impl Bialgebra for MilnorAlgebra { let xi_degrees = combinatorics::xi_degrees(self.prime()); let mut len = 1; - let p_part = &self.basis_element_from_index(op_deg, op_idx).p_part; + let p_part = self.basis_element_from_index(op_deg, op_idx).p_part; for i in p_part.iter() { len *= i + 1; @@ -2207,32 +2457,23 @@ impl Bialgebra for MilnorAlgebra { let len = len as usize; let mut result = Vec::with_capacity(len); - let mut cur_ppart: PPart = vec![0; p_part.len()]; + let n = p_part.len(); + let mut cur_ppart = PPart::zero(); loop { let mut left_degree: i32 = 0; - for i in 0..cur_ppart.len() { - left_degree += cur_ppart[i] as i32 * xi_degrees[i]; + let mut right_ppart = PPart::zero(); + for (i, &xi_degree) in xi_degrees.iter().enumerate().take(n) { + let entry = cur_ppart.get(i); + left_degree += entry as i32 * xi_degree; + // Trailing zeros are dropped by the packing, so no trimming is needed. + right_ppart.set(i, p_part.get(i) - entry); } let right_degree: i32 = op_deg - left_degree; - let mut left_ppart = cur_ppart.clone(); - while let Some(0) = left_ppart.last() { - left_ppart.pop(); - } - - let mut right_ppart = cur_ppart - .iter() - .enumerate() - .map(|(i, v)| p_part[i] - *v) - .collect::>(); - while let Some(0) = right_ppart.last() { - right_ppart.pop(); - } - let left_idx = self.basis_element_to_index(&MilnorBasisElement { degree: left_degree, q_part: 0, - p_part: left_ppart, + p_part: cur_ppart, }); let right_idx = self.basis_element_to_index(&MilnorBasisElement { degree: right_degree, @@ -2279,7 +2520,7 @@ mod tests { for i in 0..dim { let elt = algebra.basis_element_from_index(d, i); assert_eq!( - algebra.seqno(&elt.p_part), + algebra.seqno(elt.p_part), i, "seqno mismatch at degree {d}, index {i}: {elt:?}" ); @@ -2492,6 +2733,35 @@ mod tests { assert_eq!(a2.try_beps_pn(0, 8), None); } + /// `basis_element_from_string` is documented to be total. Inputs whose exponents overflow + /// intermediate arithmetic must return `None`, not panic. + #[test] + fn basis_element_from_string_rejects_overflowing_exponents() { + let algebra = MilnorAlgebra::new(fp::prime::TWO, false); + algebra.compute_basis(8); + + // `t` indexes the xi-degree table, which has exactly `MAX_LEN` entries. + assert_eq!(algebra.basis_element_from_string("P^1_10"), None); + assert_eq!(algebra.basis_element_from_string("P^1_99"), None); + // `p^s` overflows for large `s`. + assert_eq!(algebra.basis_element_from_string("P^64_2"), None); + assert_eq!(algebra.basis_element_from_string("P^4294967295_2"), None); + // ... and so does `q * x` in the `Sq`/`P` path. + assert_eq!(algebra.basis_element_from_string("Sq4294967295"), None); + } + + /// `try_beps_pn` is the non-panicking half of `beps_pn`; an out-of-range `x` must not trip + /// overflow on the way to the bounds check. + #[test] + fn try_beps_pn_rejects_overflowing_x() { + for p in [2, 3] { + let algebra = MilnorAlgebra::new(ValidPrime::new(p), false); + assert_eq!(algebra.try_beps_pn(0, PPartEntry::MAX), None); + assert_eq!(algebra.try_beps_pn(0, PPartEntry::MAX / 2), None); + assert_eq!(algebra.try_beps_pn(1, PPartEntry::MAX), None); + } + } + #[test] fn basis_element_from_string_total_milnor() { let p = ValidPrime::new(2); @@ -2506,13 +2776,17 @@ mod tests { assert_eq!(algebra.basis_element_to_string(d, i), name); } - // Syntactically-valid names that name no basis element must return `None` + // "P0"/"Sq0" name the identity. A packed p-part does not represent trailing zeros, so + // `P(0)` and `P()` are the same value, and `try_beps_pn(0, 0)` finds the degree-0 basis + // element. This matches `AdemAlgebra::try_beps_pn`, which special-cases `x == 0` to + // `Some((0, 0))`; the previous `None` here came from `vec![0]` and `vec![]` hashing + // differently, which was an artifact of the unpacked representation. + assert_eq!(algebra.basis_element_from_string("P0"), Some((0, 0))); + assert_eq!(algebra.basis_element_from_string("Sq0"), Some((0, 0))); + + // Syntactically-valid names that name no basis element must still return `None` // (they previously panicked in `basis_element_to_index`). // - // "P0"/"Sq0" parse via `try_beps_pn(0, 0)`, building the element - // {q_part: 0, p_part: [0]} in degree 0, which is not a basis element. - assert_eq!(algebra.basis_element_from_string("P0"), None); - assert_eq!(algebra.basis_element_from_string("Sq0"), None); // "Q_5" parses via the Q/P branch into a candidate element (degree 63) // whose basis lookup finds nothing at p = 2. assert_eq!(algebra.basis_element_from_string("Q_5"), None); @@ -2587,6 +2861,216 @@ mod tests { } } + /// The packing is only sound because each field is wide enough for every entry that can occur + /// at degree at most `MAX_DEGREE`. Check that against the $\xi$-degrees directly, so that + /// changing `MAX_DEGREE` or `WIDTHS` without the other fails loudly. + #[test] + fn ppart_widths_cover_max_degree() { + let xi_degrees = combinatorics::xi_degrees(fp::prime::TWO); + for (i, &xi_degree) in xi_degrees.iter().enumerate().take(PPart::MAX_LEN) { + // deg P(R) = sum_i r_i (2^i - 1) with non-negative terms, so r_i <= deg / (2^i - 1). + let bound = PPart::MAX_DEGREE / xi_degree; + assert!( + bound <= PPart::max_entry(i) as i32, + "entry {i} needs to hold {bound} but only holds up to {}", + PPart::max_entry(i), + ); + } + // There is no entry beyond `MAX_LEN` to store: the xi-degree table itself stops there, so + // `compute_ppart` cannot produce a longer p-part. If `fp` ever raises + // `MAX_MULTINOMIAL_LEN`, this fires and `WIDTHS` has to be revisited. + assert_eq!(xi_degrees.len(), PPart::MAX_LEN); + // ... and the layout uses the whole word, so `MAX_DEGREE` is as large as it can be. + assert_eq!( + PPart::shift(PPart::MAX_LEN - 1) + PPart::width(PPart::MAX_LEN - 1), + 64 + ); + } + + #[test] + fn ppart_accessors() { + let mut p = PPart::from_slice(&[3, 0, 5]); + assert_eq!(p.len(), 3); + assert_eq!(p.iter().collect::>(), vec![3, 0, 5]); + assert_eq!(p.get(1), 0); + assert_eq!(p.get(2), 5); + // Reading past the end is zero, not a panic. + assert_eq!(p.get(7), 0); + assert_eq!(p.get(PPart::MAX_LEN), 0); + + // Trailing zeros are not represented, so they do not affect equality, length or hashing. + assert_eq!(PPart::from_slice(&[3, 0, 5, 0, 0]), p); + assert_eq!(PPart::from_slice(&[]), PPart::zero()); + assert_eq!(PPart::from_slice(&[0, 0]), PPart::zero()); + assert_eq!(PPart::zero().len(), 0); + assert!(PPart::zero().is_empty()); + + assert_eq!(p.truncate(2), PPart::from_slice(&[3])); + assert_eq!(p.truncate(0), PPart::zero()); + assert_eq!(p.truncate(PPart::MAX_LEN + 3), p); + + p.set(1, 7); + assert_eq!(p, PPart::from_slice(&[3, 7, 5])); + p.set(2, 0); + assert_eq!(p, PPart::from_slice(&[3, 7])); + } + + #[test] + fn ppart_rejects_out_of_range() { + // Too many entries, and an entry too large for its field. + assert_eq!(PPart::try_from_slice(&[1; PPart::MAX_LEN + 1]), None); + assert_eq!(PPart::try_from_slice(&[0, PPart::max_entry(1) + 1]), None); + // ... but a zero past the end is only padding. + let mut padded = vec![0; PPart::MAX_LEN + 4]; + padded[0] = 2; + assert_eq!( + PPart::try_from_slice(&padded), + Some(PPart::from_slice(&[2])) + ); + } + + #[test] + #[should_panic(expected = "does not fit")] + fn ppart_set_out_of_range_panics() { + PPart::zero().set(0, PPart::max_entry(0) + 1); + } + + /// `increment_p_part` walks up to and including `max`, whose top entry may saturate its field. + /// Incrementing before carrying would overflow there. + #[test] + fn ppart_odometer_handles_saturated_field() { + let top = PPart::MAX_LEN - 1; + let mut max = PPart::from_slice(&[2]); + max.set(top, PPart::max_entry(top)); + + let mut count = 0; + let mut cur = PPart::zero(); + loop { + count += 1; + if MilnorAlgebra::increment_p_part(&mut cur, max) { + break; + } + } + assert_eq!(count, 3 * (PPart::max_entry(top) as usize + 1)); + // Wrapping leaves the odometer back at zero. + assert_eq!(cur, PPart::zero()); + } + + /// Pack every basis element the algebra actually produces and check nothing collides or is + /// lost. This is the property the whole representation rests on. + #[rstest] + #[case(2, 120)] + #[case(3, 200)] + fn ppart_packing_is_faithful(#[case] p: u32, #[case] max_degree: i32) { + let algebra = MilnorAlgebra::new(ValidPrime::new(p), false); + algebra.compute_basis(max_degree); + + for t in 0..=max_degree { + let mut seen = HashMap::default(); + for i in 0..algebra.dimension(t) { + let elt = algebra.basis_element_from_index(t, i); + // The packed value plus the q-part identifies the element within its degree. + assert!( + seen.insert((elt.p_part.bits(), elt.q_part), i).is_none(), + "collision at degree {t} for {elt}" + ); + // Round-trip through a slice, and back through the index map. + assert_eq!( + PPart::from_slice(&elt.p_part.iter().collect::>()), + elt.p_part + ); + assert_eq!(algebra.basis_element_to_index(&elt), i); + // The degree really is recoverable from the entries. + let mut recomputed = elt; + recomputed.compute_degree(ValidPrime::new(p)); + assert_eq!(recomputed.degree, t); + } + } + } + + /// The basis *order* at `p = 2` is a wire format: saved resolutions store coefficients by + /// index, so reordering silently invalidates them without `magic()` changing. Deriving the + /// basis from `ppart_table` preserves the order `generate_basis_2` produced, since the stable + /// path never sorted. Pin that down against fixed expected names, so the check does not depend + /// on `ppart_table` -- the very thing it is guarding. + #[test] + fn basis_order_at_p2_is_stable() { + let algebra = MilnorAlgebra::new(fp::prime::TWO, false); + algebra.compute_basis(8); + + let expected: [&[&str]; 9] = [ + &["1"], + &["P(1)"], + &["P(2)"], + &["P(3)", "P(0, 1)"], + &["P(4)", "P(1, 1)"], + &["P(5)", "P(2, 1)"], + &["P(6)", "P(3, 1)", "P(0, 2)"], + &["P(7)", "P(4, 1)", "P(1, 2)", "P(0, 0, 1)"], + &["P(8)", "P(5, 1)", "P(2, 2)", "P(1, 0, 1)"], + ]; + for (t, names) in expected.iter().enumerate() { + let t = t as i32; + assert_eq!(algebra.dimension(t), names.len(), "dimension in degree {t}"); + for (i, name) in names.iter().enumerate() { + assert_eq!( + &algebra.basis_element_to_string(t, i), + name, + "degree {t}, index {i}" + ); + } + } + } + + /// At `p = 2` with unstable support off, the basis is not stored: it is derived from + /// `ppart_table`. Check the derivation reproduces exactly what the table used to hold, so the + /// redundancy this relies on is asserted rather than assumed. + #[test] + fn basis_is_derived_at_p2() { + let p = fp::prime::TWO; + let algebra = MilnorAlgebra::new(p, false); + algebra.compute_basis(120); + assert!( + !algebra.stores_basis_table(), + "p = 2 stable should not be storing the basis" + ); + + for t in 0..=120 { + let pparts = algebra.ppart_table(t); + assert_eq!(algebra.dimension(t), pparts.len()); + for (i, &p_part) in pparts.iter().enumerate() { + // This is precisely what `generate_basis_2` used to store. + let expected = MilnorBasisElement { + q_part: 0, + p_part, + degree: t, + }; + let actual = algebra.basis_element_from_index(t, i); + assert_eq!(actual.p_part, expected.p_part, "degree {t}, index {i}"); + assert_eq!(actual.q_part, expected.q_part, "degree {t}, index {i}"); + assert_eq!(actual.degree, expected.degree, "degree {t}, index {i}"); + } + } + } + + /// The two configurations that still need the table really do differ from `ppart_table`, so + /// the exemption in `stores_basis_table` is not over-broad. + #[rstest] + #[case(3, false)] + #[case(2, true)] + fn basis_is_stored_when_it_must_be(#[case] p: u32, #[case] unstable: bool) { + let algebra = MilnorAlgebra::new(ValidPrime::new(p), unstable); + algebra.compute_basis(60); + assert!(algebra.stores_basis_table()); + // Every stored element still round-trips through the index map. + for t in 0..=60 { + for i in 0..algebra.dimension(t) { + let elt = algebra.basis_element_from_index(t, i); + assert_eq!(algebra.basis_element_to_index(&elt), i); + } + } + } + #[test] fn test_clone_into() { let mut other = MilnorBasisElement::default(); @@ -2598,34 +3082,34 @@ mod tests { check(&MilnorBasisElement { q_part: 3, - p_part: vec![3, 2], + p_part: PPart::from_slice(&[3, 2]), degree: 12, }); check(&MilnorBasisElement { q_part: 1, - p_part: vec![3], + p_part: PPart::from_slice(&[3]), degree: 11, }); check(&MilnorBasisElement { q_part: 5, - p_part: vec![1, 3, 5, 2], + p_part: PPart::from_slice(&[1, 3, 5, 2]), degree: 7, }); check(&MilnorBasisElement { q_part: 0, - p_part: vec![], + p_part: PPart::zero(), degree: 2, }); } #[test] fn test_ppart_multiplier_2() { - let r = vec![1, 4]; - let s = vec![2, 4]; + let r = PPart::from_slice(&[1, 4]); + let s = PPart::from_slice(&[2, 4]); let mut m = PPartMultiplier::::new_from_allocation( fp::prime::TWO, - &r, - &s, + r, + s, PPartAllocation::default(), 0, 0, @@ -2661,12 +3145,12 @@ mod tests { #[test] fn test_ppart_multiplier_3() { - let r = vec![3, 4]; - let s = vec![1, 4]; + let r = PPart::from_slice(&[3, 4]); + let s = PPart::from_slice(&[1, 4]); let mut m = PPartMultiplier::::new_from_allocation( ValidPrime::new(3), - &r, - &s, + r, + s, PPartAllocation::default(), 0, 0, diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index d93b7255e0..8fdca58a45 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -23,36 +23,461 @@ use cubecl::{ cuda::{CudaDevice, CudaRuntime}, prelude::*, }; -use cubecl_common::stream_id::StreamId; +use cubecl_common::bytes::Bytes; -/// The single CUDA stream all GPU work is pinned to (via [`StreamId::executes`]). -/// -/// CubeCL's memory pools are per-stream, and the resolution issues launches from many -/// rayon worker threads (each its own stream). Left alone, each stream's pool retains its -/// freed per-launch buffers (chiefly the hundreds-of-MB `out_h`), and across ~16 streams -/// they accumulate until the 4 GB card OOMs — `memory_cleanup` only trims the *calling* -/// stream's pool. Pinning every launch to one stream gives one pool that each launch's -/// `memory_cleanup` fully reclaims. Value 0 is a valid stream id (the first thread's). -const GPU_STREAM: StreamId = StreamId { value: 0 }; - -// Only the `#[cfg(test)]` standalone `seqno_kernel` sizes its working array by this bound; the -// production kernels use `WORKING_CAP`. -#[cfg(test)] +// Bounds the per-thread enumeration state ([`ENUM_ROW_CAP`]) and the `#[cfg(test)]` `seqno_kernel`'s +// working array; the multiply kernel uses `WORKING_CAP`. use crate::algebra::combinatorics::MAX_XI_TAU; -use crate::algebra::{Algebra, MilnorAlgebra, combinatorics::xi_degrees}; +use crate::algebra::{Algebra, MilnorAlgebra, combinatorics::xi_degrees, milnor_algebra::PPart}; /// 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; -/// Maximum `(product, matrix, term)` thread-pairs per GPU launch. The batch multiply indexes -/// threads by CubeCL's `ABSOLUTE_POS` (a `u32`), so one launch can address at most `2^32` threads; -/// a single all-rows reuse build reaches ~4.4e9 pairs at stem ~145, past that limit. Launches whose -/// pair count exceeds this cap are split into chunks each bounded by it. `1 << 30` (~1.07e9) leaves -/// >3x headroom under `2^32` even after a chunk's final product pushes it over, and keeps each -/// chunk's grid (`chunk_pairs / 256` cubes) well under CUDA's `2^31 - 1` grid-dimension limit. -const GPU_PAIR_CHUNK: usize = 1 << 30; +/// `PPart::MAX_LEN` — the number of entries the packed exponent sequence can hold. +/// +/// Positions beyond it are unreachable, not merely unused: entry `r_n` multiplies +/// `deg(xi_n) = 2^n - 1`, so length 11 requires degree >= 2047 while `PPart::MAX_DEGREE` is 2045. +/// The multiply kernel's `working` accumulator therefore packs into one `u64` with no loss. +const PPART_MAX_LEN: usize = 10; + +/// Bit offset and value mask of each packed p-part field, uploaded per launch (80 bytes) so the +/// kernel can unpack `working` without a per-thread array. Mirrors `PPart`'s private tables. +fn ppart_shift_mask() -> (Vec, Vec) { + let shift: Vec = (0..PPART_MAX_LEN).map(|i| PPart::shift(i)).collect(); + let mask: Vec = (0..PPART_MAX_LEN) + .map(|i| { + let w = PPart::shift(i + 1) - PPart::shift(i); + ((1u64 << w) - 1) as u32 + }) + .collect(); + (shift, mask) +} + +/// Per-thread local caps for the in-kernel admissible enumeration ([`enumerate_admissible_kernel`]). +/// Each `R` has `rows = |p_part| ≤ MAX_XI_TAU` and `cols ≤ ENUM_COL_CAP` (max bit-length of an entry), +/// so the enumeration's `matrix` is `rows*cols`, `col_sums` is `cols−1`, and `masks` is `rows+cols−1`. +/// These bound the fixed-size local `Array`s the kernel allocates per thread. +/// +/// `cols` is the widest bit-length of any p-part entry, so its true bound is `PPart::width(0)` — the +/// field holding `r_1`, the widest — and NOT `WORKING_CAP`, which sizes an unrelated array (the +/// multiply kernel's assembled p-part) and is nearly 3x larger. That conflation made `matrix`, the +/// hottest per-thread array, 320 u32 instead of 110. It is CUDA *local* memory: dynamically indexed, +/// so it cannot be register-allocated and every access is a real off-chip load. Deriving the cap from +/// the width table keeps it correct if `PPart`'s layout ever changes; `enum_col_cap_bounds_real_rs` +/// checks it against every actual `R`. +const ENUM_ROW_CAP: usize = MAX_XI_TAU; +const ENUM_COL_CAP: usize = PPart::width(0) as usize; +const ENUM_MATRIX_CAP: usize = ENUM_ROW_CAP * ENUM_COL_CAP; +const ENUM_MASK_CAP: usize = ENUM_ROW_CAP + ENUM_COL_CAP; + +/// Threads per block for [`enumerate_admissible_kernel`], and the stride of its shared-memory state. +/// +/// 32, not the 256 this used to launch, because the per-thread enumeration state now lives in shared +/// memory: 152 `u32` per thread (`matrix` 110 + `totals` 10 + `col_sums` 11 + `masks` 21) is 608 B, +/// so 32 threads is 19.5 KB of the 48 KB static budget while 256 would need 155 KB. +/// +/// Trading occupancy for shared memory is normally a bad deal, and when this was first tried on the +/// all-`R` benchmark it WAS: it cut warp cycles per issued instruction 8.90 -> 5.92 (-33%) but +/// dropped active warps per scheduler 2.09 -> 1.61 and lost overall. That benchmark enumerates every +/// `R` to degree 130 in ONE launch (1397 blocks), where occupancy is the binding constraint. +/// PRODUCTION IS NOT THAT: a real launch is ~6 blocks against the H200's 3168 slots +/// (`Waves Per SM = 0.002`), so there is no occupancy to lose — the SMs are empty either way. The +/// block-size sweep measured 256/64/32 threads at 561/569/542 s, i.e. flat, which is the same fact +/// from the other side: grid width does not set this kernel's time. +const ENUM_BLOCK: u32 = 32; + +/// Offsets of each per-thread array within the shared state, in units of "one element per thread". +/// Layout is `elem * ENUM_BLOCK + UNIT_POS`, so the lanes of a warp touch consecutive words — one +/// word per bank, conflict-free, and the reason the stride is the block size rather than 1. +const ENUM_ST_TOTALS: usize = ENUM_MATRIX_CAP; +const ENUM_ST_COLSUMS: usize = ENUM_ST_TOTALS + ENUM_ROW_CAP; +const ENUM_ST_MASKS: usize = ENUM_ST_COLSUMS + ENUM_COL_CAP; +const ENUM_STATE: usize = ENUM_ST_MASKS + ENUM_MASK_CAP; + +/// Target `(product, matrix, term)` thread-pairs per GPU launch. The batch multiply indexes threads +/// by CubeCL's `ABSOLUTE_POS` (a `u32`), so one launch can address at most `2^32` threads; a single +/// all-rows reuse build reaches ~4.4e9 pairs at stem ~145, past that limit. The row-block splitter +/// in [`multiply_batch_on_gpu`] closes a block once its pair count would pass this target (alongside +/// the [`gpu_block_bytes`] output budget). +/// +/// Set close to the `2^32` ceiling, not far below it: every extra split is a whole extra launch +/// (upload + kernel + blocking readback), and at record stems this — not the byte budget — is the +/// binding constraint, so a conservative value chops each giant multiply into several +/// otherwise-unnecessary launches (measured: `1 << 30` pegged the giants at ~1.07e9 pairs, ~4 +/// launches each, while their output is only ~350 MB, well under `gpu_block_bytes`). `3.9e9` leaves +/// ~0.39e9 of headroom under `2^32` for a lone over-budget row (the splitter always takes ≥1 row, +/// and a single row past `2^32` still trips the per-block `u32::try_from` assert), and keeps the +/// grid (`pairs / 256` cubes ≈ 1.5e7) far under CUDA's `2^31 - 1` grid-dimension limit. +const GPU_PAIR_CHUNK: usize = 3_900_000_000; + +/// Chunk size (log2) of the multiply kernel's coarse product index. One entry per `2^COARSE_LOG` +/// pairs, so a launch of billions of pairs needs a table of a few thousand `u32` — negligible to +/// build and upload, and it turns the per-thread product lookup from a full binary search over +/// every product into a scan bounded by how many products one chunk spans. +const COARSE_LOG: usize = 20; + +/// Terms one multiply thread handles against a single matrix. `col_sums`/`masks` depend only on the +/// matrix, so a group amortises those reads (and their address arithmetic) across `TERM_GROUP` +/// terms: loads per pair fall from 3 per column to `2/TERM_GROUP + 1`. +/// +/// Tuned, and the optimum is NOT monotonic — bigger groups amortise more but waste more lanes on +/// the ragged tail, since a product's `nt` terms need `ceil(nt / TERM_GROUP)` groups and the last +/// one is usually partial. Two interleaved rounds at the measured `nt ~ 5`: +/// +/// | TERM_GROUP | pairs/s | idle lanes at nt=5 | +/// |------------|----------------|--------------------| +/// | 2 | 1.09 / 1.06e10 | 1 of 6 (17%) | +/// | **3** | **1.31 / 1.28e10** | 1 of 6 (17%) | +/// | 4 | 1.20 / 1.18e10 | 3 of 8 (37%) | +/// | 6 | 1.27 / 1.25e10 | 1 of 6 (17%) | +/// | 8 | 1.06 / 1.07e10 | 3 of 8 (37%) | +/// +/// Every 17%-waste value beats every 37%-waste value, so tail waste dominates the choice; among +/// those, 3 amortises more than 2 and holds more registers than 6 does not need. Retune if the +/// terms-per-product regime moves: this is fitted to `nt ~ 5`, and a workload with a different +/// average would want a different divisor. +const TERM_GROUP: usize = 3; + +/// Matrices one multiply thread handles, the second axis of the tile alongside [`TERM_GROUP`]. +/// `col_sums`/`masks` are per-matrix, so an `M x T` tile costs `2M + T` loads per column for `M*T` +/// pairs. Unlike terms there is no meaningful ragged tail here -- `num_mats` runs to ~20 000, so a +/// partial tile idles a couple of lanes out of thousands. +const MATRIX_GROUP: usize = 2; + +/// Per-launch output-buffer budget in bytes (`NASSAU_GPU_BLOCK_MB`, default 512 MiB). +/// +/// A launch's transient footprint — host marshal buffers, pinned staging, device buffers, and +/// each stream's retained pool pages — scales with its output size, and with the device mutex +/// gone many workers hold such transients simultaneously; at record stems an unbounded all-rows +/// reuse build multiplies to >100 GB on both host and device. [`multiply_batch_on_gpu`] therefore +/// splits large builds into row blocks whose output buffer stays under this budget. Rows of +/// distinct products are independent (each product writes only its own row), so blocks simply +/// concatenate — the same in-between as the old per-signature builds, but with blocks big enough +/// to keep the launch amortization. Together with [`GPU_BUDGET`] this makes peak transient +/// memory a configured constant (≈ the byte budget) instead of a function of the frontier size. +fn gpu_block_bytes() -> usize { + static BYTES: LazyLock = LazyLock::new(|| { + std::env::var("NASSAU_GPU_BLOCK_MB") + .ok() + .and_then(|v| v.parse::().ok()) + .filter(|&mb| mb > 0) + .unwrap_or(512) + << 20 + }); + *BYTES +} + +/// Byte-weighted budget bounding the total *output size* of in-flight device sections +/// (`NASSAU_GPU_MEM_BUDGET_MB`, default 4096). +/// +/// A count-based cap (formerly `NASSAU_GPU_CONCURRENCY` = 8 sections) throttled exactly the +/// wrong region: low-stem launches are a few MB each and were capped at 8 concurrent (measured +/// 4x slowdown vs the uncapped code at stem 130), while the cap only exists for the +/// multi-hundred-MB frontier blocks. Weighting admission by output bytes admits dozens of +/// small launches concurrently and still bounds the frontier to ~budget / [`gpu_block_bytes`] +/// in flight. A launch heavier than the whole budget is admitted alone (when nothing else is +/// in flight), so progress is always possible. Waiters have heterogeneous weights, so release +/// notifies all. +/// +/// SAFETY INVARIANT: a permit must never be held across a rayon parallel section. A par_iter's +/// chunks execute on other threads, which do not carry the holder's thread-local +/// `ParallelGuard` flag and so can steal a resolution-step job mid-chunk; that job would park +/// on [`GpuPermit::acquire`] while the holder's permit waits on the never-finishing join — +/// a cycle (observed as a full stall on H200). [`multiply_batch_block`] therefore acquires its +/// permit only after the parallel marshal, guarding a strictly sequential section: every holder +/// makes progress, so parked acquirers always wake (priority inversion at worst, never +/// deadlock). +struct GpuBudget { + budget: usize, + used: Mutex, + freed: Condvar, +} + +static GPU_BUDGET: LazyLock = LazyLock::new(|| GpuBudget { + budget: std::env::var("NASSAU_GPU_MEM_BUDGET_MB") + .ok() + .and_then(|v| v.parse::().ok()) + .filter(|&mb| mb > 0) + .unwrap_or(4096) + << 20, + used: Mutex::new(0), + freed: Condvar::new(), +}); + +/// The single OS thread that owns the CUDA stream, and the queue that feeds it. +/// +/// # Why a dedicated thread +/// +/// Every worker used to run its own device section on the *shared* stream 0 (see +/// stream 0, so ~7 concurrent workers raced for cubecl's per-device submission path. +/// That path is not FIFO-fair, and under sustained contention a worker could be passed over for +/// minutes: measured on a stem-200 run, one `step` blocked for **370 s** inside the multiply while +/// six peers each stayed 100 % busy completing 200–576 steps at 0.6–1.6 s apiece. The same +/// `(bidegree, signature)` — bit-identical work — took 0.04 s in one run and 235 s in another; 428 +/// step pairs differed by more than 100× across two complete runs, wasting ~110–150 min of thread +/// time each. The work is uniform (within a bidegree the spread is ~5×); the outliers were purely +/// service-order artifacts. +/// +/// Funnelling every device section through one thread fed by an `mpsc` channel makes service order +/// FIFO by construction, so a worker's wait is bounded by the jobs enqueued ahead of it and the +/// starvation case cannot arise. Total device serialisation is unchanged — stream 0 already +/// serialised everything — but it is now *fair*. Workers still marshal in parallel; only the +/// (already sequential) device section moves. +/// +/// # Invariants +/// +/// - The receive loop runs inside a single `StreamId::executes`, so the stream is bound once and +/// has exactly one driver thread for the process's lifetime — what cubecl's per-stream state +/// assumes. +/// - Tasks must not need a worker thread. Nothing here enters rayon, and every resident-store lock +/// ([`resident_info`], `ensure_basis`) is taken and released *inside* one task, never held across +/// a submission — so a blocked worker can never hold a lock this thread waits on. +/// - Panics are caught per task and forwarded to the waiting worker, which resumes the unwind. A +/// panic that killed this thread would instead hang every future submission forever. +mod gpu_thread { + use std::{ + sync::{ + OnceLock, + atomic::{AtomicU64, Ordering}, + mpsc, + }, + time::Instant, + }; + + use crossbeam_channel::{Sender, unbounded}; + use cubecl_common::stream_id::StreamId; + + /// Jobs enqueued but not yet started: the shared queue's depth when a worker joins it. + static DEPTH: AtomicU64 = AtomicU64::new(0); + + type Task = Box; + + /// How long a submission waited in the queue, and how long it then took on the device. + pub(super) struct Timing { + /// Enqueue → task start. Under FIFO this is the work queued ahead of this job. + pub queue_ms: f64, + /// Task start → task end: the device section proper. + pub exec_ms: f64, + /// Queue depth observed at enqueue (including this job). + pub depth: u64, + } + + /// One queue PER DEVICE. A shared pull-queue would balance better — a worker takes the next job + /// the instant it frees up, with no need to predict job size — but the sharded master makes work + /// device-AFFINE: an `R`'s rows live on exactly one device, so its products can only run there. + /// Balance therefore comes from spreading `R`s evenly (round-robin at first sight), not from + /// letting idle workers steal. + /// + /// Each worker owns its device and its stream, and the thread-local set here is what keeps a + /// device handle from ever reaching another device's client. The resident master/basis is NOT + /// per worker — [`RESIDENT_DEV`] is indexed by device — so the [`gpu_streams`] workers sharing + /// one device's queue also share its segments rather than replicating them. + /// + /// `crossbeam-channel` because this is genuinely multi-consumer: std's `mpsc` has a single + /// receiver, so one shared queue there would mean wrapping it in a `Mutex` and serialising every + /// pop behind a lock held across a blocking `recv`. + fn senders() -> &'static Vec> { + static QUEUES: OnceLock>> = std::sync::OnceLock::new(); + QUEUES.get_or_init(|| { + let nstream = super::gpu_streams(); + let mut txs = Vec::with_capacity(super::gpu_count()); + for dev in 0..super::gpu_count() { + let (tx, rx) = unbounded::(); + txs.push(tx); + // `gpu_streams()` workers share this device's queue, so whichever frees up first + // takes the next section — the pull-queue balance the per-device split allows, now + // that there is more than one puller. Each gets its OWN `StreamId` (ids are global, + // hence `dev * nstream + w`), which is what actually lets their launches overlap on + // the device; sharing one id would re-serialise them. + for w in 0..nstream { + let rx = rx.clone(); + std::thread::Builder::new() + .name(format!("nassau-gpu{dev}s{w}")) + .spawn(move || { + super::CUR_DEVICE.with(|c| c.set(dev)); + // Bind the stream once for the whole loop: one stream per driver + // thread, a distinct id so the runtime keeps them independent. + StreamId { + value: (dev * nstream + w) as u64, + } + .executes(|| { + while let Ok(task) = rx.recv() { + task(); + } + }); + }) + .expect("failed to spawn a nassau-gpu thread"); + } + } + txs + }) + } + + /// Run `f` on the GPU thread, blocking until it returns. Panics propagate to the caller. + /// A submitted job that has not been waited on yet. + pub(super) struct Pending { + rx: mpsc::Receiver<(std::thread::Result, f64, f64)>, + depth: u64, + } + + impl Pending { + /// Block for the result. Panics in the job propagate to the caller. + pub(super) fn wait(self) -> (T, Timing) { + let (out, queue_ms, exec_ms) = + self.rx.recv().expect("the nassau-gpu thread died mid-task"); + match out { + Ok(v) => ( + v, + Timing { + queue_ms, + exec_ms, + depth: self.depth, + }, + ), + Err(payload) => std::panic::resume_unwind(payload), + } + } + } + + /// Submit `f` to `dev`'s worker WITHOUT blocking. The sharded fan-out needs every device in + /// flight at once; blocking per device would serialise exactly what the shard split parallelises, + /// and spawning a thread per device per block (the first cut) churned hundreds of OS threads a + /// second — visible as `ThreadId(867540)` in the logs. + pub(super) fn submit_on(dev: usize, f: F) -> Pending + where + F: FnOnce() -> T + Send + 'static, + T: Send + 'static, + { + let (tx, rx) = mpsc::sync_channel::<(std::thread::Result, f64, f64)>(1); + let depth = DEPTH.fetch_add(1, Ordering::Relaxed) + 1; + let enqueued = Instant::now(); + senders()[dev] + .send(Box::new(move || { + let queue_ms = enqueued.elapsed().as_secs_f64() * 1e3; + DEPTH.fetch_sub(1, Ordering::Relaxed); + let started = Instant::now(); + // `AssertUnwindSafe`: on a panic the payload is forwarded and the worker resumes + // the unwind, so no state observed after the catch is reused here. + let out = std::panic::catch_unwind(std::panic::AssertUnwindSafe(f)); + let exec_ms = started.elapsed().as_secs_f64() * 1e3; + // A send error means the worker vanished (itself panicking); drop the result. + let _ = tx.send((out, queue_ms, exec_ms)); + })) + .expect("the nassau-gpu thread died"); + Pending { rx, depth } + } + + /// Submit to `dev` and block for the result. + pub(super) fn run_on(dev: usize, f: F) -> (T, Timing) + where + F: FnOnce() -> T + Send + 'static, + T: Send + 'static, + { + submit_on(dev, f).wait() + } +} + +/// Persistent per-caller helper threads for the shard fan-out. +/// +/// # Why per-caller, and not one shared pool +/// +/// The fan-out's whole value is that each shard's pipeline — marshal, permit, submit, wait — runs +/// concurrently, so all `gpu_count()` devices receive work at once instead of each submission +/// queueing behind the previous shard's marshal. Verified stem-200s: 2412 s with concurrent +/// submission, 3138 s without, and 3358 s when only the *marshal* was parallelised (which cut +/// marshal by 840 s and still lost 220 s — marshal time does not predict wall time). +/// +/// A *shared* pool re-serialises exactly that: with every caller feeding one queue, a shard waits +/// behind other callers' shards. Measured ~3743 s for a 512-thread shared rayon pool, and 3373 s +/// routing the same work through the global rayon pool — both slower than doing nothing. So each +/// caller owns its helpers privately: no shared queue, no stealing, no cross-caller interference. +/// +/// # Sizing +/// +/// `gpu_count() - 1` threads per calling thread, created lazily on that thread's first fan-out and +/// reused forever after; the caller runs the remaining shard itself. Aggregate is therefore +/// `rayon_threads x (gpu_count - 1)` and falls out of the rayon pool size rather than being a +/// separate knob that can drift from it — shrink the rayon pool and this shrinks in proportion. +/// +/// Replaces `std::thread::scope`, which spawned `gpu_count()` OS threads per row block (~900 k per +/// stem-200 run). That churn cost ~0.01% of wall time, so this is a robustness change, not a +/// throughput one: it keeps the live thread count off `RLIMIT_NPROC` (4096 per-UID, shared +/// machine-wide, against a ~644-thread baseline) and stops the PID churn. +mod shard_pool { + use std::{cell::RefCell, sync::mpsc}; + + type Job = Box; + + thread_local! { + /// This caller's helpers. `RefCell` and not `OnceCell` only because the vector is built + /// lazily; it is never borrowed across a dispatch. + static HELPERS: RefCell>> = const { RefCell::new(Vec::new()) }; + } + + /// Hand `job` to this thread's helper `i`, spawning the helpers on first use. + /// + /// Panics in a job are contained by the helper (its result channel drops, which the caller sees + /// as a receive error) so one bad shard cannot poison a thread other callers depend on — there + /// are none, but it also keeps this caller's later blocks working. + pub(super) fn dispatch(i: usize, job: Job) { + HELPERS.with(|h| { + let mut h = h.borrow_mut(); + while h.len() <= i { + let (tx, rx) = mpsc::channel::(); + let idx = h.len(); + std::thread::Builder::new() + .name(format!("nassau-shard{idx}")) + .spawn(move || { + while let Ok(job) = rx.recv() { + let _ = std::panic::catch_unwind(std::panic::AssertUnwindSafe(job)); + } + }) + .expect("failed to spawn a shard helper thread"); + h.push(tx); + } + h[i].send(job).expect("a shard helper thread died"); + }); + } +} + +/// A/B diagnostic toggle (`NASSAU_GPU_BASIS_PASSTHROUGH=1`): when set, the batched multiply +/// marshals each term's p-part per launch and binds those buffers as the "basis" with an +/// identity index map, reproducing the pre-resident-basis behaviour through the same kernel. +/// Lets a single binary isolate a kernel-signature bug from a resident-basis host/upload bug. +fn basis_passthrough() -> bool { + static ON: LazyLock = + LazyLock::new(|| std::env::var_os("NASSAU_GPU_BASIS_PASSTHROUGH").is_some()); + *ON +} + +/// RAII reservation of `weight` output bytes from [`GPU_BUDGET`]; blocks (parked, not spinning) +/// until the budget admits it. The device section runs on the shared GPU thread (see [`gpu_thread`]), +/// so the permit no longer carries a slot. +struct GpuPermit { + weight: usize, +} + +impl GpuPermit { + fn acquire(weight: usize) -> Self { + let b = &*GPU_BUDGET; + let mut used = b.used.lock().unwrap(); + while !(*used == 0 || *used + weight <= b.budget) { + used = b.freed.wait(used).unwrap(); + } + *used += weight; + Self { weight } + } +} + +impl Drop for GpuPermit { + fn drop(&mut self) { + let b = &*GPU_BUDGET; + *b.used.lock().unwrap() -= self.weight; + b.freed.notify_all(); + } +} /// 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 @@ -62,14 +487,97 @@ fn narrow_u16(v: u32) -> u16 { u16::try_from(v).expect("admissible/term entry exceeds u16") } -use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::{ + Arc, + atomic::{AtomicBool, AtomicU64, Ordering}, +}; + +/// Set once the cubecl CUDA context has failed irrecoverably (a `CUDA_ERROR_LAUNCH_FAILED` / +/// `ServerUnhealthy` surfacing the unresolved cubecl uninit-handle bug — see +/// `~/cubecl-uninit-handle-followup.md`, tracel-ai/cubecl#1401). Such a failure **poisons the whole +/// CUDA context**: every later launch on the shared client fails too, so there is no per-call retry. +/// [`multiply_batch_on_gpu`] latches this on the first failure and then *propagates the panic* — the +/// run dies at the fault rather than silently finishing on the CPU, so a crash cannot masquerade as a +/// slow success. The flag exists purely so in-process observers (the soak test) can distinguish a +/// context death from an ordinary assertion failure. NOTE: this covers only the cubecl **multiply**; +/// the RREF path runs on a separate `fp-cuda` runtime and is not gated by this flag. +static GPU_DISABLED: AtomicBool = AtomicBool::new(false); + +/// Whether the GPU multiply has been disabled for the rest of the process (see [`GPU_DISABLED`]). +pub fn gpu_disabled() -> bool { + GPU_DISABLED.load(Ordering::Relaxed) +} + +/// Global launch counter for throttling per-launch [`memory_cleanup`] (see [`cleanup_every`]). +static CLEANUP_COUNTER: AtomicU64 = AtomicU64::new(0); + +/// How often to call `client.memory_cleanup()` after a multiply launch, via +/// `NASSAU_GPU_CLEANUP_EVERY` (default `1` = every launch). `N` cleans every Nth launch; `0` never +/// cleans. DIAGNOSTIC: the residual `CUDA_ERROR_LAUNCH_FAILED` at high stems is consistent with a +/// cross-stream pool reclaim (one stream's cleanup reclaiming a resident-master page still in flight +/// on another stream — tracel-ai/cubecl#1401). Throttling this drastically cuts that reclaim rate; if +/// the crash disappears or moves much later, cleanup is confirmed as the trigger. The tradeoff is +/// device-memory growth, since freed pages linger — watch `nvidia-smi`. +fn cleanup_every() -> u64 { + use std::sync::OnceLock; + static EVERY: OnceLock = std::sync::OnceLock::new(); + *EVERY.get_or_init(|| { + std::env::var("NASSAU_GPU_CLEANUP_EVERY") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(1) + }) +} /// 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); +/// First-sight `R`s that forced an `admissible_matrices` enumeration + a `RESIDENT_HOST` write +/// lock (see [`resident_info`]). Diffed around the pair pre-pass to attribute its cost. +static RESIDENT_MISSES: AtomicU64 = AtomicU64::new(0); +/// Enum-launch geometry, for sizing the grid against the device. `enumerate_admissible_kernel` is +/// ~99% of GPU kernel time and its benchmark measured `Waves Per SM = 0.44` — i.e. it cannot fill +/// half the machine — but that benchmark enumerates every `R` to degree 130 in ONE launch, while a +/// production launch covers only one block-segment's transient `R`s. These say what production +/// actually submits, which decides whether batching `R`s across launches is worth the scratch memory +/// it would cost. +static ENUM_LAUNCHES: AtomicU64 = AtomicU64::new(0); +static ENUM_RS: AtomicU64 = AtomicU64::new(0); +static ENUM_RS_MAX: AtomicU64 = AtomicU64::new(0); +static ENUM_BLOCKS: 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); +/// `BATCH_MARSHAL_US` split: host CPU work before any blocking, and time parked on the +/// [`GpuPermit`] / [`fp::gpu_lock`] acquisition. Conflating them hid which one dominates. +static BATCH_PREP_US: AtomicU64 = AtomicU64::new(0); +static BATCH_WAIT_US: AtomicU64 = AtomicU64::new(0); +/// `BATCH_WAIT_US` split again: the pre-existing [`GpuPermit`] (bounds in-flight output bytes) +/// versus the cross-runtime [`fp::gpu_lock`] arbitration. They have different owners and fixes. +static BATCH_PERMIT_US: AtomicU64 = AtomicU64::new(0); +static BATCH_LOCK_US: AtomicU64 = AtomicU64::new(0); +/// `BATCH_DEVICE_US` split in two by the dedicated GPU thread (see [`gpu_thread`]): time spent +/// waiting in the submission FIFO versus time the device section actually ran. The single +/// `device` figure could not tell "queued behind other work" from "computing", which is exactly +/// the distinction the 370 s stalls turned on. +static BATCH_QUEUE_US: AtomicU64 = AtomicU64::new(0); +static BATCH_EXEC_US: AtomicU64 = AtomicU64::new(0); +/// Queue depth summed over launches (÷ calls = mean depth) and its high-water mark. +static BATCH_DEPTH_SUM: AtomicU64 = AtomicU64::new(0); +static BATCH_DEPTH_MAX: AtomicU64 = AtomicU64::new(0); +/// The caller's wait, split at the point the GPU worker hands back the readback future: time until +/// this block was *launched* versus time waiting on its completion fence. `launch` counts jobs +/// queued ahead of this one on the worker, `fence` counts the device actually working — so +/// `fence >> launch` means the pipeline is full and `launch >> fence` means it is starved. Kept +/// separate from `queue`/`exec`, which measure the same run from the worker's side. +static BATCH_LAUNCH_US: AtomicU64 = AtomicU64::new(0); +static BATCH_FENCE_US: AtomicU64 = AtomicU64::new(0); +/// `prep` split three ways, to decide whether marshalling is worth a representation change: the +/// per-launch `R` intern, the `term_off`/`ensure_basis` middle, and the `term_gei` fill. Each has a +/// different fix, and the single `prep` number could not distinguish them. +static BATCH_INTERN_US: AtomicU64 = AtomicU64::new(0); +static BATCH_BASIS_US: AtomicU64 = AtomicU64::new(0); +static BATCH_TGEI_US: 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) { @@ -81,9 +589,76 @@ pub fn take_batch_stats() -> (u64, u64, u64, u64) { ) } +/// Where multiply time goes, as microsecond totals: host prep, then the GPU-thread split of +/// queue wait versus device execution, then queue depth (summed, max). See [`gpu_thread`]. +/// +/// Separate from [`take_batch_stats`] because the queue/exec split is the measurement that +/// distinguishes "waiting behind other workers" from "computing" — collapsing them into one +/// `device` figure is what made multi-minute submission stalls read as kernel time. +pub fn take_gpu_timing() -> (u64, u64, u64, u64, u64) { + ( + BATCH_PREP_US.swap(0, Ordering::Relaxed), + BATCH_QUEUE_US.swap(0, Ordering::Relaxed), + BATCH_EXEC_US.swap(0, Ordering::Relaxed), + BATCH_DEPTH_SUM.swap(0, Ordering::Relaxed), + BATCH_DEPTH_MAX.swap(0, Ordering::Relaxed), + ) +} + +/// Diagnostic (see `NASSAU_MEM_REPORT`): resident-master HOST-side heap bytes — the not-yet-uploaded +/// `col_sums`/`masks` tails, the width-padded basis `pparts`/`lens`, and the per-`R` `index` map +/// (its `Vec` keys). The bulk `col_sums`/`masks` are no longer retained (freed after +/// upload — see [`ResidentHost`]); only the pending tail + the `index` persist. Returns `(master, basis)`. +pub fn resident_host_bytes() -> (usize, usize) { + let h = RESIDENT_HOST.read().unwrap(); + let master = h.cs_pending.iter().map(|p| p.capacity()).sum::() * 2 + + h.mk_pending.iter().map(|p| p.capacity()).sum::() * 2 + + h.index.capacity() + * (std::mem::size_of::() + + std::mem::size_of::>() + + 4 * std::mem::size_of::()); + let b = RESIDENT_BASIS_HOST.read().unwrap(); + let basis = b.pparts.capacity() * 2 + b.lens.capacity() * 4 + b.global_base.capacity() * 4; + (master, basis) +} + +/// Diagnostic (see `NASSAU_MEM_REPORT`): DEVICE-side bytes of the resident master (`col_sums`+`masks`, +/// u16) and basis (`pparts` u16 + `lens` u32) — the persistent GPU buffers, from their uploaded +/// element counts. Returns `(master_bytes, basis_bytes)`. +pub fn resident_dev_bytes() -> (usize, usize) { + // Summed over every device: each holds its own replica of the resident master and basis. + let master: usize = RESIDENT_DEV + .iter() + .map(|d| { + let d = d.read().unwrap(); + d.cs.uploaded * 2 + d.mk.uploaded * 2 + }) + .sum(); + let basis: usize = RESIDENT_BASIS_DEV + .iter() + .map(|b| { + let b = b.read().unwrap(); + b.pp.uploaded * 2 + b.ln.uploaded * 4 + }) + .sum(); + (master, basis) +} + +/// Diagnostic (see `NASSAU_MEM_REPORT`): the cubecl CUDA memory pool's device usage on the default +/// device, `(bytes_in_use, bytes_reserved)`. This is the batched-multiply pool; the fp-cuda RREF runs +/// on a separate cudarc context, so `nvidia-smi total − resident_dev − reserved` estimates the RREF +/// pool. Returns `(0, 0)` if the query fails. +pub fn cubecl_device_usage() -> (u64, u64) { + let client = gpu_client(); + match client.memory_usage() { + Ok(u) => (u.bytes_in_use, u.bytes_reserved), + Err(_) => (0, 0), + } +} + use std::{ collections::HashMap, - sync::{LazyLock, Mutex}, + sync::{Condvar, LazyLock, Mutex, RwLock}, }; use cubecl::server::Handle; @@ -93,351 +668,1377 @@ 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, + /// Global element offsets into the shared master (`u64`: the master exceeds `u32::MAX` + /// u16 elements at high stems, so the offset itself must be 64-bit — `multiply_batch_kernel` + /// reads `col_sums`/`masks` at these offsets under 64-bit `address_type`). + cs_off: u64, + mk_off: u64, cs_len: u32, mk_len: u32, num_mats: u32, + /// Which device holds this `R`'s rows. The master is SHARDED, not replicated: each `R` lives on + /// exactly one device, so total device memory is one master spread over `gpu_count()` cards + /// rather than a full copy on each. That is what lifts the memory ceiling (aggregate VRAM + /// instead of per-card VRAM) and keeps the upload cost at 1x rather than Nx. + /// + /// The offsets above are therefore into THIS DEVICE's master, not a global one. + dev: u8, } -/// Process-global resident store of admissible-matrix data, both host- and device-side. +/// Process-shared host master of admissible-matrix data. /// /// Admissible-matrix enumeration is a pure function of `R`'s p-part and the same /// low-degree `R`s recur in essentially every bidegree, so the host master (`col_sums` / /// `masks`, append-only, keyed by p-part in `index`) is enumerated once per distinct `R` -/// and never recomputed. The device copies (`cs_handle` / `mk_handle`) mirror the master -/// and are re-uploaded *only when it grows* — after the `R`s saturate (early in a -/// resolution) launches upload no admissible data at all, cutting the dominant transfer. -/// -/// Guarded by a `Mutex` so the device section serializes across rayon worker threads: each -/// launch runs to its blocking readback before releasing, giving a happens-before edge and -/// no concurrent access — which is what CubeCL's single-device-thread managed-memory model -/// (its `unsafe impl Sync`) requires for a handle created on one thread to be reused on -/// another. Safe as a global because in the GPU path's regime (`p = 2`, trivial profile) -/// `admissible_matrices` depends only on the p-part, not on the algebra instance. -#[derive(Default)] -struct Resident { - col_sums: Vec, - masks: Vec, - index: HashMap, RInfo>, - cs_handle: Option, - mk_handle: Option, - cs_uploaded: usize, - mk_uploaded: usize, -} - -impl Resident { - /// Global offsets/lengths of `R`'s admissible matrices in the master, enumerating and - /// appending them on first sight (the append order fixes the offsets forever). - fn ensure(&mut self, algebra: &MilnorAlgebra, p_part: &[PPartEntry]) -> RInfo { - if let Some(info) = self.index.get(p_part) { - return *info; - } - let (cs_len, mk_len, cs, mk) = algebra.admissible_matrices(p_part); - let info = RInfo { - cs_off: self.col_sums.len() as u32, - mk_off: self.masks.len() as u32, - cs_len: cs_len as u32, - mk_len: mk_len as u32, - num_mats: (mk.len() / mk_len) as u32, - }; - self.col_sums.extend(cs.iter().map(|&v| narrow_u16(v))); - self.masks.extend(mk.iter().map(|&v| narrow_u16(v))); - self.index.insert(p_part.to_vec(), info); - info - } +/// and never recomputed. +/// +/// SHARED, not per-thread: the master reaches many GB at record stems (it grows with the +/// degree), so a thread-local copy per rayon worker multiplies it by the worker count — +/// measured at ~137 GB host / a full 143 GB H200 with 16 workers at stem 150. One copy +/// behind an `RwLock` restores the old shared-mutex footprint: lookups (the overwhelmingly +/// common case once the `R`s saturate) take the read lock, and only a first-sight append +/// takes the write lock — with the enumeration itself done *outside* the lock, so readers +/// never stall behind it. +/// The host master keeps ONLY the not-yet-uploaded tail (`*_pending`) plus a logical length +/// (`*_len`), never the full `col_sums`/`masks`. Once an `R`'s admissible data is copied to the +/// device it is dropped host-side — it is provably never read again (offsets come from `index`; +/// growth uploads only the pending tail; a capacity realloc copies the old *device* buffer, not the +/// host). This removes the multi-GB host↔device duplicate that dominated the resolver's anon RSS +/// (~27 GB at stem 130, growing). Invariant maintained by [`seg_grow`]: +/// `RESIDENT_DEV.$buf.uploaded == $len - $pending.len()`, i.e. `$pending == master[uploaded..$len]`. +/// `*_pending` is per DEVICE. The uploaded prefix is dropped host-side, so what remains is only the +/// tail each device has yet to consume; with several devices each needs its own copy of that tail, +/// because a device replicates the master rather than sharing it. This multiplies the TAIL, not the +/// master: the multi-GB host-side duplicate this design exists to avoid stays gone. +struct ResidentHost { + cs_pending: Vec>, + mk_pending: Vec>, + /// Per-device logical master lengths; an `R` extends only its own device's. + cs_len: Vec, + mk_len: Vec, + /// Master bytes admitted at each `R` degree, indexed by degree. A cumulative sum over this is + /// exactly "how large the master would be at [`resident_degree_cap`] = θ", so a run can report + /// which θ WOULD have fit in a given device budget without ever having to OOM to find out. Kept + /// unconditionally: it is one add per first-sight `R`, and guessing θ is otherwise a blind + /// extrapolation (bytes are dominated by the top degrees, so a lower-stem run cannot predict a + /// higher θ -- the degrees carrying the mass are simply absent from it). + deg_bytes: Vec, + /// Accumulated `num_mats` per device — the work proxy the shard assignment balances. + /// + /// A launch's work on a device is the sum over its products of `num_mats(R) * ceil(nt/T)`, so + /// with `R`s used at broadly similar rates the device's share is set by the `num_mats` it owns. + /// Master bytes are `num_mats * (cs_len + mk_len)`, so balancing this also tracks memory. + dev_load: Vec, + /// std's `HashMap` (SipHash), deliberately. `resident_info` is 7.7% of user cycles in a `perf` + /// profile of an uncapped stem-150 run with a further 2.2% in `DefaultHasher`/`RandomState`, so + /// swapping to `rustc_hash::FxHashMap` looks like free money for a non-adversarial integer key. + /// It is not: measured over three replicates per arm, Fx was ~8% SLOWER + /// (219.0 s mean vs 202.7 s; runs 234/210/213 against 205/208/195). Fx is a weak + /// multiply-rotate and evidently clusters on packed `PPart` keys where SipHash spreads them. + /// Do not re-try this without replicates -- the run-to-run noise floor here is ~3%. + index: HashMap, } -static RESIDENT: LazyLock> = LazyLock::new(|| Mutex::new(Resident::default())); +static RESIDENT_HOST: LazyLock> = LazyLock::new(|| { + RwLock::new(ResidentHost { + cs_pending: (0..gpu_count()).map(|_| Vec::new()).collect(), + mk_pending: (0..gpu_count()).map(|_| Vec::new()).collect(), + cs_len: vec![0; gpu_count()], + mk_len: vec![0; gpu_count()], + dev_load: vec![0; gpu_count()], + deg_bytes: Vec::new(), + index: HashMap::new(), + }) +}); -/// Elementwise F₂ addition of two bit-packed vectors: `out[i] = a[i] ^ b[i]`. +/// Hard cap on devices, so the per-device tables below are a fixed, cheap allocation. +const MAX_GPUS: usize = 8; + +/// How many CUDA devices the multiply path spreads work over. `NASSAU_GPU_DEVICES` overrides; +/// otherwise every device *visible to this process* is used. /// -/// One thread per `u32` limb. F₂ addition is XOR of the packed limbs, so this is -/// the output primitive the multiply kernels accumulate with. -#[cfg(test)] -#[cube(launch)] -fn xor_f2(a: &Array, b: &Array, out: &mut Array) { - if ABSOLUTE_POS < out.len() { - out[ABSOLUTE_POS] = a[ABSOLUTE_POS] ^ b[ABSOLUTE_POS]; - } +/// `/proc/driver/nvidia/gpus` counts the devices physically in the node, which is NOT the same thing: +/// CUDA renumbers the visible subset to `0..n`, so under `CUDA_VISIBLE_DEVICES=1,2,3` the driver +/// exposes ordinals 0..2 while the node still has four entries in `/proc`. Taking the physical count +/// there made the fourth shard open device 3 and panic with `CUDA_ERROR_INVALID_DEVICE` ("invalid +/// device ordinal"), taking the GPU worker threads down with it — so the standard way of partitioning +/// GPUs on a shared node silently broke the run. Honour the mask when it is set. +/// +/// Multi-GPU is worth it here because the single-device run is GPU-bound, not host-bound: whole-run +/// accounting on stem 200 measured 3302 s of device execution against 3931 s wall (84% duty), so +/// eliminating *all* host work would cap out at 1.19x while `629 + 3302/N` predicts 1.72x at N = 2 +/// and 2.70x at N = 4. +fn gpu_count() -> usize { + static N: LazyLock = LazyLock::new(|| { + let physical = std::fs::read_dir("/proc/driver/nvidia/gpus") + .map(|d| d.filter_map(|e| e.ok()).count()) + .unwrap_or(0) + .max(1); + // An empty mask means "no GPUs visible"; a mask listing unparseable or out-of-range entries + // truncates at the first bad one, exactly as CUDA itself does. + let visible = std::env::var("CUDA_VISIBLE_DEVICES").ok().map(|v| { + v.split(',') + .take_while(|e| { + e.trim() + .parse::() + .is_ok_and(|ord| ord < physical.max(MAX_GPUS)) + }) + .count() + }); + let detected = visible.unwrap_or(physical).max(1); + std::env::var("NASSAU_GPU_DEVICES") + .ok() + .and_then(|v| v.parse::().ok()) + .filter(|&n| n > 0) + .unwrap_or(detected) + .clamp(1, MAX_GPUS) + }); + *N } -/// Compute `a ^ b` limb-wise on the default CUDA device. +/// Worker threads — and therefore CUDA streams — per device (`NASSAU_GPU_STREAMS`, default 1). /// -/// Host-side driver for `xor_f2`: uploads both operands, launches one thread per -/// limb, and reads the result back. Panics if the operands differ in length. -#[cfg(test)] -pub fn xor_f2_on_gpu(a: &[u32], b: &[u32]) -> Vec { - assert_eq!(a.len(), b.len(), "operands must have equal limb counts"); - let n = a.len(); - let client = CudaRuntime::client(&CudaDevice::default()); - - let a_handle = client.create_from_slice(u32::as_bytes(a)); - let b_handle = client.create_from_slice(u32::as_bytes(b)); - let out_handle = client.empty(std::mem::size_of_val(a)); - - // One 1-D block of `THREADS` units, enough blocks to cover every limb. - const THREADS: u32 = 256; - let cubes = (n as u32).div_ceil(THREADS); - unsafe { - xor_f2::launch::( - &client, - CubeCount::Static(cubes, 1, 1), - CubeDim::new_1d(THREADS), - ArrayArg::from_raw_parts(a_handle, n), - ArrayArg::from_raw_parts(b_handle, n), - ArrayArg::from_raw_parts(out_handle.clone(), n), - ); - } +/// Each device's submission queue is drained by this many workers, so this many device sections run +/// CONCURRENTLY on one device instead of strictly one after another. It exists because +/// [`enumerate_admissible_kernel`] is ~99% of GPU kernel time and runs the device at +/// `Waves Per SM = 0.002`: a production launch is ~6 blocks of the 3168 an H200 can hold, and its +/// duration is set by its longest single `R` (one thread, sequential odometer), not by how many +/// blocks it occupies. Widening the grid therefore cannot help (measured: 6.5x more blocks bought +/// 3.4%) — but running independent sections *beside* each other can, because it converts a SUM of +/// serialised launches into a MAX of concurrent ones. +/// +/// Default 1 because streams were pinned to 1 for a reason: the CUDA runtime keeps a pinned staging +/// pool PER STREAM, and per-stream pools were half of the ~500 GB host-memory blowup (see the +/// resident-master notes). N streams multiply that pool count by N, so raising this trades host RSS +/// for device concurrency and must be measured, not assumed. +/// +/// MEASURED (S_2 stem 150, max_s 60, theta=125) — and it does NOT pay: +/// +/// | streams | wall | peak host RSS | peak GPU | +/// |---------|-------|---------------|----------| +/// | 1 | 528 s | 39.1 GB | 31.5 GB | +/// | 2 | 527 s | 53.3 GB | 43.9 GB | +/// | 4 | 598 s | 75.9 GB | 68.8 GB | +/// +/// Wall is flat at 2 and 13% WORSE at 4, while both memories roughly double. The growth is the +/// important part: it is proof the sections really did overlap (several sets of transient buffers in +/// flight at once), so this is not "the streams did not engage" — concurrency happened and bought +/// nothing, then started costing (4 streams push the allocator and the pinned pools hard enough to +/// lose 13%). +/// +/// The conclusion is therefore about the workload, not about streams: at this configuration the +/// resolution is NOT GPU-throughput-bound. `enumerate_admissible_kernel` is 99% of GPU *kernel* time, +/// but kernel time is not on the critical path, so the "batching turns a SUM into a MAX" argument +/// above — correct as arithmetic about the launches — optimises something that is not the limiter. +/// Measured at the same time: the process ran at ~918% CPU on a 128-core node (~7% of the machine) +/// with a wavefront only ~5 bidegrees wide, i.e. the limiter is a SERIAL DEPENDENCY CHAIN on the +/// host. Look there before spending anything more on the enum kernel. +/// +/// Kept (rather than reverted) because it is a few lines, defaults to the old behaviour exactly, and +/// is the control that makes the "not GPU-bound" claim falsifiable on a different workload. +/// +/// READ THE CAVEAT BEFORE REUSING THESE NUMBERS. Every row above was measured at +/// `NASSAU_GPU_RESIDENT_MAX_DEGREE=125`, which at stem 150 is an ARTIFICIAL WORST CASE: the same +/// run's `[MASTER-BY-DEGREE]` reports the full master as 1.5 GB (0.4 GB/GPU), so the cap forces the +/// transient re-enumeration path for a master that fits resident several times over. Same binary, +/// same session, stem 150: +/// +/// | theta | wall | enum launches | mean CPU | +/// |-------|-------|---------------|----------| +/// | 125 | 575 s | 10260 | 1110% | +/// | none | 202 s | 0 | 1618% | +/// +/// 2.85x, with the enum kernel not running at all. So roughly two thirds of the wall time in the +/// stream sweep was enum work a sensibly-configured run at this stem never does, and any conclusion +/// drawn from a theta-capped benchmark about where time goes is a conclusion about the cap. +/// +/// Set theta from `[MASTER-BY-DEGREE]`, which reports exactly how many GB each cap would cost: cap +/// only when the master genuinely does not fit, and at high stems cap as high as it does fit. +/// +/// Safe with the shared resident master: [`RESIDENT_DEV`] is indexed per DEVICE, not per thread, so +/// extra workers on one device reuse the same segments rather than replicating them, and the +/// segmented append-only store is already the cross-stream-safe shape (a stable segment written in +/// place) that replaced the churny re-upload cubecl could not synchronise. +fn gpu_streams() -> usize { + static N: LazyLock = LazyLock::new(|| { + std::env::var("NASSAU_GPU_STREAMS") + .ok() + .and_then(|v| v.parse::().ok()) + .filter(|&n| n > 0) + .unwrap_or(1) + .clamp(1, 16) + }); + *N +} - let bytes = client.read_one(out_handle).unwrap(); - u32::from_bytes(&bytes).to_vec() +thread_local! { + /// Which device the current thread's GPU work belongs to. Set once per GPU worker thread; every + /// other thread sees 0 and never touches device state directly. + static CUR_DEVICE: std::cell::Cell = const { std::cell::Cell::new(0) }; } -/// 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. +/// The device this thread's GPU work runs on. Device handles are NOT interchangeable across +/// devices, so every resident-state accessor and every client is keyed by this. +fn cur_device() -> usize { + CUR_DEVICE.with(|c| c.get()) +} + +/// The cubecl client for this thread's device, borrowed for the process's lifetime. /// -/// 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]; - } +/// `'static` on purpose: [`ComputeClient::read_async`] returns a future that borrows the client +/// (edition 2024 RPIT captures `&self`), and that future must outlive the GPU worker's task so the +/// readback can be awaited by the *caller* rather than on the worker — see [`multiply_batch_block`]. +/// A per-call clone would make the future borrow a local and pin the wait to the worker thread, +/// which is exactly the one-kernel-deep pipeline this indirection removes. Constructing each +/// device's client once also drops a `CudaRuntime::client` lookup from every launch. +fn gpu_client() -> &'static cubecl::prelude::ComputeClient { + static CLIENTS: std::sync::OnceLock>> = + std::sync::OnceLock::new(); + &CLIENTS.get_or_init(|| { + (0..gpu_count()) + .map(|d| CudaRuntime::client(&CudaDevice::new(d))) + .collect() + })[cur_device()] +} - // 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 +/// Compile-time cap on the number of fixed-size segments a resident device buffer may hold. It +/// bounds both the multiply kernel's per-buffer argument count and the [`seg_read_u16`] / +/// [`seg_read_u32`] branch depth, so it must be a constant. With `master_seg_elems()` at its default +/// `2^31`, this holds `16 × 4 GiB = 64 GiB` of u16 master per buffer — well past what fits resident +/// on one H200. Raise it (and extend the two `seg_read_*` / the kernel binding) for larger buffers. +const MASTER_MAX_SEG: usize = 16; + +/// Element count per resident segment (see [`SegBuf`]); env `NASSAU_GPU_MASTER_SEG_ELEMS`. Default +/// `2^31` — 4 GiB per u16 segment, 8 GiB per u32 — and deliberately `< u32::MAX`, so a single +/// segment's length never overflows cubecl's 32-bit array-length metadata (the truncation class of +/// bug the u64 offset addressing already guards against). Tests set it tiny to exercise many-segment +/// gathers at low degree; production leaves it large so a run needs only a handful of segments. +fn master_seg_elems() -> usize { + static N: LazyLock = LazyLock::new(|| { + std::env::var("NASSAU_GPU_MASTER_SEG_ELEMS") + .ok() + .and_then(|v| v.parse().ok()) + .filter(|&n| n > 0) + .unwrap_or(1usize << 31) + }); + *N } -/// One thread per padded p_part: `out[i] = seqno(p_parts[i])`. `p_parts` is -/// `n × width` row-major, each row a p_part zero-padded to `width` (padding entries -/// are zero and skipped, so `wlen == width` matches the CPU's trimmed loop). -#[cfg(test)] -#[cube(launch)] -fn seqno_kernel( - g: &Array, - xi: &Array, - p_parts: &Array, - out: &mut Array, - width: usize, -) { - let idx = ABSOLUTE_POS; - if idx >= out.len() { - terminate!(); - } - let base = idx * width; +/// A resident device buffer grown by APPENDING fixed-size segments — the existing segments are never +/// reallocated or copied, so the device peak is `live + one_segment`, not the `~2×` realloc-doubling +/// transient that pushed cubecl into its silent memory-corruption regime (the stem-140+ `dx != 0`). +/// Each segment holds exactly `master_seg_elems()` elements (allocated full; the last is only +/// partially written); the multiply kernel selects the segment for a global offset `o` by a static +/// branch (`o / seg_elems`), so at most [`MASTER_MAX_SEG`] segments exist. Append-only and stable — +/// a segment handle, once allocated and written, never changes identity and is never freed — which +/// is the ordinary shared-global (model-weights) pattern cubecl syncs correctly across streams; the +/// churny "swap the whole buffer on every growth" it replaced broke that sync. `uploaded` is how +/// many elements are physically resident across all segments. +#[derive(Default)] +struct SegBuf { + segs: Vec, + uploaded: usize, +} - 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); +/// Process-shared device mirror of the host master. Each buffer is a segmented, append-only +/// no-copy-growth store (see [`SegBuf`], [`seg_grow`]). This is what makes the master safe to share +/// across streams: the churny "re-upload a new handle on every growth" it replaced broke cubecl's +/// per-handle cross-stream sync (crash rate tracked re-upload frequency); a stable segment written in +/// place is the ordinary shared-global (model-weights) pattern. Reads go through `RESIDENT_DEV.read()` +/// (lock-free fan-out); growth runs outside that lock, serialized only by `RESIDENT_UPLOAD`. +#[derive(Default)] +struct ResidentDev { + cs: SegBuf, + mk: SegBuf, } -/// Run `seqno_kernel` over `n` padded p_parts and return their seqno indices. -/// -/// `g`/`xi` come from `MilnorAlgebra::seqno_table_u32` and -/// [`crate::algebra::combinatorics::xi_degrees`]; `p_parts` is `n × width` row-major, -/// each row a p_part zero-padded to `width`. -#[cfg(test)] -pub fn seqno_batch_on_gpu( - width: usize, - xi: &[u32], - g: &[u32], - p_parts: &[u32], - n: usize, -) -> Vec { - assert_eq!(xi.len(), width, "xi must have `width` entries"); - assert_eq!(p_parts.len(), n * width, "p_parts must be n × width"); - let client = CudaRuntime::client(&CudaDevice::default()); - - let g_h = client.create_from_slice(u32::as_bytes(g)); - let xi_h = client.create_from_slice(u32::as_bytes(xi)); - let pp_h = client.create_from_slice(u32::as_bytes(p_parts)); - let out_h = client.empty(n * size_of::()); - - const THREADS: u32 = 256; - let cubes = (n as u32).div_ceil(THREADS); - unsafe { - seqno_kernel::launch::( - &client, - CubeCount::Static(cubes, 1, 1), - CubeDim::new_1d(THREADS), - ArrayArg::from_raw_parts(g_h, g.len()), - ArrayArg::from_raw_parts(xi_h, xi.len()), - ArrayArg::from_raw_parts(pp_h, p_parts.len()), - ArrayArg::from_raw_parts(out_h.clone(), n), - width, - ); - } +/// One per device: a `Handle` allocated on device `i` is meaningless on device `j`, so the resident +/// master is replicated rather than shared. The HOST master ([`RESIDENT_HOST`]) stays single-copy, +/// which is what keeps the multi-GB host-side duplicate from multiplying by `gpu_count()`. +static RESIDENT_DEV: LazyLock>> = LazyLock::new(|| { + (0..gpu_count()) + .map(|_| RwLock::new(ResidentDev::default())) + .collect() +}); - let bytes = client.read_one(out_h).unwrap(); - u32::from_bytes(&bytes).to_vec() +fn resident_dev() -> &'static RwLock { + &RESIDENT_DEV[cur_device()] } -/// 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; - } +/// Serializes master device *uploads* only — never segment reads. A launch that must grow the +/// device master takes this before uploading, so at a growth point at most one grower runs (others +/// re-check and find it already done) instead of every launch piling redundant copies. Reads go +/// lock-free through `RESIDENT_DEV.read()`, so the upload no longer blocks other bidegrees' device +/// sections (the old single mutex held across the copy collapsed the whole wavefront to one +/// memcpy-ing thread). +static RESIDENT_UPLOAD: LazyLock>> = + LazyLock::new(|| (0..gpu_count()).map(|_| Mutex::new(())).collect()); - let mut working = Array::::new(WORKING_CAP); - let mut rejected = false; +fn resident_upload() -> &'static Mutex<()> { + &RESIDENT_UPLOAD[cur_device()] +} - 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]); - } +/// Shared resident device copies of the read-only seqno table `g` and the (constant) `xi` degrees. +/// These are identical across every launch at a given built degree, so re-uploading them per launch +/// (a `create_from_slice` each) was pure churn — one of the per-launch allocation/copy streams that +/// pushed cubecl's allocator into its `CUDA_ERROR_LAUNCH_FAILED` (719) failure at scale. Uploaded +/// once and re-uploaded only when `g` grows to a new max degree. Keyed by `g.len()`: `g` is a +/// deterministic function of the built degree, so equal length ⇒ identical bytes. Read-only and +/// shared cross-stream exactly like the resident master. +struct SeqnoDev { + g_len: usize, + g: Handle, + xi: Handle, +} +static RESIDENT_SEQNO: LazyLock>>> = + LazyLock::new(|| (0..gpu_count()).map(|_| RwLock::new(None)).collect()); - 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; +fn resident_seqno() -> &'static RwLock> { + &RESIDENT_SEQNO[cur_device()] +} +/// Serializes seqno-table uploads only (never reads); see [`RESIDENT_UPLOAD`]. +static RESIDENT_SEQNO_UPLOAD: LazyLock>> = + LazyLock::new(|| (0..gpu_count()).map(|_| Mutex::new(())).collect()); + +fn resident_seqno_upload() -> &'static Mutex<()> { + &RESIDENT_SEQNO_UPLOAD[cur_device()] +} + +/// Fetch the shared resident `(g, xi)` device handles, uploading only when the cached table's length +/// differs from `$g` (i.e. the built degree changed). Lock-free fast path; a burst of first-sight +/// launches coalesces behind `RESIDENT_SEQNO_UPLOAD`. The upload is synced before publishing so a +/// cross-stream reader never observes the handles before their H2D copy completes. +macro_rules! resident_seqno { + ($client:expr, $g:expr, $xi:expr) => {{ + let read_current = || { + let s = resident_seqno().read().unwrap(); + match &*s { + Some(d) if d.g_len == $g.len() => Some((d.g.clone(), d.xi.clone())), + _ => None, } - if (b & mk) != 0u32 { - rejected = true; + }; + match read_current() { + Some(h) => h, + None => { + let _upload_guard = resident_seqno_upload().lock().unwrap(); + match read_current() { + Some(h) => h, + None => { + let gh = $client.create_from_slice(u32::as_bytes(&$g)); + let xh = $client.create_from_slice(u32::as_bytes(&$xi)); + // Make the copies physically resident before publishing (cross-stream reads). + let _ = cubecl_common::reader::read_sync($client.sync()); + *resident_seqno().write().unwrap() = Some(SeqnoDev { + g_len: $g.len(), + g: gh.clone(), + xi: xh.clone(), + }); + (gh, xh) + } + } } - 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); +/// Host-side cache of cold (degree > [`resident_degree_cap`]) `R`s' admissible-matrix *shape* only — +/// `(cs_len, mk_len, num_mats)`, twelve bytes per `R`. With [in-kernel enumeration](enumerate_admissible_kernel) +/// the cold `col_sums`/`masks` are generated ON the device into transient scratch, so the host never +/// stores (nor uploads) the arrays themselves — only their sizes, needed up front to lay out the +/// scratch offsets and the pair-count prefix sum before the launch. This is the memory win over the +/// old array cache: the evicted tail of the master (tens of GB) lives neither on the device nor the +/// host. The count is computed once per distinct `R` (via `admissible_matrices`, whose arrays are +/// dropped immediately) and memoized, so the per-launch cost is an `O(1)` lookup. +static COLD_COUNT: LazyLock>> = + LazyLock::new(|| RwLock::new(HashMap::new())); + +/// Cold-`R` admissible-matrix shape `(cs_len, mk_len, num_mats)` from the [`COLD_COUNT`] cache. On a +/// miss it runs `admissible_matrices` purely to *count* (the returned arrays are dropped, not kept — +/// the device enumerates them), then memoizes the triple. Layout matches [`resident_info`]'s so the +/// kernel indexes the on-device-enumerated scratch identically to the resident master. +fn cold_count(algebra: &MilnorAlgebra, p_part: PPart) -> (u32, u32, u32) { + if let Some(&e) = COLD_COUNT.read().unwrap().get(&p_part) { + return e; } + let (cs_len, mk_len, _cs, mk) = algebra.admissible_matrices(p_part); + let e = (cs_len as u32, mk_len as u32, (mk.len() / mk_len) as u32); + COLD_COUNT.write().unwrap().entry(p_part).or_insert(e); + e } -/// Multiply `Sq(R) · s` for a single fixed operation `R` into one F₂ output vector. -/// One thread per `(matrix, term)` pair; delegates the assembly to `multiply_pair`. -#[cfg(test)] -#[cube(launch)] -#[allow(clippy::too_many_arguments)] -fn multiply_single_r_kernel( - col_sums: &Array, - masks: &Array, - term_pparts: &Array, - term_lens: &Array, - g: &Array, - xi: &Array, - out: &mut Array>, - num_terms: usize, - num_matrices: usize, - cs_len: usize, - mk_len: usize, - width: usize, -) { - let pair = ABSOLUTE_POS; - if pair >= num_matrices * num_terms { - terminate!(); +/// Grow a segmented resident device buffer ([`SegBuf`]) so it covers `$need` elements, and return +/// `(segments, uploaded)` — the per-segment handles (each `master_seg_elems()` elements) and the +/// logical resident length. NO-COPY growth: existing segments are never reallocated or copied; a +/// growth only allocates the new segment(s) it needs and stage-writes the not-yet-uploaded tail into +/// them. This replaces the `~2×` realloc-doubling transient (old+new buffer both live) that pushed +/// cubecl into memory corruption — the stem-140+ `dx != 0` — with a `live + one_segment` peak. +/// +/// Concurrency mirrors the old in-place path: a lock-free fast path returns the current segments when +/// they already cover `$need`; otherwise `$upload` (a `Mutex`) serializes growers, with a re-check +/// coalescing a burst. The long stage-write runs on a CLONE of the segment vector (segment handles +/// are refcounted, so cloning shares the buffers) and is published under a brief `$dev` write lock, +/// so readers taking `$dev.read()` always see a consistent `(segs, uploaded)` pair. Segments are +/// append-only and never freed, so a reader's cloned handle stays valid for its whole kernel with no +/// realloc barrier. `$tail` is `|uploaded| -> (Vec<$elem>, new_len)`: the owned tail +/// `master[uploaded..new_len]` (the master frees it host-side via `mem::take`; the basis copies it +/// out of its retained store) plus the new logical length. +macro_rules! seg_grow { + ($client:expr, $dev:expr, $field:ident, $upload:expr, $need:expr, + $copy:ident, $as_bytes:path, $elem:ty, $tail:expr) => {{ + let seg_elems = master_seg_elems(); + // Lock-free fast path: current segments already cover `$need`. + let read_current = || { + let dev = $dev.read().unwrap(); + if dev.$field.uploaded >= $need { + Some((dev.$field.segs.clone(), dev.$field.uploaded)) + } else { + None + } + }; + match read_current() { + Some(su) => su, + None => { + let _upload_guard = $upload.lock().unwrap(); + match read_current() { + Some(su) => su, // another grower already covered our need + None => { + // Snapshot the current segments + logical length. We extend a CLONE and + // publish it atomically, so a concurrent reader sees either the whole old + // state or the whole new one — never a half-grown vector. + let (mut segs, uploaded): (Vec, usize) = { + let dev = $dev.read().unwrap(); + (dev.$field.segs.clone(), dev.$field.uploaded) + }; + let (tail, new_len): (Vec<$elem>, usize) = ($tail)(uploaded); + debug_assert_eq!(uploaded + tail.len(), new_len); + // ABORT, do not panic. This runs on a shard's dedicated per-launch thread, + // and a panic there kills only that launch: the next one spawns a fresh + // thread and the resolution carries on having silently dropped the failed + // block's products. Observed on an uncapped stem-300 run -- four launches + // died here and the run continued for another ~900k batches, which produces + // a wrong answer rather than a stopped one. There is no recovery from + // exhausting the segment table mid-master, so take the whole process down + // where it cannot be mistaken for a slow run. + if new_len.div_ceil(seg_elems.max(1)) > MASTER_MAX_SEG { + eprintln!( + "FATAL: resident buffer needs {} segments (> MASTER_MAX_SEG={}); \ + lower NASSAU_GPU_RESIDENT_MAX_DEGREE (see `dump_master_by_degree` \ + for which theta fits), or raise MASTER_MAX_SEG / \ + NASSAU_GPU_MASTER_SEG_ELEMS", + new_len.div_ceil(seg_elems.max(1)), + MASTER_MAX_SEG + ); + ::std::process::abort(); + } + // Allocate (no copy) full-size segments until they cover `new_len`. The last + // one is allocated full even if only partially written; reads only touch + // written locals (`< uploaded`), so its uninitialized tail is never read. + while segs.len() * seg_elems < new_len { + segs.push($client.empty(seg_elems * ::core::mem::size_of::<$elem>())); + } + // Stage-write the tail `master[uploaded..new_len]` into its segments, split at + // segment boundaries and [`STAGE_CHUNK`], syncing each chunk. Existing + // segments (including the partially-filled last one) are appended into, never + // copied. The sync makes each chunk physically resident before the bumped + // `uploaded` is published, so a cross-stream reader never observes a gap + // (cubecl does not order a kernel write to a shared buffer against another + // stream's read the way `create_from_slice` does). + let mut pos = uploaded; + let mut done = 0usize; + while pos < new_len { + let seg = pos / seg_elems; + let local = pos % seg_elems; + let m = (new_len - pos).min(seg_elems - local).min(STAGE_CHUNK); + let scratch = + $client.create_from_slice($as_bytes(&tail[done..done + m])); + copy_chunked!( + $client, $copy, scratch, m, 0usize, segs[seg], seg_elems, local, m + ); + let _ = cubecl_common::reader::read_sync($client.sync()); + pos += m; + done += m; + } + { + let mut dev = $dev.write().unwrap(); + dev.$field.segs = segs.clone(); + dev.$field.uploaded = new_len; + } + (segs, new_len) + } + } + } + } + }}; +} + +/// Global offsets/lengths of `R`'s admissible matrices in the shared host master (see +/// [`ResidentHost`]), enumerating and appending them on first sight (the append order fixes +/// the offsets forever). The enumeration runs outside any lock; on a first-sight race the +/// loser rechecks under the write lock and discards its duplicate. +/// Per-`R` access statistics for the eviction probe (`NASSAU_R_STATS`): how often each distinct `R` +/// is referenced (a block that uses it counts once), its degree, and the first/last reference "time" +/// (a `BATCH_CALLS` tick). Dumped by [`dump_r_stats`] to reveal the hot/cold structure that a device +/// working-set cache would exploit. +#[derive(Clone)] +struct RStat { + count: u64, + degree: i32, + first: u64, + last: u64, + /// Admissible matrices this `R` enumerates. Enumeration cost is proportional to it, so the + /// quantity that actually matters is `count * num_mats` — total matrices enumerated for this + /// `R` over the run — not `count`. A cache policy ranked on references alone is ranking on the + /// wrong axis: a small hot `R` and a big cold one can cost exactly the same to rebuild. + num_mats: u64, +} + +static R_STATS: LazyLock>>> = + LazyLock::new(|| std::env::var_os("NASSAU_R_STATS").map(|_| Mutex::new(HashMap::new()))); + +/// Which device owns `R`'s admissible matrices, from a mixed hash of its packed representation. +/// +/// The mixing is load-bearing, not decoration. [`PPart`] packs entry `i` into a field at a fixed +/// bit offset, so the low bits are `r_1` — which correlates strongly with internal degree, and +/// degree correlates with reference rate (the `NASSAU_R_STATS` hot decile averages degree 70, the +/// cold decile 149 — low-degree `R`s are the hot core). Taking `bits() % gpu_count()` would +/// therefore partition by `r_1 mod 4` and could +/// reproduce the very skew this replaces. The splitmix64 finalizer below spreads every input bit +/// across the output, so the shard is independent of the packing's structure. +/// +/// Deterministic across runs and processes — the same `R` always lands on the same device, which +/// the sharded resident master requires and which keeps a run reproducible. +/// Pin any `R` with at least this many admissible matrices into the RESIDENT master, whatever its +/// degree (`NASSAU_GPU_PIN_MIN_MATS`, default 0 = disabled, i.e. degree is the only criterion). +/// +/// The point is an asymmetry between what costs memory and what costs time. Master BYTES are a SUM +/// over `R`s, but a transient enum launch's DURATION is the MAX over the `R`s in it — ncu at +/// production geometry shows launches of 3-104 blocks taking 7-100 ms at 1.56% occupancy, their +/// length set by the single longest odometer chain. So the long-pole `R`s dominate time while +/// contributing almost nothing to size. Measured over the 75 379 `R`s of a stem-150 run: +/// +/// | threshold | `R`s pinned | extra master | +/// |-----------|-------------|--------------| +/// | 20000 | 126 | ~0.12 GB | +/// | 10000 | 717 | ~0.42 GB | +/// | 5000 | 2297 | ~0.83 GB | +/// | 2000 | 6739 | ~1.34 GB | +/// +/// Against a 143 GB card and a master that reaches 24 GB/GPU at stem 200, that is free. This is the +/// one lever consistent with the max model: de-duplicating enum work across devices removed 36% of +/// all enumerations and bought nothing, because it shortened no chain. +/// +/// Costs one host enumeration per newly-pinned `R` (via `resident_info`), amortised over every later +/// launch that would have re-enumerated it on the device. +fn pin_min_mats() -> u64 { + static T: LazyLock = LazyLock::new(|| { + std::env::var("NASSAU_GPU_PIN_MIN_MATS") + .ok() + .and_then(|v| v.parse::().ok()) + .unwrap_or(0) + }); + *T +} + +fn shard_of(p_part: PPart) -> usize { + (shard_hash(p_part) % gpu_count() as u64) as usize +} + +/// The splitmix64 finalizer applied to `R`'s packed bits. Split out from [`shard_of`] so the +/// uniformity test can bucket a fixed device count rather than whatever the host happens to have. +fn shard_hash(p_part: PPart) -> u64 { + let mut h = p_part.bits(); + h ^= h >> 30; + h = h.wrapping_mul(0xbf58_476d_1ce4_e5b9); + h ^= h >> 27; + h = h.wrapping_mul(0x94d0_49bb_1331_11eb); + h ^= h >> 31; + h +} + +/// Internal degree of `R` from its p-part: `Σ p_part[i] · deg(ξ_{i+1})`. +fn ppart_degree(p_part: PPart) -> i32 { + let xi = xi_degrees(fp::prime::ValidPrime::new(2)); + p_part + .iter() + .zip(xi.iter()) + .map(|(e, &d)| e as i32 * d as i32) + .sum() +} + +/// Operations `R` whose internal degree exceeds this stay OUT of the resident device master and are +/// instead recomputed and uploaded to a throwaway per-launch buffer (see [`MasterMode`]). Default +/// `i32::MAX` keeps every `R` resident — byte-identical to the pre-eviction path (the caller takes a +/// fast path that never touches the transient code). The `NASSAU_R_STATS` probe found a *degree* +/// threshold, not LRU, is the right policy: low-degree `R`s are the stable hot core (reference span +/// 0.99 of the run), high-degree `R`s are scattered-recurring (0.81) *and* the biggest matrices, so +/// excluding them saves more device bytes than their count fraction. On S_2 (150,75): θ≤100 keeps +/// 17% of distinct `R`s resident and recomputes 14% of references; θ≤125 keeps 43% / recomputes 4%. +/// The resident set saturates with degree, so this bounds the master at any stem (the stem-300 lever). +/// +/// SET IT AS HIGH AS DEVICE MEMORY ALLOWS. The reference-miss rate above is a byte metric and badly +/// understates the time cost: a miss re-enumerates on the GPU in EVERY block that touches the `R` +/// (~442 times over a run at (150,75)), and `enumerate_admissible_kernel` is ~99% of GPU kernel time +/// whenever the transient path is live (nsys; `multiply_batch_kernel` is 1.1%). Measured on stem 200 +/// to max_t=310, all complete with 0 crashes: +/// +/// | θ | wall | peak GPU mem | +/// |------|---------|--------------| +/// | ∞ | 2412 s | — | +/// | 200 | 2865 s | 50.2 GB | +/// | 125 | 18566 s | 44.9 GB | +/// +/// θ=125 trades 6.5x in wall time for 5.3 GB. The knee is sharp and sits above 125: the memory curve +/// is far flatter than the time curve, so the cap is a fallback for running out of device memory, +/// not a parameter to tune down. `exec` fell 16x and `fence` 32x from 125 to 200 on identical work +/// (pairs 5.88e13 both), which is the enumeration simply not happening. +/// +/// A smarter eviction policy is not the answer and was measured: replaying a pinned count-ranked +/// cache against the reference stream saves bytes-cached LINEARLY (1% budget -> 2.4%, 25% -> 43.8%), +/// within 0.3pp of a full-hindsight oracle. cost/byte is exactly `count`, so once normalised by the +/// memory it occupies the distribution is flat and no admission rule has anything to exploit. +fn resident_degree_cap() -> i32 { + static CAP: LazyLock = LazyLock::new(|| { + std::env::var("NASSAU_GPU_RESIDENT_MAX_DEGREE") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(i32::MAX) + }); + *CAP +} + +/// Where a launch's `col_sums`/`masks` come from. A given output row's products all share one `R` +/// (the operation of its input basis element), so rows split cleanly by `R`-degree into a resident +/// group and a transient group with no row overlap — the two launches write disjoint rows. +#[derive(Clone, Copy, PartialEq, Eq)] +enum MasterMode { + /// Persist each `R`'s admissible matrices in the shared append-only device master ([`ResidentDev`]). + Resident, + /// Recompute this block's `R`s ([`MilnorAlgebra::admissible_matrices`]) into a per-block buffer, + /// uploaded fresh and freed with the launch. Keeps the resident master bounded across stems. + Transient, +} + +fn record_r_use(algebra: &MilnorAlgebra, p_part: PPart) { + if let Some(m) = R_STATS.as_ref() { + // Inside the probe guard: `cold_count` memoizes, but its first call per `R` enumerates. + // That double-enumerates once per `R` on a probe run and costs nothing when the probe is + // off, which is the right trade for a diagnostic. + let num_mats = cold_count(algebra, p_part).2 as u64; + let now = BATCH_CALLS.load(Ordering::Relaxed); + let mut map = m.lock().unwrap(); + let e = map.entry(p_part).or_insert(RStat { + count: 0, + degree: ppart_degree(p_part), + first: now, + last: now, + num_mats, + }); + e.count += 1; + e.last = now; + } +} + +/// Dump the `R`-access distribution gathered under `NASSAU_R_STATS` (see [`RStat`]) — the data that +/// decides whether/what device eviction policy helps. Prints: reference skew (top-k coverage), +/// degree-vs-frequency correlation, and reference-lifetime spans. No-op unless the probe is enabled. +pub fn dump_r_stats() { + let Some(m) = R_STATS.as_ref() else { return }; + let map = m.lock().unwrap(); + if map.is_empty() { + return; + } + let n = map.len(); + let total_refs: u64 = map.values().map(|s| s.count).sum(); + let now = BATCH_CALLS.load(Ordering::Relaxed).max(1); + let mut v: Vec<&RStat> = map.values().collect(); + // Coverage: sort by count desc, cumulative fraction of references from the top-k Rs. + v.sort_by(|a, b| b.count.cmp(&a.count)); + let cov = |frac: f64| -> f64 { + let k = ((n as f64 * frac).ceil() as usize).max(1).min(n); + let hit: u64 = v[..k].iter().map(|s| s.count).sum(); + hit as f64 / total_refs as f64 * 100.0 + }; + // used-once fraction (pure cold), and how many Rs cover 90% of refs. + let once = v.iter().filter(|s| s.count == 1).count(); + let mut acc = 0u64; + let mut k90 = 0usize; + for s in &v { + acc += s.count; + k90 += 1; + if acc as f64 >= total_refs as f64 * 0.90 { + break; + } + } + // Degree vs frequency: avg degree of the hottest decile vs the coldest decile. + let dec = (n / 10).max(1); + let avg_deg = |slice: &[&RStat]| -> f64 { + slice.iter().map(|s| s.degree as f64).sum::() / slice.len().max(1) as f64 + }; + let hot_deg = avg_deg(&v[..dec]); + let cold_deg = avg_deg(&v[n - dec..]); + // Reference lifetime: span (last-first)/now for the hot decile (are they used throughout, or windowed?). + let hot_span = v[..dec] + .iter() + .map(|s| (s.last - s.first) as f64 / now as f64) + .sum::() + / dec as f64; + let cold_span = v[n - dec..] + .iter() + .map(|s| (s.last - s.first) as f64 / now as f64) + .sum::() + / dec as f64; + let max_deg = v.iter().map(|s| s.degree).max().unwrap_or(0); + let min_deg = v.iter().map(|s| s.degree).min().unwrap_or(0); + // Degree-threshold sizing: for a resident cache that keeps Rs with degree <= θ, what fraction of + // distinct Rs it holds and what fraction of references hit it (miss rate = 100 − ref%). + let mut deg_table = String::new(); + for theta in [50, 75, 100, 125] { + let held = v.iter().filter(|s| s.degree <= theta).count(); + let refs: u64 = v + .iter() + .filter(|s| s.degree <= theta) + .map(|s| s.count) + .sum(); + deg_table += &format!( + " θ≤{theta}:[{:.0}%Rs,{:.0}%refs]", + held as f64 / n as f64 * 100.0, + refs as f64 / total_refs as f64 * 100.0 + ); + } + eprintln!( + "[R-STATS] distinct_R={n} total_refs={total_refs} used_once={once} ({:.0}%) \ + k_for_90%_refs={k90} ({:.1}% of Rs) | coverage top1%={:.0}% top5%={:.0}% top10%={:.0}% \ + top25%={:.0}% | degree hot_decile_avg={:.0} cold_decile_avg={:.0} \ + range=[{min_deg},{max_deg}] | ref_span hot={:.2} cold={:.2} (of run) | degree-threshold \ + cache sizing:{}", + once as f64 / n as f64 * 100.0, + k90 as f64 / n as f64 * 100.0, + cov(0.01), + cov(0.05), + cov(0.10), + cov(0.25), + hot_deg, + cold_deg, + hot_span, + cold_span, + deg_table, + ); + + // The full concentration curve, not just four points: "the hottest x% of `R`s carry y% of all + // references", sampled densely near the head where it bends. Four percentiles were enough to + // see that the distribution is skewed, but not to size a policy against it — a shard assignment + // or an eviction cache behaves very differently if the head is a cliff rather than a slope. + // `v` is already sorted by count descending. + let mut curve = String::new(); + let mut acc = 0u64; + let mut next = 0usize; + // Denser sampling below 10%: that is where essentially all of the curvature lives. + let marks: Vec = (1..=40) + .map(|i| i as f64 * 0.0025) + .chain((1..=18).map(|i| 0.10 + i as f64 * 0.05)) + .collect(); + for (i, s) in v.iter().enumerate() { + acc += s.count; + while next < marks.len() && (i + 1) as f64 / n as f64 >= marks[next] { + curve.push_str(&format!( + "{:.3}:{:.4} ", + marks[next], + acc as f64 / total_refs as f64 + )); + next += 1; + } + } + eprintln!("[R-LORENZ] n={n} total_refs={total_refs} points(x_frac:y_frac) {curve}"); + + // The axis that actually matters. Enumeration cost is proportional to `num_mats`, so an `R`'s + // total rebuild cost over the run is `count * num_mats` — matrices enumerated, not references. + // Ranking a cache on references alone is ranking on the wrong quantity: a small hot `R` and a + // big cold one can cost the same to rebuild, and if that trade is even, the cost distribution + // is FLAT and no partial cache has a head to exploit (only theta, the memory dial, remains). + // Optional per-`R` dump, for offline cache simulation: replaying an admission policy needs the + // individual records, not the summary curves. `first`/`last` are BATCH_CALLS ticks, so a + // simulator can place an `R`'s references across the run rather than only count them. + if let Some(path) = std::env::var_os("NASSAU_R_STATS_CSV") { + use std::io::Write; + match std::fs::File::create(&path) { + Ok(f) => { + let mut w = std::io::BufWriter::new(f); + let _ = writeln!(w, "count,degree,first,last,num_mats"); + for s in map.values() { + let _ = writeln!( + w, + "{},{},{},{},{}", + s.count, s.degree, s.first, s.last, s.num_mats + ); + } + eprintln!("[R-STATS] wrote {} rows to {:?}", map.len(), path); + } + Err(e) => eprintln!("[R-STATS] could not write {path:?}: {e}"), + } + } + + let cost = |s: &RStat| s.count * s.num_mats; + let mut c: Vec<&RStat> = map.values().collect(); + c.sort_by_key(|b| std::cmp::Reverse(cost(b))); + let total_cost: u128 = c.iter().map(|s| cost(s) as u128).sum(); + let ccov = |frac: f64| -> f64 { + let k = ((n as f64 * frac).ceil() as usize).max(1).min(n); + let hit: u128 = c[..k].iter().map(|s| cost(s) as u128).sum(); + hit as f64 / total_cost.max(1) as f64 * 100.0 + }; + let theta = 125; + let (mut t_rs, mut t_refs, mut t_cost) = (0usize, 0u64, 0u128); + for s in &c { + if s.degree > theta { + t_rs += 1; + t_refs += s.count; + t_cost += cost(s) as u128; + } + } + // Within the transient set alone: is its cost concentrated, or spread evenly? + let tv: Vec<&RStat> = c.iter().copied().filter(|s| s.degree > theta).collect(); + let tcov = |frac: f64| -> f64 { + if tv.is_empty() { + return 0.0; + } + let k = ((tv.len() as f64 * frac).ceil() as usize) + .max(1) + .min(tv.len()); + let hit: u128 = tv[..k].iter().map(|s| cost(s) as u128).sum(); + hit as f64 / t_cost.max(1) as f64 * 100.0 + }; + eprintln!( + "[R-COST] total_matrices_enumerated={total_cost} | cost coverage (all Rs) top1%={:.0}% \ + top5%={:.0}% top10%={:.0}% top25%={:.0}% top50%={:.0}% | transient(deg>{theta}): {t_rs} \ + Rs ({:.0}%) {t_refs} refs ({:.0}%) cost {:.0}% of total | within transient, cost \ + coverage top1%={:.0}% top5%={:.0}% top10%={:.0}% top25%={:.0}%", + ccov(0.01), + ccov(0.05), + ccov(0.10), + ccov(0.25), + ccov(0.50), + t_rs as f64 / n as f64 * 100.0, + t_refs as f64 / total_refs as f64 * 100.0, + t_cost as f64 / total_cost.max(1) as f64 * 100.0, + tcov(0.01), + tcov(0.05), + tcov(0.10), + tcov(0.25), + ); +} + +/// Report the master's size as a function of [`resident_degree_cap`]: for each θ, the bytes the +/// resident master would occupy if the cap were set there. Answers "what θ fits in my device +/// budget" from a single run, in place of guessing and discovering the answer as an OOM hours in. +/// +/// Bytes are sharded across [`gpu_count`] devices, so the per-GPU column is what a budget is +/// actually compared against. Note the master is typically a small share of peak device memory +/// (2.6 GB/GPU at stem 150 against a ~50 GB peak at stem 200) -- θ caps the master, not the dense +/// output matrices, so a θ that fits the master is necessary but not sufficient. +pub fn dump_master_by_degree() { + let host = RESIDENT_HOST.read().unwrap(); + if host.deg_bytes.is_empty() { + return; + } + let devs = gpu_count().max(1) as f64; + let total: u64 = host.deg_bytes.iter().sum(); + let mut acc = 0u64; + let mut out = String::new(); + for (d, &b) in host.deg_bytes.iter().enumerate() { + acc += b; + if b != 0 && (d % 25 == 0 || d + 1 == host.deg_bytes.len()) { + out += &format!( + " θ≤{d}:{:.1}GB({:.1}/GPU)", + acc as f64 / 1e9, + acc as f64 / devs / 1e9 + ); + } } - 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, + eprintln!( + "[MASTER-BY-DEGREE] full={:.1}GB ({:.1}GB/GPU over {} devices) | cumulative:{out}", + total as f64 / 1e9, + total as f64 / devs / 1e9, + gpu_count().max(1), ); } +fn resident_info(algebra: &MilnorAlgebra, p_part: PPart) -> RInfo { + record_r_use(algebra, p_part); + if let Some(info) = RESIDENT_HOST.read().unwrap().index.get(&p_part) { + return *info; + } + // Miss: enumerate the admissible matrices (expensive, CPU) and then append under the WRITE + // lock. Counted so [`multiply_batch_grouped`]'s pre-pass can report how many misses it paid + // for — the pre-pass cost is entirely a function of first-sight `R`s, which makes it depend on + // what *other* bidegrees warmed earlier, i.e. on run order rather than on this step's work. + RESIDENT_MISSES.fetch_add(1, Ordering::Relaxed); + let (cs_len, mk_len, cs, mk) = algebra.admissible_matrices(p_part); + let mut host = RESIDENT_HOST.write().unwrap(); + if let Some(info) = host.index.get(&p_part) { + return *info; + } + // Offsets are the running LOGICAL lengths (`*_len`), not the pending-buffer lengths — the + // uploaded prefix has been freed but the logical numbering is permanent (see [`ResidentHost`]). + // Assign this `R` to the least-loaded device. Its rows go there and nowhere else, so a launch + // must route products to the device owning their `R`. + // + // Assignment is a MIXED HASH of `R`, not a load heuristic. See [`shard_of`]. + // + // Two earlier policies both balanced the wrong quantity. Round-robin over first-sight order + // balances the COUNT of `R`s, and `num_mats` varies by orders of magnitude between them. + // Greedy least-loaded by accumulated `num_mats` replaced it and balances master BYTES — which + // its own doc justified by assuming `R`s are "used at broadly similar rates". The + // `NASSAU_R_STATS` probe at (150,110) says otherwise: of 173 930 distinct `R`s, the top 1% + // carry 31% of references and the top 10% carry 78%. Owning equal `num_mats` therefore says + // little about owning equal work, and the error compounds because assignment is permanent — + // measured mean SM at the frontier was 50.5% on device 0 against ~15% on the other three at + // stem 210, a 3.3x spread (1.18x at stem 200, so it worsens with scale). + // + // Hashing balances neither count nor bytes deliberately; it draws each device an INDEPENDENT + // SAMPLE of the joint (size, reference-rate) distribution. With ~174k `R`s over 4 devices and + // no single `R` above a fraction of a percent of references, both quantities concentrate. It is + // also stateless: no accumulator under the write lock, no dependence on first-sight order, and + // no tie-break bias (`min_by_key` resolved ties to device 0, which is where the hot early + // low-degree `R`s landed). + let num_mats = (mk.len() / mk_len) as u64; + let dev = shard_of(p_part); + host.dev_load[dev] += num_mats; + let info = RInfo { + cs_off: host.cs_len[dev] as u64, + mk_off: host.mk_len[dev] as u64, + cs_len: cs_len as u32, + mk_len: mk_len as u32, + num_mats: num_mats as u32, + dev: dev as u8, + }; + let deg = ppart_degree(p_part).max(0) as usize; + if host.deg_bytes.len() <= deg { + host.deg_bytes.resize(deg + 1, 0); + } + host.deg_bytes[deg] += num_mats * (cs_len + mk_len) as u64 * size_of::() as u64; + host.cs_pending[dev].extend(cs.iter().map(|&v| narrow_u16(v))); + host.mk_pending[dev].extend(mk.iter().map(|&v| narrow_u16(v))); + host.cs_len[dev] += cs.len(); + host.mk_len[dev] += mk.len(); + host.index.insert(p_part, info); + info +} + +/// Process-shared host master of the Milnor basis itself, laid out for the device. +/// +/// Every basis element's p-part is stored zero-padded to `width` at `pparts[gei*width ..]`, +/// where `gei` is the element's *global* index (elements concatenated in degree order: +/// all of degree 0, then degree 1, …). `lens[gei]` is its true (trimmed) p-part length, +/// and `global_base[d]` is the number of elements in degrees `< d`, so a term `(s_degree, +/// ti)` maps to `gei = global_base[s_degree] + ti`. +/// +/// This exists so a launch uploads only the small per-term *index* array (`term_gei`) +/// rather than re-gathering and re-uploading every term's padded p-part every launch — the +/// dominant per-launch H2D transfer. The basis is append-only and grows only when a higher +/// degree first appears, so it is uploaded to the device once and re-uploaded only on growth +/// (mirroring [`ResidentHost`]). `built_degree` is the highest degree fully appended. +#[derive(Default)] +struct ResidentBasisHost { + pparts: Vec, + lens: Vec, + global_base: Vec, + built_degree: i32, + width: usize, +} + +static RESIDENT_BASIS_HOST: LazyLock> = + LazyLock::new(|| RwLock::new(ResidentBasisHost::default())); + +/// Device mirror of [`ResidentBasisHost`], both buffers segmented no-copy-growth stores (see +/// [`SegBuf`], [`seg_grow`]). `pp` holds the width-padded p-parts (`elems * width` u16), `ln` the +/// lengths (`elems` u32); the basis element count is `ln.uploaded`. +#[derive(Default)] +struct ResidentBasisDev { + pp: SegBuf, + ln: SegBuf, +} + +static RESIDENT_BASIS_DEV: LazyLock>> = LazyLock::new(|| { + (0..gpu_count()) + .map(|_| RwLock::new(ResidentBasisDev::default())) + .collect() +}); + +fn resident_basis_dev() -> &'static RwLock { + &RESIDENT_BASIS_DEV[cur_device()] +} + +/// Serializes basis device *uploads* only (never handle reads); see [`RESIDENT_UPLOAD`]. +static RESIDENT_BASIS_UPLOAD: LazyLock>> = + LazyLock::new(|| (0..gpu_count()).map(|_| Mutex::new(())).collect()); + +fn resident_basis_upload() -> &'static Mutex<()> { + &RESIDENT_BASIS_UPLOAD[cur_device()] +} + +/// Ensure the resident basis is built through `max_degree` and return a snapshot of +/// `global_base` (so callers compute `gei = global_base[s_degree] + ti` without holding the +/// lock during the parallel marshal). `width` is the fixed p-part padding stride. +/// +/// Append-only: only the first sight of each new degree takes the write lock, and the append +/// order fixes every element's `gei` forever. Basis enumeration is a pure function of the +/// algebra, so a first-sight race just recomputes identical bytes (the loser rechecks under +/// the write lock and appends nothing already present, since we extend strictly past +/// `built_degree`). +fn ensure_basis(algebra: &MilnorAlgebra, width: usize, max_degree: i32) -> Vec { + { + let host = RESIDENT_BASIS_HOST.read().unwrap(); + // `width != 0` distinguishes an initialized store from the derived-`Default` zero state + // (where `built_degree == 0` would spuriously claim degree 0 is already built). + if host.width != 0 && host.built_degree >= max_degree { + return host.global_base.clone(); + } + } + let mut host = RESIDENT_BASIS_HOST.write().unwrap(); + if host.width == 0 { + host.width = width; + host.built_degree = -1; // nothing built yet; the loop below starts at degree 0 + host.global_base.push(0); // global_base[0] = 0 elements before degree 0 + } + debug_assert_eq!(host.width, width, "basis padding width must be stable"); + for d in (host.built_degree + 1)..=max_degree { + let dim = algebra.dimension(d); + for i in 0..dim { + let elt = algebra.basis_element_from_index(d, i); + host.lens.push(elt.p_part.len() as u32); + let base = host.pparts.len(); + host.pparts.resize(base + width, 0); + for (slot, v) in host.pparts[base..base + width] + .iter_mut() + .zip(elt.p_part.iter()) + { + *slot = narrow_u16(v); + } + } + // global_base[d+1] = total elements in degrees ≤ d. + let total = host.lens.len() as u32; + host.global_base.push(total); + } + host.built_degree = max_degree; + host.global_base.clone() +} + +/// Zero a device `u32` buffer on-device: `out[i] = 0`, one thread per limb. +/// +/// Initializes the batched multiply's XOR accumulator without allocating and uploading a host +/// zero buffer. Profiling (stem 145) showed the per-launch `create_from_slice` of a +/// hundreds-of-MB `out_h` zero vec — a host `memset` + non-pinned host→device `memcpy`, both on +/// the calling rayon worker — was the dominant serial marshaling cost, stalling the wavefront. +/// On-device zeroing is memory-bound (microseconds on an H200) and same-stream ordered before +/// the multiply kernel, so no host allocation, upload, or extra sync is needed. +#[cube(launch)] +fn zero_u32(out: &mut [u32]) { + if ABSOLUTE_POS < out.len() { + out[ABSOLUTE_POS] = 0u32; + } +} + +/// Copy `count` elements `src[src_off + i] -> dst[dst_off + i]`, one thread per element. Used to grow +/// the resident master/basis IN PLACE — new data is uploaded to a scratch buffer and copied into the +/// stable resident buffer at its append offset, so the resident device handle never changes (no +/// re-`create_from_slice` churn that would break cross-stream sync). +/// +/// Offsets are `usize` (64-bit on device under the launch's `address_type = "dynamic"`, so both the +/// `dst_off` append offset and the buffer length are safe past `u32::MAX`), and `count` bounds this +/// launch so the caller can split a copy larger than the `u32` grid/`ABSOLUTE_POS` thread limit into +/// chunks (a resident buffer exceeds 2^32 u16 elements around stem 150). See [`copy_chunked`]. +// `launch_unchecked` + dynamic addressing: `dst_off`/buffer length exceed `u32` once the resident +// master/basis passes 2^32 elements, needing 64-bit `usize`; and cubecl's checked bounds clamp emits +// `min(u64, u64)` (ambiguous for NVRTC) under u64. The `ABSOLUTE_POS < count` guard keeps it in-bounds. +#[cube(launch_unchecked, address_type = "dynamic")] +fn copy_into_u16(src: &[u16], dst: &mut [u16], src_off: usize, dst_off: usize, count: u32) { + if ABSOLUTE_POS < usize::cast_from(count) { + dst[dst_off + ABSOLUTE_POS] = src[src_off + ABSOLUTE_POS]; + } +} + +/// `u32` sibling of [`copy_into_u16`] (for the resident basis `lens`). +#[cube(launch_unchecked, address_type = "dynamic")] +fn copy_into_u32(src: &[u32], dst: &mut [u32], src_off: usize, dst_off: usize, count: u32) { + if ABSOLUTE_POS < usize::cast_from(count) { + dst[dst_off + ABSOLUTE_POS] = src[src_off + ABSOLUTE_POS]; + } +} + +/// Elements per copy-kernel launch: below the kernel's `u32` `ABSOLUTE_POS` thread limit, so copies +/// of multi-billion-element resident buffers are split into this many at a time. +const COPY_CHUNK: usize = 1 << 30; + +/// Elements per pinned host-staging chunk when uploading resident growth (see [`seg_grow`]). +/// Bounds the page-locked host buffer cubecl reserves per `create_from_slice`: those pinned pages +/// are pooled PER CUDA STREAM and never trimmed, so a single full-master `create_from_slice` (the +/// tail can be many GB) would pin that whole size on every stream — measured ~240 GB shmem at stem +/// 180 with 8 streams, the OOM driver. Staging in `STAGE_CHUNK` pieces with a sync between them +/// caps the live pinned staging at ~one chunk. 64 Mi × u16 = 128 MiB (× u32 = 256 MiB). +const STAGE_CHUNK: usize = 1 << 26; + +/// Copy `count` elements `src[src_off..] -> dst[dst_off..]` with `$kernel` (`copy_into_u16`/`_u32`), +/// splitting into [`COPY_CHUNK`]-element launches so counts past the `u32` thread limit are handled. +/// `$src_len`/`$dst_len` are the logical array lengths passed to the kernel (must cover the ranges). +macro_rules! copy_chunked { + ($client:expr, $kernel:ident, $src:expr, $src_len:expr, $src_off:expr, + $dst:expr, $dst_len:expr, $dst_off:expr, $count:expr) => {{ + const CT: u32 = 256; + let mut done: usize = 0; + while done < $count { + let n = ($count - done).min(COPY_CHUNK); + unsafe { + $kernel::launch_unchecked::( + &$client, + CubeCount::Static((n as u32).div_ceil(CT), 1, 1), + CubeDim::new_1d(CT), + // The resident dst offset ($dst_off) and buffer length exceed u32 at high stems. + AddressType::from_len(($src_len).max($dst_len).max($dst_off + $count)), + BufferArg::from_raw_parts($src.clone(), $src_len), + BufferArg::from_raw_parts($dst.clone(), $dst_len), + $src_off + done, + $dst_off + done, + n as u32, + ); + } + done += n; + } + }}; +} + +/// 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_packed( + g: &[u32], + xi: &[u32], + pp_shift: &[u32], + pp_mask: &[u32], + working: u64, + wlen: usize, + width: usize, +) -> u32 { + // cur_d = Σ working[h] · xi[h], reading entries out of the packed word. + let mut cur_d = 0u32; + for h in 0..wlen { + let e = + u32::cast_from((working >> u64::cast_from(pp_shift[h])) & u64::cast_from(pp_mask[h])); + cur_d += e * 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 = + u32::cast_from((working >> u64::cast_from(pp_shift[h])) & u64::cast_from(pp_mask[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 +} + +/// 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`). +/// How many columns one `(matrix, term)` pair actually has to visit: past the longest of the three +/// inputs, `b`, `cs` and `mk` are all zero, so [`pair_col`] returns 0 — no rejection and nothing +/// added to `working`. Stopping there is exact, not a truncation. +/// +/// This is per THREAD. The loop used to run to the launch's comptime `work_cap`, which is the max +/// `mk_len` over every `R` in the block, so one long `R` made every thread in the launch pay for +/// columns its own data does not have: measured `work_cap = 16` against a mean `mk_len` of 9.9. +/// Ablation put ~59% of kernel time in this loop (shortening it 16 -> 4 was +80%), so the waste was +/// the single largest item in the kernel. +#[cube] +fn pair_cols(term_len: usize, cs_len: usize, mk_len: usize) -> usize { + let mut cols = cs_len; + if term_len > cols { + cols = term_len; + } + if mk_len > cols { + cols = mk_len; + } + cols +} + +/// Bit [`pair_col`] sets to report that a column rejects the whole product, above the 16 bits the +/// value itself occupies (`diff | mk` and `b | mk` are all widened from `u16`). +/// +/// Packing the flag into the return value rather than signalling it out of band keeps the caller +/// branchless: it ORs the flag into an accumulator and shifts the low half into `working` +/// unconditionally, exactly as the pre-refactor code did. Guarding the accumulate on a rejection +/// test instead costs ~11% — twelve divergent branches per thread, one per column. +const PAIR_COL_REJECT: u32 = 1 << 16; + +/// The per-column rule of [`multiply_pair`], factored out so a caller that reads `b`/`cs`/`mk` +/// from somewhere other than three contiguous slices can reuse it verbatim. Returns the assembled +/// `working[j]` in the low 16 bits, with [`PAIR_COL_REJECT`] set if this column kills the product +/// (in which case the value half is meaningless — the caller discards the whole product). +#[cube] +fn pair_col(j: usize, low: usize, b: u32, cs: u32, mk: u32) -> u32 { + let mut val = 0u32; + if j < low { + if cs > b { + val = PAIR_COL_REJECT; + } else { + let diff = b - cs; + if (diff & mk) != 0u32 { + val = PAIR_COL_REJECT; + } else { + val = diff | mk; + } + } + } else { + if cs > 0u32 { + val = PAIR_COL_REJECT; + } else if (b & mk) != 0u32 { + val = PAIR_COL_REJECT; + } else { + val = b | mk; + } + } + val +} + +/// Tail of [`multiply_pair`] for an accepted product: index the assembled p-part and XOR its +/// F₂ bit into `out`. Split out alongside [`pair_col`] so both callers share it. +#[cube] +#[allow(clippy::too_many_arguments)] +fn pair_emit( + g: &[u32], + xi: &[u32], + out: &mut [Atomic], + working: u64, + row_base: usize, + out_offset: usize, + width: usize, + num_limbs: usize, + #[comptime] sq_len: usize, + pp_shift: &[u32], + pp_mask: &[u32], +) { + // `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). + // Only the first `PPART_MAX_LEN` positions can be non-zero (see `multiply_pair`'s + // accumulator), so the rank loop stops at `sq_len = min(work_cap, PPART_MAX_LEN)`, computed on + // the host (comptime arithmetic does not lower inside a `#[cube]` fn). + let idx = seqno_core_packed(g, xi, pp_shift, pp_mask, working, sq_len, width); + let global_bit = out_offset + usize::cast_from(idx); + let limb = global_bit / 32; + // Device-side mirror of the host's defensive mask: `nassau_gpu::get_partial_matrix_restricted` + // launches at the full output width but masks bits `>= target_dim` on readback because a kept + // block's `out_offset + seqno` can span past it. Skip writes past this row's `num_limbs` — they + // would otherwise overrun into the next row (silent corruption) or past the buffer (an OOB + // atomic; compute-sanitizer confirmed `Invalid __global__ atomic ... out of bounds`). + // Two independent bounds, both required: `limb < num_limbs` keeps the write inside this row + // (out_offset + seqno can span past it), and `word < out.len()` guards the row itself — a + // `row_base` that overruns the buffer (compute-sanitizer caught this as a second OOB atomic at + // higher degree, distinct from the intra-row overflow) would otherwise write past the end. + if limb < num_limbs { + let word = row_base + limb; + if word < out.len() { + let bit = u32::cast_from(global_bit % 32); + out[word].fetch_xor(1u32 << bit); + } + } +} + +#[cube] +#[allow(clippy::too_many_arguments)] +fn multiply_pair( + col_sums: &[u16], + masks: &[u16], + term_pparts: &[u16], + g: &[u32], + xi: &[u32], + out: &mut [Atomic], + 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, + num_limbs: usize, + #[comptime] sq_len: usize, + pp_shift: &[u32], + pp_mask: &[u32], +) { + let mut low = cs_len; + if term_len < cs_len { + low = term_len; + } + + // Packed accumulator instead of `Array::::new(work_cap)`: the single largest slice of + // per-thread state (work_cap x u32 ~= 16 registers of the measured 78, and registers are what + // caps occupancy at 3 blocks/SM = 37.5%). + // + // Entries at index >= PPART_MAX_LEN cannot exist, so stopping there is exact rather than a + // truncation, and it holds by the degree bound rather than by observation: at p = 2 the entry + // r_n multiplies deg(xi_n) = 2^n - 1, so a p-part of length 11 needs degree >= 2^11 - 1 = 2047, + // while `PPart::MAX_DEGREE` is 2045. `MAX_LEN = 10` is therefore forced by that bound, not a + // cap something could exceed — which is also why `PPart::set` can assert `i < MAX_LEN`. + let mut working = 0u64; + let mut rejected = 0u32; + + for j in 0..pair_cols(term_len, cs_len, mk_len) { + 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 val = pair_col(j, low, b, cs, mk); + rejected |= val & PAIR_COL_REJECT; + if j < PPART_MAX_LEN { + working |= u64::cast_from(val & 0xffffu32) << u64::cast_from(pp_shift[j]); + } + } + + if rejected == 0u32 { + pair_emit( + g, xi, out, working, row_base, out_offset, width, num_limbs, sq_len, pp_shift, pp_mask, + ); + } +} + /// Batched multiply: one launch covering all `(R, s)` products of (e.g.) a /// `get_partial_matrix` call. One thread per `(product, matrix, term)` pair. /// @@ -453,27 +2054,153 @@ fn multiply_single_r_kernel( /// /// Output is `num_rows` F₂ vectors of `num_limbs` `u32` limbs, row `r` at /// `out[r*num_limbs ..]`. -#[cube(launch)] +// 64-bit addressing (`address_type = "u64"`): the admissible master and width-padded basis exceed +// `u32::MAX` elements at high stems, and the per-`R` offsets in `r_cs_offset`/`r_mk_offset` (u64) +// index the global (across-segment) master offsets. Static u64 (not "dynamic") because dynamic would +// pick 32-bit `usize` for small blocks and then *narrow* the u64 offset arrays on read (cubecl +// `usize::cast_from(u64)` under a u32 address type), corrupting results — the (180,92) `dx != 0`. +// `launch_unchecked` because cubecl's checked-mode bounds clamp emits `min(u64, u64)`, which NVRTC +// rejects as an ambiguous overload; every access here is in-bounds by construction (the `need_*` +// prefix covers every offset, `seg_read_*` selects the owning segment, and the per-column `j` guards). +// +// The master (`cs*`/`mk*`) and basis (`pp*`/`ln*`) are each a segmented, no-copy-growth store bound as +// [`MASTER_MAX_SEG`] separate segment `Array`s (cubecl has no array-of-buffers). A thread reads its +// one matrix's `col_sums`/`masks` and its term's p-part out of the segments via `seg_read_*` (correct +// for any offset, straddle or not — no layout padding needed), one column at a time, straight into +// the arithmetic. `seg_elems` is the segment element count (`o / seg_elems` picks the segment). +// +// OCCUPANCY: this kernel was register-bound, and the register count is essentially a function of how +// much per-thread state it holds. Two successive removals took it from 78 registers (3 blocks/SM, +// 37.5% — matching a sampled 36.8% "Compute Warps in Flight") to 48 (packing `working` into a `u64`) +// to 40 (fusing the `cs_local`/`mk_local`/`term_local` gather into the column loop), i.e. 6 +// blocks/SM = 75%, with zero spill at every tier. Nothing of `work_cap` length lives in a thread any +// more, so occupancy no longer scales with the internal degree either. Measured +5.6% on the +// stem-200 bench (6.43 -> 6.79 e9 pairs/s, four paired rounds, arms non-overlapping). +// +// That retired a planned upstream cubecl change: cubecl emits only `__launch_bounds__()`, +// and forcing the second argument (`~/cubecl-min-blocks-per-sm.patch`, mirroring `cluster_dim`) was +// worth +6.2% back when ptxas settled at 78 registers. It is moot now — ptxas picks 6 blocks/SM on +// its own, past the 4 the patch would have asked for, so the floor it sets would never bind. +// +// Occupancy is no longer the constraint, and pushing it further is not obviously the next move: an +// intermediate variant reached 35 registers (7 blocks/SM) by branching on the rejection test instead +// of accumulating it, and measured ~11% SLOWER — twelve divergent branches per thread cost more than +// the extra resident warp bought. Unrolling the column loop (`#[unroll]`) was throughput-neutral and +// quadrupled the code, so it is deliberately not used. +#[cube(launch_unchecked, address_type = "u64")] #[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, + cs0: &[u16], + cs1: &[u16], + cs2: &[u16], + cs3: &[u16], + cs4: &[u16], + cs5: &[u16], + cs6: &[u16], + cs7: &[u16], + cs8: &[u16], + cs9: &[u16], + cs10: &[u16], + cs11: &[u16], + cs12: &[u16], + cs13: &[u16], + cs14: &[u16], + cs15: &[u16], + mk0: &[u16], + mk1: &[u16], + mk2: &[u16], + mk3: &[u16], + mk4: &[u16], + mk5: &[u16], + mk6: &[u16], + mk7: &[u16], + mk8: &[u16], + mk9: &[u16], + mk10: &[u16], + mk11: &[u16], + mk12: &[u16], + mk13: &[u16], + mk14: &[u16], + mk15: &[u16], + pp0: &[u16], + pp1: &[u16], + pp2: &[u16], + pp3: &[u16], + pp4: &[u16], + pp5: &[u16], + pp6: &[u16], + pp7: &[u16], + pp8: &[u16], + pp9: &[u16], + pp10: &[u16], + pp11: &[u16], + pp12: &[u16], + pp13: &[u16], + pp14: &[u16], + pp15: &[u16], + ln0: &[u32], + ln1: &[u32], + ln2: &[u32], + ln3: &[u32], + ln4: &[u32], + ln5: &[u32], + ln6: &[u32], + ln7: &[u32], + ln8: &[u32], + ln9: &[u32], + ln10: &[u32], + ln11: &[u32], + ln12: &[u32], + ln13: &[u32], + ln14: &[u32], + ln15: &[u32], + term_gei: &[u32], + g: &[u32], + xi: &[u32], + out: &mut [Atomic], + r_cs_offset: &[u64], + r_mk_offset: &[u64], + r_cs_len: &[u32], + r_mk_len: &[u32], + r_num_mats: &[u32], + prod_r_index: &[u32], + prod_term_start: &[u32], + prod_num_terms: &[u32], + prod_row_base: &[u32], + prod_out_offset: &[u32], + prod_pair_start: &[u32], + prod_coarse: &[u32], width: usize, + seg_elems: usize, + num_limbs: usize, + // Runtime scalar, deliberately NOT `#[comptime]`: it differs block to block, so specialising on + // it forces an NVRTC recompile per distinct value (measured: bench spread widened from 0.7% to + // 5.6%). With the coarse index below it is `ceil(log2(chunk span))`, typically a couple of + // steps rather than the ~15 a full search over every product needed. + search_iters: usize, + // Comptime is right here: at most `MASTER_MAX_SEG` distinct values, so the select chain folds + // to the segments that exist without a recompile storm. + #[comptime] num_segs: usize, + pp_shift: &[u32], + pp_mask: &[u32], + // `min(work_cap, PPART_MAX_LEN)`: how far the packed rank loop runs. + #[comptime] sq_len: usize, + // Per-thread column count, specialised per launch to what THIS block actually needs + // (`max(mk_len, term_len)`, rounded up), not the global worst case. + // + // This used to be the kernel's dominant cost: every thread held four arrays of this length + // (`working` u32 plus `cs_local`/`mk_local`/`term_local` u16), so the constant set register + // pressure and hence occupancy — measured 36% compute warps in flight with 64% of warp slots + // unallocated at the old fixed 32. Shrinking it to what the data needs measured +28% + // (5.97 -> 7.64 e9 pairs/s). All four arrays are gone now (`working` packs into a `u64`, the + // other three were fused away into the loop below), so what is left is the trip count of a + // loop over scalars; keeping it tight still shortens that loop, but it no longer gates + // occupancy. It stays comptime so the trip count is a literal rather than a loaded scalar. + // + // It must NOT be hardcoded. `mk_len = rows + cols - 1` grows with internal degree: 16 suffices + // to t~510, but a 9th xi appears at t>=511 and it becomes 17, then 18 past 1023. A fixed 16 + // would silently truncate at stem 300 — wrong answers, no error. Deriving it per launch keeps + // the occupancy win at every degree, and the host asserts it fits [`WORKING_CAP`]. ) { let k = ABSOLUTE_POS; let num_products = prod_pair_start.len() - 1; @@ -482,11 +2209,23 @@ fn multiply_batch_kernel( } // 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 { + // so `prod_pair_start` is strictly increasing and `p` is unique). + // + // `prod_coarse` brackets the answer before the search starts: entry `ci` is the product owning + // pair `ci << COARSE_LOG`, so `p` is in `prod_coarse[ci] ..= prod_coarse[ci + 1]`. Products are + // ordered and every product owns >= 1 pair, so the bracket is valid; the sentinel entry keeps + // `ci + 1` readable for the final chunk. + // + // Each search step is a DEPENDENT global load of `prod_pair_start[mid]` — a full latency stall + // before a thread can touch its own data — and ablation put the unbracketed search at ~12% of + // kernel time. Two cheap loads replace ~15 dependent ones. + let ci = k >> COARSE_LOG; + let mut lo = usize::cast_from(prod_coarse[ci]); + let mut hi = usize::cast_from(prod_coarse[ci + 1]) + 1; + if hi > num_products { + hi = num_products; + } + for _ in 0..search_iters { if hi - lo > 1 { let mid = (lo + hi) / 2; if usize::cast_from(prod_pair_start[mid]) <= k { @@ -499,127 +2238,217 @@ fn multiply_batch_kernel( let p = lo; let ri = usize::cast_from(prod_r_index[p]); + let p_start = usize::cast_from(prod_pair_start[p]); + let local = k - p_start; + // MATRIX varies fastest, term slowest. The obvious decode (`m = local / nt`, `t = local % nt`) + // has the opposite order, and it is the kernel's dominant cost: consecutive threads then share a + // matrix but each takes a DIFFERENT term, whose p-part sits at `gei * width` for an arbitrary + // basis index -- a 32-way scatter across a multi-GB resident basis, 2 useful bytes per 32-byte + // sector fetched. Ablation measured the whole kernel at +79% with the scatter removed (every + // thread reading offset 0) and +0.3% with the entire seqno rank loop deleted, so locality is + // essentially all of the remaining time. + // + // With `m` fastest, a warp shares one term -- its p-part read becomes a broadcast -- and the + // `col_sums`/`masks` reads become `m * cs_len` apart, i.e. one contiguous fully-used run per + // warp instead of one wasted sector per lane. + // + // This is a permutation of the same `(m, t)` set: every thread still owns exactly one pair, and + // the output is XOR-accumulated, so the result is unchanged. + // `num_mats` by lookup, not `(pair_span) / nt`: one load instead of an emulated integer + // division (I2F/MUFU.RCP/F2I plus fixups). Only ONE division remains in the decode -- `t` and + // `m` share it, since ptxas emits a single divide plus an IMAD for the remainder. + let num_mats = usize::cast_from(r_num_mats[ri]); 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; + + // A thread covers a TILE of `MATRIX_GROUP` matrices x `TERM_GROUP` terms. + // + // `col_sums`/`masks` depend only on the matrix and a term's p-part only on the term, so a + // MxT tile reads `2M + T` values per column to evaluate `M*T` pairs -- 1.17 loads per pair at + // 2x3, against 1.67 at 1x3 and 3 at 1x1. The kernel is issue-limited on integer work (ncu: + // 75% SM vs 1.9% DRAM, ALU top pipeline), so fewer loads and fewer addresses is the lever. + // + // The two axes are NOT symmetric. Terms are few (`nt ~ 5`), so the ragged tail dominates the + // choice of `TERM_GROUP` -- see its doc comment, where 4 loses to 3 purely on wasted lanes. + // Matrices are many (`num_mats ~ 20 000`), so a partial matrix tile costs a few idle lanes out + // of thousands and `MATRIX_GROUP` is free to follow the load arithmetic instead. + let mg_count = num_mats.div_ceil(MATRIX_GROUP); + let mg = local % mg_count; + let tg = local / mg_count; + let m_base = mg * MATRIX_GROUP; + let t_base = tg * TERM_GROUP; 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, - ); -} + let cs_base = usize::cast_from(r_cs_offset[ri]); + let mk_base = usize::cast_from(r_mk_offset[ri]); + let ts_base = usize::cast_from(prod_term_start[p]) + t_base; -/// Compute `Sq(R) · s` on the GPU for a single operation `R = (r_degree, r_idx)`, -/// returning the F₂ result as bit-packed `u32` limbs (bit `i` = basis index `i`). -/// -/// `term_indices` are the nonzero indices of `s` in the degree-`s_degree` basis. -/// `R` must be non-empty (`Sq(∅) = 1` is the trivial identity the caller handles). -/// Requires the algebra's basis and seqno tables built through `r_degree + s_degree`. -#[cfg(test)] -pub fn multiply_single_r_on_gpu( - algebra: &MilnorAlgebra, - r_degree: i32, - r_idx: usize, - s_degree: i32, - term_indices: &[usize], -) -> Vec { - let (width, g) = algebra.seqno_table_u32(); - // Pad `xi` to `WORKING_CAP` so the kernel's `cur_d` sum (which runs to the full - // working capacity) never reads out of bounds; padding entries multiply zero. - let mut xi: Vec = xi_degrees(algebra.prime()) - .iter() - .map(|&x| x as u32) - .collect(); - xi.resize(WORKING_CAP, 0); + // Per-term p-part offsets and lengths. Lanes past `nt` carry `term_len = 0` and are excluded + // from the output below: an all-zero term against an all-zero column does NOT reject, so they + // would otherwise emit a spurious `seqno(0)` bit. + let mut pp_off = Array::::new(TERM_GROUP); + let mut term_len = Array::::new(TERM_GROUP); + let mut cols = cs_len; + if mk_len > cols { + cols = mk_len; + } + #[unroll] + for tt in 0..TERM_GROUP { + let mut po = 0u64; + let mut tl = 0u32; + if t_base + tt < nt { + let gei = usize::cast_from(term_gei[ts_base + tt]); + po = u64::cast_from(gei * width); + tl = seg_read_u32( + ln0, ln1, ln2, ln3, ln4, ln5, ln6, ln7, ln8, ln9, ln10, ln11, ln12, ln13, ln14, + ln15, gei, seg_elems, num_segs, + ); + } + pp_off[tt] = po; + term_len[tt] = tl; + if usize::cast_from(tl) > cols { + cols = usize::cast_from(tl); + } + } - let r = algebra.basis_element_from_index(r_degree, r_idx); - assert!( - !r.p_part.is_empty(), - "R must be non-empty (Sq(∅) = 1 is the identity)" - ); - let (cs_len, mk_len, cs32, mk32) = algebra.admissible_matrices(&r.p_part); - // Ship admissible-matrix / term data as u16 (see `multiply_batch_on_gpu`). - let mut col_sums: Vec = cs32.iter().map(|&v| narrow_u16(v)).collect(); - let masks: Vec = mk32.iter().map(|&v| narrow_u16(v)).collect(); - let num_matrices = masks.len() / mk_len; - - // Terms of s, each p_part padded to `width`, with their true (trimmed) lengths. - let num_terms = term_indices.len(); - let mut term_pparts = vec![0u16; num_terms * width]; - let mut term_lens = vec![0u32; num_terms]; - for (t, &ti) in term_indices.iter().enumerate() { - let elt = algebra.basis_element_from_index(s_degree, ti); - term_lens[t] = elt.p_part.len() as u32; - for (slot, &v) in term_pparts[t * width..(t + 1) * width] - .iter_mut() - .zip(&elt.p_part) - { - *slot = narrow_u16(v); - } - } - - let out_degree = r_degree + s_degree; - let dim = algebra.dimension(out_degree); - let num_limbs = dim.div_ceil(32).max(1); - - // Device buffers must be non-empty; `cs_len == 0` (R's max entry is 1) leaves - // `col_sums` empty. The kernel never reads past the real lengths. - if col_sums.is_empty() { - col_sums.push(0); - } - - let client = CudaRuntime::client(&CudaDevice::default()); - let cs_h = client.create_from_slice(u16::as_bytes(&col_sums)); - let mk_h = client.create_from_slice(u16::as_bytes(&masks)); - let tp_h = client.create_from_slice(u16::as_bytes(&term_pparts)); - let tl_h = client.create_from_slice(u32::as_bytes(&term_lens)); - let g_h = client.create_from_slice(u32::as_bytes(&g)); - let xi_h = client.create_from_slice(u32::as_bytes(&xi)); - let zeros = vec![0u32; num_limbs]; - let out_h = client.create_from_slice(u32::as_bytes(&zeros)); - - let total_pairs = num_matrices * num_terms; - const THREADS: u32 = 256; - let cubes = (total_pairs as u32).div_ceil(THREADS).max(1); - unsafe { - multiply_single_r_kernel::launch::( - &client, - CubeCount::Static(cubes, 1, 1), - CubeDim::new_1d(THREADS), - ArrayArg::from_raw_parts(cs_h, col_sums.len()), - ArrayArg::from_raw_parts(mk_h, masks.len()), - ArrayArg::from_raw_parts(tp_h, term_pparts.len()), - ArrayArg::from_raw_parts(tl_h, term_lens.len()), - ArrayArg::from_raw_parts(g_h, g.len()), - ArrayArg::from_raw_parts(xi_h, xi.len()), - ArrayArg::from_raw_parts(out_h.clone(), num_limbs), - num_terms, - num_matrices, - cs_len, - mk_len, - width, - ); + let mut working = Array::::new(MATRIX_GROUP * TERM_GROUP); + let mut rejected = Array::::new(MATRIX_GROUP * TERM_GROUP); + #[unroll] + for i in 0..MATRIX_GROUP * TERM_GROUP { + working[i] = 0u64; + rejected[i] = 0u32; + } + + for j in 0..cols { + // One `col_sums`/`masks` pair per matrix in the tile, shared by every term. + let mut cs = Array::::new(MATRIX_GROUP); + let mut mk = Array::::new(MATRIX_GROUP); + #[unroll] + for mm in 0..MATRIX_GROUP { + let mut c = 0u32; + let mut k = 0u32; + if m_base + mm < num_mats { + if j < cs_len { + c = u32::cast_from(seg_read_u16( + cs0, + cs1, + cs2, + cs3, + cs4, + cs5, + cs6, + cs7, + cs8, + cs9, + cs10, + cs11, + cs12, + cs13, + cs14, + cs15, + cs_base + (m_base + mm) * cs_len + j, + seg_elems, + num_segs, + )); + } + if j < mk_len { + k = u32::cast_from(seg_read_u16( + mk0, + mk1, + mk2, + mk3, + mk4, + mk5, + mk6, + mk7, + mk8, + mk9, + mk10, + mk11, + mk12, + mk13, + mk14, + mk15, + mk_base + (m_base + mm) * mk_len + j, + seg_elems, + num_segs, + )); + } + } + cs[mm] = c; + mk[mm] = k; + } + + // One p-part read per term in the tile, shared by every matrix. + #[unroll] + for tt in 0..TERM_GROUP { + let tl = usize::cast_from(term_len[tt]); + let mut b = 0u32; + if j < tl { + b = u32::cast_from(seg_read_u16( + pp0, + pp1, + pp2, + pp3, + pp4, + pp5, + pp6, + pp7, + pp8, + pp9, + pp10, + pp11, + pp12, + pp13, + pp14, + pp15, + usize::cast_from(pp_off[tt]) + j, + seg_elems, + num_segs, + )); + } + let mut low = cs_len; + if tl < cs_len { + low = tl; + } + #[unroll] + for mm in 0..MATRIX_GROUP { + let val = pair_col(j, low, b, cs[mm], mk[mm]); + let i = mm * TERM_GROUP + tt; + rejected[i] |= val & PAIR_COL_REJECT; + if j < PPART_MAX_LEN { + working[i] |= u64::cast_from(val & 0xffffu32) << u64::cast_from(pp_shift[j]); + } + } + } } - let bytes = client.read_one(out_h).unwrap(); - u32::from_bytes(&bytes).to_vec() + #[unroll] + for mm in 0..MATRIX_GROUP { + #[unroll] + for tt in 0..TERM_GROUP { + let i = mm * TERM_GROUP + tt; + if m_base + mm < num_mats { + if t_base + tt < nt { + if rejected[i] == 0u32 { + pair_emit( + g, + xi, + out, + working[i], + usize::cast_from(prod_row_base[p]), + usize::cast_from(prod_out_offset[p]), + width, + num_limbs, + sq_len, + pp_shift, + pp_mask, + ); + } + } + } + } + } } /// One `Sq(R) · s` product of a batched launch, written into output row `row` at bit @@ -631,31 +2460,546 @@ pub fn multiply_single_r_on_gpu( /// 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). +#[derive(Clone)] pub struct GpuProduct { pub r_degree: i32, pub r_idx: usize, pub s_degree: i32, - pub term_indices: Vec, + /// `Arc<[usize]>`, not `Vec`, purely so cloning a `GpuProduct` is a refcount bump. + /// + /// The terms are written once at construction and only ever read afterwards, but products get + /// cloned twice on the way to the device — once to compact rows into a dense range per + /// hot/cold group, once to fan out into per-device buckets — and with a `Vec` each of those + /// duplicated every term list. A call-graph profile of an uncapped stem-150 run put 5.45% of + /// all user cycles in `_int_free` under the drop of these vectors alone (16.1% total in the + /// allocator). Sharing makes the clones free and the drops O(1). + pub term_indices: std::sync::Arc<[usize]>, pub row: usize, pub out_offset: usize, } -/// Compute a whole batch of `Sq(R) · s` products in a single GPU launch — the -/// The batched unit of one `get_partial_matrix` call. `R`s may differ (each contributes its -/// own admissible matrices). Returns `num_rows` F₂ vectors, each `⌈num_cols/32⌉` -/// bit-packed `u32` limbs. +/// Compute a whole batch of `Sq(R) · s` products on the GPU — the batched unit of one +/// `get_partial_matrix` call, split into row blocks of at most [`gpu_block_bytes`] of output +/// each (see [`multiply_batch_block`]). `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`). +/// One batch multiply's result, held as the D2H landing buffers themselves. +/// +/// The device write has to land somewhere; everything after that is waste. The original form +/// allocated a fresh `Vec` per row and copied the whole output into freshly-mapped pages right +/// after the device had written it — ~32 M allocations over a stem-200 resolution. The measured +/// best-case launch cost scaled linearly with output bytes at only 2.1-5.0 GB/s (256 MiB in +/// 130 ms), far under PCIe 5.0 x16, with the GPU idle throughout. +/// +/// So this keeps cubecl's [`Bytes`] (which may already be pinned — see `AllocationProperty`) and +/// hands out row slices as views. One block per bounded launch, in row order; the owned +/// constructor covers the eviction merge and the CPU oracle, which must accumulate. +pub struct BatchOutput { + /// One landing buffer per row-block, in row order. + blocks: Vec, + num_limbs: usize, +} + +impl BatchOutput { + /// Wrap the per-block landing buffers (zero copy). + fn from_blocks(blocks: Vec, num_limbs: usize) -> Self { + Self { blocks, num_limbs } + } + + /// Wrap owned row-major limbs (eviction merge, CPU oracle). + pub fn from_limbs(limbs: Vec, num_limbs: usize) -> Self { + Self { + blocks: vec![Bytes::from_elems(limbs)], + num_limbs, + } + } + + /// Build from per-row limb vectors (test/reference helper). + pub fn from_rows(rows: &[Vec], num_limbs: usize) -> Self { + Self::from_limbs(rows.concat(), num_limbs) + } + + /// Limbs per row. + pub fn num_limbs(&self) -> usize { + self.num_limbs + } + + /// Number of rows across all blocks. + pub fn rows(&self) -> usize { + if self.num_limbs == 0 { + return 0; + } + self.blocks.iter().map(|b| b.len() / 4).sum::() / self.num_limbs + } + + /// Row limb-slices in row order, as views into the landing buffers. + pub fn iter_rows(&self) -> impl Iterator { + let n = self.num_limbs; + self.blocks + .iter() + .flat_map(move |b| u32::from_bytes(b).chunks_exact(n)) + } +} + +impl PartialEq for BatchOutput { + fn eq(&self, other: &Self) -> bool { + self.num_limbs == other.num_limbs && self.iter_rows().eq(other.iter_rows()) + } +} + +impl Eq for BatchOutput {} + +impl std::fmt::Debug for BatchOutput { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("BatchOutput") + .field("rows", &self.rows()) + .field("num_limbs", &self.num_limbs) + .finish() + } +} + pub fn multiply_batch_on_gpu( + algebra: &Arc, + num_cols: usize, + num_rows: usize, + products: &[GpuProduct], +) -> BatchOutput { + // The CPU fallback that used to live here (catch the launch failure, latch [`GPU_DISABLED`], + // finish the run on the CPU) was removed deliberately: it turned a hard GPU fault into a silent + // ~100x slowdown, so a crashing run still reported "completed" and every A/B measurement had to + // be reconstructed by grepping stderr. A context death is now loud — the panic propagates and + // the run dies at the fault. [`GPU_DISABLED`] is still latched first so in-process observers + // (the soak test) can tell a context death from an ordinary assertion failure. + match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + multiply_batch_gpu_inner(algebra, num_cols, num_rows, products, resident_degree_cap()) + })) { + Ok(out) => out, + Err(payload) => { + // compare_exchange so exactly one thread (of the ~100 that may fail together on the + // shared poisoned context) prints the notice; the rest just resume unwinding. + if GPU_DISABLED + .compare_exchange(false, true, Ordering::Relaxed, Ordering::Relaxed) + .is_ok() + { + eprintln!( + "[nassau-gpu] GPU milnor multiply failed (CUDA context poisoned); failing the \ + run instead of falling back to the CPU. RREF (separate fp-cuda runtime) is \ + unaffected by this flag." + ); + } + std::panic::resume_unwind(payload) + } + } +} + +/// CPU reference for one [`multiply_batch_on_gpu`] batch: the exact same `num_rows × ⌈num_cols/32⌉` +/// bit-packed F₂ matrix the GPU produces, computed with [`MilnorAlgebra::multiply_basis_element_by_element_2`]. +/// Used as the stopgap fallback when the GPU context dies mid-run, and as a correctness oracle for the +/// GPU stress bench. Each product's `Sq(R)·s` lands in row `prod.row` at column offset `prod.out_offset`. +pub fn cpu_multiply_batch( algebra: &MilnorAlgebra, num_cols: usize, num_rows: usize, products: &[GpuProduct], -) -> Vec> { +) -> BatchOutput { + use fp::vector::FpVector; + let p = algebra.prime(); + let num_limbs = num_cols.div_ceil(32).max(1); + let mut rows = vec![vec![0u32; num_limbs]; num_rows]; + // (built per row, then flattened to the shared BatchOutput layout below) + for prod in products { + let out_degree = prod.r_degree + prod.s_degree; + let block_dim = algebra.dimension(out_degree); + if block_dim == 0 { + continue; + } + let s_dim = algebra.dimension(prod.s_degree); + let mut s = FpVector::new(p, s_dim); + for &ti in prod.term_indices.iter() { + s.set_entry(ti, 1); + } + let mut tmp = FpVector::new(p, block_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(), + ); + for (i, _) in tmp.iter_nonzero() { + let col = prod.out_offset + i; + rows[prod.row][col / 32] ^= 1u32 << (col % 32); + } + } + BatchOutput::from_limbs(rows.concat(), num_limbs) +} + +fn multiply_batch_gpu_inner( + algebra: &Arc, + num_cols: usize, + num_rows: usize, + products: &[GpuProduct], + // Passed in rather than read from [`resident_degree_cap`]: that is a process-wide `LazyLock` over + // an env var, so a test could only ever exercise ONE cap per process — and the eviction split has + // three distinct regimes (all-resident, all-transient, mixed) that must each be checked. + cap: i32, +) -> BatchOutput { + // Fast path (default, `cap == i32::MAX`, and any run whose `R`s are all under the cap): a single + // resident-master pass, byte-identical to the pre-eviction code. No cloning, no second launch. + let num_limbs_all = num_cols.div_ceil(32).max(1); + if cap == i32::MAX || products.iter().all(|p| p.r_degree <= cap) { + return BatchOutput::from_blocks( + multiply_batch_grouped(algebra, num_cols, num_rows, products, MasterMode::Resident), + num_limbs_all, + ); + } + // Eviction active. Each output row's products all share one `R` (see [`MasterMode`]), so the + // hot (degree ≤ cap) and cold row sets are DISJOINT. Run each group on its own rows only, + // **compacted** to a dense `0..k` range so each pass reads back just its own rows — total + // readback stays `num_rows`, not 2× (critical in the intended high-θ regime, where the cold + // set is a small tail and a full-height cold readback would be almost all zeros). Results + // scatter back to the original row indices; a row in neither group stays zero (its content, if + // any, comes from the caller's CPU identity path). + let num_limbs = num_cols.div_ceil(32).max(1); + let mut result = vec![0u32; num_rows * num_limbs]; + // Above-cap `R`s that are long enough to be pinned resident anyway (see [`pin_min_mats`]). + // Resolved once per call rather than per product: `cold_count` memoizes, but the lookup still + // costs a lock and a hash, and the same `R` recurs across most products. + let pin = pin_min_mats(); + let pinned = |d: i32, r_idx: usize| -> bool { + if pin == 0 || d <= cap { + return false; + } + let r = algebra.basis_element_from_index(d, r_idx); + cold_count(algebra, r.p_part).2 as u64 >= pin + }; + for mode in [MasterMode::Resident, MasterMode::Transient] { + let is_group = |d: i32, r_idx: usize| match mode { + MasterMode::Resident => d <= cap || pinned(d, r_idx), + MasterMode::Transient => d > cap && !pinned(d, r_idx), + }; + // Distinct rows this group touches, in order (products are row-major, so already sorted). + let mut rows: Vec = products + .iter() + .filter(|p| is_group(p.r_degree, p.r_idx)) + .map(|p| p.row) + .collect(); + rows.dedup(); + if rows.is_empty() { + continue; + } + let remap: HashMap = rows.iter().enumerate().map(|(i, &r)| (r, i)).collect(); + let compact: Vec = products + .iter() + .filter(|p| is_group(p.r_degree, p.r_idx)) + .map(|p| { + let mut q = p.clone(); + q.row = remap[&p.row]; + q + }) + .collect(); + let sub_blocks = multiply_batch_grouped(algebra, num_cols, rows.len(), &compact, mode); + // Scatter straight out of the device blocks. The first cut flattened them into one + // `Vec` first, which allocated and copied the ENTIRE group output — hundreds of MB at + // the stems that actually need eviction — only to read it once and drop it. Blocks are + // whole rows of `num_limbs` each, in compacted row order, so walking them in `num_limbs` + // chunks visits exactly the rows `rows` names, in the same order. + let mut i = 0usize; + for b in &sub_blocks { + for chunk in u32::from_bytes(b).chunks(num_limbs) { + let dst = rows[i] * num_limbs; + for (k, &v) in chunk.iter().enumerate() { + result[dst + k] ^= v; + } + i += 1; + } + } + debug_assert_eq!( + i, + rows.len(), + "block rows must cover the compacted row set exactly" + ); + } + BatchOutput::from_limbs(result, num_limbs) +} + +fn multiply_batch_grouped( + algebra: &Arc, + num_cols: usize, + num_rows: usize, + products: &[GpuProduct], + mode: MasterMode, +) -> Vec { + let num_limbs = num_cols.div_ceil(32).max(1); + let max_block_rows = (gpu_block_bytes() / (num_limbs * 4)).max(1); + // Products arrive row-major (the extract loops emit them per input row, in order; the hot/cold + // filter above preserves that order), so each block is a contiguous product slice. Rows are + // independent — every product writes only its own row — so concatenating block outputs + // reproduces the single-launch result exactly. + debug_assert!(products.windows(2).all(|w| w[0].row <= w[1].row)); + // Per-product `(matrix, term)` pair counts, i.e. kernel threads. The kernel indexes threads + // by `ABSOLUTE_POS`, a `u32`, so a block must also stay under `2^32` pairs — output bytes + // alone don't bound this (pairs per row grow with the degree; an unbounded all-rows build + // reaches ~4.4e9 pairs by stem ~145). For `Resident` this pre-pass also warms the shared + // resident master, so every block's layout lookups below are read-lock cache hits; for + // `Transient` it warms the host-side [`COLD_COUNT`] shape cache the same way (no per-block recount). + // This pre-pass is where multi-minute stalls hide: it is strictly sequential, it calls + // `admissible_matrices` for every first-sight `R`, and it serialises on the `RESIDENT_HOST` + // write lock while appending multi-GB pending buffers — all of it previously outside every + // span and every timer, so a worker parked here logged nothing at all. `new_r` distinguishes + // "paid to warm the master" from "waited for someone else's warm-up". + let prepass = tracing::debug_span!( + "pair_prepass", + products = products.len(), + new_r = tracing::field::Empty, + ); + let misses_before = RESIDENT_MISSES.load(std::sync::atomic::Ordering::Relaxed); + let prod_pairs: Vec = prepass.in_scope(|| { + // One-entry memo on `(r_degree, r_idx)`. `extract_restricted` emits one product per target + // generator block of an input row, and `R` is a property of the ROW, so consecutive products + // repeat the same `R` — a single slot catches all of it without a hash or an allocation. + // + // Worth doing because the lookup is not cheap: `resident_info` takes a read lock on the + // process-wide `RESIDENT_HOST` and SipHashes a `PPart`, and a call-graph profile of an + // uncapped stem-150 run put `resident_info` at 12.93% of user cycles (4.44% of it inside + // `RwLock::read`, plus 1.26% in `read_contended` and 3.37% in `DefaultHasher`/`RandomState`). + // Per DISTINCT `R` that cost is unavoidable; per PRODUCT it is pure repetition. + let mut memo: Option<((i32, usize), usize)> = None; + products + .iter() + .map(|prod| { + let key = (prod.r_degree, prod.r_idx); + let num_mats = match memo { + Some((k, v)) if k == key => v, + _ => { + let r = algebra.basis_element_from_index(prod.r_degree, prod.r_idx); + let v = match mode { + MasterMode::Resident => { + resident_info(algebra, r.p_part).num_mats as usize + } + MasterMode::Transient => cold_count(algebra, r.p_part).2 as usize, + }; + memo = Some((key, v)); + v + } + }; + // Threads, not pairs: one per (MATRIX_GROUP x TERM_GROUP tile). + num_mats.div_ceil(MATRIX_GROUP) * prod.term_indices.len().div_ceil(TERM_GROUP) + }) + .collect() + }); + prepass.record( + "new_r", + RESIDENT_MISSES.load(std::sync::atomic::Ordering::Relaxed) - misses_before, + ); + drop(prepass); + let mut result: Vec = Vec::new(); + let (mut r0, mut p0) = (0, 0); + while r0 < num_rows { + // Grow the block row by row until the next row would break either budget — output bytes + // ([`gpu_block_bytes`]) or kernel threads ([`GPU_PAIR_CHUNK`]) — always taking at least + // one row (a lone over-budget row still fits the kernel's `u32` limit, asserted in the + // block). + let (mut r1, mut p1) = (r0, p0); + let mut pairs = 0usize; + while r1 < num_rows && r1 - r0 < max_block_rows { + let q = p1 + products[p1..].partition_point(|p| p.row <= r1); + let row_pairs: usize = prod_pairs[p1..q].iter().sum(); + if r1 > r0 && pairs + row_pairs > GPU_PAIR_CHUNK { + break; + } + pairs += row_pairs; + (r1, p1) = (r1 + 1, q); + } + // Fan the row-block out across devices. The master is sharded, so a product can only run + // where its `R` lives; each device evaluates its own subset over the SAME rows and the + // partial outputs are XORed. That is exact, not an approximation: every product contributes + // by `fetch_xor` into the output limbs, so the contributions commute and split freely. + let block = &products[p0..p1]; + let mut by_dev: Vec> = vec![Vec::new(); gpu_count()]; + // Same one-entry memo as the pre-pass, for the same reason: `R` is a property of the row, + // so consecutive products repeat it and each repeat would otherwise cost a global read lock + // and a `PPart` hash. + let mut dev_memo: Option<((i32, usize), usize)> = None; + for (pi, prod) in block.iter().enumerate() { + let d = match mode { + // Transient blocks enumerate their own master into per-launch scratch, so they are + // device-agnostic; spread them round-robin instead of piling onto device 0. + // + // TRIED AND REVERTED: hashing `R` to a device here (`shard_of`, as `Resident` does) + // so each distinct `R` enumerates exactly ONCE per block instead of being scattered + // across up to `gpu_count()` devices that each enumerate it. The mechanism works — + // `Rs/launch` fell 1293 -> 832, i.e. 36% of all `R`-enumerations removed (13.3M -> + // 8.5M) — and it bought NOTHING: 372.0 s against 373.0 s over three interleaved + // rounds, and only 2.6% off peak GPU memory (129.4 -> 126.0 GB). + // + // That is the max-versus-sum lesson, and it generalises: an ncu profile at + // PRODUCTION geometry shows a launch is 3-104 blocks of 32 threads taking 7-100 ms + // at 1.56% achieved occupancy, i.e. its duration is set by its LONGEST single `R` + // chain, not by how many `R`s it enumerates. Deduplicating work off the non-critical + // chains cannot move a maximum. Only shortening the longest chain can — by + // splitting it (needs an unrank to jump to the k-th admissible matrix, since the + // odometer derives each from the previous) or by keeping the longest `R`s resident + // so they are never re-enumerated at all. + MasterMode::Transient => pi % gpu_count(), + MasterMode::Resident => { + let key = (prod.r_degree, prod.r_idx); + match dev_memo { + Some((k, v)) if k == key => v, + _ => { + let r = algebra.basis_element_from_index(prod.r_degree, prod.r_idx); + let v = resident_info(algebra, r.p_part).dev as usize; + dev_memo = Some((key, v)); + v + } + } + } + }; + by_dev[d].push(prod.clone()); + } + // Marshal + submit every device's share first, THEN wait. `multiply_batch_block` returns + // the wait rather than performing it, so device `d + 1`'s marshalling overlaps device `d`'s + // execution and all shards are in flight together. The first cut used `std::thread::scope` + // here, which spawned an OS thread per device per block — hundreds a second, and thread ids + // into the hundreds of thousands in the logs. + // One thread, asynchronous submissions to every device. No fan-out threads and no pool: + // `submit_on` does not block, so this marshals a shard, hands it to its device, and moves to + // the next while that device is already running. Only the final wait blocks. + // + // Sharding SPLITS the products, so marshalling the shards one after another is the same + // total host work as marshalling one unsharded block — there is nothing to parallelise here. + // And with ~32 resolution workers each submitting `gpu_count()` jobs, ~128 launches are in + // flight, which is the queue depth the devices need; a fan-out pool could only cap that + // (sized at 64 it halved depth 3.8 -> 2.2 and cost 47% end to end). + // Three phases, and the split points are load-bearing rather than stylistic. + // + // 1. MARSHAL every shard. Parallel inside (rayon), and no permit is held by this thread, so + // a stolen resolution-step job that parks on `GpuPermit::acquire` cannot wedge the join + // it needs to finish — [`GpuBudget`]'s invariant, which the previous fused shape broke + // by holding device `d`'s permit while marshalling device `d + 1`. + // 2. SUBMIT every shard. Sequential and rayon-free: takes each permit and hands the block + // to its device's worker without blocking, so all `gpu_count()` shards are executing + // together rather than one after another. + // 3. WAIT on each. `multiply_batch_block`'s readback is issued but not awaited on the + // worker, so the devices stay busy with later blocks while this thread sits on fences. + // + // In-flight work is therefore bounded by `NASSAU_GPU_MEM_BUDGET_MB` — memory, not threads. + // Each shard runs its WHOLE pipeline — marshal, permit, submit, wait — concurrently, so + // all `gpu_count()` devices receive work at once. Verified stem-200s: + // + // concurrent submission (this) 2412 s marshal 1108 s depth 2.4 + // scoped marshal, serial submit 3358 s marshal 2439 s depth 1.8 + // serial marshal, serial submit 3138 s marshal 3283 s depth 1.9 + // rayon marshal, serial submit 3373 s + // + // Row 2 pins the mechanism: it cut marshal by 840 s and came out 220 s SLOWER. Marshal time + // does not predict wall time — leaving each submit behind the previous shard's marshal + // starves the devices however fast the marshal itself is. + // + // Helpers are persistent and PRIVATE to this thread (see [`shard_pool`]); the caller takes + // the last shard itself, so `gpu_count() - 1` are dispatched. Do not consolidate them into a + // shared pool or a rayon `par_iter` without an end-to-end re-measure: both were tried + // (~3743 s and 3373 s) and both lost, because a shared queue puts a shard behind other + // callers' shards — precisely the serialisation this exists to remove. + // + // Permits are acquired from several threads at once. Safe in the default configuration + // because the marshal contains no rayon — the `term_gei` fill is deliberately sequential + // (a par_iter over it once measured a 146 s stall) — so there is no join for a + // permit-blocked steal to wedge, and every holder is on the GPU or progressing. Under the + // `NASSAU_GPU_BASIS_PASSTHROUGH` diagnostic that fill IS a par_iter, which is exactly + // [`GpuBudget`]'s documented deadlock, so that path runs the shards serially. + let mut shards: Vec<(usize, Vec)> = by_dev + .into_iter() + .enumerate() + .filter(|(_, ps)| !ps.is_empty()) + .collect(); + let block_rows = r1 - r0; + let run_shard = + move |algebra: Arc, d: usize, ps: Vec| -> Bytes { + multiply_batch_block(&algebra, num_cols, r0, block_rows, &ps, mode, d)()() + }; + let partials: Vec = if basis_passthrough() || shards.len() == 1 { + shards + .into_iter() + .map(|(d, ps)| run_shard(algebra.clone(), d, ps)) + .collect() + } else { + // Dispatch all but one, run that one here, then collect. Every helper is joined through + // its result channel before this returns, so a shard cannot outlive the block. + let mine = shards.pop().expect("shards is non-empty"); + let rxs: Vec<_> = shards + .into_iter() + .enumerate() + .map(|(i, (d, ps))| { + let (tx, rx) = std::sync::mpsc::sync_channel::(1); + let alg = algebra.clone(); + shard_pool::dispatch( + i, + Box::new(move || { + let _ = tx.send(run_shard(alg, d, ps)); + }), + ); + rx + }) + .collect(); + let here = run_shard(algebra.clone(), mine.0, mine.1); + let mut out: Vec = rxs + .into_iter() + .map(|rx| rx.recv().expect("a shard helper panicked")) + .collect(); + out.push(here); + out + }; + let mut it = partials.into_iter(); + let mut acc = it + .next() + .expect("a non-empty row block has at least one device's products"); + for part in it { + // Byte-wise XOR is identical to limb-wise here and needs no typed view. + for (x, y) in acc.iter_mut().zip(part.iter()) { + *x ^= *y; + } + } + result.push(acc); + (r0, p0) = (r1, p1); + } + result +} + +/// The wait for one launched block: blocks on its completion fence and yields its output rows. +type BlockWait = Box Bytes + Send>; + +/// A block that has been *marshalled* but not yet admitted to the device. Calling it takes the +/// [`GpuPermit`] and submits, returning the [`BlockWait`]. +/// +/// The two phases are separate so that no permit is ever held across a rayon parallel section — +/// [`GpuBudget`]'s safety invariant. Marshalling is parallel; permit acquisition and submission are +/// not. Fusing them (as this did until now) meant a caller fanning out over `gpu_count()` devices +/// held device `d`'s permit while marshalling device `d + 1`, so a par_iter chunk could steal a +/// resolution-step job that parked on `acquire` while this thread waited on a join those very +/// workers had to finish — the H200 stall the invariant exists to prevent. +type BlockSubmit<'a> = Box BlockWait + Send + 'a>; + +/// One bounded launch of [`multiply_batch_on_gpu`]: rows `row_base..row_base + num_rows` of the +/// full build, with `products` the (contiguous, row-major) slice landing in those rows. Marshals +/// on the calling thread (in parallel), then hands back the submit step; the [`GpuPermit`] taken +/// there is held until the readback completes, so the total output size of in-flight launches +/// stays under `NASSAU_GPU_MEM_BUDGET_MB` across all worker threads. +fn multiply_batch_block<'a>( + algebra: &'a MilnorAlgebra, + num_cols: usize, + row_base: usize, + num_rows: usize, + products: &'a [GpuProduct], + mode: MasterMode, + dev: usize, +) -> BlockSubmit<'a> { let (width, g) = algebra.seqno_table_u32(); let mut xi: Vec = xi_degrees(algebra.prime()) .iter() @@ -690,238 +3034,3075 @@ pub fn multiply_batch_on_gpu( }); prod_r_index.push(ri); } + // Breakdown of the `prep` figure, which was only ever a single number and so could not say + // whether marshalling is worth restructuring or merely worth accepting. `intern` is the + // per-launch `HashMap<(i32, usize), u32>` above — pure representation tax, since every `R` + // already has a resident device identity ([`RInfo`]) the caller could have carried instead of + // an algebra-local `(degree, index)` pair. If it dominates, deleting it needs no new data + // structure, only a different id in `GpuProduct`. + let intern_ms = t_marshal.elapsed().as_secs_f64() * 1e3; + let t_basis = std::time::Instant::now(); + + // Admissible-matrix data (`col_sums`/`masks` + per-`R` offsets) is resident (built in the + // thread-local `RESIDENT` store below), so nothing to enumerate or lay out here. + + // Per-product term p-parts (padded to `width`) and lengths, filled directly into two flat + // buffers rather than one `(Vec, Vec)` per product. At the frontier a launch has ~10^5-10^6 + // products; the old per-product `Vec` pair (plus the later concat-copy) was ~10^6 tiny + // allocations per launch — a dominant chunk of the marshal cost. `term_off` is the prefix sum + // of term counts, so each product owns a disjoint output range and the fill stays parallel. + let term_off: Vec = { + let mut off = Vec::with_capacity(products.len() + 1); + let mut acc = 0usize; + for prod in products { + off.push(acc); + acc += prod.term_indices.len(); + } + off.push(acc); + off + }; + let total_terms = *term_off.last().unwrap(); - // 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. + // Resident-basis path (the default): a term's p-part is not marshalled at all — it lives on + // the device (built once, grown incrementally). We upload only `term_gei[slot]`, the term's + // *global* basis-element index `global_base[s_degree] + ti`. Ensure the basis covers every + // `s_degree` in this block, then snapshot `global_base` so the parallel fill needs no lock. + let max_s_degree = products.iter().map(|p| p.s_degree).max().unwrap_or(0); + // `ensure_basis` takes the basis WRITE lock on a first-sight degree; spanned separately from + // the fill below so a wait on that lock is not misread as marshalling work. + let global_base = tracing::debug_span!("ensure_basis", max_s_degree) + .in_scope(|| ensure_basis(algebra, width, max_s_degree)); + // Everything from the intern to here: `term_off` prefix sums, the `max_s_degree` scan, and + // `ensure_basis` (which may upload). Separated from the fill because the two have different + // fixes — this shrinks by keeping the basis warm, the fill by not rebuilding indices. + let basis_ms = t_basis.elapsed().as_secs_f64() * 1e3; + let t_tgei = std::time::Instant::now(); + let mut term_gei: Vec = vec![0u32; total_terms]; + // The ONLY rayon construct inside the guarded region, hence the only place a worker can block + // at a join and enter the steal loop. The multi-minute stalls sit somewhere in this guarded + // region and every other part of it is now spanned and bounded (`extract_restricted` ≤4.8 s, + // `pair_prepass` ≤5.1 s, `gpu_submit` 40 ms), so this is what remains. `prep` was only ever + // reported as a sum, which cannot separate a few 100 s outliers from many small costs. + { + // Scoped to the fill alone: entering at function level would leave the span open across + // the permit wait and the GPU submission and attribute their time here. + let _marshal_span = tracing::debug_span!( + "marshal_terms", + products = products.len(), + terms = total_terms + ) + .entered(); + // SEQUENTIAL, deliberately. The body is one add and one store per term, so at the largest + // observed size (478 972 products / 2 300 621 terms) the whole fill is a few milliseconds + // of memory-bandwidth-bound work — rayon cannot speed that up past its own split/join + // overhead. + // + // What parallelising it DID buy was a join, and therefore rayon's steal loop, and therefore + // exposure to starvation: instrumenting this span measured a single fill at **146 s** (p99 + // 0.09 s — 99 % are fast, only the tail explodes), roughly 30 000x the work involved. That + // was the multi-hundred-second "signature step" stall, which had been misattributed in turn + // to GPU submission ordering, to the resident-master pre-pass, and to the kernel itself. + // + // With this sequential there is no join anywhere inside a bidegree's guarded region, so a + // worker cannot be parked here at all. + let tg_all = &mut term_gei[..]; + for (pi, prod) in products.iter().enumerate() { + let (off, nt) = (term_off[pi], prod.term_indices.len()); + let base = global_base[prod.s_degree as usize]; + for (slot, &ti) in tg_all[off..off + nt] + .iter_mut() + .zip(prod.term_indices.iter()) + { + *slot = base + ti as u32; + } + } + } + // The `term_gei` fill: one add and one store per term. This is the part that becomes a SLICE + // rather than a build if products are held struct-of-arrays with global basis indices, so its + // share is the ceiling on what that refactor can return. + let tgei_ms = t_tgei.elapsed().as_secs_f64() * 1e3; + // Device need: the largest `gei` any term dereferences is `< global_base[max_s_degree + 1]` + // (all elements through degree `max_s_degree`), so uploading that many covers the block. + let need_basis_elems = global_base[max_s_degree as usize + 1] as usize; - // 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| { + // A/B diagnostic (`NASSAU_GPU_BASIS_PASSTHROUGH=1`): bind the *per-launch* term buffers as the + // "basis" and set `term_gei` to the identity, so the new kernel reproduces the old behaviour + // bit-for-bit. If passthrough matches the CPU but the resident path does not, the bug is in + // the resident host/upload logic, not the kernel signature — and vice-versa. + let passthrough = basis_passthrough(); + let (mut term_pparts, mut term_lens): (Vec, Vec) = if passthrough { + let mut tp: Vec = vec![0u16; total_terms * width]; + let mut tl: Vec = vec![0u32; total_terms]; + let tp_base = tp.as_mut_ptr() as usize; + let tl_base = tl.as_mut_ptr() as usize; + (0..products.len()).into_maybe_par_iter().for_each(|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); + let (off, nt) = (term_off[pi], prod.term_indices.len()); + // SAFETY: disjoint per-product ranges, as above. + let tpp = unsafe { + std::slice::from_raw_parts_mut((tp_base as *mut u16).add(off * width), nt * width) + }; + let tll = unsafe { std::slice::from_raw_parts_mut((tl_base as *mut u32).add(off), 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) { + tll[k] = elt.p_part.len() as u32; + for (slot, v) in tpp[k * width..(k + 1) * width] + .iter_mut() + .zip(elt.p_part.iter()) + { *slot = narrow_u16(v); } } - (tp, tl) - }) - .collect(); + }); + // Identity indices, so the kernel's `gei*width` / `basis_lens[gei]` hit slot `term_slot`. + for (i, g) in term_gei.iter_mut().enumerate() { + *g = i as u32; + } + (tp, tl) + } else { + (Vec::new(), Vec::new()) + }; - // Resident admissible-matrix store: enumerate each new `R` once and reuse forever; - // the per-`R` offsets are global (into the master `col_sums`/`masks`). Taking the lock - // here also serializes the device section across rayon workers (see [`Resident`]). - let mut resident = RESIDENT.lock().unwrap(); - let mut r_cs_offset: Vec = Vec::with_capacity(distinct_r.len()); - let mut r_mk_offset: Vec = Vec::with_capacity(distinct_r.len()); - let mut r_cs_len: Vec = Vec::with_capacity(distinct_r.len()); - let mut r_mk_len: Vec = Vec::with_capacity(distinct_r.len()); - let mut r_num_matrices: Vec = Vec::with_capacity(distinct_r.len()); - for &(rd, ridx) in &distinct_r { - let r = algebra.basis_element_from_index(rd, ridx); - assert!(!r.p_part.is_empty(), "each R must be non-empty"); - let info = resident.ensure(algebra, &r.p_part); - r_cs_offset.push(info.cs_off); - r_mk_offset.push(info.mk_off); - r_cs_len.push(info.cs_len); - r_mk_len.push(info.mk_len); - r_num_matrices.push(info.num_mats as usize); - } - - // Lay out per-product term data + records + the pair-count prefix sum (sequential). - let mut term_pparts: Vec = Vec::new(); - let mut term_lens: Vec = Vec::new(); - let mut prod_term_start: Vec = Vec::with_capacity(products.len()); - let mut prod_num_terms: Vec = Vec::with_capacity(products.len()); - let mut prod_row_base: Vec = Vec::with_capacity(products.len()); - let mut prod_out_offset: Vec = Vec::with_capacity(products.len()); - // Per-product `(matrix, term)` pair count. A launch's total pair count is the sum, and can - // exceed `u32::MAX` at record degrees (a single all-rows reuse build reaches ~4.4e9 pairs at - // stem ~145). The kernel indexes threads by `ABSOLUTE_POS`, itself a `u32`, so a launch can - // address at most `2^32` threads; the device section below splits the products into chunks each - // bounded by [`GPU_PAIR_CHUNK`] so every kernel launch stays safely under that limit. Keeping the - // per-product counts (rather than a single prefix sum) lets each chunk build its own `u32` - // prefix sum locally. - let mut prod_pairs: Vec = Vec::with_capacity(products.len()); - 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); - let pairs = r_num_matrices[ri as usize] * prod.term_indices.len(); - prod_pairs.push(pairs); - pair_acc += pairs; - 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); - } - - let total_pairs = pair_acc; - let out_len = num_rows * num_limbs; - if std::env::var_os("NASSAU_GPU_DEBUG").is_some() { - let num_chunks = total_pairs.div_ceil(GPU_PAIR_CHUNK).max(1); - eprintln!( - "[gpu-batch] num_rows={num_rows} num_cols={num_cols} num_limbs={num_limbs} \ - products={} total_pairs={total_pairs} out_len={out_len} \ - chunks={num_chunks} (cap={GPU_PAIR_CHUNK})", - products.len(), + // Take the concurrency permit only now, with every rayon parallel section behind us: holding + // it across the `per_prod` par_iter above deadlocks, because that par_iter's chunks execute on + // *other* threads, which do not carry this thread's `ParallelGuard` flag and so can steal a + // bidegree job mid-chunk; the stolen job parks on `GpuPermit::acquire` while this thread's + // permit waits on the never-finishing join (observed on H200). Everything from here on is + // strictly sequential — the `ensure` calls below are cache hits (the caller's pair-count + // pre-pass already enumerated every `R`), and the device section never enters rayon — so + // every permit holder makes progress and stolen jobs waiting for a permit wake in finite + // time (priority inversion at worst, never deadlock). + // Held for the device section (RAII): bounds total in-flight output bytes across workers. It is + // MOVED INTO THE WAIT CLOSURE below and dropped only once the readback has completed, because + // that is when the output buffer (device page + pinned host landing) is actually free. A plain + // local here drops at this function's return — i.e. straight after `submit_on`, which does not + // block — so between `1639982877` (when this function started returning its wait instead of + // performing it) and now, the budget admitted every launch immediately and bounded nothing. + // With the readback no longer serialising the worker, this permit is the ONLY thing bounding + // in-flight memory, so its scope is load-bearing rather than belt-and-braces. + // Split the "marshal" figure at the point where this thread stops doing CPU work and starts + // waiting. `t_marshal` spans both, so the 80/20 marshal-vs-device headline it produced cannot + // distinguish host marshalling from time parked on our own permit / arbitration lock — and the + // two call for opposite fixes. + let prep_ms = t_marshal.elapsed().as_secs_f64() * 1e3; + + // Everything below is the submit phase: strictly sequential (no rayon), so the permit it takes + // satisfies [`GpuBudget`]'s invariant. The caller runs it only once every shard has marshalled. + Box::new(move || { + let t_wait = std::time::Instant::now(); + let permit = GpuPermit::acquire(num_rows * num_limbs * 4); + let permit_ms = t_wait.elapsed().as_secs_f64() * 1e3; + let t_lock = std::time::Instant::now(); + // Shared side of the cross-runtime GPU arbitration, taken here for the same reason as the + // permit above and never earlier: multiplies overlap each other freely but yield while an + // `fp-cuda` row reduction holds the device, so the reduction's thousands of tiny sequential + // relaunches are not stuck behind these saturating kernels (~10 000× when they are — see + // [`fp::gpu_lock`]). Taking it at function entry deadlocks exactly as described above: the + // marshalling `par_iter` runs chunks on other workers, which steal another bidegree's + // multiply, block acquiring the shared side behind a waiting reduction, and never let this + // thread's join finish (observed on H200). + // The arbitration's shared side is now taken by the GPU thread itself, around the device + // section it owns (see [`gpu_thread`]). Taking it here instead put ~100 workers through a + // writer-preferring lock to reach a stage only one of them could occupy anyway — measured at + // 10% of multiply time, pure convoy. With one submitter it is a 1-vs-1 handshake against the + // `fp-cuda` reduction, which is all the arbitration ever needed to be. + let lock_ms = t_lock.elapsed().as_secs_f64() * 1e3; + let wait_ms = t_wait.elapsed().as_secs_f64() * 1e3; + // Per-`R` offsets into the shared resident master (see [`ResidentHost`]). All read-lock + // cache hits: the caller's pair-count pre-pass already enumerated every `R` in this block. + // `need_cs`/`need_mk` track the furthest master offset this block dereferences, so the + // device section can skip the (multi-GB, mutex-serialized) master re-upload whenever the + // already-uploaded prefix covers it. + 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()); + let mut need_cs: usize = 0; + let mut need_mk: usize = 0; + // `Transient`: per-cold-`R` inputs for the on-device enumeration ([`enumerate_admissible_kernel`]). + // Instead of building this block's `col_sums`/`masks` on the host and uploading them (the H2D + // cost the eviction bench exposed), we upload only each cold `R`'s p-part + dimensions and + // generate the arrays into device scratch at the block-local `r_cs_offset`/`r_mk_offset`. Empty + // under `Resident`. + let mut enum_pp_rows: Vec> = Vec::new(); + let mut enum_rows: Vec = Vec::new(); + let mut enum_cols: Vec = Vec::new(); + 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"); + match mode { + MasterMode::Resident => { + let info = resident_info(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); + need_cs = need_cs + .max(info.cs_off as usize + info.num_mats as usize * info.cs_len as usize); + need_mk = need_mk + .max(info.mk_off as usize + info.num_mats as usize * info.mk_len as usize); + } + MasterMode::Transient => { + let (cs_len, mk_len, num_mats) = cold_count(algebra, r.p_part); + // Offsets are assigned after this loop: they depend on the matrix-count sort AND + // on segment packing, neither of which is known per-`R` in basis order. + r_cs_offset.push(0); + r_mk_offset.push(0); + r_cs_len.push(cs_len); + r_mk_len.push(mk_len); + r_num_matrices.push(num_mats as usize); + // `cols` = max bit-length of any entry, exactly as the enumeration kernel derives it; + // `cs_len == cols-1`, `mk_len == rows+cols-1` (asserted equal to `cold_count`'s below). + let cols = r + .p_part + .iter() + .map(|x| u32::BITS - x.leading_zeros()) + .max() + .unwrap(); + debug_assert_eq!( + (cs_len, mk_len), + (cols - 1, r.p_part.len() as u32 + cols - 1) + ); + enum_rows.push(r.p_part.len() as u32); + enum_cols.push(cols); + enum_pp_rows.push(r.p_part.iter().collect::>()); + } + } + } + + // (Transient) Flatten the cold p-parts (padded to the widest) for the enumeration kernel, in + // MATRIX-COUNT ORDER rather than basis order. + // + // The enumeration is an odometer, so a thread's cost is proportional to its `R`'s matrix + // count, and a warp retires only when its slowest lane does — a warp costs `32 x max`, not + // `sum`. Matrix counts are heavily skewed (the `NASSAU_R_STATS` Lorenz curve: the hottest 1% + // of `R`s carry 31% of all references), so in basis order one huge `R` idles the other 31 + // lanes for its entire run. Measured by `enum_warp_utilisation` to degree 130: 30.9% of lane + // slots do useful work in basis order against 66.8% sorted, a 2.16x headroom. + // + // Only the ENUMERATION's own inputs are permuted. The kernel writes each `R` to the absolute + // scratch offset it is handed, so reordering its inputs consistently reproduces byte-identical + // scratch — the multiply's `r_cs_offset`/`r_mk_offset`/`prod_r_index` keep basis order and are + // untouched. (`out_counts` follows the permutation, but production discards it.) + // + // The same pass lays the scratch out across as many SEGMENTS as it needs, instead of cramming + // it into one. A segment is `master_seg_elems()` (2^31 u16 = 4 GiB, deliberately under + // `u32::MAX` so a segment length cannot overflow cubecl's 32-bit array-length metadata — + // raising THAT is what reopens the truncation bug class). Binding transient scratch as + // segment 0 alone therefore capped a whole block at 4 GiB, and a stem-200 block wants 2.17e9 + // u16 of `masks`: over the line, hard assert, run dead. The resident path already spans + // `MASTER_MAX_SEG` segments through the same `seg_read`, so the ceiling here is 16x higher for + // free. + // + // Each `R` is placed WHOLLY inside one segment, with its `col_sums` and `masks` in the + // same-numbered segment, by advancing both cursors to the next boundary together whenever + // either run would straddle. That is what lets the enumeration run as one launch per segment + // against the existing single-buffer kernel signature, rather than needing a segmented-write + // kernel. The padding costs address space in an allocation already rounded to segments. + let seg_elems_layout = master_seg_elems(); + let mut enum_seg_ranges: Vec<(usize, usize)> = Vec::new(); + let (enum_pp, enum_width, enum_rows, enum_cols, enum_cs_out, enum_mk_out) = if mode + == MasterMode::Transient + { + let w = enum_rows.iter().copied().max().unwrap_or(1) as usize; + let mut order: Vec = (0..enum_pp_rows.len()).collect(); + order.sort_unstable_by_key(|&i| r_num_matrices[i]); + + let (mut cs_out, mut mk_out) = (vec![0u64; order.len()], vec![0u64; order.len()]); + let (mut seg, mut seg_start) = (0usize, 0usize); + for (slot, &i) in order.iter().enumerate() { + let cs_span = r_num_matrices[i] * r_cs_len[i] as usize; + let mk_span = r_num_matrices[i] * r_mk_len[i] as usize; + assert!( + cs_span <= seg_elems_layout && mk_span <= seg_elems_layout, + "one R's transient scratch ({cs_span}/{mk_span} u16) exceeds a whole segment \ + ({seg_elems_layout}); raise NASSAU_GPU_MASTER_SEG_ELEMS" + ); + let base = seg * seg_elems_layout; + if need_cs - base + cs_span > seg_elems_layout + || need_mk - base + mk_span > seg_elems_layout + { + enum_seg_ranges.push((seg_start, slot)); + seg += 1; + seg_start = slot; + need_cs = seg * seg_elems_layout; + need_mk = seg * seg_elems_layout; + } + cs_out[slot] = need_cs as u64; + mk_out[slot] = need_mk as u64; + r_cs_offset[i] = need_cs as u64; + r_mk_offset[i] = need_mk as u64; + need_cs += cs_span; + need_mk += mk_span; + } + enum_seg_ranges.push((seg_start, order.len())); + assert!( + seg < MASTER_MAX_SEG, + "transient scratch needs {} segments (> MASTER_MAX_SEG={MASTER_MAX_SEG}); raise \ + MASTER_MAX_SEG or NASSAU_GPU_MASTER_SEG_ELEMS", + seg + 1 + ); + + let mut pp = vec![0u32; enum_pp_rows.len() * w]; + for (slot, &i) in order.iter().enumerate() { + let row = &enum_pp_rows[i]; + for (dst, &v) in pp[slot * w..slot * w + row.len()].iter_mut().zip(row) { + *dst = v; + } + } + let pick_u32 = |src: &[u32]| -> Vec { order.iter().map(|&i| src[i]).collect() }; + ( + pp, + w, + pick_u32(&enum_rows), + pick_u32(&enum_cols), + cs_out, + mk_out, + ) + } else { + ( + Vec::new(), + 1usize, + Vec::new(), + Vec::new(), + Vec::new(), + Vec::new(), + ) + }; + + // Lay out per-product records + the pair-count prefix sum (sequential). Term data is already + // in `term_pparts`/`term_lens` (filled in parallel above); `term_off` gives each product's + // start, so nothing is copied here. + 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()); + // The pair prefix sum: entry `pi` is the number of `(matrix, term)` pairs before product + // `pi`, with the sentinel total at the end — the kernel binary-searches it to decode its + // thread index. The caller splits blocks near [`GPU_PAIR_CHUNK`], so every entry fits + // `u32` (a lone over-budget row can exceed the target but stays far below the kernel's + // `2^32` `ABSOLUTE_POS` limit; asserted below before the values are used). + let mut pps: Vec = Vec::with_capacity(products.len() + 1); + let mut pair_acc: usize = 0; + let mut real_pairs: usize = 0; + for (pi, prod) in products.iter().enumerate() { + let ri = prod_r_index[pi]; + prod_term_start.push(term_off[pi] as u32); + pps.push(pair_acc as u32); + // One thread per (matrix, TERM_GROUP-sized term group), not per (matrix, term). `pair_acc` + // sizes the grid, so it counts THREADS; `real_pairs` stays the count of `(matrix, term)` + // products actually evaluated, which is what the throughput stat must report. + prod_num_terms.push(prod.term_indices.len() as u32); + pair_acc += r_num_matrices[ri as usize].div_ceil(MATRIX_GROUP) + * prod.term_indices.len().div_ceil(TERM_GROUP); + real_pairs += r_num_matrices[ri as usize] * prod.term_indices.len(); + prod_row_base.push(((prod.row - row_base) * num_limbs) as u32); + prod_out_offset.push(prod.out_offset as u32); + } + + let total_pairs = pair_acc; + assert!( + u32::try_from(total_pairs).is_ok(), + "block pair count {total_pairs} exceeds the kernel's u32 thread limit" ); - } - if total_pairs == 0 { - return vec![vec![0u32; num_limbs]; num_rows]; - } + pps.push(total_pairs as u32); - // 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); - } + // Coarse index over the pair space: `coarse[i]` is the product owning pair `i << COARSE_LOG`, + // so the product for a thread at pair `k` lies in `coarse[ci] ..= coarse[ci + 1]` for + // `ci = k >> COARSE_LOG`. Ablation put the unaided binary search at ~12% of kernel time, and it + // is the worst kind of work: `ceil(log2(num_products))` *dependent* global loads, each a full + // latency stall, before a thread can touch any of its own data. + let mut coarse: Vec = Vec::with_capacity((total_pairs >> COARSE_LOG) + 2); + { + let mut pi = 0usize; + let mut k = 0usize; + while k <= total_pairs { + while pi + 1 < products.len() && (pps[pi + 1] as usize) <= k { + pi += 1; + } + coarse.push(pi as u32); + k += 1 << COARSE_LOG; + } + // Sentinel: `ci + 1` must be readable for threads in the final chunk. + coarse.push(products.len().saturating_sub(1) as u32); + } + // Widest product span any chunk covers, so the in-kernel scan has a static iteration bound. + let coarse_span = coarse + .windows(2) + .map(|w| (w[1] - w[0]) as usize) + .max() + .unwrap_or(0); + + let out_len = num_rows * num_limbs; + // Output offsets (`prod_out_offset`/`prod_row_base`) are `u32` values indexing `out_h`; the + // row-block splitter caps `out_len` well under `u32::MAX` (its output-byte budget is far below + // 16 GiB), so these never truncate. Assert it loudly rather than silently corrupt if a future + // budget is set absurdly high. (The `out_h` *length* itself is bound with dynamic addressing.) + assert!( + u32::try_from(out_len).is_ok(), + "block output length {out_len} exceeds u32; lower NASSAU_GPU_BLOCK_MB / row-block \ + budget" + ); + if std::env::var_os("NASSAU_GPU_DEBUG").is_some() { + let kb = |n: usize, sz: usize| n * sz / 1024; + eprintln!( + "[gpu-batch] rows={num_rows} cols={num_cols} products={} \ + total_pairs={total_pairs} out_len={out_len} | UPLOAD-KB: g={} xi={} term_gei={} \ + prod_arrays={} pps={} | resident cs={} mk={} basis_elems={need_basis_elems}", + products.len(), + kb(g.len(), 4), + kb(xi.len(), 4), + kb(term_gei.len(), 4), + kb(products.len() * 5, 4), + kb(pps.len(), 4), + kb(need_cs, 2), + kb(need_mk, 2), + ); + } + if total_pairs == 0 { + // Nothing to launch: hand back a wait that yields the zero block, so the caller's + // submit-then-wait shape is uniform. + let empty = Bytes::from_elems(vec![0u32; num_rows * num_limbs]); + return Box::new(move || empty) as BlockWait; + } + + // The resident `col_sums`/`masks` and basis are non-empty once any `R`/term is present + // (guaranteed here, since `total_pairs > 0`); only `term_gei` (and, in passthrough, the + // per-launch term buffers) needs the non-empty guard `create_from_slice` requires. + if term_gei.is_empty() { + term_gei.push(0); + } + if passthrough && term_pparts.is_empty() { + term_pparts.push(0); + term_lens.push(0); + } + + let term_gei_len = term_gei.len(); + let pps_len = pps.len(); + let marshal_ms = t_marshal.elapsed().as_secs_f64() * 1e3; + + let t_device = std::time::Instant::now(); + + // `products` is borrowed; the device section only needs its length, and everything else it + // touches is owned, so hoisting this makes the closure `'static` and thus sendable. + let num_products = products.len(); + + // Hand the whole device section to the single GPU thread (see [`gpu_thread`]) and block for the + // result. FIFO service order bounds this wait by the work already queued, replacing the + // unbounded starvation that the shared-stream free-for-all allowed (370 s observed). + // + // The `gpu_submit` span makes that wait *visible*: a worker stuck here previously logged + // nothing at all for the whole stall, which is why the multi-minute steps looked like compute. + // `dev` is the point of this field: the span is entered on the SUBMITTING (rayon) thread, not + // inside the `nassau-gpu` worker, so neither the thread id nor its name says which device a + // job went to. Without it there is no way to check shard balance from a log. + let submit_span = tracing::debug_span!( + "gpu_submit", + dev = dev, + rows = num_rows, + pairs = total_pairs, + out = out_len + ); + // Submit and return the wait: the caller launches every device's share before blocking on any + // of them, so the shards actually overlap. + let pending = submit_span.in_scope(|| { + gpu_thread::submit_on(dev, move || { + // Arbitrate against the `fp-cuda` row reduction from the one thread that submits (see the + // note where the permit is taken). Dropped at the end of this task. + let _shared = fp::gpu_lock::shared(); + let client = gpu_client(); + // Bind the segmented resident master/basis (see [`SegBuf`], [`seg_grow`]). Each store is + // `MASTER_MAX_SEG` segment handles padded with a never-indexed 1-element dummy; a + // single-buffer store (transient enum scratch or the passthrough diagnostic) is bound as + // segment 0, which the kernel resolves correctly because `seg_elems` exceeds its length so + // every offset lands in segment 0. `seg_grow!` re-uploads only the tail past the resident + // prefix (`need_*`), never copying existing segments — the no-`~2×`-spike growth that keeps + // cubecl out of its memory-corruption regime. + let seg_elems = master_seg_elems(); + let dummy16 = client.create_from_slice(u16::as_bytes(&[0u16])); + let dummy32 = client.create_from_slice(u32::as_bytes(&[0u32])); + let pad_u16 = |mut v: Vec<(Handle, usize)>| -> Vec<(Handle, usize)> { + assert!( + v.len() <= MASTER_MAX_SEG, + "segment count exceeds MASTER_MAX_SEG" + ); + while v.len() < MASTER_MAX_SEG { + v.push((dummy16.clone(), 1)); + } + v + }; + let pad_u32 = |mut v: Vec<(Handle, usize)>| -> Vec<(Handle, usize)> { + assert!( + v.len() <= MASTER_MAX_SEG, + "segment count exceeds MASTER_MAX_SEG" + ); + while v.len() < MASTER_MAX_SEG { + v.push((dummy32.clone(), 1)); + } + v + }; + let full = |segs: Vec| -> Vec<(Handle, usize)> { + segs.into_iter().map(|h| (h, seg_elems)).collect() + }; - let marshal_ms = t_marshal.elapsed().as_secs_f64() * 1e3; + // `Transient` (degree > cap `R`s): enumerate this block's cold master ON the device into + // scratch, freed with the launch. `Resident` (default): grow + reuse the shared master. + // The `Transient` enumeration launch's own inputs need the same lifetime extension as + // the multiply's: they are consumed by `from_raw_parts` and would otherwise die with + // the match arm, while the enum kernel is still running on this stream. + // (`cs_scratch`/`mk_scratch` survive into `cs_seg`/`mk_seg`, so they are already + // covered by the main keepalive below.) + let mut enum_keep: Vec = Vec::new(); + let (cs_seg, mk_seg) = match mode { + MasterMode::Resident => { + let (cs_segs, _) = seg_grow!( + client, + resident_dev(), + cs, + resident_upload(), + need_cs, + copy_into_u16, + u16::as_bytes, + u16, + |_up: usize| { + let mut h = RESIDENT_HOST.write().unwrap(); + let dev = cur_device(); + let nl = h.cs_len[dev]; + (std::mem::take(&mut h.cs_pending[dev]), nl) + } + ); + let (mk_segs, _) = seg_grow!( + client, + resident_dev(), + mk, + resident_upload(), + need_mk, + copy_into_u16, + u16::as_bytes, + u16, + |_up: usize| { + let mut h = RESIDENT_HOST.write().unwrap(); + let dev = cur_device(); + let nl = h.mk_len[dev]; + (std::mem::take(&mut h.mk_pending[dev]), nl) + } + ); + (pad_u16(full(cs_segs)), pad_u16(full(mk_segs))) + } + MasterMode::Transient => { + // The enumeration launch is issued before the multiply on this same stream, so the + // scratch is fully written when the multiply reads it (one-stream launches are + // ordered, as with `zero_u32` below). + // Must match the kernel's shared-memory stride exactly. + const ENUM_THREADS: u32 = ENUM_BLOCK; + // One allocation per segment, and one enumeration launch per segment over the + // `R`s the layout pass placed there. Every `R` sits wholly inside its segment + // with `col_sums` and `masks` in the same-numbered one, so each launch keeps + // the kernel's plain single-buffer signature and just works in segment-local + // offsets. The last segment is allocated to what it actually holds; the rest + // are full. + let nseg = enum_seg_ranges.len(); + let mut cs_segs: Vec<(Handle, usize)> = Vec::with_capacity(nseg); + let mut mk_segs: Vec<(Handle, usize)> = Vec::with_capacity(nseg); + for s in 0..nseg { + let base = s * seg_elems; + let cs_len_s = (need_cs - base).min(seg_elems).max(1); + let mk_len_s = (need_mk - base).min(seg_elems).max(1); + cs_segs.push((client.empty(cs_len_s * size_of::()), cs_len_s)); + mk_segs.push((client.empty(mk_len_s * size_of::()), mk_len_s)); + } + for (s, &(lo, hi)) in enum_seg_ranges.iter().enumerate() { + let n_s = hi - lo; + if n_s == 0 { + continue; + } + let base = (s * seg_elems) as u64; + // Segment-local offsets: the kernel indexes this segment's buffer alone. + let cs_loc: Vec = + enum_cs_out[lo..hi].iter().map(|&o| o - base).collect(); + let mk_loc: Vec = + enum_mk_out[lo..hi].iter().map(|&o| o - base).collect(); + let pp_s = &enum_pp[lo * enum_width..hi * enum_width]; + let cnt_scratch = client.empty(n_s * size_of::()); + let epp_h = client.create_from_slice(u32::as_bytes(pp_s)); + let er_h = client.create_from_slice(u32::as_bytes(&enum_rows[lo..hi])); + let ec_h = client.create_from_slice(u32::as_bytes(&enum_cols[lo..hi])); + let eco_h = client.create_from_slice(u64::as_bytes(&cs_loc)); + let emo_h = client.create_from_slice(u64::as_bytes(&mk_loc)); + enum_keep.extend([ + cnt_scratch.clone(), + epp_h.clone(), + er_h.clone(), + ec_h.clone(), + eco_h.clone(), + emo_h.clone(), + ]); + let enum_blocks = (n_s as u32).div_ceil(ENUM_THREADS).max(1); + ENUM_LAUNCHES.fetch_add(1, Ordering::Relaxed); + ENUM_RS.fetch_add(n_s as u64, Ordering::Relaxed); + ENUM_RS_MAX.fetch_max(n_s as u64, Ordering::Relaxed); + ENUM_BLOCKS.fetch_add(enum_blocks as u64, Ordering::Relaxed); + unsafe { + enumerate_admissible_kernel::launch_unchecked::( + &client, + CubeCount::Static(enum_blocks, 1, 1), + CubeDim::new_1d(ENUM_THREADS), + BufferArg::from_raw_parts(epp_h, pp_s.len()), + BufferArg::from_raw_parts(er_h, n_s), + BufferArg::from_raw_parts(ec_h, n_s), + BufferArg::from_raw_parts(eco_h, n_s), + BufferArg::from_raw_parts(emo_h, n_s), + BufferArg::from_raw_parts(cs_segs[s].0.clone(), cs_segs[s].1), + BufferArg::from_raw_parts(mk_segs[s].0.clone(), mk_segs[s].1), + BufferArg::from_raw_parts(cnt_scratch, n_s), + enum_width, + n_s, + 0u32, + 1, + ); + } + } + (pad_u16(cs_segs), pad_u16(mk_segs)) + } + }; + // Resident basis segments (default) or per-launch passthrough buffers (A/B diagnostic) bound + // as segment 0. Every `gei` a thread dereferences is `< need_basis_elems`, so growing the + // basis to `need_basis_elems` (pp: `× width`) covers it. + let (pp_seg, ln_seg) = if passthrough { + assert!( + term_pparts.len() <= seg_elems && term_lens.len() <= seg_elems, + "passthrough basis exceeds one segment; raise NASSAU_GPU_MASTER_SEG_ELEMS" + ); + let bp = client.create_from_slice(u16::as_bytes(&term_pparts)); + let bl = client.create_from_slice(u32::as_bytes(&term_lens)); + ( + pad_u16(vec![(bp, term_pparts.len())]), + pad_u32(vec![(bl, term_lens.len())]), + ) + } else { + let (pp_segs, _) = seg_grow!( + client, + resident_basis_dev(), + pp, + resident_basis_upload(), + need_basis_elems * width, + copy_into_u16, + u16::as_bytes, + u16, + |up: usize| { + let h = RESIDENT_BASIS_HOST.read().unwrap(); + let nl = h.lens.len() * h.width; + (h.pparts[up..nl].to_vec(), nl) + } + ); + let (ln_segs, _) = seg_grow!( + client, + resident_basis_dev(), + ln, + resident_basis_upload(), + need_basis_elems, + copy_into_u32, + u32::as_bytes, + u32, + |up: usize| { + let h = RESIDENT_BASIS_HOST.read().unwrap(); + let nl = h.lens.len(); + (h.lens[up..nl].to_vec(), nl) + } + ); + (pad_u16(full(pp_segs)), pad_u32(full(ln_segs))) + }; + // Upload the block's data — term data, seqno/xi tables, per-`R` offsets, per-product + // records, the pair prefix sum, and the (zeroed) output buffer — and launch once: the + // caller has already bounded this block's pair count and output size. + // Hand the marshalled buffers over (`create`) rather than have cubecl copy out of a + // borrowed slice (`create_from_slice`): the marshal already built exactly the bytes the + // upload wants, so the extra staging copy is pure waste. Mirrors what [`BatchOutput`] + // does on the way back. NOT using `client.staging()` to pin these: it consumes the + // `Bytes` by value (so a buffer cannot be pinned once and reused across launches) and + // its own docs note it blocks the compute queue. + let tg_h = client.create(Bytes::from_elems(term_gei)); + // `g`/`xi` are identical every launch at this degree: fetch the shared resident copies + // (uploaded once, re-uploaded only on a degree bump) instead of re-uploading them here. + let (g_h, xi_h) = resident_seqno!(client, g, xi); + let rco_h = client.create_from_slice(u64::as_bytes(&r_cs_offset)); + let rmo_h = client.create_from_slice(u64::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)); + // Per-`R` matrix count, so the kernel reads it instead of dividing for it. Integer + // division by a runtime value is emulated on the GPU (I2F/MUFU.RCP/F2I plus fixups, + // ~20 instructions), and this kernel is issue-limited on integer work. + let r_num_mats_u32: Vec = r_num_matrices.iter().map(|&n| n as u32).collect(); + let rnm_h = client.create_from_slice(u32::as_bytes(&r_num_mats_u32)); + const THREADS: u32 = 256; + // No realloc barrier needed: the resident master/basis are append-only segmented stores whose + // segments, once allocated and written, never change identity and are never freed (see + // [`seg_grow`]). This block cloned their segment handles above, so each stays alive (refcount + // > 0) for the whole kernel even if another thread grows the store concurrently by appending + // a new segment — the churny whole-buffer swap that needed quiescing is gone. + // Allocate the XOR accumulator uninitialized and zero it on-device (see [`zero_u32`]), + // instead of uploading a hundreds-of-MB host zero buffer — the former dominant serial + // marshaling cost. Bounded by the caller's row-batching (see `get_partial_matrix`), so it + // stays small and is returned to the pool by `memory_cleanup` below. Same stream as the + // multiply, so the zero is ordered before it. + let out_h = client.empty(out_len * size_of::()); + unsafe { + zero_u32::launch::( + &client, + CubeCount::Static((out_len as u32).div_ceil(THREADS).max(1), 1, 1), + CubeDim::new_1d(THREADS), + BufferArg::from_raw_parts(out_h.clone(), out_len), + ); + } + + let pri_h = client.create(Bytes::from_elems(prod_r_index)); + let pts_h = client.create(Bytes::from_elems(prod_term_start)); + let pnt_h = client.create(Bytes::from_elems(prod_num_terms)); + let prb_h = client.create(Bytes::from_elems(prod_row_base)); + let poo_h = client.create(Bytes::from_elems(prod_out_offset)); + let pps_h = client.create(Bytes::from_elems(pps)); + let coarse_len = coarse.len(); + let coarse_h = client.create(Bytes::from_elems(coarse)); + let cubes = (total_pairs as u32).div_ceil(THREADS).max(1); + // Search depth over a single coarse chunk's product span, not over every product: the + // coarse index brackets the answer first, so this is `ceil(log2(span))` rather than + // `ceil(log2(num_products))`. + let search_iters = + usize::BITS as usize - (coarse_span + 1).max(1).leading_zeros() as usize; + // 80 bytes per launch; lets the kernel unpack `working` without a per-thread array. + let (pp_shift_h, pp_mask_h) = ppart_shift_mask(); + let pp_shift_len = pp_shift_h.len(); + let psh_h = client.create(Bytes::from_elems(pp_shift_h)); + let pms_h = client.create(Bytes::from_elems(pp_mask_h)); + // Segments actually populated across the four segmented stores; the rest are the + // never-indexed 1-element dummies. Passed bare so the kernel's select chain specialises to + // this many arms instead of all `MASTER_MAX_SEG`. + let num_segs = need_cs + .max(need_mk) + .max(need_basis_elems * width) + .div_ceil(seg_elems) + .max(1); + // Per-thread working size this block actually needs. `mk_len` bounds the assembled + // p-part, and a term's own p-part is at most `MAX_XI_TAU` long. Rounded to a multiple + // of 4 so the number of distinct comptime values (hence NVRTC recompiles) stays small + // while still tracking the degree — a hardcoded 16 would be right to t~510 and + // silently truncate past stem ~300. + let work_cap = (r_mk_len + .iter() + .copied() + .max() + .unwrap_or(0) + .max(MAX_XI_TAU as u32) as usize) + .div_ceil(4) + * 4; + assert!( + work_cap <= WORKING_CAP, + "block needs a working array of {work_cap} > WORKING_CAP {WORKING_CAP}; raise \ + the cap (it also bounds the host-side `xi` padding)" + ); + if launch_log_enabled() { + let mk_max = r_mk_len.iter().copied().max().unwrap_or(0); + let mk_sum: u64 = r_mk_len.iter().map(|&x| x as u64).sum(); + eprintln!( + "[launch] work_cap={work_cap} mk_max={mk_max} mk_mean={:.1} n_r={} \ + products={} pairs={} cubes={cubes}", + mk_sum as f64 / r_mk_len.len().max(1) as f64, + r_mk_len.len(), + num_products, + total_pairs, + ); + } + // Bind one `BufferArg` per `(segment vector, index)` — the `.0` handle, `.1` element length. + macro_rules! sa { + ($v:expr, $i:expr) => { + BufferArg::from_raw_parts($v[$i].0.clone(), $v[$i].1) + }; + } + // Keep every buffer this launch reads alive until the READBACK completes, not merely + // until this closure returns. With the blocking `read_one` those two coincided; with + // `read_async` this closure returns while the kernel may still be running, and a + // dropped handle lets cubecl hand its pages to a later allocation that the running + // kernel is still reading — observed as `CUDA_ERROR_LAUNCH_FAILED` at `max_t=304`. + // + // `BufferArg::from_raw_parts` consuming the handles below does NOT keep them alive: + // it takes them by value and the argument dies with the launch call. Nor do the `sa!` + // segment clones — those are temporaries too. So clone every handle here, before the + // launch consumes the originals, and hand the vec back with the future. + let keepalive: Vec = [ + tg_h.clone(), + g_h.clone(), + xi_h.clone(), + out_h.clone(), + rco_h.clone(), + rmo_h.clone(), + rcl_h.clone(), + rml_h.clone(), + rnm_h.clone(), + pri_h.clone(), + pts_h.clone(), + pnt_h.clone(), + prb_h.clone(), + poo_h.clone(), + pps_h.clone(), + coarse_h.clone(), + psh_h.clone(), + pms_h.clone(), + ] + .into_iter() + // The resident segment stores, including the padding dummies. Their segments are + // never freed while a handle lives, which is exactly the guarantee being extended. + .chain( + [&cs_seg, &mk_seg, &pp_seg, &ln_seg] + .into_iter() + .flat_map(|v| v.iter().map(|(h, _)| h.clone())), + ) + .chain(enum_keep) + .collect(); + // SAFETY: `launch_unchecked` — see the kernel's `address_type = "u64"` note. Every device + // read is in-bounds by construction (uploaded `need_*` prefix, per-segment select, `j` guards). + unsafe { + multiply_batch_kernel::launch_unchecked::( + &client, + CubeCount::Static(cubes, 1, 1), + CubeDim::new_1d(THREADS), + sa!(cs_seg, 0), + sa!(cs_seg, 1), + sa!(cs_seg, 2), + sa!(cs_seg, 3), + sa!(cs_seg, 4), + sa!(cs_seg, 5), + sa!(cs_seg, 6), + sa!(cs_seg, 7), + sa!(cs_seg, 8), + sa!(cs_seg, 9), + sa!(cs_seg, 10), + sa!(cs_seg, 11), + sa!(cs_seg, 12), + sa!(cs_seg, 13), + sa!(cs_seg, 14), + sa!(cs_seg, 15), + sa!(mk_seg, 0), + sa!(mk_seg, 1), + sa!(mk_seg, 2), + sa!(mk_seg, 3), + sa!(mk_seg, 4), + sa!(mk_seg, 5), + sa!(mk_seg, 6), + sa!(mk_seg, 7), + sa!(mk_seg, 8), + sa!(mk_seg, 9), + sa!(mk_seg, 10), + sa!(mk_seg, 11), + sa!(mk_seg, 12), + sa!(mk_seg, 13), + sa!(mk_seg, 14), + sa!(mk_seg, 15), + sa!(pp_seg, 0), + sa!(pp_seg, 1), + sa!(pp_seg, 2), + sa!(pp_seg, 3), + sa!(pp_seg, 4), + sa!(pp_seg, 5), + sa!(pp_seg, 6), + sa!(pp_seg, 7), + sa!(pp_seg, 8), + sa!(pp_seg, 9), + sa!(pp_seg, 10), + sa!(pp_seg, 11), + sa!(pp_seg, 12), + sa!(pp_seg, 13), + sa!(pp_seg, 14), + sa!(pp_seg, 15), + sa!(ln_seg, 0), + sa!(ln_seg, 1), + sa!(ln_seg, 2), + sa!(ln_seg, 3), + sa!(ln_seg, 4), + sa!(ln_seg, 5), + sa!(ln_seg, 6), + sa!(ln_seg, 7), + sa!(ln_seg, 8), + sa!(ln_seg, 9), + sa!(ln_seg, 10), + sa!(ln_seg, 11), + sa!(ln_seg, 12), + sa!(ln_seg, 13), + sa!(ln_seg, 14), + sa!(ln_seg, 15), + BufferArg::from_raw_parts(tg_h, term_gei_len), + BufferArg::from_raw_parts(g_h, g.len()), + BufferArg::from_raw_parts(xi_h, xi.len()), + BufferArg::from_raw_parts(out_h.clone(), out_len), + BufferArg::from_raw_parts(rco_h, r_cs_offset.len()), + BufferArg::from_raw_parts(rmo_h, r_mk_offset.len()), + BufferArg::from_raw_parts(rcl_h, r_cs_len.len()), + BufferArg::from_raw_parts(rml_h, r_mk_len.len()), + BufferArg::from_raw_parts(rnm_h, r_num_mats_u32.len()), + BufferArg::from_raw_parts(pri_h, num_products), + BufferArg::from_raw_parts(pts_h, num_products), + BufferArg::from_raw_parts(pnt_h, num_products), + BufferArg::from_raw_parts(prb_h, num_products), + BufferArg::from_raw_parts(poo_h, num_products), + BufferArg::from_raw_parts(pps_h, pps_len), + BufferArg::from_raw_parts(coarse_h, coarse_len), + width, + seg_elems, + num_limbs, + search_iters, + num_segs, + BufferArg::from_raw_parts(psh_h, pp_shift_len), + BufferArg::from_raw_parts(pms_h, pp_shift_len), + work_cap.min(PPART_MAX_LEN), + ); + } + + // Issue the readback but DO NOT wait for it. `read_async` enqueues the device→host copy + // into pinned memory and records a CUDA event, then hands back a future whose entire body + // is that event's wait (`cubecl-cuda` `command.rs`, `Fence::wait_sync`). Returning it + // un-awaited is what makes the pipeline deeper than one kernel: this worker goes straight + // back to `rx.recv()` and launches the next block while this one is still executing, + // whereas the previous `read_one` (= `block_on(read_async(..))`) pinned the worker here + // until the kernel retired, so each device ran exactly ONE launch at a time no matter how + // many callers were queued behind it. + // + // The caller awaits it in the wait closure below — cubecl's own `Fence` doc names this + // the intended pattern ("allows the server to continue accepting other tasks"). No + // executor is involved: the future never yields `Pending` (it wraps a blocking + // `cuEventSynchronize`), so `block_on` polls it exactly once. The buffer itself is still + // handed back with no copy (see [`BatchOutput`]); `out_h` stays alive inside the future. + let result = client.read_async(vec![out_h]); + + // Trim this stream's transient pool. Historically this per-launch cleanup RENUMBERED the + // exclusive pool's page indices (`update_page`), which under ~100-way concurrency corrupted + // cached page handles on other streams → `ManagedMemoryDescriptor` id-mismatch / + // `CUDA_ERROR_LAUNCH_FAILED` at high stems (tracel-ai/cubecl#1401). The generational-slot pool + // fix (JoeyBF/cubecl@claude/pool-slot-map-v0.10.0) gives pages stable ids so cleanup no longer + // renumbers, making this safe again — and it keeps the retained pool bounded (freed pages + // returned to the driver) so device memory tracks the working set instead of ratcheting. + // Throttled by `NASSAU_GPU_CLEANUP_EVERY` (see [`cleanup_every`]) to probe whether the residual + // high-stem `LAUNCH_FAILED` is a cross-stream cleanup-reclaim race. + // + // This now runs with this launch's work still IN FLIGHT (the readback above is not + // awaited). That is safe for the output buffer specifically — the future holds a `Handle` + // clone of `out_h`, so its page cannot be reclaimed — and it does not change the input + // buffers' exposure, which the launch already consumed and dropped before any wait even + // in the blocking version. Production runs set `NASSAU_GPU_CLEANUP_EVERY=0` regardless. + let every = cleanup_every(); + if every != 0 && CLEANUP_COUNTER.fetch_add(1, Ordering::Relaxed) % every == 0 { + client.memory_cleanup(); + } + + // The keepalive rides back with the future so the caller's wait, not this closure's + // return, is what finally releases the launch's buffers. + (result, keepalive) + }) + }); + + Box::new(move || { + // Two waits, deliberately measured apart. `pending.wait()` returns as soon as the worker has + // *launched* this block and issued its readback; `block_on` then waits for the device to + // finish. Splitting them is how a log can tell a full pipeline from an empty one: if + // `launch_ms` is most of the total the worker is the bottleneck (jobs queued behind other + // launches), whereas if `fence_ms` dominates the device is genuinely busy, which is the + // regime we want. Conflated into one figure — as they were when the worker did the readback + // — the two are indistinguishable, which is how the one-kernel-deep pipeline stayed hidden + // through three separate fan-out rewrites. + let t_launch = std::time::Instant::now(); + let ((fut, keepalive), timing) = pending.wait(); + let launch_ms = t_launch.elapsed().as_secs_f64() * 1e3; + + let t_fence = std::time::Instant::now(); + let result = cubecl_common::future::block_on(fut) + .expect("GPU readback failed") + .remove(0); + let fence_ms = t_fence.elapsed().as_secs_f64() * 1e3; + + // Only now is the output buffer free — device page and pinned host landing both — so this is + // where the byte budget must be released. Explicit rather than implicit: the whole point of + // moving it here is that dropping it earlier silently unbounds memory (see its acquisition). + // Only now can the launch's input buffers be reclaimed: the fence above is the first + // moment the kernel is known to be done reading them. Dropping these when the device + // closure returned — as the first cut of this pipelining did — let cubecl reuse pages + // under a running kernel, which crashed a stem-200 at `max_t=304`. + drop(keepalive); + drop(permit); + + BATCH_LAUNCH_US.fetch_add( + (launch_ms * 1e3) as u64, + std::sync::atomic::Ordering::Relaxed, + ); + BATCH_FENCE_US.fetch_add( + (fence_ms * 1e3) as u64, + std::sync::atomic::Ordering::Relaxed, + ); + BATCH_INTERN_US.fetch_add( + (intern_ms * 1e3) as u64, + std::sync::atomic::Ordering::Relaxed, + ); + BATCH_BASIS_US.fetch_add( + (basis_ms * 1e3) as u64, + std::sync::atomic::Ordering::Relaxed, + ); + BATCH_TGEI_US.fetch_add((tgei_ms * 1e3) as u64, std::sync::atomic::Ordering::Relaxed); + + // 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; + // Keep the value this call was assigned: with ~100 workers incrementing, a separate `load` + // races past exact multiples, so a `% every == 0` test on it can fire never (observed: zero + // reports over 12 minutes). `fetch_add` returns a unique ticket per call, so exactly one + // caller sees each multiple. + let call_no = BATCH_CALLS.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1; + 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(real_pairs as u64, std::sync::atomic::Ordering::Relaxed); + BATCH_PREP_US.fetch_add((prep_ms * 1e3) as u64, std::sync::atomic::Ordering::Relaxed); + BATCH_WAIT_US.fetch_add((wait_ms * 1e3) as u64, std::sync::atomic::Ordering::Relaxed); + BATCH_PERMIT_US.fetch_add( + (permit_ms * 1e3) as u64, + std::sync::atomic::Ordering::Relaxed, + ); + BATCH_LOCK_US.fetch_add((lock_ms * 1e3) as u64, std::sync::atomic::Ordering::Relaxed); + BATCH_QUEUE_US.fetch_add( + (timing.queue_ms * 1e3) as u64, + std::sync::atomic::Ordering::Relaxed, + ); + BATCH_EXEC_US.fetch_add( + (timing.exec_ms * 1e3) as u64, + std::sync::atomic::Ordering::Relaxed, + ); + BATCH_DEPTH_SUM.fetch_add(timing.depth, std::sync::atomic::Ordering::Relaxed); + BATCH_DEPTH_MAX.fetch_max(timing.depth, std::sync::atomic::Ordering::Relaxed); + + // Periodic split of where multiply time actually goes. The counters above were being collected + // and never read ([`take_batch_stats`] had no callers), which left the dominant cost of a + // resolution unattributed: profiling a stem-200 run showed ~96% of the slow bidegrees' time + // inside the per-signature parallel section (row reduction was ~2%), but nothing said whether + // that is host marshalling or device execution. Non-resetting reads so the totals stay + // cumulative; `NASSAU_BATCH_REPORT_EVERY=0` disables. + let every = batch_report_every(); + if every != 0 && call_no % every == 0 { + let calls = call_no; + let marshal_s = + BATCH_MARSHAL_US.load(std::sync::atomic::Ordering::Relaxed) as f64 / 1e6; + let device_s = + BATCH_DEVICE_US.load(std::sync::atomic::Ordering::Relaxed) as f64 / 1e6; + let pairs = BATCH_PAIRS.load(std::sync::atomic::Ordering::Relaxed); + let prep_s = BATCH_PREP_US.load(std::sync::atomic::Ordering::Relaxed) as f64 / 1e6; + let wait_s = BATCH_WAIT_US.load(std::sync::atomic::Ordering::Relaxed) as f64 / 1e6; + let permit_s = + BATCH_PERMIT_US.load(std::sync::atomic::Ordering::Relaxed) as f64 / 1e6; + let lock_s = BATCH_LOCK_US.load(std::sync::atomic::Ordering::Relaxed) as f64 / 1e6; + let queue_s = + BATCH_QUEUE_US.load(std::sync::atomic::Ordering::Relaxed) as f64 / 1e6; + let exec_s = BATCH_EXEC_US.load(std::sync::atomic::Ordering::Relaxed) as f64 / 1e6; + let depth_sum = BATCH_DEPTH_SUM.load(std::sync::atomic::Ordering::Relaxed); + let depth_max = BATCH_DEPTH_MAX.load(std::sync::atomic::Ordering::Relaxed); + let launch_s = + BATCH_LAUNCH_US.load(std::sync::atomic::Ordering::Relaxed) as f64 / 1e6; + let fence_s = + BATCH_FENCE_US.load(std::sync::atomic::Ordering::Relaxed) as f64 / 1e6; + let el = ENUM_LAUNCHES.load(Ordering::Relaxed); + let erm = ENUM_RS_MAX.load(Ordering::Relaxed); + let total = (prep_s + wait_s + device_s).max(1e-9); + // Emit the theta->bytes curve alongside the periodic stats, not only at the end: a + // run that dies on an allocation must still leave behind the answer to "which theta + // would have fit", which is the whole point of collecting it. + dump_master_by_degree(); + eprintln!( + "[batch-stats] calls={calls} prep={prep_s:.1}s permit={permit_s:.1}s \ + lock={lock_s:.1}s device={device_s:.1}s | prep={:.0}% permit={:.0}% \ + lock={:.0}% device={:.0}% pairs={pairs} (marshal={marshal_s:.1}s \ + wait={wait_s:.1}s) queue={queue_s:.1}s exec={exec_s:.1}s | queue={:.0}% \ + exec={:.0}% depth mean={:.1} max={depth_max} | launch={launch_s:.1}s \ + fence={fence_s:.1}s pipeline={:.0}% | intern={:.1}s basis={:.1}s tgei={:.1}s \ + | enum launches={el} Rs/launch mean={:.0} max={erm} blocks/launch mean={:.0} \ + waves/SM={:.3}", + 100.0 * prep_s / total, + 100.0 * permit_s / total, + 100.0 * lock_s / total, + 100.0 * device_s / total, + 100.0 * queue_s / total, + 100.0 * exec_s / total, + depth_sum as f64 / calls as f64, + // Share of the caller's wait spent on the device rather than queued behind other + // launches. ~100% is a full pipeline; the pre-change one-kernel-deep behaviour + // drives this toward 0 as callers pile up. + 100.0 * fence_s / (launch_s + fence_s).max(1e-9), + BATCH_INTERN_US.load(std::sync::atomic::Ordering::Relaxed) as f64 / 1e6, + BATCH_BASIS_US.load(std::sync::atomic::Ordering::Relaxed) as f64 / 1e6, + BATCH_TGEI_US.load(std::sync::atomic::Ordering::Relaxed) as f64 / 1e6, + ENUM_RS.load(Ordering::Relaxed) as f64 / el.max(1) as f64, + ENUM_BLOCKS.load(Ordering::Relaxed) as f64 / el.max(1) as f64, + // 132 SMs x 24 resident blocks (the measured Block Limit Registers at 40 + // regs/thread) = 3168 block slots on an H200. + ENUM_BLOCKS.load(Ordering::Relaxed) as f64 / el.max(1) as f64 / 3168.0, + ); + } + + result + }) as BlockWait + }) +} + +/// How often [`multiply_batch_block`] prints the cumulative marshal/device split, in launches. +/// `NASSAU_BATCH_REPORT_EVERY` (default 2000; `0` disables). +/// `NASSAU_GPU_LAUNCH_LOG=1` dumps each multiply launch's shape (`work_cap`, the `mk_len` spread, +/// product count, pair count, grid size) to stderr. +/// +/// This exists because launch shape is otherwise invisible, and inferring it from aggregates is how +/// several wrong conclusions got made here. It settled one directly: two bench processes each report +/// ~7x the solo `pairs/s`, which looked like enormous headroom, and the log showed both were issuing +/// *identical* launches (same `work_cap`, same ~2.5e9 pairs, same ~9.6e6 blocks). Same launch, same +/// device, "7x faster" — so the throughput figure, not the kernel, was the thing that changed. The +/// bench's fixed-work warm-up (4 jobs, ~190 s regardless of co-tenancy) confirmed real throughput is +/// flat. Read `pairs/s` from a contended run as meaningless, not as headroom. +fn launch_log_enabled() -> bool { + use std::sync::OnceLock; + static ON: OnceLock = std::sync::OnceLock::new(); + *ON.get_or_init(|| std::env::var_os("NASSAU_GPU_LAUNCH_LOG").is_some()) +} - let t_device = std::time::Instant::now(); +fn batch_report_every() -> u64 { + use std::sync::OnceLock; + static EVERY: OnceLock = std::sync::OnceLock::new(); + *EVERY.get_or_init(|| { + std::env::var("NASSAU_BATCH_REPORT_EVERY") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(2000) + }) +} - // Pin the whole device section to one CUDA stream (see [`GPU_STREAM`]) so a single - // memory pool is reclaimed by `memory_cleanup`. Held under the `resident` lock, so this - // stream is used by at most one thread at a time. - let result = GPU_STREAM.executes(|| { - let client = CudaRuntime::client(&CudaDevice::default()); - // Resident admissible buffers: (re-)upload the master only when it grew this - // launch; otherwise reuse the handle from a previous launch and upload nothing. - if resident.cs_handle.is_none() || resident.cs_uploaded != resident.col_sums.len() { - resident.cs_handle = Some(client.create_from_slice(u16::as_bytes(&resident.col_sums))); - resident.cs_uploaded = resident.col_sums.len(); +/// Read `data[o]` from a master split into up to [`MASTER_MAX_SEG`] fixed-size segments of +/// `seg_elems` elements each: segment `o / seg_elems`, local index `o % seg_elems`. This is the +/// no-copy-growth replacement for a single contiguous `Array` — appending a segment never +/// reallocates/copies the existing ones, so the device peak is `live_size + one_segment` instead of +/// the `~2×` realloc-doubling transient that pushes cubecl into its memory-corruption regime. The +/// per-segment `Array`s are separate kernel args because cubecl has no array-of-buffers; the branch +/// is the price of staying inside cubecl (vs raw CUDA VMM). `seg_elems` is runtime so tests can use +/// tiny segments; production sets it large. Keep the branch chain length equal to [`MASTER_MAX_SEG`]. +#[cube] +#[allow(clippy::too_many_arguments)] +fn seg_read_u16( + s0: &[u16], + s1: &[u16], + s2: &[u16], + s3: &[u16], + s4: &[u16], + s5: &[u16], + s6: &[u16], + s7: &[u16], + s8: &[u16], + s9: &[u16], + s10: &[u16], + s11: &[u16], + s12: &[u16], + s13: &[u16], + s14: &[u16], + s15: &[u16], + o: usize, + seg_elems: usize, + #[comptime] num_segs: usize, +) -> u16 { + // `num_segs` is passed bare, so cubecl bakes it in as a compile-time constant and every + // `num_segs > k` below folds away — the chain specialises to the segments that actually exist + // instead of always testing all [`MASTER_MAX_SEG`] of them. This is the hot path: the gather + // loop calls a `seg_read` 3 x WORKING_CAP = 96 times per thread, so a 16-way compare/select + // chain per call dominated the instruction stream. Measured on an H200 with the kernel + // resident: SM Active 100%, SM Issue 68%, DRAM read 1.1% of peak — pure instruction cost, not + // bandwidth. The single-segment case skips the division too. + let mut v = 0u16; + if num_segs == 1 { + v = s0[o]; + } else { + let seg = o / seg_elems; + let local = o % seg_elems; + if seg == 0 { + v = s0[local]; + } else if num_segs > 1 && seg == 1 { + v = s1[local]; + } else if num_segs > 2 && seg == 2 { + v = s2[local]; + } else if num_segs > 3 && seg == 3 { + v = s3[local]; + } else if num_segs > 4 && seg == 4 { + v = s4[local]; + } else if num_segs > 5 && seg == 5 { + v = s5[local]; + } else if num_segs > 6 && seg == 6 { + v = s6[local]; + } else if num_segs > 7 && seg == 7 { + v = s7[local]; + } else if num_segs > 8 && seg == 8 { + v = s8[local]; + } else if num_segs > 9 && seg == 9 { + v = s9[local]; + } else if num_segs > 10 && seg == 10 { + v = s10[local]; + } else if num_segs > 11 && seg == 11 { + v = s11[local]; + } else if num_segs > 12 && seg == 12 { + v = s12[local]; + } else if num_segs > 13 && seg == 13 { + v = s13[local]; + } else if num_segs > 14 && seg == 14 { + v = s14[local]; + } else { + v = s15[local]; } - 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(); + } + v +} + +/// `u32` sibling of [`seg_read_u16`] (the resident basis `lens` are u32). Same static-branch segment +/// select; see [`seg_read_u16`] for the layout and rationale. +#[cube] +#[allow(clippy::too_many_arguments)] +fn seg_read_u32( + s0: &[u32], + s1: &[u32], + s2: &[u32], + s3: &[u32], + s4: &[u32], + s5: &[u32], + s6: &[u32], + s7: &[u32], + s8: &[u32], + s9: &[u32], + s10: &[u32], + s11: &[u32], + s12: &[u32], + s13: &[u32], + s14: &[u32], + s15: &[u32], + o: usize, + seg_elems: usize, + #[comptime] num_segs: usize, +) -> u32 { + // `num_segs` is passed bare, so cubecl bakes it in as a compile-time constant and every + // `num_segs > k` below folds away — the chain specialises to the segments that actually exist + // instead of always testing all [`MASTER_MAX_SEG`] of them. This is the hot path: the gather + // loop calls a `seg_read` 3 x WORKING_CAP = 96 times per thread, so a 16-way compare/select + // chain per call dominated the instruction stream. Measured on an H200 with the kernel + // resident: SM Active 100%, SM Issue 68%, DRAM read 1.1% of peak — pure instruction cost, not + // bandwidth. The single-segment case skips the division too. + let mut v = 0u32; + if num_segs == 1 { + v = s0[o]; + } else { + let seg = o / seg_elems; + let local = o % seg_elems; + if seg == 0 { + v = s0[local]; + } else if num_segs > 1 && seg == 1 { + v = s1[local]; + } else if num_segs > 2 && seg == 2 { + v = s2[local]; + } else if num_segs > 3 && seg == 3 { + v = s3[local]; + } else if num_segs > 4 && seg == 4 { + v = s4[local]; + } else if num_segs > 5 && seg == 5 { + v = s5[local]; + } else if num_segs > 6 && seg == 6 { + v = s6[local]; + } else if num_segs > 7 && seg == 7 { + v = s7[local]; + } else if num_segs > 8 && seg == 8 { + v = s8[local]; + } else if num_segs > 9 && seg == 9 { + v = s9[local]; + } else if num_segs > 10 && seg == 10 { + v = s10[local]; + } else if num_segs > 11 && seg == 11 { + v = s11[local]; + } else if num_segs > 12 && seg == 12 { + v = s12[local]; + } else if num_segs > 13 && seg == 13 { + v = s13[local]; + } else if num_segs > 14 && seg == 14 { + v = s14[local]; + } else { + v = s15[local]; } - 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(); - // Shared across every chunk: term data, seqno/xi tables, per-`R` offsets, and the output - // buffer. `prod_term_start` values index the global `term_*` arrays, so each chunk reuses - // these handles unchanged; only the per-product record slices and the pair prefix sum are - // rebuilt per chunk. - 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 zeros = vec![0u32; out_len]; - let out_h = client.create_from_slice(u32::as_bytes(&zeros)); - const THREADS: u32 = 256; + } + v +} - // Launch the products in chunks each holding at most `GPU_PAIR_CHUNK` pairs, so every - // kernel's thread count (and thus `ABSOLUTE_POS`) stays under `2^32`. Each product writes - // its F₂ bits into `out_h` with atomic XOR keyed by its global `row`/`out_offset`, so - // splitting the product set across launches and accumulating into the shared buffer is - // exact (XOR is associative and order-independent). - let mut c0 = 0usize; - while c0 < products.len() { - // Grow the chunk product-by-product until the next one would exceed the cap; always take - // at least one product (a single product's pair count is far below the cap). - let mut c1 = c0; - let mut chunk_pairs = 0usize; - while c1 < products.len() - && (c1 == c0 || chunk_pairs + prod_pairs[c1] <= GPU_PAIR_CHUNK) - { - chunk_pairs += prod_pairs[c1]; - c1 += 1; - } - - // Chunk-local pair prefix sum (values < cap, fit `u32`), sentinel at the end. - let mut pps_chunk: Vec = Vec::with_capacity(c1 - c0 + 1); - let mut acc = 0u32; - for &pairs in &prod_pairs[c0..c1] { - pps_chunk.push(acc); - acc += pairs as u32; - } - pps_chunk.push(acc); - - let pri_h = client.create_from_slice(u32::as_bytes(&prod_r_index[c0..c1])); - let pts_h = client.create_from_slice(u32::as_bytes(&prod_term_start[c0..c1])); - let pnt_h = client.create_from_slice(u32::as_bytes(&prod_num_terms[c0..c1])); - let prb_h = client.create_from_slice(u32::as_bytes(&prod_row_base[c0..c1])); - let poo_h = client.create_from_slice(u32::as_bytes(&prod_out_offset[c0..c1])); - let pps_h = client.create_from_slice(u32::as_bytes(&pps_chunk)); - let cubes = (chunk_pairs as u32).div_ceil(THREADS).max(1); - unsafe { - multiply_batch_kernel::launch::( - &client, +/// PROFILE (ncu, idle H200, degree <= 130: 89392 `R`s / 76M matrices). This kernel is ~99% of GPU +/// kernel time whenever the transient path is live (nsys; `multiply_batch_kernel` is 1.1%), and it is +/// LATENCY-bound, not bandwidth- or compute-bound: +/// +/// ```text +/// Compute (SM) Throughput 5.16 % Avg. Active Threads Per Warp 5.88 / 32 +/// Memory Throughput 13.92 % Active Warps Per Scheduler 2.09 +/// DRAM Throughput 0.65 % No Eligible 76.57 % +/// L2 Hit Rate 98.83 % Warp Cycles Per Issued Instr 8.90 +/// ``` +/// +/// Nothing is saturated; the scheduler simply has nothing to issue 3 cycles in 4. Two causes, in +/// order of size: +/// +/// 1. DIVERGENCE — 5.88 of 32 lanes active. One thread per `R`, and `R`s have wildly different +/// matrix counts, so a warp runs at its longest member and ~82% of lanes idle. Sorting by +/// `num_mats` (dbe1b49e85) won 8% and left this untouched. Unaddressed; the biggest number here. +/// 2. LOCAL MEMORY — the state below is indexed by runtime values, so it cannot be registers and +/// lands in local memory; ncu attributes 45.6% of the 8.9 stall cycles to waiting on it. +/// +/// Moving the state to `Shared<[u32]>` (stride `elem * BLOCK + tid`, bank-conflict-free) WAS tried +/// and is bit-exact, but is a net LOSS: it cuts warp cycles per issued instruction 8.90 -> 5.92 +/// (-33%, so the mechanism is real) while forcing the block from 256 to 64 threads (152 u32/thread +/// = 38.9 KB/block, at the 48 KB static shared limit), which drops active warps per scheduler +/// 2.09 -> 1.61. Net 109.63 ms -> 112.72 ms. The 33% is only reachable if the shared footprint +/// shrinks enough to keep occupancy — e.g. `matrix` as u16 (values are <= 2^11). +/// +/// Do not read ncu's "Est. Local Speedup: 76.57%" as a forecast: it is the prize if the stall +/// vanishes at zero cost, and here the cost was occupancy. +/// +/// ROOT CAUSE (ncu Occupancy/LaunchStats, same run): the kernel is STARVED OF WORK, and everything +/// above is a symptom of it. +/// +/// ```text +/// Registers Per Thread 40 Block Limit Registers 24 blocks +/// Theoretical Occupancy 75 % Waves Per SM 0.44 +/// Achieved Occupancy 8.16 % Grid Size 1397 blocks x 64 threads +/// ``` +/// +/// Registers are NOT a limiter (40/thread allows 24 blocks/SM, 75% theoretical) — do not spend +/// effort shrinking them. The binding fact is `Waves Per SM = 0.44`: an H200 has 132 SMs x 24 block +/// slots = 3168, and the grid supplies 1397. One thread per `R` over 89392 `R`s is ~89k threads +/// against the device's ~270k capacity, so the kernel cannot fill half the machine however well it +/// runs — and achieved occupancy is then 8.16% of that 75% ceiling because those few blocks drain +/// raggedly (the same divergence as the 5.88/32 lanes). +/// +/// Production is likely worse: this measurement enumerates every `R` to degree 130 at once, while a +/// production launch covers only one block's transient `R`s, so those grids are smaller still. +/// +/// So the granularity is the bug: one thread per `R` gives too FEW threads and wildly UNEQUAL ones. +/// That is why every local tweak measured between -3% and +8% — they adjust resources that are not +/// the constraint. The fixes that would matter are structural: batch far more `R`s per launch so the +/// grid fills the device, or drop thread-per-`R` for lane cooperation with dynamically pulled work. +/// +/// MEASURED (S_2 stem 150, max_s 60, theta=125), and it narrows the choice further: +/// +/// | ENUM_THREADS | blocks/launch | waves/SM | wall | +/// |--------------|---------------|----------|-------| +/// | 256 | 6 | 0.002 | 561 s | +/// | 64 | 21 | 0.007 | 569 s | +/// | 32 | 41 | 0.013 | 542 s | +/// +/// 6.5x more blocks buys 3.4% — noise. Spreading a launch over more SMs cannot help, because a +/// launch's duration is set by its LONGEST SINGLE `R` (one thread, sequential odometer), not by how +/// many blocks it occupies; the extra SMs just idle beside the one thread still grinding. `Rs/launch` +/// averages 1293 and peaks at 34157, so that spread is enormous. +/// +/// 10260 launches over a ~550 s run is ~50 ms each, and [`gpu_thread`] runs every device section on +/// ONE thread and ONE stream — so they are strictly serial. That is what makes batching worth a +/// factor of hundreds rather than a few percent: merging launches turns a SUM into a MAX. Ten +/// serialised 50 ms launches cost 500 ms; the same `R`s in one launch cost ~50 ms, because the short +/// `R`s run beside the long one instead of queueing behind it. +/// +/// So the two candidates are (a) stream concurrency — the launches are already independent, and +/// [`gpu_thread`]'s FIFO order could dispatch round-robin across N streams and stay fair — or +/// (b) aggregating `R`s across calls, which needs enumeration decoupled from the multiply that +/// consumes it. (a) is far cheaper, but streams were pinned to 1 to fix a host-memory blowup from +/// per-stream pinned pools, so that constraint has to be re-examined, not ignored. +/// +/// In-kernel admissible-matrix enumeration: one thread per distinct `R`, generating that `R`'s +/// `col_sums`/`masks` for *every* admissible matrix directly into device scratch — the on-GPU +/// replacement for the resident/uploaded master (the stem-300 memory wall + the eviction re-upload +/// cost). This is the cubecl transcription of [`enumerate_admissible_ref`] (validated bit-exact on +/// the CPU), using only flag-guarded control flow — `while … && !found` in place of the odometer's +/// `break`/early `return`, `handled` in place of its `continue`. All per-thread state lives in +/// fixed-size local `Array`s ([`ENUM_MATRIX_CAP`] etc.); the values are small (≤ `u16`), stored as +/// `u16` exactly like the uploaded master so the multiply kernel reads them unchanged. +/// +/// Inputs are per-`R`: `p_parts` (`n_r × width`, zero-padded), `r_rows`/`r_cols` (its dimensions), +/// and `r_cs_out`/`r_mk_out` (its base offset, in `u16` units, into the shared `out_cs`/`out_mk` +/// scratch — a host prefix-sum of `num_mats × cs_len` / `num_mats × mk_len`). `out_counts[ri]` +/// receives the number of matrices the thread emitted, so a count-only pre-pass can drive the +/// prefix-sum without any host enumeration. Runtime-agnostic: the same kernel lowers to CUDA (the +/// H200 path) and to the `cpu` backend (used by `admissible_enum_gpu_matches` to cross-check the +/// device lowering against `enumerate_admissible_ref` without a GPU). +// +// 64-bit addressing (`address_type = "u64"`, like `multiply_batch_kernel`): a big all-rows block's +// `out_cs`/`out_mk` scratch reaches multiple GB at high stems, so the flat element index (and the +// byte offset cubecl derives from it) overflows the default `u32` address type — the write lands at +// a wild address → `CUDA_ERROR_LAUNCH_FAILED`. `launch_unchecked` because checked u64 mode emits a +// `min(u64, u64)` NVRTC rejects; every access here is in-bounds by construction (the write offset is +// `r_cs_out[ri] + mat*cs_len + j < need_cs`, the scratch length, and reads are bounded by `n_r`). +#[cube(launch_unchecked, address_type = "u64")] +#[allow(clippy::too_many_arguments)] +fn enumerate_admissible_kernel( + p_parts: &[u32], + r_rows: &[u32], + r_cols: &[u32], + // u64: the scratch offsets index buffers that reach billions of elements in a big block, past + // `u32::MAX` (same reason the multiply's `r_cs_offset`/`r_mk_offset` are u64 — these ARE those). + r_cs_out: &[u64], + r_mk_out: &[u64], + out_cs: &mut [u16], + out_mk: &mut [u16], + out_counts: &mut [u32], + width: usize, + n_r: usize, + // Comptime: `false` compiles the emit stores OUT entirely (not merely branches around them), + // isolating the odometer's cost from the cost of writing its results. Production passes `true`; + // only the `enum_emit_store_cost` diagnostic passes `false`. + // Wrap mask for `emit == 2` only (power-of-two - 1), keeping the probe's rewritten indices in + // bounds; ignored by every other mode. + wrap: u32, + #[comptime] emit: u32, +) { + let ri = ABSOLUTE_POS; + if ri >= n_r { + terminate!(); + } + let rows = usize::cast_from(r_rows[ri]); + let cols = usize::cast_from(r_cols[ri]); + let cs_len = cols - 1; + let mk_len = rows + cols - 1; + let pbase = ri * width; + let cs_base = usize::cast_from(r_cs_out[ri]); + let mk_base = usize::cast_from(r_mk_out[ri]); + + // Per-thread local state, mirroring `AdmissibleMatrix` / `enumerate_admissible_ref`. CUDA local + // arrays are uninitialized, so every slot up to the comptime cap is explicitly zeroed first. + // Shared memory, not `Array` — `Array` is CUDA *local* memory, and because every index here is + // a runtime value none of it can be register-allocated. ncu attributed 45.6% of this kernel's + // 8.9 stall cycles to waiting on those local accesses. See [`ENUM_BLOCK`] for why the occupancy + // this costs is free in production even though it was not on the benchmark. + let tid = usize::cast_from(UNIT_POS); + let bs = ENUM_BLOCK as usize; + let mut st = Shared::<[u32]>::new_slice(ENUM_STATE * ENUM_BLOCK as usize); + // Zero only the region this `R` can actually reach, not the whole comptime cap. Every index the + // enumeration below forms is bounded by these: `matrix` by `row*cols+col < rows*cols`, `totals` + // by `rows`, `col_sums` by `cols-1 == cs_len`, and `masks` by `(rows-1)+(cols-1) < mk_len`. A + // typical `R` is far smaller than the cap (rows ~4-8 against 10, cols ~6-9 against the cap), so + // clearing the full arrays spent most of these stores on slots no read ever touches — and they + // are local-memory stores, not register writes. + for i in 0..rows * cols { + st[(i) * bs + tid] = 0u32; + } + for i in 0..rows { + st[(ENUM_ST_TOTALS + i) * bs + tid] = 0u32; + } + for i in 0..cs_len { + st[(ENUM_ST_COLSUMS + i) * bs + tid] = 0u32; + } + for i in 0..mk_len { + st[(ENUM_ST_MASKS + i) * bs + tid] = 0u32; + } + // Column 0 of the matrix (and the initial masks) is the padded p_part. + for i in 0..rows { + let x = p_parts[pbase + i]; + st[(i * cols) * bs + tid] = x; + st[(ENUM_ST_MASKS + i) * bs + tid] = x; + } + + let mut mat = 0usize; + let mut more = true; + while more { + // Emit the current matrix's col_sums/masks into this R's scratch slot. + // 0 = no stores at all (isolates the odometer). 1 = production: this R's slot, contiguous + // per thread but scattered across the warp. 2 = lane-adjacent indices, which coalesces the + // warp but writes garbage layout -- a TIMING PROBE ONLY, never correct output. 3 = every + // other entry, halving both store count and bytes. + if emit == 1 { + let co = cs_base + mat * cs_len; + for j in 0..cs_len { + out_cs[co + j] = u16::cast_from(st[(ENUM_ST_COLSUMS + j) * bs + tid]); + } + let mo = mk_base + mat * mk_len; + for j in 0..mk_len { + out_mk[mo + j] = u16::cast_from(st[(ENUM_ST_MASKS + j) * bs + tid]); + } + } else if emit == 2 { + let w = usize::cast_from(wrap); + for j in 0..cs_len { + out_cs[(ri + n_r * (mat * cs_len + j)) & w] = + u16::cast_from(st[(ENUM_ST_COLSUMS + j) * bs + tid]); + } + for j in 0..mk_len { + out_mk[(ri + n_r * (mat * mk_len + j)) & w] = + u16::cast_from(st[(ENUM_ST_MASKS + j) * bs + tid]); + } + } else if emit == 3 { + let co = cs_base + mat * cs_len; + let mut j = 0usize; + while j < cs_len { + out_cs[co + j] = u16::cast_from(st[(ENUM_ST_COLSUMS + j) * bs + tid]); + j += 2; + } + let mo = mk_base + mat * mk_len; + let mut j2 = 0usize; + while j2 < mk_len { + out_mk[mo + j2] = u16::cast_from(st[(ENUM_ST_MASKS + j2) * bs + tid]); + j2 += 2; + } + } + mat += 1; + + // One `next()` step: `found` = produced a new matrix (the ref's `return true`); `handled` + // = this column already updated `totals` (the ref's `continue`). Loops guard on `!found`. + let mut found = false; + let mut row = 0usize; + while row < rows && !found { + let mut p_to_the_j = 1u32; + st[(ENUM_ST_TOTALS + row) * bs + tid] = st[(row * cols) * bs + tid]; + let mut col = 1usize; + while col < cols && !found { + p_to_the_j *= 2u32; + let mut handled = false; + if p_to_the_j <= st[(ENUM_ST_TOTALS + row) * bs + tid] { + // Bitsum along the anti-diagonal to the bottom-left (saturating start index). + let mut d = 0u32; + let mut c = 0usize; + if row + col + 1 > rows { + c = row + col + 1 - rows; + } + while c < col { + d |= st[((row + col - c) * cols + c) * bs + tid]; + c += 1; + } + let cur = st[(row * cols + col) * bs + tid]; + let new_entry = ((cur | d) + 1u32) & !d; + let inc = new_entry - cur; + let sub = inc * p_to_the_j; + if st[(ENUM_ST_TOTALS + row) * bs + tid] < sub { + st[(ENUM_ST_TOTALS + row) * bs + tid] = + st[(ENUM_ST_TOTALS + row) * bs + tid] + p_to_the_j * cur; + handled = true; + } else { + st[(row * cols) * bs + tid] = st[(ENUM_ST_TOTALS + row) * bs + tid] - sub; + st[(ENUM_ST_MASKS + row) * bs + tid] = st[(row * cols) * bs + tid]; + st[(ENUM_ST_COLSUMS + col - 1) * bs + tid] = + st[(ENUM_ST_COLSUMS + col - 1) * bs + tid] + inc; + let mut j = 1usize; + while j < col { + st[(ENUM_ST_MASKS + row + j) * bs + tid] = st + [(ENUM_ST_MASKS + row + j) * bs + tid] + & !st[(row * cols + j) * bs + tid]; + st[(ENUM_ST_COLSUMS + j - 1) * bs + tid] = st + [(ENUM_ST_COLSUMS + j - 1) * bs + tid] + - st[(row * cols + j) * bs + tid]; + st[(row * cols + j) * bs + tid] = 0u32; + j += 1; + } + st[(row * cols + col) * bs + tid] = new_entry; + let mut i = 0usize; + while i < row { + st[(i * cols) * bs + tid] = st[(ENUM_ST_TOTALS + i) * bs + tid]; + st[(ENUM_ST_MASKS + i) * bs + tid] = + st[(ENUM_ST_TOTALS + i) * bs + tid]; + let mut j2 = 1usize; + while j2 < cols { + if i + j2 > row { + st[(ENUM_ST_MASKS + i + j2) * bs + tid] = st + [(ENUM_ST_MASKS + i + j2) * bs + tid] + & !st[(i * cols + j2) * bs + tid]; + } + st[(ENUM_ST_COLSUMS + j2 - 1) * bs + tid] = st + [(ENUM_ST_COLSUMS + j2 - 1) * bs + tid] + - st[(i * cols + j2) * bs + tid]; + st[(i * cols + j2) * bs + tid] = 0u32; + j2 += 1; + } + i += 1; + } + st[(ENUM_ST_MASKS + row + col) * bs + tid] = d | new_entry; + found = true; + handled = true; + } + } + if !handled { + st[(ENUM_ST_TOTALS + row) * bs + tid] = st[(ENUM_ST_TOTALS + row) * bs + tid] + + p_to_the_j * st[(row * cols + col) * bs + tid]; + } + col += 1; + } + row += 1; + } + more = found; + } + + out_counts[ri] = u32::cast_from(mat); +} + +#[cfg(test)] +mod tests { + use super::*; + + /// 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: &[u32], b: &[u32], out: &mut [u32]) { + 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 = gpu_client(); + + let a_handle = client.create_from_slice(u32::as_bytes(a)); + let b_handle = client.create_from_slice(u32::as_bytes(b)); + let out_handle = client.empty(std::mem::size_of_val(a)); + + // One 1-D block of `THREADS` units, enough blocks to cover every limb. + const THREADS: u32 = 256; + let cubes = (n as u32).div_ceil(THREADS); + unsafe { + xor_f2::launch::( + &client, + CubeCount::Static(cubes, 1, 1), + CubeDim::new_1d(THREADS), + BufferArg::from_raw_parts(a_handle, n), + BufferArg::from_raw_parts(b_handle, n), + BufferArg::from_raw_parts(out_handle.clone(), n), + ); + } + + let bytes = client.read_one(out_handle).unwrap(); + u32::from_bytes(&bytes).to_vec() + } + + /// 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: &[u32], + xi: &[u32], + p_parts: &[u32], + out: &mut [u32], + width: usize, + pp_shift: &[u32], + pp_mask: &[u32], + #[comptime] sq_len: usize, + ) { + let idx = ABSOLUTE_POS; + if idx >= out.len() { + terminate!(); + } + let base = idx * width; + + let mut working = 0u64; + for h in 0..PPART_MAX_LEN { + if h < width { + working |= u64::cast_from(p_parts[base + h]) << u64::cast_from(pp_shift[h]); + } + } + out[idx] = seqno_core_packed(g, xi, pp_shift, pp_mask, working, sq_len, 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 = gpu_client(); + + 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::()); + + let (psh, pms) = ppart_shift_mask(); + let pp_len = psh.len(); + let psh_h = client.create(Bytes::from_elems(psh)); + let pms_h = client.create(Bytes::from_elems(pms)); + 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), + BufferArg::from_raw_parts(g_h, g.len()), + BufferArg::from_raw_parts(xi_h, xi.len()), + BufferArg::from_raw_parts(pp_h, p_parts.len()), + BufferArg::from_raw_parts(out_h.clone(), n), + width, + BufferArg::from_raw_parts(psh_h, pp_len), + BufferArg::from_raw_parts(pms_h, pp_len), + width.min(PPART_MAX_LEN), + ); + } + + let bytes = client.read_one(out_h).unwrap(); + u32::from_bytes(&bytes).to_vec() + } + + /// 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: &[u16], + masks: &[u16], + term_pparts: &[u16], + term_lens: &[u32], + g: &[u32], + xi: &[u32], + out: &mut [Atomic], + num_terms: usize, + num_matrices: usize, + cs_len: usize, + mk_len: usize, + width: usize, + num_limbs: usize, + pp_shift: &[u32], + pp_mask: &[u32], + ) { + 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, + num_limbs, + PPART_MAX_LEN, + pp_shift, + pp_mask, + ); + } + + /// 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| narrow_u16(v)).collect(); + let masks: Vec = mk32.iter().map(|&v| narrow_u16(v)).collect(); + let num_matrices = masks.len() / mk_len; + + // Terms of s, each p_part padded to `width`, with their true (trimmed) lengths. + let num_terms = term_indices.len(); + let mut term_pparts = vec![0u16; num_terms * width]; + let mut term_lens = vec![0u32; num_terms]; + for (t, &ti) in term_indices.iter().enumerate() { + let elt = algebra.basis_element_from_index(s_degree, ti); + term_lens[t] = elt.p_part.len() as u32; + for (slot, v) in term_pparts[t * width..(t + 1) * width] + .iter_mut() + .zip(elt.p_part.iter()) + { + *slot = narrow_u16(v); + } + } + + let out_degree = r_degree + s_degree; + let dim = algebra.dimension(out_degree); + let num_limbs = dim.div_ceil(32).max(1); + + // Device buffers must be non-empty; `cs_len == 0` (R's max entry is 1) leaves + // `col_sums` empty. The kernel never reads past the real lengths. + if col_sums.is_empty() { + col_sums.push(0); + } + + let client = gpu_client(); + 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 (psh, pms) = ppart_shift_mask(); + let pp_len = psh.len(); + let psh_h = client.create(Bytes::from_elems(psh)); + let pms_h = client.create(Bytes::from_elems(pms)); + 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), + BufferArg::from_raw_parts(cs_h, col_sums.len()), + BufferArg::from_raw_parts(mk_h, masks.len()), + BufferArg::from_raw_parts(tp_h, term_pparts.len()), + BufferArg::from_raw_parts(tl_h, term_lens.len()), + BufferArg::from_raw_parts(g_h, g.len()), + BufferArg::from_raw_parts(xi_h, xi.len()), + BufferArg::from_raw_parts(out_h.clone(), num_limbs), + num_terms, + num_matrices, + cs_len, + mk_len, + width, + num_limbs, + BufferArg::from_raw_parts(psh_h, pp_len), + BufferArg::from_raw_parts(pms_h, pp_len), + ); + } + + let bytes = client.read_one(out_h).unwrap(); + u32::from_bytes(&bytes).to_vec() + } + + /// Validation kernel for [`seg_read_u16`]: `out[i] = segmented[idx[i]]`. Lets a test assert the + /// segmented read reproduces a contiguous buffer bit-for-bit (see `seg_read_matches_contiguous`). + #[cube(launch)] + #[allow(clippy::too_many_arguments)] + fn seg_gather_kernel( + s0: &[u16], + s1: &[u16], + s2: &[u16], + s3: &[u16], + s4: &[u16], + s5: &[u16], + s6: &[u16], + s7: &[u16], + s8: &[u16], + s9: &[u16], + s10: &[u16], + s11: &[u16], + s12: &[u16], + s13: &[u16], + s14: &[u16], + s15: &[u16], + idx: &[u32], + out: &mut [u16], + seg_elems: usize, + #[comptime] num_segs: usize, + ) { + let i = ABSOLUTE_POS; + if i >= out.len() { + terminate!(); + } + out[i] = seg_read_u16( + s0, + s1, + s2, + s3, + s4, + s5, + s6, + s7, + s8, + s9, + s10, + s11, + s12, + s13, + s14, + s15, + usize::cast_from(idx[i]), + seg_elems, + num_segs, + ); + } + + /// Backend-agnostic host driver for [`enumerate_admissible_kernel`]. Lays out each `R`'s scratch + /// slot from the supplied per-`R` `num_mats` (a prefix-sum of `num_mats·cs_len` / `num_mats·mk_len`), + /// uploads the compact per-`R` inputs, launches one thread per `R` on `device`, and reads back the + /// packed `(out_cs, out_mk, counts)`. Generic over [`Runtime`] so the *same* kernel can be run on + /// CUDA (the H200 path) and on the `cpu` backend — the cross-lowering check the caller uses to + /// confirm the device semantics match [`enumerate_admissible_ref`] without needing a GPU. + fn enumerate_admissible_on_runtime( + device: &R::Device, + p_parts: &[Vec], + num_mats: &[u32], + ) -> (Vec, Vec, Vec) { + let n_r = p_parts.len(); + let width = p_parts.iter().map(Vec::len).max().unwrap(); + + let mut pp_flat = vec![0u32; n_r * width]; + let mut r_rows = vec![0u32; n_r]; + let mut r_cols = vec![0u32; n_r]; + let mut r_cs_out = vec![0u64; n_r]; + let mut r_mk_out = vec![0u64; n_r]; + let mut cs_total = 0u64; + let mut mk_total = 0u64; + for (i, pp) in p_parts.iter().enumerate() { + let rows = pp.len(); + let cols = pp + .iter() + .map(|&x| (u32::BITS - x.leading_zeros()) as usize) + .max() + .unwrap(); + let cs_len = (cols - 1) as u64; + let mk_len = (rows + cols - 1) as u64; + for (slot, &v) in pp_flat[i * width..i * width + rows].iter_mut().zip(pp) { + *slot = v; + } + r_rows[i] = rows as u32; + r_cols[i] = cols as u32; + r_cs_out[i] = cs_total; + r_mk_out[i] = mk_total; + cs_total += num_mats[i] as u64 * cs_len; + mk_total += num_mats[i] as u64 * mk_len; + } + + let client = R::client(device); + let pp_h = client.create_from_slice(u32::as_bytes(&pp_flat)); + let rr_h = client.create_from_slice(u32::as_bytes(&r_rows)); + let rc_h = client.create_from_slice(u32::as_bytes(&r_cols)); + let rco_h = client.create_from_slice(u64::as_bytes(&r_cs_out)); + let rmo_h = client.create_from_slice(u64::as_bytes(&r_mk_out)); + // `empty` needs a non-zero size even when a batch happens to have no matrices. + let cs_cap = (cs_total.max(1)) as usize; + let mk_cap = (mk_total.max(1)) as usize; + let ocs_h = client.empty(cs_cap * size_of::()); + let omk_h = client.empty(mk_cap * size_of::()); + let cnt_h = client.empty(n_r * size_of::()); + + const THREADS: u32 = ENUM_BLOCK; // must match the kernel's shared stride + let cubes = (n_r as u32).div_ceil(THREADS); + unsafe { + enumerate_admissible_kernel::launch_unchecked::( + &client, + CubeCount::Static(cubes, 1, 1), + CubeDim::new_1d(THREADS), + BufferArg::from_raw_parts(pp_h, pp_flat.len()), + BufferArg::from_raw_parts(rr_h, n_r), + BufferArg::from_raw_parts(rc_h, n_r), + BufferArg::from_raw_parts(rco_h, n_r), + BufferArg::from_raw_parts(rmo_h, n_r), + BufferArg::from_raw_parts(ocs_h.clone(), cs_cap), + BufferArg::from_raw_parts(omk_h.clone(), mk_cap), + BufferArg::from_raw_parts(cnt_h.clone(), n_r), + width, + n_r, + 0u32, + 1, + ); + } + // Truncate off the `max(1)` padding element present when a batch has zero col_sums / masks (an + // all-ones p_part gives `cs_len == 0`); the caller compares against the exact packed reference. + let mut cs = u16::from_bytes(&client.read_one(ocs_h).unwrap()).to_vec(); + let mut mk = u16::from_bytes(&client.read_one(omk_h).unwrap()).to_vec(); + cs.truncate(cs_total as usize); + mk.truncate(mk_total as usize); + let counts = u32::from_bytes(&client.read_one(cnt_h).unwrap()).to_vec(); + (cs, mk, counts) + } + + /// Diagnostic: how much of [`enumerate_admissible_kernel`]'s time is the EMIT, not the odometer? + /// + /// Two prior null results cleared the enumeration arithmetic — shrinking the per-thread local + /// state 3x was noise (1.884s -> 1.872s), and the anti-diagonal bitsum averages 1.19 iterations + /// per check (see [`masks_anti_diagonal_carries_d`]). What remains proportional to the work is + /// the emit: `cs_len + mk_len` (~20) scattered 2-byte global stores per matrix, at ~1.7e10 + /// matrices for a high-degree block. Adjacent lanes write to unrelated `r_cs_out[ri]` offsets, + /// so essentially every store is its own transaction moving 16 bits. + /// + /// Runs the identical odometer with the stores compiled out (`emit = false` is comptime, so this + /// is not a predicted branch — the stores are absent from the generated code) against the real + /// kernel. `mat` still increments, so the enumeration and its trip count are unchanged; only the + /// writes differ. The gap IS the store cost. + #[test] + #[ignore = "diagnostic: needs a CUDA device; run explicitly"] + fn enum_emit_store_cost() { + use std::time::Instant; + + use fp::prime::ValidPrime; + + let algebra = MilnorAlgebra::new(ValidPrime::new(2), false); + let max_degree = 130; + algebra.compute_basis(max_degree); + + let mut p_parts: Vec> = Vec::new(); + for deg in 1..=max_degree { + for idx in 0..algebra.dimension(deg) { + let pp: Vec = algebra + .basis_element_from_index(deg, idx) + .p_part + .iter() + .map(|x| x as u32) + .collect(); + if !pp.is_empty() { + p_parts.push(pp); + } + } + } + let num_mats: Vec = p_parts + .iter() + .map(|pp| { + let (_cs_len, _mk_len, _cs, mk) = enumerate_admissible_ref(pp); + let mk_len = pp.len() + + pp.iter() + .map(|&x| (u32::BITS - x.leading_zeros()) as usize) + .max() + .unwrap() + - 1; + (mk.len() / mk_len) as u32 + }) + .collect(); + let total_mats: u64 = num_mats.iter().map(|&m| m as u64).sum(); + + let d = Default::default(); + let (cs, mk, t_emit) = time_enum::(&d, &p_parts, &num_mats, 1); + let (_, _, t_noemit) = time_enum::(&d, &p_parts, &num_mats, 0); + let (_, _, t_coal) = time_enum::(&d, &p_parts, &num_mats, 2); + let (_, _, t_half) = time_enum::(&d, &p_parts, &num_mats, 3); + let stores = cs as u64 + mk as u64; + let share = |t: f64| 100.0 * (t - t_noemit) / (t_emit - t_noemit); + eprintln!( + " emit=2 coalesced (lane-adjacent, garbage layout) {t_coal:.4}s -> {:.1}% of the \ + emit cost remains\n emit=3 half the stores {t_half:.4}s -> {:.1}% remains", + share(t_coal), + share(t_half) + ); + eprintln!( + "enum emit cost (degree <= {max_degree}, {} R's, {total_mats} matrices, {stores} u16 \ + stores = {:.2} GB):\n emit=true {:.4}s\n emit=false {:.4}s\n stores are {:.1}% \ + of kernel time ({:.2}x speedup if free)", + p_parts.len(), + stores as f64 * 2.0 / 1e9, + t_emit, + t_noemit, + 100.0 * (t_emit - t_noemit) / t_emit, + t_emit / t_noemit, + ); + } + + /// Launch [`enumerate_admissible_kernel`] with `emit` on or off and return + /// `(cs_total, mk_total, seconds)` — the median of several timed launches, sync'd by a small + /// readback so the measurement covers the kernel rather than the submission. + fn time_enum( + device: &R::Device, + p_parts: &[Vec], + num_mats: &[u32], + emit: u32, + ) -> (u64, u64, f64) { + use std::time::Instant; + + let n_r = p_parts.len(); + let width = p_parts.iter().map(Vec::len).max().unwrap(); + let mut pp_flat = vec![0u32; n_r * width]; + let mut r_rows = vec![0u32; n_r]; + let mut r_cols = vec![0u32; n_r]; + let mut r_cs_out = vec![0u64; n_r]; + let mut r_mk_out = vec![0u64; n_r]; + let (mut cs_total, mut mk_total) = (0u64, 0u64); + for (i, pp) in p_parts.iter().enumerate() { + let rows = pp.len(); + let cols = pp + .iter() + .map(|&x| (u32::BITS - x.leading_zeros()) as usize) + .max() + .unwrap(); + for (slot, &v) in pp_flat[i * width..i * width + rows].iter_mut().zip(pp) { + *slot = v; + } + r_rows[i] = rows as u32; + r_cols[i] = cols as u32; + r_cs_out[i] = cs_total; + r_mk_out[i] = mk_total; + cs_total += num_mats[i] as u64 * (cols - 1) as u64; + mk_total += num_mats[i] as u64 * (rows + cols - 1) as u64; + } + + let client = R::client(device); + let pp_h = client.create_from_slice(u32::as_bytes(&pp_flat)); + let rr_h = client.create_from_slice(u32::as_bytes(&r_rows)); + let rc_h = client.create_from_slice(u32::as_bytes(&r_cols)); + let rco_h = client.create_from_slice(u64::as_bytes(&r_cs_out)); + let rmo_h = client.create_from_slice(u64::as_bytes(&r_mk_out)); + let cs_cap = cs_total.max(1) as usize; + let mk_cap = mk_total.max(1) as usize; + let ocs_h = client.empty(cs_cap * size_of::()); + let omk_h = client.empty(mk_cap * size_of::()); + let cnt_h = client.empty(n_r * size_of::()); + + // Largest power of two that keeps `emit == 2`'s rewritten indices inside both buffers. + let wrap = { + let lim = cs_cap.min(mk_cap) as u64; + (1u64 << (63 - lim.leading_zeros().min(62))).min(lim) as u32 - 1 + }; + + const THREADS: u32 = ENUM_BLOCK; // must match the kernel's shared stride + let cubes = (n_r as u32).div_ceil(THREADS); + let mut times: Vec = Vec::new(); + for _ in 0..5 { + let t0 = Instant::now(); + unsafe { + enumerate_admissible_kernel::launch_unchecked::( + &client, CubeCount::Static(cubes, 1, 1), CubeDim::new_1d(THREADS), - ArrayArg::from_raw_parts(cs_h.clone(), cs_len_master), - ArrayArg::from_raw_parts(mk_h.clone(), mk_len_master), - ArrayArg::from_raw_parts(tp_h.clone(), term_pparts.len()), - ArrayArg::from_raw_parts(tl_h.clone(), term_lens.len()), - ArrayArg::from_raw_parts(g_h.clone(), g.len()), - ArrayArg::from_raw_parts(xi_h.clone(), xi.len()), - ArrayArg::from_raw_parts(out_h.clone(), out_len), - ArrayArg::from_raw_parts(rco_h.clone(), r_cs_offset.len()), - ArrayArg::from_raw_parts(rmo_h.clone(), r_mk_offset.len()), - ArrayArg::from_raw_parts(rcl_h.clone(), r_cs_len.len()), - ArrayArg::from_raw_parts(rml_h.clone(), r_mk_len.len()), - ArrayArg::from_raw_parts(pri_h, c1 - c0), - ArrayArg::from_raw_parts(pts_h, c1 - c0), - ArrayArg::from_raw_parts(pnt_h, c1 - c0), - ArrayArg::from_raw_parts(prb_h, c1 - c0), - ArrayArg::from_raw_parts(poo_h, c1 - c0), - ArrayArg::from_raw_parts(pps_h, pps_chunk.len()), + BufferArg::from_raw_parts(pp_h.clone(), pp_flat.len()), + BufferArg::from_raw_parts(rr_h.clone(), n_r), + BufferArg::from_raw_parts(rc_h.clone(), n_r), + BufferArg::from_raw_parts(rco_h.clone(), n_r), + BufferArg::from_raw_parts(rmo_h.clone(), n_r), + BufferArg::from_raw_parts(ocs_h.clone(), cs_cap), + BufferArg::from_raw_parts(omk_h.clone(), mk_cap), + BufferArg::from_raw_parts(cnt_h.clone(), n_r), width, + n_r, + wrap, + emit, + ); + } + let _ = client.read_one(cnt_h.clone()).unwrap(); + times.push(t0.elapsed().as_secs_f64()); + } + times.sort_by(|a, b| a.partial_cmp(b).unwrap()); + (cs_total, mk_total, times[times.len() / 2]) + } + + /// CPU reference for the planned *in-kernel* admissible-matrix enumeration — the direction that + /// replaces the resident/uploaded master (the stem-300 memory wall + the eviction re-upload cost) + /// by generating each `R`'s `col_sums`/`masks` ON THE GPU into a transient scratch buffer, never + /// storing or uploading them. This reimplements [`MilnorAlgebra::admissible_matrices`] / + /// `AdmissibleMatrix` using ONLY flag-guarded control flow — no `break`, `continue`, or early + /// `return` — because that is the subset the cubecl DSL compiles cleanly (cf. `multiply_pair`, + /// which tracks `rejected` rather than breaking). The eventual `#[cube]` kernel is then a mechanical + /// transcription of this function onto per-thread local `Array`s (state is tiny: `rows = |p_part|`, + /// `cols ≤ 32`). Returns the same `(cs_len, mk_len, col_sums, masks)` row-major flattening as + /// `admissible_matrices`; `admissible_enum_ref_matches` asserts bit-exact equivalence over every + /// real `R` up to degree 60, validating the flag-based restructuring before the hard-to-debug port. + fn enumerate_admissible_ref(p_part: &[u32]) -> (usize, usize, Vec, Vec) { + let rows = p_part.len(); + let cols = p_part + .iter() + .map(|&x| (u32::BITS - x.leading_zeros()) as usize) + .max() + .unwrap(); + let cs_len = cols - 1; + let mk_len = rows + cols - 1; + + // State mirrors `AdmissibleMatrix`: `matrix` row-major `rows*cols` (column 0 = `p_part`), + // `totals[rows]`, `col_sums[cs_len]`, `masks[mk_len]` (masks starts as the padded `p_part`). + let mut matrix = vec![0u32; rows * cols]; + for (i, &x) in p_part.iter().enumerate() { + matrix[i * cols] = x; + } + let mut totals = vec![0u32; rows]; + let mut col_sums = vec![0u32; cs_len]; + let mut masks = vec![0u32; mk_len]; + for (i, &x) in p_part.iter().enumerate() { + masks[i] = x; + } + + let mut out_cs: Vec = Vec::new(); + let mut out_mk: Vec = Vec::new(); + + // Emit the current matrix, then advance; `more` is `AdmissibleMatrix::next`'s return value. + let mut more = true; + while more { + out_cs.extend_from_slice(&col_sums); + out_mk.extend_from_slice(&masks); + + // One `next()` step, flag-based: `found` = "produced a new matrix" (the original's + // `return true`); `handled` = "this column already updated `totals`" (the original's + // `continue 'mid`, which skips the trailing add). Loops are guarded by `!found` instead + // of breaking. + let mut found = false; + let mut row = 0; + while row < rows && !found { + let mut p_to_the_j: u32 = 1; + totals[row] = matrix[row * cols]; // get(row, 0) + let mut col = 1; + while col < cols && !found { + p_to_the_j *= 2; + let mut handled = false; + if p_to_the_j <= totals[row] { + // Bitsum along the anti-diagonal to the bottom-left. + let mut d = 0u32; + let mut c = (row + col + 1).saturating_sub(rows); + while c < col { + d |= matrix[(row + col - c) * cols + c]; + c += 1; + } + let cur = matrix[row * cols + col]; + let new_entry = ((cur | d) + 1) & !d; + let inc = new_entry - cur; + let sub = inc * p_to_the_j; + if totals[row] < sub { + totals[row] += p_to_the_j * cur; + handled = true; + } else { + matrix[row * cols] = totals[row] - sub; // set(row, 0, ..) + masks[row] = matrix[row * cols]; + col_sums[col - 1] += inc; + let mut j = 1; + while j < col { + masks[row + j] &= !matrix[row * cols + j]; + col_sums[j - 1] -= matrix[row * cols + j]; + matrix[row * cols + j] = 0; + j += 1; + } + matrix[row * cols + col] = new_entry; + let mut i = 0; + while i < row { + matrix[i * cols] = totals[i]; + masks[i] = totals[i]; + let mut j = 1; + while j < cols { + if i + j > row { + masks[i + j] &= !matrix[i * cols + j]; + } + col_sums[j - 1] -= matrix[i * cols + j]; + matrix[i * cols + j] = 0; + j += 1; + } + i += 1; + } + masks[row + col] = d | new_entry; + found = true; + handled = true; + } + } + if !handled { + totals[row] += p_to_the_j * matrix[row * cols + col]; + } + col += 1; + } + row += 1; + } + more = found; + } + + (cs_len, mk_len, out_cs, out_mk) + } + + /// Host driver for [`seg_gather_kernel`]: splits `data` into ≤ [`MASTER_MAX_SEG`] segments of + /// `seg_elems`, uploads each as its own device buffer (no contiguous copy — the whole point), and + /// returns `data[idx[i]]` gathered through the segmented read. Unused segments get a 1-element dummy + /// (never indexed). Proves the segmented master reads identically to a contiguous one. + fn seg_gather_on_gpu(data: &[u16], seg_elems: usize, indices: &[u32]) -> Vec { + let n = data.len(); + let nseg = n.div_ceil(seg_elems).max(1); + assert!( + nseg <= MASTER_MAX_SEG, + "prototype caps at {MASTER_MAX_SEG} segments" + ); + let client = gpu_client(); + + // One handle per segment slot; real segments hold their slice, unused slots a 1-elem dummy. + let dummy = [0u16]; + let mut handles = Vec::with_capacity(MASTER_MAX_SEG); + let mut lens = Vec::with_capacity(MASTER_MAX_SEG); + for s in 0..MASTER_MAX_SEG { + let lo = s * seg_elems; + if lo < n { + let hi = (lo + seg_elems).min(n); + handles.push(client.create_from_slice(u16::as_bytes(&data[lo..hi]))); + lens.push(hi - lo); + } else { + handles.push(client.create_from_slice(u16::as_bytes(&dummy))); + lens.push(1); + } + } + let idx_h = client.create_from_slice(u32::as_bytes(indices)); + let out_h = client.empty(indices.len() * size_of::()); + + const THREADS: u32 = 256; + let cubes = (indices.len() as u32).div_ceil(THREADS).max(1); + let arg = |i: usize| unsafe { BufferArg::from_raw_parts(handles[i].clone(), lens[i]) }; + unsafe { + seg_gather_kernel::launch::( + &client, + CubeCount::Static(cubes, 1, 1), + CubeDim::new_1d(THREADS), + arg(0), + arg(1), + arg(2), + arg(3), + arg(4), + arg(5), + arg(6), + arg(7), + arg(8), + arg(9), + arg(10), + arg(11), + arg(12), + arg(13), + arg(14), + arg(15), + BufferArg::from_raw_parts(idx_h, indices.len()), + BufferArg::from_raw_parts(out_h.clone(), indices.len()), + seg_elems, + nseg, + ); + } + u16::from_bytes(&client.read_one(out_h).unwrap()).to_vec() + } + + /// Shared body for the per-backend enumeration tests: builds a batch of every real `R` up to + /// `max_degree`, computes the expected packed `col_sums`/`masks` (and per-`R` `num_mats`) from the + /// CPU-validated [`enumerate_admissible_ref`], runs [`enumerate_admissible_kernel`] on `R`'s + /// `device`, and asserts the device output is bit-exact — values *and* per-`R` counts. Generic so + /// CUDA (H200) and the `cpu` backend run the identical kernel through it. + fn check_enum_backend(device: &Rt::Device, max_degree: i32) { + use fp::prime::ValidPrime; + + let p = ValidPrime::new(2); + let algebra = Arc::new(MilnorAlgebra::new(p, false)); + algebra.compute_basis(max_degree); + + // Process one degree per launch. At high degree the full master is tens of GB, so batching every + // R together would OOM the host; per-degree keeps the expected arrays bounded AND pinpoints the + // exact degree if the device lowering ever diverges from the CPU reference. + let mut total_r = 0usize; + let mut total_mats = 0u64; + for deg in 1..=max_degree { + let mut p_parts: Vec> = Vec::new(); + let mut num_mats: Vec = Vec::new(); + let mut exp_cs: Vec = Vec::new(); + let mut exp_mk: Vec = Vec::new(); + for idx in 0..algebra.dimension(deg) { + let pp: Vec = algebra + .basis_element_from_index(deg, idx) + .p_part + .iter() + .collect(); + if pp.is_empty() { + continue; + } + let (_cs_len, mk_len, cs, mk) = enumerate_admissible_ref(&pp); + // `mk_len = rows+cols-1 ≥ 1` always, so it recovers the matrix count even when + // `cs_len == 0` (an all-ones p_part contributes no col_sums). + num_mats.push((mk.len() / mk_len) as u32); + exp_cs.extend(cs.iter().map(|&v| narrow_u16(v))); + exp_mk.extend(mk.iter().map(|&v| narrow_u16(v))); + p_parts.push(pp); + } + if p_parts.is_empty() { + continue; + } + let (got_cs, got_mk, counts) = + enumerate_admissible_on_runtime::(device, &p_parts, &num_mats); + assert_eq!( + counts, num_mats, + "per-R matrix counts diverged at degree {deg}" + ); + assert_eq!(got_cs, exp_cs, "device col_sums diverged at degree {deg}"); + assert_eq!(got_mk, exp_mk, "device masks diverged at degree {deg}"); + total_r += p_parts.len(); + total_mats += num_mats.iter().map(|&m| m as u64).sum::(); + } + assert!(total_r > 0, "no R's exercised"); + eprintln!( + "enum backend: {total_r} R's, {total_mats} matrices bit-exact vs \ + enumerate_admissible_ref (degrees 1..={max_degree})" + ); + } + + /// Time one enumeration launch on `device`, split into (marshal+upload, kernel, full readback). + /// `kernel` reads only the tiny `counts` buffer to force a stream sync (so it captures kernel wall + /// time without the big transfer); `readback` then pulls the full `col_sums`/`masks`. Used by + /// `bench_admissible_cpu_vs_gpu` — production never reads the arrays back (the multiply consumes the + /// scratch on-device), so `kernel` is the production-relevant cost and `readback` is bench-only. + fn enum_launch_timed( + device: &R::Device, + p_parts: &[Vec], + num_mats: &[u32], + ) -> (f64, f64, f64) { + use std::time::Instant; + let n_r = p_parts.len(); + let width = p_parts.iter().map(Vec::len).max().unwrap(); + + let t_marshal = Instant::now(); + let mut pp_flat = vec![0u32; n_r * width]; + let mut r_rows = vec![0u32; n_r]; + let mut r_cols = vec![0u32; n_r]; + let mut r_cs_out = vec![0u64; n_r]; + let mut r_mk_out = vec![0u64; n_r]; + let (mut cs_total, mut mk_total) = (0u64, 0u64); + for (i, pp) in p_parts.iter().enumerate() { + let rows = pp.len(); + let cols = pp + .iter() + .map(|&x| (u32::BITS - x.leading_zeros()) as usize) + .max() + .unwrap(); + for (slot, &v) in pp_flat[i * width..i * width + rows].iter_mut().zip(pp) { + *slot = v; + } + r_rows[i] = rows as u32; + r_cols[i] = cols as u32; + r_cs_out[i] = cs_total; + r_mk_out[i] = mk_total; + cs_total += num_mats[i] as u64 * (cols - 1) as u64; + mk_total += num_mats[i] as u64 * (rows + cols - 1) as u64; + } + let client = R::client(device); + let pp_h = client.create_from_slice(u32::as_bytes(&pp_flat)); + let rr_h = client.create_from_slice(u32::as_bytes(&r_rows)); + let rc_h = client.create_from_slice(u32::as_bytes(&r_cols)); + let rco_h = client.create_from_slice(u64::as_bytes(&r_cs_out)); + let rmo_h = client.create_from_slice(u64::as_bytes(&r_mk_out)); + let cs_cap = cs_total.max(1) as usize; + let mk_cap = mk_total.max(1) as usize; + let ocs_h = client.empty(cs_cap * size_of::()); + let omk_h = client.empty(mk_cap * size_of::()); + let cnt_h = client.empty(n_r * size_of::()); + let marshal_s = t_marshal.elapsed().as_secs_f64(); + + const THREADS: u32 = ENUM_BLOCK; // must match the kernel's shared stride + let cubes = (n_r as u32).div_ceil(THREADS); + let t_kernel = Instant::now(); + unsafe { + enumerate_admissible_kernel::launch_unchecked::( + &client, + CubeCount::Static(cubes, 1, 1), + CubeDim::new_1d(THREADS), + BufferArg::from_raw_parts(pp_h, pp_flat.len()), + BufferArg::from_raw_parts(rr_h, n_r), + BufferArg::from_raw_parts(rc_h, n_r), + BufferArg::from_raw_parts(rco_h, n_r), + BufferArg::from_raw_parts(rmo_h, n_r), + BufferArg::from_raw_parts(ocs_h.clone(), cs_cap), + BufferArg::from_raw_parts(omk_h.clone(), mk_cap), + BufferArg::from_raw_parts(cnt_h.clone(), n_r), + width, + n_r, + 0u32, + 1, + ); + } + // Reading the tiny counts buffer blocks until the kernel completes: kernel wall time, ~no transfer. + let _ = client.read_one(cnt_h).unwrap(); + let kernel_s = t_kernel.elapsed().as_secs_f64(); + + let t_read = Instant::now(); + let _ = client.read_one(ocs_h).unwrap(); + let _ = client.read_one(omk_h).unwrap(); + let readback_s = t_read.elapsed().as_secs_f64(); + + (marshal_s, kernel_s, readback_s) + } + + /// How much of the enumeration kernel's throughput the one-thread-per-`R` mapping actually gets. + /// + /// The kernel walks an odometer, so a thread's cost is proportional to its `R`'s matrix count, and + /// a warp cannot retire until its slowest lane does — warp cost is `32 x max`, not `sum`. The + /// `NASSAU_R_STATS` Lorenz curve says matrix counts are extremely skewed (the hottest 1% of `R`s + /// carry 31% of all references), so if `R`s land in warps in basis order, a single huge `R` can + /// idle 31 lanes for its whole run. + /// + /// Reports achieved utilisation against the same work sorted by matrix count, which is the win + /// available from a host-side reorder before launch. CPU-only. + /// + /// TREAT THE RATIO AS AN UPPER BOUND, NOT A FORECAST: it models a warp in isolation, but an SM + /// keeps many warps resident and runs others while a lane-heavy warp grinds, so most of the + /// modelled stall is hidden. The 2.16x this reports at degree 130 cashed out as 1.08x measured + /// end-to-end (`bench_admissible_cpu_vs_gpu`). Sorting is still worth it — a host-side sort of a + /// few thousand keys per block is free next to the launch — but do not size a decision on the + /// model without measuring. + #[test] + #[ignore = "diagnostic, not a correctness check; run explicitly with --ignored --nocapture"] + fn enum_warp_utilisation() { + use fp::prime::ValidPrime; + + let algebra = MilnorAlgebra::new(ValidPrime::new(2), false); + let max_degree = 130; + algebra.compute_basis(max_degree); + + // Warp cost model: lanes run in lockstep, so a warp costs 32 x its largest lane. + let warp_cost = + |mats: &[u64]| -> u64 { mats.chunks(32).map(|w| 32 * w.iter().max().unwrap()).sum() }; + + let (mut tot_work, mut tot_natural, mut tot_sorted) = (0u64, 0u64, 0u64); + for deg in 1..=max_degree { + let mut mats: Vec = Vec::new(); + for idx in 0..algebra.dimension(deg) { + let pp: Vec = algebra + .basis_element_from_index(deg, idx) + .p_part + .iter() + .collect(); + if pp.is_empty() { + continue; + } + let (_cs_len, mk_len, _cs, mk) = enumerate_admissible_ref(&pp); + mats.push((mk.len() / mk_len) as u64); + } + if mats.is_empty() { + continue; + } + tot_work += mats.iter().sum::(); + tot_natural += warp_cost(&mats); + let mut sorted = mats.clone(); + sorted.sort_unstable(); + tot_sorted += warp_cost(&sorted); + } + + let pct = |c: u64| 100.0 * tot_work as f64 / c as f64; + eprintln!( + "enum warp utilisation to degree {max_degree}: work {tot_work} matrices\n basis \ + order : {tot_natural} lane-slots ({:.1}% utilised)\n sorted : {tot_sorted} \ + lane-slots ({:.1}% utilised) -> {:.2}x headroom", + pct(tot_natural), + pct(tot_sorted), + tot_natural as f64 / tot_sorted as f64, + ); + } + + /// [`ENUM_COL_CAP`] / [`ENUM_ROW_CAP`] bound fixed-size local arrays the enumeration kernel + /// indexes with runtime values, so an `R` exceeding either would write out of bounds — a silent + /// `CUDA_ERROR_LAUNCH_FAILED` at some high stem, not a clean failure. The caps are derived from + /// `PPart`'s layout rather than measured, so this checks the derivation against reality: no real + /// `R` may exceed them. + /// + /// The caps stay deliberately loose: at degree 400 the real maxima are 8 rows and 9 cols, but + /// `cols` is bounded by the `r_1` field width (11) and only approaches it near `PPart::MAX_DEGREE`. + /// Tightening to the observed 9 would trade a further 18% of `matrix` for a cap that silently + /// breaks at a degree nobody is watching, so the structural bound is the one worth encoding. + /// + /// CPU-only, so it costs nothing to run in CI alongside the GPU tests. + #[test] + fn enum_col_cap_bounds_real_rs() { + use fp::prime::ValidPrime; + + let algebra = MilnorAlgebra::new(ValidPrime::new(2), false); + let max_degree = 400; + algebra.compute_basis(max_degree); + + let (mut max_rows, mut max_cols) = (0usize, 0usize); + for deg in 1..=max_degree { + for idx in 0..algebra.dimension(deg) { + let pp = algebra.basis_element_from_index(deg, idx).p_part; + if pp.is_empty() { + continue; + } + max_rows = max_rows.max(pp.len()); + max_cols = max_cols.max( + pp.iter() + .map(|x| (u32::BITS - x.leading_zeros()) as usize) + .max() + .unwrap(), + ); + } + } + assert!( + max_rows <= ENUM_ROW_CAP, + "an R has {max_rows} rows > ENUM_ROW_CAP {ENUM_ROW_CAP}" + ); + assert!( + max_cols <= ENUM_COL_CAP, + "an R has {max_cols} cols > ENUM_COL_CAP {ENUM_COL_CAP}" + ); + eprintln!( + "enum caps to degree {max_degree}: max rows {max_rows}/{ENUM_ROW_CAP}, max cols \ + {max_cols}/{ENUM_COL_CAP} (matrix {ENUM_MATRIX_CAP} u32/thread)" + ); + } + + /// Diagnostic: is the anti-diagonal bitsum `d` already available in `masks`? + /// + /// Both odometers (`AdmissibleMatrix::next` and [`enumerate_admissible_kernel`]) recompute + /// + /// ```text + /// for c in (row+col+1).saturating_sub(rows)..col { d |= matrix[(row+col-c)*cols + c] } + /// ``` + /// + /// on every visited `(row, col)`, inside the hottest triple-nested loop. But the accept path + /// stores `masks[row+col] = d | new_entry` and the clears do `masks[i+j] &= !matrix[i*cols+j]`, + /// so `masks` is itself an accumulator over the anti-diagonal `row+col`. If it coincided with + /// `d` at read time, an O(cols) scan would collapse to one load in BOTH implementations. + /// + /// VERDICT (measured, degree <= 120, 46,995,344 anti-diagonal computations): the shortcut is + /// NOT worth taking, and the reason is not the one the conjecture is about. + /// + /// * `d == masks` only 28.3% of the time — `masks[row+col]` folds in the cell AT column `col` + /// (`masks[row+col] = d | new_entry`) while `d` stops short of it, so they differ whenever + /// that cell is set. + /// * `d` SUBSET-OF `masks` in 100.000% of cases (0 violations, assert armed below). So + /// `masks[row+col] == 0` soundly implies `d == 0` and the scan can be skipped outright. + /// * But that skip fires on 5.8% of checks and avoids 6.9% of scan iterations, and the scan + /// averages **1.19 iterations per check** (55,846,521 over 46,995,344) — the range + /// `(row+col+1).saturating_sub(rows)..col` is nearly always empty or one step for real R's. + /// + /// So the "O(cols) bitsum in the hottest loop" is not O(cols) in practice and not hot. Do not + /// re-propose caching, precomputing, or incrementally maintaining `d`; the arithmetic is not + /// where enumeration spends its time. Together with the earlier null result from shrinking the + /// per-thread local state 3x (1.884s -> 1.872s, noise), that leaves the kernel's ~20 scattered + /// 2-byte output stores per matrix as the remaining candidate. + /// + /// This does not assert the conjecture — it measures it, reporting the containment direction so + /// a mismatch says what `masks` is missing rather than merely that it differs. Ignored by + /// default: it is an investigation tool, not a regression gate. + #[test] + #[ignore = "diagnostic: run explicitly to measure the d-vs-masks relationship"] + fn masks_anti_diagonal_carries_d() { + use fp::prime::ValidPrime; + + let algebra = MilnorAlgebra::new(ValidPrime::new(2), false); + let max_degree = 120; + algebra.compute_basis(max_degree); + + let (mut checks, mut equal, mut d_subset, mut masks_subset) = (0u64, 0u64, 0u64, 0u64); + let (mut scan_iters, mut saved_iters, mut masks_zero) = (0u64, 0u64, 0u64); + let mut sample: Vec = Vec::new(); + + for deg in 1..=max_degree { + for idx in 0..algebra.dimension(deg) { + let pp = &algebra.basis_element_from_index(deg, idx).p_part; + if pp.is_empty() { + continue; + } + let p_part: Vec = pp.iter().map(|x| x as u32).collect(); + let rows = p_part.len(); + let cols = p_part + .iter() + .map(|&x| (u32::BITS - x.leading_zeros()) as usize) + .max() + .unwrap(); + if cols < 2 { + continue; + } + + let mut matrix = vec![0u32; rows * cols]; + let mut masks = vec![0u32; rows + cols - 1]; + for (i, &x) in p_part.iter().enumerate() { + matrix[i * cols] = x; + masks[i] = x; + } + let mut totals = vec![0u32; rows]; + let mut col_sums = vec![0u32; cols - 1]; + + let mut more = true; + while more { + let mut found = false; + let mut row = 0; + while row < rows && !found { + let mut p_to_the_j: u32 = 1; + totals[row] = matrix[row * cols]; + let mut col = 1; + while col < cols && !found { + p_to_the_j *= 2; + let mut handled = false; + if p_to_the_j <= totals[row] { + let mut d = 0u32; + let mut c = (row + col + 1).saturating_sub(rows); + while c < col { + d |= matrix[(row + col - c) * cols + c]; + c += 1; + } + + // The measurement: compare against the maintained accumulator. + let m = masks[row + col]; + checks += 1; + scan_iters += (col - (row + col + 1).saturating_sub(rows)) as u64; + if m == 0 { + masks_zero += 1; + saved_iters += + (col - (row + col + 1).saturating_sub(rows)) as u64; + assert_eq!(d, 0, "masks==0 but d!=0 — containment violated"); + } + if d == m { + equal += 1; + } + if d | m == m { + d_subset += 1; + } + if d | m == d { + masks_subset += 1; + } + if d != m && sample.len() < 8 { + sample.push(format!( + "R={p_part:?} rows={rows} cols={cols} \ + (row={row},col={col}) d={d:#x} masks[{}]={m:#x}", + row + col + )); + } + + let cur = matrix[row * cols + col]; + let new_entry = ((cur | d) + 1) & !d; + let inc = new_entry - cur; + let sub = inc * p_to_the_j; + if totals[row] < sub { + totals[row] += p_to_the_j * cur; + handled = true; + } else { + matrix[row * cols] = totals[row] - sub; + masks[row] = matrix[row * cols]; + col_sums[col - 1] += inc; + let mut j = 1; + while j < col { + masks[row + j] &= !matrix[row * cols + j]; + col_sums[j - 1] -= matrix[row * cols + j]; + matrix[row * cols + j] = 0; + j += 1; + } + matrix[row * cols + col] = new_entry; + let mut i = 0; + while i < row { + matrix[i * cols] = totals[i]; + masks[i] = totals[i]; + let mut j = 1; + while j < cols { + if i + j > row { + masks[i + j] &= !matrix[i * cols + j]; + } + col_sums[j - 1] -= matrix[i * cols + j]; + matrix[i * cols + j] = 0; + j += 1; + } + i += 1; + } + masks[row + col] = d | new_entry; + found = true; + handled = true; + } + } + if !handled { + totals[row] += p_to_the_j * matrix[row * cols + col]; + } + col += 1; + } + row += 1; + } + more = found; + } + } + } + + let pct = |n: u64| 100.0 * n as f64 / checks.max(1) as f64; + eprintln!( + "d-vs-masks over {checks} anti-diagonal computations (degree <= {max_degree}):\n d \ + == masks : {equal} ({:.3}%)\n d subset-of masks : {d_subset} ({:.3}%)\n masks \ + subset-of d : {masks_subset} ({:.3}%)", + pct(equal), + pct(d_subset), + pct(masks_subset) + ); + eprintln!( + " masks==0 (sound skip): {masks_zero} ({:.3}%)\n d-scan iterations: {scan_iters}, \ + avoidable by the skip: {saved_iters} ({:.3}%)", + pct(masks_zero), + 100.0 * saved_iters as f64 / scan_iters.max(1) as f64 + ); + for s in &sample { + eprintln!(" mismatch: {s}"); + } + } + + /// The in-kernel [`enumerate_admissible_kernel`], run on the CUDA backend, must reproduce the + /// CPU-validated [`enumerate_admissible_ref`] bit-for-bit over every real `R` up to degree 145 — + /// validating the cubecl lowering of the flag-based enumeration (local arrays, bitops, u16 + /// stores) across the FULL degree range the eviction path exercises (cold R's reach ~144 at + /// stem 150), not just the low degrees. Requires a live GPU + the CUDA toolkit env. + /// + /// 145 is a runtime compromise, not a known ceiling: the degree-240 sweep (1,593,460 `R`s, + /// 16,949,543,206 matrices) also passes bit-exact, but takes 942 s against a few seconds here, + /// and the cost is in the CPU reference, so it grows steeply. Raise the bound to re-check after + /// touching the kernel — a `LAUNCH_FAILED` at high stem is NOT evidence against this kernel. + #[test] + fn admissible_enum_gpu_matches() { + check_enum_backend::(&CudaDevice::default(), 145); + } + + /// The segmented master read ([`seg_read_u16`]) must reproduce a contiguous buffer bit-for-bit, + /// including reads that land in every one of the [`MASTER_MAX_SEG`] segments and at segment + /// boundaries. This validates the no-copy-growth mechanic before it is wired into the multiply's + /// hot path. Requires a live GPU + the CUDA toolkit env. + #[test] + fn seg_read_matches_contiguous() { + // 1000 elements over seg_elems=137 → 8 segments (0..137, 137..274, …, 959..1000), so every + // segment slot is exercised, including the ragged last one and the boundaries between them. + let data: Vec = (0..1000u16).collect(); + let seg_elems = 137usize; + // Gather in a scrambled order so a segment-selection bug can't hide behind sequential access. + let indices: Vec = (0..1000u32).map(|i| (i * 613) % 1000).collect(); + let got = seg_gather_on_gpu(&data, seg_elems, &indices); + let want: Vec = indices.iter().map(|&i| data[i as usize]).collect(); + assert_eq!( + got, want, + "segmented read diverged from contiguous indexing" + ); + } + + /// Throughput comparison, CPU `admissible_matrices` vs the in-kernel [`enumerate_admissible_kernel`], + /// for enumerating every `R`'s admissible matrices up to a degree. Reports GPU kernel-only time + /// (the production-relevant cost — the multiply consumes the scratch on-device, no readback) and + /// the full-readback time separately. Run with `--nocapture --ignored`; needs a live GPU. + #[test] + #[ignore = "benchmark, not a correctness check; run explicitly with --ignored --nocapture"] + fn bench_admissible_cpu_vs_gpu() { + use std::time::Instant; + + use fp::prime::ValidPrime; + + let p = ValidPrime::new(2); + let algebra = Arc::new(MilnorAlgebra::new(p, false)); + let max_degree = 130; + algebra.compute_basis(max_degree); + + // Gather every non-empty R, grouped by degree (per-degree GPU launches keep host arrays bounded). + // Per degree keep the packed u16 col_sums/masks (what upload-based eviction uploads H->D), + // so we can time that upload directly and compare it against on-device enumeration. Both are + // production paths: the matrices end up on-device either way, so NEITHER counts a readback. + let mut by_degree: Vec<(Vec>, Vec, Vec, Vec)> = Vec::new(); + let mut cpu_secs = 0.0f64; + let mut total_r = 0usize; + let mut total_mats = 0u64; + for deg in 1..=max_degree { + let (mut pps, mut nms, mut cs_all, mut mk_all) = + (Vec::new(), Vec::new(), Vec::new(), Vec::new()); + for idx in 0..algebra.dimension(deg) { + let pp: Vec = algebra + .basis_element_from_index(deg, idx) + .p_part + .iter() + .collect(); + if pp.is_empty() { + continue; + } + // Time the CPU enumeration (`admissible_matrices`, the call the CPU multiply makes). + let t = Instant::now(); + let (_cs_len, mk_len, cs, mk) = + algebra.admissible_matrices(PPart::try_from_slice(&pp).unwrap()); + cpu_secs += t.elapsed().as_secs_f64(); + let nm = (mk.len() / mk_len) as u32; + nms.push(nm); + total_mats += nm as u64; + cs_all.extend(cs.iter().map(|&v| narrow_u16(v))); + mk_all.extend(mk.iter().map(|&v| narrow_u16(v))); + pps.push(pp); + } + total_r += pps.len(); + if !pps.is_empty() { + by_degree.push((pps, nms, cs_all, mk_all)); + } + } + + let device = CudaDevice::default(); + let client = CudaRuntime::client(&device); + // Warm up the runtime/JIT so the first degree's compile doesn't skew the GPU timing. + { + let (pps, nms, _, _) = &by_degree[0]; + let _ = enum_launch_timed::(&device, pps, nms); + } + let mut g_kernel = 0.0f64; + let mut g_kernel_sorted = 0.0f64; + let mut upload_secs = 0.0f64; + for (pps, nms, cs_all, mk_all) in &by_degree { + // (a) On-device enumeration, kernel only (no readback — the multiply consumes the scratch), + // in basis order and in matrix-count order. Production sorts (see the `enum_pp` marshal); + // the unsorted timing is kept as the baseline that motivates it. + let (_m, k, _r) = enum_launch_timed::(&device, pps, nms); + g_kernel += k; + let mut order: Vec = (0..pps.len()).collect(); + order.sort_unstable_by_key(|&i| nms[i]); + let pps_s: Vec> = order.iter().map(|&i| pps[i].clone()).collect(); + let nms_s: Vec = order.iter().map(|&i| nms[i]).collect(); + let (_m, ks, _r) = enum_launch_timed::(&device, &pps_s, &nms_s); + g_kernel_sorted += ks; + // (b) What upload-based eviction does instead: upload the host-built arrays H->D. Force the + // (possibly async) copies to complete by syncing the stream via a tiny throwaway readback + // (4 bytes back, negligible) — NOT by reading the big arrays back, so this is upload-only. + let t = Instant::now(); + let ch = client.create_from_slice(u16::as_bytes(cs_all)); + let mh = client.create_from_slice(u16::as_bytes(mk_all)); + let sync = client.empty(size_of::()); + let _ = client.read_one(sync).unwrap(); + upload_secs += t.elapsed().as_secs_f64(); + drop((ch, mh)); + } + + eprintln!( + "\n=== admissible matrices onto the device: enumerate vs upload (degrees \ + 1..={max_degree}) ===\nR's: {total_r} matrices: {total_mats} (both paths leave \ + the arrays ON-DEVICE, no readback)\nCPU admissible_matrices (enumerate, 1 core) : \ + {cpu_secs:.3} s\nGPU enumerate in-kernel (basis order) : {g_kernel:.3} s\nGPU \ + enumerate in-kernel (matrix-count order): {g_kernel_sorted:.3} s ({:.2}x)\nH->D \ + upload of host-built arrays : {upload_secs:.3} s\n--> in-kernel enum is \ + {:.2}x the cost of just uploading the same arrays", + g_kernel / g_kernel_sorted, + g_kernel_sorted / upload_secs, + ); + } + + /// [`shard_of`] must split the real `R`s evenly across devices, and must stay even when they are + /// grouped by degree — because assignment is permanent and the frontier walks degree upward, a + /// hash that is uniform overall but skewed within a degree band would still starve devices for + /// long stretches, which is the failure this replaced. + /// + /// Also asserts the unmixed key is NOT usable, so nobody "simplifies" `shard_of` back to a bare + /// modulo: [`PPart`] packs `r_1` in the low bits and `r_1` tracks degree, so `bits() % 4` + /// partitions by degree — precisely the correlation that produced a 3.3x device imbalance. + /// Pure CPU: no GPU needed. + #[test] + fn shard_of_is_uniform_over_real_rs() { + use fp::prime::ValidPrime; + + let algebra = MilnorAlgebra::new(ValidPrime::new(2), false); + let max_degree = 150; + algebra.compute_basis(max_degree); + + let n_dev = gpu_count().max(4); + let mut overall = vec![0usize; n_dev]; + let mut by_band: Vec> = Vec::new(); + let mut raw = vec![0usize; n_dev]; + let mut total = 0usize; + + for band in 0..5 { + let (lo, hi) = (1 + band * 30, (band + 1) * 30); + let mut counts = vec![0usize; n_dev]; + for deg in lo..=hi.min(max_degree) { + for idx in 0..algebra.dimension(deg as i32) { + let pp = algebra.basis_element_from_index(deg as i32, idx).p_part; + let d = (shard_hash(pp) % n_dev as u64) as usize; + counts[d] += 1; + overall[d] += 1; + raw[(pp.bits() % n_dev as u64) as usize] += 1; + total += 1; + } + } + by_band.push(counts); + } + + assert!(total > 10_000, "need a meaningful sample, got {total}"); + let ideal = total as f64 / n_dev as f64; + for (d, &c) in overall.iter().enumerate() { + let dev = (c as f64 - ideal).abs() / ideal; + assert!( + dev < 0.05, + "device {d} off by {:.1}% overall ({c} vs {ideal:.0})", + dev * 100.0 + ); + } + for (b, counts) in by_band.iter().enumerate() { + let n: usize = counts.iter().sum(); + if n < 500 { + continue; + } + let ideal = n as f64 / n_dev as f64; + for (d, &c) in counts.iter().enumerate() { + let dev = (c as f64 - ideal).abs() / ideal; + assert!( + dev < 0.15, + "degree band {b}, device {d} off by {:.1}%", + dev * 100.0 ); } - - c0 = c1; } + // The unmixed key must be visibly worse, else the mixing is not earning its place. + let raw_worst = raw + .iter() + .map(|&c| ((c as f64 - ideal).abs() / ideal)) + .fold(0.0_f64, f64::max); + let mixed_worst = overall + .iter() + .map(|&c| ((c as f64 - ideal).abs() / ideal)) + .fold(0.0_f64, f64::max); + eprintln!( + "shard_of over {total} real R's, {n_dev} devices: worst deviation {:.2}% mixed vs \ + {:.2}% for a bare bits() % n", + mixed_worst * 100.0, + raw_worst * 100.0 + ); + assert!( + raw_worst > mixed_worst, + "unmixed bits() % n was as balanced as the hash ({raw_worst:.3} vs {mixed_worst:.3}); \ + if the packing changed so this no longer holds, revisit shard_of's rationale" + ); + } - let bytes = client.read_one(out_h).unwrap(); - let flat = u32::from_bytes(&bytes); - let result: Vec> = (0..num_rows) - .map(|r| flat[r * num_limbs..(r + 1) * num_limbs].to_vec()) - .collect(); - - // `out_h` alone is `num_rows × num_limbs` u32 — hundreds of MB at record degrees. - // It (and the small per-launch buffers, now dropped) varies in size launch to - // launch, so CubeCL's pool cannot reuse the slab and would accumulate them until - // the 4 GB card OOMs. Return the freed memory to the driver each launch; the - // resident admissible handles stay alive (refcount > 0) so cleanup skips them. - client.memory_cleanup(); - - result - }); - - // Aggregate marshal/device totals across every launch (cheap, always on) so a whole - // resolution's GPU overhead can be split host-vs-device via [`take_batch_stats`]. - let device_ms = t_device.elapsed().as_secs_f64() * 1e3; - BATCH_CALLS.fetch_add(1, std::sync::atomic::Ordering::Relaxed); - BATCH_MARSHAL_US.fetch_add( - (marshal_ms * 1e3) as u64, - std::sync::atomic::Ordering::Relaxed, - ); - BATCH_DEVICE_US.fetch_add( - (device_ms * 1e3) as u64, - std::sync::atomic::Ordering::Relaxed, - ); - BATCH_PAIRS.fetch_add(total_pairs as u64, std::sync::atomic::Ordering::Relaxed); - - result -} + /// The flag-based [`enumerate_admissible_ref`] must reproduce `admissible_matrices` bit-for-bit + /// on every real `R` — this validates the no-break/continue/return restructuring (the tricky + /// part of the future cubecl in-kernel port) purely on the CPU, where it is fast to debug. + /// Pure CPU: no GPU needed. + #[test] + fn admissible_enum_ref_matches() { + use fp::prime::ValidPrime; -#[cfg(test)] -mod tests { - use super::*; + let p = ValidPrime::new(2); + let algebra = Arc::new(MilnorAlgebra::new(p, false)); + // To 150: the eviction bench faults on cold (high-degree) R's at internal degree ~144, above + // the degree-40/60 originally checked — extend the CPU reference to that range. + let max_degree = 150; + algebra.compute_basis(max_degree); + let mut checked = 0usize; + for deg in 1..=max_degree { + for idx in 0..algebra.dimension(deg) { + let p_part: Vec = algebra + .basis_element_from_index(deg, idx) + .p_part + .iter() + .collect(); + if p_part.is_empty() { + continue; + } + let want = algebra.admissible_matrices(PPart::try_from_slice(&p_part).unwrap()); + let got = enumerate_admissible_ref(&p_part); + assert_eq!(got, want, "R degree {deg} idx {idx} p_part {p_part:?}"); + checked += 1; + } + } + assert!(checked > 0, "no R's exercised"); + eprintln!("admissible_enum_ref: {checked} R's matched admissible_matrices"); + } /// Smoke test proving the CubeCL `cuda` runtime launches and returns correct /// results. Requires a live GPU + the CUDA toolkit env (run under the `gpu` @@ -943,7 +6124,7 @@ mod tests { use fp::prime::ValidPrime; let p = ValidPrime::new(2); - let algebra = MilnorAlgebra::new(p, false); + let algebra = Arc::new(MilnorAlgebra::new(p, false)); let max_degree = 60; algebra.compute_basis(max_degree); algebra.compute_seqno_tables(max_degree); @@ -961,7 +6142,7 @@ mod tests { 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) { + for (slot, v) in row.iter_mut().zip(elt.p_part.iter()) { *slot = v; } p_parts.extend_from_slice(&row); @@ -984,7 +6165,7 @@ mod tests { use fp::{prime::ValidPrime, vector::FpVector}; let p = ValidPrime::new(2); - let algebra = MilnorAlgebra::new(p, false); + let algebra = Arc::new(MilnorAlgebra::new(p, false)); let max_degree = 40; algebra.compute_basis(max_degree); algebra.compute_seqno_tables(max_degree); @@ -1057,7 +6238,7 @@ mod tests { use fp::{prime::ValidPrime, vector::FpVector}; let p = ValidPrime::new(2); - let algebra = MilnorAlgebra::new(p, false); + let algebra = Arc::new(MilnorAlgebra::new(p, false)); let max_degree = 40; algebra.compute_basis(max_degree); algebra.compute_seqno_tables(max_degree); @@ -1102,7 +6283,7 @@ mod tests { 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 { + for &ti in prod.term_indices.iter() { s.set_entry(ti, 1); } let mut tmp = FpVector::new(p, out_dim); @@ -1127,6 +6308,7 @@ mod tests { packed }) .collect(); + let golden = BatchOutput::from_rows(&golden, num_limbs); let got = multiply_batch_on_gpu(&algebra, out_dim, num_rows, &products); assert_eq!( @@ -1138,4 +6320,722 @@ mod tests { products.len() ); } + + /// The same batch, run through every regime of the eviction split ([`MasterMode`]) — all-resident, + /// all-transient, and the mixed two-launch path — must give the identical matrix. + /// + /// This is the coverage `multiply_batch_matches_reference` does not provide: it only ever runs the + /// default `cap == i32::MAX` fast path, so *nothing* exercised on-device enumeration feeding the + /// multiply until now. `admissible_enum_gpu_matches` validates the enumeration kernel in isolation + /// through a test-only harness, which is a different launch path — it says nothing about whether + /// the scratch it writes is laid out the way `multiply_batch_kernel` indexes it. + /// + /// The caps are chosen against the batch's actual `R` degrees (1..`out_degree`): `MAX` keeps every + /// `R` resident, `0` pushes every `R` transient, and the interior ones straddle the split so both + /// launches run and their disjoint row sets have to reassemble correctly. + #[test] + fn multiply_batch_matches_reference_under_eviction() { + // Low degree pins the *logic* of the split; the high-degree pass is where a cap actually bites + // in production (the stem-210 θ=125 run faulted only once the frontier reached t≈185, with + // every low-degree block before it clean), so a small-batch-only test would prove nothing + // about the regime the knob exists for. + check_eviction_regimes(24); + check_eviction_regimes(72); + } + + fn check_eviction_regimes(out_degree: i32) { + use fp::{prime::ValidPrime, vector::FpVector}; + + let p = ValidPrime::new(2); + let algebra = Arc::new(MilnorAlgebra::new(p, false)); + let max_degree = out_degree + 16; + algebra.compute_basis(max_degree); + algebra.compute_seqno_tables(max_degree); + + let out_dim = algebra.dimension(out_degree); + let num_rows = 8; + + 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, + }); + } + } + + 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.iter() { + 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 golden = BatchOutput::from_rows(&golden, num_limbs); + + for cap in [ + i32::MAX, + 0, + 1, + out_degree / 3, + out_degree / 2, + out_degree - 1, + ] { + let got = multiply_batch_gpu_inner(&algebra, out_dim, num_rows, &products, cap); + let transient = products.iter().filter(|p| p.r_degree > cap).count(); + assert_eq!( + got, + golden, + "batched GPU multiply diverged from reference at cap {cap} ({transient}/{} \ + products transient)", + products.len() + ); + eprintln!( + "multiply_batch cap={cap}: {transient}/{} products transient, matched reference", + products.len() + ); + } + } + + /// The stopgap CPU fallback ([`cpu_multiply_batch`]) must produce byte-identical output to the GPU + /// batch multiply — otherwise a mid-run GPU-context death would silently corrupt the resolution. + /// Uses TWO generator blocks at distinct `out_offset`s in one wide row, the module-row layout the + /// single-block `multiply_batch_matches_reference` does not exercise. + #[test] + fn cpu_multiply_batch_matches_gpu() { + use fp::prime::ValidPrime; + + let p = ValidPrime::new(2); + let algebra = Arc::new(MilnorAlgebra::new(p, false)); + let max_degree = 44; + algebra.compute_basis(max_degree); + algebra.compute_seqno_tables(max_degree); + + let num_rows = 6; + // Block A (degree 24) at offset 0, block B (degree 20) immediately after — a two-generator + // row of width dim(A) + dim(B), so products carry a nonzero `out_offset`. + let (deg_a, deg_b) = (24, 20); + let (dim_a, dim_b) = (algebra.dimension(deg_a), algebra.dimension(deg_b)); + let num_cols = dim_a + dim_b; + + let mut products = Vec::new(); + for (out_deg, out_offset) in [(deg_a, 0usize), (deg_b, dim_a)] { + for r_degree in 1..out_deg { + let s_degree = out_deg - 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, + }); + } + } + } + + let gpu = multiply_batch_on_gpu(&algebra, num_cols, num_rows, &products); + let cpu = cpu_multiply_batch(&algebra, num_cols, num_rows, &products); + assert_eq!( + gpu, cpu, + "cpu_multiply_batch diverged from the GPU batch multiply (out_offset path)" + ); + eprintln!( + "cpu_multiply_batch matches GPU: {} products, {num_rows} rows, num_cols={num_cols}", + products.len() + ); + } + + /// Drive several batched multiplies at INCREASING output degree so the shared segmented resident + /// master (see [`SegBuf`], [`seg_grow`]) grows ACROSS launches — each later call appends into the + /// partially-filled last segment and, at a small `NASSAU_GPU_MASTER_SEG_ELEMS`, allocates fresh + /// segments. Run with e.g. `NASSAU_GPU_MASTER_SEG_ELEMS=64` to force many segments and exercise + /// the cross-launch append (the one growth sub-path the single-shot test above does not reach); + /// with the default large segment it still checks correctness on the single-segment path. + #[test] + fn multiply_batch_incremental_growth() { + use fp::{prime::ValidPrime, vector::FpVector}; + + let p = ValidPrime::new(2); + let algebra = Arc::new(MilnorAlgebra::new(p, false)); + let max_degree = 48; + algebra.compute_basis(max_degree); + algebra.compute_seqno_tables(max_degree); + let num_rows = 8; + + // One batched multiply at `out_degree`, checked against the CPU reference. Reused across a + // sequence of growing degrees; the resident master persists (process-global) and grows + // monotonically between calls. + let check = |out_degree: i32| { + let out_dim = algebra.dimension(out_degree); + 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; + } + for r_idx in 0..algebra.dimension(r_degree) { + 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, + }); + } + } + let mut cpu_rows: Vec = + (0..num_rows).map(|_| FpVector::new(p, out_dim)).collect(); + for prod in &products { + let mut s = FpVector::new(p, algebra.dimension(prod.s_degree)); + for &ti in prod.term_indices.iter() { + 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 golden = BatchOutput::from_rows(&golden, num_limbs); + let got = multiply_batch_on_gpu(&algebra, out_dim, num_rows, &products); + assert_eq!( + got, golden, + "incremental-growth GPU multiply diverged from reference at \ + out_degree={out_degree}" + ); + }; + + // Strictly increasing so every call grows the resident master past the previous one. + for out_degree in [14, 20, 26, 32, 38] { + check(out_degree); + } + eprintln!( + "multiply_batch_incremental_growth: 5 growing launches matched reference \ + (seg_elems={})", + master_seg_elems() + ); + } + + /// Concurrency + growth soak on the GPU Milnor multiply: many worker threads hammer + /// [`multiply_batch_on_gpu`] against ONE shared resident master while it grows, across many + /// streams with per-launch `memory_cleanup` on — the same access pattern that trips the cubecl + /// cross-stream pool-reclaim race (tracel-ai/cubecl#1401) in the stem-200 resolution (one + /// stream's cleanup reclaiming a pool page another stream's in-flight launch still reads). + /// + /// It is BOTH a fast correctness guard and a fast #1401 reproducer: + /// - **Correctness:** every GPU result is compared against the bit-identical + /// [`cpu_multiply_batch`] oracle (up to `verify_max`), catching cross-stream renumber/identity + /// races; any mid-soak context death also flips `GPU_DISABLED` and fails the final assert. + /// - **#1401 (historic):** this used to reproduce the cross-stream pool-reclaim race — at + /// `max_degree=160` with `NASSAU_GPU_STREAMS=48` the `never initialized` / `ServerUnhealthy` + /// cascade fired within ~45 s at only ~28 GB host / ~22 GB GPU, a low-memory stand-in for the + /// 40-min stem-200 crash. The dedicated-GPU-thread redesign (see [`gpu_thread`]) deleted the + /// multi-stream mode outright: every device section now runs on one thread on stream 0, so + /// there are no cross-stream reclaims left to race. This config must now be GREEN, and this + /// test is the gate that says so. + /// + /// Ignored by default (needs a CUDA device). Run the ex-reproducer config with: + /// ```text + /// NASSAU_GPU_CLEANUP_EVERY=1 NASSAU_SOAK_MAX_DEGREE=160 \ + /// cargo test -p algebra --release --features gpu -- --ignored --nocapture concurrent_growth_soak + /// ``` + /// Tunables (env): `NASSAU_SOAK_THREADS` (64), `NASSAU_SOAK_SECS` (60), `NASSAU_SOAK_MAX_DEGREE` + /// (60), `NASSAU_SOAK_VERIFY_MAX` (44, the degree ceiling for the CPU-oracle correctness check). + #[test] + #[ignore = "GPU concurrency soak: needs a CUDA device; run explicitly"] + fn concurrent_growth_soak() { + use std::{ + sync::{ + Arc, + atomic::{AtomicU64, Ordering}, + }, + time::{Duration, Instant}, + }; + + let env_num = |key: &str, default: u64| -> u64 { + std::env::var(key) + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(default) + }; + let threads = env_num("NASSAU_SOAK_THREADS", 64) as usize; + let secs = env_num("NASSAU_SOAK_SECS", 60); + let max_degree = env_num("NASSAU_SOAK_MAX_DEGREE", 60) as i32; + // Correctness is checked only up to this degree — the CPU-oracle precompute + // ([`cpu_multiply_batch`]) explodes past ~degree 48, so keep it modest while letting + // `max_degree` run far higher to widen the resident master (the race window scales with + // buffer size). Degrees above the cap still launch on the GPU for the stability/#1401 axis; + // their output just isn't compared. + let verify_max = env_num("NASSAU_SOAK_VERIFY_MAX", 44) as i32; + let num_rows = 32usize; + + let p = fp::prime::ValidPrime::new(2); + let algebra = Arc::new(MilnorAlgebra::new(p, false)); + algebra.compute_basis(max_degree); + algebra.compute_seqno_tables(max_degree); + + // One `get_partial_matrix`-shaped batch per output degree: every non-empty R of degree + // 1..out_degree times a dense complementary element, round-robin across rows, single block. + // Mirrors `multiply_batch_incremental_growth` / the throughput bench, so it hits the exact + // kernel + resident-master path Nassau drives. + let build_batch = |out_degree: i32| -> (usize, Vec) { + let num_cols = algebra.dimension(out_degree); + 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; + } + for r_idx in 0..algebra.dimension(r_degree) { + 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, + }); + } + } + (num_cols, products) + }; + + // Ascending degrees so the sweep drives resident-master growth; precompute each batch and + // its CPU golden once (shared, read-only) so worker threads only launch + compare. + struct Job { + num_cols: usize, + products: Vec, + golden: Option, + } + let jobs: Arc> = Arc::new( + (12..=max_degree) + .step_by(2) + .filter_map(|d| { + let (num_cols, products) = build_batch(d); + if products.is_empty() { + return None; + } + let golden = (d <= verify_max) + .then(|| cpu_multiply_batch(&algebra, num_cols, num_rows, &products)); + Some(Job { + num_cols, + products, + golden, + }) + }) + .collect(), + ); + assert!( + !jobs.is_empty(), + "no non-empty batches built up to degree {max_degree}" + ); + + let launches = AtomicU64::new(0); + let mismatches = AtomicU64::new(0); + let started = Instant::now(); + let deadline = started + Duration::from_secs(secs); + + std::thread::scope(|scope| { + for t in 0..threads { + let jobs = Arc::clone(&jobs); + let algebra = &algebra; + let launches = &launches; + let mismatches = &mismatches; + scope.spawn(move || { + // Desynchronize threads across the degree sweep so some GROW the master (first + // touch of a high degree) while others READ lower resident pages + cleanup. + let mut i = t % jobs.len(); + while Instant::now() < deadline && !gpu_disabled() { + let job = &jobs[i]; + let got = + multiply_batch_on_gpu(algebra, job.num_cols, num_rows, &job.products); + launches.fetch_add(1, Ordering::Relaxed); + // A mismatch while the GPU is still enabled is a real concurrency bug (a + // cross-stream renumber/identity race). Once disabled, results come from the + // bit-identical CPU oracle, so they still match — no false alarm. Degrees + // above `verify_max` have no golden and only exercise stability. + if let Some(golden) = &job.golden { + if got != *golden && !gpu_disabled() { + mismatches.fetch_add(1, Ordering::Relaxed); + } + } + i = (i + 1) % jobs.len(); + } + }); + } + }); + + let elapsed = started.elapsed(); + let n = launches.load(Ordering::Relaxed); + let mm = mismatches.load(Ordering::Relaxed); + eprintln!( + "[soak] {threads} threads × {secs}s, 1 gpu thread: {n} launches ({:.0}/s over {} \ + degrees, verified ≤{verify_max}), {mm} correctness mismatches, gpu_disabled={}", + n as f64 / elapsed.as_secs_f64().max(1e-3), + jobs.len(), + gpu_disabled(), + ); + assert_eq!( + mm, 0, + "GPU multiply diverged from the CPU oracle under concurrency (renumber/identity race)" + ); + assert!( + !gpu_disabled(), + "cubecl GPU multiply was disabled mid-soak — the cross-stream pool-reclaim race \ + (tracel-ai/cubecl#1401) fired. This is the crash the submission-thread redesign \ + closes." + ); + } + + /// Benchmark of the **hard stem-200 regime**: the GPU submission path under the contention + /// shape a record-stem Nassau resolution actually produces. + /// + /// # Why this exists + /// + /// Every change to the GPU path was previously validated by a ~3 h stem-200 resolution, so the + /// iteration loop was measured in hours and each answer arrived with run-to-run variance mixed + /// in. This reproduces the regime in minutes. + /// + /// # Calibration + /// + /// The shape below is not invented; it is the measured distribution of `gpu_submit` spans from + /// a complete stem-200 run, restricted to the hard tail (stems ≥ 190, n = 11 287 launches): + /// + /// ```text + /// p10 p50 p90 max + /// rows 122 158 260 83 702 + /// pairs 190 962 1 963 056 14 770 242 632 386 884 + /// out_u32 326 106 379 516 665 860 42 269 510 + /// ``` + /// + /// Two properties matter as much as the sizes: + /// - **Worker count ≈ 7**, the real wavefront (time-weighted mean 6.6, and 87 % of the run sits + /// at 6–7 bidegrees in flight). The soak's 64 threads are deliberately *wrong* here: queue + /// contention is the thing under test, so the number of contenders must match the resolution. + /// - **Steady state, not growth.** The timed phase runs after a warm-up sweep has grown the + /// resident master, because at stem 200 the master is long since built; timing the growth + /// transient would measure a phase the hard regime is not in. + /// + /// `num_cols` is chosen by searching for the output degree whose dimension is closest to + /// `NASSAU_BENCH_COLS`, rather than hard-coding a degree — the degree that yields a given + /// matrix width is an artifact of the algebra, and pinning the *width* is what keeps this + /// comparable to the measured run. + /// + /// Reports launches/s, pairs/s and the prep / queue / exec split with queue depth, plus the + /// achieved workload distribution so drift from the calibration above is visible rather than + /// silent. + /// + /// Ignored by default (needs a CUDA device). Run with: + /// ```text + /// cargo test -p algebra --release --features gpu -- --ignored --nocapture stem200_regime_bench + /// ``` + /// Tunables (env): `NASSAU_BENCH_WORKERS` (7), `NASSAU_BENCH_ROWS` (158), + /// `NASSAU_BENCH_COLS` (77 000), `NASSAU_BENCH_SECS` (60), `NASSAU_BENCH_SPREAD` (4). + #[test] + #[ignore = "GPU perf bench: needs a CUDA device; run explicitly"] + fn stem200_regime_bench() { + use std::{ + sync::{ + Arc, Mutex, + atomic::{AtomicU64, Ordering}, + }, + time::{Duration, Instant}, + }; + + let env_num = |key: &str, default: u64| -> u64 { + std::env::var(key) + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(default) + }; + let workers = env_num("NASSAU_BENCH_WORKERS", 7) as usize; + let num_rows = env_num("NASSAU_BENCH_ROWS", 158) as usize; + let target_cols = env_num("NASSAU_BENCH_COLS", 77_000) as usize; + let secs = env_num("NASSAU_BENCH_SECS", 60); + // How many neighbouring degrees to sweep. A single degree would let every launch reuse one + // resident-master prefix; the real run interleaves several bidegrees at once. + let spread = env_num("NASSAU_BENCH_SPREAD", 4) as i32; + + let p = fp::prime::ValidPrime::new(2); + let algebra = Arc::new(MilnorAlgebra::new(p, false)); + + // Grow the basis until it brackets `target_cols`, then take the closest degree. Doubling + // the probe keeps this from computing a far larger basis than the bench needs. + let mut probe = 32; + loop { + algebra.compute_basis(probe); + if algebra.dimension(probe) >= target_cols || probe > 512 { + break; + } + probe *= 2; + } + let out_degree = (1..=probe) + .min_by_key(|&d| algebra.dimension(d).abs_diff(target_cols)) + .expect("non-empty degree range"); + let max_degree = out_degree; + algebra.compute_basis(max_degree); + algebra.compute_seqno_tables(max_degree); + eprintln!( + "[bench] target_cols={target_cols} -> out_degree={out_degree} (num_cols={}), \ + workers={workers} rows={num_rows} spread={spread} secs={secs}", + algebra.dimension(out_degree), + ); + + // Same `get_partial_matrix`-shaped batch the soak builds, but with the product count + // BOUNDED. + // + // Taking every `(R, s)` pair the way the soak does is fine to degree ~160 and impossible + // above it: the batch holds `sum_r dim(r) * dim(out_degree - r)` term indices, which at the + // degree that yields stem-200-scale matrices is astronomically large — the first version of + // this bench died in setup there, never reaching a launch. Real batches are bounded too + // (measured p50 ~24k products, max ~479k), so sampling `R`s on a stride reproduces the + // regime while the exhaustive build merely runs out of memory. + let max_products = env_num("NASSAU_BENCH_PRODUCTS", 24_000) as usize; + // Terms per product. Real products carry the nonzeros of a sparse vector, not a whole + // basis: run I measured `products=478972 terms=2300621`, i.e. ~4.8 terms each. Using the + // full s-degree basis (as the low-degree soak does) makes every product hundreds of + // thousands of terms here, which blows past the kernel's u32 pair limit before it can run. + let terms_per_product = env_num("NASSAU_BENCH_TERMS", 5) as usize; + let build_batch = |out_degree: i32| -> (usize, Vec) { + let num_cols = algebra.dimension(out_degree); + // Count the candidate `R`s first so the stride spreads the sample over the whole + // degree range instead of truncating at low `r_degree` (which would bias every launch + // toward small, cheap operations). + let mut candidates: Vec<(i32, usize, i32)> = Vec::new(); + for r_degree in 1..out_degree { + let s_degree = out_degree - r_degree; + if algebra.dimension(s_degree) == 0 { + continue; + } + for r_idx in 0..algebra.dimension(r_degree) { + if algebra + .basis_element_from_index(r_degree, r_idx) + .p_part + .is_empty() + { + continue; + } + candidates.push((r_degree, r_idx, s_degree)); + } + } + let stride = candidates.len().div_ceil(max_products.max(1)).max(1); + let mut products = Vec::new(); + for (r_degree, r_idx, s_degree) in candidates.into_iter().step_by(stride) { + let s_dim = algebra.dimension(s_degree); + let nt = terms_per_product.min(s_dim); + let t_stride = s_dim.div_ceil(nt.max(1)).max(1); + let row = products.len() % num_rows; + products.push(GpuProduct { + r_degree, + r_idx, + s_degree, + term_indices: (0..s_dim).step_by(t_stride).take(nt).collect(), + row, + out_offset: 0, + }); + } + (num_cols, products) + }; + + struct Job { + num_cols: usize, + products: Vec, + } + let jobs: Arc> = Arc::new( + (out_degree - spread + 1..=out_degree) + .filter(|&d| d > 1) + .filter_map(|d| { + let (num_cols, products) = build_batch(d); + (!products.is_empty()).then_some(Job { num_cols, products }) + }) + .collect(), + ); + assert!( + !jobs.is_empty(), + "no non-empty batches at degree {out_degree}" + ); + + // Warm-up: one pass per job grows the resident master to its steady-state extent, so the + // timed phase below measures the regime rather than the growth transient. + let warm = Instant::now(); + for job in jobs.iter() { + let _ = multiply_batch_on_gpu(&algebra, job.num_cols, num_rows, &job.products); + } + eprintln!( + "[bench] warm-up: {} jobs in {:.1}s", + jobs.len(), + warm.elapsed().as_secs_f64() + ); + assert!(!gpu_disabled(), "GPU died during warm-up"); + + // Discard warm-up from the counters; the timed phase starts from zero. + let _ = take_batch_stats(); + let _ = take_gpu_timing(); + + let launches = AtomicU64::new(0); + // Per-launch wall time: the tail is the starvation signal the aggregate mean hides. + let waits: Mutex> = Mutex::new(Vec::new()); + let started = Instant::now(); + let deadline = started + Duration::from_secs(secs); + + std::thread::scope(|scope| { + for t in 0..workers { + let jobs = Arc::clone(&jobs); + let algebra = &algebra; + let launches = &launches; + let waits = &waits; + scope.spawn(move || { + let mut local = Vec::new(); + let mut i = t % jobs.len(); + while Instant::now() < deadline && !gpu_disabled() { + let job = &jobs[i]; + let t0 = Instant::now(); + let _ = + multiply_batch_on_gpu(algebra, job.num_cols, num_rows, &job.products); + local.push(t0.elapsed().as_secs_f64()); + launches.fetch_add(1, Ordering::Relaxed); + i = (i + 1) % jobs.len(); + } + waits.lock().unwrap().extend(local); + }); + } + }); + + let elapsed = started.elapsed().as_secs_f64().max(1e-3); + let n = launches.load(Ordering::Relaxed); + let (calls, _marshal_us, device_us, pairs) = take_batch_stats(); + let (prep_us, queue_us, exec_us, depth_sum, depth_max) = take_gpu_timing(); + let us = |v: u64| v as f64 / 1e6; + let total = (us(prep_us) + us(device_us)).max(1e-9); + + let mut w = waits.into_inner().unwrap(); + w.sort_by(|a, b| a.partial_cmp(b).unwrap()); + let q = |pc: usize| w[(w.len() * pc / 100).min(w.len().saturating_sub(1))]; + + eprintln!( + "[bench] {n} calls ({calls} blocks) in {elapsed:.1}s: {:.1} calls/s, {:.2e} pairs/s", + n as f64 / elapsed, + pairs as f64 / elapsed, + ); + eprintln!( + "[bench] prep={:.1}s queue={:.1}s exec={:.1}s | prep={:.0}% queue={:.0}% exec={:.0}% \ + | depth mean={:.1} max={depth_max}", + us(prep_us), + us(queue_us), + us(exec_us), + 100.0 * us(prep_us) / total, + 100.0 * us(queue_us) / total, + 100.0 * us(exec_us) / total, + depth_sum as f64 / calls.max(1) as f64, + ); + eprintln!( + "[bench] per-call wall: p50={:.3}s p90={:.3}s p99={:.3}s max={:.3}s (ratio \ + max/p50={:.0}x)", + q(50), + q(90), + q(99), + w[w.len() - 1], + w[w.len() - 1] / q(50).max(1e-9), + ); + eprintln!( + "[bench] gpu-thread duty cycle: {:.0}% ({:.1}s exec of {elapsed:.1}s wall)", + 100.0 * us(exec_us) / elapsed, + us(exec_us), + ); + + assert!(!gpu_disabled(), "GPU context died during the bench"); + assert!(n > 0, "no launches completed"); + } } diff --git a/ext/crates/algebra/src/algebra/milnor_rank.rs b/ext/crates/algebra/src/algebra/milnor_rank.rs new file mode 100644 index 0000000000..65c6f0a7e4 --- /dev/null +++ b/ext/crates/algebra/src/algebra/milnor_rank.rs @@ -0,0 +1,267 @@ +//! An arithmetic replacement for the `MilnorBasisElement -> index` hash map. +//! +//! Behind the off-by-default `milnor-rank` feature, and **not wired into [`MilnorAlgebra`]** even +//! when enabled -- `basis_element_to_index` still goes through the hash map. It is here so the +//! design and its measurements survive, ready to switch on when the renumbering below is worth +//! taking on. +//! +//! It computes a *different* numbering of each degree's basis than [`MilnorAlgebra::compute_basis`] +//! produces, so adopting it would renumber the basis and invalidate every saved resolution. See +//! [`PPartRanker`] for why the numbering cannot simply be made to match, and the crate benchmarks +//! (`milnor_rank`) for what it costs relative to the hash map it would replace. +//! +//! [`MilnorAlgebra`]: crate::MilnorAlgebra +//! [`MilnorAlgebra::compute_basis`]: crate::Algebra::compute_basis + +use fp::prime::ValidPrime; + +use crate::algebra::{combinatorics, milnor_algebra::PPart}; + +/// Computes a p-part's index within its degree arithmetically, instead of looking it up. +/// +/// # How it works +/// +/// `counts[i][d]` is the number of exponent sequences of degree `d` (in units of `q`) that use +/// only $\xi_1, \ldots, \xi_i$. Splitting on whether $r_i$ is zero gives the coin-change +/// recurrence +/// +/// ```text +/// counts[i][d] = counts[i - 1][d] + counts[i][d - xi_i] +/// ``` +/// +/// so the table costs `O(MAX_LEN * max_degree)` to build, and `counts[MAX_LEN][d]` is the +/// dimension of the algebra in degree `d`. +/// +/// Ranking then rests on one identity. Among the sequences of degree `d` using +/// $\xi_1, \ldots, \xi_i$, those with $r_i \ge u$ are in bijection with *all* sequences of degree +/// `d - u * xi_i` using $\xi_1, \ldots, \xi_i$, via $r_i \mapsto r_i - u$. So the number of them +/// with $r_i > v$ is `counts[i][d - (v + 1) * xi_i]`: a single lookup, with no summation. Walking +/// the entries from the top down therefore ranks a p-part in [`PPart::MAX_LEN`] lookups and adds, +/// against a table of a few hundred KiB that serves *every* degree at once — where the hash map +/// it replaces stores one entry per basis element. +/// +/// # Speed: it depends entirely on scale +/// +/// Measured against the hash map it would replace (`cargo bench --bench milnor_rank`), at p = 2: +/// +/// ```text +/// degree per-degree map hashmap ranker ratio +/// 120 0.10 MB 11.6us 26.7us 0.43x +/// 300 3.12 MB 792us 1384us 0.57x +/// 400 12.50 MB 4490us 5310us 0.85x +/// 500 37.50 MB 33118us 15865us 2.09x +/// ``` +/// +/// The crossover is a cache effect, and it is not marginal. A lookup probes only its own degree's +/// map. While that fits in cache the map wins easily: one hash round and one probe, against six to +/// ten dependent table reads for a rank. Once it does not — the map is 37 MB in degree 500 — every +/// probe misses to DRAM at ~33 ns, whereas the ranker's whole table is ~43 KB, stays in L1, and +/// costs ~16 ns regardless of degree. The map's cost grows with the basis; the ranker's does not. +/// +/// So the ranker is the wrong tool for small degrees and the right one for large. That matters +/// because large degrees are exactly where the algebra's memory becomes the problem worth solving. +/// +/// The reverse direction does not share this, and is deliberately absent. An `unrank` -- recovering +/// the p-part at a given index, which would let `basis_element_from_index` drop `ppart_table` +/// altogether -- was written, tested and benchmarked, and removed again: it ran ~15x slower than +/// the array read it would replace, at every degree measured, with no sign of the crossover that +/// makes `rank` worthwhile. The reason is that `ppart_table` is only 8 bytes per element, against +/// ~43 for the hash map, so it stays cache-resident where the map does not. The bit-packing that +/// makes `rank` worth having is the same thing that makes `unrank` not. (See the history around +/// "Speed up unrank 1.5x" if it needs revisiting; the likelier route is enumerating the basis in +/// index order, which is O(1) amortised and matches how callers actually walk it.) +/// +/// Tuning does not move this. Five arrangements were measured — nested vs flat count table, with +/// and without a zero-padded prefix to drop the branch, and one- vs two-pass to break the +/// dependency chain. None changed the small-degree verdict; the padded variant was worst, because +/// doubling the table pushed it out of L1. +/// +/// # Why it is still not wired in +/// +/// Not speed, but numbering. This ranks in colex order on $(r_{10}, \ldots, r_1)$. `compute_ppart` emits a different order: +/// it groups by the highest non-zero entry, and recurses by decrementing one entry at a time. +/// That order *is* rankable in principle, but its natural recursion has depth $\sum_i r_i$ — up to +/// `MAX_DEGREE` — which is far worse than hashing. Getting the `O(MAX_LEN)` cost requires adopting +/// the colex order, i.e. renumbering the basis. +/// +/// A renumbering is not intrinsically hard — [`crate::Algebra::magic`] already exists to +/// discriminate save files — but it invalidates stored resolutions, so it is a migration rather +/// than a drop-in change. Note also that [`MilnorAlgebra`] re-sorts each degree by excess when +/// unstable support is enabled, which this does not model. +/// +/// [`MilnorAlgebra`]: crate::MilnorAlgebra +/// [`MilnorAlgebra::compute_basis`]: crate::Algebra::compute_basis +pub struct PPartRanker { + /// `counts[i][d]` flattened to `counts[i * stride + d]`, for `i` in `0..=PPart::MAX_LEN`. + /// + /// Flat rather than `Vec>`, so a lookup is not a dependent pointer chase. + counts: Vec, + stride: usize, + /// `effective_len[d]` is the number of $\xi_i$ of degree at most `d`. + /// + /// Entries beyond it cannot contribute to a rank in degree `d`: such an entry must be zero, and + /// its `cut` is then `d - xi_i < 0`. At degree 120 this is 6 rather than 10, so it removes + /// roughly a third of the work. + effective_len: Vec, + /// `xi[i]` is the degree of $\xi_{i+1}$, divided by `q`. + xi: [i32; PPart::MAX_LEN], + max_degree: i32, +} + +impl PPartRanker { + /// Build the table for degrees `0..=max_degree`, where `max_degree` is measured in units of + /// `q` (so it is the internal degree at `p = 2`, and the internal degree divided by + /// `2(p - 1)` otherwise). + pub fn new(p: ValidPrime, max_degree: i32) -> Self { + assert!(max_degree >= 0); + let mut xi = [0; PPart::MAX_LEN]; + xi.copy_from_slice(&combinatorics::xi_degrees(p)[..PPart::MAX_LEN]); + + let stride = max_degree as usize + 1; + + let mut counts = vec![0; (PPart::MAX_LEN + 1) * stride]; + // The empty sequence is the unique sequence of degree 0 using no generators. + counts[0] = 1; + for i in 1..=PPart::MAX_LEN { + for d in 0..stride { + // Either r_i is zero, or we can subtract one from it. + counts[i * stride + d] = counts[(i - 1) * stride + d]; + if d >= xi[i - 1] as usize { + counts[i * stride + d] += counts[i * stride + d - xi[i - 1] as usize]; + } + } + } + + let effective_len = (0..=max_degree) + .map(|d| xi.iter().filter(|&&x| x <= d).count() as u8) + .collect(); + + Self { + counts, + stride, + effective_len, + xi, + max_degree, + } + } + + /// The number of p-parts of degree `degree`, i.e. what `MilnorAlgebra::dimension` returns for + /// the p-part factor of the basis. + pub fn dimension(&self, degree: i32) -> u64 { + if degree < 0 || degree > self.max_degree { + 0 + } else { + self.counts[PPart::MAX_LEN * self.stride + degree as usize] + } + } + + /// The index of `p_part` among the p-parts of degree `degree`, in the colex order described on + /// [`PPartRanker`]. + /// + /// `degree` must be the degree of `p_part` (in units of `q`), and at most the `max_degree` + /// this was built with. + /// + /// # Cost + /// + /// One table read per entry, serialised through the running `remaining`. That is the reason + /// this loses to the hash map it was meant to replace, and no arrangement of the table fixes + /// it — see the module docs. + #[inline] + pub fn rank(&self, p_part: PPart, degree: i32) -> usize { + debug_assert!(degree >= 0 && degree <= self.max_degree); + let mut rank = 0; + let mut remaining = degree; + // Only the entries with `xi_i <= degree` can contribute; the rest are zero with a negative + // cut. At degree 120 that is 6 iterations rather than 10. + for i in (0..self.effective_len[degree as usize] as usize).rev() { + let entry = p_part.get(i) as i32; + // Everything with a larger entry here sorts earlier, and there are exactly + // `counts[i + 1][remaining - (entry + 1) * xi_i]` of them. + let cut = remaining - (entry + 1) * self.xi[i]; + if cut >= 0 { + rank += self.counts[(i + 1) * self.stride + cut as usize]; + } + remaining -= entry * self.xi[i]; + } + debug_assert_eq!(remaining, 0, "degree does not match the p-part"); + rank as usize + } +} + +#[cfg(test)] +mod tests { + use fp::prime::Prime; + use rstest::rstest; + + use super::*; + use crate::{Algebra, MilnorAlgebra}; + + /// `counts[MAX_LEN]` must agree with the algebra's own count of p-parts in each degree. + #[rstest] + #[case(2, 120)] + #[case(3, 40)] + fn table_matches_ppart_table(#[case] p: u32, #[case] max_degree: i32) { + let p = ValidPrime::new(p); + let algebra = MilnorAlgebra::new(p, false); + let q = if p == 2 { 1 } else { 2 * (p.as_i32() - 1) }; + algebra.compute_basis(max_degree * q); + + let ranker = PPartRanker::new(p, max_degree); + for d in 0..=max_degree { + assert_eq!( + ranker.dimension(d), + algebra.ppart_table(d).len() as u64, + "dimension mismatch in degree {d}" + ); + } + } + + /// The whole point: `rank` must be a bijection from the p-parts of each degree onto + /// `0..dimension`. If it is, it is a valid numbering and could replace the hash map. + #[rstest] + #[case(2, 120)] + #[case(3, 40)] + fn rank_is_a_bijection(#[case] p: u32, #[case] max_degree: i32) { + let p = ValidPrime::new(p); + let algebra = MilnorAlgebra::new(p, false); + let q = if p == 2 { 1 } else { 2 * (p.as_i32() - 1) }; + algebra.compute_basis(max_degree * q); + + let ranker = PPartRanker::new(p, max_degree); + for d in 0..=max_degree { + let table = algebra.ppart_table(d); + let mut seen = vec![false; table.len()]; + for &p_part in table { + let rank = ranker.rank(p_part, d); + assert!(rank < table.len(), "rank {rank} out of range in degree {d}"); + assert!(!seen[rank], "rank {rank} hit twice in degree {d}"); + seen[rank] = true; + } + } + } + + /// Ranking in colex order really is a different numbering than the one the algebra uses. This + /// is the reason the ranker is not wired in, so pin it down rather than leave it to prose. + #[test] + fn rank_disagrees_with_the_current_basis_order() { + let p = ValidPrime::new(2); + let algebra = MilnorAlgebra::new(p, false); + algebra.compute_basis(120); + let ranker = PPartRanker::new(p, 120); + + let mut agree = 0; + let mut total = 0; + for d in 0..=120 { + for (i, &p_part) in algebra.ppart_table(d).iter().enumerate() { + total += 1; + if ranker.rank(p_part, d) == i { + agree += 1; + } + } + } + assert!( + agree * 100 < total, + "expected the orders to differ on almost everything, but {agree}/{total} agreed" + ); + } +} diff --git a/ext/crates/algebra/src/algebra/mod.rs b/ext/crates/algebra/src/algebra/mod.rs index 4513884e2a..52091df340 100644 --- a/ext/crates/algebra/src/algebra/mod.rs +++ b/ext/crates/algebra/src/algebra/mod.rs @@ -20,6 +20,10 @@ pub use milnor_algebra::MilnorAlgebra; #[cfg(feature = "gpu")] pub mod milnor_gpu; +// Opt-in: an arithmetic alternative to the Milnor basis index map. Not wired in; see the module +// docs for what it costs and what it would take to adopt. +#[cfg(feature = "milnor-rank")] +pub mod milnor_rank; mod steenrod_algebra; pub use steenrod_algebra::{AlgebraType, SteenrodAlgebra}; diff --git a/ext/crates/algebra/src/algebra/pair_algebra.rs b/ext/crates/algebra/src/algebra/pair_algebra.rs index e080d34116..34815baaa0 100644 --- a/ext/crates/algebra/src/algebra/pair_algebra.rs +++ b/ext/crates/algebra/src/algebra/pair_algebra.rs @@ -95,16 +95,17 @@ use std::cell::RefCell; use crate::{ MilnorAlgebra, - milnor_algebra::{MilnorBasisElement as MilnorElt, PPartAllocation, PPartMultiplier}, + milnor_algebra::{MilnorBasisElement as MilnorElt, PPart, PPartAllocation, PPartMultiplier}, }; macro_rules! sub { ($elt:ident, $k:expr, $n:expr) => { if $k > 0 { - if $elt.p_part[$k - 1] < (1 << $n) { + let entry = $elt.p_part.get($k - 1); + if entry < (1 << $n) { continue; } - $elt.p_part[$k - 1] -= 1 << $n; + $elt.p_part.set($k - 1, entry - (1 << $n)); $elt.degree -= combinatorics::xi_degrees(TWO)[$k - 1] * (1 << $n); } }; @@ -112,7 +113,7 @@ macro_rules! sub { macro_rules! unsub { ($elt:ident, $k:expr, $n:expr) => { if $k > 0 { - $elt.p_part[$k - 1] += 1 << $n; + $elt.p_part.set($k - 1, $elt.p_part.get($k - 1) + (1 << $n)); $elt.degree += combinatorics::xi_degrees(TWO)[$k - 1] * (1 << $n); } }; @@ -191,8 +192,8 @@ impl PairAlgebra for MilnorAlgebra { assert_eq!(r_degree + s_degree, result.degree); // First write the Y terms - let mut r = self.basis_element_from_index(r_degree, r_idx).clone(); - let mut s = self.basis_element_from_index(s_degree, s_idx).clone(); + let mut r = self.basis_element_from_index(r_degree, r_idx); + let mut s = self.basis_element_from_index(s_degree, s_idx); PPartAllocation::with_local(|mut allocation| { for k in 0..s.p_part.len() { @@ -204,8 +205,8 @@ impl PairAlgebra for MilnorAlgebra { allocation = self.multiply_with_allocation( result.ys[m + k][n + k].as_slice_mut(), coeff, - &r, - &s, + r, + s, i32::MAX, allocation, ); @@ -219,8 +220,8 @@ impl PairAlgebra for MilnorAlgebra { // Now the product terms let mut multiplier = PPartMultiplier::::new_from_allocation( TWO, - &r.p_part, - &s.p_part, + r.p_part, + s.p_part, allocation, 0, r.degree + s.degree, @@ -262,16 +263,17 @@ impl PairAlgebra for MilnorAlgebra { // The twos terms for (r_idx, c) in r.iter_nonzero() { - let mut r = self.basis_element_from_index(r_degree, r_idx).clone(); + let mut r = self.basis_element_from_index(r_degree, r_idx); sub!(r, 1, 0); self.multiply_basis_by_element( result.copy(), coeff * c, - &r, + r, s_degree, s.twos.as_slice(), ); - unsub!(r, 1, 0); + // No matching `unsub!`: unlike the loops above, `r` is a fresh copy of the basis + // element on each iteration, so there is nothing to restore. } // The Y terms @@ -364,7 +366,7 @@ thread_local! { /// [`a_y_inner`] if not available. fn a_y_cached( algebra: &MilnorAlgebra, - a: &MilnorElt, + a: MilnorElt, k: usize, l: usize, f: impl FnOnce(&FpVector), @@ -377,7 +379,7 @@ fn a_y_cached( let raw_entry = cache.raw_entry(); let result = raw_entry - .from_hash(hasher.finish(), |v| &v.0 == a && v.1 == (k, l)) + .from_hash(hasher.finish(), |v| v.0 == a && v.1 == (k, l)) .map(|(_, y)| y); match result { @@ -385,19 +387,19 @@ fn a_y_cached( None => { let v = a_y_inner(algebra, a, k, l); f(&v); - cache.insert((a.clone(), (k, l)), v); + cache.insert((a, (k, l)), v); } } }) } /// Actually computes $A(a, Y_{k, l})$ and returns the result. -fn a_y_inner(algebra: &MilnorAlgebra, a: &MilnorElt, k: usize, l: usize) -> FpVector { - let mut a = a.clone(); +fn a_y_inner(algebra: &MilnorAlgebra, a: MilnorElt, k: usize, l: usize) -> FpVector { + let mut a = a; let mut result = FpVector::new(TWO, algebra.dimension(a.degree + (1 << k) + (1 << l) - 2)); let mut t = MilnorElt { q_part: 0, - p_part: vec![], + p_part: PPart::zero(), degree: 0, }; @@ -410,17 +412,15 @@ fn a_y_inner(algebra: &MilnorAlgebra, a: &MilnorElt, k: usize, l: usize) -> FpVe for j in 0..=std::cmp::min(i + k - l, a.p_part.len()) { sub!(a, j, l); - t.p_part.clear(); - t.p_part.resize(k + i, 0); - - t.p_part[k + i - 1] += 1; - t.p_part[l + j - 1] += 1; + t.p_part = PPart::zero(); + t.p_part.set(k + i - 1, 1); + t.p_part.set(l + j - 1, t.p_part.get(l + j - 1) + 1); t.degree = (1 << (k + i)) + (1 << (l + j)) - 2; // We can just read off the value of the product instead of passing through the // algorithm, but this is cached so problem for another day... - algebra.multiply(result.as_slice_mut(), 1, &t, &a); + algebra.multiply(result.as_slice_mut(), 1, t, a); unsub!(a, j, l); } @@ -445,7 +445,7 @@ mod tests { MilnorElt { q_part: 0, - p_part: p_part.into(), + p_part: PPart::from_slice(p_part), degree, } } @@ -462,7 +462,7 @@ mod tests { let target_deg = a.degree + (1 << k) + (1 << l) - 2; algebra.compute_basis(target_deg + 1); result.set_scratch_vector_size(algebra.dimension(target_deg)); - a_y_cached(&algebra, &a, k, l, |v| result.add(v, 1)); + a_y_cached(&algebra, a, k, l, |v| result.add(v, 1)); ans.assert_eq(&algebra.element_to_string(target_deg, result.as_slice())); }; diff --git a/ext/crates/algebra/src/module/free_module.rs b/ext/crates/algebra/src/module/free_module.rs index 9d878f6ce0..74fa0d8c54 100644 --- a/ext/crates/algebra/src/module/free_module.rs +++ b/ext/crates/algebra/src/module/free_module.rs @@ -237,6 +237,21 @@ impl> ZeroModule for MuFreeModule { } impl> MuFreeModule { + /// Diagnostic (see `NASSAU_MEM_REPORT`): total heap bytes held by this module's internal + /// tables — one `OperationGeneratorPair` per basis element per degree (`basis_element_to_opgen`) + /// plus the `generator_to_index` inverse. These scale with Σ dim(t), the same order as a + /// differential's `outputs`, so they are the other half of the resolution's retained RAM. + pub fn table_heap_bytes(&self) -> usize { + let mut bytes = 0usize; + for (_, row) in self.basis_element_to_opgen.iter() { + bytes += row.len() * std::mem::size_of::(); + } + for (_, row) in self.generator_to_index.iter() { + bytes += row.len() * std::mem::size_of::(); + } + bytes + } + pub fn gen_names(&self) -> &OnceBiVec> { &self.gen_names } diff --git a/ext/crates/algebra/src/module/homomorphism/free_module_homomorphism.rs b/ext/crates/algebra/src/module/homomorphism/free_module_homomorphism.rs index 5b349abf05..cd6dc22e39 100644 --- a/ext/crates/algebra/src/module/homomorphism/free_module_homomorphism.rs +++ b/ext/crates/algebra/src/module/homomorphism/free_module_homomorphism.rs @@ -88,6 +88,29 @@ where } } +impl MuFreeModuleHomomorphism +where + M::Algebra: MuAlgebra, +{ + /// Diagnostic (see `NASSAU_MEM_REPORT`): total heap bytes held by the stored `outputs` + /// matrices — for every degree, every generator's image `FpVector`'s limb storage. This is + /// the differential's own retained footprint. `images`/`kernels`/`quasi_inverses` are also + /// summed (they are `None` in Nassau, so contribute ~0) to prove they are not the consumer. + pub fn output_heap_bytes(&self) -> usize { + let mut bytes = 0usize; + for (_, row) in self.outputs.iter() { + bytes += row.capacity() * std::mem::size_of::(); + for v in row { + // p=2: one 64-bit limb per 64 entries. + bytes += v.len().div_ceil(64) * 8; + } + } + // Option slot overhead only (these are `None` in Nassau). + let opt = |n: i32| n.max(0) as usize * std::mem::size_of::(); + bytes + opt(self.images.len()) + opt(self.kernels.len()) + opt(self.quasi_inverses.len()) + } +} + impl MuFreeModuleHomomorphism where M::Algebra: MuAlgebra, diff --git a/ext/crates/algebra/src/module/rpn.rs b/ext/crates/algebra/src/module/rpn.rs index 6d42a1f9f8..1c4aab1d0b 100644 --- a/ext/crates/algebra/src/module/rpn.rs +++ b/ext/crates/algebra/src/module/rpn.rs @@ -157,7 +157,7 @@ fn coef_milnor(algebra: &MilnorAlgebra, op_deg: i32, op_idx: usize, mut mod_degr return false; } - let elt: &MilnorBasisElement = algebra.basis_element_from_index(op_deg, op_idx); + let elt: MilnorBasisElement = algebra.basis_element_from_index(op_deg, op_idx); let sum: PPartEntry = elt.p_part.iter().sum(); if mod_degree < 0 { @@ -170,7 +170,7 @@ fn coef_milnor(algebra: &MilnorAlgebra, op_deg: i32, op_idx: usize, mut mod_degr let mut list = Vec::with_capacity(elt.p_part.len() + 1); list.push(mod_degree - sum); - list.extend_from_slice(&elt.p_part); + list.extend(elt.p_part.iter()); PPartEntry::multinomial2(&list) == 1 } diff --git a/ext/crates/algebra/src/steenrod_evaluator.rs b/ext/crates/algebra/src/steenrod_evaluator.rs index c5b9c56135..27239cb626 100644 --- a/ext/crates/algebra/src/steenrod_evaluator.rs +++ b/ext/crates/algebra/src/steenrod_evaluator.rs @@ -8,7 +8,7 @@ use fp::{ use crate::{ algebra::{AdemAlgebra, Algebra, MilnorAlgebra, adem_algebra::AdemBasisElement}, - milnor_algebra::{MilnorBasisElement, PPartEntry}, + milnor_algebra::{MilnorBasisElement, PPart, PPartEntry}, steenrod_parser::*, }; @@ -157,7 +157,7 @@ impl SteenrodEvaluator { * q; let elt = MilnorBasisElement { degree, - p_part: p_list, + p_part: PPart::from_slice(&p_list), q_part: 0, }; @@ -270,9 +270,9 @@ impl SteenrodEvaluator { return; } let mut t: Vec = vec![0; elt.p_part.len()]; - t[elt.p_part.len() - 1] = elt.p_part[elt.p_part.len() - 1]; + t[elt.p_part.len() - 1] = elt.p_part.get(elt.p_part.len() - 1); for i in (0..elt.p_part.len() - 1).rev() { - t[i] = elt.p_part[i] + 2 * t[i + 1]; + t[i] = elt.p_part.get(i) + 2 * t[i + 1]; } let t_idx = self.adem.basis_element_to_index(&AdemBasisElement { degree, @@ -307,19 +307,10 @@ impl SteenrodEvaluator { (31u32.saturating_sub(elt.q_part.leading_zeros())) as usize, ); let mut t = vec![0; t_len]; - let last_p_part = if t_len <= elt.p_part.len() { - elt.p_part[t_len - 1] - } else { - 0 - }; - t[t_len - 1] = last_p_part + ((elt.q_part >> (t_len)) & 1); + // `PPart::get` already reads past the end as zero. + t[t_len - 1] = elt.p_part.get(t_len - 1) + ((elt.q_part >> (t_len)) & 1); for i in (0..t_len - 1).rev() { - let p_part = if i < elt.p_part.len() { - elt.p_part[i] - } else { - 0 - }; - t[i] = p_part + ((elt.q_part >> (i + 1)) & 1) + p * t[i + 1]; + t[i] = elt.p_part.get(i) + ((elt.q_part >> (i + 1)) & 1) + p * t[i + 1]; } let t_idx = self.adem.basis_element_to_index(&AdemBasisElement { degree, @@ -344,11 +335,11 @@ impl SteenrodEvaluator { MilnorBasisElement { degree, q_part: 1 << qi, - p_part: vec![], + p_part: PPart::zero(), } } else { - let mut p_part = vec![0; qi as usize + 1]; - p_part[qi as usize] = 1; + let mut p_part = PPart::zero(); + p_part.set(qi as usize, 1); MilnorBasisElement { degree, q_part: 0, diff --git a/ext/crates/algebra/src/steenrod_parser.rs b/ext/crates/algebra/src/steenrod_parser.rs index 2d03718653..a5f719571f 100644 --- a/ext/crates/algebra/src/steenrod_parser.rs +++ b/ext/crates/algebra/src/steenrod_parser.rs @@ -15,14 +15,14 @@ use nom::{ sequence::{delimited, pair, preceded}, }; -use crate::{adem_algebra::AdemBasisElement, algebra::milnor_algebra::PPart}; +use crate::{adem_algebra::AdemBasisElement, algebra::milnor_algebra::PPartEntry}; type IResult = IResultBase>; #[derive(Debug, Clone)] pub enum AlgebraBasisElt { AList(Vec), // Admissible list. - PList(PPart), + PList(Vec), P(u32), Q(u32), } diff --git a/ext/crates/fp-cuda/cuda_kernels/matmul_b1.cu b/ext/crates/fp-cuda/cuda_kernels/matmul_b1.cu index 7dd9334b54..01cff6fe8f 100644 --- a/ext/crates/fp-cuda/cuda_kernels/matmul_b1.cu +++ b/ext/crates/fp-cuda/cuda_kernels/matmul_b1.cu @@ -107,6 +107,15 @@ __device__ __forceinline__ void arrive_cluster(uint64_t* b, uint32_t cta_id) { "}\n" :: "r"(local), "r"(cta_id) : "memory"); } +// Single-CTA counterpart of `arrive_cluster`: arrive (count 1) on a *local* +// mbarrier. Used by the cluster-free variant, where the only CTA that ever +// releases a stage is the one that consumed it, so no `mapa` translation is +// needed and no cross-CTA co-residency is implied. +__device__ __forceinline__ void arrive_local(uint64_t* b) { + asm volatile("mbarrier.arrive.shared::cta.b64 _, [%0], 1;\n" + :: "r"((uint32_t)__cvta_generic_to_shared(b)) : "memory"); +} + // TMA load with cluster multicast: one HBM read of the source tile is fanned // out into the SMEM of every CTA whose bit is set in `mask` (same `dst` SMEM // offset and `b` mbarrier offset in each), and counts complete_tx bytes against @@ -243,10 +252,23 @@ constexpr uint32_t DESC_SWIZ = 1; // The output block (TM rows × NG limbs) is packed row-major into sC and written // back with a single TMA bulk store (S2G). C is padded to whole NG-limb column // groups on the host so every stored tile is complete. -extern "C" __global__ void __cluster_dims__(CLUSTER, 1, 1) matmul_b1_kernel( - const __grid_constant__ CUtensorMap tma_a, - const __grid_constant__ CUtensorMap tma_b, - const __grid_constant__ CUtensorMap tma_c, +// The body is templated on the cluster width so one source produces both +// variants (see the two `extern "C"` entry points below). `CLU == 1` is not a +// degenerate special case bolted on: every cluster-dependent construct here has +// an exact single-CTA counterpart, and the compiler discards the other branch. +// rank -> 0 (no %cluster_ctarank read) +// cluster_sync -> nothing (the preceding __syncthreads already orders it) +// arrive_cluster -> arrive_local (no mapa into a cluster-mate's SMEM) +// tma_2d_multicast -> tma_2d (each CTA reads its own B tile) +// mbar_empty count -> 1 instead of CLUSTER +// What remains is identical arithmetic on an identical tile schedule, so the +// two kernels are bit-for-bit equivalent; they differ only in HBM traffic for B +// and in whether the launch demands co-resident CTAs. +template +__device__ __forceinline__ void matmul_b1_body( + const CUtensorMap& tma_a, + const CUtensorMap& tma_b, + const CUtensorMap& tma_c, uint32_t m_tiles, uint32_t n_groups, uint32_t M, uint32_t K) @@ -272,12 +294,14 @@ extern "C" __global__ void __cluster_dims__(CLUSTER, 1, 1) matmul_b1_kernel( // Cluster geometry: CLUSTER CTAs along M share one B-panel via multicast, // so the schedule walks "M-super-rows" of CLUSTER M-tiles. The host pads // m_tiles to a multiple of CLUSTER, so m_super divides exactly. - const uint32_t rank = cluster_ctarank(); // 0..CLUSTER-1 (= M offset) - const uint32_t cluster_id = blockIdx.x / CLUSTER; - const uint32_t num_clusters = gridDim.x / CLUSTER; - const uint32_t m_super = m_tiles / CLUSTER; + uint32_t rank = 0; // 0..CLU-1 (= M offset) + if constexpr (CLU > 1) rank = cluster_ctarank(); + const uint32_t cluster_id = blockIdx.x / CLU; + const uint32_t num_clusters = gridDim.x / CLU; + const uint32_t m_super = m_tiles / CLU; const uint32_t total_cl = m_super * n_groups; - const uint16_t bmask = (uint16_t)((1u << CLUSTER) - 1u); // all ranks + const uint16_t bmask = (uint16_t)((1u << CLU) - 1u); // all ranks + (void)bmask; // unused at CLU == 1 // Register reallocation is a one-time per-warpgroup action. if (wg == 0) SET_MAXNREG_DEC(PRODUCER_REGS); @@ -291,17 +315,22 @@ extern "C" __global__ void __cluster_dims__(CLUSTER, 1, 1) matmul_b1_kernel( #pragma unroll for (int s = 0; s < STAGES; ++s) { mbar_init(&mbar_full[s], 1); - mbar_init(&mbar_empty[s], CLUSTER); + mbar_init(&mbar_empty[s], CLU); } } __syncthreads(); - cluster_sync(); // all CTAs' barriers initialized before any cross-CTA arrive + // All CTAs' barriers initialized before any cross-CTA arrive. Only needed + // when arrivals actually cross CTAs; at CLU == 1 __syncthreads is sufficient. + if constexpr (CLU > 1) cluster_sync(); - // Pre-arrive every empty barrier cluster-wide so the producer's first - // STAGES `mbar_wait(empty, 0)` succeed immediately (stages logically free). - if (wg == 1 && t_wg < CLUSTER) { + // Pre-arrive every empty barrier so the producer's first STAGES + // `mbar_wait(empty, 0)` succeed immediately (stages logically free). + if (wg == 1 && t_wg < CLU) { #pragma unroll - for (int s = 0; s < STAGES; ++s) arrive_cluster(&mbar_empty[s], t_wg); + for (int s = 0; s < STAGES; ++s) { + if constexpr (CLU > 1) arrive_cluster(&mbar_empty[s], t_wg); + else arrive_local(&mbar_empty[s]); + } } // ===================== PERSISTENT CLUSTER LOOP ===================== @@ -321,7 +350,7 @@ extern "C" __global__ void __cluster_dims__(CLUSTER, 1, 1) matmul_b1_kernel( const uint32_t local = ct - gid * GROUP_M * n_groups; const uint32_t sbi = firstm + local % curm; const int bj = (int)(local / curm); - const int bi = (int)(sbi * CLUSTER + rank); // this CTA's M-tile + const int bi = (int)(sbi * CLU + rank); // this CTA's M-tile const int row0 = bi * TM, col0 = bj * NG; uint64_t* sCb = sC + (titer & 1) * SC_STRIDE; // this tile's sC buffer @@ -353,10 +382,18 @@ extern "C" __global__ void __cluster_dims__(CLUSTER, 1, 1) matmul_b1_kernel( // B: one HBM read, multicast into every cluster member's sB // and counted against every member's full barrier. Issued by // rank 0 only (its mask bit is set, so it fills itself too). - if (rank == 0) { - tma_2d_multicast(&sB[s * TILE_B], &tma_b, 0, - (kk * n_groups + bj) * NB, &mbar_full[s], - bmask); + // Without a cluster there is nobody to share with, so each + // CTA simply loads its own copy — the extra HBM traffic is + // exactly what the cluster variant buys back. + if constexpr (CLU > 1) { + if (rank == 0) { + tma_2d_multicast(&sB[s * TILE_B], &tma_b, 0, + (kk * n_groups + bj) * NB, &mbar_full[s], + bmask); + } + } else { + tma_2d(&sB[s * TILE_B], &tma_b, 0, + (kk * n_groups + bj) * NB, &mbar_full[s]); } } if (++qidx == STAGES) { qidx = 0; p ^= 1; } @@ -402,8 +439,13 @@ extern "C" __global__ void __cluster_dims__(CLUSTER, 1, 1) matmul_b1_kernel( wgmma_wait(); // Release this stage cluster-wide: arrive on every CTA's empty - // barrier (so rank 0 may overwrite their multicast sB). - if (t_wg < CLUSTER) arrive_cluster(&mbar_empty[s], t_wg); + // barrier (so rank 0 may overwrite their multicast sB). Without + // a cluster the stage is this CTA's alone, so a local arrive is + // the whole of the release. + if (t_wg < CLU) { + if constexpr (CLU > 1) arrive_cluster(&mbar_empty[s], t_wg); + else arrive_local(&mbar_empty[s]); + } if (++qidx == STAGES) { qidx = 0; p ^= 1; } } @@ -462,6 +504,44 @@ extern "C" __global__ void __cluster_dims__(CLUSTER, 1, 1) matmul_b1_kernel( if (t == 0) tma_store_wait(); } +// ── Entry points ──────────────────────────────────────────────────────────── +// +// Two kernels, same body, differing only in cluster width: +// +// matmul_b1_kernel CLUSTER-wide clusters + TMA multicast of B. Max +// throughput (8674 binary TOPS at 16384^3 on an idle +// H200), but `__cluster_dims__` makes the cluster's CTAs +// co-resident BY CONSTRUCTION, so the launch demands a +// placement rather than queueing for one. On a GPU shared +// with another tenant that surfaces as a bare +// CUDA_ERROR_LAUNCH_FAILED. Use when this process owns +// the device. +// +// matmul_b1_kernel_nc No clusters, no multicast: an ordinary grid whose CTAs +// are independent, so the launch queues like any other +// and composes with a co-tenant at any grid size. Pays +// for B once per CTA instead of once per cluster. +// +// The host picks between them (`run_gemm_kernel`); the choice is a throughput / +// composability trade, never a correctness one. +extern "C" __global__ void __cluster_dims__(CLUSTER, 1, 1) matmul_b1_kernel( + const __grid_constant__ CUtensorMap tma_a, + const __grid_constant__ CUtensorMap tma_b, + const __grid_constant__ CUtensorMap tma_c, + uint32_t m_tiles, uint32_t n_groups, uint32_t M, uint32_t K) +{ + matmul_b1_body(tma_a, tma_b, tma_c, m_tiles, n_groups, M, K); +} + +extern "C" __global__ void matmul_b1_kernel_nc( + const __grid_constant__ CUtensorMap tma_a, + const __grid_constant__ CUtensorMap tma_b, + const __grid_constant__ CUtensorMap tma_c, + uint32_t m_tiles, uint32_t n_groups, uint32_t M, uint32_t K) +{ + matmul_b1_body<1>(tma_a, tma_b, tma_c, m_tiles, n_groups, M, K); +} + // ── Device-resident packing kernels (BLAS3 GPU row-reduction port) ─────────── // // These reproduce, on device, the host operand pre-arrangement in src/lib.rs @@ -757,6 +837,208 @@ extern "C" __global__ void panel_factor_coop( if (gtid == 0) *pr_out = *g_pr; } +// ── Streamed (kernel-boundary) panel factorization ─────────────────────────── +// +// Same all-SM parallelism as panel_factor_coop, but WITHOUT a cooperative launch: +// each of the ≤ bl·64 sequential bit-steps is three ordinary grid-wide kernels +// (pf_find → pf_swap → pf_xor), and the *kernel boundary* — stream ordering — +// replaces the in-grid `grid_sync`. This is how cuSOLVER/cuBLAS build grid-wide +// multi-step algorithms: no all-CTAs-co-resident requirement, so it composes with +// a concurrent kernel from another CUDA runtime (cubecl's Milnor multiply) instead +// of deadlocking its grid barrier (the intermittent stem-150 wedge). +// +// All state stays on the device — g_pr (pivots so far), g_min (find-first result), +// g_pivpos (this step's pivot position, for pf_xor's guard), g_pivword (the pivot +// row's bl panel limbs). The host never reads back inside the loop, so it races +// ahead queuing launches and their latency hides behind the GPU work. Bit-for-bit +// identical to panel_factor_coop; g_min must be INF and g_pr 0 at entry. + +// find-first + swap, fused into one launch. Every CTA reduces its row slice and +// atomicMin's into g_min; then a threadfence "last-CTA finalize" (the CTA whose +// leader increments the arrival counter last) reads the grid-wide minimum and does +// the swap. This is a grid-wide *reduction*, not a barrier — the last CTA to run +// finalizes, so it needs NO co-residency (unlike a spin barrier) and cannot +// deadlock against concurrent GPU work; `arrival` self-resets to 0 via atomicInc's +// wrap at gridDim-1. On a pivot: read its bl panel limbs into g_pivword, swap it up +// to r+g_pr (perm swap), record the column, bump g_pr, reset g_min for the next +// step. Publishes the pivot position (or INF) to g_pivpos so pf_xor knows whether +// to run. `arrival` and g_min must be 0 / INF at the first step. q ≥ n ⇒ no-op +// (lets a fixed-length step sequence cover a short final panel). +extern "C" __global__ void pf_find_swap( + u64_t* __restrict__ m_buf, + unsigned* __restrict__ perm, + unsigned* __restrict__ pivcols, + u64_t* __restrict__ g_pivword, + int* __restrict__ g_min, + unsigned* __restrict__ g_pr, + int* __restrict__ g_pivpos, + unsigned* __restrict__ arrival, + unsigned ppanel, unsigned bl, unsigned j, unsigned plimb, unsigned cc, + unsigned r, unsigned m, unsigned stride, unsigned n) +{ + extern __shared__ int s_red[]; + const int tid = threadIdx.x; + const int nt = blockDim.x; + const unsigned gtid = blockIdx.x * blockDim.x + threadIdx.x; + const unsigned gnt = gridDim.x * blockDim.x; + const unsigned q = ppanel * 64 + cc; + const unsigned pr = *g_pr; + + int local_min = 0x7fffffff; + if (q < n) { + for (unsigned p = r + pr + gtid; p < m; p += gnt) { + unsigned row = perm[p]; + if ((m_buf[(u64_t)row * stride + plimb] >> j) & 1ULL) + local_min = min(local_min, (int)p); + } + } + s_red[tid] = local_min; + __syncthreads(); + for (int off = nt / 2; off > 0; off >>= 1) { + if (tid < off) s_red[tid] = min(s_red[tid], s_red[tid + off]); + __syncthreads(); + } + if (tid == 0) atomicMin(g_min, s_red[0]); + __threadfence(); + + __shared__ bool am_last; + if (tid == 0) am_last = (atomicInc(arrival, gridDim.x - 1) == gridDim.x - 1); + __syncthreads(); + if (!am_last || tid != 0) return; + + int pivpos = *g_min; + *g_pivpos = pivpos; + if (pivpos == 0x7fffffff) return; // free column: g_min stays INF for next step + unsigned pivrow = perm[pivpos]; + for (unsigned t = 0; t < bl; ++t) + g_pivword[t] = m_buf[(u64_t)pivrow * stride + ppanel + t]; + unsigned a = r + pr; + perm[pivpos] = perm[a]; + perm[a] = pivrow; + pivcols[pr] = q; + *g_min = 0x7fffffff; // reset for the next column + *g_pr = pr + 1; +} + +// Fused lookahead step: clear the PREVIOUS column (cc-1) from the below rows AND +// find+swap the pivot of the CURRENT column (cc), in one launch. Because the +// forward sweep alternates xor(col j) then find(col j+1) over the *same* below-row +// range, fusing them halves the panel factor's launches and — since each thread +// owns the same rows in both phases (grid-stride) — lets it read each row once and +// see its own XOR before scanning, cutting memory traffic. Correctness rests on: +// (A) the previous pivot sits above the shared below-row range, (B) g_pivword / +// g_pivpos / g_pr are read by every CTA in phase A before the last-CTA finalize +// overwrites them (the arrival counter orders all phase-A reads before the single +// finalize write). g_min INF, arrival 0 on entry; q ≥ n ⇒ find is skipped. +extern "C" __global__ void pf_step( + u64_t* __restrict__ m_buf, + unsigned* __restrict__ perm, + u64_t* __restrict__ l_buf, + unsigned* __restrict__ pivcols, + u64_t* __restrict__ g_pivword, + int* __restrict__ g_min, + unsigned* __restrict__ g_pr, + int* __restrict__ g_pivpos, + unsigned* __restrict__ arrival, + unsigned ppanel, unsigned bl, unsigned cc, + unsigned r, unsigned m, unsigned stride, unsigned l_stride, unsigned n) +{ + extern __shared__ int s_red[]; + const int tid = threadIdx.x; + const int nt = blockDim.x; + const unsigned gtid = blockIdx.x * blockDim.x + threadIdx.x; + const unsigned gnt = gridDim.x * blockDim.x; + + const unsigned pr_now = *g_pr; // pivots found through column cc-1 + const int prev_pivpos = *g_pivpos; // pivot position of column cc-1 (INF = free) + + // ── Phase A: clear column cc-1 from the below rows [r+pr_now, m) ── + // (skipped if cc-1 was a free column). The previous pivot is at r+pr_now-1, + // above this range, so it is untouched. + if (prev_pivpos != 0x7fffffff) { + const unsigned prev_cc = cc - 1; + const unsigned pj = prev_cc & 63; + const unsigned prev_pr = pr_now - 1; // L index of the previous pivot + for (unsigned p = r + pr_now + gtid; p < m; p += gnt) { + unsigned row = perm[p]; + u64_t* base = &m_buf[(u64_t)row * stride + ppanel]; + if ((base[prev_cc / 64] >> pj) & 1ULL) { + l_buf[(u64_t)row * l_stride + (prev_pr >> 6)] |= (1ULL << (prev_pr & 63)); + for (unsigned t = 0; t < bl; ++t) + base[t] ^= g_pivword[t]; + } + } + } + + // ── Phase B: find-first for column cc over the same below rows ── + // Same thread owns the same rows as phase A, so its XORs are visible here. + const unsigned q = ppanel * 64 + cc; + const unsigned plimb = ppanel + cc / 64; + const unsigned j = cc & 63; + int local_min = 0x7fffffff; + if (q < n) { + for (unsigned p = r + pr_now + gtid; p < m; p += gnt) { + unsigned row = perm[p]; + if ((m_buf[(u64_t)row * stride + plimb] >> j) & 1ULL) + local_min = min(local_min, (int)p); + } + } + s_red[tid] = local_min; + __syncthreads(); + for (int off = nt / 2; off > 0; off >>= 1) { + if (tid < off) s_red[tid] = min(s_red[tid], s_red[tid + off]); + __syncthreads(); + } + if (tid == 0) atomicMin(g_min, s_red[0]); + __threadfence(); + + __shared__ bool am_last; + if (tid == 0) am_last = (atomicInc(arrival, gridDim.x - 1) == gridDim.x - 1); + __syncthreads(); + if (!am_last || tid != 0) return; + + int pivpos = *g_min; + *g_pivpos = pivpos; + if (pivpos == 0x7fffffff) return; // free column + unsigned pivrow = perm[pivpos]; + for (unsigned t = 0; t < bl; ++t) + g_pivword[t] = m_buf[(u64_t)pivrow * stride + ppanel + t]; + unsigned a = r + pr_now; + perm[pivpos] = perm[a]; + perm[a] = pivrow; + pivcols[pr_now] = q; + *g_min = 0x7fffffff; + *g_pr = pr_now + 1; +} + +// masked XOR of the pivot row into the rows *below* it, across all bl panel limbs, +// recording the multiplier bit into L. No-op on a free column (g_pivpos == INF). +// g_pr has already been bumped by pf_swap, so this pivot's index is *g_pr - 1. +extern "C" __global__ void pf_xor( + u64_t* __restrict__ m_buf, + const unsigned* __restrict__ perm, + u64_t* __restrict__ l_buf, + const u64_t* __restrict__ g_pivword, + const int* __restrict__ g_pivpos, + const unsigned* __restrict__ g_pr, + unsigned ppanel, unsigned bl, unsigned cc, unsigned j, + unsigned r, unsigned m, unsigned stride, unsigned l_stride) +{ + if (*g_pivpos == 0x7fffffff) return; // free column: nothing to clear + const unsigned pr = *g_pr - 1; // index of the pivot just placed + const unsigned gtid = blockIdx.x * blockDim.x + threadIdx.x; + const unsigned gnt = gridDim.x * blockDim.x; + for (unsigned p = r + pr + 1 + gtid; p < m; p += gnt) { + unsigned row = perm[p]; + u64_t* base = &m_buf[(u64_t)row * stride + ppanel]; + if ((base[cc / 64] >> j) & 1ULL) { + l_buf[(u64_t)row * l_stride + (pr >> 6)] |= (1ULL << (pr & 63)); + for (unsigned t = 0; t < bl; ++t) + base[t] ^= g_pivword[t]; + } + } +} + // ── Active-row compaction (design §8.2) ────────────────────────────────────── // // A below row that is entirely zero across the remaining columns [start_limb, @@ -976,6 +1258,50 @@ extern "C" __global__ void block_reduce_coop( } } +// ── Streamed (kernel-boundary) block reduction ─────────────────────────────── +// +// Non-cooperative equivalent of block_reduce_coop: the same grid-wide per-pivot +// clear, but each of block_reduce_coop's two grid_syncs becomes a kernel boundary +// (br_cond → br_xor per pivot k, high-to-low). No cooperative launch, so it +// composes with concurrent GPU work. `cond` holds ≥ (e-s) unsigned; both are +// launched per pivot with the same block-relative index k. Bit-identical to +// block_reduce_coop / block_reduce_rref. + +// Gather the pivot-k bit of every earlier block row j ∈ [s, k) into cond[j-s], +// *before* any XOR clears it. Grid-strided over the ≤64 earlier rows. +extern "C" __global__ void br_cond( + const u64_t* __restrict__ m_buf, const unsigned* __restrict__ perm, + const unsigned* __restrict__ pivcols, + unsigned s, unsigned k, unsigned stride, unsigned* __restrict__ cond) +{ + unsigned qk = pivcols[k]; + unsigned qlimb = qk >> 6, qbit = qk & 63; + unsigned nj = k - s; + const unsigned gtid = blockIdx.x * blockDim.x + threadIdx.x; + const unsigned gnt = gridDim.x * blockDim.x; + for (unsigned j = gtid; j < nj; j += gnt) + cond[j] = (unsigned)((m_buf[(u64_t)perm[s + j] * stride + qlimb] >> qbit) & 1ULL); +} + +// XOR row k into every flagged earlier block row across all limbs, flattened over +// (j, limb) across the grid. +extern "C" __global__ void br_xor( + u64_t* __restrict__ m_buf, const unsigned* __restrict__ perm, + unsigned s, unsigned k, unsigned stride, const unsigned* __restrict__ cond) +{ + unsigned rowk = perm[k]; + unsigned nj = k - s; + unsigned total = nj * stride; + const unsigned gtid = blockIdx.x * blockDim.x + threadIdx.x; + const unsigned gnt = gridDim.x * blockDim.x; + for (unsigned idx = gtid; idx < total; idx += gnt) { + unsigned j = idx / stride; + if (!cond[j]) continue; + unsigned c = idx - j * stride; + m_buf[(u64_t)perm[s + j] * stride + c] ^= m_buf[(u64_t)rowk * stride + c]; + } +} + // (2a) Gather X: for rows at perm positions [0, s), the bits at the `count` // block pivot columns pivcols[col_start .. col_start+count). One thread per // (row, dst-limb) builds a full limb, so no atomics. dst is s × dst_stride. diff --git a/ext/crates/fp-cuda/src/lib.rs b/ext/crates/fp-cuda/src/lib.rs index 9a574ddcc9..82f6ef7cbe 100644 --- a/ext/crates/fp-cuda/src/lib.rs +++ b/ext/crates/fp-cuda/src/lib.rs @@ -40,6 +40,57 @@ fn adaptive_bl(stride: usize) -> usize { (stride / div).clamp(1, 16) } +/// Whether the row reduction uses its **cooperative** kernels — `panel_factor_coop`, +/// `promote_coop`, `block_reduce_coop` — launched with `cuLaunchCooperativeKernel` +/// and synchronized by a hand-rolled grid-wide spin barrier. +/// +/// The cooperative launch requires **all** the grid's CTAs to be co-resident at once +/// (the barrier spins waiting for every CTA to arrive). That holds only when this +/// process owns the whole GPU: a kernel from another CUDA runtime sharing the device +/// — e.g. `cubecl`'s Milnor multiply in the `algebra` crate — can occupy SMs and +/// prevent co-residency, so the missing CTAs never reach the barrier and the resident +/// ones spin forever (the intermittent stem-150 wedge, flat sm=100%). +/// +/// **Off by default**, so the reduction composes safely with concurrent GPU work. +/// The default path keeps the cooperative kernels' all-SM parallelism but replaces +/// their in-grid `grid_sync` with kernel-boundary (stream-ordered) synchronization: +/// the forward pass runs `pf_find_swap` then a fused `pf_step` per column (each a +/// grid-wide reduction finalized by the last CTA to arrive — no barrier, no co- +/// residency), promotion uses the grid-strided `promote_pivots`, and back- +/// substitution the streamed `br_cond`/`br_xor` pair (single-CTA `block_reduce_rref` +/// below stride 1024). None launch cooperatively, so none can deadlock. The residual +/// cost is the forward pass's per-column relaunch vs the persistent cooperative grid: +/// ~2× at n≈2¹⁷, shrinking with size (~1.4× total at 2¹⁸, converging as the O(cols) +/// launch term is dwarfed by the O(cols²) work). Set `FP_CUDA_RR_COOP=1` to opt into +/// the cooperative path on a dedicated GPU. +fn rr_coop() -> bool { + std::env::var("FP_CUDA_RR_COOP") + .map(|v| v != "0" && !v.is_empty()) + .unwrap_or(false) +} + +/// Whether the GEMM uses its **cluster** kernel, `matmul_b1_kernel` +/// (`__cluster_dims__(CLUSTER,1,1)` plus TMA multicast of the B panel). +/// +/// This is the same trade as [`rr_coop`], one layer down. A thread-block cluster is +/// co-resident *by construction*: the hardware will not place one CTA of a cluster +/// without placing them all, because rank 0 multicasts B directly into its mates' +/// shared memory and every consumer arrives on their empty barriers through `mapa`. +/// So when another runtime holds SMs, the launch does not queue for a slot the way an +/// ordinary grid does — it fails outright with `CUDA_ERROR_LAUNCH_FAILED`. +/// +/// **Off by default**, so the GEMM composes with concurrent GPU work at any grid size. +/// The default `matmul_b1_kernel_nc` runs the identical tile schedule and arithmetic +/// with independent CTAs, paying for B once per CTA instead of once per cluster — at +/// most 2× B's HBM traffic, on top of the ~8× that GROUP_M rasterization already saves. +/// Set `FP_CUDA_GEMM_COOP=1` to opt into the cluster kernel on a dedicated GPU, where +/// it is the faster of the two. +fn gemm_coop() -> bool { + std::env::var("FP_CUDA_GEMM_COOP") + .map(|v| v != "0" && !v.is_empty()) + .unwrap_or(false) +} + /// Lets us pass a `CUtensorMap` by value as a (grid-constant) kernel argument /// through cudarc's typed launch builder. `repr(transparent)` so the pointer /// cudarc pushes is the address of the 128-byte descriptor itself. @@ -72,12 +123,17 @@ pub struct GpuContext { #[allow(dead_code)] module: Arc, kernel: CudaFunction, + /// Cluster-free GEMM (`matmul_b1_kernel_nc`): same arithmetic, ordinary grid. + kernel_nc: CudaFunction, // Device-resident packing/epilogue kernels for the row-reduction port. pack_a: CudaFunction, pack_b: CudaFunction, xor_into: CudaFunction, panel_factor: CudaFunction, panel_factor_coop: CudaFunction, + pf_find_swap: CudaFunction, + pf_step: CudaFunction, + pf_xor: CudaFunction, mark_live: CudaFunction, promote_pivots: CudaFunction, promote_coop: CudaFunction, @@ -85,21 +141,32 @@ pub struct GpuContext { gather_rows: CudaFunction, block_reduce_rref: CudaFunction, block_reduce_coop: CudaFunction, + br_cond: CudaFunction, + br_xor: CudaFunction, gather_cols: CudaFunction, xor_into_perm: CudaFunction, } impl GpuContext { pub fn new(device_id: usize) -> Result> { + // NOTE: this retains the device *primary* context, which the cubecl Milnor-multiply runtime + // also retains (`cubecl-cuda/src/runtime.rs`, `primary_ctx::retain`), so both CUDA consumers + // share one context. Giving the row reduction its own non-primary context was tried as a fix + // for the cross-runtime `CUDA_ERROR_LAUNCH_FAILED` and did NOT help (3/3 runs still died): + // the fault is device contention, not shared context state. See [`fp::gpu_lock`]. let ctx = CudaContext::new(device_id)?; let ptx = Ptx::from_src(String::from_utf8(PTX_IMAGE.to_vec())?); let module = ctx.load_module(ptx)?; let kernel = module.load_function("matmul_b1_kernel")?; + let kernel_nc = module.load_function("matmul_b1_kernel_nc")?; let pack_a = module.load_function("pack_a")?; let pack_b = module.load_function("pack_b")?; let xor_into = module.load_function("xor_into")?; let panel_factor = module.load_function("panel_factor")?; let panel_factor_coop = module.load_function("panel_factor_coop")?; + let pf_find_swap = module.load_function("pf_find_swap")?; + let pf_step = module.load_function("pf_step")?; + let pf_xor = module.load_function("pf_xor")?; let mark_live = module.load_function("mark_live")?; let promote_pivots = module.load_function("promote_pivots")?; let promote_coop = module.load_function("promote_coop")?; @@ -107,17 +174,23 @@ impl GpuContext { let gather_rows = module.load_function("gather_rows")?; let block_reduce_rref = module.load_function("block_reduce_rref")?; let block_reduce_coop = module.load_function("block_reduce_coop")?; + let br_cond = module.load_function("br_cond")?; + let br_xor = module.load_function("br_xor")?; let gather_cols = module.load_function("gather_cols")?; let xor_into_perm = module.load_function("xor_into_perm")?; Ok(Self { ctx, module, kernel, + kernel_nc, pack_a, pack_b, xor_into, panel_factor, panel_factor_coop, + pf_find_swap, + pf_step, + pf_xor, mark_live, promote_pivots, promote_coop, @@ -125,6 +198,8 @@ impl GpuContext { gather_rows, block_reduce_rref, block_reduce_coop, + br_cond, + br_xor, gather_cols, xor_into_perm, }) @@ -144,6 +219,36 @@ impl GpuContext { self.ctx.default_stream() } + /// A CUDA stream **private to the calling OS thread**, created lazily on first use and reused + /// thereafter. Every row-reduction method submits through this instead of the context's single + /// `default_stream()`, so work from different rayon workers runs on distinct streams — + /// overlapping transfers and kernels concurrently instead of serializing — while every + /// sub-launch of one reduce shares one stream (correct ordering within a thread). This is what + /// lets `try_row_reduce` run lock-free from many threads at once. + /// + /// Assumes a single process-wide `GpuContext` (the `OnceLock` in `fp::blas::cuda`): the + /// thread-local caches the stream by thread, not by context, so the first context to call this + /// on a given thread owns that thread's stream. With one context that is always correct. + pub fn stream(&self) -> Arc { + use std::cell::RefCell; + thread_local! { + static TLS: RefCell>> = const { RefCell::new(None) }; + } + TLS.with(|cell| { + let mut slot = cell.borrow_mut(); + // If stream creation fails (e.g. the context is already poisoned by another runtime's + // launch failure), fall back to the context default stream rather than panicking: the + // subsequent op then fails as a normal `Err`, which `try_row_reduce` turns into a CPU + // fallback instead of crashing the process. + slot.get_or_insert_with(|| { + self.ctx + .new_stream() + .unwrap_or_else(|_| self.ctx.default_stream()) + }) + .clone() + }) + } + pub fn kernel(&self) -> &CudaFunction { &self.kernel } @@ -218,7 +323,7 @@ fn matmul_b1_inner( let n_groups = n_lim.div_ceil(NG as usize); let n_padded_lim = n_groups * NG as usize; - let stream = gpu.ctx.default_stream(); + let stream = gpu.stream(); let a_padded = pad_2d(a, m, k.div_ceil(64), m_padded, k_padded / 64); let b_padded = pad_2d(b, k, n_lim, k_padded, n_lim); @@ -315,8 +420,17 @@ fn run_gemm_kernel( let smem_u64 = STAGES * tile_a + STAGES * tile_b + 2 * NG as usize * TILE_M + 2 * STAGES; let smem_bytes = (smem_u64 * std::mem::size_of::()) as u32; + // Which GEMM variant runs. The cluster kernel is faster but its + // `__cluster_dims__` requires CLUSTER co-resident CTAs, which a launch onto a + // GPU somebody else is using cannot get -- see [`gemm_coop`]. Default is the + // composable one. + let coop = gemm_coop(); + let kf = if coop { &gpu.kernel } else { &gpu.kernel_nc }; + // CTA granule the grid must be a multiple of: a whole cluster, or nothing. + let gran = if coop { CLUSTER as u32 } else { 1 }; + // Opt in to >48 KB shared memory (Hopper static default cap). - gpu.kernel.set_attribute( + kf.set_attribute( sys::CUfunction_attribute_enum::CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, smem_bytes as i32, )?; @@ -340,11 +454,31 @@ fn run_gemm_kernel( .ctx .attribute(sys::CUdevice_attribute_enum::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)? as u32; - let occ = gpu - .kernel + let occ = kf .occupancy_max_active_blocks_per_multiprocessor(THREADS, smem_bytes as usize, None)? .max(1); - let mut num_ctas = (occ * sms / CLUSTER as u32).max(1) * CLUSTER as u32; + // How much of the machine to ask for. Under the cluster kernel a grid sized to full occupancy + // is only placeable on a GPU this process owns outright: when anything else holds SMs — the + // cubecl Milnor multiply in `algebra`, or simply another tenant — the launch is not queued, it + // fails, as a bare `CUDA_ERROR_LAUNCH_FAILED` that compute-sanitizer cannot attribute (0 invalid + // accesses across a whole run: it was never a memory bug). Shrinking the grid reduces the + // collision probability but never proves it to zero — the safe size is not a sharp threshold + // (on the theta=125 stem-200 workload 1/16 ran clean while 1/8 failed 74 times), and it costs + // most of the throughput (`bench_kernel_only`, 16384^3, idle H200: 1062 binary TOPS at 16 CTAs + // and 2107 at 32, against 8674 at full grid). + // + // The cluster-free kernel has no such constraint — its CTAs are independent, so the launch + // queues like any other — and therefore takes the whole machine by default. That is the point of + // it: composability without paying the share. + // + // `FP_CUDA_GEMM_DEVICE_FRAC` overrides either default (1 = whole machine). Under the cluster + // kernel treat any value as a risk setting, not a guarantee. + let frac = std::env::var("FP_CUDA_GEMM_DEVICE_FRAC") + .ok() + .and_then(|v| v.parse::().ok()) + .filter(|&f| f > 0) + .unwrap_or(if coop { 16 } else { 1 }); + let mut num_ctas = ((occ * sms / frac) / gran).max(1) * gran; // Diagnostic: cap the persistent grid to probe how much of a small GEMM's // time is the persistent-grid startup (cluster sync + mbar init + pipeline // fill across occ×SMs CTAs). The persistent loop handles any multiple of @@ -353,7 +487,7 @@ fn run_gemm_kernel( .ok() .and_then(|v| v.parse::().ok()) { - num_ctas = (cap / CLUSTER as u32).max(1) * CLUSTER as u32; + num_ctas = (cap / gran).max(1) * gran; } let ta = TmaArg(tma_a); @@ -370,7 +504,7 @@ fn run_gemm_kernel( block_dim: (THREADS, 1, 1), shared_mem_bytes: smem_bytes, }; - let mut lb = stream.launch_builder(&gpu.kernel); + let mut lb = stream.launch_builder(kf); lb.arg(&ta) .arg(&tb) .arg(&tc) @@ -463,7 +597,7 @@ impl GpuContext { let n_groups = n_lim.div_ceil(NG as usize); let n_padded_lim = n_groups * NG as usize; - let stream = self.ctx.default_stream(); + let stream = self.stream(); // Pack A → interleaved row-major K-major tiles (m_padded × k_padded/64). // pack_a/pack_b/the GEMM fully overwrite these buffers (padding written as @@ -578,7 +712,7 @@ impl GpuContext { n: usize, ) -> Result, Box> { let n_lim = n.div_ceil(64); - let stream = self.ctx.default_stream(); + let stream = self.stream(); let a_dev = stream.clone_htod(a)?; let b_dev = stream.clone_htod(b)?; let (c_dev, n_padded_lim) = self.matmul_b1_dev(&a_dev, m, k, &b_dev, n)?; @@ -600,7 +734,7 @@ impl GpuContext { ) -> Result> { let stride = cols.div_ceil(64); assert_eq!(data.len(), rows * stride, "limb count mismatch"); - let buf = self.ctx.default_stream().clone_htod(data)?; + let buf = self.stream().clone_htod(data)?; Ok(DeviceMatrix { buf, rows, @@ -611,12 +745,12 @@ impl GpuContext { /// Download a [`DeviceMatrix`] back to host limbs (natural layout). One D2H. pub fn download(&self, dm: &DeviceMatrix) -> Result, Box> { - Ok(self.ctx.default_stream().clone_dtoh(&dm.buf)?) + Ok(self.stream().clone_dtoh(&dm.buf)?) } /// Download a device `u32` buffer (e.g. a `perm` vector) to host. pub fn download_u32(&self, s: &CudaSlice) -> Result, Box> { - Ok(self.ctx.default_stream().clone_dtoh(s)?) + Ok(self.stream().clone_dtoh(s)?) } /// The fused trailing-update / back-substitution epilogue over persistent @@ -653,7 +787,7 @@ impl GpuContext { } let (c_dev, _n_padded_lim) = self.matmul_b1_dev(&l.buf, m, k, &u.buf, t)?; let width = t.div_ceil(64); // == dst.stride - col_off/64 - let stream = self.ctx.default_stream(); + let stream = self.stream(); self.xor_into_region( &stream, &mut dst.buf, @@ -673,7 +807,7 @@ impl GpuContext { /// are `perm` swaps, so the matrix bytes never move. pub fn identity_perm(&self, m: usize) -> Result, Box> { let host: Vec = (0..m as u32).collect(); - Ok(self.ctx.default_stream().clone_htod(&host)?) + Ok(self.stream().clone_htod(&host)?) } /// Factor one 64-bit column panel (limb `plimb`) in place over the @@ -698,7 +832,7 @@ impl GpuContext { ) -> Result<(usize, Vec), Box> { assert_eq!(perm.len(), m.rows, "perm length must equal rows"); assert_eq!(l.rows, m.rows, "L rows must equal M rows"); - let stream = self.ctx.default_stream(); + let stream = self.stream(); const THREADS: u32 = 256; let pivcols = stream.alloc_zeros::(64)?; @@ -758,7 +892,7 @@ impl GpuContext { assert_eq!(l.rows, m.rows, "L rows must equal M rows"); assert!(l.stride >= bl, "L stride must be at least bl"); assert!(m_active <= m.rows && m_active >= r, "m_active out of range"); - let stream = self.ctx.default_stream(); + let stream = self.stream(); const THREADS: u32 = 256; let smem = THREADS * std::mem::size_of::() as u32; @@ -830,6 +964,169 @@ impl GpuContext { Ok((pr, cols[..pr].to_vec())) } + /// **Streamed** (kernel-boundary) equivalent of + /// [`panel_factor_coop`](Self::panel_factor_coop): identical math and the same + /// all-SM parallelism, but each of the ≤ `bl·64` sequential bit-steps is three + /// ordinary grid-wide launches (`pf_find` → `pf_swap` → `pf_xor`) whose stream + /// ordering replaces the cooperative kernel's in-grid `grid_sync`. No + /// `cuLaunchCooperativeKernel`, so no all-CTAs-co-resident requirement — it + /// composes with a concurrent kernel from another CUDA runtime instead of + /// deadlocking the grid barrier. All per-step state (`g_pr`, `g_min`, + /// `g_pivpos`, `g_pivword`) lives on the device, so the host issues every launch + /// without a readback and their latency hides behind the GPU work; only the + /// final `(pr, pivcols)` is copied back. Bit-for-bit equal to the coop kernel. + #[allow(clippy::too_many_arguments)] + pub fn panel_factor_streamed( + &self, + m: &mut DeviceMatrix, + perm: &mut CudaSlice, + l: &mut DeviceMatrix, + ppanel: usize, + bl: usize, + r: usize, + m_active: usize, + ) -> Result<(usize, Vec), Box> { + assert_eq!(perm.len(), m.rows, "perm length must equal rows"); + assert_eq!(l.rows, m.rows, "L rows must equal M rows"); + assert!(l.stride >= bl, "L stride must be at least bl"); + assert!(m_active <= m.rows && m_active >= r, "m_active out of range"); + let stream = self.stream(); + + const THREADS: u32 = 256; + const INF: i32 = 0x7fff_ffff; + let smem = THREADS * std::mem::size_of::() as u32; + + // Device-resident per-step state: pivot count, find-first result (INF = + // none), this step's pivot position (pf_xor's guard), the pivot row's bl + // panel limbs, and the last-CTA-finalize arrival counter. g_min starts INF; + // g_pr and arrival start 0 (arrival self-resets to 0 each step). + let pivcols = stream.alloc_zeros::(bl * 64)?; + let mut g_pr = stream.alloc_zeros::(1)?; + let mut g_min = stream.clone_htod(&[INF])?; + let g_pivpos = stream.alloc_zeros::(1)?; + let mut g_pivword = stream.alloc_zeros::(bl)?; + let mut arrival = stream.alloc_zeros::(1)?; + + // Regular launches wave-schedule, so the grid can be sized purely to cover + // the active rows once; cap at occ×SMs for launch efficiency. + let sms = self + .ctx + .attribute(sys::CUdevice_attribute_enum::CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT)? + as u32; + let occ = self + .pf_xor + .occupancy_max_active_blocks_per_multiprocessor(THREADS, 0, None)? + .max(1); + let rows_worth = (m_active as u32).div_ceil(THREADS).max(1); + let num_ctas = (occ * sms).min(rows_worth).max(1); + // Each step's grid-wide min-reduce + last-CTA finalize contends on g_min / + // arrival across all CTAs, so — like the cooperative kernel's FP_CUDA_PF_CTAS + // — a smaller grid makes every step cheaper once it still covers the rows. + // Cap at 128 (H200 sweet spot); overridable. + let num_ctas = std::env::var("FP_CUDA_PF_CTAS") + .ok() + .and_then(|v| v.parse::().ok()) + .unwrap_or(128) + .clamp(1, num_ctas); + + let (r_u, m_u, n_u, stride_u, l_stride_u, ppanel_u, bl_u) = ( + r as u32, + m_active as u32, + m.cols as u32, + m.stride as u32, + l.stride as u32, + ppanel as u32, + bl as u32, + ); + let find_cfg = LaunchConfig { + grid_dim: (num_ctas, 1, 1), + block_dim: (THREADS, 1, 1), + shared_mem_bytes: smem, + }; + let grid_cfg = LaunchConfig { + grid_dim: (num_ctas, 1, 1), + block_dim: (THREADS, 1, 1), + shared_mem_bytes: 0, + }; + + // Valid columns in this panel (the last panel may be short). + let ncols = (bl * 64).min(m.cols - ppanel * 64); + + // (0) find+swap the first column's pivot (no trailing XOR yet). + { + let (plimb_u, j_u, cc_u) = (ppanel_u, 0u32, 0u32); + let mut lb = stream.launch_builder(&self.pf_find_swap); + lb.arg(&mut m.buf) + .arg(&mut *perm) + .arg(&pivcols) + .arg(&mut g_pivword) + .arg(&mut g_min) + .arg(&mut g_pr) + .arg(&g_pivpos) + .arg(&mut arrival) + .arg(&ppanel_u) + .arg(&bl_u) + .arg(&j_u) + .arg(&plimb_u) + .arg(&cc_u) + .arg(&r_u) + .arg(&m_u) + .arg(&stride_u) + .arg(&n_u); + unsafe { lb.launch(find_cfg) }?; + } + + // (1) fused lookahead: each step clears column cc-1 and finds+swaps column + // cc — one launch per column instead of two, with the row read once. + for cc in 1..ncols { + let cc_u = cc as u32; + let mut lb = stream.launch_builder(&self.pf_step); + lb.arg(&mut m.buf) + .arg(&mut *perm) + .arg(&mut l.buf) + .arg(&pivcols) + .arg(&mut g_pivword) + .arg(&mut g_min) + .arg(&mut g_pr) + .arg(&g_pivpos) + .arg(&mut arrival) + .arg(&ppanel_u) + .arg(&bl_u) + .arg(&cc_u) + .arg(&r_u) + .arg(&m_u) + .arg(&stride_u) + .arg(&l_stride_u) + .arg(&n_u); + unsafe { lb.launch(find_cfg) }?; + } + + // (2) clear the final column's pivot from the rows below. + { + let (cc_u, j_u) = ((ncols - 1) as u32, ((ncols - 1) & 63) as u32); + let mut lb = stream.launch_builder(&self.pf_xor); + lb.arg(&mut m.buf) + .arg(&*perm) + .arg(&mut l.buf) + .arg(&g_pivword) + .arg(&g_pivpos) + .arg(&g_pr) + .arg(&ppanel_u) + .arg(&bl_u) + .arg(&cc_u) + .arg(&j_u) + .arg(&r_u) + .arg(&m_u) + .arg(&stride_u) + .arg(&l_stride_u); + unsafe { lb.launch(grid_cfg) }?; + } + + let pr = stream.clone_dtoh(&g_pr)?[0] as usize; + let cols = stream.clone_dtoh(&pivcols)?; + Ok((pr, cols[..pr].to_vec())) + } + /// Active-row compaction (design §8.2): mark the below rows [r, m_active) /// that are entirely zero across the remaining columns [start_limb·64, n) — /// permanently dead (they can never pivot and carry no multiplier) — and @@ -849,7 +1146,7 @@ impl GpuContext { if m_active <= r { return Ok(m_active); } - let stream = self.ctx.default_stream(); + let stream = self.stream(); let n_scan = m_active - r; let live = unsafe { stream.alloc::(n_scan) }?; { @@ -911,7 +1208,7 @@ impl GpuContext { if pr == 0 || trailing_limbs == 0 { return Ok(()); } - let stream = self.ctx.default_stream(); + let stream = self.stream(); stream.memset_zeros(pc_barrier)?; let (r_u, pr_u, fl, tl, st, ls, llo, tc) = ( r_piv as u32, @@ -968,7 +1265,7 @@ impl GpuContext { pc_cond: &CudaSlice, pc_ctas: u32, ) -> Result<(), Box> { - let stream = self.ctx.default_stream(); + let stream = self.stream(); let (rows, stride, n) = (m.rows, m.stride, m.cols); let trailing_limbs = end_limb - first_limb; if pr == 0 || trailing_limbs == 0 { @@ -1071,15 +1368,23 @@ impl GpuContext { &self, m: &mut DeviceMatrix, ) -> Result<(CudaSlice, usize, Vec), Box> { - let stream = self.ctx.default_stream(); + let stream = self.stream(); let (rows, stride) = (m.rows, m.stride); let mut perm = self.identity_perm(rows)?; let mut r = 0usize; let mut pivot_cols = Vec::new(); + // Cooperative vs. composable kernels (see [`rr_coop`]). The non-cooperative + // path never launches a cooperative grid, so it composes with concurrent GPU + // work at the cost of the single-CTA panel factor; the cooperative path is the + // faster exclusive-GPU mode. + let coop = rr_coop(); + // Panel width in limbs (b = 64·bl columns). Wider panels raise the // trailing GEMM's contraction dimension pr toward b, reclaiming the ~16× // K-padding waste. Override with FP_CUDA_BL; otherwise adaptive_bl picks - // the measured optimum (flat at bl≈12–16). + // the measured optimum (flat at bl≈12–16). Both the cooperative and the + // streamed (non-cooperative) panel factor handle wide panels, so bl is + // chosen the same way in either mode. let bl = if let Some(v) = std::env::var("FP_CUDA_BL") .ok() .and_then(|v| v.parse::().ok()) @@ -1098,7 +1403,7 @@ impl GpuContext { // Cooperative multi-CTA promotion (right-looking) replaces the single-CTA // triangular replay when the matrix is wide enough to amortize the grid // barriers; otherwise the grid-strided promote_pivots kernel is used. - let use_promote_coop = stride >= 1024; + let use_promote_coop = coop && stride >= 1024; let (pc_ctas, mut pc_barrier, pc_cond) = if use_promote_coop { let sms = self .ctx @@ -1147,8 +1452,13 @@ impl GpuContext { cols: bl_eff * 64, stride: bl_eff, }; - let (pr, pivcols) = - self.panel_factor_coop(m, &mut perm, &mut l, ppanel, bl_eff, r, m_active)?; + let (pr, pivcols) = if coop { + self.panel_factor_coop(m, &mut perm, &mut l, ppanel, bl_eff, r, m_active)? + } else { + // Multi-SM factor via kernel-boundary sync — same math and grid + // parallelism as the coop kernel, no cooperative launch. + self.panel_factor_streamed(m, &mut perm, &mut l, ppanel, bl_eff, r, m_active)? + }; if pr > 0 { for &q in &pivcols { pivot_cols.push(q as usize); @@ -1202,7 +1512,7 @@ impl GpuContext { if above_count == 0 || block_e <= block_s { return Ok(()); } - let stream = self.ctx.default_stream(); + let stream = self.stream(); let (stride, n) = (m.stride, m.cols); let bp_eff = block_e - block_s; let start_limb = pivot_cols[block_s] / 64; @@ -1296,6 +1606,7 @@ impl GpuContext { s: usize, e: usize, use_coop: bool, + streamed: bool, br_barrier: &mut CudaSlice, br_cond: &CudaSlice, br_ctas: u32, @@ -1303,7 +1614,7 @@ impl GpuContext { if e <= s { return Ok(()); } - let stream = self.ctx.default_stream(); + let stream = self.stream(); let stride = m.stride; { if use_coop { @@ -1325,6 +1636,44 @@ impl GpuContext { .arg(br_cond) .arg(&tc); unsafe { lb.launch_cooperative(cfg) }?; + } else if streamed { + // Kernel-boundary equivalent of block_reduce_coop: per pivot k + // (high-to-low), br_cond gathers the clear-conditions, then br_xor + // clears row k from the flagged block rows across the grid. Stream + // order replaces the cooperative grid barrier. + let st = stride as u32; + let xor_cfg = LaunchConfig { + grid_dim: (br_ctas, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }; + let mut k = e; + while k > s { + k -= 1; + let (s_u, k_u) = (s as u32, k as u32); + let nj = (k - s) as u32; + { + let mut lb = stream.launch_builder(&self.br_cond); + lb.arg(&m.buf) + .arg(perm) + .arg(piv_dev) + .arg(&s_u) + .arg(&k_u) + .arg(&st) + .arg(br_cond); + unsafe { lb.launch(cfg_1d((nj.max(1)) as usize)) }?; + } + { + let mut lb = stream.launch_builder(&self.br_xor); + lb.arg(&mut m.buf) + .arg(perm) + .arg(&s_u) + .arg(&k_u) + .arg(&st) + .arg(br_cond); + unsafe { lb.launch(xor_cfg) }?; + } + } } else { let (s_u, e_u, st) = (s as u32, e as u32, stride as u32); let cfg = LaunchConfig { @@ -1368,24 +1717,27 @@ impl GpuContext { e: usize, base_bp: usize, use_coop: bool, + streamed: bool, br_barrier: &mut CudaSlice, br_cond: &CudaSlice, br_ctas: u32, ) -> Result<(), Box> { if e - s <= base_bp { return self.block_reduce_elem( - m, perm, piv_dev, s, e, use_coop, br_barrier, br_cond, br_ctas, + m, perm, piv_dev, s, e, use_coop, streamed, br_barrier, br_cond, br_ctas, ); } let mid = s + (e - s) / 2; // Right half to RREF, then clear its pivots from the left half (K = e-mid). self.block_reduce_rec( - m, perm, piv_dev, pivot_cols, mid, e, base_bp, use_coop, br_barrier, br_cond, br_ctas, + m, perm, piv_dev, pivot_cols, mid, e, base_bp, use_coop, streamed, br_barrier, br_cond, + br_ctas, )?; self.bs_clear_above(m, perm, piv_dev, pivot_cols, mid, e, s, mid - s)?; // Left half to RREF (its rows now carry no right-half pivot bits). self.block_reduce_rec( - m, perm, piv_dev, pivot_cols, s, mid, base_bp, use_coop, br_barrier, br_cond, br_ctas, + m, perm, piv_dev, pivot_cols, s, mid, base_bp, use_coop, streamed, br_barrier, br_cond, + br_ctas, )?; Ok(()) } @@ -1407,25 +1759,32 @@ impl GpuContext { if r == 0 { return Ok(()); } - let stream = self.ctx.default_stream(); + let stream = self.stream(); let stride = m.stride; let piv_dev = stream.clone_htod(&pivot_cols.iter().map(|&q| q as u32).collect::>())?; - // Cooperative multi-CTA block reduction: spreads each block's per-pivot - // clear across the whole grid. The per-block cooperative launch + grid - // barriers only pay once the block work (≈ bp·stride) is large, so gate on - // a wide matrix; below that the single-CTA kernel wins. Measured (H200): - // neutral at n=2¹⁵ (stride 512), +6% at 2¹⁶, +18% at 2¹⁷. - let use_coop = stride >= 1024; + // Multi-CTA block reduction spreads each block's per-pivot clear across the + // whole grid. It only pays once the block work (≈ bp·stride) is large, so + // gate on a wide matrix; below that the single-CTA kernel wins. Measured + // (H200): neutral at n=2¹⁵ (stride 512), +6% at 2¹⁶, +18% at 2¹⁷. + // + // Two grid-parallel variants (see [`rr_coop`]): `use_coop` = the cooperative + // block_reduce_coop (dedicated GPU); `streamed` = the kernel-boundary + // br_cond/br_xor pair, which composes with concurrent GPU work and is the + // default. Below the width gate both fall back to the single-CTA + // block_reduce_rref. + let wide = stride >= 1024; + let use_coop = rr_coop() && wide; + let streamed = !rr_coop() && wide; // Pivots per back-substitution block. Wider blocks raise the X·U GEMM's // contraction dimension bp toward TILE_K, cutting its K-padding waste - // (K=64 pads 16×). block_reduce_coop's cost is ~bp-independent (its - // compute and barrier counts both scale with r, not bp), so on the coop - // path we widen bp for free; the single-CTA fallback keeps bp=64 (its - // shared cond[] is sized 64). Override with FP_CUDA_BP. - let bp = if use_coop { + // (K=64 pads 16×). The grid-parallel base reduces are ~bp-independent (their + // compute and barrier counts scale with r, not bp), so when either fires we + // widen bp for free; the single-CTA fallback keeps bp=64 (its shared cond[] + // is sized 64). Override with FP_CUDA_BP. + let bp = if use_coop || streamed { // K=1024 makes the X·U GEMM's contraction an exact TILE_K multiple — // zero K-padding — and block_reduce_coop is bp-independent. std::env::var("FP_CUDA_BP") @@ -1446,9 +1805,10 @@ impl GpuContext { .block_reduce_coop .occupancy_max_active_blocks_per_multiprocessor(BR_THREADS, 0, None)? .max(1); - // Blocked-TRSM within-block reduce (coop path): recurse each block to - // narrow base blocks + X·U GEMMs. - let use_trsm = use_coop; + // Blocked-TRSM within-block reduce: recurse each block to narrow base blocks + // + X·U GEMMs. The GEMM path composes regardless, so use it whenever a + // grid-parallel base reduce is in play (coop or streamed). + let use_trsm = use_coop || streamed; // Base ≤ 64: the single-CTA block_reduce_rref's shared cond[] is sized 64. let base_bp: usize = std::env::var("FP_CUDA_BS_BASE") .ok() @@ -1488,6 +1848,7 @@ impl GpuContext { e, base_bp, use_coop, + streamed, &mut br_barrier, &br_cond, br_ctas, @@ -1500,6 +1861,7 @@ impl GpuContext { s, e, use_coop, + streamed, &mut br_barrier, &br_cond, br_ctas, @@ -1525,6 +1887,19 @@ impl GpuContext { &self, m: &mut DeviceMatrix, ) -> Result<(CudaSlice, usize, Vec), Box> { + if std::env::var_os("FP_CUDA_RR_TIMING").is_some() { + let t0 = std::time::Instant::now(); + let (perm, r, pivot_cols) = self.forward_reduce(m)?; + let t1 = std::time::Instant::now(); + self.back_substitute(m, &perm, r, &pivot_cols)?; + let t2 = std::time::Instant::now(); + eprintln!( + "[rr_timing] forward={:.3}s back={:.3}s (r={r})", + (t1 - t0).as_secs_f64(), + (t2 - t1).as_secs_f64(), + ); + return Ok((perm, r, pivot_cols)); + } let (perm, r, pivot_cols) = self.forward_reduce(m)?; self.back_substitute(m, &perm, r, &pivot_cols)?; Ok((perm, r, pivot_cols)) diff --git a/ext/crates/fp/Cargo.toml b/ext/crates/fp/Cargo.toml index 52dd296e2a..4f5b106e06 100644 --- a/ext/crates/fp/Cargo.toml +++ b/ext/crates/fp/Cargo.toml @@ -17,6 +17,8 @@ paste = "1.0.15" proptest = { version = "1.7", optional = true } serde = { version = "1.0.219", features = ["derive"] } serde_json = "1.0.141" +# Only used for the GPU row-reduce dispatch instrumentation; pulled in by `gpu`. +tracing = { version = "0.1.41", optional = true } maybe-rayon = { path = "../maybe-rayon" } query = { path = "../query" } @@ -45,7 +47,7 @@ default = ["odd-primes"] concurrent = ["maybe-rayon/concurrent"] odd-primes = [] # Dispatch large p=2 matrix products to the Hopper GPU backend (`fp-cuda`). -gpu = ["dep:fp-cuda"] +gpu = ["dep:fp-cuda", "dep:tracing"] [[bench]] name = "mul" diff --git a/ext/crates/fp/src/blas/cuda.rs b/ext/crates/fp/src/blas/cuda.rs index 0fba7154ec..a02253a42b 100644 --- a/ext/crates/fp/src/blas/cuda.rs +++ b/ext/crates/fp/src/blas/cuda.rs @@ -14,7 +14,7 @@ //! CPU path is used. Defaults to 2048; the GPU only wins once the kernel work //! dwarfs the H2D/D2H + TMA-layout marshalling, which dominates small sizes. -use std::sync::{Mutex, OnceLock}; +use std::sync::OnceLock; use fp_cuda::GpuContext; @@ -52,22 +52,162 @@ fn rr_threshold() -> usize { /// The process-wide GPU context, created lazily on first use. `None` if no /// usable device is present (no driver, no Hopper GPU, or the kernel PTX is the -/// nvcc-absent build stub). Wrapped in a `Mutex` because a single CUDA context -/// serialises submission anyway and `GpuContext` is not shared concurrently. -fn context() -> Option<&'static Mutex> { - static GPU: OnceLock>> = OnceLock::new(); +/// nvcc-absent build stub). Shared as `&'static` — no lock: `GpuContext` is +/// `Send + Sync` (its cudarc handles are), and every submission goes through a +/// per-thread stream ([`GpuContext::stream`]), so concurrent rayon workers run +/// on independent streams (overlapping transfers + kernels) instead of +/// serializing on one mutex. Buffers are per-call and thread-local, so there is +/// no shared mutable device state to guard. +fn context() -> Option<&'static GpuContext> { + static GPU: OnceLock> = OnceLock::new(); GPU.get_or_init(|| { if std::env::var_os("FP_CUDA_DISABLE").is_some() { return None; } - match GpuContext::new(0) { - Ok(ctx) => Some(Mutex::new(ctx)), - Err(_) => None, - } + // `FP_CUDA_DEVICE` puts the row reduction on its own GPU. On a single device the two CUDA + // consumers contend: the reduction's thousands of tiny sequential relaunches queue behind + // the multiply's saturating kernels (1.8-9.7 ms standalone vs 8.6-96.8 s co-running), which + // is why [`crate::gpu_lock`] exists at all — and that arbitration then costs ~47% of + // multiply time. Separate devices remove the contention by construction, so the lock + // becomes a no-op (see [`crate::gpu_lock::set_devices_shared`]). + let device = std::env::var("FP_CUDA_DEVICE") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(0); + let mult_devices = multiply_devices(); + let shared = device < mult_devices; + crate::gpu_lock::set_devices_shared(shared); + // Log it: whether arbitration is live decides whether the reduction's thousands of tiny + // launches overlap the multiply's saturating ones, and getting it wrong is invisible in a + // normal run until something fails far away. `[batch-stats] lock=` alone cannot distinguish + // "arbitration off" from "arbitration on but uncontended". + eprintln!( + "[fp-cuda] row reduction on device {device}; multiply spans devices \ + 0..{mult_devices}; multiply yields to reductions: {}; reductions serialize against \ + each other: always", + if shared { + "yes" + } else { + "no (separate devices)" + } + ); + GpuContext::new(device).ok() }) .as_ref() } +/// How many GPUs the cubecl Milnor multiply spreads over — it shards across ALL visible devices, so +/// the row reduction shares a device with it whenever `FP_CUDA_DEVICE < multiply_devices()`. +/// +/// This used to ask which single device the multiply ran on, reading `NASSAU_GPU_DEVICE` — a +/// variable nothing in `algebra` reads any more, left behind when the multiply became multi-GPU. It +/// therefore answered "device 0" no matter how many GPUs the multiply was actually saturating, and +/// `FP_CUDA_DEVICE=2` on a 4-GPU node would silently conclude "separate devices, no arbitration +/// needed" while the multiply was hammering device 2 as well. Arbitration exists to keep the +/// reduction's thousands of tiny sequential relaunches from queueing behind saturating multiply +/// kernels (1.8-9.7 ms standalone vs 8.6-96.8 s co-running); losing it is not a small regression. +/// +/// Mirrors `algebra::algebra::milnor_gpu::gpu_count` — `fp` cannot call it (`algebra` depends on +/// `fp`, not the reverse), so the two must be kept in step. Both honour `CUDA_VISIBLE_DEVICES`, +/// since CUDA renumbers the visible subset to `0..n`. +fn multiply_devices() -> usize { + const MAX_GPUS: usize = 8; + let physical = std::fs::read_dir("/proc/driver/nvidia/gpus") + .map(|d| d.filter_map(|e| e.ok()).count()) + .unwrap_or(0) + .max(1); + let visible = std::env::var("CUDA_VISIBLE_DEVICES").ok().map(|v| { + v.split(',') + .take_while(|e| { + e.trim() + .parse::() + .is_ok_and(|ord| ord < physical.max(MAX_GPUS)) + }) + .count() + }); + std::env::var("NASSAU_GPU_DEVICES") + .ok() + .and_then(|v| v.parse::().ok()) + .filter(|&n| n > 0) + .unwrap_or_else(|| visible.unwrap_or(physical).max(1)) + .clamp(1, MAX_GPUS) +} + +/// The single thread every `fp-cuda` submission goes through, so this process has exactly one +/// owner of the reduction GPU. +/// +/// # Why a thread and not a lock +/// +/// Both `fp-cuda` entry points launch **persistent whole-device grids** (`num_ctas = occupancy x +/// SMs`, cluster-aligned): the row reduction's trailing GEMM and the standalone [`try_mul`]. Two of +/// those cannot be placed at once, and on Hopper the loser does not queue — it fails with a bare +/// `CUDA_ERROR_LAUNCH_FAILED` that compute-sanitizer cannot attribute (0 invalid accesses across a +/// whole run: it was never a memory bug). The cooperative reduction path fails worse, spinning +/// forever at a grid-wide barrier for CTAs that were never scheduled. +/// +/// A lock could serialize this, and [`crate::gpu_lock::exclusive`] did for `row_reduce` — but +/// `try_mul` was deliberately lock-free, so the device still had two independent whole-device +/// consumers and a dedicated GPU was not actually owned by anything. Routing *both* through one +/// thread makes single-ownership structural rather than a discipline every new call site has to +/// remember. +/// +/// # Completion, not submission +/// +/// The job runs to completion on this thread, and both jobs end in a device-to-host download, which +/// synchronizes. That is the property that matters: serializing *submission* is not enough, because +/// kernels outlive the call that launched them — the mistake `gpu_lock::shared` still makes on the +/// multiply side (it is taken inside the submit closure and dropped when submission returns, which +/// is why `[batch-stats] lock=` reads 0.0s in every run). +/// +/// # Not yet done +/// +/// Transfers are serialized along with compute. They need not be: copy engines do not consume SMs, +/// so the next job's H2D upload could overlap the current job's kernels without touching +/// co-residency. That requires splitting each job into upload / compute / download stages on +/// separate streams and pipelining them here; the correctness property above does not depend on it. +mod driver { + use std::sync::{Mutex, OnceLock, mpsc}; + + type Job = Box; + + fn sender() -> &'static Mutex> { + static TX: OnceLock>> = OnceLock::new(); + TX.get_or_init(|| { + let (tx, rx) = mpsc::channel::(); + std::thread::Builder::new() + .name("fp-cuda-driver".into()) + .spawn(move || { + for job in rx { + // NO `gpu_lock::exclusive()` here. Serialization among fp-cuda jobs is + // already structural — this is the only thread that submits them — so the + // guard would be redundant, and taking it deadlocked the run: it waits for + // the multiply's readers to drain while worker threads block on `run` + // waiting for this loop. Yielding to the multiply on a SHARED device has to + // be arranged without a guard held across a blocking job. + job(); + } + }) + .expect("failed to spawn the fp-cuda driver thread"); + Mutex::new(tx) + }) + } + + /// Run `f` on the driver thread and block for its result. `f` owns everything it touches (both + /// call sites have already marshalled to owned limb buffers), so nothing borrows across threads. + pub(super) fn run(f: impl FnOnce() -> T + Send + 'static) -> T { + let (tx, rx) = mpsc::channel(); + sender() + .lock() + .unwrap_or_else(|e| e.into_inner()) + .send(Box::new(move || { + // A send failure means the caller gave up; the job still ran, so just drop it. + let _ = tx.send(f()); + })) + .expect("the fp-cuda driver thread died"); + rx.recv().expect("the fp-cuda driver thread dropped a job") + } +} + /// Row-major, K-major `u64` limbs — the exact layout `fp_cuda::matmul_b1_raw` /// expects (`rows × columns.div_ceil(64)` limbs, no inter-row padding). Uses /// `Matrix::to_bytes`, which already strips the physical row stride. @@ -100,10 +240,12 @@ pub(super) fn try_mul(a: &Matrix, b: &Matrix) -> Option { let a_limbs = to_limbs(a); let b_limbs = to_limbs(b); - let c = { - let guard = ctx.lock().ok()?; - fp_cuda::matmul_b1_raw(&guard, &a_limbs, m, k, &b_limbs, n).ok()? - }; + // Through the driver: this is a persistent whole-device grid, so "concurrent callers do not + // interfere" was wrong — two at once cannot both be placed (see [`driver`]). + // `.ok()` inside the closure: the error is a `Box`, which is not `Send`, so it + // cannot cross back from the driver thread. The caller only distinguishes success from + // fall-back-to-CPU anyway. + let c = driver::run(move || fp_cuda::matmul_b1_raw(ctx, &a_limbs, m, k, &b_limbs, n).ok())?; Some(Matrix::from_data(TWO, m, n, c)) } @@ -127,14 +269,27 @@ pub(crate) fn try_row_reduce(m: &mut Matrix) -> Option { let stride = cols.div_ceil(64); let limbs = to_limbs(m); - let (dev_limbs, perm, r, pivot_cols) = { - let gpu = ctx.lock().ok()?; - let mut dm = gpu.upload(&limbs, rows, cols).ok()?; - let (perm, r, pivot_cols) = gpu.row_reduce_dev(&mut dm).ok()?; - let dev_limbs = gpu.download(&dm).ok()?; - let perm = gpu.download_u32(&perm).ok()?; - (dev_limbs, perm, r, pivot_cols) - }; + // Lock-free, per-thread stream (see [`context`]): the default row-reduce is composable (no + // cooperative launch) and allocates its device buffers per call, so concurrent rayon workers + // reduce different matrices on independent streams — overlapping instead of serializing. + // + // The claim that this "needs no cross-runtime exclusion against the cubecl multiply" is exactly + // backwards. Composability (no cooperative launch) means this path *can* overlap other GPU work + // without deadlocking — not that it should. This reduction is a chain of thousands of tiny + // sequential per-column relaunches, so overlapping it with the multiply's saturating kernels + // makes every launch queue: 1.8–9.7 ms standalone becomes 8.6–96.8 s co-running. Take the + // device exclusively for the duration; see [`fp::gpu_lock`] for the measurements and the cost + // (~5 s of multiply pause across a whole stem-200 resolution). + // The exclusive guard now lives on the driver thread, which holds it for the whole job — see + // [`driver`]. Taking it here as well would deadlock: the driver would wait on a guard this + // thread holds while this thread waits on the driver. + let (dev_limbs, perm, r, pivot_cols) = driver::run(move || { + let mut dm = ctx.upload(&limbs, rows, cols).ok()?; + let (perm, r, pivot_cols) = ctx.row_reduce_dev(&mut dm).ok()?; + let dev_limbs = ctx.download(&dm).ok()?; + let perm = ctx.download_u32(&perm).ok()?; + Some((dev_limbs, perm, r, pivot_cols)) + })?; // Materialize the canonical RREF: pivot k (column pivot_cols[k], ascending) // at row k, taken from device row perm[k]; rows [r, rows) zero. diff --git a/ext/crates/fp/src/gpu_lock.rs b/ext/crates/fp/src/gpu_lock.rs new file mode 100644 index 0000000000..7df9d94324 --- /dev/null +++ b/ext/crates/fp/src/gpu_lock.rs @@ -0,0 +1,267 @@ +//! Process-wide arbitration between the two CUDA consumers that share this GPU: the cubecl Milnor +//! multiply (`algebra::algebra::milnor_gpu`) and the `fp-cuda` row reduction ([`crate::blas::cuda`]). +//! +//! # Why this exists +//! +//! The two runtimes have opposite performance shapes. The multiply is *throughput* work: large, +//! long-running kernels that saturate the SMs. The composable (non-cooperative) row reduction is +//! *latency* work: thousands of tiny, strictly sequential per-column relaunches. Run them at the +//! same time and every one of those thousands of launches queues behind a saturating multiply +//! kernel, so a reduction that takes single-digit milliseconds on an unshared GPU takes tens of +//! seconds — measured 1.8–9.7 ms standalone versus 8.6–96.8 s co-running, a ~10 000× loss, with +//! `nvidia-smi` showing 99 % SM and 9 % memory utilisation (queueing, not compute). +//! +//! Being *composable* (no cooperative launch, so no co-residency requirement) means the reduction +//! **can** run alongside other GPU work without deadlocking. It does not mean it should: overlap is +//! precisely what destroys it. +//! +//! # The trade +//! +//! Giving a large reduction brief exclusive use of the device costs almost nothing: a whole +//! stem-200 resolution runs ~440 GPU reductions of ~10 ms each, so the multiply pauses for ~5 s in +//! total. Multiplies still overlap freely with each other — they take the shared side. +//! +//! Writer preference is deliberate. Multiplies are continuous and readers are many; a plain +//! `RwLock` would let the reduction starve indefinitely behind the stream of multiplies, which is +//! the failure this lock exists to prevent. + +use std::{ + sync::{Condvar, Mutex}, + time::{Duration, Instant}, +}; + +#[derive(Default)] +struct State { + /// Multiplies currently submitting. + readers: usize, + /// A reduction currently holds the device. + writer: bool, + /// Reductions blocked waiting; new multiplies yield to them (writer preference). + writers_waiting: usize, +} + +/// Whether the two CUDA runtimes share one device. Arbitration is only needed when they do: +/// measured at ~47% of multiply time (`[batch-stats] lock=`), which is pure waste once the row +/// reduction has its own GPU. Defaults to `true` (single-device, the safe assumption) until +/// [`crate::blas::cuda`] resolves the device ids on first GPU use. +static SHARED_DEVICE: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(true); + +/// Record whether the multiply and the row reduction target the same GPU. +pub fn set_devices_shared(shared: bool) { + SHARED_DEVICE.store(shared, std::sync::atomic::Ordering::Relaxed); +} + +fn arbitration_needed() -> bool { + SHARED_DEVICE.load(std::sync::atomic::Ordering::Relaxed) +} + +fn state() -> &'static (Mutex, Condvar) { + use std::sync::OnceLock; + static STATE: OnceLock<(Mutex, Condvar)> = OnceLock::new(); + STATE.get_or_init(|| (Mutex::new(State::default()), Condvar::new())) +} + +/// Shared access, held by the Milnor multiply while it submits and reads back. Many may be held at +/// once; all are excluded by an [`exclusive`] holder. +pub struct SharedGuard(()); + +/// Exclusive access, held by a large GPU row reduction for the duration of its launch chain. +pub struct ExclusiveGuard(()); + +impl Drop for SharedGuard { + fn drop(&mut self) { + if !arbitration_needed() { + return; + } + let (lock, cv) = state(); + let mut s = lock.lock().unwrap_or_else(|e| e.into_inner()); + s.readers -= 1; + if s.readers == 0 { + cv.notify_all(); + } + } +} + +impl Drop for ExclusiveGuard { + fn drop(&mut self) { + if !arbitration_needed() { + return; + } + let (lock, cv) = state(); + let mut s = lock.lock().unwrap_or_else(|e| e.into_inner()); + s.writer = false; + cv.notify_all(); + } +} + +/// How long a multiply defers to a waiting reduction before going ahead anyway. +/// +/// This is a **safety valve, not the mechanism**. It must exceed the time a reduction holds the +/// device, or exclusivity evaporates precisely when it matters: at 25 ms, multiplies barged back in +/// partway through every multi-second reduction, which both kept reductions ~1000× slow and put the +/// overlap back that crashes the run. Correctness against deadlock comes from *where* the shared +/// guard is taken (`milnor_gpu.rs`, past every rayon section), not from this timeout firing. +const SHARED_YIELD: Duration = Duration::from_secs(60); +/// How long a reduction waits for in-flight multiplies to drain before going ahead anyway. +const EXCLUSIVE_DRAIN: Duration = Duration::from_secs(10); + +/// Acquire shared (multiply) access, briefly yielding to any waiting reduction. +/// +/// The yield is **bounded**, and that bound is load-bearing rather than a tuning choice. Callers +/// reach this from inside rayon parallel sections: a multiply that blocks here can be holding a +/// join that another worker's stolen multiply needs, so an unbounded yield deadlocks (observed on +/// H200 — a reduction waiting on `readers == 0` while every reader waited on a join that could only +/// finish once a blocked reader proceeded). Timing out costs the reduction some exclusivity; never +/// timing out costs the whole resolution. +pub fn shared() -> SharedGuard { + if !arbitration_needed() { + return SharedGuard(()); + } + let (lock, cv) = state(); + let mut s = lock.lock().unwrap_or_else(|e| e.into_inner()); + let deadline = Instant::now() + SHARED_YIELD; + while s.writer || s.writers_waiting > 0 { + let Some(remaining) = deadline.checked_duration_since(Instant::now()) else { + break; + }; + s = cv + .wait_timeout(s, remaining) + .unwrap_or_else(|e| e.into_inner()) + .0; + } + s.readers += 1; + SharedGuard(()) +} + +/// Acquire exclusive (row-reduction) access, waiting for in-flight multiplies to drain. +/// +/// Waiting out another *reduction* is unbounded and safe: reductions never block on rayon work, so +/// they always finish in finite time, and letting two run at once is what the concurrency cap was +/// added to prevent. Waiting for *multiplies* to drain is bounded for the reason in [`shared`] — +/// past the deadline this proceeds without full exclusivity, which is slow, not wrong. +pub fn exclusive() -> ExclusiveGuard { + // NOT gated on `arbitration_needed()`. That flag answers "does the MULTIPLY share this device", + // which is the only question [`shared`] cares about. Reductions must serialize against each + // OTHER regardless, because the row reduction's GEMM is a persistent whole-device grid + // (`num_ctas = occupancy x SMs`, cluster-aligned): two concurrent reductions each demand the + // entire GPU and neither can be placed, which surfaces as `CUDA_ERROR_LAUNCH_FAILED`. + // + // Skipping this on a dedicated reduction GPU is what made `FP_CUDA_DEVICE=3` with the multiply + // on 0..2 fail 63 times in 300 s — an isolated device removes multiply contention but leaves + // reduction-vs-reduction contention untouched. + let (lock, cv) = state(); + let mut s = lock.lock().unwrap_or_else(|e| e.into_inner()); + s.writers_waiting += 1; + while s.writer { + s = cv.wait(s).unwrap_or_else(|e| e.into_inner()); + } + // Claim the slot *before* draining readers. Both waits below release the mutex, so a writer + // that only set this flag afterwards could race another writer through the check above and let + // two reductions run at once (caught by the test in this module). + s.writer = true; + s.writers_waiting -= 1; + let deadline = Instant::now() + EXCLUSIVE_DRAIN; + while s.readers > 0 { + let Some(remaining) = deadline.checked_duration_since(Instant::now()) else { + break; + }; + s = cv + .wait_timeout(s, remaining) + .unwrap_or_else(|e| e.into_inner()) + .0; + } + ExclusiveGuard(()) +} + +#[cfg(test)] +mod tests { + use std::{ + sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }, + thread, + }; + + use super::*; + + /// What the arbitration actually guarantees: every acquisition terminates under heavy + /// contention (no deadlock — the property the first version got wrong), two reductions never + /// overlap, and multiplies still overlap each other. Exclusion against multiplies is + /// deliberately best-effort (see [`EXCLUSIVE_DRAIN`]), so it is not asserted here. + #[test] + fn contended_acquisition_terminates_and_writers_are_exclusive() { + let live_shared = Arc::new(AtomicUsize::new(0)); + let live_exclusive = Arc::new(AtomicUsize::new(0)); + let violations = Arc::new(AtomicUsize::new(0)); + let max_shared = Arc::new(AtomicUsize::new(0)); + + // Readers overlap when no reduction is demanding the device. This has to be measured on its + // own: `exclusive` is now unconditional (reductions must serialize against each other — + // their GEMM is a whole-device grid), so the contended phase below keeps a writer queued + // essentially always, and writer preference then correctly holds readers off. Asserting + // overlap *during* that phase measured the scheduler, not the lock. + let mut warmup = Vec::new(); + for _ in 0..8 { + let (live, max) = (Arc::clone(&live_shared), Arc::clone(&max_shared)); + warmup.push(thread::spawn(move || { + for _ in 0..200 { + let _g = shared(); + let n = live.fetch_add(1, Ordering::SeqCst) + 1; + max.fetch_max(n, Ordering::SeqCst); + thread::yield_now(); + live.fetch_sub(1, Ordering::SeqCst); + } + })); + } + for h in warmup { + h.join().unwrap(); + } + assert!( + max_shared.load(Ordering::SeqCst) > 1, + "multiplies never overlapped — the shared side is serialising, which defeats the point" + ); + + let mut handles = Vec::new(); + for _ in 0..8 { + let (live, bad, max) = ( + Arc::clone(&live_shared), + Arc::clone(&violations), + Arc::clone(&max_shared), + ); + handles.push(thread::spawn(move || { + for _ in 0..200 { + let _g = shared(); + let n = live.fetch_add(1, Ordering::SeqCst) + 1; + max.fetch_max(n, Ordering::SeqCst); + thread::yield_now(); + live.fetch_sub(1, Ordering::SeqCst); + let _ = &bad; + } + })); + } + for _ in 0..3 { + let (live_w, bad) = (Arc::clone(&live_exclusive), Arc::clone(&violations)); + handles.push(thread::spawn(move || { + for _ in 0..100 { + let _g = exclusive(); + if live_w.fetch_add(1, Ordering::SeqCst) != 0 { + bad.fetch_add(1, Ordering::SeqCst); + } + thread::yield_now(); + live_w.fetch_sub(1, Ordering::SeqCst); + } + })); + } + // Joining at all is the deadlock assertion: the previous design hung here forever. + for h in handles { + h.join().unwrap(); + } + + assert_eq!( + violations.load(Ordering::SeqCst), + 0, + "two reductions held the device at once" + ); + } +} diff --git a/ext/crates/fp/src/lib.rs b/ext/crates/fp/src/lib.rs index 8d971da2a8..cfacf73b5c 100644 --- a/ext/crates/fp/src/lib.rs +++ b/ext/crates/fp/src/lib.rs @@ -11,6 +11,7 @@ pub mod prime; pub mod vector; pub mod blas; +pub mod gpu_lock; pub(crate) mod simd; diff --git a/ext/crates/fp/src/matrix/matrix_inner.rs b/ext/crates/fp/src/matrix/matrix_inner.rs index 99b014e631..fcf3261923 100644 --- a/ext/crates/fp/src/matrix/matrix_inner.rs +++ b/ext/crates/fp/src/matrix/matrix_inner.rs @@ -677,11 +677,53 @@ impl Matrix { // For large p = 2 matrices, try the device-resident GPU reduction; it // produces the identical canonical RREF + pivots. Falls back to the CPU // M4RI path below when the GPU is unavailable or below threshold. + // + // Instrumentation: for every p=2 reduction of a non-trivial matrix + // (min(rows,cols) >= 1024) we emit a `fp::rr` tracing event recording the + // dimensions and whether the GPU path was taken (`path="gpu"`) or it fell + // back to CPU M4RI (`path="cpu"` — either below the 8192 threshold or a + // launch failure; the logged dims disambiguate). The event inherits the + // active nassau span, so it carries the bidegree/signature context. #[cfg(feature = "gpu")] - if p == 2 - && let Some(rank) = crate::blas::cuda::try_row_reduce(self) - { - return rank; + if p == 2 { + let (rr_rows, rr_cols) = (self.rows(), self.columns()); + let rr_big = rr_rows.min(rr_cols) >= 1024; + // Wrap the GPU reduce in an ENTERED span (not just a completion event) so a wedge + // *inside* `try_row_reduce` leaves an open `gpu_row_reduce` span with no `close` in the + // log — distinguishing an RREF hang from a milnor-multiply hang. Inherits the active + // nassau step span, so it carries the bidegree/signature. Dropped on return/fallthrough. + let _rr_span = rr_big.then(|| { + tracing::info_span!(target: "fp::rr", "gpu_row_reduce", rows = rr_rows, cols = rr_cols) + .entered() + }); + match crate::blas::cuda::try_row_reduce(self) { + Some(rank) => { + if rr_big { + tracing::info!( + target: "fp::rr", + rows = rr_rows, + cols = rr_cols, + min = rr_rows.min(rr_cols), + path = "gpu", + rank, + "row_reduce" + ); + } + return rank; + } + None => { + if rr_big { + tracing::info!( + target: "fp::rr", + rows = rr_rows, + cols = rr_cols, + min = rr_rows.min(rr_cols), + path = "cpu", + "row_reduce" + ); + } + } + } } self.initialize_pivots(); diff --git a/ext/crates/fp/tests/cuda_dispatch.rs b/ext/crates/fp/tests/cuda_dispatch.rs index d339339277..3cad7e8415 100644 --- a/ext/crates/fp/tests/cuda_dispatch.rs +++ b/ext/crates/fp/tests/cuda_dispatch.rs @@ -95,3 +95,39 @@ fn gpu_row_reduce_matches_cpu() { assert_eq!(gpu, cpu, "RREF mismatch at {rows}x{cols} rank={rank}"); } } + +/// Many threads row-reducing on the GPU AT ONCE must each stay bit-identical to the CPU — the +/// concurrency the per-thread-stream refactor enables. Isolates the GPU RREF path from the cubecl +/// multiply: if concurrent reductions share any device state (a `__device__` global, a fixed +/// scratch), this corrupts or LAUNCH_FAILEDs; if they're truly independent per-stream, it passes. +#[test] +fn gpu_row_reduce_concurrent() { + // SAFETY: set once before any threshold() read; same value as the sibling test. + unsafe { std::env::set_var("FP_CUDA_RR_THRESHOLD", "2048") }; + const THREADS: usize = 16; + const ITERS: usize = 8; + std::thread::scope(|s| { + for t in 0..THREADS { + s.spawn(move || { + for i in 0..ITERS { + // Vary shapes per thread/iter so streams don't run identical work in lockstep. + let rows = 2048 + 256 * (t % 8); + let cols = 2048 + 256 * (i % 6); + let base = clean_matrix(rows, cols, 0); + let mut gpu = base.clone(); + let rank_gpu = gpu.row_reduce(); + let mut cpu = base.clone(); + let rank_cpu = cpu.row_reduce_blas3(); + assert_eq!( + rank_gpu, rank_cpu, + "concurrent rank mismatch {rows}x{cols} (t{t} i{i})" + ); + assert_eq!( + gpu, cpu, + "concurrent RREF mismatch {rows}x{cols} (t{t} i{i})" + ); + } + }); + } + }); +} diff --git a/ext/crates/maybe-rayon/src/concurrent.rs b/ext/crates/maybe-rayon/src/concurrent.rs index 9426cf7f84..bbc6bb4d8f 100644 --- a/ext/crates/maybe-rayon/src/concurrent.rs +++ b/ext/crates/maybe-rayon/src/concurrent.rs @@ -86,3 +86,34 @@ where pub fn empty() -> rayon::iter::Empty { rayon::iter::empty() } + +/// A private thread pool, so a caller can run work on threads that are NOT the global pool's. +/// +/// The motivating case: a task that blocks (e.g. waiting on a GPU) must not be able to steal an +/// unrelated large job while it waits, which is how priority inversion happens. A private pool +/// bounds what its workers can pick up to the tasks submitted to it. +pub struct MaybeThreadPool(rayon::ThreadPool); + +impl MaybeThreadPool { + /// Build a pool with `num_threads` workers named ``. + pub fn new(num_threads: usize, name_prefix: &'static str) -> Self { + Self( + rayon::ThreadPoolBuilder::new() + .num_threads(num_threads) + .thread_name(move |i| format!("{name_prefix}{i}")) + .build() + .expect("failed to build a MaybeThreadPool"), + ) + } + + /// Run `f` inside the pool; parallel iterators created within it use only this pool's workers. + pub fn install R + Send>(&self, f: F) -> R { + self.0.install(f) + } +} + +/// Width of the global pool — the number of callers that can be submitting work at once. Sizing a +/// private pool that those callers block on requires knowing this. +pub fn max_num_threads() -> usize { + rayon::current_num_threads() +} diff --git a/ext/crates/maybe-rayon/src/sequential.rs b/ext/crates/maybe-rayon/src/sequential.rs index a2493a7e43..767c912065 100644 --- a/ext/crates/maybe-rayon/src/sequential.rs +++ b/ext/crates/maybe-rayon/src/sequential.rs @@ -129,3 +129,23 @@ impl Iterator for Empty { pub fn empty() -> Empty { Empty(std::marker::PhantomData) } + +/// Sequential proxy for the concurrent module's `MaybeThreadPool`: holds no threads and runs +/// `install` inline, so a build without `concurrent` keeps the same call sites and stays +/// single-threaded for debugging. +pub struct MaybeThreadPool; + +impl MaybeThreadPool { + pub fn new(_num_threads: usize, _name_prefix: &'static str) -> Self { + Self + } + + pub fn install R + Send>(&self, f: F) -> R { + f() + } +} + +/// Sequential proxy: one caller. +pub fn max_num_threads() -> usize { + 1 +} diff --git a/ext/examples/bruner.rs b/ext/examples/bruner.rs index 71c5c0e22a..6e4411b1fd 100644 --- a/ext/examples/bruner.rs +++ b/ext/examples/bruner.rs @@ -25,7 +25,7 @@ use std::{ use algebra::{ Algebra, MilnorAlgebra, - milnor_algebra::MilnorBasisElement, + milnor_algebra::{MilnorBasisElement, PPartEntry}, module::{FreeModule as FM, Module, homomorphism::FreeModuleHomomorphism as FMH}, }; use anyhow::{Context, Error, Result}; @@ -95,7 +95,10 @@ fn get_algebra_element<'a>( let entry = &entry[1..]; let elt = MilnorBasisElement { q_part: 0, - p_part: entry.split(',').map(|x| x.parse().unwrap()).collect(), + p_part: entry + .split(',') + .map(|x| x.parse::().unwrap()) + .collect(), degree: t, }; a.basis_element_to_index(&elt) diff --git a/ext/examples/sq0.rs b/ext/examples/sq0.rs index dc671c159b..9836e18c00 100644 --- a/ext/examples/sq0.rs +++ b/ext/examples/sq0.rs @@ -77,8 +77,9 @@ mod double { mod double_algebra { use algebra::{ - AdemAlgebra, Algebra, MilnorAlgebra, SteenrodAlgebra, adem_algebra::AdemBasisElement, - milnor_algebra::MilnorBasisElement, + AdemAlgebra, Algebra, MilnorAlgebra, SteenrodAlgebra, + adem_algebra::AdemBasisElement, + milnor_algebra::{MilnorBasisElement, PPart}, }; pub trait DoubleAlgebra: Algebra { @@ -92,14 +93,14 @@ mod double { let p_part = elt .p_part .iter() - .map(|&x| { + .map(|x| { if x.is_multiple_of(2) { Some(x / 2) } else { None } }) - .collect::>>()?; + .collect::>()?; Some(self.basis_element_to_index(&MilnorBasisElement { degree: degree / 2, p_part, diff --git a/ext/src/nassau.rs b/ext/src/nassau.rs index 28e1886e91..3dbb0211f3 100644 --- a/ext/src/nassau.rs +++ b/ext/src/nassau.rs @@ -19,7 +19,7 @@ use std::{ use algebra::{ Algebra, combinatorics, - milnor_algebra::{MilnorAlgebra, PPartEntry}, + milnor_algebra::{MilnorAlgebra, PPart, PPartEntry}, module::{ FreeModule, GeneratorData, Module, ZeroModule, homomorphism::{FreeModuleHomomorphism, FullModuleHomomorphism, ModuleHomomorphism}, @@ -28,7 +28,7 @@ use algebra::{ use anyhow::anyhow; use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt}; use fp::{ - matrix::{AugmentedMatrix, Matrix}, + matrix::{AugmentedMatrix, Matrix, Subspace}, prime::{Prime, TWO, ValidPrime}, vector::{FpSlice, FpSliceMut, FpVector}, }; @@ -46,11 +46,13 @@ use sseq::coordinates::{Bidegree, BidegreeGenerator}; use crate::{ chain_complex::{AugmentedChainComplex, ChainComplex, FiniteChainComplex, FreeChainComplex}, save::{NassauCommand, NassauQiWriter, SaveDirectory, SaveKind}, + utils::parallel::ParallelGuard, }; /// See [`resolution::SenderData`](../resolution/struct.SenderData.html). This differs by not having the `new` field. struct SenderData { b: Bidegree, + retry: bool, sender: mpsc::Sender, } @@ -59,6 +61,18 @@ impl SenderData { sender .send(Self { b, + retry: false, + sender: sender.clone(), + }) + .unwrap() + } + + pub(crate) fn send_retry(b: Bidegree, sender: mpsc::Sender) { + tracing::info!(%b, "retrying"); + sender + .send(Self { + b, + retry: true, sender: sender.clone(), }) .unwrap() @@ -92,15 +106,40 @@ impl MilnorSubalgebra { Self { profile: vec![] } } - /// Computes the signature of an element - fn has_signature(&self, ppart: &[PPartEntry], signature: &[PPartEntry]) -> bool { - for (i, (&profile, &signature)) in self.profile.iter().zip(signature).enumerate() { - let ppart = ppart.get(i).copied().unwrap_or(0); - if ppart & ((1 << profile) - 1) != signature { - return false; + /// The test "does this element have this signature" compiled into a `(mask, value)` pair to + /// match against the packed p-part. + /// + /// The per-entry test is `ppart[i] & ((1 << profile[i]) - 1) == signature[i]`. Because each + /// entry occupies a fixed field of the packed word, the low `profile[i]` bits of entry `i` are + /// a fixed bit range of that word, so the whole conjunction is a single `&` and `==`. Entries + /// past the end of the p-part read as zero, which the packing already gives us for free. + /// Returns `None` if no element can have this signature, which the caller turns into an empty + /// result. A profile is not bounded by [`PPart::MAX_LEN`] -- `SubalgebraIterator` grows one + /// without limit and `from_bytes` reads whatever length a file gives -- so both ways a + /// signature can fail to be representable have to be handled here rather than assumed away. + fn packed_signature(&self, signature: &[PPartEntry]) -> Option<(u64, u64)> { + let mut mask = 0; + let mut value = 0; + for (i, (&profile, &entry)) in self.profile.iter().zip(signature).enumerate() { + if i >= PPart::MAX_LEN { + // No p-part of a representable degree has an entry this far out, so it reads as + // zero: a non-zero constraint is unsatisfiable and a zero one is vacuous. + if entry != 0 { + return None; + } + continue; + } + // A profile wider than the field constrains the whole field. + let width = std::cmp::min(profile as u32, PPart::width(i)); + // The masked entry has only `width` bits, so a signature wanting more matches nothing. + // Packing it anyway would spill into the neighbouring field. + if (entry as u64) >> width != 0 { + return None; } + mask |= ((1u64 << width) - 1) << PPart::shift(i); + value |= (entry as u64) << PPart::shift(i); } - true + Some((mask, value)) } fn zero_signature(&self) -> Vec { @@ -124,28 +163,39 @@ impl MilnorSubalgebra { signature: &'a [PPartEntry], max_gen_degree: i32, ) -> impl Iterator + 'a { - module - .iter_gen_offsets([degree]) - .take_while(move |gen_data| gen_data.gen_deg < max_gen_degree) - .flat_map( - move |GeneratorData { - gen_deg, - start: [offset], - end: _, - }| { - algebra - .ppart_table(degree - gen_deg) - .iter() - .enumerate() - .filter_map(move |(n, op)| { - if self.has_signature(op, signature) { - Some(offset + n) - } else { - None - } - }) - }, - ) + // The mask depends only on the signature, so compute it once for the whole sweep (PR #280). + // An unrepresentable signature yields no elements at all. + // + // `take_while` is retained from this branch and is NOT optional: `max_gen_degree` exists so + // a reader can ignore the generators of the current internal degree, which another thread + // may be adding concurrently. Generators are laid out in increasing degree, so this is + // exactly a prefix. + self.packed_signature(signature) + .into_iter() + .flat_map(move |(mask, value)| { + module + .iter_gen_offsets([degree]) + .take_while(move |gen_data| gen_data.gen_deg < max_gen_degree) + .flat_map( + move |GeneratorData { + gen_deg, + start: [offset], + end: _, + }| { + algebra + .ppart_table(degree - gen_deg) + .iter() + .enumerate() + .filter_map(move |(n, op)| { + if op.bits() & mask == value { + Some(offset + n) + } else { + None + } + }) + }, + ) + }) } /// The number of basis elements in `degree` coming from generators of `module` of degree @@ -167,53 +217,6 @@ impl MilnorSubalgebra { .unwrap_or(0) } - /// Get the matrix of a free module homomorphism when restricted to the subquotient given by - /// the signature. - /// - /// Only generators of the target of degree strictly less than `target_max_gen_degree` are used - /// (see [`Self::signature_mask`]). - fn signature_matrix( - &self, - hom: &FreeModuleHomomorphism>, - degree: i32, - signature: &[PPartEntry], - target_max_gen_degree: i32, - ) -> Matrix { - let p = hom.prime(); - let source = hom.source(); - let target = hom.target(); - let algebra = target.algebra(); - let target_degree = degree - hom.degree_shift(); - - let target_mask: Vec = self - .signature_mask( - &algebra, - &target, - target_degree, - signature, - target_max_gen_degree, - ) - .collect(); - - let source_mask: Vec = self - .signature_mask(&algebra, &source, degree, signature, i32::MAX) - .collect(); - - let mut scratch = FpVector::new( - p, - Self::restricted_dimension(&target, target_degree, target_max_gen_degree), - ); - let mut result = Matrix::new(p, source_mask.len(), target_mask.len()); - - for (mut row, &masked_index) in std::iter::zip(result.iter_mut(), &source_mask) { - scratch.set_to_zero(); - hom.apply_to_basis_element_restricted(scratch.as_slice_mut(), 1, degree, masked_index); - - row.add_masked(scratch.as_slice(), 1, &target_mask); - } - result - } - /// Iterate through all signatures of this algebra that contain elements of degree at most /// `degree` (inclusive). This skips the initial zero signature. fn iter_signatures(&self, degree: i32) -> impl Iterator> + '_ { @@ -361,6 +364,10 @@ fn restricted_partial_matrix( inputs: &[usize], target_dim: usize, ) -> Matrix { + // Spanned because this is the fallback every build below `NASSAU_GPU_MIN_WORK` takes, and it + // was invisible: the trace attributed 902.8 s to `extract_restricted` over 4915 GPU builds, but + // a stem-150 run issues ~21 500 builds, so most of them landed here and were never counted. + let _s = tracing::trace_span!("cpu_restricted", rows = inputs.len(), target_dim).entered(); let mut matrix = Matrix::new(hom.prime(), inputs.len(), target_dim); if target_dim > 0 { matrix @@ -430,6 +437,38 @@ fn reuse_full_matrix(_diff: &FreeModuleHomomorphism>) } } +/// Max `rows × cols` of the full restricted matrix for which [`reuse_full_matrix`] builds it all at +/// once. Above this the all-rows build (and its dense GPU readback, both held across the whole +/// signature loop) dominates host memory at high stems — the ~12 GB dense regions behind the stem-180 +/// OOM. Past the cap we fall back to per-signature builds (each a bounded row subset, like the CPU), +/// trading a little launch amortization for a peak that scales with the largest single signature +/// rather than the whole bidegree. `NASSAU_GPU_REUSE_MAX_WORK` overrides (0 = never reuse). Default +/// ~1e10 (rows×cols) ≈ a ~1.2 GB restricted matrix. +#[cfg(feature = "gpu")] +fn gpu_reuse_max_work() -> u64 { + static W: std::sync::LazyLock = std::sync::LazyLock::new(|| { + std::env::var("NASSAU_GPU_REUSE_MAX_WORK") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(10_000_000_000) + }); + *W +} + +/// Whether the full restricted matrix (`rows × cols`) is small enough to build all at once (see +/// [`gpu_reuse_max_work`]). Always true without the `gpu` feature (the reuse path is off anyway). +fn reuse_within_cap(_rows: usize, _cols: usize) -> bool { + #[cfg(feature = "gpu")] + { + let w = gpu_reuse_max_work(); + w > 0 && (_rows as u64).saturating_mul(_cols as u64) <= w + } + #[cfg(not(feature = "gpu"))] + { + true + } +} + /// Extract `rows` of `full` into a fresh matrix (`out.row(i) = full.row(rows[i])`), preserving the /// column layout. Slices a precomputed full (restricted-column) differential matrix into one /// signature's partial matrix (see [`reuse_full_matrix`]). `rows` must index within `full` — the @@ -742,16 +781,20 @@ impl> Resolution { }; let guard = tracing::info_span!("step", signature = ?zero_sig).entered(); - let next_mask: Vec = subalgebra - .signature_mask(&algebra, next, b.t(), &zero_sig, next_bound) - .collect(); + let next_mask: Vec = tracing::trace_span!("zs_masks").in_scope(|| { + subalgebra + .signature_mask(&algebra, next, b.t(), &zero_sig, next_bound) + .collect() + }); let next_masked_dim = next_mask.len(); // When GPU reuse is active, build ONE full restricted matrix over every (restricted) source // row at degree `b.t()` in a single launch, then slice each signature's rows out of it. // `target_dim` is the restricted source dimension, so `0..target_dim` is exactly the row set // the per-signature masks partition; `next_dim` is the restricted column count. - let full_reuse: Option = if reuse_full_matrix(&self.differentials[b.s() - 1]) { + let full_reuse: Option = if reuse_full_matrix(&self.differentials[b.s() - 1]) + && reuse_within_cap(target_dim, next_dim) + { let all_rows: Vec = (0..target_dim).collect(); Some(restricted_partial_matrix_maybe_gpu( &self.differentials[b.s() - 1], @@ -763,29 +806,41 @@ impl> Resolution { None }; - let full_matrix = match &full_reuse { - Some(full) => { - debug_assert!(target_mask.iter().all(|&r| r < full.rows())); - select_rows(full, &target_mask) - } - None => { - restricted_partial_matrix_maybe_gpu( - &self.differentials[b.s() - 1], - b.t(), - &target_mask, - next_dim, - ) - } - }; - let mut masked_matrix = - AugmentedMatrix::new(p, target_masked_dim, [next_masked_dim, target_masked_dim]); + let full_matrix = + tracing::trace_span!("zs_select", rows = target_mask.len()).in_scope(|| { + match &full_reuse { + Some(full) => { + debug_assert!(target_mask.iter().all(|&r| r < full.rows())); + select_rows(full, &target_mask) + } + None => restricted_partial_matrix_maybe_gpu( + &self.differentials[b.s() - 1], + b.t(), + &target_mask, + next_dim, + ), + } + }); + let mut masked_matrix = tracing::trace_span!( + "zs_assemble", + rows = target_masked_dim, + cols = next_masked_dim + ) + .in_scope(|| { + let mut m = + AugmentedMatrix::new(p, target_masked_dim, [next_masked_dim, target_masked_dim]); + m.segment(0, 0).add_masked(&full_matrix, &next_mask); + m.segment(1, 1).add_identity(); + m + }); - masked_matrix - .segment(0, 0) - .add_masked(&full_matrix, &next_mask); - masked_matrix.segment(1, 1).add_identity(); - masked_matrix.row_reduce(); - let kernel = masked_matrix.compute_kernel(); + tracing::trace_span!( + "zs_row_reduce", + rows = target_masked_dim, + cols = next_masked_dim + ) + .in_scope(|| masked_matrix.row_reduce()); + let kernel = tracing::trace_span!("zs_kernel").in_scope(|| masked_matrix.compute_kernel()); Self::write_qi( &mut f, @@ -806,13 +861,37 @@ impl> Resolution { f.write_fix()?; } - // Compute image - let mut n = - subalgebra.signature_matrix(&self.differentials[b.s()], b.t(), &zero_sig, target_bound); - n.row_reduce(); + // Compute image: d_s applied to the zero-signature source basis, column-masked to the + // zero-signature target basis. This is the same restricted multiply as `full_matrix` above + // (on d_s = differentials[b.s()] rather than d_{s-1}), and its target mask/dimension are + // exactly the `target_mask`/`target_dim` already computed for this bidegree — d_s and + // d_{s-1} share the target module `modules[b.s() - 1]`. So route it through the same + // (GPU-offloaded, work-gated) restricted-matrix path and apply the column mask on CPU, + // instead of `signature_matrix`'s serial per-row CPU multiply. + let source_mask: Vec = tracing::trace_span!("zs_source_mask").in_scope(|| { + subalgebra + .signature_mask(&algebra, &self.modules[b.s()], b.t(), &zero_sig, i32::MAX) + .collect() + }); + let img_full = restricted_partial_matrix_maybe_gpu( + &self.differentials[b.s()], + b.t(), + &source_mask, + target_dim, + ); + let mut n = tracing::trace_span!("img_assemble", rows = source_mask.len()).in_scope(|| { + let mut n = Matrix::new(p, source_mask.len(), target_masked_dim); + for (mut row, full_row) in std::iter::zip(n.iter_mut(), img_full.iter()) { + row.add_masked(full_row, 1, &target_mask); + } + n + }); + tracing::trace_span!("img_row_reduce", rows = source_mask.len()) + .in_scope(|| n.row_reduce()); let next_row = n.rows(); - let num_new_gens = n.extend_image(0, n.columns(), &kernel, 0).len(); + let num_new_gens = tracing::trace_span!("extend_image") + .in_scope(|| n.extend_image(0, n.columns(), &kernel, 0).len()); if b.t() < b.s() { assert_eq!(num_new_gens, 0, "Adding generators at {b}"); @@ -823,14 +902,17 @@ impl> Resolution { let mut xs = vec![FpVector::new(p, target_dim); num_new_gens]; let mut dxs = vec![FpVector::new(p, next_dim); num_new_gens]; - for ((x, x_masked), dx) in xs - .iter_mut() - .zip_eq(n.iter().skip(next_row)) - .zip_eq(&mut dxs) { - x.as_slice_mut().add_unmasked(x_masked, 1, &target_mask); - for (i, _) in x_masked.iter_nonzero() { - dx.as_slice_mut().add(full_matrix.row(i), 1); + let _s = tracing::trace_span!("zs_dx_init", gens = xs.len()).entered(); + for ((x, x_masked), dx) in xs + .iter_mut() + .zip_eq(n.iter().skip(next_row)) + .zip_eq(&mut dxs) + { + x.as_slice_mut().add_unmasked(x_masked, 1, &target_mask); + for (i, _) in x_masked.iter_nonzero() { + dx.as_slice_mut().add(full_matrix.row(i), 1); + } } } @@ -840,8 +922,48 @@ impl> Resolution { drop(guard); + // Probe (`NASSAU_PROBE_SIG_INDEP=1`): are the signature steps independent? + // + // Each step reads `dx.entry(v)` for `v` in its own `next_mask`, then writes `dx` with rows + // of the *unmasked* `full_matrix`, whose support can extend outside that mask. If those + // writes never land on a column another step later reads, the steps are solving against an + // unchanging `dx` and the loop is parallelisable (solve independently, combine). If they do, + // the loop is a forward substitution and must stay ordered. Comparing each read against a + // pre-loop snapshot answers exactly that, without altering what the loop computes. + // + // MEASURED (S_2 stem 60, max_s 30, 2233 bidegrees): 27 028 of 140 614 reads — 19.2% — see a + // value an earlier signature wrote. So the steps are NOT independent and the loop must stay + // ordered: it is a forward substitution, and "solve every signature separately, then + // combine" would be wrong, not merely racy. + // + // That is a negative result about the LIFT only, and the lift is the cheap end. Everything + // above it — `sig_masks`, `sig_select` (where ~91% of `gpu_submit` lives), `sig_assemble`, + // `sig_row_reduce`, `sig_quasi_inverse` — reads only `full_reuse`/the differentials and + // never touches `dxs` or `xs`, so it CAN legally run several signatures at a time. + // + // DO NOT BOTHER: that was built (a windowed prepare stage feeding an ordered lift) and it is + // worthless, because the work is not spread across the signatures. Per-bidegree `step` span + // times over a stem-150 run, 10 738 bidegrees with >= 2 signatures: + // + // sum of per-bidegree TOTAL signature time 4064.7 s + // sum of per-bidegree MAX signature time 3895.7 s -> ceiling = 1.04x + // + // One signature is ~96% of its bidegree, so the ideal speedup from parallelising the loop is + // 4%. Measured end to end it was worse than that: a window of 4 ran 537 s against a 528 s + // baseline while raising mean CPU 918% -> 1306%, i.e. it burned 42% more CPU to lose 1.7%. + // The parallelism this resolution is missing is NOT inside a bidegree. + let dx_snapshot: Option> = + std::env::var_os("NASSAU_PROBE_SIG_INDEP").map(|_| dxs.clone()); + let mut probe_reads = 0usize; + let mut probe_perturbed = 0usize; + for signature in subalgebra.iter_signatures(b.t()) { let _guard = tracing::info_span!("step", ?signature).entered(); + // Spans below split what used to be one opaque `step`: the run's own accounting put + // ~26% of worker time inside `step` but outside any named region, which is exactly the + // shape that produced several wrong diagnoses earlier. One span per signature is cheap + // (the bodies are substantial); do NOT push spans inside these loops. + let _sm = tracing::trace_span!("sig_masks").entered(); target_mask.clear(); next_mask.clear(); target_mask.extend(subalgebra.signature_mask( @@ -858,34 +980,64 @@ impl> Resolution { &signature, next_bound, )); + drop(_sm); - let full_matrix = match &full_reuse { - Some(full) => { - debug_assert!(target_mask.iter().all(|&r| r < full.rows())); - select_rows(full, &target_mask) - } - None => { - restricted_partial_matrix_maybe_gpu( + let full_matrix = tracing::trace_span!("sig_select", rows = target_mask.len()) + .in_scope(|| match &full_reuse { + Some(full) => { + debug_assert!(target_mask.iter().all(|&r| r < full.rows())); + select_rows(full, &target_mask) + } + None => restricted_partial_matrix_maybe_gpu( &self.differentials[b.s() - 1], b.t(), &target_mask, next_dim, - ) - } - }; + ), + }); - let mut masked_matrix = - AugmentedMatrix::new(p, target_mask.len(), [next_mask.len(), target_mask.len()]); - masked_matrix - .segment(0, 0) - .add_masked(&full_matrix, &next_mask); - masked_matrix.segment(1, 1).add_identity(); - masked_matrix.row_reduce(); + let mut masked_matrix = tracing::trace_span!( + "sig_assemble", + rows = target_mask.len(), + cols = next_mask.len() + ) + .in_scope(|| { + let mut m = AugmentedMatrix::new( + p, + target_mask.len(), + [next_mask.len(), target_mask.len()], + ); + m.segment(0, 0).add_masked(&full_matrix, &next_mask); + m.segment(1, 1).add_identity(); + m + }); + + // The CPU row reduction, once per signature. `gpu_row_reduce` only takes over at + // >= 8192^2, so every one of these is host work. + tracing::trace_span!( + "sig_row_reduce", + rows = target_mask.len(), + cols = next_mask.len() + ) + .in_scope(|| masked_matrix.row_reduce()); - let qi = masked_matrix.compute_quasi_inverse(); + let qi = tracing::trace_span!("sig_quasi_inverse") + .in_scope(|| masked_matrix.compute_quasi_inverse()); let pivots = qi.pivots().unwrap(); let preimage = qi.preimage(); + if let Some(snap) = &dx_snapshot { + for (dx, dx0) in dxs.iter().zip(snap) { + for &v in &next_mask { + probe_reads += 1; + if dx.entry(v) != dx0.entry(v) { + probe_perturbed += 1; + } + } + } + } + + let _lift = tracing::trace_span!("sig_lift", gens = xs.len()).entered(); for (x, dx) in xs.iter_mut().zip(&mut dxs) { scratch.set_scratch_vector_size(target_mask.len()); let mut row = 0; @@ -903,15 +1055,25 @@ impl> Resolution { dx.as_slice_mut().add(full_matrix.row(i), 1); } } - Self::write_qi( - &mut f, - &mut scratch, - &signature, - &next_mask, - &full_matrix, - &masked_matrix, - )?; + drop(_lift); + tracing::trace_span!("sig_write_qi").in_scope(|| { + Self::write_qi( + &mut f, + &mut scratch, + &signature, + &next_mask, + &full_matrix, + &masked_matrix, + ) + })?; + } + if dx_snapshot.is_some() { + eprintln!( + "[sig-probe] b={b} signatures_read_positions={probe_reads} \ + perturbed_by_earlier_signature={probe_perturbed}" + ); } + for dx in &dxs { assert!(dx.is_zero(), "dx non-zero at {b}"); } @@ -955,9 +1117,7 @@ impl> Resolution { source_dim + target_dim, 0, ); - { - chain_map.get_matrix(matrix.segment(0, 0), t); - } + chain_map.get_matrix(matrix.segment(0, 0), t); matrix.segment(1, 1).add_identity(); matrix.row_reduce(); @@ -993,14 +1153,31 @@ impl> Resolution { let source_dim = source_module.dimension(t); let target_dim = target_module.dimension(t); - let mut matrix = - AugmentedMatrix::<2>::new(p, target_dim, [cc_module.dimension(t), target_dim]); - { + // The desired image is the kernel of the augmentation `target_module -> cc_module` in this + // degree. Whenever the target complex is empty in degree `t` that map has a zero-dimensional + // codomain, so its kernel is the whole space and no computation can discover otherwise. + // Taking it directly skips building and row-reducing a `target_dim x target_dim` augmented + // identity, where `target_dim` is `dim(A_t)` -- large enough at high `t` to reach the GPU + // RREF path. + // + // The guard is on the codomain being empty, not on which module is being resolved, so it is + // not a sphere special case: every finite target module is concentrated in finitely many + // degrees, so past its top cell this holds for all `t`, which is almost the whole + // resolution. The sphere is only the extreme of it, firing from `t = 1`. + // + // When the codomain is NON-empty the reduction is still needed, but note it is wasteful + // there too: the kernel has codimension at most `cc_module.dimension(t)`, typically a + // handful, yet we materialise a full `(target_dim - c) x target_dim` basis for it. + let desired_image = if cc_module.dimension(t) == 0 { + Subspace::entire_space(p, target_dim) + } else { + let mut matrix = + AugmentedMatrix::<2>::new(p, target_dim, [cc_module.dimension(t), target_dim]); self.chain_maps[0].get_matrix(matrix.segment(0, 0), t); - } - matrix.segment(1, 1).add_identity(); - matrix.row_reduce(); - let desired_image = matrix.compute_kernel(); + matrix.segment(1, 1).add_identity(); + matrix.row_reduce(); + matrix.compute_kernel() + }; let mut matrix = AugmentedMatrix::<2>::new_with_capacity( p, @@ -1009,9 +1186,7 @@ impl> Resolution { source_dim + MAX_NEW_GENS, 0, ); - { - self.differentials[1].get_matrix(matrix.segment(0, 0), t); - } + self.differentials[1].get_matrix(matrix.segment(0, 0), t); matrix.segment(1, 1).add_identity(); matrix.row_reduce(); @@ -1099,6 +1274,24 @@ impl> Resolution { } fn step_resolution(&self, b: Bidegree) { + // One guard for the whole bidegree, rather than one per inner parallel section. + // + // This is correct by construction rather than by audit. A `step_resolution` job can only be + // stolen onto a thread that is in rayon's steal loop, i.e. blocked at a join — and running + // it there is exactly the priority inversion. So a bounce never discards useful work: it + // declines precisely the runs that would invert. Holding the guard for the whole bidegree + // therefore costs nothing and removes the need to know which callee happens to enter rayon + // today, which the narrow per-section guards did depend on. + // + // Nesting is free: [`ParallelGuard`] counts depth, so the inner guards become depth 1+ and + // keep their spans. + // + // Measurement note, for whoever revisits this: a single stem-200 A/B showed the worst step + // improving (224 s -> 87 s) but the count of steps >=20 s rising (41 -> 67) and retries + // rising (1363 -> 1613). That comparison is NOT conclusive — two runs differing in nothing + // relevant to guarding moved 31 -> 41 and 1271 s -> 1969 s, so the noise floor is the same + // size as the effect. Do not "fix" this on one run's numbers. + let _guard = ParallelGuard::new(); self.step_resolution_with_result(b) .unwrap_or_else(|e| panic!("Error computing bidegree {b}: {e}")); } @@ -1155,13 +1348,17 @@ impl> Resolution { let (sender, receiver) = mpsc::channel(); - let f = |b: Bidegree, sender: mpsc::Sender| { + let spawn_bidegree = |b: Bidegree, sender: mpsc::Sender| { if self.has_computed_bidegree(b) { SenderData::send(b, sender); } else { let tracing_span = tracing_span.clone(); scope.spawn(move |_| { let _tracing_guard = tracing_span.enter(); + if crate::utils::parallel::is_in_parallel() { + SenderData::send_retry(b, sender); + return; + } self.step_resolution(b); SenderData::send(b, sender); }); @@ -1173,32 +1370,142 @@ impl> Resolution { // diagonal predecessor `(0, min_degree)` is in region, so we let it be spawned instead. for s in 0..=max_s { if s != 1 { - f(Bidegree::s_t(s, min_degree), sender.clone()); + spawn_bidegree(Bidegree::s_t(s, min_degree), sender.clone()); } } drop(sender); - while let Ok(SenderData { b, sender }) = receiver.recv() { - assert!(progress[b.s() as usize] == b.t() - 1); - progress[b.s() as usize] = b.t(); - - // Completing `b` can only make ready its same-row successor `(s, t + 1)` and one - // diagonal successor. `ready` requires *both* predecessors, so of the two - // completions that could spawn a given bidegree, only the later one does. - let same_row = b + Bidegree::s_t(0, 1); - let diagonal = if b.s() == 0 { - Bidegree::s_t(1, b.t()) + // Bidegrees whose spawned job was stolen onto a worker already inside a critical section + // (`is_in_parallel` set on that worker) and so bounced back a retry rather than causing + // a priority inversion. Because the check is per-thread, a job is only ever bounced when + // its worker is a blocked guard holder; a job picked up by a free worker just runs. Such + // bounces are therefore rare, but when the pool is momentarily saturated we still must + // avoid re-spawning immediately in a tight loop, so we park bounced bidegrees here. + // + // The scheduler thread never holds a guard, so it cannot itself observe when a worker + // frees; instead, while anything is parked we wait on the channel with a short timeout + // and retry the parked work whenever a completion arrives (a worker likely just freed) + // or the timeout elapses (periodic re-check). Incoming messages are still handled the + // instant they arrive; the timeout only governs how promptly we retry while otherwise + // idle. This cannot deadlock: parked entries keep their senders, so the channel stays + // open, and the timeout guarantees parked work is retried until a free worker takes it. + let mut deferred: Vec<(Bidegree, mpsc::Sender)> = Vec::new(); + // Diagnostic (`NASSAU_MEM_REPORT`): count committed bidegrees so we can periodically + // report the retained-data heap split (differentials' `outputs` vs modules' tables). + let mem_report = std::env::var_os("NASSAU_MEM_REPORT").is_some(); + let mut commit_count = 0usize; + // How long to wait for a message before retrying parked bidegrees. Small enough that a + // freed worker is used promptly, large enough that the poll is negligible; it only ticks + // while something is parked. + const RETRY_POLL_INTERVAL: std::time::Duration = std::time::Duration::from_micros(100); + + loop { + let event = if deferred.is_empty() { + // Nothing parked: block until a message arrives or all senders drop. + match receiver.recv() { + Ok(data) => Some(data), + Err(_) => break, + } } else { - b + Bidegree::s_t(1, 1) + // Something parked: wake periodically to retry it. Parked entries hold senders, + // so the channel cannot be disconnected here. + match receiver.recv_timeout(RETRY_POLL_INTERVAL) { + Ok(data) => Some(data), + Err(mpsc::RecvTimeoutError::Timeout) => None, + Err(mpsc::RecvTimeoutError::Disconnected) => break, + } }; - for cand in [same_row, diagonal] { - if ready(cand.s(), cand.t(), &progress) { - f(cand, sender.clone()); + if let Some(SenderData { b, retry, sender }) = event { + if retry { + // Park until a worker frees; retried below on a completion or timeout. + deferred.push((b, sender)); + continue; + } + assert!(progress[b.s() as usize] == b.t() - 1); + progress[b.s() as usize] = b.t(); + + if mem_report { + commit_count += 1; + if commit_count % 400 == 0 { + let diff_b: usize = self + .differentials + .iter() + .map(|(_, d)| d.output_heap_bytes()) + .sum(); + let mod_b: usize = + self.modules.iter().map(|(_, m)| m.table_heap_bytes()).sum(); + #[cfg(feature = "gpu")] + let (res_master, res_basis) = + algebra::milnor_gpu::resident_host_bytes(); + #[cfg(not(feature = "gpu"))] + let (res_master, res_basis) = (0usize, 0usize); + #[cfg(feature = "gpu")] + let (dev_master, dev_basis) = algebra::milnor_gpu::resident_dev_bytes(); + #[cfg(feature = "gpu")] + let (dev_pool_use, dev_pool_res) = + algebra::milnor_gpu::cubecl_device_usage(); + #[cfg(not(feature = "gpu"))] + let ((dev_master, dev_basis), (dev_pool_use, dev_pool_res)) = + ((0usize, 0usize), (0u64, 0u64)); + let gb = |x: usize| x as f64 / (1u64 << 30) as f64; + let gbu = |x: u64| x as f64 / (1u64 << 30) as f64; + eprintln!( + "[MEM] commits={commit_count} last_b=({},{}) HOST[diff={:.1} \ + mod={:.1} res_master={:.1} res_basis={:.1}]GB DEV[master={:.1} \ + basis={:.1} cubecl_use={:.1} cubecl_reserved={:.1}]GB", + b.n(), + b.s(), + gb(diff_b), + gb(mod_b), + gb(res_master), + gb(res_basis), + gb(dev_master), + gb(dev_basis), + gbu(dev_pool_use), + gbu(dev_pool_res), + ); + } + } + + // Completing `b` can only make ready its same-row successor `(s, t + 1)` and one + // diagonal successor. `ready` requires *both* predecessors, so of the two + // completions that could spawn a given bidegree, only the later one does. + let same_row = b + Bidegree::s_t(0, 1); + let diagonal = if b.s() == 0 { + Bidegree::s_t(1, b.t()) + } else { + b + Bidegree::s_t(1, 1) + }; + + for cand in [same_row, diagonal] { + if ready(cand.s(), cand.t(), &progress) { + spawn_bidegree(cand, sender.clone()); + } + } + } + + // Retry parked bidegrees — reached after a completion (a worker likely just freed) + // or a timeout (periodic re-check), but not after a retry (which `continue`s above, + // so a bounced job waits out the timeout before being retried). Each re-spawned job + // re-checks its own worker's flag: those on a free worker run, those stolen onto a + // blocked guard holder bounce and are re-parked. This stays cheap because per-thread + // bounces are rare, so `deferred` is normally empty. + if !deferred.is_empty() { + for (b, sender) in std::mem::take(&mut deferred) { + spawn_bidegree(b, sender); } } } }); + + // Eviction probe (`NASSAU_R_STATS`): dump the R-access distribution once the wavefront is done. + #[cfg(feature = "gpu")] + { + algebra::milnor_gpu::dump_r_stats(); + // Which theta would have fit: see `resident_degree_cap`. + algebra::milnor_gpu::dump_master_by_degree(); + } } } @@ -1408,9 +1715,13 @@ impl> Resolution { mask.clear(); // At apply time the resolution is fully computed, so we read the full mask // (no concurrently-growing generators to exclude). - mask.extend( - subalgebra.signature_mask(&algebra, source, b.t(), &signature, i32::MAX), - ); + mask.extend(subalgebra.signature_mask( + &algebra, + source, + b.t(), + &signature, + i32::MAX, + )); scratch0.set_scratch_vector_size(mask.len()); } NassauCommand::Fix => { @@ -1611,6 +1922,11 @@ impl<'a, M: ZeroModule> RecomputeReader<'a, M> { .collect(); let full_matrix = { + // Kept while the per-section guards elsewhere were removed: this path is NOT under the + // whole-bidegree guard in `step_resolution`. `commands_for_signature` runs from + // `RecomputeReader::next`, driven by `apply_quasi_inverse_fallible` — an accessor that + // consumers call outside any bidegree, so nothing upstream has raised the depth. + let _guard = ParallelGuard::new(); restricted_partial_matrix(&self.res.differentials[s], t, &src_mask, self.next_dim) }; let mut masked_matrix = @@ -1818,4 +2134,55 @@ mod tests { vec![vec![0, 1, 0, 0], vec![0, 2, 0, 0], vec![0, 0, 1, 0],] ); } + + /// The packed signature test must agree with the per-entry comparison it replaced, including + /// on signatures that no element can have. Packing those naively would spill bits into the + /// neighbouring field and select unrelated elements. + #[test] + fn packed_signature_matches_per_entry_test() { + // The comparison the packed mask replaced, kept here as the reference. + fn has_signature(profile: &[u8], ppart: PPart, signature: &[PPartEntry]) -> bool { + for (i, (&profile, &signature)) in profile.iter().zip(signature).enumerate() { + if ppart.get(i) & ((1u64 << profile) - 1) as PPartEntry != signature { + return false; + } + } + true + } + + let algebra = MilnorAlgebra::new(TWO, false); + algebra.compute_basis(60); + + for profile in [ + vec![1u8, 1, 1], + vec![4, 3, 2, 1], + vec![2, 0, 3], + // Wider than the fields they constrain. + vec![9, 9, 9, 9], + // Longer than a p-part can be, so the tail entries can never be non-zero. + vec![1; PPart::MAX_LEN + 3], + ] { + let subalgebra = MilnorSubalgebra::new(profile.clone()); + for signature in [ + vec![0; profile.len()], + (0..profile.len()).map(|i| (i % 3) as PPartEntry).collect(), + // An entry too wide for its field, which must match nothing. + (0..profile.len()) + .map(|i| if i == profile.len() - 1 { 255 } else { 0 }) + .collect(), + ] { + let packed = subalgebra.packed_signature(&signature); + for t in 0..=60 { + for &op in algebra.ppart_table(t) { + let expected = has_signature(&profile, op, &signature); + let actual = packed.is_some_and(|(mask, value)| op.bits() & mask == value); + assert_eq!( + actual, expected, + "profile {profile:?}, signature {signature:?}, element {op:?}" + ); + } + } + } + } + } } diff --git a/ext/src/nassau_gpu.rs b/ext/src/nassau_gpu.rs index b1315e7d9b..0c4ebce67a 100644 --- a/ext/src/nassau_gpu.rs +++ b/ext/src/nassau_gpu.rs @@ -11,7 +11,8 @@ //! `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. +//! same layout the CPU path produces — as a limb-wise XOR, since the kernel's little-endian `u32` +//! limbs are byte-identical to `fp`'s `u64` limbs. //! //! Gated behind the `gpu` feature. Callers must ensure //! [`MilnorAlgebra::gpu_multiply_applicable`] (`p = 2`, trivial profile, stable) — the @@ -25,10 +26,48 @@ use algebra::{ homomorphism::{FreeModuleHomomorphism, ModuleHomomorphism}, }, }; -use fp::matrix::Matrix; +use fp::{matrix::Matrix, vector::FpVector}; type NassauDifferential = FreeModuleHomomorphism>; +/// Reinterpret a GPU output row's `u32` limbs as their little-endian bytes. +/// +/// The kernel's `u32` limbs and `fp`'s `u64` limbs are the same bit-vector in the same byte order, +/// so this is a view, not a conversion. `fp`'s own `limb::from_bytes`/`to_bytes` take exactly this +/// shortcut under the same `cfg`; the fallback keeps a big-endian target correct rather than +/// silently wrong. +/// +/// One bulk `memcpy` per row. The first cut wrote `w.to_le_bytes()` into `buf` one `u32` at a time, +/// which a call-graph profile of an uncapped stem-150 run showed as 12.15% of ALL user cycles under +/// `copy_from_slice` — a bounds-checked 4-byte copy per limb, thousands per row, for a region that +/// is already byte-identical. +fn fill_limb_bytes(buf: &mut [u8], limbs: &[u32]) { + #[cfg(target_endian = "little")] + { + // SAFETY: `u32` has no padding or invalid bit patterns, `u8` has alignment 1 (so a `u32` + // pointer is suitably aligned), and the length is the same region measured in bytes. The + // view borrows `limbs` and does not outlive it. + let src: &[u8] = unsafe { + std::slice::from_raw_parts(limbs.as_ptr().cast::(), std::mem::size_of_val(limbs)) + }; + let n = src.len().min(buf.len()); + buf[..n].copy_from_slice(&src[..n]); + buf[n..].fill(0); + } + #[cfg(not(target_endian = "little"))] + { + buf.fill(0); + for (k, &w) in limbs.iter().enumerate() { + let o = k * size_of::(); + if o >= buf.len() { + break; + } + let n = size_of::().min(buf.len() - o); + buf[o..o + n].copy_from_slice(&w.to_le_bytes()[..n]); + } + } +} + /// 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`]. @@ -47,22 +86,33 @@ pub fn applicable(hom: &NassauDifferential) -> bool { 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 p = hom.prime(); 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; - } - } + let out = multiply_batch_on_gpu(&algebra, num_cols, inputs.len(), &products); + // Limb-wise readback; see the equivalent (truncating) loop in + // [`get_partial_matrix_restricted`] for why the byte copy is valid. Here the widths already + // agree, so only the partial final limb needs masking. + let num_limbs = FpVector::num_limbs(p, num_cols); + let nbytes = num_limbs * size_of::(); + let mut scratch = FpVector::new(p, num_cols); + let mut buf: Vec = vec![0; nbytes]; + let tail_mask: u64 = match num_cols % 64 { + 0 => u64::MAX, + r => (1u64 << r) - 1, + }; + for (row, limbs) in out.iter_rows().enumerate() { + fill_limb_bytes(&mut buf, limbs); + let last = nbytes - size_of::(); + let masked = u64::from_le_bytes(buf[last..].try_into().unwrap()) & tail_mask; + buf[last..].copy_from_slice(&masked.to_le_bytes()); + scratch + .update_from_bytes(&mut &buf[..]) + .expect("readback scratch is exactly num_limbs * 8 bytes"); + matrix.row_mut(row).add(scratch.as_slice(), 1); } } matrix @@ -157,14 +207,38 @@ pub fn get_partial_matrix_verified( /// at/after `target_dim` are dropped (blocks are generator-major and contiguous, and `target_dim` /// falls on a generator boundary, so the whole block is outside), and the kernel is launched with /// `num_cols = target_dim`. Any returned bit `>= target_dim` is masked out defensively. +/// Rows per GPU multiply batch, chosen so the dense readback (`rows × ceil(cols/32) × 4` bytes) +/// stays under `NASSAU_GPU_MAX_READBACK_MB` (default 1024). Bounds the transient host memory of one +/// build regardless of how many rows the bidegree has; ≥ 1. `0` MB disables batching (one call). +fn gpu_rows_per_batch(cols: usize, num_rows: usize) -> usize { + static CAP_BYTES: std::sync::LazyLock = std::sync::LazyLock::new(|| { + std::env::var("NASSAU_GPU_MAX_READBACK_MB") + .ok() + .and_then(|v| v.parse::().ok()) + .unwrap_or(1024) + * (1 << 20) + }); + if *CAP_BYTES == 0 { + return num_rows.max(1); + } + let bytes_per_row = cols.div_ceil(32) * 4; // one row's readback in bytes + (*CAP_BYTES / bytes_per_row.max(1)).clamp(1, num_rows.max(1)) +} + pub fn get_partial_matrix_restricted( hom: &NassauDifferential, degree: i32, inputs: &[usize], target_dim: usize, ) -> Matrix { - let (mut matrix, products) = extract_restricted(hom, degree, inputs, target_dim); + // Spanned because the 279 s stalls land somewhere in this function *before* the multiply + // (which itself measured 40 ms), and neither this build nor the pair pre-pass inside + // `multiply_batch_on_gpu` was previously visible to the log. + let (mut matrix, mut products) = + tracing::info_span!("extract_restricted", inputs = inputs.len(), target_dim) + .in_scope(|| extract_restricted(hom, degree, inputs, target_dim)); if !products.is_empty() { + let p = hom.prime(); let target = hom.target(); let algebra = target.algebra(); // Idempotent + cheap (O(degree · width)); returns immediately once built. @@ -175,22 +249,66 @@ pub fn get_partial_matrix_restricted( // Passing `target_dim` there would truncate `num_limbs` and corrupt the row layout. We // truncate afterwards by masking bits `>= target_dim` when XORing into the matrix. let full_cols = target.dimension(degree); - let rows = multiply_batch_on_gpu(&algebra, full_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; - let col = limb_idx * 32 + b; - // Minimality should keep every bit within the restricted prefix, but mask - // defensively so a stray high bit can never write out of bounds. - if col < target_dim { - target_row.add_basis_element(col, 1); - } - bits &= bits - 1; + // Cap how large a single multiply we hand the GPU: the dense readback (num_rows × num_limbs + // u32) plus the matrix would otherwise both be held for the whole all-rows / zero-signature + // build (~12 GB dense regions at stem 180). Process the rows in batches of ≤ `rows_per_batch` + // so the readback stays bounded and is freed between batches. `products` is built in row + // order (`extract_restricted`), so each batch's products are a contiguous slice; we remap + // their `row` to batch-local (0-based) for the kernel and write back to the global rows. + let rows_per_batch = gpu_rows_per_batch(full_cols, inputs.len()); + // Readback scratch, allocated once for the whole call (`target_dim` is fixed): the GPU's + // per-row output is XORed into the matrix through a limb-wise `add` rather than bit by bit. + // + // Both sides are little-endian packed F_2 bitvectors with bit `i` = column `i`, so four of + // the kernel's `u32` limbs ARE one of `fp`'s `u64` limbs, byte for byte — no transposition, + // just a truncating copy. `update_from_bytes` fills the existing limbs in place (no + // allocation, no resize), and `read_exact` demands exactly `num_limbs * 8` bytes, which is + // why `buf` is sized once and refilled rather than sliced per row. + // + // The bit-at-a-time loop this replaces called `add_basis_element` once per set bit: at the + // logged ~26% density that is ~0.26 * cols read-modify-writes per row against cols/64 limb + // XORs here, and each one was a bounds-checked entry write rather than a word XOR. + let num_limbs = FpVector::num_limbs(p, target_dim); + let nbytes = num_limbs * size_of::(); + let mut scratch = FpVector::new(p, target_dim); + let mut buf: Vec = vec![0; nbytes]; + // Bits at or past `target_dim` inside the final limb must not survive into the vector — + // `FpVector` requires them zero, and dropping them is exactly what the old `col < target_dim` + // guard did. Whole limbs past the end are dropped by `buf` being only `nbytes` long. + let tail_mask: u64 = match target_dim % 64 { + 0 => u64::MAX, + r => (1u64 << r) - 1, + }; + let mut p0 = 0usize; + let mut r0 = 0usize; + while r0 < inputs.len() { + let r1 = (r0 + rows_per_batch).min(inputs.len()); + let mut p1 = p0; + while p1 < products.len() && products[p1].row < r1 { + p1 += 1; + } + if p1 > p0 { + for pr in &mut products[p0..p1] { + pr.row -= r0; // batch-local row index for the kernel's output layout + } + let out = multiply_batch_on_gpu(&algebra, full_cols, r1 - r0, &products[p0..p1]); + let _scatter = tracing::info_span!("gpu_readback", rows = r1 - r0).entered(); + for (bi, limbs) in out.iter_rows().enumerate() { + // Reinterpret this row's `u32` limbs as the vector's little-endian limb bytes, + // truncated at `target_dim` (both directions: partial final limb, and whole + // limbs past the restricted prefix). + fill_limb_bytes(&mut buf, limbs); + let last = nbytes - size_of::(); + let masked = u64::from_le_bytes(buf[last..].try_into().unwrap()) & tail_mask; + buf[last..].copy_from_slice(&masked.to_le_bytes()); + scratch + .update_from_bytes(&mut &buf[..]) + .expect("readback scratch is exactly num_limbs * 8 bytes"); + matrix.row_mut(r0 + bi).add(scratch.as_slice(), 1); } } + p0 = p1; + r0 = r1; } } matrix @@ -277,9 +395,10 @@ pub fn get_partial_matrix_restricted_verified( 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 restricted get_partial_matrix mismatch at degree {degree}, row {row} \ - (input {}, target_dim {target_dim}, num_rows {})", + g, + c, + "GPU/CPU restricted get_partial_matrix mismatch at degree {degree}, row {row} (input \ + {}, target_dim {target_dim}, num_rows {})", inputs[row], inputs.len(), ); diff --git a/ext/src/resolution.rs b/ext/src/resolution.rs index a8b3c14e01..6f6c711194 100644 --- a/ext/src/resolution.rs +++ b/ext/src/resolution.rs @@ -24,6 +24,7 @@ use sseq::coordinates::{Bidegree, BidegreeGenerator}; use crate::{ chain_complex::{AugmentedChainComplex, ChainComplex}, save::{SaveDirectory, SaveKind}, + utils::parallel::ParallelGuard, }; #[derive(Serialize, Deserialize)] @@ -42,6 +43,8 @@ struct SenderData { b: Bidegree, /// Whether this bidegree was newly calculated or have already been calculated. new: bool, + /// Whether this job should be retried due to priority inversion avoidance. + retry: bool, /// The sender object used to send the `SenderData`. We put this in the struct and pass it /// around the mpsc, so that when all senders are dropped, we know the computation has /// completed. Compared to keeping track of calculations manually, this has the advantage of @@ -55,6 +58,18 @@ impl SenderData { .send(Self { b, new, + retry: false, + sender: sender.clone(), + }) + .unwrap() + } + + fn send_retry(b: Bidegree, sender: mpsc::Sender) { + sender + .send(Self { + b, + new: false, + retry: true, sender: sender.clone(), }) .unwrap() @@ -258,6 +273,7 @@ where ); { + let _guard = ParallelGuard::new(); current_chain_map.get_matrix(matrix.segment(0, 0), b.t()); current_differential.get_matrix(matrix.segment(1, 1), b.t()); } @@ -487,6 +503,7 @@ where // Get the map (d, f) : X_{s, t} -> X_{s-1, t} (+) C_{s, t} into matrix { + let _guard = ParallelGuard::new(); current_chain_map.get_matrix(matrix.segment(0, 0), b.t()); current_differential.get_matrix(matrix.segment(1, 1), b.t()); } @@ -745,13 +762,27 @@ where let tracing_span = tracing_span.clone(); scope.spawn(move |_| { let _tracing_guard = tracing_span.enter(); + if crate::utils::parallel::is_in_parallel() { + SenderData::send_retry(b, sender); + return; + } self.step_resolution(b); SenderData::send(b, true, sender); }); } }; - while let Ok(SenderData { b, new, sender }) = receiver.recv() { + while let Ok(SenderData { + b, + new, + retry, + sender, + }) = receiver.recv() + { + if retry { + f(b, sender); + continue; + } assert!(progress[b.s() as usize] == b.t() - 1); progress[b.s() as usize] = b.t(); @@ -803,13 +834,27 @@ where let tracing_span = tracing_span.clone(); scope.spawn(move |_| { let _tracing_guard = tracing_span.enter(); + if crate::utils::parallel::is_in_parallel() { + SenderData::send_retry(b, sender); + return; + } self.step_resolution(b); SenderData::send(b, true, sender); }); } }; - while let Ok(SenderData { b, new, sender }) = receiver.recv() { + while let Ok(SenderData { + b, + new, + retry, + sender, + }) = receiver.recv() + { + if retry { + f(b, sender); + continue; + } assert!(progress[b.s() as usize] == b.t() - 1); progress[b.s() as usize] = b.t(); diff --git a/ext/src/utils.rs b/ext/src/utils.rs index 494b870a91..c077a95587 100644 --- a/ext/src/utils.rs +++ b/ext/src/utils.rs @@ -619,6 +619,11 @@ mod logging { .with_max_level(tracing::Level::INFO) .with_span_events(FmtSpan::NEW | FmtSpan::CLOSE) .with_thread_ids(true) + // Names too, not just ids. The GPU workers are named `nassau-gpu`, so without + // this a line from the device path reads `ThreadId(37)` and there is no way to tell + // WHICH device it came from — exactly what you need when checking shard balance. + // Rayon's pool threads are unnamed and simply print an empty name. + .with_thread_names(true) .with_env_filter(EnvFilter::try_from_default_env().unwrap_or_default()) .finish() } @@ -638,6 +643,100 @@ mod logging { pub use logging::{LogWriter, ext_tracing_subscriber, init_logging}; +pub(crate) mod parallel { + + use std::cell::Cell; + + thread_local! { + /// Depth of `par_iter_mut` critical sections currently entered *on this thread*. + /// + /// The priority inversion we guard against is narrow: a `step_resolution` job initiates a + /// `par_iter` on some rayon worker, which then blocks in the join and work-steals to stay + /// busy. If that blocked worker steals another (heavy, itself nested-parallel) resolution + /// step, that step runs on it and stalls the critical section it is blocked on. A stolen + /// job runs on the *same* OS thread as the worker that stole it, so a per-thread depth is + /// exactly the right signal: [`is_in_parallel`] reports whether *this* worker is a blocked + /// guard holder. A job picked up by any other worker — idle, or busy on non-critical work — + /// reads zero and is free to run, which is what lets independent bidegrees resolve + /// concurrently. + /// + /// Deliberately per-thread rather than a global count of active critical sections: a global + /// flag blocks *all* new work whenever *any* thread is in a critical section, which under + /// the relaxed wavefront (many bidegrees in flight) is nearly always, producing a retry + /// storm that pegs every core doing no useful work. + static PARALLEL_DEPTH: Cell = const { Cell::new(0) }; + } + + /// RAII guard that increments this thread's [`PARALLEL_DEPTH`] on creation and decrements it on + /// drop. Used to mark regions where a stolen `step_resolution` job would cause a priority + /// inversion, so it can be bounced back instead (see `nassau::step_resolution`). + /// + /// Carries no tracing span: one is taken per bidegree (and per recompute), not per inner + /// parallel section, so the span added log volume proportional to the signature count — over a + /// thousand span pairs per bidegree — for no diagnostic value the enclosing `step` span does + /// not already provide. + pub(crate) struct ParallelGuard { + _private: (), + } + + impl ParallelGuard { + pub(crate) fn new() -> Self { + PARALLEL_DEPTH.with(|d| d.set(d.get() + 1)); + Self { _private: () } + } + } + + impl Drop for ParallelGuard { + fn drop(&mut self) { + PARALLEL_DEPTH.with(|d| d.set(d.get() - 1)); + } + } + + /// Whether the *current* thread is inside a `par_iter_mut` critical section, i.e. whether it is + /// a blocked guard holder onto which stealing a resolution step would cause a priority + /// inversion. See [`PARALLEL_DEPTH`]. + pub(crate) fn is_in_parallel() -> bool { + PARALLEL_DEPTH.with(|d| d.get() > 0) + } + + #[cfg(test)] + mod tests { + use std::sync::mpsc; + + use super::{ParallelGuard, is_in_parallel}; + + /// A [`ParallelGuard`] held on one thread must not be visible on another: the whole point of + /// making `PARALLEL_DEPTH` thread-local is that a resolution step stolen onto a free worker + /// reads zero. Guards against a regression to a shared counter. + #[test] + fn parallel_guard_is_thread_local() { + assert!(!is_in_parallel()); + + let (held_tx, held_rx) = mpsc::channel(); + let (release_tx, release_rx) = mpsc::channel(); + + let handle = std::thread::spawn(move || { + let guard = ParallelGuard::new(); + // Visible on the thread that holds it. + assert!(is_in_parallel()); + held_tx.send(()).unwrap(); + // Keep the guard alive until the main thread has checked. + release_rx.recv().unwrap(); + drop(guard); + // Cleared once dropped. + assert!(!is_in_parallel()); + }); + + // Once the other thread holds the guard, it must be invisible here. + held_rx.recv().unwrap(); + assert!(!is_in_parallel()); + release_tx.send(()).unwrap(); + + handle.join().unwrap(); + } + } +} + /// The value of the SECONDARY_JOB environment variable. /// /// This is used for distributing the `secondary`. If set, only data with `s = SECONDARY_JOB` will diff --git a/ext/src/yoneda.rs b/ext/src/yoneda.rs index 2fffe713c9..c2abd2b23f 100644 --- a/ext/src/yoneda.rs +++ b/ext/src/yoneda.rs @@ -60,7 +60,7 @@ fn rate_milnor_operation(algebra: &MilnorAlgebra, deg: i32, idx: usize) -> i32 { elt.p_part .iter() .enumerate() - .map(|(i, &r)| r.count_ones() << i) + .map(|(i, r)| r.count_ones() << i) .sum::() as i32 }