From a8fff1b8ff5d0f6452a6b479227241aee388c88c Mon Sep 17 00:00:00 2001 From: Alon Haramati Date: Tue, 17 Sep 2024 15:09:08 +0300 Subject: [PATCH] Fall back to CPU in small constraint eval. --- crates/prover/Cargo.toml | 2 +- .../src/constraint_framework/component.rs | 32 ++++++- .../src/constraint_framework/cpu_domain.rs | 91 +++++++++++++++++++ crates/prover/src/constraint_framework/mod.rs | 1 + crates/prover/src/core/backend/simd/circle.rs | 4 +- crates/prover/src/core/backend/simd/column.rs | 10 ++ .../prover/src/examples/wide_fibonacci/mod.rs | 89 +++++++++--------- 7 files changed, 181 insertions(+), 48 deletions(-) create mode 100644 crates/prover/src/constraint_framework/cpu_domain.rs diff --git a/crates/prover/Cargo.toml b/crates/prover/Cargo.toml index 62c516253..04594cfba 100644 --- a/crates/prover/Cargo.toml +++ b/crates/prover/Cargo.toml @@ -38,7 +38,7 @@ wasm-bindgen-test = "0.3.43" features = ["html_reports"] version = "0.5.1" -# Default features cause compile error: +# Default features cause compile error: # "Rayon cannot be used when targeting wasi32. Try disabling default features." [target.'cfg(target_arch = "wasm32")'.dev-dependencies.criterion] default-features = false diff --git a/crates/prover/src/constraint_framework/component.rs b/crates/prover/src/constraint_framework/component.rs index c0d8319fa..49b97ae80 100644 --- a/crates/prover/src/constraint_framework/component.rs +++ b/crates/prover/src/constraint_framework/component.rs @@ -5,6 +5,7 @@ use std::ops::Deref; use itertools::Itertools; use tracing::{span, Level}; +use super::cpu_domain::CpuDomainEvaluator; use super::{EvalAtRow, InfoEvaluator, PointEvaluator, SimdDomainEvaluator}; use crate::core::air::accumulation::{DomainEvaluationAccumulator, PointEvaluationAccumulator}; use crate::core::air::{Component, ComponentProver, Trace}; @@ -16,6 +17,7 @@ use crate::core::circle::CirclePoint; use crate::core::constraints::coset_vanishing; use crate::core::fields::m31::BaseField; use crate::core::fields::qm31::SecureField; +use crate::core::fields::secure_column::SecureColumnByCoords; use crate::core::fields::FieldExpOps; use crate::core::pcs::{TreeSubspan, TreeVec}; use crate::core::poly::circle::{CanonicCoset, CircleEvaluation, PolyOps}; @@ -173,7 +175,35 @@ impl ComponentProver for FrameworkComponent { evaluation_accumulator.columns([(eval_domain.log_size(), self.n_constraints())]); accum.random_coeff_powers.reverse(); - let _span = span!(Level::INFO, "Constraint pointwise eval").entered(); + let _span = span!(Level::INFO, "Constraint point-wise eval").entered(); + + if trace_domain.log_size() < LOG_N_LANES + LOG_N_VERY_PACKED_ELEMS { + // Fall back to CPU if the trace is too small. + let mut col = accum.col.to_cpu(); + + for row in 0..(1 << eval_domain.log_size()) { + let trace_cols = trace.as_cols_ref().map_cols(|c| c.to_cpu()); + let trace_cols = trace_cols.as_cols_ref(); + + // Evaluate constrains at row. + let eval = CpuDomainEvaluator::new( + &trace_cols, + row, + &accum.random_coeff_powers, + trace_domain.log_size(), + eval_domain.log_size(), + ); + let row_res = self.eval.evaluate(eval).row_res; + + // Finalize row. + let denom_inv = denom_inv[row >> trace_domain.log_size()]; + col.set(row, col.at(row) + row_res * denom_inv) + } + let col = SecureColumnByCoords::from_cpu(col); + *accum.col = col; + return; + } + let col = unsafe { VeryPackedSecureColumnByCoords::transform_under_mut(accum.col) }; for vec_row in 0..(1 << (eval_domain.log_size() - LOG_N_LANES - LOG_N_VERY_PACKED_ELEMS)) { diff --git a/crates/prover/src/constraint_framework/cpu_domain.rs b/crates/prover/src/constraint_framework/cpu_domain.rs new file mode 100644 index 000000000..d8aad0b8b --- /dev/null +++ b/crates/prover/src/constraint_framework/cpu_domain.rs @@ -0,0 +1,91 @@ +use std::ops::Mul; + +use num_traits::Zero; + +use super::EvalAtRow; +use crate::core::backend::CpuBackend; +use crate::core::fields::m31::BaseField; +use crate::core::fields::qm31::SecureField; +use crate::core::fields::secure_column::SECURE_EXTENSION_DEGREE; +use crate::core::pcs::TreeVec; +use crate::core::poly::circle::CircleEvaluation; +use crate::core::poly::BitReversedOrder; +use crate::core::utils::offset_bit_reversed_circle_domain_index; + +/// Evaluates constraints at an evaluation domain points. +pub struct CpuDomainEvaluator<'a> { + pub trace_eval: &'a TreeVec>>, + pub column_index_per_interaction: Vec, + pub row: usize, + pub random_coeff_powers: &'a [SecureField], + pub row_res: SecureField, + pub constraint_index: usize, + pub domain_log_size: u32, + pub eval_domain_log_size: u32, +} + +impl<'a> CpuDomainEvaluator<'a> { + #[allow(dead_code)] + pub fn new( + trace_eval: &'a TreeVec>>, + row: usize, + random_coeff_powers: &'a [SecureField], + domain_log_size: u32, + eval_log_size: u32, + ) -> Self { + Self { + trace_eval, + column_index_per_interaction: vec![0; trace_eval.len()], + row, + random_coeff_powers, + row_res: SecureField::zero(), + constraint_index: 0, + domain_log_size, + eval_domain_log_size: eval_log_size, + } + } +} + +impl<'a> EvalAtRow for CpuDomainEvaluator<'a> { + type F = BaseField; + type EF = SecureField; + + // TODO(spapini): Remove all boundary checks. + fn next_interaction_mask( + &mut self, + interaction: usize, + offsets: [isize; N], + ) -> [Self::F; N] { + let col_index = self.column_index_per_interaction[interaction]; + self.column_index_per_interaction[interaction] += 1; + offsets.map(|off| { + // If the offset is 0, we can just return the value directly from this row. + if off == 0 { + let col = &self.trace_eval[interaction][col_index]; + return col[self.row]; + } + // Otherwise, we need to look up the value at the offset. + // Since the domain is bit-reversed circle domain ordered, we need to look up the value + // at the bit-reversed natural order index at an offset. + let row = offset_bit_reversed_circle_domain_index( + self.row, + self.domain_log_size, + self.eval_domain_log_size, + off, + ); + self.trace_eval[interaction][col_index][row] + }) + } + + fn add_constraint(&mut self, constraint: G) + where + Self::EF: Mul, + { + self.row_res += self.random_coeff_powers[self.constraint_index] * constraint; + self.constraint_index += 1; + } + + fn combine_ef(values: [Self::F; SECURE_EXTENSION_DEGREE]) -> Self::EF { + SecureField::from_m31_array(values) + } +} diff --git a/crates/prover/src/constraint_framework/mod.rs b/crates/prover/src/constraint_framework/mod.rs index 87069d344..7eea9eb73 100644 --- a/crates/prover/src/constraint_framework/mod.rs +++ b/crates/prover/src/constraint_framework/mod.rs @@ -2,6 +2,7 @@ mod assert; mod component; pub mod constant_columns; +mod cpu_domain; mod info; pub mod logup; mod point; diff --git a/crates/prover/src/core/backend/simd/circle.rs b/crates/prover/src/core/backend/simd/circle.rs index 7ad0ab149..215be84b1 100644 --- a/crates/prover/src/core/backend/simd/circle.rs +++ b/crates/prover/src/core/backend/simd/circle.rs @@ -1,7 +1,7 @@ use std::iter::zip; use std::mem::transmute; -use bytemuck::{cast_slice, Zeroable}; +use bytemuck::Zeroable; use num_traits::One; use super::fft::{ifft, rfft, CACHED_FFT_LOG_SIZE, MIN_FFT_LOG_SIZE}; @@ -334,7 +334,7 @@ fn slow_eval_at_point( // Swap content of a,c. a.swap_with_slice(&mut c[0..n0]); } - fold(cast_slice::<_, BaseField>(&poly.coeffs.data), &mappings) + fold(poly.coeffs.as_slice(), &mappings) } #[cfg(test)] diff --git a/crates/prover/src/core/backend/simd/column.rs b/crates/prover/src/core/backend/simd/column.rs index 64869405a..b5ace46fd 100644 --- a/crates/prover/src/core/backend/simd/column.rs +++ b/crates/prover/src/core/backend/simd/column.rs @@ -58,6 +58,10 @@ impl BaseColumn { res } + pub fn from_cpu(values: Vec) -> Self { + values.into_iter().collect() + } + /// Returns a vector of `BaseColumnMutSlice`s, each mutably owning /// `chunk_size` `PackedBaseField`s (i.e, `chuck_size` * `N_LANES` elements). pub fn chunks_mut(&mut self, chunk_size: usize) -> Vec> { @@ -400,6 +404,12 @@ impl SecureColumnByCoords { .map(|(a, b, c, d)| SecureColumnByCoordsMutSlice([a, b, c, d])) .collect_vec() } + + pub fn from_cpu(cpu: SecureColumnByCoords) -> Self { + Self { + columns: cpu.columns.map(BaseColumn::from_cpu), + } + } } impl FromIterator for SecureColumnByCoords { diff --git a/crates/prover/src/examples/wide_fibonacci/mod.rs b/crates/prover/src/examples/wide_fibonacci/mod.rs index 7b0e9b766..a5acde297 100644 --- a/crates/prover/src/examples/wide_fibonacci/mod.rs +++ b/crates/prover/src/examples/wide_fibonacci/mod.rs @@ -150,59 +150,60 @@ mod tests { } #[test_log::test] - fn test_wide_fib_prove() { - const LOG_N_INSTANCES: u32 = 6; - let config = PcsConfig::default(); - // Precompute twiddles. - let twiddles = SimdBackend::precompute_twiddles( - CanonicCoset::new(LOG_N_INSTANCES + 1 + config.fri_config.log_blowup_factor) - .circle_domain() - .half_coset, - ); - - // Setup protocol. - let prover_channel = &mut Blake2sChannel::default(); - let commitment_scheme = - &mut CommitmentSchemeProver::::new( - config, &twiddles, + fn test_wide_fib_prove_with_blake() { + for log_n_instances in 4..=6 { + let config = PcsConfig::default(); + // Precompute twiddles. + let twiddles = SimdBackend::precompute_twiddles( + CanonicCoset::new(log_n_instances + 1 + config.fri_config.log_blowup_factor) + .circle_domain() + .half_coset, ); - // Trace. - let trace = generate_test_trace(LOG_N_INSTANCES); - let mut tree_builder = commitment_scheme.tree_builder(); - tree_builder.extend_evals(trace); - tree_builder.commit(prover_channel); - - // Prove constraints. - let component = WideFibonacciComponent::new( - &mut TraceLocationAllocator::default(), - WideFibonacciEval:: { - log_n_rows: LOG_N_INSTANCES, - }, - ); - - let proof = prove::( - &[&component], - prover_channel, - commitment_scheme, - ) - .unwrap(); - - // Verify. - let verifier_channel = &mut Blake2sChannel::default(); - let commitment_scheme = &mut CommitmentSchemeVerifier::::new(config); + // Setup protocol. + let prover_channel = &mut Blake2sChannel::default(); + let commitment_scheme = + &mut CommitmentSchemeProver::::new( + config, &twiddles, + ); + + // Trace. + let trace = generate_test_trace(log_n_instances); + let mut tree_builder = commitment_scheme.tree_builder(); + tree_builder.extend_evals(trace); + tree_builder.commit(prover_channel); + + // Prove constraints. + let component = WideFibonacciComponent::new( + &mut TraceLocationAllocator::default(), + WideFibonacciEval:: { + log_n_rows: log_n_instances, + }, + ); - // Retrieve the expected column sizes in each commitment interaction, from the AIR. - let sizes = component.trace_log_degree_bounds(); - commitment_scheme.commit(proof.commitments[0], &sizes[0], verifier_channel); - verify(&[&component], verifier_channel, commitment_scheme, proof).unwrap(); + let proof = prove::( + &[&component], + prover_channel, + commitment_scheme, + ) + .unwrap(); + + // Verify. + let verifier_channel = &mut Blake2sChannel::default(); + let commitment_scheme = + &mut CommitmentSchemeVerifier::::new(config); + + // Retrieve the expected column sizes in each commitment interaction, from the AIR. + let sizes = component.trace_log_degree_bounds(); + commitment_scheme.commit(proof.commitments[0], &sizes[0], verifier_channel); + verify(&[&component], verifier_channel, commitment_scheme, proof).unwrap(); + } } #[test] #[cfg(not(target_arch = "wasm32"))] fn test_wide_fib_prove_with_poseidon() { const LOG_N_INSTANCES: u32 = 6; - let config = PcsConfig::default(); // Precompute twiddles. let twiddles = SimdBackend::precompute_twiddles(