metal: depthwise convolution via a ported MLX kernel (2-14x) - #2550
Open
czoli1976 wants to merge 2 commits into
Open
metal: depthwise convolution via a ported MLX kernel (2-14x)#2550czoli1976 wants to merge 2 commits into
czoli1976 wants to merge 2 commits into
Conversation
The Metal convolution was a direct kernel computing one output position per thread, which leaves most of the GPU idle: on this M1 Pro it runs a 56x56x64 -> 128 3x3 layer at about 20 GFLOP/s. Port mlx's tiled implicit-GEMM conv as owned .metal source and route NHWC f16/f32 single-group 2D convolutions to it, leaving every other shape on the direct kernel. The shared rewrite puts kernels in OIHW, which the direct kernel indexes, so a metal-local rule reorders eligible ones into the OHWI layout the ported kernel wants - from whichever layout the exporter used, and into a constant, since the metal transform does not declutter afterwards. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Depthwise convolutions were left on the direct kernel by the implicit-GEMM path, which is most of a mobile vision model: MobileNet v2 is 17 depthwise convolutions out of 18. Port mlx's depthwise kernel and route the shapes it covers - one channel per group in and out, kernel up to 7x7, stride up to 2, channels a multiple of 16. The OIHW kernel tract already hands over is laid out the way this kernel indexes it, so no reorder is needed. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Stacked on #2549 — its commit is the first of the two here, so this is only
reviewable once that lands.
#2549 routes regular convolutions to mlx's implicit-GEMM kernel but leaves
depthwise ones on the direct kernel, which is most of a mobile vision model:
MobileNet v2 is 17 depthwise convolutions out of 18. This ports mlx's depthwise
kernel and routes the shapes it covers — one channel per group in and out,
kernel up to 7×7, stride up to 2, channels a multiple of 16, dilation 1 — using
mlx's own gate.
No weight reorder is needed. Once the shared rule has put the kernel in
OIHW,a depthwise kernel is
[C, 1, kH, kW], which is the same bytes as the[C, kH, kW, 1]this kernel indexes, so depthwise is deliberately excluded fromthe
OHWImove that #2549 does for regular convolutions.Numbers
bench_depthwise(added here), M1 Pro, ms per dispatch, best of 10, SameUpper.MobileNet-shaped layers:
f16 tracks within a few percent (1.3×–9.5× over the same set).
End to end, MobileNet v2 (TF, NHWC, 224×224) on
--metal:1.34×, three interleaved runs each (14.8/14.8/15.1 against 11.1/11.8/12.1),
same predicted class as the CPU path. Worth being clear that the end-to-end gain
is much smaller than the kernel gain: MobileNet spends most of its time in the
pointwise 1×1 convolutions, which already lower to
MetalMlxGemm, not in thedepthwise layers this touches.
Validation
cargo test -p tract-metal --release: 93 passed, one failure —test_mfa_attention_causal_const_is_noop, pre-existing on main and unrelated(#2546). New tests cover 3×3 same, 3×3 stride 2, 5×5 valid, f16, a non-square
input, and an end-to-end case through
MetalTransform. fmt and clippy clean.Measured on an M1 Pro only — I have an M4 available but it was asleep when this
was finished, so unlike #2549 this one carries no second-device numbers yet. The
gate is mlx's own and the kernel is verbatim, but if you would rather see M4
numbers before it lands, say so and I will add them.
mlx (ml-explore/mlx) is MIT, Copyright (c) 2023-2025 Apple Inc.; attributed in
the source header, flattened from a pinned commit. The same flatten also brings
mlx's Winograd kernels into the file unused — I measured that path separately and
it does not beat the implicit-GEMM one, so nothing dispatches to it.
🍍