tract-core 0.23.9

use crate::internal::*;
use ndarray::prelude::*;

use crate::ops::cnn::pools::{ConcretePoolGeometry, PoolGeometry, PoolSpec};

#[derive(Debug, Clone, new, Hash, PartialEq, Eq)]
pub struct MaxPool {
    pub pool_spec: PoolSpec,
    pub with_index_outputs: Option<DatumType>,
}

impl Op for MaxPool {
    fn name(&self) -> StaticName {
        "MaxPool".into()
    }

    fn info(&self) -> TractResult<Vec<String>> {
        Ok(self.pool_spec.info())
    }

    op_as_typed_op!();
}

impl EvalOp for MaxPool {
    op_out_of_plan!();

    fn eval(&self, _ctx: &EvalContext, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
        let shape: TVec<TDim> = inputs[0].shape().iter().map(|d| d.to_dim()).collect();
        self.to_optimized(&shape)?.eval(_ctx, inputs)
    }
}

impl TypedOp for MaxPool {
    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
        let mut facts = self.pool_spec.output_facts(inputs)?;
        if let Some(idt) = self.with_index_outputs {
            facts.push(facts[0].clone());
            facts[1].datum_type = idt;
        }
        Ok(facts)
    }

    fn declutter(
        &self,
        model: &TypedModel,
        node: &TypedNode,
    ) -> TractResult<Option<TypedModelPatch>> {
        if self.with_index_outputs.is_some()
            && node.outputs[1].successors.len() == 0
            && !model.output_outlets()?.contains(&OutletId::new(node.id, 1))
        {
            let op = Self { with_index_outputs: None, ..self.clone() };
            let mut patch = TypedModelPatch::default();
            let mut wire = patch.tap_model(model, node.inputs[0])?;
            wire = patch.wire_node(&node.name, op, &[wire])?[0];
            patch.shunt_outside(model, node.id.into(), wire)?;
            return Ok(Some(patch));
        }
        let fact = model.outlet_fact(node.inputs[0])?;
        if let Some(pool_spec) = self.pool_spec.declutter(&fact.shape)? {
            return Ok(Some(TypedModelPatch::replace_single_op(
                model,
                node,
                &node.inputs,
                Self { pool_spec, ..self.clone() },
            )?));
        }
        Ok(None)
    }

    /// Lower to `OptMaxPool` with the geometry pre-resolved to `Concrete` when the
    /// input shape is fixed, so the `Patch` is built once here rather than per eval.
    /// Symbolic shapes are left as `MaxPool`.
    fn codegen(
        &self,
        model: &TypedModel,
        node: &TypedNode,
    ) -> TractResult<Option<TypedModelPatch>> {
        let fact = model.outlet_fact(node.inputs[0])?;
        if fact.shape.as_concrete().is_none() {
            return Ok(None);
        }
        let mut op = self.to_optimized(&fact.shape.to_tvec())?;
        op.geometry = op.geometry.optimize_if(fact.shape.as_concrete())?;
        Ok(Some(TypedModelPatch::replace_single_op(model, node, &node.inputs, op)?))
    }

    as_op!();
}

impl MaxPool {
    fn to_optimized(&self, input_shape: &[TDim]) -> TractResult<OptMaxPool> {
        Ok(OptMaxPool {
            pool_spec: self.pool_spec.clone(),
            with_index_outputs: self.with_index_outputs,
            geometry: self.pool_spec.compute_geo(input_shape)?,
        })
    }
}

#[derive(Debug, Clone, new, Hash, PartialEq, Eq)]
pub struct OptMaxPool {
    pub pool_spec: PoolSpec,
    pub with_index_outputs: Option<DatumType>,
    pub geometry: PoolGeometry,
}

impl Op for OptMaxPool {
    fn name(&self) -> StaticName {
        "OptMaxPool".into()
    }

    fn info(&self) -> TractResult<Vec<String>> {
        Ok(self.pool_spec.info())
    }

    op_as_typed_op!();
}

impl EvalOp for OptMaxPool {
    op_out_of_plan!();

    fn eval(&self, _ctx: &EvalContext, inputs: TVec<TValue>) -> TractResult<TVec<TValue>> {
        let input = args_1!(inputs);
        let geo = self.geometry.to_concrete(input.shape())?;
        dispatch_numbers!(Self::eval_t(input.datum_type())(self, &*input, geo.as_ref()))
    }
}

