feat: add shared base multiexp

halo2_proofs/src/arithmetic.rs

@@ -10,6 +10,7 @@ use group::{
     prime::PrimeCurveAffine,
     Curve, Group, GroupOpsOwned, ScalarMulOwned,
 };
+use multicore::{IntoParallelIterator, ParallelIterator};
 
 pub use halo2curves::{CurveAffine, CurveExt};
 
@@ -153,6 +154,47 @@ pub fn small_multiexp<C: CurveAffine>(coeffs: &[C::Scalar], bases: &[C]) -> C::C
     acc
 }
 
+/// Computes sum_i multi_coeffs[j][i] * bases[i] for each j.
+/// Tries to optimize memory locality from the shared bases.
+pub fn best_multiexp_shared_bases<C: CurveAffine, A, B>(
+    multi_coeffs: &[A],
+    bases: &[C],
+) -> Vec<C::Curve>
+where
+    A: AsRef<Polynomial<C::Scalar, B>> + Sync, // can't get around this trait in commit_lagrange_many...
+{
+    let bases_len = bases.len();
+    let num_threads = multicore::current_num_threads();
+    let chunk = if bases.len() < num_threads {
+        bases.len()
+    } else {
+        bases.len() / num_threads
+    };
+    let num_chunks = bases.chunks(chunk).len();
+    let mut multi_results = vec![vec![C::Curve::identity(); multi_coeffs.len()]; num_chunks];
+    multicore::scope(|scope| {
+        for (chunk_idx, (bases, results)) in bases
+            .chunks(chunk)
+            .zip(multi_results.iter_mut())
+            .enumerate()
+        {
+            let chunk_start = chunk_idx * chunk;
+            let chunk_end = std::cmp::min((chunk_idx + 1) * chunk, bases_len);
+            scope.spawn(move |_| {
+                for (coeffs, acc) in multi_coeffs.iter().zip(results.iter_mut()) {
+                    let coeffs = coeffs.as_ref();
+                    debug_assert_eq!(coeffs.len(), bases_len);
+                    multiexp_serial(&coeffs[chunk_start..chunk_end], bases, acc);
+                }
+            });
+        }
+    });
+    (0..multi_coeffs.len())
+        .into_par_iter()
+        .map(|j| multi_results.iter().map(|results| results[j]).sum())
+        .collect()
+}
+
 /// Performs a multi-exponentiation operation.
 ///
 /// This function will panic if coeffs and bases have a different length.
@@ -161,11 +203,8 @@ pub fn small_multiexp<C: CurveAffine>(coeffs: &[C::Scalar], bases: &[C]) -> C::C
 pub fn best_multiexp<C: CurveAffine>(coeffs: &[C::Scalar], bases: &[C]) -> C::Curve {
     assert_eq!(coeffs.len(), bases.len());
 
-    //println!("msm: {}", coeffs.len());
-
-    // let start = get_time();
     let num_threads = multicore::current_num_threads();
-    let res = if coeffs.len() > num_threads {
+    if coeffs.len() > num_threads {
         let chunk = coeffs.len() / num_threads;
         let num_chunks = coeffs.chunks(chunk).len();
         let mut results = vec![C::Curve::identity(); num_chunks];
@@ -187,14 +226,7 @@ pub fn best_multiexp<C: CurveAffine>(coeffs: &[C::Scalar], bases: &[C]) -> C::Cu
         let mut acc = C::Curve::identity();
         multiexp_serial(coeffs, bases, &mut acc);
         acc
-    };
-
-    // let duration = get_duration(start);
-    #[allow(unsafe_code)]
-    // unsafe {
-    //     MULTIEXP_TOTAL_TIME += duration;
-    // }
-    res
+    }
 }
 
 /// Dispatcher
