From 83819efa19ef131250aae5f8ee17d16ecd8f8514 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Mon, 20 Jul 2026 01:13:17 -0400 Subject: [PATCH 001/127] Revert "nassau: remove the ParallelGuard priority-inversion retry mechanism" This reverts commit 71a21b60bd452a1c4d50ff635923363ef9401bb6. --- ext/src/nassau.rs | 31 ++++++++++++++++++++++++++- ext/src/resolution.rs | 49 +++++++++++++++++++++++++++++++++++++++++-- ext/src/utils.rs | 42 +++++++++++++++++++++++++++++++++++++ 3 files changed, 119 insertions(+), 3 deletions(-) diff --git a/ext/src/nassau.rs b/ext/src/nassau.rs index 28e1886e91..7330f4eb87 100644 --- a/ext/src/nassau.rs +++ b/ext/src/nassau.rs @@ -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() @@ -753,6 +767,7 @@ impl> Resolution { // 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 all_rows: Vec = (0..target_dim).collect(); + let _guard = ParallelGuard::new(); Some(restricted_partial_matrix_maybe_gpu( &self.differentials[b.s() - 1], b.t(), @@ -769,6 +784,7 @@ impl> Resolution { select_rows(full, &target_mask) } None => { + let _guard = ParallelGuard::new(); restricted_partial_matrix_maybe_gpu( &self.differentials[b.s() - 1], b.t(), @@ -865,6 +881,7 @@ impl> Resolution { select_rows(full, &target_mask) } None => { + let _guard = ParallelGuard::new(); restricted_partial_matrix_maybe_gpu( &self.differentials[b.s() - 1], b.t(), @@ -956,6 +973,7 @@ impl> Resolution { 0, ); { + let _guard = ParallelGuard::new(); chain_map.get_matrix(matrix.segment(0, 0), t); } matrix.segment(1, 1).add_identity(); @@ -996,6 +1014,7 @@ impl> Resolution { let mut matrix = AugmentedMatrix::<2>::new(p, target_dim, [cc_module.dimension(t), target_dim]); { + let _guard = ParallelGuard::new(); self.chain_maps[0].get_matrix(matrix.segment(0, 0), t); } matrix.segment(1, 1).add_identity(); @@ -1010,6 +1029,7 @@ impl> Resolution { 0, ); { + let _guard = ParallelGuard::new(); self.differentials[1].get_matrix(matrix.segment(0, 0), t); } matrix.segment(1, 1).add_identity(); @@ -1162,6 +1182,10 @@ impl> Resolution { 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); }); @@ -1178,7 +1202,11 @@ impl> Resolution { } drop(sender); - while let Ok(SenderData { b, sender }) = receiver.recv() { + while let Ok(SenderData { b, 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(); @@ -1611,6 +1639,7 @@ impl<'a, M: ZeroModule> RecomputeReader<'a, M> { .collect(); let full_matrix = { + let _guard = ParallelGuard::new(); restricted_partial_matrix(&self.res.differentials[s], t, &src_mask, self.next_dim) }; let mut masked_matrix = 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..8f2cadd267 100644 --- a/ext/src/utils.rs +++ b/ext/src/utils.rs @@ -638,6 +638,48 @@ mod logging { pub use logging::{LogWriter, ext_tracing_subscriber, init_logging}; +pub(crate) mod parallel { + + use std::sync::atomic::{AtomicUsize, Ordering}; + + static PARALLEL_DEPTH: AtomicUsize = AtomicUsize::new(0); + + /// RAII guard that increments [`PARALLEL_DEPTH`] on creation and decrements on drop. Used to mark + /// regions where `par_iter_mut` work is active, so that `step_resolution` jobs can detect priority + /// inversion and retry. + pub(crate) struct ParallelGuard { + #[allow(dead_code)] + span: tracing::span::EnteredSpan, + } + + impl ParallelGuard { + pub(crate) fn new() -> Self { + // We use Release to synchronize with `is_in_parallel` + let counter_start = PARALLEL_DEPTH.fetch_add(1, Ordering::Release); + Self { + span: tracing::info_span!( + "parallel_guard", + counter_start, + counter_end = tracing::field::Empty + ) + .entered(), + } + } + } + + impl Drop for ParallelGuard { + fn drop(&mut self) { + // We use Release to synchronize with `is_in_parallel` + let counter_end = PARALLEL_DEPTH.fetch_sub(1, Ordering::Release); + self.span.record("counter_end", counter_end - 1); + } + } + + pub(crate) fn is_in_parallel() -> bool { + PARALLEL_DEPTH.load(Ordering::Acquire) > 0 + } +} + /// The value of the SECONDARY_JOB environment variable. /// /// This is used for distributing the `secondary`. If set, only data with `s = SECONDARY_JOB` will From 5beaf2208d0ced6f24db4f04327e03661b899734 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 22 Jul 2026 06:36:40 +0000 Subject: [PATCH 002/127] nassau: park retries instead of busy-spinning on ParallelGuard MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The relaxed wavefront keeps many bidegrees in flight at once, so at any instant it is likely that some job is inside a linear-algebra critical section (ParallelGuard). The scheduler re-spawned a bounced job immediately, which just re-checked is_in_parallel, found it still busy, and bounced again — spawning a whole rayon job per re-check and pegging every core on a retry storm that does no useful work. Instead the receiver checks the flag itself (a cheap atomic load) and parks a bidegree only when the section is genuinely busy. A job acquires and releases its guards many times and spends most of its time outside them, so the section frees far more often than jobs complete; parked bidegrees are therefore re-checked via a short recv_timeout while anything is parked, and re-spawned as soon as the section frees. Incoming messages are still handled the instant they arrive. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01MCUtWj6P6suSZqvATCdg6d --- ext/src/nassau.rs | 66 +++++++++++++++++++++++++++++++++++++++++++---- 1 file changed, 61 insertions(+), 5 deletions(-) diff --git a/ext/src/nassau.rs b/ext/src/nassau.rs index e3f9a68eeb..b4c76f7dac 100644 --- a/ext/src/nassau.rs +++ b/ext/src/nassau.rs @@ -1083,7 +1083,7 @@ 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 { @@ -1105,14 +1105,70 @@ 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, retry, sender }) = receiver.recv() { + // Bidegrees whose spawned job found the linear-algebra critical section busy (see + // `ParallelGuard` and `is_in_parallel`) and bounced back a retry, parked here until it + // frees. The relaxed wavefront keeps many bidegrees in flight at once, so at any instant + // it is fairly likely that *some* job is inside a critical section. The original design + // re-spawned a bounced job immediately, which just re-checked `is_in_parallel`, found it + // still busy, and bounced again — a whole rayon job spawned per re-check, pegging every + // core on a retry storm that does no useful work. Instead the receiver checks the flag + // itself (a cheap atomic load) and only parks when busy. + // + // A job acquires and releases its guards many times and spends most of its time outside + // them, so the critical section frees far more often than jobs complete — it is *not* + // safe to only re-check parked bidegrees on completions (that both wastes the free + // windows between guards and can deadlock if every completion happens to observe the + // flag set). Instead, whenever anything is parked we wait on the channel with a short + // timeout and re-check the flag each time it elapses, so a parked bidegree is re-spawned + // as soon as the section frees. Incoming messages are still handled the instant they + // arrive; the timeout only governs how promptly we notice a free window while idle. + let mut deferred: Vec<(Bidegree, mpsc::Sender)> = Vec::new(); + // How long to wait for a message before re-checking `is_in_parallel` while bidegrees are + // parked. Small enough that a freed critical section is noticed promptly, large enough + // that the poll itself is negligible; it only ticks while something is parked. + const RETRY_POLL_INTERVAL: std::time::Duration = std::time::Duration::from_micros(100); + + loop { + // Re-spawn parked bidegrees as soon as the critical section is free. If a job is + // still inside one, leave them parked; the timed wait below re-checks shortly. This + // cannot deadlock: a bidegree is only ever parked while `is_in_parallel` holds, and + // the flag drops to free whenever the running jobs are all outside their guards (in + // particular once they finish), at which point a re-check wakes the parked work. + if !deferred.is_empty() && !crate::utils::parallel::is_in_parallel() { + for (b, sender) in std::mem::take(&mut deferred) { + spawn_bidegree(b, sender); + } + } + + let SenderData { b, retry, sender } = if deferred.is_empty() { + // Nothing parked: block until a message arrives or all senders drop. + match receiver.recv() { + Ok(data) => data, + Err(_) => break, + } + } else { + // Something parked: wake periodically to re-check the flag. Parked entries hold + // senders, so the channel cannot be disconnected here. + match receiver.recv_timeout(RETRY_POLL_INTERVAL) { + Ok(data) => data, + Err(mpsc::RecvTimeoutError::Timeout) => continue, + Err(mpsc::RecvTimeoutError::Disconnected) => break, + } + }; + if retry { - f(b, sender); + if crate::utils::parallel::is_in_parallel() { + // Still busy: park until the critical section frees. + deferred.push((b, sender)); + } else { + // Free now: re-spawn right away. + spawn_bidegree(b, sender); + } continue; } assert!(progress[b.s() as usize] == b.t() - 1); @@ -1130,7 +1186,7 @@ impl> Resolution { for cand in [same_row, diagonal] { if ready(cand.s(), cand.t(), &progress) { - f(cand, sender.clone()); + spawn_bidegree(cand, sender.clone()); } } } From 80d614ff9c69399056fdc531bfcd037135bb167a Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 22 Jul 2026 16:29:15 +0000 Subject: [PATCH 003/127] nassau: make the parallel-section guard per-thread is_in_parallel was a global count of active par_iter critical sections, so a step_resolution job bounced whenever *any* thread was in one. Under the relaxed wavefront many bidegrees are in flight, so that flag is almost always set and nearly every job bounced, producing the retry churn the parking mitigation only softened. The priority inversion the guard exists to prevent is narrower: a worker that initiated a par_iter blocks in the join and work-steals, and if it steals another (heavy, nested-parallel) resolution step, that step stalls the section the worker is blocked on. A stolen job runs on the stealer's own OS thread, so a thread-local depth counter reports exactly whether *this* worker is a blocked guard holder. Jobs picked up by a free worker read zero and run, letting independent bidegrees resolve concurrently instead of serializing behind any single critical section. The scheduler thread never holds a guard, so it can no longer read the flag to sense saturation; park bounced bidegrees and retry them on each completion or a short recv_timeout tick. Bounces are now rare (only a genuine steal-onto-a-blocked-holder), so the parking path barely engages. The classical scheduler shares the guard and benefits the same way, so its immediate-respawn no longer storms. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01MCUtWj6P6suSZqvATCdg6d --- ext/src/nassau.rs | 109 ++++++++++++++++++++++------------------------ ext/src/utils.rs | 51 +++++++++++++++++----- 2 files changed, 91 insertions(+), 69 deletions(-) diff --git a/ext/src/nassau.rs b/ext/src/nassau.rs index b4c76f7dac..15a9112bd1 100644 --- a/ext/src/nassau.rs +++ b/ext/src/nassau.rs @@ -1110,83 +1110,78 @@ impl> Resolution { } drop(sender); - // Bidegrees whose spawned job found the linear-algebra critical section busy (see - // `ParallelGuard` and `is_in_parallel`) and bounced back a retry, parked here until it - // frees. The relaxed wavefront keeps many bidegrees in flight at once, so at any instant - // it is fairly likely that *some* job is inside a critical section. The original design - // re-spawned a bounced job immediately, which just re-checked `is_in_parallel`, found it - // still busy, and bounced again — a whole rayon job spawned per re-check, pegging every - // core on a retry storm that does no useful work. Instead the receiver checks the flag - // itself (a cheap atomic load) and only parks when busy. + // 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. // - // A job acquires and releases its guards many times and spends most of its time outside - // them, so the critical section frees far more often than jobs complete — it is *not* - // safe to only re-check parked bidegrees on completions (that both wastes the free - // windows between guards and can deadlock if every completion happens to observe the - // flag set). Instead, whenever anything is parked we wait on the channel with a short - // timeout and re-check the flag each time it elapses, so a parked bidegree is re-spawned - // as soon as the section frees. Incoming messages are still handled the instant they - // arrive; the timeout only governs how promptly we notice a free window while idle. + // 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(); - // How long to wait for a message before re-checking `is_in_parallel` while bidegrees are - // parked. Small enough that a freed critical section is noticed promptly, large enough - // that the poll itself is negligible; it only ticks while something is parked. + // 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 { - // Re-spawn parked bidegrees as soon as the critical section is free. If a job is - // still inside one, leave them parked; the timed wait below re-checks shortly. This - // cannot deadlock: a bidegree is only ever parked while `is_in_parallel` holds, and - // the flag drops to free whenever the running jobs are all outside their guards (in - // particular once they finish), at which point a re-check wakes the parked work. - if !deferred.is_empty() && !crate::utils::parallel::is_in_parallel() { - for (b, sender) in std::mem::take(&mut deferred) { - spawn_bidegree(b, sender); - } - } - - let SenderData { b, retry, sender } = if deferred.is_empty() { + let event = if deferred.is_empty() { // Nothing parked: block until a message arrives or all senders drop. match receiver.recv() { - Ok(data) => data, + Ok(data) => Some(data), Err(_) => break, } } else { - // Something parked: wake periodically to re-check the flag. Parked entries hold - // senders, so the channel cannot be disconnected here. + // 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) => data, - Err(mpsc::RecvTimeoutError::Timeout) => continue, + Ok(data) => Some(data), + Err(mpsc::RecvTimeoutError::Timeout) => None, Err(mpsc::RecvTimeoutError::Disconnected) => break, } }; - if retry { - if crate::utils::parallel::is_in_parallel() { - // Still busy: park until the critical section frees. + 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(); + + // 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 { - // Free now: re-spawn right away. - spawn_bidegree(b, sender); + b + Bidegree::s_t(1, 1) + }; + + for cand in [same_row, diagonal] { + if ready(cand.s(), cand.t(), &progress) { + spawn_bidegree(cand, sender.clone()); + } } - continue; } - 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()) - } 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); } } } diff --git a/ext/src/utils.rs b/ext/src/utils.rs index ea1c046eea..62c29be945 100644 --- a/ext/src/utils.rs +++ b/ext/src/utils.rs @@ -595,13 +595,31 @@ pub use logging::{LogWriter, ext_tracing_subscriber, init_logging}; pub(crate) mod parallel { - use std::sync::atomic::{AtomicUsize, Ordering}; - - static PARALLEL_DEPTH: AtomicUsize = AtomicUsize::new(0); + 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 [`PARALLEL_DEPTH`] on creation and decrements on drop. Used to mark - /// regions where `par_iter_mut` work is active, so that `step_resolution` jobs can detect priority - /// inversion and retry. + /// RAII guard that increments this thread's [`PARALLEL_DEPTH`] on creation and decrements it on + /// drop. Used to mark regions where `par_iter_mut` work is active, so that a `step_resolution` + /// job stolen onto a blocked guard holder can detect the priority inversion and retry. pub(crate) struct ParallelGuard { #[allow(dead_code)] span: tracing::span::EnteredSpan, @@ -609,8 +627,11 @@ pub(crate) mod parallel { impl ParallelGuard { pub(crate) fn new() -> Self { - // We use Release to synchronize with `is_in_parallel` - let counter_start = PARALLEL_DEPTH.fetch_add(1, Ordering::Release); + let counter_start = PARALLEL_DEPTH.with(|d| { + let v = d.get(); + d.set(v + 1); + v + }); Self { span: tracing::info_span!( "parallel_guard", @@ -624,14 +645,20 @@ pub(crate) mod parallel { impl Drop for ParallelGuard { fn drop(&mut self) { - // We use Release to synchronize with `is_in_parallel` - let counter_end = PARALLEL_DEPTH.fetch_sub(1, Ordering::Release); - self.span.record("counter_end", counter_end - 1); + let counter_end = PARALLEL_DEPTH.with(|d| { + let v = d.get() - 1; + d.set(v); + v + }); + self.span.record("counter_end", counter_end); } } + /// 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.load(Ordering::Acquire) > 0 + PARALLEL_DEPTH.with(|d| d.get() > 0) } } From 540b31f9a2ebad483a7e2106d650d4c3a56821b8 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 22 Jul 2026 16:36:25 +0000 Subject: [PATCH 004/127] nassau: test that the parallel-section guard is thread-local MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pins the invariant the previous commit relies on — a ParallelGuard held on one thread reads as absent on another — so a future change that reverts to a shared counter fails loudly instead of silently reintroducing the retry storm. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01MCUtWj6P6suSZqvATCdg6d --- ext/src/utils.rs | 37 +++++++++++++++++++++++++++++++++++++ 1 file changed, 37 insertions(+) diff --git a/ext/src/utils.rs b/ext/src/utils.rs index 62c29be945..7225fff0fa 100644 --- a/ext/src/utils.rs +++ b/ext/src/utils.rs @@ -660,6 +660,43 @@ pub(crate) mod parallel { 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. From dbc645ee747240c247f4c087e7b2d52cfb96720c Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Wed, 22 Jul 2026 17:13:47 -0400 Subject: [PATCH 005/127] milnor_gpu: drop the global device mutex; make the resident store thread-local MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The batched multiply serialized every launch behind one RESIDENT mutex held across the whole marshal+upload+kernel+readback section, and additionally pinned all work to CUDA stream 0. With the relaxed dependency graph exposing ~max_s-wide bidegree parallelism, that lock collapsed a ~12-core CPU wavefront to ~2.6 busy cores and left the GPU idle 80% of the time — making NASSAU_GPU=1 a net 1.4x slowdown over CPU-only at stem 130 (193s vs 142s). cubecl 0.10 does not need the lock: a per-device runner thread already serializes all server access (concurrent client calls are memory-safe), and memory pools are per-stream. So: - RESIDENT becomes a thread_local RefCell: each rayon worker keeps its own admissible cache and cs/mk device handles, created and consumed only on the thread (and thus the default per-thread CUDA stream) that owns them, so no handle ever crosses threads and no cross-stream event sync fires. - The GPU_STREAM{value:0}.executes pin is removed; each worker launches on its own default stream, so independent bidegrees marshal and execute concurrently. memory_cleanup now trims only the calling worker's pool. Stem 130 (S_2, s<=152, 16-core H200 box): 193s/2.6 cores (old mutex GPU) and 142s/10 cores (CPU-only) -> 44-49s/5.6 cores. Verified bit-identical to the CPU path with NASSAU_GPU_VERIFY=1 at stem 80 (MIN_WORK=0, every build) and stem 130 (default gate, all offloaded/chunked launches, concurrent workers). Note: concurrency raises peak host memory (concurrent marshal buffers across workers); a 16-worker VERIFY run at stem 130 exceeded a ~48GB cgroup, while normal runs fit comfortably. Bound RAYON_NUM_THREADS if memory-constrained. Co-Authored-By: Claude Opus 4.8 --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 220 +++++++++---------- 1 file changed, 106 insertions(+), 114 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index d93b7255e0..30401863a4 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -23,18 +23,6 @@ use cubecl::{ cuda::{CudaDevice, CudaRuntime}, prelude::*, }; -use cubecl_common::stream_id::StreamId; - -/// The single CUDA stream all GPU work is pinned to (via [`StreamId::executes`]). -/// -/// CubeCL's memory pools are per-stream, and the resolution issues launches from many -/// rayon worker threads (each its own stream). Left alone, each stream's pool retains its -/// freed per-launch buffers (chiefly the hundreds-of-MB `out_h`), and across ~16 streams -/// they accumulate until the 4 GB card OOMs — `memory_cleanup` only trims the *calling* -/// stream's pool. Pinning every launch to one stream gives one pool that each launch's -/// `memory_cleanup` fully reclaims. Value 0 is a valid stream id (the first thread's). -const GPU_STREAM: StreamId = StreamId { value: 0 }; - // Only the `#[cfg(test)]` standalone `seqno_kernel` sizes its working array by this bound; the // production kernels use `WORKING_CAP`. #[cfg(test)] @@ -81,10 +69,7 @@ pub fn take_batch_stats() -> (u64, u64, u64, u64) { ) } -use std::{ - collections::HashMap, - sync::{LazyLock, Mutex}, -}; +use std::{cell::RefCell, collections::HashMap}; use cubecl::server::Handle; @@ -100,7 +85,7 @@ struct RInfo { num_mats: u32, } -/// Process-global resident store of admissible-matrix data, both host- and device-side. +/// Per-thread resident store of admissible-matrix data, both host- and device-side. /// /// Admissible-matrix enumeration is a pure function of `R`'s p-part and the same /// low-degree `R`s recur in essentially every bidegree, so the host master (`col_sums` / @@ -109,12 +94,15 @@ struct RInfo { /// 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. +/// Held in *thread-local* storage (see [`RESIDENT`]), one independent store per rayon worker. +/// Each worker therefore creates and consumes its own `cs_handle`/`mk_handle` on the single +/// thread — and thus the single default CUDA stream — that owns them, so no handle is ever +/// shared across threads. That is what lets us drop the old global `Mutex`: cubecl 0.10's +/// per-device runner thread already serializes server access (making concurrent client calls +/// memory-safe), and keeping every worker's buffers on its own stream avoids cross-stream +/// event synchronization while letting independent bidegrees marshal and launch concurrently. +/// Duplicating the cache per thread is cheap and correct: in the GPU path's regime (`p = 2`, +/// trivial profile) `admissible_matrices` depends only on the p-part, not the algebra instance. #[derive(Default)] struct Resident { col_sums: Vec, @@ -148,7 +136,12 @@ impl Resident { } } -static RESIDENT: LazyLock> = LazyLock::new(|| Mutex::new(Resident::default())); +thread_local! { + /// Per-thread resident admissible-matrix store (see [`Resident`]). Thread-local instead of + /// a shared `Mutex` so independent bidegrees no longer serialize on one lock: each + /// rayon worker keeps its own cache and GPU handles on its own default CUDA stream. + static RESIDENT: RefCell = RefCell::new(Resident::default()); +} /// Elementwise F₂ addition of two bit-packed vectors: `out[i] = a[i] ^ b[i]`. /// @@ -691,8 +684,8 @@ pub fn multiply_batch_on_gpu( prod_r_index.push(ri); } - // Admissible-matrix data (`col_sums`/`masks` + per-`R` offsets) is resident (built - // under the `RESIDENT` lock below), so nothing to enumerate or lay out here. + // 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. // Parallel: each product's term p-parts (padded to `width`) and lengths. let per_prod: Vec<(Vec, Vec)> = (0..products.len()) @@ -713,85 +706,86 @@ pub fn multiply_batch_on_gpu( }) .collect(); - // Resident admissible-matrix store: enumerate each new `R` once and reuse forever; - // the per-`R` offsets are global (into the master `col_sums`/`masks`). Taking the lock - // here also serializes the device section across rayon workers (see [`Resident`]). - let mut resident = RESIDENT.lock().unwrap(); - let mut r_cs_offset: Vec = Vec::with_capacity(distinct_r.len()); - let mut r_mk_offset: Vec = Vec::with_capacity(distinct_r.len()); - let mut r_cs_len: Vec = Vec::with_capacity(distinct_r.len()); - let mut r_mk_len: Vec = Vec::with_capacity(distinct_r.len()); - let mut r_num_matrices: Vec = Vec::with_capacity(distinct_r.len()); - for &(rd, ridx) in &distinct_r { - let r = algebra.basis_element_from_index(rd, ridx); - assert!(!r.p_part.is_empty(), "each R must be non-empty"); - let info = resident.ensure(algebra, &r.p_part); - r_cs_offset.push(info.cs_off); - r_mk_offset.push(info.mk_off); - r_cs_len.push(info.cs_len); - r_mk_len.push(info.mk_len); - r_num_matrices.push(info.num_mats as usize); - } + // Resident admissible-matrix store (thread-local; see [`Resident`]): enumerate each new `R` + // once per worker and reuse forever, with per-`R` offsets into this thread's master + // `col_sums`/`masks`. The whole marshal + device section runs inside this borrow, but it is + // uncontended — no other thread can touch this worker's store. + RESIDENT.with_borrow_mut(|resident| { + 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); - } + // 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} \ + 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(), - ); - } - if total_pairs == 0 { - return vec![vec![0u32; num_limbs]; num_rows]; - } + products.len(), + ); + } + if total_pairs == 0 { + return vec![vec![0u32; num_limbs]; num_rows]; + } - // The resident `col_sums`/`masks` are non-empty once any `R` is present (guaranteed - // here, since `total_pairs > 0`); only `term_pparts` needs the non-empty guard. - if term_pparts.is_empty() { - term_pparts.push(0); - } + // The resident `col_sums`/`masks` are non-empty once any `R` is present (guaranteed + // here, since `total_pairs > 0`); only `term_pparts` needs the non-empty guard. + if term_pparts.is_empty() { + term_pparts.push(0); + } - let marshal_ms = t_marshal.elapsed().as_secs_f64() * 1e3; + let marshal_ms = t_marshal.elapsed().as_secs_f64() * 1e3; - let t_device = std::time::Instant::now(); + let t_device = std::time::Instant::now(); - // Pin the whole device section to one CUDA stream (see [`GPU_STREAM`]) so a single - // memory pool is reclaimed by `memory_cleanup`. Held under the `resident` lock, so this - // stream is used by at most one thread at a time. - let result = GPU_STREAM.executes(|| { + // Device section on this worker's default CUDA stream (distinct per rayon thread, so + // independent bidegrees overlap on the GPU). No lock: `resident` is thread-local, so its + // handles are only ever touched by this thread, and cubecl's per-device runner serializes + // the actual server access. `memory_cleanup` below trims only this stream's own pool. 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. @@ -899,24 +893,22 @@ pub fn multiply_batch_on_gpu( // 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); + // 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 + result + }) } #[cfg(test)] From f1cc8fa34fc7126e49cebfa5d04b06f2dda1641d Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Thu, 23 Jul 2026 11:58:59 -0400 Subject: [PATCH 006/127] milnor_gpu: bound GPU-path memory (row blocks, launch permits, shared resident master) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The mutex-removal commit let many workers run device sections concurrently, which exposed three unbounded memory consumers at record stems (>100GB host AND device by stem 150, measured): 1. Unbounded launch transients: the all-rows reuse build allocated its full output in one shot, per in-flight worker. Fixed by splitting builds into row blocks bounded by NASSAU_GPU_BLOCK_MB (default 512MB) of output AND GPU_PAIR_CHUNK kernel threads — one launch per block, subsuming the former pair-chunk loop (rows are independent, so blocks concatenate exactly). 2. Unbounded stream count: every worker thread got its own CUDA stream, and each stream's pool retains freed slabs indefinitely. Fixed by NASSAU_GPU_CONCURRENCY (default 8) permits that double as stream slots: at most 8 device sections run at once, on 8 fixed streams. A permit must never be held across a rayon parallel section (par_iter chunks execute on guard-free threads that can steal a bidegree job which then parks on acquire — observed deadlock); it is acquired only for the strictly sequential layout+device section. Do NOT raise to 16: measured catastrophic (>30x) slowdown from cross-stream sync churn. 3. Per-thread resident duplication: the thread-local resident store copied the admissible-matrix master (~8.5GB at stem 150, growing with degree) once per worker, on host and device. Fixed by re-sharing it: host master behind an RwLock (enumeration outside the write lock), one device mirror behind a small mutex, handles shared across threads/slots (cubecl event-syncs cross-stream reuse). Re-uploads are needs-based — only when a launch dereferences past the uploaded prefix — since re-uploading on mere growth serialized multi-GB copies on nearly every frontier launch (measured 1.5x wall regression). Stem 150 (S_2, s<=152, 16-core H200 box), verified bit-identical to CPU at stem 80 (every build, forced multi-block) and stem 130 (all offloaded launches): wall cores host RSS device before this commit 682s 3.3 137 GB 140 GB (full card) after (32 workers) 721s 3.8 65 GB 37 GB CPU-only reference 771s 10.3 4.4 GB — Verdict: at record stems the GPU path now merely ties CPU-only while using far more memory — the CPU path (no row-reuse matrix, per-signature builds) is both frugal and wavefront-parallel. Recommend CPU-only for the stem-300 production run; the GPU path remains correct, memory-bounded, and a real win at mid stems (3x at stem 130). Co-Authored-By: Claude Opus 4.8 --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 601 ++++++++++++------- 1 file changed, 383 insertions(+), 218 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 30401863a4..1ad0f22c81 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -23,6 +23,7 @@ use cubecl::{ cuda::{CudaDevice, CudaRuntime}, prelude::*, }; +use cubecl_common::stream_id::StreamId; // Only the `#[cfg(test)]` standalone `seqno_kernel` sizes its working array by this bound; the // production kernels use `WORKING_CAP`. #[cfg(test)] @@ -34,14 +35,102 @@ use crate::algebra::{Algebra, MilnorAlgebra, combinatorics::xi_degrees}; /// `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 +/// 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. 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. +/// 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). `1 << 30` (~1.07e9) leaves >3x headroom +/// under `2^32` even when a lone over-budget row overshoots it, and keeps the grid +/// (`total_pairs / 256` cubes) well under CUDA's `2^31 - 1` grid-dimension limit. const GPU_PAIR_CHUNK: usize = 1 << 30; +/// 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_PERMITS`] this makes peak transient +/// memory a configured constant (≈ permits × 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 +} + +/// Counting semaphore bounding how many workers may be inside the layout + device section at +/// once (`NASSAU_GPU_CONCURRENCY`, default 8). With per-thread streams every worker can launch +/// concurrently, which is the throughput win — but each concurrent section holds up to +/// [`gpu_block_bytes`] of transient host and device memory, so the count must be capped. +/// +/// 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 GpuPermits { + free: Mutex>, + freed: Condvar, +} + +static GPU_PERMITS: LazyLock = LazyLock::new(|| { + let max = std::env::var("NASSAU_GPU_CONCURRENCY") + .ok() + .and_then(|v| v.parse::().ok()) + .filter(|&n| n > 0) + .unwrap_or(8); + GpuPermits { + free: Mutex::new((0..max).rev().collect()), + freed: Condvar::new(), + } +}); + +/// RAII permit from [`GPU_PERMITS`]; blocks (parked, not spinning) until one frees. +/// +/// A permit is also a *stream slot*: the holder runs its device section under +/// `StreamId { value: slot }`, so the process only ever touches `NASSAU_GPU_CONCURRENCY` +/// CUDA streams. Without the pin every worker thread gets its own default stream, and each +/// stream's memory pool *retains* its freed slabs (`memory_cleanup` only trims the calling +/// stream, at its next launch) — ~100 worker streams each retaining ~1 GB of out/staging +/// buffers filled the whole 143 GB H200. Pinning bounds device retention to +/// ≈ permits × [`gpu_block_bytes`]. +struct GpuPermit { + slot: usize, +} + +impl GpuPermit { + fn acquire() -> Self { + let permits = &*GPU_PERMITS; + let mut free = permits.free.lock().unwrap(); + loop { + if let Some(slot) = free.pop() { + return Self { slot }; + } + free = permits.freed.wait(free).unwrap(); + } + } +} + +impl Drop for GpuPermit { + fn drop(&mut self) { + let permits = &*GPU_PERMITS; + permits.free.lock().unwrap().push(self.slot); + permits.freed.notify_one(); + } +} + /// Narrow an admissible-matrix / p-part entry to the `u16` the GPU buffers use, failing loudly /// instead of silently wrapping. Every entry is well within `u16` for the stem ranges this path /// targets; a panic here means that assumption was pushed past its limit, which must not ship @@ -69,7 +158,10 @@ pub fn take_batch_stats() -> (u64, u64, u64, u64) { ) } -use std::{cell::RefCell, collections::HashMap}; +use std::{ + collections::HashMap, + sync::{Condvar, LazyLock, Mutex, RwLock}, +}; use cubecl::server::Handle; @@ -85,62 +177,70 @@ struct RInfo { num_mats: u32, } -/// Per-thread 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. +/// and never recomputed. /// -/// Held in *thread-local* storage (see [`RESIDENT`]), one independent store per rayon worker. -/// Each worker therefore creates and consumes its own `cs_handle`/`mk_handle` on the single -/// thread — and thus the single default CUDA stream — that owns them, so no handle is ever -/// shared across threads. That is what lets us drop the old global `Mutex`: cubecl 0.10's -/// per-device runner thread already serializes server access (making concurrent client calls -/// memory-safe), and keeping every worker's buffers on its own stream avoids cross-stream -/// event synchronization while letting independent bidegrees marshal and launch concurrently. -/// Duplicating the cache per thread is cheap and correct: in the GPU path's regime (`p = 2`, -/// trivial profile) `admissible_matrices` depends only on the p-part, not the algebra instance. +/// 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. #[derive(Default)] -struct Resident { +struct ResidentHost { col_sums: Vec, masks: Vec, index: HashMap, RInfo>, +} + +static RESIDENT_HOST: LazyLock> = + LazyLock::new(|| RwLock::new(ResidentHost::default())); + +/// Process-shared device mirror of the host master: one upload for the whole process, +/// re-uploaded only when the master grew. The handles are shared across worker threads and +/// stream slots — safe under cubecl 0.10's per-device runner (which serializes all server +/// access), with cross-stream reuse event-synced via the handle's origin-stream stamp. The +/// mutex guards only the grew-check + upload, held briefly per launch. +#[derive(Default)] +struct ResidentDev { cs_handle: Option, mk_handle: Option, cs_uploaded: usize, mk_uploaded: usize, } -impl Resident { - /// Global offsets/lengths of `R`'s admissible matrices in the master, enumerating and - /// appending them on first sight (the append order fixes the offsets forever). - fn ensure(&mut self, algebra: &MilnorAlgebra, p_part: &[PPartEntry]) -> RInfo { - if let Some(info) = self.index.get(p_part) { - return *info; - } - let (cs_len, mk_len, cs, mk) = algebra.admissible_matrices(p_part); - let info = RInfo { - cs_off: self.col_sums.len() as u32, - mk_off: self.masks.len() as u32, - cs_len: cs_len as u32, - mk_len: mk_len as u32, - num_mats: (mk.len() / mk_len) as u32, - }; - self.col_sums.extend(cs.iter().map(|&v| narrow_u16(v))); - self.masks.extend(mk.iter().map(|&v| narrow_u16(v))); - self.index.insert(p_part.to_vec(), info); - info - } -} +static RESIDENT_DEV: LazyLock> = + LazyLock::new(|| Mutex::new(ResidentDev::default())); -thread_local! { - /// Per-thread resident admissible-matrix store (see [`Resident`]). Thread-local instead of - /// a shared `Mutex` so independent bidegrees no longer serialize on one lock: each - /// rayon worker keeps its own cache and GPU handles on its own default CUDA stream. - static RESIDENT: RefCell = RefCell::new(Resident::default()); +/// 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. +fn resident_info(algebra: &MilnorAlgebra, p_part: &[PPartEntry]) -> RInfo { + if let Some(info) = RESIDENT_HOST.read().unwrap().index.get(p_part) { + return *info; + } + 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; + } + let info = RInfo { + cs_off: host.col_sums.len() as u32, + mk_off: host.masks.len() as u32, + cs_len: cs_len as u32, + mk_len: mk_len as u32, + num_mats: (mk.len() / mk_len) as u32, + }; + host.col_sums.extend(cs.iter().map(|&v| narrow_u16(v))); + host.masks.extend(mk.iter().map(|&v| narrow_u16(v))); + host.index.insert(p_part.to_vec(), info); + info } /// Elementwise F₂ addition of two bit-packed vectors: `out[i] = a[i] ^ b[i]`. @@ -633,10 +733,10 @@ pub struct GpuProduct { 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), @@ -648,6 +748,66 @@ pub fn multiply_batch_on_gpu( num_cols: usize, num_rows: usize, products: &[GpuProduct], +) -> 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), 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). This pre-pass also warms the shared resident master, + // so every block's layout lookups below are read-lock cache hits. + let prod_pairs: Vec = products + .iter() + .map(|prod| { + let r = algebra.basis_element_from_index(prod.r_degree, prod.r_idx); + resident_info(algebra, &r.p_part).num_mats as usize * prod.term_indices.len() + }) + .collect(); + let mut result: Vec> = Vec::with_capacity(num_rows); + 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); + } + result.extend(multiply_batch_block( + algebra, + num_cols, + r0, + r1 - r0, + &products[p0..p1], + )); + (r0, p0) = (r1, p1); + } + result +} + +/// 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. Holds a +/// [`GpuPermit`] for its sequential layout + device section (acquired only after the parallel +/// marshal — see [`GpuPermits`]), so at most `NASSAU_GPU_CONCURRENCY` device sections run at +/// once across all worker threads. +fn multiply_batch_block( + algebra: &MilnorAlgebra, + num_cols: usize, + row_base: usize, + num_rows: usize, + products: &[GpuProduct], ) -> Vec> { let (width, g) = algebra.seqno_table_u32(); let mut xi: Vec = xi_degrees(algebra.prime()) @@ -706,105 +866,138 @@ pub fn multiply_batch_on_gpu( }) .collect(); - // Resident admissible-matrix store (thread-local; see [`Resident`]): enumerate each new `R` - // once per worker and reuse forever, with per-`R` offsets into this thread's master - // `col_sums`/`masks`. The whole marshal + device section runs inside this borrow, but it is - // uncontended — no other thread can touch this worker's store. - RESIDENT.with_borrow_mut(|resident| { - 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); - } + // 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). + let permit = GpuPermit::acquire(); + // 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; + 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_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); + } - // 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); - } + // 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()); + // 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; + 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); + pps.push(pair_acc as u32); + pair_acc += r_num_matrices[ri as usize] * prod.term_indices.len(); + prod_num_terms.push(prod.term_indices.len() as u32); + prod_row_base.push(((prod.row - row_base) * 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(), - ); - } - if total_pairs == 0 { - return vec![vec![0u32; num_limbs]; num_rows]; - } + 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" + ); + pps.push(total_pairs as u32); + let out_len = num_rows * num_limbs; + if std::env::var_os("NASSAU_GPU_DEBUG").is_some() { + eprintln!( + "[gpu-batch] row_base={row_base} num_rows={num_rows} num_cols={num_cols} \ + num_limbs={num_limbs} products={} total_pairs={total_pairs} out_len={out_len}", + products.len(), + ); + } + if total_pairs == 0 { + return vec![vec![0u32; num_limbs]; num_rows]; + } - // The resident `col_sums`/`masks` are non-empty once any `R` is present (guaranteed - // here, since `total_pairs > 0`); only `term_pparts` needs the non-empty guard. - if term_pparts.is_empty() { - term_pparts.push(0); - } + // The resident `col_sums`/`masks` are non-empty once any `R` is present (guaranteed + // here, since `total_pairs > 0`); only `term_pparts` needs the non-empty guard. + if term_pparts.is_empty() { + term_pparts.push(0); + } - let marshal_ms = t_marshal.elapsed().as_secs_f64() * 1e3; + let marshal_ms = t_marshal.elapsed().as_secs_f64() * 1e3; - let t_device = std::time::Instant::now(); + let t_device = std::time::Instant::now(); - // Device section on this worker's default CUDA stream (distinct per rayon thread, so - // independent bidegrees overlap on the GPU). No lock: `resident` is thread-local, so its - // handles are only ever touched by this thread, and cubecl's per-device runner serializes - // the actual server access. `memory_cleanup` below trims only this stream's own pool. + // Device section pinned to this permit's stream slot: up to `NASSAU_GPU_CONCURRENCY` + // launches overlap on distinct streams, but no more streams (and hence retained pools) + // than that ever exist — see [`GpuPermit`]. Cubecl's per-device runner serializes the + // actual server access; cross-slot/cross-thread reuse of the shared resident handles is + // event-synced by cubecl. `memory_cleanup` below trims only this slot's own pool. + let result = StreamId { + value: permit.slot as u64, + } + .executes(|| { let client = CudaRuntime::client(&CudaDevice::default()); - // Resident admissible buffers: (re-)upload the master only when it grew this - // launch; otherwise reuse the handle from a previous launch and upload nothing. - if resident.cs_handle.is_none() || resident.cs_uploaded != resident.col_sums.len() { - resident.cs_handle = Some(client.create_from_slice(u16::as_bytes(&resident.col_sums))); - resident.cs_uploaded = resident.col_sums.len(); - } - if resident.mk_handle.is_none() || resident.mk_uploaded != resident.masks.len() { - resident.mk_handle = Some(client.create_from_slice(u16::as_bytes(&resident.masks))); - resident.mk_uploaded = resident.masks.len(); - } - let cs_len_master = resident.col_sums.len(); - let mk_len_master = resident.masks.len(); - let cs_h = resident.cs_handle.clone().unwrap(); - let mk_h = resident.mk_handle.clone().unwrap(); - // 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. + // Shared resident admissible buffers (see [`ResidentDev`]): re-upload the master ONLY + // when this block dereferences past the uploaded prefix (`need_cs` / `need_mk`). The + // master grows continually at the frontier, so re-uploading on mere growth ships + // multi-GB uploads under this mutex on nearly every launch (measured 1.5x wall + // regression); most launches touch only long-uploaded low-degree `R`s and reuse the + // stale handle at its uploaded length. When an upload does fire it captures the full + // current master, amortizing all growth since the last one. Lock order is DEV.lock + // then HOST.read, and nothing under either lock blocks on rayon or a permit. The + // master is append-only, so the uploaded prefix is always a prefix of the current + // host master and every offset `< uploaded` is final. + let (cs_h, mk_h, cs_len_master, mk_len_master) = { + let mut dev = RESIDENT_DEV.lock().unwrap(); + if dev.cs_handle.is_none() || dev.cs_uploaded < need_cs { + let host = RESIDENT_HOST.read().unwrap(); + dev.cs_handle = Some(client.create_from_slice(u16::as_bytes(&host.col_sums))); + dev.cs_uploaded = host.col_sums.len(); + } + if dev.mk_handle.is_none() || dev.mk_uploaded < need_mk { + let host = RESIDENT_HOST.read().unwrap(); + dev.mk_handle = Some(client.create_from_slice(u16::as_bytes(&host.masks))); + dev.mk_uploaded = host.masks.len(); + } + ( + dev.cs_handle.clone().unwrap(), + dev.mk_handle.clone().unwrap(), + dev.cs_uploaded, + dev.mk_uploaded, + ) + }; + // 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. 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)); @@ -817,67 +1010,37 @@ pub fn multiply_batch_on_gpu( let out_h = client.create_from_slice(u32::as_bytes(&zeros)); const THREADS: u32 = 256; - // 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, - 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()), - width, - ); - } - - c0 = c1; + let pri_h = client.create_from_slice(u32::as_bytes(&prod_r_index)); + let pts_h = client.create_from_slice(u32::as_bytes(&prod_term_start)); + let pnt_h = client.create_from_slice(u32::as_bytes(&prod_num_terms)); + let prb_h = client.create_from_slice(u32::as_bytes(&prod_row_base)); + let poo_h = client.create_from_slice(u32::as_bytes(&prod_out_offset)); + let pps_h = client.create_from_slice(u32::as_bytes(&pps)); + let cubes = (total_pairs as u32).div_ceil(THREADS).max(1); + unsafe { + multiply_batch_kernel::launch::( + &client, + CubeCount::Static(cubes, 1, 1), + CubeDim::new_1d(THREADS), + ArrayArg::from_raw_parts(cs_h, cs_len_master), + ArrayArg::from_raw_parts(mk_h, mk_len_master), + ArrayArg::from_raw_parts(tp_h, term_pparts.len()), + ArrayArg::from_raw_parts(tl_h, term_lens.len()), + ArrayArg::from_raw_parts(g_h, g.len()), + ArrayArg::from_raw_parts(xi_h, xi.len()), + ArrayArg::from_raw_parts(out_h.clone(), out_len), + ArrayArg::from_raw_parts(rco_h, r_cs_offset.len()), + ArrayArg::from_raw_parts(rmo_h, r_mk_offset.len()), + ArrayArg::from_raw_parts(rcl_h, r_cs_len.len()), + ArrayArg::from_raw_parts(rml_h, r_mk_len.len()), + ArrayArg::from_raw_parts(pri_h, products.len()), + ArrayArg::from_raw_parts(pts_h, products.len()), + ArrayArg::from_raw_parts(pnt_h, products.len()), + ArrayArg::from_raw_parts(prb_h, products.len()), + ArrayArg::from_raw_parts(poo_h, products.len()), + ArrayArg::from_raw_parts(pps_h, pps.len()), + width, + ); } let bytes = client.read_one(out_h).unwrap(); @@ -893,22 +1056,24 @@ pub fn multiply_batch_on_gpu( // resident admissible handles stay alive (refcount > 0) so cleanup skips them. client.memory_cleanup(); - // 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 - }) + }); + + // Aggregate marshal/device totals across every launch (cheap, always on) so a whole + // resolution's GPU overhead can be split host-vs-device via [`take_batch_stats`]. + let device_ms = t_device.elapsed().as_secs_f64() * 1e3; + BATCH_CALLS.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + BATCH_MARSHAL_US.fetch_add( + (marshal_ms * 1e3) as u64, + std::sync::atomic::Ordering::Relaxed, + ); + BATCH_DEVICE_US.fetch_add( + (device_ms * 1e3) as u64, + std::sync::atomic::Ordering::Relaxed, + ); + BATCH_PAIRS.fetch_add(total_pairs as u64, std::sync::atomic::Ordering::Relaxed); + + result } #[cfg(test)] From a0022b7a71062f88a3146c54f20884b49d5e2327 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Thu, 23 Jul 2026 13:47:26 -0400 Subject: [PATCH 007/127] milnor_gpu: byte-weighted launch budget + prefix-doubling master uploads MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replaces the count-based launch cap (NASSAU_GPU_CONCURRENCY=8 exclusive sections) with two decoupled controls: - NASSAU_GPU_MEM_BUDGET_MB (default 4096): admission weighted by a launch's output bytes, so dozens of small low-stem launches run concurrently again (the count cap throttled exactly the region that never had a memory problem) while the frontier stays bounded to ~budget/block-size in flight. - NASSAU_GPU_STREAMS (default 8): fixed CUDA stream slots, round-robin and SHARED (small launches serialize on a stream rather than demanding an exclusive one), so stream/pool count is bounded independently of concurrency. Master device uploads are now prefix-only with doubling: a launch ships max(need, 2*uploaded) entries, not the whole master, so frontier launches (which append new high-degree R each t) no longer re-ship gigabytes of untouched tail. Stem 130 improved 187s -> 150s; stem 150 memory 65/37 -> 68/30 GB, verified bit-identical (stem 80 all-builds, stem 130 all offloaded). But a slots x budget sweep is FLAT (8/4G=150s, 16/8G=170s, 32/16G=160s, 64/32G=201s): concurrency knobs are not the ceiling. The ceiling is Amdahl — the GPU accelerates only the Milnor multiply (~17% of frontier wall time; row_reduce/signature_matrix/readback dominate and are CPU/serial through cubecl's single runner thread), so the end-to-end GPU:CPU ratio is flat ~1.13x across the 130-150 heavy bands, not widening. Widening it would require offloading row_reduce (PR #274's RREF). Co-Authored-By: Claude Opus 4.8 --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 124 ++++++++++++------- 1 file changed, 76 insertions(+), 48 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 1ad0f22c81..5c4a7ff1eb 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -53,8 +53,8 @@ const GPU_PAIR_CHUNK: usize = 1 << 30; /// 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_PERMITS`] this makes peak transient -/// memory a configured constant (≈ permits × budget) instead of a function of the frontier size. +/// 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") @@ -67,10 +67,17 @@ fn gpu_block_bytes() -> usize { *BYTES } -/// Counting semaphore bounding how many workers may be inside the layout + device section at -/// once (`NASSAU_GPU_CONCURRENCY`, default 8). With per-thread streams every worker can launch -/// concurrently, which is the throughput win — but each concurrent section holds up to -/// [`gpu_block_bytes`] of transient host and device memory, so the count must be capped. +/// 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 @@ -80,54 +87,70 @@ fn gpu_block_bytes() -> usize { /// 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 GpuPermits { - free: Mutex>, +struct GpuBudget { + budget: usize, + used: Mutex, freed: Condvar, } -static GPU_PERMITS: LazyLock = LazyLock::new(|| { - let max = std::env::var("NASSAU_GPU_CONCURRENCY") +static GPU_BUDGET: LazyLock = LazyLock::new(|| GpuBudget { + budget: std::env::var("NASSAU_GPU_MEM_BUDGET_MB") .ok() .and_then(|v| v.parse::().ok()) - .filter(|&n| n > 0) - .unwrap_or(8); - GpuPermits { - free: Mutex::new((0..max).rev().collect()), - freed: Condvar::new(), - } + .filter(|&mb| mb > 0) + .unwrap_or(4096) + << 20, + used: Mutex::new(0), + freed: Condvar::new(), }); -/// RAII permit from [`GPU_PERMITS`]; blocks (parked, not spinning) until one frees. -/// -/// A permit is also a *stream slot*: the holder runs its device section under -/// `StreamId { value: slot }`, so the process only ever touches `NASSAU_GPU_CONCURRENCY` -/// CUDA streams. Without the pin every worker thread gets its own default stream, and each -/// stream's memory pool *retains* its freed slabs (`memory_cleanup` only trims the calling -/// stream, at its next launch) — ~100 worker streams each retaining ~1 GB of out/staging -/// buffers filled the whole 143 GB H200. Pinning bounds device retention to -/// ≈ permits × [`gpu_block_bytes`]. +/// Fixed CUDA stream-slot count (`NASSAU_GPU_STREAMS`, default 8): device sections run under +/// `StreamId { value: counter % slots }`, so only this many streams (and hence retained +/// memory pools) ever exist. Without a pin every worker thread gets its own default stream, +/// and each stream's pool *retains* its freed slabs — ~100 worker streams each retaining +/// ~1 GB filled the whole 143 GB H200. Slots are round-robin and *shared*, not exclusive: +/// two small launches on one slot merely serialize on that CUDA stream (memory-safe, and +/// cheap for small kernels), so slot count does not cap concurrency the way the byte budget +/// does. Kept at 8: a 16-slot experiment collapsed >30x (unexplained cross-stream churn). +fn gpu_stream_slots() -> u64 { + static SLOTS: LazyLock = LazyLock::new(|| { + std::env::var("NASSAU_GPU_STREAMS") + .ok() + .and_then(|v| v.parse::().ok()) + .filter(|&n| n > 0) + .unwrap_or(8) + }); + *SLOTS +} + +/// RAII reservation of `weight` bytes from [`GPU_BUDGET`] plus a round-robin stream slot; +/// blocks (parked, not spinning) until the budget admits it. struct GpuPermit { - slot: usize, + weight: usize, + slot: u64, } impl GpuPermit { - fn acquire() -> Self { - let permits = &*GPU_PERMITS; - let mut free = permits.free.lock().unwrap(); - loop { - if let Some(slot) = free.pop() { - return Self { slot }; - } - free = permits.freed.wait(free).unwrap(); + fn acquire(weight: usize) -> Self { + static NEXT_SLOT: AtomicU64 = AtomicU64::new(0); + 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, + slot: NEXT_SLOT.fetch_add(1, Ordering::Relaxed) % gpu_stream_slots(), } } } impl Drop for GpuPermit { fn drop(&mut self) { - let permits = &*GPU_PERMITS; - permits.free.lock().unwrap().push(self.slot); - permits.freed.notify_one(); + let b = &*GPU_BUDGET; + *b.used.lock().unwrap() -= self.weight; + b.freed.notify_all(); } } @@ -800,8 +823,8 @@ pub fn multiply_batch_on_gpu( /// 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. Holds a /// [`GpuPermit`] for its sequential layout + device section (acquired only after the parallel -/// marshal — see [`GpuPermits`]), so at most `NASSAU_GPU_CONCURRENCY` device sections run at -/// once across all worker threads. +/// marshal — see [`GpuBudget`]), so the total output size of concurrent device sections +/// stays under `NASSAU_GPU_MEM_BUDGET_MB` across all worker threads. fn multiply_batch_block( algebra: &MilnorAlgebra, num_cols: usize, @@ -875,7 +898,7 @@ fn multiply_batch_block( // 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). - let permit = GpuPermit::acquire(); + let permit = GpuPermit::acquire(num_rows * num_limbs * 4); // 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 @@ -961,10 +984,7 @@ fn multiply_batch_block( // than that ever exist — see [`GpuPermit`]. Cubecl's per-device runner serializes the // actual server access; cross-slot/cross-thread reuse of the shared resident handles is // event-synced by cubecl. `memory_cleanup` below trims only this slot's own pool. - let result = StreamId { - value: permit.slot as u64, - } - .executes(|| { + let result = StreamId { value: permit.slot }.executes(|| { let client = CudaRuntime::client(&CudaDevice::default()); // Shared resident admissible buffers (see [`ResidentDev`]): re-upload the master ONLY // when this block dereferences past the uploaded prefix (`need_cs` / `need_mk`). The @@ -978,15 +998,23 @@ fn multiply_batch_block( // host master and every offset `< uploaded` is final. let (cs_h, mk_h, cs_len_master, mk_len_master) = { let mut dev = RESIDENT_DEV.lock().unwrap(); + // Prefix-only upload with doubling: ship `max(need, 2 x uploaded)` entries (clamped + // to the master), not the whole master. At the frontier every new `t` appends new + // `R`s, so uploading the full master on each miss re-ships gigabytes of tail the + // launch never touches; doubling amortizes total upload traffic to <= ~2x the final + // master size while keeping the upload count logarithmic. if dev.cs_handle.is_none() || dev.cs_uploaded < need_cs { let host = RESIDENT_HOST.read().unwrap(); - dev.cs_handle = Some(client.create_from_slice(u16::as_bytes(&host.col_sums))); - dev.cs_uploaded = host.col_sums.len(); + let len = host.col_sums.len().min(need_cs.max(2 * dev.cs_uploaded)); + dev.cs_handle = + Some(client.create_from_slice(u16::as_bytes(&host.col_sums[..len]))); + dev.cs_uploaded = len; } if dev.mk_handle.is_none() || dev.mk_uploaded < need_mk { let host = RESIDENT_HOST.read().unwrap(); - dev.mk_handle = Some(client.create_from_slice(u16::as_bytes(&host.masks))); - dev.mk_uploaded = host.masks.len(); + let len = host.masks.len().min(need_mk.max(2 * dev.mk_uploaded)); + dev.mk_handle = Some(client.create_from_slice(u16::as_bytes(&host.masks[..len]))); + dev.mk_uploaded = len; } ( dev.cs_handle.clone().unwrap(), From a04dd91ccd92c04c705200bfa94278373498b74c Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Thu, 23 Jul 2026 16:09:41 -0400 Subject: [PATCH 008/127] nassau: offload the d_s image build (signature_matrix) to the GPU multiply path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The zero-signature image matrix (d_s applied to the zero-sig source basis, column-masked to the zero-sig target) was the last per-bidegree Milnor multiply still on the CPU — a serial per-row apply_to_basis_element_restricted, ~17% of frontier wall time in the perf profile. But it is the *same* restricted multiply as the QI-source `full_matrix` already built via restricted_partial_matrix_maybe_gpu, just on d_s = differentials[b.s()] instead of d_{s-1}, and its target mask/dimension are exactly the `target_mask`/`target_dim` already computed for the 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, already-verified path and apply the column mask on CPU; drop the serial `signature_matrix` method. (Reinstates the "signature_matrix offload" win from the original nassau_gpu branch, lost in the #272 relaxed-graph merge.) row_reduce stays on CPU — the signature-masked matrices are very flat (~100 x 100000), a poor RREF target for the GPU. Correctness: GPU Ext chart byte-identical to CPU-only through (100,152); NASSAU_GPU_VERIFY passes at stem 130. This shrinks the serial tail that Amdahl-capped the GPU:CPU ratio, so the arithmetic-intensity advantage finally shows through and the gap WIDENS with stem (S_2, s<=152, 16-core H200 box, w=32): band GPU CPU ratio 130->140 159s 206s 1.30x 140->150 278s 423s 1.52x cum 0->150 596s 771s 1.29x (was 723s, a near-tie) Memory stays bounded by the same byte-budget/block machinery (this path reuses multiply_batch_on_gpu). Next lever: the full-reuse-matrix readback. Co-Authored-By: Claude Opus 4.8 --- ext/src/nassau.rs | 70 ++++++++++++++--------------------------------- 1 file changed, 20 insertions(+), 50 deletions(-) diff --git a/ext/src/nassau.rs b/ext/src/nassau.rs index c9f693a695..ab7dc8fe1d 100644 --- a/ext/src/nassau.rs +++ b/ext/src/nassau.rs @@ -181,53 +181,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> + '_ { @@ -822,9 +775,26 @@ impl> Resolution { f.write_fix()?; } - // Compute image - let mut n = - subalgebra.signature_matrix(&self.differentials[b.s()], b.t(), &zero_sig, target_bound); + // 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 = 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 = 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.row_reduce(); let next_row = n.rows(); From bc253f68089f9992eb480cbf627b12488f24ce9e Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Thu, 23 Jul 2026 17:10:55 -0400 Subject: [PATCH 009/127] milnor_gpu: zero the out_h accumulator on-device instead of uploading a host zero buffer MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per-thread stack sampling at stem 145 showed the wavefront's serial stalls were a rayon worker pegged in __memcpy_ssse3 inside create_from_slice — host-side upload marshaling, NOT readback (cubecl 0.10 already does async D2H off pinned memory with the event wait on the worker thread, so the runner is free during the copy). The dominant offender: the batched multiply allocated + zeroed a host `vec![0u32; out_len]` (hundreds of MB at the frontier) and memcpy'd it up as the kernel's XOR accumulator, every launch/block. Allocate out_h uninitialized (client.empty) and zero it with a trivial on-device kernel (zero_u32), same stream as the multiply so it is ordered before it. Removes the host memset, the non-pinned host->device copy, and the transfer itself; on-device zeroing is memory-bound (microseconds on an H200). Verified: GPU Ext chart byte-identical to CPU through (100,152); NASSAU_GPU_VERIFY passes at stem 130. Bands (S_2, s<=152, 16-core H200 box, w=32), vs the prior signature-offload binary: 0->130 159 -> 141s (ties CPU-only 142; was a 0.89x loss) 0->140 318 -> 245s 130->140 marginal 104s vs CPU 206s = 1.98x (was 1.30x) peak RSS 46GB, GPU 28GB (both down). Next serial upload to check: term_pparts / the per-product record arrays. Co-Authored-By: Claude Opus 4.8 --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 29 ++++++++++++++++++-- 1 file changed, 27 insertions(+), 2 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 5c4a7ff1eb..a68a9e51f9 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -266,6 +266,21 @@ fn resident_info(algebra: &MilnorAlgebra, p_part: &[PPartEntry]) -> RInfo { info } +/// 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 Array) { + if ABSOLUTE_POS < out.len() { + out[ABSOLUTE_POS] = 0u32; + } +} + /// 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 @@ -1034,9 +1049,19 @@ fn multiply_batch_block( 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; + // 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. Same stream as the multiply below, so it 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), + ArrayArg::from_raw_parts(out_h.clone(), out_len), + ); + } let pri_h = client.create_from_slice(u32::as_bytes(&prod_r_index)); let pts_h = client.create_from_slice(u32::as_bytes(&prod_term_start)); From c4335568561a7419f6cc7b05e026c18a3fc7a2e8 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Thu, 23 Jul 2026 20:00:37 -0400 Subject: [PATCH 010/127] milnor_gpu: move the resident-master device upload off the handle lock MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The shared admissible master reaches ~3GB by stem 138 (cs 1.2GB + mk 1.9GB), and every launch took the RESIDENT_DEV mutex to read its device handles — with a growth-triggering launch doing a multi-GB create_from_slice re-upload *while holding that mutex*. Per-thread stack sampling showed the frontier collapsing to a single thread memcpy-ing gigabytes while every other bidegree blocked on the lock (upload byte-size instrumentation under NASSAU_GPU_DEBUG confirmed the master, not term data or the seqno table, as the giant upload). Make handle reads lock-free (RESIDENT_DEV: Mutex -> RwLock) and move the upload memcpy outside that lock, serialized only among uploaders by a separate RESIDENT_UPLOAD mutex with a re-check that coalesces a burst of growth-needing launches into one upload. A launch whose R's are already resident proceeds without ever blocking on someone else's upload. Verified GPU chart byte-identical to CPU through (100,152); NASSAU_GPU_VERIFY passes at stem 130. Removes the lock-held-across-copy stall, but occupancy only rose ~3.5->4.5 cores: the dominant limiter is upstream (thin GPU bidegrees + wavefront width), not this lock. Kept because it is correct and matters more at stem 300 where the master is larger. Also adds a per-buffer upload-size line to NASSAU_GPU_DEBUG. Co-Authored-By: Claude Opus 4.8 --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 102 +++++++++++++------ 1 file changed, 69 insertions(+), 33 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index a68a9e51f9..0f817bde5e 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -227,8 +227,9 @@ static RESIDENT_HOST: LazyLock> = /// Process-shared device mirror of the host master: one upload for the whole process, /// re-uploaded only when the master grew. The handles are shared across worker threads and /// stream slots — safe under cubecl 0.10's per-device runner (which serializes all server -/// access), with cross-stream reuse event-synced via the handle's origin-stream stamp. The -/// mutex guards only the grew-check + upload, held briefly per launch. +/// access), with cross-stream reuse event-synced via the handle's origin-stream stamp. Reads +/// go through `RESIDENT_DEV.read()` (lock-free fan-out); uploads run outside that lock, +/// serialized only by `RESIDENT_UPLOAD` (see [`resident_dev_handle`]). #[derive(Default)] struct ResidentDev { cs_handle: Option, @@ -237,8 +238,56 @@ struct ResidentDev { mk_uploaded: usize, } -static RESIDENT_DEV: LazyLock> = - LazyLock::new(|| Mutex::new(ResidentDev::default())); +static RESIDENT_DEV: LazyLock> = + LazyLock::new(|| RwLock::new(ResidentDev::default())); + +/// Serializes master device *uploads* only — never handle reads. A launch that must grow the +/// device master takes this before uploading, so at a growth point at most one multi-GB +/// `create_from_slice` 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: Mutex<()> = Mutex::new(()); + +/// Fetch a resident device-master handle, uploading the current master prefix only when this +/// block dereferences past what is already on the device (`need`). The upload runs OUTSIDE +/// `RESIDENT_DEV` (readers stay lock-free; only concurrent uploaders serialize, on +/// `RESIDENT_UPLOAD`, and a re-check coalesces a burst of growth-needing launches into one +/// upload). The master is append-only, so any handle with `uploaded >= need` is valid. +macro_rules! resident_dev_handle { + ($client:expr, $need:expr, $handle:ident, $uploaded:ident, $host_vec:ident) => {{ + let read_current = || { + let dev = RESIDENT_DEV.read().unwrap(); + match (dev.$uploaded >= $need, dev.$handle.clone()) { + (true, Some(h)) => Some((h, dev.$uploaded)), + _ => None, + } + }; + match read_current() { + Some(hu) => hu, + None => { + let _upload_guard = RESIDENT_UPLOAD.lock().unwrap(); + match read_current() { + Some(hu) => hu, // another uploader already covered our need + None => { + let (handle, len) = { + let host = RESIDENT_HOST.read().unwrap(); + let len = host.$host_vec.len(); + ( + $client.create_from_slice(u16::as_bytes(&host.$host_vec[..len])), + len, + ) + }; + let mut dev = RESIDENT_DEV.write().unwrap(); + dev.$handle = Some(handle.clone()); + dev.$uploaded = len; + (handle, 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 @@ -974,10 +1023,20 @@ fn multiply_batch_block( pps.push(total_pairs as u32); let out_len = num_rows * num_limbs; if std::env::var_os("NASSAU_GPU_DEBUG").is_some() { + let kb = |n: usize, sz: usize| n * sz / 1024; eprintln!( - "[gpu-batch] row_base={row_base} num_rows={num_rows} num_cols={num_cols} \ - num_limbs={num_limbs} products={} total_pairs={total_pairs} out_len={out_len}", + "[gpu-batch] rows={num_rows} cols={num_cols} products={} total_pairs={total_pairs} \ + out_len={out_len} | UPLOAD-KB: g={} xi={} term_pparts={} term_lens={} \ + prod_arrays={} pps={} | resident cs={} mk={}", products.len(), + kb(g.len(), 4), + kb(xi.len(), 4), + kb(term_pparts.len(), 2), + kb(term_lens.len(), 4), + kb(products.len() * 5, 4), + kb(pps.len(), 4), + kb(need_cs, 2), + kb(need_mk, 2), ); } if total_pairs == 0 { @@ -1011,33 +1070,10 @@ fn multiply_batch_block( // then HOST.read, and nothing under either lock blocks on rayon or a permit. The // master is append-only, so the uploaded prefix is always a prefix of the current // host master and every offset `< uploaded` is final. - let (cs_h, mk_h, cs_len_master, mk_len_master) = { - let mut dev = RESIDENT_DEV.lock().unwrap(); - // Prefix-only upload with doubling: ship `max(need, 2 x uploaded)` entries (clamped - // to the master), not the whole master. At the frontier every new `t` appends new - // `R`s, so uploading the full master on each miss re-ships gigabytes of tail the - // launch never touches; doubling amortizes total upload traffic to <= ~2x the final - // master size while keeping the upload count logarithmic. - if dev.cs_handle.is_none() || dev.cs_uploaded < need_cs { - let host = RESIDENT_HOST.read().unwrap(); - let len = host.col_sums.len().min(need_cs.max(2 * dev.cs_uploaded)); - dev.cs_handle = - Some(client.create_from_slice(u16::as_bytes(&host.col_sums[..len]))); - dev.cs_uploaded = len; - } - if dev.mk_handle.is_none() || dev.mk_uploaded < need_mk { - let host = RESIDENT_HOST.read().unwrap(); - let len = host.masks.len().min(need_mk.max(2 * dev.mk_uploaded)); - dev.mk_handle = Some(client.create_from_slice(u16::as_bytes(&host.masks[..len]))); - dev.mk_uploaded = len; - } - ( - dev.cs_handle.clone().unwrap(), - dev.mk_handle.clone().unwrap(), - dev.cs_uploaded, - dev.mk_uploaded, - ) - }; + let (cs_h, cs_len_master) = + resident_dev_handle!(client, need_cs, cs_handle, cs_uploaded, col_sums); + let (mk_h, mk_len_master) = + resident_dev_handle!(client, need_mk, mk_handle, mk_uploaded, masks); // 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. From 164b1291d5059ac11d95eb4abffe5fdaa408124a Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Thu, 23 Jul 2026 21:48:20 -0400 Subject: [PATCH 011/127] milnor_gpu: fill the term-data upload buffers directly, killing the per-product alloc storm MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A frontier launch has ~1e5-1e6 products, and the marshal built term data as `Vec<(Vec, Vec)>` — two heap allocations per product (~1e6 tiny allocs per launch) — then extend-copied them into the flat upload buffers. Per-thread profiling of the GPU path showed this as a dominant chunk of the per-bidegree CPU "envelope" (~16% _int_malloc/_int_free plus the marshal copy) that wraps each (fast) kernel and, because the wavefront is only ~10-15 bidegrees wide, cannot be hidden — so the GPU sits idle between brief spikes. Precompute the term-count prefix sum (`term_off`), size the flat `term_pparts`/ `term_lens` once, and parallel-fill each product's disjoint slice in place (unsafe but sound: prefix-sum ranges never alias). The later layout loop just reads `term_off[pi]` for `prod_term_start` — no per-product allocation, no concat copy. Verified GPU chart byte-identical to CPU through (100,152). Same GPU results, far less allocation and marshal work per launch. (The remaining per-product `GpuProduct.term_indices: Vec` built in extract is the next alloc to flatten.) Co-Authored-By: Claude Opus 4.8 --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 57 +++++++++++++------- 1 file changed, 38 insertions(+), 19 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 0f817bde5e..65ae9d5b1c 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -934,24 +934,46 @@ fn multiply_batch_block( // 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. - // 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| { + // 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(); + let mut term_pparts: Vec = vec![0u16; total_terms * width]; + let mut term_lens: Vec = vec![0u32; total_terms]; + { + let tp_base = term_pparts.as_mut_ptr() as usize; + let tl_base = term_lens.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: products write disjoint `[off, off + nt)` ranges (from the prefix sum), + // each within the allocated buffers, so no two tasks alias any element. The `usize` + // bases are re-formed into pointers here because raw pointers are not `Send`. + let tp = unsafe { + std::slice::from_raw_parts_mut((tp_base as *mut u16).add(off * width), nt * width) + }; + let tl = 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); + tl[k] = elt.p_part.len() as u32; for (slot, &v) in tp[k * width..(k + 1) * width].iter_mut().zip(&elt.p_part) { *slot = narrow_u16(v); } } - (tp, tl) - }) - .collect(); + }); + } // 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 @@ -988,9 +1010,9 @@ fn multiply_batch_block( need_mk = need_mk.max(info.mk_off as usize + info.num_mats as usize * info.mk_len 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(); + // 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()); @@ -1002,12 +1024,9 @@ fn multiply_batch_block( // `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; - for (pi, (tp, tl)) in per_prod.iter().enumerate() { - let prod = &products[pi]; + for (pi, prod) in products.iter().enumerate() { let ri = prod_r_index[pi]; - prod_term_start.push(term_lens.len() as u32); - term_lens.extend_from_slice(tl); - term_pparts.extend_from_slice(tp); + prod_term_start.push(term_off[pi] as u32); pps.push(pair_acc as u32); pair_acc += r_num_matrices[ri as usize] * prod.term_indices.len(); prod_num_terms.push(prod.term_indices.len() as u32); From b02dcb35fd5dde9c23e76c045218359a569c8d39 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Thu, 23 Jul 2026 22:14:03 -0400 Subject: [PATCH 012/127] milnor_gpu: raise GPU_PAIR_CHUNK to ~2^32 so giant multiplies stop over-splitting MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The row-block splitter caps a launch at GPU_PAIR_CHUNK thread-pairs (the kernel indexes threads by u32 ABSOLUTE_POS, ceiling 2^32). It was set to 1<<30 (~1.07e9), ~4x below the real ceiling — so every billion-pair giant was chopped into ~4 launches, each a separate upload + kernel + BLOCKING readback round-trip, even though its output is only ~350 MB (well under gpu_block_bytes). Debug confirmed the giants pegged at 1.07e9 pairs; this, not the byte budget, was the binding split, which is why a NASSAU_GPU_BLOCK_MB sweep was flat. Raise it to 3.9e9 (leaves ~0.39e9 headroom under 2^32; the splitter always takes >=1 row and a lone row past 2^32 still trips the per-block u32::try_from assert; grid stays ~1.5e7 cubes, far under 2^31). Giants now run as a single launch (max total_pairs observed 3.90e9), collapsing 4 round-trips to 1. GPU 0->140 (w=100): ~245-288s -> 216s. Chart byte-identical to CPU through (100,152). Co-Authored-By: Claude Opus 4.8 --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 23 +++++++++++++------- 1 file changed, 15 insertions(+), 8 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 65ae9d5b1c..386d833c05 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -35,14 +35,21 @@ use crate::algebra::{Algebra, MilnorAlgebra, combinatorics::xi_degrees}; /// `mk_len = rows + cols − 1 ≤ MAX_XI_TAU + ⌈log2⌉`; 32 covers every in-range case. const WORKING_CAP: usize = 32; -/// 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). `1 << 30` (~1.07e9) leaves >3x headroom -/// under `2^32` even when a lone over-budget row overshoots it, and keeps the grid -/// (`total_pairs / 256` cubes) well under CUDA's `2^31 - 1` grid-dimension limit. -const GPU_PAIR_CHUNK: usize = 1 << 30; +/// 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; /// Per-launch output-buffer budget in bytes (`NASSAU_GPU_BLOCK_MB`, default 512 MiB). /// From e6b85fd9190e3e617b5aed00c2b7959c234f2155 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Fri, 24 Jul 2026 01:05:29 -0400 Subject: [PATCH 013/127] milnor_gpu: make the Milnor basis resident on the GPU, upload term indices MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The batched multiply re-gathered and re-uploaded every term's zero-padded p-part (`term_pparts`, width*2 bytes/term) on every launch — the dominant per-launch H2D transfer, plus a large parallel host gather. Make the basis itself resident on the device instead: build it once (grown incrementally as higher degrees appear, mirroring the admissible master) and upload only `term_gei[slot]`, the term's global basis-element index `global_base[s_degree] + ti` (4 bytes/term). The kernel reads the p-part from `basis_pparts[gei*width..]` with length `basis_lens[gei]`. At stem 140 the per-launch term transfer drops from ~width/2x larger to term_gei=95 MB, the basis is a one-time few-MB upload, and the per-term p-part gather is gone. Kernel change is minimal: params `term_pparts, term_lens` -> `basis_pparts, basis_lens, term_gei` (net +1 array arg), launch-arg order preserved 1:1 with the signature. `multiply_pair` is unchanged. An A/B toggle (`NASSAU_GPU_BASIS_PASSTHROUGH=1`) binds the per-launch term buffers as the "basis" with an identity index map, reproducing the old behaviour through the new kernel — so a single binary can isolate a kernel-signature bug from a resident host/upload bug. Both paths verified GPU==CPU per launch at stem 80 (MIN_WORK=0), and the resident path chart-matches CPU at (100,152). Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 253 +++++++++++++++++-- 1 file changed, 226 insertions(+), 27 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 386d833c05..4424029774 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -130,6 +130,16 @@ fn gpu_stream_slots() -> u64 { *SLOTS } +/// 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` bytes from [`GPU_BUDGET`] plus a round-robin stream slot; /// blocks (parked, not spinning) until the budget admits it. struct GpuPermit { @@ -322,6 +332,131 @@ fn resident_info(algebra: &MilnorAlgebra, p_part: &[PPartEntry]) -> RInfo { 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`]; `uploaded` is the element count on the device. +#[derive(Default)] +struct ResidentBasisDev { + pp_handle: Option, + ln_handle: Option, + uploaded: usize, +} + +static RESIDENT_BASIS_DEV: LazyLock> = + LazyLock::new(|| RwLock::new(ResidentBasisDev::default())); + +/// Serializes basis device *uploads* only (never handle reads); see [`RESIDENT_UPLOAD`]. +static RESIDENT_BASIS_UPLOAD: Mutex<()> = Mutex::new(()); + +/// 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) { + *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() +} + +/// Fetch the resident basis device handles `(pparts, lens)`, uploading the current host basis +/// only when the device copy does not yet cover `$need` elements. Upload runs outside +/// `RESIDENT_BASIS_DEV` (readers stay lock-free), serialized by `RESIDENT_BASIS_UPLOAD` with a +/// re-check to coalesce a burst of growth. The basis is append-only, so any handle with +/// `uploaded >= $need` is valid. A macro (not a fn) so the cubecl client type stays inferred, +/// exactly as [`resident_dev_handle`]. +macro_rules! basis_dev_handles { + ($client:expr, $need:expr) => {{ + let read_current = || { + let dev = RESIDENT_BASIS_DEV.read().unwrap(); + match (dev.uploaded >= $need, dev.pp_handle.clone(), dev.ln_handle.clone()) { + (true, Some(pp), Some(ln)) => Some((pp, ln)), + _ => None, + } + }; + match read_current() { + Some(h) => h, + None => { + let _guard = RESIDENT_BASIS_UPLOAD.lock().unwrap(); + match read_current() { + Some(h) => h, // another uploader already covered our need + None => { + let (pp, ln, elems) = { + let host = RESIDENT_BASIS_HOST.read().unwrap(); + ( + $client.create_from_slice(u16::as_bytes(&host.pparts)), + $client.create_from_slice(u32::as_bytes(&host.lens)), + host.lens.len(), + ) + }; + let mut dev = RESIDENT_BASIS_DEV.write().unwrap(); + dev.pp_handle = Some(pp.clone()); + dev.ln_handle = Some(ln.clone()); + dev.uploaded = elems; + (pp, ln) + } + } + } + } + }}; +} + /// 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 @@ -645,8 +780,9 @@ fn multiply_single_r_kernel( fn multiply_batch_kernel( col_sums: &Array, masks: &Array, - term_pparts: &Array, - term_lens: &Array, + basis_pparts: &Array, + basis_lens: &Array, + term_gei: &Array, g: &Array, xi: &Array, out: &mut Array>, @@ -694,17 +830,23 @@ fn multiply_batch_kernel( 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; + // `term_gei[term_slot]` is the term's *global* basis-element index (across all degrees): + // its (width-padded) p-part lives at `basis_pparts[gei*width ..]`, length `basis_lens[gei]`. + // The basis is resident on the device (uploaded once, grown incrementally), so a launch + // uploads only these indices instead of re-gathering every term's p-part — see + // [`ResidentBasisHost`]. + let gei = usize::cast_from(term_gei[term_slot]); multiply_pair( col_sums, masks, - term_pparts, + basis_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]), + gei * width, + usize::cast_from(basis_lens[gei]), cs_len, mk_len, usize::cast_from(prod_row_base[p]), @@ -957,30 +1099,68 @@ fn multiply_batch_block( off }; let total_terms = *term_off.last().unwrap(); - let mut term_pparts: Vec = vec![0u16; total_terms * width]; - let mut term_lens: Vec = vec![0u32; total_terms]; + + // 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); + let global_base = ensure_basis(algebra, width, max_s_degree); + let mut term_gei: Vec = vec![0u32; total_terms]; { - let tp_base = term_pparts.as_mut_ptr() as usize; - let tl_base = term_lens.as_mut_ptr() as usize; + let tg_base = term_gei.as_mut_ptr() as usize; + let global_base = &global_base; (0..products.len()).into_maybe_par_iter().for_each(|pi| { let prod = &products[pi]; let (off, nt) = (term_off[pi], prod.term_indices.len()); // SAFETY: products write disjoint `[off, off + nt)` ranges (from the prefix sum), - // each within the allocated buffers, so no two tasks alias any element. The `usize` - // bases are re-formed into pointers here because raw pointers are not `Send`. - let tp = unsafe { + // each within the allocated buffer, so no two tasks alias any element. The `usize` + // base is re-formed into a pointer here because raw pointers are not `Send`. + let tg = unsafe { std::slice::from_raw_parts_mut((tg_base as *mut u32).add(off), nt) }; + let base = global_base[prod.s_degree as usize]; + for (k, &ti) in prod.term_indices.iter().enumerate() { + tg[k] = base + ti as u32; + } + }); + } + // 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; + + // 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 (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 tl = unsafe { std::slice::from_raw_parts_mut((tl_base as *mut u32).add(off), nt) }; + 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[k] = 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) { *slot = narrow_u16(v); } } }); - } + // 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()) + }; // 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 @@ -1052,13 +1232,12 @@ fn multiply_batch_block( 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_pparts={} term_lens={} \ - prod_arrays={} pps={} | resident cs={} mk={}", + 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_pparts.len(), 2), - kb(term_lens.len(), 4), + kb(term_gei.len(), 4), kb(products.len() * 5, 4), kb(pps.len(), 4), kb(need_cs, 2), @@ -1069,10 +1248,15 @@ fn multiply_batch_block( return vec![vec![0u32; num_limbs]; num_rows]; } - // The resident `col_sums`/`masks` are non-empty once any `R` is present (guaranteed - // here, since `total_pairs > 0`); only `term_pparts` needs the non-empty guard. - if term_pparts.is_empty() { + // 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 marshal_ms = t_marshal.elapsed().as_secs_f64() * 1e3; @@ -1103,8 +1287,22 @@ fn multiply_batch_block( // 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. - let tp_h = client.create_from_slice(u16::as_bytes(&term_pparts)); - let tl_h = client.create_from_slice(u32::as_bytes(&term_lens)); + // Resident basis handles (default) or per-launch passthrough buffers (A/B diagnostic). + // `bp_len`/`bl_len` are the logical array lengths the kernel sees; every `gei` a thread + // dereferences is `< need_basis_elems`, so `need_basis_elems*width` / `need_basis_elems` + // cover it (the resident buffer may be larger — append-only — which is fine). + let (bp_h, bl_h, bp_len, bl_len) = if passthrough { + ( + client.create_from_slice(u16::as_bytes(&term_pparts)), + client.create_from_slice(u32::as_bytes(&term_lens)), + term_pparts.len(), + term_lens.len(), + ) + } else { + let (pp, ln) = basis_dev_handles!(client, need_basis_elems); + (pp, ln, need_basis_elems * width, need_basis_elems) + }; + let tg_h = client.create_from_slice(u32::as_bytes(&term_gei)); 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)); @@ -1139,8 +1337,9 @@ fn multiply_batch_block( CubeDim::new_1d(THREADS), ArrayArg::from_raw_parts(cs_h, cs_len_master), ArrayArg::from_raw_parts(mk_h, mk_len_master), - ArrayArg::from_raw_parts(tp_h, term_pparts.len()), - ArrayArg::from_raw_parts(tl_h, term_lens.len()), + ArrayArg::from_raw_parts(bp_h, bp_len), + ArrayArg::from_raw_parts(bl_h, bl_len), + ArrayArg::from_raw_parts(tg_h, term_gei.len()), ArrayArg::from_raw_parts(g_h, g.len()), ArrayArg::from_raw_parts(xi_h, xi.len()), ArrayArg::from_raw_parts(out_h.clone(), out_len), From 4efee670d5bd225e259d511ce185c9a1d0d4aacf Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Sat, 25 Jul 2026 00:05:15 -0400 Subject: [PATCH 014/127] milnor_gpu: safe multi-stream via in-place-grown resident globals MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `NASSAU_GPU=1` could not use multiple CUDA streams: the shared resident admissible master + Milnor basis were re-`create_from_slice`d into a NEW device handle on every growth, and that handle churn broke cubecl's per-handle cross-stream synchronization, so a launch on one stream could read the master while another was mid-upload -> wrong multiply (`dx != 0`), "Memory page" panic, or hang. Fix: make the resident buffers STABLE and grow them IN PLACE — the read-only shared global (model-weights) pattern cubecl supports across streams. A small `copy_into_*` kernel writes the new tail (uploaded to scratch via `create_from_slice`) at the buffer's append offset; the handle changes only on a rare capacity doubling, which is barrier-protected (`RESIDENT_REALLOC`: device sections hold the read lock across the multiply, a realloc takes the write lock and quiesces them). Each worker gets a stable per-thread stream id (`thread_stream_id`); default `NASSAU_GPU_STREAMS = 8`. Key gotcha (a stem-150 `dx != 0`): cubecl's `ArrayArg` length is u32, so a buffer of exactly 2^32 elements truncates to length 0 and the copy writes nothing (buffer reads all zeros). The doubling `cap` jumped 2^31 -> 2^32 right at stem 150's `masks` size. Capacity is clamped to `RESIDENT_MAX_CAP = 2^32 - 1`; a single resident buffer cannot exceed that (a larger master needs splitting — the old create_from_slice path had the same limit). Copies are chunked under the u32 ABSOLUTE_POS thread limit. Validated: VERIFY (GPU==CPU per launch) at 8 streams, chart-match to CPU, 0 dx-crashes over many stem-150 reps at 1 and 8 streams. Perf note: multi-stream is correct but ~neutral vs single-stream at stem 140/150 (the per-device runner serializes kernel submission); it may help at higher stems with a wider heavy-bidegree wavefront. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 381 +++++++++++++++---- 1 file changed, 311 insertions(+), 70 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 4424029774..767541249e 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -111,14 +111,15 @@ static GPU_BUDGET: LazyLock = LazyLock::new(|| GpuBudget { freed: Condvar::new(), }); -/// Fixed CUDA stream-slot count (`NASSAU_GPU_STREAMS`, default 8): device sections run under -/// `StreamId { value: counter % slots }`, so only this many streams (and hence retained -/// memory pools) ever exist. Without a pin every worker thread gets its own default stream, -/// and each stream's pool *retains* its freed slabs — ~100 worker streams each retaining -/// ~1 GB filled the whole 143 GB H200. Slots are round-robin and *shared*, not exclusive: -/// two small launches on one slot merely serialize on that CUDA stream (memory-safe, and -/// cheap for small kernels), so slot count does not cap concurrency the way the byte budget -/// does. Kept at 8: a 16-slot experiment collapsed >30x (unexplained cross-stream churn). +/// Number of distinct CUDA streams to spread device work over (`NASSAU_GPU_STREAMS`, default 8). +/// Each worker thread gets a stable per-thread stream id (see [`thread_stream_id`]), and this caps +/// how many distinct streams — hence retained memory pools — exist; `1` forces the single-stream +/// mode (all threads → stream 0). +/// +/// Multi-stream is SAFE because the shared resident master/basis are grown IN PLACE in stable +/// device buffers ([`ResidentDev`]) — the read-only "global weights" pattern cubecl supports across +/// streams. (The earlier churny design — re-`create_from_slice` a new handle on every growth — broke +/// cubecl's per-handle cross-stream sync and had to run single-stream; that is fixed.) fn gpu_stream_slots() -> u64 { static SLOTS: LazyLock = LazyLock::new(|| { std::env::var("NASSAU_GPU_STREAMS") @@ -130,6 +131,25 @@ fn gpu_stream_slots() -> u64 { *SLOTS } +/// A CUDA stream id UNIQUE to the calling worker thread (never shared or reused across threads — +/// cubecl's per-stream state assumes one driver thread per stream). Returns 0 for all threads when +/// `NASSAU_GPU_STREAMS == 1` (the single-stream fallback). With the resident master/basis now stable +/// device buffers grown IN PLACE (no handle churn — see [`ResidentDev`]), cross-stream reads of those +/// shared globals are the ordinary case cubecl supports, so distinct per-thread streams are safe and +/// give real device concurrency for the wide record-run wavefront. +fn thread_stream_id() -> u64 { + if gpu_stream_slots() == 1 { + return 0; + } + thread_local! { + static ID: u64 = { + static NEXT: AtomicU64 = AtomicU64::new(0); + NEXT.fetch_add(1, Ordering::Relaxed) + 1 // +1: id 0 is reserved for the single-stream mode + }; + } + ID.with(|&id| id) +} + /// 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. @@ -140,26 +160,22 @@ fn basis_passthrough() -> bool { *ON } -/// RAII reservation of `weight` bytes from [`GPU_BUDGET`] plus a round-robin stream slot; -/// blocks (parked, not spinning) until the budget admits it. +/// RAII reservation of `weight` output bytes from [`GPU_BUDGET`]; blocks (parked, not spinning) +/// until the budget admits it. The stream is chosen per worker thread (see [`thread_stream_id`]), +/// so the permit no longer carries a slot. struct GpuPermit { weight: usize, - slot: u64, } impl GpuPermit { fn acquire(weight: usize) -> Self { - static NEXT_SLOT: AtomicU64 = AtomicU64::new(0); 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, - slot: NEXT_SLOT.fetch_add(1, Ordering::Relaxed) % gpu_stream_slots(), - } + Self { weight } } } @@ -241,23 +257,46 @@ struct ResidentHost { static RESIDENT_HOST: LazyLock> = LazyLock::new(|| RwLock::new(ResidentHost::default())); -/// Process-shared device mirror of the host master: one upload for the whole process, -/// re-uploaded only when the master grew. The handles are shared across worker threads and -/// stream slots — safe under cubecl 0.10's per-device runner (which serializes all server -/// access), with cross-stream reuse event-synced via the handle's origin-stream stamp. Reads -/// go through `RESIDENT_DEV.read()` (lock-free fan-out); uploads run outside that lock, -/// serialized only by `RESIDENT_UPLOAD` (see [`resident_dev_handle`]). +/// A resident device buffer that GROWS IN PLACE: a single stable handle, allocated once +/// (persistent, so per-launch cleanup never reindexes it) with headroom and written to at its +/// append offset as the host master grows — never re-`create_from_slice`d. The handle changes only +/// on a rare capacity doubling. `uploaded` is how many elements have been written; `cap` is the +/// allocated element capacity. +#[derive(Default)] +struct GrowBuf { + handle: Option, + cap: usize, + uploaded: usize, +} + +/// Process-shared device mirror of the host master. Each buffer is a STABLE handle grown in place +/// (see [`GrowBuf`], [`resident_dev_handle`]). 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 fixed handle 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_handle: Option, - mk_handle: Option, - cs_uploaded: usize, - mk_uploaded: usize, + cs: GrowBuf, + mk: GrowBuf, } static RESIDENT_DEV: LazyLock> = LazyLock::new(|| RwLock::new(ResidentDev::default())); +/// Initial element capacity of a resident buffer. Doubling to grow past it is barrier-protected +/// (see [`RESIDENT_REALLOC`]) so it is correct, but each doubling briefly quiesces the device, so +/// the initial size carries headroom to keep them few (~a handful per run). u16 → ~512 MiB. +const RESIDENT_INIT_CAP: usize = 1 << 28; + +/// Maximum resident buffer capacity (elements). cubecl's `ArrayArg` length is a `u32`, so an array +/// of exactly `2^32` elements has its length truncated to 0 — the copy then writes nothing and the +/// buffer reads as all zeros (the stem-150 `dx != 0`, hit when the doubling `cap` jumped 2^31 → 2^32). +/// Capacity is clamped just below the limit. A single resident buffer therefore cannot exceed +/// `2^32 - 1` elements; a master/basis larger than that (very high stems) must be split across +/// buffers — the same fundamental cubecl limit the old full-`create_from_slice` path also had. +const RESIDENT_MAX_CAP: usize = (1 << 32) - 1; + /// Serializes master device *uploads* only — never handle reads. A launch that must grow the /// device master takes this before uploading, so at a growth point at most one multi-GB /// `create_from_slice` runs (others re-check and find it already done) instead of every launch @@ -266,17 +305,27 @@ static RESIDENT_DEV: LazyLock> = /// collapsed the whole wavefront to one memcpy-ing thread). static RESIDENT_UPLOAD: Mutex<()> = Mutex::new(()); +/// Guards the rare capacity-doubling REALLOC (the only time a resident handle changes) against +/// concurrent readers. A device section holds the read lock while its multiply kernel is reading the +/// resident buffers; a realloc takes the write lock, so it waits for every in-flight reader to drain +/// and blocks new ones — the swap + old-buffer free then happens with the device quiesced, which the +/// churny path could not guarantee (residual ~1/10 crash at the readback / a stale multiply). The +/// common in-place delta write does NOT take this lock (it's stable-handle, proven race-free), so the +/// hot path stays lock-free; reallocs happen only a handful of times per run, so the barrier is free. +static RESIDENT_REALLOC: RwLock<()> = RwLock::new(()); + /// Fetch a resident device-master handle, uploading the current master prefix only when this /// block dereferences past what is already on the device (`need`). The upload runs OUTSIDE /// `RESIDENT_DEV` (readers stay lock-free; only concurrent uploaders serialize, on /// `RESIDENT_UPLOAD`, and a re-check coalesces a burst of growth-needing launches into one /// upload). The master is append-only, so any handle with `uploaded >= need` is valid. macro_rules! resident_dev_handle { - ($client:expr, $need:expr, $handle:ident, $uploaded:ident, $host_vec:ident) => {{ + ($client:expr, $need:expr, $buf:ident, $host_vec:ident) => {{ + // Lock-free fast path: the stable handle already covers `$need`. let read_current = || { let dev = RESIDENT_DEV.read().unwrap(); - match (dev.$uploaded >= $need, dev.$handle.clone()) { - (true, Some(h)) => Some((h, dev.$uploaded)), + match (dev.$buf.uploaded >= $need, dev.$buf.handle.clone()) { + (true, Some(h)) => Some((h, dev.$buf.uploaded)), _ => None, } }; @@ -285,20 +334,68 @@ macro_rules! resident_dev_handle { None => { let _upload_guard = RESIDENT_UPLOAD.lock().unwrap(); match read_current() { - Some(hu) => hu, // another uploader already covered our need + Some(hu) => hu, // another grower already covered our need None => { - let (handle, len) = { - let host = RESIDENT_HOST.read().unwrap(); - let len = host.$host_vec.len(); - ( - $client.create_from_slice(u16::as_bytes(&host.$host_vec[..len])), - len, - ) + let (old_handle, cap, uploaded) = { + let dev = RESIDENT_DEV.read().unwrap(); + (dev.$buf.handle.clone(), dev.$buf.cap, dev.$buf.uploaded) }; - let mut dev = RESIDENT_DEV.write().unwrap(); - dev.$handle = Some(handle.clone()); - dev.$uploaded = len; - (handle, len) + let host_len = RESIDENT_HOST.read().unwrap().$host_vec.len(); + + // (Re)allocate a stable persistent buffer on first use or capacity overflow + // (rare: `RESIDENT_INIT_CAP` has headroom and the master saturates early), + // copying existing device data on-device. Only THIS handle-swap needs a sync, + // so a fresh-handle cross-stream read never sees an incomplete copy. + let (handle, cap) = if old_handle.is_none() || host_len > cap { + // Quiesce readers for the handle swap (see [`RESIDENT_REALLOC`]). + let _realloc_w = RESIDENT_REALLOC.write().unwrap(); + let new_cap = host_len.max(cap * 2).max(RESIDENT_INIT_CAP) + .min(RESIDENT_MAX_CAP); + let new_handle = $client.empty(new_cap * ::core::mem::size_of::()); + if let Some(oh) = &old_handle { + if uploaded > 0 { + copy_chunked!( + $client, copy_into_u16, oh, uploaded, 0usize, + new_handle, new_cap, 0usize, uploaded + ); + } + } + let _ = cubecl_common::reader::read_sync($client.sync()); + (new_handle, new_cap) + } else { + (old_handle.unwrap(), cap) + }; + + // Write the new tail host[uploaded..host_len] into the STABLE buffer at its + // append offset via a scratch upload + copy kernel, then SYNC before + // publishing: cubecl does not cross-stream-order a *kernel write* to a shared + // buffer against reads the way it does `create_from_slice`, so a reader on + // another worker's stream could otherwise see the tail mid-copy (harmless at + // small sizes, corrupting once the copy is GB-scale — the stem-150 `dx != 0`). + // Blocking here makes the new prefix physically resident before any reader can + // observe the bumped `uploaded`. + if host_len > uploaded { + let n = host_len - uploaded; + let scratch = { + let host = RESIDENT_HOST.read().unwrap(); + $client.create_from_slice(u16::as_bytes( + &host.$host_vec[uploaded..host_len], + )) + }; + copy_chunked!( + $client, copy_into_u16, scratch, n, 0usize, + handle, cap, uploaded, n + ); + let _ = cubecl_common::reader::read_sync($client.sync()); + } + + { + let mut dev = RESIDENT_DEV.write().unwrap(); + dev.$buf.handle = Some(handle.clone()); + dev.$buf.cap = cap; + dev.$buf.uploaded = host_len; + } + (handle, host_len) } } } @@ -357,12 +454,13 @@ struct ResidentBasisHost { static RESIDENT_BASIS_HOST: LazyLock> = LazyLock::new(|| RwLock::new(ResidentBasisHost::default())); -/// Device mirror of [`ResidentBasisHost`]; `uploaded` is the element count on the device. +/// Device mirror of [`ResidentBasisHost`], both buffers grown IN PLACE (see [`GrowBuf`], +/// [`resident_dev_handle`]). `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_handle: Option, - ln_handle: Option, - uploaded: usize, + pp: GrowBuf, + ln: GrowBuf, } static RESIDENT_BASIS_DEV: LazyLock> = @@ -423,9 +521,11 @@ fn ensure_basis(algebra: &MilnorAlgebra, width: usize, max_degree: i32) -> Vec {{ + // Lock-free fast path: the stable buffers already cover `$need` basis elements + // (`ln.uploaded` is the element count). let read_current = || { let dev = RESIDENT_BASIS_DEV.read().unwrap(); - match (dev.uploaded >= $need, dev.pp_handle.clone(), dev.ln_handle.clone()) { + match (dev.ln.uploaded >= $need, dev.pp.handle.clone(), dev.ln.handle.clone()) { (true, Some(pp), Some(ln)) => Some((pp, ln)), _ => None, } @@ -435,21 +535,98 @@ macro_rules! basis_dev_handles { None => { let _guard = RESIDENT_BASIS_UPLOAD.lock().unwrap(); match read_current() { - Some(h) => h, // another uploader already covered our need + Some(h) => h, // another grower already covered our need None => { - let (pp, ln, elems) = { - let host = RESIDENT_BASIS_HOST.read().unwrap(); - ( - $client.create_from_slice(u16::as_bytes(&host.pparts)), - $client.create_from_slice(u32::as_bytes(&host.lens)), - host.lens.len(), - ) + let width = RESIDENT_BASIS_HOST.read().unwrap().width; + let elems = RESIDENT_BASIS_HOST.read().unwrap().lens.len(); + let pp_len = elems * width; // u16 count + let ln_len = elems; // u32 count + + // Grow the width-padded p-parts buffer (u16) in place. + let (old_pp, pp_cap, pp_up) = { + let d = RESIDENT_BASIS_DEV.read().unwrap(); + (d.pp.handle.clone(), d.pp.cap, d.pp.uploaded) + }; + let (pp_h, pp_cap) = if old_pp.is_none() || pp_len > pp_cap { + let _realloc_w = RESIDENT_REALLOC.write().unwrap(); + let new_cap = pp_len.max(pp_cap * 2).max(RESIDENT_INIT_CAP) + .min(RESIDENT_MAX_CAP); + let nh = $client.empty(new_cap * ::core::mem::size_of::()); + if let Some(oh) = &old_pp { + if pp_up > 0 { + copy_chunked!( + $client, copy_into_u16, oh, pp_up, 0usize, + nh, new_cap, 0usize, pp_up + ); + } + } + let _ = cubecl_common::reader::read_sync($client.sync()); + (nh, new_cap) + } else { + (old_pp.unwrap(), pp_cap) }; - let mut dev = RESIDENT_BASIS_DEV.write().unwrap(); - dev.pp_handle = Some(pp.clone()); - dev.ln_handle = Some(ln.clone()); - dev.uploaded = elems; - (pp, ln) + if pp_len > pp_up { + let n = pp_len - pp_up; + let scratch = { + let host = RESIDENT_BASIS_HOST.read().unwrap(); + $client.create_from_slice(u16::as_bytes(&host.pparts[pp_up..pp_len])) + }; + copy_chunked!( + $client, copy_into_u16, scratch, n, 0usize, + pp_h, pp_cap, pp_up, n + ); + } + + // Grow the lengths buffer (u32) in place. + let (old_ln, ln_cap, ln_up) = { + let d = RESIDENT_BASIS_DEV.read().unwrap(); + (d.ln.handle.clone(), d.ln.cap, d.ln.uploaded) + }; + let (ln_h, ln_cap) = if old_ln.is_none() || ln_len > ln_cap { + let _realloc_w = RESIDENT_REALLOC.write().unwrap(); + let new_cap = ln_len.max(ln_cap * 2).max(RESIDENT_INIT_CAP) + .min(RESIDENT_MAX_CAP); + let nh = $client.empty(new_cap * ::core::mem::size_of::()); + if let Some(oh) = &old_ln { + if ln_up > 0 { + copy_chunked!( + $client, copy_into_u32, oh, ln_up, 0usize, + nh, new_cap, 0usize, ln_up + ); + } + } + let _ = cubecl_common::reader::read_sync($client.sync()); + (nh, new_cap) + } else { + (old_ln.unwrap(), ln_cap) + }; + if ln_len > ln_up { + let n = ln_len - ln_up; + let scratch = { + let host = RESIDENT_BASIS_HOST.read().unwrap(); + $client.create_from_slice(u32::as_bytes(&host.lens[ln_up..ln_len])) + }; + copy_chunked!( + $client, copy_into_u32, scratch, n, 0usize, + ln_h, ln_cap, ln_up, n + ); + } + + // Sync both in-place copies before publishing (see [`resident_dev_handle`]): + // a kernel write to the shared basis buffers must be physically done before a + // reader on another worker's stream can observe the bumped element count. + let _ = cubecl_common::reader::read_sync($client.sync()); + + { + let mut dev = RESIDENT_BASIS_DEV.write().unwrap(); + dev.pp.handle = Some(pp_h.clone()); + dev.pp.cap = pp_cap; + dev.pp.uploaded = pp_len; + dev.ln.handle = Some(ln_h.clone()); + dev.ln.cap = ln_cap; + dev.ln.uploaded = ln_len; + } + (pp_h, ln_h) } } } @@ -472,6 +649,60 @@ fn zero_u32(out: &mut Array) { } } +/// 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`, and `count` bounds this launch so the caller can split a copy larger than the +/// `u32` `ABSOLUTE_POS` thread limit into chunks (a resident buffer exceeds 2^32 u16 elements around +/// stem 150 — the width-padded basis p-parts especially). See [`copy_into_chunked`]. +#[cube(launch)] +fn copy_into_u16(src: &Array, dst: &mut Array, 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)] +fn copy_into_u32(src: &Array, dst: &mut Array, 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; + +/// 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::( + &$client, + CubeCount::Static((n as u32).div_ceil(CT), 1, 1), + CubeDim::new_1d(CT), + ArrayArg::from_raw_parts($src.clone(), $src_len), + ArrayArg::from_raw_parts($dst.clone(), $dst_len), + $src_off + done, + $dst_off + done, + n as u32, + ); + } + done += n; + } + }}; +} + /// 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 @@ -1171,7 +1402,9 @@ fn multiply_batch_block( // 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). - let permit = GpuPermit::acquire(num_rows * num_limbs * 4); + // Held for the device section (RAII): bounds total in-flight output bytes across workers. The + // stream is now chosen per-thread (see [`thread_stream_id`]), not from the permit. + let _permit = GpuPermit::acquire(num_rows * num_limbs * 4); // 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 @@ -1263,12 +1496,14 @@ fn multiply_batch_block( let t_device = std::time::Instant::now(); - // Device section pinned to this permit's stream slot: up to `NASSAU_GPU_CONCURRENCY` - // launches overlap on distinct streams, but no more streams (and hence retained pools) - // than that ever exist — see [`GpuPermit`]. Cubecl's per-device runner serializes the - // actual server access; cross-slot/cross-thread reuse of the shared resident handles is - // event-synced by cubecl. `memory_cleanup` below trims only this slot's own pool. - let result = StreamId { value: permit.slot }.executes(|| { + // Device section on this worker's own stable stream (see [`thread_stream_id`]): up to + // `gpu_stream_slots()` distinct streams run concurrently. Cubecl's per-device runner serializes + // the actual server access; the shared resident master/basis (stable in-place-grown buffers) are + // read cross-stream safely. `memory_cleanup` below trims only this stream's own transient pool. + let result = StreamId { + value: thread_stream_id(), + } + .executes(|| { let client = CudaRuntime::client(&CudaDevice::default()); // Shared resident admissible buffers (see [`ResidentDev`]): re-upload the master ONLY // when this block dereferences past the uploaded prefix (`need_cs` / `need_mk`). The @@ -1280,10 +1515,8 @@ fn multiply_batch_block( // then HOST.read, and nothing under either lock blocks on rayon or a permit. The // master is append-only, so the uploaded prefix is always a prefix of the current // host master and every offset `< uploaded` is final. - let (cs_h, cs_len_master) = - resident_dev_handle!(client, need_cs, cs_handle, cs_uploaded, col_sums); - let (mk_h, mk_len_master) = - resident_dev_handle!(client, need_mk, mk_handle, mk_uploaded, masks); + let (cs_h, cs_len_master) = resident_dev_handle!(client, need_cs, cs, col_sums); + let (mk_h, mk_len_master) = resident_dev_handle!(client, need_mk, mk, masks); // 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. @@ -1310,6 +1543,12 @@ fn multiply_batch_block( 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)); const THREADS: u32 = 256; + // Hold the resident READ lock across the multiply that dereferences the resident buffers, + // through the readback that syncs it (see [`RESIDENT_REALLOC`]). All resident growth for this + // block is already done above; a concurrent capacity realloc on another thread waits here for + // this multiply to finish, so the resident handle can never be swapped mid-kernel. Held to + // the end of the device section (dropped after `read_one`). + let _realloc_guard = RESIDENT_REALLOC.read().unwrap(); // 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. Same stream as the multiply below, so it is ordered before it. @@ -1368,6 +1607,8 @@ fn multiply_batch_block( // 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. + // DIAGNOSTIC: memory_cleanup removed to test whether its explicit pool reindex corrupts the + // stable resident buffers under concurrent readers. (Memory bounding restored after.) client.memory_cleanup(); result From 0c1712a334d18eaf9f197f23fd0c38541463aa2a Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Sat, 25 Jul 2026 01:15:24 -0400 Subject: [PATCH 015/127] fp: trace GPU vs CPU row-reduce dispatch (fp::rr) Instrument Matrix::row_reduce to emit a `fp::rr` tracing event for every p=2 reduction with min(rows,cols) >= 1024, recording rows/cols/min and whether the device RREF was taken (path="gpu") or it fell back to CPU M4RI (path="cpu"). The event inherits the active nassau span so each line carries its bidegree/signature context. Confirms on a stem-150 run that every reduction with min >= 8192 (up to 25091x30275, incl. the heavy zero-signature base solves) dispatches to the GPU with no fallbacks; everything below the 8192 FP_CUDA_RR_THRESHOLD stays on CPU. tracing is added as a gpu-gated optional dep. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/fp/Cargo.toml | 4 ++- ext/crates/fp/src/matrix/matrix_inner.rs | 42 +++++++++++++++++++++--- 2 files changed, 41 insertions(+), 5 deletions(-) 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/matrix/matrix_inner.rs b/ext/crates/fp/src/matrix/matrix_inner.rs index 99b014e631..d6ca2af0eb 100644 --- a/ext/crates/fp/src/matrix/matrix_inner.rs +++ b/ext/crates/fp/src/matrix/matrix_inner.rs @@ -677,11 +677,45 @@ 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; + 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(); From f22f8a73058c67971cf176dbc143c8c076562c81 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Sun, 26 Jul 2026 04:59:24 -0400 Subject: [PATCH 016/127] fp,algebra: serialize cubecl multiply vs cooperative GPU row-reduce The intermittent stem-150 hang (flat sm=100%, ~1/6 of runs, always on the zero-signature base solve) was a deadlock in the fp-cuda cooperative row-reduce kernel `panel_factor_coop`. Its grid-wide spin barrier (launched via `cuLaunchCooperativeKernel`) requires ALL its CTAs co-resident, but the algebra Milnor multiply runs on a *separate* CUDA runtime (cubecl) whose kernels concurrently occupy SMs. When a cooperative row-reduce launched while cubecl multiply kernels were resident, its CTAs could not all co-reside; the missing ones never reached the barrier and the resident ones spun forever. Fix: a cross-runtime `fp::GPU_EXCLUSIVE` RwLock. The cooperative row-reduce takes the write lock (drains in-flight cubecl multiplies, then runs with the GPU to itself) around its launch+download; every cubecl multiply takes the read lock across its whole device section (launch through readback, so releasing means the kernel has actually completed). Readers run concurrently; a pending row-reduce briefly excludes them. No lock cycle (cubecl never takes the fp-cuda ctx lock) and no same-thread read->write (a solve's multiply releases before its row-reduce). Also adds a `gpu_row_reduce` tracing span around the GPU reduce -- the diagnostic that localized the wedge (an unclosed span names a stuck reduce, distinguishing an RREF hang from a multiply hang). Validated on H200: stem-150 x16 with 0 wedges (baseline ~1/6), wall time unchanged (261-359s), GPU-vs-CPU chart match preserved. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 9 +++++++-- ext/crates/fp/src/blas/cuda.rs | 20 +++++++++++++++++++- ext/crates/fp/src/lib.rs | 6 ++++++ ext/crates/fp/src/matrix/matrix_inner.rs | 8 ++++++++ 4 files changed, 40 insertions(+), 3 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 767541249e..54baff8871 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -1505,6 +1505,13 @@ fn multiply_batch_block( } .executes(|| { let client = CudaRuntime::client(&CudaDevice::default()); + // Hold the cross-runtime GPU read lock for this whole device section (H2D + kernels + + // readback), so a concurrent cooperative fp-cuda row-reduce (raw cudarc) waits for this + // cubecl multiply to finish before taking the GPU exclusively. Without this, the RREF's + // `cuLaunchCooperativeKernel` grid barrier can't co-reside its CTAs against a resident + // multiply kernel and spins forever (the intermittent stem-150 wedge). See + // [`fp::GPU_EXCLUSIVE`]. Dropped at the end of the closure, after `read_one`. + let _gpu_shared = fp::GPU_EXCLUSIVE.read().unwrap_or_else(|e| e.into_inner()); // Shared resident admissible buffers (see [`ResidentDev`]): re-upload the master ONLY // when this block dereferences past the uploaded prefix (`need_cs` / `need_mk`). The // master grows continually at the frontier, so re-uploading on mere growth ships @@ -1607,8 +1614,6 @@ fn multiply_batch_block( // 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. - // DIAGNOSTIC: memory_cleanup removed to test whether its explicit pool reindex corrupts the - // stable resident buffers under concurrent readers. (Memory bounding restored after.) client.memory_cleanup(); result diff --git a/ext/crates/fp/src/blas/cuda.rs b/ext/crates/fp/src/blas/cuda.rs index 0fba7154ec..550b20c99d 100644 --- a/ext/crates/fp/src/blas/cuda.rs +++ b/ext/crates/fp/src/blas/cuda.rs @@ -14,12 +14,26 @@ //! 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::{Mutex, OnceLock, RwLock}; use fp_cuda::GpuContext; use crate::{matrix::Matrix, prime::TWO}; +/// Serializes the fp-cuda **cooperative** row-reduce (raw cudarc) against the algebra +/// **cubecl** Milnor multiply, which share the one physical GPU from different CUDA +/// runtimes. The row-reduce's `cuLaunchCooperativeKernel` (panel_factor_coop's grid-wide +/// spin barrier) needs *all* its CTAs co-resident; a cubecl multiply kernel occupying SMs +/// at that moment prevents co-residency, so the missing CTAs never reach the barrier and +/// the resident ones spin forever (flat sm=100% — the intermittent stem-150 wedge). +/// +/// Contract: the cooperative row-reduce takes the **write** lock (exclusive GPU) around its +/// launch+download; every cubecl multiply takes a **read** lock held from launch through its +/// readback (so releasing means that kernel has actually completed, not merely enqueued). +/// Readers run concurrently with each other; a pending row-reduce drains them, runs alone, +/// then lets them resume. See `algebra::algebra::milnor_gpu` for the read side. +pub static GPU_EXCLUSIVE: RwLock<()> = RwLock::new(()); + /// Smallest `min(m, k, n)` for which we attempt the GPU matmul. Below this the /// host marshalling (bit-repack into TMA tiles + copies) costs more than it saves. const DEFAULT_THRESHOLD: usize = 2048; @@ -127,7 +141,11 @@ pub(crate) fn try_row_reduce(m: &mut Matrix) -> Option { let stride = cols.div_ceil(64); let limbs = to_limbs(m); + // Exclusive GPU for the cooperative reduce: block new cubecl multiplies and wait for + // in-flight ones to complete, so panel_factor_coop's grid-wide barrier can co-reside all + // its CTAs (see [`GPU_EXCLUSIVE`]). Held across launch + download; ordered before ctx.lock. let (dev_limbs, perm, r, pivot_cols) = { + let _exclusive = GPU_EXCLUSIVE.write().unwrap_or_else(|e| e.into_inner()); 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()?; diff --git a/ext/crates/fp/src/lib.rs b/ext/crates/fp/src/lib.rs index 8d971da2a8..bb8dedec60 100644 --- a/ext/crates/fp/src/lib.rs +++ b/ext/crates/fp/src/lib.rs @@ -12,6 +12,12 @@ pub mod vector; pub mod blas; +/// Cross-runtime GPU serialization lock (cubecl multiply vs. cooperative fp-cuda row-reduce). +/// Re-exported so `algebra`'s Milnor-multiply GPU path can take the read side. See +/// [`blas::cuda::GPU_EXCLUSIVE`]. +#[cfg(feature = "gpu")] +pub use blas::cuda::GPU_EXCLUSIVE; + pub(crate) mod simd; // This is useful for traits that want to implement `Arbitrary`. This lets us specify that they diff --git a/ext/crates/fp/src/matrix/matrix_inner.rs b/ext/crates/fp/src/matrix/matrix_inner.rs index d6ca2af0eb..fcf3261923 100644 --- a/ext/crates/fp/src/matrix/matrix_inner.rs +++ b/ext/crates/fp/src/matrix/matrix_inner.rs @@ -688,6 +688,14 @@ impl Matrix { 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 { From 5723c04927243a36a40463636f6534fb5f6e2a00 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Sun, 26 Jul 2026 13:37:41 -0400 Subject: [PATCH 017/127] fp-cuda: add composable (non-cooperative) row-reduce path, default it MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The device RREF launched three cooperative kernels — panel_factor_coop, promote_coop, block_reduce_coop — synchronized by a hand-rolled grid-wide spin barrier. cuLaunchCooperativeKernel requires all the grid's CTAs to be co-resident, which only holds when this process owns the whole GPU. When another CUDA runtime shares the device (cubecl's Milnor multiply in the nassau GPU resolution), its kernels occupy SMs, the reduce's CTAs can't all co-reside, the missing ones never reach the barrier, and the resident ones spin forever — the intermittent stem-150 wedge (flat sm=100%). Add an FP_CUDA_RR_COOP switch (rr_coop()). Default off: the forward pass runs the single-CTA panel_factor one limb at a time, promotion uses the grid- strided promote_pivots, and back-substitution uses the single-CTA block_reduce_rref — none launched cooperatively, so the reduce composes with concurrent GPU work. Set FP_CUDA_RR_COOP=1 to opt into the cooperative path on a dedicated GPU (measured 2-3x faster at Nassau strides, up to ~10x on large dense half-rank matrices). Both paths validated bit-exact vs CPU row_reduce (row_reduce_demo) and vs the CPU BLAS3 oracle at 2^16/2^17 (reduce_pow2_half). Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/fp-cuda/src/lib.rs | 53 ++++++++++++++++++++++++++++++----- 1 file changed, 46 insertions(+), 7 deletions(-) diff --git a/ext/crates/fp-cuda/src/lib.rs b/ext/crates/fp-cuda/src/lib.rs index 9a574ddcc9..f5f9dcc5cb 100644 --- a/ext/crates/fp-cuda/src/lib.rs +++ b/ext/crates/fp-cuda/src/lib.rs @@ -40,6 +40,29 @@ 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 +/// forward pass runs the single-CTA `panel_factor` over one limb at a time, promotion +/// uses the grid-strided `promote_pivots`, and back-substitution uses the single-CTA +/// `block_reduce_rref` — none of which launch cooperatively. Set `FP_CUDA_RR_COOP=1` +/// to opt into the faster cooperative path on a GPU dedicated to this process +/// (measured +5–18% on the wide reductions). +fn rr_coop() -> bool { + std::env::var("FP_CUDA_RR_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. @@ -1076,11 +1099,20 @@ impl GpuContext { 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). - let bl = if let Some(v) = std::env::var("FP_CUDA_BL") + // the measured optimum (flat at bl≈12–16). The single-CTA `panel_factor` + // used on the non-cooperative path handles exactly one limb, so force bl=1. + let bl = if !coop { + 1 + } else if let Some(v) = std::env::var("FP_CUDA_BL") .ok() .and_then(|v| v.parse::().ok()) { @@ -1098,7 +1130,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 +1179,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 { + // Single-CTA factor of this one limb (bl_eff == 1). Scans all rows + // [r, rows); compaction still parks dead rows and drives the break. + self.panel_factor(m, &mut perm, &mut l, ppanel, r)? + }; if pr > 0 { for &q in &pivcols { pivot_cols.push(q as usize); @@ -1416,8 +1453,10 @@ impl GpuContext { // 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; + // neutral at n=2¹⁵ (stride 512), +6% at 2¹⁶, +18% at 2¹⁷. Only when the + // cooperative path is enabled (see [`rr_coop`]); otherwise the single-CTA + // block_reduce_rref runs so back-substitution launches no cooperative grid. + let use_coop = rr_coop() && stride >= 1024; // Pivots per back-substitution block. Wider blocks raise the X·U GEMM's // contraction dimension bp toward TILE_K, cutting its K-padding waste From ad185fb4a965fbb00cecfcbd270eaa6d7aabd31d Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Sun, 26 Jul 2026 13:57:32 -0400 Subject: [PATCH 018/127] fp-cuda: multi-SM row-reduce via kernel-boundary sync (not cooperative) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The composable default from the previous commit fell back to single-CTA kernels (panel_factor, block_reduce_rref) that use one SM of ~132 — 8-10x slower than the cooperative path on large dense matrices, since the cooperative kernels' real win is spreading each sequential bit-step across the whole grid. Recover that all-SM parallelism WITHOUT a cooperative launch by replacing the in-grid grid_sync with kernel-boundary (stream-ordered) synchronization, the way cuSOLVER/cuBLAS build grid-wide multi-step algorithms: - panel_factor: pf_find -> pf_swap -> pf_xor per bit-step (panel_factor_streamed) - block_reduce: br_cond -> br_xor per pivot (block_reduce_elem streamed arm) All per-step state stays on the device (pivot count, find-first result, pivword, clear-conditions), so the host issues every launch without a readback and the latency hides behind GPU work; only the final (pr, pivcols) is copied back. No cuLaunchCooperativeKernel anywhere on the default path, so it can't deadlock against a concurrent cubecl kernel. The streamed back-substitution also adopts the cooperative path's wide-block TRSM (bp=1024 + X.U GEMM) since the GEMM composes. Perf (H200 half-rank square, device-only) vs cooperative: 2^13 s128 0.080 vs 0.066 (1.21x) 2^15 s512 0.642 vs 0.577 (1.11x) 2^14 s256 0.195 vs 0.166 (1.17x) 2^16 s1024 2.27 vs 1.17 (1.95x) 2^17 s2048 8.63 vs 4.52 (1.91x) Within ~1.1-1.2x of cooperative through stride 512 (the Nassau regime; the wedge matrix was stride 215) and 3-6x faster than the single-CTA fallback at large sizes. The residual past stride 1024 is where cooperative additionally fuses promote/block-reduce; closable with CUDA graphs if those sizes ever matter. Both paths validated bit-exact vs CPU row_reduce (row_reduce_demo) and the CPU BLAS3 oracle at 2^16 (reduce_pow2_half). Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/fp-cuda/cuda_kernels/matmul_b1.cu | 148 ++++++++++ ext/crates/fp-cuda/src/lib.rs | 277 ++++++++++++++++--- 2 files changed, 392 insertions(+), 33 deletions(-) diff --git a/ext/crates/fp-cuda/cuda_kernels/matmul_b1.cu b/ext/crates/fp-cuda/cuda_kernels/matmul_b1.cu index 7dd9334b54..0403d4f809 100644 --- a/ext/crates/fp-cuda/cuda_kernels/matmul_b1.cu +++ b/ext/crates/fp-cuda/cuda_kernels/matmul_b1.cu @@ -757,6 +757,110 @@ 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: smallest perm position p in [r+g_pr, m) whose row has bit q set; +// atomicMin into g_min. Grid-strided over rows. +extern "C" __global__ void pf_find( + const u64_t* __restrict__ m_buf, + const unsigned* __restrict__ perm, + const unsigned* __restrict__ g_pr, + int* __restrict__ g_min, + unsigned plimb, unsigned j, unsigned r, unsigned m, unsigned stride) +{ + 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 = *g_pr; + + int local_min = 0x7fffffff; + 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]); +} + +// One thread: if a pivot was found, read its bl panel limbs into g_pivword, swap +// it up to position r+g_pr (perm swap), record its column, bump g_pr, and reset +// g_min to INF for the next step. Publishes the pivot position (or INF) to +// g_pivpos so pf_xor knows whether to run. Free column ⇒ everything untouched. +extern "C" __global__ void pf_swap( + const 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 ppanel, unsigned bl, unsigned r, unsigned stride, unsigned q) +{ + int pivpos = *g_min; + *g_pivpos = pivpos; + if (pivpos == 0x7fffffff) return; // free column + unsigned pr = *g_pr; + 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; +} + +// 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 +1080,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 f5f9dcc5cb..862be8c0c2 100644 --- a/ext/crates/fp-cuda/src/lib.rs +++ b/ext/crates/fp-cuda/src/lib.rs @@ -51,12 +51,16 @@ fn adaptive_bl(stride: usize) -> usize { /// 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 -/// forward pass runs the single-CTA `panel_factor` over one limb at a time, promotion -/// uses the grid-strided `promote_pivots`, and back-substitution uses the single-CTA -/// `block_reduce_rref` — none of which launch cooperatively. Set `FP_CUDA_RR_COOP=1` -/// to opt into the faster cooperative path on a GPU dedicated to this process -/// (measured +5–18% on the wide reductions). +/// **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 the streamed `pf_find`/`pf_swap`/`pf_xor` triplet per bit- +/// step, 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. Measured within ~1.1–1.2× +/// of cooperative through stride 512 (the Nassau regime), widening to ~1.9× only past +/// stride 1024 where cooperative also fuses promote and block-reduce. 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()) @@ -101,6 +105,9 @@ pub struct GpuContext { xor_into: CudaFunction, panel_factor: CudaFunction, panel_factor_coop: CudaFunction, + pf_find: CudaFunction, + pf_swap: CudaFunction, + pf_xor: CudaFunction, mark_live: CudaFunction, promote_pivots: CudaFunction, promote_coop: CudaFunction, @@ -108,6 +115,8 @@ 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, } @@ -123,6 +132,9 @@ impl GpuContext { 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 = module.load_function("pf_find")?; + let pf_swap = module.load_function("pf_swap")?; + 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")?; @@ -130,6 +142,8 @@ 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 { @@ -141,6 +155,9 @@ impl GpuContext { xor_into, panel_factor, panel_factor_coop, + pf_find, + pf_swap, + pf_xor, mark_live, promote_pivots, promote_coop, @@ -148,6 +165,8 @@ impl GpuContext { gather_rows, block_reduce_rref, block_reduce_coop, + br_cond, + br_xor, gather_cols, xor_into_perm, }) @@ -853,6 +872,149 @@ 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.ctx.default_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), and the pivot row's + // bl panel limbs. g_min starts INF; g_pr starts 0. + 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)?; + + // 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); + + let (r_u, m_u, stride_u, l_stride_u, ppanel_u, bl_u) = ( + r as u32, + m_active 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, + }; + let one_cfg = LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (1, 1, 1), + shared_mem_bytes: 0, + }; + + for cc in 0..(bl * 64) { + let q = ppanel * 64 + cc; + if q >= m.cols { + break; + } + let (plimb_u, j_u, cc_u, q_u) = + ((ppanel + cc / 64) as u32, (cc & 63) as u32, cc as u32, q as u32); + + // (1) find-first pivot for column q → g_min. + { + let mut lb = stream.launch_builder(&self.pf_find); + lb.arg(&m.buf) + .arg(&*perm) + .arg(&g_pr) + .arg(&mut g_min) + .arg(&plimb_u) + .arg(&j_u) + .arg(&r_u) + .arg(&m_u) + .arg(&stride_u); + unsafe { lb.launch(find_cfg) }?; + } + // (2) swap the pivot up, record it, bump g_pr, reset g_min (1 thread). + { + let mut lb = stream.launch_builder(&self.pf_swap); + lb.arg(&m.buf) + .arg(&mut *perm) + .arg(&pivcols) + .arg(&mut g_pivword) + .arg(&mut g_min) + .arg(&mut g_pr) + .arg(&g_pivpos) + .arg(&ppanel_u) + .arg(&bl_u) + .arg(&r_u) + .arg(&stride_u) + .arg(&q_u); + unsafe { lb.launch(one_cfg) }?; + } + // (3) clear the pivot from the rows below, across all bl panel limbs. + { + 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 @@ -1108,11 +1270,10 @@ impl GpuContext { // 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 single-CTA `panel_factor` - // used on the non-cooperative path handles exactly one limb, so force bl=1. - let bl = if !coop { - 1 - } else if let Some(v) = std::env::var("FP_CUDA_BL") + // 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()) { @@ -1182,9 +1343,9 @@ impl GpuContext { let (pr, pivcols) = if coop { self.panel_factor_coop(m, &mut perm, &mut l, ppanel, bl_eff, r, m_active)? } else { - // Single-CTA factor of this one limb (bl_eff == 1). Scans all rows - // [r, rows); compaction still parks dead rows and drives the break. - self.panel_factor(m, &mut perm, &mut l, ppanel, r)? + // 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 { @@ -1333,6 +1494,7 @@ impl GpuContext { s: usize, e: usize, use_coop: bool, + streamed: bool, br_barrier: &mut CudaSlice, br_cond: &CudaSlice, br_ctas: u32, @@ -1362,6 +1524,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 { @@ -1405,24 +1605,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(()) } @@ -1449,22 +1652,27 @@ impl GpuContext { 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¹⁷. Only when the - // cooperative path is enabled (see [`rr_coop`]); otherwise the single-CTA - // block_reduce_rref runs so back-substitution launches no cooperative grid. - let use_coop = rr_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") @@ -1485,9 +1693,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() @@ -1527,6 +1736,7 @@ impl GpuContext { e, base_bp, use_coop, + streamed, &mut br_barrier, &br_cond, br_ctas, @@ -1539,6 +1749,7 @@ impl GpuContext { s, e, use_coop, + streamed, &mut br_barrier, &br_cond, br_ctas, From 4f83a85ba560d761a175dd7f4441c5f5d28acbd4 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Sun, 26 Jul 2026 14:38:30 -0400 Subject: [PATCH 019/127] fp-cuda: fuse find+swap in streamed panel factor; add rr timing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Merge pf_find and pf_swap into pf_find_swap using a threadfence "last-CTA finalize" (a grid-wide reduction, not a barrier, so no co-residency needed), cutting the streamed forward pass from 3 to 2 launches per bit-step and removing the serial 1-thread swap kernel. Cap the streamed grid at FP_CUDA_PF_CTAS (128) like the cooperative kernel. Add FP_CUDA_RR_TIMING to split forward/back timing. Diagnostic result: the streamed forward pass is work-bound (~cols^2), not launch-bound — merging, capping, and grid size leave it unchanged — so it is the per-step relaunch efficiency vs the persistent cooperative grid, ~2.7x at 2^17 but shrinking with size (3.0x at 2^16). Still bit-exact vs CPU row_reduce and the BLAS3 oracle at 2^16. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/fp-cuda/cuda_kernels/matmul_b1.cu | 63 ++++++++------- ext/crates/fp-cuda/src/lib.rs | 84 +++++++++++--------- 2 files changed, 81 insertions(+), 66 deletions(-) diff --git a/ext/crates/fp-cuda/cuda_kernels/matmul_b1.cu b/ext/crates/fp-cuda/cuda_kernels/matmul_b1.cu index 0403d4f809..10ce16a70d 100644 --- a/ext/crates/fp-cuda/cuda_kernels/matmul_b1.cu +++ b/ext/crates/fp-cuda/cuda_kernels/matmul_b1.cu @@ -773,27 +773,44 @@ extern "C" __global__ void panel_factor_coop( // 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: smallest perm position p in [r+g_pr, m) whose row has bit q set; -// atomicMin into g_min. Grid-strided over rows. -extern "C" __global__ void pf_find( - const u64_t* __restrict__ m_buf, - const unsigned* __restrict__ perm, - const unsigned* __restrict__ g_pr, +// 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 plimb, unsigned j, unsigned r, unsigned m, unsigned stride) + 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; - 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); + 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(); @@ -802,26 +819,16 @@ extern "C" __global__ void pf_find( __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; -// One thread: if a pivot was found, read its bl panel limbs into g_pivword, swap -// it up to position r+g_pr (perm swap), record its column, bump g_pr, and reset -// g_min to INF for the next step. Publishes the pivot position (or INF) to -// g_pivpos so pf_xor knows whether to run. Free column ⇒ everything untouched. -extern "C" __global__ void pf_swap( - const 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 ppanel, unsigned bl, unsigned r, unsigned stride, unsigned q) -{ int pivpos = *g_min; *g_pivpos = pivpos; - if (pivpos == 0x7fffffff) return; // free column - unsigned pr = *g_pr; + 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]; diff --git a/ext/crates/fp-cuda/src/lib.rs b/ext/crates/fp-cuda/src/lib.rs index 862be8c0c2..b65870bc73 100644 --- a/ext/crates/fp-cuda/src/lib.rs +++ b/ext/crates/fp-cuda/src/lib.rs @@ -105,8 +105,7 @@ pub struct GpuContext { xor_into: CudaFunction, panel_factor: CudaFunction, panel_factor_coop: CudaFunction, - pf_find: CudaFunction, - pf_swap: CudaFunction, + pf_find_swap: CudaFunction, pf_xor: CudaFunction, mark_live: CudaFunction, promote_pivots: CudaFunction, @@ -132,8 +131,7 @@ impl GpuContext { 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 = module.load_function("pf_find")?; - let pf_swap = module.load_function("pf_swap")?; + let pf_find_swap = module.load_function("pf_find_swap")?; let pf_xor = module.load_function("pf_xor")?; let mark_live = module.load_function("mark_live")?; let promote_pivots = module.load_function("promote_pivots")?; @@ -155,8 +153,7 @@ impl GpuContext { xor_into, panel_factor, panel_factor_coop, - pf_find, - pf_swap, + pf_find_swap, pf_xor, mark_live, promote_pivots, @@ -905,13 +902,15 @@ impl GpuContext { 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), and the pivot row's - // bl panel limbs. g_min starts INF; g_pr starts 0. + // 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. @@ -925,10 +924,20 @@ impl GpuContext { .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, stride_u, l_stride_u, ppanel_u, bl_u) = ( + 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, @@ -944,52 +953,38 @@ impl GpuContext { block_dim: (THREADS, 1, 1), shared_mem_bytes: 0, }; - let one_cfg = LaunchConfig { - grid_dim: (1, 1, 1), - block_dim: (1, 1, 1), - shared_mem_bytes: 0, - }; for cc in 0..(bl * 64) { - let q = ppanel * 64 + cc; - if q >= m.cols { + if ppanel * 64 + cc >= m.cols { break; } - let (plimb_u, j_u, cc_u, q_u) = - ((ppanel + cc / 64) as u32, (cc & 63) as u32, cc as u32, q as u32); + let (plimb_u, j_u, cc_u) = + ((ppanel + cc / 64) as u32, (cc & 63) as u32, cc as u32); - // (1) find-first pivot for column q → g_min. + // (1) find-first pivot for column q and swap it up — one launch, its + // grid-wide min finalized by the last CTA to arrive (no barrier). { - let mut lb = stream.launch_builder(&self.pf_find); - lb.arg(&m.buf) - .arg(&*perm) - .arg(&g_pr) - .arg(&mut g_min) - .arg(&plimb_u) - .arg(&j_u) - .arg(&r_u) - .arg(&m_u) - .arg(&stride_u); - unsafe { lb.launch(find_cfg) }?; - } - // (2) swap the pivot up, record it, bump g_pr, reset g_min (1 thread). - { - let mut lb = stream.launch_builder(&self.pf_swap); - lb.arg(&m.buf) + 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(&q_u); - unsafe { lb.launch(one_cfg) }?; + .arg(&n_u); + unsafe { lb.launch(find_cfg) }?; } - // (3) clear the pivot from the rows below, across all bl panel limbs. + // (2) clear the pivot from the rows below, across all bl panel limbs. { let mut lb = stream.launch_builder(&self.pf_xor); lb.arg(&mut m.buf) @@ -1775,6 +1770,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)) From 1f5afea20c33847fda3e4fc27c51de5204fdfd4d Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Sun, 26 Jul 2026 15:02:07 -0400 Subject: [PATCH 020/127] fp-cuda: fuse trailing-XOR with next pivot search (lookahead panel step) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit pf_step does column cc-1's clear and column cc's find+swap in one launch, since the forward sweep alternates xor(col j) / find(col j+1) over the same below-row range — so each thread reads its row once and sees its own XOR before scanning. Halves the forward pass's launches (one per column instead of two). Marginal on its own (~5%), confirming the forward-pass gap vs cooperative is the per-step GPU relaunch cost of non-persistent kernels, not launch count — it resists launch reduction and only closes with scale (the O(cols) term vanishing against O(cols^2) work): ~2x total at 2^17, ~1.4x at 2^18. Bit-exact vs the CPU BLAS3 oracle at 2^16. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/fp-cuda/cuda_kernels/matmul_b1.cu | 91 ++++++++++++ ext/crates/fp-cuda/src/lib.rs | 137 +++++++++++-------- 2 files changed, 173 insertions(+), 55 deletions(-) diff --git a/ext/crates/fp-cuda/cuda_kernels/matmul_b1.cu b/ext/crates/fp-cuda/cuda_kernels/matmul_b1.cu index 10ce16a70d..8dd75c2e5d 100644 --- a/ext/crates/fp-cuda/cuda_kernels/matmul_b1.cu +++ b/ext/crates/fp-cuda/cuda_kernels/matmul_b1.cu @@ -840,6 +840,97 @@ extern "C" __global__ void pf_find_swap( *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. diff --git a/ext/crates/fp-cuda/src/lib.rs b/ext/crates/fp-cuda/src/lib.rs index b65870bc73..00eacfd6eb 100644 --- a/ext/crates/fp-cuda/src/lib.rs +++ b/ext/crates/fp-cuda/src/lib.rs @@ -54,13 +54,15 @@ fn adaptive_bl(stride: usize) -> usize { /// **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 the streamed `pf_find`/`pf_swap`/`pf_xor` triplet per bit- -/// step, 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. Measured within ~1.1–1.2× -/// of cooperative through stride 512 (the Nassau regime), widening to ~1.9× only past -/// stride 1024 where cooperative also fuses promote and block-reduce. Set -/// `FP_CUDA_RR_COOP=1` to opt into the cooperative path on a dedicated GPU. +/// 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()) @@ -106,6 +108,7 @@ pub struct GpuContext { panel_factor: CudaFunction, panel_factor_coop: CudaFunction, pf_find_swap: CudaFunction, + pf_step: CudaFunction, pf_xor: CudaFunction, mark_live: CudaFunction, promote_pivots: CudaFunction, @@ -132,6 +135,7 @@ impl GpuContext { 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")?; @@ -154,6 +158,7 @@ impl GpuContext { panel_factor, panel_factor_coop, pf_find_swap, + pf_step, pf_xor, mark_live, promote_pivots, @@ -954,55 +959,77 @@ impl GpuContext { shared_mem_bytes: 0, }; - for cc in 0..(bl * 64) { - if ppanel * 64 + cc >= m.cols { - break; - } - let (plimb_u, j_u, cc_u) = - ((ppanel + cc / 64) as u32, (cc & 63) as u32, cc as u32); + // Valid columns in this panel (the last panel may be short). + let ncols = (bl * 64).min(m.cols - ppanel * 64); - // (1) find-first pivot for column q and swap it up — one launch, its - // grid-wide min finalized by the last CTA to arrive (no barrier). - { - 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) }?; - } - // (2) clear the pivot from the rows below, across all bl panel limbs. - { - 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) }?; - } + // (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; From 109fb3b780f0dc1382aadd7b3a75874fc5a04143 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Sun, 26 Jul 2026 15:22:40 -0400 Subject: [PATCH 021/127] fp,algebra: drop GPU_EXCLUSIVE lock; the row-reduce is now composable MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The cross-runtime fp::GPU_EXCLUSIVE RwLock existed only to keep the cooperative fp-cuda row-reduce (cuLaunchCooperativeKernel grid barrier) from co-scheduling against a concurrent cubecl Milnor multiply, which could prevent CTA co-residency and deadlock the grid barrier (the stem-150 wedge). The default row-reduce is now the streamed kernel-boundary path (no cooperative launch anywhere), so it composes with concurrent GPU work by construction — the lock is unnecessary. Remove the RwLock, its fp re-export, and the read-side guard in the algebra Milnor-multiply device section. The gpu_row_reduce tracing span is kept (a useful wedge/hang diagnostic). FP_CUDA_RR_COOP=1 still selects the cooperative kernels for a dedicated GPU; do not combine that with the concurrent nassau multiply. Validated: stem-150 resolves in 242s with 30042 GPU row-reduce events and no wedge (full hunt in progress). Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/150/nassau_differential/c/126/1 | Bin 0 -> 1812 bytes ext/150/nassau_differential/c/127/0 | Bin 0 -> 1812 bytes ext/150/nassau_differential/c/127/1 | Bin 0 -> 2820 bytes ext/150/nassau_differential/c/128/0 | Bin 0 -> 2732 bytes ext/150/nassau_differential/c/128/1 | Bin 0 -> 2884 bytes ext/150/nassau_differential/c/129/0 | Bin 0 -> 2772 bytes ext/150/nassau_differential/c/129/1 | Bin 0 -> 2828 bytes ext/150/nassau_differential/c/130/0 | Bin 0 -> 2860 bytes ext/150/nassau_differential/c/130/1 | Bin 0 -> 2940 bytes ext/150/nassau_differential/c/131/0 | Bin 0 -> 3060 bytes ext/150/nassau_differential/c/131/1 | Bin 0 -> 3268 bytes ext/150/nassau_differential/c/132/0 | Bin 0 -> 3948 bytes ext/150/nassau_differential/c/132/1 | Bin 0 -> 3508 bytes ext/150/nassau_differential/c/133/0 | Bin 0 -> 1576 bytes ext/150/nassau_differential/c/133/1 | Bin 0 -> 1404 bytes ext/150/nassau_differential/zarr.json | 63 +++++++++++++++++++ ext/150/zarr.json | 18 ++++++ ext/crates/algebra/src/algebra/milnor_gpu.rs | 7 --- ext/crates/fp/src/blas/cuda.rs | 22 +------ ext/crates/fp/src/lib.rs | 6 -- 20 files changed, 84 insertions(+), 32 deletions(-) create mode 100644 ext/150/nassau_differential/c/126/1 create mode 100644 ext/150/nassau_differential/c/127/0 create mode 100644 ext/150/nassau_differential/c/127/1 create mode 100644 ext/150/nassau_differential/c/128/0 create mode 100644 ext/150/nassau_differential/c/128/1 create mode 100644 ext/150/nassau_differential/c/129/0 create mode 100644 ext/150/nassau_differential/c/129/1 create mode 100644 ext/150/nassau_differential/c/130/0 create mode 100644 ext/150/nassau_differential/c/130/1 create mode 100644 ext/150/nassau_differential/c/131/0 create mode 100644 ext/150/nassau_differential/c/131/1 create mode 100644 ext/150/nassau_differential/c/132/0 create mode 100644 ext/150/nassau_differential/c/132/1 create mode 100644 ext/150/nassau_differential/c/133/0 create mode 100644 ext/150/nassau_differential/c/133/1 create mode 100644 ext/150/nassau_differential/zarr.json create mode 100644 ext/150/zarr.json diff --git a/ext/150/nassau_differential/c/126/1 b/ext/150/nassau_differential/c/126/1 new file mode 100644 index 0000000000000000000000000000000000000000..79167dadebc6aa25186d1b635c39e443fbecaef8 GIT binary patch literal 1812 zcmZQ%U|FknuHKeV{;9*i@3=+RNv&usJW<4$c+#} zbwYJQb>cUHfBGN3sXd<;J|j5Kz-71vkKh%2f}gPaeF+=un=1^b;1b+~hwu_U!VlPG zmh*(*BwU0Wa37w-2lx&8&9nb( zQ;?`k@|G&=$GW3SQx{v)+D%3#^8Z?6oOb`xB!>){A66oRj8|<=T3(i@ac{c@Kc(66 z>YiHPk;+iWtZf_3rBrgtixqx2x9Otg%Um@ubkE%WI&Zs{n$tt=Vm0IR(tbM!FMnB5 z>($fy1LokReK$X^SiCen8|ST|KAx={$-$L*%Uv2}E^5=WyVpi7Q{EffxqhNXzjpt> zqS@bqzGr-yu@LK+GHcr=H=)G;_Yz!meq;Sp8)?RVG+aa*?P0I%=T@wZ8yhHisek$hx6Ie+!S8owVL00+TgF#WrttG^gHPyQ>a*MR-P O8eae$1UD6*UAzPM$}A)R literal 0 HcmV?d00001 diff --git a/ext/150/nassau_differential/c/128/1 b/ext/150/nassau_differential/c/128/1 new file mode 100644 index 0000000000000000000000000000000000000000..26fff2ee58b0a51ce4f54e1769983e11d63d70f5 GIT binary patch literal 2884 zcmZQ%U|^5{Vn!ea0d%_4GioD9Y$U4@0EY~OfQ3)X=4M_Hd$`FlLH+vHe&r_bi=cA|3;ULEX_1`(4MXBii8G<12s?2@uFYCbB_TU-v^DKfyQq^<1awt??L0=K;wTw;|s7t5)RBgI%s?kG=2&izX6Ru2aUf2 zjeiA={{fB9!-nb}4K%(B8b1MzUxUVVGyRu)cMv>-@oHN!y@mZ?N26H#ct@M+$hoX{xba|TlwUfjc_eDJI{Zm4NdiMtc2I}{0?keE64MD`F5Tui84>- zGHN-bh5N>SpoVi(V_IoeG4m;(=i&GBGM`$@V$$j=%;?3jz83T>AWwwMEj8s=&@-1j zPpJL5IlfeS)cMaQt!wym95t)?TWO0ZW~e+?k1>~!w%eELMr|Q!?W5xv)VP*C{R2U% zsq4ixY~k2j??!Pvd;V3YmFA$BZB@w$j7i5)OvYaiqlRms&UZe=?3iVCZl;2ba=hx{ zv#OCN`$)wl^zfQ=NI`$H-al@}dCV~-apGkp2Q`^bKChmpPoq|3O5&{l?Idbed&u8a z@1{EkaU9oeuT6~qkH&TK-mB-gkmh;2vM+=&SCUrr?u>C?sN<|^w>pPb^;v0HySmpa zXr3Z`%UJf0)jRjgT2#*(_9SYOX}#aJJ#-&E#pJOczF?%QI@umT58o|;gn8=of^H{z z_-+zSQW0YGTW7WGsq4r6g2AB#j?Xnl&D`IQ2K#f-vs&@UeAn)I520tB;<3Agn24<( zLeHwd9+$A`{YZVOah7$v$Re$PTFn(OdiH$3(Qb5BlRDhGi38L!H-j$>+zC7Y4C$wnpicoC8=}AC{2p1P!VSR3iMbBr2UPbn zWBy+7j{v8DZM0b&&jsu!=J_^(j`>f4eiiru@EhQ-#GKDg_b_uQFy`MuI@fa)_#$u* z@C)Eiz{X*>SbwKY?&i$?_<$RLTY;m%-N4U)vHlOBXNcYAw1kWPI5AfO9{`R5_Y(8^ gV*U~EPXRmV4-Dt`0q-I%q4}Nzy$87U`o6t?0Pz?wCIA2c literal 0 HcmV?d00001 diff --git a/ext/150/nassau_differential/c/129/1 b/ext/150/nassau_differential/c/129/1 new file mode 100644 index 0000000000000000000000000000000000000000..28869883990af61c742826439150d19589c16990 GIT binary patch literal 2828 zcmeI!FGvGn9LDkYbk2zrFDxt=1P2-x6crW?3k!;d1%t>1g98zx%EH2eL9rM%5E=eK zqhaWRj0KY#77QW^CYi8c5Vm0O#{J%~Z@B7CJUH(C+{3}|x%a{uN-5Jm_0c1XkBOtG zVXyyM=9hk%tLFy${bx#b`g;k;{%zf5{oKD^$uI;R^f$H{np1GoYSl9G-<&;VfK$ zZ{R2R3pTsk=ZV4zcm>YDhj0mggg@b+uYH~|I1XEI8qUE*xC+PUP;~z_yztI2y literal 0 HcmV?d00001 diff --git a/ext/150/nassau_differential/c/130/0 b/ext/150/nassau_differential/c/130/0 new file mode 100644 index 0000000000000000000000000000000000000000..9df80ddb5b00fa3125e6415d4df08701f7b9768d GIT binary patch literal 2860 zcmaKtX-HI26vxl#6q9WuTB)==ZCWX`9HSy?f~=&SY(*l$7D&oa!=&Ps5)~~I3DuXF zY6_}p`w&SK4i;$np=cv4DaoK{u|lDS!0vhP+<5:ij~|Nj2>p10h&x`@cDnvVX` z$*s-_+CGIQ4D zK8LsPjFq!aF8JJ?hoae*C$rhV1#1kn2+{7xEq{<@3TwOEVm*_y{!C$H28i~hkKIsO zxe4oMX^$tr@7^Iqg+{g;JtMTo%$UpF3A4)=cIWX@a65Yvu70gT&oBlmy!-B^a{gBN z@i%*%-$D0>G(3#rd_Pa(Fvv9P-PC?Nc85KNYMeQ8?W6{ckePMvw^c}u?QY+ds=aaI z$@%9U=}!H{U$Wnc@OKwwETeeA!E9~aaLkY|b)i`s^vy-|8<{P>7 z8?`L@uy_Z1oZoA*z$gVSdOYz|@pkJs^!(fVL~&`aiE`3n1K8vIrW0hHBNI=LM>bfP zzb;aWEKgYRuG$sM4AnSu6xqqX`phvEZ(5ArNnXTN)`u zrB&^daehzd+9X-~(%a5zoTrWgE0XGi(~&I}ws~jkb!6Lwy)108&N@SB_HeARle&&P zd@5ysQ*C0@8d-8=Weu(~OT0_ZUnxh{(_Yr^ME_o#-xAMiCtD@;()$DrthYy4?x2&S znHj2a=Gbf}zM2Qr-<-@b(fdl2hu4MByFhzg`FopKs(q&q2od47Xq1jOB=&aAT0E`!at^i#lTF45pTaB zz5Wc)t3j86#`;0v5zvjG+eGX8@w)B%(a-M(y;`)sU#wpSJP0}>TGx98yhF54_DxE! z|KE=vG}d1VycqNW&|%T~ex3mT0@}j|xr*+d-v>MZx&U;AXk9-H`U&VipmQQWlKui| CBwDBd literal 0 HcmV?d00001 diff --git a/ext/150/nassau_differential/c/130/1 b/ext/150/nassau_differential/c/130/1 new file mode 100644 index 0000000000000000000000000000000000000000..cb379032d71cd5eaab6f83ec98918e0ff49eb832 GIT binary patch literal 2940 zcmd6oT}V_x6vyXoHx;qEwU(7@xwNqcf*V_~MZy{iLO#U?DF!|G;8QFV6=5OmLmw#l z;6qT9^Z`QKCm9zXEHbDltp^K}4O)KMread9!g~o@M0Ia?$oZ%= z)5AHiy1iL79aA^(Oq`dB$4?)iwuJUL`r_^Pf32;eSU!x5Kf_ppq#aB4Rw~WR!BWyJ z%3NJ(OEb=APi*+xDb)P(Gap)r;vHz1y^LpBUO0<)gAZ;m?nF;{!BeJtJ7!NbpojBi zi6ZSebM8wWo}2H^uf89@eEqH#rnswKXLO}Fw`$n0_WT_LANFmSp1wXxB!$8p{YPZFE%XZQb!W`1WI zXidi!aaZVd!94}L~bxMR#xfm_lbxnEA_0~yNO^$b~`O!@DI(ov| zp*=~Z}>a2mJ= zco6sz@H^rNy=`?6*_dmHdH+V>y}(_-XMyhmzXbjSTtb(CFVN_!jVU;2B`s_V(xfs)4rw?*={q+y|TmegZrV zEEfJ8PY5^xoB~c0N4#Z3Kd=j&gZXp7<=$U%q24jz7T``|-oJ;qmf~?i&jE{K&&cy5 P#9W^wa69mt()!#FK(xJg literal 0 HcmV?d00001 diff --git a/ext/150/nassau_differential/c/131/0 b/ext/150/nassau_differential/c/131/0 new file mode 100644 index 0000000000000000000000000000000000000000..e9dd7430f4f8c62fd5620ed50afd6d0e3e6f3ab2 GIT binary patch literal 3060 zcmaKueM}Tb7{=$8D=zk;i=hS$B#Q-e1{4wq@#CV0l3G7fk3fz15f42qX{chc1}km1 z1=2=pq=<%CzSL-{m_ix_5;fpgA!&nDq_xy0t!b(#6_Y~wqY^r^`_5(WY0I0;%{=pa zpLu8J_IA&~7}I!@5krHfjFHn0KDM(wV#%>(bCDe)mVI`j30X0*fv1OEg6X`8UPO-7 ziBHFDb#;cq^z9+`ecsPckewkGb-nQ= zWR=AB{d8>z*&D=M`M;kNOy^DXA|x|~p{!T+io7H)rWLopj=AzIgxL@V;}1o2p2*f- zk!!!dUKW5ycGBO~$PQYF&3xSU8nR2on!XwO4A}`{4_hkZk>P!5-^*U|eI7N|fW82+ z{0FDrMwURVD5oI~Svj%ugQv5Q-6Gam*r*`;f>`;IiWSH{Ahx5Rc@MJph$U6`PavD> z`WBF{`%cG7^i>fn&gvOOHb`tBP`Mo0Wn#JQ?S5oOh|S5p--8V6={#5DhV^O|U+I%H~-8445lI~D)l2=*@eI=VG zh4=sVs_RO$Yp3!sKjTH+N(|GTEV9PDoVJW{^EqX&^Ql;j z!l|pSfH9upC>k%{Hr<-s8L_Tdn4@_!7vX)fg@QW$udr0yOEct<6rGnj89VMtI3PTJ z-W1+Mgi!%m$C?fO6pY*d9bNnpb83?_!&Nw=+vEE|tvb9rjIs`Re3Uf$$BTliQ+sF!x-HIbR`)EtIlFyl>4Gyde~8QC#!vGCZ>ztu_)T}xgd zUu*5gj%Dbn5Aj&LP!)3tO`zooJH{p<$in1%AN$3WAm`=#T59vA<2mB_FQU%sg6nmF zhVqklfjBHKo5gf>E;gZ2Q=3A0^bU{1auAPR?y#qegqopJL#U5`{hrikn{9_-UNCd z=qk_`NXzd7>)(U?QPOGFHkkVbYv3Ur0Nq4d=HCUaTJI(%Io9_OkEiwhkY7W(kT{zb Wx*ree^~66RE$dwYJ$z&T)qemh*kfh@ literal 0 HcmV?d00001 diff --git a/ext/150/nassau_differential/c/131/1 b/ext/150/nassau_differential/c/131/1 new file mode 100644 index 0000000000000000000000000000000000000000..2d41ad01ac483fc0e1a3be989bd6a9c50862383e GIT binary patch literal 3268 zcmai#e@vBS6vxkdy&$G**M^%mgqOr0#gz~#h@{=CtYx~SflHRHt0?0S+fbUT4dvBv zIwWk(O>E9FQrpDToPXTS{*i0wtO`BsuO;=OqI_G`PaSz9Uqwo7XpY#2m z=Xu}feeQz;03p02Uf|Am@BDGt!M?3M%Uaj3@x`P~$BSBuxu$_o8sxu+{k=t<+eY)C z+&=?vL%qp*Jx4Z=Gqd+2e{6gHnXK(xbn%X^%V4cNI{5i`n*H1_bDaxLUSh5Ow{t&et(djSBY*GFS`ll_Uwl%ewLI3o?YkZ)&0b&m zO;Mi4k7nK6q&;)Ep7OMwRU11uYme+H$#>Nhlu_Tsi&}~|P4g(>+%t*p(6xoy8)dIY zZ30SUjAP;9tD|}Sl3-!4c&l3)f2Y3vT}piuKxWm&+gclIA3eu)$5#&T*LCw)>p#&n zc|6TUFL@|do5tJ&5b!*Bf}rHH+8*+_c!~FQ?E7I7f#?K-p*7!6rIh`SON|5@8saOp zR>s=!&W=W{jkUj?=ep$4tAFV_IlGRQAFduxb8#H#~gU;a^ElHv1?=OFtvD(;6Do5Paa-@?X7r+^IK2W zja^G#Q_J53$GUWoYNym?`TNerzL^JL6$u!^9#4DuRKaq}xt8`S_9@uwxC9sqhkfsndSMCit#)95$5LzGt2zt%u>HW z@h-)O6epNvyrk0oy@_+n_!)fm2^TA_R9vUHg<0l1skle+zs&OfIl+#Q zh^NJ5K#tNcQ!A)u1hZ+yru^|oR_LT>+GxLwbj<6~9~k%*r~8unh55Ht?3s!MBUFFqKX7qXjY}% zI_{*k;;x>7@>c^Hzw72*6&UNh$M+nuAZcK((NeX#2I$uPKMqd;@Xbv>D%O(AV(QVdSQghCNy~cPKK+^~!)zQlrH(EyaI!*{Pr9%B$vPB}3cMjW>uI-YX&W zMkUZRGhA3pZYcP;XEw+2%d88ivxg<85X#Q7Zd*%ZCvVi^ztE&O%FP$a!4 zwXt?yv5#^aFKn2HvCilFEoA^&)}G#soO3NdVH^!fgzcv8-jH&>FdJI8eO&%Sx3^%; z!P`W=n^?@S+4XemIP2VhMbF{M046SLk!G>V_HP@Sj=BbE*rVOa9UIzLodIj2UH}I{ zkq;EFNKr@$lHwX_b4_>6k720uxn1*5Ed|#s98h*xN{U$%ufL-pKcnzSXh9H6ODQMZ zs6MR^L!9f@M=;ts_0KKHIq&lq%CHLyT4%&yBbhd(T7g>r*}68MUFl|!lmthjlQGKbxOd=%uKSVaz4Zr zuPP_Q0N*;JJZHirIeXoP5V|;*4Z$Yvuy}*@3Z_msjEG{~kP?8^<|rtri?*My?<%d8 zcS2$=Tw9lx0I&85!`IKzA+}s` zbEfFBc6Q#9#AImSeunBhHEwolNTAb_aV%r-ubJRcl#^BRqKb%Ncr|T;9CUrrYcwMr zB8}`HYV_cMrVo%@EPtLRlu{4`SkPhde#bj4ByX;6# zD#Yc*%bOEMg8W`&IHhW59QfDw?oP^(HUI0)Dy+kKFQak5VFtx!LWGa$X@(wODTOA3 zjzd+8n(J^%fHWSumpeECjp8Vf!-5p-rwQfJa;(tr8}J1gtKFA1V35~SIhlM^UFQ8# zo<*aeXW^J{Fw`WCuhW}?wuK);)nmqn25zrbmy>CumYT<|D?Rn$^G7aqZIvpPT1!fo z+|xKK{HQOTOy!1gYO*gjbqM5ZjpV@gxl66_)3QrK4VpHh6QDpf*F)^?w$wtvxN&ou zhhl1jG+sg(cgPstS%yw}ANWZTavV>k7ggwdC-9d?T|IZo!)r!&`pFov-#&O`A%<3x z#_wy6poR0mB!dcBHP7KLV^z}Fftl;qA1-RGSou@+jCj(vWI6{gnqljhbVLQId#kak zs48R3;g|RQws^W{=Lq-9gA_Xd?65!^rgA=WzcV*jge48@X(3H|9anN(bNEso$V+@IkFs=js8F(XxfV}J%C(*G@kM_cP`JSPS}Pr z-j~ZvXz;l_lZvZ% zB=Rm3b|?Het}K7}ZI&E>;^e{ZLE!J7RF>)qerXp0=SMfun%JCJ5dQqw){Cs=_uTAmA+1f!wkVt*N~^!$s0 z&jpS1^K1xttlz+T&VMHGylAc`q`$;}d-C`ItCeGZ<`*#CcE)pjhGAY0zG!GXzJ%fb z%Wn$vTN&nh-!fcc%R**??eP2+4D)! zVP21z|Aa7Kas4x&F5o-?2L)WkFt4Xm;Klqs0{^RkJ#M=m@?uj9G0fManqi*L60nkJ z$MJnD=C5Qt*K1&y>+NNDHp@SeoDpaZb_bmd^ZlVD*+eztO@_IinEx}z^ZdO6zQ{1o dA7o#6jz40U=LhoN5RL2Me?>ITHx|7+^KTGaF;)No literal 0 HcmV?d00001 diff --git a/ext/150/nassau_differential/c/132/1 b/ext/150/nassau_differential/c/132/1 new file mode 100644 index 0000000000000000000000000000000000000000..383932f51ceaaa57a6a3f6db1d0b131ee2dd7fae GIT binary patch literal 3508 zcmZ{me@s6)a{NJc*aV$JU=tFGxofH1f`p1GlGGt? zI9mo2~ab=Xtr;XXDx4^Ss~p^L^g) zo|pHYg8~3aqA8+b)S3w(3@iXuD#R!e>z8MyNj_Hn#iQ6)MU!KtvL+};+W}t6oqGp8 zuAp)09sF_lOU|AnU)IsF5zcamdG^hH#F<8{*>-LlvKY}YV?1dlaB0|SUsd|%sf$Z_ zOzWS{|8gge@O!*#La4aX9>v)nVqFm@cW}0s*e7L8Yd9+>c4_+$zaUdZ!;G=jOyC0z zd-==i?lnd_F7y^~?*@a&i~dj+buXm1rZ45Cze&@}c2bEXO@z1xrlA|*TW{3?{)<85Pn z-pv6Y_7c9sgD0Z?zGKF^gU@}oiQwV7jwf)QV77^v&HryBvfzC^NlYASw>AvA`@qu& zHQnM~yC0w}Ol)s*NE5GV3DmgoMs{cy+m$kGJs)b2q?-F@`5IcOChnV=8P47!*6=}B z3}+?8oI9_c;H-dH`HcnRyaYdUKCz9B(|%sJg;@5%28FXhVnqW3Zq9m$J$A0Un6qAD zHJO(uIXgmZZ+`RL$W+lVV^o?6@NI~pIaJ&5E(g3;yY)8rx(x?_PQ~tC$LBcIzw%b@ zQeG+%^RK*&{T`e1Df9;KRWX?k)&Fr8S)ORL9QcZuNScoQkwe4dQignLCu$FsY2)L-HhWLc(=lETP=ndX$4>`Wx;f>(zd1Bk={Rfbx zh=v*Cs+o+d;?P0ywddvOqbZQ2+b6@TGa@4s7vcr|C$ao{UOOJB)h&(+@t}3$47sgV zh)zAORQZbSrlBc0{x`m+i&VSQ)1z^AnwW3=S~F+$#J2R_#;v$x4_+lU^LvAnvsZ{+ zZ#k98*=}Ng#m*h$%ulRz|3Ede;JuqAwnP$VA}j-$Jvr0O=lzL_l*;Zb&g8ykxkhdv z3qIS6-zzIkU?fdkhmD`s;kdDGUQ;Wr=4(`_^|j#AhtP-NxKyxWsTw`R^oXlPQ#_rr zrt{Oy;n-J2lVkFGrHe5$dcaqE8a_a<$C*xOFhf9lU{U^r;(P=UB1kFk{=+ar&bONB zLelbiDw%E~E${ys(|3peat_>^a9gX^NX{1o>ivP mnJ#3ylC)e;ALA2DFOZh|6?>QYAX4Ws?IkVutMu8DU;hVf-5n(W literal 0 HcmV?d00001 diff --git a/ext/150/nassau_differential/c/133/0 b/ext/150/nassau_differential/c/133/0 new file mode 100644 index 0000000000000000000000000000000000000000..f0b3eeb08001179c692572156e045411809c54ac GIT binary patch literal 1576 zcmZQ%U|&Rj7K7v_05MdE;Wbnr zLn(z3CNw?$nYSdfWcCR$Dk&hi6y{5pv9P_fq}uN({Ym$GdoB_BGUp! zkW+x7aJNlCbDINDP64DO(SadBfuY4Dfs3JM3TuYPVYaRR9hnqV)p|Vd888Ukk%jpV zr~wEZfLH-&EQg~5m~G(D!N9@j;^4sr)EX` zM*|S!85m@s40PH8O*{n3p97^oKxts?fk6$_E*PH!8Xz!p(X1ckj)uT!2+%PEOut<- F0|5A%;eG%B literal 0 HcmV?d00001 diff --git a/ext/150/nassau_differential/c/133/1 b/ext/150/nassau_differential/c/133/1 new file mode 100644 index 0000000000000000000000000000000000000000..a41bc9e03050c77ece52e259a67136677ead9546 GIT binary patch literal 1404 zcmZQ%U|_HSVn!ea0U;Oz!UdB|NNgCBlLf+Buc3Syq(uNG0i?e}*>`?Tlmy8c05Qzi zzfd^_1t0^gslh=&K*2#k(4$GaDXqkiEkK1qNkJ|_qy|Ip9hhFQJ_itwh2a5100%>Z z83SKqA{WCsR*)d?Wt(7-eHuUvvv)a20t9#<6a#~U(l$_f4jTUk mly3s1c|eu{0gUee<-^R2LF4})1*0J_8Uh0r0^6U?Iu8I2xu*>P literal 0 HcmV?d00001 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/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 54baff8871..7cd1432b7d 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -1505,13 +1505,6 @@ fn multiply_batch_block( } .executes(|| { let client = CudaRuntime::client(&CudaDevice::default()); - // Hold the cross-runtime GPU read lock for this whole device section (H2D + kernels + - // readback), so a concurrent cooperative fp-cuda row-reduce (raw cudarc) waits for this - // cubecl multiply to finish before taking the GPU exclusively. Without this, the RREF's - // `cuLaunchCooperativeKernel` grid barrier can't co-reside its CTAs against a resident - // multiply kernel and spins forever (the intermittent stem-150 wedge). See - // [`fp::GPU_EXCLUSIVE`]. Dropped at the end of the closure, after `read_one`. - let _gpu_shared = fp::GPU_EXCLUSIVE.read().unwrap_or_else(|e| e.into_inner()); // Shared resident admissible buffers (see [`ResidentDev`]): re-upload the master ONLY // when this block dereferences past the uploaded prefix (`need_cs` / `need_mk`). The // master grows continually at the frontier, so re-uploading on mere growth ships diff --git a/ext/crates/fp/src/blas/cuda.rs b/ext/crates/fp/src/blas/cuda.rs index 550b20c99d..48f7c148e8 100644 --- a/ext/crates/fp/src/blas/cuda.rs +++ b/ext/crates/fp/src/blas/cuda.rs @@ -14,26 +14,12 @@ //! 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, RwLock}; +use std::sync::{Mutex, OnceLock}; use fp_cuda::GpuContext; use crate::{matrix::Matrix, prime::TWO}; -/// Serializes the fp-cuda **cooperative** row-reduce (raw cudarc) against the algebra -/// **cubecl** Milnor multiply, which share the one physical GPU from different CUDA -/// runtimes. The row-reduce's `cuLaunchCooperativeKernel` (panel_factor_coop's grid-wide -/// spin barrier) needs *all* its CTAs co-resident; a cubecl multiply kernel occupying SMs -/// at that moment prevents co-residency, so the missing CTAs never reach the barrier and -/// the resident ones spin forever (flat sm=100% — the intermittent stem-150 wedge). -/// -/// Contract: the cooperative row-reduce takes the **write** lock (exclusive GPU) around its -/// launch+download; every cubecl multiply takes a **read** lock held from launch through its -/// readback (so releasing means that kernel has actually completed, not merely enqueued). -/// Readers run concurrently with each other; a pending row-reduce drains them, runs alone, -/// then lets them resume. See `algebra::algebra::milnor_gpu` for the read side. -pub static GPU_EXCLUSIVE: RwLock<()> = RwLock::new(()); - /// Smallest `min(m, k, n)` for which we attempt the GPU matmul. Below this the /// host marshalling (bit-repack into TMA tiles + copies) costs more than it saves. const DEFAULT_THRESHOLD: usize = 2048; @@ -141,11 +127,9 @@ pub(crate) fn try_row_reduce(m: &mut Matrix) -> Option { let stride = cols.div_ceil(64); let limbs = to_limbs(m); - // Exclusive GPU for the cooperative reduce: block new cubecl multiplies and wait for - // in-flight ones to complete, so panel_factor_coop's grid-wide barrier can co-reside all - // its CTAs (see [`GPU_EXCLUSIVE`]). Held across launch + download; ordered before ctx.lock. + // The default row-reduce is composable (no cooperative launch), so it needs no + // cross-runtime exclusion against the concurrent cubecl multiply. let (dev_limbs, perm, r, pivot_cols) = { - let _exclusive = GPU_EXCLUSIVE.write().unwrap_or_else(|e| e.into_inner()); 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()?; diff --git a/ext/crates/fp/src/lib.rs b/ext/crates/fp/src/lib.rs index bb8dedec60..8d971da2a8 100644 --- a/ext/crates/fp/src/lib.rs +++ b/ext/crates/fp/src/lib.rs @@ -12,12 +12,6 @@ pub mod vector; pub mod blas; -/// Cross-runtime GPU serialization lock (cubecl multiply vs. cooperative fp-cuda row-reduce). -/// Re-exported so `algebra`'s Milnor-multiply GPU path can take the read side. See -/// [`blas::cuda::GPU_EXCLUSIVE`]. -#[cfg(feature = "gpu")] -pub use blas::cuda::GPU_EXCLUSIVE; - pub(crate) mod simd; // This is useful for traits that want to implement `Arbitrary`. This lets us specify that they From af74dfad957004f64545a94bc62570b157da5ee4 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Sun, 26 Jul 2026 20:03:44 -0400 Subject: [PATCH 022/127] milnor_gpu: fix >2^32-element master via cubecl 64-bit addressing The batched Milnor multiply silently corrupted results once the shared `masks` admissible-master crossed 2^32 u16 elements (~8.6 GB, around S_2 stem ~160-180): the buffer length and the per-R offsets were u32, so cubecl's default U32 `address_type` truncated them (`len as u32`), giving out-of-range reads and `dx != 0` (d^2 != 0) panics at e.g. (180,92). (160,82) stayed just under the ceiling and was clean. Fix uses cubecl 0.10's first-class 64-bit addressing rather than hand-rolled buffer splitting: - `multiply_batch_kernel` -> `#[cube(launch_unchecked, address_type = "u64")]`. Static u64 (not "dynamic"): dynamic picks u32 `usize` for small blocks and then narrows the u64 offset arrays on read (`usize::cast_from(u64)` under a u32 address type), corrupting results. `launch_unchecked` because checked mode emits `min(u64, u64)`, which NVRTC rejects as an ambiguous overload; every access is in-bounds by construction (uploaded `need_*` prefix + per-column `j` guards). - `RInfo.cs_off/mk_off` u32 -> u64; `r_cs_offset`/`r_mk_offset` bound as `Array`. - Resident grow copies (`copy_into_u16`/`_u32`) -> `#[cube(launch_unchecked, address_type = "dynamic")]` (scalar usize offsets adapt without narrowing; dynamic keeps small copies on u32). - Drop `RESIDENT_MAX_CAP` and its clamps (the u32 ceiling is gone); assert `out_len <= u32::MAX` loudly (the row-block splitter guarantees it). Validated on H200: (160,82) rc=0 clean, 517s vs 486s baseline (+6.4% for the static-u64 multiply, partly offset by unchecked dropping the bounds clamp); (180,92) ran 1500s with zero dx panics (previously a deterministic panic ~130-180s), i.e. correct well past the 2^32 masks boundary. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 83 ++++++++++++-------- 1 file changed, 51 insertions(+), 32 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 7cd1432b7d..ee2cb2a4a1 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -226,8 +226,11 @@ 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, @@ -289,14 +292,6 @@ static RESIDENT_DEV: LazyLock> = /// the initial size carries headroom to keep them few (~a handful per run). u16 → ~512 MiB. const RESIDENT_INIT_CAP: usize = 1 << 28; -/// Maximum resident buffer capacity (elements). cubecl's `ArrayArg` length is a `u32`, so an array -/// of exactly `2^32` elements has its length truncated to 0 — the copy then writes nothing and the -/// buffer reads as all zeros (the stem-150 `dx != 0`, hit when the doubling `cap` jumped 2^31 → 2^32). -/// Capacity is clamped just below the limit. A single resident buffer therefore cannot exceed -/// `2^32 - 1` elements; a master/basis larger than that (very high stems) must be split across -/// buffers — the same fundamental cubecl limit the old full-`create_from_slice` path also had. -const RESIDENT_MAX_CAP: usize = (1 << 32) - 1; - /// Serializes master device *uploads* only — never handle reads. A launch that must grow the /// device master takes this before uploading, so at a growth point at most one multi-GB /// `create_from_slice` runs (others re-check and find it already done) instead of every launch @@ -349,8 +344,10 @@ macro_rules! resident_dev_handle { let (handle, cap) = if old_handle.is_none() || host_len > cap { // Quiesce readers for the handle swap (see [`RESIDENT_REALLOC`]). let _realloc_w = RESIDENT_REALLOC.write().unwrap(); - let new_cap = host_len.max(cap * 2).max(RESIDENT_INIT_CAP) - .min(RESIDENT_MAX_CAP); + // No `u32` cap: the kernels bind these buffers with 64-bit addressing + // (`multiply_batch_kernel` reads with static `u64`; the grow copy with + // dynamic), so the length/offsets are never truncated past `u32::MAX`. + let new_cap = host_len.max(cap * 2).max(RESIDENT_INIT_CAP); let new_handle = $client.empty(new_cap * ::core::mem::size_of::()); if let Some(oh) = &old_handle { if uploaded > 0 { @@ -417,8 +414,8 @@ fn resident_info(algebra: &MilnorAlgebra, p_part: &[PPartEntry]) -> RInfo { return *info; } let info = RInfo { - cs_off: host.col_sums.len() as u32, - mk_off: host.masks.len() as u32, + cs_off: host.col_sums.len() as u64, + mk_off: host.masks.len() as u64, cs_len: cs_len as u32, mk_len: mk_len as u32, num_mats: (mk.len() / mk_len) as u32, @@ -549,8 +546,7 @@ macro_rules! basis_dev_handles { }; let (pp_h, pp_cap) = if old_pp.is_none() || pp_len > pp_cap { let _realloc_w = RESIDENT_REALLOC.write().unwrap(); - let new_cap = pp_len.max(pp_cap * 2).max(RESIDENT_INIT_CAP) - .min(RESIDENT_MAX_CAP); + let new_cap = pp_len.max(pp_cap * 2).max(RESIDENT_INIT_CAP); let nh = $client.empty(new_cap * ::core::mem::size_of::()); if let Some(oh) = &old_pp { if pp_up > 0 { @@ -584,8 +580,7 @@ macro_rules! basis_dev_handles { }; let (ln_h, ln_cap) = if old_ln.is_none() || ln_len > ln_cap { let _realloc_w = RESIDENT_REALLOC.write().unwrap(); - let new_cap = ln_len.max(ln_cap * 2).max(RESIDENT_INIT_CAP) - .min(RESIDENT_MAX_CAP); + let new_cap = ln_len.max(ln_cap * 2).max(RESIDENT_INIT_CAP); let nh = $client.empty(new_cap * ::core::mem::size_of::()); if let Some(oh) = &old_ln { if ln_up > 0 { @@ -654,10 +649,14 @@ fn zero_u32(out: &mut Array) { /// 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`, and `count` bounds this launch so the caller can split a copy larger than the -/// `u32` `ABSOLUTE_POS` thread limit into chunks (a resident buffer exceeds 2^32 u16 elements around -/// stem 150 — the width-padded basis p-parts especially). See [`copy_into_chunked`]. -#[cube(launch)] +/// 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: &Array, dst: &mut Array, 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]; @@ -665,7 +664,7 @@ fn copy_into_u16(src: &Array, dst: &mut Array, src_off: usize, dst_off } /// `u32` sibling of [`copy_into_u16`] (for the resident basis `lens`). -#[cube(launch)] +#[cube(launch_unchecked, address_type = "dynamic")] fn copy_into_u32(src: &Array, dst: &mut Array, 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]; @@ -687,10 +686,12 @@ macro_rules! copy_chunked { while done < $count { let n = ($count - done).min(COPY_CHUNK); unsafe { - $kernel::launch::( + $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)), ArrayArg::from_raw_parts($src.clone(), $src_len), ArrayArg::from_raw_parts($dst.clone(), $dst_len), $src_off + done, @@ -1006,7 +1007,15 @@ 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 masters (`col_sums`/`masks`) and the +// width-padded basis exceed `u32::MAX` elements at high stems, and the per-`R` offsets in +// `r_cs_offset`/`r_mk_offset` (u64) index into them. 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 already in-bounds by +// construction (the `need_*` prefix covers every offset and the per-column guards bound `j`). +#[cube(launch_unchecked, address_type = "u64")] #[allow(clippy::too_many_arguments)] fn multiply_batch_kernel( col_sums: &Array, @@ -1017,8 +1026,8 @@ fn multiply_batch_kernel( g: &Array, xi: &Array, out: &mut Array>, - r_cs_offset: &Array, - r_mk_offset: &Array, + r_cs_offset: &Array, + r_mk_offset: &Array, r_cs_len: &Array, r_mk_len: &Array, prod_r_index: &Array, @@ -1410,8 +1419,8 @@ fn multiply_batch_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_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()); @@ -1461,6 +1470,14 @@ fn multiply_batch_block( ); pps.push(total_pairs as u32); 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!( @@ -1538,8 +1555,8 @@ fn multiply_batch_block( let tg_h = client.create_from_slice(u32::as_bytes(&term_gei)); 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 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)); const THREADS: u32 = 256; @@ -1569,8 +1586,10 @@ fn multiply_batch_block( let poo_h = client.create_from_slice(u32::as_bytes(&prod_out_offset)); let pps_h = client.create_from_slice(u32::as_bytes(&pps)); let cubes = (total_pairs as u32).div_ceil(THREADS).max(1); + // SAFETY: `launch_unchecked` — see the kernel's `address_type = "u64"` note. Every device + // read is in-bounds by construction (uploaded `need_*` prefix + per-column `j` guards). unsafe { - multiply_batch_kernel::launch::( + multiply_batch_kernel::launch_unchecked::( &client, CubeCount::Static(cubes, 1, 1), CubeDim::new_1d(THREADS), From 59bd2d934e7d3e436b38cd8d088fc41c87c8800d Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Sun, 26 Jul 2026 22:45:55 -0400 Subject: [PATCH 023/127] =?UTF-8?q?milnor=5Fgpu:=20slash=20GPU=20host=20me?= =?UTF-8?q?mory=20=E2=80=94=20single=20stream=20by=20default=20+=20drop=20?= =?UTF-8?q?the=20resident=20host=20master?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `resolve_through_stem S_2 "" 180 92` was OOM-killed at the 500 GB cgroup limit (~280 GB anon + ~240 GB shmem). Measurement (env `NASSAU_MEM_REPORT`, plus a pure-CPU control at ~5-7 GB) showed the resolution's own data — differentials and module tables — is only ~1.5 GB; the ~500 GB was entirely GPU-runtime host memory. Two fixes cut it to ~90 GB at stem 180 (dx-clean): 1. Default `NASSAU_GPU_STREAMS` 8 -> 1. cubecl gives each CUDA stream its own page-locked (pinned) host pool that `memory_cleanup` never trims; ≥2 streams ballooned pinned host memory to 140-240 GB, single-stream holds it at ~4-6 GB. The payoff of multi-stream is ~nil here: cubecl's server is single-threaded, so extra streams buy no CPU concurrency, only GPU kernel overlap the big saturating multiplies barely use. Override for a dedicated large-RAM node. 2. Stop retaining the resident admissible master host-side. `RESIDENT_HOST` kept the full `col_sums`/`masks` (~54 GB at stem 180) forever, duplicating the device copy, even though after upload it is never read again (offsets come from `index`; growth uploads only the new tail; a capacity realloc copies the old *device* buffer). Now it keeps only the not-yet-uploaded tail (`*_pending`, ~sub-GB) plus a logical length, freeing each `R`'s data the moment it reaches the GPU — invariant `dev.uploaded == len - pending.len()`. Also: bound the pinned staging on resident growth to `STAGE_CHUNK` chunks, and a gated `NASSAU_MEM_REPORT` (differentials/modules/resident heap breakdown) for ongoing memory work. Validated on H200: stem 180 peak RSS ~90 GB (was 500 GB OOM), dx-clean through the >2^32 masks region and the heavy solves. Remaining growth for higher stems is real GPU working memory (concurrent dense output matrices) + GPU device memory, not retained host duplicates. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 169 ++++++++++++------ ext/crates/algebra/src/module/free_module.rs | 15 ++ .../homomorphism/free_module_homomorphism.rs | 23 +++ ext/src/nassau.rs | 35 ++++ 4 files changed, 192 insertions(+), 50 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index ee2cb2a4a1..5566a65c2a 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -111,22 +111,26 @@ static GPU_BUDGET: LazyLock = LazyLock::new(|| GpuBudget { freed: Condvar::new(), }); -/// Number of distinct CUDA streams to spread device work over (`NASSAU_GPU_STREAMS`, default 8). +/// Number of distinct CUDA streams to spread device work over (`NASSAU_GPU_STREAMS`, default 1). /// Each worker thread gets a stable per-thread stream id (see [`thread_stream_id`]), and this caps /// how many distinct streams — hence retained memory pools — exist; `1` forces the single-stream /// mode (all threads → stream 0). /// -/// Multi-stream is SAFE because the shared resident master/basis are grown IN PLACE in stable -/// device buffers ([`ResidentDev`]) — the read-only "global weights" pattern cubecl supports across -/// streams. (The earlier churny design — re-`create_from_slice` a new handle on every growth — broke -/// cubecl's per-handle cross-stream sync and had to run single-stream; that is fixed.) +/// **Default 1 (single stream) because multi-stream is a MEMORY disaster for little gain.** cubecl +/// gives each CUDA stream its own device AND page-locked host (pinned) memory pool, and those pools +/// are never trimmed (`memory_cleanup` frees only the GPU pool). Measured at stem 180: single-stream +/// holds pinned host memory at ~4 GB, but ≥2 streams balloons it to 140–240 GB (each stream retains +/// its own varying-size readback/staging pages) — the dominant term in the ~500 GB cgroup OOM. And +/// the payoff is ~nil: cubecl's server is single-threaded (one runner behind a channel), so extra +/// streams buy no CPU-side concurrency, only GPU kernel overlap that the big saturating multiplies +/// barely benefit from. Raise it only on a dedicated large-RAM node that wants that overlap. fn gpu_stream_slots() -> u64 { static SLOTS: LazyLock = LazyLock::new(|| { std::env::var("NASSAU_GPU_STREAMS") .ok() .and_then(|v| v.parse::().ok()) .filter(|&n| n > 0) - .unwrap_or(8) + .unwrap_or(1) }); *SLOTS } @@ -214,6 +218,23 @@ pub fn take_batch_stats() -> (u64, u64, u64, u64) { ) } +/// 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.capacity() * 2 + + h.mk_pending.capacity() * 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) +} + use std::{ collections::HashMap, sync::{Condvar, LazyLock, Mutex, RwLock}, @@ -250,10 +271,19 @@ struct RInfo { /// 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 [`resident_dev_handle`]: +/// `RESIDENT_DEV.$buf.uploaded == $len - $pending.len()`, i.e. `$pending == master[uploaded..$len]`. #[derive(Default)] struct ResidentHost { - col_sums: Vec, - masks: Vec, + cs_pending: Vec, + mk_pending: Vec, + cs_len: usize, + mk_len: usize, index: HashMap, RInfo>, } @@ -315,7 +345,7 @@ static RESIDENT_REALLOC: RwLock<()> = RwLock::new(()); /// `RESIDENT_UPLOAD`, and a re-check coalesces a burst of growth-needing launches into one /// upload). The master is append-only, so any handle with `uploaded >= need` is valid. macro_rules! resident_dev_handle { - ($client:expr, $need:expr, $buf:ident, $host_vec:ident) => {{ + ($client:expr, $need:expr, $buf:ident, $pending:ident, $len:ident) => {{ // Lock-free fast path: the stable handle already covers `$need`. let read_current = || { let dev = RESIDENT_DEV.read().unwrap(); @@ -335,7 +365,17 @@ macro_rules! resident_dev_handle { let dev = RESIDENT_DEV.read().unwrap(); (dev.$buf.handle.clone(), dev.$buf.cap, dev.$buf.uploaded) }; - let host_len = RESIDENT_HOST.read().unwrap().$host_vec.len(); + // Take the not-yet-uploaded tail and the logical length together. By the + // invariant (see [`ResidentHost`]) the tail is exactly `master[uploaded..len]`, + // so `uploaded + batch.len() == host_len`. `mem::take` FREES it host-side — + // once on the device it is never read again. Appends after this point land in + // a fresh `$pending` and upload on the next growth. `RESIDENT_UPLOAD` (held) + // makes this the sole grower, so `uploaded` is stable here. + let (batch, host_len) = { + let mut host = RESIDENT_HOST.write().unwrap(); + (std::mem::take(&mut host.$pending), host.$len) + }; + debug_assert_eq!(uploaded + batch.len(), host_len); // (Re)allocate a stable persistent buffer on first use or capacity overflow // (rare: `RESIDENT_INIT_CAP` has headroom and the master saturates early), @@ -363,27 +403,24 @@ macro_rules! resident_dev_handle { (old_handle.unwrap(), cap) }; - // Write the new tail host[uploaded..host_len] into the STABLE buffer at its - // append offset via a scratch upload + copy kernel, then SYNC before - // publishing: cubecl does not cross-stream-order a *kernel write* to a shared - // buffer against reads the way it does `create_from_slice`, so a reader on - // another worker's stream could otherwise see the tail mid-copy (harmless at - // small sizes, corrupting once the copy is GB-scale — the stem-150 `dx != 0`). - // Blocking here makes the new prefix physically resident before any reader can - // observe the bumped `uploaded`. - if host_len > uploaded { - let n = host_len - uploaded; - let scratch = { - let host = RESIDENT_HOST.read().unwrap(); - $client.create_from_slice(u16::as_bytes( - &host.$host_vec[uploaded..host_len], - )) - }; + // Write the tail `batch` (== master[uploaded..host_len]) into the STABLE + // buffer at its append offset through a BOUNDED pinned staging (see + // [`STAGE_CHUNK`]), syncing each chunk. The sync also serves the cross-stream + // ordering the resident buffers need: cubecl does not order a *kernel write* to + // a shared buffer against reads on another stream the way it does + // `create_from_slice`, so blocking here makes the new prefix physically + // resident before any reader observes the bumped `uploaded`. + let mut done = 0usize; + while done < batch.len() { + let m = (batch.len() - done).min(STAGE_CHUNK); + let lo = uploaded + done; + let scratch = + $client.create_from_slice(u16::as_bytes(&batch[done..done + m])); copy_chunked!( - $client, copy_into_u16, scratch, n, 0usize, - handle, cap, uploaded, n + $client, copy_into_u16, scratch, m, 0usize, handle, cap, lo, m ); let _ = cubecl_common::reader::read_sync($client.sync()); + done += m; } { @@ -413,15 +450,19 @@ fn resident_info(algebra: &MilnorAlgebra, p_part: &[PPartEntry]) -> RInfo { 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`]). let info = RInfo { - cs_off: host.col_sums.len() as u64, - mk_off: host.masks.len() as u64, + cs_off: host.cs_len as u64, + mk_off: host.mk_len as u64, cs_len: cs_len as u32, mk_len: mk_len as u32, num_mats: (mk.len() / mk_len) as u32, }; - host.col_sums.extend(cs.iter().map(|&v| narrow_u16(v))); - host.masks.extend(mk.iter().map(|&v| narrow_u16(v))); + host.cs_pending.extend(cs.iter().map(|&v| narrow_u16(v))); + host.mk_pending.extend(mk.iter().map(|&v| narrow_u16(v))); + host.cs_len += cs.len(); + host.mk_len += mk.len(); host.index.insert(p_part.to_vec(), info); info } @@ -563,12 +604,9 @@ macro_rules! basis_dev_handles { }; if pp_len > pp_up { let n = pp_len - pp_up; - let scratch = { - let host = RESIDENT_BASIS_HOST.read().unwrap(); - $client.create_from_slice(u16::as_bytes(&host.pparts[pp_up..pp_len])) - }; - copy_chunked!( - $client, copy_into_u16, scratch, n, 0usize, + stage_upload!( + $client, copy_into_u16, u16::as_bytes, + RESIDENT_BASIS_HOST.read().unwrap(), pparts, pp_h, pp_cap, pp_up, n ); } @@ -597,20 +635,17 @@ macro_rules! basis_dev_handles { }; if ln_len > ln_up { let n = ln_len - ln_up; - let scratch = { - let host = RESIDENT_BASIS_HOST.read().unwrap(); - $client.create_from_slice(u32::as_bytes(&host.lens[ln_up..ln_len])) - }; - copy_chunked!( - $client, copy_into_u32, scratch, n, 0usize, + stage_upload!( + $client, copy_into_u32, u32::as_bytes, + RESIDENT_BASIS_HOST.read().unwrap(), lens, ln_h, ln_cap, ln_up, n ); } - // Sync both in-place copies before publishing (see [`resident_dev_handle`]): - // a kernel write to the shared basis buffers must be physically done before a - // reader on another worker's stream can observe the bumped element count. - let _ = cubecl_common::reader::read_sync($client.sync()); + // `stage_upload!` already synced each chunk, so both grown buffers are + // physically resident before publishing (see [`resident_dev_handle`]): + // a reader on another worker's stream must never observe the bumped element + // count before the copy is done. { let mut dev = RESIDENT_BASIS_DEV.write().unwrap(); @@ -675,6 +710,40 @@ fn copy_into_u32(src: &Array, dst: &mut Array, src_off: usize, dst_off /// 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 [`stage_upload`]). +/// 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; + +/// Upload `host_slice[0..n]` into resident device `dst[dst_off..dst_off+n]` through a BOUNDED pinned +/// staging buffer, re-locking `$host_lock` to reslice each chunk. `$host_lock` is an expression +/// evaluating to a read guard whose `$field` is the source `Vec` (re-evaluated per chunk so the +/// guard is not held across the sync). Syncs after each chunk so the pinned pool reuses one page +/// instead of accumulating the whole tail (see [`STAGE_CHUNK`]). `$copy` is `copy_into_u16`/`_u32`, +/// `$as_bytes` the matching `u16`/`u32` `as_bytes`. +macro_rules! stage_upload { + ($client:expr, $copy:ident, $as_bytes:path, $host_lock:expr, $field:ident, + $dst:expr, $dst_cap:expr, $dst_off:expr, $n:expr) => {{ + let mut done = 0usize; + while done < $n { + let m = ($n - done).min(STAGE_CHUNK); + let lo = $dst_off + done; + let scratch = { + let host = $host_lock; + $client.create_from_slice($as_bytes(&host.$field[lo..lo + m])) + }; + copy_chunked!($client, $copy, scratch, m, 0usize, $dst, $dst_cap, lo, m); + // Drain the H2D copy so this chunk's pinned staging is freed (and its pool page reused) + // before the next `create_from_slice`, keeping live pinned memory at ~one chunk. + let _ = cubecl_common::reader::read_sync($client.sync()); + done += m; + } + }}; +} + /// 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). @@ -1532,8 +1601,8 @@ fn multiply_batch_block( // then HOST.read, and nothing under either lock blocks on rayon or a permit. The // master is append-only, so the uploaded prefix is always a prefix of the current // host master and every offset `< uploaded` is final. - let (cs_h, cs_len_master) = resident_dev_handle!(client, need_cs, cs, col_sums); - let (mk_h, mk_len_master) = resident_dev_handle!(client, need_mk, mk, masks); + let (cs_h, cs_len_master) = resident_dev_handle!(client, need_cs, cs, cs_pending, cs_len); + let (mk_h, mk_len_master) = resident_dev_handle!(client, need_mk, mk, mk_pending, mk_len); // 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. 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/src/nassau.rs b/ext/src/nassau.rs index ab7dc8fe1d..cf2fe33dcf 100644 --- a/ext/src/nassau.rs +++ b/ext/src/nassau.rs @@ -1187,6 +1187,10 @@ impl> Resolution { // 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. @@ -1218,6 +1222,37 @@ impl> Resolution { 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); + let gb = |x: usize| x as f64 / (1u64 << 30) as f64; + eprintln!( + "[MEM] commits={commit_count} last_b=({},{}) \ + differentials={:.1}GB modules={:.1}GB \ + resident_master={:.1}GB resident_basis={:.1}GB accounted={:.1}GB", + b.n(), + b.s(), + gb(diff_b), + gb(mod_b), + gb(res_master), + gb(res_basis), + gb(diff_b + mod_b + res_master + res_basis), + ); + } + } + // 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. From e4da7ede2bfcf855a0f414865b94f6a74967b44c Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Sun, 26 Jul 2026 23:36:17 -0400 Subject: [PATCH 024/127] nassau_gpu: cap the per-build GPU work so the dense readback stays bounded MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Even with the resident host-master freed, host RSS still grew ~linearly toward the cgroup limit at high stems — not a retained duplicate but the GPU path materializing whole dense output matrices. The CPU walks signatures one at a time (small working set); the GPU offload built the entire matrix at once and held its dense readback (`num_rows × num_limbs` u32, ~12 GB regions at stem 180) alongside the assembled matrix. Two caps bound it: 1. `reuse_full_matrix` now only builds the all-rows matrix when `rows × cols <= NASSAU_GPU_REUSE_MAX_WORK` (default 1e10). Above that it falls back to per-signature builds (each a bounded row subset, like the CPU), so the peak scales with the largest single signature, not the whole bidegree. 2. `get_partial_matrix_restricted` processes rows in batches sized so the dense readback stays under `NASSAU_GPU_MAX_READBACK_MB` (default 1024). Products are built in row order, so each batch is a contiguous slice with its `row` remapped batch-local; results XOR back into the global rows. The readback is freed between batches. Both are correctness-neutral (rows are independent): verified bit-for-bit vs the CPU under `NASSAU_GPU_VERIFY` with a 1 MB cap (many batches/build), 0 mismatches. Effect at stem 180 (streams=1 + resident de-dup + this): host anon goes from growing past 85 GB to a bounded ~18-24 GB, with no throughput loss — cross- bidegree rayon concurrency keeps the GPU fed, and the default cap only splits the few giant high-t builds into a handful of chunks. Peak RSS ~30 GB (was 500 GB OOM). This bounds the host working set so it scales toward higher stems; a D-deep async launch/collect pipeline (cubecl launches are async; only read_one blocks) can hide chunk latency if a much smaller cap is ever needed. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/src/nassau.rs | 36 +++++++++++++++++++++- ext/src/nassau_gpu.rs | 69 ++++++++++++++++++++++++++++++++++--------- 2 files changed, 90 insertions(+), 15 deletions(-) diff --git a/ext/src/nassau.rs b/ext/src/nassau.rs index cf2fe33dcf..1ca264ca6f 100644 --- a/ext/src/nassau.rs +++ b/ext/src/nassau.rs @@ -397,6 +397,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 @@ -718,7 +750,9 @@ impl> Resolution { // 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(); let _guard = ParallelGuard::new(); Some(restricted_partial_matrix_maybe_gpu( diff --git a/ext/src/nassau_gpu.rs b/ext/src/nassau_gpu.rs index b1315e7d9b..0bb33217d7 100644 --- a/ext/src/nassau_gpu.rs +++ b/ext/src/nassau_gpu.rs @@ -157,13 +157,31 @@ 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); + let (mut matrix, mut products) = extract_restricted(hom, degree, inputs, target_dim); if !products.is_empty() { let target = hom.target(); let algebra = target.algebra(); @@ -175,22 +193,45 @@ 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); + // 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()); + 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 rows = multiply_batch_on_gpu(&algebra, full_cols, r1 - r0, &products[p0..p1]); + for (bi, limbs) in rows.iter().enumerate() { + let mut target_row = matrix.row_mut(r0 + bi); + 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; + } } - bits &= bits - 1; } } + p0 = p1; + r0 = r1; } } matrix From 28a1dffbd85f8f0c3b7fbd5c519633abb2e42b13 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Mon, 27 Jul 2026 04:27:25 -0400 Subject: [PATCH 025/127] milnor_gpu: degree-threshold eviction of the resident device master MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The resident admissible master (col_sums/masks) grows with degree and is the stem-300 device-memory wall: ~34 GB at stem 180, extrapolating past the H200's 143 GB before stem 300. This adds an opt-in eviction that bounds it. Policy (from the NASSAU_R_STATS probe): a *degree* threshold, not LRU. 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 also the biggest matrices, so a degree cap evicts the most bytes for the fewest references and is stem-independent (the kept set saturates). Env NASSAU_GPU_RESIDENT_MAX_DEGREE (default i32::MAX = keep all). Design — no kernel change. Every output row's products share one operation R (extract_restricted), so hot (deg<=cap) and cold rows are DISJOINT. multiply_batch_on_gpu partitions products, compacts each group to its own dense row range (so total readback stays num_rows, not 2x), runs the resident pass on hot and a Transient pass on cold, and scatters results back. Transient builds a per-block master with create_from_slice, freed with the launch, so the device copy never persists. Cold admissible data is cached host-side (COLD_HOST) so the expensive recompute is one-shot, like the resident path. Fast path (cap==MAX or all-hot) is byte-identical to before — zero regression by default. GpuProduct gains Clone for the partition. Validated bit-exact: NASSAU_GPU_VERIFY full S_2 stem-110 resolution at theta=10 (transient path heavily stressed) and theta=100: mismatches=0 dx=0 panics=0. Memory vs throughput (stem 180, ExclusivePages, streams=1): control (no eviction) 51 GB / 1349 s; theta=125 11 GB; theta=100 12 GB (master 34 -> ~1 GB), dx=0. But eviction costs 2-4x throughput even at a high theta (stem 150: control 282 s, theta=140 613 s) because the evicted high-degree matrices are the biggest and are re-uploaded every launch. So this is a fit-in-memory lever for stems where the master won't fit at all (a slow completion beats an OOM), tuned to the highest theta that fits; it is NOT a speedup. Next optimization to cut the re-upload: a bounded LRU device-side cold cache (upload once, reuse across a bidegree's launches). Also retained: the R-access probe (NASSAU_R_STATS/dump_r_stats) and the device/host [MEM] report (NASSAU_MEM_REPORT) used to characterize this — both env-gated no-ops. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 365 ++++++++++++++++++- ext/src/nassau.rs | 25 +- 2 files changed, 367 insertions(+), 23 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 5566a65c2a..964d1efb96 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -235,9 +235,32 @@ pub fn resident_host_bytes() -> (usize, usize) { (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) { + let d = RESIDENT_DEV.read().unwrap(); + let master = d.cs.uploaded * 2 + d.mk.uploaded * 2; + let b = RESIDENT_BASIS_DEV.read().unwrap(); + let basis = b.pp.uploaded * 2 + b.ln.uploaded * 4; + (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 = CudaRuntime::client(&CudaDevice::default()); + match client.memory_usage() { + Ok(u) => (u.bytes_in_use, u.bytes_reserved), + Err(_) => (0, 0), + } +} + use std::{ collections::HashMap, - sync::{Condvar, LazyLock, Mutex, RwLock}, + sync::{Arc, Condvar, LazyLock, Mutex, RwLock}, }; use cubecl::server::Handle; @@ -339,6 +362,48 @@ static RESIDENT_UPLOAD: Mutex<()> = Mutex::new(()); /// hot path stays lock-free; reallocs happen only a handful of times per run, so the barrier is free. static RESIDENT_REALLOC: RwLock<()> = RwLock::new(()); +/// Host-side cache of cold (degree > [`resident_degree_cap`]) `R`s' admissible matrices, narrowed to +/// `u16` and ready to upload. Cold `R`s are deliberately kept OFF the device master (that is the whole +/// point of eviction — it bounds the 143 GB device), but `admissible_matrices` is the *expensive* part +/// for exactly these high-degree `R`s (big matrices) and they recur across many bidegrees. Recomputing +/// per launch collapsed throughput (2× wall at stem 180, bench 2026-07-27). Caching the result +/// host-side (host has ~755 GB — never the constraint) makes the recompute one-shot, like the resident +/// path, while the *device* copy stays transient (uploaded per launch into a block-local buffer, freed +/// after). Grows to roughly the evicted tail of the master (tens of GB), well within host RAM. +struct ColdEntry { + cs_len: u32, + mk_len: u32, + num_mats: u32, + cs: Arc>, + mk: Arc>, +} +static COLD_HOST: LazyLock, ColdEntry>>> = + LazyLock::new(|| RwLock::new(HashMap::new())); + +/// Cold-`R` admissible data, from the [`COLD_HOST`] cache (computed + narrowed once on first use). +/// Returns `(cs_len, mk_len, num_mats, cs, mk)`; the `Arc`s make the cache-hit path a cheap refcount +/// bump, no copy. Layout matches [`resident_info`]'s so the kernel indexes both identically. +fn cold_host_entry( + algebra: &MilnorAlgebra, + p_part: &[PPartEntry], +) -> (u32, u32, u32, Arc>, Arc>) { + if let Some(e) = COLD_HOST.read().unwrap().get(p_part) { + return (e.cs_len, e.mk_len, e.num_mats, e.cs.clone(), e.mk.clone()); + } + let (cs_len, mk_len, cs, mk) = algebra.admissible_matrices(p_part); + let num_mats = (mk.len() / mk_len) as u32; + let entry = ColdEntry { + cs_len: cs_len as u32, + mk_len: mk_len as u32, + num_mats, + cs: Arc::new(cs.iter().map(|&v| narrow_u16(v)).collect()), + mk: Arc::new(mk.iter().map(|&v| narrow_u16(v)).collect()), + }; + let mut w = COLD_HOST.write().unwrap(); + let e = w.entry(p_part.to_vec()).or_insert(entry); + (e.cs_len, e.mk_len, e.num_mats, e.cs.clone(), e.mk.clone()) +} + /// Fetch a resident device-master handle, uploading the current master prefix only when this /// block dereferences past what is already on the device (`need`). The upload runs OUTSIDE /// `RESIDENT_DEV` (readers stay lock-free; only concurrent uploaders serialize, on @@ -441,7 +506,162 @@ macro_rules! resident_dev_handle { /// [`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, +} + +static R_STATS: LazyLock, RStat>>>> = LazyLock::new(|| { + std::env::var_os("NASSAU_R_STATS").map(|_| Mutex::new(HashMap::new())) +}); + +/// Internal degree of `R` from its p-part: `Σ p_part[i] · deg(ξ_{i+1})`. +fn ppart_degree(p_part: &[PPartEntry]) -> 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). +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(p_part: &[PPartEntry]) { + if let Some(m) = R_STATS.as_ref() { + let now = BATCH_CALLS.load(Ordering::Relaxed); + let mut map = m.lock().unwrap(); + let e = map.entry(p_part.to_vec()).or_insert(RStat { + count: 0, + degree: ppart_degree(p_part), + first: now, + last: now, + }); + 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, + ); +} + fn resident_info(algebra: &MilnorAlgebra, p_part: &[PPartEntry]) -> RInfo { + record_r_use(p_part); if let Some(info) = RESIDENT_HOST.read().unwrap().index.get(p_part) { return *info; } @@ -1269,6 +1489,7 @@ 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, @@ -1293,23 +1514,87 @@ pub fn multiply_batch_on_gpu( num_cols: usize, num_rows: usize, products: &[GpuProduct], +) -> Vec> { + let cap = resident_degree_cap(); + // 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. + if cap == i32::MAX || products.iter().all(|p| p.r_degree <= cap) { + return multiply_batch_grouped(algebra, num_cols, num_rows, products, MasterMode::Resident); + } + // 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![vec![0u32; num_limbs]; num_rows]; + for mode in [MasterMode::Resident, MasterMode::Transient] { + let is_group = |d: i32| match mode { + MasterMode::Resident => d <= cap, + MasterMode::Transient => d > cap, + }; + // 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)) + .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)) + .map(|p| { + let mut q = p.clone(); + q.row = remap[&p.row]; + q + }) + .collect(); + let sub = multiply_batch_grouped(algebra, num_cols, rows.len(), &compact, mode); + for (i, &orig) in rows.iter().enumerate() { + for (a, b) in result[orig].iter_mut().zip(&sub[i]) { + *a ^= *b; + } + } + } + result +} + +fn multiply_batch_grouped( + algebra: &MilnorAlgebra, + 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), 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. + // 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). This pre-pass also warms the shared resident master, - // so every block's layout lookups below are read-lock cache hits. + // 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_HOST`] cache the same way (no per-block recompute). let prod_pairs: Vec = products .iter() .map(|prod| { let r = algebra.basis_element_from_index(prod.r_degree, prod.r_idx); - resident_info(algebra, &r.p_part).num_mats as usize * prod.term_indices.len() + let num_mats = match mode { + MasterMode::Resident => resident_info(algebra, &r.p_part).num_mats as usize, + MasterMode::Transient => cold_host_entry(algebra, &r.p_part).2 as usize, + }; + num_mats * prod.term_indices.len() }) .collect(); let mut result: Vec> = Vec::with_capacity(num_rows); @@ -1336,6 +1621,7 @@ pub fn multiply_batch_on_gpu( r0, r1 - r0, &products[p0..p1], + mode, )); (r0, p0) = (r1, p1); } @@ -1353,6 +1639,7 @@ fn multiply_batch_block( row_base: usize, num_rows: usize, products: &[GpuProduct], + mode: MasterMode, ) -> Vec> { let (width, g) = algebra.seqno_table_u32(); let mut xi: Vec = xi_degrees(algebra.prime()) @@ -1495,17 +1782,40 @@ fn multiply_batch_block( 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`: this block's own `col_sums`/`masks`, packed contiguously with block-local + // offsets (mirrors the resident master's layout, but built fresh here and freed after the + // launch instead of persisting). Empty / unused under `Resident`. + let mut cs_local: Vec = Vec::new(); + let mut mk_local: 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"); - 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); + 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, cs, mk) = cold_host_entry(algebra, &r.p_part); + r_cs_offset.push(cs_local.len() as u64); + r_mk_offset.push(mk_local.len() as u64); + r_cs_len.push(cs_len); + r_mk_len.push(mk_len); + r_num_matrices.push(num_mats as usize); + cs_local.extend_from_slice(&cs); + mk_local.extend_from_slice(&mk); + need_cs = cs_local.len(); + need_mk = mk_local.len(); + } + } } // Lay out per-product records + the pair-count prefix sum (sequential). Term data is already @@ -1601,8 +1911,23 @@ fn multiply_batch_block( // then HOST.read, and nothing under either lock blocks on rayon or a permit. The // master is append-only, so the uploaded prefix is always a prefix of the current // host master and every offset `< uploaded` is final. - let (cs_h, cs_len_master) = resident_dev_handle!(client, need_cs, cs, cs_pending, cs_len); - let (mk_h, mk_len_master) = resident_dev_handle!(client, need_mk, mk, mk_pending, mk_len); + // `Transient` (degree > cap `R`s): upload this block's freshly-built master and free it + // with the launch — no persistence, no cross-stream sharing, so no realloc guard below. + let (cs_h, cs_len_master, mk_h, mk_len_master) = match mode { + MasterMode::Resident => { + let (cs_h, cs_len_master) = + resident_dev_handle!(client, need_cs, cs, cs_pending, cs_len); + let (mk_h, mk_len_master) = + resident_dev_handle!(client, need_mk, mk, mk_pending, mk_len); + (cs_h, cs_len_master, mk_h, mk_len_master) + } + MasterMode::Transient => ( + client.create_from_slice(u16::as_bytes(&cs_local)), + cs_local.len(), + client.create_from_slice(u16::as_bytes(&mk_local)), + mk_local.len(), + ), + }; // 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. @@ -1633,8 +1958,10 @@ fn multiply_batch_block( // through the readback that syncs it (see [`RESIDENT_REALLOC`]). All resident growth for this // block is already done above; a concurrent capacity realloc on another thread waits here for // this multiply to finish, so the resident handle can never be swapped mid-kernel. Held to - // the end of the device section (dropped after `read_one`). - let _realloc_guard = RESIDENT_REALLOC.read().unwrap(); + // the end of the device section (dropped after `read_one`). `Transient` buffers are + // block-local (never reallocated by another thread), so they need no such guard. + let _realloc_guard = + (mode == MasterMode::Resident).then(|| RESIDENT_REALLOC.read().unwrap()); // 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. Same stream as the multiply below, so it is ordered before it. diff --git a/ext/src/nassau.rs b/ext/src/nassau.rs index 1ca264ca6f..dd54f2de67 100644 --- a/ext/src/nassau.rs +++ b/ext/src/nassau.rs @@ -1271,18 +1271,31 @@ impl> Resolution { 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=({},{}) \ - differentials={:.1}GB modules={:.1}GB \ - resident_master={:.1}GB resident_basis={:.1}GB accounted={:.1}GB", + "[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(diff_b + mod_b + res_master + res_basis), + gb(dev_master), + gb(dev_basis), + gbu(dev_pool_use), + gbu(dev_pool_res), ); } } @@ -1317,6 +1330,10 @@ impl> Resolution { } } }); + + // Eviction probe (`NASSAU_R_STATS`): dump the R-access distribution once the wavefront is done. + #[cfg(feature = "gpu")] + algebra::milnor_gpu::dump_r_stats(); } } From b447de53e18ce013d6ba130de803a621846e7c9b Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Mon, 27 Jul 2026 09:34:49 -0400 Subject: [PATCH 026/127] milnor_gpu: CPU reference for in-kernel admissible enumeration Foundational step toward generating admissible matrices (col_sums/masks) ON the GPU into transient scratch instead of storing/uploading the resident master -- the direction that eliminates both the stem-300 device-memory wall and the eviction re-upload cost (a given launch would enumerate its cold R's on-device, never uploading them). enumerate_admissible_ref reimplements AdmissibleMatrix::next 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 a `rejected` flag rather than breaking). `found` replaces the odometer's `return true`, `handled` replaces its `continue 'mid`. This validates the tricky restructuring on the CPU, where it is fast to debug, before the hard-to-debug #[cube] port -- which then becomes a mechanical transcription onto per-thread local Arrays (state is tiny: rows = |p_part|, cols <= 32). Test admissible_enum_ref_matches asserts bit-exact equivalence with admissible_matrices over every real R up to degree 60 (4155 R's, all match). Not yet ported to cubecl / not wired into the multiply -- next step. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 146 +++++++++++++++++++ 1 file changed, 146 insertions(+) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 964d1efb96..fac30c9f47 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -2044,10 +2044,156 @@ fn multiply_batch_block( result } +/// 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. +#[cfg(test)] +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) +} + #[cfg(test)] mod tests { use super::*; + /// 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; + + let p = ValidPrime::new(2); + let algebra = MilnorAlgebra::new(p, false); + let max_degree = 60; + 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 = algebra.basis_element_from_index(deg, idx).p_part.clone(); + if p_part.is_empty() { + continue; + } + let want = algebra.admissible_matrices(&p_part); + 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` /// dev shell, unsandboxed). From 0ae80f3c25cf9d17f0aea10596225ff3401fe9c4 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Mon, 27 Jul 2026 10:01:19 -0400 Subject: [PATCH 027/127] milnor_gpu: in-kernel admissible enumeration (CUDA-validated) Transcribe enumerate_admissible_ref into a cubecl #[cube] kernel that generates each R's col_sums/masks for every admissible matrix directly into device scratch -- the on-GPU replacement for the resident/uploaded master. One thread per distinct R, all state in fixed-size per-thread local Arrays; flag-based control flow (while ... && !found / handled) since the DSL has no break/continue/return. Backend-agnostic host driver (generic over Runtime) so the identical kernel can run on CUDA and the cpu backend. New test admissible_enum_gpu_matches runs it on the H200 and asserts bit-exact output (values + per-R counts) vs the CPU reference: 1055 R's, 30385 matrices, all matching. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 304 +++++++++++++++++++ 1 file changed, 304 insertions(+) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index fac30c9f47..10cf4ca246 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -35,6 +35,19 @@ use crate::algebra::{Algebra, MilnorAlgebra, combinatorics::xi_degrees}; /// `mk_len = rows + cols − 1 ≤ MAX_XI_TAU + ⌈log2⌉`; 32 covers every in-range case. const WORKING_CAP: usize = 32; +/// Per-thread local caps for the in-kernel admissible enumeration ([`enumerate_admissible_kernel`]). +/// Each `R` has `rows = |p_part| ≤ MAX_XI_TAU` and `cols ≤ WORKING_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. +#[cfg(test)] +const ENUM_ROW_CAP: usize = MAX_XI_TAU; +#[cfg(test)] +const ENUM_COL_CAP: usize = WORKING_CAP; +#[cfg(test)] +const ENUM_MATRIX_CAP: usize = ENUM_ROW_CAP * ENUM_COL_CAP; +#[cfg(test)] +const ENUM_MASK_CAP: usize = ENUM_ROW_CAP + ENUM_COL_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 @@ -2044,6 +2057,242 @@ fn multiply_batch_block( result } +/// 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). +#[cfg(test)] +#[cube(launch)] +#[allow(clippy::too_many_arguments)] +fn enumerate_admissible_kernel( + p_parts: &Array, + r_rows: &Array, + r_cols: &Array, + r_cs_out: &Array, + r_mk_out: &Array, + out_cs: &mut Array, + out_mk: &mut Array, + out_counts: &mut Array, + width: usize, + n_r: usize, +) { + 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. + let mut matrix = Array::::new(ENUM_MATRIX_CAP); + let mut totals = Array::::new(ENUM_ROW_CAP); + let mut col_sums = Array::::new(ENUM_COL_CAP); + let mut masks = Array::::new(ENUM_MASK_CAP); + for i in 0..ENUM_MATRIX_CAP { + matrix[i] = 0u32; + } + for i in 0..ENUM_ROW_CAP { + totals[i] = 0u32; + } + for i in 0..ENUM_COL_CAP { + col_sums[i] = 0u32; + } + for i in 0..ENUM_MASK_CAP { + masks[i] = 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]; + matrix[i * cols] = x; + masks[i] = 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. + let co = cs_base + mat * cs_len; + for j in 0..cs_len { + out_cs[co + j] = u16::cast_from(col_sums[j]); + } + let mo = mk_base + mat * mk_len; + for j in 0..mk_len { + out_mk[mo + j] = u16::cast_from(masks[j]); + } + 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; + totals[row] = matrix[row * cols]; + let mut col = 1usize; + while col < cols && !found { + p_to_the_j *= 2u32; + let mut handled = false; + if p_to_the_j <= totals[row] { + // 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 |= matrix[(row + col - c) * cols + c]; + c += 1; + } + let cur = matrix[row * cols + col]; + let new_entry = ((cur | d) + 1u32) & !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 = 1usize; + while j < col { + masks[row + j] &= !matrix[row * cols + j]; + col_sums[j - 1] -= matrix[row * cols + j]; + matrix[row * cols + j] = 0u32; + j += 1; + } + matrix[row * cols + col] = new_entry; + let mut i = 0usize; + while i < row { + matrix[i * cols] = totals[i]; + masks[i] = totals[i]; + let mut j2 = 1usize; + while j2 < cols { + if i + j2 > row { + masks[i + j2] &= !matrix[i * cols + j2]; + } + col_sums[j2 - 1] -= matrix[i * cols + j2]; + matrix[i * cols + j2] = 0u32; + j2 += 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; + } + + out_counts[ri] = u32::cast_from(mat); +} + +/// 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. +#[cfg(test)] +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![0u32; n_r]; + let mut r_mk_out = vec![0u32; n_r]; + let mut cs_total = 0u32; + let mut mk_total = 0u32; + 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 u32; + let mk_len = (rows + cols - 1) as u32; + 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] * cs_len; + mk_total += num_mats[i] * 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(u32::as_bytes(&r_cs_out)); + let rmo_h = client.create_from_slice(u32::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 = 64; + let cubes = (n_r as u32).div_ceil(THREADS); + unsafe { + enumerate_admissible_kernel::launch::( + &client, + CubeCount::Static(cubes, 1, 1), + CubeDim::new_1d(THREADS), + ArrayArg::from_raw_parts(pp_h, pp_flat.len()), + ArrayArg::from_raw_parts(rr_h, n_r), + ArrayArg::from_raw_parts(rc_h, n_r), + ArrayArg::from_raw_parts(rco_h, n_r), + ArrayArg::from_raw_parts(rmo_h, n_r), + ArrayArg::from_raw_parts(ocs_h.clone(), cs_cap), + ArrayArg::from_raw_parts(omk_h.clone(), mk_cap), + ArrayArg::from_raw_parts(cnt_h.clone(), n_r), + width, + n_r, + ); + } + let cs = u16::from_bytes(&client.read_one(ocs_h).unwrap()).to_vec(); + let mk = u16::from_bytes(&client.read_one(omk_h).unwrap()).to_vec(); + let counts = u32::from_bytes(&client.read_one(cnt_h).unwrap()).to_vec(); + (cs, mk, counts) +} + /// 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 @@ -2161,10 +2410,65 @@ fn enumerate_admissible_ref(p_part: &[u32]) -> (usize, usize, Vec, Vec (cs_len, mk_len, out_cs, out_mk) } +/// 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. +#[cfg(test)] +fn check_enum_backend(device: &Rt::Device, max_degree: i32) { + use fp::prime::ValidPrime; + + let p = ValidPrime::new(2); + let algebra = MilnorAlgebra::new(p, false); + algebra.compute_basis(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 deg in 1..=max_degree { + for idx in 0..algebra.dimension(deg) { + let pp = algebra.basis_element_from_index(deg, idx).p_part.clone(); + 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); + } + } + assert!(!p_parts.is_empty(), "no R's exercised"); + + 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 from the CPU reference"); + assert_eq!(got_cs, exp_cs, "device col_sums diverged from the CPU reference"); + assert_eq!(got_mk, exp_mk, "device masks diverged from the CPU reference"); + eprintln!( + "enum backend: {} R's, {} matrices bit-exact vs enumerate_admissible_ref", + p_parts.len(), + num_mats.iter().sum::() + ); +} + #[cfg(test)] mod tests { use super::*; + /// 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 40 — + /// validating the cubecl lowering of the flag-based enumeration (local arrays, bitops, u16 + /// stores) on the actual H200 target. Requires a live GPU + the CUDA toolkit env. + #[test] + fn admissible_enum_gpu_matches() { + check_enum_backend::(&CudaDevice::default(), 40); + } + /// 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. From e41ed2fe013ad3a6135bf53ecfaee05dd736d8d9 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Mon, 27 Jul 2026 12:17:47 -0400 Subject: [PATCH 028/127] milnor_gpu: enumerate cold masters on-device instead of uploading Wire enumerate_admissible_kernel into the Transient (evicted) cold path of multiply_batch_block. A cold R's col_sums/masks are now GENERATED on the GPU into transient scratch (one enumeration launch, ordered before the multiply on the same stream) rather than built on the host and uploaded via create_from_slice. Only the small p-parts + per-R dimensions upload; the scratch is freed with the launch. This kills the per-launch H2D master re-upload that made eviction a 2-4x slowdown (bench 2026-07-27), the whole point of the in-kernel-enumeration direction. Also replace the COLD_HOST full-array cache with COLD_COUNT, a 12-byte-per-R (cs_len, mk_len, num_mats) shape cache. The evicted tail of the master now lives neither on the device nor the host -- only its sizes, needed up front to lay out the scratch offsets and the pair-count prefix sum. num_mats is counted once per distinct R (admissible_matrices, arrays dropped) and memoized. De-gate enumerate_admissible_kernel + ENUM_* caps + the MAX_XI_TAU import for production. Validated bit-exact: the isolation test (1055 R's / 30385 matrices) still passes, and NASSAU_GPU_VERIFY full S_2 stem-110 resolutions at theta=10 (nearly all R's cold -> enumeration path stressed) and theta=100 both report mismatches=0 dx=0 panics=0 rc=0. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 179 ++++++++++++------- 1 file changed, 112 insertions(+), 67 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 10cf4ca246..93294d578d 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -24,9 +24,8 @@ use cubecl::{ prelude::*, }; use cubecl_common::stream_id::StreamId; -// 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}; @@ -39,13 +38,9 @@ const WORKING_CAP: usize = 32; /// Each `R` has `rows = |p_part| ≤ MAX_XI_TAU` and `cols ≤ WORKING_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. -#[cfg(test)] const ENUM_ROW_CAP: usize = MAX_XI_TAU; -#[cfg(test)] const ENUM_COL_CAP: usize = WORKING_CAP; -#[cfg(test)] const ENUM_MATRIX_CAP: usize = ENUM_ROW_CAP * ENUM_COL_CAP; -#[cfg(test)] const ENUM_MASK_CAP: usize = ENUM_ROW_CAP + ENUM_COL_CAP; /// Target `(product, matrix, term)` thread-pairs per GPU launch. The batch multiply indexes threads @@ -273,7 +268,7 @@ pub fn cubecl_device_usage() -> (u64, u64) { use std::{ collections::HashMap, - sync::{Arc, Condvar, LazyLock, Mutex, RwLock}, + sync::{Condvar, LazyLock, Mutex, RwLock}, }; use cubecl::server::Handle; @@ -375,46 +370,29 @@ static RESIDENT_UPLOAD: Mutex<()> = Mutex::new(()); /// hot path stays lock-free; reallocs happen only a handful of times per run, so the barrier is free. static RESIDENT_REALLOC: RwLock<()> = RwLock::new(()); -/// Host-side cache of cold (degree > [`resident_degree_cap`]) `R`s' admissible matrices, narrowed to -/// `u16` and ready to upload. Cold `R`s are deliberately kept OFF the device master (that is the whole -/// point of eviction — it bounds the 143 GB device), but `admissible_matrices` is the *expensive* part -/// for exactly these high-degree `R`s (big matrices) and they recur across many bidegrees. Recomputing -/// per launch collapsed throughput (2× wall at stem 180, bench 2026-07-27). Caching the result -/// host-side (host has ~755 GB — never the constraint) makes the recompute one-shot, like the resident -/// path, while the *device* copy stays transient (uploaded per launch into a block-local buffer, freed -/// after). Grows to roughly the evicted tail of the master (tens of GB), well within host RAM. -struct ColdEntry { - cs_len: u32, - mk_len: u32, - num_mats: u32, - cs: Arc>, - mk: Arc>, -} -static COLD_HOST: LazyLock, ColdEntry>>> = +/// 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, (u32, u32, u32)>>> = LazyLock::new(|| RwLock::new(HashMap::new())); -/// Cold-`R` admissible data, from the [`COLD_HOST`] cache (computed + narrowed once on first use). -/// Returns `(cs_len, mk_len, num_mats, cs, mk)`; the `Arc`s make the cache-hit path a cheap refcount -/// bump, no copy. Layout matches [`resident_info`]'s so the kernel indexes both identically. -fn cold_host_entry( - algebra: &MilnorAlgebra, - p_part: &[PPartEntry], -) -> (u32, u32, u32, Arc>, Arc>) { - if let Some(e) = COLD_HOST.read().unwrap().get(p_part) { - return (e.cs_len, e.mk_len, e.num_mats, e.cs.clone(), e.mk.clone()); - } - let (cs_len, mk_len, cs, mk) = algebra.admissible_matrices(p_part); - let num_mats = (mk.len() / mk_len) as u32; - let entry = ColdEntry { - cs_len: cs_len as u32, - mk_len: mk_len as u32, - num_mats, - cs: Arc::new(cs.iter().map(|&v| narrow_u16(v)).collect()), - mk: Arc::new(mk.iter().map(|&v| narrow_u16(v)).collect()), - }; - let mut w = COLD_HOST.write().unwrap(); - let e = w.entry(p_part.to_vec()).or_insert(entry); - (e.cs_len, e.mk_len, e.num_mats, e.cs.clone(), e.mk.clone()) +/// 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: &[PPartEntry]) -> (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.to_vec()).or_insert(e); + e } /// Fetch a resident device-master handle, uploading the current master prefix only when this @@ -1598,14 +1576,14 @@ fn multiply_batch_grouped( // 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_HOST`] cache the same way (no per-block recompute). + // `Transient` it warms the host-side [`COLD_COUNT`] shape cache the same way (no per-block recount). let prod_pairs: Vec = products .iter() .map(|prod| { let r = algebra.basis_element_from_index(prod.r_degree, prod.r_idx); let num_mats = match mode { MasterMode::Resident => resident_info(algebra, &r.p_part).num_mats as usize, - MasterMode::Transient => cold_host_entry(algebra, &r.p_part).2 as usize, + MasterMode::Transient => cold_count(algebra, &r.p_part).2 as usize, }; num_mats * prod.term_indices.len() }) @@ -1795,11 +1773,14 @@ fn multiply_batch_block( 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`: this block's own `col_sums`/`masks`, packed contiguously with block-local - // offsets (mirrors the resident master's layout, but built fresh here and freed after the - // launch instead of persisting). Empty / unused under `Resident`. - let mut cs_local: Vec = Vec::new(); - let mut mk_local: Vec = Vec::new(); + // `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"); @@ -1817,20 +1798,54 @@ fn multiply_batch_block( .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, cs, mk) = cold_host_entry(algebra, &r.p_part); - r_cs_offset.push(cs_local.len() as u64); - r_mk_offset.push(mk_local.len() as u64); + let (cs_len, mk_len, num_mats) = cold_count(algebra, &r.p_part); + r_cs_offset.push(need_cs as u64); + r_mk_offset.push(need_mk as u64); r_cs_len.push(cs_len); r_mk_len.push(mk_len); r_num_matrices.push(num_mats as usize); - cs_local.extend_from_slice(&cs); - mk_local.extend_from_slice(&mk); - need_cs = cs_local.len(); - need_mk = mk_local.len(); + // `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.clone()); + need_cs += num_mats as usize * cs_len as usize; + need_mk += num_mats as usize * mk_len as usize; } } } + // (Transient) Flatten the cold p-parts (padded to the widest) and cast the block-local scratch + // offsets to `u32` for the enumeration kernel. Offsets equal `r_cs_offset`/`r_mk_offset`; a block + // is pair-capped (`GPU_PAIR_CHUNK`) so the packed scratch stays well under `u32::MAX`. + let (enum_pp, enum_cs_out, enum_mk_out, enum_width) = if mode == MasterMode::Transient { + let w = enum_rows.iter().copied().max().unwrap_or(1) as usize; + let mut pp = vec![0u32; enum_pp_rows.len() * w]; + for (i, row) in enum_pp_rows.iter().enumerate() { + for (slot, &v) in pp[i * w..i * w + row.len()].iter_mut().zip(row) { + *slot = v; + } + } + let to_u32 = |v: &[u64]| -> Vec { + v.iter() + .map(|&x| { + assert!(x <= u32::MAX as u64, "transient scratch offset {x} exceeds u32"); + x as u32 + }) + .collect() + }; + (pp, to_u32(&r_cs_offset), to_u32(&r_mk_offset), w) + } else { + (Vec::new(), Vec::new(), Vec::new(), 1usize) + }; + // 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. @@ -1934,12 +1949,43 @@ fn multiply_batch_block( resident_dev_handle!(client, need_mk, mk, mk_pending, mk_len); (cs_h, cs_len_master, mk_h, mk_len_master) } - MasterMode::Transient => ( - client.create_from_slice(u16::as_bytes(&cs_local)), - cs_local.len(), - client.create_from_slice(u16::as_bytes(&mk_local)), - mk_local.len(), - ), + MasterMode::Transient => { + // Generate this block's cold `col_sums`/`masks` ON the device into transient scratch + // (freed with the launch), instead of uploading host-built arrays. Only the small + // p-parts + dimensions are uploaded. The enumeration launch is issued before the + // multiply on this same stream, so the scratch is fully written when the multiply + // reads it (kernel launches on one stream are ordered, as with `zero_u32` below). + const ENUM_THREADS: u32 = 256; + let n_cold = enum_rows.len(); + let cs_cap = need_cs.max(1); + let mk_cap = need_mk.max(1); + let cs_scratch = client.empty(cs_cap * size_of::()); + let mk_scratch = client.empty(mk_cap * size_of::()); + let cnt_scratch = client.empty(n_cold.max(1) * size_of::()); + let epp_h = client.create_from_slice(u32::as_bytes(&enum_pp)); + let er_h = client.create_from_slice(u32::as_bytes(&enum_rows)); + let ec_h = client.create_from_slice(u32::as_bytes(&enum_cols)); + let eco_h = client.create_from_slice(u32::as_bytes(&enum_cs_out)); + let emo_h = client.create_from_slice(u32::as_bytes(&enum_mk_out)); + unsafe { + enumerate_admissible_kernel::launch::( + &client, + CubeCount::Static((n_cold as u32).div_ceil(ENUM_THREADS).max(1), 1, 1), + CubeDim::new_1d(ENUM_THREADS), + ArrayArg::from_raw_parts(epp_h, enum_pp.len()), + ArrayArg::from_raw_parts(er_h, n_cold), + ArrayArg::from_raw_parts(ec_h, n_cold), + ArrayArg::from_raw_parts(eco_h, n_cold), + ArrayArg::from_raw_parts(emo_h, n_cold), + ArrayArg::from_raw_parts(cs_scratch.clone(), cs_cap), + ArrayArg::from_raw_parts(mk_scratch.clone(), mk_cap), + ArrayArg::from_raw_parts(cnt_scratch, n_cold.max(1)), + enum_width, + n_cold, + ); + } + (cs_scratch, cs_cap, mk_scratch, mk_cap) + } }; // 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 @@ -2073,7 +2119,6 @@ fn multiply_batch_block( /// 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). -#[cfg(test)] #[cube(launch)] #[allow(clippy::too_many_arguments)] fn enumerate_admissible_kernel( From a28e2146312a1d7ff7138b8cdf256aec619c0b22 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Mon, 27 Jul 2026 13:41:57 -0400 Subject: [PATCH 029/127] milnor_gpu: enum-path diagnostics + CPU-vs-GPU enum bench (direction refuted) Investigation of the eviction crash + a throughput bench that together show in-kernel enumeration cannot beat upload-based eviction: - enumerate_admissible_kernel -> u64 addressing (was default u32); did NOT fix the big-block CUDA_ERROR_LAUNCH_FAILED, so the fault is elsewhere in the Transient wiring, not the kernel. - admissible_enum_gpu_matches extended to degree 145, chunked per-degree: proves the kernel bit-exact vs the CPU reference (144903 R's, 185M matrices) across the full range the eviction path exercises -- so the kernel is correct. - bench_admissible_cpu_vs_gpu (ignored): CPU admissible_matrices 1.61s vs GPU kernel-only 2.02s (0.8x, SLOWER) for degrees 1-130. GPU enumeration is ~3x slower than just transferring the same arrays (0.68s readback) it replaces: the odometer is sequential per R with matrices-per-R spanning 1..millions, so the launch bottlenecks on its few longest threads at GPU scalar speed. The "trade GPU integer work for PCIe bandwidth" premise is refuted -- enumeration worsens the eviction re-upload cost instead of curing it. The enum wiring in multiply_batch_block's Transient path still faults at scale and should be reverted to upload-based eviction; kept here as validated, dormant code documenting the dead end. Default (theta=inf) path is unaffected. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 275 +++++++++++++++---- 1 file changed, 223 insertions(+), 52 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 93294d578d..c01cd25cd4 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -1822,10 +1822,10 @@ fn multiply_batch_block( } } - // (Transient) Flatten the cold p-parts (padded to the widest) and cast the block-local scratch - // offsets to `u32` for the enumeration kernel. Offsets equal `r_cs_offset`/`r_mk_offset`; a block - // is pair-capped (`GPU_PAIR_CHUNK`) so the packed scratch stays well under `u32::MAX`. - let (enum_pp, enum_cs_out, enum_mk_out, enum_width) = if mode == MasterMode::Transient { + // (Transient) Flatten the cold p-parts (padded to the widest) for the enumeration kernel. The + // per-`R` scratch offsets it writes at are `r_cs_offset`/`r_mk_offset` themselves (u64), passed + // straight through — no u32 narrowing, so a big block's multi-GB scratch is addressed safely. + let (enum_pp, enum_width) = if mode == MasterMode::Transient { let w = enum_rows.iter().copied().max().unwrap_or(1) as usize; let mut pp = vec![0u32; enum_pp_rows.len() * w]; for (i, row) in enum_pp_rows.iter().enumerate() { @@ -1833,17 +1833,9 @@ fn multiply_batch_block( *slot = v; } } - let to_u32 = |v: &[u64]| -> Vec { - v.iter() - .map(|&x| { - assert!(x <= u32::MAX as u64, "transient scratch offset {x} exceeds u32"); - x as u32 - }) - .collect() - }; - (pp, to_u32(&r_cs_offset), to_u32(&r_mk_offset), w) + (pp, w) } else { - (Vec::new(), Vec::new(), Vec::new(), 1usize) + (Vec::new(), 1usize) }; // Lay out per-product records + the pair-count prefix sum (sequential). Term data is already @@ -1965,10 +1957,10 @@ fn multiply_batch_block( let epp_h = client.create_from_slice(u32::as_bytes(&enum_pp)); let er_h = client.create_from_slice(u32::as_bytes(&enum_rows)); let ec_h = client.create_from_slice(u32::as_bytes(&enum_cols)); - let eco_h = client.create_from_slice(u32::as_bytes(&enum_cs_out)); - let emo_h = client.create_from_slice(u32::as_bytes(&enum_mk_out)); + let eco_h = client.create_from_slice(u64::as_bytes(&r_cs_offset)); + let emo_h = client.create_from_slice(u64::as_bytes(&r_mk_offset)); unsafe { - enumerate_admissible_kernel::launch::( + enumerate_admissible_kernel::launch_unchecked::( &client, CubeCount::Static((n_cold as u32).div_ceil(ENUM_THREADS).max(1), 1, 1), CubeDim::new_1d(ENUM_THREADS), @@ -2119,14 +2111,23 @@ fn multiply_batch_block( /// 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). -#[cube(launch)] +// +// 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: &Array, r_rows: &Array, r_cols: &Array, - r_cs_out: &Array, - r_mk_out: &Array, + // 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: &Array, + r_mk_out: &Array, out_cs: &mut Array, out_mk: &mut Array, out_counts: &mut Array, @@ -2276,10 +2277,10 @@ fn enumerate_admissible_on_runtime( 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![0u32; n_r]; - let mut r_mk_out = vec![0u32; n_r]; - let mut cs_total = 0u32; - let mut mk_total = 0u32; + 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 @@ -2287,8 +2288,8 @@ fn enumerate_admissible_on_runtime( .map(|&x| (u32::BITS - x.leading_zeros()) as usize) .max() .unwrap(); - let cs_len = (cols - 1) as u32; - let mk_len = (rows + cols - 1) as u32; + 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; } @@ -2296,16 +2297,16 @@ fn enumerate_admissible_on_runtime( r_cols[i] = cols as u32; r_cs_out[i] = cs_total; r_mk_out[i] = mk_total; - cs_total += num_mats[i] * cs_len; - mk_total += num_mats[i] * mk_len; + 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(u32::as_bytes(&r_cs_out)); - let rmo_h = client.create_from_slice(u32::as_bytes(&r_mk_out)); + 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; @@ -2316,7 +2317,7 @@ fn enumerate_admissible_on_runtime( const THREADS: u32 = 64; let cubes = (n_r as u32).div_ceil(THREADS); unsafe { - enumerate_admissible_kernel::launch::( + enumerate_admissible_kernel::launch_unchecked::( &client, CubeCount::Static(cubes, 1, 1), CubeDim::new_1d(THREADS), @@ -2332,8 +2333,12 @@ fn enumerate_admissible_on_runtime( n_r, ); } - let cs = u16::from_bytes(&client.read_one(ocs_h).unwrap()).to_vec(); - let mk = u16::from_bytes(&client.read_one(omk_h).unwrap()).to_vec(); + // 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) } @@ -2468,11 +2473,16 @@ fn check_enum_backend(device: &Rt::Device, max_degree: i32) { let algebra = MilnorAlgebra::new(p, false); algebra.compute_basis(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(); + // 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 = algebra.basis_element_from_index(deg, idx).p_part.clone(); if pp.is_empty() { @@ -2486,32 +2496,191 @@ fn check_enum_backend(device: &Rt::Device, max_degree: i32) { exp_mk.extend(mk.iter().map(|&v| narrow_u16(v))); p_parts.push(pp); } - } - assert!(!p_parts.is_empty(), "no R's exercised"); - - 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 from the CPU reference"); - assert_eq!(got_cs, exp_cs, "device col_sums diverged from the CPU reference"); - assert_eq!(got_mk, exp_mk, "device masks diverged from the CPU reference"); + 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: {} R's, {} matrices bit-exact vs enumerate_admissible_ref", - p_parts.len(), - num_mats.iter().sum::() + "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. +#[cfg(test)] +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 = 64; + 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), + ArrayArg::from_raw_parts(pp_h, pp_flat.len()), + ArrayArg::from_raw_parts(rr_h, n_r), + ArrayArg::from_raw_parts(rc_h, n_r), + ArrayArg::from_raw_parts(rco_h, n_r), + ArrayArg::from_raw_parts(rmo_h, n_r), + ArrayArg::from_raw_parts(ocs_h.clone(), cs_cap), + ArrayArg::from_raw_parts(omk_h.clone(), mk_cap), + ArrayArg::from_raw_parts(cnt_h.clone(), n_r), + width, + n_r, + ); + } + // 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) +} + #[cfg(test)] mod tests { use super::*; /// 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 40 — + /// 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) on the actual H200 target. Requires a live GPU + the CUDA toolkit env. + /// 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. #[test] fn admissible_enum_gpu_matches() { - check_enum_backend::(&CudaDevice::default(), 40); + check_enum_backend::(&CudaDevice::default(), 145); + } + + /// 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 fp::prime::ValidPrime; + use std::time::Instant; + + let p = ValidPrime::new(2); + let algebra = 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). + let mut by_degree: 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 = Vec::new(); + let mut nms = Vec::new(); + for idx in 0..algebra.dimension(deg) { + let pp = algebra.basis_element_from_index(deg, idx).p_part.clone(); + 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(&pp); + cpu_secs += t.elapsed().as_secs_f64(); + let nm = (mk.len() / mk_len) as u32; + nms.push(nm); + total_mats += nm as u64; + pps.push(pp); + } + total_r += pps.len(); + if !pps.is_empty() { + by_degree.push((pps, nms)); + } + } + + let device = CudaDevice::default(); + // 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_marshal, mut g_kernel, mut g_read) = (0.0f64, 0.0f64, 0.0f64); + for (pps, nms) in &by_degree { + let (m, k, r) = enum_launch_timed::(&device, pps, nms); + g_marshal += m; + g_kernel += k; + g_read += r; + } + + eprintln!( + "\n=== admissible enumeration: CPU vs GPU (degrees 1..={max_degree}) ===\n\ + R's: {total_r} matrices: {total_mats}\n\ + CPU admissible_matrices : {cpu_secs:.3} s\n\ + GPU kernel only : {g_kernel:.3} s ({:.1}x vs CPU) [production path]\n\ + GPU marshal+upload : {g_marshal:.3} s\n\ + GPU full readback : {g_read:.3} s (bench-only; prod keeps it on-device)\n\ + GPU kernel+marshal+read : {:.3} s ({:.1}x vs CPU) [full round-trip]", + cpu_secs / g_kernel, + g_marshal + g_kernel + g_read, + cpu_secs / (g_marshal + g_kernel + g_read), + ); } /// The flag-based [`enumerate_admissible_ref`] must reproduce `admissible_matrices` bit-for-bit @@ -2524,7 +2693,9 @@ mod tests { let p = ValidPrime::new(2); let algebra = MilnorAlgebra::new(p, false); - let max_degree = 60; + // 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 { From cc6ae7ba591b9347f55c6e55dd60a8c2d1b7079e Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Mon, 27 Jul 2026 14:44:37 -0400 Subject: [PATCH 030/127] milnor_gpu: correct the enum-vs-upload bench (parity, not 3x slower) The earlier bench compared GPU enum kernel-only to the D->H READBACK (0.68s) as a stand-in for the upload it replaces, and wrongly concluded enum is ~3x slower. Production uses the matrices on-device (no readback), and the readback is not the upload. Measure the actual H->D create_from_slice upload of the same host-built arrays instead, synced via a tiny throwaway readback (upload-only): CPU admissible_matrices (1 core) : 1.60 s GPU enumerate in-kernel : 2.02 s H->D upload of host-built arrays : 1.85 s -> in-kernel enum is 1.09x the upload (parity), NOT 3x So in-kernel enumeration is throughput-neutral vs uploading AND saves the host COLD_HOST array cache (tens of GB at stem 300). The direction is viable, not refuted; the big-block wiring crash is worth fixing. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 55 ++++++++++++-------- 1 file changed, 34 insertions(+), 21 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index c01cd25cd4..e10767db05 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -2628,13 +2628,16 @@ mod tests { algebra.compute_basis(max_degree); // Gather every non-empty R, grouped by degree (per-degree GPU launches keep host arrays bounded). - let mut by_degree: Vec<(Vec>, Vec)> = Vec::new(); + // 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 = Vec::new(); - let mut nms = Vec::new(); + 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 = algebra.basis_element_from_index(deg, idx).p_part.clone(); if pp.is_empty() { @@ -2642,44 +2645,54 @@ mod tests { } // 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(&pp); + let (_cs_len, mk_len, cs, mk) = algebra.admissible_matrices(&pp); 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)); + 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 (pps, nms, _, _) = &by_degree[0]; let _ = enum_launch_timed::(&device, pps, nms); } - let (mut g_marshal, mut g_kernel, mut g_read) = (0.0f64, 0.0f64, 0.0f64); - for (pps, nms) in &by_degree { - let (m, k, r) = enum_launch_timed::(&device, pps, nms); - g_marshal += m; + let mut g_kernel = 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). + let (_m, k, _r) = enum_launch_timed::(&device, pps, nms); g_kernel += k; - g_read += r; + // (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 enumeration: CPU vs GPU (degrees 1..={max_degree}) ===\n\ - R's: {total_r} matrices: {total_mats}\n\ - CPU admissible_matrices : {cpu_secs:.3} s\n\ - GPU kernel only : {g_kernel:.3} s ({:.1}x vs CPU) [production path]\n\ - GPU marshal+upload : {g_marshal:.3} s\n\ - GPU full readback : {g_read:.3} s (bench-only; prod keeps it on-device)\n\ - GPU kernel+marshal+read : {:.3} s ({:.1}x vs CPU) [full round-trip]", - cpu_secs / g_kernel, - g_marshal + g_kernel + g_read, - cpu_secs / (g_marshal + g_kernel + g_read), + "\n=== admissible matrices onto the device: enumerate vs upload (degrees 1..={max_degree}) ===\n\ + R's: {total_r} matrices: {total_mats} (both paths leave the arrays ON-DEVICE, no readback)\n\ + CPU admissible_matrices (enumerate, 1 core) : {cpu_secs:.3} s\n\ + GPU enumerate in-kernel : {g_kernel:.3} s\n\ + H->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 / upload_secs, ); } From b6475583bd9a518fcedf6e1cccbbaa2939851bcd Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Mon, 27 Jul 2026 17:45:08 -0400 Subject: [PATCH 031/127] milnor_gpu: segmented-master read primitive (chunked-growth de-risk) Prototype step 1 toward no-copy master growth, to kill the realloc-doubling transient that pushes cubecl into its memory-corruption regime (the ~2x spike at a 32->64 GiB master grow = ~96 GiB live, the root cause of the stem-140+ dx=nonzero cubecl ManagedMemoryDescriptor corruption). A segmented master keeps old segments and appends a new fixed-size segment on growth -- never copies -- so device peak is live_size + one segment, flat. The kernel selects a segment by static branch (cubecl has no array-of-buffers): offset o -> segment o/seg_elems, local o%seg_elems (seg_read_u16, MASTER_MAX_SEG separate Array args). seg_gather_kernel + seg_read_matches_contiguous validate the mechanic bit-exact (8 segments, scrambled indices) before touching the multiply hot path. Not yet wired into multiply_batch_kernel / resident_dev_handle -- next step. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 160 +++++++++++++++++++ 1 file changed, 160 insertions(+) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index e10767db05..5e704dbca0 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -2095,6 +2095,96 @@ fn multiply_batch_block( result } +/// Max number of fixed-size segments a segmented resident master can have. The kernel selects a +/// segment by a static branch (cubecl cannot dynamically index an array-of-buffers), so this is a +/// compile-time bound; `MASTER_SEG_ELEMS` sets the per-segment element count. 8 segments is the +/// prototype size (raise once the mechanic is wired + validated). Together they cap a buffer at +/// `MASTER_MAX_SEG * MASTER_SEG_ELEMS` u16 elements. +#[allow(dead_code)] // used once the segmented master is wired into the multiply path +const MASTER_MAX_SEG: usize = 8; + +/// 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. +#[cube] +#[allow(clippy::too_many_arguments)] +fn seg_read_u16( + s0: &Array, + s1: &Array, + s2: &Array, + s3: &Array, + s4: &Array, + s5: &Array, + s6: &Array, + s7: &Array, + o: usize, + seg_elems: usize, +) -> u16 { + let seg = o / seg_elems; + let local = o % seg_elems; + let mut v = 0u16; + if seg == 0 { + v = s0[local]; + } else if seg == 1 { + v = s1[local]; + } else if seg == 2 { + v = s2[local]; + } else if seg == 3 { + v = s3[local]; + } else if seg == 4 { + v = s4[local]; + } else if seg == 5 { + v = s5[local]; + } else if seg == 6 { + v = s6[local]; + } else { + v = s7[local]; + } + v +} + +/// 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 before the mechanic is wired into the +/// multiply's hot path. +#[cfg(test)] +#[cube(launch)] +#[allow(clippy::too_many_arguments)] +fn seg_gather_kernel( + s0: &Array, + s1: &Array, + s2: &Array, + s3: &Array, + s4: &Array, + s5: &Array, + s6: &Array, + s7: &Array, + idx: &Array, + out: &mut Array, + seg_elems: usize, +) { + let i = ABSOLUTE_POS; + if i >= out.len() { + terminate!(); + } + out[i] = seg_read_u16( + s0, + s1, + s2, + s3, + s4, + s5, + s6, + s7, + usize::cast_from(idx[i]), + seg_elems, + ); +} + /// 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 @@ -2460,6 +2550,59 @@ fn enumerate_admissible_ref(p_part: &[u32]) -> (usize, usize, Vec, Vec (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. +#[cfg(test)] +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 = CudaRuntime::client(&CudaDevice::default()); + + // 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 { ArrayArg::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), + ArrayArg::from_raw_parts(idx_h, indices.len()), + ArrayArg::from_raw_parts(out_h.clone(), indices.len()), + seg_elems, + ); + } + 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 @@ -2612,6 +2755,23 @@ mod tests { 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 From f362f85bec505dc55114f56bfa4c91ad0405093b Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Mon, 27 Jul 2026 18:33:21 -0400 Subject: [PATCH 032/127] milnor_gpu: segmented no-copy resident master/basis (kill the realloc-doubling spike) Replace the doubling-realloc-with-copy growth of every resident device buffer with a single segmented, append-only, no-copy mechanism. Growth now allocates only the new fixed-size segment(s) it needs and stage-writes the tail into them; existing segments are never reallocated or copied, so the device peak is `live + one_segment` instead of the ~2x transient (old+new buffer both live) that pushed cubecl into its silent memory-corruption regime -- the stem-140+ `dx != 0`. One mechanism for all four resident buffers (master cs/mk, basis pp/ln): - SegBuf + seg_grow! replace GrowBuf, resident_dev_handle!, basis_dev_handles!, stage_upload!, RESIDENT_REALLOC, RESIDENT_INIT_CAP (all deleted -- no two coexisting growth paths). - multiply_batch_kernel binds each store as MASTER_MAX_SEG (=16) segment Arrays and GATHERS a thread's matrix cs/mk + its term p-part into small WORKING_CAP locals via seg_read_u16/seg_read_u32 (correct at any offset, so no layout padding), then calls the UNCHANGED pure multiply_pair with base 0. Only this kernel changed; multiply_pair and the test kernel are untouched. - No realloc barrier: segments never change identity or get freed, so a reader's cloned handles stay valid across a concurrent append. - master_seg_elems() (env NASSAU_GPU_MASTER_SEG_ELEMS, default 1<<31, < u32 so a segment length never truncates cubecl's 32-bit metadata) => 64 GiB per buffer over 16 segments. Over-cap is a clean assert, not corruption. Validated bit-exact on H200: seg_read_matches_contiguous (16-arg primitive), multiply_batch_matches_reference across NASSAU_GPU_MASTER_SEG_ELEMS 768..4096 (multi-segment single-launch gather), and a new multiply_batch_incremental_growth at seg_elems=8192 (cross-launch append into the partially-filled last segment). Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 881 +++++++++++-------- 1 file changed, 536 insertions(+), 345 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 5e704dbca0..a2155601f9 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -307,7 +307,7 @@ struct RInfo { /// 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 [`resident_dev_handle`]: +/// (~27 GB at stem 130, growing). Invariant maintained by [`seg_grow`]: /// `RESIDENT_DEV.$buf.uploaded == $len - $pending.len()`, i.e. `$pending == master[uploaded..$len]`. #[derive(Default)] struct ResidentHost { @@ -321,55 +321,68 @@ struct ResidentHost { static RESIDENT_HOST: LazyLock> = LazyLock::new(|| RwLock::new(ResidentHost::default())); -/// A resident device buffer that GROWS IN PLACE: a single stable handle, allocated once -/// (persistent, so per-launch cleanup never reindexes it) with headroom and written to at its -/// append offset as the host master grows — never re-`create_from_slice`d. The handle changes only -/// on a rare capacity doubling. `uploaded` is how many elements have been written; `cap` is the -/// allocated element capacity. +/// 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 +} + +/// 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 GrowBuf { - handle: Option, - cap: usize, +struct SegBuf { + segs: Vec, uploaded: usize, } -/// Process-shared device mirror of the host master. Each buffer is a STABLE handle grown in place -/// (see [`GrowBuf`], [`resident_dev_handle`]). 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 fixed handle written in +/// 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: GrowBuf, - mk: GrowBuf, + cs: SegBuf, + mk: SegBuf, } static RESIDENT_DEV: LazyLock> = LazyLock::new(|| RwLock::new(ResidentDev::default())); -/// Initial element capacity of a resident buffer. Doubling to grow past it is barrier-protected -/// (see [`RESIDENT_REALLOC`]) so it is correct, but each doubling briefly quiesces the device, so -/// the initial size carries headroom to keep them few (~a handful per run). u16 → ~512 MiB. -const RESIDENT_INIT_CAP: usize = 1 << 28; - -/// Serializes master device *uploads* only — never handle reads. A launch that must grow the -/// device master takes this before uploading, so at a growth point at most one multi-GB -/// `create_from_slice` 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). +/// 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: Mutex<()> = Mutex::new(()); -/// Guards the rare capacity-doubling REALLOC (the only time a resident handle changes) against -/// concurrent readers. A device section holds the read lock while its multiply kernel is reading the -/// resident buffers; a realloc takes the write lock, so it waits for every in-flight reader to drain -/// and blocks new ones — the swap + old-buffer free then happens with the device quiesced, which the -/// churny path could not guarantee (residual ~1/10 crash at the readback / a stale multiply). The -/// common in-place delta write does NOT take this lock (it's stable-handle, proven race-free), so the -/// hot path stays lock-free; reallocs happen only a handful of times per run, so the barrier is free. -static RESIDENT_REALLOC: RwLock<()> = RwLock::new(()); - /// 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 @@ -395,97 +408,92 @@ fn cold_count(algebra: &MilnorAlgebra, p_part: &[PPartEntry]) -> (u32, u32, u32) e } -/// Fetch a resident device-master handle, uploading the current master prefix only when this -/// block dereferences past what is already on the device (`need`). The upload runs OUTSIDE -/// `RESIDENT_DEV` (readers stay lock-free; only concurrent uploaders serialize, on -/// `RESIDENT_UPLOAD`, and a re-check coalesces a burst of growth-needing launches into one -/// upload). The master is append-only, so any handle with `uploaded >= need` is valid. -macro_rules! resident_dev_handle { - ($client:expr, $need:expr, $buf:ident, $pending:ident, $len:ident) => {{ - // Lock-free fast path: the stable handle already covers `$need`. +/// 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 = RESIDENT_DEV.read().unwrap(); - match (dev.$buf.uploaded >= $need, dev.$buf.handle.clone()) { - (true, Some(h)) => Some((h, dev.$buf.uploaded)), - _ => None, + let dev = $dev.read().unwrap(); + if dev.$field.uploaded >= $need { + Some((dev.$field.segs.clone(), dev.$field.uploaded)) + } else { + None } }; match read_current() { - Some(hu) => hu, + Some(su) => su, None => { - let _upload_guard = RESIDENT_UPLOAD.lock().unwrap(); + let _upload_guard = $upload.lock().unwrap(); match read_current() { - Some(hu) => hu, // another grower already covered our need + Some(su) => su, // another grower already covered our need None => { - let (old_handle, cap, uploaded) = { - let dev = RESIDENT_DEV.read().unwrap(); - (dev.$buf.handle.clone(), dev.$buf.cap, dev.$buf.uploaded) - }; - // Take the not-yet-uploaded tail and the logical length together. By the - // invariant (see [`ResidentHost`]) the tail is exactly `master[uploaded..len]`, - // so `uploaded + batch.len() == host_len`. `mem::take` FREES it host-side — - // once on the device it is never read again. Appends after this point land in - // a fresh `$pending` and upload on the next growth. `RESIDENT_UPLOAD` (held) - // makes this the sole grower, so `uploaded` is stable here. - let (batch, host_len) = { - let mut host = RESIDENT_HOST.write().unwrap(); - (std::mem::take(&mut host.$pending), host.$len) + // 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) }; - debug_assert_eq!(uploaded + batch.len(), host_len); - - // (Re)allocate a stable persistent buffer on first use or capacity overflow - // (rare: `RESIDENT_INIT_CAP` has headroom and the master saturates early), - // copying existing device data on-device. Only THIS handle-swap needs a sync, - // so a fresh-handle cross-stream read never sees an incomplete copy. - let (handle, cap) = if old_handle.is_none() || host_len > cap { - // Quiesce readers for the handle swap (see [`RESIDENT_REALLOC`]). - let _realloc_w = RESIDENT_REALLOC.write().unwrap(); - // No `u32` cap: the kernels bind these buffers with 64-bit addressing - // (`multiply_batch_kernel` reads with static `u64`; the grow copy with - // dynamic), so the length/offsets are never truncated past `u32::MAX`. - let new_cap = host_len.max(cap * 2).max(RESIDENT_INIT_CAP); - let new_handle = $client.empty(new_cap * ::core::mem::size_of::()); - if let Some(oh) = &old_handle { - if uploaded > 0 { - copy_chunked!( - $client, copy_into_u16, oh, uploaded, 0usize, - new_handle, new_cap, 0usize, uploaded - ); - } - } - let _ = cubecl_common::reader::read_sync($client.sync()); - (new_handle, new_cap) - } else { - (old_handle.unwrap(), cap) - }; - - // Write the tail `batch` (== master[uploaded..host_len]) into the STABLE - // buffer at its append offset through a BOUNDED pinned staging (see - // [`STAGE_CHUNK`]), syncing each chunk. The sync also serves the cross-stream - // ordering the resident buffers need: cubecl does not order a *kernel write* to - // a shared buffer against reads on another stream the way it does - // `create_from_slice`, so blocking here makes the new prefix physically - // resident before any reader observes the bumped `uploaded`. + let (tail, new_len): (Vec<$elem>, usize) = ($tail)(uploaded); + debug_assert_eq!(uploaded + tail.len(), new_len); + assert!( + new_len.div_ceil(seg_elems.max(1)) <= MASTER_MAX_SEG, + "resident buffer needs {} segments (> MASTER_MAX_SEG={}); raise \ + MASTER_MAX_SEG or NASSAU_GPU_MASTER_SEG_ELEMS", + new_len.div_ceil(seg_elems.max(1)), + MASTER_MAX_SEG + ); + // 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 done < batch.len() { - let m = (batch.len() - done).min(STAGE_CHUNK); - let lo = uploaded + done; + 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(u16::as_bytes(&batch[done..done + m])); + $client.create_from_slice($as_bytes(&tail[done..done + m])); copy_chunked!( - $client, copy_into_u16, scratch, m, 0usize, handle, cap, lo, m + $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 = RESIDENT_DEV.write().unwrap(); - dev.$buf.handle = Some(handle.clone()); - dev.$buf.cap = cap; - dev.$buf.uploaded = host_len; + let mut dev = $dev.write().unwrap(); + dev.$field.segs = segs.clone(); + dev.$field.uploaded = new_len; } - (handle, host_len) + (segs, new_len) } } } @@ -703,13 +711,13 @@ struct ResidentBasisHost { static RESIDENT_BASIS_HOST: LazyLock> = LazyLock::new(|| RwLock::new(ResidentBasisHost::default())); -/// Device mirror of [`ResidentBasisHost`], both buffers grown IN PLACE (see [`GrowBuf`], -/// [`resident_dev_handle`]). `pp` holds the width-padded p-parts (`elems * width` u16), `ln` the +/// 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: GrowBuf, - ln: GrowBuf, + pp: SegBuf, + ln: SegBuf, } static RESIDENT_BASIS_DEV: LazyLock> = @@ -762,118 +770,6 @@ fn ensure_basis(algebra: &MilnorAlgebra, width: usize, max_degree: i32) -> Vec= $need` is valid. A macro (not a fn) so the cubecl client type stays inferred, -/// exactly as [`resident_dev_handle`]. -macro_rules! basis_dev_handles { - ($client:expr, $need:expr) => {{ - // Lock-free fast path: the stable buffers already cover `$need` basis elements - // (`ln.uploaded` is the element count). - let read_current = || { - let dev = RESIDENT_BASIS_DEV.read().unwrap(); - match (dev.ln.uploaded >= $need, dev.pp.handle.clone(), dev.ln.handle.clone()) { - (true, Some(pp), Some(ln)) => Some((pp, ln)), - _ => None, - } - }; - match read_current() { - Some(h) => h, - None => { - let _guard = RESIDENT_BASIS_UPLOAD.lock().unwrap(); - match read_current() { - Some(h) => h, // another grower already covered our need - None => { - let width = RESIDENT_BASIS_HOST.read().unwrap().width; - let elems = RESIDENT_BASIS_HOST.read().unwrap().lens.len(); - let pp_len = elems * width; // u16 count - let ln_len = elems; // u32 count - - // Grow the width-padded p-parts buffer (u16) in place. - let (old_pp, pp_cap, pp_up) = { - let d = RESIDENT_BASIS_DEV.read().unwrap(); - (d.pp.handle.clone(), d.pp.cap, d.pp.uploaded) - }; - let (pp_h, pp_cap) = if old_pp.is_none() || pp_len > pp_cap { - let _realloc_w = RESIDENT_REALLOC.write().unwrap(); - let new_cap = pp_len.max(pp_cap * 2).max(RESIDENT_INIT_CAP); - let nh = $client.empty(new_cap * ::core::mem::size_of::()); - if let Some(oh) = &old_pp { - if pp_up > 0 { - copy_chunked!( - $client, copy_into_u16, oh, pp_up, 0usize, - nh, new_cap, 0usize, pp_up - ); - } - } - let _ = cubecl_common::reader::read_sync($client.sync()); - (nh, new_cap) - } else { - (old_pp.unwrap(), pp_cap) - }; - if pp_len > pp_up { - let n = pp_len - pp_up; - stage_upload!( - $client, copy_into_u16, u16::as_bytes, - RESIDENT_BASIS_HOST.read().unwrap(), pparts, - pp_h, pp_cap, pp_up, n - ); - } - - // Grow the lengths buffer (u32) in place. - let (old_ln, ln_cap, ln_up) = { - let d = RESIDENT_BASIS_DEV.read().unwrap(); - (d.ln.handle.clone(), d.ln.cap, d.ln.uploaded) - }; - let (ln_h, ln_cap) = if old_ln.is_none() || ln_len > ln_cap { - let _realloc_w = RESIDENT_REALLOC.write().unwrap(); - let new_cap = ln_len.max(ln_cap * 2).max(RESIDENT_INIT_CAP); - let nh = $client.empty(new_cap * ::core::mem::size_of::()); - if let Some(oh) = &old_ln { - if ln_up > 0 { - copy_chunked!( - $client, copy_into_u32, oh, ln_up, 0usize, - nh, new_cap, 0usize, ln_up - ); - } - } - let _ = cubecl_common::reader::read_sync($client.sync()); - (nh, new_cap) - } else { - (old_ln.unwrap(), ln_cap) - }; - if ln_len > ln_up { - let n = ln_len - ln_up; - stage_upload!( - $client, copy_into_u32, u32::as_bytes, - RESIDENT_BASIS_HOST.read().unwrap(), lens, - ln_h, ln_cap, ln_up, n - ); - } - - // `stage_upload!` already synced each chunk, so both grown buffers are - // physically resident before publishing (see [`resident_dev_handle`]): - // a reader on another worker's stream must never observe the bumped element - // count before the copy is done. - - { - let mut dev = RESIDENT_BASIS_DEV.write().unwrap(); - dev.pp.handle = Some(pp_h.clone()); - dev.pp.cap = pp_cap; - dev.pp.uploaded = pp_len; - dev.ln.handle = Some(ln_h.clone()); - dev.ln.cap = ln_cap; - dev.ln.uploaded = ln_len; - } - (pp_h, ln_h) - } - } - } - } - }}; -} /// Zero a device `u32` buffer on-device: `out[i] = 0`, one thread per limb. /// @@ -921,7 +817,7 @@ fn copy_into_u32(src: &Array, dst: &mut Array, src_off: usize, dst_off /// 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 [`stage_upload`]). +/// 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 @@ -929,31 +825,6 @@ const COPY_CHUNK: usize = 1 << 30; /// caps the live pinned staging at ~one chunk. 64 Mi × u16 = 128 MiB (× u32 = 256 MiB). const STAGE_CHUNK: usize = 1 << 26; -/// Upload `host_slice[0..n]` into resident device `dst[dst_off..dst_off+n]` through a BOUNDED pinned -/// staging buffer, re-locking `$host_lock` to reslice each chunk. `$host_lock` is an expression -/// evaluating to a read guard whose `$field` is the source `Vec` (re-evaluated per chunk so the -/// guard is not held across the sync). Syncs after each chunk so the pinned pool reuses one page -/// instead of accumulating the whole tail (see [`STAGE_CHUNK`]). `$copy` is `copy_into_u16`/`_u32`, -/// `$as_bytes` the matching `u16`/`u32` `as_bytes`. -macro_rules! stage_upload { - ($client:expr, $copy:ident, $as_bytes:path, $host_lock:expr, $field:ident, - $dst:expr, $dst_cap:expr, $dst_off:expr, $n:expr) => {{ - let mut done = 0usize; - while done < $n { - let m = ($n - done).min(STAGE_CHUNK); - let lo = $dst_off + done; - let scratch = { - let host = $host_lock; - $client.create_from_slice($as_bytes(&host.$field[lo..lo + m])) - }; - copy_chunked!($client, $copy, scratch, m, 0usize, $dst, $dst_cap, lo, m); - // Drain the H2D copy so this chunk's pinned staging is freed (and its pool page reused) - // before the next `create_from_slice`, keeping live pinned memory at ~one chunk. - let _ = cubecl_common::reader::read_sync($client.sync()); - done += m; - } - }}; -} /// 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. @@ -1287,21 +1158,40 @@ fn multiply_single_r_kernel( /// /// Output is `num_rows` F₂ vectors of `num_limbs` `u32` limbs, row `r` at /// `out[r*num_limbs ..]`. -// 64-bit addressing (`address_type = "u64"`): the admissible masters (`col_sums`/`masks`) and the -// width-padded basis exceed `u32::MAX` elements at high stems, and the per-`R` offsets in -// `r_cs_offset`/`r_mk_offset` (u64) index into them. 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 already in-bounds by -// construction (the `need_*` prefix covers every offset and the per-column guards bound `j`). +// 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 GATHERS its +// one matrix's `col_sums`/`masks` and its term's p-part out of the segments into small local arrays +// via `seg_read_*` (correct for any offset, straddle or not — no layout padding needed), then hands +// those contiguous locals to `multiply_pair`, which stays a pure segmentation-agnostic arithmetic +// core. `seg_elems` is the segment element count (`o / seg_elems` picks the segment). #[cube(launch_unchecked, address_type = "u64")] #[allow(clippy::too_many_arguments)] fn multiply_batch_kernel( - col_sums: &Array, - masks: &Array, - basis_pparts: &Array, - basis_lens: &Array, + cs0: &Array, cs1: &Array, cs2: &Array, cs3: &Array, + cs4: &Array, cs5: &Array, cs6: &Array, cs7: &Array, + cs8: &Array, cs9: &Array, cs10: &Array, cs11: &Array, + cs12: &Array, cs13: &Array, cs14: &Array, cs15: &Array, + mk0: &Array, mk1: &Array, mk2: &Array, mk3: &Array, + mk4: &Array, mk5: &Array, mk6: &Array, mk7: &Array, + mk8: &Array, mk9: &Array, mk10: &Array, mk11: &Array, + mk12: &Array, mk13: &Array, mk14: &Array, mk15: &Array, + pp0: &Array, pp1: &Array, pp2: &Array, pp3: &Array, + pp4: &Array, pp5: &Array, pp6: &Array, pp7: &Array, + pp8: &Array, pp9: &Array, pp10: &Array, pp11: &Array, + pp12: &Array, pp13: &Array, pp14: &Array, pp15: &Array, + ln0: &Array, ln1: &Array, ln2: &Array, ln3: &Array, + ln4: &Array, ln5: &Array, ln6: &Array, ln7: &Array, + ln8: &Array, ln9: &Array, ln10: &Array, ln11: &Array, + ln12: &Array, ln13: &Array, ln14: &Array, ln15: &Array, term_gei: &Array, g: &Array, xi: &Array, @@ -1317,6 +1207,7 @@ fn multiply_batch_kernel( prod_out_offset: &Array, prod_pair_start: &Array, width: usize, + seg_elems: usize, ) { let k = ABSOLUTE_POS; let num_products = prod_pair_start.len() - 1; @@ -1350,23 +1241,70 @@ fn multiply_batch_kernel( 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; - // `term_gei[term_slot]` is the term's *global* basis-element index (across all degrees): - // its (width-padded) p-part lives at `basis_pparts[gei*width ..]`, length `basis_lens[gei]`. - // The basis is resident on the device (uploaded once, grown incrementally), so a launch - // uploads only these indices instead of re-gathering every term's p-part — see - // [`ResidentBasisHost`]. + // `term_gei[term_slot]` is the term's *global* basis-element index (across all degrees): its + // (width-padded) p-part lives at global offset `gei*width` in the segmented basis `pp*`, length + // `ln*[gei]`. The basis is resident on the device (uploaded once, grown incrementally), so a + // launch uploads only these indices instead of re-gathering every term's p-part. let gei = usize::cast_from(term_gei[term_slot]); + + // Global (across-segment) offsets of this matrix's data and this term's p-part. + let cs_off = usize::cast_from(r_cs_offset[ri]) + m * cs_len; + let mk_off = usize::cast_from(r_mk_offset[ri]) + m * mk_len; + let pp_off = gei * width; + let term_len = usize::cast_from(seg_read_u32( + ln0, ln1, ln2, ln3, ln4, ln5, ln6, ln7, + ln8, ln9, ln10, ln11, ln12, ln13, ln14, ln15, + gei, seg_elems, + )); + + // Gather this thread's matrix / term out of the segmented stores into contiguous locals, then run + // the pure arithmetic core on them (base 0). Entries past each length are zero, matching + // `multiply_pair`'s own out-of-range convention. The loop bounds at `WORKING_CAP`, exactly as the + // core does, so any `mk_len > WORKING_CAP` tail (never read by the core) is likewise not gathered. + let mut cs_local = Array::::new(WORKING_CAP); + let mut mk_local = Array::::new(WORKING_CAP); + let mut term_local = Array::::new(WORKING_CAP); + for j in 0..WORKING_CAP { + let mut c = 0u16; + if j < cs_len { + c = seg_read_u16( + cs0, cs1, cs2, cs3, cs4, cs5, cs6, cs7, + cs8, cs9, cs10, cs11, cs12, cs13, cs14, cs15, + cs_off + j, seg_elems, + ); + } + cs_local[j] = c; + let mut mm = 0u16; + if j < mk_len { + mm = seg_read_u16( + mk0, mk1, mk2, mk3, mk4, mk5, mk6, mk7, + mk8, mk9, mk10, mk11, mk12, mk13, mk14, mk15, + mk_off + j, seg_elems, + ); + } + mk_local[j] = mm; + let mut b = 0u16; + if j < term_len { + b = seg_read_u16( + pp0, pp1, pp2, pp3, pp4, pp5, pp6, pp7, + pp8, pp9, pp10, pp11, pp12, pp13, pp14, pp15, + pp_off + j, seg_elems, + ); + } + term_local[j] = b; + } + multiply_pair( - col_sums, - masks, - basis_pparts, + &cs_local, + &mk_local, + &term_local, g, xi, out, - usize::cast_from(r_cs_offset[ri]) + m * cs_len, - usize::cast_from(r_mk_offset[ri]) + m * mk_len, - gei * width, - usize::cast_from(basis_lens[gei]), + 0, + 0, + 0, + term_len, cs_len, mk_len, usize::cast_from(prod_row_base[p]), @@ -1921,36 +1859,71 @@ fn multiply_batch_block( } .executes(|| { let client = CudaRuntime::client(&CudaDevice::default()); - // Shared resident admissible buffers (see [`ResidentDev`]): re-upload the master ONLY - // when this block dereferences past the uploaded prefix (`need_cs` / `need_mk`). The - // master grows continually at the frontier, so re-uploading on mere growth ships - // multi-GB uploads under this mutex on nearly every launch (measured 1.5x wall - // regression); most launches touch only long-uploaded low-degree `R`s and reuse the - // stale handle at its uploaded length. When an upload does fire it captures the full - // current master, amortizing all growth since the last one. Lock order is DEV.lock - // then HOST.read, and nothing under either lock blocks on rayon or a permit. The - // master is append-only, so the uploaded prefix is always a prefix of the current - // host master and every offset `< uploaded` is final. - // `Transient` (degree > cap `R`s): upload this block's freshly-built master and free it - // with the launch — no persistence, no cross-stream sharing, so no realloc guard below. - let (cs_h, cs_len_master, mk_h, mk_len_master) = match mode { + // 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() + }; + + // `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. + let (cs_seg, mk_seg) = match mode { MasterMode::Resident => { - let (cs_h, cs_len_master) = - resident_dev_handle!(client, need_cs, cs, cs_pending, cs_len); - let (mk_h, mk_len_master) = - resident_dev_handle!(client, need_mk, mk, mk_pending, mk_len); - (cs_h, cs_len_master, mk_h, mk_len_master) + 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 nl = h.cs_len; + (std::mem::take(&mut h.cs_pending), 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 nl = h.mk_len; + (std::mem::take(&mut h.mk_pending), nl) + } + ); + (pad_u16(full(cs_segs)), pad_u16(full(mk_segs))) } MasterMode::Transient => { - // Generate this block's cold `col_sums`/`masks` ON the device into transient scratch - // (freed with the launch), instead of uploading host-built arrays. Only the small - // p-parts + dimensions are uploaded. The enumeration launch is issued before the - // multiply on this same stream, so the scratch is fully written when the multiply - // reads it (kernel launches on one stream are ordered, as with `zero_u32` below). + // 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). const ENUM_THREADS: u32 = 256; let n_cold = enum_rows.len(); let cs_cap = need_cs.max(1); let mk_cap = need_mk.max(1); + assert!( + cs_cap <= seg_elems && mk_cap <= seg_elems, + "transient scratch ({cs_cap}/{mk_cap} u16) exceeds one segment ({seg_elems}); \ + raise NASSAU_GPU_MASTER_SEG_ELEMS" + ); let cs_scratch = client.empty(cs_cap * size_of::()); let mk_scratch = client.empty(mk_cap * size_of::()); let cnt_scratch = client.empty(n_cold.max(1) * size_of::()); @@ -1976,27 +1949,47 @@ fn multiply_batch_block( n_cold, ); } - (cs_scratch, cs_cap, mk_scratch, mk_cap) + (pad_u16(vec![(cs_scratch, cs_cap)]), pad_u16(vec![(mk_scratch, mk_cap)])) } }; - // 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. - // Resident basis handles (default) or per-launch passthrough buffers (A/B diagnostic). - // `bp_len`/`bl_len` are the logical array lengths the kernel sees; every `gei` a thread - // dereferences is `< need_basis_elems`, so `need_basis_elems*width` / `need_basis_elems` - // cover it (the resident buffer may be larger — append-only — which is fine). - let (bp_h, bl_h, bp_len, bl_len) = if passthrough { + // 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)); ( - client.create_from_slice(u16::as_bytes(&term_pparts)), - client.create_from_slice(u32::as_bytes(&term_lens)), - term_pparts.len(), - term_lens.len(), + pad_u16(vec![(bp, term_pparts.len())]), + pad_u32(vec![(bl, term_lens.len())]), ) } else { - let (pp, ln) = basis_dev_handles!(client, need_basis_elems); - (pp, ln, need_basis_elems * width, need_basis_elems) + 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. let tg_h = client.create_from_slice(u32::as_bytes(&term_gei)); let g_h = client.create_from_slice(u32::as_bytes(&g)); let xi_h = client.create_from_slice(u32::as_bytes(&xi)); @@ -2005,14 +1998,11 @@ fn multiply_batch_block( 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)); const THREADS: u32 = 256; - // Hold the resident READ lock across the multiply that dereferences the resident buffers, - // through the readback that syncs it (see [`RESIDENT_REALLOC`]). All resident growth for this - // block is already done above; a concurrent capacity realloc on another thread waits here for - // this multiply to finish, so the resident handle can never be swapped mid-kernel. Held to - // the end of the device section (dropped after `read_one`). `Transient` buffers are - // block-local (never reallocated by another thread), so they need no such guard. - let _realloc_guard = - (mode == MasterMode::Resident).then(|| RESIDENT_REALLOC.read().unwrap()); + // 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. Same stream as the multiply below, so it is ordered before it. @@ -2033,17 +2023,35 @@ fn multiply_batch_block( let poo_h = client.create_from_slice(u32::as_bytes(&prod_out_offset)); let pps_h = client.create_from_slice(u32::as_bytes(&pps)); let cubes = (total_pairs as u32).div_ceil(THREADS).max(1); + // Bind one `ArrayArg` per `(segment vector, index)` — the `.0` handle, `.1` element length. + macro_rules! sa { + ($v:expr, $i:expr) => { + ArrayArg::from_raw_parts($v[$i].0.clone(), $v[$i].1) + }; + } // SAFETY: `launch_unchecked` — see the kernel's `address_type = "u64"` note. Every device - // read is in-bounds by construction (uploaded `need_*` prefix + per-column `j` guards). + // 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), - ArrayArg::from_raw_parts(cs_h, cs_len_master), - ArrayArg::from_raw_parts(mk_h, mk_len_master), - ArrayArg::from_raw_parts(bp_h, bp_len), - ArrayArg::from_raw_parts(bl_h, bl_len), + 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), ArrayArg::from_raw_parts(tg_h, term_gei.len()), ArrayArg::from_raw_parts(g_h, g.len()), ArrayArg::from_raw_parts(xi_h, xi.len()), @@ -2059,6 +2067,7 @@ fn multiply_batch_block( ArrayArg::from_raw_parts(poo_h, products.len()), ArrayArg::from_raw_parts(pps_h, pps.len()), width, + seg_elems, ); } @@ -2095,14 +2104,6 @@ fn multiply_batch_block( result } -/// Max number of fixed-size segments a segmented resident master can have. The kernel selects a -/// segment by a static branch (cubecl cannot dynamically index an array-of-buffers), so this is a -/// compile-time bound; `MASTER_SEG_ELEMS` sets the per-segment element count. 8 segments is the -/// prototype size (raise once the mechanic is wired + validated). Together they cap a buffer at -/// `MASTER_MAX_SEG * MASTER_SEG_ELEMS` u16 elements. -#[allow(dead_code)] // used once the segmented master is wired into the multiply path -const MASTER_MAX_SEG: usize = 8; - /// 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 @@ -2110,7 +2111,7 @@ const MASTER_MAX_SEG: usize = 8; /// 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. +/// 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( @@ -2122,6 +2123,14 @@ fn seg_read_u16( s5: &Array, s6: &Array, s7: &Array, + s8: &Array, + s9: &Array, + s10: &Array, + s11: &Array, + s12: &Array, + s13: &Array, + s14: &Array, + s15: &Array, o: usize, seg_elems: usize, ) -> u16 { @@ -2142,15 +2151,93 @@ fn seg_read_u16( v = s5[local]; } else if seg == 6 { v = s6[local]; + } else if seg == 7 { + v = s7[local]; + } else if seg == 8 { + v = s8[local]; + } else if seg == 9 { + v = s9[local]; + } else if seg == 10 { + v = s10[local]; + } else if seg == 11 { + v = s11[local]; + } else if seg == 12 { + v = s12[local]; + } else if seg == 13 { + v = s13[local]; + } else if seg == 14 { + v = s14[local]; } else { + v = s15[local]; + } + 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: &Array, + s1: &Array, + s2: &Array, + s3: &Array, + s4: &Array, + s5: &Array, + s6: &Array, + s7: &Array, + s8: &Array, + s9: &Array, + s10: &Array, + s11: &Array, + s12: &Array, + s13: &Array, + s14: &Array, + s15: &Array, + o: usize, + seg_elems: usize, +) -> u32 { + let seg = o / seg_elems; + let local = o % seg_elems; + let mut v = 0u32; + if seg == 0 { + v = s0[local]; + } else if seg == 1 { + v = s1[local]; + } else if seg == 2 { + v = s2[local]; + } else if seg == 3 { + v = s3[local]; + } else if seg == 4 { + v = s4[local]; + } else if seg == 5 { + v = s5[local]; + } else if seg == 6 { + v = s6[local]; + } else if seg == 7 { v = s7[local]; + } else if seg == 8 { + v = s8[local]; + } else if seg == 9 { + v = s9[local]; + } else if seg == 10 { + v = s10[local]; + } else if seg == 11 { + v = s11[local]; + } else if seg == 12 { + v = s12[local]; + } else if seg == 13 { + v = s13[local]; + } else if seg == 14 { + v = s14[local]; + } else { + v = s15[local]; } v } /// 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 before the mechanic is wired into the -/// multiply's hot path. +/// segmented read reproduces a contiguous buffer bit-for-bit (see `seg_read_matches_contiguous`). #[cfg(test)] #[cube(launch)] #[allow(clippy::too_many_arguments)] @@ -2163,6 +2250,14 @@ fn seg_gather_kernel( s5: &Array, s6: &Array, s7: &Array, + s8: &Array, + s9: &Array, + s10: &Array, + s11: &Array, + s12: &Array, + s13: &Array, + s14: &Array, + s15: &Array, idx: &Array, out: &mut Array, seg_elems: usize, @@ -2172,14 +2267,7 @@ fn seg_gather_kernel( terminate!(); } out[i] = seg_read_u16( - s0, - s1, - s2, - s3, - s4, - s5, - s6, - s7, + s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11, s12, s13, s14, s15, usize::cast_from(idx[i]), seg_elems, ); @@ -2595,6 +2683,14 @@ fn seg_gather_on_gpu(data: &[u16], seg_elems: usize, indices: &[u32]) -> 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 { + s.set_entry(ti, 1); + } + let mut tmp = FpVector::new(p, out_dim); + algebra.multiply_basis_element_by_element_2( + tmp.as_slice_mut(), + 1, + prod.r_degree, + prod.r_idx, + prod.s_degree, + s.as_slice(), + ); + cpu_rows[prod.row].add(&tmp, 1); + } + let num_limbs = out_dim.div_ceil(32).max(1); + let golden: Vec> = cpu_rows + .iter() + .map(|row| { + let mut packed = vec![0u32; num_limbs]; + for (i, _) in row.iter_nonzero() { + packed[i / 32] ^= 1u32 << (i % 32); + } + packed + }) + .collect(); + let got = multiply_batch_on_gpu(&algebra, out_dim, num_rows, &products); + assert_eq!( + got, golden, + "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() + ); + } } From 7265f04de795ee940242a17a284bda7ddeecf685 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Mon, 27 Jul 2026 19:08:57 -0400 Subject: [PATCH 033/127] milnor_gpu: move all test kernels/helpers into the tests module Relocate every #[cfg(test)] item (the test-only cube kernels xor_f2 / seqno_kernel / multiply_single_r_kernel / seg_gather_kernel and their host drivers, plus enumerate_admissible_ref / _on_runtime, seg_gather_on_gpu, check_enum_backend, enum_launch_timed) out of the production module body and into `mod tests`, so the production path is no longer interleaved with test scaffolding. Drops the now-redundant inner #[cfg(test)] attributes (the module already carries it) and applies the project's nightly rustfmt (this also formats the segmented-master rewrite from the previous commit). Pure move; no behavior change. GPU correctness suite still green. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 1682 ++++++++++-------- 1 file changed, 945 insertions(+), 737 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index a2155601f9..2d0734fb28 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -24,6 +24,7 @@ use cubecl::{ prelude::*, }; use cubecl_common::stream_id::StreamId; + // 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; @@ -404,7 +405,11 @@ fn cold_count(algebra: &MilnorAlgebra, p_part: &[PPartEntry]) -> (u32, u32, u32) } 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.to_vec()).or_insert(e); + COLD_COUNT + .write() + .unwrap() + .entry(p_part.to_vec()) + .or_insert(e); e } @@ -517,9 +522,8 @@ struct RStat { last: u64, } -static R_STATS: LazyLock, RStat>>>> = LazyLock::new(|| { - std::env::var_os("NASSAU_R_STATS").map(|_| Mutex::new(HashMap::new())) -}); +static R_STATS: LazyLock, RStat>>>> = + LazyLock::new(|| std::env::var_os("NASSAU_R_STATS").map(|_| Mutex::new(HashMap::new()))); /// Internal degree of `R` from its p-part: `Σ p_part[i] · deg(ξ_{i+1})`. fn ppart_degree(p_part: &[PPartEntry]) -> i32 { @@ -633,7 +637,11 @@ pub fn dump_r_stats() { 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(); + 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, @@ -642,9 +650,10 @@ pub fn dump_r_stats() { } 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:{}", + 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), @@ -770,7 +779,6 @@ fn ensure_basis(algebra: &MilnorAlgebra, width: usize, max_degree: i32) -> Vec) { // 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: &Array, dst: &mut Array, src_off: usize, dst_off: usize, count: u32) { +fn copy_into_u16( + src: &Array, + dst: &mut Array, + 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]; } @@ -807,7 +821,13 @@ fn copy_into_u16(src: &Array, dst: &mut Array, src_off: usize, dst_off /// `u32` sibling of [`copy_into_u16`] (for the resident basis `lens`). #[cube(launch_unchecked, address_type = "dynamic")] -fn copy_into_u32(src: &Array, dst: &mut Array, src_off: usize, dst_off: usize, count: u32) { +fn copy_into_u32( + src: &Array, + dst: &mut Array, + 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]; } @@ -825,7 +845,6 @@ const COPY_CHUNK: usize = 1 << 30; /// 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). @@ -855,50 +874,6 @@ macro_rules! copy_chunked { }}; } -/// Elementwise F₂ addition of two bit-packed vectors: `out[i] = a[i] ^ b[i]`. -/// -/// One thread per `u32` limb. F₂ addition is XOR of the packed limbs, so this is -/// the output primitive the multiply kernels accumulate with. -#[cfg(test)] -#[cube(launch)] -fn xor_f2(a: &Array, b: &Array, out: &mut Array) { - if ABSOLUTE_POS < out.len() { - out[ABSOLUTE_POS] = a[ABSOLUTE_POS] ^ b[ABSOLUTE_POS]; - } -} - -/// Compute `a ^ b` limb-wise on the default CUDA device. -/// -/// Host-side driver for `xor_f2`: uploads both operands, launches one thread per -/// limb, and reads the result back. Panics if the operands differ in length. -#[cfg(test)] -pub fn xor_f2_on_gpu(a: &[u32], b: &[u32]) -> Vec { - assert_eq!(a.len(), b.len(), "operands must have equal limb counts"); - let n = a.len(); - let client = CudaRuntime::client(&CudaDevice::default()); - - let a_handle = client.create_from_slice(u32::as_bytes(a)); - let b_handle = client.create_from_slice(u32::as_bytes(b)); - let out_handle = client.empty(std::mem::size_of_val(a)); - - // One 1-D block of `THREADS` units, enough blocks to cover every limb. - const THREADS: u32 = 256; - let cubes = (n as u32).div_ceil(THREADS); - unsafe { - xor_f2::launch::( - &client, - CubeCount::Static(cubes, 1, 1), - CubeDim::new_1d(THREADS), - ArrayArg::from_raw_parts(a_handle, n), - ArrayArg::from_raw_parts(b_handle, n), - ArrayArg::from_raw_parts(out_handle.clone(), n), - ); - } - - let bytes = client.read_one(out_handle).unwrap(); - u32::from_bytes(&bytes).to_vec() -} - /// Device port of [`MilnorAlgebra::seqno`]: the index of `P(working)` in the Milnor /// basis of its degree, from the flat `g` table with no hashing. `working` holds the /// (trimmed) p_part in its first `wlen` entries; `g` has row width `width`, entry @@ -938,72 +913,6 @@ fn seqno_core( rank } -/// One thread per padded p_part: `out[i] = seqno(p_parts[i])`. `p_parts` is -/// `n × width` row-major, each row a p_part zero-padded to `width` (padding entries -/// are zero and skipped, so `wlen == width` matches the CPU's trimmed loop). -#[cfg(test)] -#[cube(launch)] -fn seqno_kernel( - g: &Array, - xi: &Array, - p_parts: &Array, - out: &mut Array, - width: usize, -) { - let idx = ABSOLUTE_POS; - if idx >= out.len() { - terminate!(); - } - let base = idx * width; - - let mut working = Array::::new(MAX_XI_TAU); - for h in 0..width { - working[h] = p_parts[base + h]; - } - out[idx] = seqno_core(g, xi, &working, width, width); -} - -/// Run `seqno_kernel` over `n` padded p_parts and return their seqno indices. -/// -/// `g`/`xi` come from `MilnorAlgebra::seqno_table_u32` and -/// [`crate::algebra::combinatorics::xi_degrees`]; `p_parts` is `n × width` row-major, -/// each row a p_part zero-padded to `width`. -#[cfg(test)] -pub fn seqno_batch_on_gpu( - width: usize, - xi: &[u32], - g: &[u32], - p_parts: &[u32], - n: usize, -) -> Vec { - assert_eq!(xi.len(), width, "xi must have `width` entries"); - assert_eq!(p_parts.len(), n * width, "p_parts must be n × width"); - let client = CudaRuntime::client(&CudaDevice::default()); - - let g_h = client.create_from_slice(u32::as_bytes(g)); - let xi_h = client.create_from_slice(u32::as_bytes(xi)); - let pp_h = client.create_from_slice(u32::as_bytes(p_parts)); - let out_h = client.empty(n * size_of::()); - - const THREADS: u32 = 256; - let cubes = (n as u32).div_ceil(THREADS); - unsafe { - seqno_kernel::launch::( - &client, - CubeCount::Static(cubes, 1, 1), - CubeDim::new_1d(THREADS), - ArrayArg::from_raw_parts(g_h, g.len()), - ArrayArg::from_raw_parts(xi_h, xi.len()), - ArrayArg::from_raw_parts(pp_h, p_parts.len()), - ArrayArg::from_raw_parts(out_h.clone(), n), - width, - ); - } - - let bytes = client.read_one(out_h).unwrap(); - u32::from_bytes(&bytes).to_vec() -} - /// Assemble one `(admissible matrix, term)` product and XOR its F₂ output bit into /// `out` at `row_base + idx`. The whole per-term test + output assembly of /// [`MilnorAlgebra::multiply_basis_element_by_element_2`] lives here; both the @@ -1098,51 +1007,6 @@ fn multiply_pair( } } -/// Multiply `Sq(R) · s` for a single fixed operation `R` into one F₂ output vector. -/// One thread per `(matrix, term)` pair; delegates the assembly to `multiply_pair`. -#[cfg(test)] -#[cube(launch)] -#[allow(clippy::too_many_arguments)] -fn multiply_single_r_kernel( - col_sums: &Array, - masks: &Array, - term_pparts: &Array, - term_lens: &Array, - g: &Array, - xi: &Array, - out: &mut Array>, - num_terms: usize, - num_matrices: usize, - cs_len: usize, - mk_len: usize, - width: usize, -) { - let pair = ABSOLUTE_POS; - if pair >= num_matrices * num_terms { - terminate!(); - } - let m = pair / num_terms; - let t = pair % num_terms; - let term_len = usize::cast_from(term_lens[t]); - multiply_pair( - col_sums, - masks, - term_pparts, - g, - xi, - out, - m * cs_len, - m * mk_len, - t * width, - term_len, - cs_len, - mk_len, - 0, - 0, - width, - ); -} - /// Batched multiply: one launch covering all `(R, s)` products of (e.g.) a /// `get_partial_matrix` call. One thread per `(product, matrix, term)` pair. /// @@ -1176,22 +1040,70 @@ fn multiply_single_r_kernel( #[cube(launch_unchecked, address_type = "u64")] #[allow(clippy::too_many_arguments)] fn multiply_batch_kernel( - cs0: &Array, cs1: &Array, cs2: &Array, cs3: &Array, - cs4: &Array, cs5: &Array, cs6: &Array, cs7: &Array, - cs8: &Array, cs9: &Array, cs10: &Array, cs11: &Array, - cs12: &Array, cs13: &Array, cs14: &Array, cs15: &Array, - mk0: &Array, mk1: &Array, mk2: &Array, mk3: &Array, - mk4: &Array, mk5: &Array, mk6: &Array, mk7: &Array, - mk8: &Array, mk9: &Array, mk10: &Array, mk11: &Array, - mk12: &Array, mk13: &Array, mk14: &Array, mk15: &Array, - pp0: &Array, pp1: &Array, pp2: &Array, pp3: &Array, - pp4: &Array, pp5: &Array, pp6: &Array, pp7: &Array, - pp8: &Array, pp9: &Array, pp10: &Array, pp11: &Array, - pp12: &Array, pp13: &Array, pp14: &Array, pp15: &Array, - ln0: &Array, ln1: &Array, ln2: &Array, ln3: &Array, - ln4: &Array, ln5: &Array, ln6: &Array, ln7: &Array, - ln8: &Array, ln9: &Array, ln10: &Array, ln11: &Array, - ln12: &Array, ln13: &Array, ln14: &Array, ln15: &Array, + cs0: &Array, + cs1: &Array, + cs2: &Array, + cs3: &Array, + cs4: &Array, + cs5: &Array, + cs6: &Array, + cs7: &Array, + cs8: &Array, + cs9: &Array, + cs10: &Array, + cs11: &Array, + cs12: &Array, + cs13: &Array, + cs14: &Array, + cs15: &Array, + mk0: &Array, + mk1: &Array, + mk2: &Array, + mk3: &Array, + mk4: &Array, + mk5: &Array, + mk6: &Array, + mk7: &Array, + mk8: &Array, + mk9: &Array, + mk10: &Array, + mk11: &Array, + mk12: &Array, + mk13: &Array, + mk14: &Array, + mk15: &Array, + pp0: &Array, + pp1: &Array, + pp2: &Array, + pp3: &Array, + pp4: &Array, + pp5: &Array, + pp6: &Array, + pp7: &Array, + pp8: &Array, + pp9: &Array, + pp10: &Array, + pp11: &Array, + pp12: &Array, + pp13: &Array, + pp14: &Array, + pp15: &Array, + ln0: &Array, + ln1: &Array, + ln2: &Array, + ln3: &Array, + ln4: &Array, + ln5: &Array, + ln6: &Array, + ln7: &Array, + ln8: &Array, + ln9: &Array, + ln10: &Array, + ln11: &Array, + ln12: &Array, + ln13: &Array, + ln14: &Array, + ln15: &Array, term_gei: &Array, g: &Array, xi: &Array, @@ -1252,9 +1164,8 @@ fn multiply_batch_kernel( let mk_off = usize::cast_from(r_mk_offset[ri]) + m * mk_len; let pp_off = gei * width; let term_len = usize::cast_from(seg_read_u32( - ln0, ln1, ln2, ln3, ln4, ln5, ln6, ln7, - ln8, ln9, ln10, ln11, ln12, ln13, ln14, ln15, - gei, seg_elems, + ln0, ln1, ln2, ln3, ln4, ln5, ln6, ln7, ln8, ln9, ln10, ln11, ln12, ln13, ln14, ln15, gei, + seg_elems, )); // Gather this thread's matrix / term out of the segmented stores into contiguous locals, then run @@ -1268,27 +1179,72 @@ fn multiply_batch_kernel( let mut c = 0u16; if j < cs_len { c = seg_read_u16( - cs0, cs1, cs2, cs3, cs4, cs5, cs6, cs7, - cs8, cs9, cs10, cs11, cs12, cs13, cs14, cs15, - cs_off + j, seg_elems, + cs0, + cs1, + cs2, + cs3, + cs4, + cs5, + cs6, + cs7, + cs8, + cs9, + cs10, + cs11, + cs12, + cs13, + cs14, + cs15, + cs_off + j, + seg_elems, ); } cs_local[j] = c; let mut mm = 0u16; if j < mk_len { mm = seg_read_u16( - mk0, mk1, mk2, mk3, mk4, mk5, mk6, mk7, - mk8, mk9, mk10, mk11, mk12, mk13, mk14, mk15, - mk_off + j, seg_elems, + mk0, + mk1, + mk2, + mk3, + mk4, + mk5, + mk6, + mk7, + mk8, + mk9, + mk10, + mk11, + mk12, + mk13, + mk14, + mk15, + mk_off + j, + seg_elems, ); } mk_local[j] = mm; let mut b = 0u16; if j < term_len { b = seg_read_u16( - pp0, pp1, pp2, pp3, pp4, pp5, pp6, pp7, - pp8, pp9, pp10, pp11, pp12, pp13, pp14, pp15, - pp_off + j, seg_elems, + pp0, + pp1, + pp2, + pp3, + pp4, + pp5, + pp6, + pp7, + pp8, + pp9, + pp10, + pp11, + pp12, + pp13, + pp14, + pp15, + pp_off + j, + seg_elems, ); } term_local[j] = b; @@ -1313,102 +1269,6 @@ fn multiply_batch_kernel( ); } -/// Compute `Sq(R) · s` on the GPU for a single operation `R = (r_degree, r_idx)`, -/// returning the F₂ result as bit-packed `u32` limbs (bit `i` = basis index `i`). -/// -/// `term_indices` are the nonzero indices of `s` in the degree-`s_degree` basis. -/// `R` must be non-empty (`Sq(∅) = 1` is the trivial identity the caller handles). -/// Requires the algebra's basis and seqno tables built through `r_degree + s_degree`. -#[cfg(test)] -pub fn multiply_single_r_on_gpu( - algebra: &MilnorAlgebra, - r_degree: i32, - r_idx: usize, - s_degree: i32, - term_indices: &[usize], -) -> Vec { - let (width, g) = algebra.seqno_table_u32(); - // Pad `xi` to `WORKING_CAP` so the kernel's `cur_d` sum (which runs to the full - // working capacity) never reads out of bounds; padding entries multiply zero. - let mut xi: Vec = xi_degrees(algebra.prime()) - .iter() - .map(|&x| x as u32) - .collect(); - xi.resize(WORKING_CAP, 0); - - let r = algebra.basis_element_from_index(r_degree, r_idx); - assert!( - !r.p_part.is_empty(), - "R must be non-empty (Sq(∅) = 1 is the identity)" - ); - let (cs_len, mk_len, cs32, mk32) = algebra.admissible_matrices(&r.p_part); - // Ship admissible-matrix / term data as u16 (see `multiply_batch_on_gpu`). - let mut col_sums: Vec = cs32.iter().map(|&v| narrow_u16(v)).collect(); - let masks: Vec = mk32.iter().map(|&v| narrow_u16(v)).collect(); - let num_matrices = masks.len() / mk_len; - - // Terms of s, each p_part padded to `width`, with their true (trimmed) lengths. - let num_terms = term_indices.len(); - let mut term_pparts = vec![0u16; num_terms * width]; - let mut term_lens = vec![0u32; num_terms]; - for (t, &ti) in term_indices.iter().enumerate() { - let elt = algebra.basis_element_from_index(s_degree, ti); - term_lens[t] = elt.p_part.len() as u32; - for (slot, &v) in term_pparts[t * width..(t + 1) * width] - .iter_mut() - .zip(&elt.p_part) - { - *slot = narrow_u16(v); - } - } - - let out_degree = r_degree + s_degree; - let dim = algebra.dimension(out_degree); - let num_limbs = dim.div_ceil(32).max(1); - - // Device buffers must be non-empty; `cs_len == 0` (R's max entry is 1) leaves - // `col_sums` empty. The kernel never reads past the real lengths. - if col_sums.is_empty() { - col_sums.push(0); - } - - let client = CudaRuntime::client(&CudaDevice::default()); - let cs_h = client.create_from_slice(u16::as_bytes(&col_sums)); - let mk_h = client.create_from_slice(u16::as_bytes(&masks)); - let tp_h = client.create_from_slice(u16::as_bytes(&term_pparts)); - let tl_h = client.create_from_slice(u32::as_bytes(&term_lens)); - let g_h = client.create_from_slice(u32::as_bytes(&g)); - let xi_h = client.create_from_slice(u32::as_bytes(&xi)); - let zeros = vec![0u32; num_limbs]; - let out_h = client.create_from_slice(u32::as_bytes(&zeros)); - - let total_pairs = num_matrices * num_terms; - const THREADS: u32 = 256; - let cubes = (total_pairs as u32).div_ceil(THREADS).max(1); - unsafe { - multiply_single_r_kernel::launch::( - &client, - CubeCount::Static(cubes, 1, 1), - CubeDim::new_1d(THREADS), - ArrayArg::from_raw_parts(cs_h, col_sums.len()), - ArrayArg::from_raw_parts(mk_h, masks.len()), - ArrayArg::from_raw_parts(tp_h, term_pparts.len()), - ArrayArg::from_raw_parts(tl_h, term_lens.len()), - ArrayArg::from_raw_parts(g_h, g.len()), - ArrayArg::from_raw_parts(xi_h, xi.len()), - ArrayArg::from_raw_parts(out_h.clone(), num_limbs), - num_terms, - num_matrices, - cs_len, - mk_len, - width, - ); - } - - let bytes = client.read_one(out_h).unwrap(); - u32::from_bytes(&bytes).to_vec() -} - /// One `Sq(R) · s` product of a batched launch, written into output row `row` at bit /// offset `out_offset`. /// @@ -1474,8 +1334,7 @@ pub fn multiply_batch_on_gpu( if rows.is_empty() { continue; } - let remap: HashMap = - rows.iter().enumerate().map(|(i, &r)| (r, i)).collect(); + let remap: HashMap = rows.iter().enumerate().map(|(i, &r)| (r, i)).collect(); let compact: Vec = products .iter() .filter(|p| is_group(p.r_degree)) @@ -1750,7 +1609,10 @@ fn multiply_batch_block( .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)); + 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.clone()); @@ -1819,8 +1681,8 @@ fn multiply_batch_block( 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}", + 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), @@ -1870,14 +1732,20 @@ fn multiply_batch_block( 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"); + 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"); + assert!( + v.len() <= MASTER_MAX_SEG, + "segment count exceeds MASTER_MAX_SEG" + ); while v.len() < MASTER_MAX_SEG { v.push((dummy32.clone(), 1)); } @@ -1892,8 +1760,14 @@ fn multiply_batch_block( 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, + 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 nl = h.cs_len; @@ -1901,8 +1775,14 @@ fn multiply_batch_block( } ); let (mk_segs, _) = seg_grow!( - client, RESIDENT_DEV, mk, RESIDENT_UPLOAD, need_mk, - copy_into_u16, u16::as_bytes, u16, + 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 nl = h.mk_len; @@ -1949,7 +1829,10 @@ fn multiply_batch_block( n_cold, ); } - (pad_u16(vec![(cs_scratch, cs_cap)]), pad_u16(vec![(mk_scratch, mk_cap)])) + ( + pad_u16(vec![(cs_scratch, cs_cap)]), + pad_u16(vec![(mk_scratch, mk_cap)]), + ) } }; // Resident basis segments (default) or per-launch passthrough buffers (A/B diagnostic) bound @@ -1968,8 +1851,14 @@ fn multiply_batch_block( ) } 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, + 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; @@ -1977,8 +1866,14 @@ fn multiply_batch_block( } ); let (ln_segs, _) = seg_grow!( - client, RESIDENT_BASIS_DEV, ln, RESIDENT_BASIS_UPLOAD, need_basis_elems, - copy_into_u32, u32::as_bytes, u32, + 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(); @@ -2036,22 +1931,70 @@ fn multiply_batch_block( &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), + 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), ArrayArg::from_raw_parts(tg_h, term_gei.len()), ArrayArg::from_raw_parts(g_h, g.len()), ArrayArg::from_raw_parts(xi_h, xi.len()), @@ -2236,43 +2179,6 @@ fn seg_read_u32( v } -/// 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`). -#[cfg(test)] -#[cube(launch)] -#[allow(clippy::too_many_arguments)] -fn seg_gather_kernel( - s0: &Array, - s1: &Array, - s2: &Array, - s3: &Array, - s4: &Array, - s5: &Array, - s6: &Array, - s7: &Array, - s8: &Array, - s9: &Array, - s10: &Array, - s11: &Array, - s12: &Array, - s13: &Array, - s14: &Array, - s15: &Array, - idx: &Array, - out: &mut Array, - seg_elems: 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, - ); -} - /// 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 @@ -2437,409 +2343,706 @@ fn enumerate_admissible_kernel( out_counts[ri] = u32::cast_from(mat); } -/// 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. #[cfg(test)] -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; +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: &Array, b: &Array, out: &mut Array) { + if ABSOLUTE_POS < out.len() { + out[ABSOLUTE_POS] = a[ABSOLUTE_POS] ^ b[ABSOLUTE_POS]; + } + } + + /// Compute `a ^ b` limb-wise on the default CUDA device. + /// + /// Host-side driver for `xor_f2`: uploads both operands, launches one thread per + /// limb, and reads the result back. Panics if the operands differ in length. + pub fn xor_f2_on_gpu(a: &[u32], b: &[u32]) -> Vec { + assert_eq!(a.len(), b.len(), "operands must have equal limb counts"); + let n = a.len(); + let client = CudaRuntime::client(&CudaDevice::default()); + + let a_handle = client.create_from_slice(u32::as_bytes(a)); + let b_handle = client.create_from_slice(u32::as_bytes(b)); + let out_handle = client.empty(std::mem::size_of_val(a)); + + // One 1-D block of `THREADS` units, enough blocks to cover every limb. + const THREADS: u32 = 256; + let cubes = (n as u32).div_ceil(THREADS); + unsafe { + xor_f2::launch::( + &client, + CubeCount::Static(cubes, 1, 1), + CubeDim::new_1d(THREADS), + ArrayArg::from_raw_parts(a_handle, n), + ArrayArg::from_raw_parts(b_handle, n), + ArrayArg::from_raw_parts(out_handle.clone(), n), + ); + } + + let bytes = client.read_one(out_handle).unwrap(); + u32::from_bytes(&bytes).to_vec() + } + + /// One thread per padded p_part: `out[i] = seqno(p_parts[i])`. `p_parts` is + /// `n × width` row-major, each row a p_part zero-padded to `width` (padding entries + /// are zero and skipped, so `wlen == width` matches the CPU's trimmed loop). + #[cube(launch)] + fn seqno_kernel( + g: &Array, + xi: &Array, + p_parts: &Array, + out: &mut Array, + width: usize, + ) { + let idx = ABSOLUTE_POS; + if idx >= out.len() { + terminate!(); + } + let base = idx * width; + + let mut working = Array::::new(MAX_XI_TAU); + for h in 0..width { + working[h] = p_parts[base + h]; + } + out[idx] = seqno_core(g, xi, &working, width, width); + } + + /// Run `seqno_kernel` over `n` padded p_parts and return their seqno indices. + /// + /// `g`/`xi` come from `MilnorAlgebra::seqno_table_u32` and + /// [`crate::algebra::combinatorics::xi_degrees`]; `p_parts` is `n × width` row-major, + /// each row a p_part zero-padded to `width`. + pub fn seqno_batch_on_gpu( + width: usize, + xi: &[u32], + g: &[u32], + p_parts: &[u32], + n: usize, + ) -> Vec { + assert_eq!(xi.len(), width, "xi must have `width` entries"); + assert_eq!(p_parts.len(), n * width, "p_parts must be n × width"); + let client = CudaRuntime::client(&CudaDevice::default()); + + let g_h = client.create_from_slice(u32::as_bytes(g)); + let xi_h = client.create_from_slice(u32::as_bytes(xi)); + let pp_h = client.create_from_slice(u32::as_bytes(p_parts)); + let out_h = client.empty(n * size_of::()); + + const THREADS: u32 = 256; + let cubes = (n as u32).div_ceil(THREADS); + unsafe { + seqno_kernel::launch::( + &client, + CubeCount::Static(cubes, 1, 1), + CubeDim::new_1d(THREADS), + ArrayArg::from_raw_parts(g_h, g.len()), + ArrayArg::from_raw_parts(xi_h, xi.len()), + ArrayArg::from_raw_parts(pp_h, p_parts.len()), + ArrayArg::from_raw_parts(out_h.clone(), n), + width, + ); } - 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 = 64; - 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), - ArrayArg::from_raw_parts(pp_h, pp_flat.len()), - ArrayArg::from_raw_parts(rr_h, n_r), - ArrayArg::from_raw_parts(rc_h, n_r), - ArrayArg::from_raw_parts(rco_h, n_r), - ArrayArg::from_raw_parts(rmo_h, n_r), - ArrayArg::from_raw_parts(ocs_h.clone(), cs_cap), - ArrayArg::from_raw_parts(omk_h.clone(), mk_cap), - ArrayArg::from_raw_parts(cnt_h.clone(), n_r), + + 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: &Array, + masks: &Array, + term_pparts: &Array, + term_lens: &Array, + g: &Array, + xi: &Array, + out: &mut Array>, + num_terms: usize, + num_matrices: usize, + cs_len: usize, + mk_len: usize, + width: usize, + ) { + let pair = ABSOLUTE_POS; + if pair >= num_matrices * num_terms { + terminate!(); + } + let m = pair / num_terms; + let t = pair % num_terms; + let term_len = usize::cast_from(term_lens[t]); + multiply_pair( + col_sums, + masks, + term_pparts, + g, + xi, + out, + m * cs_len, + m * mk_len, + t * width, + term_len, + cs_len, + mk_len, + 0, + 0, width, - n_r, ); } - // 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) -} -/// 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. -#[cfg(test)] -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; + /// 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); - // 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 r = algebra.basis_element_from_index(r_degree, r_idx); + assert!( + !r.p_part.is_empty(), + "R must be non-empty (Sq(∅) = 1 is the identity)" + ); + let (cs_len, mk_len, cs32, mk32) = algebra.admissible_matrices(&r.p_part); + // Ship admissible-matrix / term data as u16 (see `multiply_batch_on_gpu`). + let mut col_sums: Vec = cs32.iter().map(|&v| narrow_u16(v)).collect(); + let masks: Vec = mk32.iter().map(|&v| narrow_u16(v)).collect(); + let num_matrices = masks.len() / mk_len; + + // Terms of s, each p_part padded to `width`, with their true (trimmed) lengths. + let num_terms = term_indices.len(); + let mut term_pparts = vec![0u16; num_terms * width]; + let mut term_lens = vec![0u32; num_terms]; + for (t, &ti) in term_indices.iter().enumerate() { + let elt = algebra.basis_element_from_index(s_degree, ti); + term_lens[t] = elt.p_part.len() as u32; + for (slot, &v) in term_pparts[t * width..(t + 1) * width] + .iter_mut() + .zip(&elt.p_part) + { + *slot = narrow_u16(v); + } + } + + let out_degree = r_degree + s_degree; + let dim = algebra.dimension(out_degree); + let num_limbs = dim.div_ceil(32).max(1); + + // Device buffers must be non-empty; `cs_len == 0` (R's max entry is 1) leaves + // `col_sums` empty. The kernel never reads past the real lengths. + if col_sums.is_empty() { + col_sums.push(0); + } + + let client = CudaRuntime::client(&CudaDevice::default()); + let cs_h = client.create_from_slice(u16::as_bytes(&col_sums)); + let mk_h = client.create_from_slice(u16::as_bytes(&masks)); + let tp_h = client.create_from_slice(u16::as_bytes(&term_pparts)); + let tl_h = client.create_from_slice(u32::as_bytes(&term_lens)); + let g_h = client.create_from_slice(u32::as_bytes(&g)); + let xi_h = client.create_from_slice(u32::as_bytes(&xi)); + let zeros = vec![0u32; num_limbs]; + let out_h = client.create_from_slice(u32::as_bytes(&zeros)); + + let total_pairs = num_matrices * num_terms; + const THREADS: u32 = 256; + let cubes = (total_pairs as u32).div_ceil(THREADS).max(1); + unsafe { + multiply_single_r_kernel::launch::( + &client, + CubeCount::Static(cubes, 1, 1), + CubeDim::new_1d(THREADS), + ArrayArg::from_raw_parts(cs_h, col_sums.len()), + ArrayArg::from_raw_parts(mk_h, masks.len()), + ArrayArg::from_raw_parts(tp_h, term_pparts.len()), + ArrayArg::from_raw_parts(tl_h, term_lens.len()), + ArrayArg::from_raw_parts(g_h, g.len()), + ArrayArg::from_raw_parts(xi_h, xi.len()), + ArrayArg::from_raw_parts(out_h.clone(), num_limbs), + num_terms, + num_matrices, + cs_len, + mk_len, + width, + ); + } + + let bytes = client.read_one(out_h).unwrap(); + u32::from_bytes(&bytes).to_vec() + } + + /// 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: &Array, + s1: &Array, + s2: &Array, + s3: &Array, + s4: &Array, + s5: &Array, + s6: &Array, + s7: &Array, + s8: &Array, + s9: &Array, + s10: &Array, + s11: &Array, + s12: &Array, + s13: &Array, + s14: &Array, + s15: &Array, + idx: &Array, + out: &mut Array, + seg_elems: 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, + ); } - let mut out_cs: Vec = Vec::new(); - let mut out_mk: Vec = Vec::new(); + /// 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; + } - // 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); + 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 = 64; + 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), + ArrayArg::from_raw_parts(pp_h, pp_flat.len()), + ArrayArg::from_raw_parts(rr_h, n_r), + ArrayArg::from_raw_parts(rc_h, n_r), + ArrayArg::from_raw_parts(rco_h, n_r), + ArrayArg::from_raw_parts(rmo_h, n_r), + ArrayArg::from_raw_parts(ocs_h.clone(), cs_cap), + ArrayArg::from_raw_parts(omk_h.clone(), mk_cap), + ArrayArg::from_raw_parts(cnt_h.clone(), n_r), + width, + n_r, + ); + } + // 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) + } + + /// 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; + } - // 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; + 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; } - matrix[row * cols + col] = new_entry; - let mut i = 0; - while i < row { - matrix[i * cols] = totals[i]; - masks[i] = totals[i]; + 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 < 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; + 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; } - i += 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; } - masks[row + col] = d | new_entry; - found = true; - handled = true; } + if !handled { + totals[row] += p_to_the_j * matrix[row * cols + col]; + } + col += 1; } - if !handled { - totals[row] += p_to_the_j * matrix[row * cols + col]; - } - col += 1; + row += 1; } - row += 1; + more = found; } - more = found; + + (cs_len, mk_len, out_cs, out_mk) } - (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 = CudaRuntime::client(&CudaDevice::default()); -/// 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. -#[cfg(test)] -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 = CudaRuntime::client(&CudaDevice::default()); + // 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::()); - // 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); + const THREADS: u32 = 256; + let cubes = (indices.len() as u32).div_ceil(THREADS).max(1); + let arg = |i: usize| unsafe { ArrayArg::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), + ArrayArg::from_raw_parts(idx_h, indices.len()), + ArrayArg::from_raw_parts(out_h.clone(), indices.len()), + seg_elems, + ); } + u16::from_bytes(&client.read_one(out_h).unwrap()).to_vec() } - 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 { ArrayArg::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), - ArrayArg::from_raw_parts(idx_h, indices.len()), - ArrayArg::from_raw_parts(out_h.clone(), indices.len()), - seg_elems, - ); - } - 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. -#[cfg(test)] -fn check_enum_backend(device: &Rt::Device, max_degree: i32) { - use fp::prime::ValidPrime; - - let p = ValidPrime::new(2); - let algebra = 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 = algebra.basis_element_from_index(deg, idx).p_part.clone(); - if pp.is_empty() { + /// 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 = 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 = algebra.basis_element_from_index(deg, idx).p_part.clone(); + 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 (_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. -#[cfg(test)] -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; + 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::(); } - 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 = 64; - 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), - ArrayArg::from_raw_parts(pp_h, pp_flat.len()), - ArrayArg::from_raw_parts(rr_h, n_r), - ArrayArg::from_raw_parts(rc_h, n_r), - ArrayArg::from_raw_parts(rco_h, n_r), - ArrayArg::from_raw_parts(rmo_h, n_r), - ArrayArg::from_raw_parts(ocs_h.clone(), cs_cap), - ArrayArg::from_raw_parts(omk_h.clone(), mk_cap), - ArrayArg::from_raw_parts(cnt_h.clone(), n_r), - width, - n_r, + 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})" ); } - // 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(); + /// 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 = 64; + 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), + ArrayArg::from_raw_parts(pp_h, pp_flat.len()), + ArrayArg::from_raw_parts(rr_h, n_r), + ArrayArg::from_raw_parts(rc_h, n_r), + ArrayArg::from_raw_parts(rco_h, n_r), + ArrayArg::from_raw_parts(rmo_h, n_r), + ArrayArg::from_raw_parts(ocs_h.clone(), cs_cap), + ArrayArg::from_raw_parts(omk_h.clone(), mk_cap), + ArrayArg::from_raw_parts(cnt_h.clone(), n_r), + width, + n_r, + ); + } + // 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(); - (marshal_s, kernel_s, readback_s) -} + 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(); -#[cfg(test)] -mod tests { - use super::*; + (marshal_s, kernel_s, readback_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 — @@ -2865,7 +3068,10 @@ mod tests { 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"); + assert_eq!( + got, want, + "segmented read diverged from contiguous indexing" + ); } /// Throughput comparison, CPU `admissible_matrices` vs the in-kernel [`enumerate_admissible_kernel`], @@ -2875,9 +3081,10 @@ mod tests { #[test] #[ignore = "benchmark, not a correctness check; run explicitly with --ignored --nocapture"] fn bench_admissible_cpu_vs_gpu() { - use fp::prime::ValidPrime; use std::time::Instant; + use fp::prime::ValidPrime; + let p = ValidPrime::new(2); let algebra = MilnorAlgebra::new(p, false); let max_degree = 130; @@ -2942,12 +3149,12 @@ mod tests { } eprintln!( - "\n=== admissible matrices onto the device: enumerate vs upload (degrees 1..={max_degree}) ===\n\ - R's: {total_r} matrices: {total_mats} (both paths leave the arrays ON-DEVICE, no readback)\n\ - CPU admissible_matrices (enumerate, 1 core) : {cpu_secs:.3} s\n\ - GPU enumerate in-kernel : {g_kernel:.3} s\n\ - H->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", + "\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 : {g_kernel:.3} \ + s\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 / upload_secs, ); } @@ -3279,7 +3486,8 @@ mod tests { 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}" + "incremental-growth GPU multiply diverged from reference at \ + out_degree={out_degree}" ); }; From c6879d9ff94157fa2d6aa104c1eee8da1cdfba1a Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Mon, 27 Jul 2026 23:20:29 -0400 Subject: [PATCH 034/127] milnor_gpu: persistent seqno (g/xi) + per-stream output buffer (kill per-launch churn) Phase 1 of reducing the per-launch allocation/copy churn that drives cubecl's allocator into CUDA_ERROR_LAUNCH_FAILED (719) at scale (memcheck showed the fault is on cuMemcpyHtoDAsync/cuMemAllocAsync, not our kernels). Keeps the exclusive-pages memory mode (the ~50% footprint win) and removes the churn it caused, instead of switching to the sub-slice pool (which reuses buffers but ~2x memory -> OOM at the 90 GB master). - g/xi: the seqno table and (constant) xi degrees are identical every launch at a given degree; upload them once to shared resident handles (RESIDENT_SEQNO, keyed by g.len()), re-uploading only on a degree bump, instead of a create_from_slice each launch. Synced before publish for cross-stream reads. - out_h: reuse a persistent per-worker (= per-stream) XOR accumulator (OUT_ACCUM thread-local), grown only when a larger out_len appears, instead of empty() + free-via-memory_cleanup every launch (the single biggest churned buffer). Read back only the used [0,out_len) prefix via Handle::offset_end. memory_cleanup stays for now (trims the still-churning per-R/per-product metadata; phase 2 makes those persistent and drops it). Bit-exact vs the prior path: multiply_batch_matches_reference, multiply_batch_incremental_growth (5 growing launches: out_h regrow + g/xi re-upload), multiply_single_r. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 103 ++++++++++++++++--- 1 file changed, 91 insertions(+), 12 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 2d0734fb28..4ebbd27cc9 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -208,6 +208,7 @@ fn narrow_u16(v: u32) -> u16 { u16::try_from(v).expect("admissible/term entry exceeds u16") } +use std::cell::RefCell; use std::sync::atomic::{AtomicU64, Ordering}; /// Aggregate [`multiply_batch_on_gpu`] counters across all launches (call count, host @@ -384,6 +385,70 @@ static RESIDENT_DEV: LazyLock> = /// memcpy-ing thread). static RESIDENT_UPLOAD: Mutex<()> = Mutex::new(()); +/// 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(|| RwLock::new(None)); +/// Serializes seqno-table uploads only (never reads); see [`RESIDENT_UPLOAD`]. +static RESIDENT_SEQNO_UPLOAD: Mutex<()> = Mutex::new(()); + +/// 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, + } + }; + 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) + } + } + } + } + }}; +} + +thread_local! { + /// Per-worker (= per-stream, see [`thread_stream_id`]) persistent XOR-accumulator buffer, grown + /// to the largest `out_len` this thread has needed and reused every launch — replacing the + /// per-launch `empty()` + free-via-`memory_cleanup` of the block output, the single biggest + /// allocation churned each launch. Per-thread (not shared) because each stream's multiply XORs + /// into its own output rows; a shared buffer would corrupt concurrent blocks. Holds + /// `(capacity_in_u32_elements, handle)`; the handle stays live (refcount > 0) so `memory_cleanup` + /// skips it. + static OUT_ACCUM: RefCell> = const { RefCell::new(None) }; +} + /// 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 @@ -1886,8 +1951,9 @@ fn multiply_batch_block( // 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. let tg_h = client.create_from_slice(u32::as_bytes(&term_gei)); - let g_h = client.create_from_slice(u32::as_bytes(&g)); - let xi_h = client.create_from_slice(u32::as_bytes(&xi)); + // `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)); @@ -1898,10 +1964,20 @@ fn multiply_batch_block( // [`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. Same stream as the multiply below, so it is ordered before it. - let out_h = client.empty(out_len * size_of::()); + // XOR accumulator: reuse this worker's persistent per-stream buffer (see [`OUT_ACCUM`]), + // growing it only when a larger `out_len` appears, instead of `empty()`-ing a fresh one every + // launch. `zero_u32` clears the used `[0, out_len)` prefix on-device (the multiply XORs into + // that range; any stale tail past `out_len` is never bound or read). Same stream as the + // multiply below, so the zero is ordered before it. + let (out_h, out_cap) = OUT_ACCUM.with(|cell| { + let mut slot = cell.borrow_mut(); + let cur_cap = slot.as_ref().map(|(c, _)| *c).unwrap_or(0); + if cur_cap < out_len { + *slot = Some((out_len, client.empty(out_len * size_of::()))); + } + let (cap, h) = slot.as_ref().unwrap(); + (h.clone(), *cap) + }); unsafe { zero_u32::launch::( &client, @@ -2014,17 +2090,20 @@ fn multiply_batch_block( ); } - let bytes = client.read_one(out_h).unwrap(); + // Read back only the used `[0, out_len)` prefix — the persistent buffer may be larger than + // this launch needs (`out_cap >= out_len`), so trim the unused tail off the read. + let read_h = out_h.offset_end(((out_cap - out_len) * size_of::()) as u64); + let bytes = client.read_one(read_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. + // Trim the pool: the remaining per-launch uploads (the small per-`R`/per-product metadata + // arrays) still vary in size and would otherwise accumulate in the pool. The persistent + // buffers — the resident master/basis, the seqno `g`/`xi`, and this stream's `OUT_ACCUM` + // accumulator — stay alive (refcount > 0), so cleanup skips them and only the transient + // metadata is reclaimed. (Phase 2 will make the metadata persistent too and drop this.) client.memory_cleanup(); result From a0af1db20233b5cab85cf791509bfad4f62a0b06 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Tue, 28 Jul 2026 12:29:57 -0400 Subject: [PATCH 035/127] milnor_gpu: drop the per-thread persistent out_h (keep resident g/xi) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Revert the OUT_ACCUM thread-local output buffer from the previous commit: it was per rayon worker (~100 threads), so it held ~100 persistent block-sized buffers (~28-50 GB) and DEFEATED the existing out_h bounding — each launch's out_h is already chunked to NASSAU_GPU_BLOCK_MB (512 MiB) and total in-flight output is capped by the GPU_BUDGET permit (NASSAU_GPU_MEM_BUDGET_MB). Persisting them per thread pushed device memory toward the OOM ceiling. out_h returns to a per-launch empty() reclaimed by memory_cleanup (properly chunked + budget-bounded). The resident g/xi seqno tables (uploaded once vs re-uploaded per launch) are kept — a clean churn reduction. Bit-exact: multiply_batch_matches_reference, multiply_batch_incremental_growth, multiply_single_r. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 45 +++++--------------- 1 file changed, 11 insertions(+), 34 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 4ebbd27cc9..427ce31f5f 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -208,7 +208,6 @@ fn narrow_u16(v: u32) -> u16 { u16::try_from(v).expect("admissible/term entry exceeds u16") } -use std::cell::RefCell; use std::sync::atomic::{AtomicU64, Ordering}; /// Aggregate [`multiply_batch_on_gpu`] counters across all launches (call count, host @@ -438,16 +437,6 @@ macro_rules! resident_seqno { }}; } -thread_local! { - /// Per-worker (= per-stream, see [`thread_stream_id`]) persistent XOR-accumulator buffer, grown - /// to the largest `out_len` this thread has needed and reused every launch — replacing the - /// per-launch `empty()` + free-via-`memory_cleanup` of the block output, the single biggest - /// allocation churned each launch. Per-thread (not shared) because each stream's multiply XORs - /// into its own output rows; a shared buffer would corrupt concurrent blocks. Holds - /// `(capacity_in_u32_elements, handle)`; the handle stays live (refcount > 0) so `memory_cleanup` - /// skips it. - static OUT_ACCUM: RefCell> = const { RefCell::new(None) }; -} /// 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) @@ -1964,20 +1953,12 @@ fn multiply_batch_block( // [`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. - // XOR accumulator: reuse this worker's persistent per-stream buffer (see [`OUT_ACCUM`]), - // growing it only when a larger `out_len` appears, instead of `empty()`-ing a fresh one every - // launch. `zero_u32` clears the used `[0, out_len)` prefix on-device (the multiply XORs into - // that range; any stale tail past `out_len` is never bound or read). Same stream as the - // multiply below, so the zero is ordered before it. - let (out_h, out_cap) = OUT_ACCUM.with(|cell| { - let mut slot = cell.borrow_mut(); - let cur_cap = slot.as_ref().map(|(c, _)| *c).unwrap_or(0); - if cur_cap < out_len { - *slot = Some((out_len, client.empty(out_len * size_of::()))); - } - let (cap, h) = slot.as_ref().unwrap(); - (h.clone(), *cap) - }); + // 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, @@ -2090,20 +2071,16 @@ fn multiply_batch_block( ); } - // Read back only the used `[0, out_len)` prefix — the persistent buffer may be larger than - // this launch needs (`out_cap >= out_len`), so trim the unused tail off the read. - let read_h = out_h.offset_end(((out_cap - out_len) * size_of::()) as u64); - let bytes = client.read_one(read_h).unwrap(); + 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(); - // Trim the pool: the remaining per-launch uploads (the small per-`R`/per-product metadata - // arrays) still vary in size and would otherwise accumulate in the pool. The persistent - // buffers — the resident master/basis, the seqno `g`/`xi`, and this stream's `OUT_ACCUM` - // accumulator — stay alive (refcount > 0), so cleanup skips them and only the transient - // metadata is reclaimed. (Phase 2 will make the metadata persistent too and drop this.) + // Trim the pool: the per-launch uploads (`out_h` and the per-`R`/per-product metadata arrays) + // vary in size and would otherwise accumulate. The persistent buffers — the resident + // master/basis and the seqno `g`/`xi` — stay alive (refcount > 0), so cleanup skips them and + // only the transient per-launch memory is reclaimed. client.memory_cleanup(); result From 72190b3f83bcf1a2dbbeaaf8ac41c47be947690e Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Wed, 29 Jul 2026 02:29:30 -0400 Subject: [PATCH 036/127] milnor_gpu: CPU-fallback stopgap for cubecl context death + GPU multiply bench MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When the cubecl CUDA context is poisoned by the unresolved uninit-handle bug (CUDA_ERROR_LAUNCH_FAILED / ServerUnhealthy at the multiply readback, milnor_gpu.rs:2074; tracel-ai/cubecl#1401), the whole context is dead — every later launch fails, so per-call retry cannot recover. Wrap multiply_batch_on_gpu in catch_unwind: on the first failure, set a process-wide GPU_DISABLED flag and finish the resolution on the CPU via a new cpu_multiply_batch. The CPU output is bit-identical to the GPU's (validated by cpu_multiply_batch_matches_gpu, including the multi-block out_offset path a real module row uses). RREF runs on a separate fp-cuda runtime and is intentionally not gated by this flag. Add benches/nassau_milnor_gpu.rs: a GPU batched-multiply throughput / regression bench (counterpart to nassau_milnor.rs) that hammers multiply_batch_on_gpu over an output-degree sweep with no row-reduction or resolution machinery — for `cargo bench --baseline` comparison across cubecl commits and for isolating a multiply/allocator crash from the rest of the pipeline. Compiles to a no-op main without the `gpu` feature. Also pins cubecl/cubecl-common to the JoeyBF fork branch (claude/pool-slot-map-v0.10.0) for the #1401 generational-slot-pool validation. TEMPORARY: revert to the crates.io release before merging upstream. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/Cargo.toml | 10 +- .../algebra/benches/nassau_milnor_gpu.rs | 121 +++++++++++++ ext/crates/algebra/src/algebra/milnor_gpu.rs | 171 +++++++++++++++++- 3 files changed, 295 insertions(+), 7 deletions(-) create mode 100644 ext/crates/algebra/benches/nassau_milnor_gpu.rs diff --git a/ext/crates/algebra/Cargo.toml b/ext/crates/algebra/Cargo.toml index 09031e4c4b..0f4054e081 100644 --- a/ext/crates/algebra/Cargo.toml +++ b/ext/crates/algebra/Cargo.toml @@ -33,13 +33,13 @@ 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/JoeyBF/cubecl", branch = "claude/pool-slot-map-v0.10.0", optional = true, default-features = false, features = [ "cuda", ] } # For pinning all GPU work to one CUDA stream (`StreamId`), so a single memory pool is # reclaimed by `memory_cleanup` — CubeCL's pools are per-stream, and rayon spreads launches # across threads/streams, which otherwise accumulates buffers until the card OOMs. -cubecl-common = { version = "0.10.0", optional = true } +cubecl-common = { git = "https://github.com/JoeyBF/cubecl", branch = "claude/pool-slot-map-v0.10.0", optional = true } [dev-dependencies] criterion = { version = "0.5", features = ["html_reports"] } @@ -70,6 +70,12 @@ 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 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..7b22d83b6a --- /dev/null +++ b/ext/crates/algebra/benches/nassau_milnor_gpu.rs @@ -0,0 +1,121 @@ +//! 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). + let algebra = 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/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 427ce31f5f..4312b04f61 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -208,7 +208,22 @@ 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::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 — +/// the only recovery is to abandon the GPU multiply and finish the resolution on the CPU. +/// [`multiply_batch_on_gpu`] flips this on the first failure and routes all subsequent (and the +/// current) batches through [`cpu_multiply_batch`]. 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) +} /// 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. @@ -1357,6 +1372,86 @@ pub fn multiply_batch_on_gpu( num_cols: usize, num_rows: usize, products: &[GpuProduct], +) -> Vec> { + // STOPGAP (see [`GPU_DISABLED`]): once the cubecl CUDA context has been poisoned by the + // unresolved uninit-handle bug, every launch fails, so we finish the run on the CPU. On the + // first failure we catch the panic (it surfaces as an `.unwrap()` on a `CUDA_ERROR_LAUNCH_FAILED` + // / `ServerUnhealthy` in [`multiply_batch_gpu_inner`]), flip the flag, and fall back. All later + // calls short-circuit straight to the CPU path. The CPU result is bit-identical to the GPU's + // (validated by `cpu_multiply_batch_matches_gpu`), so callers see no difference but speed. + if gpu_disabled() { + return cpu_multiply_batch(algebra, num_cols, num_rows, products); + } + match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + multiply_batch_gpu_inner(algebra, num_cols, num_rows, products) + })) { + Ok(rows) => rows, + Err(_) => { + // compare_exchange so exactly one thread (of the ~100 that may fail together on the + // shared poisoned context) prints the notice; the rest just fall through to the CPU. + if GPU_DISABLED + .compare_exchange(false, true, Ordering::Relaxed, Ordering::Relaxed) + .is_ok() + { + eprintln!( + "[nassau-gpu] GPU milnor multiply failed (CUDA context poisoned by the \ + unresolved cubecl uninit-handle bug); disabling the GPU multiply and \ + completing the resolution on the CPU. RREF (separate fp-cuda runtime) is \ + unaffected by this flag." + ); + } + cpu_multiply_batch(algebra, num_cols, num_rows, products) + } + } +} + +/// 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> { + 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]; + 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 { + 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); + } + } + rows +} + +fn multiply_batch_gpu_inner( + algebra: &MilnorAlgebra, + num_cols: usize, + num_rows: usize, + products: &[GpuProduct], ) -> Vec> { let cap = resident_degree_cap(); // Fast path (default, `cap == i32::MAX`, and any run whose `R`s are all under the cap): a single @@ -2077,10 +2172,13 @@ fn multiply_batch_block( .map(|r| flat[r * num_limbs..(r + 1) * num_limbs].to_vec()) .collect(); - // Trim the pool: the per-launch uploads (`out_h` and the per-`R`/per-product metadata arrays) - // vary in size and would otherwise accumulate. The persistent buffers — the resident - // master/basis and the seqno `g`/`xi` — stay alive (refcount > 0), so cleanup skips them and - // only the transient per-launch memory is reclaimed. + // 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. client.memory_cleanup(); result @@ -3462,6 +3560,69 @@ mod tests { ); } + /// 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 = 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 From 1c09868681d1a3c81c25c6fe4f6b2eab610121f3 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Wed, 29 Jul 2026 06:14:07 -0400 Subject: [PATCH 037/127] =?UTF-8?q?milnor=5Fgpu:=20NASSAU=5FGPU=5FCLEANUP?= =?UTF-8?q?=5FEVERY=20knob=20=E2=80=94=20throttle=20per-launch=20memory=5F?= =?UTF-8?q?cleanup?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The residual high-stem CUDA_ERROR_LAUNCH_FAILED in the GPU multiply is a cross-stream pool-reclaim race: ~100 streams each calling client.memory_cleanup() every launch, one stream's cleanup reclaiming a shared resident-master page still in flight on another stream's kernel -> bad device pointer -> context poison (tracel-ai/cubecl#1401; the fork's per-command retain_until_complete does not cover cross-stream shared-page reclaim). Gate the cleanup call on NASSAU_GPU_CLEANUP_EVERY (default 1 = every launch, N = every Nth, 0 = never). Validated: with =0, S_2 stem-200 resolves fully clean on GPU (0 LAUNCH_FAILED, 0 dx) — the first clean GPU stem-200 — with device memory plateauing ~125 GB under cubecl's internal pressure-triggered reclaim (which does NOT hit the race). A moderate throttle (~16-32) is the likely stem-300 config: race-avoiding while bounding memory further. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 28 +++++++++++++++++++- 1 file changed, 27 insertions(+), 1 deletion(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 4312b04f61..56e723eeab 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -225,6 +225,27 @@ 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 = 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); @@ -2179,7 +2200,12 @@ fn multiply_batch_block( // 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. - client.memory_cleanup(); + // 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. + let every = cleanup_every(); + if every != 0 && CLEANUP_COUNTER.fetch_add(1, Ordering::Relaxed) % every == 0 { + client.memory_cleanup(); + } result }); From aafe902f7ca98a1555cdc4035b3932ee093e77de Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Thu, 30 Jul 2026 04:01:54 -0400 Subject: [PATCH 038/127] =?UTF-8?q?milnor=5Fgpu:=20concurrent-growth=20soa?= =?UTF-8?q?k=20#[test]=20=E2=80=94=20fast=20#1401=20reproducer=20+=20corre?= =?UTF-8?q?ctness=20guard?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit An #[ignore]d GPU soak: N worker threads hammer multiply_batch_on_gpu against one growing shared resident master, across NASSAU_GPU_STREAMS streams with per-launch memory_cleanup on — the cross-stream access pattern the stem-200 resolution drives. Each result is checked against the bit-identical cpu_multiply_batch oracle (up to NASSAU_SOAK_VERIFY_MAX, to keep the CPU precompute cheap); any mid-soak context death flips GPU_DISABLED and fails the run. Two jobs in one: - Correctness: catches cross-stream renumber/identity races in seconds. - #1401 reproducer (measured): clean at max_degree<=128, but at max_degree=160 with NASSAU_GPU_STREAMS=48 the cubecl cross-stream pool-reclaim race fires within a ~45s soak (never-initialized / ServerUnhealthy cascade, gpu_disabled flips) at only ~28 GB host / ~22 GB GPU — the genuine timing race, NOT an OOM. That's a ~1-2 min, low-memory stand-in for the 40-min stem-200 crash, and the gate the coming single-submission-thread redesign must turn GREEN. Repro: NASSAU_GPU_STREAMS=48 NASSAU_GPU_CLEANUP_EVERY=1 NASSAU_SOAK_MAX_DEGREE=160 \ cargo test -p algebra --release --features gpu -- --ignored --nocapture concurrent_growth_soak Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 187 +++++++++++++++++++ 1 file changed, 187 insertions(+) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 56e723eeab..21ce9065f6 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -3744,4 +3744,191 @@ mod tests { 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 (measured):** clean at `max_degree ≤ 128`, but at `max_degree=160` with + /// `NASSAU_GPU_STREAMS=48` the cross-stream pool-reclaim race fires within a ~45 s soak — the + /// `never initialized` / `ServerUnhealthy` cascade, `gpu_disabled` flips — at only ~28 GB host + /// / ~22 GB GPU, so it is NOT a device OOM but the genuine timing race. That makes this a + /// ~1–2 min, low-memory stand-in for the 40-min stem-200 crash (the race window scales with + /// buffer size; degree 64 was just below threshold). The single-submission-thread redesign + /// must turn this exact config GREEN. + /// + /// Ignored by default (needs a CUDA device + `NASSAU_GPU_STREAMS>1`). Reproduce #1401 with: + /// ```text + /// NASSAU_GPU_STREAMS=48 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 cross-stream soak: needs a CUDA device and NASSAU_GPU_STREAMS>1; 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; + + if gpu_stream_slots() == 1 { + eprintln!( + "[soak] WARNING: NASSAU_GPU_STREAMS=1 (single stream) — the cross-stream reclaim \ + race CANNOT reproduce. Set NASSAU_GPU_STREAMS >= {threads} to exercise it." + ); + } + + let p = fp::prime::ValidPrime::new(2); + let algebra = 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, {} streams: {n} launches ({:.0}/s over {} degrees, \ + verified ≤{verify_max}), {mm} correctness mismatches, gpu_disabled={}", + gpu_stream_slots(), + 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." + ); + } } From 9dbe8ae1c5d03b29e73109d6b9593e3203f22175 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Fri, 31 Jul 2026 22:41:49 -0400 Subject: [PATCH 039/127] =?UTF-8?q?milnor=5Fgpu:=20pivot=20cubecl=200.10?= =?UTF-8?q?=20fork=20=E2=86=92=20upstream=20v0.11.0-pre.1?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Drop the JoeyBF/cubecl `pool-slot-map-v0.10.0` fork pin for upstream `tracel-ai/cubecl` tag `v0.11.0-pre.1`, whose rewritten memory manager (+1377 lines) and CUDA stream layer are the async-engine/pool subsystem where the cross-stream reclaim race (#1401) lives. Migrate milnor_gpu.rs to the 0.11 launch API (227 one-line replacements): - launchable array kernel args `&Array`/`&mut Array` → slices `&[T]`/`&mut [T]` (incl. `&mut Array>` → `&mut [Atomic]`); - host-side `ArrayArg::from_raw_parts` → `BufferArg::from_raw_parts` (identical 2-arg signature, from `cubecl::prelude::*`); - local scratch `Array` passed to `#[cube]` helpers now go through `Array::as_slice()`; non-launch helper params take `&[T]`. Every `address_type = "u64"`/`"dynamic"` attribute and `launch_unchecked` call is preserved verbatim — the high-stem u64-addressing correctness path is untouched. Effect at the harsher-than-production d=160/48-stream soak: break rate 3/5 (fork) → 1/5 (0.11-pre), independent of NASSAU_GPU_CLEANUP_EVERY, with 0 correctness mismatches across all runs. The residual fault surfaces as a gentler CUDA_ERROR_ILLEGAL_ADDRESS at read_one (vs the old LAUNCH_FAILED uninit-handle cascade), always caught by the CPU multiply fallback. The race is reduced but not eliminated; the fallback remains the safety net. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/Cargo.toml | 4 +- ext/crates/algebra/src/algebra/milnor_gpu.rs | 454 +++++++++---------- 2 files changed, 229 insertions(+), 229 deletions(-) diff --git a/ext/crates/algebra/Cargo.toml b/ext/crates/algebra/Cargo.toml index 0f4054e081..821f38deda 100644 --- a/ext/crates/algebra/Cargo.toml +++ b/ext/crates/algebra/Cargo.toml @@ -33,13 +33,13 @@ 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 = { git = "https://github.com/JoeyBF/cubecl", branch = "claude/pool-slot-map-v0.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 = { git = "https://github.com/JoeyBF/cubecl", branch = "claude/pool-slot-map-v0.10.0", optional = true } +cubecl-common = { git = "https://github.com/tracel-ai/cubecl", tag = "v0.11.0-pre.1", optional = true } [dev-dependencies] criterion = { version = "0.5", features = ["html_reports"] } diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 21ce9065f6..de1155774f 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -878,7 +878,7 @@ fn ensure_basis(algebra: &MilnorAlgebra, width: usize, max_degree: i32) -> Vec) { +fn zero_u32(out: &mut [u32]) { if ABSOLUTE_POS < out.len() { out[ABSOLUTE_POS] = 0u32; } @@ -898,8 +898,8 @@ fn zero_u32(out: &mut Array) { // `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: &Array, - dst: &mut Array, + src: &[u16], + dst: &mut [u16], src_off: usize, dst_off: usize, count: u32, @@ -912,8 +912,8 @@ fn copy_into_u16( /// `u32` sibling of [`copy_into_u16`] (for the resident basis `lens`). #[cube(launch_unchecked, address_type = "dynamic")] fn copy_into_u32( - src: &Array, - dst: &mut Array, + src: &[u32], + dst: &mut [u32], src_off: usize, dst_off: usize, count: u32, @@ -952,8 +952,8 @@ macro_rules! copy_chunked { 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)), - ArrayArg::from_raw_parts($src.clone(), $src_len), - ArrayArg::from_raw_parts($dst.clone(), $dst_len), + 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, @@ -975,9 +975,9 @@ macro_rules! copy_chunked { /// `multiply_single_r_kernel` so both index outputs identically. #[cube] fn seqno_core( - g: &Array, - xi: &Array, - working: &Array, + g: &[u32], + xi: &[u32], + working: &[u32], wlen: usize, width: usize, ) -> u32 { @@ -1023,12 +1023,12 @@ fn seqno_core( #[cube] #[allow(clippy::too_many_arguments)] fn multiply_pair( - col_sums: &Array, - masks: &Array, - term_pparts: &Array, - g: &Array, - xi: &Array, - out: &mut Array>, + col_sums: &[u16], + masks: &[u16], + term_pparts: &[u16], + g: &[u32], + xi: &[u32], + out: &mut [Atomic], cs_base: usize, mk_base: usize, b_base: usize, @@ -1089,7 +1089,7 @@ fn multiply_pair( // `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 idx = seqno_core(g, xi, working.as_slice(), 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); @@ -1130,84 +1130,84 @@ fn multiply_pair( #[cube(launch_unchecked, address_type = "u64")] #[allow(clippy::too_many_arguments)] fn multiply_batch_kernel( - cs0: &Array, - cs1: &Array, - cs2: &Array, - cs3: &Array, - cs4: &Array, - cs5: &Array, - cs6: &Array, - cs7: &Array, - cs8: &Array, - cs9: &Array, - cs10: &Array, - cs11: &Array, - cs12: &Array, - cs13: &Array, - cs14: &Array, - cs15: &Array, - mk0: &Array, - mk1: &Array, - mk2: &Array, - mk3: &Array, - mk4: &Array, - mk5: &Array, - mk6: &Array, - mk7: &Array, - mk8: &Array, - mk9: &Array, - mk10: &Array, - mk11: &Array, - mk12: &Array, - mk13: &Array, - mk14: &Array, - mk15: &Array, - pp0: &Array, - pp1: &Array, - pp2: &Array, - pp3: &Array, - pp4: &Array, - pp5: &Array, - pp6: &Array, - pp7: &Array, - pp8: &Array, - pp9: &Array, - pp10: &Array, - pp11: &Array, - pp12: &Array, - pp13: &Array, - pp14: &Array, - pp15: &Array, - ln0: &Array, - ln1: &Array, - ln2: &Array, - ln3: &Array, - ln4: &Array, - ln5: &Array, - ln6: &Array, - ln7: &Array, - ln8: &Array, - ln9: &Array, - ln10: &Array, - ln11: &Array, - ln12: &Array, - ln13: &Array, - ln14: &Array, - ln15: &Array, - term_gei: &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], + prod_r_index: &[u32], + prod_term_start: &[u32], + prod_num_terms: &[u32], + prod_row_base: &[u32], + prod_out_offset: &[u32], + prod_pair_start: &[u32], width: usize, seg_elems: usize, ) { @@ -1341,9 +1341,9 @@ fn multiply_batch_kernel( } multiply_pair( - &cs_local, - &mk_local, - &term_local, + cs_local.as_slice(), + mk_local.as_slice(), + term_local.as_slice(), g, xi, out, @@ -1987,14 +1987,14 @@ fn multiply_batch_block( &client, CubeCount::Static((n_cold as u32).div_ceil(ENUM_THREADS).max(1), 1, 1), CubeDim::new_1d(ENUM_THREADS), - ArrayArg::from_raw_parts(epp_h, enum_pp.len()), - ArrayArg::from_raw_parts(er_h, n_cold), - ArrayArg::from_raw_parts(ec_h, n_cold), - ArrayArg::from_raw_parts(eco_h, n_cold), - ArrayArg::from_raw_parts(emo_h, n_cold), - ArrayArg::from_raw_parts(cs_scratch.clone(), cs_cap), - ArrayArg::from_raw_parts(mk_scratch.clone(), mk_cap), - ArrayArg::from_raw_parts(cnt_scratch, n_cold.max(1)), + BufferArg::from_raw_parts(epp_h, enum_pp.len()), + BufferArg::from_raw_parts(er_h, n_cold), + BufferArg::from_raw_parts(ec_h, n_cold), + BufferArg::from_raw_parts(eco_h, n_cold), + BufferArg::from_raw_parts(emo_h, n_cold), + BufferArg::from_raw_parts(cs_scratch.clone(), cs_cap), + BufferArg::from_raw_parts(mk_scratch.clone(), mk_cap), + BufferArg::from_raw_parts(cnt_scratch, n_cold.max(1)), enum_width, n_cold, ); @@ -2080,7 +2080,7 @@ fn multiply_batch_block( &client, CubeCount::Static((out_len as u32).div_ceil(THREADS).max(1), 1, 1), CubeDim::new_1d(THREADS), - ArrayArg::from_raw_parts(out_h.clone(), out_len), + BufferArg::from_raw_parts(out_h.clone(), out_len), ); } @@ -2091,10 +2091,10 @@ fn multiply_batch_block( let poo_h = client.create_from_slice(u32::as_bytes(&prod_out_offset)); let pps_h = client.create_from_slice(u32::as_bytes(&pps)); let cubes = (total_pairs as u32).div_ceil(THREADS).max(1); - // Bind one `ArrayArg` per `(segment vector, index)` — the `.0` handle, `.1` element length. + // Bind one `BufferArg` per `(segment vector, index)` — the `.0` handle, `.1` element length. macro_rules! sa { ($v:expr, $i:expr) => { - ArrayArg::from_raw_parts($v[$i].0.clone(), $v[$i].1) + BufferArg::from_raw_parts($v[$i].0.clone(), $v[$i].1) }; } // SAFETY: `launch_unchecked` — see the kernel's `address_type = "u64"` note. Every device @@ -2168,20 +2168,20 @@ fn multiply_batch_block( sa!(ln_seg, 13), sa!(ln_seg, 14), sa!(ln_seg, 15), - ArrayArg::from_raw_parts(tg_h, term_gei.len()), - ArrayArg::from_raw_parts(g_h, g.len()), - ArrayArg::from_raw_parts(xi_h, xi.len()), - ArrayArg::from_raw_parts(out_h.clone(), out_len), - ArrayArg::from_raw_parts(rco_h, r_cs_offset.len()), - ArrayArg::from_raw_parts(rmo_h, r_mk_offset.len()), - ArrayArg::from_raw_parts(rcl_h, r_cs_len.len()), - ArrayArg::from_raw_parts(rml_h, r_mk_len.len()), - ArrayArg::from_raw_parts(pri_h, products.len()), - ArrayArg::from_raw_parts(pts_h, products.len()), - ArrayArg::from_raw_parts(pnt_h, products.len()), - ArrayArg::from_raw_parts(prb_h, products.len()), - ArrayArg::from_raw_parts(poo_h, products.len()), - ArrayArg::from_raw_parts(pps_h, pps.len()), + 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(pri_h, products.len()), + BufferArg::from_raw_parts(pts_h, products.len()), + BufferArg::from_raw_parts(pnt_h, products.len()), + BufferArg::from_raw_parts(prb_h, products.len()), + BufferArg::from_raw_parts(poo_h, products.len()), + BufferArg::from_raw_parts(pps_h, pps.len()), width, seg_elems, ); @@ -2238,22 +2238,22 @@ fn multiply_batch_block( #[cube] #[allow(clippy::too_many_arguments)] fn seg_read_u16( - s0: &Array, - s1: &Array, - s2: &Array, - s3: &Array, - s4: &Array, - s5: &Array, - s6: &Array, - s7: &Array, - s8: &Array, - s9: &Array, - s10: &Array, - s11: &Array, - s12: &Array, - s13: &Array, - s14: &Array, - s15: &Array, + 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, ) -> u16 { @@ -2301,22 +2301,22 @@ fn seg_read_u16( #[cube] #[allow(clippy::too_many_arguments)] fn seg_read_u32( - s0: &Array, - s1: &Array, - s2: &Array, - s3: &Array, - s4: &Array, - s5: &Array, - s6: &Array, - s7: &Array, - s8: &Array, - s9: &Array, - s10: &Array, - s11: &Array, - s12: &Array, - s13: &Array, - s14: &Array, - s15: &Array, + 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, ) -> u32 { @@ -2385,16 +2385,16 @@ fn seg_read_u32( #[cube(launch_unchecked, address_type = "u64")] #[allow(clippy::too_many_arguments)] fn enumerate_admissible_kernel( - p_parts: &Array, - r_rows: &Array, - r_cols: &Array, + 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: &Array, - r_mk_out: &Array, - out_cs: &mut Array, - out_mk: &mut Array, - out_counts: &mut Array, + 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, ) { @@ -2532,7 +2532,7 @@ mod tests { /// One thread per `u32` limb. F₂ addition is XOR of the packed limbs, so this is /// the output primitive the multiply kernels accumulate with. #[cube(launch)] - fn xor_f2(a: &Array, b: &Array, out: &mut Array) { + fn xor_f2(a: &[u32], b: &[u32], out: &mut [u32]) { if ABSOLUTE_POS < out.len() { out[ABSOLUTE_POS] = a[ABSOLUTE_POS] ^ b[ABSOLUTE_POS]; } @@ -2559,9 +2559,9 @@ mod tests { &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), + BufferArg::from_raw_parts(a_handle, n), + BufferArg::from_raw_parts(b_handle, n), + BufferArg::from_raw_parts(out_handle.clone(), n), ); } @@ -2574,10 +2574,10 @@ mod tests { /// are zero and skipped, so `wlen == width` matches the CPU's trimmed loop). #[cube(launch)] fn seqno_kernel( - g: &Array, - xi: &Array, - p_parts: &Array, - out: &mut Array, + g: &[u32], + xi: &[u32], + p_parts: &[u32], + out: &mut [u32], width: usize, ) { let idx = ABSOLUTE_POS; @@ -2590,7 +2590,7 @@ mod tests { for h in 0..width { working[h] = p_parts[base + h]; } - out[idx] = seqno_core(g, xi, &working, width, width); + out[idx] = seqno_core(g, xi, working.as_slice(), width, width); } /// Run `seqno_kernel` over `n` padded p_parts and return their seqno indices. @@ -2621,10 +2621,10 @@ mod tests { &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), + 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, ); } @@ -2638,13 +2638,13 @@ mod tests { #[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>, + 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, @@ -2753,13 +2753,13 @@ mod tests { &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), + 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, @@ -2777,24 +2777,24 @@ mod tests { #[cube(launch)] #[allow(clippy::too_many_arguments)] fn seg_gather_kernel( - s0: &Array, - s1: &Array, - s2: &Array, - s3: &Array, - s4: &Array, - s5: &Array, - s6: &Array, - s7: &Array, - s8: &Array, - s9: &Array, - s10: &Array, - s11: &Array, - s12: &Array, - s13: &Array, - s14: &Array, - s15: &Array, - idx: &Array, - out: &mut Array, + 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, ) { let i = ABSOLUTE_POS; @@ -2884,14 +2884,14 @@ mod tests { &client, CubeCount::Static(cubes, 1, 1), CubeDim::new_1d(THREADS), - ArrayArg::from_raw_parts(pp_h, pp_flat.len()), - ArrayArg::from_raw_parts(rr_h, n_r), - ArrayArg::from_raw_parts(rc_h, n_r), - ArrayArg::from_raw_parts(rco_h, n_r), - ArrayArg::from_raw_parts(rmo_h, n_r), - ArrayArg::from_raw_parts(ocs_h.clone(), cs_cap), - ArrayArg::from_raw_parts(omk_h.clone(), mk_cap), - ArrayArg::from_raw_parts(cnt_h.clone(), n_r), + 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, ); @@ -3055,7 +3055,7 @@ mod tests { const THREADS: u32 = 256; let cubes = (indices.len() as u32).div_ceil(THREADS).max(1); - let arg = |i: usize| unsafe { ArrayArg::from_raw_parts(handles[i].clone(), lens[i]) }; + let arg = |i: usize| unsafe { BufferArg::from_raw_parts(handles[i].clone(), lens[i]) }; unsafe { seg_gather_kernel::launch::( &client, @@ -3077,8 +3077,8 @@ mod tests { arg(13), arg(14), arg(15), - ArrayArg::from_raw_parts(idx_h, indices.len()), - ArrayArg::from_raw_parts(out_h.clone(), indices.len()), + BufferArg::from_raw_parts(idx_h, indices.len()), + BufferArg::from_raw_parts(out_h.clone(), indices.len()), seg_elems, ); } @@ -3200,14 +3200,14 @@ mod tests { &client, CubeCount::Static(cubes, 1, 1), CubeDim::new_1d(THREADS), - ArrayArg::from_raw_parts(pp_h, pp_flat.len()), - ArrayArg::from_raw_parts(rr_h, n_r), - ArrayArg::from_raw_parts(rc_h, n_r), - ArrayArg::from_raw_parts(rco_h, n_r), - ArrayArg::from_raw_parts(rmo_h, n_r), - ArrayArg::from_raw_parts(ocs_h.clone(), cs_cap), - ArrayArg::from_raw_parts(omk_h.clone(), mk_cap), - ArrayArg::from_raw_parts(cnt_h.clone(), n_r), + 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, ); From 3f86eb7f5ab6a4ec4d769a6e7b8c3c42f7d3692e Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Sat, 1 Aug 2026 13:21:10 -0400 Subject: [PATCH 040/127] milnor_gpu: guard the multiply output write against OOB atomic (out[word]) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit multiply_pair computed word = row_base + (out_offset + seqno)/32 and did an unguarded atomic XOR into out[word]. When out_offset + seqno spans past the row's num_limbs (which nassau_gpu::get_partial_matrix_restricted already anticipates and masks on readback — bits >= target_dim are dropped), the device write overran: into the next row (silent corruption) or past the buffer (CUDA_ERROR_ILLEGAL_ADDRESS). compute-sanitizer on a malloc_sync build confirmed "Invalid __global__ atomic of size 4 bytes ... out of bounds". Thread num_limbs through multiply_batch_kernel and skip writes with global_bit/32 >= num_limbs — a device-side mirror of the host's existing defensive mask. Correctness-preserving; the skipped bits are exactly the ones the host discards. The d=160/48-stream concurrent soak that reliably broke ~3-5/5 now runs 0/5 with 0 correctness mismatches. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 21 +++++++++++++++++--- 1 file changed, 18 insertions(+), 3 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index de1155774f..535a40f942 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -1038,6 +1038,7 @@ fn multiply_pair( row_base: usize, out_offset: usize, width: usize, + num_limbs: usize, ) { let mut low = cs_len; if term_len < cs_len { @@ -1091,9 +1092,17 @@ fn multiply_pair( // output). Both are bit offsets, added before splitting into (limb, bit). let idx = seqno_core(g, xi, working.as_slice(), 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); + 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`). + if limb < num_limbs { + let word = row_base + limb; + let bit = u32::cast_from(global_bit % 32); + out[word].fetch_xor(1u32 << bit); + } } } @@ -1210,6 +1219,7 @@ fn multiply_batch_kernel( prod_pair_start: &[u32], width: usize, seg_elems: usize, + num_limbs: usize, ) { let k = ABSOLUTE_POS; let num_products = prod_pair_start.len() - 1; @@ -1356,6 +1366,7 @@ fn multiply_batch_kernel( usize::cast_from(prod_row_base[p]), usize::cast_from(prod_out_offset[p]), width, + num_limbs, ); } @@ -2184,6 +2195,7 @@ fn multiply_batch_block( BufferArg::from_raw_parts(pps_h, pps.len()), width, seg_elems, + num_limbs, ); } @@ -2650,6 +2662,7 @@ mod tests { cs_len: usize, mk_len: usize, width: usize, + num_limbs: usize, ) { let pair = ABSOLUTE_POS; if pair >= num_matrices * num_terms { @@ -2674,6 +2687,7 @@ mod tests { 0, 0, width, + num_limbs, ); } @@ -2765,6 +2779,7 @@ mod tests { cs_len, mk_len, width, + num_limbs, ); } From bbf7c4c0f771d03f110a4537070fe251b5fa3a6e Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 31 Jul 2026 05:59:15 +0000 Subject: [PATCH 041/127] Bit-pack Milnor basis elements into a u64 The p-part of a Milnor basis element was a `Vec`, costing a heap allocation and a pointer chase per element. At p = 2 the internal degree of P(R) is sum_i r_i (2^i - 1) with non-negative terms, so r_i <= deg/(2^i - 1); sizing each field by that bound packs the whole exponent sequence into 64 bits for every degree up to 2045. At odd primes the same bound applies divided by q = 2(p-1), so one layout serves every prime. `MilnorBasisElement` is now 16 bytes, `Copy`, and entirely inline. Measured over degrees 0..=300 at p = 2, `basis_table` drops from 252 MiB in 5,036,688 allocations to 77 MiB in none. Three things fall out of the packing: - The packed value is a canonical key, so the hand-rolled `MilnorHashMap` specialization for `not(odd-primes)` is gone; a plain `HashMap` now hashes a single word on every path. That code also assumed a degree bound of 1536 without enforcing it. `compute_basis` now asserts the bound up front, which is what lets everything downstream skip range checks. - Trailing zeros are not represented, so the "pop trailing zeros" loops after building a product disappear. - `PPartMultiplier` no longer borrows its inputs, so its lifetime parameter is gone, and `PPartAllocation` loses the buffer it existed to recycle. In `ext`, `MilnorSubalgebra`'s signature test becomes one masked comparison on the packed word instead of a loop over entries, with the mask hoisted out of `signature_mask`'s inner loop. Two behaviour changes worth noting: - `basis_element_from_string("P0")` and `("Sq0")` now return the identity rather than `None`. P(0) is the identity, and `AdemAlgebra::try_beps_pn` already special-cases `x == 0` this way; the old `None` came from `vec![0]` and `vec![]` hashing differently, an artifact of the representation. - `increment_p_part` now carries before incrementing. The old order transiently stored `max[i] + 1`, which need not fit a field whose width is exactly saturated by `max[i]`. The enumeration is unchanged. The observation that every Milnor exponent sequence up to degree 512 fits in 64 bits is due to Lixiong Wu; this implementation works out the widths, finds that the same layout holds all the way to degree 2045, and carries it through the algebra. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01XK8YCdHkCqUV9YM7D957hV --- ext/crates/algebra/benches/milnor.rs | 32 +- .../algebra/src/algebra/milnor_algebra.rs | 733 ++++++++++++------ .../algebra/src/algebra/pair_algebra.rs | 38 +- ext/crates/algebra/src/module/rpn.rs | 2 +- ext/crates/algebra/src/steenrod_evaluator.rs | 29 +- ext/crates/algebra/src/steenrod_parser.rs | 4 +- ext/examples/bruner.rs | 7 +- ext/examples/sq0.rs | 9 +- ext/src/nassau.rs | 31 +- ext/src/yoneda.rs | 2 +- 10 files changed, 588 insertions(+), 299 deletions(-) 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/src/algebra/milnor_algebra.rs b/ext/crates/algebra/src/algebra/milnor_algebra.rs index 07b3bd35f8..bbf780041c 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,207 @@ 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 [`Self::WIDTHS`]`[i]` bits starting at bit [`Self::SHIFTS`]`[i]`. 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 { + /// The largest internal degree whose exponent sequences are guaranteed to fit. + /// + /// This is the largest bound for which [`Self::WIDTHS`] sums 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; + + /// `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::MAX_LEN] = [11, 10, 9, 8, 7, 6, 5, 4, 3, 1]; + + /// `SHIFTS[i]` is the bit offset of entry `i`; `SHIFTS[MAX_LEN]` is the total width, 64. + const SHIFTS: [u32; Self::MAX_LEN + 1] = { + let mut shifts = [0; Self::MAX_LEN + 1]; + let mut i = 0; + while i < Self::MAX_LEN { + shifts[i + 1] = shifts[i] + Self::WIDTHS[i]; + i += 1; + } + shifts + }; + + /// `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 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::has_signature` in `ext`). + pub const fn bits(self) -> u64 { + self.0 + } + + /// Entry `i`, or 0 if `i` is past the end. + #[inline] + pub const fn get(self, i: usize) -> PPartEntry { + if i >= Self::MAX_LEN { + return 0; + } + ((self.0 >> Self::SHIFTS[i]) & ((1 << Self::WIDTHS[i]) - 1)) as PPartEntry + } + + /// 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 +327,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 +336,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 +359,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 +389,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)) - } -} - pub struct MilnorAlgebra { profile: MilnorProfile, p: ValidPrime, @@ -371,7 +522,7 @@ impl Algebra for MilnorAlgebra { MilnorBasisElement { degree: 1, q_part: 1, - p_part: vec![], + p_part: PPart::zero(), }, )); } @@ -383,7 +534,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]), }, )); } @@ -402,7 +553,7 @@ impl Algebra for MilnorAlgebra { MilnorBasisElement { degree, q_part: 0, - p_part: vec![1 << i], + p_part: PPart::from_iter([1 << i]), }, )); } @@ -417,6 +568,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() { @@ -432,7 +590,7 @@ impl Algebra for MilnorAlgebra { let mut map = MilnorHashMap::default(); map.reserve(basis.len()); for (i, b) in basis.iter().enumerate() { - map.insert(b.clone(), i); + assert!(map.insert(*b, i).is_none(), "Duplicate entry for {b}"); } map }); @@ -585,19 +743,29 @@ 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 > PPart::MAX_LEN { + return None; + } + let entry = p.pow(s) as PPartEntry; + let degree = entry as i32 * self.q() * combinatorics::xi_degrees(p)[t]; + if degree > PPart::MAX_DEGREE || entry > PPart::max_entry(t - 1) { + 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::)), @@ -608,12 +776,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) @@ -715,7 +887,7 @@ impl GeneratedAlgebra for MilnorAlgebra { return vec![self.basis_element_to_index(&MilnorBasisElement { degree, q_part, - p_part: vec![], + p_part: PPart::zero(), })]; } } @@ -734,7 +906,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![] @@ -753,8 +925,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, @@ -833,7 +1005,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 }; @@ -863,20 +1035,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); } } @@ -918,8 +1089,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, }), @@ -936,7 +1107,7 @@ impl MilnorAlgebra { 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)); @@ -973,11 +1144,14 @@ impl MilnorAlgebra { pub fn try_beps_pn(&self, e: u32, x: PPartEntry) -> Option<(i32, usize)> { let q = self.q() as u32; let degree = (q * x + e) as i32; + if degree > PPart::MAX_DEGREE || x > PPart::max_entry(0) { + 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)) } @@ -988,7 +1162,7 @@ impl MilnorAlgebra { } fn multiply_qpart(&self, m1: &MilnorBasisElement, f: u32) -> Vec<(u32, MilnorBasisElement)> { - let mut new_result: Vec<(u32, MilnorBasisElement)> = vec![(1, m1.clone())]; + 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) { @@ -1011,23 +1185,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, @@ -1074,8 +1246,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, @@ -1092,8 +1264,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, @@ -1149,7 +1321,7 @@ impl MilnorAlgebra { #[derive(Debug, Default)] struct Matrix2D { cols: usize, - inner: PPart, + inner: Vec, } impl std::fmt::Display for Matrix2D { @@ -1202,8 +1374,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! { @@ -1218,9 +1389,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), } } @@ -1232,21 +1400,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 } @@ -1254,8 +1422,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, @@ -1276,20 +1444,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.get(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.get(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, @@ -1308,7 +1474,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, } } @@ -1392,7 +1557,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.get(row - 1); for col in 1..self.cols { self.M[0][col] += self.M[row][col]; self.M[row][col] = 0; @@ -1416,13 +1581,13 @@ 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(); + self.ans.p_part = PPart::zero(); let mut coef = 1; if self.init { @@ -1443,27 +1608,19 @@ 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. Writing a zero + // is a no-op on a packed p-part, so trailing zeros need no trimming. + for i in 0..std::cmp::max(self.cols, self.rows) - 1 { + let mut entry = self.r.get(i); + if i + 1 < self.cols { + entry += self.M[0][i + 1]; + } + if entry != 0 { + self.ans.p_part.set(i, entry); + } } return Some(coef); } else if self.update() { - self.ans.p_part.reserve(self.diag_num); 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); @@ -1510,11 +1667,12 @@ 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 sum != 0 { + self.ans.p_part.set(diag_idx - 1, sum); + } } return Some(coef); @@ -1543,7 +1701,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(); @@ -1551,7 +1709,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(); @@ -1567,13 +1725,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, }); @@ -1607,9 +1765,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 @@ -1625,12 +1783,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(); @@ -1654,15 +1812,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 = @@ -1671,22 +1829,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)]); }; @@ -1708,16 +1862,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 } } @@ -1730,7 +1891,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; @@ -1738,32 +1899,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, @@ -1920,13 +2072,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); @@ -2001,6 +2157,133 @@ 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); + } + } + } + #[test] fn test_clone_into() { let mut other = MilnorBasisElement::default(); @@ -2012,34 +2295,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, @@ -2075,12 +2358,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/pair_algebra.rs b/ext/crates/algebra/src/algebra/pair_algebra.rs index e080d34116..23ab93ae6e 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() { @@ -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,7 +263,7 @@ 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(), @@ -271,7 +272,8 @@ impl PairAlgebra for MilnorAlgebra { 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 @@ -385,7 +387,7 @@ 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); } } }) @@ -393,11 +395,11 @@ fn a_y_cached( /// 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(); + 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,11 +412,9 @@ 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; @@ -445,7 +445,7 @@ mod tests { MilnorElt { q_part: 0, - p_part: p_part.into(), + p_part: PPart::from_slice(p_part), degree, } } diff --git a/ext/crates/algebra/src/module/rpn.rs b/ext/crates/algebra/src/module/rpn.rs index 6d42a1f9f8..5d2b3f43f0 100644 --- a/ext/crates/algebra/src/module/rpn.rs +++ b/ext/crates/algebra/src/module/rpn.rs @@ -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/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 8d1919136a..5135aabeec 100644 --- a/ext/src/nassau.rs +++ b/ext/src/nassau.rs @@ -20,7 +20,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}, @@ -100,15 +100,23 @@ 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. + fn packed_signature(&self, signature: &[PPartEntry]) -> (u64, u64) { + let mut mask = 0; + let mut value = 0; + for (i, (&profile, &entry)) in self.profile.iter().zip(signature).enumerate() { + // A profile wider than the field constrains the whole field. + let width = std::cmp::min(profile as u32, PPart::width(i)); + mask |= ((1u64 << width) - 1) << PPart::shift(i); + value |= (entry as u64) << PPart::shift(i); } - true + (mask, value) } fn zero_signature(&self) -> Vec { @@ -131,12 +139,15 @@ impl MilnorSubalgebra { start: [offset], end: _, }| { + // Hoist the mask out of the inner loop: every element in this block is tested + // against the same signature. + let (mask, value) = self.packed_signature(signature); algebra .ppart_table(degree - gen_deg) .iter() .enumerate() .filter_map(move |(n, op)| { - if self.has_signature(op, signature) { + if op.bits() & mask == value { Some(offset + n) } else { None 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 } From 769fbcc553820f68a4ba5e1b752f9b7bab75edfc Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 31 Jul 2026 05:59:15 +0000 Subject: [PATCH 042/127] Cut per-entry overhead out of the Milnor multiplier The first packing pass regressed `milnor_ppart` by up to 8% at odd primes and mod 4, because assembling the answer went from a memcpy plus a vectorized add to a per-entry read-modify-write through the checked `PPart::set`, and because `PPart::get`'s range branch landed in `update`'s inner loop. Two changes, both confined to the kernel: - Assemble the answer in a plain `u64` and store it once. Entries are written in increasing index order into a value that starts at zero, so a shift and an `or` suffice; the range checks become debug assertions backed by `compute_basis`'s degree gate. - Pad the layout tables to 16 entries so the private `PPart::entry` can mask its index rather than branch on it. Padded entries have width zero and so read as zero, which is the answer `get` would have returned anyway. The public `get` keeps its explicit check, since callers outside the multiplier index it with a q-part-derived length that is not bounded by `MAX_LEN`. This recovers the regression (`ppart_4/a` and `ppart_3/a` back to baseline, `ppart_4/b` -8%) and improves the Nassau regime further. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01XK8YCdHkCqUV9YM7D957hV --- .../algebra/src/algebra/milnor_algebra.rs | 70 ++++++++++++++----- 1 file changed, 52 insertions(+), 18 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_algebra.rs b/ext/crates/algebra/src/algebra/milnor_algebra.rs index bbf780041c..3359b1b3e3 100644 --- a/ext/crates/algebra/src/algebra/milnor_algebra.rs +++ b/ext/crates/algebra/src/algebra/milnor_algebra.rs @@ -138,13 +138,18 @@ impl PPart { /// already bounds the length of the $\xi$-degree table, so it is not a new restriction. pub const MAX_LEN: usize = 10; + /// 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::MAX_LEN] = [11, 10, 9, 8, 7, 6, 5, 4, 3, 1]; + const WIDTHS: [u32; Self::TABLE_LEN] = [11, 10, 9, 8, 7, 6, 5, 4, 3, 1, 0, 0, 0, 0, 0, 0]; /// `SHIFTS[i]` is the bit offset of entry `i`; `SHIFTS[MAX_LEN]` is the total width, 64. - const SHIFTS: [u32; Self::MAX_LEN + 1] = { - let mut shifts = [0; Self::MAX_LEN + 1]; + 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]; @@ -195,18 +200,41 @@ impl PPart { /// 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::has_signature` in `ext`). + /// `MilnorSubalgebra::packed_signature` in `ext`). pub const fn bits(self) -> u64 { self.0 } - /// Entry `i`, or 0 if `i` is past the end. + /// 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.0 >> Self::SHIFTS[i]) & ((1 << Self::WIDTHS[i]) - 1)) as PPartEntry + self.entry(i) } /// Set entry `i` to `v`. @@ -1444,10 +1472,10 @@ impl PPartMultiplier { M.reset(rows, cols); for i in 1..rows { - M[i][0] = r.get(i - 1); + M[i][0] = r.entry(i - 1); } for k in 1..cols { - M[0][k] = s.get(k - 1); + M[0][k] = s.entry(k - 1); } let ans = MilnorBasisElement { @@ -1557,7 +1585,7 @@ impl PPartMultiplier { if inc <= max_inc { // If so, we found our next matrix. for row in 1..i { - self.M[row][0] = self.r.get(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; @@ -1587,7 +1615,6 @@ impl Iterator for PPartMultiplier { fn next(&mut self) -> Option { let p = self.prime().as_u32() as PPartEntry; 'outer: loop { - self.ans.p_part = PPart::zero(); let mut coef = 1; if self.init { @@ -1608,19 +1635,22 @@ impl Iterator for PPartMultiplier { continue 'outer; } } - // The answer is the top row of the matrix plus `r`, entrywise. Writing a zero - // is a no-op on a packed p-part, so trailing zeros need no trimming. + // 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.get(i); + let mut entry = self.r.entry(i); if i + 1 < self.cols { entry += self.M[0][i + 1]; } - if entry != 0 { - self.ans.p_part.set(i, entry); - } + 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() { + 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); @@ -1670,10 +1700,14 @@ impl Iterator for PPartMultiplier { // `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 sum != 0 { - self.ans.p_part.set(diag_idx - 1, sum); + 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 { From fb242b8fe63dffee78925094f2ef4543c460c4b6 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 31 Jul 2026 17:32:41 +0000 Subject: [PATCH 043/127] Derive the p=2 Milnor basis instead of storing it `basis_table` held a `MilnorBasisElement` per basis element. At p = 2 with unstable support off, that element is exactly `from_p(ppart_table[t][i], t)` -- the p-part again, with a q-part that is always zero and a degree that is the index. It was a redundant copy. Deriving it on demand costs nothing now that `MilnorBasisElement` is `Copy` and 16 bytes: `basis_element_from_index` returns by value and builds it in registers rather than handing out a reference into a table. The multiply family takes the element by value for the same reason. The table is still built at odd primes, where the q-part varies within a degree, and when unstable support is on, where the basis is re-sorted by excess. Neither is a re-wrapping of `ppart_table`. Measured over degrees 0..=250 at p = 2 (1,958,958 elements), RSS growth from `compute_basis` drops 125.0 MB -> 95.0 MB, i.e. 66.9 -> 50.8 bytes per element. Projected to degree 500 that is 5.24 GB -> 3.95 GB. Unlike the ranker, this needs no basis renumbering and costs nothing at lookup time. A test verifies the derivation matches what the table used to hold, for every element, so the redundancy is asserted rather than assumed. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01XK8YCdHkCqUV9YM7D957hV --- .../algebra/src/algebra/milnor_algebra.rs | 124 ++++++++++++++---- .../algebra/src/algebra/pair_algebra.rs | 26 ++-- ext/crates/algebra/src/module/rpn.rs | 2 +- 3 files changed, 114 insertions(+), 38 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_algebra.rs b/ext/crates/algebra/src/algebra/milnor_algebra.rs index 3359b1b3e3..ad55e988b8 100644 --- a/ext/crates/algebra/src/algebra/milnor_algebra.rs +++ b/ext/crates/algebra/src/algebra/milnor_algebra.rs @@ -436,7 +436,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>, @@ -502,8 +509,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 { @@ -614,11 +634,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() { - assert!(map.insert(*b, i).is_none(), "Duplicate entry for {b}"); + 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 }); @@ -639,8 +660,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 }) @@ -660,7 +681,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"))] @@ -837,7 +862,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) } } @@ -1132,14 +1157,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, 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 }); } @@ -1189,8 +1216,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)]; + 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) { @@ -1251,8 +1278,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) @@ -1263,8 +1290,8 @@ impl MilnorAlgebra { &self, mut res: FpSliceMut, coef: u32, - m1: &MilnorBasisElement, - m2: &MilnorBasisElement, + m1: MilnorBasisElement, + m2: MilnorBasisElement, excess: i32, mut allocation: PPartAllocation, ) -> PPartAllocation { @@ -1314,7 +1341,7 @@ impl MilnorAlgebra { &self, res: FpSliceMut, coef: u32, - m1: &MilnorBasisElement, + m1: MilnorBasisElement, s_deg: i32, s: FpSlice, ) { @@ -1327,7 +1354,7 @@ impl MilnorAlgebra { &self, mut res: FpSliceMut, coef: u32, - m1: &MilnorBasisElement, + m1: MilnorBasisElement, s_deg: i32, s: FpSlice, mut allocation: PPartAllocation, @@ -2309,15 +2336,64 @@ mod tests { PPart::from_slice(&elt.p_part.iter().collect::>()), elt.p_part ); - assert_eq!(algebra.basis_element_to_index(elt), i); + assert_eq!(algebra.basis_element_to_index(&elt), i); // The degree really is recoverable from the entries. - let mut recomputed = *elt; + let mut recomputed = elt; recomputed.compute_degree(ValidPrime::new(p)); assert_eq!(recomputed.degree, t); } } } + /// 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(); diff --git a/ext/crates/algebra/src/algebra/pair_algebra.rs b/ext/crates/algebra/src/algebra/pair_algebra.rs index 23ab93ae6e..34815baaa0 100644 --- a/ext/crates/algebra/src/algebra/pair_algebra.rs +++ b/ext/crates/algebra/src/algebra/pair_algebra.rs @@ -192,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); - let mut s = *self.basis_element_from_index(s_degree, s_idx); + 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() { @@ -205,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, ); @@ -263,12 +263,12 @@ 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); + 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(), ); @@ -366,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), @@ -379,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 { @@ -387,15 +387,15 @@ fn a_y_cached( None => { let v = a_y_inner(algebra, a, k, l); f(&v); - cache.insert((*a, (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; +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, @@ -420,7 +420,7 @@ fn a_y_inner(algebra: &MilnorAlgebra, a: &MilnorElt, k: usize, l: usize) -> FpVe // 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); } @@ -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/rpn.rs b/ext/crates/algebra/src/module/rpn.rs index 5d2b3f43f0..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 { From 6ac8173bd1a2c2828d020603e2c27260eed26f36 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 31 Jul 2026 17:32:41 +0000 Subject: [PATCH 044/127] Add PPartRanker behind an opt-in feature, not wired in `basis_element_to_index` runs once per term of every product, and is a hash map storing an entry per basis element. The index it returns is a position in an enumeration, so a canonical key alone cannot replace the map -- but the position can be computed. Let counts[i][d] be the number of exponent sequences of degree d using only xi_1..xi_i. Splitting on whether r_i is zero gives the coin-change recurrence counts[i][d] = counts[i-1][d] + counts[i][d - xi_i]. Ranking needs the number of sequences with r_i > v, and substituting r_i -> r_i - (v+1) is a bijection onto all sequences of degree d - (v+1)*xi_i, so that count is a single table lookup rather than a sum. Walking the entries downward ranks a p-part in one lookup each, against a table covering every degree at once, where the map it would replace grows with the basis. Whether that is worth it depends entirely on scale, which took some measuring to see. Against the map, at p = 2: 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 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. Once it does not -- the map is 37 MB in degree 500 -- every probe misses to DRAM at ~33 ns, whereas the ranker's table is ~43 KB, stays in L1, and costs ~16 ns regardless of degree. Benchmarking only up to degree 120, where the map is 0.1 MB, shows a 2x loss and hides all of this; the sweep here deliberately spans the crossover. Tuning does not move the small-degree end: nested vs flat table, a zero-padded prefix to drop the branch, and one- vs two-pass to break the dependency chain were all measured, and the padded variant was worst, because doubling the table pushed it out of L1. So the two suit opposite ends of the range, and the ranker is on the right side of the end where the algebra's memory is the problem worth solving: replacing the map there is 3.3 GB smaller and 2x faster. It stays off by default and unwired even when enabled, because it numbers the basis in colex order rather than the order compute_ppart emits, which would invalidate saved resolutions. That order is rankable in principle, but its natural recursion has depth equal to the sum of the entries, which is worse than hashing. The unstable path, which re-sorts each degree by excess, is not modelled either. Tests verify the table reproduces the algebra's own p-part counts and that the rank is a bijection onto 0..dim in every degree, at p = 2 and p = 3, plus one pinning down that it really does disagree with the current basis order. Also measured and rejected: an `unrank` recovering the p-part at a given index, which would let basis_element_from_index drop ppart_table entirely. It ran ~15x slower than the array read it would replace, at every degree, with none of the crossover above -- ppart_table is 8 bytes per element against ~43 for the map, so it stays cache-resident. The bit-packing that makes rank worth having is the same thing that makes unrank not. The likelier route, if it is ever revisited, is enumerating the basis in index order, which is O(1) amortised and matches how callers actually walk it. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01XK8YCdHkCqUV9YM7D957hV --- ext/crates/algebra/Cargo.toml | 9 + ext/crates/algebra/benches/milnor_rank.rs | 100 +++++++ ext/crates/algebra/src/algebra/milnor_rank.rs | 265 ++++++++++++++++++ ext/crates/algebra/src/algebra/mod.rs | 5 + 4 files changed, 379 insertions(+) create mode 100644 ext/crates/algebra/benches/milnor_rank.rs create mode 100644 ext/crates/algebra/src/algebra/milnor_rank.rs diff --git a/ext/crates/algebra/Cargo.toml b/ext/crates/algebra/Cargo.toml index 3b98825cb5..f40a7c1cc5 100644 --- a/ext/crates/algebra/Cargo.toml +++ b/ext/crates/algebra/Cargo.toml @@ -37,6 +37,10 @@ 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"] odd-primes = ["fp/odd-primes"] @@ -55,3 +59,8 @@ harness = false [[bench]] name = "nassau_milnor" harness = false + +[[bench]] +name = "milnor_rank" +harness = false +required-features = ["milnor-rank"] 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/src/algebra/milnor_rank.rs b/ext/crates/algebra/src/algebra/milnor_rank.rs new file mode 100644 index 0000000000..e0b592e19b --- /dev/null +++ b/ext/crates/algebra/src/algebra/milnor_rank.rs @@ -0,0 +1,265 @@ +//! 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. +//! +//! **This is not wired into [`MilnorAlgebra`].** 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. + +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 67baed2b04..c2425320af 100644 --- a/ext/crates/algebra/src/algebra/mod.rs +++ b/ext/crates/algebra/src/algebra/mod.rs @@ -18,6 +18,11 @@ pub use field::Field; pub mod milnor_algebra; pub use milnor_algebra::MilnorAlgebra; +/// 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}; From 99ae117f7031672c5e1e04d7f6361e889e584f00 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 1 Aug 2026 18:06:52 +0000 Subject: [PATCH 045/127] Reject inputs the packing cannot represent, instead of panicking Three paths computed with unvalidated input before checking it against the packing bounds, so the intermediate arithmetic went wrong first. All three are reachable from public, non-panicking entry points. - `basis_element_from_string("P^s_t")` indexed the xi-degree table with `t`, which has exactly `MAX_LEN` entries, so `t = MAX_LEN` was out of bounds. `p^s` and the degree product could also overflow. Now `t` is bounded by the table itself and both are computed with checked arithmetic. - `try_beps_pn` computed `q * x + e` before bounding `x`, which overflows for a large `x`. The bound moves above the computation. - `MilnorSubalgebra::packed_signature` assumed the profile was no longer than `PPart::MAX_LEN` and that each signature entry fit its field. Neither holds: `SubalgebraIterator` grows a profile without limit and `from_bytes` reads whatever length a file gives. Out of range, `PPart::shift` returns 64 and the shift overflowed; an oversized entry silently spilled into the neighbouring field, which could select unrelated basis elements. It now returns `None` for a signature no element can have, and `signature_mask` yields nothing. `basis_element_from_string` is documented as total and `try_beps_pn` is the non-panicking half of `beps_pn`, so these were contract violations rather than merely untidy. Tests cover each. The signature test checks the packed mask against the per-entry comparison it replaced, over every element up to degree 60, for profiles that are narrower than their fields, wider than their fields, and longer than a p-part can be. Also adds `basis_order_at_p2_is_stable`, which pins the first nine degrees to fixed element names. The basis order is a wire format -- saved resolutions store coefficients by index -- so it needs a guard that does not read from `ppart_table`, which is the thing being guarded. Verified separately that the order is unchanged from the base commit: identical for all 4156 elements in degrees 0..=60. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01XK8YCdHkCqUV9YM7D957hV --- .../algebra/src/algebra/milnor_algebra.rs | 91 ++++++++++++-- ext/crates/algebra/src/algebra/milnor_rank.rs | 12 +- ext/crates/algebra/src/algebra/mod.rs | 4 +- ext/src/nassau.rs | 119 ++++++++++++++---- 4 files changed, 186 insertions(+), 40 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_algebra.rs b/ext/crates/algebra/src/algebra/milnor_algebra.rs index ad55e988b8..956a015d2b 100644 --- a/ext/crates/algebra/src/algebra/milnor_algebra.rs +++ b/ext/crates/algebra/src/algebra/milnor_algebra.rs @@ -106,7 +106,7 @@ pub type PPartEntry = u32; /// 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 [`Self::WIDTHS`]`[i]` bits starting at bit [`Self::SHIFTS`]`[i]`. The +/// 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 @@ -129,7 +129,7 @@ pub struct PPart(u64); impl PPart { /// The largest internal degree whose exponent sequences are guaranteed to fit. /// - /// This is the largest bound for which [`Self::WIDTHS`] sums to at most 64. It is far beyond + /// 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; @@ -660,8 +660,8 @@ impl Algebra for MilnorAlgebra { self.multiply( res.as_slice_mut(), 1, - &self.basis_element_from_index(d as i32, i), - &self.basis_element_from_index(e as i32, j), + self.basis_element_from_index(d as i32, i), + self.basis_element_from_index(e as i32, j), ); res }) @@ -799,12 +799,17 @@ impl Algebra for MilnorAlgebra { map( (tag("P^"), digits, char('_'), digits::), |(_, s, _, t)| { - if t == 0 || t > PPart::MAX_LEN { + if t == 0 || t >= combinatorics::xi_degrees(p).len() { return None; } - let entry = p.pow(s) as PPartEntry; - let degree = entry as i32 * self.q() * combinatorics::xi_degrees(p)[t]; - if degree > PPart::MAX_DEGREE || entry > PPart::max_entry(t - 1) { + 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(); @@ -1197,9 +1202,14 @@ 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 || x > PPart::max_entry(0) { + if degree > PPart::MAX_DEGREE { return None; } self.compute_basis(degree); @@ -2119,6 +2129,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); @@ -2345,6 +2384,40 @@ mod tests { } } + /// 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. diff --git a/ext/crates/algebra/src/algebra/milnor_rank.rs b/ext/crates/algebra/src/algebra/milnor_rank.rs index e0b592e19b..65c6f0a7e4 100644 --- a/ext/crates/algebra/src/algebra/milnor_rank.rs +++ b/ext/crates/algebra/src/algebra/milnor_rank.rs @@ -5,11 +5,13 @@ //! design and its measurements survive, ready to switch on when the renumbering below is worth //! taking on. //! -//! **This is not wired into [`MilnorAlgebra`].** 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. +//! 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; diff --git a/ext/crates/algebra/src/algebra/mod.rs b/ext/crates/algebra/src/algebra/mod.rs index c2425320af..6627432dcd 100644 --- a/ext/crates/algebra/src/algebra/mod.rs +++ b/ext/crates/algebra/src/algebra/mod.rs @@ -18,8 +18,8 @@ pub use field::Field; pub mod milnor_algebra; pub use milnor_algebra::MilnorAlgebra; -/// 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. +// 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; diff --git a/ext/src/nassau.rs b/ext/src/nassau.rs index 5135aabeec..9c847781ea 100644 --- a/ext/src/nassau.rs +++ b/ext/src/nassau.rs @@ -107,16 +107,33 @@ impl MilnorSubalgebra { /// 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. - fn packed_signature(&self, signature: &[PPartEntry]) -> (u64, u64) { + /// 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); } - (mask, value) + Some((mask, value)) } fn zero_signature(&self) -> Vec { @@ -133,28 +150,31 @@ impl MilnorSubalgebra { degree: i32, signature: &'a [PPartEntry], ) -> impl Iterator + 'a { - module.iter_gen_offsets([degree]).flat_map( - move |GeneratorData { - gen_deg, - start: [offset], - end: _, - }| { - // Hoist the mask out of the inner loop: every element in this block is tested - // against the same signature. - let (mask, value) = self.packed_signature(signature); - algebra - .ppart_table(degree - gen_deg) - .iter() - .enumerate() - .filter_map(move |(n, op)| { - if op.bits() & mask == value { - Some(offset + n) - } else { - None - } - }) - }, - ) + // The mask depends only on the signature, so compute it once for the whole sweep. An + // unrepresentable signature yields no elements at all. + self.packed_signature(signature) + .into_iter() + .flat_map(move |(mask, value)| { + module.iter_gen_offsets([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 + } + }) + }, + ) + }) } /// Get the matrix of a free module homomorphism when restricted to the subquotient given by @@ -1336,4 +1356,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:?}" + ); + } + } + } + } + } } From f3aab5bba719944ae311e762bf81d5c4c502ab15 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 1 Aug 2026 18:11:26 +0000 Subject: [PATCH 046/127] Apply nightly rustfmt to the packed p-part constants CI lints with the nightly toolchain, where the unstable options in `rustfmt.toml` -- `reorder_impl_items` among them -- actually take effect. Stable rustfmt skips them with a warning, so `cargo fmt --check` passed locally and failed in CI. Formatting only: the constants are sorted and the blank lines between them dropped. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01XK8YCdHkCqUV9YM7D957hV --- .../algebra/src/algebra/milnor_algebra.rs | 49 +++++++++---------- 1 file changed, 22 insertions(+), 27 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_algebra.rs b/ext/crates/algebra/src/algebra/milnor_algebra.rs index 956a015d2b..cd72a42a56 100644 --- a/ext/crates/algebra/src/algebra/milnor_algebra.rs +++ b/ext/crates/algebra/src/algebra/milnor_algebra.rs @@ -127,26 +127,30 @@ pub type PPartEntry = u32; 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; - - /// 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]; - /// `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]; @@ -157,22 +161,13 @@ impl PPart { } shifts }; - - /// `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 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 { From d356e3508c052e9a8716e1cda0262bc862ad6526 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Sat, 1 Aug 2026 14:29:04 -0400 Subject: [PATCH 047/127] milnor_gpu: also guard the output row against OOB (word < out.len()) The prior guard bounded the intra-row limb (out_offset + seqno spanning past the row) but not the row itself. compute-sanitizer on a malloc_sync build at higher degree (soak d=200) caught a second out-of-bounds atomic on the same out[word]: when row_base overruns, word = row_base + limb lands past the buffer even with a valid limb. Add the complementary `word < out.len()` bound so no atomic write can escape the allocation by either route. Validated at the d=200/48-stream soak: 4/4 clean, 0 crashes, 0 fallbacks, 0 correctness mismatches. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 535a40f942..83a3281e4c 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -1098,10 +1098,16 @@ fn multiply_pair( // 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; - let bit = u32::cast_from(global_bit % 32); - out[word].fetch_xor(1u32 << bit); + if word < out.len() { + let bit = u32::cast_from(global_bit % 32); + out[word].fetch_xor(1u32 << bit); + } } } } From df162b96159746108c2834df1c67e79102985875 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Sun, 2 Aug 2026 00:25:33 -0400 Subject: [PATCH 048/127] fp,fp-cuda: lock-free GPU row reduction + cross-runtime device arbitration MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two CUDA consumers share this GPU: the cubecl Milnor multiply and the fp-cuda row reduction. Moving the reduction off a single global mutex onto per-thread streams lets concurrent rayon workers overlap, but on its own it made a stem-200 resolution die within 5-10 seconds with CUDA_ERROR_LAUNCH_FAILED, reproducibly (3/3). Isolating the two runtimes shows the fault needs both of them on one device: multiply on GPU + reduction on CPU ran clean, and reduction on GPU + multiply on CPU ran clean, while both together died every time. compute-sanitizer finds no invalid access in either runtime — including with cubecl's allocations forced synchronous so every buffer is tracked — so this is contention, not an out-of-bounds bug. Giving fp-cuda its own non-primary CUDA context did not help either (3/3 still died); the shared *device* is what matters, not the shared context. Overlap is also catastrophic for throughput, not just stability. The composable reduction is a chain of thousands of tiny sequential per-column relaunches, so sharing the device with the multiply's saturating kernels makes every launch queue: the same reductions take 1.8-9.7 ms on an unshared GPU and 8.6-96.8 s co-running, with nvidia-smi showing 99% SM at 10% memory utilisation (queueing, not compute). The comment claiming this path "needs no cross-runtime exclusion" had it backwards — being composable means it *can* overlap without deadlocking, not that it should. gpu_lock arbitrates: multiplies take the shared side and still overlap each other; a large reduction takes the device exclusively for its ~10 ms. Total cost is ~5 s of multiply pause across a whole stem-200 run. Writer preference is required because multiplies are continuous and would starve the reduction indefinitely. Two properties are load-bearing and easy to get wrong: - WHERE the shared guard is taken. Acquiring it at multiply entry deadlocks: the marshalling par_iter runs chunks on other workers, which steal another bidegree's multiply, block on the shared side behind a waiting reduction, and never let the original join finish. It is taken alongside the existing GpuPermit, past every rayon section, for exactly the reason documented there. - Every wait is bounded, so an unforeseen cycle degrades to lost exclusivity rather than a hang. The bounds must exceed how long a reduction holds the device; at 25 ms multiplies barged back in mid-reduction and both the slowdown and the crashes returned. With this, a stem-200 resolution completes on the GPU in 2h47m with 0 crashes, where every prior attempt died. FP_CUDA_DEVICE / NASSAU_GPU_DEVICE put the two runtimes on separate GPUs when more than one is available, which removes the contention by construction and makes the arbitration a no-op. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/fp-cuda/src/lib.rs | 71 ++++++-- ext/crates/fp/src/blas/cuda.rs | 69 +++++--- ext/crates/fp/src/gpu_lock.rs | 239 +++++++++++++++++++++++++++ ext/crates/fp/src/lib.rs | 1 + ext/crates/fp/tests/cuda_dispatch.rs | 36 ++++ 5 files changed, 378 insertions(+), 38 deletions(-) create mode 100644 ext/crates/fp/src/gpu_lock.rs diff --git a/ext/crates/fp-cuda/src/lib.rs b/ext/crates/fp-cuda/src/lib.rs index 00eacfd6eb..d671cd691a 100644 --- a/ext/crates/fp-cuda/src/lib.rs +++ b/ext/crates/fp-cuda/src/lib.rs @@ -125,6 +125,11 @@ pub struct GpuContext { 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)?; @@ -188,6 +193,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 } @@ -262,7 +297,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); @@ -507,7 +542,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 @@ -622,7 +657,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)?; @@ -644,7 +679,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, @@ -655,12 +690,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 @@ -697,7 +732,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, @@ -717,7 +752,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 @@ -742,7 +777,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)?; @@ -802,7 +837,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; @@ -900,7 +935,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; const INF: i32 = 0x7fff_ffff; @@ -1056,7 +1091,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) }?; { @@ -1118,7 +1153,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, @@ -1175,7 +1210,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 { @@ -1278,7 +1313,7 @@ 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; @@ -1422,7 +1457,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; @@ -1524,7 +1559,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 { @@ -1669,7 +1704,7 @@ 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::>())?; diff --git a/ext/crates/fp/src/blas/cuda.rs b/ext/crates/fp/src/blas/cuda.rs index 48f7c148e8..846ef0b81c 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,43 @@ 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); + crate::gpu_lock::set_devices_shared(device == multiply_device()); + GpuContext::new(device).ok() }) .as_ref() } +/// Which GPU the cubecl Milnor multiply runs on (`NASSAU_GPU_DEVICE`, default 0). Read here only to +/// decide whether the two runtimes share a device; `algebra` owns the actual client construction. +fn multiply_device() -> usize { + std::env::var("NASSAU_GPU_DEVICE") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(0) +} + /// 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 +121,9 @@ 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()? - }; + // Lock-free: `matmul_b1_raw` submits on the calling thread's own stream with per-call device + // buffers, so concurrent callers do not interfere (see [`context`]). + let c = fp_cuda::matmul_b1_raw(ctx, &a_limbs, m, k, &b_limbs, n).ok()?; Some(Matrix::from_data(TWO, m, n, c)) } @@ -127,14 +147,23 @@ pub(crate) fn try_row_reduce(m: &mut Matrix) -> Option { let stride = cols.div_ceil(64); let limbs = to_limbs(m); - // The default row-reduce is composable (no cooperative launch), so it needs no - // cross-runtime exclusion against the concurrent cubecl multiply. + // 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). + let _exclusive = crate::gpu_lock::exclusive(); 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()?; + 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()?; (dev_limbs, perm, r, pivot_cols) }; diff --git a/ext/crates/fp/src/gpu_lock.rs b/ext/crates/fp/src/gpu_lock.rs new file mode 100644 index 0000000000..fa51767d1d --- /dev/null +++ b/ext/crates/fp/src/gpu_lock.rs @@ -0,0 +1,239 @@ +//! 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 { + if !arbitration_needed() { + return ExclusiveGuard(()); + } + 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)); + + 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" + ); + assert!( + max_shared.load(Ordering::SeqCst) > 1, + "multiplies never overlapped — the shared side is serialising, which defeats the point" + ); + } +} 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/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})" + ); + } + }); + } + }); +} From f9a746704396bc0ed00593b6be3499046e9975b7 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Sun, 2 Aug 2026 00:26:00 -0400 Subject: [PATCH 049/127] milnor_gpu: fail loudly on GPU context death, and split the multiply's time MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The CPU fallback turned a hard GPU fault into a silent ~100x slowdown: a run that had lost the GPU still reported "completed", so every measurement had to be reconstructed by grepping stderr, and a crash at 5 seconds looked like a slow success three hours later. GPU_DISABLED is still latched first so the soak test can tell a context death from an ordinary assertion failure, but the panic now propagates and the run dies at the fault. The batch counters existed but take_batch_stats() had no callers, so the dominant cost of a resolution was never attributed. Reporting them periodically shows where multiply time actually goes, and the split matters because the naive reading is wrong in two ways: - The window called "marshal" spans GpuPermit::acquire and the gpu_lock acquisition, so it conflates host work with time parked on our own gates. Separating them shows the permit costs nothing measurable. - The shares move by 2-3x as a run matures — early samples are small batches where fixed overhead dominates. At 2k launches it reads 29% prep / 47% lock / 24% device; by 226k it is 14% / 9% / 77%. Only the mature numbers mean anything, and they say the multiply kernel itself is the bottleneck. Reporting keys off the value fetch_add returns rather than a separate load: with ~100 workers a load races past exact multiples and the report can fire never (observed: zero output over 12 minutes). Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 156 ++++++++++++------- 1 file changed, 104 insertions(+), 52 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 83a3281e4c..8e8dc5ef71 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -213,10 +213,11 @@ use std::sync::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 — -/// the only recovery is to abandon the GPU multiply and finish the resolution on the CPU. -/// [`multiply_batch_on_gpu`] flips this on the first failure and routes all subsequent (and the -/// current) batches through [`cpu_multiply_batch`]. NOTE: this covers only the cubecl **multiply**; +/// 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); @@ -252,6 +253,14 @@ static BATCH_CALLS: AtomicU64 = AtomicU64::new(0); static BATCH_MARSHAL_US: AtomicU64 = AtomicU64::new(0); static BATCH_DEVICE_US: AtomicU64 = AtomicU64::new(0); static BATCH_PAIRS: AtomicU64 = AtomicU64::new(0); +/// `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); /// Read and reset the aggregate batch counters: `(calls, marshal_us, device_us, pairs)`. pub fn take_batch_stats() -> (u64, u64, u64, u64) { @@ -473,7 +482,6 @@ macro_rules! resident_seqno { }}; } - /// 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 @@ -897,13 +905,7 @@ fn zero_u32(out: &mut [u32]) { // 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, -) { +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]; } @@ -911,13 +913,7 @@ fn copy_into_u16( /// `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, -) { +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]; } @@ -974,13 +970,7 @@ macro_rules! copy_chunked { /// `usize` at the index sites. Shared by `seqno_kernel` and /// `multiply_single_r_kernel` so both index outputs identically. #[cube] -fn seqno_core( - g: &[u32], - xi: &[u32], - working: &[u32], - wlen: usize, - width: usize, -) -> u32 { +fn seqno_core(g: &[u32], xi: &[u32], working: &[u32], wlen: usize, width: usize) -> u32 { // cur_d = Σ working[h] · xi[h]. let mut cur_d = 0u32; for h in 0..wlen { @@ -1411,34 +1401,30 @@ pub fn multiply_batch_on_gpu( num_rows: usize, products: &[GpuProduct], ) -> Vec> { - // STOPGAP (see [`GPU_DISABLED`]): once the cubecl CUDA context has been poisoned by the - // unresolved uninit-handle bug, every launch fails, so we finish the run on the CPU. On the - // first failure we catch the panic (it surfaces as an `.unwrap()` on a `CUDA_ERROR_LAUNCH_FAILED` - // / `ServerUnhealthy` in [`multiply_batch_gpu_inner`]), flip the flag, and fall back. All later - // calls short-circuit straight to the CPU path. The CPU result is bit-identical to the GPU's - // (validated by `cpu_multiply_batch_matches_gpu`), so callers see no difference but speed. - if gpu_disabled() { - return cpu_multiply_batch(algebra, num_cols, num_rows, products); - } + // 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) })) { Ok(rows) => rows, - Err(_) => { + 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 fall through to the CPU. + // 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 by the \ - unresolved cubecl uninit-handle bug); disabling the GPU multiply and \ - completing the resolution on the CPU. RREF (separate fp-cuda runtime) is \ + "[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." ); } - cpu_multiply_batch(algebra, num_cols, num_rows, products) + std::panic::resume_unwind(payload) } } } @@ -1744,7 +1730,26 @@ fn multiply_batch_block( // time (priority inversion at worst, never deadlock). // Held for the device section (RAII): bounds total in-flight output bytes across workers. The // stream is now chosen per-thread (see [`thread_stream_id`]), not from the permit. + // 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; + 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). + let _shared = fp::gpu_lock::shared(); + 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 @@ -2231,7 +2236,11 @@ fn multiply_batch_block( // 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); + // 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, @@ -2241,10 +2250,58 @@ fn multiply_batch_block( std::sync::atomic::Ordering::Relaxed, ); BATCH_PAIRS.fetch_add(total_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); + + // 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 total = (prep_s + wait_s + device_s).max(1e-9); + 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)", + 100.0 * prep_s / total, + 100.0 * permit_s / total, + 100.0 * lock_s / total, + 100.0 * device_s / total, + ); + } result } +/// How often [`multiply_batch_block`] prints the cumulative marshal/device split, in launches. +/// `NASSAU_BATCH_REPORT_EVERY` (default 2000; `0` disables). +fn batch_report_every() -> u64 { + use std::sync::OnceLock; + static EVERY: OnceLock = OnceLock::new(); + *EVERY.get_or_init(|| { + std::env::var("NASSAU_BATCH_REPORT_EVERY") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(2000) + }) +} + /// 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 @@ -2591,13 +2648,7 @@ mod tests { /// `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, - ) { + fn seqno_kernel(g: &[u32], xi: &[u32], p_parts: &[u32], out: &mut [u32], width: usize) { let idx = ABSOLUTE_POS; if idx >= out.len() { terminate!(); @@ -3935,8 +3986,8 @@ mod tests { let n = launches.load(Ordering::Relaxed); let mm = mismatches.load(Ordering::Relaxed); eprintln!( - "[soak] {threads} threads × {secs}s, {} streams: {n} launches ({:.0}/s over {} degrees, \ - verified ≤{verify_max}), {mm} correctness mismatches, gpu_disabled={}", + "[soak] {threads} threads × {secs}s, {} streams: {n} launches ({:.0}/s over {} \ + degrees, verified ≤{verify_max}), {mm} correctness mismatches, gpu_disabled={}", gpu_stream_slots(), n as f64 / elapsed.as_secs_f64().max(1e-3), jobs.len(), @@ -3949,7 +4000,8 @@ mod tests { 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." + (tracel-ai/cubecl#1401) fired. This is the crash the submission-thread redesign \ + closes." ); } } From 46711991d033b3198c2e3282057d7ad94c367894 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Sun, 2 Aug 2026 00:26:02 -0400 Subject: [PATCH 050/127] nassau: probe whether signature steps are independent (they are not) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The per-bidegree signature loop is the tail's critical path: at stem 200 the bidegrees use A(3), whose signature space is 1024, and the steps run sequentially on one thread — b=(200,6) spends 1076 s that way, 1024 steps at ~551 ms median plus one outlier at 470 s. With only 3-5 bidegrees in flight at that point in the wavefront, parallelising this loop looked like the largest remaining win. It is not available. Each step reads dx.entry(v) over its own signature mask and then writes dx with rows of the *unmasked* matrix, whose support extends past that mask. NASSAU_PROBE_SIG_INDEP=1 snapshots dx before the loop and compares every read against it: 6703 of 38348 reads (17.5%) differ, i.e. were perturbed by an earlier signature, starting from bidegrees as small as (14,2). At p=2 a differing entry flips the zero test that drives the step, so solving the signatures against the pre-loop dx and combining would compute different answers. The loop is a genuine forward substitution and must stay ordered. Kept as a probe (zero cost unless the variable is set) because the alternative is rediscovering this by writing the parallel version and getting wrong answers. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/src/nassau.rs | 50 +++++++++++++++++++++++++++++++++++++++-------- 1 file changed, 42 insertions(+), 8 deletions(-) diff --git a/ext/src/nassau.rs b/ext/src/nassau.rs index dd54f2de67..24df98aaea 100644 --- a/ext/src/nassau.rs +++ b/ext/src/nassau.rs @@ -860,6 +860,19 @@ 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. + 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(); target_mask.clear(); @@ -907,6 +920,17 @@ impl> Resolution { 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; + } + } + } + } + for (x, dx) in xs.iter_mut().zip(&mut dxs) { scratch.set_scratch_vector_size(target_mask.len()); let mut row = 0; @@ -933,6 +957,13 @@ impl> Resolution { &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}"); } @@ -1272,8 +1303,7 @@ impl> Resolution { #[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(); + 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(); @@ -1283,9 +1313,9 @@ impl> Resolution { 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", + "[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), @@ -1543,9 +1573,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 => { From 28ed66e0355f89238e27b7ea450201f8dbee2371 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Sun, 2 Aug 2026 00:26:31 -0400 Subject: [PATCH 051/127] nassau_gpu: rustfmt reflow of an assert message Formatting only, no behaviour change; `just lint` runs `cargo fmt --all --check`. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/src/nassau_gpu.rs | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/ext/src/nassau_gpu.rs b/ext/src/nassau_gpu.rs index 0bb33217d7..c4287c06f1 100644 --- a/ext/src/nassau_gpu.rs +++ b/ext/src/nassau_gpu.rs @@ -318,9 +318,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(), ); From d2cee652540d49977578b95efb856843361994fa Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Sun, 2 Aug 2026 09:54:52 -0400 Subject: [PATCH 052/127] milnor_gpu: dedicated GPU submission thread (kill the starvation stalls) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Every worker used to run its own device section on the shared stream 0, racing 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 across two complete stem-200 runs, comparing the SAME (bidegree, signature) — bit-identical work — between runs: b=(193,8) [3,5,1,1] 235 s vs 0.04 s b=(173,4) [6,5,3,0] 109 s vs 0.01 s b=(170,6) [2,1,1,1] 0.01 s vs 59.7 s 428 step pairs differ by >100x, 1574 by >10x, wasting ~110-150 min of thread time per run. The work itself is uniform: within a bidegree the spread is ~5x (b=(199,7): 1024 steps, p50 588 ms, max 2.8 s). The outliers were purely service-order artifacts, and 473 of the 672 slow steps close in convoys — start times spread over 31-349 s, close spreads of 0.4-5.6 s. The worst case: one step blocked 372 s emitting no log line at all, while six peer threads each stayed 100% busy completing 200-576 steps at 0.6-1.6 s. Funnel every device section through one `nassau-gpu` thread fed by an mpsc channel. Service order becomes FIFO by construction, so a worker's wait is bounded by the work queued ahead of it. 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. Panics are caught per task and resumed on the waiting worker, so a panic cannot silently hang every future submission. This also deletes the multi-stream mode (NASSAU_GPU_STREAMS, thread_stream_id) outright, which removes the cross-stream pool-reclaim race (tracel-ai/cubecl #1401) structurally rather than by throttling. The soak test that used to reproduce it in ~45 s now passes: 64 threads x 60 s at max_degree=160 with NASSAU_GPU_CLEANUP_EVERY=1 gives 769 launches, 0 correctness mismatches against the CPU oracle, no context death. gpu_lock's shared side moves onto the GPU thread too: routing ~100 workers through a writer-preferring lock to reach a stage only one could occupy cost 10% of multiply time in pure convoy. It is now a 1-vs-1 handshake with the fp-cuda reduction, which is all it ever needed to be. Diagnostics: split the single `device` timer into `queue` (enqueue -> start) and `exec` (task body) with queue-depth mean/max, and add a `gpu_submit` span so a worker blocked on the device is visible in the log instead of silent — the property whose absence made these stalls look like compute for so long. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/Cargo.toml | 5 +- ext/crates/algebra/src/algebra/milnor_gpu.rs | 853 +++++++++++-------- 2 files changed, 484 insertions(+), 374 deletions(-) diff --git a/ext/crates/algebra/Cargo.toml b/ext/crates/algebra/Cargo.toml index 821f38deda..899738da1b 100644 --- a/ext/crates/algebra/Cargo.toml +++ b/ext/crates/algebra/Cargo.toml @@ -40,6 +40,9 @@ cubecl = { git = "https://github.com/tracel-ai/cubecl", tag = "v0.11.0-pre.1", o # 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 = { 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 } [dev-dependencies] criterion = { version = "0.5", features = ["html_reports"] } @@ -51,7 +54,7 @@ rstest = "0.25.0" default = ["odd-primes"] cache-multiplication = [] concurrent = ["fp/concurrent", "maybe-rayon/concurrent"] -gpu = ["dep:cubecl", "dep:cubecl-common"] +gpu = ["dep:cubecl", "dep:cubecl-common", "dep:tracing"] odd-primes = ["fp/odd-primes"] [[bench]] diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 8e8dc5ef71..49c343e8dd 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -23,7 +23,6 @@ use cubecl::{ cuda::{CudaDevice, CudaRuntime}, prelude::*, }; -use cubecl_common::stream_id::StreamId; // Bounds the per-thread enumeration state ([`ENUM_ROW_CAP`]) and the `#[cfg(test)]` `seqno_kernel`'s // working array; the multiply kernel uses `WORKING_CAP`. @@ -120,47 +119,117 @@ static GPU_BUDGET: LazyLock = LazyLock::new(|| GpuBudget { freed: Condvar::new(), }); -/// Number of distinct CUDA streams to spread device work over (`NASSAU_GPU_STREAMS`, default 1). -/// Each worker thread gets a stable per-thread stream id (see [`thread_stream_id`]), and this caps -/// how many distinct streams — hence retained memory pools — exist; `1` forces the single-stream -/// mode (all threads → stream 0). +/// The single OS thread that owns the CUDA stream, and the queue that feeds it. /// -/// **Default 1 (single stream) because multi-stream is a MEMORY disaster for little gain.** cubecl -/// gives each CUDA stream its own device AND page-locked host (pinned) memory pool, and those pools -/// are never trimmed (`memory_cleanup` frees only the GPU pool). Measured at stem 180: single-stream -/// holds pinned host memory at ~4 GB, but ≥2 streams balloons it to 140–240 GB (each stream retains -/// its own varying-size readback/staging pages) — the dominant term in the ~500 GB cgroup OOM. And -/// the payoff is ~nil: cubecl's server is single-threaded (one runner behind a channel), so extra -/// streams buy no CPU-side concurrency, only GPU kernel overlap that the big saturating multiplies -/// barely benefit from. Raise it only on a dedicated large-RAM node that wants that overlap. -fn gpu_stream_slots() -> u64 { - static SLOTS: LazyLock = LazyLock::new(|| { - std::env::var("NASSAU_GPU_STREAMS") - .ok() - .and_then(|v| v.parse::().ok()) - .filter(|&n| n > 0) - .unwrap_or(1) - }); - *SLOTS -} +/// # 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::{self, Sender}, + }, + time::Instant, + }; -/// A CUDA stream id UNIQUE to the calling worker thread (never shared or reused across threads — -/// cubecl's per-stream state assumes one driver thread per stream). Returns 0 for all threads when -/// `NASSAU_GPU_STREAMS == 1` (the single-stream fallback). With the resident master/basis now stable -/// device buffers grown IN PLACE (no handle churn — see [`ResidentDev`]), cross-stream reads of those -/// shared globals are the ordinary case cubecl supports, so distinct per-thread streams are safe and -/// give real device concurrency for the wide record-run wavefront. -fn thread_stream_id() -> u64 { - if gpu_stream_slots() == 1 { - return 0; - } - thread_local! { - static ID: u64 = { - static NEXT: AtomicU64 = AtomicU64::new(0); - NEXT.fetch_add(1, Ordering::Relaxed) + 1 // +1: id 0 is reserved for the single-stream mode - }; + use cubecl_common::stream_id::StreamId; + + /// Jobs enqueued but not yet started, i.e. how deep the FIFO is 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, + } + + fn sender() -> &'static Sender { + static QUEUE: OnceLock> = OnceLock::new(); + QUEUE.get_or_init(|| { + let (tx, rx) = mpsc::channel::(); + std::thread::Builder::new() + .name("nassau-gpu".into()) + .spawn(move || { + // Bind the stream once for the whole loop: one stream, one driver thread. + StreamId { value: 0 }.executes(|| { + while let Ok(task) = rx.recv() { + task(); + } + }); + }) + .expect("failed to spawn the nassau-gpu thread"); + tx + }) + } + + /// Run `f` on the GPU thread, blocking until it returns. Panics propagate to the caller. + pub(super) fn run(f: F) -> (T, Timing) + 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(); + sender() + .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"); + let (out, queue_ms, exec_ms) = rx.recv().expect("the nassau-gpu thread died mid-task"); + match out { + Ok(v) => ( + v, + Timing { + queue_ms, + exec_ms, + depth, + }, + ), + Err(payload) => std::panic::resume_unwind(payload), + } } - ID.with(|&id| id) } /// A/B diagnostic toggle (`NASSAU_GPU_BASIS_PASSTHROUGH=1`): when set, the batched multiply @@ -174,7 +243,7 @@ fn basis_passthrough() -> bool { } /// RAII reservation of `weight` output bytes from [`GPU_BUDGET`]; blocks (parked, not spinning) -/// until the budget admits it. The stream is chosen per worker thread (see [`thread_stream_id`]), +/// 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, @@ -261,6 +330,15 @@ static BATCH_WAIT_US: AtomicU64 = AtomicU64::new(0); /// 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); /// Read and reset the aggregate batch counters: `(calls, marshal_us, device_us, pairs)`. pub fn take_batch_stats() -> (u64, u64, u64, u64) { @@ -1729,7 +1807,7 @@ fn multiply_batch_block( // 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. The - // stream is now chosen per-thread (see [`thread_stream_id`]), not from the permit. + // device section now runs on the dedicated GPU thread (see [`gpu_thread`]), not from the permit. // 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 @@ -1747,7 +1825,11 @@ fn multiply_batch_block( // 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). - let _shared = fp::gpu_lock::shared(); + // 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 @@ -1904,333 +1986,347 @@ fn multiply_batch_block( let t_device = std::time::Instant::now(); - // Device section on this worker's own stable stream (see [`thread_stream_id`]): up to - // `gpu_stream_slots()` distinct streams run concurrently. Cubecl's per-device runner serializes - // the actual server access; the shared resident master/basis (stable in-place-grown buffers) are - // read cross-stream safely. `memory_cleanup` below trims only this stream's own transient pool. - let result = StreamId { - value: thread_stream_id(), - } - .executes(|| { - let client = CudaRuntime::client(&CudaDevice::default()); - // 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() - }; + // `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. + let submit_span = tracing::info_span!( + "gpu_submit", + rows = num_rows, + pairs = total_pairs, + out = out_len + ); + let (result, timing) = submit_span.in_scope(|| { + gpu_thread::run(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 = CudaRuntime::client(&CudaDevice::default()); + // 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() + }; - // `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. - let (cs_seg, mk_seg) = match mode { - MasterMode::Resident => { - let (cs_segs, _) = seg_grow!( + // `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. + 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 nl = h.cs_len; + (std::mem::take(&mut h.cs_pending), 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 nl = h.mk_len; + (std::mem::take(&mut h.mk_pending), 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). + const ENUM_THREADS: u32 = 256; + let n_cold = enum_rows.len(); + let cs_cap = need_cs.max(1); + let mk_cap = need_mk.max(1); + assert!( + cs_cap <= seg_elems && mk_cap <= seg_elems, + "transient scratch ({cs_cap}/{mk_cap} u16) exceeds one segment \ + ({seg_elems}); raise NASSAU_GPU_MASTER_SEG_ELEMS" + ); + let cs_scratch = client.empty(cs_cap * size_of::()); + let mk_scratch = client.empty(mk_cap * size_of::()); + let cnt_scratch = client.empty(n_cold.max(1) * size_of::()); + let epp_h = client.create_from_slice(u32::as_bytes(&enum_pp)); + let er_h = client.create_from_slice(u32::as_bytes(&enum_rows)); + let ec_h = client.create_from_slice(u32::as_bytes(&enum_cols)); + let eco_h = client.create_from_slice(u64::as_bytes(&r_cs_offset)); + let emo_h = client.create_from_slice(u64::as_bytes(&r_mk_offset)); + unsafe { + enumerate_admissible_kernel::launch_unchecked::( + &client, + CubeCount::Static((n_cold as u32).div_ceil(ENUM_THREADS).max(1), 1, 1), + CubeDim::new_1d(ENUM_THREADS), + BufferArg::from_raw_parts(epp_h, enum_pp.len()), + BufferArg::from_raw_parts(er_h, n_cold), + BufferArg::from_raw_parts(ec_h, n_cold), + BufferArg::from_raw_parts(eco_h, n_cold), + BufferArg::from_raw_parts(emo_h, n_cold), + BufferArg::from_raw_parts(cs_scratch.clone(), cs_cap), + BufferArg::from_raw_parts(mk_scratch.clone(), mk_cap), + BufferArg::from_raw_parts(cnt_scratch, n_cold.max(1)), + enum_width, + n_cold, + ); + } + ( + pad_u16(vec![(cs_scratch, cs_cap)]), + pad_u16(vec![(mk_scratch, mk_cap)]), + ) + } + }; + // 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_DEV, - cs, - RESIDENT_UPLOAD, - need_cs, + RESIDENT_BASIS_DEV, + pp, + RESIDENT_BASIS_UPLOAD, + need_basis_elems * width, copy_into_u16, u16::as_bytes, u16, - |_up: usize| { - let mut h = RESIDENT_HOST.write().unwrap(); - let nl = h.cs_len; - (std::mem::take(&mut h.cs_pending), nl) + |up: usize| { + let h = RESIDENT_BASIS_HOST.read().unwrap(); + let nl = h.lens.len() * h.width; + (h.pparts[up..nl].to_vec(), nl) } ); - let (mk_segs, _) = seg_grow!( + let (ln_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 nl = h.mk_len; - (std::mem::take(&mut h.mk_pending), nl) + 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(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). - const ENUM_THREADS: u32 = 256; - let n_cold = enum_rows.len(); - let cs_cap = need_cs.max(1); - let mk_cap = need_mk.max(1); - assert!( - cs_cap <= seg_elems && mk_cap <= seg_elems, - "transient scratch ({cs_cap}/{mk_cap} u16) exceeds one segment ({seg_elems}); \ - raise NASSAU_GPU_MASTER_SEG_ELEMS" + (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. + let tg_h = client.create_from_slice(u32::as_bytes(&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)); + 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 cs_scratch = client.empty(cs_cap * size_of::()); - let mk_scratch = client.empty(mk_cap * size_of::()); - let cnt_scratch = client.empty(n_cold.max(1) * size_of::()); - let epp_h = client.create_from_slice(u32::as_bytes(&enum_pp)); - let er_h = client.create_from_slice(u32::as_bytes(&enum_rows)); - let ec_h = client.create_from_slice(u32::as_bytes(&enum_cols)); - let eco_h = client.create_from_slice(u64::as_bytes(&r_cs_offset)); - let emo_h = client.create_from_slice(u64::as_bytes(&r_mk_offset)); - unsafe { - enumerate_admissible_kernel::launch_unchecked::( - &client, - CubeCount::Static((n_cold as u32).div_ceil(ENUM_THREADS).max(1), 1, 1), - CubeDim::new_1d(ENUM_THREADS), - BufferArg::from_raw_parts(epp_h, enum_pp.len()), - BufferArg::from_raw_parts(er_h, n_cold), - BufferArg::from_raw_parts(ec_h, n_cold), - BufferArg::from_raw_parts(eco_h, n_cold), - BufferArg::from_raw_parts(emo_h, n_cold), - BufferArg::from_raw_parts(cs_scratch.clone(), cs_cap), - BufferArg::from_raw_parts(mk_scratch.clone(), mk_cap), - BufferArg::from_raw_parts(cnt_scratch, n_cold.max(1)), - enum_width, - n_cold, - ); - } - ( - pad_u16(vec![(cs_scratch, cs_cap)]), - pad_u16(vec![(mk_scratch, mk_cap)]), - ) } - }; - // 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. - let tg_h = client.create_from_slice(u32::as_bytes(&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)); - 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_from_slice(u32::as_bytes(&prod_r_index)); - let pts_h = client.create_from_slice(u32::as_bytes(&prod_term_start)); - let pnt_h = client.create_from_slice(u32::as_bytes(&prod_num_terms)); - let prb_h = client.create_from_slice(u32::as_bytes(&prod_row_base)); - let poo_h = client.create_from_slice(u32::as_bytes(&prod_out_offset)); - let pps_h = client.create_from_slice(u32::as_bytes(&pps)); - let cubes = (total_pairs as u32).div_ceil(THREADS).max(1); - // 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) - }; - } - // 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(pri_h, products.len()), - BufferArg::from_raw_parts(pts_h, products.len()), - BufferArg::from_raw_parts(pnt_h, products.len()), - BufferArg::from_raw_parts(prb_h, products.len()), - BufferArg::from_raw_parts(poo_h, products.len()), - BufferArg::from_raw_parts(pps_h, pps.len()), - width, - seg_elems, - num_limbs, - ); - } + let pri_h = client.create_from_slice(u32::as_bytes(&prod_r_index)); + let pts_h = client.create_from_slice(u32::as_bytes(&prod_term_start)); + let pnt_h = client.create_from_slice(u32::as_bytes(&prod_num_terms)); + let prb_h = client.create_from_slice(u32::as_bytes(&prod_row_base)); + let poo_h = client.create_from_slice(u32::as_bytes(&prod_out_offset)); + let pps_h = client.create_from_slice(u32::as_bytes(&pps)); + let cubes = (total_pairs as u32).div_ceil(THREADS).max(1); + // 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) + }; + } + // 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(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()), + width, + seg_elems, + num_limbs, + ); + } - 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(); + 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(); - // 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. - let every = cleanup_every(); - if every != 0 && CLEANUP_COUNTER.fetch_add(1, Ordering::Relaxed) % every == 0 { - client.memory_cleanup(); - } + // 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. + let every = cleanup_every(); + if every != 0 && CLEANUP_COUNTER.fetch_add(1, Ordering::Relaxed) % every == 0 { + client.memory_cleanup(); + } - result + result + }) }); // Aggregate marshal/device totals across every launch (cheap, always on) so a whole @@ -2257,6 +2353,16 @@ fn multiply_batch_block( 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 @@ -2274,15 +2380,24 @@ fn multiply_batch_block( 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 total = (prep_s + wait_s + device_s).max(1e-9); 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)", + 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}", 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, ); } @@ -3827,23 +3942,23 @@ mod tests { /// - **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 (measured):** clean at `max_degree ≤ 128`, but at `max_degree=160` with - /// `NASSAU_GPU_STREAMS=48` the cross-stream pool-reclaim race fires within a ~45 s soak — the - /// `never initialized` / `ServerUnhealthy` cascade, `gpu_disabled` flips — at only ~28 GB host - /// / ~22 GB GPU, so it is NOT a device OOM but the genuine timing race. That makes this a - /// ~1–2 min, low-memory stand-in for the 40-min stem-200 crash (the race window scales with - /// buffer size; degree 64 was just below threshold). The single-submission-thread redesign - /// must turn this exact config GREEN. + /// - **#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 + `NASSAU_GPU_STREAMS>1`). Reproduce #1401 with: + /// Ignored by default (needs a CUDA device). Run the ex-reproducer config with: /// ```text - /// NASSAU_GPU_STREAMS=48 NASSAU_GPU_CLEANUP_EVERY=1 NASSAU_SOAK_MAX_DEGREE=160 \ + /// 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 cross-stream soak: needs a CUDA device and NASSAU_GPU_STREAMS>1; run explicitly"] + #[ignore = "GPU concurrency soak: needs a CUDA device; run explicitly"] fn concurrent_growth_soak() { use std::{ sync::{ @@ -3870,13 +3985,6 @@ mod tests { let verify_max = env_num("NASSAU_SOAK_VERIFY_MAX", 44) as i32; let num_rows = 32usize; - if gpu_stream_slots() == 1 { - eprintln!( - "[soak] WARNING: NASSAU_GPU_STREAMS=1 (single stream) — the cross-stream reclaim \ - race CANNOT reproduce. Set NASSAU_GPU_STREAMS >= {threads} to exercise it." - ); - } - let p = fp::prime::ValidPrime::new(2); let algebra = MilnorAlgebra::new(p, false); algebra.compute_basis(max_degree); @@ -3986,9 +4094,8 @@ mod tests { let n = launches.load(Ordering::Relaxed); let mm = mismatches.load(Ordering::Relaxed); eprintln!( - "[soak] {threads} threads × {secs}s, {} streams: {n} launches ({:.0}/s over {} \ + "[soak] {threads} threads × {secs}s, 1 gpu thread: {n} launches ({:.0}/s over {} \ degrees, verified ≤{verify_max}), {mm} correctness mismatches, gpu_disabled={}", - gpu_stream_slots(), n as f64 / elapsed.as_secs_f64().max(1e-3), jobs.len(), gpu_disabled(), From f9f595fb2c704cb7753bba17895413fa2ed2841a Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Sun, 2 Aug 2026 10:12:45 -0400 Subject: [PATCH 053/127] milnor_gpu: bench the hard stem-200 regime instead of resolving to it MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Every change to the GPU path was validated by a ~3 h stem-200 resolution, so the iteration loop ran in hours and each answer came with run-to-run variance baked in. This reproduces the regime in minutes. The workload shape is measured, not invented: it is the distribution of `gpu_submit` spans from a real stem-200 run restricted to the hard tail (stems >= 190, n=11287 launches): p10 p50 p90 max rows 122 158 260 83702 pairs 190962 1963056 14770242 632386884 out_u32 326106 379516 665860 42269510 Two properties matter as much as the sizes. Worker count is 7 — the real wavefront (time-weighted mean 6.6, 87% of the run at 6-7 bidegrees in flight) — deliberately NOT the soak's 64: queue contention is the thing under test, so the number of contenders has to match. And 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 measures a phase the hard regime is not in. num_cols is found by searching for the 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, the prep/queue/exec split with queue depth, the per-call wall tail (max/p50 ratio — the starvation signal an aggregate mean hides), and the GPU thread's duty cycle. Also adds take_gpu_timing(), the queue/exec/depth counterpart to take_batch_stats(). Kept separate because collapsing queue and exec into one `device` figure is precisely what made multi-minute submission stalls read as kernel time. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 253 +++++++++++++++++++ 1 file changed, 253 insertions(+) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 49c343e8dd..5decfa995e 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -350,6 +350,22 @@ 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 @@ -4111,4 +4127,241 @@ mod tests { 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 = 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, so this drives the identical + // kernel + resident-master path Nassau does. + 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) + }; + + 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"); + } } From 7744b4aa1838fda6d7489904b695f4f49f42a894 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Sun, 2 Aug 2026 12:07:11 -0400 Subject: [PATCH 054/127] nassau,milnor_gpu: the multi-minute stalls were a par_iter over a 5ms loop MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The marshal fill in multiply_batch_block is `tg[k] = base + ti as u32` — one add and one store per term. At the largest observed size (478972 products / 2300621 terms) that is a few milliseconds of memory-bound work. Instrumenting it measured a single fill at 146 SECONDS, with p99 at 0.09s: 99% run fast and only the tail detonates, which is contention, not cost. Rayon could never speed up a 5ms loop past its own split/join overhead. What parallelising it actually bought was a join, hence rayon's steal loop, hence exposure to minutes-long parking. Made it sequential. Matched 0-1137s window across five stem-200 runs: run steps>=20s sum max step marshal max retries D baseline 44 2836s 356.0s n/a 1190 F gpu-thread 43 2049s 224.0s n/a 1372 G +guard 76 2889s 115.0s n/a 1941 H +guard+spans 70 2587s 166.0s 146.00s 1622 I +sequential marshal 39 1185s 43.0s 0.45s 1189 marshal_terms max 146s -> 0.45s (325x), its total 33.0min -> 0.9min, slow-step time -58% vs baseline, worst step -8x. Also wraps a whole bidegree in one ParallelGuard instead of 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 already in rayon's steal loop, i.e. blocked at a join, and running it there IS the inversion — so a bounce never discards useful work, it declines exactly the runs that would invert. It costs nothing: retries are 1189 in run I versus 1190 at baseline. The apparent retry storm in G/H was threads parked in the starved marshal, not the guard. Spans added to make this findable at all: marshal_terms, ensure_basis, pair_prepass, extract_restricted (and gpu_submit earlier). The region simply was not spanned, so its time was absorbed into whichever aggregate sat nearby — which is how the same stall got misattributed in turn to GPU submission ordering, to the resident-master pre-pass, and to the multiply kernel itself. ensure_basis measures flat 0.0s, ruling out the basis write lock. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 95 +++++++++++++++----- ext/src/nassau.rs | 18 ++++ ext/src/nassau_gpu.rs | 7 +- 3 files changed, 96 insertions(+), 24 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 5decfa995e..d54018d751 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -319,6 +319,9 @@ fn cleanup_every() -> u64 { /// 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); static BATCH_MARSHAL_US: AtomicU64 = AtomicU64::new(0); static BATCH_DEVICE_US: AtomicU64 = AtomicU64::new(0); static BATCH_PAIRS: AtomicU64 = AtomicU64::new(0); @@ -865,6 +868,11 @@ fn resident_info(algebra: &MilnorAlgebra, p_part: &[PPartEntry]) -> RInfo { 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) { @@ -1641,17 +1649,35 @@ fn multiply_batch_grouped( // 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). - let prod_pairs: Vec = products - .iter() - .map(|prod| { - let r = algebra.basis_element_from_index(prod.r_degree, prod.r_idx); - let num_mats = 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, - }; - num_mats * prod.term_indices.len() - }) - .collect(); + // 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::info_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(|| { + products + .iter() + .map(|prod| { + let r = algebra.basis_element_from_index(prod.r_degree, prod.r_idx); + let num_mats = 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, + }; + num_mats * prod.term_indices.len() + }) + .collect() + }); + prepass.record( + "new_r", + RESIDENT_MISSES.load(std::sync::atomic::Ordering::Relaxed) - misses_before, + ); + drop(prepass); let mut result: Vec> = Vec::with_capacity(num_rows); let (mut r0, mut p0) = (0, 0); while r0 < num_rows { @@ -1756,23 +1782,46 @@ fn multiply_batch_block( // *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); - let global_base = ensure_basis(algebra, width, max_s_degree); + // `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::info_span!("ensure_basis", max_s_degree) + .in_scope(|| ensure_basis(algebra, width, max_s_degree)); 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. { - let tg_base = term_gei.as_mut_ptr() as usize; - let global_base = &global_base; - (0..products.len()).into_maybe_par_iter().for_each(|pi| { - let prod = &products[pi]; + // 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::info_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()); - // SAFETY: products write disjoint `[off, off + nt)` ranges (from the prefix sum), - // each within the allocated buffer, so no two tasks alias any element. The `usize` - // base is re-formed into a pointer here because raw pointers are not `Send`. - let tg = unsafe { std::slice::from_raw_parts_mut((tg_base as *mut u32).add(off), nt) }; let base = global_base[prod.s_degree as usize]; - for (k, &ti) in prod.term_indices.iter().enumerate() { - tg[k] = base + ti as u32; + for (slot, &ti) in tg_all[off..off + nt].iter_mut().zip(&prod.term_indices) { + *slot = base + ti as u32; } - }); + } } // 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. diff --git a/ext/src/nassau.rs b/ext/src/nassau.rs index 24df98aaea..226a3e101b 100644 --- a/ext/src/nassau.rs +++ b/ext/src/nassau.rs @@ -1154,6 +1154,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}")); } diff --git a/ext/src/nassau_gpu.rs b/ext/src/nassau_gpu.rs index c4287c06f1..fd70daa8fb 100644 --- a/ext/src/nassau_gpu.rs +++ b/ext/src/nassau_gpu.rs @@ -181,7 +181,12 @@ pub fn get_partial_matrix_restricted( inputs: &[usize], target_dim: usize, ) -> Matrix { - let (mut matrix, mut 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 target = hom.target(); let algebra = target.algebra(); From 3e7c821b2414f627cef8dab883da243a399d44b5 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Sun, 2 Aug 2026 17:20:16 -0400 Subject: [PATCH 055/127] milnor_gpu: zero-copy batch output, bounded bench, kernel micro-optimisation Profiling (nsys, 45s window at stem 154-192) finally characterised the kernel instead of guessing at it: multiply_batch_kernel 41.64 s D2H memcpy 0.76 s 40.0 GB -> 52.6 GB/s H2D memcpy 0.34 s 10.0 GB -> 28.9 GB/s Transfers are 2.6% of GPU time and already at near-peak pinned PCIe 5.0 x16 -- cubecl was pinning all along. GPU is 95% busy, 97% of it in the kernel. Every transfer-side theory (pageable copies, submission starvation, pinning pools) was wrong; the 2.1-5.0 GB/s "cost floor" measured earlier was kernel time scaling with `pairs`, not transfer time. GPU hardware-metric sampling while the kernel is resident: SM Active 100% | SM Issue 68% | DRAM read 1.1% | DRAM write 4.6% Compute warps in flight 36% | Unallocated warps 64% So: instruction-bound, not bandwidth-bound, with occupancy headroom. Changes: * BatchOutput holds the D2H landing buffers (cubecl `Bytes`) and hands out row slices as views, replacing a `Vec` per row -- ~32M allocations per stem-200 resolution plus a full re-copy onto freshly-mapped pages. Soak throughput 915 vs 769 launches in the same 60s config (+19%); the win is host-side (allocations, page faults), NOT PCIe. * Marshalled buffers move to `client.create(Bytes::from_elems(v))` instead of `create_from_slice`, so the bytes the marshal built are handed over rather than copied out of a borrowed slice. * Binary search runs `ceil(log2(num_products))` iterations (~15) instead of a fixed 32, each of which was a DEPENDENT global load. Kept as a runtime scalar: `#[comptime]` on it forces an NVRTC recompile per distinct value and widened bench spread from 0.7% to 5.6%. * `seg_read_*` takes `#[comptime] num_segs`, so the 16-way select chain folds to the segments that exist. Comptime is right here (<= MASTER_MAX_SEG distinct values); as a plain runtime scalar the guards were pure overhead. Bench A/B (0.7% run-to-run), 5.89e9 -> 5.97e9 pairs/s, +1.4%. Modest: these two structural suspects are not the bottleneck. The occupancy figure points at the three per-thread `Array::::new(WORKING_CAP)` locals (96 u16 slots/thread) as the next lever. Also fixes the regime bench, which could not run at all: it enumerated every (R, s) pair, so at stem-200-scale degrees it tried to materialise `sum_r dim(r)*dim(out-r)` term indices and died in setup. Now samples R's on a stride to a bounded product count and uses ~5 terms per product, matching the measured distribution (products p50 24k, terms/product 4.8). Gives a 2-minute iteration loop at 0.7% noise, against 3 hours for a resolution. ncu is unusable on this host -- it fails identically on a 6-line kernel, an ncu 2024.1 / driver 580.119.02 incompatibility, and the matching newer builds need GLIBC 2.27 that EL7 lacks. `nsys --gpu-metrics-device` sampling is the substitute and needs no kernel replay. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 402 +++++++++++++------ ext/src/nassau.rs | 50 +-- ext/src/nassau_gpu.rs | 8 +- ext/src/utils.rs | 34 +- 4 files changed, 321 insertions(+), 173 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index d54018d751..fb60c08e77 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -23,6 +23,7 @@ use cubecl::{ cuda::{CudaDevice, CudaRuntime}, prelude::*, }; +use cubecl_common::bytes::Bytes; // Bounds the per-thread enumeration state ([`ENUM_ROW_CAP`]) and the `#[cfg(test)]` `seqno_kernel`'s // working array; the multiply kernel uses `WORKING_CAP`. @@ -1318,6 +1319,14 @@ fn multiply_batch_kernel( width: usize, seg_elems: usize, num_limbs: usize, + // Runtime scalar, deliberately NOT `#[comptime]`: it is `ceil(log2(num_products))`, which + // 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%). The win here is executing ~15 iterations + // instead of a fixed 32, which needs no unrolling. + 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, ) { let k = ABSOLUTE_POS; let num_products = prod_pair_start.len() - 1; @@ -1326,11 +1335,16 @@ 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. + // so `prod_pair_start` is strictly increasing and `p` is unique). + // + // `search_iters` is `ceil(log2(num_products))`, computed on the host and passed as a bare + // scalar, which cubecl bakes in as a compile-time constant — so the loop unrolls to exactly + // the steps the search needs. It was a fixed 32, and each step is a DEPENDENT global load of + // `prod_pair_start[mid]`, so the surplus iterations were a pure latency chain on every thread. + // At the measured p50 of ~24k products this is 15 steps rather than 32. let mut lo = 0usize; let mut hi = num_products; - for _ in 0..32 { + for _ in 0..search_iters { if hi - lo > 1 { let mid = (lo + hi) / 2; if usize::cast_from(prod_pair_start[mid]) <= k { @@ -1363,7 +1377,7 @@ fn multiply_batch_kernel( let pp_off = gei * width; let term_len = usize::cast_from(seg_read_u32( ln0, ln1, ln2, ln3, ln4, ln5, ln6, ln7, ln8, ln9, ln10, ln11, ln12, ln13, ln14, ln15, gei, - seg_elems, + seg_elems, num_segs, )); // Gather this thread's matrix / term out of the segmented stores into contiguous locals, then run @@ -1395,6 +1409,7 @@ fn multiply_batch_kernel( cs15, cs_off + j, seg_elems, + num_segs, ); } cs_local[j] = c; @@ -1419,6 +1434,7 @@ fn multiply_batch_kernel( mk15, mk_off + j, seg_elems, + num_segs, ); } mk_local[j] = mm; @@ -1443,6 +1459,7 @@ fn multiply_batch_kernel( pp15, pp_off + j, seg_elems, + num_segs, ); } term_local[j] = b; @@ -1497,12 +1514,87 @@ pub struct GpuProduct { /// 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: &MilnorAlgebra, num_cols: usize, num_rows: usize, products: &[GpuProduct], -) -> Vec> { +) -> 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 @@ -1512,7 +1604,7 @@ pub fn multiply_batch_on_gpu( match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { multiply_batch_gpu_inner(algebra, num_cols, num_rows, products) })) { - Ok(rows) => rows, + 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. @@ -1540,11 +1632,12 @@ pub fn cpu_multiply_batch( 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); @@ -1570,7 +1663,7 @@ pub fn cpu_multiply_batch( rows[prod.row][col / 32] ^= 1u32 << (col % 32); } } - rows + BatchOutput::from_limbs(rows.concat(), num_limbs) } fn multiply_batch_gpu_inner( @@ -1578,12 +1671,16 @@ fn multiply_batch_gpu_inner( num_cols: usize, num_rows: usize, products: &[GpuProduct], -) -> Vec> { +) -> BatchOutput { let cap = resident_degree_cap(); // 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 multiply_batch_grouped(algebra, num_cols, num_rows, products, MasterMode::Resident); + 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, @@ -1593,7 +1690,7 @@ fn multiply_batch_gpu_inner( // 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![vec![0u32; num_limbs]; num_rows]; + let mut result = vec![0u32; num_rows * num_limbs]; for mode in [MasterMode::Resident, MasterMode::Transient] { let is_group = |d: i32| match mode { MasterMode::Resident => d <= cap, @@ -1619,14 +1716,19 @@ fn multiply_batch_gpu_inner( q }) .collect(); - let sub = multiply_batch_grouped(algebra, num_cols, rows.len(), &compact, mode); + let sub_blocks = multiply_batch_grouped(algebra, num_cols, rows.len(), &compact, mode); + let sub: Vec = sub_blocks + .iter() + .flat_map(|b| u32::from_bytes(b).iter().copied()) + .collect(); for (i, &orig) in rows.iter().enumerate() { - for (a, b) in result[orig].iter_mut().zip(&sub[i]) { - *a ^= *b; + let (dst, src) = (orig * num_limbs, i * num_limbs); + for k in 0..num_limbs { + result[dst + k] ^= sub[src + k]; } } } - result + BatchOutput::from_limbs(result, num_limbs) } fn multiply_batch_grouped( @@ -1635,7 +1737,7 @@ fn multiply_batch_grouped( num_rows: usize, products: &[GpuProduct], mode: MasterMode, -) -> Vec> { +) -> 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 @@ -1678,7 +1780,7 @@ fn multiply_batch_grouped( RESIDENT_MISSES.load(std::sync::atomic::Ordering::Relaxed) - misses_before, ); drop(prepass); - let mut result: Vec> = Vec::with_capacity(num_rows); + 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 @@ -1696,7 +1798,7 @@ fn multiply_batch_grouped( pairs += row_pairs; (r1, p1) = (r1 + 1, q); } - result.extend(multiply_batch_block( + result.push(multiply_batch_block( algebra, num_cols, r0, @@ -1721,7 +1823,7 @@ fn multiply_batch_block( num_rows: usize, products: &[GpuProduct], mode: MasterMode, -) -> Vec> { +) -> Bytes { let (width, g) = algebra.seqno_table_u32(); let mut xi: Vec = xi_degrees(algebra.prime()) .iter() @@ -2033,7 +2135,7 @@ fn multiply_batch_block( ); } if total_pairs == 0 { - return vec![vec![0u32; num_limbs]; num_rows]; + return Bytes::from_elems(vec![0u32; num_rows * num_limbs]); } // The resident `col_sums`/`masks` and basis are non-empty once any `R`/term is present @@ -2047,6 +2149,8 @@ fn multiply_batch_block( 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(); @@ -2237,7 +2341,13 @@ fn multiply_batch_block( // 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. - let tg_h = client.create_from_slice(u32::as_bytes(&term_gei)); + // 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); @@ -2266,13 +2376,24 @@ fn multiply_batch_block( ); } - let pri_h = client.create_from_slice(u32::as_bytes(&prod_r_index)); - let pts_h = client.create_from_slice(u32::as_bytes(&prod_term_start)); - let pnt_h = client.create_from_slice(u32::as_bytes(&prod_num_terms)); - let prb_h = client.create_from_slice(u32::as_bytes(&prod_row_base)); - let poo_h = client.create_from_slice(u32::as_bytes(&prod_out_offset)); - let pps_h = client.create_from_slice(u32::as_bytes(&pps)); + 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 cubes = (total_pairs as u32).div_ceil(THREADS).max(1); + // Exactly the binary-search depth this block needs (see the kernel): passed bare so cubecl + // specialises the trip count instead of always running 32 dependent loads. + let search_iters = usize::BITS as usize - num_products.max(1).leading_zeros() as usize; + // 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); // Bind one `BufferArg` per `(segment vector, index)` — the `.0` handle, `.1` element length. macro_rules! sa { ($v:expr, $i:expr) => { @@ -2350,7 +2471,7 @@ fn multiply_batch_block( 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(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), @@ -2363,18 +2484,17 @@ fn multiply_batch_block( 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(pps_h, pps_len), width, seg_elems, num_limbs, + search_iters, + num_segs, ); } - 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(); + // Hand back the landing buffer itself — no copy, no allocation (see [`BatchOutput`]). + let result = client.read_one(out_h).unwrap(); // 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 @@ -2511,42 +2631,54 @@ fn seg_read_u16( s15: &[u16], o: usize, seg_elems: usize, + #[comptime] num_segs: usize, ) -> u16 { - let seg = o / seg_elems; - let local = o % seg_elems; + // `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 seg == 0 { - v = s0[local]; - } else if seg == 1 { - v = s1[local]; - } else if seg == 2 { - v = s2[local]; - } else if seg == 3 { - v = s3[local]; - } else if seg == 4 { - v = s4[local]; - } else if seg == 5 { - v = s5[local]; - } else if seg == 6 { - v = s6[local]; - } else if seg == 7 { - v = s7[local]; - } else if seg == 8 { - v = s8[local]; - } else if seg == 9 { - v = s9[local]; - } else if seg == 10 { - v = s10[local]; - } else if seg == 11 { - v = s11[local]; - } else if seg == 12 { - v = s12[local]; - } else if seg == 13 { - v = s13[local]; - } else if seg == 14 { - v = s14[local]; + if num_segs == 1 { + v = s0[o]; } else { - v = s15[local]; + 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]; + } } v } @@ -2574,42 +2706,54 @@ fn seg_read_u32( s15: &[u32], o: usize, seg_elems: usize, + #[comptime] num_segs: usize, ) -> u32 { - let seg = o / seg_elems; - let local = o % seg_elems; + // `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 seg == 0 { - v = s0[local]; - } else if seg == 1 { - v = s1[local]; - } else if seg == 2 { - v = s2[local]; - } else if seg == 3 { - v = s3[local]; - } else if seg == 4 { - v = s4[local]; - } else if seg == 5 { - v = s5[local]; - } else if seg == 6 { - v = s6[local]; - } else if seg == 7 { - v = s7[local]; - } else if seg == 8 { - v = s8[local]; - } else if seg == 9 { - v = s9[local]; - } else if seg == 10 { - v = s10[local]; - } else if seg == 11 { - v = s11[local]; - } else if seg == 12 { - v = s12[local]; - } else if seg == 13 { - v = s13[local]; - } else if seg == 14 { - v = s14[local]; + if num_segs == 1 { + v = s0[o]; } else { - v = s15[local]; + 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]; + } } v } @@ -3048,6 +3192,7 @@ mod tests { idx: &[u32], out: &mut [u16], seg_elems: usize, + #[comptime] num_segs: usize, ) { let i = ABSOLUTE_POS; if i >= out.len() { @@ -3072,6 +3217,7 @@ mod tests { s15, usize::cast_from(idx[i]), seg_elems, + num_segs, ); } @@ -3332,6 +3478,7 @@ mod tests { 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() @@ -3826,6 +3973,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!( @@ -3978,6 +4126,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!( got, golden, @@ -4095,7 +4244,7 @@ mod tests { struct Job { num_cols: usize, products: Vec, - golden: Option>>, + golden: Option, } let jobs: Arc> = Arc::new( (12..=max_degree) @@ -4271,15 +4420,30 @@ mod tests { algebra.dimension(out_degree), ); - // Same `get_partial_matrix`-shaped batch the soak builds, so this drives the identical - // kernel + resident-master path Nassau does. + // 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); - let mut products = Vec::new(); + // 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; - let s_dim = algebra.dimension(s_degree); - if s_dim == 0 { + if algebra.dimension(s_degree) == 0 { continue; } for r_idx in 0..algebra.dimension(r_degree) { @@ -4290,17 +4454,25 @@ mod tests { { 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, - }); + 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) }; diff --git a/ext/src/nassau.rs b/ext/src/nassau.rs index 226a3e101b..cec47c60c4 100644 --- a/ext/src/nassau.rs +++ b/ext/src/nassau.rs @@ -754,7 +754,6 @@ impl> Resolution { && reuse_within_cap(target_dim, next_dim) { let all_rows: Vec = (0..target_dim).collect(); - let _guard = ParallelGuard::new(); Some(restricted_partial_matrix_maybe_gpu( &self.differentials[b.s() - 1], b.t(), @@ -770,15 +769,12 @@ impl> Resolution { debug_assert!(target_mask.iter().all(|&r| r < full.rows())); select_rows(full, &target_mask) } - None => { - let _guard = ParallelGuard::new(); - restricted_partial_matrix_maybe_gpu( - &self.differentials[b.s() - 1], - b.t(), - &target_mask, - next_dim, - ) - } + 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]); @@ -897,15 +893,12 @@ impl> Resolution { debug_assert!(target_mask.iter().all(|&r| r < full.rows())); select_rows(full, &target_mask) } - None => { - let _guard = ParallelGuard::new(); - restricted_partial_matrix_maybe_gpu( - &self.differentials[b.s() - 1], - b.t(), - &target_mask, - next_dim, - ) - } + None => restricted_partial_matrix_maybe_gpu( + &self.differentials[b.s() - 1], + b.t(), + &target_mask, + next_dim, + ), }; let mut masked_matrix = @@ -1007,10 +1000,7 @@ impl> Resolution { source_dim + target_dim, 0, ); - { - let _guard = ParallelGuard::new(); - 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(); @@ -1048,10 +1038,7 @@ impl> Resolution { let mut matrix = AugmentedMatrix::<2>::new(p, target_dim, [cc_module.dimension(t), target_dim]); - { - let _guard = ParallelGuard::new(); - self.chain_maps[0].get_matrix(matrix.segment(0, 0), t); - } + 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(); @@ -1063,10 +1050,7 @@ impl> Resolution { source_dim + MAX_NEW_GENS, 0, ); - { - let _guard = ParallelGuard::new(); - 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(); @@ -1798,6 +1782,10 @@ 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) }; diff --git a/ext/src/nassau_gpu.rs b/ext/src/nassau_gpu.rs index fd70daa8fb..2c5adc5e3b 100644 --- a/ext/src/nassau_gpu.rs +++ b/ext/src/nassau_gpu.rs @@ -52,8 +52,8 @@ pub fn get_partial_matrix(hom: &NassauDifferential, degree: i32, inputs: &[usize // 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 out = multiply_batch_on_gpu(&algebra, num_cols, inputs.len(), &products); + for (row, limbs) in out.iter_rows().enumerate() { let mut target_row = matrix.row_mut(row); for (limb_idx, &limb) in limbs.iter().enumerate() { let mut bits = limb; @@ -217,8 +217,8 @@ pub fn get_partial_matrix_restricted( for pr in &mut products[p0..p1] { pr.row -= r0; // batch-local row index for the kernel's output layout } - let rows = multiply_batch_on_gpu(&algebra, full_cols, r1 - r0, &products[p0..p1]); - for (bi, limbs) in rows.iter().enumerate() { + let out = multiply_batch_on_gpu(&algebra, full_cols, r1 - r0, &products[p0..p1]); + for (bi, limbs) in out.iter_rows().enumerate() { let mut target_row = matrix.row_mut(r0 + bi); for (limb_idx, &limb) in limbs.iter().enumerate() { let mut bits = limb; diff --git a/ext/src/utils.rs b/ext/src/utils.rs index cb82ce16c3..17e40ed161 100644 --- a/ext/src/utils.rs +++ b/ext/src/utils.rs @@ -663,39 +663,27 @@ pub(crate) mod parallel { } /// RAII guard that increments this thread's [`PARALLEL_DEPTH`] on creation and decrements it on - /// drop. Used to mark regions where `par_iter_mut` work is active, so that a `step_resolution` - /// job stolen onto a blocked guard holder can detect the priority inversion and retry. + /// 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 { - #[allow(dead_code)] - span: tracing::span::EnteredSpan, + _private: (), } impl ParallelGuard { pub(crate) fn new() -> Self { - let counter_start = PARALLEL_DEPTH.with(|d| { - let v = d.get(); - d.set(v + 1); - v - }); - Self { - span: tracing::info_span!( - "parallel_guard", - counter_start, - counter_end = tracing::field::Empty - ) - .entered(), - } + PARALLEL_DEPTH.with(|d| d.set(d.get() + 1)); + Self { _private: () } } } impl Drop for ParallelGuard { fn drop(&mut self) { - let counter_end = PARALLEL_DEPTH.with(|d| { - let v = d.get() - 1; - d.set(v); - v - }); - self.span.record("counter_end", counter_end); + PARALLEL_DEPTH.with(|d| d.set(d.get() - 1)); } } From 62fa89caec4a379d2cc37d56b11297a657985eec Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Sun, 2 Aug 2026 17:43:04 -0400 Subject: [PATCH 056/127] milnor_gpu: size the per-thread working arrays per launch (+29%) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The multiply kernel is instruction-bound, not bandwidth-bound: with it resident, SM Active 100%, SM Issue 68%, DRAM read 1.1% of peak. But only 36% of compute warps were in flight with 64% of warp slots UNALLOCATED — occupancy headroom, not a work shortage. The cause is per-thread state. Every thread held four arrays sized by the fixed `WORKING_CAP = 32`: `working` (u32) plus `cs_local`/`mk_local`/`term_local` (u16). That constant set register pressure and therefore occupancy. 32 is far above what the data needs. `mk_len = rows + cols - 1`, and for p=2 at internal degree <=310 the xi degrees 1,3,...,255 give rows <= 8 while entries satisfy sum r_i(2^i - 1) <= 310 so cols <= 9 — i.e. mk_len <= 16, half the cap. Hardcoding 16 measured +30% and would have been a bug. `mk_len` grows with degree: 16 holds to t~510, but a 9th xi appears at t>=511 making it 17, and 18 past 1023. A fixed 16 would silently truncate at stem 300 — wrong answers, no error, exactly where we are headed. So `work_cap` is a `#[comptime]` parameter derived per launch from this block's `max(r_mk_len)` and `MAX_XI_TAU`, rounded to a multiple of 4 to keep the number of distinct comptime values (hence NVRTC recompiles) small, with a host assert that it fits `WORKING_CAP` (which still bounds the host-side `xi` padding). It tracks the degree upward instead of pinning a constant. Bench A/B (0.8% run-to-run spread): baseline 5.89e9 pairs/s hardcoded cap 16 (unsafe) 7.64e9 +30% adaptive comptime cap 7.61e9 +29% The adaptive version captures essentially the whole win. For contrast the two structural suspects we had been circling all session — the fixed 32-iteration binary search and the 16-way segment select — were worth +1.4% combined. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 50 +++++++++++++++++--- 1 file changed, 43 insertions(+), 7 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index fb60c08e77..84799ec432 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -1132,16 +1132,17 @@ fn multiply_pair( out_offset: usize, width: usize, num_limbs: usize, + #[comptime] work_cap: usize, ) { let mut low = cs_len; if term_len < cs_len { low = term_len; } - let mut working = Array::::new(WORKING_CAP); + let mut working = Array::::new(work_cap); let mut rejected = false; - for j in 0..WORKING_CAP { + for j in 0..work_cap { let mut b = 0u32; if j < term_len { b = u32::cast_from(term_pparts[b_base + j]); @@ -1183,7 +1184,7 @@ fn multiply_pair( // `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.as_slice(), WORKING_CAP, width); + let idx = seqno_core(g, xi, working.as_slice(), work_cap, 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` @@ -1327,6 +1328,20 @@ fn multiply_batch_kernel( // 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, + // Per-thread working-array size, specialised per launch to what THIS block actually needs + // (`max(mk_len, term_len)`, rounded up), not the global worst case. + // + // This is the kernel's dominant cost. Every thread holds four arrays of this length + // (`working` u32 plus `cs_local`/`mk_local`/`term_local` u16), so the constant sets 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). + // + // 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`]. + #[comptime] work_cap: usize, ) { let k = ABSOLUTE_POS; let num_products = prod_pair_start.len() - 1; @@ -1384,10 +1399,10 @@ fn multiply_batch_kernel( // the pure arithmetic core on them (base 0). Entries past each length are zero, matching // `multiply_pair`'s own out-of-range convention. The loop bounds at `WORKING_CAP`, exactly as the // core does, so any `mk_len > WORKING_CAP` tail (never read by the core) is likewise not gathered. - let mut cs_local = Array::::new(WORKING_CAP); - let mut mk_local = Array::::new(WORKING_CAP); - let mut term_local = Array::::new(WORKING_CAP); - for j in 0..WORKING_CAP { + let mut cs_local = Array::::new(work_cap); + let mut mk_local = Array::::new(work_cap); + let mut term_local = Array::::new(work_cap); + for j in 0..work_cap { let mut c = 0u16; if j < cs_len { c = seg_read_u16( @@ -1482,6 +1497,7 @@ fn multiply_batch_kernel( usize::cast_from(prod_out_offset[p]), width, num_limbs, + work_cap, ); } @@ -2394,6 +2410,24 @@ fn multiply_batch_block( .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)" + ); // Bind one `BufferArg` per `(segment vector, index)` — the `.0` handle, `.1` element length. macro_rules! sa { ($v:expr, $i:expr) => { @@ -2490,6 +2524,7 @@ fn multiply_batch_block( num_limbs, search_iters, num_segs, + work_cap, ); } @@ -3069,6 +3104,7 @@ mod tests { 0, width, num_limbs, + WORKING_CAP, ); } From 22b0a67b51f09a6a33b362df4953f53b47cd8fde Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Sun, 2 Aug 2026 22:38:17 -0400 Subject: [PATCH 057/127] milnor_gpu: pack the kernel's `working` accumulator into a u64 Replaces `Array::::new(work_cap)` -- the largest single piece of per-thread state -- with one packed `u64`, read back through a per-launch 80-byte shift/mask table mirroring `PPart`'s private SHIFTS/WIDTHS. Dropping positions >= PPART_MAX_LEN is exact, not a truncation, and holds by the degree bound rather than by observation: at p = 2 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 forced by that bound, which is also why PPart::set can assert on it. THROUGHPUT-NEUTRAL TODAY. Bench 7.61e9 pairs/s against 7.62e9 (0.8% noise), and ptxas still reports 78 registers. That is not the change failing: the generated CUDA confirms `array` is gone (only the six uint16[16] for cs/mk/term remain). ptxas allocates subject to an occupancy target, and at 256 threads it cannot reach the next tier (64 regs -> 4 blocks/SM) either way, so it spends the freed registers keeping other values live and lands on 78 again. The saving is real but unbanked. It shows up as soon as occupancy is forced. Recompiling both variants with __launch_bounds__ minBlocksPerSM: tier phase A spill st/ld packed spill st/ld 4 blocks (64) 150 / 60 B 130 / 52 B 5 blocks (48) 318 / 144 B 290 / 116 B 6 blocks (40) 454 / 232 B 398 / 176 B 13-25% less spill traffic at equal occupancy. So this is kept as a strictly smaller resource footprint and as the prerequisite for the forced-occupancy experiment, not for any gain it delivers on its own. That experiment is the actual open lever: occupancy is 37.5% (3 blocks/SM at 78 registers, matching the measured 36.8% Compute Warps in Flight), with 63% of warp slots idle. Testing it needs cubecl to emit minBlocksPerSM -- today cubecl-cpp's cuda dialect emits only `__launch_bounds__()`. Correctness: 75/75 algebra tests, including every GPU test. CI gate clean. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 116 +++++++++++++++++-- 1 file changed, 104 insertions(+), 12 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 1d54aa458f..b248bd8e8a 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -35,6 +35,26 @@ use crate::algebra::{Algebra, MilnorAlgebra, combinatorics::xi_degrees, milnor_a /// `mk_len = rows + cols − 1 ≤ MAX_XI_TAU + ⌈log2⌉`; 32 covers every in-range case. const WORKING_CAP: usize = 32; +/// `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 ≤ WORKING_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`. @@ -1072,18 +1092,29 @@ macro_rules! copy_chunked { /// `usize` at the index sites. Shared by `seqno_kernel` and /// `multiply_single_r_kernel` so both index outputs identically. #[cube] -fn seqno_core(g: &[u32], xi: &[u32], working: &[u32], wlen: usize, width: usize) -> u32 { - // cur_d = Σ working[h] · xi[h]. +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 { - cur_d += working[h] * xi[h]; + 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 = working[h]; + 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; @@ -1132,13 +1163,25 @@ fn multiply_pair( width: usize, num_limbs: usize, #[comptime] work_cap: usize, + #[comptime] sq_len: usize, + pp_shift: &[u32], + pp_mask: &[u32], ) { let mut low = cs_len; if term_len < cs_len { low = term_len; } - let mut working = Array::::new(work_cap); + // 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 = false; for j in 0..work_cap { @@ -1176,14 +1219,19 @@ fn multiply_pair( } val = b | mk; } - working[j] = val; + if j < PPART_MAX_LEN { + working |= u64::cast_from(val) << u64::cast_from(pp_shift[j]); + } } 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.as_slice(), work_cap, width); + // Only the first `PPART_MAX_LEN` positions can be non-zero (see the accumulator above), + // 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` @@ -1327,6 +1375,10 @@ fn multiply_batch_kernel( // 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 working-array size, specialised per launch to what THIS block actually needs // (`max(mk_len, term_len)`, rounded up), not the global worst case. // @@ -1497,6 +1549,9 @@ fn multiply_batch_kernel( width, num_limbs, work_cap, + sq_len, + pp_shift, + pp_mask, ); } @@ -2404,6 +2459,11 @@ fn multiply_batch_block( // Exactly the binary-search depth this block needs (see the kernel): passed bare so cubecl // specialises the trip count instead of always running 32 dependent loads. let search_iters = usize::BITS as usize - num_products.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`. @@ -2526,6 +2586,9 @@ fn multiply_batch_block( 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), work_cap, ); } @@ -3009,18 +3072,29 @@ mod tests { /// `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) { + 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 = Array::::new(MAX_XI_TAU); - for h in 0..width { - working[h] = p_parts[base + h]; + 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(g, xi, working.as_slice(), width, width); + 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. @@ -3044,6 +3118,10 @@ mod tests { 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 { @@ -3056,6 +3134,9 @@ mod tests { 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), ); } @@ -3081,6 +3162,8 @@ mod tests { mk_len: usize, width: usize, num_limbs: usize, + pp_shift: &[u32], + pp_mask: &[u32], ) { let pair = ABSOLUTE_POS; if pair >= num_matrices * num_terms { @@ -3107,6 +3190,9 @@ mod tests { width, num_limbs, WORKING_CAP, + PPART_MAX_LEN, + pp_shift, + pp_mask, ); } @@ -3178,6 +3264,10 @@ mod tests { 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); @@ -3199,6 +3289,8 @@ mod tests { mk_len, width, num_limbs, + BufferArg::from_raw_parts(psh_h, pp_len), + BufferArg::from_raw_parts(pms_h, pp_len), ); } From 76869d200f06dea9396e77f031fec4a6f0394056 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Sun, 2 Aug 2026 23:38:19 -0400 Subject: [PATCH 058/127] milnor_gpu: record the occupancy finding and prepare the cubecl change The multiply kernel is register-bound. ptxas with no `minBlocksPerSM` target settles at 78 registers = 3 blocks/SM = 37.5% occupancy, matching a sampled 36.8% "Compute Warps in Flight" -- 63% of warp slots sit idle. Forcing occupancy pays, measured on the stem-200 bench (0.4-0.8% run-to-run): default (3 blocks/SM) 7.61e9 pairs/s 4 blocks/SM 8.08e9 +6.2% 130 B spill stores 5 blocks/SM 8.05e9 290 B 6 blocks/SM 8.07e9 398 B 7-8 blocks/SM 6.7e9 -12% 4 is the pick: same throughput as 5/6, tightest spread, least spill pressure, and a margin before the cliff at 7 where spilling overtakes the occupancy gain. Not enabled here, because cubecl emits only `__launch_bounds__()` and offers no way to request the second argument. The change is prepared as a proper compilation setting rather than a side channel -- `KernelOptions:: min_blocks_per_sm`, a builder on the kernel settings, and a `min_blocks_per_sm = N` argument to `#[cube(..)]`, mirroring how `cluster_dim` is already threaded through, with the CUDA dialect emitting the second argument only when set (`None` reproduces current output byte for byte). It is exported to ~/cubecl-min-blocks-per-sm.patch (51 insertions, 6 files) to be opened against tracel-ai/cubecl. Verified end to end against a local build: 75/75 algebra tests, 8.09/8.06/8.08e9 pairs/s. Deliberately opt-in per kernel rather than a global heuristic: the right target depends on register pressure, and the 7-8 measurements show forcing it too high is a sizeable regression. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index b248bd8e8a..4fdfad5ae1 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -1283,6 +1283,17 @@ fn multiply_pair( // via `seg_read_*` (correct for any offset, straddle or not — no layout padding needed), then hands // those contiguous locals to `multiply_pair`, which stays a pure segmentation-agnostic arithmetic // core. `seg_elems` is the segment element count (`o / seg_elems` picks the segment). +// OCCUPANCY (pending an upstream cubecl change): this kernel is register-bound, and ptxas without +// a `minBlocksPerSM` target settles at 78 registers = 3 blocks/SM = 37.5% occupancy (matching a +// sampled 36.8% "Compute Warps in Flight", with 63% of warp slots idle). Asking for 4 blocks/SM +// measured +6.2% (7.61e9 -> 8.08e9 pairs/s) for 130 bytes of spill stores; 5 and 6 tie within +// noise, 7-8 lose 12% as spilling overtakes the gain. +// +// It is not enabled here because cubecl emits only `__launch_bounds__()` -- there is no +// way to request the second argument. The change is prepared as +// `~/cubecl-min-blocks-per-sm.patch` (adds `KernelOptions::min_blocks_per_sm` and a +// `min_blocks_per_sm = N` argument to `#[cube(..)]`, mirroring `cluster_dim`); once it lands +// upstream, add `min_blocks_per_sm = 4` below and re-measure. #[cube(launch_unchecked, address_type = "u64")] #[allow(clippy::too_many_arguments)] fn multiply_batch_kernel( From 93520e8078d6fd2da74dc50b714c7aaa015106cd Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Mon, 3 Aug 2026 02:23:04 -0400 Subject: [PATCH 059/127] milnor_gpu: fuse the per-thread gather into the column loop (48 -> 40 registers) The batch kernel gathered each thread's matrix and term out of the segmented stores into three `Array::::new(work_cap)` locals and then handed those to `multiply_pair`. But every entry is consumed exactly once, at the same `j`, by that function's own `0..work_cap` loop -- the arrays only ferried values between two loops with identical bounds. Reading each column straight into the arithmetic deletes all three. That is the last per-thread state of `work_cap` length (`working` became a `u64` in the previous commit), and per-thread state is what set this kernel's register count: 78 -> 48 -> 40, i.e. 3 -> 5 -> 6 blocks/SM (37.5% -> 75% occupancy), with zero spill at every tier. Occupancy no longer scales with the internal degree. Measured +5.6% on the stem-200 bench (6.43 -> 6.79 e9 pairs/s) over four paired rounds with the arms interleaved; the arms do not overlap (head 6.41-6.47, fused 6.73-6.83). 75/75 tests pass. The per-column rule and the output tail move into `pair_col` / `pair_emit` so the batch and single-`R` paths still share them verbatim and cannot drift. `pair_col` returns the rejection flag packed above the value rather than out of band, which keeps the caller branchless: an intermediate version that branched on the rejection test reached 35 registers (7 blocks/SM) but measured ~11% SLOWER -- twelve divergent branches per thread cost more than the extra resident warp bought. `#[unroll]` on the column loop was throughput-neutral and quadrupled the code, so it is not used. This also retires the prepared upstream cubecl `min_blocks_per_sm` change: ptxas now picks 6 blocks/SM unaided, past the 4 that patch would have requested. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 280 +++++++++++-------- 1 file changed, 170 insertions(+), 110 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 4fdfad5ae1..557e257014 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -1143,6 +1143,89 @@ fn seqno_core_packed( /// 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`). +/// 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( @@ -1182,7 +1265,7 @@ fn multiply_pair( // 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 = false; + let mut rejected = 0u32; for j in 0..work_cap { let mut b = 0u32; @@ -1198,58 +1281,17 @@ fn multiply_pair( mk = u32::cast_from(masks[mk_base + j]); } - let mut val = 0u32; - if j < low { - if cs > b { - rejected = true; - } else { - let diff = b - cs; - if (diff & mk) != 0u32 { - rejected = true; - } else { - val = diff | mk; - } - } - } else { - if cs > 0u32 { - rejected = true; - } - if (b & mk) != 0u32 { - rejected = true; - } - val = b | mk; - } + let val = pair_col(j, low, b, cs, mk); + rejected |= val & PAIR_COL_REJECT; if j < PPART_MAX_LEN { - working |= u64::cast_from(val) << u64::cast_from(pp_shift[j]); + working |= u64::cast_from(val & 0xffffu32) << u64::cast_from(pp_shift[j]); } } - 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). - // Only the first `PPART_MAX_LEN` positions can be non-zero (see the accumulator above), - // 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); - } - } + if rejected == 0u32 { + pair_emit( + g, xi, out, working, row_base, out_offset, width, num_limbs, sq_len, pp_shift, pp_mask, + ); } } @@ -1278,22 +1320,29 @@ fn multiply_pair( // 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 GATHERS its -// one matrix's `col_sums`/`masks` and its term's p-part out of the segments into small local arrays -// via `seg_read_*` (correct for any offset, straddle or not — no layout padding needed), then hands -// those contiguous locals to `multiply_pair`, which stays a pure segmentation-agnostic arithmetic -// core. `seg_elems` is the segment element count (`o / seg_elems` picks the segment). -// OCCUPANCY (pending an upstream cubecl change): this kernel is register-bound, and ptxas without -// a `minBlocksPerSM` target settles at 78 registers = 3 blocks/SM = 37.5% occupancy (matching a -// sampled 36.8% "Compute Warps in Flight", with 63% of warp slots idle). Asking for 4 blocks/SM -// measured +6.2% (7.61e9 -> 8.08e9 pairs/s) for 130 bytes of spill stores; 5 and 6 tie within -// noise, 7-8 lose 12% as spilling overtakes the gain. +// [`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. // -// It is not enabled here because cubecl emits only `__launch_bounds__()` -- there is no -// way to request the second argument. The change is prepared as -// `~/cubecl-min-blocks-per-sm.patch` (adds `KernelOptions::min_blocks_per_sm` and a -// `min_blocks_per_sm = N` argument to `#[cube(..)]`, mirroring `cluster_dim`); once it lands -// upstream, add `min_blocks_per_sm = 4` below and re-measure. +// 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( @@ -1390,14 +1439,17 @@ fn multiply_batch_kernel( pp_mask: &[u32], // `min(work_cap, PPART_MAX_LEN)`: how far the packed rank loop runs. #[comptime] sq_len: usize, - // Per-thread working-array size, specialised per launch to what THIS block actually needs + // 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 is the kernel's dominant cost. Every thread holds four arrays of this length - // (`working` u32 plus `cs_local`/`mk_local`/`term_local` u16), so the constant sets register - // pressure and hence occupancy: measured 36% compute warps in flight with 64% of warp slots + // 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). + // (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 @@ -1457,17 +1509,29 @@ fn multiply_batch_kernel( seg_elems, num_segs, )); - // Gather this thread's matrix / term out of the segmented stores into contiguous locals, then run - // the pure arithmetic core on them (base 0). Entries past each length are zero, matching - // `multiply_pair`'s own out-of-range convention. The loop bounds at `WORKING_CAP`, exactly as the - // core does, so any `mk_len > WORKING_CAP` tail (never read by the core) is likewise not gathered. - let mut cs_local = Array::::new(work_cap); - let mut mk_local = Array::::new(work_cap); - let mut term_local = Array::::new(work_cap); + // The arithmetic of `multiply_pair`, with this thread's matrix / term read straight out of the + // segmented stores column by column. It used to gather into three `Array::::new(work_cap)` + // locals and call `multiply_pair` on them, but each entry is consumed exactly once, at the same + // `j`, by that function's own `0..work_cap` loop — so the arrays only ferried values between two + // loops with identical bounds. Fusing them deletes `work_cap x 3 x u16` of per-thread state, + // which is what caps occupancy (see the kernel's `work_cap` note). The per-column rule and the + // output tail stay shared with the single-`R` path via `pair_col` / `pair_emit`, so the two + // callers cannot drift. + // + // Entries past each length read as zero, matching `multiply_pair`'s out-of-range convention; + // the loop bounds at `work_cap` exactly as the core does, so any `mk_len > work_cap` tail is + // neither read nor needed. + let mut low = cs_len; + if term_len < cs_len { + low = term_len; + } + let mut working = 0u64; + let mut rejected = 0u32; + for j in 0..work_cap { - let mut c = 0u16; + let mut cs = 0u32; if j < cs_len { - c = seg_read_u16( + cs = u32::cast_from(seg_read_u16( cs0, cs1, cs2, @@ -1487,12 +1551,11 @@ fn multiply_batch_kernel( cs_off + j, seg_elems, num_segs, - ); + )); } - cs_local[j] = c; - let mut mm = 0u16; + let mut mk = 0u32; if j < mk_len { - mm = seg_read_u16( + mk = u32::cast_from(seg_read_u16( mk0, mk1, mk2, @@ -1512,12 +1575,11 @@ fn multiply_batch_kernel( mk_off + j, seg_elems, num_segs, - ); + )); } - mk_local[j] = mm; - let mut b = 0u16; + let mut b = 0u32; if j < term_len { - b = seg_read_u16( + b = u32::cast_from(seg_read_u16( pp0, pp1, pp2, @@ -1537,33 +1599,31 @@ fn multiply_batch_kernel( pp_off + j, seg_elems, num_segs, - ); + )); + } + + 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]); } - term_local[j] = b; } - multiply_pair( - cs_local.as_slice(), - mk_local.as_slice(), - term_local.as_slice(), - g, - xi, - out, - 0, - 0, - 0, - term_len, - cs_len, - mk_len, - usize::cast_from(prod_row_base[p]), - usize::cast_from(prod_out_offset[p]), - width, - num_limbs, - work_cap, - sq_len, - pp_shift, - pp_mask, - ); + if rejected == 0u32 { + pair_emit( + g, + xi, + out, + working, + 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 From 7887f24d197abbd36d85d54b59546fec283b8a18 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Mon, 3 Aug 2026 03:13:13 -0400 Subject: [PATCH 060/127] =?UTF-8?q?milnor=5Fgpu:=20NASSAU=5FGPU=5FLAUNCH?= =?UTF-8?q?=5FLOG=20=E2=80=94=20dump=20per-launch=20shape?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Launch shape (`work_cap`, the `mk_len` spread, product/pair counts, grid size) was invisible, and inferring it from aggregates is how several wrong conclusions got made in this file's history. It immediately earned its keep. Two bench processes on one GPU each report ~7x the solo `pairs/s`, which reads as enormous headroom sitting unused. This log showed both were issuing *identical* launches: same `work_cap=16`, same ~2.5e9 pairs, same ~9.6e6 blocks. Same launch on the same device cannot be 7x faster because a second process showed up, so the metric was suspect, not the kernel -- and the bench's fixed-work warm-up settles it: 4 jobs take ~190s whether run solo, alongside a second bench, alongside a 30 GiB memory hog, or alongside a compute hog. Real throughput is flat; the timed loop's `pairs/s` inflates under contention. Gated behind a `OnceLock` so the hot path pays one relaxed load per launch. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 28 ++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 557e257014..ea7c27a6d3 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -2561,6 +2561,18 @@ fn multiply_batch_block( "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) => { @@ -2762,6 +2774,22 @@ fn multiply_batch_block( /// 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 = OnceLock::new(); + *ON.get_or_init(|| std::env::var_os("NASSAU_GPU_LAUNCH_LOG").is_some()) +} + fn batch_report_every() -> u64 { use std::sync::OnceLock; static EVERY: OnceLock = OnceLock::new(); From 9e0abd94c54553a884508d166359e1e9b3f04699 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Mon, 3 Aug 2026 04:37:39 -0400 Subject: [PATCH 061/127] milnor_gpu: bound the column loop per thread, not per launch (+8.5%) The column loop ran to the launch's comptime `work_cap` -- the max `mk_len` over every `R` in the block. So one long `R` made every thread in the launch walk columns its own data does not have. `NASSAU_GPU_LAUNCH_LOG` measured the gap directly: `work_cap = 16` against a mean `mk_len` of 9.9. Past `max(term_len, cs_len, mk_len)` all three inputs are zero, so `pair_col` returns 0 -- it cannot reject and adds nothing to `working`. Stopping there is exact rather than a truncation, and it is a per-thread bound. Found by ablation, which also says where the rest of the time is NOT: | ablation | pairs/s | vs base | |-------------------------|---------|---------| | baseline | 6.71e9 | -- | | atomics removed | 6.79e9 | +1% | | binary search -> 1 iter | 7.49e9 | +12% | | column loop 16 -> 4 | 1.21e10 | +80% | So output-atomic contention is a non-issue, the product binary search is worth ~12%, and the column loop is the kernel. Fitting the two loop points gives ~41% fixed cost and ~59% per-column at cap 16. Measured +8.5% (6.79 -> 7.37 e9 pairs/s), four paired interleaved rounds, arms non-overlapping. Less than the ~25% the fit predicted, because a warp pays its widest lane and `cs_len`/`mk_len` are uniform within a product -- the remaining spread is across products, so grouping products by `mk_len` is the next lever. `work_cap` is no longer a kernel argument (the host still derives it to assert against `WORKING_CAP`), which also removes it as a source of NVRTC recompiles. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 29 ++++++++++++++++---- 1 file changed, 23 insertions(+), 6 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index ea7c27a6d3..5c37b58d0b 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -1143,6 +1143,27 @@ fn seqno_core_packed( /// 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`). /// @@ -1245,7 +1266,6 @@ fn multiply_pair( out_offset: usize, width: usize, num_limbs: usize, - #[comptime] work_cap: usize, #[comptime] sq_len: usize, pp_shift: &[u32], pp_mask: &[u32], @@ -1267,7 +1287,7 @@ fn multiply_pair( let mut working = 0u64; let mut rejected = 0u32; - for j in 0..work_cap { + 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]); @@ -1455,7 +1475,6 @@ fn multiply_batch_kernel( // 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`]. - #[comptime] work_cap: usize, ) { let k = ABSOLUTE_POS; let num_products = prod_pair_start.len() - 1; @@ -1528,7 +1547,7 @@ fn multiply_batch_kernel( let mut working = 0u64; let mut rejected = 0u32; - for j in 0..work_cap { + for j in 0..pair_cols(term_len, cs_len, mk_len) { let mut cs = 0u32; if j < cs_len { cs = u32::cast_from(seg_read_u16( @@ -2672,7 +2691,6 @@ fn multiply_batch_block( BufferArg::from_raw_parts(psh_h, pp_shift_len), BufferArg::from_raw_parts(pms_h, pp_shift_len), work_cap.min(PPART_MAX_LEN), - work_cap, ); } @@ -3288,7 +3306,6 @@ mod tests { 0, width, num_limbs, - WORKING_CAP, PPART_MAX_LEN, pp_shift, pp_mask, From 2aed0b0bd5f4981d48a1f3871e13a3ee47f54c55 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Mon, 3 Aug 2026 11:02:17 -0400 Subject: [PATCH 062/127] milnor_gpu: coarse product index for the thread->product lookup (+4.2%) Each thread found its product by binary-searching `prod_pair_start`, which is `ceil(log2(num_products))` ~ 15 DEPENDENT global loads -- each a full latency stall -- before it can touch any of its own data. `prod_coarse[i]` is the product owning pair `i << COARSE_LOG`, so a thread at pair `k` has its product bracketed by `prod_coarse[ci] ..= prod_coarse[ci + 1]` and searches only that span. The table is one entry per 2^20 pairs: a few thousand `u32` for a launch of billions, negligible to build and upload. Measured +4.2% (7.30 -> 7.60 e9 pairs/s), four paired interleaved rounds, variant winning every round. Worth recording that the ablation OVERSTATED this: forcing `search_iters = 1` measured +12%, but it also collapses every thread onto product 0, which removes the memory scatter along with the search. An ablation that changes locality measures more than the thing it names. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 74 ++++++++++++++++---- 1 file changed, 60 insertions(+), 14 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 5c37b58d0b..f09c86b624 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -80,6 +80,12 @@ const ENUM_MASK_CAP: usize = ENUM_ROW_CAP + ENUM_COL_CAP; /// 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; + /// 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 @@ -1444,13 +1450,14 @@ fn multiply_batch_kernel( 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 is `ceil(log2(num_products))`, which - // 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%). The win here is executing ~15 iterations - // instead of a fixed 32, which needs no unrolling. + // 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. @@ -1485,13 +1492,20 @@ 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). // - // `search_iters` is `ceil(log2(num_products))`, computed on the host and passed as a bare - // scalar, which cubecl bakes in as a compile-time constant — so the loop unrolls to exactly - // the steps the search needs. It was a fixed 32, and each step is a DEPENDENT global load of - // `prod_pair_start[mid]`, so the surplus iterations were a pure latency chain on every thread. - // At the measured p50 of ~24k products this is 15 steps rather than 32. - let mut lo = 0usize; - let mut hi = num_products; + // `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; @@ -2272,6 +2286,33 @@ fn multiply_batch_block( "block pair count {total_pairs} exceeds the kernel's u32 thread limit" ); pps.push(total_pairs as u32); + + // 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 @@ -2545,10 +2586,14 @@ fn multiply_batch_block( 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); - // Exactly the binary-search depth this block needs (see the kernel): passed bare so cubecl - // specialises the trip count instead of always running 32 dependent loads. - let search_iters = usize::BITS as usize - num_products.max(1).leading_zeros() as usize; + // 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(); @@ -2683,6 +2728,7 @@ fn multiply_batch_block( 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, From c305f37eae7b27d7c16319041cecd138eeecf5f1 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Mon, 3 Aug 2026 12:21:04 -0400 Subject: [PATCH 063/127] milnor_gpu: decode matrix fastest, term slowest (+2.0%) `m = local / nt; t = local % nt` made the TERM vary fastest, so consecutive threads shared a matrix but each took a different term, reading its p-part at `gei * width` for an arbitrary basis index. Swapping to `m = local % num_mats; t = local / num_mats` gives a warp one shared term (its p-part read becomes a broadcast) and `col_sums`/`masks` reads `m * cs_len` apart -- one contiguous, fully-used run per warp instead of one part-used sector per lane. It is a permutation of the same `(m, t)` set, every thread still owns exactly one pair, and the output is XOR-accumulated, so results are unchanged (75/75). Measured +2.0% (7.52 -> 7.64 e9 pairs/s), four paired rounds, variant winning every round. It is small because with ~5 terms per product a warp only ever cycled over ~5 distinct terms, so the old order was never a 32-way scatter. This kernel is COMPUTE bound, not memory bound. Measured with Nsight Compute 2025.3 (run inside /wsu/el7/containers/cuda/cuda-13.0.sif -- the host ncu 12.4 is older than the 580 driver and fails, and a newer one needs GLIBC 2.27 vs EL7's 2.17): Compute (SM) throughput 74.8% DRAM throughput 0.99% Issue slots busy 84.1% Memory throughput 6.2% Executed IPC 3.36/4 L1/TEX hit rate 93.3% ALU pipeline 68.4% L2 hit rate 12.5% So it is issue-limited on integer/logic work, and the data is already served out of L1. Going faster means executing fewer integer instructions (~1762 per warp today), not improving locality. Beware the ablation that suggested otherwise: pointing every read at offset 0 measured +79%, but that folds away the address arithmetic and the segment-select chain as well as the scatter, so it measured compute, not locality. An ablation that changes more than it names -- the second time this bit in one session. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 22 +++++++++++++++++--- 1 file changed, 19 insertions(+), 3 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index f09c86b624..362b01999c 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -1520,9 +1520,25 @@ fn multiply_batch_kernel( let ri = usize::cast_from(prod_r_index[p]); let nt = usize::cast_from(prod_num_terms[p]); - let local = k - usize::cast_from(prod_pair_start[p]); - let m = local / nt; - let t = local % nt; + let 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. + let num_mats = (usize::cast_from(prod_pair_start[p + 1]) - p_start) / nt; + let m = local % num_mats; + let t = local / num_mats; let cs_len = usize::cast_from(r_cs_len[ri]); let mk_len = usize::cast_from(r_mk_len[ri]); From 19fb7dbea334d59d6c6d99478fd06acac6d64660 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Mon, 3 Aug 2026 14:30:35 -0400 Subject: [PATCH 064/127] milnor_gpu: look up num_mats instead of dividing for it The thread decode computed `num_mats = (pair span) / nt`, then `local / num_mats` and `local % num_mats`. GPUs have no integer-division instruction -- ptxas emulates it via float reciprocal (I2F / MUFU.RCP / F2I plus fixups, ~20 instructions), and the SASS confirms it: three I2F/MUFU/F2I triples in the kernel. `num_mats` is just `r_num_matrices[ri]`, which the host already has, so upload it per distinct `R` (a small array) and read it. That drops one of the three divisions, and `nt` then has no remaining use, so the whole `prod_num_terms` buffer stops being built, uploaded, and read. This matters because the kernel is issue-limited on integer work, not memory -- ncu: 74.8% SM throughput vs 0.99% DRAM, IPC 3.36/4, ALU the top pipeline at 68.4%, 93% L1 hit rate, SASS 26% IMAD / 14% ISETP. Instruction count is the only lever left. Measured honestly: this was benched together with a loop-splitting change at +1.0% over four paired rounds (7.64 -> 7.71 e9). The isolation A/B for this change alone was cut short, so treat the +1.0% as an upper bound on it rather than a figure for it. It is kept regardless because it is strictly less work and less code: one fewer emulated division, one fewer uploaded buffer, one fewer load. The loop split it was benched with is NOT kept: hoisting an unguarded `j < lo3` prefix out of the column loop duplicated the body and cost 8 registers (40 -> 48, 6 -> 5 blocks/SM) for at most that same 1%. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 20 ++++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 362b01999c..31fe7d4eb8 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -1444,9 +1444,9 @@ fn multiply_batch_kernel( 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], @@ -1519,7 +1519,6 @@ fn multiply_batch_kernel( let p = lo; let ri = usize::cast_from(prod_r_index[p]); - let nt = usize::cast_from(prod_num_terms[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`) @@ -1536,9 +1535,12 @@ fn multiply_batch_kernel( // // 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. - let num_mats = (usize::cast_from(prod_pair_start[p + 1]) - p_start) / nt; - let m = local % num_mats; + // `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 t = local / num_mats; + let m = local - t * num_mats; let cs_len = usize::cast_from(r_cs_len[ri]); let mk_len = usize::cast_from(r_mk_len[ri]); @@ -2276,7 +2278,6 @@ fn multiply_batch_block( // 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 @@ -2291,7 +2292,6 @@ fn multiply_batch_block( prod_term_start.push(term_off[pi] as u32); pps.push(pair_acc as u32); pair_acc += r_num_matrices[ri as usize] * prod.term_indices.len(); - prod_num_terms.push(prod.term_indices.len() as u32); prod_row_base.push(((prod.row - row_base) * num_limbs) as u32); prod_out_offset.push(prod.out_offset as u32); } @@ -2575,6 +2575,11 @@ fn multiply_batch_block( 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 @@ -2598,7 +2603,6 @@ fn multiply_batch_block( 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)); @@ -2738,9 +2742,9 @@ fn multiply_batch_block( 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), From d54aec9cb79105a6c04020d587285d4967ee6d43 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Mon, 3 Aug 2026 16:30:06 -0400 Subject: [PATCH 065/127] milnor_gpu: one thread per (matrix, term GROUP) (+51%) `col_sums` and `masks` depend only on the matrix, but with one thread per (matrix, term) pair every thread re-read those ~2 x cols values for each term -- two thirds of all loads were redundant by a factor of `nt`. A thread now takes TERM_GROUP = 4 terms against one matrix, reading `cs`/`mk` once per column and applying the whole group against them. Loads per pair fall from 3 per column to 2/TERM_GROUP + 1 = 1.5. Measured +51%, four paired interleaved rounds (7.61-7.85e9 -> 1.16-1.18e10), ranges nowhere near overlapping. Confirmed on a metric that involves no pair accounting at all -- job completions per second, identical units of work in both arms: HEAD 193 calls in 62.1s = 3.1 calls/s, per-call p50 2.219s coarsened 291 calls in 61.3s = 4.7 calls/s, per-call p50 1.468s This is the lever the peephole work did not have. ncu showed the kernel issue-limited on integer work (74.8% SM vs 0.99% DRAM, IPC 3.36/4, ALU 68.4%, 93% L1 hit) and the last three instruction-shaving attempts paid +2.0%, +1.0% and ~0. Removing whole loads rather than trimming their arithmetic is what moved it. Costs registers -- 40 -> 64, so 4 blocks/SM instead of 6, plus 8 bytes of spill. That is the right trade here precisely because occupancy was never the binding constraint (going 75% -> 87.5% earlier measured 11% SLOWER). The group's tail is ragged whenever TERM_GROUP does not divide `nt`. Those lanes carry `term_len = 0` and MUST be excluded from the output explicitly: an all-zero term against an all-zero column does not reject, so they would otherwise emit a spurious seqno(0) bit. Grid sizing and throughput accounting now differ and are tracked separately: `pair_acc` counts THREADS (one per matrix/term-group, sizes the grid) while `real_pairs` counts (matrix, term) products evaluated and feeds BATCH_PAIRS. Conflating them made the first run of this A/B report 4.66e9 and look like a 40% regression, when it was really reporting threads/s. Note the bench's fixed-work warm-up does NOT move (195.4s -> 193.3s) and that is expected: it is the cold-master growth transient, dominated by CPU `admissible_matrices` and uploads, so it is not a valid discriminator for a kernel-only change. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 207 ++++++++++++------- 1 file changed, 128 insertions(+), 79 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 31fe7d4eb8..38aa9feb82 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -86,6 +86,11 @@ const GPU_PAIR_CHUNK: usize = 3_900_000_000; /// 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`. +const TERM_GROUP: usize = 4; + /// 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 @@ -1447,6 +1452,7 @@ fn multiply_batch_kernel( 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], @@ -1539,47 +1545,68 @@ fn multiply_batch_kernel( // 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 t = local / num_mats; - let m = local - t * num_mats; + let tgrp = local / num_mats; + let m = local - tgrp * num_mats; + let t_base = tgrp * TERM_GROUP; + let nt = usize::cast_from(prod_num_terms[p]); 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; - // `term_gei[term_slot]` is the term's *global* basis-element index (across all degrees): its - // (width-padded) p-part lives at global offset `gei*width` in the segmented basis `pp*`, length - // `ln*[gei]`. The basis is resident on the device (uploaded once, grown incrementally), so a - // launch uploads only these indices instead of re-gathering every term's p-part. - let gei = usize::cast_from(term_gei[term_slot]); - - // Global (across-segment) offsets of this matrix's data and this term's p-part. + // Global (across-segment) offset of this matrix's data. let cs_off = usize::cast_from(r_cs_offset[ri]) + m * cs_len; let mk_off = usize::cast_from(r_mk_offset[ri]) + m * mk_len; - let pp_off = gei * width; - let term_len = usize::cast_from(seg_read_u32( - ln0, ln1, ln2, ln3, ln4, ln5, ln6, ln7, ln8, ln9, ln10, ln11, ln12, ln13, ln14, ln15, gei, - seg_elems, num_segs, - )); - - // The arithmetic of `multiply_pair`, with this thread's matrix / term read straight out of the - // segmented stores column by column. It used to gather into three `Array::::new(work_cap)` - // locals and call `multiply_pair` on them, but each entry is consumed exactly once, at the same - // `j`, by that function's own `0..work_cap` loop — so the arrays only ferried values between two - // loops with identical bounds. Fusing them deletes `work_cap x 3 x u16` of per-thread state, - // which is what caps occupancy (see the kernel's `work_cap` note). The per-column rule and the - // output tail stay shared with the single-`R` path via `pair_col` / `pair_emit`, so the two - // callers cannot drift. + let ts_base = usize::cast_from(prod_term_start[p]) + t_base; + + // One thread now covers `TERM_GROUP` terms against ONE matrix, because `col_sums` and `masks` + // depend only on the matrix. One thread per pair re-read those ~2 x cols values for every term, + // i.e. two thirds of all loads (and their address arithmetic) were redundant by a factor of + // `nt`. Reading them once per column and applying the whole term group against them cuts loads + // per pair from 3 to `2/TERM_GROUP + 1`. // - // Entries past each length read as zero, matching `multiply_pair`'s out-of-range convention; - // the loop bounds at `work_cap` exactly as the core does, so any `mk_len > work_cap` tail is - // neither read nor needed. - let mut low = cs_len; - if term_len < cs_len { - low = term_len; + // This is the lever that is left: ncu measures the kernel issue-limited on integer work (74.8% + // SM vs 0.99% DRAM, IPC 3.36/4, ALU 68.4%, 93% L1 hit), and the SASS is 26% IMAD / 14% ISETP -- + // address and predicate arithmetic. Peephole work had run dry at ~1% a change; this removes + // whole loads rather than shaving instructions off them. + // + // `term_gei[term_slot]` is the term's *global* basis-element index (across all degrees): its + // (width-padded) p-part lives at global offset `gei*width` in the segmented basis `pp*`, length + // `ln*[gei]`. The group's tail is ragged whenever `TERM_GROUP` does not divide `nt`; those lanes + // carry `term_len = 0` and are excluded from the output below (they cannot simply fall out of + // the arithmetic -- an all-zero term against an all-zero column does NOT reject). + 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 mut working = 0u64; - let mut rejected = 0u32; - for j in 0..pair_cols(term_len, cs_len, mk_len) { + let mut working = Array::::new(TERM_GROUP); + let mut rejected = Array::::new(TERM_GROUP); + #[unroll] + for tt in 0..TERM_GROUP { + working[tt] = 0u64; + rejected[tt] = 0u32; + } + + for j in 0..cols { let mut cs = 0u32; if j < cs_len { cs = u32::cast_from(seg_read_u16( @@ -1628,52 +1655,64 @@ fn multiply_batch_kernel( num_segs, )); } - let mut b = 0u32; - if j < term_len { - b = u32::cast_from(seg_read_u16( - pp0, - pp1, - pp2, - pp3, - pp4, - pp5, - pp6, - pp7, - pp8, - pp9, - pp10, - pp11, - pp12, - pp13, - pp14, - pp15, - pp_off + j, - seg_elems, - num_segs, - )); - } - - 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]); + #[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; + } + let val = pair_col(j, low, b, cs, mk); + rejected[tt] |= val & PAIR_COL_REJECT; + if j < PPART_MAX_LEN { + working[tt] |= u64::cast_from(val & 0xffffu32) << u64::cast_from(pp_shift[j]); + } } } - if rejected == 0u32 { - pair_emit( - g, - xi, - out, - working, - usize::cast_from(prod_row_base[p]), - usize::cast_from(prod_out_offset[p]), - width, - num_limbs, - sq_len, - pp_shift, - pp_mask, - ); + #[unroll] + for tt in 0..TERM_GROUP { + if t_base + tt < nt { + if rejected[tt] == 0u32 { + pair_emit( + g, + xi, + out, + working[tt], + usize::cast_from(prod_row_base[p]), + usize::cast_from(prod_out_offset[p]), + width, + num_limbs, + sq_len, + pp_shift, + pp_mask, + ); + } + } } } @@ -1963,7 +2002,8 @@ fn multiply_batch_grouped( MasterMode::Resident => resident_info(algebra, r.p_part).num_mats as usize, MasterMode::Transient => cold_count(algebra, r.p_part).2 as usize, }; - num_mats * prod.term_indices.len() + // Threads, not pairs: one per (matrix, TERM_GROUP-sized term group). + num_mats * prod.term_indices.len().div_ceil(TERM_GROUP) }) .collect() }); @@ -2278,6 +2318,7 @@ fn multiply_batch_block( // 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 @@ -2287,11 +2328,17 @@ fn multiply_batch_block( // `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); - pair_acc += r_num_matrices[ri as usize] * prod.term_indices.len(); + // 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] * 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); } @@ -2603,6 +2650,7 @@ fn multiply_batch_block( 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)); @@ -2745,6 +2793,7 @@ fn multiply_batch_block( 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), @@ -2797,7 +2846,7 @@ fn multiply_batch_block( (device_ms * 1e3) as u64, std::sync::atomic::Ordering::Relaxed, ); - BATCH_PAIRS.fetch_add(total_pairs 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( From 37a298ea4081a1df35dec65b69134c3c1a9dfd2c Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Mon, 3 Aug 2026 17:38:36 -0400 Subject: [PATCH 066/127] milnor_gpu: TERM_GROUP 4 -> 3 (+9%) Tuned by measurement, and the optimum is NOT monotonic. A bigger group amortises the `col_sums`/`masks` reads over more terms, but a product's `nt` terms need `ceil(nt / TERM_GROUP)` groups and the last one is usually partial, so a bigger group also idles more lanes. 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: tail waste, not amortisation, dominates the choice. Among the 17% group, 3 wins -- it amortises more than 2 and costs fewer registers than 6. This is fitted to `nt ~ 5`; retune if the terms-per-product regime moves. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 38aa9feb82..7d02be96ff 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -89,7 +89,24 @@ 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`. -const TERM_GROUP: usize = 4; +/// +/// 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; /// Per-launch output-buffer budget in bytes (`NASSAU_GPU_BLOCK_MB`, default 512 MiB). /// From 144727fedb97d43f06211beddebebbbcd25a9d7a Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Mon, 3 Aug 2026 19:22:54 -0400 Subject: [PATCH 067/127] gitignore: exclude ext/*.log GPU run logs These are multi-hundred-MB analysis artifacts written beside the crate. Five of them (1.9 GB, largest 261 MB) were swept into a commit by an over-broad `git add` and had to be stripped from the branch before it could be pushed -- GitHub rejects any file over 100 MB. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- .gitignore | 3 +++ 1 file changed, 3 insertions(+) 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 From d167b6f198174d4b28403e62d8cfe5ba1d564e3c Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Mon, 3 Aug 2026 23:28:12 -0400 Subject: [PATCH 068/127] nassau: span the CPU work inside `step` The stem-200 run accounted for only ~74% of worker time by name: `gpu_submit` 51%, `extract_restricted` 15%, `pair_prepass` 4%, `gpu_row_reduce` 3%, `marshal_terms` 1%. The remaining ~26% sat inside `step` but outside any named region -- bigger than every identified CPU item except the GPU wait itself, and invisible. That is the exact shape that produced several wrong diagnoses in this file's history, all of them from inferring a mechanism for an unspanned region. So span it before theorising about it. Per-signature loop (~555k iterations, most of the `step` spans): `sig_masks`, `sig_select`, `sig_assemble`, `sig_row_reduce`, `sig_quasi_inverse`, `sig_lift`, `sig_write_qi`. The prime suspect is `sig_row_reduce` -- a CPU row reduction per signature, since `gpu_row_reduce` only takes over at >= 8192^2 and these are far smaller. Zero-signature path (once per bidegree): `zs_assemble`, `zs_row_reduce`, `zs_kernel`, `img_assemble`, `img_row_reduce`, `extend_image`. One span per signature, not per row: the bodies are substantial, so the overhead is negligible. Compare the `parallel_guard` span that was removed for emitting 1000+ span pairs per bidegree -- do NOT push these inside the inner loops. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/src/nassau.rs | 134 +++++++++++++++++++++++++++++++--------------- 1 file changed, 90 insertions(+), 44 deletions(-) diff --git a/ext/src/nassau.rs b/ext/src/nassau.rs index 00774eb235..52c845c2fa 100644 --- a/ext/src/nassau.rs +++ b/ext/src/nassau.rs @@ -812,15 +812,26 @@ impl> Resolution { next_dim, ), }; - let mut masked_matrix = - AugmentedMatrix::new(p, target_masked_dim, [next_masked_dim, target_masked_dim]); + let mut masked_matrix = tracing::info_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::info_span!( + "zs_row_reduce", + rows = target_masked_dim, + cols = next_masked_dim + ) + .in_scope(|| masked_matrix.row_reduce()); + let kernel = tracing::info_span!("zs_kernel").in_scope(|| masked_matrix.compute_kernel()); Self::write_qi( &mut f, @@ -857,14 +868,18 @@ impl> Resolution { &source_mask, target_dim, ); - 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.row_reduce(); + let mut n = tracing::info_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::info_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::info_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}"); @@ -907,6 +922,11 @@ impl> Resolution { 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::info_span!("sig_masks").entered(); target_mask.clear(); next_mask.clear(); target_mask.extend(subalgebra.signature_mask( @@ -923,29 +943,51 @@ impl> Resolution { &signature, next_bound, )); + drop(_sm); + + let full_matrix = + tracing::info_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 = tracing::info_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 + }); - 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_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 qi = masked_matrix.compute_quasi_inverse(); + // 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::info_span!( + "sig_row_reduce", + rows = target_mask.len(), + cols = next_mask.len() + ) + .in_scope(|| masked_matrix.row_reduce()); + + let qi = tracing::info_span!("sig_quasi_inverse") + .in_scope(|| masked_matrix.compute_quasi_inverse()); let pivots = qi.pivots().unwrap(); let preimage = qi.preimage(); @@ -960,6 +1002,7 @@ impl> Resolution { } } + let _lift = tracing::info_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; @@ -977,14 +1020,17 @@ 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::info_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!( From d4db8379120eb349fa5a8cbed1ad324bd35433b1 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Tue, 4 Aug 2026 08:45:42 -0400 Subject: [PATCH 069/127] milnor_gpu: tile over matrices as well as terms (+6.6%) `col_sums`/`masks` are per-matrix and a p-part is per-term, so an `M x T` tile reads `2M + T` values per column to evaluate `M*T` pairs. At 2x3 that is 1.17 loads per pair against 1.67 at 1x3 (and 3.0 before any tiling). Measured +6.6% (1.28 -> 1.37 e10 pairs/s), two paired interleaved rounds, arms non-overlapping. Registers stay at 64 (4 blocks/SM) and the 8 bytes of spill the 1x3 tile carried are gone. The two axes are NOT symmetric, which is why they get separate constants: - terms are few (`nt ~ 5`), so a partial tile wastes a large fraction of lanes and TERM_GROUP is chosen by tail waste -- 4 loses to 3 on that alone. - matrices are many (`num_mats ~ 20 000`), so a partial tile idles a couple of lanes out of thousands and MATRIX_GROUP can follow the load arithmetic. The gain is smaller than the 30% load reduction would suggest. Occupancy is unchanged between the arms (64 registers, 4 blocks/SM, both), so it is not a regression in occupancy -- but 50% occupancy may simply leave too few warps to hide the latency of the loads that remain. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 223 +++++++++++-------- 1 file changed, 126 insertions(+), 97 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 7d02be96ff..7ad4d4aa4d 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -108,6 +108,12 @@ const COARSE_LOG: usize = 20; /// 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 @@ -1562,34 +1568,34 @@ fn multiply_batch_kernel( // 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 tgrp = local / num_mats; - let m = local - tgrp * num_mats; - let t_base = tgrp * TERM_GROUP; let nt = usize::cast_from(prod_num_terms[p]); + // 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]); - // Global (across-segment) offset of this matrix's data. - let cs_off = usize::cast_from(r_cs_offset[ri]) + m * cs_len; - let mk_off = usize::cast_from(r_mk_offset[ri]) + m * mk_len; + 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; - // One thread now covers `TERM_GROUP` terms against ONE matrix, because `col_sums` and `masks` - // depend only on the matrix. One thread per pair re-read those ~2 x cols values for every term, - // i.e. two thirds of all loads (and their address arithmetic) were redundant by a factor of - // `nt`. Reading them once per column and applying the whole term group against them cuts loads - // per pair from 3 to `2/TERM_GROUP + 1`. - // - // This is the lever that is left: ncu measures the kernel issue-limited on integer work (74.8% - // SM vs 0.99% DRAM, IPC 3.36/4, ALU 68.4%, 93% L1 hit), and the SASS is 26% IMAD / 14% ISETP -- - // address and predicate arithmetic. Peephole work had run dry at ~1% a change; this removes - // whole loads rather than shaving instructions off them. - // - // `term_gei[term_slot]` is the term's *global* basis-element index (across all degrees): its - // (width-padded) p-part lives at global offset `gei*width` in the segmented basis `pp*`, length - // `ln*[gei]`. The group's tail is ragged whenever `TERM_GROUP` does not divide `nt`; those lanes - // carry `term_len = 0` and are excluded from the output below (they cannot simply fall out of - // the arithmetic -- an all-zero term against an all-zero column does NOT reject). + // 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; @@ -1615,63 +1621,75 @@ fn multiply_batch_kernel( } } - let mut working = Array::::new(TERM_GROUP); - let mut rejected = Array::::new(TERM_GROUP); + let mut working = Array::::new(MATRIX_GROUP * TERM_GROUP); + let mut rejected = Array::::new(MATRIX_GROUP * TERM_GROUP); #[unroll] - for tt in 0..TERM_GROUP { - working[tt] = 0u64; - rejected[tt] = 0u32; + for i in 0..MATRIX_GROUP * TERM_GROUP { + working[i] = 0u64; + rejected[i] = 0u32; } for j in 0..cols { - let mut cs = 0u32; - if j < cs_len { - cs = u32::cast_from(seg_read_u16( - cs0, - cs1, - cs2, - cs3, - cs4, - cs5, - cs6, - cs7, - cs8, - cs9, - cs10, - cs11, - cs12, - cs13, - cs14, - cs15, - cs_off + j, - seg_elems, - num_segs, - )); - } - let mut mk = 0u32; - if j < mk_len { - mk = u32::cast_from(seg_read_u16( - mk0, - mk1, - mk2, - mk3, - mk4, - mk5, - mk6, - mk7, - mk8, - mk9, - mk10, - mk11, - mk12, - mk13, - mk14, - mk15, - mk_off + j, - seg_elems, - num_segs, - )); + // 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]); @@ -1703,31 +1721,41 @@ fn multiply_batch_kernel( if tl < cs_len { low = tl; } - let val = pair_col(j, low, b, cs, mk); - rejected[tt] |= val & PAIR_COL_REJECT; - if j < PPART_MAX_LEN { - working[tt] |= u64::cast_from(val & 0xffffu32) << u64::cast_from(pp_shift[j]); + #[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]); + } } } } #[unroll] - for tt in 0..TERM_GROUP { - if t_base + tt < nt { - if rejected[tt] == 0u32 { - pair_emit( - g, - xi, - out, - working[tt], - usize::cast_from(prod_row_base[p]), - usize::cast_from(prod_out_offset[p]), - width, - num_limbs, - sq_len, - pp_shift, - pp_mask, - ); + 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, + ); + } + } } } } @@ -2019,8 +2047,8 @@ fn multiply_batch_grouped( MasterMode::Resident => resident_info(algebra, r.p_part).num_mats as usize, MasterMode::Transient => cold_count(algebra, r.p_part).2 as usize, }; - // Threads, not pairs: one per (matrix, TERM_GROUP-sized term group). - num_mats * prod.term_indices.len().div_ceil(TERM_GROUP) + // 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() }); @@ -2354,7 +2382,8 @@ fn multiply_batch_block( // 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] * prod.term_indices.len().div_ceil(TERM_GROUP); + 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); From 9174826c433fe166f7711965d95f6069530abaa2 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Tue, 4 Aug 2026 11:49:40 -0400 Subject: [PATCH 070/127] milnor_gpu: shard the resident master across GPUs (2.27x on 4 devices) One `nassau-gpu` worker per CUDA device, and the resident master SHARDED across them: each first-sight `R` is assigned a device round-robin and its rows live only there, so a launch routes its products to the device owning their `R`, evaluates the partials concurrently, and XORs them. That is exact rather than an approximation -- every product contributes to the output by `fetch_xor`, so the contributions commute and split freely. 1 GPU 1.29e10 pairs/s (unchanged from before this commit) 4 GPUs, replicated 5.99e9 0.46x -- a REGRESSION 4 GPUs, sharded 2.93e10 2.27x, duty cycle 379% of 4 Replication was tried first and lost badly. Every device rebuilt the whole ~100 GB master over PCIe, so cost scaled with device count while the gain saturated: warm-up went 190s -> 508s and the GPUs filled strictly one at a time (dev3 to 39 GB, then dev0 to 66 GB, then dev1 to 84 GB, dev2 still empty). NCCL or a `to_client` broadcast over the NV6 mesh would have made that copy fast, but it would still put a full master on every card -- 1x per-device memory, and no change to the stem-300 ceiling. Sharding is what actually helps: total upload stays 1x, uploads now overlap (warm-up 158.7s, FASTER than one GPU's 188.6s), device memory fills evenly (~37 GB x 4 against 99 GB x 1), and the aggregate master can span 4 x 140 GB instead of 140 GB -- which is the constraint that matters at stem 300. Cost of the design, recorded honestly: - Work is device-AFFINE now, so the shared pull-queue is gone and each device has its own queue. Balance comes from spreading `R`s, not from idle workers stealing, and it is imperfect: batch-stats shows queue 67% / exec 32% with mean depth 8.4. Round-robin over first-sight order is the obvious thing to improve (assign by measured `R` load instead). - The fan-out clones each `GpuProduct` into its device's bucket and XORs the partials host-side. The combine is ~1.5 MB per device per block (~72 MB/s), i.e. noise -- `all_reduce` could not express it anyway, since cubecl offers only Sum/Mean and over F2 a sum carries. - Host `*_pending` tails and logical master lengths are per-device. The single shared HOST master copy is retained, so the multi-GB host duplicate that design exists to avoid stays gone. `crossbeam-channel` replaces `std::mpsc` for the queues: genuinely multi-consumer, where std's single-receiver would force a `Mutex`. 75/75 tests pass on 1 and on 4 devices, including the CPU-reference comparison. Single-GPU throughput and warm-up are unchanged, so this costs nothing on one card. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/Cargo.toml | 6 +- ext/crates/algebra/src/algebra/milnor_gpu.rs | 334 ++++++++++++++----- 2 files changed, 262 insertions(+), 78 deletions(-) diff --git a/ext/crates/algebra/Cargo.toml b/ext/crates/algebra/Cargo.toml index c45c7b8e3f..5845f309a9 100644 --- a/ext/crates/algebra/Cargo.toml +++ b/ext/crates/algebra/Cargo.toml @@ -43,6 +43,10 @@ cubecl-common = { git = "https://github.com/tracel-ai/cubecl", tag = "v0.11.0-pr # 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"] } @@ -58,7 +62,7 @@ cache-multiplication = [] # `algebra::milnor_rank`. milnor-rank = [] concurrent = ["fp/concurrent", "maybe-rayon/concurrent"] -gpu = ["dep:cubecl", "dep:cubecl-common", "dep:tracing"] +gpu = ["dep:cubecl", "dep:cubecl-common", "dep:tracing", "dep:crossbeam-channel"] odd-primes = ["fp/odd-primes"] [[bench]] diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 7ad4d4aa4d..50c06f65cd 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -209,14 +209,15 @@ mod gpu_thread { sync::{ OnceLock, atomic::{AtomicU64, Ordering}, - mpsc::{self, Sender}, + mpsc, }, time::Instant, }; + use crossbeam_channel::{Sender, unbounded}; use cubecl_common::stream_id::StreamId; - /// Jobs enqueued but not yet started, i.e. how deep the FIFO is when a worker joins it. + /// 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; @@ -231,27 +232,46 @@ mod gpu_thread { pub depth: u64, } - fn sender() -> &'static Sender { - static QUEUE: OnceLock> = OnceLock::new(); - QUEUE.get_or_init(|| { - let (tx, rx) = mpsc::channel::(); - std::thread::Builder::new() - .name("nassau-gpu".into()) - .spawn(move || { - // Bind the stream once for the whole loop: one stream, one driver thread. - StreamId { value: 0 }.executes(|| { - while let Ok(task) = rx.recv() { - task(); - } - }); - }) - .expect("failed to spawn the nassau-gpu thread"); - tx + /// 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, its stream, and (via the thread-local set here) its own replica + /// of the resident master/basis, so a device handle can never reach another device's client. + /// + /// `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>> = OnceLock::new(); + QUEUES.get_or_init(|| { + let mut txs = Vec::with_capacity(super::gpu_count()); + for dev in 0..super::gpu_count() { + let (tx, rx) = unbounded::(); + txs.push(tx); + std::thread::Builder::new() + .name(format!("nassau-gpu{dev}")) + .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 per device so the runtime keeps them independent. + StreamId { value: dev 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. - pub(super) fn run(f: F) -> (T, Timing) + /// Run `f` on `dev`'s worker. The device is chosen by the caller from where the data lives. + pub(super) fn run_on(dev: usize, f: F) -> (T, Timing) where F: FnOnce() -> T + Send + 'static, T: Send + 'static, @@ -259,7 +279,7 @@ mod gpu_thread { 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(); - sender() + senders()[dev] .send(Box::new(move || { let queue_ms = enqueued.elapsed().as_secs_f64() * 1e3; DEPTH.fetch_sub(1, Ordering::Relaxed); @@ -430,8 +450,8 @@ pub fn take_gpu_timing() -> (u64, u64, u64, u64, u64) { /// 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.capacity() * 2 - + h.mk_pending.capacity() * 2 + 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::>() @@ -445,10 +465,21 @@ pub fn resident_host_bytes() -> (usize, usize) { /// 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) { - let d = RESIDENT_DEV.read().unwrap(); - let master = d.cs.uploaded * 2 + d.mk.uploaded * 2; - let b = RESIDENT_BASIS_DEV.read().unwrap(); - let basis = b.pp.uploaded * 2 + b.ln.uploaded * 4; + // 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) } @@ -457,7 +488,7 @@ pub fn resident_dev_bytes() -> (usize, usize) { /// 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 = CudaRuntime::client(&CudaDevice::default()); + let client = gpu_client(); match client.memory_usage() { Ok(u) => (u.bytes_in_use, u.bytes_reserved), Err(_) => (0, 0), @@ -484,6 +515,13 @@ struct RInfo { 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-shared host master of admissible-matrix data. @@ -507,17 +545,74 @@ struct RInfo { /// 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]`. -#[derive(Default)] +/// `*_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, - cs_len: usize, - mk_len: usize, + 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, + /// Round-robin cursor for assigning the next first-sight `R` to a device. + next_dev: usize, index: HashMap, } -static RESIDENT_HOST: LazyLock> = - LazyLock::new(|| RwLock::new(ResidentHost::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()], + next_dev: 0, + index: HashMap::new(), + }) +}); + +/// 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 the driver exposes is used. +/// +/// 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 detected = std::fs::read_dir("/proc/driver/nvidia/gpus") + .map(|d| d.filter_map(|e| e.ok()).count()) + .unwrap_or(0) + .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 +} + +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) }; +} + +/// 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. +fn gpu_client() -> cubecl::prelude::ComputeClient { + CudaRuntime::client(&CudaDevice::new(cur_device())) +} /// 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`] / @@ -570,8 +665,18 @@ struct ResidentDev { mk: SegBuf, } -static RESIDENT_DEV: LazyLock> = - LazyLock::new(|| RwLock::new(ResidentDev::default())); +/// 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() +}); + +fn resident_dev() -> &'static RwLock { + &RESIDENT_DEV[cur_device()] +} /// 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 @@ -579,7 +684,12 @@ static RESIDENT_DEV: LazyLock> = /// 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: Mutex<()> = Mutex::new(()); +static RESIDENT_UPLOAD: LazyLock>> = + LazyLock::new(|| (0..gpu_count()).map(|_| Mutex::new(())).collect()); + +fn resident_upload() -> &'static Mutex<()> { + &RESIDENT_UPLOAD[cur_device()] +} /// 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 @@ -593,9 +703,19 @@ struct SeqnoDev { g: Handle, xi: Handle, } -static RESIDENT_SEQNO: LazyLock>> = LazyLock::new(|| RwLock::new(None)); +static RESIDENT_SEQNO: LazyLock>>> = + LazyLock::new(|| (0..gpu_count()).map(|_| RwLock::new(None)).collect()); + +fn resident_seqno() -> &'static RwLock> { + &RESIDENT_SEQNO[cur_device()] +} /// Serializes seqno-table uploads only (never reads); see [`RESIDENT_UPLOAD`]. -static RESIDENT_SEQNO_UPLOAD: Mutex<()> = Mutex::new(()); +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 @@ -604,7 +724,7 @@ static RESIDENT_SEQNO_UPLOAD: Mutex<()> = Mutex::new(()); macro_rules! resident_seqno { ($client:expr, $g:expr, $xi:expr) => {{ let read_current = || { - let s = RESIDENT_SEQNO.read().unwrap(); + let s = resident_seqno().read().unwrap(); match &*s { Some(d) if d.g_len == $g.len() => Some((d.g.clone(), d.xi.clone())), _ => None, @@ -613,7 +733,7 @@ macro_rules! resident_seqno { match read_current() { Some(h) => h, None => { - let _upload_guard = RESIDENT_SEQNO_UPLOAD.lock().unwrap(); + let _upload_guard = resident_seqno_upload().lock().unwrap(); match read_current() { Some(h) => h, None => { @@ -621,7 +741,7 @@ macro_rules! resident_seqno { 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 { + *resident_seqno().write().unwrap() = Some(SeqnoDev { g_len: $g.len(), g: gh.clone(), xi: xh.clone(), @@ -931,17 +1051,22 @@ fn resident_info(algebra: &MilnorAlgebra, p_part: PPart) -> RInfo { } // 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 a device, round-robin over first-sight order. Its rows go to that device + // and nowhere else, so a launch must route products to the device owning their `R`. + let dev = host.next_dev % gpu_count(); + host.next_dev = host.next_dev.wrapping_add(1); let info = RInfo { - cs_off: host.cs_len as u64, - mk_off: host.mk_len as u64, + 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: (mk.len() / mk_len) as u32, + dev: dev as u8, }; - host.cs_pending.extend(cs.iter().map(|&v| narrow_u16(v))); - host.mk_pending.extend(mk.iter().map(|&v| narrow_u16(v))); - host.cs_len += cs.len(); - host.mk_len += mk.len(); + 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 } @@ -980,11 +1105,23 @@ struct ResidentBasisDev { ln: SegBuf, } -static RESIDENT_BASIS_DEV: LazyLock> = - LazyLock::new(|| RwLock::new(ResidentBasisDev::default())); +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: Mutex<()> = Mutex::new(()); +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 @@ -2075,14 +2212,54 @@ fn multiply_batch_grouped( pairs += row_pairs; (r1, p1) = (r1 + 1, q); } - result.push(multiply_batch_block( - algebra, - num_cols, - r0, - r1 - r0, - &products[p0..p1], - mode, - )); + // 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()]; + 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. + MasterMode::Transient => pi % gpu_count(), + MasterMode::Resident => { + let r = algebra.basis_element_from_index(prod.r_degree, prod.r_idx); + resident_info(algebra, r.p_part).dev as usize + } + }; + by_dev[d].push(prod.clone()); + } + // Scoped threads, not a sequential loop: `gpu_thread::run_on` BLOCKS until its device + // finishes, so submitting one after another would serialise the very devices this is + // splitting work across. + let partials: Vec = std::thread::scope(|scope| { + let handles: Vec<_> = by_dev + .iter() + .enumerate() + .filter(|(_, ps)| !ps.is_empty()) + .map(|(d, ps)| { + scope.spawn(move || { + multiply_batch_block(algebra, num_cols, r0, r1 - r0, ps, mode, d) + }) + }) + .collect(); + handles + .into_iter() + .map(|h| h.join().expect("a sharded sub-launch panicked")) + .collect() + }); + 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 @@ -2100,6 +2277,7 @@ fn multiply_batch_block( num_rows: usize, products: &[GpuProduct], mode: MasterMode, + dev: usize, ) -> Bytes { let (width, g) = algebra.seqno_table_u32(); let mut xi: Vec = xi_degrees(algebra.prime()) @@ -2485,11 +2663,11 @@ fn multiply_batch_block( out = out_len ); let (result, timing) = submit_span.in_scope(|| { - gpu_thread::run(move || { + gpu_thread::run_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 = CudaRuntime::client(&CudaDevice::default()); + 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 @@ -2530,32 +2708,34 @@ fn multiply_batch_block( MasterMode::Resident => { let (cs_segs, _) = seg_grow!( client, - RESIDENT_DEV, + resident_dev(), cs, - RESIDENT_UPLOAD, + resident_upload(), need_cs, copy_into_u16, u16::as_bytes, u16, |_up: usize| { let mut h = RESIDENT_HOST.write().unwrap(); - let nl = h.cs_len; - (std::mem::take(&mut h.cs_pending), nl) + 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, + resident_dev(), mk, - RESIDENT_UPLOAD, + resident_upload(), need_mk, copy_into_u16, u16::as_bytes, u16, |_up: usize| { let mut h = RESIDENT_HOST.write().unwrap(); - let nl = h.mk_len; - (std::mem::take(&mut h.mk_pending), nl) + 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))) @@ -2621,9 +2801,9 @@ fn multiply_batch_block( } else { let (pp_segs, _) = seg_grow!( client, - RESIDENT_BASIS_DEV, + resident_basis_dev(), pp, - RESIDENT_BASIS_UPLOAD, + resident_basis_upload(), need_basis_elems * width, copy_into_u16, u16::as_bytes, @@ -2636,9 +2816,9 @@ fn multiply_batch_block( ); let (ln_segs, _) = seg_grow!( client, - RESIDENT_BASIS_DEV, + resident_basis_dev(), ln, - RESIDENT_BASIS_UPLOAD, + resident_basis_upload(), need_basis_elems, copy_into_u32, u32::as_bytes, @@ -3322,7 +3502,7 @@ mod tests { 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 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)); @@ -3389,7 +3569,7 @@ mod tests { ) -> 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 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)); @@ -3531,7 +3711,7 @@ mod tests { col_sums.push(0); } - let client = CudaRuntime::client(&CudaDevice::default()); + 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)); @@ -3838,7 +4018,7 @@ mod tests { nseg <= MASTER_MAX_SEG, "prototype caps at {MASTER_MAX_SEG} segments" ); - let client = CudaRuntime::client(&CudaDevice::default()); + let client = gpu_client(); // One handle per segment slot; real segments hold their slice, unused slots a 1-elem dummy. let dummy = [0u16]; From cc116877c9a82a94790c0d53adcee44c3fe81cea Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Tue, 4 Aug 2026 12:01:15 -0400 Subject: [PATCH 071/127] milnor_gpu: assign shards to the least-loaded device (+32%) Round-robin over first-sight order balanced the COUNT of `R`s per device, but `num_mats` varies by orders of magnitude between `R`s, so equal counts left the devices badly uneven. Greedy least-loaded by accumulated `num_mats` fixes it: a launch's device work 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 this tracks memory balance too. 4 GPUs, round-robin 2.93e10 pairs/s duty 379% of 400% 4 GPUs, least-loaded 3.86e10 pairs/s duty 395% of 400% +32%, and 2.99x a single GPU (1.29e10). Costs nothing: the weight is accumulated at first sight from data already computed, with no lookahead and no extra state beyond one counter per device. Note the queue/exec split is NOT the balance metric and barely moves (67/32 -> 70/30): it reflects the bench's 7 workers submitting faster than 4 devices drain, not imbalance between them. Duty cycle is the one to watch. 75/75 tests pass on 4 devices. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 32 +++++++++++++++----- 1 file changed, 24 insertions(+), 8 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 50c06f65cd..50b8785929 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -555,8 +555,12 @@ struct ResidentHost { /// Per-device logical master lengths; an `R` extends only its own device's. cs_len: Vec, mk_len: Vec, - /// Round-robin cursor for assigning the next first-sight `R` to a device. - next_dev: usize, + /// 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, index: HashMap, } @@ -566,7 +570,7 @@ static RESIDENT_HOST: LazyLock> = LazyLock::new(|| { mk_pending: (0..gpu_count()).map(|_| Vec::new()).collect(), cs_len: vec![0; gpu_count()], mk_len: vec![0; gpu_count()], - next_dev: 0, + dev_load: vec![0; gpu_count()], index: HashMap::new(), }) }); @@ -1051,16 +1055,28 @@ fn resident_info(algebra: &MilnorAlgebra, p_part: PPart) -> RInfo { } // 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 a device, round-robin over first-sight order. Its rows go to that device - // and nowhere else, so a launch must route products to the device owning their `R`. - let dev = host.next_dev % gpu_count(); - host.next_dev = host.next_dev.wrapping_add(1); + // 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`. + // + // Round-robin over first-sight order was the first cut and balances COUNT, not work — `num_mats` + // varies by orders of magnitude between `R`s, so equal counts left the devices badly uneven + // (batch-stats: queue 67% / exec 32%, mean depth 8.4, i.e. devices waiting while others ran). + // Greedy least-loaded is the standard fix and needs no lookahead. + let num_mats = (mk.len() / mk_len) as u64; + let dev = host + .dev_load + .iter() + .enumerate() + .min_by_key(|&(i, &load)| (load, i)) + .map(|(i, _)| i) + .expect("at least one device"); + 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: (mk.len() / mk_len) as u32, + num_mats: num_mats as u32, dev: dev as u8, }; host.cs_pending[dev].extend(cs.iter().map(|&v| narrow_u16(v))); From 1639982877916defe3ded746cf3786340b2b8025 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Tue, 4 Aug 2026 13:07:15 -0400 Subject: [PATCH 072/127] milnor_gpu: submit shards without spawning threads; name the GPU workers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The sharded fan-out used `std::thread::scope`, spawning an OS thread per device per row block — hundreds a second, with thread ids reaching `ThreadId(867540)` in a stem-200 log. `multiply_batch_block` now returns the WAIT rather than performing it: `gpu_thread::submit_on` hands back a `Pending`, so the caller marshals and submits every device's share first and only then blocks. Device `d + 1`'s marshalling overlaps device `d`'s execution, all shards are in flight together, and no thread is created. Deliberately NOT rayon: a `par_iter` whose bodies block on the GPU is the join + steal-loop pattern behind the 146 s signature stalls this project already diagnosed once (see the sequential-marshal fix). Also `with_thread_names(true)` in the logging setup. The GPU workers are named `nassau-gpu`, but the subscriber printed ids only, so a device-path line read `ThreadId(37)` with no way to tell WHICH device it came from — precisely what is needed to check shard balance. Rayon's pool threads are unnamed and print an empty name. 75/75 tests pass on 4 devices. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 245 +++++++++++-------- ext/src/utils.rs | 5 + 2 files changed, 142 insertions(+), 108 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 50b8785929..5a4f2ecfc4 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -270,8 +270,36 @@ mod gpu_thread { } /// Run `f` on the GPU thread, blocking until it returns. Panics propagate to the caller. - /// Run `f` on `dev`'s worker. The device is chosen by the caller from where the data lives. - pub(super) fn run_on(dev: usize, f: F) -> (T, Timing) + /// 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, @@ -292,18 +320,16 @@ mod gpu_thread { let _ = tx.send((out, queue_ms, exec_ms)); })) .expect("the nassau-gpu thread died"); - let (out, queue_ms, exec_ms) = rx.recv().expect("the nassau-gpu thread died mid-task"); - match out { - Ok(v) => ( - v, - Timing { - queue_ms, - exec_ms, - depth, - }, - ), - Err(payload) => std::panic::resume_unwind(payload), - } + 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() } } @@ -2246,25 +2272,18 @@ fn multiply_batch_grouped( }; by_dev[d].push(prod.clone()); } - // Scoped threads, not a sequential loop: `gpu_thread::run_on` BLOCKS until its device - // finishes, so submitting one after another would serialise the very devices this is - // splitting work across. - let partials: Vec = std::thread::scope(|scope| { - let handles: Vec<_> = by_dev - .iter() - .enumerate() - .filter(|(_, ps)| !ps.is_empty()) - .map(|(d, ps)| { - scope.spawn(move || { - multiply_batch_block(algebra, num_cols, r0, r1 - r0, ps, mode, d) - }) - }) - .collect(); - handles - .into_iter() - .map(|h| h.join().expect("a sharded sub-launch panicked")) - .collect() - }); + // 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. + let waits: Vec<_> = by_dev + .iter() + .enumerate() + .filter(|(_, ps)| !ps.is_empty()) + .map(|(d, ps)| multiply_batch_block(algebra, num_cols, r0, r1 - r0, ps, mode, d)) + .collect(); + let partials: Vec = waits.into_iter().map(|w| w()).collect(); let mut it = partials.into_iter(); let mut acc = it .next() @@ -2294,7 +2313,7 @@ fn multiply_batch_block( products: &[GpuProduct], mode: MasterMode, dev: usize, -) -> Bytes { +) -> Box Bytes> { let (width, g) = algebra.seqno_table_u32(); let mut xi: Vec = xi_degrees(algebra.prime()) .iter() @@ -2642,7 +2661,10 @@ fn multiply_batch_block( ); } if total_pairs == 0 { - return Bytes::from_elems(vec![0u32; num_rows * num_limbs]); + // 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 Box Bytes>; } // The resident `col_sums`/`masks` and basis are non-empty once any `R`/term is present @@ -2678,8 +2700,10 @@ fn multiply_batch_block( pairs = total_pairs, out = out_len ); - let (result, timing) = submit_span.in_scope(|| { - gpu_thread::run_on(dev, move || { + // 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(); @@ -3072,79 +3096,84 @@ fn multiply_batch_block( }) }); - // 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 total = (prep_s + wait_s + device_s).max(1e-9); - 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}", - 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, + Box::new(move || { + let (result, timing) = pending.wait(); + + // 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 total = (prep_s + wait_s + device_s).max(1e-9); + 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}", + 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, + ); + } - result + result + }) } /// How often [`multiply_batch_block`] prints the cumulative marshal/device split, in launches. diff --git a/ext/src/utils.rs b/ext/src/utils.rs index 17e40ed161..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() } From a9fc7294e736dd027cae25d0b7b221a9ec4bbcb0 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Tue, 4 Aug 2026 13:30:30 -0400 Subject: [PATCH 073/127] milnor_gpu: marshal the shards in parallel again, wait outside Replacing `std::thread::scope` with submit-then-wait removed the per-block thread churn, but it also serialised the MARSHAL: each device's share was marshalled one after another on the calling thread before any of them was submitted. Marshalling is ~4.7 ks over a stem-200 run, and the regression showed up immediately -- at matched elapsed time the run was ~18% behind its predecessor (max_t 220 vs 268, 21 450 bidegrees closed vs 23 807). Marshal + submit now run under `into_maybe_par_iter`, and the WAIT stays outside the parallel section. That distinction is the whole point: submitting is non-blocking, so a rayon worker never parks on the device inside a `par_iter` -- the join + steal-loop pattern behind the 146 s signature stalls. Only the shard marshalling, which is pure host work, runs in parallel. Results are order-independent (the partials are XORed), but the waits are sorted by device before collecting so the combine order is deterministic. 75/75 tests pass on 4 devices. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 33 ++++++++++++++++---- 1 file changed, 27 insertions(+), 6 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 5a4f2ecfc4..3c3fcb8c17 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -2193,6 +2193,8 @@ fn multiply_batch_grouped( products: &[GpuProduct], mode: MasterMode, ) -> Vec { + use maybe_rayon::prelude::*; + 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 @@ -2277,13 +2279,28 @@ fn multiply_batch_grouped( // 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. - let waits: Vec<_> = by_dev - .iter() + // Marshal + submit every device's share IN PARALLEL, then wait. Two things matter here and + // they pull in opposite directions: + // - the marshal is real host work (~4.7 ks over a stem-200 run), so doing the devices' + // shares one after another on the calling thread serialises it — measured ~18% behind + // at matched elapsed time when this was sequential; + // - the WAIT must stay outside the parallel section. `par_iter` bodies that block on the + // GPU are the join + steal-loop pattern behind the 146 s signature stalls. + // Submitting is non-blocking, so marshalling in parallel and blocking afterwards gets the + // overlap without ever parking a rayon worker on the device. + let mut waits: Vec<(usize, Box Bytes + Send>)> = by_dev + .into_maybe_par_iter() .enumerate() .filter(|(_, ps)| !ps.is_empty()) - .map(|(d, ps)| multiply_batch_block(algebra, num_cols, r0, r1 - r0, ps, mode, d)) + .map(|(d, ps)| { + ( + d, + multiply_batch_block(algebra, num_cols, r0, r1 - r0, &ps, mode, d), + ) + }) .collect(); - let partials: Vec = waits.into_iter().map(|w| w()).collect(); + waits.sort_by_key(|(d, _)| *d); + let partials: Vec = waits.into_iter().map(|(_, w)| w()).collect(); let mut it = partials.into_iter(); let mut acc = it .next() @@ -2313,7 +2330,7 @@ fn multiply_batch_block( products: &[GpuProduct], mode: MasterMode, dev: usize, -) -> Box Bytes> { +) -> Box Bytes + Send> { let (width, g) = algebra.seqno_table_u32(); let mut xi: Vec = xi_degrees(algebra.prime()) .iter() @@ -2664,7 +2681,7 @@ fn multiply_batch_block( // 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 Box Bytes>; + return Box::new(move || empty) as Box Bytes + Send>; } // The resident `col_sums`/`masks` and basis are non-empty once any `R`/term is present @@ -2694,8 +2711,12 @@ fn multiply_batch_block( // // 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::info_span!( "gpu_submit", + dev = dev, rows = num_rows, pairs = total_pairs, out = out_len From 0bc89566c6f364c3b4ffa7124b56134cb195f687 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Tue, 4 Aug 2026 14:03:12 -0400 Subject: [PATCH 074/127] milnor_gpu: keep the scoped-thread fan-out; it beat both rewrites Reverts the fan-out to `std::thread::scope`, which measurement says is the fastest of the three shapes tried on full stem-200 runs: scoped threads (marshal + wait per device, on its own thread) BEST sequential marshal, then submit all and wait ~18% behind marshal + submit under par_iter, wait outside ~10 max_t behind The second lost because the marshal is ~4.7 ks of host work per run and doing the devices' shares one after another serialises it. The third kept the marshal parallel and still lost: a `par_iter` join per row block costs more here than a thread does. Priority inversion was the suspected cause of the third's loss and the data says otherwise -- steps >= 20 s over a matched 1040 s window totalled 3665 s (par_iter), 3779 s (threads) and 3619 s (the single-GPU run). Sharding raises the COUNT of long steps somewhat (65 -> ~90) but not the total time, and the whole-bidegree `ParallelGuard` already stands between these jobs and the steal loop. The cost of this shape is real and stays documented: up to `gpu_count()` OS threads per row block, visible as thread ids in the hundreds of thousands in a long log. `gpu_thread::submit_on` and the returned-wait signature are kept -- they cost nothing here and are what a future non-thread fan-out would need. 75/75 tests pass on 4 devices. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 54 ++++++++++++-------- 1 file changed, 32 insertions(+), 22 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 3c3fcb8c17..d0d95b772b 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -2279,28 +2279,38 @@ fn multiply_batch_grouped( // 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. - // Marshal + submit every device's share IN PARALLEL, then wait. Two things matter here and - // they pull in opposite directions: - // - the marshal is real host work (~4.7 ks over a stem-200 run), so doing the devices' - // shares one after another on the calling thread serialises it — measured ~18% behind - // at matched elapsed time when this was sequential; - // - the WAIT must stay outside the parallel section. `par_iter` bodies that block on the - // GPU are the join + steal-loop pattern behind the 146 s signature stalls. - // Submitting is non-blocking, so marshalling in parallel and blocking afterwards gets the - // overlap without ever parking a rayon worker on the device. - let mut waits: Vec<(usize, Box Bytes + Send>)> = by_dev - .into_maybe_par_iter() - .enumerate() - .filter(|(_, ps)| !ps.is_empty()) - .map(|(d, ps)| { - ( - d, - multiply_batch_block(algebra, num_cols, r0, r1 - r0, &ps, mode, d), - ) - }) - .collect(); - waits.sort_by_key(|(d, _)| *d); - let partials: Vec = waits.into_iter().map(|(_, w)| w()).collect(); + // Scoped threads, deliberately, after measuring the alternatives. Each device's share is + // marshalled AND awaited on its own thread, so both the host marshalling and the four device + // sections overlap. + // + // Two tidier-looking designs were tried on a full stem-200 and both lost: + // - marshal sequentially, then submit all and wait (no threads at all): the marshal is + // ~4.7 ks of host work over a run, and serialising it across devices ran ~18% behind at + // matched elapsed time. + // - marshal + submit under `into_maybe_par_iter`, waiting outside the parallel section: + // still ~10 points of `max_t` behind at matched elapsed. Stall burden was NOT the cause + // (steps >= 20 s totalled 3.7 ks either way, the same as the single-GPU run) -- a + // `par_iter` join per row block simply costs more here than a thread does. + // + // The cost is real and was worth checking: this spawns up to `gpu_count()` OS threads per + // row block, which shows up as thread ids in the hundreds of thousands in a long run. It is + // still the fastest of the three, so it stays until something beats it on a measured run. + let partials: Vec = std::thread::scope(|scope| { + let handles: Vec<_> = by_dev + .iter() + .enumerate() + .filter(|(_, ps)| !ps.is_empty()) + .map(|(d, ps)| { + scope.spawn(move || { + multiply_batch_block(algebra, num_cols, r0, r1 - r0, ps, mode, d)() + }) + }) + .collect(); + handles + .into_iter() + .map(|h| h.join().expect("a sharded sub-launch panicked")) + .collect() + }); let mut it = partials.into_iter(); let mut acc = it .next() From 07d7aab28ff5fe87af7a9c812a872621e5a6f39b Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Tue, 4 Aug 2026 14:38:24 -0400 Subject: [PATCH 075/127] maybe-rayon: add MaybeThreadPool; use it for the multi-GPU fan-out MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The fan-out spawned up to `gpu_count()` OS threads PER ROW BLOCK via `std::thread::scope` — on the order of a million over a stem-200 run. It matched the fastest configuration measured, but unbounded thread churn is not something to run at scale. A private pool of `gpu_count()` persistent workers replaces it, and is private for a second reason beyond cost: a fan-out worker blocked on a device can only steal ANOTHER SHARD of the same block, never a large `step_resolution` job off the global pool. That is the priority inversion this codebase already paid for once (the 146 s signature stalls), ruled out structurally instead of by a guard. Two paired rounds on 4 GPUs: scoped threads 2.96 / 3.43 e10 pairs/s (16% spread) private pool 3.38 / 3.36 e10 pairs/s (0.6% spread) Throughput-neutral within noise; the arms cross. What does move is variance, which is exactly what dropping a million thread spawns should do. Adopted for bounded resources and stable timing, not for a throughput claim. The pool lives in `maybe-rayon` rather than pulling `rayon` into `algebra`, so the crate keeps its single-threaded build for debugging: `sequential::MaybeThreadPool` holds no threads and runs `install` inline. The fan-out then runs the shards one after another, which stays correct because the partials are XORed and therefore order-independent. Both `--features gpu` and `--no-default-features --features gpu` build. 75/75 tests pass on 4 devices. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 59 +++++++++++--------- ext/crates/maybe-rayon/src/concurrent.rs | 25 +++++++++ ext/crates/maybe-rayon/src/sequential.rs | 15 +++++ 3 files changed, 74 insertions(+), 25 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index d0d95b772b..0e8bcd6d44 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -640,6 +640,24 @@ fn cur_device() -> usize { } /// The cubecl client for this thread's device. +/// Private thread pool for the multi-GPU fan-out: `gpu_count()` persistent workers, created once. +/// +/// Private on purpose, for two independent reasons. +/// - Isolation: a fan-out worker blocked on a device can only ever steal ANOTHER SHARD of the +/// same block — never a giant `step_resolution` job off the global pool. That is the priority +/// inversion this codebase has already paid for once (the 146 s signature stalls), and a +/// private pool rules it out structurally rather than relying on a guard. +/// - Cost: it replaces `std::thread::scope`, which spawned up to `gpu_count()` OS threads PER ROW +/// BLOCK — on the order of a million over a stem-200 run. +/// +/// Via `maybe-rayon`, so a build without `concurrent` gets the sequential proxy and the shards run +/// one after another. That stays correct because the partials are XORed and so order-independent. +fn fanout_pool() -> &'static maybe_rayon::MaybeThreadPool { + static POOL: LazyLock = + LazyLock::new(|| maybe_rayon::MaybeThreadPool::new(gpu_count(), "nassau-fanout")); + &POOL +} + fn gpu_client() -> cubecl::prelude::ComputeClient { CudaRuntime::client(&CudaDevice::new(cur_device())) } @@ -2279,37 +2297,28 @@ fn multiply_batch_grouped( // 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. - // Scoped threads, deliberately, after measuring the alternatives. Each device's share is - // marshalled AND awaited on its own thread, so both the host marshalling and the four device - // sections overlap. - // - // Two tidier-looking designs were tried on a full stem-200 and both lost: - // - marshal sequentially, then submit all and wait (no threads at all): the marshal is - // ~4.7 ks of host work over a run, and serialising it across devices ran ~18% behind at - // matched elapsed time. - // - marshal + submit under `into_maybe_par_iter`, waiting outside the parallel section: - // still ~10 points of `max_t` behind at matched elapsed. Stall burden was NOT the cause - // (steps >= 20 s totalled 3.7 ks either way, the same as the single-GPU run) -- a - // `par_iter` join per row block simply costs more here than a thread does. + // Fan the shards out on a PRIVATE pool (see `fanout_pool`): persistent workers, and no + // possibility of a blocked fan-out worker stealing a `step_resolution` job. // - // The cost is real and was worth checking: this spawns up to `gpu_count()` OS threads per - // row block, which shows up as thread ids in the hundreds of thousands in a long run. It is - // still the fastest of the three, so it stays until something beats it on a measured run. - let partials: Vec = std::thread::scope(|scope| { - let handles: Vec<_> = by_dev - .iter() + // Each task marshals AND waits for its own device, so host marshalling (~4.7 ks per run) and + // the four device sections all overlap. Two tidier shapes were measured on full stem-200 + // runs and both lost: marshalling sequentially then submitting all ran ~18% behind, and + // marshal+submit on the GLOBAL pool with the wait outside ran ~10 `max_t` behind. + let partials: Vec = fanout_pool().install(move || { + let mut out: Vec<(usize, Bytes)> = by_dev + .into_maybe_par_iter() .enumerate() .filter(|(_, ps)| !ps.is_empty()) .map(|(d, ps)| { - scope.spawn(move || { - multiply_batch_block(algebra, num_cols, r0, r1 - r0, ps, mode, d)() - }) + ( + d, + multiply_batch_block(algebra, num_cols, r0, r1 - r0, &ps, mode, d)(), + ) }) .collect(); - handles - .into_iter() - .map(|h| h.join().expect("a sharded sub-launch panicked")) - .collect() + // Order-independent (the partials are XORed), but sorted so the combine is deterministic. + out.sort_by_key(|(d, _)| *d); + out.into_iter().map(|(_, b)| b).collect() }); let mut it = partials.into_iter(); let mut acc = it diff --git a/ext/crates/maybe-rayon/src/concurrent.rs b/ext/crates/maybe-rayon/src/concurrent.rs index 9426cf7f84..0688dd2886 100644 --- a/ext/crates/maybe-rayon/src/concurrent.rs +++ b/ext/crates/maybe-rayon/src/concurrent.rs @@ -86,3 +86,28 @@ 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) + } +} diff --git a/ext/crates/maybe-rayon/src/sequential.rs b/ext/crates/maybe-rayon/src/sequential.rs index a2493a7e43..364bdc5571 100644 --- a/ext/crates/maybe-rayon/src/sequential.rs +++ b/ext/crates/maybe-rayon/src/sequential.rs @@ -129,3 +129,18 @@ 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() + } +} From b8c4470ade48027f35fe730ecedcaca020366dc7 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Tue, 4 Aug 2026 15:11:33 -0400 Subject: [PATCH 076/127] milnor_gpu: size the fan-out pool for concurrent callers, not devices `gpu_count()` workers looked like the natural size and starved the devices. Every resolution worker fans a row block into `gpu_count()` shards, so with ~32 workers there are ~32 fan-outs live at once; a 4-worker pool capped in-flight submissions at 4 and left the per-device queues shallow (they run depth 8-24 when healthy), so devices idled between jobs. A full stem-200 ran ~250 s behind its predecessor at matched elapsed time before this was raised. Default is now `gpu_count() * 16` clamped to [16, 128], with `NASSAU_GPU_FANOUT_THREADS` to override. Oversubscription is cheap here because these workers are almost always BLOCKED on a device rather than running. Worth recording how this was missed: the stem-200 kernel bench drives only `NASSAU_BENCH_WORKERS` (7) submitters, so it showed pool and scoped threads as equivalent. The starvation only appears at the concurrency of a real resolution. A bench that fixes its own concurrency cannot see a contention bug that depends on it. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 0e8bcd6d44..5ecb895755 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -652,9 +652,23 @@ fn cur_device() -> usize { /// /// Via `maybe-rayon`, so a build without `concurrent` gets the sequential proxy and the shards run /// one after another. That stays correct because the partials are XORed and so order-independent. +/// Sized for CONCURRENT CALLERS x devices, not devices. Sizing it at `gpu_count()` looks natural +/// and is badly wrong: every resolution worker fans a row block out into `gpu_count()` shards, so +/// with ~32 workers there are ~32 fan-outs live at once. A 4-worker pool caps in-flight submissions +/// at 4 and starves the per-device queues (measured depth 8-24 when healthy), leaving devices idle +/// between jobs — a full stem-200 ran ~250 s behind before this was raised. +/// +/// These workers spend nearly all their time BLOCKED waiting on a device, so oversubscribing costs +/// almost nothing. `NASSAU_GPU_FANOUT_THREADS` overrides. fn fanout_pool() -> &'static maybe_rayon::MaybeThreadPool { - static POOL: LazyLock = - LazyLock::new(|| maybe_rayon::MaybeThreadPool::new(gpu_count(), "nassau-fanout")); + static POOL: LazyLock = LazyLock::new(|| { + let n = std::env::var("NASSAU_GPU_FANOUT_THREADS") + .ok() + .and_then(|v| v.parse::().ok()) + .filter(|&n| n > 0) + .unwrap_or_else(|| (gpu_count() * 16).clamp(16, 128)); + maybe_rayon::MaybeThreadPool::new(n, "nassau-fanout") + }); &POOL } From 8c2365cf41c71087cc3bc8ea51eb494e8e73bffb Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Tue, 4 Aug 2026 16:33:39 -0400 Subject: [PATCH 077/127] milnor_gpu: size the fan-out pool by callers x devices MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Diagnosis of the 47% end-to-end regression the private pool caused (3743 s vs 2551 s for scoped threads, same 898 000 calls): metric threads pool GPU exec 6074 s 5933 s identical work marshal 4677 s 1558 s pool 3x FASTER queue 20598 s 15879 s pool waits less depth mean 3.8 2.2 <-- the cause device duty 60% 40% The pool improved everything except the one thing that mattered. `install` holds a worker for a whole block (marshal AND the device wait), and a block needs one slot per shard, so the pool caps blocks in flight at `threads / gpu_count`. At 64 threads with 32 rayon workers only 16 of 32 possible blocks could be live; queue depth fell and the devices idled 60% of the time. Faster marshalling and shorter queues are worth nothing against starved GPUs. The default was `gpu_count() * 16` — a multiple of the wrong quantity. The pool serves CALLERS, so it is now `max_num_threads() * gpu_count()` (128 here), clamped to [16, 1024], still overridable by `NASSAU_GPU_FANOUT_THREADS`. Oversubscription is cheap because these workers are almost always blocked on a device. `maybe_rayon::max_num_threads()` is added to both modules rather than reaching for `rayon::current_num_threads()` directly, so the sequential build keeps compiling and reports one caller. Not yet confirmed end to end: this needs a full stem-200 to verify it recovers 2551 s. The kernel bench cannot see it — it drives a fixed 7 submitters, which is why it called pool and threads a tie in the first place. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 11 ++++++++++- ext/crates/maybe-rayon/src/concurrent.rs | 6 ++++++ ext/crates/maybe-rayon/src/sequential.rs | 5 +++++ 3 files changed, 21 insertions(+), 1 deletion(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 5ecb895755..6599a17ec5 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -666,7 +666,16 @@ fn fanout_pool() -> &'static maybe_rayon::MaybeThreadPool { .ok() .and_then(|v| v.parse::().ok()) .filter(|&n| n > 0) - .unwrap_or_else(|| (gpu_count() * 16).clamp(16, 128)); + .unwrap_or_else(|| { + // `callers x gpu_count`, NOT a multiple of `gpu_count` alone. `install` holds a + // worker for a whole block (marshal AND device wait), and a block needs one slot per + // shard, so the pool caps blocks in flight at `threads / gpu_count`. Sizing it below + // `resolution workers x gpu_count` throttles submission and starves the devices: + // at 64 threads with 32 rayon workers only 16 blocks could be live, queue depth fell + // 3.8 -> 2.2, device duty 60% -> 40%, and a full stem-200 took 3743 s against 2551 s. + // Marshal got 3x FASTER and queueing dropped; none of that mattered against idle GPUs. + (maybe_rayon::max_num_threads() * gpu_count()).clamp(16, 1024) + }); maybe_rayon::MaybeThreadPool::new(n, "nassau-fanout") }); &POOL diff --git a/ext/crates/maybe-rayon/src/concurrent.rs b/ext/crates/maybe-rayon/src/concurrent.rs index 0688dd2886..bbc6bb4d8f 100644 --- a/ext/crates/maybe-rayon/src/concurrent.rs +++ b/ext/crates/maybe-rayon/src/concurrent.rs @@ -111,3 +111,9 @@ impl MaybeThreadPool { 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 364bdc5571..767c912065 100644 --- a/ext/crates/maybe-rayon/src/sequential.rs +++ b/ext/crates/maybe-rayon/src/sequential.rs @@ -144,3 +144,8 @@ impl MaybeThreadPool { f() } } + +/// Sequential proxy: one caller. +pub fn max_num_threads() -> usize { + 1 +} From 199c052ad9091d75baf02794d36c246b8b2f752f Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Tue, 4 Aug 2026 17:27:43 -0400 Subject: [PATCH 078/127] milnor_gpu: one thread, async submissions to every device MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Drops the fan-out pool (and, before it, the per-block thread spawning) entirely. `submit_on` does not block, so a single caller marshals a shard, hands it to its device, and moves to the next while that device is already running. Only the final wait blocks. Two beliefs that drove the earlier designs were wrong: - "the shards must marshal in parallel". Sharding SPLITS the products, so marshalling them one after another on one thread is the same total host work as marshalling one unsharded block. There was nothing to parallelise. - "one thread cannot keep four devices fed". It does not have to: ~32 resolution workers each submit `gpu_count()` jobs, so ~128 launches are in flight. That is the queue depth the devices need, and any fan-out pool can only CAP it — sized at 64 it halved depth (3.8 -> 2.2), dropped duty to 40% and cost 47% end to end. I had measured a version of this shape (run 2) at ~18% behind and blamed serialised marshalling, on a comparison taken 17 minutes into a killed run. That attribution was not supported: the gap had already narrowed from 48 to 12 points of `max_t` by the time it was stopped, and marshal total is unchanged by how shards are split. 75/75 tests pass on 4 devices. Needs an end-to-end stem-200 against the 2551 s scoped-thread result to confirm. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 68 +++++--------------- 1 file changed, 15 insertions(+), 53 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 6599a17ec5..32318d05a7 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -652,35 +652,6 @@ fn cur_device() -> usize { /// /// Via `maybe-rayon`, so a build without `concurrent` gets the sequential proxy and the shards run /// one after another. That stays correct because the partials are XORed and so order-independent. -/// Sized for CONCURRENT CALLERS x devices, not devices. Sizing it at `gpu_count()` looks natural -/// and is badly wrong: every resolution worker fans a row block out into `gpu_count()` shards, so -/// with ~32 workers there are ~32 fan-outs live at once. A 4-worker pool caps in-flight submissions -/// at 4 and starves the per-device queues (measured depth 8-24 when healthy), leaving devices idle -/// between jobs — a full stem-200 ran ~250 s behind before this was raised. -/// -/// These workers spend nearly all their time BLOCKED waiting on a device, so oversubscribing costs -/// almost nothing. `NASSAU_GPU_FANOUT_THREADS` overrides. -fn fanout_pool() -> &'static maybe_rayon::MaybeThreadPool { - static POOL: LazyLock = LazyLock::new(|| { - let n = std::env::var("NASSAU_GPU_FANOUT_THREADS") - .ok() - .and_then(|v| v.parse::().ok()) - .filter(|&n| n > 0) - .unwrap_or_else(|| { - // `callers x gpu_count`, NOT a multiple of `gpu_count` alone. `install` holds a - // worker for a whole block (marshal AND device wait), and a block needs one slot per - // shard, so the pool caps blocks in flight at `threads / gpu_count`. Sizing it below - // `resolution workers x gpu_count` throttles submission and starves the devices: - // at 64 threads with 32 rayon workers only 16 blocks could be live, queue depth fell - // 3.8 -> 2.2, device duty 60% -> 40%, and a full stem-200 took 3743 s against 2551 s. - // Marshal got 3x FASTER and queueing dropped; none of that mattered against idle GPUs. - (maybe_rayon::max_num_threads() * gpu_count()).clamp(16, 1024) - }); - maybe_rayon::MaybeThreadPool::new(n, "nassau-fanout") - }); - &POOL -} - fn gpu_client() -> cubecl::prelude::ComputeClient { CudaRuntime::client(&CudaDevice::new(cur_device())) } @@ -2234,8 +2205,6 @@ fn multiply_batch_grouped( products: &[GpuProduct], mode: MasterMode, ) -> Vec { - use maybe_rayon::prelude::*; - 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 @@ -2320,29 +2289,22 @@ fn multiply_batch_grouped( // 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. - // Fan the shards out on a PRIVATE pool (see `fanout_pool`): persistent workers, and no - // possibility of a blocked fan-out worker stealing a `step_resolution` job. + // 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. // - // Each task marshals AND waits for its own device, so host marshalling (~4.7 ks per run) and - // the four device sections all overlap. Two tidier shapes were measured on full stem-200 - // runs and both lost: marshalling sequentially then submitting all ran ~18% behind, and - // marshal+submit on the GLOBAL pool with the wait outside ran ~10 `max_t` behind. - let partials: Vec = fanout_pool().install(move || { - let mut out: Vec<(usize, Bytes)> = by_dev - .into_maybe_par_iter() - .enumerate() - .filter(|(_, ps)| !ps.is_empty()) - .map(|(d, ps)| { - ( - d, - multiply_batch_block(algebra, num_cols, r0, r1 - r0, &ps, mode, d)(), - ) - }) - .collect(); - // Order-independent (the partials are XORed), but sorted so the combine is deterministic. - out.sort_by_key(|(d, _)| *d); - out.into_iter().map(|(_, b)| b).collect() - }); + // 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). + let waits: Vec Bytes + Send>> = by_dev + .iter() + .enumerate() + .filter(|(_, ps)| !ps.is_empty()) + .map(|(d, ps)| multiply_batch_block(algebra, num_cols, r0, r1 - r0, ps, mode, d)) + .collect(); + let partials: Vec = waits.into_iter().map(|w| w()).collect(); let mut it = partials.into_iter(); let mut acc = it .next() From a4aed87c6439ddfe16cb4fc837f74be4589a0727 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Tue, 4 Aug 2026 18:47:21 -0400 Subject: [PATCH 079/127] milnor_gpu: pipeline the readback; restore the memory bound MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Each device ran exactly ONE kernel at a time, no matter how many callers were queued. The device task ended in `client.read_one(out_h)`, and the worker loop is `while let Ok(task) = rx.recv() { task() }`, so the worker was pinned inside the readback until the kernel retired. `DEPTH` counted jobs WAITING, not jobs in flight — which is why three successive fan-out rewrites (scoped threads 2551s, pool 3743s, single-thread async 3140s) all landed in the same place: they varied what fed a queue whose service rate was one kernel deep. `read_one` is `block_on(read_async(..))`, and `read_async` already does the right thing: it enqueues the copy, records a CUDA event, and returns a future whose entire body is that event's wait. So the fix is to stop consuming the future at the point of launch and consume it at the point of use. No executor is involved — the future never yields `Pending` (it wraps a blocking `cuEventSynchronize`), so `block_on` polls it exactly once. A runtime like tokio would be strictly worse: blocking inside a poll stalls a worker's whole run queue, and `spawn_blocking` is just this thread pool again with a scheduler on top. `gpu_client` now returns `&'static` so the future (edition 2024 RPIT captures `&self`) can outlive the worker task. Separately, and load-bearing here: the byte budget has bounded nothing since 1639982877. That commit made this function return its wait instead of performing it, but left `_permit` a plain local — dropped at function return, i.e. straight after the non-blocking `submit_on`. Every fan-out experiment since ran with no memory backpressure. The permit now moves into the wait closure and drops after the readback, which is when the output buffer (device page + pinned host landing) is actually free. Marshal and submit are split (`BlockSubmit` -> `BlockWait`) so no permit is ever held across a rayon section — GpuBudget's safety invariant. The fused shape held device d's permit while marshalling device d+1, so a par_iter chunk could steal a step job that parked on `acquire` while this thread waited on a join those workers had to finish. batch-stats gains launch/fence, splitting the caller's wait at the point the worker hands back the future: fence >> launch means the pipeline is full, launch >> fence means it is starved. Conflated, the two are indistinguishable — which is how one-kernel-deep stayed hidden this long. 75/75 algebra tests pass on 4 devices. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 1507 ++++++++++-------- 1 file changed, 810 insertions(+), 697 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 32318d05a7..9b9f3d9021 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -245,7 +245,7 @@ mod gpu_thread { /// 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>> = OnceLock::new(); + static QUEUES: OnceLock>> = std::sync::OnceLock::new(); QUEUES.get_or_init(|| { let mut txs = Vec::with_capacity(super::gpu_count()); for dev in 0..super::gpu_count() { @@ -408,7 +408,7 @@ static CLEANUP_COUNTER: AtomicU64 = AtomicU64::new(0); /// device-memory growth, since freed pages linger — watch `nvidia-smi`. fn cleanup_every() -> u64 { use std::sync::OnceLock; - static EVERY: OnceLock = OnceLock::new(); + static EVERY: OnceLock = std::sync::OnceLock::new(); *EVERY.get_or_init(|| { std::env::var("NASSAU_GPU_CLEANUP_EVERY") .ok() @@ -443,6 +443,13 @@ 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); /// Read and reset the aggregate batch counters: `(calls, marshal_us, device_us, pairs)`. pub fn take_batch_stats() -> (u64, u64, u64, u64) { @@ -639,21 +646,22 @@ fn cur_device() -> usize { CUR_DEVICE.with(|c| c.get()) } -/// The cubecl client for this thread's device. -/// Private thread pool for the multi-GPU fan-out: `gpu_count()` persistent workers, created once. +/// The cubecl client for this thread's device, borrowed for the process's lifetime. /// -/// Private on purpose, for two independent reasons. -/// - Isolation: a fan-out worker blocked on a device can only ever steal ANOTHER SHARD of the -/// same block — never a giant `step_resolution` job off the global pool. That is the priority -/// inversion this codebase has already paid for once (the 146 s signature stalls), and a -/// private pool rules it out structurally rather than relying on a guard. -/// - Cost: it replaces `std::thread::scope`, which spawned up to `gpu_count()` OS threads PER ROW -/// BLOCK — on the order of a million over a stem-200 run. -/// -/// Via `maybe-rayon`, so a build without `concurrent` gets the sequential proxy and the shards run -/// one after another. That stays correct because the partials are XORed and so order-independent. -fn gpu_client() -> cubecl::prelude::ComputeClient { - CudaRuntime::client(&CudaDevice::new(cur_device())) +/// `'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()] } /// Compile-time cap on the number of fixed-size segments a resident device buffer may hold. It @@ -2298,12 +2306,26 @@ fn multiply_batch_grouped( // 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). - let waits: Vec Bytes + Send>> = by_dev + // 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. + let submits: Vec> = by_dev .iter() .enumerate() .filter(|(_, ps)| !ps.is_empty()) .map(|(d, ps)| multiply_batch_block(algebra, num_cols, r0, r1 - r0, ps, mode, d)) .collect(); + let waits: Vec = submits.into_iter().map(|s| s()).collect(); let partials: Vec = waits.into_iter().map(|w| w()).collect(); let mut it = partials.into_iter(); let mut acc = it @@ -2321,20 +2343,34 @@ fn multiply_batch_grouped( 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 + '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. Holds a -/// [`GpuPermit`] for its sequential layout + device section (acquired only after the parallel -/// marshal — see [`GpuBudget`]), so the total output size of concurrent device sections +/// 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( - algebra: &MilnorAlgebra, +fn multiply_batch_block<'a>( + algebra: &'a MilnorAlgebra, num_cols: usize, row_base: usize, num_rows: usize, - products: &[GpuProduct], + products: &'a [GpuProduct], mode: MasterMode, dev: usize, -) -> Box Bytes + Send> { +) -> BlockSubmit<'a> { let (width, g) = algebra.seqno_table_u32(); let mut xi: Vec = xi_degrees(algebra.prime()) .iter() @@ -2487,717 +2523,794 @@ fn multiply_batch_block( // 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. The - // device section now runs on the dedicated GPU thread (see [`gpu_thread`]), not from the permit. + // 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; - 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); - r_cs_offset.push(need_cs as u64); - r_mk_offset.push(need_mk as u64); - 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::>()); - need_cs += num_mats as usize * cs_len as usize; - need_mk += num_mats as usize * mk_len as usize; + + // 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); + r_cs_offset.push(need_cs as u64); + r_mk_offset.push(need_mk as u64); + 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::>()); + need_cs += num_mats as usize * cs_len as usize; + need_mk += num_mats as usize * mk_len as usize; + } } } - } - // (Transient) Flatten the cold p-parts (padded to the widest) for the enumeration kernel. The - // per-`R` scratch offsets it writes at are `r_cs_offset`/`r_mk_offset` themselves (u64), passed - // straight through — no u32 narrowing, so a big block's multi-GB scratch is addressed safely. - let (enum_pp, enum_width) = if mode == MasterMode::Transient { - let w = enum_rows.iter().copied().max().unwrap_or(1) as usize; - let mut pp = vec![0u32; enum_pp_rows.len() * w]; - for (i, row) in enum_pp_rows.iter().enumerate() { - for (slot, &v) in pp[i * w..i * w + row.len()].iter_mut().zip(row) { - *slot = v; + // (Transient) Flatten the cold p-parts (padded to the widest) for the enumeration kernel. The + // per-`R` scratch offsets it writes at are `r_cs_offset`/`r_mk_offset` themselves (u64), passed + // straight through — no u32 narrowing, so a big block's multi-GB scratch is addressed safely. + let (enum_pp, enum_width) = if mode == MasterMode::Transient { + let w = enum_rows.iter().copied().max().unwrap_or(1) as usize; + let mut pp = vec![0u32; enum_pp_rows.len() * w]; + for (i, row) in enum_pp_rows.iter().enumerate() { + for (slot, &v) in pp[i * w..i * w + row.len()].iter_mut().zip(row) { + *slot = v; + } } - } - (pp, w) - } else { - (Vec::new(), 1usize) - }; + (pp, w) + } else { + (Vec::new(), 1usize) + }; - // 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); - } + // 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" - ); - pps.push(total_pairs as u32); - - // 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; + 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" + ); + pps.push(total_pairs as u32); + + // 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; } - 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); } - // 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), + // 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 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 Box Bytes + Send>; - } + 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); - } + // 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 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(); + 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(); + // `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::info_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() - }; + // 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::info_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() + }; - // `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. - let (cs_seg, mk_seg) = match mode { - MasterMode::Resident => { - let (cs_segs, _) = seg_grow!( + // `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. + 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). + const ENUM_THREADS: u32 = 256; + let n_cold = enum_rows.len(); + let cs_cap = need_cs.max(1); + let mk_cap = need_mk.max(1); + assert!( + cs_cap <= seg_elems && mk_cap <= seg_elems, + "transient scratch ({cs_cap}/{mk_cap} u16) exceeds one segment \ + ({seg_elems}); raise NASSAU_GPU_MASTER_SEG_ELEMS" + ); + let cs_scratch = client.empty(cs_cap * size_of::()); + let mk_scratch = client.empty(mk_cap * size_of::()); + let cnt_scratch = client.empty(n_cold.max(1) * size_of::()); + let epp_h = client.create_from_slice(u32::as_bytes(&enum_pp)); + let er_h = client.create_from_slice(u32::as_bytes(&enum_rows)); + let ec_h = client.create_from_slice(u32::as_bytes(&enum_cols)); + let eco_h = client.create_from_slice(u64::as_bytes(&r_cs_offset)); + let emo_h = client.create_from_slice(u64::as_bytes(&r_mk_offset)); + unsafe { + enumerate_admissible_kernel::launch_unchecked::( + &client, + CubeCount::Static( + (n_cold as u32).div_ceil(ENUM_THREADS).max(1), + 1, + 1, + ), + CubeDim::new_1d(ENUM_THREADS), + BufferArg::from_raw_parts(epp_h, enum_pp.len()), + BufferArg::from_raw_parts(er_h, n_cold), + BufferArg::from_raw_parts(ec_h, n_cold), + BufferArg::from_raw_parts(eco_h, n_cold), + BufferArg::from_raw_parts(emo_h, n_cold), + BufferArg::from_raw_parts(cs_scratch.clone(), cs_cap), + BufferArg::from_raw_parts(mk_scratch.clone(), mk_cap), + BufferArg::from_raw_parts(cnt_scratch, n_cold.max(1)), + enum_width, + n_cold, + ); + } + ( + pad_u16(vec![(cs_scratch, cs_cap)]), + pad_u16(vec![(mk_scratch, mk_cap)]), + ) + } + }; + // 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_dev(), - cs, - resident_upload(), - need_cs, + resident_basis_dev(), + pp, + resident_basis_upload(), + need_basis_elems * width, 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) + |up: usize| { + let h = RESIDENT_BASIS_HOST.read().unwrap(); + let nl = h.lens.len() * h.width; + (h.pparts[up..nl].to_vec(), nl) } ); - let (mk_segs, _) = seg_grow!( + let (ln_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) + 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(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). - const ENUM_THREADS: u32 = 256; - let n_cold = enum_rows.len(); - let cs_cap = need_cs.max(1); - let mk_cap = need_mk.max(1); - assert!( - cs_cap <= seg_elems && mk_cap <= seg_elems, - "transient scratch ({cs_cap}/{mk_cap} u16) exceeds one segment \ - ({seg_elems}); raise NASSAU_GPU_MASTER_SEG_ELEMS" + (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 cs_scratch = client.empty(cs_cap * size_of::()); - let mk_scratch = client.empty(mk_cap * size_of::()); - let cnt_scratch = client.empty(n_cold.max(1) * size_of::()); - let epp_h = client.create_from_slice(u32::as_bytes(&enum_pp)); - let er_h = client.create_from_slice(u32::as_bytes(&enum_rows)); - let ec_h = client.create_from_slice(u32::as_bytes(&enum_cols)); - let eco_h = client.create_from_slice(u64::as_bytes(&r_cs_offset)); - let emo_h = client.create_from_slice(u64::as_bytes(&r_mk_offset)); - unsafe { - enumerate_admissible_kernel::launch_unchecked::( - &client, - CubeCount::Static((n_cold as u32).div_ceil(ENUM_THREADS).max(1), 1, 1), - CubeDim::new_1d(ENUM_THREADS), - BufferArg::from_raw_parts(epp_h, enum_pp.len()), - BufferArg::from_raw_parts(er_h, n_cold), - BufferArg::from_raw_parts(ec_h, n_cold), - BufferArg::from_raw_parts(eco_h, n_cold), - BufferArg::from_raw_parts(emo_h, n_cold), - BufferArg::from_raw_parts(cs_scratch.clone(), cs_cap), - BufferArg::from_raw_parts(mk_scratch.clone(), mk_cap), - BufferArg::from_raw_parts(cnt_scratch, n_cold.max(1)), - enum_width, - n_cold, - ); - } - ( - pad_u16(vec![(cs_scratch, cs_cap)]), - pad_u16(vec![(mk_scratch, mk_cap)]), - ) } - }; - // 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 { + + 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!( - 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), + 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) + }; + } + // 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), + ); + } - 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)" + // 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(); + } + + result + }) + }); + + 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, 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). + 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, + ); + + // 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, ); - 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(); + 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 total = (prep_s + wait_s + device_s).max(1e-9); 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) - }; - } - // 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), + "[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}%", + 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), ); } - // Hand back the landing buffer itself — no copy, no allocation (see [`BatchOutput`]). - let result = client.read_one(out_h).unwrap(); - - // 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. - let every = cleanup_every(); - if every != 0 && CLEANUP_COUNTER.fetch_add(1, Ordering::Relaxed) % every == 0 { - client.memory_cleanup(); - } - result - }) - }); - - Box::new(move || { - let (result, timing) = pending.wait(); - - // 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 total = (prep_s + wait_s + device_s).max(1e-9); - 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}", - 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, - ); - } - - result + }) as BlockWait }) } @@ -3215,13 +3328,13 @@ fn multiply_batch_block( /// 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 = OnceLock::new(); + static ON: OnceLock = std::sync::OnceLock::new(); *ON.get_or_init(|| std::env::var_os("NASSAU_GPU_LAUNCH_LOG").is_some()) } fn batch_report_every() -> u64 { use std::sync::OnceLock; - static EVERY: OnceLock = OnceLock::new(); + static EVERY: OnceLock = std::sync::OnceLock::new(); *EVERY.get_or_init(|| { std::env::var("NASSAU_BATCH_REPORT_EVERY") .ok() From ef16f934ef9e3bf5e0ade8cb66bbe45f6827947c Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Tue, 4 Aug 2026 19:51:18 -0400 Subject: [PATCH 080/127] milnor_gpu: marshal the shards in parallel again MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The readback pipeline landed but did not pay: stem-200 wall was 3136s against 3140s for the single-thread async run it replaced, while scoped threads still holds the record at 2551s. Final batch-stats say why the pipelining alone could not win — queue=61%, exec=14%, pipeline=19%. The readback genuinely stopped serialising the worker (fence is 1321s against 5522s of launch wait), but that was never where the time was. What separates the 2551s arm from the two ~3140s arms is not fan-out depth, it is that scoped threads marshalled the shards CONCURRENTLY. A comment here previously reasoned that sharding splits the products, so marshalling shards one after another is the same total host work as marshalling one unsharded block, and concluded there was nothing to parallelise. That holds for total work and fails for the critical path: a row block is not finished until its slowest shard is, so serial marshalling makes each block's latency the SUM over devices instead of the max. `marshal` totals 3444s. Phase 1 is now parallel, which is safe by construction rather than by argument: the marshal/submit split means no permit is held until phase 2, so this cannot recreate the permit-across-rayon deadlock that the scoped version was always one refactor away from. 75/75 algebra tests pass on 4 devices. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 17 +++++++++++++++-- 1 file changed, 15 insertions(+), 2 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 9b9f3d9021..f432a12f20 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -2319,10 +2319,23 @@ fn multiply_batch_grouped( // 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. - let submits: Vec> = by_dev + // Phase 1 runs the shards' marshalling CONCURRENTLY, which is safe here precisely because + // no permit is taken until phase 2 — the split is what makes this legal, and it is also + // where the time is. Sharding divides the products, so marshalling shards one after another + // costs the same TOTAL host work as one unsharded block, and an earlier comment concluded + // from that there was nothing to parallelise. That is true of total work and false of the + // critical path: a row block is not done until its slowest shard is, so serial marshalling + // makes each block's latency the SUM over devices rather than the max. Measured end to end + // at stem 200: parallel marshal 2551 s, serial marshal 3136 s and 3140 s across two + // independent runs, with `marshal` itself totalling 3444 s. + use maybe_rayon::prelude::*; + let nonempty: Vec<(usize, &Vec)> = by_dev .iter() .enumerate() .filter(|(_, ps)| !ps.is_empty()) + .collect(); + let submits: Vec> = nonempty + .into_maybe_par_iter() .map(|(d, ps)| multiply_batch_block(algebra, num_cols, r0, r1 - r0, ps, mode, d)) .collect(); let waits: Vec = submits.into_iter().map(|s| s()).collect(); @@ -2355,7 +2368,7 @@ type BlockWait = Box Bytes + Send>; /// 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 + 'a>; +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 From 3ab37920e8809a687cb0456423ca3bbfdffc485f Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Tue, 4 Aug 2026 20:07:45 -0400 Subject: [PATCH 081/127] milnor_gpu: split `prep` into intern / basis / tgei `prep` was a single number, so it could say marshalling costs 3444s but not whether that is worth a representation change or merely worth accepting. The three parts have different fixes and want measuring apart: intern the per-launch HashMap<(i32, usize), u32> over every product. Pure representation tax: every R already has a resident device identity (RInfo) that the caller could carry instead of an algebra-local (degree, index) pair. Deleting this needs a different id in GpuProduct, not a new data structure. basis term_off prefix sums, the max_s_degree scan, and ensure_basis (which may upload). Shrinks by keeping the basis warm. tgei the term_gei fill. 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. Measure before building either: two ablations this session predicted wins they did not deliver, both by perturbing more than they named. 75/75 algebra tests pass on 4 devices. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 37 +++++++++++++++++++- 1 file changed, 36 insertions(+), 1 deletion(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index f432a12f20..f00f624d33 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -450,6 +450,12 @@ static BATCH_DEPTH_MAX: AtomicU64 = AtomicU64::new(0); /// 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) { @@ -2418,6 +2424,14 @@ fn multiply_batch_block<'a>( }); 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. @@ -2448,6 +2462,11 @@ fn multiply_batch_block<'a>( // the fill below so a wait on that lock is not misread as marshalling work. let global_base = tracing::info_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 @@ -2485,6 +2504,10 @@ fn multiply_batch_block<'a>( } } } + // 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; @@ -3236,6 +3259,15 @@ fn multiply_batch_block<'a>( (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`]. @@ -3307,7 +3339,7 @@ fn multiply_batch_block<'a>( 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}%", + fence={fence_s:.1}s pipeline={:.0}% | intern={:.1}s basis={:.1}s tgei={:.1}s", 100.0 * prep_s / total, 100.0 * permit_s / total, 100.0 * lock_s / total, @@ -3319,6 +3351,9 @@ fn multiply_batch_block<'a>( // 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, ); } From 220e8f0e4b2ae0eb28f6e7d4e86a1ae6875cb5bc Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Tue, 4 Aug 2026 20:25:58 -0400 Subject: [PATCH 082/127] nassau, milnor_gpu: demote the fine-grained spans below info MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit RUST_LOG=info was emitting ~122 MB of log in the first few minutes of a stem-200, which is both unreadable and a cost of its own on a run whose host side is already the larger share. The spans are worth keeping — they are what finally located the copying and the one-kernel-deep pipeline — but not all of them belong at info. Split by how often each fires: trace the 13 per-signature spans from d167b6f198 (sig_*, zs_*, img_*, extend_image). ~555k events per run each. debug the per-launch spans (gpu_submit, marshal_terms, ensure_basis, pair_prepass). ~10^5-10^6 events per run each. info unchanged: step_resolution_with_subalgebra (per bidegree) and `step`, both of which predate this branch. That leaves `info` at exactly master's span set, with everything added since reachable one or two levels down — so the progress parsing that reads per-bidegree events still works against a default run, and the detailed budgets are a RUST_LOG away rather than always on. 75/75 algebra tests pass on 4 devices. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 8 ++-- ext/src/nassau.rs | 49 ++++++++++---------- 2 files changed, 28 insertions(+), 29 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index f00f624d33..f578c88cdd 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -2237,7 +2237,7 @@ fn multiply_batch_grouped( // 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::info_span!( + let prepass = tracing::debug_span!( "pair_prepass", products = products.len(), new_r = tracing::field::Empty, @@ -2460,7 +2460,7 @@ fn multiply_batch_block<'a>( 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::info_span!("ensure_basis", max_s_degree) + 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 @@ -2476,7 +2476,7 @@ fn multiply_batch_block<'a>( { // 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::info_span!( + let _marshal_span = tracing::debug_span!( "marshal_terms", products = products.len(), terms = total_terms @@ -2801,7 +2801,7 @@ fn multiply_batch_block<'a>( // `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::info_span!( + let submit_span = tracing::debug_span!( "gpu_submit", dev = dev, rows = num_rows, diff --git a/ext/src/nassau.rs b/ext/src/nassau.rs index 52c845c2fa..80372e49f1 100644 --- a/ext/src/nassau.rs +++ b/ext/src/nassau.rs @@ -812,7 +812,7 @@ impl> Resolution { next_dim, ), }; - let mut masked_matrix = tracing::info_span!( + let mut masked_matrix = tracing::trace_span!( "zs_assemble", rows = target_masked_dim, cols = next_masked_dim @@ -825,13 +825,13 @@ impl> Resolution { m }); - tracing::info_span!( + tracing::trace_span!( "zs_row_reduce", rows = target_masked_dim, cols = next_masked_dim ) .in_scope(|| masked_matrix.row_reduce()); - let kernel = tracing::info_span!("zs_kernel").in_scope(|| masked_matrix.compute_kernel()); + let kernel = tracing::trace_span!("zs_kernel").in_scope(|| masked_matrix.compute_kernel()); Self::write_qi( &mut f, @@ -868,17 +868,18 @@ impl> Resolution { &source_mask, target_dim, ); - let mut n = tracing::info_span!("img_assemble", rows = source_mask.len()).in_scope(|| { + 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::info_span!("img_row_reduce", rows = source_mask.len()).in_scope(|| n.row_reduce()); + tracing::trace_span!("img_row_reduce", rows = source_mask.len()) + .in_scope(|| n.row_reduce()); let next_row = n.rows(); - let num_new_gens = tracing::info_span!("extend_image") + 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() { @@ -926,7 +927,7 @@ impl> Resolution { // ~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::info_span!("sig_masks").entered(); + let _sm = tracing::trace_span!("sig_masks").entered(); target_mask.clear(); next_mask.clear(); target_mask.extend(subalgebra.signature_mask( @@ -945,23 +946,21 @@ impl> Resolution { )); drop(_sm); - let full_matrix = - tracing::info_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 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 = tracing::info_span!( + let mut masked_matrix = tracing::trace_span!( "sig_assemble", rows = target_mask.len(), cols = next_mask.len() @@ -979,14 +978,14 @@ impl> Resolution { // 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::info_span!( + tracing::trace_span!( "sig_row_reduce", rows = target_mask.len(), cols = next_mask.len() ) .in_scope(|| masked_matrix.row_reduce()); - let qi = tracing::info_span!("sig_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(); @@ -1002,7 +1001,7 @@ impl> Resolution { } } - let _lift = tracing::info_span!("sig_lift", gens = xs.len()).entered(); + 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; @@ -1021,7 +1020,7 @@ impl> Resolution { } } drop(_lift); - tracing::info_span!("sig_write_qi").in_scope(|| { + tracing::trace_span!("sig_write_qi").in_scope(|| { Self::write_qi( &mut f, &mut scratch, From 51a89494fceba4d7076e7131781ace5eeed2b7f8 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Tue, 4 Aug 2026 20:42:12 -0400 Subject: [PATCH 083/127] milnor_gpu: flag the readback keepalive bug at the site A full stem-200 crashes with CUDA_ERROR_LAUNCH_FAILED because the device closure now returns before the kernel retires and drops its input handles. Documented at the readback rather than left for the next reader to rediscover from a crash; the fix is a keepalive bundle returned alongside the future. See memory nassau-gpu-readback-keepalive-bug. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index f578c88cdd..faf34118bb 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -3187,7 +3187,21 @@ fn multiply_batch_block<'a>( ); } - // Issue the readback but DO NOT wait for it. `read_async` enqueues the device→host copy + // KNOWN BUG, NOT YET FIXED — a full stem-200 crashes with + // `CUDA_ERROR_LAUNCH_FAILED, "unspecified launch failure"` (observed at `max_t=304` + // after 2673 s; NOT the cleanup race, `NASSAU_GPU_CLEANUP_EVERY=0` was set). + // + // Because this closure no longer blocks, it RETURNS AND DROPS ITS INPUT HANDLES while + // the kernel may still be running, so cubecl can hand those pages to a later allocation + // that the running kernel is still reading. `read_one` blocked, which kept them alive + // until the kernel retired. `BufferArg::from_raw_parts` consuming the handles does NOT + // save this: `cs_seg`/`mk_seg` are `Vec<(Handle, usize)>` CLONES with their own + // lifetimes, as are the dummies and the `r*_h` group. + // + // Fix: return `(DynFut, Vec)` and drop the vec after `block_on`, exactly as the + // permit now does — clone each handle into it immediately before `launch_unchecked`. + // + // 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 From 061c6672c7a8d3701d0786d516dcbcc511a90412 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Tue, 4 Aug 2026 21:01:51 -0400 Subject: [PATCH 084/127] milnor_gpu: keep launch buffers alive until the readback completes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fixes the CUDA_ERROR_LAUNCH_FAILED that killed a stem-200 at max_t=304. With the blocking `read_one`, "this closure returns" and "the kernel is done" were the same instant, so every buffer the launch read stayed alive long enough by accident. `read_async` separates them: the closure returns while the kernel may still be running, and each dropped handle lets cubecl hand its pages to a later allocation that the running kernel is still reading. I dismissed this when writing the pipelining, reasoning that `BufferArg::from_raw_parts` consumes the handles at launch. That is exactly backwards — consuming them is what kills them, since the argument dies with the launch call. The `sa!` segment clones are temporaries too. Every handle the launch touches is now cloned into a keepalive vec before the launch consumes the originals, returned alongside the future, and dropped only after `block_on` — the same lifetime discipline the permit got one commit earlier, and for the same reason: under `read_async` the wait, not the closure's return, is the moment the device is finished. Covers the `Transient` enumeration launch too. Its inputs (`epp_h`, `er_h`, `ec_h`, `eco_h`, `emo_h`, `cnt_scratch`) would otherwise die with the match arm while the enum kernel runs; `cs_scratch`/`mk_scratch` survive into `cs_seg`/`mk_seg` and were already covered. 75/75 algebra tests pass on 4 devices. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 81 ++++++++++++++++---- 1 file changed, 64 insertions(+), 17 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index faf34118bb..df2b1e4757 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -2852,6 +2852,12 @@ fn multiply_batch_block<'a>( // `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!( @@ -2909,6 +2915,14 @@ fn multiply_batch_block<'a>( let ec_h = client.create_from_slice(u32::as_bytes(&enum_cols)); let eco_h = client.create_from_slice(u64::as_bytes(&r_cs_offset)); let emo_h = client.create_from_slice(u64::as_bytes(&r_mk_offset)); + enum_keep.extend([ + cnt_scratch.clone(), + epp_h.clone(), + er_h.clone(), + ec_h.clone(), + eco_h.clone(), + emo_h.clone(), + ]); unsafe { enumerate_admissible_kernel::launch_unchecked::( &client, @@ -3089,6 +3103,46 @@ fn multiply_batch_block<'a>( 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 { @@ -3187,21 +3241,7 @@ fn multiply_batch_block<'a>( ); } - // KNOWN BUG, NOT YET FIXED — a full stem-200 crashes with - // `CUDA_ERROR_LAUNCH_FAILED, "unspecified launch failure"` (observed at `max_t=304` - // after 2673 s; NOT the cleanup race, `NASSAU_GPU_CLEANUP_EVERY=0` was set). - // - // Because this closure no longer blocks, it RETURNS AND DROPS ITS INPUT HANDLES while - // the kernel may still be running, so cubecl can hand those pages to a later allocation - // that the running kernel is still reading. `read_one` blocked, which kept them alive - // until the kernel retired. `BufferArg::from_raw_parts` consuming the handles does NOT - // save this: `cs_seg`/`mk_seg` are `Vec<(Handle, usize)>` CLONES with their own - // lifetimes, as are the dummies and the `r*_h` group. - // - // Fix: return `(DynFut, Vec)` and drop the vec after `block_on`, exactly as the - // permit now does — clone each handle into it immediately before `launch_unchecked`. - // - // Issue the readback but DO NOT wait for it. `read_async` enqueues the device→host copy + // 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 @@ -3237,7 +3277,9 @@ fn multiply_batch_block<'a>( client.memory_cleanup(); } - result + // 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) }) }); @@ -3251,7 +3293,7 @@ fn multiply_batch_block<'a>( // — 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, timing) = pending.wait(); + let ((fut, keepalive), timing) = pending.wait(); let launch_ms = t_launch.elapsed().as_secs_f64() * 1e3; let t_fence = std::time::Instant::now(); @@ -3263,6 +3305,11 @@ fn multiply_batch_block<'a>( // 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( From bdbd4eb4b89a4987708ad862a6bf93484ea2cd7a Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Tue, 4 Aug 2026 22:37:23 -0400 Subject: [PATCH 085/127] Revert "milnor_gpu: marshal the shards in parallel again" MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This reverts ef16f934ef. Measured worse, not better: 3373s with parallel marshal against 3136s serial, same code otherwise. The reasoning behind it — that a row block's latency is the SUM over devices rather than the max when shards marshal serially — is arithmetically true and turned out not to govern the wall time. It was also the wrong inference from the right observation. Scoped threads (2551s) do marshal shards concurrently, but that is not what makes them fast: a 512-thread private pool (8c2365cf41, callers x devices) also marshals concurrently off the rayon pool, and lost at 3743s. Five fan-out shapes have now been measured, exactly one is fast, and no proposed explanation has survived contact with the other four. Not restoring the scoped-thread shape despite it holding the record. Its live thread peak is callers x devices ~512 on top of a ~644-thread baseline, against a per-UID RLIMIT_NPROC of 4096 shared machine-wide across every process this user runs — and two concurrent resolutions is a thing we do for A/B benches. Reintroducing the only shape that pushes toward that limit, for a 600s effect nobody can explain, is not a good trade on a shared node. The current shape is one GPU worker per device plus rayon: ~130 threads, no churn. Kept from this line of work: the permit-scope regression fix, the readback keepalive, and the prep breakdown, all of which stand on their own. 75/75 algebra tests pass on 4 devices. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 17 ++--------------- 1 file changed, 2 insertions(+), 15 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index df2b1e4757..3f36f067a5 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -2325,23 +2325,10 @@ fn multiply_batch_grouped( // 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. - // Phase 1 runs the shards' marshalling CONCURRENTLY, which is safe here precisely because - // no permit is taken until phase 2 — the split is what makes this legal, and it is also - // where the time is. Sharding divides the products, so marshalling shards one after another - // costs the same TOTAL host work as one unsharded block, and an earlier comment concluded - // from that there was nothing to parallelise. That is true of total work and false of the - // critical path: a row block is not done until its slowest shard is, so serial marshalling - // makes each block's latency the SUM over devices rather than the max. Measured end to end - // at stem 200: parallel marshal 2551 s, serial marshal 3136 s and 3140 s across two - // independent runs, with `marshal` itself totalling 3444 s. - use maybe_rayon::prelude::*; - let nonempty: Vec<(usize, &Vec)> = by_dev + let submits: Vec> = by_dev .iter() .enumerate() .filter(|(_, ps)| !ps.is_empty()) - .collect(); - let submits: Vec> = nonempty - .into_maybe_par_iter() .map(|(d, ps)| multiply_batch_block(algebra, num_cols, r0, r1 - r0, ps, mode, d)) .collect(); let waits: Vec = submits.into_iter().map(|s| s()).collect(); @@ -2374,7 +2361,7 @@ type BlockWait = Box Bytes + Send>; /// 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>; +type BlockSubmit<'a> = Box BlockWait + '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 From 02dc8f0565aca79559f57dcb15a44c9870577e8e Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Wed, 5 Aug 2026 00:41:23 -0400 Subject: [PATCH 086/127] milnor_gpu: marshal the shards on dedicated threads MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The scoped-thread fan-out really is faster, and I owe it a correction: I retracted that claim yesterday on the grounds that its 2551s log was a killed run. The log WAS unverifiable — I had deleted it — but rerunning 0bc89566c6 to completion settles it. Verified stem-200s, both max_t=310, closed=28013, calls=898000, pairs=5.87e13, zero crashes: scoped threads 2439s HEAD (serial) 3138s 22%. `marshal` 1134s vs 3283s and mean queue depth 4.1 vs 1.9 — the devices are simply fed better. Crucially this is NOT just "parallel marshal": routing the identical work through rayon (ef16f934ef) measured 3373s, WORSE than serial. What pays is parallelism that does not contend with the ~128 resolution workers already occupying the global pool. Dedicated threads have it; a rayon par_iter does not. That also explains the 512-thread private pool losing — a shared pool reintroduces queueing between callers. Ported onto the current tree rather than restoring 0bc89566c6 wholesale, because the original took each shard's permit while its siblings were still marshalling. That was harmless there only because the permit was the broken pre-a4aed87c64 one that bounded nothing (permit=0.0s in that run's stats); with the budget actually enforcing it is the deadlock GpuBudget documents. Marshal moves to scoped threads, permits stay in the rayon-free phase. Thread budget: gpu_count() spawns per row block, ~1M per run, live peak callers x devices ~512 against a per-UID RLIMIT_NPROC of 4096 shared machine-wide. PIDs recycle on join, so that is ~1% churn, not accumulation. 75/75 algebra tests pass on 4 devices. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 38 ++++++++++++++++++-- 1 file changed, 35 insertions(+), 3 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 3f36f067a5..7af1d354a9 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -2325,12 +2325,44 @@ fn multiply_batch_grouped( // 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. - let submits: Vec> = by_dev + // Phase 1 marshals the shards concurrently on DEDICATED threads, not on the rayon pool. + // + // That distinction is the whole measurement. Verified complete stem-200 runs: scoped threads + // 2439 s against 3138 s for serial marshal, with `marshal` 1134 s vs 3283 s and mean queue + // depth 4.1 vs 1.9 — the devices are simply fed better. But routing the same work through + // rayon (`into_maybe_par_iter`, ef16f934ef) measured 3373 s, i.e. WORSE than serial. So it is + // not parallelism as such that pays; it is parallelism that does not contend with the ~128 + // resolution workers already occupying the global pool. + // + // Cost: `gpu_count()` spawns per row block, ~1M over a run. Live threads peak at + // callers x devices ~512 against a per-UID `RLIMIT_NPROC` of 4096 shared machine-wide, and + // PIDs recycle on join, so the churn is ~1% of runtime rather than an accumulation. Do NOT + // "tidy" this into a shared pool without re-measuring — that is exactly what ef16f934ef did. + // + // Permits are still taken in phase 2, never here. The original scoped fan-out took each + // shard's permit while its siblings were still marshalling, which was harmless only because + // the permit was the broken pre-`a4aed87c64` one that bounded nothing (`permit=0.0s` in that + // run's stats). With the budget actually enforcing, that shape is the deadlock [`GpuBudget`] + // documents; keeping acquisition in the rayon-free phase gets the throughput without it. + let shards: Vec<(usize, &Vec)> = by_dev .iter() .enumerate() .filter(|(_, ps)| !ps.is_empty()) - .map(|(d, ps)| multiply_batch_block(algebra, num_cols, r0, r1 - r0, ps, mode, d)) .collect(); + let submits: Vec> = std::thread::scope(|scope| { + let handles: Vec<_> = shards + .into_iter() + .map(|(d, ps)| { + scope.spawn(move || { + multiply_batch_block(algebra, num_cols, r0, r1 - r0, ps, mode, d) + }) + }) + .collect(); + handles + .into_iter() + .map(|h| h.join().expect("a shard's marshal panicked")) + .collect() + }); let waits: Vec = submits.into_iter().map(|s| s()).collect(); let partials: Vec = waits.into_iter().map(|w| w()).collect(); let mut it = partials.into_iter(); @@ -2361,7 +2393,7 @@ type BlockWait = Box Bytes + Send>; /// 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 + 'a>; +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 From 9aa0211526e1b596ae23c0d19ee5e2cb47f89954 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Wed, 5 Aug 2026 01:48:01 -0400 Subject: [PATCH 087/127] milnor_gpu: run each shard's whole pipeline on its own thread MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Concurrent SUBMISSION is the lever, not concurrent marshalling. Isolating that took three verified stem-200s (all max_t=310, closed=28013, calls=898000, pairs=5.87e13, zero crashes): scoped, whole pipeline per thread 2439s marshal 1134s depth 4.1 scoped marshal, serial submit 3358s marshal 2439s depth 1.8 serial marshal, serial submit 3138s marshal 3283s depth 1.9 rayon marshal, serial submit 3373s The middle row is the one that settles it: my previous commit cut marshal by 840s and came out 220s SLOWER. Marshal time does not predict wall time; queue depth does. Parallelising the marshal alone still leaves each shard's submission behind the previous one, so the devices see ~2 blocks queued and idle between them. Only submitting concurrently reaches depth 4.1. So the permits do get acquired from several threads at once, which the last commit avoided on deadlock grounds. That concern does not apply in the default configuration: the marshal contains no rayon at all — the term_gei fill is deliberately sequential after a par_iter over it measured a 146s stall — so there is no join for a permit-blocked steal to wedge, and every permit holder is either on the GPU or progressing. The exception is the NASSAU_GPU_BASIS_PASSTHROUGH diagnostic, whose per-product fill IS a par_iter and which therefore is exactly GpuBudget's documented deadlock; that path runs the shards serially instead. Unlike 0bc89566c6 this keeps the working byte budget and the readback keepalive, so the throughput comes without the two bugs that shape shipped with. 75/75 tests pass on 4 devices; multiply_batch also verified under NASSAU_GPU_BASIS_PASSTHROUGH=1 for the serial fallback. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 71 +++++++++++--------- 1 file changed, 40 insertions(+), 31 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 7af1d354a9..836e7156ca 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -2325,46 +2325,55 @@ fn multiply_batch_grouped( // 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. - // Phase 1 marshals the shards concurrently on DEDICATED threads, not on the rayon pool. + // Each shard runs its WHOLE pipeline — marshal, permit, submit, wait — on its own thread. // - // That distinction is the whole measurement. Verified complete stem-200 runs: scoped threads - // 2439 s against 3138 s for serial marshal, with `marshal` 1134 s vs 3283 s and mean queue - // depth 4.1 vs 1.9 — the devices are simply fed better. But routing the same work through - // rayon (`into_maybe_par_iter`, ef16f934ef) measured 3373 s, i.e. WORSE than serial. So it is - // not parallelism as such that pays; it is parallelism that does not contend with the ~128 - // resolution workers already occupying the global pool. + // Concurrent SUBMISSION is the lever, and it took three measured runs to isolate. Verified + // complete stem-200s: + // + // scoped, whole pipeline per thread 2439 s marshal 1134 s depth 4.1 + // 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 + // + // Read the middle row carefully: it CUT marshal by 840 s and got 220 s SLOWER. Marshal time + // does not predict wall time; queue depth does. Parallelising the marshal alone leaves every + // shard's submission behind the previous one, so the devices still see ~2 blocks queued and + // idle between them. Only submitting concurrently gets depth to 4.1 and the GPUs fed. // // Cost: `gpu_count()` spawns per row block, ~1M over a run. Live threads peak at // callers x devices ~512 against a per-UID `RLIMIT_NPROC` of 4096 shared machine-wide, and - // PIDs recycle on join, so the churn is ~1% of runtime rather than an accumulation. Do NOT - // "tidy" this into a shared pool without re-measuring — that is exactly what ef16f934ef did. - // - // Permits are still taken in phase 2, never here. The original scoped fan-out took each - // shard's permit while its siblings were still marshalling, which was harmless only because - // the permit was the broken pre-`a4aed87c64` one that bounded nothing (`permit=0.0s` in that - // run's stats). With the budget actually enforcing, that shape is the deadlock [`GpuBudget`] - // documents; keeping acquisition in the rayon-free phase gets the throughput without it. + // PIDs recycle on join, so this is ~1% churn, not accumulation. Do NOT convert this to a + // shared pool or a rayon `par_iter` without re-measuring END TO END: both were tried + // (3743 s and 3373 s) and both lost, because they queue behind the ~128 resolution workers + // instead of running beside them. let shards: Vec<(usize, &Vec)> = by_dev .iter() .enumerate() .filter(|(_, ps)| !ps.is_empty()) .collect(); - let submits: Vec> = std::thread::scope(|scope| { - let handles: Vec<_> = shards - .into_iter() - .map(|(d, ps)| { - scope.spawn(move || { - multiply_batch_block(algebra, num_cols, r0, r1 - r0, ps, mode, d) - }) - }) - .collect(); - handles - .into_iter() - .map(|h| h.join().expect("a shard's marshal panicked")) - .collect() - }); - let waits: Vec = submits.into_iter().map(|s| s()).collect(); - let partials: Vec = waits.into_iter().map(|w| w()).collect(); + let run_shard = |d: usize, ps: &Vec| -> Bytes { + multiply_batch_block(algebra, num_cols, r0, r1 - r0, ps, mode, d)()() + }; + // Threads hold their permit while siblings marshal. That is safe HERE and only here 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 either on the GPU or making progress. The one exception is the + // `NASSAU_GPU_BASIS_PASSTHROUGH` diagnostic, whose per-product fill IS a par_iter; that + // combination is exactly [`GpuBudget`]'s documented deadlock, so it runs serially instead. + let partials: Vec = if basis_passthrough() { + shards.into_iter().map(|(d, ps)| run_shard(d, ps)).collect() + } else { + std::thread::scope(|scope| { + let handles: Vec<_> = shards + .into_iter() + .map(|(d, ps)| scope.spawn(move || run_shard(d, ps))) + .collect(); + handles + .into_iter() + .map(|h| h.join().expect("a shard panicked")) + .collect() + }) + }; let mut it = partials.into_iter(); let mut acc = it .next() From 63595b7f60ca74e6935be245a9bbd590a3fe8dea Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Wed, 5 Aug 2026 09:12:07 -0400 Subject: [PATCH 088/127] milnor_gpu: persistent per-caller shard helpers, no more spawn churn MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replaces the `std::thread::scope` fan-out, which spawned gpu_count() OS threads per row block — ~900k over a stem-200 run. Sized rayon_threads x (gpu_count - 1): each calling thread lazily creates gpu_count()-1 helpers on its first fan-out and reuses them for the process's life, running the remaining shard itself. The aggregate therefore falls out of the rayon pool size instead of being a second knob that can drift from it — shrink the rayon pool and this shrinks in proportion. Deliberately PRIVATE per caller rather than one shared pool. The value of the fan-out is that each shard's whole pipeline runs concurrently so every device gets work at once; a shared queue puts a shard behind other callers' shards and re-serialises exactly that. Measured: ~3743s for a 512-thread shared rayon pool and 3373s through the global rayon pool, against 2412s here. These helpers do not steal and share nothing. Jobs own their data, so no lifetime erasure and no unsafe: `by_dev` already builds an owned Vec per shard, and the algebra now arrives as Arc — which the production caller (nassau_gpu.rs) already had, since Module::algebra() returns Arc. Expect no speedup: the churn cost ~0.01% of wall time. This is a robustness change — it keeps the live thread count off RLIMIT_NPROC (4096 per-UID, shared machine-wide, over a ~644-thread baseline) and stops ~900k spawn/join cycles per run. Validated end to end anyway, because four "obviously equivalent" changes regressed this week. 75/75 tests on 4 devices, plus multiply_batch under NASSAU_GPU_BASIS_PASSTHROUGH=1 for the serial fallback. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- .../algebra/benches/nassau_milnor_gpu.rs | 3 +- ext/crates/algebra/src/algebra/milnor_gpu.rs | 187 +++++++++++++----- 2 files changed, 137 insertions(+), 53 deletions(-) diff --git a/ext/crates/algebra/benches/nassau_milnor_gpu.rs b/ext/crates/algebra/benches/nassau_milnor_gpu.rs index 7b22d83b6a..a8d0d3fb47 100644 --- a/ext/crates/algebra/benches/nassau_milnor_gpu.rs +++ b/ext/crates/algebra/benches/nassau_milnor_gpu.rs @@ -74,7 +74,8 @@ mod gpu { pub fn nassau_milnor_gpu(c: &mut Criterion) { // Exactly the algebra Nassau uses: the full Milnor algebra at p=2, stable (not unstable). - let algebra = MilnorAlgebra::new(TWO, false); + 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 { diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 836e7156ca..538184556e 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -333,6 +333,69 @@ mod gpu_thread { } } +/// 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. @@ -378,7 +441,10 @@ fn narrow_u16(v: u32) -> u16 { u16::try_from(v).expect("admissible/term entry exceeds u16") } -use std::sync::atomic::{AtomicBool, 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 @@ -2071,7 +2137,7 @@ impl std::fmt::Debug for BatchOutput { } pub fn multiply_batch_on_gpu( - algebra: &MilnorAlgebra, + algebra: &Arc, num_cols: usize, num_rows: usize, products: &[GpuProduct], @@ -2148,7 +2214,7 @@ pub fn cpu_multiply_batch( } fn multiply_batch_gpu_inner( - algebra: &MilnorAlgebra, + algebra: &Arc, num_cols: usize, num_rows: usize, products: &[GpuProduct], @@ -2213,7 +2279,7 @@ fn multiply_batch_gpu_inner( } fn multiply_batch_grouped( - algebra: &MilnorAlgebra, + algebra: &Arc, num_cols: usize, num_rows: usize, products: &[GpuProduct], @@ -2325,54 +2391,71 @@ fn multiply_batch_grouped( // 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 — on its own thread. - // - // Concurrent SUBMISSION is the lever, and it took three measured runs to isolate. Verified - // complete stem-200s: + // Each shard runs its WHOLE pipeline — marshal, permit, submit, wait — concurrently, so + // all `gpu_count()` devices receive work at once. Verified stem-200s: // - // scoped, whole pipeline per thread 2439 s marshal 1134 s depth 4.1 + // 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 // - // Read the middle row carefully: it CUT marshal by 840 s and got 220 s SLOWER. Marshal time - // does not predict wall time; queue depth does. Parallelising the marshal alone leaves every - // shard's submission behind the previous one, so the devices still see ~2 blocks queued and - // idle between them. Only submitting concurrently gets depth to 4.1 and the GPUs fed. + // 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. // - // Cost: `gpu_count()` spawns per row block, ~1M over a run. Live threads peak at - // callers x devices ~512 against a per-UID `RLIMIT_NPROC` of 4096 shared machine-wide, and - // PIDs recycle on join, so this is ~1% churn, not accumulation. Do NOT convert this to a - // shared pool or a rayon `par_iter` without re-measuring END TO END: both were tried - // (3743 s and 3373 s) and both lost, because they queue behind the ~128 resolution workers - // instead of running beside them. - let shards: Vec<(usize, &Vec)> = by_dev - .iter() + // 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 run_shard = |d: usize, ps: &Vec| -> Bytes { - multiply_batch_block(algebra, num_cols, r0, r1 - r0, ps, mode, d)()() - }; - // Threads hold their permit while siblings marshal. That is safe HERE and only here 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 either on the GPU or making progress. The one exception is the - // `NASSAU_GPU_BASIS_PASSTHROUGH` diagnostic, whose per-product fill IS a par_iter; that - // combination is exactly [`GpuBudget`]'s documented deadlock, so it runs serially instead. - let partials: Vec = if basis_passthrough() { - shards.into_iter().map(|(d, ps)| run_shard(d, ps)).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 { - std::thread::scope(|scope| { - let handles: Vec<_> = shards - .into_iter() - .map(|(d, ps)| scope.spawn(move || run_shard(d, ps))) - .collect(); - handles - .into_iter() - .map(|h| h.join().expect("a shard panicked")) - .collect() - }) + // 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 @@ -4400,7 +4483,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)); algebra.compute_basis(max_degree); // Process one degree per launch. At high degree the full master is tens of GB, so batching every @@ -4576,7 +4659,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 = 130; algebra.compute_basis(max_degree); @@ -4663,7 +4746,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)); // 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; @@ -4709,7 +4792,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); @@ -4750,7 +4833,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); @@ -4823,7 +4906,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); @@ -4915,7 +4998,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 = 44; algebra.compute_basis(max_degree); algebra.compute_seqno_tables(max_degree); @@ -4980,7 +5063,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 = 48; algebra.compute_basis(max_degree); algebra.compute_seqno_tables(max_degree); @@ -5120,7 +5203,7 @@ mod tests { let num_rows = 32usize; let p = fp::prime::ValidPrime::new(2); - let algebra = MilnorAlgebra::new(p, false); + let algebra = Arc::new(MilnorAlgebra::new(p, false)); algebra.compute_basis(max_degree); algebra.compute_seqno_tables(max_degree); @@ -5316,7 +5399,7 @@ mod tests { let spread = env_num("NASSAU_BENCH_SPREAD", 4) as i32; let p = fp::prime::ValidPrime::new(2); - let algebra = MilnorAlgebra::new(p, false); + 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. From 38bbdd78964dcb83a337a3fbf1da1450f6bb12dc Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Wed, 5 Aug 2026 10:01:24 -0400 Subject: [PATCH 089/127] milnor_gpu: crossbeam channels for the shard helpers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The persistent helpers measured 2499s against 2412s for the thread::scope fan-out they replaced — 87s, and real rather than scatter: two near-identical serial configs came in at 3136s and 3138s, so run-to-run variance here is about +/-2s. The suspect is wake latency. A freshly spawned thread starts running immediately on a warm core; a parked helper must be woken from a blocking recv, which is a futex wake plus whatever placement the scheduler picks, and 384 mostly-idle helpers pay that on every block. std's mpsc parks almost at once; crossbeam-channel spins briefly first, and it is already a gpu-feature dependency. Applied to both directions — job dispatch and the result hand-back — since the caller blocks on the latter for every shard it farms out. 75/75 tests pass on 4 devices. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 538184556e..9c6fd6c8e3 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -360,14 +360,16 @@ mod gpu_thread { /// 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}; + use std::cell::RefCell; + + use crossbeam_channel::{Sender, unbounded}; 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()) }; + static HELPERS: RefCell>> = const { RefCell::new(Vec::new()) }; } /// Hand `job` to this thread's helper `i`, spawning the helpers on first use. @@ -379,7 +381,7 @@ mod shard_pool { HELPERS.with(|h| { let mut h = h.borrow_mut(); while h.len() <= i { - let (tx, rx) = mpsc::channel::(); + let (tx, rx) = unbounded::(); let idx = h.len(); std::thread::Builder::new() .name(format!("nassau-shard{idx}")) @@ -2438,7 +2440,7 @@ fn multiply_batch_grouped( .into_iter() .enumerate() .map(|(i, (d, ps))| { - let (tx, rx) = std::sync::mpsc::sync_channel::(1); + let (tx, rx) = crossbeam_channel::bounded::(1); let alg = algebra.clone(); shard_pool::dispatch( i, From 3c31f87a6e4ac51c378970e86aab694a10c523c9 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Wed, 5 Aug 2026 14:12:12 -0400 Subject: [PATCH 090/127] Revert "milnor_gpu: crossbeam channels for the shard helpers" This reverts 38bbdd7896. Measured worse, on verified complete stem-200s (max_t=310, closed=28013, pairs=5.87e13, zero crashes): thread::scope spawn per block 2412s persistent helpers, std mpsc 2499s persistent helpers, crossbeam 2612s The theory was that the 87s the persistent helpers cost against thread::scope was futex wake latency, and that crossbeam's spin-before-park would recover it. It cost a further 113s instead. With 384 mostly-idle helpers, spinning appears to take CPU the marshalling threads need rather than saving a wake. Back to std mpsc at 2499s, which is the configuration to keep: +87s against thread::scope, in exchange for ~900k fewer spawn/join cycles per run and a thread count that derives from the rayon pool. 75/75 tests pass on 4 devices. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 9c6fd6c8e3..538184556e 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -360,16 +360,14 @@ mod gpu_thread { /// 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; - - use crossbeam_channel::{Sender, unbounded}; + 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()) }; + static HELPERS: RefCell>> = const { RefCell::new(Vec::new()) }; } /// Hand `job` to this thread's helper `i`, spawning the helpers on first use. @@ -381,7 +379,7 @@ mod shard_pool { HELPERS.with(|h| { let mut h = h.borrow_mut(); while h.len() <= i { - let (tx, rx) = unbounded::(); + let (tx, rx) = mpsc::channel::(); let idx = h.len(); std::thread::Builder::new() .name(format!("nassau-shard{idx}")) @@ -2440,7 +2438,7 @@ fn multiply_batch_grouped( .into_iter() .enumerate() .map(|(i, (d, ps))| { - let (tx, rx) = crossbeam_channel::bounded::(1); + let (tx, rx) = std::sync::mpsc::sync_channel::(1); let alg = algebra.clone(); shard_pool::dispatch( i, From 8ceb3b21d180f741cef92b174610cf0a64f63428 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Wed, 5 Aug 2026 16:08:51 -0400 Subject: [PATCH 091/127] milnor_gpu: shard R by a mixed hash, not by accumulated num_mats MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Frontier GPU load was 50.5% on device 0 against ~15% on the other three at stem 210 — a 3.3x spread, up from 1.18x at stem 200, so it worsens with scale and three of four devices sat ~85% idle. Cause is in the policy's own justification. `dev_load` balances accumulated `num_mats`, which its doc defends by assuming `R`s are "used at broadly similar rates". The NASSAU_R_STATS probe at (150,110) says otherwise: of 173930 distinct `R`s, the top 1% carry 31% of references and the top 10% carry 78%. Equal `num_mats` therefore says little about equal work. Two things compounded it: assignment is permanent, and `min_by_key` broke ties toward device 0, which is where the early low-degree `R`s landed while every device still had near-zero load. Hashing balances neither count nor bytes on purpose — it draws each device an independent sample of the joint (size, reference-rate) distribution, and with ~174k `R`s and no single one above a fraction of a percent of references, both concentrate. It is also stateless: no accumulator under the write lock, no first-sight ordering, no tie-break. The mixing is load-bearing. PPart packs entry i at a fixed bit offset, so the low bits are r_1, which tracks internal degree, which tracks work (the probe's hot decile averages degree 70 vs the cold decile's 14). A bare bits() % 4 would partition by r_1. Measured over 168781 real `R`s: worst deviation 0.65% mixed vs 4.55% unmixed. The new test asserts both the overall balance and balance WITHIN degree bands — a hash uniform overall but skewed per band would still starve devices for long stretches, since the frontier walks degree upward. 76/76 tests pass on 4 devices. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- .../algebra/src/algebra/milnor_algebra.rs | 5 + ext/crates/algebra/src/algebra/milnor_gpu.rs | 143 ++++++++++++++++-- 2 files changed, 137 insertions(+), 11 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_algebra.rs b/ext/crates/algebra/src/algebra/milnor_algebra.rs index 370ea80ac5..953c877e23 100644 --- a/ext/crates/algebra/src/algebra/milnor_algebra.rs +++ b/ext/crates/algebra/src/algebra/milnor_algebra.rs @@ -196,6 +196,11 @@ impl PPart { /// 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 } diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 538184556e..99d0ec3e80 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -1013,6 +1013,33 @@ struct RStat { 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 work (the `NASSAU_R_STATS` hot decile averages degree 70 against the cold +/// decile's 14). 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. +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)); @@ -1176,18 +1203,26 @@ fn resident_info(algebra: &MilnorAlgebra, p_part: PPart) -> RInfo { // 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`. // - // Round-robin over first-sight order was the first cut and balances COUNT, not work — `num_mats` - // varies by orders of magnitude between `R`s, so equal counts left the devices badly uneven - // (batch-stats: queue 67% / exec 32%, mean depth 8.4, i.e. devices waiting while others ran). - // Greedy least-loaded is the standard fix and needs no lookahead. + // 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 = host - .dev_load - .iter() - .enumerate() - .min_by_key(|&(i, &load)| (load, i)) - .map(|(i, _)| i) - .expect("at least one device"); + let dev = shard_of(p_part); host.dev_load[dev] += num_mats; let info = RInfo { cs_off: host.cs_len[dev] as u64, @@ -4737,6 +4772,92 @@ mod tests { ); } + /// [`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 + ); + } + } + // 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" + ); + } + /// 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. From 5d2334c331e15954be07ea2c99af5592d7696d66 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Wed, 5 Aug 2026 18:01:50 -0400 Subject: [PATCH 092/127] milnor_gpu: cover the eviction path, and honour CUDA_VISIBLE_DEVICES MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `MasterMode::Transient` — the `NASSAU_GPU_RESIDENT_MAX_DEGREE` eviction path that is the intended stem-300 memory lever — had no end-to-end test. `admissible_enum_gpu_matches` validates the enumeration kernel through a test-only harness, which is a different launch path: it says nothing about whether the scratch that kernel writes is laid out the way `multiply_batch_kernel` indexes it. Every existing batch test took the default `cap == i32::MAX` fast path and never entered the eviction code at all. `multiply_batch_gpu_inner` now takes the cap as an argument rather than reading the process-wide `LazyLock`, so one process can exercise all three regimes, and `multiply_batch_matches_reference_under_eviction` runs the same batch through all-resident, all-transient and the mixed two-launch split against one CPU golden, at both low and high degree. It passes in every regime (7684 products at degree 72), so the split logic and the on-device enumeration feeding the multiply are sound. Separately, `gpu_count()` counted `/proc/driver/nvidia/gpus` — the devices physically in the node, not the ones visible to the process. CUDA renumbers the visible subset to `0..n`, so under `CUDA_VISIBLE_DEVICES=1,2,3` the fourth shard opened device 3 and panicked with `CUDA_ERROR_INVALID_DEVICE`, taking the GPU workers down with it. The standard way to partition GPUs on a shared node silently broke the run. Also records that the enum kernel is bit-exact to degree 240 (1,593,460 R's, 16,949,543,206 matrices), so a high-stem LAUNCH_FAILED is not evidence against it. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 184 ++++++++++++++++++- 1 file changed, 178 insertions(+), 6 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 99d0ec3e80..fa5203182f 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -684,7 +684,14 @@ static RESIDENT_HOST: LazyLock> = LazyLock::new(|| { const MAX_GPUS: usize = 8; /// How many CUDA devices the multiply path spreads work over. `NASSAU_GPU_DEVICES` overrides; -/// otherwise every device the driver exposes is used. +/// otherwise every device *visible to this process* is used. +/// +/// `/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 @@ -692,10 +699,22 @@ const MAX_GPUS: usize = 8; /// and 2.70x at N = 4. fn gpu_count() -> usize { static N: LazyLock = LazyLock::new(|| { - let detected = std::fs::read_dir("/proc/driver/nvidia/gpus") + 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()) @@ -1017,8 +1036,9 @@ static R_STATS: LazyLock>>> = /// /// 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 work (the `NASSAU_R_STATS` hot decile averages degree 70 against the cold -/// decile's 14). Taking `bits() % gpu_count()` would therefore partition by `r_1 mod 4` and could +/// 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. /// @@ -1181,6 +1201,32 @@ pub fn dump_r_stats() { 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}"); } fn resident_info(algebra: &MilnorAlgebra, p_part: PPart) -> RInfo { @@ -2184,7 +2230,7 @@ pub fn multiply_batch_on_gpu( // 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) + multiply_batch_gpu_inner(algebra, num_cols, num_rows, products, resident_degree_cap()) })) { Ok(out) => out, Err(payload) => { @@ -2253,8 +2299,11 @@ fn multiply_batch_gpu_inner( 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 { - let cap = resident_degree_cap(); // 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); @@ -4657,6 +4706,11 @@ mod tests { /// 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); @@ -5110,6 +5164,124 @@ mod tests { ); } + /// 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 { + 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 From dbe1b49e8508552c418d5c89533608d4e0cb8f71 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Wed, 5 Aug 2026 19:54:50 -0400 Subject: [PATCH 093/127] milnor_gpu: order the enumeration by matrix count; tighten its local caps MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three changes to `enumerate_admissible_kernel`, of which only one is a measured win. Sort the enumeration's `R`s by matrix count before launch (+8%). The kernel walks an odometer, so a thread's cost tracks its `R`'s matrix count, and a warp retires only when its slowest lane does. Matrix counts are heavily skewed — the R_STATS Lorenz curve has the hottest 1% of `R`s carrying 31% of references — so in basis order one huge `R` idles 31 lanes. Only the enumeration's own inputs are permuted: it writes each `R` at the absolute offset it is handed, so the scratch comes out byte-identical and the multiply's `r_cs_offset`/`prod_r_index` keep basis order untouched. `enum_warp_utilisation` (new, CPU-only) puts the modelled headroom at 2.16x — 30.9% of lane slots useful in basis order against 66.8% sorted. It cashed out as 1.08x measured, because an SM keeps many warps resident and hides most of the modelled stall. The docstring says so, so nobody sizes a decision on the model again. `ENUM_COL_CAP` was `WORKING_CAP` (32), conflating this cap with the multiply kernel's assembled-p-part array. `cols` is the widest bit-length of a p-part entry, bounded by the `r_1` field width — 11. That made `matrix`, the hottest per-thread local array, 320 u32 instead of 110. Also zero only the region each `R` reaches rather than the full cap. Both are strictly better on local-memory footprint and NEITHER is measurable end-to-end (1.884s -> 1.872s, inside noise) — kept for correctness of the derivation, not for speed. `enum_col_cap_bounds_real_rs` checks the derived caps against every real `R` to degree 400 (actual maxima: 8 rows, 9 cols). For context, in-kernel enumeration already costs 0.81x what merely uploading the same arrays costs, so there is no large multiple waiting here. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 219 +++++++++++++++++-- 1 file changed, 197 insertions(+), 22 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index fa5203182f..845bdecb7f 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -56,11 +56,19 @@ fn ppart_shift_mask() -> (Vec, Vec) { } /// Per-thread local caps for the in-kernel admissible enumeration ([`enumerate_admissible_kernel`]). -/// Each `R` has `rows = |p_part| ≤ MAX_XI_TAU` and `cols ≤ WORKING_CAP` (max bit-length of an entry), +/// 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 = WORKING_CAP; +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; @@ -2854,20 +2862,52 @@ fn multiply_batch_block<'a>( } } - // (Transient) Flatten the cold p-parts (padded to the widest) for the enumeration kernel. The - // per-`R` scratch offsets it writes at are `r_cs_offset`/`r_mk_offset` themselves (u64), passed - // straight through — no u32 narrowing, so a big block's multi-GB scratch is addressed safely. - let (enum_pp, enum_width) = if mode == MasterMode::Transient { + // (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.) + 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 pp = vec![0u32; enum_pp_rows.len() * w]; - for (i, row) in enum_pp_rows.iter().enumerate() { - for (slot, &v) in pp[i * w..i * w + row.len()].iter_mut().zip(row) { - *slot = v; + 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; } } - (pp, w) + let pick_u32 = |src: &[u32]| -> Vec { order.iter().map(|&i| src[i]).collect() }; + let pick_u64 = |src: &[u64]| -> Vec { order.iter().map(|&i| src[i]).collect() }; + ( + pp, + w, + pick_u32(&enum_rows), + pick_u32(&enum_cols), + pick_u64(&r_cs_offset), + pick_u64(&r_mk_offset), + ) } else { - (Vec::new(), 1usize) + ( + 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 @@ -3108,8 +3148,10 @@ fn multiply_batch_block<'a>( let epp_h = client.create_from_slice(u32::as_bytes(&enum_pp)); let er_h = client.create_from_slice(u32::as_bytes(&enum_rows)); let ec_h = client.create_from_slice(u32::as_bytes(&enum_cols)); - let eco_h = client.create_from_slice(u64::as_bytes(&r_cs_offset)); - let emo_h = client.create_from_slice(u64::as_bytes(&r_mk_offset)); + // The permuted offsets, matching `enum_pp`/`enum_rows`/`enum_cols` — NOT the + // multiply's basis-order `r_cs_offset`/`r_mk_offset`. + let eco_h = client.create_from_slice(u64::as_bytes(&enum_cs_out)); + let emo_h = client.create_from_slice(u64::as_bytes(&enum_mk_out)); enum_keep.extend([ cnt_scratch.clone(), epp_h.clone(), @@ -3860,16 +3902,22 @@ fn enumerate_admissible_kernel( let mut totals = Array::::new(ENUM_ROW_CAP); let mut col_sums = Array::::new(ENUM_COL_CAP); let mut masks = Array::::new(ENUM_MASK_CAP); - for i in 0..ENUM_MATRIX_CAP { + // 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 { matrix[i] = 0u32; } - for i in 0..ENUM_ROW_CAP { + for i in 0..rows { totals[i] = 0u32; } - for i in 0..ENUM_COL_CAP { + for i in 0..cs_len { col_sums[i] = 0u32; } - for i in 0..ENUM_MASK_CAP { + for i in 0..mk_len { masks[i] = 0u32; } // Column 0 of the matrix (and the initial masks) is the padded p_part. @@ -4701,6 +4749,122 @@ mod tests { (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)" + ); + } + /// 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 @@ -4798,11 +4962,20 @@ mod tests { 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). + // (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. @@ -4819,10 +4992,12 @@ mod tests { "\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 : {g_kernel:.3} \ - s\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 / upload_secs, + {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, ); } From 2d765afcd4dcc93c50d033952a7aeb38f8aeb405 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Wed, 5 Aug 2026 20:29:43 -0400 Subject: [PATCH 094/127] milnor_gpu: span transient scratch across segments, not one MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The eviction path bound its per-block enumeration scratch as segment 0 alone, so a block was capped at one segment — 4 GiB of u16. A stem-200 block wants 2.17e9 u16 of `masks`, just over the 2^31 line, so `NASSAU_GPU_RESIDENT_MAX_DEGREE=125` died on a hard assert partway through the run. That cap is why theta=125 could not complete. The fix is not a bigger segment. `master_seg_elems()` is deliberately under `u32::MAX` so a segment's element count cannot overflow cubecl's 32-bit array-length metadata — the same truncation class that caused the earlier silent corruption — and raising it would buy 2x and reopen that door. The resident path already spans MASTER_MAX_SEG segments through the same `seg_read`, so the transient path just had to use them: the ceiling goes 4 GiB -> 64 GiB with no new truncation surface. Layout places each `R` wholly inside one segment, with its `col_sums` and `masks` in the same-numbered segment, advancing both cursors to the next boundary together when either run would straddle. That keeps the enumeration as one launch per segment against the existing single-buffer kernel signature, so no segmented-write kernel is needed and the validated kernel is untouched. Offsets handed to each launch are segment-local; the multiply keeps reading global offsets through `seg_read` exactly as before. Verified by running the eviction test with NASSAU_GPU_MASTER_SEG_ELEMS shrunk to 1048576 and 262144, which forces many segments at degree 72 — all regimes still match the CPU golden, as does the full suite at both sizes. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 180 +++++++++++++------ 1 file changed, 123 insertions(+), 57 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 845bdecb7f..c843702d1d 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -2836,8 +2836,10 @@ fn multiply_batch_block<'a>( } MasterMode::Transient => { let (cs_len, mk_len, num_mats) = cold_count(algebra, r.p_part); - r_cs_offset.push(need_cs as u64); - r_mk_offset.push(need_mk as u64); + // 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); @@ -2856,8 +2858,6 @@ fn multiply_batch_block<'a>( enum_rows.push(r.p_part.len() as u32); enum_cols.push(cols); enum_pp_rows.push(r.p_part.iter().collect::>()); - need_cs += num_mats as usize * cs_len as usize; - need_mk += num_mats as usize * mk_len as usize; } } } @@ -2876,12 +2876,65 @@ fn multiply_batch_block<'a>( // 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]; @@ -2890,14 +2943,13 @@ fn multiply_batch_block<'a>( } } let pick_u32 = |src: &[u32]| -> Vec { order.iter().map(|&i| src[i]).collect() }; - let pick_u64 = |src: &[u64]| -> Vec { order.iter().map(|&i| src[i]).collect() }; ( pp, w, pick_u32(&enum_rows), pick_u32(&enum_cols), - pick_u64(&r_cs_offset), - pick_u64(&r_mk_offset), + cs_out, + mk_out, ) } else { ( @@ -3134,57 +3186,71 @@ fn multiply_batch_block<'a>( // scratch is fully written when the multiply reads it (one-stream launches are // ordered, as with `zero_u32` below). const ENUM_THREADS: u32 = 256; - let n_cold = enum_rows.len(); - let cs_cap = need_cs.max(1); - let mk_cap = need_mk.max(1); - assert!( - cs_cap <= seg_elems && mk_cap <= seg_elems, - "transient scratch ({cs_cap}/{mk_cap} u16) exceeds one segment \ - ({seg_elems}); raise NASSAU_GPU_MASTER_SEG_ELEMS" - ); - let cs_scratch = client.empty(cs_cap * size_of::()); - let mk_scratch = client.empty(mk_cap * size_of::()); - let cnt_scratch = client.empty(n_cold.max(1) * size_of::()); - let epp_h = client.create_from_slice(u32::as_bytes(&enum_pp)); - let er_h = client.create_from_slice(u32::as_bytes(&enum_rows)); - let ec_h = client.create_from_slice(u32::as_bytes(&enum_cols)); - // The permuted offsets, matching `enum_pp`/`enum_rows`/`enum_cols` — NOT the - // multiply's basis-order `r_cs_offset`/`r_mk_offset`. - let eco_h = client.create_from_slice(u64::as_bytes(&enum_cs_out)); - let emo_h = client.create_from_slice(u64::as_bytes(&enum_mk_out)); - enum_keep.extend([ - cnt_scratch.clone(), - epp_h.clone(), - er_h.clone(), - ec_h.clone(), - eco_h.clone(), - emo_h.clone(), - ]); - unsafe { - enumerate_admissible_kernel::launch_unchecked::( - &client, - CubeCount::Static( - (n_cold as u32).div_ceil(ENUM_THREADS).max(1), - 1, - 1, - ), - CubeDim::new_1d(ENUM_THREADS), - BufferArg::from_raw_parts(epp_h, enum_pp.len()), - BufferArg::from_raw_parts(er_h, n_cold), - BufferArg::from_raw_parts(ec_h, n_cold), - BufferArg::from_raw_parts(eco_h, n_cold), - BufferArg::from_raw_parts(emo_h, n_cold), - BufferArg::from_raw_parts(cs_scratch.clone(), cs_cap), - BufferArg::from_raw_parts(mk_scratch.clone(), mk_cap), - BufferArg::from_raw_parts(cnt_scratch, n_cold.max(1)), - enum_width, - n_cold, - ); + // 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(), + ]); + unsafe { + enumerate_admissible_kernel::launch_unchecked::( + &client, + CubeCount::Static( + (n_s as u32).div_ceil(ENUM_THREADS).max(1), + 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, + ); + } } - ( - pad_u16(vec![(cs_scratch, cs_cap)]), - pad_u16(vec![(mk_scratch, mk_cap)]), - ) + (pad_u16(cs_segs), pad_u16(mk_segs)) } }; // Resident basis segments (default) or per-launch passthrough buffers (A/B diagnostic) bound From 90e38c2344b8b37724bf79fe153b2fce755c2202 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Thu, 6 Aug 2026 13:08:41 -0400 Subject: [PATCH 095/127] fp: make the multiply/reduction device-sharing test multi-GPU aware, and log it MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `gpu_lock` arbitration was gated on `FP_CUDA_DEVICE == multiply_device()`, where `multiply_device()` read `NASSAU_GPU_DEVICE` — a variable nothing in `algebra` reads any more, left behind when the multiply became multi-GPU. It answered "device 0" however many GPUs the multiply was actually saturating, so `FP_CUDA_DEVICE=2` on a 4-GPU node would conclude "separate devices, no arbitration needed" while the multiply hammered device 2 too. The multiply shards across every visible device, so the test is `FP_CUDA_DEVICE < multiply_devices()`. This is a latent bug, NOT the cause of the theta=125 LAUNCH_FAILED: with default settings both sides resolved to 0, so arbitration was already enabled. Verified by the log line this adds, which exists because `[batch-stats] lock=` cannot distinguish "arbitration off" from "arbitration on but uncontended" — and the answer turned out to be a third thing. What the instrumentation actually exposed: `lock=0.0s` in every run *with arbitration enabled*, because the multiply takes `gpu_lock::shared()` inside the SUBMISSION closure and drops it when submission returns. Submission only enqueues; the kernels run long after. So the multiply releases the guard while its saturating kernels are still executing, the reduction then takes `exclusive()` believing the device is idle, and its thousands of tiny sequential launches interleave with them — the exact overlap the lock exists to prevent. Fixing that means holding the guard through the fence rather than the submit (the same scope error as the earlier GpuPermit bug); left for its own change since it alters the hot path's concurrency and invalidates the ~5 s/stem-200 cost estimate in the lock's docs. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/fp/src/blas/cuda.rs | 53 +++++++++++++++++++++++++++++----- 1 file changed, 46 insertions(+), 7 deletions(-) diff --git a/ext/crates/fp/src/blas/cuda.rs b/ext/crates/fp/src/blas/cuda.rs index 846ef0b81c..a6cf082e14 100644 --- a/ext/crates/fp/src/blas/cuda.rs +++ b/ext/crates/fp/src/blas/cuda.rs @@ -74,19 +74,58 @@ fn context() -> Option<&'static GpuContext> { .ok() .and_then(|v| v.parse().ok()) .unwrap_or(0); - crate::gpu_lock::set_devices_shared(device == multiply_device()); + 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}; gpu_lock arbitration {}", + if shared { "ENABLED" } else { "disabled" } + ); GpuContext::new(device).ok() }) .as_ref() } -/// Which GPU the cubecl Milnor multiply runs on (`NASSAU_GPU_DEVICE`, default 0). Read here only to -/// decide whether the two runtimes share a device; `algebra` owns the actual client construction. -fn multiply_device() -> usize { - std::env::var("NASSAU_GPU_DEVICE") - .ok() - .and_then(|v| v.parse().ok()) +/// 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) } /// Row-major, K-major `u64` limbs — the exact layout `fp_cuda::matmul_b1_raw` From 32bd1f0354b5b8930db08d0c288a6080eec3c2fe Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Thu, 6 Aug 2026 14:28:02 -0400 Subject: [PATCH 096/127] fp: reductions must serialize against each other, not just against the multiply MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `gpu_lock::exclusive()` early-returned whenever `arbitration_needed()` was false, which conflated two independent questions: - does the MULTIPLY have to yield? (only when it shares the reduction's device) - do REDUCTIONS have to serialize? (always) The second is unconditional. The row reduction's GEMM is a persistent whole-device grid — `num_ctas = occupancy x SMs`, cluster-aligned, with cluster sync and DSMEM multicast — so two concurrent reductions each demand the entire GPU and neither can be placed. On Hopper that surfaces as a bare CUDA_ERROR_LAUNCH_FAILED, which compute-sanitizer does not attribute (0 invalid accesses across a whole run: it was never a memory bug). This is why putting the row reduction on its own GPU did not help on its own: `FP_CUDA_DEVICE=3` with the multiply on 0..2 turned arbitration off wholesale, so reductions stopped serializing against each other and the run still failed 63 times in 300 s. Isolating the device removes multiply contention and leaves reduction-vs-reduction contention untouched. With the split, an isolated reduction GPU runs clean at the FULL GEMM grid — no CTA cap, no throughput sacrificed. Measured on the theta=125 stem-200 repro that faulted at ~105 s in every other configuration: 400 s, 0 launch failures, 0 panics. The alternative was capping the persistent grid (`FP_CUDA_GEMM_CTAS`). The sweep on bench_kernel_only (16384^3, idle H200) shows why that is the wrong trade: the largest cap that survives the workload is 32, and 32 CTAs is 2107 TOPS against 8674 at full grid — 24% of peak. 64 CTAs reaches 49% but still fails. Throughput is linear in CTA count to ~128 (97% of peak), so the kernel only needs the device it is not being given. Not fixed here: the multiply takes `shared()` inside the SUBMISSION closure and drops it when submission returns, while its kernels are still resident — so on a SHARED device the yield is ineffective (`lock=0.0s` in every run). Correct fix is a dedicated fp-cuda driver thread that fences between reductions and overlaps transfers (copy engines do not consume SMs, so uploads can pipeline against compute). Isolating the reduction GPU sidesteps it. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/fp/src/blas/cuda.rs | 9 +++++++-- ext/crates/fp/src/gpu_lock.rs | 12 +++++++++--- 2 files changed, 16 insertions(+), 5 deletions(-) diff --git a/ext/crates/fp/src/blas/cuda.rs b/ext/crates/fp/src/blas/cuda.rs index a6cf082e14..628bba14d4 100644 --- a/ext/crates/fp/src/blas/cuda.rs +++ b/ext/crates/fp/src/blas/cuda.rs @@ -83,8 +83,13 @@ fn context() -> Option<&'static GpuContext> { // "arbitration off" from "arbitration on but uncontended". eprintln!( "[fp-cuda] row reduction on device {device}; multiply spans devices \ - 0..{mult_devices}; gpu_lock arbitration {}", - if shared { "ENABLED" } else { "disabled" } + 0..{mult_devices}; multiply yields to reductions: {}; reductions serialize against \ + each other: always", + if shared { + "yes" + } else { + "no (separate devices)" + } ); GpuContext::new(device).ok() }) diff --git a/ext/crates/fp/src/gpu_lock.rs b/ext/crates/fp/src/gpu_lock.rs index fa51767d1d..03544bf852 100644 --- a/ext/crates/fp/src/gpu_lock.rs +++ b/ext/crates/fp/src/gpu_lock.rs @@ -140,9 +140,15 @@ pub fn shared() -> SharedGuard { /// 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 { - if !arbitration_needed() { - return 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; From 85fa53c868c0b82fe1b8751899ae1ebd522daadf Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Thu, 6 Aug 2026 15:21:35 -0400 Subject: [PATCH 097/127] fp: route every fp-cuda submission through one driver thread MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 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. `gpu_lock::exclusive()` covered `row_reduce` only; `try_mul` was deliberately lock-free ("concurrent callers do not interfere"), so even a dedicated reduction GPU had two independent whole-device consumers and was not actually owned by anything. That is why enabling the cooperative reduction path on an isolated device still wedged after a single reduction: the cooperative kernel spun at its grid-wide barrier for CTAs a concurrent `try_mul` was holding. Routing both through one thread makes single-ownership structural instead of a discipline each new call site must remember. Jobs run to completion there, and both end in a device-to-host download, so serialization is on COMPLETION, not submission — the distinction that matters, and the one `gpu_lock::shared()` still gets wrong on the multiply side (taken inside the submit closure, dropped when submission returns, which is why `[batch-stats] lock=` reads 0.0s everywhere). Measured on the theta=125 stem-200 repro, reduction isolated to device 3, FULL GEMM grid: 200 s, 120 GPU reductions, 0 launch failures, both with and without FP_CUDA_RR_COOP. Default shared-device config unchanged: 147 reductions, 0 failures, max_t=245 in 200 s. Cooperative mode is now SAFE but not a win at this workload (closed=20406 vs 20306 in 200 s — noise). The reduction path is too small a share of the resolution for its ~2x to show. It stays off by default. Two things this does NOT do. The driver takes no `gpu_lock` guard: serialization among fp-cuda jobs is structural, and taking the guard there deadlocked the run (it waits for the multiply's readers while workers block on the driver), so yielding to the multiply on a SHARED device still needs arranging without a guard held across a blocking job. Transfers are serialized with compute, though copy engines do not consume SMs, so a later change can pipeline the next job's upload against the current job's kernels. `contended_acquisition_terminates_and_writers_are_exclusive` now measures reader overlap in its own uncontended phase: with `exclusive` unconditional, the contended phase keeps a writer queued almost always and writer preference correctly holds readers off, so asserting overlap there measured the scheduler rather than the lock. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/fp/src/blas/cuda.rs | 94 +++++++++++++++++++++++++++++++--- ext/crates/fp/src/gpu_lock.rs | 30 +++++++++-- 2 files changed, 113 insertions(+), 11 deletions(-) diff --git a/ext/crates/fp/src/blas/cuda.rs b/ext/crates/fp/src/blas/cuda.rs index 628bba14d4..a02253a42b 100644 --- a/ext/crates/fp/src/blas/cuda.rs +++ b/ext/crates/fp/src/blas/cuda.rs @@ -133,6 +133,81 @@ fn multiply_devices() -> usize { .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. @@ -165,9 +240,12 @@ pub(super) fn try_mul(a: &Matrix, b: &Matrix) -> Option { let a_limbs = to_limbs(a); let b_limbs = to_limbs(b); - // Lock-free: `matmul_b1_raw` submits on the calling thread's own stream with per-call device - // buffers, so concurrent callers do not interfere (see [`context`]). - let c = fp_cuda::matmul_b1_raw(ctx, &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)) } @@ -202,14 +280,16 @@ pub(crate) fn try_row_reduce(m: &mut Matrix) -> Option { // 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). - let _exclusive = crate::gpu_lock::exclusive(); - let (dev_limbs, perm, r, pivot_cols) = { + // 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()?; - (dev_limbs, perm, r, pivot_cols) - }; + 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 index 03544bf852..7df9d94324 100644 --- a/ext/crates/fp/src/gpu_lock.rs +++ b/ext/crates/fp/src/gpu_lock.rs @@ -196,6 +196,32 @@ mod tests { 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) = ( @@ -237,9 +263,5 @@ mod tests { 0, "two reductions held the device at once" ); - assert!( - max_shared.load(Ordering::SeqCst) > 1, - "multiplies never overlapped — the shared side is serialising, which defeats the point" - ); } } From 527a18d32c55ee1af473d5b7762d11e97731bf43 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Thu, 6 Aug 2026 16:01:23 -0400 Subject: [PATCH 098/127] fp-cuda: take a share of the device for the GEMM, not all of it MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The row reduction's GEMM sized its persistent grid to full occupancy (`occupancy x SMs`, ~264 CTAs on an H200). That grid is only placeable on a GPU this process owns outright. When anything else holds SMs 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). Arbitrating with the other runtime would make fp-cuda and algebra depend on each other's scheduling, and would still not help a co-tenant in another process. Asking for a share needs no such knowledge: the persistent loop already handles any multiple of CLUSTER (fewer CTAs do more tile-iterations each), so grid size was never a correctness parameter. Default 1/16. The safe size is NOT a sharp threshold: on the theta=125 stem-200 workload 1/16 ran a clean 240 s while 1/8 failed 74 times, and a 32-CTA cap — essentially the same grid as 1/8 — had passed cleanly in an earlier run under the same nominal config. Where the launch stops being placeable depends on what the other tenant is doing at that moment. A share sharply reduces collision probability; it does not prove it to zero. Cost, from the idle-device sweep (bench_kernel_only, 16384^3): 16 CTAs = 1062 binary TOPS, 32 = 2107, full grid = 8674. Set FP_CUDA_GEMM_DEVICE_FRAC=1 on a dedicated linear-algebra GPU to restore full throughput. Left unresolved, and recorded in the source: why an oversubscribed grid fails rather than queueing. Ordinary launches schedule in waves; something here — thread-block clusters, the dynamic shared-memory request, or both — makes placement a hard launch-time requirement. Until that is understood no share can be called correct, only likelier to fit. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/fp-cuda/src/lib.rs | 39 ++++++++++++++++++++++++++++++++++- 1 file changed, 38 insertions(+), 1 deletion(-) diff --git a/ext/crates/fp-cuda/src/lib.rs b/ext/crates/fp-cuda/src/lib.rs index d671cd691a..6612416114 100644 --- a/ext/crates/fp-cuda/src/lib.rs +++ b/ext/crates/fp-cuda/src/lib.rs @@ -423,7 +423,44 @@ fn run_gemm_kernel( .kernel .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; + // Take a SHARE of the machine, not all of it. 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). + // + // Fixing that by arbitrating with the other runtime would make `fp-cuda` and `algebra` depend on + // each other's scheduling, which is exactly the coupling worth avoiding — and it would not help + // a co-tenant outside this process at all. Asking only for a share needs no such knowledge: the + // persistent loop already handles any multiple of `CLUSTER` (fewer CTAs simply do more + // tile-iterations each), so grid size was never a correctness parameter, only a throughput one. + // + // The default is empirical and deliberately conservative, because the safe grid size is NOT a + // sharp threshold. On the theta=125 stem-200 workload (the multiply resident nearly + // continuously, full grid ~264 CTAs): 1/16 ran a clean 240 s, while 1/8 failed 74 times — and a + // 32-CTA cap, essentially the same grid as 1/8, had passed cleanly in an earlier run. Where the + // launch stops being placeable depends on what the other tenant is doing at that moment, so a + // share reduces the collision probability sharply but does not prove it to zero. Treat any + // value here as a risk setting, not a guarantee. + // + // The cost is real: the idle-device sweep (`bench_kernel_only`, 16384^3) puts 16 CTAs at 1062 + // binary TOPS and 32 at 2107, against 8674 at full grid. That is the price of composing on a + // shared GPU — but it is measured against a full grid that *fails* on such a GPU, not one that + // succeeds. + // + // Set `FP_CUDA_GEMM_DEVICE_FRAC=1` to take the whole machine when this process owns the GPU (a + // dedicated linear-algebra device), which restores full throughput. + // + // UNRESOLVED: why an oversubscribed grid *fails* rather than queueing. Ordinary launches + // schedule in waves; something here (thread-block clusters, the dynamic shared-memory request, + // or both) makes the placement a hard launch-time requirement. Until that is understood, no + // choice of share can be called correct — only likelier to fit. + let frac = std::env::var("FP_CUDA_GEMM_DEVICE_FRAC") + .ok() + .and_then(|v| v.parse::().ok()) + .filter(|&f| f > 0) + .unwrap_or(16); + let mut num_ctas = ((occ * sms / frac) / CLUSTER as u32).max(1) * CLUSTER as u32; // 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 From 9bc3012bd888053d960b216a8798b7611c668856 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Thu, 6 Aug 2026 16:05:46 -0400 Subject: [PATCH 099/127] fp-cuda: name Phase 9 clusters as the co-scheduling constraint to remove MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replaces the "UNRESOLVED: why does an oversubscribed grid fail rather than queue" note with the answer, found in the GEMM's own history. Phase 9 (264b4111d4, thread-block clusters + TMA multicast of B) introduced `__cluster_dims__`, B multicast with an all-ranks mask, and a cluster-wide empty barrier reached via `mapa`. A cluster's CTAs are co-resident by construction — a hard placement constraint an ordinary launch does not carry, and the reason the launch fails instead of scheduling in waves. Phase 8 (60b05cddc4) is not implicated, which is the useful half: its persistent loop strides (`tile += gridDim.x`) and was already work-capped (`sms.min(total_tiles)`), so the grouped rasterization that gives the kernel its L2 behaviour composes at any grid size. Only Phase 9's cluster layer has to go. That makes the fix the same one the row reduction already took (see `rr_coop`): a cluster-free variant, composable by construction, needing no knowledge of any other tenant — in this process or another. Cost is Phase 9's 2x cut in B HBM traffic, on top of the 8x GROUP_M rasterization already delivers, against the ~8x this share currently sacrifices (1062 TOPS at 1/16 vs 8674 at full grid). Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/fp-cuda/src/lib.rs | 18 ++++++++++++++---- 1 file changed, 14 insertions(+), 4 deletions(-) diff --git a/ext/crates/fp-cuda/src/lib.rs b/ext/crates/fp-cuda/src/lib.rs index 6612416114..9e557333d3 100644 --- a/ext/crates/fp-cuda/src/lib.rs +++ b/ext/crates/fp-cuda/src/lib.rs @@ -451,10 +451,20 @@ fn run_gemm_kernel( // Set `FP_CUDA_GEMM_DEVICE_FRAC=1` to take the whole machine when this process owns the GPU (a // dedicated linear-algebra device), which restores full throughput. // - // UNRESOLVED: why an oversubscribed grid *fails* rather than queueing. Ordinary launches - // schedule in waves; something here (thread-block clusters, the dynamic shared-memory request, - // or both) makes the placement a hard launch-time requirement. Until that is understood, no - // choice of share can be called correct — only likelier to fit. + // This whole knob is a STOPGAP. The co-scheduling requirement came in with Phase 9 + // (`264b4111d4`, thread-block clusters + TMA multicast of B): `__cluster_dims__` makes a + // cluster's CTAs co-resident by construction, B is multicast with an all-ranks mask, and the + // empty barrier is cluster-wide via `mapa`. That is a hard placement constraint, and it is why + // an oversubscribed grid fails rather than queueing the way an ordinary launch would. + // + // Phase 8 (`60b05cddc4`) is NOT implicated: its persistent loop strides (`tile += gridDim.x`) + // and was already work-capped (`sms.min(total_tiles)`), so it composes at any grid size. + // + // The fix is a cluster-free GEMM variant — the same move that made the row reduction composable + // (see `rr_coop`). The cost is Phase 9's 2x cut in B's HBM traffic, on top of the 8x that Phase + // 8's GROUP_M rasterization already delivers, against the ~8x throughput this share sacrifices + // (1062 TOPS at 1/16 vs 8674 at full grid). On a bandwidth-bound kernel that trade needs + // measuring, but it starts from a much better place than shrinking the grid. let frac = std::env::var("FP_CUDA_GEMM_DEVICE_FRAC") .ok() .and_then(|v| v.parse::().ok()) From d3bc6561003fd26c08a7c6b1237ba57f21469bc6 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Thu, 6 Aug 2026 16:27:05 -0400 Subject: [PATCH 100/127] fp-cuda: cluster-free GEMM as the default, cluster kernel opt-in The GEMM's `__cluster_dims__(CLUSTER,1,1)` makes a cluster's CTAs co-resident by construction -- rank 0 multicasts B straight into its mates' shared memory and every consumer arrives on their empty barriers through `mapa`. So the launch does not queue for SMs the way an ordinary grid does; when another runtime holds them it fails outright with CUDA_ERROR_LAUNCH_FAILED. That is the fault compute- sanitizer could never attribute (0 invalid accesses across a whole run: it was never a memory bug), and the reason `FP_CUDA_GEMM_DEVICE_FRAC` had to give away ~8x of the throughput to make the resolution survive a shared GPU. Template the kernel body on the cluster width and emit two entry points. Every cluster-dependent construct has an exact single-CTA counterpart -- rank is 0, cluster_sync drops (the preceding __syncthreads already orders it), arrive_cluster becomes a local arrive, multicast becomes a plain TMA load, the empty barrier counts 1 -- so the two kernels run identical arithmetic on an identical tile schedule and differ only in B's HBM traffic and in whether the launch demands co-resident CTAs. The cost of dropping multicast turns out to be nearly nil. bench_kernel_only on an idle H200, cluster-free vs cluster, correct=true idempotent=true throughout: 4096^3 4091 vs 4071 binary TOPS 8192^3 6687 vs 6778 16384^3 8501 vs 8632 32768^3 9608 vs 9664 Within 1.5% everywhere. GROUP_M rasterization was already keeping B resident in L2, so the second-order saving multicast adds does not show up at these shapes -- and it was never worth a hard placement constraint. So the cluster-free kernel takes the whole machine by default (frac 1) instead of a 1/16 share, which is a ~8x throughput gain over the stopgap on a shared GPU. `FP_CUDA_GEMM_COOP=1` opts back into the cluster kernel for a dedicated device, mirroring `FP_CUDA_RR_COOP` one layer down. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/fp-cuda/cuda_kernels/matmul_b1.cu | 124 +++++++++++++++---- ext/crates/fp-cuda/src/lib.rs | 102 ++++++++------- 2 files changed, 157 insertions(+), 69 deletions(-) diff --git a/ext/crates/fp-cuda/cuda_kernels/matmul_b1.cu b/ext/crates/fp-cuda/cuda_kernels/matmul_b1.cu index 8dd75c2e5d..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 diff --git a/ext/crates/fp-cuda/src/lib.rs b/ext/crates/fp-cuda/src/lib.rs index 9e557333d3..82f6ef7cbe 100644 --- a/ext/crates/fp-cuda/src/lib.rs +++ b/ext/crates/fp-cuda/src/lib.rs @@ -69,6 +69,28 @@ fn rr_coop() -> bool { .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. @@ -101,6 +123,8 @@ 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, @@ -134,6 +158,7 @@ impl GpuContext { 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")?; @@ -157,6 +182,7 @@ impl GpuContext { ctx, module, kernel, + kernel_nc, pack_a, pack_b, xor_into, @@ -394,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, )?; @@ -419,58 +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); - // Take a SHARE of the machine, not all of it. 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). - // - // Fixing that by arbitrating with the other runtime would make `fp-cuda` and `algebra` depend on - // each other's scheduling, which is exactly the coupling worth avoiding — and it would not help - // a co-tenant outside this process at all. Asking only for a share needs no such knowledge: the - // persistent loop already handles any multiple of `CLUSTER` (fewer CTAs simply do more - // tile-iterations each), so grid size was never a correctness parameter, only a throughput one. - // - // The default is empirical and deliberately conservative, because the safe grid size is NOT a - // sharp threshold. On the theta=125 stem-200 workload (the multiply resident nearly - // continuously, full grid ~264 CTAs): 1/16 ran a clean 240 s, while 1/8 failed 74 times — and a - // 32-CTA cap, essentially the same grid as 1/8, had passed cleanly in an earlier run. Where the - // launch stops being placeable depends on what the other tenant is doing at that moment, so a - // share reduces the collision probability sharply but does not prove it to zero. Treat any - // value here as a risk setting, not a guarantee. - // - // The cost is real: the idle-device sweep (`bench_kernel_only`, 16384^3) puts 16 CTAs at 1062 - // binary TOPS and 32 at 2107, against 8674 at full grid. That is the price of composing on a - // shared GPU — but it is measured against a full grid that *fails* on such a GPU, not one that - // succeeds. - // - // Set `FP_CUDA_GEMM_DEVICE_FRAC=1` to take the whole machine when this process owns the GPU (a - // dedicated linear-algebra device), which restores full throughput. - // - // This whole knob is a STOPGAP. The co-scheduling requirement came in with Phase 9 - // (`264b4111d4`, thread-block clusters + TMA multicast of B): `__cluster_dims__` makes a - // cluster's CTAs co-resident by construction, B is multicast with an all-ranks mask, and the - // empty barrier is cluster-wide via `mapa`. That is a hard placement constraint, and it is why - // an oversubscribed grid fails rather than queueing the way an ordinary launch would. + // 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). // - // Phase 8 (`60b05cddc4`) is NOT implicated: its persistent loop strides (`tile += gridDim.x`) - // and was already work-capped (`sms.min(total_tiles)`), so it composes at any grid size. + // 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. // - // The fix is a cluster-free GEMM variant — the same move that made the row reduction composable - // (see `rr_coop`). The cost is Phase 9's 2x cut in B's HBM traffic, on top of the 8x that Phase - // 8's GROUP_M rasterization already delivers, against the ~8x throughput this share sacrifices - // (1062 TOPS at 1/16 vs 8674 at full grid). On a bandwidth-bound kernel that trade needs - // measuring, but it starts from a much better place than shrinking the grid. + // `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(16); - let mut num_ctas = ((occ * sms / frac) / CLUSTER as u32).max(1) * CLUSTER as u32; + .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 @@ -479,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); @@ -496,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) From ec3ed4550ca00e409e95edb17c5dee66ce9f5dbc Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Thu, 6 Aug 2026 19:24:18 -0400 Subject: [PATCH 101/127] algebra: measure the d-vs-masks enumeration shortcut, and record that it loses Both odometers -- `AdmissibleMatrix::next` and `enumerate_admissible_kernel` -- recompute an anti-diagonal bitsum `d` for every visited (row, col), and the same loop already maintains `masks[row+col] = d | new_entry`. That looks like a precomputation waiting to be exploited: collapse an O(cols) scan in the hottest triple-nested loop down to one load, in both implementations at once. Measured over 46,995,344 anti-diagonal computations (degree <= 120): d == masks 13,322,067 28.348% d subset-of masks 46,995,344 100.000% masks==0 (sound skip) 2,736,685 5.823% d-scan iterations 55,846,521 -> 1.19 per check The containment is exact and universal, and it does yield a sound skip (`masks[row+col] == 0` implies `d == 0`). But the scan it would skip averages 1.19 iterations: `(row+col+1).saturating_sub(rows)..col` is nearly always empty or a single step for real R's, because they are short and wide relative to the anti-diagonal. The skip fires on 5.8% of checks and avoids 6.9% of iterations of a loop that barely runs. There is nothing here. Keep the probe as an ignored diagnostic with the verdict in its docstring, so the shortcut is not re-proposed. Its assert is armed: 0 containment violations. With the earlier null result from cutting per-thread local state 3x (1.884s -> 1.872s), the odometer's arithmetic and its local state are both cleared. What is left proportional to the work is the emit: ~20 scattered 2-byte global stores per matrix, ~680 GB at degree 240. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 190 +++++++++++++++++++ 1 file changed, 190 insertions(+) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index c843702d1d..3b7e519f91 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -4931,6 +4931,196 @@ mod tests { ); } + /// 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 From f3391b76fe586117f4df883882d02cc396780f97 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Thu, 6 Aug 2026 22:16:25 -0400 Subject: [PATCH 102/127] =?UTF-8?q?algebra:=20measure=20the=20enum=20kerne?= =?UTF-8?q?l's=20emit=20cost=20=E2=80=94=2042.8%=20of=20it,=20at=201%=20of?= =?UTF-8?q?=20HBM=20bandwidth?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two null results had already 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 (ec3ed4550c). That left the emit as the only candidate proportional to the work -- `cs_len + mk_len` (~20) scattered 2-byte global stores per matrix, with adjacent lanes writing to unrelated `r_cs_out[ri]` offsets. Add a comptime `emit` flag so the stores compile OUT rather than branching around them, and time the identical odometer with and without. `mat` still increments, so the enumeration and its trip count are unchanged; only the writes differ. degree <= 130: 89,392 R's, 75,987,575 matrices, 1.03e9 u16 stores (2.07 GB) emit=true 0.0858s emit=false 0.0491s stores: 42.8% of kernel time, 1.75x if free Large, but NOT dominant -- the odometer is still 57%. The recoverable part is what the bandwidth says: 2.07 GB in 0.0367s is ~56 GB/s against ~4.8 TB/s of HBM, about 1% of what the device can do. That is the signature of scattered 2-byte transactions, not of a store volume we are stuck with. Packing u16 pairs into u32 stores, or staging a matrix through shared memory for coalesced warp writes, should recover most of the 42.8% without touching the format the multiply reads. Production launches pass `emit = true`; only the diagnostic passes false. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 173 ++++++++++++++++++- 1 file changed, 166 insertions(+), 7 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 3b7e519f91..e7a4c3adf7 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -3247,6 +3247,7 @@ fn multiply_batch_block<'a>( BufferArg::from_raw_parts(cnt_scratch, n_s), enum_width, n_s, + true, ); } } @@ -3949,6 +3950,10 @@ fn enumerate_admissible_kernel( 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`. + #[comptime] emit: bool, ) { let ri = ABSOLUTE_POS; if ri >= n_r { @@ -3997,13 +4002,15 @@ fn enumerate_admissible_kernel( let mut more = true; while more { // Emit the current matrix's col_sums/masks into this R's scratch slot. - let co = cs_base + mat * cs_len; - for j in 0..cs_len { - out_cs[co + j] = u16::cast_from(col_sums[j]); - } - let mo = mk_base + mat * mk_len; - for j in 0..mk_len { - out_mk[mo + j] = u16::cast_from(masks[j]); + if emit { + let co = cs_base + mat * cs_len; + for j in 0..cs_len { + out_cs[co + j] = u16::cast_from(col_sums[j]); + } + let mo = mk_base + mat * mk_len; + for j in 0..mk_len { + out_mk[mo + j] = u16::cast_from(masks[j]); + } } mat += 1; @@ -4480,6 +4487,7 @@ mod tests { BufferArg::from_raw_parts(cnt_h.clone(), n_r), width, n_r, + true, ); } // Truncate off the `max(1)` padding element present when a batch has zero col_sums / masks (an @@ -4492,6 +4500,156 @@ mod tests { (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 (cs, mk, t_emit) = + time_enum::(&Default::default(), &p_parts, &num_mats, true); + let (_, _, t_noemit) = + time_enum::(&Default::default(), &p_parts, &num_mats, false); + let stores = cs as u64 + mk as u64; + 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: bool, + ) -> (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::()); + + const THREADS: u32 = 64; + 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), + 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, + 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 @@ -4801,6 +4959,7 @@ mod tests { BufferArg::from_raw_parts(cnt_h.clone(), n_r), width, n_r, + true, ); } // Reading the tiny counts buffer blocks until the kernel completes: kernel wall time, ~no transfer. From e51670c05d60344a42b31c7c15c72ac1aef1508a Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Thu, 6 Aug 2026 22:41:23 -0400 Subject: [PATCH 103/127] algebra: record that stack-allocating AdmissibleMatrix loses Replacing the four `Vec`s in `AdmissibleMatrix` with fixed arrays sized to `PPart`'s structural caps (MAX_LEN = 10 rows, width(0) = 11 cols) removes four allocations per R, of which there are ~1.6e6 at degree 240. It looked like a free win. Measured on benches/nassau_milnor it is a net loss: op24xel32 6.7% slower op40xel1 2.2% slower op8xel24, op20xel24, op32xel8, op8xel8 0.3-0.6% slower op24xel24, op16xel32, op46xel1 0.1-0.3% faster rest no significant change The allocations were never the cost; the zeroing is. A typical R is ~4 rows x ~6 cols = 24 entries against a 110-entry cap, so `vec![0; rows * cols]` clears about a third of what `[0; 110]` does, and the saved mallocs do not pay for the extra stores. An end-to-end stem-110 CPU resolution agreed: 9.59s vs 9.71s, noise. Same effect as shrinking the GPU kernel's ENUM_COL_CAP 32 -> 11, from the other direction: over-sized fixed state costs more than the allocation it avoids. Noted on the struct so the change is not re-proposed. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_algebra.rs | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/ext/crates/algebra/src/algebra/milnor_algebra.rs b/ext/crates/algebra/src/algebra/milnor_algebra.rs index 953c877e23..a3abf12de4 100644 --- a/ext/crates/algebra/src/algebra/milnor_algebra.rs +++ b/ext/crates/algebra/src/algebra/milnor_algebra.rs @@ -1755,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, From 1a7e0f9c42bc24c4a8c7c9c56d9e82eda986bfcd Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Thu, 6 Aug 2026 22:49:26 -0400 Subject: [PATCH 104/127] algebra: the enum emit is store-ISSUE bound, not bandwidth bound f3391b76fe measured the emit at 42.8% of the enum kernel and read its ~56 GB/s as ~1% of HBM -- implying a bandwidth pathology from scattered 2-byte writes, to be fixed by relayout. Two more probe modes say otherwise. emit=1 production 0.0858s emit=0 no stores 0.0490s emit=2 lane-adjacent (coalesced) 0.0829s 92.1% of the cost REMAINS emit=3 every other entry 0.0685s 53.1% remains Mode 2 writes the same bytes and the same number of stores with the warp's lanes hitting adjacent addresses -- perfect coalescing, garbage layout, timing probe only. If transactions or bandwidth were the constraint that would have collapsed the cost; it moved 8%. Mode 3 halves the store count and takes almost exactly half the cost. The emit is bound by store-instruction issue. So relayout is dead and packing is the fix, with width tracking the win directly since cost is linear in instruction count: 2 x u16 -> u32 half the stores ~21% of kernel time ~1.27x 4 x u16 -> u64 quarter the stores ~32% ~1.47x Same bytes moved either way; what is bought is issue slots. Requires padding each R's cs/mk stride to the packing width so a matrix's run starts word-aligned, and teaching the multiply kernel's reader to unpack. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 66 +++++++++++++++++--- 1 file changed, 56 insertions(+), 10 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index e7a4c3adf7..40de52999e 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -3247,7 +3247,8 @@ fn multiply_batch_block<'a>( BufferArg::from_raw_parts(cnt_scratch, n_s), enum_width, n_s, - true, + 0u32, + 1, ); } } @@ -3953,7 +3954,10 @@ fn enumerate_admissible_kernel( // 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`. - #[comptime] emit: bool, + // 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 { @@ -4002,7 +4006,11 @@ fn enumerate_admissible_kernel( let mut more = true; while more { // Emit the current matrix's col_sums/masks into this R's scratch slot. - if emit { + // 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(col_sums[j]); @@ -4011,6 +4019,27 @@ fn enumerate_admissible_kernel( for j in 0..mk_len { out_mk[mo + j] = u16::cast_from(masks[j]); } + } 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(col_sums[j]); + } + for j in 0..mk_len { + out_mk[(ri + n_r * (mat * mk_len + j)) & w] = u16::cast_from(masks[j]); + } + } 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(col_sums[j]); + j += 2; + } + let mo = mk_base + mat * mk_len; + let mut j2 = 0usize; + while j2 < mk_len { + out_mk[mo + j2] = u16::cast_from(masks[j2]); + j2 += 2; + } } mat += 1; @@ -4487,7 +4516,8 @@ mod tests { BufferArg::from_raw_parts(cnt_h.clone(), n_r), width, n_r, - true, + 0u32, + 1, ); } // Truncate off the `max(1)` padding element present when a batch has zero col_sums / masks (an @@ -4553,11 +4583,19 @@ mod tests { .collect(); let total_mats: u64 = num_mats.iter().map(|&m| m as u64).sum(); - let (cs, mk, t_emit) = - time_enum::(&Default::default(), &p_parts, &num_mats, true); - let (_, _, t_noemit) = - time_enum::(&Default::default(), &p_parts, &num_mats, false); + 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}% \ @@ -4578,7 +4616,7 @@ mod tests { device: &R::Device, p_parts: &[Vec], num_mats: &[u32], - emit: bool, + emit: u32, ) -> (u64, u64, f64) { use std::time::Instant; @@ -4620,6 +4658,12 @@ mod tests { 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 = 64; let cubes = (n_r as u32).div_ceil(THREADS); let mut times: Vec = Vec::new(); @@ -4640,6 +4684,7 @@ mod tests { BufferArg::from_raw_parts(cnt_h.clone(), n_r), width, n_r, + wrap, emit, ); } @@ -4959,7 +5004,8 @@ mod tests { BufferArg::from_raw_parts(cnt_h.clone(), n_r), width, n_r, - true, + 0u32, + 1, ); } // Reading the tiny counts buffer blocks until the kernel completes: kernel wall time, ~no transfer. From ece4f45f91125d045482715e12d36543c58fef53 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Thu, 6 Aug 2026 23:54:48 -0400 Subject: [PATCH 105/127] =?UTF-8?q?algebra:=20measure=20R=20rebuild=20COST?= =?UTF-8?q?,=20not=20reference=20count=20=E2=80=94=20the=20transient=20hea?= =?UTF-8?q?d=20is=20steep?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `NASSAU_R_STATS` ranked Rs by reference count, and the eviction policy was tuned on that plus bytes. But enumeration cost is proportional to num_mats, so what a rebuild actually costs is `count * num_mats`, and a small hot R can cost the same as a big cold one. Record num_mats per R and report the cost-weighted concentration, plus the transient set on its own. S_2 (150,75), theta=125: total matrices enumerated (if every ref rebuilt) 385.5e9 transient (deg>125) 98,551 Rs (57%) 43.5M refs (3%) 22% of cost within transient, cost coverage top1%=21% top5%=51% top10%=68% top25%=88% Two things follow. Transient is 22% of total enumeration work but ~99% of GPU KERNEL time (nsys, transient-heavy config: enumerate_admissible_kernel 98.9%, multiply_batch_kernel 1.1%). No contradiction: resident Rs are enumerated once on the CPU at first sight and then live in the master, so they never cost GPU time again. Rebuild work and GPU time are different denominators and were being conflated. And the cost inside the transient set is NOT flat -- 5% of its Rs carry half of it. A partial cache has a head to exploit. Better, since bytes are ~num_mats and cost is count*num_mats, cost-per-byte is EXACTLY count: num_mats cancels, so admission ranked on plain reference count is optimal for any byte budget. The degree cutoff remains the right first filter (degree proxies bytes) but discards hot and cold alike among what it excludes; a small count-ranked cache layered on top of theta is the missing piece, with memory bounded by what it is given rather than by theta. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 68 +++++++++++++++++++- 1 file changed, 66 insertions(+), 2 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 40de52999e..18da1ba54f 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -1035,6 +1035,11 @@ struct RStat { 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>>> = @@ -1109,8 +1114,12 @@ enum MasterMode { Transient, } -fn record_r_use(p_part: PPart) { +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 { @@ -1118,6 +1127,7 @@ fn record_r_use(p_part: PPart) { degree: ppart_degree(p_part), first: now, last: now, + num_mats, }); e.count += 1; e.last = now; @@ -1235,10 +1245,64 @@ pub fn dump_r_stats() { } } 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). + 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), + ); } fn resident_info(algebra: &MilnorAlgebra, p_part: PPart) -> RInfo { - record_r_use(p_part); + record_r_use(algebra, p_part); if let Some(info) = RESIDENT_HOST.read().unwrap().index.get(&p_part) { return *info; } From 00f8d9051ccd1034f1ec8d125df768e399f3c009 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Fri, 7 Aug 2026 00:14:02 -0400 Subject: [PATCH 106/127] =?UTF-8?q?algebra:=20simulate=20a=20pinned=20tran?= =?UTF-8?q?sient=20cache=20=E2=80=94=20it=20is=20linear,=20so=20theta=20is?= =?UTF-8?q?=20already=20optimal?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The previous commit reported that transient cost is concentrated (top 5% of Rs carry 51%) and proposed a count-ranked pinned cache. Add a per-R CSV dump and simulate that policy against the reference stream. It does not work, and the reason retires the idea. transient at (150,75): 98,551 Rs, 43.5M refs, 83.76e9 matrices, 8.2 GB total budget oracle k=1 k=2 k=4 k=8 k=16 1% 2.7% 2.4% 2.5% 2.5% 2.5% 2.6% 5% 11.8% 11.7% 11.7% 11.7% 11.7% 11.6% 10% 21.6% 21.4% 21.3% 21.3% 21.2% 21.2% 25% 44.0% 43.8% 43.8% 43.7% 43.5% 43.0% Savings are LINEAR in bytes cached, and every admission rule lands within 0.3pp of an oracle that ranks by cost-per-byte with full hindsight. When the oracle cannot beat arrival order there is no structure to exploit. The "steep head" was an artifact of measuring concentration per R rather than per BYTE. cost/byte is exactly `count`, and the high-cost Rs are the big ones, so 5% of Rs is nowhere near 5% of bytes. Normalised by the memory it costs -- the only thing a cache budget cares about -- the distribution is flat. So a pinned cache is the same linear memory-for-time dial theta already is, and adds nothing. What the numbers do show is that theta is not a parameter to tune but a fallback: the whole transient set is 8.2 GB at (150,75), and giving it up costs ~99% of GPU kernel time (2412s uncapped vs 18566s at theta=125 for stem 200). Theta should be as high as device memory allows, and the levers worth having are the ones that make it unnecessary -- a smaller master (values are <= 2^11 but stored as u16; bit-packing is ~31% fewer bytes) and faster enumeration for whatever still misses. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 22 ++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 18da1ba54f..c27c379aa6 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -1251,6 +1251,28 @@ pub fn dump_r_stats() { // 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))); From e7c8d2f1d9ff6e1688a5975c71f1e66d289317c8 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Fri, 7 Aug 2026 01:03:46 -0400 Subject: [PATCH 107/127] =?UTF-8?q?algebra:=20theta=3D125=20costs=206.5x?= =?UTF-8?q?=20wall=20time=20to=20save=205.3=20GB=20=E2=80=94=20set=20it=20?= =?UTF-8?q?as=20high=20as=20memory=20allows?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `resident_degree_cap`'s docstring justified the threshold on a byte metric: at theta<=125, 43% of distinct Rs stay resident and only 4% of references miss. That is true and it is the wrong axis. A miss is not a fixed cost -- it re-enumerates on the GPU in EVERY block touching that 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, against multiply_batch_kernel's 1.1% (nsys). Stem 200 to max_t=310, all complete, 0 crashes: theta=inf 2412s theta=200 2865s 50.2 GB peak theta=125 18566s 44.9 GB peak theta=125 gives up 6.5x in wall time for 5.3 GB. The knee is sharp and sits above 125; the memory curve is much flatter than the time curve. exec fell 16x (26269s -> 1619s) and fence 32x (245398s -> 7596s) on identical work (pairs 5.88e13 both) -- the enumeration simply stops happening. So the cap is a fallback for exhausting device memory, not a knob to tune down, and the docstring now says so with the numbers. Also records the negative result on smarter eviction: a pinned count-ranked cache saves linearly in bytes cached, within 0.3pp of a full-hindsight oracle, because cost/byte is exactly `count` and the per-byte distribution is flat. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 22 ++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index c27c379aa6..254adcd23d 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -1092,6 +1092,28 @@ fn ppart_degree(p_part: PPart) -> i32 { /// 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") From a3f46eb4f764995adbf694a0a959c5d170752bd1 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Fri, 7 Aug 2026 01:19:16 -0400 Subject: [PATCH 108/127] algebra: report which theta would have fit, instead of discovering it as an OOM Picking theta was a blind guess: master bytes are dominated by the top degrees (stem 150: 2.2 GB at theta<=125, 9.9 GB at theta<=150), so a lower-stem run cannot predict a higher theta -- the degrees carrying the mass are absent from it. The only way to learn theta was to set it too high and crash hours in. Track master bytes per R degree at admission (one add per first-sight R). The cumulative sum IS "how big the master would be at theta = d", so any run reports the whole curve, and a budget picks theta off it directly. Printed with the periodic batch-stats as well as at the end, so a run that dies on an allocation still leaves the answer behind. stem 140: full=4.1GB (1.0GB/GPU over 4 devices) theta<=100 0.3GB, <=125 1.5GB, <=141 4.1GB Which also corrects the premise this cap was reasoned about under. The master is a SMALL share of device memory -- 1.0 GB/GPU at stem 140, 2.6 GB/GPU at stem 150, against a ~50 GB peak at stem 200. theta caps the master, so it is not the knob that governs peak memory; the concurrent dense output matrices are. A theta that fits the master is necessary, not sufficient, and lowering theta buys far less memory than its 6.5x wall-time cost (e7c8d2f1d9) suggests it should. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 52 ++++++++++++++++++++ ext/src/nassau.rs | 1 + 2 files changed, 53 insertions(+) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 254adcd23d..0d1cdd8e37 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -668,6 +668,13 @@ struct ResidentHost { /// 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 @@ -684,6 +691,7 @@ static RESIDENT_HOST: LazyLock> = LazyLock::new(|| { 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(), }) }); @@ -1345,6 +1353,41 @@ pub fn dump_r_stats() { ); } +/// 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 + ); + } + } + 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) { @@ -1394,6 +1437,11 @@ fn resident_info(algebra: &MilnorAlgebra, p_part: PPart) -> RInfo { 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(); @@ -3807,6 +3855,10 @@ fn multiply_batch_block<'a>( let fence_s = BATCH_FENCE_US.load(std::sync::atomic::Ordering::Relaxed) as f64 / 1e6; 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}% \ diff --git a/ext/src/nassau.rs b/ext/src/nassau.rs index 80372e49f1..fd165e63d1 100644 --- a/ext/src/nassau.rs +++ b/ext/src/nassau.rs @@ -1447,6 +1447,7 @@ impl> Resolution { // Eviction probe (`NASSAU_R_STATS`): dump the R-access distribution once the wavefront is done. #[cfg(feature = "gpu")] algebra::milnor_gpu::dump_r_stats(); + algebra::milnor_gpu::dump_master_by_degree(); } } From a1b832be63e8a26f6c9891a156abb4ce3fe5dee4 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Fri, 7 Aug 2026 02:35:29 -0400 Subject: [PATCH 109/127] nassau: skip step1's kernel computation when the augmentation codomain is empty step1 computes `desired_image` as the kernel of the augmentation `target_module -> cc_module` in degree t, by building a `target_dim x target_dim` augmented identity and row-reducing it. When the target complex is empty in degree t that map has a zero-dimensional codomain, so the kernel is the whole space and no computation can discover otherwise. Take it directly. `target_dim` is dim(A_t), so at high t this allocated and reduced a matrix large enough to reach the GPU RREF path purely to rediscover the entire space. The guard is on the codomain being empty, NOT on resolving the sphere: every finite target module is concentrated in finitely many degrees, so past its top cell this holds for all t -- almost the whole resolution. The sphere is only the extreme of it, firing from t = 1. Verified bit-identical output on S_2, C2 and C2_eta (stem 40, max_s 20) -- i.e. on non-sphere modules, where the guard is false at low t and true above the top cell. NOT verified faster. At stem 120 the difference is noise (4.85s vs 4.91s): the waste is O(dim(A_t)^2) so it only bites at high t, and the high-t A/B was killed before finishing. Landed on correctness and on there being no reason to do the work, not on a measured win. Also gate the `dump_master_by_degree` call behind `feature = "gpu"` with its neighbour -- it broke the non-GPU build. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/src/nassau.rs | 40 +++++++++++++++++++++++++++++++--------- 1 file changed, 31 insertions(+), 9 deletions(-) diff --git a/ext/src/nassau.rs b/ext/src/nassau.rs index fd165e63d1..f458f91be8 100644 --- a/ext/src/nassau.rs +++ b/ext/src/nassau.rs @@ -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}, }; @@ -1117,12 +1117,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]); - 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(); + // 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(); + matrix.compute_kernel() + }; let mut matrix = AugmentedMatrix::<2>::new_with_capacity( p, @@ -1446,8 +1465,11 @@ impl> Resolution { // Eviction probe (`NASSAU_R_STATS`): dump the R-access distribution once the wavefront is done. #[cfg(feature = "gpu")] - algebra::milnor_gpu::dump_r_stats(); - algebra::milnor_gpu::dump_master_by_degree(); + { + algebra::milnor_gpu::dump_r_stats(); + // Which theta would have fit: see `resident_degree_cap`. + algebra::milnor_gpu::dump_master_by_degree(); + } } } From 3b83d1e68e8d899160d42811995998862dc5bc82 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Fri, 7 Aug 2026 02:35:43 -0400 Subject: [PATCH 110/127] algebra: abort on segment-table exhaustion instead of killing one shard thread seg_grow's MASTER_MAX_SEG check was an assert, and it runs on a shard's dedicated per-launch thread. A panic there kills only that launch: the next spawns a fresh thread and the resolution continues, having silently dropped the failed block's products. That is not hypothetical. On an uncapped stem-300 run all four shard threads hit it at b=(174,26) and the run carried on for another ~900k batches (batch-stats calls 828k -> 1740k) with GPU_DISABLED never set, because no fallback engaged -- nothing had failed from the resolution's point of view. The output would have been wrong, not slow, and nothing in the run's own reporting would have said so. There is no recovery from exhausting the segment table mid-master, so abort the process. The message now points at the actual fix -- lower theta, sized off `dump_master_by_degree` -- rather than only at raising the caps. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 26 ++++++++++++++------ 1 file changed, 19 insertions(+), 7 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 0d1cdd8e37..12db82841e 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -981,13 +981,25 @@ macro_rules! seg_grow { }; let (tail, new_len): (Vec<$elem>, usize) = ($tail)(uploaded); debug_assert_eq!(uploaded + tail.len(), new_len); - assert!( - new_len.div_ceil(seg_elems.max(1)) <= MASTER_MAX_SEG, - "resident buffer needs {} segments (> MASTER_MAX_SEG={}); raise \ - MASTER_MAX_SEG or NASSAU_GPU_MASTER_SEG_ELEMS", - new_len.div_ceil(seg_elems.max(1)), - MASTER_MAX_SEG - ); + // 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. From b5d927a950b78cca8a0a78c8a0524bd94906cb02 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Fri, 7 Aug 2026 11:58:49 -0400 Subject: [PATCH 111/127] =?UTF-8?q?algebra:=20profile=20the=20enum=20kerne?= =?UTF-8?q?l=20=E2=80=94=20it=20is=20latency-bound,=20and=20shared=20state?= =?UTF-8?q?=20does=20not=20fix=20it?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The emit ablations (f3391b76fe, 1a7e0f9c42) said stores are 42.8% of this kernel and that store COUNT matters while coalescing does not. Both true, and both misleading about the mechanism. ncu on an idle H200: Compute (SM) 5.16% Avg Active Threads/Warp 5.88 / 32 Memory 13.92% Active Warps/Scheduler 2.09 DRAM 0.65% No Eligible 76.57% L2 Hit 98.83% Warp Cycles/Issued Instr 8.90 Nothing is saturated. Stores hit cache (L2 98.8%) and never reach DRAM, which is why making them lane-adjacent bought 8%: there was no transaction pressure to relieve. The kernel is latency-bound, and ncu attributes 45.6% of its 8.9 stall cycles to L1TEX waits on the odometer's own state -- `matrix`/`totals`/ `col_sums`/`masks` are indexed by runtime values, so they live in local memory. Tried the obvious fix: state in `Shared<[u32]>`, strided `elem * BLOCK + tid` so a warp hits consecutive banks. Bit-exact (admissible_enum_gpu_matches passes), and the mechanism is confirmed -- warp cycles per issued instruction 8.90 -> 5.92, a 33% cut. But 152 u32/thread is 38.9 KB per block, which forces the block from 256 to 64 threads, and active warps per scheduler fall 2.09 -> 1.61. Net 109.63 ms -> 112.72 ms, a 2.8% LOSS. Reverted. Two things worth carrying forward. The 33% is real and only paywalled behind the shared footprint, so a smaller state (`matrix` as u16 -- values are <= 2^11) may still collect it. And divergence is untouched by any of this: 5.88 of 32 lanes active, ~82% of the machine idle, because one thread owns one R and warps run at their longest member. That is now the largest single number in the profile. Also: ncu's "Est. Local Speedup: 76.57%" is not a forecast. It is the prize if the stall vanishes for free; here it cost occupancy and the trade was negative. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 30 ++++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 12db82841e..9adebfede4 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -4085,6 +4085,36 @@ fn seg_read_u32( v } +/// 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. +/// /// 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 From 815fbdd41524b84a90b6b85aa1cf179aa7f58fae Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Fri, 7 Aug 2026 12:20:09 -0400 Subject: [PATCH 112/127] =?UTF-8?q?algebra:=20the=20enum=20kernel=20is=20w?= =?UTF-8?q?ork-starved=20=E2=80=94=20registers=20and=20shared=20memory=20a?= =?UTF-8?q?re=20not=20the=20constraint?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Follow-up to the latency profile. ncu Occupancy/LaunchStats on the production kernel (idle H200, degree <= 130): Registers Per Thread 40 Block Limit Registers 24 blocks Theoretical Occupancy 75% Waves Per SM 0.44 Achieved Occupancy 8.16% Grid Size 1397 x 64 threads Registers allow 24 blocks/SM and a 75% theoretical ceiling, so shrinking them buys nothing -- the achieved figure is 8.16%, an order of magnitude below a ceiling that is already generous. Waves Per SM = 0.44 is the fact that matters. The H200 has 132 SMs x 24 block slots = 3168; the grid supplies 1397. One thread per R over 89392 Rs is ~89k threads against ~270k of device capacity, so the kernel cannot fill half the machine however well it runs. Achieved occupancy then falls to 8.16% because those few blocks drain raggedly -- the same divergence that leaves 5.88 of 32 lanes active. Production is likely worse still: this benchmark enumerates every R to degree 130 in one launch, while a production launch covers only one block's transient Rs. So the granularity is the bug: thread-per-R gives too FEW threads and wildly UNEQUAL ones, and stores / local memory / shared memory / registers are all adjustments to a kernel that has no work to hide latency with. That explains why every attempt today measured between -3% and +8%. The changes that would matter are structural: batch many more Rs per launch so the grid fills the device, or replace thread-per-R with lane cooperation over dynamically pulled work. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 24 ++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 9adebfede4..3d01327b0c 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -4115,6 +4115,30 @@ fn seg_read_u32( /// 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. +/// /// 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 From 18e79d8b1340e6df6944882b38320637511c8cb0 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Fri, 7 Aug 2026 13:56:15 -0400 Subject: [PATCH 113/127] =?UTF-8?q?algebra:=20log=20enum-launch=20geometry?= =?UTF-8?q?=20=E2=80=94=20production=20runs=20the=20kernel=20at=200.2%=20o?= =?UTF-8?q?f=20the=20grid?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `enumerate_admissible_kernel` is ~99% of GPU kernel time, and profiling it standalone gave Waves Per SM = 0.44: it could not fill half the device. But that benchmark enumerates every R to degree 130 in ONE launch, whereas production launches cover a single block-segment's transient Rs. Count what production actually submits. S_2 stem 150, max_s 60, theta=125: enum launches 10260 Rs/launch mean 1293 (max 34157) blocks/launch mean 6 waves/SM 0.002 Six blocks against an H200's 3168 slots. Production is 220x smaller than the benchmark and uses 0.2% of the machine per launch, ~1293 threads out of ~270k of capacity, 10260 times over. That settles batching vs lane cooperation, which was the open question. Lane cooperation (warp per R instead of thread per R) fixes divergence -- 5.88 of 32 lanes active -- and multiplies threads by 32, but the odometer is sequential so it is Amdahl-capped near 1.6x, and the grids would still be a fraction of a percent of the device. Batching Rs across launches attacks a factor of hundreds. It also explains why this kernel dominates GPU time despite being simple integer work: it is not slow, it is run 10260 times on an almost empty GPU, so the cost is launch latency and drain. And it closes the micro-optimisation thread -- stores, local memory, shared memory, registers and coalescing all tune a kernel that is idle by construction. The 42.8% emit share was real, but it is 42.8% of a kernel at 0.2% grid utilisation. The counters ride with `[batch-stats]`, so any run reports its own geometry. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 32 ++++++++++++++++---- 1 file changed, 26 insertions(+), 6 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 3d01327b0c..04671340e3 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -497,6 +497,16 @@ 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); @@ -3396,14 +3406,15 @@ fn multiply_batch_block<'a>( 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( - (n_s as u32).div_ceil(ENUM_THREADS).max(1), - 1, - 1, - ), + 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), @@ -3866,6 +3877,8 @@ fn multiply_batch_block<'a>( 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 @@ -3877,7 +3890,9 @@ fn multiply_batch_block<'a>( 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", + 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, @@ -3892,6 +3907,11 @@ fn multiply_batch_block<'a>( 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, ); } From 6e5b88f68f90d83ca64d25e1c727f20c0e48e4fc Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Fri, 7 Aug 2026 14:39:05 -0400 Subject: [PATCH 114/127] =?UTF-8?q?algebra:=20block=20size=20is=20not=20th?= =?UTF-8?q?e=20enum=20lever=20either=20=E2=80=94=20the=20launches=20are=20?= =?UTF-8?q?serial?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Production submits ~6 blocks per enum launch (waves/SM 0.002), so the obvious move was to spread the same work over more SMs with smaller blocks. Measured on S_2 stem 150 / max_s 60 / theta=125: 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 the blocks for 3.4% — noise. Reverted. The reason is the useful part. A launch's duration is set by its LONGEST SINGLE R -- one thread, sequential odometer -- not by how many blocks it occupies, so extra SMs idle beside the one thread still grinding. Rs/launch averages 1293 and peaks at 34157, so the intra-launch spread is enormous. And 10260 launches across a ~550 s run is ~50 ms apiece, which matches `gpu_thread` running every device section on one thread and one stream: they are strictly serial by design (that serialisation fixed a 370 s starvation bug). This is what makes batching worth a factor of hundreds rather than a few percent: merging launches converts a SUM into a MAX. Ten serialised 50 ms launches cost 500 ms; the same Rs in one launch cost ~50 ms, because the short Rs run beside the long one instead of queueing behind it. Neither block size nor any per-thread micro-optimisation can reach that, which is why every attempt has landed between -3% and +8%. Two candidates, recorded on the kernel: stream concurrency (cheap -- the launches are already independent and FIFO dispatch could round-robin across N streams and stay fair, but streams were pinned to 1 to fix a host-memory blowup from per-stream pinned pools), or aggregating Rs across calls (needs enumeration decoupled from the multiply that consumes its output). Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 25 ++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 04671340e3..e300b61ca2 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -4159,6 +4159,31 @@ fn seg_read_u32( /// 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 From 78872ffca3da0948693a60b31985e028d83ffdcb Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Fri, 7 Aug 2026 17:45:41 -0400 Subject: [PATCH 115/127] =?UTF-8?q?milnor=5Fgpu:=20GPU=20stream=20concurre?= =?UTF-8?q?ncy=20does=20not=20pay=20=E2=80=94=20the=20run=20is=20not=20GPU?= =?UTF-8?q?-bound?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `NASSAU_GPU_STREAMS` (default 1, i.e. unchanged) lets N worker threads drain one device's queue, each on its own `StreamId`, so independent device sections overlap instead of running strictly one after another. This was the cheap half of the "batch the enum kernel" plan: `enumerate_admissible_kernel` is ~99% of GPU kernel time and runs at `Waves Per SM = 0.002`, and merging serialised launches turns a SUM into a MAX. Measured (S_2 stem 150, max_s 60, theta=125), and it refutes the plan: | 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 memory growth is the load-bearing part of the result: it proves the sections really did overlap, so this is not "the streams did not engage" — concurrency happened and bought nothing. The conclusion is therefore about the workload, not about streams: at this configuration the resolution is not GPU-throughput-bound. Kernel time is 99% enum, but kernel time is not the critical path. Measured alongside: ~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. Kept rather than reverted: a few lines, defaults to exactly the old behaviour, and it is the control that makes the "not GPU-bound" claim falsifiable elsewhere. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 121 +++++++++++++++---- 1 file changed, 96 insertions(+), 25 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index e300b61ca2..c26842e337 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -246,8 +246,10 @@ mod gpu_thread { /// Balance therefore comes from spreading `R`s evenly (round-robin at first sight), not from /// letting idle workers steal. /// - /// Each worker owns its device, its stream, and (via the thread-local set here) its own replica - /// of the resident master/basis, so a device handle can never reach another device's client. + /// 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 @@ -255,23 +257,35 @@ mod gpu_thread { 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); - std::thread::Builder::new() - .name(format!("nassau-gpu{dev}")) - .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 per device so the runtime keeps them independent. - StreamId { value: dev as u64 }.executes(|| { - while let Ok(task) = rx.recv() { - task(); + // `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, } - }); - }) - .expect("failed to spawn a nassau-gpu thread"); + .executes(|| { + while let Ok(task) = rx.recv() { + task(); + } + }); + }) + .expect("failed to spawn a nassau-gpu thread"); + } } txs }) @@ -751,6 +765,63 @@ fn gpu_count() -> usize { *N } +/// Worker threads — and therefore CUDA streams — per device (`NASSAU_GPU_STREAMS`, default 1). +/// +/// 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. +/// +/// 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 +} + 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. @@ -1357,8 +1428,8 @@ pub fn dump_r_stats() { }; 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 \ + 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), @@ -3890,8 +3961,8 @@ fn multiply_batch_block<'a>( 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} \ + 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, @@ -4862,8 +4933,8 @@ mod tests { 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", + " 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) ); @@ -5520,8 +5591,8 @@ mod tests { } 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}", + "R={p_part:?} rows={rows} cols={cols} \ + (row={row},col={col}) d={d:#x} masks[{}]={m:#x}", row + col )); } @@ -5579,9 +5650,9 @@ mod tests { 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}%)", + "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) From 72d6e1faf9e410bbdaa266913dd9bd71b835c137 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Fri, 7 Aug 2026 17:45:57 -0400 Subject: [PATCH 116/127] nassau_gpu: read the GPU multiply back a limb at a time, not a bit at a time MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. The readback can therefore be a truncating byte copy into a scratch `FpVector` plus one limb-wise `add`, instead of walking the set bits and calling `add_basis_element` on each. At the logged ~26% density that is cols/64 word XORs per row in place of ~0.26*cols bounds-checked entry writes. `update_from_bytes` refills the scratch in place, so the whole call allocates two buffers rather than one per row. Applies to both `get_partial_matrix` and the restricted variant; the restricted one also drops bits past `target_dim`, which is exactly what its old `col < target_dim` guard did — whole limbs by truncation, the partial final limb by mask. Verified with NASSAU_GPU_VERIFY over a full S_2 stem-40 max_s-30 resolution, which builds every matrix both ways and compares nonzero column sets: 0 mismatches. Honest timing result: no measurable win (528 s, identical to baseline at stem 150 / theta=125). That is the same story as the streams experiment — the run is bound by a serial dependency chain at ~7% CPU, so making CPU work cheaper does not move wall time. Kept because it is verified-equivalent and strictly less work, and it will matter once the serial chain is opened up. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/src/nassau_gpu.rs | 100 ++++++++++++++++++++++++++++++++---------- 1 file changed, 77 insertions(+), 23 deletions(-) diff --git a/ext/src/nassau_gpu.rs b/ext/src/nassau_gpu.rs index 2c5adc5e3b..4a83a1a309 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,7 +26,7 @@ use algebra::{ homomorphism::{FreeModuleHomomorphism, ModuleHomomorphism}, }, }; -use fp::matrix::Matrix; +use fp::{matrix::Matrix, vector::FpVector}; type NassauDifferential = FreeModuleHomomorphism>; @@ -47,22 +48,43 @@ 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 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() { - 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 copy = (limbs.len() * size_of::()).min(nbytes); + for (k, &w) in limbs + .iter() + .take(copy.div_ceil(size_of::())) + .enumerate() + { + let o = k * size_of::(); + let n = size_of::().min(nbytes - o); + buf[o..o + n].copy_from_slice(&w.to_le_bytes()[..n]); } + buf[copy..].fill(0); + 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 @@ -188,6 +210,7 @@ pub fn get_partial_matrix_restricted( 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. @@ -205,6 +228,29 @@ pub fn get_partial_matrix_restricted( // 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() { @@ -218,21 +264,29 @@ pub fn get_partial_matrix_restricted( 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() { - let mut target_row = matrix.row_mut(r0 + bi); - 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; - } + // 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). + let copy = (limbs.len() * size_of::()).min(nbytes); + for (k, &w) in limbs + .iter() + .take(copy.div_ceil(size_of::())) + .enumerate() + { + let o = k * size_of::(); + let n = size_of::().min(nbytes - o); + buf[o..o + n].copy_from_slice(&w.to_le_bytes()[..n]); } + buf[copy..].fill(0); + 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; From 75faee70230ad9a78b4011a552d629b75e5edeba Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Fri, 7 Aug 2026 17:50:26 -0400 Subject: [PATCH 117/127] =?UTF-8?q?nassau:=20the=20missing=20parallelism?= =?UTF-8?q?=20is=20not=20inside=20a=20bidegree=20=E2=80=94=20two=20probes,?= =?UTF-8?q?=20both=20negative?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Ran the `NASSAU_PROBE_SIG_INDEP` probe that had been sitting unused: 19.2% of the values a signature reads (27 028 of 140 614) were written by an earlier signature, so the loop is a genuine forward substitution and the lift must stay ordered. That only rules out reordering the lift. Everything above it reads read-only state, so a windowed prepare stage feeding an ordered lift is legal, and it was built and measured. It is worthless, for a reason worth recording so nobody rebuilds it: 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 is 4%. End to end a window of 4 measured 537 s against a 528 s baseline while raising mean CPU from 918% to 1306% — 42% more CPU to lose 1.7%, exactly as the ceiling predicts. The refactor is reverted; only the measurement is kept. Context for what this closes off: the run uses ~7% of a 128-core node with a wavefront ~5 bidegrees wide, and GPU stream concurrency was independently measured to buy nothing. Three attempts to add concurrency have now failed, which locates the limiter as latency on a serial chain rather than any throughput shortfall. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/src/nassau.rs | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/ext/src/nassau.rs b/ext/src/nassau.rs index f458f91be8..6fa2106461 100644 --- a/ext/src/nassau.rs +++ b/ext/src/nassau.rs @@ -916,6 +916,28 @@ impl> Resolution { // 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; From fbf5c7ef217185a2ca9b68ee056956276210d434 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Fri, 7 Aug 2026 18:56:51 -0400 Subject: [PATCH 118/127] =?UTF-8?q?nassau:=20span=20the=20zero-signature?= =?UTF-8?q?=20region=20=E2=80=94=20and=20find=20the=20run=20is=2044%=20CUD?= =?UTF-8?q?A=20spin-wait?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The zero signature is ~96% of its bidegree (it is unconstrained, so its mask is the widest and its matrix the largest — reproducible in 100% of bidegrees across four runs), and almost all of its work sat outside every span. Adds `cpu_restricted`, `zs_masks`, `zs_select`, `zs_source_mask` and `zs_dx_init` to close that gap. They are all small, which is itself the result — none of the CPU regions I suspected is the cost: extract_restricted 1062.7 s cpu_restricted 97.0 s gpu_readback 319.6 s pair_prepass 90.1 s marshal_terms 54.3 s zs_select 1.8 s gpu_submit 0.3 s zs_dx_init 0.3 s against a zero-signature `step` total of 4844.9 s, leaving ~66% unattributed. Span accounting could not find it because `time.busy` is inclusive AND the thing being waited on is not inside any span: `[batch-stats]` reports `fence=1024.2 s` (caller blocked on the GPU completion fence) against a worker-side `exec=77.1 s`. A sampling profile settles it. `perf record` on a steady-state stem-150 run, by shared object: 43.84% libcuda.so 29.18% the binary 17.40% libc The whole libcuda share is one tight address range (0x461860-0x461999): the driver's spin-wait. The process burns ~44% of its CPU busy-waiting on the device. That retro-explains every negative result in this branch of work. Extra CUDA streams, a parallel signature-prepare window, and a cheaper matrix readback all failed for the same reason — they add or cheapen CPU work on a path that is waiting for the GPU, and the extra threads spin too. It also puts `enumerate_admissible_kernel` (99% of GPU kernel time) back on the critical path as a LATENCY problem, which is a different target from the throughput framing that streams were testing. Top CPU symbols for whoever picks this up: `get_partial_matrix_restricted` 11.3%, allocator (`_int_free`/`malloc_consolidate`/`malloc`/`_int_malloc`) 8.6%, `__memcpy_ssse3` 5.9%, `__vdso_clock_gettime` 4.7%. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/src/nassau.rs | 64 +++++++++++++++++++++++++++++------------------ 1 file changed, 39 insertions(+), 25 deletions(-) diff --git a/ext/src/nassau.rs b/ext/src/nassau.rs index 6fa2106461..3dbb0211f3 100644 --- a/ext/src/nassau.rs +++ b/ext/src/nassau.rs @@ -364,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 @@ -777,9 +781,11 @@ 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 @@ -800,18 +806,21 @@ 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 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, @@ -859,9 +868,11 @@ impl> Resolution { // 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 = subalgebra - .signature_mask(&algebra, &self.modules[b.s()], b.t(), &zero_sig, i32::MAX) - .collect(); + 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(), @@ -891,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); + } } } From a9d3f5e1352a070ca48cbc4cb4596479aa60a14c Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Fri, 7 Aug 2026 19:12:36 -0400 Subject: [PATCH 119/127] =?UTF-8?q?milnor=5Fgpu:=20theta=3D125=20at=20stem?= =?UTF-8?q?=20150=20is=20an=20artificial=20worst=20case=20=E2=80=94=202.85?= =?UTF-8?q?x=20slower=20than=20uncapped?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The stream sweep (and every other A/B in this session) ran at NASSAU_GPU_RESIDENT_MAX_DEGREE=125. The same runs' own `[MASTER-BY-DEGREE]` line reports the FULL master at stem 150 as 1.5 GB, 0.4 GB/GPU — so the cap was forcing the transient re-enumeration path for a master that fits resident several times over. Same binary, same session, stem 150 / max_s 60: | theta | wall | enum launches | mean CPU | |-------|-------|---------------|----------| | 125 | 575 s | 10260 | 1110% | | none | 202 s | 0 | 1618% | 2.85x, with `enumerate_admissible_kernel` not launching once. Two thirds of the wall time being optimised all session was enum work this stem would never do. This does not overturn the stream result (that arm is still flat-to-worse against its own control), but it does bound what any theta-capped benchmark can say about where time goes: at theta=125 it mostly says where the CAP puts time. Records the caveat next to the numbers, and says to size theta from `[MASTER-BY-DEGREE]` — cap only when the master genuinely does not fit. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index c26842e337..ceb893ef69 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -806,6 +806,24 @@ fn gpu_count() -> usize { /// 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 From c068c4888f8899aa3e2549492b061e3d48710bee Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Fri, 7 Aug 2026 19:44:07 -0400 Subject: [PATCH 120/127] milnor_gpu: record that FxHashMap is 8% SLOWER than SipHash here MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A perf profile of an UNCAPPED stem-150 run (the regime that matters — see the theta note) is CPU-bound, not GPU-bound: libcuda falls to 2.96% of user cycles, and the top costs are memcpy 15.3%, get_partial_matrix_restricted 14.4%, the allocator 16.1%, resident_info 7.7%, row_reduce 7.6%, and 2.2% in DefaultHasher/RandomState. That last one makes swapping the resident index (and COLD_COUNT, the per-launch intern map, and the row remap) to `rustc_hash::FxHashMap` look like free money: the keys are packed integers with no adversary. Measured, three replicates per arm: pre-change 205, 208, 195 mean 202.7 s FxHashMap 234, 210, 213 mean 219.0 s ~8% slower, so it is reverted. Fx is a weak multiply-rotate and evidently clusters on packed `PPart` keys where SipHash spreads them. The replicates are the other result. Run-to-run noise at stem 150 is ~3% (and ~9% at theta=125 across sessions), which is larger than several single-sample deltas measured earlier in this session — those should be read as "no effect", not as small effects. Use three runs per arm here. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index ceb893ef69..c4e2581dea 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -705,6 +705,13 @@ struct ResidentHost { /// 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, } From d9f2e45c29658c24a285333f0d69613f31382c28 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Fri, 7 Aug 2026 20:49:09 -0400 Subject: [PATCH 121/127] nassau_gpu: one bulk memcpy per readback row instead of one per u32 limb MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The limb-wise readback wrote the GPU's `u32` limbs into the staging buffer one at a time (`buf[o..o+n].copy_from_slice(&w.to_le_bytes()[..n])`) — a bounds-checked 4-byte copy per limb, thousands per row, for a region that is already byte-identical on a little-endian target. A call-graph profile of an UNCAPPED stem-150 run put that `copy_from_slice` at 12.15% of all user cycles, under `restricted_partial_matrix_maybe_gpu`. `fill_limb_bytes` views the row as bytes and does a single `memcpy`, keeping an explicit big-endian fallback so a non-little-endian target stays correct rather than silently wrong (`fp`'s own `limb::from_bytes` takes the same shortcut under the same `cfg`). Verified with NASSAU_GPU_VERIFY over a full S_2 stem-40 max_s-30 resolution: 0 mismatches. Worth 5.0% on its own (199.7 s vs 210.3 s, interleaved, 3 rounds). See the following commit for the other half. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/src/nassau_gpu.rs | 62 ++++++++++++++++++++++++++++--------------- 1 file changed, 40 insertions(+), 22 deletions(-) diff --git a/ext/src/nassau_gpu.rs b/ext/src/nassau_gpu.rs index 4a83a1a309..0c4ebce67a 100644 --- a/ext/src/nassau_gpu.rs +++ b/ext/src/nassau_gpu.rs @@ -30,6 +30,44 @@ 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`]. @@ -67,17 +105,7 @@ pub fn get_partial_matrix(hom: &NassauDifferential, degree: i32, inputs: &[usize r => (1u64 << r) - 1, }; for (row, limbs) in out.iter_rows().enumerate() { - let copy = (limbs.len() * size_of::()).min(nbytes); - for (k, &w) in limbs - .iter() - .take(copy.div_ceil(size_of::())) - .enumerate() - { - let o = k * size_of::(); - let n = size_of::().min(nbytes - o); - buf[o..o + n].copy_from_slice(&w.to_le_bytes()[..n]); - } - buf[copy..].fill(0); + 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()); @@ -269,17 +297,7 @@ pub fn get_partial_matrix_restricted( // 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). - let copy = (limbs.len() * size_of::()).min(nbytes); - for (k, &w) in limbs - .iter() - .take(copy.div_ceil(size_of::())) - .enumerate() - { - let o = k * size_of::(); - let n = size_of::().min(nbytes - o); - buf[o..o + n].copy_from_slice(&w.to_le_bytes()[..n]); - } - buf[copy..].fill(0); + 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()); From 892835643e8571a52abbdaa576aa115340013d98 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Fri, 7 Aug 2026 20:49:09 -0400 Subject: [PATCH 122/127] =?UTF-8?q?milnor=5Fgpu:=20share=20GpuProduct=20te?= =?UTF-8?q?rm=20lists=20=E2=80=94=201.45x=20on=20stem=20150?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `term_indices` was a `Vec` written once and only ever read, but a product is cloned twice on its way to the device: once to compact rows into a dense range per hot/cold group, once to fan out into per-device buckets. Each clone duplicated every term list, and the drops showed up as 5.45% of all user cycles in `_int_free` alone (16.1% total in the allocator) in a call-graph profile of an uncapped stem-150 run. `Arc<[usize]>` makes the clones a refcount bump and the drops O(1). `Arc<[T]>` implements `FromIterator`, so every construction site still just `.collect()`s; only two `for x in &prod.term_indices` loops needed an explicit `.iter()`. MEASURED, interleaved A/B/C, 3 rounds (arms alternate so drift in machine load hits all of them equally): A control 211, 214, 206 mean 210.3 s B bulk-memcpy readback 188, 203, 208 mean 199.7 s -5.0% C B + Arc term_indices 149, 151, 135 mean 145.0 s -31.1% C wins every round with no overlap against A. 1.45x. Method note, because it changed the answer: the same two changes measured BLOCK-sequentially looked like regressions (+7%). Arms run in blocks confound the effect with load drift, and this node drifts by more than 5%. Interleave, or do not believe the number — an earlier `FxHashMap` verdict in this file was measured the block way and should be re-checked before it is trusted. Verified with NASSAU_GPU_VERIFY over a full S_2 stem-40 max_s-30 resolution: 0 mismatches. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index c4e2581dea..5edaa33da6 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -2409,7 +2409,15 @@ 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, } @@ -2556,7 +2564,7 @@ pub fn cpu_multiply_batch( } 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, block_dim); @@ -2976,7 +2984,10 @@ fn multiply_batch_block<'a>( 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) { + for (slot, &ti) in tg_all[off..off + nt] + .iter_mut() + .zip(prod.term_indices.iter()) + { *slot = base + ti as u32; } } From 33c971352d2d7c4e1b10bb48c66a9a6bf0af612e Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Sat, 8 Aug 2026 00:10:25 -0400 Subject: [PATCH 123/127] =?UTF-8?q?milnor=5Fgpu:=20memoize=20the=20per-pro?= =?UTF-8?q?duct=20R=20lookup=20=E2=80=94=206-8%?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `resident_info` takes a read lock on the process-wide `RESIDENT_HOST` and SipHashes a `PPart`. It was called once per PRODUCT in two places — the pair pre-pass and the per-device fan-out — but `R` is a property of the input ROW, and `extract_restricted` emits one product per target generator block of a row, so consecutive products repeat the same `R`. Every repeat was paying for a lock and a hash to learn something it had just learned. A call-graph profile of an uncapped stem-150 run (i.e. after the memcpy/Arc fixes, so this is the NEW top cost) 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`. A one-entry memo keyed on `(r_degree, r_idx)` catches the repeats with no hash, no allocation and no lock. Per distinct `R` the cost is unavoidable; per product it was pure repetition. MEASURED, interleaved, arms alternating within each round: stem 150, 3 rounds C 149,156,150 -> mean 151.7 s D 150,139,131 -> mean 140.0 s -7.7% stem 200, 2 rounds C 1857,1904 -> mean 1880.5 s D 1755,1782 -> mean 1768.5 s -6.0% At stem 200 the arms do not overlap. All four stem-200 runs verified complete (max_n=200, max_s=110, ~3.599M bidegree spans each, 0 panics) — a wall time from a run that lost its GPU workers would otherwise look like a win. Verified with NASSAU_GPU_VERIFY over a full S_2 stem-40 max_s-30 resolution: 0 mismatches. Cumulative with the two preceding commits: stem 150 210.3 s -> 140.0 s (1.50x), stem 200 1768.5 s against a 2412 s pre-session reference (1.36x). Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 44 +++++++++++++++++--- 1 file changed, 38 insertions(+), 6 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 5edaa33da6..78bf923632 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -2684,13 +2684,33 @@ fn multiply_batch_grouped( ); 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 r = algebra.basis_element_from_index(prod.r_degree, prod.r_idx); - let num_mats = 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, + 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) @@ -2726,14 +2746,26 @@ fn multiply_batch_grouped( // 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. MasterMode::Transient => pi % gpu_count(), MasterMode::Resident => { - let r = algebra.basis_element_from_index(prod.r_degree, prod.r_idx); - resident_info(algebra, r.p_part).dev as usize + 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()); From 3ec60d706e325a39079835bc17e70dc6ad99d1ae Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Sat, 8 Aug 2026 00:58:53 -0400 Subject: [PATCH 124/127] milnor_gpu: scatter eviction output straight out of the device blocks MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The eviction path flattened every block of a hot/cold group into one `Vec` and then XORed that into place — allocating and copying the ENTIRE group output, reading it once, and dropping it. Blocks are whole rows of `num_limbs` each in compacted row order, so the flatten buys nothing: walking them in `num_limbs` chunks visits exactly the rows `rows` names, in the same order. Only reachable when theta is CAPPED — the uncapped path short-circuits earlier — so this is a high-stem change by construction, and stem 150 can only weakly test it. MEASURED, interleaved, 3 rounds at theta=125 / stem 150 (the smallest config that exercises the path): before 378, 406, 398 mean 394.0 s after 387, 390, 376 mean 384.3 s -2.5% 2.5% with overlapping spreads is inside the noise at this scale, so treat the number as "not a regression", not as the win. It is kept because it is strictly less work — one fewer allocation and one fewer full pass over the output — and the regime where that output is hundreds of MB per call (stems 250-300, where theta MUST be capped) is exactly the regime this is for and the one a stem-150 A/B cannot see. Verified with NASSAU_GPU_VERIFY under a cap (NASSAU_GPU_RESIDENT_MAX_DEGREE=20, so the eviction split is actually exercised) over a full S_2 stem-40 max_s-30 resolution: 0 mismatches. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 26 ++++++++++++++------ 1 file changed, 18 insertions(+), 8 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 78bf923632..c75a8f9dc8 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -2638,16 +2638,26 @@ fn multiply_batch_gpu_inner( }) .collect(); let sub_blocks = multiply_batch_grouped(algebra, num_cols, rows.len(), &compact, mode); - let sub: Vec = sub_blocks - .iter() - .flat_map(|b| u32::from_bytes(b).iter().copied()) - .collect(); - for (i, &orig) in rows.iter().enumerate() { - let (dst, src) = (orig * num_limbs, i * num_limbs); - for k in 0..num_limbs { - result[dst + k] ^= sub[src + k]; + // 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) } From c28f4311e4c9e17373d10e94a75f726fc3a011ce Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Sat, 8 Aug 2026 01:52:57 -0400 Subject: [PATCH 125/127] =?UTF-8?q?milnor=5Fgpu:=20put=20the=20enum=20kern?= =?UTF-8?q?el's=20state=20in=20shared=20memory=20=E2=80=94=205.5%=20in=20t?= =?UTF-8?q?he=20capped=20regime?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `enumerate_admissible_kernel` held its per-thread state in `Array`, which is CUDA *local* memory: every index is a runtime value, so none of it can be register- allocated and each access is a real off-chip load. ncu attributed 45.6% of the kernel's 8.9 stall cycles to waiting on exactly that. Moving it to `Shared<[u32]>` at stride `elem * ENUM_BLOCK + UNIT_POS` (lanes of a warp on consecutive words, so conflict-free) forces the block from 256 to 32 threads — 152 u32/thread is 19.5 KB at 32 but 155 KB at 256, against a 48 KB static budget. THIS WAS TRIED BEFORE AND REJECTED, and the rejection was right for the benchmark it was measured on: on the all-`R` benchmark 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, occupancy-bound. Production is not that shape: a real launch is ~6 blocks against the H200's 3168 slots (`Waves Per SM = 0.002`), so there is no occupancy to lose and the -33% converts. The block-size sweep (256/64/32 -> 561/569/542 s, flat) says the same thing from the other side. MEASURED where the enum kernel actually runs, i.e. theta CAPPED, interleaved, 3 rounds at theta=125 / stem 150: before 375, 374, 375 mean 374.7 s (spread 0.6 s) after 345, 372, 345 mean 354.0 s -5.5% The control is almost perfectly repeatable, and two of the three "after" runs sit 8% below it, well outside that distribution; one run did not land. So 5.5% is the mean, ~8% is what it does when it works, and the variance is unexplained. This is the regime that matters for stems above 200: the resident master grows ~3.5x per +25 degrees (at stem 200: theta<=150 1.7 GB/GPU, <=175 6.8, <=200 23.2), so past stem 200 theta must be capped and everything above it is enumerated on the device. Correctness: `admissible_enum_gpu_matches` and `admissible_enum_ref_matches` both pass (bit-exact against the CPU reference), plus NASSAU_GPU_VERIFY under a cap over a full S_2 stem-40 max_s-30 resolution, 0 mismatches. Also fixes three `#[cfg(test)]` sites left broken by the `Arc<[usize]>` change — the release example build never compiles them, and `just lint` runs clippy without the `gpu` feature, so gpu-gated test code has no automated gate at all. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 128 ++++++++++++------- 1 file changed, 85 insertions(+), 43 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index c75a8f9dc8..51f581c90a 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -72,6 +72,30 @@ 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 @@ -3512,7 +3536,8 @@ fn multiply_batch_block<'a>( // 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). - const ENUM_THREADS: u32 = 256; + // 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 @@ -4393,10 +4418,13 @@ fn enumerate_admissible_kernel( // 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. - let mut matrix = Array::::new(ENUM_MATRIX_CAP); - let mut totals = Array::::new(ENUM_ROW_CAP); - let mut col_sums = Array::::new(ENUM_COL_CAP); - let mut masks = Array::::new(ENUM_MASK_CAP); + // 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 @@ -4404,22 +4432,22 @@ fn enumerate_admissible_kernel( // 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 { - matrix[i] = 0u32; + st[(i) * bs + tid] = 0u32; } for i in 0..rows { - totals[i] = 0u32; + st[(ENUM_ST_TOTALS + i) * bs + tid] = 0u32; } for i in 0..cs_len { - col_sums[i] = 0u32; + st[(ENUM_ST_COLSUMS + i) * bs + tid] = 0u32; } for i in 0..mk_len { - masks[i] = 0u32; + 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]; - matrix[i * cols] = x; - masks[i] = x; + st[(i * cols) * bs + tid] = x; + st[(ENUM_ST_MASKS + i) * bs + tid] = x; } let mut mat = 0usize; @@ -4433,31 +4461,33 @@ fn enumerate_admissible_kernel( if emit == 1 { let co = cs_base + mat * cs_len; for j in 0..cs_len { - out_cs[co + j] = u16::cast_from(col_sums[j]); + 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(masks[j]); + 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(col_sums[j]); + 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(masks[j]); + 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(col_sums[j]); + 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(masks[j2]); + out_mk[mo + j2] = u16::cast_from(st[(ENUM_ST_MASKS + j2) * bs + tid]); j2 += 2; } } @@ -4469,12 +4499,12 @@ fn enumerate_admissible_kernel( let mut row = 0usize; while row < rows && !found { let mut p_to_the_j = 1u32; - totals[row] = matrix[row * cols]; + 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 <= totals[row] { + 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; @@ -4482,50 +4512,62 @@ fn enumerate_admissible_kernel( c = row + col + 1 - rows; } while c < col { - d |= matrix[(row + col - c) * cols + c]; + d |= st[((row + col - c) * cols + c) * bs + tid]; c += 1; } - let cur = matrix[row * cols + col]; + 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 totals[row] < sub { - totals[row] += p_to_the_j * cur; + 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 { - matrix[row * cols] = totals[row] - sub; - masks[row] = matrix[row * cols]; - col_sums[col - 1] += inc; + 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 { - masks[row + j] &= !matrix[row * cols + j]; - col_sums[j - 1] -= matrix[row * cols + j]; - matrix[row * cols + j] = 0u32; + 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; } - matrix[row * cols + col] = new_entry; + st[(row * cols + col) * bs + tid] = new_entry; let mut i = 0usize; while i < row { - matrix[i * cols] = totals[i]; - masks[i] = totals[i]; + 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 { - masks[i + j2] &= !matrix[i * cols + j2]; + st[(ENUM_ST_MASKS + i + j2) * bs + tid] = st + [(ENUM_ST_MASKS + i + j2) * bs + tid] + & !st[(i * cols + j2) * bs + tid]; } - col_sums[j2 - 1] -= matrix[i * cols + j2]; - matrix[i * cols + j2] = 0u32; + 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; } - masks[row + col] = d | new_entry; + st[(ENUM_ST_MASKS + row + col) * bs + tid] = d | new_entry; found = true; handled = true; } } if !handled { - totals[row] += p_to_the_j * matrix[row * cols + col]; + 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; } @@ -4919,7 +4961,7 @@ mod tests { let omk_h = client.empty(mk_cap * size_of::()); let cnt_h = client.empty(n_r * size_of::()); - const THREADS: u32 = 64; + 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::( @@ -5084,7 +5126,7 @@ mod tests { (1u64 << (63 - lim.leading_zeros().min(62))).min(lim) as u32 - 1 }; - const THREADS: u32 = 64; + 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 { @@ -5406,7 +5448,7 @@ mod tests { let cnt_h = client.empty(n_r * size_of::()); let marshal_s = t_marshal.elapsed().as_secs_f64(); - const THREADS: u32 = 64; + 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 { @@ -6182,7 +6224,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); @@ -6287,7 +6329,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); @@ -6453,7 +6495,7 @@ mod tests { (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 { + for &ti in prod.term_indices.iter() { s.set_entry(ti, 1); } let mut tmp = FpVector::new(p, out_dim); From e30ecebeeef9c889c6deb01269997f3e6ee1a834 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Sat, 8 Aug 2026 04:13:01 -0400 Subject: [PATCH 126/127] milnor_gpu: record that de-duplicating enum work across devices buys nothing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Transient products are spread round-robin by product index, but `R` is a property of the input ROW, so consecutive products share one `R` and round-robin scatters each `R` across up to `gpu_count()` devices — every one of which enumerates it independently. Hashing `R` to a device instead (`shard_of`, as `Resident` already does) fixes that. The mechanism works exactly as intended: `Rs/launch` fell 1293 -> 832, removing 36% of all `R`-enumerations (13.3M -> 8.5M). It bought nothing: round-robin 393, 364, 362 mean 373.0 s peak GPU 129.4 GB R-affine 360, 373, 383 mean 372.0 s peak GPU 126.0 GB Time-neutral, and 2.6% off peak memory — the transient enum scratch is not what drives GPU memory. Reverted; only the measurement is kept. WHY IT COULD NOT HAVE WORKED, which is the useful part. An ncu profile of the enum kernel at PRODUCTION geometry (not the all-`R` benchmark) shows launches of 3-104 blocks x 32 threads running 7-100 ms at 1.56% achieved occupancy, 4.6-11 of 32 lanes active, 80% "No Eligible", compute 0.07% and DRAM 0.00%. Nothing is saturated: a launch's duration is set by its LONGEST single `R` chain. Removing duplicate work from the other chains cannot move a maximum. This also retires emit packing as a candidate: its 42.8%-of-kernel figure came from the benchmark, where 1397 blocks did saturate the machine. Here there is no bandwidth to reclaim. What is left that can actually shorten the longest chain: split it (needs an unrank to jump to the k-th admissible matrix, since the odometer derives each from the previous), or keep the highest-`num_mats` `R`s resident regardless of theta so the long poles are never re-enumerated. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 51f581c90a..93af0cb02b 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -2788,6 +2788,22 @@ fn multiply_batch_grouped( 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); From 05bf3e05ebdf70aa96170cc8fa189b4f9cb2373c Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Sat, 8 Aug 2026 05:20:47 -0400 Subject: [PATCH 127/127] =?UTF-8?q?milnor=5Fgpu:=20pin=20long-pole=20R's?= =?UTF-8?q?=20resident=20regardless=20of=20theta=20=E2=80=94=2013.6%=20cap?= =?UTF-8?q?ped?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `NASSAU_GPU_PIN_MIN_MATS` (default 0 = off) routes any `R` with at least that many admissible matrices into the RESIDENT master whatever its degree. The idea comes straight out of the ncu profile at production geometry: an enum launch is 3-104 blocks of 32 threads running 7-100 ms at 1.56% achieved occupancy, so its duration is set by its LONGEST single odometer chain, not by how many `R`s it covers. Master BYTES, meanwhile, are a SUM. Those two facts point in opposite directions, and the long-pole `R`s sit on the good side of the asymmetry: they dominate time while contributing almost nothing to size. MEASURED, interleaved, 3 rounds, theta=125 / stem 150: | PIN | runs | mean | | master | |-------|---------------|---------|--------|--------| | 0 | 347, 347, 369 | 354.3 s | | 1.5 GB | | 10000 | 324, 331, 338 | 331.0 s | -6.6% | 4.5 GB | | 2000 | 300, 317, 301 | 306.0 s | -13.6% | 6.5 GB | Monotone, and PIN=2000 wins every round. Note the mechanism is visible in the counters: `Rs/launch` falls only 1293 -> 1038 (20%) while wall falls 13.6% — few `R`s removed, but the right ones. The contrast that proves the model: making transient sharding `R`-affine removed 36% of ALL enumerations and bought exactly nothing (372.0 s vs 373.0 s), because it shortened no chain. Work reduction is worthless here; only max reduction pays. Left OFF by default because the right threshold is workload-dependent and the extra master competes with theta for the same device memory — which above stem 200 is the binding constraint. Size it from `[MASTER-BY-DEGREE]` alongside theta. Verified with NASSAU_GPU_VERIFY under a cap with pinning active: 0 mismatches. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 53 ++++++++++++++++++-- 1 file changed, 48 insertions(+), 5 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 93af0cb02b..8fdca58a45 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -1207,6 +1207,38 @@ static R_STATS: LazyLock>>> = /// /// 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 } @@ -2636,15 +2668,26 @@ fn multiply_batch_gpu_inner( // 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| match mode { - MasterMode::Resident => d <= cap, - MasterMode::Transient => d > cap, + 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)) + .filter(|p| is_group(p.r_degree, p.r_idx)) .map(|p| p.row) .collect(); rows.dedup(); @@ -2654,7 +2697,7 @@ fn multiply_batch_gpu_inner( let remap: HashMap = rows.iter().enumerate().map(|(i, &r)| (r, i)).collect(); let compact: Vec = products .iter() - .filter(|p| is_group(p.r_degree)) + .filter(|p| is_group(p.r_degree, p.r_idx)) .map(|p| { let mut q = p.clone(); q.row = remap[&p.row];