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
151        if self.with_index_outputs.is_none()
152            && let Some(values) = self.try_nchw_2x2_f32::<T>(input, geo)?
153        {
154            return Ok(tvec!(values.into_tvalue()));
155        }
156
157        let input_plain = input.try_as_plain_ram()?;
158        let input: ArrayViewD<T> = input_plain.to_array_view()?;
159        let input_ptr = input.as_ptr();
160
161        let mut values = unsafe { ArrayD::<T>::uninit(&*geo.output_shape.shape).assume_init() };
162        let mut indices = if self.with_index_outputs.is_some() {
163            Some(unsafe { ArrayD::<i32>::uninit(&*geo.output_shape.shape).assume_init() })
164        } else {
165            None
166        };
167        let n = *geo.input_shape.n().unwrap_or(&1);
168        let n_stride_i = geo.input_shape.n_stride().unwrap_or(&0);
169        let n_stride_o = geo.output_shape.n_stride().unwrap_or(&0);
170        unsafe {
171            geo.patch.visit_output(|visitor| {
172                for n in 0..n {
173                    let input_offset = n * n_stride_i;
174                    let output_offset = n * n_stride_o;
175                    for c in 0..*geo.input_shape.c() {
176                        let input_offset = input_offset + geo.input_shape.c_stride() * c;
177                        let output_offset = output_offset + geo.output_shape.c_stride() * c;
178                        let max = visitor
179                            .valid_offsets()
180                            .map(|v| (v, *input_ptr.offset(v + input_offset as isize)))
181                            .fold((0, T::min_value()), |acc, v| if acc.1 < v.1 { v } else { acc });
182                        *values
183                            .as_mut_ptr()
184                            .offset(output_offset as isize + visitor.output_offset) = max.1;
185                        if let Some(ref mut indices) = indices {
186                            *indices
187                                .as_mut_ptr()
188                                .offset(output_offset as isize + visitor.output_offset) =
189                                max.0 as i32 / geo.patch.spec.output_inner_stride as i32;
190                        }
191                    }
192                }
193            });
194        }
195        let mut values = values.into_tensor();
196        unsafe {
197            values.set_datum_type(input_dt);
198        }
199        if let Some(dt) = self.with_index_outputs {
200            Ok(tvec!(
201                values.into_tvalue(),
202                indices.unwrap().into_tensor().cast_to_dt(dt)?.into_owned().into_tvalue()
203            ))
204        } else {
205            Ok(tvec!(values.into_tvalue()))
206        }
207    }
208
209    /// 2x2 f32 NCHW, walked (n, c, y, x) so W is contiguous and the stride-1
210    /// interior vectorises along x. Every other pool keeps `visit_output`.
211    fn try_nchw_2x2_f32<T: Datum>(
212        &self,
213        input: &Tensor,
214        geo: &ConcretePoolGeometry,
215    ) -> TractResult<Option<Tensor>> {
216        let patch = &geo.patch;
217        if T::datum_type() != f32::datum_type()
218            || self.pool_spec.data_format != crate::ops::nn::DataFormat::NCHW
219            || patch.rank() != 2
220            || *geo.input_shape.w_stride() != 1
221            || patch.spec.kernel_shape[..] != [2, 2]
222            || patch.spec.dilations[..] != [1, 1]
223        {
224            return Ok(None);
225        }
226        let mut values =
227            unsafe { Tensor::uninitialized_dt(input.datum_type(), &geo.output_shape.shape)? };
228        unsafe {
229            maxpool_2x2_f32(input.as_ptr::<f32>()?, values.as_ptr_mut::<f32>()?, geo);
230        }
231        Ok(Some(values))
232    }
233}
234
235unsafe fn maxpool_2x2_f32(iptr: *const f32, optr: *mut f32, geo: &ConcretePoolGeometry) {
236    unsafe {
237        let ish = &geo.input_shape;
238        let osh = &geo.output_shape;
239        let (h, w) = (ish.hw_dims()[0] as isize, ish.hw_dims()[1] as isize);
240        let (oh, ow) = (geo.patch.output_shape[0], geo.patch.output_shape[1]);
241        let sh = geo.patch.spec.strides[0] as isize;
242        let sw = geo.patch.spec.strides[1] as isize;
243        let pt = geo.patch.pad_before[0] as isize;
244        let pl = geo.patch.pad_before[1] as isize;
245        let ih_stride = *ish.h_stride() as isize;
246        let oh_stride = *osh.h_stride() as isize;
247        let n = *ish.n().unwrap_or(&1) as isize;
248        let in_stride = *ish.n_stride().unwrap_or(&0) as isize;
249        let on_stride = *osh.n_stride().unwrap_or(&0) as isize;
250        let c = *ish.c() as isize;
251        let ic_stride = *ish.c_stride() as isize;
252        let oc_stride = *osh.c_stride() as isize;
253        // Fully-valid 2×2 windows (both taps in-bounds). SameUpper 2×2 s=1
254        // keeps H/W and pads after, so the last row/col are partial.
255        let simd_s1 = sh == 1 && sw == 1;
256        let y0 = pt.max(0) as usize;
257        let y1 = ((h - 1 + pt).max(0) as usize).min(oh);
258        let x0 = pl.max(0) as usize;
259        let x1 = ((w - 1 + pl).max(0) as usize).min(ow);
260        for nn in 0..n {
261            for cc in 0..c {
262                let in_base = nn * in_stride + cc * ic_stride;
263                let out_base = nn * on_stride + cc * oc_stride;
264                if simd_s1 && y1 > y0 && x1 > x0 {
265                    maxpool_2x2_s1_valid_f32(
266                        iptr.offset(in_base + (y0 as isize - pt) * ih_stride + (x0 as isize - pl)),
267                        optr.offset(out_base + y0 as isize * oh_stride + x0 as isize),
268                        y1 - y0,
269                        x1 - x0,
270                        ih_stride,
271                        oh_stride,
272                    );
273                }
274                for oy in 0..oh {
275                    let interior_y = simd_s1 && oy >= y0 && oy < y1;
276                    for ox in 0..ow {
277                        if interior_y && ox >= x0 && ox < x1 {
278                            continue;
279                        }
280                        let mut m = f32::MIN;
281                        for ky in 0..2 {
282                            let iy = oy as isize * sh + ky - pt;
283                            if iy < 0 || iy >= h {
284                                continue;
285                            }
286                            let row = iptr.offset(in_base + iy * ih_stride);
287                            for kx in 0..2 {
288                                let ix = ox as isize * sw + kx - pl;
289                                if ix < 0 || ix >= w {
290                                    continue;
291                                }
292                                m = m.max(*row.offset(ix));
293                            }
294                        }
295                        *optr.offset(out_base + oy as isize * oh_stride + ox as isize) = m;
296                    }
297                }
298            }
299        }
300    }
301}
302
303unsafe fn maxpool_2x2_s1_valid_f32(
304    iptr: *const f32,
305    optr: *mut f32,
306    oh: usize,
307    ow: usize,
308    ih_stride: isize,
309    oh_stride: isize,
310) {
311    unsafe {
312        for oy in 0..oh {
313            let row0 = iptr.offset(oy as isize * ih_stride);
314            let row1 = iptr.offset((oy as isize + 1) * ih_stride);
315            let dst = optr.offset(oy as isize * oh_stride);
316            let mut ox = 0usize;
317            #[cfg(target_arch = "aarch64")]
318            {
319                use std::arch::aarch64::*;
320                while ox + 4 <= ow {
321                    let a = vld1q_f32(row0.add(ox));
322                    let b = vld1q_f32(row0.add(ox + 1));
323                    let c = vld1q_f32(row1.add(ox));
324                    let d = vld1q_f32(row1.add(ox + 1));
325                    vst1q_f32(dst.add(ox), vmaxq_f32(vmaxq_f32(a, b), vmaxq_f32(c, d)));
326                    ox += 4;
327                }
328            }
329            #[cfg(target_arch = "x86_64")]
330            {
331                if is_x86_feature_detected!("avx") {
332                    use std::arch::x86_64::*;
333                    while ox + 8 <= ow {
334                        let a = _mm256_loadu_ps(row0.add(ox));
335                        let b = _mm256_loadu_ps(row0.add(ox + 1));
336                        let c = _mm256_loadu_ps(row1.add(ox));
337                        let d = _mm256_loadu_ps(row1.add(ox + 1));
338                        _mm256_storeu_ps(
339                            dst.add(ox),
340                            _mm256_max_ps(_mm256_max_ps(a, b), _mm256_max_ps(c, d)),
341                        );
342                        ox += 8;
343                    }
344                }
345            }
346            while ox < ow {
347                let m = (*row0.add(ox))
348                    .max(*row0.add(ox + 1))
349                    .max(*row1.add(ox))
350                    .max(*row1.add(ox + 1));
351                *dst.add(ox) = m;
352                ox += 1;
353            }
354        }
355    }
356}
357
358#[cfg(test)]
359mod tests {
360    use super::*;
361    use crate::ops::cnn::PaddingSpec;
362    use crate::ops::nn::DataFormat;
363
364    fn test_case() -> (TypedModel, TVec<TValue>) {
365        let mut model = TypedModel::default();
366        let source = model.add_source("data", f32::fact([1, 3, 8, 8])).unwrap();
367        let pool_spec = PoolSpec::new(
368            DataFormat::NCHW,
369            tvec![2, 2],
370            PaddingSpec::Valid,
371            None,
372            Some(tvec![2, 2]),
373            3,
374            3,
375        );
376        let op = MaxPool { pool_spec, with_index_outputs: None };
377        let out = model.wire_node("pool", op, &[source]).unwrap();
378        model.select_output_outlets(&out).unwrap();
379        let input = ndarray::Array4::from_shape_fn((1, 3, 8, 8), |(_, c, y, x)| {
380            (c * 64 + y * 8 + x) as f32
381        })
382        .into_tensor()
383        .into_tvalue();
384        (model, tvec!(input))
385    }
386
387    #[test]
388    fn optimized_maxpool_has_concrete_geometry() {
389        let (model, input) = test_case();
390        let plain = model.clone().into_runnable().unwrap().run(input.clone()).unwrap();
391
392        let optimized = model.into_optimized().unwrap();
393        let pool = optimized
394            .nodes
395            .iter()
396            .find_map(|n| n.op_as::<OptMaxPool>())
397            .expect("optimized model should contain an OptMaxPool");
398        assert!(
399            pool.geometry.is_concrete(),
400            "OptMaxPool geometry should be concrete after optimization"
401        );
402
403        let opt = optimized.into_runnable().unwrap().run(input).unwrap();
404        assert_eq!(*opt[0], *plain[0]);
405    }
406}
tract_core/ops/cnn/maxpool.rs

tract_core/ops/cnn/
maxpool.rs
