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}
82
83impl Default for CompileOptions {
84 fn default() -> Self {
85 Self {
86 precision: Precision::F32,
87 policy: None,
88 scaled_quant: None,
89 rng: rlx_ir::RngOptions::default(),
90 dce: true,
91 constant_folding: true,
92 verbose: false,
93 fusion_target: None,
94 fusion_opts: FusionOptions::default(),
95 arena_alignment: 64,
96 assert_fusion_clean: false,
97 supported_ops: None,
98 dim_binding: None,
99 param_bindings: None,
100 quant_param_bindings: None,
101 kernel_dispatch: KernelDispatchConfig::from_env(),
102 }
103 }
104}
105
106impl CompileOptions {
107 pub fn new() -> Self {
108 Self::default()
109 }
110
111 pub fn precision(mut self, p: Precision) -> Self {
112 self.precision = p;
113 self
114 }
115 pub fn policy(mut self, p: PrecisionPolicy) -> Self {
116 self.policy = Some(p);
117 self
118 }
119 /// Enable native low-precision GEMM for every 2-D matmul in the graph, in
120 /// the given element format + scale layout. Accepts any [`ScaledFormat`] via
121 /// [`ScaledQuantConfig`] — e.g. `ScaledQuantConfig::fp8_e4m3()`, or a
122 /// parameterized `ScaledFormat::custom(3, 0)` (`f4e3m0`).
123 pub fn scaled_quant(mut self, cfg: ScaledQuantConfig) -> Self {
124 self.scaled_quant = Some(cfg);
125 self
126 }
127 pub fn rng(mut self, rng: rlx_ir::RngOptions) -> Self {
128 self.rng = rng;
129 self
130 }
131 pub fn rng_backend(mut self, backend: rlx_ir::RngBackend) -> Self {
132 self.rng.backend = backend;
133 self
134 }
135 pub fn rng_seed(mut self, seed: u64) -> Self {
136 self.rng.seed = seed;
137 self
138 }
139 pub fn no_policy(mut self) -> Self {
140 self.policy = None;
141 self
142 }
143 pub fn with_dce(mut self, on: bool) -> Self {
144 self.dce = on;
145 self
146 }
147 pub fn with_constant_folding(mut self, on: bool) -> Self {
148 self.constant_folding = on;
149 self
150 }
151 pub fn with_verbose(mut self, on: bool) -> Self {
152 self.verbose = on;
153 self
154 }
155 pub fn fusion_target(mut self, target: FusionTarget) -> Self {
156 self.fusion_target = Some(target);
157 self
158 }
159 pub fn fusion_opts(mut self, opts: FusionOptions) -> Self {
160 self.fusion_opts = opts;
161 self
162 }
163 pub fn arena_alignment(mut self, bytes: usize) -> Self {
164 self.arena_alignment = bytes;
165 self
166 }
167 pub fn supported_ops(mut self, ops: &'static [OpKind]) -> Self {
168 self.supported_ops = Some(ops);
169 self
170 }
171 pub fn assert_fusion_clean(mut self, on: bool) -> Self {
172 self.assert_fusion_clean = on;
173 self
174 }
175 pub fn dim_binding(mut self, binding: rlx_ir::DimBinding) -> Self {
176 self.dim_binding = Some(binding);
177 self
178 }
179 pub fn param_bindings(mut self, bindings: HashMap<String, Vec<f32>>) -> Self {
180 self.param_bindings = Some(bindings);
181 self
182 }
183 /// Packed U8 weights for TPU GGUF `DequantMatMul` param bake at compile time.
184 pub fn quant_param_bindings(mut self, bindings: HashMap<String, Vec<u8>>) -> Self {
185 self.quant_param_bindings = Some(bindings);
186 self
187 }
188 pub fn kernel_dispatch(mut self, policy: KernelDispatchPolicy) -> Self {
189 self.kernel_dispatch.policy = policy;
190 self
191 }
192
193 pub fn kernel_dispatch_config(mut self, config: KernelDispatchConfig) -> Self {
194 self.kernel_dispatch = config;
195 self
196 }
197
198 /// Force listed logical kernels to use common IR even when native is in `supported_ops`.
199 pub fn force_common_kinds(mut self, kinds: &'static [OpKind]) -> Self {
200 self.kernel_dispatch.force_common_kinds = kinds;
201 self
202 }
203
204 /// Keep listed logical kernels native even under `ForceCommon` / missing from `supported_ops`.
205 pub fn force_native_kinds(mut self, kinds: &'static [OpKind]) -> Self {
206 self.kernel_dispatch.force_native_kinds = kinds;
207 self
208 }
209}