rlx-metal 0.2.13

// RLX — versatile ML compiler + runtime.
// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, version 3.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program. If not, see <https://www.gnu.org/licenses/>.

//! GPU arena allocator — single Metal buffer with sub-region offsets.
//!
//! Mirrors rlx-cpu's arena: one big allocation, all intermediate tensors
//! get byte offsets within it. Apple Silicon unified memory means the same
//! buffer is accessible from both CPU and GPU with zero copy.

use crate::device::metal_device;
use metal::Buffer;
use rlx_ir::{DType, Graph, NodeId};
use rlx_opt::memory::MemoryPlan;
use std::collections::HashMap;

pub struct Arena {
    pub buffer: Buffer,
    pub size_bytes: usize,
    pub offsets: HashMap<NodeId, usize>, // byte offsets per node
    pub element_counts: HashMap<NodeId, usize>, // element counts per node
    pub dtypes: HashMap<NodeId, DType>,  // per-node dtype (for f16 vs f32 dispatch)
}

impl Arena {
    pub fn from_plan(plan: MemoryPlan) -> Self {
        Self::from_plan_with_graph(plan, None)
    }

    /// Build arena from memory plan, recording per-node dtype from the graph.
    /// If `graph` is None, all buffers are assumed F32.
    pub fn from_plan_with_graph(plan: MemoryPlan, graph: Option<&Graph>) -> Self {
        let dev = metal_device().expect("Metal device required for rlx-metal arena");
        let buffer = dev.alloc_shared(plan.arena_size.max(64));

        let mut offsets = HashMap::with_capacity(plan.assignments.len());
        let mut element_counts = HashMap::with_capacity(plan.assignments.len());
        let mut dtypes = HashMap::with_capacity(plan.assignments.len());
        for (node_id, slot) in &plan.assignments {
            offsets.insert(*node_id, slot.offset);
            // Element count derived from byte size and dtype
            let dt = graph
                .map(|g| g.node(*node_id).shape.dtype())
                .unwrap_or(DType::F32);
            let elem_size = dt.size_bytes();
            element_counts.insert(*node_id, slot.size / elem_size.max(1));
            dtypes.insert(*node_id, dt);
        }
        Self {
            buffer,
            size_bytes: plan.arena_size,
            offsets,
            element_counts,
            dtypes,
        }
    }

    pub fn has_buffer(&self, id: NodeId) -> bool {
        self.offsets.contains_key(&id)
    }

    pub fn byte_offset(&self, id: NodeId) -> usize {
        *self.offsets.get(&id).expect("node not in arena")
    }

    pub fn dtype(&self, id: NodeId) -> DType {
        self.dtypes.get(&id).copied().unwrap_or(DType::F32)
    }

    /// Get a CPU-side mutable slice for the node's region as f32. Only valid
    /// when the node's dtype is F32 (debug-asserted).
    pub fn slice_mut(&mut self, id: NodeId) -> &mut [f32] {
        debug_assert_eq!(self.dtype(id), DType::F32);
        let off = self.byte_offset(id);
        let len = *self.element_counts.get(&id).unwrap_or(&0);
        unsafe {
            let ptr = self.buffer.contents() as *mut u8;
            std::slice::from_raw_parts_mut(ptr.add(off) as *mut f32, len)
        }
    }

    pub fn slice(&self, id: NodeId) -> &[f32] {
        debug_assert_eq!(self.dtype(id), DType::F32);
        let off = self.byte_offset(id);
        let len = *self.element_counts.get(&id).unwrap_or(&0);
        unsafe {
            let ptr = self.buffer.contents() as *const u8;
            std::slice::from_raw_parts(ptr.add(off) as *const f32, len)
        }
    }

