onnx_runtime_optimizer/
dead_node.rs

1//! Dead-node elimination: remove nodes whose outputs are not transitively
2//! consumed by any graph output (see `docs/ORT2.md` §18.1).
3
4use std::collections::HashSet;
5
6use onnx_runtime_ir::{Graph, NodeId, ValueId};
7
8use crate::error::Result;
9use crate::pass::{OptimizationPass, PassContext};
10
11/// Removes nodes that do not contribute to any graph output.
12///
13/// A node is *live* if any of its outputs is (transitively) required to compute
14/// a value in [`Graph::outputs`]. Liveness is found by walking backwards from
15/// the graph outputs through producer edges; every node not reached is removed
16/// via [`Graph::remove_nodes`], which keeps producer/consumer edges consistent
17/// and garbage-collects orphaned values. Graph inputs and initializers are left
18/// untouched.
19#[derive(Clone, Copy, Debug, Default)]
20pub struct DeadNodeElimination;
21
22impl DeadNodeElimination {
23    /// Node ids reachable (via producer edges) from the graph outputs.
24    fn live_nodes(graph: &Graph) -> HashSet<NodeId> {
25        let mut live: HashSet<NodeId> = HashSet::new();
26        let mut seen_values: HashSet<ValueId> = HashSet::new();
27        let mut stack: Vec<ValueId> = graph.outputs.clone();
28
29        while let Some(v) = stack.pop() {
30            if !seen_values.insert(v) {
31                continue;
32            }
33            let Some(val) = graph.try_value(v) else {
34                continue;
35            };
36            if let Some(prod) = val.producer
37                && live.insert(prod)
38            {
39                for iv in graph.node(prod).input_values() {
40                    stack.push(iv);
41                }
42            }
43        }
44        live
45    }
46}
47
48impl OptimizationPass for DeadNodeElimination {
49    fn name(&self) -> &str {
50        "DeadNodeElimination"
51    }
52
53    fn run(&self, graph: &mut Graph, _ctx: &PassContext) -> Result<()> {
54        let live = Self::live_nodes(graph);
55        let dead: Vec<NodeId> = graph
56            .nodes
57            .keys()
58            .filter(|nid| !live.contains(nid))
59            .collect();
60        graph.remove_nodes(&dead);
61        Ok(())
62    }
63}
64
65#[cfg(test)]
66mod tests {
67    use super::*;
68    use onnx_runtime_ir::{DataType, Node, NodeId, static_shape};
69
70    /// `a -> Relu -> b -> Add(b, c) -> d(out)`, plus a dangling
71    /// `a -> Neg -> e` branch whose result `e` feeds nothing.
72    fn graph_with_dangling_branch() -> (Graph, NodeId, NodeId) {
73        let mut g = Graph::new();
74        g.opset_imports.insert(String::new(), 17);
75        let a = g.create_named_value("a", DataType::Float32, static_shape([4]));
76        let c = g.create_named_value("c", DataType::Float32, static_shape([4]));
77        g.add_input(a);
78        g.add_input(c);
79
80        let b = g.create_value(DataType::Float32, static_shape([4]));
81        let relu = g.insert_node(Node::new(NodeId(0), "Relu", vec![Some(a)], vec![b]));
82
83        let d = g.create_named_value("d", DataType::Float32, static_shape([4]));
84        g.insert_node(Node::new(NodeId(0), "Add", vec![Some(b), Some(c)], vec![d]));
85        g.add_output(d);
86
87        // Dangling branch: consumes `a`, produces `e`, which nothing reads.
88        let e = g.create_value(DataType::Float32, static_shape([4]));
89        let neg = g.insert_node(Node::new(NodeId(0), "Neg", vec![Some(a)], vec![e]));
90
91        (g, relu, neg)
92    }
93
94    #[test]
95    fn removes_dangling_branch() {
96        let (mut g, _relu, neg) = graph_with_dangling_branch();
97        assert_eq!(g.num_nodes(), 3);
98        DeadNodeElimination
99            .run(&mut g, &PassContext::new())
100            .unwrap();
101        assert_eq!(g.num_nodes(), 2);
102        assert!(g.try_node(neg).is_none());
103        assert!(g.validate().is_ok());
104    }
105
106    #[test]
107    fn keeps_live_nodes() {
108        let (mut g, relu, _neg) = graph_with_dangling_branch();
109        DeadNodeElimination
110            .run(&mut g, &PassContext::new())
111            .unwrap();
112        // Relu and Add remain.
113        assert!(g.try_node(relu).is_some());
114        assert!(g.validate().is_ok());
115    }
116
117    #[test]
118    fn dead_value_is_garbage_collected() {
119        let (mut g, _relu, _neg) = graph_with_dangling_branch();
120        let before = g.num_values();
121        DeadNodeElimination
122            .run(&mut g, &PassContext::new())
123            .unwrap();
124        // `e` had no other consumer, so removing Neg should GC it.
125        assert!(g.num_values() < before);
126    }
127
128    #[test]
129    fn noop_on_all_live_graph() {
130        let mut g = Graph::new();
131        g.opset_imports.insert(String::new(), 17);
132        let a = g.create_named_value("a", DataType::Float32, static_shape([2]));
133        g.add_input(a);
134        let b = g.create_named_value("b", DataType::Float32, static_shape([2]));
135        g.insert_node(Node::new(NodeId(0), "Relu", vec![Some(a)], vec![b]));
136        g.add_output(b);
137
138        DeadNodeElimination
139            .run(&mut g, &PassContext::new())
140            .unwrap();
141        assert_eq!(g.num_nodes(), 1);
142        assert!(g.validate().is_ok());
143    }
144
145    #[test]
146    fn removes_transitively_dead_chain() {
147        // a(in) -> Relu -> b -> Neg -> e ; and a -> Abs -> out. b/e are dead.
148        let mut g = Graph::new();
149        g.opset_imports.insert(String::new(), 17);
150        let a = g.create_named_value("a", DataType::Float32, static_shape([2]));
151        g.add_input(a);
152
153        let b = g.create_value(DataType::Float32, static_shape([2]));
154        g.insert_node(Node::new(NodeId(0), "Relu", vec![Some(a)], vec![b]));
155        let e = g.create_value(DataType::Float32, static_shape([2]));
156        g.insert_node(Node::new(NodeId(0), "Neg", vec![Some(b)], vec![e]));
157
158        let out = g.create_named_value("out", DataType::Float32, static_shape([2]));
159        g.insert_node(Node::new(NodeId(0), "Abs", vec![Some(a)], vec![out]));
160        g.add_output(out);
161
162        DeadNodeElimination
163            .run(&mut g, &PassContext::new())
164            .unwrap();
165        assert_eq!(g.num_nodes(), 1); // only Abs survives
166        assert!(g.validate().is_ok());
167    }
168}
onnx_runtime_optimizer/dead_node.rs

onnx_runtime_optimizer/
dead_node.rs
