GKR backend for LogUp-GKR

.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,13 +3,11 @@
 // 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.
 ///
@@ -18,12 +16,12 @@ use super::lagrange::LagrangeKernelRandElements;
 /// - 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
 #[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> {
+impl<E: FieldElement> 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 {
@@ -32,7 +30,7 @@ impl<E> AuxRandElements<E> {
 
     /// Creates a new [`AuxRandElements`], where the auxiliary trace contains 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 {
+    pub fn new_with_gkr(rand_elements: Vec<E>, gkr: GkrData<E>) -> Self {
         Self { rand_elements, gkr: Some(gkr) }
     }
 
@@ -43,7 +41,7 @@ impl<E> AuxRandElements<E> {
 
     /// 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,6 +50,10 @@ 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())
     }
+
+    pub fn gkr_data(&self) -> Option<GkrData<E>> {
+        self.gkr.clone()
+    }
 }
 
 /// Holds all the random elements needed when using GKR to accelerate LogUp.
@@ -65,70 +67,66 @@ impl<E> AuxRandElements<E> {
 /// the Lagrange kernel random elements. Those claims are (linearly) combined into one using the
 /// openings combining randomness.
 #[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/coefficients.rs

@@ -32,6 +32,7 @@ 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
@@ -109,4 +110,6 @@ pub struct DeepCompositionCoefficients<E: FieldElement> {
     pub constraints: Vec<E>,
     /// Lagrange kernel trace polynomial composition coefficient $\gamma$.
     pub lagrange: Option<E>,
+    /// TODO
+    pub s_col: Option<E>,
 }

air/src/air/context.rs

@@ -21,11 +21,11 @@ pub struct AirContext<B: StarkField> {
     pub(super) aux_transition_constraint_degrees: Vec<TransitionConstraintDegree>,
     pub(super) num_main_assertions: usize,
     pub(super) num_aux_assertions: usize,
-    pub(super) lagrange_kernel_aux_column_idx: Option<usize>,
     pub(super) ce_blowup_factor: usize,
     pub(super) trace_domain_generator: B,
     pub(super) lde_domain_generator: B,
     pub(super) num_transition_exemptions: usize,
+    pub(super) logup_gkr: bool,
 }
 
 impl<B: StarkField> AirContext<B> {
@@ -62,7 +62,6 @@ impl<B: StarkField> AirContext<B> {
             Vec::new(),
             num_assertions,
             0,
-            None,
             options,
         )
     }
@@ -95,7 +94,6 @@ impl<B: StarkField> AirContext<B> {
         aux_transition_constraint_degrees: Vec<TransitionConstraintDegree>,
         num_main_assertions: usize,
         num_aux_assertions: usize,
-        lagrange_kernel_aux_column_idx: Option<usize>,
         options: ProofOptions,
     ) -> Self {
         assert!(
@@ -104,11 +102,11 @@ impl<B: StarkField> AirContext<B> {
         );
         assert!(num_main_assertions > 0, "at least one assertion must be specified");
 
-        if trace_info.is_multi_segment() {
+        if trace_info.is_multi_segment() && !trace_info.is_with_logup_gkr() {
             assert!(
-                !aux_transition_constraint_degrees.is_empty(),
-                "at least one transition constraint degree must be specified for the auxiliary trace segment"
-                );
+             !aux_transition_constraint_degrees.is_empty(),
+            "at least one transition constraint degree must be specified for the auxiliary trace segment"
+            );
             assert!(
                 num_aux_assertions > 0,
                 "at least one assertion must be specified against the auxiliary trace segment"
@@ -124,15 +122,6 @@ impl<B: StarkField> AirContext<B> {
             );
         }
 
-        // validate Lagrange kernel aux column, if any
-        if let Some(lagrange_kernel_aux_column_idx) = lagrange_kernel_aux_column_idx {
-            assert!(
-            lagrange_kernel_aux_column_idx == trace_info.get_aux_segment_width() - 1,
-            "Lagrange kernel column should be the last column of the auxiliary trace: index={}, but aux trace width is {}",
-            lagrange_kernel_aux_column_idx, trace_info.get_aux_segment_width()
-            );
-        }
-
         // determine minimum blowup factor needed to evaluate transition constraints by taking
         // the blowup factor of the highest degree constraint
         let mut ce_blowup_factor = 0;
@@ -165,14 +154,34 @@ impl<B: StarkField> AirContext<B> {
             aux_transition_constraint_degrees,
             num_main_assertions,
             num_aux_assertions,
-            lagrange_kernel_aux_column_idx,
             ce_blowup_factor,
             trace_domain_generator: B::get_root_of_unity(trace_length.ilog2()),
             lde_domain_generator: B::get_root_of_unity(lde_domain_size.ilog2()),
             num_transition_exemptions: 1,
+            logup_gkr: false,
         }
     }
 
+    pub fn with_logup_gkr(
+        trace_info: TraceInfo,
+        main_transition_constraint_degrees: Vec<TransitionConstraintDegree>,
+        aux_transition_constraint_degrees: Vec<TransitionConstraintDegree>,
+        num_main_assertions: usize,
+        num_aux_assertions: usize,
+        options: ProofOptions,
+    ) -> Self {
+        let mut air_context = Self::new_multi_segment(
+            trace_info,
+            main_transition_constraint_degrees,
+            aux_transition_constraint_degrees,
+            num_main_assertions,
+            num_aux_assertions,
+            options,
+        );
+        air_context.logup_gkr = true;
+        air_context
+    }
+
     // PUBLIC ACCESSORS
     // --------------------------------------------------------------------------------------------
 
@@ -232,12 +241,16 @@ impl<B: StarkField> AirContext<B> {
 
     /// Returns the index of the auxiliary column which implements the Lagrange kernel, if any
     pub fn lagrange_kernel_aux_column_idx(&self) -> Option<usize> {
-        self.lagrange_kernel_aux_column_idx
+        if self.uses_logup_gkr() {
+            Some(self.trace_info().aux_segment_width() - 1)
+        } else {
+            None
+        }
     }
 
-    /// Returns true if the auxiliary trace segment contains a Lagrange kernel column
-    pub fn has_lagrange_kernel_aux_column(&self) -> bool {
-        self.lagrange_kernel_aux_column_idx().is_some()
+    /// Returns true if LogUp-GKR is enabled.
+    pub fn uses_logup_gkr(&self) -> bool {
+        self.logup_gkr
     }
 
     /// Returns the total number of assertions defined for a computation, excluding the Lagrange

air/src/air/logup_gkr.rs

@@ -5,7 +5,7 @@
 
 use alloc::vec::Vec;
 
-use math::{ExtensionOf, FieldElement, StarkField, ToElements};
+use math::{fields::f64::BaseElement, ExtensionOf, FieldElement, StarkField, ToElements};
 
 use super::EvaluationFrame;
 
@@ -25,7 +25,7 @@ pub trait LogUpGkrEvaluator: Clone + Sync {
 
     /// Gets a list of all oracles involved in LogUp-GKR; this is intended to be used in construction of
     /// MLEs.
-    fn get_oracles(&self) -> Vec<LogUpGkrOracle<Self::BaseField>>;
+    fn get_oracles(&self) -> &[LogUpGkrOracle<Self::BaseField>];
 
     /// Returns the number of random values needed to evaluate a query.
     fn get_num_rand_values(&self) -> usize;
@@ -81,9 +81,55 @@ pub trait LogUpGkrEvaluator: Clone + Sync {
     }
 }
 
+impl LogUpGkrEvaluator for () {
+    type BaseField = BaseElement;
+
+    type PublicInputs = ();
+
+    fn get_oracles(&self) -> &[LogUpGkrOracle<Self::BaseField>] {
+        panic!("LogUpGkrEvaluator method called but LogUp-GKR is not implemented")
+    }
+
+    fn get_num_rand_values(&self) -> usize {
+        panic!("LogUpGkrEvaluator method called but LogUp-GKR is not implemented")
+    }
+
+    fn get_num_fractions(&self) -> usize {
+        panic!("LogUpGkrEvaluator method called but LogUp-GKR is not implemented")
+    }
+
+    fn max_degree(&self) -> usize {
+        panic!("LogUpGkrEvaluator method called but LogUp-GKR is not implemented")
+    }
+
+    fn build_query<E>(&self, _frame: &EvaluationFrame<E>, _periodic_values: &[E], _query: &mut [E])
+    where
+        E: FieldElement<BaseField = Self::BaseField>,
+    {
+        panic!("LogUpGkrEvaluator method called but LogUp-GKR is not implemented")
+    }
+
+    fn evaluate_query<F, E>(
+        &self,
+        _query: &[F],
+        _rand_values: &[E],
+        _numerator: &mut [E],
+        _denominator: &mut [E],
+    ) where
+        F: FieldElement<BaseField = Self::BaseField>,
+        E: FieldElement<BaseField = Self::BaseField> + ExtensionOf<F>,
+    {
+        panic!("LogUpGkrEvaluator method called but LogUp-GKR is not implemented")
+    }
+}
+
 #[derive(Clone, Debug, PartialEq, PartialOrd, Eq, Ord)]
 pub enum LogUpGkrOracle<B: StarkField> {
+    // a column with a given index in the main trace segment
     CurrentRow(usize),
+    // a column with a given index in the main trace segment but shifted upwards
     NextRow(usize),
+    // a virtual periodic column defined by its values in a given cycle. Note that the cycle length
+    // must be a power of 2.
     PeriodicValue(Vec<B>),
 }