    /// Read the node's data as f32 regardless of native precision (converts
    /// f16 → f32 on the fly). Used at graph output boundary.
    pub fn read_as_f32(&self, id: NodeId) -> Vec<f32> {
        let dt = self.dtype(id);
        let off = self.byte_offset(id);
        let len = *self.element_counts.get(&id).unwrap_or(&0);
        unsafe {
            let base = (self.buffer.contents() as *const u8).add(off);
            match dt {
                DType::F32 => std::slice::from_raw_parts(base as *const f32, len).to_vec(),
                DType::F16 => {
                    let src = std::slice::from_raw_parts(base as *const half::f16, len);
                    src.iter().map(|h| h.to_f32()).collect()
                }
                _ => std::slice::from_raw_parts(base as *const f32, len).to_vec(),
            }
        }
    }

    /// Write f32 data, converting to the node's native dtype.
    /// Used at graph input/param boundary.
    pub fn write_from_f32(&mut self, id: NodeId, data: &[f32]) {
        let dt = self.dtype(id);
        let off = self.byte_offset(id);
        let cap = *self.element_counts.get(&id).unwrap_or(&0);
        let len = data.len().min(cap);
        unsafe {
            let base = (self.buffer.contents() as *mut u8).add(off);
            match dt {
                DType::F32 => {
                    std::ptr::copy_nonoverlapping(data.as_ptr(), base as *mut f32, len);
                }
                DType::F16 => {
                    let dst = std::slice::from_raw_parts_mut(base as *mut half::f16, len);
                    for (i, &v) in data.iter().take(len).enumerate() {
                        dst[i] = half::f16::from_f32(v);
                    }
                }
                DType::BF16 => {
                    let dst = std::slice::from_raw_parts_mut(base as *mut half::bf16, len);
                    for (i, &v) in data.iter().take(len).enumerate() {
                        dst[i] = half::bf16::from_f32(v);
                    }
                }
                // Integer-typed inputs (token IDs, position indices) get
                // cast from f32 → int. The previous fallthrough memcpy
                // bit-pattern-reinterpreted the floats as ints, which
                // produced stable garbled-token streams from gather/take.
                DType::I32 => {
                    let dst = std::slice::from_raw_parts_mut(base as *mut i32, len);
                    for (i, &v) in data.iter().take(len).enumerate() {
                        dst[i] = v as i32;
                    }
                }
                DType::I64 => {
                    let dst = std::slice::from_raw_parts_mut(base as *mut i64, len);
                    for (i, &v) in data.iter().take(len).enumerate() {
                        dst[i] = v as i64;
                    }
                }
                DType::U32 => {
                    let dst = std::slice::from_raw_parts_mut(base as *mut u32, len);
                    for (i, &v) in data.iter().take(len).enumerate() {
                        dst[i] = v as u32;
                    }
                }
                DType::I16 => {
                    let dst = std::slice::from_raw_parts_mut(base as *mut i16, len);
                    for (i, &v) in data.iter().take(len).enumerate() {
                        dst[i] = v as i16;
                    }
                }
                DType::I8 => {
                    let dst = std::slice::from_raw_parts_mut(base as *mut i8, len);
                    for (i, &v) in data.iter().take(len).enumerate() {
                        dst[i] = v as i8;
                    }
                }
                DType::U8 => {
                    let dst = std::slice::from_raw_parts_mut(base, len);
                    for (i, &v) in data.iter().take(len).enumerate() {
                        dst[i] = v as u8;
                    }
                }
                DType::Bool => {
                    let dst = std::slice::from_raw_parts_mut(base, len);
                    for (i, &v) in data.iter().take(len).enumerate() {
                        dst[i] = if v != 0.0 { 1 } else { 0 };
                    }
                }
                // F64 / Complex64 don't appear in user input feeds today.
                _ => {
                    std::ptr::copy_nonoverlapping(data.as_ptr(), base as *mut f32, len);
                }
            }
        }
    }

