tokitai-operator 0.1.0

//! Normalization and activation operators.
//!
//! These are the activation and normalization ops that are table
//! stakes for any neural-network-shaped workload. Most of them are
//! pointwise scalar functions; `Softmax` and `LayerNorm` reduce over
//! the last axis only (no full config object yet).
//!
//! i64 design notes:
//!
//! On the i64 scalar backend these ops are *approximations* of their
//! floating-point counterparts. ReLU is exact (max(0, x)). The
//! Sigmoid, Tanh, GELU ops scale their result by 1_000_000 (so the
//! output is micro-sigmoid / micro-tanh in the range roughly
//! [-1_000_000, 1_000_000] for Tanh and GELU, and [0, 1_000_000]
//! for Sigmoid). Softmax uses a fixed-point exp2 lookup and scales
//! the result by 1_000_000 to give a probability in [0, 1_000_000].
//! LayerNorm uses integer mean and stddev with floor division.
//!
//! These scaling choices are recorded in the operator-level
//! `provided_contracts` as `Approximate` (Deterministic + Approximate)
//! so the planner emits the correct obligation text for downstream
//! consumers that need floating-point fidelity.
//!
//! The CPU scalar backend in `src/backend/cpu.rs` uses the same
//! scale factor (named `NN_APPROX_SCALE` there). Any consumer that
//! wants the floating-point value should divide by 1_000_000.

use crate::domain::{Claim, Contract, ContractId, ContractSet, Evidence, Scope};
use crate::object::{Dim, ObjectKind, ObjectMeta, Shape};
use crate::{Error, Result};

use super::{LayerBehavior, OpInput, OpOutput, OpSignature, Operator};

fn unary_signature() -> OpSignature {
    OpSignature {
        inputs: vec![OpInput {
            name: "input".to_string(),
        }],
        outputs: vec![OpOutput {
            name: "out".to_string(),
        }],
    }
}

fn tensor_input_check(op: &str, inputs: &[ObjectMeta], expected: usize) -> Result<()> {
    if inputs.len() != expected {
        return Err(Error::operator(format!(
            "{op} expects {expected} input(s), got {}",
            inputs.len()
        )));
    }
    for (i, m) in inputs.iter().enumerate() {
        if m.object_kind != ObjectKind::Tensor {
            return Err(Error::operator(format!(
                "{op} only supports tensor inputs, input {i} is {:?}",
                m.object_kind
            )));
        }
    }
    Ok(())
}

/// Deterministic + Exact contracts, used for ops that have no
/// floating-point approximation (ReLU on i64 is a clamp; Softmax /
/// LayerNorm declare Approximate because of i64 scaling).
fn deterministic_exact_contracts(scope: Scope) -> ContractSet {
    ContractSet::from_iter([
        Contract::new(
            ContractId(40),
            Claim::Deterministic,
            scope.clone(),
            Evidence::Axiom,
        ),
        Contract::new(ContractId(41), Claim::Exact, scope, Evidence::Axiom),
    ])
}

/// Deterministic + Approximate contracts, used for the
/// floating-point-style i64 ops. The approximation comes from the
/// fixed-point scaling at the i64 scalar backend.
fn deterministic_approximate_contracts(scope: Scope) -> ContractSet {
    ContractSet::from_iter([
        Contract::new(
            ContractId(42),
            Claim::Deterministic,
            scope.clone(),
            Evidence::Axiom,
        ),
        Contract::new(ContractId(43), Claim::Approximate, scope, Evidence::Axiom),
    ])
}

// ---------------------------------------------------------------------------
// ReLU: out = max(0, x). Exact on i64.
// ---------------------------------------------------------------------------

/// ReLU: `out = max(0, x)`, applied element-wise.
#[derive(Debug, Clone, Copy, Default)]
pub struct ReluOp;

