onnx_runtime_memory_api/
capability.rs

1//! Optional capabilities discovered from an already-selected allocator.
2//!
3//! Eager allocation needs neither trait. A VMM allocator exposes
4//! [`VirtualBacking`]; a mechanism with reusable physical handles independently
5//! exposes [`SharedMapping`]. Neither capability owns whole-allocation release.
6
7use std::fmt::Debug;
8use std::ptr::NonNull;
9
10use crate::{AllocationCommitRange, MemoryError, SharedDevicePrefix, SharedPrefixCommitInfo};
11
12/// Lazy virtual reservation and physical commit/decommit.
13///
14/// All pointers accepted or returned here belong to the coherent
15/// [`crate::DeviceAllocator`] from which this capability was discovered.
16/// Terminal release must go through that allocator's canonical release path.
17pub trait VirtualBacking: Send + Sync + Debug {
18    /// Reserve one allocation while committing only the ranges immediately live.
19    fn allocate_committed(
20        &self,
21        bytes: usize,
22        align: usize,
23        committed_ranges: &[std::ops::Range<usize>],
24    ) -> Result<NonNull<u8>, MemoryError>;
25
26    fn commit_allocation_range(
27        &self,
28        ptr: NonNull<u8>,
29        allocation_bytes: usize,
30        align: usize,
31        offset: usize,
32        bytes: usize,
33    ) -> Result<(), MemoryError>;
34
35    fn commit_allocation_ranges(
36        &self,
37        ranges: &[AllocationCommitRange],
38    ) -> Result<(), MemoryError> {
39        for range in ranges {
40            self.commit_allocation_range(
41                range.ptr,
42                range.allocation_bytes,
43                range.align,
44                range.offset,
45                range.bytes,
46            )?;
47        }
48        Ok(())
49    }
50
51    /// Conservative mapped-byte estimate for a batched commit.
52    ///
53    /// Required because only the concrete mechanism knows its granularity and
54    /// whether disjoint tiny ranges share physical granules.
55    fn mapped_bytes_for_allocation_ranges(
56        &self,
57        ranges: &[AllocationCommitRange],
58    ) -> Result<u64, MemoryError>;
59
60    fn mapped_bytes_for_allocation(&self, bytes: usize, align: usize) -> Result<u64, MemoryError>;
61
62    /// Release physical backing while retaining the virtual allocation.
63    ///
64    /// Returns actual newly unmapped bytes after granularity and shared
65    /// references are applied. This is not whole-allocation release.
66    fn decommit_allocation_range(
67        &self,
68        ptr: NonNull<u8>,
69        allocation_bytes: usize,
70        align: usize,
71        offset: usize,
72        bytes: usize,
73    ) -> Result<u64, MemoryError>;
74
75    fn allocation_committed_bytes(
76        &self,
77        ptr: NonNull<u8>,
78        allocation_bytes: usize,
79        align: usize,
80    ) -> usize;
81}
82
83/// Reusable shared physical handles and read-only prefix mappings.
84///
85/// This capability is independent of [`VirtualBacking`]. A pool-less VMM may
86/// expose virtual backing but no shared mapping, and another coherent
87/// mechanism may expose shared mapping without virtual backing.
88pub trait SharedMapping: Send + Sync + Debug {
89    fn create_shared_prefix(
90        &self,
91        bytes: usize,
92    ) -> Result<Box<dyn SharedDevicePrefix>, MemoryError>;
93
94    /// Incremental owned physical cost of admitting another mapping.
95    ///
96    /// Zero is valid only for a prefix this capability can actually map. A
97    /// wrong-device, wrong-authority, or foreign prefix must be rejected before
98    /// any cost is reported, matching
99    /// [`commit_shared_prefix`](Self::commit_shared_prefix).
100    fn incremental_owned_bytes_for_shared_prefix(
101        &self,
102        prefix: &dyn SharedDevicePrefix,
103    ) -> Result<u64, MemoryError>;
104
105    /// Map the prefix transactionally.
106    ///
107    /// Returning `Err` must leave no mapping from this call in the target
108    /// allocation. Admission coordinators rely on that invariant to reject the
109    /// first half of a K/V pair without exposing partially shared state.
110    fn commit_shared_prefix(
111        &self,
112        prefix: &dyn SharedDevicePrefix,
113        ptr: NonNull<u8>,
114        allocation_bytes: usize,
115        byte_offset: usize,
116    ) -> Result<SharedPrefixCommitInfo, MemoryError>;
117}
118
119#[cfg(test)]
120mod tests {
121    use super::*;
122    use crate::{DeviceAllocator, DeviceKey, HostAllocator};
123    use std::any::Any;
124    use std::sync::atomic::{AtomicU64, Ordering};
125
126    #[derive(Debug, Default)]
127    struct VirtualOnly {
128        commits: AtomicU64,
129        decommits: AtomicU64,
130    }
131
132    impl DeviceAllocator for VirtualOnly {
133        fn allocate(&self, bytes: usize, align: usize) -> Result<NonNull<u8>, MemoryError> {
134            HostAllocator.allocate(bytes, align)
135        }
136
137        unsafe fn deallocate(&self, ptr: NonNull<u8>, bytes: usize, align: usize) {
138            // SAFETY: canonical release forwards the exact allocation.
139            unsafe { HostAllocator.deallocate(ptr, bytes, align) };
140        }
141
142        fn device(&self) -> DeviceKey {
143            DeviceKey::HOST
144        }
145
146        fn as_virtual_backing(&self) -> Option<&dyn VirtualBacking> {
147            Some(self)
148        }
149    }
150
151    impl VirtualBacking for VirtualOnly {
152        fn allocate_committed(
153            &self,
154            bytes: usize,
155            align: usize,
156            _ranges: &[std::ops::Range<usize>],
157        ) -> Result<NonNull<u8>, MemoryError> {
158            self.commits.fetch_add(1, Ordering::Relaxed);
159            HostAllocator.allocate(bytes, align)
160        }
161
162        fn commit_allocation_range(
163            &self,
164            _ptr: NonNull<u8>,
165            _allocation_bytes: usize,
166            _align: usize,
167            _offset: usize,
168            _bytes: usize,
169        ) -> Result<(), MemoryError> {
170            self.commits.fetch_add(1, Ordering::Relaxed);
171            Ok(())
172        }
173
174        fn mapped_bytes_for_allocation_ranges(
175            &self,
176            ranges: &[AllocationCommitRange],
177        ) -> Result<u64, MemoryError> {
178            Ok(ranges.iter().map(|range| range.bytes as u64).sum())
179        }
180
181        fn mapped_bytes_for_allocation(
182            &self,
183            bytes: usize,
184            _align: usize,
185        ) -> Result<u64, MemoryError> {
186            Ok(bytes as u64)
187        }
188
189        fn decommit_allocation_range(
190            &self,
191            _ptr: NonNull<u8>,
192            _allocation_bytes: usize,
193            _align: usize,
194            _offset: usize,
195            bytes: usize,
196        ) -> Result<u64, MemoryError> {
197            self.decommits.fetch_add(1, Ordering::Relaxed);
198            Ok(bytes as u64)
199        }
200
201        fn allocation_committed_bytes(
202            &self,
203            _ptr: NonNull<u8>,
204            allocation_bytes: usize,
205            _align: usize,
206        ) -> usize {
207            allocation_bytes
208        }
209    }
210
211    #[derive(Debug)]
212    struct SharedOnly;
213
214    #[derive(Debug)]
215    struct ForeignZeroPrefix;
216
217    impl SharedDevicePrefix for ForeignZeroPrefix {
218        fn device_ptr(&self) -> u64 {
219            0
220        }
221
222        fn committed_physical_bytes(&self) -> u64 {
223            0
224        }
225
226        fn mapped_bytes(&self) -> usize {
227            0
228        }
229
230        fn requested_bytes(&self) -> usize {
231            0
232        }
233
234        fn as_any(&self) -> &dyn Any {
235            self
236        }
237    }
238
239    impl DeviceAllocator for SharedOnly {
240        fn allocate(&self, bytes: usize, align: usize) -> Result<NonNull<u8>, MemoryError> {
241            HostAllocator.allocate(bytes, align)
242        }
243
244        unsafe fn deallocate(&self, ptr: NonNull<u8>, bytes: usize, align: usize) {
245            // SAFETY: canonical release forwards the exact allocation.
246            unsafe { HostAllocator.deallocate(ptr, bytes, align) };
247        }
248
249        fn device(&self) -> DeviceKey {
250            DeviceKey::HOST
251        }
252
253        fn as_shared_mapping(&self) -> Option<&dyn SharedMapping> {
254            Some(self)
255        }
256    }
257
258    impl SharedMapping for SharedOnly {
259        fn create_shared_prefix(
260            &self,
261            bytes: usize,
262        ) -> Result<Box<dyn SharedDevicePrefix>, MemoryError> {
263            Err(MemoryError::InvalidRequest {
264                tier: "host",
265                requested: bytes as u64,
266                reason: "test shared prefix construction is intentionally unsupported",
267            })
268        }
269
270        fn incremental_owned_bytes_for_shared_prefix(
271            &self,
272            prefix: &dyn SharedDevicePrefix,
273        ) -> Result<u64, MemoryError> {
274            Err(MemoryError::InvalidRequest {
275                tier: "host",
276                requested: prefix.requested_bytes() as u64,
277                reason: "test shared prefix is foreign to this capability",
278            })
279        }
280
281        fn commit_shared_prefix(
282            &self,
283            _prefix: &dyn SharedDevicePrefix,
284            _ptr: NonNull<u8>,
285            allocation_bytes: usize,
286            _byte_offset: usize,
287        ) -> Result<SharedPrefixCommitInfo, MemoryError> {
288            Err(MemoryError::InvalidRequest {
289                tier: "host",
290                requested: allocation_bytes as u64,
291                reason: "test shared mapping is intentionally unsupported",
292            })
293        }
294    }
295
296    #[test]
297    fn virtual_backing_is_discovered_and_used_before_canonical_release() {
298        let allocator = VirtualOnly::default();
299        let ordinary: &dyn DeviceAllocator = &allocator;
300        let backing = ordinary
301            .as_virtual_backing()
302            .expect("virtual backing capability");
303        let initial = 0..8;
304        let ptr = backing
305            .allocate_committed(64, 16, std::slice::from_ref(&initial))
306            .expect("reserve and commit");
307        backing
308            .commit_allocation_range(ptr, 64, 16, 8, 8)
309            .expect("additional commit");
310        assert_eq!(
311            backing
312                .decommit_allocation_range(ptr, 64, 16, 8, 8)
313                .expect("partial decommit"),
314            8
315        );
316        assert_eq!(backing.allocation_committed_bytes(ptr, 64, 16), 64);
317        assert_eq!(allocator.commits.load(Ordering::Relaxed), 2);
318        assert_eq!(allocator.decommits.load(Ordering::Relaxed), 1);
319        // SAFETY: whole-allocation release remains on the ordinary allocator.
320        unsafe { ordinary.deallocate(ptr, 64, 16) };
321    }
322
323    #[test]
324    fn shared_mapping_discovery_is_independent_of_virtual_backing() {
325        let allocator: &dyn DeviceAllocator = &SharedOnly;
326        assert!(allocator.as_virtual_backing().is_none());
327        let mapping = allocator
328            .as_shared_mapping()
329            .expect("independent shared mapping capability");
330        assert!(
331            mapping
332                .incremental_owned_bytes_for_shared_prefix(&ForeignZeroPrefix)
333                .is_err(),
334            "foreign input must be rejected before even a zero cost is reported"
335        );
336    }
337}
onnx_runtime_memory_api/capability.rs

onnx_runtime_memory_api/
capability.rs
