onnxruntime-ep-mlx 0.29.4

//! Safe RAII wrappers over the raw `sys::mlx` bindgen bindings.
//!
//! This is where the memory-safety win of the Rust rewrite lives: every MLX handle is owned by a
//! wrapper whose `Drop` calls the matching `mlx_*_free`, so op handlers never free by hand and a
//! leaked / double-freed `mlx_array` (a class of bug the C++ EP hit repeatedly) is impossible by
//! construction. Raw `unsafe`/FFI stays confined to `sys::mlx`; the engine and ops use these types.

use crate::sys::mlx;

/// Owning wrapper over an `mlx_stream` (freed once on drop).
pub struct Stream {
    raw: mlx::mlx_stream,
}

impl Stream {
    /// The default GPU stream (what every op in a plan runs on).
    pub fn new_default_gpu() -> Self {
        Stream {
            raw: unsafe { mlx::mlx_default_gpu_stream_new() },
        }
    }

    /// The default CPU stream. MLX hard-errors on float64 on the GPU
    /// ("float64 is not supported on the GPU"), so an fp64 subgraph must run here.
    pub fn new_default_cpu() -> Self {
        Stream {
            raw: unsafe { mlx::mlx_default_cpu_stream_new() },
        }
    }

    pub fn new_gpu() -> Self {
        unsafe {
            let device = mlx::mlx_device_new_type(mlx::mlx_device_type__MLX_GPU, 0);
            let raw = mlx::mlx_stream_new_device(device);
            mlx::mlx_device_free(device);
            Stream { raw }
        }
    }

    #[inline]
    pub fn as_raw(&self) -> mlx::mlx_stream {
        self.raw
    }
}

impl Drop for Stream {
    fn drop(&mut self) {
        unsafe { mlx::mlx_stream_free(self.raw) };
    }
}

/// Owning wrapper over an `mlx_array`. Holds exactly one reference; `Drop` releases it.
///
/// MLX ops do NOT consume their operands — they take their own internal references — so a handler
/// resolves an input to a borrowed raw handle (`as_raw`) and only the wrapper owns the reference.
/// Freshly produced arrays are wrapped with `from_raw` and kept alive (in the run arena or the plan
/// cache) until they are no longer needed.
pub struct Array {
    raw: mlx::mlx_array,
}

impl Array {
    /// Take ownership of a raw handle returned by an `mlx_*` call (e.g. the `res` out-param).
    #[inline]
    pub fn from_raw(raw: mlx::mlx_array) -> Self {
        Array { raw }
    }

    /// A fresh, empty array handle (the `mlx_array_new()` out-param sink for op results).
    #[inline]
    pub fn new() -> Self {
        Array {
            raw: unsafe { mlx::mlx_array_new() },
        }
    }

    /// Wrap host bytes into a new MLX array of `dtype` and the given shape (row-major). MLX copies
    /// the data (managed lifetime), so the source buffer need not outlive the array.
    pub fn from_data(
        data: *const std::os::raw::c_void,
        shape: &[i32],
        dtype: mlx::mlx_dtype,
    ) -> Self {
        let arr = Array {
            raw: unsafe {
                mlx::mlx_array_new_data(data, shape.as_ptr(), shape.len() as i32, dtype)
            },
        };
        // Memory view: a COPY-wrap (MLX copies the bytes into managed memory). Gated so a
        // traced-off run pays a single atomic load.
        let tr = crate::trace::tracer();
        if tr.active() {
            tr.record_copy_wrap((arr.size() * arr.itemsize()) as u64);
        }
        arr
    }

