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

rlx_runtime/
options.rs
