mlx-rs 0.34.0

use std::{cell::RefCell, ffi::CStr, thread::LocalKey};

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

thread_local! {
    static THREAD_LOCAL_DEFAULT_STREAM: RefCell<Option<Stream>> = const { RefCell::new(None) };
}

struct ScopedValueGuard<T: 'static> {
    local: &'static LocalKey<RefCell<Option<T>>>,
    previous: Option<T>,
}

impl<T: 'static> Drop for ScopedValueGuard<T> {
    fn drop(&mut self) {
        self.local.with_borrow_mut(|stream| {
            *stream = self.previous.take();
        });
    }
}

fn with_scoped_value<T: 'static, R>(
    local: &'static LocalKey<RefCell<Option<T>>>,
    value: T,
    f: impl FnOnce() -> R,
) -> R {
    let previous = local.with_borrow_mut(|current| current.replace(value));
    let _guard = ScopedValueGuard { local, previous };
    f()
}

/// Gets the thread-local scoped default stream.
///
/// The value does not propagate across asynchronous task suspension or between operating-system
/// threads.
pub fn thread_local_default_stream() -> Option<Stream> {
    THREAD_LOCAL_DEFAULT_STREAM.with_borrow(|s| s.clone())
}

/// Uses `stream` for operations constructed during `f`.
///
/// Scopes are synchronous, thread-local, nestable, and restore the previous stream if `f` panics.
/// To select a stream for one operation, put only that operation in the closure:
///
/// ```rust
/// use mlx_rs::{with_stream, Array, Stream};
///
/// let input = Array::from_slice(&[1.0_f32, 2.0, 3.0, 4.0], &[4]);
/// let stream = Stream::cpu();
/// let output = with_stream(&stream, || mlx_rs::fft::fft(&input, None, None)).unwrap();
/// assert_eq!(output.shape(), &[4]);
/// ```
pub fn with_stream<F, T>(stream: &Stream, f: F) -> T
where
    F: FnOnce() -> T,
{
    with_scoped_value(&THREAD_LOCAL_DEFAULT_STREAM, stream.clone(), f)
}

/// Uses the default stream on `device` for operations constructed during `f`.
///
/// This is equivalent to creating a stream for the device and passing it to [`with_stream`].
/// Scopes are synchronous, thread-local, nestable, and panic-safe.
pub fn with_device<F, T>(device: Device, f: F) -> T
where
    F: FnOnce() -> T,
{
    let stream = Stream::new_with_device(&device);
    with_stream(&stream, f)
}

/// Gets the thread-local scoped default stream.
#[deprecated(since = "0.26.0", note = "use `thread_local_default_stream`")]
pub fn task_local_default_stream() -> Option<Stream> {
    thread_local_default_stream()
}

/// Uses a given default stream for the duration of `f`.
#[deprecated(since = "0.26.0", note = "use `with_stream(&stream, f)`")]
pub fn with_new_default_stream<F, T>(default_stream: Stream, f: F) -> T
where
    F: FnOnce() -> T,
{
    with_stream(&default_stream, f)
}

/// Parameter type for all MLX operations.
///
/// Use this to control where operations are evaluated:
///
/// If omitted it will use the [Default::default()], which will be [Device::gpu()] unless
/// set otherwise.
#[derive(PartialEq)]
pub struct StreamOrDevice {
    pub(crate) stream: Stream,
}

impl StreamOrDevice {
    /// Create a new [`StreamOrDevice`] with a [`Stream`].
    pub fn new(stream: Stream) -> StreamOrDevice {
        StreamOrDevice { stream }
    }

    /// Create a new [`StreamOrDevice`] with a [`Device`].
    pub fn new_with_device(device: &Device) -> StreamOrDevice {
        StreamOrDevice {
            stream: Stream::new_with_device(device),
        }
    }

    /// Current default CPU stream.
    pub fn cpu() -> StreamOrDevice {
        StreamOrDevice {
            stream: Stream::cpu(),
        }
    }

    /// Current default GPU stream.
    pub fn gpu() -> StreamOrDevice {
        StreamOrDevice {
            stream: Stream::gpu(),
        }
    }
}

