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/16] 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/16] 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/16] 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/16] 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/16] 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/16] 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/16] 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/16] 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/16] 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/16] 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/16] 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/16] milnor_gpu: raise GPU_PAIR_CHUNK to ~2^32 so giant multiplies stop over-splitting MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The row-block splitter caps a launch at GPU_PAIR_CHUNK thread-pairs (the kernel indexes threads by u32 ABSOLUTE_POS, ceiling 2^32). It was set to 1<<30 (~1.07e9), ~4x below the real ceiling — so every billion-pair giant was chopped into ~4 launches, each a separate upload + kernel + BLOCKING readback round-trip, even though its output is only ~350 MB (well under gpu_block_bytes). Debug confirmed the giants pegged at 1.07e9 pairs; this, not the byte budget, was the binding split, which is why a NASSAU_GPU_BLOCK_MB sweep was flat. Raise it to 3.9e9 (leaves ~0.39e9 headroom under 2^32; the splitter always takes >=1 row and a lone row past 2^32 still trips the per-block u32::try_from assert; grid stays ~1.5e7 cubes, far under 2^31). Giants now run as a single launch (max total_pairs observed 3.90e9), collapsing 4 round-trips to 1. GPU 0->140 (w=100): ~245-288s -> 216s. Chart byte-identical to CPU through (100,152). Co-Authored-By: Claude Opus 4.8 --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 23 +++++++++++++------- 1 file changed, 15 insertions(+), 8 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 65ae9d5b1c..386d833c05 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -35,14 +35,21 @@ use crate::algebra::{Algebra, MilnorAlgebra, combinatorics::xi_degrees}; /// `mk_len = rows + cols − 1 ≤ MAX_XI_TAU + ⌈log2⌉`; 32 covers every in-range case. const WORKING_CAP: usize = 32; -/// Target `(product, matrix, term)` thread-pairs per GPU launch. The batch multiply indexes -/// threads by CubeCL's `ABSOLUTE_POS` (a `u32`), so one launch can address at most `2^32` threads; -/// a single all-rows reuse build reaches ~4.4e9 pairs at stem ~145, past that limit. The row-block -/// splitter in [`multiply_batch_on_gpu`] closes a block once its pair count would pass this target -/// (alongside the [`gpu_block_bytes`] output budget). `1 << 30` (~1.07e9) leaves >3x headroom -/// under `2^32` even when a lone over-budget row overshoots it, and keeps the grid -/// (`total_pairs / 256` cubes) well under CUDA's `2^31 - 1` grid-dimension limit. -const GPU_PAIR_CHUNK: usize = 1 << 30; +/// Target `(product, matrix, term)` thread-pairs per GPU launch. The batch multiply indexes threads +/// by CubeCL's `ABSOLUTE_POS` (a `u32`), so one launch can address at most `2^32` threads; a single +/// all-rows reuse build reaches ~4.4e9 pairs at stem ~145, past that limit. The row-block splitter +/// in [`multiply_batch_on_gpu`] closes a block once its pair count would pass this target (alongside +/// the [`gpu_block_bytes`] output budget). +/// +/// Set close to the `2^32` ceiling, not far below it: every extra split is a whole extra launch +/// (upload + kernel + blocking readback), and at record stems this — not the byte budget — is the +/// binding constraint, so a conservative value chops each giant multiply into several +/// otherwise-unnecessary launches (measured: `1 << 30` pegged the giants at ~1.07e9 pairs, ~4 +/// launches each, while their output is only ~350 MB, well under `gpu_block_bytes`). `3.9e9` leaves +/// ~0.39e9 of headroom under `2^32` for a lone over-budget row (the splitter always takes ≥1 row, +/// and a single row past `2^32` still trips the per-block `u32::try_from` assert), and keeps the +/// grid (`pairs / 256` cubes ≈ 1.5e7) far under CUDA's `2^31 - 1` grid-dimension limit. +const GPU_PAIR_CHUNK: usize = 3_900_000_000; /// Per-launch output-buffer budget in bytes (`NASSAU_GPU_BLOCK_MB`, default 512 MiB). /// From e6b85fd9190e3e617b5aed00c2b7959c234f2155 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Fri, 24 Jul 2026 01:05:29 -0400 Subject: [PATCH 13/16] milnor_gpu: make the Milnor basis resident on the GPU, upload term indices MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The batched multiply re-gathered and re-uploaded every term's zero-padded p-part (`term_pparts`, width*2 bytes/term) on every launch — the dominant per-launch H2D transfer, plus a large parallel host gather. Make the basis itself resident on the device instead: build it once (grown incrementally as higher degrees appear, mirroring the admissible master) and upload only `term_gei[slot]`, the term's global basis-element index `global_base[s_degree] + ti` (4 bytes/term). The kernel reads the p-part from `basis_pparts[gei*width..]` with length `basis_lens[gei]`. At stem 140 the per-launch term transfer drops from ~width/2x larger to term_gei=95 MB, the basis is a one-time few-MB upload, and the per-term p-part gather is gone. Kernel change is minimal: params `term_pparts, term_lens` -> `basis_pparts, basis_lens, term_gei` (net +1 array arg), launch-arg order preserved 1:1 with the signature. `multiply_pair` is unchanged. An A/B toggle (`NASSAU_GPU_BASIS_PASSTHROUGH=1`) binds the per-launch term buffers as the "basis" with an identity index map, reproducing the old behaviour through the new kernel — so a single binary can isolate a kernel-signature bug from a resident host/upload bug. Both paths verified GPU==CPU per launch at stem 80 (MIN_WORK=0), and the resident path chart-matches CPU at (100,152). Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 253 +++++++++++++++++-- 1 file changed, 226 insertions(+), 27 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 386d833c05..4424029774 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -130,6 +130,16 @@ fn gpu_stream_slots() -> u64 { *SLOTS } +/// A/B diagnostic toggle (`NASSAU_GPU_BASIS_PASSTHROUGH=1`): when set, the batched multiply +/// marshals each term's p-part per launch and binds those buffers as the "basis" with an +/// identity index map, reproducing the pre-resident-basis behaviour through the same kernel. +/// Lets a single binary isolate a kernel-signature bug from a resident-basis host/upload bug. +fn basis_passthrough() -> bool { + static ON: LazyLock = + LazyLock::new(|| std::env::var_os("NASSAU_GPU_BASIS_PASSTHROUGH").is_some()); + *ON +} + /// RAII reservation of `weight` bytes from [`GPU_BUDGET`] plus a round-robin stream slot; /// blocks (parked, not spinning) until the budget admits it. struct GpuPermit { @@ -322,6 +332,131 @@ fn resident_info(algebra: &MilnorAlgebra, p_part: &[PPartEntry]) -> RInfo { info } +/// Process-shared host master of the Milnor basis itself, laid out for the device. +/// +/// Every basis element's p-part is stored zero-padded to `width` at `pparts[gei*width ..]`, +/// where `gei` is the element's *global* index (elements concatenated in degree order: +/// all of degree 0, then degree 1, …). `lens[gei]` is its true (trimmed) p-part length, +/// and `global_base[d]` is the number of elements in degrees `< d`, so a term `(s_degree, +/// ti)` maps to `gei = global_base[s_degree] + ti`. +/// +/// This exists so a launch uploads only the small per-term *index* array (`term_gei`) +/// rather than re-gathering and re-uploading every term's padded p-part every launch — the +/// dominant per-launch H2D transfer. The basis is append-only and grows only when a higher +/// degree first appears, so it is uploaded to the device once and re-uploaded only on growth +/// (mirroring [`ResidentHost`]). `built_degree` is the highest degree fully appended. +#[derive(Default)] +struct ResidentBasisHost { + pparts: Vec, + lens: Vec, + global_base: Vec, + built_degree: i32, + width: usize, +} + +static RESIDENT_BASIS_HOST: LazyLock> = + LazyLock::new(|| RwLock::new(ResidentBasisHost::default())); + +/// Device mirror of [`ResidentBasisHost`]; `uploaded` is the element count on the device. +#[derive(Default)] +struct ResidentBasisDev { + pp_handle: Option, + ln_handle: Option, + uploaded: usize, +} + +static RESIDENT_BASIS_DEV: LazyLock> = + LazyLock::new(|| RwLock::new(ResidentBasisDev::default())); + +/// Serializes basis device *uploads* only (never handle reads); see [`RESIDENT_UPLOAD`]. +static RESIDENT_BASIS_UPLOAD: Mutex<()> = Mutex::new(()); + +/// Ensure the resident basis is built through `max_degree` and return a snapshot of +/// `global_base` (so callers compute `gei = global_base[s_degree] + ti` without holding the +/// lock during the parallel marshal). `width` is the fixed p-part padding stride. +/// +/// Append-only: only the first sight of each new degree takes the write lock, and the append +/// order fixes every element's `gei` forever. Basis enumeration is a pure function of the +/// algebra, so a first-sight race just recomputes identical bytes (the loser rechecks under +/// the write lock and appends nothing already present, since we extend strictly past +/// `built_degree`). +fn ensure_basis(algebra: &MilnorAlgebra, width: usize, max_degree: i32) -> Vec { + { + let host = RESIDENT_BASIS_HOST.read().unwrap(); + // `width != 0` distinguishes an initialized store from the derived-`Default` zero state + // (where `built_degree == 0` would spuriously claim degree 0 is already built). + if host.width != 0 && host.built_degree >= max_degree { + return host.global_base.clone(); + } + } + let mut host = RESIDENT_BASIS_HOST.write().unwrap(); + if host.width == 0 { + host.width = width; + host.built_degree = -1; // nothing built yet; the loop below starts at degree 0 + host.global_base.push(0); // global_base[0] = 0 elements before degree 0 + } + debug_assert_eq!(host.width, width, "basis padding width must be stable"); + for d in (host.built_degree + 1)..=max_degree { + let dim = algebra.dimension(d); + for i in 0..dim { + let elt = algebra.basis_element_from_index(d, i); + host.lens.push(elt.p_part.len() as u32); + let base = host.pparts.len(); + host.pparts.resize(base + width, 0); + for (slot, &v) in host.pparts[base..base + width].iter_mut().zip(&elt.p_part) { + *slot = narrow_u16(v); + } + } + // global_base[d+1] = total elements in degrees ≤ d. + let total = host.lens.len() as u32; + host.global_base.push(total); + } + host.built_degree = max_degree; + host.global_base.clone() +} + +/// Fetch the resident basis device handles `(pparts, lens)`, uploading the current host basis +/// only when the device copy does not yet cover `$need` elements. Upload runs outside +/// `RESIDENT_BASIS_DEV` (readers stay lock-free), serialized by `RESIDENT_BASIS_UPLOAD` with a +/// re-check to coalesce a burst of growth. The basis is append-only, so any handle with +/// `uploaded >= $need` is valid. A macro (not a fn) so the cubecl client type stays inferred, +/// exactly as [`resident_dev_handle`]. +macro_rules! basis_dev_handles { + ($client:expr, $need:expr) => {{ + let read_current = || { + let dev = RESIDENT_BASIS_DEV.read().unwrap(); + match (dev.uploaded >= $need, dev.pp_handle.clone(), dev.ln_handle.clone()) { + (true, Some(pp), Some(ln)) => Some((pp, ln)), + _ => None, + } + }; + match read_current() { + Some(h) => h, + None => { + let _guard = RESIDENT_BASIS_UPLOAD.lock().unwrap(); + match read_current() { + Some(h) => h, // another uploader already covered our need + None => { + let (pp, ln, elems) = { + let host = RESIDENT_BASIS_HOST.read().unwrap(); + ( + $client.create_from_slice(u16::as_bytes(&host.pparts)), + $client.create_from_slice(u32::as_bytes(&host.lens)), + host.lens.len(), + ) + }; + let mut dev = RESIDENT_BASIS_DEV.write().unwrap(); + dev.pp_handle = Some(pp.clone()); + dev.ln_handle = Some(ln.clone()); + dev.uploaded = elems; + (pp, ln) + } + } + } + } + }}; +} + /// Zero a device `u32` buffer on-device: `out[i] = 0`, one thread per limb. /// /// Initializes the batched multiply's XOR accumulator without allocating and uploading a host @@ -645,8 +780,9 @@ fn multiply_single_r_kernel( fn multiply_batch_kernel( col_sums: &Array, masks: &Array, - term_pparts: &Array, - term_lens: &Array, + basis_pparts: &Array, + basis_lens: &Array, + term_gei: &Array, g: &Array, xi: &Array, out: &mut Array>, @@ -694,17 +830,23 @@ fn multiply_batch_kernel( let cs_len = usize::cast_from(r_cs_len[ri]); let mk_len = usize::cast_from(r_mk_len[ri]); let term_slot = usize::cast_from(prod_term_start[p]) + t; + // `term_gei[term_slot]` is the term's *global* basis-element index (across all degrees): + // its (width-padded) p-part lives at `basis_pparts[gei*width ..]`, length `basis_lens[gei]`. + // The basis is resident on the device (uploaded once, grown incrementally), so a launch + // uploads only these indices instead of re-gathering every term's p-part — see + // [`ResidentBasisHost`]. + let gei = usize::cast_from(term_gei[term_slot]); multiply_pair( col_sums, masks, - term_pparts, + basis_pparts, g, xi, out, usize::cast_from(r_cs_offset[ri]) + m * cs_len, usize::cast_from(r_mk_offset[ri]) + m * mk_len, - term_slot * width, - usize::cast_from(term_lens[term_slot]), + gei * width, + usize::cast_from(basis_lens[gei]), cs_len, mk_len, usize::cast_from(prod_row_base[p]), @@ -957,30 +1099,68 @@ fn multiply_batch_block( off }; let total_terms = *term_off.last().unwrap(); - let mut term_pparts: Vec = vec![0u16; total_terms * width]; - let mut term_lens: Vec = vec![0u32; total_terms]; + + // Resident-basis path (the default): a term's p-part is not marshalled at all — it lives on + // the device (built once, grown incrementally). We upload only `term_gei[slot]`, the term's + // *global* basis-element index `global_base[s_degree] + ti`. Ensure the basis covers every + // `s_degree` in this block, then snapshot `global_base` so the parallel fill needs no lock. + let max_s_degree = products.iter().map(|p| p.s_degree).max().unwrap_or(0); + let global_base = ensure_basis(algebra, width, max_s_degree); + let mut term_gei: Vec = vec![0u32; total_terms]; { - let tp_base = term_pparts.as_mut_ptr() as usize; - let tl_base = term_lens.as_mut_ptr() as usize; + let tg_base = term_gei.as_mut_ptr() as usize; + let global_base = &global_base; (0..products.len()).into_maybe_par_iter().for_each(|pi| { let prod = &products[pi]; let (off, nt) = (term_off[pi], prod.term_indices.len()); // SAFETY: products write disjoint `[off, off + nt)` ranges (from the prefix sum), - // each within the allocated buffers, so no two tasks alias any element. The `usize` - // bases are re-formed into pointers here because raw pointers are not `Send`. - let tp = unsafe { + // each within the allocated buffer, so no two tasks alias any element. The `usize` + // base is re-formed into a pointer here because raw pointers are not `Send`. + let tg = unsafe { std::slice::from_raw_parts_mut((tg_base as *mut u32).add(off), nt) }; + let base = global_base[prod.s_degree as usize]; + for (k, &ti) in prod.term_indices.iter().enumerate() { + tg[k] = base + ti as u32; + } + }); + } + // Device need: the largest `gei` any term dereferences is `< global_base[max_s_degree + 1]` + // (all elements through degree `max_s_degree`), so uploading that many covers the block. + let need_basis_elems = global_base[max_s_degree as usize + 1] as usize; + + // A/B diagnostic (`NASSAU_GPU_BASIS_PASSTHROUGH=1`): bind the *per-launch* term buffers as the + // "basis" and set `term_gei` to the identity, so the new kernel reproduces the old behaviour + // bit-for-bit. If passthrough matches the CPU but the resident path does not, the bug is in + // the resident host/upload logic, not the kernel signature — and vice-versa. + let passthrough = basis_passthrough(); + let (mut term_pparts, mut term_lens): (Vec, Vec) = if passthrough { + let mut tp: Vec = vec![0u16; total_terms * width]; + let mut tl: Vec = vec![0u32; total_terms]; + let tp_base = tp.as_mut_ptr() as usize; + let tl_base = tl.as_mut_ptr() as usize; + (0..products.len()).into_maybe_par_iter().for_each(|pi| { + let prod = &products[pi]; + let (off, nt) = (term_off[pi], prod.term_indices.len()); + // SAFETY: disjoint per-product ranges, as above. + let tpp = unsafe { std::slice::from_raw_parts_mut((tp_base as *mut u16).add(off * width), nt * width) }; - let tl = unsafe { std::slice::from_raw_parts_mut((tl_base as *mut u32).add(off), nt) }; + let tll = unsafe { std::slice::from_raw_parts_mut((tl_base as *mut u32).add(off), nt) }; for (k, &ti) in prod.term_indices.iter().enumerate() { let elt = algebra.basis_element_from_index(prod.s_degree, ti); - tl[k] = elt.p_part.len() as u32; - for (slot, &v) in tp[k * width..(k + 1) * width].iter_mut().zip(&elt.p_part) { + tll[k] = elt.p_part.len() as u32; + for (slot, &v) in tpp[k * width..(k + 1) * width].iter_mut().zip(&elt.p_part) { *slot = narrow_u16(v); } } }); - } + // Identity indices, so the kernel's `gei*width` / `basis_lens[gei]` hit slot `term_slot`. + for (i, g) in term_gei.iter_mut().enumerate() { + *g = i as u32; + } + (tp, tl) + } else { + (Vec::new(), Vec::new()) + }; // Take the concurrency permit only now, with every rayon parallel section behind us: holding // it across the `per_prod` par_iter above deadlocks, because that par_iter's chunks execute on @@ -1052,13 +1232,12 @@ fn multiply_batch_block( let kb = |n: usize, sz: usize| n * sz / 1024; eprintln!( "[gpu-batch] rows={num_rows} cols={num_cols} products={} total_pairs={total_pairs} \ - out_len={out_len} | UPLOAD-KB: g={} xi={} term_pparts={} term_lens={} \ - prod_arrays={} pps={} | resident cs={} mk={}", + out_len={out_len} | UPLOAD-KB: g={} xi={} term_gei={} \ + prod_arrays={} pps={} | resident cs={} mk={} basis_elems={need_basis_elems}", products.len(), kb(g.len(), 4), kb(xi.len(), 4), - kb(term_pparts.len(), 2), - kb(term_lens.len(), 4), + kb(term_gei.len(), 4), kb(products.len() * 5, 4), kb(pps.len(), 4), kb(need_cs, 2), @@ -1069,10 +1248,15 @@ fn multiply_batch_block( return vec![vec![0u32; num_limbs]; num_rows]; } - // The resident `col_sums`/`masks` are non-empty once any `R` is present (guaranteed - // here, since `total_pairs > 0`); only `term_pparts` needs the non-empty guard. - if term_pparts.is_empty() { + // The resident `col_sums`/`masks` and basis are non-empty once any `R`/term is present + // (guaranteed here, since `total_pairs > 0`); only `term_gei` (and, in passthrough, the + // per-launch term buffers) needs the non-empty guard `create_from_slice` requires. + if term_gei.is_empty() { + term_gei.push(0); + } + if passthrough && term_pparts.is_empty() { term_pparts.push(0); + term_lens.push(0); } let marshal_ms = t_marshal.elapsed().as_secs_f64() * 1e3; @@ -1103,8 +1287,22 @@ fn multiply_batch_block( // Upload the block's data — term data, seqno/xi tables, per-`R` offsets, per-product // records, the pair prefix sum, and the (zeroed) output buffer — and launch once: the // caller has already bounded this block's pair count and output size. - let tp_h = client.create_from_slice(u16::as_bytes(&term_pparts)); - let tl_h = client.create_from_slice(u32::as_bytes(&term_lens)); + // Resident basis handles (default) or per-launch passthrough buffers (A/B diagnostic). + // `bp_len`/`bl_len` are the logical array lengths the kernel sees; every `gei` a thread + // dereferences is `< need_basis_elems`, so `need_basis_elems*width` / `need_basis_elems` + // cover it (the resident buffer may be larger — append-only — which is fine). + let (bp_h, bl_h, bp_len, bl_len) = if passthrough { + ( + client.create_from_slice(u16::as_bytes(&term_pparts)), + client.create_from_slice(u32::as_bytes(&term_lens)), + term_pparts.len(), + term_lens.len(), + ) + } else { + let (pp, ln) = basis_dev_handles!(client, need_basis_elems); + (pp, ln, need_basis_elems * width, need_basis_elems) + }; + let tg_h = client.create_from_slice(u32::as_bytes(&term_gei)); let g_h = client.create_from_slice(u32::as_bytes(&g)); let xi_h = client.create_from_slice(u32::as_bytes(&xi)); let rco_h = client.create_from_slice(u32::as_bytes(&r_cs_offset)); @@ -1139,8 +1337,9 @@ fn multiply_batch_block( CubeDim::new_1d(THREADS), ArrayArg::from_raw_parts(cs_h, cs_len_master), ArrayArg::from_raw_parts(mk_h, mk_len_master), - ArrayArg::from_raw_parts(tp_h, term_pparts.len()), - ArrayArg::from_raw_parts(tl_h, term_lens.len()), + ArrayArg::from_raw_parts(bp_h, bp_len), + ArrayArg::from_raw_parts(bl_h, bl_len), + ArrayArg::from_raw_parts(tg_h, term_gei.len()), ArrayArg::from_raw_parts(g_h, g.len()), ArrayArg::from_raw_parts(xi_h, xi.len()), ArrayArg::from_raw_parts(out_h.clone(), out_len), From 4efee670d5bd225e259d511ce185c9a1d0d4aacf Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Sat, 25 Jul 2026 00:05:15 -0400 Subject: [PATCH 14/16] milnor_gpu: safe multi-stream via in-place-grown resident globals MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `NASSAU_GPU=1` could not use multiple CUDA streams: the shared resident admissible master + Milnor basis were re-`create_from_slice`d into a NEW device handle on every growth, and that handle churn broke cubecl's per-handle cross-stream synchronization, so a launch on one stream could read the master while another was mid-upload -> wrong multiply (`dx != 0`), "Memory page" panic, or hang. Fix: make the resident buffers STABLE and grow them IN PLACE — the read-only shared global (model-weights) pattern cubecl supports across streams. A small `copy_into_*` kernel writes the new tail (uploaded to scratch via `create_from_slice`) at the buffer's append offset; the handle changes only on a rare capacity doubling, which is barrier-protected (`RESIDENT_REALLOC`: device sections hold the read lock across the multiply, a realloc takes the write lock and quiesces them). Each worker gets a stable per-thread stream id (`thread_stream_id`); default `NASSAU_GPU_STREAMS = 8`. Key gotcha (a stem-150 `dx != 0`): cubecl's `ArrayArg` length is u32, so a buffer of exactly 2^32 elements truncates to length 0 and the copy writes nothing (buffer reads all zeros). The doubling `cap` jumped 2^31 -> 2^32 right at stem 150's `masks` size. Capacity is clamped to `RESIDENT_MAX_CAP = 2^32 - 1`; a single resident buffer cannot exceed that (a larger master needs splitting — the old create_from_slice path had the same limit). Copies are chunked under the u32 ABSOLUTE_POS thread limit. Validated: VERIFY (GPU==CPU per launch) at 8 streams, chart-match to CPU, 0 dx-crashes over many stem-150 reps at 1 and 8 streams. Perf note: multi-stream is correct but ~neutral vs single-stream at stem 140/150 (the per-device runner serializes kernel submission); it may help at higher stems with a wider heavy-bidegree wavefront. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 381 +++++++++++++++---- 1 file changed, 311 insertions(+), 70 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 4424029774..767541249e 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -111,14 +111,15 @@ static GPU_BUDGET: LazyLock = LazyLock::new(|| GpuBudget { freed: Condvar::new(), }); -/// Fixed CUDA stream-slot count (`NASSAU_GPU_STREAMS`, default 8): device sections run under -/// `StreamId { value: counter % slots }`, so only this many streams (and hence retained -/// memory pools) ever exist. Without a pin every worker thread gets its own default stream, -/// and each stream's pool *retains* its freed slabs — ~100 worker streams each retaining -/// ~1 GB filled the whole 143 GB H200. Slots are round-robin and *shared*, not exclusive: -/// two small launches on one slot merely serialize on that CUDA stream (memory-safe, and -/// cheap for small kernels), so slot count does not cap concurrency the way the byte budget -/// does. Kept at 8: a 16-slot experiment collapsed >30x (unexplained cross-stream churn). +/// Number of distinct CUDA streams to spread device work over (`NASSAU_GPU_STREAMS`, default 8). +/// Each worker thread gets a stable per-thread stream id (see [`thread_stream_id`]), and this caps +/// how many distinct streams — hence retained memory pools — exist; `1` forces the single-stream +/// mode (all threads → stream 0). +/// +/// Multi-stream is SAFE because the shared resident master/basis are grown IN PLACE in stable +/// device buffers ([`ResidentDev`]) — the read-only "global weights" pattern cubecl supports across +/// streams. (The earlier churny design — re-`create_from_slice` a new handle on every growth — broke +/// cubecl's per-handle cross-stream sync and had to run single-stream; that is fixed.) fn gpu_stream_slots() -> u64 { static SLOTS: LazyLock = LazyLock::new(|| { std::env::var("NASSAU_GPU_STREAMS") @@ -130,6 +131,25 @@ fn gpu_stream_slots() -> u64 { *SLOTS } +/// A CUDA stream id UNIQUE to the calling worker thread (never shared or reused across threads — +/// cubecl's per-stream state assumes one driver thread per stream). Returns 0 for all threads when +/// `NASSAU_GPU_STREAMS == 1` (the single-stream fallback). With the resident master/basis now stable +/// device buffers grown IN PLACE (no handle churn — see [`ResidentDev`]), cross-stream reads of those +/// shared globals are the ordinary case cubecl supports, so distinct per-thread streams are safe and +/// give real device concurrency for the wide record-run wavefront. +fn thread_stream_id() -> u64 { + if gpu_stream_slots() == 1 { + return 0; + } + thread_local! { + static ID: u64 = { + static NEXT: AtomicU64 = AtomicU64::new(0); + NEXT.fetch_add(1, Ordering::Relaxed) + 1 // +1: id 0 is reserved for the single-stream mode + }; + } + ID.with(|&id| id) +} + /// A/B diagnostic toggle (`NASSAU_GPU_BASIS_PASSTHROUGH=1`): when set, the batched multiply /// marshals each term's p-part per launch and binds those buffers as the "basis" with an /// identity index map, reproducing the pre-resident-basis behaviour through the same kernel. @@ -140,26 +160,22 @@ fn basis_passthrough() -> bool { *ON } -/// RAII reservation of `weight` bytes from [`GPU_BUDGET`] plus a round-robin stream slot; -/// blocks (parked, not spinning) until the budget admits it. +/// RAII reservation of `weight` output bytes from [`GPU_BUDGET`]; blocks (parked, not spinning) +/// until the budget admits it. The stream is chosen per worker thread (see [`thread_stream_id`]), +/// so the permit no longer carries a slot. struct GpuPermit { weight: usize, - slot: u64, } impl GpuPermit { fn acquire(weight: usize) -> Self { - static NEXT_SLOT: AtomicU64 = AtomicU64::new(0); let b = &*GPU_BUDGET; let mut used = b.used.lock().unwrap(); while !(*used == 0 || *used + weight <= b.budget) { used = b.freed.wait(used).unwrap(); } *used += weight; - Self { - weight, - slot: NEXT_SLOT.fetch_add(1, Ordering::Relaxed) % gpu_stream_slots(), - } + Self { weight } } } @@ -241,23 +257,46 @@ struct ResidentHost { static RESIDENT_HOST: LazyLock> = LazyLock::new(|| RwLock::new(ResidentHost::default())); -/// Process-shared device mirror of the host master: one upload for the whole process, -/// re-uploaded only when the master grew. The handles are shared across worker threads and -/// stream slots — safe under cubecl 0.10's per-device runner (which serializes all server -/// access), with cross-stream reuse event-synced via the handle's origin-stream stamp. Reads -/// go through `RESIDENT_DEV.read()` (lock-free fan-out); uploads run outside that lock, -/// serialized only by `RESIDENT_UPLOAD` (see [`resident_dev_handle`]). +/// A resident device buffer that GROWS IN PLACE: a single stable handle, allocated once +/// (persistent, so per-launch cleanup never reindexes it) with headroom and written to at its +/// append offset as the host master grows — never re-`create_from_slice`d. The handle changes only +/// on a rare capacity doubling. `uploaded` is how many elements have been written; `cap` is the +/// allocated element capacity. +#[derive(Default)] +struct GrowBuf { + handle: Option, + cap: usize, + uploaded: usize, +} + +/// Process-shared device mirror of the host master. Each buffer is a STABLE handle grown in place +/// (see [`GrowBuf`], [`resident_dev_handle`]). This is what makes the master safe to share across +/// streams: the churny "re-upload a new handle on every growth" it replaced broke cubecl's +/// per-handle cross-stream sync (crash rate tracked re-upload frequency); a fixed handle written in +/// place is the ordinary shared-global (model-weights) pattern. Reads go through `RESIDENT_DEV.read()` +/// (lock-free fan-out); growth runs outside that lock, serialized only by `RESIDENT_UPLOAD`. #[derive(Default)] struct ResidentDev { - cs_handle: Option, - mk_handle: Option, - cs_uploaded: usize, - mk_uploaded: usize, + cs: GrowBuf, + mk: GrowBuf, } static RESIDENT_DEV: LazyLock> = LazyLock::new(|| RwLock::new(ResidentDev::default())); +/// Initial element capacity of a resident buffer. Doubling to grow past it is barrier-protected +/// (see [`RESIDENT_REALLOC`]) so it is correct, but each doubling briefly quiesces the device, so +/// the initial size carries headroom to keep them few (~a handful per run). u16 → ~512 MiB. +const RESIDENT_INIT_CAP: usize = 1 << 28; + +/// Maximum resident buffer capacity (elements). cubecl's `ArrayArg` length is a `u32`, so an array +/// of exactly `2^32` elements has its length truncated to 0 — the copy then writes nothing and the +/// buffer reads as all zeros (the stem-150 `dx != 0`, hit when the doubling `cap` jumped 2^31 → 2^32). +/// Capacity is clamped just below the limit. A single resident buffer therefore cannot exceed +/// `2^32 - 1` elements; a master/basis larger than that (very high stems) must be split across +/// buffers — the same fundamental cubecl limit the old full-`create_from_slice` path also had. +const RESIDENT_MAX_CAP: usize = (1 << 32) - 1; + /// Serializes master device *uploads* only — never handle reads. A launch that must grow the /// device master takes this before uploading, so at a growth point at most one multi-GB /// `create_from_slice` runs (others re-check and find it already done) instead of every launch @@ -266,17 +305,27 @@ static RESIDENT_DEV: LazyLock> = /// collapsed the whole wavefront to one memcpy-ing thread). static RESIDENT_UPLOAD: Mutex<()> = Mutex::new(()); +/// Guards the rare capacity-doubling REALLOC (the only time a resident handle changes) against +/// concurrent readers. A device section holds the read lock while its multiply kernel is reading the +/// resident buffers; a realloc takes the write lock, so it waits for every in-flight reader to drain +/// and blocks new ones — the swap + old-buffer free then happens with the device quiesced, which the +/// churny path could not guarantee (residual ~1/10 crash at the readback / a stale multiply). The +/// common in-place delta write does NOT take this lock (it's stable-handle, proven race-free), so the +/// hot path stays lock-free; reallocs happen only a handful of times per run, so the barrier is free. +static RESIDENT_REALLOC: RwLock<()> = RwLock::new(()); + /// Fetch a resident device-master handle, uploading the current master prefix only when this /// block dereferences past what is already on the device (`need`). The upload runs OUTSIDE /// `RESIDENT_DEV` (readers stay lock-free; only concurrent uploaders serialize, on /// `RESIDENT_UPLOAD`, and a re-check coalesces a burst of growth-needing launches into one /// upload). The master is append-only, so any handle with `uploaded >= need` is valid. macro_rules! resident_dev_handle { - ($client:expr, $need:expr, $handle:ident, $uploaded:ident, $host_vec:ident) => {{ + ($client:expr, $need:expr, $buf:ident, $host_vec:ident) => {{ + // Lock-free fast path: the stable handle already covers `$need`. let read_current = || { let dev = RESIDENT_DEV.read().unwrap(); - match (dev.$uploaded >= $need, dev.$handle.clone()) { - (true, Some(h)) => Some((h, dev.$uploaded)), + match (dev.$buf.uploaded >= $need, dev.$buf.handle.clone()) { + (true, Some(h)) => Some((h, dev.$buf.uploaded)), _ => None, } }; @@ -285,20 +334,68 @@ macro_rules! resident_dev_handle { None => { let _upload_guard = RESIDENT_UPLOAD.lock().unwrap(); match read_current() { - Some(hu) => hu, // another uploader already covered our need + Some(hu) => hu, // another grower already covered our need None => { - let (handle, len) = { - let host = RESIDENT_HOST.read().unwrap(); - let len = host.$host_vec.len(); - ( - $client.create_from_slice(u16::as_bytes(&host.$host_vec[..len])), - len, - ) + let (old_handle, cap, uploaded) = { + let dev = RESIDENT_DEV.read().unwrap(); + (dev.$buf.handle.clone(), dev.$buf.cap, dev.$buf.uploaded) }; - let mut dev = RESIDENT_DEV.write().unwrap(); - dev.$handle = Some(handle.clone()); - dev.$uploaded = len; - (handle, len) + let host_len = RESIDENT_HOST.read().unwrap().$host_vec.len(); + + // (Re)allocate a stable persistent buffer on first use or capacity overflow + // (rare: `RESIDENT_INIT_CAP` has headroom and the master saturates early), + // copying existing device data on-device. Only THIS handle-swap needs a sync, + // so a fresh-handle cross-stream read never sees an incomplete copy. + let (handle, cap) = if old_handle.is_none() || host_len > cap { + // Quiesce readers for the handle swap (see [`RESIDENT_REALLOC`]). + let _realloc_w = RESIDENT_REALLOC.write().unwrap(); + let new_cap = host_len.max(cap * 2).max(RESIDENT_INIT_CAP) + .min(RESIDENT_MAX_CAP); + let new_handle = $client.empty(new_cap * ::core::mem::size_of::()); + if let Some(oh) = &old_handle { + if uploaded > 0 { + copy_chunked!( + $client, copy_into_u16, oh, uploaded, 0usize, + new_handle, new_cap, 0usize, uploaded + ); + } + } + let _ = cubecl_common::reader::read_sync($client.sync()); + (new_handle, new_cap) + } else { + (old_handle.unwrap(), cap) + }; + + // Write the new tail host[uploaded..host_len] into the STABLE buffer at its + // append offset via a scratch upload + copy kernel, then SYNC before + // publishing: cubecl does not cross-stream-order a *kernel write* to a shared + // buffer against reads the way it does `create_from_slice`, so a reader on + // another worker's stream could otherwise see the tail mid-copy (harmless at + // small sizes, corrupting once the copy is GB-scale — the stem-150 `dx != 0`). + // Blocking here makes the new prefix physically resident before any reader can + // observe the bumped `uploaded`. + if host_len > uploaded { + let n = host_len - uploaded; + let scratch = { + let host = RESIDENT_HOST.read().unwrap(); + $client.create_from_slice(u16::as_bytes( + &host.$host_vec[uploaded..host_len], + )) + }; + copy_chunked!( + $client, copy_into_u16, scratch, n, 0usize, + handle, cap, uploaded, n + ); + let _ = cubecl_common::reader::read_sync($client.sync()); + } + + { + let mut dev = RESIDENT_DEV.write().unwrap(); + dev.$buf.handle = Some(handle.clone()); + dev.$buf.cap = cap; + dev.$buf.uploaded = host_len; + } + (handle, host_len) } } } @@ -357,12 +454,13 @@ struct ResidentBasisHost { static RESIDENT_BASIS_HOST: LazyLock> = LazyLock::new(|| RwLock::new(ResidentBasisHost::default())); -/// Device mirror of [`ResidentBasisHost`]; `uploaded` is the element count on the device. +/// Device mirror of [`ResidentBasisHost`], both buffers grown IN PLACE (see [`GrowBuf`], +/// [`resident_dev_handle`]). `pp` holds the width-padded p-parts (`elems * width` u16), `ln` the +/// lengths (`elems` u32); the basis element count is `ln.uploaded`. #[derive(Default)] struct ResidentBasisDev { - pp_handle: Option, - ln_handle: Option, - uploaded: usize, + pp: GrowBuf, + ln: GrowBuf, } static RESIDENT_BASIS_DEV: LazyLock> = @@ -423,9 +521,11 @@ fn ensure_basis(algebra: &MilnorAlgebra, width: usize, max_degree: i32) -> Vec {{ + // Lock-free fast path: the stable buffers already cover `$need` basis elements + // (`ln.uploaded` is the element count). let read_current = || { let dev = RESIDENT_BASIS_DEV.read().unwrap(); - match (dev.uploaded >= $need, dev.pp_handle.clone(), dev.ln_handle.clone()) { + match (dev.ln.uploaded >= $need, dev.pp.handle.clone(), dev.ln.handle.clone()) { (true, Some(pp), Some(ln)) => Some((pp, ln)), _ => None, } @@ -435,21 +535,98 @@ macro_rules! basis_dev_handles { None => { let _guard = RESIDENT_BASIS_UPLOAD.lock().unwrap(); match read_current() { - Some(h) => h, // another uploader already covered our need + Some(h) => h, // another grower already covered our need None => { - let (pp, ln, elems) = { - let host = RESIDENT_BASIS_HOST.read().unwrap(); - ( - $client.create_from_slice(u16::as_bytes(&host.pparts)), - $client.create_from_slice(u32::as_bytes(&host.lens)), - host.lens.len(), - ) + let width = RESIDENT_BASIS_HOST.read().unwrap().width; + let elems = RESIDENT_BASIS_HOST.read().unwrap().lens.len(); + let pp_len = elems * width; // u16 count + let ln_len = elems; // u32 count + + // Grow the width-padded p-parts buffer (u16) in place. + let (old_pp, pp_cap, pp_up) = { + let d = RESIDENT_BASIS_DEV.read().unwrap(); + (d.pp.handle.clone(), d.pp.cap, d.pp.uploaded) + }; + let (pp_h, pp_cap) = if old_pp.is_none() || pp_len > pp_cap { + let _realloc_w = RESIDENT_REALLOC.write().unwrap(); + let new_cap = pp_len.max(pp_cap * 2).max(RESIDENT_INIT_CAP) + .min(RESIDENT_MAX_CAP); + let nh = $client.empty(new_cap * ::core::mem::size_of::()); + if let Some(oh) = &old_pp { + if pp_up > 0 { + copy_chunked!( + $client, copy_into_u16, oh, pp_up, 0usize, + nh, new_cap, 0usize, pp_up + ); + } + } + let _ = cubecl_common::reader::read_sync($client.sync()); + (nh, new_cap) + } else { + (old_pp.unwrap(), pp_cap) }; - let mut dev = RESIDENT_BASIS_DEV.write().unwrap(); - dev.pp_handle = Some(pp.clone()); - dev.ln_handle = Some(ln.clone()); - dev.uploaded = elems; - (pp, ln) + if pp_len > pp_up { + let n = pp_len - pp_up; + let scratch = { + let host = RESIDENT_BASIS_HOST.read().unwrap(); + $client.create_from_slice(u16::as_bytes(&host.pparts[pp_up..pp_len])) + }; + copy_chunked!( + $client, copy_into_u16, scratch, n, 0usize, + pp_h, pp_cap, pp_up, n + ); + } + + // Grow the lengths buffer (u32) in place. + let (old_ln, ln_cap, ln_up) = { + let d = RESIDENT_BASIS_DEV.read().unwrap(); + (d.ln.handle.clone(), d.ln.cap, d.ln.uploaded) + }; + let (ln_h, ln_cap) = if old_ln.is_none() || ln_len > ln_cap { + let _realloc_w = RESIDENT_REALLOC.write().unwrap(); + let new_cap = ln_len.max(ln_cap * 2).max(RESIDENT_INIT_CAP) + .min(RESIDENT_MAX_CAP); + let nh = $client.empty(new_cap * ::core::mem::size_of::()); + if let Some(oh) = &old_ln { + if ln_up > 0 { + copy_chunked!( + $client, copy_into_u32, oh, ln_up, 0usize, + nh, new_cap, 0usize, ln_up + ); + } + } + let _ = cubecl_common::reader::read_sync($client.sync()); + (nh, new_cap) + } else { + (old_ln.unwrap(), ln_cap) + }; + if ln_len > ln_up { + let n = ln_len - ln_up; + let scratch = { + let host = RESIDENT_BASIS_HOST.read().unwrap(); + $client.create_from_slice(u32::as_bytes(&host.lens[ln_up..ln_len])) + }; + copy_chunked!( + $client, copy_into_u32, scratch, n, 0usize, + ln_h, ln_cap, ln_up, n + ); + } + + // Sync both in-place copies before publishing (see [`resident_dev_handle`]): + // a kernel write to the shared basis buffers must be physically done before a + // reader on another worker's stream can observe the bumped element count. + let _ = cubecl_common::reader::read_sync($client.sync()); + + { + let mut dev = RESIDENT_BASIS_DEV.write().unwrap(); + dev.pp.handle = Some(pp_h.clone()); + dev.pp.cap = pp_cap; + dev.pp.uploaded = pp_len; + dev.ln.handle = Some(ln_h.clone()); + dev.ln.cap = ln_cap; + dev.ln.uploaded = ln_len; + } + (pp_h, ln_h) } } } @@ -472,6 +649,60 @@ fn zero_u32(out: &mut Array) { } } +/// Copy `count` elements `src[src_off + i] -> dst[dst_off + i]`, one thread per element. Used to grow +/// the resident master/basis IN PLACE — new data is uploaded to a scratch buffer and copied into the +/// stable resident buffer at its append offset, so the resident device handle never changes (no +/// re-`create_from_slice` churn that would break cross-stream sync). +/// +/// Offsets are `usize`, and `count` bounds this launch so the caller can split a copy larger than the +/// `u32` `ABSOLUTE_POS` thread limit into chunks (a resident buffer exceeds 2^32 u16 elements around +/// stem 150 — the width-padded basis p-parts especially). See [`copy_into_chunked`]. +#[cube(launch)] +fn copy_into_u16(src: &Array, dst: &mut Array, src_off: usize, dst_off: usize, count: u32) { + if ABSOLUTE_POS < usize::cast_from(count) { + dst[dst_off + ABSOLUTE_POS] = src[src_off + ABSOLUTE_POS]; + } +} + +/// `u32` sibling of [`copy_into_u16`] (for the resident basis `lens`). +#[cube(launch)] +fn copy_into_u32(src: &Array, dst: &mut Array, src_off: usize, dst_off: usize, count: u32) { + if ABSOLUTE_POS < usize::cast_from(count) { + dst[dst_off + ABSOLUTE_POS] = src[src_off + ABSOLUTE_POS]; + } +} + +/// Elements per copy-kernel launch: below the kernel's `u32` `ABSOLUTE_POS` thread limit, so copies +/// of multi-billion-element resident buffers are split into this many at a time. +const COPY_CHUNK: usize = 1 << 30; + +/// Copy `count` elements `src[src_off..] -> dst[dst_off..]` with `$kernel` (`copy_into_u16`/`_u32`), +/// splitting into [`COPY_CHUNK`]-element launches so counts past the `u32` thread limit are handled. +/// `$src_len`/`$dst_len` are the logical array lengths passed to the kernel (must cover the ranges). +macro_rules! copy_chunked { + ($client:expr, $kernel:ident, $src:expr, $src_len:expr, $src_off:expr, + $dst:expr, $dst_len:expr, $dst_off:expr, $count:expr) => {{ + const CT: u32 = 256; + let mut done: usize = 0; + while done < $count { + let n = ($count - done).min(COPY_CHUNK); + unsafe { + $kernel::launch::( + &$client, + CubeCount::Static((n as u32).div_ceil(CT), 1, 1), + CubeDim::new_1d(CT), + ArrayArg::from_raw_parts($src.clone(), $src_len), + ArrayArg::from_raw_parts($dst.clone(), $dst_len), + $src_off + done, + $dst_off + done, + n as u32, + ); + } + done += n; + } + }}; +} + /// Elementwise F₂ addition of two bit-packed vectors: `out[i] = a[i] ^ b[i]`. /// /// One thread per `u32` limb. F₂ addition is XOR of the packed limbs, so this is @@ -1171,7 +1402,9 @@ fn multiply_batch_block( // pre-pass already enumerated every `R`), and the device section never enters rayon — so // every permit holder makes progress and stolen jobs waiting for a permit wake in finite // time (priority inversion at worst, never deadlock). - let permit = GpuPermit::acquire(num_rows * num_limbs * 4); + // Held for the device section (RAII): bounds total in-flight output bytes across workers. The + // stream is now chosen per-thread (see [`thread_stream_id`]), not from the permit. + let _permit = GpuPermit::acquire(num_rows * num_limbs * 4); // Per-`R` offsets into the shared resident master (see [`ResidentHost`]). All read-lock // cache hits: the caller's pair-count pre-pass already enumerated every `R` in this block. // `need_cs`/`need_mk` track the furthest master offset this block dereferences, so the @@ -1263,12 +1496,14 @@ fn multiply_batch_block( let t_device = std::time::Instant::now(); - // Device section pinned to this permit's stream slot: up to `NASSAU_GPU_CONCURRENCY` - // launches overlap on distinct streams, but no more streams (and hence retained pools) - // than that ever exist — see [`GpuPermit`]. Cubecl's per-device runner serializes the - // actual server access; cross-slot/cross-thread reuse of the shared resident handles is - // event-synced by cubecl. `memory_cleanup` below trims only this slot's own pool. - let result = StreamId { value: permit.slot }.executes(|| { + // Device section on this worker's own stable stream (see [`thread_stream_id`]): up to + // `gpu_stream_slots()` distinct streams run concurrently. Cubecl's per-device runner serializes + // the actual server access; the shared resident master/basis (stable in-place-grown buffers) are + // read cross-stream safely. `memory_cleanup` below trims only this stream's own transient pool. + let result = StreamId { + value: thread_stream_id(), + } + .executes(|| { let client = CudaRuntime::client(&CudaDevice::default()); // Shared resident admissible buffers (see [`ResidentDev`]): re-upload the master ONLY // when this block dereferences past the uploaded prefix (`need_cs` / `need_mk`). The @@ -1280,10 +1515,8 @@ fn multiply_batch_block( // then HOST.read, and nothing under either lock blocks on rayon or a permit. The // master is append-only, so the uploaded prefix is always a prefix of the current // host master and every offset `< uploaded` is final. - let (cs_h, cs_len_master) = - resident_dev_handle!(client, need_cs, cs_handle, cs_uploaded, col_sums); - let (mk_h, mk_len_master) = - resident_dev_handle!(client, need_mk, mk_handle, mk_uploaded, masks); + let (cs_h, cs_len_master) = resident_dev_handle!(client, need_cs, cs, col_sums); + let (mk_h, mk_len_master) = resident_dev_handle!(client, need_mk, mk, masks); // Upload the block's data — term data, seqno/xi tables, per-`R` offsets, per-product // records, the pair prefix sum, and the (zeroed) output buffer — and launch once: the // caller has already bounded this block's pair count and output size. @@ -1310,6 +1543,12 @@ fn multiply_batch_block( let rcl_h = client.create_from_slice(u32::as_bytes(&r_cs_len)); let rml_h = client.create_from_slice(u32::as_bytes(&r_mk_len)); const THREADS: u32 = 256; + // Hold the resident READ lock across the multiply that dereferences the resident buffers, + // through the readback that syncs it (see [`RESIDENT_REALLOC`]). All resident growth for this + // block is already done above; a concurrent capacity realloc on another thread waits here for + // this multiply to finish, so the resident handle can never be swapped mid-kernel. Held to + // the end of the device section (dropped after `read_one`). + let _realloc_guard = RESIDENT_REALLOC.read().unwrap(); // Allocate the XOR accumulator uninitialized and zero it on-device (see [`zero_u32`]), // instead of uploading a hundreds-of-MB host zero buffer — the former dominant serial // marshaling cost. Same stream as the multiply below, so it is ordered before it. @@ -1368,6 +1607,8 @@ fn multiply_batch_block( // launch, so CubeCL's pool cannot reuse the slab and would accumulate them until // the 4 GB card OOMs. Return the freed memory to the driver each launch; the // resident admissible handles stay alive (refcount > 0) so cleanup skips them. + // DIAGNOSTIC: memory_cleanup removed to test whether its explicit pool reindex corrupts the + // stable resident buffers under concurrent readers. (Memory bounding restored after.) client.memory_cleanup(); result From 0c1712a334d18eaf9f197f23fd0c38541463aa2a Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Sat, 25 Jul 2026 01:15:24 -0400 Subject: [PATCH 15/16] fp: trace GPU vs CPU row-reduce dispatch (fp::rr) Instrument Matrix::row_reduce to emit a `fp::rr` tracing event for every p=2 reduction with min(rows,cols) >= 1024, recording rows/cols/min and whether the device RREF was taken (path="gpu") or it fell back to CPU M4RI (path="cpu"). The event inherits the active nassau span so each line carries its bidegree/signature context. Confirms on a stem-150 run that every reduction with min >= 8192 (up to 25091x30275, incl. the heavy zero-signature base solves) dispatches to the GPU with no fallbacks; everything below the 8192 FP_CUDA_RR_THRESHOLD stays on CPU. tracing is added as a gpu-gated optional dep. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/fp/Cargo.toml | 4 ++- ext/crates/fp/src/matrix/matrix_inner.rs | 42 +++++++++++++++++++++--- 2 files changed, 41 insertions(+), 5 deletions(-) diff --git a/ext/crates/fp/Cargo.toml b/ext/crates/fp/Cargo.toml index 52dd296e2a..4f5b106e06 100644 --- a/ext/crates/fp/Cargo.toml +++ b/ext/crates/fp/Cargo.toml @@ -17,6 +17,8 @@ paste = "1.0.15" proptest = { version = "1.7", optional = true } serde = { version = "1.0.219", features = ["derive"] } serde_json = "1.0.141" +# Only used for the GPU row-reduce dispatch instrumentation; pulled in by `gpu`. +tracing = { version = "0.1.41", optional = true } maybe-rayon = { path = "../maybe-rayon" } query = { path = "../query" } @@ -45,7 +47,7 @@ default = ["odd-primes"] concurrent = ["maybe-rayon/concurrent"] odd-primes = [] # Dispatch large p=2 matrix products to the Hopper GPU backend (`fp-cuda`). -gpu = ["dep:fp-cuda"] +gpu = ["dep:fp-cuda", "dep:tracing"] [[bench]] name = "mul" diff --git a/ext/crates/fp/src/matrix/matrix_inner.rs b/ext/crates/fp/src/matrix/matrix_inner.rs index 99b014e631..d6ca2af0eb 100644 --- a/ext/crates/fp/src/matrix/matrix_inner.rs +++ b/ext/crates/fp/src/matrix/matrix_inner.rs @@ -677,11 +677,45 @@ impl Matrix { // For large p = 2 matrices, try the device-resident GPU reduction; it // produces the identical canonical RREF + pivots. Falls back to the CPU // M4RI path below when the GPU is unavailable or below threshold. + // + // Instrumentation: for every p=2 reduction of a non-trivial matrix + // (min(rows,cols) >= 1024) we emit a `fp::rr` tracing event recording the + // dimensions and whether the GPU path was taken (`path="gpu"`) or it fell + // back to CPU M4RI (`path="cpu"` — either below the 8192 threshold or a + // launch failure; the logged dims disambiguate). The event inherits the + // active nassau span, so it carries the bidegree/signature context. #[cfg(feature = "gpu")] - if p == 2 - && let Some(rank) = crate::blas::cuda::try_row_reduce(self) - { - return rank; + if p == 2 { + let (rr_rows, rr_cols) = (self.rows(), self.columns()); + let rr_big = rr_rows.min(rr_cols) >= 1024; + match crate::blas::cuda::try_row_reduce(self) { + Some(rank) => { + if rr_big { + tracing::info!( + target: "fp::rr", + rows = rr_rows, + cols = rr_cols, + min = rr_rows.min(rr_cols), + path = "gpu", + rank, + "row_reduce" + ); + } + return rank; + } + None => { + if rr_big { + tracing::info!( + target: "fp::rr", + rows = rr_rows, + cols = rr_cols, + min = rr_rows.min(rr_cols), + path = "cpu", + "row_reduce" + ); + } + } + } } self.initialize_pivots(); From f22f8a73058c67971cf176dbc143c8c076562c81 Mon Sep 17 00:00:00 2001 From: Joey Beauvais-Feisthauer Date: Sun, 26 Jul 2026 04:59:24 -0400 Subject: [PATCH 16/16] fp,algebra: serialize cubecl multiply vs cooperative GPU row-reduce The intermittent stem-150 hang (flat sm=100%, ~1/6 of runs, always on the zero-signature base solve) was a deadlock in the fp-cuda cooperative row-reduce kernel `panel_factor_coop`. Its grid-wide spin barrier (launched via `cuLaunchCooperativeKernel`) requires ALL its CTAs co-resident, but the algebra Milnor multiply runs on a *separate* CUDA runtime (cubecl) whose kernels concurrently occupy SMs. When a cooperative row-reduce launched while cubecl multiply kernels were resident, its CTAs could not all co-reside; the missing ones never reached the barrier and the resident ones spun forever. Fix: a cross-runtime `fp::GPU_EXCLUSIVE` RwLock. The cooperative row-reduce takes the write lock (drains in-flight cubecl multiplies, then runs with the GPU to itself) around its launch+download; every cubecl multiply takes the read lock across its whole device section (launch through readback, so releasing means the kernel has actually completed). Readers run concurrently; a pending row-reduce briefly excludes them. No lock cycle (cubecl never takes the fp-cuda ctx lock) and no same-thread read->write (a solve's multiply releases before its row-reduce). Also adds a `gpu_row_reduce` tracing span around the GPU reduce -- the diagnostic that localized the wedge (an unclosed span names a stuck reduce, distinguishing an RREF hang from a multiply hang). Validated on H200: stem-150 x16 with 0 wedges (baseline ~1/6), wall time unchanged (261-359s), GPU-vs-CPU chart match preserved. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01SmhK9czwJj2dZeemDhXtBf --- ext/crates/algebra/src/algebra/milnor_gpu.rs | 9 +++++++-- ext/crates/fp/src/blas/cuda.rs | 20 +++++++++++++++++++- ext/crates/fp/src/lib.rs | 6 ++++++ ext/crates/fp/src/matrix/matrix_inner.rs | 8 ++++++++ 4 files changed, 40 insertions(+), 3 deletions(-) diff --git a/ext/crates/algebra/src/algebra/milnor_gpu.rs b/ext/crates/algebra/src/algebra/milnor_gpu.rs index 767541249e..54baff8871 100644 --- a/ext/crates/algebra/src/algebra/milnor_gpu.rs +++ b/ext/crates/algebra/src/algebra/milnor_gpu.rs @@ -1505,6 +1505,13 @@ fn multiply_batch_block( } .executes(|| { let client = CudaRuntime::client(&CudaDevice::default()); + // Hold the cross-runtime GPU read lock for this whole device section (H2D + kernels + + // readback), so a concurrent cooperative fp-cuda row-reduce (raw cudarc) waits for this + // cubecl multiply to finish before taking the GPU exclusively. Without this, the RREF's + // `cuLaunchCooperativeKernel` grid barrier can't co-reside its CTAs against a resident + // multiply kernel and spins forever (the intermittent stem-150 wedge). See + // [`fp::GPU_EXCLUSIVE`]. Dropped at the end of the closure, after `read_one`. + let _gpu_shared = fp::GPU_EXCLUSIVE.read().unwrap_or_else(|e| e.into_inner()); // Shared resident admissible buffers (see [`ResidentDev`]): re-upload the master ONLY // when this block dereferences past the uploaded prefix (`need_cs` / `need_mk`). The // master grows continually at the frontier, so re-uploading on mere growth ships @@ -1607,8 +1614,6 @@ fn multiply_batch_block( // launch, so CubeCL's pool cannot reuse the slab and would accumulate them until // the 4 GB card OOMs. Return the freed memory to the driver each launch; the // resident admissible handles stay alive (refcount > 0) so cleanup skips them. - // DIAGNOSTIC: memory_cleanup removed to test whether its explicit pool reindex corrupts the - // stable resident buffers under concurrent readers. (Memory bounding restored after.) client.memory_cleanup(); result diff --git a/ext/crates/fp/src/blas/cuda.rs b/ext/crates/fp/src/blas/cuda.rs index 0fba7154ec..550b20c99d 100644 --- a/ext/crates/fp/src/blas/cuda.rs +++ b/ext/crates/fp/src/blas/cuda.rs @@ -14,12 +14,26 @@ //! CPU path is used. Defaults to 2048; the GPU only wins once the kernel work //! dwarfs the H2D/D2H + TMA-layout marshalling, which dominates small sizes. -use std::sync::{Mutex, OnceLock}; +use std::sync::{Mutex, OnceLock, RwLock}; use fp_cuda::GpuContext; use crate::{matrix::Matrix, prime::TWO}; +/// Serializes the fp-cuda **cooperative** row-reduce (raw cudarc) against the algebra +/// **cubecl** Milnor multiply, which share the one physical GPU from different CUDA +/// runtimes. The row-reduce's `cuLaunchCooperativeKernel` (panel_factor_coop's grid-wide +/// spin barrier) needs *all* its CTAs co-resident; a cubecl multiply kernel occupying SMs +/// at that moment prevents co-residency, so the missing CTAs never reach the barrier and +/// the resident ones spin forever (flat sm=100% — the intermittent stem-150 wedge). +/// +/// Contract: the cooperative row-reduce takes the **write** lock (exclusive GPU) around its +/// launch+download; every cubecl multiply takes a **read** lock held from launch through its +/// readback (so releasing means that kernel has actually completed, not merely enqueued). +/// Readers run concurrently with each other; a pending row-reduce drains them, runs alone, +/// then lets them resume. See `algebra::algebra::milnor_gpu` for the read side. +pub static GPU_EXCLUSIVE: RwLock<()> = RwLock::new(()); + /// Smallest `min(m, k, n)` for which we attempt the GPU matmul. Below this the /// host marshalling (bit-repack into TMA tiles + copies) costs more than it saves. const DEFAULT_THRESHOLD: usize = 2048; @@ -127,7 +141,11 @@ pub(crate) fn try_row_reduce(m: &mut Matrix) -> Option { let stride = cols.div_ceil(64); let limbs = to_limbs(m); + // Exclusive GPU for the cooperative reduce: block new cubecl multiplies and wait for + // in-flight ones to complete, so panel_factor_coop's grid-wide barrier can co-reside all + // its CTAs (see [`GPU_EXCLUSIVE`]). Held across launch + download; ordered before ctx.lock. let (dev_limbs, perm, r, pivot_cols) = { + let _exclusive = GPU_EXCLUSIVE.write().unwrap_or_else(|e| e.into_inner()); let gpu = ctx.lock().ok()?; let mut dm = gpu.upload(&limbs, rows, cols).ok()?; let (perm, r, pivot_cols) = gpu.row_reduce_dev(&mut dm).ok()?; diff --git a/ext/crates/fp/src/lib.rs b/ext/crates/fp/src/lib.rs index 8d971da2a8..bb8dedec60 100644 --- a/ext/crates/fp/src/lib.rs +++ b/ext/crates/fp/src/lib.rs @@ -12,6 +12,12 @@ pub mod vector; pub mod blas; +/// Cross-runtime GPU serialization lock (cubecl multiply vs. cooperative fp-cuda row-reduce). +/// Re-exported so `algebra`'s Milnor-multiply GPU path can take the read side. See +/// [`blas::cuda::GPU_EXCLUSIVE`]. +#[cfg(feature = "gpu")] +pub use blas::cuda::GPU_EXCLUSIVE; + pub(crate) mod simd; // This is useful for traits that want to implement `Arbitrary`. This lets us specify that they diff --git a/ext/crates/fp/src/matrix/matrix_inner.rs b/ext/crates/fp/src/matrix/matrix_inner.rs index d6ca2af0eb..fcf3261923 100644 --- a/ext/crates/fp/src/matrix/matrix_inner.rs +++ b/ext/crates/fp/src/matrix/matrix_inner.rs @@ -688,6 +688,14 @@ impl Matrix { if p == 2 { let (rr_rows, rr_cols) = (self.rows(), self.columns()); let rr_big = rr_rows.min(rr_cols) >= 1024; + // Wrap the GPU reduce in an ENTERED span (not just a completion event) so a wedge + // *inside* `try_row_reduce` leaves an open `gpu_row_reduce` span with no `close` in the + // log — distinguishing an RREF hang from a milnor-multiply hang. Inherits the active + // nassau step span, so it carries the bidegree/signature. Dropped on return/fallthrough. + let _rr_span = rr_big.then(|| { + tracing::info_span!(target: "fp::rr", "gpu_row_reduce", rows = rr_rows, cols = rr_cols) + .entered() + }); match crate::blas::cuda::try_row_reduce(self) { Some(rank) => { if rr_big {