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
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141

use std::ffi::CStr;

use crate::{
    error::Result,
    utils::{guard::Guarded, SUCCESS},
};

///Type of device.
#[derive(num_enum::IntoPrimitive, Debug, Clone, Copy)]
#[repr(u32)]
pub enum DeviceType {
    /// CPU device
    Cpu = mlx_sys::mlx_device_type__MLX_CPU,

    /// GPU device
    Gpu = mlx_sys::mlx_device_type__MLX_GPU,
}

/// Representation of a Device in MLX.
pub struct Device {
    pub(crate) c_device: mlx_sys::mlx_device,
}

impl PartialEq for Device {
    fn eq(&self, other: &Self) -> bool {
        unsafe { mlx_sys::mlx_device_equal(self.c_device, other.c_device) }
    }
}

impl Device {
    /// Create a new [`Device`]
    pub fn new(device_type: DeviceType, index: i32) -> Device {
        crate::error::infallible(|| {
            let c_device = unsafe { mlx_sys::mlx_device_new_type(device_type.into(), index) };
            Device { c_device }
        })
    }

    /// Try to get the default device.
    pub fn try_default() -> Result<Self> {
        Device::try_from_op(|res| unsafe { mlx_sys::mlx_get_default_device(res) })
    }

    /// Create a default CPU device.
    pub fn cpu() -> Device {
        Device::new(DeviceType::Cpu, 0)
    }

    /// Create a default GPU device.
    pub fn gpu() -> Device {
        Device::new(DeviceType::Gpu, 0)
    }

    /// Get the device index
    pub fn get_index(&self) -> Result<i32> {
        i32::try_from_op(|res| unsafe { mlx_sys::mlx_device_get_index(res, self.c_device) })
    }

    /// Get the device type
    pub fn get_type(&self) -> Result<DeviceType> {
        DeviceType::try_from_op(|res| unsafe { mlx_sys::mlx_device_get_type(res, self.c_device) })
    }

    /// Set the process-wide default device.
    ///
    /// Prefer [with_device](crate::with_device) for safe, thread-local device selection.
    ///
    /// # Safety
    ///
    /// No other thread may be using MLX while this function runs. The change is visible
    /// to every thread that has not scoped its own device or stream.
    ///
    /// # Example
    ///
    /// ```rust
    /// use mlx_rs::Device;
    ///
    /// // SAFETY: No other thread is using MLX.
    /// unsafe { Device::set_default(&Device::cpu())? };
    /// # Ok::<(), mlx_rs::error::Exception>(())
    /// ```
    pub unsafe fn set_default(device: &Device) -> Result<()> {
        <()>::try_from_op(|_| unsafe { mlx_sys::mlx_set_default_device(device.c_device) })
    }

    fn describe(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        unsafe {
            let mut mlx_str = mlx_sys::mlx_string_new();
            let result = match mlx_sys::mlx_device_tostring(&mut mlx_str as *mut _, self.c_device) {
                SUCCESS => {
                    let ptr = mlx_sys::mlx_string_data(mlx_str);
                    let c_str = CStr::from_ptr(ptr);
                    write!(f, "{}", c_str.to_string_lossy())
                }
                _ => Err(std::fmt::Error),
            };
            mlx_sys::mlx_string_free(mlx_str);
            result
        }
    }
}

impl Drop for Device {
    fn drop(&mut self) {
        let status = unsafe { mlx_sys::mlx_device_free(self.c_device) };
        debug_assert_eq!(status, SUCCESS);
    }
}

impl Default for Device {
    fn default() -> Self {
        Self::try_default().unwrap()
    }
}

impl std::fmt::Debug for Device {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        self.describe(f)
    }
}

impl std::fmt::Display for Device {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        self.describe(f)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_fmt() {
        let device = Device::default();
        let description = format!("{device}");
        #[cfg(feature = "metal")]
        assert_eq!(description, "Device(gpu, 0)");
        #[cfg(not(feature = "metal"))]
        assert_eq!(description, "Device(cpu, 0)");
    }
}