diff --git a/air/src/air/aux.rs b/air/src/air/aux.rs index 7dc9c4f48..33d7d8539 100644 --- a/air/src/air/aux.rs +++ b/air/src/air/aux.rs @@ -5,9 +5,9 @@ use alloc::vec::Vec; -use math::FieldElement; +use math::{ExtensionOf, FieldElement}; -use super::{lagrange::LagrangeKernelRandElements, LogUpGkrOracle}; +use super::{LagrangeKernelRandElements, LogUpGkrOracle}; /// Holds the randomly generated elements used in defining the auxiliary segment of the trace. /// @@ -130,7 +130,11 @@ impl GkrData { .fold(E::ZERO, |acc, (a, b)| acc + *a * *b) } - pub fn compute_batched_query(&self, query: &[E::BaseField]) -> E { + pub fn compute_batched_query(&self, query: &[F]) -> E + where + F: FieldElement, + E: ExtensionOf, + { E::from(query[0]) + query .iter() diff --git a/air/src/air/context.rs b/air/src/air/context.rs index c36173ca3..8998e64e0 100644 --- a/air/src/air/context.rs +++ b/air/src/air/context.rs @@ -250,13 +250,9 @@ impl AirContext { self.aux_transition_constraint_degrees.len() } - /// Returns the index of the auxiliary column which implements the Lagrange kernel, if any - pub fn lagrange_kernel_aux_column_idx(&self) -> Option { - if self.logup_gkr_enabled() { - Some(self.trace_info().aux_segment_width() - 1) - } else { - None - } + /// Returns the index of the auxiliary column which implements the Lagrange kernel, if any. + pub fn lagrange_kernel_column_idx(&self) -> Option { + self.trace_info.lagrange_kernel_column_idx() } /// Returns true if LogUp-GKR is enabled. diff --git a/air/src/air/lagrange/boundary.rs b/air/src/air/logup_gkr/lagrange/boundary.rs similarity index 97% rename from air/src/air/lagrange/boundary.rs rename to air/src/air/logup_gkr/lagrange/boundary.rs index 5d1954615..f3cc886aa 100644 --- a/air/src/air/lagrange/boundary.rs +++ b/air/src/air/logup_gkr/lagrange/boundary.rs @@ -5,7 +5,7 @@ use math::FieldElement; -use crate::{LagrangeKernelEvaluationFrame, LagrangeKernelRandElements}; +use super::{LagrangeKernelEvaluationFrame, LagrangeKernelRandElements}; #[derive(Debug, Clone, Eq, PartialEq)] pub struct LagrangeKernelBoundaryConstraint diff --git a/air/src/air/lagrange/frame.rs b/air/src/air/logup_gkr/lagrange/frame.rs similarity index 100% rename from air/src/air/lagrange/frame.rs rename to air/src/air/logup_gkr/lagrange/frame.rs diff --git a/air/src/air/lagrange/mod.rs b/air/src/air/logup_gkr/lagrange/mod.rs similarity index 95% rename from air/src/air/lagrange/mod.rs rename to air/src/air/logup_gkr/lagrange/mod.rs index fed5897f3..9d80b4437 100644 --- a/air/src/air/lagrange/mod.rs +++ b/air/src/air/logup_gkr/lagrange/mod.rs @@ -3,17 +3,18 @@ // This source code is licensed under the MIT license found in the // LICENSE file in the root directory of this source tree. -mod boundary; use alloc::vec::Vec; use core::ops::Deref; +use math::FieldElement; + +mod boundary; pub use boundary::LagrangeKernelBoundaryConstraint; mod frame; pub use frame::LagrangeKernelEvaluationFrame; mod transition; -use math::FieldElement; pub use transition::LagrangeKernelTransitionConstraints; use crate::LagrangeConstraintsCompositionCoefficients; @@ -22,7 +23,6 @@ use crate::LagrangeConstraintsCompositionCoefficients; pub struct LagrangeKernelConstraints { pub transition: LagrangeKernelTransitionConstraints, pub boundary: LagrangeKernelBoundaryConstraint, - pub lagrange_kernel_col_idx: usize, } impl LagrangeKernelConstraints { @@ -30,7 +30,6 @@ impl LagrangeKernelConstraints { pub fn new( lagrange_composition_coefficients: LagrangeConstraintsCompositionCoefficients, lagrange_kernel_rand_elements: &LagrangeKernelRandElements, - lagrange_kernel_col_idx: usize, ) -> Self { Self { transition: LagrangeKernelTransitionConstraints::new( @@ -40,7 +39,6 @@ impl LagrangeKernelConstraints { lagrange_composition_coefficients.boundary, lagrange_kernel_rand_elements, ), - lagrange_kernel_col_idx, } } } diff --git a/air/src/air/lagrange/transition.rs b/air/src/air/logup_gkr/lagrange/transition.rs similarity index 100% rename from air/src/air/lagrange/transition.rs rename to air/src/air/logup_gkr/lagrange/transition.rs diff --git a/air/src/air/logup_gkr.rs b/air/src/air/logup_gkr/mod.rs similarity index 87% rename from air/src/air/logup_gkr.rs rename to air/src/air/logup_gkr/mod.rs index 0438064d9..a907fad40 100644 --- a/air/src/air/logup_gkr.rs +++ b/air/src/air/logup_gkr/mod.rs @@ -9,7 +9,15 @@ use core::marker::PhantomData; use crypto::{ElementHasher, RandomCoin}; use math::{ExtensionOf, FieldElement, StarkField, ToElements}; -use super::{EvaluationFrame, GkrData, LagrangeKernelRandElements}; +use super::{EvaluationFrame, GkrData, LagrangeConstraintsCompositionCoefficients}; +mod s_column; +use s_column::SColumnConstraint; + +mod lagrange; +pub use lagrange::{ + LagrangeKernelBoundaryConstraint, LagrangeKernelConstraints, LagrangeKernelEvaluationFrame, + LagrangeKernelRandElements, LagrangeKernelTransitionConstraints, +}; /// A trait containing the necessary information in order to run the LogUp-GKR protocol of [1]. /// @@ -116,6 +124,27 @@ pub trait LogUpGkrEvaluator: Clone + Sync { self.get_oracles().to_vec(), ) } + + /// Returns a new [`LagrangeKernelConstraints`]. + fn get_lagrange_kernel_constraints>( + &self, + lagrange_composition_coefficients: LagrangeConstraintsCompositionCoefficients, + lagrange_kernel_rand_elements: &LagrangeKernelRandElements, + ) -> LagrangeKernelConstraints { + LagrangeKernelConstraints::new( + lagrange_composition_coefficients, + lagrange_kernel_rand_elements, + ) + } + + /// Returns a new [`SColumnConstraints`]. + fn get_s_column_constraints>( + &self, + gkr_data: GkrData, + composition_coefficient: E, + ) -> SColumnConstraint { + SColumnConstraint::new(gkr_data, composition_coefficient) + } } #[derive(Clone, Default)] diff --git a/air/src/air/logup_gkr/s_column.rs b/air/src/air/logup_gkr/s_column.rs new file mode 100644 index 000000000..29848ceeb --- /dev/null +++ b/air/src/air/logup_gkr/s_column.rs @@ -0,0 +1,56 @@ +// Copyright (c) Facebook, Inc. and its affiliates. +// +// This source code is licensed under the MIT license found in the +// LICENSE file in the root directory of this source tree. + +use math::FieldElement; + +use super::{super::Air, EvaluationFrame, GkrData}; +use crate::LogUpGkrEvaluator; + +/// Represents the transition constraint for the s-column, as well as the random coefficient used +/// to linearly combine the constraint into the constraint composition polynomial. +/// +/// The s-column implements the cohomological sum-check argument of [1] and the constraint in +/// [`SColumnConstraint`] is exactly Eq (4) in Lemma 1 in [1]. +/// +/// +/// [1]: https://eprint.iacr.org/2021/930 +pub struct SColumnConstraint { + gkr_data: GkrData, + composition_coefficient: E, +} + +impl SColumnConstraint { + pub fn new(gkr_data: GkrData, composition_coefficient: E) -> Self { + Self { gkr_data, composition_coefficient } + } + + /// Evaluates the transition constraint over the specificed main trace segment, s-column, + /// and Lagrange kernel evaluation frames. + pub fn evaluate( + &self, + air: &A, + main_trace_frame: &EvaluationFrame, + s_cur: E, + s_nxt: E, + l_cur: E, + x: E, + ) -> E + where + A: Air, + { + let batched_claim = self.gkr_data.compute_batched_claim(); + let mean = batched_claim + .mul_base(E::BaseField::ONE / E::BaseField::from(air.trace_length() as u32)); + + let mut query = vec![E::ZERO; air.get_logup_gkr_evaluator().get_oracles().len()]; + air.get_logup_gkr_evaluator().build_query(main_trace_frame, &[], &mut query); + let batched_claim_at_query = self.gkr_data.compute_batched_query::(&query); + let rhs = s_cur - mean + batched_claim_at_query * l_cur; + let lhs = s_nxt; + + let divisor = x.exp((air.trace_length() as u32).into()) - E::ONE; + self.composition_coefficient * (rhs - lhs) / divisor + } +} diff --git a/air/src/air/mod.rs b/air/src/air/mod.rs index 5dcee0717..bedfa5e35 100644 --- a/air/src/air/mod.rs +++ b/air/src/air/mod.rs @@ -6,7 +6,6 @@ use alloc::{collections::BTreeMap, vec::Vec}; use crypto::{RandomCoin, RandomCoinError}; -use logup_gkr::PhantomLogUpGkrEval; use math::{fft, ExtensibleField, ExtensionOf, FieldElement, StarkField, ToElements}; use crate::ProofOptions; @@ -29,15 +28,14 @@ pub use boundary::{BoundaryConstraint, BoundaryConstraintGroup, BoundaryConstrai mod transition; pub use transition::{EvaluationFrame, TransitionConstraintDegree, TransitionConstraints}; -mod lagrange; -pub use lagrange::{ +mod logup_gkr; +use logup_gkr::PhantomLogUpGkrEval; +pub use logup_gkr::{ LagrangeKernelBoundaryConstraint, LagrangeKernelConstraints, LagrangeKernelEvaluationFrame, - LagrangeKernelRandElements, LagrangeKernelTransitionConstraints, + LagrangeKernelRandElements, LagrangeKernelTransitionConstraints, LogUpGkrEvaluator, + LogUpGkrOracle, }; -mod logup_gkr; -pub use logup_gkr::{LogUpGkrEvaluator, LogUpGkrOracle}; - mod coefficients; pub use coefficients::{ ConstraintCompositionCoefficients, DeepCompositionCoefficients, @@ -331,25 +329,6 @@ pub trait Air: Send + Sync { Ok(rand_elements) } - /// Returns a new [`LagrangeKernelConstraints`] if a Lagrange kernel auxiliary column is present - /// in the trace, or `None` otherwise. - fn get_lagrange_kernel_constraints>( - &self, - lagrange_composition_coefficients: LagrangeConstraintsCompositionCoefficients, - lagrange_kernel_rand_elements: &LagrangeKernelRandElements, - ) -> Option> { - if self.context().logup_gkr_enabled() { - let col_idx = self.context().trace_info().aux_segment_width() - 1; - Some(LagrangeKernelConstraints::new( - lagrange_composition_coefficients, - lagrange_kernel_rand_elements, - col_idx, - )) - } else { - None - } - } - /// Returns values for all periodic columns used in the computation. /// /// These values will be used to compute column values at specific states of the computation @@ -600,7 +579,7 @@ pub trait Air: Send + Sync { None }; - let s_col = if self.context().logup_gkr_enabled() { + let s_col_cc = if self.context().logup_gkr_enabled() { Some(public_coin.draw()?) } else { None @@ -610,7 +589,7 @@ pub trait Air: Send + Sync { trace: t_coefficients, constraints: c_coefficients, lagrange: lagrange_cc, - s_col, + s_col: s_col_cc, }) } } diff --git a/air/src/air/tests.rs b/air/src/air/tests.rs index 5e9871ca5..2400cb883 100644 --- a/air/src/air/tests.rs +++ b/air/src/air/tests.rs @@ -225,7 +225,6 @@ impl MockAir { impl Air for MockAir { type BaseField = BaseElement; type PublicInputs = (); - //type LogUpGkrEvaluator = DummyLogUpGkrEval; fn new(trace_info: TraceInfo, _pub_inputs: (), _options: ProofOptions) -> Self { let num_assertions = trace_info.meta()[0] as usize; diff --git a/air/src/air/trace_info.rs b/air/src/air/trace_info.rs index 44aa0a7ea..8c3545539 100644 --- a/air/src/air/trace_info.rs +++ b/air/src/air/trace_info.rs @@ -39,6 +39,10 @@ impl TraceInfo { pub const MAX_META_LENGTH: usize = 65535; /// Maximum number of random elements in the auxiliary trace segment; currently set to 255. pub const MAX_RAND_SEGMENT_ELEMENTS: usize = 255; + /// The Lagrange kernel, if present, is the last column of the auxiliary trace. + pub const LAGRANGE_KERNEL_OFFSET: usize = 1; + /// The s-column, if present, is the second to last column of the auxiliary trace. + pub const S_COLUMN_OFFSET: usize = 2; // CONSTRUCTORS // -------------------------------------------------------------------------------------------- @@ -112,7 +116,7 @@ impl TraceInfo { // validate trace segment widths assert!(main_segment_width > 0, "main trace segment must consist of at least one column"); - let full_width = main_segment_width + aux_segment_width; + let full_width = main_segment_width + aux_segment_width + 2 * logup_gkr as usize; assert!( full_width <= TraceInfo::MAX_TRACE_WIDTH, "total number of columns in the trace cannot be greater than {}, but was {}", @@ -170,9 +174,9 @@ impl TraceInfo { &self.trace_meta } - /// Returns true if an execution trace contains the auxiliary trace segment. + /// Returns true if an execution trace contains an auxiliary trace segment. pub fn is_multi_segment(&self) -> bool { - self.aux_segment_width > 0 + self.aux_segment_width > 0 || self.logup_gkr } /// Returns the number of columns in the main segment of an execution trace. @@ -210,6 +214,24 @@ impl TraceInfo { self.logup_gkr } + /// Returns the index of the auxiliary column which implements the Lagrange kernel, if any. + pub fn lagrange_kernel_column_idx(&self) -> Option { + if self.logup_gkr_enabled() { + Some(self.aux_segment_width() - TraceInfo::LAGRANGE_KERNEL_OFFSET) + } else { + None + } + } + + /// Returns the index of the auxiliary column which implements the s-column, if any. + pub fn s_column_idx(&self) -> Option { + if self.logup_gkr_enabled() { + Some(self.aux_segment_width() - TraceInfo::S_COLUMN_OFFSET) + } else { + None + } + } + /// Returns the number of random elements needed to build all auxiliary columns, except for the /// Lagrange kernel column. pub fn get_num_aux_segment_rand_elements(&self) -> usize { diff --git a/prover/src/constraints/evaluator/default.rs b/prover/src/constraints/evaluator/default.rs index ea02b41d4..4373494f9 100644 --- a/prover/src/constraints/evaluator/default.rs +++ b/prover/src/constraints/evaluator/default.rs @@ -13,9 +13,9 @@ use utils::iter_mut; use utils::{iterators::*, rayon}; use super::{ - super::EvaluationTableFragment, lagrange::LagrangeKernelConstraintsBatchEvaluator, - BoundaryConstraints, CompositionPolyTrace, ConstraintEvaluationTable, ConstraintEvaluator, - PeriodicValueTable, StarkDomain, TraceLde, + super::EvaluationTableFragment, logup_gkr::LogUpGkrConstraintsEvaluator, BoundaryConstraints, + CompositionPolyTrace, ConstraintEvaluationTable, ConstraintEvaluator, PeriodicValueTable, + StarkDomain, TraceLde, }; // CONSTANTS @@ -40,7 +40,7 @@ pub struct DefaultConstraintEvaluator<'a, A: Air, E: FieldElement, transition_constraints: TransitionConstraints, - lagrange_constraints_evaluator: Option>, + logup_gkr_constraints_evaluator: Option>, aux_rand_elements: Option>, periodic_values: PeriodicValueTable, } @@ -117,10 +117,10 @@ where evaluation_table.validate_transition_degrees(); // combine all constraint evaluations into a single column, including the evaluations of the - // Lagrange kernel constraints (if present) + // LogUp-GKR constraints (if present) let combined_evaluations = { let mut constraints_evaluations = evaluation_table.combine(); - self.evaluate_lagrange_kernel_constraints(trace, domain, &mut constraints_evaluations); + self.evaluate_logup_gkr_constraints(trace, domain, &mut constraints_evaluations); constraints_evaluations }; @@ -158,18 +158,21 @@ where &composition_coefficients.boundary, ); - let lagrange_constraints_evaluator = if air.context().logup_gkr_enabled() { + let logup_gkr_constraints_evaluator = if air.context().logup_gkr_enabled() { let aux_rand_elements = aux_rand_elements.as_ref().expect("expected aux rand elements to be present"); - let lagrange_rand_elements = aux_rand_elements - .lagrange() - .expect("expected lagrange rand elements to be present"); - Some(LagrangeKernelConstraintsBatchEvaluator::new( + + Some(LogUpGkrConstraintsEvaluator::new( air, - lagrange_rand_elements.clone(), + aux_rand_elements + .gkr_data() + .expect("expected LogUp-GKR randomness to be present"), composition_coefficients .lagrange .expect("expected Lagrange kernel composition coefficients to be present"), + composition_coefficients + .s_col + .expect("expected s-column composition coefficient to be present"), )) } else { None @@ -179,7 +182,7 @@ where air, boundary_constraints, transition_constraints, - lagrange_constraints_evaluator, + logup_gkr_constraints_evaluator, aux_rand_elements, periodic_values, } @@ -295,7 +298,7 @@ where } } - /// If present, evaluates the Lagrange kernel constraints over the constraint evaluation domain. + /// If present, evaluates the LogUp-GKR constraints over the constraint evaluation domain. /// The evaluation of each constraint (both boundary and transition) is divided by its divisor, /// multiplied by its composition coefficient, the result of which is added to /// `combined_evaluations_accumulator`. @@ -303,14 +306,14 @@ where /// Specifically, `combined_evaluations_accumulator` is a buffer whose length is the size of the /// constraint evaluation domain, where each index contains combined evaluations of other /// constraints in the system. - fn evaluate_lagrange_kernel_constraints>( + fn evaluate_logup_gkr_constraints>( &self, trace: &T, domain: &StarkDomain, combined_evaluations_accumulator: &mut [E], ) { - if let Some(ref lagrange_constraints_evaluator) = self.lagrange_constraints_evaluator { - lagrange_constraints_evaluator.evaluate_constraints( + if let Some(ref logup_gkr_constraints_evaluator) = self.logup_gkr_constraints_evaluator { + logup_gkr_constraints_evaluator.evaluate_constraints( trace, domain, combined_evaluations_accumulator, diff --git a/prover/src/constraints/evaluator/lagrange.rs b/prover/src/constraints/evaluator/logup_gkr.rs similarity index 70% rename from prover/src/constraints/evaluator/lagrange.rs rename to prover/src/constraints/evaluator/logup_gkr.rs index 89d07f62c..f8fa3ae36 100644 --- a/prover/src/constraints/evaluator/lagrange.rs +++ b/prover/src/constraints/evaluator/logup_gkr.rs @@ -6,46 +6,50 @@ use alloc::vec::Vec; use air::{ - Air, LagrangeConstraintsCompositionCoefficients, LagrangeKernelConstraints, - LagrangeKernelEvaluationFrame, LagrangeKernelRandElements, + Air, EvaluationFrame, GkrData, LagrangeConstraintsCompositionCoefficients, + LagrangeKernelConstraints, LagrangeKernelEvaluationFrame, LogUpGkrEvaluator, }; use math::{batch_inversion, FieldElement}; use crate::{StarkDomain, TraceLde}; -/// Contains a specific strategy for evaluating the Lagrange kernel boundary and transition -/// constraints where the divisors' evaluation is batched. -/// -/// Specifically, [`batch_inversion`] is used to reduce the number of divisions performed. -pub struct LagrangeKernelConstraintsBatchEvaluator { +/// Contains a specific strategy for evaluating the Lagrange kernel and s-column boundary and +/// transition constraints. +pub struct LogUpGkrConstraintsEvaluator<'a, E: FieldElement, A: Air> { + air: &'a A, lagrange_kernel_constraints: LagrangeKernelConstraints, - rand_elements: LagrangeKernelRandElements, + gkr_data: GkrData, + s_col_composition_coefficient: E, } -impl LagrangeKernelConstraintsBatchEvaluator { - /// Constructs a new [`LagrangeConstraintsBatchEvaluator`]. - pub fn new( - air: &A, - lagrange_kernel_rand_elements: LagrangeKernelRandElements, +impl<'a, E, A> LogUpGkrConstraintsEvaluator<'a, E, A> +where + E: FieldElement, + A: Air, +{ + /// Constructs a new [`LogUpGkrConstraintsEvaluator`]. + pub fn new( + air: &'a A, + gkr_data: GkrData, lagrange_composition_coefficients: LagrangeConstraintsCompositionCoefficients, - ) -> Self - where - E: FieldElement, - { + s_col_composition_coefficient: E, + ) -> Self { Self { lagrange_kernel_constraints: air + .get_logup_gkr_evaluator() .get_lagrange_kernel_constraints( lagrange_composition_coefficients, - &lagrange_kernel_rand_elements, - ) - .expect("expected Lagrange kernel constraints to be present"), - rand_elements: lagrange_kernel_rand_elements, + gkr_data.lagrange_kernel_rand_elements(), + ), + air, + gkr_data, + s_col_composition_coefficient, } } /// Evaluates the transition and boundary constraints. Specifically, the constraint evaluations /// are divided by their corresponding divisors, and the resulting terms are linearly combined - /// using the composition coefficients. + /// using the constraint composition coefficients. /// /// Writes the evaluations in `combined_evaluations_acc` at the corresponding (constraint /// evaluation) domain index. @@ -64,28 +68,44 @@ impl LagrangeKernelConstraintsBatchEvaluator { ); let boundary_divisors_inv = self.compute_boundary_divisors_inv(domain); - let mut frame = LagrangeKernelEvaluationFrame::new_empty(); + let mut lagrange_frame = LagrangeKernelEvaluationFrame::new_empty(); + + let evaluator = self.air.get_logup_gkr_evaluator(); + let s_col_constraint_divisor = compute_s_col_divisor::(domain, self.air.trace_length()); + let s_col_idx = trace.trace_info().s_column_idx().expect("S-column should be present"); + let l_col_idx = trace + .trace_info() + .lagrange_kernel_column_idx() + .expect("Lagrange kernel should be present"); + let mut main_frame = EvaluationFrame::new(trace.trace_info().main_segment_width()); + let mut aux_frame = EvaluationFrame::new(trace.trace_info().aux_segment_width()); + + let c = self.gkr_data.compute_batched_claim(); + let mean = c / E::from(E::BaseField::from(trace.trace_info().length() as u32)); + let mut query = vec![E::BaseField::ZERO; evaluator.get_oracles().len()]; for step in 0..domain.ce_domain_size() { // compute Lagrange kernel frame trace.read_lagrange_kernel_frame_into( step << lde_shift, - self.lagrange_kernel_constraints.lagrange_kernel_col_idx, - &mut frame, + l_col_idx, + &mut lagrange_frame, ); // Compute the combined transition and boundary constraints evaluations for this row - let combined_evaluations = { + let lagrange_combined_evaluations = { let mut combined_evaluations = E::ZERO; // combine transition constraints for trans_constraint_idx in 0..self.lagrange_kernel_constraints.transition.num_constraints() { - let numerator = self - .lagrange_kernel_constraints - .transition - .evaluate_ith_numerator(&frame, &self.rand_elements, trans_constraint_idx); + let numerator = + self.lagrange_kernel_constraints.transition.evaluate_ith_numerator( + &lagrange_frame, + &self.gkr_data.lagrange_kernel_eval_point, + trans_constraint_idx, + ); let inv_divisor = trans_constraints_divisors .get_inverse_divisor_eval(trans_constraint_idx, step); @@ -94,8 +114,10 @@ impl LagrangeKernelConstraintsBatchEvaluator { // combine boundary constraints { - let boundary_numerator = - self.lagrange_kernel_constraints.boundary.evaluate_numerator_at(&frame); + let boundary_numerator = self + .lagrange_kernel_constraints + .boundary + .evaluate_numerator_at(&lagrange_frame); combined_evaluations += boundary_numerator * boundary_divisors_inv[step]; } @@ -103,7 +125,33 @@ impl LagrangeKernelConstraintsBatchEvaluator { combined_evaluations }; - combined_evaluations_acc[step] += combined_evaluations; + // compute and combine the transition constraints for the s-column. + // The s-column implements the cohomological sum-check argument of [1] and + // the constraint we enfore is exactly Eq (4) in Lemma 1 in [1]. + // + // [1]: https://eprint.iacr.org/2021/930 + let s_col_combined_evaluation = { + trace.read_main_trace_frame_into(step << lde_shift, &mut main_frame); + trace.read_aux_trace_frame_into(step << lde_shift, &mut aux_frame); + + let l_cur = aux_frame.current()[l_col_idx]; + let s_cur = aux_frame.current()[s_col_idx]; + let s_nxt = aux_frame.next()[s_col_idx]; + + evaluator.build_query(&main_frame, &[], &mut query); + let batched_query = self.gkr_data.compute_batched_query(&query); + + let rhs = s_cur - mean + batched_query * l_cur; + let lhs = s_nxt; + + let divisor_at_step = + s_col_constraint_divisor[step % (domain.trace_to_ce_blowup())]; + + (rhs - lhs) * self.s_col_composition_coefficient.mul_base(divisor_at_step) + }; + + combined_evaluations_acc[step] += + lagrange_combined_evaluations + s_col_combined_evaluation; } } @@ -293,3 +341,22 @@ impl TransitionDivisorEvaluator { - E::BaseField::ONE } } + +/// Computes the evaluations of the s-column divisor. +/// +/// The divisor for the s-column is $X^n - 1$ where $n$ is the trace length. This means that +/// we need only compute `ce_blowup` many values and thus only that many exponentiations. +fn compute_s_col_divisor( + domain: &StarkDomain, + trace_length: usize, +) -> Vec { + let degree = trace_length as u32; + let mut result = Vec::with_capacity(domain.trace_to_ce_blowup()); + + for row in 0..domain.trace_to_ce_blowup() { + let x = domain.get_ce_x_at(row).exp(degree.into()) - E::BaseField::ONE; + + result.push(x); + } + batch_inversion(&result) +} diff --git a/prover/src/constraints/evaluator/mod.rs b/prover/src/constraints/evaluator/mod.rs index da8a166c2..0ff6916f8 100644 --- a/prover/src/constraints/evaluator/mod.rs +++ b/prover/src/constraints/evaluator/mod.rs @@ -14,7 +14,7 @@ pub use default::DefaultConstraintEvaluator; mod boundary; use boundary::BoundaryConstraints; -mod lagrange; +mod logup_gkr; mod periodic_table; use periodic_table::PeriodicValueTable; diff --git a/prover/src/lib.rs b/prover/src/lib.rs index 703f19d8c..b62df14d8 100644 --- a/prover/src/lib.rs +++ b/prover/src/lib.rs @@ -346,7 +346,7 @@ pub trait Prover { }; trace_polys - .add_aux_segment(aux_segment_polys, air.context().lagrange_kernel_aux_column_idx()); + .add_aux_segment(aux_segment_polys, air.context().lagrange_kernel_column_idx()); Some(AuxTraceWithMetadata { aux_trace, aux_rand_elements, gkr_proof }) } else { diff --git a/prover/src/trace/mod.rs b/prover/src/trace/mod.rs index 8d7c999bf..5fef475e6 100644 --- a/prover/src/trace/mod.rs +++ b/prover/src/trace/mod.rs @@ -3,7 +3,10 @@ // This source code is licensed under the MIT license found in the // LICENSE file in the root directory of this source tree. -use air::{Air, AuxRandElements, EvaluationFrame, LagrangeKernelBoundaryConstraint, TraceInfo}; +use air::{ + Air, AuxRandElements, EvaluationFrame, LagrangeKernelBoundaryConstraint, LogUpGkrEvaluator, + TraceInfo, +}; use math::{polynom, FieldElement, StarkField}; use sumcheck::GkrCircuitProof; @@ -139,7 +142,7 @@ pub trait Trace: Sized { } // then, check the Lagrange kernel assertion, if any - if let Some(lagrange_kernel_col_idx) = air.context().lagrange_kernel_aux_column_idx() { + if let Some(lagrange_kernel_col_idx) = air.context().lagrange_kernel_column_idx() { let boundary_constraint_assertion_value = LagrangeKernelBoundaryConstraint::assertion_value( aux_rand_elements @@ -222,19 +225,27 @@ pub trait Trace: Sized { x *= g; } - // evaluate transition constraints for Lagrange kernel column (if any) and make sure - // they all evaluate to zeros - if let Some(col_idx) = air.context().lagrange_kernel_aux_column_idx() { + // evaluate transition constraints for Lagrange kernel column and s-column, when LogUp-GKR + // is enabled, and make sure they all evaluate to zeros + if air.context().logup_gkr_enabled() { let aux_trace_with_metadata = aux_trace_with_metadata.expect("expected aux trace to be present"); let aux_trace = &aux_trace_with_metadata.aux_trace; let aux_rand_elements = &aux_trace_with_metadata.aux_rand_elements; - - let c = aux_trace.get_column(col_idx); - let v = self.length().ilog2() as usize; - let r = aux_rand_elements.lagrange().expect("expected Lagrange column to be present"); - - // Loop over every constraint + let l_col_idx = air + .context() + .trace_info() + .lagrange_kernel_column_idx() + .expect("should not be None"); + let s_col_idx = air.context().trace_info().s_column_idx().expect("should not be None"); + + let c = aux_trace.get_column(l_col_idx); + let trace_length = self.length(); + let v = trace_length.ilog2() as usize; + let gkr_data = aux_rand_elements.gkr_data().expect("should not be None"); + let r = gkr_data.lagrange_kernel_rand_elements(); + + // Loop over every Lagrange kernel constraint for constraint_idx in 1..v + 1 { let domain_step = 2_usize.pow((v - constraint_idx + 1) as u32); let domain_half_step = 2_usize.pow((v - constraint_idx) as u32); @@ -254,6 +265,36 @@ pub trait Trace: Sized { ); } } + + // Validate the s-column constraint + let evaluator = air.get_logup_gkr_evaluator(); + let mut aux_frame = EvaluationFrame::new(self.aux_trace_width()); + + let c = gkr_data.compute_batched_claim(); + let mean = c / E::from(E::BaseField::from(trace_length as u32)); + + let mut query = vec![E::BaseField::ZERO; evaluator.get_oracles().len()]; + for step in 0..self.length() { + self.read_main_frame(step, &mut main_frame); + read_aux_frame(aux_trace, step, &mut aux_frame); + + let l_cur = aux_frame.current()[l_col_idx]; + let s_cur = aux_frame.current()[s_col_idx]; + let s_nxt = aux_frame.next()[s_col_idx]; + + evaluator.build_query(&main_frame, &[], &mut query); + let batched_query = gkr_data.compute_batched_query(&query); + + let rhs = s_cur - mean + batched_query * l_cur; + let lhs = s_nxt; + + let evaluation = rhs - lhs; + + assert!( + evaluation == E::ZERO, + "s-column transition constraint did not evaluate to ZERO at step {step}" + ); + } } } } diff --git a/sumcheck/src/lib.rs b/sumcheck/src/lib.rs index 5beac9bb9..b7f670a9d 100644 --- a/sumcheck/src/lib.rs +++ b/sumcheck/src/lib.rs @@ -26,7 +26,7 @@ mod univariate; pub use univariate::{CompressedUnivariatePoly, CompressedUnivariatePolyEvals}; mod multilinear; -pub use multilinear::{EqFunction, MultiLinearPoly}; +pub use multilinear::{inner_product, EqFunction, MultiLinearPoly}; /// Represents an opening claim at an evaluation point against a batch of oracles. /// diff --git a/verifier/src/composer.rs b/verifier/src/composer.rs index 5f10ef79f..ae20f4586 100644 --- a/verifier/src/composer.rs +++ b/verifier/src/composer.rs @@ -43,7 +43,7 @@ impl DeepComposer { x_coordinates, z: [z, z * E::from(g_trace)], g_trace, - lagrange_kernel_column_idx: air.context().lagrange_kernel_aux_column_idx(), + lagrange_kernel_column_idx: air.context().lagrange_kernel_column_idx(), } } diff --git a/verifier/src/evaluator.rs b/verifier/src/evaluator.rs index b26f7b926..a226ec9c8 100644 --- a/verifier/src/evaluator.rs +++ b/verifier/src/evaluator.rs @@ -7,7 +7,7 @@ use alloc::vec::Vec; use air::{ Air, AuxRandElements, ConstraintCompositionCoefficients, EvaluationFrame, - LagrangeKernelEvaluationFrame, + LagrangeKernelEvaluationFrame, LogUpGkrEvaluator, }; use math::{polynom, FieldElement}; @@ -89,35 +89,50 @@ pub fn evaluate_constraints>( // 3 ----- evaluate Lagrange kernel constraints ------------------------------------ if let Some(lagrange_kernel_column_frame) = lagrange_kernel_frame { + let logup_gkr_evaluator = air.get_logup_gkr_evaluator(); + let lagrange_coefficients = composition_coefficients .lagrange .expect("expected Lagrange kernel composition coefficients to be present"); - let air::GkrData { - lagrange_kernel_eval_point: lagrange_kernel_evaluation_point, - openings_combining_randomness: _, - openings: _, - oracles: _, - } = aux_rand_elements + + let gkr_data = aux_rand_elements .expect("expected aux rand elements to be present") .gkr_data() .expect("expected LogUp-GKR rand elements to be present"); // Lagrange kernel constraints - let lagrange_constraints = air - .get_lagrange_kernel_constraints( - lagrange_coefficients, - &lagrange_kernel_evaluation_point, - ) - .expect("expected Lagrange kernel constraints to be present"); + let lagrange_constraints = logup_gkr_evaluator.get_lagrange_kernel_constraints( + lagrange_coefficients, + &gkr_data.lagrange_kernel_eval_point, + ); result += lagrange_constraints.transition.evaluate_and_combine::( lagrange_kernel_column_frame, - &lagrange_kernel_evaluation_point, + &gkr_data.lagrange_kernel_eval_point, x, ); - result += lagrange_constraints.boundary.evaluate_at(x, lagrange_kernel_column_frame); + + // s-column constraints + + let s_col_idx = air.trace_info().s_column_idx().expect("s-column should be present"); + + let aux_trace_frame = + aux_trace_frame.as_ref().expect("expected aux rand elements to be present"); + + let s_cur = aux_trace_frame.current()[s_col_idx]; + let s_nxt = aux_trace_frame.next()[s_col_idx]; + let l_cur = lagrange_kernel_column_frame.inner()[0]; + + let s_column_cc = composition_coefficients + .s_col + .expect("expected constraint composition coefficient for s-column to be present"); + + let s_column_constraint = + logup_gkr_evaluator.get_s_column_constraints(gkr_data, s_column_cc); + + result += s_column_constraint.evaluate(air, main_trace_frame, s_cur, s_nxt, l_cur, x); } result diff --git a/winterfell/src/tests.rs b/winterfell/src/tests.rs index 858f35574..99e52971f 100644 --- a/winterfell/src/tests.rs +++ b/winterfell/src/tests.rs @@ -22,7 +22,7 @@ use crate::{ #[test] fn test_logup_gkr() { - let aux_trace_width = 2; + let aux_trace_width = 1; let trace = LogUpGkrSimple::new(2_usize.pow(7), aux_trace_width); let prover = LogUpGkrSimpleProver::new(aux_trace_width); @@ -269,7 +269,7 @@ impl LogUpGkrSimpleProver { fn new(aux_trace_width: usize) -> Self { Self { aux_trace_width, - options: ProofOptions::new(1, 8, 0, FieldExtension::None, 2, 1), + options: ProofOptions::new(1, 8, 0, FieldExtension::Quadratic, 2, 1), } } }