    /// Wrap an externally-owned buffer WITHOUT copying (zero-copy). MLX takes the raw pointer and,
    /// on Apple unified memory, hands it straight to Metal via `newBufferWithBytesNoCopy` — no host
    /// memcpy. If the pointer is not page-aligned (so Metal refuses the no-copy buffer) MLX silently
    /// falls back to allocate+copy, so correctness is preserved unconditionally; only the perf win is
    /// conditional on alignment.
    ///
    /// SAFETY / LIFETIME: `data` is owned by the caller. The registered deallocator is a NO-OP, so
    /// MLX never frees `data`. Runtime input wrappers must be dropped after the synchronous eval;
    /// cached initializer wrappers may live longer because ORT keeps initializer storage alive for
    /// the owning session. In either case the caller must keep the buffer valid and immutable until
    /// the last MLX array referencing it is dropped.
    pub fn from_data_managed(
        data: *const std::os::raw::c_void,
        shape: &[i32],
        dtype: mlx::mlx_dtype,
    ) -> Self {
        // ORT owns the buffer; MLX must never free it. A no-op dtor makes the wrap purely borrowing.
        unsafe extern "C" fn noop_dtor(_: *mut std::os::raw::c_void) {}
        let arr = Array {
            raw: unsafe {
                mlx::mlx_array_new_data_managed(
                    data as *mut std::os::raw::c_void,
                    shape.as_ptr(),
                    shape.len() as i32,
                    dtype,
                    Some(noop_dtor),
                )
            },
        };
        // Memory view: the boundary zero-copy managed-wrap. A 16 KB page-aligned buffer takes MLX's
        // true `newBufferWithBytesNoCopy` no-copy path; an unaligned one silently falls back to an
        // internal allocate+copy — record which, plus the bytes borrowed. Gated (one atomic load off).
        let tr = crate::trace::tracer();
        if tr.active() {
            let aligned = (data as usize).is_multiple_of(16384);
            tr.record_managed_wrap((arr.size() * arr.itemsize()) as u64, aligned);
        }
        arr
    }

    /// The raw handle, for passing to `mlx_*` calls. Ownership is NOT transferred.
    #[inline]
    pub fn as_raw(&self) -> mlx::mlx_array {
        self.raw
    }

    pub fn ndim(&self) -> usize {
        unsafe { mlx::mlx_array_ndim(self.raw) }
    }

    pub fn shape(&self) -> Vec<i64> {
        let nd = self.ndim();
        let sh = unsafe { mlx::mlx_array_shape(self.raw) };
        (0..nd).map(|i| unsafe { *sh.add(i) } as i64).collect()
    }

    #[allow(dead_code)]
    pub fn size(&self) -> usize {
        unsafe { mlx::mlx_array_size(self.raw) }
    }

    pub fn itemsize(&self) -> usize {
        unsafe { mlx::mlx_array_itemsize(self.raw) }
    }

    #[allow(dead_code)]
    pub fn dtype(&self) -> mlx::mlx_dtype {
        unsafe { mlx::mlx_array_dtype(self.raw) }
    }

    /// Force evaluation of this (single) array.
    #[allow(dead_code)]
    pub fn eval(&self) {
        unsafe { mlx::mlx_array_eval(self.raw) };
    }

    /// Raw byte pointer to the (evaluated) contiguous buffer, for the unified-memory copy-out.
    pub fn data_bytes(&self) -> *const u8 {
        unsafe { mlx::mlx_array_data_uint8(self.raw) }
    }
}

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

impl Drop for Array {
    fn drop(&mut self) {
        unsafe { mlx::mlx_array_free(self.raw) };
    }
}

/// Owning wrapper over an `mlx_vector_array` (the input list passed to a single `mlx_eval`).
pub struct VectorArray {
    raw: mlx::mlx_vector_array,
}

impl VectorArray {
    pub fn new() -> Self {
        VectorArray {
            raw: unsafe { mlx::mlx_vector_array_new() },
        }
    }

    /// Take ownership of a raw `mlx_vector_array` handle (e.g. a `mlx_split` out-param).
    #[inline]
    pub fn from_raw(raw: mlx::mlx_vector_array) -> Self {
        VectorArray { raw }
    }

    /// Append a borrowed array handle (the vector takes its own reference).
    pub fn append(&mut self, a: mlx::mlx_array) {
        unsafe { mlx::mlx_vector_array_append_value(self.raw, a) };
    }

