Run 2D MaxPool on the Metal GPU

core/src/ops/cnn/mod.rs

@@ -11,7 +11,7 @@ mod sumpool;
 
 pub use self::conv::{Conv, KernelFormat};
 pub use self::deconv::Deconv;
-pub use self::maxpool::MaxPool;
+pub use self::maxpool::{MaxPool, OptMaxPool};
 pub use self::padding::PaddingSpec;
 pub use self::patch_axis::PatchAxis;
 pub use self::patches::{Patch, PatchSpec};

gpu/src/ops/max_pool.rs

@@ -0,0 +1,107 @@
+use crate::tensor::{DeviceTensor, DeviceTensorExt};
+use derive_new::new;
+use tract_core::internal::*;
+
+/// Geometry of a two-spatial-axis max pooling, with every axis stride given
+/// explicitly so one kernel serves both NCHW and NHWC.
+#[derive(Clone, Debug, PartialEq, Eq, Hash)]
+pub struct MaxPool2dGeometry {
+    pub batch: usize,
+    pub channels: usize,
+    pub input_hw: (usize, usize),
+    pub output_hw: (usize, usize),
+    pub kernel: (usize, usize),
+    pub strides: (usize, usize),
+    pub dilations: (usize, usize),
+    pub padding: (usize, usize),
+    pub input_strides: [usize; 4],
+    pub output_strides: [usize; 4],
+}
+
+pub type DispatchMaxPool2dFn =
+    fn(input: &DeviceTensor, geo: &MaxPool2dGeometry, output: &DeviceTensor) -> TractResult<()>;
+
+/// Max pooling over two spatial axes, geometry resolved at translation time.
+/// Windows lying entirely in the padding take the datum type's lowest value,
+/// matching the CPU op.
+#[derive(Clone, new)]
+pub struct GpuMaxPool {
+    pub geometry: MaxPool2dGeometry,
+    pub output_shape: TVec<usize>,
+    pub backend_name: &'static str,
+    pub dispatch: DispatchMaxPool2dFn,
+}
+
+impl std::fmt::Debug for GpuMaxPool {
+    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+        write!(f, "{}MaxPool", self.backend_name)
+    }
+}
+
+impl PartialEq for GpuMaxPool {
+    fn eq(&self, other: &Self) -> bool {
+        self.backend_name == other.backend_name
+            && self.geometry == other.geometry
+            && self.output_shape == other.output_shape
+    }
+}
+impl Eq for GpuMaxPool {}
+
+impl std::hash::Hash for GpuMaxPool {
+    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+        self.backend_name.hash(state);
+        self.geometry.hash(state);
+        self.output_shape.hash(state);
+    }
+}
+
+impl Op for GpuMaxPool {
+    fn name(&self) -> StaticName {
+        format!("{}MaxPool", self.backend_name).into()
+    }
+    fn info(&self) -> TractResult<Vec<String>> {
+        Ok(vec![format!(
+            "kernel={:?} strides={:?} dilations={:?} padding={:?}",
+            self.geometry.kernel,
+            self.geometry.strides,
+            self.geometry.dilations,
+            self.geometry.padding
+        )])
+    }
+    op_as_typed_op!();
+}
+
+impl EvalOp for GpuMaxPool {
+    fn is_stateless(&self) -> bool {
+        true
+    }
+
+    fn eval_with_session(
+        &self,
+        node_id: usize,
+        session: &TurnState,
+        inputs: TVec<TValue>,
+    ) -> TractResult<TVec<TValue>> {
+        let input = inputs[0].to_device_tensor()?;
+        let output = crate::session_handler::make_tensor_for_node(
+            session,
+            node_id,
+            input.datum_type(),
+            &self.output_shape,
+        )?;
+        (self.dispatch)(input, &self.geometry, &output)?;
+        Ok(tvec!(output.into_tensor().into_tvalue()))
+    }
+}
+
+impl TypedOp for GpuMaxPool {
+    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
+        crate::utils::facts_to_device_facts(inputs, |facts| {
+            ensure!(facts.len() == 1);
+            let shape: TVec<TDim> = self.output_shape.iter().map(|d| d.to_dim()).collect();
+            Ok(tvec!(facts[0].datum_type.fact(&shape)))
+        })
+        .with_context(|| format!("Error while computing facts for {:?}", self.name()))
+    }
+    as_op!();
+}

