rlx_ir/
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//! Reflection over HIR/MIR/LIR — layout and structure without executing.
6//!
7//! Host code can introspect unspecialized templates (Slang front-end / reflection API)
8//! and specialized layouts independently of backend codegen.
9
10use crate::Shape;
11use crate::binding_manifest::BindingManifest;
12use crate::component::ModelComponent;
13use crate::hir::{HirModule, HirNodeId, HirOp};
14use crate::lir::LirModule;
15use crate::mir::MirModule;
16use crate::shape::DimBinding;
17
18/// Introspection of an unspecialized [`HirModule`] (loadModule analogue).
19#[derive(Debug, Clone, PartialEq, Eq)]
20pub struct HirReflection {
21    pub name: String,
22    pub node_count: usize,
23    pub fusion_policy: String,
24    pub inputs: Vec<(String, Shape)>,
25    pub params: Vec<(String, Shape)>,
26    pub outputs: Vec<Shape>,
27    pub block_labels: Vec<String>,
28}
29
30impl HirReflection {
31    pub fn from_hir(hir: &HirModule) -> Self {
32        let mut inputs = Vec::new();
33        let mut params = Vec::new();
34        let mut block_labels = Vec::new();
35        for node in hir.nodes().iter() {
36            let label = node
37                .name
38                .clone()
39                .unwrap_or_else(|| format!("{:?}", node.op));
40            match &node.op {
41                HirOp::Input { name } => inputs.push((name.clone(), node.shape.clone())),
42                HirOp::Param { name } => params.push((name.clone(), node.shape.clone())),
43                HirOp::LlamaDecoderBlock { .. }
44                | HirOp::SwiGLU
45                | HirOp::Attention { .. }
46                | HirOp::GatedDeltaNet { .. }
47                | HirOp::Qwen35MtpHead { .. } => block_labels.push(label),
48                _ => {}
49            }
50        }
51        let outputs = hir
52            .outputs
53            .iter()
54            .map(|&id| hir.node(id).shape.clone())
55            .collect();
56        HirReflection {
57            name: hir.name.clone(),
58            node_count: hir.nodes().len(),
59            fusion_policy: format!("{:?}", hir.fusion_policy),
60            inputs,
61            params,
62            outputs,
63            block_labels,
64        }
65    }
66}
67
68/// MIR-level summary after HIR lower (specializeType / graph shape probe).
69#[derive(Debug, Clone, PartialEq, Eq)]
70pub struct MirReflection {
71    pub name: String,
72    pub node_count: usize,
73    pub op_kinds: Vec<(String, usize)>,
74}
75
76impl MirReflection {
77    pub fn from_mir(mir: &MirModule) -> Self {
78        let g = mir.as_graph();
79        let mut counts = std::collections::HashMap::new();
80        for node in g.nodes() {
81            *counts.entry(format!("{:?}", node.op.kind())).or_default() += 1;
82        }
83        let mut op_kinds: Vec<_> = counts.into_iter().collect();
84        op_kinds.sort_by(|a, b| a.0.cmp(&b.0));
85        MirReflection {
86            name: g.name.clone(),
87            node_count: g.nodes().len(),
88            op_kinds,
89        }
90    }
91}
92
93/// Layout reflection from specialized LIR (getTypeLayout / parameter block).
94pub fn layout_from_lir(lir: &LirModule) -> BindingManifest {
95    BindingManifest::from_lir(lir)
96}
97
98/// Layout for a concrete [`ModelComponent`] binding without retaining the graph.
99pub fn layout_for_binding(lir: &LirModule, _component: &ModelComponent) -> BindingManifest {
100    layout_from_lir(lir)
101}
102
103/// Compare template vs specialized manifests (dims / arena may differ).
104#[derive(Debug, Clone, PartialEq, Eq)]
105pub struct ManifestDiff {
106    pub template_arena: usize,
107    pub specialized_arena: usize,
108    pub params_only_in_template: Vec<String>,
109    pub params_only_in_specialized: Vec<String>,
110}
111
112impl ManifestDiff {
113    pub fn compare(template: &BindingManifest, specialized: &BindingManifest) -> Self {
114        let t: std::collections::HashSet<_> = template.param_names().collect();
115        let s: std::collections::HashSet<_> = specialized.param_names().collect();
116        Self {
117            template_arena: template.arena_size,
118            specialized_arena: specialized.arena_size,
119            params_only_in_template: t.difference(&s).map(|x| (*x).to_string()).collect(),
120            params_only_in_specialized: s.difference(&t).map(|x| (*x).to_string()).collect(),
121        }
122    }
123}
124
125/// Block specialization choice (coarse-grained type argument).
126#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
127pub enum BlockSpecialization {
128    Default,
129    FusedTransformerLayer,
130    UnfusedPrimitives,
131}
132
133/// Record of a specialization decision for tooling (specializeType analogue).
134#[derive(Debug, Clone, PartialEq, Eq)]
135pub struct SpecializeBlockRecord {
136    pub node: HirNodeId,
137    pub label: String,
138    pub choice: BlockSpecialization,
139}
140
141/// Probe which HIR blocks would take a given specialization (static; no MIR mutate).
142pub fn probe_block_specialization(
143    hir: &HirModule,
144    choice: BlockSpecialization,
145) -> Vec<SpecializeBlockRecord> {
146    hir.nodes()
147        .iter()
148        .filter_map(|node| {
149            let fused = matches!(
150                node.op,
151                HirOp::LlamaDecoderBlock { .. } | HirOp::SwiGLU | HirOp::GatedDeltaNet { .. }
152            );
153            if !fused {
154                return None;
155            }
156            let effective = choice;
157            Some(SpecializeBlockRecord {
158                node: node.id,
159                label: node
160                    .name
161                    .clone()
162                    .unwrap_or_else(|| format!("{:?}", node.op)),
163                choice: effective,
164            })
165        })
166        .collect()
167}
168
169/// Binding-only layout probe when only [`DimBinding`] is known (no full compile).
170pub fn symbolic_layout_hint(binding: &DimBinding) -> String {
171    format!("DimBinding({:?})", binding)
172}
173
174#[cfg(test)]
175mod tests {
176    use super::*;
177    use crate::hir::HirMut;
178    use crate::{DType, HirModule};
179
180    #[test]
181    fn hir_reflection_lists_inputs() {
182        let mut hir = HirModule::new("t");
183        let mut gb = HirMut::new(&mut hir);
184        let _x = gb.input("x", Shape::new(&[1, 4], DType::F32));
185        let _w = gb.param("w", Shape::new(&[4, 2], DType::F32));
186        let r = HirReflection::from_hir(&hir);
187        assert_eq!(r.inputs.len(), 1);
188        assert_eq!(r.params.len(), 1);
189    }
190}
rlx_ir/reflect.rs

rlx_ir/
reflect.rs