impl Operator for ReluOp {
    fn name(&self) -> &'static str {
        "relu"
    }

    fn signature(&self) -> OpSignature {
        unary_signature()
    }

    fn infer(&self, inputs: &[ObjectMeta]) -> Result<Vec<ObjectMeta>> {
        tensor_input_check(self.name(), inputs, 1)?;
        Ok(vec![inputs[0].clone()])
    }

    fn required_contracts(&self) -> ContractSet {
        ContractSet::new()
    }
    fn provided_contracts(&self) -> ContractSet {
        deterministic_exact_contracts(Scope::Operator(self.name().to_string()))
    }
    fn layer_behavior(&self) -> Vec<LayerBehavior> {
        vec![LayerBehavior::Pointwise]
    }
}

// ---------------------------------------------------------------------------
// Sigmoid: out = 1 / (1 + exp(-x)), scaled by 1_000_000.
// ---------------------------------------------------------------------------

/// Sigmoid: `out = sigmoid(x)`, scaled by 1_000_000 so the
/// output lives in `[0, 1_000_000]`. The fixed-point polynomial
/// approximation is monotonic, saturates for large |x|, and
/// evaluates to 500_000 at x = 0.
#[derive(Debug, Clone, Copy, Default)]
pub struct SigmoidOp;

impl Operator for SigmoidOp {
    fn name(&self) -> &'static str {
        "sigmoid"
    }

    fn signature(&self) -> OpSignature {
        unary_signature()
    }

    fn infer(&self, inputs: &[ObjectMeta]) -> Result<Vec<ObjectMeta>> {
        tensor_input_check(self.name(), inputs, 1)?;
        Ok(vec![inputs[0].clone()])
    }

    fn required_contracts(&self) -> ContractSet {
        ContractSet::new()
    }
    fn provided_contracts(&self) -> ContractSet {
        deterministic_approximate_contracts(Scope::Operator(self.name().to_string()))
    }
    fn layer_behavior(&self) -> Vec<LayerBehavior> {
        vec![LayerBehavior::Pointwise]
    }
}

// ---------------------------------------------------------------------------
// Tanh: out = tanh(x), scaled by 1_000_000.
// ---------------------------------------------------------------------------

/// Tanh: `out = tanh(x)`, scaled by 1_000_000 so the output
/// lives in roughly `[-1_000_000, 1_000_000]`. The fixed-point
/// approximation saturates for large |x|.
#[derive(Debug, Clone, Copy, Default)]
pub struct TanhOp;

impl Operator for TanhOp {
    fn name(&self) -> &'static str {
        "tanh"
    }

    fn signature(&self) -> OpSignature {
        unary_signature()
    }

    fn infer(&self, inputs: &[ObjectMeta]) -> Result<Vec<ObjectMeta>> {
        tensor_input_check(self.name(), inputs, 1)?;
        Ok(vec![inputs[0].clone()])
    }

    fn required_contracts(&self) -> ContractSet {
        ContractSet::new()
    }
    fn provided_contracts(&self) -> ContractSet {
        deterministic_approximate_contracts(Scope::Operator(self.name().to_string()))
    }
    fn layer_behavior(&self) -> Vec<LayerBehavior> {
        vec![LayerBehavior::Pointwise]
    }
}

// ---------------------------------------------------------------------------
// GELU: out = 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3))).
// ---------------------------------------------------------------------------

/// GELU: Gaussian Error Linear Unit approximation using the
/// `0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))` form.
/// Output is scaled by 1_000_000; the final integer is
/// approximately `gelu(x) * 1_000_000`.
#[derive(Debug, Clone, Copy, Default)]
pub struct GeluOp;

