milnor_gpu: segmented no-copy resident master/basis (kill the realloc-doubling spike)

.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

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"
+    }
+  }
+}

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"
+  }
+}

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<Receiver>`; 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"]

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<const MOD4: bool>(
     g: &mut BenchmarkGroup<WallTime>,
     name: &str,
     p: u32,
-    r: Vec<PPartEntry>,
-    s: Vec<PPartEntry>,
+    r: PPart,
+    s: PPart,
 ) {
     let p = ValidPrime::new(p);
     g.bench_function(name, |bench| {
@@ -21,7 +21,7 @@ fn bench_ppart<const MOD4: bool>(
         bench.iter_batched(
             PPartAllocation::default,
             |alloc| {
-                let m = PPartMultiplier::<MOD4>::new_from_allocation(p, &r, &s, alloc, 0, 0);
+                let m = PPartMultiplier::<MOD4>::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::<false>(
         &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::<true>(
         &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::<true>(
         &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::<false>(
             &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]),
         );
     }
 

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<usize> {
+    let mut order: Vec<usize> = (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 &degree 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);