gpu/src/ops/mod.rs

@@ -12,9 +12,11 @@ pub mod gather;
 pub mod gelu_approximate;
 pub mod iff;
 pub mod leaky_relu;
+pub mod max_pool;
 pub mod pad;
 pub mod pulse;
 pub mod reduce;
+pub mod resize;
 pub mod rms_norm;
 pub mod rotate_half;
 pub mod scaled_masked_softmax;

gpu/src/ops/resize.rs

@@ -0,0 +1,109 @@
+use crate::tensor::{DeviceTensor, DeviceTensorExt, IntoDevice};
+use derive_new::new;
+use tract_core::internal::*;
+
+/// Resamples one axis: `output[.., x, ..] = sum_k weights[x, k] * input[.., indices[x, k], ..]`,
+/// with `indices` already clamped into the axis by the host-built plan.
+pub type DispatchResizeAxisFn = fn(
+    input: &DeviceTensor,
+    axis: usize,
+    indices: &DeviceTensor,
+    weights: &DeviceTensor,
+    window: usize,
+    output: &DeviceTensor,
+) -> TractResult<()>;
+
+/// Resize against a plan baked at translation time, one dispatch per resampled
+/// axis. The plan makes the op independent of the interpolator: nearest, linear
+/// and cubic differ only in window size and weights, so translation is limited
+/// to nodes whose shapes and scales are known then. The scales/sizes input is
+/// kept for arity but no longer read.
+#[derive(Clone, new)]
+pub struct GpuResize {
+    pub axes: TVec<usize>,
+    pub windows: TVec<usize>,
+    pub plans: TVec<(Arc<Tensor>, Arc<Tensor>)>,
+    pub output_shape: TVec<usize>,
+    pub backend_name: &'static str,
+    pub dispatch: DispatchResizeAxisFn,
+}
+
+impl std::fmt::Debug for GpuResize {
+    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+        write!(f, "{}Resize", self.backend_name)
+    }
+}
+
+impl PartialEq for GpuResize {
+    fn eq(&self, other: &Self) -> bool {
+        self.backend_name == other.backend_name
+            && self.axes == other.axes
+            && self.windows == other.windows
+            && self.output_shape == other.output_shape
+    }
+}
+impl Eq for GpuResize {}
+
+impl std::hash::Hash for GpuResize {
+    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+        self.backend_name.hash(state);
+        self.axes.hash(state);
+        self.windows.hash(state);
+        self.output_shape.hash(state);
+    }
+}
+
+impl Op for GpuResize {
+    fn name(&self) -> StaticName {
+        format!("{}Resize", self.backend_name).into()
+    }
+    fn info(&self) -> TractResult<Vec<String>> {
+        Ok(vec![format!("axes={:?} windows={:?}", self.axes, self.windows)])
+    }
+    op_as_typed_op!();
+}
+
+impl EvalOp for GpuResize {
+    fn is_stateless(&self) -> bool {
+        true
+    }
+
+    fn eval_with_session(
+        &self,
+        node_id: usize,
+        session: &TurnState,
+        inputs: TVec<TValue>,
+    ) -> TractResult<TVec<TValue>> {
+        let data = inputs[0].to_device_tensor()?;
+        let dt = data.datum_type();
+        let mut shape: TVec<usize> = data.shape().into();
+        let mut current = data.clone();
+        for (step, (&axis, &window)) in self.axes.iter().zip(&self.windows).enumerate() {
+            let (indices, weights) = &self.plans[step];
+            let indices = indices.as_ref().clone().into_device()?;
+            let weights = weights.as_ref().clone().into_device()?;
+            shape[axis] = self.output_shape[axis];
+            let last = step + 1 == self.axes.len();
+            let output = if last {
+                crate::session_handler::make_tensor_for_node(session, node_id, dt, &shape)?
+            } else {
+                DeviceTensor::uninitialized_dt(dt, &shape)?
+            };
+            (self.dispatch)(&current, axis, &indices, &weights, window, &output)?;
+            current = output;
+        }
+        Ok(tvec!(current.into_tensor().into_tvalue()))
+    }
+}
+
+impl TypedOp for GpuResize {
+    fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult<TVec<TypedFact>> {
+        crate::utils::facts_to_device_facts(inputs, |facts| {
+            ensure!(facts.len() == 1);
+            let shape: TVec<TDim> = self.output_shape.iter().map(|d| d.to_dim()).collect();
+            Ok(tvec!(facts[0].datum_type.fact(&shape)))
+        })
+        .with_context(|| format!("Error while computing facts for {:?}", self.name()))
+    }
+    as_op!();
+}

