stylus_benchmark

.dockerignore

@@ -36,6 +36,7 @@ arbitrator/tools/wasmer/target/
 arbitrator/tools/wasm-tools/
 arbitrator/tools/pricers/
 arbitrator/tools/module_roots/
+arbitrator/tools/stylus_benchmark
 arbitrator/langs/rust/target/
 arbitrator/langs/bf/target/
 

.github/workflows/arbitrator-ci.yml

@@ -171,6 +171,9 @@ jobs:
       - name: Rustfmt - langs/rust
         run: cargo fmt --all --manifest-path arbitrator/langs/rust/Cargo.toml -- --check
 
+      - name: Rustfmt - tools/stylus_benchmark
+        run: cargo fmt --all --manifest-path arbitrator/tools/stylus_benchmark/Cargo.toml -- --check
+
       - name: Make proofs from test cases
         run: make -j test-gen-proofs
 

arbitrator/Cargo.toml

@@ -12,6 +12,7 @@ members = [
 exclude = [
         "stylus/tests/",
         "tools/wasmer/",
+        "tools/stylus_benchmark",
 ]
 resolver = "2"
 

arbitrator/arbutil/src/benchmark.rs

@@ -0,0 +1,14 @@
+// Copyright 2024, Offchain Labs, Inc.
+// For license information, see https://github.com/nitro/blob/master/LICENSE
+
+use crate::evm::api::Ink;
+use std::time::{Duration, Instant};
+
+// Benchmark is used to track the performance of blocks of code in stylus
+#[derive(Clone, Copy, Debug, Default)]
+pub struct Benchmark {
+    pub timer: Option<Instant>,
+    pub elapsed_total: Duration,
+    pub ink_start: Option<Ink>,
+    pub ink_total: Ink,
+}

arbitrator/arbutil/src/lib.rs

@@ -1,6 +1,7 @@
 // Copyright 2022-2024, Offchain Labs, Inc.
 // For license information, see https://github.com/OffchainLabs/nitro/blob/master/LICENSE
 
+pub mod benchmark;
 /// cbindgen:ignore
 pub mod color;
 pub mod crypto;

arbitrator/jit/src/lib.rs

@@ -0,0 +1,51 @@
+// Copyright 2021-2024, Offchain Labs, Inc.
+// For license information, see https://github.com/OffchainLabs/nitro/blob/master/LICENSE
+
+use std::path::PathBuf;
+use structopt::StructOpt;
+
+mod arbcompress;
+mod caller_env;
+pub mod machine;
+mod prepare;
+pub mod program;
+mod socket;
+pub mod stylus_backend;
+mod test;
+mod wasip1_stub;
+mod wavmio;
+
+#[derive(StructOpt)]
+#[structopt(name = "jit-prover")]
+pub struct Opts {
+    #[structopt(short, long)]
+    binary: PathBuf,
+    #[structopt(long, default_value = "0")]
+    inbox_position: u64,
+    #[structopt(long, default_value = "0")]
+    delayed_inbox_position: u64,
+    #[structopt(long, default_value = "0")]
+    position_within_message: u64,
+    #[structopt(long)]
+    last_block_hash: Option<String>,
+    #[structopt(long)]
+    last_send_root: Option<String>,
+    #[structopt(long)]
+    inbox: Vec<PathBuf>,
+    #[structopt(long)]
+    delayed_inbox: Vec<PathBuf>,
+    #[structopt(long)]
+    preimages: Option<PathBuf>,
+    #[structopt(long)]
+    cranelift: bool,
+    #[structopt(long)]
+    forks: bool,
+    #[structopt(long)]
+    pub debug: bool,
+    #[structopt(long)]
+    pub require_success: bool,
+    // JSON inputs supercede any of the command-line inputs which could
+    // be specified in the JSON file.
+    #[structopt(long)]
+    json_inputs: Option<PathBuf>,
+}

arbitrator/jit/src/main.rs

@@ -1,58 +1,13 @@
 // Copyright 2022-2024, Offchain Labs, Inc.
 // For license information, see https://github.com/nitro/blob/master/LICENSE
 
-use crate::machine::{Escape, WasmEnv};
 use arbutil::{color, Color};
 use eyre::Result;
-use std::path::PathBuf;
+use jit::machine;
+use jit::machine::{Escape, WasmEnv};
+use jit::Opts;
 use structopt::StructOpt;
 
-mod arbcompress;
-mod caller_env;
-mod machine;
-mod prepare;
-mod program;
-mod socket;
-mod stylus_backend;
-mod test;
-mod wasip1_stub;
-mod wavmio;
-
-#[derive(StructOpt)]
-#[structopt(name = "jit-prover")]
-pub struct Opts {
-    #[structopt(short, long)]
-    binary: PathBuf,
-    #[structopt(long, default_value = "0")]
-    inbox_position: u64,
-    #[structopt(long, default_value = "0")]
-    delayed_inbox_position: u64,
-    #[structopt(long, default_value = "0")]
-    position_within_message: u64,
-    #[structopt(long)]
-    last_block_hash: Option<String>,
-    #[structopt(long)]
-    last_send_root: Option<String>,
-    #[structopt(long)]
-    inbox: Vec<PathBuf>,
-    #[structopt(long)]
-    delayed_inbox: Vec<PathBuf>,
-    #[structopt(long)]
-    preimages: Option<PathBuf>,
-    #[structopt(long)]
-    cranelift: bool,
-    #[structopt(long)]
-    forks: bool,
-    #[structopt(long)]
-    debug: bool,
-    #[structopt(long)]
-    require_success: bool,
-    // JSON inputs supercede any of the command-line inputs which could
-    // be specified in the JSON file.
-    #[structopt(long)]
-    json_inputs: Option<PathBuf>,
-}
-
 fn main() -> Result<()> {
     let opts = Opts::from_args();
     let env = match WasmEnv::cli(&opts) {

arbitrator/jit/src/prepare.rs

@@ -1,7 +1,7 @@
 // Copyright 2022-2024, Offchain Labs, Inc.
 // For license information, see https://github.com/nitro/blob/master/LICENSE
 
-use crate::WasmEnv;
+use crate::machine::WasmEnv;
 use arbutil::{Bytes32, PreimageType};
 use eyre::Ok;
 use prover::parse_input::FileData;

arbitrator/jit/src/program.rs

@@ -4,8 +4,8 @@
 #![allow(clippy::too_many_arguments)]
 
 use crate::caller_env::JitEnv;
-use crate::machine::{Escape, MaybeEscape, WasmEnvMut};
-use crate::stylus_backend::exec_wasm;
+use crate::machine::{Escape, MaybeEscape, WasmEnv, WasmEnvMut};
+use crate::stylus_backend::{exec_wasm, MessageFromCothread};
 use arbutil::evm::api::Gas;
 use arbutil::Bytes32;
 use arbutil::{evm::EvmData, format::DebugBytes, heapify};
@@ -16,6 +16,7 @@ use prover::{
     machine::Module,
     programs::{config::PricingParams, prelude::*},
 };
+use std::sync::Arc;
 
 const DEFAULT_STYLUS_ARBOS_VERSION: u64 = 31;
 
@@ -130,10 +131,6 @@ pub fn new_program(
     let evm_data: EvmData = unsafe { *Box::from_raw(evm_data_handler as *mut EvmData) };
     let config: JitConfig = unsafe { *Box::from_raw(stylus_config_handler as *mut JitConfig) };
 
-    // buy ink
-    let pricing = config.stylus.pricing;
-    let ink = pricing.gas_to_ink(Gas(gas));
-
     let Some(module) = exec.module_asms.get(&compiled_hash).cloned() else {
         return Err(Escape::Failure(format!(
             "module hash {:?} not found in {:?}",
@@ -142,6 +139,21 @@ pub fn new_program(
         )));
     };
 
+    exec_program(exec, module, calldata, config, evm_data, gas)
+}
+
+pub fn exec_program(
+    exec: &mut WasmEnv,
+    module: Arc<[u8]>,
+    calldata: Vec<u8>,
+    config: JitConfig,
+    evm_data: EvmData,
+    gas: u64,
+) -> Result<u32, Escape> {
+    // buy ink
+    let pricing = config.stylus.pricing;
+    let ink = pricing.gas_to_ink(Gas(gas));
+
     let cothread = exec_wasm(
         module,
         calldata,
@@ -162,7 +174,10 @@ pub fn new_program(
 /// returns request_id for the first request from the program
 pub fn start_program(mut env: WasmEnvMut, module: u32) -> Result<u32, Escape> {
     let (_, exec) = env.jit_env();
+    start_program_with_wasm_env(exec, module)
+}
 
+pub fn start_program_with_wasm_env(exec: &mut WasmEnv, module: u32) -> Result<u32, Escape> {
     if exec.threads.len() as u32 != module || module == 0 {
         return Escape::hostio(format!(
             "got request for thread {module} but len is {}",
@@ -179,26 +194,27 @@ pub fn start_program(mut env: WasmEnvMut, module: u32) -> Result<u32, Escape> {
 /// request_id MUST be last request id returned from start_program or send_response
 pub fn get_request(mut env: WasmEnvMut, id: u32, len_ptr: GuestPtr) -> Result<u32, Escape> {
     let (mut mem, exec) = env.jit_env();
+    let msg = get_last_msg(exec, id)?;
+    mem.write_u32(len_ptr, msg.req_data.len() as u32);
+    Ok(msg.req_type)
+}
+
+pub fn get_last_msg(exec: &mut WasmEnv, id: u32) -> Result<MessageFromCothread, Escape> {
     let thread = exec.threads.last_mut().unwrap();
     let msg = thread.last_message()?;
     if msg.1 != id {
         return Escape::hostio("get_request id doesn't match");
     };
-    mem.write_u32(len_ptr, msg.0.req_data.len() as u32);
-    Ok(msg.0.req_type)
+    Ok(msg.0)
 }
 
 // gets data associated with last request.
 // request_id MUST be last request receieved
 // data_ptr MUST point to a buffer of at least the length returned by get_request
 pub fn get_request_data(mut env: WasmEnvMut, id: u32, data_ptr: GuestPtr) -> MaybeEscape {
     let (mut mem, exec) = env.jit_env();
-    let thread = exec.threads.last_mut().unwrap();
-    let msg = thread.last_message()?;
-    if msg.1 != id {
-        return Escape::hostio("get_request id doesn't match");
-    };
-    mem.write_slice(data_ptr, &msg.0.req_data);
+    let msg = get_last_msg(exec, id)?;
+    mem.write_slice(data_ptr, &msg.req_data);
     Ok(())
 }
 
@@ -217,11 +233,21 @@ pub fn set_response(
     let result = mem.read_slice(result_ptr, result_len as usize);
     let raw_data = mem.read_slice(raw_data_ptr, raw_data_len as usize);
 
+    set_response_with_wasm_env(exec, id, gas, result, raw_data)
+}
+
+pub fn set_response_with_wasm_env(
+    exec: &mut WasmEnv,
+    id: u32,
+    gas: u64,
+    result: Vec<u8>,
+    raw_data: Vec<u8>,
+) -> MaybeEscape {
     let thread = exec.threads.last_mut().unwrap();
     thread.set_response(id, result, raw_data, Gas(gas))
 }
 
-/// sends previos response
+/// sends previous response
 /// MUST be called right after set_response to the same id
 /// returns request_id for the next request
 pub fn send_response(mut env: WasmEnvMut, req_id: u32) -> Result<u32, Escape> {
@@ -239,16 +265,19 @@ pub fn send_response(mut env: WasmEnvMut, req_id: u32) -> Result<u32, Escape> {
 /// removes the last created program
 pub fn pop(mut env: WasmEnvMut) -> MaybeEscape {
     let (_, exec) = env.jit_env();
+    pop_with_wasm_env(exec)
+}
 
+pub fn pop_with_wasm_env(exec: &mut WasmEnv) -> MaybeEscape {
     match exec.threads.pop() {
         None => Err(Escape::Child(eyre!("no child"))),
         Some(mut thread) => thread.wait_done(),
     }
 }
 
 pub struct JitConfig {
-    stylus: StylusConfig,
-    compile: CompileConfig,
+    pub stylus: StylusConfig,
+    pub compile: CompileConfig,
 }
 
 /// Creates a `StylusConfig` from its component parts.

arbitrator/jit/src/stylus_backend.rs

@@ -4,6 +4,7 @@
 #![allow(clippy::too_many_arguments)]
 
 use crate::machine::{Escape, MaybeEscape};
+use arbutil::benchmark::Benchmark;
 use arbutil::evm::api::{Gas, Ink, VecReader};
 use arbutil::evm::{
     api::{EvmApiMethod, EVM_API_METHOD_REQ_OFFSET},
@@ -35,6 +36,7 @@ struct MessageToCothread {
 pub struct MessageFromCothread {
     pub req_type: u32,
     pub req_data: Vec<u8>,
+    pub benchmark: Benchmark,
 }
 
 struct CothreadRequestor {
@@ -51,6 +53,7 @@ impl RequestHandler<VecReader> for CothreadRequestor {
         let msg = MessageFromCothread {
             req_type: req_type as u32 + EVM_API_METHOD_REQ_OFFSET,
             req_data: req_data.as_ref().to_vec(),
+            benchmark: Benchmark::default(),
         };
 
         if let Err(error) = self.tx.send(msg) {
@@ -169,6 +172,7 @@ pub fn exec_wasm(
         let msg = MessageFromCothread {
             req_data: output,
             req_type: out_kind as u32,
+            benchmark: instance.env().benchmark,
         };
         instance
             .env_mut()

arbitrator/stylus/src/env.rs

@@ -2,6 +2,7 @@
 // For license information, see https://github.com/nitro/blob/master/LICENSE
 
 use arbutil::{
+    benchmark::Benchmark,
     evm::{
         api::{DataReader, EvmApi, Ink},
         EvmData,
@@ -48,6 +49,8 @@ pub struct WasmEnv<D: DataReader, E: EvmApi<D>> {
     pub compile: CompileConfig,
     /// The runtime config
     pub config: Option<StylusConfig>,
+    // Used to benchmark execution blocks of code
+    pub benchmark: Benchmark,
     // Using the unused generic parameter D in a PhantomData field
     _data_reader_marker: PhantomData<D>,
 }
@@ -68,6 +71,7 @@ impl<D: DataReader, E: EvmApi<D>> WasmEnv<D, E> {
             outs: vec![],
             memory: None,
             meter: None,
+            benchmark: Benchmark::default(),
             _data_reader_marker: PhantomData,
         }
     }