oxfuse_core/passes/
dce.rs

1//! Dead-code elimination: keep what the outputs transitively use, plus every
2//! graph input. Inputs are never removed because they are the interface.
3
4use super::rebuild::Rebuilder;
5use super::{Pass, PassOutput};
6use crate::{Error, Graph, Limits, Source};
7
8pub struct DeadCodeElimination {
9    limits: Limits,
10}
11
12impl DeadCodeElimination {
13    pub fn new(limits: Limits) -> Self {
14        Self { limits }
15    }
16}
17
18impl Pass for DeadCodeElimination {
19    fn name(&self) -> &'static str {
20        "dce"
21    }
22
23    fn run(&self, graph: &Graph) -> Result<PassOutput, Error> {
24        let count = graph.values().len();
25        let mut live = vec![false; count];
26        for id in graph.outputs().iter().chain(graph.inputs()) {
27            live[id.index()] = true;
28        }
29        // Inputs of a node have smaller ids, so one backward sweep suffices.
30        for index in (0..count).rev() {
31            if !live[index] {
32                continue;
33            }
34            if let Source::Node { inputs, .. } = &graph.values()[index].source {
35                for input in inputs {
36                    live[input.index()] = true;
37                }
38            }
39        }
40        let mut rebuild = Rebuilder::new(graph, &self.limits);
41        let mut removed = 0;
42        for (index, &is_live) in live.iter().enumerate() {
43            if is_live {
44                rebuild.copy(crate::ValueId(index as u32))?;
45            } else {
46                removed += 1;
47            }
48        }
49        Ok(PassOutput {
50            graph: rebuild.finish()?,
51            changed: removed,
52        })
53    }
54}
55
56#[cfg(test)]
57mod tests {
58    use super::*;
59    use crate::{GraphBuilder, Op};
60
61    fn builder() -> GraphBuilder {
62        GraphBuilder::new(Limits::default())
63    }
64
65    fn run(graph: &Graph) -> PassOutput {
66        DeadCodeElimination::new(Limits::default())
67            .run(graph)
68            .unwrap()
69    }
70
71    #[test]
72    fn removes_dead_nodes_and_unused_constants() {
73        let mut b = builder();
74        let x = b.input("x", vec![2]).unwrap();
75        let c = b.constant("c", vec![2], vec![1.0, 2.0]).unwrap();
76        let live = b.node("live", Op::Relu, &[x]).unwrap();
77        let dead = b.node("dead", Op::Mul, &[x, c]).unwrap();
78        b.node("deader", Op::Relu, &[dead]).unwrap();
79        let g = b.finish(&[live]).unwrap();
80        let out = run(&g);
81        assert_eq!(out.changed, 3);
82        let names: Vec<String> = (0..out.graph.values().len())
83            .map(|i| out.graph.label(crate::ValueId(i as u32)))
84            .collect();
85        assert_eq!(names, ["x", "live"]);
86        assert_eq!(out.graph.output_names(), ["live"]);
87    }
88
89    #[test]
90    fn unused_inputs_survive() {
91        let mut b = builder();
92        let x = b.input("x", vec![2]).unwrap();
93        let z = b.input("z", vec![3]).unwrap();
94        let y = b.node("y", Op::Relu, &[x]).unwrap();
95        let g = b.finish(&[y]).unwrap();
96        let out = run(&g);
97        assert_eq!(out.changed, 0);
98        assert_eq!(out.graph.inputs().len(), 2);
99        assert_eq!(out.graph.label(z), "z");
100        assert_eq!(out.graph, g);
101    }
102
103    #[test]
104    fn outputs_that_are_inputs_or_constants_survive() {
105        let mut b = builder();
106        let x = b.input("x", vec![1]).unwrap();
107        let c = b.constant("c", vec![1], vec![5.0]).unwrap();
108        let g = b.finish(&[x, c, x]).unwrap();
109        let out = run(&g);
110        assert_eq!(out.changed, 0);
111        assert_eq!(out.graph, g);
112    }
113
114    #[test]
115    fn a_graph_with_no_dead_code_is_rebuilt_identically() {
116        let mut b = builder();
117        let x = b.input("x", vec![2, 2]).unwrap();
118        let w = b
119            .constant("w", vec![2, 2], vec![1.0, 2.0, 3.0, 4.0])
120            .unwrap();
121        let p = b.node("p", Op::MatMul, &[x, w]).unwrap();
122        let r = b.node("r", Op::Reshape { shape: vec![4] }, &[p]).unwrap();
123        let g = b.finish(&[r]).unwrap();
124        let out = run(&g);
125        assert_eq!(out.changed, 0);
126        assert_eq!(out.graph, g);
127    }
128
129    #[test]
130    fn a_node_used_by_live_and_dead_code_stays() {
131        let mut b = builder();
132        let x = b.input("x", vec![2]).unwrap();
133        let shared = b.node("shared", Op::Relu, &[x]).unwrap();
134        let live = b.node("live", Op::Sigmoid, &[shared]).unwrap();
135        b.node("dead", Op::Sigmoid, &[shared]).unwrap();
136        let g = b.finish(&[live]).unwrap();
137        let out = run(&g);
138        assert_eq!(out.changed, 1);
139        assert_eq!(out.graph.values().len(), 3);
140    }
141
142    #[test]
143    fn constants_are_shared_not_copied() {
144        let mut b = builder();
145        let c = b.constant("c", vec![2], vec![1.0, 2.0]).unwrap();
146        let g = b.finish(&[c]).unwrap();
147        let out = run(&g);
148        match (&g.value(c).source, &out.graph.value(c).source) {
149            (crate::Source::Constant(a), crate::Source::Constant(b)) => {
150                assert!(std::sync::Arc::ptr_eq(a, b))
151            }
152            _ => panic!("expected constants"),
153        }
154    }
155}
oxfuse_core/passes/dce.rs

oxfuse_core/passes/
dce.rs
