rlx-runtime 0.2.16

// RLX — versatile ML compiler + runtime.
// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
// SPDX-License-Identifier: MIT OR Apache-2.0
use super::*;
use rlx_cpu::{arena::Arena, thunk};
use rlx_ir::{DType, NodeId, Op};
use rlx_opt::memory::{self, MemoryPlan};
// Arena typed read/write helpers live in `crate::arena` so every
// backend (CPU, Metal, future CUDA/wgpu/WASM) shares one implementation.
use rlx_driver::arena::{read_typed_to_f32, write_typed_from_f32};

// `madvise(2)` — used by `release_scratch_pages` to return idle scratch pages to
// the OS without pulling in the whole `libc` crate.
#[cfg(unix)]
unsafe extern "C" {
    fn madvise(
        addr: *mut core::ffi::c_void,
        len: usize,
        advice: core::ffi::c_int,
    ) -> core::ffi::c_int;
}

pub struct CpuBackend;

impl Backend for CpuBackend {
    fn supported_ops(&self) -> &'static [rlx_ir::OpKind] {
        rlx_cpu::SUPPORTED_OPS
    }

    fn compile(&self, graph: Graph, options: &CompileOptions) -> Box<dyn ExecutableGraph> {
        use rlx_opt::pass::Pass as _;
        static ONNX_KERNELS: std::sync::Once = std::sync::Once::new();
        ONNX_KERNELS.call_once(rlx_cpu::onnx_ref::register_onnx_reference_kernels);
        // Claimed `Op::If` / `Op::While` → primitives before legalize (CPU has
        // no native sub-graph executor). Also runs again in the fusion
        // pipeline; no-op when neither op is present.
        let graph = rlx_opt::LowerControlFlow.run(graph);
        // Lower f32 SPD-manifold ops (ReEig/LogEig/BiMap/SpdBatchNorm) to the
        // graph-primitive Jacobi eigensolver BEFORE legalize — like If/While,
        // the CPU backend has no f32 SPD kernel (the native ones are f64
        // LAPACK). No-op unless an f32 SPD op is present; f64 SPD nodes are
        // left untouched for the native LAPACK path.
        let graph = rlx_opt::LowerSpectral.run(graph);
        // PLAN L4: legalize against the backend's claimed op set
        // BEFORE running fusion (so the diagnostic points at the
        // user's IR, not at a fused-away node).
        if let Err(errors) = rlx_opt::legalize_for_backend(&graph, rlx_cpu::SUPPORTED_OPS) {
            panic!("{}", rlx_opt::format_legalize_error("cpu", &errors));
        }
        let policy = options.policy.clone();
        let _precision = options.precision;
        let cfg = rlx_cpu::config::RuntimeConfig::global();

        let graph = crate::precompile::precompile_cleanup(graph, options);

        // Run fusion pipeline (HIR/MIR/LIR ideology — fusion is first-class).
        let mut compile_opts = options.clone();
        compile_opts.arena_alignment = cfg.arena_alignment;
        let compile_result = crate::stages::compile_graph_stages_for_backend(
            rlx_driver::Device::Cpu,
            graph,
            &compile_opts,
            rlx_cpu::SUPPORTED_OPS,
        );
        crate::stages::maybe_log_fusion(&compile_result.fusion);
        let fused = compile_result.lir.into_graph();
        // FusedConvBiasAct / PartitionedConv / FusedTransformerLayer /
        // TransformRegion / BatchElementwiseRegion may survive fusion now
        // that they are claimed — expand to primitives the thunk path runs.
        let fused = rlx_cpu::prepare_graph_for_thunks(fused);

        // Apply precision policy AFTER fusion — Cast nodes don't disrupt
        // the now-flattened fused ops.
        let fused = match policy {
            Some(p) => rlx_opt::AutoMixedPrecision::new(p).run(fused),
            None => fused,
        };

        let io_manifest = cpu_low_precision::IoDtypeManifest::from_graph(&fused);
        let exec_graph = if cpu_low_precision::needs_f32_exec(&fused) {
            cpu_low_precision::promote_to_f32(fused)
        } else {
            fused
        };

        // Re-plan after precision rewrites (may change dtypes / sizes).
        // CPU runs the schedule strictly in-order, so it does NOT need the
        // conservative `pin_output_ancestors` guard (which pins the whole output
        // ancestor DAG to the last step and destroys slot reuse on deep
        // feed-forward graphs — the supertonic vector_estimator/vocoder arenas
        // ballooned to 1.36 GB / 1.17 GB from it). `plan_memory_native` keeps the
        // same Native dtype widths but only pins when the graph has host indexing
        // (Scatter/Gather), so ordinary graphs get full liveness-based reuse.
        //
        // `_in_order` additionally drops the `dequant_host_fallback` pin for the
        // same in-order reason: that guard protects GPU backends whose dequant
        // matmuls flush on a DEFERRED host path, which CPU never takes (it runs
        // `Thunk::DequantMatMulGguf` inline). Left on, its chain-walk pins nearly
        // every activation of a packed K-quant prefill to the graph end — the
        // difference between a ~4 GB reused arena and a 60 GB pinned one on a
        // 30B-class GGUF.
        let plan = memory::plan_memory_native_in_order(&exec_graph, cfg.arena_alignment);
        if rlx_ir::env::flag("RLX_CPU_ARENA_REPORT") {
            eprint!("[rlx-arena-report]\n{}", plan.report());
        }
        if cfg.verbose >= 1 {
            eprintln!(
                "[rlx] arena: {} bytes, {} buffers, alignment: {}",
                plan.arena_size,
                plan.assignments.len(),
                cfg.arena_alignment
            );
        }
        Box::new(build_cpu_executable(
            exec_graph,
            plan,
            io_manifest,
            options.rng,
        ))
    }

    fn compile_lir(&self, lir: LirModule, options: &CompileOptions) -> Box<dyn ExecutableGraph> {
        // `Instant` is unimplemented on wasm32 — never touch it there.
        #[cfg(not(target_arch = "wasm32"))]
        let prof = rlx_ir::env::var_os("RLX_PROFILE_COMPILE").is_some();
        #[cfg(not(target_arch = "wasm32"))]
        let tt = if prof {
            Some(std::time::Instant::now())
        } else {
            None
        };
        let alignment = lir.buffers.alignment.max(options.arena_alignment);
        let mut graph = lir.into_graph();
        // Session path: fusion may leave claimed fused/region ops in LIR;
        // expand before broadcast legalize + arena plan + thunks.
        graph = rlx_cpu::prepare_graph_for_thunks(graph);
        {
            use rlx_opt::pass::Pass as _;
            graph = rlx_opt::LegalizeBroadcast.run(graph);
        }
        if let Some(p) = options.policy.clone() {
            use rlx_opt::pass::Pass;
            graph = rlx_opt::AutoMixedPrecision::new(p).run(graph);
        }
        let io_manifest = cpu_low_precision::IoDtypeManifest::from_graph(&graph);
        let promote = cpu_low_precision::needs_f32_exec(&graph);
        let exec_graph = if promote {
            cpu_low_precision::promote_to_f32(graph)
        } else {
            graph
        };
        #[cfg(not(target_arch = "wasm32"))]
        let t_prep = tt.map(|t| t.elapsed());
        #[cfg(not(target_arch = "wasm32"))]
        let t1 = if prof {
            Some(std::time::Instant::now())
        } else {
            None
        };
        // LegalizeBroadcast may insert Expand nodes — must replan; the
        // embedded LIR buffer map is from before legalization. In-order CPU
        // executor → `plan_memory_native_in_order` (unpinned unless host
        // indexing, and no deferred-host dequant pin) for full liveness slot
        // reuse; see the note at the other plan site above.
        let plan = memory::plan_memory_native_in_order(&exec_graph, alignment);
        #[cfg(not(target_arch = "wasm32"))]
        let t_plan = t1.map(|t| t.elapsed());
        #[cfg(not(target_arch = "wasm32"))]
        if prof {
            eprintln!(
                "[compile_lir] {} nodes: into_graph+passes={:?} plan_memory={:?}",
                exec_graph.nodes().len(),
                t_prep.unwrap(),
                t_plan.unwrap(),
            );
        }
        let cfg = rlx_cpu::config::RuntimeConfig::global();
        if cfg.verbose >= 1 {
            eprintln!(
                "[rlx] compile_lir: arena {} bytes ({} buffers, alignment {})",
                plan.arena_size,
                plan.assignments.len(),
                alignment,
            );
        }
        #[cfg(not(target_arch = "wasm32"))]
        let t2 = if prof {
            Some(std::time::Instant::now())
        } else {
            None
        };
        let exe = build_cpu_executable(exec_graph, plan, io_manifest, options.rng);
        #[cfg(not(target_arch = "wasm32"))]
        if let Some(t2) = t2 {
            eprintln!("[compile_lir] build_thunks={:?}", t2.elapsed());
        }
        Box::new(exe)
    }
}

