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

oxfuse_core/passes/
mod.rs
