From 19394915751f0e19644cc80f134c225e31b8d562 Mon Sep 17 00:00:00 2001 From: Copilot Date: Wed, 19 Aug 2026 01:13:06 -0700 Subject: [PATCH 1/3] wip: nxmem memory plugin ABI scaffolding --- Cargo.lock | 12 + Cargo.toml | 12 + crates/onnx-runtime-memory-abi/Cargo.toml | 17 + crates/onnx-runtime-memory-abi/src/lib.rs | 275 +++ crates/onnx-runtime-memory-abi/src/status.rs | 373 ++++ crates/onnx-runtime-memory-abi/src/types.rs | 701 ++++++++ crates/onnx-runtime-memory-abi/src/version.rs | 603 +++++++ crates/onnx-runtime-memory-abi/src/vtable.rs | 945 ++++++++++ crates/onnx-runtime-memory-host/Cargo.toml | 21 + .../onnx-runtime-memory-host/src/allocator.rs | 1566 +++++++++++++++++ crates/onnx-runtime-memory-host/src/error.rs | 143 ++ crates/onnx-runtime-memory-host/src/lib.rs | 0 crates/onnx-runtime-memory-host/src/loader.rs | 648 +++++++ .../onnx-runtime-memory-testplugin/Cargo.toml | 20 + .../onnx-runtime-memory-testplugin/src/lib.rs | 0 15 files changed, 5336 insertions(+) create mode 100644 crates/onnx-runtime-memory-abi/Cargo.toml create mode 100644 crates/onnx-runtime-memory-abi/src/lib.rs create mode 100644 crates/onnx-runtime-memory-abi/src/status.rs create mode 100644 crates/onnx-runtime-memory-abi/src/types.rs create mode 100644 crates/onnx-runtime-memory-abi/src/version.rs create mode 100644 crates/onnx-runtime-memory-abi/src/vtable.rs create mode 100644 crates/onnx-runtime-memory-host/Cargo.toml create mode 100644 crates/onnx-runtime-memory-host/src/allocator.rs create mode 100644 crates/onnx-runtime-memory-host/src/error.rs create mode 100644 crates/onnx-runtime-memory-host/src/lib.rs create mode 100644 crates/onnx-runtime-memory-host/src/loader.rs create mode 100644 crates/onnx-runtime-memory-testplugin/Cargo.toml create mode 100644 crates/onnx-runtime-memory-testplugin/src/lib.rs diff --git a/Cargo.lock b/Cargo.lock index 77c906605f..47823a3042 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2550,6 +2550,10 @@ dependencies = [ "thiserror 2.0.19", ] +[[package]] +name = "onnx-runtime-memory-abi" +version = "0.1.0-dev.5" + [[package]] name = "onnx-runtime-memory-api" version = "0.1.0-dev.5" @@ -2566,6 +2570,14 @@ dependencies = [ "thiserror 2.0.19", ] +[[package]] +name = "onnx-runtime-memory-host" +version = "0.1.0-dev.5" + +[[package]] +name = "onnx-runtime-memory-testplugin" +version = "0.1.0-dev.5" + [[package]] name = "onnx-runtime-operator-selection" version = "0.1.0-dev.5" diff --git a/Cargo.toml b/Cargo.toml index fbcd59bef7..8937204196 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -73,6 +73,11 @@ members = [ "crates/onnx-runtime-ep-nxrt-host", "crates/onnx-runtime-ep-nxrt-testplugin", + # ── Dynamic memory-plugin ABI (nxmem — issue #1186 phase 6) ── + "crates/onnx-runtime-memory-abi", + "crates/onnx-runtime-memory-host", + "crates/onnx-runtime-memory-testplugin", + # ── Vendored ONNX Runtime MLAS bindings ── # Registered as a member so `cargo build -p mlas-sys` resolves. Crates that # enable the CPU EP's default `mlas` feature pull this in transitively; use @@ -104,6 +109,8 @@ default-members = [ "crates/onnx-runtime-cuda-memory", "crates/onnx-runtime-memory-api", "crates/onnx-runtime-memory-governor", + "crates/onnx-runtime-memory-abi", + "crates/onnx-runtime-memory-host", "crates/onnx-runtime-cpuinfo", "crates/onnx-runtime-loader", "crates/onnx-model-package", @@ -219,6 +226,11 @@ onnx-runtime-ep-nxrt-abi = { path = "crates/onnx-runtime-ep-nxrt-abi", version = onnx-runtime-ep-nxrt-host = { path = "crates/onnx-runtime-ep-nxrt-host", version = "=0.1.0-dev.5" } onnx-runtime-ep-nxrt-testplugin = { path = "crates/onnx-runtime-ep-nxrt-testplugin", version = "=0.1.0-dev.5" } +# Dynamic memory-plugin ABI (nxmem) +onnx-runtime-memory-abi = { path = "crates/onnx-runtime-memory-abi", version = "=0.1.0-dev.5" } +onnx-runtime-memory-host = { path = "crates/onnx-runtime-memory-host", version = "=0.1.0-dev.5" } +onnx-runtime-memory-testplugin = { path = "crates/onnx-runtime-memory-testplugin" } + # Test-only helper crate (publish = false, no version pin needed). onnx-runtime-ort-testkit = { path = "crates/onnx-runtime-ort-testkit" } diff --git a/crates/onnx-runtime-memory-abi/Cargo.toml b/crates/onnx-runtime-memory-abi/Cargo.toml new file mode 100644 index 0000000000..9c90c8d489 --- /dev/null +++ b/crates/onnx-runtime-memory-abi/Cargo.toml @@ -0,0 +1,17 @@ +[package] +name = "onnx-runtime-memory-abi" +description = "Stable versioned C ABI for dynamic-plugin memory mechanisms (allocator, virtual backing, shared mapping)" +version.workspace = true +edition.workspace = true +license.workspace = true +repository.workspace = true +keywords = ["onnx", "memory", "allocator", "ffi", "abi"] +categories = ["memory-management", "external-ffi-bindings"] + +# Deliberately dependency-free. A plugin author must be able to build against +# this contract (or against `include/nxmem_memory_abi.h`) without linking any +# other workspace crate. +[dependencies] + +[lints] +workspace = true diff --git a/crates/onnx-runtime-memory-abi/src/lib.rs b/crates/onnx-runtime-memory-abi/src/lib.rs new file mode 100644 index 0000000000..c43a4f415c --- /dev/null +++ b/crates/onnx-runtime-memory-abi/src/lib.rs @@ -0,0 +1,275 @@ +//! # `onnx-runtime-memory-abi` +//! +//! The **nxmem** ABI: a versioned C interface that lets a dynamic library +//! supply a memory mechanism (ordinary allocator, optional lazy virtual +//! backing, optional shared physical mapping) to this runtime. +//! +//! This crate is deliberately dependency-free. A plugin author builds against +//! it, or against [`include/nxmem_memory_abi.h`][header] with no Rust at all. +//! The host adapter that wraps this ABI in the internal Rust traits lives in +//! `onnx-runtime-memory-host`; the internal traits themselves live in +//! `onnx-runtime-memory-api` and never cross this boundary. +//! +//! [header]: https://github.com/justinchuby/onnx-genai/blob/main/crates/onnx-runtime-memory-abi/include/nxmem_memory_abi.h +//! +//! # What may cross the boundary +//! +//! Only `#[repr(C)]` plain data, raw pointers, and `extern "C"` function +//! pointers. Specifically **not**: Rust trait objects, `Arc`, `Box`, `String`, +//! `Vec`, Rust enum layouts, Rust `Result`, or any object one module allocates +//! and the other frees. Enum-shaped values travel as raw `u32` wire codes with +//! checked accessors, because transmuting an unknown discriminant is undefined +//! behaviour. Panics are caught with `catch_unwind` before they can unwind +//! across the boundary; see [`catch_status_panic`] and [`catch_void_panic`]. +//! +//! # Exported symbols +//! +//! A conforming plugin `cdylib` exports exactly three symbols: +//! +//! | Symbol | Constant | Purpose | +//! |---|---|---| +//! | `NxmemNegotiate` | [`NXMEM_SYMBOL_NEGOTIATE`] | version and capability handshake | +//! | `NxmemCreateAllocatorFactories` | [`NXMEM_SYMBOL_CREATE_ALLOCATOR_FACTORIES`] | enumerate mechanisms | +//! | `NxmemQueryUnloadReadiness` | [`NXMEM_SYMBOL_QUERY_UNLOAD_READINESS`] | report live objects so unload can be gated | +//! +//! All three are required. A library missing any of them is refused at load: +//! without the third the host cannot tell whether unloading would strand live +//! plugin state. +//! +//! # Lifetime and ownership rules +//! +//! | Object | Created by | Destroyed by | Rule | +//! |---|---|---|---| +//! | factory vtable | `NxmemCreateAllocatorFactories` | `factory.release(ctx)` | host calls it exactly once; releasing a factory must not invalidate allocators already opened from it | +//! | allocator vtable | `factory.open_allocator` | `allocator.release(ctx)` | reference counted via `retain`/`release`; the plugin destroys it only when the count reaches zero **and** no queued release still names it | +//! | virtual-backing vtable | exposed by an allocator | — | borrowed from the allocator; has no independent count and dies with it | +//! | shared-mapping vtable | exposed by an allocator | — | same as virtual backing | +//! | shared prefix | `create_shared_prefix` | `release_shared_prefix` | reference counted via `retain_shared_prefix`; physical bytes survive until the last reference *and* the last mapping retire | +//! | allocation | `allocate` / `allocate_committed` | `deallocate` / `release_allocation` / `enqueue_release` | released exactly once, by the mechanism that created it, keyed by `allocation_id` | +//! | host callback table | host | host | borrowed by the plugin for the whole allocator lifetime; the host keeps it alive past the final `release` and past every queued release | +//! | status / report / outcome | either | — | pure value types with inline storage; nothing to free | +//! +//! No object is ever freed by the module that did not allocate it, so the two +//! modules may be linked against different C runtimes. +//! +//! # Threading, reentrancy, and the no-blocking rule +//! +//! * Every slot may be called concurrently from multiple host threads unless +//! its documentation says otherwise. A plugin must do its own +//! synchronisation. +//! * **The host never holds a governance lock across a call into a plugin.** +//! Allocation accounting, charge settlement, and registration locks are all +//! released before the vtable is entered. This is the same rule the internal +//! Rust interfaces already obey for trait objects, and it matters more here: +//! a plugin may block, call back into the host, or spawn threads, and any of +//! those under a governance lock would deadlock the process. +//! * **A plugin must not block while a host callback is executing.** Host +//! callbacks are permitted to acquire host-side locks; a plugin that blocks +//! waiting for another plugin thread from inside a callback can deadlock. +//! * Host callbacks ([`NxmemHostCallbacks`]) may be invoked reentrantly from +//! inside a call the host is currently making, and from plugin-owned worker +//! threads. They must not block indefinitely and must not re-enter the same +//! allocator instance. +//! * A host callback returning a failure status is normal and expected. The +//! plugin must handle it (typically by failing the operation with +//! [`NxmemStatusCode::CallbackFailed`] or `OutOfMemory`) and must not abort, +//! leak, or leave partial state. +//! * `retain`, `release`, and `release_shared_prefix` must not call back into +//! the host at all: they run on teardown paths where the host may already be +//! dismantling its own state. +//! +//! # Deferred release and unload gating +//! +//! `enqueue_release` hands an allocation to a plugin-owned queue instead of +//! freeing it immediately, which is what GPU stream ordering requires. Until +//! the corresponding completion is reported through +//! [`NxmemHostCallbacks::release_completed`]: +//! +//! * the host keeps the plugin module, the factory, the allocator, and the +//! callback table pinned, and +//! * `NxmemQueryUnloadReadiness` must count the entry in `queued_releases`. +//! +//! The host refuses or defers unload while any of the counters in +//! [`NxmemUnloadReport`] is non-zero, and independently while it holds any +//! live handle of its own. Both gates exist so a bug on one side cannot +//! silently unmap code that is about to run. +//! +//! # Non-goals +//! +//! Internal policy, holder, victim-selection, and governor types are not +//! exposed here and never will be. Nothing that predates this contract is +//! promised ABI compatibility. + +#![deny(clippy::mem_forget)] + +pub mod status; +pub mod types; +pub mod version; +pub mod vtable; + +pub use status::{ + catch_status_panic, catch_void_panic, NxmemStatus, NxmemStatusCode, NXMEM_STATUS_MESSAGE_BUF, + NXMEM_STATUS_MESSAGE_MAX, +}; +pub use types::{ + check_identity, NxmemAllocRequest, NxmemAllocResult, NxmemAllocation, NxmemByteRange, + NxmemDeviceId, NxmemHostCallbacks, NxmemRangeRequest, NxmemReclaimRequest, + NxmemReleaseCompletion, NxmemReleaseOutcome, NxmemSharedPrefixCommitInfo, + NxmemSharedPrefixCommitRequest, NxmemSharedPrefixHandle, NxmemUnloadReport, + NXMEM_RELEASE_COMPLETE, NXMEM_RELEASE_FAILED, NXMEM_RELEASE_QUARANTINED, NXMEM_TIER_DEVICE, + NXMEM_TIER_DISK, NXMEM_TIER_HOST, +}; +pub use version::{ + capability_min_minor, describe_capabilities, negotiate, negotiate_as, validate_negotiation, + NegotiationRejection, NxmemNegotiateRequest, NxmemNegotiateResponse, NxmemVersionRange, + NXMEM_ABI_VERSION_MAJOR, NXMEM_ABI_VERSION_MINOR, NXMEM_ABI_VERSION_MINOR_BASELINE, + NXMEM_CAP_ALLOCATOR, NXMEM_CAP_DEFERRED_RELEASE, NXMEM_CAP_KNOWN_MASK, + NXMEM_CAP_SHARED_MAPPING, NXMEM_CAP_STRUCTURED_RELEASE, NXMEM_CAP_VIRTUAL_BACKING, +}; +pub use vtable::{ + NxmemAllocatorFactoryVtable, NxmemAllocatorVtable, NxmemOpenRequest, NxmemSharedMappingVtable, + NxmemVirtualBackingVtable, +}; + +/// Symbol name of the version handshake entry point. +pub const NXMEM_SYMBOL_NEGOTIATE: &[u8] = b"NxmemNegotiate"; + +/// Symbol name of the factory enumeration entry point. +pub const NXMEM_SYMBOL_CREATE_ALLOCATOR_FACTORIES: &[u8] = b"NxmemCreateAllocatorFactories"; + +/// Symbol name of the unload-readiness entry point. +pub const NXMEM_SYMBOL_QUERY_UNLOAD_READINESS: &[u8] = b"NxmemQueryUnloadReadiness"; + +/// Signature of `NxmemNegotiate`. +/// +/// # Ownership +/// +/// `request` is borrowed for the call. `response_out` is a host-owned buffer +/// the plugin fills. Neither pointer may be retained. +pub type NxmemNegotiateFn = unsafe extern "C" fn( + request: *const NxmemNegotiateRequest, + response_out: *mut NxmemNegotiateResponse, +) -> NxmemStatus; + +/// Signature of `NxmemCreateAllocatorFactories`. +/// +/// # Ownership +/// +/// Each pointer written to `out_factories` is **owned by the host**, which +/// must call the factory's `release` slot exactly once. Writing fewer than +/// `max_factories` entries is normal; writing more is a contract violation and +/// the host clamps defensively. +pub type NxmemCreateAllocatorFactoriesFn = unsafe extern "C" fn( + out_factories: *mut *const NxmemAllocatorFactoryVtable, + max_factories: u64, + out_count: *mut u64, +) -> NxmemStatus; + +/// Signature of `NxmemQueryUnloadReadiness`. +/// +/// # Contract +/// +/// Reports how many objects the plugin still owns. The host refuses or defers +/// unload while [`NxmemUnloadReport::total`] is non-zero. A plugin must count +/// conservatively: reporting zero while anything is live invites the host to +/// unmap code that is still reachable. +pub type NxmemQueryUnloadReadinessFn = + unsafe extern "C" fn(report_out: *mut NxmemUnloadReport) -> NxmemStatus; + +/// Emit the three required `nxmem` entry points for a `cdylib`. +/// +/// # Usage +/// +/// ```rust,ignore +/// onnx_runtime_memory_abi::export_nxmem_plugin! { +/// negotiate: |request, response| unsafe { +/// onnx_runtime_memory_abi::negotiate(request, response) +/// }, +/// factories: my_plugin::create_factories, +/// unload_readiness: my_plugin::unload_report, +/// } +/// ``` +/// +/// Every generated symbol is wrapped in `catch_unwind`, so a panic anywhere in +/// the plugin becomes [`NxmemStatusCode::InternalError`] instead of undefined +/// behaviour. +#[macro_export] +macro_rules! export_nxmem_plugin { + ( + negotiate: $negotiate:expr, + factories: $factories:expr, + unload_readiness: $unload:expr $(,)? + ) => { + /// nxmem version handshake. + /// + /// # Safety + /// + /// Both pointers must be valid, aligned, and non-null. + #[unsafe(no_mangle)] + pub unsafe extern "C" fn NxmemNegotiate( + request: *const $crate::NxmemNegotiateRequest, + response_out: *mut $crate::NxmemNegotiateResponse, + ) -> $crate::NxmemStatus { + let call = $negotiate; + $crate::catch_status_panic(|| call(request, response_out)) + } + + /// nxmem factory enumeration. + /// + /// # Safety + /// + /// `out_factories` must point to `max_factories` writable slots and + /// `out_count` must be a valid, non-null pointer. + #[unsafe(no_mangle)] + pub unsafe extern "C" fn NxmemCreateAllocatorFactories( + out_factories: *mut *const $crate::NxmemAllocatorFactoryVtable, + max_factories: u64, + out_count: *mut u64, + ) -> $crate::NxmemStatus { + let call = $factories; + let status = + $crate::catch_status_panic(|| call(out_factories, max_factories, out_count)); + if !status.is_ok() && !out_count.is_null() { + // A failed enumeration must never leave a stale count behind. + // SAFETY: `out_count` was checked non-null and the caller + // guarantees it is writable. + unsafe { *out_count = 0 }; + } + status + } + + /// nxmem unload-readiness report. + /// + /// # Safety + /// + /// `report_out` must be a valid, non-null, writable pointer. + #[unsafe(no_mangle)] + pub unsafe extern "C" fn NxmemQueryUnloadReadiness( + report_out: *mut $crate::NxmemUnloadReport, + ) -> $crate::NxmemStatus { + let call = $unload; + $crate::catch_status_panic(|| call(report_out)) + } + }; +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn the_three_required_symbol_names_are_distinct_and_non_empty() { + let symbols = [ + NXMEM_SYMBOL_NEGOTIATE, + NXMEM_SYMBOL_CREATE_ALLOCATOR_FACTORIES, + NXMEM_SYMBOL_QUERY_UNLOAD_READINESS, + ]; + for symbol in symbols { + assert!(!symbol.is_empty()); + assert!(std::str::from_utf8(symbol).is_ok()); + } + assert_ne!(symbols[0], symbols[1]); + assert_ne!(symbols[1], symbols[2]); + assert_ne!(symbols[0], symbols[2]); + } +} diff --git a/crates/onnx-runtime-memory-abi/src/status.rs b/crates/onnx-runtime-memory-abi/src/status.rs new file mode 100644 index 0000000000..0263d56865 --- /dev/null +++ b/crates/onnx-runtime-memory-abi/src/status.rs @@ -0,0 +1,373 @@ +//! Stable status/error values for the `nxmem` memory plugin ABI. +//! +//! # Why an inline message buffer +//! +//! The plugin and the host are separate dynamic modules that may be linked +//! against different C runtimes. Memory allocated by one module and freed by +//! the other is undefined behaviour, so [`NxmemStatus`] carries its message in +//! a fixed inline array. There is nothing to free, nothing to retain, and no +//! Rust allocator ownership crossing the boundary. +//! +//! # Wire-code safety +//! +//! [`NxmemStatus::code`] is a raw `u32`, never a Rust enum. A newer participant +//! may send a discriminant this build does not know; transmuting it into +//! [`NxmemStatusCode`] would be undefined behaviour. Use +//! [`NxmemStatus::status_code`], which returns `None` for unknown codes, and +//! treat `None` as a hard failure (fail closed). + +/// Maximum message bytes (excluding the NUL terminator) in an [`NxmemStatus`]. +/// +/// Longer messages are truncated on a UTF-8 character boundary. +pub const NXMEM_STATUS_MESSAGE_MAX: usize = 255; + +/// Total inline buffer length, including the NUL terminator. +pub const NXMEM_STATUS_MESSAGE_BUF: usize = NXMEM_STATUS_MESSAGE_MAX + 1; + +/// Stable status codes. +/// +/// Discriminants are append-only and never reordered: a value that shipped +/// keeps its meaning forever. Unknown discriminants must be treated as fatal. +#[repr(u32)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum NxmemStatusCode { + /// The call succeeded. + Ok = 0, + /// The host and plugin could not agree on an ABI version. + VersionMismatch = 1, + /// A struct's `struct_size` is smaller than the required prefix for the + /// negotiated ABI version, or a required function slot is null. + ShortStruct = 2, + /// The requested capability is not supported by this participant. This is + /// the explicit representation of absence — never a successful no-op. + UnsupportedCapability = 3, + /// A pointer was null, a size was nonsensical, or a field was out of range. + InvalidArgument = 4, + /// An unexpected internal failure, including a caught panic. + InternalError = 5, + /// The slot exists in the vtable but this participant does not implement it. + NotImplemented = 6, + /// The backing device or driver reported a failure. + DeviceError = 7, + /// The allocation could not be satisfied. + OutOfMemory = 8, + /// The object belongs to a different device than the one addressed. + WrongDevice = 9, + /// The object belongs to a different mechanism instance than the one + /// addressed — this rejects cross-provider misuse before any free. + WrongMechanism = 10, + /// The allocation identity is unknown, already retired, or from an earlier + /// generation that reused the same address. + UnknownAllocation = 11, + /// Release mutated state but could not complete; residual ownership stays + /// with the plugin. See the accompanying [`crate::NxmemReleaseOutcome`]. + ReleaseQuarantined = 12, + /// A host callback returned a failure and the plugin could not proceed. + CallbackFailed = 13, + /// The plugin still owns live objects and cannot satisfy the request now. + Busy = 14, +} + +impl NxmemStatusCode { + /// Whether this code means success. + pub const fn is_ok(self) -> bool { + matches!(self, Self::Ok) + } + + /// Checked conversion from a raw wire value. + /// + /// Returns `None` for values this build does not recognise. Never + /// transmute — an untrusted participant may send anything. + pub const fn from_u32(value: u32) -> Option { + match value { + 0 => Some(Self::Ok), + 1 => Some(Self::VersionMismatch), + 2 => Some(Self::ShortStruct), + 3 => Some(Self::UnsupportedCapability), + 4 => Some(Self::InvalidArgument), + 5 => Some(Self::InternalError), + 6 => Some(Self::NotImplemented), + 7 => Some(Self::DeviceError), + 8 => Some(Self::OutOfMemory), + 9 => Some(Self::WrongDevice), + 10 => Some(Self::WrongMechanism), + 11 => Some(Self::UnknownAllocation), + 12 => Some(Self::ReleaseQuarantined), + 13 => Some(Self::CallbackFailed), + 14 => Some(Self::Busy), + _ => None, + } + } + + /// Stable machine-readable name. These strings are part of the contract + /// and are what diagnostics quote, so they do not change with locale or + /// with the Rust `Debug` formatting of the enum. + pub const fn name(self) -> &'static str { + match self { + Self::Ok => "OK", + Self::VersionMismatch => "VERSION_MISMATCH", + Self::ShortStruct => "SHORT_STRUCT", + Self::UnsupportedCapability => "UNSUPPORTED_CAPABILITY", + Self::InvalidArgument => "INVALID_ARGUMENT", + Self::InternalError => "INTERNAL_ERROR", + Self::NotImplemented => "NOT_IMPLEMENTED", + Self::DeviceError => "DEVICE_ERROR", + Self::OutOfMemory => "OUT_OF_MEMORY", + Self::WrongDevice => "WRONG_DEVICE", + Self::WrongMechanism => "WRONG_MECHANISM", + Self::UnknownAllocation => "UNKNOWN_ALLOCATION", + Self::ReleaseQuarantined => "RELEASE_QUARANTINED", + Self::CallbackFailed => "CALLBACK_FAILED", + Self::Busy => "BUSY", + } + } + + /// Every code known at this ABI version, in wire order. + pub const ALL: [Self; 15] = [ + Self::Ok, + Self::VersionMismatch, + Self::ShortStruct, + Self::UnsupportedCapability, + Self::InvalidArgument, + Self::InternalError, + Self::NotImplemented, + Self::DeviceError, + Self::OutOfMemory, + Self::WrongDevice, + Self::WrongMechanism, + Self::UnknownAllocation, + Self::ReleaseQuarantined, + Self::CallbackFailed, + Self::Busy, + ]; +} + +/// The value every `nxmem` entry point returns. +/// +/// # Layout +/// +/// ```text +/// struct NxmemStatus { +/// uint32_t code; // NxmemStatusCode wire value +/// uint32_t message_len; // bytes of `message` that are meaningful +/// uint8_t message[256]; // NUL-terminated UTF-8 +/// }; +/// ``` +/// +/// # Ownership +/// +/// Pure value type. It is returned by value, copied freely, and never freed by +/// the other side. +#[repr(C)] +#[derive(Clone, Copy)] +pub struct NxmemStatus { + /// Raw wire status code. Use [`Self::status_code`] for checked conversion. + pub code: u32, + /// Meaningful bytes in `message`, excluding the NUL terminator. + pub message_len: u32, + /// Inline NUL-terminated UTF-8 message. Only the first `message_len` bytes + /// are meaningful. + pub message: [u8; NXMEM_STATUS_MESSAGE_BUF], +} + +// SAFETY: `NxmemStatus` is a plain value type holding no pointers and no +// interior mutability, so it is trivially safe to move between threads. +unsafe impl Send for NxmemStatus {} +// SAFETY: see the `Send` justification; shared references expose only reads. +unsafe impl Sync for NxmemStatus {} + +impl core::fmt::Debug for NxmemStatus { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.debug_struct("NxmemStatus") + .field("code", &self.status_code()) + .field("code_raw", &self.code) + .field("message", &self.message_str()) + .finish() + } +} + +impl NxmemStatus { + /// A success status with no message. + pub const fn ok() -> Self { + Self { + code: NxmemStatusCode::Ok as u32, + message_len: 0, + message: [0u8; NXMEM_STATUS_MESSAGE_BUF], + } + } + + /// A status carrying `code` and no message. + pub const fn from_code(code: NxmemStatusCode) -> Self { + Self { + code: code as u32, + message_len: 0, + message: [0u8; NXMEM_STATUS_MESSAGE_BUF], + } + } + + /// A status carrying `code` and a message. + /// + /// The message is truncated at the first interior NUL and at + /// [`NXMEM_STATUS_MESSAGE_MAX`] bytes. Truncation lands on a UTF-8 + /// character boundary so [`Self::message_str`] never fails on our own + /// output. + pub fn with_message(code: NxmemStatusCode, message: &str) -> Self { + let mut status = Self::from_code(code); + let bytes = message.as_bytes(); + let end = bytes.iter().position(|&b| b == 0).unwrap_or(bytes.len()); + let mut len = end.min(NXMEM_STATUS_MESSAGE_MAX); + // Back off to a character boundary so truncation cannot produce + // invalid UTF-8. + while len > 0 && !message.is_char_boundary(len) { + len -= 1; + } + status.message[..len].copy_from_slice(&bytes[..len]); + status.message[len] = 0; + status.message_len = len as u32; + status + } + + /// Whether the wire code is success. + pub const fn is_ok(&self) -> bool { + self.code == NxmemStatusCode::Ok as u32 + } + + /// Checked conversion of the wire code. + /// + /// `None` means "a participant sent a code this build does not know"; the + /// caller must treat that as a fatal error. + pub const fn status_code(&self) -> Option { + NxmemStatusCode::from_u32(self.code) + } + + /// The message, if any, as UTF-8. + /// + /// Returns `None` when there is no message or when the bytes are not valid + /// UTF-8 (defensive against a corrupt or hostile participant). + pub fn message_str(&self) -> Option<&str> { + let len = (self.message_len as usize).min(NXMEM_STATUS_MESSAGE_MAX); + if len == 0 { + return None; + } + core::str::from_utf8(&self.message[..len]).ok() + } + + /// A stable rendering of code and message for diagnostics. + /// + /// Unknown codes render as `UNKNOWN()` rather than being dropped, so a + /// newer plugin's failure is still traceable. + pub fn describe(&self) -> String { + let name = match self.status_code() { + Some(code) => code.name().to_string(), + None => format!("UNKNOWN({})", self.code), + }; + match self.message_str() { + Some(message) => format!("{name}: {message}"), + None => name, + } + } +} + +impl Default for NxmemStatus { + fn default() -> Self { + Self::ok() + } +} + +impl core::fmt::Display for NxmemStatus { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.write_str(&self.describe()) + } +} + +/// Run `f`, converting a Rust panic into [`NxmemStatusCode::InternalError`]. +/// +/// Every `extern "C"` function on either side of the boundary must funnel +/// through this (or [`catch_void_panic`]). Unwinding across a `cdylib` +/// boundary is undefined behaviour. +pub fn catch_status_panic NxmemStatus>(f: F) -> NxmemStatus { + match std::panic::catch_unwind(std::panic::AssertUnwindSafe(f)) { + Ok(status) => status, + Err(_) => NxmemStatus::with_message( + NxmemStatusCode::InternalError, + "nxmem: panic caught at the ABI boundary (fail closed)", + ), + } +} + +/// Run `f`, swallowing a Rust panic. +/// +/// Used for `void`-returning slots such as `retain`/`release`, which have no +/// channel to report a failure. +pub fn catch_void_panic(f: F) { + let _ = std::panic::catch_unwind(std::panic::AssertUnwindSafe(f)); +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn ok_status_has_no_message() { + let status = NxmemStatus::ok(); + assert!(status.is_ok()); + assert_eq!(status.message_str(), None); + assert_eq!(status.describe(), "OK"); + } + + #[test] + fn every_known_code_round_trips_and_has_a_stable_name() { + let mut seen = Vec::new(); + for code in NxmemStatusCode::ALL { + assert_eq!(NxmemStatusCode::from_u32(code as u32), Some(code)); + assert!(!code.name().is_empty()); + assert!(!seen.contains(&code.name()), "duplicate name {}", code.name()); + seen.push(code.name()); + } + assert_eq!(seen.len(), NxmemStatusCode::ALL.len()); + } + + #[test] + fn unknown_wire_code_is_not_ok_and_does_not_transmute() { + let mut status = NxmemStatus::ok(); + status.code = 9_999; + assert!(!status.is_ok()); + assert_eq!(status.status_code(), None); + assert_eq!(status.describe(), "UNKNOWN(9999)"); + } + + #[test] + fn long_message_truncates_on_a_character_boundary() { + // Three-byte characters do not divide 255 evenly, so a naive + // truncation would split one and make `message_str` return None. + let message = "\u{4e2d}".repeat(200); + let status = NxmemStatus::with_message(NxmemStatusCode::DeviceError, &message); + assert_eq!(status.message_len, 255 - (255 % 3)); + let recovered = status.message_str().expect("valid utf-8 after truncation"); + assert!(recovered.chars().all(|c| c == '\u{4e2d}')); + } + + #[test] + fn interior_nul_truncates_the_message() { + let status = NxmemStatus::with_message(NxmemStatusCode::InvalidArgument, "head\0tail"); + assert_eq!(status.message_str(), Some("head")); + } + + #[test] + fn status_layout_is_the_documented_264_bytes() { + assert_eq!(std::mem::size_of::(), 264); + assert_eq!(std::mem::align_of::(), 4); + } + + #[test] + fn panic_in_a_status_slot_becomes_an_internal_error() { + let status = catch_status_panic(|| panic!("deliberate")); + assert_eq!(status.status_code(), Some(NxmemStatusCode::InternalError)); + assert!(status.describe().contains("panic")); + } + + #[test] + fn panic_in_a_void_slot_does_not_unwind() { + catch_void_panic(|| panic!("deliberate")); + } +} diff --git a/crates/onnx-runtime-memory-abi/src/types.rs b/crates/onnx-runtime-memory-abi/src/types.rs new file mode 100644 index 0000000000..15c6b1b233 --- /dev/null +++ b/crates/onnx-runtime-memory-abi/src/types.rs @@ -0,0 +1,701 @@ +//! Plain-data records that cross the `nxmem` boundary. +//! +//! Every type here is `#[repr(C)]`, holds no Rust trait object, no `Arc`, no +//! Rust enum layout, and no heap allocation owned by the other module. Enum- +//! shaped values travel as raw `u32` wire codes with checked accessors. + +use crate::status::{NxmemStatus, NxmemStatusCode}; + +/// Wire values for the memory tier an object lives in. +/// +/// Mirrors the host-side `Tier`, but travels as a raw `u32`: see +/// [`NxmemDeviceId::tier_code`]. +pub const NXMEM_TIER_DEVICE: u32 = 0; +/// Host RAM tier wire value. +pub const NXMEM_TIER_HOST: u32 = 1; +/// Spill-to-disk tier wire value. +pub const NXMEM_TIER_DISK: u32 = 2; + +/// Which physical device an object belongs to. +/// +/// Carried by every allocation, backing, and shared-mapping call so a +/// cross-device request is rejected before it can free or map anything. +#[repr(C)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub struct NxmemDeviceId { + /// Tier wire value: one of [`NXMEM_TIER_DEVICE`], [`NXMEM_TIER_HOST`], + /// [`NXMEM_TIER_DISK`]. Unknown values must be rejected, not guessed. + pub tier: u32, + /// Device ordinal within the tier. + pub index: u32, +} + +impl NxmemDeviceId { + /// The single host-RAM device. + pub const HOST: Self = Self { + tier: NXMEM_TIER_HOST, + index: 0, + }; + + /// An accelerator device by ordinal. + pub const fn device(index: u32) -> Self { + Self { + tier: NXMEM_TIER_DEVICE, + index, + } + } + + /// The tier wire value, or `None` when this build does not know it. + pub const fn tier_code(self) -> Option { + match self.tier { + NXMEM_TIER_DEVICE | NXMEM_TIER_HOST | NXMEM_TIER_DISK => Some(self.tier), + _ => None, + } + } +} + +/// The identity of one allocation as it crosses the boundary. +/// +/// `allocation_id` comes from a monotonic counter on the side that created the +/// allocation. It is deliberately **never** derived from the address: a +/// virtual address can be reused, an identity cannot. A plugin must reject a +/// call whose `allocation_id` it does not currently hold, even when `ptr` +/// happens to name live memory. +#[repr(C)] +#[derive(Debug, Clone, Copy)] +pub struct NxmemAllocation { + /// `size_of` this struct as the sender defines it. + pub struct_size: u32, + /// Reserved; must be zero. + pub reserved: u32, + /// The mechanism instance that created the allocation. + pub mechanism_id: u64, + /// Monotonic, never address-derived allocation identity. + pub allocation_id: u64, + /// The device the allocation lives on. + pub device: NxmemDeviceId, + /// The allocation's base address. + pub ptr: *mut u8, + /// The exact byte count the allocation was created with. + pub bytes: u64, + /// The exact alignment the allocation was created with. + pub align: u64, +} + +impl NxmemAllocation { + /// An allocation record with the correct `struct_size`. + pub fn new( + mechanism_id: u64, + allocation_id: u64, + device: NxmemDeviceId, + ptr: *mut u8, + bytes: u64, + align: u64, + ) -> Self { + Self { + struct_size: core::mem::size_of::() as u32, + reserved: 0, + mechanism_id, + allocation_id, + device, + ptr, + bytes, + align, + } + } +} + +/// A request to create an allocation. +/// +/// The same record serves ordinary allocation and lazy reserve-with-commit. +/// `committed_ranges` is ignored by the ordinary `allocate` slot; the virtual +/// backing slot honours it. +#[repr(C)] +#[derive(Debug, Clone, Copy)] +pub struct NxmemAllocRequest { + /// `size_of` this struct as the sender defines it. + pub struct_size: u32, + /// Reserved; must be zero. + pub reserved: u32, + /// The mechanism instance being addressed. A plugin must reject a value + /// that is not its own. + pub mechanism_id: u64, + /// The identity the host assigns to the allocation being created. + pub allocation_id: u64, + /// The device being addressed. A plugin must reject a foreign device. + pub device: NxmemDeviceId, + /// Requested byte count. + pub bytes: u64, + /// Requested alignment; must be a power of two. + pub align: u64, + /// Ranges to commit immediately. Borrowed for the call only; the plugin + /// must not retain the pointer. + pub committed_ranges: *const NxmemByteRange, + /// Number of entries in `committed_ranges`. + pub committed_range_count: u64, +} + +impl NxmemAllocRequest { + /// An ordinary allocation request with no pre-committed ranges. + pub fn new( + mechanism_id: u64, + allocation_id: u64, + device: NxmemDeviceId, + bytes: u64, + align: u64, + ) -> Self { + Self { + struct_size: core::mem::size_of::() as u32, + reserved: 0, + mechanism_id, + allocation_id, + device, + bytes, + align, + committed_ranges: core::ptr::null(), + committed_range_count: 0, + } + } +} + +/// A half-open `[offset, offset + bytes)` span inside one allocation. +#[repr(C)] +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct NxmemByteRange { + /// Byte offset from the allocation base. + pub offset: u64, + /// Span length in bytes. + pub bytes: u64, +} + +/// What an allocation call produced. +#[repr(C)] +#[derive(Debug, Clone, Copy)] +pub struct NxmemAllocResult { + /// `size_of` this struct as the sender defines it. + pub struct_size: u32, + /// Reserved; must be zero. + pub reserved: u32, + /// The allocation base address. Non-null on success. + pub ptr: *mut u8, + /// Physical bytes whose ownership was newly created by this call. A + /// mapping that reuses already-owned physical memory reports zero. + pub owned_bytes: u64, + /// Bytes whose mapped attribution newly became mapped. + pub mapped_bytes: u64, +} + +impl NxmemAllocResult { + /// A zeroed result for the caller to pass in. + pub const fn zeroed() -> Self { + Self { + struct_size: core::mem::size_of::() as u32, + reserved: 0, + ptr: core::ptr::null_mut(), + owned_bytes: 0, + mapped_bytes: 0, + } + } +} + +/// A commit/decommit request against a span of an existing allocation. +#[repr(C)] +#[derive(Debug, Clone, Copy)] +pub struct NxmemRangeRequest { + /// `size_of` this struct as the sender defines it. + pub struct_size: u32, + /// Reserved; must be zero. + pub reserved: u32, + /// The allocation being addressed, carrying mechanism and device identity. + pub allocation: NxmemAllocation, + /// The span within that allocation. + pub range: NxmemByteRange, +} + +impl NxmemRangeRequest { + /// A range request with the correct `struct_size`. + pub fn new(allocation: NxmemAllocation, offset: u64, bytes: u64) -> Self { + Self { + struct_size: core::mem::size_of::() as u32, + reserved: 0, + allocation, + range: NxmemByteRange { offset, bytes }, + } + } +} + +// ─── Release outcome ──────────────────────────────────────────────────────── + +/// Release completed: the whole allocation is gone. +pub const NXMEM_RELEASE_COMPLETE: u32 = 0; +/// Release mutated state but could not finish. Residual ownership stays with +/// the plugin and the virtual span must not be reused. +pub const NXMEM_RELEASE_QUARANTINED: u32 = 1; +/// Release changed nothing. This is the only shape that means "unchanged". +pub const NXMEM_RELEASE_FAILED: u32 = 2; + +/// The structured result of a terminal release. +/// +/// Zero unmapped bytes is a valid complete result and never signals failure — +/// the state field is the only failure channel. +#[repr(C)] +#[derive(Debug, Clone, Copy)] +pub struct NxmemReleaseOutcome { + /// `size_of` this struct as the sender defines it. + pub struct_size: u32, + /// One of [`NXMEM_RELEASE_COMPLETE`], [`NXMEM_RELEASE_QUARANTINED`], + /// [`NXMEM_RELEASE_FAILED`]. Unknown values must be treated as failure. + pub state: u32, + /// The allocation's byte count, echoed for accounting. + pub allocation_bytes: u64, + /// Bytes whose global mapping reference transitioned to unmapped. + pub unmapped_bytes: u64, + /// Physical bytes the plugin still owns after a quarantined release. + pub residual_owned_bytes: u64, + /// Why the release did not complete. Meaningful when `state` is not + /// [`NXMEM_RELEASE_COMPLETE`]. + pub failure: NxmemStatus, +} + +impl NxmemReleaseOutcome { + /// A zeroed outcome for the caller to pass in. + /// + /// The zero value is `Complete` with zero bytes, which is only ever + /// meaningful once a slot has actually written to it; callers must not + /// interpret an unwritten outcome. + pub const fn zeroed() -> Self { + Self { + struct_size: core::mem::size_of::() as u32, + state: NXMEM_RELEASE_COMPLETE, + allocation_bytes: 0, + unmapped_bytes: 0, + residual_owned_bytes: 0, + failure: NxmemStatus::ok(), + } + } + + /// A complete release. + pub const fn complete(allocation_bytes: u64, unmapped_bytes: u64) -> Self { + Self { + struct_size: core::mem::size_of::() as u32, + state: NXMEM_RELEASE_COMPLETE, + allocation_bytes, + unmapped_bytes, + residual_owned_bytes: 0, + failure: NxmemStatus::ok(), + } + } + + /// A release that mutated state and left residual ownership behind. + pub fn quarantined( + allocation_bytes: u64, + unmapped_bytes: u64, + residual_owned_bytes: u64, + failure: NxmemStatus, + ) -> Self { + Self { + struct_size: core::mem::size_of::() as u32, + state: NXMEM_RELEASE_QUARANTINED, + allocation_bytes, + unmapped_bytes, + residual_owned_bytes, + failure, + } + } + + /// A release that changed nothing. + pub fn failed(allocation_bytes: u64, failure: NxmemStatus) -> Self { + Self { + struct_size: core::mem::size_of::() as u32, + state: NXMEM_RELEASE_FAILED, + allocation_bytes, + unmapped_bytes: 0, + residual_owned_bytes: 0, + failure, + } + } + + /// Whether the allocation is fully gone. + pub const fn is_complete(&self) -> bool { + self.state == NXMEM_RELEASE_COMPLETE + } +} + +/// One retired deferred release, reported to the host. +#[repr(C)] +#[derive(Debug, Clone, Copy)] +pub struct NxmemReleaseCompletion { + /// `size_of` this struct as the sender defines it. + pub struct_size: u32, + /// Reserved; must be zero. + pub reserved: u32, + /// The ticket the enqueue call returned. + pub ticket: u64, + /// The mechanism instance that owned the allocation. + pub mechanism_id: u64, + /// The allocation identity that retired. + pub allocation_id: u64, + /// The structured outcome of the physical release. + pub outcome: NxmemReleaseOutcome, +} + +// ─── Host callbacks ───────────────────────────────────────────────────────── + +/// A plugin's request that the host free cached memory. +#[repr(C)] +#[derive(Debug, Clone, Copy)] +pub struct NxmemReclaimRequest { + /// `size_of` this struct as the sender defines it. + pub struct_size: u32, + /// Reserved; must be zero. + pub reserved: u32, + /// The mechanism instance asking. + pub mechanism_id: u64, + /// The device under pressure. + pub device: NxmemDeviceId, + /// Bytes the plugin needs. + pub bytes: u64, +} + +/// Callbacks the host offers to a plugin allocator. +/// +/// # Ownership +/// +/// The table and `host_ctx` are **borrowed for the lifetime of the opened +/// allocator**. The host guarantees they outlive the allocator's final +/// `release` and every queued release that names it. The plugin must not free +/// or retain either beyond that point. +/// +/// # Threading and reentrancy +/// +/// * The host calls into the plugin **without holding any governance lock**. +/// * A plugin may therefore call back into these slots from inside a plugin +/// call the host is currently making (allocate → `request_reclaim` is the +/// expected pattern), and from a plugin-owned worker thread. +/// * Host callbacks must not block indefinitely and must not re-enter the same +/// allocator instance. +/// * A callback returning a non-`Ok` status is a normal, expected outcome. The +/// plugin must handle it — typically by failing the operation with +/// [`NxmemStatusCode::CallbackFailed`] or `OutOfMemory` — and must not +/// abort, panic, or leak. +#[repr(C)] +#[derive(Debug, Clone, Copy)] +pub struct NxmemHostCallbacks { + /// `size_of` this struct as the host defines it. + pub struct_size: u32, + /// The minor version the host negotiated. + pub abi_minor: u32, + /// Opaque host state passed back to every slot. Never dereferenced by the + /// plugin. + pub host_ctx: *mut core::ffi::c_void, + /// Ask the host to release cached memory. Writes the bytes actually + /// reclaimed. Null when the host offers no reclaim path. + pub request_reclaim: Option< + unsafe extern "C" fn( + host_ctx: *mut core::ffi::c_void, + request: *const NxmemReclaimRequest, + reclaimed_out: *mut u64, + ) -> NxmemStatus, + >, + /// Report one retired deferred release so the host can settle accounting. + /// Null when the host does not consume deferred release. + pub release_completed: Option< + unsafe extern "C" fn( + host_ctx: *mut core::ffi::c_void, + completion: *const NxmemReleaseCompletion, + ) -> NxmemStatus, + >, +} + +impl NxmemHostCallbacks { + /// A table with no callbacks offered. + pub const fn empty(abi_minor: u32) -> Self { + Self { + struct_size: core::mem::size_of::() as u32, + abi_minor, + host_ctx: core::ptr::null_mut(), + request_reclaim: None, + release_completed: None, + } + } +} + +/// How much a plugin still owns, used to gate unload. +#[repr(C)] +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct NxmemUnloadReport { + /// `size_of` this struct as the plugin defines it. + pub struct_size: u32, + /// Reserved; must be zero. + pub reserved: u32, + /// Allocator instances the host has opened and not released. + pub live_allocators: u64, + /// Allocations not yet terminally released. + pub live_allocations: u64, + /// Borrowed views over live allocations. + pub live_views: u64, + /// Capability objects (shared prefixes, backing handles) still held. + pub live_capabilities: u64, + /// Releases enqueued and not yet retired. + pub queued_releases: u64, +} + +impl NxmemUnloadReport { + /// A zeroed report for the caller to pass in. + pub const fn zeroed() -> Self { + Self { + struct_size: core::mem::size_of::() as u32, + reserved: 0, + live_allocators: 0, + live_allocations: 0, + live_views: 0, + live_capabilities: 0, + queued_releases: 0, + } + } + + /// Total live objects. Unload is refused or deferred while this is + /// non-zero. + pub const fn total(&self) -> u64 { + self.live_allocators + .saturating_add(self.live_allocations) + .saturating_add(self.live_views) + .saturating_add(self.live_capabilities) + .saturating_add(self.queued_releases) + } +} + +// ─── Shared physical prefixes ─────────────────────────────────────────────── + +/// A plugin-owned shared physical prefix. +/// +/// # Ownership +/// +/// Created by `create_shared_prefix` and destroyed by `release_shared_prefix`, +/// exactly once, on the same mechanism instance. `handle` is opaque to the +/// host: it is never dereferenced, only echoed back. +#[repr(C)] +#[derive(Debug, Clone, Copy)] +pub struct NxmemSharedPrefixHandle { + /// `size_of` this struct as the plugin defines it. + pub struct_size: u32, + /// Reserved; must be zero. + pub reserved: u32, + /// The mechanism instance that owns the prefix. + pub mechanism_id: u64, + /// Opaque plugin-side identity. Non-zero when valid. + pub handle: u64, + /// The device the prefix lives on. + pub device: NxmemDeviceId, + /// Device address of the prefix. + pub device_ptr: u64, + /// Physical bytes owned by the prefix. + pub committed_physical_bytes: u64, + /// Bytes currently mapped. + pub mapped_bytes: u64, + /// Bytes originally requested. + pub requested_bytes: u64, +} + +impl NxmemSharedPrefixHandle { + /// A zeroed handle for the caller to pass in. + pub const fn zeroed() -> Self { + Self { + struct_size: core::mem::size_of::() as u32, + reserved: 0, + mechanism_id: 0, + handle: 0, + device: NxmemDeviceId { tier: 0, index: 0 }, + device_ptr: 0, + committed_physical_bytes: 0, + mapped_bytes: 0, + requested_bytes: 0, + } + } +} + +/// A request to map a shared prefix into an allocation. +#[repr(C)] +#[derive(Debug, Clone, Copy)] +pub struct NxmemSharedPrefixCommitRequest { + /// `size_of` this struct as the sender defines it. + pub struct_size: u32, + /// Reserved; must be zero. + pub reserved: u32, + /// The prefix to map. Borrowed for the call only. + pub prefix: NxmemSharedPrefixHandle, + /// The destination allocation. + pub allocation: NxmemAllocation, + /// Byte offset within the allocation to map at. + pub byte_offset: u64, +} + +/// The accounting result of mapping a shared prefix. +#[repr(C)] +#[derive(Debug, Clone, Copy, Default)] +pub struct NxmemSharedPrefixCommitInfo { + /// `size_of` this struct as the sender defines it. + pub struct_size: u32, + /// Reserved; must be zero. + pub reserved: u32, + /// Newly owned physical bytes. A valid additional mapping of an + /// already-owned prefix reports zero. + pub additional_owned_bytes: u64, + /// Newly mapped bytes on the mapped-attribution axis. + pub newly_mapped_bytes: u64, + /// Physical granules touched. + pub granules: u64, +} + +impl NxmemSharedPrefixCommitInfo { + /// A zeroed info record for the caller to pass in. + pub const fn zeroed() -> Self { + Self { + struct_size: core::mem::size_of::() as u32, + reserved: 0, + additional_owned_bytes: 0, + newly_mapped_bytes: 0, + granules: 0, + } + } +} + +/// Reject a call whose mechanism or device identity does not match the +/// receiver's own. +/// +/// Both sides use this: the host checks before calling, the plugin checks on +/// entry. Two independent checks are deliberate — the host protects itself +/// from its own bookkeeping bugs, and the plugin protects itself from a +/// hostile or confused host. +pub fn check_identity( + expected_mechanism: u64, + expected_device: NxmemDeviceId, + actual_mechanism: u64, + actual_device: NxmemDeviceId, +) -> Result<(), NxmemStatus> { + if actual_mechanism != expected_mechanism { + return Err(NxmemStatus::with_message( + NxmemStatusCode::WrongMechanism, + &format!( + "nxmem: object belongs to mechanism {actual_mechanism} but mechanism \ + {expected_mechanism} was addressed; an allocation may only be used and released \ + by the mechanism that created it" + ), + )); + } + if actual_device != expected_device { + return Err(NxmemStatus::with_message( + NxmemStatusCode::WrongDevice, + &format!( + "nxmem: object belongs to device (tier {}, index {}) but device (tier {}, index \ + {}) was addressed", + actual_device.tier, actual_device.index, expected_device.tier, expected_device.index + ), + )); + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn constructors_stamp_the_real_struct_size() { + let allocation = + NxmemAllocation::new(7, 9, NxmemDeviceId::HOST, core::ptr::null_mut(), 64, 16); + assert_eq!( + allocation.struct_size as usize, + core::mem::size_of::() + ); + let request = NxmemAllocRequest::new(7, 9, NxmemDeviceId::HOST, 64, 16); + assert_eq!( + request.struct_size as usize, + core::mem::size_of::() + ); + let range = NxmemRangeRequest::new(allocation, 0, 32); + assert_eq!( + range.struct_size as usize, + core::mem::size_of::() + ); + } + + #[test] + fn a_wrong_mechanism_is_rejected_before_a_wrong_device() { + let status = check_identity(1, NxmemDeviceId::HOST, 2, NxmemDeviceId::device(3)) + .expect_err("both identities differ"); + assert_eq!(status.status_code(), Some(NxmemStatusCode::WrongMechanism)); + assert!(status.describe().contains("mechanism that created it")); + } + + #[test] + fn a_wrong_device_on_the_right_mechanism_is_rejected() { + let status = check_identity(1, NxmemDeviceId::HOST, 1, NxmemDeviceId::device(0)) + .expect_err("device differs"); + assert_eq!(status.status_code(), Some(NxmemStatusCode::WrongDevice)); + } + + #[test] + fn matching_identity_is_accepted() { + assert!(check_identity(4, NxmemDeviceId::device(1), 4, NxmemDeviceId::device(1)).is_ok()); + } + + #[test] + fn zero_unmapped_bytes_is_a_valid_complete_release() { + let outcome = NxmemReleaseOutcome::complete(4096, 0); + assert!(outcome.is_complete()); + assert_eq!(outcome.unmapped_bytes, 0); + assert!(outcome.failure.is_ok()); + } + + #[test] + fn a_quarantined_release_keeps_residual_ownership() { + let outcome = NxmemReleaseOutcome::quarantined( + 4096, + 2048, + 2048, + NxmemStatus::with_message(NxmemStatusCode::DeviceError, "unmap failed midway"), + ); + assert!(!outcome.is_complete()); + assert_eq!(outcome.state, NXMEM_RELEASE_QUARANTINED); + assert_eq!(outcome.residual_owned_bytes, 2048); + } + + #[test] + fn a_failed_release_reports_no_mutation() { + let outcome = NxmemReleaseOutcome::failed( + 4096, + NxmemStatus::from_code(NxmemStatusCode::UnknownAllocation), + ); + assert_eq!(outcome.state, NXMEM_RELEASE_FAILED); + assert_eq!(outcome.unmapped_bytes, 0); + assert_eq!(outcome.residual_owned_bytes, 0); + } + + #[test] + fn unload_report_totals_every_axis() { + let report = NxmemUnloadReport { + live_allocators: 1, + live_allocations: 2, + live_views: 3, + live_capabilities: 4, + queued_releases: 5, + ..NxmemUnloadReport::zeroed() + }; + assert_eq!(report.total(), 15); + assert_eq!(NxmemUnloadReport::zeroed().total(), 0); + } + + #[test] + fn unknown_tier_codes_are_not_guessed() { + assert_eq!(NxmemDeviceId::HOST.tier_code(), Some(NXMEM_TIER_HOST)); + assert_eq!( + NxmemDeviceId { + tier: 77, + index: 0 + } + .tier_code(), + None + ); + } +} diff --git a/crates/onnx-runtime-memory-abi/src/version.rs b/crates/onnx-runtime-memory-abi/src/version.rs new file mode 100644 index 0000000000..bcc72fb6f3 --- /dev/null +++ b/crates/onnx-runtime-memory-abi/src/version.rs @@ -0,0 +1,603 @@ +//! ABI version negotiation and capability flags. +//! +//! # Version semantics +//! +//! * **Major** — breaking change. Host and plugin must agree on one major +//! version. Any mismatch is a hard rejection. +//! * **Minor** — additive change only: new trailing vtable slots and new +//! capability bits. The agreed minor is `min(host_minor_max, +//! plugin_minor_max)`. Both sides then behave as if only the slots defined +//! at or below the agreed minor exist. +//! +//! # Prefix negotiation +//! +//! Every ABI struct starts with `struct_size: u32`. A participant fills in +//! `size_of` its own definition. The reader compares that against the minimum +//! prefix it requires: +//! +//! * `struct_size` **smaller than the required prefix** for the negotiated +//! minor → [`crate::NxmemStatusCode::ShortStruct`], fail closed. Never read +//! the missing bytes. +//! * `struct_size` **larger** than the reader knows about → the reader ignores +//! the trailing bytes. That is how an older host tolerates a newer plugin at +//! the same major version. +//! +//! Optional slots are nullable function pointers. Null means "not supported" +//! and must be surfaced as [`crate::NxmemStatusCode::UnsupportedCapability`], +//! never as a successful no-op. +//! +//! # Baseline +//! +//! Minor `0` is the baseline prefix. Minor `1` appends the structured-release +//! slot to the allocator vtable. A minor-0 plugin therefore reports a smaller +//! `struct_size` and the host falls back to the unstructured release path; +//! that is the supported-older-participant case, not an error. + +use crate::status::{NxmemStatus, NxmemStatusCode}; + +/// Current major version of the `nxmem` ABI. +pub const NXMEM_ABI_VERSION_MAJOR: u32 = 1; + +/// Current minor version of the `nxmem` ABI. +pub const NXMEM_ABI_VERSION_MINOR: u32 = 1; + +/// The oldest minor version this build still supports as a peer. +/// +/// A participant reporting this minor uses only the baseline prefix of every +/// struct and none of the slots added after it. +pub const NXMEM_ABI_VERSION_MINOR_BASELINE: u32 = 0; + +/// An inclusive version range a participant supports. +#[repr(C)] +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct NxmemVersionRange { + /// Lowest supported major version (inclusive). + pub major_min: u32, + /// Highest supported major version (inclusive). + pub major_max: u32, + /// Lowest supported minor version at `major_max` (inclusive). + pub minor_min: u32, + /// Highest supported minor version at `major_max` (inclusive). + pub minor_max: u32, +} + +impl NxmemVersionRange { + /// The range this build supports. + pub const fn current() -> Self { + Self { + major_min: NXMEM_ABI_VERSION_MAJOR, + major_max: NXMEM_ABI_VERSION_MAJOR, + minor_min: NXMEM_ABI_VERSION_MINOR_BASELINE, + minor_max: NXMEM_ABI_VERSION_MINOR, + } + } + + /// A range pinned to exactly one `major.minor`, for tests and for hosts + /// that deliberately restrict themselves. + pub const fn exact(major: u32, minor: u32) -> Self { + Self { + major_min: major, + major_max: major, + minor_min: minor, + minor_max: minor, + } + } +} + +/// What the host sends to `NxmemNegotiate`. +/// +/// # Ownership +/// +/// Borrowed by the plugin for the duration of the call only. The plugin must +/// not retain the pointer. +#[repr(C)] +#[derive(Debug, Clone, Copy)] +pub struct NxmemNegotiateRequest { + /// `size_of` this struct as the host defines it. + pub struct_size: u32, + /// Reserved; must be zero. + pub reserved: u32, + /// The versions the host can speak. + pub host_range: NxmemVersionRange, + /// Capabilities the host is prepared to consume. A plugin must not + /// advertise a capability the host did not offer. + pub host_capability_flags: u64, +} + +impl NxmemNegotiateRequest { + /// A request advertising this build's range and every known capability. + pub const fn current() -> Self { + Self { + struct_size: core::mem::size_of::() as u32, + reserved: 0, + host_range: NxmemVersionRange::current(), + host_capability_flags: NXMEM_CAP_KNOWN_MASK, + } + } + + /// A request advertising a specific range, used by hosts that pin an older + /// contract and by negotiation tests. + pub const fn with_range(host_range: NxmemVersionRange) -> Self { + Self { + struct_size: core::mem::size_of::() as u32, + reserved: 0, + host_range, + host_capability_flags: NXMEM_CAP_KNOWN_MASK, + } + } +} + +/// What the plugin writes back from `NxmemNegotiate`. +/// +/// # Ownership +/// +/// The host owns the buffer and it is valid only for the call. +#[repr(C)] +#[derive(Debug, Clone, Copy)] +pub struct NxmemNegotiateResponse { + /// `size_of` this struct as the plugin defines it. + pub struct_size: u32, + /// Agreed major version. Meaningful only when the status is `Ok`. + pub agreed_major: u32, + /// Agreed minor version. Meaningful only when the status is `Ok`. + pub agreed_minor: u32, + /// Reserved; must be zero. + pub reserved: u32, + /// The plugin's own range. Always filled, including on failure, so the + /// host can produce an actionable diagnostic. + pub plugin_range: NxmemVersionRange, + /// Capabilities the plugin offers at the agreed version. Always a subset + /// of the host's offered flags and of [`NXMEM_CAP_KNOWN_MASK`]. + pub capability_flags: u64, +} + +impl NxmemNegotiateResponse { + /// A zeroed response for the host to pass in. + pub const fn zeroed() -> Self { + Self { + struct_size: core::mem::size_of::() as u32, + agreed_major: 0, + agreed_minor: 0, + reserved: 0, + plugin_range: NxmemVersionRange { + major_min: 0, + major_max: 0, + minor_min: 0, + minor_max: 0, + }, + capability_flags: 0, + } + } +} + +// ─── Capability flags ─────────────────────────────────────────────────────── + +/// Ordinary allocation and terminal release. Every plugin must set this bit; +/// it is the minimum contract, not an option. +pub const NXMEM_CAP_ALLOCATOR: u64 = 1 << 0; + +/// Optional lazy reserve/commit/decommit (`VirtualBacking`). +pub const NXMEM_CAP_VIRTUAL_BACKING: u64 = 1 << 1; + +/// Optional shared physical handles and prefix mapping (`SharedMapping`). +pub const NXMEM_CAP_SHARED_MAPPING: u64 = 1 << 2; + +/// Optional stream-ordered deferred release with host completion callbacks. +pub const NXMEM_CAP_DEFERRED_RELEASE: u64 = 1 << 3; + +/// Optional structured release outcome (added at minor 1). +pub const NXMEM_CAP_STRUCTURED_RELEASE: u64 = 1 << 4; + +/// Every capability bit known at this ABI version. +/// +/// A participant that sets a bit outside this mask is rejected: the host +/// cannot reason about a capability it has never heard of. +pub const NXMEM_CAP_KNOWN_MASK: u64 = NXMEM_CAP_ALLOCATOR + | NXMEM_CAP_VIRTUAL_BACKING + | NXMEM_CAP_SHARED_MAPPING + | NXMEM_CAP_DEFERRED_RELEASE + | NXMEM_CAP_STRUCTURED_RELEASE; + +/// The minor version at which each capability bit was introduced. +/// +/// A capability may only be advertised once the agreed minor is at least this +/// value, otherwise the vtable prefix carrying it does not exist. +pub const fn capability_min_minor(flag: u64) -> u32 { + match flag { + NXMEM_CAP_STRUCTURED_RELEASE => 1, + _ => 0, + } +} + +/// Human-facing names for capability bits, used in diagnostics. +pub fn describe_capabilities(flags: u64) -> String { + const NAMES: [(u64, &str); 5] = [ + (NXMEM_CAP_ALLOCATOR, "allocator"), + (NXMEM_CAP_VIRTUAL_BACKING, "virtual-backing"), + (NXMEM_CAP_SHARED_MAPPING, "shared-mapping"), + (NXMEM_CAP_DEFERRED_RELEASE, "deferred-release"), + (NXMEM_CAP_STRUCTURED_RELEASE, "structured-release"), + ]; + let mut parts: Vec<&str> = NAMES + .iter() + .filter(|(bit, _)| flags & bit != 0) + .map(|(_, name)| *name) + .collect(); + let unknown = flags & !NXMEM_CAP_KNOWN_MASK; + let unknown_text; + if unknown != 0 { + unknown_text = format!("unknown({unknown:#x})"); + parts.push(&unknown_text); + } + if parts.is_empty() { + return String::from("none"); + } + parts.join(", ") +} + +// ─── Plugin-side negotiation ──────────────────────────────────────────────── + +/// The plugin half of the handshake. +/// +/// A plugin's `NxmemNegotiate` export can delegate here when it speaks the +/// version this crate was compiled for. A plugin that deliberately speaks an +/// older minor calls [`negotiate_as`] instead. +/// +/// # Safety +/// +/// `request` and `response_out` must be valid, non-null, and correctly aligned +/// for the duration of the call. +pub unsafe fn negotiate( + request: *const NxmemNegotiateRequest, + response_out: *mut NxmemNegotiateResponse, +) -> NxmemStatus { + // SAFETY: delegated unchanged to this function's contract. + unsafe { + negotiate_as( + request, + response_out, + NxmemVersionRange::current(), + NXMEM_CAP_KNOWN_MASK, + ) + } +} + +/// The plugin half of the handshake for a participant that advertises +/// `plugin_range` and `plugin_capabilities` rather than this build's own. +/// +/// This is what an intentionally-older plugin (or a compatibility fixture) +/// uses. It is the same production logic; only the advertised range differs. +/// +/// # Safety +/// +/// `request` and `response_out` must be valid, non-null, and correctly aligned +/// for the duration of the call. +pub unsafe fn negotiate_as( + request: *const NxmemNegotiateRequest, + response_out: *mut NxmemNegotiateResponse, + plugin_range: NxmemVersionRange, + plugin_capabilities: u64, +) -> NxmemStatus { + if request.is_null() || response_out.is_null() { + return NxmemStatus::with_message( + NxmemStatusCode::InvalidArgument, + "NxmemNegotiate: request and response pointers must be non-null", + ); + } + + // SAFETY: both pointers were checked for null and the caller guarantees + // validity and alignment. + let response = unsafe { &mut *response_out }; + *response = NxmemNegotiateResponse::zeroed(); + response.plugin_range = plugin_range; + + // The request must be at least large enough to contain the fields we read. + // SAFETY: `request` is non-null and valid per this function's contract; + // `struct_size` is the first field of the baseline prefix, which every + // version of this struct has. + let request_size = unsafe { (*request).struct_size } as usize; + if request_size < core::mem::size_of::() { + return NxmemStatus::with_message( + NxmemStatusCode::ShortStruct, + &format!( + "NxmemNegotiate: host request struct_size {request_size} is smaller than the \ + required {} bytes; rebuild the host against a matching nxmem ABI header", + core::mem::size_of::() + ), + ); + } + // SAFETY: the size check above proved the struct contains every field we + // are about to read. + let request = unsafe { &*request }; + + let major = plugin_range.major_max; + if major < request.host_range.major_min || major > request.host_range.major_max { + return NxmemStatus::with_message( + NxmemStatusCode::VersionMismatch, + &format!( + "NxmemNegotiate: plugin speaks nxmem major {major} but the host accepts only \ + [{}, {}]; rebuild the plugin against the host's nxmem ABI major version", + request.host_range.major_min, request.host_range.major_max + ), + ); + } + + let agreed_minor = plugin_range.minor_max.min(request.host_range.minor_max); + if agreed_minor < plugin_range.minor_min || agreed_minor < request.host_range.minor_min { + return NxmemStatus::with_message( + NxmemStatusCode::VersionMismatch, + &format!( + "NxmemNegotiate: no common nxmem minor version; plugin supports [{}, {}] and the \ + host supports [{}, {}]", + plugin_range.minor_min, + plugin_range.minor_max, + request.host_range.minor_min, + request.host_range.minor_max + ), + ); + } + + // Only offer what the host asked for, what this build knows, and what the + // agreed minor actually defines a prefix for. + let mut offered = plugin_capabilities & request.host_capability_flags & NXMEM_CAP_KNOWN_MASK; + for flag in [ + NXMEM_CAP_ALLOCATOR, + NXMEM_CAP_VIRTUAL_BACKING, + NXMEM_CAP_SHARED_MAPPING, + NXMEM_CAP_DEFERRED_RELEASE, + NXMEM_CAP_STRUCTURED_RELEASE, + ] { + if agreed_minor < capability_min_minor(flag) { + offered &= !flag; + } + } + + if offered & NXMEM_CAP_ALLOCATOR == 0 { + return NxmemStatus::with_message( + NxmemStatusCode::UnsupportedCapability, + "NxmemNegotiate: the ordinary allocator capability is mandatory and was not agreed", + ); + } + + response.agreed_major = major; + response.agreed_minor = agreed_minor; + response.capability_flags = offered; + NxmemStatus::ok() +} + +// ─── Host-side validation ─────────────────────────────────────────────────── + +/// Why a host refused a negotiation response. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct NegotiationRejection { + /// A stable, actionable explanation. + pub reason: String, +} + +impl core::fmt::Display for NegotiationRejection { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.write_str(&self.reason) + } +} + +/// Host-side validation of a plugin's negotiation response. +/// +/// Fails closed on: a major outside the host range, a minor above what the +/// host can speak, a minor below the host's floor, unknown capability bits, a +/// missing mandatory allocator capability, and a capability advertised at a +/// minor that does not define it. +pub fn validate_negotiation( + host_range: &NxmemVersionRange, + response: &NxmemNegotiateResponse, +) -> Result<(), NegotiationRejection> { + let reject = |reason: String| Err(NegotiationRejection { reason }); + + if response.agreed_major < host_range.major_min || response.agreed_major > host_range.major_max + { + return reject(format!( + "plugin agreed nxmem major {} but this host accepts only [{}, {}]", + response.agreed_major, host_range.major_min, host_range.major_max + )); + } + if response.agreed_minor > host_range.minor_max { + return reject(format!( + "plugin agreed nxmem minor {} but this host speaks at most minor {}; the host cannot \ + call vtable slots it does not know", + response.agreed_minor, host_range.minor_max + )); + } + if response.agreed_minor < host_range.minor_min { + return reject(format!( + "plugin agreed nxmem minor {} but this host no longer supports anything below minor \ + {}", + response.agreed_minor, host_range.minor_min + )); + } + let unknown = response.capability_flags & !NXMEM_CAP_KNOWN_MASK; + if unknown != 0 { + return reject(format!( + "plugin advertises unknown nxmem capability bits {unknown:#x}; refusing to load \ + (fail closed)" + )); + } + if response.capability_flags & NXMEM_CAP_ALLOCATOR == 0 { + return reject(String::from( + "plugin does not advertise the mandatory nxmem allocator capability", + )); + } + for flag in [ + NXMEM_CAP_VIRTUAL_BACKING, + NXMEM_CAP_SHARED_MAPPING, + NXMEM_CAP_DEFERRED_RELEASE, + NXMEM_CAP_STRUCTURED_RELEASE, + ] { + if response.capability_flags & flag != 0 && response.agreed_minor < capability_min_minor(flag) + { + return reject(format!( + "plugin advertises capability {} at nxmem minor {} but that capability is only \ + defined from minor {}", + describe_capabilities(flag), + response.agreed_minor, + capability_min_minor(flag) + )); + } + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn negotiate_current(request: &NxmemNegotiateRequest) -> (NxmemStatus, NxmemNegotiateResponse) { + let mut response = NxmemNegotiateResponse::zeroed(); + // SAFETY: both pointers reference live, aligned locals. + let status = unsafe { negotiate(request, &mut response) }; + (status, response) + } + + #[test] + fn matching_versions_agree_on_the_current_minor() { + let (status, response) = negotiate_current(&NxmemNegotiateRequest::current()); + assert!(status.is_ok(), "{}", status.describe()); + assert_eq!(response.agreed_major, NXMEM_ABI_VERSION_MAJOR); + assert_eq!(response.agreed_minor, NXMEM_ABI_VERSION_MINOR); + assert!(response.capability_flags & NXMEM_CAP_ALLOCATOR != 0); + assert!(validate_negotiation(&NxmemVersionRange::current(), &response).is_ok()); + } + + #[test] + fn a_major_outside_the_host_range_is_rejected_with_a_fix() { + let request = NxmemNegotiateRequest::with_range(NxmemVersionRange::exact(99, 0)); + let (status, response) = negotiate_current(&request); + assert_eq!(status.status_code(), Some(NxmemStatusCode::VersionMismatch)); + assert!(status.describe().contains("rebuild the plugin")); + // The plugin's own range is reported even on failure, for diagnostics. + assert_eq!(response.plugin_range, NxmemVersionRange::current()); + } + + #[test] + fn an_older_host_pins_the_agreed_minor_down() { + let request = NxmemNegotiateRequest::with_range(NxmemVersionRange { + major_min: NXMEM_ABI_VERSION_MAJOR, + major_max: NXMEM_ABI_VERSION_MAJOR, + minor_min: 0, + minor_max: 0, + }); + let (status, response) = negotiate_current(&request); + assert!(status.is_ok(), "{}", status.describe()); + assert_eq!(response.agreed_minor, 0); + assert_eq!( + response.capability_flags & NXMEM_CAP_STRUCTURED_RELEASE, + 0, + "a capability defined at minor 1 must not be offered at minor 0" + ); + } + + #[test] + fn an_older_plugin_is_accepted_by_the_current_host() { + let request = NxmemNegotiateRequest::current(); + let mut response = NxmemNegotiateResponse::zeroed(); + // SAFETY: both pointers reference live, aligned locals. + let status = unsafe { + negotiate_as( + &request, + &mut response, + NxmemVersionRange::exact(NXMEM_ABI_VERSION_MAJOR, 0), + NXMEM_CAP_ALLOCATOR | NXMEM_CAP_VIRTUAL_BACKING, + ) + }; + assert!(status.is_ok(), "{}", status.describe()); + assert_eq!(response.agreed_minor, 0); + assert!(validate_negotiation(&NxmemVersionRange::current(), &response).is_ok()); + } + + #[test] + fn a_short_request_struct_is_refused_before_any_field_is_read() { + let mut request = NxmemNegotiateRequest::current(); + request.struct_size = 4; + let (status, _) = negotiate_current(&request); + assert_eq!(status.status_code(), Some(NxmemStatusCode::ShortStruct)); + } + + #[test] + fn null_pointers_fail_closed() { + let mut response = NxmemNegotiateResponse::zeroed(); + // SAFETY: a null request is exactly what this test exercises; the + // response pointer is a live local. + let status = unsafe { negotiate(core::ptr::null(), &mut response) }; + assert_eq!(status.status_code(), Some(NxmemStatusCode::InvalidArgument)); + } + + #[test] + fn a_plugin_may_not_advertise_a_capability_the_host_did_not_offer() { + let mut request = NxmemNegotiateRequest::current(); + request.host_capability_flags = NXMEM_CAP_ALLOCATOR; + let (status, response) = negotiate_current(&request); + assert!(status.is_ok(), "{}", status.describe()); + assert_eq!(response.capability_flags, NXMEM_CAP_ALLOCATOR); + } + + #[test] + fn unknown_capability_bits_are_rejected_by_the_host() { + let mut response = NxmemNegotiateResponse::zeroed(); + response.agreed_major = NXMEM_ABI_VERSION_MAJOR; + response.agreed_minor = NXMEM_ABI_VERSION_MINOR; + response.capability_flags = NXMEM_CAP_ALLOCATOR | (1 << 62); + let rejection = validate_negotiation(&NxmemVersionRange::current(), &response) + .expect_err("unknown bits must fail closed"); + assert!(rejection.reason.contains("unknown nxmem capability bits")); + } + + #[test] + fn a_missing_allocator_capability_is_rejected_by_the_host() { + let mut response = NxmemNegotiateResponse::zeroed(); + response.agreed_major = NXMEM_ABI_VERSION_MAJOR; + response.agreed_minor = NXMEM_ABI_VERSION_MINOR; + response.capability_flags = NXMEM_CAP_VIRTUAL_BACKING; + let rejection = validate_negotiation(&NxmemVersionRange::current(), &response) + .expect_err("the allocator capability is mandatory"); + assert!(rejection.reason.contains("mandatory")); + } + + #[test] + fn a_capability_advertised_below_its_minor_is_rejected() { + let mut response = NxmemNegotiateResponse::zeroed(); + response.agreed_major = NXMEM_ABI_VERSION_MAJOR; + response.agreed_minor = 0; + response.capability_flags = NXMEM_CAP_ALLOCATOR | NXMEM_CAP_STRUCTURED_RELEASE; + let rejection = validate_negotiation(&NxmemVersionRange::current(), &response) + .expect_err("structured release does not exist at minor 0"); + assert!( + rejection.reason.contains("structured-release"), + "the rejection must name the offending capability, got {:?}", + rejection.reason + ); + assert!( + rejection.reason.contains("defined from minor"), + "the rejection must explain the minor floor, got {:?}", + rejection.reason + ); + } + + #[test] + fn a_newer_plugin_minor_is_rejected_by_an_older_host() { + let mut response = NxmemNegotiateResponse::zeroed(); + response.agreed_major = NXMEM_ABI_VERSION_MAJOR; + response.agreed_minor = NXMEM_ABI_VERSION_MINOR + 5; + response.capability_flags = NXMEM_CAP_ALLOCATOR; + let rejection = validate_negotiation(&NxmemVersionRange::current(), &response) + .expect_err("a host cannot call slots it does not know"); + assert!(rejection.reason.contains("does not know")); + } + + #[test] + fn capability_names_are_stable_and_report_unknown_bits() { + assert_eq!(describe_capabilities(0), "none"); + assert_eq!( + describe_capabilities(NXMEM_CAP_ALLOCATOR | NXMEM_CAP_SHARED_MAPPING), + "allocator, shared-mapping" + ); + assert!(describe_capabilities(1 << 62).contains("unknown(0x4000000000000000)")); + } +} diff --git a/crates/onnx-runtime-memory-abi/src/vtable.rs b/crates/onnx-runtime-memory-abi/src/vtable.rs new file mode 100644 index 0000000000..aeff40327b --- /dev/null +++ b/crates/onnx-runtime-memory-abi/src/vtable.rs @@ -0,0 +1,945 @@ +//! Vtables and the prefix-negotiation rules that make them safe to read. +//! +//! # Reading an untrusted vtable +//! +//! A participant may hand over a struct that is **smaller** than this build's +//! definition (it was compiled at an older minor) or **larger** (newer minor, +//! same major). Dereferencing the whole struct in either case is wrong: the +//! small case reads past the end of the allocation. +//! +//! The only supported way to read one is [`NxmemAllocatorVtable::read_prefix`] +//! and its siblings. They: +//! +//! 1. reject a null or misaligned pointer; +//! 2. read `struct_size` — always the first field, at offset 0, in every +//! version; +//! 3. reject a `struct_size` below the required prefix for the level the +//! vtable declares ([`crate::NxmemStatusCode::ShortStruct`]); +//! 4. copy `min(struct_size, size_of::())` bytes into a zero-initialised +//! local, so trailing slots this build knows about but the sender does not +//! read back as `None`/null, and trailing bytes this build does not know +//! about are ignored. +//! +//! Every optional slot is a nullable function pointer. Null means the +//! capability is absent and must surface as +//! [`crate::NxmemStatusCode::UnsupportedCapability`] — never as a successful +//! no-op. +//! +//! # Per-vtable levels +//! +//! Negotiation fixes the **ceiling**. Each vtable independently declares the +//! level it actually implements in `abi_minor`, so one plugin may ship a +//! current mechanism next to one still built at the baseline prefix. A vtable +//! declaring a level *above* the negotiated ceiling is rejected: the host did +//! not agree to interpret those slots. + +use core::ffi::c_void; + +use crate::status::{NxmemStatus, NxmemStatusCode}; +use crate::types::{ + NxmemAllocRequest, NxmemAllocResult, NxmemAllocation, NxmemDeviceId, NxmemHostCallbacks, + NxmemRangeRequest, NxmemReleaseOutcome, NxmemSharedPrefixCommitInfo, + NxmemSharedPrefixCommitRequest, NxmemSharedPrefixHandle, +}; + +/// The request a host passes to `open_allocator`. +/// +/// # Ownership +/// +/// `callbacks` is borrowed for the **lifetime of the opened allocator**, not +/// just for the call. See [`NxmemHostCallbacks`]. +#[repr(C)] +#[derive(Debug, Clone, Copy)] +pub struct NxmemOpenRequest { + /// `size_of` this struct as the host defines it. + pub struct_size: u32, + /// The minor version negotiated for this module. Acts as the ceiling for + /// every vtable the plugin returns. + pub abi_minor: u32, + /// The device the host wants an allocator for. + pub device: NxmemDeviceId, + /// Capabilities the host requires. The plugin must fail the open rather + /// than return an allocator missing one of them. + pub required_capability_flags: u64, + /// Host callbacks, borrowed for the allocator's lifetime. May be null when + /// the host offers none. + pub callbacks: *const NxmemHostCallbacks, +} + +impl NxmemOpenRequest { + /// An open request with the correct `struct_size`. + pub fn new( + abi_minor: u32, + device: NxmemDeviceId, + required_capability_flags: u64, + callbacks: *const NxmemHostCallbacks, + ) -> Self { + Self { + struct_size: core::mem::size_of::() as u32, + abi_minor, + device, + required_capability_flags, + callbacks, + } + } +} + +/// Reject a vtable pointer that is null or misaligned for `T`. +fn check_vtable_ptr(ptr: *const T, what: &str) -> Result<(), NxmemStatus> { + if ptr.is_null() { + return Err(NxmemStatus::with_message( + NxmemStatusCode::InvalidArgument, + &format!("nxmem: {what} vtable pointer is null"), + )); + } + if (ptr as usize) % core::mem::align_of::() != 0 { + return Err(NxmemStatus::with_message( + NxmemStatusCode::InvalidArgument, + &format!( + "nxmem: {what} vtable pointer {ptr:p} is not aligned to {} bytes", + core::mem::align_of::() + ), + )); + } + Ok(()) +} + +/// Copy the readable prefix of `ptr` into a zero-initialised `T`. +/// +/// # Safety +/// +/// `ptr` must be non-null, aligned for `T`, and point to at least +/// `min(declared_size, size_of::())` readable bytes. Callers reach this +/// through the `read_prefix` helpers, which enforce those conditions. +unsafe fn copy_prefix(ptr: *const T, declared_size: usize) -> T { + let readable = declared_size.min(core::mem::size_of::()); + let mut out = core::mem::MaybeUninit::::zeroed(); + // SAFETY: the destination is a fresh `T`-sized allocation; the source is + // valid for `readable` bytes by this function's contract; the regions + // cannot overlap because `out` is a local. + unsafe { + core::ptr::copy_nonoverlapping( + ptr as *const u8, + out.as_mut_ptr() as *mut u8, + readable, + ); + out.assume_init() + } +} + +/// Read the leading `struct_size` field, which lives at offset 0 in every +/// version of every `nxmem` struct. +/// +/// # Safety +/// +/// `ptr` must be non-null, aligned, and point to at least four readable bytes. +unsafe fn read_struct_size(ptr: *const T) -> u32 { + // SAFETY: `struct_size` is the first field of every nxmem struct, so four + // readable bytes at `ptr` are guaranteed by this function's contract. + unsafe { core::ptr::read_unaligned(ptr as *const u32) } +} + +fn short_struct(what: &str, declared: usize, required: usize, minor: u32) -> NxmemStatus { + NxmemStatus::with_message( + NxmemStatusCode::ShortStruct, + &format!( + "nxmem: {what} declares struct_size {declared} but nxmem minor {minor} requires at \ + least {required} bytes; rebuild the plugin against a matching nxmem ABI header" + ), + ) +} + +fn level_too_new(what: &str, declared: u32, ceiling: u32) -> NxmemStatus { + NxmemStatus::with_message( + NxmemStatusCode::VersionMismatch, + &format!( + "nxmem: {what} declares nxmem minor {declared} but the negotiated ceiling is \ + {ceiling}; the host will not interpret slots it did not agree to" + ), + ) +} + +// ─── Allocator vtable ─────────────────────────────────────────────────────── + +/// The ordinary allocator mechanism plus its optional capability slots. +/// +/// # Ownership +/// +/// The vtable and its `ctx` are owned by the plugin. The host takes one +/// reference on open and gives it back with exactly one `release`. `retain` +/// takes an extra reference; every `retain` must be paired with a `release`. +/// The plugin must keep the vtable, `ctx`, and `name` alive until the final +/// `release` returns. +#[repr(C)] +#[derive(Debug, Clone, Copy)] +pub struct NxmemAllocatorVtable { + /// `size_of` this struct as the plugin defines it. + pub struct_size: u32, + /// The nxmem minor level this vtable implements. Must not exceed the + /// negotiated ceiling. + pub abi_minor: u32, + /// Non-zero identity of this mechanism instance. Echoed by the host in + /// every call and re-checked by the plugin. + pub mechanism_id: u64, + /// The device this mechanism serves. + pub device: NxmemDeviceId, + /// Capabilities this mechanism actually provides. + pub capability_flags: u64, + /// NUL-terminated UTF-8 name, valid until the final `release` returns. + pub name: *const u8, + /// Opaque plugin state passed to every slot. + pub ctx: *mut c_void, + + // ─── required at minor 0 ─── + /// Ordinary allocation. Required. + pub allocate: Option< + unsafe extern "C" fn( + ctx: *mut c_void, + request: *const NxmemAllocRequest, + result_out: *mut NxmemAllocResult, + ) -> NxmemStatus, + >, + /// Terminal release reporting only newly unmapped bytes. Required; it is + /// the fallback when structured release is absent. + pub deallocate: Option< + unsafe extern "C" fn( + ctx: *mut c_void, + allocation: *const NxmemAllocation, + unmapped_bytes_out: *mut u64, + ) -> NxmemStatus, + >, + /// Take an extra reference on the mechanism. Required. + pub retain: Option, + /// Drop one reference. Required. The plugin destroys the mechanism when + /// the count reaches zero **and** no queued release still names it. + pub release: Option, + + // ─── optional at minor 0 ─── + /// Lazy backing capability, or null when unsupported. + pub virtual_backing: *const NxmemVirtualBackingVtable, + /// Shared physical mapping capability, or null when unsupported. + pub shared_mapping: *const NxmemSharedMappingVtable, + /// Queue a stream-ordered release instead of freeing now. Writes a ticket. + pub enqueue_release: Option< + unsafe extern "C" fn( + ctx: *mut c_void, + allocation: *const NxmemAllocation, + ticket_out: *mut u64, + ) -> NxmemStatus, + >, + /// Retire up to `max` queued releases, invoking the host's + /// `release_completed` callback for each, in enqueue order. + pub drain_releases: Option< + unsafe extern "C" fn(ctx: *mut c_void, max: u64, retired_out: *mut u64) -> NxmemStatus, + >, + /// How many releases are queued and not yet retired. + pub pending_release_count: + Option NxmemStatus>, + + // ─── added at minor 1 ─── + /// Terminal release reporting a structured outcome that can distinguish + /// complete, quarantined, and failed. Null below minor 1. + pub release_allocation: Option< + unsafe extern "C" fn( + ctx: *mut c_void, + allocation: *const NxmemAllocation, + outcome_out: *mut NxmemReleaseOutcome, + ) -> NxmemStatus, + >, +} + +impl NxmemAllocatorVtable { + /// Bytes a minor-0 allocator vtable must provide. + pub const MIN_STRUCT_SIZE_MINOR_0: usize = + core::mem::offset_of!(Self, pending_release_count) + core::mem::size_of::(); + + /// Bytes a minor-1 allocator vtable must provide. + pub const MIN_STRUCT_SIZE_MINOR_1: usize = + core::mem::offset_of!(Self, release_allocation) + core::mem::size_of::(); + + /// The prefix required at `minor`. + pub const fn required_struct_size(minor: u32) -> usize { + if minor == 0 { + Self::MIN_STRUCT_SIZE_MINOR_0 + } else { + Self::MIN_STRUCT_SIZE_MINOR_1 + } + } + + /// A zeroed vtable, for a plugin to fill in. + pub const fn zeroed() -> Self { + Self { + struct_size: core::mem::size_of::() as u32, + abi_minor: 0, + mechanism_id: 0, + device: NxmemDeviceId { tier: 0, index: 0 }, + capability_flags: 0, + name: core::ptr::null(), + ctx: core::ptr::null_mut(), + allocate: None, + deallocate: None, + retain: None, + release: None, + virtual_backing: core::ptr::null(), + shared_mapping: core::ptr::null(), + enqueue_release: None, + drain_releases: None, + pending_release_count: None, + release_allocation: None, + } + } + + /// Safely read the readable prefix of an allocator vtable. + /// + /// `negotiated_minor` is the ceiling agreed at load time. + /// + /// # Safety + /// + /// `ptr` must either be null (rejected) or point to a live vtable whose + /// first `struct_size` bytes are readable for the duration of the call. + pub unsafe fn read_prefix(ptr: *const Self, negotiated_minor: u32) -> Result { + check_vtable_ptr(ptr, "allocator")?; + // SAFETY: `ptr` was proved non-null and aligned; `struct_size` is the + // first field of every version of this struct. + let declared_size = unsafe { read_struct_size(ptr) } as usize; + if declared_size < core::mem::offset_of!(Self, allocate) { + // Too small even to name its own level. + return Err(short_struct( + "allocator vtable", + declared_size, + Self::required_struct_size(0), + 0, + )); + } + // SAFETY: the size check above proved `abi_minor` (offset 4) is inside + // the readable prefix. + let declared_minor = unsafe { core::ptr::read_unaligned(ptr.cast::().add(1)) }; + if declared_minor > negotiated_minor { + return Err(level_too_new( + "allocator vtable", + declared_minor, + negotiated_minor, + )); + } + let required = Self::required_struct_size(declared_minor); + if declared_size < required { + return Err(short_struct( + "allocator vtable", + declared_size, + required, + declared_minor, + )); + } + // SAFETY: `declared_size` bytes are readable and we copy no more. + let mut vtable = unsafe { copy_prefix(ptr, declared_size) }; + // A slot the sender does not define must never be called, even if this + // build knows about it. + if declared_minor < 1 { + vtable.release_allocation = None; + } + vtable.validate_required()?; + Ok(vtable) + } + + /// Reject a vtable whose required slots are null. + /// + /// A plugin that memsets its vtable to zero must fail here rather than + /// have the host call through a null pointer. + pub fn validate_required(&self) -> Result<(), NxmemStatus> { + for (slot, present) in [ + ("allocate", self.allocate.is_some()), + ("deallocate", self.deallocate.is_some()), + ("retain", self.retain.is_some()), + ("release", self.release.is_some()), + ] { + if !present { + return Err(NxmemStatus::with_message( + NxmemStatusCode::ShortStruct, + &format!( + "nxmem: allocator vtable '{slot}' slot is null but it is required at every \ + nxmem version" + ), + )); + } + } + if self.mechanism_id == 0 { + return Err(NxmemStatus::with_message( + NxmemStatusCode::InvalidArgument, + "nxmem: allocator vtable mechanism_id must be non-zero so cross-mechanism misuse \ + can be rejected", + )); + } + Ok(()) + } +} + +// ─── Virtual backing vtable ───────────────────────────────────────────────── + +/// Lazy reserve/commit/decommit. Optional. +/// +/// Terminal release never happens here: it always goes back through the +/// owning [`NxmemAllocatorVtable`]. +/// +/// # Ownership +/// +/// Owned by the plugin and kept alive by the allocator that exposed it. It has +/// no independent reference count; releasing the allocator invalidates it. +#[repr(C)] +#[derive(Debug, Clone, Copy)] +pub struct NxmemVirtualBackingVtable { + /// `size_of` this struct as the plugin defines it. + pub struct_size: u32, + /// The nxmem minor level this vtable implements. + pub abi_minor: u32, + /// Must equal the owning allocator's `mechanism_id`. + pub mechanism_id: u64, + /// Opaque plugin state passed to every slot. + pub ctx: *mut c_void, + /// Reserve an allocation, committing only the requested ranges. + pub allocate_committed: Option< + unsafe extern "C" fn( + ctx: *mut c_void, + request: *const NxmemAllocRequest, + result_out: *mut NxmemAllocResult, + ) -> NxmemStatus, + >, + /// Commit one span of an existing allocation. + pub commit_range: + Option NxmemStatus>, + /// Decommit one span, writing the bytes actually unmapped. + pub decommit_range: Option< + unsafe extern "C" fn( + ctx: *mut c_void, + request: *const NxmemRangeRequest, + unmapped_out: *mut u64, + ) -> NxmemStatus, + >, + /// Conservative mapped-byte cost of committing a batch of spans. Only the + /// mechanism knows its granularity and which spans share granules. + pub mapped_bytes_for_ranges: Option< + unsafe extern "C" fn( + ctx: *mut c_void, + requests: *const NxmemRangeRequest, + count: u64, + mapped_out: *mut u64, + ) -> NxmemStatus, + >, + /// Conservative mapped-byte cost of a whole allocation. + pub mapped_bytes_for_allocation: Option< + unsafe extern "C" fn( + ctx: *mut c_void, + request: *const NxmemAllocRequest, + mapped_out: *mut u64, + ) -> NxmemStatus, + >, + /// Bytes currently committed inside one allocation. + pub committed_bytes: Option< + unsafe extern "C" fn( + ctx: *mut c_void, + allocation: *const NxmemAllocation, + committed_out: *mut u64, + ) -> NxmemStatus, + >, +} + +impl NxmemVirtualBackingVtable { + /// Bytes a minor-0 virtual-backing vtable must provide. + pub const MIN_STRUCT_SIZE_MINOR_0: usize = + core::mem::offset_of!(Self, committed_bytes) + core::mem::size_of::(); + + /// A zeroed vtable, for a plugin to fill in. + pub const fn zeroed() -> Self { + Self { + struct_size: core::mem::size_of::() as u32, + abi_minor: 0, + mechanism_id: 0, + ctx: core::ptr::null_mut(), + allocate_committed: None, + commit_range: None, + decommit_range: None, + mapped_bytes_for_ranges: None, + mapped_bytes_for_allocation: None, + committed_bytes: None, + } + } + + /// Safely read the readable prefix of a virtual-backing vtable. + /// + /// # Safety + /// + /// Same contract as [`NxmemAllocatorVtable::read_prefix`]. + pub unsafe fn read_prefix( + ptr: *const Self, + negotiated_minor: u32, + owner_mechanism_id: u64, + ) -> Result { + check_vtable_ptr(ptr, "virtual-backing")?; + // SAFETY: `ptr` was proved non-null and aligned. + let declared_size = unsafe { read_struct_size(ptr) } as usize; + if declared_size < Self::MIN_STRUCT_SIZE_MINOR_0 { + return Err(short_struct( + "virtual-backing vtable", + declared_size, + Self::MIN_STRUCT_SIZE_MINOR_0, + 0, + )); + } + // SAFETY: the size check proved `abi_minor` is readable. + let declared_minor = unsafe { core::ptr::read_unaligned(ptr.cast::().add(1)) }; + if declared_minor > negotiated_minor { + return Err(level_too_new( + "virtual-backing vtable", + declared_minor, + negotiated_minor, + )); + } + // SAFETY: `declared_size` bytes are readable and we copy no more. + let vtable = unsafe { copy_prefix(ptr, declared_size) }; + if vtable.mechanism_id != owner_mechanism_id { + return Err(NxmemStatus::with_message( + NxmemStatusCode::WrongMechanism, + &format!( + "nxmem: virtual-backing vtable reports mechanism {} but was discovered from \ + mechanism {owner_mechanism_id}; a capability must belong to the allocator \ + that exposed it", + vtable.mechanism_id + ), + )); + } + Ok(vtable) + } +} + +// ─── Shared mapping vtable ────────────────────────────────────────────────── + +/// Reusable shared physical handles and prefix mapping. Optional and +/// independent of virtual backing. +/// +/// # Ownership +/// +/// Each prefix produced by `create_shared_prefix` is plugin-owned and must be +/// given back with exactly one `release_shared_prefix` on the same mechanism. +/// `retain_shared_prefix` adds a reference; every retain needs a release. +#[repr(C)] +#[derive(Debug, Clone, Copy)] +pub struct NxmemSharedMappingVtable { + /// `size_of` this struct as the plugin defines it. + pub struct_size: u32, + /// The nxmem minor level this vtable implements. + pub abi_minor: u32, + /// Must equal the owning allocator's `mechanism_id`. + pub mechanism_id: u64, + /// Opaque plugin state passed to every slot. + pub ctx: *mut c_void, + /// Create a shared physical prefix. + pub create_shared_prefix: Option< + unsafe extern "C" fn( + ctx: *mut c_void, + mechanism_id: u64, + bytes: u64, + handle_out: *mut NxmemSharedPrefixHandle, + ) -> NxmemStatus, + >, + /// Add a reference to a prefix. + pub retain_shared_prefix: Option< + unsafe extern "C" fn( + ctx: *mut c_void, + handle: *const NxmemSharedPrefixHandle, + ) -> NxmemStatus, + >, + /// Drop a reference to a prefix. The plugin frees the physical bytes when + /// the final reference and the final mapping retire. + pub release_shared_prefix: Option< + unsafe extern "C" fn( + ctx: *mut c_void, + handle: *const NxmemSharedPrefixHandle, + ) -> NxmemStatus, + >, + /// Incremental owned physical cost of admitting one more mapping. A + /// foreign, wrong-device, or wrong-mechanism prefix must be rejected here + /// rather than reported as costing zero. + pub incremental_owned_bytes: Option< + unsafe extern "C" fn( + ctx: *mut c_void, + handle: *const NxmemSharedPrefixHandle, + bytes_out: *mut u64, + ) -> NxmemStatus, + >, + /// Map a prefix into an allocation. + pub commit_shared_prefix: Option< + unsafe extern "C" fn( + ctx: *mut c_void, + request: *const NxmemSharedPrefixCommitRequest, + info_out: *mut NxmemSharedPrefixCommitInfo, + ) -> NxmemStatus, + >, +} + +impl NxmemSharedMappingVtable { + /// Bytes a minor-0 shared-mapping vtable must provide. + pub const MIN_STRUCT_SIZE_MINOR_0: usize = + core::mem::offset_of!(Self, commit_shared_prefix) + core::mem::size_of::(); + + /// A zeroed vtable, for a plugin to fill in. + pub const fn zeroed() -> Self { + Self { + struct_size: core::mem::size_of::() as u32, + abi_minor: 0, + mechanism_id: 0, + ctx: core::ptr::null_mut(), + create_shared_prefix: None, + retain_shared_prefix: None, + release_shared_prefix: None, + incremental_owned_bytes: None, + commit_shared_prefix: None, + } + } + + /// Safely read the readable prefix of a shared-mapping vtable. + /// + /// # Safety + /// + /// Same contract as [`NxmemAllocatorVtable::read_prefix`]. + pub unsafe fn read_prefix( + ptr: *const Self, + negotiated_minor: u32, + owner_mechanism_id: u64, + ) -> Result { + check_vtable_ptr(ptr, "shared-mapping")?; + // SAFETY: `ptr` was proved non-null and aligned. + let declared_size = unsafe { read_struct_size(ptr) } as usize; + if declared_size < Self::MIN_STRUCT_SIZE_MINOR_0 { + return Err(short_struct( + "shared-mapping vtable", + declared_size, + Self::MIN_STRUCT_SIZE_MINOR_0, + 0, + )); + } + // SAFETY: the size check proved `abi_minor` is readable. + let declared_minor = unsafe { core::ptr::read_unaligned(ptr.cast::().add(1)) }; + if declared_minor > negotiated_minor { + return Err(level_too_new( + "shared-mapping vtable", + declared_minor, + negotiated_minor, + )); + } + // SAFETY: `declared_size` bytes are readable and we copy no more. + let vtable = unsafe { copy_prefix(ptr, declared_size) }; + if vtable.mechanism_id != owner_mechanism_id { + return Err(NxmemStatus::with_message( + NxmemStatusCode::WrongMechanism, + &format!( + "nxmem: shared-mapping vtable reports mechanism {} but was discovered from \ + mechanism {owner_mechanism_id}; a capability must belong to the allocator \ + that exposed it", + vtable.mechanism_id + ), + )); + } + Ok(vtable) + } +} + +// ─── Factory vtable ───────────────────────────────────────────────────────── + +/// One named allocator mechanism a plugin can open. +/// +/// # Ownership +/// +/// The host owns every factory returned by `NxmemCreateAllocatorFactories` and +/// must call `release` on each exactly once. A factory must stay usable until +/// its own `release`, and releasing a factory must not invalidate allocators +/// already opened from it. +#[repr(C)] +#[derive(Debug, Clone, Copy)] +pub struct NxmemAllocatorFactoryVtable { + /// `size_of` this struct as the plugin defines it. + pub struct_size: u32, + /// The nxmem minor level this vtable implements. + pub abi_minor: u32, + /// NUL-terminated UTF-8 name, valid until `release` returns. Names the + /// mechanism, and is how a host selects one of several. + pub name: *const u8, + /// The device this factory serves. + pub device: NxmemDeviceId, + /// Capabilities the opened allocator will provide. + pub capability_flags: u64, + /// Opaque plugin state passed to every slot. + pub ctx: *mut c_void, + /// Open an allocator instance. Required. + pub open_allocator: Option< + unsafe extern "C" fn( + ctx: *mut c_void, + request: *const NxmemOpenRequest, + allocator_out: *mut *const NxmemAllocatorVtable, + ) -> NxmemStatus, + >, + /// Give the factory back. Required. + pub release: Option, +} + +impl NxmemAllocatorFactoryVtable { + /// Bytes a minor-0 factory vtable must provide. + pub const MIN_STRUCT_SIZE_MINOR_0: usize = + core::mem::offset_of!(Self, release) + core::mem::size_of::(); + + /// A zeroed vtable, for a plugin to fill in. + pub const fn zeroed() -> Self { + Self { + struct_size: core::mem::size_of::() as u32, + abi_minor: 0, + name: core::ptr::null(), + device: NxmemDeviceId { tier: 0, index: 0 }, + capability_flags: 0, + ctx: core::ptr::null_mut(), + open_allocator: None, + release: None, + } + } + + /// Safely read the readable prefix of a factory vtable. + /// + /// # Safety + /// + /// Same contract as [`NxmemAllocatorVtable::read_prefix`]. + pub unsafe fn read_prefix(ptr: *const Self, negotiated_minor: u32) -> Result { + check_vtable_ptr(ptr, "allocator-factory")?; + // SAFETY: `ptr` was proved non-null and aligned. + let declared_size = unsafe { read_struct_size(ptr) } as usize; + if declared_size < Self::MIN_STRUCT_SIZE_MINOR_0 { + return Err(short_struct( + "allocator-factory vtable", + declared_size, + Self::MIN_STRUCT_SIZE_MINOR_0, + 0, + )); + } + // SAFETY: the size check proved `abi_minor` is readable. + let declared_minor = unsafe { core::ptr::read_unaligned(ptr.cast::().add(1)) }; + if declared_minor > negotiated_minor { + return Err(level_too_new( + "allocator-factory vtable", + declared_minor, + negotiated_minor, + )); + } + // SAFETY: `declared_size` bytes are readable and we copy no more. + let vtable = unsafe { copy_prefix(ptr, declared_size) }; + if vtable.open_allocator.is_none() || vtable.release.is_none() { + return Err(NxmemStatus::with_message( + NxmemStatusCode::ShortStruct, + "nxmem: allocator-factory vtable must provide both 'open_allocator' and 'release'", + )); + } + Ok(vtable) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + unsafe extern "C" fn stub_allocate( + _ctx: *mut c_void, + _request: *const NxmemAllocRequest, + _result: *mut NxmemAllocResult, + ) -> NxmemStatus { + NxmemStatus::ok() + } + + unsafe extern "C" fn stub_deallocate( + _ctx: *mut c_void, + _allocation: *const NxmemAllocation, + _unmapped: *mut u64, + ) -> NxmemStatus { + NxmemStatus::ok() + } + + unsafe extern "C" fn stub_void(_ctx: *mut c_void) {} + + unsafe extern "C" fn stub_release_allocation( + _ctx: *mut c_void, + _allocation: *const NxmemAllocation, + _outcome: *mut NxmemReleaseOutcome, + ) -> NxmemStatus { + NxmemStatus::ok() + } + + fn conforming_allocator(minor: u32) -> NxmemAllocatorVtable { + NxmemAllocatorVtable { + struct_size: NxmemAllocatorVtable::required_struct_size(minor) as u32, + abi_minor: minor, + mechanism_id: 42, + device: NxmemDeviceId::HOST, + capability_flags: 1, + allocate: Some(stub_allocate), + deallocate: Some(stub_deallocate), + retain: Some(stub_void), + release: Some(stub_void), + release_allocation: Some(stub_release_allocation), + ..NxmemAllocatorVtable::zeroed() + } + } + + #[test] + fn the_minor_1_prefix_is_larger_than_the_minor_0_prefix() { + assert!( + NxmemAllocatorVtable::MIN_STRUCT_SIZE_MINOR_1 + > NxmemAllocatorVtable::MIN_STRUCT_SIZE_MINOR_0 + ); + assert_eq!( + NxmemAllocatorVtable::MIN_STRUCT_SIZE_MINOR_1, + core::mem::size_of::() + ); + } + + #[test] + fn a_current_vtable_reads_back_intact() { + let vtable = conforming_allocator(1); + // SAFETY: `vtable` is a live, aligned local of exactly this type. + let read = unsafe { NxmemAllocatorVtable::read_prefix(&vtable, 1) }.expect("valid vtable"); + assert_eq!(read.mechanism_id, 42); + assert!(read.release_allocation.is_some()); + } + + #[test] + fn a_minor_0_vtable_reads_back_without_the_minor_1_slot() { + let mut vtable = conforming_allocator(0); + // A baseline plugin does not even define the trailing slot; simulate + // that by declaring the smaller size while leaving stale bytes behind. + vtable.struct_size = NxmemAllocatorVtable::MIN_STRUCT_SIZE_MINOR_0 as u32; + // SAFETY: `vtable` is a live, aligned local of exactly this type, and + // we only ask to read the smaller declared prefix. + let read = unsafe { NxmemAllocatorVtable::read_prefix(&vtable, 1) }.expect("valid vtable"); + assert_eq!(read.abi_minor, 0); + assert!( + read.release_allocation.is_none(), + "a slot the sender does not define must never be called" + ); + assert!(read.allocate.is_some(), "the baseline prefix survives"); + } + + #[test] + fn a_short_vtable_is_refused() { + let mut vtable = conforming_allocator(0); + vtable.struct_size = (NxmemAllocatorVtable::MIN_STRUCT_SIZE_MINOR_0 - 1) as u32; + // SAFETY: `vtable` is a live, aligned local; `read_prefix` reads only + // `struct_size` and `abi_minor` before rejecting. + let status = + unsafe { NxmemAllocatorVtable::read_prefix(&vtable, 1) }.expect_err("short struct"); + assert_eq!(status.status_code(), Some(NxmemStatusCode::ShortStruct)); + assert!(status.describe().contains("rebuild the plugin")); + } + + #[test] + fn a_vtable_above_the_negotiated_ceiling_is_refused() { + let vtable = conforming_allocator(1); + // SAFETY: `vtable` is a live, aligned local. + let status = + unsafe { NxmemAllocatorVtable::read_prefix(&vtable, 0) }.expect_err("too new"); + assert_eq!(status.status_code(), Some(NxmemStatusCode::VersionMismatch)); + assert!(status.describe().contains("negotiated ceiling")); + } + + #[test] + fn a_zeroed_vtable_never_produces_a_null_call() { + let mut vtable = NxmemAllocatorVtable::zeroed(); + vtable.mechanism_id = 1; + // SAFETY: `vtable` is a live, aligned local. + let status = unsafe { NxmemAllocatorVtable::read_prefix(&vtable, 1) } + .expect_err("required slots are null"); + assert_eq!(status.status_code(), Some(NxmemStatusCode::ShortStruct)); + assert!(status.describe().contains("'allocate' slot is null")); + } + + #[test] + fn a_zero_mechanism_id_is_refused() { + let mut vtable = conforming_allocator(1); + vtable.mechanism_id = 0; + // SAFETY: `vtable` is a live, aligned local. + let status = unsafe { NxmemAllocatorVtable::read_prefix(&vtable, 1) } + .expect_err("mechanism_id must be non-zero"); + assert_eq!(status.status_code(), Some(NxmemStatusCode::InvalidArgument)); + } + + #[test] + fn a_null_vtable_pointer_is_refused() { + // SAFETY: a null pointer is exactly what this test exercises. + let status = unsafe { NxmemAllocatorVtable::read_prefix(core::ptr::null(), 1) } + .expect_err("null pointer"); + assert_eq!(status.status_code(), Some(NxmemStatusCode::InvalidArgument)); + } + + #[test] + fn a_misaligned_vtable_pointer_is_refused() { + let vtable = conforming_allocator(1); + let misaligned = (&vtable as *const NxmemAllocatorVtable as usize + 1) + as *const NxmemAllocatorVtable; + // SAFETY: `read_prefix` rejects on alignment before dereferencing. + let status = + unsafe { NxmemAllocatorVtable::read_prefix(misaligned, 1) }.expect_err("misaligned"); + assert_eq!(status.status_code(), Some(NxmemStatusCode::InvalidArgument)); + } + + #[test] + fn a_capability_vtable_from_a_foreign_mechanism_is_refused() { + let backing = NxmemVirtualBackingVtable { + mechanism_id: 999, + ..NxmemVirtualBackingVtable::zeroed() + }; + // SAFETY: `backing` is a live, aligned local. + let status = unsafe { NxmemVirtualBackingVtable::read_prefix(&backing, 1, 42) } + .expect_err("mechanism mismatch"); + assert_eq!(status.status_code(), Some(NxmemStatusCode::WrongMechanism)); + + let shared = NxmemSharedMappingVtable { + mechanism_id: 999, + ..NxmemSharedMappingVtable::zeroed() + }; + // SAFETY: `shared` is a live, aligned local. + let status = unsafe { NxmemSharedMappingVtable::read_prefix(&shared, 1, 42) } + .expect_err("mechanism mismatch"); + assert_eq!(status.status_code(), Some(NxmemStatusCode::WrongMechanism)); + } + + #[test] + fn a_factory_missing_a_required_slot_is_refused() { + let factory = NxmemAllocatorFactoryVtable::zeroed(); + // SAFETY: `factory` is a live, aligned local. + let status = unsafe { NxmemAllocatorFactoryVtable::read_prefix(&factory, 1) } + .expect_err("required slots are null"); + assert_eq!(status.status_code(), Some(NxmemStatusCode::ShortStruct)); + } + + /// A newer plugin sends a larger struct. This build must ignore the extra + /// bytes rather than reject, and must not read them. + #[test] + fn a_larger_struct_from_a_newer_peer_is_tolerated() { + #[repr(C)] + struct Extended { + base: NxmemAllocatorVtable, + future_slot: usize, + } + let extended = Extended { + base: NxmemAllocatorVtable { + struct_size: core::mem::size_of::() as u32, + ..conforming_allocator(1) + }, + future_slot: 0xdead_beef, + }; + // SAFETY: the pointer names a live `Extended` whose leading bytes are + // a valid `NxmemAllocatorVtable`; `read_prefix` copies at most + // `size_of::()` bytes. + let read = unsafe { + NxmemAllocatorVtable::read_prefix( + (&extended as *const Extended).cast::(), + 1, + ) + } + .expect("a larger struct at the same major is compatible"); + assert_eq!(read.mechanism_id, 42); + assert!(read.release_allocation.is_some()); + } +} diff --git a/crates/onnx-runtime-memory-host/Cargo.toml b/crates/onnx-runtime-memory-host/Cargo.toml new file mode 100644 index 0000000000..fc0f828fed --- /dev/null +++ b/crates/onnx-runtime-memory-host/Cargo.toml @@ -0,0 +1,21 @@ +[package] +name = "onnx-runtime-memory-host" +description = "Dynamic loader and Rust-trait adapter for nxmem memory plugins" +version.workspace = true +edition.workspace = true +license.workspace = true +repository.workspace = true +keywords = ["onnx", "memory", "allocator", "plugin", "ffi"] +categories = ["memory-management"] + +[dependencies] +libloading = { workspace = true } +onnx-runtime-memory-abi = { workspace = true } +onnx-runtime-memory-api = { workspace = true } +thiserror = { workspace = true } + +[dev-dependencies] +onnx-runtime-memory-testplugin = { workspace = true } + +[lints] +workspace = true diff --git a/crates/onnx-runtime-memory-host/src/allocator.rs b/crates/onnx-runtime-memory-host/src/allocator.rs new file mode 100644 index 0000000000..ff5f5709be --- /dev/null +++ b/crates/onnx-runtime-memory-host/src/allocator.rs @@ -0,0 +1,1566 @@ +//! Wrapping an nxmem allocator vtable in the internal Rust interfaces. +//! +//! # Lock discipline +//! +//! This adapter keeps one small `Mutex` mapping a live pointer to the +//! `allocation_id` the plugin knows it by. **No ABI call is ever made while +//! that lock is held.** Every path takes the lock, mutates the map, drops the +//! guard, and only then enters the plugin. That mirrors the rule the rest of +//! the memory stack already obeys for trait objects — no trait-object call +//! under a governance lock — and matters more here, because a plugin may +//! block, call back into the host, or spawn threads. +//! +//! # Identity +//! +//! `allocation_id` comes from a monotonic counter and is **never** +//! pointer-derived, so a freed-and-reused address can never be mistaken for +//! the allocation that previously lived there. A pointer the host does not +//! recognise is refused outright rather than retried by address. + +use std::collections::HashMap; +use std::ptr::NonNull; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::{Arc, Mutex}; + +use onnx_runtime_memory_abi::{ + NXMEM_CAP_DEFERRED_RELEASE, NXMEM_CAP_SHARED_MAPPING, NXMEM_CAP_STRUCTURED_RELEASE, + NXMEM_CAP_VIRTUAL_BACKING, NXMEM_RELEASE_COMPLETE, NXMEM_RELEASE_FAILED, + NXMEM_RELEASE_QUARANTINED, NxmemAllocRequest, NxmemAllocResult, NxmemAllocation, + NxmemAllocatorVtable, NxmemByteRange, NxmemHostCallbacks, NxmemOpenRequest, NxmemRangeRequest, + NxmemReclaimRequest, NxmemReleaseCompletion, NxmemReleaseOutcome, NxmemSharedMappingVtable, + NxmemSharedPrefixCommitInfo, NxmemSharedPrefixCommitRequest, NxmemSharedPrefixHandle, + NxmemStatus, NxmemStatusCode, NxmemVirtualBackingVtable, catch_status_panic, +}; +use onnx_runtime_memory_api::{ + AllocationCommitRange, AllocationReleaseOutcome, DeviceAllocator, DeviceKey, MemoryError, + QuarantineReason, ReleaseAccounting, ResidualOwnership, SharedDevicePrefix, + SharedMapping, SharedPrefixCommitInfo, Tier, VirtualBacking, +}; + +use crate::error::{PluginError, status_to_memory_error}; +use crate::loader::{PluginFactory, PluginModule, device_id, device_key}; + +/// The host's reclaim hook, offered to a plugin under memory pressure. +/// +/// # Threading and reentrancy +/// +/// * Called from inside a plugin call the host is currently making, and from +/// plugin-owned worker threads. +/// * **Must not be invoked, and must not itself acquire, any governance lock.** +/// The host releases every accounting, charge, and registration lock before +/// entering a plugin, so a reclaim hook that re-acquires one is safe only if +/// it never blocks on a thread that is itself inside a plugin call. +/// * Returning an error is a normal outcome. The plugin must handle it and +/// must not abort. +pub trait HostReclaim: Send + Sync + std::fmt::Debug { + /// Free up to `bytes` of host-cached memory on `device`. + /// + /// Returns the bytes actually reclaimed. Zero is a valid answer. + fn request_reclaim(&self, device: DeviceKey, bytes: u64) -> Result; +} + +/// One retired deferred release, as reported by the plugin. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct RetiredRelease { + /// The ticket the plugin issued when the release was queued. + pub ticket: u64, + /// The allocation the ticket refers to. + pub allocation_id: u64, + /// The allocation size the release was prepared for. + pub allocation_bytes: u64, + /// Bytes whose mapping reference transitioned to unmapped. + pub unmapped_bytes: u64, + /// Whether the release completed, quarantined, or failed. + pub state: u32, +} + +/// Host state reachable from a plugin callback. +/// +/// Lives behind a stable heap address for the whole life of the allocator: the +/// plugin stores a raw pointer to it in its own state, and the contract +/// promises the table outlives the allocator's final `release` and every +/// queued release naming it. +#[derive(Debug)] +struct HostBridge { + device: DeviceKey, + reclaim: Option>, + reclaim_calls: AtomicU64, + reclaim_failures: AtomicU64, + retired: Mutex>, + module: Arc, +} + +/// `request_reclaim` trampoline. +/// +/// # Safety +/// +/// `host_ctx` must be the pointer the host installed in the callback table and +/// `request` / `reclaimed_out` must be valid for the call. +unsafe extern "C" fn host_request_reclaim( + host_ctx: *mut core::ffi::c_void, + request: *const NxmemReclaimRequest, + reclaimed_out: *mut u64, +) -> NxmemStatus { + catch_status_panic(|| { + if host_ctx.is_null() || request.is_null() || reclaimed_out.is_null() { + return NxmemStatus::with_message( + NxmemStatusCode::InvalidArgument, + "nxmem: request_reclaim was called with a null pointer", + ); + } + // SAFETY: checked non-null; the host installed this pointer and keeps + // the bridge alive past the allocator's final release. + let bridge = unsafe { &*(host_ctx as *const HostBridge) }; + // SAFETY: checked non-null; the plugin passes a valid request. + let request = unsafe { &*request }; + bridge.reclaim_calls.fetch_add(1, Ordering::AcqRel); + + let Some(reclaim) = bridge.reclaim.as_ref() else { + bridge.reclaim_failures.fetch_add(1, Ordering::AcqRel); + return NxmemStatus::with_message( + NxmemStatusCode::UnsupportedCapability, + "nxmem: this host offers no reclaim path", + ); + }; + let Some(device) = device_key(request.device) else { + bridge.reclaim_failures.fetch_add(1, Ordering::AcqRel); + return NxmemStatus::with_message( + NxmemStatusCode::InvalidArgument, + "nxmem: the reclaim request named a tier this host does not know", + ); + }; + if device != bridge.device { + bridge.reclaim_failures.fetch_add(1, Ordering::AcqRel); + return NxmemStatus::with_message( + NxmemStatusCode::WrongDevice, + "nxmem: the reclaim request named a device this allocator does not serve", + ); + } + + match reclaim.request_reclaim(device, request.bytes) { + Ok(reclaimed) => { + // SAFETY: checked non-null above. + unsafe { *reclaimed_out = reclaimed }; + NxmemStatus::ok() + } + Err(reason) => { + bridge.reclaim_failures.fetch_add(1, Ordering::AcqRel); + // SAFETY: checked non-null above. A failing callback must still + // leave the out-parameter defined. + unsafe { *reclaimed_out = 0 }; + NxmemStatus::with_message(NxmemStatusCode::CallbackFailed, &reason) + } + } + }) +} + +/// `release_completed` trampoline. +/// +/// # Safety +/// +/// `host_ctx` must be the pointer the host installed and `completion` must be +/// valid for the call. +unsafe extern "C" fn host_release_completed( + host_ctx: *mut core::ffi::c_void, + completion: *const NxmemReleaseCompletion, +) -> NxmemStatus { + catch_status_panic(|| { + if host_ctx.is_null() || completion.is_null() { + return NxmemStatus::with_message( + NxmemStatusCode::InvalidArgument, + "nxmem: release_completed was called with a null pointer", + ); + } + // SAFETY: checked non-null; the host keeps the bridge alive past every + // queued release that can name it. + let bridge = unsafe { &*(host_ctx as *const HostBridge) }; + // SAFETY: checked non-null. + let completion = unsafe { &*completion }; + + let record = RetiredRelease { + ticket: completion.ticket, + allocation_id: completion.allocation_id, + allocation_bytes: completion.outcome.allocation_bytes, + unmapped_bytes: completion.outcome.unmapped_bytes, + state: completion.outcome.state, + }; + match bridge.retired.lock() { + Ok(mut retired) => retired.push(record), + Err(_) => { + return NxmemStatus::with_message( + NxmemStatusCode::InternalError, + "nxmem: the host completion log is poisoned", + ); + } + } + bridge.module.release_retired(); + bridge.module.allocation_closed(); + NxmemStatus::ok() + }) +} + +/// A live allocation as the host knows it. +#[derive(Debug, Clone, Copy)] +struct LiveAllocation { + allocation_id: u64, + bytes: usize, + align: usize, +} + +/// The shared core of a plugin-backed mechanism. +/// +/// Every capability view holds an `Arc` of this, so the plugin allocator is +/// released only once the allocator **and** all of its capability views are +/// gone. No back-reference is stored anywhere, so this introduces no reference +/// cycle. +#[derive(Debug)] +pub struct AllocatorCore { + vtable: NxmemAllocatorVtable, + backing: Option, + shared: Option, + device: DeviceKey, + mechanism_id: u64, + name: String, + negotiated_minor: u32, + capability_flags: u64, + next_allocation_id: AtomicU64, + live: Mutex>, + /// Boxed so the plugin's stored pointer stays valid across moves of the + /// core. Declared before `module` and dropped after `Drop::drop` has + /// already made the final `release` call, so the plugin can still touch + /// the table from inside that call. + bridge: Box, + callbacks: Box, + module: Arc, +} + +// SAFETY: the vtables are plain-data copies of function pointers plus an +// opaque `ctx`. The nxmem contract requires every slot to be callable from any +// host thread and requires the plugin to synchronise its own state. +unsafe impl Send for AllocatorCore {} +// SAFETY: as above. +unsafe impl Sync for AllocatorCore {} + +impl AllocatorCore { + /// The mechanism instance id, used to reject cross-provider misuse. + pub fn mechanism_id(&self) -> u64 { + self.mechanism_id + } + + /// The mechanism name published by the plugin. + pub fn name(&self) -> &str { + &self.name + } + + /// The capabilities this mechanism actually provides. + pub fn capability_flags(&self) -> u64 { + self.capability_flags + } + + /// The nxmem minor level this allocator's vtable was read at. + pub fn abi_minor(&self) -> u32 { + self.vtable.abi_minor.min(self.negotiated_minor) + } + + /// How many times the plugin has called back into the host for reclaim. + pub fn reclaim_calls(&self) -> u64 { + self.bridge.reclaim_calls.load(Ordering::Acquire) + } + + /// How many of those reclaim calls the host refused. + pub fn reclaim_failures(&self) -> u64 { + self.bridge.reclaim_failures.load(Ordering::Acquire) + } + + /// Deferred releases the plugin has reported as retired, in report order. + pub fn retired_releases(&self) -> Vec { + self.bridge + .retired + .lock() + .map(|retired| retired.clone()) + .unwrap_or_default() + } + + /// How many allocations the host still tracks for this mechanism. + pub fn live_allocation_count(&self) -> usize { + self.live.lock().map(|live| live.len()).unwrap_or(0) + } + + fn next_allocation_id(&self) -> u64 { + // Monotonic and explicitly not pointer-derived, so a reused address + // can never alias a previous allocation's identity. + self.next_allocation_id.fetch_add(1, Ordering::AcqRel) + 1 + } + + fn allocation_record(&self, ptr: NonNull) -> Option { + self.live.lock().ok()?.get(&(ptr.as_ptr() as usize)).copied() + } + + /// Remove a live record, **releasing the lock before any ABI call**. + fn take_allocation(&self, ptr: NonNull) -> Option { + let mut live = self.live.lock().ok()?; + let record = live.remove(&(ptr.as_ptr() as usize)); + drop(live); + record + } + + fn insert_allocation(&self, ptr: NonNull, record: LiveAllocation) { + if let Ok(mut live) = self.live.lock() { + live.insert(ptr.as_ptr() as usize, record); + } + } + + fn abi_allocation(&self, ptr: NonNull, record: LiveAllocation) -> NxmemAllocation { + NxmemAllocation::new( + self.mechanism_id, + record.allocation_id, + device_id(self.device), + ptr.as_ptr(), + record.bytes as u64, + record.align as u64, + ) + } + + fn alloc_request(&self, allocation_id: u64, bytes: usize, align: usize) -> NxmemAllocRequest { + NxmemAllocRequest::new( + self.mechanism_id, + allocation_id, + device_id(self.device), + bytes as u64, + align as u64, + ) + } + + fn tier(&self) -> Tier { + self.device.tier + } + + /// Common tail of every allocation path. + fn finish_allocation( + &self, + request: &NxmemAllocRequest, + result: NxmemAllocResult, + bytes: usize, + align: usize, + ) -> Result, MemoryError> { + let ptr = NonNull::new(result.ptr).ok_or_else(|| MemoryError::AllocationFailed { + tier: self.tier().name(), + requested: bytes as u64, + reason: String::from( + "the memory plugin reported success but returned a null pointer; a successful \ + allocation must return a live, unique address", + ), + })?; + if ptr.as_ptr() as usize % align.max(1) != 0 { + // The plugin broke the alignment contract. Hand the bytes straight + // back rather than letting a misaligned pointer escape. + let record = LiveAllocation { + allocation_id: request.allocation_id, + bytes, + align, + }; + let allocation = self.abi_allocation(ptr, record); + self.raw_release(&allocation); + return Err(MemoryError::AllocationFailed { + tier: self.tier().name(), + requested: bytes as u64, + reason: format!( + "the memory plugin returned an address that is not aligned to {align}; the \ + bytes were handed straight back" + ), + }); + } + self.insert_allocation( + ptr, + LiveAllocation { + allocation_id: request.allocation_id, + bytes, + align, + }, + ); + self.module.allocation_opened(); + Ok(ptr) + } + + /// Terminal release of an allocation the host is no longer tracking. + /// + /// Used on the error paths above, where the record was never installed. + fn raw_release(&self, allocation: &NxmemAllocation) { + if let Some(release) = self.structured_release_slot() { + let mut outcome = NxmemReleaseOutcome::zeroed(); + // SAFETY: `ctx` came from this vtable; `allocation` and `outcome` + // are valid locals. No host lock is held. + let _ = unsafe { release(self.vtable.ctx, allocation, &raw mut outcome) }; + return; + } + if let Some(deallocate) = self.vtable.deallocate { + let mut unmapped = 0u64; + // SAFETY: as above. + let _ = unsafe { deallocate(self.vtable.ctx, allocation, &raw mut unmapped) }; + } + } + + /// The structured-release slot, if it was both negotiated and provided. + /// + /// Negotiation fixes a ceiling; the vtable declares its own level. Both + /// must permit the slot before it is called, which is what makes an older + /// mechanism inside a newer module safe. + fn structured_release_slot( + &self, + ) -> Option< + unsafe extern "C" fn( + *mut core::ffi::c_void, + *const NxmemAllocation, + *mut NxmemReleaseOutcome, + ) -> NxmemStatus, + > { + if self.capability_flags & NXMEM_CAP_STRUCTURED_RELEASE == 0 { + return None; + } + if self.negotiated_minor < 1 || self.vtable.abi_minor < 1 { + return None; + } + self.vtable.release_allocation + } + + /// The canonical release path. Never called with any lock held. + fn release_allocation( + &self, + ptr: NonNull, + bytes: usize, + align: usize, + ) -> AllocationReleaseOutcome { + // Take the record and drop the guard before entering the plugin. + let Some(record) = self.take_allocation(ptr) else { + // An address the host does not recognise is refused outright. It + // is never retried by pointer, because a freed-and-reused address + // would then free somebody else's live allocation. + return AllocationReleaseOutcome::failed(format!( + "the memory plugin host does not recognise address {:p} as a live allocation of \ + mechanism {}; it was already released, or it belongs to another mechanism", + ptr.as_ptr(), + self.mechanism_id + )); + }; + + if record.bytes != bytes || record.align != align { + // Put the record back: nothing was mutated, so the allocation is + // exactly as live as the caller left it. + self.insert_allocation(ptr, record); + return AllocationReleaseOutcome::failed(format!( + "release of address {:p} named {bytes} bytes at align {align} but the allocation \ + is {} bytes at align {}; nothing was released", + ptr.as_ptr(), + record.bytes, + record.align + )); + } + + let allocation = self.abi_allocation(ptr, record); + + if let Some(release) = self.structured_release_slot() { + let mut outcome = NxmemReleaseOutcome::zeroed(); + // SAFETY: `ctx` came from this vtable; both pointers address valid + // locals that outlive the call. No host lock is held. + let status = unsafe { release(self.vtable.ctx, &raw const allocation, &raw mut outcome) }; + if !status.is_ok() { + self.insert_allocation(ptr, record); + return AllocationReleaseOutcome::failed(format!( + "the memory plugin refused release without mutating anything: {}", + status.describe() + )); + } + return self.interpret_outcome(ptr, record, outcome); + } + + let Some(deallocate) = self.vtable.deallocate else { + self.insert_allocation(ptr, record); + return AllocationReleaseOutcome::failed(String::from( + "the memory plugin exposes no release slot at all; the allocation cannot be \ + handed back", + )); + }; + let mut unmapped = 0u64; + // SAFETY: as above. + let status = + unsafe { deallocate(self.vtable.ctx, &raw const allocation, &raw mut unmapped) }; + if !status.is_ok() { + self.insert_allocation(ptr, record); + return AllocationReleaseOutcome::failed(format!( + "the memory plugin refused release without mutating anything: {}", + status.describe() + )); + } + self.module.allocation_closed(); + AllocationReleaseOutcome::complete(ReleaseAccounting { + allocation_bytes: record.bytes as u64, + unmapped_bytes: unmapped, + }) + } + + fn interpret_outcome( + &self, + ptr: NonNull, + record: LiveAllocation, + outcome: NxmemReleaseOutcome, + ) -> AllocationReleaseOutcome { + let accounting = ReleaseAccounting { + allocation_bytes: outcome.allocation_bytes.max(record.bytes as u64), + unmapped_bytes: outcome.unmapped_bytes, + }; + match outcome.state { + NXMEM_RELEASE_COMPLETE => { + self.module.allocation_closed(); + AllocationReleaseOutcome::complete(accounting) + } + NXMEM_RELEASE_QUARANTINED => { + // The allocation is gone from the host's live map and the + // plugin kept residual ownership. It must never be reused. + self.module.allocation_closed(); + AllocationReleaseOutcome::quarantined( + accounting, + ResidualOwnership { + state: onnx_runtime_memory_api::AllocationReleaseState::PartiallyUnmapped, + reason: QuarantineReason::PartialRelease, + retained_bytes: outcome.residual_owned_bytes, + address: ptr.as_ptr() as usize, + align: record.align, + }, + ) + } + NXMEM_RELEASE_FAILED => { + // The plugin promises nothing was mutated, so the allocation is + // as live as the caller left it. Restore the record. + self.insert_allocation(ptr, record); + AllocationReleaseOutcome::failed(format!( + "the memory plugin refused release without mutating anything: {}", + outcome.failure.describe() + )) + } + other => { + // An unknown state is never guessed. Treat it as residual + // ownership so nothing is reused or double-refunded. + self.module.allocation_closed(); + AllocationReleaseOutcome::quarantined( + ReleaseAccounting { + allocation_bytes: record.bytes as u64, + unmapped_bytes: 0, + }, + ResidualOwnership { + state: onnx_runtime_memory_api::AllocationReleaseState::PartiallyUnmapped, + reason: QuarantineReason::AllocatorRefused, + retained_bytes: record.bytes as u64, + address: ptr.as_ptr() as usize, + align: record.align, + }, + ) + .tap_unknown(other) + } + } + } +} + +/// Small helper so the unknown-state branch above stays readable. +trait TapUnknown: Sized { + fn tap_unknown(self, _state: u32) -> Self { + self + } +} + +impl TapUnknown for AllocationReleaseOutcome {} + +impl Drop for AllocatorCore { + fn drop(&mut self) { + // Any allocation still in the map is a host-side leak, not something to + // silently free: the plugin may already have been told about it. Report + // through the plugin's terminal slot so the bytes are not stranded. + let leaked: Vec<(usize, LiveAllocation)> = match self.live.lock() { + Ok(mut live) => live.drain().collect(), + Err(_) => Vec::new(), + }; + for (address, record) in leaked { + let Some(ptr) = NonNull::new(address as *mut u8) else { + continue; + }; + let allocation = self.abi_allocation(ptr, record); + self.raw_release(&allocation); + self.module.allocation_closed(); + } + + if let Some(release) = self.vtable.release { + // SAFETY: `ctx` came from this vtable and `release` is called + // exactly once, here. The callback table and bridge are still + // alive: they are fields of `self`, which is only destructured + // after this `drop` body returns. + unsafe { release(self.vtable.ctx) }; + } + self.module.allocator_closed(); + } +} + +/// A plugin-backed [`DeviceAllocator`]. +/// +/// Field order is load-bearing. The capability views are declared **before** +/// `core` so they drop first; each holds its own `Arc`, so the +/// plugin's `release` runs only after every view is gone. +#[derive(Debug)] +pub struct PluginAllocator { + backing: Option, + shared: Option, + core: Arc, +} + +impl PluginAllocator { + /// The shared core, for callers that need to inspect plugin-side state. + pub fn core(&self) -> &Arc { + &self.core + } + + /// Queue a stream-ordered release instead of freeing immediately. + /// + /// Returns the plugin's ticket. Until the matching completion arrives + /// through `release_completed`, the module, factory, allocator, and + /// callback table all stay pinned and unload stays refused. + /// + /// # Safety + /// + /// `ptr` must identify one live allocation from this mechanism with + /// exactly this `bytes` and `align`, and must not be released twice. + pub unsafe fn enqueue_release( + &self, + ptr: NonNull, + bytes: usize, + align: usize, + ) -> Result { + let core = &self.core; + if core.capability_flags & NXMEM_CAP_DEFERRED_RELEASE == 0 { + return Err(MemoryError::AllocationFailed { + tier: core.tier().name(), + requested: bytes as u64, + reason: String::from( + "this mechanism does not advertise deferred release; use the canonical \ + release path instead", + ), + }); + } + let Some(enqueue) = core.vtable.enqueue_release else { + return Err(MemoryError::AllocationFailed { + tier: core.tier().name(), + requested: bytes as u64, + reason: String::from( + "this mechanism advertises deferred release but provides no enqueue slot", + ), + }); + }; + // Lock taken and released before the ABI call. + let Some(record) = core.take_allocation(ptr) else { + return Err(MemoryError::AllocationFailed { + tier: core.tier().name(), + requested: bytes as u64, + reason: format!( + "address {:p} is not a live allocation of this mechanism", + ptr.as_ptr() + ), + }); + }; + if record.bytes != bytes || record.align != align { + core.insert_allocation(ptr, record); + return Err(MemoryError::AllocationFailed { + tier: core.tier().name(), + requested: bytes as u64, + reason: format!( + "deferred release named {bytes} bytes at align {align} but the allocation is \ + {} bytes at align {}", + record.bytes, record.align + ), + }); + } + + let allocation = core.abi_allocation(ptr, record); + let mut ticket = 0u64; + // Count the queued release *before* the call so a completion that + // arrives synchronously from inside `enqueue_release` cannot decrement + // a counter that was never incremented. + core.module.release_queued(); + // SAFETY: `ctx` came from this vtable; both pointers address valid + // locals. No host lock is held. + let status = + unsafe { enqueue(core.vtable.ctx, &raw const allocation, &raw mut ticket) }; + if !status.is_ok() { + core.module.release_retired(); + core.insert_allocation(ptr, record); + return Err(status_to_memory_error( + "enqueue_release", + core.tier(), + bytes as u64, + &status, + )); + } + Ok(ticket) + } + + /// Retire up to `max` queued releases, driving the completion callbacks. + pub fn drain_releases(&self, max: u64) -> Result { + let core = &self.core; + let Some(drain) = core.vtable.drain_releases else { + return Ok(0); + }; + let mut retired = 0u64; + // SAFETY: `ctx` came from this vtable and `retired` is a valid local. + // No host lock is held; the plugin will call back into the host from + // inside this call, which is exactly why no lock may be held. + let status = unsafe { drain(core.vtable.ctx, max, &raw mut retired) }; + if !status.is_ok() { + return Err(status_to_memory_error( + "drain_releases", + core.tier(), + 0, + &status, + )); + } + Ok(retired) + } + + /// How many releases the plugin still has queued. + pub fn pending_release_count(&self) -> Result { + let core = &self.core; + let Some(pending) = core.vtable.pending_release_count else { + return Ok(0); + }; + let mut count = 0u64; + // SAFETY: `ctx` came from this vtable and `count` is a valid local. + let status = unsafe { pending(core.vtable.ctx, &raw mut count) }; + if !status.is_ok() { + return Err(status_to_memory_error( + "pending_release_count", + core.tier(), + 0, + &status, + )); + } + Ok(count) + } +} + +impl DeviceAllocator for PluginAllocator { + fn allocate(&self, bytes: usize, align: usize) -> Result, MemoryError> { + let core = &self.core; + let Some(allocate) = core.vtable.allocate else { + return Err(MemoryError::AllocationFailed { + tier: core.tier().name(), + requested: bytes as u64, + reason: String::from("this mechanism provides no allocate slot"), + }); + }; + let request = core.alloc_request(core.next_allocation_id(), bytes, align); + let mut result = NxmemAllocResult::zeroed(); + // SAFETY: `ctx` came from this vtable; both pointers address valid + // locals that outlive the call. No host lock is held, so the plugin may + // call back into the host from inside this call. + let status = + unsafe { allocate(core.vtable.ctx, &raw const request, &raw mut result) }; + if !status.is_ok() { + return Err(status_to_memory_error( + "allocate", + core.tier(), + bytes as u64, + &status, + )); + } + core.finish_allocation(&request, result, bytes, align) + } + + unsafe fn deallocate(&self, ptr: NonNull, bytes: usize, align: usize) { + // Route through the canonical structured path so there is exactly one + // production release entry point to reason about and to test. + let _ = self.core.release_allocation(ptr, bytes, align); + } + + unsafe fn deallocate_with_unmapped(&self, ptr: NonNull, bytes: usize, align: usize) -> u64 { + match self.core.release_allocation(ptr, bytes, align) { + AllocationReleaseOutcome::Complete { accounting } + | AllocationReleaseOutcome::Quarantined { accounting, .. } => accounting.unmapped_bytes, + AllocationReleaseOutcome::Failed { .. } => 0, + } + } + + unsafe fn release( + &self, + ptr: NonNull, + bytes: usize, + align: usize, + ) -> AllocationReleaseOutcome { + self.core.release_allocation(ptr, bytes, align) + } + + fn device(&self) -> DeviceKey { + self.core.device + } + + fn commits_on_demand(&self) -> bool { + // Only a mechanism that both reserves lazily *and* reports its commits + // to the host may claim this. Virtual backing alone is not enough. + self.core.capability_flags & NXMEM_CAP_VIRTUAL_BACKING != 0 && self.backing.is_some() + } + + fn as_virtual_backing(&self) -> Option<&dyn VirtualBacking> { + self.backing.as_ref().map(|backing| backing as &dyn VirtualBacking) + } + + fn as_shared_mapping(&self) -> Option<&dyn SharedMapping> { + self.shared.as_ref().map(|shared| shared as &dyn SharedMapping) + } +} + +/// The lazy-backing capability of a plugin mechanism. +#[derive(Debug)] +pub struct PluginVirtualBacking { + core: Arc, +} + +impl PluginVirtualBacking { + fn vtable(&self) -> &NxmemVirtualBackingVtable { + self.core + .backing + .as_ref() + .expect("a virtual-backing view is only built when the vtable is present") + } + + fn range_request( + &self, + ptr: NonNull, + allocation_bytes: usize, + align: usize, + offset: usize, + bytes: usize, + ) -> Result { + let record = self.core.allocation_record(ptr).ok_or_else(|| { + MemoryError::AllocationFailed { + tier: self.core.tier().name(), + requested: bytes as u64, + reason: format!( + "address {:p} is not a live allocation of this mechanism, so no range of it \ + can be committed", + ptr.as_ptr() + ), + } + })?; + if record.bytes != allocation_bytes || record.align != align { + return Err(MemoryError::AllocationFailed { + tier: self.core.tier().name(), + requested: bytes as u64, + reason: format!( + "the range names a {allocation_bytes}-byte allocation at align {align} but \ + that address is {} bytes at align {}", + record.bytes, record.align + ), + }); + } + Ok(NxmemRangeRequest::new( + self.core.abi_allocation(ptr, record), + offset as u64, + bytes as u64, + )) + } +} + +impl VirtualBacking for PluginVirtualBacking { + fn allocate_committed( + &self, + bytes: usize, + align: usize, + committed_ranges: &[std::ops::Range], + ) -> Result, MemoryError> { + let core = &self.core; + let Some(allocate) = self.vtable().allocate_committed else { + return Err(MemoryError::AllocationFailed { + tier: core.tier().name(), + requested: bytes as u64, + reason: String::from("this mechanism provides no allocate_committed slot"), + }); + }; + let ranges: Vec = committed_ranges + .iter() + .map(|range| NxmemByteRange { + offset: range.start as u64, + bytes: range.len() as u64, + }) + .collect(); + let mut request = core.alloc_request(core.next_allocation_id(), bytes, align); + request.committed_ranges = ranges.as_ptr(); + request.committed_range_count = ranges.len() as u64; + + let mut result = NxmemAllocResult::zeroed(); + // SAFETY: `ctx` came from the backing vtable; `request` borrows `ranges` + // which outlives the call, and `result` is a valid local. No host lock + // is held. + let status = unsafe { + allocate(self.vtable().ctx, &raw const request, &raw mut result) + }; + drop(ranges); + if !status.is_ok() { + return Err(status_to_memory_error( + "allocate_committed", + core.tier(), + bytes as u64, + &status, + )); + } + core.finish_allocation(&request, result, bytes, align) + } + + fn commit_allocation_range( + &self, + ptr: NonNull, + allocation_bytes: usize, + align: usize, + offset: usize, + bytes: usize, + ) -> Result<(), MemoryError> { + let Some(commit) = self.vtable().commit_range else { + return Err(MemoryError::AllocationFailed { + tier: self.core.tier().name(), + requested: bytes as u64, + reason: String::from("this mechanism provides no commit_range slot"), + }); + }; + let request = self.range_request(ptr, allocation_bytes, align, offset, bytes)?; + // SAFETY: `ctx` came from the backing vtable and `request` is a valid + // local that outlives the call. No host lock is held. + let status = unsafe { commit(self.vtable().ctx, &raw const request) }; + if !status.is_ok() { + return Err(status_to_memory_error( + "commit_range", + self.core.tier(), + bytes as u64, + &status, + )); + } + Ok(()) + } + + fn mapped_bytes_for_allocation_ranges( + &self, + ranges: &[AllocationCommitRange], + ) -> Result { + let Some(mapped) = self.vtable().mapped_bytes_for_ranges else { + return Err(MemoryError::AllocationFailed { + tier: self.core.tier().name(), + requested: 0, + reason: String::from("this mechanism provides no mapped_bytes_for_ranges slot"), + }); + }; + let requests = ranges + .iter() + .map(|range| { + self.range_request( + range.ptr, + range.allocation_bytes, + range.align, + range.offset, + range.bytes, + ) + }) + .collect::, _>>()?; + let mut out = 0u64; + // SAFETY: `requests` outlives the call and its length is passed + // alongside it; `out` is a valid local. No host lock is held. + let status = unsafe { + mapped( + self.vtable().ctx, + requests.as_ptr(), + requests.len() as u64, + &raw mut out, + ) + }; + drop(requests); + if !status.is_ok() { + return Err(status_to_memory_error( + "mapped_bytes_for_ranges", + self.core.tier(), + 0, + &status, + )); + } + Ok(out) + } + + fn mapped_bytes_for_allocation(&self, bytes: usize, align: usize) -> Result { + let Some(mapped) = self.vtable().mapped_bytes_for_allocation else { + return Err(MemoryError::AllocationFailed { + tier: self.core.tier().name(), + requested: bytes as u64, + reason: String::from( + "this mechanism provides no mapped_bytes_for_allocation slot", + ), + }); + }; + // A sizing query carries no allocation identity yet, so id zero is + // used deliberately: it can never match a live allocation. + let request = self.core.alloc_request(0, bytes, align); + let mut out = 0u64; + // SAFETY: both pointers address valid locals; no host lock is held. + let status = + unsafe { mapped(self.vtable().ctx, &raw const request, &raw mut out) }; + if !status.is_ok() { + return Err(status_to_memory_error( + "mapped_bytes_for_allocation", + self.core.tier(), + bytes as u64, + &status, + )); + } + Ok(out) + } + + fn decommit_allocation_range( + &self, + ptr: NonNull, + allocation_bytes: usize, + align: usize, + offset: usize, + bytes: usize, + ) -> Result { + let Some(decommit) = self.vtable().decommit_range else { + return Err(MemoryError::AllocationFailed { + tier: self.core.tier().name(), + requested: bytes as u64, + reason: String::from("this mechanism provides no decommit_range slot"), + }); + }; + let request = self.range_request(ptr, allocation_bytes, align, offset, bytes)?; + let mut unmapped = 0u64; + // SAFETY: both pointers address valid locals; no host lock is held. + let status = + unsafe { decommit(self.vtable().ctx, &raw const request, &raw mut unmapped) }; + if !status.is_ok() { + return Err(status_to_memory_error( + "decommit_range", + self.core.tier(), + bytes as u64, + &status, + )); + } + Ok(unmapped) + } + + fn allocation_committed_bytes( + &self, + ptr: NonNull, + allocation_bytes: usize, + align: usize, + ) -> usize { + let Some(committed) = self.vtable().committed_bytes else { + return 0; + }; + let Some(record) = self.core.allocation_record(ptr) else { + return 0; + }; + if record.bytes != allocation_bytes || record.align != align { + return 0; + } + let allocation = self.core.abi_allocation(ptr, record); + let mut out = 0u64; + // SAFETY: both pointers address valid locals; no host lock is held. + let status = unsafe { + committed(self.vtable().ctx, &raw const allocation, &raw mut out) + }; + if !status.is_ok() { + return 0; + } + out as usize + } +} + +/// The shared-physical-mapping capability of a plugin mechanism. +#[derive(Debug)] +pub struct PluginSharedMapping { + core: Arc, +} + +impl PluginSharedMapping { + fn vtable(&self) -> &NxmemSharedMappingVtable { + self.core + .shared + .as_ref() + .expect("a shared-mapping view is only built when the vtable is present") + } + + /// Recover a prefix this capability actually created. + /// + /// A foreign, wrong-device, or wrong-mechanism prefix is refused here + /// rather than reported as costing zero. + fn own_prefix<'a>( + &self, + prefix: &'a dyn SharedDevicePrefix, + ) -> Result<&'a PluginSharedPrefix, MemoryError> { + let plugin_prefix = prefix + .as_any() + .downcast_ref::() + .ok_or_else(|| MemoryError::AllocationFailed { + tier: self.core.tier().name(), + requested: 0, + reason: String::from( + "the shared prefix was not created by a memory plugin, so this mechanism \ + cannot map it", + ), + })?; + if plugin_prefix.handle.mechanism_id != self.core.mechanism_id { + return Err(MemoryError::AllocationFailed { + tier: self.core.tier().name(), + requested: 0, + reason: format!( + "the shared prefix belongs to mechanism {} but this mechanism is {}; a \ + cross-mechanism mapping is refused rather than costed as free", + plugin_prefix.handle.mechanism_id, self.core.mechanism_id + ), + }); + } + if device_key(plugin_prefix.handle.device) != Some(self.core.device) { + return Err(MemoryError::AllocationFailed { + tier: self.core.tier().name(), + requested: 0, + reason: String::from( + "the shared prefix lives on a different device than this mechanism serves", + ), + }); + } + Ok(plugin_prefix) + } +} + +impl SharedMapping for PluginSharedMapping { + fn create_shared_prefix( + &self, + bytes: usize, + ) -> Result, MemoryError> { + let Some(create) = self.vtable().create_shared_prefix else { + return Err(MemoryError::AllocationFailed { + tier: self.core.tier().name(), + requested: bytes as u64, + reason: String::from("this mechanism provides no create_shared_prefix slot"), + }); + }; + let mut handle = NxmemSharedPrefixHandle::zeroed(); + // SAFETY: `handle` is a valid writable local; no host lock is held. + let status = unsafe { + create( + self.vtable().ctx, + self.core.mechanism_id, + bytes as u64, + &raw mut handle, + ) + }; + if !status.is_ok() { + return Err(status_to_memory_error( + "create_shared_prefix", + self.core.tier(), + bytes as u64, + &status, + )); + } + if handle.handle == 0 || handle.mechanism_id != self.core.mechanism_id { + return Err(MemoryError::AllocationFailed { + tier: self.core.tier().name(), + requested: bytes as u64, + reason: String::from( + "the memory plugin reported success but returned a prefix with no identity", + ), + }); + } + Ok(Box::new(PluginSharedPrefix { + handle, + core: Arc::clone(&self.core), + })) + } + + fn incremental_owned_bytes_for_shared_prefix( + &self, + prefix: &dyn SharedDevicePrefix, + ) -> Result { + let prefix = self.own_prefix(prefix)?; + let Some(incremental) = self.vtable().incremental_owned_bytes else { + return Err(MemoryError::AllocationFailed { + tier: self.core.tier().name(), + requested: 0, + reason: String::from("this mechanism provides no incremental_owned_bytes slot"), + }); + }; + let mut out = 0u64; + // SAFETY: both pointers address valid locals; no host lock is held. + let status = unsafe { + incremental(self.vtable().ctx, &raw const prefix.handle, &raw mut out) + }; + if !status.is_ok() { + return Err(status_to_memory_error( + "incremental_owned_bytes", + self.core.tier(), + 0, + &status, + )); + } + Ok(out) + } + + fn commit_shared_prefix( + &self, + prefix: &dyn SharedDevicePrefix, + ptr: NonNull, + allocation_bytes: usize, + byte_offset: usize, + ) -> Result { + let prefix = self.own_prefix(prefix)?; + let Some(commit) = self.vtable().commit_shared_prefix else { + return Err(MemoryError::AllocationFailed { + tier: self.core.tier().name(), + requested: 0, + reason: String::from("this mechanism provides no commit_shared_prefix slot"), + }); + }; + let record = self.core.allocation_record(ptr).ok_or_else(|| { + MemoryError::AllocationFailed { + tier: self.core.tier().name(), + requested: 0, + reason: format!( + "address {:p} is not a live allocation of this mechanism", + ptr.as_ptr() + ), + } + })?; + if record.bytes != allocation_bytes { + return Err(MemoryError::AllocationFailed { + tier: self.core.tier().name(), + requested: 0, + reason: format!( + "the mapping names a {allocation_bytes}-byte allocation but that address is \ + {} bytes", + record.bytes + ), + }); + } + + let request = NxmemSharedPrefixCommitRequest::new( + prefix.handle, + self.core.abi_allocation(ptr, record), + byte_offset as u64, + ); + let mut info = NxmemSharedPrefixCommitInfo::zeroed(); + // SAFETY: both pointers address valid locals; no host lock is held. + let status = unsafe { + commit(self.vtable().ctx, &raw const request, &raw mut info) + }; + if !status.is_ok() { + return Err(status_to_memory_error( + "commit_shared_prefix", + self.core.tier(), + 0, + &status, + )); + } + Ok(SharedPrefixCommitInfo { + additional_owned_bytes: info.additional_owned_bytes, + newly_mapped_bytes: info.newly_mapped_bytes, + granules: info.granules as usize, + }) + } +} + +/// A plugin-owned shared physical prefix. +/// +/// Holds an `Arc` so the mechanism cannot be released while a +/// prefix it owns is still alive. +#[derive(Debug)] +pub struct PluginSharedPrefix { + handle: NxmemSharedPrefixHandle, + core: Arc, +} + +impl PluginSharedPrefix { + /// The plugin-side handle, for tests that need the raw identity. + pub fn handle(&self) -> NxmemSharedPrefixHandle { + self.handle + } +} + +impl SharedDevicePrefix for PluginSharedPrefix { + fn device_ptr(&self) -> u64 { + self.handle.device_ptr + } + + fn committed_physical_bytes(&self) -> u64 { + self.handle.committed_physical_bytes + } + + fn mapped_bytes(&self) -> usize { + self.handle.mapped_bytes as usize + } + + fn requested_bytes(&self) -> usize { + self.handle.requested_bytes as usize + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } +} + +impl Drop for PluginSharedPrefix { + fn drop(&mut self) { + let Some(shared) = self.core.shared.as_ref() else { + return; + }; + let Some(release) = shared.release_shared_prefix else { + return; + }; + // SAFETY: `ctx` came from this vtable; the handle is a valid local and + // `release_shared_prefix` is called exactly once, here. No host lock is + // held and this slot may not call back into the host. + let _ = unsafe { release(shared.ctx, &raw const self.handle) }; + } +} + +/// Open an allocator from a factory, wiring up the host callback table. +pub(crate) fn open_allocator( + factory: &PluginFactory, + required_capability_flags: u64, + reclaim: Option>, +) -> Result { + let module = Arc::clone(factory.module()); + let minor = module.negotiated().minor; + let factory_vtable = factory.vtable(); + let Some(open) = factory_vtable.open_allocator else { + return Err(PluginError::Contract { + path: module.path().display().to_string(), + reason: format!( + "factory `{}` provides no open_allocator slot", + factory.name() + ), + }); + }; + + // The bridge and the callback table must live at stable heap addresses + // *before* the plugin sees them, and must outlive the allocator's final + // release. Boxing them here and moving the boxes into the core satisfies + // both: moving a `Box` never moves the heap allocation. + let bridge = Box::new(HostBridge { + device: factory.device(), + reclaim, + reclaim_calls: AtomicU64::new(0), + reclaim_failures: AtomicU64::new(0), + retired: Mutex::new(Vec::new()), + module: Arc::clone(&module), + }); + let callbacks = Box::new(NxmemHostCallbacks { + struct_size: core::mem::size_of::() as u32, + abi_minor: minor, + host_ctx: (&raw const *bridge) as *mut core::ffi::c_void, + request_reclaim: Some(host_request_reclaim), + release_completed: Some(host_release_completed), + }); + + let request = NxmemOpenRequest::new( + minor, + factory_vtable.device, + required_capability_flags, + &raw const *callbacks, + ); + let mut raw_allocator = core::ptr::null::(); + // SAFETY: `ctx` came from the factory vtable; `request` and + // `raw_allocator` are valid locals that outlive the call. No host lock is + // held. + let status = unsafe { + open( + factory_vtable.ctx, + &raw const request, + &raw mut raw_allocator, + ) + }; + if !status.is_ok() { + return Err(PluginError::call("open_allocator", status)); + } + + // SAFETY: the plugin wrote this pointer in response to the call above. + // `read_prefix` null-, alignment- and size-checks it before reading any + // field, and copies rather than borrowing. + let vtable = unsafe { NxmemAllocatorVtable::read_prefix(raw_allocator, minor) } + .map_err(|status| PluginError::call("allocator vtable", status))?; + vtable + .validate_required() + .map_err(|status| PluginError::call("allocator vtable", status))?; + + let device = device_key(vtable.device).ok_or_else(|| PluginError::Contract { + path: module.path().display().to_string(), + reason: format!( + "allocator `{}` declared tier code {} which this host does not know", + factory.name(), + vtable.device.tier + ), + })?; + if device != factory.device() { + return Err(PluginError::Contract { + path: module.path().display().to_string(), + reason: format!( + "factory `{}` serves {:?} but opened an allocator for {device:?}; a mechanism \ + must stay on one device", + factory.name(), + factory.device() + ), + }); + } + if vtable.mechanism_id == 0 { + return Err(PluginError::Contract { + path: module.path().display().to_string(), + reason: format!( + "allocator `{}` published mechanism id zero, which can never be matched against \ + an allocation", + factory.name() + ), + }); + } + let missing = required_capability_flags & !vtable.capability_flags; + if missing != 0 { + return Err(PluginError::Contract { + path: module.path().display().to_string(), + reason: format!( + "allocator `{}` was opened requiring capabilities {missing:#x} that it does not \ + provide", + factory.name() + ), + }); + } + + // Capability vtables are read at the allocator's own level, and each must + // name the same mechanism instance. + let backing = read_capability( + vtable.virtual_backing, + vtable.capability_flags & NXMEM_CAP_VIRTUAL_BACKING != 0, + minor, + vtable.mechanism_id, + "virtual backing", + &module, + |ptr, minor| { + // SAFETY: delegated to `read_capability`, which only calls this + // with a pointer the plugin published for this capability. + unsafe { NxmemVirtualBackingVtable::read_prefix(ptr, minor) } + }, + |vtable| vtable.mechanism_id, + )?; + let shared = read_capability( + vtable.shared_mapping, + vtable.capability_flags & NXMEM_CAP_SHARED_MAPPING != 0, + minor, + vtable.mechanism_id, + "shared mapping", + &module, + |ptr, minor| { + // SAFETY: as above. + unsafe { NxmemSharedMappingVtable::read_prefix(ptr, minor) } + }, + |vtable| vtable.mechanism_id, + )?; + + // SAFETY: the contract requires `name` to stay valid until the final + // `release`, which has not been called. + let name = unsafe { read_optional_c_string(vtable.name) }.unwrap_or_else(|| { + factory.name().to_string() + }); + + module.allocator_opened(); + let core = Arc::new(AllocatorCore { + vtable, + backing, + shared, + device, + mechanism_id: vtable.mechanism_id, + name, + negotiated_minor: minor, + capability_flags: vtable.capability_flags, + next_allocation_id: AtomicU64::new(0), + live: Mutex::new(HashMap::new()), + bridge, + callbacks, + module, + }); + + let backing_view = core + .backing + .is_some() + .then(|| PluginVirtualBacking { + core: Arc::clone(&core), + }); + let shared_view = core + .shared + .is_some() + .then(|| PluginSharedMapping { + core: Arc::clone(&core), + }); + + Ok(PluginAllocator { + backing: backing_view, + shared: shared_view, + core, + }) +} + +#[allow(clippy::too_many_arguments)] +fn read_capability( + ptr: *const T, + advertised: bool, + minor: u32, + mechanism_id: u64, + label: &str, + module: &Arc, + read: impl Fn(*const T, u32) -> Result, + id_of: impl Fn(&T) -> u64, +) -> Result, PluginError> { + if !advertised { + // An unsupported capability is represented explicitly: the flag is + // clear and the pointer is expected to be null. A non-null pointer + // behind a clear flag is a contract violation, not a bonus feature. + if !ptr.is_null() { + return Err(PluginError::Contract { + path: module.path().display().to_string(), + reason: format!( + "the allocator published a {label} vtable without advertising the capability; \ + an unsupported capability must be a null pointer and a clear flag" + ), + }); + } + return Ok(None); + } + if ptr.is_null() { + return Err(PluginError::Contract { + path: module.path().display().to_string(), + reason: format!( + "the allocator advertised {label} but published a null vtable for it" + ), + }); + } + let vtable = read(ptr, minor).map_err(|status| PluginError::call("capability vtable", status))?; + if id_of(&vtable) != mechanism_id { + return Err(PluginError::Contract { + path: module.path().display().to_string(), + reason: format!( + "the {label} vtable names mechanism {} but the allocator is mechanism \ + {mechanism_id}; a capability must belong to one coherent mechanism", + id_of(&vtable) + ), + }); + } + Ok(Some(vtable)) +} + +/// # Safety +/// +/// `ptr` must be null or point to a NUL-terminated byte string valid for the +/// duration of the call. +unsafe fn read_optional_c_string(ptr: *const u8) -> Option { + if ptr.is_null() { + return None; + } + // SAFETY: delegated to this function's contract. + unsafe { core::ffi::CStr::from_ptr(ptr.cast()) } + .to_str() + .ok() + .map(str::to_owned) +} diff --git a/crates/onnx-runtime-memory-host/src/error.rs b/crates/onnx-runtime-memory-host/src/error.rs new file mode 100644 index 0000000000..89ba425ca7 --- /dev/null +++ b/crates/onnx-runtime-memory-host/src/error.rs @@ -0,0 +1,143 @@ +//! Translating nxmem statuses into the runtime's error vocabulary. + +use onnx_runtime_memory_abi::{NxmemStatus, NxmemStatusCode}; +use onnx_runtime_memory_api::{MemoryError, Tier}; + +/// Something went wrong loading or talking to a memory plugin. +#[derive(Debug, thiserror::Error)] +pub enum PluginError { + /// The dynamic library could not be opened. + #[error("cannot open memory plugin {path}: {source}")] + Open { + path: String, + #[source] + source: libloading::Error, + }, + + /// A required entry point is missing. + /// + /// All three entry points are required. In particular a library without + /// `NxmemQueryUnloadReadiness` is refused rather than loaded, because the + /// host would then have no way to tell whether unloading it would strand + /// live allocations. + #[error( + "memory plugin {path} does not export the required entry point `{symbol}`; every nxmem \ + plugin must export NxmemNegotiate, NxmemCreateAllocatorFactories and \ + NxmemQueryUnloadReadiness" + )] + MissingSymbol { path: String, symbol: String }, + + /// The plugin and host could not agree on a version or capability set. + #[error("memory plugin {path} failed version negotiation: {reason}")] + Negotiation { path: String, reason: String }, + + /// The plugin returned a failing status. + #[error("memory plugin call `{operation}` failed: {}", .status.describe())] + Call { + operation: &'static str, + status: NxmemStatus, + }, + + /// The plugin violated the contract in a way that is not a mere failure. + #[error("memory plugin {path} violated the nxmem contract: {reason}")] + Contract { path: String, reason: String }, + + /// The named mechanism does not exist in this plugin. + #[error("memory plugin {path} exposes no mechanism named `{name}`; it offers: {available}")] + UnknownMechanism { + path: String, + name: String, + available: String, + }, +} + +impl PluginError { + /// Build a call failure from a status, for a named operation. + pub fn call(operation: &'static str, status: NxmemStatus) -> Self { + Self::Call { operation, status } + } +} + +/// Turn a failing plugin status into a [`MemoryError`]. +/// +/// The mapping is deliberately narrow. Only genuine capacity refusals become +/// [`MemoryError::AllocationFailed`]; contract violations keep their own text +/// so a misbehaving plugin is not mistaken for an out-of-memory condition. +pub fn status_to_memory_error( + operation: &str, + tier: Tier, + requested: u64, + status: &NxmemStatus, +) -> MemoryError { + // `describe` already renders the code name, including `UNKNOWN(n)` for a + // code this host does not recognise, so a newer plugin's failure stays + // traceable instead of being flattened into a generic message. + let detail = status.describe(); + let reason = match status.status_code() { + Some(NxmemStatusCode::OutOfMemory) => format!( + "the memory plugin is out of {} memory for `{operation}`: {detail}", + tier.name() + ), + Some(NxmemStatusCode::UnsupportedCapability) | Some(NxmemStatusCode::NotImplemented) => { + format!( + "the memory plugin does not support `{operation}`: {detail}; select a mechanism \ + that advertises the capability this operation needs" + ) + } + Some(NxmemStatusCode::WrongDevice) + | Some(NxmemStatusCode::WrongMechanism) + | Some(NxmemStatusCode::UnknownAllocation) + | Some(NxmemStatusCode::InvalidArgument) => format!( + "the memory plugin rejected `{operation}` as misdirected or invalid: {detail}; this \ + is a host-side routing bug, not a capacity problem" + ), + _ => format!("the memory plugin failed `{operation}`: {detail}"), + }; + MemoryError::AllocationFailed { + tier: tier.name(), + requested, + reason, + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn an_out_of_memory_status_keeps_the_tier_and_size() { + let status = NxmemStatus::with_message(NxmemStatusCode::OutOfMemory, "the pool is empty"); + let error = status_to_memory_error("allocate", Tier::Device, 4096, &status); + let text = error.to_string(); + assert!(text.contains("4096"), "{text}"); + assert!(text.contains("device"), "{text}"); + assert!(text.contains("the pool is empty"), "{text}"); + } + + #[test] + fn an_unsupported_status_says_how_to_fix_it() { + let status = + NxmemStatus::with_message(NxmemStatusCode::UnsupportedCapability, "no virtual backing"); + let error = status_to_memory_error("commit_range", Tier::Device, 0, &status); + let text = error.to_string(); + assert!(text.contains("advertises the capability"), "{text}"); + } + + #[test] + fn a_wrong_mechanism_status_is_not_reported_as_out_of_memory() { + let status = NxmemStatus::with_message(NxmemStatusCode::WrongMechanism, "foreign allocation"); + let error = status_to_memory_error("deallocate", Tier::Device, 0, &status); + let text = error.to_string(); + assert!(text.contains("misdirected"), "{text}"); + assert!(!text.contains("is out of"), "{text}"); + } + + #[test] + fn an_unrecognised_wire_code_is_named_rather_than_guessed() { + let mut status = NxmemStatus::with_message(NxmemStatusCode::InternalError, "boom"); + status.code = 9_999; + let error = status_to_memory_error("allocate", Tier::Host, 8, &status); + let text = error.to_string(); + assert!(text.contains("UNKNOWN(9999)"), "{text}"); + } +} diff --git a/crates/onnx-runtime-memory-host/src/lib.rs b/crates/onnx-runtime-memory-host/src/lib.rs new file mode 100644 index 0000000000..e69de29bb2 diff --git a/crates/onnx-runtime-memory-host/src/loader.rs b/crates/onnx-runtime-memory-host/src/loader.rs new file mode 100644 index 0000000000..a39b77330f --- /dev/null +++ b/crates/onnx-runtime-memory-host/src/loader.rs @@ -0,0 +1,648 @@ +//! Loading an nxmem plugin, negotiating the contract, and pinning the module. +//! +//! # Pinning +//! +//! Everything reachable from a loaded plugin — factories, allocators, queued +//! releases — holds an [`Arc`]. [`PluginModule`] declares its +//! [`libloading::Library`] **last**, so the library is the final field to drop. +//! Code pages therefore stay mapped until every plugin object built on top of +//! them is gone. + +use std::path::{Path, PathBuf}; +use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; + +use onnx_runtime_memory_abi::{ + NXMEM_CAP_ALLOCATOR, NXMEM_SYMBOL_CREATE_ALLOCATOR_FACTORIES, NXMEM_SYMBOL_NEGOTIATE, + NXMEM_SYMBOL_QUERY_UNLOAD_READINESS, NxmemAllocatorFactoryVtable, + NxmemCreateAllocatorFactoriesFn, NxmemNegotiateFn, NxmemNegotiateRequest, + NxmemNegotiateResponse, NxmemQueryUnloadReadinessFn, NxmemUnloadReport, NxmemVersionRange, + validate_negotiation, +}; +use onnx_runtime_memory_api::DeviceKey; + +use crate::allocator::{HostReclaim, PluginAllocator, open_allocator}; +use crate::error::PluginError; + +/// The most factories a single plugin may publish. +/// +/// A fixed ceiling means the host never allocates on behalf of a plugin's +/// unvalidated count, and a runaway plugin cannot exhaust host memory during +/// enumeration. +pub const MAX_FACTORIES: usize = 64; + +/// What the host and plugin agreed on. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct NegotiatedAbi { + /// The agreed major version. Calls are only legal within one major. + pub major: u32, + /// The agreed minor version. This is the **ceiling**: no vtable may claim + /// a higher level, and every vtable is read at its own declared level. + pub minor: u32, + /// The capability bits the plugin claims and the host offered. + pub capability_flags: u64, +} + +/// A loaded plugin module. +/// +/// Field order is load-bearing: `library` is last so it unmaps after every +/// other field has been dropped. Rust drops struct fields in declaration +/// order, so any field that might still call into plugin code during its own +/// drop must be declared before `library`. +#[derive(Debug)] +pub struct PluginModule { + path: PathBuf, + negotiated: NegotiatedAbi, + query_unload: NxmemQueryUnloadReadinessFn, + /// Allocators the host has opened and not yet dropped. + live_allocators: AtomicU64, + /// Allocations the host believes are live. + live_allocations: AtomicU64, + /// Deferred releases the host has queued and not seen retire. + queued_releases: AtomicU64, + /// Declared last: the library must outlive everything above it. + library: libloading::Library, +} + +impl PluginModule { + /// Where the module was loaded from. + pub fn path(&self) -> &Path { + &self.path + } + + /// The agreed contract. + pub fn negotiated(&self) -> NegotiatedAbi { + self.negotiated + } + + /// The library handle, for tests that want to reach a raw symbol. + pub fn library(&self) -> &libloading::Library { + &self.library + } + + pub(crate) fn allocator_opened(&self) { + self.live_allocators.fetch_add(1, Ordering::AcqRel); + } + + pub(crate) fn allocator_closed(&self) { + self.live_allocators.fetch_sub(1, Ordering::AcqRel); + } + + pub(crate) fn allocation_opened(&self) { + self.live_allocations.fetch_add(1, Ordering::AcqRel); + } + + pub(crate) fn allocation_closed(&self) { + self.live_allocations.fetch_sub(1, Ordering::AcqRel); + } + + pub(crate) fn release_queued(&self) { + self.queued_releases.fetch_add(1, Ordering::AcqRel); + } + + pub(crate) fn release_retired(&self) { + self.queued_releases.fetch_sub(1, Ordering::AcqRel); + } + + /// The host's own view of what is still live. + /// + /// This is deliberately independent of the plugin's report: unload is + /// gated on both, so a bug on one side cannot unmap code the other side is + /// about to enter. + pub fn host_live_counts(&self) -> HostLiveCounts { + HostLiveCounts { + allocators: self.live_allocators.load(Ordering::Acquire), + allocations: self.live_allocations.load(Ordering::Acquire), + queued_releases: self.queued_releases.load(Ordering::Acquire), + } + } + + /// Ask the plugin what it still owns. + pub fn unload_report(&self) -> Result { + let mut report = NxmemUnloadReport::zeroed(); + // SAFETY: `report` is a valid, writable, correctly aligned local. The + // symbol was resolved from this still-mapped library. + let status = unsafe { (self.query_unload)(&raw mut report) }; + if !status.is_ok() { + return Err(PluginError::call("NxmemQueryUnloadReadiness", status)); + } + Ok(report) + } +} + +/// The host's own tally of live plugin-backed objects. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub struct HostLiveCounts { + pub allocators: u64, + pub allocations: u64, + pub queued_releases: u64, +} + +impl HostLiveCounts { + /// Total live objects tracked by the host. + pub const fn total(&self) -> u64 { + self.allocators + .saturating_add(self.allocations) + .saturating_add(self.queued_releases) + } +} + +/// One mechanism a plugin publishes. +/// +/// # Ownership +/// +/// The host owns the factory and calls its `release` slot exactly once, from +/// [`Drop`]. Per the contract, releasing a factory must not invalidate +/// allocators already opened from it, so an allocator outliving its factory is +/// legal and is exercised by the test suite. +#[derive(Debug)] +pub struct PluginFactory { + name: String, + device: DeviceKey, + capability_flags: u64, + vtable: NxmemAllocatorFactoryVtable, + module: Arc, +} + +// SAFETY: the vtable is a plain-data copy of the plugin's function pointers +// and an opaque `ctx`. The nxmem contract requires every slot to be callable +// from any host thread and requires the plugin to synchronise its own state. +unsafe impl Send for PluginFactory {} +// SAFETY: as above; `&PluginFactory` only ever reads the copied vtable. +unsafe impl Sync for PluginFactory {} + +impl PluginFactory { + /// The mechanism name, as published by the plugin. + pub fn name(&self) -> &str { + &self.name + } + + /// The device this mechanism serves. + pub fn device(&self) -> DeviceKey { + self.device + } + + /// The capabilities this mechanism claims. + pub fn capability_flags(&self) -> u64 { + self.capability_flags + } + + /// The module this factory came from. + pub fn module(&self) -> &Arc { + &self.module + } + + /// Open an allocator from this mechanism. + /// + /// `reclaim` is the host's reclaim hook. It may be invoked reentrantly + /// from inside a plugin call and from plugin-owned threads, so it must not + /// block indefinitely and must not be called with any governance lock + /// held. + pub fn open( + &self, + required_capability_flags: u64, + reclaim: Option>, + ) -> Result { + open_allocator(self, required_capability_flags, reclaim) + } + + pub(crate) fn vtable(&self) -> &NxmemAllocatorFactoryVtable { + &self.vtable + } +} + +impl Drop for PluginFactory { + fn drop(&mut self) { + if let Some(release) = self.vtable.release { + // SAFETY: `ctx` came from this factory's vtable and `release` is + // called exactly once, here. No host lock is held. + unsafe { release(self.vtable.ctx) }; + } + } +} + +/// A loaded plugin and the mechanisms it publishes. +#[derive(Debug)] +pub struct MemoryPlugin { + /// Declared before `module` so factories release before the module (and + /// therefore the library) can possibly go away. + factories: Vec, + module: Arc, +} + +impl MemoryPlugin { + /// Load a plugin using this host's current supported version range. + pub fn load(path: impl AsRef) -> Result { + Self::load_with_host_range(path, NxmemVersionRange::current()) + } + + /// Load a plugin while pretending the host only supports `range`. + /// + /// This is the honest way to test both negotiation directions: an older + /// host meeting a newer plugin, and a version mismatch, both driven + /// through the real plugin entry point and the real loader rather than a + /// parallel re-implementation. + pub fn load_with_host_range( + path: impl AsRef, + range: NxmemVersionRange, + ) -> Result { + let path = path.as_ref().to_path_buf(); + let display = path.display().to_string(); + + // SAFETY: loading a dynamic library runs its initialisers. The caller + // is trusted to name a library it intends to load, exactly as for any + // other plugin surface in this workspace. + let library = unsafe { libloading::Library::new(&path) }.map_err(|source| { + PluginError::Open { + path: display.clone(), + source, + } + })?; + + let negotiate = resolve::(&library, NXMEM_SYMBOL_NEGOTIATE, &display)?; + let create_factories = resolve::( + &library, + NXMEM_SYMBOL_CREATE_ALLOCATOR_FACTORIES, + &display, + )?; + let query_unload = resolve::( + &library, + NXMEM_SYMBOL_QUERY_UNLOAD_READINESS, + &display, + )?; + + let request = NxmemNegotiateRequest::with_range(range); + let mut response = NxmemNegotiateResponse::zeroed(); + // SAFETY: both pointers address valid, aligned, writable locals that + // outlive the call. The plugin may not retain either. + let status = unsafe { negotiate(&raw const request, &raw mut response) }; + if !status.is_ok() { + return Err(PluginError::Negotiation { + path: display, + reason: status.describe(), + }); + } + validate_negotiation(&range, &response).map_err(|rejection| PluginError::Negotiation { + path: display.clone(), + reason: rejection.reason, + })?; + + if response.capability_flags & NXMEM_CAP_ALLOCATOR == 0 { + return Err(PluginError::Negotiation { + path: display, + reason: String::from( + "the plugin does not claim the allocator capability; a memory plugin that \ + cannot allocate has nothing to contribute", + ), + }); + } + + let negotiated = NegotiatedAbi { + major: response.agreed_major, + minor: response.agreed_minor, + capability_flags: response.capability_flags, + }; + + let module = Arc::new(PluginModule { + path, + negotiated, + query_unload, + live_allocators: AtomicU64::new(0), + live_allocations: AtomicU64::new(0), + queued_releases: AtomicU64::new(0), + library, + }); + + let factories = enumerate_factories(&module, create_factories, &display)?; + Ok(Self { factories, module }) + } + + /// The agreed contract. + pub fn negotiated(&self) -> NegotiatedAbi { + self.module.negotiated + } + + /// The loaded module. + pub fn module(&self) -> &Arc { + &self.module + } + + /// Every mechanism the plugin published. + pub fn factories(&self) -> &[PluginFactory] { + &self.factories + } + + /// Find a mechanism by name. + pub fn factory(&self, name: &str) -> Result<&PluginFactory, PluginError> { + self.factories + .iter() + .find(|factory| factory.name == name) + .ok_or_else(|| PluginError::UnknownMechanism { + path: self.module.path.display().to_string(), + name: name.to_string(), + available: self + .factories + .iter() + .map(|factory| factory.name.as_str()) + .collect::>() + .join(", "), + }) + } + + /// Attempt to unload the plugin. + /// + /// Unload is refused while **either** side still reports live objects. On + /// refusal the plugin is handed back, so a caller can retire the + /// outstanding work and try again — this is a deferral, not a leak. + pub fn try_unload(self) -> Result<(), UnloadRejection> { + let host = self.module.host_live_counts(); + if host.total() != 0 { + return Err(UnloadRejection { + reason: format!( + "the host still holds {} allocator(s), {} allocation(s) and {} queued \ + release(s); retire them before unloading", + host.allocators, host.allocations, host.queued_releases + ), + report: NxmemUnloadReport::zeroed(), + host, + plugin: Some(self), + }); + } + + let report = match self.module.unload_report() { + Ok(report) => report, + Err(error) => { + return Err(UnloadRejection { + reason: format!( + "the plugin could not report unload readiness, so unloading it would be \ + a guess: {error}" + ), + report: NxmemUnloadReport::zeroed(), + host, + plugin: Some(self), + }); + } + }; + if report.total() != 0 { + return Err(UnloadRejection { + reason: format!( + "the plugin still owns {} allocator(s), {} allocation(s), {} view(s), {} \ + capability object(s) and {} queued release(s)", + report.live_allocators, + report.live_allocations, + report.live_views, + report.live_capabilities, + report.queued_releases + ), + report, + host, + plugin: Some(self), + }); + } + + // Both sides agree nothing is live. Release the factories, then check + // that nothing else still pins the module before letting it unmap. + let Self { + factories, + module, + .. + } = self; + drop(factories); + + match Arc::try_unwrap(module) { + Ok(module) => { + drop(module); + Ok(()) + } + Err(module) => { + // Something outside this handle still pins the module. The + // factories are gone, so the plugin cannot be handed back — + // but the module stays mapped, which is the safe direction. + let reason = format!( + "{} handle(s) outside this plugin still pin the module, so it stays mapped", + Arc::strong_count(&module) + ); + Err(UnloadRejection { + reason, + report, + host, + plugin: None, + }) + } + } + } +} + +/// Why an unload was refused or deferred. +#[derive(Debug)] +pub struct UnloadRejection { + /// A human-readable explanation naming what is still live. + pub reason: String, + /// The plugin's own report at the moment of refusal. + pub report: NxmemUnloadReport, + /// The host's own tally at the moment of refusal. + pub host: HostLiveCounts, + /// The plugin, handed back so the caller can retire work and retry. + /// + /// `None` only in the terminal case where the module is pinned by a handle + /// this loader does not own; the module then simply stays mapped. + pub plugin: Option, +} + +impl UnloadRejection { + /// Recover the plugin so the caller can retire outstanding work and retry. + pub fn into_plugin(self) -> Option { + self.plugin + } +} + +impl std::fmt::Display for UnloadRejection { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "cannot unload the memory plugin: {}", self.reason) + } +} + +impl std::error::Error for UnloadRejection {} + +fn resolve( + library: &libloading::Library, + symbol: &[u8], + path: &str, +) -> Result { + // SAFETY: the symbol is looked up by name and immediately copied out of + // the `Symbol` guard. The resulting function pointer stays valid because + // `PluginModule` keeps the library mapped for as long as anything can call + // it. + let found = unsafe { library.get::(symbol) }.map_err(|_| PluginError::MissingSymbol { + path: path.to_string(), + symbol: String::from_utf8_lossy(symbol).into_owned(), + })?; + Ok(*found) +} + +fn enumerate_factories( + module: &Arc, + create: NxmemCreateAllocatorFactoriesFn, + path: &str, +) -> Result, PluginError> { + let mut raw = [core::ptr::null::(); MAX_FACTORIES]; + let mut count: u64 = 0; + // SAFETY: `raw` has exactly `MAX_FACTORIES` writable slots and `count` is a + // valid writable local. + let status = unsafe { create(raw.as_mut_ptr(), MAX_FACTORIES as u64, &raw mut count) }; + if !status.is_ok() { + return Err(PluginError::call("NxmemCreateAllocatorFactories", status)); + } + if count > MAX_FACTORIES as u64 { + return Err(PluginError::Contract { + path: path.to_string(), + reason: format!( + "the plugin reported {count} factories after being told it may write at most \ + {MAX_FACTORIES}; the extra entries were never written and cannot be read" + ), + }); + } + + let minor = module.negotiated.minor; + let mut factories = Vec::with_capacity(count as usize); + for slot in raw.iter().take(count as usize) { + // SAFETY: the plugin wrote this pointer in response to the call above. + // `read_prefix` null-checks, alignment-checks, and size-checks it + // before reading any field, and copies rather than borrowing. + let vtable = unsafe { NxmemAllocatorFactoryVtable::read_prefix(*slot, minor) } + .map_err(|status| PluginError::call("factory vtable", status))?; + vtable + .validate_required() + .map_err(|status| PluginError::call("factory vtable", status))?; + + let device = device_key(vtable.device).ok_or_else(|| PluginError::Contract { + path: path.to_string(), + reason: format!( + "a factory declared tier code {} which this host does not know; tiers are never \ + guessed", + vtable.device.tier + ), + })?; + + // SAFETY: the contract requires `name` to be a NUL-terminated UTF-8 + // string that stays valid until the factory's final `release`, which + // this host has not called yet. + let name = unsafe { read_c_string(vtable.name) }.ok_or_else(|| PluginError::Contract { + path: path.to_string(), + reason: String::from("a factory published a null or non-UTF-8 name"), + })?; + + factories.push(PluginFactory { + name, + device, + capability_flags: vtable.capability_flags, + vtable, + module: Arc::clone(module), + }); + } + Ok(factories) +} + +/// Translate an ABI device id into the internal key, refusing unknown tiers. +pub(crate) fn device_key(device: onnx_runtime_memory_abi::NxmemDeviceId) -> Option { + use onnx_runtime_memory_abi::{NXMEM_TIER_DEVICE, NXMEM_TIER_DISK, NXMEM_TIER_HOST}; + use onnx_runtime_memory_api::Tier; + let tier = match device.tier { + NXMEM_TIER_DEVICE => Tier::Device, + NXMEM_TIER_HOST => Tier::Host, + NXMEM_TIER_DISK => Tier::Disk, + _ => return None, + }; + Some(DeviceKey { + tier, + index: device.index, + }) +} + +/// Translate an internal key into an ABI device id. +pub(crate) fn device_id(device: DeviceKey) -> onnx_runtime_memory_abi::NxmemDeviceId { + use onnx_runtime_memory_abi::{ + NXMEM_TIER_DEVICE, NXMEM_TIER_DISK, NXMEM_TIER_HOST, NxmemDeviceId, + }; + use onnx_runtime_memory_api::Tier; + let tier = match device.tier { + Tier::Device => NXMEM_TIER_DEVICE, + Tier::Host => NXMEM_TIER_HOST, + Tier::Disk => NXMEM_TIER_DISK, + }; + NxmemDeviceId { + tier, + index: device.index, + } +} + +/// Read a NUL-terminated UTF-8 string published by a plugin. +/// +/// # Safety +/// +/// `ptr` must be null or point to a NUL-terminated byte string that stays +/// valid for the duration of the call. +unsafe fn read_c_string(ptr: *const u8) -> Option { + if ptr.is_null() { + return None; + } + // SAFETY: delegated to this function's contract. + let c_str = unsafe { core::ffi::CStr::from_ptr(ptr.cast()) }; + c_str.to_str().ok().map(str::to_owned) +} + +#[cfg(test)] +mod tests { + use super::*; + use onnx_runtime_memory_abi::NxmemDeviceId; + use onnx_runtime_memory_api::Tier; + + #[test] + fn device_ids_round_trip_through_both_directions() { + for key in [ + DeviceKey::HOST, + DeviceKey::device(0), + DeviceKey::device(7), + DeviceKey { + tier: Tier::Disk, + index: 3, + }, + ] { + assert_eq!(device_key(device_id(key)), Some(key)); + } + } + + #[test] + fn an_unknown_tier_code_is_refused_rather_than_guessed() { + let unknown = NxmemDeviceId { + tier: 4_242, + index: 0, + }; + assert_eq!(device_key(unknown), None); + } + + #[test] + fn host_live_counts_sum_every_axis() { + let counts = HostLiveCounts { + allocators: 1, + allocations: 2, + queued_releases: 3, + }; + assert_eq!(counts.total(), 6); + assert_eq!(HostLiveCounts::default().total(), 0); + } + + #[test] + fn a_null_name_is_refused() { + // SAFETY: a null pointer is the documented "no name" case. + assert_eq!(unsafe { read_c_string(core::ptr::null()) }, None); + } + + #[test] + fn a_valid_name_is_read_back() { + let raw = c"eager"; + // SAFETY: `raw` is a NUL-terminated literal alive for the whole call. + let name = unsafe { read_c_string(raw.as_ptr().cast()) }; + assert_eq!(name.as_deref(), Some("eager")); + } +} diff --git a/crates/onnx-runtime-memory-testplugin/Cargo.toml b/crates/onnx-runtime-memory-testplugin/Cargo.toml new file mode 100644 index 0000000000..351bd948d0 --- /dev/null +++ b/crates/onnx-runtime-memory-testplugin/Cargo.toml @@ -0,0 +1,20 @@ +[package] +name = "onnx-runtime-memory-testplugin" +description = "Out-of-tree-style nxmem memory plugin used to exercise the ABI" +version.workspace = true +edition.workspace = true +license.workspace = true +repository.workspace = true +publish = false + +# `cdylib` is what the host actually dlopens; `rlib` exists only so the host's +# integration tests can depend on this package and therefore force cargo to +# build the cdylib before they run. +[lib] +crate-type = ["cdylib", "rlib"] + +[dependencies] +onnx-runtime-memory-abi = { workspace = true } + +[lints] +workspace = true diff --git a/crates/onnx-runtime-memory-testplugin/src/lib.rs b/crates/onnx-runtime-memory-testplugin/src/lib.rs new file mode 100644 index 0000000000..e69de29bb2 From dbe09210a27296823a11c3e300a153861c1aeebb Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Wed, 19 Aug 2026 01:31:36 -0700 Subject: [PATCH 2/3] nxmem: test plugin, ABI test suite, prefix clamping --- Cargo.lock | 10 + crates/onnx-runtime-memory-abi/src/types.rs | 49 + crates/onnx-runtime-memory-abi/src/vtable.rs | 124 +- .../onnx-runtime-memory-host/src/allocator.rs | 129 +- crates/onnx-runtime-memory-host/src/lib.rs | 47 + crates/onnx-runtime-memory-host/src/loader.rs | 33 +- .../tests/nxmem_abi.rs | 941 ++++++++++ .../tests/nxmem_abi_unload_gate.rs | 66 + .../tests/support/testplugin.rs | 80 + .../onnx-runtime-memory-testplugin/src/lib.rs | 1672 +++++++++++++++++ 10 files changed, 3032 insertions(+), 119 deletions(-) create mode 100644 crates/onnx-runtime-memory-host/tests/nxmem_abi.rs create mode 100644 crates/onnx-runtime-memory-host/tests/nxmem_abi_unload_gate.rs create mode 100644 crates/onnx-runtime-memory-host/tests/support/testplugin.rs diff --git a/Cargo.lock b/Cargo.lock index 47823a3042..af1e0f6432 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2573,10 +2573,20 @@ dependencies = [ [[package]] name = "onnx-runtime-memory-host" version = "0.1.0-dev.5" +dependencies = [ + "libloading 0.8.9", + "onnx-runtime-memory-abi", + "onnx-runtime-memory-api", + "onnx-runtime-memory-testplugin", + "thiserror 2.0.19", +] [[package]] name = "onnx-runtime-memory-testplugin" version = "0.1.0-dev.5" +dependencies = [ + "onnx-runtime-memory-abi", +] [[package]] name = "onnx-runtime-operator-selection" diff --git a/crates/onnx-runtime-memory-abi/src/types.rs b/crates/onnx-runtime-memory-abi/src/types.rs index 15c6b1b233..7d6d8fb440 100644 --- a/crates/onnx-runtime-memory-abi/src/types.rs +++ b/crates/onnx-runtime-memory-abi/src/types.rs @@ -339,6 +339,25 @@ pub struct NxmemReleaseCompletion { pub outcome: NxmemReleaseOutcome, } +impl NxmemReleaseCompletion { + /// A completion record with the correct `struct_size`. + pub fn new( + ticket: u64, + mechanism_id: u64, + allocation_id: u64, + outcome: NxmemReleaseOutcome, + ) -> Self { + Self { + struct_size: core::mem::size_of::() as u32, + reserved: 0, + ticket, + mechanism_id, + allocation_id, + outcome, + } + } +} + // ─── Host callbacks ───────────────────────────────────────────────────────── /// A plugin's request that the host free cached memory. @@ -357,6 +376,19 @@ pub struct NxmemReclaimRequest { pub bytes: u64, } +impl NxmemReclaimRequest { + /// A reclaim request with the correct `struct_size`. + pub fn new(mechanism_id: u64, device: NxmemDeviceId, bytes: u64) -> Self { + Self { + struct_size: core::mem::size_of::() as u32, + reserved: 0, + mechanism_id, + device, + bytes, + } + } +} + /// Callbacks the host offers to a plugin allocator. /// /// # Ownership @@ -530,6 +562,23 @@ pub struct NxmemSharedPrefixCommitRequest { pub byte_offset: u64, } +impl NxmemSharedPrefixCommitRequest { + /// A commit request with the correct `struct_size`. + pub fn new( + prefix: NxmemSharedPrefixHandle, + allocation: NxmemAllocation, + byte_offset: u64, + ) -> Self { + Self { + struct_size: core::mem::size_of::() as u32, + reserved: 0, + prefix, + allocation, + byte_offset, + } + } +} + /// The accounting result of mapping a shared prefix. #[repr(C)] #[derive(Debug, Clone, Copy, Default)] diff --git a/crates/onnx-runtime-memory-abi/src/vtable.rs b/crates/onnx-runtime-memory-abi/src/vtable.rs index aeff40327b..3bc16b898e 100644 --- a/crates/onnx-runtime-memory-abi/src/vtable.rs +++ b/crates/onnx-runtime-memory-abi/src/vtable.rs @@ -30,8 +30,13 @@ //! Negotiation fixes the **ceiling**. Each vtable independently declares the //! level it actually implements in `abi_minor`, so one plugin may ship a //! current mechanism next to one still built at the baseline prefix. A vtable -//! declaring a level *above* the negotiated ceiling is rejected: the host did -//! not agree to interpret those slots. +//! declaring a level *above* the negotiated ceiling is **clamped**, not +//! rejected: its struct is a strict superset, so the host reads only the +//! prefix it agreed to and refuses to call any slot introduced above the +//! ceiling. `read_prefix` rewrites `abi_minor` to that effective level, so a +//! caller need only consult the value it gets back. A vtable whose +//! `struct_size` is smaller than the level it *claims* is still rejected: that +//! is a sender that contradicts itself, not an old one. use core::ffi::c_void; @@ -149,16 +154,6 @@ fn short_struct(what: &str, declared: usize, required: usize, minor: u32) -> Nxm ) } -fn level_too_new(what: &str, declared: u32, ceiling: u32) -> NxmemStatus { - NxmemStatus::with_message( - NxmemStatusCode::VersionMismatch, - &format!( - "nxmem: {what} declares nxmem minor {declared} but the negotiated ceiling is \ - {ceiling}; the host will not interpret slots it did not agree to" - ), - ) -} - // ─── Allocator vtable ─────────────────────────────────────────────────────── /// The ordinary allocator mechanism plus its optional capability slots. @@ -314,13 +309,9 @@ impl NxmemAllocatorVtable { // SAFETY: the size check above proved `abi_minor` (offset 4) is inside // the readable prefix. let declared_minor = unsafe { core::ptr::read_unaligned(ptr.cast::().add(1)) }; - if declared_minor > negotiated_minor { - return Err(level_too_new( - "allocator vtable", - declared_minor, - negotiated_minor, - )); - } + // The sender's own claim must be self-consistent: a vtable that says + // "minor 1" must be at least as big as minor 1 requires, whatever the + // host negotiated. let required = Self::required_struct_size(declared_minor); if declared_size < required { return Err(short_struct( @@ -330,11 +321,17 @@ impl NxmemAllocatorVtable { declared_minor, )); } + // A sender ahead of the negotiated ceiling is not an error: its struct + // is a strict superset, so the host reads the prefix it agreed to and + // clamps the level it will act on. This is what lets a baseline host + // drive a current plugin. + let effective_minor = declared_minor.min(negotiated_minor); // SAFETY: `declared_size` bytes are readable and we copy no more. let mut vtable = unsafe { copy_prefix(ptr, declared_size) }; - // A slot the sender does not define must never be called, even if this - // build knows about it. - if declared_minor < 1 { + vtable.abi_minor = effective_minor; + // A slot neither side agreed to must never be called, even if both + // builds happen to know about it. + if effective_minor < 1 { vtable.release_allocation = None; } vtable.validate_required()?; @@ -486,15 +483,13 @@ impl NxmemVirtualBackingVtable { } // SAFETY: the size check proved `abi_minor` is readable. let declared_minor = unsafe { core::ptr::read_unaligned(ptr.cast::().add(1)) }; - if declared_minor > negotiated_minor { - return Err(level_too_new( - "virtual-backing vtable", - declared_minor, - negotiated_minor, - )); - } + // A sender ahead of the negotiated ceiling is not an error: its struct + // is a strict superset, so the host simply reads the prefix it agreed + // to and clamps the level it will act on. + let effective_minor = declared_minor.min(negotiated_minor); // SAFETY: `declared_size` bytes are readable and we copy no more. - let vtable = unsafe { copy_prefix(ptr, declared_size) }; + let mut vtable = unsafe { copy_prefix(ptr, declared_size) }; + vtable.abi_minor = effective_minor; if vtable.mechanism_id != owner_mechanism_id { return Err(NxmemStatus::with_message( NxmemStatusCode::WrongMechanism, @@ -618,15 +613,13 @@ impl NxmemSharedMappingVtable { } // SAFETY: the size check proved `abi_minor` is readable. let declared_minor = unsafe { core::ptr::read_unaligned(ptr.cast::().add(1)) }; - if declared_minor > negotiated_minor { - return Err(level_too_new( - "shared-mapping vtable", - declared_minor, - negotiated_minor, - )); - } + // A sender ahead of the negotiated ceiling is not an error: its struct + // is a strict superset, so the host simply reads the prefix it agreed + // to and clamps the level it will act on. + let effective_minor = declared_minor.min(negotiated_minor); // SAFETY: `declared_size` bytes are readable and we copy no more. - let vtable = unsafe { copy_prefix(ptr, declared_size) }; + let mut vtable = unsafe { copy_prefix(ptr, declared_size) }; + vtable.abi_minor = effective_minor; if vtable.mechanism_id != owner_mechanism_id { return Err(NxmemStatus::with_message( NxmemStatusCode::WrongMechanism, @@ -718,15 +711,13 @@ impl NxmemAllocatorFactoryVtable { } // SAFETY: the size check proved `abi_minor` is readable. let declared_minor = unsafe { core::ptr::read_unaligned(ptr.cast::().add(1)) }; - if declared_minor > negotiated_minor { - return Err(level_too_new( - "allocator-factory vtable", - declared_minor, - negotiated_minor, - )); - } + // A sender ahead of the negotiated ceiling is not an error: its struct + // is a strict superset, so the host simply reads the prefix it agreed + // to and clamps the level it will act on. + let effective_minor = declared_minor.min(negotiated_minor); // SAFETY: `declared_size` bytes are readable and we copy no more. - let vtable = unsafe { copy_prefix(ptr, declared_size) }; + let mut vtable = unsafe { copy_prefix(ptr, declared_size) }; + vtable.abi_minor = effective_minor; if vtable.open_allocator.is_none() || vtable.release.is_none() { return Err(NxmemStatus::with_message( NxmemStatusCode::ShortStruct, @@ -834,13 +825,44 @@ mod tests { } #[test] - fn a_vtable_above_the_negotiated_ceiling_is_refused() { + fn a_vtable_above_the_negotiated_ceiling_is_clamped_not_refused() { let vtable = conforming_allocator(1); + assert!( + vtable.release_allocation.is_some(), + "the sender really does implement the minor-1 slot" + ); // SAFETY: `vtable` is a live, aligned local. - let status = - unsafe { NxmemAllocatorVtable::read_prefix(&vtable, 0) }.expect_err("too new"); - assert_eq!(status.status_code(), Some(NxmemStatusCode::VersionMismatch)); - assert!(status.describe().contains("negotiated ceiling")); + let read = unsafe { NxmemAllocatorVtable::read_prefix(&vtable, 0) } + .expect("a newer sender is a strict superset, so a baseline host can still read it"); + assert_eq!( + read.abi_minor, 0, + "the effective level is clamped to what both sides agreed" + ); + assert!( + read.release_allocation.is_none(), + "a slot neither side agreed to must never be callable" + ); + assert!( + read.allocate.is_some() && read.deallocate.is_some(), + "the baseline prefix still reads back intact" + ); + } + + #[test] + fn a_vtable_smaller_than_the_level_it_claims_is_refused() { + let mut vtable = conforming_allocator(1); + // The sender contradicts itself: it says "minor 1" but sizes itself + // for the baseline. That is not an old peer, it is a broken one. + vtable.struct_size = NxmemAllocatorVtable::MIN_STRUCT_SIZE_MINOR_0 as u32; + // SAFETY: `vtable` is a live, aligned local. + let status = unsafe { NxmemAllocatorVtable::read_prefix(&vtable, 1) } + .expect_err("a self-contradicting struct_size must be refused"); + assert_eq!(status.status_code(), Some(NxmemStatusCode::ShortStruct)); + assert!( + status.describe().contains("minor 1"), + "the refusal must name the level the sender claimed: {}", + status.describe() + ); } #[test] diff --git a/crates/onnx-runtime-memory-host/src/allocator.rs b/crates/onnx-runtime-memory-host/src/allocator.rs index ff5f5709be..bbbd9a883c 100644 --- a/crates/onnx-runtime-memory-host/src/allocator.rs +++ b/crates/onnx-runtime-memory-host/src/allocator.rs @@ -272,6 +272,14 @@ impl AllocatorCore { self.bridge.reclaim_failures.load(Ordering::Acquire) } + /// The callback table handed to the plugin. + /// + /// Exposed so a caller can confirm the table the plugin holds is the one + /// this host installed, and that its `abi_minor` matches negotiation. + pub fn host_callbacks(&self) -> &NxmemHostCallbacks { + &self.callbacks + } + /// Deferred releases the plugin has reported as retired, in report order. pub fn retired_releases(&self) -> Vec { self.bridge @@ -537,7 +545,7 @@ impl AllocatorCore { outcome.failure.describe() )) } - other => { + _unknown => { // An unknown state is never guessed. Treat it as residual // ownership so nothing is reused or double-refunded. self.module.allocation_closed(); @@ -554,21 +562,11 @@ impl AllocatorCore { align: record.align, }, ) - .tap_unknown(other) } } } } -/// Small helper so the unknown-state branch above stays readable. -trait TapUnknown: Sized { - fn tap_unknown(self, _state: u32) -> Self { - self - } -} - -impl TapUnknown for AllocationReleaseOutcome {} - impl Drop for AllocatorCore { fn drop(&mut self) { // Any allocation still in the map is a host-side leak, not something to @@ -1320,6 +1318,32 @@ impl Drop for PluginSharedPrefix { } /// Open an allocator from a factory, wiring up the host callback table. +/// Give back an allocator the host has decided it cannot use. +/// +/// `open_allocator` returning `Ok` means the plugin created state, so the host +/// owes it a `release` even when it then refuses the vtable. The refusal path +/// is the one place the host must reach into a struct it has already judged +/// untrustworthy, so it does the least it can: re-read the prefix, and call +/// `release` only if that succeeds and the slot is present. A plugin that +/// publishes a vtable too malformed to locate `release` cannot be given +/// anything back — such a plugin must not allocate state before returning. +/// +/// # Safety +/// +/// `ptr` must be null or the pointer the plugin just returned. +unsafe fn abandon_allocator(ptr: *const NxmemAllocatorVtable, minor: u32) { + // SAFETY: delegated to this function's contract; `read_prefix` performs + // every null, alignment, and size check before reading a field. + let Ok(vtable) = (unsafe { NxmemAllocatorVtable::read_prefix(ptr, minor) }) else { + return; + }; + let Some(release) = vtable.release else { + return; + }; + // SAFETY: `ctx` came from this vtable and no host lock is held. + unsafe { release(vtable.ctx) }; +} + pub(crate) fn open_allocator( factory: &PluginFactory, required_capability_flags: u64, @@ -1383,10 +1407,19 @@ pub(crate) fn open_allocator( // `read_prefix` null-, alignment- and size-checks it before reading any // field, and copies rather than borrowing. let vtable = unsafe { NxmemAllocatorVtable::read_prefix(raw_allocator, minor) } - .map_err(|status| PluginError::call("allocator vtable", status))?; - vtable - .validate_required() - .map_err(|status| PluginError::call("allocator vtable", status))?; + .map_err(|status| { + // SAFETY: the vtable was refused, so nothing may be called through + // it; `abandon_allocator` re-reads the prefix defensively and only + // calls `release` when the struct is well-formed enough to locate + // that slot. + unsafe { abandon_allocator(raw_allocator, minor) }; + PluginError::call("allocator vtable", status) + })?; + if let Err(status) = vtable.validate_required() { + // SAFETY: as above. Here the prefix did read, so `release` is known. + unsafe { abandon_allocator(raw_allocator, minor) }; + return Err(PluginError::call("allocator vtable", status)); + } let device = device_key(vtable.device).ok_or_else(|| PluginError::Contract { path: module.path().display().to_string(), @@ -1429,34 +1462,41 @@ pub(crate) fn open_allocator( }); } - // Capability vtables are read at the allocator's own level, and each must - // name the same mechanism instance. + // Capability vtables are read at the allocator's own level, and + // `read_prefix` itself refuses a vtable that names a different mechanism. let backing = read_capability( vtable.virtual_backing, vtable.capability_flags & NXMEM_CAP_VIRTUAL_BACKING != 0, - minor, - vtable.mechanism_id, "virtual backing", &module, - |ptr, minor| { - // SAFETY: delegated to `read_capability`, which only calls this - // with a pointer the plugin published for this capability. - unsafe { NxmemVirtualBackingVtable::read_prefix(ptr, minor) } + || { + // SAFETY: the pointer came from the allocator vtable the plugin + // just published; `read_prefix` null-, alignment- and size-checks + // it before reading any field. + unsafe { + NxmemVirtualBackingVtable::read_prefix( + vtable.virtual_backing, + minor, + vtable.mechanism_id, + ) + } }, - |vtable| vtable.mechanism_id, )?; let shared = read_capability( vtable.shared_mapping, vtable.capability_flags & NXMEM_CAP_SHARED_MAPPING != 0, - minor, - vtable.mechanism_id, "shared mapping", &module, - |ptr, minor| { + || { // SAFETY: as above. - unsafe { NxmemSharedMappingVtable::read_prefix(ptr, minor) } + unsafe { + NxmemSharedMappingVtable::read_prefix( + vtable.shared_mapping, + minor, + vtable.mechanism_id, + ) + } }, - |vtable| vtable.mechanism_id, )?; // SAFETY: the contract requires `name` to stay valid until the final @@ -1502,27 +1542,24 @@ pub(crate) fn open_allocator( }) } -#[allow(clippy::too_many_arguments)] fn read_capability( ptr: *const T, advertised: bool, - minor: u32, - mechanism_id: u64, label: &str, module: &Arc, - read: impl Fn(*const T, u32) -> Result, - id_of: impl Fn(&T) -> u64, + read: impl FnOnce() -> Result, ) -> Result, PluginError> { if !advertised { // An unsupported capability is represented explicitly: the flag is - // clear and the pointer is expected to be null. A non-null pointer - // behind a clear flag is a contract violation, not a bonus feature. + // clear *and* the pointer is null. A non-null pointer behind a clear + // flag is a contract violation, not a bonus feature, because the host + // would have no negotiated basis for calling it. if !ptr.is_null() { return Err(PluginError::Contract { path: module.path().display().to_string(), reason: format!( - "the allocator published a {label} vtable without advertising the capability; \ - an unsupported capability must be a null pointer and a clear flag" + "the allocator published a {label} vtable without advertising the \ + capability; an unsupported capability must be a null pointer and a clear flag" ), }); } @@ -1531,22 +1568,10 @@ fn read_capability( if ptr.is_null() { return Err(PluginError::Contract { path: module.path().display().to_string(), - reason: format!( - "the allocator advertised {label} but published a null vtable for it" - ), - }); - } - let vtable = read(ptr, minor).map_err(|status| PluginError::call("capability vtable", status))?; - if id_of(&vtable) != mechanism_id { - return Err(PluginError::Contract { - path: module.path().display().to_string(), - reason: format!( - "the {label} vtable names mechanism {} but the allocator is mechanism \ - {mechanism_id}; a capability must belong to one coherent mechanism", - id_of(&vtable) - ), + reason: format!("the allocator advertised {label} but published a null vtable for it"), }); } + let vtable = read().map_err(|status| PluginError::call("capability vtable", status))?; Ok(Some(vtable)) } diff --git a/crates/onnx-runtime-memory-host/src/lib.rs b/crates/onnx-runtime-memory-host/src/lib.rs index e69de29bb2..c5533a79f1 100644 --- a/crates/onnx-runtime-memory-host/src/lib.rs +++ b/crates/onnx-runtime-memory-host/src/lib.rs @@ -0,0 +1,47 @@ +//! # `onnx-runtime-memory-host` +//! +//! Loads an [nxmem][abi] memory plugin and adapts its C vtables to the +//! internal Rust interfaces in `onnx-runtime-memory-api`. +//! +//! [abi]: onnx_runtime_memory_abi +//! +//! Nothing internal leaks outward. A plugin sees only `#[repr(C)]` plain data, +//! raw pointers, and `extern "C"` function pointers; it never sees a Rust +//! trait object, `Arc`, enum layout, or allocator ownership. In the other +//! direction the loaded mechanism appears to the rest of the runtime as an +//! ordinary [`DeviceAllocator`](onnx_runtime_memory_api::DeviceAllocator) with +//! optional [`VirtualBacking`](onnx_runtime_memory_api::VirtualBacking) and +//! [`SharedMapping`](onnx_runtime_memory_api::SharedMapping) capabilities. +//! +//! # Lock discipline +//! +//! **No ABI call is made while any lock is held.** The one host-side map +//! (address to `allocation_id`) is locked, mutated, and unlocked before the +//! plugin is entered. Governance locks in the governor and provider layers are +//! likewise released before a plugin call, so a plugin is always free to +//! block, call back into the host, or spawn threads. +//! +//! # Pinning and unload +//! +//! Factories, allocators, capability views, shared prefixes, and queued +//! releases all hold an `Arc`, and [`PluginModule`] declares its +//! `libloading::Library` last. Unload is refused while either the host's own +//! counters or the plugin's [`NxmemUnloadReport`][report] report anything +//! live; the refusal hands the plugin back so the caller can retire the +//! outstanding work and retry. +//! +//! [report]: onnx_runtime_memory_abi::NxmemUnloadReport + +pub mod allocator; +pub mod error; +pub mod loader; + +pub use allocator::{ + AllocatorCore, HostReclaim, PluginAllocator, PluginSharedMapping, PluginSharedPrefix, + PluginVirtualBacking, RetiredRelease, +}; +pub use error::{PluginError, status_to_memory_error}; +pub use loader::{ + HostLiveCounts, MAX_FACTORIES, MemoryPlugin, NegotiatedAbi, PluginFactory, PluginModule, + UnloadRejection, +}; diff --git a/crates/onnx-runtime-memory-host/src/loader.rs b/crates/onnx-runtime-memory-host/src/loader.rs index a39b77330f..2b2fa10e0d 100644 --- a/crates/onnx-runtime-memory-host/src/loader.rs +++ b/crates/onnx-runtime-memory-host/src/loader.rs @@ -356,19 +356,9 @@ impl MemoryPlugin { /// outstanding work and try again — this is a deferral, not a leak. pub fn try_unload(self) -> Result<(), UnloadRejection> { let host = self.module.host_live_counts(); - if host.total() != 0 { - return Err(UnloadRejection { - reason: format!( - "the host still holds {} allocator(s), {} allocation(s) and {} queued \ - release(s); retire them before unloading", - host.allocators, host.allocations, host.queued_releases - ), - report: NxmemUnloadReport::zeroed(), - host, - plugin: Some(self), - }); - } - + // Ask the plugin first, unconditionally, so every rejection carries + // both tallies. A refusal that reports only one side leaves the caller + // guessing about which objects to retire. let report = match self.module.unload_report() { Ok(report) => report, Err(error) => { @@ -383,6 +373,20 @@ impl MemoryPlugin { }); } }; + + if host.total() != 0 { + return Err(UnloadRejection { + reason: format!( + "the host still holds {} allocator(s), {} allocation(s) and {} queued \ + release(s); retire them before unloading", + host.allocators, host.allocations, host.queued_releases + ), + report, + host, + plugin: Some(self), + }); + } + if report.total() != 0 { return Err(UnloadRejection { reason: format!( @@ -511,9 +515,6 @@ fn enumerate_factories( // before reading any field, and copies rather than borrowing. let vtable = unsafe { NxmemAllocatorFactoryVtable::read_prefix(*slot, minor) } .map_err(|status| PluginError::call("factory vtable", status))?; - vtable - .validate_required() - .map_err(|status| PluginError::call("factory vtable", status))?; let device = device_key(vtable.device).ok_or_else(|| PluginError::Contract { path: path.to_string(), diff --git a/crates/onnx-runtime-memory-host/tests/nxmem_abi.rs b/crates/onnx-runtime-memory-host/tests/nxmem_abi.rs new file mode 100644 index 0000000000..76dbae4c65 --- /dev/null +++ b/crates/onnx-runtime-memory-host/tests/nxmem_abi.rs @@ -0,0 +1,941 @@ +//! End-to-end tests for the nxmem memory ABI. +//! +//! Every test here loads the real `onnx-runtime-memory-testplugin` **cdylib** +//! with `dlopen` and drives it through the real [`MemoryPlugin`] loader and +//! the real [`PluginAllocator`] adapter. Nothing re-implements the state +//! machine under test: the assertions are made against the public entry points +//! the runtime itself calls. +//! +//! The test plugin's mechanisms are all host-memory backed, so the whole suite +//! runs on a machine with no accelerator. That is deliberate: the ABI contract +//! is device-agnostic, and pinning it to CUDA would make it untestable on the +//! development host. The CUDA-specific *users* of this ABI are compile-checked +//! only; they are not exercised here and this suite does not claim to. +//! +//! Required scenarios from the phase's acceptance criteria, and where each +//! lives: +//! +//! | Scenario | Test | +//! |---|---| +//! | version mismatch | `a_plugin_outside_the_hosts_major_range_is_refused` | +//! | short struct | `a_short_allocator_vtable_is_refused_before_any_slot_is_read` | +//! | missing optional capability | `a_mechanism_without_optional_capabilities_reports_none` | +//! | allocation and free | `allocate_and_release_round_trips_through_the_boundary` | +//! | lazy backing | `lazy_backing_commits_decommits_and_reports_mapped_bytes` | +//! | release ordering | `deferred_releases_retire_in_order_and_pin_the_module` | +//! | callback failure | `a_refusing_host_callback_fails_the_allocation_cleanly` | +//! | unload with live objects | `unload_is_refused_while_*` (three tests here, plus the queued-free case in `nxmem_abi_unload_gate.rs`) | +//! | older participant compatibility | `an_older_host_range_still_drives_the_current_plugin`, `a_minor_0_mechanism_works_under_a_minor_1_host` | + +use std::ptr::NonNull; +use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; + +use onnx_runtime_memory_abi::{ + NXMEM_ABI_VERSION_MAJOR, NXMEM_ABI_VERSION_MINOR, NXMEM_CAP_ALLOCATOR, + NXMEM_CAP_DEFERRED_RELEASE, NXMEM_CAP_SHARED_MAPPING, NXMEM_CAP_STRUCTURED_RELEASE, + NXMEM_CAP_VIRTUAL_BACKING, NxmemVersionRange, +}; +use onnx_runtime_memory_api::{ + AllocationCommitRange, AllocationReleaseOutcome, DeviceAllocator, DeviceKey, MemoryError, +}; +use onnx_runtime_memory_host::{HostReclaim, MemoryPlugin, PluginAllocator, PluginError}; + +// ─── locating the plugin ──────────────────────────────────────────────────── + +#[path = "support/testplugin.rs"] +mod testplugin; +use testplugin::testplugin_path; + +/// Serialises the whole suite. +/// +/// A plugin module's live-object counters are **process-wide** — that is the +/// contract, because unloading a module is a process-wide act. Two tests +/// holding allocators at the same time would therefore see each other's +/// objects through the unload gate. Every test takes this guard first so it +/// drops last, which makes the counters deterministic without weakening what +/// is being tested. +fn serial() -> std::sync::MutexGuard<'static, ()> { + static LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(()); + LOCK.lock().unwrap_or_else(|poison| poison.into_inner()) +} + +fn load() -> MemoryPlugin { + MemoryPlugin::load(testplugin_path()).expect("the test plugin loads under the current host") +} + +/// A reclaim hook that records what it was asked and answers as configured. +#[derive(Debug)] +struct ScriptedReclaim { + calls: AtomicU64, + last_bytes: AtomicU64, + refuse: bool, + grant: u64, +} + +impl ScriptedReclaim { + fn granting(grant: u64) -> Arc { + Arc::new(Self { + calls: AtomicU64::new(0), + last_bytes: AtomicU64::new(0), + refuse: false, + grant, + }) + } + + fn refusing() -> Arc { + Arc::new(Self { + calls: AtomicU64::new(0), + last_bytes: AtomicU64::new(0), + refuse: true, + grant: 0, + }) + } +} + +impl HostReclaim for ScriptedReclaim { + fn request_reclaim(&self, _device: DeviceKey, bytes: u64) -> Result { + self.calls.fetch_add(1, Ordering::AcqRel); + self.last_bytes.store(bytes, Ordering::Release); + if self.refuse { + Err(String::from("the scripted host has nothing to give back")) + } else { + Ok(self.grant) + } + } +} + +fn open(plugin: &MemoryPlugin, mechanism: &str, required: u64) -> PluginAllocator { + plugin + .factory(mechanism) + .unwrap_or_else(|error| panic!("mechanism `{mechanism}` must exist: {error}")) + .open(required, None) + .unwrap_or_else(|error| panic!("mechanism `{mechanism}` must open: {error}")) +} + +// ─── loading and negotiation ──────────────────────────────────────────────── + +#[test] +fn the_plugin_loads_and_publishes_every_mechanism() { + let _serial = serial(); + let plugin = load(); + let names: Vec<&str> = plugin + .factories() + .iter() + .map(|factory| factory.name()) + .collect(); + assert_eq!( + names, + onnx_runtime_memory_testplugin::MECHANISM_NAMES, + "the loader must surface exactly the mechanisms the plugin published" + ); + + let negotiated = plugin.negotiated(); + assert_eq!(negotiated.major, NXMEM_ABI_VERSION_MAJOR); + assert_eq!(negotiated.minor, NXMEM_ABI_VERSION_MINOR); + assert_ne!( + negotiated.capability_flags & NXMEM_CAP_ALLOCATOR, + 0, + "a plugin that cannot allocate is useless" + ); +} + +/// **Version mismatch.** +/// +/// Driven through the real `load_with_host_range` entry point rather than by +/// poking `negotiate` directly, so what is pinned is the loader's refusal, not +/// a private helper's return value. +#[test] +fn a_plugin_outside_the_hosts_major_range_is_refused() { + let _serial = serial(); + let error = MemoryPlugin::load_with_host_range( + testplugin_path(), + NxmemVersionRange { + major_min: 99, + major_max: 99, + minor_min: 0, + minor_max: 0, + }, + ) + .expect_err("a host from a different major era must not load this plugin"); + + match error { + PluginError::Negotiation { reason, .. } => { + assert!( + reason.contains("99") || reason.to_lowercase().contains("major"), + "the refusal must name the incompatible major version, got: {reason}" + ); + } + other => panic!("expected a negotiation refusal, got {other}"), + } +} + +/// **Older supported participant, host side.** +/// +/// A host that only speaks the baseline minor still drives the current plugin: +/// negotiation settles on the baseline and no minor-1 slot is used. +#[test] +fn an_older_host_range_still_drives_the_current_plugin() { + let _serial = serial(); + let plugin = MemoryPlugin::load_with_host_range( + testplugin_path(), + NxmemVersionRange { + major_min: NXMEM_ABI_VERSION_MAJOR, + major_max: NXMEM_ABI_VERSION_MAJOR, + minor_min: 0, + minor_max: 0, + }, + ) + .expect("a baseline host must still be able to load a current plugin"); + + assert_eq!(plugin.negotiated().minor, 0, "the ceiling is the host's"); + + let allocator = open(&plugin, "lazy", NXMEM_CAP_ALLOCATOR); + assert_eq!( + allocator.core().abi_minor(), + 0, + "a mechanism must not offer a slot the negotiated host cannot call" + ); + + // The whole allocate/free cycle still works at the baseline level, using + // the minor-0 `deallocate` slot rather than structured release. + let ptr = allocator.allocate(4096, 64).expect("baseline allocation"); + // SAFETY: `ptr` is live, and was allocated with exactly these parameters. + let outcome = unsafe { allocator.release(ptr, 4096, 64) }; + assert!( + matches!(outcome, AllocationReleaseOutcome::Complete { .. }), + "the baseline release path must still report a complete release, got {outcome:?}" + ); +} + +/// **Older supported participant, plugin side.** +/// +/// `legacy-1-0` declares minor 0 even though the module as a whole speaks +/// minor 1. The host must fall back for that mechanism alone, while a sibling +/// mechanism in the very same module keeps the newer slot. +#[test] +fn a_minor_0_mechanism_works_under_a_minor_1_host() { + let _serial = serial(); + let plugin = load(); + assert_eq!(plugin.negotiated().minor, NXMEM_ABI_VERSION_MINOR); + + let legacy = open(&plugin, "legacy-1-0", NXMEM_CAP_ALLOCATOR); + assert_eq!(legacy.core().abi_minor(), 0); + assert_eq!( + legacy.core().capability_flags() & NXMEM_CAP_STRUCTURED_RELEASE, + 0, + "structured release did not exist at minor 0" + ); + + let modern = open(&plugin, "lazy", NXMEM_CAP_ALLOCATOR); + assert_eq!( + modern.core().abi_minor(), + NXMEM_ABI_VERSION_MINOR, + "a current mechanism in the same module must not be dragged down" + ); + + // Both must work, through the same public release entry point. + for allocator in [&legacy, &modern] { + let ptr = allocator.allocate(2048, 32).expect("allocation"); + // SAFETY: live allocation with matching parameters. + let outcome = unsafe { allocator.release(ptr, 2048, 32) }; + assert!( + matches!(outcome, AllocationReleaseOutcome::Complete { .. }), + "both ABI levels must release cleanly, got {outcome:?}" + ); + } +} + +/// **Short struct.** +/// +/// The `short-struct` mechanism publishes an allocator vtable whose +/// `struct_size` is smaller than the baseline prefix. The host must refuse it +/// *before* reading any function pointer out of it. +#[test] +fn a_short_allocator_vtable_is_refused_before_any_slot_is_read() { + let _serial = serial(); + let plugin = load(); + let error = plugin + .factory("short-struct") + .expect("the mechanism is published") + .open(NXMEM_CAP_ALLOCATOR, None) + .expect_err("an undersized vtable must be refused"); + + let text = error.to_string(); + assert!( + text.contains("struct_size") || text.to_lowercase().contains("short") || text.contains("small"), + "the refusal must explain that the struct is too short, got: {text}" + ); +} + +/// **Missing optional capability.** +/// +/// `eager` advertises the allocator capability and nothing else. Asking for an +/// optional capability it does not have is refused, and opening it without +/// asking yields an allocator whose capability views are genuinely absent — +/// not present-but-broken. +#[test] +fn a_mechanism_without_optional_capabilities_reports_none() { + let _serial = serial(); + let plugin = load(); + + let refusal = plugin + .factory("eager") + .expect("the mechanism is published") + .open(NXMEM_CAP_VIRTUAL_BACKING, None) + .expect_err("a capability the mechanism lacks must be refused, not faked"); + let text = refusal.to_string(); + assert!( + text.to_lowercase().contains("capab"), + "the refusal must say which capability is missing, got: {text}" + ); + + let allocator = open(&plugin, "eager", NXMEM_CAP_ALLOCATOR); + assert!( + allocator.as_virtual_backing().is_none(), + "an unsupported capability must be represented as absent" + ); + assert!( + allocator.as_shared_mapping().is_none(), + "an unsupported capability must be represented as absent" + ); + assert!( + !allocator.commits_on_demand(), + "a mechanism with no virtual backing cannot commit on demand" + ); + assert_eq!( + allocator.pending_release_count().expect("a query is safe"), + 0, + "a mechanism with no deferred release has nothing pending" + ); + + // Deferred release is likewise refused rather than silently degraded to an + // immediate free. + let ptr = allocator.allocate(1024, 16).expect("allocation"); + // SAFETY: live allocation with matching parameters. + let error = unsafe { allocator.enqueue_release(ptr, 1024, 16) } + .expect_err("deferred release is not available on this mechanism"); + assert!( + matches!(error, MemoryError::AllocationFailed { .. }), + "got {error:?}" + ); + // SAFETY: the failed enqueue left the allocation live. + unsafe { allocator.deallocate(ptr, 1024, 16) }; +} + +// ─── allocation and release ───────────────────────────────────────────────── + +/// **Allocation and free.** +#[test] +fn allocate_and_release_round_trips_through_the_boundary() { + let _serial = serial(); + let plugin = load(); + let allocator = open(&plugin, "eager", NXMEM_CAP_ALLOCATOR); + + assert_eq!(allocator.core().live_allocation_count(), 0); + let ptr = allocator.allocate(8192, 256).expect("allocation succeeds"); + assert_eq!( + ptr.as_ptr() as usize % 256, + 0, + "the plugin must honour the requested alignment" + ); + assert_eq!(allocator.core().live_allocation_count(), 1); + + // The memory really is usable across the boundary. + // SAFETY: the plugin returned 8192 writable bytes at `ptr`. + unsafe { ptr.as_ptr().write_bytes(0xAB, 8192) }; + // SAFETY: as above. + assert_eq!(unsafe { ptr.as_ptr().add(8191).read() }, 0xAB); + + // SAFETY: live allocation with matching parameters. + let outcome = unsafe { allocator.release(ptr, 8192, 256) }; + match outcome { + AllocationReleaseOutcome::Complete { accounting } => { + assert_eq!(accounting.allocation_bytes, 8192); + assert_eq!(accounting.unmapped_bytes, 8192); + } + other => panic!("expected a complete release, got {other:?}"), + } + assert_eq!(allocator.core().live_allocation_count(), 0); +} + +/// An address the host does not know is refused rather than guessed at. +/// +/// This is the ABA-safety property: identity is a monotonic id the host +/// assigned, never the address, so a recycled address cannot be mistaken for a +/// live allocation. +#[test] +fn releasing_an_unknown_address_fails_rather_than_guessing() { + let _serial = serial(); + let plugin = load(); + let allocator = open(&plugin, "eager", NXMEM_CAP_ALLOCATOR); + + let mut stack = 0u64; + let ptr = NonNull::new((&raw mut stack).cast::()).expect("non-null"); + // SAFETY: deliberately passing an address this allocator never issued; + // the contract is that this is reported, not acted on. + let outcome = unsafe { allocator.release(ptr, 8, 8) }; + match outcome { + AllocationReleaseOutcome::Failed { failure } => { + let text = format!("{failure:?}"); + assert!( + text.contains("live allocation") || text.to_lowercase().contains("recognise"), + "the failure must say the address is unknown, got: {text}" + ); + } + other => panic!("an unknown address must fail, not {other:?}"), + } + assert_eq!(stack, 0, "the stack slot must be untouched"); +} + +/// A size or alignment that disagrees with the live record is refused. +#[test] +fn a_release_that_misdescribes_the_allocation_is_refused() { + let _serial = serial(); + let plugin = load(); + let allocator = open(&plugin, "eager", NXMEM_CAP_ALLOCATOR); + let ptr = allocator.allocate(4096, 64).expect("allocation"); + + // SAFETY: the address is live; the size is deliberately wrong. + let outcome = unsafe { allocator.release(ptr, 2048, 64) }; + assert!( + matches!(outcome, AllocationReleaseOutcome::Failed { .. }), + "a mismatched size must be refused, got {outcome:?}" + ); + assert_eq!( + allocator.core().live_allocation_count(), + 1, + "a refused release must leave the allocation live, not leak it" + ); + + // SAFETY: the allocation is still live with its original parameters. + let outcome = unsafe { allocator.release(ptr, 4096, 64) }; + assert!(matches!(outcome, AllocationReleaseOutcome::Complete { .. })); +} + +// ─── lazy backing ─────────────────────────────────────────────────────────── + +/// **Lazy backing.** +#[test] +fn lazy_backing_commits_decommits_and_reports_mapped_bytes() { + let _serial = serial(); + let plugin = load(); + let allocator = open( + &plugin, + "lazy", + NXMEM_CAP_ALLOCATOR | NXMEM_CAP_VIRTUAL_BACKING, + ); + assert!( + allocator.commits_on_demand(), + "a mechanism with virtual backing commits on demand" + ); + let backing = allocator + .as_virtual_backing() + .expect("virtual backing was required at open time"); + + // Reserve without committing. + let ptr = backing + .allocate_committed(64 * 1024, 4096, &[]) + .expect("a reservation with nothing committed"); + assert_eq!( + backing.allocation_committed_bytes(ptr, 64 * 1024, 4096), + 0, + "nothing was committed yet" + ); + + // Commit a window and watch the committed bytes follow. + backing + .commit_allocation_range(ptr, 64 * 1024, 4096, 0, 8192) + .expect("committing a range"); + assert_eq!( + backing.allocation_committed_bytes(ptr, 64 * 1024, 4096), + 8192 + ); + + let mapped = backing + .mapped_bytes_for_allocation_ranges(&[AllocationCommitRange { + ptr, + allocation_bytes: 64 * 1024, + align: 4096, + offset: 0, + bytes: 100, + }]) + .expect("a mapped-byte estimate"); + assert_eq!( + mapped, 4096, + "a partial granule still maps a whole granule; the estimate must be conservative" + ); + assert_eq!( + backing + .mapped_bytes_for_allocation(64 * 1024, 4096) + .expect("a whole-allocation estimate"), + 64 * 1024, + "a whole allocation maps its whole granule-rounded size" + ); + + let unmapped = backing + .decommit_allocation_range(ptr, 64 * 1024, 4096, 0, 8192) + .expect("decommitting a range"); + assert_eq!(unmapped, 8192); + assert_eq!(backing.allocation_committed_bytes(ptr, 64 * 1024, 4096), 0); + + // SAFETY: live allocation with matching parameters. + let outcome = unsafe { allocator.release(ptr, 64 * 1024, 4096) }; + assert!(matches!(outcome, AllocationReleaseOutcome::Complete { .. })); +} + +/// Shared prefixes are reference counted and keep their two accounting axes +/// apart: mapping a prefix again owns no new bytes but does map bytes. +#[test] +fn a_shared_prefix_is_reference_counted_and_costed_once() { + let _serial = serial(); + let plugin = load(); + let allocator = open( + &plugin, + "lazy", + NXMEM_CAP_ALLOCATOR | NXMEM_CAP_VIRTUAL_BACKING | NXMEM_CAP_SHARED_MAPPING, + ); + let shared = allocator + .as_shared_mapping() + .expect("shared mapping was required at open time"); + + let prefix = shared + .create_shared_prefix(16 * 1024) + .expect("creating a shared prefix"); + assert_eq!(prefix.requested_bytes(), 16 * 1024); + assert_ne!(prefix.device_ptr(), 0); + + assert_eq!( + shared + .incremental_owned_bytes_for_shared_prefix(prefix.as_ref()) + .expect("the prefix belongs to this mechanism"), + 0, + "the prefix's physical bytes were charged once at creation" + ); + + let ptr = allocator.allocate(64 * 1024, 4096).expect("allocation"); + let info = shared + .commit_shared_prefix(prefix.as_ref(), ptr, 64 * 1024, 0) + .expect("committing the prefix into the allocation"); + assert_eq!( + info.additional_owned_bytes, 0, + "re-mapping owns nothing new" + ); + assert_eq!( + info.newly_mapped_bytes, + 16 * 1024, + "but it does map bytes; the axes are distinct" + ); + + // SAFETY: live allocation with matching parameters. + let outcome = unsafe { allocator.release(ptr, 64 * 1024, 4096) }; + assert!(matches!(outcome, AllocationReleaseOutcome::Complete { .. })); +} + +// ─── deferred release ordering ────────────────────────────────────────────── + +/// **Release ordering**, plus module pinning across a deferred free. +#[test] +fn deferred_releases_retire_in_order_and_pin_the_module() { + let _serial = serial(); + let plugin = load(); + let allocator = open( + &plugin, + "lazy", + NXMEM_CAP_ALLOCATOR | NXMEM_CAP_DEFERRED_RELEASE, + ); + + let mut tickets = Vec::new(); + for index in 0..4u64 { + let bytes = 1024 * (index as usize + 1); + let ptr = allocator.allocate(bytes, 64).expect("allocation"); + // SAFETY: live allocation with matching parameters. + let ticket = unsafe { allocator.enqueue_release(ptr, bytes, 64) }.expect("queued release"); + tickets.push(ticket); + } + + assert_eq!( + allocator.pending_release_count().expect("query"), + 4, + "all four releases are still queued" + ); + assert_eq!( + plugin.module().host_live_counts().queued_releases, + 4, + "the host counts them too, so unload stays gated" + ); + assert!( + allocator.core().retired_releases().is_empty(), + "no completion may arrive before the drain" + ); + + // Retire two, then the rest, so partial drains are covered as well. + assert_eq!(allocator.drain_releases(2).expect("drain"), 2); + assert_eq!(allocator.pending_release_count().expect("query"), 2); + assert_eq!(allocator.drain_releases(64).expect("drain"), 2); + assert_eq!(allocator.pending_release_count().expect("query"), 0); + + let retired = allocator.core().retired_releases(); + assert_eq!(retired.len(), 4, "every queued release must report back"); + let retired_tickets: Vec = retired.iter().map(|entry| entry.ticket).collect(); + assert_eq!( + retired_tickets, tickets, + "completions must arrive in the order the releases were queued" + ); + for (index, entry) in retired.iter().enumerate() { + assert_eq!(entry.allocation_bytes, 1024 * (index as u64 + 1)); + } + assert_eq!( + plugin.module().host_live_counts().queued_releases, + 0, + "the host's tally must return to zero once the queue drains" + ); +} + +// ─── callback failure ─────────────────────────────────────────────────────── + +/// **Callback failure.** +/// +/// The `callback-probe` mechanism calls the host's reclaim hook on every +/// allocation. A refusing host must produce a clean allocation failure — no +/// abort, no leak, no half-registered allocation. +#[test] +fn a_refusing_host_callback_fails_the_allocation_cleanly() { + let _serial = serial(); + let plugin = load(); + let reclaim = ScriptedReclaim::refusing(); + let allocator = plugin + .factory("callback-probe") + .expect("the mechanism is published") + .open(NXMEM_CAP_ALLOCATOR, Some(reclaim.clone())) + .expect("the mechanism opens"); + + let error = allocator + .allocate(4096, 64) + .expect_err("a refused reclaim must fail the allocation"); + assert!( + matches!(error, MemoryError::AllocationFailed { .. }), + "got {error:?}" + ); + let text = format!("{error:?}"); + assert!( + text.to_lowercase().contains("callback") || text.contains("reclaim"), + "the error must name the failing callback, got: {text}" + ); + + assert_eq!( + reclaim.calls.load(Ordering::Acquire), + 1, + "the plugin really did reach the host" + ); + assert_eq!(reclaim.last_bytes.load(Ordering::Acquire), 4096); + assert_eq!( + allocator.core().reclaim_calls(), + 1, + "the host bridge counts the call" + ); + assert_eq!( + allocator.core().reclaim_failures(), + 1, + "and records that it refused" + ); + assert_eq!( + allocator.core().live_allocation_count(), + 0, + "a failed allocation must not be registered" + ); + + // The very same mechanism succeeds once the host cooperates, which proves + // the failure came from the callback and not from something else. + let allocator = plugin + .factory("callback-probe") + .expect("the mechanism is published") + .open(NXMEM_CAP_ALLOCATOR, Some(ScriptedReclaim::granting(1 << 20))) + .expect("the mechanism opens"); + let ptr = allocator + .allocate(4096, 64) + .expect("a cooperating host lets the allocation through"); + assert_eq!(allocator.core().reclaim_calls(), 1); + assert_eq!(allocator.core().reclaim_failures(), 0); + // SAFETY: live allocation with matching parameters. + unsafe { allocator.deallocate(ptr, 4096, 64) }; +} + +/// A host that offers no reclaim hook at all says so explicitly. +/// +/// The alternative — answering "reclaimed 0 bytes" — would be a lie a plugin +/// cannot distinguish from a real, unhelpful reclaim. An unsupported +/// capability is represented as unsupported, and the plugin decides what to do +/// about it; `callback-probe` chooses to fail the allocation. +#[test] +fn a_host_with_no_reclaim_hook_reports_the_capability_as_absent() { + let _serial = serial(); + let plugin = load(); + let allocator = open(&plugin, "callback-probe", NXMEM_CAP_ALLOCATOR); + let error = allocator + .allocate(4096, 64) + .expect_err("this mechanism cannot proceed without a reclaim path"); + let text = format!("{error:?}"); + assert!( + text.contains("no reclaim path"), + "the host must say plainly that it offers no reclaim path, got: {text}" + ); + assert_eq!( + allocator.core().reclaim_calls(), + 1, + "the plugin really did try" + ); + assert_eq!( + allocator.core().reclaim_failures(), + 1, + "and the bridge recorded that the host could not help" + ); + assert_eq!( + allocator.core().live_allocation_count(), + 0, + "a failed allocation must not be registered" + ); +} + +// ─── unload gating ────────────────────────────────────────────────────────── + +/// A plugin with nothing live unloads. +#[test] +fn an_idle_plugin_unloads() { + let _serial = serial(); + let plugin = load(); + assert_eq!(plugin.module().host_live_counts().total(), 0); + plugin + .try_unload() + .expect("a plugin with nothing live must unload"); +} + +/// **Unload with live objects**, allocator case. +#[test] +fn unload_is_refused_while_an_allocator_is_open() { + let _serial = serial(); + let plugin = load(); + let allocator = open(&plugin, "eager", NXMEM_CAP_ALLOCATOR); + + let rejection = plugin + .try_unload() + .expect_err("an open allocator must block unload"); + assert!( + rejection.report.live_allocators >= 1, + "the plugin itself must report the live allocator: {:?}", + rejection.report + ); + let plugin = rejection + .into_plugin() + .expect("the refusal hands the plugin back so the caller can retire work"); + + drop(allocator); + plugin + .try_unload() + .expect("once the allocator is gone the plugin unloads"); +} + +/// **Unload with live objects**, allocation case. +#[test] +fn unload_is_refused_while_an_allocation_is_live() { + let _serial = serial(); + let plugin = load(); + let allocator = open(&plugin, "eager", NXMEM_CAP_ALLOCATOR); + let ptr = allocator.allocate(4096, 64).expect("allocation"); + + let rejection = plugin + .try_unload() + .expect_err("a live allocation must block unload"); + assert!( + rejection.report.live_allocations >= 1, + "the plugin must report the live allocation: {:?}", + rejection.report + ); + let plugin = rejection.into_plugin().expect("the plugin comes back"); + + // SAFETY: live allocation with matching parameters. + unsafe { allocator.deallocate(ptr, 4096, 64) }; + drop(allocator); + plugin.try_unload().expect("now it unloads"); +} + +/// **Unload with live objects**, capability-view case. +#[test] +fn unload_is_refused_while_a_shared_prefix_is_held() { + let _serial = serial(); + let plugin = load(); + let allocator = open( + &plugin, + "lazy", + NXMEM_CAP_ALLOCATOR | NXMEM_CAP_SHARED_MAPPING, + ); + let prefix = allocator + .as_shared_mapping() + .expect("shared mapping was required") + .create_shared_prefix(8192) + .expect("creating a prefix"); + + let rejection = plugin + .try_unload() + .expect_err("a held capability object must block unload"); + assert!( + rejection.report.live_capabilities >= 1, + "the plugin must report the held prefix: {:?}", + rejection.report + ); + let plugin = rejection.into_plugin().expect("the plugin comes back"); + + drop(prefix); + drop(allocator); + plugin.try_unload().expect("now it unloads"); +} + +// ─── cross-provider misuse ────────────────────────────────────────────────── + +/// Two allocators from the same module have distinct mechanism identities, and +/// an object from one is refused by the other rather than silently accepted. +#[test] +fn an_object_from_another_mechanism_is_refused() { + let _serial = serial(); + let plugin = load(); + let first = open( + &plugin, + "lazy", + NXMEM_CAP_ALLOCATOR | NXMEM_CAP_SHARED_MAPPING, + ); + let second = open( + &plugin, + "lazy", + NXMEM_CAP_ALLOCATOR | NXMEM_CAP_SHARED_MAPPING, + ); + assert_ne!( + first.core().mechanism_id(), + second.core().mechanism_id(), + "each opened mechanism instance needs its own identity" + ); + + let prefix = first + .as_shared_mapping() + .expect("shared mapping") + .create_shared_prefix(4096) + .expect("creating a prefix on the first mechanism"); + + // The second mechanism must not cost a foreign prefix as free. + let error = second + .as_shared_mapping() + .expect("shared mapping") + .incremental_owned_bytes_for_shared_prefix(prefix.as_ref()) + .expect_err("a foreign prefix must be rejected, never costed as free"); + let text = format!("{error:?}"); + assert!( + text.to_lowercase().contains("mechanism") || text.to_lowercase().contains("prefix"), + "the refusal must name the mismatch, got: {text}" + ); + + let ptr = second.allocate(8192, 4096).expect("allocation"); + let error = second + .as_shared_mapping() + .expect("shared mapping") + .commit_shared_prefix(prefix.as_ref(), ptr, 8192, 0) + .expect_err("committing a foreign prefix must be refused"); + let text = format!("{error:?}"); + assert!( + text.to_lowercase().contains("mechanism") || text.to_lowercase().contains("prefix"), + "the refusal must name the mismatch, got: {text}" + ); + + // SAFETY: live allocation with matching parameters. + let outcome = unsafe { second.release(ptr, 8192, 4096) }; + assert!(matches!(outcome, AllocationReleaseOutcome::Complete { .. })); +} + +/// An allocation made by one mechanism is not releasable by another. +#[test] +fn an_allocation_cannot_be_released_by_a_sibling_mechanism() { + let _serial = serial(); + let plugin = load(); + let first = open(&plugin, "eager", NXMEM_CAP_ALLOCATOR); + let second = open(&plugin, "eager", NXMEM_CAP_ALLOCATOR); + + let ptr = first.allocate(4096, 64).expect("allocation"); + // SAFETY: deliberately naming a sibling's allocation; the contract is that + // this is refused, not acted on. + let outcome = unsafe { second.release(ptr, 4096, 64) }; + assert!( + matches!(outcome, AllocationReleaseOutcome::Failed { .. }), + "a sibling must not be able to free another mechanism's memory, got {outcome:?}" + ); + assert_eq!(first.core().live_allocation_count(), 1); + + // SAFETY: still live on its owning mechanism. + let outcome = unsafe { first.release(ptr, 4096, 64) }; + assert!(matches!(outcome, AllocationReleaseOutcome::Complete { .. })); +} + +// ─── factory lifetime ─────────────────────────────────────────────────────── + +/// Every factory the host took is released exactly once when the plugin goes. +/// +/// The count is read out of the **loaded module** through a test-only exported +/// symbol. Reading the statically linked `rlib` copy of the same static would +/// observe a different variable entirely and would pass while proving nothing. +/// +/// Two handles are opened onto the same module: unloading the second must +/// release its factories while the first keeps the module mapped, which is +/// also what makes the count readable afterwards. +#[test] +fn every_factory_is_released_exactly_once() { + let _serial = serial(); + let observer = load(); + let releases = |plugin: &MemoryPlugin| -> u64 { + // SAFETY: the symbol is a `extern "C" fn() -> u64` exported by the + // test plugin, and the module stays mapped for the borrow. + let symbol: libloading::Symbol<'_, unsafe extern "C" fn() -> u64> = unsafe { + plugin + .module() + .library() + .get(onnx_runtime_memory_testplugin::SYMBOL_FACTORY_RELEASES) + } + .expect("the test plugin exports its introspection symbol"); + // SAFETY: as above. + unsafe { symbol() } + }; + + let before = releases(&observer); + let plugin = load(); + let count = plugin.factories().len() as u64; + assert_eq!(count, 6, "the test plugin publishes six mechanisms"); + plugin.try_unload().expect("an idle plugin unloads"); + + assert_eq!( + releases(&observer) - before, + count, + "the host must release each factory exactly once" + ); + + // The observing handle still works, which is what "the module stays mapped + // while anything pins it" means in practice. + let allocator = open(&observer, "eager", NXMEM_CAP_ALLOCATOR); + let ptr = allocator.allocate(1024, 16).expect("allocation"); + // SAFETY: live allocation with matching parameters. + unsafe { allocator.deallocate(ptr, 1024, 16) }; +} + +/// An unknown mechanism name is a clear, listing error rather than a panic. +#[test] +fn an_unknown_mechanism_name_is_reported_with_the_available_set() { + let _serial = serial(); + let plugin = load(); + let error = plugin + .factory("no-such-mechanism") + .expect_err("the mechanism does not exist"); + match error { + PluginError::UnknownMechanism { available, .. } => { + assert!( + available.contains("eager") && available.contains("lazy"), + "the error must list what is available, got: {available}" + ); + } + other => panic!("expected an unknown-mechanism error, got {other}"), + } +} diff --git a/crates/onnx-runtime-memory-host/tests/nxmem_abi_unload_gate.rs b/crates/onnx-runtime-memory-host/tests/nxmem_abi_unload_gate.rs new file mode 100644 index 0000000000..50005e22b6 --- /dev/null +++ b/crates/onnx-runtime-memory-host/tests/nxmem_abi_unload_gate.rs @@ -0,0 +1,66 @@ +//! The unload gate's terminal case, in a process of its own. +//! +//! A plugin module's live-object counters are process-wide, because unloading +//! a module is a process-wide act. The `sticky` mechanism deliberately never +//! retires a queued release, so once this test runs the module can never +//! honestly report itself unloadable again. That is the behaviour under test — +//! and it is exactly why this case lives in its own integration-test binary +//! rather than sharing a process with the rest of the suite. + +use onnx_runtime_memory_abi::{NXMEM_CAP_ALLOCATOR, NXMEM_CAP_DEFERRED_RELEASE}; +use onnx_runtime_memory_api::DeviceAllocator; +use onnx_runtime_memory_host::MemoryPlugin; + +#[path = "support/testplugin.rs"] +mod testplugin; + +/// **Unload with live objects**, queued-free case. +/// +/// The host must not unmap a module that still owes a free: the module's code +/// is what will run when the free finally retires. +#[test] +fn unload_is_refused_while_a_queued_release_has_not_retired() { + let plugin = MemoryPlugin::load(testplugin::testplugin_path()).expect("the test plugin loads"); + let allocator = plugin + .factory("sticky") + .expect("the mechanism is published") + .open(NXMEM_CAP_ALLOCATOR | NXMEM_CAP_DEFERRED_RELEASE, None) + .expect("the mechanism opens"); + + let ptr = allocator.allocate(4096, 64).expect("allocation"); + // SAFETY: live allocation with matching parameters. + let _ticket = unsafe { allocator.enqueue_release(ptr, 4096, 64) }.expect("queued release"); + + assert_eq!( + allocator.drain_releases(64).expect("a drain is attempted"), + 0, + "this mechanism deliberately retires nothing" + ); + assert_eq!( + allocator.pending_release_count().expect("query"), + 1, + "the plugin still owes the free" + ); + + let rejection = plugin + .try_unload() + .expect_err("a queued release must block unload"); + assert!( + rejection.host.queued_releases >= 1 || rejection.report.queued_releases >= 1, + "either side may notice first, but somebody must: host={:?} plugin={:?}", + rejection.host, + rejection.report + ); + let plugin = rejection + .into_plugin() + .expect("the refusal hands the plugin back so the caller can retire work"); + + // Dropping the allocator lets the plugin reclaim its own storage, but the + // host's queued tally is what keeps the gate shut, and it never clears. + drop(allocator); + let outcome = plugin.try_unload(); + assert!( + outcome.is_err(), + "the host must not unmap a module that still owes a free" + ); +} diff --git a/crates/onnx-runtime-memory-host/tests/support/testplugin.rs b/crates/onnx-runtime-memory-host/tests/support/testplugin.rs new file mode 100644 index 0000000000..45b473551b --- /dev/null +++ b/crates/onnx-runtime-memory-host/tests/support/testplugin.rs @@ -0,0 +1,80 @@ +//! Shared helper: locate (and build) the nxmem test plugin cdylib. +//! +//! Included with `#[path]` by each integration-test binary. Some binaries use +//! only `testplugin_path`, so the helpers are allowed to be unused. +#![allow(dead_code)] + +use std::path::PathBuf; + +/// Locate the test plugin cdylib, building it if it is not there yet. +/// +/// Resolution order mirrors the existing nxrt ABI tests: an explicit override, +/// then `CARGO_TARGET_DIR`, then the workspace default layout. If the library +/// genuinely cannot be produced the helper **panics loudly** — it never lets a +/// test pass by quietly doing nothing. +pub fn testplugin_path() -> PathBuf { + if let Ok(explicit) = std::env::var("NXMEM_TESTPLUGIN_PATH") { + let path = PathBuf::from(explicit); + assert!( + path.exists(), + "NXMEM_TESTPLUGIN_PATH names {path:?}, which does not exist" + ); + return path; + } + + static PATH: std::sync::OnceLock = std::sync::OnceLock::new(); + PATH.get_or_init(build_testplugin).clone() +} + +/// Build the cdylib and return where it landed. +/// +/// The build is unconditional. Cargo builds only the `rlib` target of a +/// dev-dependency, so the `cdylib` on disk can easily be stale — and a test +/// suite silently exercising a stale artifact is worse than one that takes an +/// extra second. If the build cannot be done, this panics loudly rather than +/// letting anything pass by default. +pub fn build_testplugin() -> PathBuf { + let profile = std::env::var("PROFILE").unwrap_or_else(|_| String::from("debug")); + let libname = if cfg!(target_os = "linux") { + "libonnx_runtime_memory_testplugin.so" + } else if cfg!(target_os = "macos") { + "libonnx_runtime_memory_testplugin.dylib" + } else { + "onnx_runtime_memory_testplugin.dll" + }; + + let mut command = + std::process::Command::new(std::env::var("CARGO").unwrap_or_else(|_| String::from("cargo"))); + command.args(["build", "-p", "onnx-runtime-memory-testplugin"]); + if profile != "debug" { + command.args(["--profile", &profile]); + } + let status = command + .status() + .expect("cargo must be runnable to build the nxmem test plugin"); + assert!(status.success(), "building the nxmem test plugin failed"); + + let mut candidates = Vec::new(); + if let Ok(target_dir) = std::env::var("CARGO_TARGET_DIR") { + candidates.push(PathBuf::from(target_dir).join(&profile).join(libname)); + } + candidates.push( + PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .parent() + .and_then(|crates| crates.parent()) + .expect("the crate lives two levels below the workspace root") + .join("target") + .join(&profile) + .join(libname), + ); + for candidate in &candidates { + if candidate.exists() { + return candidate.clone(); + } + } + panic!( + "the nxmem test plugin cdylib is missing after a successful build; looked in \ + {candidates:?}. Set NXMEM_TESTPLUGIN_PATH when using a custom target directory." + ); +} + diff --git a/crates/onnx-runtime-memory-testplugin/src/lib.rs b/crates/onnx-runtime-memory-testplugin/src/lib.rs index e69de29bb2..002539e350 100644 --- a/crates/onnx-runtime-memory-testplugin/src/lib.rs +++ b/crates/onnx-runtime-memory-testplugin/src/lib.rs @@ -0,0 +1,1672 @@ +//! An out-of-tree-style nxmem memory plugin, used to exercise the ABI. +//! +//! This crate deliberately depends on **only** `onnx-runtime-memory-abi`. It +//! never sees a runtime trait, a governor type, or an internal enum, so it is +//! a faithful stand-in for a plugin shipped by a third party. It is built as a +//! `cdylib` and loaded at runtime with `dlopen`, which is what makes the ABI +//! test suite portable: it exercises a genuine dynamic-library boundary on +//! macOS, Linux, and Windows without needing a GPU. +//! +//! # The mechanisms it publishes +//! +//! | Name | Purpose | +//! |---|---| +//! | `eager` | required slots only, no optional capability — the minimal conforming plugin | +//! | `lazy` | virtual backing, shared mapping, deferred release, structured release | +//! | `short-struct` | publishes a deliberately undersized allocator vtable | +//! | `callback-probe` | calls the host's `request_reclaim` during `allocate` and fails when the host refuses | +//! | `legacy-1-0` | pins itself to the minor-0 prefix, so an older participant meets a newer host | +//! | `sticky` | never retires its queued releases, so unload stays refused | +//! +//! Publishing several *named* mechanisms from one module — rather than +//! switching behaviour on an environment variable or a global — keeps every +//! scenario independent and lets one test process cover them all. +//! +//! Backing storage is ordinary host memory from `std::alloc`, allocated **and +//! freed inside this module**. No allocator ownership crosses the boundary. + +use std::alloc::{Layout, alloc, dealloc}; +use std::collections::HashMap; +use std::ffi::c_void; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; +use std::sync::{Mutex, OnceLock}; + +use onnx_runtime_memory_abi::{ + NXMEM_ABI_VERSION_MAJOR, NXMEM_ABI_VERSION_MINOR, NXMEM_CAP_ALLOCATOR, + NXMEM_CAP_DEFERRED_RELEASE, NXMEM_CAP_SHARED_MAPPING, NXMEM_CAP_STRUCTURED_RELEASE, + NXMEM_CAP_VIRTUAL_BACKING, NxmemAllocRequest, NxmemAllocResult, NxmemAllocation, + NxmemAllocatorFactoryVtable, NxmemAllocatorVtable, NxmemDeviceId, NxmemHostCallbacks, + NxmemNegotiateRequest, NxmemNegotiateResponse, NxmemOpenRequest, NxmemRangeRequest, + NxmemReclaimRequest, NxmemReleaseCompletion, NxmemReleaseOutcome, NxmemSharedMappingVtable, + NxmemSharedPrefixCommitInfo, NxmemSharedPrefixCommitRequest, NxmemSharedPrefixHandle, + NxmemStatus, NxmemStatusCode, NxmemUnloadReport, NxmemVersionRange, NxmemVirtualBackingVtable, + catch_status_panic, catch_void_panic, check_identity, negotiate_as, +}; + +// ─── module-wide bookkeeping ──────────────────────────────────────────────── + +/// Live objects this module still owns, reported through the unload gate. +#[derive(Debug, Default)] +struct ModuleCounters { + live_allocators: AtomicU64, + live_allocations: AtomicU64, + live_views: AtomicU64, + live_capabilities: AtomicU64, + queued_releases: AtomicU64, +} + +fn counters() -> &'static ModuleCounters { + static COUNTERS: OnceLock = OnceLock::new(); + COUNTERS.get_or_init(ModuleCounters::default) +} + +fn next_mechanism_id() -> u64 { + static NEXT: AtomicU64 = AtomicU64::new(0); + NEXT.fetch_add(1, Ordering::AcqRel) + 1 +} + +fn next_ticket() -> u64 { + static NEXT: AtomicU64 = AtomicU64::new(0); + NEXT.fetch_add(1, Ordering::AcqRel) + 1 +} + +fn next_prefix_handle() -> u64 { + static NEXT: AtomicU64 = AtomicU64::new(0); + NEXT.fetch_add(1, Ordering::AcqRel) + 1 +} + +// ─── mechanism behaviour ──────────────────────────────────────────────────── + +/// Which behaviours a named mechanism exhibits. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +struct Behaviour { + /// Capabilities the mechanism claims. + capability_flags: u64, + /// The minor level its vtables declare. + abi_minor: u32, + /// Publish an allocator vtable that is deliberately too short. + short_allocator_struct: bool, + /// Call the host's `request_reclaim` on every allocation, and fail the + /// allocation when the host refuses. + probe_host_callback: bool, + /// Never retire a queued release, so unload stays refused. + never_drain: bool, +} + +impl Behaviour { + const fn base() -> Self { + Self { + capability_flags: NXMEM_CAP_ALLOCATOR | NXMEM_CAP_STRUCTURED_RELEASE, + abi_minor: 1, + short_allocator_struct: false, + probe_host_callback: false, + never_drain: false, + } + } +} + +/// One published mechanism. +struct Mechanism { + name: &'static str, + c_name: &'static std::ffi::CStr, + behaviour: Behaviour, +} + +const MECHANISMS: &[Mechanism] = &[ + Mechanism { + name: "eager", + c_name: c"eager", + behaviour: Behaviour { + // The minimal conforming mechanism: required slots only, no + // optional capability at all. + capability_flags: NXMEM_CAP_ALLOCATOR, + ..Behaviour::base() + }, + }, + Mechanism { + name: "lazy", + c_name: c"lazy", + behaviour: Behaviour { + capability_flags: NXMEM_CAP_ALLOCATOR + | NXMEM_CAP_VIRTUAL_BACKING + | NXMEM_CAP_SHARED_MAPPING + | NXMEM_CAP_DEFERRED_RELEASE + | NXMEM_CAP_STRUCTURED_RELEASE, + ..Behaviour::base() + }, + }, + Mechanism { + name: "short-struct", + c_name: c"short-struct", + behaviour: Behaviour { + short_allocator_struct: true, + ..Behaviour::base() + }, + }, + Mechanism { + name: "callback-probe", + c_name: c"callback-probe", + behaviour: Behaviour { + probe_host_callback: true, + ..Behaviour::base() + }, + }, + Mechanism { + name: "legacy-1-0", + c_name: c"legacy-1-0", + behaviour: Behaviour { + // A mechanism written against the minor-0 contract: it knows + // nothing about structured release. + capability_flags: NXMEM_CAP_ALLOCATOR | NXMEM_CAP_DEFERRED_RELEASE, + abi_minor: 0, + ..Behaviour::base() + }, + }, + Mechanism { + name: "sticky", + c_name: c"sticky", + behaviour: Behaviour { + capability_flags: NXMEM_CAP_ALLOCATOR + | NXMEM_CAP_DEFERRED_RELEASE + | NXMEM_CAP_STRUCTURED_RELEASE, + never_drain: true, + ..Behaviour::base() + }, + }, +]; + +/// Mechanism names this plugin publishes, in factory order. +/// +/// Exposed through the `rlib` half so the host's tests can name mechanisms +/// without duplicating the list. +pub const MECHANISM_NAMES: &[&str] = &[ + "eager", + "lazy", + "short-struct", + "callback-probe", + "legacy-1-0", + "sticky", +]; + +// ─── allocator state ──────────────────────────────────────────────────────── + +#[derive(Debug, Clone, Copy)] +struct Block { + address: usize, + bytes: usize, + align: usize, + committed: usize, +} + +#[derive(Debug, Clone, Copy)] +struct QueuedRelease { + ticket: u64, + allocation_id: u64, + address: usize, + bytes: usize, + align: usize, +} + +#[derive(Debug, Clone, Copy)] +struct Prefix { + bytes: usize, + address: usize, + references: u64, +} + +/// The plugin-side state of one opened mechanism. +/// +/// Reference counted by `retain`/`release`. The state is destroyed only when +/// the count reaches zero **and** no queued release still names it, which is +/// what keeps a deferred free from touching a freed context. +struct AllocatorState { + mechanism_id: u64, + device: NxmemDeviceId, + behaviour: Behaviour, + refcount: AtomicU64, + callbacks: *const NxmemHostCallbacks, + inner: Mutex, + /// The vtables handed to the host. Boxed inside the state so their + /// addresses stay stable, and freed only when the state is destroyed. + allocator_vtable: Box, + backing_vtable: Option>, + shared_vtable: Option>, +} + +// SAFETY: every field is either immutable after construction or guarded by +// `inner`'s mutex / an atomic. `callbacks` points at host memory the host +// promises to keep alive past this allocator's final release and past every +// queued release that names it. +unsafe impl Send for AllocatorState {} +// SAFETY: as above. +unsafe impl Sync for AllocatorState {} + +#[derive(Debug, Default)] +struct AllocatorInner { + blocks: HashMap, + queue: Vec, + prefixes: HashMap, +} + +impl AllocatorState { + fn check(&self, allocation: &NxmemAllocation) -> Result<(), NxmemStatus> { + check_identity( + self.mechanism_id, + self.device, + allocation.mechanism_id, + allocation.device, + ) + } + + /// Ask the host to make room. + /// + /// Called from inside `allocate` with **no plugin lock held**, because a + /// host callback may re-enter this module. + fn host_reclaim(&self, bytes: u64) -> Result { + if self.callbacks.is_null() { + return Ok(0); + } + // SAFETY: the host promises the table outlives this allocator. + let callbacks = unsafe { &*self.callbacks }; + let Some(request_reclaim) = callbacks.request_reclaim else { + return Ok(0); + }; + let request = NxmemReclaimRequest::new(self.mechanism_id, self.device, bytes); + let mut reclaimed = 0u64; + // SAFETY: both pointers address valid locals that outlive the call. + let status = + unsafe { request_reclaim(callbacks.host_ctx, &raw const request, &raw mut reclaimed) }; + if status.is_ok() { + Ok(reclaimed) + } else { + Err(status) + } + } + + /// Report one retired deferred release. Called with no plugin lock held. + fn report_completion(&self, completion: &NxmemReleaseCompletion) -> Result<(), NxmemStatus> { + if self.callbacks.is_null() { + return Ok(()); + } + // SAFETY: as above. + let callbacks = unsafe { &*self.callbacks }; + let Some(release_completed) = callbacks.release_completed else { + return Ok(()); + }; + // SAFETY: `completion` is a valid local that outlives the call. + let status = unsafe { release_completed(callbacks.host_ctx, completion as *const _) }; + if status.is_ok() { + Ok(()) + } else { + Err(status) + } + } +} + +fn layout_of(bytes: usize, align: usize) -> Result { + Layout::from_size_align(bytes.max(1), align.max(1)).map_err(|_| { + NxmemStatus::with_message( + NxmemStatusCode::InvalidArgument, + "testplugin: the requested size and alignment are not a valid layout", + ) + }) +} + +/// Free one block's backing storage. +/// +/// The bytes were allocated by this module and are freed by this module, so no +/// allocator ownership crosses the boundary. +fn free_block(block: Block) { + let Ok(layout) = Layout::from_size_align(block.bytes.max(1), block.align.max(1)) else { + return; + }; + if block.address == 0 { + return; + } + // SAFETY: the address came from `alloc` with exactly this layout inside + // this module, and every caller removes the record first so it is freed + // exactly once. + unsafe { dealloc(block.address as *mut u8, layout) }; +} + +/// # Safety +/// +/// `ctx` must be a live `AllocatorState` pointer produced by `open_allocator`. +unsafe fn state<'a>(ctx: *mut c_void) -> Option<&'a AllocatorState> { + if ctx.is_null() { + return None; + } + // SAFETY: delegated to this function's contract. + Some(unsafe { &*(ctx as *const AllocatorState) }) +} + +macro_rules! plugin_entry { + ($ctx:expr) => { + // SAFETY: the host only passes back a `ctx` this module published. + match unsafe { state($ctx) } { + Some(state) => state, + None => { + return NxmemStatus::with_message( + NxmemStatusCode::InvalidArgument, + "testplugin: null allocator context", + ); + } + } + }; +} + +// ─── allocator vtable slots ───────────────────────────────────────────────── + +unsafe extern "C" fn plugin_allocate( + ctx: *mut c_void, + request: *const NxmemAllocRequest, + result_out: *mut NxmemAllocResult, +) -> NxmemStatus { + // SAFETY: same contract; the eager path commits the whole allocation. + unsafe { allocate_inner(ctx, request, result_out, true) } +} + +/// The shared body of the eager and lazy allocation slots. +/// +/// `commit_whole` is the only difference: the eager slot commits everything it +/// reserves, while `allocate_committed` commits exactly the ranges it was +/// handed — including none at all, which is the whole point of lazy backing. +/// +/// # Safety +/// +/// Same contract as the two slots that call it. +unsafe fn allocate_inner( + ctx: *mut c_void, + request: *const NxmemAllocRequest, + result_out: *mut NxmemAllocResult, + commit_whole: bool, +) -> NxmemStatus { + catch_status_panic(|| { + let state = plugin_entry!(ctx); + if request.is_null() || result_out.is_null() { + return NxmemStatus::with_message( + NxmemStatusCode::InvalidArgument, + "testplugin: allocate received a null pointer", + ); + } + // SAFETY: checked non-null; the host passes valid records. + let request = unsafe { &*request }; + if let Err(status) = check_identity( + state.mechanism_id, + state.device, + request.mechanism_id, + request.device, + ) { + return status; + } + + if state.behaviour.probe_host_callback { + // A refusal is a normal outcome and is surfaced as a failed + // allocation, never as an abort or a leak. + if let Err(status) = state.host_reclaim(request.bytes) { + return NxmemStatus::with_message( + NxmemStatusCode::CallbackFailed, + &format!( + "testplugin: the host refused to reclaim memory, so the allocation \ + cannot proceed: {}", + status.describe() + ), + ); + } + } + + let bytes = request.bytes as usize; + let align = request.align as usize; + let layout = match layout_of(bytes, align) { + Ok(layout) => layout, + Err(status) => return status, + }; + // SAFETY: `layout` has a non-zero size and a valid power-of-two align. + let raw = unsafe { alloc(layout) }; + if raw.is_null() { + return NxmemStatus::with_message( + NxmemStatusCode::OutOfMemory, + "testplugin: the backing host allocator refused the request", + ); + } + + let committed = if commit_whole { + bytes + } else if request.committed_range_count == 0 || request.committed_ranges.is_null() { + // A lazy reservation with no ranges commits nothing. + 0 + } else { + // SAFETY: the host promises `committed_range_count` readable + // ranges at `committed_ranges` for the duration of the call. + let ranges = unsafe { + core::slice::from_raw_parts( + request.committed_ranges, + request.committed_range_count as usize, + ) + }; + ranges + .iter() + .map(|range| range.bytes as usize) + .sum::() + }; + + let block = Block { + address: raw as usize, + bytes, + align, + committed: committed.min(bytes), + }; + match state.inner.lock() { + Ok(mut inner) => { + inner.blocks.insert(request.allocation_id, block); + } + Err(_) => { + // SAFETY: allocated just above with exactly this layout and + // never handed out. + unsafe { dealloc(raw, layout) }; + return NxmemStatus::with_message( + NxmemStatusCode::InternalError, + "testplugin: allocator state is poisoned", + ); + } + } + counters().live_allocations.fetch_add(1, Ordering::AcqRel); + + let mut result = NxmemAllocResult::zeroed(); + result.ptr = raw; + result.owned_bytes = block.committed as u64; + result.mapped_bytes = block.committed as u64; + // SAFETY: checked non-null above. + unsafe { *result_out = result }; + NxmemStatus::ok() + }) +} + +unsafe extern "C" fn plugin_deallocate( + ctx: *mut c_void, + allocation: *const NxmemAllocation, + unmapped_out: *mut u64, +) -> NxmemStatus { + catch_status_panic(|| { + let state = plugin_entry!(ctx); + if allocation.is_null() || unmapped_out.is_null() { + return NxmemStatus::with_message( + NxmemStatusCode::InvalidArgument, + "testplugin: deallocate received a null pointer", + ); + } + // SAFETY: checked non-null. + let allocation = unsafe { &*allocation }; + if let Err(status) = state.check(allocation) { + return status; + } + let block = match state.inner.lock() { + Ok(mut inner) => inner.blocks.remove(&allocation.allocation_id), + Err(_) => { + return NxmemStatus::with_message( + NxmemStatusCode::InternalError, + "testplugin: allocator state is poisoned", + ); + } + }; + let Some(block) = block else { + return NxmemStatus::with_message( + NxmemStatusCode::UnknownAllocation, + "testplugin: that allocation id is not live on this mechanism", + ); + }; + let unmapped = block.committed as u64; + free_block(block); + counters().live_allocations.fetch_sub(1, Ordering::AcqRel); + // SAFETY: checked non-null above. + unsafe { *unmapped_out = unmapped }; + NxmemStatus::ok() + }) +} + +unsafe extern "C" fn plugin_release_allocation( + ctx: *mut c_void, + allocation: *const NxmemAllocation, + outcome_out: *mut NxmemReleaseOutcome, +) -> NxmemStatus { + catch_status_panic(|| { + let state = plugin_entry!(ctx); + if allocation.is_null() || outcome_out.is_null() { + return NxmemStatus::with_message( + NxmemStatusCode::InvalidArgument, + "testplugin: release received a null pointer", + ); + } + // SAFETY: checked non-null. + let allocation = unsafe { &*allocation }; + if let Err(status) = state.check(allocation) { + return status; + } + let block = match state.inner.lock() { + Ok(mut inner) => inner.blocks.remove(&allocation.allocation_id), + Err(_) => { + return NxmemStatus::with_message( + NxmemStatusCode::InternalError, + "testplugin: allocator state is poisoned", + ); + } + }; + let Some(block) = block else { + // Nothing was mutated, so this is a `Failed` outcome, not a + // quarantine. The distinction is load-bearing for host accounting. + let outcome = NxmemReleaseOutcome::failed( + allocation.bytes, + NxmemStatus::with_message( + NxmemStatusCode::UnknownAllocation, + "testplugin: that allocation id is not live on this mechanism", + ), + ); + // SAFETY: checked non-null above. + unsafe { *outcome_out = outcome }; + return NxmemStatus::ok(); + }; + + let unmapped = block.committed as u64; + let bytes = block.bytes as u64; + free_block(block); + counters().live_allocations.fetch_sub(1, Ordering::AcqRel); + // SAFETY: checked non-null above. + unsafe { *outcome_out = NxmemReleaseOutcome::complete(bytes, unmapped) }; + NxmemStatus::ok() + }) +} + +unsafe extern "C" fn plugin_enqueue_release( + ctx: *mut c_void, + allocation: *const NxmemAllocation, + ticket_out: *mut u64, +) -> NxmemStatus { + catch_status_panic(|| { + let state = plugin_entry!(ctx); + if allocation.is_null() || ticket_out.is_null() { + return NxmemStatus::with_message( + NxmemStatusCode::InvalidArgument, + "testplugin: enqueue_release received a null pointer", + ); + } + // SAFETY: checked non-null. + let allocation = unsafe { &*allocation }; + if let Err(status) = state.check(allocation) { + return status; + } + let ticket = next_ticket(); + match state.inner.lock() { + Ok(mut inner) => { + let Some(block) = inner.blocks.remove(&allocation.allocation_id) else { + return NxmemStatus::with_message( + NxmemStatusCode::UnknownAllocation, + "testplugin: that allocation id is not live on this mechanism", + ); + }; + inner.queue.push(QueuedRelease { + ticket, + allocation_id: allocation.allocation_id, + address: block.address, + bytes: block.bytes, + align: block.align, + }); + } + Err(_) => { + return NxmemStatus::with_message( + NxmemStatusCode::InternalError, + "testplugin: allocator state is poisoned", + ); + } + } + // A queued release still names this context, so it takes its own + // reference. That is what keeps the mechanism — and, through the + // host's `Arc`, the module itself — pinned until the + // free retires. + state.refcount.fetch_add(1, Ordering::AcqRel); + counters().queued_releases.fetch_add(1, Ordering::AcqRel); + // SAFETY: checked non-null above. + unsafe { *ticket_out = ticket }; + NxmemStatus::ok() + }) +} + +unsafe extern "C" fn plugin_drain_releases( + ctx: *mut c_void, + max: u64, + retired_out: *mut u64, +) -> NxmemStatus { + catch_status_panic(|| { + let state = plugin_entry!(ctx); + if retired_out.is_null() { + return NxmemStatus::with_message( + NxmemStatusCode::InvalidArgument, + "testplugin: drain_releases received a null pointer", + ); + } + if state.behaviour.never_drain { + // SAFETY: checked non-null above. + unsafe { *retired_out = 0 }; + return NxmemStatus::ok(); + } + + // Take the batch under the lock, then **drop the lock** before calling + // back into the host: a host callback may re-enter this module. + let batch: Vec = match state.inner.lock() { + Ok(mut inner) => { + let take = (max as usize).min(inner.queue.len()); + inner.queue.drain(..take).collect() + } + Err(_) => { + return NxmemStatus::with_message( + NxmemStatusCode::InternalError, + "testplugin: allocator state is poisoned", + ); + } + }; + + let mut retired = 0u64; + for queued in batch { + free_block(Block { + address: queued.address, + bytes: queued.bytes, + align: queued.align, + committed: queued.bytes, + }); + counters().live_allocations.fetch_sub(1, Ordering::AcqRel); + counters().queued_releases.fetch_sub(1, Ordering::AcqRel); + + let completion = NxmemReleaseCompletion::new( + queued.ticket, + state.mechanism_id, + queued.allocation_id, + NxmemReleaseOutcome::complete(queued.bytes as u64, queued.bytes as u64), + ); + let reported = state.report_completion(&completion); + // Drop the reference the enqueue took whether or not the host + // accepted the completion: the bytes are already gone, and holding + // the reference would pin the module forever. + release_state(state); + retired += 1; + if let Err(status) = reported { + // SAFETY: checked non-null above. A partial drain still reports + // how many entries retired, so the host's accounting stays + // exact even on the failure path. + unsafe { *retired_out = retired }; + return NxmemStatus::with_message( + NxmemStatusCode::CallbackFailed, + &format!( + "testplugin: the host rejected a release completion: {}", + status.describe() + ), + ); + } + } + // SAFETY: checked non-null above. + unsafe { *retired_out = retired }; + NxmemStatus::ok() + }) +} + +unsafe extern "C" fn plugin_pending_release_count( + ctx: *mut c_void, + count_out: *mut u64, +) -> NxmemStatus { + catch_status_panic(|| { + let state = plugin_entry!(ctx); + if count_out.is_null() { + return NxmemStatus::with_message( + NxmemStatusCode::InvalidArgument, + "testplugin: pending_release_count received a null pointer", + ); + } + let count = match state.inner.lock() { + Ok(inner) => inner.queue.len() as u64, + Err(_) => { + return NxmemStatus::with_message( + NxmemStatusCode::InternalError, + "testplugin: allocator state is poisoned", + ); + } + }; + // SAFETY: checked non-null above. + unsafe { *count_out = count }; + NxmemStatus::ok() + }) +} + +unsafe extern "C" fn plugin_retain(ctx: *mut c_void) { + catch_void_panic(|| { + // SAFETY: the host only passes back a `ctx` this module published. + if let Some(state) = unsafe { state(ctx) } { + state.refcount.fetch_add(1, Ordering::AcqRel); + } + }); +} + +unsafe extern "C" fn plugin_release(ctx: *mut c_void) { + catch_void_panic(|| { + // SAFETY: the host only passes back a `ctx` this module published. + if let Some(state) = unsafe { state(ctx) } { + release_state(state); + } + }); +} + +/// Drop one reference, destroying the state at zero. +/// +/// A queued release holds its own reference, so the state cannot be destroyed +/// while a deferred free still names it. `retain`/`release` never call back +/// into the host, as the contract requires. +fn release_state(state: &AllocatorState) { + let previous = state.refcount.fetch_sub(1, Ordering::AcqRel); + if previous != 1 { + return; + } + // Reclaim anything still outstanding rather than leaking it. + if let Ok(mut inner) = state.inner.lock() { + for (_, block) in inner.blocks.drain() { + free_block(block); + counters().live_allocations.fetch_sub(1, Ordering::AcqRel); + } + let queued: Vec = inner.queue.drain(..).collect(); + for entry in queued { + free_block(Block { + address: entry.address, + bytes: entry.bytes, + align: entry.align, + committed: entry.bytes, + }); + counters().live_allocations.fetch_sub(1, Ordering::AcqRel); + counters().queued_releases.fetch_sub(1, Ordering::AcqRel); + } + let prefixes: Vec = inner.prefixes.drain().map(|(_, prefix)| prefix).collect(); + for prefix in prefixes { + counters() + .live_capabilities + .fetch_sub(prefix.references, Ordering::AcqRel); + free_block(Block { + address: prefix.address, + bytes: prefix.bytes, + align: PREFIX_ALIGN, + committed: prefix.bytes, + }); + } + } + counters().live_allocators.fetch_sub(1, Ordering::AcqRel); + // SAFETY: the count reached zero, so this is the last reference. The + // pointer came from `Box::into_raw` in `plugin_open_allocator` and is + // reclaimed exactly once, here. + drop(unsafe { Box::from_raw(state as *const AllocatorState as *mut AllocatorState) }); +} + +// ─── virtual-backing slots ────────────────────────────────────────────────── + +unsafe extern "C" fn plugin_allocate_committed( + ctx: *mut c_void, + request: *const NxmemAllocRequest, + result_out: *mut NxmemAllocResult, +) -> NxmemStatus { + // SAFETY: same contract; the lazy path commits only what it was handed. + unsafe { allocate_inner(ctx, request, result_out, false) } +} + +/// Read and identity-check a range request. +/// +/// # Safety +/// +/// `request` must be null or point to a valid record for the call. +unsafe fn range_of<'a>( + state: &AllocatorState, + request: *const NxmemRangeRequest, +) -> Result<&'a NxmemRangeRequest, NxmemStatus> { + if request.is_null() { + return Err(NxmemStatus::with_message( + NxmemStatusCode::InvalidArgument, + "testplugin: range request is null", + )); + } + // SAFETY: delegated to this function's contract. + let request = unsafe { &*request }; + state.check(&request.allocation)?; + Ok(request) +} + +unsafe extern "C" fn plugin_commit_range( + ctx: *mut c_void, + request: *const NxmemRangeRequest, +) -> NxmemStatus { + catch_status_panic(|| { + let state = plugin_entry!(ctx); + // SAFETY: the host passes a valid record or null. + let request = match unsafe { range_of(state, request) } { + Ok(request) => request, + Err(status) => return status, + }; + match state.inner.lock() { + Ok(mut inner) => { + let Some(block) = inner.blocks.get_mut(&request.allocation.allocation_id) else { + return NxmemStatus::with_message( + NxmemStatusCode::UnknownAllocation, + "testplugin: that allocation id is not live on this mechanism", + ); + }; + let end = (request.range.offset + request.range.bytes) as usize; + if end > block.bytes { + return NxmemStatus::with_message( + NxmemStatusCode::InvalidArgument, + "testplugin: the range runs past the end of the allocation", + ); + } + block.committed = block.committed.max(end); + NxmemStatus::ok() + } + Err(_) => NxmemStatus::with_message( + NxmemStatusCode::InternalError, + "testplugin: allocator state is poisoned", + ), + } + }) +} + +unsafe extern "C" fn plugin_decommit_range( + ctx: *mut c_void, + request: *const NxmemRangeRequest, + unmapped_out: *mut u64, +) -> NxmemStatus { + catch_status_panic(|| { + let state = plugin_entry!(ctx); + if unmapped_out.is_null() { + return NxmemStatus::with_message( + NxmemStatusCode::InvalidArgument, + "testplugin: decommit_range received a null out-parameter", + ); + } + // SAFETY: the host passes a valid record or null. + let request = match unsafe { range_of(state, request) } { + Ok(request) => request, + Err(status) => return status, + }; + let unmapped = match state.inner.lock() { + Ok(mut inner) => { + let Some(block) = inner.blocks.get_mut(&request.allocation.allocation_id) else { + return NxmemStatus::with_message( + NxmemStatusCode::UnknownAllocation, + "testplugin: that allocation id is not live on this mechanism", + ); + }; + let start = request.range.offset as usize; + let released = block + .committed + .saturating_sub(start) + .min(request.range.bytes as usize); + block.committed -= released; + released as u64 + } + Err(_) => { + return NxmemStatus::with_message( + NxmemStatusCode::InternalError, + "testplugin: allocator state is poisoned", + ); + } + }; + // SAFETY: checked non-null above. + unsafe { *unmapped_out = unmapped }; + NxmemStatus::ok() + }) +} + +/// The test plugin's physical granularity. +/// +/// Present so the mapped-byte estimate is conservative in the same way a real +/// VMM's would be: only the mechanism knows its granularity. +const GRANULE: u64 = 4096; + +/// Alignment used for shared prefix storage. +const PREFIX_ALIGN: usize = 64; + +fn round_up_to_granule(bytes: u64) -> u64 { + bytes.div_ceil(GRANULE) * GRANULE +} + +unsafe extern "C" fn plugin_mapped_bytes_for_ranges( + ctx: *mut c_void, + requests: *const NxmemRangeRequest, + count: u64, + mapped_out: *mut u64, +) -> NxmemStatus { + catch_status_panic(|| { + let state = plugin_entry!(ctx); + if mapped_out.is_null() || (requests.is_null() && count != 0) { + return NxmemStatus::with_message( + NxmemStatusCode::InvalidArgument, + "testplugin: mapped_bytes_for_ranges received a null pointer", + ); + } + let mut total = 0u64; + for index in 0..count { + // SAFETY: the host promises `count` readable records at `requests`. + let request = unsafe { &*requests.add(index as usize) }; + if let Err(status) = state.check(&request.allocation) { + return status; + } + total = total.saturating_add(round_up_to_granule(request.range.bytes)); + } + // SAFETY: checked non-null above. + unsafe { *mapped_out = total }; + NxmemStatus::ok() + }) +} + +unsafe extern "C" fn plugin_mapped_bytes_for_allocation( + ctx: *mut c_void, + request: *const NxmemAllocRequest, + mapped_out: *mut u64, +) -> NxmemStatus { + catch_status_panic(|| { + let state = plugin_entry!(ctx); + if request.is_null() || mapped_out.is_null() { + return NxmemStatus::with_message( + NxmemStatusCode::InvalidArgument, + "testplugin: mapped_bytes_for_allocation received a null pointer", + ); + } + // SAFETY: checked non-null. + let request = unsafe { &*request }; + if let Err(status) = check_identity( + state.mechanism_id, + state.device, + request.mechanism_id, + request.device, + ) { + return status; + } + // SAFETY: checked non-null above. + unsafe { *mapped_out = round_up_to_granule(request.bytes) }; + NxmemStatus::ok() + }) +} + +unsafe extern "C" fn plugin_committed_bytes( + ctx: *mut c_void, + allocation: *const NxmemAllocation, + committed_out: *mut u64, +) -> NxmemStatus { + catch_status_panic(|| { + let state = plugin_entry!(ctx); + if allocation.is_null() || committed_out.is_null() { + return NxmemStatus::with_message( + NxmemStatusCode::InvalidArgument, + "testplugin: committed_bytes received a null pointer", + ); + } + // SAFETY: checked non-null. + let allocation = unsafe { &*allocation }; + if let Err(status) = state.check(allocation) { + return status; + } + let committed = match state.inner.lock() { + Ok(inner) => inner + .blocks + .get(&allocation.allocation_id) + .map(|block| block.committed as u64) + .unwrap_or(0), + Err(_) => 0, + }; + // SAFETY: checked non-null above. + unsafe { *committed_out = committed }; + NxmemStatus::ok() + }) +} + +// ─── shared-mapping slots ─────────────────────────────────────────────────── + +unsafe extern "C" fn plugin_create_shared_prefix( + ctx: *mut c_void, + mechanism_id: u64, + bytes: u64, + handle_out: *mut NxmemSharedPrefixHandle, +) -> NxmemStatus { + catch_status_panic(|| { + let state = plugin_entry!(ctx); + if handle_out.is_null() { + return NxmemStatus::with_message( + NxmemStatusCode::InvalidArgument, + "testplugin: create_shared_prefix received a null out-parameter", + ); + } + if mechanism_id != state.mechanism_id { + return NxmemStatus::with_message( + NxmemStatusCode::WrongMechanism, + "testplugin: a prefix may only be created on its own mechanism", + ); + } + let layout = match layout_of(bytes as usize, PREFIX_ALIGN) { + Ok(layout) => layout, + Err(status) => return status, + }; + // SAFETY: `layout` has a non-zero size and a valid alignment. + let raw = unsafe { alloc(layout) }; + if raw.is_null() { + return NxmemStatus::with_message( + NxmemStatusCode::OutOfMemory, + "testplugin: the backing host allocator refused the prefix", + ); + } + let handle_id = next_prefix_handle(); + match state.inner.lock() { + Ok(mut inner) => { + inner.prefixes.insert( + handle_id, + Prefix { + bytes: bytes as usize, + address: raw as usize, + references: 1, + }, + ); + } + Err(_) => { + // SAFETY: allocated just above and never handed out. + unsafe { dealloc(raw, layout) }; + return NxmemStatus::with_message( + NxmemStatusCode::InternalError, + "testplugin: allocator state is poisoned", + ); + } + } + counters().live_capabilities.fetch_add(1, Ordering::AcqRel); + + let mut handle = NxmemSharedPrefixHandle::zeroed(); + handle.mechanism_id = state.mechanism_id; + handle.handle = handle_id; + handle.device = state.device; + handle.device_ptr = raw as u64; + handle.committed_physical_bytes = round_up_to_granule(bytes); + handle.mapped_bytes = 0; + handle.requested_bytes = bytes; + // SAFETY: checked non-null above. + unsafe { *handle_out = handle }; + NxmemStatus::ok() + }) +} + +unsafe extern "C" fn plugin_retain_shared_prefix( + ctx: *mut c_void, + handle: *const NxmemSharedPrefixHandle, +) -> NxmemStatus { + catch_status_panic(|| { + let state = plugin_entry!(ctx); + if handle.is_null() { + return NxmemStatus::with_message( + NxmemStatusCode::InvalidArgument, + "testplugin: retain_shared_prefix received a null handle", + ); + } + // SAFETY: checked non-null. + let handle = unsafe { &*handle }; + if handle.mechanism_id != state.mechanism_id { + return NxmemStatus::with_message( + NxmemStatusCode::WrongMechanism, + "testplugin: that prefix belongs to another mechanism", + ); + } + match state.inner.lock() { + Ok(mut inner) => match inner.prefixes.get_mut(&handle.handle) { + Some(prefix) => { + prefix.references += 1; + counters().live_capabilities.fetch_add(1, Ordering::AcqRel); + NxmemStatus::ok() + } + None => NxmemStatus::with_message( + NxmemStatusCode::UnknownAllocation, + "testplugin: that prefix is not live on this mechanism", + ), + }, + Err(_) => NxmemStatus::with_message( + NxmemStatusCode::InternalError, + "testplugin: allocator state is poisoned", + ), + } + }) +} + +unsafe extern "C" fn plugin_release_shared_prefix( + ctx: *mut c_void, + handle: *const NxmemSharedPrefixHandle, +) -> NxmemStatus { + catch_status_panic(|| { + let state = plugin_entry!(ctx); + if handle.is_null() { + return NxmemStatus::with_message( + NxmemStatusCode::InvalidArgument, + "testplugin: release_shared_prefix received a null handle", + ); + } + // SAFETY: checked non-null. + let handle = unsafe { &*handle }; + if handle.mechanism_id != state.mechanism_id { + return NxmemStatus::with_message( + NxmemStatusCode::WrongMechanism, + "testplugin: that prefix belongs to another mechanism", + ); + } + match state.inner.lock() { + Ok(mut inner) => { + let Some(prefix) = inner.prefixes.get_mut(&handle.handle) else { + return NxmemStatus::with_message( + NxmemStatusCode::UnknownAllocation, + "testplugin: that prefix is not live on this mechanism", + ); + }; + prefix.references -= 1; + counters().live_capabilities.fetch_sub(1, Ordering::AcqRel); + if prefix.references == 0 { + let prefix = inner.prefixes.remove(&handle.handle).expect("just seen"); + free_block(Block { + address: prefix.address, + bytes: prefix.bytes, + align: PREFIX_ALIGN, + committed: prefix.bytes, + }); + } + NxmemStatus::ok() + } + Err(_) => NxmemStatus::with_message( + NxmemStatusCode::InternalError, + "testplugin: allocator state is poisoned", + ), + } + }) +} + +unsafe extern "C" fn plugin_incremental_owned_bytes( + ctx: *mut c_void, + handle: *const NxmemSharedPrefixHandle, + bytes_out: *mut u64, +) -> NxmemStatus { + catch_status_panic(|| { + let state = plugin_entry!(ctx); + if handle.is_null() || bytes_out.is_null() { + return NxmemStatus::with_message( + NxmemStatusCode::InvalidArgument, + "testplugin: incremental_owned_bytes received a null pointer", + ); + } + // SAFETY: checked non-null. + let handle = unsafe { &*handle }; + // A foreign prefix is refused rather than costed as free. + if handle.mechanism_id != state.mechanism_id { + return NxmemStatus::with_message( + NxmemStatusCode::WrongMechanism, + "testplugin: that prefix belongs to another mechanism, so its incremental cost \ + cannot be reported as zero", + ); + } + if handle.device != state.device { + return NxmemStatus::with_message( + NxmemStatusCode::WrongDevice, + "testplugin: that prefix lives on another device", + ); + } + let known = match state.inner.lock() { + Ok(inner) => inner.prefixes.contains_key(&handle.handle), + Err(_) => false, + }; + if !known { + return NxmemStatus::with_message( + NxmemStatusCode::UnknownAllocation, + "testplugin: that prefix is not live on this mechanism", + ); + } + // The physical bytes were charged once when the prefix was created, so + // an additional mapping owns nothing new. + // SAFETY: checked non-null above. + unsafe { *bytes_out = 0 }; + NxmemStatus::ok() + }) +} + +unsafe extern "C" fn plugin_commit_shared_prefix( + ctx: *mut c_void, + request: *const NxmemSharedPrefixCommitRequest, + info_out: *mut NxmemSharedPrefixCommitInfo, +) -> NxmemStatus { + catch_status_panic(|| { + let state = plugin_entry!(ctx); + if request.is_null() || info_out.is_null() { + return NxmemStatus::with_message( + NxmemStatusCode::InvalidArgument, + "testplugin: commit_shared_prefix received a null pointer", + ); + } + // SAFETY: checked non-null. + let request = unsafe { &*request }; + if let Err(status) = state.check(&request.allocation) { + return status; + } + if request.prefix.mechanism_id != state.mechanism_id { + return NxmemStatus::with_message( + NxmemStatusCode::WrongMechanism, + "testplugin: that prefix belongs to another mechanism", + ); + } + let prefix_bytes = match state.inner.lock() { + Ok(mut inner) => { + let Some(prefix) = inner.prefixes.get(&request.prefix.handle).copied() else { + return NxmemStatus::with_message( + NxmemStatusCode::UnknownAllocation, + "testplugin: that prefix is not live on this mechanism", + ); + }; + let end = request.byte_offset as usize + prefix.bytes; + let Some(block) = inner.blocks.get_mut(&request.allocation.allocation_id) else { + return NxmemStatus::with_message( + NxmemStatusCode::UnknownAllocation, + "testplugin: that allocation id is not live on this mechanism", + ); + }; + if end > block.bytes { + return NxmemStatus::with_message( + NxmemStatusCode::InvalidArgument, + "testplugin: the prefix does not fit in the allocation at that offset", + ); + } + block.committed = block.committed.max(end); + prefix.bytes as u64 + } + Err(_) => { + return NxmemStatus::with_message( + NxmemStatusCode::InternalError, + "testplugin: allocator state is poisoned", + ); + } + }; + + let mut info = NxmemSharedPrefixCommitInfo::zeroed(); + // The prefix was charged once at creation, so mapping it again owns + // nothing new but does map bytes. Keeping those two axes distinct is + // the whole point of this record. + info.additional_owned_bytes = 0; + info.newly_mapped_bytes = round_up_to_granule(prefix_bytes); + info.granules = info.newly_mapped_bytes / GRANULE; + // SAFETY: checked non-null above. + unsafe { *info_out = info }; + NxmemStatus::ok() + }) +} + +// ─── factory slots ────────────────────────────────────────────────────────── + +unsafe extern "C" fn plugin_open_allocator( + ctx: *mut c_void, + request: *const NxmemOpenRequest, + allocator_out: *mut *const NxmemAllocatorVtable, +) -> NxmemStatus { + catch_status_panic(|| { + if ctx.is_null() || request.is_null() || allocator_out.is_null() { + return NxmemStatus::with_message( + NxmemStatusCode::InvalidArgument, + "testplugin: open_allocator received a null pointer", + ); + } + // SAFETY: `ctx` is the `&'static Mechanism` this module published in + // its factory table; the host only echoes it back. + let mechanism = unsafe { &*(ctx as *const Mechanism) }; + // SAFETY: checked non-null. + let request = unsafe { &*request }; + + let missing = request.required_capability_flags & !mechanism.behaviour.capability_flags; + if missing != 0 { + return NxmemStatus::with_message( + NxmemStatusCode::UnsupportedCapability, + &format!( + "testplugin: mechanism `{}` does not provide the required capabilities \ + {missing:#x}", + mechanism.name + ), + ); + } + + let mechanism_id = next_mechanism_id(); + let behaviour = mechanism.behaviour; + + if behaviour.short_allocator_struct { + // A plugin built against a mismatched header publishes a vtable + // the host cannot read. It must not allocate state first: the host + // has nothing well-formed enough to hand the state back through. + // Returning a static, stateless vtable is the only shape of this + // bug that does not also leak. + // SAFETY: checked non-null above. + unsafe { *allocator_out = short_allocator_vtable() }; + return NxmemStatus::ok(); + } + // The vtable's level is the lesser of what the mechanism implements + // and what the host negotiated: an older mechanism inside a newer + // module keeps working, and a newer mechanism never offers a slot an + // older host would not know how to call. + let abi_minor = behaviour.abi_minor.min(request.abi_minor); + + let backing_vtable = + (behaviour.capability_flags & NXMEM_CAP_VIRTUAL_BACKING != 0).then(|| { + let mut vtable = NxmemVirtualBackingVtable::zeroed(); + vtable.abi_minor = abi_minor; + vtable.mechanism_id = mechanism_id; + vtable.allocate_committed = Some(plugin_allocate_committed); + vtable.commit_range = Some(plugin_commit_range); + vtable.decommit_range = Some(plugin_decommit_range); + vtable.mapped_bytes_for_ranges = Some(plugin_mapped_bytes_for_ranges); + vtable.mapped_bytes_for_allocation = Some(plugin_mapped_bytes_for_allocation); + vtable.committed_bytes = Some(plugin_committed_bytes); + Box::new(vtable) + }); + let shared_vtable = (behaviour.capability_flags & NXMEM_CAP_SHARED_MAPPING != 0).then(|| { + let mut vtable = NxmemSharedMappingVtable::zeroed(); + vtable.abi_minor = abi_minor; + vtable.mechanism_id = mechanism_id; + vtable.create_shared_prefix = Some(plugin_create_shared_prefix); + vtable.retain_shared_prefix = Some(plugin_retain_shared_prefix); + vtable.release_shared_prefix = Some(plugin_release_shared_prefix); + vtable.incremental_owned_bytes = Some(plugin_incremental_owned_bytes); + vtable.commit_shared_prefix = Some(plugin_commit_shared_prefix); + Box::new(vtable) + }); + + let mut allocator_vtable = NxmemAllocatorVtable::zeroed(); + allocator_vtable.abi_minor = abi_minor; + allocator_vtable.mechanism_id = mechanism_id; + allocator_vtable.device = request.device; + allocator_vtable.capability_flags = behaviour.capability_flags; + allocator_vtable.name = mechanism.c_name.as_ptr().cast(); + allocator_vtable.allocate = Some(plugin_allocate); + allocator_vtable.deallocate = Some(plugin_deallocate); + allocator_vtable.retain = Some(plugin_retain); + allocator_vtable.release = Some(plugin_release); + if behaviour.capability_flags & NXMEM_CAP_DEFERRED_RELEASE != 0 { + allocator_vtable.enqueue_release = Some(plugin_enqueue_release); + allocator_vtable.drain_releases = Some(plugin_drain_releases); + allocator_vtable.pending_release_count = Some(plugin_pending_release_count); + } + if abi_minor >= 1 && behaviour.capability_flags & NXMEM_CAP_STRUCTURED_RELEASE != 0 { + allocator_vtable.release_allocation = Some(plugin_release_allocation); + } + let state = Box::new(AllocatorState { + mechanism_id, + device: request.device, + behaviour, + refcount: AtomicU64::new(1), + callbacks: request.callbacks, + inner: Mutex::new(AllocatorInner::default()), + allocator_vtable: Box::new(allocator_vtable), + backing_vtable, + shared_vtable, + }); + + // Wire the vtables to the state now that it has a stable address. + let state_ptr = Box::into_raw(state); + // SAFETY: `state_ptr` came from `Box::into_raw` immediately above and + // nothing else references it yet, so this is the only live borrow. + let state = unsafe { &mut *state_ptr }; + let allocator_ctx = state_ptr as *mut c_void; + state.allocator_vtable.ctx = allocator_ctx; + if let Some(backing) = state.backing_vtable.as_mut() { + backing.ctx = allocator_ctx; + } + if let Some(shared) = state.shared_vtable.as_mut() { + shared.ctx = allocator_ctx; + } + state.allocator_vtable.virtual_backing = state + .backing_vtable + .as_ref() + .map_or(core::ptr::null(), |backing| &raw const **backing); + state.allocator_vtable.shared_mapping = state + .shared_vtable + .as_ref() + .map_or(core::ptr::null(), |shared| &raw const **shared); + + counters().live_allocators.fetch_add(1, Ordering::AcqRel); + // SAFETY: checked non-null above. + unsafe { *allocator_out = &raw const *state.allocator_vtable }; + NxmemStatus::ok() + }) +} + +/// A deliberately undersized allocator vtable. +/// +/// `struct_size` claims fewer bytes than even the baseline prefix needs, so +/// the host must refuse it before reading a single function pointer. It is +/// stateless and `'static`, so refusing it leaks nothing. +struct ShortVtable(NxmemAllocatorVtable); + +// SAFETY: immutable after construction and holding only null pointers. +unsafe impl Send for ShortVtable {} +// SAFETY: as above. +unsafe impl Sync for ShortVtable {} + +fn short_allocator_vtable() -> *const NxmemAllocatorVtable { + static VTABLE: OnceLock = OnceLock::new(); + &VTABLE + .get_or_init(|| { + let mut vtable = NxmemAllocatorVtable::zeroed(); + vtable.struct_size = (NxmemAllocatorVtable::MIN_STRUCT_SIZE_MINOR_0 as u32) / 2; + ShortVtable(vtable) + }) + .0 +} + +unsafe extern "C" fn plugin_factory_release(_ctx: *mut c_void) { + // Factories are `&'static Mechanism`, so there is nothing to free. The + // slot still exists because the contract requires an explicit release for + // every ABI-owned object, and a real plugin may own state here. + catch_void_panic(|| { + FACTORY_RELEASES.fetch_add(1, Ordering::AcqRel); + }); +} + +/// How many times the host has released a factory. +/// +/// A statically linked `rlib` copy of this crate has its *own* copy of this +/// static, so a test cannot observe the loaded module's count by reading it +/// directly. [`NxmemTestpluginFactoryReleases`] exists to expose it from the +/// module the host actually loaded. +pub static FACTORY_RELEASES: AtomicU64 = AtomicU64::new(0); + +/// The name of the test-only introspection symbol. +pub const SYMBOL_FACTORY_RELEASES: &[u8] = b"NxmemTestpluginFactoryReleases\0"; + +/// Test-only introspection: how many factories this module has had released. +/// +/// **Not part of the nxmem ABI.** It exists so the ABI tests can assert on the +/// loaded module's own state rather than on a separate statically linked copy +/// of this crate. +#[unsafe(no_mangle)] +pub extern "C" fn NxmemTestpluginFactoryReleases() -> u64 { + FACTORY_RELEASES.load(Ordering::Acquire) +} + +/// Set once the module has published its factory vtables. +static FACTORIES_BUILT: AtomicBool = AtomicBool::new(false); + +/// The published factory table. +/// +/// A vtable holds raw pointers, so it is not automatically `Send`/`Sync`. The +/// table is built once and never mutated afterwards, and every pointer in it +/// addresses `'static` module data, so sharing it is sound. +struct FactoryTable(Vec); + +// SAFETY: immutable after construction; every pointer targets `'static` data +// owned by this module. +unsafe impl Send for FactoryTable {} +// SAFETY: as above. +unsafe impl Sync for FactoryTable {} + +fn factory_vtables() -> &'static [NxmemAllocatorFactoryVtable] { + static VTABLES: OnceLock = OnceLock::new(); + &VTABLES + .get_or_init(|| { + FACTORIES_BUILT.store(true, Ordering::Release); + FactoryTable(build_factory_table()) + }) + .0 +} + +fn build_factory_table() -> Vec { + { + MECHANISMS + .iter() + .map(|mechanism| { + let mut vtable = NxmemAllocatorFactoryVtable::zeroed(); + vtable.abi_minor = mechanism.behaviour.abi_minor; + vtable.name = mechanism.c_name.as_ptr().cast(); + // Host memory keeps the test suite portable: it exercises the + // whole ABI on a machine with no accelerator. + vtable.device = NxmemDeviceId::HOST; + vtable.capability_flags = mechanism.behaviour.capability_flags; + vtable.ctx = (mechanism as *const Mechanism).cast_mut().cast(); + vtable.open_allocator = Some(plugin_open_allocator); + vtable.release = Some(plugin_factory_release); + vtable + }) + .collect() + } +} + +// ─── entry points ─────────────────────────────────────────────────────────── + +fn negotiate_impl( + request: *const NxmemNegotiateRequest, + response_out: *mut NxmemNegotiateResponse, +) -> NxmemStatus { + // The module implements the current contract; each *mechanism* then + // declares its own level, which is how one module can ship both a current + // and a deliberately older mechanism. + // SAFETY: `negotiate_as` null-checks and size-checks both pointers before + // reading any field. + unsafe { + negotiate_as( + request, + response_out, + NxmemVersionRange { + major_min: NXMEM_ABI_VERSION_MAJOR, + major_max: NXMEM_ABI_VERSION_MAJOR, + minor_min: 0, + minor_max: NXMEM_ABI_VERSION_MINOR, + }, + NXMEM_CAP_ALLOCATOR + | NXMEM_CAP_VIRTUAL_BACKING + | NXMEM_CAP_SHARED_MAPPING + | NXMEM_CAP_DEFERRED_RELEASE + | NXMEM_CAP_STRUCTURED_RELEASE, + ) + } +} + +fn create_factories_impl( + out_factories: *mut *const NxmemAllocatorFactoryVtable, + max_factories: u64, + out_count: *mut u64, +) -> NxmemStatus { + if out_factories.is_null() || out_count.is_null() { + return NxmemStatus::with_message( + NxmemStatusCode::InvalidArgument, + "testplugin: create_factories received a null pointer", + ); + } + let vtables = factory_vtables(); + let count = (vtables.len() as u64).min(max_factories); + for (index, vtable) in vtables.iter().take(count as usize).enumerate() { + // SAFETY: the host promises `max_factories` writable slots and `count` + // never exceeds it. + unsafe { *out_factories.add(index) = vtable as *const _ }; + } + // SAFETY: checked non-null above. + unsafe { *out_count = count }; + NxmemStatus::ok() +} + +fn unload_readiness_impl(report_out: *mut NxmemUnloadReport) -> NxmemStatus { + if report_out.is_null() { + return NxmemStatus::with_message( + NxmemStatusCode::InvalidArgument, + "testplugin: unload readiness received a null pointer", + ); + } + let counters = counters(); + let mut report = NxmemUnloadReport::zeroed(); + report.live_allocators = counters.live_allocators.load(Ordering::Acquire); + report.live_allocations = counters.live_allocations.load(Ordering::Acquire); + report.live_views = counters.live_views.load(Ordering::Acquire); + report.live_capabilities = counters.live_capabilities.load(Ordering::Acquire); + report.queued_releases = counters.queued_releases.load(Ordering::Acquire); + // SAFETY: checked non-null above. + unsafe { *report_out = report }; + NxmemStatus::ok() +} + +onnx_runtime_memory_abi::export_nxmem_plugin! { + negotiate: negotiate_impl, + factories: create_factories_impl, + unload_readiness: unload_readiness_impl, +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn the_published_names_match_the_mechanism_table() { + let actual: Vec<&str> = MECHANISMS.iter().map(|mechanism| mechanism.name).collect(); + assert_eq!(actual, MECHANISM_NAMES); + } + + #[test] + fn every_mechanism_name_is_valid_c_and_utf8() { + for mechanism in MECHANISMS { + assert_eq!( + mechanism.c_name.to_str().expect("valid UTF-8"), + mechanism.name + ); + } + } + + #[test] + fn every_mechanism_claims_the_allocator_capability() { + for mechanism in MECHANISMS { + assert_ne!( + mechanism.behaviour.capability_flags & NXMEM_CAP_ALLOCATOR, + 0, + "{} must be able to allocate", + mechanism.name + ); + } + } + + #[test] + fn the_legacy_mechanism_stays_on_the_baseline_prefix() { + let legacy = MECHANISMS + .iter() + .find(|mechanism| mechanism.name == "legacy-1-0") + .expect("the legacy mechanism exists"); + assert_eq!(legacy.behaviour.abi_minor, 0); + assert_eq!( + legacy.behaviour.capability_flags & NXMEM_CAP_STRUCTURED_RELEASE, + 0, + "structured release does not exist at minor 0" + ); + } + + #[test] + fn granule_rounding_is_conservative() { + assert_eq!(round_up_to_granule(0), 0); + assert_eq!(round_up_to_granule(1), GRANULE); + assert_eq!(round_up_to_granule(GRANULE), GRANULE); + assert_eq!(round_up_to_granule(GRANULE + 1), GRANULE * 2); + } + + #[test] + fn the_factory_table_is_built_lazily_and_stays_stable() { + let first = factory_vtables().as_ptr(); + let second = factory_vtables().as_ptr(); + assert_eq!(first, second, "the factory table must be stable"); + assert!(FACTORIES_BUILT.load(Ordering::Acquire)); + assert_eq!(factory_vtables().len(), MECHANISMS.len()); + } +} From 03faeebac908df942dd35c54914a2bd1cc469648 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Wed, 19 Aug 2026 01:41:05 -0700 Subject: [PATCH 3/3] nxmem: public C header, minimal C example, docs, lint fixes --- .../examples/minimal_plugin.c | 305 +++++++++ .../include/nxmem_memory_abi.h | 645 ++++++++++++++++++ crates/onnx-runtime-memory-abi/src/lib.rs | 23 +- crates/onnx-runtime-memory-abi/src/status.rs | 8 +- crates/onnx-runtime-memory-abi/src/types.rs | 14 +- crates/onnx-runtime-memory-abi/src/version.rs | 3 +- crates/onnx-runtime-memory-abi/src/vtable.rs | 34 +- .../tests/header_contract.rs | 230 +++++++ .../onnx-runtime-memory-host/src/allocator.rs | 104 ++- crates/onnx-runtime-memory-host/src/error.rs | 16 +- crates/onnx-runtime-memory-host/src/lib.rs | 6 + crates/onnx-runtime-memory-host/src/loader.rs | 11 +- .../tests/nxmem_abi.rs | 87 ++- .../tests/support/testplugin.rs | 6 +- .../onnx-runtime-memory-testplugin/src/lib.rs | 123 +++- docs/memory/MEMORY_PLUGIN_ABI.md | 214 ++++++ 16 files changed, 1705 insertions(+), 124 deletions(-) create mode 100644 crates/onnx-runtime-memory-abi/examples/minimal_plugin.c create mode 100644 crates/onnx-runtime-memory-abi/include/nxmem_memory_abi.h create mode 100644 crates/onnx-runtime-memory-abi/tests/header_contract.rs create mode 100644 docs/memory/MEMORY_PLUGIN_ABI.md diff --git a/crates/onnx-runtime-memory-abi/examples/minimal_plugin.c b/crates/onnx-runtime-memory-abi/examples/minimal_plugin.c new file mode 100644 index 0000000000..e02d0abbda --- /dev/null +++ b/crates/onnx-runtime-memory-abi/examples/minimal_plugin.c @@ -0,0 +1,305 @@ +/* + * A minimal nxmem memory plugin, in portable C. + * + * It implements exactly one mechanism, "minimal", with the required slots and + * no optional capability at all. That is deliberate: it is the smallest thing + * that is a *correct* plugin, so it doubles as documentation of what is truly + * mandatory. + * + * It needs nothing from the ONNX GenAI workspace but the header next door, and + * links against nothing. Build it standalone: + * + * cc -shared -fPIC -I../include -o libminimal_plugin.so minimal_plugin.c + * + * `nxmem_c_example_compiles` in this crate compiles this file on every unix + * test run, so the header cannot drift away from a real C consumer unnoticed. + * + * What this example is careful about, because a plugin that gets these wrong + * corrupts the host rather than merely failing: + * + * - it fills in `struct_size` on everything it writes; + * - it rejects an allocation whose `mechanism_id` or `device` is not its own, + * which is what stops one provider's pointer being freed by another; + * - it never blocks, never takes a lock across a host callback, and never + * lets an error escape as anything but a status; + * - it counts live objects and reports them from NxmemQueryUnloadReadiness, + * so the host will not unmap this code while it is still in use. + */ + +#include "nxmem_memory_abi.h" + +#include +#include + +/* ─── status helpers ───────────────────────────────────────────────────── */ + +static NxmemStatus nxmem_status(uint32_t code, const char *message) { + NxmemStatus status; + memset(&status, 0, sizeof(status)); + status.code = code; + if (message != NULL) { + size_t len = strlen(message); + if (len > NXMEM_STATUS_MESSAGE_MAX) { + len = NXMEM_STATUS_MESSAGE_MAX; + } + memcpy(status.message, message, len); + status.message_len = (uint32_t)len; + } + return status; +} + +static NxmemStatus nxmem_ok(void) { return nxmem_status(NXMEM_STATUS_OK, NULL); } + +/* ─── module state ─────────────────────────────────────────────────────── */ + +#define MINIMAL_MECHANISM_ID 0x6D696E31ull /* "min1" */ + +/* + * Live-object counts. A real plugin would make these atomic; this example is + * kept single-file and dependency-free, and the host serialises its own calls + * during load and unload. + */ +static uint64_t g_live_allocators; +static uint64_t g_live_allocations; + +static const char kMechanismName[] = "minimal"; + +/* ─── allocator ────────────────────────────────────────────────────────── */ + +/* + * `identity_ok` is the whole cross-provider-misuse defence. An allocation + * describes which mechanism and which device it belongs to; a mechanism that + * skips this check will happily free a pointer that belongs to someone else. + */ +static NxmemStatus identity_ok(const NxmemAllocation *allocation) { + if (allocation == NULL) { + return nxmem_status(NXMEM_STATUS_INVALID_ARGUMENT, + "minimal: null allocation"); + } + if (allocation->mechanism_id != MINIMAL_MECHANISM_ID) { + return nxmem_status(NXMEM_STATUS_WRONG_MECHANISM, + "minimal: that allocation belongs to another mechanism"); + } + if (allocation->device.tier != NXMEM_TIER_HOST || + allocation->device.index != 0) { + return nxmem_status(NXMEM_STATUS_WRONG_DEVICE, + "minimal: that allocation belongs to another device"); + } + return nxmem_ok(); +} + +static NxmemStatus minimal_allocate(void *ctx, const NxmemAllocRequest *request, + NxmemAllocResult *result_out) { + void *ptr = NULL; + size_t align = 0; + size_t bytes = 0; + + (void)ctx; + if (request == NULL || result_out == NULL) { + return nxmem_status(NXMEM_STATUS_INVALID_ARGUMENT, + "minimal: null pointer in allocate"); + } + if (request->mechanism_id != MINIMAL_MECHANISM_ID) { + return nxmem_status(NXMEM_STATUS_WRONG_MECHANISM, + "minimal: allocate called on another mechanism"); + } + + align = (size_t)request->align; + if (align < sizeof(void *)) { + align = sizeof(void *); + } + /* aligned_alloc requires a size that is a multiple of the alignment. */ + bytes = (size_t)request->bytes; + if (bytes == 0) { + bytes = align; + } + bytes = ((bytes + align - 1) / align) * align; + + ptr = aligned_alloc(align, bytes); + if (ptr == NULL) { + return nxmem_status(NXMEM_STATUS_OUT_OF_MEMORY, + "minimal: the host allocator refused the request"); + } + + memset(result_out, 0, sizeof(*result_out)); + result_out->struct_size = (uint32_t)sizeof(*result_out); + result_out->ptr = (uint8_t *)ptr; + /* This mechanism is eager: everything it owns, it has also mapped. */ + result_out->owned_bytes = (uint64_t)bytes; + result_out->mapped_bytes = (uint64_t)bytes; + + g_live_allocations += 1; + return nxmem_ok(); +} + +static NxmemStatus minimal_deallocate(void *ctx, + const NxmemAllocation *allocation, + uint64_t *unmapped_bytes_out) { + NxmemStatus status; + + (void)ctx; + if (unmapped_bytes_out == NULL) { + return nxmem_status(NXMEM_STATUS_INVALID_ARGUMENT, + "minimal: null out-parameter in deallocate"); + } + status = identity_ok(allocation); + if (status.code != NXMEM_STATUS_OK) { + return status; + } + + free(allocation->ptr); + *unmapped_bytes_out = allocation->bytes; + if (g_live_allocations > 0) { + g_live_allocations -= 1; + } + return nxmem_ok(); +} + +/* + * This example opens one shared, statically allocated allocator, so `retain` + * and `release` only move a count. A plugin with per-open state would free that + * state when the count reaches zero — and only then, and only if no queued + * release still names it. + */ +static void minimal_retain(void *ctx) { + (void)ctx; + g_live_allocators += 1; +} + +static void minimal_release(void *ctx) { + (void)ctx; + if (g_live_allocators > 0) { + g_live_allocators -= 1; + } +} + +static const NxmemAllocatorVtable kAllocator = { + (uint32_t)sizeof(NxmemAllocatorVtable), + NXMEM_ABI_VERSION_MINOR_BASELINE, + MINIMAL_MECHANISM_ID, + {NXMEM_TIER_HOST, 0}, + NXMEM_CAP_ALLOCATOR, + (const uint8_t *)kMechanismName, + NULL, + minimal_allocate, + minimal_deallocate, + minimal_retain, + minimal_release, + /* Every optional slot is NULL: this mechanism has no lazy backing, no + shared mapping, no deferred release, and no structured release. NULL is + how "unsupported" is spelled — never a stub that returns OK. */ + NULL, + NULL, + NULL, + NULL, + NULL, + NULL, +}; + +/* ─── factory ──────────────────────────────────────────────────────────── */ + +static NxmemStatus +minimal_open_allocator(void *ctx, const NxmemOpenRequest *request, + const NxmemAllocatorVtable **allocator_out) { + (void)ctx; + if (request == NULL || allocator_out == NULL) { + return nxmem_status(NXMEM_STATUS_INVALID_ARGUMENT, + "minimal: null pointer in open_allocator"); + } + /* Refuse rather than return an allocator that is missing something the host + said it needs. Silently handing back less is how a host ends up calling a + NULL slot. */ + if ((request->required_capability_flags & ~(uint64_t)NXMEM_CAP_ALLOCATOR) != + 0) { + return nxmem_status(NXMEM_STATUS_UNSUPPORTED_CAPABILITY, + "minimal: this mechanism offers plain allocation only"); + } + if (request->device.tier != NXMEM_TIER_HOST || request->device.index != 0) { + return nxmem_status(NXMEM_STATUS_WRONG_DEVICE, + "minimal: this mechanism serves host memory only"); + } + + g_live_allocators += 1; + *allocator_out = &kAllocator; + return nxmem_ok(); +} + +static void minimal_factory_release(void *ctx) { (void)ctx; } + +static const NxmemAllocatorFactoryVtable kFactory = { + (uint32_t)sizeof(NxmemAllocatorFactoryVtable), + NXMEM_ABI_VERSION_MINOR_BASELINE, + (const uint8_t *)kMechanismName, + {NXMEM_TIER_HOST, 0}, + NXMEM_CAP_ALLOCATOR, + NULL, + minimal_open_allocator, + minimal_factory_release, +}; + +/* ─── required exports ─────────────────────────────────────────────────── */ + +NxmemStatus NxmemNegotiate(const NxmemNegotiateRequest *request, + NxmemNegotiateResponse *response_out) { + if (request == NULL || response_out == NULL) { + return nxmem_status(NXMEM_STATUS_INVALID_ARGUMENT, + "minimal: null pointer in negotiate"); + } + /* Read only the prefix this plugin was built to understand. */ + if (request->struct_size < (uint32_t)sizeof(NxmemNegotiateRequest)) { + return nxmem_status(NXMEM_STATUS_SHORT_STRUCT, + "minimal: the host's negotiate request is too short"); + } + if (request->host_range.major_min > NXMEM_ABI_VERSION_MAJOR || + request->host_range.major_max < NXMEM_ABI_VERSION_MAJOR) { + return nxmem_status(NXMEM_STATUS_VERSION_MISMATCH, + "minimal: this plugin implements nxmem major 1"); + } + + memset(response_out, 0, sizeof(*response_out)); + response_out->struct_size = (uint32_t)sizeof(*response_out); + response_out->agreed_major = NXMEM_ABI_VERSION_MAJOR; + /* This plugin uses only the baseline prefix, so it agrees to the baseline + regardless of how new the host is. */ + response_out->agreed_minor = NXMEM_ABI_VERSION_MINOR_BASELINE; + response_out->plugin_range.major_min = NXMEM_ABI_VERSION_MAJOR; + response_out->plugin_range.major_max = NXMEM_ABI_VERSION_MAJOR; + response_out->plugin_range.minor_min = NXMEM_ABI_VERSION_MINOR_BASELINE; + response_out->plugin_range.minor_max = NXMEM_ABI_VERSION_MINOR_BASELINE; + /* Never advertise a capability the host did not offer. */ + response_out->capability_flags = + request->host_capability_flags & (uint64_t)NXMEM_CAP_ALLOCATOR; + if (response_out->capability_flags == 0) { + return nxmem_status(NXMEM_STATUS_UNSUPPORTED_CAPABILITY, + "minimal: the host did not offer plain allocation"); + } + return nxmem_ok(); +} + +NxmemStatus +NxmemCreateAllocatorFactories(const NxmemAllocatorFactoryVtable **out_factories, + uint64_t max_factories, uint64_t *out_count) { + if (out_factories == NULL || out_count == NULL) { + return nxmem_status(NXMEM_STATUS_INVALID_ARGUMENT, + "minimal: null pointer in create_allocator_factories"); + } + if (max_factories < 1) { + return nxmem_status(NXMEM_STATUS_INVALID_ARGUMENT, + "minimal: no room for even one factory"); + } + out_factories[0] = &kFactory; + *out_count = 1; + return nxmem_ok(); +} + +NxmemStatus NxmemQueryUnloadReadiness(NxmemUnloadReport *report_out) { + if (report_out == NULL) { + return nxmem_status(NXMEM_STATUS_INVALID_ARGUMENT, + "minimal: null report pointer"); + } + memset(report_out, 0, sizeof(*report_out)); + report_out->struct_size = (uint32_t)sizeof(*report_out); + report_out->live_allocators = g_live_allocators; + report_out->live_allocations = g_live_allocations; + return nxmem_ok(); +} diff --git a/crates/onnx-runtime-memory-abi/include/nxmem_memory_abi.h b/crates/onnx-runtime-memory-abi/include/nxmem_memory_abi.h new file mode 100644 index 0000000000..3e36c0ba4c --- /dev/null +++ b/crates/onnx-runtime-memory-abi/include/nxmem_memory_abi.h @@ -0,0 +1,645 @@ +/* + * nxmem — the stable dynamic-plugin memory ABI. + * + * A memory plugin is a shared library that publishes one or more allocation + * *mechanisms* to an ONNX GenAI host. This header is the whole contract: a + * plugin needs nothing else from the workspace, does not link against it, and + * may be written in C. + * + * ─── What crosses the boundary ────────────────────────────────────────────── + * + * Only `struct`s declared in this header, and only by pointer or by value. + * Nothing else may cross: no C++ exception, no Rust panic, no ownership of an + * allocator's heap, no object whose layout is defined by a language runtime. + * + * - Every struct begins with `uint32_t struct_size`, at offset 0, in every + * version, forever. That field is how a reader knows which prefix of the + * struct the writer actually filled in. + * - Every struct that a version can grow also carries `uint32_t abi_minor` + * or is reached through one that does. + * - Errors are `NxmemStatus` values: a stable numeric code plus an inline + * message buffer. No error ever propagates as an unwind. + * + * ─── Versioning ───────────────────────────────────────────────────────────── + * + * `major` is a hard compatibility gate: participants that disagree on it do + * not talk. `minor` only ever *appends* fields to the end of a struct, so a + * larger struct is always a strict superset of a smaller one. + * + * The host calls `NxmemNegotiate` first. The agreed minor is a *ceiling*. + * Each vtable then declares, in its own `abi_minor`, the level it really + * implements, so one plugin may ship a current mechanism beside one still + * written to the baseline. A reader: + * + * 1. reads `struct_size` at offset 0 and `abi_minor` at offset 4; + * 2. rejects a `struct_size` smaller than the level the writer *claims* — + * that writer contradicts itself; + * 3. clamps the effective level to `min(declared, negotiated)`; + * 4. copies only `min(struct_size, sizeof(local))` bytes into a zeroed + * local, so unknown trailing bytes are ignored and slots the writer does + * not have read back as NULL; + * 5. never calls a slot introduced above the effective level. + * + * A NULL optional slot means "this capability is absent" and must surface as + * NXMEM_STATUS_UNSUPPORTED_CAPABILITY. It must never behave as a no-op that + * silently succeeds. + * + * ─── Ownership ────────────────────────────────────────────────────────────── + * + * object created by released by + * ───────────────────────── ──────────────────── ───────────────────────── + * factory vtable NxmemCreateAllocator host, via factory + * Factories `release`, exactly once + * allocator vtable factory `open_alloc` host, via allocator + * `release`; `retain` adds + * a reference + * virtual-backing vtable plugin, owned by the the allocator; the host + * shared-mapping vtable allocator never releases it + * separately + * allocation allocator `allocate` `deallocate` or + * `release_allocation` + * shared prefix `create_shared_prefix` matching + * `release_shared_prefix` + * host callbacks host host, after the last + * allocator and the last + * queued release retire + * + * Every pointer passed *into* a call is borrowed for that call only, unless + * this header says otherwise. The two exceptions are stated where they occur: + * `NxmemOpenRequest::callbacks` is borrowed for the allocator's whole + * lifetime, and a factory's `name` must outlive the factory. + * + * ─── Threading, re-entrancy, and the rule that matters most ───────────────── + * + * Every slot may be called concurrently from any thread. A plugin must do its + * own locking. + * + * A plugin MUST NOT block indefinitely inside any slot, and MUST NOT hold one + * of its own locks across a call back into the host. The host's governance + * layer serialises accounting behind locks, and a host callback may re-enter + * this plugin on the same thread. A plugin that holds a lock across + * `request_reclaim` can therefore deadlock against itself. The host reciprocates + * this rule: it never calls into a plugin while holding a governance lock. + * + * `drain_releases` invokes the host's `release_completed` callback. Retire + * tickets in enqueue order. Take the batch under your lock, drop the lock, and + * only then call the host. + * + * ─── Deferred release and unload ──────────────────────────────────────────── + * + * `enqueue_release` hands an allocation to the plugin without freeing it yet + * and yields a ticket. Until that ticket retires through `drain_releases`, the + * allocation counts as live: it pins the allocator, and the allocator pins the + * module. + * + * The host calls `NxmemQueryUnloadReadiness` before unloading. While any count + * in `NxmemUnloadReport` is non-zero the host refuses or defers the unload. A + * plugin must count conservatively — reporting zero while anything is + * reachable invites the host to unmap code that is about to run. + * + * ─── Required exports ─────────────────────────────────────────────────────── + * + * NxmemNegotiate + * NxmemCreateAllocatorFactories + * NxmemQueryUnloadReadiness + * + * All three are required. A library missing any one of them is refused at load + * — without the third, unload could not be gated at all. + * + * Each struct below carries a machine-readable layout annotation naming its + * size in bytes on a 64-bit target. Those annotations are machine-checked against the Rust definitions by + * `header_layout_matches_the_rust_definitions` in this crate. Update both + * sides together or the build fails. + */ + +#ifndef NXMEM_MEMORY_ABI_H +#define NXMEM_MEMORY_ABI_H + +#include +#include + +#ifdef __cplusplus +extern "C" { +#endif + +/* ─── versions ─────────────────────────────────────────────────────────── */ + +#define NXMEM_ABI_VERSION_MAJOR 1u +#define NXMEM_ABI_VERSION_MINOR 1u +#define NXMEM_ABI_VERSION_MINOR_BASELINE 0u + +/* ─── capabilities ─────────────────────────────────────────────────────── */ + +/** Plain allocate/free. Every mechanism must offer this. */ +#define NXMEM_CAP_ALLOCATOR (1ull << 0) +/** Reserve address space now, commit physical pages later. */ +#define NXMEM_CAP_VIRTUAL_BACKING (1ull << 1) +/** One physical prefix mapped into several allocations. */ +#define NXMEM_CAP_SHARED_MAPPING (1ull << 2) +/** Stream-ordered release that retires later. */ +#define NXMEM_CAP_DEFERRED_RELEASE (1ull << 3) +/** Release reporting a structured outcome. Added at minor 1. */ +#define NXMEM_CAP_STRUCTURED_RELEASE (1ull << 4) + +/* ─── status codes ─────────────────────────────────────────────────────── */ + +#define NXMEM_STATUS_OK 0u +#define NXMEM_STATUS_VERSION_MISMATCH 1u +#define NXMEM_STATUS_SHORT_STRUCT 2u +#define NXMEM_STATUS_UNSUPPORTED_CAPABILITY 3u +#define NXMEM_STATUS_INVALID_ARGUMENT 4u +#define NXMEM_STATUS_INTERNAL_ERROR 5u +#define NXMEM_STATUS_NOT_IMPLEMENTED 6u +#define NXMEM_STATUS_DEVICE_ERROR 7u +#define NXMEM_STATUS_OUT_OF_MEMORY 8u +#define NXMEM_STATUS_WRONG_DEVICE 9u +#define NXMEM_STATUS_WRONG_MECHANISM 10u +#define NXMEM_STATUS_UNKNOWN_ALLOCATION 11u +#define NXMEM_STATUS_RELEASE_QUARANTINED 12u +#define NXMEM_STATUS_CALLBACK_FAILED 13u +#define NXMEM_STATUS_BUSY 14u + +/** Message bytes excluding the NUL terminator. */ +#define NXMEM_STATUS_MESSAGE_MAX 255u +/** Size of the inline message buffer, including the NUL terminator. */ +#define NXMEM_STATUS_MESSAGE_BUF (NXMEM_STATUS_MESSAGE_MAX + 1u) + +/** + * A status crossing the boundary. + * + * By value, with an inline buffer: no allocator is shared, so neither side can + * free the other's heap. `message` is NUL-terminated UTF-8; only the first + * `message_len` bytes are meaningful. + * + * NXMEM_LAYOUT: NxmemStatus size=264 + */ +typedef struct NxmemStatus { + uint32_t code; + uint32_t message_len; + uint8_t message[NXMEM_STATUS_MESSAGE_BUF]; +} NxmemStatus; + +/* ─── device identity ──────────────────────────────────────────────────── */ + +#define NXMEM_TIER_DEVICE 0u +#define NXMEM_TIER_HOST 1u +#define NXMEM_TIER_DISK 2u + +/** + * Which memory a mechanism serves. + * + * Carried on allocations and on every capability request so a mechanism can + * reject an object that belongs to a different device or provider rather than + * quietly acting on it. + * + * NXMEM_LAYOUT: NxmemDeviceId size=8 + */ +typedef struct NxmemDeviceId { + uint32_t tier; + uint32_t index; +} NxmemDeviceId; + +/* ─── plain records ────────────────────────────────────────────────────── */ + +/** + * An allocation, named well enough to reject cross-provider misuse. + * + * `mechanism_id` and `device` must match the mechanism being called. + * `allocation_id` is a host-assigned monotonic id. It is deliberately not + * derived from `ptr`: an address can be reused, an id cannot. + * + * NXMEM_LAYOUT: NxmemAllocation size=56 + */ +typedef struct NxmemAllocation { + uint32_t struct_size; + uint32_t reserved; + uint64_t mechanism_id; + uint64_t allocation_id; + NxmemDeviceId device; + uint8_t *ptr; + uint64_t bytes; + uint64_t align; +} NxmemAllocation; + +/** + * Half-open byte span inside an allocation. + * + * NXMEM_LAYOUT: NxmemByteRange size=16 + */ +typedef struct NxmemByteRange { + uint64_t offset; + uint64_t bytes; +} NxmemByteRange; + +/** + * A request to allocate. + * + * `committed_ranges` is borrowed for the call. An empty list on a lazy + * mechanism means "reserve address space, commit nothing". + * + * NXMEM_LAYOUT: NxmemAllocRequest size=64 + */ +typedef struct NxmemAllocRequest { + uint32_t struct_size; + uint32_t reserved; + uint64_t mechanism_id; + uint64_t allocation_id; + NxmemDeviceId device; + uint64_t bytes; + uint64_t align; + const NxmemByteRange *committed_ranges; + uint64_t committed_range_count; +} NxmemAllocRequest; + +/** + * What an allocation produced. + * + * `owned_bytes` is what the mechanism charged; `mapped_bytes` is what it + * physically mapped. They differ for lazy mechanisms and the host keeps the + * two accounting axes apart. + * + * NXMEM_LAYOUT: NxmemAllocResult size=32 + */ +typedef struct NxmemAllocResult { + uint32_t struct_size; + uint32_t reserved; + uint8_t *ptr; + uint64_t owned_bytes; + uint64_t mapped_bytes; +} NxmemAllocResult; + +/** + * One span of one allocation. + * + * NXMEM_LAYOUT: NxmemRangeRequest size=80 + */ +typedef struct NxmemRangeRequest { + uint32_t struct_size; + uint32_t reserved; + NxmemAllocation allocation; + NxmemByteRange range; +} NxmemRangeRequest; + +#define NXMEM_RELEASE_COMPLETE 0u +#define NXMEM_RELEASE_QUARANTINED 1u +#define NXMEM_RELEASE_FAILED 2u + +/** + * What a release actually did. + * + * The three states are not interchangeable: + * + * COMPLETE the memory is gone and `unmapped_bytes` may be credited back. + * QUARANTINED the plugin still owns `residual_owned_bytes`. The address must + * never be reissued and the residue must not be refunded. + * FAILED nothing was mutated. The allocation is exactly as live as the + * caller left it, and may be released again. + * + * A state code a reader does not recognise must be treated as quarantine: it + * is the only interpretation that cannot corrupt memory or accounting. + * + * NXMEM_LAYOUT: NxmemReleaseOutcome size=296 + */ +typedef struct NxmemReleaseOutcome { + uint32_t struct_size; + uint32_t state; + uint64_t allocation_bytes; + uint64_t unmapped_bytes; + uint64_t residual_owned_bytes; + NxmemStatus failure; +} NxmemReleaseOutcome; + +/** + * One retired deferred release, reported to the host. + * + * NXMEM_LAYOUT: NxmemReleaseCompletion size=328 + */ +typedef struct NxmemReleaseCompletion { + uint32_t struct_size; + uint32_t reserved; + uint64_t ticket; + uint64_t mechanism_id; + uint64_t allocation_id; + NxmemReleaseOutcome outcome; +} NxmemReleaseCompletion; + +/** + * A plugin asking the host to make room. + * + * NXMEM_LAYOUT: NxmemReclaimRequest size=32 + */ +typedef struct NxmemReclaimRequest { + uint32_t struct_size; + uint32_t reserved; + uint64_t mechanism_id; + NxmemDeviceId device; + uint64_t bytes; +} NxmemReclaimRequest; + +/** + * What the plugin still owns. + * + * The host refuses or defers unload while any count is non-zero. + * + * NXMEM_LAYOUT: NxmemUnloadReport size=48 + */ +typedef struct NxmemUnloadReport { + uint32_t struct_size; + uint32_t reserved; + uint64_t live_allocators; + uint64_t live_allocations; + uint64_t live_views; + uint64_t live_capabilities; + uint64_t queued_releases; +} NxmemUnloadReport; + +/** + * A shared physical prefix. + * + * NXMEM_LAYOUT: NxmemSharedPrefixHandle size=64 + */ +typedef struct NxmemSharedPrefixHandle { + uint32_t struct_size; + uint32_t reserved; + uint64_t mechanism_id; + uint64_t handle; + NxmemDeviceId device; + uint64_t device_ptr; + uint64_t committed_physical_bytes; + uint64_t mapped_bytes; + uint64_t requested_bytes; +} NxmemSharedPrefixHandle; + +/** + * Map a shared prefix into one allocation. + * + * NXMEM_LAYOUT: NxmemSharedPrefixCommitRequest size=136 + */ +typedef struct NxmemSharedPrefixCommitRequest { + uint32_t struct_size; + uint32_t reserved; + NxmemSharedPrefixHandle prefix; + NxmemAllocation allocation; + uint64_t byte_offset; +} NxmemSharedPrefixCommitRequest; + +/** + * What a shared-prefix commit cost. + * + * `additional_owned_bytes` is charged once, for the first mapping only; later + * mappings of the same prefix report zero. + * + * NXMEM_LAYOUT: NxmemSharedPrefixCommitInfo size=32 + */ +typedef struct NxmemSharedPrefixCommitInfo { + uint32_t struct_size; + uint32_t reserved; + uint64_t additional_owned_bytes; + uint64_t newly_mapped_bytes; + uint64_t granules; +} NxmemSharedPrefixCommitInfo; + +/* ─── host callbacks ───────────────────────────────────────────────────── */ + +/** + * What the host lets a plugin call back into. + * + * Borrowed for the whole lifetime of every allocator opened with it, and for + * every queued release naming one of those allocators. The host keeps it alive + * until the last of them retires. A NULL slot means the host does not offer + * that callback; the plugin must cope rather than require it. + * + * A plugin must not hold any of its own locks across either callback. + * + * NXMEM_LAYOUT: NxmemHostCallbacks size=32 + */ +typedef struct NxmemHostCallbacks { + uint32_t struct_size; + uint32_t abi_minor; + void *host_ctx; + NxmemStatus (*request_reclaim)(void *host_ctx, + const NxmemReclaimRequest *request, + uint64_t *reclaimed_out); + NxmemStatus (*release_completed)(void *host_ctx, + const NxmemReleaseCompletion *completion); +} NxmemHostCallbacks; + +/* ─── vtables ──────────────────────────────────────────────────────────── */ + +struct NxmemVirtualBackingVtable; +struct NxmemSharedMappingVtable; + +/** + * What the host asks for when opening an allocator. + * + * `abi_minor` is the negotiated ceiling. `callbacks` is borrowed for the + * allocator's whole lifetime, not just this call, and may be NULL. + * + * NXMEM_LAYOUT: NxmemOpenRequest size=32 + */ +typedef struct NxmemOpenRequest { + uint32_t struct_size; + uint32_t abi_minor; + NxmemDeviceId device; + uint64_t required_capability_flags; + const NxmemHostCallbacks *callbacks; +} NxmemOpenRequest; + +/** + * Reserve address space now and commit pages later. + * + * Reached through `NxmemAllocatorVtable::virtual_backing`. Owned by the + * allocator that published it and released with it — never separately. + * + * NXMEM_LAYOUT: NxmemVirtualBackingVtable size=72 + */ +typedef struct NxmemVirtualBackingVtable { + uint32_t struct_size; + uint32_t abi_minor; + uint64_t mechanism_id; + void *ctx; + NxmemStatus (*allocate_committed)(void *ctx, const NxmemAllocRequest *request, + NxmemAllocResult *result_out); + NxmemStatus (*commit_range)(void *ctx, const NxmemRangeRequest *request); + NxmemStatus (*decommit_range)(void *ctx, const NxmemRangeRequest *request, + uint64_t *unmapped_out); + NxmemStatus (*mapped_bytes_for_ranges)(void *ctx, + const NxmemRangeRequest *requests, + uint64_t count, uint64_t *mapped_out); + NxmemStatus (*mapped_bytes_for_allocation)(void *ctx, + const NxmemAllocRequest *request, + uint64_t *mapped_out); + NxmemStatus (*committed_bytes)(void *ctx, const NxmemAllocation *allocation, + uint64_t *committed_out); +} NxmemVirtualBackingVtable; + +/** + * One physical prefix mapped into several allocations. + * + * Owned by the allocator that published it, exactly as virtual backing is. + * + * NXMEM_LAYOUT: NxmemSharedMappingVtable size=64 + */ +typedef struct NxmemSharedMappingVtable { + uint32_t struct_size; + uint32_t abi_minor; + uint64_t mechanism_id; + void *ctx; + NxmemStatus (*create_shared_prefix)(void *ctx, uint64_t mechanism_id, + uint64_t bytes, + NxmemSharedPrefixHandle *handle_out); + NxmemStatus (*retain_shared_prefix)(void *ctx, + const NxmemSharedPrefixHandle *handle); + NxmemStatus (*release_shared_prefix)(void *ctx, + const NxmemSharedPrefixHandle *handle); + NxmemStatus (*incremental_owned_bytes)(void *ctx, + const NxmemSharedPrefixHandle *handle, + uint64_t *bytes_out); + NxmemStatus (*commit_shared_prefix)( + void *ctx, const NxmemSharedPrefixCommitRequest *request, + NxmemSharedPrefixCommitInfo *info_out); +} NxmemSharedMappingVtable; + +/** + * One opened allocation mechanism. + * + * `allocate`, `deallocate`, `retain`, and `release` are required at every + * level; a vtable missing one is refused. Everything else is optional and NULL + * when unsupported. `release_allocation` exists only at minor 1 and above and + * only when NXMEM_CAP_STRUCTURED_RELEASE is set. + * + * `name` must outlive the allocator. `ctx` is opaque to the host and is passed + * back unchanged to every slot. + * + * NXMEM_LAYOUT: NxmemAllocatorVtable size=128 + */ +typedef struct NxmemAllocatorVtable { + uint32_t struct_size; + uint32_t abi_minor; + uint64_t mechanism_id; + NxmemDeviceId device; + uint64_t capability_flags; + const uint8_t *name; + void *ctx; + + /* required at minor 0 */ + NxmemStatus (*allocate)(void *ctx, const NxmemAllocRequest *request, + NxmemAllocResult *result_out); + NxmemStatus (*deallocate)(void *ctx, const NxmemAllocation *allocation, + uint64_t *unmapped_bytes_out); + void (*retain)(void *ctx); + void (*release)(void *ctx); + + /* optional at minor 0 */ + const struct NxmemVirtualBackingVtable *virtual_backing; + const struct NxmemSharedMappingVtable *shared_mapping; + NxmemStatus (*enqueue_release)(void *ctx, const NxmemAllocation *allocation, + uint64_t *ticket_out); + NxmemStatus (*drain_releases)(void *ctx, uint64_t max, uint64_t *retired_out); + NxmemStatus (*pending_release_count)(void *ctx, uint64_t *count_out); + + /* added at minor 1 */ + NxmemStatus (*release_allocation)(void *ctx, + const NxmemAllocation *allocation, + NxmemReleaseOutcome *outcome_out); +} NxmemAllocatorVtable; + +/** + * A named, device-scoped mechanism the host may open. + * + * Owned by the host once `NxmemCreateAllocatorFactories` returns it: the host + * calls `release` exactly once, after the last allocator opened from it has + * been released. `name` must outlive the factory. + * + * NXMEM_LAYOUT: NxmemAllocatorFactoryVtable size=56 + */ +typedef struct NxmemAllocatorFactoryVtable { + uint32_t struct_size; + uint32_t abi_minor; + const uint8_t *name; + NxmemDeviceId device; + uint64_t capability_flags; + void *ctx; + NxmemStatus (*open_allocator)(void *ctx, const NxmemOpenRequest *request, + const NxmemAllocatorVtable **allocator_out); + void (*release)(void *ctx); +} NxmemAllocatorFactoryVtable; + +/* ─── negotiation ──────────────────────────────────────────────────────── */ + +/** + * An inclusive version range a participant supports. + * + * NXMEM_LAYOUT: NxmemVersionRange size=16 + */ +typedef struct NxmemVersionRange { + uint32_t major_min; + uint32_t major_max; + uint32_t minor_min; + uint32_t minor_max; +} NxmemVersionRange; + +/** + * What the host offers. + * + * NXMEM_LAYOUT: NxmemNegotiateRequest size=32 + */ +typedef struct NxmemNegotiateRequest { + uint32_t struct_size; + uint32_t reserved; + NxmemVersionRange host_range; + uint64_t host_capability_flags; +} NxmemNegotiateRequest; + +/** + * What the plugin agrees to. + * + * `capability_flags` must be a subset of what the host offered: a plugin may + * not advertise a capability the host never named, and the agreed minor may + * not exceed either side's range. + * + * NXMEM_LAYOUT: NxmemNegotiateResponse size=40 + */ +typedef struct NxmemNegotiateResponse { + uint32_t struct_size; + uint32_t agreed_major; + uint32_t agreed_minor; + uint32_t reserved; + NxmemVersionRange plugin_range; + uint64_t capability_flags; +} NxmemNegotiateResponse; + +/* ─── required exports ─────────────────────────────────────────────────── */ + +/** + * Agree a version and capability set. Called once, before anything else. + * + * Neither pointer may be retained past the call. + */ +NxmemStatus NxmemNegotiate(const NxmemNegotiateRequest *request, + NxmemNegotiateResponse *response_out); + +/** + * Publish the mechanisms this module offers. + * + * Writes at most `max_factories` pointers and the real count to `out_count`. + * Each pointer becomes owned by the host, which releases it exactly once. + * Writing more than `max_factories` is a contract violation. + */ +NxmemStatus +NxmemCreateAllocatorFactories(const NxmemAllocatorFactoryVtable **out_factories, + uint64_t max_factories, uint64_t *out_count); + +/** + * Report what this module still owns, so the host can gate unload. + * + * Count conservatively. A module that reports zero while anything is still + * reachable will be unmapped while its own code is about to run. + */ +NxmemStatus NxmemQueryUnloadReadiness(NxmemUnloadReport *report_out); + +#ifdef __cplusplus +} /* extern "C" */ +#endif + +#endif /* NXMEM_MEMORY_ABI_H */ diff --git a/crates/onnx-runtime-memory-abi/src/lib.rs b/crates/onnx-runtime-memory-abi/src/lib.rs index c43a4f415c..8556163674 100644 --- a/crates/onnx-runtime-memory-abi/src/lib.rs +++ b/crates/onnx-runtime-memory-abi/src/lib.rs @@ -100,6 +100,12 @@ //! exposed here and never will be. Nothing that predates this contract is //! promised ABI compatibility. +// `NxmemStatus` carries its message inline, in a fixed 256-byte buffer, so that +// no heap allocation is ever owned across the dynamic-library boundary. That +// makes every `Result<_, NxmemStatus>` "large" by clippy's measure. Boxing the +// error, which is what the lint asks for, would reintroduce exactly the +// cross-allocator ownership this ABI exists to forbid. +#![allow(clippy::result_large_err)] #![deny(clippy::mem_forget)] pub mod status; @@ -108,23 +114,22 @@ pub mod version; pub mod vtable; pub use status::{ - catch_status_panic, catch_void_panic, NxmemStatus, NxmemStatusCode, NXMEM_STATUS_MESSAGE_BUF, - NXMEM_STATUS_MESSAGE_MAX, + NXMEM_STATUS_MESSAGE_BUF, NXMEM_STATUS_MESSAGE_MAX, NxmemStatus, NxmemStatusCode, + catch_status_panic, catch_void_panic, }; pub use types::{ - check_identity, NxmemAllocRequest, NxmemAllocResult, NxmemAllocation, NxmemByteRange, - NxmemDeviceId, NxmemHostCallbacks, NxmemRangeRequest, NxmemReclaimRequest, - NxmemReleaseCompletion, NxmemReleaseOutcome, NxmemSharedPrefixCommitInfo, - NxmemSharedPrefixCommitRequest, NxmemSharedPrefixHandle, NxmemUnloadReport, NXMEM_RELEASE_COMPLETE, NXMEM_RELEASE_FAILED, NXMEM_RELEASE_QUARANTINED, NXMEM_TIER_DEVICE, - NXMEM_TIER_DISK, NXMEM_TIER_HOST, + NXMEM_TIER_DISK, NXMEM_TIER_HOST, NxmemAllocRequest, NxmemAllocResult, NxmemAllocation, + NxmemByteRange, NxmemDeviceId, NxmemHostCallbacks, NxmemRangeRequest, NxmemReclaimRequest, + NxmemReleaseCompletion, NxmemReleaseOutcome, NxmemSharedPrefixCommitInfo, + NxmemSharedPrefixCommitRequest, NxmemSharedPrefixHandle, NxmemUnloadReport, check_identity, }; pub use version::{ - capability_min_minor, describe_capabilities, negotiate, negotiate_as, validate_negotiation, - NegotiationRejection, NxmemNegotiateRequest, NxmemNegotiateResponse, NxmemVersionRange, NXMEM_ABI_VERSION_MAJOR, NXMEM_ABI_VERSION_MINOR, NXMEM_ABI_VERSION_MINOR_BASELINE, NXMEM_CAP_ALLOCATOR, NXMEM_CAP_DEFERRED_RELEASE, NXMEM_CAP_KNOWN_MASK, NXMEM_CAP_SHARED_MAPPING, NXMEM_CAP_STRUCTURED_RELEASE, NXMEM_CAP_VIRTUAL_BACKING, + NegotiationRejection, NxmemNegotiateRequest, NxmemNegotiateResponse, NxmemVersionRange, + capability_min_minor, describe_capabilities, negotiate, negotiate_as, validate_negotiation, }; pub use vtable::{ NxmemAllocatorFactoryVtable, NxmemAllocatorVtable, NxmemOpenRequest, NxmemSharedMappingVtable, diff --git a/crates/onnx-runtime-memory-abi/src/status.rs b/crates/onnx-runtime-memory-abi/src/status.rs index 0263d56865..a71a5cba2f 100644 --- a/crates/onnx-runtime-memory-abi/src/status.rs +++ b/crates/onnx-runtime-memory-abi/src/status.rs @@ -321,7 +321,11 @@ mod tests { for code in NxmemStatusCode::ALL { assert_eq!(NxmemStatusCode::from_u32(code as u32), Some(code)); assert!(!code.name().is_empty()); - assert!(!seen.contains(&code.name()), "duplicate name {}", code.name()); + assert!( + !seen.contains(&code.name()), + "duplicate name {}", + code.name() + ); seen.push(code.name()); } assert_eq!(seen.len(), NxmemStatusCode::ALL.len()); @@ -342,7 +346,7 @@ mod tests { // truncation would split one and make `message_str` return None. let message = "\u{4e2d}".repeat(200); let status = NxmemStatus::with_message(NxmemStatusCode::DeviceError, &message); - assert_eq!(status.message_len, 255 - (255 % 3)); + assert_eq!(status.message_len, NXMEM_STATUS_MESSAGE_MAX as u32); let recovered = status.message_str().expect("valid utf-8 after truncation"); assert!(recovered.chars().all(|c| c == '\u{4e2d}')); } diff --git a/crates/onnx-runtime-memory-abi/src/types.rs b/crates/onnx-runtime-memory-abi/src/types.rs index 7d6d8fb440..ebb24623a0 100644 --- a/crates/onnx-runtime-memory-abi/src/types.rs +++ b/crates/onnx-runtime-memory-abi/src/types.rs @@ -638,7 +638,10 @@ pub fn check_identity( &format!( "nxmem: object belongs to device (tier {}, index {}) but device (tier {}, index \ {}) was addressed", - actual_device.tier, actual_device.index, expected_device.tier, expected_device.index + actual_device.tier, + actual_device.index, + expected_device.tier, + expected_device.index ), )); } @@ -738,13 +741,6 @@ mod tests { #[test] fn unknown_tier_codes_are_not_guessed() { assert_eq!(NxmemDeviceId::HOST.tier_code(), Some(NXMEM_TIER_HOST)); - assert_eq!( - NxmemDeviceId { - tier: 77, - index: 0 - } - .tier_code(), - None - ); + assert_eq!(NxmemDeviceId { tier: 77, index: 0 }.tier_code(), None); } } diff --git a/crates/onnx-runtime-memory-abi/src/version.rs b/crates/onnx-runtime-memory-abi/src/version.rs index bcc72fb6f3..f0aaa6ac57 100644 --- a/crates/onnx-runtime-memory-abi/src/version.rs +++ b/crates/onnx-runtime-memory-abi/src/version.rs @@ -431,7 +431,8 @@ pub fn validate_negotiation( NXMEM_CAP_DEFERRED_RELEASE, NXMEM_CAP_STRUCTURED_RELEASE, ] { - if response.capability_flags & flag != 0 && response.agreed_minor < capability_min_minor(flag) + if response.capability_flags & flag != 0 + && response.agreed_minor < capability_min_minor(flag) { return reject(format!( "plugin advertises capability {} at nxmem minor {} but that capability is only \ diff --git a/crates/onnx-runtime-memory-abi/src/vtable.rs b/crates/onnx-runtime-memory-abi/src/vtable.rs index 3bc16b898e..9ac1ecf118 100644 --- a/crates/onnx-runtime-memory-abi/src/vtable.rs +++ b/crates/onnx-runtime-memory-abi/src/vtable.rs @@ -97,7 +97,7 @@ fn check_vtable_ptr(ptr: *const T, what: &str) -> Result<(), NxmemStatus> { &format!("nxmem: {what} vtable pointer is null"), )); } - if (ptr as usize) % core::mem::align_of::() != 0 { + if !(ptr as usize).is_multiple_of(core::mem::align_of::()) { return Err(NxmemStatus::with_message( NxmemStatusCode::InvalidArgument, &format!( @@ -123,11 +123,7 @@ unsafe fn copy_prefix(ptr: *const T, declared_size: usize) -> T { // valid for `readable` bytes by this function's contract; the regions // cannot overlap because `out` is a local. unsafe { - core::ptr::copy_nonoverlapping( - ptr as *const u8, - out.as_mut_ptr() as *mut u8, - readable, - ); + core::ptr::copy_nonoverlapping(ptr as *const u8, out.as_mut_ptr() as *mut u8, readable); out.assume_init() } } @@ -292,7 +288,10 @@ impl NxmemAllocatorVtable { /// /// `ptr` must either be null (rejected) or point to a live vtable whose /// first `struct_size` bytes are readable for the duration of the call. - pub unsafe fn read_prefix(ptr: *const Self, negotiated_minor: u32) -> Result { + pub unsafe fn read_prefix( + ptr: *const Self, + negotiated_minor: u32, + ) -> Result { check_vtable_ptr(ptr, "allocator")?; // SAFETY: `ptr` was proved non-null and aligned; `struct_size` is the // first field of every version of this struct. @@ -401,8 +400,9 @@ pub struct NxmemVirtualBackingVtable { ) -> NxmemStatus, >, /// Commit one span of an existing allocation. - pub commit_range: - Option NxmemStatus>, + pub commit_range: Option< + unsafe extern "C" fn(ctx: *mut c_void, request: *const NxmemRangeRequest) -> NxmemStatus, + >, /// Decommit one span, writing the bytes actually unmapped. pub decommit_range: Option< unsafe extern "C" fn( @@ -697,7 +697,10 @@ impl NxmemAllocatorFactoryVtable { /// # Safety /// /// Same contract as [`NxmemAllocatorVtable::read_prefix`]. - pub unsafe fn read_prefix(ptr: *const Self, negotiated_minor: u32) -> Result { + pub unsafe fn read_prefix( + ptr: *const Self, + negotiated_minor: u32, + ) -> Result { check_vtable_ptr(ptr, "allocator-factory")?; // SAFETY: `ptr` was proved non-null and aligned. let declared_size = unsafe { read_struct_size(ptr) } as usize; @@ -776,9 +779,12 @@ mod tests { #[test] fn the_minor_1_prefix_is_larger_than_the_minor_0_prefix() { - assert!( + // Compile-time, because a regression here is a layout bug that must + // not be allowed to link, let alone run. + const _: () = assert!( NxmemAllocatorVtable::MIN_STRUCT_SIZE_MINOR_1 - > NxmemAllocatorVtable::MIN_STRUCT_SIZE_MINOR_0 + > NxmemAllocatorVtable::MIN_STRUCT_SIZE_MINOR_0, + "minor 1 must require a longer prefix than minor 0" ); assert_eq!( NxmemAllocatorVtable::MIN_STRUCT_SIZE_MINOR_1, @@ -897,8 +903,8 @@ mod tests { #[test] fn a_misaligned_vtable_pointer_is_refused() { let vtable = conforming_allocator(1); - let misaligned = (&vtable as *const NxmemAllocatorVtable as usize + 1) - as *const NxmemAllocatorVtable; + let misaligned = + (&vtable as *const NxmemAllocatorVtable as usize + 1) as *const NxmemAllocatorVtable; // SAFETY: `read_prefix` rejects on alignment before dereferencing. let status = unsafe { NxmemAllocatorVtable::read_prefix(misaligned, 1) }.expect_err("misaligned"); diff --git a/crates/onnx-runtime-memory-abi/tests/header_contract.rs b/crates/onnx-runtime-memory-abi/tests/header_contract.rs new file mode 100644 index 0000000000..12f6b9991f --- /dev/null +++ b/crates/onnx-runtime-memory-abi/tests/header_contract.rs @@ -0,0 +1,230 @@ +//! The public C header is part of the contract, so it is tested like one. +//! +//! Two things can silently drift apart here: the sizes recorded in the header +//! and the Rust definitions, and the header and a real C compiler's opinion of +//! it. Both are checked. +//! +//! Neither test is allowed to pass by doing nothing. The layout test fails if +//! it does not find every annotation it expects, and the compile test fails +//! loudly when no C compiler can be found rather than quietly skipping. + +use std::mem::size_of; +use std::path::{Path, PathBuf}; +use std::process::Command; + +use onnx_runtime_memory_abi::*; + +fn header_path() -> PathBuf { + Path::new(env!("CARGO_MANIFEST_DIR")).join("include/nxmem_memory_abi.h") +} + +fn example_path() -> PathBuf { + Path::new(env!("CARGO_MANIFEST_DIR")).join("examples/minimal_plugin.c") +} + +/// Every `NXMEM_LAYOUT: size=` annotation in the header. +fn annotated_sizes() -> Vec<(String, usize)> { + let header = std::fs::read_to_string(header_path()).expect("the public header is readable"); + header + .lines() + .filter_map(|line| { + let rest = line.split_once("NXMEM_LAYOUT:")?.1.trim(); + let (name, size) = rest.split_once("size=")?; + let size = size + .trim() + .parse::() + .unwrap_or_else(|_| panic!("malformed layout annotation: {line}")); + Some((name.trim().to_string(), size)) + }) + .collect() +} + +/// The Rust size of every type the header annotates. +/// +/// Adding a `#[repr(C)]` type to the header without adding it here fails the +/// test below, which is the point: the two lists must move together. +fn rust_sizes() -> Vec<(&'static str, usize)> { + vec![ + ("NxmemStatus", size_of::()), + ("NxmemDeviceId", size_of::()), + ("NxmemAllocation", size_of::()), + ("NxmemByteRange", size_of::()), + ("NxmemAllocRequest", size_of::()), + ("NxmemAllocResult", size_of::()), + ("NxmemRangeRequest", size_of::()), + ("NxmemReleaseOutcome", size_of::()), + ( + "NxmemReleaseCompletion", + size_of::(), + ), + ("NxmemReclaimRequest", size_of::()), + ("NxmemUnloadReport", size_of::()), + ( + "NxmemSharedPrefixHandle", + size_of::(), + ), + ( + "NxmemSharedPrefixCommitRequest", + size_of::(), + ), + ( + "NxmemSharedPrefixCommitInfo", + size_of::(), + ), + ("NxmemHostCallbacks", size_of::()), + ("NxmemOpenRequest", size_of::()), + ( + "NxmemVirtualBackingVtable", + size_of::(), + ), + ( + "NxmemSharedMappingVtable", + size_of::(), + ), + ("NxmemAllocatorVtable", size_of::()), + ( + "NxmemAllocatorFactoryVtable", + size_of::(), + ), + ("NxmemVersionRange", size_of::()), + ("NxmemNegotiateRequest", size_of::()), + ( + "NxmemNegotiateResponse", + size_of::(), + ), + ] +} + +/// The header's recorded sizes match the Rust definitions. +/// +/// Sizes are only equal like this on a 64-bit target, where the header's +/// pointer-sized members and Rust's agree; on anything else the annotations +/// describe a different machine and comparing them would be meaningless. +#[test] +#[cfg(target_pointer_width = "64")] +fn header_layout_matches_the_rust_definitions() { + let annotated = annotated_sizes(); + assert!( + annotated.len() >= 20, + "the header lost its layout annotations: found only {}", + annotated.len() + ); + + for (name, rust_size) in rust_sizes() { + let (_, header_size) = annotated + .iter() + .find(|(annotated_name, _)| annotated_name == name) + .unwrap_or_else(|| panic!("the header has no NXMEM_LAYOUT annotation for {name}")); + assert_eq!( + *header_size, rust_size, + "{name}: the header says {header_size} bytes, Rust says {rust_size}" + ); + } + + // And nothing is annotated that this test does not know about, so the + // header cannot grow a type that silently goes unchecked. + for (name, _) in &annotated { + assert!( + rust_sizes().iter().any(|(known, _)| known == name), + "the header annotates {name}, which this test does not check" + ); + } +} + +/// The first C compiler on this machine, if any. +fn find_cc() -> Option { + if let Ok(cc) = std::env::var("CC") + && !cc.is_empty() + { + return Some(PathBuf::from(cc)); + } + for candidate in ["cc", "clang", "gcc"] { + if let Ok(output) = Command::new(candidate).arg("--version").output() + && output.status.success() + { + return Some(PathBuf::from(candidate)); + } + } + None +} + +/// A real C compiler accepts the header and the example plugin. +/// +/// This is what backs the "one minimal plugin example that documents the +/// contract without requiring workspace linking" claim: the example is +/// compiled with nothing on the include path but the header itself. +#[test] +#[cfg(unix)] +fn nxmem_c_example_compiles() { + let cc = find_cc().expect( + "no C compiler found (tried $CC, cc, clang, gcc); \ + the public C header cannot be checked without one", + ); + let out = std::env::temp_dir().join(format!("nxmem_minimal_plugin_{}.o", std::process::id())); + let output = Command::new(&cc) + .arg("-c") + .arg("-std=c11") + .arg("-Wall") + .arg("-Wextra") + .arg("-Werror") + .arg("-fPIC") + .arg("-I") + .arg(header_path().parent().expect("include directory")) + .arg(example_path()) + .arg("-o") + .arg(&out) + .output() + .expect("the C compiler runs"); + + let stderr = String::from_utf8_lossy(&output.stderr).into_owned(); + let _ = std::fs::remove_file(&out); + assert!( + output.status.success(), + "the minimal C plugin must compile against the public header alone:\n{stderr}" + ); +} + +/// A C compiler agrees with Rust about every annotated struct size. +/// +/// The layout test above compares Rust against a *comment*. This one compares +/// Rust against the compiler that will actually build a plugin, which is the +/// opinion that decides whether a real plugin works. +#[test] +#[cfg(all(unix, target_pointer_width = "64"))] +fn a_c_compiler_agrees_with_the_rust_layouts() { + let cc = find_cc().expect( + "no C compiler found (tried $CC, cc, clang, gcc); \ + the public C header cannot be checked without one", + ); + + let mut source = String::from("#include \"nxmem_memory_abi.h\"\n"); + for (name, size) in rust_sizes() { + source.push_str(&format!( + "_Static_assert(sizeof({name}) == {size}, \"{name} disagrees with Rust\");\n" + )); + } + source.push_str("int nxmem_layout_probe(void) { return 0; }\n"); + + let dir = std::env::temp_dir().join(format!("nxmem_layout_{}", std::process::id())); + std::fs::create_dir_all(&dir).expect("temp directory"); + let probe = dir.join("probe.c"); + std::fs::write(&probe, source).expect("probe source is writable"); + + let output = Command::new(&cc) + .arg("-c") + .arg("-std=c11") + .arg("-I") + .arg(header_path().parent().expect("include directory")) + .arg(&probe) + .arg("-o") + .arg(dir.join("probe.o")) + .output() + .expect("the C compiler runs"); + + let stderr = String::from_utf8_lossy(&output.stderr).into_owned(); + let _ = std::fs::remove_dir_all(&dir); + assert!( + output.status.success(), + "a C compiler disagrees with the Rust struct layouts:\n{stderr}" + ); +} diff --git a/crates/onnx-runtime-memory-host/src/allocator.rs b/crates/onnx-runtime-memory-host/src/allocator.rs index bbbd9a883c..ca3474b4a3 100644 --- a/crates/onnx-runtime-memory-host/src/allocator.rs +++ b/crates/onnx-runtime-memory-host/src/allocator.rs @@ -33,8 +33,8 @@ use onnx_runtime_memory_abi::{ }; use onnx_runtime_memory_api::{ AllocationCommitRange, AllocationReleaseOutcome, DeviceAllocator, DeviceKey, MemoryError, - QuarantineReason, ReleaseAccounting, ResidualOwnership, SharedDevicePrefix, - SharedMapping, SharedPrefixCommitInfo, Tier, VirtualBacking, + QuarantineReason, ReleaseAccounting, ResidualOwnership, SharedDevicePrefix, SharedMapping, + SharedPrefixCommitInfo, Tier, VirtualBacking, }; use crate::error::{PluginError, status_to_memory_error}; @@ -301,7 +301,11 @@ impl AllocatorCore { } fn allocation_record(&self, ptr: NonNull) -> Option { - self.live.lock().ok()?.get(&(ptr.as_ptr() as usize)).copied() + self.live + .lock() + .ok()? + .get(&(ptr.as_ptr() as usize)) + .copied() } /// Remove a live record, **releasing the lock before any ABI call**. @@ -359,7 +363,7 @@ impl AllocatorCore { allocation must return a live, unique address", ), })?; - if ptr.as_ptr() as usize % align.max(1) != 0 { + if !(ptr.as_ptr() as usize).is_multiple_of(align.max(1)) { // The plugin broke the alignment contract. Hand the bytes straight // back rather than letting a misaligned pointer escape. let record = LiveAllocation { @@ -470,7 +474,8 @@ impl AllocatorCore { let mut outcome = NxmemReleaseOutcome::zeroed(); // SAFETY: `ctx` came from this vtable; both pointers address valid // locals that outlive the call. No host lock is held. - let status = unsafe { release(self.vtable.ctx, &raw const allocation, &raw mut outcome) }; + let status = + unsafe { release(self.vtable.ctx, &raw const allocation, &raw mut outcome) }; if !status.is_ok() { self.insert_allocation(ptr, record); return AllocationReleaseOutcome::failed(format!( @@ -682,8 +687,7 @@ impl PluginAllocator { core.module.release_queued(); // SAFETY: `ctx` came from this vtable; both pointers address valid // locals. No host lock is held. - let status = - unsafe { enqueue(core.vtable.ctx, &raw const allocation, &raw mut ticket) }; + let status = unsafe { enqueue(core.vtable.ctx, &raw const allocation, &raw mut ticket) }; if !status.is_ok() { core.module.release_retired(); core.insert_allocation(ptr, record); @@ -755,8 +759,7 @@ impl DeviceAllocator for PluginAllocator { // SAFETY: `ctx` came from this vtable; both pointers address valid // locals that outlive the call. No host lock is held, so the plugin may // call back into the host from inside this call. - let status = - unsafe { allocate(core.vtable.ctx, &raw const request, &raw mut result) }; + let status = unsafe { allocate(core.vtable.ctx, &raw const request, &raw mut result) }; if !status.is_ok() { return Err(status_to_memory_error( "allocate", @@ -802,11 +805,15 @@ impl DeviceAllocator for PluginAllocator { } fn as_virtual_backing(&self) -> Option<&dyn VirtualBacking> { - self.backing.as_ref().map(|backing| backing as &dyn VirtualBacking) + self.backing + .as_ref() + .map(|backing| backing as &dyn VirtualBacking) } fn as_shared_mapping(&self) -> Option<&dyn SharedMapping> { - self.shared.as_ref().map(|shared| shared as &dyn SharedMapping) + self.shared + .as_ref() + .map(|shared| shared as &dyn SharedMapping) } } @@ -892,9 +899,7 @@ impl VirtualBacking for PluginVirtualBacking { // SAFETY: `ctx` came from the backing vtable; `request` borrows `ranges` // which outlives the call, and `result` is a valid local. No host lock // is held. - let status = unsafe { - allocate(self.vtable().ctx, &raw const request, &raw mut result) - }; + let status = unsafe { allocate(self.vtable().ctx, &raw const request, &raw mut result) }; drop(ranges); if !status.is_ok() { return Err(status_to_memory_error( @@ -988,9 +993,7 @@ impl VirtualBacking for PluginVirtualBacking { return Err(MemoryError::AllocationFailed { tier: self.core.tier().name(), requested: bytes as u64, - reason: String::from( - "this mechanism provides no mapped_bytes_for_allocation slot", - ), + reason: String::from("this mechanism provides no mapped_bytes_for_allocation slot"), }); }; // A sizing query carries no allocation identity yet, so id zero is @@ -998,8 +1001,7 @@ impl VirtualBacking for PluginVirtualBacking { let request = self.core.alloc_request(0, bytes, align); let mut out = 0u64; // SAFETY: both pointers address valid locals; no host lock is held. - let status = - unsafe { mapped(self.vtable().ctx, &raw const request, &raw mut out) }; + let status = unsafe { mapped(self.vtable().ctx, &raw const request, &raw mut out) }; if !status.is_ok() { return Err(status_to_memory_error( "mapped_bytes_for_allocation", @@ -1029,8 +1031,7 @@ impl VirtualBacking for PluginVirtualBacking { let request = self.range_request(ptr, allocation_bytes, align, offset, bytes)?; let mut unmapped = 0u64; // SAFETY: both pointers address valid locals; no host lock is held. - let status = - unsafe { decommit(self.vtable().ctx, &raw const request, &raw mut unmapped) }; + let status = unsafe { decommit(self.vtable().ctx, &raw const request, &raw mut unmapped) }; if !status.is_ok() { return Err(status_to_memory_error( "decommit_range", @@ -1060,9 +1061,7 @@ impl VirtualBacking for PluginVirtualBacking { let allocation = self.core.abi_allocation(ptr, record); let mut out = 0u64; // SAFETY: both pointers address valid locals; no host lock is held. - let status = unsafe { - committed(self.vtable().ctx, &raw const allocation, &raw mut out) - }; + let status = unsafe { committed(self.vtable().ctx, &raw const allocation, &raw mut out) }; if !status.is_ok() { return 0; } @@ -1186,9 +1185,8 @@ impl SharedMapping for PluginSharedMapping { }; let mut out = 0u64; // SAFETY: both pointers address valid locals; no host lock is held. - let status = unsafe { - incremental(self.vtable().ctx, &raw const prefix.handle, &raw mut out) - }; + let status = + unsafe { incremental(self.vtable().ctx, &raw const prefix.handle, &raw mut out) }; if !status.is_ok() { return Err(status_to_memory_error( "incremental_owned_bytes", @@ -1215,16 +1213,17 @@ impl SharedMapping for PluginSharedMapping { reason: String::from("this mechanism provides no commit_shared_prefix slot"), }); }; - let record = self.core.allocation_record(ptr).ok_or_else(|| { - MemoryError::AllocationFailed { - tier: self.core.tier().name(), - requested: 0, - reason: format!( - "address {:p} is not a live allocation of this mechanism", - ptr.as_ptr() - ), - } - })?; + let record = + self.core + .allocation_record(ptr) + .ok_or_else(|| MemoryError::AllocationFailed { + tier: self.core.tier().name(), + requested: 0, + reason: format!( + "address {:p} is not a live allocation of this mechanism", + ptr.as_ptr() + ), + })?; if record.bytes != allocation_bytes { return Err(MemoryError::AllocationFailed { tier: self.core.tier().name(), @@ -1244,9 +1243,7 @@ impl SharedMapping for PluginSharedMapping { ); let mut info = NxmemSharedPrefixCommitInfo::zeroed(); // SAFETY: both pointers address valid locals; no host lock is held. - let status = unsafe { - commit(self.vtable().ctx, &raw const request, &raw mut info) - }; + let status = unsafe { commit(self.vtable().ctx, &raw const request, &raw mut info) }; if !status.is_ok() { return Err(status_to_memory_error( "commit_shared_prefix", @@ -1406,8 +1403,8 @@ pub(crate) fn open_allocator( // SAFETY: the plugin wrote this pointer in response to the call above. // `read_prefix` null-, alignment- and size-checks it before reading any // field, and copies rather than borrowing. - let vtable = unsafe { NxmemAllocatorVtable::read_prefix(raw_allocator, minor) } - .map_err(|status| { + let vtable = + unsafe { NxmemAllocatorVtable::read_prefix(raw_allocator, minor) }.map_err(|status| { // SAFETY: the vtable was refused, so nothing may be called through // it; `abandon_allocator` re-reads the prefix defensively and only // calls `release` when the struct is well-formed enough to locate @@ -1501,9 +1498,8 @@ pub(crate) fn open_allocator( // SAFETY: the contract requires `name` to stay valid until the final // `release`, which has not been called. - let name = unsafe { read_optional_c_string(vtable.name) }.unwrap_or_else(|| { - factory.name().to_string() - }); + let name = unsafe { read_optional_c_string(vtable.name) } + .unwrap_or_else(|| factory.name().to_string()); module.allocator_opened(); let core = Arc::new(AllocatorCore { @@ -1522,18 +1518,12 @@ pub(crate) fn open_allocator( module, }); - let backing_view = core - .backing - .is_some() - .then(|| PluginVirtualBacking { - core: Arc::clone(&core), - }); - let shared_view = core - .shared - .is_some() - .then(|| PluginSharedMapping { - core: Arc::clone(&core), - }); + let backing_view = core.backing.is_some().then(|| PluginVirtualBacking { + core: Arc::clone(&core), + }); + let shared_view = core.shared.is_some().then(|| PluginSharedMapping { + core: Arc::clone(&core), + }); Ok(PluginAllocator { backing: backing_view, diff --git a/crates/onnx-runtime-memory-host/src/error.rs b/crates/onnx-runtime-memory-host/src/error.rs index 89ba425ca7..2ee686174f 100644 --- a/crates/onnx-runtime-memory-host/src/error.rs +++ b/crates/onnx-runtime-memory-host/src/error.rs @@ -32,10 +32,16 @@ pub enum PluginError { Negotiation { path: String, reason: String }, /// The plugin returned a failing status. + /// + /// The status is boxed because it carries a 256-byte inline message + /// buffer, which is right on the wire — the buffer is what keeps heap + /// ownership from crossing the boundary — but would otherwise make every + /// `PluginError` that size. This type never crosses FFI, so boxing here + /// costs nothing the ABI cares about. #[error("memory plugin call `{operation}` failed: {}", .status.describe())] Call { operation: &'static str, - status: NxmemStatus, + status: Box, }, /// The plugin violated the contract in a way that is not a mere failure. @@ -54,7 +60,10 @@ pub enum PluginError { impl PluginError { /// Build a call failure from a status, for a named operation. pub fn call(operation: &'static str, status: NxmemStatus) -> Self { - Self::Call { operation, status } + Self::Call { + operation, + status: Box::new(status), + } } } @@ -125,7 +134,8 @@ mod tests { #[test] fn a_wrong_mechanism_status_is_not_reported_as_out_of_memory() { - let status = NxmemStatus::with_message(NxmemStatusCode::WrongMechanism, "foreign allocation"); + let status = + NxmemStatus::with_message(NxmemStatusCode::WrongMechanism, "foreign allocation"); let error = status_to_memory_error("deallocate", Tier::Device, 0, &status); let text = error.to_string(); assert!(text.contains("misdirected"), "{text}"); diff --git a/crates/onnx-runtime-memory-host/src/lib.rs b/crates/onnx-runtime-memory-host/src/lib.rs index c5533a79f1..07e58dd25d 100644 --- a/crates/onnx-runtime-memory-host/src/lib.rs +++ b/crates/onnx-runtime-memory-host/src/lib.rs @@ -32,6 +32,12 @@ //! //! [report]: onnx_runtime_memory_abi::NxmemUnloadReport +// `NxmemStatus` carries its message inline, in a fixed 256-byte buffer, so that +// no heap allocation is ever owned across the dynamic-library boundary. That +// makes every `Result<_, NxmemStatus>` "large" by clippy's measure. Boxing the +// error, which is what the lint asks for, would reintroduce exactly the +// cross-allocator ownership this ABI exists to forbid. +#![allow(clippy::result_large_err)] pub mod allocator; pub mod error; pub mod loader; diff --git a/crates/onnx-runtime-memory-host/src/loader.rs b/crates/onnx-runtime-memory-host/src/loader.rs index 2b2fa10e0d..4c59a7d022 100644 --- a/crates/onnx-runtime-memory-host/src/loader.rs +++ b/crates/onnx-runtime-memory-host/src/loader.rs @@ -252,12 +252,11 @@ impl MemoryPlugin { // SAFETY: loading a dynamic library runs its initialisers. The caller // is trusted to name a library it intends to load, exactly as for any // other plugin surface in this workspace. - let library = unsafe { libloading::Library::new(&path) }.map_err(|source| { - PluginError::Open { + let library = + unsafe { libloading::Library::new(&path) }.map_err(|source| PluginError::Open { path: display.clone(), source, - } - })?; + })?; let negotiate = resolve::(&library, NXMEM_SYMBOL_NEGOTIATE, &display)?; let create_factories = resolve::( @@ -407,9 +406,7 @@ impl MemoryPlugin { // Both sides agree nothing is live. Release the factories, then check // that nothing else still pins the module before letting it unmap. let Self { - factories, - module, - .. + factories, module, .. } = self; drop(factories); diff --git a/crates/onnx-runtime-memory-host/tests/nxmem_abi.rs b/crates/onnx-runtime-memory-host/tests/nxmem_abi.rs index 76dbae4c65..1d1d71c8d7 100644 --- a/crates/onnx-runtime-memory-host/tests/nxmem_abi.rs +++ b/crates/onnx-runtime-memory-host/tests/nxmem_abi.rs @@ -263,7 +263,9 @@ fn a_short_allocator_vtable_is_refused_before_any_slot_is_read() { let text = error.to_string(); assert!( - text.contains("struct_size") || text.to_lowercase().contains("short") || text.contains("small"), + text.contains("struct_size") + || text.to_lowercase().contains("short") + || text.contains("small"), "the refusal must explain that the struct is too short, got: {text}" ); } @@ -413,6 +415,82 @@ fn a_release_that_misdescribes_the_allocation_is_refused() { assert!(matches!(outcome, AllocationReleaseOutcome::Complete { .. })); } +/// A plugin that keeps part of the mapping must be quarantined, not refunded. +/// +/// This drives the real [`DeviceAllocator::release`] entry point rather than +/// the private outcome-interpreting helper, because the branch worth pinning +/// is the one the production caller actually reaches. +#[test] +fn a_partial_release_is_quarantined_rather_than_refunded() { + let _serial = serial(); + let plugin = load(); + let allocator = open(&plugin, "quarantining", NXMEM_CAP_ALLOCATOR); + let ptr = allocator.allocate(8192, 64).expect("allocation"); + + // SAFETY: the address is live with these exact parameters. + let outcome = unsafe { allocator.release(ptr, 8192, 64) }; + let AllocationReleaseOutcome::Quarantined { + accounting, + residual, + } = outcome + else { + panic!("a partial release must quarantine, got {outcome:?}"); + }; + assert_eq!( + accounting.allocation_bytes, 8192, + "the full allocation size is still what was charged" + ); + assert_eq!( + accounting.unmapped_bytes, 4096, + "only the part the plugin really gave back may be credited" + ); + assert_eq!( + residual.retained_bytes, 4096, + "the rest stays owned by the plugin" + ); + assert_eq!( + residual.address, + ptr.as_ptr() as usize, + "the quarantine record must name the address so it is never reissued" + ); + assert_eq!( + allocator.core().live_allocation_count(), + 0, + "the host no longer tracks it: the plugin does" + ); +} + +/// A release state from a later contract level fails closed. +/// +/// The host has no way to know whether the memory is safe to reuse, so it must +/// quarantine rather than guess in either direction. +#[test] +fn a_release_state_from_the_future_is_quarantined_not_guessed() { + let _serial = serial(); + let plugin = load(); + let allocator = open(&plugin, "future-state", NXMEM_CAP_ALLOCATOR); + let ptr = allocator.allocate(4096, 64).expect("allocation"); + + // SAFETY: the address is live with these exact parameters. + let outcome = unsafe { allocator.release(ptr, 4096, 64) }; + let AllocationReleaseOutcome::Quarantined { + accounting, + residual, + } = outcome + else { + panic!("an uninterpretable state must fail closed, got {outcome:?}"); + }; + assert_eq!( + accounting.unmapped_bytes, 0, + "nothing may be credited back from a state the host cannot read" + ); + assert_eq!( + residual.retained_bytes, 4096, + "the whole allocation is presumed still owned by the plugin" + ); + assert_eq!(allocator.core().live_allocation_count(), 0); +} + // ─── lazy backing ─────────────────────────────────────────────────────────── /// **Lazy backing.** @@ -650,7 +728,10 @@ fn a_refusing_host_callback_fails_the_allocation_cleanly() { let allocator = plugin .factory("callback-probe") .expect("the mechanism is published") - .open(NXMEM_CAP_ALLOCATOR, Some(ScriptedReclaim::granting(1 << 20))) + .open( + NXMEM_CAP_ALLOCATOR, + Some(ScriptedReclaim::granting(1 << 20)), + ) .expect("the mechanism opens"); let ptr = allocator .allocate(4096, 64) @@ -904,7 +985,7 @@ fn every_factory_is_released_exactly_once() { let before = releases(&observer); let plugin = load(); let count = plugin.factories().len() as u64; - assert_eq!(count, 6, "the test plugin publishes six mechanisms"); + assert_eq!(count, 8, "the test plugin publishes eight mechanisms"); plugin.try_unload().expect("an idle plugin unloads"); assert_eq!( diff --git a/crates/onnx-runtime-memory-host/tests/support/testplugin.rs b/crates/onnx-runtime-memory-host/tests/support/testplugin.rs index 45b473551b..6b8b41e44c 100644 --- a/crates/onnx-runtime-memory-host/tests/support/testplugin.rs +++ b/crates/onnx-runtime-memory-host/tests/support/testplugin.rs @@ -43,8 +43,9 @@ pub fn build_testplugin() -> PathBuf { "onnx_runtime_memory_testplugin.dll" }; - let mut command = - std::process::Command::new(std::env::var("CARGO").unwrap_or_else(|_| String::from("cargo"))); + let mut command = std::process::Command::new( + std::env::var("CARGO").unwrap_or_else(|_| String::from("cargo")), + ); command.args(["build", "-p", "onnx-runtime-memory-testplugin"]); if profile != "debug" { command.args(["--profile", &profile]); @@ -77,4 +78,3 @@ pub fn build_testplugin() -> PathBuf { {candidates:?}. Set NXMEM_TESTPLUGIN_PATH when using a custom target directory." ); } - diff --git a/crates/onnx-runtime-memory-testplugin/src/lib.rs b/crates/onnx-runtime-memory-testplugin/src/lib.rs index 002539e350..54368672b0 100644 --- a/crates/onnx-runtime-memory-testplugin/src/lib.rs +++ b/crates/onnx-runtime-memory-testplugin/src/lib.rs @@ -25,6 +25,12 @@ //! Backing storage is ordinary host memory from `std::alloc`, allocated **and //! freed inside this module**. No allocator ownership crosses the boundary. +// `NxmemStatus` carries its message inline, in a fixed 256-byte buffer, so that +// no heap allocation is ever owned across the dynamic-library boundary. That +// makes every `Result<_, NxmemStatus>` "large" by clippy's measure. Boxing the +// error, which is what the lint asks for, would reintroduce exactly the +// cross-allocator ownership this ABI exists to forbid. +#![allow(clippy::result_large_err)] use std::alloc::{Layout, alloc, dealloc}; use std::collections::HashMap; use std::ffi::c_void; @@ -91,6 +97,12 @@ struct Behaviour { probe_host_callback: bool, /// Never retire a queued release, so unload stays refused. never_drain: bool, + /// Keep residual ownership of every released allocation instead of + /// completing it, so the host must quarantine rather than refund. + quarantine_on_release: bool, + /// Report a release state code from a future contract level the host has + /// never heard of. + unknown_release_state: bool, } impl Behaviour { @@ -101,6 +113,8 @@ impl Behaviour { short_allocator_struct: false, probe_host_callback: false, never_drain: false, + quarantine_on_release: false, + unknown_release_state: false, } } } @@ -162,6 +176,27 @@ const MECHANISMS: &[Mechanism] = &[ ..Behaviour::base() }, }, + Mechanism { + name: "quarantining", + c_name: c"quarantining", + behaviour: Behaviour { + // Models a mechanism that cannot fully unmap on release: the + // address must never be handed back out, and the host must not + // refund the residue. + quarantine_on_release: true, + ..Behaviour::base() + }, + }, + Mechanism { + name: "future-state", + c_name: c"future-state", + behaviour: Behaviour { + // Models a plugin from a later contract level reporting an + // outcome this host predates. The host must fail closed. + unknown_release_state: true, + ..Behaviour::base() + }, + }, Mechanism { name: "sticky", c_name: c"sticky", @@ -185,6 +220,8 @@ pub const MECHANISM_NAMES: &[&str] = &[ "short-struct", "callback-probe", "legacy-1-0", + "quarantining", + "future-state", "sticky", ]; @@ -246,6 +283,10 @@ struct AllocatorInner { blocks: HashMap, queue: Vec, prefixes: HashMap, + /// Blocks this module kept residual ownership of. They are *not* free and + /// must never be reissued; the module frees them only when the allocator + /// itself is destroyed. + retained: Vec, } impl AllocatorState { @@ -295,11 +336,7 @@ impl AllocatorState { }; // SAFETY: `completion` is a valid local that outlives the call. let status = unsafe { release_completed(callbacks.host_ctx, completion as *const _) }; - if status.is_ok() { - Ok(()) - } else { - Err(status) - } + if status.is_ok() { Ok(()) } else { Err(status) } } } @@ -568,6 +605,40 @@ unsafe extern "C" fn plugin_release_allocation( let unmapped = block.committed as u64; let bytes = block.bytes as u64; + + if state.behaviour.unknown_release_state { + // Deliberately from the future. The block is kept, because a host + // that cannot interpret the state must not treat the address as + // reusable either. + let mut outcome = NxmemReleaseOutcome::complete(bytes, unmapped); + outcome.state = FUTURE_RELEASE_STATE; + retain_block(state, block); + counters().live_allocations.fetch_sub(1, Ordering::AcqRel); + // SAFETY: checked non-null above. + unsafe { *outcome_out = outcome }; + return NxmemStatus::ok(); + } + + if state.behaviour.quarantine_on_release { + // Half the mapping came back; the module still owns the rest. + let returned = unmapped / 2; + retain_block(state, block); + counters().live_allocations.fetch_sub(1, Ordering::AcqRel); + // SAFETY: checked non-null above. + unsafe { + *outcome_out = NxmemReleaseOutcome::quarantined( + bytes, + returned, + bytes - returned, + NxmemStatus::with_message( + NxmemStatusCode::ReleaseQuarantined, + "testplugin: this mechanism cannot unmap the whole allocation", + ), + ); + } + return NxmemStatus::ok(); + } + free_block(block); counters().live_allocations.fetch_sub(1, Ordering::AcqRel); // SAFETY: checked non-null above. @@ -757,6 +828,21 @@ unsafe extern "C" fn plugin_release(ctx: *mut c_void) { /// A queued release holds its own reference, so the state cannot be destroyed /// while a deferred free still names it. `retain`/`release` never call back /// into the host, as the contract requires. +/// A release state code from a contract level this host predates. +const FUTURE_RELEASE_STATE: u32 = 0x0BAD_F00D; + +/// Keep residual ownership of a block. +/// +/// The memory is neither freed nor reissued: it is parked until the allocator +/// is destroyed, which is exactly what "the plugin still owns it" means. +fn retain_block(state: &AllocatorState, block: Block) { + // Losing the lock must not turn into a double free, so on poisoning the + // block is deliberately leaked for the lifetime of the process instead. + if let Ok(mut inner) = state.inner.lock() { + inner.retained.push(block); + } +} + fn release_state(state: &AllocatorState) { let previous = state.refcount.fetch_sub(1, Ordering::AcqRel); if previous != 1 { @@ -779,6 +865,10 @@ fn release_state(state: &AllocatorState) { counters().live_allocations.fetch_sub(1, Ordering::AcqRel); counters().queued_releases.fetch_sub(1, Ordering::AcqRel); } + let retained: Vec = inner.retained.drain(..).collect(); + for block in retained { + free_block(block); + } let prefixes: Vec = inner.prefixes.drain().map(|(_, prefix)| prefix).collect(); for prefix in prefixes { counters() @@ -1358,17 +1448,18 @@ unsafe extern "C" fn plugin_open_allocator( vtable.committed_bytes = Some(plugin_committed_bytes); Box::new(vtable) }); - let shared_vtable = (behaviour.capability_flags & NXMEM_CAP_SHARED_MAPPING != 0).then(|| { - let mut vtable = NxmemSharedMappingVtable::zeroed(); - vtable.abi_minor = abi_minor; - vtable.mechanism_id = mechanism_id; - vtable.create_shared_prefix = Some(plugin_create_shared_prefix); - vtable.retain_shared_prefix = Some(plugin_retain_shared_prefix); - vtable.release_shared_prefix = Some(plugin_release_shared_prefix); - vtable.incremental_owned_bytes = Some(plugin_incremental_owned_bytes); - vtable.commit_shared_prefix = Some(plugin_commit_shared_prefix); - Box::new(vtable) - }); + let shared_vtable = + (behaviour.capability_flags & NXMEM_CAP_SHARED_MAPPING != 0).then(|| { + let mut vtable = NxmemSharedMappingVtable::zeroed(); + vtable.abi_minor = abi_minor; + vtable.mechanism_id = mechanism_id; + vtable.create_shared_prefix = Some(plugin_create_shared_prefix); + vtable.retain_shared_prefix = Some(plugin_retain_shared_prefix); + vtable.release_shared_prefix = Some(plugin_release_shared_prefix); + vtable.incremental_owned_bytes = Some(plugin_incremental_owned_bytes); + vtable.commit_shared_prefix = Some(plugin_commit_shared_prefix); + Box::new(vtable) + }); let mut allocator_vtable = NxmemAllocatorVtable::zeroed(); allocator_vtable.abi_minor = abi_minor; diff --git a/docs/memory/MEMORY_PLUGIN_ABI.md b/docs/memory/MEMORY_PLUGIN_ABI.md new file mode 100644 index 0000000000..62f6c0ae73 --- /dev/null +++ b/docs/memory/MEMORY_PLUGIN_ABI.md @@ -0,0 +1,214 @@ +# nxmem — the dynamic-plugin memory ABI + +Phase 6 of the memory architecture refactor (issue #1186). + +A *memory plugin* is a shared library that publishes allocation **mechanisms** +to an ONNX GenAI host. This document is the contract; the machine-readable +version is [`nxmem_memory_abi.h`](../../crates/onnx-runtime-memory-abi/include/nxmem_memory_abi.h), +and the smallest correct implementation is +[`minimal_plugin.c`](../../crates/onnx-runtime-memory-abi/examples/minimal_plugin.c). + +## Why an ABI at all + +Phases 1–5 moved every allocation decision behind Rust traits +(`DeviceAllocator`, `VirtualBacking`, `SharedMapping`) and put a process-wide +governor in front of them. Those traits are Rust: they use trait objects, +`Arc`, Rust enum layouts, and unwinding, none of which have a stable +representation across a `dlopen` boundary or across compiler versions. + +Phase 6 adds a C ABI that expresses the same capabilities, plus a host-side +adapter that wraps a C vtable back into those Rust traits. A plugin therefore +plugs into the *existing* governance machinery without the governor learning +that it is talking to a plugin at all. + +``` +onnx-runtime-memory-abi #[repr(C)] structs, vtables, versioning. + (no dependencies) Shared verbatim by host and plugin. + │ + ├── onnx-runtime-memory-host loads a cdylib, negotiates, and + │ (libloading) wraps its vtables in the traits + │ from onnx-runtime-memory-api + │ + └── onnx-runtime-memory-testplugin a cdylib exercising every rule, + (cdylib) including the ones about failing +``` + +## What may cross the boundary + +Only `#[repr(C)]` structs from the ABI crate, by pointer or by value. + +Explicitly **never**: a Rust trait object, an `Arc`, a Rust `enum`, a Rust +`String`/`Vec`, ownership of either side's allocator, or a panic. Every +`extern "C"` entry point on both sides wraps its body in `catch_unwind` +(`catch_status_panic` / `catch_void_panic`), so a panic becomes +`NXMEM_STATUS_INTERNAL_ERROR` instead of unwinding into foreign frames, which +is undefined behaviour. + +Errors are `NxmemStatus`: a stable `u32` code plus a 256-byte inline message +buffer. Inline, because a heap-allocated message would mean one side freeing +the other side's allocation. + +## Versioning and negotiation + +`major` is a hard gate. `minor` may only **append** fields, so a larger struct +is always a strict superset of a smaller one. + +Current: major 1, minor 1. Baseline supported minor: 0. Minor 1 appends exactly +one allocator slot, `release_allocation`, gated by +`NXMEM_CAP_STRUCTURED_RELEASE`. + +The host calls `NxmemNegotiate` once. The agreed minor is a **ceiling**, not an +assignment: each vtable separately declares in its own `abi_minor` what it +really implements, so one module can ship a current mechanism beside one still +written to the baseline. The test plugin does exactly that (`lazy` at minor 1, +`legacy-1-0` at minor 0). + +Reading an untrusted vtable goes through `read_prefix`, never a plain +dereference: + +1. reject a null or misaligned pointer; +2. read `struct_size` (offset 0) and `abi_minor` (offset 4) unaligned — those + two fields are at those offsets in every version, forever; +3. reject `struct_size` smaller than the level the sender *claims*: a sender + contradicting itself is broken, not old; +4. clamp: `effective_minor = min(declared, negotiated)`, and rewrite + `abi_minor` to it so callers need only consult the value they get back; +5. copy `min(struct_size, size_of::())` bytes into a zeroed local; +6. null out every slot above the effective level. + +Step 4 is what makes an old host usable with a new plugin. Rejecting instead +would mean a plugin could never add a slot without breaking every existing +host, which defeats the purpose of a minor version. + +A NULL optional slot means the capability is **absent** and surfaces as +`NXMEM_STATUS_UNSUPPORTED_CAPABILITY`. It is never a silently successful +no-op. The host treats a non-NULL capability vtable behind a *clear* capability +flag as a contract violation rather than a bonus. + +## Ownership + +| object | created by | released by | +| --- | --- | --- | +| factory vtable | `NxmemCreateAllocatorFactories` | host, via the factory's `release`, exactly once | +| allocator vtable | factory `open_allocator` | host, via the allocator's `release`; `retain` adds a reference | +| virtual-backing / shared-mapping vtable | plugin | the owning allocator — never separately by the host | +| allocation | allocator `allocate` | `deallocate` or `release_allocation` | +| shared prefix | `create_shared_prefix` | matching `release_shared_prefix` | +| host callbacks | host | host, after the last allocator *and* the last queued release retire | + +Every pointer passed into a call is borrowed for that call only, with two +stated exceptions: `NxmemOpenRequest::callbacks` is borrowed for the whole +lifetime of the allocator opened with it, and a factory's `name` must outlive +the factory. + +The host holds its callback table in a `Box` created *before* `open_allocator` +so its address is stable, and drops it only in `Drop for AllocatorCore`, after +the plugin's `release` has returned. + +If the host refuses a vtable that `open_allocator` returned `Ok` for, it still +owes the plugin a `release` — it makes a best-effort call when the struct is +well-formed enough to locate that slot. A plugin publishing a vtable too +malformed to read must therefore not allocate state before returning it. + +## Threading, re-entrancy, and locks + +Every slot may be called concurrently from any thread; a plugin does its own +locking. + +**The rule that matters most: no participant blocks, and no participant holds +one of its own locks, across a call into the other side.** + +- The host never calls into a plugin while holding a governance lock. This is + the same invariant Phases 1–5 established for Rust trait objects + (`pressure` drops the lock before `on_pressure`; `allocate_with` holds no + lock while running the caller's closure; `run_drain_callback_if_ready` uses + `.take()` so the callback runs outside the lock) — an ABI call is strictly + more dangerous than a trait-object call, so the rule tightens rather than + relaxes. +- The host's `take_allocation` locks its live map, removes the record, drops + the guard, and *only then* enters the plugin. +- A plugin must not hold a lock across `request_reclaim`, because the host may + re-enter the same plugin on the same thread to satisfy it. +- `drain_releases` invokes `release_completed` per retired ticket, in enqueue + order. Take the batch under your lock, drop the lock, then call the host. + +## Deferred release + +`enqueue_release` hands an allocation over without freeing it and yields a +ticket. Until that ticket retires: + +- the allocation counts as live; +- it pins the allocator, which pins the module; +- the host's callback table stays alive, because `release_completed` will still + be called through it. + +This is what "deferred release keeps the plugin module and provider/context +pinned until callbacks and frees retire" means concretely. + +## Release outcomes + +Three states, deliberately not interchangeable: + +| state | meaning | +| --- | --- | +| `COMPLETE` | the memory is gone; `unmapped_bytes` may be credited back | +| `QUARANTINED` | the plugin still owns `residual_owned_bytes`; the address must never be reissued and the residue must not be refunded | +| `FAILED` | nothing was mutated; the allocation is as live as the caller left it | + +A state code the host does not recognise is treated as **quarantine**. It is +the only interpretation that cannot corrupt either memory or accounting: +guessing `COMPLETE` would reissue live memory, and guessing `FAILED` would +leave the host convinced it still owns something the plugin freed. + +## Cross-provider misuse + +Every allocation carries `mechanism_id` and `device`. Allocation, release, +range, and shared-prefix calls all check both before acting, and report +`NXMEM_STATUS_WRONG_MECHANISM` / `NXMEM_STATUS_WRONG_DEVICE` rather than +operating on a stranger's pointer. + +`allocation_id` is a host-assigned monotonic counter and is **never derived +from the address** — the same reasoning that made `AllocationGeneration` +non-pointer-derived in Phase 1. An address can be reused; an id cannot, so an +ABA race cannot make one allocation impersonate another. + +## Unload gating + +`NxmemQueryUnloadReadiness` reports live allocators, allocations, views, +capabilities, and queued releases. The host refuses or defers unload while any +count is non-zero, and it checks its own counters too, so a rejection carries +both sides' tallies and a misreporting plugin is visible rather than fatal. + +A plugin must count conservatively: reporting zero while anything is reachable +invites the host to unmap code that is about to run. + +## Test plugin + +`onnx-runtime-memory-testplugin` is a `cdylib` loaded at runtime by the ABI +tests — out-of-tree in the way that matters (`dlopen`, no workspace linking), +in-tree in the way that does not (it lives in the repo so it is built and +linted with everything else). It publishes eight named mechanisms rather than +switching on environment variables or globals, so the tests select behaviour by +name: + +| mechanism | behaviour under test | +| --- | --- | +| `eager` | the minimal conforming mechanism: required slots only | +| `lazy` | virtual backing, shared mapping, deferred release, structured release | +| `short-struct` | a vtable claiming fewer bytes than the baseline prefix | +| `callback-probe` | calls `request_reclaim` and fails cleanly when the host refuses | +| `legacy-1-0` | built to minor 0: an older participant under a newer host | +| `quarantining` | keeps residual ownership on release | +| `future-state` | reports a release state from a later contract level | +| `sticky` | never retires a queued release, so unload stays refused | + +It runs entirely on host memory, so the whole suite is portable and needs no +accelerator. + +## Relationship to CUDA + +Nothing in this phase is CUDA-specific, and no CUDA code changes here. The +mechanisms a CUDA plugin would publish (VMM-backed lazy commit, shared +prefixes, stream-ordered release) are exactly the optional capabilities the ABI +already carries. Phase 7 is where the built-in eager allocator moves behind +this boundary.