    /// Number of arrays held.
    pub fn size(&self) -> usize {
        unsafe { mlx::mlx_vector_array_size(self.raw) }
    }

    /// A fresh owning reference to element `i` (the vector keeps its own; the returned `Array` owns
    /// the new reference and frees it on drop).
    pub fn get(&self, i: usize) -> Array {
        let mut a = unsafe { mlx::mlx_array_new() };
        unsafe { mlx::mlx_vector_array_get(&mut a, self.raw, i) };
        Array::from_raw(a)
    }

    #[inline]
    pub fn as_raw(&self) -> mlx::mlx_vector_array {
        self.raw
    }

    /// Consume the wrapper WITHOUT freeing, returning the raw handle (ownership transferred to the
    /// caller — e.g. handing a trace result to mlx via the closure's `out` param).
    #[inline]
    pub fn into_raw(self) -> mlx::mlx_vector_array {
        let raw = self.raw;
        std::mem::forget(self);
        raw
    }

    #[inline]
    pub fn as_mut_ptr(&mut self) -> *mut mlx::mlx_vector_array {
        &mut self.raw
    }
}

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

impl Drop for VectorArray {
    fn drop(&mut self) {
        unsafe { mlx::mlx_vector_array_free(self.raw) };
    }
}

/// Evaluate the whole boundary graph in one shot (mirrors the C++ single-`mlx_eval` boundary).
pub fn eval(outputs: &VectorArray) -> Result<(), String> {
    let rc = unsafe { mlx::mlx_eval(outputs.as_raw()) };
    if rc != 0 {
        return Err("mlx_eval failed".to_string());
    }
    Ok(())
}

/// Owning wrapper over an `mlx_closure` (a captured/compiled callable), freed once on drop.
///
/// Two flavours are used by the compiled-decode fast path:
///   * [`Closure::new_func_payload`] wraps a Rust `extern "C"` trace thunk plus an opaque payload
///     pointer — the *base* (un-compiled) closure whose body traces the whole decode subgraph.
///   * [`Closure::compile`] runs `mlx_compile` (shapeless) on a base closure and returns the
///     *compiled* closure that fuses the traced graph into far fewer kernel launches.
///     [`Closure::apply`] runs the closure over an input vector, returning the output arrays.
pub struct Closure {
    raw: mlx::mlx_closure,
}

impl Closure {
    /// Wrap a trace thunk + opaque payload as a base closure. The payload pointer must stay valid
    /// (and point at a stable allocation) for as long as this closure — and any closure compiled
    /// from it — may be applied. No destructor is registered (`dtor = None`); the payload is owned
    /// elsewhere (the plan).
    pub fn new_func_payload(
        fun: unsafe extern "C" fn(
            *mut mlx::mlx_vector_array,
            mlx::mlx_vector_array,
            *mut std::os::raw::c_void,
        ) -> std::os::raw::c_int,
        payload: *mut std::os::raw::c_void,
    ) -> Self {
        let raw = unsafe { mlx::mlx_closure_new_func_payload(Some(fun), payload, None) };
        Closure { raw }
    }

    /// Compile `base` shapeless (so a growing KV length never triggers a recompile) into a fused
    /// closure. Returns `Err` if `mlx_compile` fails (caller falls back to the eager path).
    pub fn compile(base: &Closure, shapeless: bool) -> Result<Closure, String> {
        let mut res = unsafe { mlx::mlx_closure_new() };
        let rc = unsafe { mlx::mlx_compile(&mut res, base.raw, shapeless) };
        if rc != 0 {
            unsafe { mlx::mlx_closure_free(res) };
            return Err("mlx_compile failed".to_string());
        }
        Ok(Closure { raw: res })
    }

