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    /// Serializes the tests below.
84    ///
85    /// `clear_for_tests` wipes the process-global registry, so a sibling's
86    /// clear can land between this test's `register` and its `execute` — which
87    /// is exactly how `re_register_replaces` came to unwrap a `None` that the
88    /// code under test never produced.
89    static REGISTRY_TEST_LOCK: Mutex<()> = Mutex::new(());
90
91    /// Take [`REGISTRY_TEST_LOCK`], ignoring poisoning: it orders registry
92    /// access and guards no invariant a panicking test could leave broken.
93    fn registry_lock() -> std::sync::MutexGuard<'static, ()> {
94        REGISTRY_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner())
95    }
96
97    #[test]
98    fn register_then_execute() {
99        let _lock = registry_lock();
100        clear_for_tests();
101        register("test.identity", |ins| ins[0].to_vec());
102        let out = execute("test.identity", &[&[1.0, 2.0, 3.0]]).unwrap();
103        assert_eq!(out, vec![1.0, 2.0, 3.0]);
104    }
105
106    #[test]
107    fn unknown_op_returns_none() {
108        let _lock = registry_lock();
109        clear_for_tests();
110        assert!(execute("nope", &[]).is_none());
111    }
112
113    #[test]
114    fn re_register_replaces() {
115        let _lock = registry_lock();
116        clear_for_tests();
117        register("test.f", |_| vec![1.0]);
118        register("test.f", |_| vec![2.0]);
119        assert_eq!(execute("test.f", &[]).unwrap(), vec![2.0]);
120    }
121}
rlx_runtime/custom_ops.rs

rlx_runtime/
custom_ops.rs
