diff --git a/crates/onnx-runtime-cuda-memory/src/vmm_allocator.rs b/crates/onnx-runtime-cuda-memory/src/vmm_allocator.rs index 2b6832f7f5..ad1b8aa155 100644 --- a/crates/onnx-runtime-cuda-memory/src/vmm_allocator.rs +++ b/crates/onnx-runtime-cuda-memory/src/vmm_allocator.rs @@ -57,9 +57,8 @@ use std::sync::{Arc, Mutex}; use onnx_runtime_memory_governor::{ AllocationCommitRange, DeviceAllocator, DeviceKey, HolderId, MappedPhysicalCapacityToken, - MechanismId, MemoryAuthorityId, MemoryError, MemoryGovernor, MemoryLease, MemoryRole, - SharedDevicePrefix, SharedMapping, SharedPrefixCommitInfo, Tier, - VirtualBacking as MemoryVirtualBacking, + MemoryAuthorityId, MemoryError, MemoryGovernor, MemoryLease, MemoryRole, SharedDevicePrefix, + SharedMapping, SharedPrefixCommitInfo, Tier, VirtualBacking as MemoryVirtualBacking, }; use onnx_runtime_virtual_memory::{PhysicalMemoryAccounting, VirtualBacking}; @@ -1778,16 +1777,6 @@ impl DeviceAllocator for CudaVmmAllocator { fn as_shared_mapping(&self) -> Option<&dyn SharedMapping> { self.backing.physical_pool().is_some().then_some(self) } - - /// `Some(MechanismId::of(self))`: `CudaVmmAllocator` is a concrete, - /// `'static` type, and every capability it advertises through - /// `as_virtual_backing`/`as_shared_mapping` returns `self` — so its own - /// identity is exactly what a caller needs to prove a capability - /// genuinely came from this allocator (#1186 Phase 2 review, round 4 - /// finding 2). - fn mechanism_id(&self) -> Option { - Some(MechanismId::of(self)) - } } // SAFETY: every pointer handed to these methods lies inside this allocator's @@ -2149,10 +2138,7 @@ impl SharedMapping for CudaVmmAllocator { /// (`impl DeviceAllocator for CudaVmmAllocator::deallocate_with_unmapped` /// overrides the base default entirely, above). This exists for a caller /// that reaches release accounting through `as_shared_mapping` directly, - /// bypassing that override (#1186 Phase 2 review, round 3 finding 2), and - /// for that caller it must still perform a genuine, atomic, exact-once - /// release — not merely answer a query (#1186 Phase 2 review, round 4 - /// finding 1). + /// bypassing that override (#1186 Phase 2 review, round 3 finding 2). /// /// This allocator's granule tracking does not record whether a committed /// granule became committed through an ordinary @@ -2162,17 +2148,38 @@ impl SharedMapping for CudaVmmAllocator { /// the physical page is mapped and the address range is live. So /// delegating straight to [`deallocate_span`](Self::deallocate_span) — /// the same machinery [`DeviceAllocator::deallocate`] and - /// [`MemoryVirtualBacking::deallocate_committed`] both use — is both the + /// [`MemoryVirtualBacking::deallocate_committed`] both use — is the /// conservative, always-correct accounting answer (this allocation's /// whole committed footprint, exactly what `deallocate_committed` would - /// itself report if invoked instead) *and* a genuine atomic release: the - /// span's removal from the live-allocation table and its granule release - /// happen under one lock acquisition inside `deallocate_span`, and - /// removing an already-removed span is a safe no-op that reports `0` — - /// so this method is exact-once, refunds zero on a repeated or otherwise - /// unmatched call, and is safe to call whether or not a caller also - /// later calls [`DeviceAllocator::deallocate`]/`deallocate_with_unmapped` - /// on the same pointer (which would simply find nothing left to release). + /// itself report if invoked instead), and is exact-once against repeated + /// calls from this crate's own bookkeeping: the span's removal from the + /// live-allocation table happens under one lock acquisition inside + /// `deallocate_span`, and removing an already-removed span is a safe + /// no-op that reports `0`, so a repeated or unmatched call on the same + /// `ptr` refunds zero rather than double-counting. + /// + /// That bookkeeping-level guarantee is **not** the same as a genuine + /// atomic release against the underlying CUDA driver calls. + /// [`deallocate_span`](Self::deallocate_span) removes the span (and hands + /// its address range back to the free list) unconditionally, even when + /// the CUDA unmap/handle release this triggers only partially succeeds: + /// if releasing one contiguous granule run fails, this allocator's + /// internal granule-release bookkeeping leaves that run's reference + /// counts exactly as they were (so the + /// bytes it could not release are simply never counted in this method's + /// return value, rather than being double-refunded), but the *address + /// range* is still given back to the free list either way. A subsequent + /// allocation can then be granted a virtual-address range that overlaps + /// granules this allocator's own bookkeeping still believes are + /// committed. This is a pre-existing property of this allocator's + /// internal `deallocate_span`/granule-release bookkeeping that predates + /// the + /// Phase 2 capability split — no Phase 2 change introduced or altered + /// it — and closing it needs an owning allocation token or generation + /// counter that can keep a VA range out of the free list until its + /// release genuinely reaches a consistent terminal state, which is + /// #1186 Phase 4's scope (deferred, reconciled release), not something + /// this method can honestly claim to already provide. /// /// [`DeviceAllocator::deallocate`]: onnx_runtime_memory_governor::DeviceAllocator::deallocate /// diff --git a/crates/onnx-runtime-ep-cuda/src/provider.rs b/crates/onnx-runtime-ep-cuda/src/provider.rs index 2605af9538..c30acaa1ce 100644 --- a/crates/onnx-runtime-ep-cuda/src/provider.rs +++ b/crates/onnx-runtime-ep-cuda/src/provider.rs @@ -471,13 +471,20 @@ impl CudaExecutionProvider { /// A caller who needs a specific base allocator must build without VMM /// (`ONNX_GENAI_CUDA_VMM` unset and dynamic lending disabled). /// - If `memory` advertises a `VirtualBacking`/`SharedMapping` capability - /// whose `device()` disagrees with `memory.device()`, or whose - /// underlying object is not `memory` itself (see - /// [`capability_shares_mechanism`](onnx_runtime_memory_governor::capability_shares_mechanism)). - /// A composite allocator whose ordinary allocate/deallocate and whose - /// advertised capability are two different objects cannot prove that a - /// pointer produced through the capability is ever released through the - /// mechanism that produced it (#1186 Phase 2 review). + /// whose `device()` disagrees with `memory.device()`. A capability that + /// disagrees with its own allocator about which device it serves + /// cannot be trusted to belong to it. This device check is a sanity + /// check on `memory`'s own self-reported answers, not a proof that the + /// capability object is genuinely `memory` itself as opposed to some + /// other object `memory` merely composes: no self-reported identity + /// token can distinguish those two cases from outside the type (see + /// the `capability` module's docs in `onnx_runtime_memory_api` for why + /// #1186 Phase 2 review's round 6 removed the earlier attempt at that + /// proof). A composite allocator whose ordinary allocate/deallocate and + /// whose advertised capability are two different objects is expected + /// to implement `as_virtual_backing`/`as_shared_mapping` honestly + /// enough that this does not arise; proving it structurally is #1186 + /// Phase 3's scope. pub fn with_memory( mut self, memory: Arc, @@ -503,72 +510,40 @@ impl CudaExecutionProvider { expected.index, key.tier, key.index, expected.index ))); } - // Any capability this allocator advertises must genuinely be *this* - // allocator, not a different inner object it merely composes. - // Release for a `VirtualBacking`-produced pointer is reached through - // this same `as_virtual_backing()` call at both allocate and - // deallocate time (see `deallocate_with_unmapped`), so a foreign - // capability object can never actually be freed through a different - // mechanism through this EP — but rejecting a malformed allocator - // here, rather than merely tolerating it, keeps "which mechanism - // owns this pointer" an identity a caller can trust rather than an - // accident of how `deallocate_with_unmapped` happens to be wired - // (#1186 Phase 2 review, finding 1). - if let Some(virtual_backing) = memory.as_virtual_backing() { - if virtual_backing.device() != key { - return Err(EpError::KernelFailed(format!( - "cuda_ep: the injected allocator's VirtualBacking capability reports \ - device {:?} {}, but the allocator itself reports {:?} {}; a capability \ - that disagrees with its own allocator about device identity cannot be \ - trusted to belong to it", - virtual_backing.device().tier, - virtual_backing.device().index, - key.tier, - key.index - ))); - } - if !onnx_runtime_memory_governor::capability_shares_mechanism( - memory.as_ref(), - virtual_backing.as_any(), - ) { - return Err(EpError::KernelFailed( - "cuda_ep: the injected allocator's VirtualBacking capability is not the \ - allocator itself. A composite allocator whose ordinary \ - allocate/deallocate and whose VirtualBacking capability are two \ - different objects sharing only a matching device cannot prove that \ - release always reaches the mechanism that allocated; implement \ - `as_virtual_backing` to return `Some(self)` on the same type instead." - .into(), - )); - } + // A capability's own `device()` must at least agree with the + // allocator's `device()`. This is a sanity check on self-reported + // values, not proof the capability object is genuinely `memory` + // itself rather than a distinct object `memory` composes — see the + // doc comment above and `onnx_runtime_memory_api::capability`'s + // module docs for why that stronger proof was removed rather than + // patched further (#1186 Phase 2 review, round 6). + if let Some(virtual_backing) = memory.as_virtual_backing() + && virtual_backing.device() != key + { + return Err(EpError::KernelFailed(format!( + "cuda_ep: the injected allocator's VirtualBacking capability reports \ + device {:?} {}, but the allocator itself reports {:?} {}; a capability \ + that disagrees with its own allocator about device identity cannot be \ + trusted to belong to it", + virtual_backing.device().tier, + virtual_backing.device().index, + key.tier, + key.index + ))); } - if let Some(shared_mapping) = memory.as_shared_mapping() { - if shared_mapping.device() != key { - return Err(EpError::KernelFailed(format!( - "cuda_ep: the injected allocator's SharedMapping capability reports \ - device {:?} {}, but the allocator itself reports {:?} {}; a capability \ - that disagrees with its own allocator about device identity cannot be \ - trusted to belong to it", - shared_mapping.device().tier, - shared_mapping.device().index, - key.tier, - key.index - ))); - } - if !onnx_runtime_memory_governor::capability_shares_mechanism( - memory.as_ref(), - shared_mapping.as_any(), - ) { - return Err(EpError::KernelFailed( - "cuda_ep: the injected allocator's SharedMapping capability is not the \ - allocator itself. A composite allocator whose ordinary \ - allocate/deallocate and whose SharedMapping capability are two \ - different objects sharing only a matching device cannot be trusted to \ - mix them safely; implement `as_shared_mapping` to return `Some(self)` \ - on the same type instead." - .into(), - )); - } + if let Some(shared_mapping) = memory.as_shared_mapping() + && shared_mapping.device() != key + { + return Err(EpError::KernelFailed(format!( + "cuda_ep: the injected allocator's SharedMapping capability reports \ + device {:?} {}, but the allocator itself reports {:?} {}; a capability \ + that disagrees with its own allocator about device identity cannot be \ + trusted to belong to it", + shared_mapping.device().tier, + shared_mapping.device().index, + key.tier, + key.index + ))); } self.memory = memory; Ok(self) @@ -2170,29 +2145,36 @@ extern "C" __global__ void write_after_delay(unsigned int* out, long long spin) ); } - /// A hypothetical composite allocator whose ordinary allocate/deallocate - /// are backed by one object but whose `VirtualBacking` capability is a - /// totally different (real) VMM arena, sharing only a matching - /// `DeviceKey` — exactly the "same-device wrapper exposing foreign - /// backing" scenario the #1186 Phase 2 review flagged as unable to prove - /// release always reaches the mechanism that allocated. `with_memory` - /// must reject it outright rather than accept it and rely on downstream - /// code happening to be safe (finding 1). + /// `with_memory`'s capability checks only validate that a + /// `VirtualBacking`/`SharedMapping` capability's own `device()` agrees + /// with the allocator's `device()` — a sanity check on self-reported + /// values, not a structural proof that the capability object is + /// genuinely the allocator itself rather than some other object the + /// allocator composes. An earlier design (#1186 Phase 2 review, rounds + /// 3-4) attempted that stronger proof through a self-reported + /// `MechanismId`/`TypeId`+pointer identity token; coordinator review + /// (round 6) found a type can honestly return `Some(self)` from both its + /// capability accessor and its identity accessor while its capability's + /// method *bodies* still operate on unrelated state, which no + /// self-reported token can detect from outside the type. That apparatus + /// was removed rather than patched further; see + /// `onnx_runtime_memory_api::capability`'s module docs for the full + /// rationale, and #1186 Phase 3 for where a real proof belongs. This + /// test proves the surviving, real (if narrower) device-mismatch check. #[cfg_attr( not(feature = "gpu-tests"), ignore = "requires CUDA device; enable the gpu-tests feature on a CUDA runner" )] #[test] - fn with_memory_rejects_an_allocator_whose_capability_is_a_different_mechanism() { + fn with_memory_rejects_a_capability_whose_device_disagrees_with_its_allocator() { // SAFETY: single-process test; ensure VMM is not force-enabled by a - // previous test in this binary, so the identity-validation path below - // (not the "VMM already selected" path) is what actually runs. + // previous test in this binary, so the capability-validation path + // below (not the "VMM already selected" path) is what actually runs. unsafe { std::env::remove_var(crate::vmm_allocator::CUDA_VMM_ENV) }; let Ok(provider) = CudaExecutionProvider::new(0) else { eprintln!( - "SKIPPED (no CUDA runtime): the with_memory capability-identity proof did NOT \ - run." + "SKIPPED (no CUDA runtime): the with_memory device-mismatch proof did NOT run." ); panic!("CUDA test path did not run; report as a failed GPU test, not a pass"); }; @@ -2205,29 +2187,18 @@ extern "C" __global__ void write_after_delay(unsigned int* out, long long spin) let base = Arc::new(crate::device_allocator::CudaDeviceAllocator::new( context.clone(), )); - let foreign = Arc::new( - crate::vmm_allocator::CudaVmmAllocator::standalone( - context, - onnx_runtime_memory_governor::DeviceKey::device(0), - 0, - 64 << 20, - onnx_runtime_memory_governor::HolderId::new(1186), - onnx_runtime_memory_governor::MemoryRole::Workspace { step_scoped: false }, - ) - .expect("a second, unrelated VMM arena over the same device"), - ); - /// Deliberately composes two different objects behind one - /// `DeviceAllocator`: ordinary allocate/deallocate go through `base`, - /// but the advertised `VirtualBacking` is `foreign` — a different - /// mechanism than the one that would service `allocate`/`deallocate`. + /// A `DeviceAllocator` that honestly returns `Some(self)` as its own + /// `VirtualBacking` capability, but whose `VirtualBacking::device()` + /// disagrees with its own `DeviceAllocator::device()` — the one + /// mismatch `with_memory` can still detect without a structural + /// identity proof, since it only compares self-reported values. #[derive(Debug)] - struct CompositeAllocator { + struct DeviceMismatchedAllocator { base: Arc, - foreign: Arc, } - impl onnx_runtime_memory_governor::DeviceAllocator for CompositeAllocator { + impl onnx_runtime_memory_governor::DeviceAllocator for DeviceMismatchedAllocator { fn allocate( &self, bytes: usize, @@ -2252,155 +2223,107 @@ extern "C" __global__ void write_after_delay(unsigned int* out, long long spin) fn as_virtual_backing( &self, ) -> Option<&dyn onnx_runtime_memory_governor::VirtualBacking> { - // Deliberately a *different* object than `self`/`self.base` — - // the exact composite the review flagged. - Some(self.foreign.as_ref()) - } - - fn mechanism_id(&self) -> Option { - // The wrapper's own identity, not `self.foreign`'s: this - // proves the rejection below comes from a genuine identity - // mismatch, not from `CompositeAllocator` having opted out of - // identity entirely (#1186 Phase 2 review, round 4 finding 2). - Some(onnx_runtime_memory_governor::MechanismId::of(self)) + Some(self) } } - let composite: Arc = - Arc::new(CompositeAllocator { base, foreign }); - - let error = provider.with_memory(composite).expect_err( - "with_memory must reject an allocator whose VirtualBacking capability is not \ - itself, not accept it and hope release stays consistent by convention", - ); - let message = error.to_string(); - assert!( - message.contains("VirtualBacking capability is not the allocator itself"), - "the rejection must name what happened, not just fail: {message}" - ); - } - - /// #1186 Phase 2 review, round 3 finding 1: pointer equality alone is not - /// unforgeable. `#[repr(C)]` guarantees a struct's first field starts at - /// the struct's own address, so a wrapper that embeds a foreign - /// `CudaVmmAllocator` **by value** as its first field and advertises it - /// as its `VirtualBacking` produces the exact same address as `&self` — - /// the identity check must still reject it. This is the non-vacuous, - /// real-type counterpart to the synthetic proof in - /// `onnx_runtime_memory_api::capability::tests`. - #[cfg_attr( - not(feature = "gpu-tests"), - ignore = "requires CUDA device; enable the gpu-tests feature on a CUDA runner" - )] - #[test] - fn with_memory_rejects_a_foreign_vmm_embedded_at_offset_zero() { - // SAFETY: single-process test; ensure VMM is not force-enabled by a - // previous test in this binary. - unsafe { std::env::remove_var(crate::vmm_allocator::CUDA_VMM_ENV) }; - - let Ok(provider) = CudaExecutionProvider::new(0) else { - eprintln!( - "SKIPPED (no CUDA runtime): the offset-zero identity proof did NOT run." - ); - panic!("CUDA test path did not run; report as a failed GPU test, not a pass"); - }; - assert!( - provider.vmm.get().is_none(), - "this proof requires the plain (non-VMM) construction path" - ); - - let context = provider.runtime().cuda_context(); - let base = Arc::new(crate::device_allocator::CudaDeviceAllocator::new( - context.clone(), - )); - let foreign = crate::vmm_allocator::CudaVmmAllocator::standalone( - context, - onnx_runtime_memory_governor::DeviceKey::device(0), - 0, - 64 << 20, - onnx_runtime_memory_governor::HolderId::new(1186), - onnx_runtime_memory_governor::MemoryRole::Workspace { step_scoped: false }, - ) - .expect("a second, unrelated VMM arena over the same device"); - - /// `#[repr(C)]` with `foreign` as the first field: `&self` and - /// `&self.foreign` are numerically the same address. Ordinary - /// allocate/deallocate go through `base`; the advertised - /// `VirtualBacking` is the embedded `foreign`, not `self`. - #[repr(C)] - #[derive(Debug)] - struct OffsetZeroWrapper { - foreign: crate::vmm_allocator::CudaVmmAllocator, - base: Arc, - } + impl onnx_runtime_memory_governor::VirtualBacking for DeviceMismatchedAllocator { + fn device(&self) -> onnx_runtime_memory_governor::DeviceKey { + // Deliberately a different device than `DeviceAllocator::device` + // above reports. + onnx_runtime_memory_governor::DeviceKey::device( + onnx_runtime_memory_governor::DeviceAllocator::device(self.base.as_ref()) + .index + + 1, + ) + } - impl onnx_runtime_memory_governor::DeviceAllocator for OffsetZeroWrapper { - fn allocate( + fn allocate_committed( &self, - bytes: usize, - align: usize, + _bytes: usize, + _align: usize, + _committed_ranges: &[std::ops::Range], ) -> std::result::Result< std::ptr::NonNull, onnx_runtime_memory_governor::MemoryError, > { - self.base.allocate(bytes, align) + unimplemented!("test double: never called; with_memory must reject before use") } - unsafe fn deallocate(&self, ptr: std::ptr::NonNull, bytes: usize, align: usize) { - // SAFETY: forwarded verbatim; our own caller upholds the same - // contract `DeviceAllocator::deallocate` documents. - unsafe { self.base.deallocate(ptr, bytes, align) } + unsafe fn deallocate_committed( + &self, + _ptr: std::ptr::NonNull, + _allocation_bytes: usize, + _align: usize, + ) -> u64 { + unimplemented!("test double: never called") } - fn device(&self) -> onnx_runtime_memory_governor::DeviceKey { - onnx_runtime_memory_governor::DeviceAllocator::device(self.base.as_ref()) + fn commit_allocation_range( + &self, + _ptr: std::ptr::NonNull, + _allocation_bytes: usize, + _align: usize, + _offset: usize, + _bytes: usize, + ) -> std::result::Result<(), onnx_runtime_memory_governor::MemoryError> { + unimplemented!("test double: never called") } - fn as_virtual_backing( + fn mapped_bytes_for_allocation_ranges( &self, - ) -> Option<&dyn onnx_runtime_memory_governor::VirtualBacking> { - // Deliberately the embedded first field, reached through - // `self`'s own address, not a separately allocated object. - Some(&self.foreign) + _ranges: &[onnx_runtime_memory_governor::AllocationCommitRange], + ) -> std::result::Result { + unimplemented!("test double: never called") } - fn mechanism_id(&self) -> Option { - // The wrapper's own identity, not `self.foreign`'s: this - // proves the rejection below comes from a genuine `TypeId` - // mismatch defeating the offset-zero address collision, not - // from `OffsetZeroWrapper` having opted out of identity - // entirely (#1186 Phase 2 review, round 4 finding 2). - Some(onnx_runtime_memory_governor::MechanismId::of(self)) + fn mapped_bytes_for_allocation( + &self, + _bytes: usize, + _align: usize, + ) -> std::result::Result { + unimplemented!("test double: never called") } - } - let wrapper = OffsetZeroWrapper { foreign, base }; + fn decommit_allocation_range( + &self, + _ptr: std::ptr::NonNull, + _allocation_bytes: usize, + _align: usize, + _offset: usize, + _bytes: usize, + ) -> std::result::Result { + unimplemented!("test double: never called") + } - // Non-vacuous: prove the address collision this test exists to guard - // against is real before asserting on the rejection it causes. - let wrapper_addr = &wrapper as *const OffsetZeroWrapper as *const (); - let foreign_addr = &wrapper.foreign as *const crate::vmm_allocator::CudaVmmAllocator - as *const (); - assert!( - std::ptr::eq(wrapper_addr, foreign_addr), - "test setup did not reproduce the offset-zero address collision `#[repr(C)]` is \ - supposed to guarantee" - ); + fn allocation_committed_bytes( + &self, + _ptr: std::ptr::NonNull, + _allocation_bytes: usize, + _align: usize, + ) -> usize { + unimplemented!("test double: never called") + } - let composite: Arc = Arc::new(wrapper); + fn as_any(&self) -> &dyn std::any::Any { + self + } + } - let error = provider.with_memory(composite).expect_err( - "with_memory must reject a foreign VirtualBacking embedded at offset zero even \ - though its address coincides with the wrapper's own address (#1186 Phase 2 \ - review, round 3 finding 1)", + let mismatched: Arc = + Arc::new(DeviceMismatchedAllocator { base }); + + let error = provider.with_memory(mismatched).expect_err( + "with_memory must reject a VirtualBacking capability whose device() disagrees with its own allocator's device()", ); let message = error.to_string(); assert!( - message.contains("VirtualBacking capability is not the allocator itself"), + message.contains("VirtualBacking capability reports"), "the rejection must name what happened, not just fail: {message}" ); } + /// #956 contrast: the default `cuMemAlloc` path makes exactly one driver /// allocation per request, so it scales one-for-one with decode steps — /// which is the residual the VMM arena removes. Fully isolated (per-instance diff --git a/crates/onnx-runtime-memory-api/src/allocator.rs b/crates/onnx-runtime-memory-api/src/allocator.rs index 9573cd8b4e..3f045d9e2c 100644 --- a/crates/onnx-runtime-memory-api/src/allocator.rs +++ b/crates/onnx-runtime-memory-api/src/allocator.rs @@ -12,7 +12,7 @@ use std::any::Any; use std::fmt::Debug; use std::ptr::NonNull; -use crate::capability::{MechanismId, SharedMapping, VirtualBacking}; +use crate::capability::{SharedMapping, VirtualBacking}; use crate::{MemoryError, Tier}; #[derive(Clone, Copy, Debug)] @@ -169,19 +169,26 @@ pub struct SharedPrefixCommitInfo { /// # Why this does not require [`Any`] /// /// An earlier design (#1186 Phase 2 review, round 3) added `Any` as a -/// supertrait here so [`capability_shares_mechanism`](crate::capability::capability_shares_mechanism) -/// could upcast a `&dyn DeviceAllocator` to `&dyn Any` and compare data -/// pointers and [`TypeId`](std::any::TypeId)s, rejecting a composing wrapper -/// whose first field happens to share its own address. `Any: 'static`, so -/// that supertrait transitively forced every implementation to be `'static` -/// — a real regression against "minimal ordinary contract", since it -/// rejected an otherwise valid allocator built over borrowed, non-`'static` -/// data (#1186 Phase 2 review, round 4 finding 2). -/// -/// [`mechanism_id`](Self::mechanism_id) replaces that: it is a plain, -/// `'static`-free default method that only a concrete type opting into -/// identity-checked capabilities needs to override, so no lifetime -/// restriction reaches an implementation that never does. +/// supertrait here so a self-reported mechanism identity could be compared +/// across a `DeviceAllocator` and its `VirtualBacking`/`SharedMapping` +/// capabilities. `Any: 'static`, so that supertrait transitively forced +/// every implementation to be `'static` — a real regression against +/// "minimal ordinary contract", since it rejected an otherwise valid +/// allocator built over borrowed, non-`'static` data (#1186 Phase 2 review, +/// round 4 finding 2). A later design (round 4) replaced `Any` with a plain, +/// `'static`-free `mechanism_id()` method instead, but coordinator review +/// (round 6) found that a self-reported identity — no matter how it is +/// carried — cannot prove what it was introduced to prove: a type can +/// honestly return `Some(self)` from both `as_virtual_backing()` and +/// whatever identity accessor exists, while its `VirtualBacking` method +/// bodies silently operate on unrelated state (see the module docs on +/// [`capability`](crate::capability) for the full accepted trust boundary). +/// Since no version of the identity check actually closed that gap, this +/// trait carries no identity method at all: it only ever bought a false +/// sense of safety at the cost of a real `'static` regression, so removing +/// it is a strict simplification, not a lost guarantee. A caller that needs +/// this class of proof needs the binding/provenance work in #1186 Phase 3, +/// not something layered onto this trait. pub trait DeviceAllocator: Send + Sync + Debug { /// Take `bytes` aligned to `align`. fn allocate(&self, bytes: usize, align: usize) -> Result, MemoryError>; @@ -203,17 +210,18 @@ pub trait DeviceAllocator: Send + Sync + Debug { /// /// The default forwards to [`deallocate`](Self::deallocate) and then, if /// [`as_shared_mapping`](Self::as_shared_mapping) reports a capability, - /// atomically releases and reports how many of `ptr`'s bytes were backed - /// by a shared-prefix mapping — an allocator that can create mapped cost + /// releases and reports how many of `ptr`'s bytes were backed by a + /// shared-prefix mapping — an allocator that can create mapped cost /// through [`SharedMapping::commit_shared_prefix`] must never lose that /// refund by silently inheriting a plain zero the way an eager allocator /// correctly does (#1186 Phase 2 review, round 3 finding 2), and must /// never report it from a stale snapshot taken before the actual release - /// (#1186 Phase 2 review, round 4 finding 1). It has no - /// [`VirtualBacking`] equivalent here on purpose: an allocator with that - /// capability instead releases through - /// [`VirtualBacking::deallocate_committed`], never through this default — - /// see that method's documentation. + /// (#1186 Phase 2 review, round 4 finding 1) — see + /// [`SharedMapping::release_shared_mapping`] for exactly what is and is + /// not guaranteed about that release. It has no [`VirtualBacking`] + /// equivalent here on purpose: an allocator with that capability instead + /// releases through [`VirtualBacking::deallocate_committed`], never + /// through this default — see that method's documentation. /// /// This is the unambiguously correct answer for any allocator with /// neither capability: nothing here is ever committed lazily or shared, @@ -264,30 +272,6 @@ pub trait DeviceAllocator: Send + Sync + Debug { fn as_shared_mapping(&self) -> Option<&dyn SharedMapping> { None } - - /// This mechanism's stable, unforgeable identity, if it has one. - /// - /// [`capability_shares_mechanism`](crate::capability::capability_shares_mechanism) - /// uses this to bind a `VirtualBacking`/`SharedMapping` capability to the - /// allocator that advertises it, rejecting a composing wrapper whose - /// capability is actually a different, foreign object (#1186 Phase 2 - /// review, round 1 finding 1) — including one whose address merely - /// coincides with the wrapper's own (round 3 finding 1). See - /// [`MechanismId`]'s own documentation for the full identity design. - /// - /// The default `None` is correct and imposes nothing on an implementation - /// that never overrides it: it simply cannot participate in - /// identity-checked capability binding, exactly like a type with no - /// capabilities at all. A concrete, `'static` type that wants - /// `as_virtual_backing`/`as_shared_mapping` to be trusted overrides this - /// with `Some(MechanismId::of(self))` — computed from its own `&Self` - /// inside its own `impl` block, never through a trait-object upcast, so - /// this method (and this trait) never has to require `Any` or `'static` - /// of every implementation the way an earlier design did (#1186 Phase 2 - /// review, round 4 finding 2). - fn mechanism_id(&self) -> Option { - None - } } /// Host memory from the global allocator. @@ -333,17 +317,6 @@ impl DeviceAllocator for HostAllocator { fn device(&self) -> DeviceKey { DeviceKey::HOST } - - /// `HostAllocator` has no capability to bind an identity to, but it is a - /// concrete, owned, `'static` type like any other, so it overrides this - /// as a simple illustration of the opt-in pattern every capability- - /// bearing mechanism follows (see [`capability_shares_mechanism`]'s - /// doctest). - /// - /// [`capability_shares_mechanism`]: crate::capability::capability_shares_mechanism - fn mechanism_id(&self) -> Option { - Some(MechanismId::of(self)) - } } #[cfg(test)] diff --git a/crates/onnx-runtime-memory-api/src/capability.rs b/crates/onnx-runtime-memory-api/src/capability.rs index e972cc4613..58ea1d5fc2 100644 --- a/crates/onnx-runtime-memory-api/src/capability.rs +++ b/crates/onnx-runtime-memory-api/src/capability.rs @@ -23,146 +23,62 @@ //! successful-looking no-op. A caller that needs the capability gets an //! `Option` to match on, not a method to call and hope was meaningful. //! -//! # Capability identity +//! # Capability identity — what Phase 2 does and does not prove //! //! [`VirtualBacking::device`] and [`SharedMapping::device`] must equal the -//! owning allocator's [`DeviceAllocator::device`]. +//! owning allocator's [`DeviceAllocator::device`], and a caller must check +//! that before trusting the two belong together — two unrelated arenas on the +//! same device compare equal by [`DeviceKey`] alone. +//! //! Discovery always goes through the same allocator reference a caller -//! already holds (`allocator.as_virtual_backing()`), so there is no second -//! handle that could go stale or be swapped independently — but a caller that -//! stores the returned `&dyn VirtualBacking`/`&dyn SharedMapping` alongside a -//! *different* allocator's identity (for instance, mixing up two devices, or -//! mixing a capability obtained from one allocator with a pointer allocated by -//! another) can still check `device()` itself before trusting the two belong -//! together. - -use std::any::{Any, TypeId}; +//! already holds (`allocator.as_virtual_backing()`), so there is no *second* +//! handle that could independently go stale or be swapped out from under a +//! caller holding the first one. +//! +//! That is the full extent of what this crate structurally proves. It does +//! **not** prove that a type's `VirtualBacking`/`SharedMapping` method bodies +//! actually operate on the same backing state as its own +//! [`DeviceAllocator::allocate`]/[`DeviceAllocator::deallocate`]. An earlier +//! design (#1186 Phase 2 review, rounds 1/3/4) tried to prove exactly that +//! with a self-reported identity token (compare a data pointer, then a data +//! pointer plus [`TypeId`](std::any::TypeId)), rejecting a composite type +//! whose advertised capability was a visibly different object. Coordinator +//! review (round 6) found that design fundamentally unsound in the general +//! case: nothing stops a single, honestly-self-identifying type from +//! internally forwarding its `DeviceAllocator` methods to one backing object +//! and its `VirtualBacking`/`SharedMapping` methods to an unrelated one — +//! the check only ever compares *which object* implements a trait, never +//! *what its methods do*, and no such self-attested token can close that gap +//! (a hostile or merely buggy implementor can report whatever identity it +//! likes; a legitimate delegating wrapper can just as easily be rejected by +//! the same check). The token apparatus was removed rather than kept as a +//! false assurance. +//! +//! Phase 2 therefore treats "a mechanism's `VirtualBacking`/`SharedMapping` +//! genuinely shares state with its own ordinary allocation/release" as a +//! **trusted-implementor invariant**, scoped to the mechanisms this +//! workspace ships and reviews (`HostAllocator`, `CudaVmmAllocator`) — not +//! something proven against an arbitrary or adversarial implementor. Closing +//! that gap for real requires a registry that *issues* capability access +//! rather than trusting a mechanism to self-report it, which is #1186 +//! Phase 3's "provider/context pinning and binding identity" work, not +//! Phase 2's. + +use std::any::Any; use std::ptr::NonNull; use crate::MemoryError; +// `DeviceAllocator` itself is never named in this module's actual code — only +// in intra-doc links (`[`DeviceAllocator::allocate`]` and friends) above and +// on individual items below, which need it in scope to resolve. Keep the +// import and silence the resulting unused-import lint rather than dropping +// it, since dropping it turns those doc links into rustdoc errors. +#[allow(unused_imports)] +use crate::allocator::DeviceAllocator; use crate::allocator::{ - AllocationCommitRange, DeviceAllocator, DeviceKey, SharedDevicePrefix, SharedPrefixCommitInfo, + AllocationCommitRange, DeviceKey, SharedDevicePrefix, SharedPrefixCommitInfo, }; -/// A concrete mechanism's stable, unforgeable identity. -/// -/// Two different [`DeviceAllocator`] instances must never compare equal -/// here, and one concrete instance must always compare equal to itself — -/// including across the different trait-object references a caller might -/// hold to it (its own `&dyn DeviceAllocator`, and any `&dyn -/// VirtualBacking`/`&dyn SharedMapping` it returns from -/// `as_virtual_backing`/`as_shared_mapping`). This is exactly what -/// [`capability_shares_mechanism`] compares, through -/// [`DeviceAllocator::mechanism_id`]. -/// -/// # Why not upcast `&dyn DeviceAllocator` to `&dyn Any` directly -/// -/// An earlier design (#1186 Phase 2 review, round 3) added `Any` as a -/// supertrait of [`DeviceAllocator`] so a `&dyn DeviceAllocator` could be -/// upcast to `&dyn Any` and compared by data pointer and [`TypeId`]. -/// `Any: 'static`, so that supertrait transitively forced *every* -/// `DeviceAllocator` implementation to be `'static` — a real regression -/// against "minimal ordinary contract", since it rejected an otherwise -/// valid allocator that borrows non-`'static` data, such as one built over -/// a caller-owned arena reference (#1186 Phase 2 review, round 4 finding 2). -/// -/// `MechanismId` itself carries no lifetime parameter: it stores an untyped -/// data pointer and a [`TypeId`], both plain, lifetime-free values. -/// Constructing one does still require the value it is computed from to be -/// `'static` (`TypeId::of` requires it), but that requirement now falls only -/// on the specific, concrete [`DeviceAllocator::mechanism_id`] override that -/// opts in — never on the trait itself, and never on an implementation that -/// has no capability to identify and simply keeps the default `None`. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub struct MechanismId { - ptr: *const (), - type_id: TypeId, -} - -impl MechanismId { - /// Compute the identity of a concrete, `'static` mechanism from its own - /// reference to itself. - /// - /// Call this as `MechanismId::of(self)` from inside a concrete type's - /// own [`DeviceAllocator::mechanism_id`] override (or an equivalent - /// override on [`VirtualBacking`]/[`SharedMapping`] for a capability - /// object that is a distinct value from the allocator) — never from - /// outside through a trait object, which is the one thing this design - /// exists to avoid requiring. - pub fn of(value: &T) -> Self { - Self { - ptr: (value as *const T).cast(), - type_id: TypeId::of::(), - } - } - - /// Compute the identity of a capability object already reached through - /// [`VirtualBacking::as_any`]/[`SharedMapping::as_any`]. - pub fn of_dyn(value: &dyn Any) -> Self { - Self { - ptr: (value as *const dyn Any).cast(), - type_id: value.type_id(), - } - } -} - -/// Whether `capability`, discovered from `allocator`, is genuinely the same -/// concrete mechanism as `allocator` itself — not a different internal object -/// that merely reports a matching [`DeviceKey`]. -/// -/// [`DeviceKey`] alone is not strong enough identity to bind a capability to -/// the allocator it came from: two unrelated arenas on the same device -/// compare equal by `DeviceKey`. This instead compares each side's -/// [`MechanismId`] — a data pointer paired with a concrete [`TypeId`], which -/// together cannot be produced by accident (see `MechanismId`'s own -/// documentation for why the pointer alone is forgeable and the `TypeId` -/// closes that gap). -/// -/// `allocator.mechanism_id()` must be `Some` and equal `capability`'s own -/// `MechanismId` (via [`MechanismId::of_dyn`]) for this to return `true`. An -/// allocator that never overrides [`DeviceAllocator::mechanism_id`] (the -/// default `None`) can never share a mechanism with anything through this -/// check — which is the correct, conservative answer for a type that has -/// opted out of identity-checked capabilities entirely, not a false -/// rejection. -/// -/// Every capability implementation this crate ships (`CudaVmmAllocator`'s -/// `as_virtual_backing`/`as_shared_mapping`, paired with its own -/// `mechanism_id` override) reports the same identity for the allocator and -/// for the capability object it hands out, so this always holds for them; -/// this exists to let a caller reject, rather than trust, an allocator -/// composed so that its ordinary [`DeviceAllocator::allocate`]/ -/// [`DeviceAllocator::deallocate`] and its advertised capability are backed -/// by two different objects — which would let a pointer produced through -/// the capability be freed through a mechanism that never produced it -/// (#1186 Phase 2 review, round 1 finding 1; round 3 finding 1 closed the -/// pointer-only version of this check; round 4 finding 2 moved the identity -/// off a blanket `Any` supertrait and onto this opt-in method). -/// -/// [`TypeId`]: std::any::TypeId -/// -/// # Example -/// -/// ``` -/// use std::any::Any; -/// use onnx_runtime_memory_api::allocator::{DeviceAllocator, DeviceKey, HostAllocator}; -/// use onnx_runtime_memory_api::capability::capability_shares_mechanism; -/// -/// let allocator = HostAllocator; -/// let allocator_ref: &dyn DeviceAllocator = &allocator; -/// let other: &dyn Any = &allocator; // a different reference to the *same* value -/// assert!(capability_shares_mechanism(allocator_ref, other)); -/// -/// let unrelated = HostAllocator; -/// let unrelated_ref: &dyn Any = &unrelated; -/// assert!(!capability_shares_mechanism(allocator_ref, unrelated_ref)); -/// ``` -pub fn capability_shares_mechanism(allocator: &dyn DeviceAllocator, capability: &dyn Any) -> bool { - allocator - .mechanism_id() - .is_some_and(|id| id == MechanismId::of_dyn(capability)) -} - /// Lazy commit/decommit over an allocator's own reservations. /// /// Implement this only when allocation and physical commitment are genuinely @@ -185,10 +101,9 @@ pub fn capability_shares_mechanism(allocator: &dyn DeviceAllocator, capability: /// [`DeviceAllocator::deallocate`]/[`DeviceAllocator::deallocate_with_unmapped`]: /// routing release through the same capability reference that produced the /// pointer (rather than re-deriving a release path from the owning -/// allocator) is what makes "the mechanism that allocated this is the -/// mechanism that frees it" true by construction instead of by convention — -/// see [`capability_shares_mechanism`] for why `DeviceKey` alone cannot -/// prove that on its own. Reports the number of bytes whose physical +/// allocator) is the discipline a caller must follow — this crate does not +/// structurally verify it (see this module's "Capability identity" notes). +/// Reports the number of bytes whose physical /// mapping actually transitioned from committed to uncommitted, the same /// accounting [`decommit_allocation_range`](Self::decommit_allocation_range) /// reports, so a mapped-zone refund is never silently lost the way an @@ -360,12 +275,15 @@ pub trait VirtualBacking: Send + Sync { /// * `commit_shared_prefix` maps `prefix` read-only into a live allocation /// and never mis-maps: it errors rather than mapping over an already /// committed region, mixing devices, or mixing pool authorities. -/// * `release_shared_mapping` atomically clears the shared-mapping bookkeeping -/// for one allocation and reports exactly the bytes that transitioned from -/// mapped to unmapped in that single call, so that mapped cost this -/// capability created is never lost the way an always-zero default would -/// lose it, and is never reported from a stale snapshot the way a -/// query-then-clear pair would risk under concurrency — symmetric with +/// * `release_shared_mapping` clears the shared-mapping bookkeeping for one +/// allocation as a single, non-interleavable operation from the mechanism's +/// own perspective (a lock, a compare-and-swap loop, or equivalent — see +/// its own documentation for exactly what this does and does not +/// guarantee), and reports exactly the bytes that transitioned from mapped +/// to unmapped in that call, so that mapped cost this capability created is +/// never lost the way an always-zero default would lose it, and is never +/// reported from a stale snapshot the way a query-then-clear pair would +/// risk under concurrency — symmetric with /// [`VirtualBacking::deallocate_committed`]'s release report, and required /// for the same reason: `SharedMapping` and `VirtualBacking` are /// independent capabilities, so a mechanism that implements only this one @@ -400,9 +318,9 @@ pub trait SharedMapping: Send + Sync { byte_offset: usize, ) -> Result; - /// Atomically release the shared-mapping bookkeeping for the allocation - /// at `ptr` and report exactly the bytes that transitioned from mapped - /// to unmapped as a result of **this** call. + /// Release the shared-mapping bookkeeping for the allocation at `ptr` + /// and report exactly the bytes that transitioned from mapped to + /// unmapped as a result of **this** call. /// /// [`DeviceAllocator::deallocate_with_unmapped`]'s default calls this /// whenever [`DeviceAllocator::as_shared_mapping`] returns `Some`, so a @@ -421,22 +339,35 @@ pub trait SharedMapping: Send + Sync { /// /// # Contract /// - /// This is a release operation, not a query: it must atomically observe - /// and clear whatever bookkeeping tracks this allocation's shared-mapped - /// bytes, under one mechanism synchronization boundary (a lock, a - /// compare-and-swap loop, or equivalent) so no concurrent caller can - /// observe the accounting between the read and the clear. A prior - /// version of this contract split those two steps into a pure query - /// method and a separate `deallocate` call; that gap let a concurrent - /// mapping change, a cached free, or a failed release be reported as if - /// it had actually happened (#1186 Phase 2 review, round 4 finding 1). + /// This is a release operation, not a query: from the perspective of this + /// mechanism's own in-memory bookkeeping, it must observe and clear + /// whatever tracks this allocation's shared-mapped bytes as one + /// non-interleavable step (a lock, a compare-and-swap loop, or + /// equivalent) so no concurrent caller can observe the accounting between + /// the read and the clear. A prior version of this contract split those + /// two steps into a pure query method and a separate `deallocate` call; + /// that gap let a concurrent mapping change, a cached free, or a failed + /// release be reported as if it had actually happened (#1186 Phase 2 + /// review, round 4 finding 1). The guarantees below are about that + /// in-memory bookkeeping; they are **not** a claim that an underlying + /// multi-step OS/driver release (for instance, unmap-then-release-handle) + /// is itself indivisible. If such an operation can succeed partway and + /// then fail, that partial, possibly irreversible mutation is a property + /// of the backend, not something this contract can retroactively undo — + /// implementations must document any such gap explicitly rather than + /// imply it does not exist. A caller that needs a release to be provably + /// safe against that class of partial hardware/driver failure needs an + /// owning allocation token and a deferred, reconciled release state + /// machine, which is #1186 Phase 4's scope, not this trait's. /// - /// * **Exact once**: a `ptr` whose shared mapping this call has already - /// released must report `0` on every subsequent call — it must not be - /// possible to refund the same bytes twice. - /// * **Failure refunds zero and preserves state**: if the underlying - /// release cannot complete, this returns `0` and leaves the mapping - /// bookkeeping exactly as it was, so a caller may safely retry. + /// * **Exact once** (bookkeeping): a `ptr` whose shared mapping this call + /// has already released must report `0` on every subsequent call — it + /// must not be possible to refund the same bookkept bytes twice. + /// * **Failure refunds zero and preserves state** (bookkeeping): if the + /// underlying release cannot complete, this returns `0` and leaves the + /// mapping bookkeeping exactly as it was, so a caller may safely retry + /// — except insofar as the backend itself already made an irreversible + /// partial change outside this bookkeeping, per the paragraph above. /// * **Concurrency**: multiple threads calling this for the same `ptr` /// must divide the true mapped total exactly once among them (any /// ordering is fine; double-counting or losing bytes is not). @@ -463,91 +394,25 @@ mod tests { use std::collections::HashMap; use std::sync::Mutex; - use crate::allocator::HostAllocator; - - /// A trivial `VirtualBacking` that is never actually called: it exists - /// only so its address can be compared, never so its methods run. + use crate::allocator::{DeviceAllocator, HostAllocator}; + + /// A `DeviceAllocator` whose `VirtualBacking` capability is *itself* + /// (`as_virtual_backing` returns `Some(self)`, the idiomatic, honestly + /// self-identifying pattern) but whose `VirtualBacking` method bodies + /// silently forward to a second, unrelated backing (`foreign`) rather + /// than the same state its own `allocate`/`deallocate` use. This is the + /// case coordinator round-6 review found no self-reported identity token + /// can distinguish from a genuinely consistent implementation: both + /// present the exact same object as "the" `VirtualBacking`, so nothing + /// short of inspecting method bodies (impossible from outside the type) + /// tells them apart. The test below documents this honestly as an + /// accepted trust boundary rather than pretending it is checked. #[derive(Debug, Default)] - struct Marker; - - impl VirtualBacking for Marker { - fn device(&self) -> DeviceKey { - DeviceKey::device(0) - } - - fn allocate_committed( - &self, - _bytes: usize, - _align: usize, - _committed_ranges: &[std::ops::Range], - ) -> Result, MemoryError> { - unimplemented!("test double: never called") - } - - unsafe fn deallocate_committed( - &self, - _ptr: NonNull, - _allocation_bytes: usize, - _align: usize, - ) -> u64 { - unimplemented!("test double: never called") - } - - fn commit_allocation_range( - &self, - _ptr: NonNull, - _allocation_bytes: usize, - _align: usize, - _offset: usize, - _bytes: usize, - ) -> Result<(), MemoryError> { - unimplemented!("test double: never called") - } - - fn mapped_bytes_for_allocation_ranges( - &self, - _ranges: &[AllocationCommitRange], - ) -> Result { - unimplemented!("test double: never called") - } - - fn mapped_bytes_for_allocation(&self, _bytes: usize, _align: usize) -> Result { - unimplemented!("test double: never called") - } - - fn decommit_allocation_range( - &self, - _ptr: NonNull, - _allocation_bytes: usize, - _align: usize, - _offset: usize, - _bytes: usize, - ) -> Result { - unimplemented!("test double: never called") - } - - fn allocation_committed_bytes( - &self, - _ptr: NonNull, - _allocation_bytes: usize, - _align: usize, - ) -> usize { - unimplemented!("test double: never called") - } - - fn as_any(&self) -> &dyn Any { - self - } + struct SplitInnerWrapper { + foreign_committed: std::sync::atomic::AtomicU64, } - /// Genuinely returns itself as its own `VirtualBacking`, the way a real - /// combined allocator/capability type does — as opposed to - /// `OffsetZeroComposite` below, which returns an *embedded* object that - /// merely shares its address. - #[derive(Debug, Default)] - struct HonestSelfReporter; - - impl DeviceAllocator for HonestSelfReporter { + impl DeviceAllocator for SplitInnerWrapper { fn allocate(&self, bytes: usize, align: usize) -> Result, MemoryError> { HostAllocator.allocate(bytes, align) } @@ -562,35 +427,40 @@ mod tests { } fn as_virtual_backing(&self) -> Option<&dyn VirtualBacking> { + // Idiomatic and honest at the object-identity level: returns + // `self`, not a visibly different object. Some(self) } - - fn mechanism_id(&self) -> Option { - Some(MechanismId::of(self)) - } } - impl VirtualBacking for HonestSelfReporter { + impl VirtualBacking for SplitInnerWrapper { fn device(&self) -> DeviceKey { DeviceKey::device(0) } fn allocate_committed( &self, - _bytes: usize, + bytes: usize, _align: usize, _committed_ranges: &[std::ops::Range], ) -> Result, MemoryError> { - unimplemented!("test double: never called") + // Deliberately tracks state that `DeviceAllocator::allocate` + // above never touches, standing in for "a totally different + // backing object" without needing a second real allocator. + self.foreign_committed + .fetch_add(bytes as u64, std::sync::atomic::Ordering::AcqRel); + HostAllocator.allocate(bytes, 8) } unsafe fn deallocate_committed( &self, _ptr: NonNull, - _allocation_bytes: usize, + allocation_bytes: usize, _align: usize, ) -> u64 { - unimplemented!("test double: never called") + self.foreign_committed + .fetch_sub(allocation_bytes as u64, std::sync::atomic::Ordering::AcqRel); + allocation_bytes as u64 } fn commit_allocation_range( @@ -640,86 +510,49 @@ mod tests { } } - /// The exact composite the round-3 review flagged: `#[repr(C)]` puts - /// `foreign` at byte offset zero, so `&composite as *const _ as *const - /// ()` and `&composite.foreign as *const _ as *const ()` are the same - /// address, even though `foreign` is a distinct, embedded object with its - /// own type — not `composite` itself. Overrides `mechanism_id` with its - /// own (wrapper) identity so the rejection test below is non-vacuous: it - /// proves the `TypeId` mismatch defeats the address collision, not merely - /// that the wrapper lacks identity support at all. - #[repr(C)] - #[derive(Debug)] - struct OffsetZeroComposite { - foreign: Marker, - _extra: u64, - } - - impl DeviceAllocator for OffsetZeroComposite { - fn allocate(&self, bytes: usize, align: usize) -> Result, MemoryError> { - HostAllocator.allocate(bytes, align) - } - - unsafe fn deallocate(&self, ptr: NonNull, bytes: usize, align: usize) { - // SAFETY: forwarded under this method's identical contract. - unsafe { HostAllocator.deallocate(ptr, bytes, align) }; - } - - fn device(&self) -> DeviceKey { - DeviceKey::device(0) - } - - fn as_virtual_backing(&self) -> Option<&dyn VirtualBacking> { - // Deliberately the embedded first field, not `self`. - Some(&self.foreign) - } - - fn mechanism_id(&self) -> Option { - // Deliberately the composite's own identity, not `self.foreign`'s - // — proving the rejection below comes from a genuine `TypeId` - // mismatch, not from `OffsetZeroComposite` having opted out of - // identity entirely. - Some(MechanismId::of(self)) - } - } - - #[test] - fn genuine_self_reference_still_shares_mechanism() { - let allocator = HonestSelfReporter; - let allocator_ref: &dyn DeviceAllocator = &allocator; - let virtual_backing = allocator_ref.as_virtual_backing().expect("advertised"); - assert!( - capability_shares_mechanism(allocator_ref, virtual_backing.as_any()), - "an allocator that honestly returns `Some(self)` (through its own field, reached by \ - address) must still be recognized as sharing its own mechanism" - ); - } - #[test] - fn offset_zero_embedded_capability_is_not_treated_as_the_same_mechanism() { - let composite = OffsetZeroComposite { - foreign: Marker, - _extra: 0, - }; - - // Non-vacuous: prove the address collision this test exists to guard - // against is real, not a hypothetical. - let composite_addr = &composite as *const OffsetZeroComposite as *const (); - let foreign_addr = &composite.foreign as *const Marker as *const (); - assert!( - std::ptr::eq(composite_addr, foreign_addr), - "test setup did not reproduce the offset-zero address collision `#[repr(C)]` is \ - supposed to guarantee" + fn documented_limitation_a_same_self_capability_can_still_track_unrelated_state() { + // This is not a desirable property; it is a regression guard on the + // documented trust boundary in this module's top-level docs. If this + // test starts failing, either `VirtualBacking`/`DeviceAllocator` grew + // a way to detect this (update the docs to say so) or `allocate` + // stopped forwarding to `HostAllocator` (unrelated breakage). + let wrapper = SplitInnerWrapper::default(); + let allocator_ref: &dyn DeviceAllocator = &wrapper; + let virtual_backing = allocator_ref + .as_virtual_backing() + .expect("SplitInnerWrapper advertises VirtualBacking"); + + // `committed_ranges` is `&[std::ops::Range]`, an actual list + // of ranges this capability tracks, not a `Vec` in disguise — + // a single-element array literal here is a deliberate, correct call, + // not the ambiguous pattern the lint exists to catch. + #[allow(clippy::single_range_in_vec_init)] + let committed_ranges = [0..64]; + let ptr = virtual_backing + .allocate_committed(64, 8, &committed_ranges) + .expect("allocate_committed granted"); + assert_eq!( + wrapper + .foreign_committed + .load(std::sync::atomic::Ordering::Acquire), + 64, + "allocate_committed must have run — proving it really is disjoint bookkeeping from DeviceAllocator::allocate, not merely a hypothetical" ); - let allocator_ref: &dyn DeviceAllocator = &composite; - let virtual_backing = allocator_ref.as_virtual_backing().expect("advertised"); - assert!( - !capability_shares_mechanism(allocator_ref, virtual_backing.as_any()), - "a foreign `VirtualBacking` embedded at a wrapper's offset zero must never be \ - treated as the wrapper's own mechanism just because their addresses coincide \ - (#1186 Phase 2 review, round 3 finding 1)" + // SAFETY: `ptr` came from this capability's own `allocate_committed` + // above with matching bytes/align, released exactly once here. + unsafe { virtual_backing.deallocate_committed(ptr, 64, 8) }; + assert_eq!( + wrapper + .foreign_committed + .load(std::sync::atomic::Ordering::Acquire), + 0, + "deallocate_committed must have run against the same disjoint bookkeeping" ); + // SAFETY: `ptr` also came from `HostAllocator::allocate` (what our + // `DeviceAllocator::allocate` forwards to) with matching bytes/align. + unsafe { allocator_ref.deallocate(ptr, 64, 8) }; } /// A `SharedMapping`-only allocator: no `VirtualBacking` at all, so its @@ -1042,10 +875,10 @@ mod tests { DeviceKey::device(0) } - // Deliberately does not override `mechanism_id`, `as_virtual_backing`, - // or `as_shared_mapping`: a borrowed, non-`'static` allocator is - // exactly the kind of type that should be able to implement only the - // minimal ordinary contract and nothing more. + // Deliberately does not override `as_virtual_backing` or + // `as_shared_mapping`: a borrowed, non-`'static` allocator is exactly + // the kind of type that should be able to implement only the minimal + // ordinary contract and nothing more. } #[test] @@ -1053,8 +886,9 @@ mod tests { let arena = BorrowedArena::default(); let allocator = BorrowedAllocator { arena: &arena }; - // This line is the actual proof: `DeviceAllocator` no longer requires - // `Any`, so a trait object over a non-`'static` reference compiles. + // This line is the actual proof: `DeviceAllocator` does not require + // `Any`/`'static`, so a trait object over a non-`'static` reference + // compiles. let as_dyn: &dyn DeviceAllocator = &allocator; let ptr = as_dyn.allocate(64, 8).expect("borrowed allocation"); @@ -1064,11 +898,6 @@ mod tests { 64, "allocate must still run through the borrowed arena" ); - assert!( - as_dyn.mechanism_id().is_none(), - "an allocator that never overrides `mechanism_id` must report `None`, not be forced \ - to fabricate an identity (#1186 Phase 2 review, round 4 finding 2)" - ); // SAFETY: `ptr` came from this allocator's own `allocate` above with // matching `bytes`/`align`, and is released exactly once here. diff --git a/crates/onnx-runtime-memory-api/src/lib.rs b/crates/onnx-runtime-memory-api/src/lib.rs index ecd350b6c4..af16cd3d49 100644 --- a/crates/onnx-runtime-memory-api/src/lib.rs +++ b/crates/onnx-runtime-memory-api/src/lib.rs @@ -47,7 +47,7 @@ pub use allocator::{ AllocationCommitRange, DeviceAllocator, DeviceKey, HostAllocator, MappedAllocation, SharedDevicePrefix, SharedPrefixCommitInfo, }; -pub use capability::{SharedMapping, VirtualBacking, capability_shares_mechanism}; +pub use capability::{SharedMapping, VirtualBacking}; /// Where the bytes physically live. /// diff --git a/crates/onnx-runtime-memory-governor/src/allocator.rs b/crates/onnx-runtime-memory-governor/src/allocator.rs index 62f795e3ab..daada9fd1c 100644 --- a/crates/onnx-runtime-memory-governor/src/allocator.rs +++ b/crates/onnx-runtime-memory-governor/src/allocator.rs @@ -32,9 +32,7 @@ pub use onnx_runtime_memory_api::allocator::{ AllocationCommitRange, DeviceAllocator, DeviceKey, HostAllocator, MappedAllocation, SharedDevicePrefix, SharedPrefixCommitInfo, }; -pub use onnx_runtime_memory_api::capability::{ - MechanismId, SharedMapping, VirtualBacking, capability_shares_mechanism, -}; +pub use onnx_runtime_memory_api::capability::{SharedMapping, VirtualBacking}; #[cfg(test)] mod tests { diff --git a/crates/onnx-runtime-memory-governor/src/lib.rs b/crates/onnx-runtime-memory-governor/src/lib.rs index 8c175dc208..01c2cf3bae 100644 --- a/crates/onnx-runtime-memory-governor/src/lib.rs +++ b/crates/onnx-runtime-memory-governor/src/lib.rs @@ -76,8 +76,7 @@ pub mod shareability; pub use allocator::{ AllocationCommitRange, DeviceAllocator, DeviceKey, HostAllocator, MappedAllocation, - MechanismId, SharedDevicePrefix, SharedMapping, SharedPrefixCommitInfo, VirtualBacking, - capability_shares_mechanism, + SharedDevicePrefix, SharedMapping, SharedPrefixCommitInfo, VirtualBacking, }; pub use large_alloc_cache::{ DEFAULT_CACHE_BUDGET_BYTES, FALLBACK_FLOOR_BYTES, LargeAllocCache, LargeAllocCacheStats,