packing for Grand product sumcheck for Execution

Author: TomWambsgans

crates/lean_prover/src/common.rs

@@ -1,4 +1,4 @@
-use p3_field::BasedVectorSpace;
+use p3_field::{Algebra, BasedVectorSpace};
 use p3_field::{ExtensionField, PrimeCharacteristicRing};
 use p3_util::log2_ceil_usize;
 use packed_pcs::ColDims;
@@ -157,46 +157,57 @@ pub struct PrecompileFootprint {
     pub multilinear_eval_powers: [EF; 6],
 }
 
-impl<N: ExtensionField<PF<EF>>> SumcheckComputation<N, EF> for PrecompileFootprint
-where
-    EF: ExtensionField<N>,
-{
-    fn degree(&self) -> usize {
-        3
-    }
-    fn eval(&self, point: &[N], _: &[EF]) -> EF {
-        // TODO not all columns are used
-
-        let nu_a = (EF::ONE - point[COL_INDEX_FLAG_A]) * point[COL_INDEX_MEM_VALUE_A]
+impl PrecompileFootprint {
+    fn air_eval<
+        PointF: PrimeCharacteristicRing + Copy,
+        ResultF: Algebra<EF> + Algebra<PointF> + Copy,
+    >(
+        &self,
+        point: &[PointF],
+        mul_point_f_and_ef: impl Fn(PointF, EF) -> ResultF,
+    ) -> ResultF {
+        let nu_a = (ResultF::ONE - point[COL_INDEX_FLAG_A]) * point[COL_INDEX_MEM_VALUE_A]
             + point[COL_INDEX_FLAG_A] * point[COL_INDEX_OPERAND_A];
-        let nu_b = (EF::ONE - point[COL_INDEX_FLAG_B]) * point[COL_INDEX_MEM_VALUE_B]
+        let nu_b = (ResultF::ONE - point[COL_INDEX_FLAG_B]) * point[COL_INDEX_MEM_VALUE_B]
             + point[COL_INDEX_FLAG_B] * point[COL_INDEX_OPERAND_B];
-        let nu_c = (EF::ONE - point[COL_INDEX_FLAG_C]) * point[COL_INDEX_MEM_VALUE_C]
+        let nu_c = (ResultF::ONE - point[COL_INDEX_FLAG_C]) * point[COL_INDEX_MEM_VALUE_C]
             + point[COL_INDEX_FLAG_C] * point[COL_INDEX_FP];
 
-        self.global_challenge
-            + (self.p16_powers[1]
-                + self.p16_powers[2] * nu_a
-                + self.p16_powers[3] * nu_b
-                + self.p16_powers[4] * nu_c)
-                * point[COL_INDEX_POSEIDON_16]
-            + (self.p24_powers[1]
-                + self.p24_powers[2] * nu_a
-                + self.p24_powers[3] * nu_b
-                + self.p24_powers[4] * nu_c)
+        (nu_a * self.p16_powers[2]
+            + nu_b * self.p16_powers[3]
+            + nu_c * self.p16_powers[4]
+            + self.p16_powers[1])
+            * point[COL_INDEX_POSEIDON_16]
+            + (nu_a * self.p24_powers[2]
+                + nu_b * self.p24_powers[3]
+                + nu_c * self.p24_powers[4]
+                + self.p24_powers[1])
                 * point[COL_INDEX_POSEIDON_24]
-            + (self.dot_product_powers[1]
-                + self.dot_product_powers[2] * nu_a
-                + self.dot_product_powers[3] * nu_b
-                + self.dot_product_powers[4] * nu_c
-                + self.dot_product_powers[5] * point[COL_INDEX_AUX])
+            + (nu_a * self.dot_product_powers[2]
+                + nu_b * self.dot_product_powers[3]
+                + nu_c * self.dot_product_powers[4]
+                + mul_point_f_and_ef(point[COL_INDEX_AUX], self.dot_product_powers[5])
+                + self.dot_product_powers[1])
                 * point[COL_INDEX_DOT_PRODUCT]
-            + (self.multilinear_eval_powers[1]
-                + self.multilinear_eval_powers[2] * nu_a
-                + self.multilinear_eval_powers[3] * nu_b
-                + self.multilinear_eval_powers[4] * nu_c
-                + self.multilinear_eval_powers[5] * point[COL_INDEX_AUX])
+            + (nu_a * self.multilinear_eval_powers[2]
+                + nu_b * self.multilinear_eval_powers[3]
+                + nu_c * self.multilinear_eval_powers[4]
+                + mul_point_f_and_ef(point[COL_INDEX_AUX], self.multilinear_eval_powers[5])
+                + self.multilinear_eval_powers[1])
                 * point[COL_INDEX_MULTILINEAR_EVAL]
+            + self.global_challenge
+    }
+}
+
+impl<N: ExtensionField<F>> SumcheckComputation<N, EF> for PrecompileFootprint
+where
+    EF: ExtensionField<N>,
+{
+    fn degree(&self) -> usize {
+        3
+    }
+    fn eval(&self, point: &[N], _: &[EF]) -> EF {
+        self.air_eval(point, |p, c| c * p)
     }
 }
 
@@ -205,12 +216,12 @@ impl SumcheckComputationPacked<EF> for PrecompileFootprint {
         3
     }
 
-    fn eval_packed_extension(&self, _point: &[EFPacking<EF>], _: &[EF]) -> EFPacking<EF> {
-        todo!()
+    fn eval_packed_extension(&self, point: &[EFPacking<EF>], _: &[EF]) -> EFPacking<EF> {
+        self.air_eval(point, |p, c| p * c)
     }
 
-    fn eval_packed_base(&self, _point: &[utils::PFPacking<EF>], _: &[EF]) -> EFPacking<EF> {
-        todo!()
+    fn eval_packed_base(&self, point: &[PFPacking<EF>], _: &[EF]) -> EFPacking<EF> {
+        self.air_eval::<PFPacking<EF>, EFPacking<EF>>(point, |p, c| EFPacking::<EF>::from(p) * c)
     }
 }
 

crates/lean_prover/src/prove_execution.rs

@@ -408,23 +408,25 @@ pub fn prove_execution(
         grand_product_dot_product_sumcheck_point,
         grand_product_dot_product_sumcheck_inner_evals,
         _,
-    ) = sumcheck::prove(
-        1, // TODO univariate skip?
-        MleGroupRef::Extension(
-            dot_product_columns[..5]
-                .iter()
-                .map(|c| c.as_slice())
-                .collect::<Vec<_>>(),
-        ), // TODO packing
-        &dot_product_footprint_computation,
-        &dot_product_footprint_computation,
-        &[],
-        Some((grand_product_dot_product_statement.point.0.clone(), None)),
-        false,
-        &mut prover_state,
-        grand_product_dot_product_statement.value,
-        None,
-    );
+    ) = info_span!("Grand product sumcheck for Dot Product").in_scope(|| {
+        sumcheck::prove(
+            1, // TODO univariate skip?
+            MleGroupRef::Extension(
+                dot_product_columns[..5]
+                    .iter()
+                    .map(|c| c.as_slice())
+                    .collect::<Vec<_>>(),
+            ), // TODO packing
+            &dot_product_footprint_computation,
+            &dot_product_footprint_computation,
+            &[],
+            Some((grand_product_dot_product_statement.point.0.clone(), None)),
+            false,
+            &mut prover_state,
+            grand_product_dot_product_statement.value,
+            None,
+        )
+    });
     assert_eq!(grand_product_dot_product_sumcheck_inner_evals.len(), 5);
     prover_state.add_extension_scalars(&grand_product_dot_product_sumcheck_inner_evals);
 
@@ -460,21 +462,24 @@ pub fn prove_execution(
     };
 
     let (grand_product_exec_sumcheck_point, grand_product_exec_sumcheck_inner_evals, _) =
-        sumcheck::prove(
-            1, // TODO univariate skip?
-            MleGroupRef::Base(
-                // TODO not all columns re required
-                full_trace.iter().map(|c| c.as_slice()).collect::<Vec<_>>(),
-            ), // TODO packing
-            &precompile_foot_print_computation,
-            &precompile_foot_print_computation,
-            &[],
-            Some((grand_product_exec_statement.point.0.clone(), None)),
-            false,
-            &mut prover_state,
-            grand_product_exec_statement.value,
-            None,
-        );
+        info_span!("Grand product sumcheck for Execution").in_scope(|| {
+            sumcheck::prove(
+                1, // TODO univariate skip
+                MleGroupRef::Base(
+                    // TODO not all columns re required
+                    full_trace.iter().map(|c| c.as_slice()).collect::<Vec<_>>(),
+                )
+                .pack(),
+                &precompile_foot_print_computation,
+                &precompile_foot_print_computation,
+                &[],
+                Some((grand_product_exec_statement.point.0.clone(), None)),
+                false,
+                &mut prover_state,
+                grand_product_exec_statement.value,
+                None,
+            )
+        });
 
     prover_state.add_extension_scalars(&grand_product_exec_sumcheck_inner_evals);
     assert_eq!(

crates/rec_aggregation/src/recursion.rs

@@ -183,7 +183,7 @@ fn run_recursion_benchmark() -> RecursionBenchStats {
     // #[rustfmt::skip] // debug
     // std::fs::write("public_input.txt", build_public_memory(&public_input).chunks_exact(8).enumerate().map(|(i, chunk)| { format!("{} - {}: {}\n", i, i * 8, chunk.iter().map(|x| x.to_string()).collect::<Vec<_>>().join(", ")) }).collect::<String>(),).unwrap();
 
-    // utils::init_tracing();
+    utils::init_tracing();
     let (bytecode, function_locations) = compile_program(&program_str);
     let time = Instant::now();
     let (proof_data, proof_size) = prove_execution(