Skip to content

Commit

Permalink
feat: parallelize (cpu) shplonk prover
Browse files Browse the repository at this point in the history
  • Loading branch information
jonathanpwang committed Dec 26, 2022
1 parent 50ee8ad commit d6f562a
Show file tree
Hide file tree
Showing 2 changed files with 15 additions and 8 deletions.
8 changes: 4 additions & 4 deletions halo2_proofs/src/poly/kzg/multiopen/shplonk.rs
Original file line number Diff line number Diff line change
Expand Up @@ -80,7 +80,7 @@ where
assert_eq!(*point, query.get_point());
}
// All points appear in queries
let super_point_set: Vec<F> = rotation_point_map.values().cloned().collect();
let super_point_set: Vec<F> = rotation_point_map.values().copied().collect();

// Collect rotation sets for each commitment
// Example elements in the vector:
Expand All @@ -90,7 +90,7 @@ where
// (C_3, {r_2, r_3, r_4}),
// ...
let mut commitment_rotation_set_map: Vec<(Q::Commitment, Vec<F>)> = vec![];
for query in queries.clone() {
for query in queries.iter() {
let rotation = query.get_point();
if let Some(pos) = commitment_rotation_set_map
.iter()
Expand All @@ -114,8 +114,7 @@ where
let mut rotation_set_commitment_map = Vec::<(Vec<_>, Vec<Q::Commitment>)>::new();
for (commitment, rotation_set) in commitment_rotation_set_map.iter() {
if let Some(pos) = rotation_set_commitment_map.iter().position(|(set, _)| {
BTreeSet::<F>::from_iter(set.iter().cloned())
== BTreeSet::<F>::from_iter(rotation_set.iter().cloned())
BTreeSet::<&F>::from_iter(set.iter()) == BTreeSet::<&F>::from_iter(rotation_set.iter())
}) {
let (_, commitments) = &mut rotation_set_commitment_map[pos];
if !commitments.contains(commitment) {
Expand All @@ -126,6 +125,7 @@ where
}
}

// TODO: parallelize
let rotation_sets = rotation_set_commitment_map
.into_iter()
.map(|(rotations, commitments)| {
Expand Down
15 changes: 11 additions & 4 deletions halo2_proofs/src/poly/kzg/multiopen/shplonk/prover.rs
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,8 @@ use ff::Field;
use group::Curve;
use halo2curves::pairing::Engine;
use rand_core::RngCore;
use rayon::iter::ParallelIterator;
use rayon::prelude::{IntoParallelIterator, IntoParallelRefIterator};
use std::fmt::Debug;
use std::io::{self, Write};
use std::marker::PhantomData;
Expand Down Expand Up @@ -171,11 +173,11 @@ impl<'params, E: Engine + Debug> Prover<'params, KZGCommitmentScheme<E>>
);

let rotation_sets: Vec<RotationSetExtension<E::G1Affine>> = rotation_sets
.iter()
.par_iter()
.map(|rotation_set| {
let commitments: Vec<CommitmentExtension<E::G1Affine>> = rotation_set
.commitments
.iter()
.par_iter()
.map(|commitment_data| commitment_data.extend(rotation_set.points.clone()))
.collect();
rotation_set.extend(commitments)
Expand All @@ -184,9 +186,13 @@ impl<'params, E: Engine + Debug> Prover<'params, KZGCommitmentScheme<E>>

let v: ChallengeV<_> = transcript.squeeze_challenge_scalar();

let quotient_polynomials = rotation_sets.iter().map(quotient_contribution);
let quotient_polynomials = rotation_sets
.par_iter()
.map(quotient_contribution)
.collect::<Vec<_>>();

let h_x: Polynomial<E::Scalar, Coeff> = quotient_polynomials
.into_iter()
.zip(powers(*v))
.map(|(poly, power_of_v)| poly * power_of_v)
.reduce(|acc, poly| acc + &poly)
Expand Down Expand Up @@ -235,7 +241,7 @@ impl<'params, E: Engine + Debug> Prover<'params, KZGCommitmentScheme<E>>
Vec<Polynomial<E::Scalar, Coeff>>,
Vec<E::Scalar>,
) = rotation_sets
.into_iter()
.into_par_iter()
.map(linearisation_contribution)
.unzip();

Expand All @@ -249,6 +255,7 @@ impl<'params, E: Engine + Debug> Prover<'params, KZGCommitmentScheme<E>>
let l_x = l_x - &(h_x * zt_eval);

// sanity check
#[cfg(debug_assertions)]
{
let must_be_zero = eval_polynomial(&l_x.values[..], *u);
assert_eq!(must_be_zero, E::Scalar::zero());
Expand Down

0 comments on commit d6f562a

Please sign in to comment.