Skip to content

perf: parallelize FRI fold with Rayon #597

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

Open
MauroToscano wants to merge 1 commit into
base: main
Choose a base branch
from
Open

perf: parallelize FRI fold with Rayon #597

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
35 changes: 33 additions & 2 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 @@ -5,19 +5,50 @@ 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)
/// = 2 * (p_even(x_j²) + zeta * p_odd(x_j²))
///
/// 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<F: IsSubFieldOf<E>, E: IsField>(
pub fn fold_evaluations_in_place<F, E>(
evals: &mut Vec<FieldElement<E>>,
zeta: &FieldElement<E>,
inv_twiddles: &[FieldElement<F>],
) {
) where
F: IsSubFieldOf<E> + Send + Sync,
E: IsField + Send + Sync,
FieldElement<E>: Send + Sync,
FieldElement<F>: Send + Sync,
Comment on lines +27 to +31
Copy link
Copy Markdown

@claude claude Bot May 19, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Low – Send + Sync bounds leak into non-parallel builds

These four bounds are unconditional, so they apply even when the parallel feature is off. Any caller with a field element type that doesn't implement Send + Sync (e.g. a type containing Rc or Cell) will fail to compile, even in a single-threaded build. In practice, all concrete field element types in this crate are Send + Sync, but the bounds still unnecessarily restrict the API surface.

The cleanest fix without duplicating the function body is to gate them with a custom supertrait:

#[cfg(feature = "parallel")]
pub trait MaybeParallel: Send + Sync {}
#[cfg(feature = "parallel")]
impl<T: Send + Sync> MaybeParallel for T {}
#[cfg(not(feature = "parallel"))]
pub trait MaybeParallel {}
#[cfg(not(feature = "parallel"))]
impl<T> MaybeParallel for T {}

then replace Send + Sync with MaybeParallel throughout. Alternatively, if all callers satisfy the bounds unconditionally, document that assumption and leave it as-is — the important thing is the intent is clear.

{
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<FieldElement<E>> = evals
.par_chunks_exact(2)
.zip(inv_twiddles.par_iter())
Comment on lines +40 to +41
Copy link
Copy Markdown

@claude claude Bot May 19, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Low – zip silently truncates if lengths diverge; sequential panics

.par_chunks_exact(2).zip(inv_twiddles.par_iter()) stops at the shorter iterator, so if inv_twiddles.len() < evals.len() / 2 the parallel path silently produces a shorter-than-expected vector and returns with a wrong result. The sequential path would catch that with an out-of-bounds index panic.

The caller (commit_phase_from_evaluations) always passes a correctly-sized inv_twiddles slice, so this won't trigger in practice. But a debug-mode assert would make the invariant explicit and catch misuse during development:

debug_assert_eq!(inv_twiddles.len(), evals.len() / 2);

.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];
Expand Down
2 changes: 1 addition & 1 deletion crypto/stark/src/fri/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<F: IsFFTField + IsSubFieldOf<E>, E: IsField>(
pub fn commit_phase_from_evaluations<F: IsFFTField + IsSubFieldOf<E> + Send + Sync, E: IsField + Send + Sync>(
number_layers: usize,
mut evals: Vec<FieldElement<E>>,
transcript: &mut impl IsStarkTranscript<E, F>,
Expand Down
Loading