onnxruntime-ep-mlx 0.27.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() },
        }
    }

    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) };
    }
}

/// Owning wrapper over an `mlx_vector_string` — the `input_names`/`output_names` lists a
/// [`FastMetalKernel`] is built from. Freed once on drop.
pub struct VectorString {
    raw: mlx::mlx_vector_string,
}

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

    /// Append a name (MLX copies the bytes; `s` need not outlive the call).
    pub fn append(&mut self, s: &std::ffi::CStr) {
        unsafe { mlx::mlx_vector_string_append_value(self.raw, s.as_ptr()) };
    }

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

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

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

/// Owning wrapper over an `mlx_fast_metal_kernel_config` — the per-call dispatch/output-shape
/// description passed to [`FastMetalKernel::apply`]. Cheap to build (no compilation happens here);
/// unlike the kernel object itself this is NOT meant to be cached, since grid/threadgroup/output
/// shape vary per call (per M/N).
pub struct FastMetalKernelConfig {
    raw: mlx::mlx_fast_metal_kernel_config,
}

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

    /// Declare one output array's shape + dtype (call once per name in the kernel's `output_names`,
    /// in order).
    pub fn add_output_arg(&mut self, shape: &[i32], dtype: mlx::mlx_dtype) -> Result<(), String> {
        let rc = unsafe {
            mlx::mlx_fast_metal_kernel_config_add_output_arg(
                self.raw,
                shape.as_ptr(),
                shape.len(),
                dtype,
            )
        };
        if rc != 0 {
            return Err("mlx_fast_metal_kernel_config_add_output_arg failed".to_string());
        }
        Ok(())
    }

    /// Total thread count per grid dimension (Metal's `dispatchThreads:`, NOT threadgroup count).
    pub fn set_grid(&mut self, x: i32, y: i32, z: i32) -> Result<(), String> {
        let rc = unsafe { mlx::mlx_fast_metal_kernel_config_set_grid(self.raw, x, y, z) };
        if rc != 0 {
            return Err("mlx_fast_metal_kernel_config_set_grid failed".to_string());
        }
        Ok(())
    }

    /// Threads per threadgroup (Metal's `threadsPerThreadgroup:`).
    pub fn set_thread_group(&mut self, x: i32, y: i32, z: i32) -> Result<(), String> {
        let rc = unsafe { mlx::mlx_fast_metal_kernel_config_set_thread_group(self.raw, x, y, z) };
        if rc != 0 {
            return Err("mlx_fast_metal_kernel_config_set_thread_group failed".to_string());
        }
        Ok(())
    }

    /// Print the generated Metal source (once, at first `apply`) to stderr — debug-only.
    #[allow(dead_code)]
    pub fn set_verbose(&mut self, verbose: bool) -> Result<(), String> {
        let rc = unsafe { mlx::mlx_fast_metal_kernel_config_set_verbose(self.raw, verbose) };
        if rc != 0 {
            return Err("mlx_fast_metal_kernel_config_set_verbose failed".to_string());
        }
        Ok(())
    }

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

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

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

/// Owning wrapper over a compiled `mlx_fast_metal_kernel` (an `mx.fast.metal_kernel`-equivalent
/// custom kernel object). Building one is cheap (compilation is deferred to the first `apply`), but
/// once built it should be **cached and reused** across calls with the same source — recompiling per
/// node/call would repeatedly hit the Metal shader compiler.
///
/// `Send + Sync`: the handle is an MLX-internal reference-counted context (mirroring `Array`'s own
/// safety story); MLX's C++ core reference-counts these with atomics, so sharing one behind a
/// `OnceLock`/`Mutex` across the plugin's (effectively single-threaded-per-call) usage is sound.
pub struct FastMetalKernel {
    raw: mlx::mlx_fast_metal_kernel,
}

unsafe impl Send for FastMetalKernel {}
unsafe impl Sync for FastMetalKernel {}

impl FastMetalKernel {
    /// Build (but do not yet compile) a named custom Metal kernel. `source` is ONLY the kernel body
    /// (MLX auto-generates the `[[kernel]] void ...(...)` signature from `input_names`/
    /// `output_names` plus whichever `<name>_shape`/`_strides`/`_ndim` identifiers and Metal
    /// attribute names are textually referenced in `source`); `header` is prepended verbatim before
    /// the generated signature (includes/helper functions). Any Metal compile error surfaces as an
    /// `Err` from the FIRST [`Self::apply`] call, not here.
    #[allow(clippy::too_many_arguments)]
    pub fn new(
        name: &std::ffi::CStr,
        input_names: &VectorString,
        output_names: &VectorString,
        source: &std::ffi::CStr,
        header: &std::ffi::CStr,
        ensure_row_contiguous: bool,
        atomic_outputs: bool,
    ) -> Self {
        let raw = unsafe {
            mlx::mlx_fast_metal_kernel_new(
                name.as_ptr(),
                input_names.as_raw(),
                output_names.as_raw(),
                source.as_ptr(),
                header.as_ptr(),
                ensure_row_contiguous,
                atomic_outputs,
            )
        };
        FastMetalKernel { raw }
    }

    /// Dispatch the kernel over `inputs` (in the same order as the `input_names` it was built
    /// with), producing one output array per `add_output_arg` call on `config`. `Err` on any MLX
    /// failure (a bad Metal source, an eligibility/shape mismatch caught by `ensure_row_contiguous`,
    /// etc.) — callers must treat this as "fall back to the slow path", never as a hard error.
    pub fn apply(
        &self,
        inputs: &VectorArray,
        config: &FastMetalKernelConfig,
        stream: mlx::mlx_stream,
    ) -> Result<VectorArray, String> {
        let mut res = unsafe { mlx::mlx_vector_array_new() };
        let rc = unsafe {
            mlx::mlx_fast_metal_kernel_apply(
                &mut res,
                self.raw,
                inputs.as_raw(),
                config.as_raw(),
                stream,
            )
        };
        if rc != 0 {
            unsafe { mlx::mlx_vector_array_free(res) };
            return Err("mlx_fast_metal_kernel_apply failed".to_string());
        }
        Ok(VectorArray::from_raw(res))
    }
}

impl Drop for FastMetalKernel {
    fn drop(&mut self) {
        unsafe { mlx::mlx_fast_metal_kernel_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) };
    }
}