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

tract_core/ops/cnn/
maxpool.rs
