diff --git a/crates/onnx-genai-bench/Cargo.toml b/crates/onnx-genai-bench/Cargo.toml index 0cd506e684..a2557d48cd 100644 --- a/crates/onnx-genai-bench/Cargo.toml +++ b/crates/onnx-genai-bench/Cargo.toml @@ -161,6 +161,10 @@ name = "bench_generic" # arm with `--features mlas`. required-features = ["bench-native"] +[[bin]] +name = "bench_decode_gap" +required-features = ["bench-native"] + [[bin]] name = "cuda_place" required-features = ["bench-native"] diff --git a/crates/onnx-genai-bench/src/bin/bench_decode_gap.rs b/crates/onnx-genai-bench/src/bin/bench_decode_gap.rs new file mode 100644 index 0000000000..66e72401ff --- /dev/null +++ b/crates/onnx-genai-bench/src/bin/bench_decode_gap.rs @@ -0,0 +1,986 @@ +//! Decode-shaped, gap-aware model-level benchmark for the CPU task runtime. +//! +//! `bench_generic` measures a model in a tight loop. That answers "how fast is +//! this graph when it is run back-to-back forever", which is the right question +//! for prefill and the wrong one for decode. Decode issues a parallel region, +//! does a stretch of serial host work, and issues the next one -- and the cost +//! of that pattern is dominated by whether the pool's workers were still +//! spinning when the next fan-out arrived. A tight loop never lets them park, +//! so it reports a best case that production never sees; a short loop never +//! leaves pool construction, so it reports a worst case that production also +//! never sees. +//! +//! This harness runs the model with a configurable gap between iterations, and +//! reports the steady-state distribution *plus* the counters that say which +//! regime it was in: +//! +//! * `vol-ctxt/iter` -- voluntary context switches per iteration. This is the +//! park/wake count. Near zero means the pool spun through the gap; near the +//! worker count means it parked and paid a futex wake on every dispatch. +//! * `cpu/wall` -- CPU-seconds burned per wall-second. A pool that spins +//! through gaps converts idle time into CPU here. +//! * `rss` -- so a scheduler change that trades memory for latency is visible. +//! +//! Steady state is *detected*, not assumed: iterations are recorded from the +//! first one, and the reported window starts where the series stops trending +//! (see [`steady_state_start`]). A run whose series never settles says so +//! rather than quietly reporting its own warm-up. +//! +//! # Examples +//! +//! ```text +//! # A decode-shaped gap, native alone, steady state detected automatically. +//! bench_decode_gap --model m.onnx --gap-us 20 --iters 600 --native-threads 16 +//! +//! # The A/A null control: the same binary twice, to size the noise band +//! # before believing any A/B ratio. +//! bench_decode_gap --model m.onnx --arm null --gap-us 20 --iters 600 +//! +//! # Four concurrent sessions, the oversubscription shape. +//! bench_decode_gap --model m.onnx --sessions 4 --gap-us 50 +//! +//! # Thread ownership census. +//! bench_decode_gap --model m.onnx --census --iters 64 +//! ``` +//! +//! Syscall-level attribution (futex and `sched_yield` counts) is not +//! self-reported, because Linux exposes no per-syscall counter to a process +//! about itself. The harness prints a ready-to-run `strace -c -f` line for the +//! exact invocation instead; voluntary context switches, which *are* +//! self-reported, already capture the park/wake behaviour those syscalls +//! implement. + +use std::path::PathBuf; +use std::time::Instant; + +use anyhow::{Context, Result, bail}; +use clap::Parser; +use onnx_genai_bench::decode_gap::{ + GapDistribution, GapKind, ProcessMetrics, Summary, ThreadInfo, census_by_name, census_delta, + sample_process_metrics, spend_gap, steady_state_start, thread_census, +}; +use onnx_genai_bench::model_io::{ + F16_EPSILON, build_arm, build_inputs, build_ort_inputs, compare_outputs, parse_shape, +}; +use onnx_genai_ort::{Environment, Session, SessionOptions, ep_selection}; +use onnx_runtime_ep_cpu::{dispatch_ledger, task_runtime}; +use onnx_runtime_session::InferenceSession; + +/// The native CPU EP's decode-pool width knob, read once into a `OnceLock` by +/// the EP, so it must be set before the first session is built. +const NATIVE_DECODE_THREADS_ENV: &str = "ONNX_GENAI_CPU_DECODE_THREADS"; + +#[derive(Debug, Parser)] +#[command(about = "Decode-shaped, gap-aware CPU benchmark with park/wake accounting")] +struct Args { + /// ONNX model file. + #[arg(long)] + model: PathBuf, + /// Which runtime to time. `null` runs the native arm twice and reports the + /// second over the first -- the A/A control that sizes the host's noise + /// band. No A/B ratio from this harness means anything until the null + /// control has been read. + #[arg(long, default_value = "native")] + arm: String, + /// Mean gap between iterations, in microseconds. `0` reproduces the + /// tight-loop shape, which is useful only as a reference point. + #[arg(long, default_value_t = 20)] + gap_us: u64, + /// Fractional half-width of the uniform spread around `--gap-us`. A fixed + /// gap can sit permanently just inside or just outside the pool's spin + /// window and report a clean number for a bimodal reality. + #[arg(long, default_value_t = 0.25)] + gap_jitter: f64, + /// How the gap is spent: `busy` (spin, holding the core, like host-side + /// compute), `sleep` (release the core, like blocking on a tokenizer), or + /// `mixed` (alternate). + #[arg(long, default_value = "busy")] + gap_kind: String, + /// Seed for the gap sequence. Printed with the results so a run can be + /// reproduced exactly. + #[arg(long, default_value_t = 0x5EB_A571_A2)] + gap_seed: u64, + /// Total iterations recorded per session, including the warm-up transient. + /// Must be large enough to contain a transient *and* two steady windows; + /// a 32-wide pool needs several hundred. + #[arg(long, default_value_t = 600)] + iters: usize, + /// Force the steady-state window to start here instead of detecting it. + /// `usize::MAX` sentinel is avoided: `--warmup-iters 0` means "detect". + #[arg(long, default_value_t = 0)] + warmup_iters: usize, + /// Samples per window used by steady-state detection. + #[arg(long, default_value_t = 32)] + steady_window: usize, + /// Relative tolerance for steady-state detection. + #[arg(long, default_value_t = 0.05)] + steady_tolerance: f64, + /// Concurrent sessions, each with its own model instance and gap sequence. + /// This is the oversubscription shape: N independent decoders sharing one + /// host, which is what a server does. + #[arg(long, default_value_t = 1)] + sessions: usize, + /// Native CPU decode-pool width. `0` leaves `ONNX_GENAI_CPU_DECODE_THREADS` + /// exactly as inherited, which is what a user gets out of the box. + #[arg(long, default_value_t = 0)] + native_threads: usize, + /// ORT `intra_op_num_threads` for the ORT arm. `0` keeps ORT's default. + #[arg(long, default_value_t = 0)] + ort_intra_threads: i32, + /// Override the first model input shape, for example 1,3,416,416. + #[arg(long)] + input_shape: Option, + /// Skip the ORT parity check. Parity is checked once, before timing, and + /// the ORT session is then dropped so it cannot spin against the timed + /// native arm -- so this flag only matters when ORT cannot load the model. + #[arg(long)] + no_parity: bool, + /// Print a per-thread census (name, CPU, context switches) after timing. + #[arg(long)] + census: bool, + /// Print which threads each lifecycle phase created: native session load, + /// ORT parity session, first inference, and the timed loop. This is how an + /// unnamed thread gets an owner -- a census alone cannot tell an + /// unnamed Rayon worker of ours from one of ORT's, but a delta can. + #[arg(long)] + census_phases: bool, + /// Write every per-iteration sample to this CSV, warm-up included, so the + /// transient can be inspected rather than inferred. Columns: + /// `arm,iteration,wall_ms,in_steady_window`. + #[arg(long)] + dump_csv: Option, + /// Relative tolerance used for Float32 output parity. + #[arg(long, default_value_t = 1e-3)] + rel_tolerance: f32, + /// Absolute tolerance used for Float32 output parity. + #[arg(long, default_value_t = 1e-4)] + abs_tolerance: f32, +} + +/// Prints which threads each lifecycle phase created. +/// +/// Thread `comm` truncates to 15 bytes and defaults to the parent's, so a +/// thread nobody named shows up wearing the *process* name. That makes a flat +/// census actively misleading: the anonymous workers look like they belong to +/// the binary. Bracketing each construction step tells you who really made +/// them. +struct PhaseCensus { + enabled: bool, + previous: Vec, +} + +impl PhaseCensus { + fn new(enabled: bool) -> Self { + Self { + enabled, + previous: if enabled { thread_census() } else { Vec::new() }, + } + } + + fn mark(&mut self, phase: &str) { + if !self.enabled { + return; + } + let now = thread_census(); + let created = census_delta(&self.previous, &now); + let removed = census_delta(&now, &self.previous); + if created.is_empty() && removed.is_empty() { + println!("phase[{phase}]: no thread change ({} live)", now.len()); + } else { + let names = census_by_name(&created) + .into_iter() + .map(|(name, count, _)| format!("{count}x {name}")) + .collect::>() + .join(", "); + println!( + "phase[{phase}]: +{} -{} threads ({} live){}", + created.len(), + removed.len(), + now.len(), + if names.is_empty() { + String::new() + } else { + format!(" -> {names}") + } + ); + } + self.previous = now; + } +} + +/// One point in time: process counters, the native pool's counters, and the +/// instant they were taken at. +#[derive(Clone, Copy)] +struct Snapshot { + iteration: usize, + metrics: ProcessMetrics, + pool: task_runtime::PoolCounters, + at: Instant, +} + +fn snapshot(iteration: usize) -> Snapshot { + Snapshot { + iteration, + metrics: sample_process_metrics(), + pool: task_runtime::testing::counters(), + at: Instant::now(), + } +} + +/// Field-wise difference of two pool counter snapshots. +fn pool_delta( + after: task_runtime::PoolCounters, + before: task_runtime::PoolCounters, +) -> task_runtime::PoolCounters { + task_runtime::PoolCounters { + dispatches: after.dispatches.saturating_sub(before.dispatches), + tasks: after.tasks.saturating_sub(before.tasks), + tasks_by_dispatcher: after + .tasks_by_dispatcher + .saturating_sub(before.tasks_by_dispatcher), + slot_exhausted: after.slot_exhausted.saturating_sub(before.slot_exhausted), + parks: after.parks.saturating_sub(before.parks), + spin_hits: after.spin_hits.saturating_sub(before.spin_hits), + panics: after.panics.saturating_sub(before.panics), + straggler_waits: after.straggler_waits.saturating_sub(before.straggler_waits), + straggler_yields: after + .straggler_yields + .saturating_sub(before.straggler_yields), + } +} + +/// Which kernel families ran on which backend, and with what parallel degree. +/// +/// The pool counters answer "how hard did `task_runtime` work"; they cannot +/// answer "did this model use `task_runtime` at all". Those two produce the +/// same output -- a row of zeroes -- and only one of them means the harness is +/// broken. Attribution is what separates them, so it is collected rather than +/// assumed: an fp32 `MatMul` fans out on rayon, not on the native task pool, so +/// a model built from fp32 `MatMul`s reports zero dispatches while saturating +/// eight cores, and that reading is correct. +/// +/// Recording is on only for the single attribution inference, never during the +/// timed window, because `record_with` builds an `Observation` per dispatch. +fn attribute_routes(run: impl FnOnce() -> Result<()>) -> Result> { + dispatch_ledger::reset(); + dispatch_ledger::enable(); + let outcome = run(); + dispatch_ledger::disable(); + outcome?; + + let mut rows: Vec = Vec::new(); + for observation in dispatch_ledger::snapshot() { + let key = ( + format!("{:?}", observation.family), + format!("{:?}", observation.backend), + observation.dtype, + ); + match rows.iter_mut().find(|row| row.key() == key) { + Some(row) => { + row.calls += 1; + row.max_threads = row.max_threads.max(observation.threads); + } + None => rows.push(RouteRow { + family: key.0, + backend: key.1, + dtype: observation.dtype, + calls: 1, + max_threads: observation.threads, + }), + } + } + rows.sort_by(|a, b| b.calls.cmp(&a.calls).then_with(|| a.family.cmp(&b.family))); + Ok(rows) +} + +/// One `(family, backend, dtype)` route observed during attribution. +struct RouteRow { + family: String, + backend: String, + dtype: &'static str, + calls: usize, + max_threads: usize, +} + +impl RouteRow { + fn key(&self) -> (String, String, &'static str) { + (self.family.clone(), self.backend.clone(), self.dtype) + } +} + +/// Whether any observed route is one the native task pool would drive. +/// +/// Used only to phrase the zero-dispatch case: "this model does not use the +/// pool" is a different statement from "the pool did nothing this run", and the +/// report should not silently pick one. +fn any_route_observed(rows: &[RouteRow]) -> bool { + !rows.is_empty() +} + +/// What one timed arm produced. +struct ArmResult { + /// Every iteration, in order, including the warm-up transient. + samples_ms: Vec, + /// Where steady state began, and how it was decided. + steady_start: Option, + detected: bool, + /// Counters accumulated across the steady window only. + steady_metrics: ProcessMetrics, + /// Wall-clock seconds spanned by the steady window. + steady_wall_s: f64, + /// Counters accumulated across the whole run, warm-up included. + total_metrics: ProcessMetrics, + total_wall_s: f64, + /// Native pool counters accumulated across the steady window. + steady_pool: task_runtime::PoolCounters, + /// Iterations the counter deltas above actually span. + /// + /// Counters are only sampled every `--steady-window` iterations, so the + /// nearest snapshot at or before the steady boundary is generally *earlier* + /// than the boundary. Dividing a counter delta by the sample-window length + /// would then overstate every per-iteration figure, so per-iteration + /// counters divide by this instead. It is reported, because a counter + /// window wider than the sample window is exactly the kind of thing that + /// silently turns a per-iteration number into a slightly wrong one. + counter_iters: usize, + /// Whether the counter window had to fall back to iteration zero, i.e. + /// whether these "steady" counters in fact include the warm-up transient. + counters_include_warmup: bool, + /// Routes observed during the attribution inference, if one was taken. + routes: Vec, +} + +impl ArmResult { + fn steady_samples(&self) -> &[f64] { + match self.steady_start { + Some(start) => &self.samples_ms[start..], + None => &self.samples_ms, + } + } + + fn summary(&self) -> Summary { + Summary::from(self.steady_samples()) + } + + /// Voluntary context switches per iteration over the counter window: the + /// park/wake count per dispatch. + fn parks_per_iter(&self) -> f64 { + self.steady_metrics.voluntary_ctxt_switches as f64 / self.counter_iters.max(1) as f64 + } + + /// CPU-seconds burned per wall-second over the steady window. A spinning + /// pool converts the gap into CPU and shows up here. + fn cpu_per_wall(&self) -> f64 { + if self.steady_wall_s <= 0.0 { + return 0.0; + } + (self.steady_metrics.cpu_us() as f64 / 1e6) / self.steady_wall_s + } +} + +/// Runs `iterations` of `body`, spending a gap from `gaps` before each one. +/// +/// Metrics are sampled twice: once at the start, and once at the steady-state +/// boundary. Because the boundary is not known until the series is complete, +/// the loop records a metrics snapshot every `snapshot_every` iterations and +/// the caller picks the one nearest the detected boundary. That keeps the +/// per-iteration cost of measurement off the timed path -- reading +/// `/proc/self/status` costs tens of microseconds, which is the same order as +/// the gaps being measured. +fn timed_loop( + iterations: usize, + gaps: &mut GapDistribution, + snapshot_every: usize, + mut body: F, +) -> Result<(Vec, Vec)> +where + F: FnMut() -> Result<()>, +{ + let mut samples = Vec::with_capacity(iterations); + let mut snapshots = vec![snapshot(0)]; + for index in 0..iterations { + let kind = gaps.next_kind(); + spend_gap(gaps.next_gap(), kind); + let started = Instant::now(); + body()?; + samples.push(started.elapsed().as_secs_f64() * 1_000.0); + if snapshot_every > 0 && (index + 1).is_multiple_of(snapshot_every) { + snapshots.push(snapshot(index + 1)); + } + } + snapshots.push(snapshot(iterations)); + Ok((samples, snapshots)) +} + +/// Turns a sample series and its metric snapshots into an [`ArmResult`]. +fn assemble(samples: Vec, snapshots: Vec, args: &Args) -> ArmResult { + let detected = args.warmup_iters == 0; + let steady_start = if detected { + steady_state_start(&samples, args.steady_window, args.steady_tolerance) + } else { + (args.warmup_iters < samples.len()).then_some(args.warmup_iters) + }; + let first = snapshots.first().expect("a snapshot at iteration zero"); + let last = snapshots.last().expect("a snapshot at the final iteration"); + // The latest snapshot at or before the steady boundary. Counters are only + // sampled every `--steady-window` iterations, so this is generally + // *earlier* than the boundary and the counter window is correspondingly + // wider than the sample window -- and when steady state is reached before + // the first interior snapshot there is no qualifying snapshot at all, so + // the window falls back to iteration zero and the counters do include the + // warm-up transient. Neither case is avoidable without sampling counters + // every iteration (which would perturb what is being measured), so both are + // recorded and reported instead of being asserted away. + let boundary = steady_start + .and_then(|start| { + snapshots + .iter() + .rev() + .find(|point| point.iteration <= start && point.iteration > 0) + }) + .unwrap_or(first); + let counter_iters = last.iteration.saturating_sub(boundary.iteration); + let counters_include_warmup = boundary.iteration == 0 && steady_start.unwrap_or(0) > 0; + ArmResult { + samples_ms: samples, + steady_start, + detected, + steady_metrics: last.metrics.since(&boundary.metrics), + steady_wall_s: last.at.duration_since(boundary.at).as_secs_f64(), + total_metrics: last.metrics.since(&first.metrics), + total_wall_s: last.at.duration_since(first.at).as_secs_f64(), + steady_pool: pool_delta(last.pool, boundary.pool), + counter_iters, + counters_include_warmup, + routes: Vec::new(), + } +} + +fn run_native_arm(args: &Args, label: &str) -> Result { + let mut phases = PhaseCensus::new(args.census_phases); + phases.mark("process-start"); + let mut sessions = (0..args.sessions.max(1)) + .map(|_| { + InferenceSession::load(&args.model) + .with_context(|| format!("load native session from {}", args.model.display())) + }) + .collect::>>()?; + phases.mark("native-session-load"); + + let input_shape = args + .input_shape + .as_deref() + .map(parse_shape) + .transpose() + .map_err(anyhow::Error::msg)?; + + // Parity is judged once, against ORT, *before* the timed loop -- and the + // ORT session is dropped immediately afterwards. A co-resident ORT session + // spin-waits long after its last op and depresses a native arm measured + // beside it by several times; building it, using it and dropping it makes + // the timed window solo by construction rather than by remembering a flag. + if !args.no_parity { + let environment = Environment::new("bench-decode-gap")?; + let options = + SessionOptions::with_execution_provider(ep_selection("cpu")).with_intra_op_threads(1); + let ort_session = Session::new(&environment, &args.model, options) + .with_context(|| format!("load ORT CPU session from {}", args.model.display()))?; + let inputs = build_inputs(&sessions[0], &ort_session, input_shape.as_deref())?; + let native_inputs = inputs + .iter() + .map(|input| (input.name.as_str(), &input.native)) + .collect::>(); + let ort_inputs = inputs + .iter() + .map(|input| (input.name.as_str(), &input.ort)) + .collect::>(); + phases.mark("ort-parity-session-built"); + let native_reference = sessions[0] + .run(&native_inputs) + .context("native parity run")?; + phases.mark("first-native-inference"); + let ort_reference = ort_session.run(&ort_inputs).context("ORT parity run")?; + let diffs = compare_outputs( + &native_reference, + &ort_reference, + args.abs_tolerance, + args.rel_tolerance, + F16_EPSILON, + 4.0 * F16_EPSILON, + )?; + let failures = diffs.iter().filter(|diff| !diff.pass).count(); + println!( + "parity[{label}]: {} outputs checked, {failures} failed, max_rel={:.3e}", + diffs.len(), + diffs.iter().map(|diff| diff.max_rel).fold(0.0f32, f32::max) + ); + if failures > 0 { + bail!( + "native/ORT parity failed on {failures} output(s); timing a wrong kernel is not a measurement" + ); + } + } + // The ORT session is out of scope here. Anything it created that is still + // alive shows up as a *negative* delta failing to appear. + phases.mark("ort-session-dropped"); + + // Inputs for the timed loop, built without ORT in the picture. + let owned_inputs = build_timed_inputs(&sessions[0], &args.model, input_shape.as_deref())?; + phases.mark("inputs-built"); + + if args.sessions <= 1 { + let mut session = sessions.pop().expect("at least one session"); + let refs = owned_inputs + .iter() + .map(|(name, tensor)| (name.as_str(), tensor)) + .collect::>(); + let mut gaps = GapDistribution::new( + args.gap_us, + args.gap_jitter, + GapKind::parse(&args.gap_kind).map_err(anyhow::Error::msg)?, + args.gap_seed, + ); + let routes = attribute_routes(|| { + std::hint::black_box(session.run(&refs).context("native warm run")?); + Ok(()) + })?; + phases.mark("timed-session-warm"); + let (samples, snapshots) = + timed_loop(args.iters, &mut gaps, args.steady_window.max(1), || { + std::hint::black_box(session.run(&refs).context("native measured run")?); + Ok(()) + })?; + phases.mark("timed-loop"); + let mut result = assemble(samples, snapshots, args); + result.routes = routes; + return Ok(result); + } + + // Attribution runs on one session before the concurrent region, not inside + // it: the ledger takes a mutex per observation, so recording across N + // barrier-synchronised threads would serialise the very contention this arm + // exists to measure. Routing does not depend on how many sessions are live, + // so one session's routes describe them all -- and without this the + // concurrent arm reports zero routes, which the report cannot distinguish + // from a ledger that saw nothing and would wrongly call a valid run a dead + // instrument. + let routes = { + let refs = owned_inputs + .iter() + .map(|(name, tensor)| (name.as_str(), tensor)) + .collect::>(); + let session = sessions.first_mut().expect("at least one session"); + attribute_routes(|| { + std::hint::black_box(session.run(&refs).context("native attribution run")?); + Ok(()) + })? + }; + let mut result = run_concurrent_native(args, sessions, owned_inputs)?; + result.routes = routes; + phases.mark("timed-loop"); + Ok(result) +} + +/// The concurrent-session arm. +/// +/// Each session gets its own thread, its own model instance and its own gap +/// sequence (seeded per session, so the sessions do not march in lockstep and +/// manufacture an artificial thundering herd). Process counters are sampled on +/// the parent thread around the whole barrier-synchronised region, because the +/// interesting quantity is what the *process* costs when N decoders share it. +fn run_concurrent_native( + args: &Args, + sessions: Vec, + owned_inputs: Vec<(String, onnx_runtime_session::Tensor)>, +) -> Result { + use std::sync::{Arc, Barrier}; + + let gap_kind = GapKind::parse(&args.gap_kind).map_err(anyhow::Error::msg)?; + let barrier = Arc::new(Barrier::new(sessions.len() + 1)); + let inputs = Arc::new(owned_inputs); + let iters = args.iters; + let (gap_us, jitter, seed) = (args.gap_us, args.gap_jitter, args.gap_seed); + + let started_all = Instant::now(); + let before = snapshot(0); + let mut handles = Vec::new(); + for (index, mut session) in sessions.into_iter().enumerate() { + let barrier = Arc::clone(&barrier); + let inputs = Arc::clone(&inputs); + handles.push( + std::thread::Builder::new() + .name(format!("gapbench-{index}")) + .spawn(move || -> Result> { + let refs = inputs + .iter() + .map(|(name, tensor)| (name.as_str(), tensor)) + .collect::>(); + let mut gaps = GapDistribution::new( + gap_us, + jitter, + gap_kind, + seed.wrapping_add(index as u64 * 0x9E37_79B9), + ); + barrier.wait(); + let mut samples = Vec::with_capacity(iters); + for _ in 0..iters { + let kind = gaps.next_kind(); + spend_gap(gaps.next_gap(), kind); + let started = Instant::now(); + std::hint::black_box(session.run(&refs).context("native measured run")?); + samples.push(started.elapsed().as_secs_f64() * 1_000.0); + } + Ok(samples) + }) + .context("spawn concurrent session thread")?, + ); + } + barrier.wait(); + + let mut per_session = Vec::new(); + for handle in handles { + per_session.push( + handle + .join() + .map_err(|_| anyhow::anyhow!("concurrent session thread panicked"))??, + ); + } + let after = snapshot(0); + let elapsed = started_all.elapsed().as_secs_f64(); + + for (index, samples) in per_session.iter().enumerate() { + let summary = Summary::from(samples); + println!( + " session[{index}]: p50={:.3} ms p90={:.3} ms n={}", + summary.p50, summary.p90, summary.count + ); + } + + // Pool the sessions' samples: the question a server asks is what a request + // costs, not what a particular worker's stream of requests costs. + let samples = per_session.concat(); + let snapshots = vec![ + Snapshot { + iteration: 0, + metrics: before.metrics, + pool: before.pool, + at: started_all, + }, + Snapshot { + iteration: samples.len(), + metrics: after.metrics, + pool: after.pool, + at: Instant::now(), + }, + ]; + let mut result = assemble(samples, snapshots, args); + // Steady-state detection is meaningless across a concatenation of + // independent series, so the concurrent arm reports the whole window and + // says so rather than pretending it detected a boundary. + result.steady_start = None; + result.detected = false; + result.steady_metrics = after.metrics.since(&before.metrics); + result.steady_pool = pool_delta(after.pool, before.pool); + result.steady_wall_s = elapsed; + Ok(result) +} + +/// Synthetic inputs for the native session alone. +/// +/// Reuses the shared generators so the tensors the timed loop feeds are +/// byte-identical to the ones the parity check judged. The throwaway ORT +/// session exists only to supply declared dtypes to the shared builder; it is +/// dropped before it returns, so nothing of ORT's is alive during timing. +fn build_timed_inputs( + session: &InferenceSession, + model: &std::path::Path, + override_shape: Option<&[usize]>, +) -> Result> { + let environment = Environment::new("bench-decode-gap-inputs")?; + let options = + SessionOptions::with_execution_provider(ep_selection("cpu")).with_intra_op_threads(1); + let probe = Session::new(&environment, model, options) + .with_context(|| format!("load ORT CPU session from {}", model.display()))?; + let pairs = build_inputs(session, &probe, override_shape)?; + Ok(pairs + .into_iter() + .map(|pair| (pair.name, pair.native)) + .collect()) +} + +fn run_ort_arm(args: &Args) -> Result { + let environment = Environment::new("bench-decode-gap")?; + let intra = if args.ort_intra_threads > 0 { + args.ort_intra_threads + } else { + 0 + }; + let options = + SessionOptions::with_execution_provider(ep_selection("cpu")).with_intra_op_threads(intra); + let session = Session::new(&environment, &args.model, options) + .with_context(|| format!("load ORT CPU session from {}", args.model.display()))?; + let input_shape = args + .input_shape + .as_deref() + .map(parse_shape) + .transpose() + .map_err(anyhow::Error::msg)?; + let inputs = build_ort_inputs(&session, input_shape.as_deref())?; + let refs = inputs + .iter() + .map(|(name, value)| (name.as_str(), value)) + .collect::>(); + let mut gaps = GapDistribution::new( + args.gap_us, + args.gap_jitter, + GapKind::parse(&args.gap_kind).map_err(anyhow::Error::msg)?, + args.gap_seed, + ); + let (samples, snapshots) = + timed_loop(args.iters, &mut gaps, args.steady_window.max(1), || { + std::hint::black_box(session.run(&refs).context("ORT measured run")?); + Ok(()) + })?; + Ok(assemble(samples, snapshots, args)) +} + +fn report(label: &str, result: &ArmResult, args: &Args) { + let summary = result.summary(); + let window = match (result.steady_start, result.detected) { + (Some(start), true) => format!("detected at iter {start}"), + (Some(start), false) => format!("forced at iter {start}"), + (None, true) => "NOT REACHED (series never settled)".to_string(), + (None, false) => "whole run".to_string(), + }; + println!("\n== {label} =="); + println!( + " steady window: {window} ({} of {} iters)", + summary.count, + result.samples_ms.len() + ); + println!( + " wall/iter: p50={:.4} ms p90={:.4} ms p99={:.4} ms min={:.4} ms spread={:.2}x", + summary.p50, + summary.p90, + summary.p99, + summary.min, + summary.spread() + ); + println!( + " cpu: {:.3} s over steady window ({:.2} cpu-s per wall-s)", + result.steady_metrics.cpu_us() as f64 / 1e6, + result.cpu_per_wall() + ); + println!( + " ctxt switches: {:.1} vol/iter {:.1} invol/iter ({} vol total)", + result.parks_per_iter(), + result.steady_metrics.involuntary_ctxt_switches as f64 / result.counter_iters.max(1) as f64, + result.steady_metrics.voluntary_ctxt_switches + ); + println!( + " rss: {} kB now, {} kB peak, {} threads", + result.steady_metrics.rss_kb, + result.steady_metrics.peak_rss_kb, + result.steady_metrics.threads + ); + let iterations = result.counter_iters.max(1) as f64; + println!( + " native pool: {:.2} dispatches/iter {:.2} parks/iter {:.2} spin-hits/iter \ + {} slot-exhausted", + result.steady_pool.dispatches as f64 / iterations, + result.steady_pool.parks as f64 / iterations, + result.steady_pool.spin_hits as f64 / iterations, + result.steady_pool.slot_exhausted + ); + let sample_iters = result.steady_samples().len(); + if result.counter_iters == sample_iters { + println!( + " counter window: {} iters, same span as the samples above", + result.counter_iters + ); + } else { + println!( + " counter window: {} iters, WIDER than the {} the samples above cover -- \ + counters are only sampled every --steady-window iters, so per-iter figures \ + here are divided by {}, not by {}", + result.counter_iters, sample_iters, result.counter_iters, sample_iters + ); + } + if result.counters_include_warmup { + println!( + " WARNING: steady state was reached before the first interior counter \ + snapshot, so every counter and cpu figure above spans the whole run and \ + INCLUDES the warm-up transient. Lower --steady-window to separate them." + ); + } + // A zero-dispatch row has two very different causes and the reader cannot + // tell them apart from the row alone, so say which one it was. + if result.steady_pool.dispatches == 0 { + if any_route_observed(&result.routes) { + println!( + " ATTRIBUTION: this model never dispatched to task_runtime -- the pool \ + row above is a true zero, not a dead counter. Routes below say what ran \ + instead; pool numbers here describe a pool this model does not use." + ); + } else { + println!( + " ATTRIBUTION: zero dispatches AND zero routes recorded. The ledger saw \ + nothing, so this is a dead instrument, not a measurement -- do not quote \ + the pool row." + ); + } + } + for route in &result.routes { + println!( + " route: {} -> {} ({}, {} calls, up to {} threads)", + route.family, route.backend, route.dtype, route.calls, route.max_threads + ); + } + if dispatch_ledger::dropped() > 0 { + println!( + " route: WARNING {} observations dropped (ledger full); the route list \ + is truncated and call counts are lower bounds", + dispatch_ledger::dropped() + ); + } + println!( + " whole run: {:.3} cpu-s over {:.3} wall-s (warm-up included)", + result.total_metrics.cpu_us() as f64 / 1e6, + result.total_wall_s + ); + if result.steady_start.is_none() && result.detected { + println!( + " WARNING: no steady state in {} iterations at gap {}us. Every number \ + above describes the warm-up transient. Raise --iters.", + result.samples_ms.len(), + args.gap_us + ); + } +} + +fn main() -> Result<()> { + let args = Args::parse(); + if args.iters == 0 { + bail!("--iters must be greater than zero"); + } + if args.sessions == 0 { + bail!("--sessions must be greater than zero"); + } + GapKind::parse(&args.gap_kind).map_err(anyhow::Error::msg)?; + + if args.native_threads > 0 { + // SAFETY: single-threaded startup, before any session, thread pool or + // other reader of the process environment exists. + unsafe { std::env::set_var(NATIVE_DECODE_THREADS_ENV, args.native_threads.to_string()) }; + } + let native_threads = std::env::var(NATIVE_DECODE_THREADS_ENV) + .ok() + .filter(|value| !value.trim().is_empty()) + .unwrap_or_else(|| "default".to_string()); + + println!("model: {}", args.model.display()); + println!( + "arm: {} (build={}, decode_threads={native_threads}, sessions={})", + args.arm, + build_arm(), + args.sessions + ); + println!( + "gap: {}us +/-{:.0}% {} (seed {})", + args.gap_us, + args.gap_jitter * 100.0, + args.gap_kind, + args.gap_seed + ); + + match args.arm.as_str() { + "native" => { + let result = run_native_arm(&args, "native")?; + report("native", &result, &args); + emit_result_line("native", &result); + if let Some(path) = &args.dump_csv { + dump_csv(path, "native", &result)?; + } + } + "ort" => { + let result = run_ort_arm(&args)?; + report("ort", &result, &args); + emit_result_line("ort", &result); + if let Some(path) = &args.dump_csv { + dump_csv(path, "ort", &result)?; + } + } + "null" => { + let first = run_native_arm(&args, "null-a")?; + report("null-a", &first, &args); + let second = run_native_arm(&args, "null-b")?; + report("null-b", &second, &args); + if let Some(path) = &args.dump_csv { + dump_csv(path, "null-a", &first)?; + dump_csv(path, "null-b", &second)?; + } + let ratio = second.summary().p50 / first.summary().p50; + println!( + "\nresult: null_b_over_a={ratio:.3}x (a={:.4} ms, b={:.4} ms) \ + -- any A/B ratio inside this band is noise", + first.summary().p50, + second.summary().p50 + ); + } + other => bail!("unknown --arm '{other}'; expected native, ort or null"), + } + + if args.census { + let threads = thread_census(); + println!("\n== thread census ({} threads) ==", threads.len()); + println!("{:>28} {:>6} {:>12}", "name", "count", "cpu_ms"); + for (name, count, cpu_us) in census_by_name(&threads) { + println!("{name:>28} {count:>6} {:>12.1}", cpu_us as f64 / 1000.0); + } + } + + println!( + "\nfor futex/sched_yield attribution, re-run under:\n strace -c -f -e trace=futex,sched_yield " + ); + Ok(()) +} + +/// Appends this arm's whole sample series to `path`, creating it with a +/// header if it does not exist. +fn dump_csv(path: &std::path::Path, arm: &str, result: &ArmResult) -> Result<()> { + use std::io::Write; + let fresh = !path.exists(); + let mut file = std::fs::OpenOptions::new() + .create(true) + .append(true) + .open(path) + .with_context(|| format!("open {} for append", path.display()))?; + if fresh { + writeln!(file, "arm,iteration,wall_ms,in_steady_window")?; + } + let steady_start = result.steady_start.unwrap_or(usize::MAX); + for (index, sample) in result.samples_ms.iter().enumerate() { + writeln!( + file, + "{arm},{index},{sample:.6},{}", + u8::from(index >= steady_start) + )?; + } + Ok(()) +} + +fn emit_result_line(arm: &str, result: &ArmResult) { + let summary = result.summary(); + println!( + "result: arm={arm} p50={:.4} ms p90={:.4} ms cpu_per_wall={:.2} vol_ctxt_per_iter={:.2} \ + parks_per_iter={:.2} spin_hits_per_iter={:.2} steady_iters={} rss_kb={}", + summary.p50, + summary.p90, + result.cpu_per_wall(), + result.parks_per_iter(), + result.steady_pool.parks as f64 / summary.count.max(1) as f64, + result.steady_pool.spin_hits as f64 / summary.count.max(1) as f64, + summary.count, + result.steady_metrics.rss_kb + ); +} diff --git a/crates/onnx-genai-bench/src/decode_gap.rs b/crates/onnx-genai-bench/src/decode_gap.rs new file mode 100644 index 0000000000..3f41751f90 --- /dev/null +++ b/crates/onnx-genai-bench/src/decode_gap.rs @@ -0,0 +1,908 @@ +//! Measurement primitives for a *decode-shaped* model-level benchmark. +//! +//! # Why this exists +//! +//! Every model-level number this campaign published came from one of two +//! harness shapes, and both of them are dishonest about decode in opposite +//! directions: +//! +//! * **Few runs in a tight loop** (the `bench_generic` default of 7-10) sits +//! entirely inside the thread-pool warm-up transient. A 32-wide decode pool +//! needs roughly 60 inferences before its marginal cost stops falling, so a +//! 7-run measurement reports pool construction and calls it inference. +//! * **Many runs in a tight loop** (400+) leaves the transient but replaces it +//! with the opposite bias: with no gap between iterations the workers never +//! park, so every dispatch hits an already-spinning pool. That is the one +//! regime real decode never occupies. +//! +//! Real decode is neither. It has microsecond-to-millisecond serial stretches +//! between parallel regions -- sampling, KV bookkeeping, the host side of the +//! next token -- and the interesting scheduler behaviour is precisely what +//! happens across those gaps: whether a worker is still spinning when the next +//! fan-out arrives, or has parked and must be woken through the kernel. +//! +//! This module supplies the pieces to measure that shape at the model level: +//! a configurable gap distribution, warm-up that is defined by observed +//! steady state rather than by a guessed constant, and the process counters +//! that distinguish "spun" from "parked". +//! +//! The equivalent micro-benchmark already exists in the CPU EP's +//! `task_runtime_latency` integration test, which sweeps `GAPS_US` against a +//! bare fan-out. That test answers "what does a dispatch cost after a gap of +//! N microseconds"; this module answers "what does a *model* cost when it is +//! run the way decode runs it", and the two are cross-validated against each +//! other in [`crate::decode_gap::tests`] and in the harness's `--validate` +//! mode. + +use std::time::{Duration, Instant}; + +/// How the harness spends the gap between two model iterations. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum GapKind { + /// Spin without yielding, the way a serial stretch of host-side decode + /// work does. The pool's workers keep their cores warm and are free to + /// keep spinning, so this measures the spin-window hit path. + Busy, + /// Sleep, releasing the core. This is the pessimistic shape: the OS is + /// free to migrate or deschedule the workers, so the next dispatch is + /// likely to pay a real wake. + Sleep, + /// Alternate `Busy` and `Sleep` per iteration. Real decode is a mixture -- + /// some inter-token gaps are pure compute, others block on a tokenizer, + /// a sampler allocation or a detokenizer write -- and a harness that only + /// ever does one of the two will tune the spin window to the wrong shape. + Mixed, +} + +impl GapKind { + /// Parses the `--gap-kind` flag. + pub fn parse(value: &str) -> Result { + match value { + "busy" => Ok(Self::Busy), + "sleep" => Ok(Self::Sleep), + "mixed" => Ok(Self::Mixed), + other => Err(format!( + "unknown gap kind '{other}'; expected busy, sleep or mixed" + )), + } + } +} + +/// A deterministic inter-iteration gap generator. +/// +/// Jitter matters. A fixed gap can sit permanently just inside or just outside +/// the spin window and produce a bimodal result that looks like a clean number; +/// spreading the gap across the window boundary reports the mixture that real +/// decode actually pays. The generator is a seeded xorshift rather than a real +/// RNG so a run is reproducible from its printed seed. +#[derive(Debug, Clone)] +pub struct GapDistribution { + mean_us: u64, + jitter: f64, + kind: GapKind, + state: u64, + issued: u64, +} + +impl GapDistribution { + /// `jitter` is the fractional half-width of a uniform spread around + /// `mean_us`: `0.0` is a fixed gap, `0.5` spreads over `[0.5x, 1.5x]`. + /// It is clamped to `[0, 1]` so the gap can never go negative. + pub fn new(mean_us: u64, jitter: f64, kind: GapKind, seed: u64) -> Self { + Self { + mean_us, + jitter: jitter.clamp(0.0, 1.0), + kind, + // A zero seed is a fixed point of xorshift, which would emit an + // endless run of zeros and silently turn jitter off. + state: if seed == 0 { + 0x9E37_79B9_7F4A_7C15 + } else { + seed + }, + issued: 0, + } + } + + fn next_u64(&mut self) -> u64 { + let mut x = self.state; + x ^= x << 13; + x ^= x >> 7; + x ^= x << 17; + self.state = x; + x + } + + /// The next gap length. Deterministic given the seed. + pub fn next_gap(&mut self) -> Duration { + let mean = self.mean_us; + if mean == 0 { + self.issued = self.issued.wrapping_add(1); + return Duration::ZERO; + } + let micros = if self.jitter == 0.0 { + mean + } else { + // Uniform in [mean*(1-jitter), mean*(1+jitter)], computed in + // floating point and rounded, so small means still jitter. + let unit = (self.next_u64() >> 11) as f64 / (1u64 << 53) as f64; + let low = mean as f64 * (1.0 - self.jitter); + let high = mean as f64 * (1.0 + self.jitter); + (low + unit * (high - low)).round().max(0.0) as u64 + }; + self.issued = self.issued.wrapping_add(1); + Duration::from_micros(micros) + } + + /// Which of [`GapKind::Busy`] / [`GapKind::Sleep`] the next gap uses. + /// [`GapKind::Mixed`] alternates on the issue counter, so a run of N + /// iterations gets as close to an even split as N allows. + pub fn next_kind(&self) -> GapKind { + match self.kind { + GapKind::Mixed => { + if self.issued.is_multiple_of(2) { + GapKind::Busy + } else { + GapKind::Sleep + } + } + other => other, + } + } +} + +/// Spends `gap` the way `kind` says to. +/// +/// The busy arm deliberately does not call `yield_now`: the point is to hold +/// the core the way host-side decode work would, so the pool's spin window +/// sees a realistic neighbour. +pub fn spend_gap(gap: Duration, kind: GapKind) { + if gap.is_zero() { + return; + } + match kind { + GapKind::Busy | GapKind::Mixed => { + let until = Instant::now() + gap; + while Instant::now() < until { + std::hint::spin_loop(); + } + } + GapKind::Sleep => std::thread::sleep(gap), + } +} + +/// Process-wide counters sampled from `/proc/self`. +/// +/// `voluntary_ctxt_switches` is the load-bearing one: a worker that parks and +/// is later woken through a futex records exactly one voluntary switch, so the +/// delta across a measured window is a direct count of park/wake round trips. +/// A pool that spins through every gap moves this number by almost nothing; +/// one that parks on every gap moves it by roughly `iterations x workers`. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct ProcessMetrics { + pub user_us: u64, + pub sys_us: u64, + pub rss_kb: u64, + pub peak_rss_kb: u64, + pub threads: u64, + pub voluntary_ctxt_switches: u64, + pub involuntary_ctxt_switches: u64, +} + +impl ProcessMetrics { + /// Field-wise `self - earlier`, saturating so a counter that wrapped or + /// was re-read out of order reports zero rather than a huge bogus delta. + /// RSS is a level, not a counter, so it is carried through rather than + /// differenced. + pub fn since(&self, earlier: &Self) -> Self { + Self { + user_us: self.user_us.saturating_sub(earlier.user_us), + sys_us: self.sys_us.saturating_sub(earlier.sys_us), + rss_kb: self.rss_kb, + peak_rss_kb: self.peak_rss_kb, + threads: self.threads, + voluntary_ctxt_switches: self + .voluntary_ctxt_switches + .saturating_sub(earlier.voluntary_ctxt_switches), + involuntary_ctxt_switches: self + .involuntary_ctxt_switches + .saturating_sub(earlier.involuntary_ctxt_switches), + } + } + + pub fn cpu_us(&self) -> u64 { + self.user_us + self.sys_us + } +} + +/// Parses the subset of `/proc/self/status` this harness reports. +/// +/// Split out from the read so it can be tested against a captured sample -- +/// the format is stable but the field set differs across kernels, and a silent +/// parse failure here would show up as a suspiciously flat counter rather than +/// as an error. +pub fn parse_status(text: &str) -> ProcessMetrics { + let mut metrics = ProcessMetrics::default(); + for line in text.lines() { + let Some((key, value)) = line.split_once(':') else { + continue; + }; + let number = value + .split_whitespace() + .next() + .and_then(|token| token.parse::().ok()) + .unwrap_or(0); + match key { + "VmRSS" => metrics.rss_kb = number, + "VmHWM" => metrics.peak_rss_kb = number, + "Threads" => metrics.threads = number, + "voluntary_ctxt_switches" => metrics.voluntary_ctxt_switches = number, + "nonvoluntary_ctxt_switches" => metrics.involuntary_ctxt_switches = number, + _ => {} + } + } + metrics +} + +/// Parses `utime`/`stime` (fields 14 and 15) out of a `/proc/.../stat` line. +/// +/// The comm field is parenthesized and may itself contain spaces and +/// parentheses, so the parse has to start after the *last* `)` rather than +/// splitting the whole line on whitespace. +pub fn parse_stat_cpu_ticks(line: &str) -> Option<(u64, u64)> { + let rest = &line[line.rfind(')')? + 1..]; + let fields = rest.split_whitespace().collect::>(); + // `rest` starts at field 3 (state), so utime/stime are indices 11 and 12. + let utime = fields.get(11)?.parse().ok()?; + let stime = fields.get(12)?.parse().ok()?; + Some((utime, stime)) +} + +fn clock_ticks_per_second() -> u64 { + // `sysconf(_SC_CLK_TCK)` is 100 on every Linux target this runs on. Rather + // than link libc for one constant, use the value and let the CPU numbers + // be reported in the same units the kernel already uses. + 100 +} + +/// Samples this process's counters now. Returns defaults off Linux or if +/// `/proc` is not mounted, so the harness degrades to timing-only rather than +/// failing. +/// +/// Three different files, because Linux reports these three things in three +/// different scopes and getting that wrong produces confident nonsense: +/// +/// * CPU time comes from `/proc/self/stat`, whose `utime`/`stime` *are* +/// summed over the whole thread group. `status` carries no CPU accounting at +/// all, and reading only it is how a harness reports `0.00 cpu-s`. +/// * RSS and the thread count come from `/proc/self/status`, which is +/// process-wide for those fields. +/// * Context switches are summed over `/proc/self/task/*/status`, because the +/// counters in `/proc/self/status` describe **only the leader thread**. A +/// pool of sixteen workers parking and waking on every dispatch moves the +/// leader's counter by approximately nothing, so reading the process file +/// reports "nothing ever parks" no matter what the pool does. This is the +/// measurement that the park/spin question turns on. The sum only covers +/// threads that are *still alive*: a thread that exits between two samples +/// takes its counters with it, so a teardown-heavy workload undercounts. +/// Pool workers live for the whole measured window, so this does not affect +/// the numbers this harness reports. +pub fn sample_process_metrics() -> ProcessMetrics { + let mut metrics = std::fs::read_to_string("/proc/self/status") + .map(|text| parse_status(&text)) + .unwrap_or_default(); + if let Some((utime, stime)) = std::fs::read_to_string("/proc/self/stat") + .ok() + .and_then(|line| parse_stat_cpu_ticks(&line)) + { + let ticks = clock_ticks_per_second(); + metrics.user_us = utime * 1_000_000 / ticks; + metrics.sys_us = stime * 1_000_000 / ticks; + } + let (voluntary, involuntary) = sum_thread_ctxt_switches(); + metrics.voluntary_ctxt_switches = voluntary; + metrics.involuntary_ctxt_switches = involuntary; + metrics +} + +/// Sums `(voluntary, nonvoluntary)` context switches over every thread. +fn sum_thread_ctxt_switches() -> (u64, u64) { + let Ok(entries) = std::fs::read_dir("/proc/self/task") else { + return (0, 0); + }; + let mut voluntary = 0; + let mut involuntary = 0; + for entry in entries.flatten() { + if let Ok(text) = std::fs::read_to_string(entry.path().join("status")) { + let metrics = parse_status(&text); + voluntary += metrics.voluntary_ctxt_switches; + involuntary += metrics.involuntary_ctxt_switches; + } + } + (voluntary, involuntary) +} + +/// One OS thread's identity and accumulated cost. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ThreadInfo { + pub tid: u64, + /// `comm`, which is what shows up in `top -H` and in a debugger. Rust's + /// `std::thread::Builder::name` sets it; Rayon's default pools do not, + /// which is exactly why unnamed threads are hard to attribute. + pub name: String, + pub cpu_us: u64, + pub voluntary_ctxt_switches: u64, + pub involuntary_ctxt_switches: u64, +} + +/// Every thread in this process, with its name and CPU time. +/// +/// This is the tool for attributing pool threads to their owner. A thread with +/// no name is either a Rayon worker from a pool built without a thread-name +/// callback, or a thread created by a C library (ORT's intra-op pool, the +/// allocator's background threads). Naming our own pools is what makes the +/// remainder identifiable by elimination. +pub fn thread_census() -> Vec { + let Ok(entries) = std::fs::read_dir("/proc/self/task") else { + return Vec::new(); + }; + let ticks = clock_ticks_per_second(); + let mut threads = Vec::new(); + for entry in entries.flatten() { + let Ok(tid) = entry.file_name().to_string_lossy().parse::() else { + continue; + }; + let path = entry.path(); + let name = std::fs::read_to_string(path.join("comm")) + .map(|text| text.trim().to_string()) + .unwrap_or_default(); + let cpu_us = std::fs::read_to_string(path.join("stat")) + .ok() + .and_then(|line| parse_stat_cpu_ticks(&line)) + .map(|(utime, stime)| (utime + stime) * 1_000_000 / ticks) + .unwrap_or(0); + let status = std::fs::read_to_string(path.join("status")) + .map(|text| parse_status(&text)) + .unwrap_or_default(); + threads.push(ThreadInfo { + tid, + name, + cpu_us, + voluntary_ctxt_switches: status.voluntary_ctxt_switches, + involuntary_ctxt_switches: status.involuntary_ctxt_switches, + }); + } + threads.sort_by(|a, b| b.cpu_us.cmp(&a.cpu_us).then(a.tid.cmp(&b.tid))); + threads +} + +/// Groups a census by thread name, returning `(name, count, total_cpu_us)` +/// sorted by descending count. Unnamed threads are grouped under `""`. +pub fn census_by_name(threads: &[ThreadInfo]) -> Vec<(String, usize, u64)> { + let mut groups: Vec<(String, usize, u64)> = Vec::new(); + for thread in threads { + let key = if thread.name.is_empty() { + "" + } else { + thread.name.as_str() + }; + match groups.iter_mut().find(|(name, _, _)| name == key) { + Some(group) => { + group.1 += 1; + group.2 += thread.cpu_us; + } + None => groups.push((key.to_string(), 1, thread.cpu_us)), + } + } + groups.sort_by(|a, b| b.1.cmp(&a.1).then(a.0.cmp(&b.0))); + groups +} + +/// Finds where a timing series stops trending, i.e. where warm-up ends. +/// +/// Returns the first index `s` such that every non-overlapping block of +/// `window` samples in `samples[s..]` has a median within `tolerance` +/// (relative) of the median of the whole tail `samples[s..]`. `None` means the +/// series never settled -- which is itself the answer: the run was too short +/// to have a steady state, and any number taken from it describes the +/// transient. +/// +/// Three deliberate choices, each of which was wrong in an earlier version: +/// +/// * **Medians, not means.** A shared host delivers occasional multi- +/// millisecond outliers that are another tenant's scheduling decision, not +/// ours. One of those moves a window's mean past any sane tolerance and +/// reports "never settled" for a series that plainly settled. +/// * **Non-overlapping blocks, not a sliding window.** Requiring every one of +/// ~n overlapping windows to pass makes the verdict hostage to a single +/// blip anywhere in the run, because that blip appears in `window` +/// consecutive windows. Blocks reduce the test to `n/window` weakly +/// correlated checks. +/// * **Compared against the whole tail, not against the final block.** The +/// final block cannot be its own reference -- it always matches, which is +/// how a monotonically rising series gets declared steady at the very end. +/// +/// At least two full blocks must fit, so a run too short to demonstrate +/// stability reports `None` rather than a confident answer. +/// +/// This replaces the guessed `--warmups` constant. A 32-wide pool and a 4-wide +/// pool have warm-ups that differ by more than an order of magnitude, so one +/// constant cannot be right for both, and picking the constant per-arm by hand +/// is how an A/B ends up comparing a warm arm against a cold one. +pub fn steady_state_start(samples: &[f64], window: usize, tolerance: f64) -> Option { + if window == 0 || samples.len() < window * 2 { + return None; + } + let median_of = |slice: &[f64]| -> f64 { + let mut sorted = slice.to_vec(); + sorted.sort_by(f64::total_cmp); + sorted[sorted.len() / 2] + }; + let settled_from = |start: usize| -> bool { + let tail = &samples[start..]; + if tail.len() < window * 2 { + return false; + } + let reference = median_of(tail); + if reference <= 0.0 { + return false; + } + tail.chunks_exact(window) + .all(|block| ((median_of(block) - reference) / reference).abs() <= tolerance) + }; + (0..=samples.len() - window * 2).find(|&start| settled_from(start)) +} + +/// Threads present in `after` but not in `before`, i.e. created during the +/// phase the two censuses bracket. +/// +/// This is the attribution tool. An unnamed thread is anonymous in a census +/// but not in a *delta*: bracketing "build the native session", "build the ORT +/// session" and "run the first inference" separately says which component +/// created it, which is the difference between an unexplained thread count and +/// an owned one. +pub fn census_delta(before: &[ThreadInfo], after: &[ThreadInfo]) -> Vec { + after + .iter() + .filter(|thread| !before.iter().any(|earlier| earlier.tid == thread.tid)) + .cloned() + .collect() +} + +/// Nearest-rank percentile summary of a timing series. +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct Summary { + pub count: usize, + pub min: f64, + pub p50: f64, + pub p90: f64, + pub p99: f64, + pub max: f64, + pub mean: f64, +} + +impl Summary { + /// `samples` must be non-empty. + pub fn from(samples: &[f64]) -> Self { + assert!(!samples.is_empty(), "Summary::from needs a sample"); + let mut sorted = samples.to_vec(); + sorted.sort_by(f64::total_cmp); + let rank = |fraction: f64| -> f64 { + let index = ((sorted.len() as f64) * fraction).ceil().max(1.0) as usize; + sorted[index.min(sorted.len()) - 1] + }; + Self { + count: sorted.len(), + min: sorted[0], + p50: sorted[sorted.len() / 2], + p90: rank(0.9), + p99: rank(0.99), + max: sorted[sorted.len() - 1], + mean: sorted.iter().sum::() / sorted.len() as f64, + } + } + + /// p90/p50. A quiet host sits near 1.0; this run's host does not, and the + /// ratio is printed so a reader can tell which. + pub fn spread(&self) -> f64 { + if self.p50 == 0.0 { + 0.0 + } else { + self.p90 / self.p50 + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn a_zero_mean_gap_is_always_zero() { + let mut gaps = GapDistribution::new(0, 0.5, GapKind::Busy, 7); + for _ in 0..16 { + assert_eq!(gaps.next_gap(), Duration::ZERO); + } + } + + #[test] + fn zero_jitter_repeats_the_mean_exactly() { + let mut gaps = GapDistribution::new(25, 0.0, GapKind::Busy, 7); + for _ in 0..16 { + assert_eq!(gaps.next_gap(), Duration::from_micros(25)); + } + } + + #[test] + fn jittered_gaps_stay_inside_the_requested_band_and_actually_vary() { + let mut gaps = GapDistribution::new(100, 0.5, GapKind::Busy, 12345); + let drawn = (0..512).map(|_| gaps.next_gap()).collect::>(); + for gap in &drawn { + let micros = gap.as_micros() as u64; + assert!( + (50..=150).contains(µs), + "gap {micros}us escaped the +/-50% band" + ); + } + let distinct = drawn.iter().collect::>(); + assert!( + distinct.len() > 16, + "jitter produced only {} distinct gaps", + distinct.len() + ); + } + + #[test] + fn jitter_is_clamped_so_a_gap_can_never_be_negative() { + let mut gaps = GapDistribution::new(10, 9.0, GapKind::Busy, 3); + for _ in 0..256 { + let micros = gaps.next_gap().as_micros() as u64; + assert!((0..=20).contains(µs), "gap {micros}us escaped clamp"); + } + } + + #[test] + fn a_seed_reproduces_its_sequence() { + let draw = || { + let mut gaps = GapDistribution::new(80, 0.4, GapKind::Sleep, 99); + (0..32).map(|_| gaps.next_gap()).collect::>() + }; + assert_eq!(draw(), draw()); + } + + #[test] + fn a_zero_seed_still_jitters() { + let mut gaps = GapDistribution::new(100, 0.5, GapKind::Busy, 0); + let drawn = (0..64).map(|_| gaps.next_gap()).collect::>(); + let distinct = drawn.iter().collect::>(); + assert!(distinct.len() > 8, "seed 0 collapsed to a constant gap"); + } + + #[test] + fn mixed_alternates_between_spinning_and_sleeping() { + let mut gaps = GapDistribution::new(10, 0.0, GapKind::Mixed, 1); + let mut kinds = Vec::new(); + for _ in 0..4 { + kinds.push(gaps.next_kind()); + let _ = gaps.next_gap(); + } + assert_eq!( + kinds, + vec![GapKind::Busy, GapKind::Sleep, GapKind::Busy, GapKind::Sleep] + ); + } + + #[test] + fn a_fixed_kind_never_alternates() { + let mut gaps = GapDistribution::new(10, 0.0, GapKind::Sleep, 1); + for _ in 0..4 { + assert_eq!(gaps.next_kind(), GapKind::Sleep); + let _ = gaps.next_gap(); + } + } + + #[test] + fn gap_kinds_parse_and_reject_junk() { + assert_eq!(GapKind::parse("busy"), Ok(GapKind::Busy)); + assert_eq!(GapKind::parse("sleep"), Ok(GapKind::Sleep)); + assert_eq!(GapKind::parse("mixed"), Ok(GapKind::Mixed)); + assert!(GapKind::parse("spin").is_err()); + } + + const SAMPLE_STATUS: &str = "\ +Name:\tbench_decode_gap +Umask:\t0002 +State:\tR (running) +Tgid:\t4242 +Threads:\t33 +VmHWM:\t 184320 kB +VmRSS:\t 151552 kB +voluntary_ctxt_switches:\t1611 +nonvoluntary_ctxt_switches:\t97 +"; + + #[test] + fn status_parses_the_fields_the_harness_reports() { + let metrics = parse_status(SAMPLE_STATUS); + assert_eq!(metrics.threads, 33); + assert_eq!(metrics.rss_kb, 151_552); + assert_eq!(metrics.peak_rss_kb, 184_320); + assert_eq!(metrics.voluntary_ctxt_switches, 1611); + assert_eq!(metrics.involuntary_ctxt_switches, 97); + } + + #[test] + fn status_parsing_survives_a_kernel_without_those_fields() { + let metrics = parse_status("Name:\tx\nState:\tS (sleeping)\n"); + assert_eq!(metrics, ProcessMetrics::default()); + } + + #[test] + fn deltas_difference_counters_and_carry_levels() { + let before = ProcessMetrics { + user_us: 100, + sys_us: 40, + rss_kb: 900, + peak_rss_kb: 950, + threads: 4, + voluntary_ctxt_switches: 10, + involuntary_ctxt_switches: 2, + }; + let after = ProcessMetrics { + user_us: 350, + sys_us: 90, + rss_kb: 1200, + peak_rss_kb: 1300, + threads: 33, + voluntary_ctxt_switches: 610, + involuntary_ctxt_switches: 5, + }; + let delta = after.since(&before); + assert_eq!(delta.user_us, 250); + assert_eq!(delta.sys_us, 50); + assert_eq!(delta.cpu_us(), 300); + assert_eq!(delta.voluntary_ctxt_switches, 600); + assert_eq!(delta.involuntary_ctxt_switches, 3); + // Levels are carried, not differenced. + assert_eq!(delta.rss_kb, 1200); + assert_eq!(delta.peak_rss_kb, 1300); + assert_eq!(delta.threads, 33); + } + + #[test] + fn a_counter_read_out_of_order_reports_zero_rather_than_wrapping() { + let high = ProcessMetrics { + user_us: 500, + voluntary_ctxt_switches: 900, + ..ProcessMetrics::default() + }; + let delta = ProcessMetrics::default().since(&high); + assert_eq!(delta.user_us, 0); + assert_eq!(delta.voluntary_ctxt_switches, 0); + } + + #[test] + fn stat_cpu_parsing_survives_a_comm_containing_spaces_and_parens() { + let line = "4242 (weird ) name) R 1 4242 4242 0 -1 4194560 100 0 0 0 \ + 731 219 0 0 20 0 33 0 900"; + assert_eq!(parse_stat_cpu_ticks(line), Some((731, 219))); + } + + #[test] + fn stat_cpu_parsing_rejects_a_truncated_line() { + assert_eq!(parse_stat_cpu_ticks("4242 (x) R 1 2 3"), None); + assert_eq!(parse_stat_cpu_ticks("no parens here"), None); + } + + fn thread(tid: u64, name: &str, cpu_us: u64) -> ThreadInfo { + ThreadInfo { + tid, + name: name.to_string(), + cpu_us, + voluntary_ctxt_switches: 0, + involuntary_ctxt_switches: 0, + } + } + + #[test] + fn a_census_groups_by_name_and_buckets_the_unnamed() { + let threads = vec![ + thread(1, "bench_decode_gap", 500), + thread(2, "nxrt-decode-0", 300), + thread(3, "nxrt-decode-1", 280), + thread(4, "", 40), + thread(5, "", 35), + thread(6, "", 30), + ]; + let groups = census_by_name(&threads); + assert_eq!(groups[0], ("".to_string(), 3, 105)); + assert!(groups.contains(&("nxrt-decode-0".to_string(), 1, 300))); + assert!(groups.contains(&("bench_decode_gap".to_string(), 1, 500))); + } + + #[test] + fn an_empty_census_groups_to_nothing() { + assert!(census_by_name(&[]).is_empty()); + } + + #[test] + fn a_census_delta_reports_only_threads_created_in_the_phase() { + let before = vec![thread(1, "main", 10), thread(2, "", 5)]; + let after = vec![ + thread(1, "main", 20), + thread(2, "", 9), + thread(7, "nxrt-task-0", 3), + thread(8, "", 1), + ]; + let created = census_delta(&before, &after); + assert_eq!(created.len(), 2); + assert!(created.iter().any(|t| t.tid == 7)); + assert!(created.iter().any(|t| t.tid == 8)); + } + + #[test] + fn a_census_delta_over_a_phase_that_created_nothing_is_empty() { + let census = vec![thread(1, "main", 10), thread(2, "", 5)]; + assert!(census_delta(&census, &census).is_empty()); + } + + #[test] + fn a_census_delta_ignores_threads_that_exited() { + let before = vec![thread(1, "main", 10), thread(2, "", 5)]; + let after = vec![thread(1, "main", 20)]; + assert!(census_delta(&before, &after).is_empty()); + } + + #[cfg(target_os = "linux")] + #[test] + fn process_metrics_report_real_cpu_and_threads_on_linux() { + // Burn a little CPU so utime is certain to be non-zero, then check the + // sampler actually observed it. This is the regression guard for + // reading CPU out of `/proc/self/status`, which does not carry it. + let started = Instant::now(); + let mut sink = 0u64; + while started.elapsed() < Duration::from_millis(60) { + sink = sink.wrapping_add(started.elapsed().as_nanos() as u64); + } + std::hint::black_box(sink); + let metrics = sample_process_metrics(); + assert!(metrics.threads >= 1, "census saw no threads"); + assert!(metrics.rss_kb > 0, "no RSS reported"); + assert!(metrics.cpu_us() > 0, "no CPU time reported"); + } + + #[cfg(target_os = "linux")] + #[test] + fn ctxt_switches_are_summed_over_every_thread_not_just_the_leader() { + // A parked-then-woken worker records a voluntary switch on *its own* + // task file, not the leader's. Park a spawned thread deliberately and + // check the process-wide sampler sees it: reading only + // `/proc/self/status` would report zero here, which is exactly the bug + // that made a spinning pool and a parking pool look identical. + // + // The worker is sampled while it is still alive, because a thread that + // has exited takes its `/proc/self/task/` entry -- and its + // counters -- with it. + use std::sync::Arc; + use std::sync::atomic::{AtomicBool, Ordering}; + + let slept = Arc::new(AtomicBool::new(false)); + let stop = Arc::new(AtomicBool::new(false)); + let before = sample_process_metrics(); + let handle = { + let slept = Arc::clone(&slept); + let stop = Arc::clone(&stop); + std::thread::spawn(move || { + for _ in 0..40 { + std::thread::sleep(Duration::from_micros(200)); + } + slept.store(true, Ordering::SeqCst); + while !stop.load(Ordering::SeqCst) { + std::hint::spin_loop(); + } + }) + }; + while !slept.load(Ordering::SeqCst) { + std::hint::spin_loop(); + } + let delta = sample_process_metrics().since(&before); + stop.store(true, Ordering::SeqCst); + handle.join().expect("worker thread"); + assert!( + delta.voluntary_ctxt_switches >= 20, + "sampler saw only {} voluntary switches for a thread that slept 40 times; \ + counters are probably being read from the leader alone", + delta.voluntary_ctxt_switches + ); + } + + #[cfg(target_os = "linux")] + #[test] + fn a_live_census_sees_at_least_this_thread() { + let threads = thread_census(); + assert!(!threads.is_empty(), "live census was empty"); + let groups = census_by_name(&threads); + assert_eq!( + groups.iter().map(|(_, count, _)| count).sum::(), + threads.len() + ); + } + + #[test] + fn steady_state_is_found_after_a_decaying_transient() { + // A warm-up that decays over the first 20 samples, then flat at 1.0. + let mut samples = (0..20).map(|i| 5.0 - 0.2 * i as f64).collect::>(); + samples.extend(std::iter::repeat_n(1.0, 40)); + let start = steady_state_start(&samples, 8, 0.05).expect("series settles"); + assert!( + (16..=24).contains(&start), + "steady state reported at {start}, expected near the end of the transient" + ); + } + + #[test] + fn an_already_flat_series_is_steady_from_the_first_sample() { + let samples = vec![2.0; 40]; + assert_eq!(steady_state_start(&samples, 8, 0.05), Some(0)); + } + + #[test] + fn a_series_that_never_settles_reports_no_steady_state() { + // Monotonically rising: no early window matches the final one, and the + // final window is not allowed to qualify by matching itself. + let samples = (0..40).map(|i| 1.0 + i as f64).collect::>(); + assert_eq!(steady_state_start(&samples, 8, 0.01), None); + } + + #[test] + fn a_dip_inside_the_transient_is_not_mistaken_for_steady_state() { + // Decays from 5.0, touches the eventual steady value across indices + // 8..16, then climbs back before finally settling at index 24. A + // detector that matched a single window would report the dip. + let mut samples = vec![5.0; 8]; + samples.extend(std::iter::repeat_n(1.0, 8)); + samples.extend(std::iter::repeat_n(4.0, 8)); + samples.extend(std::iter::repeat_n(1.0, 24)); + let start = steady_state_start(&samples, 8, 0.05).expect("series settles eventually"); + // Past the dip is the property under test. A robust estimator is + // allowed a sample or two of slack at the final step, so this asserts + // the dip was rejected rather than pinning an exact index. + assert!( + start >= 16, + "steady state reported at {start}, inside the transient dip at 8..16" + ); + } + + #[test] + fn exactly_two_windows_of_flat_samples_is_the_minimum_settled_series() { + assert_eq!(steady_state_start(&[1.0; 16], 8, 0.05), Some(0)); + assert_eq!(steady_state_start(&[1.0; 15], 8, 0.05), None); + } + + #[test] + fn a_series_shorter_than_two_windows_has_no_steady_state() { + let samples = vec![1.0; 10]; + assert_eq!(steady_state_start(&samples, 8, 0.05), None); + assert_eq!(steady_state_start(&samples, 0, 0.05), None); + } + + #[test] + fn summary_percentiles_use_nearest_rank() { + let samples = (1..=100).map(|i| i as f64).collect::>(); + let summary = Summary::from(&samples); + assert_eq!(summary.count, 100); + assert_eq!(summary.min, 1.0); + assert_eq!(summary.max, 100.0); + assert_eq!(summary.p90, 90.0); + assert_eq!(summary.p99, 99.0); + assert!((summary.mean - 50.5).abs() < 1e-9); + } + + #[test] + fn summary_of_one_sample_is_that_sample() { + let summary = Summary::from(&[3.5]); + assert_eq!(summary.p50, 3.5); + assert_eq!(summary.p99, 3.5); + assert_eq!(summary.spread(), 1.0); + } +} diff --git a/crates/onnx-genai-bench/src/lib.rs b/crates/onnx-genai-bench/src/lib.rs index b6b164a927..4a3d7ce71d 100644 --- a/crates/onnx-genai-bench/src/lib.rs +++ b/crates/onnx-genai-bench/src/lib.rs @@ -7,6 +7,9 @@ use onnx_genai_engine::logits::{ use onnx_genai_engine::{ProcessorContext, TokenId}; use tokenizers::Tokenizer; +pub mod decode_gap; +#[cfg(feature = "bench-native")] +pub mod model_io; #[cfg(feature = "bench-native")] pub mod synthetic_decoder; diff --git a/crates/onnx-genai-bench/src/model_io.rs b/crates/onnx-genai-bench/src/model_io.rs new file mode 100644 index 0000000000..4c58b83cca --- /dev/null +++ b/crates/onnx-genai-bench/src/model_io.rs @@ -0,0 +1,532 @@ +//! Synthetic model inputs and native-versus-ORT output comparison. +//! +//! Extracted from the `bench_generic` binary so every model-level benchmark in +//! this crate builds its inputs and judges parity the same way. Two harnesses +//! that synthesize inputs differently are not comparable, and two that apply +//! different tolerances can disagree about whether the same kernel is correct. + +use anyhow::{Context, Result, bail}; +use onnx_genai_ort::{DataType as OrtDataType, Session, Value}; +use onnx_runtime_ir::{DataType as NativeDataType, Dim}; +use onnx_runtime_session::{InferenceSession, Tensor}; + +/// Machine epsilon of IEEE binary16 (`2^-10`). f16 carries a 10-bit mantissa, +/// so two f16 values that differ by one ULP near 1.0 differ by this much. +pub const F16_EPSILON: f32 = 9.765_625e-4; + +pub struct InputPair { + pub name: String, + pub shape: Vec, + pub native: Tensor, + pub ort: Value, +} + +#[derive(Debug)] +pub struct OutputDiff { + pub index: usize, + pub max_abs: f32, + pub max_rel: f32, + pub pass: bool, +} + +pub fn parse_shape(value: &str) -> std::result::Result, String> { + let shape = value + .split([',', 'x', 'X']) + .map(|dim| { + dim.trim() + .parse::() + .map_err(|error| format!("invalid dimension '{dim}': {error}")) + }) + .collect::, _>>()?; + if shape.is_empty() || shape.contains(&0) { + return Err("input shape must contain only positive dimensions".to_string()); + } + Ok(shape) +} + +pub fn validate_tolerance(name: &str, value: f32) -> Result<()> { + if !value.is_finite() || value < 0.0 { + bail!("--{name} must be finite and non-negative"); + } + Ok(()) +} + +pub fn resolved_shape(declared: &[Dim], override_shape: Option<&[usize]>) -> Result> { + if let Some(shape) = override_shape { + if shape.len() != declared.len() { + bail!( + "--input-shape rank {} does not match declared input rank {}", + shape.len(), + declared.len() + ); + } + return Ok(shape.to_vec()); + } + + let rank = declared.len(); + Ok(declared + .iter() + .enumerate() + .map(|(axis, dim)| { + dim.as_static().unwrap_or_else(|| { + if rank >= 4 && axis >= rank - 2 { + 224 + } else { + 1 + } + }) + }) + .collect()) +} + +pub fn resolved_ort_shape( + declared: &[i64], + override_shape: Option<&[usize]>, +) -> Result> { + if let Some(shape) = override_shape { + if shape.len() != declared.len() { + bail!( + "--input-shape rank {} does not match declared input rank {}", + shape.len(), + declared.len() + ); + } + return Ok(shape.to_vec()); + } + let rank = declared.len(); + declared + .iter() + .enumerate() + .map(|(axis, &dim)| { + if dim > 0 { + usize::try_from(dim).context("declared ORT input dimension exceeds usize") + } else if rank >= 4 && axis >= rank - 2 { + Ok(224) + } else { + Ok(1) + } + }) + .collect() +} + +pub fn element_count(shape: &[usize]) -> Result { + shape.iter().try_fold(1usize, |count, &dim| { + count + .checked_mul(dim) + .context("input shape element count overflow") + }) +} + +pub fn synthetic_f32(count: usize) -> Vec { + (0..count) + .map(|index| ((index.wrapping_mul(37) % 257) as f32 - 128.0) / 128.0) + .collect() +} + +pub fn synthetic_i64(count: usize) -> Vec { + (0..count).map(|index| (index % 17) as i64).collect() +} + +/// Float16 bit patterns for the same values [`synthetic_f32`] produces, so a +/// Float16 graph is fed the numerically closest version of the f32 input. +pub fn synthetic_f16_bits(count: usize) -> Vec { + synthetic_f32(count) + .into_iter() + .map(|value| half::f16::from_f32(value).to_bits()) + .collect() +} + +/// Unsigned 8-bit inputs spread over the whole quantized range (QLinearMatMul +/// and friends interpret these through a scale/zero-point, so the raw spread +/// matters more than the float value). +pub fn synthetic_u8(count: usize) -> Vec { + (0..count) + .map(|index| (index.wrapping_mul(37) % 251) as u8) + .collect() +} + +pub fn synthetic_i8_bytes(count: usize) -> Vec { + (0..count) + .map(|index| (((index.wrapping_mul(37) % 251) as i32 - 125) as i8) as u8) + .collect() +} + +pub fn synthetic_i32(count: usize) -> Vec { + (0..count).map(|index| (index % 17) as i32).collect() +} + +pub fn i32_bytes(values: &[i32]) -> Vec { + values + .iter() + .flat_map(|value| value.to_le_bytes()) + .collect() +} + +pub fn build_inputs( + native_session: &InferenceSession, + ort_session: &Session, + override_shape: Option<&[usize]>, +) -> Result> { + if native_session.inputs().len() != ort_session.inputs().len() { + bail!( + "runtime input-count mismatch: native={} ORT={}", + native_session.inputs().len(), + ort_session.inputs().len() + ); + } + + native_session + .inputs() + .iter() + .enumerate() + .map(|(index, input)| { + let ort_input = &ort_session.inputs()[index]; + if input.name != ort_input.name { + bail!( + "runtime input-name mismatch at index {index}: native='{}' ORT='{}'", + input.name, + ort_input.name + ); + } + let shape = resolved_shape( + &input.shape, + (index == 0).then_some(override_shape).flatten(), + )?; + let ort_shape = shape + .iter() + .map(|&dim| i64::try_from(dim).context("input dimension exceeds i64")) + .collect::>>()?; + let count = element_count(&shape)?; + let (native, ort) = match (input.dtype, ort_input.dtype) { + (NativeDataType::Float32, OrtDataType::Float32) => { + let data = synthetic_f32(count); + ( + Tensor::from_f32(&shape, &data)?, + Value::from_slice_f32(&data, &ort_shape)?, + ) + } + (NativeDataType::Int64, OrtDataType::Int64) => { + let data = synthetic_i64(count); + ( + Tensor::from_i64(&shape, &data)?, + Value::from_slice_i64(&data, &ort_shape)?, + ) + } + (NativeDataType::Int32, OrtDataType::Int32) => { + let bytes = i32_bytes(&synthetic_i32(count)); + ( + Tensor::from_raw(NativeDataType::Int32, shape.clone(), &bytes)?, + Value::from_raw_bytes(bytes, &ort_shape, OrtDataType::Int32)?, + ) + } + (NativeDataType::Float16, OrtDataType::Float16) => { + let bits = synthetic_f16_bits(count); + let bytes = bits.iter().flat_map(|v| v.to_le_bytes()).collect::>(); + ( + Tensor::from_raw(NativeDataType::Float16, shape.clone(), &bytes)?, + Value::from_slice_f16_bits(&bits, &ort_shape)?, + ) + } + (NativeDataType::Uint8, OrtDataType::Uint8) => { + let bytes = synthetic_u8(count); + ( + Tensor::from_raw(NativeDataType::Uint8, shape.clone(), &bytes)?, + Value::from_raw_bytes(bytes, &ort_shape, OrtDataType::Uint8)?, + ) + } + (NativeDataType::Int8, OrtDataType::Int8) => { + let bytes = synthetic_i8_bytes(count); + ( + Tensor::from_raw(NativeDataType::Int8, shape.clone(), &bytes)?, + Value::from_raw_bytes(bytes, &ort_shape, OrtDataType::Int8)?, + ) + } + (native, ort) => bail!( + "input '{}' has unsupported or mismatched dtype: native={native:?} ORT={ort:?}; \ + bench_generic currently synthesizes Float32, Float16, Int32, Int64, Uint8, \ + and Int8 inputs", + input.name + ), + }; + Ok(InputPair { + name: input.name.clone(), + shape, + native, + ort, + }) + }) + .collect() +} + +/// Largest absolute and relative gap between two f32 sequences, plus whether +/// every element is inside `abs_tolerance + rel_tolerance * max(|a|, |b|)`. +pub fn compare_f32( + native: &[f32], + ort: &[f32], + abs_tolerance: f32, + rel_tolerance: f32, +) -> (f32, f32, bool) { + let mut max_abs = 0.0_f32; + let mut max_rel = 0.0_f32; + let mut pass = true; + for (&native, &ort) in native.iter().zip(ort) { + if native == ort { + continue; + } + if !native.is_finite() || !ort.is_finite() { + max_abs = f32::INFINITY; + max_rel = f32::INFINITY; + pass = false; + continue; + } + let abs = (native - ort).abs(); + let rel = abs / native.abs().max(ort.abs()).max(f32::MIN_POSITIVE); + max_abs = max_abs.max(abs); + max_rel = max_rel.max(rel); + pass &= abs <= abs_tolerance + rel_tolerance * native.abs().max(ort.abs()); + } + (max_abs, max_rel, pass) +} + +pub fn compare_outputs( + native: &[Tensor], + ort: &[Value], + abs_tolerance: f32, + rel_tolerance: f32, + f16_abs_tolerance: f32, + f16_rel_tolerance: f32, +) -> Result> { + if native.len() != ort.len() { + bail!( + "runtime output-count mismatch: native={} ORT={}", + native.len(), + ort.len() + ); + } + native + .iter() + .zip(ort) + .enumerate() + .map(|(index, (native, ort))| { + if native + .shape + .iter() + .copied() + .map(|dim| dim as i64) + .ne(ort.shape().iter().copied()) + { + bail!( + "output {index} shape mismatch: native={:?} ORT={:?}", + native.shape, + ort.shape() + ); + } + match (native.dtype, ort.dtype()) { + (NativeDataType::Float32, OrtDataType::Float32) => { + let (max_abs, max_rel, pass) = compare_f32( + &native.to_vec_f32(), + &ort.to_vec_f32()?, + abs_tolerance, + rel_tolerance, + ); + Ok(OutputDiff { + index, + max_abs, + max_rel, + pass, + }) + } + (NativeDataType::Float16, OrtDataType::Float16) => { + let widen = |bits: &[u16]| -> Vec { + bits.iter() + .map(|&bits| half::f16::from_bits(bits).to_f32()) + .collect() + }; + let native_bits: Vec = native + .as_bytes() + .chunks_exact(2) + .map(|pair| u16::from_le_bytes([pair[0], pair[1]])) + .collect(); + // f16-scaled tolerances: the f32 defaults are ~1 f16 ULP, + // which would pass almost any pair of f16 values and make + // the parity check meaningless. + let (max_abs, max_rel, pass) = compare_f32( + &widen(&native_bits), + &widen(&ort.to_vec_f16_bits()?), + f16_abs_tolerance, + f16_rel_tolerance, + ); + Ok(OutputDiff { + index, + max_abs, + max_rel, + pass, + }) + } + (NativeDataType::Uint8, OrtDataType::Uint8) + | (NativeDataType::Int8, OrtDataType::Int8) => { + // Quantized outputs are exact integers: any mismatch is a + // real disagreement, so report the largest code-unit gap and + // require zero of them. + let ort_bytes = ort.to_raw_bytes()?; + let signed = native.dtype == NativeDataType::Int8; + let max_abs = native + .as_bytes() + .iter() + .zip(&ort_bytes) + .map(|(&native, &ort)| { + if signed { + ((native as i8) as i32 - (ort as i8) as i32).unsigned_abs() + } else { + (native as i32 - ort as i32).unsigned_abs() + } + }) + .max() + .unwrap_or(0); + Ok(OutputDiff { + index, + max_abs: max_abs as f32, + max_rel: 0.0, + pass: max_abs == 0, + }) + } + (NativeDataType::Int64, OrtDataType::Int64) => { + let pass = native.to_vec_i64() == ort.to_vec_i64()?; + Ok(OutputDiff { + index, + max_abs: if pass { 0.0 } else { f32::INFINITY }, + max_rel: if pass { 0.0 } else { f32::INFINITY }, + pass, + }) + } + (native_dtype, ort_dtype) => bail!( + "output {index} has unsupported or mismatched dtype: \ + native={native_dtype:?} ORT={ort_dtype:?}; parity supports Float32 and Int64" + ), + } + }) + .collect() +} + +pub fn classifier_top1_native(output: &Tensor) -> Option { + (output.dtype == NativeDataType::Float32 + && output.shape.len() == 2 + && output.shape[0] == 1 + && output.shape[1] > 1) + .then(|| argmax(&output.to_vec_f32())) +} + +pub fn classifier_top1_ort(output: &Value) -> Result> { + Ok((output.dtype() == OrtDataType::Float32 + && output.shape().len() == 2 + && output.shape()[0] == 1 + && output.shape()[1] > 1) + .then(|| output.to_vec_f32().map(|values| argmax(&values))) + .transpose()?) +} + +pub fn argmax(values: &[f32]) -> usize { + values + .iter() + .enumerate() + .max_by(|(_, left), (_, right)| left.total_cmp(right)) + .map_or(0, |(index, _)| index) +} + +pub fn median_ms(mut samples: Vec) -> f64 { + samples.sort_by(f64::total_cmp); + samples[samples.len() / 2] +} + +/// p50/p90/min of one runtime's samples. p90 uses the nearest-rank definition +/// (`ceil(0.9 * n)`-th smallest), so a 10-run comparison reports the 9th +/// sample rather than interpolating; dispersion is reported as p90/p50 so a +/// noisy shared host is visible in the record instead of hidden by the median. +#[derive(Clone, Copy)] +pub struct Stats { + pub p50: f64, + pub p90: f64, + pub min: f64, +} + +impl Stats { + /// Nearest-rank percentiles over `samples`, which must be non-empty (the + /// caller bails on `--runs 0`; this asserts rather than panicking on an + /// out-of-bounds index if a future caller filters samples down to nothing). + pub fn from(mut samples: Vec) -> Self { + assert!( + !samples.is_empty(), + "Stats::from requires at least one timing sample" + ); + samples.sort_by(f64::total_cmp); + let rank = ((samples.len() as f64) * 0.9).ceil().max(1.0) as usize; + Self { + p50: samples[samples.len() / 2], + p90: samples[rank.min(samples.len()) - 1], + min: samples[0], + } + } + + pub fn spread(&self) -> f64 { + self.p90 / self.p50 + } +} + +/// Which CPU-kernel arm this binary was built with. +/// +/// Printed on every result line because the distinction is not cosmetic: `mlas` +/// is not a default feature of `onnx-runtime-ep-cpu`, so an MLAS-linked build +/// does not measure what ships. This binary used to *require* the `mlas` +/// feature, which meant every ratio ever published from it came from the +/// research arm while being read as a production number. Labelling the arm in +/// the output makes that impossible to do again by accident. +pub fn build_arm() -> &'static str { + if cfg!(feature = "mlas") { + "mlas-reference" + } else { + "native" + } +} + +/// Synthetic inputs for an ORT session alone, for arms that never build a +/// native session (an ORT baseline, or the ORT half of a solo-arm comparison). +/// +/// Shares [`resolved_ort_shape`] and the `synthetic_*` generators with +/// [`build_inputs`], so an ORT-only arm is fed byte-identical tensors to the +/// ones the paired arm would have fed it. +pub fn build_ort_inputs( + session: &Session, + override_shape: Option<&[usize]>, +) -> Result> { + session + .inputs() + .iter() + .enumerate() + .map(|(index, input)| { + let shape = resolved_ort_shape( + &input.shape, + (index == 0).then_some(override_shape).flatten(), + )?; + let ort_shape = shape + .iter() + .map(|&dim| i64::try_from(dim).context("input dimension exceeds i64")) + .collect::>>()?; + let count = element_count(&shape)?; + let value = match input.dtype { + OrtDataType::Float32 => Value::from_slice_f32(&synthetic_f32(count), &ort_shape)?, + OrtDataType::Int64 => Value::from_slice_i64(&synthetic_i64(count), &ort_shape)?, + OrtDataType::Int32 => Value::from_raw_bytes( + i32_bytes(&synthetic_i32(count)), + &ort_shape, + OrtDataType::Int32, + )?, + dtype => bail!( + "input '{}' has unsupported dtype {dtype:?}; only Float32, Int32, and \ + Int64 inputs are synthesized", + input.name + ), + }; + Ok((input.name.clone(), value)) + }) + .collect() +}