Uniform table cap + Twiddle dedup + FRI fold parallelism

crypto/stark/src/fri/fri_functions.rs

@@ -14,18 +14,45 @@ use math::field::{
 /// in bit-reversed order, preserving conjugate pairing for the next fold.
 pub fn fold_evaluations_in_place<F: IsSubFieldOf<E>, E: IsField>(
     evals: &mut Vec<FieldElement<E>>,
+    #[cfg_attr(not(feature = "parallel"), allow(unused_variables))] scratch: &mut Vec<
+        FieldElement<E>,
+    >,
     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]).
+        // 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.
+        (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_into_vec(scratch);
+        std::mem::swap(evals, scratch);
+    }
+
+    #[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.

crypto/stark/src/fri/mod.rs

@@ -39,6 +39,8 @@ where
     let mut fri_layer_list = Vec::with_capacity(number_layers);
     let mut current_coset_offset = coset_offset.clone();
     let mut current_domain_size = domain_size;
+    // Scratch buffer reused across fold levels to avoid reallocating once per level.
+    let mut scratch: Vec<FieldElement<E>> = Vec::with_capacity(evals.len() / 2);
 
     for _ in 1..number_layers {
         // <<<< Receive challenge 𝜁ₖ₋₁
@@ -47,7 +49,7 @@ where
         current_domain_size /= 2;
 
         // Fold evaluations in-place (no FFT needed)
-        fold_evaluations_in_place(&mut evals, &zeta, &inv_twiddles);
+        fold_evaluations_in_place(&mut evals, &mut scratch, &zeta, &inv_twiddles);
 
         // Build Merkle tree from consecutive pairs
         let leaves: Vec<[FieldElement<E>; 2]> = evals
@@ -75,7 +77,7 @@ where
     let zeta = transcript.sample_field_element();
 
     // Final fold
-    fold_evaluations_in_place(&mut evals, &zeta, &inv_twiddles);
+    fold_evaluations_in_place(&mut evals, &mut scratch, &zeta, &inv_twiddles);
 
     let last_value = evals.first().unwrap_or(&FieldElement::zero()).clone();
 

crypto/stark/src/prover.rs

@@ -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;
@@ -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
@@ -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);
@@ -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)));
 
             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.
@@ -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(
@@ -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;
@@ -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();

crypto/stark/src/tests/fri_tests.rs

@@ -127,7 +127,8 @@ fn test_eval_fold_matches_coeff_fold() {
     let mut path_b_evals = evals_fft;
     in_place_bit_reverse_permute(&mut path_b_evals);
     let inv_twiddles = compute_coset_twiddles_inv::<GoldilocksField>(&coset_offset, n);
-    fold_evaluations_in_place(&mut path_b_evals, &beta, &inv_twiddles);
+    let mut scratch = Vec::new();
+    fold_evaluations_in_place(&mut path_b_evals, &mut scratch, &beta, &inv_twiddles);
 
     assert_eq!(path_a_evals, path_b_evals);
 }

prover/src/tables/mod.rs

@@ -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.
+///
+/// | 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)
 }
 
 /// Per-table maximum row limits, configurable for different environments.