rlx-metal 0.2.14

// RLX — versatile ML compiler + runtime.
// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
// SPDX-License-Identifier: MIT OR Apache-2.0

//! MPSMatrixMultiplication bridge — Apple's per-chip-tuned matmul.
//!
//! For large matmuls (M·K·N above a threshold) Apple's MPS sgemm routinely
//! beats hand-rolled MSL because it has private knowledge of the GPU's
//! tensor-unit scheduling per chip generation. We bridge it via objc.
//!
//! Trade-off: per-call objc bridging is ~5–20µs, so MPS only wins above a
//! threshold (rough rule of thumb: M·K·N ≥ 16M FLOPs). The cost model in
//! `cost.rs` decides.
//!
//! Note: `MPSMatrixMultiplication::encode` allocates and submits its own
//! compute encoder internally — callers must end any open compute encoder
//! on the same command buffer before invoking us. The shared-encoder split
//! is wired in `backend::encode_and_run`.

use crate::mtl::{Buffer, CommandBufferRef};
use objc::runtime::{BOOL, NO, Object, YES};
use objc::{class, msg_send, sel, sel_impl};
use std::collections::HashMap;
use std::sync::{Mutex, OnceLock, RwLock};

/// Lifecycle guard over the process-global MPSMatrix / kernel caches.
///
/// The caches hand out retained Objective-C pointers that an *encode* uses
/// while building a matmul, but `invalidate_caches()` (run on every compile
/// and on `MetalExecutable::drop`) releases all of them. With two
/// `MetalExecutable`s live on separate threads, one's invalidate would free
/// the `MPSMatrix` the other is mid-encode with → `objc_msgSend` on a freed
/// object → SIGSEGV.
///
/// An `RwLock` resolves it without serializing the hot path: every MPS encode
/// holds a *read* lock (concurrent encodes still overlap), while
/// `invalidate_caches` takes the *write* lock so it can only run once no
/// encode is in flight.
static CACHE_GUARD: RwLock<()> = RwLock::new(());

// Link the MetalPerformanceShaders framework. metal-rs gates its own
// MPS link directive behind a feature we don't enable.
#[link(name = "MetalPerformanceShaders", kind = "framework")]
unsafe extern "C" {}

/// Wrap an in-flight `MTLCommandBuffer` for MPSGraph / MPS kernel encode APIs.
pub fn mps_command_buffer_wrap(cmd_buf: &CommandBufferRef) -> *mut Object {
    unsafe {
        let cls = class!(MPSCommandBuffer);
        msg_send![cls, commandBufferWithCommandBuffer: cmd_buf]
    }
}

/// MPSDataType values (from MPSCore.h).
#[allow(non_upper_case_globals, dead_code)]
mod mps_dtype {
    pub const Float32: u32 = 0x10000000 | 32;
    pub const Float16: u32 = 0x10000000 | 16;
}

/// True iff MPSMatrixMultiplication is available (any modern macOS Metal device).
pub fn mps_supports_matmul() -> bool {
    static AVAIL: OnceLock<bool> = OnceLock::new();
    *AVAIL.get_or_init(|| objc::runtime::Class::get("MPSMatrixMultiplication").is_some())
}

/// Cache of `(m,k,n)` → retained MPSMatrixMultiplication kernel.
///
/// **Bridge-cost mitigation #1.** Building the kernel involves ~6 objc
/// messages (alloc, init, set transposes, set α/β). For typical inference
/// the same shapes recur every layer, so caching reduces per-call objc
/// overhead.
struct KernelCache {
    map: Mutex<HashMap<(usize, usize, usize, bool, bool), usize>>,
}
unsafe impl Send for KernelCache {}
unsafe impl Sync for KernelCache {}

fn kernel_cache() -> &'static KernelCache {
    static CACHE: OnceLock<KernelCache> = OnceLock::new();
    CACHE.get_or_init(|| KernelCache {
        map: Mutex::new(HashMap::new()),
    })
}

