Merge pull request #2937 from fitzgen/bench-api-stdio-and-repeated-instantiations

bench-api: pass in explicit stdio files, allow repeated instantiations per compilation

crates/bench-api/.gitignore

@@ -0,0 +1,2 @@
+stdout.log
+stderr.log

crates/bench-api/Cargo.toml

@@ -31,4 +31,4 @@ wat = "1.0"
 default = ["shuffling-allocator"]
 wasi-crypto = ["wasmtime-wasi-crypto"]
 wasi-nn = ["wasmtime-wasi-nn"]
-
+old-x86-backend = ["wasmtime/old-x86-backend"]

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`
+//! ```text
-//!  - `wasm_bench_compile`
+//!               |
-//!  - `wasm_bench_instantiate`
+//!               |
-//!  - `wasm_bench_execute`
+//!               V
-//!  - `wasm_bench_free`
+//! .---> wasm_bench_create
-//!
+//! |        |        |
-//! You may repeat this sequence of calls multiple times to take multiple
+//! |        |        |
-//! measurements of compilation, instantiation, and execution time within a
+//! |        |        V
-//! single process.
+//! |        |   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.
 //!
@@ -25,9 +42,59 @@
 //! use wasmtime_bench_api::*;
 //!
 //! let working_dir = std::env::current_dir().unwrap().display().to_string();
+//! 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(),
+//!     stdout_path_ptr: stdout_path.as_ptr(),
+//!     stdout_path_len: stdout_path.len(),
+//!     stderr_path_ptr: stderr_path.as_ptr(),
+//!     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();
 //! unsafe {
-//!     let code = wasm_bench_create(working_dir.as_ptr(), working_dir.len(), &mut bench_api);
+//!     let code = wasm_bench_create(config, &mut bench_api);
 //!     assert_eq!(code, OK);
 //!     assert!(!bench_api.is_null());
 //! };
@@ -47,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"
+//! // This will call the `instantiation_{start,end}` timing functions on success.
-//! // "start") and stop ("bench" "stop") the measurement counters/timers during
+//! let code = unsafe { wasm_bench_instantiate(bench_api) };
-//! // 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.
 //! assert_eq!(code, OK);
 //!
-//! // No need to start timers for the execution since, by convention, the timer
+//! // This will call the `execution_{start,end}` timing functions on success.
-//! // functions we passed during instantiation will be called by the benchmark
-//! // at the appropriate time (before and after the benchmarked section).
 //! let code = unsafe { wasm_bench_execute(bench_api) };
 //! assert_eq!(code, OK);
 //!
@@ -78,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 std::{env, path::PathBuf};
-use wasmtime_wasi::sync::{Wasi, WasiCtxBuilder};
+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;
@@ -98,6 +156,83 @@ pub const ERR: ExitCode = -1;
 static ALLOC: shuffling_allocator::ShufflingAllocator<std::alloc::System> =
     shuffling_allocator::wrap!(&std::alloc::System);
 
+/// Configuration options for the benchmark.
+#[repr(C)]
+pub struct WasmBenchConfig {
+    /// The working directory where benchmarks should be executed.
+    pub working_dir_ptr: *const u8,
+    pub working_dir_len: usize,
+
+    /// The file path that should be created and used as `stdout`.
+    pub stdout_path_ptr: *const u8,
+    pub stdout_path_len: usize,
+
+    /// The file path that should be created and used as `stderr`.
+    pub stderr_path_ptr: *const u8,
+    pub stderr_path_len: usize,
+
+    /// The (optional) file path that should be opened and used as `stdin`. If
+    /// 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<PathBuf> {
+        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.into())
+    }
+
+    fn stdout_path(&self) -> Result<PathBuf> {
+        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.into())
+    }
+
+    fn stderr_path(&self) -> Result<PathBuf> {
+        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.into())
+    }
+
+    fn stdin_path(&self) -> Result<Option<PathBuf>> {
+        if self.stdin_path_ptr.is_null() {
+            return Ok(None);
+        }
+
+        let stdin_path =
+            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.into()))
+    }
+}
+
 /// Exposes a C-compatible way of creating the engine from the bytes of a single
 /// Wasm module.
 ///
