diff --git a/crates/prover/src/core/air/air_ext.rs b/crates/prover/src/core/air/air_ext.rs index 577497595..441e82ea7 100644 --- a/crates/prover/src/core/air/air_ext.rs +++ b/crates/prover/src/core/air/air_ext.rs @@ -21,6 +21,14 @@ pub trait AirExt: Air { .unwrap() } + fn n_interaction_phases(&self) -> u32 { + self.components() + .iter() + .map(|component| component.n_interaction_phases()) + .max() + .unwrap() + } + fn trace_commitment_domains(&self) -> Vec { self.column_log_sizes() .iter() diff --git a/crates/prover/src/core/air/mod.rs b/crates/prover/src/core/air/mod.rs index ba43d3ef8..f75ecd666 100644 --- a/crates/prover/src/core/air/mod.rs +++ b/crates/prover/src/core/air/mod.rs @@ -50,6 +50,9 @@ pub trait Component { fn max_constraint_log_degree_bound(&self) -> u32; + /// Returns the number of interaction phases done by the component. + fn n_interaction_phases(&self) -> u32; + /// Returns the degree bounds of each trace column. fn trace_log_degree_bounds(&self) -> Vec; diff --git a/crates/prover/src/core/constraints.rs b/crates/prover/src/core/constraints.rs index f525f50b3..cc7523e2a 100644 --- a/crates/prover/src/core/constraints.rs +++ b/crates/prover/src/core/constraints.rs @@ -61,10 +61,10 @@ pub fn pair_vanishing>( /// Evaluates a vanishing polynomial of the vanish_point at a point. /// Note that this function has a pole on the antipode of the vanish_point. pub fn point_vanishing, EF: ExtensionOf>( - vanish_point: CirclePoint, - p: CirclePoint, + vanish_point: CirclePoint, + p: CirclePoint, ) -> EF { - let h = p.into_ef() - vanish_point; + let h = p - vanish_point.into_ef(); h.y / (EF::one() + h.x) } diff --git a/crates/prover/src/core/prover/mod.rs b/crates/prover/src/core/prover/mod.rs index 078a16810..8a04aa3c9 100644 --- a/crates/prover/src/core/prover/mod.rs +++ b/crates/prover/src/core/prover/mod.rs @@ -80,8 +80,7 @@ pub fn evaluate_and_commit_on_trace>( .map(|poly| poly.interpolate_with_twiddles(twiddles)) }) .collect_vec(); - let n_interaction_traces = interaction_trace_polys.len(); - if n_interaction_traces > 0 { + if !interaction_trace_polys.is_empty() { commitment_scheme.commit(interaction_trace_polys, channel, twiddles); } @@ -119,7 +118,7 @@ pub fn generate_proof>( // TODO(spapini): Change when we support multiple interactions. // First tree - trace. let mut sample_points = TreeVec::new(vec![sample_points.flatten()]); - if commitment_scheme.trees.len() > 2 { + if air.n_interaction_phases() == 2 { // Second tree - interaction trace. sample_points.push(vec![ vec![oods_point]; @@ -210,7 +209,7 @@ pub fn verify( commitment_scheme.commit(proof.commitments[0], air.column_log_sizes(), channel); let interaction_elements = air.interaction_elements(channel); - if proof.commitments.len() > 2 { + if air.n_interaction_phases() == 2 { commitment_scheme.commit( proof.commitments[1], air.column_log_sizes()[..1].to_vec(), @@ -236,7 +235,7 @@ pub fn verify( // TODO(spapini): Change when we support multiple interactions. // First tree - trace. let mut sample_points = TreeVec::new(vec![trace_sample_points.flatten()]); - if proof.commitments.len() > 2 { + if air.n_interaction_phases() == 2 { // Second tree - interaction trace. // TODO(AlonH): Get the number of interaction traces from the air. sample_points.push(vec![vec![oods_point]; 1]); @@ -288,7 +287,7 @@ fn sampled_values_to_mask( ) }); - if sampled_values.len() > 2 { + if air.n_interaction_phases() == 2 { let interaction_values = &mut sampled_values .get(1) .ok_or(InvalidOodsSampleStructure)? @@ -429,6 +428,10 @@ mod tests { self.max_constraint_log_degree_bound } + fn n_interaction_phases(&self) -> u32 { + 1 + } + fn trace_log_degree_bounds(&self) -> Vec { vec![self.log_size] } diff --git a/crates/prover/src/core/utils.rs b/crates/prover/src/core/utils.rs index 46ca0bddb..7b033e6e7 100644 --- a/crates/prover/src/core/utils.rs +++ b/crates/prover/src/core/utils.rs @@ -1,9 +1,10 @@ use std::iter::Peekable; -use num_traits::{One, Zero}; +use num_traits::One; use super::fields::m31::BaseField; use super::fields::qm31::SecureField; +use super::fields::ExtensionOf; pub trait IteratorMutExt<'a, T: 'a>: Iterator { fn assign(self, other: impl IntoIterator) @@ -89,14 +90,14 @@ pub fn generate_secure_powers(felt: SecureField, n_powers: usize) -> Vec>( + values: &[F], alpha: BaseField, z: BaseField, -) -> BaseField { +) -> F { let res = values .iter() - .fold(BaseField::zero(), |acc, &value| acc * alpha + value); + .fold(F::zero(), |acc, &value| acc * alpha + value); res - z } diff --git a/crates/prover/src/examples/fibonacci/component.rs b/crates/prover/src/examples/fibonacci/component.rs index 243669d99..74ce51235 100644 --- a/crates/prover/src/examples/fibonacci/component.rs +++ b/crates/prover/src/examples/fibonacci/component.rs @@ -80,6 +80,10 @@ impl Component for FibonacciComponent { self.log_size + 1 } + fn n_interaction_phases(&self) -> u32 { + 1 + } + fn trace_log_degree_bounds(&self) -> Vec { vec![self.log_size] } diff --git a/crates/prover/src/examples/wide_fibonacci/component.rs b/crates/prover/src/examples/wide_fibonacci/component.rs index b8d6ed19b..0ba4c32db 100644 --- a/crates/prover/src/examples/wide_fibonacci/component.rs +++ b/crates/prover/src/examples/wide_fibonacci/component.rs @@ -5,20 +5,21 @@ use crate::core::air::mask::fixed_mask_points; use crate::core::air::{Air, Component, ComponentTraceWriter}; use crate::core::backend::CpuBackend; use crate::core::circle::CirclePoint; -use crate::core::constraints::coset_vanishing; +use crate::core::constraints::{coset_vanishing, point_vanishing}; use crate::core::fields::m31::BaseField; use crate::core::fields::qm31::SecureField; use crate::core::fields::FieldExpOps; use crate::core::poly::circle::{CanonicCoset, CircleEvaluation}; use crate::core::poly::BitReversedOrder; +use crate::core::utils::shifted_secure_combination; use crate::core::{ColumnVec, InteractionElements}; use crate::examples::wide_fibonacci::trace_gen::write_lookup_column; pub const LOG_N_COLUMNS: usize = 8; pub const N_COLUMNS: usize = 1 << LOG_N_COLUMNS; -const ALPHA_ID: &str = "wide_fibonacci_alpha"; -const Z_ID: &str = "wide_fibonacci_z"; +pub const ALPHA_ID: &str = "wide_fibonacci_alpha"; +pub const Z_ID: &str = "wide_fibonacci_z"; /// Component that computes 2^`self.log_n_instances` instances of fibonacci sequences of size /// 2^`self.log_fibonacci_size`. The numbers are computes over [N_COLUMNS] trace columns. The @@ -57,13 +58,17 @@ impl Air for WideFibAir { impl Component for WideFibComponent { fn n_constraints(&self) -> usize { - self.n_columns() - 2 + self.n_columns() - 1 } fn max_constraint_log_degree_bound(&self) -> u32 { self.log_column_size() + 1 } + fn n_interaction_phases(&self) -> u32 { + 2 + } + fn trace_log_degree_bounds(&self) -> Vec { vec![self.log_column_size(); self.n_columns()] } @@ -84,9 +89,20 @@ impl Component for WideFibComponent { point: CirclePoint, mask: &ColumnVec>, evaluation_accumulator: &mut PointEvaluationAccumulator, - _interaction_elements: &InteractionElements, + interaction_elements: &InteractionElements, ) { let constraint_zero_domain = CanonicCoset::new(self.log_column_size()).coset; + let (alpha, z) = (interaction_elements[ALPHA_ID], interaction_elements[Z_ID]); + let lookup_numerator = (mask[self.n_columns()][0] + * shifted_secure_combination( + &[mask[self.n_columns() - 2][0], mask[self.n_columns() - 1][0]], + alpha, + z, + )) + - shifted_secure_combination(&[mask[0][0], mask[1][0]], alpha, z); + let lookup_denom = point_vanishing(constraint_zero_domain.at(0), point); + evaluation_accumulator.accumulate(lookup_numerator / lookup_denom); + let denom = coset_vanishing(constraint_zero_domain, point); let denom_inverse = denom.inverse(); for i in 0..self.n_columns() - 2 { diff --git a/crates/prover/src/examples/wide_fibonacci/constraint_eval.rs b/crates/prover/src/examples/wide_fibonacci/constraint_eval.rs index 18fc41b71..9f309d844 100644 --- a/crates/prover/src/examples/wide_fibonacci/constraint_eval.rs +++ b/crates/prover/src/examples/wide_fibonacci/constraint_eval.rs @@ -3,7 +3,7 @@ use std::collections::BTreeMap; use itertools::{zip_eq, Itertools}; use num_traits::Zero; -use super::component::{Input, WideFibAir, WideFibComponent}; +use super::component::{Input, WideFibAir, WideFibComponent, ALPHA_ID, Z_ID}; use super::trace_gen::write_trace_row; use crate::core::air::accumulation::DomainEvaluationAccumulator; use crate::core::air::{ @@ -12,13 +12,13 @@ use crate::core::air::{ }; use crate::core::backend::CpuBackend; use crate::core::channel::{Blake2sChannel, Channel}; -use crate::core::constraints::coset_vanishing; +use crate::core::constraints::{coset_vanishing, point_vanishing}; use crate::core::fields::m31::BaseField; use crate::core::fields::qm31::SecureField; use crate::core::fields::FieldExpOps; use crate::core::poly::circle::{CanonicCoset, CircleEvaluation}; use crate::core::poly::BitReversedOrder; -use crate::core::utils::bit_reverse; +use crate::core::utils::{bit_reverse, shifted_secure_combination}; use crate::core::{ColumnVec, ComponentVec, InteractionElements}; use crate::examples::wide_fibonacci::component::LOG_N_COLUMNS; @@ -59,24 +59,30 @@ impl ComponentProver for WideFibComponent { &self, trace: &ComponentTrace<'_, CpuBackend>, evaluation_accumulator: &mut DomainEvaluationAccumulator, - _interaction_elements: &InteractionElements, + interaction_elements: &InteractionElements, ) { let max_constraint_degree = self.max_constraint_log_degree_bound(); let trace_eval_domain = CanonicCoset::new(max_constraint_degree).circle_domain(); let trace_evals = &trace.evals; let zero_domain = CanonicCoset::new(self.log_column_size()).coset; let mut denoms = vec![]; + let mut lookup_denoms = vec![]; for point in trace_eval_domain.iter() { denoms.push(coset_vanishing(zero_domain, point)); + lookup_denoms.push(point_vanishing(zero_domain.at(0), point)); } bit_reverse(&mut denoms); let mut denom_inverses = vec![BaseField::zero(); 1 << (max_constraint_degree)]; BaseField::batch_inverse(&denoms, &mut denom_inverses); + bit_reverse(&mut lookup_denoms); + let mut lookup_denom_inverses = vec![BaseField::zero(); 1 << (max_constraint_degree)]; + BaseField::batch_inverse(&lookup_denoms, &mut lookup_denom_inverses); let mut numerators = vec![SecureField::zero(); 1 << (max_constraint_degree)]; + let mut lookup_numerators = vec![SecureField::zero(); 1 << (max_constraint_degree)]; let [mut accum] = evaluation_accumulator.columns([(max_constraint_degree, self.n_constraints())]); + let (alpha, z) = (interaction_elements[ALPHA_ID], interaction_elements[Z_ID]); - #[allow(clippy::needless_range_loop)] for i in 0..trace_eval_domain.size() { // Step constraints. for j in 0..self.n_columns() - 2 { @@ -84,9 +90,33 @@ impl ComponentProver for WideFibComponent { * (trace_evals[0][j][i].square() + trace_evals[0][j + 1][i].square() - trace_evals[0][j + 2][i]); } + + // Lookup constraints. + lookup_numerators[i] = accum.random_coeff_powers[self.n_columns() - 2] + * ((trace_evals[1][0][i] + * shifted_secure_combination( + &[ + trace_evals[0][self.n_columns() - 2][i], + trace_evals[0][self.n_columns() - 1][i], + ], + alpha, + z, + )) + - shifted_secure_combination( + &[trace_evals[0][0][i], trace_evals[0][1][i]], + alpha, + z, + )); + } + for (i, (num, denom_inverse)) in numerators.iter().zip(denom_inverses.iter()).enumerate() { + accum.accumulate(i, *num * *denom_inverse); } - for (i, (num, denom)) in numerators.iter().zip(denom_inverses.iter()).enumerate() { - accum.accumulate(i, *num * *denom); + for (i, (num, denom_inverse)) in lookup_numerators + .iter() + .zip(lookup_denom_inverses.iter()) + .enumerate() + { + accum.accumulate(i, *num * *denom_inverse); } } } diff --git a/crates/prover/src/examples/wide_fibonacci/simd.rs b/crates/prover/src/examples/wide_fibonacci/simd.rs index 114a7fee2..c919f4f05 100644 --- a/crates/prover/src/examples/wide_fibonacci/simd.rs +++ b/crates/prover/src/examples/wide_fibonacci/simd.rs @@ -2,25 +2,69 @@ use itertools::Itertools; use num_traits::{One, Zero}; use tracing::{span, Level}; -use super::component::{WideFibAir, WideFibComponent}; -use crate::core::air::accumulation::DomainEvaluationAccumulator; +use super::component::LOG_N_COLUMNS; +use crate::core::air::accumulation::{DomainEvaluationAccumulator, PointEvaluationAccumulator}; +use crate::core::air::mask::fixed_mask_points; use crate::core::air::{ - AirProver, AirTraceWriter, Component, ComponentProver, ComponentTrace, ComponentTraceWriter, + Air, AirProver, AirTraceVerifier, AirTraceWriter, Component, ComponentProver, ComponentTrace, + ComponentTraceWriter, }; use crate::core::backend::simd::column::BaseFieldVec; use crate::core::backend::simd::m31::{PackedBaseField, LOG_N_LANES}; use crate::core::backend::simd::qm31::PackedSecureField; use crate::core::backend::simd::SimdBackend; use crate::core::backend::{Col, Column, ColumnOps}; +use crate::core::channel::Blake2sChannel; +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::{FieldExpOps, FieldOps}; use crate::core::poly::circle::{CanonicCoset, CircleEvaluation}; use crate::core::poly::BitReversedOrder; use crate::core::{ColumnVec, ComponentVec, InteractionElements}; -use crate::examples::wide_fibonacci::component::N_COLUMNS; +use crate::examples::wide_fibonacci::component::{ALPHA_ID, N_COLUMNS, Z_ID}; -impl AirTraceWriter for WideFibAir { +// TODO(AlonH): Remove this once the Cpu and Simd implementations are aligned. +pub struct SimdWideFibComponent { + pub log_fibonacci_size: u32, + pub log_n_instances: u32, +} + +impl SimdWideFibComponent { + /// Returns the log of the size of the columns in the trace (which could also be looked at as + /// the log number of rows). + pub fn log_column_size(&self) -> u32 { + self.log_n_instances + self.log_fibonacci_size - LOG_N_COLUMNS as u32 + } + + pub fn log_n_columns(&self) -> usize { + LOG_N_COLUMNS + } + + pub fn n_columns(&self) -> usize { + N_COLUMNS + } +} + +// TODO(AlonH): Remove this once the Cpu and Simd implementations are aligned. +pub struct SimdWideFibAir { + pub component: SimdWideFibComponent, +} + +impl Air for SimdWideFibAir { + fn components(&self) -> Vec<&dyn Component> { + vec![&self.component] + } +} + +impl AirTraceVerifier for SimdWideFibAir { + fn interaction_elements(&self, _channel: &mut Blake2sChannel) -> InteractionElements { + InteractionElements::default() + } +} + +impl AirTraceWriter for SimdWideFibAir { fn interact( &self, _trace: &ColumnVec>, @@ -34,7 +78,52 @@ impl AirTraceWriter for WideFibAir { } } -impl AirProver for WideFibAir { +impl Component for SimdWideFibComponent { + fn n_constraints(&self) -> usize { + self.n_columns() - 2 + } + + fn max_constraint_log_degree_bound(&self) -> u32 { + self.log_column_size() + 1 + } + + fn n_interaction_phases(&self) -> u32 { + 1 + } + + fn trace_log_degree_bounds(&self) -> Vec { + vec![self.log_column_size(); self.n_columns()] + } + + fn mask_points( + &self, + point: CirclePoint, + ) -> ColumnVec>> { + fixed_mask_points(&vec![vec![0_usize]; self.n_columns()], point) + } + + fn interaction_element_ids(&self) -> Vec { + vec![ALPHA_ID.to_string(), Z_ID.to_string()] + } + + fn evaluate_constraint_quotients_at_point( + &self, + point: CirclePoint, + mask: &ColumnVec>, + evaluation_accumulator: &mut PointEvaluationAccumulator, + _interaction_elements: &InteractionElements, + ) { + let constraint_zero_domain = CanonicCoset::new(self.log_column_size()).coset; + let denom = coset_vanishing(constraint_zero_domain, point); + let denom_inverse = denom.inverse(); + for i in 0..self.n_columns() - 2 { + let numerator = mask[i][0].square() + mask[i + 1][0].square() - mask[i + 2][0]; + evaluation_accumulator.accumulate(numerator * denom_inverse); + } + } +} + +impl AirProver for SimdWideFibAir { fn prover_components(&self) -> Vec<&dyn ComponentProver> { vec![&self.component] } @@ -66,7 +155,8 @@ pub fn gen_trace( .collect_vec() } -impl ComponentTraceWriter for WideFibComponent { +// TODO(AlonH): Implement. +impl ComponentTraceWriter for SimdWideFibComponent { fn write_interaction_trace( &self, _trace: &ColumnVec<&CircleEvaluation>, @@ -76,7 +166,7 @@ impl ComponentTraceWriter for WideFibComponent { } } -impl ComponentProver for WideFibComponent { +impl ComponentProver for SimdWideFibComponent { fn evaluate_constraint_quotients_on_domain( &self, trace: &ComponentTrace<'_, SimdBackend>, @@ -140,7 +230,6 @@ impl ComponentProver for WideFibComponent { mod tests { use tracing::{span, Level}; - use super::{gen_trace, WideFibAir}; use crate::core::backend::simd::SimdBackend; use crate::core::channel::{Blake2sChannel, Channel}; use crate::core::fields::m31::BaseField; @@ -148,7 +237,8 @@ mod tests { use crate::core::prover::{prove, verify}; use crate::core::vcs::blake2_hash::Blake2sHasher; use crate::core::vcs::hasher::Hasher; - use crate::examples::wide_fibonacci::component::{WideFibComponent, LOG_N_COLUMNS}; + use crate::examples::wide_fibonacci::component::LOG_N_COLUMNS; + use crate::examples::wide_fibonacci::simd::{gen_trace, SimdWideFibAir, SimdWideFibComponent}; #[test_log::test] fn test_simd_wide_fib_prove() { @@ -159,7 +249,7 @@ mod tests { // Note: 17 means 128MB of trace. const LOG_N_ROWS: u32 = 12; - let component = WideFibComponent { + let component = SimdWideFibComponent { log_fibonacci_size: LOG_N_COLUMNS as u32, log_n_instances: LOG_N_ROWS, }; @@ -167,7 +257,7 @@ mod tests { let trace = gen_trace(component.log_column_size()); span.exit(); let channel = &mut Blake2sChannel::new(Blake2sHasher::hash(BaseField::into_slice(&[]))); - let air = WideFibAir { component }; + let air = SimdWideFibAir { component }; let proof = prove::(&air, channel, trace).unwrap(); let channel = &mut Blake2sChannel::new(Blake2sHasher::hash(BaseField::into_slice(&[])));