nxpu_ir/
graph.rs

1//! Graph-level intermediate representation for multi-operation models.
2//!
3//! Extends the single-kernel IR to support directed acyclic graphs (DAGs)
4//! of operations, enabling transpilation of production ML models with
5//! 50-500+ operations.
6
7use std::collections::{BTreeSet, HashMap};
8
9use crate::IrError;
10use crate::types::{Scalar, TensorShape};
11
12/// A unique identifier for a node in the computation graph.
13#[derive(Clone, Copy, Debug, Hash, Eq, PartialEq, Ord, PartialOrd)]
14pub struct NodeId(pub u32);
15
16/// A unique identifier for an edge (tensor) in the computation graph.
17#[derive(Clone, Copy, Debug, Hash, Eq, PartialEq)]
18pub struct EdgeId(pub u32);
19
20/// Activation function that can be fused into a preceding operation.
21#[derive(Clone, Copy, Debug, PartialEq, Eq)]
22pub enum ActivationFunction {
23    /// No activation (identity).
24    None,
25    /// Rectified Linear Unit.
26    Relu,
27    /// Sigmoid activation.
28    Sigmoid,
29    /// Hyperbolic tangent.
30    Tanh,
31}
32
33/// The operation type for a graph node.
34#[derive(Clone, Debug, PartialEq, Eq)]
35pub enum GraphOp {
36    /// Matrix multiplication.
37    MatMul,
38    /// Convolution 2D.
39    Conv2d,
40    /// Element-wise addition.
41    Add,
42    /// Element-wise subtraction.
43    Sub,
44    /// Element-wise multiplication.
45    Mul,
46    /// Element-wise division.
47    Div,
48    /// Rectified Linear Unit activation.
49    Relu,
50    /// Sigmoid activation.
51    Sigmoid,
52    /// Softmax activation.
53    Softmax,
54    /// Batch normalization.
55    BatchNorm,
56    /// Layer normalization.
57    LayerNorm,
58    /// Max pooling 2D.
59    MaxPool2d,
60    /// Average pooling 2D.
61    AvgPool2d,
62    /// Reshape/view.
63    Reshape,
64    /// Transpose/permute dimensions.
65    Transpose,
66    /// Concatenation along an axis.
67    Concat { axis: i32 },
68    /// Custom/vendor-specific operation.
69    Custom { op_type: String },
70
71    // --- Fused operations ---
72    /// Fused Conv2D with optional bias addition and activation.
73    /// Inputs: [input, weight, bias?]
74    FusedConv2d { activation: ActivationFunction },
75    /// General matrix multiply: alpha * A @ B + beta * C (ONNX Gemm).
76    /// Inputs: [A, B, C]
77    Gemm {
78        /// Scaling factor for the matrix product.
79        alpha: i32,
80        /// Scaling factor for the bias term.
81        beta: i32,
82    },
83    /// Element-wise binary operation with fused activation.
84    /// Inputs: [lhs, rhs]
85    FusedElementWise {
86        /// The base element-wise operation (Add, Sub, Mul, Div).
87        base_op: Box<GraphOp>,
88        activation: ActivationFunction,
89    },
90}
91
92impl GraphOp {
93    /// Returns the ONNX operator type string.
94    pub fn onnx_op_type(&self) -> &str {
95        match self {
96            Self::MatMul => "MatMul",
97            Self::Conv2d => "Conv",
98            Self::Add => "Add",
99            Self::Sub => "Sub",
100            Self::Mul => "Mul",
101            Self::Div => "Div",
102            Self::Relu => "Relu",
103            Self::Sigmoid => "Sigmoid",
104            Self::Softmax => "Softmax",
105            Self::BatchNorm => "BatchNormalization",
106            Self::LayerNorm => "LayerNormalization",
107            Self::MaxPool2d => "MaxPool",
108            Self::AvgPool2d => "AveragePool",
109            Self::Reshape => "Reshape",
110            Self::Transpose => "Transpose",
111            Self::Concat { .. } => "Concat",
112            Self::Custom { op_type } => op_type,
113            Self::FusedConv2d { .. } => "Conv",
114            Self::Gemm { .. } => "Gemm",
115            Self::FusedElementWise { base_op, .. } => base_op.onnx_op_type(),
116        }
117    }
118}
119
120/// Metadata about a tensor edge in the graph.
121#[derive(Clone, Debug)]
122pub struct TensorInfo {
123    /// Human-readable name.
124    pub name: String,
125    /// Element scalar type.
126    pub scalar: Scalar,
127    /// Shape (may contain dynamic dimensions).
128    pub shape: TensorShape,
129}
130
131/// A node in the computation graph.
132#[derive(Clone, Debug)]
133pub struct GraphNode {
134    /// Unique identifier for this node.
135    pub id: NodeId,
136    /// The operation this node performs.
137    pub op: GraphOp,
138    /// Input edge identifiers (ordered).
139    pub inputs: Vec<EdgeId>,
140    /// Output edge identifiers (ordered).
141    pub outputs: Vec<EdgeId>,
142    /// Human-readable name for this node.
143    pub name: String,
144}
145
146/// A computation graph representing a multi-operation model.
147///
148/// This is a DAG where nodes are operations and edges are tensors
149/// flowing between operations.
150#[derive(Clone, Debug, Default)]
151pub struct ComputeGraph {
152    /// All nodes in the graph, keyed by NodeId.
153    pub nodes: Vec<GraphNode>,
154    /// All tensor edges, keyed by EdgeId.
155    pub edges: HashMap<EdgeId, TensorInfo>,
156    /// Graph-level input edge ids (model inputs).
157    pub inputs: Vec<EdgeId>,
158    /// Graph-level output edge ids (model outputs).
159    pub outputs: Vec<EdgeId>,
160    /// Next available node id.
161    next_node_id: u32,
162    /// Next available edge id.
163    next_edge_id: u32,
164}
165
166impl ComputeGraph {
167    /// Create a new empty graph.
168    pub fn new() -> Self {
169        Self::default()
170    }
171
172    /// Add a tensor edge to the graph and return its id.
173    pub fn add_edge(&mut self, info: TensorInfo) -> EdgeId {
174        let id = EdgeId(self.next_edge_id);
175        self.next_edge_id += 1;
176        self.edges.insert(id, info);
177        id
178    }
179
180    /// Add a node to the graph and return its id.
181    ///
182    /// # Errors
183    ///
184    /// Returns [`IrError::UnknownEdge`] if any input or output [`EdgeId`] has
185    /// not been previously registered via [`add_edge`](Self::add_edge), or
186    /// [`IrError::DuplicateEdgeProducer`] if an output edge already has a
187    /// producer node (each edge may have at most one producer).
188    pub fn add_node(
189        &mut self,
190        op: GraphOp,
191        inputs: Vec<EdgeId>,
192        outputs: Vec<EdgeId>,
193        name: impl Into<String>,
194    ) -> Result<NodeId, IrError> {
195        let name = name.into();
196
197        // Validate that all referenced edges exist.
198        for &e in inputs.iter().chain(outputs.iter()) {
199            if !self.edges.contains_key(&e) {
200                return Err(IrError::UnknownEdge {
201                    edge_id: e.0,
202                    node_name: name,
203                });
204            }
205        }
206
207        // Enforce single-producer-per-edge invariant.
208        for &out in &outputs {
209            if let Some(existing) = self.nodes.iter().find(|n| n.outputs.contains(&out)) {
210                return Err(IrError::DuplicateEdgeProducer {
211                    edge_id: out.0,
212                    existing_producer: existing.name.clone(),
213                    new_producer: name,
214                });
215            }
216        }
217
218        let id = NodeId(self.next_node_id);
219        self.next_node_id += 1;
220        self.nodes.push(GraphNode {
221            id,
222            op,
223            inputs,
224            outputs,
225            name,
226        });
227        Ok(id)
228    }
229
230    /// Number of nodes in the graph.
231    pub fn node_count(&self) -> usize {
232        self.nodes.len()
233    }
234
235    /// Number of edges (tensors) in the graph.
236    pub fn edge_count(&self) -> usize {
237        self.edges.len()
238    }
239
240    /// Returns nodes in topological order.
241    ///
242    /// The ordering is deterministic: among nodes with the same in-degree,
243    /// the one with the smaller [`NodeId`] is emitted first.
244    ///
245    /// # Errors
246    ///
247    /// Returns [`IrError::CycleDetected`] if the graph contains a cycle.
248    pub fn topological_order(&self) -> Result<Vec<&GraphNode>, IrError> {
249        // Build adjacency: which node produces each edge?
250        let mut edge_producer: HashMap<EdgeId, usize> = HashMap::new();
251        for (i, node) in self.nodes.iter().enumerate() {
252            for &out in &node.outputs {
253                edge_producer.insert(out, i);
254            }
255        }
256
257        // Build per-node consumer lists and in-degree (O(V+E))
258        let n = self.nodes.len();
259        let mut in_degree = vec![0usize; n];
260        let mut consumers: Vec<Vec<usize>> = vec![Vec::new(); n];
261
262        for (ci, node) in self.nodes.iter().enumerate() {
263            for &inp in &node.inputs {
264                if let Some(&pi) = edge_producer.get(&inp) {
265                    in_degree[ci] += 1;
266                    consumers[pi].push(ci);
267                }
268            }
269        }
270
271        // Kahn's algorithm with deterministic BTreeSet (ordered by NodeId)
272        let mut ready: BTreeSet<(NodeId, usize)> = BTreeSet::new();
273        for (i, &deg) in in_degree.iter().enumerate() {
274            if deg == 0 {
275                ready.insert((self.nodes[i].id, i));
276            }
277        }
278
279        let mut result: Vec<&GraphNode> = Vec::with_capacity(n);
280
281        while let Some(&(_, idx)) = ready.iter().next() {
282            ready.remove(&(self.nodes[idx].id, idx));
283            result.push(&self.nodes[idx]);
284
285            for &ci in &consumers[idx] {
286                in_degree[ci] -= 1;
287                if in_degree[ci] == 0 {
288                    ready.insert((self.nodes[ci].id, ci));
289                }
290            }
291        }
292
293        if result.len() != n {
294            return Err(IrError::CycleDetected {
295                visited: result.len(),
296                total: n,
297            });
298        }
299
300        Ok(result)
301    }
302
303    /// Find all nodes that consume the given edge.
304    pub fn edge_consumers(&self, edge: EdgeId) -> Vec<&GraphNode> {
305        self.nodes
306            .iter()
307            .filter(|n| n.inputs.contains(&edge))
308            .collect()
309    }
310
311    /// Find the node that produces the given edge, if any.
312    pub fn edge_producer(&self, edge: EdgeId) -> Option<&GraphNode> {
313        self.nodes.iter().find(|n| n.outputs.contains(&edge))
314    }
315}
316
317/// Extends [`crate::Module`] with an optional computation graph.
318///
319/// When present, the graph describes how multiple entry points
320/// compose into a single model.
321#[derive(Clone, Debug, Default)]
322pub struct GraphModule {
323    /// The base IR module with types, globals, and entry points.
324    pub module: crate::Module,
325    /// Optional multi-operation graph overlay.
326    pub graph: Option<ComputeGraph>,
327}
328
329#[cfg(test)]
330mod tests {
331    use super::*;
332    use crate::types::{Dimension, TensorShape};
333
334    fn make_tensor_info(name: &str, shape: &[i64]) -> TensorInfo {
335        TensorInfo {
336            name: name.into(),
337            scalar: Scalar::F32,
338            shape: TensorShape {
339                dims: shape
340                    .iter()
341                    .map(|&d| {
342                        if d < 0 {
343                            Dimension::Dynamic(None)
344                        } else {
345                            Dimension::Fixed(d as u32)
346                        }
347                    })
348                    .collect(),
349            },
350        }
351    }
352
353    #[test]
354    fn build_simple_graph() {
355        let mut graph = ComputeGraph::new();
356
357        // MatMul → Add → ReLU
358        let a = graph.add_edge(make_tensor_info("A", &[-1, 768]));
359        let b = graph.add_edge(make_tensor_info("B", &[768, 768]));
360        let matmul_out = graph.add_edge(make_tensor_info("matmul_out", &[-1, 768]));
361        let bias = graph.add_edge(make_tensor_info("bias", &[768]));
362        let add_out = graph.add_edge(make_tensor_info("add_out", &[-1, 768]));
363        let relu_out = graph.add_edge(make_tensor_info("relu_out", &[-1, 768]));
364
365        graph.inputs = vec![a, b, bias];
366        graph.outputs = vec![relu_out];
367
368        graph
369            .add_node(GraphOp::MatMul, vec![a, b], vec![matmul_out], "matmul_0")
370            .unwrap();
371        graph
372            .add_node(GraphOp::Add, vec![matmul_out, bias], vec![add_out], "add_0")
373            .unwrap();
374        graph
375            .add_node(GraphOp::Relu, vec![add_out], vec![relu_out], "relu_0")
376            .unwrap();
377
378        assert_eq!(graph.node_count(), 3);
379        assert_eq!(graph.edge_count(), 6);
380    }
381
382    #[test]
383    fn topological_order() {
384        let mut graph = ComputeGraph::new();
385
386        let a = graph.add_edge(make_tensor_info("A", &[-1, 10]));
387        let b = graph.add_edge(make_tensor_info("B", &[10, 10]));
388        let c = graph.add_edge(make_tensor_info("C", &[-1, 10]));
389        let d = graph.add_edge(make_tensor_info("D", &[-1, 10]));
390
391        graph.inputs = vec![a, b];
392        graph.outputs = vec![d];
393
394        graph
395            .add_node(GraphOp::MatMul, vec![a, b], vec![c], "matmul")
396            .unwrap();
397        graph
398            .add_node(GraphOp::Relu, vec![c], vec![d], "relu")
399            .unwrap();
400
401        let order = graph.topological_order().unwrap();
402        assert_eq!(order.len(), 2);
403        assert_eq!(order[0].name, "matmul");
404        assert_eq!(order[1].name, "relu");
405    }
406
407    #[test]
408    fn edge_producer_consumer() {
409        let mut graph = ComputeGraph::new();
410
411        let a = graph.add_edge(make_tensor_info("A", &[10]));
412        let b = graph.add_edge(make_tensor_info("B", &[10]));
413        let c = graph.add_edge(make_tensor_info("C", &[10]));
414
415        graph
416            .add_node(GraphOp::Add, vec![a], vec![b], "add")
417            .unwrap();
418        graph
419            .add_node(GraphOp::Relu, vec![b], vec![c], "relu")
420            .unwrap();
421
422        let producer = graph.edge_producer(b).unwrap();
423        assert_eq!(producer.name, "add");
424
425        let consumers = graph.edge_consumers(b);
426        assert_eq!(consumers.len(), 1);
427        assert_eq!(consumers[0].name, "relu");
428
429        // Graph input has no producer
430        assert!(graph.edge_producer(a).is_none());
431    }
432
433    #[test]
434    fn graph_op_onnx_names() {
435        assert_eq!(GraphOp::MatMul.onnx_op_type(), "MatMul");
436        assert_eq!(GraphOp::Conv2d.onnx_op_type(), "Conv");
437        assert_eq!(GraphOp::Relu.onnx_op_type(), "Relu");
438        assert_eq!(GraphOp::Softmax.onnx_op_type(), "Softmax");
439        assert_eq!(
440            GraphOp::Custom {
441                op_type: "MyOp".into()
442            }
443            .onnx_op_type(),
444            "MyOp"
445        );
446    }
447
448    #[test]
449    fn graph_module() {
450        let gm = GraphModule::default();
451        assert!(gm.graph.is_none());
452        assert_eq!(gm.module.entry_points.len(), 0);
453    }
454
455    #[test]
456    fn topological_order_empty_graph() {
457        let graph = ComputeGraph::new();
458        let order = graph.topological_order().unwrap();
459        assert_eq!(order.len(), 0);
460    }
461
462    #[test]
463    fn topological_order_diamond_dag() {
464        // A → B, A → C, B → D, C → D
465        let mut graph = ComputeGraph::new();
466
467        let e_in = graph.add_edge(make_tensor_info("in", &[10]));
468        let e_ab = graph.add_edge(make_tensor_info("ab", &[10]));
469        let e_ac = graph.add_edge(make_tensor_info("ac", &[10]));
470        let e_bd = graph.add_edge(make_tensor_info("bd", &[10]));
471        let e_cd = graph.add_edge(make_tensor_info("cd", &[10]));
472        let e_out = graph.add_edge(make_tensor_info("out", &[10]));
473
474        graph
475            .add_node(GraphOp::Relu, vec![e_in], vec![e_ab, e_ac], "A")
476            .unwrap();
477        graph
478            .add_node(GraphOp::Relu, vec![e_ab], vec![e_bd], "B")
479            .unwrap();
480        graph
481            .add_node(GraphOp::Relu, vec![e_ac], vec![e_cd], "C")
482            .unwrap();
483        graph
484            .add_node(GraphOp::Add, vec![e_bd, e_cd], vec![e_out], "D")
485            .unwrap();
486
487        let order = graph.topological_order().unwrap();
488        assert_eq!(order.len(), 4);
489        assert_eq!(order[0].name, "A");
490        // B before C (deterministic by NodeId)
491        assert_eq!(order[1].name, "B");
492        assert_eq!(order[2].name, "C");
493        assert_eq!(order[3].name, "D");
494    }
495
496    #[test]
497    fn topological_order_detects_cycle() {
498        let mut graph = ComputeGraph::new();
499
500        let e0 = graph.add_edge(make_tensor_info("e0", &[10]));
501        let e1 = graph.add_edge(make_tensor_info("e1", &[10]));
502
503        // Manually build a cycle by pushing nodes directly
504        // (bypassing add_node validation which checks single-producer)
505        graph.nodes.push(GraphNode {
506            id: NodeId(0),
507            op: GraphOp::Relu,
508            inputs: vec![e1],
509            outputs: vec![e0],
510            name: "A".into(),
511        });
512        graph.nodes.push(GraphNode {
513            id: NodeId(1),
514            op: GraphOp::Relu,
515            inputs: vec![e0],
516            outputs: vec![e1],
517            name: "B".into(),
518        });
519        graph.next_node_id = 2;
520
521        let err = graph.topological_order().unwrap_err();
522        assert!(
523            matches!(err, IrError::CycleDetected { .. }),
524            "expected CycleDetected, got {err:?}"
525        );
526    }
527
528    #[test]
529    fn add_node_rejects_unknown_edge() {
530        let mut graph = ComputeGraph::new();
531        let fake_edge = EdgeId(999);
532        let out = graph.add_edge(make_tensor_info("out", &[10]));
533        let err = graph
534            .add_node(GraphOp::Relu, vec![fake_edge], vec![out], "bad")
535            .unwrap_err();
536        assert!(
537            matches!(err, IrError::UnknownEdge { .. }),
538            "expected UnknownEdge, got {err:?}"
539        );
540    }
541
542    #[test]
543    fn add_node_rejects_duplicate_producer() {
544        let mut graph = ComputeGraph::new();
545        let a = graph.add_edge(make_tensor_info("a", &[10]));
546        let b = graph.add_edge(make_tensor_info("b", &[10]));
547        let c = graph.add_edge(make_tensor_info("c", &[10]));
548
549        graph
550            .add_node(GraphOp::Relu, vec![a], vec![b], "first")
551            .unwrap();
552        let err = graph
553            .add_node(GraphOp::Relu, vec![c], vec![b], "second")
554            .unwrap_err();
555        assert!(
556            matches!(err, IrError::DuplicateEdgeProducer { .. }),
557            "expected DuplicateEdgeProducer, got {err:?}"
558        );
559    }
560}
nxpu_ir/graph.rs

nxpu_ir/
graph.rs
