Skip to content

Commit

Permalink
Rework batches in ppsnark (#305)
Browse files Browse the repository at this point in the history 
* refactor: Streamline matrix multiply & some zip_with instances

- Reworked vector and matrix multiplication operations for improved memory use in `sparse.rs`.
- Modified `batch_eval_prove` function in `snark.rs` for more cleaner assertion checks related to polynomial sizes.
- Enhanced `InnerSumcheckInstance` within `setup` function for parallel computation in `ppsnark.rs`.

* refactor: Streamline transcript absorption in Spartan batched code

- Optimized the transcript insertion process in `batched.rs` by replacing loop-absorbing RelaxedR1CSInstances with a single absorb statement

* refactor: Refactor 'claims_outer' and related functions in BatchedRelaxedR1CSSNARK

- Refactored the structure of `claims_outer` in `BatchedRelaxedR1CSSNARK` from a tuple of vectors to a vector of tuples, streamlining its usage and unpacking.
- Adjusted the `prove` and `verify` functions' code in accordance with the new `claims_outer` implementation.
- Enhanced the extraction and manipulation of `ABCE_evals` in the `verify` function resulting in simpler and more efficient usage.

* fix: simplify inner final claim

* refactor: Optimize Spartan batch evaluation and polynomial computation

- Refined tau polynomial definitions using `PowPolynomial::squares` function to improve efficiency.
- uses in verify the same efficient method as exists in proving.

* fix: streamline usage of PowPolynomial

* refactor: Enhance efficiency of generating power polynomials

- Introduced `#[must_use]` annotation to the `squares` function in `power.rs` to ensure the returned result is handled appropriately.
- avoid re-computing powers of tau over and over in batched_ppsnark

* fix: one forgotten instance of the poly_X refactor

* refactor: slightly simplify padding

* fix: remove misleading comment
  • Loading branch information
@huitseeker
huitseeker authored Feb 12, 2024
1 parent efaf628 commit f1c6c19
Show file tree
Hide file tree
Showing 10 changed files with 116 additions and 148 deletions.
2 changes: 1 addition & 1 deletion src/provider/util/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -68,7 +68,7 @@ pub mod test_utils {
.collect::<Vec<_>>();

// Calculation evaluation of point over polynomial.
let eval = MultilinearPolynomial::evaluate_with(poly.evaluations(), &point);
let eval = poly.evaluate(&point);

(poly, point, eval)
}
Expand Down
8 changes: 7 additions & 1 deletion src/r1cs/sparse.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -135,6 +135,12 @@ impl<F: PrimeField> SparseMatrix<F> {
name = "SparseMatrix::multiply_vec_unchecked"
)]
pub fn multiply_vec_unchecked(&self, vector: &[F]) -> Vec<F> {
let mut sink: Vec<F> = Vec::with_capacity(self.indptr.len() - 1);
self.multiply_vec_into_unchecked(vector, &mut sink);
sink
}

pub fn multiply_vec_into_unchecked(&self, vector: &[F], sink: &mut Vec<F>) {
self
.indptr
.par_windows(2)
Expand All @@ -144,7 +150,7 @@ impl<F: PrimeField> SparseMatrix<F> {
.map(|(val, col_idx)| *val * vector[*col_idx])
.sum()
})
.collect()
.collect_into_vec(sink);
}

/// Multiply by a witness representing a dense vector; uses rayon to parallelize.
Expand Down
93 changes: 32 additions & 61 deletions src/spartan/batched.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ use serde::{Deserialize, Serialize};
use itertools::Itertools;
use once_cell::sync::OnceCell;
use rayon::prelude::*;
use std::sync::Arc;
use std::{iter, sync::Arc};

