nuvec_gpu/metal/
mod.rs

1//! The Metal runtime on `objc2-metal` (macOS only).
2
3mod allocator;
4mod buffer;
5mod command;
6pub mod profile;
7mod residency;
8
9use std::collections::{HashMap, HashSet};
10use std::ptr::NonNull;
11use std::sync::{Arc, Condvar, Mutex, OnceLock, PoisonError};
12
13use allocator::Pool;
14pub use buffer::Buffer;
15pub use command::{Arg, CommandBuffer, Pending, Submitted, Timeline};
16use nuvec_kernels::MetallibMode;
17use objc2::rc::Retained;
18use objc2::runtime::ProtocolObject;
19use objc2_foundation::NSString;
20use objc2_metal::{
21    MTLCommandQueue, MTLComputePipelineState, MTLCreateSystemDefaultDevice, MTLDataType, MTLDevice,
22    MTLFunctionConstantValues, MTLLibrary,
23};
24
25use crate::info::parse_architecture;
26use crate::{DeviceInfo, GpuError, MemoryStats, Result};
27
28type Library = Retained<ProtocolObject<dyn MTLLibrary>>;
29
30/// A function constant value for [`Device::pipeline_with_constants`].
31#[derive(Debug, Clone, Copy, PartialEq, Eq)]
32pub enum Constant {
33    /// A `bool` constant.
34    Bool(bool),
35    /// An `int` constant.
36    Int(i32),
37}
38
39/// A Metal GPU with its command queue, kernel library and pipeline cache. Cheap to clone.
40#[derive(Clone)]
41pub struct Device {
42    inner: Arc<Inner>,
43}
44
45struct Inner {
46    device: Retained<ProtocolObject<dyn MTLDevice>>,
47    queue: Retained<ProtocolObject<dyn MTLCommandQueue>>,
48    info: DeviceInfo,
49    library: OnceLock<Result<Library>>,
50    pipelines: Mutex<HashMap<String, Pipeline>>,
51    pool: Arc<Pool>,
52    in_flight: InFlight,
53}
54
55/// The number of committed command buffers that have not completed yet.
56#[derive(Default)]
57pub(crate) struct InFlight {
58    count: Mutex<usize>,
59    done: Condvar,
60}
61
62impl InFlight {
63    pub(crate) fn begin(&self) {
64        *self.count.lock().unwrap_or_else(PoisonError::into_inner) += 1;
65    }
66
67    pub(crate) fn end(&self) {
68        let mut count = self.count.lock().unwrap_or_else(PoisonError::into_inner);
69        *count = count.saturating_sub(1);
70        self.done.notify_all();
71    }
72
73    fn get(&self) -> usize {
74        *self.count.lock().unwrap_or_else(PoisonError::into_inner)
75    }
76
77    fn wait_below(&self, max: usize) {
78        let mut count = self.count.lock().unwrap_or_else(PoisonError::into_inner);
79        while *count >= max.max(1) {
80            count = self
81                .done
82                .wait(count)
83                .unwrap_or_else(PoisonError::into_inner);
84        }
85    }
86}
87
88impl Device {
89    /// The system's default Metal device, or `None` if there is none (for example in some
90    /// virtual machines).
91    pub fn system_default() -> Option<Device> {
92        let device = MTLCreateSystemDefaultDevice()?;
93        let queue = device.newCommandQueue()?;
94        let architecture = device.architecture().name().to_string();
95        let (generation, class) = parse_architecture(&architecture);
96        let info = DeviceInfo {
97            name: device.name().to_string(),
98            architecture,
99            generation,
100            class,
101            recommended_working_set: device.recommendedMaxWorkingSetSize(),
102            max_threads_per_threadgroup: device.maxThreadsPerThreadgroup().width,
103        };
104        let ram = objc2_foundation::NSProcessInfo::processInfo().physicalMemory();
105        let options = objc2_metal::MTLResourceOptions::StorageModeShared;
106        let placeholder = device.newBufferWithLength_options(16, options);
107        let residency = residency::Residency::new(residency::attach(&device, &queue));
108        let pool = Arc::new(Pool::new(
109            info.recommended_working_set,
110            ram,
111            placeholder,
112            residency,
113        ));
114        Some(Device {
115            inner: Arc::new(Inner {
116                device,
117                queue,
118                info,
119                library: OnceLock::new(),
120                pipelines: Mutex::new(HashMap::new()),
121                pool,
122                in_flight: InFlight::default(),
123            }),
124        })
125    }
126
127    /// The process-wide default device, created on first use: one allocator, kernel cache
128    /// and queue shared by everything that uses the GPU.
129    pub fn shared() -> Option<&'static Device> {
130        static SHARED: OnceLock<Option<Device>> = OnceLock::new();
131        SHARED.get_or_init(Device::system_default).as_ref()
132    }
133
134    /// Memory use and limits of this device's allocator.
135    pub fn memory(&self) -> MemoryStats {
136        self.inner.pool.stats()
137    }
138
139    /// Sets the memory limit (allocations beyond it fail) and returns the previous one.
140    pub fn set_memory_limit(&self, bytes: usize) -> usize {
141        self.inner.pool.set_memory_limit(bytes)
142    }
143
144    /// Sets the cache limit, trimming the cache to it, and returns the previous one.
145    pub fn set_cache_limit(&self, bytes: usize) -> usize {
146        self.inner.pool.set_cache_limit(bytes)
147    }
148
149    /// Sets the wired limit and returns the previous one (MLX's `set_wired_limit`; 0 by
150    /// default). Buffers, mapped files included, are kept resident up to this many bytes
151    /// while GPU work runs, instead of being made resident per command buffer, which can
152    /// page memory in again under pressure. Lowering it unwires buffers until the wired set
153    /// fits. A limit above the recommended working set is an error. Without residency sets
154    /// (before macOS 15, or on GPUs before Metal 3) the limit is recorded but nothing is
155    /// wired.
156    pub fn set_wired_limit(&self, bytes: usize) -> Result<usize> {
157        if bytes as u64 > self.inner.info.recommended_working_set {
158            return Err(GpuError::InvalidArgument(format!(
159                "wired limit {bytes} exceeds the recommended working set ({} bytes)",
160                self.inner.info.recommended_working_set
161            )));
162        }
163        Ok(self.inner.pool.set_wired_limit(bytes))
164    }
165
166    /// Whether buffers can be wired on this system (residency sets: macOS 15 and a Metal 3
167    /// GPU).
168    pub fn can_wire(&self) -> bool {
169        self.inner.pool.can_wire()
170    }
171
172    /// Releases every cached buffer.
173    pub fn clear_cache(&self) {
174        self.inner.pool.clear_cache();
175    }
176
177    /// Command buffers committed on this device that have not completed yet.
178    pub fn in_flight(&self) -> usize {
179        self.inner.in_flight.get()
180    }
181
182    /// Blocks until fewer than `max` command buffers are in flight (back-pressure for
183    /// callers that submit without waiting).
184    pub fn wait_for_capacity(&self, max: usize) {
185        self.inner.in_flight.wait_below(max);
186    }
187
188    pub(crate) fn flights(&self) -> &InFlight {
189        &self.inner.in_flight
190    }
191
192    /// Resets the peak to the current active memory.
193    pub fn reset_peak(&self) {
194        self.inner.pool.reset_peak();
195    }
196
197    pub(crate) fn pool(&self) -> &Arc<Pool> {
198        &self.inner.pool
199    }
200
201    /// What this GPU is and can do.
202    pub fn info(&self) -> &DeviceInfo {
203        &self.inner.info
204    }
205
206    /// The compute pipeline for the kernel `name` from nuvec's kernel library, compiled on
207    /// first use and cached.
208    pub fn pipeline(&self, name: &str) -> Result<Pipeline> {
209        self.pipeline_with(name, &[])
210    }
211
212    /// The pipeline for kernel `name` specialised with boolean function constants
213    /// `(index, value)`, compiled on first use and cached per name and constants.
214    pub fn pipeline_with(&self, name: &str, constants: &[(usize, bool)]) -> Result<Pipeline> {
215        let constants: Vec<(usize, Constant)> = constants
216            .iter()
217            .map(|&(i, v)| (i, Constant::Bool(v)))
218            .collect();
219        self.pipeline_with_constants(name, &constants)
220    }
221
222    /// The pipeline for kernel `name` specialised with boolean or integer function constants
223    /// `(index, value)`, compiled on first use and cached per name and constants.
224    pub fn pipeline_with_constants(
225        &self,
226        name: &str,
227        constants: &[(usize, Constant)],
228    ) -> Result<Pipeline> {
229        let key = constants.iter().fold(name.to_owned(), |mut k, (i, v)| {
230            match v {
231                Constant::Bool(b) => k.push_str(&format!("#{i}={}", u8::from(*b))),
232                Constant::Int(n) => k.push_str(&format!("#{i}=i{n}")),
233            }
234            k
235        });
236        let mut cache = self
237            .inner
238            .pipelines
239            .lock()
240            .unwrap_or_else(PoisonError::into_inner);
241        if let Some(p) = cache.get(&key) {
242            return Ok(p.clone());
243        }
244        let library = self.library()?;
245        let ns_name = NSString::from_str(name);
246        let function = if constants.is_empty() {
247            library
248                .newFunctionWithName(&ns_name)
249                .ok_or_else(|| GpuError::MissingKernel(name.to_owned()))?
250        } else {
251            let values = MTLFunctionConstantValues::new();
252            for &(index, value) in constants {
253                match value {
254                    // SAFETY: `b` is a live local, which Metal copies during the call; `bool`
255                    // matches `MTLDataType::Bool` (one byte).
256                    Constant::Bool(b) => unsafe {
257                        values.setConstantValue_type_atIndex(
258                            NonNull::from(&b).cast(),
259                            MTLDataType::Bool,
260                            index,
261                        );
262                    },
263                    // SAFETY: as above; `i32` matches `MTLDataType::Int`.
264                    Constant::Int(n) => unsafe {
265                        values.setConstantValue_type_atIndex(
266                            NonNull::from(&n).cast(),
267                            MTLDataType::Int,
268                            index,
269                        );
270                    },
271                }
272            }
273            library
274                .newFunctionWithName_constantValues_error(&ns_name, &values)
275                .map_err(|e| GpuError::Pipeline {
276                    name: key.clone(),
277                    message: e.localizedDescription().to_string(),
278                })?
279        };
280        let state = self
281            .inner
282            .device
283            .newComputePipelineStateWithFunction_error(&function)
284            .map_err(|e| GpuError::Pipeline {
285                name: key.clone(),
286                message: e.localizedDescription().to_string(),
287            })?;
288        let pipeline = Pipeline {
289            state,
290            name: key.as_str().into(),
291        };
292        cache.insert(key, pipeline.clone());
293        Ok(pipeline)
294    }
295
296    /// The pipeline for kernel `name` from Metal source that `source` generates, compiled
297    /// at runtime on first use and cached by name (so a name must always mean the same
298    /// source). The source may include nuvec's headers (`#include "common.h"`); it is
299    /// compiled with IEEE-conformant math, as nuvec's own kernels are.
300    pub fn pipeline_from_source(
301        &self,
302        name: &str,
303        source: impl FnOnce() -> String,
304    ) -> Result<Pipeline> {
305        let mut cache = self
306            .inner
307            .pipelines
308            .lock()
309            .unwrap_or_else(PoisonError::into_inner);
310        if let Some(p) = cache.get(name) {
311            return Ok(p.clone());
312        }
313        let failed = |message: String| GpuError::Pipeline {
314            name: name.to_owned(),
315            message,
316        };
317        let source = inline_sources(&[("generated.metal", &source())]);
318        let options = objc2_metal::MTLCompileOptions::new();
319        options.setMathMode(objc2_metal::MTLMathMode::Safe);
320        let library = self
321            .inner
322            .device
323            .newLibraryWithSource_options_error(&NSString::from_str(&source), Some(&options))
324            .map_err(|e| failed(e.localizedDescription().to_string()))?;
325        let function = library
326            .newFunctionWithName(&NSString::from_str(name))
327            .ok_or_else(|| GpuError::MissingKernel(name.to_owned()))?;
328        let state = self
329            .inner
330            .device
331            .newComputePipelineStateWithFunction_error(&function)
332            .map_err(|e| failed(e.localizedDescription().to_string()))?;
333        let pipeline = Pipeline {
334            state,
335            name: name.into(),
336        };
337        cache.insert(name.to_owned(), pipeline.clone());
338        Ok(pipeline)
339    }
340
341    /// nuvec's kernel library, loaded once.
342    fn library(&self) -> Result<&Library> {
343        self.inner
344            .library
345            .get_or_init(|| load_library(&self.inner.device))
346            .as_ref()
347            .map_err(Clone::clone)
348    }
349
350    pub(crate) fn raw(&self) -> &ProtocolObject<dyn MTLDevice> {
351        &self.inner.device
352    }
353
354    pub(crate) fn queue(&self) -> &ProtocolObject<dyn MTLCommandQueue> {
355        &self.inner.queue
356    }
357}
358
359impl std::fmt::Debug for Device {
360    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
361        f.debug_struct("Device")
362            .field("info", &self.inner.info)
363            .finish_non_exhaustive()
364    }
365}
366
367/// Loads the precompiled library, or compiles the embedded sources (with their headers
368/// inlined, as the runtime compiler has no include path for them).
369fn load_library(device: &ProtocolObject<dyn MTLDevice>) -> Result<Library> {
370    let error = |e: Retained<objc2_foundation::NSError>| {
371        GpuError::LibraryLoad(e.localizedDescription().to_string())
372    };
373    match nuvec_kernels::MODE {
374        MetallibMode::Precompiled => {
375            let bytes = nuvec_kernels::metallib().ok_or(GpuError::NoKernels)?;
376            device
377                .newLibraryWithData_error(&dispatch2::DispatchData::from_static_bytes(bytes))
378                .map_err(error)
379        }
380        MetallibMode::RuntimeSource => {
381            let source = inline_sources(nuvec_kernels::sources());
382            // IEEE-conformant math, matching the precompiled library's `-fno-fast-math`.
383            let options = objc2_metal::MTLCompileOptions::new();
384            options.setMathMode(objc2_metal::MTLMathMode::Safe);
385            device
386                .newLibraryWithSource_options_error(&NSString::from_str(&source), Some(&options))
387                .map_err(error)
388        }
389        _ => Err(GpuError::NoKernels),
390    }
391}
392
393/// All `.metal` sources as one translation unit, each `#include "header.h"` of one of
394/// nuvec's headers replaced by the header's text on first use.
395fn inline_sources(sources: &[(&str, &str)]) -> String {
396    let headers: HashMap<&str, &str> = nuvec_kernels::sources()
397        .iter()
398        .chain(sources)
399        .copied()
400        .chain([("common.h", nuvec_kernels::COMMON_H)])
401        .filter(|(p, _)| p.ends_with(".h"))
402        .collect();
403    let mut included = HashSet::new();
404    let mut out = String::new();
405    for (_, source) in sources.iter().filter(|(p, _)| p.ends_with(".metal")) {
406        for line in source.lines() {
407            let header = line
408                .trim()
409                .strip_prefix("#include \"")
410                .and_then(|r| r.strip_suffix('"'));
411            match header.and_then(|h| headers.get(h).map(|text| (h, text))) {
412                Some((h, text)) => {
413                    if included.insert(h) {
414                        out.push_str(text);
415                        out.push('\n');
416                    }
417                }
418                None => {
419                    out.push_str(line);
420                    out.push('\n');
421                }
422            }
423        }
424    }
425    out
426}
427
428/// A compiled compute kernel. Cheap to clone.
429#[derive(Clone)]
430pub struct Pipeline {
431    state: Retained<ProtocolObject<dyn MTLComputePipelineState>>,
432    name: Arc<str>,
433}
434
435impl Pipeline {
436    /// The kernel's name.
437    pub fn name(&self) -> &str {
438        &self.name
439    }
440
441    /// The most threads one threadgroup of this kernel may have.
442    pub fn max_threads_per_threadgroup(&self) -> usize {
443        self.state.maxTotalThreadsPerThreadgroup()
444    }
445
446    /// The SIMD-group width (threads that execute in lockstep).
447    pub fn execution_width(&self) -> usize {
448        self.state.threadExecutionWidth()
449    }
450
451    /// Whether two handles refer to the same compiled pipeline.
452    pub fn same_as(&self, other: &Pipeline) -> bool {
453        Retained::as_ptr(&self.state) == Retained::as_ptr(&other.state)
454    }
455
456    pub(crate) fn raw(&self) -> &ProtocolObject<dyn MTLComputePipelineState> {
457        &self.state
458    }
459}
460
461impl std::fmt::Debug for Pipeline {
462    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
463        f.debug_struct("Pipeline")
464            .field("name", &self.name)
465            .finish_non_exhaustive()
466    }
467}
468
469#[cfg(test)]
470mod tests {
471    use super::inline_sources;
472
473    #[test]
474    fn inlines_each_header_once() {
475        let sources = [
476            ("a.metal", "#include \"h.h\"\nkernel A;"),
477            ("b.metal", "  #include \"h.h\"\nkernel B;"),
478            ("h.h", "#pragma once\ninline int f();"),
479        ];
480        let out = inline_sources(&sources);
481        assert_eq!(out.matches("inline int f();").count(), 1);
482        assert!(out.contains("kernel A;") && out.contains("kernel B;"));
483        assert!(!out.contains("#include \"h.h\""));
484    }
485}
nuvec_gpu/metal/mod.rs

nuvec_gpu/metal/
mod.rs
