diff --git a/crypto/stark/src/fri/fri_functions.rs b/crypto/stark/src/fri/fri_functions.rs index 8bd355ec4..b8c512d52 100644 --- a/crypto/stark/src/fri/fri_functions.rs +++ b/crypto/stark/src/fri/fri_functions.rs @@ -5,6 +5,14 @@ use math::field::{ element::FieldElement, traits::{IsFFTField, IsField, IsSubFieldOf}, }; +#[cfg(feature = "parallel")] +use rayon::prelude::*; + +/// Minimum number of elements for parallel FRI fold. Below this threshold +/// the overhead of Rayon dispatch exceeds the benefit. +#[cfg(feature = "parallel")] +const FRI_FOLD_PARALLEL_THRESHOLD: usize = 4096; + /// Evaluation-form FRI fold: given evaluations in bit-reversed order where /// consecutive pairs (2j, 2j+1) are conjugates (p(x_j), p(-x_j)), compute /// the folded evaluations: (lo + hi) + inv_twiddle[j] * zeta * (lo - hi) @@ -12,12 +20,35 @@ use math::field::{ /// /// After folding, the N/2 results are evaluations on the squared coset /// in bit-reversed order, preserving conjugate pairing for the next fold. -pub fn fold_evaluations_in_place, E: IsField>( +pub fn fold_evaluations_in_place( evals: &mut Vec>, zeta: &FieldElement, inv_twiddles: &[FieldElement], -) { +) where + F: IsSubFieldOf + Send + Sync, + E: IsField + Send + Sync, + FieldElement: Send + Sync, + FieldElement: Send + Sync, +{ let half = evals.len() / 2; + + #[cfg(feature = "parallel")] + if half >= FRI_FOLD_PARALLEL_THRESHOLD { + // Parallel: compute folded values into a new buffer to avoid + // read/write aliasing on the same slice. + let folded: Vec> = evals + .par_chunks_exact(2) + .zip(inv_twiddles.par_iter()) + .map(|(pair, tw)| { + let sum = &pair[0] + &pair[1]; + let diff = &pair[0] - &pair[1]; + &sum + &(tw * &(zeta * &diff)) + }) + .collect(); + *evals = folded; + return; + } + for j in 0..half { let lo = &evals[2 * j]; let hi = &evals[2 * j + 1]; diff --git a/crypto/stark/src/fri/mod.rs b/crypto/stark/src/fri/mod.rs index 87ab66a5b..cfae88ab1 100644 --- a/crypto/stark/src/fri/mod.rs +++ b/crypto/stark/src/fri/mod.rs @@ -19,7 +19,7 @@ use self::fri_functions::{ /// FRI commit phase from pre-computed bit-reversed evaluations. /// skipping the initial FFT. Use this when the caller already has the evaluation /// vector (e.g. from a fused LDE pipeline). -pub fn commit_phase_from_evaluations, E: IsField>( +pub fn commit_phase_from_evaluations + Send + Sync, E: IsField + Send + Sync>( number_layers: usize, mut evals: Vec>, transcript: &mut impl IsStarkTranscript,