fn build_cpu_executable(
    graph: Graph,
    plan: MemoryPlan,
    io_manifest: cpu_low_precision::IoDtypeManifest,
    rng: rlx_ir::RngOptions,
) -> CpuExecutable {
    let mut arena = Arena::from_plan(plan);
    let mut input_ids = HashMap::new();
    let mut param_ids = HashMap::new();
    let mut node_dtypes: HashMap<NodeId, DType> = HashMap::new();
    for node in graph.nodes() {
        node_dtypes.insert(node.id, node.shape.dtype());
        match &node.op {
            Op::Input { name } => {
                input_ids.insert(name.clone(), node.id);
            }
            Op::Param { name } => {
                param_ids.insert(name.clone(), node.id);
            }
            _ => {}
        }
    }

    let schedule = thunk::compile_thunks_with_rng(&graph, &arena, rng);

    let mut input_slots = Vec::new();
    for node in graph.nodes() {
        if let Op::Input { name } = &node.op {
            let off = arena.byte_offset(node.id);
            let len = node.shape.num_elements().unwrap_or(0);
            input_slots.push((name.clone(), off, len, node.shape.dtype()));
        }
    }

    let output_slots: Vec<(usize, usize)> = graph
        .outputs
        .iter()
        .map(|&id| {
            let off = arena.byte_offset(id);
            let len = graph.node(id).shape.num_elements().unwrap_or(0);
            (off, len)
        })
        .collect();

    for node in graph.nodes() {
        if let Op::Constant { data } = &node.op
            && arena.has_buffer(node.id)
            && !data.is_empty()
        {
            match node.shape.dtype() {
                // True-width dtypes (their arena slot is sized to the real
                // element width, not f32): copy the raw bytes. I64/I32/U32
                // constants were previously caught by the f32-reinterpret
                // branch below, which read them in 4-byte chunks as f32 —
                // corrupting e.g. the VITS sequence-mask `arange` constant
                // (i64 [0..T-1]) so the downstream i64 `Compare` read garbage.
                // Bool/U8/I8 must also raw-copy: a 1-byte `ConstantOfShape(true)`
                // has `data.len()/4 == 0` in the f32 branch, so the slot stayed
                // zero and Soprano's causal `And(ones, mask)` collapsed → Softmax NaN.
                DType::F64
                | DType::F16
                | DType::BF16
                | DType::I64
                | DType::I32
                | DType::U32
                | DType::Bool
                | DType::U8
                | DType::I8 => {
                    let off = arena.byte_offset(node.id);
                    let buf = arena.raw_buf_mut();
                    let n = buf.len().saturating_sub(off).min(data.len());
                    buf[off..off + n].copy_from_slice(&data[..n]);
                }
                _ => {
                    let buf = arena.slice_mut(node.id);
                    let n_floats = data.len() / 4;
                    let n = buf.len().min(n_floats);
                    for i in 0..n {
                        let bytes = [
                            data[i * 4],
                            data[i * 4 + 1],
                            data[i * 4 + 2],
                            data[i * 4 + 3],
                        ];
                        buf[i] = f32::from_le_bytes(bytes);
                    }
                }
            }
        }
    }

    CpuExecutable {
        graph,
        arena,
        input_ids,
        param_ids,
        node_dtypes,
        io_manifest,
        schedule,
        input_slots,
        output_slots,
        handles: HashMap::new(),
        active_extent: None,
        moe_resident: None,
        moe_resident_layers: None,
        moe_topk_capture: None,
        baseline_written: false,
    }
}

