From 3150e167b874e12f0727de3c9e95983e0d9bd3f4 Mon Sep 17 00:00:00 2001 From: Alon Haramati <91828241+alonh5@users.noreply.github.com> Date: Thu, 27 Jun 2024 15:38:46 +0300 Subject: [PATCH] Use SecureField in interaction. (#642) --- crates/prover/src/core/air/mod.rs | 4 +-- crates/prover/src/core/mod.rs | 8 +++--- .../src/core/poly/circle/secure_poly.rs | 1 + crates/prover/src/core/prover/mod.rs | 27 ++++++++---------- crates/prover/src/core/utils.rs | 14 ++++++---- crates/prover/src/examples/fibonacci/air.rs | 10 +++---- .../src/examples/wide_fibonacci/component.rs | 28 ++++++++++++++----- .../wide_fibonacci/constraint_eval.rs | 19 +++++++------ .../prover/src/examples/wide_fibonacci/mod.rs | 26 ++++++++--------- .../src/examples/wide_fibonacci/simd.rs | 6 ++-- .../src/examples/wide_fibonacci/trace_gen.rs | 10 +++---- 11 files changed, 86 insertions(+), 67 deletions(-) diff --git a/crates/prover/src/core/air/mod.rs b/crates/prover/src/core/air/mod.rs index 8cf8c725e..b559780df 100644 --- a/crates/prover/src/core/air/mod.rs +++ b/crates/prover/src/core/air/mod.rs @@ -7,7 +7,7 @@ use super::fields::qm31::SecureField; use super::pcs::TreeVec; use super::poly::circle::{CircleEvaluation, CirclePoly}; use super::poly::BitReversedOrder; -use super::{ColumnVec, ComponentVec, InteractionElements}; +use super::{ColumnVec, InteractionElements}; pub mod accumulation; mod air_ext; @@ -34,7 +34,7 @@ pub trait AirTraceWriter: AirTraceVerifier { &self, trace: &ColumnVec>, elements: &InteractionElements, - ) -> ComponentVec>; + ) -> Vec>; fn to_air_prover(&self) -> &impl AirProver; } diff --git a/crates/prover/src/core/mod.rs b/crates/prover/src/core/mod.rs index 3278e60cd..499738c18 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; @@ -62,10 +62,10 @@ impl DerefMut for ComponentVec { } #[derive(Default)] -pub struct InteractionElements(BTreeMap); +pub struct InteractionElements(BTreeMap); impl InteractionElements { - pub fn new(elements: BTreeMap) -> Self { + pub fn new(elements: BTreeMap) -> Self { Self(elements) } @@ -75,7 +75,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/poly/circle/secure_poly.rs b/crates/prover/src/core/poly/circle/secure_poly.rs index f9ff38107..931e283e8 100644 --- a/crates/prover/src/core/poly/circle/secure_poly.rs +++ b/crates/prover/src/core/poly/circle/secure_poly.rs @@ -35,6 +35,7 @@ impl SecureCirclePoly { /// Evaluates the polynomial at a point, given evaluations of its composing base field /// polynomials at that point. + // TODO(AlonH): Move to SecureField and rename. pub fn eval_from_partial_evals(evals: [SecureField; SECURE_EXTENSION_DEGREE]) -> SecureField { let mut res = evals[0]; res += evals[1] * SecureField::from_u32_unchecked(0, 1, 0, 0); diff --git a/crates/prover/src/core/prover/mod.rs b/crates/prover/src/core/prover/mod.rs index f297ae505..41cb83146 100644 --- a/crates/prover/src/core/prover/mod.rs +++ b/crates/prover/src/core/prover/mod.rs @@ -61,7 +61,7 @@ pub fn evaluate_and_commit_on_trace>( let trace_polys = trace .clone() .into_iter() - .map(|poly| poly.interpolate_with_twiddles(twiddles)) + .map(|eval| eval.interpolate_with_twiddles(twiddles)) .collect(); span.exit(); @@ -71,17 +71,14 @@ 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(); - if !interaction_trace_polys.is_empty() { + let interaction_trace = air.interact(&trace, &interaction_elements); + if !interaction_trace.is_empty() { + let span = span!(Level::INFO, "Interaction trace interpolation").entered(); + let interaction_trace_polys = interaction_trace + .into_iter() + .map(|eval| eval.interpolate_with_twiddles(twiddles)) + .collect(); + span.exit(); commitment_scheme.commit(interaction_trace_polys, channel, twiddles); } @@ -356,7 +353,7 @@ mod tests { use crate::core::poly::BitReversedOrder; use crate::core::prover::{prove, ProvingError}; use crate::core::test_utils::test_channel; - use crate::core::{ColumnVec, ComponentVec, InteractionElements}; + use crate::core::{ColumnVec, InteractionElements}; use crate::qm31; struct TestAir> { @@ -380,8 +377,8 @@ mod tests { &self, _trace: &ColumnVec>, _elements: &InteractionElements, - ) -> ComponentVec> { - ComponentVec(vec![vec![]]) + ) -> Vec> { + vec![] } fn to_air_prover(&self) -> &impl AirProver { 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/air.rs b/crates/prover/src/examples/fibonacci/air.rs index dbec91ecf..c8c1b8c39 100644 --- a/crates/prover/src/examples/fibonacci/air.rs +++ b/crates/prover/src/examples/fibonacci/air.rs @@ -9,7 +9,7 @@ use crate::core::channel::Blake2sChannel; use crate::core::fields::m31::BaseField; use crate::core::poly::circle::CircleEvaluation; use crate::core::poly::BitReversedOrder; -use crate::core::{ColumnVec, ComponentVec, InteractionElements}; +use crate::core::{ColumnVec, InteractionElements}; pub struct FibonacciAir { pub component: FibonacciComponent, @@ -38,8 +38,8 @@ impl AirTraceWriter for FibonacciAir { &self, _trace: &ColumnVec>, _elements: &InteractionElements, - ) -> ComponentVec> { - ComponentVec(vec![vec![]]) + ) -> Vec> { + vec![] } fn to_air_prover(&self) -> &impl AirProver { @@ -87,8 +87,8 @@ impl AirTraceWriter for MultiFibonacciAir { &self, _trace: &ColumnVec>, _elements: &InteractionElements, - ) -> ComponentVec> { - ComponentVec(vec![vec![]]) + ) -> Vec> { + vec![] } fn to_air_prover(&self) -> &impl AirProver { diff --git a/crates/prover/src/examples/wide_fibonacci/component.rs b/crates/prover/src/examples/wide_fibonacci/component.rs index 758559abd..6a0c0e685 100644 --- a/crates/prover/src/examples/wide_fibonacci/component.rs +++ b/crates/prover/src/examples/wide_fibonacci/component.rs @@ -8,9 +8,10 @@ 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}; use crate::core::poly::BitReversedOrder; use crate::core::utils::shifted_secure_combination; use crate::core::{ColumnVec, InteractionElements}; @@ -73,7 +74,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], ]) } @@ -83,7 +84,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], ]) } @@ -100,7 +101,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, @@ -125,12 +130,21 @@ impl ComponentTraceWriter for WideFibComponent { trace: &ColumnVec<&CircleEvaluation>, elements: &InteractionElements, ) -> 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]); + // TODO(AlonH): Return a secure column directly. let values = write_lookup_column(&trace_values, alpha, z); - let eval = CircleEvaluation::new(interaction_trace_domain, values); - vec![eval] + let secure_column: SecureColumn = values.into_iter().collect(); + secure_column + .columns + .into_iter() + .map(|eval| { + CircleEvaluation::::new( + trace[0].domain, + eval, + ) + }) + .collect_vec() } } diff --git a/crates/prover/src/examples/wide_fibonacci/constraint_eval.rs b/crates/prover/src/examples/wide_fibonacci/constraint_eval.rs index 9f309d844..841fe90f1 100644 --- a/crates/prover/src/examples/wide_fibonacci/constraint_eval.rs +++ b/crates/prover/src/examples/wide_fibonacci/constraint_eval.rs @@ -16,10 +16,10 @@ 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::circle::{CanonicCoset, CircleEvaluation, SecureCirclePoly}; use crate::core::poly::BitReversedOrder; use crate::core::utils::{bit_reverse, shifted_secure_combination}; -use crate::core::{ColumnVec, ComponentVec, InteractionElements}; +use crate::core::{ColumnVec, InteractionElements}; use crate::examples::wide_fibonacci::component::LOG_N_COLUMNS; // TODO(AlonH): Rename file to `cpu.rs`. @@ -27,7 +27,7 @@ use crate::examples::wide_fibonacci::component::LOG_N_COLUMNS; impl AirTraceVerifier for WideFibAir { fn interaction_elements(&self, channel: &mut Blake2sChannel) -> InteractionElements { let ids = self.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::new(BTreeMap::from_iter(zip_eq(ids, elements))) } } @@ -37,10 +37,9 @@ impl AirTraceWriter for WideFibAir { &self, trace: &ColumnVec>, elements: &InteractionElements, - ) -> ComponentVec> { - ComponentVec(vec![self - .component - .write_interaction_trace(&trace.iter().collect(), elements)]) + ) -> Vec> { + self.component + .write_interaction_trace(&trace.iter().collect(), elements) } fn to_air_prover(&self) -> &impl AirProver { @@ -92,8 +91,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..678e6ca22 100644 --- a/crates/prover/src/examples/wide_fibonacci/mod.rs +++ b/crates/prover/src/examples/wide_fibonacci/mod.rs @@ -18,6 +18,7 @@ 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; @@ -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,14 @@ 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() + .map(|eval| eval.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 cd7698f3e..877b03ba8 100644 --- a/crates/prover/src/examples/wide_fibonacci/simd.rs +++ b/crates/prover/src/examples/wide_fibonacci/simd.rs @@ -23,7 +23,7 @@ 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::{ColumnVec, ComponentVec, InteractionElements}; +use crate::core::{ColumnVec, InteractionElements}; use crate::examples::wide_fibonacci::component::{ALPHA_ID, N_COLUMNS, Z_ID}; // TODO(AlonH): Remove this once the Cpu and Simd implementations are aligned. @@ -70,8 +70,8 @@ impl AirTraceWriter for SimdWideFibAir { &self, _trace: &ColumnVec>, _elements: &InteractionElements, - ) -> ComponentVec> { - ComponentVec(vec![vec![]]) + ) -> Vec> { + vec![] } fn to_air_prover(&self) -> &impl AirProver { 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 =