From 83819efa19ef131250aae5f8ee17d16ecd8f8514 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Mon, 20 Jul 2026 01:13:17 -0400 Subject: [PATCH 01/12] 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 02/12] 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 03/12] 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 04/12] 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 05/12] 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 06/12] 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 07/12] 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 08/12] 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 09/12] 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 10/12] 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 11/12] 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 12/12] 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). ///