From 90407e60e1e329308c9851e59374b751b20f5a4c Mon Sep 17 00:00:00 2001 From: Shahar Samocha Date: Tue, 27 Aug 2024 11:17:46 +0300 Subject: [PATCH] Refactor wide fibonacci Remove lookup - each row contain independent instance --- .../src/examples/wide_fibonacci/component.rs | 297 ------------ .../wide_fibonacci/constraint_eval.rs | 369 -------------- .../prover/src/examples/wide_fibonacci/mod.rs | 454 +++++++++--------- .../src/examples/wide_fibonacci/simd.rs | 294 ------------ .../src/examples/wide_fibonacci/trace_gen.rs | 79 --- 5 files changed, 217 insertions(+), 1276 deletions(-) delete mode 100644 crates/prover/src/examples/wide_fibonacci/component.rs delete mode 100644 crates/prover/src/examples/wide_fibonacci/constraint_eval.rs delete mode 100644 crates/prover/src/examples/wide_fibonacci/simd.rs delete mode 100644 crates/prover/src/examples/wide_fibonacci/trace_gen.rs diff --git a/crates/prover/src/examples/wide_fibonacci/component.rs b/crates/prover/src/examples/wide_fibonacci/component.rs deleted file mode 100644 index 4eb08740f..000000000 --- a/crates/prover/src/examples/wide_fibonacci/component.rs +++ /dev/null @@ -1,297 +0,0 @@ -use itertools::Itertools; - -use crate::core::air::accumulation::PointEvaluationAccumulator; -use crate::core::air::mask::fixed_mask_points; -use crate::core::air::{Air, Component}; -use crate::core::backend::cpu::CpuCircleEvaluation; -use crate::core::backend::CpuBackend; -use crate::core::circle::{CirclePoint, Coset}; -use crate::core::constraints::{coset_vanishing, point_excluder, point_vanishing}; -use crate::core::fields::m31::BaseField; -use crate::core::fields::qm31::SecureField; -use crate::core::fields::secure_column::{SecureColumnByCoords, SECURE_EXTENSION_DEGREE}; -use crate::core::fields::FieldExpOps; -use crate::core::pcs::TreeVec; -use crate::core::poly::circle::{CanonicCoset, CircleEvaluation}; -use crate::core::poly::BitReversedOrder; -use crate::core::utils::shifted_secure_combination; -use crate::core::{ColumnVec, InteractionElements, LookupValues}; -use crate::examples::wide_fibonacci::trace_gen::write_lookup_column; -use crate::trace_generation::registry::ComponentGenerationRegistry; -use crate::trace_generation::ComponentTraceGenerator; - -pub const LOG_N_COLUMNS: usize = 8; -pub const N_COLUMNS: usize = 1 << LOG_N_COLUMNS; - -pub const ALPHA_ID: &str = "wide_fibonacci_alpha"; -pub const Z_ID: &str = "wide_fibonacci_z"; -pub const LOOKUP_VALUE_0_ID: &str = "wide_fibonacci_0"; -pub const LOOKUP_VALUE_1_ID: &str = "wide_fibonacci_1"; -pub const LOOKUP_VALUE_N_MINUS_2_ID: &str = "wide_fibonacci_n-2"; -pub const LOOKUP_VALUE_N_MINUS_1_ID: &str = "wide_fibonacci_n-1"; - -/// 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 -/// number of rows (i.e the size of the columns) is determined by the parameters above (see -/// [WideFibComponent::log_column_size()]). -#[derive(Clone)] -pub struct WideFibComponent { - pub log_fibonacci_size: u32, - pub log_n_instances: u32, -} - -impl WideFibComponent { - /// 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 - } - - pub fn interaction_element_ids(&self) -> Vec { - vec![ALPHA_ID.to_string(), Z_ID.to_string()] - } - - fn evaluate_trace_boundary_constraints_at_point( - &self, - point: CirclePoint, - mask: &TreeVec>>, - evaluation_accumulator: &mut PointEvaluationAccumulator, - constraint_zero_domain: Coset, - lookup_values: &LookupValues, - ) { - let numerator = mask[0][0][0] - lookup_values[LOOKUP_VALUE_0_ID]; - let denom = point_vanishing(constraint_zero_domain.at(0), point); - evaluation_accumulator.accumulate(numerator / denom); - let numerator = mask[0][1][0] - lookup_values[LOOKUP_VALUE_1_ID]; - evaluation_accumulator.accumulate(numerator / denom); - - let numerator = mask[0][self.n_columns() - 2][0] - lookup_values[LOOKUP_VALUE_N_MINUS_2_ID]; - let denom = point_vanishing( - constraint_zero_domain.at(constraint_zero_domain.size() - 1), - point, - ); - evaluation_accumulator.accumulate(numerator / denom); - let numerator = mask[0][self.n_columns() - 1][0] - lookup_values[LOOKUP_VALUE_N_MINUS_1_ID]; - evaluation_accumulator.accumulate(numerator / denom); - } - - fn evaluate_trace_step_constraints_at_point( - &self, - point: CirclePoint, - mask: &ColumnVec>, - evaluation_accumulator: &mut PointEvaluationAccumulator, - constraint_zero_domain: 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); - } - } - - fn evaluate_lookup_boundary_constraints_at_point( - &self, - point: CirclePoint, - mask: &TreeVec>>, - evaluation_accumulator: &mut PointEvaluationAccumulator, - constraint_zero_domain: Coset, - interaction_elements: &InteractionElements, - lookup_values: &LookupValues, - ) { - let (alpha, z) = (interaction_elements[ALPHA_ID], interaction_elements[Z_ID]); - let value = SecureField::from_partial_evals(std::array::from_fn(|i| mask[1][i][0])); - let numerator = (value - * shifted_secure_combination( - &[ - mask[0][self.n_columns() - 2][0], - mask[0][self.n_columns() - 1][0], - ], - alpha, - z, - )) - - shifted_secure_combination(&[mask[0][0][0], mask[0][1][0]], alpha, z); - let denom = point_vanishing(constraint_zero_domain.at(0), point); - evaluation_accumulator.accumulate(numerator / denom); - - let numerator = (value - * shifted_secure_combination( - &[ - lookup_values[LOOKUP_VALUE_N_MINUS_2_ID], - lookup_values[LOOKUP_VALUE_N_MINUS_1_ID], - ], - alpha, - z, - )) - - shifted_secure_combination( - &[ - lookup_values[LOOKUP_VALUE_0_ID], - lookup_values[LOOKUP_VALUE_1_ID], - ], - alpha, - z, - ); - let denom = point_vanishing( - constraint_zero_domain.at(constraint_zero_domain.size() - 1), - point, - ); - evaluation_accumulator.accumulate(numerator / denom); - } - - fn evaluate_lookup_step_constraints_at_point( - &self, - point: CirclePoint, - mask: &TreeVec>>, - evaluation_accumulator: &mut PointEvaluationAccumulator, - constraint_zero_domain: Coset, - interaction_elements: &InteractionElements, - ) { - let (alpha, z) = (interaction_elements[ALPHA_ID], interaction_elements[Z_ID]); - let value = SecureField::from_partial_evals(std::array::from_fn(|i| mask[1][i][0])); - let prev_value = SecureField::from_partial_evals(std::array::from_fn(|i| mask[1][i][1])); - let numerator = (value - * shifted_secure_combination( - &[ - mask[0][self.n_columns() - 2][0], - mask[0][self.n_columns() - 1][0], - ], - alpha, - z, - )) - - (prev_value * shifted_secure_combination(&[mask[0][0][0], mask[0][1][0]], alpha, z)); - let denom = coset_vanishing(constraint_zero_domain, point) - / point_excluder(constraint_zero_domain.at(0), point); - evaluation_accumulator.accumulate(numerator / denom); - } -} - -#[derive(Clone)] -pub struct WideFibAir { - pub component: WideFibComponent, -} - -impl Air for WideFibAir { - fn components(&self) -> Vec<&dyn Component> { - vec![&self.component] - } -} - -impl Component for WideFibComponent { - fn n_constraints(&self) -> usize { - self.n_columns() + 5 - } - - fn max_constraint_log_degree_bound(&self) -> u32 { - self.log_column_size() + 1 - } - - fn trace_log_degree_bounds(&self) -> TreeVec> { - TreeVec::new(vec![ - vec![self.log_column_size(); self.n_columns()], - vec![self.log_column_size(); SECURE_EXTENSION_DEGREE], - ]) - } - - fn mask_points( - &self, - point: CirclePoint, - ) -> TreeVec>>> { - let domain = CanonicCoset::new(self.log_column_size()); - TreeVec::new(vec![ - fixed_mask_points(&vec![vec![0_usize]; self.n_columns()], point), - vec![vec![point, point - domain.step().into_ef()]; SECURE_EXTENSION_DEGREE], - ]) - } - - fn evaluate_constraint_quotients_at_point( - &self, - point: CirclePoint, - mask: &TreeVec>>, - evaluation_accumulator: &mut PointEvaluationAccumulator, - interaction_elements: &InteractionElements, - lookup_values: &LookupValues, - ) { - let constraint_zero_domain = CanonicCoset::new(self.log_column_size()).coset; - self.evaluate_trace_boundary_constraints_at_point( - point, - mask, - evaluation_accumulator, - constraint_zero_domain, - lookup_values, - ); - self.evaluate_lookup_step_constraints_at_point( - point, - mask, - evaluation_accumulator, - constraint_zero_domain, - interaction_elements, - ); - self.evaluate_lookup_boundary_constraints_at_point( - point, - mask, - evaluation_accumulator, - constraint_zero_domain, - interaction_elements, - lookup_values, - ); - self.evaluate_trace_step_constraints_at_point( - point, - &mask[0], - evaluation_accumulator, - constraint_zero_domain, - ); - } -} - -impl ComponentTraceGenerator for WideFibComponent { - type Component = Self; - type Inputs = (); - - fn add_inputs(&mut self, _inputs: &Self::Inputs) {} - - fn write_trace( - _component_id: &str, - _registry: &mut ComponentGenerationRegistry, - ) -> ColumnVec> { - vec![] - } - - fn write_interaction_trace( - &self, - trace: &ColumnVec<&CircleEvaluation>, - elements: &InteractionElements, - ) -> ColumnVec> { - let trace_values = trace.iter().map(|eval| &eval.values[..]).collect_vec(); - let (alpha, z) = (elements[ALPHA_ID], elements[Z_ID]); - // TODO(AlonH): Return a secure column directly. - let values = write_lookup_column(&trace_values, alpha, z); - let secure_column: SecureColumnByCoords = values.into_iter().collect(); - secure_column - .columns - .into_iter() - .map(|eval| { - let coset = CanonicCoset::new(trace[0].domain.log_size()); - CpuCircleEvaluation::new_canonical_ordered(coset, eval) - }) - .collect_vec() - } - - fn component(&self) -> Self::Component { - self.clone() - } -} - -// Input for the fibonacci claim. -#[derive(Debug, Clone, Copy)] -pub struct Input { - pub a: BaseField, - pub b: BaseField, -} diff --git a/crates/prover/src/examples/wide_fibonacci/constraint_eval.rs b/crates/prover/src/examples/wide_fibonacci/constraint_eval.rs deleted file mode 100644 index bd3f9025f..000000000 --- a/crates/prover/src/examples/wide_fibonacci/constraint_eval.rs +++ /dev/null @@ -1,369 +0,0 @@ -use std::collections::BTreeMap; - -use itertools::{zip_eq, Itertools}; -use num_traits::Zero; - -use super::component::{ - Input, WideFibAir, WideFibComponent, ALPHA_ID, LOOKUP_VALUE_0_ID, LOOKUP_VALUE_1_ID, - LOOKUP_VALUE_N_MINUS_1_ID, LOOKUP_VALUE_N_MINUS_2_ID, Z_ID, -}; -use super::trace_gen::write_trace_row; -use crate::core::air::accumulation::{ColumnAccumulator, DomainEvaluationAccumulator}; -use crate::core::air::{AirProver, Component, ComponentProver, Trace}; -use crate::core::backend::CpuBackend; -use crate::core::channel::Channel; -use crate::core::circle::Coset; -use crate::core::constraints::{coset_vanishing, point_excluder}; -use crate::core::fields::m31::BaseField; -use crate::core::fields::qm31::SecureField; -use crate::core::fields::FieldExpOps; -use crate::core::pcs::TreeVec; -use crate::core::poly::circle::{CanonicCoset, CircleDomain, CircleEvaluation}; -use crate::core::poly::BitReversedOrder; -use crate::core::prover::VerificationError; -use crate::core::utils::{ - bit_reverse, point_vanish_denominator_inverses, previous_bit_reversed_circle_domain_index, - shifted_secure_combination, -}; -use crate::core::{ColumnVec, InteractionElements, LookupValues}; -use crate::examples::wide_fibonacci::component::LOG_N_COLUMNS; -use crate::trace_generation::{ - AirTraceGenerator, AirTraceVerifier, ComponentTraceGenerator, BASE_TRACE, INTERACTION_TRACE, -}; - -// TODO(AlonH): Rename file to `cpu.rs`. - -impl AirTraceVerifier for WideFibAir { - fn interaction_elements(&self, channel: &mut impl Channel) -> InteractionElements { - let ids = self.component.interaction_element_ids(); - let elements = channel.draw_felts(ids.len()); - InteractionElements::new(BTreeMap::from_iter(zip_eq(ids, elements))) - } - - fn verify_lookups(&self, _lookup_values: &LookupValues) -> Result<(), VerificationError> { - Ok(()) - } -} - -impl AirTraceGenerator for WideFibAir { - fn interact( - &self, - trace: &ColumnVec>, - elements: &InteractionElements, - ) -> Vec> { - self.component - .write_interaction_trace(&trace.iter().collect(), elements) - } - - fn to_air_prover(&self) -> impl AirProver { - self.clone() - } - - fn composition_log_degree_bound(&self) -> u32 { - self.component.max_constraint_log_degree_bound() - } -} - -impl AirProver for WideFibAir { - fn component_provers(&self) -> Vec<&dyn ComponentProver> { - vec![&self.component] - } -} - -impl WideFibComponent { - fn evaluate_trace_boundary_constraints( - &self, - trace_evals: &TreeVec>>, - trace_eval_domain: CircleDomain, - zero_domain: Coset, - accum: &mut ColumnAccumulator<'_, CpuBackend>, - lookup_values: &LookupValues, - ) { - let first_point_denom_inverses = - point_vanish_denominator_inverses(trace_eval_domain, zero_domain.at(0)); - let last_point_denom_inverses = point_vanish_denominator_inverses( - trace_eval_domain, - zero_domain.at(zero_domain.size() - 1), - ); - let (lookup_value_0, lookup_value_1, lookup_value_n_minus_2, lookup_value_n_minus_1) = ( - lookup_values[LOOKUP_VALUE_0_ID], - lookup_values[LOOKUP_VALUE_1_ID], - lookup_values[LOOKUP_VALUE_N_MINUS_2_ID], - lookup_values[LOOKUP_VALUE_N_MINUS_1_ID], - ); - - for (i, (first_point_denom_inverse, last_point_denom_inverse)) in - zip_eq(first_point_denom_inverses, last_point_denom_inverses).enumerate() - { - let first_point_numerator = accum.random_coeff_powers[self.n_columns() + 4] - * (trace_evals[BASE_TRACE][0][i] - lookup_value_0) - + accum.random_coeff_powers[self.n_columns() + 3] - * (trace_evals[BASE_TRACE][1][i] - lookup_value_1); - let last_point_numerator = accum.random_coeff_powers[self.n_columns() + 2] - * (trace_evals[BASE_TRACE][self.n_columns() - 2][i] - lookup_value_n_minus_2) - + accum.random_coeff_powers[self.n_columns() + 1] - * (trace_evals[BASE_TRACE][self.n_columns() - 1][i] - lookup_value_n_minus_1); - accum.accumulate( - i, - first_point_numerator * first_point_denom_inverse - + last_point_numerator * last_point_denom_inverse, - ); - } - } - - fn evaluate_trace_step_constraints( - &self, - trace_evals: &TreeVec>>, - trace_eval_domain: CircleDomain, - zero_domain: Coset, - accum: &mut ColumnAccumulator<'_, CpuBackend>, - ) { - let max_constraint_degree = self.max_constraint_log_degree_bound(); - let mut denoms = vec![]; - for point in trace_eval_domain.iter() { - denoms.push(coset_vanishing(zero_domain, point)); - } - bit_reverse(&mut denoms); - let mut denom_inverses = vec![BaseField::zero(); 1 << (max_constraint_degree)]; - BaseField::batch_inverse(&denoms, &mut denom_inverses); - - for (i, denom_inverse) in denom_inverses.iter().enumerate() { - let mut numerator = SecureField::zero(); - for j in 0..self.n_columns() - 2 { - numerator += accum.random_coeff_powers[self.n_columns() - 3 - j] - * (trace_evals[BASE_TRACE][j][i].square() - + trace_evals[BASE_TRACE][j + 1][i].square() - - trace_evals[BASE_TRACE][j + 2][i]); - } - accum.accumulate(i, numerator * *denom_inverse) - } - } - - fn evaluate_lookup_boundary_constraints( - &self, - trace_evals: &TreeVec>>, - trace_eval_domain: CircleDomain, - zero_domain: Coset, - accum: &mut ColumnAccumulator<'_, CpuBackend>, - interaction_elements: &InteractionElements, - lookup_values: &LookupValues, - ) { - let first_point_denom_inverses = - point_vanish_denominator_inverses(trace_eval_domain, zero_domain.at(0)); - let last_point_denom_inverses = point_vanish_denominator_inverses( - trace_eval_domain, - zero_domain.at(zero_domain.size() - 1), - ); - let (alpha, z) = (interaction_elements[ALPHA_ID], interaction_elements[Z_ID]); - let (lookup_value_0, lookup_value_1, lookup_value_n_minus_2, lookup_value_n_minus_1) = ( - lookup_values[LOOKUP_VALUE_0_ID], - lookup_values[LOOKUP_VALUE_1_ID], - lookup_values[LOOKUP_VALUE_N_MINUS_2_ID], - lookup_values[LOOKUP_VALUE_N_MINUS_1_ID], - ); - - for (i, (first_point_denom_inverse, last_point_denom_inverse)) in - zip_eq(first_point_denom_inverses, last_point_denom_inverses).enumerate() - { - let value = SecureField::from_m31_array(std::array::from_fn(|j| { - trace_evals[INTERACTION_TRACE][j][i] - })); - let first_point_numerator = accum.random_coeff_powers[self.n_columns() - 1] - * ((value - * shifted_secure_combination( - &[ - trace_evals[BASE_TRACE][self.n_columns() - 2][i], - trace_evals[BASE_TRACE][self.n_columns() - 1][i], - ], - alpha, - z, - )) - - shifted_secure_combination( - &[trace_evals[BASE_TRACE][0][i], trace_evals[BASE_TRACE][1][i]], - alpha, - z, - )); - let last_point_numerator = accum.random_coeff_powers[self.n_columns() - 2] - * ((value - * shifted_secure_combination( - &[lookup_value_n_minus_2, lookup_value_n_minus_1], - alpha, - z, - )) - - shifted_secure_combination(&[lookup_value_0, lookup_value_1], alpha, z)); - accum.accumulate( - i, - first_point_numerator * first_point_denom_inverse - + last_point_numerator * last_point_denom_inverse, - ); - } - } - - // TODO(AlonH): Simplify this function by using utility functions. - fn evaluate_lookup_step_constraints( - &self, - trace_evals: &TreeVec>>, - trace_eval_domain: CircleDomain, - zero_domain: Coset, - accum: &mut ColumnAccumulator<'_, CpuBackend>, - interaction_elements: &InteractionElements, - ) { - let max_constraint_degree = self.max_constraint_log_degree_bound(); - let mut denoms = vec![]; - for point in trace_eval_domain.iter() { - denoms.push( - coset_vanishing(zero_domain, point) / point_excluder(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); - let (alpha, z) = (interaction_elements[ALPHA_ID], interaction_elements[Z_ID]); - - for (i, denom_inverse) in denom_inverses.iter().enumerate() { - let value = SecureField::from_m31_array(std::array::from_fn(|j| { - trace_evals[INTERACTION_TRACE][j][i] - })); - let prev_index = previous_bit_reversed_circle_domain_index( - i, - zero_domain.log_size, - trace_eval_domain.log_size(), - ); - let prev_value = SecureField::from_m31_array(std::array::from_fn(|j| { - trace_evals[INTERACTION_TRACE][j][prev_index] - })); - let numerator = accum.random_coeff_powers[self.n_columns()] - * ((value - * shifted_secure_combination( - &[ - trace_evals[BASE_TRACE][self.n_columns() - 2][i], - trace_evals[BASE_TRACE][self.n_columns() - 1][i], - ], - alpha, - z, - )) - - (prev_value - * shifted_secure_combination( - &[trace_evals[BASE_TRACE][0][i], trace_evals[BASE_TRACE][1][i]], - alpha, - z, - ))); - accum.accumulate(i, numerator * *denom_inverse); - } - } -} - -impl ComponentProver for WideFibComponent { - fn evaluate_constraint_quotients_on_domain( - &self, - trace: &Trace<'_, CpuBackend>, - evaluation_accumulator: &mut DomainEvaluationAccumulator, - interaction_elements: &InteractionElements, - lookup_values: &LookupValues, - ) { - 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 accum] = - evaluation_accumulator.columns([(max_constraint_degree, self.n_constraints())]); - - // TODO(AlonH): Evaluate the numerators together and the denominators together (i.e. in the - // same for loop) - self.evaluate_trace_boundary_constraints( - trace_evals, - trace_eval_domain, - zero_domain, - &mut accum, - lookup_values, - ); - self.evaluate_lookup_step_constraints( - trace_evals, - trace_eval_domain, - zero_domain, - &mut accum, - interaction_elements, - ); - self.evaluate_lookup_boundary_constraints( - trace_evals, - trace_eval_domain, - zero_domain, - &mut accum, - interaction_elements, - lookup_values, - ); - self.evaluate_trace_step_constraints( - trace_evals, - trace_eval_domain, - zero_domain, - &mut accum, - ); - } - - fn lookup_values(&self, trace: &Trace<'_, CpuBackend>) -> LookupValues { - let domain = CanonicCoset::new(self.log_column_size()); - let trace_poly = &trace.polys[BASE_TRACE]; - let values = BTreeMap::from_iter([ - ( - LOOKUP_VALUE_0_ID.to_string(), - trace_poly[0] - .eval_at_point(domain.at(0).into_ef()) - .try_into() - .unwrap(), - ), - ( - LOOKUP_VALUE_1_ID.to_string(), - trace_poly[1] - .eval_at_point(domain.at(0).into_ef()) - .try_into() - .unwrap(), - ), - ( - LOOKUP_VALUE_N_MINUS_2_ID.to_string(), - trace_poly[self.n_columns() - 2] - .eval_at_point(domain.at(domain.size() - 1).into_ef()) - .try_into() - .unwrap(), - ), - ( - LOOKUP_VALUE_N_MINUS_1_ID.to_string(), - trace_poly[self.n_columns() - 1] - .eval_at_point(domain.at(domain.size() - 1).into_ef()) - .try_into() - .unwrap(), - ), - ]); - LookupValues::new(values) - } -} - -/// Generates the trace for the wide Fibonacci example. -pub fn gen_trace( - wide_fib: &WideFibComponent, - private_input: Vec, -) -> ColumnVec> { - let n_instances = 1 << wide_fib.log_n_instances; - assert_eq!( - private_input.len(), - n_instances, - "The number of inputs must match the number of instances" - ); - assert!( - wide_fib.log_fibonacci_size >= LOG_N_COLUMNS as u32, - "The fibonacci size must be at least equal to the length of a row" - ); - let n_rows_per_instance = 1 << (wide_fib.log_fibonacci_size - wide_fib.log_n_columns() as u32); - let n_rows = n_instances * n_rows_per_instance; - let zero_vec = vec![BaseField::zero(); n_rows]; - let mut dst = vec![zero_vec; wide_fib.n_columns()]; - (0..n_rows_per_instance).fold(private_input, |input, row| { - (0..n_instances) - .map(|instance| { - let (a, b) = - write_trace_row(&mut dst, &input[instance], row * n_instances + instance); - Input { a, b } - }) - .collect_vec() - }); - dst -} diff --git a/crates/prover/src/examples/wide_fibonacci/mod.rs b/crates/prover/src/examples/wide_fibonacci/mod.rs index 7eecee5a4..85f673aec 100644 --- a/crates/prover/src/examples/wide_fibonacci/mod.rs +++ b/crates/prover/src/examples/wide_fibonacci/mod.rs @@ -1,287 +1,267 @@ -pub mod component; -pub mod constraint_eval; -pub mod simd; -pub mod trace_gen; +use itertools::Itertools; + +use crate::constraint_framework::{EvalAtRow, FrameworkComponent, FrameworkEval}; +use crate::core::backend::simd::m31::{PackedBaseField, LOG_N_LANES}; +use crate::core::backend::simd::SimdBackend; +use crate::core::backend::{Col, Column}; +use crate::core::fields::m31::BaseField; +use crate::core::fields::FieldExpOps; +use crate::core::poly::circle::{CanonicCoset, CircleEvaluation}; +use crate::core::poly::BitReversedOrder; +use crate::core::ColumnVec; + +pub type WideFibonacciComponent = FrameworkComponent>; + +pub struct FibInput { + a: PackedBaseField, + b: PackedBaseField, +} + +/// A component that enforces the Fibonacci sequence. +/// Each row contains a seperate Fibonacci sequence of length `N`. +#[derive(Clone)] +pub struct WideFibonacciEval { + pub log_n_rows: u32, +} +impl FrameworkEval for WideFibonacciEval { + fn log_size(&self) -> u32 { + self.log_n_rows + } + fn max_constraint_log_degree_bound(&self) -> u32 { + self.log_n_rows + 1 + } + fn evaluate(&self, mut eval: E) -> E { + let mut a = eval.next_trace_mask(); + let mut b = eval.next_trace_mask(); + for _ in 2..N { + let c = eval.next_trace_mask(); + eval.add_constraint(c - (a.square() + b.square())); + a = b; + b = c; + } + eval + } +} + +pub fn generate_trace( + log_size: u32, + inputs: &[FibInput], +) -> ColumnVec> { + assert!(log_size >= LOG_N_LANES); + assert_eq!(inputs.len(), 1 << (log_size - LOG_N_LANES)); + let mut trace = (0..N) + .map(|_| Col::::zeros(1 << log_size)) + .collect_vec(); + for (vec_index, input) in inputs.iter().enumerate() { + let mut a = input.a; + let mut b = input.b; + trace[0].data[vec_index] = a; + trace[1].data[vec_index] = b; + trace.iter_mut().skip(2).for_each(|col| { + (a, b) = (b, a.square() + b.square()); + col.data[vec_index] = b; + }); + } + let domain = CanonicCoset::new(log_size).circle_domain(); + trace + .into_iter() + .map(|eval| CircleEvaluation::::new(domain, eval)) + .collect_vec() +} #[cfg(test)] mod tests { - use std::collections::BTreeMap; - use itertools::Itertools; - use num_traits::{One, Zero}; + use num_traits::One; - use super::component::{Input, WideFibAir, WideFibComponent, LOG_N_COLUMNS}; - use super::constraint_eval::gen_trace; - use crate::core::air::accumulation::DomainEvaluationAccumulator; - use crate::core::air::{Component, ComponentProver, Trace}; - use crate::core::backend::cpu::CpuCircleEvaluation; - use crate::core::backend::CpuBackend; + use super::WideFibonacciEval; + use crate::constraint_framework::{ + assert_constraints, AssertEvaluator, FrameworkEval, TraceLocationAllocator, + }; + use crate::core::air::Component; + use crate::core::backend::simd::m31::{PackedBaseField, LOG_N_LANES}; + use crate::core::backend::simd::SimdBackend; + use crate::core::backend::Column; use crate::core::channel::Blake2sChannel; #[cfg(not(target_arch = "wasm32"))] use crate::core::channel::Poseidon252Channel; use crate::core::fields::m31::BaseField; - use crate::core::fields::qm31::SecureField; - use crate::core::pcs::{PcsConfig, TreeVec}; - use crate::core::poly::circle::CanonicCoset; - use crate::core::utils::{ - bit_reverse, circle_domain_order_to_coset_order, shifted_secure_combination, - }; + 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::{prove, verify}; use crate::core::vcs::blake2_merkle::Blake2sMerkleChannel; #[cfg(not(target_arch = "wasm32"))] use crate::core::vcs::poseidon252_merkle::Poseidon252MerkleChannel; - use crate::core::InteractionElements; - use crate::examples::wide_fibonacci::trace_gen::write_lookup_column; - use crate::trace_generation::{commit_and_prove, commit_and_verify, ComponentTraceGenerator}; - use crate::{m31, qm31}; - - pub fn assert_constraints_on_row(row: &[BaseField]) { - for i in 2..row.len() { - assert_eq!( - (row[i] - (row[i - 1] * row[i - 1] + row[i - 2] * row[i - 2])), - BaseField::zero() - ); - } - } + use crate::core::{ColumnVec, InteractionElements}; + use crate::examples::wide_fibonacci::{generate_trace, FibInput, WideFibonacciComponent}; - pub fn assert_constraints_on_lookup_column( - column: &[SecureField], - input_trace: &[Vec], - alpha: SecureField, - z: SecureField, - ) { - let n_columns = input_trace.len(); - let column_length = column.len(); - assert_eq!(column_length, input_trace[0].len()); - let mut prev_value = SecureField::one(); - for (i, cell) in column.iter().enumerate() { - assert_eq!( - *cell - * shifted_secure_combination( - &[input_trace[n_columns - 2][i], input_trace[n_columns - 1][i]], - alpha, - z, - ), - shifted_secure_combination(&[input_trace[0][i], input_trace[1][i]], alpha, z) - * prev_value - ); - prev_value = *cell; - } + const FIB_SEQUENCE_LENGTH: usize = 100; - // Assert the last cell in the column is equal to the combination of the first two values - // divided by the combination of the last two values in the sequence (all other values - // should cancel out). - assert_eq!( - column[column_length - 1] - * shifted_secure_combination( - &[input_trace[n_columns - 2][1], input_trace[n_columns - 1][1]], - alpha, - z, - ), - (shifted_secure_combination(&[input_trace[0][0], input_trace[1][0]], alpha, z)) - ); + fn generate_test_trace( + log_n_instances: u32, + ) -> ColumnVec> { + let inputs = (0..(1 << (log_n_instances - LOG_N_LANES))) + .map(|i| FibInput { + a: PackedBaseField::one(), + b: PackedBaseField::from_array(std::array::from_fn(|j| { + BaseField::from_u32_unchecked((i * 16 + j) as u32) + })), + }) + .collect_vec(); + generate_trace::(log_n_instances, &inputs) } - #[test] - fn test_trace_row_constraints() { - let wide_fib = WideFibComponent { - log_fibonacci_size: LOG_N_COLUMNS as u32, - log_n_instances: 1, - }; - let input = Input { - a: m31!(0x76), - b: m31!(0x483), - }; - - let trace = gen_trace(&wide_fib, vec![input, input]); - let row_0 = trace.iter().map(|col| col[0]).collect_vec(); - let row_1 = trace.iter().map(|col| col[1]).collect_vec(); - - assert_constraints_on_row(&row_0); - assert_constraints_on_row(&row_1); + fn fibonacci_constraint_evaluator(eval: AssertEvaluator<'_>) { + WideFibonacciEval:: { log_n_rows: N }.evaluate(eval); } #[test] - fn test_lookup_column_constraints() { - let wide_fib = WideFibComponent { - log_fibonacci_size: 4 + LOG_N_COLUMNS as u32, - log_n_instances: 0, - }; - let input = Input { - a: m31!(1), - b: m31!(1), - }; - - let alpha = qm31!(7, 1, 3, 4); - let z = qm31!(11, 1, 2, 3); - let mut trace = gen_trace(&wide_fib, vec![input]); - let input_trace = trace.iter().map(|values| &values[..]).collect_vec(); - let lookup_column = write_lookup_column(&input_trace, alpha, z); + fn test_wide_fibonacci_constraints() { + const LOG_N_INSTANCES: u32 = 6; + let traces = TreeVec::new(vec![generate_test_trace(LOG_N_INSTANCES)]); + let trace_polys = + traces.map(|trace| trace.into_iter().map(|c| c.interpolate()).collect_vec()); - trace = trace - .iter_mut() - .map(|column| { - bit_reverse(column); - circle_domain_order_to_coset_order(column) - }) - .collect_vec(); - assert_constraints_on_lookup_column(&lookup_column, &trace, alpha, z) + assert_constraints( + &trace_polys, + CanonicCoset::new(LOG_N_INSTANCES), + fibonacci_constraint_evaluator::, + ); } #[test] - fn test_composition_is_low_degree() { - let wide_fib = WideFibComponent { - log_fibonacci_size: 3 + LOG_N_COLUMNS as u32, - log_n_instances: 0, - }; - let random_coeff = qm31!(1, 2, 3, 4); - let mut acc = DomainEvaluationAccumulator::new( - random_coeff, - wide_fib.max_constraint_log_degree_bound(), - wide_fib.n_constraints(), - ); - let inputs = (0..1 << wide_fib.log_n_instances) - .map(|i| Input { - a: m31!(1), - b: m31!(i + 1_u32), - }) - .collect_vec(); + #[should_panic] + fn test_wide_fibonacci_constraints_fails() { + const LOG_N_INSTANCES: u32 = 6; - let trace_values = gen_trace(&wide_fib, inputs); + let mut trace = generate_test_trace(LOG_N_INSTANCES); + // Modify the trace such that a constraint fail. + trace[17].values.set(2, BaseField::one()); + let traces = TreeVec::new(vec![trace]); + let trace_polys = + traces.map(|trace| trace.into_iter().map(|c| c.interpolate()).collect_vec()); - let trace_domain = CanonicCoset::new(wide_fib.log_column_size()); - let trace = trace_values - .into_iter() - .map(|eval| CpuCircleEvaluation::new_canonical_ordered(trace_domain, eval)) - .collect_vec(); - let trace_polys = trace - .clone() - .into_iter() - .map(|eval| eval.interpolate()) - .collect_vec(); - let eval_domain = - CanonicCoset::new(wide_fib.max_constraint_log_degree_bound()).circle_domain(); - let trace_evals = trace_polys - .iter() - .map(|poly| poly.evaluate(eval_domain)) - .collect_vec(); - - let interaction_elements = InteractionElements::new(BTreeMap::from_iter( - wide_fib - .interaction_element_ids() - .iter() - .cloned() - .enumerate() - .map(|(i, id)| (id, qm31!(43 + i as u32, 1, 2, 3))), - )); - let interaction_poly = wide_fib - .write_interaction_trace(&trace.iter().collect(), &interaction_elements) - .into_iter() - .map(|eval| eval.interpolate()) - .collect_vec(); - - let interaction_trace = interaction_poly - .iter() - .map(|poly| poly.evaluate(eval_domain)) - .collect_vec(); - let trace = Trace { - polys: TreeVec::new(vec![ - trace_polys.iter().collect_vec(), - interaction_poly.iter().collect_vec(), - ]), - evals: TreeVec::new(vec![ - trace_evals.iter().collect_vec(), - interaction_trace.iter().collect_vec(), - ]), - }; - - let lookup_values = wide_fib.lookup_values(&trace); - wide_fib.evaluate_constraint_quotients_on_domain( - &trace, - &mut acc, - &interaction_elements, - &lookup_values, + assert_constraints( + &trace_polys, + CanonicCoset::new(LOG_N_INSTANCES), + fibonacci_constraint_evaluator::, ); - - let res = acc.finalize(); - let poly = res.0[0].clone(); - for coeff in poly.coeffs[(1 << wide_fib.max_constraint_log_degree_bound()) - 1..].iter() { - assert_eq!(*coeff, BaseField::zero()); - } } #[test_log::test] - fn test_single_instance_wide_fib_prove() { - // Note: To see time measurement, run test with - // RUST_LOG_SPAN_EVENTS=enter,close RUST_LOG=info RUST_BACKTRACE=1 cargo test - // test_prove -- --nocapture - - const LOG_N_INSTANCES: u32 = 0; + fn test_wide_fib_prove() { + const LOG_N_INSTANCES: u32 = 6; let config = PcsConfig::default(); - let component = WideFibComponent { - log_fibonacci_size: 3 + LOG_N_COLUMNS as u32, - log_n_instances: LOG_N_INSTANCES, - }; - let private_input = (0..(1 << LOG_N_INSTANCES)) - .map(|i| Input { - a: m31!(1), - b: m31!(i), - }) - .collect(); - let trace = gen_trace(&component, private_input); + // Precompute twiddles. + let twiddles = SimdBackend::precompute_twiddles( + CanonicCoset::new(LOG_N_INSTANCES + 1 + config.fri_config.log_blowup_factor) + .circle_domain() + .half_coset, + ); - let trace_domain = CanonicCoset::new(component.log_column_size()); - let trace = trace - .into_iter() - .map(|eval| CpuCircleEvaluation::new_canonical_ordered(trace_domain, eval)) - .collect_vec(); - let air = WideFibAir { component }; + // Setup protocol. let prover_channel = &mut Blake2sChannel::default(); - let proof = commit_and_prove::( - &air, + 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, + }, + ); + + let proof = prove::( + &[&component], prover_channel, - trace, - config, + &InteractionElements::default(), + commitment_scheme, ) .unwrap(); + // Verify. let verifier_channel = &mut Blake2sChannel::default(); - commit_and_verify::(proof, &air, verifier_channel, config).unwrap(); + 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, + &InteractionElements::default(), + commitment_scheme, + proof, + ) + .unwrap(); } + #[test] #[cfg(not(target_arch = "wasm32"))] - #[test_log::test] - fn test_single_instance_wide_fib_prove_with_poseidon() { - const LOG_N_INSTANCES: u32 = 0; + fn test_wide_fib_prove_with_poseidon() { + const LOG_N_INSTANCES: u32 = 6; + let config = PcsConfig::default(); - let component = WideFibComponent { - log_fibonacci_size: 3 + LOG_N_COLUMNS as u32, - log_n_instances: LOG_N_INSTANCES, - }; - let private_input = (0..(1 << LOG_N_INSTANCES)) - .map(|i| Input { - a: m31!(1), - b: m31!(i), - }) - .collect(); - let trace = gen_trace(&component, private_input); + // Precompute twiddles. + let twiddles = SimdBackend::precompute_twiddles( + CanonicCoset::new(LOG_N_INSTANCES + 1 + config.fri_config.log_blowup_factor) + .circle_domain() + .half_coset, + ); - let trace_domain = CanonicCoset::new(component.log_column_size()); - let trace = trace - .into_iter() - .map(|eval| CpuCircleEvaluation::new_canonical_ordered(trace_domain, eval)) - .collect_vec(); - let air = WideFibAir { component }; + // Setup protocol. let prover_channel = &mut Poseidon252Channel::default(); - let proof = commit_and_prove::( - &air, + 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, + }, + ); + let proof = prove::( + &[&component], prover_channel, - trace, - config, + &InteractionElements::default(), + commitment_scheme, ) .unwrap(); + // Verify. let verifier_channel = &mut Poseidon252Channel::default(); - commit_and_verify::(proof, &air, verifier_channel, config) - .unwrap(); + 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, + &InteractionElements::default(), + commitment_scheme, + proof, + ) + .unwrap(); } } diff --git a/crates/prover/src/examples/wide_fibonacci/simd.rs b/crates/prover/src/examples/wide_fibonacci/simd.rs deleted file mode 100644 index a9509cdcf..000000000 --- a/crates/prover/src/examples/wide_fibonacci/simd.rs +++ /dev/null @@ -1,294 +0,0 @@ -use itertools::Itertools; -use num_traits::{One, Zero}; -use tracing::{span, Level}; - -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::{Air, AirProver, Component, ComponentProver, Trace}; -use crate::core::backend::simd::column::BaseColumn; -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::Channel; -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::pcs::TreeVec; -use crate::core::poly::circle::{CanonicCoset, CircleEvaluation}; -use crate::core::poly::BitReversedOrder; -use crate::core::prover::VerificationError; -use crate::core::{ColumnVec, InteractionElements, LookupValues}; -use crate::examples::wide_fibonacci::component::N_COLUMNS; -use crate::trace_generation::registry::ComponentGenerationRegistry; -use crate::trace_generation::{ - AirTraceGenerator, AirTraceVerifier, ComponentTraceGenerator, BASE_TRACE, -}; - -// TODO(AlonH): Remove this once the Cpu and Simd implementations are aligned. -#[derive(Clone)] -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. -#[derive(Clone)] -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 impl Channel) -> InteractionElements { - InteractionElements::default() - } - - fn verify_lookups(&self, _lookup_values: &LookupValues) -> Result<(), VerificationError> { - Ok(()) - } -} - -impl AirTraceGenerator for SimdWideFibAir { - fn interact( - &self, - _trace: &ColumnVec>, - _elements: &InteractionElements, - ) -> Vec> { - vec![] - } - - fn to_air_prover(&self) -> impl AirProver { - self.clone() - } - - fn composition_log_degree_bound(&self) -> u32 { - self.component.max_constraint_log_degree_bound() - } -} - -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 trace_log_degree_bounds(&self) -> TreeVec> { - TreeVec::new(vec![vec![self.log_column_size(); self.n_columns()]]) - } - - fn mask_points( - &self, - point: CirclePoint, - ) -> TreeVec>>> { - TreeVec::new(vec![fixed_mask_points( - &vec![vec![0_usize]; self.n_columns()], - point, - )]) - } - - fn evaluate_constraint_quotients_at_point( - &self, - point: CirclePoint, - mask: &TreeVec>>, - evaluation_accumulator: &mut PointEvaluationAccumulator, - _interaction_elements: &InteractionElements, - _lookup_values: &LookupValues, - ) { - 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[0][i][0].square() + mask[0][i + 1][0].square() - mask[0][i + 2][0]; - evaluation_accumulator.accumulate(numerator * denom_inverse); - } - } -} - -impl AirProver for SimdWideFibAir { - fn component_provers(&self) -> Vec<&dyn ComponentProver> { - vec![&self.component] - } -} - -pub fn gen_trace( - log_size: u32, -) -> ColumnVec> { - assert!(log_size >= LOG_N_LANES); - let mut trace = (0..N_COLUMNS) - .map(|_| Col::::zeros(1 << log_size)) - .collect_vec(); - for vec_index in 0..(1 << (log_size - LOG_N_LANES)) { - let mut a = PackedBaseField::one(); - let mut b = PackedBaseField::from_array(std::array::from_fn(|i| { - BaseField::from_u32_unchecked((vec_index * 16 + i) as u32) - })); - trace[0].data[vec_index] = a; - trace[1].data[vec_index] = b; - trace.iter_mut().skip(2).for_each(|col| { - (a, b) = (b, a.square() + b.square()); - col.data[vec_index] = b; - }); - } - let domain = CanonicCoset::new(log_size).circle_domain(); - trace - .into_iter() - .map(|eval| CircleEvaluation::::new(domain, eval)) - .collect_vec() -} - -impl ComponentTraceGenerator for SimdWideFibComponent { - type Component = Self; - type Inputs = (); - - fn add_inputs(&mut self, _inputs: &Self::Inputs) {} - - fn write_trace( - _component_id: &str, - _registry: &mut ComponentGenerationRegistry, - ) -> ColumnVec> { - vec![] - } - - fn write_interaction_trace( - &self, - _trace: &ColumnVec<&CircleEvaluation>, - _elements: &InteractionElements, - ) -> ColumnVec> { - vec![] - } - - fn component(&self) -> Self::Component { - self.clone() - } -} - -impl ComponentProver for SimdWideFibComponent { - fn evaluate_constraint_quotients_on_domain( - &self, - trace: &Trace<'_, SimdBackend>, - evaluation_accumulator: &mut DomainEvaluationAccumulator, - _interaction_elements: &InteractionElements, - _lookup_values: &LookupValues, - ) { - assert_eq!(trace.polys[BASE_TRACE].len(), self.n_columns()); - // TODO(spapini): Steal evaluation from commitment. - let eval_domain = CanonicCoset::new(self.log_column_size() + 1).circle_domain(); - let trace_eval = &trace.evals; - - // Denoms. - let span = span!(Level::INFO, "Constraint eval denominators").entered(); - // TODO(spapini): Make this prettier. - let zero_domain = CanonicCoset::new(self.log_column_size()).coset; - let mut denoms = - BaseColumn::from_iter(eval_domain.iter().map(|p| coset_vanishing(zero_domain, p))); - >::bit_reverse_column(&mut denoms); - let mut denom_inverses = BaseColumn::zeros(denoms.len()); - >::batch_inverse(&denoms, &mut denom_inverses); - span.exit(); - - let _span = span!(Level::INFO, "Constraint pointwise eval").entered(); - - let constraint_log_degree_bound = self.max_constraint_log_degree_bound(); - let n_constraints = self.n_constraints(); - let [accum] = - evaluation_accumulator.columns([(constraint_log_degree_bound, n_constraints)]); - - for vec_row in 0..(1 << (eval_domain.log_size() - LOG_N_LANES)) { - // Numerator. - let a = trace_eval[BASE_TRACE][0].data[vec_row]; - let mut row_res = PackedSecureField::zero(); - let mut a_sq = a.square(); - let mut b_sq = trace_eval[BASE_TRACE][1].data[vec_row].square(); - #[allow(clippy::needless_range_loop)] - for i in 0..(self.n_columns() - 2) { - unsafe { - let c = *trace_eval[BASE_TRACE] - .get_unchecked(i + 2) - .data - .get_unchecked(vec_row); - row_res += PackedSecureField::broadcast( - accum.random_coeff_powers[self.n_columns() - 3 - i], - ) * (a_sq + b_sq - c); - (a_sq, b_sq) = (b_sq, c.square()); - } - } - - unsafe { - accum.col.set_packed( - vec_row, - accum.col.packed_at(vec_row) + row_res * denom_inverses.data[vec_row], - ) - } - } - } - - fn lookup_values(&self, _trace: &Trace<'_, SimdBackend>) -> LookupValues { - LookupValues::default() - } -} - -#[cfg(test)] -mod tests { - use tracing::{span, Level}; - - use crate::core::channel::Blake2sChannel; - use crate::core::pcs::PcsConfig; - use crate::core::vcs::blake2_merkle::Blake2sMerkleChannel; - use crate::examples::wide_fibonacci::component::LOG_N_COLUMNS; - use crate::examples::wide_fibonacci::simd::{gen_trace, SimdWideFibAir, SimdWideFibComponent}; - use crate::trace_generation::{commit_and_prove, commit_and_verify}; - - #[test_log::test] - fn test_simd_wide_fib_prove() { - // Note: To see time measurement, run test with - // RUST_LOG_SPAN_EVENTS=enter,close RUST_LOG=info RUST_BACKTRACE=1 RUSTFLAGS=" - // -C target-cpu=native -C target-feature=+avx512f -C opt-level=3" cargo test - // test_simd_wide_fib_prove -- --nocapture - - // Note: 17 means 128MB of trace. - const LOG_N_ROWS: u32 = 12; - let config = PcsConfig::default(); - let component = SimdWideFibComponent { - log_fibonacci_size: LOG_N_COLUMNS as u32, - log_n_instances: LOG_N_ROWS, - }; - let span = span!(Level::INFO, "Trace generation").entered(); - let trace = gen_trace(component.log_column_size()); - span.exit(); - let channel = &mut Blake2sChannel::default(); - let air = SimdWideFibAir { component }; - let proof = - commit_and_prove::<_, Blake2sMerkleChannel>(&air, channel, trace, config).unwrap(); - - let channel = &mut Blake2sChannel::default(); - commit_and_verify::(proof, &air, channel, config).unwrap(); - } -} diff --git a/crates/prover/src/examples/wide_fibonacci/trace_gen.rs b/crates/prover/src/examples/wide_fibonacci/trace_gen.rs deleted file mode 100644 index 550aa0635..000000000 --- a/crates/prover/src/examples/wide_fibonacci/trace_gen.rs +++ /dev/null @@ -1,79 +0,0 @@ -use itertools::Itertools; -use num_traits::{One, Zero}; - -use super::component::Input; -use crate::core::fields::m31::BaseField; -use crate::core::fields::qm31::SecureField; -use crate::core::fields::FieldExpOps; -use crate::core::utils::{ - bit_reverse, circle_domain_order_to_coset_order, shifted_secure_combination, -}; - -/// Writes the trace row for the wide Fibonacci example to dst, given a private input. Returns the -/// last two elements of the row in case the sequence is continued. -pub fn write_trace_row( - dst: &mut [Vec], - private_input: &Input, - row_index: usize, -) -> (BaseField, BaseField) { - let n_columns = dst.len(); - dst[0][row_index] = private_input.a; - dst[1][row_index] = private_input.b; - for i in 2..n_columns { - dst[i][row_index] = dst[i - 1][row_index].square() + dst[i - 2][row_index].square(); - } - - (dst[n_columns - 2][row_index], dst[n_columns - 1][row_index]) -} - -/// Writes and returns the lookup column for the wide Fibonacci example, which is the partial -/// product of the shifted secure combination of the first two elements in each row divided by the -/// the shifted secure combination of the last two elements in each row. -pub fn write_lookup_column( - input_trace: &[&[BaseField]], - alpha: SecureField, - z: SecureField, -) -> Vec { - let n_rows = input_trace[0].len(); - let n_columns = input_trace.len(); - let mut prev_value = SecureField::one(); - let mut input_trace = input_trace - .iter() - .map(|column| column.to_vec()) - .collect_vec(); - let natural_ordered_trace = input_trace - .iter_mut() - .map(|column| { - bit_reverse(column); - circle_domain_order_to_coset_order(column) - }) - .collect_vec(); - - let denominators = (0..n_rows) - .map(|i| { - shifted_secure_combination( - &[ - natural_ordered_trace[n_columns - 2][i], - natural_ordered_trace[n_columns - 1][i], - ], - alpha, - z, - ) - }) - .collect_vec(); - let mut denominator_inverses = vec![SecureField::zero(); denominators.len()]; - SecureField::batch_inverse(&denominators, &mut denominator_inverses); - - (0..n_rows) - .map(|i| { - let numerator = shifted_secure_combination( - &[natural_ordered_trace[0][i], natural_ordered_trace[1][i]], - alpha, - z, - ); - let cell = (numerator * denominator_inverses[i]) * prev_value; - prev_value = cell; - cell - }) - .collect_vec() -}