feat: Memory Offline Checker (#294)

---------

Co-authored-by: Zach Langley <[email protected]>
Co-authored-by: Jonathan Wang <[email protected]>

.github/workflows/db.yml

@@ -36,12 +36,14 @@ jobs:
 
   benchmark-rw:
     uses: ./.github/workflows/benchmark-call.yml
+    if: github.event.pull_request.draft == false
     with:
       benchmark_name: single_rw
     secrets: inherit
 
   benchmark-filter:
     uses: ./.github/workflows/benchmark-call.yml
+    if: github.event.pull_request.draft == false
     with:
       benchmark_name: single_filter
     secrets: inherit

.github/workflows/recursion.yml

@@ -38,6 +38,7 @@ jobs:
 
   benchmark:
     uses: ./.github/workflows/benchmark-call.yml
+    if: github.event.pull_request.draft == false
     with:
       benchmark_name: vm_verify_fibair
     secrets: inherit

benchmark/src/commands/vm/benchmark_helpers.rs

@@ -123,10 +123,10 @@ pub fn run_recursive_test_benchmark(
     let mut witness_stream = Vec::new();
     witness_stream.extend(input.write());
 
-    vm_benchmark_execute_and_prove::<1>(program, witness_stream, benchmark_name)
+    vm_benchmark_execute_and_prove::<8, 1>(program, witness_stream, benchmark_name)
 }
 
-pub fn vm_benchmark_execute_and_prove<const WORD_SIZE: usize>(
+pub fn vm_benchmark_execute_and_prove<const NUM_WORDS: usize, const WORD_SIZE: usize>(
     program: Program<BabyBear>,
     input_stream: Vec<Vec<BabyBear>>,
     benchmark_name: &str,
@@ -138,7 +138,7 @@ pub fn vm_benchmark_execute_and_prove<const WORD_SIZE: usize>(
         ..Default::default()
     };
 
-    let mut vm = VirtualMachine::<WORD_SIZE, _>::new(vm_config, program, input_stream);
+    let mut vm = VirtualMachine::<NUM_WORDS, WORD_SIZE, _>::new(vm_config, program, input_stream);
     vm.enable_metrics_collection();
 
     let vm_execute_span = info_span!("Benchmark vm execute").entered();
@@ -152,7 +152,7 @@ pub fn vm_benchmark_execute_and_prove<const WORD_SIZE: usize>(
     } = vm.execute_and_generate_traces().unwrap();
     vm_execute_span.exit();
 
-    let chips = VirtualMachine::<WORD_SIZE, _>::get_chips(&chips);
+    let chips = VirtualMachine::<NUM_WORDS, WORD_SIZE, _>::get_chips(&chips);
 
     let perm = default_perm();
     // blowup factor 8 for poseidon2 chip

benchmark/src/commands/vm/vm_fib_program.rs

@@ -9,6 +9,9 @@ use p3_field::{extension::BinomialExtensionField, AbstractField};
 use super::benchmark_helpers::vm_benchmark_execute_and_prove;
 
 pub fn benchmark_fib_program(n: usize) -> Result<()> {
+    const NUM_WORDS: usize = 8;
+    const WORD_SIZE: usize = 1;
+
     println!("Running Fibonacci program benchmark with n = {}", n);
 
     type F = BabyBear;
@@ -33,5 +36,9 @@ pub fn benchmark_fib_program(n: usize) -> Result<()> {
 
     let fib_program = builder.compile_isa::<1>();
 
-    vm_benchmark_execute_and_prove::<1>(fib_program, vec![], "VM Fibonacci Program")
+    vm_benchmark_execute_and_prove::<NUM_WORDS, WORD_SIZE>(
+        fib_program,
+        vec![],
+        "VM Fibonacci Program",
+    )
 }

benchmark/src/commands/vm/vm_fib_verifier_program.rs

@@ -1,3 +1,5 @@
+use std::ops::Deref;
+
 use afs_compiler::{
     asm::AsmBuilder,
     ir::{Felt, Var},
@@ -14,6 +16,9 @@ use stark_vm::vm::{config::VmConfig, ExecutionAndTraceGenerationResult, VirtualM
 use super::benchmark_helpers::run_recursive_test_benchmark;
 
 pub fn benchmark_fib_verifier_program(n: usize) -> Result<()> {
+    const NUM_WORDS: usize = 8;
+    const WORD_SIZE: usize = 1;
+
     println!(
         "Running verifier program of VM STARK benchmark with n = {}",
         n
@@ -46,7 +51,7 @@ pub fn benchmark_fib_verifier_program(n: usize) -> Result<()> {
         ..Default::default()
     };
 
-    let vm = VirtualMachine::<1, _>::new(vm_config, fib_program.clone(), vec![]);
+    let vm = VirtualMachine::<8, 1, _>::new(vm_config, fib_program.clone(), vec![]);
 
     let ExecutionAndTraceGenerationResult {
         max_log_degree: _,
@@ -55,12 +60,13 @@ pub fn benchmark_fib_verifier_program(n: usize) -> Result<()> {
         nonempty_pis: pis,
         ..
     } = vm.execute_and_generate_traces().unwrap();
-    let chips = VirtualMachine::<1, _>::get_chips(&chips);
+    let chips = VirtualMachine::<NUM_WORDS, WORD_SIZE, _>::get_chips(&chips);
 
-    let dummy_vm = VirtualMachine::<1, _>::new(vm_config, fib_program, vec![]);
-    let rec_raps = get_rec_raps::<1, InnerConfig>(&dummy_vm.segments[0]);
+    let dummy_vm = VirtualMachine::<NUM_WORDS, WORD_SIZE, _>::new(vm_config, fib_program, vec![]);
+    let rec_raps = get_rec_raps::<NUM_WORDS, WORD_SIZE, InnerConfig>(&dummy_vm.segments[0]);
+    let rec_raps: Vec<_> = rec_raps.iter().map(|x| x.deref()).collect();
 
-    assert!(chips.len() == rec_raps.len());
+    assert_eq!(chips.len(), rec_raps.len());
 
     let pvs = pis;
     let (chips, rec_raps, traces, pvs) = sort_chips(chips, rec_raps, traces, pvs);

compiler/src/asm/compiler.rs

@@ -12,7 +12,7 @@ use crate::{
 };
 
 /// The memory location for the top of memory
-pub const MEMORY_TOP: i32 = (1 << 30) - 4;
+pub const MEMORY_TOP: i32 = (1 << 29) - 4;
 
 /// The heap pointer address.
 pub(crate) const HEAP_PTR: i32 = MEMORY_TOP - 4;

compiler/src/util.rs

@@ -21,7 +21,7 @@ pub fn execute_program_with_config<const WORD_SIZE: usize>(
     program: Program<BabyBear>,
     input_stream: Vec<Vec<BabyBear>>,
 ) {
-    let vm = VirtualMachine::<WORD_SIZE, _>::new(config, program, input_stream);
+    let vm = VirtualMachine::<1, WORD_SIZE, _>::new(config, program, input_stream);
     vm.execute().unwrap();
 }
 
@@ -43,7 +43,7 @@ pub fn execute_program<const WORD_SIZE: usize>(
     program: Program<BabyBear>,
     input_stream: Vec<Vec<BabyBear>>,
 ) {
-    let vm = VirtualMachine::<WORD_SIZE, _>::new(
+    let vm = VirtualMachine::<1, WORD_SIZE, _>::new(
         VmConfig {
             num_public_values: 4,
             max_segment_len: (1 << 25) - 100,
@@ -60,7 +60,7 @@ pub fn execute_program_and_generate_traces<const WORD_SIZE: usize>(
     program: Program<BabyBear>,
     input_stream: Vec<Vec<BabyBear>>,
 ) {
-    let vm = VirtualMachine::<WORD_SIZE, _>::new(
+    let vm = VirtualMachine::<1, WORD_SIZE, _>::new(
         VmConfig {
             num_public_values: 4,
             max_segment_len: (1 << 25) - 100,
@@ -77,7 +77,7 @@ pub fn execute_program_with_public_values<const WORD_SIZE: usize>(
     input_stream: Vec<Vec<BabyBear>>,
     public_values: &[(usize, BabyBear)],
 ) {
-    let mut vm = VirtualMachine::<WORD_SIZE, _>::new(
+    let mut vm = VirtualMachine::<1, WORD_SIZE, _>::new(
         VmConfig {
             num_public_values: 4,
             ..Default::default()
@@ -148,7 +148,7 @@ pub fn execute_and_prove_program<const WORD_SIZE: usize>(
     program: Program<BabyBear>,
     input_stream: Vec<Vec<BabyBear>>,
 ) {
-    let vm = VirtualMachine::<WORD_SIZE, _>::new(
+    let vm = VirtualMachine::<1, WORD_SIZE, _>::new(
         VmConfig {
             num_public_values: 4,
             ..Default::default()
@@ -164,7 +164,7 @@ pub fn execute_and_prove_program<const WORD_SIZE: usize>(
         nonempty_pis: pis,
         ..
     } = vm.execute_and_generate_traces().unwrap();
-    let chips = VirtualMachine::<WORD_SIZE, _>::get_chips(&chips);
+    let chips = VirtualMachine::<1, WORD_SIZE, _>::get_chips(&chips);
 
     let perm = random_perm();
     // blowup factor 8 for poseidon2 chip

compiler/tests/arithmetic.rs

@@ -438,7 +438,7 @@ fn assert_failed_assertion(
     builder: Builder<AsmConfig<BabyBear, BinomialExtensionField<BabyBear, 4>>>,
 ) {
     let program = builder.compile_isa::<WORD_SIZE>();
-    let vm = VirtualMachine::<WORD_SIZE, _>::new(VmConfig::default(), program, vec![]);
+    let vm = VirtualMachine::<1, WORD_SIZE, _>::new(VmConfig::default(), program, vec![]);
     let result = vm.execute_and_generate_traces();
     assert!(matches!(result, Err(Fail(_))));
 }

docs/specs/vm/ISA.md

@@ -94,7 +94,7 @@ The following notation is used throughout this document:
 
 **Operand values**: `opa`, `opb`, `opc` denote the value encoded in the operand a, b, or c of the current instruction.
 
-**CPU registers**: `pc` denote the value of the current program counter.
+**CPU registers**: `pc` denotes the value of the current program counter.
 
 **Addressing**: we support different address spaces via `as_b, as_c`.
 

poseidon2-air/src/poseidon2/air.rs

@@ -26,6 +26,7 @@ use super::{
 /// - WIDTH is multiple of 4 and >= 8
 ///
 /// Spec is at https://hackmd.io/_I1lx-6GROWbKbDi_Vz-pw?view .
+#[derive(Clone)]
 pub struct Poseidon2Air<const WIDTH: usize, F> {
     pub rounds_f: usize,
     pub external_constants: Vec<[F; WIDTH]>,

poseidon2-air/src/poseidon2/columns.rs

@@ -39,7 +39,7 @@ pub struct Poseidon2ColsIndexMap<const WIDTH: usize> {
 }
 
 impl<const WIDTH: usize, T: Clone> Poseidon2Cols<WIDTH, T> {
-    pub fn get_width(poseidon2_air: &Poseidon2Air<WIDTH, T>) -> usize {
+    pub fn get_width<F: Clone>(poseidon2_air: &Poseidon2Air<WIDTH, F>) -> usize {
         let io_width = Poseidon2IoCols::<WIDTH, T>::get_width();
         let aux_width = Poseidon2AuxCols::<WIDTH, T>::get_width(poseidon2_air);
         io_width + aux_width
@@ -136,7 +136,7 @@ impl<const WIDTH: usize, T: Clone> Poseidon2IoCols<WIDTH, T> {
 }
 
 impl<const WIDTH: usize, T: Clone> Poseidon2AuxCols<WIDTH, T> {
-    pub fn get_width(poseidon2_air: &Poseidon2Air<WIDTH, T>) -> usize {
+    pub fn get_width<F: Clone>(poseidon2_air: &Poseidon2Air<WIDTH, F>) -> usize {
         (poseidon2_air.rounds_f + poseidon2_air.rounds_p) * WIDTH
     }
 

primitives/src/is_equal_vec/columns.rs

@@ -1,3 +1,6 @@
+use derive_new::new;
+use p3_air::AirBuilder;
+
 use super::IsEqualVecAir;
 
 #[derive(Default)]
@@ -28,7 +31,7 @@ impl<T> IsEqualVecIoCols<T> {
     }
 }
 
-#[derive(Default, Debug, Clone, PartialEq, Eq)]
+#[derive(Default, Debug, Clone, PartialEq, Eq, new)]
 pub struct IsEqualVecAuxCols<T> {
     /// prods[i] indicates whether x[i] == y[i] up to the i-th index
     pub prods: Vec<T>,
@@ -52,6 +55,18 @@ impl<T: Clone> IsEqualVecAuxCols<T> {
     }
 }
 
+impl<T> IsEqualVecAuxCols<T> {
+    pub fn into_expr<AB: AirBuilder>(self) -> IsEqualVecAuxCols<AB::Expr>
+    where
+        T: Into<AB::Expr>,
+    {
+        IsEqualVecAuxCols::new(
+            self.prods.into_iter().map(|x| x.into()).collect(),
+            self.invs.into_iter().map(|x| x.into()).collect(),
+        )
+    }
+}
+
 #[derive(Default)]
 pub struct IsEqualVecCols<T> {
     pub io: IsEqualVecIoCols<T>,

primitives/src/is_less_than/air.rs

@@ -35,7 +35,7 @@ impl IsLessThanAir {
     pub(crate) fn eval_without_interactions<AB: AirBuilder>(
         &self,
         builder: &mut AB,
-        io: IsLessThanIoCols<AB::Var>,
+        io: IsLessThanIoCols<AB::Expr>,
         aux: IsLessThanAuxCols<AB::Var>,
     ) {
         let x = io.x;
@@ -54,7 +54,8 @@ impl IsLessThanAir {
         // constrain that the lower_bits + less_than * 2^limb_bits is the correct intermediate sum
         // note that the intermediate value will be >= 2^limb_bits if and only if x < y, and check_val will therefore be
         // the correct value if and only if less_than is the indicator for whether x < y
-        let check_val = lower + less_than * AB::Expr::from_canonical_u64(1 << self.max_bits);
+        let check_val =
+            lower + less_than.clone() * AB::Expr::from_canonical_u64(1 << self.max_bits);
 
         builder.assert_eq(intermed_val, check_val);
 
@@ -120,13 +121,28 @@ impl<AB: InteractionBuilder> SubAir<AB> for IsLessThanAir {
     // constrain that the result of x < y is given by less_than
     // warning: send for range check must be included for the constraints to be sound
     fn eval(&self, builder: &mut AB, io: Self::IoView, aux: Self::AuxView) {
-        // Note: every AIR that uses this sub-AIR must include these interactions for soundness
-        self.eval_interactions(builder, aux.lower_decomp.clone());
-        self.eval_without_interactions(builder, io, aux);
+        // Note: every AIR that uses this sub-AIR must include the interactions for soundness
+        self.subair_eval(
+            builder,
+            IsLessThanIoCols::<AB::Expr>::new(io.x, io.y, io.less_than),
+            aux,
+        );
     }
 }
 
 impl IsLessThanAir {
+    pub fn subair_eval<AB: InteractionBuilder>(
+        &self,
+        builder: &mut AB,
+        io: IsLessThanIoCols<AB::Expr>,
+        aux: IsLessThanAuxCols<AB::Var>,
+    ) {
+        let io_exprs = IsLessThanIoCols::<AB::Expr>::new(io.x, io.y, io.less_than);
+
+        self.eval_interactions(builder, aux.lower_decomp.clone());
+        self.eval_without_interactions(builder, io_exprs, aux);
+    }
+
     /// Imposes the non-interaction constraints on all except the last row. This is
     /// intended for use when the comparators `x, y` are on adjacent rows.
     ///
@@ -138,10 +154,12 @@ impl IsLessThanAir {
     pub fn eval_when_transition<AB: InteractionBuilder>(
         &self,
         builder: &mut AB,
-        io: IsLessThanIoCols<AB::Var>,
+        io: IsLessThanIoCols<impl Into<AB::Expr>>,
         aux: IsLessThanAuxCols<AB::Var>,
     ) {
+        let io_exprs = IsLessThanIoCols::<AB::Expr>::new(io.x, io.y, io.less_than);
+
         self.eval_interactions(builder, aux.lower_decomp.clone());
-        self.eval_without_interactions(&mut builder.when_transition(), io, aux);
+        self.eval_without_interactions(&mut builder.when_transition(), io_exprs, aux);
     }
 }