tract_hir/infer/
optim.rs

1use crate::infer::*;
2use std::fmt;
3
4pub fn incorporate() -> Vec<Box<dyn IncorporatePass>> {
5    vec![Box::new(IncorporateOps)]
6}
7
8pub trait IncorporatePass: fmt::Debug + Send + Sync {
9    fn pass(&self, model: &mut InferenceModel) -> TractResult<bool>;
10}
11
12#[derive(Debug)]
13pub struct IncorporateOps;
14
15impl IncorporatePass for IncorporateOps {
16    fn pass(&self, model: &mut InferenceModel) -> TractResult<bool> {
17        let mut done_something = false;
18        loop {
19            let mut done_something_this_time = false;
20            for id in model.eval_order()? {
21                let reduced = {
22                    let node = &model.nodes()[id];
23                    trace!("Incorporate {node}");
24                    node.op
25                        .incorporate(model, node)
26                        .with_context(|| format!("{self:?} node {node}"))?
27                };
28                if let Some(red) = reduced {
29                    {
30                        let node = &model.nodes()[id];
31                        debug!("Apply a model patch for {self:?}: {node}");
32                    }
33                    red.apply(model)?;
34                    if cfg!(debug_assertions) {
35                        model.check_edges()?;
36                    }
37                    done_something_this_time = true;
38                }
39            }
40            done_something = done_something || done_something_this_time;
41            if !done_something_this_time {
42                break;
43            }
44        }
45        Ok(done_something)
46    }
47}