rlx_ir/
mir.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//! **MIR** — mid-level IR.
6//!
7//! The fused, backend-neutral tensor DAG that [`rlx_opt`] runs fusion,
8//! precision, and legalization passes on. Today MIR is structurally
9//! identical to [`Graph`]; the newtype marks pipeline stage and gives
10//! us room to attach MIR-only metadata later (alias sets, layout hints).
11
12use crate::{Graph, Node, NodeId, Op};
13
14/// Mid-level module — optimizer input.
15#[cfg_attr(feature = "serialize", derive(serde::Serialize, serde::Deserialize))]
16#[derive(Debug, Clone, PartialEq)]
17pub struct MirModule {
18    inner: Graph,
19}
20
21/// MIR node / op aliases (same types as the legacy graph API).
22pub type MirNode = Node;
23pub type MirNodeId = NodeId;
24pub type MirOp = Op;
25
26impl MirModule {
27    pub fn new(name: impl Into<String>) -> Self {
28        Self {
29            inner: Graph::new(name),
30        }
31    }
32
33    pub fn from_graph(graph: Graph) -> Self {
34        Self { inner: graph }
35    }
36
37    pub fn into_graph(self) -> Graph {
38        self.inner
39    }
40
41    pub fn as_graph(&self) -> &Graph {
42        &self.inner
43    }
44
45    pub fn as_graph_mut(&mut self) -> &mut Graph {
46        &mut self.inner
47    }
48
49    pub fn name(&self) -> &str {
50        &self.inner.name
51    }
52
53    pub fn len(&self) -> usize {
54        self.inner.len()
55    }
56
57    pub fn is_empty(&self) -> bool {
58        self.inner.is_empty()
59    }
60
61    pub fn outputs(&self) -> &[NodeId] {
62        &self.inner.outputs
63    }
64
65    pub fn set_outputs(&mut self, outputs: Vec<NodeId>) {
66        self.inner.set_outputs(outputs);
67    }
68}
69
70impl From<Graph> for MirModule {
71    fn from(graph: Graph) -> Self {
72        Self::from_graph(graph)
73    }
74}
75
76impl From<MirModule> for Graph {
77    fn from(mir: MirModule) -> Self {
78        mir.into_graph()
79    }
80}
81
82impl std::fmt::Display for MirModule {
83    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
84        write!(f, "mir @{}", self.inner)
85    }
86}
rlx_ir/mir.rs

rlx_ir/
mir.rs
