diff --git a/crates/bench-api/src/lib.rs b/crates/bench-api/src/lib.rs index 5ae6a05ecf..711fbd11a4 100644 --- a/crates/bench-api/src/lib.rs +++ b/crates/bench-api/src/lib.rs @@ -1,17 +1,34 @@ //! A C API for benchmarking Wasmtime's WebAssembly compilation, instantiation, //! and execution. //! -//! The API expects sequential calls to: +//! The API expects calls that match the following state machine: //! -//! - `wasm_bench_create` -//! - `wasm_bench_compile` -//! - `wasm_bench_instantiate` -//! - `wasm_bench_execute` -//! - `wasm_bench_free` -//! -//! You may repeat this sequence of calls multiple times to take multiple -//! measurements of compilation, instantiation, and execution time within a -//! single process. +//! ```text +//! | +//! | +//! V +//! .---> wasm_bench_create +//! | | | +//! | | | +//! | | V +//! | | wasm_bench_compile +//! | | | | +//! | | | | .----. +//! | | | | | | +//! | | | V V | +//! | | | wasm_bench_instantiate <------. +//! | | | | | | +//! | | | | | | +//! | | | | | | +//! | | | .------' '-----> wasm_bench_execute +//! | | | | | +//! | | | | | +//! | V V V | +//! '------ wasm_bench_free <--------------------------' +//! | +//! | +//! V +//! ``` //! //! All API calls must happen on the same thread. //! @@ -28,6 +45,33 @@ //! let stdout_path = "./stdout.log"; //! let stderr_path = "./stderr.log"; //! +//! // Functions to start/end timers for compilation. +//! // +//! // The `compilation_timer` pointer configured in the `WasmBenchConfig` is +//! // passed through. +//! extern "C" fn compilation_start(timer: *mut u8) { +//! // Start your compilation timer here. +//! } +//! extern "C" fn compilation_end(timer: *mut u8) { +//! // End your compilation timer here. +//! } +//! +//! // Similar for instantiation. +//! extern "C" fn instantiation_start(timer: *mut u8) { +//! // Start your instantiation timer here. +//! } +//! extern "C" fn instantiation_end(timer: *mut u8) { +//! // End your instantiation timer here. +//! } +//! +//! // Similar for execution. +//! extern "C" fn execution_start(timer: *mut u8) { +//! // Start your execution timer here. +//! } +//! extern "C" fn execution_end(timer: *mut u8) { +//! // End your execution timer here. +//! } +//! //! let config = WasmBenchConfig { //! working_dir_ptr: working_dir.as_ptr(), //! working_dir_len: working_dir.len(), @@ -37,6 +81,15 @@ //! stderr_path_len: stderr_path.len(), //! stdin_path_ptr: ptr::null(), //! stdin_path_len: 0, +//! compilation_timer: ptr::null_mut(), +//! compilation_start, +//! compilation_end, +//! instantiation_timer: ptr::null_mut(), +//! instantiation_start, +//! instantiation_end, +//! execution_timer: ptr::null_mut(), +//! execution_start, +//! execution_end, //! }; //! //! let mut bench_api = ptr::null_mut(); @@ -61,29 +114,15 @@ //! ) //! "#).unwrap(); //! -//! // Start your compilation timer here. +//! // This will call the `compilation_{start,end}` timing functions on success. //! let code = unsafe { wasm_bench_compile(bench_api, wasm.as_ptr(), wasm.len()) }; -//! // End your compilation timer here. //! assert_eq!(code, OK); //! -//! // The Wasm benchmark will expect us to provide functions to start ("bench" -//! // "start") and stop ("bench" "stop") the measurement counters/timers during -//! // execution. -//! extern "C" fn bench_start() { -//! // Start your execution timer here. -//! } -//! extern "C" fn bench_stop() { -//! // End your execution timer here. -//! } -//! -//! // Start your instantiation timer here. -//! let code = unsafe { wasm_bench_instantiate(bench_api, bench_start, bench_stop) }; -//! // End your instantiation timer here. +//! // This will call the `instantiation_{start,end}` timing functions on success. +//! let code = unsafe { wasm_bench_instantiate(bench_api) }; //! assert_eq!(code, OK); //! -//! // No need to start timers for the execution since, by convention, the timer -//! // functions we passed during instantiation will be called by the benchmark -//! // at the appropriate time (before and after the benchmarked section). +//! // This will call the `execution_{start,end}` timing functions on success. //! let code = unsafe { wasm_bench_execute(bench_api) }; //! assert_eq!(code, OK); //! @@ -92,13 +131,18 @@ //! } //! ``` +mod unsafe_send_sync; + +use crate::unsafe_send_sync::UnsafeSendSync; use anyhow::{anyhow, Context, Result}; -use std::env; use std::os::raw::{c_int, c_void}; -use std::path::Path; use std::slice; -use wasmtime::{Config, Engine, Instance, Linker, Module, Store}; -use wasmtime_wasi::sync::{Wasi, WasiCtxBuilder}; +use std::{env, path::PathBuf}; +use wasmtime::{Config, Engine, FuncType, Instance, Linker, Module, Store}; +use wasmtime_wasi::{ + sync::{Wasi, WasiCtxBuilder}, + WasiCtx, +}; pub type ExitCode = c_int; pub const OK: ExitCode = 0; @@ -131,34 +175,52 @@ pub struct WasmBenchConfig { /// not provided, then the WASI context will not have a `stdin` initialized. pub stdin_path_ptr: *const u8, pub stdin_path_len: usize, + + /// The functions to start and stop performance timers/counters during Wasm + /// compilation. + pub compilation_timer: *mut u8, + pub compilation_start: extern "C" fn(*mut u8), + pub compilation_end: extern "C" fn(*mut u8), + + /// The functions to start and stop performance timers/counters during Wasm + /// instantiation. + pub instantiation_timer: *mut u8, + pub instantiation_start: extern "C" fn(*mut u8), + pub instantiation_end: extern "C" fn(*mut u8), + + /// The functions to start and stop performance timers/counters during Wasm + /// execution. + pub execution_timer: *mut u8, + pub execution_start: extern "C" fn(*mut u8), + pub execution_end: extern "C" fn(*mut u8), } impl WasmBenchConfig { - fn working_dir(&self) -> Result<&str> { + fn working_dir(&self) -> Result { let working_dir = unsafe { std::slice::from_raw_parts(self.working_dir_ptr, self.working_dir_len) }; let working_dir = std::str::from_utf8(working_dir) .context("given working directory is not valid UTF-8")?; - Ok(working_dir) + Ok(working_dir.into()) } - fn stdout_path(&self) -> Result<&str> { + fn stdout_path(&self) -> Result { let stdout_path = unsafe { std::slice::from_raw_parts(self.stdout_path_ptr, self.stdout_path_len) }; let stdout_path = std::str::from_utf8(stdout_path).context("given stdout path is not valid UTF-8")?; - Ok(stdout_path) + Ok(stdout_path.into()) } - fn stderr_path(&self) -> Result<&str> { + fn stderr_path(&self) -> Result { let stderr_path = unsafe { std::slice::from_raw_parts(self.stderr_path_ptr, self.stderr_path_len) }; let stderr_path = std::str::from_utf8(stderr_path).context("given stderr path is not valid UTF-8")?; - Ok(stderr_path) + Ok(stderr_path.into()) } - fn stdin_path(&self) -> Result> { + fn stdin_path(&self) -> Result> { if self.stdin_path_ptr.is_null() { return Ok(None); } @@ -167,7 +229,7 @@ impl WasmBenchConfig { unsafe { std::slice::from_raw_parts(self.stdin_path_ptr, self.stdin_path_len) }; let stdin_path = std::str::from_utf8(stdin_path).context("given stdin path is not valid UTF-8")?; - Ok(Some(stdin_path)) + Ok(Some(stdin_path.into())) } } @@ -185,14 +247,63 @@ pub extern "C" fn wasm_bench_create( ) -> ExitCode { let result = (|| -> Result<_> { let working_dir = config.working_dir()?; + let working_dir = unsafe { cap_std::fs::Dir::open_ambient_dir(&working_dir) } + .with_context(|| { + format!( + "failed to preopen the working directory: {}", + working_dir.display(), + ) + })?; + let stdout_path = config.stdout_path()?; let stderr_path = config.stderr_path()?; let stdin_path = config.stdin_path()?; + let state = Box::new(BenchState::new( - working_dir, - stdout_path, - stderr_path, - stdin_path, + config.compilation_timer, + config.compilation_start, + config.compilation_end, + config.instantiation_timer, + config.instantiation_start, + config.instantiation_end, + config.execution_timer, + config.execution_start, + config.execution_end, + move || { + let mut cx = WasiCtxBuilder::new(); + + let stdout = std::fs::File::create(&stdout_path) + .with_context(|| format!("failed to create {}", stdout_path.display()))?; + let stdout = unsafe { cap_std::fs::File::from_std(stdout) }; + let stdout = wasi_cap_std_sync::file::File::from_cap_std(stdout); + cx = cx.stdout(Box::new(stdout)); + + let stderr = std::fs::File::create(&stderr_path) + .with_context(|| format!("failed to create {}", stderr_path.display()))?; + let stderr = unsafe { cap_std::fs::File::from_std(stderr) }; + let stderr = wasi_cap_std_sync::file::File::from_cap_std(stderr); + cx = cx.stderr(Box::new(stderr)); + + if let Some(stdin_path) = &stdin_path { + let stdin = std::fs::File::open(stdin_path) + .with_context(|| format!("failed to open {}", stdin_path.display()))?; + let stdin = unsafe { cap_std::fs::File::from_std(stdin) }; + let stdin = wasi_cap_std_sync::file::File::from_cap_std(stdin); + cx = cx.stdin(Box::new(stdin)); + } + + // Allow access to the working directory so that the benchmark can read + // its input workload(s). + cx = cx.preopened_dir(working_dir.try_clone()?, ".")?; + + // Pass this env var along so that the benchmark program can use smaller + // input workload(s) if it has them and that has been requested. + if let Ok(val) = env::var("WASM_BENCH_USE_SMALL_WORKLOAD") { + cx = cx.env("WASM_BENCH_USE_SMALL_WORKLOAD", &val)?; + } + + Ok(cx.build()) + }, )?); Ok(Box::into_raw(state) as _) })(); @@ -231,15 +342,9 @@ pub extern "C" fn wasm_bench_compile( /// Instantiate the Wasm benchmark module. #[no_mangle] -pub extern "C" fn wasm_bench_instantiate( - state: *mut c_void, - bench_start: extern "C" fn(), - bench_end: extern "C" fn(), -) -> ExitCode { +pub extern "C" fn wasm_bench_instantiate(state: *mut c_void) -> ExitCode { let state = unsafe { (state as *mut BenchState).as_mut().unwrap() }; - let result = state - .instantiate(bench_start, bench_end) - .context("failed to instantiate"); + let result = state.instantiate().context("failed to instantiate"); to_exit_code(result) } @@ -268,65 +373,107 @@ fn to_exit_code(result: impl Into>) -> ExitCode { /// to manage the Wasmtime engine between calls. struct BenchState { engine: Engine, - linker: Linker, + compilation_timer: *mut u8, + compilation_start: extern "C" fn(*mut u8), + compilation_end: extern "C" fn(*mut u8), + instantiation_timer: *mut u8, + instantiation_start: extern "C" fn(*mut u8), + instantiation_end: extern "C" fn(*mut u8), + make_wasi_cx: Box Result>, module: Option, instance: Option, - did_execute: bool, } impl BenchState { fn new( - working_dir: impl AsRef, - stdout: impl AsRef, - stderr: impl AsRef, - stdin: Option>, + compilation_timer: *mut u8, + compilation_start: extern "C" fn(*mut u8), + compilation_end: extern "C" fn(*mut u8), + instantiation_timer: *mut u8, + instantiation_start: extern "C" fn(*mut u8), + instantiation_end: extern "C" fn(*mut u8), + execution_timer: *mut u8, + execution_start: extern "C" fn(*mut u8), + execution_end: extern "C" fn(*mut u8), + make_wasi_cx: impl FnMut() -> Result + 'static, ) -> Result { + // NB: do not configure a code cache. let mut config = Config::new(); config.wasm_simd(true); - // NB: do not configure a code cache. + Wasi::add_to_config(&mut config); + + // Define the benchmarking start/end functions. + let execution_timer = unsafe { + // Safe because this bench API's contract requires that its methods + // are only ever called from a single thread. + UnsafeSendSync::new(execution_timer) + }; + config.define_host_func( + "bench", + "start", + FuncType::new(vec![], vec![]), + move |_, _, _| { + execution_start(*execution_timer.get()); + Ok(()) + }, + ); + config.define_host_func( + "bench", + "end", + FuncType::new(vec![], vec![]), + move |_, _, _| { + execution_end(*execution_timer.get()); + Ok(()) + }, + ); let engine = Engine::new(&config)?; - let store = Store::new(&engine); - let mut linker = Linker::new(&store); + Ok(Self { + engine, + compilation_timer, + compilation_start, + compilation_end, + instantiation_timer, + instantiation_start, + instantiation_end, + make_wasi_cx: Box::new(make_wasi_cx) as _, + module: None, + instance: None, + }) + } - // Create a WASI environment. + fn compile(&mut self, bytes: &[u8]) -> Result<()> { + assert!( + self.module.is_none(), + "create a new engine to repeat compilation" + ); - let mut cx = WasiCtxBuilder::new(); + (self.compilation_start)(self.compilation_timer); + let module = Module::from_binary(&self.engine, bytes)?; + (self.compilation_end)(self.compilation_timer); - let stdout = std::fs::File::create(stdout.as_ref()) - .with_context(|| format!("failed to create {}", stdout.as_ref().display()))?; - let stdout = unsafe { cap_std::fs::File::from_std(stdout) }; - let stdout = wasi_cap_std_sync::file::File::from_cap_std(stdout); - cx = cx.stdout(Box::new(stdout)); + self.module = Some(module); + Ok(()) + } - let stderr = std::fs::File::create(stderr.as_ref()) - .with_context(|| format!("failed to create {}", stderr.as_ref().display()))?; - let stderr = unsafe { cap_std::fs::File::from_std(stderr) }; - let stderr = wasi_cap_std_sync::file::File::from_cap_std(stderr); - cx = cx.stderr(Box::new(stderr)); + fn instantiate(&mut self) -> Result<()> { + let module = self + .module + .as_ref() + .expect("compile the module before instantiating it"); - if let Some(stdin) = stdin { - let stdin = std::fs::File::open(stdin.as_ref()) - .with_context(|| format!("failed to open {}", stdin.as_ref().display()))?; - let stdin = unsafe { cap_std::fs::File::from_std(stdin) }; - let stdin = wasi_cap_std_sync::file::File::from_cap_std(stdin); - cx = cx.stdin(Box::new(stdin)); - } + let wasi_cx = (self.make_wasi_cx)().context("failed to create a WASI context")?; - // Allow access to the working directory so that the benchmark can read - // its input workload(s). - let working_dir = unsafe { cap_std::fs::Dir::open_ambient_dir(working_dir) } - .context("failed to preopen the working directory")?; - cx = cx.preopened_dir(working_dir, ".")?; + // NB: Start measuring instantiation time *after* we've created the WASI + // context, since that needs to do file I/O to setup + // stdin/stdout/stderr. + (self.instantiation_start)(self.instantiation_timer); - // Pass this env var along so that the benchmark program can use smaller - // input workload(s) if it has them and that has been requested. - if let Ok(val) = env::var("WASM_BENCH_USE_SMALL_WORKLOAD") { - cx = cx.env("WASM_BENCH_USE_SMALL_WORKLOAD", &val)?; - } + let store = Store::new(&self.engine); + assert!(Wasi::set_context(&store, wasi_cx).is_ok()); - Wasi::new(linker.store(), cx.build()).add_to_linker(&mut linker)?; + let linker = Linker::new(&store); #[cfg(feature = "wasi-nn")] { @@ -355,53 +502,17 @@ impl BenchState { WasiCryptoSymmetric::new(linker.store(), cx_crypto).add_to_linker(linker)?; } - Ok(Self { - engine, - linker, - module: None, - instance: None, - did_execute: false, - }) - } + let instance = linker.instantiate(&module)?; + (self.instantiation_end)(self.instantiation_timer); - fn compile(&mut self, bytes: &[u8]) -> Result<()> { - assert!( - self.module.is_none(), - "create a new engine to repeat compilation" - ); - self.module = Some(Module::from_binary(&self.engine, bytes)?); - Ok(()) - } - - fn instantiate( - &mut self, - bench_start: extern "C" fn(), - bench_end: extern "C" fn(), - ) -> Result<()> { - assert!( - self.instance.is_none(), - "create a new engine to repeat instantiation" - ); - let module = self - .module - .as_mut() - .expect("compile the module before instantiating it"); - - // Import the specialized benchmarking functions. - self.linker.func("bench", "start", move || bench_start())?; - self.linker.func("bench", "end", move || bench_end())?; - - self.instance = Some(self.linker.instantiate(&module)?); + self.instance = Some(instance); Ok(()) } fn execute(&mut self) -> Result<()> { - assert!(!self.did_execute, "create a new engine to repeat execution"); - self.did_execute = true; - let instance = self .instance - .as_ref() + .take() .expect("instantiate the module before executing it"); let start_func = instance.get_typed_func::<(), ()>("_start")?; diff --git a/crates/bench-api/src/unsafe_send_sync.rs b/crates/bench-api/src/unsafe_send_sync.rs new file mode 100644 index 0000000000..d948b62d70 --- /dev/null +++ b/crates/bench-api/src/unsafe_send_sync.rs @@ -0,0 +1,19 @@ +#[derive(Clone, Copy)] +pub struct UnsafeSendSync(T); + +impl UnsafeSendSync { + /// Create a new `UnsafeSendSync` wrapper around the given value. + /// + /// The result is a type that is `Send` and `Sync` regardless of whether `T: + /// Send + Sync`, so this constructor is unsafe. + pub unsafe fn new(val: T) -> Self { + UnsafeSendSync(val) + } + + pub fn get(&self) -> &T { + &self.0 + } +} + +unsafe impl Send for UnsafeSendSync {} +unsafe impl Sync for UnsafeSendSync {}