tract_hir/infer/
ops.rs

1use super::Factoid;
2use crate::infer::*;
3use std::fmt;
4use tract_core::optim::CONST_FOLD_MEM_BUDGET;
5use tract_data::TooEarly;
6
7tract_core::dyn_clone::clone_trait_object!(InferenceOp);
8
9fn tensor_mem(t: &Tensor) -> u64 {
10    (t.volume() * t.datum_type().size_of()) as u64
11}
12
13/// An operation with tensor type inference
14pub trait InferenceOp: Op {
15    /// Infers properties about the input and output tensors.
16    ///
17    /// The `inputs` and `outputs` arguments correspond to properties about
18    /// the input and output tensors that are already known.
19    ///
20    /// The default implementation will call the private infer_facts method,
21    /// which is usually implemented using the InferenceRulesOp trait. It will
22    /// also try to eval() the op if its a EvalOp and if the inputs are
23    /// fully determined.
24    ///
25    /// Facts carry no graph, so that eager eval cannot count a constant's
26    /// consumers: on top of holding the output within
27    /// [`CONST_FOLD_MEM_BUDGET`], it requires every input to be plain and
28    /// itself under the allowance.
29    ///
30    /// Returns Err in case of an unrecoverable error during the inference,
31    /// and the refined properties about the inputs and outputs otherwise.
32    fn infer(
33        &mut self,
34        inputs: TVec<&InferenceFact>,
35        outputs: TVec<&InferenceFact>,
36        observed: TVec<&InferenceFact>,
37    ) -> TractResult<(TVec<InferenceFact>, TVec<InferenceFact>, TVec<InferenceFact>)> {
38        let (infered_inputs, infered_outputs, observed) =
39            self.infer_facts(inputs, outputs, observed).context("Infering facts")?;
40
41        if infered_inputs.len() > 0
42            && let Some(input_values) = infered_inputs
43                .iter()
44                .map(|i| {
45                    i.value
46                        .concretize()
47                        .filter(|t| t.is_plain() && tensor_mem(t) <= CONST_FOLD_MEM_BUDGET)
48                        .map(|t| t.into_tvalue())
49                })
50                .collect::<Option<TVec<_>>>()
51        {
52            let input_mem: u64 = input_values.iter().map(|t| tensor_mem(t)).sum();
53            match self.eval_out_of_plan(input_values) {
54                Ok(None) => (),
55                Ok(Some(values)) => {
56                    let output_mem: u64 = values.iter().map(|t| tensor_mem(t)).sum();
57                    if output_mem <= input_mem.max(CONST_FOLD_MEM_BUDGET) {
58                        let output_values = values
59                            .into_iter()
60                            .map(|t| t.into_arc_tensor().try_into())
61                            .collect::<TractResult<TVec<_>>>()?;
62                        return Ok((infered_inputs, output_values, observed));
63                    }
64                }
65                Err(e) if e.root_cause().downcast_ref::<TooEarly>().is_some() => (),
66                Err(e) => return Err(e).context("Eager eval during inference"),
67            }
68        }
69
70        Ok((infered_inputs, infered_outputs, observed))
71    }
72
73    /// Allow an op to specify a supplementary list of outlets facts that
74    /// will trigger inference again.
75    fn observe_outlets(
76        &self,
77        _model: &InferenceModel,
78        _node: &InferenceNode,
79    ) -> TractResult<Vec<OutletId>> {
80        Ok(vec![])
81    }
82
83    /// Infer properties about inputs and output tensors. This method does not
84    /// need to deal with the "trivial" stateless op with fully determined
85    /// inputs cases.
86    ///
87    /// Most of the time, it is implemented using InferenceRulesOp.
88    fn infer_facts(
89        &mut self,
90        inputs: TVec<&InferenceFact>,
91        outputs: TVec<&InferenceFact>,
92        observed: TVec<&InferenceFact>,
93    ) -> TractResult<(TVec<InferenceFact>, TVec<InferenceFact>, TVec<InferenceFact>)>;
94
95    /// Early pass on inference model, after analyse, but before translation to
96    /// typed network. Meant to deal with some framework idiosyncrasies that
97    /// manifest with temporaries nodes that can run some form of inference but
98    /// require refactoring the network before it can be evaluated.
99    ///
100    /// Called after successful analyse, but before translating to typed model.
101    #[allow(unused_variables)]
102    fn incorporate(
103        &self,
104        model: &InferenceModel,
105        node: &InferenceNode,
106    ) -> TractResult<Option<InferenceModelPatch>> {
107        Ok(None)
108    }
109
110    fn nboutputs(&self) -> TractResult<usize> {
111        Ok(1)
112    }
113
114    /// Reinterpret the InferenceOp as an Op.
115    fn as_op(&self) -> &dyn Op;
116
117    /// Reinterpret the InferenceOp as an Op, mutably.
118    fn as_op_mut(&mut self) -> &mut dyn Op;
119
120    /// Called during translation to TypedModel.
121    #[allow(unused_variables)]
122    fn to_typed(
123        &self,
124        source: &InferenceModel,
125        node: &InferenceNode,
126        target: &mut TypedModel,
127        mapping: &HashMap<OutletId, OutletId>,
128    ) -> TractResult<TVec<OutletId>> {
129        bail!("Operator can not be made a TypedOp.")
130    }
131}
132
133impl std::fmt::Display for Box<dyn InferenceOp> {
134    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
135        write!(fmt, "{}", self.name())
136    }
137}
138
139impl<O: InferenceOp> From<O> for Box<dyn InferenceOp> {
140    fn from(it: O) -> Box<dyn InferenceOp> {
141        Box::new(it)
142    }
143}
144
145impl AsRef<dyn Op> for dyn InferenceOp {
146    fn as_ref(&self) -> &dyn Op {
147        self.as_op()
148    }
149}
150
151impl AsRef<dyn Op> for Box<dyn InferenceOp> {
152    fn as_ref(&self) -> &dyn Op {
153        self.as_op()
154    }
155}
156
157impl AsMut<dyn Op> for dyn InferenceOp {
158    fn as_mut(&mut self) -> &mut dyn Op {
159        self.as_op_mut()
160    }
161}
162
163impl AsMut<dyn Op> for Box<dyn InferenceOp> {
164    fn as_mut(&mut self) -> &mut dyn Op {
165        self.as_op_mut()
166    }
167}
tract_hir/infer/ops.rs

tract_hir/infer/
ops.rs
