Allow benchmarking just one particular compilation/instantiation/execution phase

This avoids recompiling the Wasm module many times when we are only interested
in benchmarking instantiation or execution. Note that it is still compiled once
per process, but you can compile it exactly once if you do `--processes 1`.

Author: fitzgen

crates/cli/src/benchmark.rs

@@ -5,7 +5,7 @@ use sightglass_data::{Format, Measurement, Phase};
 use sightglass_recorder::bench_api::Engine;
 use sightglass_recorder::cpu_affinity::bind_to_single_core;
 use sightglass_recorder::measure::Measurements;
-use sightglass_recorder::{bench_api::BenchApi, benchmark::benchmark, measure::MeasureType};
+use sightglass_recorder::{bench_api::BenchApi, benchmark, measure::MeasureType};
 use std::{
     fs,
     io::{self, BufWriter, Write},
@@ -106,9 +106,10 @@ pub struct BenchmarkCommand {
     #[structopt(short("d"), long("working-dir"), parse(from_os_str))]
     working_dir: Option<PathBuf>,
 
-    /// Stop measuring after the given phase (compilation/instantiation/execution).
-    #[structopt(long("stop-after"))]
-    stop_after_phase: Option<Phase>,
+    /// Benchmark only the given phase (compilation, instantiation, or
+    /// execution). Benchmarks all phases if omitted.
+    #[structopt(long("benchmark-phase"))]
+    benchmark_phase: Option<Phase>,
 
     /// The significance level for confidence intervals. Typical values are 0.01
     /// and 0.05, which correspond to 99% and 95% confidence respectively. This
@@ -195,6 +196,8 @@ impl BenchmarkCommand {
                 let mut measurements = Measurements::new(this_arch(), engine, wasm_file);
                 let mut measure = self.measure.build();
 
+                // Create the bench API engine and cache it for reuse across all
+                // iterations of this benchmark.
                 let engine = Engine::new(
                     &mut bench_api,
                     &working_dir,
@@ -207,35 +210,68 @@ impl BenchmarkCommand {
                 );
                 let mut engine = Some(engine);
 
+                // And if we are benchmarking just a post-compilation phase,
+                // then eagerly compile the Wasm module for reuse.
+                let mut module = None;
+                if let Some(Phase::Instantiation | Phase::Execution) = self.benchmark_phase {
+                    module = Some(engine.take().unwrap().compile(&bytes));
+                }
+
                 // Run the benchmark (compilation, instantiation, and execution) several times in
                 // this process.
                 for _ in 0..self.iterations_per_process {
-                    let new_engine = benchmark(
-                        engine.take().unwrap(),
-                        &bytes,
-                        self.stop_after_phase.clone(),
-                    )?;
-                    engine = Some(new_engine);
+                    match self.benchmark_phase {
+                        None => {
+                            let new_engine = benchmark::all(engine.take().unwrap(), &bytes)?;
+                            engine = Some(new_engine);
+                        }
+                        Some(Phase::Compilation) => {
+                            let new_engine =
+                                benchmark::compilation(engine.take().unwrap(), &bytes)?;
+                            engine = Some(new_engine);
+                        }
+                        Some(Phase::Instantiation) => {
+                            let new_module = benchmark::instantiation(module.take().unwrap())?;
+                            module = Some(new_module);
+                        }
+                        Some(Phase::Execution) => {
+                            let new_module = benchmark::execution(module.take().unwrap())?;
+                            module = Some(new_module);
+                        }
+                    }
 
                     self.check_output(Path::new(wasm_file), stdout, stderr)?;
-                    engine.as_mut().unwrap().measurements().next_iteration();
+                    engine
+                        .as_mut()
+                        .map(|e| e.measurements())
+                        .or_else(|| module.as_mut().map(|m| m.measurements()))
+                        .unwrap()
+                        .next_iteration();
                 }
 
-                drop(engine);
+                drop((engine, module));
                 all_measurements.extend(measurements.finish());
             }
         }
 
+        // If we are only benchmarking one phase then filter out any
+        // measurements for other phases. These get included because we have to
+        // compile at least once to measure instantiation, for example.
+        if let Some(phase) = self.benchmark_phase {
+            all_measurements.retain(|m| m.phase == phase);
+        }
+
         self.write_results(&all_measurements, &mut output_file)?;
         Ok(())
     }
 
     /// Assert that our actual `stdout` and `stderr` match our expectations.
     fn check_output(&self, wasm_file: &Path, stdout: &Path, stderr: &Path) -> Result<()> {
-        // If we aren't going through all phases and executing the Wasm, then we
-        // won't have any actual output to check.
-        if self.stop_after_phase.is_some() {
-            return Ok(());
+        match self.benchmark_phase {
+            None | Some(Phase::Execution) => {}
+            // If we aren't executing the Wasm, then we won't have any actual
+            // output to check.
+            Some(Phase::Compilation | Phase::Instantiation) => return Ok(()),
         }
 
         let wasm_file_dir: PathBuf = if let Some(dir) = wasm_file.parent() {
@@ -328,8 +364,8 @@ impl BenchmarkCommand {
                 command.env("WASM_BENCH_USE_SMALL_WORKLOAD", "1");
             }
 
-            if let Some(phase) = self.stop_after_phase {
-                command.arg("--stop-after").arg(phase.to_string());
+            if let Some(phase) = self.benchmark_phase {
+                command.arg("--benchmark-phase").arg(phase.to_string());
             }
 
             if let Some(flags) = &self.engine_flags {

crates/cli/tests/all/benchmark.rs

@@ -5,14 +5,14 @@ use sightglass_data::Measurement;
 use std::path::PathBuf;
 
 #[test]
-fn benchmark_stop_after_compilation() {
+fn benchmark_phase_compilation() {
     sightglass_cli_benchmark()
         .arg("--raw")
         .arg("--processes")
         .arg("2")
         .arg("--iterations-per-process")
         .arg("1")
-        .arg("--stop-after")
+        .arg("--benchmark-phase")
         .arg("compilation")
         .arg(benchmark("noop"))
         .assert()
@@ -25,25 +25,47 @@ fn benchmark_stop_after_compilation() {
 }
 
 #[test]
-fn benchmark_stop_after_instantiation() {
+fn benchmark_phase_instantiation() {
     sightglass_cli_benchmark()
         .arg("--raw")
         .arg("--processes")
         .arg("2")
         .arg("--iterations-per-process")
         .arg("1")
-        .arg("--stop-after")
+        .arg("--benchmark-phase")
         .arg("instantiation")
         .arg(benchmark("noop"))
         .assert()
         .success()
         .stdout(
             predicate::str::contains("Compilation")
+                .not()
                 .and(predicate::str::contains("Instantiation"))
                 .and(predicate::str::contains("Execution").not()),
         );
 }
 
+#[test]
+fn benchmark_phase_execution() {
+    sightglass_cli_benchmark()
+        .arg("--raw")
+        .arg("--processes")
+        .arg("2")
+        .arg("--iterations-per-process")
+        .arg("1")
+        .arg("--benchmark-phase")
+        .arg("execution")
+        .arg(benchmark("noop"))
+        .assert()
+        .success()
+        .stdout(
+            predicate::str::contains("Compilation")
+                .not()
+                .and(predicate::str::contains("Instantiation").not())
+                .and(predicate::str::contains("Execution")),
+        );
+}
+
 #[test]
 fn benchmark_json() {
     let assert = sightglass_cli_benchmark()

crates/recorder/src/bench_api.rs

@@ -74,10 +74,7 @@ pub struct Engine<'a, 'b, 'c, M> {
     engine: *mut c_void,
 }
 
-impl<'a, 'b, 'c, M> Engine<'a, 'b, 'c, M>
-where
-    M: Measure,
-{
+impl<'a, 'b, 'c, M> Engine<'a, 'b, 'c, M> {
     /// Construct a new engine from the given `BenchApi`.
     // NB: take a mutable reference to the `BenchApi` so that no one else can
     // call its API methods out of order.
@@ -90,7 +87,10 @@ where
         measurements: &'a mut Measurements<'c>,
         measure: &'a mut M,
         execution_flags: Option<&'a str>,
-    ) -> Self {
+    ) -> Self
+    where
+        M: Measure,
+    {
         let working_dir = working_dir.display().to_string();
         let stdout_path = stdout_path.display().to_string();
         let stderr_path = stderr_path.display().to_string();
@@ -148,31 +148,43 @@ where
     }
 
     /// Bench API callback for the start of compilation.
-    extern "C" fn compilation_start(data: *mut u8) {
+    extern "C" fn compilation_start(data: *mut u8)
+    where
+        M: Measure,
+    {
         log::debug!("Starting compilation measurement");
         let data = data as *mut (*mut M, *mut Measurements<'b>);
         let measure = unsafe { data.as_mut().unwrap().0.as_mut().unwrap() };
         measure.start(Phase::Compilation);
     }
 
     /// Bench API callback for the start of instantiation.
-    extern "C" fn instantiation_start(data: *mut u8) {
+    extern "C" fn instantiation_start(data: *mut u8)
+    where
+        M: Measure,
+    {
         log::debug!("Starting instantiation measurement");
         let data = data as *mut (*mut M, *mut Measurements<'b>);
         let measure = unsafe { data.as_mut().unwrap().0.as_mut().unwrap() };
         measure.start(Phase::Instantiation);
     }
 
     /// Bench API callback for the start of execution.
-    extern "C" fn execution_start(data: *mut u8) {
+    extern "C" fn execution_start(data: *mut u8)
+    where
+        M: Measure,
+    {
         log::debug!("Starting execution measurement");
         let data = data as *mut (*mut M, *mut Measurements<'b>);
         let measure = unsafe { data.as_mut().unwrap().0.as_mut().unwrap() };
         measure.start(Phase::Execution);
     }
 
     /// Bench API callback for the end of compilation.
-    extern "C" fn compilation_end(data: *mut u8) {
+    extern "C" fn compilation_end(data: *mut u8)
+    where
+        M: Measure,
+    {
         let data = data as *mut (*mut M, *mut Measurements<'b>);
         let (measure, measurements) = unsafe {
             let data = data.as_mut().unwrap();
@@ -183,7 +195,10 @@ where
     }
 
     /// Bench API callback for the end of instantiation.
-    extern "C" fn instantiation_end(data: *mut u8) {
+    extern "C" fn instantiation_end(data: *mut u8)
+    where
+        M: Measure,
+    {
         let data = data as *mut (*mut M, *mut Measurements<'b>);
         let (measure, measurements) = unsafe {
             let data = data.as_mut().unwrap();
@@ -194,7 +209,10 @@ where
     }
 
     /// Bench API callback for the end of execution.
-    extern "C" fn execution_end(data: *mut u8) {
+    extern "C" fn execution_end(data: *mut u8)
+    where
+        M: Measure,
+    {
         let data = data as *mut (*mut M, *mut Measurements<'b>);
         let (measure, measurements) = unsafe {
             let data = data.as_mut().unwrap();
@@ -225,6 +243,11 @@ impl<'a, 'b, 'c, M> Module<'a, 'b, 'c, M> {
         self.engine
     }
 
+    /// Get this engine's measurements.
+    pub fn measurements(&mut self) -> &mut Measurements<'c> {
+        self.engine.measurements()
+    }
+
     /// Instantiate this module, returning the resulting `Instance`.
     pub fn instantiate(self) -> Instance<'a, 'b, 'c, M> {
         let result = unsafe { (self.engine.bench_api.wasm_bench_instantiate)(self.engine.engine) };

crates/recorder/src/benchmark.rs

@@ -1,20 +1,13 @@
-use crate::bench_api::Engine;
+use crate::bench_api::{Engine, Module};
 use crate::measure::Measure;
 use anyhow::Result;
 use log::info;
-use sightglass_data::Phase;
-
-/// Measure various phases of a Wasm module's lifetime.
-///
-/// Provide paths to files created for logging the Wasm's `stdout` and `stderr`
-/// and (optionally) a file read and piped into the Wasm execution as `stdin`.
-///
-/// Optionally stop after the given `stop_after_phase`, rather than running all
-/// phases.
-pub fn benchmark<'a, 'b, 'c, M: Measure>(
+
+/// Measure all phases of a Wasm module's lifetime: compilation, instantiation,
+/// and execution.
+pub fn all<'a, 'b, 'c, M: Measure>(
     engine: Engine<'a, 'b, 'c, M>,
     wasm_bytes: &[u8],
-    stop_after_phase: Option<Phase>,
 ) -> Result<Engine<'a, 'b, 'c, M>> {
     #[cfg(target_os = "linux")]
     info!("Benchmark scheduled on CPU: {}", unsafe {
@@ -25,20 +18,51 @@ pub fn benchmark<'a, 'b, 'c, M: Measure>(
     let module = engine.compile(wasm_bytes);
     info!("Compiled successfully");
 
-    if stop_after_phase == Some(Phase::Compilation) {
-        return Ok(module.into_engine());
-    }
-
     // Measure the module instantiation.
     let instance = module.instantiate();
     info!("Instantiated successfully");
 
-    if stop_after_phase == Some(Phase::Instantiation) {
-        return Ok(instance.into_module().into_engine());
-    }
-
     let module = instance.execute();
     info!("Executed successfully");
 
     Ok(module.into_engine())
 }
+
+/// Measure just the compilation phase of a Wasm module's lifetime.
+pub fn compilation<'a, 'b, 'c, M: Measure>(
+    engine: Engine<'a, 'b, 'c, M>,
+    wasm_bytes: &[u8],
+) -> Result<Engine<'a, 'b, 'c, M>> {
+    #[cfg(target_os = "linux")]
+    info!("Benchmark scheduled on CPU: {}", unsafe {
+        libc::sched_getcpu()
+    });
+
+    let module = engine.compile(wasm_bytes);
+    info!("Compiled successfully");
+
+    Ok(module.into_engine())
+}
+
+/// Measure just the instantiation phase of a Wasm module's lifetime.
+pub fn instantiation<'a, 'b, 'c, M: Measure>(
+    module: Module<'a, 'b, 'c, M>,
+) -> Result<Module<'a, 'b, 'c, M>> {
+    let instance = module.instantiate();
+    info!("Instantiated successfully");
+
+    Ok(instance.into_module())
+}
+
+/// Measure just the execution phase of a Wasm module's lifetime.
+pub fn execution<'a, 'b, 'c, M: Measure>(
+    module: Module<'a, 'b, 'c, M>,
+) -> Result<Module<'a, 'b, 'c, M>> {
+    let instance = module.instantiate();
+    info!("Instantiated successfully");
+
+    let module = instance.execute();
+    info!("Executed successfully");
+
+    Ok(module)
+}