rlx_ir/
component.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3// SPDX-License-Identifier: MIT OR Apache-2.0
4
5//! Unified model component — one object drives specialization, compile cache, and binding.
6//!
7//! Mirrors Slang “shader components”: the same granularity selects **what to specialize**
8//! (dims, dispatch, compilation mode) and **how host code binds** (via [`BindingManifest`]
9//! after specialize). Works across eager, lazy, and AOT pipelines while keeping HIR/MIR/LIR.
10
11use std::collections::hash_map::DefaultHasher;
12use std::hash::{Hash, Hasher};
13
14use crate::logical_kernel::KernelDispatchConfig;
15use crate::quant::QuantScheme;
16use crate::variant::ModelVariant;
17
18/// When the backend executable is produced relative to the host loop.
19#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
20pub enum CompilationMode {
21    /// Compile before first `run` (default inference).
22    #[default]
23    Eager,
24    /// Build template at load; specialize/compile on first use per variant.
25    Lazy,
26    /// Serialize LIR / executable to disk; load without re-running fusion.
27    Aot,
28}
29
30/// Full specialization + binding bundle (Slang shader-component analogue).
31#[derive(Debug, Clone)]
32pub struct ModelComponent {
33    pub variant: ModelVariant,
34    pub kernel_dispatch: KernelDispatchConfig,
35    pub compilation_mode: CompilationMode,
36    /// Hash of tier-1 [`CompileProfile`] or arch preset (see `rlx-flow` presets).
37    pub profile_key: u64,
38    /// Optional quant scheme affecting lowers and weight layout.
39    pub quant: Option<QuantScheme>,
40    /// Composite layer-stack fingerprint (homogeneous depth, pair nesting).
41    pub layer_composition_key: u64,
42}
43
44impl ModelComponent {
45    pub fn new(variant: ModelVariant) -> Self {
46        Self {
47            variant,
48            kernel_dispatch: KernelDispatchConfig::default(),
49            compilation_mode: CompilationMode::Eager,
50            profile_key: 0,
51            quant: None,
52            layer_composition_key: 0,
53        }
54    }
55
56    pub fn with_kernel_dispatch(mut self, config: KernelDispatchConfig) -> Self {
57        self.kernel_dispatch = config;
58        self
59    }
60
61    pub fn with_compilation_mode(mut self, mode: CompilationMode) -> Self {
62        self.compilation_mode = mode;
63        self
64    }
65
66    pub fn with_profile_key(mut self, key: u64) -> Self {
67        self.profile_key = key;
68        self
69    }
70
71    pub fn with_quant(mut self, scheme: QuantScheme) -> Self {
72        self.quant = Some(scheme);
73        self
74    }
75
76    pub fn with_layer_composition_key(mut self, key: u64) -> Self {
77        self.layer_composition_key = key;
78        self
79    }
80
81    /// Stable key for compile caches (variant + dispatch + profile + composition).
82    pub fn cache_key(&self) -> u64 {
83        let mut h = DefaultHasher::new();
84        self.variant.cache_key().hash(&mut h);
85        (self.kernel_dispatch.policy as u8).hash(&mut h);
86        for k in self.kernel_dispatch.force_common_kinds.iter() {
87            k.hash(&mut h);
88        }
89        for k in self.kernel_dispatch.force_native_kinds.iter() {
90            k.hash(&mut h);
91        }
92        self.compilation_mode.hash(&mut h);
93        self.profile_key.hash(&mut h);
94        if let Some(q) = &self.quant {
95            format!("{q:?}").hash(&mut h);
96        }
97        self.layer_composition_key.hash(&mut h);
98        h.finish()
99    }
100
101    pub fn dim_binding(&self) -> crate::DimBinding {
102        self.variant.dim_binding()
103    }
104
105    /// Stable on-disk prefix for [`rlx_runtime::AotCache`] (`{base}__{binding_hash}` per variant).
106    pub fn aot_disk_base(&self) -> String {
107        format!("rlx_{:016x}", self.cache_key())
108    }
109}
110
111#[cfg(test)]
112mod tests {
113    use super::*;
114    use crate::ModelVariant;
115    use crate::logical_kernel::KernelDispatchPolicy;
116
117    #[test]
118    fn cache_key_changes_with_mode_and_profile() {
119        let v = ModelVariant::prefill(1, 8);
120        let a = ModelComponent::new(v.clone()).cache_key();
121        let b = ModelComponent::new(v.clone())
122            .with_compilation_mode(CompilationMode::Lazy)
123            .cache_key();
124        let c = ModelComponent::new(v)
125            .with_profile_key(42)
126            .with_kernel_dispatch(KernelDispatchConfig::new(KernelDispatchPolicy::ForceCommon))
127            .cache_key();
128        assert_ne!(a, b);
129        assert_ne!(a, c);
130    }
131}
rlx_ir/component.rs

rlx_ir/
component.rs