impl Default for StreamOrDevice {
    /// The default stream on the default device.
    ///
    /// Uses MLX's process-wide default device. For safe, thread-local operation
    /// selection, use [with_device] or [with_stream].
    fn default() -> Self {
        Self {
            stream: Stream::new(),
        }
    }
}

impl AsRef<Stream> for StreamOrDevice {
    fn as_ref(&self) -> &Stream {
        &self.stream
    }
}

impl std::fmt::Debug for StreamOrDevice {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "{}", self.stream)
    }
}

impl std::fmt::Display for StreamOrDevice {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "{}", self.stream)
    }
}

/// A stream of evaluation attached to a particular device.
///
/// Typically, this is used via the `stream:` parameter on a method with a [StreamOrDevice]:
pub struct Stream {
    pub(crate) c_stream: mlx_sys::mlx_stream,
}

impl AsRef<Stream> for Stream {
    fn as_ref(&self) -> &Stream {
        self
    }
}

impl Clone for Stream {
    fn clone(&self) -> Self {
        Stream::try_from_op(|res| unsafe { mlx_sys::mlx_stream_set(res, self.c_stream) })
            .expect("Failed to clone stream")
    }
}

impl Stream {
    /// Create a new stream on the default device, or return the thread-local
    /// default stream if present.
    pub fn thread_local_or_default() -> Self {
        thread_local_default_stream().unwrap_or_default()
    }

    /// Create a new stream on the default cpu device, or return the thread-local
    /// default stream if present.
    pub fn thread_local_or_cpu() -> Self {
        thread_local_default_stream().unwrap_or_else(Stream::cpu)
    }

    /// Create a new stream on the default gpu device, or return the thread-local
    /// default stream if present.
    pub fn thread_local_or_gpu() -> Self {
        thread_local_default_stream().unwrap_or_else(Stream::gpu)
    }

    /// Returns the thread-local scoped stream or the default stream.
    #[deprecated(since = "0.26.0", note = "use `Stream::thread_local_or_default`")]
    pub fn task_local_or_default() -> Self {
        Self::thread_local_or_default()
    }

    /// Returns the thread-local scoped stream or the CPU stream.
    #[deprecated(since = "0.26.0", note = "use `Stream::thread_local_or_cpu`")]
    pub fn task_local_or_cpu() -> Self {
        Self::thread_local_or_cpu()
    }

    /// Returns the thread-local scoped stream or the GPU stream.
    #[deprecated(since = "0.26.0", note = "use `Stream::thread_local_or_gpu`")]
    pub fn task_local_or_gpu() -> Self {
        Self::thread_local_or_gpu()
    }

    /// Create a new stream on the default device. Panics if fails.
    pub fn new() -> Stream {
        Self::try_default_on_device(&Device::default()).expect("Failed to create stream")
    }

    /// Try to get the default stream on the given device.
    pub fn try_default_on_device(device: &Device) -> Result<Stream> {
        Stream::try_from_op(|res| unsafe { mlx_sys::mlx_get_default_stream(res, device.c_device) })
    }

    /// Create a new stream on the given device
    pub fn new_with_device(device: &Device) -> Stream {
        crate::error::infallible(|| unsafe {
            let c_stream = mlx_sys::mlx_stream_new_device(device.c_device);
            Stream { c_stream }
        })
    }

    /// Get the underlying C pointer.
    pub fn as_ptr(&self) -> mlx_sys::mlx_stream {
        self.c_stream
    }

    /// Current default CPU stream.
    pub fn cpu() -> Self {
        crate::error::infallible(|| unsafe {
            let c_stream = mlx_sys::mlx_default_cpu_stream_new();
            Stream { c_stream }
        })
    }

    /// Current default GPU stream.
    pub fn gpu() -> Self {
        crate::error::infallible(|| unsafe {
            let c_stream = mlx_sys::mlx_default_gpu_stream_new();
            Stream { c_stream }
        })
    }

    /// Returns the device this stream executes on.
    pub fn device(&self) -> Result<Device> {
        Device::try_from_op(|res| unsafe { mlx_sys::mlx_stream_get_device(res, self.c_stream) })
    }