metal/src/kernels/array/array_ops.metal

@@ -336,6 +336,51 @@ typedef decltype(gather<float>) gather_t;
     template [[host_name(                                                      \
         "array_ops::gather_" #tname)]] [[kernel]] gather_t gather<type>;
 
+// Resample one axis against a host-built plan: every output index owns a window
+// of `window` input indices (already clamped) and their weights, so nearest,
+// linear, cubic and their antialiased variants all run through this kernel.
+// The axes around the resampled one are flattened into `outer` and `inner`.
+//
+// params layout: [outer, len_in, len_out, inner, window]
+template <typename T>
+[[kernel]] void resize_axis(device const void *input_b [[buffer(0)]],
+                            device void *output_b [[buffer(1)]],
+                            device const int32_t *indices [[buffer(2)]],
+                            device const float *weights [[buffer(3)]],
+                            constant const int32_t *params [[buffer(4)]],
+                            uint3 tpig [[thread_position_in_grid]]) {
+    const int32_t len_in = params[1];
+    const int32_t len_out = params[2];
+    const int32_t inner = params[3];
+    const int32_t window = params[4];
+
+    const int32_t i = (int32_t)tpig.x;
+    const int32_t x = (int32_t)tpig.y;
+    const int32_t o = (int32_t)tpig.z;
+
+    if (i >= inner || x >= len_out)
+        return;
+
+    device const T *input = (device const T *)input_b;
+    device T *output = (device T *)output_b;
+
+    const int32_t in_base = o * len_in * inner + i;
+    float acc = 0.0f;
+    for (int32_t k = 0; k < window; ++k) {
+        const float w = weights[x * window + k];
+        if (w != 0.0f)
+            acc += w * (float)input[in_base + indices[x * window + k] * inner];
+    }
+    output[o * len_out * inner + x * inner + i] = (T)acc;
+}
+
+typedef decltype(resize_axis<float>) resize_axis_t;
+
+#define INSTANTIATE_RESIZE_AXIS(tname, type)                                   \
+    template [[host_name(                                                      \
+        "array_ops::resize_axis_" #tname)]] [[kernel]] resize_axis_t           \
+        resize_axis<type>;
+
 // Copy kernels: only u8/u16/u32/u64 (copy is type-size based)
 INSTANTIATE_COPY(u8, uint8_t)
 INSTANTIATE_COPY(u16, uint16_t)
@@ -379,3 +424,7 @@ INSTANTIATE_DIAG_GATHER(f16, half)
 // Axis Gather: f32 and f16 only (indices are int64).
 INSTANTIATE_GATHER(f32, float)
 INSTANTIATE_GATHER(f16, half)
+
+// Axis resample: f32 and f16 only (plan indices are int32, weights f32).
+INSTANTIATE_RESIZE_AXIS(f32, float)
+INSTANTIATE_RESIZE_AXIS(f16, half)

metal/src/kernels/array/mod.rs

@@ -3,6 +3,7 @@ mod copy;
 mod diag_gather;
 mod dispatch;
 mod gather;
+mod resize;
 mod rotate_half;
 
 pub use cast::Cast;
@@ -13,6 +14,9 @@ pub use diag_gather::metal_diag_gather_dispatch;
 pub use dispatch::metal_copy_nd_dispatch;
 pub use gather::Gather;
 pub use gather::metal_gather_dispatch;
+pub use resize::ResizeAxis;
+pub use resize::metal_resize;
+pub use resize::metal_resize_axis_dispatch;
 pub use rotate_half::RotateHalf;
 pub use rotate_half::metal_rotate_half_dispatch;
 
@@ -50,5 +54,11 @@ pub fn all_functions() -> Vec<String> {
             .flat_map(|dt| Gather.kernel_name(dt).into_iter()),
     );
 
+    functions.extend(
+        tract_gpu::tensor::DeviceTensor::SUPPORTED_DT
+            .into_iter()
+            .flat_map(|dt| ResizeAxis.kernel_name(dt).into_iter()),
+    );
+
     functions.into_iter().collect()
 }