only 1 execution, at witness generation, instead of 2, using precomputed values for the "non -vectorized" memory size (idea by Lambdaclass)

Co-authored-by:  Robin Salen  <[email protected]>

Author: TomWambsgans

Cargo.lock

@@ -111,7 +111,9 @@ multilinear-toolkit.workspace = true
 # p3-poseidon2-air = { path = "../zk/Plonky3/poseidon2-air" }
 # p3-dft = { path = "../zk/Plonky3/dft" }
 # p3-challenger = { path = "../zk/Plonky3/challenger" }
-# p3-util = { path = "../zk/Plonky3/util" }
+
+# [patch."https://github.com/TomWambsgans/whir-p3.git"]
+# whir-p3 = { path = "../zk/whir/fork-whir-p3" }
 
 [dev-dependencies]
 criterion = { version = "0.7", default-features = false, features = ["cargo_bench_support"] }

Cargo.toml

@@ -28,6 +28,9 @@
 - opti compute_eval_eq when scalar = ONE
 - Dmitry's range check, bonus: we can spare 2 memory cells if the value being range check is small (using the zeros present by conventio on the public memory)
 - Make everything "padding aware" (including WHIR, logup*, AIR, etc)
+- Opti WHIR: in sumcheck we know more than f(0) + f(1), we know f(0) and f(1)
+- Opti WHIR https://github.com/tcoratger/whir-p3/issues/303 and https://github.com/tcoratger/whir-p3/issues/306
+- Avoid committing to extra columns / adding some constraints in poseidon16 AIR, for "compression", and use instead sumcheck
 
 About "the packed pcs" (similar to SP1 Jagged PCS, slightly less efficient, but simpler (no sumchecks)):
 - The best strategy is probably to pack as much as possible (the cost increasing the density = additional inner evaluations), if we can fit below a power of 2 - epsilon  (epsilon = 20% for instance, tbd), if the sum of the non zero data is just above a power of 2, no packed technique, even the best, can help us, so we should spread aniway (to reduce the pressure of inner evaluations)

TODO.md