    /// Copy one arena node's f32 payload into another (unified-memory memcpy).
    pub fn copy_node_f32(&self, dst: NodeId, src: NodeId) {
        let dst_len = *self.element_counts.get(&dst).unwrap_or(&0);
        let src_len = *self.element_counts.get(&src).unwrap_or(&0);
        self.copy_node_f32_prefix(dst, src, dst_len.min(src_len));
    }

    /// Copy `n` f32 from `src` (starting at element `src_elem`) into `dst`
    /// (starting at `dst_elem`), clamped to both node element counts. Used by the
    /// resident KV *row* feed to drop one new-token row (output row `upper`) into
    /// the resident `past_k_*` slot at the active row, in unified memory.
    pub fn copy_node_f32_range(
        &self,
        dst: NodeId,
        dst_elem: usize,
        src: NodeId,
        src_elem: usize,
        n: usize,
    ) {
        let dst_cap = *self.element_counts.get(&dst).unwrap_or(&0);
        let src_cap = *self.element_counts.get(&src).unwrap_or(&0);
        if n == 0 || dst_elem + n > dst_cap || src_elem + n > src_cap {
            return;
        }
        let dst_off = self.byte_offset(dst);
        let src_off = self.byte_offset(src);
        unsafe {
            let base = self.buffer.contents() as *mut u8;
            let src_p = (base.add(src_off) as *const f32).add(src_elem);
            let dst_p = (base.add(dst_off) as *mut f32).add(dst_elem);
            if src_p as *const () != dst_p as *const () {
                std::ptr::copy(src_p, dst_p, n);
            }
        }
    }

    /// Copy the first `elems` floats from `src` into `dst` (KV prefix after active-extent).
    pub fn copy_node_f32_prefix(&self, dst: NodeId, src: NodeId, elems: usize) {
        if elems == 0 {
            return;
        }
        let dst_off = self.byte_offset(dst);
        let src_off = self.byte_offset(src);
        let dst_cap = *self.element_counts.get(&dst).unwrap_or(&0);
        let src_cap = *self.element_counts.get(&src).unwrap_or(&0);
        let len = elems.min(dst_cap).min(src_cap);
        if len == 0 {
            return;
        }
        unsafe {
            let base = self.buffer.contents() as *mut u8;
            std::ptr::copy(
                base.add(src_off) as *const f32,
                base.add(dst_off) as *mut f32,
                len,
            );
        }
    }

    /// Copy raw bytes into the node's arena slot (U8/I8 packed weights).
    pub fn write_bytes(&mut self, id: NodeId, data: &[u8]) {
        let off = self.byte_offset(id);
        let cap = *self.element_counts.get(&id).unwrap_or(&0);
        let len = data.len().min(cap);
        unsafe {
            let base = (self.buffer.contents() as *mut u8).add(off);
            std::ptr::copy_nonoverlapping(data.as_ptr(), base, len);
        }
    }

    /// Copy a node's byte payload from another arena (packed param sharing).
    pub fn copy_node_bytes_from(&self, dst: NodeId, src_arena: &Arena, src: NodeId) {
        let dst_off = self.byte_offset(dst);
        let src_off = src_arena.byte_offset(src);
        let dst_cap = *self.element_counts.get(&dst).unwrap_or(&0);
        let src_cap = src_arena.element_counts.get(&src).copied().unwrap_or(0);
        let elems = dst_cap.min(src_cap);
        if elems == 0 {
            return;
        }
        // `element_counts` are ELEMENT counts, so scale by the dtype width to get
        // bytes — otherwise an F32 param copies only 1 of every 4 bytes (a scalar
        // scale becomes a denormal ≈0 → div-by-zero → NaN on a reused/cloned graph).
        let elem_size = self.dtype(dst).size_bytes().max(1);
        let bytes = elems * elem_size;
        unsafe {
            let dst_base = self.buffer.contents() as *mut u8;
            let src_base = src_arena.buffer.contents() as *const u8;
            std::ptr::copy(src_base.add(src_off), dst_base.add(dst_off), bytes);
        }
    }
}