feat(new-execution): metered execution

benchmarks/execute/Cargo.toml

@@ -16,6 +16,8 @@ openvm-stark-sdk.workspace = true
 openvm-transpiler.workspace = true
 openvm-rv32im-circuit.workspace = true
 openvm-rv32im-transpiler.workspace = true
+openvm-bigint-circuit.workspace = true
+openvm-bigint-transpiler.workspace = true
 # openvm-keccak256-circuit.workspace = true
 # openvm-keccak256-transpiler.workspace = true
 

benchmarks/execute/benches/execute.rs

@@ -1,10 +1,8 @@
-use std::path::PathBuf;
-
-use divan;
 use eyre::Result;
 use openvm_benchmarks_utils::{get_elf_path, get_programs_dir, read_elf_file};
+use openvm_bigint_circuit::Int256Rv32Config;
+use openvm_bigint_transpiler::Int256TranspilerExtension;
 use openvm_circuit::arch::{instructions::exe::VmExe, VmExecutor};
-use openvm_rv32im_circuit::Rv32ImConfig;
 use openvm_rv32im_transpiler::{
     Rv32ITranspilerExtension, Rv32IoTranspilerExtension, Rv32MTranspilerExtension,
 };
@@ -16,13 +14,14 @@ static AVAILABLE_PROGRAMS: &[&str] = &[
     "fibonacci_iterative",
     "quicksort",
     "bubblesort",
+    "factorial_iterative_u256",
+    "revm_snailtracer",
     // "pairing",
     // "keccak256",
     // "keccak256_iter",
     // "sha256",
     // "sha256_iter",
     // "revm_transfer",
-    // "revm_snailtracer",
 ];
 
 fn main() {
@@ -35,12 +34,13 @@ fn run_program(program: &str) -> Result<()> {
     let elf_path = get_elf_path(&program_dir);
     let elf = read_elf_file(&elf_path)?;
 
-    let vm_config = Rv32ImConfig::default();
+    let vm_config = Int256Rv32Config::default();
 
     let transpiler = Transpiler::<BabyBear>::default()
         .with_extension(Rv32ITranspilerExtension)
         .with_extension(Rv32IoTranspilerExtension)
-        .with_extension(Rv32MTranspilerExtension);
+        .with_extension(Rv32MTranspilerExtension)
+        .with_extension(Int256TranspilerExtension);
 
     let exe = VmExe::from_elf(elf, transpiler)?;
 

benchmarks/execute/src/main.rs

@@ -1,12 +1,16 @@
 use clap::{Parser, ValueEnum};
 use eyre::Result;
 use openvm_benchmarks_utils::{get_elf_path, get_programs_dir, read_elf_file};
-use openvm_circuit::arch::{instructions::exe::VmExe, VmExecutor};
-use openvm_rv32im_circuit::Rv32ImConfig;
+use openvm_bigint_circuit::Int256Rv32Config;
+use openvm_bigint_transpiler::Int256TranspilerExtension;
+use openvm_circuit::arch::{instructions::exe::VmExe, VirtualMachine, VmExecutor};
 use openvm_rv32im_transpiler::{
     Rv32ITranspilerExtension, Rv32IoTranspilerExtension, Rv32MTranspilerExtension,
 };
-use openvm_stark_sdk::{bench::run_with_metric_collection, p3_baby_bear::BabyBear};
+use openvm_stark_sdk::{
+    bench::run_with_metric_collection, config::baby_bear_poseidon2::default_engine,
+    p3_baby_bear::BabyBear,
+};
 use openvm_transpiler::{transpiler::Transpiler, FromElf};
 
 #[derive(Debug, Clone, ValueEnum)]
@@ -22,13 +26,14 @@ static AVAILABLE_PROGRAMS: &[&str] = &[
     "fibonacci_iterative",
     "quicksort",
     "bubblesort",
+    "factorial_iterative_u256",
+    "revm_snailtracer",
     // "pairing",
     // "keccak256",
     // "keccak256_iter",
     // "sha256",
     // "sha256_iter",
     // "revm_transfer",
-    // "revm_snailtracer",
 ];
 
 #[derive(Parser)]
@@ -113,19 +118,40 @@ fn main() -> Result<()> {
             // let config_path = program_dir.join(DEFAULT_APP_CONFIG_PATH);
             // let vm_config = read_config_toml_or_default(&config_path)?.app_vm_config;
             // let transpiler = vm_config.transpiler;
-            let vm_config = Rv32ImConfig::default();
+            let vm_config = Int256Rv32Config::default();
 
             let transpiler = Transpiler::<BabyBear>::default()
                 .with_extension(Rv32ITranspilerExtension)
+                .with_extension(Rv32MTranspilerExtension)
                 .with_extension(Rv32IoTranspilerExtension)
-                .with_extension(Rv32MTranspilerExtension);
+                .with_extension(Int256TranspilerExtension);
 
             let exe = VmExe::from_elf(elf, transpiler)?;
 
+            let (widths, interactions): (Vec<usize>, Vec<usize>) = {
+                let vm = VirtualMachine::new(default_engine(), vm_config.clone());
+                let pk = vm.keygen();
+                let vk = pk.get_vk();
+                vk.inner
+                    .per_air
+                    .iter()
+                    .map(|vk| {
+                        // TODO(ayush): figure out which width to use
+                        // let total_width = vk.params.width.preprocessed.unwrap_or(0)
+                        //     + vk.params.width.cached_mains.iter().sum::<usize>()
+                        //     + vk.params.width.common_main
+                        //     + vk.params.width.after_challenge.iter().sum::<usize>();
+                        let total_width = vk.params.width.main_widths().iter().sum::<usize>();
+                        (total_width, vk.symbolic_constraints.interactions.len())
+                    })
+                    .unzip()
+            };
+
             let executor = VmExecutor::new(vm_config);
             executor
-                .execute_e1(exe, vec![])
+                .execute_metered(exe.clone(), vec![], widths, interactions)
                 .expect("Failed to execute program");
+
             tracing::info!("Completed program: {}", program);
         }
         tracing::info!("All programs executed successfully");

benchmarks/guest/factorial_iterative_u256/Cargo.toml

@@ -0,0 +1,17 @@
+[workspace]
+[package]
+name = "openvm-factorial-iterative-u256-program"
+version = "0.0.0"
+edition = "2021"
+
+[dependencies]
+openvm = { path = "../../../crates/toolchain/openvm", features = ["std"] }
+openvm-bigint-guest = { path = "../../../extensions/bigint/guest" }
+
+[features]
+default = []
+
+[profile.profiling]
+inherits = "release"
+debug = 2
+strip = false

benchmarks/guest/factorial_iterative_u256/openvm.toml

@@ -0,0 +1,4 @@
+[app_vm_config.rv32i]
+[app_vm_config.rv32m]
+[app_vm_config.io]
+[app_vm_config.bigint]

benchmarks/guest/factorial_iterative_u256/src/main.rs

@@ -0,0 +1,16 @@
+use core::hint::black_box;
+use openvm as _;
+use openvm_bigint_guest::U256;
+
+// This will overflow but that is fine
+const N: u32 = 65_000;
+
+pub fn main() {
+    let mut acc = U256::from_u32(1);
+    let mut i = U256::from_u32(N);
+    while i > black_box(U256::ZERO) {
+        acc *= i.clone();
+        i -= U256::from_u32(1);
+    }
+    black_box(acc);
+}

benchmarks/guest/fibonacci_iterative/src/main.rs

@@ -1,15 +1,15 @@
 use core::hint::black_box;
-use openvm as _;
+use openvm::io::reveal_u32;
 
-const N: u64 = 500_000;
+const N: u32 = 900_000;
 
 pub fn main() {
-    let mut a: u64 = 0;
-    let mut b: u64 = 1;
+    let mut a: u32 = 0;
+    let mut b: u32 = 1;
     for _ in 0..black_box(N) {
-        let c: u64 = a.wrapping_add(b);
+        let c: u32 = a.wrapping_add(b);
         a = b;
         b = c;
     }
-    black_box(a);
+    reveal_u32(a, 0);
 }

benchmarks/guest/fibonacci_recursive/src/main.rs

@@ -1,14 +1,15 @@
 use core::hint::black_box;
-use openvm as _;
+use openvm::io::reveal_u32;
 
-const N: u64 = 26;
+const N: u32 = 27;
 
 pub fn main() {
     let n = black_box(N);
-    black_box(fibonacci(n));
+    let result = fibonacci(n);
+    reveal_u32(result, 0);
 }
 
-fn fibonacci(n: u64) -> u64 {
+fn fibonacci(n: u32) -> u32 {
     if n == 0 {
         0
     } else if n == 1 {

crates/vm/derive/src/lib.rs

@@ -144,14 +144,27 @@ pub fn ins_executor_e1_executor_derive(input: TokenStream) -> TokenStream {
                 impl #impl_generics ::openvm_circuit::arch::InsExecutorE1<F> for #name #ty_generics #where_clause {
                     fn execute_e1<Ctx>(
                         &mut self,
-                        state: ::openvm_circuit::arch::VmStateMut<::openvm_circuit::system::memory::online::GuestMemory, Ctx>,
+                        state: &mut ::openvm_circuit::arch::VmStateMut<::openvm_circuit::system::memory::online::GuestMemory, Ctx>,
                         instruction: &::openvm_circuit::arch::instructions::instruction::Instruction<F>,
                     ) -> ::openvm_circuit::arch::Result<()>
                     where
-                        F: ::openvm_stark_backend::p3_field::PrimeField32
+                        F: ::openvm_stark_backend::p3_field::PrimeField32,
+                        Ctx: ::openvm_circuit::arch::execution_mode::E1E2ExecutionCtx,
                     {
                         self.0.execute_e1(state, instruction)
                     }
+
+                    fn execute_metered(
+                        &mut self,
+                        state: &mut ::openvm_circuit::arch::VmStateMut<::openvm_circuit::system::memory::online::GuestMemory, ::openvm_circuit::arch::execution_mode::metered::MeteredCtx>,
+                        instruction: &::openvm_circuit::arch::instructions::instruction::Instruction<F>,
+                        chip_index: usize,
+                    ) -> ::openvm_circuit::arch::Result<()>
+                    where
+                        F: ::openvm_stark_backend::p3_field::PrimeField32,
+                    {
+                        self.0.execute_metered(state, instruction, chip_index)
+                    }
                 }
             }
             .into()
@@ -180,25 +193,46 @@ pub fn ins_executor_e1_executor_derive(input: TokenStream) -> TokenStream {
                 .expect("First generic must be type for Field");
             // Use full path ::openvm_circuit... so it can be used either within or outside the vm
             // crate. Assume F is already generic of the field.
-            let execute_arms = variants.iter().map(|(variant_name, field)| {
+            let execute_e1_arms = variants.iter().map(|(variant_name, field)| {
                 let field_ty = &field.ty;
                 quote! {
                     #name::#variant_name(x) => <#field_ty as ::openvm_circuit::arch::InsExecutorE1<#first_ty_generic>>::execute_e1(x, state, instruction)
                 }
             }).collect::<Vec<_>>();
+            let execute_metered_arms = variants.iter().map(|(variant_name, field)| {
+                let field_ty = &field.ty;
+                quote! {
+                    #name::#variant_name(x) => <#field_ty as ::openvm_circuit::arch::InsExecutorE1<#first_ty_generic>>::execute_metered(x, state, instruction, chip_index)
+                }
+            }).collect::<Vec<_>>();
 
             quote! {
                 impl #impl_generics ::openvm_circuit::arch::InsExecutorE1<#first_ty_generic> for #name #ty_generics {
                     fn execute_e1<Ctx>(
                         &mut self,
-                        state: ::openvm_circuit::arch::VmStateMut<::openvm_circuit::system::memory::online::GuestMemory, Ctx>,
+                        state: &mut ::openvm_circuit::arch::VmStateMut<::openvm_circuit::system::memory::online::GuestMemory, Ctx>,
                         instruction: &::openvm_circuit::arch::instructions::instruction::Instruction<#first_ty_generic>,
                     ) -> ::openvm_circuit::arch::Result<()>
+                    where
+                        #first_ty_generic: ::openvm_stark_backend::p3_field::PrimeField32,
+                        Ctx: ::openvm_circuit::arch::execution_mode::E1E2ExecutionCtx,
+                    {
+                        match self {
+                            #(#execute_e1_arms,)*
+                        }
+                    }
+
+                    fn execute_metered(
+                        &mut self,
+                        state: &mut ::openvm_circuit::arch::VmStateMut<::openvm_circuit::system::memory::online::GuestMemory, ::openvm_circuit::arch::execution_mode::metered::MeteredCtx>,
+                        instruction: &::openvm_circuit::arch::instructions::instruction::Instruction<#first_ty_generic>,
+                        chip_index: usize,
+                    ) -> ::openvm_circuit::arch::Result<()>
                     where
                         #first_ty_generic: ::openvm_stark_backend::p3_field::PrimeField32
                     {
                         match self {
-                            #(#execute_arms,)*
+                            #(#execute_metered_arms,)*
                         }
                     }
                 }

crates/vm/src/arch/execution.rs

@@ -11,7 +11,10 @@ use openvm_stark_backend::{
 use serde::{Deserialize, Serialize};
 use thiserror::Error;
 
-use super::Streams;
+use super::{
+    execution_mode::{metered::MeteredCtx, E1E2ExecutionCtx},
+    Streams,
+};
 use crate::system::{
     memory::{
         online::{GuestMemory, TracingMemory},
@@ -115,8 +118,18 @@ pub trait InstructionExecutor<F> {
 pub trait InsExecutorE1<F> {
     fn execute_e1<Ctx>(
         &mut self,
-        state: VmStateMut<GuestMemory, Ctx>,
+        state: &mut VmStateMut<GuestMemory, Ctx>,
+        instruction: &Instruction<F>,
+    ) -> Result<()>
+    where
+        F: PrimeField32,
+        Ctx: E1E2ExecutionCtx;
+
+    fn execute_metered(
+        &mut self,
+        state: &mut VmStateMut<GuestMemory, MeteredCtx>,
         instruction: &Instruction<F>,
+        chip_index: usize,
     ) -> Result<()>
     where
         F: PrimeField32;
@@ -128,14 +141,28 @@ where
 {
     fn execute_e1<Ctx>(
         &mut self,
-        state: VmStateMut<GuestMemory, Ctx>,
+        state: &mut VmStateMut<GuestMemory, Ctx>,
         instruction: &Instruction<F>,
     ) -> Result<()>
     where
         F: PrimeField32,
+        Ctx: E1E2ExecutionCtx,
     {
         self.borrow_mut().execute_e1(state, instruction)
     }
+
+    fn execute_metered(
+        &mut self,
+        state: &mut VmStateMut<GuestMemory, MeteredCtx>,
+        instruction: &Instruction<F>,
+        chip_index: usize,
+    ) -> Result<()>
+    where
+        F: PrimeField32,
+    {
+        self.borrow_mut()
+            .execute_metered(state, instruction, chip_index)
+    }
 }
 
 impl<F, C: InstructionExecutor<F>> InstructionExecutor<F> for RefCell<C> {