#[derive(Clone)]
struct CpuExecutable {
    graph: Graph,
    arena: Arena,
    input_ids: HashMap<String, NodeId>,
    param_ids: HashMap<String, NodeId>,
    /// Per-node arena dtype. Lets set_param/run cast f32 ↔ F16/BF16
    /// when AutoMixedPrecision has rewritten the graph.
    node_dtypes: HashMap<NodeId, DType>,
    /// User-facing boundary dtypes (before f32 promotion for CPU exec).
    io_manifest: cpu_low_precision::IoDtypeManifest,
    schedule: thunk::ThunkSchedule,
    // Pre-resolved: ordered list of (input_name, arena_byte_offset, max_elems, dtype)
    input_slots: Vec<(String, usize, usize, DType)>,
    /// Output (byte_offset, num_elements). dtype is in node_dtypes.
    output_slots: Vec<(usize, usize)>,
    /// Persistent buffer handles (KV-cache, optimizer state, etc.).
    /// Lives outside the arena and survives across run() calls.
    /// On run(): if a handle's name matches a graph input, the
    /// handle's data is used as the input.
    handles: HashMap<String, Vec<f32>>,
    /// Active-extent hint (`Some((actual, upper))`) for L1 bucketed
    /// dispatch. When set AND every thunk in the schedule is in
    /// `Thunk::safe_for_active_extent`, the executor processes only
    /// `actual / upper` of each kernel's work. Otherwise (or when
    /// `None`) runs at the full compiled extent. See PLAN L1.
    active_extent: Option<(usize, usize)>,
    moe_resident: Option<std::sync::Arc<[bool]>>,
    moe_resident_layers: Option<std::sync::Arc<Vec<std::sync::Arc<[bool]>>>>,
    moe_topk_capture: Option<std::sync::Arc<rlx_cpu::moe_topk_capture::MoeTopkCapture>>,
    /// Whether params + constants are already resident in the arena. While
    /// `true`, `restore_arena_baseline` zeros only the scratch buffers instead
    /// of re-zeroing + rewriting every param each run (which is O(params) and
    /// allocates a full params clone — catastrophic for multi-GB models).
    /// `set_param`/`set_param_typed` reset it to `false`.
    baseline_written: bool,
}

unsafe impl Send for CpuExecutable {}

impl CpuExecutable {
    /// `RLX_CPU_ARENA_REPORT=1`: attribute the arena's resident bytes to nodes,
    /// split persistent (Param/Input/Constant — live the whole run) vs scratch,
    /// so a bloated footprint can be traced to specific buffers.
    fn arena_report_if_requested(&self) {
        if !rlx_ir::env::flag("RLX_CPU_ARENA_REPORT") {
            return;
        }
        let mut rows: Vec<(usize, NodeId, bool)> = Vec::new();
        for node in self.graph.nodes() {
            if !self.arena.has_buffer(node.id) {
                continue;
            }
            let dt = self
                .node_dtypes
                .get(&node.id)
                .copied()
                .unwrap_or(DType::F32);
            let bytes = node.shape.num_elements().unwrap_or(0) * dt.size_bytes();
            let persistent = matches!(
                node.op,
                Op::Param { .. } | Op::Input { .. } | Op::Constant { .. }
            );
            rows.push((bytes, node.id, persistent));
        }
        rows.sort_by_key(|b| std::cmp::Reverse(b.0));
        let persist: usize = rows.iter().filter(|r| r.2).map(|r| r.0).sum();
        let scratch: usize = rows.iter().filter(|r| !r.2).map(|r| r.0).sum();
        eprintln!(
            "[arena-report] {} buffers, persistent={}MB scratch={}MB",
            rows.len(),
            persist >> 20,
            scratch >> 20
        );
        for (b, id, p) in rows.iter().take(15) {
            eprintln!("  {:>6}MB  node {:>5}  persistent={}", b >> 20, id, p);
        }
    }

