rlx_runtime/
custom_ops.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3// SPDX-License-Identifier: MIT OR Apache-2.0
4
5//! Custom-op extensibility (plan #25).
6//!
7//! Borrowed from MAX's `extensibility/compiler_internal/` /
8//! `extensibility/tensor/` pattern: downstream users register their
9//! own ops + executors without forking the framework.
10//!
11//! Today this is the data layer — a registry mapping a string op
12//! name to an executor closure. The matching `#[rlx_op]` proc macro
13//! is the syntactic sugar layer; adding it is straightforward when a
14//! real consumer needs less boilerplate.
15//!
16//! The IR doesn't model custom ops natively (there's no
17//! `Op::Custom("name")` variant) — they enter through the runtime
18//! via `CustomOpRegistry::execute` rather than as graph nodes.
19//! That's deliberate: the optimizer's fusion patterns can't reason
20//! about ops it doesn't know, so custom ops should be opaque
21//! "black box" sub-stages rather than first-class IR citizens.
22//! Promote them to real ops once a fusion pattern would benefit.
23
24use std::collections::HashMap;
25use std::sync::{Mutex, OnceLock};
26
27/// Boxed executor: takes (read-only inputs) → produces an owned output.
28/// `Vec<Vec<f32>>` for now since that matches the `rlx_runtime`
29/// `CompiledGraph::run` signature; revisit when the runtime moves to
30/// `Buffer` (plan #59) end-to-end.
31pub type CustomOpFn = Box<dyn Fn(&[&[f32]]) -> Vec<f32> + Send + Sync>;
32
33struct Registry {
34    map: Mutex<HashMap<String, CustomOpFn>>,
35}
36
37fn registry() -> &'static Registry {
38    static R: OnceLock<Registry> = OnceLock::new();
39    R.get_or_init(|| Registry {
40        map: Mutex::new(HashMap::new()),
41    })
42}
43
44/// Register a custom op under `name`. Idempotent — re-registering
45/// replaces. Names are arbitrary strings; convention: dotted
46/// namespacing like `"my-crate.my-op"`.
47pub fn register<F>(name: impl Into<String>, f: F)
48where
49    F: Fn(&[&[f32]]) -> Vec<f32> + Send + Sync + 'static,
50{
51    let r = registry();
52    let mut m = r.map.lock().expect("custom-op registry poisoned");
53    m.insert(name.into(), Box::new(f));
54}
55
56/// Execute a previously-registered op. Returns `None` if the op
57/// isn't registered.
58pub fn execute(name: &str, inputs: &[&[f32]]) -> Option<Vec<f32>> {
59    let r = registry();
60    let m = r.map.lock().expect("custom-op registry poisoned");
61    m.get(name).map(|f| f(inputs))
62}
63
64/// Snapshot of registered op names (sorted, deterministic).
65pub fn registered() -> Vec<String> {
66    let r = registry();
67    let m = r.map.lock().expect("custom-op registry poisoned");
68    let mut v: Vec<String> = m.keys().cloned().collect();
69    v.sort();
70    v
71}
72
73#[doc(hidden)]
74pub fn clear_for_tests() {
75    let r = registry();
76    r.map.lock().unwrap().clear();
77}
78
79#[cfg(test)]
80mod tests {
81    use super::*;
82
83    #[test]
84    fn register_then_execute() {
85        clear_for_tests();
86        register("test.identity", |ins| ins[0].to_vec());
87        let out = execute("test.identity", &[&[1.0, 2.0, 3.0]]).unwrap();
88        assert_eq!(out, vec![1.0, 2.0, 3.0]);
89    }
90
91    #[test]
92    fn unknown_op_returns_none() {
93        clear_for_tests();
94        assert!(execute("nope", &[]).is_none());
95    }
96
97    #[test]
98    fn re_register_replaces() {
99        clear_for_tests();
100        register("test.f", |_| vec![1.0]);
101        register("test.f", |_| vec![2.0]);
102        assert_eq!(execute("test.f", &[]).unwrap(), vec![2.0]);
103    }
104}
rlx_runtime/custom_ops.rs

rlx_runtime/
custom_ops.rs
