Skip to content

perf(stark): eliminate redundant bit-reversal permutations in prover … #1187

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
diegokingston wants to merge 2 commits into
base: main
Choose a base branch
from
Open

perf(stark): eliminate redundant bit-reversal permutations in prover … #1187

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
90166de
perf(stark): eliminate redundant bit-reversal permutations in prover …
diegokingston Feb 25, 2026
c3c8e7d
Merge branch 'main' into opt/eliminate-redundant-bit-reversals
diegokingston Feb 25, 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
18 changes: 18 additions & 0 deletions crates/math/src/fft/cpu/ops.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,3 +24,21 @@ pub fn fft<F: IsFFTField + IsSubFieldOf<E>, E: IsField>(

Ok(results)
}

/// Like [`fft`], but returns evaluations in **bit-reversed order** (skips the
/// final bit-reverse permutation). This is useful when the output will be
/// committed to a Merkle tree in bit-reversed order, avoiding a redundant
/// permutation round-trip.
pub fn fft_bitrev<F: IsFFTField + IsSubFieldOf<E>, E: IsField>(
input: &[FieldElement<E>],
twiddles: &[FieldElement<F>],
) -> Result<alloc::vec::Vec<FieldElement<E>>, FFTError> {
if !input.len().is_power_of_two() {
return Err(FFTError::InputError(input.len()));
}

let mut results = input.to_vec();
in_place_nr_2radix_fft(&mut results, twiddles);

Ok(results)
}
Copy link
Copy Markdown

@github-actions github-actions bot Feb 25, 2026

Choose a reason for hiding this comment

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

Correctness

  • FFT Implementation: Ensure that in_place_nr_2radix_fft maintains the correctness properties of the FFT, especially the handling of edge cases like inputs of size one or power-of-two constraints. The call to is_power_of_two() should ensure this, but audit the helper function for correctness in recursion and base cases.

Security

  • Timing Side-Channels: Verify that in_place_nr_2radix_fft operates in constant time. If the algorithm's operations depend on secret data, ensure it does not lead to timing side-channel vulnerabilities.
  • Sensitive Data Zeroization: After operations, if any secret data persists in results, ensure it is zeroized to prevent unintended disclosure.

Performance

  • Unnecessary Allocations: Converting input to vector with input.to_vec() results in an allocation. Ensure this conversion is necessary, and consider using a different approach if input can be operated on in-place.
  • Redundant Calculations: Confirm that twiddles are precomputed efficiently, and all inverse calculations (like inverses and field operations in the FFT) are optimized.

Bugs & Errors

  • Off-by-One Errors: Since FFT involves array indexing, double-check accesses to results and twiddles for potential off-by-one errors.
  • Error Handling: While FFTError::InputError(input.len()) is returned for non-power-of-two sizes, ensure this is a comprehensive error management strategy for other potential issues within the FFT process.

Code Simplicity

  • Abstraction Quality: The new function has specific use cases, i.e., creating Merkle trees. Ensure it leverages existing abstractions efficiently and doesn't introduce unnecessary complexity. Create clear documentation for the circumstances that fft_bitrev should be used over fft.

In summary, verify correctness in edge cases, constant-time operation of FFT logic, and efficient handling of inputs without unnecessary allocations. Enhanced documentation to aid understanding of when fft_bitrev should be preferred is also recommended before merging.

48 changes: 48 additions & 0 deletions crates/math/src/fft/polynomial.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -161,6 +161,41 @@ impl<E: IsField> Polynomial<FieldElement<E>> {
Polynomial::evaluate_fft::<F>(&scaled, blowup_factor, domain_size)
}

