rlx-coreml 0.2.19

// RLX — versatile ML compiler + runtime.
// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
// SPDX-License-Identifier: MIT OR Apache-2.0
//
// Split mixed host/CoreML graphs into alternating host and MIL segments.

use std::collections::{HashMap, HashSet};

use rlx_ir::{DType, Graph, NodeId, Op, Shape};

use crate::host_exec::is_host_node;
use crate::{CoremlError, Result};

/// One MIL subgraph plus synthetic inputs wired from host tensors.
#[derive(Debug, Clone)]
pub struct MilSegment {
    pub graph: Graph,
    pub extra_inputs: Vec<(String, Shape)>,
    /// Maps each MIL-local output node id → the ORIGINAL global node id. The MIL
    /// subgraph renumbers nodes locally, so a global output `v80` may become a
    /// local `v10`; without this map the runner would write the buffer to
    /// `env[10]` instead of `env[80]`, losing the graph output (only the first
    /// segment escapes this because its local ids coincide with the global ones).
    pub out_local_to_global: HashMap<u32, u32>,
}

/// One execution step — host ops or a CoreML-compilable subgraph.
#[derive(Debug, Clone)]
pub enum Segment {
    Host(Vec<NodeId>),
    Mil(MilSegment),
}

/// How to run a graph that mixes host ops with MIL-lowerable compute.
#[derive(Debug, Clone)]
pub enum ExecutionPlan {
    /// Entire graph lowers to one CoreML model.
    MilOnly,
    /// Alternating host / CoreML segments in topological order.
    Segmented(Vec<Segment>),
}

/// True when the CoreML body has no compute (host-only graphs like `Sample`).
pub fn mil_body_is_trivial(graph: &Graph) -> bool {
    !graph.nodes().iter().any(|n| {
        !matches!(
            n.op,
            Op::Input { .. } | Op::Param { .. } | Op::Constant { .. }
        )
    })
}

/// Classify `graph` for hybrid execution.
pub fn plan_execution(graph: &Graph) -> Result<ExecutionPlan> {
    let segments = build_segments(graph)?;
    if rlx_ir::env::var("RLX_COREML_SEG_REPORT").as_deref() == Some("1") {
        let nh = segments
            .iter()
            .filter(|s| matches!(s, Segment::Host(_)))
            .count();
        let nm = segments
            .iter()
            .filter(|s| matches!(s, Segment::Mil(_)))
            .count();
        eprintln!(
            "[coreml] hybrid segments: {nh} host + {nm} MIL ({nm} CoreML models to compile+load)"
        );
    }
    if segments.len() == 1 {
        if let Segment::Mil(ref m) = segments[0] {
            if mil_body_is_trivial(&m.graph) {
                return Ok(ExecutionPlan::Segmented(segments));
            }
            return Ok(ExecutionPlan::MilOnly);
        }
    }
    if segments.is_empty() {
        return Ok(ExecutionPlan::MilOnly);
    }
    Ok(ExecutionPlan::Segmented(segments))
}

fn build_segments(graph: &Graph) -> Result<Vec<Segment>> {
    let order: Vec<NodeId> = graph.topo_order().collect();
    let mut segments: Vec<Segment> = Vec::new();
    let mut mil_nodes: Vec<NodeId> = Vec::new();
    let mut host_chain: Vec<NodeId> = Vec::new();

    for id in order {
        if is_host_node(graph, id) {
            flush_mil(graph, &mut segments, &mut mil_nodes)?;
            host_chain.push(id);
        } else {
            flush_host(&mut segments, &mut host_chain);
            mil_nodes.push(id);
        }
    }
    flush_host(&mut segments, &mut host_chain);
    flush_mil(graph, &mut segments, &mut mil_nodes)?;
    Ok(segments)
}

fn flush_host(segments: &mut Vec<Segment>, host_chain: &mut Vec<NodeId>) {
    if !host_chain.is_empty() {
        segments.push(Segment::Host(std::mem::take(host_chain)));
    }
}

