oxmera_tensor/

backend.rs

//! The backend seam: the op vocabulary every device implements, and the
//! registry that resolves a [`Device`] handle to an implementation.
//!
//! The traits live here (rather than a separate crate) so that `Tensor`'s
//! methods and `std::ops` overloads can dispatch through them without
//! violating the orphan rule; `oxmera-ops` re-exports this vocabulary.

use std::collections::HashMap;
use std::sync::{Arc, OnceLock, RwLock};

use oxmera_core::{Device, Error, Result, Shape};

use crate::tensor::Tensor;

/// Elementwise unary operations. Float (`f32`) tensors only.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
#[non_exhaustive]
pub enum UnaryOp {
    /// `-x`
    Neg,
    /// `e^x`
    Exp,
    /// `ln(x)`
    Ln,
    /// `|x|`
    Abs,
    /// `√x`
    Sqrt,
    /// `sin(x)`
    Sin,
    /// `cos(x)`
    Cos,
    /// `tanh(x)`
    Tanh,
    /// `max(x, 0)`
    Relu,
    /// GELU with the tanh approximation.
    Gelu,
    /// `1 / (1 + e^-x)`
    Sigmoid,
}

impl UnaryOp {
    /// Stable lowercase name, used for kernel lookup and error messages.
    pub fn name(self) -> &'static str {
        match self {
            UnaryOp::Neg => "neg",
            UnaryOp::Exp => "exp",
            UnaryOp::Ln => "ln",
            UnaryOp::Abs => "abs",
            UnaryOp::Sqrt => "sqrt",
            UnaryOp::Sin => "sin",
            UnaryOp::Cos => "cos",
            UnaryOp::Tanh => "tanh",
            UnaryOp::Relu => "relu",
            UnaryOp::Gelu => "gelu",
            UnaryOp::Sigmoid => "sigmoid",
        }
    }

    /// Every unary op, for exhaustive backend tests.
    pub fn all() -> &'static [UnaryOp] {
        &[
            UnaryOp::Neg,
            UnaryOp::Exp,
            UnaryOp::Ln,
            UnaryOp::Abs,
            UnaryOp::Sqrt,
            UnaryOp::Sin,
            UnaryOp::Cos,
            UnaryOp::Tanh,
            UnaryOp::Relu,
            UnaryOp::Gelu,
            UnaryOp::Sigmoid,
        ]
    }

    /// Apply the op to one scalar — the CPU reference semantics every
    /// backend must reproduce.
    pub fn eval(self, x: f32) -> f32 {
        match self {
            UnaryOp::Neg => -x,
            UnaryOp::Exp => x.exp(),
            UnaryOp::Ln => x.ln(),
            UnaryOp::Abs => x.abs(),
            UnaryOp::Sqrt => x.sqrt(),
            UnaryOp::Sin => x.sin(),
            UnaryOp::Cos => x.cos(),
            UnaryOp::Tanh => x.tanh(),
            UnaryOp::Relu => x.max(0.0),
            UnaryOp::Gelu => {
                const SQRT_2_OVER_PI: f32 = 0.797_884_6;
                0.5 * x * (1.0 + (SQRT_2_OVER_PI * (x + 0.044_715 * x * x * x)).tanh())
            }
            UnaryOp::Sigmoid => 1.0 / (1.0 + (-x).exp()),
        }
    }
}

/// Elementwise binary operations with NumPy broadcasting. `f32` only.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
#[non_exhaustive]
pub enum BinaryOp {
    /// `a + b`
    Add,
    /// `a - b`
    Sub,
    /// `a * b`
    Mul,
    /// `a / b`
    Div,
    /// `a ^ b`
    Pow,
    /// `max(a, b)`
    Maximum,
    /// `min(a, b)`
    Minimum,
    /// `a > b` as a 0.0/1.0 mask. Not differentiable.
    Gt,
    /// `a == b` as a 0.0/1.0 mask. Not differentiable.
    Eq,
}