    /// Per-node dump (mirror of rlx-wgpu's `RLX_WGPU_DUMP_NODES`) for
    /// cross-backend divergence bisection: each F32 node's max|x| + nonzero
    /// count in topo order. Diff against the wgpu dump to find the first
    /// diverging node. `RLX_CPU_DUMP_FLAT=<i>` also prints that flat element.
    fn dump_nodes_if_requested(&self) {
        self.arena_report_if_requested();
        if !rlx_ir::env::flag("RLX_CPU_DUMP_NODES") {
            return;
        }
        let limit = rlx_ir::env::parse_or("RLX_CPU_DUMP_NODES_LIMIT", 2000usize);
        let flat_probe = rlx_ir::env::parse_or::<usize>("RLX_CPU_DUMP_FLAT", usize::MAX);
        eprintln!("[rlx-cpu-dump] per-node max |x| (topo order, limit={limit})");
        let buf = self.arena.raw_buf();
        // `RLX_CPU_DUMP_NODE_DATA=<id>[,<id>...]` writes each node's raw f32
        // buffer to `$RLX_CPU_DUMP_DIR/node_<id>.f32`. Unlike the summary loop
        // below this does not skip Param/Reshape, so a matmul's weights can be
        // pulled out and the op recomputed in f64 to decide which backend is
        // right when two disagree.
        if let Some(spec) = rlx_ir::env::var("RLX_CPU_DUMP_NODE_DATA") {
            let dir = rlx_ir::env::var("RLX_CPU_DUMP_DIR").unwrap_or_else(|| "/tmp".into());
            for want in spec.split(',').filter_map(|t| t.trim().parse::<u32>().ok()) {
                let Some(node) = self.graph.nodes().iter().find(|n| n.id.0 == want) else {
                    continue;
                };
                if !self.arena.has_buffer(node.id) {
                    continue;
                }
                let off = self.arena.byte_offset(node.id);
                let n = node.shape.num_elements().unwrap_or(0);
                let data: &[f32] =
                    unsafe { std::slice::from_raw_parts(buf.as_ptr().add(off) as *const f32, n) };
                let bytes: Vec<u8> = data.iter().flat_map(|v| v.to_le_bytes()).collect();
                let path = format!("{dir}/node_{want}.f32");
                match std::fs::write(&path, &bytes) {
                    Ok(()) => eprintln!(
                        "[rlx-cpu-dump] wrote {path} ({n} f32) byte_range=[{off}, {})",
                        off + n * 4
                    ),
                    Err(e) => eprintln!("[rlx-cpu-dump] {path}: {e}"),
                }
            }
        }
        let mut shown = 0usize;
        let (mut no_buf, mut skipped_op, mut not_f32) = (0usize, 0usize, 0usize);
        for (i, node) in self.graph.nodes().iter().enumerate() {
            if !self.arena.has_buffer(node.id) {
                no_buf += 1;
                continue;
            }
            if matches!(
                node.op,
                rlx_ir::Op::Input { .. }
                    | rlx_ir::Op::Param { .. }
                    | rlx_ir::Op::Constant { .. }
                    | rlx_ir::Op::Reshape { .. }
                    | rlx_ir::Op::Cast { .. }
            ) {
                skipped_op += 1;
                continue;
            }
            if self
                .node_dtypes
                .get(&node.id)
                .copied()
                .unwrap_or(DType::F32)
                != DType::F32
            {
                not_f32 += 1;
                continue;
            }
            let off = self.arena.byte_offset(node.id);
            let n = node.shape.num_elements().unwrap_or(0);
            let data: &[f32] =
                unsafe { std::slice::from_raw_parts(buf.as_ptr().add(off) as *const f32, n) };
            let max = data.iter().fold(0f32, |m, &v| m.max(v.abs()));
            let nz = data.iter().filter(|&&v| v != 0.0).count();
            // f64 sum: max/tail are single-element summaries and two tensors can
            // share both while differing elsewhere, which hides the real first
            // divergence. A whole-tensor sum does not.
            let sum: f64 = data.iter().map(|&v| v as f64).sum();
            // Max over the LAST row (final index of the second-to-last axis).
            // Cross-backend divergences often sit in the tail token only, where
            // the whole-tensor max is dominated by other rows and matches.
            let row = node
                .shape
                .dims()
                .last()
                .and_then(|d| match d {
                    rlx_ir::Dim::Static(n) => Some(*n),
                    rlx_ir::Dim::Dynamic(_) => None,
                })
                .unwrap_or(data.len())
                .max(1);
            let tail = data
                .len()
                .checked_sub(row)
                .map(|o| data[o..].iter().fold(0f32, |m, &v| m.max(v.abs())))
                .unwrap_or(max);
            // Sum of |v| over the tail row. A LayerNorm row sums to ~0, so a
            // plain sum cancels and hides real differences; this does not.
            let tabs: f64 = data
                .len()
                .checked_sub(row)
                .map(|o| data[o..].iter().map(|&v| v.abs() as f64).sum())
                .unwrap_or(0.0);
            let flat_s = if flat_probe < data.len() {
                format!(" flat[{flat_probe}]={:.6}", data[flat_probe])
            } else {
                String::new()
            };
            eprintln!(
                "  [{i:>3}] in={:?} id={:?} name={:?} {:?} shape={:?} max={max:.6} tail={tail:.6} sum={sum:.6} tabs={tabs:.6} nonzero={nz}/{}{flat_s}",
                node.inputs.iter().map(|n| n.0).collect::<Vec<_>>(),
                node.id,
                node.name,
                node.op,
                node.shape.dims(),
                data.len()
            );
            // `RLX_CPU_ROW_TABS=<node_id>`: per-row sum of |v| for one node.
            // Lets a GPU tail row be searched for among ALL the reference rows —
            // if it matches a different row exactly, the kernel is indexing
            // wrong rather than computing wrong.
            if rlx_ir::env::parse_or::<u32>("RLX_CPU_ROW_TABS", u32::MAX) == node.id.0 {
                for (r, chunk) in data.chunks(row).enumerate() {
                    let t: f64 = chunk.iter().map(|&v| v.abs() as f64).sum();
                    eprintln!("  rowtabs node={} row={r} tabs={t:.6}", node.id.0);
                }
            }
            shown += 1;
            if shown >= limit {
                break;
            }
        }
        eprintln!(
            "[rlx-cpu-dump] nodes={} shown={shown} no_buffer={no_buf} skipped_op={skipped_op} not_f32={not_f32}",
            self.graph.nodes().len()
        );
    }

    /// Write a f32 input slice into the arena, casting to the node's dtype.
    fn write_input(&mut self, id: NodeId, data: &[f32]) {
        let dtype = self.node_dtypes.get(&id).copied().unwrap_or(DType::F32);
        let off = self.arena.byte_offset(id);
        let buf = self.arena.raw_buf_mut();
        let elem_size = dtype.size_bytes();
        let max_elems = (buf.len() - off) / elem_size;
        unsafe {
            write_typed_from_f32(buf.as_mut_ptr().add(off), dtype, data, max_elems);
        }
    }

    /// Read a node's arena bytes back as Vec<f32>, casting from its dtype.
    fn read_output(&self, id: NodeId) -> Vec<f32> {
        let dtype = self.node_dtypes.get(&id).copied().unwrap_or(DType::F32);
        let off = self.arena.byte_offset(id);
        let buf = self.arena.raw_buf();
        let n_elems = self.graph.node(id).shape.num_elements().unwrap_or(0);
        unsafe { read_typed_to_f32(buf.as_ptr().add(off), dtype, n_elems) }
    }
}