@@ -107,15 +242,69 @@ static ALLOC: shuffling_allocator::ShufflingAllocator<std::alloc::System> =
 /// untouched.
 #[no_mangle]
 pub extern "C" fn wasm_bench_create(
-    working_dir_ptr: *const u8,
+    config: WasmBenchConfig,
-    working_dir_len: usize,
     out_bench_ptr: *mut *mut c_void,
 ) -> ExitCode {
     let result = (|| -> Result<_> {
-        let working_dir = unsafe { std::slice::from_raw_parts(working_dir_ptr, working_dir_len) };
+        let working_dir = config.working_dir()?;
-        let working_dir = std::str::from_utf8(working_dir)
+        let working_dir = unsafe { cap_std::fs::Dir::open_ambient_dir(&working_dir) }
-            .context("given working directory is not valid UTF-8")?;
+            .with_context(|| {
-        let state = Box::new(BenchState::new(working_dir)?);
+                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(
+            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 _)
     })();
 
@@ -153,15 +342,9 @@ pub extern "C" fn wasm_bench_compile(
 
 /// Instantiate the Wasm benchmark module.
 #[no_mangle]
-pub extern "C" fn wasm_bench_instantiate(
+pub extern "C" fn wasm_bench_instantiate(state: *mut c_void) -> ExitCode {
-    state: *mut c_void,
-    bench_start: extern "C" fn(),
-    bench_end: extern "C" fn(),
-) -> ExitCode {
     let state = unsafe { (state as *mut BenchState).as_mut().unwrap() };
-    let result = state
+    let result = state.instantiate().context("failed to instantiate");
-        .instantiate(bench_start, bench_end)
-        .context("failed to instantiate");
     to_exit_code(result)
 }
 
@@ -190,39 +373,107 @@ fn to_exit_code<T>(result: impl Into<Result<T>>) -> 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<dyn FnMut() -> Result<WasiCtx>>,
     module: Option<Module>,
     instance: Option<Instance>,
-    did_execute: bool,
 }
 
 impl BenchState {
-    fn new(working_dir: impl AsRef<Path>) -> Result<Self> {
+    fn new(
+        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<WasiCtx> + 'static,
+    ) -> Result<Self> {
+        // 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);
-        cx = cx.inherit_stdio();
+        let module = Module::from_binary(&self.engine, bytes)?;
-        // Allow access to the working directory so that the benchmark can read
+        (self.compilation_end)(self.compilation_timer);
-        // 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, ".")?;
-        // 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)?;
-        }
 
-        Wasi::new(linker.store(), cx.build()).add_to_linker(&mut linker)?;
+        self.module = Some(module);
+        Ok(())
+    }
+
+    fn instantiate(&mut self) -> Result<()> {
+        let module = self
+            .module
+            .as_ref()
+            .expect("compile the module before instantiating it");
+
+        let wasi_cx = (self.make_wasi_cx)().context("failed to create a WASI context")?;
+
+        // 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);
+
+        let store = Store::new(&self.engine);
+        assert!(Wasi::set_context(&store, wasi_cx).is_ok());
+
+        let linker = Linker::new(&store);
 
         #[cfg(feature = "wasi-nn")]
         {
@@ -251,53 +502,17 @@ impl BenchState {
             WasiCryptoSymmetric::new(linker.store(), cx_crypto).add_to_linker(linker)?;
         }
 
-        Ok(Self {
+        let instance = linker.instantiate(&module)?;
-            engine,
+        (self.instantiation_end)(self.instantiation_timer);
-            linker,
-            module: None,
-            instance: None,
-            did_execute: false,
-        })
-    }
 
-    fn compile(&mut self, bytes: &[u8]) -> Result<()> {
+        self.instance = Some(instance);
-        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)?);
         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")?;

crates/bench-api/src/unsafe_send_sync.rs

@@ -0,0 +1,19 @@
+#[derive(Clone, Copy)]
+pub struct UnsafeSendSync<T>(T);
+
+impl<T> UnsafeSendSync<T> {
+    /// 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<T> Send for UnsafeSendSync<T> {}
+unsafe impl<T> Sync for UnsafeSendSync<T> {}