tract_core/ops/
submodel.rs

1use std::fmt::Debug;
2
3use tract_downcast_rs::Downcast;
4
5use crate::internal::*;
6
7#[derive(Debug, Clone)]
8pub struct SubmodelOp {
9    pub model: Box<dyn InnerModel>,
10    label: String,
11    decluttered: bool,
12    codegen: bool,
13}
14
15impl PartialEq for SubmodelOp {
16    fn eq(&self, _other: &Self) -> bool {
17        false
18    }
19}
20impl Eq for SubmodelOp {}
21
22impl SubmodelOp {
23    pub fn new(model: Box<dyn InnerModel>, label: &str) -> TractResult<Self> {
24        Ok(Self { model, label: label.to_string(), decluttered: false, codegen: false })
25    }
26
27    pub fn iteration_count(&self, _inputs: &[&TypedFact]) -> Option<TDim> {
28        None
29    }
30
31    pub fn model(&self) -> &TypedModel {
32        self.model.as_typed()
33    }
34
35    pub fn label(&self) -> &str {
36        self.label.as_str()
37    }
38}
39
40impl Op for SubmodelOp {
41    fn name(&self) -> StaticName {
42        "SubmodelOp".into()
43    }
44
45    op_as_typed_op!();
46}
47
48impl EvalOp for SubmodelOp {
49    not_out_of_plan!();
50
51    fn state(&self, ctx: &EvalContext) -> TractResult<Option<Box<dyn OpState>>> {
52        self.model.state(ctx)
53    }
54}
55
56impl TypedOp for SubmodelOp {
57    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
58        let facts = self.model.output_facts(inputs)?;
59        Ok(facts)
60    }
61
62    fn declutter(
63        &self,
64        model: &TypedModel,
65        node: &TypedNode,
66    ) -> TractResult<Option<TypedModelPatch>> {
67        if !self.decluttered {
68            let mut new = self.clone();
69            new.model.declutter()?;
70            new.decluttered = true;
71            Ok(Some(TypedModelPatch::replace_single_op(model, node, &node.inputs, new)?))
72        } else {
73            Ok(None)
74        }
75    }
76
77    fn codegen(
78        &self,
79        model: &TypedModel,
80        node: &TypedNode,
81    ) -> TractResult<Option<TypedModelPatch>> {
82        if !self.codegen {
83            let mut new = self.clone();
84            new.model.codegen()?;
85            new.codegen = true;
86            Ok(Some(TypedModelPatch::replace_single_op(model, node, &node.inputs, new)?))
87        } else {
88            Ok(None)
89        }
90    }
91
92    as_op!();
93}
94
95pub trait InnerModel: Debug + dyn_clone::DynClone + Downcast + Sync + Send + 'static {
96    #[allow(unused_variables)]
97    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>>;
98
99    #[allow(unused_variables)]
100    fn state(&self, ctx: &EvalContext) -> TractResult<Option<Box<dyn OpState>>> {
101        Ok(None)
102    }
103
104    #[allow(unused_variables)]
105    fn declutter(&mut self) -> TractResult<()>;
106
107    fn codegen(&mut self) -> TractResult<()>;
108
109    fn as_typed(&self) -> &TypedModel;
110}
111
112dyn_clone::clone_trait_object!(InnerModel);
113downcast_rs::impl_downcast!(InnerModel);
114
115impl InnerModel for TypedModel {
116    fn output_facts(&self, _inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
117        let facts = self
118            .output_outlets()?
119            .iter()
120            .map(|outlet| self.outlet_fact(*outlet).cloned())
121            .collect::<TractResult<TVec<_>>>()?;
122        Ok(facts)
123    }
124    #[allow(unused_variables)]
125    fn state(&self, ctx: &EvalContext) -> TractResult<Option<Box<dyn OpState>>> {
126        let plan = self.clone().into_runnable()?;
127        let state = plan.spawn()?;
128        Ok(Some(Box::new(state)))
129    }
130
131    #[allow(unused_variables)]
132    fn declutter(&mut self) -> TractResult<()> {
133        self.declutter()
134    }
135
136    fn codegen(&mut self) -> TractResult<()> {
137        self.optimize()
138    }
139
140    fn as_typed(&self) -> &TypedModel {
141        self
142    }
143}
144
145pub type TypedModelOpState = TypedSimpleState;
146
147impl OpState for TypedModelOpState {
148    fn eval(
149        &mut self,
150        _ctx: &EvalContext,
151        _op: &dyn Op,
152        inputs: TVec<TValue>,
153    ) -> TractResult<TVec<TValue>> {
154        let inference_out = self.run(inputs)?;
155        Ok(inference_out)
156    }
157
158    fn reset_lanes(&mut self, _lanes: &[LaneId]) -> TractResult<()> {
159        bail!("Submodel is not lane-aware: its body is a nested state")
160    }
161}
tract_core/ops/submodel.rs

tract_core/ops/
submodel.rs
