rlx_ir/
graph.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//! The computation graph — a DAG of typed tensor operations.
17//!
18//! Graphs are append-only during construction (like SSA). Nodes reference
19//! inputs by [`NodeId`], forming a directed acyclic graph. The graph
20//! owns all nodes and provides traversal, printing, and validation.
21
22use crate::{Op, Shape};
23
24use crate::provenance::NodeOrigin;
25
26/// Stable identifier for a node in the graph. Indices are never reused.
27#[cfg_attr(feature = "serialize", derive(serde::Serialize, serde::Deserialize))]
28#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
29pub struct NodeId(pub u32);
30
31impl std::fmt::Display for NodeId {
32    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
33        write!(f, "%{}", self.0)
34    }
35}
36
37/// A single node in the computation graph.
38#[cfg_attr(feature = "serialize", derive(serde::Serialize, serde::Deserialize))]
39#[derive(Debug, Clone)]
40pub struct Node {
41    pub id: NodeId,
42    /// The operation this node performs.
43    pub op: Op,
44    /// Input node IDs (operands). Order matches `Op::num_inputs()`.
45    pub inputs: Vec<NodeId>,
46    /// Output tensor shape (computed at construction time).
47    pub shape: Shape,
48    /// Human-readable name for debugging.
49    pub name: Option<String>,
50    /// Cross-stage provenance (HIR block, fusion pass, …).
51    pub origin: Option<NodeOrigin>,
52}
53
54/// A computation graph — the core IR data structure.
55///
56/// # Example
57/// ```
58/// use rlx_ir::*;
59///
60/// let mut g = Graph::new("bert_layer");
61///
62/// // Inputs
63/// let x = g.input("hidden", Shape::new(&[4, 15, 384], DType::F32));
64/// let w = g.param("qkv_weight", Shape::new(&[384, 1152], DType::F32));
65/// let b = g.param("qkv_bias", Shape::new(&[1152], DType::F32));
66///
67/// // QKV projection: matmul + bias
68/// let mm = g.matmul(x, w, Shape::new(&[4, 15, 1152], DType::F32));
69/// let qkv = g.binary(op::BinaryOp::Add, mm, b, Shape::new(&[4, 15, 1152], DType::F32));
70///
71/// assert_eq!(g.len(), 5);
72/// println!("{g}");
73/// ```
74#[cfg_attr(feature = "serialize", derive(serde::Serialize, serde::Deserialize))]
75#[derive(Clone, Debug)]
76pub struct Graph {
77    pub name: String,
78    nodes: Vec<Node>,
79    /// Output node IDs (the graph's results).
80    pub outputs: Vec<NodeId>,
81}
82
83// Subgraph equality is structural: same name, same node count, same outputs.
84// Full deep equality would require comparing every node and is rarely useful;
85// this gives Op derives `PartialEq` cheap structural comparison.
86impl PartialEq for Graph {
87    fn eq(&self, other: &Self) -> bool {
88        self.name == other.name
89            && self.nodes.len() == other.nodes.len()
90            && self.outputs == other.outputs
91    }
92}
93
94impl Graph {
95    pub fn new(name: impl Into<String>) -> Self {
96        Self {
97            name: name.into(),
98            nodes: Vec::new(),
99            outputs: Vec::new(),
100        }
101    }
102
103    /// Number of nodes in the graph.
104    pub fn len(&self) -> usize {
105        self.nodes.len()
106    }
107    pub fn is_empty(&self) -> bool {
108        self.nodes.is_empty()
109    }
110
111    /// Get a node by ID.
112    pub fn node(&self, id: NodeId) -> &Node {
113        &self.nodes[id.0 as usize]
114    }
115
116    /// Iterate all nodes in topological order (insertion order = topo order).
117    pub fn nodes(&self) -> &[Node] {
118        &self.nodes
119    }
120
121    /// Get the shape of a node's output.
122    pub fn shape(&self, id: NodeId) -> &Shape {
123        &self.nodes[id.0 as usize].shape
124    }
125
126    /// Set the graph outputs.
127    pub fn set_outputs(&mut self, outputs: Vec<NodeId>) {
128        self.outputs = outputs;
129    }
130
131    /// Replace the input list of a node in place. Used by post-
132    /// construction passes (`quant_propagate`, `dce`, etc.) that
133    /// rewire consumers without inserting new nodes.
134    /// Caller is responsible for shape consistency — this does no
135    /// re-inference.
136    pub fn set_inputs(&mut self, id: NodeId, inputs: Vec<NodeId>) {
137        self.nodes[id.0 as usize].inputs = inputs;
138    }
139
140    pub fn node_mut(&mut self, id: NodeId) -> &mut Node {
141        &mut self.nodes[id.0 as usize]
142    }
143
144    pub fn nodes_mut(&mut self) -> &mut [Node] {
145        &mut self.nodes
146    }
147
148    // ── Node constructors ───────────────────────────────────────
149
150    /// Append a node to the graph. `pub(crate)` so per-op builder
151    /// files in `rlx_ir::ops::*` can call it (plan #53).
152    /// Append a node for backend graph slicing (e.g. TPU HLO segments).
153    pub fn append_node(
154        &mut self,
155        op: Op,
156        inputs: Vec<NodeId>,
157        shape: Shape,
158        name: Option<String>,
159    ) -> NodeId {
160        self.push(op, inputs, shape, name)
161    }
162
163    pub(crate) fn push(
164        &mut self,
165        op: Op,
166        inputs: Vec<NodeId>,
167        shape: Shape,
168        name: Option<String>,
169    ) -> NodeId {
170        self.push_ext(op, inputs, shape, name, None)
171    }
172
173    pub(crate) fn push_ext(
174        &mut self,
175        op: Op,
176        inputs: Vec<NodeId>,
177        shape: Shape,
178        name: Option<String>,
179        origin: Option<NodeOrigin>,
180    ) -> NodeId {
181        let id = NodeId(self.nodes.len() as u32);
182        self.nodes.push(Node {
183            id,
184            op,
185            inputs,
186            shape,
187            name,
188            origin,
189        });
190        id
191    }
192
193    // Per-op builders moved to `crate::ops::*` (plan #53).
194    // Adding new op families = drop a new file in `ops/`, no edits here.
195
196    // ── Analysis helpers ────────────────────────────────────────
197
198    /// Find all nodes that use a given node's output.
199    pub fn users(&self, id: NodeId) -> Vec<NodeId> {
200        self.nodes
201            .iter()
202            .filter(|n| n.inputs.contains(&id))
203            .map(|n| n.id)
204            .collect()
205    }
206
207    /// Count how many nodes use a given node's output.
208    pub fn use_count(&self, id: NodeId) -> usize {
209        self.nodes.iter().filter(|n| n.inputs.contains(&id)).count()
210    }
211
212    /// Find a node by the name on its [`Op::Input`] or [`Op::Param`].
213    ///
214    /// Input/Param leaves are the graph's roots and survive optimizer passes
215    /// (fusion, DCE), so a *name* is the stable way to recover a handle into a
216    /// **rewritten** graph whose `NodeId`s have been renumbered — e.g. after
217    /// [`fuse`](../../rlx_compile/fusion_pipeline/fn.fuse.html). Outputs are
218    /// already positionally stable in [`outputs`](Self::outputs). Returns the
219    /// first match (names are expected unique).
220    pub fn node_id_by_name(&self, name: &str) -> Option<NodeId> {
221        self.nodes.iter().find_map(|n| match &n.op {
222            Op::Input { name: nm } | Op::Param { name: nm } if nm == name => Some(n.id),
223            _ => None,
224        })
225    }
226
227    /// [`node_id_by_name`](Self::node_id_by_name) restricted to graph inputs.
228    pub fn input_id(&self, name: &str) -> Option<NodeId> {
229        self.nodes.iter().find_map(|n| match &n.op {
230            Op::Input { name: nm } if nm == name => Some(n.id),
231            _ => None,
232        })
233    }
234
235    /// [`node_id_by_name`](Self::node_id_by_name) restricted to parameters.
236    pub fn param_id(&self, name: &str) -> Option<NodeId> {
237        self.nodes.iter().find_map(|n| match &n.op {
238            Op::Param { name: nm } if nm == name => Some(n.id),
239            _ => None,
240        })
241    }
242
243    /// Topological order (already guaranteed by construction — just node indices).
244    pub fn topo_order(&self) -> impl Iterator<Item = NodeId> + '_ {
245        (0..self.nodes.len()).map(|i| NodeId(i as u32))
246    }
247
248    /// Reverse topological order (outputs first).
249    pub fn reverse_topo(&self) -> impl Iterator<Item = NodeId> + '_ {
250        (0..self.nodes.len()).rev().map(|i| NodeId(i as u32))
251    }
252
253    // ── HIR / MIR / LIR pipeline (higher-order DX) ─────────────────
254
255    /// Fusion-first model definition at HIR level.
256    ///
257    /// Returns a [`GraphModule`] at HIR stage; call [`GraphModule::lower`]
258    /// or pass to [`rlx_opt::CompilePipeline::compile_module`].
259    pub fn define(
260        name: impl Into<String>,
261        build: impl FnOnce(&mut crate::hir::HirModule) -> crate::hir::HirNodeId,
262    ) -> crate::GraphModule {
263        crate::GraphModule::define(name, build)
264    }
265
266    /// Start an empty HIR-stage [`GraphModule`].
267    pub fn hir(name: impl Into<String>) -> crate::GraphModule {
268        crate::GraphModule::hir(name)
269    }
270
271    /// Wrap this MIR graph in a [`GraphModule`] for pipeline operations.
272    pub fn module(self) -> crate::GraphModule {
273        crate::GraphModule::from_graph(self)
274    }
275
276    /// Lower a HIR module to a MIR graph.
277    pub fn from_hir(hir: crate::hir::HirModule) -> Result<Self, crate::hir::LowerError> {
278        hir.lower_to_mir().map(|m| m.into_graph())
279    }
280
281    /// View as [`MirModule`].
282    pub fn to_mir(self) -> crate::MirModule {
283        crate::MirModule::from_graph(self)
284    }
285
286    /// Extract the MIR graph from optimized LIR.
287    pub fn from_lir(lir: crate::LirModule) -> Self {
288        lir.into_graph()
289    }
290
291    /// Annotated text dump ([`inspect_graph`]).
292    pub fn inspect(&self) -> String {
293        crate::inspect_graph(self)
294    }
295
296    /// True if any node shape uses a [`Dim::Dynamic`] symbol.
297    pub fn has_dynamic_dims(&self) -> bool {
298        crate::dynamic::has_dynamic_dims(self)
299    }
300
301    /// All dynamic symbols referenced in this graph.
302    pub fn dynamic_symbols(&self) -> Vec<u32> {
303        crate::dynamic::collect_dynamic_symbols(self)
304    }
305
306    /// Specialize symbolic dims to concrete sizes.
307    pub fn bind(&self, bindings: &crate::DimBinding) -> Self {
308        crate::dynamic::bind_graph(self, bindings)
309    }
310
311    /// Stage-aware dump when wrapped in [`GraphModule`].
312    pub fn inspect_module(module: &crate::GraphModule) -> String {
313        module.inspect()
314    }
315}
316
317/// Pretty-print the graph in a readable IR format.
318impl std::fmt::Display for Graph {
319    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
320        writeln!(f, "graph @{} {{", self.name)?;
321        for node in &self.nodes {
322            write!(f, "  {} = {}", node.id, node.op)?;
323            if !node.inputs.is_empty() {
324                write!(f, "(")?;
325                for (i, inp) in node.inputs.iter().enumerate() {
326                    if i > 0 {
327                        write!(f, ", ")?;
328                    }
329                    write!(f, "{inp}")?;
330                }
331                write!(f, ")")?;
332            }
333            writeln!(f, " : {}", node.shape)?;
334        }
335        if !self.outputs.is_empty() {
336            write!(f, "  return ")?;
337            for (i, o) in self.outputs.iter().enumerate() {
338                if i > 0 {
339                    write!(f, ", ")?;
340                }
341                write!(f, "{o}")?;
342            }
343            writeln!(f)?;
344        }
345        writeln!(f, "}}")
346    }
347}
348
349#[cfg(test)]
350mod tests {
351    use super::*;
352    use crate::{
353        DType,
354        op::{Activation, BinaryOp},
355    };
356
357    #[test]
358    fn build_simple_graph() {
359        let mut g = Graph::new("test");
360
361        let x = g.input("x", Shape::new(&[4, 15, 384], DType::F32));
362        let w = g.param("weight", Shape::new(&[384, 1536], DType::F32));
363        let b = g.param("bias", Shape::new(&[1536], DType::F32));
364
365        let mm = g.matmul(x, w, Shape::new(&[4, 15, 1536], DType::F32));
366        let add = g.binary(BinaryOp::Add, mm, b, Shape::new(&[4, 15, 1536], DType::F32));
367        let out = g.activation(
368            Activation::Gelu,
369            add,
370            Shape::new(&[4, 15, 1536], DType::F32),
371        );
372
373        g.set_outputs(vec![out]);
374
375        assert_eq!(g.len(), 6);
376        assert_eq!(g.use_count(mm), 1); // matmul used by add
377        assert_eq!(g.use_count(x), 1); // x used by matmul
378
379        let printed = format!("{g}");
380        assert!(printed.contains("matmul(%0, %1)"));
381        assert!(printed.contains("Gelu(%4)"));
382        assert!(printed.contains("return %5"));
383    }
384
385    /// Build a BERT layer to verify the IR can represent real models.
386    #[test]
387    fn bert_layer_graph() {
388        let mut g = Graph::new("bert_layer");
389        let f = DType::F32;
390        let h = 384;
391        let int = 1536;
392
393        // Input
394        let x = g.input("hidden", Shape::new(&[4, 15, h], f));
395
396        // QKV
397        let qkv_w = g.param("qkv.weight", Shape::new(&[h, 3 * h], f));
398        let qkv_b = g.param("qkv.bias", Shape::new(&[3 * h], f));
399        let qkv = g.matmul(x, qkv_w, Shape::new(&[4, 15, 3 * h], f));
400        let _qkv = g.binary(BinaryOp::Add, qkv, qkv_b, Shape::new(&[4, 15, 3 * h], f));
401
402        // (would split Q/K/V, attention, out_proj here — simplified)
403
404        // FFN
405        let int_w = g.param("ffn.weight", Shape::new(&[h, int], f));
406        let int_b = g.param("ffn.bias", Shape::new(&[int], f));
407        let ffn = g.matmul(x, int_w, Shape::new(&[4, 15, int], f));
408        let ffn = g.binary(BinaryOp::Add, ffn, int_b, Shape::new(&[4, 15, int], f));
409        let ffn = g.activation(Activation::Gelu, ffn, Shape::new(&[4, 15, int], f));
410
411        let out_w = g.param("ffn_out.weight", Shape::new(&[int, h], f));
412        let ffn_out = g.matmul(ffn, out_w, Shape::new(&[4, 15, h], f));
413
414        g.set_outputs(vec![ffn_out]);
415
416        assert!(g.len() > 10);
417        println!("{g}");
418    }
419}
rlx_ir/graph.rs

rlx_ir/
graph.rs
