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::{FusionOptions, FusionTarget, PrecisionPolicy};
37use std::collections::HashMap;
38
39/// All knobs the compile pipeline understands.
40/// Add new fields here rather than introducing new compile entry points.
41#[derive(Debug, Clone)]
42pub struct CompileOptions {
43    /// Target numeric precision for execution. Default: F32.
44    pub precision: Precision,
45    /// Optional per-op precision policy (mixed precision rewrite).
46    pub policy: Option<PrecisionPolicy>,
47    /// RNG policy for in-graph [`Op::RngNormal`] / [`Op::RngUniform`] nodes.
48    pub rng: rlx_ir::RngOptions,
49    /// Run dead-code elimination as part of compile. Default: true.
50    pub dce: bool,
51    /// Run constant folding. Default: true (cheap, only helps).
52    pub constant_folding: bool,
53    /// Verbose pass logging. Equivalent to `RLX_VERBOSE=1` or
54    /// [`rlx_ir::env::set("RLX_VERBOSE", "1")`].
55    pub verbose: bool,
56    /// Override fusion pipeline target (default: inferred from device).
57    pub fusion_target: Option<FusionTarget>,
58    /// Per-target fusion toggles (Metal env overrides, skip fusion, …).
59    pub fusion_opts: FusionOptions,
60    /// Arena alignment for buffer planning. Default: 64.
61    pub arena_alignment: usize,
62    /// Panic at compile time if fusion diagnostics report missed patterns.
63    pub assert_fusion_clean: bool,
64    /// Backend op claim set for backend-aware fusion + post-fusion
65    /// legalization. Set by [`Backend::compile`] implementations.
66    pub supported_ops: Option<&'static [OpKind]>,
67    /// When set, specialize symbolic dims before backend lowering.
68    pub dim_binding: Option<rlx_ir::DimBinding>,
69    /// Bake fixed param tensors into constants before DCE / constant folding.
70    pub param_bindings: Option<HashMap<String, Vec<f32>>>,
71    /// Packed U8 GGUF weights for `Op::Param` nodes in `DequantMatMul` (TPU HLO bake).
72    pub quant_param_bindings: Option<HashMap<String, Vec<u8>>>,
73    /// Native vs common IR lowering ([`KernelDispatchConfig`], `RLX_KERNEL_DISPATCH=common`).
74    pub kernel_dispatch: KernelDispatchConfig,
75}
76
77impl Default for CompileOptions {
78    fn default() -> Self {
79        Self {
80            precision: Precision::F32,
81            policy: None,
82            rng: rlx_ir::RngOptions::default(),
83            dce: true,
84            constant_folding: true,
85            verbose: false,
86            fusion_target: None,
87            fusion_opts: FusionOptions::default(),
88            arena_alignment: 64,
89            assert_fusion_clean: false,
90            supported_ops: None,
91            dim_binding: None,
92            param_bindings: None,
93            quant_param_bindings: None,
94            kernel_dispatch: KernelDispatchConfig::from_env(),
95        }
96    }
97}
98
99impl CompileOptions {
100    pub fn new() -> Self {
101        Self::default()
102    }
103
104    pub fn precision(mut self, p: Precision) -> Self {
105        self.precision = p;
106        self
107    }
108    pub fn policy(mut self, p: PrecisionPolicy) -> Self {
109        self.policy = Some(p);
110        self
111    }
112    pub fn rng(mut self, rng: rlx_ir::RngOptions) -> Self {
113        self.rng = rng;
114        self
115    }
116    pub fn rng_backend(mut self, backend: rlx_ir::RngBackend) -> Self {
117        self.rng.backend = backend;
118        self
119    }
120    pub fn rng_seed(mut self, seed: u64) -> Self {
121        self.rng.seed = seed;
122        self
123    }
124    pub fn no_policy(mut self) -> Self {
125        self.policy = None;
126        self
127    }
128    pub fn with_dce(mut self, on: bool) -> Self {
129        self.dce = on;
130        self
131    }
132    pub fn with_constant_folding(mut self, on: bool) -> Self {
133        self.constant_folding = on;
134        self
135    }
136    pub fn with_verbose(mut self, on: bool) -> Self {
137        self.verbose = on;
138        self
139    }
140    pub fn fusion_target(mut self, target: FusionTarget) -> Self {
141        self.fusion_target = Some(target);
142        self
143    }
144    pub fn fusion_opts(mut self, opts: FusionOptions) -> Self {
145        self.fusion_opts = opts;
146        self
147    }
148    pub fn arena_alignment(mut self, bytes: usize) -> Self {
149        self.arena_alignment = bytes;
150        self
151    }
152    pub fn supported_ops(mut self, ops: &'static [OpKind]) -> Self {
153        self.supported_ops = Some(ops);
154        self
155    }
156    pub fn assert_fusion_clean(mut self, on: bool) -> Self {
157        self.assert_fusion_clean = on;
158        self
159    }
160    pub fn dim_binding(mut self, binding: rlx_ir::DimBinding) -> Self {
161        self.dim_binding = Some(binding);
162        self
163    }
164    pub fn param_bindings(mut self, bindings: HashMap<String, Vec<f32>>) -> Self {
165        self.param_bindings = Some(bindings);
166        self
167    }
168    /// Packed U8 weights for TPU GGUF `DequantMatMul` param bake at compile time.
169    pub fn quant_param_bindings(mut self, bindings: HashMap<String, Vec<u8>>) -> Self {
170        self.quant_param_bindings = Some(bindings);
171        self
172    }
173    pub fn kernel_dispatch(mut self, policy: KernelDispatchPolicy) -> Self {
174        self.kernel_dispatch.policy = policy;
175        self
176    }
177
178    pub fn kernel_dispatch_config(mut self, config: KernelDispatchConfig) -> Self {
179        self.kernel_dispatch = config;
180        self
181    }
182
183    /// Force listed logical kernels to use common IR even when native is in `supported_ops`.
184    pub fn force_common_kinds(mut self, kinds: &'static [OpKind]) -> Self {
185        self.kernel_dispatch.force_common_kinds = kinds;
186        self
187    }
188
189    /// Keep listed logical kernels native even under `ForceCommon` / missing from `supported_ops`.
190    pub fn force_native_kinds(mut self, kinds: &'static [OpKind]) -> Self {
191        self.kernel_dispatch.force_native_kinds = kinds;
192        self
193    }
194}
rlx_runtime/options.rs

rlx_runtime/
options.rs
