From 62301c89b6d622e133e2018123632add17caf8ac Mon Sep 17 00:00:00 2001 From: jotabulacios Date: Tue, 4 Aug 2026 11:46:41 -0300 Subject: [PATCH 1/3] fix: use CPU-aware prover scheduler --- crypto/stark/src/prover.rs | 84 +++++++++++++++++++++++++++++++++++++- 1 file changed, 82 insertions(+), 2 deletions(-) diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index 4047458bc..efd86b342 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -19,6 +19,8 @@ use math::{ polynomial::Polynomial, }; +#[cfg(all(feature = "parallel", not(feature = "cuda")))] +use rayon::prelude::IntoParallelRefIterator; #[cfg(feature = "parallel")] use rayon::prelude::{IntoParallelIterator, ParallelIterator}; @@ -640,17 +642,20 @@ fn estimate_table_vram_bytes(main_cols: usize, aux_cols: usize, lde_size: usize) /// 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. +#[cfg_attr(not(feature = "cuda"), allow(dead_code))] struct VramGate { used: std::sync::Mutex, freed: std::sync::Condvar, budget: u64, } +#[cfg_attr(not(feature = "cuda"), allow(dead_code))] struct VramPermit<'a> { gate: &'a VramGate, bytes: u64, } +#[cfg_attr(not(feature = "cuda"), allow(dead_code))] impl VramGate { fn new(budget: u64) -> Self { Self { @@ -686,12 +691,13 @@ impl Drop for VramPermit<'_> { /// 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. +#[cfg(feature = "cuda")] fn run_admitted( order: &[usize], estimates: &[u64], gate: &VramGate, workers: usize, - task: impl Fn(usize) -> T + Sync, + task: impl Fn(usize) -> T + Send + Sync, ) -> Vec> { let results: Vec>> = estimates .iter() @@ -721,6 +727,55 @@ fn run_admitted( .collect() } +/// CPU version of the table scheduler. There is no device admission to wait +/// on here, and the work underneath each table already uses Rayon internally. +/// Keep the outer scheduling inside the same Rayon pool instead of creating +/// driver OS threads: otherwise every table task launches nested Rayon work +/// from outside the pool and loses the work-stealing behavior of the original +/// phase scheduler. +#[cfg(all(not(feature = "cuda"), feature = "parallel"))] +fn run_admitted( + order: &[usize], + _estimates: &[u64], + _gate: &VramGate, + workers: usize, + task: impl Fn(usize) -> T + Send + Sync, +) -> Vec> { + let results: Vec>> = (0..order.len()) + .map(|_| std::sync::Mutex::new(None)) + .collect(); + + for chunk in order.chunks(workers.max(1)) { + let chunk_results: Vec<(usize, T)> = + chunk.par_iter().map(|&idx| (idx, task(idx))).collect(); + for (idx, result) in chunk_results { + *results[idx].lock().unwrap() = Some(result); + } + } + + results + .into_iter() + .map(|m| m.into_inner().unwrap()) + .collect() +} + +/// Sequential fallback when the prover is built without its default +/// `parallel` feature. +#[cfg(all(not(feature = "cuda"), not(feature = "parallel")))] +fn run_admitted( + order: &[usize], + _estimates: &[u64], + _gate: &VramGate, + _workers: usize, + task: impl Fn(usize) -> T + Send + Sync, +) -> Vec> { + let mut results: Vec> = (0..order.len()).map(|_| None).collect(); + for &idx in order { + results[idx] = Some(task(idx)); + } + results +} + /// Table indices sorted heaviest-first by estimate. fn heaviest_first(estimates: &[u64]) -> Vec { let mut order: Vec = (0..estimates.len()).collect(); @@ -3624,12 +3679,37 @@ pub trait IsStarkProver< // 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"))] + #[cfg(all(feature = "cuda", 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) }); + // CPU has no device wait to hide. Keep a phase barrier between the + // aux/commit work and rounds 2-4 so Rayon can finish one homogeneous + // region before starting the next one, as it did before the GPU + // scheduler was introduced. The fused path above is intentionally + // CUDA-only; interleaving these CPU-heavy regions regresses the host + // prover through cache and memory-bandwidth contention. + #[cfg(all(not(feature = "cuda"), not(feature = "debug-checks")))] + let table_results = { + let aux_outs = run_admitted(&peak_order, &peak_estimates, &vram_gate, k, aux_stage); + let mut staged = Vec::with_capacity(num_airs); + for out in aux_outs { + staged.push(std::sync::Mutex::new(Some( + out.expect("run_admitted fills every slot")?, + ))); + } + run_admitted(&peak_order, &peak_estimates, &vram_gate, k, |idx| { + let (commitment, lde) = staged[idx] + .lock() + .unwrap() + .take() + .expect("aux result consumed once per table"); + rounds_stage(idx, commitment, lde) + }) + }; + // 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. From 5bba8d4cbc6da52940c8b11f268f7e3d48e13bc8 Mon Sep 17 00:00:00 2001 From: jotabulacios Date: Wed, 2 Sep 2026 18:15:22 -0300 Subject: [PATCH 2/3] review fixes --- Makefile | 6 ++-- crypto/stark/src/instruments.rs | 11 ++++-- crypto/stark/src/prover.rs | 59 ++++++++++++++++++++++----------- prover/src/auto_storage.rs | 10 +++--- 4 files changed, 57 insertions(+), 29 deletions(-) diff --git a/Makefile b/Makefile index 3e4a88ecb..be92871a4 100644 --- a/Makefile +++ b/Makefile @@ -666,7 +666,7 @@ check: clippy: cargo clippy --workspace --all-targets -- -D warnings -A clippy::op_ref cargo clippy --workspace --all-targets --no-default-features --features lambda-vm-prover/debug-checks -- -D warnings -A clippy::op_ref - cargo clippy --workspace --all-targets --features lambda-vm-prover/disk-spill -- -D warnings -A clippy::op_ref + cargo clippy --workspace --all-targets --features lambda-vm-prover/disk-spill,lambda-vm-prover/instruments -- -D warnings -A clippy::op_ref fmt: cargo fmt --all @@ -676,7 +676,9 @@ lint: cargo fmt --check --all cargo clippy --workspace --all-targets -- -D warnings -A clippy::op_ref cargo clippy --workspace --all-targets --no-default-features --features lambda-vm-prover/debug-checks -- -D warnings -A clippy::op_ref - cargo clippy --workspace --all-targets --features lambda-vm-prover/disk-spill -- -D warnings -A clippy::op_ref + # `instruments` rides on this pass rather than adding a fifth one: it gates the per-table + # timing plumbing, which no other pass compiles, so breakage there used to reach main. + cargo clippy --workspace --all-targets --features lambda-vm-prover/disk-spill,lambda-vm-prover/instruments -- -D warnings -A clippy::op_ref # The cuda feature gates whole modules + cuda-only integration tests. build.rs emits empty # cubin stubs when nvcc is absent, so this checks on a GPU-less host (CI lint runner, dev laptop) # too — no GPU required. Catches cuda-gated breakage that the non-cuda passes above miss. diff --git a/crypto/stark/src/instruments.rs b/crypto/stark/src/instruments.rs index 0f68059f4..6ef06a387 100644 --- a/crypto/stark/src/instruments.rs +++ b/crypto/stark/src/instruments.rs @@ -362,8 +362,15 @@ pub fn take_r1_sub() -> Round1SubOps { /// Note: thread local stores (R2_SUB, R4_SUB, ROUND_SUB_OPS) are only cleared /// for the calling thread. Rayon worker threads are not reset, so stale data is /// possible if a previous run panicked without consuming stored values. -/// In practice this is safe because store/take pairs always execute within the -/// same rayon task closure. +/// +/// Each store/take pair must also stay adjacent on one thread, with no rayon +/// region in between. The CPU table scheduler runs a whole table's task as a +/// rayon item, so a worker that blocks inside one table's nested parallelism +/// can run a sibling table's task on top of that frame, consume this thread's +/// slot, and leave the outer table reporting zeros. Measured with a probe: no +/// loss while the tables in flight stay under the pool's thread count (the +/// default `k` does), 7-10 % of them lost once they exceed it — e.g. with +/// `TABLE_PARALLELISM` raised past the core count. pub fn reset_all() { R1_MAIN_LDE_US.store(0, Ordering::Relaxed); R1_MAIN_MERKLE_US.store(0, Ordering::Relaxed); diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index b19d0b282..051cb3186 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -269,12 +269,15 @@ where /// 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 the -/// scheduler's `k` coexist (O(k × aux_cols × lde_size)) — which under `cuda` -/// is `num_airs`, so there they are all-N-live like the main ones. +/// - aux: all N are live at once too in every configuration built today +/// (O(N × aux_cols × lde_size), and aux is ext3, so 24 B per element). +/// Under `cuda` the fused task holds only the scheduler's `k` of them, but +/// `k` there is `num_airs`. CPU builds keep a phase barrier between the +/// aux/commit stage and rounds 2-4 and stage every table's aux LDE across +/// it, so `k` does not bound this peak there at all. /// -/// 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. +/// Under `debug-checks` the fused task is split on `cuda` as well, around the +/// cross-table bus balance check. struct Lde { /// Row-major main LDE buffer + its column count. main: (Vec>, usize), @@ -702,6 +705,11 @@ fn estimate_table_vram_bytes(main_cols: usize, aux_cols: usize, lde_size: usize) /// 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. +/// +/// Device-only mechanism: CPU builds have no device allocation to admit +/// against, so their `run_admitted` arms take the gate and ignore it. That is +/// what the `allow(dead_code)` attributes here are for — `multi_prove` still +/// constructs the gate unconditionally so the three arms share one signature. #[cfg_attr(not(feature = "cuda"), allow(dead_code))] struct VramGate { used: std::sync::Mutex, @@ -757,7 +765,7 @@ fn run_admitted( estimates: &[u64], gate: &VramGate, workers: usize, - task: impl Fn(usize) -> T + Send + Sync, + task: impl Fn(usize) -> T + Sync, ) -> Vec> { let results: Vec>> = estimates .iter() @@ -796,27 +804,24 @@ fn run_admitted( #[cfg(all(not(feature = "cuda"), feature = "parallel"))] fn run_admitted( order: &[usize], - _estimates: &[u64], + estimates: &[u64], _gate: &VramGate, workers: usize, - task: impl Fn(usize) -> T + Send + Sync, + task: impl Fn(usize) -> T + Sync, ) -> Vec> { - let results: Vec>> = (0..order.len()) - .map(|_| std::sync::Mutex::new(None)) - .collect(); + // No interior mutability here, unlike the `cuda` arm: `collect()` joins + // each chunk, so every slot is written from this thread. + let mut results: Vec> = (0..estimates.len()).map(|_| None).collect(); for chunk in order.chunks(workers.max(1)) { let chunk_results: Vec<(usize, T)> = chunk.par_iter().map(|&idx| (idx, task(idx))).collect(); for (idx, result) in chunk_results { - *results[idx].lock().unwrap() = Some(result); + results[idx] = Some(result); } } results - .into_iter() - .map(|m| m.into_inner().unwrap()) - .collect() } /// Sequential fallback when the prover is built without its default @@ -824,12 +829,12 @@ fn run_admitted( #[cfg(all(not(feature = "cuda"), not(feature = "parallel")))] fn run_admitted( order: &[usize], - _estimates: &[u64], + estimates: &[u64], _gate: &VramGate, _workers: usize, - task: impl Fn(usize) -> T + Send + Sync, + task: impl Fn(usize) -> T + Sync, ) -> Vec> { - let mut results: Vec> = (0..order.len()).map(|_| None).collect(); + let mut results: Vec> = (0..estimates.len()).map(|_| None).collect(); for &idx in order { results[idx] = Some(task(idx)); } @@ -4037,6 +4042,12 @@ pub trait IsStarkProver< // scheduler was introduced. The fused path above is intentionally // CUDA-only; interleaving these CPU-heavy regions regresses the host // prover through cache and memory-bandwidth contention. + // + // Cost of the barrier: every table's Round 1 result is staged until its + // rounds run, so the aux LDE *and* the aux Merkle tree are live for all + // N tables where the fused path held at most `k`. Under `disk-spill`'s + // Disk mode the trees are spilled but the aux LDE is a plain `Vec` with + // no spill path, so Disk does not bound this peak. #[cfg(all(not(feature = "cuda"), not(feature = "debug-checks")))] let table_results = { let aux_outs = run_admitted(&peak_order, &peak_estimates, &vram_gate, k, aux_stage); @@ -4563,6 +4574,15 @@ pub trait IsStarkProver< &boundary_coefficients, )?; + // Round 2's sub-op timings travel in a thread-local, so take them here + // instead of after round 4: this thread re-enters rayon in rounds 3 and + // 4, and a worker that blocks there can run another table's whole + // scheduler task on top of this frame and consume the slot. Keeping + // store and take adjacent is what pins the pair to one thread — see + // `instruments::reset_all`. + #[cfg(feature = "instruments")] + let r2_sub = crate::instruments::take_r2_sub(); + // >>>> Send commitments: [H₁], [H₂] transcript.append_bytes(&round_2_result.composition_poly_root); @@ -4659,8 +4679,7 @@ pub trait IsStarkProver< #[cfg(feature = "instruments")] { let zero = Duration::ZERO; - let (r2_constraints, r2_fft, r2_merkle) = - crate::instruments::take_r2_sub().unwrap_or((zero, zero, zero)); + let (r2_constraints, r2_fft, r2_merkle) = r2_sub.unwrap_or((zero, zero, zero)); let (r4_fft, r4_merkle, r4_deep_comp, r4_queries) = crate::instruments::take_r4_sub().unwrap_or((zero, zero, zero, zero)); crate::instruments::store_round_sub_ops(crate::instruments::TableSubOps { diff --git a/prover/src/auto_storage.rs b/prover/src/auto_storage.rs index b4718974c..ddf64f0c1 100644 --- a/prover/src/auto_storage.rs +++ b/prover/src/auto_storage.rs @@ -243,11 +243,11 @@ pub fn peak_bytes(lengths: &TableLengths, blowup_factor: u8, table_parallelism: let specs = table_specs(lengths); // 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 is produced and - // consumed inside one table's fused task, so only the scheduler's k coexist - // — exactly all of them on `cuda`, fewer on CPU builds. Counted for every - // table either way, which is exact on `cuda` and an over-estimate on CPU - // rather than an unsound bound. + // Round 1 main commit is a phase-wide barrier). Counting the aux LDE and + // its Merkle tree for every table is exact rather than an over-estimate: + // the fused task on `cuda` holds only the scheduler's k of them, but k + // there is num_airs, and CPU builds stage every table's Round 1 result + // across the phase barrier between the aux/commit stage and rounds 2-4. let persistent_total: u64 = specs .iter() .map(|s| persistent_per_table(*s, blowup)) From 3b3cc869dc985f32601dc082c7c99227bd9baac8 Mon Sep 17 00:00:00 2001 From: jotabulacios Date: Fri, 4 Sep 2026 12:17:32 -0300 Subject: [PATCH 3/3] Schedule the CPU table work inside the rayon pool, with no phase barrier. --- crypto/stark/src/prover.rs | 139 ++++++++++++++++++++++++++----------- prover/src/auto_storage.rs | 10 +-- 2 files changed, 104 insertions(+), 45 deletions(-) diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index 051cb3186..6aee481b9 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -269,15 +269,12 @@ where /// 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: all N are live at once too in every configuration built today -/// (O(N × aux_cols × lde_size), and aux is ext3, so 24 B per element). -/// Under `cuda` the fused task holds only the scheduler's `k` of them, but -/// `k` there is `num_airs`. CPU builds keep a phase barrier between the -/// aux/commit stage and rounds 2-4 and stage every table's aux LDE across -/// it, so `k` does not bound this peak there at all. +/// - aux: produced and consumed inside the same fused task, so at most the +/// scheduler's `k` coexist (O(k × aux_cols × lde_size)) — which under `cuda` +/// is `num_airs`, so there they are all-N-live like the main ones. /// -/// Under `debug-checks` the fused task is split on `cuda` as well, around the -/// cross-table bus balance check. +/// 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), @@ -754,6 +751,26 @@ impl Drop for VramPermit<'_> { } } +/// Debug-only contract check for `run_admitted`: the result slots are addressed +/// by `order`'s *values*, so every value must index `estimates` and appear at +/// most once. Callers pass either a full permutation or one group of one, and a +/// grouping bug would otherwise surface as a confusing `take()` panic inside a +/// task rather than here. +fn debug_check_order(order: &[usize], estimates: &[u64]) { + if cfg!(debug_assertions) { + let mut seen = vec![false; estimates.len()]; + for &idx in order { + assert!( + idx < estimates.len(), + "run_admitted: order names table {idx}, estimates has {}", + estimates.len() + ); + assert!(!seen[idx], "run_admitted: order names table {idx} twice"); + seen[idx] = true; + } + } +} + /// 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 @@ -767,6 +784,7 @@ fn run_admitted( workers: usize, task: impl Fn(usize) -> T + Sync, ) -> Vec> { + debug_check_order(order, estimates); let results: Vec>> = estimates .iter() .map(|_| std::sync::Mutex::new(None)) @@ -809,6 +827,7 @@ fn run_admitted( workers: usize, task: impl Fn(usize) -> T + Sync, ) -> Vec> { + debug_check_order(order, estimates); // No interior mutability here, unlike the `cuda` arm: `collect()` joins // each chunk, so every slot is written from this thread. let mut results: Vec> = (0..estimates.len()).map(|_| None).collect(); @@ -834,6 +853,7 @@ fn run_admitted( _workers: usize, task: impl Fn(usize) -> T + Sync, ) -> Vec> { + debug_check_order(order, estimates); let mut results: Vec> = (0..estimates.len()).map(|_| None).collect(); for &idx in order { results[idx] = Some(task(idx)); @@ -4030,43 +4050,12 @@ pub trait IsStarkProver< // 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(all(feature = "cuda", not(feature = "debug-checks")))] + #[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) }); - // CPU has no device wait to hide. Keep a phase barrier between the - // aux/commit work and rounds 2-4 so Rayon can finish one homogeneous - // region before starting the next one, as it did before the GPU - // scheduler was introduced. The fused path above is intentionally - // CUDA-only; interleaving these CPU-heavy regions regresses the host - // prover through cache and memory-bandwidth contention. - // - // Cost of the barrier: every table's Round 1 result is staged until its - // rounds run, so the aux LDE *and* the aux Merkle tree are live for all - // N tables where the fused path held at most `k`. Under `disk-spill`'s - // Disk mode the trees are spilled but the aux LDE is a plain `Vec` with - // no spill path, so Disk does not bound this peak. - #[cfg(all(not(feature = "cuda"), not(feature = "debug-checks")))] - let table_results = { - let aux_outs = run_admitted(&peak_order, &peak_estimates, &vram_gate, k, aux_stage); - let mut staged = Vec::with_capacity(num_airs); - for out in aux_outs { - staged.push(std::sync::Mutex::new(Some( - out.expect("run_admitted fills every slot")?, - ))); - } - run_admitted(&peak_order, &peak_estimates, &vram_gate, k, |idx| { - let (commitment, lde) = staged[idx] - .lock() - .unwrap() - .take() - .expect("aux result consumed once per table"); - rounds_stage(idx, commitment, lde) - }) - }; - // 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. @@ -4839,3 +4828,73 @@ fn print_bus_balance_report( } } } + +#[cfg(test)] +mod scheduler_tests { + use super::*; + use std::sync::atomic::{AtomicUsize, Ordering}; + + /// Runs `run_admitted` over `order` with `n` tables and returns how many + /// times each table's task ran, plus the slots it wrote. + fn calls_and_results( + order: &[usize], + n: usize, + workers: usize, + ) -> (Vec, Vec>) { + let estimates = vec![1u64; n]; + let gate = VramGate::new(u64::MAX); + let calls: Vec = (0..n).map(|_| AtomicUsize::new(0)).collect(); + let out = run_admitted(order, &estimates, &gate, workers, |idx| { + calls[idx].fetch_add(1, Ordering::Relaxed); + idx + }); + ( + calls.iter().map(|c| c.load(Ordering::Relaxed)).collect(), + out, + ) + } + + /// The scheduler's contract, independent of how many tables run at once: + /// each table's task runs exactly once and its result lands on its own + /// index. Nothing else pins this, and a chunking or sizing bug would only + /// show up on hosts with a particular core count. + #[test] + fn every_table_runs_once_and_lands_on_its_own_slot() { + for n in [1usize, 2, 5, 21] { + // Heaviest-first, which is the order `multi_prove` passes. + let estimates: Vec = (0..n).map(|i| (n - i) as u64).collect(); + let order = heaviest_first(&estimates); + for workers in [1usize, 2, 3, n, n + 5] { + let (calls, out) = calls_and_results(&order, n, workers); + assert!( + calls.iter().all(|&c| c == 1), + "n={n} workers={workers} calls={calls:?}" + ); + for (idx, slot) in out.iter().enumerate() { + assert_eq!(*slot, Some(idx), "n={n} workers={workers}"); + } + } + } + } + + /// What the `Staged(depth)` arm passes once the barrier runs over groups: + /// a subset of the permutation. The untouched tables must stay empty. + #[test] + fn a_group_fills_only_its_own_slots() { + let n = 9; + let group = [7usize, 1, 4]; + let (calls, out) = calls_and_results(&group, n, 2); + for idx in 0..n { + let ran = group.contains(&idx); + assert_eq!(calls[idx], ran as usize, "table {idx}"); + assert_eq!(out[idx], ran.then_some(idx), "table {idx}"); + } + } + + #[test] + fn an_empty_order_runs_nothing() { + let (calls, out) = calls_and_results(&[], 4, 3); + assert!(calls.iter().all(|&c| c == 0)); + assert!(out.iter().all(|s| s.is_none())); + } +} diff --git a/prover/src/auto_storage.rs b/prover/src/auto_storage.rs index ddf64f0c1..b4718974c 100644 --- a/prover/src/auto_storage.rs +++ b/prover/src/auto_storage.rs @@ -243,11 +243,11 @@ pub fn peak_bytes(lengths: &TableLengths, blowup_factor: u8, table_parallelism: let specs = table_specs(lengths); // Persistent: every table's main LDE + Merkle really is alive at once (the - // Round 1 main commit is a phase-wide barrier). Counting the aux LDE and - // its Merkle tree for every table is exact rather than an over-estimate: - // the fused task on `cuda` holds only the scheduler's k of them, but k - // there is num_airs, and CPU builds stage every table's Round 1 result - // across the phase barrier between the aux/commit stage and rounds 2-4. + // Round 1 main commit is a phase-wide barrier). The aux LDE is produced and + // consumed inside one table's fused task, so only the scheduler's k coexist + // — exactly all of them on `cuda`, fewer on CPU builds. Counted for every + // table either way, which is exact on `cuda` and an over-estimate on CPU + // rather than an unsound bound. let persistent_total: u64 = specs .iter() .map(|s| persistent_per_table(*s, blowup))