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, Default)]
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    /// Residency sets holding the wired buffers (MLX's capped sets: at most 5% of the
46    /// working set each, `MLX_RESIDENCY_SET_MAX_PCT`; 0 without residency-set support).
47    pub residency_sets: usize,
48}
49
50/// The virtual-memory page size of Apple silicon (the allocator's large-buffer granule).
51const PAGE: usize = 16 * 1024;
52/// The allocator's granule for buffers up to one page.
53const SMALL: usize = 256;
54
55/// The bytes the allocator charges (in `active`) for a buffer of `len` bytes: its size
56/// class (256-byte steps up to one page, whole pages beyond); `None` if that overflows.
57pub fn allocation_size(len: usize) -> Option<usize> {
58    match len {
59        n if n > PAGE => n.checked_next_multiple_of(PAGE),
60        n => Some(n.div_ceil(SMALL) * SMALL),
61    }
62}
63
64/// The bytes mapping a file of `len` bytes counts in `mapped`: the whole file, rounded up
65/// to the system's pages.
66pub fn mapping_size(len: usize) -> usize {
67    len.div_ceil(page()) * page()
68}
69
70/// The system's page size: the granularity of file mappings (`map_file_range` offsets).
71pub fn page_size() -> usize {
72    page()
73}
74
75/// The system's page size.
76pub(crate) fn page() -> usize {
77    #[cfg(target_os = "macos")]
78    {
79        // SAFETY: `sysconf` has no preconditions.
80        let page = unsafe { libc::sysconf(libc::_SC_PAGESIZE) };
81        usize::try_from(page).unwrap_or(PAGE)
82    }
83    #[cfg(not(target_os = "macos"))]
84    PAGE
85}
nuvec_gpu/memory.rs

nuvec_gpu/
memory.rs