/// Cache of `(buf_ptr, offset, rows, cols)` → retained MPSMatrix wrapper.
///
/// **Bridge-cost mitigation #2.** Each MPSMatrix wraps an MTLBuffer + offset
/// + descriptor. Within a compiled graph the (buffer, offset, dims) triple
/// is fully static — it never changes call-to-call. Caching them eliminates
/// 9 objc messages per matmul (3 descriptor alloc + 3 matrix alloc/init +
/// 3 matrix release).
///
/// Key uses `(buf_ptr as usize, offset, rows, cols)`; descriptor is also
/// cached separately keyed on `(rows, cols)`.
struct MatrixCache {
    matrices: Mutex<HashMap<(usize, usize, usize, usize), usize>>,
    descriptors: Mutex<HashMap<(usize, usize), usize>>,
}
unsafe impl Send for MatrixCache {}
unsafe impl Sync for MatrixCache {}

fn matrix_cache() -> &'static MatrixCache {
    static CACHE: OnceLock<MatrixCache> = OnceLock::new();
    CACHE.get_or_init(|| MatrixCache {
        matrices: Mutex::new(HashMap::new()),
        descriptors: Mutex::new(HashMap::new()),
    })
}

unsafe fn get_or_build_descriptor(rows: usize, cols: usize, dtype: u32) -> *mut Object {
    let cache = matrix_cache();
    let mut map = cache.descriptors.lock().expect("descriptor cache poisoned");
    let key = (rows, cols * 8 + dtype as usize); // pack dtype into key
    if let Some(&p) = map.get(&key) {
        return p as *mut Object;
    }
    let cls = class!(MPSMatrixDescriptor);
    let bytes_per_elem = if dtype == mps_dtype::Float16 { 2 } else { 4 };
    let row_bytes = cols * bytes_per_elem;
    let desc: *mut Object = msg_send![cls,
        matrixDescriptorWithRows: rows as u64
        columns: cols as u64
        rowBytes: row_bytes as u64
        dataType: dtype];
    let _: () = msg_send![desc, retain];
    map.insert(key, desc as usize);
    desc
}

unsafe fn get_or_build_matrix(
    buf: &Buffer,
    offset: usize,
    rows: usize,
    cols: usize,
    dtype: u32,
) -> *mut Object {
    unsafe {
        let cache = matrix_cache();
        // The Buffer-wrapper address can recycle when a previous Sam is
        // dropped — relying on `&**buf as usize` as the identity led to
        // stale `MPSMatrix` lookups → GPU reads from freed memory →
        // NaN. To break the aliasing, callers that build new arenas
        // (e.g. `MetalBackend::compile_inner`) call
        // `invalidate_caches()` first so this map is empty.
        let buf_ptr = (&**buf as *const crate::mtl::BufferRef) as usize;
        let key = (buf_ptr, offset, rows, cols * 8 + dtype as usize);
        let mut map = cache.matrices.lock().expect("matrix cache poisoned");
        if let Some(&p) = map.get(&key) {
            return p as *mut Object;
        }
        let desc = get_or_build_descriptor(rows, cols, dtype);
        let cls = class!(MPSMatrix);
        let alloc: *mut Object = msg_send![cls, alloc];
        let buf_ref: &crate::mtl::BufferRef = buf;
        let mat: *mut Object = msg_send![alloc,
        initWithBuffer: buf_ref
        offset: offset as u64
        descriptor: desc];
        map.insert(key, mat as usize);
        mat
    }
}

/// Drop every cached MPSMatrix / MPSMatrixDescriptor / MPSMatrixMultiplication
/// reference. Lets a caller (e.g. backend test harness, hot reload) reset
/// MPS state explicitly.
pub fn invalidate_caches() {
    // Exclusive: block until no thread is mid-encode, so we never release a
    // cached MPSMatrix / kernel another thread is still using.
    let _guard = CACHE_GUARD.write().expect("MPS cache guard poisoned");
    let cache = matrix_cache();
    {
        let mut mats = cache.matrices.lock().expect("matrix cache poisoned");
        release_cached_ptrs(mats.drain().map(|(_, p)| p));
    }
    {
        let mut descs = cache.descriptors.lock().expect("descriptor cache poisoned");
        release_cached_ptrs(descs.drain().map(|(_, p)| p));
    }
    let kcache = kernel_cache();
    {
        let mut km = kcache.map.lock().expect("kernel cache poisoned");
        release_cached_ptrs(km.drain().map(|(_, p)| p));
    }
}

