Skip to main content

rlx_ir/
mir.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3//
4// This program is free software: you can redistribute it and/or modify
5// it under the terms of the GNU General Public License as published by
6// the Free Software Foundation, version 3.
7//
8// This program is distributed in the hope that it will be useful,
9// but WITHOUT ANY WARRANTY; without even the implied warranty of
10// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
11// GNU General Public License for more details.
12//
13// You should have received a copy of the GNU General Public License
14// along with this program. If not, see <https://www.gnu.org/licenses/>.
15
16//! **MIR** — mid-level IR.
17//!
18//! The fused, backend-neutral tensor DAG that [`rlx_opt`] runs fusion,
19//! precision, and legalization passes on. Today MIR is structurally
20//! identical to [`Graph`]; the newtype marks pipeline stage and gives
21//! us room to attach MIR-only metadata later (alias sets, layout hints).
22
23use crate::{Graph, Node, NodeId, Op};
24
25/// Mid-level module — optimizer input.
26#[cfg_attr(feature = "serialize", derive(serde::Serialize, serde::Deserialize))]
27#[derive(Debug, Clone, PartialEq)]
28pub struct MirModule {
29    inner: Graph,
30}
31
32/// MIR node / op aliases (same types as the legacy graph API).
33pub type MirNode = Node;
34pub type MirNodeId = NodeId;
35pub type MirOp = Op;
36
37impl MirModule {
38    pub fn new(name: impl Into<String>) -> Self {
39        Self {
40            inner: Graph::new(name),
41        }
42    }
43
44    pub fn from_graph(graph: Graph) -> Self {
45        Self { inner: graph }
46    }
47
48    pub fn into_graph(self) -> Graph {
49        self.inner
50    }
51
52    pub fn as_graph(&self) -> &Graph {
53        &self.inner
54    }
55
56    pub fn as_graph_mut(&mut self) -> &mut Graph {
57        &mut self.inner
58    }
59
60    pub fn name(&self) -> &str {
61        &self.inner.name
62    }
63
64    pub fn len(&self) -> usize {
65        self.inner.len()
66    }
67
68    pub fn is_empty(&self) -> bool {
69        self.inner.is_empty()
70    }
71
72    pub fn outputs(&self) -> &[NodeId] {
73        &self.inner.outputs
74    }
75
76    pub fn set_outputs(&mut self, outputs: Vec<NodeId>) {
77        self.inner.set_outputs(outputs);
78    }
79}
80
81impl From<Graph> for MirModule {
82    fn from(graph: Graph) -> Self {
83        Self::from_graph(graph)
84    }
85}
86
87impl From<MirModule> for Graph {
88    fn from(mir: MirModule) -> Self {
89        mir.into_graph()
90    }
91}
92
93impl std::fmt::Display for MirModule {
94    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
95        write!(f, "mir @{}", self.inner)
96    }
97}