feat: add native and web benchmark harness (#68)

Co-authored-by: Hiroshi Ogawa <4232207+hi-ogawa@users.noreply.github.com>
Co-authored-by: OpenCode <noreply@opencode.ai>

Author: hi-ogawa

Cargo.lock

@@ -70,6 +70,36 @@ pnpm build-model htdemucs
 
 Run `pnpm build-model --all` to build the standard model and all fine-tuned specialists.
 
+## Benchmark
+
+Compare native ONNX Runtime inference with Chromium's ONNX Runtime WASM backend using a deterministic 30-second workload:
+
+```bash
+mkdir -p data/benchmark
+ffmpeg -f lavfi -i "sine=frequency=440:sample_rate=44100:duration=30" \
+  -filter_complex "[0:a]asplit=2[left][right];[left][right]join=inputs=2:channel_layout=stereo" \
+  -c:a pcm_f32le -y data/benchmark/input-30s.wav
+cargo build --release -p demucs-cli
+pnpm build-wasm
+for threads in 0 1 2 4 8 16; do
+  label=$threads
+  test "$threads" = 0 && label=default
+  for run in 0 1 2 3; do
+    rm -rf "data/benchmark/native-threads-$label-run-$run"
+    target/release/demucs separate \
+      --models data/onnx-lean \
+      --threads "$threads" \
+      --timings-json "data/benchmark/native-threads-$label-run-$run.json" \
+      data/benchmark/input-30s.wav \
+      "data/benchmark/native-threads-$label-run-$run"
+  done
+done
+pnpm -C packages/app benchmark
+pnpm tsx tools/benchmark-summary.ts
+```
+
+See [`docs/benchmark.md`](docs/benchmark.md) for the workload, timing boundaries, and generated results.
+
 ## Web App
 
 Build the WASM binding and start the fully client-side app:

README.md

@@ -13,6 +13,8 @@ clap = { version = "4", features = ["derive"] }
 console = "0.16"
 demucs-core = { path = "../core" }
 indicatif = "0.18"
+serde = { version = "1", features = ["derive"] }
+serde_json = "1"
 # default features with tls-native swapped for tls-rustls, so the binary downloader
 # doesn't need system OpenSSL headers
 ort = { version = "2.0.0-rc.10", default-features = false, features = [

crates/cli/Cargo.toml

@@ -4,6 +4,7 @@ use clap::{Args, Parser, Subcommand};
 use console::style;
 use demucs_core as core;
 use indicatif::{MultiProgress, ProgressBar, ProgressStyle};
+use serde::Serialize;
 use std::path::PathBuf;
 use std::time::{Duration, Instant};
 
@@ -69,6 +70,14 @@ struct SeparateArgs {
     )]
     shifts: u32,
 
+    /// ONNX Runtime intra-op threads; 0 uses the runtime default
+    #[arg(long, default_value_t = 4, value_name = "N", hide_default_value = true)]
+    threads: usize,
+
+    /// Write machine-readable phase timings to this JSON file
+    #[arg(long, value_name = "FILE")]
+    timings_json: Option<PathBuf>,
+
     /// Input WAV file
     #[arg(value_name = "INPUT.WAV")]
     input: String,
@@ -112,9 +121,14 @@ fn separate(args: SeparateArgs) -> Result<()> {
 
     eprintln!("prepared audio in {}", format_duration(prepare_elapsed));
     let mut progress = CliProgress::new();
-    let outputs = ort_driver::run_all(&args.models_dir, &members, wav, opts, |event| {
-        progress.update(event)
-    })?;
+    let outputs = ort_driver::run_all(
+        &args.models_dir,
+        &members,
+        wav,
+        opts,
+        args.threads,
+        |event| progress.update(event),
+    )?;
     progress.finish();
 
     let named: Vec<(String, [Vec<f32>; core::CHANNELS])> = match outputs {
@@ -150,9 +164,50 @@ fn separate(args: SeparateArgs) -> Result<()> {
         format_duration(progress.finalize_elapsed),
         format_duration(write_elapsed),
     );
+    if let Some(path) = args.timings_json {
+        let timings = Timings {
+            prepare_ms: millis(prepare_elapsed),
+            load_ms: millis(progress.load_elapsed),
+            inference_ms: millis(progress.inference_elapsed),
+            chunks: progress.chunks,
+            finalize_ms: millis(progress.finalize_elapsed),
+            write_ms: millis(write_elapsed),
+            total_ms: millis(total_elapsed),
+        };
+        std::fs::write(&path, serde_json::to_vec_pretty(&timings)?)
+            .with_context(|| format!("write {}", path.display()))?;
+    }
     Ok(())
 }
 
+#[derive(Serialize)]
+#[serde(rename_all = "camelCase")]
+struct Timings {
+    prepare_ms: f64,
+    load_ms: f64,
+    inference_ms: f64,
+    chunks: Vec<ChunkTiming>,
+    finalize_ms: f64,
+    write_ms: f64,
+    total_ms: f64,
+}
+
+#[derive(Serialize)]
+#[serde(rename_all = "camelCase")]
+struct ChunkTiming {
+    member: usize,
+    shift: usize,
+    chunk: usize,
+    prepare_input_ms: f64,
+    ort_run_ms: f64,
+    output_copy_ms: f64,
+    process_output_ms: f64,
+}
+
+fn millis(duration: Duration) -> f64 {
+    duration.as_secs_f64() * 1000.0
+}
+
 // Interactive layouts stay model-oriented while the overall row remains stable:
 //   htdemucs.onnx
 //     Load      done | 5.2s
@@ -165,6 +220,7 @@ struct CliProgress {
     phase: Option<ProgressBar>,
     load_elapsed: Duration,
     inference_elapsed: Duration,
+    chunks: Vec<ChunkTiming>,
     finalize_elapsed: Duration,
     loaded: usize,
     eta: Option<Duration>,
@@ -181,6 +237,7 @@ impl CliProgress {
             phase: None,
             load_elapsed: Duration::ZERO,
             inference_elapsed: Duration::ZERO,
+            chunks: Vec::new(),
             finalize_elapsed: Duration::ZERO,
             loaded: 0,
             eta: None,
@@ -257,9 +314,22 @@ impl CliProgress {
                 shift,
                 shifts,
                 member_done,
+                member,
                 elapsed,
+                prepare_elapsed,
+                run_elapsed,
+                process_elapsed,
             } => {
                 self.inference_elapsed += elapsed;
+                self.chunks.push(ChunkTiming {
+                    member,
+                    shift,
+                    chunk: member_done,
+                    prepare_input_ms: millis(prepare_elapsed),
+                    ort_run_ms: millis(run_elapsed),
+                    output_copy_ms: 0.0,
+                    process_output_ms: millis(process_elapsed),
+                });
                 self.overall.set_length(total as u64);
                 self.overall.set_position(done as u64);
                 if let Some(phase) = &self.phase {

crates/cli/src/main.rs

@@ -22,10 +22,14 @@ pub enum Progress<'a> {
         done: usize,
         total: usize,
         members: usize,
+        member: usize,
         shift: usize,
         shifts: usize,
         member_done: usize,
         elapsed: Duration,
+        prepare_elapsed: Duration,
+        run_elapsed: Duration,
+        process_elapsed: Duration,
     },
     MemberFinished {
         chunks: usize,
@@ -57,6 +61,7 @@ pub fn run_all(
     members: &[core::vocab::Member],
     wav: [Vec<f32>; core::CHANNELS],
     opts: core::Options,
+    threads: usize,
     mut on_progress: impl FnMut(Progress<'_>),
 ) -> Result<core::Outputs> {
     let mut separation = core::Separation::new(wav, opts)?;
@@ -90,10 +95,11 @@ pub fn run_all(
         });
         let path = models_dir.join(file);
         let load_started = Instant::now();
-        let mut session = ort::session::Session::builder()
-            .map_err(ort_err)?
-            .with_intra_threads(4)
-            .map_err(ort_err)?
+        let mut builder = ort::session::Session::builder().map_err(ort_err)?;
+        if threads > 0 {
+            builder = builder.with_intra_threads(threads).map_err(ort_err)?;
+        }
+        let mut session = builder
             .commit_from_file(&path)
             .map_err(ort_err)
             .with_context(|| format!("load {}", path.display()))?;
@@ -107,31 +113,41 @@ pub fn run_all(
             let mut chunk_processor = separation.plan.create_chunk_processor(shift);
             for &chunk in &shift.chunks {
                 let chunk_started = Instant::now();
+                let prepare_started = Instant::now();
                 chunk_processor.prepare_input(chunk, &mut input)?;
                 let value = ort::value::TensorRef::from_array_view((
                     [1usize, core::CHANNELS, core::SEGMENT],
                     input.as_slice(),
                 ))
                 .map_err(ort_err)?;
+                let prepare_elapsed = prepare_started.elapsed();
+                let run_started = Instant::now();
                 let run = session
                     .run(ort::inputs!["input" => value])
                     .map_err(ort_err)?;
+                let run_elapsed = run_started.elapsed();
+                let process_started = Instant::now();
                 let (shape, data) = run["output"].try_extract_tensor::<f32>().map_err(ort_err)?;
                 let dims: Vec<usize> = shape.iter().map(|&d| d as usize).collect();
                 if dims != [1, core::NUM_SOURCES, core::CHANNELS, core::SEGMENT] {
                     bail!("unexpected output shape {dims:?}");
                 }
                 chunk_processor.process_output(chunk, data)?;
+                let process_elapsed = process_started.elapsed();
                 done += 1;
                 member_done += 1;
                 on_progress(Progress::Inference {
                     done,
                     total,
                     members: members.len(),
+                    member: member_index + 1,
                     shift: shift_index + 1,
                     shifts: member_plan.shifts.len(),
                     member_done,
                     elapsed: chunk_started.elapsed(),
+                    prepare_elapsed,
+                    run_elapsed,
+                    process_elapsed,
                 });
             }
             shift_merger.add(chunk_processor.finish());