nuvec_gpu/metal/
capture.rs

1//! GPU captures (MLX's `mx.metal.start_capture`): the command buffers committed between
2//! start and stop, written to a `.gputrace` document for Xcode.
3
4use std::path::Path;
5
6use objc2::rc::autoreleasepool;
7use objc2_foundation::{NSString, NSURL};
8use objc2_metal::{MTLCaptureDescriptor, MTLCaptureDestination, MTLCaptureManager};
9
10use super::Device;
11use crate::{GpuError, Result};
12
13impl Device {
14    /// Starts capturing this device's GPU work into `path`, which must end in `.gputrace`
15    /// and not exist. Metal allows captures only with `MTL_CAPTURE_ENABLED=1` in the
16    /// environment (or under Xcode); one capture runs at a time.
17    pub fn start_capture(&self, path: impl AsRef<Path>) -> Result<()> {
18        let path = path.as_ref();
19        if path.extension().is_none_or(|e| e != "gputrace") {
20            return Err(GpuError::Capture(format!(
21                "the capture file must end in .gputrace: {}",
22                path.display()
23            )));
24        }
25        if path.exists() {
26            return Err(GpuError::Capture(format!(
27                "the capture file already exists: {}",
28                path.display()
29            )));
30        }
31        autoreleasepool(|_| {
32            // SAFETY: the shared manager is a process-wide singleton, valid for the call.
33            let manager = unsafe { MTLCaptureManager::sharedCaptureManager() };
34            if !manager.supportsDestination(MTLCaptureDestination::GPUTraceDocument) {
35                return Err(GpuError::Capture(
36                    "captures to a file are disabled: set MTL_CAPTURE_ENABLED=1".into(),
37                ));
38            }
39            let descriptor = MTLCaptureDescriptor::new();
40            let device: &objc2::runtime::AnyObject = self.inner.device.as_ref();
41            // SAFETY: the capture object is a live MTLDevice, the kind Metal expects.
42            unsafe { descriptor.setCaptureObject(Some(device)) };
43            descriptor.setDestination(MTLCaptureDestination::GPUTraceDocument);
44            let url = NSURL::fileURLWithPath(&NSString::from_str(&path.to_string_lossy()));
45            descriptor.setOutputURL(Some(&url));
46            manager
47                .startCaptureWithDescriptor_error(&descriptor)
48                .map_err(|e| GpuError::Capture(e.localizedDescription().to_string()))
49        })
50    }
51
52    /// Stops the running capture, writing its document (nothing if none runs).
53    pub fn stop_capture(&self) {
54        autoreleasepool(|_| {
55            // SAFETY: as in `start_capture`.
56            let manager = unsafe { MTLCaptureManager::sharedCaptureManager() };
57            if manager.isCapturing() {
58                manager.stopCapture();
59            }
60        });
61    }
62}
63
64#[cfg(test)]
65mod tests {
66    use super::*;
67
68    #[test]
69    fn capture_paths_and_the_environment_are_checked() {
70        let Some(d) = Device::shared() else {
71            return;
72        };
73        let dir = std::env::temp_dir().join(format!("nuvec-capture-{}", std::process::id()));
74        let e = d.start_capture(dir.join("trace.txt")).unwrap_err();
75        assert!(e.to_string().contains(".gputrace"), "{e}");
76        let trace = dir.join("trace.gputrace");
77        if std::env::var_os("MTL_CAPTURE_ENABLED").is_none() {
78            let e = d.start_capture(&trace).unwrap_err();
79            assert!(e.to_string().contains("MTL_CAPTURE_ENABLED"), "{e}");
80        } else {
81            std::fs::create_dir_all(&dir).unwrap();
82            d.start_capture(&trace).unwrap();
83            let again = d.start_capture(dir.join("again.gputrace")).unwrap_err();
84            assert!(matches!(again, GpuError::Capture(_)), "{again}");
85            d.stop_capture();
86            assert!(trace.exists());
87            let e = d.start_capture(&trace).unwrap_err();
88            assert!(e.to_string().contains("already exists"), "{e}");
89            std::fs::remove_dir_all(&dir).unwrap();
90        }
91        d.stop_capture();
92    }
93}
nuvec_gpu/metal/capture.rs

nuvec_gpu/metal/
capture.rs