impl ExecutableGraph for CpuExecutable {
    fn capabilities(&self) -> crate::ExecutableCapabilities {
        crate::ExecutableCapabilities {
            clone: true,
            moe: true,
            // `bind_handle` / `read_handle` are implemented below and return
            // real data; the flag said otherwise.
            persistent_handles: true,
            typed_io: true,
            active_extent: true,
            ..crate::ExecutableCapabilities::NONE
        }
    }

    fn clone_box(&self) -> Box<dyn ExecutableGraph> {
        Box::new(self.clone())
    }
    fn set_param(&mut self, name: &str, data: &[f32]) {
        // Params live solely in the arena (dedicated, never-aliased slots, see
        // the memory planner) — no redundant CPU-side copy is kept, which would
        // double the weight footprint for multi-GB models.
        // Cast f32 → arena dtype when the param has been rewritten to F16/BF16.
        if let Some(&id) = self.param_ids.get(name)
            && self.arena.has_buffer(id)
        {
            let dtype = self.node_dtypes.get(&id).copied().unwrap_or(DType::F32);
            let off = self.arena.byte_offset(id);
            let buf = self.arena.raw_buf_mut();
            let elem_size = dtype.size_bytes();
            let max_elems = (buf.len() - off) / elem_size;
            unsafe {
                write_typed_from_f32(buf.as_mut_ptr().add(off), dtype, data, max_elems);
            }
        }
    }

    fn set_param_range(&mut self, name: &str, byte_offset: usize, data: &[u8]) -> bool {
        // Paged-MoE residency: rewrite ONE expert slot of a large bank instead
        // of re-uploading the bank. Without this the caller has to assemble the
        // whole gathered bank in a Vec and hand it over, copying every active
        // byte twice — at GLM-5.3-Flash scale ~2 GB of redundant memcpy per
        // token, and CPU is the backend the cluster planner assigns paged
        // stages to most often.
        let Some(&id) = self.param_ids.get(name) else {
            return false;
        };
        if !self.arena.has_buffer(id) {
            return false;
        }
        // The bytes go in verbatim, and the signature carries no dtype to check
        // them against. A param that AMP rewrote to F16/BF16 would be silently
        // corrupted by F32 bytes, so refuse anything but F32 storage and let
        // the caller take the whole-buffer path, which does convert.
        if self.node_dtypes.get(&id).copied().unwrap_or(DType::F32) != DType::F32 {
            return false;
        }
        // Bound to the PARAM, not to the arena: params sit in a shared buffer,
        // so an over-long write does not fault — it lands in whichever tensor
        // is next and corrupts it silently.
        let size = self.arena.byte_size(id);
        let Some(end) = byte_offset.checked_add(data.len()) else {
            return false;
        };
        if end > size {
            return false;
        }
        let off = self.arena.byte_offset(id);
        self.arena.raw_buf_mut()[off + byte_offset..off + end].copy_from_slice(data);
        true
    }

    fn run(&mut self, inputs: &[(&str, &[f32])]) -> Vec<Vec<f32>> {
        self.restore_arena_baseline();
        // 1. Apply persistent handles first — they act like default inputs.
        //    Explicit `inputs` passed to run() override matching handle names.
        let handle_names: Vec<String> = self.handles.keys().cloned().collect();
        for name in &handle_names {
            if let Some(&id) = self.input_ids.get(name)
                && self.arena.has_buffer(id)
            {
                let data = self.handles.get(name).cloned().unwrap_or_default();
                self.write_input(id, &data);
            }
        }
        // 2. Explicit per-call inputs override handles.
        for &(name, data) in inputs {
            if let Some(&id) = self.input_ids.get(name)
                && self.arena.has_buffer(id)
            {
                self.write_input(id, data);
            }
        }

        // Active-extent fast-path (PLAN L1): if hinted AND every thunk
        // in the schedule supports it, run scaled. Otherwise fall back
        // to full-extent dispatch — preserves correctness when the
        // schedule contains a thunk that hasn't yet been wired in.
        let active_used = if let Some((actual, upper)) = self.active_extent {
            thunk::execute_thunks_active(&self.schedule, self.arena.raw_buf_mut(), actual, upper)
        } else {
            false
        };
        if !active_used {
            // Execute via pre-compiled thunks (zero per-node dispatch overhead)
            thunk::execute_thunks(&self.schedule, self.arena.raw_buf_mut());
        }

        self.dump_nodes_if_requested();

        // 3. Sync any handle whose name matches a graph OUTPUT —
        //    KV-cache pattern: outputs flow back into the same-named
        //    handle for the next iteration.
        for (idx, &out_id) in self.graph.outputs.iter().enumerate() {
            let name = format!("out{idx}");
            if self.handles.contains_key(&name) {
                let v = self.read_output(out_id);
                self.handles.insert(name, v);
            }
        }

        self.graph
            .outputs
            .iter()
            .map(|&out_id| self.read_output(out_id))
            .collect()
    }

    fn run_raw(&mut self, inputs: &[(&str, &[f32])]) -> Vec<(*const f32, usize)> {
        self.restore_arena_baseline();
        // Copy inputs by name (HashMap lookup), casting to arena dtype.
        for &(name, data) in inputs {
            if let Some(&id) = self.input_ids.get(name)
                && self.arena.has_buffer(id)
            {
                self.write_input(id, data);
            }
        }
        thunk::execute_thunks(&self.schedule, self.arena.raw_buf_mut());
        // Note: pointers are raw arena bytes — for F16 outputs, callers
        // must read 2 bytes/elem, not 4. run() is the safe path for
        // mixed precision; run_raw() is only meaningful for F32.
        self.graph
            .outputs
            .iter()
            .map(|&out_id| {
                let (ptr, len) = self.arena.raw_ptr(out_id);
                (ptr as *const f32, len)
            })
            .collect()
    }

