metal: max and average pooling kernels

core/src/ops/cnn/mod.rs

@@ -11,12 +11,12 @@ 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};
 pub use self::pools::PoolSpec;
-pub use self::sumpool::SumPool;
+pub use self::sumpool::{OptSumPool, SumPool};
 
 use super::array::MultiBroadcastTo;
 

metal/src/kernels/mod.rs

@@ -31,7 +31,7 @@ const GGML: &str = include_str!("matmul/ggml_gemm/ggml_mm_mv.metal");
 const BASIC_MAT_MUL: &str = include_str!("matmul/basic/basic_mat_mul.metal");
 const ARRAY_OPS: &str = include_str!("array/array_ops.metal");
 const BIN_OPS: &str = include_str!("bin_ops.metal");
-const NN_OPS: &str = include_str!("nn/nn_ops.metal");
+const NN_OPS: &str = concat!(include_str!("nn/nn_ops.metal"), include_str!("nn/pool.metal"));
 const CONV_OPS: &str = include_str!("conv.metal");
 const ELEMENT_WISE_OPS: &str = include_str!("element_wise.metal");
 const FFT_OPS: &str = include_str!("fft.metal");

metal/src/kernels/nn/mod.rs

@@ -1,6 +1,7 @@
 pub mod apply_rope;
 pub mod gelu_approximate;
 pub mod leaky_relu;
+pub mod pool;
 pub mod reduce;
 pub mod rms_norm;
 pub mod scaled_masked_softmax;

metal/src/kernels/nn/pool.metal

