nuvec_gpu/
memory.rs

1//! Allocator statistics (shared by every target).
2
3/// Memory use and limits of a device's allocator, in bytes.
4///
5/// `peak_with_mapped` is nuvec's one peak-memory metric: compare it with MLX's
6/// `get_peak_memory()` (which counts MLX's copied weights), each reset after loading, in a
7/// fresh process per phase, cross-checked by the operating system's peak resident size.
8#[non_exhaustive]
9#[derive(Debug, Clone, Copy, PartialEq, Eq)]
10pub struct MemoryStats {
11    /// Bytes in live allocator buffers (rounded allocation sizes), weights copied from
12    /// checkpoints included; mapped files not included. MLX: `get_active_memory`.
13    pub active: usize,
14    /// The highest `active` since creation or the last `Device::reset_peak`; mapped files
15    /// not included.
16    pub peak: usize,
17    /// `active + mapped`: everything the allocator holds for live arrays, mapped
18    /// checkpoints included. MLX: `get_active_memory` (MLX copies every weight, so its
19    /// active memory includes them).
20    pub active_with_mapped: usize,
21    /// The highest `active + mapped` since creation or the last `Device::reset_peak`,
22    /// updated on every change of either: **the one peak metric** (see above). MLX:
23    /// `get_peak_memory`.
24    pub peak_with_mapped: usize,
25    /// Bytes held in the cache for reuse (not in `active`). MLX: `get_cache_memory`.
26    pub cache: usize,
27    /// Allocations and mappings that would push `active + mapped` past this fail with
28    /// `OutOfMemory` (a hard refusal; MLX's `set_memory_limit` is a guideline).
29    pub memory_limit: usize,
30    /// The cache is trimmed to at most this many bytes. MLX: `set_cache_limit`.
31    pub cache_limit: usize,
32    /// Buffers obtained from Metal (cache misses) so far.
33    pub system_allocations: u64,
34    /// Buffers and mapped files are kept resident up to this many bytes (see
35    /// `Device::set_wired_limit`; MLX: `set_wired_limit`).
36    pub wired_limit: usize,
37    /// Bytes currently wired, mapped files included (`wired ≤ wired_limit`).
38    pub wired: usize,
39    /// Bytes of live memory-mapped files (checkpoints loaded without a copy): each file's
40    /// whole length, header included, rounded up to pages. Not in `active` or `peak`. MLX
41    /// does not map files.
42    pub mapped: usize,
43    /// The part of `wired` that belongs to memory-mapped files.
44    pub wired_mapped: usize,
45}
46
47/// The virtual-memory page size of Apple silicon (the allocator's large-buffer granule).
48const PAGE: usize = 16 * 1024;
49/// The allocator's granule for buffers up to one page.
50const SMALL: usize = 256;
51
52/// The bytes the allocator charges (in `active`) for a buffer of `len` bytes: its size
53/// class (256-byte steps up to one page, whole pages beyond); `None` if that overflows.
54pub fn allocation_size(len: usize) -> Option<usize> {
55    match len {
56        n if n > PAGE => n.checked_next_multiple_of(PAGE),
57        n => Some(n.div_ceil(SMALL) * SMALL),
58    }
59}
60
61/// The bytes mapping a file of `len` bytes counts in `mapped`: the whole file, rounded up
62/// to the system's pages.
63pub fn mapping_size(len: usize) -> usize {
64    len.div_ceil(page()) * page()
65}
66
67/// The system's page size.
68pub(crate) fn page() -> usize {
69    #[cfg(target_os = "macos")]
70    {
71        // SAFETY: `sysconf` has no preconditions.
72        let page = unsafe { libc::sysconf(libc::_SC_PAGESIZE) };
73        usize::try_from(page).unwrap_or(PAGE)
74    }
75    #[cfg(not(target_os = "macos"))]
76    PAGE
77}
nuvec_gpu/memory.rs

nuvec_gpu/
memory.rs