fn flush_mil(
    graph: &Graph,
    segments: &mut Vec<Segment>,
    mil_nodes: &mut Vec<NodeId>,
) -> Result<()> {
    if mil_has_compute(graph, mil_nodes) {
        let (g, extra, out_map) = build_mil_subgraph(graph, mil_nodes)?;
        segments.push(Segment::Mil(MilSegment {
            graph: g,
            extra_inputs: extra,
            out_local_to_global: out_map,
        }));
    }
    mil_nodes.clear();
    Ok(())
}

fn mil_has_compute(graph: &Graph, mil_nodes: &[NodeId]) -> bool {
    mil_nodes.iter().any(|&id| {
        !matches!(
            graph.node(id).op,
            Op::Input { .. } | Op::Param { .. } | Op::Constant { .. }
        )
    })
}

fn build_mil_subgraph(
    graph: &Graph,
    mil_nodes: &[NodeId],
) -> Result<(Graph, Vec<(String, Shape)>, HashMap<u32, u32>)> {
    // MPSGraph requires every tensor's shape to be fully static and does not
    // report otherwise: it aborts the process from inside MPSRuntime
    // (`failed assertion 'shape for TensorData is not static'`, SIGABRT), which
    // no caller can catch and which takes any surrounding sweep down with it.
    // Refuse here instead, naming the node, so the caller gets an error and can
    // fall back to another backend.
    if let Some(id) = mil_nodes
        .iter()
        .copied()
        .find(|&id| !graph.node(id).shape.is_static())
    {
        let node = graph.node(id);
        return Err(CoremlError::Unsupported(format!(
            "dynamic shape {:?} on {:?} ({id:?}) — MPSGraph needs static shapes and \
             aborts the process rather than erroring; run this graph on another backend",
            node.shape.dims(),
            node.op.kind(),
        )));
    }
    let mil_set: HashSet<NodeId> = mil_nodes.iter().copied().collect();
    let mut g = Graph::new(format!("{}_coreml", graph.name));
    let mut map: HashMap<NodeId, NodeId> = HashMap::new();
    let mut extra_inputs: Vec<(String, Shape)> = Vec::new();

    let clone_leaf = |g: &mut Graph, map: &mut HashMap<NodeId, NodeId>, old: NodeId| -> NodeId {
        if let Some(&n) = map.get(&old) {
            return n;
        }
        let node = graph.node(old);
        let new_id = match &node.op {
            Op::Input { name } => g.input(name, node.shape.clone()),
            Op::Param { name } => g.param(name, node.shape.clone()),
            Op::Constant { data } => g.add_node(
                Op::Constant { data: data.clone() },
                vec![],
                node.shape.clone(),
            ),
            _ => unreachable!("clone_leaf on non-leaf"),
        };
        map.insert(old, new_id);
        new_id
    };

    // A plan can hand this segment an `Input` whose readers all live in other
    // segments (CLEF's tokenizer `d_output` seed). CoreML refuses to load such a
    // model ("Error in declaring input d_output with error -1"), so leave the
    // input out — unless it is the segment's only input: a model with no inputs
    // loads, but CoreML's E5 runtime throws ("Variable is not associated with a
    // name") and silently falls back to another engine.
    let unused_input = |id: NodeId| {
        matches!(graph.node(id).op, Op::Input { .. })
            && !graph.users(id).iter().any(|u| mil_set.contains(u))
    };
    let keeps_an_input = mil_nodes.iter().any(|&id| {
        let n = graph.node(id);
        (matches!(n.op, Op::Input { .. }) && !unused_input(id))
            || n.inputs.iter().any(|&i| {
                !mil_set.contains(&i)
                    && !matches!(graph.node(i).op, Op::Param { .. } | Op::Constant { .. })
            })
    });
    for &id in mil_nodes {
        if keeps_an_input && unused_input(id) {
            continue;
        }
        let node = graph.node(id);
        let mut new_inputs = Vec::with_capacity(node.inputs.len());
        for &inp in &node.inputs {
            let is_leaf = matches!(
                graph.node(inp).op,
                Op::Input { .. } | Op::Param { .. } | Op::Constant { .. }
            );
            if mil_set.contains(&inp) {
                new_inputs.push(
                    *map.get(&inp)
                        .ok_or_else(|| CoremlError::Runtime(format!("mil map missing {inp:?}")))?,
                );
            } else if is_leaf {
                new_inputs.push(clone_leaf(&mut g, &mut map, inp));
            } else {
                // Any external NON-leaf value — a host op's output OR a prior MIL
                // segment's output — is fed into this MIL model as a boundary
                // input, sourced from the global `env` (which every segment
                // writes its outputs to) via the `host_v{id}` naming the runner
                // resolves. Previously only `is_host_op` inputs took this path,
                // so a later MIL segment consuming an earlier MIL segment's
                // output hit `clone_leaf on non-leaf`.
                let name = format!("host_v{}", inp.0);
                if let std::collections::hash_map::Entry::Vacant(e) = map.entry(inp) {
                    let shape = graph.shape(inp).clone();
                    // The global `env` (and the runner) carry every value as f32,
                    // and MIL feature I/O only accepts F32/F16/I32/F64 — never Bool
                    // or I64. Declare the boundary input F32 and, when the true
                    // dtype differs, cast it back inside the MIL graph so
                    // downstream ops see the right type. (moss's cross-segment
                    // masks are Bool; its indices small enough to be exact in f32.)
                    if shape.dtype() == DType::F32 {
                        let nid = g.input(&name, shape.clone());
                        e.insert(nid);
                        extra_inputs.push((name, shape));
                    } else {
                        let f32_shape = shape.clone().with_dtype(DType::F32);
                        let in_nid = g.input(&name, f32_shape.clone());
                        let cast_nid =
                            g.add_node(Op::Cast { to: shape.dtype() }, vec![in_nid], shape);
                        e.insert(cast_nid);
                        extra_inputs.push((name, f32_shape));
                    }
                }
                new_inputs.push(map[&inp]);
            }
        }
        let new_id = g.append_node(node.op.clone(), new_inputs, node.shape.clone(), None);
        map.insert(id, new_id);
    }

    let outs = mil_segment_outputs(graph, &mil_set);
    let mapped: Vec<NodeId> = outs
        .iter()
        .map(|&oid| {
            let local = map.get(&oid).copied().ok_or_else(|| {
                CoremlError::Runtime(format!("missing mil output map for {oid:?}"))
            })?;
            // MIL feature outputs must be F32/F16/I32/F64, and the runner reads
            // them back into the f32 `env` — so cast any Bool/I64 output to F32.
            let d = graph.shape(oid).dtype();
            if d == DType::F32 {
                Ok(local)
            } else {
                let f32_shape = graph.shape(oid).clone().with_dtype(DType::F32);
                Ok(g.add_node(Op::Cast { to: DType::F32 }, vec![local], f32_shape))
            }
        })
        .collect::<Result<Vec<_>>>()?;
    if mapped.is_empty() {
        return Err(CoremlError::Unsupported(
            "empty MIL segment (no outputs)".into(),
        ));
    }
    // Reverse map: MIL-local output id → original global id, so the runner writes
    // each output buffer back to the correct global `env` slot.
    let out_local_to_global: HashMap<u32, u32> =
        outs.iter().zip(&mapped).map(|(g, l)| (l.0, g.0)).collect();
    g.set_outputs(mapped);
    Ok((g, extra_inputs, out_local_to_global))
}

