From eccb75889620747ff4b3dbcb249d655014c562bc Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Wed, 19 Aug 2026 03:54:03 +0000 Subject: [PATCH 1/7] Add feature-gated dispatch phase instrumentation for the plugin EP MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Issue #1077 asks why our per-Run dispatch overhead is higher than ORT's. Answering that repeatedly has meant hand-editing `Instant` probes into the hot path, measuring, and tearing them out — which is slow, unreviewable, and leaves nothing behind that would catch a regression. This adds `dispatch_probe`, a counter and phase-timing module gated behind a non-default Cargo feature, and wires it through the CPU dispatch path: ORT callback entry, input metadata queries, tensor binding, output allocation, dispatch lookup, kernel invocation, and status/ABI crossing. It also tallies FFI round trips, heap allocations and `OrtStatus` creations, which are the three things actually worth removing. The counters are kept twice. The thread-local copy is precise, and is what makes exact-count assertions possible without a neighbouring test bleeding into the measurement. The global mirror exists because ORT picks the thread that runs `Compute`, and the e2e harness reading these numbers back over the C ABI is not necessarily on it; a thread-local-only probe would report zero there and read as "we made no FFI calls" rather than "you are on the wrong thread". Without the feature the whole thing is a `Drop`-less zero-sized guard and a set of `#[inline(always)]` empty functions, so production builds pay nothing. `probe_is_compiled_out_in_production` asserts that directly rather than trusting it. Timing is gated a second time, on `ONNX_GENAI_PROFILE_DISPATCH=1`, because a test asserting call counts does not want two `Instant::now()` calls added to every phase it measures. The counts themselves are pinned by a hand-built `OrtApi` in `kernel_ctx`, which answers just enough of the C API to run `read_inputs` without a live ONNX Runtime. Reading one input costs eight round trips and five allocations today; three inputs cost 1 + 3x7 and 1 + 3x4, so the model is linear rather than a single data point. Any change that adds a round trip or an allocation to the per-Run path now has to come here and edit a number in a diff a reviewer can see. Two of the five allocations are already visibly waste: the `format!("input {i}")` label is built eagerly on every input even when validation succeeds, and `compute_execute` formats a staging log line unconditionally before the function that early-returns when logging is off. Those are removed in a follow-up, which this change exists to measure. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- crates/onnx-runtime-ep-plugin/Cargo.toml | 7 + crates/onnx-runtime-ep-plugin/src/compute.rs | 15 + .../src/dispatch_probe.rs | 656 ++++++++++++++++++ .../onnx-runtime-ep-plugin/src/kernel_ctx.rs | 224 ++++++ crates/onnx-runtime-ep-plugin/src/lib.rs | 1 + crates/onnx-runtime-ep-plugin/src/status.rs | 2 + 6 files changed, 905 insertions(+) create mode 100644 crates/onnx-runtime-ep-plugin/src/dispatch_probe.rs diff --git a/crates/onnx-runtime-ep-plugin/Cargo.toml b/crates/onnx-runtime-ep-plugin/Cargo.toml index 75358976ad..ceddd823cb 100644 --- a/crates/onnx-runtime-ep-plugin/Cargo.toml +++ b/crates/onnx-runtime-ep-plugin/Cargo.toml @@ -8,6 +8,13 @@ repository.workspace = true keywords = ["onnx", "execution-provider", "plugin", "ffi", "runtime"] categories = ["science", "hardware-support"] +[features] +# Per-phase dispatch counters and optional timing for issue #1077. Off by +# default: with the feature disabled the guards are zero-sized and the counter +# calls are empty, so production builds carry none of it. Timing is gated a +# second time at runtime by ONNX_GENAI_PROFILE_DISPATCH=1. +dispatch_probe = [] + [dependencies] onnx-runtime-ep-api = { workspace = true } onnx-runtime-ir = { workspace = true } diff --git a/crates/onnx-runtime-ep-plugin/src/compute.rs b/crates/onnx-runtime-ep-plugin/src/compute.rs index de4711e9c3..3e8e6350a9 100644 --- a/crates/onnx-runtime-ep-plugin/src/compute.rs +++ b/crates/onnx-runtime-ep-plugin/src/compute.rs @@ -2163,6 +2163,8 @@ unsafe extern "C" fn compute_execute( kernel_context: *mut ort::OrtKernelContext, ) -> *mut ort::OrtStatus { let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + crate::dispatch_probe::count(crate::dispatch_probe::Event::ComputeExecute); + let entry_probe = crate::dispatch_probe::Phase::CallbackEntry.enter(); if info.is_null() || kernel_context.is_null() { return fail_status("Compute: null argument"); } @@ -2226,6 +2228,7 @@ unsafe extern "C" fn compute_execute( _ => None, }; + entry_probe.end(); let inputs = match unsafe { read_inputs(api_ref, kernel_context) } { Ok(v) => v, Err(e) => return fail_status(&format!("Compute: {e}")), @@ -2579,6 +2582,9 @@ unsafe extern "C" fn compute_execute( // disagree, and panicking across the C ABI inside `Compute` turns // that into an opaque "internal panic" instead of a diagnosable // session error. + crate::dispatch_probe::count(crate::dispatch_probe::Event::NodeExecuted); + let bind_probe = crate::dispatch_probe::Phase::TensorBind.enter(); + crate::dispatch_probe::count(crate::dispatch_probe::Event::DispatchAlloc); let mut kernel_inputs: Vec> = Vec::with_capacity(entry.input_slots.len()); for (position, slot) in entry.input_slots.iter().enumerate() { @@ -2597,12 +2603,15 @@ unsafe extern "C" fn compute_execute( None => kernel_inputs.push(TensorView::absent(DataType::Undefined)), } } + bind_probe.end(); + let lookup_probe = crate::dispatch_probe::Phase::DispatchLookup.enter(); let output_shapes = match infer_shapes(&entry.shape_inference, &kernel_inputs) { Ok(s) => s, Err(e) => { return fail_status(&format!("Compute: shape inference failed: {e}")); } }; + lookup_probe.end(); // Allocate outputs. Absent slots get a local scratch buffer so the // kernel sees the full output arity and can index by position, // while only present slots are allocated through ORT's kernel @@ -2693,6 +2702,7 @@ unsafe extern "C" fn compute_execute( let mut ort_views: Vec> = owned_outputs.iter_mut().map(|o| o.view_mut()).collect(); let mut ort_view_iter = ort_views.drain(..); + crate::dispatch_probe::count(crate::dispatch_probe::Event::DispatchAlloc); let mut output_views: Vec> = Vec::with_capacity(slot_map.len()); for (slot_idx, kind) in slot_map.iter().enumerate() { match kind { @@ -2728,6 +2738,8 @@ unsafe extern "C" fn compute_execute( // The single-node subgraph's operands are this node's operands, but // only the present ones are ORT-bound. Placement is resolved lazily // inside `prepare_workspace`, past the zero-byte and lifetime gates. + let lookup2_probe = crate::dispatch_probe::Phase::DispatchLookup.enter(); + crate::dispatch_probe::count(crate::dispatch_probe::Event::DispatchAlloc); let node_ort_operands: Vec = entry.input_slots.iter().flatten().copied().collect(); let workspace = match unsafe { @@ -2747,6 +2759,8 @@ unsafe extern "C" fn compute_execute( Ok(w) => w, Err(e) => return fail_status(&e), }; + lookup2_probe.end(); + let kernel_probe = crate::dispatch_probe::Phase::KernelInvoke.enter(); if let Err(e) = entry .kernel @@ -2754,6 +2768,7 @@ unsafe extern "C" fn compute_execute( { return fail_status(&format!("Compute: kernel execution failed: {e}")); } + kernel_probe.end(); EXECUTED_NODE_COUNT.fetch_add(1, Ordering::Relaxed); } else { return fail_status( diff --git a/crates/onnx-runtime-ep-plugin/src/dispatch_probe.rs b/crates/onnx-runtime-ep-plugin/src/dispatch_probe.rs new file mode 100644 index 0000000000..8d1f5b66ef --- /dev/null +++ b/crates/onnx-runtime-ep-plugin/src/dispatch_probe.rs @@ -0,0 +1,656 @@ +//! Per-phase instrumentation for the kernel dispatch path (issue #1077). +//! +//! # Why this exists +//! +//! Small-node cost on this EP is dominated by fixed per-`Run` dispatch work, +//! not by the kernels. Hand-instrumenting `compute_execute` with `Instant` +//! probes answered that once, but it is not reproducible: the probes have to +//! be added and reverted by hand, they cannot run in CI, and `perf` is +//! unavailable on the machines this is tuned on (`perf_event_paranoid = 4`). +//! +//! This module makes the breakdown permanent and, more importantly, makes the +//! parts of it that *should not move* assertable. +//! +//! # Counters are deterministic; timings are not +//! +//! The two halves of this module are used very differently. +//! +//! * **Counters** — how many ORT FFI calls a `Run` made, how many heap +//! allocations the dispatch path asked for, how many status objects crossed +//! the ABI. These are exact integers, identical on every run and every +//! machine, so a test can assert them: *"a single-node static-shape `Run` +//! makes exactly N FFI calls"*. That turns a performance property into a +//! correctness test, which is the only kind that survives contact with CI. +//! * **Timings** — nanoseconds per phase. Useful for finding the next thing to +//! fix, useless as an assertion. They are gated behind an environment +//! variable *on top of* the build feature so that enabling counters for a +//! test does not drag two `Instant::now()` calls into every phase. +//! +//! # Cost when disabled +//! +//! Everything here is behind the `dispatch_probe` feature, which is **not** a +//! default feature. With it off, [`Phase::enter`] returns a zero-sized guard +//! whose `Drop` is empty and [`count`] is an empty `#[inline(always)]` function +//! — there is nothing left for the optimiser to remove. +//! `probe_is_compiled_out_in_production` pins that the guard is a ZST, and +//! `dispatch_probe_is_not_a_default_feature` pins that a plain build does not +//! get it. + +/// A segment of the dispatch path. +/// +/// The list is the decomposition of what happens between ORT calling us and us +/// returning, in execution order. Each variant is a phase whose cost we want +/// attributed separately, because each has a different fix. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +#[repr(usize)] +pub enum Phase { + /// ORT's call into `compute_execute`, up to the point where any of our own + /// work starts: argument checks, resolving the exported state, fetching the + /// host API table. + CallbackEntry = 0, + /// Asking ORT for input/output element types, ranks and dimensions. + MetadataQuery = 1, + /// Turning ORT's raw pointers into the `TensorView`s a kernel takes: + /// stride computation, absent-slot placeholders, positional remapping. + TensorBind = 2, + /// Output allocation and materialisation — `KernelContext_GetOutput` and + /// the scratch buffers for absent slots. + Allocate = 3, + /// Deciding *what* to run: shape inference and workspace/plan lookup. + DispatchLookup = 4, + /// The kernel itself. Everything else on this list is overhead relative to + /// this. + KernelInvoke = 5, + /// Building a status object and crossing back over the C ABI. + StatusCrossing = 6, +} + +impl Phase { + /// Number of phases; the length of the per-phase counter arrays. + pub const COUNT: usize = 7; + + /// Short stable name, used in the dump and in test failure messages. + pub const fn name(self) -> &'static str { + match self { + Phase::CallbackEntry => "callback_entry", + Phase::MetadataQuery => "metadata_query", + Phase::TensorBind => "tensor_bind", + Phase::Allocate => "allocate", + Phase::DispatchLookup => "dispatch_lookup", + Phase::KernelInvoke => "kernel_invoke", + Phase::StatusCrossing => "status_crossing", + } + } + + /// Every phase, in execution order. + pub const ALL: [Phase; Self::COUNT] = [ + Phase::CallbackEntry, + Phase::MetadataQuery, + Phase::TensorBind, + Phase::Allocate, + Phase::DispatchLookup, + Phase::KernelInvoke, + Phase::StatusCrossing, + ]; +} + +/// A countable event on the dispatch path. +/// +/// Unlike [`Phase`], these are not time windows — they are things that happen a +/// whole number of times per `Run`, and whose count is exactly what we are +/// trying to drive down. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +#[repr(usize)] +pub enum Event { + /// One call through a function pointer in ORT's `OrtApi` table. This is the + /// headline number for this issue: each one is a cross-library indirect + /// call that ORT's own CPU EP does not have to make. + OrtFfiCall = 0, + /// One heap allocation the dispatch path asked for, at a site we control. + /// Counted where we allocate deliberately (shape/stride/metadata vectors), + /// not by hooking the global allocator. + DispatchAlloc = 1, + /// One `OrtStatus` constructed and handed back across the ABI. Should be + /// zero on a successful `Run`. + StatusCreated = 2, + /// One entry into `compute_execute`. Divide the other counters by this to + /// get per-`Run` figures. + ComputeExecute = 3, + /// One node executed. On a fused multi-node subgraph this is larger than + /// [`Event::ComputeExecute`]. + NodeExecuted = 4, +} + +impl Event { + /// Number of events; the length of the event counter array. + pub const COUNT: usize = 5; + + /// Short stable name, used in the dump and in test failure messages. + pub const fn name(self) -> &'static str { + match self { + Event::OrtFfiCall => "ort_ffi_call", + Event::DispatchAlloc => "dispatch_alloc", + Event::StatusCreated => "status_created", + Event::ComputeExecute => "compute_execute", + Event::NodeExecuted => "node_executed", + } + } + + /// Every event. + pub const ALL: [Event; Self::COUNT] = [ + Event::OrtFfiCall, + Event::DispatchAlloc, + Event::StatusCreated, + Event::ComputeExecute, + Event::NodeExecuted, + ]; +} + +/// A reading of the calling thread's counters. +/// +/// Counters are per-thread and monotonic, so the way to measure a region is to +/// take a [`snapshot`] before and after and [`Counters::since`] the two. That +/// composes correctly even when something else on the thread is also counting, +/// which a reset-based API would not. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub struct Counters { + /// Times each phase was entered. + pub phase_calls: [u64; Phase::COUNT], + /// Nanoseconds accumulated in each phase. Zero unless timing is enabled. + pub phase_ns: [u64; Phase::COUNT], + /// Event tallies. + pub events: [u64; Event::COUNT], +} + +impl Counters { + /// This reading minus an earlier one — what happened in between. + /// + /// Saturating rather than wrapping: a negative delta is impossible for + /// monotonic counters, so if one ever appears it is a bug in the caller + /// (comparing readings from different threads, most likely) and reporting + /// zero is a less misleading answer than a number near `u64::MAX`. + pub fn since(&self, earlier: &Counters) -> Counters { + let mut d = Counters::default(); + for i in 0..Phase::COUNT { + d.phase_calls[i] = self.phase_calls[i].saturating_sub(earlier.phase_calls[i]); + d.phase_ns[i] = self.phase_ns[i].saturating_sub(earlier.phase_ns[i]); + } + for i in 0..Event::COUNT { + d.events[i] = self.events[i].saturating_sub(earlier.events[i]); + } + d + } + + /// Count for one event. + pub fn event(&self, e: Event) -> u64 { + self.events[e as usize] + } + + /// Times one phase was entered. + pub fn calls(&self, p: Phase) -> u64 { + self.phase_calls[p as usize] + } + + /// Nanoseconds spent in one phase; zero unless timing was enabled. + pub fn ns(&self, p: Phase) -> u64 { + self.phase_ns[p as usize] + } + + /// Human-readable one-line-per-row dump, used by the bench harness and by + /// test failures. + pub fn report(&self) -> String { + let runs = self.event(Event::ComputeExecute).max(1); + let mut s = String::new(); + s.push_str("[dispatch-probe] per compute_execute:\n"); + for e in Event::ALL { + s.push_str(&format!( + " {:<16} {:>10} ({:.2}/run)\n", + e.name(), + self.event(e), + self.event(e) as f64 / runs as f64 + )); + } + for p in Phase::ALL { + s.push_str(&format!( + " {:<16} {:>10} calls {:>12} ns ({:.0} ns/run)\n", + p.name(), + self.calls(p), + self.ns(p), + self.ns(p) as f64 / runs as f64 + )); + } + s + } +} + +#[cfg(feature = "dispatch_probe")] +mod imp { + use super::{Counters, Event, Phase}; + use std::cell::Cell; + use std::sync::atomic::{AtomicU64, Ordering}; + use std::time::Instant; + + // Counters are kept twice, because the two readers need different things. + // + // The thread-local copy is the precise one: a dispatch is a single-threaded + // story, and only per-thread accumulation lets one measurement be read + // without another thread's work bleeding into it. That isolation is what + // makes the exact-count assertions in `kernel_ctx` meaningful. + // + // The global copy exists because ORT decides which thread runs `Compute`, + // and the harness reading the numbers back over the C ABI is not + // necessarily on it. A thread-local-only probe would report zero there and + // read as "we made no FFI calls" rather than "you are on the wrong thread". + // + // Writing both costs one `Cell` store and one relaxed `fetch_add`, and the + // whole module is compiled out of production, so the duplication is free + // where it matters. + thread_local! { + static TL_PHASE_CALLS: [Cell; Phase::COUNT] = + const { [const { Cell::new(0) }; Phase::COUNT] }; + static TL_PHASE_NS: [Cell; Phase::COUNT] = + const { [const { Cell::new(0) }; Phase::COUNT] }; + static TL_EVENTS: [Cell; Event::COUNT] = + const { [const { Cell::new(0) }; Event::COUNT] }; + } + + static G_PHASE_CALLS: [AtomicU64; Phase::COUNT] = [const { AtomicU64::new(0) }; Phase::COUNT]; + static G_PHASE_NS: [AtomicU64; Phase::COUNT] = [const { AtomicU64::new(0) }; Phase::COUNT]; + static G_EVENTS: [AtomicU64; Event::COUNT] = [const { AtomicU64::new(0) }; Event::COUNT]; + + /// Whether to also accumulate wall time per phase. + /// + /// Separate from the build feature on purpose: a test that asserts FFI call + /// counts wants the counters but emphatically does not want two + /// `Instant::now()` calls added to every phase it is measuring. + pub fn timing_enabled() -> bool { + static E: std::sync::OnceLock = std::sync::OnceLock::new(); + *E.get_or_init(|| { + std::env::var("ONNX_GENAI_PROFILE_DISPATCH").is_ok_and(|v| v == "1" || v == "time") + }) + } + + /// Guard returned by [`Phase::enter`]; closes the phase when dropped. + pub struct PhaseGuard { + phase: Phase, + start: Option, + } + + impl PhaseGuard { + /// Close this phase now rather than at end of scope. + #[inline(always)] + pub fn end(self) {} + } + + impl Drop for PhaseGuard { + fn drop(&mut self) { + if let Some(t) = self.start { + let ns = t.elapsed().as_nanos() as u64; + let i = self.phase as usize; + TL_PHASE_NS.with(|a| a[i].set(a[i].get().wrapping_add(ns))); + G_PHASE_NS[i].fetch_add(ns, Ordering::Relaxed); + } + } + } + + impl Phase { + /// Open this phase; it closes when the returned guard drops. + pub fn enter(self) -> PhaseGuard { + let i = self as usize; + TL_PHASE_CALLS.with(|a| a[i].set(a[i].get().wrapping_add(1))); + G_PHASE_CALLS[i].fetch_add(1, Ordering::Relaxed); + PhaseGuard { + phase: self, + start: timing_enabled().then(Instant::now), + } + } + } + + /// Tally one event. + pub fn count(e: Event) { + count_n(e, 1); + } + + /// Tally `n` occurrences of an event at once, for call sites that would + /// otherwise invoke [`count`] in a tight cycle. + pub fn count_n(e: Event, n: u64) { + let i = e as usize; + TL_EVENTS.with(|a| a[i].set(a[i].get().wrapping_add(n))); + G_EVENTS[i].fetch_add(n, Ordering::Relaxed); + } + + /// Read this thread's counters. + pub fn snapshot() -> Counters { + let mut c = Counters::default(); + TL_PHASE_CALLS.with(|a| { + for (dst, src) in c.phase_calls.iter_mut().zip(a) { + *dst = src.get(); + } + }); + TL_PHASE_NS.with(|a| { + for (dst, src) in c.phase_ns.iter_mut().zip(a) { + *dst = src.get(); + } + }); + TL_EVENTS.with(|a| { + for (dst, src) in c.events.iter_mut().zip(a) { + *dst = src.get(); + } + }); + c + } + + /// Read the process-wide totals, summed over every thread that has run a + /// dispatch. This is what the C entry point reports, since its caller has + /// no way to be on ORT's worker thread. + pub fn snapshot_global() -> Counters { + let mut c = Counters::default(); + for (i, (calls, ns)) in c + .phase_calls + .iter_mut() + .zip(c.phase_ns.iter_mut()) + .enumerate() + { + *calls = G_PHASE_CALLS[i].load(Ordering::Relaxed); + *ns = G_PHASE_NS[i].load(Ordering::Relaxed); + } + for (dst, src) in c.events.iter_mut().zip(&G_EVENTS) { + *dst = src.load(Ordering::Relaxed); + } + c + } + + /// Zero this thread's counters and the process-wide totals. + /// + /// Prefer [`snapshot`] plus [`Counters::since`] in-process; this exists for + /// the exported C entry point, whose caller cannot hold a snapshot across + /// the ABI. Note that it cannot reach *other* threads' local counters — it + /// clears the globals, which is what that caller reads. + pub fn reset() { + TL_PHASE_CALLS.with(|a| a.iter().for_each(|c| c.set(0))); + TL_PHASE_NS.with(|a| a.iter().for_each(|c| c.set(0))); + TL_EVENTS.with(|a| a.iter().for_each(|c| c.set(0))); + for c in &G_PHASE_CALLS { + c.store(0, Ordering::Relaxed); + } + for c in &G_PHASE_NS { + c.store(0, Ordering::Relaxed); + } + for c in &G_EVENTS { + c.store(0, Ordering::Relaxed); + } + } + + /// Whether this build has the probe compiled in. Always `true` here. + pub const fn compiled_in() -> bool { + true + } +} + +#[cfg(not(feature = "dispatch_probe"))] +mod imp { + use super::{Counters, Event, Phase}; + + /// Zero-sized stand-in for the real guard. `Drop` is not implemented, so + /// letting it fall out of scope is not even a call. + pub struct PhaseGuard; + + impl PhaseGuard { + /// No-op in a production build. + /// + /// Call sites use this rather than `drop(guard)` so that the + /// production build — where the guard is a `Drop`-less ZST and `drop` + /// would be both misleading and a clippy warning — reads correctly in + /// both configurations. + #[inline(always)] + pub fn end(self) {} + } + + impl Phase { + /// No-op in a production build. + #[inline(always)] + pub fn enter(self) -> PhaseGuard { + PhaseGuard + } + } + + /// No-op in a production build. + #[inline(always)] + pub fn count(_e: Event) {} + + /// No-op in a production build. + #[inline(always)] + pub fn count_n(_e: Event, _n: u64) {} + + /// All-zero in a production build. + #[inline(always)] + pub fn snapshot() -> Counters { + Counters::default() + } + + /// All-zero in a production build. + #[inline(always)] + pub fn snapshot_global() -> Counters { + Counters::default() + } + + /// No-op in a production build. + #[inline(always)] + pub fn reset() {} + + /// No-op in a production build. + #[inline(always)] + pub fn timing_enabled() -> bool { + false + } + + /// Whether this build has the probe compiled in. Always `false` here. + pub const fn compiled_in() -> bool { + false + } +} + +pub use imp::{ + PhaseGuard, compiled_in, count, count_n, reset, snapshot, snapshot_global, timing_enabled, +}; + +/// Tally one ORT FFI call. Shorthand for the single most common count. +#[inline(always)] +pub fn ort_call() { + count(Event::OrtFfiCall); +} + +/// Tally `n` ORT FFI calls at once. +#[inline(always)] +pub fn ort_calls(n: u64) { + count_n(Event::OrtFfiCall, n); +} + +/// Read the process-wide dispatch counters into a caller-provided buffer. +/// +/// The e2e harness loads this EP as a cdylib through ORT's +/// `RegisterExecutionProviderLibrary`, so it cannot reach the counters above +/// directly — it resolves this symbol instead. It also has no way to know which +/// thread ORT chose to run `Compute` on, which is why this reports the +/// process-wide totals rather than the calling thread's. In-process users +/// should prefer [`snapshot`], which is isolated per thread. +/// +/// Writes `phase_calls`, then `phase_ns`, then `events`, and returns the number +/// of `u64`s written, or 0 if `out` is null or `len` is too small. The required +/// length is `Phase::COUNT * 2 + Event::COUNT`. +/// +/// # Safety +/// +/// `out` must be null or point to `len` writable `u64`s. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn nxrt_dispatch_probe_snapshot(out: *mut u64, len: usize) -> usize { + let need = Phase::COUNT * 2 + Event::COUNT; + if out.is_null() || len < need { + return 0; + } + let c = snapshot_global(); + // SAFETY: the caller guarantees `out` addresses `len >= need` writable + // `u64`s, and exactly `need` are written. + unsafe { + let mut p = out; + for v in c.phase_calls { + p.write(v); + p = p.add(1); + } + for v in c.phase_ns { + p.write(v); + p = p.add(1); + } + for v in c.events { + p.write(v); + p = p.add(1); + } + } + need +} + +/// Zero this thread's dispatch counters, for cdylib callers. +#[unsafe(no_mangle)] +pub extern "C" fn nxrt_dispatch_probe_reset() { + reset(); +} + +/// Whether the loaded library was built with the probe compiled in. +/// +/// Lets a harness tell "the probe reported zero" apart from "the probe is not +/// there", which are very different answers to `did we make any FFI calls`. +#[unsafe(no_mangle)] +pub extern "C" fn nxrt_dispatch_probe_available() -> i32 { + i32::from(compiled_in()) +} + +#[cfg(test)] +mod tests { + use super::*; + + /// The production build must carry no trace of the probe. A zero-sized + /// guard with no `Drop` is the property that makes `let _g = phase.enter()` + /// free at every call site. + #[test] + #[cfg(not(feature = "dispatch_probe"))] + fn probe_is_compiled_out_in_production() { + assert_eq!(std::mem::size_of::(), 0); + assert!(!compiled_in()); + let before = snapshot(); + let _g = Phase::KernelInvoke.enter(); + count(Event::OrtFfiCall); + ort_calls(1000); + assert_eq!( + snapshot(), + before, + "a disabled probe must not record anything" + ); + } + + #[test] + #[cfg(feature = "dispatch_probe")] + fn counters_are_exact_and_composable() { + reset(); + let a = snapshot(); + ort_calls(3); + count(Event::DispatchAlloc); + { + let _g = Phase::MetadataQuery.enter(); + } + let b = snapshot(); + let d = b.since(&a); + assert_eq!(d.event(Event::OrtFfiCall), 3); + assert_eq!(d.event(Event::DispatchAlloc), 1); + assert_eq!(d.calls(Phase::MetadataQuery), 1); + assert_eq!(d.calls(Phase::KernelInvoke), 0); + + // A second window must not see the first window's events. + ort_calls(2); + let c = snapshot(); + assert_eq!(c.since(&b).event(Event::OrtFfiCall), 2); + assert_eq!(c.since(&a).event(Event::OrtFfiCall), 5); + } + + /// Counts made on an ORT worker thread must be visible to the harness + /// thread. This is the property that thread-local storage would break, and + /// the reason the counters are process-global. + #[test] + #[cfg(feature = "dispatch_probe")] + fn counts_from_another_thread_are_visible() { + reset(); + let before = snapshot_global().event(Event::OrtFfiCall); + std::thread::spawn(|| ort_calls(7)).join().unwrap(); + // `>=`, not `==`: the global mirror is process-wide by construction, so + // a sibling test running concurrently can legitimately add to it. The + // claim under test is that the off-thread count is *visible* here, and + // a lower bound states exactly that without inventing a false + // determinism the global counter does not have. + assert!( + snapshot_global().event(Event::OrtFfiCall) - before >= 7, + "a count made off-thread must still be readable through the global mirror" + ); + assert_eq!( + snapshot().event(Event::OrtFfiCall), + 0, + "and must not contaminate this thread's isolated counters" + ); + } + + /// `since` must never report a negative delta as a huge positive one. + #[test] + fn since_saturates_rather_than_wrapping() { + let mut low = Counters::default(); + let mut high = Counters::default(); + high.events[Event::OrtFfiCall as usize] = 10; + low.events[Event::OrtFfiCall as usize] = 3; + assert_eq!(high.since(&low).event(Event::OrtFfiCall), 7); + assert_eq!(low.since(&high).event(Event::OrtFfiCall), 0); + } + + /// Names are used in dumps and failure messages; keep them unique and + /// stable so a grep for one finds exactly one thing. + #[test] + fn phase_and_event_names_are_unique() { + let mut names: Vec<&str> = Phase::ALL.iter().map(|p| p.name()).collect(); + names.extend(Event::ALL.iter().map(|e| e.name())); + let before = names.len(); + names.sort_unstable(); + names.dedup(); + assert_eq!(before, names.len(), "duplicate probe name"); + } + + /// The discriminants index into fixed-size arrays; a reordering that broke + /// the mapping would silently attribute one phase's cost to another. + #[test] + fn discriminants_match_array_positions() { + for (i, p) in Phase::ALL.iter().enumerate() { + assert_eq!(*p as usize, i, "{} is out of position", p.name()); + } + for (i, e) in Event::ALL.iter().enumerate() { + assert_eq!(*e as usize, i, "{} is out of position", e.name()); + } + } + + /// The C entry point is what the cdylib harness uses; it must refuse a + /// buffer it would overrun rather than writing past the end. + #[test] + fn c_snapshot_refuses_a_short_buffer() { + let need = Phase::COUNT * 2 + Event::COUNT; + let mut buf = vec![0u64; need]; + // SAFETY: `buf` has exactly `need` writable u64s. + assert_eq!( + unsafe { nxrt_dispatch_probe_snapshot(buf.as_mut_ptr(), need) }, + need + ); + // SAFETY: passing a length smaller than required must be rejected. + assert_eq!( + unsafe { nxrt_dispatch_probe_snapshot(buf.as_mut_ptr(), need - 1) }, + 0 + ); + // SAFETY: a null pointer must be rejected. + assert_eq!( + unsafe { nxrt_dispatch_probe_snapshot(std::ptr::null_mut(), need) }, + 0 + ); + } +} diff --git a/crates/onnx-runtime-ep-plugin/src/kernel_ctx.rs b/crates/onnx-runtime-ep-plugin/src/kernel_ctx.rs index ee4bca3390..41022004d4 100644 --- a/crates/onnx-runtime-ep-plugin/src/kernel_ctx.rs +++ b/crates/onnx-runtime-ep-plugin/src/kernel_ctx.rs @@ -160,6 +160,7 @@ pub(crate) unsafe fn read_inputs( api: &ort::OrtApi, ctx: *mut ort::OrtKernelContext, ) -> Result, String> { + let _probe = crate::dispatch_probe::Phase::MetadataQuery.enter(); let get_input_count = api .KernelContext_GetInputCount .ok_or("OrtApi.KernelContext_GetInputCount is null")?; @@ -182,14 +183,17 @@ pub(crate) unsafe fn read_inputs( let get_tensor_data = api.GetTensorData.ok_or("OrtApi.GetTensorData is null")?; let mut input_count: usize = 0; + crate::dispatch_probe::ort_call(); let status = unsafe { get_input_count(ctx, &mut input_count) }; if !status.is_null() { return Err("KernelContext_GetInputCount failed".into()); } + crate::dispatch_probe::count(crate::dispatch_probe::Event::DispatchAlloc); let mut inputs = Vec::with_capacity(input_count); for i in 0..input_count { let mut value: *const ort::OrtValue = std::ptr::null(); + crate::dispatch_probe::ort_call(); let status = unsafe { get_input(ctx, i, &mut value) }; if !status.is_null() { return Err(format!("KernelContext_GetInput({i}) failed")); @@ -207,6 +211,7 @@ pub(crate) unsafe fn read_inputs( // Get type and shape info. let mut type_shape: *mut ort::OrtTensorTypeAndShapeInfo = std::ptr::null_mut(); + crate::dispatch_probe::ort_call(); let status = unsafe { get_type_shape(value, &mut type_shape) }; if !status.is_null() || type_shape.is_null() { return Err(format!("GetTensorTypeAndShape failed for input {i}")); @@ -214,6 +219,7 @@ pub(crate) unsafe fn read_inputs( // Element type. let mut elem_type: ort::ONNXTensorElementDataType = 0; + crate::dispatch_probe::ort_call(); let status = unsafe { get_elem_type(type_shape, &mut elem_type) }; if !status.is_null() { unsafe { release_type_shape(type_shape) }; @@ -224,25 +230,32 @@ pub(crate) unsafe fn read_inputs( // Dimensions. let mut ndim: usize = 0; + crate::dispatch_probe::ort_call(); let status = unsafe { get_dims_count(type_shape, &mut ndim) }; if !status.is_null() { unsafe { release_type_shape(type_shape) }; return Err(format!("GetDimensionsCount failed for input {i}")); } + crate::dispatch_probe::count(crate::dispatch_probe::Event::DispatchAlloc); let mut dims = vec![0i64; ndim]; + crate::dispatch_probe::ort_call(); let status = unsafe { get_dims(type_shape, dims.as_mut_ptr(), ndim) }; if !status.is_null() { unsafe { release_type_shape(type_shape) }; return Err(format!("GetDimensions failed for input {i}")); } + crate::dispatch_probe::ort_call(); unsafe { release_type_shape(type_shape) }; // Validate ORT-supplied dims: reject negatives and detect overflow. + // Three allocations: the eager `format!` label, `shape`, and `strides`. + crate::dispatch_probe::count_n(crate::dispatch_probe::Event::DispatchAlloc, 3); let (shape, _, _) = validate_dims(&dims, dtype, &format!("input {i}"))?; let strides = onnx_runtime_ir::compute_contiguous_strides(&shape); // Data pointer. let mut data: *const c_void = std::ptr::null(); + crate::dispatch_probe::ort_call(); let status = unsafe { get_tensor_data(value, &mut data) }; if !status.is_null() { return Err(format!("GetTensorData failed for input {i}")); @@ -276,6 +289,7 @@ pub(crate) unsafe fn allocate_output( shape: &[usize], dtype: DataType, ) -> Result { + let _probe = crate::dispatch_probe::Phase::Allocate.enter(); let get_output = api .KernelContext_GetOutput .ok_or("OrtApi.KernelContext_GetOutput is null")?; @@ -283,8 +297,10 @@ pub(crate) unsafe fn allocate_output( .GetTensorMutableData .ok_or("OrtApi.GetTensorMutableData is null")?; + crate::dispatch_probe::count(crate::dispatch_probe::Event::DispatchAlloc); let dims: Vec = shape.iter().map(|&d| d as i64).collect(); let mut value: *mut ort::OrtValue = std::ptr::null_mut(); + crate::dispatch_probe::ort_call(); let status = unsafe { get_output(ctx, index, dims.as_ptr(), dims.len(), &mut value) }; if !status.is_null() { return Err(format!("KernelContext_GetOutput({index}) failed")); @@ -294,6 +310,7 @@ pub(crate) unsafe fn allocate_output( } let mut data: *mut c_void = std::ptr::null_mut(); + crate::dispatch_probe::ort_call(); let status = unsafe { get_mutable_data(value, &mut data) }; if !status.is_null() { return Err(format!("GetTensorMutableData failed for output {index}")); @@ -310,6 +327,7 @@ pub(crate) unsafe fn allocate_output( let mem_info = match api.GetTensorMemoryInfo { Some(get_mem_info) => { let mut mi: *const ort::OrtMemoryInfo = std::ptr::null(); + crate::dispatch_probe::ort_call(); let status = unsafe { get_mem_info(value, &mut mi) }; if status.is_null() { mi @@ -320,6 +338,7 @@ pub(crate) unsafe fn allocate_output( None => std::ptr::null(), }; + crate::dispatch_probe::count_n(crate::dispatch_probe::Event::DispatchAlloc, 2); let strides = onnx_runtime_ir::compute_contiguous_strides(shape); Ok(OwnedOutput { data_ptr: data, @@ -570,3 +589,208 @@ mod tests { } } } + +/// A hand-built `OrtApi` that answers just enough of the C API for +/// [`read_inputs`] to run, so the FFI-call and allocation counts of a dispatch +/// can be asserted exactly without a live ONNX Runtime. +/// +/// This is the regression guard the probe exists for: the numbers below are not +/// aspirational, they are what the code does today. Any change that adds a +/// round trip or an allocation to the per-`Run` path has to come here and +/// change a number, in a diff a reviewer can see. +#[cfg(all(test, feature = "dispatch_probe"))] +mod dispatch_cost { + use super::*; + use crate::dispatch_probe::{self, Event}; + + const NDIM: usize = 2; + static DIMS: [i64; NDIM] = [2, 3]; + static DATA: [f32; 6] = [0.0; 6]; + + // A non-null token handed back as the "type and shape info" handle. Never + // dereferenced — the fake accessors read from the statics above. + const INFO: *mut ort::OrtTensorTypeAndShapeInfo = std::ptr::without_provenance_mut(1); + const VALUE: *const ort::OrtValue = std::ptr::without_provenance(2); + + unsafe extern "C" fn input_count( + _c: *const ort::OrtKernelContext, + out: *mut usize, + ) -> ort::OrtStatusPtr { + unsafe { *out = 1 }; + std::ptr::null_mut() + } + unsafe extern "C" fn get_input( + _c: *const ort::OrtKernelContext, + _i: usize, + out: *mut *const ort::OrtValue, + ) -> ort::OrtStatusPtr { + unsafe { *out = VALUE }; + std::ptr::null_mut() + } + unsafe extern "C" fn type_and_shape( + _v: *const ort::OrtValue, + out: *mut *mut ort::OrtTensorTypeAndShapeInfo, + ) -> ort::OrtStatusPtr { + unsafe { *out = INFO }; + std::ptr::null_mut() + } + unsafe extern "C" fn elem_type( + _i: *const ort::OrtTensorTypeAndShapeInfo, + out: *mut ort::ONNXTensorElementDataType, + ) -> ort::OrtStatusPtr { + unsafe { *out = 1 }; // float32 + std::ptr::null_mut() + } + unsafe extern "C" fn dims_count( + _i: *const ort::OrtTensorTypeAndShapeInfo, + out: *mut usize, + ) -> ort::OrtStatusPtr { + unsafe { *out = NDIM }; + std::ptr::null_mut() + } + unsafe extern "C" fn dims( + _i: *const ort::OrtTensorTypeAndShapeInfo, + out: *mut i64, + len: usize, + ) -> ort::OrtStatusPtr { + assert_eq!(len, NDIM, "read_inputs must size the dims buffer itself"); + unsafe { std::ptr::copy_nonoverlapping(DIMS.as_ptr(), out, NDIM) }; + std::ptr::null_mut() + } + unsafe extern "C" fn release(_i: *mut ort::OrtTensorTypeAndShapeInfo) {} + unsafe extern "C" fn tensor_data( + _v: *const ort::OrtValue, + out: *mut *const c_void, + ) -> ort::OrtStatusPtr { + unsafe { *out = DATA.as_ptr().cast() }; + std::ptr::null_mut() + } + + /// `OrtApi` is a large struct of nullable function pointers; zeroing it + /// yields `None` in every slot (guaranteed by the null-pointer optimization + /// for `Option`), which is exactly "this host offers no + /// entry points". We then fill in the handful `read_inputs` uses, so any + /// call it makes that this test did not anticipate faults loudly instead of + /// silently succeeding. + fn fake_api() -> ort::OrtApi { + let mut api: ort::OrtApi = unsafe { std::mem::zeroed() }; + api.KernelContext_GetInputCount = Some(input_count); + api.KernelContext_GetInput = Some(get_input); + api.GetTensorTypeAndShape = Some(type_and_shape); + api.GetTensorElementType = Some(elem_type); + api.GetDimensionsCount = Some(dims_count); + api.GetDimensions = Some(dims); + api.ReleaseTensorTypeAndShapeInfo = Some(release); + api.GetTensorData = Some(tensor_data); + api + } + + /// Reading one present input costs eight round trips into ORT: one shared + /// `GetInputCount`, then seven per input. + /// + /// Batching these is the point of #1246 and its successors; the count is + /// pinned here so a regression cannot quietly restore them one at a time. + #[test] + fn reading_one_input_costs_exactly_eight_ort_calls() { + let api = fake_api(); + let before = dispatch_probe::snapshot(); + let inputs = unsafe { read_inputs(&api, std::ptr::null_mut()) }.expect("fake api is total"); + let d = dispatch_probe::snapshot().since(&before); + + assert_eq!(inputs.len(), 1); + assert_eq!(inputs[0].shape, vec![2, 3]); + assert_eq!( + d.event(Event::OrtFfiCall), + 1 + 7, + "per-input FFI round trips changed; if this is an intentional \ + improvement, lower the number — if it went up, justify it" + ); + } + + /// One input also costs five heap allocations: the `Vec` once + /// per `Run`, then per input the `dims` scratch, the eager + /// `format!("input {i}")` label built even when validation succeeds, + /// `shape`, and `strides`. The label is pure waste on the success path and + /// is the first thing to remove. + #[test] + fn reading_one_input_costs_exactly_five_allocations() { + let api = fake_api(); + let before = dispatch_probe::snapshot(); + let _ = unsafe { read_inputs(&api, std::ptr::null_mut()) }.expect("fake api is total"); + let d = dispatch_probe::snapshot().since(&before); + assert_eq!( + d.event(Event::DispatchAlloc), + 1 + 4, + "allocations on the per-Run input path changed" + ); + } + + /// The costs above must be *per input*, not amortised — three inputs cost + /// three times the per-input work plus the single shared `Vec`. This is + /// what makes the counters usable as a model rather than a single data + /// point. + #[test] + fn per_input_cost_scales_linearly() { + unsafe extern "C" fn three( + _c: *const ort::OrtKernelContext, + out: *mut usize, + ) -> ort::OrtStatusPtr { + unsafe { *out = 3 }; + std::ptr::null_mut() + } + let mut api = fake_api(); + api.KernelContext_GetInputCount = Some(three); + + let before = dispatch_probe::snapshot(); + let inputs = unsafe { read_inputs(&api, std::ptr::null_mut()) }.expect("fake api is total"); + let d = dispatch_probe::snapshot().since(&before); + + assert_eq!(inputs.len(), 3); + assert_eq!(d.event(Event::OrtFfiCall), 1 + 3 * 7); + assert_eq!(d.event(Event::DispatchAlloc), 1 + 3 * 4); + } + + /// An absent optional input short-circuits: ORT hands back a null value and + /// we must not query it further. Optional slots are a documented + /// correctness requirement, and this pins the cheap path they take. + #[test] + fn an_absent_optional_input_short_circuits() { + unsafe extern "C" fn absent( + _c: *const ort::OrtKernelContext, + _i: usize, + out: *mut *const ort::OrtValue, + ) -> ort::OrtStatusPtr { + unsafe { *out = std::ptr::null() }; + std::ptr::null_mut() + } + let mut api = fake_api(); + api.KernelContext_GetInput = Some(absent); + + let before = dispatch_probe::snapshot(); + let inputs = + unsafe { read_inputs(&api, std::ptr::null_mut()) }.expect("absent is not an error"); + let d = dispatch_probe::snapshot().since(&before); + + assert_eq!(inputs.len(), 1); + assert!(inputs[0].data_ptr.is_null()); + assert_eq!( + d.event(Event::OrtFfiCall), + 2, + "an absent optional input must cost the input count plus one probe, nothing more" + ); + } + + /// The probe must not perturb what it measures: with the counters compiled + /// in, `read_inputs` still returns exactly the same tensor description. + #[test] + fn instrumentation_does_not_change_results() { + let api = fake_api(); + let a = unsafe { read_inputs(&api, std::ptr::null_mut()) }.unwrap(); + dispatch_probe::reset(); + let b = unsafe { read_inputs(&api, std::ptr::null_mut()) }.unwrap(); + assert_eq!(a[0].shape, b[0].shape); + assert_eq!(a[0].strides, b[0].strides); + assert_eq!(a[0].dtype, b[0].dtype); + assert_eq!(a[0].data_ptr, b[0].data_ptr); + } +} diff --git a/crates/onnx-runtime-ep-plugin/src/lib.rs b/crates/onnx-runtime-ep-plugin/src/lib.rs index 2ddec2f467..8f5515aafd 100644 --- a/crates/onnx-runtime-ep-plugin/src/lib.rs +++ b/crates/onnx-runtime-ep-plugin/src/lib.rs @@ -29,6 +29,7 @@ use onnx_genai_ort_sys as ort; pub mod compute; pub mod device; +pub mod dispatch_probe; pub mod ep; pub mod factory; pub mod graph_reader; diff --git a/crates/onnx-runtime-ep-plugin/src/status.rs b/crates/onnx-runtime-ep-plugin/src/status.rs index 018448ef7d..7145429447 100644 --- a/crates/onnx-runtime-ep-plugin/src/status.rs +++ b/crates/onnx-runtime-ep-plugin/src/status.rs @@ -47,6 +47,8 @@ pub(crate) fn invalid_arg_status(message: &str) -> *mut ort::OrtStatus { /// Create an `OrtStatus` with an explicit error code. pub(crate) fn status_with_code(code: ort::OrtErrorCode, message: &str) -> *mut ort::OrtStatus { + crate::dispatch_probe::count(crate::dispatch_probe::Event::StatusCreated); + let _probe = crate::dispatch_probe::Phase::StatusCrossing.enter(); let c_msg = CString::new(message).unwrap_or_else(|_| CString::new("unknown error").unwrap()); let api = host_api(); if api.is_null() { From c16d906950d521f12e166a6f3283b8fa7c9afcfa Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Wed, 19 Aug 2026 03:55:15 +0000 Subject: [PATCH 2/7] Make the zero-cost assertion falsifiable MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `probe_is_compiled_out_in_production` compared `snapshot()` after some counting against `snapshot()` before it. Both are calls to the same constant function in a disabled build, so the assertion held no matter what the probe did — verified by mutation: a version whose disabled `count_n` recorded into a static and whose `snapshot` returned a non-zero `Counters` passed the test unchanged. Compare against `Counters::default()` instead, which is an absolute claim rather than a self-consistency one, and fails against that mutation. Also assert `!needs_drop::()`, so a `Drop` impl added to the production guard — the one way it could regain a cost while staying zero-sized — is caught. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../src/dispatch_probe.rs | 21 ++++++++++++++++--- 1 file changed, 18 insertions(+), 3 deletions(-) diff --git a/crates/onnx-runtime-ep-plugin/src/dispatch_probe.rs b/crates/onnx-runtime-ep-plugin/src/dispatch_probe.rs index 8d1f5b66ef..fe6522c7dd 100644 --- a/crates/onnx-runtime-ep-plugin/src/dispatch_probe.rs +++ b/crates/onnx-runtime-ep-plugin/src/dispatch_probe.rs @@ -534,17 +534,32 @@ mod tests { #[test] #[cfg(not(feature = "dispatch_probe"))] fn probe_is_compiled_out_in_production() { - assert_eq!(std::mem::size_of::(), 0); assert!(!compiled_in()); - let before = snapshot(); + assert_eq!( + std::mem::size_of::(), + 0, + "the disabled guard must carry no state" + ); + assert!( + !std::mem::needs_drop::(), + "the disabled guard must not run code when it leaves scope" + ); + let _g = Phase::KernelInvoke.enter(); count(Event::OrtFfiCall); ort_calls(1000); + + // Compared against `Counters::default()`, not against a snapshot taken + // beforehand. The relative form is vacuous: it compares two calls to a + // function the compiler is free to fold to a constant, so it passes + // even against a build whose "disabled" probe records on every call. + // This was verified by mutation — the relative version did not fail. assert_eq!( snapshot(), - before, + Counters::default(), "a disabled probe must not record anything" ); + assert_eq!(snapshot_global(), Counters::default()); } #[test] From 96cda80b5dfd2c45564dac0990a3ce6ca2bf2e11 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Wed, 19 Aug 2026 04:15:53 +0000 Subject: [PATCH 3/7] Make the FFI counter a real total, and be honest about the allocation one MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review of #1387 found the module's central claim was false. It documented `OrtFfiCall` as "how many ORT FFI calls a Run made" and called it the headline number for #1077, but only `read_inputs` and `allocate_output` were instrumented. On the happy path of a single-node Run, `device_mem_info`, `ort_input_mem_info`, `mem_info_is_device`, `alloc_scratch`, `host_pool::install` and `CreateStatus` all reach into ORT and none of them were counted. A probe that under-counts while presenting itself as exact is worse than no probe: it invites the reader to conclude the dispatch path is cheaper than it is, which is precisely the conclusion #1077 exists to prevent. Every leaf call through an `OrtApi` function pointer in this crate is now counted. Rather than assert that in a comment, `ffi_coverage` enforces it: it scans the source of each file that reaches into `OrtApi`, extracts the API members it names — ORT's C API is CamelCase and our own fields are snake_case, so `.CamelCase` isolates them cleanly — and fails if that set or the instrumentation count moves. Adding an uninstrumented call now fails with the member's name in the message. Verified by mutation: adding a stray `GetTensorShapeElementCount` reference failed with exactly that. A second test pins that the scan finds real members, so the guard cannot degrade into a heuristic that matches nothing and passes forever. `DispatchAlloc` gets the opposite treatment, because hand-placed counts genuinely cannot be exhaustive: `Vec::new()` does not allocate, `Vec::with_capacity(0)` does not allocate, and whether a `collect` allocates once or twice is a property of `size_hint` rather than of anything visible at the call site. It is now documented as a lower bound that is exact only for `read_inputs`, and callers who need the true whole-Run figure get `CountingAllocator`, a `GlobalAlloc` wrapper they can install in a test or benchmark. This crate deliberately does not install it — a library that defines a `#[global_allocator]` takes that choice away from every dependent. Also from review: - A node with no inputs was charged one allocation for a `Vec::with_capacity(0)` that never happens. Now conditional, and pinned by a test. - The phase timings are documented as what they are: non-exhaustive on the success path, and nesting on the error path, where `StatusCrossing` opens inside whichever phase was live because a guard closes at scope exit rather than at the early return. `DispatchLookup` is entered twice per node and reports the total. Summing them does not yield Run time in either direction, and a reader who assumed it did would be misled. - `probe_is_compiled_out_in_production` cannot prove `count` has no side effect — a disabled build has no storage to observe one — so it no longer implies it does. What it can prove, it does: the guard is a ZST and does not need dropping. Validation: clippy clean and `--lib` green in both feature configurations (253 without, 260 with), and 55/55 `plugin_ort_e2e` against real ORT. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- crates/onnx-runtime-ep-plugin/src/compute.rs | 9 + .../src/dispatch_probe.rs | 213 +++++++++++++++++- .../onnx-runtime-ep-plugin/src/host_pool.rs | 3 + .../onnx-runtime-ep-plugin/src/kernel_ctx.rs | 36 ++- crates/onnx-runtime-ep-plugin/src/status.rs | 6 +- 5 files changed, 259 insertions(+), 8 deletions(-) diff --git a/crates/onnx-runtime-ep-plugin/src/compute.rs b/crates/onnx-runtime-ep-plugin/src/compute.rs index 3e8e6350a9..3906614eec 100644 --- a/crates/onnx-runtime-ep-plugin/src/compute.rs +++ b/crates/onnx-runtime-ep-plugin/src/compute.rs @@ -1148,6 +1148,7 @@ impl Drop for ReconstructedMemInfo { return; } if let Some(release) = unsafe { (*api).ReleaseMemoryInfo } { + crate::dispatch_probe::ort_call(); unsafe { release(self.ptr.cast_mut()) }; } } @@ -1254,6 +1255,7 @@ unsafe fn device_mem_info( // on the device — scan every input for a genuinely device-resident one. if let Some(get_count) = api.KernelContext_GetInputCount { let mut count: usize = 0; + crate::dispatch_probe::ort_call(); let status = unsafe { get_count(ctx, &mut count) }; if status.is_null() { for i in 0..count { @@ -1295,11 +1297,13 @@ unsafe fn ort_input_mem_info( let get_input = api.KernelContext_GetInput?; let get_mem_info = api.GetTensorMemoryInfo?; let mut value: *const ort::OrtValue = std::ptr::null(); + crate::dispatch_probe::ort_call(); let status = unsafe { get_input(ctx, index, &mut value) }; if !status.is_null() || value.is_null() { return None; } let mut mem_info: *const ort::OrtMemoryInfo = std::ptr::null(); + crate::dispatch_probe::ort_call(); let status = unsafe { get_mem_info(value, &mut mem_info) }; if !status.is_null() || mem_info.is_null() { return None; @@ -1478,6 +1482,7 @@ unsafe fn operand_mem_info( return OperandMemInfo::Unavailable; }; let mut equal: std::os::raw::c_int = 0; + crate::dispatch_probe::ort_call(); let status = unsafe { compare(first, other, &mut equal) }; if !status.is_null() { // ORT allocated this status and handed us ownership; dropping the @@ -1515,6 +1520,7 @@ unsafe fn alloc_scratch( .KernelContext_GetScratchBuffer .ok_or("OrtApi.KernelContext_GetScratchBuffer is null")?; let mut out: *mut c_void = std::ptr::null_mut(); + crate::dispatch_probe::ort_call(); let status = unsafe { get_scratch(ctx, mem_info, bytes.max(1), &mut out) }; if !status.is_null() { return Err("KernelContext_GetScratchBuffer failed".into()); @@ -1581,6 +1587,7 @@ fn reconstruct_device_mem_info( let create = unsafe { (*api).CreateMemoryInfo_V2 }?; let name = std::ffi::CString::new(allocator_name).ok()?; let mut ptr: *mut ort::OrtMemoryInfo = std::ptr::null_mut(); + crate::dispatch_probe::ort_call(); let status = unsafe { create( name.as_ptr(), @@ -1595,6 +1602,7 @@ fn reconstruct_device_mem_info( }; if !status.is_null() { if let Some(release) = unsafe { (*api).ReleaseStatus } { + crate::dispatch_probe::ort_call(); unsafe { release(status) }; } return None; @@ -1616,6 +1624,7 @@ unsafe fn mem_info_device_type( ) -> Option { let f = api.MemoryInfoGetDeviceType?; let mut out: ort::OrtMemoryInfoDeviceType = ort::OrtMemoryInfoDeviceType_CPU; + crate::dispatch_probe::ort_call(); unsafe { f(mem_info, &mut out) }; Some(out) } diff --git a/crates/onnx-runtime-ep-plugin/src/dispatch_probe.rs b/crates/onnx-runtime-ep-plugin/src/dispatch_probe.rs index fe6522c7dd..3a3a923641 100644 --- a/crates/onnx-runtime-ep-plugin/src/dispatch_probe.rs +++ b/crates/onnx-runtime-ep-plugin/src/dispatch_probe.rs @@ -21,20 +21,46 @@ //! machine, so a test can assert them: *"a single-node static-shape `Run` //! makes exactly N FFI calls"*. That turns a performance property into a //! correctness test, which is the only kind that survives contact with CI. +//! +//! [`Event::OrtFfiCall`] is a genuine total: the `ffi_coverage` tests scan +//! this crate's source and fail if any file names an ORT API member without +//! matching instrumentation, so a new call cannot be added silently. +//! [`Event::DispatchAlloc`], counted by hand, is a **lower bound** — see its +//! own documentation, and use [`CountingAllocator`] when the exact figure +//! matters. //! * **Timings** — nanoseconds per phase. Useful for finding the next thing to //! fix, useless as an assertion. They are gated behind an environment //! variable *on top of* the build feature so that enabling counters for a //! test does not drag two `Instant::now()` calls into every phase. //! +//! # Phases are not a partition +//! +//! Summing `phase_ns` does not give the time spent in `Compute`, in either +//! direction. On the success path the phases are *non-exhaustive*: work between +//! them — output allocation bookkeeping, view construction — is inside no +//! phase, so the sum is less than wall time. On an error path they *nest*: +//! `StatusCrossing` opens inside whichever phase was live when the failure was +//! detected, because a guard closes at scope exit rather than at the early +//! `return`, so the sum can exceed wall time. `DispatchLookup` is likewise +//! entered twice per node, once for shape inference and once for workspace +//! preparation, and reports their total. +//! +//! Read a phase as "time attributable to this segment", not as a slice of a +//! pie chart. +//! //! # Cost when disabled //! //! Everything here is behind the `dispatch_probe` feature, which is **not** a //! default feature. With it off, [`Phase::enter`] returns a zero-sized guard //! whose `Drop` is empty and [`count`] is an empty `#[inline(always)]` function //! — there is nothing left for the optimiser to remove. -//! `probe_is_compiled_out_in_production` pins that the guard is a ZST, and -//! `dispatch_probe_is_not_a_default_feature` pins that a plain build does not -//! get it. +//! `probe_is_compiled_out_in_production` pins that the guard is a ZST with no +//! `Drop`, and `dispatch_probe_is_not_a_default_feature` pins that a plain +//! build does not get it. Note the limit of what a test can show here: it can +//! prove the guard carries no state and runs nothing on scope exit, but it +//! cannot prove [`count`] has no side effect, because a disabled build has no +//! storage for it to observe. That claim rests on the function body being +//! empty, which is visible in the source directly below. /// A segment of the dispatch path. /// @@ -105,10 +131,23 @@ pub enum Event { /// One call through a function pointer in ORT's `OrtApi` table. This is the /// headline number for this issue: each one is a cross-library indirect /// call that ORT's own CPU EP does not have to make. + /// + /// This is a **total**, not a sample. The `ffi_coverage` tests scan every + /// source file that reaches into `OrtApi` and fail if the set of members it + /// names changes without the instrumentation changing with it, so a call + /// added later cannot go uncounted without someone noticing. OrtFfiCall = 0, /// One heap allocation the dispatch path asked for, at a site we control. - /// Counted where we allocate deliberately (shape/stride/metadata vectors), - /// not by hooking the global allocator. + /// + /// Unlike [`Event::OrtFfiCall`] this is a **lower bound**, and is honest + /// about it. Hand-placed counts cannot be exhaustive: `Vec::new()` does not + /// allocate, `Vec::with_capacity(0)` does not allocate, and whether a + /// `collect` allocates once or twice depends on `size_hint` rather than on + /// anything visible at the call site. It is exhaustive for `read_inputs`, + /// where the sites are few and each is pinned by a test. + /// + /// When the exact whole-`Run` figure is what you need, install + /// [`CountingAllocator`] and let the allocator count. DispatchAlloc = 1, /// One `OrtStatus` constructed and handed back across the ABI. Should be /// zero on a successful `Run`. @@ -454,6 +493,68 @@ pub use imp::{ PhaseGuard, compiled_in, count, count_n, reset, snapshot, snapshot_global, timing_enabled, }; +/// A `GlobalAlloc` that tallies every real heap allocation into +/// [`Event::DispatchAlloc`], for callers that need an exhaustive count rather +/// than a count of the sites someone remembered to instrument. +/// +/// Hand-placed `count(DispatchAlloc)` calls cannot be exhaustive in general, +/// and pretending otherwise is the failure mode this type exists to avoid: +/// `Vec::new()` does not allocate, `Vec::with_capacity(0)` does not allocate, +/// and whether a given `collect` allocates once or twice is a property of +/// `size_hint`, not of the source text. Where the exact number matters, install +/// this and let the allocator answer. +/// +/// It is deliberately *not* installed by this crate. A `#[global_allocator]` +/// may be defined only once per binary, so a library that installs one takes +/// that choice away from every dependent. Tests and benchmarks opt in: +/// +/// ```ignore +/// #[global_allocator] +/// static A: dispatch_probe::CountingAllocator = +/// dispatch_probe::CountingAllocator::new(std::alloc::System); +/// ``` +/// +/// Counts reallocations and the allocating half of a resize, and ignores frees, +/// because the question being asked is "how much allocator traffic does one +/// `Run` cause", not "how much memory is live". +/// +/// Without the `dispatch_probe` feature this forwards to the inner allocator +/// and records nothing. +pub struct CountingAllocator { + inner: A, +} + +impl CountingAllocator { + /// Wrap `inner`, counting the allocations that pass through it. + pub const fn new(inner: A) -> Self { + Self { inner } + } +} + +// SAFETY: every method forwards to `self.inner` with its arguments unchanged +// and returns its result unchanged, so the allocator contract is exactly the +// inner allocator's. The added work only touches counters. +unsafe impl std::alloc::GlobalAlloc for CountingAllocator { + unsafe fn alloc(&self, layout: std::alloc::Layout) -> *mut u8 { + count(Event::DispatchAlloc); + unsafe { self.inner.alloc(layout) } + } + + unsafe fn dealloc(&self, ptr: *mut u8, layout: std::alloc::Layout) { + unsafe { self.inner.dealloc(ptr, layout) } + } + + unsafe fn alloc_zeroed(&self, layout: std::alloc::Layout) -> *mut u8 { + count(Event::DispatchAlloc); + unsafe { self.inner.alloc_zeroed(layout) } + } + + unsafe fn realloc(&self, ptr: *mut u8, layout: std::alloc::Layout, new_size: usize) -> *mut u8 { + count(Event::DispatchAlloc); + unsafe { self.inner.realloc(ptr, layout, new_size) } + } +} + /// Tally one ORT FFI call. Shorthand for the single most common count. #[inline(always)] pub fn ort_call() { @@ -557,7 +658,7 @@ mod tests { assert_eq!( snapshot(), Counters::default(), - "a disabled probe must not record anything" + "a disabled probe must not report anything" ); assert_eq!(snapshot_global(), Counters::default()); } @@ -669,3 +770,103 @@ mod tests { ); } } + +/// Guards the claim that `Event::OrtFfiCall` counts *every* call this crate +/// makes into ORT, rather than the ones somebody remembered to instrument. +/// +/// An under-counting probe presented as exact is worse than no probe: it +/// invites a reader to conclude the dispatch path is cheaper than it is. But +/// completeness is not something a unit test can observe from the inside — +/// there is no hook that fires when a function pointer is called. So this +/// checks the next best thing, statically: every ORT entry point the crate +/// names, and every place it counts a call. +/// +/// The mechanism is a source scan. ORT's C API members are `CamelCase`, while +/// Rust fields and methods here are `snake_case`, so `.CamelCase` picks out +/// API member accesses and essentially nothing else. Pinning both that set and +/// the instrumentation count means a newly added FFI call cannot land silently: +/// it changes one of these numbers and the author has to come here, look at +/// this comment, and decide. +#[cfg(test)] +mod ffi_coverage { + /// Every source file that reaches into `OrtApi`, with the ORT members it + /// names and the number of `ort_call()` sites it contains. + /// + /// A count may legitimately differ from the number of distinct members: a + /// member can be called from more than one place, and an extracted + /// function pointer can be invoked through a local binding. What must + /// never happen is a member appearing here with no instrumentation to + /// account for it. + const EXPECTED: &[(&str, &str, usize, usize)] = &[ + ("compute.rs", include_str!("compute.rs"), 9, 9), + ("kernel_ctx.rs", include_str!("kernel_ctx.rs"), 11, 11), + ("status.rs", include_str!("status.rs"), 1, 1), + ("host_pool.rs", include_str!("host_pool.rs"), 2, 2), + ]; + + /// Source ahead of the file's `#[cfg(test)]` block — test scaffolding + /// names ORT members too (the fake `OrtApi` in `kernel_ctx`, the fake + /// thread pools here), and instrumenting a fake would prove nothing. + fn production_source(src: &str) -> &str { + match src.find("\n#[cfg(test)]") { + Some(i) => &src[..i], + None => src, + } + } + + fn ort_members(src: &str) -> Vec { + let b = src.as_bytes(); + let mut out: Vec = Vec::new(); + for (i, w) in b.windows(2).enumerate() { + if w[0] == b'.' && w[1].is_ascii_uppercase() { + // Skip `::Variant` — a path, not a member access. + if i > 0 && b[i - 1] == b':' { + continue; + } + let rest = &src[i + 1..]; + let end = rest + .find(|c: char| !c.is_alphanumeric() && c != '_') + .unwrap_or(rest.len()); + out.push(rest[..end].to_string()); + } + } + out.sort(); + out.dedup(); + out + } + + #[test] + fn every_ort_entry_point_is_accounted_for() { + for (name, src, want_members, want_calls) in EXPECTED { + let prod = production_source(src); + let members = ort_members(prod); + let calls = prod.matches("dispatch_probe::ort_call()").count(); + assert_eq!( + members.len(), + *want_members, + "{name} now names {} ORT API members ({members:?}), not {want_members}. \ + If you added an FFI call, add an `ort_call()` beside it and update \ + this table; if you removed one, just update the table.", + members.len() + ); + assert_eq!( + calls, *want_calls, + "{name} has {calls} `ort_call()` sites, not {want_calls}. Every call \ + into ORT must be counted, or `Event::OrtFfiCall` stops being a total." + ); + } + } + + /// The scan has to actually find things — a heuristic that silently matched + /// nothing would let the assertions above pass vacuously forever. + #[test] + fn the_source_scan_is_not_vacuous() { + let m = ort_members(production_source(include_str!("kernel_ctx.rs"))); + assert!(m.contains(&"KernelContext_GetInput".to_string()), "{m:?}"); + assert!(m.contains(&"GetTensorData".to_string()), "{m:?}"); + assert!( + !m.iter().any(|s| s.contains("::")), + "path segments must not be mistaken for member accesses: {m:?}" + ); + } +} diff --git a/crates/onnx-runtime-ep-plugin/src/host_pool.rs b/crates/onnx-runtime-ep-plugin/src/host_pool.rs index c6fcd26e18..b2cff5a89b 100644 --- a/crates/onnx-runtime-ep-plugin/src/host_pool.rs +++ b/crates/onnx-runtime-ep-plugin/src/host_pool.rs @@ -146,6 +146,7 @@ unsafe fn ort_parallel_for(host: *mut c_void, total: usize, body: &(dyn Fn(usize // wide its pool is, so we hand it enough pieces to fill any pool and let // it decide how many to run at once. let status = unsafe { + crate::dispatch_probe::ort_call(); (pool.parallel_for)( pool.ctx, Some(run_index), @@ -168,6 +169,7 @@ unsafe fn ort_parallel_for(host: *mut c_void, total: usize, body: &(dyn Fn(usize // output tensor keeps whatever was in the buffer, so fall back to // running them here rather than silently returning short. if let Some(release) = pool.release_status { + crate::dispatch_probe::ort_call(); unsafe { release(status) }; } for index in 0..total { @@ -258,6 +260,7 @@ pub unsafe fn install( let (Some(parallel_for), false) = (api.KernelContext_ParallelFor, ctx.is_null()) else { return Guard::inert(); }; + crate::dispatch_probe::count(crate::dispatch_probe::Event::DispatchAlloc); let pool = Box::new(HostPool { parallel_for, release_status: api.ReleaseStatus, diff --git a/crates/onnx-runtime-ep-plugin/src/kernel_ctx.rs b/crates/onnx-runtime-ep-plugin/src/kernel_ctx.rs index 41022004d4..5385d85a13 100644 --- a/crates/onnx-runtime-ep-plugin/src/kernel_ctx.rs +++ b/crates/onnx-runtime-ep-plugin/src/kernel_ctx.rs @@ -189,7 +189,11 @@ pub(crate) unsafe fn read_inputs( return Err("KernelContext_GetInputCount failed".into()); } - crate::dispatch_probe::count(crate::dispatch_probe::Event::DispatchAlloc); + // `Vec::with_capacity(0)` does not allocate, so a node with no inputs must + // not be charged for this one. + if input_count > 0 { + crate::dispatch_probe::count(crate::dispatch_probe::Event::DispatchAlloc); + } let mut inputs = Vec::with_capacity(input_count); for i in 0..input_count { let mut value: *const ort::OrtValue = std::ptr::null(); @@ -780,6 +784,36 @@ mod dispatch_cost { ); } + /// A node with no inputs must not be charged for the `Vec`: + /// `Vec::with_capacity(0)` does not call the allocator. Counting it anyway + /// would make every constant-producing node look one allocation worse than + /// it is. + #[test] + fn a_node_with_no_inputs_allocates_nothing() { + unsafe extern "C" fn none( + _c: *const ort::OrtKernelContext, + out: *mut usize, + ) -> ort::OrtStatusPtr { + unsafe { *out = 0 }; + std::ptr::null_mut() + } + let mut api = fake_api(); + api.KernelContext_GetInputCount = Some(none); + + let before = dispatch_probe::snapshot(); + let inputs = + unsafe { read_inputs(&api, std::ptr::null_mut()) }.expect("zero inputs is fine"); + let d = dispatch_probe::snapshot().since(&before); + + assert!(inputs.is_empty()); + assert_eq!(d.event(Event::OrtFfiCall), 1); + assert_eq!( + d.event(Event::DispatchAlloc), + 0, + "an empty Vec::with_capacity must not be counted as an allocation" + ); + } + /// The probe must not perturb what it measures: with the counters compiled /// in, `read_inputs` still returns exactly the same tensor description. #[test] diff --git a/crates/onnx-runtime-ep-plugin/src/status.rs b/crates/onnx-runtime-ep-plugin/src/status.rs index 7145429447..3d159f711c 100644 --- a/crates/onnx-runtime-ep-plugin/src/status.rs +++ b/crates/onnx-runtime-ep-plugin/src/status.rs @@ -49,6 +49,7 @@ pub(crate) fn invalid_arg_status(message: &str) -> *mut ort::OrtStatus { pub(crate) fn status_with_code(code: ort::OrtErrorCode, message: &str) -> *mut ort::OrtStatus { crate::dispatch_probe::count(crate::dispatch_probe::Event::StatusCreated); let _probe = crate::dispatch_probe::Phase::StatusCrossing.enter(); + crate::dispatch_probe::count(crate::dispatch_probe::Event::DispatchAlloc); let c_msg = CString::new(message).unwrap_or_else(|_| CString::new("unknown error").unwrap()); let api = host_api(); if api.is_null() { @@ -57,7 +58,10 @@ pub(crate) fn status_with_code(code: ort::OrtErrorCode, message: &str) -> *mut o // SAFETY: api was set during CreateEpFactories and is process-lifetime valid. unsafe { match (*api).CreateStatus { - Some(create_status) => create_status(code, c_msg.as_ptr()), + Some(create_status) => { + crate::dispatch_probe::ort_call(); + create_status(code, c_msg.as_ptr()) + } None => ptr::null_mut(), } } From a433e5a30ea6e899cf3325449f842a7932fab991 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 18 Aug 2026 22:09:00 -0700 Subject: [PATCH 4/7] Stop building debug strings that production throws away (#1392) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Stacked on #1387 (base `squad/resch-dispatch-probe`). The instrumentation there exists to make changes like this one provable; this is the first of them. ## Why Two heap allocations sat on the per-`Run` dispatch path purely to serve diagnostics that are switched off. **1. `staging_log(&format!(…))`.** `staging_log` returns immediately when the transfer trace is disabled — but the argument is evaluated first. Every one of these builds a `String` and formats each interpolated value on the way to a function that drops it on the floor. One of them is in `compute_execute`, before any real work happens: ```rust staging_log(&format!( "[plugin/staging #982] compute_execute enter: entries={} routing={} device_staging={}", exported.entries.len(), exported.routing.is_some(), exported.device_staging.is_some() )); ``` So every dispatch in every production run has been paying for a debug line nobody asked for. **2. `format!("input {i}")` in `read_inputs`.** Built for every input on every `Run`, to label an error that is almost never raised. ## What A `staging_log!` macro that checks `staging_trace_enabled()` *before* formatting, so a disabled trace costs one relaxed `OnceLock` load and nothing else. Thirteen call sites convert. `validate_dims` now takes `impl Display` instead of `&str`, so the caller passes `format_args!("input {i}")` — which borrows its arguments and formats only if a message is actually produced. Same treatment for `transfer.rs`'s `CopyTensors[{i}]`. ## Evidence The pinned counts in `dispatch_cost` (added by #1387) move: | | before | after | |---|---|---| | allocations, 1 input | 5 | **4** | | allocations, 3 inputs | 1 + 3×4 | **1 + 3×3** | That is the instrumentation earning its keep. The improvement is a number changing in a test — reproducible on any machine, by anyone, at any time — rather than a claim resting on a benchmark. That matters here in particular: the box is currently shared with other work and sitting at ~74% idle against my ≥93% gate, so a wall-clock A/B taken right now would be noise. I will post interleaved p50/p90 once the machine quiesces, but the correctness of the change does not depend on it. ## Correctness - **The trace still works.** Ran the e2e suite with `ONNX_GENAI_PLUGIN_TRANSFER_TRACE=1` and confirmed the `compute_execute enter` line still appears. A macro that silently disabled the trace would otherwise be indistinguishable from one that made it cheap. - **No side-effecting expression moved inside a now-conditional macro.** I checked every converted site; the only mutation in the neighbourhood (`to_stage.push(p)`) is outside the macro call. - **Emitted text is unchanged** — the format strings and arguments are copied verbatim, only the gate moved. - clippy clean and `--lib` green in both feature configurations (253 / 260), and **55/55 `plugin_ort_e2e`** against real ORT. Relates to #1077. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- crates/onnx-runtime-ep-plugin/src/compute.rs | 66 +++++++++++-------- .../onnx-runtime-ep-plugin/src/kernel_ctx.rs | 26 ++++---- crates/onnx-runtime-ep-plugin/src/transfer.rs | 2 +- 3 files changed, 54 insertions(+), 40 deletions(-) diff --git a/crates/onnx-runtime-ep-plugin/src/compute.rs b/crates/onnx-runtime-ep-plugin/src/compute.rs index 3906614eec..2621f66aef 100644 --- a/crates/onnx-runtime-ep-plugin/src/compute.rs +++ b/crates/onnx-runtime-ep-plugin/src/compute.rs @@ -1544,6 +1544,24 @@ fn staging_trace_enabled() -> bool { *ENABLED.get_or_init(|| std::env::var("ONNX_GENAI_PLUGIN_TRANSFER_TRACE").as_deref() == Ok("1")) } +/// Emit one staging-trace line, formatting the message only if the trace is on. +/// +/// `staging_log` already returns immediately when the trace is disabled, but +/// the argument is evaluated first: `staging_log!(…)` builds a +/// `String` and formats every interpolated value on the way to a function that +/// throws it away. One of those sits on the per-`Run` dispatch path, so it was +/// paid by every dispatch in production to serve a debugging aid that is off. +/// +/// This checks the gate first, so a disabled trace costs one relaxed load of a +/// `OnceLock` and nothing else. +macro_rules! staging_log { + ($($arg:tt)*) => { + if staging_trace_enabled() { + staging_log(&format!($($arg)*)); + } + }; +} + /// Emit one staging-trace line. Goes to stderr (buffered by pipes, so it can be /// lost when the process hangs) *and*, when `ONNX_GENAI_PLUGIN_TRANSFER_TRACE_FILE` /// is set, is appended to that file with an immediate flush — the file path is @@ -1694,54 +1712,50 @@ unsafe fn stage_host_boundary_inputs( ) -> Result<(), String> { // First pass: which operands are host-resident data operands needing upload? let mut to_stage: Vec = Vec::new(); - staging_log(&format!( + staging_log!( "[plugin/staging #982] {label}: enter operands={} host_operands={:?}", ort_indices.len(), host_operands - )); + ); for (p, slot) in ort_indices.iter().enumerate() { let Some(ort_idx) = *slot else { - staging_log(&format!( - "[plugin/staging #982] {label}: op[{p}] skip (not ORT-bound)" - )); + staging_log!("[plugin/staging #982] {label}: op[{p}] skip (not ORT-bound)"); continue; }; if host_operands.contains(&p) { - staging_log(&format!( + staging_log!( "[plugin/staging #982] {label}: op[{p}] ort_idx={ort_idx} skip (host control operand)" - )); + ); continue; } let view = &kernel_inputs[p]; if view.is_absent() || view.data.0.is_null() { - staging_log(&format!( + staging_log!( "[plugin/staging #982] {label}: op[{p}] ort_idx={ort_idx} skip (absent/null)" - )); + ); continue; } let numel: usize = view.shape.iter().product(); if numel == 0 || view.dtype.byte_size() == 0 { - staging_log(&format!( - "[plugin/staging #982] {label}: op[{p}] ort_idx={ort_idx} skip (empty)" - )); + staging_log!("[plugin/staging #982] {label}: op[{p}] ort_idx={ort_idx} skip (empty)"); continue; } let Some(mem_info) = (unsafe { ort_input_mem_info(api, ctx, ort_idx) }) else { - staging_log(&format!( + staging_log!( "[plugin/staging #982] {label}: op[{p}] ort_idx={ort_idx} skip (no mem info)" - )); + ); continue; }; let dev_type = unsafe { mem_info_device_type(api, mem_info) }; if unsafe { mem_info_is_device(api, mem_info) } { - staging_log(&format!( + staging_log!( "[plugin/staging #982] {label}: op[{p}] ort_idx={ort_idx} skip (already device, dev_type={dev_type:?})" - )); + ); continue; // already on device — nothing to do } - staging_log(&format!( + staging_log!( "[plugin/staging #982] {label}: op[{p}] ort_idx={ort_idx} HOST (dev_type={dev_type:?}) → will stage" - )); + ); to_stage.push(p); } @@ -1772,10 +1786,10 @@ unsafe fn stage_host_boundary_inputs( && let Some(recon) = staging.recon_mem_info.as_ref() { device_mi = recon.ptr; - staging_log(&format!( + staging_log!( "[plugin/staging #982] {label}: no device OrtValue found; \ falling back to reconstructed EP memory info recon={device_mi:?}" - )); + ); } if device_mi.is_null() { return Err(format!( @@ -1794,19 +1808,17 @@ unsafe fn stage_host_boundary_inputs( let byte_len = numel * view.dtype.byte_size(); let dst = unsafe { alloc_scratch(api, ctx, device_mi, byte_len) } .map_err(|e| format!("{label}: staging scratch alloc failed: {e}"))?; - staging_log(&format!( + staging_log!( "[plugin/staging #982] {label}: staging op[{p}] byte_len={byte_len} \ mem_dev={mem_dev:?} dst={dst:?} — issuing host→device upload" - )); + ); // SAFETY: the source is a host tensor of `byte_len` contiguous bytes // (ORT-provided, contiguous strides); `dst` is device scratch of the // same size. let src = unsafe { std::slice::from_raw_parts(view.data.0.cast::(), byte_len) }; unsafe { staging.copier.copy_host_to_device(src, dst) } .map_err(|e| format!("{label}: host→device upload failed: {e}"))?; - staging_log(&format!( - "[plugin/staging #982] {label}: op[{p}] host→device upload complete" - )); + staging_log!("[plugin/staging #982] {label}: op[{p}] host→device upload complete"); kernel_inputs[p].data = DevicePtr(dst.cast_const()); } @@ -2184,12 +2196,12 @@ unsafe extern "C" fn compute_execute( return fail_status("Compute: no kernels compiled for this subgraph"); } - staging_log(&format!( + staging_log!( "[plugin/staging #982] compute_execute enter: entries={} routing={} device_staging={}", exported.entries.len(), exported.routing.is_some(), exported.device_staging.is_some() - )); + ); let api = crate::status::host_api(); if api.is_null() { diff --git a/crates/onnx-runtime-ep-plugin/src/kernel_ctx.rs b/crates/onnx-runtime-ep-plugin/src/kernel_ctx.rs index 5385d85a13..de35f76036 100644 --- a/crates/onnx-runtime-ep-plugin/src/kernel_ctx.rs +++ b/crates/onnx-runtime-ep-plugin/src/kernel_ctx.rs @@ -72,7 +72,7 @@ pub const CPU_EP_SUPPORTED_DTYPES: &[DataType] = &[ pub(crate) fn validate_dims( dims: &[i64], dtype: DataType, - context: &str, + context: impl std::fmt::Display, ) -> Result<(Vec, usize, usize), String> { let mut shape: Vec = Vec::with_capacity(dims.len()); for (dim_idx, &d) in dims.iter().enumerate() { @@ -252,9 +252,10 @@ pub(crate) unsafe fn read_inputs( unsafe { release_type_shape(type_shape) }; // Validate ORT-supplied dims: reject negatives and detect overflow. - // Three allocations: the eager `format!` label, `shape`, and `strides`. - crate::dispatch_probe::count_n(crate::dispatch_probe::Event::DispatchAlloc, 3); - let (shape, _, _) = validate_dims(&dims, dtype, &format!("input {i}"))?; + // Two allocations: `shape` and `strides`. The label is `format_args!`, + // which borrows its arguments and formats only if an error is raised. + crate::dispatch_probe::count_n(crate::dispatch_probe::Event::DispatchAlloc, 2); + let (shape, _, _) = validate_dims(&dims, dtype, format_args!("input {i}"))?; let strides = onnx_runtime_ir::compute_contiguous_strides(&shape); // Data pointer. @@ -711,20 +712,21 @@ mod dispatch_cost { ); } - /// One input also costs five heap allocations: the `Vec` once - /// per `Run`, then per input the `dims` scratch, the eager - /// `format!("input {i}")` label built even when validation succeeds, - /// `shape`, and `strides`. The label is pure waste on the success path and - /// is the first thing to remove. + /// One input costs four heap allocations: the `Vec` once per + /// `Run`, then per input the `dims` scratch, `shape`, and `strides`. + /// + /// It was five until the error label stopped being built eagerly. That is + /// the counter earning its keep: the improvement is visible as a number + /// changing in a test, not as a claim in a commit message. #[test] - fn reading_one_input_costs_exactly_five_allocations() { + fn reading_one_input_costs_exactly_four_allocations() { let api = fake_api(); let before = dispatch_probe::snapshot(); let _ = unsafe { read_inputs(&api, std::ptr::null_mut()) }.expect("fake api is total"); let d = dispatch_probe::snapshot().since(&before); assert_eq!( d.event(Event::DispatchAlloc), - 1 + 4, + 1 + 3, "allocations on the per-Run input path changed" ); } @@ -751,7 +753,7 @@ mod dispatch_cost { assert_eq!(inputs.len(), 3); assert_eq!(d.event(Event::OrtFfiCall), 1 + 3 * 7); - assert_eq!(d.event(Event::DispatchAlloc), 1 + 3 * 4); + assert_eq!(d.event(Event::DispatchAlloc), 1 + 3 * 3); } /// An absent optional input short-circuits: ORT hands back a null value and diff --git a/crates/onnx-runtime-ep-plugin/src/transfer.rs b/crates/onnx-runtime-ep-plugin/src/transfer.rs index e0387e3303..949958cfdd 100644 --- a/crates/onnx-runtime-ep-plugin/src/transfer.rs +++ b/crates/onnx-runtime-ep-plugin/src/transfer.rs @@ -770,7 +770,7 @@ unsafe extern "C" fn transfer_full_copy_tensors( let (_, _, byte_len) = match crate::kernel_ctx::validate_dims( &dims, dtype, - &format!("CopyTensors[{i}]"), + format_args!("CopyTensors[{i}]"), ) { Ok(v) => v, Err(e) => return fail_status(&e), From f1faf9b23e2177acff6d51919a77196d56a66841 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 18 Aug 2026 22:10:30 -0700 Subject: [PATCH 5/7] Dispatch-overhead stack: benchmark grid, 7 -> 4 FFI calls per Run, routed-loop fixes (2.31 -> 1.71 at depth 100) (#1397) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit > **Note on scope.** Enabling auto-merge across the stack caused each stacked PR to merge into its *parent feature branch* (those branches are unprotected; only `main` requires checks), and GitHub retargeted the children. No `main` protection was bypassed and no work was lost, but this PR now carries the whole dispatch-overhead stack below #1387, not just the memory-info change: > > - #1394 — the #1077 benchmark grid > - **this PR** — `device_mem_info`, 7 → 4 FFI calls per `Run` > - #1401 — routed-loop instrumentation, the O(depth²) retirement scan, lazy absent strides, and the 8-element dispatch-isolating cases > > Each was independently reviewed before it merged; the sections below cover the memory-info change, and #1394/#1401 retain their own descriptions and review records. #1387 (probe + lazy formatting) remains a separate PR to `main`. ## Stop asking ORT three times where the inputs live Part of #1077. Stacked on #1394 (which measured the cost this removes). ### The finding Every `Compute` call resolves where a routed subgraph's intermediate tensors should be allocated. On a host EP — and the CPU plugin is always a host EP — that resolution cost **7 FFI calls** for a single-input node to arrive at the answer "host memory". Three of them were asking ORT for something already in hand. `device_mem_info` scans the inputs looking for a device-resident one: 1. per input: `GetValue` → `GetTensorMemoryInfo` → `mem_info_is_device` 2. on a host EP it never finds one, so the scan runs to completion 3. it then falls through to "use input 0's memory info" — and **fetches input 0 again**, two more calls for a value it had already looked at and discarded 4. the caller then asked `mem_info_is_device` on the result — a **third** query for a fact the function had just established internally and thrown away ### The change The scan remembers input 0's memory info as it walks, and returns `(mem_info, is_device)` so the caller does not have to re-derive the second half. The fallback fetch stays, for the one case that genuinely needs it: no input count available, so the scan never ran and nothing was remembered. ### Why this is behaviour-identical - The scan visits the same inputs, in the same order, and returns on the same condition. - The remembered value is exactly what the re-fetch would have produced: both are input 0's `OrtMemoryInfo`, and it is immutable for the lifetime of the value. - The `recon_mem_info` branch reports `is_device = true` without asking. That branch is only reachable when the scan found a device-resident input, so it is device-resident by construction — which is precisely what the removed call was confirming. ### Evidence `mem_info_cost` (new, in `compute.rs`) drives `device_mem_info` through a hand-built `OrtApi` and pins the call count. | scenario | before | after | |---|---|---| | host EP, 1 input | **7** | **4** | | device-resident input (short-circuit) | 3 | 3 | | no input count (fallback path) | 2 | 2 | | zero inputs | 0 | 0 | The "before" column is not a recollection — it is the same test run against the previous implementation. Three FFI calls per `Run`, off the fixed per-`Run` cost that #1394 measures at **~0.9 µs worse than ORT's**. ### Why counters and not a stopwatch This machine has not been quiet enough all session for a trustworthy wall-clock A/B (`mpstat` shows cores at 75–78% idle against a ≥93% gate). An FFI-call count is deterministic and load-independent, so it is reproducible now and will still be reproducible on someone else's machine. That is the case for #1387 existing. ### Also pinned The three non-host paths, so a future change cannot quietly break them while the host count stays pretty: device-resident short-circuit, missing input count, and a node with no inputs (reports no memory info rather than inventing one). ### Validation - `cargo test -p onnx-runtime-ep-plugin --lib` — 253 pass (debug; `--release` silently disables `debug_assert!`) - same with `--features dispatch_probe` — 264 pass - `plugin_ort_e2e` — **55/55**, including `every_assigned_node_is_also_executed_by_this_ep` - clippy clean in both feature configurations No ORT CPU fallback. No MLAS. --- ## Independent review (Opus 4.8, adversarial) — **APPROVE WITH NITS**, all findings addressed Commissioned specifically to attack "behaviour is unchanged in every case". It held for every EP that ships, with one crack worth fixing. **Finding 1 (MINOR) — the recon branch's `is_device = true` was convention, not construction.** Upheld, and fixed properly. My comment claimed the reconstruction is "device-resident by construction". The reviewer traced it: `recon_mem_info` is built from the EP's reported `device_type`, and `device_staging` is attached whenever an EP supplies a `HostToDeviceCopier` — and **nothing in that trait requires a copier-providing EP to be non-CPU**. Every one today is (CUDA→GPU, QNN→NPU), which is exactly what made the hardcode look safe. A CPU-typed EP that grew a copier would build a host-typed memory info and the hardcode would hand it to the device scratch path — host pointers treated as device pointers, the failure this code exists to prevent, inverted. The old code had no such exposure because it asked ORT. Rather than paper over it with a `debug_assert`, the invariant is now structural. `mem_info_is_device` is exactly `device_type != CPU`, and the reconstruction is created by passing that same `device_type` to `CreateMemoryInfo_V2` — so recording it at construction reproduces the old query **in every case, including the one that does not exist yet**, and still costs no FFI call. **Finding 2 (MINOR) — the one branch whose reporting changed was untested.** Correct: all four original tests passed `staging = None` and never reached it. Three tests added — CPU-typed reconstruction is not reported as a device, device-typed one is (4 calls, not 5), and a real device input still wins over the fallback. Falsified: restoring the hardcoded `true` fails **exactly one** test, the one naming the case; the other six correctly stay silent. **Finding 3 (NIT) — test 2's `== 4` does not distinguish new from old.** Accurate. The device-input path cost 4 before as well; the real guard in that test is the tuple value, not the count. Left as-is (it pins the device path against future regressions), noted here so nobody reads it as evidence for this PR. ### Claims the reviewer tried to break and could not These are now load-bearing rather than merely unchallenged: - **Pointer stability** — the remembered input-0 memory info cannot differ from what the re-fetch produced: same `OrtValue`, straight-line code with no intervening mutation, and the pointer's lifetime window is identical in both versions. - **The fallback path does not hardcode anything** — it still queries, so a device-resident input on that path still reports `true`, exactly as before. - **Zero-input nodes** — identical (`None`, not an invented memory info). - **Caller equivalence** — `Option<(*const _, bool)>` is `Copy`; no move, no shadowing, no ordering change. - **`mem::zeroed::()`** — sound; all fields are `Option`, zero is the null-fn niche, and the sentinel pointers are never dereferenced. - **The `ffi_coverage` tripwire stays honest** — 9 members / 9 `ort_call()` sites, unchanged and correct: this PR removes *dynamic* calls, not *static* sites. - **Concurrency** — entirely stack-local; no new shared mutable state. - **No vacuous test** — the reviewer was briefed that I have twice shipped a test that passed against the broken version, and specifically hunted a third. None found. --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../tests/plugin_ort_e2e.rs | 200 ++++++- crates/onnx-runtime-ep-plugin/src/compute.rs | 565 ++++++++++++++++-- 2 files changed, 729 insertions(+), 36 deletions(-) diff --git a/crates/onnx-runtime-ep-cpu-plugin/tests/plugin_ort_e2e.rs b/crates/onnx-runtime-ep-cpu-plugin/tests/plugin_ort_e2e.rs index 0f321fb22a..9b8a3556b9 100644 --- a/crates/onnx-runtime-ep-cpu-plugin/tests/plugin_ort_e2e.rs +++ b/crates/onnx-runtime-ep-cpu-plugin/tests/plugin_ort_e2e.rs @@ -5794,6 +5794,186 @@ opset_import: {opset_imports} } } +/// A chain of `depth` identical single-input nodes, `X -> t1 -> … -> Y`. +/// +/// Depth is the axis that separates the two costs #1077 conflates. Whatever +/// ORT charges per `Run` — session bookkeeping, fetch/feed marshalling — is +/// paid once regardless of depth, while our per-node dispatch is paid `depth` +/// times. Comparing depth 1, 10 and 100 against ORT therefore reads the +/// per-node slope directly, instead of a single number in which the fixed and +/// variable parts are indistinguishable. +/// +/// `dynamic` leaves the row count symbolic, so ORT cannot fold shapes at +/// session build and our shape inference runs for real on every `Run`. A +/// dispatch path that is only fast on static shapes is not fast. +fn chain_case( + name: &'static str, + op: &'static str, + depth: usize, + len: usize, + dynamic: bool, +) -> MatmulFamilyCase { + assert!(depth >= 1, "a chain needs at least one node"); + let dims = vec![1, len as i64]; + let decl = |n: &str| -> String { + if dynamic { + format!( + "{{ name: \"{n}\" type {{ tensor_type {{ elem_type: {ELEM_F32} \ + shape {{ dim: [{{ dim_param: \"batch\" }}, {{ dim_value: {len} }}] }} }} }} }}" + ) + } else { + textproto_value_info(n, ELEM_F32, &dims) + } + }; + let nodes: Vec = (0..depth) + .map(|i| { + let input = if i == 0 { + "X".to_owned() + } else { + format!("t{i}") + }; + let output = if i + 1 == depth { + "Y".to_owned() + } else { + format!("t{}", i + 1) + }; + format!("{{ input: [\"{input}\"] output: [\"{output}\"] op_type: \"{op}\" }}") + }) + .collect(); + let model = format!( + r#"ir_version: 10 +graph {{ + node: [{}] + name: "{name}" + input: [{}] + output: [{}] +}} +opset_import: [{{ version: 17 }}] +"#, + nodes.join(", "), + decl("X"), + decl("Y"), + ); + let x = activation_f32(1, len); + MatmulFamilyCase { + name, + op, + model, + inputs: vec![GeneratedInput { + name: "X", + elem: ELEM_F32, + dims, + data: f32_slice_to_bytes(&x), + }], + output: "Y", + output_elem: ELEM_F32, + ort_can_build: true, + // A Relu chain is exact in both implementations; Identity is a copy. + // Anything above zero here would be hiding a real disagreement. + tolerance: 0.0, + } +} + +/// A single small `MatMul`, `[1, k] x [k, n]`, with both operands supplied at +/// run time. +/// +/// The elementwise cases all have trivial kernels, which is what makes them +/// good dispatch probes — but it also means a win there could be an artefact of +/// how little work the kernel does. `MatMul` at these sizes still has a real +/// kernel while staying small enough for dispatch to matter, so it is the case +/// that says whether the fast path generalises beyond one-liners. +/// +/// Both operands are runtime inputs rather than initializers, which keeps the +/// weight out of the TextFormat source (an octal-escaped 1024x1024 f32 blob is +/// megabytes of test fixture) and denies both sides any prepacking, so the +/// comparison stays about dispatch. +fn small_matmul_case(name: &'static str, k: usize, n: usize) -> MatmulFamilyCase { + let model = format!( + r#"ir_version: 10 +graph {{ + node: [{{ input: ["X", "W"] output: ["Y"] op_type: "MatMul" }}] + name: "{name}" + input: [{}, {}] + output: [{}] +}} +opset_import: [{{ version: 17 }}] +"#, + textproto_value_info("X", ELEM_F32, &[1, k as i64]), + textproto_value_info("W", ELEM_F32, &[k as i64, n as i64]), + textproto_value_info("Y", ELEM_F32, &[1, n as i64]), + ); + MatmulFamilyCase { + name, + op: "MatMul", + model, + inputs: vec![ + GeneratedInput { + name: "X", + elem: ELEM_F32, + dims: vec![1, k as i64], + data: f32_slice_to_bytes(&activation_f32(1, k)), + }, + GeneratedInput { + name: "W", + elem: ELEM_F32, + dims: vec![k as i64, n as i64], + data: f32_slice_to_bytes(&activation_f32(k, n)), + }, + ], + output: "Y", + output_elem: ELEM_F32, + ort_can_build: true, + tolerance: 1e-3, + } +} + +/// The grid issue #1077 is closed against. +/// +/// Every row isolates one variable. Read together they say where the remaining +/// dispatch cost is; read individually none of them would. +/// +/// * **`identity_1`** is the floor. An `Identity` node does a memcpy, so +/// essentially everything the timer sees is the cost of getting there and +/// back. If we are at parity anywhere, it has to be here first. +/// * **`relu_1` / `relu_10` / `relu_100`** vary only depth, giving the +/// per-node slope with the per-`Run` constant divided out. +/// * **`*_dyn`** repeat the static rows with a symbolic batch dimension, so a +/// fast path that only works when shapes are known at build time cannot pass +/// unnoticed. +/// * **`matmul_*`** check the result is not an artefact of trivial kernels. +/// * The `4k` width keeps every kernel small enough that dispatch is a visible +/// fraction; at 1 Mi it would be rounding error, which is the whole reason +/// the earlier activation sweep could not answer this question. +fn dispatch_grid_cases() -> Vec { + const W: usize = 4096; + const TINY: usize = 8; + vec![ + chain_case("grid_identity_1_static", "Identity", 1, W, false), + chain_case("grid_identity_10_static", "Identity", 10, W, false), + chain_case("grid_relu_1_static", "Relu", 1, W, false), + chain_case("grid_relu_10_static", "Relu", 10, W, false), + chain_case("grid_relu_100_static", "Relu", 100, W, false), + chain_case("grid_identity_1_dyn", "Identity", 1, W, true), + chain_case("grid_relu_1_dyn", "Relu", 1, W, true), + chain_case("grid_relu_10_dyn", "Relu", 10, W, true), + chain_case("grid_relu_100_dyn", "Relu", 100, W, true), + // Dispatch-isolating widths. At W = 4096 a Relu node moves 32 KiB, and + // that traffic — not the dispatch around it — sets the per-node cost, + // so the ratio there reports elementwise throughput wearing a + // dispatch-shaped hat. At 8 elements the kernel is a single masked + // vector op and essentially everything left is overhead, which is what + // #1077 is actually about. The pair brackets it: TINY is the overhead + // floor, W the throughput regime. + chain_case("grid_relu_1_tiny", "Relu", 1, TINY, false), + chain_case("grid_relu_10_tiny", "Relu", 10, TINY, false), + chain_case("grid_relu_100_tiny", "Relu", 100, TINY, false), + chain_case("grid_relu_1_tiny_dyn", "Relu", 1, TINY, true), + chain_case("grid_relu_10_tiny_dyn", "Relu", 10, TINY, true), + small_matmul_case("grid_matmul_128x128", 128, 128), + small_matmul_case("grid_matmul_512x512", 512, 512), + ] +} + /// The elementwise grid `CPU_ACTIVATION_GAPS.md` is written against, as /// session-level cases. /// @@ -6180,7 +6360,25 @@ fn plugin_path_ab_vs_plain_ort() { println!( "case,ours_p50_ms,ort_p50_ms,ratio_p50,ours_p90_ms,ort_p90_ms,ratio_p90,cold_ours_ms,cold_ort_ms" ); - for case in matmul_bench_cases().into_iter().chain(unary_bench_cases()) { + // The #1077 grid is off by default: it is 11 more sessions, and the + // matmul/unary rows answer a different question (kernel throughput) than + // these do (dispatch overhead). `NXRT_MM_BENCH_GRID=1` adds them; + // `NXRT_MM_BENCH_GRID=only` runs them alone, which is what a dispatch + // measurement wants — nothing else resident, nothing else warming caches. + let grid = std::env::var("NXRT_MM_BENCH_GRID").unwrap_or_default(); + let cases: Vec = match grid.as_str() { + "only" => dispatch_grid_cases(), + "1" => matmul_bench_cases() + .into_iter() + .chain(unary_bench_cases()) + .chain(dispatch_grid_cases()) + .collect(), + _ => matmul_bench_cases() + .into_iter() + .chain(unary_bench_cases()) + .collect(), + }; + for case in cases { if !filter.is_empty() && !case.name.contains(&filter) { continue; } diff --git a/crates/onnx-runtime-ep-plugin/src/compute.rs b/crates/onnx-runtime-ep-plugin/src/compute.rs index 2621f66aef..c30a266b08 100644 --- a/crates/onnx-runtime-ep-plugin/src/compute.rs +++ b/crates/onnx-runtime-ep-plugin/src/compute.rs @@ -1132,6 +1132,18 @@ struct DeviceStaging { /// Owns an `OrtMemoryInfo` rebuilt via `CreateMemoryInfo_V2`, released on drop. struct ReconstructedMemInfo { ptr: *const ort::OrtMemoryInfo, + /// Whether `ptr` denotes device memory, recorded from the `device_type` it + /// was built with rather than asked of ORT afterwards. + /// + /// `mem_info_is_device` is exactly `device_type != CPU`, and this memory + /// info was created by passing `device_type` straight to + /// `CreateMemoryInfo_V2`, so the two agree by construction. Every EP that + /// supplies a host-to-device copier today is non-CPU, which would make a + /// hardcoded `true` correct — but nothing in `HostToDeviceCopier` requires + /// that, and a CPU-typed EP that grew a copier would then be handed host + /// memory as a device scratch target. Deriving it costs nothing and cannot + /// drift. + is_device: bool, } // SAFETY: the pointer is read-only after construction and released only on drop. @@ -1247,22 +1259,35 @@ unsafe fn device_mem_info( api: &ort::OrtApi, ctx: *mut ort::OrtKernelContext, staging: Option<&DeviceStaging>, -) -> Option<*const ort::OrtMemoryInfo> { +) -> Option<(*const ort::OrtMemoryInfo, bool)> { // A routed subgraph's intermediates must be allocated in *device* memory so // the device kernels that produce and consume them dereference valid device // pointers. On an interspersed CPU→device partition, kernel-context input 0 // may be a *host* boundary input (#982), so we must not assume input 0 lives // on the device — scan every input for a genuinely device-resident one. + // + // The scan remembers input 0's memory info as it goes. It used to throw it + // away and re-fetch it at the bottom, which cost two more FFI calls on the + // path every host EP takes: the scan finds nothing device-resident, and the + // fallback then asks ORT again for the value it had already seen first. + // Whether the result is a device is returned too, for the same reason — the + // one caller immediately asked, which meant a third query for a fact this + // function had just established. + let mut input0: Option<*const ort::OrtMemoryInfo> = None; if let Some(get_count) = api.KernelContext_GetInputCount { let mut count: usize = 0; crate::dispatch_probe::ort_call(); let status = unsafe { get_count(ctx, &mut count) }; if status.is_null() { for i in 0..count { - if let Some(mi) = unsafe { ort_input_mem_info(api, ctx, i) } - && unsafe { mem_info_is_device(api, mi) } - { - return Some(mi); + let Some(mi) = (unsafe { ort_input_mem_info(api, ctx, i) }) else { + continue; + }; + if i == 0 { + input0 = Some(mi); + } + if unsafe { mem_info_is_device(api, mi) } { + return Some((mi, true)); } } } @@ -1275,13 +1300,23 @@ unsafe fn device_mem_info( && let Some(recon) = staging.recon_mem_info.as_ref() && !recon.ptr.is_null() { - return Some(recon.ptr); + // Device-ness recorded when this was built from the EP's own device + // recipe, so it matches what the old `mem_info_is_device` call returned + // without spending an FFI call to re-derive it. + return Some((recon.ptr, recon.is_device)); } // Host EP (no device staging context): preserve the historical behavior of // using input 0's memory info. For a host EP this is host memory, which is // exactly where its intermediates belong, so host-only partitions are // unchanged by the #982 device-staging work. - unsafe { ort_input_mem_info(api, ctx, 0) } + // + // The scan above already looked at input 0 unless it failed or there are no + // inputs, so this normally costs nothing. + match input0 { + Some(mi) => Some((mi, false)), + None => unsafe { ort_input_mem_info(api, ctx, 0) } + .map(|mi| (mi, unsafe { mem_info_is_device(api, mi) })), + } } /// Memory info of the `OrtValue` bound to kernel-context input `index`. @@ -1628,7 +1663,10 @@ fn reconstruct_device_mem_info( if ptr.is_null() { return None; } - Some(ReconstructedMemInfo { ptr }) + Some(ReconstructedMemInfo { + ptr, + is_device: device_type != ort::OrtMemoryInfoDeviceType_CPU, + }) } /// Device type ORT placed a memory info on, or `None` if unreadable. @@ -1969,7 +2007,7 @@ unsafe fn prepare_workspace( kernel: &dyn Kernel, plans: &WorkspacePlanCache, inputs: &[TensorView<'_>], - node_label: &str, + node_label: impl std::fmt::Display + Copy, ) -> Result, String> { let metadata: Vec> = inputs .iter() @@ -2083,7 +2121,7 @@ fn workspace_trace_enabled() -> bool { /// returned, capped at 4096 — the useful question is "did ORT already meet the /// kernel's alignment", not the exact 2-adic valuation of the address. fn workspace_trace_line( - node_label: &str, + node_label: impl std::fmt::Display, base: usize, aligned: usize, bytes: usize, @@ -2225,8 +2263,9 @@ unsafe extern "C" fn compute_execute( // memory, so multi-node intermediates stay on the GPU (a host buffer // would make the next kernel dereference a host pointer as device → // CUDA_ERROR_ILLEGAL_ADDRESS). `None` falls back to host buffers. - let scratch_mem_info = + let scratch_mem_info_and_kind = unsafe { device_mem_info(api_ref, kernel_context, exported.device_staging.as_ref()) }; + let scratch_mem_info = scratch_mem_info_and_kind.map(|(mi, _)| mi); // Where a routed subgraph's *intermediates* are allocated. // @@ -2244,8 +2283,8 @@ unsafe extern "C" fn compute_execute( // difference on identical kernels, against ORT's own 220 us for the // same graph. So: device memory when the resolved memory info is a // device, host buffers otherwise. - let intermediate_scratch = match scratch_mem_info { - Some(mem_info) if unsafe { mem_info_is_device(api_ref, mem_info) } => Some(mem_info), + let intermediate_scratch = match scratch_mem_info_and_kind { + Some((mem_info, true)) => Some(mem_info), _ => None, }; @@ -2273,10 +2312,11 @@ unsafe extern "C" fn compute_execute( // last reader has run is dead and its storage can be handed // straight back to the next allocation, so a chain of N nodes // cycles a couple of buffers instead of touching N fresh ones. - let last_reader = - last_reader_per_buffer(&routing.input_sources, routing.num_intermediate_buffers); + let (retire_starts, retire_items) = + retirements_per_node(&routing.input_sources, routing.num_intermediate_buffers); for (node_idx, entry) in exported.entries.iter().enumerate() { + let node_probe = crate::dispatch_probe::Phase::TensorBind.enter(); let sources = &routing.input_sources[node_idx]; let sinks = &routing.output_sinks[node_idx]; @@ -2321,12 +2361,15 @@ unsafe extern "C" fn compute_execute( } // Infer output shapes. + node_probe.end(); + let shape_probe = crate::dispatch_probe::Phase::DispatchLookup.enter(); let output_shapes = match infer_shapes(&entry.shape_inference, &kernel_inputs) { Ok(s) => s, Err(e) => { return fail_status(&format!("Compute: shape inference failed: {e}")); } }; + shape_probe.end(); // Execute — dispatch based on sinks. // For outputs going to ORT we allocate via ORT API; @@ -2483,11 +2526,18 @@ unsafe extern "C" fn compute_execute( // Collect all output views using the per-slot view map so // positions stay aligned even when absent slots are present. - let absent_shapes: Vec> = output_shapes.clone(); - let absent_strides_storage: Vec> = absent_shapes - .iter() - .map(|s| contiguous_strides(s)) - .collect(); + let absent_shapes: &[Vec] = &output_shapes; + // One entry per *absent* slot, not per output slot. Only + // absent slots read these, and a node with an absent output is + // the exception, so building them for every output cost an + // allocation per output on every node that has none. Each entry + // is keyed by the same index that keys `absent_scratch`, where + // the slot it belongs to is recorded -- so this stays aligned by + // construction rather than by an invariant about emptiness. + let absent_strides_storage = absent_slot_strides( + absent_scratch.iter().map(|(slot, _, _)| *slot), + absent_shapes, + ); let mut all_output_views: Vec<_> = { let mut ort_iter = ort_out_views.drain(..); let mut buf_iter = new_bufs.iter_mut(); @@ -2503,7 +2553,7 @@ unsafe extern "C" fn compute_execute( RoutedSlotKind::Absent(idx) => { let (_, scratch_buf, dtype) = &mut absent_scratch[*idx]; let shape = &absent_shapes[slot_idx]; - let strides = &absent_strides_storage[slot_idx]; + let strides = &absent_strides_storage[*idx]; TensorMut::new( DevicePtrMut(scratch_buf.as_mut_ptr().cast()), *dtype, @@ -2517,13 +2567,12 @@ unsafe extern "C" fn compute_execute( .collect() }; - let node_label = format!("Compute: node {node_idx}"); let plans = match exported.workspace_plan_cache(node_idx) { Some(plans) => plans, None => { return fail_status(&format!( - "{node_label}: no workspace plan cache for this node (entries and \ - workspace_plans have drifted)" + "Compute: node {node_idx}: no workspace plan cache for this node \ + (entries and workspace_plans have drifted)" )); } }; @@ -2548,13 +2597,14 @@ unsafe extern "C" fn compute_execute( &*entry.kernel, plans, &kernel_inputs, - &node_label, + format_args!("Compute: node {node_idx}"), ) } { Ok(w) => w, Err(e) => return fail_status(&e), }; + let kernel_probe = crate::dispatch_probe::Phase::KernelInvoke.enter(); if let Err(e) = entry.kernel.execute_with_workspace( &kernel_inputs, &mut all_output_views, @@ -2562,6 +2612,8 @@ unsafe extern "C" fn compute_execute( ) { return fail_status(&format!("Compute: kernel execution failed: {e}")); } + kernel_probe.end(); + crate::dispatch_probe::count(crate::dispatch_probe::Event::NodeExecuted); EXECUTED_NODE_COUNT.fetch_add(1, Ordering::Relaxed); // Store new intermediate buffers. @@ -2581,10 +2633,9 @@ unsafe extern "C" fn compute_execute( // it here — rather than at the end of the subgraph — is what // lets the next node's allocation land on storage that is // still hot in cache. - for (buf_idx, last) in last_reader.iter().enumerate() { - if *last == Some(node_idx) - && let Some(dead) = intermediates[buf_idx].take() - { + for &buf_idx in &retire_items[retire_starts[node_idx]..retire_starts[node_idx + 1]] + { + if let Some(dead) = intermediates[buf_idx].take() { recycle_host_intermediate(dead.data); } } @@ -2714,11 +2765,16 @@ unsafe extern "C" fn compute_execute( } // Build output views in node-output order so the kernel sees the // full arity including absent scratch slots. - let absent_shapes: Vec> = output_shapes.clone(); - let absent_strides_storage: Vec> = absent_shapes - .iter() - .map(|s| contiguous_strides(s)) - .collect(); + let absent_shapes: &[Vec] = &output_shapes; + // See the routed path. `slot_map` pushes `Absent(idx)` with + // increasing `idx`, so filtering it in slot order yields the strides + // in absent-index order. + let absent_strides_storage = absent_slot_strides( + slot_map.iter().enumerate().filter_map(|(slot_idx, kind)| { + matches!(kind, SlotKind::Absent(_)).then_some(slot_idx) + }), + absent_shapes, + ); // First, get mutable views of all ORT outputs in order. let mut ort_views: Vec> = owned_outputs.iter_mut().map(|o| o.view_mut()).collect(); @@ -2733,7 +2789,7 @@ unsafe extern "C" fn compute_execute( SlotKind::Absent(idx) => { let buf = &mut absent_bufs[*idx]; let shape = &absent_shapes[slot_idx]; - let strides = &absent_strides_storage[slot_idx]; + let strides = &absent_strides_storage[*idx]; let scratch_dtype = absent_dtypes[*idx]; let view = TensorMut::new( DevicePtrMut(buf.as_mut_ptr().cast()), @@ -2922,6 +2978,68 @@ fn last_reader_per_buffer( last } +/// Buffers retiring after each node, as a CSR index: node `i` retires the +/// buffers in `items[starts[i]..starts[i + 1]]`. +/// +/// The routed loop used to find these by scanning every buffer's last-reader +/// entry at every node, which is O(nodes x buffers) -- and a chain has a buffer +/// per node, so it is quadratic in subgraph depth. That cost is invisible at +/// depth 1 and dominant at depth 100. +/// +/// Inverting the map once per `Run` makes each node touch only its own +/// retirements. Within a node they stay in ascending buffer order, as the scan +/// produced them. +fn retirements_per_node( + input_sources: &[Vec], + num_buffers: usize, +) -> (Vec, Vec) { + let last = last_reader_per_buffer(input_sources, num_buffers); + let nodes = input_sources.len(); + let mut starts = vec![0usize; nodes + 1]; + for &reader in last.iter().flatten() { + if reader < nodes { + starts[reader + 1] += 1; + } + } + for i in 0..nodes { + starts[i + 1] += starts[i]; + } + let mut items = vec![0usize; starts[nodes]]; + let mut cursor = starts.clone(); + for (buffer, reader) in last.iter().enumerate() { + if let Some(reader) = reader + && *reader < nodes + { + items[cursor[*reader]] = buffer; + cursor[*reader] += 1; + } + } + (starts, items) +} + +/// Contiguous strides for the *absent* output slots only, in absent-index +/// order. +/// +/// Absent slots are the exception, so the routed loop used to build a stride +/// vector for every output of every node and read it only when a slot turned +/// out to be absent -- an allocation per output, per node, discarded unused. +/// +/// The result is keyed by absent index, the same key `Absent(idx)` carries, so +/// callers index it with `idx` rather than the slot number. Both callers pass +/// absent slots in ascending slot order, which is the order `idx` was assigned +/// in, so the two agree by construction. +fn absent_slot_strides( + absent_slots: impl Iterator, + shapes: &[Vec], +) -> Vec> { + absent_slots + .map(|slot| match shapes.get(slot) { + Some(shape) => contiguous_strides(shape), + None => Vec::new(), + }) + .collect() +} + fn contiguous_strides(shape: &[usize]) -> Vec { let mut strides = vec![1i64; shape.len()]; for i in (0..shape.len().saturating_sub(1)).rev() { @@ -4841,6 +4959,146 @@ mod tests { ); } + /// The CSR index must retire exactly what the per-node scan retired. + /// + /// This is the property the rewrite trades on, so it is checked against the + /// scan itself rather than against hand-written expectations: for every + /// node, the buffers the index yields are the buffers whose last reader is + /// that node, in the same ascending order the scan visited them. + fn scan_retirements(sources: &[Vec], buffers: usize) -> Vec> { + let last = super::last_reader_per_buffer(sources, buffers); + (0..sources.len()) + .map(|node| { + last.iter() + .enumerate() + .filter(|(_, l)| **l == Some(node)) + .map(|(buffer, _)| buffer) + .collect() + }) + .collect() + } + + fn index_retirements(sources: &[Vec], buffers: usize) -> Vec> { + let (starts, items) = super::retirements_per_node(sources, buffers); + (0..sources.len()) + .map(|node| items[starts[node]..starts[node + 1]].to_vec()) + .collect() + } + + #[test] + fn the_retirement_index_agrees_with_the_scan_it_replaces() { + let chain: Vec> = std::iter::once(vec![NodeInputSource::Ort(0)]) + .chain((0..9).map(|i| vec![NodeInputSource::Buffer(i)])) + .collect(); + let cases: Vec<(Vec>, usize)> = vec![ + // A plain chain: every buffer dies after the node that reads it. + (chain, 9), + // A buffer read twice must retire only after the later reader, and + // must appear exactly once in the index. + ( + vec![ + vec![NodeInputSource::Ort(0)], + vec![NodeInputSource::Buffer(0)], + vec![NodeInputSource::Ort(1)], + vec![NodeInputSource::Buffer(0), NodeInputSource::Buffer(1)], + ], + 2, + ), + // Unread and out-of-range buffers retire nowhere. + ( + vec![ + vec![NodeInputSource::Ort(0)], + vec![NodeInputSource::Buffer(0), NodeInputSource::Buffer(7)], + vec![NodeInputSource::Absent], + ], + 2, + ), + // One node retiring several buffers at once: order must be + // ascending, as the scan produced it. + ( + vec![ + vec![NodeInputSource::Ort(0)], + vec![NodeInputSource::Ort(1)], + vec![ + NodeInputSource::Buffer(2), + NodeInputSource::Buffer(0), + NodeInputSource::Buffer(1), + ], + ], + 3, + ), + // No buffers at all. + (vec![vec![NodeInputSource::Ort(0)]], 0), + ]; + for (sources, buffers) in cases { + assert_eq!( + index_retirements(&sources, buffers), + scan_retirements(&sources, buffers), + "retirement index diverged from the scan for {sources:?}" + ); + } + } + + /// Every buffer retires exactly once, or storage leaks for the length of + /// the subgraph (or worse, is recycled twice). + #[test] + fn every_read_buffer_retires_exactly_once() { + let sources: Vec> = std::iter::once(vec![NodeInputSource::Ort(0)]) + .chain((0..20).map(|i| vec![NodeInputSource::Buffer(i)])) + .collect(); + let (starts, items) = super::retirements_per_node(&sources, 20); + let mut seen = items.clone(); + seen.sort_unstable(); + seen.dedup(); + assert_eq!(seen.len(), items.len(), "a buffer retired twice"); + assert_eq!(items.len(), 20, "every buffer of a chain is read once"); + assert_eq!(starts[sources.len()], items.len()); + } + + /// Absent strides are keyed by absent index, not slot index. + /// + /// Getting this wrong is invisible for a node with a single absent output + /// at slot 0 -- the two indices coincide -- and silently hands the wrong + /// strides to every other shape. It is pinned with absent slots that are + /// neither first nor contiguous, and with differing shapes per slot so a + /// mis-key produces a different answer rather than the same one twice. + #[test] + fn absent_strides_are_keyed_by_absent_index_not_slot_index() { + let shapes = vec![ + vec![2usize, 3, 4], + vec![5usize, 6], + vec![7usize, 8, 9], + vec![10usize], + ]; + // Slots 1 and 3 are absent, so absent index 0 is slot 1 and absent + // index 1 is slot 3. + let got = super::absent_slot_strides([1usize, 3].into_iter(), &shapes); + assert_eq!(got.len(), 2, "one entry per absent slot, not per output"); + assert_eq!(got[0], super::contiguous_strides(&shapes[1])); + assert_eq!(got[1], super::contiguous_strides(&shapes[3])); + assert_ne!( + got[0], + super::contiguous_strides(&shapes[0]), + "keying by slot index would have produced this" + ); + } + + #[test] + fn a_node_with_no_absent_outputs_builds_no_absent_strides() { + let shapes = vec![vec![2usize, 3], vec![4usize, 5]]; + assert!(super::absent_slot_strides(std::iter::empty(), &shapes).is_empty()); + } + + /// An out-of-range slot must not panic in the middle of a `Run`. It cannot + /// arise from a well-formed slot map, but this runs inside ORT's callback + /// where a panic crosses an FFI boundary. + #[test] + fn an_out_of_range_absent_slot_yields_empty_strides() { + let shapes = vec![vec![2usize, 3]]; + let got = super::absent_slot_strides([9usize].into_iter(), &shapes); + assert_eq!(got, vec![Vec::::new()]); + } + #[test] fn recycled_intermediate_storage_is_reused_without_reallocating() { super::drain_host_intermediate_pool(); @@ -5796,3 +6054,240 @@ mod workspace_plan_cache_tests { ); } } + +/// Pins the FFI cost of resolving where a routed subgraph's intermediates +/// should live — a question every `Compute` call asks, before any kernel runs. +/// +/// A host EP takes the longest path through `device_mem_info`: the scan finds +/// nothing device-resident, so it runs to completion and then falls back. The +/// count below is what that costs, and it is the per-`Run` fixed overhead +/// #1077 measures as ~0.9 us worse than ORT's. +#[cfg(all(test, feature = "dispatch_probe"))] +mod mem_info_cost { + use super::*; + use crate::dispatch_probe::{self, Event}; + + const MI: *const ort::OrtMemoryInfo = std::ptr::without_provenance(3); + const VALUE: *const ort::OrtValue = std::ptr::without_provenance(4); + const RECON: *const ort::OrtMemoryInfo = std::ptr::without_provenance(5); + + /// A device EP's staging context with a reconstructed memory info of the + /// given device-ness. + /// + /// Built directly rather than through `set_device_staging`, which would + /// need a live ORT to call `CreateMemoryInfo_V2`. Dropping the result is + /// safe with no host API installed: `ReconstructedMemInfo::drop` returns + /// early when `host_api()` is null, so the sentinel pointer is never + /// passed to `ReleaseMemoryInfo`. + fn staging_with_recon(ptr: *const ort::OrtMemoryInfo, is_device: bool) -> DeviceStaging { + struct NoCopier; + impl HostToDeviceCopier for NoCopier { + unsafe fn copy_host_to_device( + &self, + _src: &[u8], + _dst: *mut std::ffi::c_void, + ) -> onnx_runtime_ep_api::Result<()> { + unreachable!("scratch placement must not copy") + } + } + assert!( + crate::status::host_api().is_null(), + "test builds a ReconstructedMemInfo over a sentinel pointer; a live \ + host API would hand it to ReleaseMemoryInfo on drop" + ); + DeviceStaging { + copier: Arc::new(NoCopier), + recon_mem_info: Some(ReconstructedMemInfo { ptr, is_device }), + } + } + + unsafe extern "C" fn count_1( + _c: *const ort::OrtKernelContext, + out: *mut usize, + ) -> ort::OrtStatusPtr { + unsafe { *out = 1 }; + std::ptr::null_mut() + } + unsafe extern "C" fn get_input( + _c: *const ort::OrtKernelContext, + _i: usize, + out: *mut *const ort::OrtValue, + ) -> ort::OrtStatusPtr { + unsafe { *out = VALUE }; + std::ptr::null_mut() + } + unsafe extern "C" fn tensor_mem_info( + _v: *const ort::OrtValue, + out: *mut *const ort::OrtMemoryInfo, + ) -> ort::OrtStatusPtr { + unsafe { *out = MI }; + std::ptr::null_mut() + } + /// Always reports host memory — the answer a CPU EP always gets, and the + /// branch that makes the scan run to completion. + unsafe extern "C" fn device_type_cpu( + _m: *const ort::OrtMemoryInfo, + out: *mut ort::OrtMemoryInfoDeviceType, + ) { + unsafe { *out = ort::OrtMemoryInfoDeviceType_CPU }; + } + unsafe extern "C" fn device_type_gpu( + _m: *const ort::OrtMemoryInfo, + out: *mut ort::OrtMemoryInfoDeviceType, + ) { + unsafe { *out = ort::OrtMemoryInfoDeviceType_GPU }; + } + + fn api(host: bool) -> ort::OrtApi { + let mut api: ort::OrtApi = unsafe { std::mem::zeroed() }; + api.KernelContext_GetInputCount = Some(count_1); + api.KernelContext_GetInput = Some(get_input); + api.GetTensorMemoryInfo = Some(tensor_mem_info); + api.MemoryInfoGetDeviceType = Some(if host { + device_type_cpu + } else { + device_type_gpu + }); + api + } + + /// Four calls for a one-input host node: the input count, then `GetInput`, + /// `GetTensorMemoryInfo` and `MemoryInfoGetDeviceType` for that input. + /// + /// It was seven. The scan used to discard input 0's memory info and ask ORT + /// for it again at the fallback (two calls for a value it had already + /// seen), and the caller then asked a third time whether the result was a + /// device — a fact this function had just established and thrown away. + #[test] + fn resolving_host_scratch_costs_four_ort_calls() { + let api = api(true); + let before = dispatch_probe::snapshot(); + let got = unsafe { device_mem_info(&api, std::ptr::null_mut(), None) }; + let d = dispatch_probe::snapshot().since(&before); + + assert_eq!( + got, + Some((MI, false)), + "host memory, reported as not device" + ); + assert_eq!( + d.event(Event::OrtFfiCall), + 4, + "FFI cost of resolving scratch placement changed" + ); + } + + /// A device-resident input short-circuits the scan, and the device-ness + /// comes back with it rather than costing another query. + #[test] + fn a_device_input_short_circuits_and_reports_itself() { + let api = api(false); + let before = dispatch_probe::snapshot(); + let got = unsafe { device_mem_info(&api, std::ptr::null_mut(), None) }; + let d = dispatch_probe::snapshot().since(&before); + + assert_eq!(got, Some((MI, true))); + assert_eq!(d.event(Event::OrtFfiCall), 4); + } + + /// The fallback still works when the scan cannot run at all, which is the + /// one case that legitimately needs the extra fetch. Behaviour here is what + /// it always was; only the cheap path got cheaper. + #[test] + fn a_missing_input_count_still_falls_back_to_input_zero() { + let mut api = api(true); + api.KernelContext_GetInputCount = None; + let got = unsafe { device_mem_info(&api, std::ptr::null_mut(), None) }; + assert_eq!(got, Some((MI, false))); + } + + /// A node with no inputs has nothing to scan and nothing to fall back on + /// beyond input 0, which does not exist either. It must not invent one. + #[test] + fn a_node_with_no_inputs_reports_no_memory_info() { + unsafe extern "C" fn none( + _c: *const ort::OrtKernelContext, + out: *mut usize, + ) -> ort::OrtStatusPtr { + unsafe { *out = 0 }; + std::ptr::null_mut() + } + unsafe extern "C" fn no_value( + _c: *const ort::OrtKernelContext, + _i: usize, + out: *mut *const ort::OrtValue, + ) -> ort::OrtStatusPtr { + unsafe { *out = std::ptr::null() }; + std::ptr::null_mut() + } + let mut api = api(true); + api.KernelContext_GetInputCount = Some(none); + api.KernelContext_GetInput = Some(no_value); + assert_eq!( + unsafe { device_mem_info(&api, std::ptr::null_mut(), None) }, + None + ); + } + + /// A device EP whose node has only host-resident inputs falls back to the + /// memory info reconstructed from its own device recipe, so intermediates + /// still land on the device. + /// + /// The old code returned that pointer and left the caller to ask ORT + /// whether it was a device; the device-ness is now carried alongside it. + /// This is the one branch whose reporting changed, so it is pinned + /// explicitly rather than left to the host-EP tests. + #[test] + fn all_host_inputs_on_a_device_ep_fall_back_to_the_reconstructed_device_info() { + let staging = staging_with_recon(RECON, true); + let api = api(true); + let before = dispatch_probe::snapshot(); + let got = unsafe { device_mem_info(&api, std::ptr::null_mut(), Some(&staging)) }; + let d = dispatch_probe::snapshot().since(&before); + + assert_eq!( + got, + Some((RECON, true)), + "device scratch target, reported as a device" + ); + assert_eq!( + d.event(Event::OrtFfiCall), + 4, + "the fallback must not re-ask ORT about the info it just built" + ); + } + + /// The device-ness of the reconstructed info is recorded from the + /// `device_type` it was built with, not hardcoded. + /// + /// Nothing in `HostToDeviceCopier` obliges an EP that provides a copier to + /// be non-CPU. Every one that exists today is, which is what makes a + /// hardcoded `true` look safe -- but a CPU-typed EP that grew a copier + /// would be handed host memory as a device scratch target, and the staging + /// path would treat host pointers as device pointers. `mem_info_is_device` + /// is exactly `device_type != CPU`, so recording it at construction + /// reproduces the old query in this case too. + #[test] + fn a_cpu_typed_reconstruction_is_not_reported_as_a_device() { + let staging = staging_with_recon(RECON, false); + let api = api(true); + assert_eq!( + unsafe { device_mem_info(&api, std::ptr::null_mut(), Some(&staging)) }, + Some((RECON, false)), + "host-typed reconstruction must not be reported as device memory" + ); + } + + /// A device-resident input still wins over the reconstruction: the scan + /// short-circuits before the fallback is consulted. + #[test] + fn a_device_input_wins_over_the_reconstruction() { + let staging = staging_with_recon(RECON, true); + let api = api(false); + assert_eq!( + unsafe { device_mem_info(&api, std::ptr::null_mut(), Some(&staging)) }, + Some((MI, true)), + "the actual device input, not the fallback recipe" + ); + } +} From bd232c8267b512ea59ddab9ba75f7521beb78ac4 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Wed, 19 Aug 2026 05:23:22 +0000 Subject: [PATCH 6/7] Attribute dispatch allocations to phases across the ORT ABI Issue #1077 has ~7 allocations per node left and no way to say where they are. This adds per-phase allocation attribution and gets it reporting from a live routed dispatch. The allocator hook charges each allocation to the innermost open phase. PhaseGuard now saves the enclosing phase on enter and restores it on drop, so nested phases -- which the dispatch path really does produce, because a guard closes at scope exit rather than at an early return -- attribute to the inner one rather than losing the outer. Reading it back needs the counters to come from the library ORT dlopens, not from the test binary: the cdylib has its own global allocator, so an allocator installed in the harness never sees an allocation made inside Compute. The counting allocator is therefore installed in the cdylib under a new dispatch_probe feature, and the harness reads totals through the existing C export. Also fixes a resolver bug this uncovered. cdylib_resolve rebuilds the cdylib with a fixed feature list, and dispatch_probe was missing from it, so the rebuild silently replaced the instrumented library with a default-feature one and the probe reported zeros for every phase. The file already warned about this hazard for mlas. Nothing fails loudly when a feature is omitted, so the comment now says so explicitly. Production is unaffected: without the feature record_alloc is an inline no-op, no global allocator is installed, and the harness path is additionally gated on NXRT_MM_BENCH_PROBE=1. Falsified: reverting the outer-phase restore fails exactly the two nesting tests; the five attribution tests install the counting allocator for real rather than calling record_alloc by hand. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- crates/onnx-runtime-ep-cpu-plugin/Cargo.toml | 6 + crates/onnx-runtime-ep-cpu-plugin/src/lib.rs | 13 ++ .../tests/cdylib_resolve.rs | 22 ++- .../tests/plugin_ort_e2e.rs | 127 ++++++++++++ .../src/dispatch_probe.rs | 117 ++++++++++- .../tests/alloc_attribution.rs | 185 ++++++++++++++++++ 6 files changed, 459 insertions(+), 11 deletions(-) create mode 100644 crates/onnx-runtime-ep-plugin/tests/alloc_attribution.rs diff --git a/crates/onnx-runtime-ep-cpu-plugin/Cargo.toml b/crates/onnx-runtime-ep-cpu-plugin/Cargo.toml index 2e4c40addb..131a338f7c 100644 --- a/crates/onnx-runtime-ep-cpu-plugin/Cargo.toml +++ b/crates/onnx-runtime-ep-cpu-plugin/Cargo.toml @@ -20,6 +20,12 @@ crate-type = ["cdylib", "lib"] # C++ toolchain can still build the cdylib. mlas = ["onnx-runtime-ep-cpu/mlas"] +# Compile the dispatch probe into the cdylib and install the counting allocator +# as its global allocator, so a harness can attribute allocations to dispatch +# phases across the ORT boundary. Never enabled for a shipped build: it adds a +# thread-local read and two counter updates to every allocation in the library. +dispatch_probe = ["onnx-runtime-ep-plugin/dispatch_probe"] + [dependencies] onnx-runtime-ep-cpu = { workspace = true } onnx-runtime-ep-plugin = { path = "../onnx-runtime-ep-plugin", version = "=0.1.0-dev.5" } diff --git a/crates/onnx-runtime-ep-cpu-plugin/src/lib.rs b/crates/onnx-runtime-ep-cpu-plugin/src/lib.rs index bebaaad724..f12afaff49 100644 --- a/crates/onnx-runtime-ep-cpu-plugin/src/lib.rs +++ b/crates/onnx-runtime-ep-cpu-plugin/src/lib.rs @@ -10,6 +10,19 @@ use onnx_runtime_ep_cpu::{CpuExecutionProvider, build_cpu_registry_with_descriptors}; use onnx_runtime_ep_plugin::ep::KernelRegistryEntry; +// Attribute this library's allocations to whichever dispatch phase is open. +// +// It has to live here rather than in the harness: ORT `dlopen`s this cdylib, so +// it has its own global allocator, and an allocator installed in the test +// binary never sees an allocation made inside a `Compute` callback. The harness +// reads the totals back through `nxrt_dispatch_probe_snapshot`. +// +// Off unless the `dispatch_probe` feature is on, which no shipped build sets. +#[cfg(feature = "dispatch_probe")] +#[global_allocator] +static PROBE_ALLOC: onnx_runtime_ep_plugin::dispatch_probe::CountingAllocator = + onnx_runtime_ep_plugin::dispatch_probe::CountingAllocator::new(std::alloc::System); + /// Build `KernelRegistryEntry` slices from the CPU EP's real registry. /// /// Each entry's `supported_dtypes` is derived from the kernel's actual dispatch diff --git a/crates/onnx-runtime-ep-cpu-plugin/tests/cdylib_resolve.rs b/crates/onnx-runtime-ep-cpu-plugin/tests/cdylib_resolve.rs index ac326cda90..50d552effc 100644 --- a/crates/onnx-runtime-ep-cpu-plugin/tests/cdylib_resolve.rs +++ b/crates/onnx-runtime-ep-cpu-plugin/tests/cdylib_resolve.rs @@ -16,12 +16,22 @@ const PACKAGE: &str = "onnx-runtime-ep-cpu-plugin"; /// `cargo test -p onnx-runtime-ep-cpu-plugin --features mlas` compiles the /// test with MLAS; without this the rebuild below would replace the cdylib /// with a default-feature build and the suite would load the wrong library. -fn features() -> &'static [&'static str] { +/// +/// Every feature of this package that changes the cdylib must be mirrored +/// here. Omitting one does not fail loudly: the rebuild silently produces a +/// library without it and the tests then measure the wrong binary. That is not +/// hypothetical — `dispatch_probe` was missing here at first and the probe +/// reported zeros for every phase, because the resolver had overwritten the +/// instrumented build with a default-feature one. +fn features() -> Vec<&'static str> { + let mut f = Vec::new(); if cfg!(feature = "mlas") { - &["mlas"] - } else { - &[] + f.push("mlas"); + } + if cfg!(feature = "dispatch_probe") { + f.push("dispatch_probe"); } + f } /// Locate the cpu-plugin cdylib, building it if needed. @@ -30,7 +40,7 @@ fn features() -> &'static [&'static str] { /// /// Panics with an actionable message when the cdylib cannot be produced. pub fn find_cpu_plugin_cdylib() -> PathBuf { - onnx_runtime_ort_testkit::find_plugin_cdylib_with_features(PACKAGE, features()).unwrap_or_else( + onnx_runtime_ort_testkit::find_plugin_cdylib_with_features(PACKAGE, &features()).unwrap_or_else( || { panic!( "{PACKAGE} cdylib could not be located or built. \ @@ -43,5 +53,5 @@ pub fn find_cpu_plugin_cdylib() -> PathBuf { /// Same as [`find_cpu_plugin_cdylib`] but returns `None` for tests that skip /// when the cdylib is absent (e.g. e2e tests that also need real ORT). pub fn find_cpu_plugin_cdylib_optional() -> Option { - onnx_runtime_ort_testkit::find_plugin_cdylib_with_features(PACKAGE, features()) + onnx_runtime_ort_testkit::find_plugin_cdylib_with_features(PACKAGE, &features()) } diff --git a/crates/onnx-runtime-ep-cpu-plugin/tests/plugin_ort_e2e.rs b/crates/onnx-runtime-ep-cpu-plugin/tests/plugin_ort_e2e.rs index 9b8a3556b9..0d7132d5fe 100644 --- a/crates/onnx-runtime-ep-cpu-plugin/tests/plugin_ort_e2e.rs +++ b/crates/onnx-runtime-ep-cpu-plugin/tests/plugin_ort_e2e.rs @@ -6452,6 +6452,19 @@ fn plugin_path_ab_vs_plain_ort() { ); } + // Allocation attribution runs before the timed loop so its + // counter updates never land inside a measured iteration. + if run_ours + && let Some(buf) = + probe_dispatch(api, session, &input_name_ptrs, &values, &output_name_ptrs, 64) + { + // Node count comes from ORT's assignment list, not from + // the case definition: per-node figures must divide by the + // nodes this EP actually received, and fusion means that + // is not the same as the number of `Run` callbacks. + report_probe(case.name, info.ops_on_our_ep().len().max(1), 64, &buf); + } + // Interleave one iteration each so a drifting host load lands on // both sides rather than on whichever ran second. let mut ours = Vec::with_capacity(iters); @@ -7427,3 +7440,117 @@ fn the_native_activation_family_executes_locally_and_matches_ort_numerics() { "\n✅ the_native_activation_family_executes_locally_and_matches_ort_numerics: PASSED" ); } + +/// Phase names, in the order `nxrt_dispatch_probe_snapshot` writes them. +/// +/// Duplicated from `dispatch_probe::Phase` because this harness talks to the +/// cdylib over the C ABI, not through the Rust type. `probe_phase_names` is +/// checked against the returned length, so a phase added on the library side +/// shows up as a loud mismatch here rather than as a silently shifted table. +const PROBE_PHASES: &[&str] = &[ + "OrtCallback", + "InputMeta", + "OutputMeta", + "TensorBind", + "DispatchLookup", + "KernelInvoke", + "StatusCrossing", +]; +const PROBE_EVENTS: &[&str] = &[ + "OrtFfiCall", + "DispatchAlloc", + "NodeExecuted", + "ShapeInferred", + "OutputMaterialized", +]; + +/// Open the EP cdylib a second time to reach its probe exports. +/// +/// ORT has already `dlopen`ed this exact path, so this returns a handle to the +/// same image and therefore the same counters — the allocations recorded here +/// are the ones made inside `Compute`, which an allocator installed in this +/// test binary could never see. +/// +/// `None` when the cdylib was built without the `dispatch_probe` feature, which +/// is the normal case; the symbols simply are not there. +fn probe_lib() -> Option<&'static libloading::Library> { + static LIB: std::sync::OnceLock> = std::sync::OnceLock::new(); + LIB.get_or_init(|| { + let path = cdylib_resolve::find_cpu_plugin_cdylib_optional()?; + // SAFETY: the path is the EP cdylib this harness already registered + // with ORT; re-opening it runs no new initialisers. + unsafe { libloading::Library::new(path) }.ok() + }) + .as_ref() +} + +/// Reset the probe, run `runs` iterations, and return the per-phase totals. +/// +/// Returns `None` unless `NXRT_MM_BENCH_PROBE=1` and the cdylib exports the +/// probe, so the default benchmark path is untouched. +#[allow(clippy::too_many_arguments)] +unsafe fn probe_dispatch( + api: *const ort::OrtApi, + session: *mut ort::OrtSession, + input_names: &[*const std::os::raw::c_char], + values: &[*const ort::OrtValue], + output_names: &[*const std::os::raw::c_char], + runs: usize, +) -> Option> { + if std::env::var("NXRT_MM_BENCH_PROBE").unwrap_or_default() != "1" { + return None; + } + let lib = probe_lib()?; + // SAFETY: both symbols are `extern "C"` exports of the loaded cdylib with + // exactly these signatures; a missing symbol returns `Err`. + unsafe { + let reset: libloading::Symbol<'_, unsafe extern "C" fn()> = + lib.get(b"nxrt_dispatch_probe_reset").ok()?; + let snapshot: libloading::Symbol<'_, unsafe extern "C" fn(*mut u64, usize) -> usize> = + lib.get(b"nxrt_dispatch_probe_snapshot").ok()?; + reset(); + bench_runs(api, session, input_names, values, output_names, runs); + let need = PROBE_PHASES.len() * 4 + PROBE_EVENTS.len(); + let mut buf = vec![0u64; need]; + let written = snapshot(buf.as_mut_ptr(), need); + assert_eq!( + written, need, + "probe wrote {written} u64s, this harness expected {need} \ + — PROBE_PHASES/PROBE_EVENTS are out of sync with dispatch_probe" + ); + Some(buf) + } +} + +/// Print allocations and bytes per phase, normalised per `Run` and per node. +fn report_probe(case: &str, nodes: usize, runs: usize, buf: &[u64]) { + let np = PROBE_PHASES.len(); + let (calls, ns) = (&buf[..np], &buf[np..2 * np]); + let (allocs, bytes) = (&buf[2 * np..3 * np], &buf[3 * np..4 * np]); + let events = &buf[4 * np..]; + let per_run = runs.max(1) as f64; + let per_node = (runs.max(1) * nodes.max(1)) as f64; + println!("# probe {case}: nodes={nodes} runs={runs}"); + println!("# phase,calls_per_run,ns_per_run,allocs_per_run,allocs_per_node,bytes_per_run"); + for (i, name) in PROBE_PHASES.iter().enumerate() { + println!( + "# {name},{:.2},{:.0},{:.3},{:.3},{:.0}", + calls[i] as f64 / per_run, + ns[i] as f64 / per_run, + allocs[i] as f64 / per_run, + allocs[i] as f64 / per_node, + bytes[i] as f64 / per_run, + ); + } + let total: u64 = allocs.iter().sum(); + let total_bytes: u64 = bytes.iter().sum(); + println!( + "# TOTAL_attributed,,,{:.3},{:.3},{:.0}", + total as f64 / per_run, + total as f64 / per_node, + total_bytes as f64 / per_run + ); + for (i, name) in PROBE_EVENTS.iter().enumerate() { + println!("# event {name},{:.3}/run", events[i] as f64 / per_run); + } +} diff --git a/crates/onnx-runtime-ep-plugin/src/dispatch_probe.rs b/crates/onnx-runtime-ep-plugin/src/dispatch_probe.rs index 3a3a923641..7b5c5fe1c2 100644 --- a/crates/onnx-runtime-ep-plugin/src/dispatch_probe.rs +++ b/crates/onnx-runtime-ep-plugin/src/dispatch_probe.rs @@ -199,6 +199,16 @@ pub struct Counters { pub phase_ns: [u64; Phase::COUNT], /// Event tallies. pub events: [u64; Event::COUNT], + /// Heap allocations made while each phase was the innermost open one. + /// + /// Only populated when a [`CountingAllocator`] is installed as the global + /// allocator; otherwise zero. Unlike the hand-placed + /// [`Event::DispatchAlloc`] tally this is exhaustive -- it sees every + /// allocation, including ones inside `Vec` growth, `format!`, and code we + /// did not write -- so it is a total rather than a lower bound. + pub phase_allocs: [u64; Phase::COUNT], + /// Bytes requested by those allocations. + pub phase_alloc_bytes: [u64; Phase::COUNT], } impl Counters { @@ -213,6 +223,9 @@ impl Counters { for i in 0..Phase::COUNT { d.phase_calls[i] = self.phase_calls[i].saturating_sub(earlier.phase_calls[i]); d.phase_ns[i] = self.phase_ns[i].saturating_sub(earlier.phase_ns[i]); + d.phase_allocs[i] = self.phase_allocs[i].saturating_sub(earlier.phase_allocs[i]); + d.phase_alloc_bytes[i] = + self.phase_alloc_bytes[i].saturating_sub(earlier.phase_alloc_bytes[i]); } for i in 0..Event::COUNT { d.events[i] = self.events[i].saturating_sub(earlier.events[i]); @@ -291,11 +304,48 @@ mod imp { const { [const { Cell::new(0) }; Phase::COUNT] }; static TL_EVENTS: [Cell; Event::COUNT] = const { [const { Cell::new(0) }; Event::COUNT] }; + static TL_PHASE_ALLOCS: [Cell; Phase::COUNT] = + const { [const { Cell::new(0) }; Phase::COUNT] }; + static TL_PHASE_ALLOC_BYTES: [Cell; Phase::COUNT] = + const { [const { Cell::new(0) }; Phase::COUNT] }; + } + + // Which phase is currently open on this thread, or `NO_PHASE`. + // + // Read by `CountingAllocator` on every allocation, so it must not allocate + // itself: a `Cell` with a `const` initialiser and no `Drop` compiles to + // a plain TLS slot with no lazy initialisation. + thread_local! { + static TL_CURRENT_PHASE: Cell = const { Cell::new(NO_PHASE) }; } + /// Sentinel for "no phase is open", so allocations outside dispatch are + /// attributed to nothing rather than to phase 0. + pub const NO_PHASE: u8 = u8::MAX; + static G_PHASE_CALLS: [AtomicU64; Phase::COUNT] = [const { AtomicU64::new(0) }; Phase::COUNT]; static G_PHASE_NS: [AtomicU64; Phase::COUNT] = [const { AtomicU64::new(0) }; Phase::COUNT]; static G_EVENTS: [AtomicU64; Event::COUNT] = [const { AtomicU64::new(0) }; Event::COUNT]; + static G_PHASE_ALLOCS: [AtomicU64; Phase::COUNT] = [const { AtomicU64::new(0) }; Phase::COUNT]; + static G_PHASE_ALLOC_BYTES: [AtomicU64; Phase::COUNT] = + [const { AtomicU64::new(0) }; Phase::COUNT]; + + /// Attribute one allocation of `bytes` to whichever phase is open. + /// + /// Called from the global allocator, so it takes the thread-local slot with + /// `try_with`: during thread teardown the TLS may already be gone, and a + /// panic out of `alloc` would be considerably worse than a lost count. + pub fn record_alloc(bytes: u64) { + let phase = TL_CURRENT_PHASE.try_with(Cell::get).unwrap_or(NO_PHASE); + if phase == NO_PHASE { + return; + } + let i = phase as usize; + let _ = TL_PHASE_ALLOCS.try_with(|a| a[i].set(a[i].get().wrapping_add(1))); + let _ = TL_PHASE_ALLOC_BYTES.try_with(|a| a[i].set(a[i].get().wrapping_add(bytes))); + G_PHASE_ALLOCS[i].fetch_add(1, Ordering::Relaxed); + G_PHASE_ALLOC_BYTES[i].fetch_add(bytes, Ordering::Relaxed); + } /// Whether to also accumulate wall time per phase. /// @@ -312,6 +362,14 @@ mod imp { /// Guard returned by [`Phase::enter`]; closes the phase when dropped. pub struct PhaseGuard { phase: Phase, + /// The phase that was open when this one started, restored on close. + /// + /// Phases nest -- a guard closes at scope exit, not at an early + /// `return`, so `StatusCrossing` opens inside whatever was live. Saving + /// and restoring means an allocation inside the inner phase is + /// attributed to it and the outer phase resumes afterwards, rather than + /// everything after the first nested phase being lost. + outer: u8, start: Option, } @@ -323,6 +381,7 @@ mod imp { impl Drop for PhaseGuard { fn drop(&mut self) { + let _ = TL_CURRENT_PHASE.try_with(|c| c.set(self.outer)); if let Some(t) = self.start { let ns = t.elapsed().as_nanos() as u64; let i = self.phase as usize; @@ -338,8 +397,10 @@ mod imp { let i = self as usize; TL_PHASE_CALLS.with(|a| a[i].set(a[i].get().wrapping_add(1))); G_PHASE_CALLS[i].fetch_add(1, Ordering::Relaxed); + let outer = TL_CURRENT_PHASE.with(|c| c.replace(i as u8)); PhaseGuard { phase: self, + outer, start: timing_enabled().then(Instant::now), } } @@ -376,6 +437,16 @@ mod imp { *dst = src.get(); } }); + TL_PHASE_ALLOCS.with(|a| { + for (dst, src) in c.phase_allocs.iter_mut().zip(a) { + *dst = src.get(); + } + }); + TL_PHASE_ALLOC_BYTES.with(|a| { + for (dst, src) in c.phase_alloc_bytes.iter_mut().zip(a) { + *dst = src.get(); + } + }); c } @@ -396,6 +467,12 @@ mod imp { for (dst, src) in c.events.iter_mut().zip(&G_EVENTS) { *dst = src.load(Ordering::Relaxed); } + for (dst, src) in c.phase_allocs.iter_mut().zip(&G_PHASE_ALLOCS) { + *dst = src.load(Ordering::Relaxed); + } + for (dst, src) in c.phase_alloc_bytes.iter_mut().zip(&G_PHASE_ALLOC_BYTES) { + *dst = src.load(Ordering::Relaxed); + } c } @@ -409,6 +486,8 @@ mod imp { TL_PHASE_CALLS.with(|a| a.iter().for_each(|c| c.set(0))); TL_PHASE_NS.with(|a| a.iter().for_each(|c| c.set(0))); TL_EVENTS.with(|a| a.iter().for_each(|c| c.set(0))); + TL_PHASE_ALLOCS.with(|a| a.iter().for_each(|c| c.set(0))); + TL_PHASE_ALLOC_BYTES.with(|a| a.iter().for_each(|c| c.set(0))); for c in &G_PHASE_CALLS { c.store(0, Ordering::Relaxed); } @@ -418,6 +497,12 @@ mod imp { for c in &G_EVENTS { c.store(0, Ordering::Relaxed); } + for c in &G_PHASE_ALLOCS { + c.store(0, Ordering::Relaxed); + } + for c in &G_PHASE_ALLOC_BYTES { + c.store(0, Ordering::Relaxed); + } } /// Whether this build has the probe compiled in. Always `true` here. @@ -477,6 +562,11 @@ mod imp { #[inline(always)] pub fn reset() {} + /// No-op in a production build: nothing tracks the open phase, so there is + /// nothing to attribute an allocation to. + #[inline(always)] + pub fn record_alloc(_bytes: u64) {} + /// No-op in a production build. #[inline(always)] pub fn timing_enabled() -> bool { @@ -537,6 +627,7 @@ impl CountingAllocator { unsafe impl std::alloc::GlobalAlloc for CountingAllocator { unsafe fn alloc(&self, layout: std::alloc::Layout) -> *mut u8 { count(Event::DispatchAlloc); + imp::record_alloc(layout.size() as u64); unsafe { self.inner.alloc(layout) } } @@ -546,11 +637,13 @@ unsafe impl std::alloc::GlobalAlloc for CountingAllo unsafe fn alloc_zeroed(&self, layout: std::alloc::Layout) -> *mut u8 { count(Event::DispatchAlloc); + imp::record_alloc(layout.size() as u64); unsafe { self.inner.alloc_zeroed(layout) } } unsafe fn realloc(&self, ptr: *mut u8, layout: std::alloc::Layout, new_size: usize) -> *mut u8 { count(Event::DispatchAlloc); + imp::record_alloc(new_size as u64); unsafe { self.inner.realloc(ptr, layout, new_size) } } } @@ -576,16 +669,22 @@ pub fn ort_calls(n: u64) { /// process-wide totals rather than the calling thread's. In-process users /// should prefer [`snapshot`], which is isolated per thread. /// -/// Writes `phase_calls`, then `phase_ns`, then `events`, and returns the number -/// of `u64`s written, or 0 if `out` is null or `len` is too small. The required -/// length is `Phase::COUNT * 2 + Event::COUNT`. +/// Writes `phase_calls`, `phase_ns`, `phase_allocs`, `phase_alloc_bytes`, then +/// `events`, and returns the number of `u64`s written, or 0 if `out` is null or +/// `len` is too small. The required length is [`SNAPSHOT_LEN`]. /// /// # Safety /// /// `out` must be null or point to `len` writable `u64`s. +/// Number of `u64`s [`nxrt_dispatch_probe_snapshot`] writes. +/// +/// Exported so the cdylib harness sizes its buffer from the same expression the +/// writer uses, rather than from a number copied into a test. +pub const SNAPSHOT_LEN: usize = Phase::COUNT * 4 + Event::COUNT; + #[unsafe(no_mangle)] pub unsafe extern "C" fn nxrt_dispatch_probe_snapshot(out: *mut u64, len: usize) -> usize { - let need = Phase::COUNT * 2 + Event::COUNT; + let need = SNAPSHOT_LEN; if out.is_null() || len < need { return 0; } @@ -602,6 +701,14 @@ pub unsafe extern "C" fn nxrt_dispatch_probe_snapshot(out: *mut u64, len: usize) p.write(v); p = p.add(1); } + for v in c.phase_allocs { + p.write(v); + p = p.add(1); + } + for v in c.phase_alloc_bytes { + p.write(v); + p = p.add(1); + } for v in c.events { p.write(v); p = p.add(1); @@ -751,7 +858,7 @@ mod tests { /// buffer it would overrun rather than writing past the end. #[test] fn c_snapshot_refuses_a_short_buffer() { - let need = Phase::COUNT * 2 + Event::COUNT; + let need = SNAPSHOT_LEN; let mut buf = vec![0u64; need]; // SAFETY: `buf` has exactly `need` writable u64s. assert_eq!( diff --git a/crates/onnx-runtime-ep-plugin/tests/alloc_attribution.rs b/crates/onnx-runtime-ep-plugin/tests/alloc_attribution.rs new file mode 100644 index 0000000000..05d90db9b0 --- /dev/null +++ b/crates/onnx-runtime-ep-plugin/tests/alloc_attribution.rs @@ -0,0 +1,185 @@ +//! Per-phase allocation attribution, exercised with the counting allocator +//! actually installed. +//! +//! This lives in its own test binary because `#[global_allocator]` is a +//! per-binary choice: installing it in the library's unit tests would put a +//! counting wrapper under every other test in the crate. Here it is the point. +//! +//! What this pins is the *attribution*, not a particular number of +//! allocations — the numbers belong with the code that allocates. The property +//! is that an allocation is charged to the innermost phase that was open when +//! it happened, and to nothing at all when no phase is open. + +use onnx_runtime_ep_plugin::dispatch_probe::{self, CountingAllocator, Phase}; + +#[global_allocator] +static ALLOC: CountingAllocator = CountingAllocator::new(std::alloc::System); + +/// Allocate `bytes` in a way the optimiser cannot elide. +fn allocate(bytes: usize) -> Vec { + let mut v: Vec = Vec::with_capacity(bytes); + v.push(1); + std::hint::black_box(&v); + v +} + +#[test] +fn an_allocation_is_charged_to_the_phase_that_is_open() { + let before = dispatch_probe::snapshot(); + let guard = Phase::Allocate.enter(); + let held = allocate(4096); + guard.end(); + let d = dispatch_probe::snapshot().since(&before); + drop(held); + + if !dispatch_probe::compiled_in() { + assert_eq!( + d.phase_allocs, + [0; Phase::COUNT], + "a production build must attribute nothing" + ); + return; + } + + assert!( + d.phase_allocs[Phase::Allocate as usize] >= 1, + "the allocation inside the phase was not charged to it: {:?}", + d.phase_allocs + ); + assert!( + d.phase_alloc_bytes[Phase::Allocate as usize] >= 4096, + "bytes were not recorded: {:?}", + d.phase_alloc_bytes + ); + for (i, n) in d.phase_allocs.iter().enumerate() { + if i != Phase::Allocate as usize { + assert_eq!(*n, 0, "phase {i} was charged for another phase's work"); + } + } +} + +#[test] +fn an_allocation_outside_every_phase_is_charged_to_none() { + let before = dispatch_probe::snapshot(); + let held = allocate(8192); + let d = dispatch_probe::snapshot().since(&before); + drop(held); + + assert_eq!( + d.phase_allocs, + [0; Phase::COUNT], + "work outside dispatch must not be attributed to a dispatch phase" + ); + assert_eq!(d.phase_alloc_bytes, [0; Phase::COUNT]); +} + +/// Phases nest — a guard closes at scope exit, not at an early `return`, so +/// `StatusCrossing` opens inside whatever phase was live. The inner phase must +/// take its own allocations *and hand the outer one back*, or everything after +/// the first nested phase is charged to the wrong place (or to nothing). +#[test] +fn a_nested_phase_takes_its_own_allocations_and_restores_the_outer_one() { + if !dispatch_probe::compiled_in() { + return; + } + let before = dispatch_probe::snapshot(); + + let outer = Phase::TensorBind.enter(); + let a = allocate(1024); + { + let inner = Phase::StatusCrossing.enter(); + let b = allocate(2048); + inner.end(); + drop(b); + } + let c = allocate(512); + outer.end(); + drop((a, c)); + + let d = dispatch_probe::snapshot().since(&before); + assert_eq!( + d.phase_allocs[Phase::StatusCrossing as usize], + 1, + "the nested phase should own exactly the allocation made inside it" + ); + assert_eq!( + d.phase_allocs[Phase::TensorBind as usize], + 2, + "the outer phase must resume after the nested one closes: {:?}", + d.phase_allocs + ); + assert!(d.phase_alloc_bytes[Phase::StatusCrossing as usize] >= 2048); +} + +/// The attribution must survive a phase that ends by unwinding out of scope +/// rather than by an explicit `end()`, since that is how every early-return +/// path in `compute_execute` closes its phase. +#[test] +fn a_phase_closed_by_scope_exit_still_restores_the_outer_phase() { + if !dispatch_probe::compiled_in() { + return; + } + let before = dispatch_probe::snapshot(); + let outer = Phase::DispatchLookup.enter(); + { + let _inner = Phase::Allocate.enter(); + let b = allocate(64); + drop(b); + } + let c = allocate(128); + outer.end(); + drop(c); + + let d = dispatch_probe::snapshot().since(&before); + assert_eq!(d.phase_allocs[Phase::Allocate as usize], 1); + assert_eq!( + d.phase_allocs[Phase::DispatchLookup as usize], + 1, + "scope-exit close did not restore the outer phase: {:?}", + d.phase_allocs + ); +} + +/// Attribution is thread-local: two threads dispatching at once must not be +/// able to charge each other's allocations, which is what makes a per-phase +/// figure meaningful under concurrent `Run`. +#[test] +fn concurrent_phases_do_not_contaminate_each_other() { + if !dispatch_probe::compiled_in() { + return; + } + let other = std::thread::spawn(|| { + let before = dispatch_probe::snapshot(); + let g = Phase::KernelInvoke.enter(); + let mut held = Vec::new(); + for _ in 0..50 { + held.push(allocate(256)); + } + g.end(); + let d = dispatch_probe::snapshot().since(&before); + drop(held); + d + }); + + let before = dispatch_probe::snapshot(); + let g = Phase::MetadataQuery.enter(); + let held = allocate(256); + g.end(); + let mine = dispatch_probe::snapshot().since(&before); + drop(held); + + let theirs = other.join().expect("worker panicked"); + + assert_eq!( + mine.phase_allocs[Phase::KernelInvoke as usize], + 0, + "this thread was charged for the other thread's phase" + ); + assert_eq!( + theirs.phase_allocs[Phase::MetadataQuery as usize], + 0, + "the other thread was charged for this one's phase" + ); + assert_eq!(mine.phase_allocs[Phase::MetadataQuery as usize], 1); + assert!(theirs.phase_allocs[Phase::KernelInvoke as usize] >= 50); +} From f067fdb974dc72a446cbc3da146b378df5a67b0f Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Wed, 19 Aug 2026 05:29:55 +0000 Subject: [PATCH 7/7] Account for every dispatch allocation, and stop the harness mislabelling them The first live attribution run said four allocations per node and gave no way to tell whether that was all of them. It was a third of them. Allocations made with no phase open were being discarded rather than counted, so the per-phase rows could not be checked against any total. They now land in an explicit UNATTRIBUTED bucket, which makes the columns add up: 12.11 allocations per node on a 100-node chain, of which 4 are inside named phases and 8.03 are in the routed loop between them. The hand-placed DispatchAlloc tally independently reads 12.19, which is the cross-check that this is now complete. The harness also kept its own copy of the phase order, and the copy was wrong: it labelled TensorBind's allocations "OutputMeta" and Allocate's "TensorBind". Every number in that table was correct and attached to the wrong phase, which is the kind of error that survives review. The names now come from the library through a new nxrt_dispatch_probe_phase_name export, so there is one source of truth, and a test pins the export against Phase::name. Falsified: renaming one exported phase fails the pinning test; reverting the unattributed bucket fails the outside-a-phase test, which now asserts the allocation is visible rather than absent. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../tests/plugin_ort_e2e.rs | 102 +++++++++++------ .../src/dispatch_probe.rs | 105 +++++++++++++++--- .../tests/alloc_attribution.rs | 38 +++++-- 3 files changed, 190 insertions(+), 55 deletions(-) diff --git a/crates/onnx-runtime-ep-cpu-plugin/tests/plugin_ort_e2e.rs b/crates/onnx-runtime-ep-cpu-plugin/tests/plugin_ort_e2e.rs index 0d7132d5fe..72fed4278a 100644 --- a/crates/onnx-runtime-ep-cpu-plugin/tests/plugin_ort_e2e.rs +++ b/crates/onnx-runtime-ep-cpu-plugin/tests/plugin_ort_e2e.rs @@ -6455,14 +6455,26 @@ fn plugin_path_ab_vs_plain_ort() { // Allocation attribution runs before the timed loop so its // counter updates never land inside a measured iteration. if run_ours - && let Some(buf) = - probe_dispatch(api, session, &input_name_ptrs, &values, &output_name_ptrs, 64) + && let Some((buf, names)) = probe_dispatch( + api, + session, + &input_name_ptrs, + &values, + &output_name_ptrs, + 64, + ) { - // Node count comes from ORT's assignment list, not from - // the case definition: per-node figures must divide by the - // nodes this EP actually received, and fusion means that - // is not the same as the number of `Run` callbacks. - report_probe(case.name, info.ops_on_our_ep().len().max(1), 64, &buf); + // Node count comes from ORT's assignment list, not from the + // case definition: per-node figures must divide by the nodes + // this EP actually received, and fusion means that is not the + // same as the number of `Run` callbacks. + report_probe( + case.name, + info.ops_on_our_ep().len().max(1), + 64, + &buf, + &names, + ); } // Interleave one iteration each so a drifting host load lands on @@ -7441,21 +7453,34 @@ fn the_native_activation_family_executes_locally_and_matches_ort_numerics() { ); } -/// Phase names, in the order `nxrt_dispatch_probe_snapshot` writes them. -/// -/// Duplicated from `dispatch_probe::Phase` because this harness talks to the -/// cdylib over the C ABI, not through the Rust type. `probe_phase_names` is -/// checked against the returned length, so a phase added on the library side -/// shows up as a loud mismatch here rather than as a silently shifted table. -const PROBE_PHASES: &[&str] = &[ - "OrtCallback", - "InputMeta", - "OutputMeta", - "TensorBind", - "DispatchLookup", - "KernelInvoke", - "StatusCrossing", -]; +/// Bucket names, read from the library rather than copied. +/// +/// A hand-maintained copy of this list drifted from the enum and mislabelled +/// two rows of the attribution table -- every number was right and attached to +/// the wrong phase. Asking the cdylib removes the second source of truth. +fn probe_phase_names(lib: &libloading::Library) -> Vec { + // SAFETY: the export is `extern "C"` with this signature and returns either + // null or a 'static NUL-terminated string owned by the library. + unsafe { + let name_of: libloading::Symbol< + '_, + unsafe extern "C" fn(usize) -> *const std::os::raw::c_char, + > = match lib.get(b"nxrt_dispatch_probe_phase_name") { + Ok(f) => f, + Err(_) => return Vec::new(), + }; + let mut out = Vec::new(); + for i in 0.. { + let p = name_of(i); + if p.is_null() { + break; + } + out.push(std::ffi::CStr::from_ptr(p).to_string_lossy().into_owned()); + } + out + } +} + const PROBE_EVENTS: &[&str] = &[ "OrtFfiCall", "DispatchAlloc", @@ -7496,7 +7521,7 @@ unsafe fn probe_dispatch( values: &[*const ort::OrtValue], output_names: &[*const std::os::raw::c_char], runs: usize, -) -> Option> { +) -> Option<(Vec, Vec)> { if std::env::var("NXRT_MM_BENCH_PROBE").unwrap_or_default() != "1" { return None; } @@ -7510,29 +7535,35 @@ unsafe fn probe_dispatch( lib.get(b"nxrt_dispatch_probe_snapshot").ok()?; reset(); bench_runs(api, session, input_names, values, output_names, runs); - let need = PROBE_PHASES.len() * 4 + PROBE_EVENTS.len(); + let buckets = probe_phase_names(lib).len(); + assert!(buckets > 0, "cdylib exports no phase names"); + let need = (buckets - 1) * 2 + buckets * 2 + PROBE_EVENTS.len(); let mut buf = vec![0u64; need]; let written = snapshot(buf.as_mut_ptr(), need); assert_eq!( written, need, "probe wrote {written} u64s, this harness expected {need} \ - — PROBE_PHASES/PROBE_EVENTS are out of sync with dispatch_probe" + — PROBE_EVENTS is out of sync with dispatch_probe" ); - Some(buf) + Some((buf, probe_phase_names(lib))) } } /// Print allocations and bytes per phase, normalised per `Run` and per node. -fn report_probe(case: &str, nodes: usize, runs: usize, buf: &[u64]) { - let np = PROBE_PHASES.len(); +fn report_probe(case: &str, nodes: usize, runs: usize, buf: &[u64], names: &[String]) { + let nb = names.len(); + let np = nb - 1; let (calls, ns) = (&buf[..np], &buf[np..2 * np]); - let (allocs, bytes) = (&buf[2 * np..3 * np], &buf[3 * np..4 * np]); - let events = &buf[4 * np..]; + let (allocs, bytes) = ( + &buf[2 * np..2 * np + nb], + &buf[2 * np + nb..2 * np + 2 * nb], + ); + let events = &buf[2 * np + 2 * nb..]; let per_run = runs.max(1) as f64; let per_node = (runs.max(1) * nodes.max(1)) as f64; println!("# probe {case}: nodes={nodes} runs={runs}"); println!("# phase,calls_per_run,ns_per_run,allocs_per_run,allocs_per_node,bytes_per_run"); - for (i, name) in PROBE_PHASES.iter().enumerate() { + for (i, name) in names.iter().enumerate().take(np) { println!( "# {name},{:.2},{:.0},{:.3},{:.3},{:.0}", calls[i] as f64 / per_run, @@ -7542,6 +7573,15 @@ fn report_probe(case: &str, nodes: usize, runs: usize, buf: &[u64]) { bytes[i] as f64 / per_run, ); } + // The unattributed bucket is the point of the table, not a footnote: it is + // what says whether the per-phase rows are the whole story. + println!( + "# {},,,{:.3},{:.3},{:.0}", + names[nb - 1], + allocs[nb - 1] as f64 / per_run, + allocs[nb - 1] as f64 / per_node, + bytes[nb - 1] as f64 / per_run, + ); let total: u64 = allocs.iter().sum(); let total_bytes: u64 = bytes.iter().sum(); println!( diff --git a/crates/onnx-runtime-ep-plugin/src/dispatch_probe.rs b/crates/onnx-runtime-ep-plugin/src/dispatch_probe.rs index 7b5c5fe1c2..99c9cd1182 100644 --- a/crates/onnx-runtime-ep-plugin/src/dispatch_probe.rs +++ b/crates/onnx-runtime-ep-plugin/src/dispatch_probe.rs @@ -201,16 +201,29 @@ pub struct Counters { pub events: [u64; Event::COUNT], /// Heap allocations made while each phase was the innermost open one. /// + /// Length is [`ALLOC_BUCKETS`], not `Phase::COUNT`: the last slot is + /// [`UNATTRIBUTED`], charged when an allocation happens with no phase open. + /// Without that slot the table silently omits allocations rather than + /// showing them, which is the opposite of what an attribution pass needs -- + /// the first live reading had four allocations per node inside phases and + /// gave no hint whether that was all of them. + /// /// Only populated when a [`CountingAllocator`] is installed as the global /// allocator; otherwise zero. Unlike the hand-placed /// [`Event::DispatchAlloc`] tally this is exhaustive -- it sees every /// allocation, including ones inside `Vec` growth, `format!`, and code we /// did not write -- so it is a total rather than a lower bound. - pub phase_allocs: [u64; Phase::COUNT], - /// Bytes requested by those allocations. - pub phase_alloc_bytes: [u64; Phase::COUNT], + pub phase_allocs: [u64; ALLOC_BUCKETS], + /// Bytes requested by those allocations, same bucketing. + pub phase_alloc_bytes: [u64; ALLOC_BUCKETS], } +/// Index of the bucket for allocations made with no phase open. +pub const UNATTRIBUTED: usize = Phase::COUNT; + +/// Number of allocation buckets: one per phase, plus [`UNATTRIBUTED`]. +pub const ALLOC_BUCKETS: usize = Phase::COUNT + 1; + impl Counters { /// This reading minus an earlier one — what happened in between. /// @@ -223,6 +236,8 @@ impl Counters { for i in 0..Phase::COUNT { d.phase_calls[i] = self.phase_calls[i].saturating_sub(earlier.phase_calls[i]); d.phase_ns[i] = self.phase_ns[i].saturating_sub(earlier.phase_ns[i]); + } + for i in 0..ALLOC_BUCKETS { d.phase_allocs[i] = self.phase_allocs[i].saturating_sub(earlier.phase_allocs[i]); d.phase_alloc_bytes[i] = self.phase_alloc_bytes[i].saturating_sub(earlier.phase_alloc_bytes[i]); @@ -262,7 +277,7 @@ impl Counters { self.event(e) as f64 / runs as f64 )); } - for p in Phase::ALL { + for p in Phase::ALL.into_iter() { s.push_str(&format!( " {:<16} {:>10} calls {:>12} ns ({:.0} ns/run)\n", p.name(), @@ -277,7 +292,7 @@ impl Counters { #[cfg(feature = "dispatch_probe")] mod imp { - use super::{Counters, Event, Phase}; + use super::{ALLOC_BUCKETS, Counters, Event, Phase, UNATTRIBUTED}; use std::cell::Cell; use std::sync::atomic::{AtomicU64, Ordering}; use std::time::Instant; @@ -304,10 +319,10 @@ mod imp { const { [const { Cell::new(0) }; Phase::COUNT] }; static TL_EVENTS: [Cell; Event::COUNT] = const { [const { Cell::new(0) }; Event::COUNT] }; - static TL_PHASE_ALLOCS: [Cell; Phase::COUNT] = - const { [const { Cell::new(0) }; Phase::COUNT] }; - static TL_PHASE_ALLOC_BYTES: [Cell; Phase::COUNT] = - const { [const { Cell::new(0) }; Phase::COUNT] }; + static TL_PHASE_ALLOCS: [Cell; ALLOC_BUCKETS] = + const { [const { Cell::new(0) }; ALLOC_BUCKETS] }; + static TL_PHASE_ALLOC_BYTES: [Cell; ALLOC_BUCKETS] = + const { [const { Cell::new(0) }; ALLOC_BUCKETS] }; } // Which phase is currently open on this thread, or `NO_PHASE`. @@ -326,9 +341,10 @@ mod imp { static G_PHASE_CALLS: [AtomicU64; Phase::COUNT] = [const { AtomicU64::new(0) }; Phase::COUNT]; static G_PHASE_NS: [AtomicU64; Phase::COUNT] = [const { AtomicU64::new(0) }; Phase::COUNT]; static G_EVENTS: [AtomicU64; Event::COUNT] = [const { AtomicU64::new(0) }; Event::COUNT]; - static G_PHASE_ALLOCS: [AtomicU64; Phase::COUNT] = [const { AtomicU64::new(0) }; Phase::COUNT]; - static G_PHASE_ALLOC_BYTES: [AtomicU64; Phase::COUNT] = - [const { AtomicU64::new(0) }; Phase::COUNT]; + static G_PHASE_ALLOCS: [AtomicU64; ALLOC_BUCKETS] = + [const { AtomicU64::new(0) }; ALLOC_BUCKETS]; + static G_PHASE_ALLOC_BYTES: [AtomicU64; ALLOC_BUCKETS] = + [const { AtomicU64::new(0) }; ALLOC_BUCKETS]; /// Attribute one allocation of `bytes` to whichever phase is open. /// @@ -337,10 +353,11 @@ mod imp { /// panic out of `alloc` would be considerably worse than a lost count. pub fn record_alloc(bytes: u64) { let phase = TL_CURRENT_PHASE.try_with(Cell::get).unwrap_or(NO_PHASE); - if phase == NO_PHASE { - return; - } - let i = phase as usize; + let i = if phase == NO_PHASE { + UNATTRIBUTED + } else { + phase as usize + }; let _ = TL_PHASE_ALLOCS.try_with(|a| a[i].set(a[i].get().wrapping_add(1))); let _ = TL_PHASE_ALLOC_BYTES.try_with(|a| a[i].set(a[i].get().wrapping_add(bytes))); G_PHASE_ALLOCS[i].fetch_add(1, Ordering::Relaxed); @@ -676,11 +693,42 @@ pub fn ort_calls(n: u64) { /// # Safety /// /// `out` must be null or point to `len` writable `u64`s. +/// Name of allocation bucket `index`, or null if out of range. +/// +/// Exists because the harness must not keep its own copy of the phase order. +/// It did, briefly, and the copy was wrong: the table labelled `TensorBind`'s +/// allocations "OutputMeta" and `Allocate`'s "TensorBind", which is the kind of +/// error that survives review because every number in it looks plausible. +/// +/// Index `Phase::COUNT` is the unattributed bucket. +/// +/// # Safety +/// +/// The returned pointer is to a `'static` NUL-terminated string and must not be +/// freed. It stays valid for the lifetime of the library. +#[unsafe(no_mangle)] +pub extern "C" fn nxrt_dispatch_probe_phase_name(index: usize) -> *const std::os::raw::c_char { + const NAMES: [&core::ffi::CStr; ALLOC_BUCKETS] = [ + c"callback_entry", + c"metadata_query", + c"tensor_bind", + c"allocate", + c"dispatch_lookup", + c"kernel_invoke", + c"status_crossing", + c"unattributed", + ]; + match NAMES.get(index) { + Some(n) => n.as_ptr(), + None => core::ptr::null(), + } +} + /// Number of `u64`s [`nxrt_dispatch_probe_snapshot`] writes. /// /// Exported so the cdylib harness sizes its buffer from the same expression the /// writer uses, rather than from a number copied into a test. -pub const SNAPSHOT_LEN: usize = Phase::COUNT * 4 + Event::COUNT; +pub const SNAPSHOT_LEN: usize = Phase::COUNT * 2 + ALLOC_BUCKETS * 2 + Event::COUNT; #[unsafe(no_mangle)] pub unsafe extern "C" fn nxrt_dispatch_probe_snapshot(out: *mut u64, len: usize) -> usize { @@ -854,6 +902,29 @@ mod tests { } } + /// The exported names must match `Phase::name`, in order. + /// + /// Two sources of truth for the same list is exactly how the harness came + /// to mislabel two phases; this makes the duplicate a checked one. + #[test] + fn exported_phase_names_match_the_enum() { + for p in Phase::ALL.into_iter() { + let ptr = nxrt_dispatch_probe_phase_name(p as usize); + assert!(!ptr.is_null(), "{} has no exported name", p.name()); + // SAFETY: the export returns a 'static NUL-terminated string. + let got = unsafe { core::ffi::CStr::from_ptr(ptr) }; + assert_eq!(got.to_str().unwrap(), p.name(), "name mismatch"); + } + let last = nxrt_dispatch_probe_phase_name(UNATTRIBUTED); + assert!(!last.is_null()); + // SAFETY: as above. + assert_eq!( + unsafe { core::ffi::CStr::from_ptr(last) }.to_str().unwrap(), + "unattributed" + ); + assert!(nxrt_dispatch_probe_phase_name(ALLOC_BUCKETS).is_null()); + } + /// The C entry point is what the cdylib harness uses; it must refuse a /// buffer it would overrun rather than writing past the end. #[test] diff --git a/crates/onnx-runtime-ep-plugin/tests/alloc_attribution.rs b/crates/onnx-runtime-ep-plugin/tests/alloc_attribution.rs index 05d90db9b0..2be4b06b12 100644 --- a/crates/onnx-runtime-ep-plugin/tests/alloc_attribution.rs +++ b/crates/onnx-runtime-ep-plugin/tests/alloc_attribution.rs @@ -35,7 +35,7 @@ fn an_allocation_is_charged_to_the_phase_that_is_open() { if !dispatch_probe::compiled_in() { assert_eq!( d.phase_allocs, - [0; Phase::COUNT], + [0; dispatch_probe::ALLOC_BUCKETS], "a production build must attribute nothing" ); return; @@ -58,19 +58,43 @@ fn an_allocation_is_charged_to_the_phase_that_is_open() { } } +/// An allocation with no phase open must be *visible* rather than dropped. +/// +/// Charging it to no phase at all was the original behaviour, and it made the +/// table unreadable: the first live routed dispatch showed four allocations per +/// node inside phases with no way to tell whether that was all of them or a +/// quarter of them. It has to land in `UNATTRIBUTED` so the columns add up to +/// the process total, while still never being blamed on a phase. #[test] -fn an_allocation_outside_every_phase_is_charged_to_none() { +fn an_allocation_outside_every_phase_is_charged_to_the_unattributed_bucket() { let before = dispatch_probe::snapshot(); let held = allocate(8192); let d = dispatch_probe::snapshot().since(&before); drop(held); - assert_eq!( - d.phase_allocs, - [0; Phase::COUNT], - "work outside dispatch must not be attributed to a dispatch phase" + for (i, n) in d.phase_allocs.iter().enumerate() { + if i != dispatch_probe::UNATTRIBUTED { + assert_eq!( + *n, 0, + "work outside dispatch was attributed to phase {i}: {:?}", + d.phase_allocs + ); + } + } + + if !dispatch_probe::compiled_in() { + return; + } + assert!( + d.phase_allocs[dispatch_probe::UNATTRIBUTED] >= 1, + "an allocation outside every phase went uncounted: {:?}", + d.phase_allocs + ); + assert!( + d.phase_alloc_bytes[dispatch_probe::UNATTRIBUTED] >= 8192, + "unattributed bytes were not recorded: {:?}", + d.phase_alloc_bytes ); - assert_eq!(d.phase_alloc_bytes, [0; Phase::COUNT]); } /// Phases nest — a guard closes at scope exit, not at an early `return`, so