mlx_rs/
metal.rs

1//! Metal-specific runtime configuration.
2
3use std::ffi::{CStr, CString};
4
5use crate::{
6    error::{Exception, Result},
7    utils::SUCCESS,
8};
9
10struct StringHandle {
11    raw: mlx_sys::mlx_string,
12}
13
14impl StringHandle {
15    fn new() -> Self {
16        Self {
17            raw: unsafe { mlx_sys::mlx_string_new() },
18        }
19    }
20
21    fn to_string(&self) -> Result<String> {
22        let data = unsafe { mlx_sys::mlx_string_data(self.raw) };
23        if data.is_null() {
24            return Err(Exception::custom("MLX returned an empty string handle"));
25        }
26        unsafe { CStr::from_ptr(data) }
27            .to_str()
28            .map(str::to_owned)
29            .map_err(|error| Exception::custom(error.to_string()))
30    }
31}
32
33impl Drop for StringHandle {
34    fn drop(&mut self) {
35        let _ = unsafe { mlx_sys::mlx_string_free(self.raw) };
36    }
37}
38
39fn install_error_handler() {
40    crate::error::INIT_ERR_HANDLER.call_once(crate::error::setup_mlx_error_handler);
41}
42
43fn status_result(status: i32, operation: &str) -> Result<()> {
44    if status == SUCCESS {
45        Ok(())
46    } else {
47        Err(crate::error::exception_from_status(status, operation))
48    }
49}
50
51/// Returns the path used to load the default Metal library.
52pub fn metallib_path() -> Result<String> {
53    install_error_handler();
54    let mut path = StringHandle::new();
55    let status = unsafe { mlx_sys::mlx_metal_get_metallib_path(&mut path.raw) };
56    status_result(status, "reading the Metal library path")?;
57    path.to_string()
58}
59
60/// Sets the path used by subsequent Metal initialization.
61///
62/// This changes process-global MLX state and can affect every thread that initializes Metal
63/// afterward.
64pub fn set_metallib_path(path: impl AsRef<str>) -> Result<()> {
65    install_error_handler();
66    let path = CString::new(path.as_ref()).map_err(|error| Exception::custom(error.to_string()))?;
67    let status = unsafe { mlx_sys::mlx_metal_set_metallib_path(path.as_ptr()) };
68    status_result(status, "setting the Metal library path")
69}
70
71#[cfg(test)]
72mod tests {
73    use std::sync::Mutex;
74
75    use super::{metallib_path, set_metallib_path};
76
77    static METALLIB_PATH: Mutex<()> = Mutex::new(());
78
79    struct RestorePath(String);
80
81    impl Drop for RestorePath {
82        fn drop(&mut self) {
83            set_metallib_path(&self.0).expect("restore metallib path");
84        }
85    }
86
87    #[test]
88    fn metallib_path_is_non_empty() {
89        let _guard = METALLIB_PATH.lock().expect("lock metallib path");
90        let original = metallib_path().unwrap();
91        let _restore = RestorePath(original);
92        set_metallib_path("mlx-rs-metallib-path-test").unwrap();
93
94        assert!(!metallib_path().unwrap().is_empty());
95    }
96
97    #[test]
98    fn metallib_path_roundtrips() {
99        let _guard = METALLIB_PATH.lock().expect("lock metallib path");
100        let original = metallib_path().unwrap();
101        let _restore = RestorePath(original.clone());
102        let temporary = format!("{original}.mlx-rs-roundtrip");
103
104        set_metallib_path(&temporary).unwrap();
105
106        assert_eq!(metallib_path().unwrap(), temporary);
107    }
108}
mlx_rs/metal.rs

mlx_rs/
metal.rs