fn release_cached_ptrs(ptrs: impl Iterator<Item = usize>) {
    unsafe {
        for p in ptrs {
            if p != 0 {
                let obj = p as *mut Object;
                let _: () = msg_send![obj, release];
            }
        }
    }
}

unsafe fn get_or_build_kernel(
    m: usize,
    k: usize,
    n: usize,
    transpose_a: bool,
    transpose_b: bool,
) -> *mut Object {
    let cache = kernel_cache();
    let mut map = cache.map.lock().expect("kernel cache poisoned");
    if let Some(&p) = map.get(&(m, k, n, transpose_a, transpose_b)) {
        return p as *mut Object;
    }
    use crate::device::metal_device;
    let dev = metal_device().expect("Metal device required");
    let cls = class!(MPSMatrixMultiplication);
    let alloc: *mut Object = msg_send![cls, alloc];
    let dev_ref: &crate::mtl::DeviceRef = &dev.device;
    let kernel: *mut Object = msg_send![alloc,
        initWithDevice: dev_ref
        transposeLeft: if transpose_a { YES } else { NO } as BOOL
        transposeRight: if transpose_b { YES } else { NO } as BOOL
        resultRows: m as u64
        resultColumns: n as u64
        interiorColumns: k as u64
        alpha: 1.0_f64
        beta: 0.0_f64
    ];
    map.insert((m, k, n, transpose_a, transpose_b), kernel as usize);
    kernel
}

/// Encode `C = A @ B` via MPSMatrixMultiplication.
///
/// Hot path: cached kernel + cached MPSMatrix wrappers → only one objc
/// message at runtime (`encodeToCommandBuffer`). Everything else amortizes.
pub fn encode_mps_sgemm(
    cmd_buf: &CommandBufferRef,
    arena: &Buffer,
    a_off: usize,
    b_off: usize,
    c_off: usize,
    m: usize,
    k: usize,
    n: usize,
) {
    encode_mps_matmul(
        cmd_buf,
        arena,
        a_off,
        b_off,
        c_off,
        m,
        k,
        n,
        mps_dtype::Float32,
        false,
        false,
    );
}

/// `C = A @ B^T` where `B` is stored as `[n, k]` row-major (GGUF dequant layout).
pub fn encode_mps_sgemm_bt(
    cmd_buf: &CommandBufferRef,
    arena: &Buffer,
    a_off: usize,
    b_off: usize,
    c_off: usize,
    m: usize,
    k: usize,
    n: usize,
) {
    encode_mps_matmul(
        cmd_buf,
        arena,
        a_off,
        b_off,
        c_off,
        m,
        k,
        n,
        mps_dtype::Float32,
        false,
        true,
    );
}

/// Encode `C = A @ B` at half-precision via MPS.
pub fn encode_mps_hgemm(
    cmd_buf: &CommandBufferRef,
    arena: &Buffer,
    a_off: usize,
    b_off: usize,
    c_off: usize,
    m: usize,
    k: usize,
    n: usize,
) {
    encode_mps_matmul(
        cmd_buf,
        arena,
        a_off,
        b_off,
        c_off,
        m,
        k,
        n,
        mps_dtype::Float16,
        false,
        false,
    );
}

