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        if self.with_index_outputs.is_none()
157            && let Some(values) = self.try_row_spans::<T>(input, geo)?
158        {
159            return Ok(tvec!(values.into_tvalue()));
160        }
161
162        let input_plain = input.try_as_plain_ram()?;
163        let input: ArrayViewD<T> = input_plain.to_array_view()?;
164        let input_ptr = input.as_ptr();
165
166        let mut values = unsafe { ArrayD::<T>::uninit(&*geo.output_shape.shape).assume_init() };
167        let mut indices = if self.with_index_outputs.is_some() {
168            Some(unsafe { ArrayD::<i32>::uninit(&*geo.output_shape.shape).assume_init() })
169        } else {
170            None
171        };
172        let n = *geo.input_shape.n().unwrap_or(&1);
173        let n_stride_i = geo.input_shape.n_stride().unwrap_or(&0);
174        let n_stride_o = geo.output_shape.n_stride().unwrap_or(&0);
175        unsafe {
176            geo.patch.visit_output(|visitor| {
177                for n in 0..n {
178                    let input_offset = n * n_stride_i;
179                    let output_offset = n * n_stride_o;
180                    for c in 0..*geo.input_shape.c() {
181                        let input_offset = input_offset + geo.input_shape.c_stride() * c;
182                        let output_offset = output_offset + geo.output_shape.c_stride() * c;
183                        let max = visitor
184                            .valid_offsets()
185                            .map(|v| (v, *input_ptr.offset(v + input_offset as isize)))
186                            .fold((0, T::min_value()), |acc, v| if acc.1 < v.1 { v } else { acc });
187                        *values
188                            .as_mut_ptr()
189                            .offset(output_offset as isize + visitor.output_offset) = max.1;
190                        if let Some(ref mut indices) = indices {
191                            *indices
192                                .as_mut_ptr()
193                                .offset(output_offset as isize + visitor.output_offset) =
194                                max.0 as i32 / geo.patch.spec.output_inner_stride as i32;
195                        }
196                    }
197                }
198            });
199        }
200        let mut values = values.into_tensor();
201        unsafe {
202            values.set_datum_type(input_dt);
203        }
204        if let Some(dt) = self.with_index_outputs {
205            Ok(tvec!(
206                values.into_tvalue(),
207                indices.unwrap().into_tensor().cast_to_dt(dt)?.into_owned().into_tvalue()
208            ))
209        } else {
210            Ok(tvec!(values.into_tvalue()))
211        }
212    }
213
214    /// 2x2 f32 NCHW, walked (n, c, y, x) so W is contiguous and the stride-1
215    /// interior vectorises along x. Every other pool keeps `visit_output`.
216    fn try_nchw_2x2_f32<T: Datum>(
217        &self,
218        input: &Tensor,
219        geo: &ConcretePoolGeometry,
220    ) -> TractResult<Option<Tensor>> {
221        let patch = &geo.patch;
222        if T::datum_type() != f32::datum_type()
223            || self.pool_spec.data_format != crate::ops::nn::DataFormat::NCHW
224            || patch.rank() != 2
225            || *geo.input_shape.w_stride() != 1
226            || patch.spec.kernel_shape[..] != [2, 2]
227            || patch.spec.dilations[..] != [1, 1]
228        {
229            return Ok(None);
230        }
231        let mut values =
232            unsafe { Tensor::uninitialized_dt(input.datum_type(), &geo.output_shape.shape)? };
233        unsafe {
234            maxpool_2x2_f32(
235                input.as_ptr::<f32>()? as usize,
236                values.as_slice_mut_unchecked::<f32>(),
237                geo,
238            )?;
239        }
240        Ok(Some(values))
241    }
242
243    /// Any 2D pool, walked one output row at a time over images of contiguous pixels: a
244    /// channel plane of 1-element pixels when channels come first, a whole image of
245    /// C-element pixels when they come last. Each kernel tap folds into the row over the
246    /// span of pixels where it lands in bounds, so a stride-1 tap reduces two contiguous
247    /// slices. Taps fold in the generic walk's order with its comparison, so both pick the
248    /// same values.
249    fn try_row_spans<T: Datum + Copy + num_traits::Bounded + PartialOrd>(
250        &self,
251        input: &Tensor,
252        geo: &ConcretePoolGeometry,
253    ) -> TractResult<Option<Tensor>> {
254        let patch = &geo.patch;
255        let ish = &geo.input_shape;
256        if patch.rank() != 2 {
257            return Ok(None);
258        }
259        let (h, w, c) = (ish.hw_dims()[0], ish.hw_dims()[1], *ish.c());
260        let px = if self.pool_spec.data_format.c_is_last() { c } else { 1 };
261        let contiguous = if px == 1 {
262            *ish.w_stride() == 1 && *ish.h_stride() == w && *ish.c_stride() == h * w
263        } else {
264            *ish.c_stride() == 1 && *ish.w_stride() == c && *ish.h_stride() == w * c
265        };
266        if !contiguous {
267            return Ok(None);
268        }
269        let (oh, ow) = (patch.output_shape[0], patch.output_shape[1]);
270        let spec = &patch.spec;
271        let (kh, kw) = (spec.kernel_shape[0], spec.kernel_shape[1]);
272        let (sh, sw) = (spec.strides[0], spec.strides[1]);
273        let (dh, dw) = (spec.dilations[0], spec.dilations[1]);
274        let (pt, pl) = (patch.pad_before[0] as isize, patch.pad_before[1] as isize);
275        let mut output = Tensor::zero_dt(input.datum_type(), &geo.output_shape.shape)?;
276        let mut output_plain = output.try_as_plain_ram_mut()?;
277        let output_slice = output_plain.as_slice_mut::<T>()?;
278        output_slice.fill(T::min_value());
279        let input_plain = input.try_as_plain_ram()?;
280        let input_slice = input_plain.as_slice::<T>()?;
281        let column_spans: TVec<(usize, usize, usize)> = (0..kw)
282            .filter_map(|kx| {
283                let offset = (kx * dw) as isize - pl;
284                let last = w as isize - 1 - offset;
285                if last < 0 {
286                    return None;
287                }
288                let lo = if offset >= 0 { 0 } else { (-offset as usize).div_ceil(sw) };
289                let hi = (last as usize / sw + 1).min(ow);
290                (lo < hi).then(|| (lo, hi, ((lo * sw) as isize + offset) as usize))
291            })
292            .collect();
293        for (image, out_image) in
294            input_slice.chunks_exact(h * w * px).zip(output_slice.chunks_exact_mut(oh * ow * px))
295        {
296            for (oy, row) in out_image.chunks_exact_mut(ow * px).enumerate() {
297                for ky in 0..kh {
298                    let iy = (oy * sh + ky * dh) as isize - pt;
299                    if iy < 0 || iy >= h as isize {
300                        continue;
301                    }
302                    let irow = &image[iy as usize * w * px..][..w * px];
303                    for &(lo, hi, start) in &column_spans {
304                        let (row, taps) = (&mut row[lo * px..hi * px], &irow[start * px..]);
305                        if sw == 1 {
306                            fold_max(row, taps.iter());
307                        } else if px == 1 {
308                            fold_max(row, taps.iter().step_by(sw));
309                        } else {
310                            for (pixel, taps) in row.chunks_exact_mut(px).zip(taps.chunks(sw * px))
311                            {
312                                fold_max(pixel, taps.iter());
313                            }
314                        }
315                    }
316                }
317            }
318        }
319        Ok(Some(output))
320    }
321}
322
323fn fold_max<'a, T: Copy + PartialOrd + 'a>(row: &mut [T], taps: impl Iterator<Item = &'a T>) {
324    for (o, &v) in row.iter_mut().zip(taps) {
325        *o = if *o < v { v } else { *o };
326    }
327}
328
329unsafe fn maxpool_2x2_f32(
330    i_base: usize,
331    out: &mut [f32],
332    geo: &ConcretePoolGeometry,
333) -> TractResult<()> {
334    let ish = &geo.input_shape;
335    let osh = &geo.output_shape;
336    let (h, w) = (ish.hw_dims()[0] as isize, ish.hw_dims()[1] as isize);
337    let (oh, ow) = (geo.patch.output_shape[0], geo.patch.output_shape[1]);
338    let sh = geo.patch.spec.strides[0] as isize;
339    let sw = geo.patch.spec.strides[1] as isize;
340    let pt = geo.patch.pad_before[0] as isize;
341    let pl = geo.patch.pad_before[1] as isize;
342    let ih_stride = *ish.h_stride() as isize;
343    let oh_stride = *osh.h_stride() as isize;
344    let n = *ish.n().unwrap_or(&1) as isize;
345    let in_stride = *ish.n_stride().unwrap_or(&0) as isize;
346    let on_stride = *osh.n_stride().unwrap_or(&0) as isize;
347    let c = *ish.c() as isize;
348    let ic_stride = *ish.c_stride() as isize;
349    let oc_stride = *osh.c_stride() as isize;
350    // Fully-valid 2×2 windows (both taps in-bounds). SameUpper 2×2 s=1
351    // keeps H/W and pads after, so the last row/col are partial.
352    let simd_s1 = sh == 1 && sw == 1;
353    let y0 = pt.max(0) as usize;
354    let y1 = ((h - 1 + pt).max(0) as usize).min(oh);
355    let x0 = pl.max(0) as usize;
356    let x1 = ((w - 1 + pl).max(0) as usize).min(ow);
357    let plane_len = oh * ow;
358    let one_plane = |p: isize, plane_out: *mut f32| unsafe {
359        let iptr = i_base as *const f32;
360        let nn = p / c;
361        let cc = p % c;
362        let in_base = nn * in_stride + cc * ic_stride;
363        if simd_s1 && y1 > y0 && x1 > x0 {
364            maxpool_2x2_s1_valid_f32(
365                iptr.offset(in_base + (y0 as isize - pt) * ih_stride + (x0 as isize - pl)),
366                plane_out.offset(y0 as isize * oh_stride + x0 as isize),
367                y1 - y0,
368                x1 - x0,
369                ih_stride,
370                oh_stride,
371            );
372        }
373        for oy in 0..oh {
374            let interior_y = simd_s1 && oy >= y0 && oy < y1;
375            for ox in 0..ow {
376                if interior_y && ox >= x0 && ox < x1 {
377                    continue;
378                }
379                let mut m = f32::MIN;
380                for ky in 0..2 {
381                    let iy = oy as isize * sh + ky - pt;
382                    if iy < 0 || iy >= h {
383                        continue;
384                    }
385                    let row = iptr.offset(in_base + iy * ih_stride);
386                    for kx in 0..2 {
387                        let ix = ox as isize * sw + kx - pl;
388                        if ix < 0 || ix >= w {
389                            continue;
390                        }
391                        m = m.max(*row.offset(ix));
392                    }
393                }
394                *plane_out.offset(oy as isize * oh_stride + ox as isize) = m;
395            }
396        }
397    };
398    if oc_stride == plane_len as isize && on_stride == c * oc_stride {
399        tract_linalg::multithread::par_chunks_mut(out, plane_len, out.len(), |first, chunk| {
400            for (i, plane_out) in chunk.chunks_mut(plane_len).enumerate() {
401                one_plane(first as isize + i as isize, plane_out.as_mut_ptr());
402            }
403            Ok(())
404        })
405    } else {
406        for p in 0..(n * c) {
407            let out_base = (p / c) * on_stride + (p % c) * oc_stride;
408            unsafe { one_plane(p, out.as_mut_ptr().offset(out_base)) };
409        }
410        Ok(())
411    }
412}
413
414unsafe fn maxpool_2x2_s1_valid_f32(
415    iptr: *const f32,
416    optr: *mut f32,
417    oh: usize,
418    ow: usize,
419    ih_stride: isize,
420    oh_stride: isize,
421) {
422    unsafe {
423        for oy in 0..oh {
424            let row0 = iptr.offset(oy as isize * ih_stride);
425            let row1 = iptr.offset((oy as isize + 1) * ih_stride);
426            let dst = optr.offset(oy as isize * oh_stride);
427            let mut ox = 0usize;
428            #[cfg(target_arch = "aarch64")]
429            {
430                use std::arch::aarch64::*;
431                while ox + 4 <= ow {
432                    let a = vld1q_f32(row0.add(ox));
433                    let b = vld1q_f32(row0.add(ox + 1));
434                    let c = vld1q_f32(row1.add(ox));
435                    let d = vld1q_f32(row1.add(ox + 1));
436                    vst1q_f32(dst.add(ox), vmaxq_f32(vmaxq_f32(a, b), vmaxq_f32(c, d)));
437                    ox += 4;
438                }
439            }
440            #[cfg(target_arch = "x86_64")]
441            {
442                if is_x86_feature_detected!("avx") {
443                    use std::arch::x86_64::*;
444                    while ox + 8 <= ow {
445                        let a = _mm256_loadu_ps(row0.add(ox));
446                        let b = _mm256_loadu_ps(row0.add(ox + 1));
447                        let c = _mm256_loadu_ps(row1.add(ox));
448                        let d = _mm256_loadu_ps(row1.add(ox + 1));
449                        _mm256_storeu_ps(
450                            dst.add(ox),
451                            _mm256_max_ps(_mm256_max_ps(a, b), _mm256_max_ps(c, d)),
452                        );
453                        ox += 8;
454                    }
455                }
456            }
457            while ox < ow {
458                let m = (*row0.add(ox))
459                    .max(*row0.add(ox + 1))
460                    .max(*row1.add(ox))
461                    .max(*row1.add(ox + 1));
462                *dst.add(ox) = m;
463                ox += 1;
464            }
465        }
466    }
467}
468
469#[cfg(test)]
470mod tests {
471    use super::*;
472    use crate::ops::cnn::PaddingSpec;
473    use crate::ops::nn::DataFormat;
474
475    fn test_case() -> (TypedModel, TVec<TValue>) {
476        let mut model = TypedModel::default();
477        let source = model.add_source("data", f32::fact([1, 3, 8, 8])).unwrap();
478        let pool_spec = PoolSpec::new(
479            DataFormat::NCHW,
480            tvec![2, 2],
481            PaddingSpec::Valid,
482            None,
483            Some(tvec![2, 2]),
484            3,
485            3,
486        );
487        let op = MaxPool { pool_spec, with_index_outputs: None };
488        let out = model.wire_node("pool", op, &[source]).unwrap();
489        model.select_output_outlets(&out).unwrap();
490        let input = ndarray::Array4::from_shape_fn((1, 3, 8, 8), |(_, c, y, x)| {
491            (c * 64 + y * 8 + x) as f32
492        })
493        .into_tensor()
494        .into_tvalue();
495        (model, tvec!(input))
496    }
497
498    #[test]
499    fn maxpool_2x2_multithreaded_matches_reference() {
500        let cases = [
501            ((1usize, 24, 320, 320), PaddingSpec::SameUpper, tvec![1usize, 1]),
502            ((2, 8, 33, 17), PaddingSpec::Valid, tvec![2, 2]),
503        ];
504        for (shape, padding, strides) in cases {
505            let c = shape.1;
506            let pool_spec =
507                PoolSpec::new(DataFormat::NCHW, tvec![2, 2], padding, None, Some(strides), c, c);
508            let input = ndarray::Array4::from_shape_fn(shape, |(n, c, y, x)| {
509                ((n * 31 + c * 7 + y * 13 + x * 3) % 97) as f32 * 0.25 - 12.0
510            })
511            .into_tensor()
512            .into_tvalue();
513            // The index-outputs variant never takes the 2x2 fast path, so its
514            // output 0 is the visit_output reference.
515            let reference =
516                MaxPool { pool_spec: pool_spec.clone(), with_index_outputs: Some(DatumType::I32) }
517                    .eval(&EvalContext::out_of_plan(), tvec!(input.clone()))
518                    .unwrap();
519            let op = MaxPool { pool_spec, with_index_outputs: None };
520            let pool = tract_linalg::multithread::Executor::multithread(4);
521            let fast = tract_linalg::multithread::multithread_tract_scope(pool, || {
522                op.eval(&EvalContext::out_of_plan(), tvec!(input))
523            })
524            .unwrap();
525            assert_eq!(*fast[0], *reference[0]);
526        }
527    }
528
529    #[test]
530    fn optimized_maxpool_has_concrete_geometry() {
531        let (model, input) = test_case();
532        let plain = model.clone().into_runnable().unwrap().run(input.clone()).unwrap();
533
534        let optimized = model.into_optimized().unwrap();
535        let pool = optimized
536            .nodes
537            .iter()
538            .find_map(|n| n.op_as::<OptMaxPool>())
539            .expect("optimized model should contain an OptMaxPool");
540        assert!(
541            pool.geometry.is_concrete(),
542            "OptMaxPool geometry should be concrete after optimization"
543        );
544
545        let opt = optimized.into_runnable().unwrap().run(input).unwrap();
546        assert_eq!(*opt[0], *plain[0]);
547    }
548}
tract_core/ops/cnn/maxpool.rs

tract_core/ops/cnn/
maxpool.rs
