nuvec_gpu/metal/
allocator.rs

1//! The buffer cache: MLX's Metal allocator rules (`mlx/backend/metal/allocator.cpp`).
2//!
3//! Freed buffers are cached and handed back for similar sizes, so steady-state loops make
4//! almost no system allocations. The cache is bounded (least-recently-used eviction), and
5//! allocations beyond the memory limit fail instead of exhausting shared system RAM.
6
7use std::collections::{BTreeMap, HashSet};
8use std::sync::{Mutex, MutexGuard, PoisonError};
9
10use objc2::rc::Retained;
11use objc2::runtime::ProtocolObject;
12use objc2_metal::{MTLBuffer, MTLDevice, MTLResourceOptions};
13
14use crate::{GpuError, Result};
15
16/// The virtual-memory page size on Apple silicon.
17pub(crate) const PAGE: usize = 16 * 1024;
18/// Granularity of allocations up to one page.
19const SMALL: usize = 256;
20/// Number of small size classes: 256 B steps up to one page.
21const SMALL_CLASSES: usize = PAGE / SMALL + 1;
22
23pub(crate) type Raw = Retained<ProtocolObject<dyn MTLBuffer>>;
24
25/// A cached buffer. Cached buffers are never accessed, only retained and released.
26struct Cached(Raw);
27
28// SAFETY: retaining and releasing Metal buffer objects is thread-safe; a cached buffer's
29// memory is not touched until it is handed out again, under the pool's lock.
30unsafe impl Send for Cached {}
31
32/// Memory use and limits of a device's allocator, in bytes.
33#[derive(Debug, Clone, Copy, PartialEq, Eq)]
34pub struct MemoryStats {
35    /// Bytes in live buffers (rounded allocation sizes).
36    pub active: usize,
37    /// Highest `active` since creation or the last [`Device::reset_peak`](super::Device::reset_peak).
38    pub peak: usize,
39    /// Bytes held in the cache for reuse.
40    pub cache: usize,
41    /// Allocations that would push `active` past this fail.
42    pub memory_limit: usize,
43    /// The cache is trimmed to at most this many bytes.
44    pub cache_limit: usize,
45    /// Buffers obtained from Metal (cache misses) so far.
46    pub system_allocations: u64,
47}
48
49/// The size an allocation of `len` bytes is rounded up to.
50pub(crate) fn rounded(len: usize) -> usize {
51    match len {
52        0 => 0,
53        n if n > PAGE => n.div_ceil(PAGE) * PAGE,
54        n => n.div_ceil(SMALL) * SMALL,
55    }
56}
57
58pub(crate) struct Pool {
59    state: Mutex<State>,
60}
61
62struct State {
63    /// Cached buffers up to one page, by size class (`size / 256`), most recent last: O(1)
64    /// reuse for the many small outputs of a computation.
65    small: Vec<Vec<(u64, Cached)>>,
66    /// Cached larger buffers by `(size, stamp)`: best fit is the first key at or above a size.
67    large: BTreeMap<(usize, u64), Cached>,
68    /// Stamps order cached buffers by age, for least-recently-used eviction.
69    next_stamp: u64,
70    cached: usize,
71    active: usize,
72    peak: usize,
73    memory_limit: usize,
74    cache_limit: usize,
75    working_set: usize,
76    system_allocations: u64,
77}
78
79impl State {
80    /// A cached buffer of at least `size` bytes within MLX's slack (smaller than
81    /// `min(2·size, size + 2 pages)`), and its size.
82    fn take(&mut self, size: usize) -> Option<(Raw, usize)> {
83        let limit = (2 * size).min(size + 2 * PAGE);
84        let found = if size <= PAGE {
85            let last = ((limit - 1) / SMALL).min(SMALL_CLASSES - 1);
86            (size / SMALL..=last)
87                .find_map(|c| self.small[c].pop().map(|(_, Cached(raw))| (raw, c * SMALL)))
88        } else {
89            let key = self
90                .large
91                .range((size, 0)..(limit, 0))
92                .next()
93                .map(|(k, _)| *k);
94            key.and_then(|k| self.large.remove(&k).map(|Cached(raw)| (raw, k.0)))
95        };
96        if let Some((_, bucket)) = &found {
97            self.cached -= bucket;
98        }
99        found
100    }
101
102    fn put(&mut self, raw: Raw, size: usize) {
103        let stamp = self.next_stamp;
104        self.next_stamp += 1;
105        if size <= PAGE {
106            self.small[size / SMALL].push((stamp, Cached(raw)));
107        } else {
108            self.large.insert((size, stamp), Cached(raw));
109        }
110        self.cached += size;
111    }
112
113    /// Trims the cache to at most `target` bytes, oldest buffers first. Eviction is rare,
114    /// so it sorts by age on demand instead of maintaining an ordered index on every reuse.
115    fn evict_to(&mut self, target: usize) {
116        if self.cached <= target {
117            return;
118        }
119        if target == 0 {
120            self.small.iter_mut().for_each(Vec::clear);
121            self.large.clear();
122            self.cached = 0;
123            return;
124        }
125        let mut ages: Vec<(u64, usize)> = self
126            .small
127            .iter()
128            .enumerate()
129            .flat_map(|(class, v)| v.iter().map(move |(stamp, _)| (*stamp, class * SMALL)))
130            .chain(self.large.keys().map(|&(size, stamp)| (stamp, size)))
131            .collect();
132        ages.sort_unstable();
133        let mut evict = HashSet::new();
134        for (stamp, size) in ages {
135            if self.cached <= target {
136                break;
137            }
138            evict.insert(stamp);
139            self.cached -= size;
140        }
141        self.small
142            .iter_mut()
143            .for_each(|v| v.retain(|(stamp, _)| !evict.contains(stamp)));
144        self.large.retain(|&(_, stamp), _| !evict.contains(&stamp));
145    }
146
147    /// GC threshold: cached memory is released before active + cached reaches it.
148    fn gc_limit(&self) -> usize {
149        ((self.working_set as f64 * 0.95) as usize).min(self.memory_limit)
150    }
151}
152
153impl Pool {
154    /// A pool with MLX's default limits for a GPU with the given working set and RAM.
155    pub(crate) fn new(working_set: u64, ram: u64) -> Self {
156        let memory_limit = (1.5 * working_set as f64).min(0.95 * ram as f64) as usize;
157        Pool {
158            state: Mutex::new(State {
159                small: (0..SMALL_CLASSES).map(|_| Vec::new()).collect(),
160                large: BTreeMap::new(),
161                next_stamp: 0,
162                cached: 0,
163                active: 0,
164                peak: 0,
165                memory_limit,
166                cache_limit: memory_limit,
167                working_set: working_set as usize,
168                system_allocations: 0,
169            }),
170        }
171    }
172
173    fn lock(&self) -> MutexGuard<'_, State> {
174        self.state.lock().unwrap_or_else(PoisonError::into_inner)
175    }
176
177    /// A buffer of at least `len` bytes and its accounted size; `None` for zero bytes.
178    /// Reused buffers hold stale contents.
179    pub(crate) fn alloc(
180        &self,
181        device: &ProtocolObject<dyn MTLDevice>,
182        len: usize,
183    ) -> Result<Option<(Raw, usize)>> {
184        let size = rounded(len);
185        if size == 0 {
186            return Ok(None);
187        }
188        let mut s = self.lock();
189        if s.active + size > s.memory_limit {
190            s.evict_to(0);
191            return Err(GpuError::OutOfMemory(len));
192        }
193        if let Some((raw, bucket)) = s.take(size) {
194            s.active += bucket;
195            s.peak = s.peak.max(s.active);
196            return Ok(Some((raw, bucket)));
197        }
198        let gc = s.gc_limit();
199        if s.active + s.cached + size >= gc {
200            let target = gc.saturating_sub(s.active + size);
201            s.evict_to(target);
202        }
203        // Reserve before the system call so concurrent allocations respect the limit.
204        s.active += size;
205        s.peak = s.peak.max(s.active);
206        drop(s);
207        let options =
208            MTLResourceOptions::StorageModeShared | MTLResourceOptions::HazardTrackingModeUntracked;
209        match device.newBufferWithLength_options(size, options) {
210            Some(raw) => {
211                self.lock().system_allocations += 1;
212                Ok(Some((raw, size)))
213            }
214            None => {
215                self.lock().active -= size;
216                Err(GpuError::OutOfMemory(len))
217            }
218        }
219    }
220
221    /// Returns a buffer of accounted `size`. It is cached for reuse if `reusable` (no GPU
222    /// work still uses it), otherwise released.
223    pub(crate) fn free(&self, raw: Raw, size: usize, reusable: bool) {
224        let mut s = self.lock();
225        s.active -= size;
226        if reusable {
227            s.put(raw, size);
228            let limit = s.cache_limit;
229            s.evict_to(limit);
230        }
231    }
232
233    pub(crate) fn stats(&self) -> MemoryStats {
234        let s = self.lock();
235        MemoryStats {
236            active: s.active,
237            peak: s.peak,
238            cache: s.cached,
239            memory_limit: s.memory_limit,
240            cache_limit: s.cache_limit,
241            system_allocations: s.system_allocations,
242        }
243    }
244
245    pub(crate) fn set_memory_limit(&self, limit: usize) -> usize {
246        std::mem::replace(&mut self.lock().memory_limit, limit)
247    }
248
249    pub(crate) fn set_cache_limit(&self, limit: usize) -> usize {
250        let mut s = self.lock();
251        let previous = std::mem::replace(&mut s.cache_limit, limit);
252        s.evict_to(limit);
253        previous
254    }
255
256    pub(crate) fn clear_cache(&self) {
257        self.lock().evict_to(0);
258    }
259
260    pub(crate) fn reset_peak(&self) {
261        let mut s = self.lock();
262        s.peak = s.active;
263    }
264}
265
266#[cfg(test)]
267mod tests {
268    use super::*;
269
270    #[test]
271    fn rounding() {
272        assert_eq!(rounded(0), 0);
273        assert_eq!(rounded(1), 256);
274        assert_eq!(rounded(256), 256);
275        assert_eq!(rounded(257), 512);
276        assert_eq!(rounded(PAGE), PAGE);
277        assert_eq!(rounded(PAGE + 1), 2 * PAGE);
278    }
279
280    #[test]
281    fn default_limits_follow_mlx() {
282        let pool = Pool::new(40_000_000_000, 48_000_000_000);
283        let stats = pool.stats();
284        assert_eq!(stats.memory_limit, 45_600_000_000); // min(1.5 * 40e9, 0.95 * 48e9)
285        assert_eq!(stats.cache_limit, stats.memory_limit);
286    }
287}
nuvec_gpu/metal/allocator.rs

nuvec_gpu/metal/
allocator.rs
