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
// 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 model component — one object drives specialization, compile cache, and binding.
//!
//! Mirrors Slang “shader components”: the same granularity selects **what to specialize**
//! (dims, dispatch, compilation mode) and **how host code binds** (via [`BindingManifest`]
//! after specialize). Works across eager, lazy, and AOT pipelines while keeping HIR/MIR/LIR.
use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};
use crate::logical_kernel::KernelDispatchConfig;
use crate::quant::QuantScheme;
use crate::variant::ModelVariant;
/// When the backend executable is produced relative to the host loop.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
pub enum CompilationMode {
/// Compile before first `run` (default inference).
#[default]
Eager,
/// Build template at load; specialize/compile on first use per variant.
Lazy,
/// Serialize LIR / executable to disk; load without re-running fusion.
Aot,
}
/// Full specialization + binding bundle (Slang shader-component analogue).
#[derive(Debug, Clone)]
pub struct ModelComponent {
pub variant: ModelVariant,
pub kernel_dispatch: KernelDispatchConfig,
pub compilation_mode: CompilationMode,
/// Hash of tier-1 [`CompileProfile`] or arch preset (see `rlx-flow` presets).
pub profile_key: u64,
/// Optional quant scheme affecting lowers and weight layout.
pub quant: Option<QuantScheme>,
/// Composite layer-stack fingerprint (homogeneous depth, pair nesting).
pub layer_composition_key: u64,
}
impl ModelComponent {
pub fn new(variant: ModelVariant) -> Self {
Self {
variant,
kernel_dispatch: KernelDispatchConfig::default(),
compilation_mode: CompilationMode::Eager,
profile_key: 0,
quant: None,
layer_composition_key: 0,
}
}
pub fn with_kernel_dispatch(mut self, config: KernelDispatchConfig) -> Self {
self.kernel_dispatch = config;
self
}
pub fn with_compilation_mode(mut self, mode: CompilationMode) -> Self {
self.compilation_mode = mode;
self
}
pub fn with_profile_key(mut self, key: u64) -> Self {
self.profile_key = key;
self
}
pub fn with_quant(mut self, scheme: QuantScheme) -> Self {
self.quant = Some(scheme);
self
}
pub fn with_layer_composition_key(mut self, key: u64) -> Self {
self.layer_composition_key = key;
self
}
/// Stable key for compile caches (variant + dispatch + profile + composition).
pub fn cache_key(&self) -> u64 {
let mut h = DefaultHasher::new();
self.variant.cache_key().hash(&mut h);
(self.kernel_dispatch.policy as u8).hash(&mut h);
for k in self.kernel_dispatch.force_common_kinds.iter() {
k.hash(&mut h);
}
for k in self.kernel_dispatch.force_native_kinds.iter() {
k.hash(&mut h);
}
self.compilation_mode.hash(&mut h);
self.profile_key.hash(&mut h);
if let Some(q) = &self.quant {
format!("{q:?}").hash(&mut h);
}
self.layer_composition_key.hash(&mut h);
h.finish()
}
pub fn dim_binding(&self) -> crate::DimBinding {
self.variant.dim_binding()
}
/// Stable on-disk prefix for [`rlx_runtime::AotCache`] (`{base}__{binding_hash}` per variant).
pub fn aot_disk_base(&self) -> String {
format!("rlx_{:016x}", self.cache_key())
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::ModelVariant;
use crate::logical_kernel::KernelDispatchPolicy;
#[test]
fn cache_key_changes_with_mode_and_profile() {
let v = ModelVariant::prefill(1, 8);
let a = ModelComponent::new(v.clone()).cache_key();
let b = ModelComponent::new(v.clone())
.with_compilation_mode(CompilationMode::Lazy)
.cache_key();
let c = ModelComponent::new(v)
.with_profile_key(42)
.with_kernel_dispatch(KernelDispatchConfig::new(KernelDispatchPolicy::ForceCommon))
.cache_key();
assert_ne!(a, b);
assert_ne!(a, c);
}
}