Interactions prove/verify via rational sumcheck

Author: zlangley

crates/stark-backend/src/air_builders/debug/check_constraints.rs

@@ -1,15 +1,21 @@
-use itertools::izip;
+use itertools::{izip, Itertools};
 use p3_air::BaseAir;
 use p3_field::{Field, FieldAlgebra};
-use p3_matrix::{dense::RowMajorMatrixView, stack::VerticalPair, Matrix};
+use p3_matrix::{
+    dense::{RowMajorMatrix, RowMajorMatrixView},
+    stack::VerticalPair,
+    Matrix,
+};
 use p3_maybe_rayon::prelude::*;
+use p3_util::log2_strict_usize;
 
 use crate::{
     air_builders::debug::DebugConstraintBuilder,
     config::{StarkGenericConfig, Val},
     interaction::{
         debug::{generate_logical_interactions, LogicalInteractions},
-        InteractionType, RapPhaseSeqKind, SymbolicInteraction,
+        gkr_log_up::fold_multilinear_lagrange_col_constraints,
+        InteractionType, RapPhaseSeq, RapPhaseSeqKind, SymbolicInteraction,
     },
     rap::{PartitionedBaseAir, Rap},
 };
@@ -69,7 +75,7 @@ pub fn check_constraints<R, SC>(
             .iter()
             .map(|mat| (mat.row_slice(i), mat.row_slice(i_next)))
             .collect::<Vec<_>>();
-        let after_challenge = after_challenge_row_pair
+        let after_challenge_pairs = after_challenge_row_pair
             .iter()
             .map(|(local, next)| {
                 VerticalPair::new(
@@ -87,7 +93,7 @@ pub fn check_constraints<R, SC>(
                 RowMajorMatrixView::new_row(preprocessed_next.as_slice()),
             ),
             partitioned_main,
-            after_challenge,
+            after_challenge: after_challenge_pairs,
             challenges,
             public_values,
             exposed_values_after_challenge,
@@ -106,9 +112,53 @@ pub fn check_constraints<R, SC>(
         }
 
         rap.eval(&mut builder);
+
+        if matches!(SC::RapPhaseSeq::KIND, RapPhaseSeqKind::GkrLogUp) {
+            check_gkr_log_up_adapter_constraints_for_row::<SC>(after_challenge, challenges, i);
+        }
     });
 }
 
+fn check_gkr_log_up_adapter_constraints_for_row<SC: StarkGenericConfig>(
+    after_challenge: &[RowMajorMatrixView<SC::Challenge>],
+    challenges: &[Vec<SC::Challenge>],
+    i: usize,
+) {
+    if after_challenge.is_empty() {
+        return;
+    }
+    assert_eq!(after_challenge.len(), 1);
+    assert_eq!(challenges.len(), 1);
+
+    let after_challenge = &after_challenge[0];
+    let challenges = &challenges[0];
+    let height = after_challenge.height();
+
+    let log_height = log2_strict_usize(height);
+    let indices = std::iter::once(i).chain((0..log_height).map(|j| (i + (1 << j)) % height));
+    let after_challenge_window = RowMajorMatrix::new(
+        indices.flat_map(|i| after_challenge.row(i)).collect_vec(),
+        after_challenge.width(),
+    );
+
+    let r = &challenges[challenges.len() - log_height..];
+    let mut accumulator = SC::Challenge::ZERO;
+    let alpha = SC::Challenge::TWO;
+
+    let is_cyclic_row = (0..=log_height)
+        .map(|k| SC::Challenge::from_bool(i & ((1 << (log_height - k)) - 1) == 0))
+        .collect_vec();
+    fold_multilinear_lagrange_col_constraints(
+        &mut accumulator,
+        alpha,
+        &after_challenge_window,
+        &is_cyclic_row,
+        r,
+        0,
+    );
+    assert_eq!(accumulator, SC::Challenge::ZERO);
+}
+
 pub fn check_logup<F: Field>(
     air_names: &[String],
     interactions: &[&[SymbolicInteraction<F>]],

crates/stark-backend/src/air_builders/debug/mod.rs

@@ -221,8 +221,8 @@ where
 impl<SC: StarkGenericConfig> InteractionPhaseAirBuilder for DebugConstraintBuilder<'_, SC> {
     fn finalize_interactions(&mut self) {}
 
-    fn interaction_chunk_size(&self) -> usize {
-        0
+    fn interaction_chunk_size(&self) -> Option<usize> {
+        None
     }
 
     fn rap_phase_seq_kind(&self) -> RapPhaseSeqKind {

crates/stark-backend/src/air_builders/symbolic/mod.rs

@@ -43,6 +43,7 @@ pub struct SymbolicConstraints<F> {
     /// the prover for after challenge trace generation, and some partial
     /// information may be used by the verifier.
     ///
+    /// FIXME[zach]: False for GKR; there, some constraints depend on height of the trace and therefore cannot be precomputed.
     /// **However**, any contributions to the quotient polynomial from
     /// logup are already included in `constraints` and do not need to
     /// be separately calculated from `interactions`.
@@ -110,7 +111,7 @@ where
         num_challenges_to_sample,
         num_exposed_values_after_challenge,
         rap_phase_seq_kind,
-        interaction_chunk_size,
+        Some(interaction_chunk_size),
     );
     Rap::eval(rap, &mut builder);
     builder
@@ -127,7 +128,7 @@ pub struct SymbolicRapBuilder<F> {
     exposed_values_after_challenge: Vec<Vec<SymbolicVariable<F>>>,
     constraints: Vec<SymbolicExpression<F>>,
     interactions: Vec<SymbolicInteraction<F>>,
-    interaction_chunk_size: usize,
+    interaction_chunk_size: Option<usize>,
     rap_phase_seq_kind: RapPhaseSeqKind,
     trace_width: TraceWidth,
 }
@@ -142,7 +143,7 @@ impl<F: Field> SymbolicRapBuilder<F> {
         num_challenges_to_sample: &[usize],
         num_exposed_values_after_challenge: &[usize],
         rap_phase_seq_kind: RapPhaseSeqKind,
-        interaction_chunk_size: usize,
+        interaction_chunk_size: Option<usize>,
     ) -> Self {
         let preprocessed_width = width.preprocessed.unwrap_or(0);
         let prep_values = [0, 1]
@@ -398,7 +399,15 @@ impl<F: Field> InteractionPhaseAirBuilder for SymbolicRapBuilder<F> {
             assert!(self.challenges.is_empty());
             assert!(self.exposed_values_after_challenge.is_empty());
 
-            let perm_width = num_interactions.div_ceil(self.interaction_chunk_size) + 1;
+            let perm_width = match self.rap_phase_seq_kind {
+                RapPhaseSeqKind::StarkLogUp => {
+                    let interaction_chunk_size = self
+                        .interaction_chunk_size()
+                        .expect("interaction chunk size should be set for StarkLogUp");
+                    num_interactions.div_ceil(interaction_chunk_size) + 1
+                }
+                RapPhaseSeqKind::GkrLogUp => 2,
+            };
             self.after_challenge = Self::new_after_challenge(&[perm_width]);
 
             let phases_shapes = self.rap_phase_seq_kind.shape();
@@ -410,7 +419,7 @@ impl<F: Field> InteractionPhaseAirBuilder for SymbolicRapBuilder<F> {
         }
     }
 
-    fn interaction_chunk_size(&self) -> usize {
+    fn interaction_chunk_size(&self) -> Option<usize> {
         self.interaction_chunk_size
     }
 

crates/stark-backend/src/config.rs

@@ -91,15 +91,15 @@ pub type PackedChallenge<SC> =
 #[derive(Debug)]
 pub struct StarkConfig<Pcs, RapPhaseSeq, Challenge, Challenger> {
     pcs: Pcs,
-    rap_phase: RapPhaseSeq,
+    rap_phase_seq: RapPhaseSeq,
     _phantom: PhantomData<(Challenge, Challenger)>,
 }
 
 impl<Pcs, RapPhaseSeq, Challenge, Challenger> StarkConfig<Pcs, RapPhaseSeq, Challenge, Challenger> {
-    pub const fn new(pcs: Pcs, rap_phase: RapPhaseSeq) -> Self {
+    pub const fn new(pcs: Pcs, rap_phase_seq: RapPhaseSeq) -> Self {
         Self {
             pcs,
-            rap_phase,
+            rap_phase_seq,
             _phantom: PhantomData,
         }
     }
@@ -128,7 +128,7 @@ where
         &self.pcs
     }
     fn rap_phase_seq(&self) -> &Self::RapPhaseSeq {
-        &self.rap_phase
+        &self.rap_phase_seq
     }
 }
 

crates/stark-backend/src/gkr/prover.rs

@@ -7,7 +7,7 @@ use std::{
 
 use itertools::Itertools;
 use p3_challenger::FieldChallenger;
-use p3_field::Field;
+use p3_field::{ExtensionField, Field};
 use thiserror::Error;
 
 use crate::{
@@ -327,10 +327,10 @@ pub struct NotConstantPolyError;
 ///
 /// The input layers should be committed to the channel before calling this function.
 // GKR algorithm: <https://people.cs.georgetown.edu/jthaler/ProofsArgsAndZK.pdf> (page 64)
-pub fn prove_batch<F: Field>(
+pub fn prove_batch<F: Field, EF: ExtensionField<F>>(
     challenger: &mut impl FieldChallenger<F>,
-    input_layer_by_instance: Vec<Layer<F>>,
-) -> (GkrBatchProof<F>, GkrArtifact<F>) {
+    input_layer_by_instance: Vec<Layer<EF>>,
+) -> (GkrBatchProof<EF>, GkrArtifact<EF>) {
     let n_instances = input_layer_by_instance.len();
     let n_layers_by_instance = input_layer_by_instance
         .iter()
@@ -366,12 +366,14 @@ pub fn prove_batch<F: Field>(
 
         // Seed the channel with layer claims.
         for claims_to_verify in claims_to_verify_by_instance.iter().flatten() {
-            challenger.observe_slice(claims_to_verify);
+            for claim in claims_to_verify {
+                challenger.observe_ext_element(*claim);
+            }
         }
 
         let eq_evals = HypercubeEqEvals::eval(&ood_point);
-        let sumcheck_alpha = challenger.sample();
-        let instance_lambda = challenger.sample();
+        let sumcheck_alpha: EF = challenger.sample_ext_element();
+        let instance_lambda: EF = challenger.sample_ext_element();
 
         let mut sumcheck_oracles = Vec::new();
         let mut sumcheck_claims = Vec::new();
@@ -385,7 +387,7 @@ pub fn prove_batch<F: Field>(
                 sumcheck_oracles.push(GkrMultivariatePolyOracle {
                     eq_evals: &eq_evals,
                     input_layer: layer,
-                    eq_fixed_var_correction: F::ONE,
+                    eq_fixed_var_correction: EF::ONE,
                     lambda: instance_lambda,
                 });
                 sumcheck_claims.push(random_linear_combination(claims_to_verify, instance_lambda));
@@ -417,12 +419,14 @@ pub fn prove_batch<F: Field>(
         // Seed the channel with the layer masks.
         for (&instance, mask) in zip(&sumcheck_instances, &masks) {
             for column in mask.columns() {
-                challenger.observe_slice(column);
+                for el in column {
+                    challenger.observe_ext_element(*el);
+                }
             }
             layer_masks_by_instance[instance].push(mask.clone());
         }
 
-        let challenge = challenger.sample();
+        let challenge: EF = challenger.sample_ext_element();
         ood_point = sumcheck_ood_point;
         ood_point.push(challenge);
 
@@ -501,7 +505,7 @@ pub fn correct_sum_as_poly_in_first_variable<F: Field>(
     let r_at_2 = f_at_2 * hypercube_eq(&[F::TWO], &[y[n - k]]) * a_const;
 
     // Interpolate.
-    UnivariatePolynomial::from_interpolation(&[
+    UnivariatePolynomial::from_points(&[
         (F::ZERO, r_at_0),
         (F::ONE, r_at_1),
         (F::TWO, r_at_2),

crates/stark-backend/src/gkr/tests.rs

@@ -20,8 +20,8 @@ fn test_batch() -> Result<(), GkrError<BabyBear>> {
     let engine = default_engine();
     let mut rng = create_seeded_rng();
 
-    let col0 = Mle::new((0..1 << LOG_N).map(|_| rng.gen()).collect_vec());
-    let col1 = Mle::new((0..1 << LOG_N).map(|_| rng.gen()).collect_vec());
+    let col0 = Mle::from_vec((0..1 << LOG_N).map(|_| rng.gen()).collect_vec());
+    let col1 = Mle::from_vec((0..1 << LOG_N).map(|_| rng.gen()).collect_vec());
 
     let product0 = col0.iter().copied().product();
     let product1 = col1.iter().copied().product();
@@ -62,8 +62,8 @@ fn test_batch_with_different_sizes() -> Result<(), GkrError<BabyBear>> {
     const LOG_N0: usize = 5;
     const LOG_N1: usize = 7;
 
-    let col0 = Mle::new((0..1 << LOG_N0).map(|_| rng.gen()).collect());
-    let col1 = Mle::new((0..1 << LOG_N1).map(|_| rng.gen()).collect());
+    let col0 = Mle::from_vec((0..1 << LOG_N0).map(|_| rng.gen()).collect());
+    let col1 = Mle::from_vec((0..1 << LOG_N1).map(|_| rng.gen()).collect());
 
     let product0 = col0.iter().copied().product();
     let product1 = col1.iter().copied().product();
@@ -106,7 +106,7 @@ fn test_grand_product() -> Result<(), GkrError<BabyBear>> {
 
     let values = (0..N).map(|_| rng.gen()).collect_vec();
     let product = values.iter().copied().product();
-    let col = Mle::<BabyBear>::new(values);
+    let col = Mle::<BabyBear>::from_vec(values);
     let input_layer = Layer::GrandProduct(col.clone());
     let (proof, _) = gkr::prove_batch(&mut engine.new_challenger(), vec![input_layer]);
 
@@ -137,8 +137,8 @@ fn test_logup_with_generic_trace() -> Result<(), GkrError<BabyBear>> {
     let sum: Fraction<F> = zip(&numerator_values, &denominator_values)
         .map(|(&n, &d)| Fraction::new(n, d))
         .sum();
-    let numerators = Mle::<F>::new(numerator_values);
-    let denominators = Mle::<F>::new(denominator_values);
+    let numerators = Mle::<F>::from_vec(numerator_values);
+    let denominators = Mle::<F>::from_vec(denominator_values);
     let top_layer = Layer::LogUpGeneric {
         numerators: numerators.clone(),
         denominators: denominators.clone(),
@@ -178,7 +178,7 @@ fn test_logup_with_singles_trace() -> Result<(), GkrError<BabyBear>> {
         .iter()
         .map(|&d| Fraction::new(F::ONE, d))
         .sum();
-    let denominators = Mle::new(denominator_values);
+    let denominators = Mle::from_vec(denominator_values);
     let top_layer = Layer::LogUpSingles {
         denominators: denominators.clone(),
     };
@@ -214,8 +214,8 @@ fn test_logup_with_multiplicities_trace() -> Result<(), GkrError<BabyBear>> {
     let sum: Fraction<BabyBear> = zip(&numerator_values, &denominator_values)
         .map(|(&n, &d)| Fraction::new(n, d))
         .sum();
-    let numerators = Mle::new(numerator_values);
-    let denominators = Mle::new(denominator_values);
+    let numerators = Mle::from_vec(numerator_values);
+    let denominators = Mle::from_vec(denominator_values);
     let top_layer = Layer::LogUpMultiplicities {
         numerators: numerators.clone(),
         denominators: denominators.clone(),

crates/stark-backend/src/gkr/types.rs

@@ -1,6 +1,7 @@
 use std::ops::Index;
 
 use p3_field::Field;
+use serde::{Deserialize, Serialize};
 use thiserror::Error;
 
 use crate::{
@@ -12,6 +13,7 @@ use crate::{
 };
 
 /// Batch GKR proof.
+#[derive(Clone, Serialize, Deserialize)]
 pub struct GkrBatchProof<F> {
     /// Sum-check proof for each layer.
     pub sumcheck_proofs: Vec<SumcheckProof<F>>,
@@ -22,6 +24,7 @@ pub struct GkrBatchProof<F> {
 }
 
 /// Values of interest obtained from the execution of the GKR protocol.
+#[derive(Clone, Serialize, Deserialize)]
 pub struct GkrArtifact<F> {
     /// Out-of-domain (OOD) point for evaluating columns in the input layer.
     pub ood_point: Vec<F>,
@@ -32,7 +35,7 @@ pub struct GkrArtifact<F> {
 }
 
 /// Stores two evaluations of each column in a GKR layer.
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct GkrMask<F> {
     columns: Vec<[F; 2]>,
 }
@@ -169,7 +172,7 @@ impl<F: Field> Layer<F> {
             .chunks_exact(2) // Process in chunks of 2 elements
             .map(|chunk| chunk[0] * chunk[1]) // Multiply each pair
             .collect();
-        Layer::GrandProduct(Mle::new(res))
+        Layer::GrandProduct(Mle::from_vec(res))
     }
 
     fn next_logup_layer(numerators: MleExpr<'_, F>, denominators: &Mle<F>) -> Layer<F> {
@@ -186,8 +189,8 @@ impl<F: Field> Layer<F> {
         }
 
         Layer::LogUpGeneric {
-            numerators: Mle::new(next_numerators),
-            denominators: Mle::new(next_denominators),
+            numerators: Mle::from_vec(next_numerators),
+            denominators: Mle::from_vec(next_denominators),
         }
     }