metal: max and average pooling kernels - #2553
Conversation
The eval only stored a value inside `if let Some(div)`, which is None when normalize is false, so the output tensor was left untouched and the op returned zeros. NNEF's `box` fragment defaults normalize to false and the deserializer passes it straight through, so this is reachable from a model file. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
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 <noreply@anthropic.com>
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.
|
Added a second commit: channels-first max pooling. The kernels here only covered channels-last, so a channels-first model — what an ONNX or PyTorch export gives — kept every max pool on the CPU. The new kernel takes the width axis as the fastest-moving one so its reads stay coalesced the way the channels-last one's do. Sum pooling stays channels-last. On u2netp (rembg, ONNX/PyTorch so NCHW, 1x3x320x320, Output matches the CPU path to 1.8e-7 max abs. Two new unit tests cover channels-first max pooling in f32 and f16, valid and same-padding. This supersedes #2557, which I opened before noticing this PR existed; closing that one. |
1 participant
Stacked on #2552 — its commit is the first of the two here, because the tests
compare against the CPU op and one of them exercises the path that fix repairs.
Pooling has no Metal kernel, so a pooled model leaves the GPU at every pool.
On Inception v3 that is 14
DeviceSyncToHostnodes; profiling it after theconvolution work lands shows the pools themselves at 28% of runtime on the CPU
side and the syncs dominating the rest.
This adds 2D max and sum pooling over channels-last tensors — one thread per
(n, oh, ow, c), so consecutive threads walk the contiguous channel axis andevery window read is coalesced — and routes
MaxPool/SumPooland theiroptimized forms to them. NCHW, rank ≠ 4, non-float dtypes and max-pool with an
index output stay on the existing path.
After this, Inception v3 has 14 pools on the GPU and its sync count goes
14 → 1.
Numbers
Inception v3 (TF, NHWC, 299×299) on
--metal, M1 Pro:Worth being explicit about that first pair: on main this is close to neutral,
because the direct convolution kernel dominates everything and 16 ms of CPU
pooling is noise beside it. The gain only appears once convolution is off the
critical path, where it is 1.80×. If the convolution PRs are not wanted,
this one is not worth much on its own.
Output matches the CPU path (same predicted class, max difference 3e-7).
Validation
cargo test -p tract-metal --release: 83 passed, one failure —test_mfa_attention_causal_const_is_noop, pre-existing on main and unrelated(#2546). New tests cover max pooling with valid and same padding, average
pooling with and without
count_include_pad, un-normalized sum pooling, f16,and an end-to-end case asserting both pools land on the GPU and match CPU.
fmt and clippy clean.
One wart worth flagging for review: the translators live in
ops::pool, whichnothing else calls into, and the linker drops the module — and with it the
inventoryregistrations — unless something references it. There is alink_translators()call in the transform to hold it. If there is an idiom youprefer for that, say so.
M1 Pro only; my M4 was asleep.
🍍