From f54597b77c784c050e5a922f3bb9cacb73c79916 Mon Sep 17 00:00:00 2001 From: Alon Haramati Date: Sun, 26 May 2024 16:37:28 +0300 Subject: [PATCH] Use SecureField in interaction. --- crates/prover/src/core/air/air_ext.rs | 34 ++++++++++++------- crates/prover/src/core/air/mod.rs | 4 +-- crates/prover/src/core/mod.rs | 8 ++--- crates/prover/src/core/prover/mod.rs | 22 ++++-------- crates/prover/src/core/utils.rs | 14 +++++--- .../src/examples/fibonacci/component.rs | 4 +-- .../src/examples/wide_fibonacci/component.rs | 26 ++++++++++---- .../wide_fibonacci/constraint_eval.rs | 8 +++-- .../prover/src/examples/wide_fibonacci/mod.rs | 33 ++++++++++-------- .../src/examples/wide_fibonacci/simd.rs | 4 +-- .../src/examples/wide_fibonacci/trace_gen.rs | 10 +++--- 11 files changed, 97 insertions(+), 70 deletions(-) diff --git a/crates/prover/src/core/air/air_ext.rs b/crates/prover/src/core/air/air_ext.rs index c77480e6e..971582653 100644 --- a/crates/prover/src/core/air/air_ext.rs +++ b/crates/prover/src/core/air/air_ext.rs @@ -10,7 +10,7 @@ use crate::core::circle::CirclePoint; use crate::core::fields::m31::BaseField; use crate::core::fields::qm31::SecureField; use crate::core::pcs::{CommitmentTreeProver, TreeVec}; -use crate::core::poly::circle::{CircleEvaluation, SecureCirclePoly}; +use crate::core::poly::circle::{CircleEvaluation, CirclePoly, SecureCirclePoly}; use crate::core::poly::BitReversedOrder; use crate::core::vcs::blake2_merkle::Blake2sMerkleHasher; use crate::core::vcs::ops::MerkleOps; @@ -51,7 +51,7 @@ pub trait AirExt: Air { for component in self.components() { ids.extend(component.interaction_element_ids()); } - let elements = channel.draw_felts(ids.len()).into_iter().map(|e| e.0 .0); + let elements = channel.draw_felts(ids.len()); InteractionElements(BTreeMap::from_iter(zip_eq(ids, elements))) } @@ -132,18 +132,26 @@ pub trait AirProverExt: AirProver { &self, trace: &ColumnVec>, elements: &InteractionElements, - ) -> ComponentVec> { + ) -> Vec> { let trace_iter = &mut trace.iter(); - ComponentVec( - self.prover_components() - .iter() - .map(|component| { - let n_columns = component.trace_log_degree_bounds()[0].len(); - let trace_columns = trace_iter.take(n_columns).collect_vec(); - component.write_interaction_trace(&trace_columns, elements) - }) - .collect(), - ) + + self.prover_components() + .iter() + .flat_map(|component| { + let n_columns = component.trace_log_degree_bounds()[0].len(); + let trace_columns = trace_iter.take(n_columns).collect_vec(); + component + .write_interaction_trace(&trace_columns, elements) + .into_iter() + .flat_map(|eval| { + eval.values.columns.map(|c| { + CircleEvaluation::::new(eval.domain, c) + .interpolate() + }) + }) + .collect_vec() + }) + .collect() } fn compute_composition_polynomial( diff --git a/crates/prover/src/core/air/mod.rs b/crates/prover/src/core/air/mod.rs index a52d7dcb5..d24f1a89f 100644 --- a/crates/prover/src/core/air/mod.rs +++ b/crates/prover/src/core/air/mod.rs @@ -4,7 +4,7 @@ use super::circle::CirclePoint; use super::fields::m31::BaseField; use super::fields::qm31::SecureField; use super::pcs::TreeVec; -use super::poly::circle::{CircleEvaluation, CirclePoly}; +use super::poly::circle::{CircleEvaluation, CirclePoly, SecureEvaluation}; use super::poly::BitReversedOrder; use super::{ColumnVec, InteractionElements}; @@ -60,7 +60,7 @@ pub trait ComponentTraceWriter { &self, trace: &ColumnVec<&CircleEvaluation>, elements: &InteractionElements, - ) -> ColumnVec>; + ) -> ColumnVec>; } pub trait ComponentProver: Component + ComponentTraceWriter { diff --git a/crates/prover/src/core/mod.rs b/crates/prover/src/core/mod.rs index 0c21afdae..a4b869706 100644 --- a/crates/prover/src/core/mod.rs +++ b/crates/prover/src/core/mod.rs @@ -1,7 +1,7 @@ use std::collections::BTreeMap; use std::ops::{Deref, DerefMut, Index}; -use self::fields::m31::BaseField; +use self::fields::qm31::SecureField; pub mod air; pub mod backend; @@ -61,10 +61,10 @@ impl DerefMut for ComponentVec { } } -pub struct InteractionElements(BTreeMap); +pub struct InteractionElements(BTreeMap); impl InteractionElements { - pub fn new(elements: BTreeMap) -> Self { + pub fn new(elements: BTreeMap) -> Self { Self(elements) } @@ -74,7 +74,7 @@ impl InteractionElements { } impl Index<&str> for InteractionElements { - type Output = BaseField; + type Output = SecureField; fn index(&self, index: &str) -> &Self::Output { // TODO(AlonH): Return an error if the key is not found. diff --git a/crates/prover/src/core/prover/mod.rs b/crates/prover/src/core/prover/mod.rs index 477e412c9..4f92c6370 100644 --- a/crates/prover/src/core/prover/mod.rs +++ b/crates/prover/src/core/prover/mod.rs @@ -4,6 +4,7 @@ use tracing::{span, Level}; use super::air::AirProver; use super::backend::Backend; +use super::fields::secure_column::SECURE_EXTENSION_DEGREE; use super::fri::FriVerificationError; use super::pcs::{CommitmentSchemeProof, TreeVec}; use super::poly::circle::{CanonicCoset, SecureCirclePoly, MAX_CIRCLE_DOMAIN_LOG_SIZE}; @@ -69,16 +70,7 @@ pub fn evaluate_and_commit_on_trace>( span.exit(); let interaction_elements = air.interaction_elements(channel); - let interaction_traces = air.interact(&trace, &interaction_elements); - let interaction_trace_polys = interaction_traces - .0 - .into_iter() - .flat_map(|trace| { - trace - .into_iter() - .map(|poly| poly.interpolate_with_twiddles(twiddles)) - }) - .collect_vec(); + let interaction_trace_polys = air.interact(&trace, &interaction_elements); let n_interaction_traces = interaction_trace_polys.len(); if n_interaction_traces > 0 { commitment_scheme.commit(interaction_trace_polys, channel, twiddles); @@ -116,7 +108,7 @@ pub fn generate_proof>( let mut sample_points = air.mask_points(oods_point); // Get composition polynomial sample points. - sample_points.push(vec![vec![oods_point]; 4]); + sample_points.push(vec![vec![oods_point]; SECURE_EXTENSION_DEGREE]); // Prove the trace and composition OODS values, and retrieve them. let commitment_scheme_proof = commitment_scheme.prove_values(sample_points, channel, twiddles); @@ -207,7 +199,7 @@ pub fn verify( // Read composition polynomial commitment. commitment_scheme.commit( *proof.commitments.last().unwrap(), - &[air.composition_log_degree_bound(); 4], + &[air.composition_log_degree_bound(); SECURE_EXTENSION_DEGREE], channel, ); @@ -218,7 +210,7 @@ pub fn verify( let mut sample_points = air.mask_points(oods_point); // Get composition polynomial sample points. - sample_points.push(vec![vec![oods_point]; 4]); + sample_points.push(vec![vec![oods_point]; SECURE_EXTENSION_DEGREE]); // TODO(spapini): Save clone. let (trace_oods_values, composition_oods_value) = sampled_values_to_mask( @@ -350,7 +342,7 @@ mod tests { use crate::core::fields::qm31::SecureField; use crate::core::pcs::TreeVec; use crate::core::poly::circle::{ - CanonicCoset, CircleDomain, CircleEvaluation, MAX_CIRCLE_DOMAIN_LOG_SIZE, + CanonicCoset, CircleDomain, CircleEvaluation, SecureEvaluation, MAX_CIRCLE_DOMAIN_LOG_SIZE, }; use crate::core::poly::BitReversedOrder; use crate::core::prover::{prove, ProvingError}; @@ -419,7 +411,7 @@ mod tests { &self, _trace: &ColumnVec<&CircleEvaluation>, _elements: &InteractionElements, - ) -> ColumnVec> { + ) -> ColumnVec> { vec![] } } diff --git a/crates/prover/src/core/utils.rs b/crates/prover/src/core/utils.rs index 7b033e6e7..cb6cc222e 100644 --- a/crates/prover/src/core/utils.rs +++ b/crates/prover/src/core/utils.rs @@ -1,6 +1,7 @@ use std::iter::Peekable; +use std::ops::Add; -use num_traits::One; +use num_traits::{One, Zero}; use super::fields::m31::BaseField; use super::fields::qm31::SecureField; @@ -92,12 +93,15 @@ pub fn generate_secure_powers(felt: SecureField, n_powers: usize) -> Vec>( values: &[F], - alpha: BaseField, - z: BaseField, -) -> F { + alpha: SecureField, + z: SecureField, +) -> SecureField +where + SecureField: Add, +{ let res = values .iter() - .fold(F::zero(), |acc, &value| acc * alpha + value); + .fold(SecureField::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 270f2c6d1..0f45d5225 100644 --- a/crates/prover/src/examples/fibonacci/component.rs +++ b/crates/prover/src/examples/fibonacci/component.rs @@ -12,7 +12,7 @@ use crate::core::fields::m31::BaseField; use crate::core::fields::qm31::SecureField; use crate::core::fields::{ExtensionOf, FieldExpOps}; use crate::core::pcs::TreeVec; -use crate::core::poly::circle::{CanonicCoset, CircleEvaluation}; +use crate::core::poly::circle::{CanonicCoset, CircleEvaluation, SecureEvaluation}; use crate::core::poly::BitReversedOrder; use crate::core::utils::bit_reverse_index; use crate::core::{ColumnVec, InteractionElements}; @@ -127,7 +127,7 @@ impl ComponentTraceWriter for FibonacciComponent { &self, _trace: &ColumnVec<&CircleEvaluation>, _elements: &InteractionElements, - ) -> ColumnVec> { + ) -> ColumnVec> { vec![] } } diff --git a/crates/prover/src/examples/wide_fibonacci/component.rs b/crates/prover/src/examples/wide_fibonacci/component.rs index 9606ca256..647dc3af8 100644 --- a/crates/prover/src/examples/wide_fibonacci/component.rs +++ b/crates/prover/src/examples/wide_fibonacci/component.rs @@ -8,9 +8,12 @@ use crate::core::circle::CirclePoint; use crate::core::constraints::{coset_vanishing, point_vanishing}; use crate::core::fields::m31::BaseField; use crate::core::fields::qm31::SecureField; +use crate::core::fields::secure_column::{SecureColumn, SECURE_EXTENSION_DEGREE}; use crate::core::fields::FieldExpOps; use crate::core::pcs::TreeVec; -use crate::core::poly::circle::{CanonicCoset, CircleEvaluation}; +use crate::core::poly::circle::{ + CanonicCoset, CircleEvaluation, SecureCirclePoly, SecureEvaluation, +}; use crate::core::poly::BitReversedOrder; use crate::core::utils::shifted_secure_combination; use crate::core::{ColumnVec, InteractionElements}; @@ -69,7 +72,7 @@ impl Component for WideFibComponent { fn trace_log_degree_bounds(&self) -> TreeVec> { TreeVec::new(vec![ vec![self.log_column_size(); self.n_columns()], - vec![self.log_column_size(); 1], + vec![self.log_column_size(); SECURE_EXTENSION_DEGREE], ]) } @@ -79,7 +82,7 @@ impl Component for WideFibComponent { ) -> TreeVec>>> { TreeVec::new(vec![ fixed_mask_points(&vec![vec![0_usize]; self.n_columns()], point), - vec![vec![point]], + vec![vec![point]; SECURE_EXTENSION_DEGREE], ]) } @@ -96,7 +99,11 @@ impl Component for WideFibComponent { ) { 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] + let lookup_value = + SecureCirclePoly::::eval_from_partial_evals(std::array::from_fn(|i| { + mask[self.n_columns() + i][0] + })); + let lookup_numerator = (lookup_value * shifted_secure_combination( &[mask[self.n_columns() - 2][0], mask[self.n_columns() - 1][0]], alpha, @@ -120,12 +127,19 @@ impl ComponentTraceWriter for WideFibComponent { &self, trace: &ColumnVec<&CircleEvaluation>, elements: &InteractionElements, - ) -> ColumnVec> { + ) -> ColumnVec> { let interaction_trace_domain = trace[0].domain; let trace_values = trace.iter().map(|eval| &eval.values[..]).collect_vec(); let (alpha, z) = (elements[ALPHA_ID], elements[Z_ID]); let values = write_lookup_column(&trace_values, alpha, z); - let eval = CircleEvaluation::new(interaction_trace_domain, values); + let mut secure_column = SecureColumn::::zeros(values.len()); + for (i, value) in values.into_iter().enumerate() { + secure_column.set(i, value); + } + let eval = SecureEvaluation { + domain: interaction_trace_domain, + values: secure_column, + }; vec![eval] } } diff --git a/crates/prover/src/examples/wide_fibonacci/constraint_eval.rs b/crates/prover/src/examples/wide_fibonacci/constraint_eval.rs index 33042ff2a..41b9ddd96 100644 --- a/crates/prover/src/examples/wide_fibonacci/constraint_eval.rs +++ b/crates/prover/src/examples/wide_fibonacci/constraint_eval.rs @@ -10,7 +10,7 @@ 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; +use crate::core::poly::circle::{CanonicCoset, SecureCirclePoly}; use crate::core::utils::{bit_reverse, shifted_secure_combination}; use crate::core::{ColumnVec, InteractionElements}; use crate::examples::wide_fibonacci::component::LOG_N_COLUMNS; @@ -61,8 +61,12 @@ impl ComponentProver for WideFibComponent { } // Lookup constraints. + let lookup_value = + SecureCirclePoly::::eval_from_partial_evals(std::array::from_fn(|j| { + trace_evals[1][j][i].into() + })); lookup_numerators[i] = accum.random_coeff_powers[self.n_columns() - 2] - * ((trace_evals[1][0][i] + * ((lookup_value * shifted_secure_combination( &[ trace_evals[0][self.n_columns() - 2][i], diff --git a/crates/prover/src/examples/wide_fibonacci/mod.rs b/crates/prover/src/examples/wide_fibonacci/mod.rs index 809a35616..aa42d1e59 100644 --- a/crates/prover/src/examples/wide_fibonacci/mod.rs +++ b/crates/prover/src/examples/wide_fibonacci/mod.rs @@ -18,9 +18,10 @@ mod tests { use crate::core::backend::CpuBackend; use crate::core::channel::{Blake2sChannel, Channel}; use crate::core::fields::m31::BaseField; + use crate::core::fields::qm31::SecureField; use crate::core::fields::IntoSlice; use crate::core::pcs::TreeVec; - use crate::core::poly::circle::CanonicCoset; + use crate::core::poly::circle::{CanonicCoset, CircleEvaluation}; use crate::core::poly::BitReversedOrder; use crate::core::prover::{prove, verify}; use crate::core::utils::shifted_secure_combination; @@ -40,15 +41,15 @@ mod tests { } pub fn assert_constraints_on_lookup_column( - column: &[BaseField], + column: &[SecureField], input_trace: &[Vec], - alpha: BaseField, - z: BaseField, + 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 = BaseField::one(); + let mut prev_value = SecureField::one(); for (i, cell) in column.iter().enumerate() { assert_eq!( *cell @@ -110,8 +111,8 @@ mod tests { b: m31!(1), }; - let alpha = m31!(7); - let z = m31!(11); + let alpha = qm31!(7, 1, 3, 4); + let z = qm31!(11, 1, 2, 3); let 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); @@ -163,15 +164,19 @@ mod tests { .iter() .cloned() .enumerate() - .map(|(i, id)| (id, m31!(43 + i as u32))), + .map(|(i, id)| (id, qm31!(43 + i as u32, 1, 2, 3))), )); - let interaction_trace = - wide_fib.write_interaction_trace(&trace.iter().collect(), &interaction_elements); - - let interaction_poly = interaction_trace - .iter() - .map(|trace| trace.clone().interpolate()) + let interaction_poly = wide_fib + .write_interaction_trace(&trace.iter().collect(), &interaction_elements) + .into_iter() + .flat_map(|eval| { + eval.values.columns.map(|c| { + CircleEvaluation::::new(eval.domain, c) + .interpolate() + }) + }) .collect_vec(); + let interaction_trace = interaction_poly .iter() .map(|poly| poly.evaluate(eval_domain)) diff --git a/crates/prover/src/examples/wide_fibonacci/simd.rs b/crates/prover/src/examples/wide_fibonacci/simd.rs index 67fa6b865..0b7842991 100644 --- a/crates/prover/src/examples/wide_fibonacci/simd.rs +++ b/crates/prover/src/examples/wide_fibonacci/simd.rs @@ -19,7 +19,7 @@ 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::circle::{CanonicCoset, CircleEvaluation, SecureEvaluation}; use crate::core::poly::BitReversedOrder; use crate::core::{ColumnVec, InteractionElements}; use crate::examples::wide_fibonacci::component::{ALPHA_ID, N_COLUMNS, Z_ID}; @@ -139,7 +139,7 @@ impl ComponentTraceWriter for SimdWideFibComponent { &self, _trace: &ColumnVec<&CircleEvaluation>, _elements: &InteractionElements, - ) -> ColumnVec> { + ) -> ColumnVec> { vec![] } } diff --git a/crates/prover/src/examples/wide_fibonacci/trace_gen.rs b/crates/prover/src/examples/wide_fibonacci/trace_gen.rs index 2d5ecd148..9f6e9e50b 100644 --- a/crates/prover/src/examples/wide_fibonacci/trace_gen.rs +++ b/crates/prover/src/examples/wide_fibonacci/trace_gen.rs @@ -2,6 +2,7 @@ use num_traits::One; 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::shifted_secure_combination; @@ -27,13 +28,12 @@ pub fn write_trace_row( /// the shifted secure combination of the last two elements in each row. pub fn write_lookup_column( input_trace: &[&[BaseField]], - // TODO(AlonH): Change alpha and z to SecureField. - alpha: BaseField, - z: BaseField, -) -> Vec { + alpha: SecureField, + z: SecureField, +) -> Vec { let n_rows = input_trace[0].len(); let n_columns = input_trace.len(); - let mut prev_value = BaseField::one(); + let mut prev_value = SecureField::one(); (0..n_rows) .map(|i| { let numerator =