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