rlx_runtime/
options.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3// SPDX-License-Identifier: MIT OR Apache-2.0
4
5//! Unified compile options.
6//!
7//! Replaces the historical mix of `compile()`, `compile_with_precision()`,
8//! `compile_with_options()` with a single `Backend::compile(graph, &options)`.
9//! New compile-time knobs can be added to `CompileOptions` without
10//! changing the trait — backends just read what they care about.
11//!
12//! Builder-pattern API for ergonomics:
13//!
14//! ```rust,ignore
15//! let opts = CompileOptions::new()
16//!     .precision(Precision::F16)
17//!     .policy(PrecisionPolicy::AutoMixed)
18//!     .with_dce(true)
19//!     .with_constant_folding(true);
20//! ```
21
22use crate::Precision;
23use rlx_ir::OpKind;
24use rlx_ir::logical_kernel::{KernelDispatchConfig, KernelDispatchPolicy};
25use rlx_opt::rlx_compile::scaled_quant_insert::ScaledQuantConfig;
26use rlx_opt::{FusionOptions, FusionTarget, PrecisionPolicy};
27use std::collections::HashMap;
28
29/// All knobs the compile pipeline understands.
30/// Add new fields here rather than introducing new compile entry points.
31#[derive(Debug, Clone)]
32pub struct CompileOptions {
33    /// Target numeric precision for execution. Default: F32.
34    pub precision: Precision,
35    /// Optional per-op precision policy (mixed precision rewrite).
36    pub policy: Option<PrecisionPolicy>,
37    /// Opt-in native low-precision GEMM. When set, every 2-D `Op::MatMul` is
38    /// rewritten to a `ScaledMatMul` whose operands are dynamically quantized to
39    /// this element format + scale layout — any [`ScaledFormat`], including a
40    /// parameterized `Custom` (e.g. `f4e3m0`). Off by default; changes numerics.
41    pub scaled_quant: Option<ScaledQuantConfig>,
42    /// RNG policy for in-graph [`Op::RngNormal`] / [`Op::RngUniform`] nodes.
43    pub rng: rlx_ir::RngOptions,
44    /// Run dead-code elimination as part of compile. Default: true.
45    pub dce: bool,
46    /// Run constant folding. Default: true (cheap, only helps).
47    pub constant_folding: bool,
48    /// Verbose pass logging. Equivalent to `RLX_VERBOSE=1` or
49    /// [`rlx_ir::env::set("RLX_VERBOSE", "1")`].
50    pub verbose: bool,
51    /// Override fusion pipeline target (default: inferred from device).
52    pub fusion_target: Option<FusionTarget>,
53    /// Per-target fusion toggles (Metal env overrides, skip fusion, …).
54    pub fusion_opts: FusionOptions,
55    /// Arena alignment for buffer planning. Default: 64.
56    pub arena_alignment: usize,
57    /// Panic at compile time if fusion diagnostics report missed patterns.
58    pub assert_fusion_clean: bool,
59    /// Backend op claim set for backend-aware fusion + post-fusion
60    /// legalization. Set by [`Backend::compile`] implementations.
61    pub supported_ops: Option<&'static [OpKind]>,
62    /// When set, specialize symbolic dims before backend lowering.
63    pub dim_binding: Option<rlx_ir::DimBinding>,
64    /// Bake fixed param tensors into constants before DCE / constant folding.
65    pub param_bindings: Option<HashMap<String, Vec<f32>>>,
66    /// Packed U8 GGUF weights for `Op::Param` nodes in `DequantMatMul` (TPU HLO bake).
67    pub quant_param_bindings: Option<HashMap<String, Vec<u8>>>,
68    /// Native vs common IR lowering ([`KernelDispatchConfig`], `RLX_KERNEL_DISPATCH=common`).
69    pub kernel_dispatch: KernelDispatchConfig,
70    /// Prevent Metal from lowering the graph through MPSGraph.
71    ///
72    /// This keeps the regular Metal thunk schedule while retaining all other
73    /// compile options. Off by default.
74    pub disable_mpsgraph: bool,
75    /// On backends that list `OpKind::Scan` (host-fallback), Scans with
76    /// `length <= scan_unroll_max_length` are IR-unrolled into body replicas so
77    /// they run as ordinary device ops. Longer Scans keep `ScanHost`.
78    /// Default: **64**. `0` disables unroll on claiming backends; `u32::MAX`
79    /// prefers unroll for every eligible Scan.
80    pub scan_unroll_max_length: u32,
81    /// Hoist the param-invariant subgraph (compute that depends only on weights,
82    /// never on `Op::Input`) into a separate "prepare" graph run ONCE, feeding
83    /// its outputs into the main graph across forwards. Complements
84    /// `param_bindings` (which folds such compute away at compile when weight
85    /// VALUES are known); this handles run-time-only weights. Opt-in; default
86    /// off (also settable via `RLX_CACHE_PARAM_INVARIANT=1`).
87    pub cache_param_invariant: bool,
88    /// Static provable NaN/Inf lint during compile (`RLX_LINT_NUMERICS`).
89    pub lint_numerics: bool,
90    /// Print fusion before/after report (`RLX_FUSION_REPORT`).
91    pub fusion_report: bool,
92    /// Skip Conv+Bias+Act fusion (`RLX_DISABLE_CONV_BIAS_ACT_FUSION`).
93    pub disable_conv_bias_act_fusion: bool,
94    /// Disable IO-gated peaks-only fusion (`RLX_NO_IO_PEAKS_OUTPUT`).
95    pub no_io_peaks_output: bool,
96}
97
98impl Default for CompileOptions {
99    fn default() -> Self {
100        Self::from_env()
101    }
102}
103
104impl CompileOptions {
105    /// Defaults with **no** process-env overlay (tests / explicit builders).
106    pub fn baseline() -> Self {
107        Self {
108            precision: Precision::F32,
109            policy: None,
110            scaled_quant: None,
111            rng: rlx_ir::RngOptions::default(),
112            dce: true,
113            constant_folding: true,
114            verbose: false,
115            fusion_target: None,
116            fusion_opts: FusionOptions::default(),
117            arena_alignment: 64,
118            assert_fusion_clean: false,
119            supported_ops: None,
120            dim_binding: None,
121            param_bindings: None,
122            quant_param_bindings: None,
123            kernel_dispatch: KernelDispatchConfig::default(),
124            disable_mpsgraph: false,
125            scan_unroll_max_length: 64,
126            cache_param_invariant: false,
127            lint_numerics: false,
128            fusion_report: false,
129            disable_conv_bias_act_fusion: false,
130            no_io_peaks_output: false,
131        }
132    }
133
134    /// [`baseline`] plus Public/compile-layer `RLX_*` overlays from the registry.
135    pub fn from_env() -> Self {
136        let mut opts = Self::baseline();
137        opts.apply_env();
138        opts
139    }
140
141    /// Overlay compile-layer env keys (idempotent). Uses [`rlx_ir::env_registry`]
142    /// so aliases resolve.
143    pub fn apply_env(&mut self) {
144        use rlx_ir::env_registry;
145        if env_registry::flag("RLX_VERBOSE") {
146            self.verbose = true;
147        }
148        if env_registry::flag("RLX_CACHE_PARAM_INVARIANT") {
149            self.cache_param_invariant = true;
150        }
151        if env_registry::flag("RLX_DISABLE_MPSGRAPH") {
152            self.disable_mpsgraph = true;
153        }
154        self.kernel_dispatch = KernelDispatchConfig::from_env();
155        if env_registry::flag("RLX_LINT_NUMERICS") {
156            self.lint_numerics = true;
157        }
158        if env_registry::flag("RLX_FUSION_REPORT") {
159            self.fusion_report = true;
160        }
161        if env_registry::flag("RLX_DISABLE_CONV_BIAS_ACT_FUSION") {
162            self.disable_conv_bias_act_fusion = true;
163        }
164        if env_registry::flag("RLX_NO_IO_PEAKS_OUTPUT") {
165            self.no_io_peaks_output = true;
166        }
167    }
168
169    pub fn new() -> Self {
170        Self::from_env()
171    }
172
173    pub fn precision(mut self, p: Precision) -> Self {
174        self.precision = p;
175        self
176    }
177    pub fn policy(mut self, p: PrecisionPolicy) -> Self {
178        self.policy = Some(p);
179        self
180    }
181    /// Enable native low-precision GEMM for every 2-D matmul in the graph, in
182    /// the given element format + scale layout. Accepts any [`ScaledFormat`] via
183    /// [`ScaledQuantConfig`] — e.g. `ScaledQuantConfig::fp8_e4m3()`, or a
184    /// parameterized `ScaledFormat::custom(3, 0)` (`f4e3m0`).
185    pub fn scaled_quant(mut self, cfg: ScaledQuantConfig) -> Self {
186        self.scaled_quant = Some(cfg);
187        self
188    }
189    pub fn rng(mut self, rng: rlx_ir::RngOptions) -> Self {
190        self.rng = rng;
191        self
192    }
193    pub fn rng_backend(mut self, backend: rlx_ir::RngBackend) -> Self {
194        self.rng.backend = backend;
195        self
196    }
197    pub fn rng_seed(mut self, seed: u64) -> Self {
198        self.rng.seed = seed;
199        self
200    }
201    pub fn no_policy(mut self) -> Self {
202        self.policy = None;
203        self
204    }
205    pub fn with_dce(mut self, on: bool) -> Self {
206        self.dce = on;
207        self
208    }
209    pub fn with_constant_folding(mut self, on: bool) -> Self {
210        self.constant_folding = on;
211        self
212    }
213    pub fn with_verbose(mut self, on: bool) -> Self {
214        self.verbose = on;
215        self
216    }
217    pub fn fusion_target(mut self, target: FusionTarget) -> Self {
218        self.fusion_target = Some(target);
219        self
220    }
221    pub fn fusion_opts(mut self, opts: FusionOptions) -> Self {
222        self.fusion_opts = opts;
223        self
224    }
225    pub fn arena_alignment(mut self, bytes: usize) -> Self {
226        self.arena_alignment = bytes;
227        self
228    }
229    pub fn supported_ops(mut self, ops: &'static [OpKind]) -> Self {
230        self.supported_ops = Some(ops);
231        self
232    }
233    pub fn assert_fusion_clean(mut self, on: bool) -> Self {
234        self.assert_fusion_clean = on;
235        self
236    }
237    pub fn dim_binding(mut self, binding: rlx_ir::DimBinding) -> Self {
238        self.dim_binding = Some(binding);
239        self
240    }
241    pub fn param_bindings(mut self, bindings: HashMap<String, Vec<f32>>) -> Self {
242        self.param_bindings = Some(bindings);
243        self
244    }
245    /// Packed U8 weights for TPU GGUF `DequantMatMul` param bake at compile time.
246    pub fn quant_param_bindings(mut self, bindings: HashMap<String, Vec<u8>>) -> Self {
247        self.quant_param_bindings = Some(bindings);
248        self
249    }
250    pub fn kernel_dispatch(mut self, policy: KernelDispatchPolicy) -> Self {
251        self.kernel_dispatch.policy = policy;
252        self
253    }
254
255    pub fn kernel_dispatch_config(mut self, config: KernelDispatchConfig) -> Self {
256        self.kernel_dispatch = config;
257        self
258    }
259
260    /// Prefer IR-unrolling Scans with `length <= max` on Scan-claiming backends.
261    pub fn scan_unroll_max_length(mut self, max: u32) -> Self {
262        self.scan_unroll_max_length = max;
263        self
264    }
265
266    /// Force listed logical kernels to use common IR even when native is in `supported_ops`.
267    pub fn force_common_kinds(mut self, kinds: &'static [OpKind]) -> Self {
268        self.kernel_dispatch.force_common_kinds = kinds;
269        self
270    }
271
272    /// Keep listed logical kernels native even under `ForceCommon` / missing from `supported_ops`.
273    pub fn force_native_kinds(mut self, kinds: &'static [OpKind]) -> Self {
274        self.kernel_dispatch.force_native_kinds = kinds;
275        self
276    }
277
278    pub fn with_lint_numerics(mut self, on: bool) -> Self {
279        self.lint_numerics = on;
280        self
281    }
282
283    pub fn with_fusion_report(mut self, on: bool) -> Self {
284        self.fusion_report = on;
285        self
286    }
287
288    pub fn with_disable_mpsgraph(mut self, on: bool) -> Self {
289        self.disable_mpsgraph = on;
290        self
291    }
292
293    pub fn with_cache_param_invariant(mut self, on: bool) -> Self {
294        self.cache_param_invariant = on;
295        self
296    }
297
298    pub fn with_no_io_peaks_output(mut self, on: bool) -> Self {
299        self.no_io_peaks_output = on;
300        self
301    }
302
303    pub fn with_disable_conv_bias_act_fusion(mut self, on: bool) -> Self {
304        self.disable_conv_bias_act_fusion = on;
305        self
306    }
307}
rlx_runtime/options.rs

rlx_runtime/
options.rs
