metal: mlx's preferred implicit-GEMM conv kernel, plus grouped convolutions (1.2-75x)

[czoli1976]

Stacked on #2549 and #2550 — their commits are the first two here.

#2549 ported mlx's general implicit-GEMM convolution. That is mlx's fallback:
its own routing prefers a second kernel whenever the channel counts are aligned
(C ≤ 4 or C%16 == 0, and O ≤ 16 or O%16 == 0), which is most real shapes.
This ports that one too and routes the way mlx does, so the general kernel now
only takes what the specialised one declines.

The specialised kernel drops the per-output-position jump tables and is
instantiated on small channel counts (1–4) and short filters, which is where the
difference mostly comes from.

Numbers

bench_conv_specialized (added here), M1 Pro, ms per dispatch, best of 10.
Both kernels are already checked against the CPU op by the existing tests.

shape	general	specialised
299×299×3 → 32, 3×3 s2	0.950	0.213	4.46×
147×147×32 → 64, 3×3	1.564	0.618	2.53×
73×73×80 → 192, 3×3	1.229	0.601	2.05×
112×112×32 → 64, 3×3	0.384	0.218	1.76×
8×8×384 → 384, 3×3	0.367	0.255	1.44×
35×35×64 → 96, 3×3	0.154	0.128	1.20×

f16 tracks within a few percent. Every shape measured improves; the first-layer
case gains most, since three input channels hit the channel specialisation.

End to end on Inception v3 (TF, NHWC, 299×299), --metal, M1 Pro:

	ms/inference
main	477.7
#2549 + #2550	44.8
with this	40.6

So 1.10× on top of the stack, 11.8× against main. Same predicted class as
the CPU path.

Grouped convolutions

The same kernel takes grouped (non-depthwise) convolutions — group on the grid's
z axis, implicit GEMM sized per group — under mlx's gate: C/group ≤ 4 or a
multiple of 16, and O/group ≤ 16 or a multiple of 16. Those were on the direct
kernel before, so the gap is wide (M1 Pro, 3×3, ms/dispatch):

shape	direct	mlx
56×56×128 → 128, g4	11.13	0.149	74.5×
28×28×256 → 256, g8	5.562	0.089	62.5×
56×56×256 → 256, g16	11.42	0.283	40.4×
14×14×512 → 512, g16	2.571	0.064	40.1×
28×28×512 → 512, g32	5.365	0.161	33.4×

No end-to-end number for these: I have no grouped-conv model to hand, and none
of the vision models I used has them. Depthwise keeps its own kernel (#2550) and
groups outside mlx's gate stay on the direct one.

Worth knowing if you review the tests: for a grouped convolution tract's OHWI
puts the group on the I axis, so its kernel is [O/group, kH, kW, C] while
this kernel indexes [O, kH, kW, C/group]. The metal rule produces the latter
from OIHW. The grouped tests drive MetalTransform end to end rather than
hand-building a layout, because I got that distinction wrong first time and a
hand-built test just compares two different convolutions.

Validation

cargo test -p tract-metal --release: 98 passed, one failure —
test_mfa_attention_causal_const_is_noop, pre-existing on main and unrelated
(#2546). The existing conv tests now exercise both kernels: aligned channels go
to the specialised one, and the unaligned case (C=5, O=7) still covers the
general path, plus grouped cases (2/4/8 groups, stride 2, 1×1, f16) and a
gate test for the grouped shapes mlx would decline. A large-activation case
(1024×1024×32 → 64) covers the shape arithmetic mlx has been fixing upstream
(ml-explore/mlx#3938, #3893). fmt and clippy clean.

M1 Pro only, like #2550 — my M4 was asleep. Happy to add second-device numbers
before this lands if you would rather have them.

mlx (ml-explore/mlx) is MIT, Copyright (c) 2023-2025 Apple Inc.; attributed in
the source header, flattened from the same pinned commit as #2549.

🍍

[czoli1976]

Added a large-activation case (1024x1024x32 -> 64) to the conv tests. mlx has been fixing 32-bit shape arithmetic in this area upstream (ml-explore/mlx#3938, #3893); this port shares that 32-bit arithmetic, and the headroom is wide — the worst realistic shape I tried gives ~131k tiles against i32's 2.1e9 — but the case is covered now rather than assumed.