Batch all (#89)

- use sumcheck to batch open PCS
- split Prod and witness into two batches
- benchmark code

arithmetic/src/lib.rs

@@ -6,10 +6,10 @@ mod virtual_polynomial;
 
 pub use errors::ArithErrors;
 pub use multilinear_polynomial::{
-    evaluate_no_par, evaluate_opt, fix_first_variable, fix_variables, identity_permutation_mle,
-    merge_polynomials, random_mle_list, random_permutation_mle, random_zero_mle_list,
-    DenseMultilinearExtension,
+    evaluate_no_par, evaluate_opt, fix_last_variables, fix_last_variables_no_par, fix_variables,
+    identity_permutation_mle, merge_polynomials, random_mle_list, random_permutation_mle,
+    random_zero_mle_list, DenseMultilinearExtension,
 };
 pub use univariate_polynomial::{build_l, get_uni_domain};
 pub use util::{bit_decompose, gen_eval_point, get_batched_nv, get_index};
-pub use virtual_polynomial::{build_eq_x_r, VPAuxInfo, VirtualPolynomial};
+pub use virtual_polynomial::{build_eq_x_r, build_eq_x_r_vec, VPAuxInfo, VirtualPolynomial};

arithmetic/src/multilinear_polynomial.rs

@@ -123,36 +123,24 @@ pub fn fix_variables<F: Field>(
     DenseMultilinearExtension::<F>::from_evaluations_slice(nv - dim, &poly[..(1 << (nv - dim))])
 }
 
-pub fn fix_first_variable<F: Field>(
-    poly: &DenseMultilinearExtension<F>,
-    partial_point: &F,
-) -> DenseMultilinearExtension<F> {
-    assert!(poly.num_vars != 0, "invalid size of partial point");
-
-    let nv = poly.num_vars;
-    let res = fix_one_variable_helper(&poly.evaluations, nv, partial_point);
-    DenseMultilinearExtension::<F>::from_evaluations_slice(nv - 1, &res)
-}
-
 fn fix_one_variable_helper<F: Field>(data: &[F], nv: usize, point: &F) -> Vec<F> {
     let mut res = vec![F::zero(); 1 << (nv - 1)];
-    let one_minus_p = F::one() - point;
 
     // evaluate single variable of partial point from left to right
     #[cfg(not(feature = "parallel"))]
-    for b in 0..(1 << (nv - 1)) {
-        res[b] = data[b << 1] * one_minus_p + data[(b << 1) + 1] * point;
+    for i in 0..(1 << (nv - 1)) {
+        res[i] = data[i] + (data[(i << 1) + 1] - data[i << 1]) * point;
     }
 
     #[cfg(feature = "parallel")]
     if nv >= 13 {
         // on my computer we parallelization doesn't help till nv >= 13
         res.par_iter_mut().enumerate().for_each(|(i, x)| {
-            *x = data[i << 1] * one_minus_p + data[(i << 1) + 1] * point;
+            *x = data[i << 1] + (data[(i << 1) + 1] - data[i << 1]) * point;
         });
     } else {
-        for b in 0..(1 << (nv - 1)) {
-            res[b] = data[b << 1] * one_minus_p + data[(b << 1) + 1] * point;
+        for i in 0..(1 << (nv - 1)) {
+            res[i] = data[i << 1] + (data[(i << 1) + 1] - data[i << 1]) * point;
         }
     }
 
@@ -178,9 +166,8 @@ fn fix_variables_no_par<F: Field>(
     // evaluate single variable of partial point from left to right
     for i in 1..dim + 1 {
         let r = partial_point[i - 1];
-        let one_minus_r = F::one() - r;
         for b in 0..(1 << (nv - i)) {
-            poly[b] = poly[b << 1] * one_minus_r + poly[(b << 1) + 1] * r;
+            poly[b] = poly[b << 1] + (poly[(b << 1) + 1] - poly[b << 1]) * r;
         }
     }
     DenseMultilinearExtension::from_evaluations_slice(nv - dim, &poly[..(1 << (nv - dim))])
@@ -210,3 +197,71 @@ pub fn merge_polynomials<F: PrimeField>(
         merged_nv, scalars,
     )))
 }
+
+pub fn fix_last_variables_no_par<F: PrimeField>(
+    poly: &DenseMultilinearExtension<F>,
+    partial_point: &[F],
+) -> DenseMultilinearExtension<F> {
+    let mut res = fix_last_variable_no_par(poly, partial_point.last().unwrap());
+    for p in partial_point.iter().rev().skip(1) {
+        res = fix_last_variable_no_par(&res, p);
+    }
+    res
+}
+
+fn fix_last_variable_no_par<F: PrimeField>(
+    poly: &DenseMultilinearExtension<F>,
+    partial_point: &F,
+) -> DenseMultilinearExtension<F> {
+    let nv = poly.num_vars();
+    let half_len = 1 << (nv - 1);
+    let mut res = vec![F::zero(); half_len];
+    for (i, e) in res.iter_mut().enumerate().take(half_len) {
+        *e = poly.evaluations[i]
+            + *partial_point * (poly.evaluations[i + half_len] - poly.evaluations[i]);
+    }
+    DenseMultilinearExtension::from_evaluations_vec(nv - 1, res)
+}
+pub fn fix_last_variables<F: PrimeField>(
+    poly: &DenseMultilinearExtension<F>,
+    partial_point: &[F],
+) -> DenseMultilinearExtension<F> {
+    assert!(
+        partial_point.len() <= poly.num_vars,
+        "invalid size of partial point"
+    );
+    let nv = poly.num_vars;
+    let mut poly = poly.evaluations.to_vec();
+    let dim = partial_point.len();
+    // evaluate single variable of partial point from left to right
+    for (i, point) in partial_point.iter().rev().enumerate().take(dim) {
+        poly = fix_last_variable_helper(&poly, nv - i, point);
+    }
+
+    DenseMultilinearExtension::<F>::from_evaluations_slice(nv - dim, &poly[..(1 << (nv - dim))])
+}
+
+fn fix_last_variable_helper<F: Field>(data: &[F], nv: usize, point: &F) -> Vec<F> {
+    let half_len = 1 << (nv - 1);
+    let mut res = vec![F::zero(); half_len];
+
+    // evaluate single variable of partial point from left to right
+    #[cfg(not(feature = "parallel"))]
+    for b in 0..half_len {
+        res[b] = data[b] + (data[b + half_len] - data[b]) * point;
+    }
+
+    #[cfg(feature = "parallel")]
+    if nv >= 13 {
+        // on my computer we parallelization doesn't help till nv >= 13
+        res.par_iter_mut().enumerate().for_each(|(i, x)| {
+            *x = data[i] + (data[i + half_len] - data[i]) * point;
+        });
+    } else {
+        for b in 0..(1 << (nv - 1)) {
+            res[b] = data[b] + (data[b + half_len] - data[b]) * point;
+        }
+    }
+
+    res
+}

arithmetic/src/virtual_polynomial.rs

@@ -10,6 +10,7 @@ use ark_std::{
     rand::{Rng, RngCore},
     start_timer,
 };
+use rayon::prelude::*;
 use std::{cmp::max, collections::HashMap, marker::PhantomData, ops::Add, rc::Rc};
 
 #[rustfmt::skip]
@@ -324,16 +325,29 @@ impl<F: PrimeField> VirtualPolynomial<F> {
     }
 }
 
-// This function build the eq(x, r) polynomial for any given r.
-//
-// Evaluate
-//      eq(x,y) = \prod_i=1^num_var (x_i * y_i + (1-x_i)*(1-y_i))
-// over r, which is
-//      eq(x,y) = \prod_i=1^num_var (x_i * r_i + (1-x_i)*(1-r_i))
+/// This function build the eq(x, r) polynomial for any given r.
+///
+/// Evaluate
+///      eq(x,y) = \prod_i=1^num_var (x_i * y_i + (1-x_i)*(1-y_i))
+/// over r, which is
+///      eq(x,y) = \prod_i=1^num_var (x_i * r_i + (1-x_i)*(1-r_i))
 pub fn build_eq_x_r<F: PrimeField>(
     r: &[F],
 ) -> Result<Rc<DenseMultilinearExtension<F>>, ArithErrors> {
-    let start = start_timer!(|| "zero check build eq_x_r");
+    let evals = build_eq_x_r_vec(r)?;
+    let mle = DenseMultilinearExtension::from_evaluations_vec(r.len(), evals);
+
+    Ok(Rc::new(mle))
+}
+/// This function build the eq(x, r) polynomial for any given r, and output the
+/// evaluation of eq(x, r) in its vector form.
+///
+/// Evaluate
+///      eq(x,y) = \prod_i=1^num_var (x_i * y_i + (1-x_i)*(1-y_i))
+/// over r, which is
+///      eq(x,y) = \prod_i=1^num_var (x_i * r_i + (1-x_i)*(1-r_i))
+pub fn build_eq_x_r_vec<F: PrimeField>(r: &[F]) -> Result<Vec<F>, ArithErrors> {
+    let start = start_timer!(|| format!("build eq_x_r of size {}", r.len()));
 
     // we build eq(x,r) from its evaluations
     // we want to evaluate eq(x,r) over x \in {0, 1}^num_vars
@@ -349,11 +363,8 @@ pub fn build_eq_x_r<F: PrimeField>(
     let mut eval = Vec::new();
     build_eq_x_r_helper(r, &mut eval)?;
 
-    let mle = DenseMultilinearExtension::from_evaluations_vec(r.len(), eval);
-
-    let res = Rc::new(mle);
     end_timer!(start);
-    Ok(res)
+    Ok(eval)
 }
 
 /// A helper function to build eq(x, r) recursively.
@@ -373,13 +384,24 @@ fn build_eq_x_r_helper<F: PrimeField>(r: &[F], buf: &mut Vec<F>) -> Result<(), A
         // for the current step we will need
         // if x_0 = 0:   (1-r0) * [b_1, ..., b_k]
         // if x_0 = 1:   r0 * [b_1, ..., b_k]
+        // let mut res = vec![];
+        // for &b_i in buf.iter() {
+        //     let tmp = r[0] * b_i;
+        //     res.push(b_i - tmp);
+        //     res.push(tmp);
+        // }
+        // *buf = res;
 
-        let mut res = vec![];
-        for &b_i in buf.iter() {
-            let tmp = r[0] * b_i;
-            res.push(b_i - tmp);
-            res.push(tmp);
-        }
+        let mut res = vec![F::zero(); buf.len() << 1];
+        res.par_iter_mut().enumerate().for_each(|(i, val)| {
+            let bi = buf[i >> 1];
+            let tmp = r[0] * bi;
+            if i & 1 == 0 {
+                *val = bi - tmp;
+            } else {
+                *val = tmp;
+            }
+        });
         *buf = res;
     }