mlx_c_rust/types/
stream.rs

1use crate::MlxCError;
2use crate::raw;
3
4/// Owned MLX stream handle used to preserve runtime thread affinity.
5#[derive(Debug)]
6pub struct MlxStream(raw::mlx_stream);
7
8impl MlxStream {
9    pub fn default_cpu() -> Result<Self, MlxCError> {
10        // SAFETY: The runtime error handler is installed and the returned
11        // stream is placed immediately under RAII ownership.
12        let raw_stream = unsafe { raw::mlx_default_cpu_stream_new() };
13        Self::from_raw(raw_stream, "acquire the default MLX CPU stream")
14    }
15
16    pub fn default_gpu() -> Result<Self, MlxCError> {
17        // SAFETY: The runtime error handler is installed and the returned
18        // stream is placed immediately under RAII ownership.
19        let raw_stream = unsafe { raw::mlx_default_gpu_stream_new() };
20        Self::from_raw(raw_stream, "acquire the default MLX GPU stream")
21    }
22
23    pub const fn raw(&self) -> raw::mlx_stream {
24        self.0
25    }
26
27    fn from_raw(raw_stream: raw::mlx_stream, operation: &'static str) -> Result<Self, MlxCError> {
28        if raw_stream.ctx.is_null() {
29            return Err(MlxCError {
30                operation,
31                description: "MLX returned an empty stream handle".to_owned(),
32            });
33        }
34        Ok(Self(raw_stream))
35    }
36
37    pub(crate) fn from_live_raw(
38        raw_stream: raw::mlx_stream,
39        operation: &'static str,
40    ) -> Result<Self, MlxCError> {
41        Self::from_raw(raw_stream, operation)
42    }
43}
44
45impl Drop for MlxStream {
46    fn drop(&mut self) {
47        // SAFETY: This owner releases its live stream handle exactly once.
48        unsafe {
49            raw::mlx_stream_free(self.0);
50        }
51    }
52}
mlx_c_rust/types/stream.rs

mlx_c_rust/types/
stream.rs