impl BinaryOp {
    /// Stable lowercase name, used for kernel lookup and error messages.
    pub fn name(self) -> &'static str {
        match self {
            BinaryOp::Add => "add",
            BinaryOp::Sub => "sub",
            BinaryOp::Mul => "mul",
            BinaryOp::Div => "div",
            BinaryOp::Pow => "pow",
            BinaryOp::Maximum => "maximum",
            BinaryOp::Minimum => "minimum",
            BinaryOp::Gt => "gt",
            BinaryOp::Eq => "eq",
        }
    }

    /// Every binary op, for exhaustive backend tests.
    pub fn all() -> &'static [BinaryOp] {
        &[
            BinaryOp::Add,
            BinaryOp::Sub,
            BinaryOp::Mul,
            BinaryOp::Div,
            BinaryOp::Pow,
            BinaryOp::Maximum,
            BinaryOp::Minimum,
            BinaryOp::Gt,
            BinaryOp::Eq,
        ]
    }

    /// Apply the op to one scalar pair — the reference semantics.
    pub fn eval(self, a: f32, b: f32) -> f32 {
        match self {
            BinaryOp::Add => a + b,
            BinaryOp::Sub => a - b,
            BinaryOp::Mul => a * b,
            BinaryOp::Div => a / b,
            BinaryOp::Pow => a.powf(b),
            BinaryOp::Maximum => a.max(b),
            BinaryOp::Minimum => a.min(b),
            BinaryOp::Gt => f32::from(a > b),
            BinaryOp::Eq => f32::from(a == b),
        }
    }
}

/// Reductions along axes. `f32` only.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
#[non_exhaustive]
pub enum ReduceOp {
    /// Sum of the reduced elements.
    Sum,
    /// Maximum of the reduced elements.
    Max,
    /// Minimum of the reduced elements.
    Min,
}

impl ReduceOp {
    /// Stable lowercase name, used for kernel lookup and error messages.
    pub fn name(self) -> &'static str {
        match self {
            ReduceOp::Sum => "sum",
            ReduceOp::Max => "max",
            ReduceOp::Min => "min",
        }
    }

    /// The identity element the reduction starts from.
    pub fn identity(self) -> f32 {
        match self {
            ReduceOp::Sum => 0.0,
            ReduceOp::Max => f32::NEG_INFINITY,
            ReduceOp::Min => f32::INFINITY,
        }
    }

    /// Combine an accumulator with one element.
    pub fn combine(self, acc: f32, x: f32) -> f32 {
        match self {
            ReduceOp::Sum => acc + x,
            ReduceOp::Max => acc.max(x),
            ReduceOp::Min => acc.min(x),
        }
    }
}

/// The shape contract of a matmul, resolved once so every backend agrees.
///
/// Operands are rank 2 (`[m, k]`) or rank 3 (`[b, m, k]`). Batch
/// dimensions broadcast: a batch of 1 pairs with a batch of `n`, and a
/// rank-2 operand is treated as batch 1. The output is rank 2 only when
/// both inputs are; otherwise `[batch, m, n]`.
///
/// Batch strides are in elements of the operand's *contiguous* data and
/// are 0 for a broadcast operand, so no backend has to materialize the
/// broadcast — batch `i` of `a` starts at `i * a_batch_stride`.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct MatmulPlan {
    /// Output batch count (1 for a rank-2 result).
    pub batch: usize,
    /// Rows of `a` and of the output.
    pub m: usize,
    /// Shared dimension.
    pub k: usize,
    /// Columns of `b` and of the output.
    pub n: usize,
    /// Element offset between consecutive batches of `a` (0 = broadcast).
    pub a_batch_stride: usize,
    /// Element offset between consecutive batches of `b` (0 = broadcast).
    pub b_batch_stride: usize,
    /// The result shape.
    pub out_shape: Shape,
}

