feat(new-execution): rv32im tracegen and e1 execution (#1607)

Author: shuklaayush

.gitignore

@@ -36,3 +36,6 @@ metrics.json
 
 # openvm generated files
 crates/cli/openvm/
+
+# samply profile
+profile.json.gz

Cargo.lock

@@ -9,15 +9,13 @@ license.workspace = true
 
 [dependencies]
 openvm-benchmarks-utils.workspace = true
-cargo-openvm.workspace = true
 openvm-circuit.workspace = true
-openvm-sdk.workspace = true
 openvm-stark-sdk.workspace = true
 openvm-transpiler.workspace = true
 openvm-rv32im-circuit.workspace = true
 openvm-rv32im-transpiler.workspace = true
-openvm-keccak256-circuit.workspace = true
-openvm-keccak256-transpiler.workspace = true
+# openvm-keccak256-circuit.workspace = true
+# openvm-keccak256-transpiler.workspace = true
 
 clap = { version = "4.5.9", features = ["derive", "env"] }
 eyre.workspace = true
@@ -30,13 +28,13 @@ tracing-subscriber = { version = "0.3.17", features = ["std", "env-filter"] }
 criterion = { version = "0.5", features = ["html_reports"] }
 
 [features]
-default = ["mimalloc"]
-profiling = ["openvm-sdk/profiling"]
+default = ["jemalloc"]
 aggregation = []
 mimalloc = ["openvm-circuit/mimalloc"]
 jemalloc = ["openvm-circuit/jemalloc"]
 jemalloc-prof = ["openvm-circuit/jemalloc-prof"]
 nightly-features = ["openvm-circuit/nightly-features"]
+profiling = ["openvm-circuit/function-span", "openvm-transpiler/function-span"]
 
 [[bench]]
 name = "fibonacci_execute"

benchmarks/execute/Cargo.toml

@@ -1,30 +1,34 @@
-use cargo_openvm::{default::DEFAULT_APP_CONFIG_PATH, util::read_config_toml_or_default};
 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_sdk::StdIn;
-use openvm_stark_sdk::bench::run_with_metric_collection;
-use openvm_transpiler::FromElf;
+use openvm_rv32im_circuit::Rv32ImConfig;
+use openvm_rv32im_transpiler::{
+    Rv32ITranspilerExtension, Rv32IoTranspilerExtension, Rv32MTranspilerExtension,
+};
+use openvm_stark_sdk::{bench::run_with_metric_collection, p3_baby_bear::BabyBear};
+use openvm_transpiler::{transpiler::Transpiler, FromElf};
 
 #[derive(Debug, Clone, ValueEnum)]
 enum BuildProfile {
     Debug,
     Release,
 }
 
+// const DEFAULT_APP_CONFIG_PATH: &str = "./openvm.toml";
+
 static AVAILABLE_PROGRAMS: &[&str] = &[
-    "fibonacci_recursive",
-    "fibonacci_iterative",
+    // "fibonacci_recursive",
+    // "fibonacci_iterative",
     "quicksort",
-    "bubblesort",
-    "pairing",
-    "keccak256",
-    "keccak256_iter",
-    "sha256",
-    "sha256_iter",
-    "revm_transfer",
-    "revm_snailtracer",
+    // "bubblesort",
+    // "pairing",
+    // "keccak256",
+    // "keccak256_iter",
+    // "sha256",
+    // "sha256_iter",
+    // "revm_transfer",
+    // "revm_snailtracer",
 ];
 
 #[derive(Parser)]
@@ -106,13 +110,22 @@ fn main() -> Result<()> {
             let elf_path = get_elf_path(&program_dir);
             let elf = read_elf_file(&elf_path)?;
 
-            let config_path = program_dir.join(DEFAULT_APP_CONFIG_PATH);
-            let vm_config = read_config_toml_or_default(&config_path)?.app_vm_config;
+            // 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 transpiler = Transpiler::<BabyBear>::default()
+                .with_extension(Rv32ITranspilerExtension)
+                .with_extension(Rv32IoTranspilerExtension)
+                .with_extension(Rv32MTranspilerExtension);
 
-            let exe = VmExe::from_elf(elf, vm_config.transpiler())?;
+            let exe = VmExe::from_elf(elf, transpiler)?;
 
             let executor = VmExecutor::new(vm_config);
-            executor.execute(exe, StdIn::default())?;
+            executor
+                .execute_e1(exe, vec![])
+                .expect("Failed to execute program");
             tracing::info!("Completed program: {}", program);
         }
         tracing::info!("All programs executed successfully");

benchmarks/execute/src/main.rs

@@ -1,12 +1,9 @@
 use itertools::Itertools;
 use num_bigint::BigUint;
 use num_traits::Zero;
-use openvm_circuit::{
-    arch::{
-        AdapterAirContext, AdapterRuntimeContext, DynAdapterInterface, DynArray, InsExecutorE1,
-        MinimalInstruction, Result, VmAdapterInterface, VmCoreAir, VmCoreChip, VmExecutionState,
-    },
-    system::memory::online::GuestMemory,
+use openvm_circuit::arch::{
+    AdapterAirContext, AdapterRuntimeContext, DynAdapterInterface, DynArray, MinimalInstruction,
+    Result, VmAdapterInterface, VmCoreAir, VmCoreChip,
 };
 use openvm_circuit_primitives::{
     var_range::SharedVariableRangeCheckerChip, SubAir, TraceSubRowGenerator,
@@ -314,20 +311,6 @@ where
     }
 }
 
-impl<Mem, Ctx, F> InsExecutorE1<Mem, Ctx, F> for FieldExpressionCoreChip
-where
-    Mem: GuestMemory,
-    F: PrimeField32,
-{
-    fn execute_e1(
-        &mut self,
-        _state: &mut VmExecutionState<Mem, Ctx>,
-        _instruction: &Instruction<F>,
-    ) -> Result<()> {
-        todo!("Implement execute_e1")
-    }
-}
-
 impl FieldExpressionCoreChip {
     // We will be setting is_valid = 0. That forces all flags be 0 (otherwise setup will be -1).
     // We generate a dummy row with all flags set to 0, then we set is_valid = 0.

crates/circuits/mod-builder/src/core_chip.rs

@@ -113,6 +113,104 @@ pub fn instruction_executor_derive(input: TokenStream) -> TokenStream {
     }
 }
 
+#[proc_macro_derive(InsExecutorE1)]
+pub fn ins_executor_e1_executor_derive(input: TokenStream) -> TokenStream {
+    let ast: syn::DeriveInput = syn::parse(input).unwrap();
+
+    let name = &ast.ident;
+    let generics = &ast.generics;
+    let (impl_generics, ty_generics, _) = generics.split_for_impl();
+
+    match &ast.data {
+        Data::Struct(inner) => {
+            // Check if the struct has only one unnamed field
+            let inner_ty = match &inner.fields {
+                Fields::Unnamed(fields) => {
+                    if fields.unnamed.len() != 1 {
+                        panic!("Only one unnamed field is supported");
+                    }
+                    fields.unnamed.first().unwrap().ty.clone()
+                }
+                _ => panic!("Only unnamed fields are supported"),
+            };
+            // 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 mut new_generics = generics.clone();
+            let where_clause = new_generics.make_where_clause();
+            where_clause
+                .predicates
+                .push(syn::parse_quote! { #inner_ty: ::openvm_circuit::arch::InsExecutorE1<F> });
+            quote! {
+                impl #impl_generics ::openvm_circuit::arch::InsExecutorE1<F> for #name #ty_generics #where_clause {
+                    fn execute_e1<Mem, Ctx>(
+                        &mut self,
+                        state: ::openvm_circuit::arch::execution::VmStateMut<Mem, Ctx>,
+                        instruction: &::openvm_circuit::arch::instructions::instruction::Instruction<F>,
+                    ) -> ::openvm_circuit::arch::Result<()>
+                    where
+                        Mem: ::openvm_circuit::system::memory::online::GuestMemory,
+                        F: ::openvm_stark_backend::p3_field::PrimeField32
+                    {
+                        self.0.execute_e1(state, instruction)
+                    }
+                }
+            }
+            .into()
+        }
+        Data::Enum(e) => {
+            let variants = e
+                .variants
+                .iter()
+                .map(|variant| {
+                    let variant_name = &variant.ident;
+
+                    let mut fields = variant.fields.iter();
+                    let field = fields.next().unwrap();
+                    assert!(fields.next().is_none(), "Only one field is supported");
+                    (variant_name, field)
+                })
+                .collect::<Vec<_>>();
+            let first_ty_generic = ast
+                .generics
+                .params
+                .first()
+                .and_then(|param| match param {
+                    GenericParam::Type(type_param) => Some(&type_param.ident),
+                    _ => None,
+                })
+                .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 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<_>>();
+
+            quote! {
+                impl #impl_generics ::openvm_circuit::arch::InsExecutorE1<#first_ty_generic> for #name #ty_generics {
+                    fn execute_e1<Mem, Ctx>(
+                        &mut self,
+                        state: ::openvm_circuit::arch::execution::VmStateMut<Mem, Ctx>,
+                        instruction: &::openvm_circuit::arch::instructions::instruction::Instruction<#first_ty_generic>,
+                    ) -> ::openvm_circuit::arch::Result<()>
+                    where
+                        Mem: ::openvm_circuit::system::memory::online::GuestMemory,
+                        #first_ty_generic: ::openvm_stark_backend::p3_field::PrimeField32
+                    {
+                        match self {
+                            #(#execute_arms,)*
+                        }
+                    }
+                }
+            }
+            .into()
+        }
+        Data::Union(_) => unimplemented!("Unions are not supported"),
+    }
+}
+
 /// Derives `AnyEnum` trait on an enum type.
 /// By default an enum arm will just return `self` as `&dyn Any`.
 ///
@@ -347,7 +445,7 @@ pub fn vm_generic_config_derive(input: proc_macro::TokenStream) -> proc_macro::T
             let periphery_type = Ident::new(&format!("{}Periphery", name), name.span());
 
             TokenStream::from(quote! {
-                #[derive(::openvm_circuit::circuit_derive::ChipUsageGetter, ::openvm_circuit::circuit_derive::Chip, ::openvm_circuit::derive::InstructionExecutor, ::derive_more::derive::From, ::openvm_circuit::derive::AnyEnum)]
+                #[derive(::openvm_circuit::circuit_derive::ChipUsageGetter, ::openvm_circuit::circuit_derive::Chip, ::openvm_circuit::derive::InstructionExecutor, ::openvm_circuit::derive::InsExecutorE1, ::derive_more::derive::From, ::openvm_circuit::derive::AnyEnum)]
                 pub enum #executor_type<F: PrimeField32> {
                     #[any_enum]
                     #source_name_upper(#source_executor_type<F>),

crates/vm/derive/src/lib.rs

@@ -11,7 +11,7 @@ use openvm_stark_backend::{
 use serde::{Deserialize, Serialize};
 use thiserror::Error;
 
-use super::{Streams, VmExecutionState};
+use super::Streams;
 use crate::system::{
     memory::{
         online::{GuestMemory, TracingMemory},
@@ -79,6 +79,7 @@ pub enum ExecutionError {
 /// Global VM state accessible during instruction execution.
 /// The state is generic in guest memory `MEM` and additional host state `CTX`.
 /// The host state is execution context specific.
+#[derive(derive_new::new)]
 pub struct VmStateMut<'a, MEM, CTX> {
     pub pc: &'a mut u32,
     pub memory: &'a mut MEM,
@@ -111,16 +112,32 @@ pub trait InstructionExecutor<F> {
 }
 
 /// New trait for instruction execution
-pub trait InsExecutorE1<Mem, Ctx, F>
+pub trait InsExecutorE1<F> {
+    fn execute_e1<Mem, Ctx>(
+        &mut self,
+        state: VmStateMut<Mem, Ctx>,
+        instruction: &Instruction<F>,
+    ) -> Result<()>
+    where
+        Mem: GuestMemory,
+        F: PrimeField32;
+}
+
+impl<F, C> InsExecutorE1<F> for RefCell<C>
 where
-    Mem: GuestMemory,
-    F: PrimeField32,
+    C: InsExecutorE1<F>,
 {
-    fn execute_e1(
+    fn execute_e1<Mem, Ctx>(
         &mut self,
-        state: &mut VmExecutionState<Mem, Ctx>,
+        state: VmStateMut<Mem, Ctx>,
         instruction: &Instruction<F>,
-    ) -> Result<()>;
+    ) -> Result<()>
+    where
+        Mem: GuestMemory,
+        F: PrimeField32,
+    {
+        self.borrow_mut().execute_e1(state, instruction)
+    }
 }
 
 impl<F, C: InstructionExecutor<F>> InstructionExecutor<F> for RefCell<C> {