tract_core/ops/
mod.rs

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

tract_core/ops/
mod.rs