@@ -707,19 +707,24 @@ fn compile_poseidon(
     compiler: &Compiler,
     over_16: bool, // otherwise over_24
 ) -> Result<(), String> {
-    assert_eq!(args.len(), 3);
-
     let low_level_arg_a = IntermediateValue::from_simple_expr(&args[0], compiler);
     let low_level_arg_b = IntermediateValue::from_simple_expr(&args[1], compiler);
     let low_level_res = IntermediateValue::from_simple_expr(&args[2], compiler);
 
     if over_16 {
+        assert_eq!(args.len(), 4);
+        let is_compression = args[3].as_constant().unwrap().naive_eval().unwrap();
+        assert!(is_compression.is_zero() || is_compression.is_one());
+        let is_compression = is_compression.is_one();
+
         instructions.push(IntermediateInstruction::Poseidon2_16 {
             arg_a: low_level_arg_a,
             arg_b: low_level_arg_b,
             res: low_level_res,
+            is_compression,
         });
     } else {
+        assert_eq!(args.len(), 3);
         instructions.push(IntermediateInstruction::Poseidon2_24 {
             arg_a: low_level_arg_a,
             arg_b: low_level_arg_b,

crates/lean_compiler/src/b_compile_intermediate.rs

@@ -261,11 +261,17 @@ fn compile_block(
                     IntermediateValue::Constant(ConstExpression::Value(ConstantValue::Scalar(1)));
                 codegen_jump(hints, low_level_bytecode, one, dest, updated_fp)
             }
-            IntermediateInstruction::Poseidon2_16 { arg_a, arg_b, res } => {
+            IntermediateInstruction::Poseidon2_16 {
+                arg_a,
+                arg_b,
+                res,
+                is_compression,
+            } => {
                 low_level_bytecode.push(Instruction::Poseidon2_16 {
                     arg_a: try_as_mem_or_constant(&arg_a).unwrap(),
                     arg_b: try_as_mem_or_constant(&arg_b).unwrap(),
                     res: try_as_mem_or_fp(&res).unwrap(),
+                    is_compression,
                 });
             }
             IntermediateInstruction::Poseidon2_24 { arg_a, arg_b, res } => {

crates/lean_compiler/src/c_compile_final.rs

@@ -31,7 +31,8 @@ pub enum IntermediateInstruction {
     Poseidon2_16 {
         arg_a: IntermediateValue, // vectorized pointer, of size 1
         arg_b: IntermediateValue, // vectorized pointer, of size 1
-        res: IntermediateValue,   // vectorized pointer, of size 2
+        res: IntermediateValue,   // vectorized pointer, of size 1 if `is_compression`, else size 2
+        is_compression: bool,
     },
     Poseidon2_24 {
         arg_a: IntermediateValue, // vectorized pointer, of size 2 (2 first inputs)
@@ -168,8 +169,16 @@ impl Display for IntermediateInstruction {
                     write!(f, "jump_if_not_zero {condition} to {dest}")
                 }
             }
-            Self::Poseidon2_16 { arg_a, arg_b, res } => {
-                write!(f, "{res} = poseidon2_16({arg_a}, {arg_b})")
+            Self::Poseidon2_16 {
+                arg_a,
+                arg_b,
+                res,
+                is_compression,
+            } => {
+                write!(
+                    f,
+                    "{res} = poseidon2_16({arg_a}, {arg_b}, is_compression={is_compression})"
+                )
             }
             Self::Poseidon2_24 { arg_a, arg_b, res } => {
                 write!(f, "{res} = poseidon2_24({arg_a}, {arg_b})")

crates/lean_compiler/src/ir/instruction.rs

@@ -27,7 +27,7 @@ impl Display for PrecompileName {
 
 pub const POSEIDON_16: Precompile = Precompile {
     name: PrecompileName::Poseidon16,
-    n_inputs: 3,
+    n_inputs: 4,
 };
 
 pub const POSEIDON_24: Precompile = Precompile {

crates/lean_compiler/src/precompiles.rs

@@ -33,7 +33,6 @@ pub fn get_base_dims(
             ColDims::padded(n_cycles, F::ZERO),                           // mem_addr_c
             ColDims::padded(n_poseidons_16, F::from_usize(ZERO_VEC_PTR)), // poseidon16 index a
             ColDims::padded(n_poseidons_16, F::from_usize(ZERO_VEC_PTR)), // poseidon16 index b
-            ColDims::padded(n_poseidons_16, F::from_usize(POSEIDON_16_NULL_HASH_PTR)), // poseidon16 index res
             ColDims::padded(n_poseidons_24, F::from_usize(ZERO_VEC_PTR)), // poseidon24 index a
             ColDims::padded(n_poseidons_24, F::from_usize(ZERO_VEC_PTR)), // poseidon24 index b
             ColDims::padded(n_poseidons_24, F::from_usize(POSEIDON_24_NULL_HASH_PTR)), // poseidon24 index res
@@ -139,7 +138,7 @@ pub fn add_memory_statements_for_dot_product_precompile(
 
 pub struct PrecompileFootprint {
     pub global_challenge: EF,
-    pub p16_powers: [EF; 5],
+    pub p16_powers: [EF; 6],
     pub p24_powers: [EF; 5],
     pub dot_product_powers: [EF; 6],
     pub multilinear_eval_powers: [EF; 6],
@@ -164,6 +163,7 @@ impl PrecompileFootprint {
         (nu_a * self.p16_powers[2]
             + nu_b * self.p16_powers[3]
             + nu_c * self.p16_powers[4]
+            + mul_point_f_and_ef(point[COL_INDEX_AUX], self.p16_powers[5])
             + self.p16_powers[1])
             * point[COL_INDEX_POSEIDON_16]
             + (nu_a * self.p24_powers[2]
@@ -383,10 +383,10 @@ pub fn add_poseidon_lookup_statements_on_indexes(
     log_n_p24: usize,
     index_lookup_point: &MultilinearPoint<EF>,
     inner_values: &[EF],
-    p16_indexe_statements: [&mut Vec<Evaluation<EF>>; 3], // a, b, res
+    p16_indexe_statements: [&mut Vec<Evaluation<EF>>; 4], // a, b, res_1, res_2
     p24_indexe_statements: [&mut Vec<Evaluation<EF>>; 3], // a, b, res
 ) {
-    assert_eq!(inner_values.len(), 6);
+    assert_eq!(inner_values.len(), 7);
     let mut idx_point_right_p16 = MultilinearPoint(index_lookup_point[3..].to_vec());
     let mut idx_point_right_p24 =
         MultilinearPoint(index_lookup_point[3 + log_n_p16.abs_diff(log_n_p24)..].to_vec());
@@ -402,7 +402,7 @@ pub fn add_poseidon_lookup_statements_on_indexes(
     for (i, stmt) in p24_indexe_statements.into_iter().enumerate() {
         stmt.push(Evaluation::new(
             idx_point_right_p24.clone(),
-            inner_values[i + 3],
+            inner_values[i + 4],
         ));
     }
 }