mlx_rs/
stream.rs

1use std::{cell::RefCell, ffi::CStr, thread::LocalKey};
2
3use crate::{
4    device::Device,
5    error::Result,
6    utils::{guard::Guarded, SUCCESS},
7};
8
9thread_local! {
10    static THREAD_LOCAL_DEFAULT_STREAM: RefCell<Option<Stream>> = const { RefCell::new(None) };
11}
12
13struct ScopedValueGuard<T: 'static> {
14    local: &'static LocalKey<RefCell<Option<T>>>,
15    previous: Option<T>,
16}
17
18impl<T: 'static> Drop for ScopedValueGuard<T> {
19    fn drop(&mut self) {
20        self.local.with_borrow_mut(|stream| {
21            *stream = self.previous.take();
22        });
23    }
24}
25
26fn with_scoped_value<T: 'static, R>(
27    local: &'static LocalKey<RefCell<Option<T>>>,
28    value: T,
29    f: impl FnOnce() -> R,
30) -> R {
31    let previous = local.with_borrow_mut(|current| current.replace(value));
32    let _guard = ScopedValueGuard { local, previous };
33    f()
34}
35
36/// Gets the thread-local scoped default stream.
37///
38/// The value does not propagate across asynchronous task suspension or between operating-system
39/// threads.
40pub fn thread_local_default_stream() -> Option<Stream> {
41    THREAD_LOCAL_DEFAULT_STREAM.with_borrow(|s| s.clone())
42}
43
44/// Uses `stream` for operations constructed during `f`.
45///
46/// Scopes are synchronous, thread-local, nestable, and restore the previous stream if `f` panics.
47/// To select a stream for one operation, put only that operation in the closure:
48///
49/// ```rust
50/// use mlx_rs::{with_stream, Array, Stream};
51///
52/// let input = Array::from_slice(&[1.0_f32, 2.0, 3.0, 4.0], &[4]);
53/// let stream = Stream::cpu();
54/// let output = with_stream(&stream, || mlx_rs::fft::fft(&input, None, None)).unwrap();
55/// assert_eq!(output.shape(), &[4]);
56/// ```
57pub fn with_stream<F, T>(stream: &Stream, f: F) -> T
58where
59    F: FnOnce() -> T,
60{
61    with_scoped_value(&THREAD_LOCAL_DEFAULT_STREAM, stream.clone(), f)
62}
63
64/// Uses the default stream on `device` for operations constructed during `f`.
65///
66/// This is equivalent to creating a stream for the device and passing it to [`with_stream`].
67/// Scopes are synchronous, thread-local, nestable, and panic-safe.
68pub fn with_device<F, T>(device: Device, f: F) -> T
69where
70    F: FnOnce() -> T,
71{
72    let stream = Stream::new_with_device(&device);
73    with_stream(&stream, f)
74}
75
76/// Gets the thread-local scoped default stream.
77#[deprecated(since = "0.26.0", note = "use `thread_local_default_stream`")]
78pub fn task_local_default_stream() -> Option<Stream> {
79    thread_local_default_stream()
80}
81
82/// Uses a given default stream for the duration of `f`.
83#[deprecated(since = "0.26.0", note = "use `with_stream(&stream, f)`")]
84pub fn with_new_default_stream<F, T>(default_stream: Stream, f: F) -> T
85where
86    F: FnOnce() -> T,
87{
88    with_stream(&default_stream, f)
89}
90
91/// Parameter type for all MLX operations.
92///
93/// Use this to control where operations are evaluated:
94///
95/// If omitted it will use the [Default::default()], which will be [Device::gpu()] unless
96/// set otherwise.
97#[derive(PartialEq)]
98pub struct StreamOrDevice {
99    pub(crate) stream: Stream,
100}
101
102impl StreamOrDevice {
103    /// Create a new [`StreamOrDevice`] with a [`Stream`].
104    pub fn new(stream: Stream) -> StreamOrDevice {
105        StreamOrDevice { stream }
106    }
107
108    /// Create a new [`StreamOrDevice`] with a [`Device`].
109    pub fn new_with_device(device: &Device) -> StreamOrDevice {
110        StreamOrDevice {
111            stream: Stream::new_with_device(device),
112        }
113    }
114
115    /// Current default CPU stream.
116    pub fn cpu() -> StreamOrDevice {
117        StreamOrDevice {
118            stream: Stream::cpu(),
119        }
120    }
121
122    /// Current default GPU stream.
123    pub fn gpu() -> StreamOrDevice {
124        StreamOrDevice {
125            stream: Stream::gpu(),
126        }
127    }
128}
129
130impl Default for StreamOrDevice {
131    /// The default stream on the default device.
132    ///
133    /// This will be [Device::gpu()] unless [Device::set_default()]
134    /// sets it otherwise.
135    fn default() -> Self {
136        Self {
137            stream: Stream::new(),
138        }
139    }
140}
141
142impl AsRef<Stream> for StreamOrDevice {
143    fn as_ref(&self) -> &Stream {
144        &self.stream
145    }
146}
147
148impl std::fmt::Debug for StreamOrDevice {
149    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
150        write!(f, "{}", self.stream)
151    }
152}
153
154impl std::fmt::Display for StreamOrDevice {
155    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
156        write!(f, "{}", self.stream)
157    }
158}
159
160/// A stream of evaluation attached to a particular device.
161///
162/// Typically, this is used via the `stream:` parameter on a method with a [StreamOrDevice]:
163pub struct Stream {
164    pub(crate) c_stream: mlx_sys::mlx_stream,
165}
166
167impl AsRef<Stream> for Stream {
168    fn as_ref(&self) -> &Stream {
169        self
170    }
171}
172
173impl Clone for Stream {
174    fn clone(&self) -> Self {
175        Stream::try_from_op(|res| unsafe { mlx_sys::mlx_stream_set(res, self.c_stream) })
176            .expect("Failed to clone stream")
177    }
178}
179
180impl Stream {
181    /// Create a new stream on the default device, or return the thread-local
182    /// default stream if present.
183    pub fn thread_local_or_default() -> Self {
184        thread_local_default_stream().unwrap_or_default()
185    }
186
187    /// Create a new stream on the default cpu device, or return the thread-local
188    /// default stream if present.
189    pub fn thread_local_or_cpu() -> Self {
190        thread_local_default_stream().unwrap_or_else(Stream::cpu)
191    }
192
193    /// Create a new stream on the default gpu device, or return the thread-local
194    /// default stream if present.
195    pub fn thread_local_or_gpu() -> Self {
196        thread_local_default_stream().unwrap_or_else(Stream::gpu)
197    }
198
199    /// Returns the thread-local scoped stream or the default stream.
200    #[deprecated(since = "0.26.0", note = "use `Stream::thread_local_or_default`")]
201    pub fn task_local_or_default() -> Self {
202        Self::thread_local_or_default()
203    }
204
205    /// Returns the thread-local scoped stream or the CPU stream.
206    #[deprecated(since = "0.26.0", note = "use `Stream::thread_local_or_cpu`")]
207    pub fn task_local_or_cpu() -> Self {
208        Self::thread_local_or_cpu()
209    }
210
211    /// Returns the thread-local scoped stream or the GPU stream.
212    #[deprecated(since = "0.26.0", note = "use `Stream::thread_local_or_gpu`")]
213    pub fn task_local_or_gpu() -> Self {
214        Self::thread_local_or_gpu()
215    }
216
217    /// Create a new stream on the default device. Panics if fails.
218    pub fn new() -> Stream {
219        unsafe {
220            let mut dev = mlx_sys::mlx_device_new();
221            // SAFETY: mlx_get_default_device internally never throws an error
222            mlx_sys::mlx_get_default_device(&mut dev as *mut _);
223
224            let mut c_stream = mlx_sys::mlx_stream_new();
225            // SAFETY: mlx_get_default_stream internally never throws if dev is valid
226            mlx_sys::mlx_get_default_stream(&mut c_stream as *mut _, dev);
227
228            mlx_sys::mlx_device_free(dev);
229            Stream { c_stream }
230        }
231    }
232
233    /// Try to get the default stream on the given device.
234    pub fn try_default_on_device(device: &Device) -> Result<Stream> {
235        Stream::try_from_op(|res| unsafe { mlx_sys::mlx_get_default_stream(res, device.c_device) })
236    }
237
238    /// Create a new stream on the given device
239    pub fn new_with_device(device: &Device) -> Stream {
240        unsafe {
241            let c_stream = mlx_sys::mlx_stream_new_device(device.c_device);
242            Stream { c_stream }
243        }
244    }
245
246    /// Get the underlying C pointer.
247    pub fn as_ptr(&self) -> mlx_sys::mlx_stream {
248        self.c_stream
249    }
250
251    /// Current default CPU stream.
252    pub fn cpu() -> Self {
253        unsafe {
254            let c_stream = mlx_sys::mlx_default_cpu_stream_new();
255            Stream { c_stream }
256        }
257    }
258
259    /// Current default GPU stream.
260    pub fn gpu() -> Self {
261        unsafe {
262            let c_stream = mlx_sys::mlx_default_gpu_stream_new();
263            Stream { c_stream }
264        }
265    }
266
267    /// Get the index of the stream.
268    pub fn get_index(&self) -> Result<i32> {
269        i32::try_from_op(|res| unsafe { mlx_sys::mlx_stream_get_index(res, self.c_stream) })
270    }
271
272    fn describe(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
273        unsafe {
274            let mut mlx_str = mlx_sys::mlx_string_new();
275            let result = match mlx_sys::mlx_stream_tostring(&mut mlx_str as *mut _, self.c_stream) {
276                SUCCESS => {
277                    let ptr = mlx_sys::mlx_string_data(mlx_str);
278                    let c_str = CStr::from_ptr(ptr);
279                    write!(f, "{}", c_str.to_string_lossy())
280                }
281                _ => Err(std::fmt::Error),
282            };
283            mlx_sys::mlx_string_free(mlx_str);
284            result
285        }
286    }
287}
288
289impl Drop for Stream {
290    fn drop(&mut self) {
291        unsafe { mlx_sys::mlx_stream_free(self.c_stream) };
292    }
293}
294
295impl Default for Stream {
296    fn default() -> Self {
297        Stream::new()
298    }
299}
300
301impl std::fmt::Debug for Stream {
302    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
303        self.describe(f)
304    }
305}
306
307impl std::fmt::Display for Stream {
308    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
309        self.describe(f)
310    }
311}
312
313impl PartialEq for Stream {
314    fn eq(&self, other: &Self) -> bool {
315        unsafe { mlx_sys::mlx_stream_equal(self.c_stream, other.c_stream) }
316    }
317}
318
319#[cfg(test)]
320mod tests {
321    use super::*;
322
323    #[test]
324    fn canonical_scopes_nest_and_restore_after_panic() {
325        let outer = Stream::cpu();
326        with_stream(&outer, || {
327            assert_eq!(thread_local_default_stream(), Some(outer.clone()));
328
329            let panic = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
330                with_device(Device::cpu(), || panic!("scope panic"));
331            }));
332
333            assert!(panic.is_err());
334            assert_eq!(thread_local_default_stream(), Some(outer.clone()));
335        });
336
337        assert!(thread_local_default_stream().is_none());
338    }
339
340    #[test]
341    fn test_scoped_default_stream() {
342        // First set default stream to CPU
343        let cpu_device = Device::cpu();
344        Device::set_default(&cpu_device);
345        let cpu_stream = Stream::default();
346
347        let task_default_stream = Stream::gpu();
348        with_stream(&task_default_stream, || {
349            let task_local_stream_0 = Stream::thread_local_or_default();
350            let task_local_stream_1 = Stream::thread_local_or_default();
351            assert_eq!(task_local_stream_0, task_local_stream_1);
352            assert_ne!(task_local_stream_0, cpu_stream);
353        });
354    }
355
356    #[test]
357    fn test_scoped_default_stream_restored_after_panic() {
358        let cpu = Device::cpu();
359        let outer_stream = Stream::new_with_device(&cpu);
360        with_stream(&outer_stream, || {
361            let panic = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
362                let inner = Stream::new_with_device(&cpu);
363                with_stream(&inner, || panic!("stream panic"));
364            }));
365
366            assert!(panic.is_err());
367            assert_eq!(thread_local_default_stream(), Some(outer_stream.clone()));
368        });
369
370        assert!(thread_local_default_stream().is_none());
371    }
372
373    #[test]
374    fn test_stream_clone() {
375        let stream = Stream::new();
376        let cloned_stream = stream.clone();
377        assert_eq!(stream, cloned_stream);
378    }
379
380    #[test]
381    fn test_cpu_gpu_stream_not_equal() {
382        let cpu_device = Device::cpu();
383        let gpu_device = Device::gpu();
384
385        // First set default stream to CPU
386        Device::set_default(&cpu_device);
387        let cpu_stream = Stream::default();
388
389        // Then set default stream to GPU
390        Device::set_default(&gpu_device);
391        let gpu_stream = Stream::default();
392
393        // Assert that CPU and GPU streams are not equal
394        assert_ne!(cpu_stream, gpu_stream);
395    }
396}
mlx_rs/stream.rs

mlx_rs/
stream.rs