fn mil_segment_outputs(graph: &Graph, mil_set: &HashSet<NodeId>) -> Vec<NodeId> {
    // A MIL segment must export both the final graph outputs it owns AND any of
    // its nodes consumed by a *later* segment (host or another MIL segment) — the
    // latter become that segment's boundary inputs, resolved through the global
    // `env`. (Previously an early return after collecting graph outputs dropped
    // the cross-segment-consumed nodes, so a later segment's `host_v{id}` input
    // had no producer in `env`.)
    //
    // Skip Input/Param/Constant leaves: `seed_leaf_env` already places them in
    // `env`, and exporting an Input as a MIL output reuses its feature name
    // (e.g. `speed`) for both the model input and output — CoreML then fails
    // load with "Error in declaring input <name> with error -1".
    let is_leaf = |id: NodeId| {
        matches!(
            graph.node(id).op,
            Op::Input { .. } | Op::Param { .. } | Op::Constant { .. }
        )
    };
    let mut outs: Vec<NodeId> = graph
        .outputs
        .iter()
        .filter(|o| mil_set.contains(o) && !is_leaf(**o))
        .copied()
        .collect();
    for &id in mil_set {
        if is_leaf(id) {
            continue;
        }
        for user in graph.users(id) {
            if !mil_set.contains(&user) {
                outs.push(id);
                break;
            }
        }
    }
    outs.sort_by_key(|id| id.0);
    outs.dedup();
    outs
}

