oxfuse_core/passes/
mod.rs

1//! Graph-to-graph passes and the manager that runs them. A pass never
2//! mutates its input: it returns a new graph, so before and after are both
3//! intact for dumps and equivalence tests, and no stale ids can exist.
4
5mod cse;
6mod dce;
7mod fold;
8mod layout;
9mod rebuild;
10
11pub use cse::CommonSubexpressionElimination;
12pub use dce::DeadCodeElimination;
13pub use fold::{ConstEvaluator, ConstantFolding};
14pub use layout::LayoutElision;
15
16use crate::{Error, Graph, Limits, TensorType};
17use std::sync::Arc;
18
19/// What a pass returns: the new graph and how many rewrites it made.
20#[derive(Debug)]
21pub struct PassOutput {
22    pub graph: Graph,
23    pub changed: usize,
24}
25
26pub trait Pass {
27    /// Short stable name, used by reports and `--disable-pass`.
28    fn name(&self) -> &'static str;
29    fn run(&self, graph: &Graph) -> Result<PassOutput, Error>;
30}
31
32#[derive(Debug, Clone, PartialEq, Eq)]
33pub struct PassReport {
34    pub pass: &'static str,
35    pub values_before: usize,
36    pub values_after: usize,
37    pub changed: usize,
38}
39
40#[derive(Debug)]
41pub struct Optimized {
42    pub graph: Graph,
43    pub reports: Vec<PassReport>,
44    /// `(label, ir text)` per stage; empty unless dumps were requested.
45    pub stages: Vec<(String, String)>,
46}
47
48impl Optimized {
49    pub fn dump_text(&self) -> String {
50        self.stages
51            .iter()
52            .map(|(label, ir)| format!("== {label} ==\n{ir}\n"))
53            .collect()
54    }
55}
56
57pub struct PassManager {
58    passes: Vec<Box<dyn Pass>>,
59    limits: Limits,
60    dump: bool,
61}
62
63impl PassManager {
64    pub fn new(limits: Limits) -> Self {
65        Self {
66            passes: Vec::new(),
67            limits,
68            dump: false,
69        }
70    }
71
72    #[allow(clippy::should_implement_trait)]
73    pub fn add(mut self, pass: impl Pass + 'static) -> Self {
74        self.passes.push(Box::new(pass));
75        self
76    }
77
78    /// Removes every pass with this name, for testing a pipeline without it.
79    pub fn without(mut self, name: &str) -> Self {
80        self.passes.retain(|pass| pass.name() != name);
81        self
82    }
83
84    pub fn with_dumps(mut self, dump: bool) -> Self {
85        self.dump = dump;
86        self
87    }
88
89    pub fn pass_names(&self) -> Vec<&'static str> {
90        self.passes.iter().map(|pass| pass.name()).collect()
91    }
92
93    /// Runs the passes once, in order. The graph is verified and checked
94    /// against the limits before the first pass and after every pass.
95    pub fn run(&self, graph: &Graph) -> Result<Optimized, Error> {
96        check(graph, &self.limits, "input graph is invalid")?;
97        let expected = interface(graph);
98        let mut current = graph.clone();
99        let mut reports = Vec::new();
100        let mut stages = Vec::new();
101        if self.dump {
102            stages.push(("input".to_owned(), current.to_string()));
103        }
104        for pass in &self.passes {
105            let values_before = current.values().len();
106            let output = pass.run(&current)?;
107            check(
108                &output.graph,
109                &self.limits,
110                &format!("pass `{}` produced an invalid graph", pass.name()),
111            )?;
112            if interface(&output.graph) != expected {
113                return Err(Error::Verify(format!(
114                    "pass `{}` changed the graph interface",
115                    pass.name()
116                )));
117            }
118            let report = PassReport {
119                pass: pass.name(),
120                values_before,
121                values_after: output.graph.values().len(),
122                changed: output.changed,
123            };
124            if self.dump {
125                stages.push((
126                    format!(
127                        "after {} (changed {}, values {} -> {})",
128                        report.pass, report.changed, report.values_before, report.values_after
129                    ),
130                    output.graph.to_string(),
131                ));
132            }
133            reports.push(report);
134            current = output.graph;
135        }
136        Ok(Optimized {
137            graph: current,
138            reports,
139            stages,
140        })
141    }
142}
143
144/// What users see of a graph: inputs (name, type, order), output names and
145/// output types. No pass may change it.
146type Interface = (Vec<(String, TensorType)>, Vec<String>, Vec<TensorType>);
147
148fn interface(graph: &Graph) -> Interface {
149    (
150        graph
151            .inputs()
152            .iter()
153            .map(|&id| (graph.label(id), graph.value(id).ty.clone()))
154            .collect(),
155        graph.output_names().to_vec(),
156        graph
157            .outputs()
158            .iter()
159            .map(|&id| graph.value(id).ty.clone())
160            .collect(),
161    )
162}
163
164fn check(graph: &Graph, limits: &Limits, context: &str) -> Result<(), Error> {
165    graph
166        .verify()
167        .and_then(|()| graph.verify_limits(limits))
168        .map_err(|e| Error::Verify(format!("{context}: {e}")))
169}
170
171/// The standard pipeline: fold constants, remove layout round trips, merge
172/// duplicates, drop what the outputs no longer need. The order is deliberate:
173/// folding creates constants that leave their inputs dead, layout elision can
174/// make two nodes identical for CSE, and dead-code elimination runs last to
175/// collect what the others leave unused.
176pub fn default_pipeline(limits: Limits, evaluator: Arc<dyn ConstEvaluator>) -> PassManager {
177    PassManager::new(limits.clone())
178        .add(ConstantFolding::new(evaluator, limits.clone()))
179        .add(LayoutElision::new(limits.clone()))
180        .add(CommonSubexpressionElimination::new(limits.clone()))
181        .add(DeadCodeElimination::new(limits))
182}
183
184#[cfg(test)]
185mod tests {
186    use super::*;
187    use crate::{GraphBuilder, Op};
188
189    fn sample() -> Graph {
190        let mut b = GraphBuilder::new(Limits::default());
191        let x = b.input("x", vec![2]).unwrap();
192        let y = b.node("y", Op::Relu, &[x]).unwrap();
193        b.finish(&[y]).unwrap()
194    }
195
196    struct Identity;
197    impl Pass for Identity {
198        fn name(&self) -> &'static str {
199            "identity"
200        }
201        fn run(&self, graph: &Graph) -> Result<PassOutput, Error> {
202            Ok(PassOutput {
203                graph: graph.clone(),
204                changed: 0,
205            })
206        }
207    }
208
209    /// Returns a graph with no outputs, which the verifier must reject.
210    struct Breaker;
211    impl Pass for Breaker {
212        fn name(&self) -> &'static str {
213            "breaker"
214        }
215        fn run(&self, graph: &Graph) -> Result<PassOutput, Error> {
216            let mut broken = graph.clone();
217            broken.outputs.clear();
218            broken.output_names.clear();
219            Ok(PassOutput {
220                graph: broken,
221                changed: 1,
222            })
223        }
224    }
225
226    #[test]
227    fn passes_run_in_order_and_are_reported() {
228        let manager = PassManager::new(Limits::default())
229            .add(Identity)
230            .add(Identity);
231        let optimized = manager.run(&sample()).unwrap();
232        assert_eq!(optimized.graph, sample());
233        assert_eq!(
234            optimized.reports,
235            vec![
236                PassReport {
237                    pass: "identity",
238                    values_before: 2,
239                    values_after: 2,
240                    changed: 0
241                },
242                PassReport {
243                    pass: "identity",
244                    values_before: 2,
245                    values_after: 2,
246                    changed: 0
247                },
248            ]
249        );
250    }
251
252    #[test]
253    fn without_removes_every_pass_with_that_name() {
254        let manager = PassManager::new(Limits::default())
255            .add(Identity)
256            .add(Breaker)
257            .without("breaker");
258        assert_eq!(manager.pass_names(), ["identity"]);
259        assert!(manager.run(&sample()).is_ok());
260    }
261
262    #[test]
263    fn a_pass_that_breaks_the_graph_is_named() {
264        let manager = PassManager::new(Limits::default())
265            .add(Identity)
266            .add(Breaker);
267        let err = manager.run(&sample()).unwrap_err();
268        assert!(
269            matches!(&err, Error::Verify(m) if m.contains("pass `breaker` produced an invalid graph")),
270            "{err}"
271        );
272    }
273
274    /// Renames the first output, which would change what users see.
275    struct Renamer;
276    impl Pass for Renamer {
277        fn name(&self) -> &'static str {
278            "renamer"
279        }
280        fn run(&self, graph: &Graph) -> Result<PassOutput, Error> {
281            let mut renamed = graph.clone();
282            renamed.output_names[0] = "renamed".to_owned();
283            Ok(PassOutput {
284                graph: renamed,
285                changed: 1,
286            })
287        }
288    }
289
290    #[test]
291    fn a_pass_that_changes_the_interface_is_rejected() {
292        let manager = PassManager::new(Limits::default()).add(Renamer);
293        let err = manager.run(&sample()).unwrap_err();
294        assert!(
295            matches!(&err, Error::Verify(m) if m.contains("pass `renamer` changed the graph interface")),
296            "{err}"
297        );
298    }
299
300    #[test]
301    fn invalid_input_is_rejected_before_any_pass_runs() {
302        let mut broken = sample();
303        broken.outputs.clear();
304        broken.output_names.clear();
305        let manager = PassManager::new(Limits::default()).add(Identity);
306        let err = manager.run(&broken).unwrap_err();
307        assert!(
308            matches!(&err, Error::Verify(m) if m.contains("input graph is invalid")),
309            "{err}"
310        );
311    }
312
313    #[test]
314    fn an_input_graph_over_the_limits_is_rejected() {
315        let tight = Limits {
316            max_values: 1,
317            ..Limits::default()
318        };
319        let manager = PassManager::new(tight).add(Identity);
320        let err = manager.run(&sample()).unwrap_err();
321        assert!(
322            matches!(&err, Error::Verify(m) if m.contains("input graph is invalid")),
323            "{err}"
324        );
325    }
326
327    #[test]
328    fn dumps_are_recorded_only_when_asked() {
329        let quiet = PassManager::new(Limits::default()).add(Identity);
330        assert!(quiet.run(&sample()).unwrap().stages.is_empty());
331        let loud = PassManager::new(Limits::default())
332            .add(Identity)
333            .with_dumps(true);
334        let optimized = loud.run(&sample()).unwrap();
335        let labels: Vec<&str> = optimized.stages.iter().map(|(l, _)| l.as_str()).collect();
336        assert_eq!(
337            labels,
338            ["input", "after identity (changed 0, values 2 -> 2)"]
339        );
340        let text = optimized.dump_text();
341        assert!(text.starts_with("== input ==\ngraph {\n"), "{text}");
342        assert!(
343            text.contains("== after identity (changed 0, values 2 -> 2) ==\n"),
344            "{text}"
345        );
346        assert!(text.ends_with("}\n"), "{text}");
347    }
348
349    #[test]
350    fn the_default_pipeline_runs_fold_cse_dce_in_that_order() {
351        struct Never;
352        impl ConstEvaluator for Never {
353            fn eval(
354                &self,
355                _: &crate::Op,
356                _: &[(&crate::TensorType, &[f32])],
357            ) -> Result<Vec<f32>, Error> {
358                Err(Error::Verify("unused".into()))
359            }
360        }
361        let manager = default_pipeline(Limits::default(), std::sync::Arc::new(Never));
362        assert_eq!(manager.pass_names(), ["fold", "layout", "cse", "dce"]);
363        assert_eq!(
364            manager.without("cse").pass_names(),
365            ["fold", "layout", "dce"]
366        );
367    }
368}
oxfuse_core/passes/mod.rs

oxfuse_core/passes/
mod.rs
