only 1 execution, at witness generation, instead of 2, using precomputed values (idea by Lambdaclass)

Co-authored-by: Mauro Toscano <[email protected]>

Author: TomWambsgans

TODO.md

@@ -19,7 +19,6 @@
 - Lev's trick to skip some low-level modular reduction
 - Sumcheck, case z = 0, no need to fold, only keep first half of the values (done in PR 33 by Lambda) (and also in WHIR?)
 - Custom AVX2 / AVX512 / Neon implem in Plonky3 for all of the finite field operations (done for degree 4 extension, but not degree 5)
-- the 2 executions of the program, before generating the validity proof, can be merged, using some kind of placeholders
 - Many times, we evaluate different multilinear polynomials (diferent columns of the same table etc) at a common point. OPTI = compute the eq(.) once, and then dot_product with everything
 - To commit to multiple AIR table using 1 single pcs, the most general form our "packed pcs" api should accept is:
  a list of n (n not a power of 2) columns, each ending with m repeated values (in this manner we can reduce proof size when they are a lot of columns (poseidons ...))

benches/xmss.rs

@@ -1,11 +1,10 @@
 use std::{hint::black_box, time::Duration};
 
 use criterion::{Criterion, Throughput, criterion_group, criterion_main};
-use rec_aggregation::bench_xmss;
+use rec_aggregation::run_xmss_benchmark;
 
 fn bench_xmss_benchmark(c: &mut Criterion) {
     const N: usize = 500;
-    const LOG_LIFETIME: usize = 32;
 
     let mut group = c.benchmark_group("xmss");
     group.sample_size(10);
@@ -15,7 +14,7 @@ fn bench_xmss_benchmark(c: &mut Criterion) {
 
     group.bench_function("xmss", |b| {
         b.iter(|| {
-            let duration = bench_xmss(N, LOG_LIFETIME);
+            let duration = run_xmss_benchmark(N).proving_time;
             black_box(duration);
         });
     });

crates/lean_compiler/src/lib.rs

@@ -32,14 +32,21 @@ pub fn compile_program(program: &str) -> (Bytecode, BTreeMap<usize, String>) {
     (compiled, function_locations)
 }
 
-pub fn compile_and_run(program: &str, public_input: &[F], private_input: &[F], profiler: bool) {
+pub fn compile_and_run(
+    program: &str,
+    public_input: &[F],
+    private_input: &[F],
+    no_vec_runtime_memory: usize, // size of the "non-vectorized" runtime memory
+    profiler: bool,
+) {
     let (bytecode, function_locations) = compile_program(program);
     execute_bytecode(
         &bytecode,
         public_input,
         private_input,
         program,
         &function_locations,
+        no_vec_runtime_memory,
         profiler,
     );
 }

crates/lean_compiler/tests/test_compiler.rs

@@ -3,6 +3,8 @@ use lean_vm::*;
 use p3_symmetric::Permutation;
 use utils::{get_poseidon16, get_poseidon24};
 
+const DEFAULT_NO_VEC_RUNTIME_MEMORY: usize = 1 << 15;
+
 #[test]
 #[should_panic]
 fn test_duplicate_function_name() {
@@ -20,7 +22,7 @@ fn test_duplicate_function_name() {
         return;
     }
     "#;
-    compile_and_run(program, &[], &[], false);
+    compile_and_run(program, &[], &[], DEFAULT_NO_VEC_RUNTIME_MEMORY, false);
 }
 
 #[test]
@@ -34,7 +36,7 @@ fn test_duplicate_constant_name() {
         return;
     }
     "#;
-    compile_and_run(program, &[], &[], false);
+    compile_and_run(program, &[], &[], DEFAULT_NO_VEC_RUNTIME_MEMORY, false);
 }
 
 #[test]
@@ -55,7 +57,7 @@ fn test_fibonacci_program() {
         return;
     }
    "#;
-    compile_and_run(program, &[], &[], false);
+    compile_and_run(program, &[], &[], DEFAULT_NO_VEC_RUNTIME_MEMORY, false);
 }
 
 #[test]
@@ -73,7 +75,7 @@ fn test_edge_case_0() {
         return;
     }
    "#;
-    compile_and_run(program, &[], &[], false);
+    compile_and_run(program, &[], &[], DEFAULT_NO_VEC_RUNTIME_MEMORY, false);
 }
 
 #[test]
@@ -86,7 +88,7 @@ fn test_edge_case_1() {
         return;
     }
    "#;
-    compile_and_run(program, &[], &[], false);
+    compile_and_run(program, &[], &[], DEFAULT_NO_VEC_RUNTIME_MEMORY, false);
 }
 
 #[test]
@@ -104,7 +106,7 @@ fn test_edge_case_2() {
         return;
     }
    "#;
-    compile_and_run(program, &[], &[], false);
+    compile_and_run(program, &[], &[], DEFAULT_NO_VEC_RUNTIME_MEMORY, false);
 }
 
 #[test]
@@ -120,7 +122,7 @@ fn test_decompose_bits() {
         return;
     }
    "#;
-    compile_and_run(program, &[], &[], false);
+    compile_and_run(program, &[], &[], DEFAULT_NO_VEC_RUNTIME_MEMORY, false);
 }
 
 #[test]
