rlx_opt/
lib.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3//
4// This program is free software: you can redistribute it and/or modify
5// it under the terms of the GNU General Public License as published by
6// the Free Software Foundation, version 3.
7//
8// This program is distributed in the hope that it will be useful,
9// but WITHOUT ANY WARRANTY; without even the implied warranty of
10// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
11// GNU General Public License for more details.
12//
13// You should have received a copy of the GNU General Public License
14// along with this program. If not, see <https://www.gnu.org/licenses/>.
15
16//! RLX optimizer facade — re-exports [`rlx_fusion`], [`rlx_autodiff`], and
17//! [`rlx_compile`] for backward-compatible `rlx_opt::` paths.
18//!
19//! # Crates
20//!
21//! | Crate | Role |
22//! |-------|------|
23//! | [`rlx_fusion`] | Fusion passes + [`rlx_fusion::unfuse_fused_for_autodiff`] |
24//! | [`rlx_autodiff`] | `grad_with_loss`, `jvp`, `vmap`, [`prepare_graph_for_ad`] (feature `training`) |
25//! | [`rlx_compile`] | [`CompilePipeline`], memory plan, legalization (feature `compile`) |
26//!
27//! # Features
28//!
29//! - `compile` (default) — HIR → MIR → LIR pipeline
30//! - `training` (default) — autodiff transforms
31//! - `full` — both
32
33pub use rlx_fusion;
34
35#[cfg(feature = "training")]
36pub use rlx_autodiff;
37
38#[cfg(feature = "compile")]
39pub use rlx_compile;
40
41// ── Backward-compatible module paths ─────────────────────────────
42
43pub use rlx_fusion::control_flow;
44pub use rlx_fusion::fusion;
45pub use rlx_fusion::fusion_report;
46pub use rlx_fusion::lower_dot_general;
47pub use rlx_fusion::pass;
48pub use rlx_fusion::unfuse;
49
50#[cfg(feature = "training")]
51pub mod autodiff {
52    pub use rlx_autodiff::autodiff::*;
53    pub use rlx_autodiff::prepare_ad::*;
54}
55
56#[cfg(feature = "training")]
57pub mod autodiff_fwd {
58    pub use rlx_autodiff::autodiff_fwd::*;
59}
60
61#[cfg(feature = "training")]
62pub mod prepare_ad {
63    pub use rlx_autodiff::prepare_ad::*;
64}
65
66#[cfg(feature = "compile")]
67pub mod compiler {
68    pub use rlx_compile::compiler::*;
69}
70
71#[cfg(feature = "compile")]
72pub mod memory {
73    pub use rlx_compile::memory::*;
74}
75
76#[cfg(feature = "compile")]
77pub mod fusion_pipeline {
78    pub use rlx_compile::fusion_pipeline::*;
79}
80
81#[cfg(feature = "compile")]
82pub mod inspect {
83    pub use rlx_compile::inspect::*;
84}
85
86#[cfg(feature = "compile")]
87pub mod legalize {
88    pub use rlx_compile::legalize::*;
89}
90
91#[cfg(feature = "compile")]
92pub mod legalize_broadcast {
93    pub use rlx_compile::legalize_broadcast::*;
94}
95
96#[cfg(feature = "compile")]
97pub mod const_fold {
98    pub use rlx_compile::const_fold::*;
99}
100
101#[cfg(feature = "compile")]
102pub mod dce {
103    pub use rlx_compile::dce::*;
104}
105
106#[cfg(feature = "compile")]
107pub mod precision {
108    pub use rlx_compile::precision::*;
109}
110
111#[cfg(feature = "compile")]
112pub mod quant_insert {
113    pub use rlx_compile::quant_insert::*;
114}
115
116#[cfg(feature = "compile")]
117pub mod quant_propagate {
118    pub use rlx_compile::quant_propagate::*;
119}
120
121#[cfg(feature = "compile")]
122pub mod promote_params {
123    pub use rlx_compile::promote_params::*;
124}
125
126#[cfg(feature = "compile")]
127pub mod inline {
128    pub use rlx_compile::inline::*;
129}
130
131#[cfg(feature = "compile")]
132pub mod svg {
133    pub use rlx_compile::svg::*;
134}
135
136#[cfg(feature = "training")]
137pub mod vmap {
138    pub use rlx_autodiff::vmap::*;
139}
140
141// ── Root re-exports (legacy `use rlx_opt::…`) ─────────────────────
142
143pub use rlx_fusion::{
144    FuseAttentionBlock, FuseMatMulBiasAct, FuseResidualLN, FuseResidualRmsNorm, FuseRmsNormReshape,
145    FuseSharedInputMatMul, FuseSwiGLU, FuseSwiGLUDualMatmul, FuseTransformerLayer, FusionReport,
146    LowerControlFlow, LowerDotGeneral, MarkElementwiseRegions, MissReason, MissedFusion, Pass,
147    UnfuseElementwiseRegions, inline_if, inline_subgraph_into, run_passes,
148    unfuse_fused_for_autodiff, unroll_while,
149};
150
151#[cfg(feature = "training")]
152pub use rlx_autodiff::{
153    AutodiffError, ForceEnergyLossWeights, GradWithLossOptions, HigherOrderOptions, MirAutodiffExt,
154    PrepareForAutodiff, build_force_energy_loss, cse, decompose_backward_for_ad,
155    directional_nth_grad, fuse_elementwise, grad, grad_subgraph, grad_subgraph_for_jvp,
156    grad_with_loss, grad_with_loss_module, grad_with_loss_opts, hvp, hvp_module, jvp, jvp_module,
157    nth_order_grad, nth_order_grad_module, nth_order_grad_with_options, prepare_grad_graph_for_jvp,
158    prepare_graph_for_ad, prepare_mir_for_ad, prepare_module_for_ad, quantized_weight_bits,
159};
160
161#[cfg(feature = "training")]
162pub use rlx_autodiff::vmap::vmap;
163
164#[cfg(feature = "training")]
165pub use rlx_autodiff::autodiff::{convert_scans_for_ad, inline_custom_fn_for_autodiff};
166
167#[cfg(all(feature = "compile", feature = "training"))]
168pub use rlx_compile::{TrainingCompileError, TrainingCompileResult, backward_cleanup_passes};
169
170#[cfg(feature = "compile")]
171pub use rlx_compile::{
172    AlgebraicSimplify, AutoMixedPrecision, CalibrationEntry, CalibrationRecord, CastConfig,
173    CompilePipeline, CompileResult, ConstantFolding, DeadCodeElimination, DispatchPath,
174    FusionLimits, FusionOptions, FusionTarget, KernelDispatchConfig, KernelDispatchPolicy,
175    KernelDispatchReport, KindDispatchSummary, LegalizeBroadcast, LegalizeResult,
176    MemoryPlanOptions, OpKind, PipelineInspect, Precision, PrecisionPolicy, SharedWeightLayout,
177    SpecializeParams, WeightSlot, analyze_dispatch, format_dispatch_report, format_legalize_error,
178    fusion_limits_for_target, fusion_passes, fusion_passes_for_supported, inline_into, insert_q_dq,
179    inspect_compiled, inspect_fusion, inspect_pipeline, is_pure_view, legalize_for_backend,
180    legalize_or_rewrite_for_backend, legalize_or_rewrite_for_backend_with_config,
181    legalize_or_rewrite_for_backend_with_dispatch, maybe_dump_pipeline, maybe_log_dispatch_report,
182    plan_memory_backward, plan_memory_f32_uniform, plan_memory_with_options,
183    prepare_graph_for_backend_with_report, promote_params_to_inputs, rewrite_for_backend,
184    rewrite_for_backend_with_config, rewrite_for_backend_with_dispatch, specialize_params,
185    supported_for_target, supports_op,
186};
rlx_opt/lib.rs

rlx_opt/
lib.rs
