rlx_runtime/options.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//! Unified compile options.
17//!
18//! Replaces the historical mix of `compile()`, `compile_with_precision()`,
19//! `compile_with_options()` with a single `Backend::compile(graph, &options)`.
20//! New compile-time knobs can be added to `CompileOptions` without
21//! changing the trait — backends just read what they care about.
22//!
23//! Builder-pattern API for ergonomics:
24//!
25//! ```rust,ignore
26//! let opts = CompileOptions::new()
27//! .precision(Precision::F16)
28//! .policy(PrecisionPolicy::AutoMixed)
29//! .with_dce(true)
30//! .with_constant_folding(true);
31//! ```
32
33use crate::Precision;
34use rlx_ir::OpKind;
35use rlx_ir::logical_kernel::{KernelDispatchConfig, KernelDispatchPolicy};
36use rlx_opt::rlx_compile::scaled_quant_insert::ScaledQuantConfig;
37use rlx_opt::{FusionOptions, FusionTarget, PrecisionPolicy};
38use std::collections::HashMap;
39
40/// All knobs the compile pipeline understands.
41/// Add new fields here rather than introducing new compile entry points.
42#[derive(Debug, Clone)]
43pub struct CompileOptions {
44 /// Target numeric precision for execution. Default: F32.
45 pub precision: Precision,
46 /// Optional per-op precision policy (mixed precision rewrite).
47 pub policy: Option<PrecisionPolicy>,
48 /// Opt-in native low-precision GEMM. When set, every 2-D `Op::MatMul` is
49 /// rewritten to a `ScaledMatMul` whose operands are dynamically quantized to
50 /// this element format + scale layout — any [`ScaledFormat`], including a
51 /// parameterized `Custom` (e.g. `f4e3m0`). Off by default; changes numerics.
52 pub scaled_quant: Option<ScaledQuantConfig>,
53 /// RNG policy for in-graph [`Op::RngNormal`] / [`Op::RngUniform`] nodes.
54 pub rng: rlx_ir::RngOptions,
55 /// Run dead-code elimination as part of compile. Default: true.
56 pub dce: bool,
57 /// Run constant folding. Default: true (cheap, only helps).
58 pub constant_folding: bool,
59 /// Verbose pass logging. Equivalent to `RLX_VERBOSE=1` or
60 /// [`rlx_ir::env::set("RLX_VERBOSE", "1")`].
61 pub verbose: bool,
62 /// Override fusion pipeline target (default: inferred from device).
63 pub fusion_target: Option<FusionTarget>,
64 /// Per-target fusion toggles (Metal env overrides, skip fusion, …).
65 pub fusion_opts: FusionOptions,
66 /// Arena alignment for buffer planning. Default: 64.
67 pub arena_alignment: usize,
68 /// Panic at compile time if fusion diagnostics report missed patterns.
69 pub assert_fusion_clean: bool,
70 /// Backend op claim set for backend-aware fusion + post-fusion
71 /// legalization. Set by [`Backend::compile`] implementations.
72 pub supported_ops: Option<&'static [OpKind]>,
73 /// When set, specialize symbolic dims before backend lowering.
74 pub dim_binding: Option<rlx_ir::DimBinding>,
75 /// Bake fixed param tensors into constants before DCE / constant folding.
76 pub param_bindings: Option<HashMap<String, Vec<f32>>>,
77 /// Packed U8 GGUF weights for `Op::Param` nodes in `DequantMatMul` (TPU HLO bake).
78 pub quant_param_bindings: Option<HashMap<String, Vec<u8>>>,
79 /// Native vs common IR lowering ([`KernelDispatchConfig`], `RLX_KERNEL_DISPATCH=common`).
80 pub kernel_dispatch: KernelDispatchConfig,
81 /// Prevent Metal from lowering the graph through MPSGraph.
82 ///
83 /// This keeps the regular Metal thunk schedule while retaining all other
84 /// compile options. Off by default.
85 pub disable_mpsgraph: bool,
86 /// On backends that list `OpKind::Scan` (host-fallback), Scans with
87 /// `length <= scan_unroll_max_length` are IR-unrolled into body replicas so
88 /// they run as ordinary device ops. Longer Scans keep `ScanHost`.
89 /// Default: **64**. `0` disables unroll on claiming backends; `u32::MAX`
90 /// prefers unroll for every eligible Scan.
91 pub scan_unroll_max_length: u32,
92 /// Hoist the param-invariant subgraph (compute that depends only on weights,
93 /// never on `Op::Input`) into a separate "prepare" graph run ONCE, feeding
94 /// its outputs into the main graph across forwards. Complements
95 /// `param_bindings` (which folds such compute away at compile when weight
96 /// VALUES are known); this handles run-time-only weights. Opt-in; default
97 /// off (also settable via `RLX_CACHE_PARAM_INVARIANT=1`).
98 pub cache_param_invariant: bool,
99}
100
101impl Default for CompileOptions {
102 fn default() -> Self {
103 Self {
104 precision: Precision::F32,
105 policy: None,
106 scaled_quant: None,
107 rng: rlx_ir::RngOptions::default(),
108 dce: true,
109 constant_folding: true,
110 verbose: false,
111 fusion_target: None,
112 fusion_opts: FusionOptions::default(),
113 arena_alignment: 64,
114 assert_fusion_clean: false,
115 supported_ops: None,
116 dim_binding: None,
117 param_bindings: None,
118 quant_param_bindings: None,
119 kernel_dispatch: KernelDispatchConfig::from_env(),
120 disable_mpsgraph: false,
121 scan_unroll_max_length: 64,
122 cache_param_invariant: rlx_ir::env::flag("RLX_CACHE_PARAM_INVARIANT"),
123 }
124 }
125}
126
127impl CompileOptions {
128 pub fn new() -> Self {
129 Self::default()
130 }
131
132 pub fn precision(mut self, p: Precision) -> Self {
133 self.precision = p;
134 self
135 }
136 pub fn policy(mut self, p: PrecisionPolicy) -> Self {
137 self.policy = Some(p);
138 self
139 }
140 /// Enable native low-precision GEMM for every 2-D matmul in the graph, in
141 /// the given element format + scale layout. Accepts any [`ScaledFormat`] via
142 /// [`ScaledQuantConfig`] — e.g. `ScaledQuantConfig::fp8_e4m3()`, or a
143 /// parameterized `ScaledFormat::custom(3, 0)` (`f4e3m0`).
144 pub fn scaled_quant(mut self, cfg: ScaledQuantConfig) -> Self {
145 self.scaled_quant = Some(cfg);
146 self
147 }
148 pub fn rng(mut self, rng: rlx_ir::RngOptions) -> Self {
149 self.rng = rng;
150 self
151 }
152 pub fn rng_backend(mut self, backend: rlx_ir::RngBackend) -> Self {
153 self.rng.backend = backend;
154 self
155 }
156 pub fn rng_seed(mut self, seed: u64) -> Self {
157 self.rng.seed = seed;
158 self
159 }
160 pub fn no_policy(mut self) -> Self {
161 self.policy = None;
162 self
163 }
164 pub fn with_dce(mut self, on: bool) -> Self {
165 self.dce = on;
166 self
167 }
168 pub fn with_constant_folding(mut self, on: bool) -> Self {
169 self.constant_folding = on;
170 self
171 }
172 pub fn with_verbose(mut self, on: bool) -> Self {
173 self.verbose = on;
174 self
175 }
176 pub fn fusion_target(mut self, target: FusionTarget) -> Self {
177 self.fusion_target = Some(target);
178 self
179 }
180 pub fn fusion_opts(mut self, opts: FusionOptions) -> Self {
181 self.fusion_opts = opts;
182 self
183 }
184 pub fn arena_alignment(mut self, bytes: usize) -> Self {
185 self.arena_alignment = bytes;
186 self
187 }
188 pub fn supported_ops(mut self, ops: &'static [OpKind]) -> Self {
189 self.supported_ops = Some(ops);
190 self
191 }
192 pub fn assert_fusion_clean(mut self, on: bool) -> Self {
193 self.assert_fusion_clean = on;
194 self
195 }
196 pub fn dim_binding(mut self, binding: rlx_ir::DimBinding) -> Self {
197 self.dim_binding = Some(binding);
198 self
199 }
200 pub fn param_bindings(mut self, bindings: HashMap<String, Vec<f32>>) -> Self {
201 self.param_bindings = Some(bindings);
202 self
203 }
204 /// Packed U8 weights for TPU GGUF `DequantMatMul` param bake at compile time.
205 pub fn quant_param_bindings(mut self, bindings: HashMap<String, Vec<u8>>) -> Self {
206 self.quant_param_bindings = Some(bindings);
207 self
208 }
209 pub fn kernel_dispatch(mut self, policy: KernelDispatchPolicy) -> Self {
210 self.kernel_dispatch.policy = policy;
211 self
212 }
213
214 pub fn kernel_dispatch_config(mut self, config: KernelDispatchConfig) -> Self {
215 self.kernel_dispatch = config;
216 self
217 }
218
219 /// Prefer IR-unrolling Scans with `length <= max` on Scan-claiming backends.
220 pub fn scan_unroll_max_length(mut self, max: u32) -> Self {
221 self.scan_unroll_max_length = max;
222 self
223 }
224
225 /// Force listed logical kernels to use common IR even when native is in `supported_ops`.
226 pub fn force_common_kinds(mut self, kinds: &'static [OpKind]) -> Self {
227 self.kernel_dispatch.force_common_kinds = kinds;
228 self
229 }
230
231 /// Keep listed logical kernels native even under `ForceCommon` / missing from `supported_ops`.
232 pub fn force_native_kinds(mut self, kinds: &'static [OpKind]) -> Self {
233 self.kernel_dispatch.force_native_kinds = kinds;
234 self
235 }
236}