impl TypedOp for OptMaxPool {
    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
        let mut facts = self.pool_spec.output_facts(inputs)?;
        if let Some(idt) = self.with_index_outputs {
            facts.push(facts[0].clone());
            facts[1].datum_type = idt;
        }
        Ok(facts)
    }

    as_op!();
}

impl OptMaxPool {
    fn eval_t<T: Datum + Copy + num_traits::Bounded + PartialOrd>(
        &self,
        input: &Tensor,
        geo: &ConcretePoolGeometry,
    ) -> TractResult<TVec<TValue>> {
        let input_dt = input.datum_type();

        if self.with_index_outputs.is_none()
            && let Some(values) = self.try_nchw_2x2_f32::<T>(input, geo)?
        {
            return Ok(tvec!(values.into_tvalue()));
        }
        if self.with_index_outputs.is_none()
            && let Some(values) = self.try_row_spans::<T>(input, geo)?
        {
            return Ok(tvec!(values.into_tvalue()));
        }

        let input_plain = input.try_as_plain_ram()?;
        let input: ArrayViewD<T> = input_plain.to_array_view()?;
        let input_ptr = input.as_ptr();

        let mut values = unsafe { ArrayD::<T>::uninit(&*geo.output_shape.shape).assume_init() };
        let mut indices = if self.with_index_outputs.is_some() {
            Some(unsafe { ArrayD::<i32>::uninit(&*geo.output_shape.shape).assume_init() })
        } else {
            None
        };
        let n = *geo.input_shape.n().unwrap_or(&1);
        let n_stride_i = geo.input_shape.n_stride().unwrap_or(&0);
        let n_stride_o = geo.output_shape.n_stride().unwrap_or(&0);
        unsafe {
            geo.patch.visit_output(|visitor| {
                for n in 0..n {
                    let input_offset = n * n_stride_i;
                    let output_offset = n * n_stride_o;
                    for c in 0..*geo.input_shape.c() {
                        let input_offset = input_offset + geo.input_shape.c_stride() * c;
                        let output_offset = output_offset + geo.output_shape.c_stride() * c;
                        let max = visitor
                            .valid_offsets()
                            .map(|v| (v, *input_ptr.offset(v + input_offset as isize)))
                            .fold((0, T::min_value()), |acc, v| if acc.1 < v.1 { v } else { acc });
                        *values
                            .as_mut_ptr()
                            .offset(output_offset as isize + visitor.output_offset) = max.1;
                        if let Some(ref mut indices) = indices {
                            *indices
                                .as_mut_ptr()
                                .offset(output_offset as isize + visitor.output_offset) =
                                max.0 as i32 / geo.patch.spec.output_inner_stride as i32;
                        }
                    }
                }
            });
        }
        let mut values = values.into_tensor();
        unsafe {
            values.set_datum_type(input_dt);
        }
        if let Some(dt) = self.with_index_outputs {
            Ok(tvec!(
                values.into_tvalue(),
                indices.unwrap().into_tensor().cast_to_dt(dt)?.into_owned().into_tvalue()
            ))
        } else {
            Ok(tvec!(values.into_tvalue()))
        }
    }

    /// 2x2 f32 NCHW, walked (n, c, y, x) so W is contiguous and the stride-1
    /// interior vectorises along x. Every other pool keeps `visit_output`.
    fn try_nchw_2x2_f32<T: Datum>(
        &self,
        input: &Tensor,
        geo: &ConcretePoolGeometry,
    ) -> TractResult<Option<Tensor>> {
        let patch = &geo.patch;
        if T::datum_type() != f32::datum_type()
            || self.pool_spec.data_format != crate::ops::nn::DataFormat::NCHW
            || patch.rank() != 2
            || *geo.input_shape.w_stride() != 1
            || patch.spec.kernel_shape[..] != [2, 2]
            || patch.spec.dilations[..] != [1, 1]
        {
            return Ok(None);
        }
        let mut values =
            unsafe { Tensor::uninitialized_dt(input.datum_type(), &geo.output_shape.shape)? };
        unsafe {
            maxpool_2x2_f32(
                input.as_ptr::<f32>()? as usize,
                values.as_slice_mut_unchecked::<f32>(),
                geo,
            )?;
        }
        Ok(Some(values))
    }

