Implement GKR backend for LogUp-GKR (#296)

.cargo/katex-header.html

@@ -11,6 +11,7 @@
         renderMathInElement(document.body, {
             fleqn: false,
             macros: {
+                "\\B": "\\mathbb{B}",
                 "\\F": "\\mathbb{F}",
                 "\\G": "\\mathbb{G}",
                 "\\O": "\\mathcal{O}",

air/src/air/aux.rs

@@ -3,47 +3,40 @@
 // This source code is licensed under the MIT license found in the
 // LICENSE file in the root directory of this source tree.
 
-use alloc::{string::ToString, vec::Vec};
+use alloc::vec::Vec;
 
-use crypto::{ElementHasher, RandomCoin, RandomCoinError};
 use math::FieldElement;
-use utils::Deserializable;
 
-use super::lagrange::LagrangeKernelRandElements;
+use super::{lagrange::LagrangeKernelRandElements, LogUpGkrOracle};
 
-/// Holds the randomly generated elements necessary to build the auxiliary trace.
+/// Holds the randomly generated elements used in defining the auxiliary segment of the trace.
 ///
-/// Specifically, [`AuxRandElements`] currently supports 3 types of random elements:
-/// - the ones needed to build the Lagrange kernel column (when using GKR to accelerate LogUp),
-/// - the ones needed to build the "s" auxiliary column (when using GKR to accelerate LogUp),
-/// - the ones needed to build all the other auxiliary columns
+/// Specifically, [`AuxRandElements`] currently supports 2 types of random elements:
+/// - the ones needed to build all the auxiliary columns except for the ones associated
+///   to LogUp-GKR.
+/// - the ones needed to build the "s" and Lagrange kernel auxiliary columns (when using GKR to
+///   accelerate LogUp). These also include additional information needed to evaluate constraints
+///   one these two columns.
 #[derive(Debug, Clone)]
-pub struct AuxRandElements<E> {
+pub struct AuxRandElements<E: FieldElement> {
     rand_elements: Vec<E>,
-    gkr: Option<GkrRandElements<E>>,
+    gkr: Option<GkrData<E>>,
 }
 
-impl<E> AuxRandElements<E> {
-    /// Creates a new [`AuxRandElements`], where the auxiliary trace doesn't contain a Lagrange
-    /// kernel column.
-    pub fn new(rand_elements: Vec<E>) -> Self {
-        Self { rand_elements, gkr: None }
-    }
-
-    /// Creates a new [`AuxRandElements`], where the auxiliary trace contains columns needed when
+impl<E: FieldElement> AuxRandElements<E> {
+    /// Creates a new [`AuxRandElements`], where the auxiliary segment may contain columns needed when
     /// using GKR to accelerate LogUp (i.e. a Lagrange kernel column and the "s" column).
-    pub fn new_with_gkr(rand_elements: Vec<E>, gkr: GkrRandElements<E>) -> Self {
-        Self { rand_elements, gkr: Some(gkr) }
+    pub fn new(rand_elements: Vec<E>, gkr: Option<GkrData<E>>) -> Self {
+        Self { rand_elements, gkr }
     }
-
     /// Returns the random elements needed to build all columns other than the two GKR-related ones.
     pub fn rand_elements(&self) -> &[E] {
         &self.rand_elements
     }
 
     /// Returns the random elements needed to build the Lagrange kernel column.
     pub fn lagrange(&self) -> Option<&LagrangeKernelRandElements<E>> {
-        self.gkr.as_ref().map(|gkr| &gkr.lagrange)
+        self.gkr.as_ref().map(|gkr| &gkr.lagrange_kernel_eval_point)
     }
 
     /// Returns the random values used to linearly combine the openings returned from the GKR proof.
@@ -52,83 +45,97 @@ impl<E> AuxRandElements<E> {
     pub fn gkr_openings_combining_randomness(&self) -> Option<&[E]> {
         self.gkr.as_ref().map(|gkr| gkr.openings_combining_randomness.as_ref())
     }
+
+    /// Returns a collection of data necessary for implementing the univariate IOP for multi-linear
+    /// evaluations of [1] when LogUp-GKR is enabled, else returns a `None`.
+    ///
+    /// [1]: https://eprint.iacr.org/2023/1284
+    pub fn gkr_data(&self) -> Option<GkrData<E>> {
+        self.gkr.clone()
+    }
 }
 
-/// Holds all the random elements needed when using GKR to accelerate LogUp.
+/// Holds all the data needed when using LogUp-GKR in order to build and verify the correctness of
+/// two extra auxiliary columns required for running the univariate IOP for multi-linear
+/// evaluations of [1].
 ///
-/// This consists of two sets of random values:
-/// 1. The Lagrange kernel random elements (expanded on in [`LagrangeKernelRandElements`]), and
+/// This consists of:
+/// 1. The Lagrange kernel random elements (expanded on in [`LagrangeKernelRandElements`]). These
+///    make up the evaluation point of the multi-linear extension polynomials underlying the oracles
+///    in point 4 below.
 /// 2. The "openings combining randomness".
+/// 3. The openings of the multi-linear extension polynomials of the main trace columns involved
+///    in LogUp.
+/// 4. A description of the each of the oracles involved in LogUp.
 ///
-/// After the verifying the LogUp-GKR circuit, the verifier is left with unproven claims provided
-/// nondeterministically by the prover about the evaluations of the MLE of the main trace columns at
-/// the Lagrange kernel random elements. Those claims are (linearly) combined into one using the
-/// openings combining randomness.
+/// After verifying the LogUp-GKR circuit, the verifier is left with unproven claims provided
+/// by the prover about the evaluations of the MLEs of the main trace columns at the evaluation
+/// point defining the Lagrange kernel. Those claims are (linearly) batched into one using the
+/// openings combining randomness and checked against the batched oracles using univariate IOP
+/// for multi-linear evaluations of [1].
+///
+/// [1]: https://eprint.iacr.org/2023/1284
 #[derive(Clone, Debug)]
-pub struct GkrRandElements<E> {
-    lagrange: LagrangeKernelRandElements<E>,
-    openings_combining_randomness: Vec<E>,
+pub struct GkrData<E: FieldElement> {
+    pub lagrange_kernel_eval_point: LagrangeKernelRandElements<E>,
+    pub openings_combining_randomness: Vec<E>,
+    pub openings: Vec<E>,
+    pub oracles: Vec<LogUpGkrOracle<E::BaseField>>,
 }
 
-impl<E> GkrRandElements<E> {
-    /// Constructs a new [`GkrRandElements`] from [`LagrangeKernelRandElements`], and the openings
-    /// combining randomness.
+impl<E: FieldElement> GkrData<E> {
+    /// Constructs a new [`GkrData`] from [`LagrangeKernelRandElements`], the openings combining
+    /// randomness and the LogUp-GKR oracles.
     ///
-    /// See [`GkrRandElements`] for a more detailed description.
+    /// See [`GkrData`] for a more detailed description.
     pub fn new(
-        lagrange: LagrangeKernelRandElements<E>,
+        lagrange_kernel_eval_point: LagrangeKernelRandElements<E>,
         openings_combining_randomness: Vec<E>,
+        openings: Vec<E>,
+        oracles: Vec<LogUpGkrOracle<E::BaseField>>,
     ) -> Self {
-        Self { lagrange, openings_combining_randomness }
+        Self {
+            lagrange_kernel_eval_point,
+            openings_combining_randomness,
+            openings,
+            oracles,
+        }
     }
 
     /// Returns the random elements needed to build the Lagrange kernel column.
     pub fn lagrange_kernel_rand_elements(&self) -> &LagrangeKernelRandElements<E> {
-        &self.lagrange
+        &self.lagrange_kernel_eval_point
     }
 
     /// Returns the random values used to linearly combine the openings returned from the GKR proof.
     pub fn openings_combining_randomness(&self) -> &[E] {
         &self.openings_combining_randomness
     }
-}
 
-/// A trait for verifying a GKR proof.
-///
-/// Specifically, the use case in mind is proving the constraints of a LogUp bus using GKR, as
-/// described in [Improving logarithmic derivative lookups using
-/// GKR](https://eprint.iacr.org/2023/1284.pdf).
-pub trait GkrVerifier {
-    /// The GKR proof.
-    type GkrProof: Deserializable;
-    /// The error that can occur during GKR proof verification.
-    type Error: ToString;
-
-    /// Verifies the GKR proof, and returns the random elements that were used in building
-    /// the Lagrange kernel auxiliary column.
-    fn verify<E, Hasher>(
-        &self,
-        gkr_proof: Self::GkrProof,
-        public_coin: &mut impl RandomCoin<BaseField = E::BaseField, Hasher = Hasher>,
-    ) -> Result<GkrRandElements<E>, Self::Error>
-    where
-        E: FieldElement,
-        Hasher: ElementHasher<BaseField = E::BaseField>;
-}
+    pub fn openings(&self) -> &[E] {
+        &self.openings
+    }
+
+    pub fn oracles(&self) -> &[LogUpGkrOracle<E::BaseField>] {
+        &self.oracles
+    }
+
+    pub fn compute_batched_claim(&self) -> E {
+        self.openings[0]
+            + self
+                .openings
+                .iter()
+                .skip(1)
+                .zip(self.openings_combining_randomness.iter())
+                .fold(E::ZERO, |acc, (a, b)| acc + *a * *b)
+    }
 
-impl GkrVerifier for () {
-    type GkrProof = ();
-    type Error = RandomCoinError;
-
-    fn verify<E, Hasher>(
-        &self,
-        _gkr_proof: Self::GkrProof,
-        _public_coin: &mut impl RandomCoin<BaseField = E::BaseField, Hasher = Hasher>,
-    ) -> Result<GkrRandElements<E>, Self::Error>
-    where
-        E: FieldElement,
-        Hasher: ElementHasher<BaseField = E::BaseField>,
-    {
-        Ok(GkrRandElements::new(LagrangeKernelRandElements::default(), Vec::new()))
+    pub fn compute_batched_query(&self, query: &[E::BaseField]) -> E {
+        E::from(query[0])
+            + query
+                .iter()
+                .skip(1)
+                .zip(self.openings_combining_randomness.iter())
+                .fold(E::ZERO, |acc, (a, b)| acc + b.mul_base(*a))
     }
 }

air/src/air/boundary/mod.rs

@@ -58,8 +58,8 @@ impl<E: FieldElement> BoundaryConstraints<E> {
     ///   coefficients.
     /// * The specified assertions are not valid in the context of the computation (e.g., assertion
     ///   column index is out of bounds).
-    pub fn new(
-        context: &AirContext<E::BaseField>,
+    pub fn new<P>(
+        context: &AirContext<E::BaseField, P>,
         main_assertions: Vec<Assertion<E::BaseField>>,
         aux_assertions: Vec<Assertion<E>>,
         composition_coefficients: &[E],
@@ -88,7 +88,7 @@ impl<E: FieldElement> BoundaryConstraints<E> {
         );
 
         let trace_length = context.trace_info.length();
-        let main_trace_width = context.trace_info.main_trace_width();
+        let main_trace_width = context.trace_info.main_segment_width();
         let aux_trace_width = context.trace_info.aux_segment_width();
 
         // make sure the assertions are valid in the context of their respective trace segments;
@@ -151,9 +151,9 @@ impl<E: FieldElement> BoundaryConstraints<E> {
 
 /// Translates the provided assertions into boundary constraints, groups the constraints by their
 /// divisor, and sorts the resulting groups by the degree adjustment factor.
-fn group_constraints<F, E>(
+fn group_constraints<F, E, P>(
     assertions: Vec<Assertion<F>>,
-    context: &AirContext<F::BaseField>,
+    context: &AirContext<F::BaseField, P>,
     composition_coefficients: &[E],
     inv_g: F::BaseField,
     twiddle_map: &mut BTreeMap<usize, Vec<F::BaseField>>,

air/src/air/coefficients.rs

@@ -27,11 +27,19 @@ use math::FieldElement;
 ///
 /// The coefficients are separated into two lists: one for transition constraints and another one
 /// for boundary constraints. This separation is done for convenience only.
+///
+/// In addition to the above, and when LogUp-GKR is enabled, there are two extra sets of
+/// constraint composition coefficients that are used, namely for:
+///
+/// 1. Lagrange kernel constraints, which include both transition and boundary constraints.
+/// 2. S-column constraint, which is used in implementing the cohomological sum-check argument
+///    of https://eprint.iacr.org/2021/930
 #[derive(Debug, Clone)]
 pub struct ConstraintCompositionCoefficients<E: FieldElement> {
     pub transition: Vec<E>,
     pub boundary: Vec<E>,
     pub lagrange: Option<LagrangeConstraintsCompositionCoefficients<E>>,
+    pub s_col: Option<E>,
 }
 
 /// Stores the constraint composition coefficients for the Lagrange kernel transition and boundary
@@ -83,8 +91,9 @@ pub struct LagrangeConstraintsCompositionCoefficients<E: FieldElement> {
 ///    negligible increase in soundness error. The formula for the updated error can be found in
 ///    Theorem 8 of https://eprint.iacr.org/2022/1216.
 ///
-/// In the case when the trace polynomials contain a trace polynomial corresponding to a Lagrange
-/// kernel column, the above expression of $Y(x)$ includes the additional term given by
+/// In the case when LogUp-GKR is enabled, the trace polynomials contain an additional trace
+/// polynomial corresponding to a Lagrange kernel column and the above expression of $Y(x)$
+/// includes the additional term given by
 ///
 /// $$
 /// \gamma \cdot \frac{T_l(x) - p_S(x)}{Z_S(x)}
@@ -99,8 +108,13 @@ pub struct LagrangeConstraintsCompositionCoefficients<E: FieldElement> {
 /// 4. $p_S(X)$ is the polynomial of minimal degree interpolating the set ${(a, T_l(a)): a \in S}$.
 /// 5. $Z_S(X)$ is the polynomial of minimal degree vanishing over the set $S$.
 ///
-/// Note that, if a Lagrange kernel trace polynomial is present, then $\rho^{+}$ from above should
-/// be updated to be $\rho^{+} := \frac{\kappa + log_2(\nu) + 1}{\nu}$.
+/// Note that when LogUp-GKR is enabled, we also have to take into account an additional column,
+/// called s-column throughout, used in implementing the univariate IOP for multi-linear evaluation.
+/// This means that we need and additional random value, in addition to $\gamma$ above, when
+/// LogUp-GKR is enabled.
+///
+/// Note that, when LogUp-GKR is enabled, $\rho^{+}$ from above should be updated to be
+/// $\rho^{+} := \frac{\kappa + log_2(\nu) + 1}{\nu}$.
 #[derive(Debug, Clone)]
 pub struct DeepCompositionCoefficients<E: FieldElement> {
     /// Trace polynomial composition coefficients $\alpha_i$.
@@ -109,4 +123,6 @@ pub struct DeepCompositionCoefficients<E: FieldElement> {
     pub constraints: Vec<E>,
     /// Lagrange kernel trace polynomial composition coefficient $\gamma$.
     pub lagrange: Option<E>,
+    /// S-column trace polynomial composition coefficient.
+    pub s_col: Option<E>,
 }