/// Feed no `Reduce` directly from a `Transpose`.
///
/// CoreML's CPU engine (BNNS) mis-executes `transpose → reduce → transpose`:
/// it segfaults (an out-of-bounds read inside `BNNSGraphContextExecute_v2`) or
/// silently returns wrong sums — `[3,32,32]` summed over two axes came back off
/// by 98 on values near 144. Either transpose alone is fine. It reproduces with
/// coremltools and no rlx code (macOS 26.6.2), on every compute-unit setting.
/// rlx-autodiff's sum-to-shape emits exactly this chain for bias gradients, so
/// CLEF's training graphs crashed on CoreML.
///
/// A reduction commutes with a permutation: `reduce(transpose(x, p), A)` equals
/// `reduce(x, p[A])` with the surviving axes reordered by `p`. When `p` keeps
/// those axes in order, the reorder is a reshape; otherwise a transpose follows
/// the reduce, never precedes one, so the pattern cannot form. A transpose whose
/// every reader is a rewritten reduce is dropped.
pub fn hoist_transpose_out_of_reduce(g: Graph) -> Graph {
    let is_static = |s: &Shape| s.dims().iter().all(|d| matches!(d, rlx_ir::Dim::Static(_)));
    let rewritable = |g: &Graph, n: &rlx_ir::Node| -> bool {
        let Op::Reduce { axes, .. } = &n.op else {
            return false;
        };
        let t = g.node(n.inputs[0]);
        matches!(t.op, Op::Transpose { .. })
            && !axes.is_empty()
            && is_static(&n.shape)
            && is_static(&t.shape)
            && is_static(&g.node(t.inputs[0]).shape)
    };
    if !g.nodes().iter().any(|n| rewritable(&g, n)) {
        return g;
    }
    // Transposes that only feed rewritten reduces are dead afterwards.
    let mut readers: HashMap<NodeId, (usize, usize)> = HashMap::new(); // (all, rewritten)
    for n in g.nodes() {
        let rw = rewritable(&g, n);
        for &i in &n.inputs {
            let e = readers.entry(i).or_default();
            e.0 += 1;
            e.1 += usize::from(rw && i == n.inputs[0]);
        }
    }
    let outputs: HashSet<NodeId> = g.outputs.iter().copied().collect();
    let static_dims =
        |s: &Shape| -> Vec<usize> { s.dims().iter().map(|d| d.unwrap_static()).collect() };

    let mut out = Graph::new(g.name.clone());
    let mut id_map: HashMap<NodeId, NodeId> = HashMap::new();
    for node in g.nodes() {
        if matches!(node.op, Op::Transpose { .. })
            && !outputs.contains(&node.id)
            && readers
                .get(&node.id)
                .is_some_and(|&(all, rw)| all > 0 && all == rw)
        {
            continue;
        }
        let new_id = if rewritable(&g, node) {
            // `keep_dim` needs no case: the final reshape to the node's own shape
            // drops or keeps the reduced axes either way.
            let Op::Reduce { op, axes, .. } = &node.op else {
                unreachable!()
            };
            let t = g.node(node.inputs[0]);
            let Op::Transpose { perm } = &t.op else {
                unreachable!()
            };
            let x = id_map[&t.inputs[0]];
            let x_shape = g.node(t.inputs[0]).shape.clone();
            let dtype = node.shape.dtype();
            let mut x_axes: Vec<usize> = axes.iter().map(|&a| perm[a]).collect();
            x_axes.sort_unstable();
            let mut z_dims = static_dims(&x_shape);
            for &a in &x_axes {
                z_dims[a] = 1;
            }
            let z = out.add_node(
                Op::Reduce {
                    op: *op,
                    axes: x_axes,
                    keep_dim: true,
                },
                vec![x],
                Shape::new(&z_dims, dtype),
            );
            let kept: Vec<usize> = (0..perm.len())
                .filter(|i| !axes.contains(i))
                .map(|i| perm[i])
                .collect();
            let want = static_dims(&node.shape);
            let pre = if kept.windows(2).all(|w| w[0] < w[1]) {
                z
            } else {
                let w_dims: Vec<usize> = perm.iter().map(|&p| z_dims[p]).collect();
                out.add_node(
                    Op::Transpose { perm: perm.clone() },
                    vec![z],
                    Shape::new(&w_dims, dtype),
                )
            };
            if static_dims(&out.node(pre).shape) == want {
                pre
            } else {
                out.add_node(
                    Op::Reshape {
                        new_shape: want.iter().map(|&d| d as i64).collect(),
                    },
                    vec![pre],
                    node.shape.clone(),
                )
            }
        } else {
            let new_inputs: Vec<NodeId> = node.inputs.iter().map(|i| id_map[i]).collect();
            out.add_node(node.op.clone(), new_inputs, node.shape.clone())
        };
        id_map.insert(node.id, new_id);
    }
    out.set_outputs(g.outputs.iter().map(|i| id_map[i]).collect());
    out
}

