numrs/ir/
graph.rs

1use crate::ir::hlo::{HloNode, HloOp};
2
3/// Very small IR graph container. Holds nodes and entry points.
4#[derive(Clone, Debug)]
5pub struct IRGraph {
6    pub nodes: Vec<HloNode>,
7}
8
9impl IRGraph {
10    pub fn new() -> Self { Self { nodes: vec![] } }
11
12    /// Append a node and return its id
13    pub fn push(&mut self, node: HloNode) -> usize {
14        let id = node.id;
15        self.nodes.push(node);
16        id
17    }
18
19    /// Simple convenience: create a 1-node graph for an op between two shapes
20    pub fn binary_op(op: HloOp, left_shape: Vec<usize>, right_shape: Vec<usize>) -> Self {
21        // naive broadcast: pick the bigger shape
22        let shape = if left_shape.len() >= right_shape.len() { left_shape.clone() } else { right_shape.clone() };
23        let mut g = IRGraph::new();
24        // constants as id 0,1 then op id 2
25        g.nodes.push(HloNode::new(0, op.clone(), vec![], left_shape));
26        g.nodes.push(HloNode::new(1, op.clone(), vec![], right_shape));
27        g.nodes.push(HloNode::new(2, op, vec![0,1], shape));
28        g
29    }
30}
numrs/ir/graph.rs

numrs/ir/
graph.rs
