1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
// RLX — versatile ML compiler + runtime.
// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, version 3.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program. If not, see <https://www.gnu.org/licenses/>.
//! Unified compile options.
//!
//! Replaces the historical mix of `compile()`, `compile_with_precision()`,
//! `compile_with_options()` with a single `Backend::compile(graph, &options)`.
//! New compile-time knobs can be added to `CompileOptions` without
//! changing the trait — backends just read what they care about.
//!
//! Builder-pattern API for ergonomics:
//!
//! ```rust,ignore
//! let opts = CompileOptions::new()
//! .precision(Precision::F16)
//! .policy(PrecisionPolicy::AutoMixed)
//! .with_dce(true)
//! .with_constant_folding(true);
//! ```
use crate::Precision;
use rlx_ir::OpKind;
use rlx_ir::logical_kernel::{KernelDispatchConfig, KernelDispatchPolicy};
use rlx_opt::{FusionOptions, FusionTarget, PrecisionPolicy};
use std::collections::HashMap;
/// All knobs the compile pipeline understands.
/// Add new fields here rather than introducing new compile entry points.
#[derive(Debug, Clone)]
pub struct CompileOptions {
/// Target numeric precision for execution. Default: F32.
pub precision: Precision,
/// Optional per-op precision policy (mixed precision rewrite).
pub policy: Option<PrecisionPolicy>,
/// Run dead-code elimination as part of compile. Default: true.
pub dce: bool,
/// Run constant folding. Default: true (cheap, only helps).
pub constant_folding: bool,
/// Verbose pass logging. Equivalent to `RLX_VERBOSE=1` or
/// [`rlx_ir::env::set("RLX_VERBOSE", "1")`].
pub verbose: bool,
/// Override fusion pipeline target (default: inferred from device).
pub fusion_target: Option<FusionTarget>,
/// Per-target fusion toggles (Metal env overrides, skip fusion, …).
pub fusion_opts: FusionOptions,
/// Arena alignment for buffer planning. Default: 64.
pub arena_alignment: usize,
/// Panic at compile time if fusion diagnostics report missed patterns.
pub assert_fusion_clean: bool,
/// Backend op claim set for backend-aware fusion + post-fusion
/// legalization. Set by [`Backend::compile`] implementations.
pub supported_ops: Option<&'static [OpKind]>,
/// When set, specialize symbolic dims before backend lowering.
pub dim_binding: Option<rlx_ir::DimBinding>,
/// Bake fixed param tensors into constants before DCE / constant folding.
pub param_bindings: Option<HashMap<String, Vec<f32>>>,
/// Native vs common IR lowering ([`KernelDispatchConfig`], `RLX_KERNEL_DISPATCH=common`).
pub kernel_dispatch: KernelDispatchConfig,
}
impl Default for CompileOptions {
fn default() -> Self {
Self {
precision: Precision::F32,
policy: None,
dce: true,
constant_folding: true,
verbose: false,
fusion_target: None,
fusion_opts: FusionOptions::default(),
arena_alignment: 64,
assert_fusion_clean: false,
supported_ops: None,
dim_binding: None,
param_bindings: None,
kernel_dispatch: KernelDispatchConfig::from_env(),
}
}
}
impl CompileOptions {
pub fn new() -> Self {
Self::default()
}
pub fn precision(mut self, p: Precision) -> Self {
self.precision = p;
self
}
pub fn policy(mut self, p: PrecisionPolicy) -> Self {
self.policy = Some(p);
self
}
pub fn no_policy(mut self) -> Self {
self.policy = None;
self
}
pub fn with_dce(mut self, on: bool) -> Self {
self.dce = on;
self
}
pub fn with_constant_folding(mut self, on: bool) -> Self {
self.constant_folding = on;
self
}
pub fn with_verbose(mut self, on: bool) -> Self {
self.verbose = on;
self
}
pub fn fusion_target(mut self, target: FusionTarget) -> Self {
self.fusion_target = Some(target);
self
}
pub fn fusion_opts(mut self, opts: FusionOptions) -> Self {
self.fusion_opts = opts;
self
}
pub fn arena_alignment(mut self, bytes: usize) -> Self {
self.arena_alignment = bytes;
self
}
pub fn supported_ops(mut self, ops: &'static [OpKind]) -> Self {
self.supported_ops = Some(ops);
self
}
pub fn assert_fusion_clean(mut self, on: bool) -> Self {
self.assert_fusion_clean = on;
self
}
pub fn dim_binding(mut self, binding: rlx_ir::DimBinding) -> Self {
self.dim_binding = Some(binding);
self
}
pub fn param_bindings(mut self, bindings: HashMap<String, Vec<f32>>) -> Self {
self.param_bindings = Some(bindings);
self
}
pub fn kernel_dispatch(mut self, policy: KernelDispatchPolicy) -> Self {
self.kernel_dispatch.policy = policy;
self
}
pub fn kernel_dispatch_config(mut self, config: KernelDispatchConfig) -> Self {
self.kernel_dispatch = config;
self
}
/// Force listed logical kernels to use common IR even when native is in `supported_ops`.
pub fn force_common_kinds(mut self, kinds: &'static [OpKind]) -> Self {
self.kernel_dispatch.force_common_kinds = kinds;
self
}
/// Keep listed logical kernels native even under `ForceCommon` / missing from `supported_ops`.
pub fn force_native_kinds(mut self, kinds: &'static [OpKind]) -> Self {
self.kernel_dispatch.force_native_kinds = kinds;
self
}
}