Disk spill with scaling improvements

crypto/stark/src/fri/fri_functions.rs

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

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;
@@ -660,7 +676,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
@@ -676,7 +692,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);
@@ -1534,24 +1556,31 @@ 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;
 
+        // 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 lde_size = domain.interpolation_domain_size * domain.blowup_factor;
-            let twiddles = LdeTwiddles::new(&domain);
+            let twiddles = twiddle_by_size
+                .entry(lde_size)
+                .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));
         }
 
         // Spill all main trace tables to mmap before allocating pool buffers.
@@ -1631,7 +1660,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];
 
                     #[allow(unused_mut)]
                     let (mut tree, root, pre_tree, pre_root, n_pre) = if air.is_preprocessed() {
@@ -1828,7 +1857,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;
@@ -2080,7 +2109,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();

prover/src/tables/mod.rs

@@ -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.