metal: max and average pooling kernels - #2553
Conversation
|
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. |
|
🔴 Bench vs main — 7 speed regression(s) · Reference: 2026-08-08 morning nightly run (0d old) · full report → run Speed — evaltime · prefill · decode
+2 more regression(s)
|
| Δ | metric | device | main → PR |
|---|---|---|---|
| arm_ml_kws_cnn_m load · pass |
cortex-a9 |
84 ms → 120 ms | |
| arm_ml_kws_cnn_m load+optimize · pass |
cortex-a9 |
129 ms → 181 ms | |
| hey_snips_v1 load+optimize · 400ms |
cortex-a9 |
144 ms → 176 ms | |
| parakeet_tdt_600m_v3_f32f32_decoder_pass load+optimize · cpu |
apple-m1-max |
33 ms → 40 ms | |
| hey_snips_v1 load · 400ms |
cortex-a9 |
76 ms → 92 ms | |
| en_tdnn_15M_nnef load+optimize · pulse_240ms |
apple-m1-max |
43 ms → 51 ms | |
| en_tdnn_8M_nnef load+optimize · pulse_240ms |
apple-m1-max |
34 ms → 40 ms | |
| parakeet_tdt_600m_v3_f32f32_joint_pass load+optimize · cpu |
i9-11900kb_rtx-4060 |
41 ms → 48 ms | |
| hey_snips_v4_model17_nnef load · pulse8 |
apple-m1-max |
32 ms → 36 ms | |
| en_tdnn_15M_nnef load · pulse_240ms |
apple-m1-max |
31 ms → 34 ms | |
| arm_ml_kws_cnn_m load · pass |
cortex-a53 |
47 ms → 51 ms | |
| hey_snips_v4_model17_nnef load+optimize · pulse8 |
apple-m1-max |
50 ms → 54 ms | |
| en_tdnn_15M load+optimize · 2600ms |
apple-m1-max |
89 ms → 95 ms | |
| en_tdnn_15M RSS @ ready · 2600ms |
cortex-a55 |
110 MB → 118 MB | |
| voicecom_float load+optimize · 2sec |
cortex-a53 |
160 ms → 170 ms | |
| arm_ml_kws_cnn_m load+optimize · pass |
cortex-a53 |
73 ms → 77 ms | |
| voicecom_fake_quant load+optimize · 2sec |
cortex-a53 |
232 ms → 244 ms |
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.
|
|
7464c30 to
6bdc351
Compare
|
oups, sorry! |
|
🔴 Bench vs main — 1 speed regression(s) · Reference: 2026-08-09 morning nightly run (0d old) · full report → run Speed — evaltime · prefill · decode
|
| Δ | metric | device | main → PR |
|---|---|---|---|
| arm_ml_kws_cnn_m load · pass |
cortex-a9 |
84 ms → 124 ms | |
| arm_ml_kws_cnn_m load+optimize · pass |
cortex-a9 |
130 ms → 181 ms | |
| en_tdnn_15M RSS @ ready · 2600ms |
cortex-a55 |
111 MB → 118 MB |
2 participants
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.
🍍