From b1365c155fb76c80eb394451d4f21a9fa475fb9f Mon Sep 17 00:00:00 2001 From: Julian Arnesino Date: Thu, 29 Aug 2024 15:31:17 -0300 Subject: [PATCH] Add new PolyOps trait for interpolating many columns. Add a default implementation for it that parallelizes extend_evals with Rayon. --- crates/prover/src/core/pcs/prover.rs | 19 ++++--- crates/prover/src/core/poly/circle/ops.rs | 28 +++++++++- .../src/core/poly/circle/secure_poly.rs | 8 ++- crates/prover/src/examples/poseidon/mod.rs | 56 ++++++++++++++++++- 4 files changed, 98 insertions(+), 13 deletions(-) diff --git a/crates/prover/src/core/pcs/prover.rs b/crates/prover/src/core/pcs/prover.rs index 7a761eaa9..efb8f4e2a 100644 --- a/crates/prover/src/core/pcs/prover.rs +++ b/crates/prover/src/core/pcs/prover.rs @@ -15,9 +15,9 @@ use super::quotients::{compute_fri_quotients, PointSample}; use super::utils::TreeVec; use super::{PcsConfig, TreeSubspan}; use crate::core::air::Trace; -use crate::core::backend::BackendForChannel; +use crate::core::backend::{BackendForChannel, ColumnOps}; use crate::core::channel::{Channel, MerkleChannel}; -use crate::core::poly::circle::{CircleEvaluation, CirclePoly}; +use crate::core::poly::circle::{CircleEvaluation, CirclePoly, PolyOps}; use crate::core::poly::twiddles::TwiddleTree; use crate::core::vcs::ops::MerkleHasher; use crate::core::vcs::prover::{MerkleDecommitment, MerkleProver}; @@ -168,13 +168,18 @@ impl<'a, 'b, B: BackendForChannel, MC: MerkleChannel> TreeBuilder<'a, 'b, B, pub fn extend_evals( &mut self, columns: ColumnVec>, - ) -> TreeSubspan { + ) -> TreeSubspan + where + CircleEvaluation: Send + Sync, + >::Column: Send + Sync, + ::Twiddles: Send + Sync, + { let span = span!(Level::INFO, "Interpolation for commitment").entered(); let col_start = self.polys.len(); - let polys = columns - .into_iter() - .map(|eval| eval.interpolate_with_twiddles(self.commitment_scheme.twiddles)) - .collect_vec(); + + let polys: Vec> = + B::interpolate_columns(&columns, self.commitment_scheme.twiddles); + span.exit(); self.polys.extend(polys); TreeSubspan { diff --git a/crates/prover/src/core/poly/circle/ops.rs b/crates/prover/src/core/poly/circle/ops.rs index 40b86cb68..318a579fa 100644 --- a/crates/prover/src/core/poly/circle/ops.rs +++ b/crates/prover/src/core/poly/circle/ops.rs @@ -1,11 +1,15 @@ +#[cfg(feature = "parallel")] +use rayon::prelude::*; + use super::{CanonicCoset, CircleDomain, CircleEvaluation, CirclePoly}; -use crate::core::backend::Col; +use crate::core::backend::{Col, ColumnOps}; use crate::core::circle::{CirclePoint, Coset}; use crate::core::fields::m31::BaseField; use crate::core::fields::qm31::SecureField; use crate::core::fields::FieldOps; use crate::core::poly::twiddles::TwiddleTree; use crate::core::poly::BitReversedOrder; +use crate::core::ColumnVec; /// Operations on BaseField polynomials. pub trait PolyOps: FieldOps + Sized { @@ -27,6 +31,28 @@ pub trait PolyOps: FieldOps + Sized { itwiddles: &TwiddleTree, ) -> CirclePoly; + fn interpolate_columns( + columns: &ColumnVec>, + twiddles: &TwiddleTree, + ) -> Vec> + where + CircleEvaluation: Send + Sync, + >::Column: Send + Sync, + ::Twiddles: Send + Sync, + { + #[cfg(feature = "parallel")] + return columns + .into_par_iter() + .map(|eval| Self::interpolate(eval.clone(), twiddles)) + .collect(); + + #[cfg(not(feature = "parallel"))] + return columns + .into_iter() + .map(|eval| Self::interpolate(eval.clone(), twiddles)) + .collect(); + } + /// Evaluates the polynomial at a single point. /// Used by the [`CirclePoly::eval_at_point()`] function. fn eval_at_point(poly: &CirclePoly, point: CirclePoint) -> SecureField; diff --git a/crates/prover/src/core/poly/circle/secure_poly.rs b/crates/prover/src/core/poly/circle/secure_poly.rs index a503bd2c6..b573af8cd 100644 --- a/crates/prover/src/core/poly/circle/secure_poly.rs +++ b/crates/prover/src/core/poly/circle/secure_poly.rs @@ -64,7 +64,7 @@ pub struct SecureEvaluation, EvalOrder> { _eval_order: PhantomData, } -impl, EvalOrder> SecureEvaluation { +impl + PolyOps, EvalOrder> SecureEvaluation { pub fn new(domain: CircleDomain, values: SecureColumnByCoords) -> Self { assert_eq!(domain.size(), values.len()); Self { @@ -76,7 +76,11 @@ impl, EvalOrder> SecureEvaluation { pub fn into_coordinate_evals( self, - ) -> [CircleEvaluation; SECURE_EXTENSION_DEGREE] { + ) -> [CircleEvaluation; SECURE_EXTENSION_DEGREE] + where + CircleEvaluation: Send + Sync, + ::Twiddles: Send + Sync, + { let Self { domain, values, .. } = self; values.columns.map(|c| CircleEvaluation::new(domain, c)) } diff --git a/crates/prover/src/examples/poseidon/mod.rs b/crates/prover/src/examples/poseidon/mod.rs index c94f0ba1d..3ca58f04d 100644 --- a/crates/prover/src/examples/poseidon/mod.rs +++ b/crates/prover/src/examples/poseidon/mod.rs @@ -377,20 +377,25 @@ mod tests { use itertools::Itertools; use num_traits::One; + use tracing::{span, Level}; use crate::constraint_framework::assert_constraints; use crate::constraint_framework::logup::{LogupAtRow, LookupElements}; use crate::core::air::Component; + use crate::core::backend::simd::SimdBackend; + use crate::core::backend::CpuBackend; use crate::core::channel::Blake2sChannel; use crate::core::fields::m31::BaseField; use crate::core::fri::FriConfig; - use crate::core::pcs::{CommitmentSchemeVerifier, PcsConfig, TreeVec}; - use crate::core::poly::circle::CanonicCoset; + use crate::core::pcs::{CommitmentSchemeProver, CommitmentSchemeVerifier, PcsConfig, TreeVec}; + use crate::core::poly::circle::{CanonicCoset, CircleEvaluation, PolyOps}; + use crate::core::poly::BitReversedOrder; use crate::core::prover::verify; use crate::core::vcs::blake2_merkle::Blake2sMerkleChannel; + use crate::core::ColumnVec; use crate::examples::poseidon::{ apply_internal_round_matrix, apply_m4, eval_poseidon_constraints, gen_interaction_trace, - gen_trace, prove_poseidon, PoseidonElements, + gen_trace, prove_poseidon, PoseidonElements, LOG_EXPAND, N_COLUMNS }; use crate::math::matrix::{RowMajorMatrix, SquareMatrix}; @@ -505,4 +510,49 @@ mod tests { verify(&[&component], channel, commitment_scheme, proof).unwrap(); } + + #[test_log::test] + fn test_optimize_interpolation() { + let blowup_factor = 2; + let domain = + CanonicCoset::new(N_COLUMNS.ilog2() + LOG_EXPAND + blowup_factor).circle_domain(); + let config = PcsConfig { + pow_bits: 10, + fri_config: FriConfig::new(5, 1, 64), + }; + + let simd_span = span!(Level::INFO, "Test SIMD interpolation").entered(); + let simd_columns: ColumnVec> = + (0..N_COLUMNS) + .map( |_index| { + CircleEvaluation::::new( + domain, + (0..domain.size()).map(BaseField::from).collect(), + ) + }) + .collect(); + let simd_twiddles = SimdBackend::precompute_twiddles(domain.half_coset); + let mut simd_commitment_scheme_prover: CommitmentSchemeProver< + '_, SimdBackend, Blake2sMerkleChannel, + > = CommitmentSchemeProver::new(config, &simd_twiddles); + simd_commitment_scheme_prover.tree_builder().extend_evals(simd_columns); + simd_span.exit(); + + let cpu_span = span!(Level::INFO, "Test CPU interpolation").entered(); + let cpu: ColumnVec> = + (0..N_COLUMNS) + .map( |_index| { + CircleEvaluation::::new( + domain, + (0..domain.size()).map(BaseField::from).collect(), + ) + }) + .collect(); + let cpu_twiddles = CpuBackend::precompute_twiddles(domain.half_coset); + let mut simd_commitment_scheme_prover: CommitmentSchemeProver< + '_, CpuBackend, Blake2sMerkleChannel + > = CommitmentSchemeProver::new(config, &cpu_twiddles); + simd_commitment_scheme_prover.tree_builder().extend_evals(cpu); + cpu_span.exit(); + } }