perf: only do exit check for instructions that can exit (#2126)

- update `create_tco_handler` macro to `create_handler` which now
automatically sets `exit_code` based on whether the execute impl returns
`Result::Err`. it acts as a simple wrapper for execute impls that don't
return a `Result`
- only do exit checks for executors that can exit in tco mode i.e. for
execute impls that return `Result`
- feature gate all non-tco functions with `#[cfg(not(feature = "tco"))]`

[benchmark
comparison](https://github.com/axiom-crypto/openvm-reth-benchmark/actions/runs/17590682650#summary-49971529404)

Towards INT-4921

Author: shuklaayush

crates/toolchain/tests/Cargo.toml

@@ -38,6 +38,7 @@ num-bigint.workspace = true
 default = ["parallel"]
 parallel = ["openvm-circuit/parallel"]
 cuda = [ "openvm-rv32im-circuit/cuda" ]
+tco = ["openvm-circuit/tco"]
 touchemall = [ 
     "cuda",
     "openvm-circuit/touchemall",

crates/vm/derive/src/common.rs

@@ -0,0 +1,67 @@
+use quote::{format_ident, quote};
+use syn::{Ident, ItemFn};
+
+/// Extract the first two generic type parameters (F and CTX) from function generics
+pub fn extract_f_and_ctx_types(generics: &syn::Generics) -> (syn::Ident, syn::Ident) {
+    let mut type_params = generics.params.iter().filter_map(|param| {
+        if let syn::GenericParam::Type(type_param) = param {
+            Some(&type_param.ident)
+        } else {
+            None
+        }
+    });
+
+    let f_type = type_params
+        .next()
+        .expect("Function must have at least one type parameter (F)")
+        .clone();
+    let ctx_type = type_params
+        .next()
+        .expect("Function must have at least two type parameters (F and CTX)")
+        .clone();
+
+    (f_type, ctx_type)
+}
+
+/// Build a list of generic arguments for function calls
+pub fn build_generic_args(generics: &syn::Generics) -> Vec<proc_macro2::TokenStream> {
+    generics
+        .params
+        .iter()
+        .map(|param| match param {
+            syn::GenericParam::Type(type_param) => {
+                let ident = &type_param.ident;
+                quote! { #ident }
+            }
+            syn::GenericParam::Lifetime(lifetime) => {
+                let lifetime = &lifetime.lifetime;
+                quote! { #lifetime }
+            }
+            syn::GenericParam::Const(const_param) => {
+                let ident = &const_param.ident;
+                quote! { #ident }
+            }
+        })
+        .collect()
+}
+
+/// Generate handler name from function name:
+/// If original ends with `_impl`, replace with `_handler`, else append `_handler` suffix.
+pub fn handler_name_from_fn(fn_name: &Ident) -> Ident {
+    let new_name_str = fn_name
+        .to_string()
+        .strip_suffix("_impl")
+        .map(|base| format!("{base}_handler"))
+        .unwrap_or_else(|| format!("{fn_name}_handler"));
+    format_ident!("{}", new_name_str)
+}
+
+/// Check if function returns Result type
+pub fn returns_result_type(input_fn: &ItemFn) -> bool {
+    match &input_fn.sig.output {
+        syn::ReturnType::Type(_, ty) => {
+            matches!(**ty, syn::Type::Path(ref path) if path.path.segments.last().is_some_and(|seg| seg.ident == "Result"))
+        }
+        _ => false,
+    }
+}

crates/vm/derive/src/lib.rs

@@ -9,6 +9,9 @@ use syn::{
     GenericParam, Ident, Meta, Token,
 };
 
+mod common;
+#[cfg(not(feature = "tco"))]
+mod nontco;
 #[cfg(feature = "tco")]
 mod tco;
 
@@ -48,7 +51,7 @@ pub fn preflight_executor_derive(input: TokenStream) -> TokenStream {
                 _ => panic!("Only unnamed fields are supported"),
             };
             // Use full path ::openvm_circuit... so it can be used either within or outside the vm
-            // crate.
+            // crate. Assume F is already generic of the field.
             let where_clause = new_generics.make_where_clause();
             where_clause.predicates.push(
                 syn::parse_quote! { #inner_ty: ::openvm_circuit::arch::PreflightExecutor<#field_ty_generic, RA> },
@@ -85,7 +88,7 @@ pub fn preflight_executor_derive(input: TokenStream) -> TokenStream {
                 })
                 .collect::<Vec<_>>();
             // 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.
+            // crate.
             let (execute_arms, get_opcode_name_arms, where_predicates): (Vec<_>, Vec<_>, Vec<_>) =
                 multiunzip(variants.iter().map(|(variant_name, field)| {
                     let field_ty = &field.ty;
@@ -183,6 +186,7 @@ pub fn executor_derive(input: TokenStream) -> TokenStream {
                     fn pre_compute_size(&self) -> usize {
                         self.0.pre_compute_size()
                     }
+                    #[cfg(not(feature = "tco"))]
                     #[inline(always)]
                     fn pre_compute<Ctx>(
                         &self,
@@ -281,6 +285,7 @@ pub fn executor_derive(input: TokenStream) -> TokenStream {
                         }
                     }
 
+                    #[cfg(not(feature = "tco"))]
                     #[inline(always)]
                     fn pre_compute<Ctx>(
                         &self,
@@ -325,7 +330,7 @@ pub fn metered_executor_derive(input: TokenStream) -> TokenStream {
                 _ => 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.
+            // crate.
             let mut new_generics = generics.clone();
             let where_clause = new_generics.make_where_clause();
             where_clause
@@ -357,6 +362,7 @@ pub fn metered_executor_derive(input: TokenStream) -> TokenStream {
                     fn metered_pre_compute_size(&self) -> usize {
                         self.0.metered_pre_compute_size()
                     }
+                    #[cfg(not(feature = "tco"))]
                     #[inline(always)]
                     fn metered_pre_compute<Ctx>(
                         &self,
@@ -457,6 +463,7 @@ pub fn metered_executor_derive(input: TokenStream) -> TokenStream {
                         }
                     }
 
+                    #[cfg(not(feature = "tco"))]
                     #[inline(always)]
                     fn metered_pre_compute<Ctx>(
                         &self,
@@ -832,13 +839,13 @@ fn parse_executor_type(
 /// unless the execute_impl returns an error. This is done for performance to skip an exit code
 /// check.
 #[proc_macro_attribute]
-pub fn create_tco_handler(_attr: TokenStream, item: TokenStream) -> TokenStream {
+pub fn create_handler(_attr: TokenStream, item: TokenStream) -> TokenStream {
     #[cfg(feature = "tco")]
     {
         tco::tco_impl(item)
     }
     #[cfg(not(feature = "tco"))]
     {
-        item
+        nontco::nontco_impl(item)
     }
 }

crates/vm/derive/src/nontco.rs

@@ -0,0 +1,83 @@
+use proc_macro::TokenStream;
+use quote::quote;
+use syn::{parse_macro_input, ItemFn};
+
+use crate::common::{
+    build_generic_args, extract_f_and_ctx_types, handler_name_from_fn, returns_result_type,
+};
+
+/// Implementation of the non-TCO handler generation logic.
+/// This is called from the proc macro attribute in lib.rs.
+pub fn nontco_impl(item: TokenStream) -> TokenStream {
+    // Parse the input function
+    let input_fn = parse_macro_input!(item as ItemFn);
+
+    // Extract information from the function
+    let fn_name = &input_fn.sig.ident;
+    let generics = &input_fn.sig.generics;
+    let where_clause = &generics.where_clause;
+
+    // Check if function returns Result
+    let returns_result = returns_result_type(&input_fn);
+
+    // Extract the first two generic type parameters (F and CTX)
+    let (f_type, ctx_type) = extract_f_and_ctx_types(generics);
+
+    // Derive new function name:
+    // If original ends with `_impl`, replace with `_handler`, else append suffix.
+    let handler_name = handler_name_from_fn(fn_name);
+
+    // Build the function call with all the generics
+    let generic_args = build_generic_args(generics);
+    let execute_call = if generic_args.is_empty() {
+        quote! { #fn_name(pre_compute, instret, pc, arg, exec_state) }
+    } else {
+        quote! { #fn_name::<#(#generic_args),*>(pre_compute, instret, pc, arg, exec_state) }
+    };
+
+    // Generate the execute and exit check code based on return type
+    let handler_body = if returns_result {
+        quote! {
+            // Call original impl and wire errors into exit_code.
+            let __ret = { #execute_call };
+            if let ::core::result::Result::Err(e) = __ret {
+                exec_state.set_instret_and_pc(*instret, *pc);
+                exec_state.exit_code = ::core::result::Result::Err(e);
+                return;
+            }
+        }
+    } else {
+        quote! {
+            #execute_call;
+        }
+    };
+
+    // Generate the non-TCO handler function
+    let handler_fn = quote! {
+        #[inline(always)]
+        unsafe fn #handler_name #generics (
+            pre_compute: &[u8],
+            instret: &mut u64,
+            pc: &mut u32,
+            arg: u64,
+            exec_state: &mut ::openvm_circuit::arch::VmExecState<
+                #f_type,
+                ::openvm_circuit::system::memory::online::GuestMemory,
+                #ctx_type,
+            >,
+        )
+        #where_clause
+        {
+            #handler_body
+        }
+    };
+
+    // Return both the original function and the new handler
+    let output = quote! {
+        #input_fn
+
+        #handler_fn
+    };
+
+    TokenStream::from(output)
+}

crates/vm/derive/src/tco.rs

@@ -1,7 +1,11 @@
 use proc_macro::TokenStream;
-use quote::{format_ident, quote};
+use quote::quote;
 use syn::{parse_macro_input, ItemFn};
 
+use crate::common::{
+    build_generic_args, extract_f_and_ctx_types, handler_name_from_fn, returns_result_type,
+};
+
 /// Implementation of the TCO handler generation logic.
 /// This is called from the proc macro attribute in lib.rs.
 pub fn tco_impl(item: TokenStream) -> TokenStream {
@@ -13,16 +17,15 @@ pub fn tco_impl(item: TokenStream) -> TokenStream {
     let generics = &input_fn.sig.generics;
     let where_clause = &generics.where_clause;
 
+    // Check if function returns Result
+    let returns_result = returns_result_type(&input_fn);
+
     // Extract the first two generic type parameters (F and CTX)
     let (f_type, ctx_type) = extract_f_and_ctx_types(generics);
+
     // Derive new function name:
-    // If original ends with `_impl`, replace with `_tco_handler`, else append suffix.
-    let new_name_str = fn_name
-        .to_string()
-        .strip_suffix("_impl")
-        .map(|base| format!("{base}_tco_handler"))
-        .unwrap_or_else(|| format!("{fn_name}_tco_handler"));
-    let handler_name = format_ident!("{}", new_name_str);
+    // If original ends with `_impl`, replace with `_handler`, else append suffix.
+    let handler_name = handler_name_from_fn(fn_name);
 
     // Build the generic parameters for the handler, preserving all original generics
     let handler_generics = generics.clone();
@@ -35,6 +38,21 @@ pub fn tco_impl(item: TokenStream) -> TokenStream {
         quote! { #fn_name::<#(#generic_args),*>(pre_compute, &mut instret, &mut pc, arg, exec_state) }
     };
 
+    // Generate the execute and exit check code based on return type
+    let execute_stmt = if returns_result {
+        quote! {
+            // Call original impl and wire errors into exit_code.
+            let __ret = { #execute_call };
+            if let ::core::result::Result::Err(e) = __ret {
+                exec_state.set_instret_and_pc(instret, pc);
+                exec_state.exit_code = ::core::result::Result::Err(e);
+                return;
+            }
+        }
+    } else {
+        quote! { #execute_call; }
+    };
+
     // Generate the TCO handler function
     let handler_fn = quote! {
         #[inline(never)]
@@ -54,12 +72,8 @@ pub fn tco_impl(item: TokenStream) -> TokenStream {
             use ::openvm_circuit::arch::ExecutionError;
 
             let pre_compute = interpreter.get_pre_compute(pc);
-            #execute_call;
+            #execute_stmt
 
-            if ::core::intrinsics::unlikely(exec_state.exit_code.is_err()) {
-                exec_state.set_instret_and_pc(instret, pc);
-                return;
-            }
             if ::core::intrinsics::unlikely(#ctx_type::should_suspend(instret, pc, arg, exec_state)) {
                 exec_state.set_instret_and_pc(instret, pc);
                 return;
@@ -89,45 +103,3 @@ pub fn tco_impl(item: TokenStream) -> TokenStream {
 
     TokenStream::from(output)
 }
-
-fn extract_f_and_ctx_types(generics: &syn::Generics) -> (syn::Ident, syn::Ident) {
-    let mut type_params = generics.params.iter().filter_map(|param| {
-        if let syn::GenericParam::Type(type_param) = param {
-            Some(&type_param.ident)
-        } else {
-            None
-        }
-    });
-
-    let f_type = type_params
-        .next()
-        .expect("Function must have at least one type parameter (F)")
-        .clone();
-    let ctx_type = type_params
-        .next()
-        .expect("Function must have at least two type parameters (F and CTX)")
-        .clone();
-
-    (f_type, ctx_type)
-}
-
-fn build_generic_args(generics: &syn::Generics) -> Vec<proc_macro2::TokenStream> {
-    generics
-        .params
-        .iter()
-        .map(|param| match param {
-            syn::GenericParam::Type(type_param) => {
-                let ident = &type_param.ident;
-                quote! { #ident }
-            }
-            syn::GenericParam::Lifetime(lifetime) => {
-                let lifetime = &lifetime.lifetime;
-                quote! { #lifetime }
-            }
-            syn::GenericParam::Const(const_param) => {
-                let ident = &const_param.ident;
-                quote! { #ident }
-            }
-        })
-        .collect()
-}

crates/vm/src/arch/execution.rs

@@ -129,6 +129,7 @@ pub type Handler<F, CTX> = unsafe fn(
 pub trait Executor<F> {
     fn pre_compute_size(&self) -> usize;
 
+    #[cfg(not(feature = "tco"))]
     fn pre_compute<Ctx>(
         &self,
         pc: u32,
@@ -160,6 +161,7 @@ pub trait Executor<F> {
 pub trait MeteredExecutor<F> {
     fn metered_pre_compute_size(&self) -> usize;
 
+    #[cfg(not(feature = "tco"))]
     fn metered_pre_compute<Ctx>(
         &self,
         air_idx: usize,