Skip to content

Disk spill with scaling improvements #486

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
gabrielbosio wants to merge 8 commits into from
Closed

Disk spill with scaling improvements #486

Show file tree
Hide file tree
Changes from 5 commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
cdf1784
perf(tables): cap LT/LOAD/BRANCH/MEMW_R max_rows at 2^20
diegokingston Apr 9, 2026
a639710
perf(stark): deduplicate LdeTwiddles by domain size in multi_prove
diegokingston Apr 9, 2026
4c29c10
perf(fri): parallelize fold_evaluations_in_place with Rayon
diegokingston Apr 9, 2026
a809ced
perf(prover): eliminate per-row Vec alloc in commit_columns_bit_reversed
diegokingston Apr 9, 2026
0ce8a63
Fix fmt and clippy after cherry-pick
gabrielbosio Apr 9, 2026
c6058b9
Remove redundant capped annotations from max_rows
gabrielbosio Apr 9, 2026
48c1be0
Document twiddle dedup invariant
gabrielbosio Apr 10, 2026
bddf96a
Merge remote-tracking branch 'origin/perf/disk-spill-time-regression'…
gabrielbosio Apr 13, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
38 changes: 30 additions & 8 deletions crypto/stark/src/fri/fri_functions.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,15 +17,37 @@ pub fn fold_evaluations_in_place<F: IsSubFieldOf<E>, E: IsField>(
zeta: &FieldElement<E>,
inv_twiddles: &[FieldElement<F>],
) {
let half = evals.len() / 2;
for j in 0..half {
let lo = &evals[2 * j];
let hi = &evals[2 * j + 1];
let sum = lo + hi;
let diff = lo - hi;
evals[j] = &sum + &(&inv_twiddles[j] * &(zeta * &diff));
#[cfg(feature = "parallel")]
{
use rayon::prelude::*;
// Evaluations are stored interleaved: pairs (evals[2j], evals[2j+1]).
// Fold each pair in parallel, collecting into a new Vec of half length.
let folded: Vec<FieldElement<E>> = evals
.par_chunks(2)
.zip(inv_twiddles.par_iter())
Comment thread
gabrielbosio marked this conversation as resolved.
.map(|(pair, tw)| {
let lo = &pair[0];
let hi = &pair[1];
let sum = lo + hi;
let diff = lo - hi;
&sum + &(tw * &(zeta * &diff))
})
.collect();
*evals = folded;
}

#[cfg(not(feature = "parallel"))]
{
let half = evals.len() / 2;
for j in 0..half {
let lo = &evals[2 * j];
let hi = &evals[2 * j + 1];
let sum = lo + hi;
let diff = lo - hi;
evals[j] = &sum + &(&inv_twiddles[j] * &(zeta * &diff));
}
evals.truncate(half);
}
evals.truncate(half);
}

/// Compute inverse twiddle factors for evaluation-form FRI folding.
Expand Down
68 changes: 48 additions & 20 deletions crypto/stark/src/prover.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -335,19 +335,35 @@ pub trait IsStarkProver<
let num_cols = columns.len();

#[cfg(feature = "parallel")]
let iter = (0..num_rows).into_par_iter();
#[cfg(not(feature = "parallel"))]
let iter = 0..num_rows;
let hashed_leaves: Vec<Commitment> = {
(0..num_rows)
.into_par_iter()
.map_init(
|| vec![FieldElement::<E>::zero(); num_cols],
|row_buf, row_idx| {
let br_idx = reverse_index(row_idx, num_rows as u64);
for col_idx in 0..num_cols {
row_buf[col_idx] = columns[col_idx][br_idx].clone();
}
BatchedMerkleTreeBackend::<E>::hash_data(row_buf)
},
)
.collect()
};

let hashed_leaves: Vec<Commitment> = iter
.map(|row_idx| {
let br_idx = reverse_index(row_idx, num_rows as u64);
let row: Vec<FieldElement<E>> = (0..num_cols)
.map(|col_idx| columns[col_idx][br_idx].clone())
.collect();
BatchedMerkleTreeBackend::<E>::hash_data(&row)
})
.collect();
#[cfg(not(feature = "parallel"))]
let hashed_leaves: Vec<Commitment> = {
let mut row_buf = vec![FieldElement::<E>::zero(); num_cols];
(0..num_rows)
.map(|row_idx| {
let br_idx = reverse_index(row_idx, num_rows as u64);
for col_idx in 0..num_cols {
row_buf[col_idx] = columns[col_idx][br_idx].clone();
}
BatchedMerkleTreeBackend::<E>::hash_data(&row_buf)
})
.collect()
};

let tree = BatchedMerkleTree::<E>::build_from_hashed_leaves(hashed_leaves)?;
let root = tree.root;
Expand Down Expand Up @@ -660,7 +676,7 @@ pub trait IsStarkProver<
air_trace_pairs: &[AirTracePair<'_, Field, FieldExtension, PI>],
metadatas: &[Round1Metadata<Field, FieldExtension>],
domains: &[Domain<Field>],
twiddle_caches: &[LdeTwiddles<Field>],
twiddle_caches: &[Arc<LdeTwiddles<Field>>],
main_pool: &mut [Vec<FieldElement<Field>>],
aux_pool: &mut [Vec<FieldElement<FieldExtension>>],
) where
Expand All @@ -676,7 +692,13 @@ pub trait IsStarkProver<
.zip(domains.iter().zip(twiddle_caches.iter()))
{
let result = Self::reconstruct_round1(
*air, *trace, domain, metadata, twiddles, main_pool, aux_pool,
*air,
*trace,
domain,
metadata,
&**twiddles,
main_pool,
aux_pool,
)
.expect("reconstruct_round1 failed in debug-checks");
temp_results.push(result);
Expand Down Expand Up @@ -1534,24 +1556,30 @@ pub trait IsStarkProver<
let phase_start = Instant::now();

let mut domains = Vec::with_capacity(num_airs);
let mut twiddle_caches: Vec<LdeTwiddles<Field>> = Vec::with_capacity(num_airs);
let mut twiddle_caches: Vec<Arc<LdeTwiddles<Field>>> = Vec::with_capacity(num_airs);
let mut max_main_cols = 0usize;
let mut max_aux_cols = 0usize;
let mut max_lde_size = 0usize;

// Deduplicate twiddle caches: tables with the same lde_size share one Arc.
let mut twiddle_by_size: std::collections::HashMap<usize, Arc<LdeTwiddles<Field>>> =
std::collections::HashMap::new();

for (air, trace, _pub_inputs) in &*air_trace_pairs {
let trace_length = trace.num_rows();
let domain = new_domain(*air, trace_length);

let lde_size = domain.interpolation_domain_size * domain.blowup_factor;
let twiddles = LdeTwiddles::new(&domain);
let twiddles = twiddle_by_size
.entry(lde_size)
.or_insert_with(|| Arc::new(LdeTwiddles::new(&domain)));
Comment thread
gabrielbosio marked this conversation as resolved.
Comment thread
gabrielbosio marked this conversation as resolved.
Outdated

max_main_cols = max_main_cols.max(trace.num_main_columns);
max_aux_cols = max_aux_cols.max(air.num_auxiliary_rap_columns());
max_lde_size = max_lde_size.max(lde_size);

domains.push(domain);
twiddle_caches.push(twiddles);
twiddle_caches.push(Arc::clone(twiddles));
}

// Spill all main trace tables to mmap before allocating pool buffers.
Expand Down Expand Up @@ -1631,7 +1659,7 @@ pub trait IsStarkProver<
let idx = chunk_start + j;
let (air, trace, _) = &air_trace_pairs[idx];
let domain = &domains[idx];
let twiddles = &twiddle_caches[idx];
let twiddles = &*twiddle_caches[idx];

#[allow(unused_mut)]
let (mut tree, root, pre_tree, pre_root, n_pre) = if air.is_preprocessed() {
Expand Down Expand Up @@ -1828,7 +1856,7 @@ pub trait IsStarkProver<
let idx = chunk_start + j;
let (air, trace, _) = &air_trace_pairs[idx];
let domain = &domains[idx];
let twiddles = &twiddle_caches[idx];
let twiddles = &*twiddle_caches[idx];

if air.has_aux_trace() {
let num_aux_cols = trace.num_aux_columns;
Expand Down Expand Up @@ -2080,7 +2108,7 @@ pub trait IsStarkProver<
let (air, trace, pub_inputs) = &air_trace_pairs[idx];
let metadata = &metadatas[idx];
let domain = &domains[idx];
let twiddles = &twiddle_caches[idx];
let twiddles = &*twiddle_caches[idx];

#[cfg(feature = "instruments")]
let table_start = Instant::now();
Expand Down
32 changes: 16 additions & 16 deletions prover/src/tables/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,29 +49,29 @@ pub use types::BusId;
/// (* MEMW_A formula gives 2^20, but set to 2^19 to match MEMW chunk geometry;
/// benchmarks show better parallel throughput with smaller chunks.)
///
/// | Table | Main | Bus | Eff.width | Max rows |
/// |---------|------|-----|-----------|----------|
/// | MEMW | 49 | 26 | 127 | 2^19 |
/// | MEMW_A | 29 | 20 | 89 | 2^19 * |
/// | CPU | 74 | 40 | 194 | 2^19 |
/// | DVRM | 34 | 34 | 136 | 2^19 |
/// | MUL | 26 | 16 | 74 | 2^20 |
/// | LT | 15 | 9 | 42 | 2^21 |
/// | SHIFT | 27 | 15 | 72 | 2^20 |
/// | LOAD | 18 | 5 | 33 | 2^21 |
/// | BRANCH | 14 | 6 | 32 | 2^21 |
/// | MEMW_R | 10 | 7 | 31 | 2^21 |
/// | Table | Main | Bus | Eff.width | Max rows |
/// |---------|------|-----|-----------|-----------------|
/// | MEMW | 49 | 26 | 127 | 2^19 |
/// | MEMW_A | 29 | 20 | 89 | 2^19 * |
/// | CPU | 74 | 40 | 194 | 2^19 |
/// | DVRM | 34 | 34 | 136 | 2^19 |
/// | MUL | 26 | 16 | 74 | 2^20 |
/// | LT | 15 | 9 | 42 | 2^20 (capped) |
/// | SHIFT | 27 | 15 | 72 | 2^20 |
/// | LOAD | 18 | 5 | 33 | 2^20 (capped) |
/// | BRANCH | 14 | 6 | 32 | 2^20 (capped) |
/// | MEMW_R | 10 | 7 | 31 | 2^20 (capped) |
pub mod max_rows {
pub const CPU: usize = 1 << 19; // 524,288 — eff. width 194
pub const MEMW: usize = 1 << 19; // 524,288 — eff. width 127 (baseline)
pub const MEMW_A: usize = 1 << 19; // 524,288 — eff. width 89
pub const DVRM: usize = 1 << 19; // 524,288 — eff. width 136
pub const MUL: usize = 1 << 20; // 1,048,576 — eff. width 74
pub const LT: usize = 1 << 21; // 2,097,152 — eff. width 42
pub const LT: usize = 1 << 20; // 1,048,576 — eff. width 42 (capped at 2^20)
pub const SHIFT: usize = 1 << 20; // 1,048,576 — eff. width 72
pub const LOAD: usize = 1 << 21; // 2,097,152 — eff. width 33
pub const BRANCH: usize = 1 << 21; // 2,097,152 — eff. width 32
pub const MEMW_R: usize = 1 << 21; // 2,097,152 — eff. width 31
pub const LOAD: usize = 1 << 20; // 1,048,576 — eff. width 33 (capped at 2^20)
pub const BRANCH: usize = 1 << 20; // 1,048,576 — eff. width 32 (capped at 2^20)
pub const MEMW_R: usize = 1 << 20; // 1,048,576 — eff. width 31 (capped at 2^20)
}

/// Per-table maximum row limits, configurable for different environments.
Expand Down
Loading