    /// Any 2D pool, walked one output row at a time over images of contiguous pixels: a
    /// channel plane of 1-element pixels when channels come first, a whole image of
    /// C-element pixels when they come last. Each kernel tap folds into the row over the
    /// span of pixels where it lands in bounds, so a stride-1 tap reduces two contiguous
    /// slices. Taps fold in the generic walk's order with its comparison, so both pick the
    /// same values.
    fn try_row_spans<T: Datum + Copy + num_traits::Bounded + PartialOrd>(
        &self,
        input: &Tensor,
        geo: &ConcretePoolGeometry,
    ) -> TractResult<Option<Tensor>> {
        let patch = &geo.patch;
        let ish = &geo.input_shape;
        if patch.rank() != 2 {
            return Ok(None);
        }
        let (h, w, c) = (ish.hw_dims()[0], ish.hw_dims()[1], *ish.c());
        let px = if self.pool_spec.data_format.c_is_last() { c } else { 1 };
        let contiguous = if px == 1 {
            *ish.w_stride() == 1 && *ish.h_stride() == w && *ish.c_stride() == h * w
        } else {
            *ish.c_stride() == 1 && *ish.w_stride() == c && *ish.h_stride() == w * c
        };
        if !contiguous {
            return Ok(None);
        }
        let (oh, ow) = (patch.output_shape[0], patch.output_shape[1]);
        let spec = &patch.spec;
        let (kh, kw) = (spec.kernel_shape[0], spec.kernel_shape[1]);
        let (sh, sw) = (spec.strides[0], spec.strides[1]);
        let (dh, dw) = (spec.dilations[0], spec.dilations[1]);
        let (pt, pl) = (patch.pad_before[0] as isize, patch.pad_before[1] as isize);
        let mut output = Tensor::zero_dt(input.datum_type(), &geo.output_shape.shape)?;
        let mut output_plain = output.try_as_plain_ram_mut()?;
        let output_slice = output_plain.as_slice_mut::<T>()?;
        output_slice.fill(T::min_value());
        let input_plain = input.try_as_plain_ram()?;
        let input_slice = input_plain.as_slice::<T>()?;
        let column_spans: TVec<(usize, usize, usize)> = (0..kw)
            .filter_map(|kx| {
                let offset = (kx * dw) as isize - pl;
                let last = w as isize - 1 - offset;
                if last < 0 {
                    return None;
                }
                let lo = if offset >= 0 { 0 } else { (-offset as usize).div_ceil(sw) };
                let hi = (last as usize / sw + 1).min(ow);
                (lo < hi).then(|| (lo, hi, ((lo * sw) as isize + offset) as usize))
            })
            .collect();
        for (image, out_image) in
            input_slice.chunks_exact(h * w * px).zip(output_slice.chunks_exact_mut(oh * ow * px))
        {
            for (oy, row) in out_image.chunks_exact_mut(ow * px).enumerate() {
                for ky in 0..kh {
                    let iy = (oy * sh + ky * dh) as isize - pt;
                    if iy < 0 || iy >= h as isize {
                        continue;
                    }
                    let irow = &image[iy as usize * w * px..][..w * px];
                    for &(lo, hi, start) in &column_spans {
                        let (row, taps) = (&mut row[lo * px..hi * px], &irow[start * px..]);
                        if sw == 1 {
                            fold_max(row, taps.iter());
                        } else if px == 1 {
                            fold_max(row, taps.iter().step_by(sw));
                        } else {
                            for (pixel, taps) in row.chunks_exact_mut(px).zip(taps.chunks(sw * px))
                            {
                                fold_max(pixel, taps.iter());
                            }
                        }
                    }
                }
            }
        }
        Ok(Some(output))
    }
}

fn fold_max<'a, T: Copy + PartialOrd + 'a>(row: &mut [T], taps: impl Iterator<Item = &'a T>) {
    for (o, &v) in row.iter_mut().zip(taps) {
        *o = if *o < v { v } else { *o };
    }
}

