mlx-rs 0.34.0

Unofficial rust wrapper for Apple's mlx machine learning library.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108

//! Metal-specific runtime configuration.

use std::ffi::{CStr, CString};

use crate::{
    error::{Exception, Result},
    utils::SUCCESS,
};

struct StringHandle {
    raw: mlx_sys::mlx_string,
}

impl StringHandle {
    fn new() -> Self {
        Self {
            raw: unsafe { mlx_sys::mlx_string_new() },
        }
    }

    fn to_string(&self) -> Result<String> {
        let data = unsafe { mlx_sys::mlx_string_data(self.raw) };
        if data.is_null() {
            return Err(Exception::custom("MLX returned an empty string handle"));
        }
        unsafe { CStr::from_ptr(data) }
            .to_str()
            .map(str::to_owned)
            .map_err(|error| Exception::custom(error.to_string()))
    }
}

impl Drop for StringHandle {
    fn drop(&mut self) {
        let _ = unsafe { mlx_sys::mlx_string_free(self.raw) };
    }
}

fn install_error_handler() {
    crate::error::ensure_mlx_error_handler();
}

fn status_result(status: i32, operation: &str) -> Result<()> {
    if status == SUCCESS {
        Ok(())
    } else {
        Err(crate::error::exception_from_status(status, operation))
    }
}

/// Returns the path used to load the default Metal library.
pub fn metallib_path() -> Result<String> {
    install_error_handler();
    let mut path = StringHandle::new();
    let status = unsafe { mlx_sys::mlx_metal_get_metallib_path(&mut path.raw) };
    status_result(status, "reading the Metal library path")?;
    path.to_string()
}

/// Sets the path used by subsequent Metal initialization.
///
/// This changes process-global MLX state and can affect every thread that initializes Metal
/// afterward.
pub fn set_metallib_path(path: impl AsRef<str>) -> Result<()> {
    install_error_handler();
    let path = CString::new(path.as_ref()).map_err(|error| Exception::custom(error.to_string()))?;
    let status = unsafe { mlx_sys::mlx_metal_set_metallib_path(path.as_ptr()) };
    status_result(status, "setting the Metal library path")
}

#[cfg(test)]
mod tests {
    use std::sync::Mutex;

    use super::{metallib_path, set_metallib_path};

    static METALLIB_PATH: Mutex<()> = Mutex::new(());

    struct RestorePath(String);

    impl Drop for RestorePath {
        fn drop(&mut self) {
            set_metallib_path(&self.0).expect("restore metallib path");
        }
    }

    #[test]
    fn metallib_path_is_non_empty() {
        let _guard = METALLIB_PATH.lock().expect("lock metallib path");
        let original = metallib_path().unwrap();
        let _restore = RestorePath(original);
        set_metallib_path("mlx-rs-metallib-path-test").unwrap();

        assert!(!metallib_path().unwrap().is_empty());
    }

    #[test]
    fn metallib_path_roundtrips() {
        let _guard = METALLIB_PATH.lock().expect("lock metallib path");
        let original = metallib_path().unwrap();
        let _restore = RestorePath(original.clone());
        let temporary = format!("{original}.mlx-rs-roundtrip");

        set_metallib_path(&temporary).unwrap();

        assert_eq!(metallib_path().unwrap(), temporary);
    }
}