From 15ede74cba2c3dac0ea159dce6de03fb79a3abf5 Mon Sep 17 00:00:00 2001 From: czoli1976 Date: Sun, 2 Aug 2026 06:22:27 +0100 Subject: [PATCH 1/2] metal: max and average pooling kernels Pooling had no Metal kernel, so a pooled model bounced back to the host at every pool. On Inception v3 that is 14 device syncs and, once the convolutions are on the GPU, most of what is left: the pools alone were a fifth of the runtime on the CPU side. Add 2D max and sum pooling over channels-last tensors and route MaxPool/SumPool - and their optimized forms - to them. Anything else, including NCHW and rank other than 4, stays where it was. Co-Authored-By: Claude Opus 5 --- core/src/ops/cnn/mod.rs | 4 +- metal/src/kernels/mod.rs | 2 +- metal/src/kernels/nn/mod.rs | 1 + metal/src/kernels/nn/pool.metal | 127 ++++++++++++++ metal/src/kernels/nn/pool.rs | 296 ++++++++++++++++++++++++++++++++ metal/src/ops/mod.rs | 1 + metal/src/ops/pool.rs | 131 ++++++++++++++ metal/src/transform.rs | 5 + 8 files changed, 564 insertions(+), 3 deletions(-) create mode 100644 metal/src/kernels/nn/pool.metal create mode 100644 metal/src/kernels/nn/pool.rs create mode 100644 metal/src/ops/pool.rs diff --git a/core/src/ops/cnn/mod.rs b/core/src/ops/cnn/mod.rs index 540e975c03..0abc831a0a 100644 --- a/core/src/ops/cnn/mod.rs +++ b/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; diff --git a/metal/src/kernels/mod.rs b/metal/src/kernels/mod.rs index 036c030a64..8fce4884e9 100644 --- a/metal/src/kernels/mod.rs +++ b/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"); diff --git a/metal/src/kernels/nn/mod.rs b/metal/src/kernels/nn/mod.rs index 9ff4b13ea8..8f4ffd99c7 100644 --- a/metal/src/kernels/nn/mod.rs +++ b/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; diff --git a/metal/src/kernels/nn/pool.metal b/metal/src/kernels/nn/pool.metal new file mode 100644 index 0000000000..0a47c4cd61 --- /dev/null +++ b/metal/src/kernels/nn/pool.metal @@ -0,0 +1,127 @@ +#include +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 +[[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 +[[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); +} + +#define instantiate_pool(name, tname, itype) \ + template [[host_name(#name "_" #tname)]] [[kernel]] void name( \ + 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) diff --git a/metal/src/kernels/nn/pool.rs b/metal/src/kernels/nn/pool.rs new file mode 100644 index 0000000000..3834c1df3f --- /dev/null +++ b/metal/src/kernels/nn/pool.rs @@ -0,0 +1,296 @@ +//! 2D max and sum pooling on channels-last tensors (see pool.metal). +//! +//! Pooling had no Metal kernel, so a pooled model bounced back to the host at +//! every pool: on Inception v3 that was 14 device syncs and a fifth of the +//! runtime. + +use crate::encoder::EncoderExt; +use crate::{LibraryName, MetalStream}; +use anyhow::ensure; +use metal::MTLSize; +use tract_core::internal::*; +use tract_core::ops::cnn::PoolSpec; +use tract_gpu::tensor::DeviceTensor; + +/// Mirror of `PoolParams` in pool.metal — keep field order in sync. +#[repr(C)] +#[derive(Debug, Clone, Copy)] +struct PoolParams { + n: i32, + ih: i32, + iw: i32, + c: i32, + oh: i32, + ow: i32, + kh: i32, + kw: i32, + stride_h: i32, + stride_w: i32, + pad_h: i32, + pad_w: i32, + dil_h: i32, + dil_w: i32, + count_include_pad: i32, + normalize: i32, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum PoolKind { + Max, + /// `normalize` turns the sum into an average. + Sum { + count_include_pad: bool, + normalize: bool, + }, +} + +/// Whether the kernel covers this pooling: rank-4 channels-last f16/f32, with +/// static geometry and no index output. +pub fn metal_pool_supported(pool_spec: &PoolSpec, fact: &TypedFact) -> bool { + matches!(fact.datum_type, DatumType::F16 | DatumType::F32) + && pool_spec.data_format.c_is_last() + && fact.rank() == 4 + && pool_spec.kernel_shape.len() == 2 + && fact.shape.as_concrete().is_some() +} + +pub fn dispatch_metal_pool( + stream: &MetalStream, + pool_spec: &PoolSpec, + kind: PoolKind, + input: &DeviceTensor, + output: &DeviceTensor, +) -> TractResult<()> { + let dt = input.datum_type(); + let tname = match dt { + DatumType::F32 => "f32", + DatumType::F16 => "f16", + _ => bail!("Metal pool: F32/F16 only, got {dt:?}"), + }; + let in_shape = pool_spec.data_format.shape(input.shape())?; + let out_shape = pool_spec.data_format.shape(output.shape())?; + ensure!(in_shape.hw_rank() == 2, "Metal pool is 2D only"); + + let strides = pool_spec.strides(); + let dilations = pool_spec.dilations(); + let padding = pool_spec.computed_padding(in_shape.hw_dims()); + let (count_include_pad, normalize) = match kind { + PoolKind::Max => (false, false), + PoolKind::Sum { count_include_pad, normalize } => (count_include_pad, normalize), + }; + let params = PoolParams { + n: *in_shape.n().unwrap_or(&1) as i32, + ih: in_shape.hw_dims()[0] as i32, + iw: in_shape.hw_dims()[1] as i32, + c: *in_shape.c() as i32, + oh: out_shape.hw_dims()[0] as i32, + ow: out_shape.hw_dims()[1] as i32, + kh: pool_spec.kernel_shape[0] as i32, + kw: pool_spec.kernel_shape[1] as i32, + stride_h: strides[0] as i32, + stride_w: strides[1] as i32, + pad_h: padding[0].pad_before as i32, + pad_w: padding[1].pad_before as i32, + dil_h: dilations[0] as i32, + dil_w: dilations[1] as i32, + count_include_pad: count_include_pad as i32, + normalize: normalize as i32, + }; + + let base = match kind { + PoolKind::Max => "max_pool_2d", + PoolKind::Sum { .. } => "sum_pool_2d", + }; + let pipeline = stream.load_pipeline(LibraryName::NNOps, &format!("{base}_{tname}"))?; + + stream.retain_tensor(input); + stream.retain_tensor(output); + + let command_buffer = stream.command_buffer(); + command_buffer.encode(|encoder| { + encoder.set_compute_pipeline_state(&pipeline); + encoder.set_metal_tensor(0, input, metal::MTLResourceUsage::Read); + encoder.set_metal_tensor(1, output, metal::MTLResourceUsage::Write); + encoder.set_slice(2, std::slice::from_ref(¶ms)); + let group_w = 32u64.min(params.c as u64).max(1); + encoder.dispatch_threads( + MTLSize { + width: params.c as _, + height: params.ow as _, + depth: (params.oh * params.n) as _, + }, + MTLSize { width: group_w, height: 1, depth: 1 }, + ); + }); + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::MetalTransform; + use crate::utils::with_borrowed_metal_stream; + use tract_core::ops::cnn::{MaxPool, PaddingSpec, SumPool}; + use tract_core::ops::nn::DataFormat; + use tract_core::transform::ModelTransform; + use tract_gpu::tensor::IntoDevice; + + fn ramp(dt: DatumType, shape: &[usize]) -> TractResult { + let len: usize = shape.iter().product(); + let v: Vec = (0..len).map(|i| ((i * 17 % 61) as f32 - 30.0) / 8.0).collect(); + Ok(Tensor::from_shape(shape, &v)?.cast_to_dt(dt)?.into_owned()) + } + + fn spec(k: usize, stride: usize, padding: PaddingSpec, c: usize) -> PoolSpec { + PoolSpec::new( + DataFormat::NHWC, + tvec![k, k], + padding, + None, + Some(tvec![stride, stride]), + c, + c, + ) + } + + /// Run the CPU op and the Metal kernel over the same input. + fn check( + dt: DatumType, + shape: &[usize], + pool_spec: PoolSpec, + kind: PoolKind, + cpu: Box, + ) -> TractResult<()> { + let input = ramp(dt, shape)?; + let mut model = TypedModel::default(); + let i = model.add_source("i", dt.fact(shape))?; + let out = model.wire_node("pool", cpu, &[i])?; + model.select_output_outlets(&out)?; + let expected = model + .clone() + .into_optimized()? + .into_runnable()? + .run(tvec![input.clone().into_tvalue()])?[0] + .clone() + .into_tensor(); + let got = with_borrowed_metal_stream(|stream| { + let i = input.clone().into_device()?; + let o_shape = pool_spec.output_shape(&input.shape().to_vec())?; + let o = unsafe { DeviceTensor::uninitialized_dt(dt, &o_shape.shape)? }; + dispatch_metal_pool(stream, &pool_spec, kind, &i, &o)?; + stream.wait_until_completed()?; + Ok(o.to_host()?.into_tensor()) + })?; + expected + .close_enough(&got, Approximation::Approximate) + .with_context(|| format!("{kind:?} dt={dt:?} shape={shape:?}")) + } + + #[test] + fn max_pool_3x3_stride2() -> TractResult<()> { + let s = spec(3, 2, PaddingSpec::Valid, 32); + check( + DatumType::F32, + &[1, 16, 16, 32], + s.clone(), + PoolKind::Max, + Box::new(MaxPool::new(s, None)), + ) + } + + #[test] + fn max_pool_same_padding() -> TractResult<()> { + let s = spec(3, 1, PaddingSpec::SameUpper, 48); + check( + DatumType::F32, + &[1, 13, 11, 48], + s.clone(), + PoolKind::Max, + Box::new(MaxPool::new(s, None)), + ) + } + + #[test] + fn avg_pool_3x3_same() -> TractResult<()> { + let s = spec(3, 1, PaddingSpec::SameUpper, 64); + check( + DatumType::F32, + &[1, 12, 12, 64], + s.clone(), + PoolKind::Sum { count_include_pad: false, normalize: true }, + Box::new(SumPool::new(s, false, true)), + ) + } + + #[test] + fn avg_pool_count_include_pad() -> TractResult<()> { + let s = spec(3, 2, PaddingSpec::SameUpper, 16); + check( + DatumType::F32, + &[2, 9, 9, 16], + s.clone(), + PoolKind::Sum { count_include_pad: true, normalize: true }, + Box::new(SumPool::new(s, true, true)), + ) + } + + #[test] + fn sum_pool_no_normalize() -> TractResult<()> { + let s = spec(2, 2, PaddingSpec::Valid, 32); + check( + DatumType::F32, + &[1, 8, 8, 32], + s.clone(), + PoolKind::Sum { count_include_pad: false, normalize: false }, + Box::new(SumPool::new(s, false, false)), + ) + } + + #[test] + fn pool_f16() -> TractResult<()> { + let s = spec(3, 2, PaddingSpec::SameUpper, 96); + check( + DatumType::F16, + &[1, 14, 14, 96], + s.clone(), + PoolKind::Max, + Box::new(MaxPool::new(s, None)), + ) + } + + // The pools must actually land on the GPU once the transform has run. + #[test] + fn pools_route_through_metal_transform() -> TractResult<()> { + let dt = DatumType::F32; + let shape = [1usize, 14, 14, 32]; + let input = ramp(dt, &shape)?; + let mut model = TypedModel::default(); + let i = model.add_source("i", dt.fact(&shape))?; + let m = model.wire_node( + "max", + MaxPool::new(spec(3, 1, PaddingSpec::SameUpper, 32), None), + &[i], + )?[0]; + let a = model.wire_node( + "avg", + SumPool::new(spec(3, 1, PaddingSpec::SameUpper, 32), false, true), + &[m], + )?; + model.select_output_outlets(&a)?; + let cpu = model + .clone() + .into_optimized()? + .into_runnable()? + .run(tvec![input.clone().into_tvalue()])?; + let metal = MetalTransform::default().transform_into(model.into_optimized()?)?; + let n_pools = + metal.nodes().iter().filter(|n| n.op_is::()).count(); + assert_eq!(n_pools, 2, "both pools should be on the GPU"); + let got = metal.into_runnable()?.run(tvec![input.into_tvalue()])?; + cpu[0] + .clone() + .into_tensor() + .close_enough(&got[0].clone().into_tensor(), Approximation::Approximate)?; + Ok(()) + } +} diff --git a/metal/src/ops/mod.rs b/metal/src/ops/mod.rs index 4c0defe450..26306fe79e 100644 --- a/metal/src/ops/mod.rs +++ b/metal/src/ops/mod.rs @@ -1,6 +1,7 @@ pub mod conv; pub mod fused_axis_op; pub mod gemm; +pub mod pool; pub use fused_axis_op::MetalFusedAxisOp; pub use gemm::MetalGemm; diff --git a/metal/src/ops/pool.rs b/metal/src/ops/pool.rs new file mode 100644 index 0000000000..5cf06867f1 --- /dev/null +++ b/metal/src/ops/pool.rs @@ -0,0 +1,131 @@ +use crate::kernels::nn::pool::{PoolKind, dispatch_metal_pool, metal_pool_supported}; +use tract_core::internal::*; +use tract_core::ops::cnn::{MaxPool, OptMaxPool, OptSumPool, PoolSpec, SumPool}; +use tract_gpu::tensor::DeviceTensorExt; + +/// Metal device op for `OptMaxPool` / `OptSumPool` over a channels-last tensor. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct MetalPool { + pub pool_spec: PoolSpec, + pub kind: PoolKindOp, +} + +/// `PoolKind` without the borrowed geometry, so the op stays hashable. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum PoolKindOp { + Max, + Sum { count_include_pad: bool, normalize: bool }, +} + +impl From for PoolKind { + fn from(k: PoolKindOp) -> Self { + match k { + PoolKindOp::Max => PoolKind::Max, + PoolKindOp::Sum { count_include_pad, normalize } => { + PoolKind::Sum { count_include_pad, normalize } + } + } + } +} + +impl Op for MetalPool { + fn name(&self) -> StaticName { + match self.kind { + PoolKindOp::Max => "MetalMaxPool".into(), + PoolKindOp::Sum { .. } => "MetalSumPool".into(), + } + } + + fn info(&self) -> TractResult> { + Ok(self.pool_spec.info()) + } + + op_as_typed_op!(); +} + +impl EvalOp for MetalPool { + fn is_stateless(&self) -> bool { + true + } + + fn eval_with_session( + &self, + node_id: usize, + session: &TurnState, + inputs: TVec, + ) -> TractResult> { + let input = inputs[0].to_device_tensor()?; + let output_shape = self.pool_spec.output_shape(input.shape())?; + let output = tract_gpu::session_handler::make_tensor_for_node( + session, + node_id, + input.datum_type(), + &output_shape.shape, + )?; + if output.len() > 0 { + crate::with_metal_stream(|stream| { + dispatch_metal_pool(stream, &self.pool_spec, self.kind.into(), input, &output) + })?; + } + Ok(tvec!(output.into_tensor().into_tvalue())) + } +} + +impl TypedOp for MetalPool { + as_op!(); + + fn output_facts(&self, inputs: &[&TypedFact]) -> TractResult> { + tract_gpu::utils::facts_to_device_facts(inputs, |facts| { + let shape = self.pool_spec.output_shape(&facts[0].shape)?; + Ok(tvec!(facts[0].datum_type.fact(shape.shape))) + }) + .with_context(|| "Error while computing facts for MetalPool") + } +} + +crate::register_metal_op!(OptMaxPool, |source, node, op| { + let facts = source.node_input_facts(node.id)?; + // an index output would need a second buffer the kernel does not write + if op.with_index_outputs.is_some() || !metal_pool_supported(&op.pool_spec, facts[0]) { + return Ok(None); + } + Ok(Some(Box::new(MetalPool { pool_spec: op.pool_spec.clone(), kind: PoolKindOp::Max }) + as Box)) +}); + +crate::register_metal_op!(OptSumPool, |source, node, op| { + let facts = source.node_input_facts(node.id)?; + if !metal_pool_supported(&op.pool_spec, facts[0]) { + return Ok(None); + } + Ok(Some(Box::new(MetalPool { + pool_spec: op.pool_spec.clone(), + kind: PoolKindOp::Sum { count_include_pad: op.count_include_pad, normalize: op.normalize }, + }) as Box)) +}); + +// The metal transform runs before optimization, so a model still carries the +// unoptimized pools; the Opt* forms are registered too for callers that +// optimize first. +crate::register_metal_op!(MaxPool, |source, node, op| { + let facts = source.node_input_facts(node.id)?; + if op.with_index_outputs.is_some() || !metal_pool_supported(&op.pool_spec, facts[0]) { + return Ok(None); + } + Ok(Some(Box::new(MetalPool { pool_spec: op.pool_spec.clone(), kind: PoolKindOp::Max }) + as Box)) +}); + +crate::register_metal_op!(SumPool, |source, node, op| { + let facts = source.node_input_facts(node.id)?; + if !metal_pool_supported(&op.pool_spec, facts[0]) { + return Ok(None); + } + Ok(Some(Box::new(MetalPool { + pool_spec: op.pool_spec.clone(), + kind: PoolKindOp::Sum { count_include_pad: op.count_include_pad, normalize: op.normalize }, + }) as Box)) +}); + +/// Referenced from the transform so the registrations below survive linking. +pub fn link_translators() {} diff --git a/metal/src/transform.rs b/metal/src/transform.rs index c3633442be..4dfd82fb0e 100644 --- a/metal/src/transform.rs +++ b/metal/src/transform.rs @@ -146,6 +146,11 @@ impl ModelTransform for MetalTransform { } fn transform(&self, model: &mut TypedModel) -> TractResult<()> { + // The pool translators live in `ops::pool`, which nothing else calls + // into; without a reference the linker drops the module and with it the + // inventory registrations. + crate::ops::pool::link_translators(); + self.transform_up_to_phase(model, usize::MAX) } } From 6bdc35172bb71401ccbefc12d4b0e1073852ee89 Mon Sep 17 00:00:00 2001 From: czoli1976 Date: Sun, 2 Aug 2026 20:03:17 +0100 Subject: [PATCH 2/2] metal: pool channels-first tensors too MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The pooling kernels only covered channels-last tensors, so a channels-first model — what an ONNX or PyTorch export gives — kept every max pool on the CPU. A channels-first max pool now takes the width axis as the fastest-moving one, so its reads stay coalesced the way the channels-last kernel's do. Sum pooling stays channels-last. --- metal/src/kernels/nn/pool.metal | 40 ++++++++++++++++++ metal/src/kernels/nn/pool.rs | 73 ++++++++++++++++++++++++--------- metal/src/ops/pool.rs | 14 +++++-- 3 files changed, 103 insertions(+), 24 deletions(-) diff --git a/metal/src/kernels/nn/pool.metal b/metal/src/kernels/nn/pool.metal index 0a47c4cd61..daa8a128e7 100644 --- a/metal/src/kernels/nn/pool.metal +++ b/metal/src/kernels/nn/pool.metal @@ -114,6 +114,44 @@ template 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 +[[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( \ const device itype* input [[buffer(0)]], \ @@ -125,3 +163,5 @@ 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) diff --git a/metal/src/kernels/nn/pool.rs b/metal/src/kernels/nn/pool.rs index 3834c1df3f..c76336b937 100644 --- a/metal/src/kernels/nn/pool.rs +++ b/metal/src/kernels/nn/pool.rs @@ -44,11 +44,14 @@ pub enum PoolKind { }, } -/// Whether the kernel covers this pooling: rank-4 channels-last f16/f32, with -/// static geometry and no index output. -pub fn metal_pool_supported(pool_spec: &PoolSpec, fact: &TypedFact) -> bool { +/// Whether the kernel covers this pooling: rank-4 f16/f32 with static geometry +/// and no index output. Max pooling has a channels-first variant too; sum +/// pooling stays channels-last. +pub fn metal_pool_supported(pool_spec: &PoolSpec, kind: PoolKind, fact: &TypedFact) -> bool { + let layout_ok = pool_spec.data_format.c_is_last() + || (kind == PoolKind::Max && pool_spec.data_format.has_n()); matches!(fact.datum_type, DatumType::F16 | DatumType::F32) - && pool_spec.data_format.c_is_last() + && layout_ok && fact.rank() == 4 && pool_spec.kernel_shape.len() == 2 && fact.shape.as_concrete().is_some() @@ -97,10 +100,13 @@ pub fn dispatch_metal_pool( normalize: normalize as i32, }; + let channels_last = pool_spec.data_format.c_is_last(); let base = match kind { - PoolKind::Max => "max_pool_2d", + PoolKind::Max if channels_last => "max_pool_2d", + PoolKind::Max => "max_pool_2d_nchw", PoolKind::Sum { .. } => "sum_pool_2d", }; + ensure!(channels_last || kind == PoolKind::Max, "Metal sum pool is channels-last only"); let pipeline = stream.load_pipeline(LibraryName::NNOps, &format!("{base}_{tname}"))?; stream.retain_tensor(input); @@ -112,13 +118,14 @@ pub fn dispatch_metal_pool( encoder.set_metal_tensor(0, input, metal::MTLResourceUsage::Read); encoder.set_metal_tensor(1, output, metal::MTLResourceUsage::Write); encoder.set_slice(2, std::slice::from_ref(¶ms)); - let group_w = 32u64.min(params.c as u64).max(1); + let (fastest, height, depth) = if channels_last { + (params.c, params.ow, params.oh * params.n) + } else { + (params.ow, params.oh, params.c * params.n) + }; + let group_w = 32u64.min(fastest as u64).max(1); encoder.dispatch_threads( - MTLSize { - width: params.c as _, - height: params.ow as _, - depth: (params.oh * params.n) as _, - }, + MTLSize { width: fastest as _, height: height as _, depth: depth as _ }, MTLSize { width: group_w, height: 1, depth: 1 }, ); }); @@ -142,15 +149,17 @@ mod tests { } fn spec(k: usize, stride: usize, padding: PaddingSpec, c: usize) -> PoolSpec { - PoolSpec::new( - DataFormat::NHWC, - tvec![k, k], - padding, - None, - Some(tvec![stride, stride]), - c, - c, - ) + spec_fmt(DataFormat::NHWC, k, stride, padding, c) + } + + fn spec_fmt( + data_format: DataFormat, + k: usize, + stride: usize, + padding: PaddingSpec, + c: usize, + ) -> PoolSpec { + PoolSpec::new(data_format, tvec![k, k], padding, None, Some(tvec![stride, stride]), c, c) } /// Run the CPU op and the Metal kernel over the same input. @@ -198,6 +207,30 @@ mod tests { ) } + #[test] + fn max_pool_nchw() -> TractResult<()> { + let s = spec_fmt(DataFormat::NCHW, 2, 2, PaddingSpec::Valid, 24); + check( + DatumType::F32, + &[1, 24, 16, 16], + s.clone(), + PoolKind::Max, + Box::new(MaxPool::new(s, None)), + ) + } + + #[test] + fn max_pool_nchw_same_padding() -> TractResult<()> { + let s = spec_fmt(DataFormat::NCHW, 3, 2, PaddingSpec::SameUpper, 12); + check( + DatumType::F16, + &[2, 12, 11, 13], + s.clone(), + PoolKind::Max, + Box::new(MaxPool::new(s, None)), + ) + } + #[test] fn max_pool_same_padding() -> TractResult<()> { let s = spec(3, 1, PaddingSpec::SameUpper, 48); diff --git a/metal/src/ops/pool.rs b/metal/src/ops/pool.rs index 5cf06867f1..08f3c628ba 100644 --- a/metal/src/ops/pool.rs +++ b/metal/src/ops/pool.rs @@ -86,7 +86,9 @@ impl TypedOp for MetalPool { crate::register_metal_op!(OptMaxPool, |source, node, op| { let facts = source.node_input_facts(node.id)?; // an index output would need a second buffer the kernel does not write - if op.with_index_outputs.is_some() || !metal_pool_supported(&op.pool_spec, facts[0]) { + if op.with_index_outputs.is_some() + || !metal_pool_supported(&op.pool_spec, PoolKind::Max, facts[0]) + { return Ok(None); } Ok(Some(Box::new(MetalPool { pool_spec: op.pool_spec.clone(), kind: PoolKindOp::Max }) @@ -95,7 +97,8 @@ crate::register_metal_op!(OptMaxPool, |source, node, op| { crate::register_metal_op!(OptSumPool, |source, node, op| { let facts = source.node_input_facts(node.id)?; - if !metal_pool_supported(&op.pool_spec, facts[0]) { + let kind = PoolKind::Sum { count_include_pad: op.count_include_pad, normalize: op.normalize }; + if !metal_pool_supported(&op.pool_spec, kind, facts[0]) { return Ok(None); } Ok(Some(Box::new(MetalPool { @@ -109,7 +112,9 @@ crate::register_metal_op!(OptSumPool, |source, node, op| { // optimize first. crate::register_metal_op!(MaxPool, |source, node, op| { let facts = source.node_input_facts(node.id)?; - if op.with_index_outputs.is_some() || !metal_pool_supported(&op.pool_spec, facts[0]) { + if op.with_index_outputs.is_some() + || !metal_pool_supported(&op.pool_spec, PoolKind::Max, facts[0]) + { return Ok(None); } Ok(Some(Box::new(MetalPool { pool_spec: op.pool_spec.clone(), kind: PoolKindOp::Max }) @@ -118,7 +123,8 @@ crate::register_metal_op!(MaxPool, |source, node, op| { crate::register_metal_op!(SumPool, |source, node, op| { let facts = source.node_input_facts(node.id)?; - if !metal_pool_supported(&op.pool_spec, facts[0]) { + let kind = PoolKind::Sum { count_include_pad: op.count_include_pad, normalize: op.normalize }; + if !metal_pool_supported(&op.pool_spec, kind, facts[0]) { return Ok(None); } Ok(Some(Box::new(MetalPool {