diff --git a/crypto/stark/src/fri/fri_functions.rs b/crypto/stark/src/fri/fri_functions.rs index 8bd355ec4..53d140782 100644 --- a/crypto/stark/src/fri/fri_functions.rs +++ b/crypto/stark/src/fri/fri_functions.rs @@ -14,18 +14,45 @@ use math::field::{ /// in bit-reversed order, preserving conjugate pairing for the next fold. pub fn fold_evaluations_in_place, E: IsField>( evals: &mut Vec>, + #[cfg_attr(not(feature = "parallel"), allow(unused_variables))] scratch: &mut Vec< + FieldElement, + >, zeta: &FieldElement, inv_twiddles: &[FieldElement], ) { 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]). + // Index-based iteration matches the sequential path and panics on a + // short `inv_twiddles`; `par_chunks(2).zip(inv_twiddles)` would drop + // a trailing odd element and silently truncate on length mismatch. + (0..half) + .into_par_iter() + .map(|j| { + let lo = &evals[2 * j]; + let hi = &evals[2 * j + 1]; + let sum = lo + hi; + let diff = lo - hi; + &sum + &(&inv_twiddles[j] * &(zeta * &diff)) + }) + .collect_into_vec(scratch); + std::mem::swap(evals, scratch); + } + + #[cfg(not(feature = "parallel"))] + { + 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. diff --git a/crypto/stark/src/fri/mod.rs b/crypto/stark/src/fri/mod.rs index 87ab66a5b..7855c1916 100644 --- a/crypto/stark/src/fri/mod.rs +++ b/crypto/stark/src/fri/mod.rs @@ -39,6 +39,8 @@ where let mut fri_layer_list = Vec::with_capacity(number_layers); let mut current_coset_offset = coset_offset.clone(); let mut current_domain_size = domain_size; + // Scratch buffer reused across fold levels to avoid reallocating once per level. + let mut scratch: Vec> = Vec::with_capacity(evals.len() / 2); for _ in 1..number_layers { // <<<< Receive challenge 𝜁ₖ₋₁ @@ -47,7 +49,7 @@ where current_domain_size /= 2; // Fold evaluations in-place (no FFT needed) - fold_evaluations_in_place(&mut evals, &zeta, &inv_twiddles); + fold_evaluations_in_place(&mut evals, &mut scratch, &zeta, &inv_twiddles); // Build Merkle tree from consecutive pairs let leaves: Vec<[FieldElement; 2]> = evals @@ -75,7 +77,7 @@ where let zeta = transcript.sample_field_element(); // Final fold - fold_evaluations_in_place(&mut evals, &zeta, &inv_twiddles); + fold_evaluations_in_place(&mut evals, &mut scratch, &zeta, &inv_twiddles); let last_value = evals.first().unwrap_or(&FieldElement::zero()).clone(); diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index 50af95525..b866e5571 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -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 = { + (0..num_rows) + .into_par_iter() + .map_init( + || vec![FieldElement::::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::::hash_data(&*row_buf) + }, + ) + .collect() + }; - let hashed_leaves: Vec = iter - .map(|row_idx| { - let br_idx = reverse_index(row_idx, num_rows as u64); - let row: Vec> = (0..num_cols) - .map(|col_idx| columns[col_idx][br_idx].clone()) - .collect(); - BatchedMerkleTreeBackend::::hash_data(&row) - }) - .collect(); + #[cfg(not(feature = "parallel"))] + let hashed_leaves: Vec = { + let mut row_buf = vec![FieldElement::::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::::hash_data(&row_buf) + }) + .collect() + }; let tree = BatchedMerkleTree::::build_from_hashed_leaves(hashed_leaves)?; let root = tree.root; @@ -635,7 +651,7 @@ pub trait IsStarkProver< air_trace_pairs: &[AirTracePair<'_, Field, FieldExtension, PI>], metadatas: &[Round1Metadata], domains: &[Domain], - twiddle_caches: &[LdeTwiddles], + twiddle_caches: &[Arc>], main_pool: &mut [Vec>], aux_pool: &mut [Vec>], ) where @@ -651,7 +667,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); @@ -1512,24 +1534,42 @@ pub trait IsStarkProver< let phase_start = Instant::now(); let mut domains = Vec::with_capacity(num_airs); - let mut twiddle_caches: Vec> = Vec::with_capacity(num_airs); + let mut twiddle_caches: Vec>> = 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 identical domain parameters share one Arc. + // Key on (interpolation_domain_size, blowup_factor, coset_offset_bytes) because each + // affects the twiddle values: interpolation_domain_size sizes `inv` and `coset_weights`, + // blowup_factor sizes `fwd`, and coset_offset sets `coset_weights` values. + // coset_offset is serialized via AsBytes since FieldElement does not impl Hash without + // a BaseType: Hash bound. + let mut twiddle_by_domain: std::collections::HashMap< + (usize, usize, Vec), + Arc>, + > = 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 key = ( + domain.interpolation_domain_size, + domain.blowup_factor, + domain.coset_offset.as_bytes(), + ); + let twiddles = twiddle_by_domain + .entry(key) + .or_insert_with(|| Arc::new(LdeTwiddles::new(&domain))); 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)); } // Allocate K independent LDE column buffer pool sets for parallel table processing. @@ -1573,7 +1613,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.is_preprocessed() { Self::commit_preprocessed_trace( @@ -1693,7 +1733,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; @@ -1809,7 +1849,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(); diff --git a/crypto/stark/src/tests/fri_tests.rs b/crypto/stark/src/tests/fri_tests.rs index 8012d0b0c..b1eaff1db 100644 --- a/crypto/stark/src/tests/fri_tests.rs +++ b/crypto/stark/src/tests/fri_tests.rs @@ -127,7 +127,8 @@ fn test_eval_fold_matches_coeff_fold() { let mut path_b_evals = evals_fft; in_place_bit_reverse_permute(&mut path_b_evals); let inv_twiddles = compute_coset_twiddles_inv::(&coset_offset, n); - fold_evaluations_in_place(&mut path_b_evals, &beta, &inv_twiddles); + let mut scratch = Vec::new(); + fold_evaluations_in_place(&mut path_b_evals, &mut scratch, &beta, &inv_twiddles); assert_eq!(path_a_evals, path_b_evals); } diff --git a/prover/src/tables/mod.rs b/prover/src/tables/mod.rs index 551dc4aa3..ec8e2f4a9 100644 --- a/prover/src/tables/mod.rs +++ b/prover/src/tables/mod.rs @@ -49,29 +49,34 @@ 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 | +/// Capped tables use 2^20, not the 2^21 the formula gives. A 2^21 chunk +/// runs on a single thread and peaks at twice the memory of a 2^20 chunk. +/// Programs that exceed 2^20 rows in a capped table hit a table-overflow +/// error. +/// +/// | 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.