rlx-coreml 0.2.17

// RLX — versatile ML compiler + runtime.
// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
// SPDX-License-Identifier: MIT OR Apache-2.0
//
// IR → CoreML ML Program (MIL) lowering. Pure data transformation: takes
// an RLX `Graph` plus baked parameter/constant data and produces a
// `proto::Model` ready to serialise into a `.mlpackage`. No FFI, so this
// builds and unit-tests on any host.

//! `activation` — extracted from the `mil` module for navigability (see `mod.rs`).

#![allow(unused_imports)]

use super::helpers::simple_op_flex;
use super::helpers::*;
use crate::proto;
use crate::{CoremlError, Result};
use rlx_ir::op::{Activation, CmpOp, MaskKind, ReduceOp};
use rlx_ir::quant::QuantScheme;
use rlx_ir::{DType, Dim, Graph, NodeId, Op, Shape};
use std::collections::HashMap;

use super::*;

impl<'a> LowerCtx<'a> {
    /// Lower an activation, composing the ones MIL has no direct op for.
    pub(crate) fn lower_activation(
        &mut self,
        id: NodeId,
        act: Activation,
        out_name: &str,
    ) -> Result<()> {
        let node = self.graph.node(id);
        let x = self.val(node.inputs[0]);
        // (mil_type, optional ("param", value)) for the simple unary cases.
        let direct: Option<(&str, Vec<(&str, proto::Argument)>)> = match act {
            Activation::Relu => Some(("relu", vec![])),
            Activation::Sigmoid => Some(("sigmoid", vec![])),
            Activation::Tanh => Some(("tanh", vec![])),
            Activation::Exp => Some(("exp", vec![])),
            // MIL `log` is `log(x + epsilon)` and — like `rsqrt` — requires the
            // param explicitly (CoreML rejects the model otherwise). Use CoreML's
            // own default 1e-45 so the result is unperturbed (negligible for x≥1,
            // e.g. the log-sum-exp in softmax-cross-entropy that surfaced this).
            Activation::Log => Some(("log", vec![("epsilon", bind_value(scalar_f32(1e-45)))])),
            Activation::Sqrt => Some(("sqrt", vec![])),
            Activation::Rsqrt => Some(("rsqrt", vec![("epsilon", bind_value(scalar_f32(1e-12)))])),
            Activation::Abs => Some(("abs", vec![])),
            Activation::Sin => Some(("sin", vec![])),
            Activation::Cos => Some(("cos", vec![])),
            Activation::Tan => Some(("tan", vec![])),
            Activation::Atan => Some(("atan", vec![])),
            Activation::Round => Some(("round", vec![])),
            Activation::Gelu => Some(("gelu", vec![("mode", bind_value(scalar_str("EXACT")))])),
            Activation::GeluApprox => Some((
                "gelu",
                vec![("mode", bind_value(scalar_str("TANH_APPROXIMATION")))],
            )),
            // Composed below.
            Activation::Floor => Some(("floor", vec![])),
            Activation::Ceil => Some(("ceil", vec![])),
            Activation::Sign => Some(("sign", vec![])),
            Activation::Softplus => Some(("softplus", vec![])),
            Activation::Elu => Some(("elu", vec![("alpha", bind_value(scalar_f32(1.0)))])),
            Activation::Erf => Some(("erf", vec![])),
            Activation::Softsign => Some(("softsign", vec![])),
            Activation::Silu
            | Activation::Neg
            | Activation::Recip
            | Activation::HardSwish
            | Activation::HardSigmoid
            | Activation::Mish
            | Activation::LogSigmoid => None,
        };

        if let Some((ty, mut params)) = direct {
            let mut binds = vec![("x", bind_name(&x))];
            binds.append(&mut params);
            let op = self.simple_op(ty, out_name, &node.shape, binds)?;
            self.push_named(id, out_name.to_string(), op);
            return Ok(());
        }

        match act {
            // silu(x) = x * sigmoid(x)
            Activation::Silu => {
                let sig = format!("{out_name}_sig");
                let sig_op =
                    self.simple_op("sigmoid", &sig, &node.shape, vec![("x", bind_name(&x))])?;
                self.operations.push(sig_op);
                let op = self.simple_op(
                    "mul",
                    out_name,
                    &node.shape,
                    vec![("x", bind_name(&x)), ("y", bind_name(&sig))],
                )?;
                self.push_named(id, out_name.to_string(), op);
            }
            // neg(x) = mul(x, -1)
            Activation::Neg => {
                let op = self.simple_op(
                    "mul",
                    out_name,
                    &node.shape,
                    vec![("x", bind_name(&x)), ("y", bind_value(scalar_f32(-1.0)))],
                )?;
                self.push_named(id, out_name.to_string(), op);
            }
            Activation::Recip => {
                let op = self.simple_op(
                    "real_div",
                    out_name,
                    &node.shape,
                    vec![("x", bind_value(scalar_f32(1.0))), ("y", bind_name(&x))],
                )?;
                self.push_named(id, out_name.to_string(), op);
            }
            // mish(x) = x * tanh(softplus(x))
            Activation::Mish => {
                let sp = format!("{out_name}_softplus");
                let op =
                    self.simple_op("softplus", &sp, &node.shape, vec![("x", bind_name(&x))])?;
                self.operations.push(op);
                let th = format!("{out_name}_tanh");
                let op = self.simple_op("tanh", &th, &node.shape, vec![("x", bind_name(&sp))])?;
                self.operations.push(op);
                let op = self.simple_op(
                    "mul",
                    out_name,
                    &node.shape,
                    vec![("x", bind_name(&x)), ("y", bind_name(&th))],
                )?;
                self.push_named(id, out_name.to_string(), op);
            }
            // log_sigmoid(x) = -softplus(-x). Computing it as `log(sigmoid(x))`
            // instead underflows to -inf once sigmoid rounds to zero, which is
            // the whole reason the fused form exists.
            Activation::LogSigmoid => {
                let neg = format!("{out_name}_neg");
                let op = self.simple_op(
                    "mul",
                    &neg,
                    &node.shape,
                    vec![("x", bind_name(&x)), ("y", bind_value(scalar_f32(-1.0)))],
                )?;
                self.operations.push(op);
                let sp = format!("{out_name}_softplus");
                let op =
                    self.simple_op("softplus", &sp, &node.shape, vec![("x", bind_name(&neg))])?;
                self.operations.push(op);
                let op = self.simple_op(
                    "mul",
                    out_name,
                    &node.shape,
                    vec![("x", bind_name(&sp)), ("y", bind_value(scalar_f32(-1.0)))],
                )?;
                self.push_named(id, out_name.to_string(), op);
            }
            // hard_sigmoid(x) = clip(x/6 + 0.5, 0, 1). MIL's own `sigmoid_hard`
            // takes alpha/beta with exactly this parameterisation.
            Activation::HardSigmoid => {
                let op = self.simple_op(
                    "sigmoid_hard",
                    out_name,
                    &node.shape,
                    vec![
                        ("x", bind_name(&x)),
                        ("alpha", bind_value(scalar_f32(1.0 / 6.0))),
                        ("beta", bind_value(scalar_f32(0.5))),
                    ],
                )?;
                self.push_named(id, out_name.to_string(), op);
            }
            // hard_swish(x) = x * hard_sigmoid(x)
            Activation::HardSwish => {
                let hs = format!("{out_name}_hard_sigmoid");
                let op = self.simple_op(
                    "sigmoid_hard",
                    &hs,
                    &node.shape,
                    vec![
                        ("x", bind_name(&x)),
                        ("alpha", bind_value(scalar_f32(1.0 / 6.0))),
                        ("beta", bind_value(scalar_f32(0.5))),
                    ],
                )?;
                self.operations.push(op);
                let op = self.simple_op(
                    "mul",
                    out_name,
                    &node.shape,
                    vec![("x", bind_name(&x)), ("y", bind_name(&hs))],
                )?;
                self.push_named(id, out_name.to_string(), op);
            }
            // Not `unreachable!`: the `direct` table above and this match are
            // two lists that have to stay in step, and they have drifted before
            // — `Mish` reached here and panicked inside the backend rather than
            // reporting an unsupported op. A new `Activation` variant must
            // produce a diagnosable error, not a crash.
            other => {
                return Err(CoremlError::Unsupported(format!(
                    "activation {other:?} has no MIL lowering"
                )));
            }
        }
        Ok(())
    }