/// Resolve the matmul contract for two shapes, or the typed error a
/// caller sees for incompatible operands.
pub fn plan_matmul(a: &Shape, b: &Shape) -> Result<MatmulPlan> {
    let (ad, bd) = (a.dims(), b.dims());
    let (ba, m, k) = match ad {
        [m, k] => (1usize, *m, *k),
        [b, m, k] => (*b, *m, *k),
        _ => {
            return Err(Error::InvalidArgument {
                op: "matmul",
                detail: format!(
                    "supported ranks are 2 and 3 (batched); got {}x{}",
                    ad.len(),
                    bd.len()
                ),
            });
        }
    };
    let (bb, kb, n) = match bd {
        [k, n] => (1usize, *k, *n),
        [b, k, n] => (*b, *k, *n),
        _ => {
            return Err(Error::InvalidArgument {
                op: "matmul",
                detail: format!(
                    "supported ranks are 2 and 3 (batched); got {}x{}",
                    ad.len(),
                    bd.len()
                ),
            });
        }
    };
    if k != kb {
        return Err(Error::ShapeMismatch {
            expected: if bd.len() == 3 {
                Shape::from([bb, k, n])
            } else {
                Shape::from([k, n])
            },
            got: b.clone(),
            op: "matmul",
        });
    }
    let batch = match (ba, bb) {
        (x, y) if x == y => x,
        (1, y) => y,
        (x, 1) => x,
        // Batch dimensions that are neither equal nor 1 cannot broadcast.
        _ => {
            return Err(Error::BroadcastIncompatible {
                lhs: a.clone(),
                rhs: b.clone(),
            });
        }
    };
    let out_shape = if ad.len() == 2 && bd.len() == 2 {
        Shape::from([m, n])
    } else {
        Shape::from([batch, m, n])
    };
    Ok(MatmulPlan {
        batch,
        m,
        k,
        n,
        a_batch_stride: if ba == 1 { 0 } else { m * k },
        b_batch_stride: if bb == 1 { 0 } else { k * n },
        out_shape,
    })
}

/// One Adam/AdamW update for a single parameter, for backends that fuse
/// the ~12 elementwise ops of the composite step into one launch
/// ([`Backend::adam_step`]). Tensors are `f32` on the backend's device;
/// `m`/`v` are `None` on the first step.
#[derive(Debug, Clone, Copy)]
pub struct AdamStep<'a> {
    /// Current parameter value.
    pub param: &'a Tensor,
    /// Its gradient.
    pub grad: &'a Tensor,
    /// First-moment state, if any step has run.
    pub m: Option<&'a Tensor>,
    /// Second-moment state, if any step has run.
    pub v: Option<&'a Tensor>,
    /// Learning rate.
    pub lr: f32,
    /// β₁.
    pub beta1: f32,
    /// β₂.
    pub beta2: f32,
    /// ε added to the denominator.
    pub eps: f32,
    /// Weight decay; `0.0` disables it.
    pub weight_decay: f32,
    /// `true` for AdamW (decay the weights), `false` for Adam (add to the
    /// gradient).
    pub decoupled: bool,
    /// `1 - β₁ᵗ` for this step.
    pub bias_correction1: f32,
    /// `1 - β₂ᵗ` for this step.
    pub bias_correction2: f32,
}

/// A complete backend: every primitive the tensor method layer dispatches.
///
/// Composite operations (mean, softmax, losses, convolution, …) are built
/// from these primitives device-generically; only what is listed here is
/// implemented per device.
pub trait Backend: Send + Sync {
    /// The device this backend serves.
    fn device(&self) -> Device;

    /// A short stable name for reports and `oxmera doctor`.
    fn name(&self) -> &'static str;

    /// Elementwise unary op over a (possibly strided) `f32` tensor,
    /// producing a fresh contiguous tensor of the same shape.
    fn unary(&self, op: UnaryOp, a: &Tensor) -> Result<Tensor>;

    /// Elementwise binary op with broadcasting, producing a fresh
    /// contiguous tensor of the broadcast shape.
    fn binary(&self, op: BinaryOp, a: &Tensor, b: &Tensor) -> Result<Tensor>;

    /// Matrix product: rank-2 `[m, k] x [k, n] -> [m, n]`, or batched
    /// rank-3 `[b, m, k] x [b, k, n] -> [b, m, n]`.
    fn matmul(&self, a: &Tensor, b: &Tensor) -> Result<Tensor>;

    /// Reduce over `axes` (sorted, deduplicated by the caller; empty means
    /// all axes). `keepdim` keeps reduced axes as size 1.
    fn reduce(&self, op: ReduceOp, a: &Tensor, axes: &[usize], keepdim: bool) -> Result<Tensor>;

    /// Index of the maximum along `dim`, as an `I64` tensor.
    fn argmax(&self, a: &Tensor, dim: usize, keepdim: bool) -> Result<Tensor>;

    /// A fresh contiguous tensor with the same logical elements.
    fn contiguous(&self, a: &Tensor) -> Result<Tensor>;

    /// Download to a contiguous CPU tensor.
    fn download(&self, a: &Tensor) -> Result<Tensor>;

    /// Upload a contiguous CPU tensor to this backend's device.
    fn upload(&self, a: &Tensor) -> Result<Tensor>;

