rlx_cpu/
op_registry.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3//
4// This program is free software: you can redistribute it and/or modify
5// it under the terms of the GNU General Public License as published by
6// the Free Software Foundation, version 3.
7//
8// This program is distributed in the hope that it will be useful,
9// but WITHOUT ANY WARRANTY; without even the implied warranty of
10// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
11// GNU General Public License for more details.
12//
13// You should have received a copy of the GNU General Public License
14// along with this program. If not, see <https://www.gnu.org/licenses/>.
15
16//! Per-backend (CPU) kernel registry for `Op::Custom`.
17//!
18//! Companion to [`rlx_ir::op_registry`]. The IR-level [`rlx_ir::OpExtension`]
19//! covers shape inference + autodiff; this registry covers
20//! *execution* on the CPU backend. Splitting them keeps `rlx-ir`
21//! portable and lets a custom op honestly support a subset of
22//! backends — attempting to compile an `Op::Custom` for a backend
23//! whose kernel isn't registered is a hard error, not a silent no-op.
24//!
25//! ## API contract for downstream kernel authors
26//!
27//! - **One method, typed views in.** Each input arrives as a
28//!   [`CpuTensorRef`] variant matching that input's declared dtype.
29//!   The output is a [`CpuTensorMut`] matching the output dtype. No
30//!   byte reinterpretation in user code.
31//! - **Mixed-dtype inputs work directly.** A Sparse-LU op with
32//!   `(F64 values, I32 col_idx, I32 row_ptr, F64 b)` gets each input
33//!   as the right typed slice — no manual byte casts.
34//! - **Contiguous, dense buffers from the arena.** Strided / broadcast
35//!   inputs need to be materialized by the caller before reaching the
36//!   kernel; the IR's `Op::Expand` / `Op::Transpose` already cover
37//!   that.
38//! - **`attrs` is opaque** — same `Vec<u8>` as the IR variant. Decode
39//!   it however the kernel likes (typical: `bincode`, `bytemuck`, a
40//!   hand-rolled struct cast).
41//!
42//! ## Multiple logical outputs
43//!
44//! `Op::Custom` produces a single tensor by design — the IR's `Node`
45//! is one-shape-per-node. Ops that conceptually return multiple
46//! outputs (LU returning L+U, eigendecomp returning λ+V) write a
47//! *packed* output and the user follows the custom op with `Narrow`
48//! to extract each logical output. Use
49//! [`rlx_ir::Graph::custom_op_packed`] when registry-driven shape
50//! inference isn't sufficient.
51
52use std::collections::HashMap;
53use std::sync::{Arc, OnceLock, RwLock};
54
55use rlx_ir::{DType, Shape};
56
57// Why an enum, not generics? `CpuKernel` takes inputs of *mixed*
58// dtypes (e.g. Sparse-LU has `(F64 values, I32 col_idx, I32 row_ptr,
59// F64 b)`). A generic `CpuKernel<T>` couldn't express that — every
60// input would have to share the same `T`. The enum-of-typed-views
61// is the right shape for this contract; generics over `T: Pod`
62// would only buy us the per-input case, which is what `as_*` /
63// `expect_*` accessors already provide.
64//
65// One variant per `rlx_ir::DType`. The dispatcher in
66// `thunk.rs::dispatch_custom_op` enumerates all of them — adding a
67// dtype to `DType` requires adding a variant here and an arm
68// there. Single source of truth for "what's wired."
69
70macro_rules! dtype_variants {
71    (
72        $(
73            $variant:ident => $rust_ty:ty,
74            $as_method:ident, $as_mut_method:ident,
75            $expect_method:ident, $expect_mut_method:ident,
76        )*
77    ) => {
78        /// Read-only typed view of one input tensor handed to a [`CpuKernel`].
79        /// The variant matches the input's declared dtype on the IR side.
80        pub enum CpuTensorRef<'a> {
81            $(
82                $variant { data: &'a [$rust_ty], shape: &'a Shape },
83            )*
84        }
85
86        /// Mutable typed view of the output tensor handed to a [`CpuKernel`].
87        pub enum CpuTensorMut<'a> {
88            $(
89                $variant { data: &'a mut [$rust_ty], shape: &'a Shape },
90            )*
91        }
92
93        impl<'a> CpuTensorRef<'a> {
94            pub fn shape(&self) -> &Shape {
95                match self {
96                    $( Self::$variant { shape, .. } => shape, )*
97                }
98            }
99            pub fn dtype(&self) -> DType { self.shape().dtype() }
100
101            $(
102                pub fn $as_method(&self) -> Option<&[$rust_ty]> {
103                    if let Self::$variant { data, .. } = self { Some(data) } else { None }
104                }
105                pub fn $expect_method(&self, role: &str) -> Result<&[$rust_ty], String> {
106                    self.$as_method().ok_or_else(|| format!(
107                        "{role}: expected {:?}, got {:?}",
108                        DType::$variant, self.dtype()))
109                }
110            )*
111        }
112
113        impl<'a> CpuTensorMut<'a> {
114            pub fn shape(&self) -> &Shape {
115                match self {
116                    $( Self::$variant { shape, .. } => shape, )*
117                }
118            }
119            pub fn dtype(&self) -> DType { self.shape().dtype() }
120
121            $(
122                pub fn $as_mut_method(self) -> Option<&'a mut [$rust_ty]> {
123                    if let Self::$variant { data, .. } = self { Some(data) } else { None }
124                }
125                pub fn $expect_mut_method(self, role: &str) -> Result<&'a mut [$rust_ty], String> {
126                    let dt = self.dtype();
127                    self.$as_mut_method().ok_or_else(|| format!(
128                        "{role}: expected {:?}, got {dt:?}", DType::$variant))
129                }
130            )*
131        }
132    };
133}
134
135// One row per DType. Bool is stored as `u8` on the wire (one byte
136// per element, 0 = false / non-zero = true) — exposing it as a bool
137// slice directly would be UB if any byte pattern other than 0/1
138// landed there, which the IR doesn't guarantee.
139dtype_variants! {
140    F32  => f32,        as_f32,  as_f32_mut,  expect_f32,  expect_f32_mut,
141    F64  => f64,        as_f64,  as_f64_mut,  expect_f64,  expect_f64_mut,
142    F16  => half::f16,  as_f16,  as_f16_mut,  expect_f16,  expect_f16_mut,
143    BF16 => half::bf16, as_bf16, as_bf16_mut, expect_bf16, expect_bf16_mut,
144    I8   => i8,         as_i8,   as_i8_mut,   expect_i8,   expect_i8_mut,
145    I16  => i16,        as_i16,  as_i16_mut,  expect_i16,  expect_i16_mut,
146    I32  => i32,        as_i32,  as_i32_mut,  expect_i32,  expect_i32_mut,
147    I64  => i64,        as_i64,  as_i64_mut,  expect_i64,  expect_i64_mut,
148    U8   => u8,         as_u8,   as_u8_mut,   expect_u8,   expect_u8_mut,
149    U32  => u32,        as_u32,  as_u32_mut,  expect_u32,  expect_u32_mut,
150    Bool => u8,         as_bool, as_bool_mut, expect_bool, expect_bool_mut,
151}
152
153/// Trait a CPU kernel implements for one custom op. Registered under
154/// the same `name` used in `Op::Custom` and `OpExtension::name`.
155///
156/// One method, typed views in. Match on the variants you support and
157/// return `Err(...)` for anything else — the executor surfaces that
158/// as a panic naming the op + dtype, so missing support fails loudly
159/// instead of silently zeroing the output.
160pub trait CpuKernel: Send + Sync {
161    fn name(&self) -> &str;
162
163    fn execute(
164        &self,
165        inputs: &[CpuTensorRef<'_>],
166        output: CpuTensorMut<'_>,
167        attrs: &[u8],
168    ) -> Result<(), String>;
169}
170
171pub struct CpuKernelRegistry {
172    kernels: RwLock<HashMap<String, Arc<dyn CpuKernel>>>,
173}
174
175impl CpuKernelRegistry {
176    pub fn new() -> Self {
177        Self {
178            kernels: RwLock::new(HashMap::new()),
179        }
180    }
181
182    /// Register a kernel. Re-registration replaces the previous entry
183    /// and prints a one-line warning to stderr — silent overwrite has
184    /// bitten us before, the warning is cheap.
185    pub fn register(&self, k: Arc<dyn CpuKernel>) {
186        let name = k.name().to_string();
187        let mut g = self.kernels.write().unwrap();
188        if g.contains_key(&name) {
189            eprintln!(
190                "rlx-cpu: CpuKernel '{name}' was already registered — \
191                 replacing the previous entry"
192            );
193        }
194        g.insert(name, k);
195    }
196
197    pub fn lookup(&self, name: &str) -> Option<Arc<dyn CpuKernel>> {
198        self.kernels.read().unwrap().get(name).cloned()
199    }
200}
201
202impl Default for CpuKernelRegistry {
203    fn default() -> Self {
204        Self::new()
205    }
206}
207
208pub fn global_cpu_kernels() -> &'static CpuKernelRegistry {
209    static R: OnceLock<CpuKernelRegistry> = OnceLock::new();
210    R.get_or_init(CpuKernelRegistry::new)
211}
212
213pub fn register_cpu_kernel(k: Arc<dyn CpuKernel>) {
214    global_cpu_kernels().register(k);
215}
216
217pub fn lookup_cpu_kernel(name: &str) -> Option<Arc<dyn CpuKernel>> {
218    global_cpu_kernels().lookup(name)
219}
220
221/// Run a registered custom-op [`CpuKernel`] against **host byte buffers** — the
222/// generic host-fallback a GPU backend uses for a custom op it has no native
223/// kernel for. The `collective.*` ops are the motivating case: they are
224/// host/transport ops with no device kernel, so every backend runs them by
225/// staging the operands to host and calling the one registered CPU kernel here.
226///
227/// All operands are treated as `F32` (the dtype of the ops that take this path).
228/// Each `(&[u8], &Shape)` input is reinterpreted as an `f32` slice and handed to
229/// the kernel as a [`CpuTensorRef::F32`]; the output the same. Returns `Err` if
230/// no kernel is registered under `name`, if a buffer is not 4-byte aligned or a
231/// multiple of 4 bytes, or if the kernel itself errors.
232pub fn run_f32_custom_op_host(
233    name: &str,
234    inputs: &[(&[u8], &Shape)],
235    out: (&mut [u8], &Shape),
236    attrs: &[u8],
237) -> Result<(), String> {
238    let kernel = lookup_cpu_kernel(name)
239        .ok_or_else(|| format!("no CpuKernel registered for custom op '{name}'"))?;
240
241    fn as_f32<'a>(bytes: &'a [u8], role: &str) -> Result<&'a [f32], String> {
242        if !(bytes.as_ptr() as usize).is_multiple_of(std::mem::align_of::<f32>())
243            || !bytes.len().is_multiple_of(4)
244        {
245            return Err(format!("{role}: host buffer not f32-aligned/sized"));
246        }
247        // SAFETY: alignment + length checked; every bit pattern is a valid f32.
248        Ok(unsafe { std::slice::from_raw_parts(bytes.as_ptr() as *const f32, bytes.len() / 4) })
249    }
250
251    let refs: Vec<CpuTensorRef<'_>> = inputs
252        .iter()
253        .enumerate()
254        .map(|(i, &(b, sh))| {
255            as_f32(b, &format!("{name} input {i}"))
256                .map(|data| CpuTensorRef::F32 { data, shape: sh })
257        })
258        .collect::<Result<_, _>>()?;
259
260    let (out_bytes, out_shape) = out;
261    if !(out_bytes.as_ptr() as usize).is_multiple_of(std::mem::align_of::<f32>())
262        || !out_bytes.len().is_multiple_of(4)
263    {
264        return Err(format!("{name} output: host buffer not f32-aligned/sized"));
265    }
266    let out_len = out_bytes.len() / 4;
267    // SAFETY: alignment + length checked above.
268    let out_data =
269        unsafe { std::slice::from_raw_parts_mut(out_bytes.as_mut_ptr() as *mut f32, out_len) };
270    let out_view = CpuTensorMut::F32 {
271        data: out_data,
272        shape: out_shape,
273    };
274
275    kernel.execute(&refs, out_view, attrs)
276}
277
278/// Reinterpret `bytes` as a `&[T]`, checking alignment + size. SAFETY caller:
279/// every bit pattern of the raw scalar types we use here (integers/floats/half)
280/// is a valid value, so the transmute is sound once alignment + length hold.
281fn as_typed<'a, T>(bytes: &'a [u8], role: &str) -> Result<&'a [T], String> {
282    let sz = std::mem::size_of::<T>();
283    if !(bytes.as_ptr() as usize).is_multiple_of(std::mem::align_of::<T>())
284        || !bytes.len().is_multiple_of(sz)
285    {
286        return Err(format!("{role}: host buffer not {sz}-byte aligned/sized"));
287    }
288    Ok(unsafe { std::slice::from_raw_parts(bytes.as_ptr() as *const T, bytes.len() / sz) })
289}
290
291fn as_typed_mut<'a, T>(bytes: &'a mut [u8], role: &str) -> Result<&'a mut [T], String> {
292    let sz = std::mem::size_of::<T>();
293    if !(bytes.as_ptr() as usize).is_multiple_of(std::mem::align_of::<T>())
294        || !bytes.len().is_multiple_of(sz)
295    {
296        return Err(format!("{role}: host buffer not {sz}-byte aligned/sized"));
297    }
298    let n = bytes.len() / sz;
299    Ok(unsafe { std::slice::from_raw_parts_mut(bytes.as_mut_ptr() as *mut T, n) })
300}
301
302fn cpu_ref_from_bytes<'a>(
303    b: &'a [u8],
304    sh: &'a Shape,
305    role: &str,
306) -> Result<CpuTensorRef<'a>, String> {
307    Ok(match sh.dtype() {
308        DType::F32 => CpuTensorRef::F32 {
309            data: as_typed(b, role)?,
310            shape: sh,
311        },
312        DType::F64 => CpuTensorRef::F64 {
313            data: as_typed(b, role)?,
314            shape: sh,
315        },
316        DType::F16 => CpuTensorRef::F16 {
317            data: as_typed(b, role)?,
318            shape: sh,
319        },
320        DType::BF16 => CpuTensorRef::BF16 {
321            data: as_typed(b, role)?,
322            shape: sh,
323        },
324        DType::I8 => CpuTensorRef::I8 {
325            data: as_typed(b, role)?,
326            shape: sh,
327        },
328        DType::I16 => CpuTensorRef::I16 {
329            data: as_typed(b, role)?,
330            shape: sh,
331        },
332        DType::I32 => CpuTensorRef::I32 {
333            data: as_typed(b, role)?,
334            shape: sh,
335        },
336        DType::I64 => CpuTensorRef::I64 {
337            data: as_typed(b, role)?,
338            shape: sh,
339        },
340        DType::U8 => CpuTensorRef::U8 {
341            data: as_typed(b, role)?,
342            shape: sh,
343        },
344        DType::U32 => CpuTensorRef::U32 {
345            data: as_typed(b, role)?,
346            shape: sh,
347        },
348        DType::Bool => CpuTensorRef::Bool {
349            data: as_typed(b, role)?,
350            shape: sh,
351        },
352        other => {
353            return Err(format!(
354                "{role}: unsupported dtype {other:?} for host custom-op"
355            ));
356        }
357    })
358}
359
360fn cpu_mut_from_bytes<'a>(b: &'a mut [u8], sh: &'a Shape) -> Result<CpuTensorMut<'a>, String> {
361    Ok(match sh.dtype() {
362        DType::F32 => CpuTensorMut::F32 {
363            data: as_typed_mut(b, "output")?,
364            shape: sh,
365        },
366        DType::F64 => CpuTensorMut::F64 {
367            data: as_typed_mut(b, "output")?,
368            shape: sh,
369        },
370        DType::F16 => CpuTensorMut::F16 {
371            data: as_typed_mut(b, "output")?,
372            shape: sh,
373        },
374        DType::BF16 => CpuTensorMut::BF16 {
375            data: as_typed_mut(b, "output")?,
376            shape: sh,
377        },
378        DType::I8 => CpuTensorMut::I8 {
379            data: as_typed_mut(b, "output")?,
380            shape: sh,
381        },
382        DType::I16 => CpuTensorMut::I16 {
383            data: as_typed_mut(b, "output")?,
384            shape: sh,
385        },
386        DType::I32 => CpuTensorMut::I32 {
387            data: as_typed_mut(b, "output")?,
388            shape: sh,
389        },
390        DType::I64 => CpuTensorMut::I64 {
391            data: as_typed_mut(b, "output")?,
392            shape: sh,
393        },
394        DType::U8 => CpuTensorMut::U8 {
395            data: as_typed_mut(b, "output")?,
396            shape: sh,
397        },
398        DType::U32 => CpuTensorMut::U32 {
399            data: as_typed_mut(b, "output")?,
400            shape: sh,
401        },
402        DType::Bool => CpuTensorMut::Bool {
403            data: as_typed_mut(b, "output")?,
404            shape: sh,
405        },
406        other => {
407            return Err(format!(
408                "output: unsupported dtype {other:?} for host custom-op"
409            ));
410        }
411    })
412}
413
414/// Dtype-aware sibling of [`run_f32_custom_op_host`]: run a registered custom-op
415/// [`CpuKernel`] against **host byte buffers**, reinterpreting each operand per
416/// its `Shape`'s dtype (so integer ops like `onnx.Mod` on I64, or Bool masks,
417/// are handled correctly — not forced to F32). This is the generic host-delegate
418/// a GPU backend (Metal / MLX / wgpu) uses for any `onnx.*` custom op it has no
419/// native kernel for, staging operands to host over unified/mapped memory.
420pub fn run_custom_op_host(
421    name: &str,
422    inputs: &[(&[u8], &Shape)],
423    out: (&mut [u8], &Shape),
424    attrs: &[u8],
425) -> Result<(), String> {
426    let kernel = lookup_cpu_kernel(name)
427        .ok_or_else(|| format!("no CpuKernel registered for custom op '{name}'"))?;
428    let refs: Vec<CpuTensorRef<'_>> = inputs
429        .iter()
430        .enumerate()
431        .map(|(i, &(b, sh))| cpu_ref_from_bytes(b, sh, &format!("{name} input {i}")))
432        .collect::<Result<_, _>>()?;
433    let (out_bytes, out_shape) = out;
434    let out_view = cpu_mut_from_bytes(out_bytes, out_shape)?;
435    kernel.execute(&refs, out_view, attrs)
436}
rlx_cpu/op_registry.rs

rlx_cpu/
op_registry.rs
