rlx_runtime/
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 Runtime — the user-facing API.
17//!
18//! Provides a unified [`Session`] that compiles and executes IR graphs
19//! on the selected backend. Backend selection is via Cargo features:
20//!
21//! ```toml
22//! [dependencies]
23//! rlx-runtime = { version = "0.1", features = ["cpu"] }                # CPU (default)
24//! rlx-runtime = { version = "0.1", features = ["blas-accelerate"] }    # CPU + Apple Accelerate
25//! rlx-runtime = { version = "0.1", features = ["blas-mkl"] }           # CPU + Intel MKL
26//! # rlx-runtime = { version = "0.1", features = ["gpu"] }             # GPU via wgpu
27//! # rlx-runtime = { version = "0.1", features = ["cuda"] }            # GPU via CUDA
28//! ```
29//!
30//! # Example
31//! ```rust,no_run
32//! use rlx_runtime::*;
33//! use rlx_ir::*;
34//!
35//! // Build a graph
36//! let mut g = Graph::new("example");
37//! let x = g.input("x", Shape::new(&[2, 4], DType::F32));
38//! let w = g.param("w", Shape::new(&[4, 3], DType::F32));
39//! let b = g.param("b", Shape::new(&[3], DType::F32));
40//! let mm = g.matmul(x, w, Shape::new(&[2, 3], DType::F32));
41//! let out = g.binary(op::BinaryOp::Add, mm, b, Shape::new(&[2, 3], DType::F32));
42//! g.set_outputs(vec![out]);
43//!
44//! // Compile and execute
45//! let session = Session::new(Device::Cpu);
46//! let mut compiled = session.compile(g);
47//! compiled.set_param("w", &[1.0f32; 12]);
48//! compiled.set_param("b", &[0.0f32; 3]);
49//! let result = compiled.run(&[("x", &[1.0f32; 8])]);
50//! ```
51
52// Driver-layer concerns (device, arena, handle, stream, buffer)
53// live in rlx-driver as of plan #58; re-exported below so
54// existing callers compile unchanged.
55pub mod aot_cache;
56pub mod attn_mask;
57pub mod backend;
58pub mod backends_manifest;
59pub mod browser;
60pub mod browser_support;
61pub mod capabilities;
62pub mod check;
63pub mod compile_cache;
64pub mod compile_config;
65pub mod compiled;
66// On-device (ANE) gradient-based training + MLUpdateTask eligibility probe.
67#[cfg(all(feature = "coreml", feature = "training"))]
68pub mod coreml_training;
69pub mod cost;
70mod cpu_low_precision;
71mod deferred;
72pub mod device_bench;
73pub mod device_ext;
74pub mod device_parse;
75pub mod device_policy;
76/// Ship-graph distributed execution (build the worker once, run any model).
77pub mod dist;
78pub mod expert_pool;
79pub mod export;
80pub mod graph_io;
81#[cfg(feature = "hwprofile")]
82pub mod hwprofile_select;
83pub mod jacfwd;
84pub mod kernel_trace;
85pub mod kv_cache;
86pub mod lora_scheduler;
87pub mod memory_estimate;
88pub mod model_pipeline;
89pub mod moe_expert_store;
90#[cfg(feature = "cpu")]
91pub mod onnx_active;
92pub mod op_registry;
93pub mod options;
94pub mod paged_kv;
95pub mod precision;
96pub mod precompile;
97pub mod quantized_kv;
98pub mod record_replay;
99pub mod reflect;
100pub mod registry;
101pub mod router;
102pub mod samplers;
103pub mod session;
104pub mod stages;
105pub mod subgraph;
106pub mod trace;
107pub mod weight_registry;
108pub mod weights;
109pub mod worker_pool;
110/// PLAN L3 — Perfetto / chrome-trace JSON tracing. Lives in `rlx-ir`
111/// (alongside the `Tick` cycle counter it depends on) so every backend
112/// can instrument per-thunk without crate-dep gymnastics. Re-exported
113/// here so callers see one consistent `rlx_runtime::perfetto::TraceSpan`.
114pub use rlx_ir::perfetto;
115pub mod custom_ops;
116pub mod device_router;
117pub mod flexible_session;
118pub mod graph_devices;
119pub mod hetero;
120pub mod hwinfo;
121pub mod lm;
122pub mod logit_verify;
123pub mod nan_check;
124pub mod phase;
125pub mod spec_decode;
126pub mod telemetry;
127pub mod validators;
128
129// Always-available now that serde is a non-optional dep + #32
130// router consumes the OpenAI-shaped structs unconditionally.
131pub mod mock_requests;
132
133// Driver-layer types — re-exported from rlx-driver (plan #58).
134pub use rlx_driver::{Buffer, BufferHandle, CommandStream, Device, DeviceArena, SyncStream};
135// Symmetric-memory primitives (plan #49) — foundation for #12.
136pub use rlx_driver::{
137    CollectiveError, LocalTransport, Rank, SymmetricBuffer, SymmetricHeap, SymmetricTransport,
138};
139// Collective ops (plan #12).
140pub use aot_cache::{AotCache, AotCacheError};
141pub use backend::{Backend, ExecutableGraph, compile_hir, compile_module};
142pub use backends_manifest::BackendsManifest;
143pub use browser::{
144    BrowserCompiledGraph, BrowserError, BrowserSession, browser_block_reason, init_webgpu,
145    supports_browser_graph,
146};
147pub use browser_support::{
148    collective_block_reason, passes_browser_preflight, select_browser_device_for_graph,
149};
150pub use capabilities::ExecutableCapabilities;
151pub use compile_cache::{
152    BucketedCompileCache, CacheRunInput, CompileCache, DynamicDimCompileCache, pad_rows, slice_rows,
153};
154pub use compile_config::{
155    COMPILE_OUTPUT_CAP_ENV, COMPILE_OUTPUT_CAP_ENV_MLX, DEFAULT_COMPILE_OUTPUT_CAP,
156    compile_output_cap, device_has_compile_output_cap, reset_compile_output_cap,
157    set_compile_output_cap,
158};
159pub use compiled::CompiledGraph;
160#[cfg(all(feature = "coreml", feature = "training"))]
161pub use coreml_training::{
162    CoremlTrainingSession, MlUpdateEligibility, Optimizer, StepReport, TrainConfig, UpdatePath,
163};
164pub use cost::fastest_device_for;
165pub use device_bench::{DeviceBenchResult, benchmark_devices, warm_all};
166#[cfg(feature = "apple")]
167pub use device_ext::available_apple_devices;
168pub use device_ext::{
169    BROWSER_DEVICE_PRIORITY, available_browser_devices, available_devices, devices_for,
170    dispatch_report_for_device, dispatch_report_for_device_with_options, fastest_device,
171    first_unsupported_op, first_unsupported_op_with_options, full_name, is_available,
172    legalize_graph_for_device, legalize_graph_for_device_with_options,
173    legalize_graph_for_device_with_report, preferred_browser_device, supports, supports_graph,
174    supports_graph_with_options, supports_run_slots, trim_accelerator_arena_pool,
175};
176pub use device_parse::{ParseDeviceError, device_label, parse_device, parse_device_list};
177pub use device_policy::{
178    DeviceCandidate, DeviceFallbackError, DevicePickStrategy, DevicePolicy, advisory_capabilities,
179    device_chain_from_env, device_chain_from_env_key, device_from_env, device_from_env_key,
180    device_report, devices_for_with_policy, resolve_device, resolve_device_chain,
181    run_with_fallback,
182};
183pub use device_router::DeviceRouter;
184pub use expert_pool::{
185    ExpertPool, ExpertPoolConfig, ExpertPoolStats, ExpertRefreshPolicy, ExpertRefreshResult,
186    MoEExecMode, gpu_expert_budget_from_vram,
187};
188pub use export::{ExportOptions, ExportTarget, ExportedArtifacts, export_graph};
189#[cfg(feature = "fpga")]
190pub use export::{ExportSession, export_tinyconv_mnist};
191pub use flexible_session::FlexibleSession;
192pub use graph_devices::{GraphDevices, graph_param_names};
193pub use hetero::{DeviceMap, HeteroExecutable};
194#[cfg(feature = "hwprofile")]
195pub use hwprofile_select::{device_vram_bytes, fastest_device_with_vram};
196pub use kv_cache::LayerKvCache;
197pub use lm::{
198    ConfigSource, LmRunner, LmRunnerBuilder, MirostatMode, ModelRegistration,
199    PACKED_GGUF_AUTO_THRESHOLD_BYTES, SampleOpts, WeightFormat, auto_runner_name,
200    registered_models,
201};
202pub use memory_estimate::{
203    DEFAULT_SOFT_MEMORY_FRACTION, MoeOffloadEstimate, available_unified_memory,
204    estimate_moe_offload, llama_decode_bucket_compile_peak_bytes,
205    llama_decode_bucket_resident_bytes, llama_decode_oneshot_compile_peak_bytes,
206    memory_headroom_bytes, process_rss_bytes, soft_memory_budget_bytes, soft_memory_fraction,
207    would_exceed_soft_budget,
208};
209pub use model_pipeline::ModelCompilePipeline;
210pub use options::CompileOptions;
211pub use precision::Precision;
212pub use reflect::{ModelReflection, load_hir_template_with_extensions, specialize_entry};
213pub use registry::{BackendFactory, backend_for, register_backend, registered_devices};
214/// Native low-precision GEMM config for [`CompileOptions::scaled_quant`].
215pub use rlx_opt::rlx_compile::scaled_quant_insert::ScaledQuantConfig;
216
217/// Alias for [`COMPILE_OUTPUT_CAP_ENV`].
218pub const MLX_COMPILE_OUTPUT_CAP_ENV: &str = COMPILE_OUTPUT_CAP_ENV;
219
220/// Alias for [`DEFAULT_COMPILE_OUTPUT_CAP`].
221pub const DEFAULT_MLX_COMPILE_OUTPUT_CAP: usize = DEFAULT_COMPILE_OUTPUT_CAP;
222
223/// Alias for [`compile_output_cap`].
224#[inline]
225pub fn mlx_compile_output_cap() -> usize {
226    compile_output_cap()
227}
228
229/// Alias for [`set_compile_output_cap`].
230#[inline]
231pub fn set_mlx_compile_output_cap(cap: usize) {
232    set_compile_output_cap(cap);
233}
234
235/// Alias for [`reset_compile_output_cap`].
236#[inline]
237pub fn reset_mlx_compile_output_cap() {
238    reset_compile_output_cap();
239}
240
241#[cfg(feature = "cpu")]
242pub use rlx_cpu::moe_residency::MoeResidencyStats;
243/// Stub for builds without the `cpu` backend feature (e.g. downstream crates that
244/// bring their own `rlx-cpu` and set `default-features = false`). The MoE-residency
245/// backend trait methods reference `crate::MoeResidencyStats` unconditionally but only
246/// ever return `None` off the CPU path, so an empty type keeps the API compiling.
247#[cfg(not(feature = "cpu"))]
248#[derive(Debug, Clone, Default)]
249pub struct MoeResidencyStats;
250#[cfg(feature = "cpu")]
251pub use rlx_cpu::moe_topk_capture::MoeTopkCapture;
252pub use rlx_driver::{
253    DEFAULT_HEAP_BYTES, NetTransport, ProcessGroup, TcpTransport, ThunderboltTransport, Transport,
254};
255pub use rlx_driver::{
256    ReduceKind, ReduceMode, all_gather, all_reduce, env_reduce_mode, reduce_scatter,
257};
258pub use rlx_ir::env::{self, RlxEnv, RuntimeOverrides};
259pub use rlx_ir::{ENV_CATALOG, EnvVarDoc, format_env_catalog};
260pub use session::Session;
261pub use stages::{
262    compile_graph_stages, compile_graph_stages_for_backend, compile_hir_stages,
263    compile_module_stages, fusion_target_for, graph_from_lir, maybe_log_fusion,
264    options_with_supported_ops, pipeline_for,
265};
266pub use subgraph::{SubgraphCache, run_if, run_while};
267
268pub use expert_pool::{merged_resident_mask, per_layer_resident_masks};
269pub use moe_expert_store::{ExpertStackF32, LayerMoeWeights, MoeExpertStore};
270pub use weight_registry::{WeightEntry, WeightHandle, WeightKind, WeightRegistry};
271pub use weights::{BytesWeightLoader, WeightLoader};
272
273// Cycle-accurate timing primitive lives in rlx-ir (lowest crate); re-export
274// here so `rlx_runtime::Tick` works without forcing callers to add the IR
275// crate as a direct dep.
276pub use rlx_ir::{AsyncCopy, BarrierToken, DoubleBuffer, SyncCopy};
277pub use rlx_ir::{CacheBuster, Tick, time_ns};
278
279// Re-export precision policy from rlx-opt for convenience
280pub use rlx_ir::{
281    inspect_graph, inspect_hir, inspect_hir_stats, inspect_lir, inspect_mir, inspect_mir_stats,
282};
283pub use rlx_opt::{OpKind, PrecisionPolicy};
284pub use rlx_opt::{PipelineInspect, inspect_pipeline};
285
286// Re-export IR types for convenience
287pub use rlx_ir::logical_kernel::{KernelDispatchConfig, KernelDispatchPolicy};
288pub use rlx_ir::op;
289pub use rlx_ir::{
290    BindingManifest, CompilationMode, DType, Graph, HirExtensionFn, HirReflection, IoBindingEntry,
291    ManifestDiff, ModelComponent, ModelPhase, ModelVariant, Node, NodeId, Op, RngBackend,
292    RngOptions, Shape, WeightBlock, apply_hir_extensions, register_hir_extension,
293    registered_hir_extensions,
294};
295
296// Re-export proc macro
297pub use rlx_macros::pipeline_schedule;
298pub use rlx_macros::rlx_model;
rlx_runtime/lib.rs

rlx_runtime/
lib.rs
