cherenkov/
metal.rs

1use anyhow::{Context, Result};
2use objc2::rc::Retained;
3use objc2::runtime::ProtocolObject;
4use objc2_foundation::NSString;
5use objc2_metal::{
6    MTLBuffer, MTLCommandBuffer, MTLCommandEncoder, MTLCommandQueue, MTLComputeCommandEncoder,
7    MTLComputePipelineState, MTLCreateSystemDefaultDevice, MTLDevice, MTLLibrary,
8    MTLResourceOptions, MTLSize,
9};
10use std::ffi::c_void;
11use std::ptr::NonNull;
12use std::time::Instant;
13
14pub const PAGE_SIZE: usize = 16384;
15
16pub struct MetalContext {
17    pub allocation_limit: std::cell::Cell<Option<usize>>,
18    pub device: Retained<ProtocolObject<dyn MTLDevice>>,
19    pub queue: Retained<ProtocolObject<dyn MTLCommandQueue>>,
20}
21
22impl MetalContext {
23    pub fn new() -> Result<Self> {
24        let device = MTLCreateSystemDefaultDevice().context("no Metal device")?;
25        let queue = device.newCommandQueue().context("newCommandQueue failed")?;
26
27        Ok(MetalContext {
28            device,
29            queue,
30            allocation_limit: std::cell::Cell::new(None),
31        })
32    }
33
34    pub fn compile_library(
35        &self,
36        source: &str,
37    ) -> Result<Retained<ProtocolObject<dyn MTLLibrary>>> {
38        let source = NSString::from_str(source);
39
40        self.device
41            .newLibraryWithSource_options_error(&source, None)
42            .map_err(|e| anyhow::anyhow!("MSL compile failed: {e}"))
43    }
44
45    pub fn pipeline(
46        &self,
47        library: &ProtocolObject<dyn MTLLibrary>,
48        function: &str,
49    ) -> Result<Retained<ProtocolObject<dyn MTLComputePipelineState>>> {
50        let name = NSString::from_str(function);
51        let function_obj = library
52            .newFunctionWithName(&name)
53            .with_context(|| format!("kernel function {function:?} not found"))?;
54        let pso = self
55            .device
56            .newComputePipelineStateWithFunction_error(&function_obj)
57            .map_err(|e| anyhow::anyhow!("pipeline creation failed for {function}: {e}"))?;
58        // Register-cliff canary: the compiler lowers a pipeline's
59        // maxTotalThreadsPerThreadgroup below the 1024 default exactly when
60        // it hits its register budget (Apple family 9). Warn so cliffs are
61        // caught at load, not in tokens/s. Keep the warning for unexpectedly narrow pipelines.
62        let max_threads = pso.maxTotalThreadsPerThreadgroup();
63
64        if max_threads < 256 {
65            eprintln!(
66                "WARNING: pipeline {function} register-limited: \
67                 max_threads/tg = {max_threads} (< 256)"
68            );
69        }
70
71        Ok(pso)
72    }
73
74    /// Expose an existing mapping as a Metal buffer without copying it.
75    ///
76    /// # Safety
77    /// The page-rounded region must remain mapped until this buffer and
78    /// all GPU work using it are finished. The caller must synchronize
79    /// CPU writes with GPU access; the buffer does not own the mapping.
80    pub unsafe fn wrap_mmap(
81        &self,
82        bytes: &[u8],
83    ) -> Result<Retained<ProtocolObject<dyn MTLBuffer>>> {
84        let ptr = bytes.as_ptr() as *mut c_void;
85
86        anyhow::ensure!(
87            (ptr as usize).is_multiple_of(PAGE_SIZE),
88            "mmap base is not page-aligned"
89        );
90
91        let len = bytes.len().div_ceil(PAGE_SIZE) * PAGE_SIZE;
92
93        self.check_allocation(len)?;
94
95        let ptr = NonNull::new(ptr).context("null mmap pointer")?;
96
97        unsafe {
98            self.device
99                .newBufferWithBytesNoCopy_length_options_deallocator(
100                    ptr,
101                    len,
102                    MTLResourceOptions::empty(),
103                    None,
104                )
105        }
106        .context("newBufferWithBytesNoCopy failed (is the region page-aligned?)")
107    }
108
109    /// Wrap a page-aligned portion of a mapping owned by the expert pool.
110    ///
111    /// # Safety
112    /// `ptr..ptr+len` must remain mapped for the buffer's lifetime and all
113    /// submitted GPU work. CPU/GPU access must obey the pool's event
114    /// handshake, including outstanding reads before a slot is reused.
115    pub unsafe fn wrap_region(
116        &self,
117        ptr: *mut u8,
118        len: usize,
119    ) -> Result<Retained<ProtocolObject<dyn MTLBuffer>>> {
120        anyhow::ensure!(
121            (ptr as usize).is_multiple_of(PAGE_SIZE),
122            "region not page-aligned"
123        );
124        anyhow::ensure!(
125            len.is_multiple_of(PAGE_SIZE),
126            "region length not page-aligned"
127        );
128
129        let ptr = NonNull::new(ptr.cast::<c_void>()).context("null region")?;
130
131        self.check_allocation(len)?;
132
133        unsafe {
134            self.device
135                .newBufferWithBytesNoCopy_length_options_deallocator(
136                    ptr,
137                    len,
138                    MTLResourceOptions::empty(),
139                    None,
140                )
141        }
142        .context("newBufferWithBytesNoCopy failed for session region")
143    }
144
145    fn check_allocation(&self, len: usize) -> Result<()> {
146        if let Some(limit) = self.allocation_limit.get() {
147            anyhow::ensure!(
148                self.device
149                    .currentAllocatedSize()
150                    .checked_add(len)
151                    .is_some_and(|n| n <= limit),
152                "allocation would exceed the server memory budget; reduce pool or prefill chunk size"
153            );
154        }
155
156        Ok(())
157    }
158
159    pub fn new_buffer(&self, len: usize) -> Result<Retained<ProtocolObject<dyn MTLBuffer>>> {
160        self.check_allocation(len)?;
161
162        self.device
163            .newBufferWithLength_options(len, MTLResourceOptions::empty())
164            .context("newBufferWithLength failed")
165    }
166
167    pub fn throttle_probe(&self) -> Result<f64> {
168        let lib = self.compile_library(crate::kernels::CLOCK_PROBE_MSL)?;
169        let pso = self.pipeline(&lib, "clock_probe")?;
170        let out = self.new_buffer(65536 * 4)?;
171        let iters: u32 = 60_000;
172        let mut last = 0.0f64;
173
174        for _ in 0..2 {
175            let cb = self.queue.commandBuffer().context("commandBuffer")?;
176            let enc = cb.computeCommandEncoder().context("encoder")?;
177
178            enc.setComputePipelineState(&pso);
179
180            unsafe {
181                enc.setBuffer_offset_atIndex(Some(&out), 0, 0);
182                enc.setBytes_length_atIndex(
183                    NonNull::from(&iters).cast::<c_void>(),
184                    size_of::<u32>(),
185                    1,
186                );
187            }
188
189            let grid = MTLSize {
190                width: 65536,
191                height: 1,
192                depth: 1,
193            };
194            let tg = MTLSize {
195                width: 256,
196                height: 1,
197                depth: 1,
198            };
199
200            enc.dispatchThreads_threadsPerThreadgroup(grid, tg);
201            enc.endEncoding();
202
203            let start = Instant::now();
204
205            cb.commit();
206            cb.waitUntilCompleted();
207
208            last = start.elapsed().as_secs_f64() * 1e3;
209        }
210
211        Ok(last)
212    }
213}
214
215#[cfg(test)]
216#[path = "../tests/unit/metal.rs"]
217mod tests;
cherenkov/metal.rs

cherenkov/
metal.rs
