Skip to main content

Module custom_ops

Module custom_ops 

Source
Expand description

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.

Functions§

execute
Execute a previously-registered op. Returns None if the op isn’t registered.
register
Register a custom op under name. Idempotent — re-registering replaces. Names are arbitrary strings; convention: dotted namespacing like "my-crate.my-op".
registered
Snapshot of registered op names (sorted, deterministic).

Type Aliases§

CustomOpFn
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.