#[allow(clippy::too_many_arguments)]
fn encode_mps_matmul(
    cmd_buf: &CommandBufferRef,
    arena: &Buffer,
    a_off: usize,
    b_off: usize,
    c_off: usize,
    m: usize,
    k: usize,
    n: usize,
    dtype: u32,
    transpose_a: bool,
    transpose_b: bool,
) {
    // Shared: pins the cached MPSMatrix / kernel pointers alive for the whole
    // encode so a concurrent `invalidate_caches` cannot free them under us.
    let _guard = CACHE_GUARD.read().expect("MPS cache guard poisoned");
    unsafe {
        // With transposeLeft, MPS reads A stored as [k, m] and treats it as [m, k].
        let (a_rows, a_cols) = if transpose_a { (k, m) } else { (m, k) };
        let a_mat = get_or_build_matrix(arena, a_off, a_rows, a_cols, dtype);
        let (b_rows, b_cols) = if transpose_b { (n, k) } else { (k, n) };
        let b_mat = get_or_build_matrix(arena, b_off, b_rows, b_cols, dtype);
        let c_mat = get_or_build_matrix(arena, c_off, m, n, dtype);
        let kernel = get_or_build_kernel(m, k, n, transpose_a, transpose_b);
        let _: () = msg_send![kernel,
            encodeToCommandBuffer: cmd_buf
            leftMatrix: a_mat
            rightMatrix: b_mat
            resultMatrix: c_mat
        ];
    }
}

/// `C = op(A) · op(B)` where `op` transposes per flag — folds a materialized
/// last-two-swap `Transpose` on either operand into the GEMM (the autodiff VJP
/// emits `dW = Xᵀ·dY` and `dX = dY·Wᵀ`). Shapes: `m`,`k`,`n` describe the
/// logical (post-transpose) `[m,k]·[k,n]=[m,n]`; buffers hold the pre-transpose
/// operands. Arena-only (callers must ensure no operand is weight-buffer-tagged).
#[allow(clippy::too_many_arguments)]
pub fn encode_mps_sgemm_t(
    cmd_buf: &CommandBufferRef,
    arena: &Buffer,
    a_off: usize,
    b_off: usize,
    c_off: usize,
    m: usize,
    k: usize,
    n: usize,
    transpose_a: bool,
    transpose_b: bool,
) {
    encode_mps_matmul(
        cmd_buf,
        arena,
        a_off,
        b_off,
        c_off,
        m,
        k,
        n,
        mps_dtype::Float32,
        transpose_a,
        transpose_b,
    );
}

/// Buffer-aware transpose-folded GEMM: like [`encode_mps_sgemm_t`] but each
/// operand carries its **own** `MTLBuffer` (activation arena OR the weight
/// buffer). This lets a folded `matmul_t(weight)` read the weight directly with
/// `transposeRight` instead of materializing a transposed copy — the arena-only
/// [`encode_mps_sgemm_t`] can't, since it assumes a single buffer.
/// `get_or_build_matrix` already keys its cache by buffer pointer, so mixing an
/// arena operand with a weight-buffer operand is safe. Passing the same arena
/// buffer for all three is byte-identical to [`encode_mps_sgemm_t`].
#[allow(clippy::too_many_arguments)]
pub fn encode_mps_sgemm_t_bufs(
    cmd_buf: &CommandBufferRef,
    a_buf: &Buffer,
    a_off: usize,
    b_buf: &Buffer,
    b_off: usize,
    c_buf: &Buffer,
    c_off: usize,
    m: usize,
    k: usize,
    n: usize,
    transpose_a: bool,
    transpose_b: bool,
) {
    let _guard = CACHE_GUARD.read().expect("MPS cache guard poisoned");
    unsafe {
        let (a_rows, a_cols) = if transpose_a { (k, m) } else { (m, k) };
        let a_mat = get_or_build_matrix(a_buf, a_off, a_rows, a_cols, mps_dtype::Float32);
        let (b_rows, b_cols) = if transpose_b { (n, k) } else { (k, n) };
        let b_mat = get_or_build_matrix(b_buf, b_off, b_rows, b_cols, mps_dtype::Float32);
        let c_mat = get_or_build_matrix(c_buf, c_off, m, n, mps_dtype::Float32);
        let kernel = get_or_build_kernel(m, k, n, transpose_a, transpose_b);
        let _: () = msg_send![kernel,
            encodeToCommandBuffer: cmd_buf
            leftMatrix: a_mat
            rightMatrix: b_mat
            resultMatrix: c_mat
        ];
    }
}