Skip to content

Cache LDE + Scaling improvements #494

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

Closed
gabrielbosio wants to merge 7 commits into from
Closed

Cache LDE + Scaling improvements #494

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
0126881
perf(tables): cap LT/LOAD/BRANCH/MEMW_R max_rows at 2^20
diegokingston Apr 9, 2026
ada4503
perf(stark): deduplicate LdeTwiddles by domain size in multi_prove
diegokingston Apr 9, 2026
398e18f
perf(fri): parallelize fold_evaluations_in_place with Rayon
diegokingston Apr 9, 2026
d8e432f
perf(prover): eliminate per-row Vec alloc in commit_columns_bit_reversed
diegokingston Apr 9, 2026
24b44e4
Fix fmt and clippy after cherry-pick
gabrielbosio Apr 9, 2026
4ce7a24
Remove redundant capped annotations from max_rows
gabrielbosio Apr 9, 2026
ec49bf6
Document twiddle dedup invariant
gabrielbosio Apr 10, 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
38 changes: 30 additions & 8 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 @@ -17,15 +17,37 @@ pub fn fold_evaluations_in_place<F: IsSubFieldOf<E>, E: IsField>(
zeta: &FieldElement<E>,
inv_twiddles: &[FieldElement<F>],
) {
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<FieldElement<E>> = evals
.par_chunks(2)
.zip(inv_twiddles.par_iter())
Comment on lines +25 to +27

@claude claude Bot Apr 13, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Low – Parallel and sequential branches diverge on contract violation

If inv_twiddles.len() < evals.len() / 2 (broken contract):

  • parallel: zip silently truncates to the shorter iterator → folded is shorter than evals.len() / 2, and *evals = folded quietly installs a wrong-length vector. Downstream code proceeds with corrupted state.
  • sequential: evals[j] with j reaching half - 1 will index out of inv_twiddles → panic.

In practice the contract is maintained by callers, but if it breaks the parallel path gives no diagnostic while the sequential path at least terminates loudly. A debug_assert_eq!(inv_twiddles.len(), evals.len() / 2) at the top of the function would catch this in both builds without runtime cost in release.

.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.
Expand Down
69 changes: 50 additions & 19 deletions crypto/stark/src/prover.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -341,19 +341,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<Commitment> = {
(0..num_rows)
.into_par_iter()
.map_init(
|| vec![FieldElement::<E>::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::<E>::hash_data(row_buf)
},
)
.collect()
};

let hashed_leaves: Vec<Commitment> = iter
.map(|row_idx| {
let br_idx = reverse_index(row_idx, num_rows as u64);
let row: Vec<FieldElement<E>> = (0..num_cols)
.map(|col_idx| columns[col_idx][br_idx].clone())
.collect();
BatchedMerkleTreeBackend::<E>::hash_data(&row)
})
.collect();
#[cfg(not(feature = "parallel"))]
let hashed_leaves: Vec<Commitment> = {
let mut row_buf = vec![FieldElement::<E>::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::<E>::hash_data(&row_buf)
})
.collect()
};

let tree = BatchedMerkleTree::<E>::build_from_hashed_leaves(hashed_leaves)?;
let root = tree.root;
Expand Down Expand Up @@ -641,7 +657,7 @@ pub trait IsStarkProver<
air_trace_pairs: &[AirTracePair<'_, Field, FieldExtension, PI>],
metadatas: &[Round1Metadata<Field, FieldExtension>],
domains: &[Domain<Field>],
twiddle_caches: &[LdeTwiddles<Field>],
twiddle_caches: &[Arc<LdeTwiddles<Field>>],
main_pool: &mut [Vec<FieldElement<Field>>],
aux_pool: &mut [Vec<FieldElement<FieldExtension>>],
) where
Expand All @@ -657,7 +673,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);
Expand Down Expand Up @@ -1518,20 +1540,29 @@ pub trait IsStarkProver<
let phase_start = Instant::now();

let mut domains = Vec::with_capacity(num_airs);
let mut twiddle_caches: Vec<LdeTwiddles<Field>> = Vec::with_capacity(num_airs);
let mut twiddle_caches: Vec<Arc<LdeTwiddles<Field>>> = Vec::with_capacity(num_airs);
let mut max_main_cols = 0usize;
let mut max_aux_cols = 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<usize, Arc<LdeTwiddles<Field>>> =
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 twiddles = LdeTwiddles::new(&domain);

let lde_size = domain.interpolation_domain_size * domain.blowup_factor;
let twiddles = twiddle_by_size
.entry(lde_size)
.or_insert_with(|| Arc::new(LdeTwiddles::new(&domain)));
Comment on lines +1556 to +1559

@claude claude Bot Apr 13, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Medium – Incorrect cache key if AIRs have different blowup_factor or coset_offset

LdeTwiddles::new bakes three domain-specific quantities into the struct:

  • inv twiddles — sized by interpolation_domain_size
  • fwd twiddles — sized by lde_size
  • coset_weights — sized by interpolation_domain_size, scaled by coset_offset

Using only lde_size as the cache key is safe only when all AIRs share the same blowup_factor (so equal lde_size ⟹ equal interpolation_domain_size) and the same coset_offset. The comment acknowledges this but it isn't enforced. If violated, twiddles from one domain are silently reused for another, producing cryptographically wrong (but non-panicking) proofs.

Consider asserting the invariant at the cache hit, e.g.:

// or store (interpolation_domain_size, coset_offset) as key
debug_assert_eq!(
    domain.interpolation_domain_size,
    first_domain_with_this_lde_size.interpolation_domain_size
);

Or use (lde_size, coset_offset_bits) as the key to make the assumption self-enforcing.


max_main_cols = max_main_cols.max(trace.num_main_columns);
max_aux_cols = max_aux_cols.max(air.num_auxiliary_rap_columns());

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.
Expand Down Expand Up @@ -1573,7 +1604,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];

let result = if air.is_preprocessed() {
Self::commit_preprocessed_trace(
Expand Down Expand Up @@ -1705,7 +1736,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;
Expand Down
16 changes: 8 additions & 8 deletions prover/src/tables/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
Loading