/// Divide an un-folded batch norm by a reciprocal multiply instead.
///
/// CoreML fuses `conv → [+ bias] → sub(mean) → div(std) → mul(γ) → add(β)` —
/// an eval-mode BatchNorm written out elementwise over baked constants — into
/// the convolution, and BNNS then segfaults while initialising the program
/// (`EXC_BAD_ACCESS` in `libBNNS.dylib` on the `MLE5ProgramLibrary` lazy-init
/// queue, macOS 26.6.2). Whether it does depends on the shapes in no monotone
/// way: a 6×1 or 64×1 spatial kernel crashed, 65×1 did not, and 64×1 passed
/// with 8 filters but crashed with 64. Every one of the four elementwise ops is
/// needed, and so is the convolution in front; `div` → `mul(recip)` is enough
/// to stop it. LEAF's tokenizer at `dim_cnn = 8` hit it.
///
/// The rewrite applies to a `div(sub(c, k₁), k₂)` whose `c` is a convolution
/// (optionally plus a bias) and whose `k₁`, `k₂` depend on no graph input.
/// It changes that one quotient by at most an ulp.
pub fn recip_bn_div_after_conv(g: Graph) -> Graph {
    use rlx_ir::op::{Activation, BinaryOp};
    let dep = rlx_ir::head::input_dependent(&g);
    let indep = |id: NodeId| !dep[id.0 as usize];
    let is_conv = |id: NodeId| matches!(g.node(id).op, Op::Conv { .. });
    let conv_like = |id: NodeId| {
        let n = g.node(id);
        is_conv(id)
            || (matches!(n.op, Op::Binary(BinaryOp::Add))
                && ((is_conv(n.inputs[0]) && indep(n.inputs[1]))
                    || (is_conv(n.inputs[1]) && indep(n.inputs[0]))))
    };
    let rewritable = |n: &rlx_ir::Node| -> bool {
        if !matches!(n.op, Op::Binary(BinaryOp::Div)) || !indep(n.inputs[1]) {
            return false;
        }
        let s = g.node(n.inputs[0]);
        matches!(s.op, Op::Binary(BinaryOp::Sub)) && conv_like(s.inputs[0]) && indep(s.inputs[1])
    };
    if !g.nodes().iter().any(rewritable) {
        return g;
    }
    let mut out = Graph::new(g.name.clone());
    let mut id_map: HashMap<NodeId, NodeId> = HashMap::new();
    for node in g.nodes() {
        let new_inputs: Vec<NodeId> = node.inputs.iter().map(|i| id_map[i]).collect();
        let new_id = if rewritable(node) {
            let k = new_inputs[1];
            let ks = out.node(k).shape.clone();
            let inv = out.activation(Activation::Recip, k, ks);
            out.add_node(
                Op::Binary(BinaryOp::Mul),
                vec![new_inputs[0], inv],
                node.shape.clone(),
            )
        } else {
            out.add_node(node.op.clone(), new_inputs, node.shape.clone())
        };
        id_map.insert(node.id, new_id);
    }
    out.set_outputs(g.outputs.iter().map(|i| id_map[i]).collect());
    out
}

