feat(sumcheck): build h1(x) without fix mles at r0 from first round (#29)

* build round 2 univariate poly from original evaluations

* refactor

* misc: support runtime mode and falleback with legacy feature (#30)

* bug fix: sumcheck phase 2 always use legacy mode

* add memory reduction benchmark

* cargo fmt

* revise the build_uni_poly_round2 algorithm to reuse fi(r0,0,b) and fi(r0,1,b) among block (b,fi)

* add new feature s.t. sumcheck prover follows reduced peak memory flow and fix a bug

---------

Co-authored-by: Ming <hero78119@gmail.com>

Author: kunxian-xia

crates/multilinear_extensions/src/mle.rs

@@ -1,4 +1,4 @@
-use std::{any::TypeId, borrow::Cow, mem, sync::Arc};
+use std::{any::TypeId, borrow::Cow, mem, ops::Range, sync::Arc};
 
 use crate::{
     field_type_mut_map,
@@ -231,6 +231,14 @@ impl<'a, E: ExtensionField> FieldType<'a, E> {
             FieldType::Unreachable => unreachable!(),
         }
     }
+
+    pub fn as_slice(&self, range: Range<usize>) -> Either<&[E::BaseField], &[E]> {
+        match self {
+            FieldType::Base(slice) => Either::Left(&slice[range]),
+            FieldType::Ext(slice) => Either::Right(&slice[range]),
+            FieldType::Unreachable => unreachable!(),
+        }
+    }
 }
 
 impl<'a, E: ExtensionField> PartialEq for FieldType<'a, E> {
@@ -567,6 +575,85 @@ impl<'a, E: ExtensionField> MultilinearExtension<'a, E> {
         self.num_vars = nv - partial_point.len();
     }
 
+    // compute f(r0,r1,b) from block { f(0,0,b), f(1,0,b), f(0,1,b), f(1,1,b) }
+    #[inline(always)]
+    fn eval_block_2_vars_base(block: &[E::BaseField], r0: E, r1: E) -> E {
+        // f(r0,r1,b) = (1 - r1) * f(r0,0,b) + r1 * f(r0,1,b)
+        // f(r0,0,b) = (1 - r0) * f(0,0,b) + r0 * f(1,0,b)
+        // f(r0,1,b) = (1 - r0) * f(0,1,b) + r0 * f(1,1,b)
+        let y0: E = r0 * (block[1] - block[0]) + block[0];
+        let y1: E = r0 * (block[3] - block[2]) + block[2];
+        y0 + (y1 - y0) * r1
+    }
+
+    #[inline(always)]
+    fn eval_block_2_vars_ext(block: &[E], r0: E, r1: E) -> E {
+        // f(r0,r1,b) = (1 - r1) * f(r0,0,b) + r1 * f(r0,1,b)
+        // f(r0,0,b) = (1 - r0) * f(0,0,b) + r0 * f(1,0,b)
+        // f(r0,1,b) = (1 - r0) * f(0,1,b) + r0 * f(1,1,b)
+        let y0: E = block[0] + (block[1] - block[0]) * r0;
+        let y1: E = block[2] + (block[3] - block[2]) * r0;
+        y0 + (y1 - y0) * r1
+    }
+
+    /// Reduce the number of variables by 2 in one pass.
+    ///
+    /// This avoids calling `fix_variables` twice and directly computes
+    /// `f(r0, r1, ..)` from 4-point blocks.
+    pub fn fix_two_variables(&self, r0: E, r1: E) -> Self {
+        assert!(self.num_vars() >= 2, "num_vars {} < 2", self.num_vars());
+        let nv = self.num_vars();
+        match self.evaluations() {
+            FieldType::Base(slice) => MultilinearExtension::from_evaluations_ext_vec(
+                nv - 2,
+                slice
+                    .chunks(4)
+                    .map(|buf| Self::eval_block_2_vars_base(buf, r0, r1))
+                    .collect(),
+            ),
+            FieldType::Ext(slice) => MultilinearExtension::from_evaluations_ext_vec(
+                nv - 2,
+                slice
+                    .chunks(4)
+                    .map(|buf| Self::eval_block_2_vars_ext(buf, r0, r1))
+                    .collect(),
+            ),
+            FieldType::Unreachable => unreachable!(),
+        }
+    }
+
+    /// In-place variant of `fix_two_variables`.
+    pub fn fix_two_variables_in_place(&mut self, r0: E, r1: E) {
+        assert!(self.is_mut());
+        assert!(self.num_vars() >= 2, "num_vars {} < 2", self.num_vars());
+        let nv = self.num_vars();
+
+        match &mut self.evaluations {
+            FieldType::Base(slice) => {
+                let ext_vec = slice
+                    .chunks(4)
+                    .map(|buf| Self::eval_block_2_vars_base(buf, r0, r1))
+                    .collect();
+                let _ = std::mem::replace(
+                    &mut self.evaluations,
+                    FieldType::Ext(SmartSlice::Owned(ext_vec)),
+                );
+            }
+            FieldType::Ext(slice) => {
+                let slice_mut = slice.to_mut();
+                let new_len = 1 << (nv - 2);
+                for i in 0..new_len {
+                    let b = i << 2;
+                    slice_mut[i] = Self::eval_block_2_vars_ext(&slice_mut[b..b + 4], r0, r1);
+                }
+                slice.truncate_mut(new_len);
+            }
+            FieldType::Unreachable => unreachable!(),
+        }
+
+        self.num_vars = nv - 2;
+    }
+
     /// Evaluate the MLE at a give point.
     /// Returns an error if the MLE length does not match the point.
     pub fn evaluate(&self, point: &[E]) -> E {

crates/sumcheck/Cargo.toml

@@ -31,8 +31,13 @@ rand.workspace = true
 [features]
 nightly-features = ["ff_ext/nightly-features"]
 parallel = ["p3/parallel", "multilinear_extensions/parallel"]
+reduce-peak-memory = []
 
 
 [[bench]]
 harness = false
 name = "devirgo_sumcheck"
+
+[[bench]]
+harness = false
+name = "memory_usage"

crates/sumcheck/benches/devirgo_sumcheck.rs

@@ -9,15 +9,20 @@ use ff_ext::{ExtensionField, GoldilocksExt2};
 use itertools::Itertools;
 use p3::field::FieldAlgebra;
 use rand::thread_rng;
-use sumcheck::structs::IOPProverState;
+use sumcheck::structs::{IOPProverState, SumcheckProverMode};
 
 use multilinear_extensions::{
     mle::MultilinearExtension, monomial::Term, op_mle, util::max_usable_threads,
     virtual_poly::VirtualPolynomial, virtual_polys::VirtualPolynomials,
 };
 use transcript::BasicTranscript as Transcript;
 
-criterion_group!(benches, sumcheck_fn, devirgo_sumcheck_fn,);
+criterion_group!(
+    benches,
+    sumcheck_fn,
+    devirgo_sumcheck_fn,
+    devirgo_sumcheck_reduced_peak_memory_fn,
+);
 criterion_main!(benches);
 
 const NUM_SAMPLES: usize = 10;
@@ -223,3 +228,138 @@ fn devirgo_sumcheck_fn(c: &mut Criterion) {
         group.finish();
     }
 }
+
+fn devirgo_sumcheck_reduced_peak_memory_fn(c: &mut Criterion) {
+    type E = GoldilocksExt2;
+
+    let threads = max_usable_threads();
+    for nv in NV {
+        let mut group = c.benchmark_group(format!("devirgo_reduced_peak_memory_nv_{}", nv));
+        group.sample_size(NUM_SAMPLES);
+
+        group.bench_function(
+            BenchmarkId::new(
+                "prove_sumcheck",
+                format!("devirgo_reduced_peak_memory_nv_{}", nv),
+            ),
+            |b| {
+                b.iter_custom(|iters| {
+                    let mut time = Duration::new(0, 0);
+                    for _ in 0..iters {
+                        let mut prover_transcript = Transcript::new(b"test");
+                        let (_, fs) = { prepare_input(nv) };
+
+                        let virtual_poly = VirtualPolynomials::new_from_monimials(
+                            threads,
+                            nv,
+                            vec![Term {
+                                scalar: Either::Right(E::ONE),
+                                product: fs.iter().map(Either::Left).collect_vec(),
+                            }],
+                        );
+                        let instant = std::time::Instant::now();
+                        let _ = IOPProverState::<E>::prove_with_mode(
+                            virtual_poly,
+                            &mut prover_transcript,
+                            SumcheckProverMode::ReducedPeakMemory,
+                        );
+                        time += instant.elapsed();
+                    }
+                    time
+                });
+            },
+        );
+
+        group.bench_function(
+            BenchmarkId::new(
+                "prove_sumcheck_ext",
+                format!("devirgo_reduced_peak_memory_nv_{}", nv),
+            ),
+            |b| {
+                b.iter_custom(|iters| {
+                    let mut time = Duration::new(0, 0);
+                    for _ in 0..iters {
+                        let mut prover_transcript = Transcript::new(b"test");
+                        let (_, fs) = { prepare_input(nv) };
+                        let fs = fs
+                            .into_iter()
+                            .map(|mle: MultilinearExtension<'_, E>| {
+                                MultilinearExtension::from_evaluation_vec_smart(
+                                    mle.num_vars(),
+                                    mle.get_base_field_vec()
+                                        .iter()
+                                        .map(E::from_ref_base)
+                                        .collect_vec(),
+                                )
+                            })
+                            .collect_vec();
+
+                        let virtual_poly = VirtualPolynomials::new_from_monimials(
+                            threads,
+                            nv,
+                            vec![Term {
+                                scalar: Either::Right(E::ONE),
+                                product: fs.iter().map(Either::Left).collect_vec(),
+                            }],
+                        );
+                        let instant = std::time::Instant::now();
+                        let _ = IOPProverState::<E>::prove_with_mode(
+                            virtual_poly,
+                            &mut prover_transcript,
+                            SumcheckProverMode::ReducedPeakMemory,
+                        );
+                        time += instant.elapsed();
+                    }
+                    time
+                });
+            },
+        );
+
+        group.bench_function(
+            BenchmarkId::new(
+                "prove_sumcheck_ext_in_place",
+                format!("devirgo_reduced_peak_memory_nv_{}", nv),
+            ),
+            |b| {
+                b.iter_custom(|iters| {
+                    let mut time = Duration::new(0, 0);
+                    for _ in 0..iters {
+                        let mut prover_transcript = Transcript::new(b"test");
+                        let (_, fs) = { prepare_input(nv) };
+                        let mut fs = fs
+                            .into_iter()
+                            .map(|mle: MultilinearExtension<'_, E>| {
+                                MultilinearExtension::from_evaluation_vec_smart(
+                                    mle.num_vars(),
+                                    mle.get_base_field_vec()
+                                        .iter()
+                                        .map(E::from_ref_base)
+                                        .collect_vec(),
+                                )
+                            })
+                            .collect_vec();
+
+                        let virtual_poly = VirtualPolynomials::new_from_monimials(
+                            threads,
+                            nv,
+                            vec![Term {
+                                scalar: Either::Right(E::ONE),
+                                product: fs.iter_mut().map(Either::Right).collect_vec(),
+                            }],
+                        );
+                        let instant = std::time::Instant::now();
+                        let _ = IOPProverState::<E>::prove_with_mode(
+                            virtual_poly,
+                            &mut prover_transcript,
+                            SumcheckProverMode::ReducedPeakMemory,
+                        );
+                        time += instant.elapsed();
+                    }
+                    time
+                });
+            },
+        );
+
+        group.finish();
+    }
+}