    /// Fastest path: inputs by index (matching input_slots order), zero-copy output.
    /// No HashMap, no name matching, no Vec allocation. Casts f32 input
    /// to F16/BF16 if the input slot's dtype was rewritten.
    fn run_slots(&mut self, inputs: &[&[f32]]) -> &[(usize, usize)] {
        self.restore_arena_baseline();
        let buf = self.arena.raw_buf_mut();
        for (i, &data) in inputs.iter().enumerate() {
            if i < self.input_slots.len() {
                let (_, off, max_len, dtype) = &self.input_slots[i];
                unsafe {
                    write_typed_from_f32(buf.as_mut_ptr().add(*off), *dtype, data, *max_len);
                }
            }
        }
        thunk::execute_thunks(&self.schedule, self.arena.raw_buf_mut());
        &self.output_slots
    }

    fn arena_ptr(&self) -> *const u8 {
        self.arena.raw_buf_mut_ptr()
    }

    fn bind_handle(&mut self, name: &str, data: &[f32]) -> bool {
        // Persistent buffer: stored separately from arena, survives run().
        // If the name matches a graph input, run() will use this data
        // as the input. If the graph also writes back to this name (via
        // an output binding pattern), read_handle returns the latest.
        self.handles.insert(name.to_string(), data.to_vec());
        true
    }

    fn read_handle(&self, name: &str) -> Option<Vec<f32>> {
        self.handles.get(name).cloned()
    }

    fn set_active_extent(&mut self, extent: Option<(usize, usize)>) {
        self.active_extent = extent;
    }

    fn release_scratch(&mut self) {
        self.release_scratch_pages();
    }

    fn set_rng(&mut self, rng: rlx_ir::RngOptions) {
        *self.schedule.rng.write().unwrap() = rng;
    }

    fn rng(&self) -> rlx_ir::RngOptions {
        *self.schedule.rng.read().unwrap()
    }

    fn set_moe_resident_experts(&mut self, mask: &[bool]) {
        self.moe_resident_layers = None;
        self.schedule.moe_resident_layers = None;
        self.moe_resident = Some(Arc::from(mask));
        self.schedule.moe_resident = self.moe_resident.clone();
    }

    fn set_moe_resident_experts_per_layer(&mut self, masks: &[&[bool]]) {
        self.moe_resident = None;
        self.schedule.moe_resident = None;
        let layers: Vec<Arc<[bool]>> = masks.iter().map(|m| Arc::from(*m)).collect();
        let arc = Arc::new(layers);
        self.moe_resident_layers = Some(arc.clone());
        self.schedule.moe_resident_layers = Some(arc);
    }

    fn enable_moe_topk_capture(&mut self, num_experts: usize) -> bool {
        let cap = rlx_cpu::moe_topk_capture::MoeTopkCapture::new(num_experts);
        self.moe_topk_capture = Some(cap.clone());
        self.schedule.moe_topk_capture = Some(cap);
        true
    }

    fn take_moe_topk_capture(&mut self) -> Option<Vec<Vec<u32>>> {
        let cap = self.moe_topk_capture.as_ref()?;
        let layers = cap.take_layers();
        if layers.is_empty() {
            None
        } else {
            Some(layers)
        }
    }

    fn take_moe_residency_stats(&mut self) -> Option<crate::MoeResidencyStats> {
        rlx_cpu::moe_residency::take_last_forward_stats()
    }

    /// Typed param upload. F32 / F16 / BF16 go through the existing
    /// widen-to-f32 path (the CPU arena is historically f32 with
    /// optional half-precision rewrite). F64 (and any future
    /// non-widenable dtype) lands directly in the arena as bytes —
    /// the f32 path would lose precision.
    fn set_param_typed(&mut self, name: &str, data: &[u8], dtype: rlx_ir::DType) {
        if matches!(dtype, DType::F64 | DType::I64 | DType::I32 | DType::U32) {
            self.set_param_bytes(name, data, dtype);
            return;
        }
        // U8 / I8 raw byte tensors: opaque storage for the GGUF
        // K-quant `Op::DequantMatMul` path (weights stay packed
        // in the arena). One arena byte = one element.
        if matches!(dtype, DType::U8 | DType::I8) {
            self.set_param_bytes(name, data, dtype);
            return;
        }
        if dtype == DType::F32 {
            let n = data.len() / 4;
            let s = unsafe { std::slice::from_raw_parts(data.as_ptr() as *const f32, n) };
            self.set_param(name, s);
        } else {
            let f32_buf = super::widen_bytes_to_f32(data, dtype);
            self.set_param(name, &f32_buf);
        }
    }