@@ -0,0 +1,167 @@
+#include <metal_stdlib>
+using namespace metal;
+
+// 2D pooling over a channels-last tensor. One thread owns one (n, oh, ow, c),
+// so consecutive threads walk the contiguous channel axis and every window read
+// is coalesced.
+//
+// Buffer layout:
+//   0: input   [N, iH, iW, C]
+//   1: output  [N, oH, oW, C]
+//   2: params  (see PoolParams)
+struct PoolParams {
+    int n;
+    int ih;
+    int iw;
+    int c;
+    int oh;
+    int ow;
+    int kh;
+    int kw;
+    int stride_h;
+    int stride_w;
+    int pad_h;
+    int pad_w;
+    int dil_h;
+    int dil_w;
+    // Divide the sum by the window area including padding, rather than by the
+    // number of positions that actually landed inside the input.
+    int count_include_pad;
+    int normalize;
+};
+
+template <typename T>
+[[kernel]] void max_pool_2d(
+    const device T* input [[buffer(0)]],
+    device T* output [[buffer(1)]],
+    const constant PoolParams& p [[buffer(2)]],
+    uint3 gid [[thread_position_in_grid]]) {
+    const int c = int(gid.x);
+    const int ow = int(gid.y);
+    const int rest = int(gid.z);
+    const int oh = rest % p.oh;
+    const int n = rest / p.oh;
+    if (c >= p.c || ow >= p.ow || n >= p.n) {
+        return;
+    }
+
+    const int h_start = oh * p.stride_h - p.pad_h;
+    const int w_start = ow * p.stride_w - p.pad_w;
+
+    // An all-padding window has no value to take, and tract's CPU op leaves
+    // -inf there too.
+    T best = T(-INFINITY);
+    for (int kh = 0; kh < p.kh; ++kh) {
+        const int ih = h_start + kh * p.dil_h;
+        if (ih < 0 || ih >= p.ih) {
+            continue;
+        }
+        for (int kw = 0; kw < p.kw; ++kw) {
+            const int iw = w_start + kw * p.dil_w;
+            if (iw < 0 || iw >= p.iw) {
+                continue;
+            }
+            const int64_t idx =
+                ((int64_t(n) * p.ih + ih) * p.iw + iw) * p.c + c;
+            best = max(best, input[idx]);
+        }
+    }
+    const int64_t out_idx = ((int64_t(n) * p.oh + oh) * p.ow + ow) * p.c + c;
+    output[out_idx] = best;
+}
+
+template <typename T>
+[[kernel]] void sum_pool_2d(
+    const device T* input [[buffer(0)]],
+    device T* output [[buffer(1)]],
+    const constant PoolParams& p [[buffer(2)]],
+    uint3 gid [[thread_position_in_grid]]) {
+    const int c = int(gid.x);
+    const int ow = int(gid.y);
+    const int rest = int(gid.z);
+    const int oh = rest % p.oh;
+    const int n = rest / p.oh;
+    if (c >= p.c || ow >= p.ow || n >= p.n) {
+        return;
+    }
+
+    const int h_start = oh * p.stride_h - p.pad_h;
+    const int w_start = ow * p.stride_w - p.pad_w;
+
+    float acc = 0.0f;
+    int counted = 0;
+    for (int kh = 0; kh < p.kh; ++kh) {
+        const int ih = h_start + kh * p.dil_h;
+        if (ih < 0 || ih >= p.ih) {
+            continue;
+        }
+        for (int kw = 0; kw < p.kw; ++kw) {
+            const int iw = w_start + kw * p.dil_w;
+            if (iw < 0 || iw >= p.iw) {
+                continue;
+            }
+            const int64_t idx =
+                ((int64_t(n) * p.ih + ih) * p.iw + iw) * p.c + c;
+            acc += float(input[idx]);
+            counted += 1;
+        }
+    }
+    if (p.normalize) {
+        const int divisor = p.count_include_pad ? (p.kh * p.kw) : max(counted, 1);
+        acc /= float(divisor);
+    }
+    const int64_t out_idx = ((int64_t(n) * p.oh + oh) * p.ow + ow) * p.c + c;
+    output[out_idx] = T(acc);
+}
+
+// Channels-first variant: the contiguous axis is now width, so one thread owns
+// one (n, c, oh, ow) and consecutive threads walk the row.
+template <typename T>
+[[kernel]] void max_pool_2d_nchw(
+    const device T* input [[buffer(0)]],
+    device T* output [[buffer(1)]],
+    const constant PoolParams& p [[buffer(2)]],
+    uint3 gid [[thread_position_in_grid]]) {
+    const int ow = int(gid.x);
+    const int oh = int(gid.y);
+    const int rest = int(gid.z);
+    const int c = rest % p.c;
+    const int n = rest / p.c;
+    if (ow >= p.ow || oh >= p.oh || n >= p.n) {
+        return;
+    }
+
+    const int h_start = oh * p.stride_h - p.pad_h;
+    const int w_start = ow * p.stride_w - p.pad_w;
+    const int64_t plane = (int64_t(n) * p.c + c);
+
+    T best = T(-INFINITY);
+    for (int kh = 0; kh < p.kh; ++kh) {
+        const int ih = h_start + kh * p.dil_h;
+        if (ih < 0 || ih >= p.ih) {
+            continue;
+        }
+        for (int kw = 0; kw < p.kw; ++kw) {
+            const int iw = w_start + kw * p.dil_w;
+            if (iw < 0 || iw >= p.iw) {
+                continue;
+            }
+            best = max(best, input[(plane * p.ih + ih) * p.iw + iw]);
+        }
+    }
+    output[(plane * p.oh + oh) * p.ow + ow] = best;
+}
+
+#define instantiate_pool(name, tname, itype)                                  \
+    template [[host_name(#name "_" #tname)]] [[kernel]] void name<itype>(     \
+        const device itype* input [[buffer(0)]],                              \
+        device itype* output [[buffer(1)]],                                   \
+        const constant PoolParams& p [[buffer(2)]],                           \
+        uint3 gid [[thread_position_in_grid]]);
+
+instantiate_pool(max_pool_2d, f32, float)
+instantiate_pool(max_pool_2d, f16, half)
+instantiate_pool(sum_pool_2d, f32, float)
+instantiate_pool(sum_pool_2d, f16, half)
+instantiate_pool(max_pool_2d_nchw, f32, float)
+instantiate_pool(max_pool_2d_nchw, f16, half)