/// Like [`Self::evaluate_fft`], but returns evaluations in **bit-reversed order**.
/// The NR DIT FFT naturally produces bit-reversed output; this variant skips
/// the final permutation to avoid a redundant round-trip when the output is
/// destined for a Merkle tree commitment.
pub fn evaluate_fft_bitrev<F: IsFFTField + IsSubFieldOf<E>>(
poly: &Self,
blowup_factor: usize,
domain_size: Option<usize>,
) -> Result<Vec<FieldElement<E>>, FFTError> {
let domain_size = domain_size.unwrap_or_default();
let len = core::cmp::max(poly.coeff_len(), domain_size).next_power_of_two() * blowup_factor;
if len.trailing_zeros() as u64 > F::TWO_ADICITY {
return Err(FFTError::DomainSizeError(len.trailing_zeros() as usize));
}
if poly.coefficients().is_empty() {
return Ok(vec![FieldElement::zero(); len]);
}

let mut coeffs = poly.coefficients().to_vec();
coeffs.resize(len, FieldElement::zero());

evaluate_fft_bitrev_cpu::<F, E>(&coeffs)
}

/// Like [`Self::evaluate_offset_fft`], but returns evaluations in **bit-reversed order**.
pub fn evaluate_offset_fft_bitrev<F: IsFFTField + IsSubFieldOf<E>>(
poly: &Self,
blowup_factor: usize,
domain_size: Option<usize>,
offset: &FieldElement<F>,
) -> Result<Vec<FieldElement<E>>, FFTError> {
let scaled = poly.scale(offset);
Self::evaluate_fft_bitrev::<F>(&scaled, blowup_factor, domain_size)
}

/// Returns a new polynomial that interpolates `(w^i, fft_evals[i])`, with `w` being a
/// Nth primitive root of unity in a subfield F of E, and `i in 0..N`, with `N = fft_evals.len()`.
/// This is considered to be the inverse operation of [Self::evaluate_fft()].
Expand Down Expand Up @@ -313,6 +348,19 @@ where
ops::fft(coeffs, &twiddles)
}

/// Like [`evaluate_fft_cpu`], but returns evaluations in bit-reversed order.
pub fn evaluate_fft_bitrev_cpu<F, E>(
coeffs: &[FieldElement<E>],
) -> Result<Vec<FieldElement<E>>, FFTError>
where
F: IsFFTField + IsSubFieldOf<E>,
E: IsField,
{
let order = coeffs.len().trailing_zeros();
let twiddles = roots_of_unity::get_twiddles::<F>(order.into(), RootsConfig::BitReverse)?;
ops::fft_bitrev(coeffs, &twiddles)
}

pub fn interpolate_fft_cpu<F, E>(
fft_evals: &[FieldElement<E>],
) -> Result<Polynomial<FieldElement<E>>, FFTError>
Copy link
Copy Markdown

@github-actions github-actions bot Feb 25, 2026

Choose a reason for hiding this comment

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

Correctness

  • Domain Size Check: In evaluate_fft_bitrev, the check if len.trailing_zeros() as u64 > F::TWO_ADICITY should probably use >= to ensure correctness when exactly equal to TWO_ADICITY.
  • Edge Case Handling: The function evaluate_fft_bitrev handles empty polynomials correctly by returning zeros. However, ensure this behavior is intended, specifically for downstream applications.

Security

  • Timing Side-channels: Ensure that operations involved in FFT (such as permutations and bit-reversals) are constant-time, especially if working over secret data.
  • Randomness: No new randomness operations are introduced in the changes, but ensure randomness used elsewhere in FFT or cryptographic commitments is secure.

