oxfuse_core/passes/
cse.rs

1//! Common-subexpression elimination: a node with the same op and the same
2//! (remapped) inputs as an earlier node is aliased to it. Keys use the new
3//! ids, so merging is transitive. Constants and inputs are never merged, and
4//! `Add(a, b)` is not `Add(b, a)`; both are valid extensions, left out until
5//! a measured case needs them.
6
7use super::rebuild::Rebuilder;
8use super::{Pass, PassOutput};
9use crate::{Error, Graph, Limits, Op, Source, ValueId};
10use std::collections::HashMap;
11
12pub struct CommonSubexpressionElimination {
13    limits: Limits,
14}
15
16impl CommonSubexpressionElimination {
17    pub fn new(limits: Limits) -> Self {
18        Self { limits }
19    }
20}
21
22impl Pass for CommonSubexpressionElimination {
23    fn name(&self) -> &'static str {
24        "cse"
25    }
26
27    fn run(&self, graph: &Graph) -> Result<PassOutput, Error> {
28        let mut rebuild = Rebuilder::new(graph, &self.limits);
29        let mut seen: HashMap<(Op, Vec<ValueId>), ValueId> = HashMap::new();
30        let mut merged = 0;
31        for (index, value) in graph.values().iter().enumerate() {
32            let old = ValueId(index as u32);
33            let Source::Node { op, inputs } = &value.source else {
34                rebuild.copy(old)?;
35                continue;
36            };
37            let key = (
38                op.clone(),
39                inputs
40                    .iter()
41                    .map(|&id| rebuild.mapped(id))
42                    .collect::<Vec<_>>(),
43            );
44            if let Some(&existing) = seen.get(&key) {
45                rebuild.alias(old, existing);
46                merged += 1;
47            } else {
48                let new = rebuild.copy(old)?;
49                seen.insert(key, new);
50            }
51        }
52        Ok(PassOutput {
53            graph: rebuild.finish()?,
54            changed: merged,
55        })
56    }
57}
58
59#[cfg(test)]
60mod tests {
61    use super::*;
62    use crate::{GraphBuilder, Op};
63
64    fn builder() -> GraphBuilder {
65        GraphBuilder::new(Limits::default())
66    }
67
68    fn run(graph: &Graph) -> PassOutput {
69        CommonSubexpressionElimination::new(Limits::default())
70            .run(graph)
71            .unwrap()
72    }
73
74    #[test]
75    fn merges_identical_nodes() {
76        let mut b = builder();
77        let x = b.input("x", vec![2]).unwrap();
78        let p = b.node("p", Op::Relu, &[x]).unwrap();
79        let q = b.node("q", Op::Relu, &[x]).unwrap();
80        let y = b.node("y", Op::Add, &[p, q]).unwrap();
81        let g = b.finish(&[y]).unwrap();
82        let out = run(&g);
83        assert_eq!(out.changed, 1);
84        assert_eq!(out.graph.values().len(), 3);
85        let add_inputs = match &out.graph.value(out.graph.outputs()[0]).source {
86            crate::Source::Node { inputs, .. } => inputs.clone(),
87            other => panic!("{other:?}"),
88        };
89        assert_eq!(add_inputs[0], add_inputs[1]);
90    }
91
92    #[test]
93    fn merged_outputs_keep_their_names() {
94        let mut b = builder();
95        let x = b.input("x", vec![2]).unwrap();
96        let p = b.node("p", Op::Relu, &[x]).unwrap();
97        let q = b.node("q", Op::Relu, &[x]).unwrap();
98        let g = b.finish(&[p, q]).unwrap();
99        let out = run(&g);
100        assert_eq!(out.changed, 1);
101        assert_eq!(out.graph.output_names(), ["p", "q"]);
102        assert_eq!(out.graph.outputs()[0], out.graph.outputs()[1]);
103    }
104
105    #[test]
106    fn merging_is_transitive() {
107        let mut b = builder();
108        let x = b.input("x", vec![2]).unwrap();
109        let c = b.constant("c", vec![2], vec![1.0, 2.0]).unwrap();
110        let a1 = b.node("a1", Op::Relu, &[x]).unwrap();
111        let a2 = b.node("a2", Op::Relu, &[x]).unwrap();
112        let b1 = b.node("b1", Op::Add, &[a1, c]).unwrap();
113        let b2 = b.node("b2", Op::Add, &[a2, c]).unwrap();
114        let g = b.finish(&[b1, b2]).unwrap();
115        let out = run(&g);
116        assert_eq!(out.changed, 2);
117        assert_eq!(out.graph.values().len(), 4);
118    }
119
120    #[test]
121    fn different_attributes_are_not_merged() {
122        let mut b = builder();
123        let x = b.input("x", vec![6]).unwrap();
124        let r1 = b
125            .node("r1", Op::Reshape { shape: vec![2, 3] }, &[x])
126            .unwrap();
127        let r2 = b
128            .node("r2", Op::Reshape { shape: vec![3, 2] }, &[x])
129            .unwrap();
130        let r3 = b
131            .node("r3", Op::Reshape { shape: vec![2, 3] }, &[x])
132            .unwrap();
133        let g = b.finish(&[r1, r2, r3]).unwrap();
134        let out = run(&g);
135        assert_eq!(out.changed, 1);
136        assert_eq!(out.graph.outputs()[0], out.graph.outputs()[2]);
137        assert_ne!(out.graph.outputs()[0], out.graph.outputs()[1]);
138    }
139
140    #[test]
141    fn operand_order_matters() {
142        let mut b = builder();
143        let x = b.input("x", vec![2]).unwrap();
144        let z = b.input("z", vec![2]).unwrap();
145        let s1 = b.node("s1", Op::Add, &[x, z]).unwrap();
146        let s2 = b.node("s2", Op::Add, &[z, x]).unwrap();
147        let g = b.finish(&[s1, s2]).unwrap();
148        let out = run(&g);
149        assert_eq!(out.changed, 0);
150        assert_eq!(out.graph, g);
151    }
152
153    #[test]
154    fn constants_and_inputs_are_never_merged() {
155        let mut b = builder();
156        let x1 = b.input("x1", vec![2]).unwrap();
157        let x2 = b.input("x2", vec![2]).unwrap();
158        let c1 = b.constant("c1", vec![2], vec![1.0, 2.0]).unwrap();
159        let c2 = b.constant("c2", vec![2], vec![1.0, 2.0]).unwrap();
160        let g = b.finish(&[x1, x2, c1, c2]).unwrap();
161        let out = run(&g);
162        assert_eq!(out.changed, 0);
163        assert_eq!(out.graph, g);
164    }
165
166    #[test]
167    fn nothing_to_merge_rebuilds_the_same_graph() {
168        let mut b = builder();
169        let x = b.input("x", vec![2]).unwrap();
170        let p = b.node("p", Op::Relu, &[x]).unwrap();
171        let q = b.node("q", Op::Sigmoid, &[p]).unwrap();
172        let g = b.finish(&[q]).unwrap();
173        let out = run(&g);
174        assert_eq!(out.changed, 0);
175        assert_eq!(out.graph, g);
176    }
177}
oxfuse_core/passes/cse.rs

oxfuse_core/passes/
cse.rs
