Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Parallelization optimizations #105

Merged
merged 8 commits into from
Dec 19, 2022
Merged

Parallelization optimizations #105

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
648ea51
arc
chancharles92 Dec 16, 2022
81073df
batch inversion
chancharles92 Dec 16, 2022
5d6985b
more parallelization
chancharles92 Dec 16, 2022
c48c1b9
more parallelization
chancharles92 Dec 16, 2022
44833a0
remove unncessary timer
chancharles92 Dec 16, 2022
35620b9
tmp commit for benchmark multithread
chancharles92 Dec 16, 2022
c18d5d6
uncomment
chancharles92 Dec 16, 2022
65b1938
typofix
chancharles92 Dec 19, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
50 changes: 18 additions & 32 deletions arithmetic/src/multilinear_polynomial.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@ use ark_poly::MultilinearExtension;
use ark_std::{end_timer, rand::RngCore, start_timer};
#[cfg(feature = "parallel")]
use rayon::prelude::{IndexedParallelIterator, IntoParallelRefMutIterator, ParallelIterator};
use std::rc::Rc;
use std::sync::Arc;

pub use ark_poly::DenseMultilinearExtension;

Expand All @@ -16,7 +16,7 @@ pub fn random_mle_list<F: PrimeField, R: RngCore>(
nv: usize,
degree: usize,
rng: &mut R,
) -> (Vec<Rc<DenseMultilinearExtension<F>>>, F) {
) -> (Vec<Arc<DenseMultilinearExtension<F>>>, F) {
let start = start_timer!(|| "sample random mle list");
let mut multiplicands = Vec::with_capacity(degree);
for _ in 0..degree {
Expand All @@ -37,7 +37,7 @@ pub fn random_mle_list<F: PrimeField, R: RngCore>(

let list = multiplicands
.into_iter()
.map(|x| Rc::new(DenseMultilinearExtension::from_evaluations_vec(nv, x)))
.map(|x| Arc::new(DenseMultilinearExtension::from_evaluations_vec(nv, x)))
.collect();

end_timer!(start);
Expand All @@ -49,7 +49,7 @@ pub fn random_zero_mle_list<F: PrimeField, R: RngCore>(
nv: usize,
degree: usize,
rng: &mut R,
) -> Vec<Rc<DenseMultilinearExtension<F>>> {
) -> Vec<Arc<DenseMultilinearExtension<F>>> {
let start = start_timer!(|| "sample random zero mle list");

let mut multiplicands = Vec::with_capacity(degree);
Expand All @@ -65,7 +65,7 @@ pub fn random_zero_mle_list<F: PrimeField, R: RngCore>(

let list = multiplicands
.into_iter()
.map(|x| Rc::new(DenseMultilinearExtension::from_evaluations_vec(nv, x)))
.map(|x| Arc::new(DenseMultilinearExtension::from_evaluations_vec(nv, x)))
.collect();

end_timer!(start);
Expand All @@ -81,12 +81,12 @@ pub fn identity_permutation<F: PrimeField>(num_vars: usize, num_chunks: usize) -
pub fn identity_permutation_mles<F: PrimeField>(
num_vars: usize,
num_chunks: usize,
) -> Vec<Rc<DenseMultilinearExtension<F>>> {
) -> Vec<Arc<DenseMultilinearExtension<F>>> {
let mut res = vec![];
for i in 0..num_chunks {
let shift = (i * (1 << num_vars)) as u64;
let s_id_vec = (shift..shift + (1u64 << num_vars)).map(F::from).collect();
res.push(Rc::new(DenseMultilinearExtension::from_evaluations_vec(
res.push(Arc::new(DenseMultilinearExtension::from_evaluations_vec(
num_vars, s_id_vec,
)));
}
Expand All @@ -113,12 +113,12 @@ pub fn random_permutation_mles<F: PrimeField, R: RngCore>(
num_vars: usize,
num_chunks: usize,
rng: &mut R,
) -> Vec<Rc<DenseMultilinearExtension<F>>> {
) -> Vec<Arc<DenseMultilinearExtension<F>>> {
let s_perm_vec = random_permutation(num_vars, num_chunks, rng);
let mut res = vec![];
let n = 1 << num_vars;
for i in 0..num_chunks {
res.push(Rc::new(DenseMultilinearExtension::from_evaluations_vec(
res.push(Arc::new(DenseMultilinearExtension::from_evaluations_vec(
num_vars,
s_perm_vec[i * n..i * n + n].to_vec(),
)));
Expand Down Expand Up @@ -160,16 +160,9 @@ fn fix_one_variable_helper<F: Field>(data: &[F], nv: usize, point: &F) -> Vec<F>
}

#[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] + (data[(i << 1) + 1] - data[i << 1]) * point;
});
} else {
for i in 0..(1 << (nv - 1)) {
res[i] = data[i << 1] + (data[(i << 1) + 1] - data[i << 1]) * point;
}
}
res.par_iter_mut().enumerate().for_each(|(i, x)| {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should we continue to make parallel dependent on nv?

*x = data[i << 1] + (data[(i << 1) + 1] - data[i << 1]) * point;
});

res
}
Expand Down Expand Up @@ -203,8 +196,8 @@ fn fix_variables_no_par<F: Field>(
/// merge a set of polynomials. Returns an error if the
/// polynomials do not share a same number of nvs.
pub fn merge_polynomials<F: PrimeField>(
polynomials: &[Rc<DenseMultilinearExtension<F>>],
) -> Result<Rc<DenseMultilinearExtension<F>>, ArithErrors> {
polynomials: &[Arc<DenseMultilinearExtension<F>>],
) -> Result<Arc<DenseMultilinearExtension<F>>, ArithErrors> {
let nv = polynomials[0].num_vars();
for poly in polynomials.iter() {
if nv != poly.num_vars() {
Expand All @@ -220,7 +213,7 @@ pub fn merge_polynomials<F: PrimeField>(
scalars.extend_from_slice(poly.to_evaluations().as_slice());
}
scalars.extend_from_slice(vec![F::zero(); (1 << merged_nv) - scalars.len()].as_ref());
Ok(Rc::new(DenseMultilinearExtension::from_evaluations_vec(
Ok(Arc::new(DenseMultilinearExtension::from_evaluations_vec(
merged_nv, scalars,
)))
}
Expand Down Expand Up @@ -279,16 +272,9 @@ fn fix_last_variable_helper<F: Field>(data: &[F], nv: usize, point: &F) -> Vec<F
}

#[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.par_iter_mut().enumerate().for_each(|(i, x)| {
*x = data[i] + (data[i + half_len] - data[i]) * point;
});

res
}
34 changes: 14 additions & 20 deletions arithmetic/src/virtual_polynomial.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@ use ark_std::{
start_timer,
};
use rayon::prelude::*;
use std::{cmp::max, collections::HashMap, marker::PhantomData, ops::Add, rc::Rc};
use std::{cmp::max, collections::HashMap, marker::PhantomData, ops::Add, sync::Arc};

#[rustfmt::skip]
/// A virtual polynomial is a sum of products of multilinear polynomials;
Expand Down Expand Up @@ -48,7 +48,7 @@ pub struct VirtualPolynomial<F: PrimeField> {
pub products: Vec<(F, Vec<usize>)>,
/// Stores multilinear extensions in which product multiplicand can refer
/// to.
pub flattened_ml_extensions: Vec<Rc<DenseMultilinearExtension<F>>>,
pub flattened_ml_extensions: Vec<Arc<DenseMultilinearExtension<F>>>,
/// Pointers to the above poly extensions
raw_pointers_lookup_table: HashMap<*const DenseMultilinearExtension<F>, usize>,
}
Expand All @@ -71,7 +71,7 @@ impl<F: PrimeField> Add for &VirtualPolynomial<F> {
let start = start_timer!(|| "virtual poly add");
let mut res = self.clone();
for products in other.products.iter() {
let cur: Vec<Rc<DenseMultilinearExtension<F>>> = products
let cur: Vec<Arc<DenseMultilinearExtension<F>>> = products
.1
.iter()
.map(|&x| other.flattened_ml_extensions[x].clone())
Expand Down Expand Up @@ -102,8 +102,8 @@ impl<F: PrimeField> VirtualPolynomial<F> {
}

/// Creates an new virtual polynomial from a MLE and its coefficient.
pub fn new_from_mle(mle: &Rc<DenseMultilinearExtension<F>>, coefficient: F) -> Self {
let mle_ptr: *const DenseMultilinearExtension<F> = Rc::as_ptr(mle);
pub fn new_from_mle(mle: &Arc<DenseMultilinearExtension<F>>, coefficient: F) -> Self {
let mle_ptr: *const DenseMultilinearExtension<F> = Arc::as_ptr(mle);
let mut hm = HashMap::new();
hm.insert(mle_ptr, 0);

Expand All @@ -129,10 +129,10 @@ impl<F: PrimeField> VirtualPolynomial<F> {
/// `coefficient`.
pub fn add_mle_list(
&mut self,
mle_list: impl IntoIterator<Item = Rc<DenseMultilinearExtension<F>>>,
mle_list: impl IntoIterator<Item = Arc<DenseMultilinearExtension<F>>>,
coefficient: F,
) -> Result<(), ArithErrors> {
let mle_list: Vec<Rc<DenseMultilinearExtension<F>>> = mle_list.into_iter().collect();
let mle_list: Vec<Arc<DenseMultilinearExtension<F>>> = mle_list.into_iter().collect();
let mut indexed_product = Vec::with_capacity(mle_list.len());

if mle_list.is_empty() {
Expand All @@ -151,7 +151,7 @@ impl<F: PrimeField> VirtualPolynomial<F> {
)));
}

let mle_ptr: *const DenseMultilinearExtension<F> = Rc::as_ptr(&mle);
let mle_ptr: *const DenseMultilinearExtension<F> = Arc::as_ptr(&mle);
if let Some(index) = self.raw_pointers_lookup_table.get(&mle_ptr) {
indexed_product.push(*index)
} else {
Expand All @@ -171,7 +171,7 @@ impl<F: PrimeField> VirtualPolynomial<F> {
/// Returns an error if the MLE has a different `num_vars` from self.
pub fn mul_by_mle(
&mut self,
mle: Rc<DenseMultilinearExtension<F>>,
mle: Arc<DenseMultilinearExtension<F>>,
coefficient: F,
) -> Result<(), ArithErrors> {
let start = start_timer!(|| "mul by mle");
Expand All @@ -183,7 +183,7 @@ impl<F: PrimeField> VirtualPolynomial<F> {
)));
}

let mle_ptr: *const DenseMultilinearExtension<F> = Rc::as_ptr(&mle);
let mle_ptr: *const DenseMultilinearExtension<F> = Arc::as_ptr(&mle);

// check if this mle already exists in the virtual polynomial
let mle_index = match self.raw_pointers_lookup_table.get(&mle_ptr) {
Expand Down Expand Up @@ -350,11 +350,11 @@ pub fn eq_eval<F: PrimeField>(x: &[F], y: &[F]) -> Result<F, ArithErrors> {
/// 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> {
) -> Result<Arc<DenseMultilinearExtension<F>>, ArithErrors> {
let evals = build_eq_x_r_vec(r)?;
let mle = DenseMultilinearExtension::from_evaluations_vec(r.len(), evals);

Ok(Rc::new(mle))
Ok(Arc::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.
Expand All @@ -364,8 +364,6 @@ pub fn build_eq_x_r<F: PrimeField>(
/// 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
// for example, with num_vars = 4, x is a binary vector of 4, then
Expand All @@ -380,7 +378,6 @@ pub fn build_eq_x_r_vec<F: PrimeField>(r: &[F]) -> Result<Vec<F>, ArithErrors> {
let mut eval = Vec::new();
build_eq_x_r_helper(r, &mut eval)?;

end_timer!(start);
Ok(eval)
}

Expand Down Expand Up @@ -511,9 +508,7 @@ mod test {
// 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))
fn build_eq_x_r_for_test<F: PrimeField>(r: &[F]) -> Rc<DenseMultilinearExtension<F>> {
let start = start_timer!(|| "zero check naive build eq_x_r");

fn build_eq_x_r_for_test<F: PrimeField>(r: &[F]) -> Arc<DenseMultilinearExtension<F>> {
// we build eq(x,r) from its evaluations
// we want to evaluate eq(x,r) over x \in {0, 1}^num_vars
// for example, with num_vars = 4, x is a binary vector of 4, then
Expand Down Expand Up @@ -545,8 +540,7 @@ mod test {

let mle = DenseMultilinearExtension::from_evaluations_vec(num_var, eval);

let res = Rc::new(mle);
end_timer!(start);
let res = Arc::new(mle);
res
}
}
4 changes: 2 additions & 2 deletions hyperplonk/src/mock.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -238,8 +238,8 @@ mod test {
MultilinearKzgPCS::<Bls12_381>::gen_srs_for_testing(&mut rng, SUPPORTED_SIZE)?;
let nv = MAX_NUM_VARS;

let vanilla_gate = CustomizedGates::vanilla_plonk_gate();
test_mock_circuit_zkp_helper(nv, &vanilla_gate, &pcs_srs)?;
let turboplonk_gate = CustomizedGates::jellyfish_turbo_plonk_gate();
test_mock_circuit_zkp_helper(nv, &turboplonk_gate, &pcs_srs)?;

Ok(())
}
Expand Down
4 changes: 2 additions & 2 deletions hyperplonk/src/selectors.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@ use crate::{build_mle, errors::HyperPlonkErrors};
use ark_ff::PrimeField;
use ark_poly::DenseMultilinearExtension;
use ark_std::log2;
use std::rc::Rc;
use std::sync::Arc;

/// A row of selector of width `#selectors`
#[derive(Debug, Clone, PartialEq, Eq)]
Expand Down Expand Up @@ -67,7 +67,7 @@ impl<F: PrimeField> SelectorRow<F> {
/// output mle(a1,b1,c1, ...), mle(a2,b2,c2, ...), ...
pub fn build_mles(
matrix: &[Self],
) -> Result<Vec<Rc<DenseMultilinearExtension<F>>>, HyperPlonkErrors> {
) -> Result<Vec<Arc<DenseMultilinearExtension<F>>>, HyperPlonkErrors> {
build_mle!(matrix)
}
}
29 changes: 16 additions & 13 deletions hyperplonk/src/snark.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,10 @@ use arithmetic::{evaluate_opt, gen_eval_point, VPAuxInfo};
use ark_ec::PairingEngine;
use ark_poly::DenseMultilinearExtension;
use ark_std::{end_timer, log2, start_timer, One, Zero};
use std::{marker::PhantomData, rc::Rc};
use rayon::iter::IntoParallelRefIterator;
#[cfg(feature = "parallel")]
use rayon::iter::ParallelIterator;
use std::{marker::PhantomData, sync::Arc};
use subroutines::{
pcs::prelude::{Commitment, PolynomialCommitmentScheme},
poly_iop::{
Expand All @@ -28,7 +31,7 @@ where
// we cannot bound PCS::Polynomial with a property trait bound.
PCS: PolynomialCommitmentScheme<
E,
Polynomial = Rc<DenseMultilinearExtension<E::Fr>>,
Polynomial = Arc<DenseMultilinearExtension<E::Fr>>,
Point = Vec<E::Fr>,
Evaluation = E::Fr,
Commitment = Commitment<E>,
Expand Down Expand Up @@ -56,7 +59,7 @@ where
let mut perm_comms = vec![];
let chunk_size = 1 << num_vars;
for i in 0..index.num_witness_columns() {
let perm_oracle = Rc::new(DenseMultilinearExtension::from_evaluations_slice(
let perm_oracle = Arc::new(DenseMultilinearExtension::from_evaluations_slice(
num_vars,
&index.permutation[i * chunk_size..(i + 1) * chunk_size],
));
Expand All @@ -66,14 +69,14 @@ where
}

// build selector oracles and commit to it
let selector_oracles: Vec<Rc<DenseMultilinearExtension<E::Fr>>> = index
let selector_oracles: Vec<Arc<DenseMultilinearExtension<E::Fr>>> = index
.selectors
.iter()
.map(|s| Rc::new(DenseMultilinearExtension::from(s)))
.map(|s| Arc::new(DenseMultilinearExtension::from(s)))
.collect();

let selector_commitments = selector_oracles
.iter()
.par_iter()
.map(|poly| PCS::commit(&pcs_prover_param, poly))
.collect::<Result<Vec<_>, _>>()?;

Expand Down Expand Up @@ -168,13 +171,13 @@ where
// =======================================================================
let step = start_timer!(|| "commit witnesses");

let witness_polys: Vec<Rc<DenseMultilinearExtension<E::Fr>>> = witnesses
let witness_polys: Vec<Arc<DenseMultilinearExtension<E::Fr>>> = witnesses
.iter()
.map(|w| Rc::new(DenseMultilinearExtension::from(w)))
.map(|w| Arc::new(DenseMultilinearExtension::from(w)))
.collect();

let witness_commits = witness_polys
.iter()
.par_iter()
.map(|x| PCS::commit(&pk.pcs_param, x).unwrap())
.collect::<Vec<_>>();
for w_com in witness_commits.iter() {
Expand Down Expand Up @@ -663,15 +666,15 @@ mod tests {
let w1 = WitnessColumn(vec![
E::Fr::zero(),
E::Fr::one(),
E::Fr::from(2u64),
E::Fr::from(3u64),
E::Fr::from(2u128),
E::Fr::from(3u128),
]);
// w2 := [0^5, 1^5, 2^5, 3^5]
let w2 = WitnessColumn(vec![
E::Fr::zero(),
E::Fr::one(),
E::Fr::from(32u64),
E::Fr::from(243u64),
E::Fr::from(32u128),
E::Fr::from(243u128),
]);
// public input = w1
let pi = w1.clone();
Expand Down
6 changes: 3 additions & 3 deletions hyperplonk/src/structs.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@ use ark_ec::PairingEngine;
use ark_ff::PrimeField;
use ark_poly::DenseMultilinearExtension;
use ark_std::log2;
use std::rc::Rc;
use std::sync::Arc;
use subroutines::{
pcs::PolynomialCommitmentScheme,
poly_iop::prelude::{PermutationCheck, ZeroCheck},
Expand Down Expand Up @@ -126,9 +126,9 @@ pub struct HyperPlonkProvingKey<E: PairingEngine, PCS: PolynomialCommitmentSchem
/// Hyperplonk instance parameters
pub params: HyperPlonkParams,
/// The preprocessed permutation polynomials
pub permutation_oracles: Vec<Rc<DenseMultilinearExtension<E::Fr>>>,
pub permutation_oracles: Vec<Arc<DenseMultilinearExtension<E::Fr>>>,
/// The preprocessed selector polynomials
pub selector_oracles: Vec<Rc<DenseMultilinearExtension<E::Fr>>>,
pub selector_oracles: Vec<Arc<DenseMultilinearExtension<E::Fr>>>,
/// Commitments to the preprocessed selector polynomials
pub selector_commitments: Vec<PCS::Commitment>,
/// Commitments to the preprocessed permutation polynomials
Expand Down
Loading