improve batching sumcheck concurrency to dominate by max num var polynomial (#895)

Previously (devirgo) sumcheck concurrency when batching with different
number of variables, are dominated by the smallest num_vars. The reason
is because for the num_var < log_2(threads) we can NOT divide it's
evaluation into #thread region.

This PR addressed the issue by introducing a extra meta information for
each polynomial, so we are able to differentiate those small poly and
calculate its multiplicity correctly

### Design rationale
For some extreme small polynomial (num_var << log(threads)), we define a
new `PolyType::Phase2Only` to differentiate them. These type of small
polynomial will be handle by main_worker only during phase 1 sumcheck.
And when moving to phase2 sumcheck, those small poly will expand the
size to match with all. This cost is negligible giving phase2 only need
to run on #thread evaluations.

### benchmark 
There is no different with before/after change, since we dont have batch
sumcheck with different num_vars in critical path, which meet expected

e2e Fibonacci 2^20 
```
fibonacci_max_steps_1048576/prove_fibonacci/fibonacci_max_steps_1048576
                        time:   [2.8894 s 2.9028 s 2.9180 s]
                        change: [+0.1414% +0.8059% +1.5474%] (p = 0.05 < 0.05)
                        Change within noise threshold.
```

e2e Fibonacci 2^21
```
fibonacci_max_steps_2097152/prove_fibonacci/fibonacci_max_steps_2097152
                        time:   [5.1785 s 5.2061 s 5.2291 s]
                        change: [+0.7256% +1.3249% +1.9085%] (p = 0.00 < 0.05)
                        Change within noise threshold.
```

e2e Fibonacci 2^22
```
fibonacci_max_steps_4194304/prove_fibonacci/fibonacci_max_steps_4194304
                        time:   [10.308 s 10.330 s 10.353 s]
                        change: [+0.9437% +1.2216% +1.4899%] (p = 0.00 < 0.05)
                        Change within noise threshold.
```

Author: hero78119

Diff for: ceno_zkvm/src/scheme/prover.rs

@@ -572,11 +572,7 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> ZKVMProver<E, PCS> {
         }
 
         tracing::debug!("main sel sumcheck start");
-        let (main_sel_sumcheck_proofs, state) = IOPProverState::prove_batch_polys(
-            num_threads,
-            virtual_polys.get_batched_polys(),
-            transcript,
-        );
+        let (main_sel_sumcheck_proofs, state) = IOPProverState::prove(virtual_polys, transcript);
         tracing::debug!("main sel sumcheck end");
 
         let main_sel_evals = state.get_mle_final_evaluations();
@@ -1015,11 +1011,8 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> ZKVMProver<E, PCS> {
                     virtual_polys.add_mle_list(vec![eq, lk_d_wit], *alpha);
                 }
 
-                let (same_r_sumcheck_proofs, state) = IOPProverState::prove_batch_polys(
-                    num_threads,
-                    virtual_polys.get_batched_polys(),
-                    transcript,
-                );
+                let (same_r_sumcheck_proofs, state) =
+                    IOPProverState::prove(virtual_polys, transcript);
                 let evals = state.get_mle_final_evaluations();
                 let mut evals_iter = evals.into_iter();
                 let rw_in_evals = cs
@@ -1271,9 +1264,8 @@ impl TowerProver {
                 // NOTE: at the time of adding this span, visualizing it with the flamegraph layer
                 // shows it to be (inexplicably) much more time-consuming than the call to `prove_batch_polys`
                 // This is likely a bug in the tracing-flame crate.
-                let (sumcheck_proofs, state) = IOPProverState::prove_batch_polys(
-                    num_threads,
-                    virtual_polys.get_batched_polys(),
+                let (sumcheck_proofs, state) = IOPProverState::prove(
+                    virtual_polys,
                     transcript,
                 );
                 exit_span!(wrap_batch_span);

Diff for: mpcs/src/basefold/commit_phase.rs

@@ -21,13 +21,14 @@ use serde::{Serialize, de::DeserializeOwned};
 use sumcheck::{
     macros::{entered_span, exit_span},
     structs::IOPProverState,
-    util::{AdditiveVec, merge_sumcheck_polys, optimal_sumcheck_threads},
+    util::{AdditiveVec, merge_sumcheck_prover_state, optimal_sumcheck_threads},
 };
 use transcript::{Challenge, Transcript};
 
 use multilinear_extensions::{
     commutative_op_mle_pair,
     mle::{DenseMultilinearExtension, IntoMLE},
+    util::ceil_log2,
     virtual_poly::{ArcMultilinearExtension, build_eq_x_r_vec},
     virtual_polys::VirtualPolynomials,
 };
@@ -98,15 +99,23 @@ where
     end_timer!(build_eq_timer);
 
     let num_threads = optimal_sumcheck_threads(num_vars);
+    let log_num_threads = ceil_log2(num_threads);
 
     let mut polys = VirtualPolynomials::new(num_threads, num_vars);
     polys.add_mle_list(vec![&eq, &running_evals], E::ONE);
-    let batched_polys = polys.get_batched_polys();
+    let (batched_polys, poly_meta) = polys.get_batched_polys();
 
     let mut prover_states = batched_polys
         .into_iter()
-        .map(|poly| {
-            IOPProverState::prover_init_with_extrapolation_aux(poly, vec![(vec![], vec![])])
+        .enumerate()
+        .map(|(thread_id, poly)| {
+            IOPProverState::prover_init_with_extrapolation_aux(
+                thread_id == 0, // set thread_id 0 to be main worker
+                poly,
+                vec![(vec![], vec![])],
+                Some(log_num_threads),
+                Some(poly_meta.clone()),
+            )
         })
         .collect::<Vec<_>>();
 
@@ -140,13 +149,16 @@ where
     }
 
     // deal with log(#thread) basefold rounds
-    let merge_sumcheck_polys_span = entered_span!("merge_sumcheck_polys");
-    let poly = merge_sumcheck_polys(&prover_states);
+    let merge_sumcheck_prover_state_span = entered_span!("merge_sumcheck_prover_state");
+    let poly = merge_sumcheck_prover_state(prover_states);
     let mut prover_states = vec![IOPProverState::prover_init_with_extrapolation_aux(
+        true,
         poly,
         vec![(vec![], vec![])],
+        None,
+        None,
     )];
-    exit_span!(merge_sumcheck_polys_span);
+    exit_span!(merge_sumcheck_prover_state_span);
 
     let mut challenge = None;
 

Diff for: multilinear_extensions/src/virtual_poly.rs

@@ -18,6 +18,7 @@ use serde::{Deserialize, Serialize};
 
 pub type ArcMultilinearExtension<'a, E> =
     Arc<dyn MultilinearExtension<E, Output = DenseMultilinearExtension<E>> + 'a>;
+
 #[rustfmt::skip]
 /// A virtual polynomial is a sum of products of multilinear polynomials;
 /// where the multilinear polynomials are stored via their multilinear
@@ -113,21 +114,17 @@ impl<'a, E: ExtensionField> VirtualPolynomial<'a, E> {
     ///
     /// The MLEs will be multiplied together, and then multiplied by the scalar
     /// `coefficient`.
-    pub fn add_mle_list(&mut self, mle_list: Vec<ArcMultilinearExtension<'a, E>>, coefficient: E) {
+    pub fn add_mle_list(
+        &mut self,
+        mle_list: Vec<ArcMultilinearExtension<'a, E>>,
+        coefficient: E,
+    ) -> &[usize] {
         let mle_list: Vec<ArcMultilinearExtension<E>> = mle_list.into_iter().collect();
         let mut indexed_product = Vec::with_capacity(mle_list.len());
 
         assert!(!mle_list.is_empty(), "input mle_list is empty");
         // sanity check: all mle in mle_list must have same num_vars()
-        assert!(
-            mle_list
-                .iter()
-                .map(|m| {
-                    assert!(m.num_vars() <= self.aux_info.max_num_variables);
-                    m.num_vars()
-                })
-                .all_equal()
-        );
+        assert!(mle_list.iter().map(|m| { m.num_vars() }).all_equal());
 
         self.aux_info.max_degree = max(self.aux_info.max_degree, mle_list.len());
 
@@ -143,6 +140,7 @@ impl<'a, E: ExtensionField> VirtualPolynomial<'a, E> {
             }
         }
         self.products.push((coefficient, indexed_product));
+        &self.products.last().unwrap().1
     }
 
     /// in-place merge with another virtual polynomial

Diff for: multilinear_extensions/src/virtual_polys.rs

@@ -1,19 +1,31 @@
-use std::{collections::HashMap, sync::Arc};
+use std::{
+    collections::{BTreeMap, HashMap},
+    sync::Arc,
+};
 
 use crate::{
     util::ceil_log2,
     virtual_poly::{ArcMultilinearExtension, VirtualPolynomial},
 };
 use ff_ext::ExtensionField;
 use itertools::Itertools;
+use p3::util::log2_strict_usize;
 
 use crate::util::transpose;
 
+#[derive(Debug, Default, Clone, Copy)]
+pub enum PolyMeta {
+    #[default]
+    Normal,
+    Phase2Only,
+}
+
 pub struct VirtualPolynomials<'a, E: ExtensionField> {
-    num_threads: usize,
+    pub num_threads: usize,
     polys: Vec<VirtualPolynomial<'a, E>>,
     /// a storage to keep thread based mles, specific to multi-thread logic
     thread_based_mles_storage: HashMap<usize, Vec<ArcMultilinearExtension<'a, E>>>,
+    pub(crate) poly_meta: BTreeMap<usize, PolyMeta>,
 }
 
 impl<'a, E: ExtensionField> VirtualPolynomials<'a, E> {
@@ -25,6 +37,7 @@ impl<'a, E: ExtensionField> VirtualPolynomials<'a, E> {
                 .map(|_| VirtualPolynomial::new(max_num_variables - ceil_log2(num_threads)))
                 .collect_vec(),
             thread_based_mles_storage: HashMap::new(),
+            poly_meta: BTreeMap::new(),
         }
     }
 
@@ -44,32 +57,52 @@ impl<'a, E: ExtensionField> VirtualPolynomials<'a, E> {
     }
 
     pub fn add_mle_list(&mut self, polys: Vec<&'a ArcMultilinearExtension<'a, E>>, coeff: E) {
-        let polys = polys
+        let log2_num_threads = log2_strict_usize(self.num_threads);
+        let (poly_meta, polys): (Vec<PolyMeta>, Vec<Vec<ArcMultilinearExtension<E>>>) = polys
             .into_iter()
             .map(|p| {
                 let mle_ptr: usize = Arc::as_ptr(p) as *const () as usize;
-                if let Some(mles) = self.thread_based_mles_storage.get(&mle_ptr) {
+                let poly_meta = if p.num_vars() > log2_num_threads {
+                    PolyMeta::Normal
+                } else {
+                    // polynomial is too small
+                    PolyMeta::Phase2Only
+                };
+                let mles_cloned = if let Some(mles) = self.thread_based_mles_storage.get(&mle_ptr) {
                     mles.clone()
                 } else {
                     let mles = (0..self.num_threads)
-                        .map(|thread_id| {
-                            self.get_range_polys_by_thread_id(thread_id, vec![p])
-                                .remove(0)
+                        .map(|thread_id| match poly_meta {
+                            PolyMeta::Normal => self
+                                .get_range_polys_by_thread_id(thread_id, vec![p])
+                                .remove(0),
+                            PolyMeta::Phase2Only => Arc::new(p.get_ranged_mle(1, 0)),
                         })
                         .collect_vec();
                     let mles_cloned = mles.clone();
                     self.thread_based_mles_storage.insert(mle_ptr, mles);
                     mles_cloned
-                }
+                };
+                (poly_meta, mles_cloned)
             })
-            .collect_vec();
+            .unzip();
 
         // poly -> thread to thread -> poly
         let polys = transpose(polys);
-        (0..self.num_threads)
+        let poly_index: &[usize] = self
+            .polys
+            .iter_mut()
             .zip_eq(polys)
-            .for_each(|(thread_id, polys)| {
-                self.polys[thread_id].add_mle_list(polys, coeff);
+            .map(|(poly, polys)| poly.add_mle_list(polys, coeff))
+            .collect_vec()
+            .first()
+            .expect("expect to get at index from first thread");
+
+        poly_index
+            .iter()
+            .zip_eq(&poly_meta)
+            .for_each(|(index, poly_meta)| {
+                self.poly_meta.insert(*index, *poly_meta);
             });
     }
 
@@ -84,8 +117,13 @@ impl<'a, E: ExtensionField> VirtualPolynomials<'a, E> {
         }
     }
 
-    pub fn get_batched_polys(self) -> Vec<VirtualPolynomial<'a, E>> {
-        self.polys
+    /// return thread_based polynomial with its polynomial type
+    pub fn get_batched_polys(self) -> (Vec<VirtualPolynomial<'a, E>>, Vec<PolyMeta>) {
+        let mut poly_meta = vec![PolyMeta::Normal; self.polys[0].flattened_ml_extensions.len()];
+        for (index, poly_meta_by_index) in self.poly_meta {
+            poly_meta[index] = poly_meta_by_index
+        }
+        (self.polys, poly_meta)
     }
 
     pub fn degree(&self) -> usize {

Diff for: sumcheck/benches/devirgo_sumcheck.rs

@@ -1,19 +1,21 @@
 #![allow(clippy::manual_memcpy)]
 #![allow(clippy::needless_range_loop)]
 
-use std::{array, time::Duration};
+use std::time::Duration;
 
 use ark_std::test_rng;
 use criterion::*;
 use ff_ext::{ExtensionField, GoldilocksExt2};
 use itertools::Itertools;
-use sumcheck::{structs::IOPProverState, util::ceil_log2};
+use p3::field::PrimeCharacteristicRing;
+use sumcheck::structs::IOPProverState;
 
 use multilinear_extensions::{
     mle::DenseMultilinearExtension,
     op_mle,
     util::max_usable_threads,
     virtual_poly::{ArcMultilinearExtension, VirtualPolynomial},
+    virtual_polys::VirtualPolynomials,
 };
 use transcript::BasicTranscript as Transcript;
 
@@ -39,49 +41,15 @@ pub fn transpose<T>(v: Vec<Vec<T>>) -> Vec<Vec<T>> {
         .collect()
 }
 
-fn prepare_input<'a, E: ExtensionField>(
-    nv: usize,
-) -> (E, VirtualPolynomial<'a, E>, Vec<VirtualPolynomial<'a, E>>) {
+fn prepare_input<'a, E: ExtensionField>(nv: usize) -> (E, Vec<ArcMultilinearExtension<'a, E>>) {
     let mut rng = test_rng();
-    let max_thread_id = max_usable_threads();
-    let size_log2 = ceil_log2(max_thread_id);
-    let fs: [ArcMultilinearExtension<'a, E>; NUM_DEGREE] = array::from_fn(|_| {
-        let mle: ArcMultilinearExtension<'a, E> =
-            DenseMultilinearExtension::<E>::random(nv, &mut rng).into();
-        mle
-    });
-
-    let mut virtual_poly_v1 = VirtualPolynomial::new(nv);
-    virtual_poly_v1.add_mle_list(fs.to_vec(), E::ONE);
-
-    // devirgo version
-    let virtual_poly_v2: Vec<Vec<ArcMultilinearExtension<'a, E>>> = transpose(
-        fs.iter()
-            .map(|f| match &f.evaluations() {
-                multilinear_extensions::mle::FieldType::Base(evaluations) => evaluations
-                    .chunks((1 << nv) >> size_log2)
-                    .map(|chunk| {
-                        let mle: ArcMultilinearExtension<'a, E> =
-                            DenseMultilinearExtension::<E>::from_evaluations_vec(
-                                nv - size_log2,
-                                chunk.to_vec(),
-                            )
-                            .into();
-                        mle
-                    })
-                    .collect_vec(),
-                _ => unreachable!(),
-            })
-            .collect(),
-    );
-    let virtual_poly_v2: Vec<VirtualPolynomial<E>> = virtual_poly_v2
-        .into_iter()
-        .map(|fs| {
-            let mut virtual_polynomial = VirtualPolynomial::new(fs[0].num_vars());
-            virtual_polynomial.add_mle_list(fs, E::ONE);
-            virtual_polynomial
+    let fs = (0..NUM_DEGREE)
+        .map(|_| {
+            let mle: ArcMultilinearExtension<'a, E> =
+                DenseMultilinearExtension::<E>::random(nv, &mut rng).into();
+            mle
         })
-        .collect();
+        .collect_vec();
 
     let asserted_sum = fs
         .iter()
@@ -97,7 +65,7 @@ fn prepare_input<'a, E: ExtensionField>(
         .cloned()
         .sum::<E>();
 
-    (asserted_sum, virtual_poly_v1, virtual_poly_v2)
+    (asserted_sum, fs)
 }
 
 fn sumcheck_fn(c: &mut Criterion) {
@@ -116,12 +84,15 @@ fn sumcheck_fn(c: &mut Criterion) {
                     let mut time = Duration::new(0, 0);
                     for _ in 0..iters {
                         let mut prover_transcript = Transcript::new(b"test");
-                        let (_, virtual_poly, _) = { prepare_input(nv) };
+                        let (_, fs) = { prepare_input(nv) };
+
+                        let mut virtual_poly_v1 = VirtualPolynomial::new(nv);
+                        virtual_poly_v1.add_mle_list(fs.to_vec(), E::ONE);
 
                         let instant = std::time::Instant::now();
                         #[allow(deprecated)]
                         let (_sumcheck_proof_v1, _) = IOPProverState::<E>::prove_parallel(
-                            virtual_poly.clone(),
+                            virtual_poly_v1,
                             &mut prover_transcript,
                         );
                         let elapsed = instant.elapsed();
@@ -153,14 +124,14 @@ fn devirgo_sumcheck_fn(c: &mut Criterion) {
                     let mut time = Duration::new(0, 0);
                     for _ in 0..iters {
                         let mut prover_transcript = Transcript::new(b"test");
-                        let (_, _, virtual_poly_splitted) = { prepare_input(nv) };
+                        let (_, fs) = { prepare_input(nv) };
+
+                        let mut virtual_poly_v2 = VirtualPolynomials::new(threads, nv);
+                        virtual_poly_v2.add_mle_list(fs.iter().collect_vec(), E::ONE);
 
                         let instant = std::time::Instant::now();
-                        let (_sumcheck_proof_v2, _) = IOPProverState::<E>::prove_batch_polys(
-                            threads,
-                            virtual_poly_splitted,
-                            &mut prover_transcript,
-                        );
+                        let (_sumcheck_proof_v2, _) =
+                            IOPProverState::<E>::prove(virtual_poly_v2, &mut prover_transcript);
                         let elapsed = instant.elapsed();
                         time += elapsed;
                     }