oxfuse_core/
graph.rs

1//! The IR: an arena of values in definition order. Every op has one output.
2
3use crate::{Op, TensorType, ValueId};
4use std::sync::Arc;
5
6#[derive(Debug, Clone)]
7pub enum Source {
8    Input,
9    /// Shared so rebuilding passes and execution plans never copy weights.
10    Constant(Arc<[f32]>),
11    Node {
12        op: Op,
13        inputs: Vec<ValueId>,
14    },
15}
16
17impl PartialEq for Source {
18    fn eq(&self, other: &Self) -> bool {
19        match (self, other) {
20            (Self::Input, Self::Input) => true,
21            (Self::Constant(a), Self::Constant(b)) => {
22                Arc::ptr_eq(a, b)
23                    || (a.len() == b.len()
24                        && a.iter()
25                            .zip(b.iter())
26                            .all(|(x, y)| x.to_bits() == y.to_bits()))
27            }
28            (
29                Self::Node {
30                    op: op_a,
31                    inputs: in_a,
32                },
33                Self::Node {
34                    op: op_b,
35                    inputs: in_b,
36                },
37            ) => op_a == op_b && in_a == in_b,
38            _ => false,
39        }
40    }
41}
42
43impl Eq for Source {}
44
45#[derive(Debug, Clone, PartialEq, Eq)]
46pub struct Value {
47    pub ty: TensorType,
48    pub source: Source,
49    pub name: Option<String>,
50}
51
52/// A graph is only constructed through `GraphBuilder::finish`, which verifies it.
53#[derive(Debug, Clone, PartialEq, Eq)]
54pub struct Graph {
55    pub(crate) values: Vec<Value>,
56    pub(crate) inputs: Vec<ValueId>,
57    pub(crate) outputs: Vec<ValueId>,
58    /// Parallel to `outputs`. The names users see; passes never change them.
59    pub(crate) output_names: Vec<String>,
60}
61
62impl Graph {
63    pub fn values(&self) -> &[Value] {
64        &self.values
65    }
66
67    pub fn value(&self, id: ValueId) -> &Value {
68        &self.values[id.index()]
69    }
70
71    /// Graph inputs in declared order.
72    pub fn inputs(&self) -> &[ValueId] {
73        &self.inputs
74    }
75
76    pub fn outputs(&self) -> &[ValueId] {
77        &self.outputs
78    }
79
80    pub fn output_names(&self) -> &[String] {
81        &self.output_names
82    }
83
84    /// The value's name, or `%<id>` when it has none.
85    pub fn label(&self, id: ValueId) -> String {
86        match &self.value(id).name {
87            Some(name) => name.clone(),
88            None => format!("%{}", id.0),
89        }
90    }
91}
92
93#[cfg(test)]
94mod tests {
95    use crate::{GraphBuilder, Limits, Op};
96
97    fn graph_with_constant(value: f32) -> crate::Graph {
98        let mut b = GraphBuilder::new(Limits::default());
99        let c = b.constant("c", vec![1], vec![value]).unwrap();
100        b.finish(&[c]).unwrap()
101    }
102
103    #[test]
104    fn a_graph_with_a_nan_constant_equals_its_clone() {
105        let g = graph_with_constant(f32::NAN);
106        assert_eq!(g, g.clone());
107    }
108
109    #[test]
110    fn constants_compare_by_bits_so_signed_zeros_differ() {
111        assert_ne!(graph_with_constant(0.0), graph_with_constant(-0.0));
112        assert_eq!(graph_with_constant(-0.0), graph_with_constant(-0.0));
113    }
114
115    #[test]
116    fn nodes_with_the_same_op_but_different_inputs_are_not_equal() {
117        let build = |swap: bool| {
118            let mut b = GraphBuilder::new(Limits::default());
119            let x = b.input("x", vec![2]).unwrap();
120            let z = b.input("z", vec![2]).unwrap();
121            let inputs = if swap { [z, x] } else { [x, z] };
122            let y = b.node("y", Op::Add, &inputs).unwrap();
123            b.finish(&[y]).unwrap()
124        };
125        assert_eq!(build(false), build(false));
126        assert_ne!(build(false), build(true));
127    }
128
129    #[test]
130    fn nodes_with_the_same_inputs_but_different_ops_are_not_equal() {
131        let build = |op: Op| {
132            let mut b = GraphBuilder::new(Limits::default());
133            let x = b.input("x", vec![2]).unwrap();
134            let y = b.node("y", op, &[x]).unwrap();
135            b.finish(&[y]).unwrap()
136        };
137        assert_ne!(build(Op::Relu), build(Op::Sigmoid));
138    }
139}
oxfuse_core/graph.rs

oxfuse_core/
graph.rs
