nuvec_gpu/metal/
buffer.rs

1//! Shared-memory buffers.
2
3use std::ptr::NonNull;
4use std::sync::Arc;
5use std::sync::atomic::{AtomicUsize, Ordering};
6
7use objc2::runtime::ProtocolObject;
8use objc2_metal::MTLBuffer;
9
10use super::Device;
11use super::allocator::{Pool, Raw};
12use crate::{GpuError, Result};
13
14/// GPU memory in unified (shared) storage: the CPU reads and writes it in place.
15///
16/// Buffers come from the device's caching allocator and return to it when dropped. Hazard
17/// tracking is disabled at the Metal level (as in MLX); nuvec orders accesses itself. CPU
18/// access is refused while submitted GPU work uses the buffer, and GPU writes require
19/// `&mut Buffer`, so the CPU and the GPU never race on it.
20pub struct Buffer {
21    /// `None` for zero-length buffers, which own no Metal memory.
22    raw: Option<Raw>,
23    /// `raw.contents()`, fixed for the buffer's lifetime (cached: each lookup is an
24    /// Objective-C message).
25    ptr: NonNull<u8>,
26    len: usize,
27    /// Accounted (rounded) size in the pool.
28    size: usize,
29    pool: Arc<Pool>,
30    /// Number of submitted-but-unwaited command buffers that use this buffer.
31    in_flight: Arc<AtomicUsize>,
32}
33
34// SAFETY: Metal buffer objects may be retained, released and bound from any thread. The only
35// non-thread-safe part is the memory behind `contents()`, and access to it is synchronised by
36// Rust borrows (`contents_mut` needs `&mut self`; GPU writes need `&mut Buffer`) plus the
37// `in_flight` guard, which refuses CPU access while GPU work is pending.
38unsafe impl Send for Buffer {}
39// SAFETY: as above; `&Buffer` only allows shared reads of memory no GPU work is writing.
40unsafe impl Sync for Buffer {}
41
42impl Device {
43    /// A buffer of `len` bytes from the caching allocator. Its contents are unspecified
44    /// (a reused buffer holds stale bytes); fails with [`GpuError::OutOfMemory`] beyond the
45    /// memory limit.
46    pub fn new_buffer(&self, len: usize) -> Result<Buffer> {
47        let pool = self.pool().clone();
48        let (raw, size) = match pool.alloc(self.raw(), len)? {
49            Some((raw, size)) => (Some(raw), size),
50            None => (None, 0),
51        };
52        let ptr = raw
53            .as_ref()
54            .map_or(NonNull::dangling(), |r| r.contents().cast());
55        Ok(Buffer {
56            raw,
57            ptr,
58            len,
59            size,
60            pool,
61            in_flight: Arc::new(AtomicUsize::new(0)),
62        })
63    }
64
65    /// A buffer holding a copy of `bytes`.
66    pub fn buffer_from(&self, bytes: &[u8]) -> Result<Buffer> {
67        let mut buffer = self.new_buffer(bytes.len())?;
68        buffer.contents_mut()?.copy_from_slice(bytes);
69        Ok(buffer)
70    }
71}
72
73impl Buffer {
74    /// Size in bytes.
75    pub fn len(&self) -> usize {
76        self.len
77    }
78
79    /// Whether the buffer holds no bytes.
80    pub fn is_empty(&self) -> bool {
81        self.len == 0
82    }
83
84    /// Whether submitted GPU work still uses the buffer.
85    pub fn in_use(&self) -> bool {
86        self.in_flight.load(Ordering::Acquire) != 0
87    }
88
89    /// The bytes, if no submitted GPU work uses the buffer.
90    pub fn contents(&self) -> Result<&[u8]> {
91        if self.in_use() {
92            return Err(GpuError::BufferInUse);
93        }
94        if self.raw.is_none() {
95            return Ok(&[]);
96        }
97        // SAFETY: `ptr` is `raw.contents()`, at least `len` bytes (the allocation is at least the
98        // rounded size) of CPU-visible shared storage that lives as long as `raw`, which this
99        // borrow keeps alive. No GPU work uses the buffer (checked above), and new GPU writes
100        // need `&mut self`, which this shared borrow excludes.
101        Ok(unsafe { std::slice::from_raw_parts(self.ptr.as_ptr(), self.len) })
102    }
103
104    /// The bytes, mutably, if no submitted GPU work uses the buffer.
105    pub fn contents_mut(&mut self) -> Result<&mut [u8]> {
106        if self.in_use() {
107            return Err(GpuError::BufferInUse);
108        }
109        if self.raw.is_none() {
110            return Ok(&mut []);
111        }
112        // SAFETY: as in `contents`; additionally `&mut self` is exclusive, and `Buffer` is not
113        // `Clone`, so no other CPU reference to these bytes exists.
114        Ok(unsafe { std::slice::from_raw_parts_mut(self.ptr.as_ptr(), self.len) })
115    }
116
117    pub(crate) fn raw(&self) -> Option<&ProtocolObject<dyn MTLBuffer>> {
118        self.raw.as_deref()
119    }
120
121    pub(crate) fn retained(&self) -> Option<Raw> {
122        self.raw.clone()
123    }
124
125    pub(crate) fn guard(&self) -> Arc<AtomicUsize> {
126        self.in_flight.clone()
127    }
128
129    /// Identity for hazard tracking (unique per live buffer, including empty ones).
130    pub(crate) fn id(&self) -> usize {
131        Arc::as_ptr(&self.in_flight) as usize
132    }
133}
134
135impl Drop for Buffer {
136    /// Returns the memory to the allocator's cache, unless GPU work still uses it (then it
137    /// is released when the GPU finishes, never reused early).
138    fn drop(&mut self) {
139        if let Some(raw) = self.raw.take() {
140            let reusable = !self.in_use();
141            self.pool.free(raw, self.size, reusable);
142        }
143    }
144}
145
146impl std::fmt::Debug for Buffer {
147    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
148        f.debug_struct("Buffer")
149            .field("len", &self.len)
150            .field("in_use", &self.in_use())
151            .finish()
152    }
153}
nuvec_gpu/metal/buffer.rs

nuvec_gpu/metal/
buffer.rs
