rlx_runtime/
reflect.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//! Model reflection services (Slang compiler/runtime API §5).
6//!
7//! Introspect unspecialized templates and specialized layouts across eager/lazy/AOT
8//! while preserving the HIR → MIR → LIR pipeline.
9
10use rlx_ir::hir::HirModule;
11use rlx_ir::{
12    BindingManifest, HirReflection, ManifestDiff, MirReflection, ModelComponent,
13    apply_hir_extensions, layout_from_lir,
14};
15use rlx_opt::CompileResult;
16
17use crate::Device;
18use crate::model_pipeline::ModelCompilePipeline;
19use crate::options::CompileOptions;
20use crate::stages;
21
22/// Loaded template + HIR reflection (front-end load).
23pub struct ModelReflection {
24    pub hir: HirReflection,
25    template: Option<CompileResult>,
26}
27
28impl ModelReflection {
29    /// Build HIR reflection only (no compile).
30    pub fn from_hir(hir: &HirModule) -> Self {
31        Self {
32            hir: HirReflection::from_hir(hir),
33            template: None,
34        }
35    }
36
37    /// Compile symbolic template on `device` and retain for specialize/layout.
38    pub fn load_hir_template(
39        device: Device,
40        hir: HirModule,
41        options: &CompileOptions,
42    ) -> Result<Self, rlx_ir::hir::LowerError> {
43        let mut opts = options.clone();
44        opts.dim_binding = None;
45        let hir_ref = HirReflection::from_hir(&hir);
46        let pipe = stages::pipeline_for(device, &opts);
47        let template = pipe.compile_hir(hir)?;
48        Ok(Self {
49            hir: hir_ref,
50            template: Some(template),
51        })
52    }
53
54    pub fn has_template(&self) -> bool {
55        self.template.is_some()
56    }
57
58    pub fn mir_summary(&self) -> Option<MirReflection> {
59        self.template
60            .as_ref()
61            .map(|t| MirReflection::from_mir(&t.lir.mir))
62    }
63
64    /// Template layout (symbolic dims may be unresolved in arena sizes).
65    pub fn template_layout(&self) -> Option<BindingManifest> {
66        self.template.as_ref().map(|t| layout_from_lir(&t.lir))
67    }
68
69    /// Specialized layout for a [`ModelComponent`] (getTypeLayout after specialize).
70    pub fn layout_for_component(
71        &self,
72        component: &ModelComponent,
73        device: Device,
74        options: &CompileOptions,
75    ) -> Option<BindingManifest> {
76        let template = self.template.as_ref()?;
77        let mut opts = options.clone();
78        opts.dim_binding = None;
79        let pipe = stages::pipeline_for(device, &opts);
80        let specialized = template.specialize(&pipe, &component.dim_binding());
81        Some(layout_from_lir(&specialized.lir))
82    }
83
84    pub fn manifest_diff_for_component(
85        &self,
86        component: &ModelComponent,
87        device: Device,
88        options: &CompileOptions,
89    ) -> Option<ManifestDiff> {
90        let t = self.template_layout()?;
91        let s = self.layout_for_component(component, device, options)?;
92        Some(ManifestDiff::compare(&t, &s))
93    }
94}
95
96/// Full specialize + compile entry (specializeEntryPoint analogue).
97pub fn specialize_entry<'a>(
98    pipeline: &'a mut ModelCompilePipeline,
99    component: &ModelComponent,
100    build_hir: impl FnOnce() -> HirModule,
101    options: &CompileOptions,
102) -> Result<&'a mut crate::CompiledGraph, rlx_ir::hir::LowerError> {
103    let key = component.cache_key();
104    let binding = component.dim_binding();
105    pipeline.get_or_compile(key, &binding, build_hir, options)
106}
107
108/// Apply HIR extensions then load template.
109pub fn load_hir_template_with_extensions(
110    device: Device,
111    mut hir: HirModule,
112    options: &CompileOptions,
113) -> Result<ModelReflection, rlx_ir::hir::LowerError> {
114    apply_hir_extensions(&mut hir);
115    ModelReflection::load_hir_template(device, hir, options)
116}
117
118#[cfg(test)]
119mod tests {
120    use super::*;
121    use rlx_ir::hir::HirMut;
122    use rlx_ir::{DType, HirModule, ModelVariant, Shape};
123
124    #[test]
125    fn reflection_loads_template_on_cpu() {
126        let device = Device::Cpu;
127        let hir = || {
128            let mut hir = HirModule::new("refl");
129            let mut gb = HirMut::new(&mut hir);
130            let x = gb.input("x", Shape::new(&[1, 4], DType::F32));
131            let w = gb.param("w", Shape::new(&[4, 2], DType::F32));
132            let y = hir.linear(x, w, None, None, Shape::new(&[1, 2], DType::F32));
133            hir.set_outputs(vec![y]);
134            hir
135        };
136        let refl =
137            ModelReflection::load_hir_template(device, hir(), &CompileOptions::new()).unwrap();
138        assert!(refl.has_template());
139        let layout = refl.template_layout().unwrap();
140        assert_eq!(layout.params[0].name, "w");
141        let comp = ModelComponent::new(ModelVariant::prefill(1, 4));
142        let spec_layout = refl
143            .layout_for_component(&comp, device, &CompileOptions::new())
144            .unwrap();
145        assert_eq!(spec_layout.params[0].name, "w");
146    }
147}
rlx_runtime/reflect.rs

rlx_runtime/
reflect.rs
