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, FusionReport, LowerControlFlow,
146    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, MirAutodiffExt, PrepareForAutodiff, grad, grad_with_loss, grad_with_loss_module,
154    hvp, jvp, jvp_module, prepare_graph_for_ad, prepare_mir_for_ad, prepare_module_for_ad,
155    quantized_weight_bits,
156};
157
158#[cfg(feature = "training")]
159pub use rlx_autodiff::vmap::vmap;
160
161#[cfg(feature = "training")]
162pub use rlx_autodiff::autodiff::{convert_scans_for_ad, inline_custom_fn_for_autodiff};
163
164#[cfg(all(feature = "compile", feature = "training"))]
165pub use rlx_compile::{TrainingCompileError, TrainingCompileResult, backward_cleanup_passes};
166
167#[cfg(feature = "compile")]
168pub use rlx_compile::{
169    AutoMixedPrecision, CalibrationEntry, CalibrationRecord, CastConfig, CompilePipeline,
170    CompileResult, ConstantFolding, DeadCodeElimination, DispatchPath, FusionLimits, FusionOptions,
171    FusionTarget, KernelDispatchConfig, KernelDispatchPolicy, KernelDispatchReport,
172    KindDispatchSummary, LegalizeBroadcast, LegalizeResult, MemoryPlanOptions, OpKind,
173    PipelineInspect, Precision, PrecisionPolicy, SharedWeightLayout, WeightSlot, analyze_dispatch,
174    format_dispatch_report, format_legalize_error, fusion_limits_for_target, fusion_passes,
175    fusion_passes_for_supported, inline_into, insert_q_dq, inspect_compiled, inspect_fusion,
176    inspect_pipeline, is_pure_view, legalize_for_backend, legalize_or_rewrite_for_backend,
177    legalize_or_rewrite_for_backend_with_config, legalize_or_rewrite_for_backend_with_dispatch,
178    maybe_dump_pipeline, maybe_log_dispatch_report, plan_memory_backward, plan_memory_with_options,
179    prepare_graph_for_backend_with_report, promote_params_to_inputs, rewrite_for_backend,
180    rewrite_for_backend_with_config, rewrite_for_backend_with_dispatch, supported_for_target,
181    supports_op,
182};
rlx_opt/lib.rs

rlx_opt/
lib.rs
