diff --git a/crypto/stark/src/fri/fri_functions.rs b/crypto/stark/src/fri/fri_functions.rs index 8bd355ec4..b02090533 100644 --- a/crypto/stark/src/fri/fri_functions.rs +++ b/crypto/stark/src/fri/fri_functions.rs @@ -17,15 +17,37 @@ pub fn fold_evaluations_in_place, E: IsField>( 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]). + // Fold each pair in parallel, collecting into a new Vec of half length. + let folded: Vec> = evals + .par_chunks(2) + .zip(inv_twiddles.par_iter()) + .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. diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index e1f0aaf74..5fefbcbd1 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; @@ -660,7 +676,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 @@ -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); @@ -1534,24 +1556,31 @@ 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; + // Share twiddles across tables with equal lde_size. Relies on all AIRs + // using the same ProofOptions (blowup_factor, coset_offset). + let mut twiddle_by_size: std::collections::HashMap>> = + 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))); 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. @@ -1631,7 +1660,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() { @@ -1828,7 +1857,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; @@ -2080,7 +2109,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/prover/src/tables/mod.rs b/prover/src/tables/mod.rs index 551dc4aa3..19d14411d 100644 --- a/prover/src/tables/mod.rs +++ b/prover/src/tables/mod.rs @@ -56,22 +56,22 @@ pub use types::BusId; /// | CPU | 74 | 40 | 194 | 2^19 | /// | DVRM | 34 | 34 | 136 | 2^19 | /// | MUL | 26 | 16 | 74 | 2^20 | -/// | LT | 15 | 9 | 42 | 2^21 | +/// | LT | 15 | 9 | 42 | 2^20 | /// | 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 | +/// | LOAD | 18 | 5 | 33 | 2^20 | +/// | BRANCH | 14 | 6 | 32 | 2^20 | +/// | MEMW_R | 10 | 7 | 31 | 2^20 | 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 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 + pub const BRANCH: usize = 1 << 20; // 1,048,576 — eff. width 32 + pub const MEMW_R: usize = 1 << 20; // 1,048,576 — eff. width 31 } /// Per-table maximum row limits, configurable for different environments.