feat: interactions via rational sumcheck

Author: zlangley

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

@@ -1,14 +1,20 @@
-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},
+        gkr_log_up::fold_multilinear_lagrange_col_constraints,
         RapPhaseSeqKind, SymbolicInteraction,
     },
     rap::{PartitionedBaseAir, Rap},
@@ -87,9 +93,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: &[Vec<SymbolicInteraction<F>>],
@@ -98,7 +148,7 @@ pub fn check_logup<F: Field>(
     public_values: &[Vec<F>],
 ) {
     let mut logical_interactions = LogicalInteractions::<F>::default();
-    for (air_idx, (interactions, preprocessed, partitioned_main, public_values)) in
+    for (air_idx, (interactions, &preprocessed, partitioned_main, public_values)) in
         izip!(interactions, preprocessed, partitioned_main, public_values).enumerate()
     {
         generate_logical_interactions(

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

@@ -41,6 +41,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`.
@@ -399,15 +400,18 @@ impl<F: Field> InteractionPhaseAirBuilder for SymbolicRapBuilder<F> {
             assert!(self.challenges.is_empty());
             assert!(self.exposed_values_after_challenge.is_empty());
 
-            if self.rap_phase_seq_kind == RapPhaseSeqKind::FriLogUp {
-                let interaction_partitions =
-                    find_interaction_chunks(&self.interactions, self.max_constraint_degree)
-                        .interaction_partitions();
-                let num_chunks = interaction_partitions.len();
-                self.interaction_partitions.replace(interaction_partitions);
-                let perm_width = num_chunks + 1;
-                self.after_challenge = Self::new_after_challenge(&[perm_width]);
-            }
+            let perm_width = match self.rap_phase_seq_kind {
+                RapPhaseSeqKind::FriLogUp => {
+                    let interaction_partitions =
+                        find_interaction_chunks(&self.interactions, self.max_constraint_degree)
+                            .interaction_partitions();
+                    let num_chunks = interaction_partitions.len();
+                    self.interaction_partitions.replace(interaction_partitions);
+                    num_chunks + 1
+                }
+                RapPhaseSeqKind::GkrLogUp => 2,
+            };
+            self.after_challenge = Self::new_after_challenge(&[perm_width]);
 
             let phases_shapes = self.rap_phase_seq_kind.shape();
             let phase_shape = phases_shapes.first().unwrap();

crates/stark-backend/src/config.rs

@@ -93,15 +93,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,
         }
     }
@@ -132,7 +132,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),
         }
     }
 

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

@@ -9,6 +9,7 @@ use crate::{
         gate::Gate,
         types::{GkrArtifact, GkrBatchProof, GkrError},
     },
+    p3_field::ExtensionField,
     poly::{multi::hypercube_eq, uni::random_linear_combination},
     sumcheck,
 };
@@ -18,11 +19,11 @@ use crate::{
 /// On successful verification the function returns a [`GkrArtifact`] which stores the out-of-domain
 /// point and claimed evaluations in the input layer columns for each instance at the OOD point.
 /// These claimed evaluations are not checked in this function - hence partial verification.
-pub fn partially_verify_batch<F: Field>(
+pub fn partially_verify_batch<F: Field, EF: ExtensionField<F>>(
     gate_by_instance: Vec<Gate>,
-    proof: &GkrBatchProof<F>,
+    proof: &GkrBatchProof<EF>,
     challenger: &mut impl FieldChallenger<F>,
-) -> Result<GkrArtifact<F>, GkrError<F>> {
+) -> Result<GkrArtifact<EF>, GkrError<EF>> {
     let GkrBatchProof {
         sumcheck_proofs,
         layer_masks_by_instance,
@@ -64,11 +65,13 @@ pub fn partially_verify_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 sumcheck_alpha = challenger.sample();
-        let instance_lambda = challenger.sample();
+        let sumcheck_alpha = challenger.sample_ext_element();
+        let instance_lambda = challenger.sample_ext_element();
 
         let mut sumcheck_claims = Vec::new();
         let mut sumcheck_instances = Vec::new();
@@ -125,12 +128,14 @@ pub fn partially_verify_batch<F: Field>(
             let n_unused = n_layers - instance_n_layers(instance);
             let mask = &layer_masks_by_instance[instance][layer - n_unused];
             for column in mask.columns() {
-                challenger.observe_slice(column);
+                for elt in column {
+                    challenger.observe_ext_element(*elt);
+                }
             }
         }
 
         // Set the OOD evaluation point for layer above.
-        let challenge = challenger.sample();
+        let challenge = challenger.sample_ext_element();
         ood_point = sumcheck_ood_point;
         ood_point.push(challenge);