    /// `out = a * b + c` as MIL `mul` then `add` (no native fma op).
    pub(crate) fn lower_fma(&mut self, id: NodeId, out_name: &str) -> Result<()> {
        let node = self.graph.node(id);
        let a = self.val(node.inputs[0]);
        let b = self.val(node.inputs[1]);
        let c = self.val(node.inputs[2]);
        let shape = node.shape.clone();
        let prod = format!("{out_name}_mul");
        self.emit(
            "mul",
            &prod,
            &shape,
            vec![("x", bind_name(&a)), ("y", bind_name(&b))],
        )?;
        self.emit(
            "add",
            out_name,
            &shape,
            vec![("x", bind_name(&prod)), ("y", bind_name(&c))],
        )?;
        self.names.insert(id.0, out_name.to_string());
        Ok(())
    }

    /// Apply an activation to a named intermediate (for fused epilogues).
    pub(crate) fn emit_activation(
        &mut self,
        x: &str,
        act: Option<Activation>,
        out_name: &str,
        shape: &Shape,
    ) -> Result<()> {
        let Some(act) = act else {
            // Identity: rename via add-zero would waste a node; use reshape no-op
            // only when names differ — otherwise alias.
            if x == out_name {
                return Ok(());
            }
            let dims: Vec<i32> = static_dims(shape)?.iter().map(|&d| d as i32).collect();
            self.emit(
                "reshape",
                out_name,
                shape,
                vec![("x", bind_name(x)), ("shape", bind_value(vec_i32(&dims)))],
            )?;
            return Ok(());
        };
        let direct: Option<(&str, Vec<(&str, proto::Argument)>)> = match act {
            Activation::Relu => Some(("relu", vec![])),
            Activation::Sigmoid => Some(("sigmoid", vec![])),
            Activation::Tanh => Some(("tanh", vec![])),
            Activation::Exp => Some(("exp", vec![])),
            Activation::Log => Some(("log", vec![("epsilon", bind_value(scalar_f32(1e-45)))])),
            Activation::Sqrt => Some(("sqrt", vec![])),
            Activation::Rsqrt => Some(("rsqrt", vec![("epsilon", bind_value(scalar_f32(1e-12)))])),
            Activation::Abs => Some(("abs", vec![])),
            Activation::Sin => Some(("sin", vec![])),
            Activation::Cos => Some(("cos", vec![])),
            Activation::Tan => Some(("tan", vec![])),
            Activation::Atan => Some(("atan", vec![])),
            Activation::Round => Some(("round", vec![])),
            Activation::Gelu => Some(("gelu", vec![("mode", bind_value(scalar_str("EXACT")))])),
            Activation::GeluApprox => Some((
                "gelu",
                vec![("mode", bind_value(scalar_str("TANH_APPROXIMATION")))],
            )),
            Activation::Floor => Some(("floor", vec![])),
            Activation::Ceil => Some(("ceil", vec![])),
            Activation::Sign => Some(("sign", vec![])),
            Activation::Softplus => Some(("softplus", vec![])),
            Activation::Elu => Some(("elu", vec![("alpha", bind_value(scalar_f32(1.0)))])),
            Activation::Erf => Some(("erf", vec![])),
            Activation::Softsign => Some(("softsign", vec![])),
            Activation::Silu
            | Activation::Neg
            | Activation::Recip
            | Activation::HardSwish
            | Activation::HardSigmoid
            | Activation::Mish
            | Activation::LogSigmoid => None,
        };
        if let Some((ty, mut params)) = direct {
            let mut binds = vec![("x", bind_name(x))];
            binds.append(&mut params);
            self.emit(ty, out_name, shape, binds)?;
            return Ok(());
        }
        match act {
            Activation::Silu => {
                let sig = format!("{out_name}_sig");
                self.emit("sigmoid", &sig, shape, vec![("x", bind_name(x))])?;
                self.emit(
                    "mul",
                    out_name,
                    shape,
                    vec![("x", bind_name(x)), ("y", bind_name(&sig))],
                )?;
            }
            Activation::Neg => {
                self.emit(
                    "mul",
                    out_name,
                    shape,
                    vec![("x", bind_name(x)), ("y", bind_value(scalar_f32(-1.0)))],
                )?;
            }
            Activation::Recip => {
                self.emit(
                    "real_div",
                    out_name,
                    shape,
                    vec![("x", bind_value(scalar_f32(1.0))), ("y", bind_name(x))],
                )?;
            }
            _ => unreachable!("handled above"),
        }
        Ok(())
    }
}