rlx_ir/
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 Tensor IR — the intermediate representation for the RLX ML compiler.
17//!
18//! This IR is:
19//! - **Standalone**: no runtime, no backend, no framework coupling
20//! - **Serializable**: graphs can be saved/loaded for AOT compilation
21//! - **Optimizable**: designed for pattern-matching fusion and buffer planning
22//!
23//! The compiler pipeline has three named levels:
24//!
25//! - **HIR** ([`hir`]) — block-oriented IR for model builders (`Linear`,
26//!   `SwiGLU`, `ResidualRmsNorm`, …).
27//! - **MIR** ([`mir`]) — fused tensor DAG; input to [`rlx_opt`].
28//! - **LIR** ([`lir`]) — optimized MIR + arena buffer plan for backends.
29//!
30//! [`Graph`] is the primary DX surface. Use [`Graph::define`] for
31//! fusion-first HIR builders, or [`Graph::new`] / [`GraphModule::mir`]
32//! for primitive MIR. [`GraphModule`] tracks pipeline stage (HIR/MIR/LIR).
33//!
34//! - [`Graph`]: a DAG of tensor operations (like XLA's HloModule)
35//! - [`Node`]: a single operation with typed inputs/outputs
36//! - [`Op`]: the operation kind with parameters
37
38pub mod ad;
39pub mod async_copy;
40pub mod const_check;
41pub mod dtype;
42pub mod dynamic;
43pub mod env;
44pub mod fft;
45pub mod graph;
46pub mod hir;
47pub mod infer;
48pub mod infer_shape;
49pub mod inspect;
50pub mod layout;
51pub mod lir;
52pub mod logical_kernel;
53pub mod measure;
54pub mod mir;
55pub mod module;
56pub mod nvfp4;
57pub mod op;
58pub mod op_registry;
59pub mod ops;
60pub mod perfetto;
61pub mod phase;
62pub mod pretty;
63pub mod provenance;
64pub mod quant;
65pub mod rng;
66pub use nvfp4::{FP4_E2M1_LUT, NVFP4_GROUP_SIZE, fp4_e2m1_to_f32, fp8_e4m3_scale_to_f32};
67pub mod binding_manifest;
68pub mod component;
69pub mod hir_extension;
70pub mod reflect;
71#[cfg(feature = "serialize")]
72pub mod serialize;
73pub mod shape;
74pub mod target;
75pub mod variant;
76pub mod verify;
77
78pub use ad::AdPipelineStage;
79pub use async_copy::{AsyncCopy, BarrierToken, DoubleBuffer, SyncCopy};
80pub use dtype::{DType, Element, ElementSubtype};
81pub use dynamic::sym;
82pub use dynamic::{
83    DimEnv, bind_graph, collect_dynamic_symbols, has_dynamic_dims, infer_bindings_from_f32_inputs,
84    infer_bindings_from_inputs, same_binding, sync_concat_shapes, sync_graph_shapes,
85    sync_narrow_ops, sync_reshape_ops,
86};
87pub use env::{RlxEnv, RuntimeOverrides, flag, is_unset, parse_or, set, unset, var, var_os};
88pub use fft::{FftGpuPlan, FftMeta, FftNorm, fft_meta, fftn_axes_all, normalize_fftn_axes};
89pub use graph::{Graph, Node, NodeId};
90pub use hir::{FusionPolicy, HirGraphExt, HirModule, HirMut, HirNode, HirNodeId, HirOp};
91pub use infer::GraphExt;
92pub use inspect::{
93    inspect_buffer_plan, inspect_graph, inspect_graph_diff, inspect_hir, inspect_hir_stats,
94    inspect_lir, inspect_mir, inspect_mir_diff, inspect_mir_stats,
95};
96pub use layout::{Coord2, Ragged, ShapeTuple, Strides2, Strides3, Tile2, Tile3};
97pub use lir::{
98    LirBufferPlan, LirBufferSlot, LirFingerprint, LirIoManifest, LirModule, LirViewAlias,
99};
100pub use logical_kernel::{
101    KernelDispatchConfig, KernelDispatchPolicy, LogicalKernelEntry, logical_kinds_in_graph,
102    registered_logical_kernels, should_lower_to_common,
103};
104pub use measure::{CacheBuster, Tick, time_ns};
105pub use mir::{MirModule, MirNode, MirNodeId, MirOp};
106pub use module::{GraphModule, GraphStage};
107pub use op::{Op, OpKind};
108pub use op_registry::{
109    JvpContext, OpExtension, OpRegistry, VjpContext, VmapContext, global_registry, lookup_op,
110    register_op,
111};
112pub use ops::attention::attention_kind_op;
113pub use phase::{Phase, PhaseSchedule, derive_phases};
114pub use provenance::{NodeOrigin, node_label, stamp_pass_origins};
115pub use quant::{QuantMap, QuantScheme};
116pub use rng::Philox4x32;
117#[cfg(feature = "serialize")]
118pub use serialize::{hir_from_json, hir_to_json, lir_from_json, lir_to_json};
119pub use verify::{VerifyError, verify, verify_all, verify_shapes};
120
121/// Lower a HIR module to MIR, then extract the legacy [`Graph`] API surface.
122pub fn hir_to_graph(hir: HirModule) -> Result<Graph, hir::LowerError> {
123    Ok(hir.lower_to_mir()?.into_graph())
124}
125pub use binding_manifest::{BindingManifest, IoBindingEntry, WeightBlock};
126pub use component::{CompilationMode, ModelComponent};
127pub use hir_extension::{
128    HirExtensionFn, apply_hir_extensions, apply_hir_extensions_named, register_hir_extension,
129    registered_hir_extensions,
130};
131pub use reflect::{
132    BlockSpecialization, HirReflection, ManifestDiff, MirReflection, SpecializeBlockRecord,
133    layout_for_binding, layout_from_lir, probe_block_specialization, symbolic_layout_hint,
134};
135pub use shape::{Dim, DimBinding, Shape};
136pub use variant::{ModelPhase, ModelVariant};
rlx_ir/lib.rs

rlx_ir/
lib.rs