@@ -136,7 +138,7 @@ fn test_unroll() {
         return;
     }
    "#;
-    compile_and_run(program, &[], &[], false);
+    compile_and_run(program, &[], &[], DEFAULT_NO_VEC_RUNTIME_MEMORY, false);
 }
 
 #[test]
@@ -148,7 +150,7 @@ fn test_rev_unroll() {
         return;
     }
    "#;
-    compile_and_run(program, &[], &[], false);
+    compile_and_run(program, &[], &[], DEFAULT_NO_VEC_RUNTIME_MEMORY, false);
 }
 
 #[test]
@@ -168,7 +170,7 @@ fn test_mini_program_0() {
         return;
     }
    "#;
-    compile_and_run(program, &[], &[], false);
+    compile_and_run(program, &[], &[], DEFAULT_NO_VEC_RUNTIME_MEMORY, false);
 }
 
 #[test]
@@ -211,7 +213,7 @@ fn test_mini_program_1() {
         return;
     }
    "#;
-    compile_and_run(program, &[], &[], false);
+    compile_and_run(program, &[], &[], DEFAULT_NO_VEC_RUNTIME_MEMORY, false);
 }
 
 #[test]
@@ -239,7 +241,7 @@ fn test_mini_program_2() {
         return sum, product;
     }
    "#;
-    compile_and_run(program, &[], &[], false);
+    compile_and_run(program, &[], &[], DEFAULT_NO_VEC_RUNTIME_MEMORY, false);
 }
 
 #[test]
@@ -266,7 +268,13 @@ fn test_mini_program_3() {
     }
    "#;
     let mut public_input: [F; 16] = (0..16).map(F::new).collect::<Vec<F>>().try_into().unwrap();
-    compile_and_run(program, &public_input, &[], false);
+    compile_and_run(
+        program,
+        &public_input,
+        &[],
+        DEFAULT_NO_VEC_RUNTIME_MEMORY,
+        false,
+    );
 
     get_poseidon16().permute_mut(&mut public_input);
     let _ = public_input;
@@ -289,7 +297,13 @@ fn test_mini_program_4() {
     }
    "#;
     let mut public_input: [F; 24] = (0..24).map(F::new).collect::<Vec<F>>().try_into().unwrap();
-    compile_and_run(program, &public_input, &[], false);
+    compile_and_run(
+        program,
+        &public_input,
+        &[],
+        DEFAULT_NO_VEC_RUNTIME_MEMORY,
+        false,
+    );
 
     get_poseidon24().permute_mut(&mut public_input);
     dbg!(&public_input[16..]);
@@ -369,7 +383,7 @@ fn test_inlined() {
         return;
     }
    "#;
-    compile_and_run(program, &[], &[], false);
+    compile_and_run(program, &[], &[], DEFAULT_NO_VEC_RUNTIME_MEMORY, false);
 }
 
 #[test]
@@ -421,7 +435,7 @@ fn test_match() {
         return x * x * x * x * x * x;
     }
    "#;
-    compile_and_run(program, &[], &[], false);
+    compile_and_run(program, &[], &[], DEFAULT_NO_VEC_RUNTIME_MEMORY, false);
 }
 
 // #[test]
@@ -462,7 +476,7 @@ fn test_const_functions_calling_const_functions() {
     }
     "#;
 
-    compile_and_run(program, &[], &[], false);
+    compile_and_run(program, &[], &[], DEFAULT_NO_VEC_RUNTIME_MEMORY, false);
 }
 
 #[test]
@@ -485,7 +499,7 @@ fn test_inline_functions_calling_inline_functions() {
     }
     "#;
 
-    compile_and_run(program, &[], &[], false);
+    compile_and_run(program, &[], &[], DEFAULT_NO_VEC_RUNTIME_MEMORY, false);
 }
 
 #[test]
@@ -512,5 +526,5 @@ fn test_nested_inline_functions() {
     }
     "#;
 
-    compile_and_run(program, &[], &[], false);
+    compile_and_run(program, &[], &[], DEFAULT_NO_VEC_RUNTIME_MEMORY, false);
 }

crates/lean_prover/src/prove_execution.rs

@@ -29,9 +29,9 @@ pub fn prove_execution(
     bytecode: &Bytecode,
     source_code: &str,                            // debug purpose
     function_locations: &BTreeMap<usize, String>, // debug purpose
-    public_input: &[F],
-    private_input: &[F],
+    (public_input, private_input): (&[F], &[F]),
     whir_config_builder: WhirConfigBuilder,
+    no_vec_runtime_memory: usize, // size of the "non-vectorized" runtime memory
     vm_profiler: bool,
 ) -> (Vec<PF<EF>>, usize) {
     let ExecutionTrace {
@@ -52,6 +52,7 @@ pub fn prove_execution(
             private_input,
             source_code,
             function_locations,
+            no_vec_runtime_memory,
             vm_profiler,
         );
         get_execution_trace(bytecode, execution_result)

crates/lean_prover/tests/test_zkvm.rs

@@ -83,9 +83,9 @@ fn test_zk_vm() {
         &bytecode,
         &program_str,
         &function_locations,
-        &public_input,
-        &private_input,
+        (&public_input, &private_input),
         whir_config_builder(),
+        1 << 20,
         false,
     )
     .0;