unsafe fn maxpool_2x2_f32(
    i_base: usize,
    out: &mut [f32],
    geo: &ConcretePoolGeometry,
) -> TractResult<()> {
    let ish = &geo.input_shape;
    let osh = &geo.output_shape;
    let (h, w) = (ish.hw_dims()[0] as isize, ish.hw_dims()[1] as isize);
    let (oh, ow) = (geo.patch.output_shape[0], geo.patch.output_shape[1]);
    let sh = geo.patch.spec.strides[0] as isize;
    let sw = geo.patch.spec.strides[1] as isize;
    let pt = geo.patch.pad_before[0] as isize;
    let pl = geo.patch.pad_before[1] as isize;
    let ih_stride = *ish.h_stride() as isize;
    let oh_stride = *osh.h_stride() as isize;
    let n = *ish.n().unwrap_or(&1) as isize;
    let in_stride = *ish.n_stride().unwrap_or(&0) as isize;
    let on_stride = *osh.n_stride().unwrap_or(&0) as isize;
    let c = *ish.c() as isize;
    let ic_stride = *ish.c_stride() as isize;
    let oc_stride = *osh.c_stride() as isize;
    // Fully-valid 2×2 windows (both taps in-bounds). SameUpper 2×2 s=1
    // keeps H/W and pads after, so the last row/col are partial.
    let simd_s1 = sh == 1 && sw == 1;
    let y0 = pt.max(0) as usize;
    let y1 = ((h - 1 + pt).max(0) as usize).min(oh);
    let x0 = pl.max(0) as usize;
    let x1 = ((w - 1 + pl).max(0) as usize).min(ow);
    let plane_len = oh * ow;
    let one_plane = |p: isize, plane_out: *mut f32| unsafe {
        let iptr = i_base as *const f32;
        let nn = p / c;
        let cc = p % c;
        let in_base = nn * in_stride + cc * ic_stride;
        if simd_s1 && y1 > y0 && x1 > x0 {
            maxpool_2x2_s1_valid_f32(
                iptr.offset(in_base + (y0 as isize - pt) * ih_stride + (x0 as isize - pl)),
                plane_out.offset(y0 as isize * oh_stride + x0 as isize),
                y1 - y0,
                x1 - x0,
                ih_stride,
                oh_stride,
            );
        }
        for oy in 0..oh {
            let interior_y = simd_s1 && oy >= y0 && oy < y1;
            for ox in 0..ow {
                if interior_y && ox >= x0 && ox < x1 {
                    continue;
                }
                let mut m = f32::MIN;
                for ky in 0..2 {
                    let iy = oy as isize * sh + ky - pt;
                    if iy < 0 || iy >= h {
                        continue;
                    }
                    let row = iptr.offset(in_base + iy * ih_stride);
                    for kx in 0..2 {
                        let ix = ox as isize * sw + kx - pl;
                        if ix < 0 || ix >= w {
                            continue;
                        }
                        m = m.max(*row.offset(ix));
                    }
                }
                *plane_out.offset(oy as isize * oh_stride + ox as isize) = m;
            }
        }
    };
    if oc_stride == plane_len as isize && on_stride == c * oc_stride {
        tract_linalg::multithread::par_chunks_mut(out, plane_len, out.len(), |first, chunk| {
            for (i, plane_out) in chunk.chunks_mut(plane_len).enumerate() {
                one_plane(first as isize + i as isize, plane_out.as_mut_ptr());
            }
            Ok(())
        })
    } else {
        for p in 0..(n * c) {
            let out_base = (p / c) * on_stride + (p % c) * oc_stride;
            unsafe { one_plane(p, out.as_mut_ptr().offset(out_base)) };
        }
        Ok(())
    }
}