    /// Get the index of the stream.
    pub fn get_index(&self) -> Result<i32> {
        i32::try_from_op(|res| unsafe { mlx_sys::mlx_stream_get_index(res, self.c_stream) })
    }

    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_stream_tostring(&mut mlx_str as *mut _, self.c_stream) {
                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 Stream {
    fn drop(&mut self) {
        unsafe { mlx_sys::mlx_stream_free(self.c_stream) };
    }
}

impl Default for Stream {
    fn default() -> Self {
        Stream::new()
    }
}

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

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

impl PartialEq for Stream {
    fn eq(&self, other: &Self) -> bool {
        unsafe { mlx_sys::mlx_stream_equal(self.c_stream, other.c_stream) }
    }
}

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

    #[test]
    fn canonical_scopes_nest_and_restore_after_panic() {
        let outer = Stream::cpu();
        with_stream(&outer, || {
            assert_eq!(thread_local_default_stream(), Some(outer.clone()));

            let panic = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
                with_device(Device::cpu(), || panic!("scope panic"));
            }));

            assert!(panic.is_err());
            assert_eq!(thread_local_default_stream(), Some(outer.clone()));
        });

        assert!(thread_local_default_stream().is_none());
    }

    #[test]
    #[cfg(feature = "metal")]
    fn test_scoped_default_stream() {
        with_device(Device::cpu(), || {
            let cpu_stream = Stream::thread_local_or_default();
            let task_default_stream = Stream::gpu();
            with_stream(&task_default_stream, || {
                let task_local_stream_0 = Stream::thread_local_or_default();
                let task_local_stream_1 = Stream::thread_local_or_default();
                assert_eq!(task_local_stream_0, task_local_stream_1);
                assert_ne!(task_local_stream_0, cpu_stream);
            });
            assert_eq!(Stream::thread_local_or_default(), cpu_stream);
        });
    }

    #[test]
    fn test_scoped_default_stream_restored_after_panic() {
        let cpu = Device::cpu();
        let outer_stream = Stream::new_with_device(&cpu);
        with_stream(&outer_stream, || {
            let panic = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
                let inner = Stream::new_with_device(&cpu);
                with_stream(&inner, || panic!("stream panic"));
            }));

            assert!(panic.is_err());
            assert_eq!(thread_local_default_stream(), Some(outer_stream.clone()));
        });

        assert!(thread_local_default_stream().is_none());
    }

    #[test]
    fn test_stream_clone() {
        let stream = Stream::new();
        let cloned_stream = stream.clone();
        assert_eq!(stream, cloned_stream);
    }

    #[test]
    #[cfg(feature = "metal")]
    fn test_cpu_gpu_stream_not_equal() {
        const CHILD: &str = "MLX_RS_DEFAULT_DEVICE_CHILD";
        if std::env::var_os(CHILD).is_some() {
            // SAFETY: This child runs only this test, with no other MLX threads.
            unsafe { Device::set_default(&Device::cpu()).unwrap() };
            assert_eq!(Device::default(), Device::cpu());
            let cpu_stream = Stream::default();

            unsafe { Device::set_default(&Device::gpu()).unwrap() };
            assert_eq!(Device::default(), Device::gpu());
            let gpu_stream = Stream::default();
            assert_ne!(cpu_stream, gpu_stream);
            return;
        }

        let output = std::process::Command::new(std::env::current_exe().unwrap())
            .args([
                "--exact",
                "stream::tests::test_cpu_gpu_stream_not_equal",
                "--nocapture",
                "--test-threads=1",
            ])
            .env(CHILD, "1")
            .output()
            .unwrap();

        assert!(
            output.status.success(),
            "child status: {:?}\nstdout:\n{}\nstderr:\n{}",
            output.status,
            String::from_utf8_lossy(&output.stdout),
            String::from_utf8_lossy(&output.stderr)
        );
    }
}

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

    #[test]
    fn stream_device_survives_stream_drop_and_respects_scopes() -> Result<()> {
        for device in [
            Device::cpu(),
            #[cfg(feature = "metal")]
            Device::gpu(),
        ] {
            let stream = Stream::new_with_device(&device);
            let queried = stream.device()?;
            assert_eq!(queried, device);
            with_stream(&stream, || -> Result<()> {
                assert_eq!(Stream::thread_local_or_default().device()?, device);
                Ok(())
            })?;
            drop(stream);
            assert_eq!(queried, device);
        }
        Ok(())
    }
}