Poseidon GKR: Move the matrix multiplications out of the sumchecks

crates/lean_prover/src/common.rs

@@ -9,8 +9,8 @@ use poseidon_circuit::{PoseidonGKRLayers, default_cube_layers};
 use crate::*;
 use lean_vm::*;
 
-pub(crate) const N_COMMITED_CUBES_P16: usize = KOALABEAR_RC16_INTERNAL.len() - 1;
-pub(crate) const N_COMMITED_CUBES_P24: usize = KOALABEAR_RC24_INTERNAL.len() - 1;
+pub(crate) const N_COMMITED_CUBES_P16: usize = KOALABEAR_RC16_INTERNAL.len() - 2;
+pub(crate) const N_COMMITED_CUBES_P24: usize = KOALABEAR_RC24_INTERNAL.len() - 2;
 
 pub fn get_base_dims(
     n_cycles: usize,

crates/lean_prover/witness_generation/src/poseidon_tables.rs

@@ -6,6 +6,7 @@ use p3_field::PrimeCharacteristicRing;
 use p3_koala_bear::{KoalaBearInternalLayerParameters, KoalaBearParameters};
 use p3_monty_31::InternalLayerBaseParameters;
 use poseidon_circuit::{PoseidonGKRLayers, PoseidonWitness, generate_poseidon_witness};
+use tracing::instrument;
 use utils::{padd_with_zero_to_next_power_of_two, transposed_par_iter_mut};
 
 pub fn all_poseidon_16_indexes(poseidons_16: &[WitnessPoseidon16]) -> [Vec<F>; 3] {
@@ -84,6 +85,7 @@ pub fn full_poseidon_indexes_poly(
     all_poseidon_indexes
 }
 
+#[instrument(skip_all)]
 pub fn generate_poseidon_witness_helper<
     const WIDTH: usize,
     const N_COMMITED_CUBES: usize,

crates/poseidon_circuit/src/gkr_layers/compression.rs

@@ -0,0 +1,69 @@
+use multilinear_toolkit::prelude::*;
+use p3_field::ExtensionField;
+
+use crate::{EF, F};
+
+#[derive(Debug)]
+pub struct CompressionComputation<const WIDTH: usize> {
+    pub compressed_output: usize,
+}
+
+impl<NF: ExtensionField<F>, const WIDTH: usize> SumcheckComputation<NF, EF>
+    for CompressionComputation<WIDTH>
+where
+    EF: ExtensionField<NF>,
+{
+    fn degree(&self) -> usize {
+        2
+    }
+
+    fn eval(&self, point: &[NF], alpha_powers: &[EF]) -> EF {
+        debug_assert_eq!(point.len(), WIDTH + 1);
+        let mut res = EF::ZERO;
+        let compressed = point[WIDTH];
+        for i in 0..self.compressed_output {
+            res += alpha_powers[i] * point[i];
+        }
+        for i in self.compressed_output..WIDTH {
+            res += alpha_powers[i] * point[i] * (EF::ONE - compressed);
+        }
+
+        res
+    }
+}
+
+impl<const WIDTH: usize> SumcheckComputationPacked<EF> for CompressionComputation<WIDTH> {
+    fn degree(&self) -> usize {
+        2
+    }
+
+    fn eval_packed_base(&self, point: &[PFPacking<EF>], alpha_powers: &[EF]) -> EFPacking<EF> {
+        debug_assert_eq!(point.len(), WIDTH + 1);
+        let mut res = EFPacking::<EF>::ZERO;
+        let compressed = point[WIDTH];
+        for i in 0..self.compressed_output {
+            res += EFPacking::<EF>::from(alpha_powers[i]) * point[i];
+        }
+        for i in self.compressed_output..WIDTH {
+            res += EFPacking::<EF>::from(alpha_powers[i])
+                * point[i]
+                * (PFPacking::<EF>::ONE - compressed);
+        }
+
+        res
+    }
+
+    fn eval_packed_extension(&self, point: &[EFPacking<EF>], alpha_powers: &[EF]) -> EFPacking<EF> {
+        debug_assert_eq!(point.len(), WIDTH + 1);
+        let mut res = EFPacking::<EF>::ZERO;
+        let compressed = point[WIDTH];
+        for i in 0..self.compressed_output {
+            res += point[i] * alpha_powers[i];
+        }
+        for i in self.compressed_output..WIDTH {
+            res += point[i] * (EFPacking::<EF>::ONE - compressed) * alpha_powers[i];
+        }
+
+        res
+    }
+}

crates/poseidon_circuit/src/gkr_layers/full_round.rs

@@ -1,144 +1,55 @@
-use std::array;
-
 use multilinear_toolkit::prelude::*;
 use p3_field::ExtensionField;
-use p3_koala_bear::{
-    GenericPoseidon2LinearLayersKoalaBear, KoalaBearInternalLayerParameters, KoalaBearParameters,
-};
+use p3_koala_bear::{KoalaBearInternalLayerParameters, KoalaBearParameters};
 use p3_monty_31::InternalLayerBaseParameters;
-use p3_poseidon2::GenericPoseidon2LinearLayers;
 
 use crate::{EF, F};
 
 #[derive(Debug)]
-pub struct FullRoundComputation<const WIDTH: usize, const FIRST: bool> {
-    pub constants: [F; WIDTH],
-    pub compressed_output: Option<usize>,
-}
+pub struct FullRoundComputation<const WIDTH: usize> {}
 
-impl<NF: ExtensionField<F>, const WIDTH: usize, const FIRST: bool> SumcheckComputation<NF, EF>
-    for FullRoundComputation<WIDTH, FIRST>
+impl<NF: ExtensionField<F>, const WIDTH: usize> SumcheckComputation<NF, EF>
+    for FullRoundComputation<WIDTH>
 where
     KoalaBearInternalLayerParameters: InternalLayerBaseParameters<KoalaBearParameters, WIDTH>,
     EF: ExtensionField<NF>,
 {
     fn degree(&self) -> usize {
-        3 + self.compressed_output.is_some() as usize
+        3
     }
 
     fn eval(&self, point: &[NF], alpha_powers: &[EF]) -> EF {
-        debug_assert_eq!(
-            point.len(),
-            WIDTH
-                + if self.compressed_output.is_some() {
-                    1
-                } else {
-                    0
-                }
-        );
-        let mut buff: [NF; WIDTH] = array::from_fn(|j| point[j]);
-        if FIRST {
-            GenericPoseidon2LinearLayersKoalaBear::external_linear_layer(&mut buff);
-        }
-        buff.iter_mut().enumerate().for_each(|(j, val)| {
-            *val = (*val + self.constants[j]).cube();
-        });
-        GenericPoseidon2LinearLayersKoalaBear::external_linear_layer(&mut buff);
+        debug_assert_eq!(point.len(), WIDTH);
         let mut res = EF::ZERO;
-        if let Some(compression_output_width) = self.compressed_output {
-            let compressed = point[WIDTH];
-            for i in 0..compression_output_width {
-                res += alpha_powers[i] * buff[i];
-            }
-            for i in compression_output_width..WIDTH {
-                res += alpha_powers[i] * buff[i] * (EF::ONE - compressed);
-            }
-        } else {
-            for i in 0..WIDTH {
-                res += alpha_powers[i] * buff[i];
-            }
+        for i in 0..WIDTH {
+            res += alpha_powers[i] * point[i].cube();
         }
         res
     }
 }
 
-impl<const WIDTH: usize, const FIRST: bool> SumcheckComputationPacked<EF>
-    for FullRoundComputation<WIDTH, FIRST>
+impl<const WIDTH: usize> SumcheckComputationPacked<EF> for FullRoundComputation<WIDTH>
 where
     KoalaBearInternalLayerParameters: InternalLayerBaseParameters<KoalaBearParameters, WIDTH>,
 {
     fn degree(&self) -> usize {
-        3 + self.compressed_output.is_some() as usize
+        3
     }
 
     fn eval_packed_base(&self, point: &[PFPacking<EF>], alpha_powers: &[EF]) -> EFPacking<EF> {
-        debug_assert_eq!(
-            point.len(),
-            WIDTH
-                + if self.compressed_output.is_some() {
-                    1
-                } else {
-                    0
-                }
-        );
-        let mut buff: [PFPacking<EF>; WIDTH] = array::from_fn(|j| point[j]);
-        if FIRST {
-            GenericPoseidon2LinearLayersKoalaBear::external_linear_layer(&mut buff);
-        }
-        buff.iter_mut().enumerate().for_each(|(j, val)| {
-            *val = (*val + self.constants[j]).cube();
-        });
-        GenericPoseidon2LinearLayersKoalaBear::external_linear_layer(&mut buff);
+        debug_assert_eq!(point.len(), WIDTH);
         let mut res = EFPacking::<EF>::ZERO;
-        if let Some(compression_output_width) = self.compressed_output {
-            let compressed = point[WIDTH];
-            for i in 0..compression_output_width {
-                res += EFPacking::<EF>::from(alpha_powers[i]) * buff[i];
-            }
-            for i in compression_output_width..WIDTH {
-                res += EFPacking::<EF>::from(alpha_powers[i])
-                    * buff[i]
-                    * (PFPacking::<EF>::ONE - compressed);
-            }
-        } else {
-            for j in 0..WIDTH {
-                res += EFPacking::<EF>::from(alpha_powers[j]) * buff[j];
-            }
+        for i in 0..WIDTH {
+            res += EFPacking::<EF>::from(alpha_powers[i]) * point[i].cube();
         }
         res
     }
 
     fn eval_packed_extension(&self, point: &[EFPacking<EF>], alpha_powers: &[EF]) -> EFPacking<EF> {
-        debug_assert_eq!(
-            point.len(),
-            WIDTH
-                + if self.compressed_output.is_some() {
-                    1
-                } else {
-                    0
-                }
-        );
-        let mut buff: [EFPacking<EF>; WIDTH] = array::from_fn(|j| point[j]);
-        if FIRST {
-            GenericPoseidon2LinearLayersKoalaBear::external_linear_layer(&mut buff);
-        }
-        buff.iter_mut().enumerate().for_each(|(j, val)| {
-            *val = (*val + PFPacking::<EF>::from(self.constants[j])).cube();
-        });
-        GenericPoseidon2LinearLayersKoalaBear::external_linear_layer(&mut buff);
+        debug_assert_eq!(point.len(), WIDTH);
         let mut res = EFPacking::<EF>::ZERO;
-        if let Some(compression_output_width) = self.compressed_output {
-            let compressed = point[WIDTH];
-            for i in 0..compression_output_width {
-                res += buff[i] * alpha_powers[i];
-            }
-            for i in compression_output_width..WIDTH {
-                res += buff[i] * (EFPacking::<EF>::ONE - compressed) * alpha_powers[i];
-            }
-        } else {
-            for j in 0..WIDTH {
-                res += buff[j] * alpha_powers[j];
-            }
+        for i in 0..WIDTH {
+            res += point[i].cube() * alpha_powers[i];
         }
         res
     }

crates/poseidon_circuit/src/gkr_layers/mod.rs

@@ -7,6 +7,9 @@ pub use partial_round::*;
 mod batch_partial_rounds;
 pub use batch_partial_rounds::*;
 
+mod compression;
+pub use compression::*;
+
 use p3_koala_bear::{
     KOALABEAR_RC16_EXTERNAL_FINAL, KOALABEAR_RC16_EXTERNAL_INITIAL, KOALABEAR_RC16_INTERNAL,
     KOALABEAR_RC24_EXTERNAL_FINAL, KOALABEAR_RC24_EXTERNAL_INITIAL, KOALABEAR_RC24_INTERNAL,
@@ -16,11 +19,11 @@ use crate::F;
 
 #[derive(Debug)]
 pub struct PoseidonGKRLayers<const WIDTH: usize, const N_COMMITED_CUBES: usize> {
-    pub initial_full_round: FullRoundComputation<WIDTH, true>,
-    pub initial_full_rounds_remaining: Vec<FullRoundComputation<WIDTH, false>>,
-    pub batch_partial_rounds: BatchPartialRounds<WIDTH, N_COMMITED_CUBES>,
-    pub partial_rounds_remaining: Vec<PartialRoundComputation<WIDTH>>,
-    pub final_full_rounds: Vec<FullRoundComputation<WIDTH, false>>,
+    pub initial_full_rounds: Vec<[F; WIDTH]>,
+    pub batch_partial_rounds: Option<BatchPartialRounds<WIDTH, N_COMMITED_CUBES>>,
+    pub partial_rounds_remaining: Vec<F>,
+    pub final_full_rounds: Vec<[F; WIDTH]>,
+    pub compressed_output: Option<usize>,
 }
 
 impl<const WIDTH: usize, const N_COMMITED_CUBES: usize> PoseidonGKRLayers<WIDTH, N_COMMITED_CUBES> {
@@ -54,43 +57,26 @@ impl<const WIDTH: usize, const N_COMMITED_CUBES: usize> PoseidonGKRLayers<WIDTH,
         final_constants: &[[F; WIDTH]],
         compressed_output: Option<usize>,
     ) -> Self {
-        let initial_full_round = FullRoundComputation {
-            constants: initial_constants[0],
-            compressed_output: None,
-        };
-        let initial_full_rounds_remaining = initial_constants[1..]
-            .iter()
-            .map(|&constants| FullRoundComputation {
-                constants,
-                compressed_output: None,
-            })
-            .collect::<Vec<_>>();
-        let batch_partial_rounds = BatchPartialRounds {
-            constants: internal_constants[..N_COMMITED_CUBES].try_into().unwrap(),
-            last_constant: internal_constants[N_COMMITED_CUBES],
+        assert!(N_COMMITED_CUBES < internal_constants.len() - 1); // TODO we could go up to internal_constants.len() in theory
+        let initial_full_rounds = initial_constants.to_vec();
+        let (batch_partial_rounds, partial_rounds_remaining) = if N_COMMITED_CUBES == 0 {
+            (None, internal_constants.to_vec())
+        } else {
+            (
+                Some(BatchPartialRounds {
+                    constants: internal_constants[..N_COMMITED_CUBES].try_into().unwrap(),
+                    last_constant: internal_constants[N_COMMITED_CUBES],
+                }),
+                internal_constants[N_COMMITED_CUBES + 1..].to_vec(),
+            )
         };
-        let partial_rounds_remaining = internal_constants[N_COMMITED_CUBES + 1..]
-            .iter()
-            .map(|&constant| PartialRoundComputation { constant })
-            .collect::<Vec<_>>();
-        let final_full_rounds = final_constants
-            .iter()
-            .enumerate()
-            .map(|(i, &constants)| FullRoundComputation {
-                constants,
-                compressed_output: if i == final_constants.len() - 1 {
-                    compressed_output
-                } else {
-                    None
-                },
-            })
-            .collect::<Vec<_>>();
+        let final_full_rounds = final_constants.to_vec();
         Self {
-            initial_full_round,
-            initial_full_rounds_remaining,
+            initial_full_rounds,
             batch_partial_rounds,
             partial_rounds_remaining,
             final_full_rounds,
+            compressed_output,
         }
     }
 }