    /// Typed run with mixed-dtype inputs/outputs.
    ///
    /// For each input: if its declared graph dtype matches the
    /// caller's bytes, we write directly into the arena (zero
    /// precision loss — F64 stays F64). For F32 with a half-precision
    /// arena rewrite, we widen as before. F16/BF16 callers go
    /// through the existing widen path.
    ///
    /// Outputs are read straight from the arena in the graph node's
    /// declared dtype — F64 outputs come back as 8 bytes/element,
    /// F32 as 4, etc.
    fn run_typed(
        &mut self,
        inputs: &[(&str, &[u8], rlx_ir::DType)],
    ) -> Vec<(Vec<u8>, rlx_ir::DType)> {
        // Decide: are *all* inputs F64? If so, use the direct-byte
        // path for everything and skip the f32 widening machinery
        // entirely. Mixed dtype graphs (F32 + F64) take the
        // per-input dispatch route below.
        let all_f64 = !inputs.is_empty() && inputs.iter().all(|(_, _, dt)| *dt == DType::F64);

        if all_f64 {
            for (name, data, _) in inputs {
                if let Some(&id) = self.input_ids.get(*name) {
                    if !self.arena.has_buffer(id) {
                        continue;
                    }
                    let off = self.arena.byte_offset(id);
                    let buf = self.arena.raw_buf_mut();
                    let n = data.len();
                    debug_assert!(
                        off + n <= buf.len(),
                        "run_typed: input '{name}' overflows arena slot"
                    );
                    buf[off..off + n].copy_from_slice(data);
                }
            }
            thunk::execute_thunks(&self.schedule, self.arena.raw_buf_mut());
        } else {
            // Mixed-dtype path: dtypes that survive untouched
            // through the f32-aliased arena (F64, I32, I64, U32)
            // go in as bytes; F32 and the half-precision family
            // route through widen-to-f32 + run.
            let mut f32_owned: Vec<(String, Vec<f32>)> = Vec::new();
            for (name, data, dt) in inputs {
                let direct = matches!(
                    *dt,
                    DType::F64 | DType::I32 | DType::I64 | DType::U32 | DType::C64 | DType::C128
                );
                if direct {
                    if let Some(&id) = self.input_ids.get(*name) {
                        if !self.arena.has_buffer(id) {
                            continue;
                        }
                        let off = self.arena.byte_offset(id);
                        let buf = self.arena.raw_buf_mut();
                        buf[off..off + data.len()].copy_from_slice(data);
                    }
                } else {
                    let v = super::widen_bytes_to_f32(data, *dt);
                    f32_owned.push((name.to_string(), v));
                }
            }
            for (name, data) in &f32_owned {
                if let Some(&id) = self.input_ids.get(name.as_str()) {
                    if self.arena.has_buffer(id) {
                        self.write_input(id, data);
                    }
                }
            }
            let active_used = if let Some((actual, upper)) = self.active_extent {
                thunk::execute_thunks_active(
                    &self.schedule,
                    self.arena.raw_buf_mut(),
                    actual,
                    upper,
                )
            } else {
                false
            };
            if !active_used {
                thunk::execute_thunks(&self.schedule, self.arena.raw_buf_mut());
            }
        }

        self.dump_nodes_if_requested();

        // Read outputs in declared boundary dtypes.
        self.graph
            .outputs
            .iter()
            .enumerate()
            .map(|(idx, &id)| {
                let exec_dtype = self.graph.node(id).shape.dtype();
                let declared = self.io_manifest.output_dtype(idx, exec_dtype);
                if matches!(
                    exec_dtype,
                    DType::F64
                        | DType::F16
                        | DType::BF16
                        | DType::I32
                        | DType::I64
                        | DType::U32
                        | DType::C64
                        | DType::C128
                ) {
                    let n_elems = self.graph.node(id).shape.num_elements().unwrap_or(0);
                    let n_bytes = n_elems * exec_dtype.size_bytes();
                    let off = self.arena.byte_offset(id);
                    let bytes = self.arena.raw_buf()[off..off + n_bytes].to_vec();
                    return (bytes, declared);
                }
                let f32_vals = self.read_output(id);
                if declared != exec_dtype {
                    return (super::narrow_f32_to_bytes(&f32_vals, declared), declared);
                }
                let bytes = f32_vals.iter().flat_map(|v| v.to_le_bytes()).collect();
                (bytes, declared)
            })
            .collect()
    }
}

impl CpuExecutable {
    /// Clear ephemeral (scratch) arena slots before each `run()`. Params are
    /// written into their dedicated, never-aliased arena slots by `set_param`
    /// and live for the whole execution, so they are NOT re-zeroed/rewritten
    /// here — only the intermediate buffers (which carry stale data from the
    /// previous pass) are zeroed. Compile-time constants are written once.
    ///
    /// This keeps the per-run cost O(scratch) instead of O(params): a previous
    /// version cloned + rewrote the entire (multi-GB) weight region every run,
    /// which made large models swap-thrash.
    /// Sorted, merged persistent (param + constant) byte ranges in the arena.
    /// These live for the whole execution; everything else is ephemeral scratch.
    /// Used by `release_scratch_pages` (returns the complement's pages to the
    /// OS), which is `cfg(unix)` — hence the same gate here.
    /// `restore_arena_baseline` computes its own persistent set inline.
    #[cfg(unix)]
    fn persistent_keep_ranges(&self) -> Vec<(usize, usize)> {
        let persistent: std::collections::HashSet<NodeId> = {
            let mut s: std::collections::HashSet<NodeId> =
                self.param_ids.values().copied().collect();
            for node in self.graph.nodes() {
                if matches!(node.op, Op::Constant { .. }) {
                    s.insert(node.id);
                }
            }
            s
        };
        let mut keep: Vec<(usize, usize)> = self
            .graph
            .nodes()
            .iter()
            .filter_map(|node| {
                let id = node.id;
                if !persistent.contains(&id) || !self.arena.has_buffer(id) {
                    return None;
                }
                let dtype = self.node_dtypes.get(&id).copied().unwrap_or(DType::F32);
                let nbytes = node.shape.num_elements().unwrap_or(0) * dtype.size_bytes();
                let off = self.arena.byte_offset(id);
                Some((off, off + nbytes))
            })
            .collect();
        keep.sort_unstable();
        keep
    }