    /// Apply the closure to `input`, returning the produced output arrays (owning). `Err` on any
    /// MLX failure inside the (traced or replayed) body.
    pub fn apply(&self, input: &VectorArray) -> Result<VectorArray, String> {
        let mut res = unsafe { mlx::mlx_vector_array_new() };
        let rc = unsafe { mlx::mlx_closure_apply(&mut res, self.raw, input.as_raw()) };
        if rc != 0 {
            unsafe { mlx::mlx_vector_array_free(res) };
            return Err("mlx_closure_apply failed".to_string());
        }
        Ok(VectorArray::from_raw(res))
    }
}

impl Drop for Closure {
    fn drop(&mut self) {
        unsafe { mlx::mlx_closure_free(self.raw) };
    }
}

#[cfg(test)]
mod float64_primitive_tests {
    use super::*;
    use crate::sys::mlx as sys;

    fn scalar_f64(v: f64) -> Array {
        let shape: [i32; 1] = [1];
        Array::from_data(
            &v as *const f64 as *const std::ffi::c_void,
            &shape,
            sys::mlx_dtype__MLX_FLOAT64,
        )
    }

    fn eval_unary(
        op: unsafe extern "C" fn(*mut sys::mlx_array, sys::mlx_array, sys::mlx_stream) -> i32,
        x: f64,
        stream: &Stream,
    ) -> f64 {
        let a = scalar_f64(x);
        let mut raw = unsafe { sys::mlx_array_new() };
        let rc = unsafe { op(&mut raw, a.as_raw(), stream.as_raw()) };
        assert_eq!(rc, 0, "mlx unary op failed on a float64 CPU-stream array");
        let out = Array::from_raw(raw);
        unsafe { sys::mlx_array_eval(out.as_raw()) };
        assert_eq!(
            out.dtype(),
            sys::mlx_dtype__MLX_FLOAT64,
            "float64 input must produce a float64 result"
        );
        unsafe { *sys::mlx_array_data_float64(out.as_raw()) }
    }

    fn eval_softmax<const N: usize>(x: &[f64; N], stream: &Stream) -> Vec<f64> {
        let shape: [i32; 1] = [N as i32];
        let a = Array::from_data(
            x.as_ptr() as *const std::ffi::c_void,
            &shape,
            sys::mlx_dtype__MLX_FLOAT64,
        );
        let mut raw = unsafe { sys::mlx_array_new() };
        let rc = unsafe { sys::mlx_softmax_axis(&mut raw, a.as_raw(), 0, false, stream.as_raw()) };
        assert_eq!(
            rc, 0,
            "mlx_softmax_axis failed on a float64 CPU-stream array"
        );
        let out = Array::from_raw(raw);
        unsafe { sys::mlx_array_eval(out.as_raw()) };
        unsafe { std::slice::from_raw_parts(sys::mlx_array_data_float64(out.as_raw()), N).to_vec() }
    }

    fn eval_logsumexp<const N: usize>(x: &[f64; N], stream: &Stream) -> f64 {
        let shape: [i32; 1] = [N as i32];
        let a = Array::from_data(
            x.as_ptr() as *const std::ffi::c_void,
            &shape,
            sys::mlx_dtype__MLX_FLOAT64,
        );
        let mut raw = unsafe { sys::mlx_array_new() };
        let rc = unsafe { sys::mlx_logsumexp(&mut raw, a.as_raw(), false, stream.as_raw()) };
        assert_eq!(rc, 0, "mlx_logsumexp failed on a float64 CPU-stream array");
        let out = Array::from_raw(raw);
        unsafe { sys::mlx_array_eval(out.as_raw()) };
        unsafe { *sys::mlx_array_data_float64(out.as_raw()) }
    }

    fn relative_error(got: f64, want: f64) -> f64 {
        if want == 0.0 {
            got.abs()
        } else {
            ((got - want) / want).abs()
        }
    }

