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

tract_core/ops/
mod.rs