    /// Rows of `a` along `dim` selected by `indices` (`I64`).
    ///
    /// Backends may return [`Error::NotImplemented`]; the method layer
    /// then falls back to the CPU backend with a device round-trip.
    fn index_select(&self, a: &Tensor, dim: usize, indices: &Tensor) -> Result<Tensor> {
        let _ = (dim, indices);
        Err(Error::NotImplemented {
            op: "index_select",
            detail: format!("backend {}", a.device().kind_name()),
        })
    }

    /// `out[indices[i]] += src[i]` along `dim`, on a fresh copy of `a`.
    ///
    /// Same fallback contract as [`Backend::index_select`].
    fn index_add(&self, a: &Tensor, dim: usize, indices: &Tensor, src: &Tensor) -> Result<Tensor> {
        let _ = (dim, indices, src);
        Err(Error::NotImplemented {
            op: "index_add",
            detail: format!("backend {}", a.device().kind_name()),
        })
    }

    /// Lower-triangular Cholesky factors of a `[.., n, n]` batch of SPD
    /// matrices (lower triangle read; a non-PD matrix is a typed error).
    ///
    /// Same fallback contract as [`Backend::index_select`].
    fn cholesky(&self, a: &Tensor) -> Result<Tensor> {
        Err(Error::NotImplemented {
            op: "cholesky",
            detail: format!("backend {}", a.device().kind_name()),
        })
    }

    /// Symmetric eigen-decomposition of a `[.., n, n]` batch: eigenvalues
    /// ascending (`[.., n]`) and eigenvectors as columns (`[.., n, n]`).
    ///
    /// Same fallback contract as [`Backend::index_select`].
    fn eigh(&self, a: &Tensor) -> Result<(Tensor, Tensor)> {
        Err(Error::NotImplemented {
            op: "eigh",
            detail: format!("backend {}", a.device().kind_name()),
        })
    }

    /// One fused Adam/AdamW update: returns the new `(param, m, v)`.
    ///
    /// The optimizer calls this for parameters on a non-CPU device and
    /// falls back to the composite elementwise step when the backend
    /// declines, so the observable update is the same either way (within
    /// `f32` rounding of the same formula).
    fn adam_step(&self, step: &AdamStep<'_>) -> Result<(Tensor, Tensor, Tensor)> {
        Err(Error::NotImplemented {
            op: "adam_step",
            detail: format!("backend {}", step.param.device().kind_name()),
        })
    }
}

type Registry = RwLock<HashMap<Device, Arc<dyn Backend>>>;

fn registry() -> &'static Registry {
    static REGISTRY: OnceLock<Registry> = OnceLock::new();
    REGISTRY.get_or_init(|| RwLock::new(HashMap::new()))
}

/// Register a backend for its device, replacing any previous registration.
///
/// Backend crates call this from their load-time constructors; linking a
/// backend crate is what makes its device usable.
pub fn register_backend(backend: Arc<dyn Backend>) {
    registry()
        .write()
        .expect("backend registry poisoned")
        .insert(backend.device(), backend);
}

/// The backend serving `device`, or a typed error when none is registered.
///
/// The CPU backend is registered lazily on first use, so it can never be
/// unavailable; GPU backends register at load time or via
/// `oxmera_runtime::init()`.
pub fn backend_for(device: Device) -> Result<Arc<dyn Backend>> {
    let found = registry()
        .read()
        .expect("backend registry poisoned")
        .get(&device)
        .cloned();
    match found {
        Some(b) => Ok(b),
        None if device == Device::Cpu => {
            crate::cpu::register();
            registry()
                .read()
                .expect("backend registry poisoned")
                .get(&device)
                .cloned()
                .ok_or(Error::BackendUnavailable { device })
        }
        None => Err(Error::BackendUnavailable { device }),
    }
}

/// Every registered device, sorted for stable reporting.
pub fn registered_devices() -> Vec<Device> {
    let mut devices: Vec<Device> = registry()
        .read()
        .expect("backend registry poisoned")
        .keys()
        .copied()
        .collect();
    devices.sort_by_key(|d| (d.kind_name(), device_index(*d)));
    devices
}

fn device_index(d: Device) -> usize {
    match d {
        Device::Cpu => 0,
        Device::Metal { index } | Device::Cuda { index } => index,
        _ => 0,
    }
}