Performance

  • Redundant Operations: The function evaluate_fft_bitrev duplicates some logic of evaluate_fft. Consider DRY (Don't Repeat Yourself) principle by refactoring common logic into a helper function.
  • Unnecessary Allocations: The coeffs vector is resized with zeros before evaluate_fft_bitrev_cpu is called, which might incur unnecessary allocation if coeffs is already of the needed length.

Bugs & Errors

  • Handling Errors: In evaluate_fft_bitrev_cpu, ensure that the construction of twiddles correctly handles potential errors from get_twiddles, which seems covered but should be verified.

Code Simplicity

  • Complexity: Overall, the functions added are simple, though there is some duplication with existing functionality. Consider refactoring to improve maintainability.

In summary, correctness around domain size check and constant-time assurance for security are the primary concerns. Some refactoring could resolve performance and DRY issues. Address these concerns before considering merging.

Expand Down
13 changes: 7 additions & 6 deletions crates/provers/stark/src/fri/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,11 +3,9 @@ pub mod fri_decommit;
mod fri_functions;

use lambdaworks_crypto::fiat_shamir::is_transcript::IsStarkTranscript;
use lambdaworks_math::field::traits::IsSubFieldOf;
use lambdaworks_math::field::traits::{IsFFTField, IsField};
use lambdaworks_math::traits::AsBytes;
use lambdaworks_math::{
fft::cpu::bit_reversing::in_place_bit_reverse_permute, field::traits::IsSubFieldOf,
};
pub use lambdaworks_math::{
field::{element::FieldElement, fields::u64_prime_field::U64PrimeField},
polynomial::Polynomial,
Expand Down Expand Up @@ -136,9 +134,12 @@ where
FieldElement<F>: AsBytes + Sync + Send,
FieldElement<E>: AsBytes + Sync + Send,
{
let mut evaluation = Polynomial::evaluate_offset_fft(poly, 1, Some(domain_size), coset_offset)?;

in_place_bit_reverse_permute(&mut evaluation);
// The NR DIT FFT naturally produces bit-reversed output. Use the bitrev
// variant directly to avoid the natural→bitrev round-trip that previously
// happened: evaluate_offset_fft (FFT+bitrev→natural) then
// in_place_bit_reverse_permute (natural→bitrev).
let evaluation =
Polynomial::evaluate_offset_fft_bitrev(poly, 1, Some(domain_size), coset_offset)?;

let mut to_commit = Vec::new();
for chunk in evaluation.chunks(2) {
Copy link
Copy Markdown

@github-actions github-actions bot Feb 25, 2026

Choose a reason for hiding this comment

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

Correctness

  • The change of using evaluate_offset_fft_bitrev directly instead of performing manual bit reversal after FFT suggests a more optimized path, avoiding an unnecessary round-trip. However, ensure that evaluate_offset_fft_bitrev is thoroughly tested, producing the same results in different contexts, particularly with edge cases like the zero polynomial and others.

Security

  • Ensure that operations over sensitive data, such as those involving secret keys or cryptographic randomness, are constant-time to prevent timing side-channel attacks. Double-check that no secret-dependent branching remains in these operations.
  • Evaluate the use of cryptographically secure randomness where applicable, especially in zero-knowledge proof systems.
  • If not already handled, sensitive data must be zeroized after use to prevent memory scraping attacks.

Performance

  • The change improves performance slightly by avoiding redundant permutation steps in FFT processing. The consolidation into a single call (evaluate_offset_fft_bitrev) is a beneficial optimization.

Bugs & Errors

  • No new Rust language-specific errors related to unwrap or potential panics introduced in this section of the code. Ensure that other parts using this code handle possible errors without panics.

Code Simplicity

  • The reduction of unnecessary bit-reversing simplifies the code, making the implementation more streamlined and easier to maintain. No other issues with code complexity observed in the specific section provided.

Overall

  • The changes here look mostly beneficial, but more comprehensive testing should confirm that similar optimizations apply across the board without edge case regressions. Ensure all mathematical operations remain valid, especially in zero-knowledge proof and polynomial operations.

Expand Down
122 changes: 90 additions & 32 deletions crates/provers/stark/src/prover.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@ use std::marker::PhantomData;
use std::time::Instant;

use lambdaworks_crypto::fiat_shamir::is_transcript::IsStarkTranscript;
use lambdaworks_math::fft::cpu::bit_reversing::{in_place_bit_reverse_permute, reverse_index};
use lambdaworks_math::fft::cpu::bit_reversing::reverse_index;
use lambdaworks_math::fft::errors::FFTError;

use lambdaworks_math::field::traits::{IsField, IsSubFieldOf};
Expand All @@ -23,7 +23,7 @@ use crate::domain::new_domain;
use crate::fri;
use crate::proof::stark::{DeepPolynomialOpenings, PolynomialOpenings};
use crate::table::Table;
use crate::trace::{columns2rows_bit_reversed, LDETraceTable};
use crate::trace::{columns2rows, LDETraceTable};

use super::config::{BatchedMerkleTree, Commitment};
use super::constraints::evaluator::ConstraintEvaluator;
Expand Down Expand Up @@ -212,6 +212,28 @@ where
}
}

/// Like [`evaluate_polynomial_on_lde_domain`], but returns evaluations in **bit-reversed order**.
/// Used for Merkle tree commitments where data must be bit-reversed, avoiding the
/// natural→bitrev round-trip.
fn evaluate_polynomial_on_lde_domain_bitrev<F, E>(
p: &Polynomial<FieldElement<E>>,
blowup_factor: usize,
domain_size: usize,
offset: &FieldElement<F>,
) -> Result<Vec<FieldElement<E>>, FFTError>
where
F: IsFFTField + IsSubFieldOf<E>,
E: IsField,
{
let evaluations =
Polynomial::evaluate_offset_fft_bitrev(p, blowup_factor, Some(domain_size), offset)?;
let step = evaluations.len() / (domain_size * blowup_factor);
match step {
1 => Ok(evaluations),
_ => Ok(evaluations.into_iter().step_by(step).collect()),
}
}

/// The functionality of a STARK prover providing methods to run the STARK Prove protocol
/// https://lambdaclass.github.io/lambdaworks/starks/protocol.html
/// The default implementation is complete and is compatible with Stone prover
Expand Down Expand Up @@ -278,12 +300,15 @@ pub trait IsStarkProver<
// Interpolate columns of `trace`.
let trace_polys = trace.compute_trace_polys_main::<Field>()?;

// Evaluate those polynomials t_j on the large domain D_LDE.
// Evaluate polynomials on LDE domain in natural order (for constraint evaluation).
let lde_trace_evaluations =
Self::compute_lde_trace_evaluations::<Field>(&trace_polys, domain)?;

// Compute commitment using fused bit-reverse + transpose (avoids cloning)
let lde_trace_permuted_rows = columns2rows_bit_reversed(&lde_trace_evaluations);
// Evaluate in bit-reversed order directly (for Merkle tree commitment).
// This avoids the natural→bitrev round-trip that columns2rows_bit_reversed did.
let lde_trace_bitrev =
Self::compute_lde_trace_evaluations_bitrev::<Field>(&trace_polys, domain)?;
let lde_trace_permuted_rows = columns2rows(lde_trace_bitrev);

let (lde_trace_merkle_tree, lde_trace_merkle_root) =
Self::batch_commit_main(&lde_trace_permuted_rows)
Expand Down Expand Up @@ -329,11 +354,12 @@ pub trait IsStarkProver<
// Interpolate columns of `trace`.
let trace_polys = trace.compute_trace_polys_aux::<Field>()?;

// Evaluate those polynomials t_j on the large domain D_LDE.
// Evaluate polynomials on LDE domain in natural order (for constraint evaluation).
let lde_trace_evaluations = Self::compute_lde_trace_evaluations(&trace_polys, domain)?;

// Compute commitment using fused bit-reverse + transpose (avoids cloning)
let lde_trace_permuted_rows = columns2rows_bit_reversed(&lde_trace_evaluations);
// Evaluate in bit-reversed order directly (for Merkle tree commitment).
let lde_trace_bitrev = Self::compute_lde_trace_evaluations_bitrev(&trace_polys, domain)?;
let lde_trace_permuted_rows = columns2rows(lde_trace_bitrev);

let (lde_trace_merkle_tree, lde_trace_merkle_root) =
Self::batch_commit_extension(&lde_trace_permuted_rows)
Expand Down Expand Up @@ -377,6 +403,33 @@ pub trait IsStarkProver<
.collect::<Result<Vec<Vec<FieldElement<E>>>, FFTError>>()
}

/// Like [`Self::compute_lde_trace_evaluations`], but returns columns in bit-reversed order.
/// Used for Merkle tree commitment where the leaves must be in bit-reversed order.
fn compute_lde_trace_evaluations_bitrev<E>(
trace_polys: &[Polynomial<FieldElement<E>>],
domain: &Domain<Field>,
) -> Result<Vec<Vec<FieldElement<E>>>, FFTError>
where
E: IsSubFieldOf<FieldExtension>,
Field: IsSubFieldOf<E>,
{
#[cfg(not(feature = "parallel"))]
let trace_polys_iter = trace_polys.iter();
#[cfg(feature = "parallel")]
let trace_polys_iter = trace_polys.par_iter();

trace_polys_iter
.map(|poly| {
evaluate_polynomial_on_lde_domain_bitrev(
poly,
domain.blowup_factor,
domain.interpolation_domain_size,
&domain.coset_offset,
)
})
.collect::<Result<Vec<Vec<FieldElement<E>>>, FFTError>>()
}

/// Returns the result of the first round of the STARK Prove protocol.
fn round_1_randomized_air_with_preprocessing(
air: &dyn AIR<Field = Field, FieldExtension = FieldExtension, PublicInputs = PI>,
Expand Down Expand Up @@ -429,54 +482,45 @@ pub trait IsStarkProver<
})
}

