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
//! Unified custom-operator extension system.
//!
//! This module is the single entry point for defining operators that plug into
//! all three RSSN-Advanced pipelines simultaneously:
//!
//! | Pipeline | Integration method |
//! |---|---|
//! | **JIT compilation** | [`CustomOpRegistry::apply_to_jit`](crate::custom::descriptor::CustomOpRegistry::apply_to_jit) — registers `eval_fn` pointers; enables `compile_batch_f64x2` for pure operators |
//! | **Heuristic simplifier** | [`CustomOpRegistry::build_rule_registry`](crate::custom::descriptor::CustomOpRegistry::build_rule_registry) — produces a [`RuleRegistry`](crate::heuristic::RuleRegistry) from attached [`SimplifyRule`](crate::custom::descriptor::SimplifyRule)s |
//! | **E-graph saturation** | [`CustomOpRegistry::apply_to_egraph`](crate::custom::descriptor::CustomOpRegistry::apply_to_egraph) — injects [`EGraphRule`](crate::custom::descriptor::EGraphRule)s into an [`EGraph`][crate::egraph::egraph::EGraph] |
//!
//! ## Previous fragmented API vs. this module
//!
//! Before this module, extending the pipeline required three independent
//! registration calls that could easily fall out of sync:
//!
//! ```text
//! // OLD — three disconnected registrations:
//! compiler.register_custom_function(fn_id, fn_ptr); // JIT
//! rule_registry.register_named("name", closure, priority, …); // simplifier
//! egraph.add_rule(closure); // e-graph
//! ```
//!
//! With `CustomOpRegistry` the same operator is described once:
//!
//! ```rust,ignore
//! extern "C" fn my_sigmoid(x: f64) -> f64 { 1.0 / (1.0 + (-x).exp()) }
//!
//! let desc = CustomOpDescriptor::builder(FnId(10), "sigmoid", EvalFn::Arity1(my_sigmoid))
//! .vectorizable()
//! .simplify_rule("sigmoid identity", 5, |_b, _k, _c| None)
//! .build();
//!
//! let mut reg = CustomOpRegistry::new();
//! reg.register(desc)?;
//! let reg = Arc::new(reg);
//!
//! reg.apply_to_jit(&mut compiler); // JIT + batch vectorisation
//! let rule_reg = reg.build_rule_registry(); // heuristic simplifier
//! reg.apply_to_egraph(&mut egraph); // e-graph rules
//! ```
//!
//! ## C FFI
//!
//! C/C++ callers use the `rssn_custom_op_*` family of functions in
//! [`crate::ffi::c_api`]:
//!
//! ```c
//! RssnCustomOpRegistry* reg = rssn_custom_op_registry_new();
//! rssn_custom_op_register_fn1(reg, 10, "sigmoid", my_sigmoid_c, /*vectorizable=*/1);
//! rssn_custom_op_add_simplify_rule(reg, 10, "no-op rule", 5, my_rule_cb, NULL);
//!
//! // Use registry in each pipeline step:
//! rssn_dag_compile_with_custom_ops(builder, root, reg, &fn_ptr);
//! rssn_dag_simplify_with_custom_ops(builder, root, reg, &out_id);
//! rssn_dag_egraph_with_custom_ops(builder, root, config, reg, &out_id);
//!
//! rssn_custom_op_registry_free(reg);
//! ```
pub use ;