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    /// Whether evaluating this op can *bind* runtime symbols (symbols that only
199    /// get a value when tensors are produced: the length of a dynamic Range, a
200    /// NonZero count...). The plan's shape-feedback loop resolves such symbols
201    /// as these nodes evaluate, and uses this marker to order other nodes that
202    /// mention them (in their shapes or in symbolic tensor values) after their
203    /// defining node.
204    ///
205    /// Sources must be marked even though `set_input` binds their symbols
206    /// before evaluation: they come first in dataflow, so they claim their
207    /// shape's symbols in the definer map before an over-eager marker (say, a
208    /// Scan propagating a source-defined symbol in its output shape) would.
209    fn mints_runtime_symbols(&self) -> bool {
210        false
211    }
212
213    /// Symbols that must be bound before this op's evaluation can succeed,
214    /// beyond those visible in its output facts. The default collects the
215    /// symbols carried by symbolic tensor *values* on the input side (a TDim
216    /// konst or uniform, e.g. a folded Shape feeding a Cast): the plan's
217    /// static ordering uses these to place the node after whatever binds
218    /// them, so hidden mentions are not missed. Ops whose needs are fully
219    /// covered by their output facts need not override this.
220    fn required_bound_symbols(&self, inputs: &[&TypedFact]) -> TVec<Symbol> {
221        let mut syms: TVec<Symbol> = tvec![];
222        for fact in inputs {
223            for t in [fact.konst.as_deref(), fact.uniform.as_deref()].into_iter().flatten() {
224                if t.datum_type() == TDim::datum_type()
225                    && let Some(slice) =
226                        t.try_as_plain_ram().ok().and_then(|v| v.as_slice::<TDim>().ok())
227                {
228                    syms.extend(slice.iter().flat_map(|d| d.symbols()));
229                }
230            }
231        }
232        syms
233    }
234
235    fn as_typed(&self) -> Option<&dyn TypedOp>;
236}
237
238impl_downcast!(Op);
239dyn_clone::clone_trait_object!(Op);
240dyn_eq::eq_trait_object!(Op);
241
242pub trait TypedOp:
243    Op + fmt::Debug + dyn_clone::DynClone + Send + Sync + 'static + Downcast + EvalOp
244{
245    /// Reinterpret the TypedOp as an Op.
246    fn as_op(&self) -> &dyn Op;
247
248    /// Reinterpret the TypedOp as an Op, mutably.
249    fn as_op_mut(&mut self) -> &mut dyn Op;
250
251    /// Deduce output facts from input facts.
252    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>>;
253
254    #[allow(unused_variables)]
255    fn axes_mapping(
256        &self,
257        inputs: &[&TypedFact],
258        outputs: &[&TypedFact],
259    ) -> TractResult<AxesMapping> {
260        AxesMapping::disconnected(inputs, outputs)
261    }
262
263    /// Fuse op after codegen to deal with local optimisations.
264    fn fuse(&self, _model: &TypedModel, _node: &TypedNode) -> TractResult<Option<TypedModelPatch>> {
265        Ok(None)
266    }
267
268    /// Declutter the op to the tract_core operator set as much as possible.
269    #[allow(unused_variables)]
270    fn declutter_with_session(
271        &self,
272        session: &mut OptimizerSession,
273        model: &TypedModel,
274        node: &TypedNode,
275    ) -> TractResult<Option<TypedModelPatch>> {
276        self.declutter(model, node)
277    }
278
279    /// Declutter the op to the tract_core operator set as much as possible.
280    #[allow(unused_variables)]
281    fn declutter(
282        &self,
283        model: &TypedModel,
284        node: &TypedNode,
285    ) -> TractResult<Option<TypedModelPatch>> {
286        Ok(None)
287    }
288
289    /// Computes a cost hint of the operation.
290    ///
291    /// Each pair is a type of operation and a number per call on eval.
292    fn cost(&self, _inputs: &[&TypedFact]) -> TractResult<TVec<(Cost, TDim)>> {
293        Ok(tvec!())
294    }
295
296    /// Derive ROI (region of interest) expressions for this node's inputs.
297    /// Called by the PropagateRoi pass. Default returns None (no propagation).
298    /// Override to introduce ROIs or bubble them through.
299    #[allow(unused_variables)]
300    fn input_roi(
301        &self,
302        model: &TypedModel,
303        node: &TypedNode,
304    ) -> TractResult<Option<TVec<Option<TDim>>>> {
305        Ok(None)
306    }
307
308    #[allow(unused_variables)]
309    fn suggested_axis_changes(&self) -> TractResult<TVec<(InOut, AxisOp)>> {
310        Ok(tvec!())
311    }
312
313    #[allow(unused_variables)]
314    fn change_axes(
315        &self,
316        model: &TypedModel,
317        node: &TypedNode,
318        io: InOut,
319        change: &AxisOp,
320    ) -> TractResult<Option<AxisChangeConsequence>> {
321        Ok(None)
322    }
323
324    #[allow(unused_variables)]
325    #[allow(clippy::too_many_arguments)]
326    fn slice(
327        &self,
328        patch: &mut TypedModelPatch,
329        model: &TypedModel,
330        node: &TypedNode,
331        prefix: &str,
332        inputs: &[OutletId],
333        output_axis: usize,
334        start: &TDim,
335        end: &TDim,
336    ) -> TractResult<Option<TVec<OutletId>>> {
337        Ok(None)
338    }
339
340    /// Transforms the op in an equivalent one, operating on dt (i8 or u8).
341    ///
342    /// Returns None if the op can not be translated.
343    #[allow(unused_variables)]
344    fn quantize(
345        &self,
346        model: &TypedModel,
347        node: &TypedNode,
348        dt: DatumType,
349        scale: f32,
350        zero_point: i32,
351    ) -> TractResult<Option<Box<dyn TypedOp>>> {
352        Ok(None)
353    }
354
355    /// Transform the op by substituting one or more symbols with TDim
356    /// expressions (a concrete integer is `TDim::Val(v)`; an expression
357    /// can be any other TDim, including symbolic ones).
358    #[allow(unused_variables)]
359    fn set_symbols(
360        &self,
361        source: &TypedModel,
362        node: &TypedNode,
363        target: &mut TypedModel,
364        mapping: &HashMap<OutletId, OutletId>,
365        subs: &HashMap<Symbol, TDim>,
366    ) -> TractResult<TVec<OutletId>> {
367        let inputs = node.inputs.iter().map(|i| mapping[i]).collect::<TVec<_>>();
368        target.wire_node(&node.name, node.op.clone(), &inputs)
369    }
370
371    /// Translate the op into the most efficient form possible for execution.
372    ///
373    /// This transformation is supposed to be final, no more pass are expected
374    /// to be run on the codegen networks.
375    #[allow(unused_variables)]
376    fn codegen(
377        &self,
378        model: &TypedModel,
379        node: &TypedNode,
380    ) -> TractResult<Option<TypedModelPatch>> {
381        Ok(None)
382    }
383
384    /// Nested model multipliers, with label (for profiling).
385    #[allow(unused_variables)]
386    fn nested_model_multipliers(&self, inputs: &[&TypedFact]) -> Vec<(StaticName, TDim)> {
387        vec![]
388    }
389}
390
391impl_downcast!(TypedOp);
392dyn_clone::clone_trait_object!(TypedOp);
393dyn_eq::eq_trait_object!(TypedOp);
394
395impl<O: Op> From<O> for Box<dyn Op> {
396    fn from(it: O) -> Box<dyn Op> {
397        Box::new(it)
398    }
399}
400
401impl<O: TypedOp> From<O> for Box<dyn TypedOp> {
402    fn from(it: O) -> Box<dyn TypedOp> {
403        Box::new(it)
404    }
405}
406
407impl<'a> From<&'a Box<dyn TypedOp>> for Box<dyn TypedOp> {
408    fn from(it: &'a Box<dyn TypedOp>) -> Box<dyn TypedOp> {
409        it.clone()
410    }
411}
412
413impl AsRef<dyn Op> for dyn TypedOp {
414    fn as_ref(&self) -> &dyn Op {
415        self.as_op()
416    }
417}
418
419impl AsRef<dyn Op> for Box<dyn TypedOp> {
420    fn as_ref(&self) -> &dyn Op {
421        self.as_op()
422    }
423}
424
425impl AsMut<dyn Op> for dyn TypedOp {
426    fn as_mut(&mut self) -> &mut dyn Op {
427        self.as_op_mut()
428    }
429}
430
431impl AsMut<dyn Op> for Box<dyn TypedOp> {
432    fn as_mut(&mut self) -> &mut dyn Op {
433        self.as_op_mut()
434    }
435}
436
437impl std::fmt::Display for Box<dyn Op> {
438    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
439        write!(fmt, "{}", self.name())
440    }
441}
442
443impl std::fmt::Display for Box<dyn TypedOp> {
444    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
445        write!(fmt, "{}", self.name())
446    }
447}
tract_core/ops/mod.rs

tract_core/ops/
mod.rs
