nuvec_gpu/metal/
mod.rs

1//! The Metal runtime on `objc2-metal` (macOS only).
2
3mod allocator;
4mod buffer;
5mod command;
6
7use std::collections::{HashMap, HashSet};
8use std::sync::{Arc, Mutex, OnceLock, PoisonError};
9
10pub use allocator::MemoryStats;
11use allocator::Pool;
12pub use buffer::Buffer;
13pub use command::{Arg, CommandBuffer, Submitted};
14use nuvec_kernels::MetallibMode;
15use objc2::rc::Retained;
16use objc2::runtime::ProtocolObject;
17use objc2_foundation::NSString;
18use objc2_metal::{
19    MTLCommandQueue, MTLComputePipelineState, MTLCreateSystemDefaultDevice, MTLDevice, MTLLibrary,
20};
21
22use crate::info::parse_architecture;
23use crate::{DeviceInfo, GpuError, Result};
24
25type Library = Retained<ProtocolObject<dyn MTLLibrary>>;
26
27/// A Metal GPU with its command queue, kernel library and pipeline cache. Cheap to clone.
28#[derive(Clone)]
29pub struct Device {
30    inner: Arc<Inner>,
31}
32
33struct Inner {
34    device: Retained<ProtocolObject<dyn MTLDevice>>,
35    queue: Retained<ProtocolObject<dyn MTLCommandQueue>>,
36    info: DeviceInfo,
37    library: OnceLock<Result<Library>>,
38    pipelines: Mutex<HashMap<String, Pipeline>>,
39    pool: Arc<Pool>,
40}
41
42impl Device {
43    /// The system's default Metal device, or `None` if there is none (for example in some
44    /// virtual machines).
45    pub fn system_default() -> Option<Device> {
46        let device = MTLCreateSystemDefaultDevice()?;
47        let queue = device.newCommandQueue()?;
48        let architecture = device.architecture().name().to_string();
49        let (generation, class) = parse_architecture(&architecture);
50        let info = DeviceInfo {
51            name: device.name().to_string(),
52            architecture,
53            generation,
54            class,
55            recommended_working_set: device.recommendedMaxWorkingSetSize(),
56            max_threads_per_threadgroup: device.maxThreadsPerThreadgroup().width,
57        };
58        let ram = objc2_foundation::NSProcessInfo::processInfo().physicalMemory();
59        let pool = Arc::new(Pool::new(info.recommended_working_set, ram));
60        Some(Device {
61            inner: Arc::new(Inner {
62                device,
63                queue,
64                info,
65                library: OnceLock::new(),
66                pipelines: Mutex::new(HashMap::new()),
67                pool,
68            }),
69        })
70    }
71
72    /// The process-wide default device, created on first use: one allocator, kernel cache
73    /// and queue shared by everything that uses the GPU.
74    pub fn shared() -> Option<&'static Device> {
75        static SHARED: OnceLock<Option<Device>> = OnceLock::new();
76        SHARED.get_or_init(Device::system_default).as_ref()
77    }
78
79    /// Memory use and limits of this device's allocator.
80    pub fn memory(&self) -> MemoryStats {
81        self.inner.pool.stats()
82    }
83
84    /// Sets the memory limit (allocations beyond it fail) and returns the previous one.
85    pub fn set_memory_limit(&self, bytes: usize) -> usize {
86        self.inner.pool.set_memory_limit(bytes)
87    }
88
89    /// Sets the cache limit, trimming the cache to it, and returns the previous one.
90    pub fn set_cache_limit(&self, bytes: usize) -> usize {
91        self.inner.pool.set_cache_limit(bytes)
92    }
93
94    /// Releases every cached buffer.
95    pub fn clear_cache(&self) {
96        self.inner.pool.clear_cache();
97    }
98
99    /// Resets the peak to the current active memory.
100    pub fn reset_peak(&self) {
101        self.inner.pool.reset_peak();
102    }
103
104    pub(crate) fn pool(&self) -> &Arc<Pool> {
105        &self.inner.pool
106    }
107
108    /// What this GPU is and can do.
109    pub fn info(&self) -> &DeviceInfo {
110        &self.inner.info
111    }
112
113    /// The compute pipeline for the kernel `name` from nuvec's kernel library, compiled on
114    /// first use and cached.
115    pub fn pipeline(&self, name: &str) -> Result<Pipeline> {
116        let mut cache = self
117            .inner
118            .pipelines
119            .lock()
120            .unwrap_or_else(PoisonError::into_inner);
121        if let Some(p) = cache.get(name) {
122            return Ok(p.clone());
123        }
124        let function = self
125            .library()?
126            .newFunctionWithName(&NSString::from_str(name))
127            .ok_or_else(|| GpuError::MissingKernel(name.to_owned()))?;
128        let state = self
129            .inner
130            .device
131            .newComputePipelineStateWithFunction_error(&function)
132            .map_err(|e| GpuError::Pipeline {
133                name: name.to_owned(),
134                message: e.localizedDescription().to_string(),
135            })?;
136        let pipeline = Pipeline {
137            state,
138            name: name.into(),
139        };
140        cache.insert(name.to_owned(), pipeline.clone());
141        Ok(pipeline)
142    }
143
144    /// nuvec's kernel library, loaded once.
145    fn library(&self) -> Result<&Library> {
146        self.inner
147            .library
148            .get_or_init(|| load_library(&self.inner.device))
149            .as_ref()
150            .map_err(Clone::clone)
151    }
152
153    pub(crate) fn raw(&self) -> &ProtocolObject<dyn MTLDevice> {
154        &self.inner.device
155    }
156
157    pub(crate) fn queue(&self) -> &ProtocolObject<dyn MTLCommandQueue> {
158        &self.inner.queue
159    }
160}
161
162impl std::fmt::Debug for Device {
163    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
164        f.debug_struct("Device")
165            .field("info", &self.inner.info)
166            .finish_non_exhaustive()
167    }
168}
169
170/// Loads the precompiled library, or compiles the embedded sources (with their headers
171/// inlined, as the runtime compiler has no include path for them).
172fn load_library(device: &ProtocolObject<dyn MTLDevice>) -> Result<Library> {
173    let error = |e: Retained<objc2_foundation::NSError>| {
174        GpuError::LibraryLoad(e.localizedDescription().to_string())
175    };
176    match nuvec_kernels::MODE {
177        MetallibMode::Precompiled => {
178            let bytes = nuvec_kernels::metallib().ok_or(GpuError::NoKernels)?;
179            device
180                .newLibraryWithData_error(&dispatch2::DispatchData::from_static_bytes(bytes))
181                .map_err(error)
182        }
183        MetallibMode::RuntimeSource => {
184            let source = inline_sources(nuvec_kernels::sources());
185            device
186                .newLibraryWithSource_options_error(&NSString::from_str(&source), None)
187                .map_err(error)
188        }
189        _ => Err(GpuError::NoKernels),
190    }
191}
192
193/// All `.metal` sources as one translation unit, each `#include "header.h"` of an embedded
194/// header replaced by the header's text on first use.
195fn inline_sources(sources: &[(&str, &str)]) -> String {
196    let headers: HashMap<&str, &str> = sources
197        .iter()
198        .copied()
199        .filter(|(p, _)| p.ends_with(".h"))
200        .collect();
201    let mut included = HashSet::new();
202    let mut out = String::new();
203    for (_, source) in sources.iter().filter(|(p, _)| p.ends_with(".metal")) {
204        for line in source.lines() {
205            let header = line
206                .trim()
207                .strip_prefix("#include \"")
208                .and_then(|r| r.strip_suffix('"'));
209            match header.and_then(|h| headers.get(h).map(|text| (h, text))) {
210                Some((h, text)) => {
211                    if included.insert(h) {
212                        out.push_str(text);
213                        out.push('\n');
214                    }
215                }
216                None => {
217                    out.push_str(line);
218                    out.push('\n');
219                }
220            }
221        }
222    }
223    out
224}
225
226/// A compiled compute kernel. Cheap to clone.
227#[derive(Clone)]
228pub struct Pipeline {
229    state: Retained<ProtocolObject<dyn MTLComputePipelineState>>,
230    name: Arc<str>,
231}
232
233impl Pipeline {
234    /// The kernel's name.
235    pub fn name(&self) -> &str {
236        &self.name
237    }
238
239    /// The most threads one threadgroup of this kernel may have.
240    pub fn max_threads_per_threadgroup(&self) -> usize {
241        self.state.maxTotalThreadsPerThreadgroup()
242    }
243
244    /// The SIMD-group width (threads that execute in lockstep).
245    pub fn execution_width(&self) -> usize {
246        self.state.threadExecutionWidth()
247    }
248
249    /// Whether two handles refer to the same compiled pipeline.
250    pub fn same_as(&self, other: &Pipeline) -> bool {
251        Retained::as_ptr(&self.state) == Retained::as_ptr(&other.state)
252    }
253
254    pub(crate) fn raw(&self) -> &ProtocolObject<dyn MTLComputePipelineState> {
255        &self.state
256    }
257}
258
259impl std::fmt::Debug for Pipeline {
260    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
261        f.debug_struct("Pipeline")
262            .field("name", &self.name)
263            .finish_non_exhaustive()
264    }
265}
266
267#[cfg(test)]
268mod tests {
269    use super::inline_sources;
270
271    #[test]
272    fn inlines_each_header_once() {
273        let sources = [
274            ("a.metal", "#include \"h.h\"\nkernel A;"),
275            ("b.metal", "  #include \"h.h\"\nkernel B;"),
276            ("h.h", "#pragma once\ninline int f();"),
277        ];
278        let out = inline_sources(&sources);
279        assert_eq!(out.matches("inline int f();").count(), 1);
280        assert!(out.contains("kernel A;") && out.contains("kernel B;"));
281        assert!(!out.contains("#include \"h.h\""));
282    }
283}
nuvec_gpu/metal/mod.rs

nuvec_gpu/metal/
mod.rs
