Skip to content

Uniform table cap + Twiddle dedup + FRI fold parallelism #497

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 5 commits into from
Closed

Uniform table cap + Twiddle dedup + FRI fold parallelism #497

Show file tree
Hide file tree
Changes from 4 commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
0571e45
perf(tables): cap LT/LOAD/BRANCH/MEMW_R max_rows at 2^20
diegokingston Apr 9, 2026
1e217ad
perf(stark): deduplicate LdeTwiddles by domain size in multi_prove
diegokingston Apr 9, 2026
1674ed1
perf(fri): parallelize fold_evaluations_in_place with Rayon
diegokingston Apr 9, 2026
cf1c8d9
perf(prover): eliminate per-row Vec alloc in commit_columns_bit_reversed
diegokingston Apr 9, 2026
eeb860c
Reuse FRI fold scratch buffer across levels
gabrielbosio Apr 15, 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: 31 additions & 7 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 @@ -18,14 +18,38 @@ pub fn fold_evaluations_in_place<F: IsSubFieldOf<E>, E: IsField>(
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]).
// Index-based iteration matches the sequential path and panics on a
// short `inv_twiddles`; `par_chunks(2).zip(inv_twiddles)` would drop
// a trailing odd element and silently truncate on length mismatch.
let folded: Vec<FieldElement<E>> = (0..half)
.into_par_iter()
.map(|j| {
let lo = &evals[2 * j];
let hi = &evals[2 * j + 1];
let sum = lo + hi;
let diff = lo - hi;
&sum + &(&inv_twiddles[j] * &(zeta * &diff))
})
.collect();
*evals = folded;

@claude claude Bot Apr 14, 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 path always allocates a fresh Vec per fold level

*evals = folded drops the old allocation and replaces it. FRI folding runs ~log₂(N) times on the hot path, so this adds ~log₂(N) allocations compared to the sequential path.

One option is to reuse the buffer:

Suggested change
*evals = folded;
evals.clear();
evals.extend(folded);

Or pass in a mutable scratch buffer so the caller can pre-allocate once and reuse across levels. The current approach is correct but leaves a performance opportunity on the table given this is the inner loop of FRI.

}

#[cfg(not(feature = "parallel"))]
{
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
80 changes: 60 additions & 20 deletions crypto/stark/src/prover.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<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 @@ -635,7 +651,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 @@ -651,7 +667,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 @@ -1512,24 +1534,42 @@ 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;
let mut max_lde_size = 0usize;

// Deduplicate twiddle caches: tables with identical domain parameters share one Arc.
// Key on (interpolation_domain_size, blowup_factor, coset_offset_bytes) because each
// affects the twiddle values: interpolation_domain_size sizes `inv` and `coset_weights`,
// blowup_factor sizes `fwd`, and coset_offset sets `coset_weights` values.
// coset_offset is serialized via AsBytes since FieldElement<F> does not impl Hash without
// a BaseType: Hash bound.
let mut twiddle_by_domain: std::collections::HashMap<
(usize, usize, Vec<u8>),
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 lde_size = domain.interpolation_domain_size * domain.blowup_factor;
let twiddles = LdeTwiddles::new(&domain);
let key = (
domain.interpolation_domain_size,
domain.blowup_factor,
domain.coset_offset.as_bytes(),
);
let twiddles = twiddle_by_domain
.entry(key)
.or_insert_with(|| Arc::new(LdeTwiddles::new(&domain)));
Comment thread
gabrielbosio marked this conversation as resolved.

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));
}

// Allocate K independent LDE column buffer pool sets for parallel table processing.
Expand Down Expand Up @@ -1573,7 +1613,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.is_preprocessed() {
Self::commit_preprocessed_trace(
Expand Down Expand Up @@ -1693,7 +1733,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 Expand Up @@ -1809,7 +1849,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();
Expand Down
37 changes: 21 additions & 16 deletions prover/src/tables/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,29 +49,34 @@ pub use types::BusId;
/// (* MEMW_A formula gives 2^20, but set to 2^19 to match MEMW chunk geometry;
/// benchmarks show better parallel throughput with smaller chunks.)
///
/// | Table | Main | Bus | Eff.width | Max rows |
/// |---------|------|-----|-----------|----------|
/// | MEMW | 49 | 26 | 127 | 2^19 |
/// | MEMW_A | 29 | 20 | 89 | 2^19 * |
/// | CPU | 74 | 40 | 194 | 2^19 |
/// | DVRM | 34 | 34 | 136 | 2^19 |
/// | MUL | 26 | 16 | 74 | 2^20 |
/// | LT | 15 | 9 | 42 | 2^21 |
/// | 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 |
/// Capped tables use 2^20, not the 2^21 the formula gives. A 2^21 chunk
/// runs on a single thread and peaks at twice the memory of a 2^20 chunk.
/// Programs that exceed 2^20 rows in a capped table hit a table-overflow
/// error.
///
Comment thread
gabrielbosio marked this conversation as resolved.
/// | Table | Main | Bus | Eff.width | Max rows |
/// |---------|------|-----|-----------|-----------------|
/// | MEMW | 49 | 26 | 127 | 2^19 |
/// | MEMW_A | 29 | 20 | 89 | 2^19 * |
/// | CPU | 74 | 40 | 194 | 2^19 |
/// | DVRM | 34 | 34 | 136 | 2^19 |
/// | MUL | 26 | 16 | 74 | 2^20 |
/// | LT | 15 | 9 | 42 | 2^20 (capped) |
/// | SHIFT | 27 | 15 | 72 | 2^20 |
/// | LOAD | 18 | 5 | 33 | 2^20 (capped) |
/// | BRANCH | 14 | 6 | 32 | 2^20 (capped) |
/// | MEMW_R | 10 | 7 | 31 | 2^20 (capped) |
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 (capped at 2^20)
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 (capped at 2^20)
pub const BRANCH: usize = 1 << 20; // 1,048,576 — eff. width 32 (capped at 2^20)
pub const MEMW_R: usize = 1 << 20; // 1,048,576 — eff. width 31 (capped at 2^20)
Comment thread
gabrielbosio marked this conversation as resolved.
}

/// Per-table maximum row limits, configurable for different environments.
Expand Down
Loading