@@ -464,9 +496,12 @@ pub fn bitreverse(mut n: usize, l: usize) -> usize {
 #[cfg(test)]
 use rand_core::OsRng;
 
-use crate::fft::{self, recursive::FFTData};
 #[cfg(test)]
 use crate::halo2curves::pasta::Fp;
+use crate::{
+    fft::{self, recursive::FFTData},
+    poly::Polynomial,
+};
 // use crate::plonk::{get_duration, get_time, start_measure, stop_measure};
 
 #[test]

halo2_proofs/src/plonk/prover.rs

@@ -7,10 +7,6 @@ use rand_core::RngCore;
 use std::hash::Hash;
 use std::marker::PhantomData;
 use std::ops::RangeTo;
-
-#[cfg(feature = "multicore")]
-use crate::multicore::IndexedParallelIterator;
-use crate::multicore::{IntoParallelIterator, ParallelIterator};
 use std::{collections::HashMap, iter};
 
 use super::{
@@ -300,11 +296,10 @@ where
                         }
                     }
                 } else {
-                    let instance_commitments_projective: Vec<_> =
-                        (&self.instance_single.instance_values)
-                            .into_par_iter()
-                            .map(|poly| self.params.commit_lagrange(poly, Blind::default()))
-                            .collect();
+                    let instance_commitments_projective = self.params.commit_lagrange_many(
+                        &self.instance_single.instance_values,
+                        vec![Blind::default(); self.instance_single.instance_values.len()],
+                    );
                     let mut instance_commitments =
                         vec![C::identity(); instance_commitments_projective.len()];
                     C::CurveExt::batch_normalize(
@@ -341,11 +336,9 @@ where
                 .iter()
                 .map(|_| Blind(F::random(&mut self.rng)))
                 .collect();
-            let advice_commitments_projective: Vec<_> = (&advice_values)
-                .into_par_iter()
-                .zip((&blinds).into_par_iter())
-                .map(|(poly, blind)| self.params.commit_lagrange(poly, *blind))
-                .collect();
+            let advice_commitments_projective: Vec<_> = self
+                .params
+                .commit_lagrange_many(&advice_values, blinds.clone());
             let mut advice_commitments = vec![C::identity(); advice_commitments_projective.len()];
             C::CurveExt::batch_normalize(&advice_commitments_projective, &mut advice_commitments);
             let advice_commitments = advice_commitments;

halo2_proofs/src/poly.rs

@@ -144,6 +144,12 @@ impl<F, B> DerefMut for Polynomial<F, B> {
     }
 }
 
+impl<F, B> AsRef<Polynomial<F, B>> for Polynomial<F, B> {
+    fn as_ref(&self) -> &Polynomial<F, B> {
+        self
+    }
+}
+
 impl<F, B> Polynomial<F, B> {
     /// Iterate over the values, which are either in coefficient or evaluation
     /// form depending on the basis `B`.

halo2_proofs/src/poly/commitment.rs

@@ -66,6 +66,19 @@ pub trait Params<'params, C: CurveAffine>: Sized + Clone {
         r: Blind<C::ScalarExt>,
     ) -> C::CurveExt;
 
+    fn commit_lagrange_many(
+        &self,
+        polys: &[impl AsRef<Polynomial<C::ScalarExt, LagrangeCoeff>> + Sync],
+        r: Vec<Blind<C::ScalarExt>>,
+    ) -> Vec<C::CurveExt> {
+        assert_eq!(polys.len(), r.len());
+        polys
+            .iter()
+            .zip(r)
+            .map(|(poly, r)| self.commit_lagrange(poly.as_ref(), r))
+            .collect()
+    }
+
     /// Writes params to a buffer.
     fn write<W: io::Write>(&self, writer: &mut W) -> io::Result<()>;
 

halo2_proofs/src/poly/kzg/commitment.rs

@@ -1,4 +1,4 @@
-use crate::arithmetic::{best_multiexp, g_to_lagrange, parallelize};
+use crate::arithmetic::{best_multiexp, best_multiexp_shared_bases, g_to_lagrange, parallelize};
 use crate::helpers::SerdeCurveAffine;
 use crate::poly::commitment::{Blind, CommitmentScheme, Params, ParamsProver, ParamsVerifier};
 use crate::poly::{Coeff, LagrangeCoeff, Polynomial};
@@ -305,6 +305,19 @@ where
         best_multiexp(poly, &self.g_lagrange[0..size])
     }
 
+    fn commit_lagrange_many(
+        &self,
+        polys: &[impl AsRef<Polynomial<E::Fr, LagrangeCoeff>> + Sync],
+        _: Vec<Blind<<E::G1Affine as halo2curves::CurveAffine>::ScalarExt>>,
+    ) -> Vec<E::G1> {
+        if polys.is_empty() {
+            return vec![];
+        }
+        let size = polys[0].as_ref().len();
+        assert!(self.n() >= size as u64);
+        best_multiexp_shared_bases(polys, &self.g_lagrange[0..size])
+    }
+
     /// Writes params to a buffer.
     fn write<W: io::Write>(&self, writer: &mut W) -> io::Result<()> {
         self.write_custom(writer, SerdeFormat::RawBytesUnchecked)