diff --git a/crates/onnx-runtime-ep-cpu-plugin/Cargo.toml b/crates/onnx-runtime-ep-cpu-plugin/Cargo.toml index 3c51cd2381..79c4e2c900 100644 --- a/crates/onnx-runtime-ep-cpu-plugin/Cargo.toml +++ b/crates/onnx-runtime-ep-cpu-plugin/Cargo.toml @@ -20,10 +20,12 @@ crate-type = ["cdylib", "lib"] # resolved graph and the built cdylib's symbol table. mlas = ["onnx-runtime-ep-cpu/mlas"] -# Compile the dispatch probe into the cdylib, so a harness can resolve its -# entry points across the ORT boundary. Never enabled for a shipped build: the -# probe's `no_mangle` exports exist only under this feature precisely so the -# shipped cdylib keeps exporting the ORT plugin ABI and nothing else +# 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, +# and the probe's `no_mangle` exports exist only under this feature precisely so +# the shipped cdylib keeps exporting the ORT plugin ABI and nothing else # (`plugin_export_abi::l1_no_symbol_leakage` pins that). dispatch_probe = ["onnx-runtime-ep-plugin/dispatch_probe"] 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_export_abi.rs b/crates/onnx-runtime-ep-cpu-plugin/tests/plugin_export_abi.rs index 9366f76745..d13e2ec0a8 100644 --- a/crates/onnx-runtime-ep-cpu-plugin/tests/plugin_export_abi.rs +++ b/crates/onnx-runtime-ep-cpu-plugin/tests/plugin_export_abi.rs @@ -627,7 +627,7 @@ fn l1_no_symbol_leakage() { && *name != "nxrt_ep_reset_executed_node_count" && *name != "nxrt_ep_build_features" && *name != "nxrt_ep_persistent_decode_pool_built" - // The dispatch probe is a research build. Its three exports + // The dispatch probe is a research build. Its four exports // exist only under the `dispatch_probe` feature, which no // shipped build sets, so in a production cdylib these are // absent and this arm never fires. Gated on the same `cfg` as @@ -639,6 +639,7 @@ fn l1_no_symbol_leakage() { "nxrt_dispatch_probe_snapshot" | "nxrt_dispatch_probe_reset" | "nxrt_dispatch_probe_available" + | "nxrt_dispatch_probe_phase_name" )) && !name.starts_with("_Z") && !name.starts_with("__rust") 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..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 @@ -6452,6 +6452,31 @@ 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, 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, + &names, + ); + } + // 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 +7452,145 @@ fn the_native_activation_family_executes_locally_and_matches_ort_numerics() { "\n✅ the_native_activation_family_executes_locally_and_matches_ort_numerics: PASSED" ); } + +/// 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", + "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<(Vec, Vec)> { + 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 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_EVENTS is out of sync with dispatch_probe" + ); + 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], 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..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 names.iter().enumerate().take(np) { + 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, + ); + } + // 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!( + "# 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 ab25b7a11d..75d4f352fa 100644 --- a/crates/onnx-runtime-ep-plugin/src/dispatch_probe.rs +++ b/crates/onnx-runtime-ep-plugin/src/dispatch_probe.rs @@ -199,8 +199,31 @@ 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. + /// + /// 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; 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. /// @@ -214,6 +237,11 @@ impl Counters { 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]); + } for i in 0..Event::COUNT { d.events[i] = self.events[i].saturating_sub(earlier.events[i]); } @@ -249,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(), @@ -264,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; @@ -291,11 +319,50 @@ 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; 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`. + // + // 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; 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. + /// + /// 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); + 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); + G_PHASE_ALLOC_BYTES[i].fetch_add(bytes, Ordering::Relaxed); + } /// Whether to also accumulate wall time per phase. /// @@ -312,6 +379,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 +398,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 +414,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 +454,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 +484,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 +503,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 +514,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 +579,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 +644,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 +654,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,24 +686,67 @@ 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`. -/// -/// Only exported when the `dispatch_probe` feature is on. A shipped plugin must -/// export the ORT plugin ABI and nothing else, and a `no_mangle` symbol is not -/// free just because the code behind it is: it survives `--gc-sections`, is -/// interposable, and appears in every dynamic symbol table. Absence *is* the -/// "probe not compiled in" answer, which is what a `dlsym` caller already has to -/// handle. +/// 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. +/// 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. +/// Only compiled under the `dispatch_probe` feature. The shipped cdylib must +/// export the ORT plugin ABI and nothing else, and a `no_mangle` symbol is not +/// free just because its body is: it survives `--gc-sections`, it is +/// interposable, and it lands in every dynamic symbol table. Absence *is* the +/// "not compiled in" answer, which `libloading`'s `Option`-returning `get` +/// already models. +#[cfg(feature = "dispatch_probe")] +#[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 * 2 + ALLOC_BUCKETS * 2 + Event::COUNT; + +/// Only compiled under the `dispatch_probe` feature. The shipped cdylib must +/// export the ORT plugin ABI and nothing else, and a `no_mangle` symbol is not +/// free just because its body is: it survives `--gc-sections`, it is +/// interposable, and it lands in every dynamic symbol table. Absence *is* the +/// "not compiled in" answer, which `libloading`'s `Option`-returning `get` +/// already models. #[cfg(feature = "dispatch_probe")] #[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; } @@ -610,6 +763,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); @@ -619,9 +780,12 @@ pub unsafe extern "C" fn nxrt_dispatch_probe_snapshot(out: *mut u64, len: usize) } /// Zero this thread's dispatch counters, for cdylib callers. -/// -/// Feature-gated for the same reason as -/// [`nxrt_dispatch_probe_snapshot`]: production exports the ORT plugin ABI only. +/// Only compiled under the `dispatch_probe` feature. The shipped cdylib must +/// export the ORT plugin ABI and nothing else, and a `no_mangle` symbol is not +/// free just because its body is: it survives `--gc-sections`, it is +/// interposable, and it lands in every dynamic symbol table. Absence *is* the +/// "not compiled in" answer, which `libloading`'s `Option`-returning `get` +/// already models. #[cfg(feature = "dispatch_probe")] #[unsafe(no_mangle)] pub extern "C" fn nxrt_dispatch_probe_reset() { @@ -632,12 +796,12 @@ pub extern "C" fn nxrt_dispatch_probe_reset() { /// /// 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`. -/// -/// This symbol only exists in a probe build, so resolving it *at all* already -/// answers the question and it always returns 1. It is kept so a caller that -/// resolved it can read a value rather than having to special-case a symbol it -/// looked up successfully, and it still reads `compiled_in()` rather than -/// hard-coding the answer so the two cannot drift apart. +/// Only compiled under the `dispatch_probe` feature. The shipped cdylib must +/// export the ORT plugin ABI and nothing else, and a `no_mangle` symbol is not +/// free just because its body is: it survives `--gc-sections`, it is +/// interposable, and it lands in every dynamic symbol table. Absence *is* the +/// "not compiled in" answer, which `libloading`'s `Option`-returning `get` +/// already models. #[cfg(feature = "dispatch_probe")] #[unsafe(no_mangle)] pub extern "C" fn nxrt_dispatch_probe_available() -> i32 { @@ -766,14 +930,36 @@ mod tests { } } - /// The C entry point is what the cdylib harness uses; it must refuse a - /// buffer it would overrun rather than writing past the end. + /// The exported names must match `Phase::name`, in order. /// - /// Only exists in a probe build, because the entry point only exists there. + /// 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. + #[cfg(feature = "dispatch_probe")] #[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. #[cfg(feature = "dispatch_probe")] + #[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..2be4b06b12 --- /dev/null +++ b/crates/onnx-runtime-ep-plugin/tests/alloc_attribution.rs @@ -0,0 +1,209 @@ +//! 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; dispatch_probe::ALLOC_BUCKETS], + "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"); + } + } +} + +/// 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_the_unattributed_bucket() { + let before = dispatch_probe::snapshot(); + let held = allocate(8192); + let d = dispatch_probe::snapshot().since(&before); + drop(held); + + 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 + ); +} + +/// 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); +}