tract_core/ops/
mod.rs

1//! Ops
2use std::fmt;
3
4use downcast_rs::Downcast;
5
6use dyn_clone;
7use dyn_eq::DynEq;
8
9#[macro_use]
10pub mod macros;
11#[macro_use]
12pub mod element_wise;
13#[macro_use]
14pub mod binary;
15
16pub mod array;
17pub mod cast;
18pub mod change_axes;
19pub mod cnn;
20pub mod downsample;
21pub mod dummy;
22pub mod einsum;
23pub mod fft;
24pub mod identity;
25pub mod konst;
26pub mod logic;
27pub mod lstm_cell;
28pub mod math;
29pub mod matmul;
30// pub mod memory;
31pub mod nn;
32pub mod quant;
33pub mod scan;
34pub mod source;
35pub mod submodel;
36pub mod unimpl;
37
38pub use downsample::Downsample;
39pub use memory::*;
40
41use crate::internal::*;
42use crate::optim::OptimizerSession;
43
44/// Level of precision to be expected in implementations comparisons.
45#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
46pub enum Validation {
47    /// Output is random
48    Random,
49    /// Implementation may induce rounding errors
50    Rounding,
51    /// Implementation must be accurate
52    Accurate,
53}
54
55#[derive(Clone, PartialEq, Eq, Hash, Ord, PartialOrd)]
56pub enum Cost {
57    Div(DatumType),
58    FMA(DatumType),
59    Buffer(DatumType),
60    Params(DatumType),
61    Custom(bool, String),
62}
63
64impl Cost {
65    pub fn is_compute(&self) -> bool {
66        use Cost::*;
67        match self {
68            FMA(_) | Div(_) => true,
69            Buffer(_) | Params(_) => false,
70            Custom(compute, _) => *compute,
71        }
72    }
73}
74
75impl std::fmt::Debug for Cost {
76    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
77        use Cost::*;
78        match self {
79            Div(dt) => write!(f, "Div({dt:?})"),
80            FMA(dt) => write!(f, "FMA({dt:?})"),
81            Buffer(dt) => write!(f, "Buffer({dt:?})"),
82            Params(dt) => write!(f, "Params({dt:?})"),
83            Custom(_, name) => write!(f, "{name}"),
84        }
85    }
86}
87
88pub trait FrozenOpState: fmt::Debug + dyn_clone::DynClone + Send + 'static {
89    fn unfreeze(&self) -> Box<dyn OpState>;
90}
91
92pub trait OpStateFreeze {
93    fn freeze(&self) -> Box<dyn FrozenOpState>;
94    /// Consuming freeze: moves data instead of cloning. Default delegates to freeze().
95    fn freeze_into(self: Box<Self>) -> Box<dyn FrozenOpState> {
96        self.freeze()
97    }
98}
99
100dyn_clone::clone_trait_object!(FrozenOpState);
101
102pub trait OpState: fmt::Debug + dyn_clone::DynClone + OpStateFreeze + Downcast {
103    fn load_from(
104        &mut self,
105        _: &mut TurnState,
106        _: &mut dyn Iterator<Item = TValue>,
107    ) -> TractResult<()> {
108        Ok(())
109    }
110
111    fn save_to(&self, _: &mut Vec<TValue>) -> TractResult<()> {
112        Ok(())
113    }
114
115    fn init_tensor_fact(&self) -> Option<(String, TypedFact)> {
116        None
117    }
118
119    /// Allocation-free predicate mirroring whether [`OpState::init_tensor_fact`]
120    /// returns `Some`. The per-run symbol-resolution path queries this once for
121    /// every stateful op on every `run`, so it must not call `init_tensor_fact`
122    /// (which clones a `String` and a `TypedFact`) merely to test for presence.
123    /// Any impl that overrides `init_tensor_fact` to return `Some` must override
124    /// this to return `true` (and delegate it wherever `init_tensor_fact` is
125    /// delegated), or its `resolve_symbols` will not run.
126    fn has_init_tensor_fact(&self) -> bool {
127        false
128    }
129
130    fn resolve_symbols(&mut self, _: &mut TurnState) -> TractResult<()> {
131        Ok(())
132    }
133
134    fn eval(
135        &mut self,
136        session: &mut TurnState,
137        op: &dyn Op,
138        inputs: TVec<TValue>,
139    ) -> TractResult<TVec<TValue>>;
140}
141dyn_clone::clone_trait_object!(OpState);
142impl_downcast!(OpState);
143
144pub trait EvalOp {
145    #[allow(unused_variables)]
146    fn eval(&self, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
147        bail!("stateless evaluation not implemented")
148    }
149
150    #[allow(unused_variables)]
151    fn eval_with_session(
152        &self,
153        node_id: usize,
154        session: &TurnState,
155        inputs: TVec<TValue>,
156    ) -> TractResult<TVec<TValue>> {
157        self.eval(inputs).context("Running legacy eval")
158    }
159
160    #[allow(unused_variables)]
161    fn state(&self, session: &TurnState, node_id: usize) -> TractResult<Option<Box<dyn OpState>>> {
162        Ok(None)
163    }
164
165    fn is_stateless(&self) -> bool;
166}
167
168/// A base operation
169pub trait Op:
170    fmt::Debug + dyn_clone::DynClone + dyn_eq::DynEq + Send + Sync + 'static + Downcast + EvalOp
171{
172    fn name(&self) -> StaticName;
173
174    /// The kind of accuracy check that should be performed on operation when
175    /// testing them.
176    fn validation(&self) -> Validation {
177        Validation::Accurate
178    }
179
180    /// Short (one-line) strings giving hints on internal implementation or
181    /// important configuration details to be displayed in dumps.
182    fn info(&self) -> TractResult<Vec<String>> {
183        Ok(vec![])
184    }
185
186    fn as_typed(&self) -> Option<&dyn TypedOp>;
187}
188
189impl_downcast!(Op);
190dyn_clone::clone_trait_object!(Op);
191dyn_eq::eq_trait_object!(Op);
192
193pub trait TypedOp:
194    Op + fmt::Debug + dyn_clone::DynClone + Send + Sync + 'static + Downcast + EvalOp
195{
196    /// Reinterpret the TypedOp as an Op.
197    fn as_op(&self) -> &dyn Op;
198
199    /// Reinterpret the TypedOp as an Op, mutably.
200    fn as_op_mut(&mut self) -> &mut dyn Op;
201
202    /// Deduce output facts from input facts.
203    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>>;
204
205    #[allow(unused_variables)]
206    fn axes_mapping(
207        &self,
208        inputs: &[&TypedFact],
209        outputs: &[&TypedFact],
210    ) -> TractResult<AxesMapping> {
211        AxesMapping::disconnected(inputs, outputs)
212    }
213
214    /// Fuse op after codegen to deal with local optimisations.
215    fn fuse(&self, _model: &TypedModel, _node: &TypedNode) -> TractResult<Option<TypedModelPatch>> {
216        Ok(None)
217    }
218
219    /// Declutter the op to the tract_core operator set as much as possible.
220    #[allow(unused_variables)]
221    fn declutter_with_session(
222        &self,
223        session: &mut OptimizerSession,
224        model: &TypedModel,
225        node: &TypedNode,
226    ) -> TractResult<Option<TypedModelPatch>> {
227        self.declutter(model, node)
228    }
229
230    /// Declutter the op to the tract_core operator set as much as possible.
231    #[allow(unused_variables)]
232    fn declutter(
233        &self,
234        model: &TypedModel,
235        node: &TypedNode,
236    ) -> TractResult<Option<TypedModelPatch>> {
237        Ok(None)
238    }
239
240    /// Computes a cost hint of the operation.
241    ///
242    /// Each pair is a type of operation and a number per call on eval.
243    fn cost(&self, _inputs: &[&TypedFact]) -> TractResult<TVec<(Cost, TDim)>> {
244        Ok(tvec!())
245    }
246
247    /// Derive ROI (region of interest) expressions for this node's inputs.
248    /// Called by the PropagateRoi pass. Default returns None (no propagation).
249    /// Override to introduce ROIs or bubble them through.
250    #[allow(unused_variables)]
251    fn input_roi(
252        &self,
253        model: &TypedModel,
254        node: &TypedNode,
255    ) -> TractResult<Option<TVec<Option<TDim>>>> {
256        Ok(None)
257    }
258
259    #[allow(unused_variables)]
260    fn suggested_axis_changes(&self) -> TractResult<TVec<(InOut, AxisOp)>> {
261        Ok(tvec!())
262    }
263
264    #[allow(unused_variables)]
265    fn change_axes(
266        &self,
267        model: &TypedModel,
268        node: &TypedNode,
269        io: InOut,
270        change: &AxisOp,
271    ) -> TractResult<Option<AxisChangeConsequence>> {
272        Ok(None)
273    }
274
275    #[allow(unused_variables)]
276    #[allow(clippy::too_many_arguments)]
277    fn slice(
278        &self,
279        patch: &mut TypedModelPatch,
280        model: &TypedModel,
281        node: &TypedNode,
282        prefix: &str,
283        inputs: &[OutletId],
284        output_axis: usize,
285        start: &TDim,
286        end: &TDim,
287    ) -> TractResult<Option<TVec<OutletId>>> {
288        Ok(None)
289    }
290
291    /// Transforms the op in an equivalent one, operating on dt (i8 or u8).
292    ///
293    /// Returns None if the op can not be translated.
294    #[allow(unused_variables)]
295    fn quantize(
296        &self,
297        model: &TypedModel,
298        node: &TypedNode,
299        dt: DatumType,
300        scale: f32,
301        zero_point: i32,
302    ) -> TractResult<Option<Box<dyn TypedOp>>> {
303        Ok(None)
304    }
305
306    /// Transform the op by substituting one or more symbols with TDim
307    /// expressions (a concrete integer is `TDim::Val(v)`; an expression
308    /// can be any other TDim, including symbolic ones).
309    #[allow(unused_variables)]
310    fn set_symbols(
311        &self,
312        source: &TypedModel,
313        node: &TypedNode,
314        target: &mut TypedModel,
315        mapping: &HashMap<OutletId, OutletId>,
316        subs: &HashMap<Symbol, TDim>,
317    ) -> TractResult<TVec<OutletId>> {
318        let inputs = node.inputs.iter().map(|i| mapping[i]).collect::<TVec<_>>();
319        target.wire_node(&node.name, node.op.clone(), &inputs)
320    }
321
322    /// Translate the op into the most efficient form possible for execution.
323    ///
324    /// This transformation is supposed to be final, no more pass are expected
325    /// to be run on the codegen networks.
326    #[allow(unused_variables)]
327    fn codegen(
328        &self,
329        model: &TypedModel,
330        node: &TypedNode,
331    ) -> TractResult<Option<TypedModelPatch>> {
332        Ok(None)
333    }
334
335    /// Nested model multipliers, with label (for profiling).
336    #[allow(unused_variables)]
337    fn nested_model_multipliers(&self, inputs: &[&TypedFact]) -> Vec<(StaticName, TDim)> {
338        vec![]
339    }
340}
341
342impl_downcast!(TypedOp);
343dyn_clone::clone_trait_object!(TypedOp);
344dyn_eq::eq_trait_object!(TypedOp);
345
346impl<O: Op> From<O> for Box<dyn Op> {
347    fn from(it: O) -> Box<dyn Op> {
348        Box::new(it)
349    }
350}
351
352impl<O: TypedOp> From<O> for Box<dyn TypedOp> {
353    fn from(it: O) -> Box<dyn TypedOp> {
354        Box::new(it)
355    }
356}
357
358impl<'a> From<&'a Box<dyn TypedOp>> for Box<dyn TypedOp> {
359    fn from(it: &'a Box<dyn TypedOp>) -> Box<dyn TypedOp> {
360        it.clone()
361    }
362}
363
364impl AsRef<dyn Op> for dyn TypedOp {
365    fn as_ref(&self) -> &dyn Op {
366        self.as_op()
367    }
368}
369
370impl AsRef<dyn Op> for Box<dyn TypedOp> {
371    fn as_ref(&self) -> &dyn Op {
372        self.as_op()
373    }
374}
375
376impl AsMut<dyn Op> for dyn TypedOp {
377    fn as_mut(&mut self) -> &mut dyn Op {
378        self.as_op_mut()
379    }
380}
381
382impl AsMut<dyn Op> for Box<dyn TypedOp> {
383    fn as_mut(&mut self) -> &mut dyn Op {
384        self.as_op_mut()
385    }
386}
387
388impl std::fmt::Display for Box<dyn Op> {
389    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
390        write!(fmt, "{}", self.name())
391    }
392}
393
394impl std::fmt::Display for Box<dyn TypedOp> {
395    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
396        write!(fmt, "{}", self.name())
397    }
398}
tract_core/ops/mod.rs

tract_core/ops/
mod.rs