use super::{
compute_eval_table_sparse,
Expand Down Expand Up @@ -45,7 +45,7 @@ use crate::{
pub struct BatchedRelaxedR1CSSNARK<E: Engine, EE: EvaluationEngineTrait<E>> {
sc_proof_outer: SumcheckProof<E>,
// Claims ([Azᵢ(τᵢ)], [Bzᵢ(τᵢ)], [Czᵢ(τᵢ)])
claims_outer: (Vec<E::Scalar>, Vec<E::Scalar>, Vec<E::Scalar>),
claims_outer: Vec<(E::Scalar, E::Scalar, E::Scalar)>,
// [Eᵢ(r_x)]
evals_E: Vec<E::Scalar>,
sc_proof_inner: SumcheckProof<E>,
Expand Down Expand Up @@ -134,17 +134,7 @@ impl<E: Engine, EE: EvaluationEngineTrait<E>> BatchedRelaxedR1CSSNARKTrait<E>
) -> Result<Self, NovaError> {
let num_instances = U.len();
// Pad shapes and ensure their sizes are correct
let S = S
.iter()
.map(|s| {
let s = s.pad();
if s.is_regular_shape() {
Ok(s)
} else {
Err(NovaError::InternalError)
}
})
.collect::<Result<Vec<_>, _>>()?;
let S = S.iter().map(|s| s.pad()).collect::<Vec<_>>();

// Pad (W,E) for each instance
let W = zip_with!(iter, (W, S), |w, s| w.pad(s)).collect::<Vec<RelaxedR1CSWitness<E>>>();
Expand All @@ -156,9 +146,7 @@ impl<E: Engine, EE: EvaluationEngineTrait<E>> BatchedRelaxedR1CSSNARKTrait<E>
let num_instances_field = E::Scalar::from(num_instances as u64);
transcript.absorb(b"n", &num_instances_field);
}
for u in U.iter() {
transcript.absorb(b"U", u);
}
transcript.absorb(b"U", &U);

let (polys_W, polys_E): (Vec<_>, Vec<_>) = W.into_iter().map(|w| (w.W, w.E)).unzip();

Expand Down Expand Up @@ -371,12 +359,9 @@ impl<E: Engine, EE: EvaluationEngineTrait<E>> BatchedRelaxedR1CSSNARKTrait<E>
&batched_u.e,
)?;

let (evals_Az, evals_Bz, evals_Cz): (Vec<_>, Vec<_>, Vec<_>) =
evals_Az_Bz_Cz.into_iter().multiunzip();

Ok(Self {
sc_proof_outer,
claims_outer: (evals_Az, evals_Bz, evals_Cz),
claims_outer: evals_Az_Bz_Cz,
evals_E,
sc_proof_inner,
evals_W,
Expand All @@ -395,9 +380,7 @@ impl<E: Engine, EE: EvaluationEngineTrait<E>> BatchedRelaxedR1CSSNARKTrait<E>
let num_instances_field = E::Scalar::from(num_instances as u64);
transcript.absorb(b"n", &num_instances_field);
}
for u in U.iter() {
transcript.absorb(b"U", u);
}
transcript.absorb(b"U", &U);

let num_instances = U.len();

Expand All @@ -410,14 +393,14 @@ impl<E: Engine, EE: EvaluationEngineTrait<E>> BatchedRelaxedR1CSSNARKTrait<E>
let num_rounds_y_max = *num_rounds_y.iter().max().unwrap();

// Define τ polynomials of the appropriate size for each instance
let polys_tau = {
let tau = transcript.squeeze(b"t")?;
let tau = transcript.squeeze(b"t")?;
let all_taus = PowPolynomial::squares(&tau, num_rounds_x_max);

num_rounds_x
.iter()
.map(|&num_rounds| PowPolynomial::new(&tau, num_rounds))
.collect::<Vec<_>>()
};
let polys_tau = num_rounds_x
.iter()
.map(|&num_rounds_x| PowPolynomial::evals_with_powers(&all_taus, num_rounds_x))
.map(MultilinearPolynomial::new)
.collect::<Vec<_>>();

// Sample challenge for random linear-combination of outer claims
let outer_r = transcript.squeeze(b"out_r")?;
Expand All @@ -441,20 +424,10 @@ impl<E: Engine, EE: EvaluationEngineTrait<E>> BatchedRelaxedR1CSSNARKTrait<E>

// Extract evaluations into a vector [(Azᵢ, Bzᵢ, Czᵢ, Eᵢ)]
// TODO: This is a multizip, simplify
let ABCE_evals = zip_with!(
iter,
(
self.evals_E,
self.claims_outer.0,
self.claims_outer.1,
self.claims_outer.2
),
|eval_E, claim_Az, claim_Bz, claim_Cz| (*claim_Az, *claim_Bz, *claim_Cz, *eval_E)
)
.collect::<Vec<_>>();
let ABCE_evals = || self.claims_outer.iter().zip_eq(self.evals_E.iter());

// Add evaluations of Az, Bz, Cz, E to transcript
for (claim_Az, claim_Bz, claim_Cz, eval_E) in ABCE_evals.iter() {
for ((claim_Az, claim_Bz, claim_Cz), eval_E) in ABCE_evals() {
transcript.absorb(
b"claims_outer",
&[*claim_Az, *claim_Bz, *claim_Cz, *eval_E].as_slice(),
Expand All @@ -467,15 +440,10 @@ impl<E: Engine, EE: EvaluationEngineTrait<E>> BatchedRelaxedR1CSSNARKTrait<E>

// Compute expected claim for all instances ∑ᵢ rⁱ⋅τ(rₓ)⋅(Azᵢ⋅Bzᵢ − uᵢ⋅Czᵢ − Eᵢ)
let claim_outer_final_expected = zip_with!(
(
ABCE_evals.iter().copied(),
U.iter(),
evals_tau,
outer_r_powers.iter()
),
(ABCE_evals(), U.iter(), evals_tau, outer_r_powers.iter()),
|ABCE_eval, u, eval_tau, r| {
let (claim_Az, claim_Bz, claim_Cz, eval_E) = ABCE_eval;
*r * eval_tau * (claim_Az * claim_Bz - u.u * claim_Cz - eval_E)
let ((claim_Az, claim_Bz, claim_Cz), eval_E) = ABCE_eval;
*r * eval_tau * (*claim_Az * claim_Bz - u.u * claim_Cz - eval_E)
}
)
.sum::<E::Scalar>();
Expand All @@ -491,10 +459,11 @@ impl<E: Engine, EE: EvaluationEngineTrait<E>> BatchedRelaxedR1CSSNARKTrait<E>

// Compute inner claims Mzᵢ = (Azᵢ + r⋅Bzᵢ + r²⋅Czᵢ),
// which are batched by Sumcheck into one claim: ∑ᵢ r³ⁱ⋅Mzᵢ
let claims_inner = ABCE_evals
.into_iter()
.map(|(claim_Az, claim_Bz, claim_Cz, _)| {
claim_Az + inner_r * claim_Bz + inner_r_square * claim_Cz
let claims_inner = self
.claims_outer
.iter()
.map(|(claim_Az, claim_Bz, claim_Cz)| {
*claim_Az + inner_r * claim_Bz + inner_r_square * claim_Cz
})
.collect::<Vec<_>>();

Expand All @@ -515,15 +484,17 @@ impl<E: Engine, EE: EvaluationEngineTrait<E>> BatchedRelaxedR1CSSNARKTrait<E>
let evals_Z = zip_with!(iter, (self.evals_W, U, r_y), |eval_W, U, r_y| {
let eval_X = {
// constant term
let mut poly_X = vec![(0, U.u)];
//remaining inputs
poly_X.extend(
U.X
let poly_X = iter::once((0, U.u))
.chain(
//remaining inputs
U.X
.iter()
.enumerate()
.map(|(i, x_i)| (i + 1, *x_i))
.collect::<Vec<(usize, E::Scalar)>>(),
);
// filter_map uses the sparsity of the polynomial, if irrelevant
// we should replace by UniPoly
.filter_map(|(i, x_i)| (!x_i.is_zero_vartime()).then_some((i + 1, *x_i))),
)
.collect();
SparsePolynomial::new(r_y.len() - 1, poly_X).evaluate(&r_y[1..])
};
(E::Scalar::ONE - r_y[0]) * eval_W + r_y[0] * eval_X
Expand Down
78 changes: 32 additions & 46 deletions src/spartan/batched_ppsnark.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -177,29 +177,20 @@ impl<E: Engine, EE: EvaluationEngineTrait<E>> BatchedRelaxedR1CSSNARKTrait<E>
W: &[RelaxedR1CSWitness<E>],
) -> Result<Self, NovaError> {
// Pad shapes so that num_vars = num_cons = Nᵢ and check the sizes are correct
let S = S
.par_iter()
.map(|s| {
let s = s.pad();
if s.is_regular_shape() {
Ok(s)
} else {
Err(NovaError::InternalError)
}
})
.collect::<Result<Vec<_>, _>>()?;
let S = S.par_iter().map(|s| s.pad()).collect::<Vec<_>>();

// N[i] = max(|Aᵢ|+|Bᵢ|+|Cᵢ|, 2*num_varsᵢ, num_consᵢ)
let N = pk.S_repr.iter().map(|s| s.N).collect::<Vec<_>>();
assert!(N.iter().all(|&Ni| Ni.is_power_of_two()));
let Nis = pk.S_repr.iter().map(|s| s.N).collect::<Vec<_>>();
assert!(Nis.iter().all(|&Ni| Ni.is_power_of_two()));
let N_max = *Nis.iter().max().unwrap();

let num_instances = U.len();

// Pad [(Wᵢ,Eᵢ)] to the next power of 2 (not to Ni)
let W = zip_with!(par_iter, (W, S), |w, s| w.pad(s)).collect::<Vec<RelaxedR1CSWitness<E>>>();

// number of rounds of sum-check
let num_rounds_sc = N.iter().max().unwrap().log_2();
let num_rounds_sc = N_max.log_2();

// Initialize transcript with vk || [Uᵢ]
let mut transcript = E::TE::new(b"BatchedRelaxedR1CSSNARK");
Expand All @@ -208,12 +199,10 @@ impl<E: Engine, EE: EvaluationEngineTrait<E>> BatchedRelaxedR1CSSNARKTrait<E>
let num_instances_field = E::Scalar::from(num_instances as u64);
transcript.absorb(b"n", &num_instances_field);
}
for u in U.iter() {
transcript.absorb(b"U", u);
}
transcript.absorb(b"U", &U);

// Append public inputs to Wᵢ: Zᵢ = [Wᵢ, uᵢ, Xᵢ]
let polys_Z = zip_with!(par_iter, (W, U, N), |W, U, Ni| {
let polys_Z = zip_with!(par_iter, (W, U, Nis), |W, U, Ni| {
// poly_Z will be resized later, so we preallocate the correct capacity
let mut poly_Z = Vec::with_capacity(*Ni);
poly_Z.extend(W.W.iter().chain([&U.u]).chain(U.X.iter()));
Expand Down Expand Up @@ -249,21 +238,23 @@ impl<E: Engine, EE: EvaluationEngineTrait<E>> BatchedRelaxedR1CSSNARKTrait<E>

// Compute eq(tau) for each instance in log2(Ni) variables
let tau = transcript.squeeze(b"t")?;
let (polys_tau, coords_tau): (Vec<_>, Vec<_>) = N
.iter()
let all_taus = PowPolynomial::squares(&tau, N_max.log_2());

let (polys_tau, coords_tau): (Vec<_>, Vec<_>) = Nis
.par_iter()
.map(|&N_i| {
let log_Ni = N_i.log_2();
let poly = PowPolynomial::new(&tau, log_Ni);
let evals = poly.evals();
let coords = poly.coordinates();
let eqp: EqPolynomial<_> = all_taus[..log_Ni].iter().cloned().collect();
let evals = eqp.evals();
let coords = eqp.r;
(evals, coords)
})
.unzip();

// Pad [Az, Bz, Cz] to Ni
polys_Az_Bz_Cz
.par_iter_mut()
.zip_eq(N.par_iter())
.zip_eq(Nis.par_iter())
.for_each(|(az_bz_cz, &Ni)| {
az_bz_cz
.par_iter_mut()
Expand Down Expand Up @@ -298,7 +289,7 @@ impl<E: Engine, EE: EvaluationEngineTrait<E>> BatchedRelaxedR1CSSNARKTrait<E>
// Pad Zᵢ, E to Nᵢ
let polys_Z = polys_Z
.into_par_iter()
.zip_eq(N.par_iter())
.zip_eq(Nis.par_iter())
.map(|(mut poly_Z, &Ni)| {
poly_Z.resize(Ni, E::Scalar::ZERO);
poly_Z
Expand All @@ -309,7 +300,7 @@ impl<E: Engine, EE: EvaluationEngineTrait<E>> BatchedRelaxedR1CSSNARKTrait<E>
// but it makes it easier to handle the batching at the end.
let polys_E = polys_E
.into_par_iter()
.zip_eq(N.par_iter())
.zip_eq(Nis.par_iter())
.map(|(mut poly_E, &Ni)| {
poly_E.resize(Ni, E::Scalar::ZERO);
poly_E
Expand All @@ -318,7 +309,7 @@ impl<E: Engine, EE: EvaluationEngineTrait<E>> BatchedRelaxedR1CSSNARKTrait<E>

let polys_W = polys_W
.into_par_iter()
.zip_eq(N.par_iter())
.zip_eq(Nis.par_iter())
.map(|(mut poly_W, &Ni)| {
poly_W.resize(Ni, E::Scalar::ZERO);
poly_W
Expand All @@ -330,7 +321,7 @@ impl<E: Engine, EE: EvaluationEngineTrait<E>> BatchedRelaxedR1CSSNARKTrait<E>
// L_col(i) = z(col(i)) for all i in [0..Nᵢ]
let polys_L_row_col = zip_with!(
par_iter,
(S, N, polys_Z, polys_tau),
(S, Nis, polys_Z, polys_tau),
|S, Ni, poly_Z, poly_tau| {
let mut L_row = vec![poly_tau[0]; *Ni]; // we place mem_row[0] since resized row is appended with 0s
let mut L_col = vec![poly_Z[Ni - 1]; *Ni]; // we place mem_col[Ni-1] since resized col is appended with Ni-1
Expand Down Expand Up @@ -495,13 +486,12 @@ impl<E: Engine, EE: EvaluationEngineTrait<E>> BatchedRelaxedR1CSSNARKTrait<E>

// Sample new random variable for eq polynomial
let rho = transcript.squeeze(b"r")?;
let N_max = N.iter().max().unwrap();
let all_rhos = PowPolynomial::squares(&rho, N_max.log_2());

let instances = zip_with!(
(
pk.S_repr.par_iter(),
N.par_iter(),
Nis.par_iter(),
polys_mem_oracles.par_iter(),
mem_aux.into_par_iter()
),
Expand Down Expand Up @@ -765,9 +755,7 @@ impl<E: Engine, EE: EvaluationEngineTrait<E>> BatchedRelaxedR1CSSNARKTrait<E>
let num_instances_field = E::Scalar::from(num_instances as u64);
transcript.absorb(b"n", &num_instances_field);
}
for u in U.iter() {
transcript.absorb(b"U", u);
}
transcript.absorb(b"U", &U);

// Decompress commitments
let comms_Az_Bz_Cz = self
Expand Down Expand Up @@ -907,14 +895,10 @@ impl<E: Engine, EE: EvaluationEngineTrait<E>> BatchedRelaxedR1CSSNARKTrait<E>
let num_rounds_i = rand_sc.len();
let num_vars_log = num_vars.log_2();

let eq_rho = {
let rho_coords = PowPolynomial::new(&rho, num_rounds_i).coordinates();
EqPolynomial::new(rho_coords).evaluate(rand_sc)
};
let eq_rho = PowPolynomial::new(&rho, num_rounds_i).evaluate(rand_sc);

let (eq_tau, eq_masked_tau) = {
let tau_coords = PowPolynomial::new(&tau, num_rounds_i).coordinates();
let eq_tau = EqPolynomial::new(tau_coords);
let eq_tau: EqPolynomial<_> = PowPolynomial::new(&tau, num_rounds_i).into();

let eq_tau_at_rand = eq_tau.evaluate(rand_sc);
let eq_masked_tau = MaskedEqPolynomial::new(&eq_tau, num_vars_log).evaluate(rand_sc);
Expand All @@ -941,13 +925,15 @@ impl<E: Engine, EE: EvaluationEngineTrait<E>> BatchedRelaxedR1CSSNARKTrait<E>

let X = {
// constant term
let mut poly_X = vec![(0, U.u)];
//remaining inputs
poly_X.extend(
(0..U.X.len())
.map(|i| (i + 1, U.X[i]))
.collect::<Vec<(usize, E::Scalar)>>(),
);
let poly_X = std::iter::once((0, U.u))
.chain(
//remaining inputs
(0..U.X.len())
// filter_map uses the sparsity of the polynomial, if irrelevant
// we should replace by UniPoly
.filter_map(|i| (!U.X[i].is_zero_vartime()).then_some((i + 1, U.X[i]))),
)
.collect();
SparsePolynomial::new(num_vars_log, poly_X).evaluate(&rand_sc_unpad[1..])
};

Expand Down
Loading

1 comment on commit f1c6c19

@github-actions
Copy link
Contributor

Choose a reason for hiding this comment

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

Benchmarks

Table of Contents

Overview

This benchmark report shows the Arecibo GPU benchmarks.
NVIDIA L4
Intel(R) Xeon(R) CPU @ 2.20GHz
32 vCPUs
125 GB RAM
Workflow run: https://github.com/lurk-lab/arecibo/actions/runs/7874209603

Benchmark Results

RecursiveSNARK-NIVC-2

ref=efaf628 ref=f1c6c19
Prove-NumCons-6540 53.04 ms (✅ 1.00x) 53.10 ms (✅ 1.00x slower)
Verify-NumCons-6540 33.00 ms (✅ 1.00x) 32.75 ms (✅ 1.01x faster)
Prove-NumCons-1028888 344.76 ms (✅ 1.00x) 342.84 ms (✅ 1.01x faster)
Verify-NumCons-1028888 256.85 ms (✅ 1.00x) 252.76 ms (✅ 1.02x faster)

CompressedSNARK-NIVC-Commitments-2

ref=efaf628 ref=f1c6c19
Prove-NumCons-6540 13.98 s (✅ 1.00x) 13.93 s (✅ 1.00x faster)
Verify-NumCons-6540 78.31 ms (✅ 1.00x) 77.39 ms (✅ 1.01x faster)
Prove-NumCons-1028888 110.33 s (✅ 1.00x) 110.97 s (✅ 1.01x slower)
Verify-NumCons-1028888 772.37 ms (✅ 1.00x) 770.79 ms (✅ 1.00x faster)

Made with criterion-table

Please sign in to comment.