oxmera_tensor/
backend.rs

1//! The backend seam: the op vocabulary every device implements, and the
2//! registry that resolves a [`Device`] handle to an implementation.
3//!
4//! The traits live here (rather than a separate crate) so that `Tensor`'s
5//! methods and `std::ops` overloads can dispatch through them without
6//! violating the orphan rule; `oxmera-ops` re-exports this vocabulary.
7
8use std::collections::HashMap;
9use std::sync::{Arc, OnceLock, RwLock};
10
11use oxmera_core::{Device, Error, Result, Shape};
12
13use crate::tensor::Tensor;
14
15/// Elementwise unary operations. Float (`f32`) tensors only.
16#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
17#[non_exhaustive]
18pub enum UnaryOp {
19    /// `-x`
20    Neg,
21    /// `e^x`
22    Exp,
23    /// `ln(x)`
24    Ln,
25    /// `|x|`
26    Abs,
27    /// `√x`
28    Sqrt,
29    /// `sin(x)`
30    Sin,
31    /// `cos(x)`
32    Cos,
33    /// `tanh(x)`
34    Tanh,
35    /// `max(x, 0)`
36    Relu,
37    /// GELU with the tanh approximation.
38    Gelu,
39    /// `1 / (1 + e^-x)`
40    Sigmoid,
41}
42
43impl UnaryOp {
44    /// Stable lowercase name, used for kernel lookup and error messages.
45    pub fn name(self) -> &'static str {
46        match self {
47            UnaryOp::Neg => "neg",
48            UnaryOp::Exp => "exp",
49            UnaryOp::Ln => "ln",
50            UnaryOp::Abs => "abs",
51            UnaryOp::Sqrt => "sqrt",
52            UnaryOp::Sin => "sin",
53            UnaryOp::Cos => "cos",
54            UnaryOp::Tanh => "tanh",
55            UnaryOp::Relu => "relu",
56            UnaryOp::Gelu => "gelu",
57            UnaryOp::Sigmoid => "sigmoid",
58        }
59    }
60
61    /// Every unary op, for exhaustive backend tests.
62    pub fn all() -> &'static [UnaryOp] {
63        &[
64            UnaryOp::Neg,
65            UnaryOp::Exp,
66            UnaryOp::Ln,
67            UnaryOp::Abs,
68            UnaryOp::Sqrt,
69            UnaryOp::Sin,
70            UnaryOp::Cos,
71            UnaryOp::Tanh,
72            UnaryOp::Relu,
73            UnaryOp::Gelu,
74            UnaryOp::Sigmoid,
75        ]
76    }
77
78    /// Apply the op to one scalar — the CPU reference semantics every
79    /// backend must reproduce.
80    pub fn eval(self, x: f32) -> f32 {
81        match self {
82            UnaryOp::Neg => -x,
83            UnaryOp::Exp => x.exp(),
84            UnaryOp::Ln => x.ln(),
85            UnaryOp::Abs => x.abs(),
86            UnaryOp::Sqrt => x.sqrt(),
87            UnaryOp::Sin => x.sin(),
88            UnaryOp::Cos => x.cos(),
89            UnaryOp::Tanh => x.tanh(),
90            UnaryOp::Relu => x.max(0.0),
91            UnaryOp::Gelu => {
92                const SQRT_2_OVER_PI: f32 = 0.797_884_6;
93                0.5 * x * (1.0 + (SQRT_2_OVER_PI * (x + 0.044_715 * x * x * x)).tanh())
94            }
95            UnaryOp::Sigmoid => 1.0 / (1.0 + (-x).exp()),
96        }
97    }
98}
99
100/// Elementwise binary operations with NumPy broadcasting. `f32` only.
101#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
102#[non_exhaustive]
103pub enum BinaryOp {
104    /// `a + b`
105    Add,
106    /// `a - b`
107    Sub,
108    /// `a * b`
109    Mul,
110    /// `a / b`
111    Div,
112    /// `a ^ b`
113    Pow,
114    /// `max(a, b)`
115    Maximum,
116    /// `min(a, b)`
117    Minimum,
118    /// `a > b` as a 0.0/1.0 mask. Not differentiable.
119    Gt,
120    /// `a == b` as a 0.0/1.0 mask. Not differentiable.
121    Eq,
122}
123
124impl BinaryOp {
125    /// Stable lowercase name, used for kernel lookup and error messages.
126    pub fn name(self) -> &'static str {
127        match self {
128            BinaryOp::Add => "add",
129            BinaryOp::Sub => "sub",
130            BinaryOp::Mul => "mul",
131            BinaryOp::Div => "div",
132            BinaryOp::Pow => "pow",
133            BinaryOp::Maximum => "maximum",
134            BinaryOp::Minimum => "minimum",
135            BinaryOp::Gt => "gt",
136            BinaryOp::Eq => "eq",
137        }
138    }
139
140    /// Every binary op, for exhaustive backend tests.
141    pub fn all() -> &'static [BinaryOp] {
142        &[
143            BinaryOp::Add,
144            BinaryOp::Sub,
145            BinaryOp::Mul,
146            BinaryOp::Div,
147            BinaryOp::Pow,
148            BinaryOp::Maximum,
149            BinaryOp::Minimum,
150            BinaryOp::Gt,
151            BinaryOp::Eq,
152        ]
153    }
154
155    /// Apply the op to one scalar pair — the reference semantics.
156    pub fn eval(self, a: f32, b: f32) -> f32 {
157        match self {
158            BinaryOp::Add => a + b,
159            BinaryOp::Sub => a - b,
160            BinaryOp::Mul => a * b,
161            BinaryOp::Div => a / b,
162            BinaryOp::Pow => a.powf(b),
163            BinaryOp::Maximum => a.max(b),
164            BinaryOp::Minimum => a.min(b),
165            BinaryOp::Gt => f32::from(a > b),
166            BinaryOp::Eq => f32::from(a == b),
167        }
168    }
169}
170
171/// Reductions along axes. `f32` only.
172#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
173#[non_exhaustive]
174pub enum ReduceOp {
175    /// Sum of the reduced elements.
176    Sum,
177    /// Maximum of the reduced elements.
178    Max,
179    /// Minimum of the reduced elements.
180    Min,
181}
182
183impl ReduceOp {
184    /// Stable lowercase name, used for kernel lookup and error messages.
185    pub fn name(self) -> &'static str {
186        match self {
187            ReduceOp::Sum => "sum",
188            ReduceOp::Max => "max",
189            ReduceOp::Min => "min",
190        }
191    }
192
193    /// The identity element the reduction starts from.
194    pub fn identity(self) -> f32 {
195        match self {
196            ReduceOp::Sum => 0.0,
197            ReduceOp::Max => f32::NEG_INFINITY,
198            ReduceOp::Min => f32::INFINITY,
199        }
200    }
201
202    /// Combine an accumulator with one element.
203    pub fn combine(self, acc: f32, x: f32) -> f32 {
204        match self {
205            ReduceOp::Sum => acc + x,
206            ReduceOp::Max => acc.max(x),
207            ReduceOp::Min => acc.min(x),
208        }
209    }
210}
211
212/// The shape contract of a matmul, resolved once so every backend agrees.
213///
214/// Operands are rank 2 (`[m, k]`) or rank 3 (`[b, m, k]`). Batch
215/// dimensions broadcast: a batch of 1 pairs with a batch of `n`, and a
216/// rank-2 operand is treated as batch 1. The output is rank 2 only when
217/// both inputs are; otherwise `[batch, m, n]`.
218///
219/// Batch strides are in elements of the operand's *contiguous* data and
220/// are 0 for a broadcast operand, so no backend has to materialize the
221/// broadcast — batch `i` of `a` starts at `i * a_batch_stride`.
222#[derive(Debug, Clone, PartialEq, Eq)]
223pub struct MatmulPlan {
224    /// Output batch count (1 for a rank-2 result).
225    pub batch: usize,
226    /// Rows of `a` and of the output.
227    pub m: usize,
228    /// Shared dimension.
229    pub k: usize,
230    /// Columns of `b` and of the output.
231    pub n: usize,
232    /// Element offset between consecutive batches of `a` (0 = broadcast).
233    pub a_batch_stride: usize,
234    /// Element offset between consecutive batches of `b` (0 = broadcast).
235    pub b_batch_stride: usize,
236    /// The result shape.
237    pub out_shape: Shape,
238}
239
240/// Resolve the matmul contract for two shapes, or the typed error a
241/// caller sees for incompatible operands.
242pub fn plan_matmul(a: &Shape, b: &Shape) -> Result<MatmulPlan> {
243    let (ad, bd) = (a.dims(), b.dims());
244    let (ba, m, k) = match ad {
245        [m, k] => (1usize, *m, *k),
246        [b, m, k] => (*b, *m, *k),
247        _ => {
248            return Err(Error::InvalidArgument {
249                op: "matmul",
250                detail: format!(
251                    "supported ranks are 2 and 3 (batched); got {}x{}",
252                    ad.len(),
253                    bd.len()
254                ),
255            });
256        }
257    };
258    let (bb, kb, n) = match bd {
259        [k, n] => (1usize, *k, *n),
260        [b, k, n] => (*b, *k, *n),
261        _ => {
262            return Err(Error::InvalidArgument {
263                op: "matmul",
264                detail: format!(
265                    "supported ranks are 2 and 3 (batched); got {}x{}",
266                    ad.len(),
267                    bd.len()
268                ),
269            });
270        }
271    };
272    if k != kb {
273        return Err(Error::ShapeMismatch {
274            expected: if bd.len() == 3 {
275                Shape::from([bb, k, n])
276            } else {
277                Shape::from([k, n])
278            },
279            got: b.clone(),
280            op: "matmul",
281        });
282    }
283    let batch = match (ba, bb) {
284        (x, y) if x == y => x,
285        (1, y) => y,
286        (x, 1) => x,
287        // Batch dimensions that are neither equal nor 1 cannot broadcast.
288        _ => {
289            return Err(Error::BroadcastIncompatible {
290                lhs: a.clone(),
291                rhs: b.clone(),
292            });
293        }
294    };
295    let out_shape = if ad.len() == 2 && bd.len() == 2 {
296        Shape::from([m, n])
297    } else {
298        Shape::from([batch, m, n])
299    };
300    Ok(MatmulPlan {
301        batch,
302        m,
303        k,
304        n,
305        a_batch_stride: if ba == 1 { 0 } else { m * k },
306        b_batch_stride: if bb == 1 { 0 } else { k * n },
307        out_shape,
308    })
309}
310
311/// A complete backend: every primitive the tensor method layer dispatches.
312///
313/// Composite operations (mean, softmax, losses, convolution, …) are built
314/// from these primitives device-generically; only what is listed here is
315/// implemented per device.
316pub trait Backend: Send + Sync {
317    /// The device this backend serves.
318    fn device(&self) -> Device;
319
320    /// A short stable name for reports and `oxmera doctor`.
321    fn name(&self) -> &'static str;
322
323    /// Elementwise unary op over a (possibly strided) `f32` tensor,
324    /// producing a fresh contiguous tensor of the same shape.
325    fn unary(&self, op: UnaryOp, a: &Tensor) -> Result<Tensor>;
326
327    /// Elementwise binary op with broadcasting, producing a fresh
328    /// contiguous tensor of the broadcast shape.
329    fn binary(&self, op: BinaryOp, a: &Tensor, b: &Tensor) -> Result<Tensor>;
330
331    /// Matrix product: rank-2 `[m, k] x [k, n] -> [m, n]`, or batched
332    /// rank-3 `[b, m, k] x [b, k, n] -> [b, m, n]`.
333    fn matmul(&self, a: &Tensor, b: &Tensor) -> Result<Tensor>;
334
335    /// Reduce over `axes` (sorted, deduplicated by the caller; empty means
336    /// all axes). `keepdim` keeps reduced axes as size 1.
337    fn reduce(&self, op: ReduceOp, a: &Tensor, axes: &[usize], keepdim: bool) -> Result<Tensor>;
338
339    /// Index of the maximum along `dim`, as an `I64` tensor.
340    fn argmax(&self, a: &Tensor, dim: usize, keepdim: bool) -> Result<Tensor>;
341
342    /// A fresh contiguous tensor with the same logical elements.
343    fn contiguous(&self, a: &Tensor) -> Result<Tensor>;
344
345    /// Download to a contiguous CPU tensor.
346    fn download(&self, a: &Tensor) -> Result<Tensor>;
347
348    /// Upload a contiguous CPU tensor to this backend's device.
349    fn upload(&self, a: &Tensor) -> Result<Tensor>;
350
351    /// Rows of `a` along `dim` selected by `indices` (`I64`).
352    ///
353    /// Backends may return [`Error::NotImplemented`]; the method layer
354    /// then falls back to the CPU backend with a device round-trip.
355    fn index_select(&self, a: &Tensor, dim: usize, indices: &Tensor) -> Result<Tensor> {
356        let _ = (dim, indices);
357        Err(Error::NotImplemented {
358            op: "index_select",
359            detail: format!("backend {}", a.device().kind_name()),
360        })
361    }
362
363    /// `out[indices[i]] += src[i]` along `dim`, on a fresh copy of `a`.
364    ///
365    /// Same fallback contract as [`Backend::index_select`].
366    fn index_add(&self, a: &Tensor, dim: usize, indices: &Tensor, src: &Tensor) -> Result<Tensor> {
367        let _ = (dim, indices, src);
368        Err(Error::NotImplemented {
369            op: "index_add",
370            detail: format!("backend {}", a.device().kind_name()),
371        })
372    }
373}
374
375type Registry = RwLock<HashMap<Device, Arc<dyn Backend>>>;
376
377fn registry() -> &'static Registry {
378    static REGISTRY: OnceLock<Registry> = OnceLock::new();
379    REGISTRY.get_or_init(|| RwLock::new(HashMap::new()))
380}
381
382/// Register a backend for its device, replacing any previous registration.
383///
384/// Backend crates call this from their load-time constructors; linking a
385/// backend crate is what makes its device usable.
386pub fn register_backend(backend: Arc<dyn Backend>) {
387    registry()
388        .write()
389        .expect("backend registry poisoned")
390        .insert(backend.device(), backend);
391}
392
393/// The backend serving `device`, or a typed error when none is registered.
394///
395/// The CPU backend is registered lazily on first use, so it can never be
396/// unavailable; GPU backends register at load time or via
397/// `oxmera_runtime::init()`.
398pub fn backend_for(device: Device) -> Result<Arc<dyn Backend>> {
399    let found = registry()
400        .read()
401        .expect("backend registry poisoned")
402        .get(&device)
403        .cloned();
404    match found {
405        Some(b) => Ok(b),
406        None if device == Device::Cpu => {
407            crate::cpu::register();
408            registry()
409                .read()
410                .expect("backend registry poisoned")
411                .get(&device)
412                .cloned()
413                .ok_or(Error::BackendUnavailable { device })
414        }
415        None => Err(Error::BackendUnavailable { device }),
416    }
417}
418
419/// Every registered device, sorted for stable reporting.
420pub fn registered_devices() -> Vec<Device> {
421    let mut devices: Vec<Device> = registry()
422        .read()
423        .expect("backend registry poisoned")
424        .keys()
425        .copied()
426        .collect();
427    devices.sort_by_key(|d| (d.kind_name(), device_index(*d)));
428    devices
429}
430
431fn device_index(d: Device) -> usize {
432    match d {
433        Device::Cpu => 0,
434        Device::Metal { index } | Device::Cuda { index } => index,
435        _ => 0,
436    }
437}
oxmera_tensor/backend.rs

oxmera_tensor/
backend.rs