impl Operator for GeluOp {
    fn name(&self) -> &'static str {
        "gelu"
    }

    fn signature(&self) -> OpSignature {
        unary_signature()
    }

    fn infer(&self, inputs: &[ObjectMeta]) -> Result<Vec<ObjectMeta>> {
        tensor_input_check(self.name(), inputs, 1)?;
        Ok(vec![inputs[0].clone()])
    }

    fn required_contracts(&self) -> ContractSet {
        ContractSet::new()
    }
    fn provided_contracts(&self) -> ContractSet {
        deterministic_approximate_contracts(Scope::Operator(self.name().to_string()))
    }
    fn layer_behavior(&self) -> Vec<LayerBehavior> {
        vec![LayerBehavior::Pointwise]
    }
}

// ---------------------------------------------------------------------------
// Softmax: out = exp(x - max) / sum(exp(x - max)), last axis only.
// ---------------------------------------------------------------------------

/// Softmax over the last axis. The output shape equals the input
/// shape; per-row (last-axis) the values are normalized so the row
/// sums to 1_000_000. Uses a fixed-point exp2 lookup so the
/// computation stays in i64.
#[derive(Debug, Clone, Copy, Default)]
pub struct SoftmaxOp;

impl Operator for SoftmaxOp {
    fn name(&self) -> &'static str {
        "softmax"
    }

    fn signature(&self) -> OpSignature {
        unary_signature()
    }

    fn infer(&self, inputs: &[ObjectMeta]) -> Result<Vec<ObjectMeta>> {
        tensor_input_check(self.name(), inputs, 1)?;
        Ok(vec![inputs[0].clone()])
    }

    fn required_contracts(&self) -> ContractSet {
        ContractSet::new()
    }
    fn provided_contracts(&self) -> ContractSet {
        deterministic_approximate_contracts(Scope::Operator(self.name().to_string()))
    }
    fn layer_behavior(&self) -> Vec<LayerBehavior> {
        vec![LayerBehavior::Pointwise, LayerBehavior::Global]
    }
}

// ---------------------------------------------------------------------------
// LayerNorm: out = (x - mean) / stddev * gamma + beta, last axis.
// ---------------------------------------------------------------------------

/// LayerNorm over the last axis. `gamma` defaults to 1 and `beta`
/// defaults to 0; the lowering accepts them as optional inputs
/// (gamma and beta are 1-D tensors of size `dim_size`, or omitted).
#[derive(Debug, Clone, Copy, Default)]
pub struct LayerNormOp;

impl Operator for LayerNormOp {
    fn name(&self) -> &'static str {
        "layer_norm"
    }

    fn signature(&self) -> OpSignature {
        // We declare 3 inputs: input, gamma, beta. Tests can pass
        // `gamma` and `beta` as the same input tensor twice when no
        // affine transform is desired; the CPU lowering will detect
        // that and treat them as identity.
        OpSignature {
            inputs: vec![
                OpInput {
                    name: "input".to_string(),
                },
                OpInput {
                    name: "gamma".to_string(),
                },
                OpInput {
                    name: "beta".to_string(),
                },
            ],
            outputs: vec![OpOutput {
                name: "out".to_string(),
            }],
        }
    }

    fn infer(&self, inputs: &[ObjectMeta]) -> Result<Vec<ObjectMeta>> {
        tensor_input_check(self.name(), inputs, 3)?;
        let rank = inputs[0].shape.rank();
        if rank == 0 {
            return Err(Error::operator(
                "layer_norm requires at least 1-D input (last axis is the normalization axis)",
            ));
        }
        // The output shape equals the input shape (LayerNorm is per-row).
        let mut output = inputs[0].clone();
        output.shape = Shape::new(
            (0..rank)
                .map(|i| inputs[0].shape.dims[i].clone())
                .collect::<Vec<Dim>>(),
        );
        Ok(vec![output])
    }

    fn required_contracts(&self) -> ContractSet {
        ContractSet::new()
    }
    fn provided_contracts(&self) -> ContractSet {
        deterministic_approximate_contracts(Scope::Operator(self.name().to_string()))
    }
    fn layer_behavior(&self) -> Vec<LayerBehavior> {
        vec![LayerBehavior::Pointwise, LayerBehavior::Global]
    }
}