    /// Return the ephemeral (scratch) pages of the arena to the OS via
    /// `madvise(MADV_FREE)`, dropping resident RSS while this cached graph sits
    /// idle between pipeline stages. The arena's virtual allocation and its param
    /// bytes stay intact; only whole pages that fall ENTIRELY inside the scratch
    /// complement of the param ranges are advised, so weight bytes are never
    /// touched. The next run's `restore_arena_baseline` re-zeros scratch, so this
    /// is purely advisory and cannot affect correctness. Call it only when the
    /// graph won't run again immediately (e.g. a TTS pipeline finished with a
    /// subgraph) — never between the steps of a hot loop, or the pages just fault
    /// back in. No-op on non-unix.
    #[allow(unused_variables)]
    fn release_scratch_pages(&self) {
        #[cfg(unix)]
        {
            // 16 KiB covers Apple-silicon pages and is a multiple of 4 KiB, so
            // rounding to it yields valid page-aligned spans on 4 KiB hosts too.
            const PAGE: usize = 16 * 1024;
            let keep = self.persistent_keep_ranges();
            let base = self.arena.raw_buf_mut_ptr() as usize;
            let total = self.arena.size();
            // Walk the complement of `keep` = the scratch spans.
            let mut cursor = 0usize;
            let mut spans: Vec<(usize, usize)> = Vec::new();
            for &(s, e) in &keep {
                let s = s.min(total);
                if cursor < s {
                    spans.push((cursor, s));
                }
                cursor = cursor.max(e.min(total));
            }
            if cursor < total {
                spans.push((cursor, total));
            }
            for (s, e) in spans {
                let a = (base + s).div_ceil(PAGE) * PAGE;
                let b = (base + e) / PAGE * PAGE;
                if b > a {
                    // macOS MADV_FREE (5) is lazy — it flags pages freeable but does
                    // NOT drop them from the physical footprint until memory
                    // pressure. MADV_FREE_REUSABLE (7) returns them to the OS now and
                    // decrements phys_footprint (the next access refaults). Linux
                    // MADV_DONTNEED (4) drops anonymous pages immediately and
                    // zero-fills on refault. Override with RLX_MADV to trade footprint
                    // for refault latency (e.g. RLX_MADV=5 for lazy/pressure-only).
                    #[cfg(target_os = "macos")]
                    const ADVICE: i32 = 7;
                    #[cfg(not(target_os = "macos"))]
                    const ADVICE: i32 = 4;
                    let advice = rlx_ir::env::var("RLX_MADV")
                        .and_then(|s| s.parse().ok())
                        .unwrap_or(ADVICE);
                    unsafe { madvise(a as *mut core::ffi::c_void, b - a, advice) };
                }
            }
        }
    }

    fn restore_arena_baseline(&mut self) {
        // Persistent slots (params + constants) — never zeroed.
        let persistent: std::collections::HashSet<NodeId> = {
            let mut s: std::collections::HashSet<NodeId> =
                self.param_ids.values().copied().collect();
            for node in self.graph.nodes() {
                if matches!(node.op, Op::Constant { .. }) {
                    s.insert(node.id);
                }
            }
            s
        };

        // Write compile-time constants into the arena once (a fresh arena is
        // zero-initialized; params are already resident via set_param).
        if !self.baseline_written {
            let constants: Vec<(NodeId, DType, Vec<u8>)> = self
                .graph
                .nodes()
                .iter()
                .filter_map(|node| {
                    if let Op::Constant { data } = &node.op
                        && self.arena.has_buffer(node.id)
                        && !data.is_empty()
                    {
                        Some((node.id, node.shape.dtype(), data.clone()))
                    } else {
                        None
                    }
                })
                .collect();
            for (id, dtype, data) in constants {
                self.write_constant_to_arena(id, dtype, &data);
            }
            self.baseline_written = true;
        }

        // Zero everything EXCEPT the persistent (param + constant) byte ranges.
        //
        // We zero the *complement* of the persistent ranges rather than each
        // scratch node's exact byte span. That covers inter-slot padding and
        // arena gaps too — a kernel that over-reads its input into adjacent
        // alignment padding (common in SIMD reductions) would otherwise pick up
        // stale bytes from a previous run, since per-node zeroing only clears
        // `num_elements` and leaves the padding dirty. The cost stays O(arena −
        // params): for a 7B the params dominate the arena and are skipped, so
        // the swept region is tiny.
        let mut keep: Vec<(usize, usize)> = self
            .graph
            .nodes()
            .iter()
            .filter_map(|node| {
                let id = node.id;
                if !persistent.contains(&id) || !self.arena.has_buffer(id) {
                    return None;
                }
                let dtype = self.node_dtypes.get(&id).copied().unwrap_or(DType::F32);
                let nbytes = node.shape.num_elements().unwrap_or(0) * dtype.size_bytes();
                let off = self.arena.byte_offset(id);
                Some((off, off + nbytes))
            })
            .collect();
        keep.sort_unstable();

        let buf = self.arena.raw_buf_mut();
        let len = buf.len();
        let mut cursor = 0usize;
        for (start, end) in keep {
            let start = start.min(len);
            if cursor < start {
                buf[cursor..start].fill(0);
            }
            cursor = cursor.max(end.min(len));
        }
        if cursor < len {
            buf[cursor..len].fill(0);
        }
    }

    fn write_constant_to_arena(&mut self, id: NodeId, dtype: DType, data: &[u8]) {
        match dtype {
            DType::F64
            | DType::F16
            | DType::BF16
            | DType::U8
            | DType::I8
            | DType::Bool
            | DType::I64
            | DType::I32
            | DType::U32 => {
                let off = self.arena.byte_offset(id);
                let buf = self.arena.raw_buf_mut();
                let n = buf.len().saturating_sub(off).min(data.len());
                buf[off..off + n].copy_from_slice(&data[..n]);
            }
            _ => {
                let buf = self.arena.slice_mut(id);
                let n_floats = data.len() / 4;
                let n = buf.len().min(n_floats);
                for i in 0..n {
                    let bytes = [
                        data[i * 4],
                        data[i * 4 + 1],
                        data[i * 4 + 2],
                        data[i * 4 + 3],
                    ];
                    buf[i] = f32::from_le_bytes(bytes);
                }
            }
        }
    }

    /// Direct-byte param upload — copies caller's bytes into the
    /// arena slot for the named param without any dtype conversion.
    /// Used by `set_param_typed` for dtypes that f32-widening would
    /// corrupt (F64). Caller is responsible for matching the param's
    /// declared graph dtype.
    fn set_param_bytes(&mut self, name: &str, data: &[u8], _dtype: rlx_ir::DType) {
        // Byte-backed params also live solely in the arena (no CPU-side copy).
        self.write_param_bytes_to_arena(name, data);
    }

    fn write_param_bytes_to_arena(&mut self, name: &str, data: &[u8]) {
        if let Some(&id) = self.param_ids.get(name)
            && self.arena.has_buffer(id)
        {
            let off = self.arena.byte_offset(id);
            let buf = self.arena.raw_buf_mut();
            debug_assert!(
                off + data.len() <= buf.len(),
                "set_param_bytes: '{name}' would overflow arena slot"
            );
            buf[off..off + data.len()].copy_from_slice(data);
        }
    }
}