/// Expand fused ops claimed for coverage that CPU HostOp would Nop:
/// `FusedConvBiasAct`, `PartitionedConv`, `FusedTransformerLayer`.
pub fn lower_cpu_nop_fused_for_coreml(g: Graph) -> Graph {
    let needs = g.nodes().iter().any(|n| {
        matches!(
            n.op,
            Op::FusedConvBiasAct { .. }
                | Op::PartitionedConv { .. }
                | Op::FusedTransformerLayer { .. }
        )
    });
    if !needs {
        return g;
    }
    let mut out = Graph::new(g.name.clone());
    let mut id_map: HashMap<NodeId, NodeId> = HashMap::new();
    let nodes: Vec<rlx_ir::Node> = g.nodes().to_vec();
    for node in &nodes {
        let new_inputs: Vec<NodeId> = node.inputs.iter().map(|i| id_map[i]).collect();
        let new_id = match &node.op {
            Op::FusedConvBiasAct { .. }
            | Op::PartitionedConv { .. }
            | Op::FusedTransformerLayer { .. } => {
                inline_unfused_fused_op(&mut out, &node.op, &new_inputs, &node.shape)
            }
            _ => out.add_node(node.op.clone(), new_inputs, node.shape.clone()),
        };
        id_map.insert(node.id, new_id);
    }
    out.set_outputs(g.outputs.iter().map(|i| id_map[i]).collect());
    out
}

fn inline_unfused_fused_op(out: &mut Graph, op: &Op, inputs: &[NodeId], shape: &Shape) -> NodeId {
    let mut mini = Graph::new("coreml_unfuse");
    let mut mini_ins = Vec::with_capacity(inputs.len());
    for (i, &src) in inputs.iter().enumerate() {
        let sh = out.node(src).shape.clone();
        mini_ins.push(mini.append_node(
            Op::Input {
                name: format!("in{i}"),
            },
            vec![],
            sh,
            None,
        ));
    }
    let out_id = mini.append_node(op.clone(), mini_ins, shape.clone(), None);
    mini.set_outputs(vec![out_id]);
    let expanded = rlx_opt::unfuse_fused_for_autodiff(mini);
    let mut map: HashMap<NodeId, NodeId> = HashMap::new();
    for n in expanded.nodes() {
        if let Op::Input { name } = &n.op {
            if let Some(rest) = name.strip_prefix("in") {
                if let Ok(i) = rest.parse::<usize>() {
                    map.insert(n.id, inputs[i]);
                    continue;
                }
            }
        }
        let mapped_ins: Vec<NodeId> = n.inputs.iter().map(|i| map[i]).collect();
        let id = out.add_node(n.op.clone(), mapped_ins, n.shape.clone());
        map.insert(n.id, id);
    }
    map[&expanded.outputs[0]]
}