unsafe fn maxpool_2x2_s1_valid_f32(
    iptr: *const f32,
    optr: *mut f32,
    oh: usize,
    ow: usize,
    ih_stride: isize,
    oh_stride: isize,
) {
    unsafe {
        for oy in 0..oh {
            let row0 = iptr.offset(oy as isize * ih_stride);
            let row1 = iptr.offset((oy as isize + 1) * ih_stride);
            let dst = optr.offset(oy as isize * oh_stride);
            let mut ox = 0usize;
            #[cfg(target_arch = "aarch64")]
            {
                use std::arch::aarch64::*;
                while ox + 4 <= ow {
                    let a = vld1q_f32(row0.add(ox));
                    let b = vld1q_f32(row0.add(ox + 1));
                    let c = vld1q_f32(row1.add(ox));
                    let d = vld1q_f32(row1.add(ox + 1));
                    vst1q_f32(dst.add(ox), vmaxq_f32(vmaxq_f32(a, b), vmaxq_f32(c, d)));
                    ox += 4;
                }
            }
            #[cfg(target_arch = "x86_64")]
            {
                if is_x86_feature_detected!("avx") {
                    use std::arch::x86_64::*;
                    while ox + 8 <= ow {
                        let a = _mm256_loadu_ps(row0.add(ox));
                        let b = _mm256_loadu_ps(row0.add(ox + 1));
                        let c = _mm256_loadu_ps(row1.add(ox));
                        let d = _mm256_loadu_ps(row1.add(ox + 1));
                        _mm256_storeu_ps(
                            dst.add(ox),
                            _mm256_max_ps(_mm256_max_ps(a, b), _mm256_max_ps(c, d)),
                        );
                        ox += 8;
                    }
                }
            }
            while ox < ow {
                let m = (*row0.add(ox))
                    .max(*row0.add(ox + 1))
                    .max(*row1.add(ox))
                    .max(*row1.add(ox + 1));
                *dst.add(ox) = m;
                ox += 1;
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::ops::cnn::PaddingSpec;
    use crate::ops::nn::DataFormat;

    fn test_case() -> (TypedModel, TVec<TValue>) {
        let mut model = TypedModel::default();
        let source = model.add_source("data", f32::fact([1, 3, 8, 8])).unwrap();
        let pool_spec = PoolSpec::new(
            DataFormat::NCHW,
            tvec![2, 2],
            PaddingSpec::Valid,
            None,
            Some(tvec![2, 2]),
            3,
            3,
        );
        let op = MaxPool { pool_spec, with_index_outputs: None };
        let out = model.wire_node("pool", op, &[source]).unwrap();
        model.select_output_outlets(&out).unwrap();
        let input = ndarray::Array4::from_shape_fn((1, 3, 8, 8), |(_, c, y, x)| {
            (c * 64 + y * 8 + x) as f32
        })
        .into_tensor()
        .into_tvalue();
        (model, tvec!(input))
    }

    #[test]
    fn maxpool_2x2_multithreaded_matches_reference() {
        let cases = [
            ((1usize, 24, 320, 320), PaddingSpec::SameUpper, tvec![1usize, 1]),
            ((2, 8, 33, 17), PaddingSpec::Valid, tvec![2, 2]),
        ];
        for (shape, padding, strides) in cases {
            let c = shape.1;
            let pool_spec =
                PoolSpec::new(DataFormat::NCHW, tvec![2, 2], padding, None, Some(strides), c, c);
            let input = ndarray::Array4::from_shape_fn(shape, |(n, c, y, x)| {
                ((n * 31 + c * 7 + y * 13 + x * 3) % 97) as f32 * 0.25 - 12.0
            })
            .into_tensor()
            .into_tvalue();
            // The index-outputs variant never takes the 2x2 fast path, so its
            // output 0 is the visit_output reference.
            let reference =
                MaxPool { pool_spec: pool_spec.clone(), with_index_outputs: Some(DatumType::I32) }
                    .eval(&EvalContext::out_of_plan(), tvec!(input.clone()))
                    .unwrap();
            let op = MaxPool { pool_spec, with_index_outputs: None };
            let pool = tract_linalg::multithread::Executor::multithread(4);
            let fast = tract_linalg::multithread::multithread_tract_scope(pool, || {
                op.eval(&EvalContext::out_of_plan(), tvec!(input))
            })
            .unwrap();
            assert_eq!(*fast[0], *reference[0]);
        }
    }

    #[test]
    fn optimized_maxpool_has_concrete_geometry() {
        let (model, input) = test_case();
        let plain = model.clone().into_runnable().unwrap().run(input.clone()).unwrap();

        let optimized = model.into_optimized().unwrap();
        let pool = optimized
            .nodes
            .iter()
            .find_map(|n| n.op_as::<OptMaxPool>())
            .expect("optimized model should contain an OptMaxPool");
        assert!(
            pool.geometry.is_concrete(),
            "OptMaxPool geometry should be concrete after optimization"
        );

        let opt = optimized.into_runnable().unwrap().run(input).unwrap();
        assert_eq!(*opt[0], *plain[0]);
    }
}