tract_core/ops/cnn/
maxpool.rs

1use crate::internal::*;
2use ndarray::prelude::*;
3
4use crate::ops::cnn::pools::{ConcretePoolGeometry, PoolGeometry, PoolSpec};
5
6#[derive(Debug, Clone, new, Hash, PartialEq, Eq)]
7pub struct MaxPool {
8    pub pool_spec: PoolSpec,
9    pub with_index_outputs: Option<DatumType>,
10}
11
12impl Op for MaxPool {
13    fn name(&self) -> StaticName {
14        "MaxPool".into()
15    }
16
17    fn info(&self) -> TractResult<Vec<String>> {
18        Ok(self.pool_spec.info())
19    }
20
21    op_as_typed_op!();
22}
23
24impl EvalOp for MaxPool {
25    fn is_stateless(&self) -> bool {
26        true
27    }
28
29    fn eval(&self, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
30        let shape: TVec<TDim> = inputs[0].shape().iter().map(|d| d.to_dim()).collect();
31        self.to_optimized(&shape)?.eval(inputs)
32    }
33}
34
35impl TypedOp for MaxPool {
36    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
37        let mut facts = self.pool_spec.output_facts(inputs)?;
38        if let Some(idt) = self.with_index_outputs {
39            facts.push(facts[0].clone());
40            facts[1].datum_type = idt;
41        }
42        Ok(facts)
43    }
44
45    fn declutter(
46        &self,
47        model: &TypedModel,
48        node: &TypedNode,
49    ) -> TractResult<Option<TypedModelPatch>> {
50        if self.with_index_outputs.is_some()
51            && node.outputs[1].successors.len() == 0
52            && !model.output_outlets()?.contains(&OutletId::new(node.id, 1))
53        {
54            let op = Self { with_index_outputs: None, ..self.clone() };
55            let mut patch = TypedModelPatch::default();
56            let mut wire = patch.tap_model(model, node.inputs[0])?;
57            wire = patch.wire_node(&node.name, op, &[wire])?[0];
58            patch.shunt_outside(model, node.id.into(), wire)?;
59            return Ok(Some(patch));
60        }
61        let fact = model.outlet_fact(node.inputs[0])?;
62        if let Some(pool_spec) = self.pool_spec.declutter(&fact.shape)? {
63            return Ok(Some(TypedModelPatch::replace_single_op(
64                model,
65                node,
66                &node.inputs,
67                Self { pool_spec, ..self.clone() },
68            )?));
69        }
70        Ok(None)
71    }
72
73    /// Lower to `OptMaxPool` with the geometry pre-resolved to `Concrete` when the
74    /// input shape is fixed, so the `Patch` is built once here rather than per eval.
75    /// Symbolic shapes are left as `MaxPool`.
76    fn codegen(
77        &self,
78        model: &TypedModel,
79        node: &TypedNode,
80    ) -> TractResult<Option<TypedModelPatch>> {
81        let fact = model.outlet_fact(node.inputs[0])?;
82        if fact.shape.as_concrete().is_none() {
83            return Ok(None);
84        }
85        let mut op = self.to_optimized(&fact.shape.to_tvec())?;
86        op.geometry = op.geometry.optimize_if(fact.shape.as_concrete())?;
87        Ok(Some(TypedModelPatch::replace_single_op(model, node, &node.inputs, op)?))
88    }
89
90    as_op!();
91}
92
93impl MaxPool {
94    fn to_optimized(&self, input_shape: &[TDim]) -> TractResult<OptMaxPool> {
95        Ok(OptMaxPool {
96            pool_spec: self.pool_spec.clone(),
97            with_index_outputs: self.with_index_outputs,
98            geometry: self.pool_spec.compute_geo(input_shape)?,
99        })
100    }
101}
102
103#[derive(Debug, Clone, new, Hash, PartialEq, Eq)]
104pub struct OptMaxPool {
105    pub pool_spec: PoolSpec,
106    pub with_index_outputs: Option<DatumType>,
107    pub geometry: PoolGeometry,
108}
109
110impl Op for OptMaxPool {
111    fn name(&self) -> StaticName {
112        "OptMaxPool".into()
113    }
114
115    fn info(&self) -> TractResult<Vec<String>> {
116        Ok(self.pool_spec.info())
117    }
118
119    op_as_typed_op!();
120}
121
122impl EvalOp for OptMaxPool {
123    fn is_stateless(&self) -> bool {
124        true
125    }
126
127    fn eval(&self, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
128        let input = args_1!(inputs);
129        let geo = self.geometry.to_concrete(input.shape())?;
130        dispatch_numbers!(Self::eval_t(input.datum_type())(self, &*input, geo.as_ref()))
131    }
132}
133
134impl TypedOp for OptMaxPool {
135    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
136        let mut facts = self.pool_spec.output_facts(inputs)?;
137        if let Some(idt) = self.with_index_outputs {
138            facts.push(facts[0].clone());
139            facts[1].datum_type = idt;
140        }
141        Ok(facts)
142    }
143
144    as_op!();
145}
146
147impl OptMaxPool {
148    fn eval_t<T: Datum + Copy + num_traits::Bounded + PartialOrd>(
149        &self,
150        input: &Tensor,
151        geo: &ConcretePoolGeometry,
152    ) -> TractResult<TVec<TValue>> {
153        let input_dt = input.datum_type();
154        let input_plain = input.try_as_plain()?;
155        let input: ArrayViewD<T> = input_plain.to_array_view()?;
156        let input_ptr = input.as_ptr();
157
158        let mut values = unsafe { ArrayD::<T>::uninit(&*geo.output_shape.shape).assume_init() };
159        let mut indices = if self.with_index_outputs.is_some() {
160            Some(unsafe { ArrayD::<i32>::uninit(&*geo.output_shape.shape).assume_init() })
161        } else {
162            None
163        };
164        let n = *geo.input_shape.n().unwrap_or(&1);
165        let n_stride_i = geo.input_shape.n_stride().unwrap_or(&0);
166        let n_stride_o = geo.output_shape.n_stride().unwrap_or(&0);
167        unsafe {
168            geo.patch.visit_output(|visitor| {
169                for n in 0..n {
170                    let input_offset = n * n_stride_i;
171                    let output_offset = n * n_stride_o;
172                    for c in 0..*geo.input_shape.c() {
173                        let input_offset = input_offset + geo.input_shape.c_stride() * c;
174                        let output_offset = output_offset + geo.output_shape.c_stride() * c;
175                        let max = visitor
176                            .valid_offsets()
177                            .map(|v| (v, *input_ptr.offset(v + input_offset as isize)))
178                            .fold((0, T::min_value()), |acc, v| if acc.1 < v.1 { v } else { acc });
179                        *values
180                            .as_mut_ptr()
181                            .offset(output_offset as isize + visitor.output_offset) = max.1;
182                        if let Some(ref mut indices) = indices {
183                            *indices
184                                .as_mut_ptr()
185                                .offset(output_offset as isize + visitor.output_offset) =
186                                max.0 as i32 / geo.patch.spec.output_inner_stride as i32;
187                        }
188                    }
189                }
190            });
191        }
192        let mut values = values.into_tensor();
193        unsafe {
194            values.set_datum_type(input_dt);
195        }
196        if let Some(dt) = self.with_index_outputs {
197            Ok(tvec!(
198                values.into_tvalue(),
199                indices.unwrap().into_tensor().cast_to_dt(dt)?.into_owned().into_tvalue()
200            ))
201        } else {
202            Ok(tvec!(values.into_tvalue()))
203        }
204    }
205}
206
207#[cfg(test)]
208mod tests {
209    use super::*;
210    use crate::ops::cnn::PaddingSpec;
211    use crate::ops::nn::DataFormat;
212
213    fn test_case() -> (TypedModel, TVec<TValue>) {
214        let mut model = TypedModel::default();
215        let source = model.add_source("data", f32::fact([1, 3, 8, 8])).unwrap();
216        let pool_spec = PoolSpec::new(
217            DataFormat::NCHW,
218            tvec![2, 2],
219            PaddingSpec::Valid,
220            None,
221            Some(tvec![2, 2]),
222            3,
223            3,
224        );
225        let op = MaxPool { pool_spec, with_index_outputs: None };
226        let out = model.wire_node("pool", op, &[source]).unwrap();
227        model.select_output_outlets(&out).unwrap();
228        let input = ndarray::Array4::from_shape_fn((1, 3, 8, 8), |(_, c, y, x)| {
229            (c * 64 + y * 8 + x) as f32
230        })
231        .into_tensor()
232        .into_tvalue();
233        (model, tvec!(input))
234    }
235
236    #[test]
237    fn optimized_maxpool_has_concrete_geometry() {
238        let (model, input) = test_case();
239        let plain = model.clone().into_runnable().unwrap().run(input.clone()).unwrap();
240
241        let optimized = model.into_optimized().unwrap();
242        let pool = optimized
243            .nodes
244            .iter()
245            .find_map(|n| n.op_as::<OptMaxPool>())
246            .expect("optimized model should contain an OptMaxPool");
247        assert!(
248            pool.geometry.is_concrete(),
249            "OptMaxPool geometry should be concrete after optimization"
250        );
251
252        let opt = optimized.into_runnable().unwrap().run(input).unwrap();
253        assert_eq!(*opt[0], *plain[0]);
254    }
255}
tract_core/ops/cnn/maxpool.rs

tract_core/ops/cnn/
maxpool.rs