/// Returns the Merkle tree and the commitment to the evaluations of the parts of the
/// composition polynomial.
/// Commits to composition polynomial part evaluations that are already in bit-reversed order.
/// Transposes columns to rows, pairs consecutive rows, and builds the Merkle tree.
///
/// Returns `None` if:
/// - `lde_composition_poly_parts_evaluations` is empty
/// - Any part has different length than others
/// - The LDE length is not even (required for pairing evaluations)
fn commit_composition_polynomial(
lde_composition_poly_parts_evaluations: &[Vec<FieldElement<FieldExtension>>],
/// Returns `None` if the input is empty, has zero-length parts, odd-length parts, or
/// parts of different lengths.
fn commit_composition_polynomial_bitrev(
lde_composition_poly_parts_bitrev: &[Vec<FieldElement<FieldExtension>>],
) -> Option<(BatchedMerkleTree<FieldExtension>, Commitment)>
where
FieldElement<Field>: AsBytes + Sync + Send,
FieldElement<FieldExtension>: AsBytes + Sync + Send,
{
// Validate input: must have at least one part with at least one evaluation
let first_part = lde_composition_poly_parts_evaluations.first()?;
let first_part = lde_composition_poly_parts_bitrev.first()?;
let lde_len = first_part.len();

if lde_len == 0 {
return None;
}

// LDE length must be even for the chunking below to work correctly
if lde_len % 2 != 0 {
if lde_len == 0 || lde_len % 2 != 0 {
return None;
}

// All parts must have the same length
if !lde_composition_poly_parts_evaluations
if !lde_composition_poly_parts_bitrev
.iter()
.all(|part| part.len() == lde_len)
{
return None;
}

let num_parts = lde_composition_poly_parts_evaluations.len();
let mut lde_composition_poly_evaluations: Vec<Vec<_>> = (0..lde_len)
// Transpose columns to rows — data is already bit-reversed, no permutation needed.
let num_parts = lde_composition_poly_parts_bitrev.len();
let lde_composition_poly_evaluations: Vec<Vec<_>> = (0..lde_len)
.map(|i| {
let mut row = Vec::with_capacity(num_parts);
for evaluation in lde_composition_poly_parts_evaluations.iter() {
for evaluation in lde_composition_poly_parts_bitrev.iter() {
row.push(evaluation[i].clone());
}
row
})
.collect();

in_place_bit_reverse_permute(&mut lde_composition_poly_evaluations);

// Pair consecutive rows for Merkle leaves.
let mut lde_composition_poly_evaluations_merged = Vec::with_capacity(lde_len / 2);
let mut iter = lde_composition_poly_evaluations.into_iter();
while let (Some(mut chunk0), Some(chunk1)) = (iter.next(), iter.next()) {
Expand Down Expand Up @@ -526,6 +570,7 @@ pub trait IsStarkProver<
}
let composition_poly_parts = composition_poly.break_in_parts(number_of_parts);

// Natural-order evaluations (used for opening queries via reverse_index).
let lde_composition_poly_parts_evaluations: Vec<_> = composition_poly_parts
.iter()
.map(|part| {
Expand All @@ -538,8 +583,21 @@ pub trait IsStarkProver<
})
.collect::<Result<Vec<_>, _>>()?;

// Bit-reversed evaluations (for Merkle tree commitment — avoids natural→bitrev round-trip).
let lde_composition_poly_parts_bitrev: Vec<_> = composition_poly_parts
.iter()
.map(|part| {
evaluate_polynomial_on_lde_domain_bitrev(
part,
domain.blowup_factor,
domain.interpolation_domain_size,
&domain.coset_offset,
)
})
.collect::<Result<Vec<_>, _>>()?;

let Some((composition_poly_merkle_tree, composition_poly_root)) =
Self::commit_composition_polynomial(&lde_composition_poly_parts_evaluations)
Self::commit_composition_polynomial_bitrev(&lde_composition_poly_parts_bitrev)
else {
return Err(ProvingError::EmptyCommitment);
};
Copy link
Copy Markdown

@github-actions github-actions bot Feb 25, 2026

Choose a reason for hiding this comment

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

Correctness:

  1. Polynomial Operations: Ensure that evaluate_offset_fft_bitrev and related polynomial evaluation functions are correctly implemented and verified for correctness.

  2. Domain-Size Edge Cases: The method evaluate_polynomial_on_lde_domain_bitrev uses a division by domain_size * blowup_factor. Make sure that these values are never zero to avoid division by zero errors.

  3. Transposition Logic: In commit_composition_polynomial_bitrev, ensure that transposition logic correctly handles cases for different domain sizes and polynomial lengths, confirming that no data is inadvertently lost or incorrectly rearranged.

  4. Handling of Bit-Reversed Data: The function commit_composition_polynomial_bitrev assumes data is already in bit-reversed format. Double-check that all inputs meet this assumption to avoid logic errors in FFT operations.

Security:

  1. Constant-Time Operation: Review all FFT and modular arithmetic operations for constant-time implementations to mitigate timing attacks, especially those near secret values.

  2. Zeroization: Ensure that sensitive data, particularly the evaluations and results related to cryptographic operations, are zeroized after use.

  3. Randomness: Verify that cryptographically secure randomness is used in all areas requiring randomness, such as field element generation.

Bugs & Errors:

  1. Potential Panics: Ensure code involving indexing and slicing safely handles all possible ranges and does not lead to panics.

  2. Error Handling: The new functions returning Result types should properly handle and propagate errors instead of potentially ignoring them.

  3. Integer Overflow/Underflow: Rust's built-in checks mitigate this risk, but manual verification can confirm there are no logical issues leading to unintended overflows.

Performance:

  1. Redundant Computations: Check if there is any unnecessary duplication of polynomial evaluation where bit-reversing could be unified to reduce computational overhead (e.g., avoid double evaluations when mapping directly to desired order output).

  2. Memory Usage: Look for optimization opportunities in evaluate_polynomial_on_lde_domain_bitrev where vector manipulations and storage could be streamlined to reduce allocations.

Code Simplicity:

  1. Abstraction Consistency: The introduction of bit-reversed specific functions increases complexity slightly. Consider whether these can be integrated more harmoniously with existing functions to minimize the expansion of the API surface.

Overall, the proposed changes introduce modifications in FFT and polynomial evaluations while aiming to optimize bit-reversal requirements. These need careful verification in practice to ensure no edge cases or unintended errors arise during both typical and boundary-case usage.

Expand Down
Loading