rlx_ir/
module.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//! [`GraphModule`] — unified higher-order DX over HIR / MIR / LIR.
6
7use std::ops::{Deref, DerefMut};
8
9use crate::hir::{FusionPolicy, HirModule, HirNodeId, LowerError};
10use crate::inspect::{inspect_hir, inspect_lir, inspect_mir};
11use crate::lir::LirModule;
12use crate::mir::MirModule;
13use crate::op::Activation;
14use crate::op::MaskKind;
15use crate::quant::QuantScheme;
16use crate::{Graph, NodeId, Op, Shape};
17
18/// Which stage of the HIR → MIR → LIR pipeline a [`GraphModule`] holds.
19#[derive(Debug, Clone, Copy, PartialEq, Eq)]
20pub enum GraphStage {
21    Hir,
22    Mir,
23    Lir,
24}
25
26#[derive(Debug, Clone)]
27enum Stage {
28    Hir(HirModule),
29    Mir(MirModule),
30    Lir(LirModule),
31}
32
33/// Unified model module — primary builder surface above HIR/MIR/LIR.
34#[derive(Debug, Clone)]
35pub struct GraphModule {
36    stage: Stage,
37}
38
39impl GraphModule {
40    pub fn define(
41        name: impl Into<String>,
42        build: impl FnOnce(&mut HirModule) -> HirNodeId,
43    ) -> Self {
44        let mut hir = HirModule::new(name);
45        let out = build(&mut hir);
46        hir.set_outputs(vec![out]);
47        Self {
48            stage: Stage::Hir(hir),
49        }
50    }
51
52    /// Start an empty HIR-stage module (like [`Graph::new`] for MIR).
53    pub fn hir(name: impl Into<String>) -> Self {
54        Self {
55            stage: Stage::Hir(HirModule::new(name)),
56        }
57    }
58
59    /// Start an empty MIR-stage module.
60    pub fn mir(name: impl Into<String>) -> Self {
61        Self {
62            stage: Stage::Mir(MirModule::new(name)),
63        }
64    }
65
66    pub fn from_hir(hir: HirModule) -> Self {
67        Self {
68            stage: Stage::Hir(hir),
69        }
70    }
71
72    pub fn from_graph(graph: Graph) -> Self {
73        Self {
74            stage: Stage::Mir(MirModule::from_graph(graph)),
75        }
76    }
77
78    pub fn from_mir(mir: MirModule) -> Self {
79        Self {
80            stage: Stage::Mir(mir),
81        }
82    }
83
84    pub fn from_lir(lir: LirModule) -> Self {
85        Self {
86            stage: Stage::Lir(lir),
87        }
88    }
89
90    pub fn block(
91        hir: &mut HirModule,
92        name: impl Into<String>,
93        build: impl FnOnce(&mut HirModule) -> HirNodeId,
94    ) -> HirNodeId {
95        hir.named(name, build)
96    }
97
98    /// Declared dtype of every named `Input` / `Param`, when this module is at
99    /// a stage that has a graph.
100    ///
101    /// The f32-shaped entry points (`CompiledGraph::run`,
102    /// `CompiledGraph::set_param`) need this to tell a boundary that is
103    /// genuinely `f32` from one that is `F16`/`BF16` and merely being *fed* as
104    /// f32. Without it those bytes are handed to the backend verbatim and read
105    /// as pairs of halves — no error, plausible-looking numbers.
106    ///
107    /// Empty at HIR stage, where names are not yet bound to graph nodes.
108    pub fn io_dtypes(&self) -> Vec<(String, crate::DType)> {
109        let graph = match &self.stage {
110            Stage::Mir(m) => m.as_graph(),
111            Stage::Lir(l) => l.as_graph(),
112            Stage::Hir(_) => return Vec::new(),
113        };
114        graph
115            .nodes()
116            .iter()
117            .filter_map(|n| match &n.op {
118                crate::Op::Input { name } | crate::Op::Param { name } => {
119                    Some((name.clone(), n.shape.dtype()))
120                }
121                _ => None,
122            })
123            .collect()
124    }
125
126    pub fn fusion_policy(&self) -> Option<FusionPolicy> {
127        self.as_hir().map(|h| h.fusion_policy)
128    }
129
130    pub fn with_fusion_policy(mut self, policy: FusionPolicy) -> Self {
131        if let Stage::Hir(h) = &mut self.stage {
132            h.fusion_policy = policy;
133        } else {
134            panic!("GraphModule::with_fusion_policy requires HIR stage");
135        }
136        self
137    }
138
139    /// Set graph outputs at the current stage.
140    ///
141    /// At HIR stage accepts [`HirNodeId`] values; at MIR/LIR the same
142    /// indices map to [`NodeId`] (both are insertion-order node ids).
143    pub fn set_outputs(&mut self, outputs: Vec<HirNodeId>) {
144        match &mut self.stage {
145            Stage::Hir(h) => h.set_outputs(outputs),
146            Stage::Mir(m) => m.set_outputs(outputs.into_iter().map(|h| NodeId(h.0)).collect()),
147            Stage::Lir(l) => l
148                .mir
149                .set_outputs(outputs.into_iter().map(|h| NodeId(h.0)).collect()),
150        }
151    }
152
153    pub fn set_hir_outputs(&mut self, outputs: Vec<HirNodeId>) {
154        self.set_outputs(outputs);
155    }
156
157    /// Finish HIR construction and set the module output.
158    pub fn finish_hir(mut self, output: HirNodeId) -> Self {
159        self.set_hir_outputs(vec![output]);
160        self
161    }
162
163    fn hir_mut(&mut self) -> &mut HirModule {
164        self.as_hir_mut()
165            .expect("GraphModule: HIR builder methods require HIR stage — use GraphModule::hir() or Graph::define()")
166    }
167
168    // ── HIR block builders (forward to HirModule) ─────────────────
169
170    pub fn input(&mut self, name: impl Into<String>, shape: Shape) -> HirNodeId {
171        match &mut self.stage {
172            Stage::Hir(h) => h.input(name, shape),
173            Stage::Mir(m) => {
174                let id = m.as_graph_mut().input(name, shape);
175                HirNodeId(id.0)
176            }
177            Stage::Lir(l) => {
178                let id = l.mir.as_graph_mut().input(name, shape);
179                HirNodeId(id.0)
180            }
181        }
182    }
183
184    pub fn param(&mut self, name: impl Into<String>, shape: Shape) -> HirNodeId {
185        match &mut self.stage {
186            Stage::Hir(h) => h.param(name, shape),
187            Stage::Mir(m) => {
188                let id = m.as_graph_mut().param(name, shape);
189                HirNodeId(id.0)
190            }
191            Stage::Lir(l) => {
192                let id = l.mir.as_graph_mut().param(name, shape);
193                HirNodeId(id.0)
194            }
195        }
196    }
197
198    pub fn linear(
199        &mut self,
200        x: HirNodeId,
201        weight: HirNodeId,
202        bias: Option<HirNodeId>,
203        activation: Option<Activation>,
204        out_shape: Shape,
205    ) -> HirNodeId {
206        self.hir_mut()
207            .linear(x, weight, bias, activation, out_shape)
208    }
209
210    pub fn linear_fused(
211        &mut self,
212        x: HirNodeId,
213        weight: HirNodeId,
214        bias: HirNodeId,
215        activation: Option<Activation>,
216        out_shape: Shape,
217    ) -> HirNodeId {
218        self.hir_mut()
219            .linear_fused(x, weight, bias, activation, out_shape)
220    }
221
222    pub fn shared_linear_pair(
223        &mut self,
224        x: HirNodeId,
225        w_first: HirNodeId,
226        w_second: HirNodeId,
227        out_shape: Shape,
228    ) -> (HirNodeId, HirNodeId) {
229        self.hir_mut()
230            .shared_linear_pair(x, w_first, w_second, out_shape)
231    }
232
233    pub fn swiglu_ffn(
234        &mut self,
235        x: HirNodeId,
236        up_w: HirNodeId,
237        gate_w: HirNodeId,
238        down_w: HirNodeId,
239        out_shape: Shape,
240    ) -> HirNodeId {
241        self.hir_mut()
242            .swiglu_ffn(x, up_w, gate_w, down_w, out_shape)
243    }
244
245    pub fn residual_rms_norm(
246        &mut self,
247        x: HirNodeId,
248        residual: HirNodeId,
249        gamma: HirNodeId,
250        beta: HirNodeId,
251        eps: f32,
252        out_shape: Shape,
253    ) -> HirNodeId {
254        self.hir_mut()
255            .residual_rms_norm(x, residual, gamma, beta, eps, out_shape)
256    }
257
258    pub fn attention(
259        &mut self,
260        q: HirNodeId,
261        k: HirNodeId,
262        v: HirNodeId,
263        mask: Option<HirNodeId>,
264        num_heads: usize,
265        head_dim: usize,
266        mask_kind: MaskKind,
267        out_shape: Shape,
268    ) -> HirNodeId {
269        self.hir_mut()
270            .attention(q, k, v, mask, num_heads, head_dim, mask_kind, out_shape)
271    }
272
273    pub fn depthwise_conv1d_causal(
274        &mut self,
275        input: HirNodeId,
276        weight: HirNodeId,
277        left_pad: HirNodeId,
278        kernel_size: usize,
279        out_shape: Shape,
280    ) -> HirNodeId {
281        self.hir_mut()
282            .depthwise_conv1d_causal(input, weight, left_pad, kernel_size, out_shape)
283    }
284
285    pub fn dequant_matmul(
286        &mut self,
287        x: HirNodeId,
288        w: HirNodeId,
289        scale: Option<HirNodeId>,
290        zp: Option<HirNodeId>,
291        scheme: QuantScheme,
292        out_shape: Shape,
293    ) -> HirNodeId {
294        self.hir_mut()
295            .dequant_matmul(x, w, scale, zp, scheme, out_shape)
296    }
297
298    pub fn gated_delta_net(
299        &mut self,
300        q: HirNodeId,
301        k: HirNodeId,
302        v: HirNodeId,
303        g: HirNodeId,
304        beta: HirNodeId,
305        state_size: usize,
306        out_shape: Shape,
307    ) -> HirNodeId {
308        self.hir_mut()
309            .gated_delta_net(q, k, v, g, beta, state_size, out_shape)
310    }
311
312    pub fn gated_delta_net_carry(
313        &mut self,
314        q: HirNodeId,
315        k: HirNodeId,
316        v: HirNodeId,
317        g: HirNodeId,
318        beta: HirNodeId,
319        state: HirNodeId,
320        state_size: usize,
321        out_shape: Shape,
322    ) -> HirNodeId {
323        self.hir_mut()
324            .gated_delta_net_carry(q, k, v, g, beta, state, state_size, out_shape)
325    }
326
327    pub fn rope(
328        &mut self,
329        x: HirNodeId,
330        cos: HirNodeId,
331        sin: HirNodeId,
332        head_dim: usize,
333        n_rot: usize,
334        out_shape: Shape,
335    ) -> HirNodeId {
336        self.hir_mut().rope(x, cos, sin, head_dim, n_rot, out_shape)
337    }
338
339    pub fn rms_norm(
340        &mut self,
341        x: HirNodeId,
342        gamma: HirNodeId,
343        beta: HirNodeId,
344        eps: f32,
345        out_shape: Shape,
346    ) -> HirNodeId {
347        self.hir_mut().rms_norm(x, gamma, beta, eps, out_shape)
348    }
349
350    pub fn hir_mir(&mut self, op: Op, inputs: Vec<HirNodeId>, shape: Shape) -> HirNodeId {
351        self.hir_mut().mir(op, inputs, shape)
352    }
353
354    pub fn named(
355        &mut self,
356        name: impl Into<String>,
357        build: impl FnOnce(&mut HirModule) -> HirNodeId,
358    ) -> HirNodeId {
359        self.hir_mut().named(name, build)
360    }
361
362    pub fn stage(&self) -> GraphStage {
363        match &self.stage {
364            Stage::Hir(_) => GraphStage::Hir,
365            Stage::Mir(_) => GraphStage::Mir,
366            Stage::Lir(_) => GraphStage::Lir,
367        }
368    }
369
370    pub fn name(&self) -> &str {
371        match &self.stage {
372            Stage::Hir(h) => &h.name,
373            Stage::Mir(m) => m.name(),
374            Stage::Lir(l) => l.name(),
375        }
376    }
377
378    pub fn lower(self) -> Result<Self, LowerError> {
379        match self.stage {
380            Stage::Hir(hir) => Ok(Self {
381                stage: Stage::Mir(hir.lower_to_mir()?),
382            }),
383            other => Ok(Self { stage: other }),
384        }
385    }
386
387    pub fn into_hir(self) -> Option<HirModule> {
388        match self.stage {
389            Stage::Hir(h) => Some(h),
390            _ => None,
391        }
392    }
393
394    pub fn into_mir(self) -> Result<MirModule, LowerError> {
395        match self.stage {
396            Stage::Hir(hir) => hir.lower_to_mir(),
397            Stage::Mir(m) => Ok(m),
398            Stage::Lir(l) => Ok(l.mir),
399        }
400    }
401
402    pub fn into_lir(self) -> Option<LirModule> {
403        match self.stage {
404            Stage::Lir(l) => Some(l),
405            _ => None,
406        }
407    }
408
409    pub fn into_graph(self) -> Result<Graph, LowerError> {
410        Ok(self.into_mir()?.into_graph())
411    }
412
413    pub fn as_hir(&self) -> Option<&HirModule> {
414        match &self.stage {
415            Stage::Hir(h) => Some(h),
416            _ => None,
417        }
418    }
419
420    pub fn as_hir_mut(&mut self) -> Option<&mut HirModule> {
421        match &mut self.stage {
422            Stage::Hir(h) => Some(h),
423            _ => None,
424        }
425    }
426
427    pub fn as_mir(&self) -> Option<&MirModule> {
428        match &self.stage {
429            Stage::Mir(m) => Some(m),
430            Stage::Lir(l) => Some(&l.mir),
431            _ => None,
432        }
433    }
434
435    pub fn as_lir(&self) -> Option<&LirModule> {
436        match &self.stage {
437            Stage::Lir(l) => Some(l),
438            _ => None,
439        }
440    }
441
442    pub fn as_graph(&self) -> Option<&Graph> {
443        match &self.stage {
444            Stage::Mir(m) => Some(m.as_graph()),
445            Stage::Lir(l) => Some(l.as_graph()),
446            Stage::Hir(_) => None,
447        }
448    }
449
450    pub fn inspect(&self) -> String {
451        match &self.stage {
452            Stage::Hir(h) => inspect_hir(h),
453            Stage::Mir(m) => inspect_mir(m),
454            Stage::Lir(l) => inspect_lir(l),
455        }
456    }
457}
458
459impl Deref for GraphModule {
460    type Target = Graph;
461
462    fn deref(&self) -> &Graph {
463        self.as_graph()
464            .expect("GraphModule: HIR stage — call lower() before accessing MIR Graph")
465    }
466}
467
468impl DerefMut for GraphModule {
469    fn deref_mut(&mut self) -> &mut Graph {
470        match &mut self.stage {
471            Stage::Mir(m) => m.as_graph_mut(),
472            Stage::Lir(l) => l.mir.as_graph_mut(),
473            Stage::Hir(_) => panic!("GraphModule: HIR stage — use as_hir_mut() or lower() first"),
474        }
475    }
476}
477
478impl From<Graph> for GraphModule {
479    fn from(graph: Graph) -> Self {
480        Self::from_graph(graph)
481    }
482}
483
484impl TryFrom<GraphModule> for Graph {
485    type Error = LowerError;
486
487    fn try_from(module: GraphModule) -> Result<Self, LowerError> {
488        module.into_graph()
489    }
490}
491
492impl From<MirModule> for GraphModule {
493    fn from(mir: MirModule) -> Self {
494        Self::from_mir(mir)
495    }
496}
497
498impl From<HirModule> for GraphModule {
499    fn from(hir: HirModule) -> Self {
500        Self::from_hir(hir)
501    }
502}
503
504impl std::fmt::Display for GraphModule {
505    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
506        match &self.stage {
507            Stage::Hir(h) => write!(f, "{h}"),
508            Stage::Mir(m) => write!(f, "{m}"),
509            Stage::Lir(l) => write!(f, "lir @{}", l.name()),
510        }
511    }
512}
513
514#[cfg(test)]
515mod tests {
516    use super::*;
517    use crate::DType;
518    use crate::Graph;
519    use crate::Shape;
520
521    fn f32_shape(d: &[usize]) -> Shape {
522        Shape::new(d, DType::F32)
523    }
524
525    #[test]
526    fn define_lowers_to_mir_graph() {
527        let module = GraphModule::define("m", |m| {
528            let x = m.input("x", f32_shape(&[2, 8]));
529            let w = m.param("w", f32_shape(&[8, 8]));
530            m.linear(x, w, None, None, f32_shape(&[2, 8]))
531        });
532        assert_eq!(module.stage(), GraphStage::Hir);
533        let module = module.lower().expect("lower");
534        assert_eq!(module.stage(), GraphStage::Mir);
535        assert!(module.len() >= 3);
536    }
537
538    #[test]
539    fn mir_module_deref_builds_graph() {
540        let mut module = GraphModule::mir("raw");
541        let x = module.input("x", f32_shape(&[4]));
542        module.set_outputs(vec![x]);
543        assert_eq!(module.len(), 1);
544    }
545
546    #[test]
547    fn hir_module_block_builders_via_graph_module() {
548        use crate::quant::QuantScheme;
549
550        let mut module = GraphModule::hir("layer");
551        let x = module.input("x", f32_shape(&[2, 128]));
552        let w = module.param("w", f32_shape(&[128, 128]));
553        let y = module.dequant_matmul(x, w, None, None, QuantScheme::GgufQ4K, f32_shape(&[2, 128]));
554        module.set_outputs(vec![y]);
555        assert_eq!(module.stage(), GraphStage::Hir);
556
557        let module = module.lower().expect("lower");
558        assert_eq!(module.stage(), GraphStage::Mir);
559        assert!(module.len() >= 3);
560    }
561
562    #[test]
563    fn graph_hir_entry_matches_define() {
564        let via_graph = Graph::hir("m");
565        let via_define = Graph::define("m", |m| {
566            let x = m.input("x", f32_shape(&[4]));
567            m.rms_norm(x, x, x, 1e-5, f32_shape(&[4]))
568        });
569        assert_eq!(via_graph.stage(), GraphStage::Hir);
570        assert_eq!(via_define.stage(), GraphStage::Hir);
571    }
572}
rlx_ir/module.rs

rlx_ir/
module.rs
