diff --git a/crypto/math-cuda/src/device.rs b/crypto/math-cuda/src/device.rs index 8bc140f21..a7c129cc8 100644 --- a/crypto/math-cuda/src/device.rs +++ b/crypto/math-cuda/src/device.rs @@ -509,6 +509,12 @@ impl Backend { /// Map `rayon::current_thread_index()` to a slot index, with a defensive /// clamp in case the rayon pool grew past the Vec we sized at init. + /// + /// The per-table scheduler's driver threads are not rayon workers: they + /// all resolve to slot 0 and deliberately share one slab. Spreading them + /// over per-driver slots costs more in repeated pinned allocation than + /// the shared mutex does — the staged transfers are already hidden by + /// cross-table overlap. fn worker_slot(&self, len: usize) -> usize { let idx = rayon::current_thread_index().unwrap_or(0); // Should be unreachable with rayon's fixed default pool, but if a diff --git a/crypto/stark/src/instruments.rs b/crypto/stark/src/instruments.rs index 21866c465..790c06cb0 100644 --- a/crypto/stark/src/instruments.rs +++ b/crypto/stark/src/instruments.rs @@ -6,12 +6,20 @@ use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; // Wall clock span timeline: the trustworthy per step latency breakdown. // -// Spans open and close on the main thread at phase boundaries. They do not -// overlap and sum to their parent, so the tree is a true latency breakdown -// (unlike the accum_* thread local sub timers below, which sum per worker CPU -// time across rayon threads and can exceed 100%). A parallel region is one span -// around the blocking call; its internal split is reported separately as CPU -// time, never mixed into the wall tree. +// Top level phase spans open and close on the main thread at phase boundaries. +// They do not overlap and sum to their parent, so that part of the tree is a +// true latency breakdown (unlike the accum_* thread local sub timers below, +// which sum per worker CPU time across rayon threads and can exceed 100%). A +// parallel region is one span around the blocking call; its internal split is +// reported separately as CPU time, never mixed into the wall tree. +// +// Per table spans are the exception. `multi_prove` runs one driver thread per +// in flight table (`r1_aux_build`, `r1_aux_commit`, `rounds_2to4`), so up to +// `table_parallelism()` of them are open at once: they DO overlap in wall time +// and their sum exceeds their parent. They still nest correctly, because each +// driver seeds its span depth from the spawning thread (`current_depth` / +// `enter_depth`) instead of starting a fresh thread at depth 0 — but read them +// as per table wall time, not as a share of the enclosing phase. // // let _s = instruments::span("trace_build"); // RAII, stops on drop // @@ -36,6 +44,32 @@ thread_local! { static SPAN_DEPTH: std::cell::Cell = const { std::cell::Cell::new(0) }; } +/// Depth the next span opened on this thread would be stamped with. +/// +/// Read this on the thread that spawns workers and hand the value to +/// [`enter_depth`] inside each worker: a freshly spawned OS thread starts at +/// depth 0, so without it every span opened off the main thread is recorded as +/// a root sibling of `prove_total` and the tree stops being a breakdown. +pub fn current_depth() -> u16 { + SPAN_DEPTH.with(|d| d.get()) +} + +/// Restores the span depth this thread had before [`enter_depth`]. +#[must_use] +pub struct DepthGuard(u16); + +/// Seed this thread's span depth from a parent thread, restoring the previous +/// value when the returned guard drops. +pub fn enter_depth(depth: u16) -> DepthGuard { + DepthGuard(SPAN_DEPTH.with(|d| d.replace(depth))) +} + +impl Drop for DepthGuard { + fn drop(&mut self) { + SPAN_DEPTH.with(|d| d.set(self.0)); + } +} + #[must_use] pub struct SpanGuard { label: &'static str, @@ -263,8 +297,15 @@ pub struct Round1SubOps { pub struct MultiProveTiming { pub prepass: Duration, pub main_commits: Duration, + /// Aux build wall time summed over the concurrent per-table drivers. It is + /// no longer a phase of its own — the fused chain runs it inside + /// `rounds_2_4`, so it overlaps itself and is a subset of that wall time, + /// not an addend. pub aux_build: Duration, + /// Aux commit, same accounting as `aux_build`. pub aux_commit: Duration, + /// Wall clock of the fused per-table region: aux build + aux commit + + /// rounds 2-4, all of it concurrent across `table_parallelism()` drivers. pub rounds_2_4: Duration, /// Sub-op breakdown for Round 1 (main + aux LDE vs Merkle). pub round1_sub: Round1SubOps, @@ -278,6 +319,11 @@ static R1_MAIN_LDE_US: AtomicU64 = AtomicU64::new(0); static R1_MAIN_MERKLE_US: AtomicU64 = AtomicU64::new(0); static R1_AUX_LDE_US: AtomicU64 = AtomicU64::new(0); static R1_AUX_MERKLE_US: AtomicU64 = AtomicU64::new(0); +// Aux build / aux commit wall time per table, summed across the concurrent +// per-table drivers of the fused chain (so, like the sub-timers, this is a +// CPU-style total that can exceed the fused region's wall clock). +static AUX_BUILD_US: AtomicU64 = AtomicU64::new(0); +static AUX_COMMIT_US: AtomicU64 = AtomicU64::new(0); // Aux build (LogUp) sub-phases, CPU time accumulated across tables/chunks. static AUX_FINGERPRINT_US: AtomicU64 = AtomicU64::new(0); static AUX_INVERT_US: AtomicU64 = AtomicU64::new(0); @@ -326,6 +372,20 @@ pub fn accum_aux_accumulate(d: Duration) { AUX_ACCUM_US.fetch_add(d.as_micros() as u64, Ordering::Relaxed); } +/// One table's aux build and aux commit wall time, from its fused-chain driver. +pub fn accum_aux_phases(build: Duration, commit: Duration) { + AUX_BUILD_US.fetch_add(build.as_micros() as u64, Ordering::Relaxed); + AUX_COMMIT_US.fetch_add(commit.as_micros() as u64, Ordering::Relaxed); +} + +/// Drain the summed per-table aux build / aux commit times. +pub fn take_aux_phases() -> (Duration, Duration) { + ( + Duration::from_micros(AUX_BUILD_US.swap(0, Ordering::Relaxed)), + Duration::from_micros(AUX_COMMIT_US.swap(0, Ordering::Relaxed)), + ) +} + pub fn take_r1_sub() -> Round1SubOps { Round1SubOps { main_lde: Duration::from_micros(R1_MAIN_LDE_US.swap(0, Ordering::Relaxed)), @@ -352,6 +412,8 @@ pub fn reset_all() { R1_MAIN_MERKLE_US.store(0, Ordering::Relaxed); R1_AUX_LDE_US.store(0, Ordering::Relaxed); R1_AUX_MERKLE_US.store(0, Ordering::Relaxed); + AUX_BUILD_US.store(0, Ordering::Relaxed); + AUX_COMMIT_US.store(0, Ordering::Relaxed); AUX_FINGERPRINT_US.store(0, Ordering::Relaxed); AUX_INVERT_US.store(0, Ordering::Relaxed); AUX_TERM_US.store(0, Ordering::Relaxed); diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index 42142f770..4199d0a80 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -20,10 +20,7 @@ use math::{ }; #[cfg(feature = "parallel")] -use rayon::prelude::{ - IndexedParallelIterator, IntoParallelIterator, IntoParallelRefIterator, - IntoParallelRefMutIterator, ParallelIterator, -}; +use rayon::prelude::{IntoParallelIterator, ParallelIterator}; #[cfg(feature = "debug-checks")] use crate::debug::validate_trace; @@ -249,7 +246,7 @@ type MainCommitTuple = ( type MainCommitTuple = (TableCommit, (Vec>, usize)); /// Round 1 commitment artifacts — Merkle trees, roots, challenges, and bus inputs. -/// Borrowed (not consumed) when building `Round1` in Phase D. +/// Borrowed (not consumed) when building `Round1`. pub(crate) struct Round1Commitments where Field: IsFFTField + IsSubFieldOf, @@ -263,10 +260,18 @@ where bus_public_inputs: Option>, } -/// LDE columns for main (Phase A) and auxiliary (Phase C) traces, consumed by value in Phase D. +/// Main and auxiliary LDE columns, consumed by value when the table's `Round1` +/// is assembled. +/// +/// Memory trade-off, asymmetric since the per-table scheduler fused aux build, +/// aux commit and rounds 2-4 into one task: +/// - main: produced by the Round 1 main commit, which is a phase-wide barrier, +/// so all N tables' main LDEs are live at once (O(N × main_cols × lde_size)). +/// - aux: produced and consumed inside the same fused task, so at most +/// `table_parallelism()` of them coexist (O(k × aux_cols × lde_size)). /// -/// Memory trade-off: all N tables' LDE columns are live simultaneously between Phase A/C -/// and Phase D (O(N × cols × lde_size)). +/// Under `debug-checks` the fused task is split around the cross-table bus +/// balance check, so there the aux LDEs are all-N-live like the main ones. struct Lde { /// Row-major main LDE buffer + its column count. main: (Vec>, usize), @@ -569,8 +574,17 @@ where } /// Number of tables to process concurrently in `multi_prove`. -/// Default: num_cores / 3 (benchmarked optimal on both M3 Pro and EPYC 9454P). -/// Override with `TABLE_PARALLELISM` env var. +/// +/// Defaults: `num_cores / 3` on CPU builds (benchmarked optimal on both M3 Pro +/// and EPYC 9454P — every table there is pure host work), `num_cores * 2 / 3` +/// under `cuda`, where most in-flight tables sit in GPU waits so more of them +/// pay (swept flat at ~2/3 of the cores on a 16-core/RTX 5090 box). Both arms +/// are overridden by the `TABLE_PARALLELISM` env var. Without the `parallel` +/// feature this is hardcoded to 1 and the env var is ignored. +/// +/// Not only the prover's `k`: `auto_storage::decide` feeds this into the +/// RAM-vs-Disk storage estimate, so the `cuda` arm also doubles that transient +/// term (see `peak_bytes`). pub fn table_parallelism() -> usize { #[cfg(feature = "parallel")] { @@ -581,7 +595,18 @@ pub fn table_parallelism() -> usize { let cores = std::thread::available_parallelism() .map(|n| n.get()) .unwrap_or(4); - (cores / 3).max(1) + // GPU builds: with the admission scheduler most in-flight + // tables sit in GPU waits, so more of them pay (swept flat at + // ~2/3 of the cores on a 16-core/RTX 5090 box). CPU builds + // stay at cores/3 — every table is pure host work there. + #[cfg(feature = "cuda")] + { + (cores * 2 / 3).max(1) + } + #[cfg(not(feature = "cuda"))] + { + (cores / 3).max(1) + } }) } #[cfg(not(feature = "parallel"))] @@ -607,34 +632,108 @@ fn estimate_table_vram_bytes(main_cols: usize, aux_cols: usize, lde_size: usize) lde_term.saturating_add(tree_term) } -/// Plan contiguous table chunks for parallel proving. A chunk grows until it -/// hits `k` tables or its summed VRAM estimate would exceed `budget`; a single -/// table larger than `budget` runs solo. With `budget == u64::MAX` (non-cuda, -/// or VRAM not binding) chunks fall back to fixed size `k`, identical to the -/// old `step_by(k)`, so scheduling and the proof are unchanged. Returns -/// `(start, end)` half open ranges covering `0..estimates.len()` in order. -fn plan_table_chunks(estimates: &[u64], k: usize, budget: u64) -> Vec<(usize, usize)> { - let n = estimates.len(); - let k = k.max(1); - let budget = budget as u128; - let mut chunks = Vec::new(); - let mut start = 0; - while start < n { - let mut end = start; - let mut acc: u128 = 0; - while end < n { - let next = estimates[end] as u128; - // Always admit at least one table per chunk (oversized → solo). - if end > start && (end - start >= k || acc + next > budget) { - break; +/// Byte-budget admission gate for concurrently proven tables. `acquire` +/// blocks until the requested bytes fit under the budget, releasing on +/// permit drop. An oversized request is admitted alone (when nothing else +/// holds bytes), so tables larger than the whole budget still prove. +/// +/// Only OS driver threads block here (see `run_admitted`) — never rayon +/// workers, whose pool the admitted tables use internally and which a +/// blocked worker would starve. +struct VramGate { + used: std::sync::Mutex, + freed: std::sync::Condvar, + budget: u64, +} + +struct VramPermit<'a> { + gate: &'a VramGate, + bytes: u64, +} + +impl VramGate { + fn new(budget: u64) -> Self { + Self { + used: std::sync::Mutex::new(0), + freed: std::sync::Condvar::new(), + budget, + } + } + + fn acquire(&self, bytes: u64) -> VramPermit<'_> { + let mut used = self.used.lock().unwrap(); + loop { + if *used == 0 || used.saturating_add(bytes) <= self.budget { + *used = used.saturating_add(bytes); + return VramPermit { gate: self, bytes }; } - acc += next; - end += 1; + used = self.freed.wait(used).unwrap(); } - chunks.push((start, end)); - start = end; } - chunks +} + +impl Drop for VramPermit<'_> { + fn drop(&mut self) { + let mut used = self.gate.used.lock().unwrap(); + *used = used.saturating_sub(self.bytes); + drop(used); + self.gate.freed.notify_all(); + } +} + +/// Run `task` once per table index on `workers` OS driver threads, admitting +/// each index through `gate` with its estimated bytes. `order` fixes the +/// start order (heaviest table first, so the long pole starts early and small +/// tables fill around it — the fixed chunks this replaces made every table +/// wait for the slowest of its chunk). Returns one slot per original index. +fn run_admitted( + order: &[usize], + estimates: &[u64], + gate: &VramGate, + workers: usize, + task: impl Fn(usize) -> T + Sync, +) -> Vec> { + let results: Vec>> = estimates + .iter() + .map(|_| std::sync::Mutex::new(None)) + .collect(); + let cursor = std::sync::atomic::AtomicUsize::new(0); + // Spans opened inside `task` run on these worker threads, and a fresh OS + // thread's span depth starts at 0 — which would record every per-table + // span as a root instead of a child of the phase that spawned it. Carry + // the spawning thread's depth across the scope boundary. + #[cfg(feature = "instruments")] + let parent_depth = crate::instruments::current_depth(); + std::thread::scope(|scope| { + for _ in 0..workers.max(1).min(order.len().max(1)) { + scope.spawn(|| { + #[cfg(feature = "instruments")] + let _depth = crate::instruments::enter_depth(parent_depth); + loop { + let pos = cursor.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + if pos >= order.len() { + return; + } + let idx = order[pos]; + let permit = gate.acquire(estimates[idx]); + let out = task(idx); + *results[idx].lock().unwrap() = Some(out); + drop(permit); + } + }); + } + }); + results + .into_iter() + .map(|m| m.into_inner().unwrap()) + .collect() +} + +/// Table indices sorted heaviest-first by estimate. +fn heaviest_first(estimates: &[u64]) -> Vec { + let mut order: Vec = (0..estimates.len()).collect(); + order.sort_by_key(|&i| std::cmp::Reverse(estimates[i])); + order } /// A container for the results of the second round of the STARK Prove protocol. @@ -955,9 +1054,9 @@ pub trait IsStarkProver< } /// Compute the main-trace LDE and commit. Returns a `TableCommit` along - /// with the owned LDE columns (consumed later in Phase D) and (under - /// cuda) the optional device LDE buffer kept alive for downstream rounds - /// when the R1 fused GPU pipeline ran. + /// with the owned LDE columns (consumed later by the table's fused task) + /// and (under cuda) the optional device LDE buffer kept alive for + /// downstream rounds when the R1 fused GPU pipeline ran. /// /// `precomputed`: if present, the leading `num_cols` columns are committed /// as a separate Merkle tree (the precomputed split for preprocessed @@ -1241,8 +1340,8 @@ pub trait IsStarkProver< /// Recompute Round1 from the trace, reusing the Merkle trees stored in commitments. /// - /// Only used by `run_debug_checks` — Phase D consumes the cached LDE - /// directly and does not go through this path. + /// Only used by `run_debug_checks` — the production path consumes the + /// cached LDE directly and does not go through here. #[cfg(feature = "debug-checks")] fn reconstruct_round1( air: &dyn AIR, @@ -1318,10 +1417,18 @@ pub trait IsStarkProver< } /// Reconstruct Round1 for every table, print the bus balance report, and - /// validate each trace. Called once after Phase C commits. + /// validate each trace. + /// + /// Cross-table (bus balance) checks need every table's commitments at once, + /// so under `debug-checks` the fused per-table chain is split into two + /// admitted passes and this runs once, on the main thread, between them. + /// + /// `pair_cells` is the same per-table slot vector the drivers use. Each + /// driver only ever locks its own index, so locking them here — after the + /// aux pass has joined and before the rounds pass starts — is uncontended. #[cfg(feature = "debug-checks")] fn run_debug_checks( - air_trace_pairs: &[AirTracePair<'_, Field, FieldExtension, PI>], + pair_cells: &[std::sync::Mutex>], commitments: &[Round1Commitments], domains: &[Arc>], twiddle_caches: &[Arc>], @@ -1331,13 +1438,15 @@ pub trait IsStarkProver< PI: Send + Sync + Clone, { let mut temp_results: Vec> = - Vec::with_capacity(air_trace_pairs.len()); - for (((air, trace, _), commitment), (domain, twiddles)) in air_trace_pairs + Vec::with_capacity(pair_cells.len()); + for ((cell, commitment), (domain, twiddles)) in pair_cells .iter() .zip(commitments.iter()) .zip(domains.iter().zip(twiddle_caches.iter())) { - let result = Self::reconstruct_round1(*air, *trace, domain, commitment, twiddles) + let pair = cell.lock().unwrap(); + let (air, trace, _) = &*pair; + let result = Self::reconstruct_round1(*air, trace, domain, commitment, twiddles) .expect("reconstruct_round1 failed in debug-checks"); temp_results.push(result); } @@ -1348,15 +1457,17 @@ pub trait IsStarkProver< .collect(); print_bus_balance_report(&all_bus_public_inputs); - for (((air, trace, pub_inputs), round_1_result), domain) in air_trace_pairs + for ((cell, round_1_result), domain) in pair_cells .iter() .zip(temp_results.iter()) .zip(domains.iter()) { + let pair = cell.lock().unwrap(); + let (air, trace, pub_inputs) = &*pair; validate_trace( *air, *pub_inputs, - *trace, + trace, domain, &round_1_result.rap_challenges, round_1_result.bus_public_inputs.as_ref(), @@ -2933,7 +3044,7 @@ pub trait IsStarkProver< /// /// The transcript must be safely initialized before passing it to this method. fn multi_prove( - mut air_trace_pairs: Vec>, + #[allow(unused_mut)] mut air_trace_pairs: Vec>, transcript: &mut (impl IsStarkTranscript + Clone + Send), #[cfg(feature = "disk-spill")] storage_mode: StorageMode, ) -> Result, ProvingError> @@ -2984,7 +3095,7 @@ pub trait IsStarkProver< // of the tables proved concurrently so large blocks don't exhaust VRAM. // It is an extra ceiling on top of `k` (it never raises concurrency). On // non-cuda builds, or when the budget can't be queried, it is `u64::MAX` - // and chunking falls back to fixed size `k`. + // and the gate is inert — concurrency is then bounded by `k` alone. #[cfg(feature = "cuda")] let vram_budget = math_cuda::device::backend() .map(|b| b.vram_budget_bytes()) @@ -3000,20 +3111,18 @@ pub trait IsStarkProver< // don't re-add pre-sizing without a shared-slab design that bounds the // number of allocations. + let vram_gate = VramGate::new(vram_budget); + // R1 main commit: only the main LDE and its Merkle scratch are resident, // so the aux columns add nothing to this phase's working set. - let main_chunks = { - let estimates: Vec = air_trace_pairs - .iter() - .enumerate() - .map(|(idx, (_, trace, _))| { - let lde_size = - domains[idx].interpolation_domain_size * domains[idx].blowup_factor; - estimate_table_vram_bytes(trace.num_main_columns, 0, lde_size) - }) - .collect(); - plan_table_chunks(&estimates, k, vram_budget) - }; + let main_estimates: Vec = air_trace_pairs + .iter() + .enumerate() + .map(|(idx, (_, trace, _))| { + let lde_size = domains[idx].interpolation_domain_size * domains[idx].blowup_factor; + estimate_table_vram_bytes(trace.num_main_columns, 0, lde_size) + }) + .collect(); // Spill main traces to mmap before Round 1 LDE. #[cfg(feature = "disk-spill")] @@ -3036,7 +3145,7 @@ pub trait IsStarkProver< } // ===================================================================== - // Round 1, Phase A: Commit all main traces (parallel in chunks of K) + // Round 1: Commit all main traces (VRAM-admitted, up to K concurrent) // ===================================================================== // All main trace commitments must be in the transcript before sampling // LogUp challenges. @@ -3049,57 +3158,62 @@ pub trait IsStarkProver< let mut main_commits: Vec> = Vec::with_capacity(num_airs); let mut main_ldes: Vec<(Vec>, usize)> = Vec::with_capacity(num_airs); // Optional device-side LDE handle per table, populated only when the - // R1 fused GPU pipeline produced one. Threaded through Phase D's zip - // chain so each handle stays paired with its table by construction. + // R1 fused GPU pipeline produced one. Indexed by table, and moved into + // the per-table `gpu_main_cells` slots below so each handle stays + // paired with its table across the fused chain. #[cfg(feature = "cuda")] let mut main_gpu_handles: Vec> = Vec::with_capacity(num_airs); - for &(chunk_start, chunk_end) in &main_chunks { - let chunk_range = chunk_start..chunk_end; - - let chunk_results: Vec> = - crate::par::par_map_collect(chunk_range, |idx| { - let (air, trace, _) = &air_trace_pairs[idx]; - let domain = &domains[idx]; - let twiddles = &twiddle_caches[idx]; - - let precomputed = air - .is_preprocessed() - .then(|| (air.precomputed_commitment(), air.num_precomputed_columns())); + // All main commits with continuous VRAM admission (no chunk barriers); + // the transcript only needs the roots absorbed in index order, done + // sequentially below once every commit completed — the one ordering + // Fiat-Shamir requires before sampling the shared challenges. + let main_results = run_admitted( + &heaviest_first(&main_estimates), + &main_estimates, + &vram_gate, + k, + |idx| { + let (air, trace, _) = &air_trace_pairs[idx]; + let domain = &domains[idx]; + let twiddles = &twiddle_caches[idx]; + + let precomputed = air + .is_preprocessed() + .then(|| (air.precomputed_commitment(), air.num_precomputed_columns())); + + // Stage-3 device-only gate: when it holds, `commit_main_trace` + // keeps the R1 LDE device-resident and skips the host D2H. + #[cfg(feature = "cuda")] + let device_only = Self::device_only_for(*air, domain); - // Stage-3 device-only gate: when it holds, `commit_main_trace` - // keeps the R1 LDE device-resident and skips the host D2H. + Self::commit_main_trace( + *trace, + domain, + twiddles, + precomputed, #[cfg(feature = "cuda")] - let device_only = Self::device_only_for(*air, domain); - - Self::commit_main_trace( - *trace, - domain, - twiddles, - precomputed, - #[cfg(feature = "cuda")] - device_only, - #[cfg(feature = "disk-spill")] - storage_mode, - ) - }); - - // Sequential: append roots to shared transcript (Fiat-Shamir ordering) - for result in chunk_results { - #[cfg(feature = "cuda")] - let (commit, cached_main, gpu_main) = result?; - #[cfg(not(feature = "cuda"))] - let (commit, cached_main) = result?; - if let Some(ref pre_root) = commit.precomputed_root { - transcript.append_bytes(pre_root); - } - transcript.append_bytes(&commit.root); - main_commits.push(commit); - main_ldes.push(cached_main); - #[cfg(feature = "cuda")] - main_gpu_handles.push(gpu_main); + device_only, + #[cfg(feature = "disk-spill")] + storage_mode, + ) + }, + ); + for result in main_results { + let result = result.expect("run_admitted fills every slot"); + #[cfg(feature = "cuda")] + let (commit, cached_main, gpu_main) = result?; + #[cfg(not(feature = "cuda"))] + let (commit, cached_main) = result?; + if let Some(ref pre_root) = commit.precomputed_root { + transcript.append_bytes(pre_root); } + transcript.append_bytes(&commit.root); + main_commits.push(commit); + main_ldes.push(cached_main); + #[cfg(feature = "cuda")] + main_gpu_handles.push(gpu_main); } #[cfg(feature = "instruments")] @@ -3112,7 +3226,7 @@ pub trait IsStarkProver< } // ===================================================================== - // Round 1, Phase B: Sample shared LogUp challenges + // Round 1: Sample shared LogUp challenges // ===================================================================== let lookup_challenges: Vec> = if needs_lookup_challenges { @@ -3124,23 +3238,16 @@ pub trait IsStarkProver< }; // ===================================================================== - // Phase C + Rounds 2-4: Forked per table + // Aux build + aux commit + Rounds 2-4: fused per table // ===================================================================== // Each table gets an independent transcript fork (cloned from the shared - // state after Phase B, domain-separated by table index). This matches - // the verifier's forking and makes per-table proving independent. + // state after the LogUp challenges, domain-separated by table index). + // This matches the verifier's forking and makes per-table proving + // independent. // - // Split into two passes for parallelism: - // Pass 1 (parallel): Build all auxiliary traces (fingerprint + batch inversion) - // Pass 2 (parallel): Fork transcript → extract → LDE → commit - - // Pass 1: Build aux traces in parallel. - // Each build_auxiliary_trace has internal parallelism (batch_inverse, par_chunks), - // but outer parallelism over 12 tables also helps on high-core-count machines. - #[cfg(feature = "instruments")] - let phase_start = Instant::now(); - #[cfg(feature = "instruments")] - let __sp = crate::instruments::span("r1_aux_build"); + // Aux build, aux commit and rounds 2-4 run FUSED per table below (one + // driver chains all three for its table, so tables never wait on a + // phase barrier); only this sequential prep runs here. // Disk-spill needs the aux columns in the host trace to spill them, so // disable the GPU-resident aux build (it would keep them device-only). @@ -3165,67 +3272,8 @@ pub trait IsStarkProver< } } - #[cfg(feature = "parallel")] - let aux_iter = air_trace_pairs.par_iter_mut(); - #[cfg(not(feature = "parallel"))] - let aux_iter = air_trace_pairs.iter_mut(); - let bus_inputs_vec: Vec>> = aux_iter - .map(|(air, trace, _)| { - if air.has_aux_trace() { - air.build_auxiliary_trace(*trace, &lookup_challenges) - } else { - None - } - }) - .collect(); - - // The trace-domain snapshots retained by the R1 main LDE (both Arcs: - // trace.main_trace_dev and GpuLdeBase.trace_dev) have exactly one - // consumer — the aux build above. Drop them now so the main-trace-sized - // device buffers are reclaimed before the aux-commit + DEEP/FRI VRAM - // peak instead of living to the end of the proof. - #[cfg(feature = "cuda")] - { - for (_, trace, _) in air_trace_pairs.iter_mut() { - trace.clear_main_trace_dev(); - } - for handle in main_gpu_handles.iter_mut().flatten() { - handle.trace_dev = None; - handle.trace_rows = 0; - } - } - - // Spill all aux trace tables to mmap before any Round 1 aux LDE work. - #[cfg(feature = "disk-spill")] - if storage_mode == StorageMode::Disk { - crate::par::par_try_for_each_mut(&mut air_trace_pairs, |(air, trace, _)| { - if air.has_aux_trace() { - trace - .spill_aux_to_disk() - .map_err(|e| ProvingError::DiskSpill(format!("aux trace: {e}")))?; - } - Ok::<(), ProvingError>(()) - })?; - } - - #[cfg(feature = "instruments")] - drop(__sp); - #[cfg(feature = "instruments")] - let aux_build_elapsed = phase_start.elapsed(); - #[cfg(feature = "instruments")] - if let Some(s) = crate::instruments::snap("After aux build") { - heap_snaps.push(s); - } - - // Pass 2: Parallel fork transcript → extract → LDE → commit in chunks of K. - // Each table gets its own transcript fork. - #[cfg(feature = "instruments")] - let phase_start = Instant::now(); - #[cfg(feature = "instruments")] - let __sp = crate::instruments::span("r1_aux_commit"); - // Pre-fork all transcripts (cheap, sequential — must match verifier ordering) - let mut table_transcripts: Vec<_> = (0..num_airs) + let table_transcripts: Vec<_> = (0..num_airs) .map(|idx| { let mut t = transcript.clone(); if num_airs > 1 { @@ -3235,10 +3283,10 @@ pub trait IsStarkProver< }) .collect(); - // Parallel aux commit in chunks of K. The closure returns a cfg-gated - // AuxResult. Under cuda it carries the optional ext3 GPU LDE handle as - // a third element, so Phase D's zip chain keeps it paired with its - // table without a separate handle vector. + // The aux stage of the fused chain returns a cfg-gated AuxResult. Under + // cuda it carries the optional ext3 GPU LDE handle as a third element, + // so the handle stays inside its own table's task and never needs a + // separate handle vector. #[cfg(feature = "cuda")] type AuxResult = ( Option>, @@ -3247,44 +3295,110 @@ pub trait IsStarkProver< ); #[cfg(not(feature = "cuda"))] type AuxResult = (Option>, (Vec>, usize)); - #[allow(clippy::type_complexity)] - let mut aux_results: Vec> = Vec::with_capacity(num_airs); - // R1 aux commit and rounds 2 to 4 share the peak working set: the main // and aux LDEs are co-resident, plus the composition and Merkle - // transients (in the scratch factor). `num_aux_columns` is populated by - // the aux build above, so this estimate is accurate for both phases. - let peak_chunks = { - let estimates: Vec = air_trace_pairs - .iter() - .enumerate() - .map(|(idx, (_, trace, _))| { - let lde_size = - domains[idx].interpolation_domain_size * domains[idx].blowup_factor; - estimate_table_vram_bytes( - trace.num_main_columns, - trace.num_aux_columns, - lde_size, - ) - }) + // transients (in the scratch factor). The aux width comes from the AIR + // layout (the aux build itself runs inside the admitted chain below). + let peak_estimates: Vec = air_trace_pairs + .iter() + .enumerate() + .map(|(idx, (air, trace, _))| { + let lde_size = domains[idx].interpolation_domain_size * domains[idx].blowup_factor; + let (_, aux_cols) = air.trace_layout(); + estimate_table_vram_bytes(trace.num_main_columns, aux_cols, lde_size) + }) + .collect(); + + // Per-table slots for the fused chain: each driver takes or locks only + // its own index, so every mutex is uncontended by construction. + let pair_cells: Vec>> = + air_trace_pairs + .into_iter() + .map(std::sync::Mutex::new) .collect(); - plan_table_chunks(&estimates, k, vram_budget) - }; + let main_commit_cells: Vec>>> = main_commits + .into_iter() + .map(|c| std::sync::Mutex::new(Some(c))) + .collect(); + #[allow(clippy::type_complexity)] + let main_lde_cells: Vec< + std::sync::Mutex>, usize)>>, + > = main_ldes + .into_iter() + .map(|l| std::sync::Mutex::new(Some(l))) + .collect(); + #[cfg(feature = "cuda")] + let gpu_main_cells: Vec>> = + main_gpu_handles + .into_iter() + .map(std::sync::Mutex::new) + .collect(); + let transcript_cells: Vec<_> = table_transcripts + .into_iter() + .map(std::sync::Mutex::new) + .collect(); + #[cfg(feature = "instruments")] + #[allow(clippy::type_complexity)] + let table_timings_mx: std::sync::Mutex< + Vec<(String, usize, Duration, crate::instruments::TableSubOps)>, + > = std::sync::Mutex::new(Vec::new()); - for &(chunk_start, chunk_end) in &peak_chunks { - let chunk_range = chunk_start..chunk_end; + // Fused chain, stage 1: aux build → aux commit → aux root into the + // table's transcript fork → Round1 assembly. + #[allow(clippy::type_complexity)] + let aux_stage = |idx: usize| -> Result< + ( + Round1Commitments, + Lde, + ), + ProvingError, + > { + let mut pair = pair_cells[idx].lock().unwrap(); + let (air, trace, _) = &mut *pair; + let domain = &domains[idx]; + let twiddles = &twiddle_caches[idx]; - #[allow(clippy::type_complexity)] - let chunk_aux: Vec, ProvingError>> = - crate::par::par_map_collect(chunk_range, |idx| { - let (air, trace, _) = &air_trace_pairs[idx]; - let domain = &domains[idx]; - let twiddles = &twiddle_caches[idx]; + #[cfg(feature = "instruments")] + let __sp = crate::instruments::span("r1_aux_build"); + #[cfg(feature = "instruments")] + let t_aux_build = Instant::now(); + let bus_public_inputs = if air.has_aux_trace() { + air.build_auxiliary_trace(*trace, &lookup_challenges) + } else { + None + }; + // The trace-domain snapshot retained by the R1 main LDE has exactly + // one consumer — the aux build above. Reclaim it before this + // table's aux-commit + DEEP/FRI VRAM peak. + #[cfg(feature = "cuda")] + { + trace.clear_main_trace_dev(); + if let Some(handle) = gpu_main_cells[idx].lock().unwrap().as_mut() { + handle.trace_dev = None; + handle.trace_rows = 0; + } + } + #[cfg(feature = "disk-spill")] + if storage_mode == StorageMode::Disk && air.has_aux_trace() { + trace + .spill_aux_to_disk() + .map_err(|e| ProvingError::DiskSpill(format!("aux trace: {e}")))?; + } + #[cfg(feature = "instruments")] + let aux_build_dur = t_aux_build.elapsed(); + #[cfg(feature = "instruments")] + drop(__sp); + #[cfg(feature = "instruments")] + let __sp = crate::instruments::span("r1_aux_commit"); + #[cfg(feature = "instruments")] + let t_aux_commit = Instant::now(); + let aux_full: AuxResult = + (|| -> Result, ProvingError> { if air.has_aux_trace() { let lde_size = domain.interpolation_domain_size * domain.blowup_factor; - // Same gate as the main commit (Phase A): skip the aux + // Same gate as the Round 1 main commit: skip the aux // host D2H when device-only, so both buffers are left // empty together for this table. #[cfg(feature = "cuda")] @@ -3418,181 +3532,161 @@ pub trait IsStarkProver< #[cfg(not(feature = "cuda"))] Ok((None, (Vec::new(), 0))) } - }); - - // Sequential: append aux roots to forked transcripts. - for (j, result) in chunk_aux.into_iter().enumerate() { - let aux_full = result?; - // Tuple shape is cfg-gated; `.0` is the optional TableCommit - // in both variants. - if let Some(ref c) = aux_full.0 { - table_transcripts[chunk_start + j].append_bytes(&c.root); - } - aux_results.push(aux_full); + })()?; + // Tuple shape is cfg-gated; `.0` is the optional TableCommit in + // both variants. Aux roots go to the table's OWN fork, so no + // cross-table ordering is needed here. + if let Some(ref c) = aux_full.0 { + transcript_cells[idx].lock().unwrap().append_bytes(&c.root); } - } - - // Build commitments and cached LDEs as separate vecs: - // commitments are borrowed in Phase D, LDEs are consumed by value. - let mut commitments: Vec> = - Vec::with_capacity(num_airs); - let mut cached_ldes: Vec> = Vec::with_capacity(num_airs); - // Under cuda, fold main_gpu_handles into the zip chain so each handle - // stays paired with its table by construction. - #[cfg(feature = "cuda")] - let main_iter = main_commits - .into_iter() - .zip(main_ldes) - .zip(main_gpu_handles); - #[cfg(not(feature = "cuda"))] - let main_iter = main_commits.into_iter().zip(main_ldes); + #[cfg(feature = "instruments")] + crate::instruments::accum_aux_phases(aux_build_dur, t_aux_commit.elapsed()); + #[cfg(feature = "instruments")] + drop(__sp); - for ((main_pack, aux_full), bus_public_inputs) in - main_iter.zip(aux_results).zip(bus_inputs_vec) - { - #[cfg(feature = "cuda")] - let ((main_commit, main_lde), gpu_main) = main_pack; - #[cfg(not(feature = "cuda"))] - let (main_commit, main_lde) = main_pack; #[cfg(feature = "cuda")] let (aux_commit, cached_aux, gpu_aux) = aux_full; #[cfg(not(feature = "cuda"))] let (aux_commit, cached_aux) = aux_full; - commitments.push(Round1Commitments { + let main_commit = main_commit_cells[idx] + .lock() + .unwrap() + .take() + .expect("main commit consumed once per table"); + let main_lde = main_lde_cells[idx] + .lock() + .unwrap() + .take() + .expect("main lde consumed once per table"); + #[cfg(feature = "cuda")] + let gpu_main = gpu_main_cells[idx].lock().unwrap().take(); + let commitment = Round1Commitments { main: main_commit, aux: aux_commit, rap_challenges: lookup_challenges.clone(), bus_public_inputs, - }); + }; #[cfg(feature = "cuda")] - cached_ldes.push(Lde { + let lde = Lde { main: main_lde, aux: cached_aux, gpu_main, gpu_aux, - }); + }; #[cfg(not(feature = "cuda"))] - cached_ldes.push(Lde { + let lde = Lde { main: main_lde, aux: cached_aux, - }); - } + }; + Ok((commitment, lde)) + }; - #[cfg(feature = "instruments")] - drop(__sp); - #[cfg(feature = "instruments")] - let aux_commit_elapsed = phase_start.elapsed(); - #[cfg(feature = "instruments")] - if let Some(s) = crate::instruments::snap("After aux commit") { - heap_snaps.push(s); - } + // Fused chain, stage 2: Round1 from the cached LDE (consumed by value, + // no recomputation) → rounds 2-4 against the table's transcript fork. + let rounds_stage = |idx: usize, + commitment: Round1Commitments, + lde: Lde| + -> Result, ProvingError> { + let pair = pair_cells[idx].lock().unwrap(); + let (air, trace, pub_inputs) = &*pair; + let _ = trace; // used by instruments + let domain = &domains[idx]; - #[cfg(feature = "debug-checks")] - Self::run_debug_checks(&air_trace_pairs, &commitments, &domains, &twiddle_caches); + #[cfg(feature = "instruments")] + let __sp = crate::instruments::span("rounds_2to4"); + #[cfg(feature = "instruments")] + let table_start = Instant::now(); - // ===================================================================== - // Rounds 2-4: Parallel per-table proving in chunks of K - // ===================================================================== - // Each chunk of K tables is processed in parallel. Cached LDE columns - // from Phase A/C are consumed here (zero-copy move), eliminating the - // expensive reconstruct_round1 recomputation. + let mut round_1_result = + commitment.build_round1(lde, air.step_size(), domain.blowup_factor); + + let mut tguard = transcript_cells[idx].lock().unwrap(); + if let Some(ref bpi) = round_1_result.bus_public_inputs { + tguard.append_field_element(&bpi.table_contribution); + } + + let proof = Self::prove_rounds_2_to_4( + *air, + *pub_inputs, + &mut round_1_result, + &mut *tguard, + domain, + &twiddle_caches[idx], + )?; + + #[cfg(feature = "instruments")] + { + let sub_ops = crate::instruments::take_round_sub_ops().unwrap_or_default(); + table_timings_mx.lock().unwrap().push(( + air.name().to_string(), + trace.num_rows(), + table_start.elapsed(), + sub_ops, + )); + } + Ok(proof) + }; #[cfg(feature = "instruments")] let phase_start = Instant::now(); - #[cfg(feature = "instruments")] - let __sp = crate::instruments::span("rounds_2to4"); - #[cfg(feature = "instruments")] - let mut table_timings: Vec<( - String, - usize, - Duration, - crate::instruments::TableSubOps, - )> = Vec::with_capacity(num_airs); - let mut proofs = Vec::with_capacity(num_airs); - let mut lde_drain = cached_ldes.into_iter(); - for &(chunk_start, chunk_end) in &peak_chunks { - let chunk_size = chunk_end - chunk_start; - - let chunk_ldes: Vec> = - lde_drain.by_ref().take(chunk_size).collect(); - let chunk_commitments = &commitments[chunk_start..chunk_end]; - let chunk_transcripts = &mut table_transcripts[chunk_start..chunk_end]; - - #[cfg(feature = "parallel")] - let iter = chunk_ldes - .into_par_iter() - .zip(chunk_commitments.par_iter()) - .zip(chunk_transcripts.par_iter_mut()) - .enumerate(); - #[cfg(not(feature = "parallel"))] - let iter = chunk_ldes - .into_iter() - .zip(chunk_commitments.iter()) - .zip(chunk_transcripts.iter_mut()) - .enumerate(); - - let chunk_results: Vec> = iter - .map(|(j, ((lde, commitment), table_transcript))| { - let idx = chunk_start + j; - let (air, trace, pub_inputs) = &air_trace_pairs[idx]; - let _ = trace; // used by instruments - let domain = &domains[idx]; - - #[cfg(feature = "instruments")] - let table_start = Instant::now(); - - // Build Round1 from cached LDE (consumed by value, no recomputation). - let mut round_1_result = - commitment.build_round1(lde, air.step_size(), domain.blowup_factor); - - if let Some(ref bpi) = round_1_result.bus_public_inputs { - table_transcript.append_field_element(&bpi.table_contribution); - } + let peak_order = heaviest_first(&peak_estimates); - let proof = Self::prove_rounds_2_to_4( - *air, - *pub_inputs, - &mut round_1_result, - table_transcript, - domain, - &twiddle_caches[idx], - )?; - - #[cfg(feature = "instruments")] - let table_timing = { - let sub_ops = crate::instruments::take_round_sub_ops().unwrap_or_default(); - ( - air.name().to_string(), - trace.num_rows(), - table_start.elapsed(), - sub_ops, - ) - }; + // One fused task per table: while a heavy table works through a + // host-bound stretch, the others' GPU stages fill the device. The + // shared transcript is untouched past this point (each fork is + // per-table), so any order is sound; proofs are drained in index order. + #[cfg(not(feature = "debug-checks"))] + let table_results = run_admitted(&peak_order, &peak_estimates, &vram_gate, k, |idx| { + let (commitment, lde) = aux_stage(idx)?; + rounds_stage(idx, commitment, lde) + }); - #[cfg(feature = "instruments")] - return Ok((proof, table_timing)); - #[cfg(not(feature = "instruments"))] - Ok(proof) - }) + // debug-checks needs every table's commitments and traces between the + // aux and rounds stages (cross-table bus balance), so it splits the + // fused chain into two admitted passes around the check. + #[cfg(feature = "debug-checks")] + let table_results = { + let aux_outs = run_admitted(&peak_order, &peak_estimates, &vram_gate, k, aux_stage); + let mut commitments = Vec::with_capacity(num_airs); + let mut ldes = Vec::with_capacity(num_airs); + for out in aux_outs { + let (c, l) = out.expect("run_admitted fills every slot")?; + commitments.push(c); + ldes.push(l); + } + Self::run_debug_checks(&pair_cells, &commitments, &domains, &twiddle_caches); + #[allow(clippy::type_complexity)] + let staged: Vec< + std::sync::Mutex< + Option<( + Round1Commitments, + Lde, + )>, + >, + > = commitments + .into_iter() + .zip(ldes) + .map(|p| std::sync::Mutex::new(Some(p))) .collect(); + run_admitted(&peak_order, &peak_estimates, &vram_gate, k, |idx| { + let (c, l) = staged[idx].lock().unwrap().take().unwrap(); + rounds_stage(idx, c, l) + }) + }; - for result in chunk_results { - #[cfg(feature = "instruments")] - { - let (proof, timing) = result?; - proofs.push(proof); - table_timings.push(timing); - } - #[cfg(not(feature = "instruments"))] - proofs.push(result?); - } + let mut proofs = Vec::with_capacity(num_airs); + for result in table_results { + proofs.push(result.expect("run_admitted fills every slot")?); } - #[cfg(feature = "instruments")] - drop(__sp); + let table_timings = table_timings_mx.into_inner().unwrap(); #[cfg(feature = "instruments")] { + // Aux build/commit are no longer phases of their own: each table's + // driver runs them inside the fused region, so these are sums over + // concurrent drivers and are a subset of `rounds_2_4`, not addends. + let (aux_build_elapsed, aux_commit_elapsed) = crate::instruments::take_aux_phases(); // Store timing data for the top-level report in prove_with_options. // Uses a thread-local to avoid changing multi_prove's return type. crate::instruments::store(crate::instruments::MultiProveTiming { @@ -3918,3 +4012,198 @@ fn print_bus_balance_report( } } } + +/// Scheduling primitives only. These are free functions over `&[u64]` with no +/// AIR, field or device dependency, so they are testable without a GPU — which +/// matters because CI never exercises them concurrently: `ubuntu-latest` has +/// 2-4 vCPU, so `table_parallelism()` floors to 1, and on non-cuda builds the +/// budget is `u64::MAX`, which makes `VramGate` inert. +#[cfg(test)] +mod scheduler_tests { + use super::{VramGate, heaviest_first, run_admitted}; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::time::Duration; + + /// Bound on how long a correct implementation may take to wake a waiter. + /// Only a failure deadline — never used to sequence the test. + const WAKE_DEADLINE: Duration = Duration::from_secs(10); + + #[test] + fn heaviest_first_is_a_descending_permutation() { + let empty: [u64; 0] = []; + assert!(heaviest_first(&empty).is_empty()); + assert_eq!(heaviest_first(&[3, 1, 2]), vec![0, 2, 1]); + + let estimates = [7u64, 0, 7, 3, 100, 1]; + let order = heaviest_first(&estimates); + let mut seen = order.clone(); + seen.sort_unstable(); + assert_eq!( + seen, + (0..estimates.len()).collect::>(), + "must be a permutation of 0..n" + ); + for w in order.windows(2) { + assert!( + estimates[w[0]] >= estimates[w[1]], + "must be descending by estimate, got {order:?}" + ); + } + } + + #[test] + fn heaviest_first_breaks_ties_by_index() { + // `sort_by_key` is stable, so equal estimates keep ascending index + // order: the admission order is a pure function of the estimates, and + // does not vary run to run. + assert_eq!(heaviest_first(&[5, 5, 5]), vec![0, 1, 2]); + assert_eq!(heaviest_first(&[1, 5, 5, 1]), vec![1, 2, 0, 3]); + } + + #[test] + fn run_admitted_fills_every_slot_exactly_once() { + // Includes the degenerate shapes: no work, one worker, more workers + // than tables, and `workers == 0` (clamped to 1 inside). + for (n, workers) in [(0, 4), (1, 1), (1, 8), (5, 0), (5, 1), (5, 8), (9, 4)] { + let estimates: Vec = (0..n).map(|i| i as u64 + 1).collect(); + let runs: Vec = (0..n).map(|_| AtomicUsize::new(0)).collect(); + let gate = VramGate::new(u64::MAX); + let out = run_admitted( + &heaviest_first(&estimates), + &estimates, + &gate, + workers, + |idx| { + runs[idx].fetch_add(1, Ordering::SeqCst); + idx * 10 + }, + ); + assert_eq!(out.len(), n, "one slot per table (n={n})"); + for i in 0..n { + assert_eq!( + runs[i].load(Ordering::SeqCst), + 1, + "table {i} ran exactly once (n={n}, workers={workers})" + ); + assert_eq!( + out[i], + Some(i * 10), + "slot {i} holds its own result (n={n}, workers={workers})" + ); + } + assert_eq!(*gate.used.lock().unwrap(), 0, "all permits released"); + } + } + + #[test] + fn concurrent_admissions_stay_under_budget() { + const BUDGET: u64 = 100; + const EACH: u64 = 40; // 2 fit, 3 do not + let estimates = vec![EACH; 16]; + let gate = VramGate::new(BUDGET); + let in_flight = AtomicUsize::new(0); + let max_in_flight = AtomicUsize::new(0); + run_admitted(&heaviest_first(&estimates), &estimates, &gate, 8, |_| { + let now = in_flight.fetch_add(1, Ordering::SeqCst) + 1; + max_in_flight.fetch_max(now, Ordering::SeqCst); + // Read under a held permit: the gate's own counter must never + // exceed the budget while anything is admitted. + assert!( + *gate.used.lock().unwrap() <= BUDGET, + "admitted bytes exceeded the budget" + ); + in_flight.fetch_sub(1, Ordering::SeqCst); + }); + assert!( + max_in_flight.load(Ordering::SeqCst) <= (BUDGET / EACH) as usize, + "more tables were in flight than the budget allows" + ); + assert_eq!(*gate.used.lock().unwrap(), 0); + } + + #[test] + fn oversized_request_is_admitted_alone_and_wakes_waiters() { + // A table bigger than the whole budget must still prove: it is admitted + // when nothing else holds bytes, rather than deadlocking forever. + let gate = VramGate::new(10); + let big = gate.acquire(100); + assert_eq!( + *gate.used.lock().unwrap(), + 100, + "oversized request admitted alone" + ); + + let (ready_tx, ready_rx) = std::sync::mpsc::channel(); + let (done_tx, done_rx) = std::sync::mpsc::channel(); + std::thread::scope(|s| { + s.spawn(|| { + ready_tx.send(()).unwrap(); + let permit = gate.acquire(5); + done_tx.send(()).unwrap(); + drop(permit); + }); + ready_rx.recv().unwrap(); + drop(big); + assert!( + done_rx.recv_timeout(WAKE_DEADLINE).is_ok(), + "dropping a permit must wake a waiter" + ); + }); + assert_eq!( + *gate.used.lock().unwrap(), + 0, + "permits release their bytes on drop" + ); + } + + /// Guards the instruments span tree: a driver thread starts at span depth + /// 0, so without the seeding in `run_admitted` every per-table span is + /// recorded as a root sibling of `prove_total` and the "% of total" column + /// stops meaning anything. Reads the depth directly instead of the global + /// span timeline, which other tests in this binary also write to. + #[cfg(feature = "instruments")] + #[test] + fn run_admitted_seeds_worker_span_depth() { + const PARENT_DEPTH: u16 = 3; + let outer = crate::instruments::enter_depth(PARENT_DEPTH); + let n = 6; + let estimates = vec![1u64; n]; + let gate = VramGate::new(u64::MAX); + let seen = run_admitted(&heaviest_first(&estimates), &estimates, &gate, 4, |_| { + crate::instruments::current_depth() + }); + for (i, depth) in seen.iter().enumerate() { + assert_eq!( + *depth, + Some(PARENT_DEPTH), + "table {i}'s driver must inherit the spawning thread's span depth" + ); + } + drop(outer); + assert_eq!( + crate::instruments::current_depth(), + 0, + "the calling thread's depth is restored" + ); + } + + #[test] + fn max_budget_gate_never_blocks() { + // The non-cuda / unqueryable-VRAM configuration. Two acquires that + // saturate `u64` must both be admitted. + let gate = VramGate::new(u64::MAX); + let held = gate.acquire(u64::MAX / 2); + let (tx, rx) = std::sync::mpsc::channel(); + std::thread::scope(|s| { + s.spawn(|| { + let _permit = gate.acquire(u64::MAX / 2 + 1000); + tx.send(()).unwrap(); + }); + assert!( + rx.recv_timeout(WAKE_DEADLINE).is_ok(), + "a u64::MAX budget must never block an acquire" + ); + }); + drop(held); + } +} diff --git a/prover/src/auto_storage.rs b/prover/src/auto_storage.rs index 49707cb4c..6b5ed8a5d 100644 --- a/prover/src/auto_storage.rs +++ b/prover/src/auto_storage.rs @@ -48,7 +48,8 @@ pub const SAFETY_FRACTION_DEN: u64 = 10; /// `(rows, main_cols, aux_cols, num_main_merkle_trees)` for a single table. type TableSpec = (u64, u64, u64, u64); -/// Bytes alive for the duration of phase D (LDE columns + main/aux Merkle). +/// Bytes counted as alive for the whole proof (LDE columns + main/aux Merkle). +/// Deliberately an over-estimate for the aux half — see `peak_bytes`. fn persistent_per_table(spec: TableSpec, blowup: u64) -> u64 { let (rows, main_cols, aux_cols, main_trees) = spec; let main_lde = rows @@ -228,19 +229,31 @@ pub fn decide(lengths: &TableLengths, blowup_factor: u8) -> StorageMode { } /// Peak RAM estimate in bytes for a proof whose trace shape matches `lengths`. +/// +/// `table_parallelism` is the prover's `k` (`stark::prover::table_parallelism`), +/// and it is not only a prover knob: `decide` feeds it in here, so the `cuda` +/// arm's `cores * 2 / 3` doubles the transient term below versus the CPU arm's +/// `cores / 3` and makes `Disk` more likely. That direction is safe (it +/// over-estimates), but it means a change to `k` changes the storage decision. pub fn peak_bytes(lengths: &TableLengths, blowup_factor: u8, table_parallelism: usize) -> u64 { let blowup = blowup_factor as u64; let k = table_parallelism.max(1); let specs = table_specs(lengths); - // Persistent: every table's LDE + main/aux Merkle is alive across phase D. + // Persistent: every table's main LDE + Merkle really is alive at once (the + // Round 1 main commit is a phase-wide barrier). The aux LDE no longer is — + // it is produced and consumed inside one table's fused task, so at most k + // coexist — but it is still counted for every table here, which keeps this + // an over-estimate rather than making the bound unsound. let persistent_total: u64 = specs .iter() .map(|s| persistent_per_table(*s, blowup)) .fold(0u64, u64::saturating_add); - // Transient: only k tables run round 2-4 in parallel. Conservative bound is - // the top-k tables by transient bytes (worst possible chunk assignment). + // Transient: only k tables run the fused aux+rounds task at a time. The + // top-k tables by transient bytes bound it; with the scheduler's + // heaviest-first admission that top-k is also the set actually admitted + // first, so this is the realistic peak, not a worst case. let mut transient_per: Vec = specs .iter() .map(|s| transient_per_table(*s, blowup)) diff --git a/prover/src/instruments.rs b/prover/src/instruments.rs index 0ea28273b..4663f092c 100644 --- a/prover/src/instruments.rs +++ b/prover/src/instruments.rs @@ -71,11 +71,13 @@ pub fn print_report( row_top("AIR construction", air_construction, total); if let Some(ref mp) = mp { - let round1 = mp.main_commits + mp.aux_build + mp.aux_commit; - + // Round 1's main commits are the last phase-level barrier: every main + // root must be in the transcript before the shared LogUp challenges are + // sampled. Aux build, aux commit and rounds 2-4 are fused per table, so + // they are reported under one wall-clock parent below, with their + // components as concurrent (summed-over-drivers) sub-rows. row_top("Pre-pass (domains/twiddles)", mp.prepass, total); - row_top("Round 1", round1, total); - row_sub(" Main trace commits", mp.main_commits, total); + row_top("Round 1 (main trace commits)", mp.main_commits, total); row_sub( " Main LDE (fused GPU: LDE+Keccak+Merkle / CPU: LDE only)", mp.round1_sub.main_lde, @@ -86,7 +88,15 @@ pub fn print_report( mp.round1_sub.main_merkle, total, ); - row_sub(" Aux trace build (parallel)", mp.aux_build, total); + row_top( + "Rounds 2\u{2013}4 (aux build+commit fused)", + mp.rounds_2_4, + total, + ); + eprintln!( + " \u{2500}\u{2500} below: summed across concurrent drivers \u{2500}\u{2500}" + ); + row_sub(" Aux trace build", mp.aux_build, total); row_sub( " LogUp fingerprint (CPU)", mp.round1_sub.aux_fingerprint, @@ -118,7 +128,7 @@ pub fn print_report( mp.round1_sub.aux_merkle, total, ); - row_top("Rounds 2\u{2013}4", mp.rounds_2_4, total); + eprintln!(" \u{2500}\u{2500} per table (R2\u{2013}4) \u{2500}\u{2500}"); // Merge split tables: MEMW[0..4] → MEMW x5 let mut merged: BTreeMap = BTreeMap::new(); @@ -209,10 +219,7 @@ pub fn print_report( ("R4 queries & openings", total_queries), ]; sub_ops.sort_by(|a, b| b.1.cmp(&a.1)); - eprintln!( - " {}", - " \u{2500}\u{2500} sub-operation totals (all tables) \u{2500}\u{2500}", - ); + eprintln!(" \u{2500}\u{2500} sub-operation totals (all tables) \u{2500}\u{2500}"); for (label, dur) in &sub_ops { row_sub(&format!(" {label}"), *dur, total); } diff --git a/scripts/bench_prover_scaling.sh b/scripts/bench_prover_scaling.sh index c1196d76e..520d13492 100755 --- a/scripts/bench_prover_scaling.sh +++ b/scripts/bench_prover_scaling.sh @@ -74,15 +74,17 @@ parse_run() { /^ Trace build/ { v = secs(); if (v) print "t_trace_build=" v } /^ AIR construction/ { v = secs(); if (v) print "t_air=" v } /^ Pre-pass/ { v = secs(); if (v) print "t_prepass=" v } + # "Round 1 (main trace commits)" is the whole of round 1 now; aux build and + # aux commit are fused into the per-table region and reported under + # "Rounds 2-4" as sums over the concurrent drivers, so they can exceed it. /^ Round 1 / { v = secs(); if (v) print "t_round1=" v } - /Main trace commits/ { v = secs(); if (v) print "t_main_commits="v } /Aux trace build/ { v = secs(); if (v) print "t_aux_build=" v } /Aux trace commit/ { v = secs(); if (v) print "t_aux_commit=" v } /Rounds 2/ { v = secs(); if (v) print "t_rounds24=" v } - /Main expand_columns_to_lde/{ v = secs(); if (v) print "t_main_lde=" v } - /Aux expand_columns_to_lde/ { v = secs(); if (v) print "t_aux_lde=" v } - /Main commit \(Merkle\)/ { v = secs(); if (v) print "t_main_merkle=" v } - /Aux commit \(Merkle\)/ { v = secs(); if (v) print "t_aux_merkle=" v } + /Main LDE/ { v = secs(); if (v) print "t_main_lde=" v } + /Aux LDE/ { v = secs(); if (v) print "t_aux_lde=" v } + /Main commit \(Merkle/ { v = secs(); if (v) print "t_main_merkle=" v } + /Aux commit \(Merkle/ { v = secs(); if (v) print "t_aux_merkle=" v } /^ Total FFT/ { v = secs(); if (v) print "t_total_fft=" v } /^ Total Merkle/ { v = secs(); if (v) print "t_total_merkle="v } /^ TOTAL / { v = secs(); if (v) print "t_total=" v } @@ -90,9 +92,13 @@ parse_run() { /After trace build/ { print "h_trace_build=" $(NF-1) } /After AIR/ { print "h_air=" $(NF-1) } /After pool alloc/ { print "h_pool_alloc=" $(NF-1) } + # NOTE: the "After aux build" / "After aux commit" heap snapshots were + # removed when aux build/commit were fused into the per-table scheduler -- + # with k tables in flight there is no single moment at which either has + # finished, so the snapshot had no meaning. "After main commits" is the + # last phase-wide barrier, and "Peak heap" still guards the region below + # it, so the heap-growth regressions keep coverage of the fused region. /After main commits/ { print "h_main_commits=" $(NF-1) } - /After aux build/ { print "h_aux_build=" $(NF-1) } - /After aux commit/ { print "h_aux_commit=" $(NF-1) } ' "$stderr" grep -o 'Peak heap: [0-9]*' "$stdout" | awk '{print "peak=" $3}' @@ -184,15 +190,14 @@ print_row "Execute" t_execute s print_row "Trace build" t_trace_build s print_row "AIR construction" t_air s print_row "Pre-pass" t_prepass s -print_row "Round 1" t_round1 s -print_row " Main trace commits" t_main_commits s +print_row "Round 1 (main commits)" t_round1 s print_row " Main LDE" t_main_lde s print_row " Main Merkle" t_main_merkle s +print_row "Rounds 2-4 (fused)" t_rounds24 s print_row " Aux trace build" t_aux_build s print_row " Aux trace commit" t_aux_commit s print_row " Aux LDE" t_aux_lde s print_row " Aux Merkle" t_aux_merkle s -print_row "Rounds 2-4" t_rounds24 s print_row "Total FFT (all rounds)" t_total_fft s print_row "Total Merkle" t_total_merkle s print_row "TOTAL" t_total s @@ -206,8 +211,8 @@ if [[ "$MODE" == "heap" ]]; then print_row "After AIR construction" h_air mb print_row "After pool alloc" h_pool_alloc mb print_row "After main commits" h_main_commits mb - print_row "After aux build" h_aux_build mb - print_row "After aux commit" h_aux_commit mb + # "After aux build" / "After aux commit" intentionally absent: see the NOTE + # in parse_run. Peak heap covers the fused region they used to bracket. print_row "Peak heap" peak mb fi @@ -270,8 +275,9 @@ if [[ "$MODE" == "heap" ]]; then regress "After AIR construction" h_air mb regress "After pool alloc" h_pool_alloc mb regress "After main commits" h_main_commits mb - regress "After aux build" h_aux_build mb - regress "After aux commit" h_aux_commit mb + # The "After aux build" / "After aux commit" heap-growth guards were dropped + # with their snapshots (see the NOTE in parse_run). "Peak heap" is the + # remaining regression guard over the fused per-table region. regress "Peak heap" peak mb fi