    /// Pins the measured float64 behaviour of the MLX primitives the fp64 opt-in relies on.
    ///
    /// MLX keeps the `float64` dtype but computes some primitives in float32 and widens back, so a
    /// result can *look* double-precision and carry only ~7 significant digits. The op claim table
    /// (`is_mlx_cpu_float`) is derived from exactly this split, so if an MLX upgrade moves a
    /// primitive between the two columns this test fails and the claim table must be revisited —
    /// rather than the EP silently returning float32-accurate answers for a float64 model.
    #[test]
    fn mlx_float64_primitives() {
        let cpu = Stream::new_default_cpu();
        let compound_input: [f64; 16] = [
            0.1257302210933933,
            -0.1321048632913019,
            0.6404226504432821,
            0.10490011715303971,
            -0.535669373161111,
            0.36159505490948474,
            1.3040000451301372,
            0.9470809631292422,
            -0.7037352358069926,
            -1.2654214710460525,
            -0.6232744625373522,
            0.0413259793472436,
            -2.3250307746388343,
            -0.21879166393254573,
            -1.2459109472530652,
            -0.7322673547034516,
        ];
        let compound_sum = compound_input.iter().map(|x| x.exp()).sum::<f64>();

        // Exact in float64: safe for ops to opt in.
        for (name, got, want) in [
            ("log", eval_unary(sys::mlx_log, 0.7, &cpu), 0.7f64.ln()),
            ("sqrt", eval_unary(sys::mlx_sqrt, 2.0, &cpu), 2.0f64.sqrt()),
            ("tanh", eval_unary(sys::mlx_tanh, 0.7, &cpu), 0.7f64.tanh()),
            (
                "expm1",
                eval_unary(sys::mlx_expm1, -1e-9, &cpu),
                (-1e-9f64).exp_m1(),
            ),
            (
                "reciprocal",
                eval_unary(sys::mlx_reciprocal, 3.0, &cpu),
                1.0 / 3.0,
            ),
        ] {
            assert!(
                relative_error(got, want) <= 1e-15,
                "mlx_{name} was float64-exact when the claim table was written but now differs \
                 (got {got:.17e}, want {want:.17e}); re-check which ops may claim float64"
            );
        }

        // Silently float32-accurate: ops built on these must NOT claim float64.
        for (name, got, want) in [
            ("exp", eval_unary(sys::mlx_exp, -3.5, &cpu), (-3.5f64).exp()),
            ("sin", eval_unary(sys::mlx_sin, -3.5, &cpu), (-3.5f64).sin()),
            ("cos", eval_unary(sys::mlx_cos, -3.5, &cpu), (-3.5f64).cos()),
            (
                "sigmoid",
                eval_unary(sys::mlx_sigmoid, -3.5, &cpu),
                1.0 / (1.0 + 3.5f64.exp()),
            ),
        ] {
            let rel = relative_error(got, want);
            assert!(
                rel > 1e-15,
                "mlx_{name} is now float64-exact (got {got:.17e}, want {want:.17e}) — the ops \
                 built on it may finally claim float64; update the claim table"
            );
            assert!(
                rel <= 1e-6,
                "mlx_{name} is neither float64-exact nor float32-accurate (rel={rel:.3e}) — \
                 something is badly wrong with the fp64 CPU path"
            );
        }

        let softmax = eval_softmax(&compound_input, &cpu);
        let softmax_rel = softmax
            .iter()
            .zip(compound_input)
            .map(|(&got, x)| relative_error(got, x.exp() / compound_sum))
            .fold(0.0, f64::max);
        let logsumexp_rel =
            relative_error(eval_logsumexp(&compound_input, &cpu), compound_sum.ln());
        for (name, rel) in [("softmax", softmax_rel), ("logsumexp", logsumexp_rel)] {
            assert!(
                rel > 1e-15,
                "mlx_{name} is now float64-exact (max rel={rel:.3e}) — the ops built on it may \
                  finally claim float64; update the claim table"
            );
            assert!(
                rel <= 1e-6,
                "mlx_{name} is neither float64-exact nor float32-accurate (max rel={rel:.3e}) — \
                  something is badly wrong with the fp64 CPU path"
            );
        }
    }
}