metal: implicit-GEMM convolution via a ported MLX kernel (9-66x)

[czoli1976]

The Metal convolution is a direct kernel computing one output position per
thread. On an M1 Pro that runs a 56×56×64 → 128 3×3 layer at roughly
20 GFLOP/s; the same layer through mlx's tiled implicit-GEMM conv runs at about
1.2 TFLOP/s. You called convolution the remaining Metal gap on #2320 — this is
the mlx kernel for it, vendored the same way as MlxGemm/MlxGemv.

NHWC, f16/f32, single group, 2D convolutions route to the ported kernel; every
other shape stays on the direct kernel, so nothing regresses by construction.

The shared rewrite puts every conv kernel in OIHW, which the direct kernel
indexes, while the ported kernel wants OHWI. A metal-local rule moves eligible
convs to OHWI and leaves the rest to the shared rule — the kernel is normally
a constant, so the reorder folds away.

Numbers

bench_conv (added here), ms per dispatch, best of 5×10, SameUpper padding.
Both kernels are checked against the CPU op first, each with the weight layout
it expects, so this compares two correct implementations.

M1 Pro, 14-core GPU:

shape	direct	mlx
1×112×112×32 → 64, 3×3	24.84	0.386	64×
1×56×56×64 → 128, 3×3	23.52	0.391	60×
1×28×28×128 → 256, 3×3	23.48	0.421	56×
1×56×56×64 → 64, 1×1	1.442	0.061	24×
1×224×224×3 → 32, 3×3 s2	0.504	0.187	2.7×
8×28×28×128 → 128, 3×3	92.43	1.407	66×

M4, 8-core GPU (macOS 26.6) — its direct kernel is much less bad, so the ratios
are smaller, but the ported kernel is still 9–15×:

shape	direct	mlx
1×112×112×32 → 64, 3×3	5.312	0.407	13.1×
1×56×56×64 → 128, 3×3	5.338	0.396	13.5×
1×56×56×64 → 64, 1×1	0.498	0.056	8.9×
1×224×224×3 → 32, 3×3 s2	0.305	0.178	1.7×
8×28×28×128 → 128, 3×3	21.70	1.438	15.1×

f16 tracks f32 within a few percent on both machines. The first-layer case
(3 input channels) gains least, as expected — there is little to tile.

End to end

Inception v3 (TF, NHWC, 299×299) on --metal, M1 Pro. All 54 convolutions are
regular and route to the ported kernel; MobileNet v2, for contrast, is 17/18
depthwise and gets almost nothing from this — worth knowing which models benefit.

	main	this PR
M1 Pro	477.7 ms	44.8	10.7×
M4	158.4 ms	35.9	4.4×

Stable across three interleaved runs on each — M1 Pro 477.7/477.2/477.1 against
44.8/44.9/45.7; M4 165.2/158.4/159.3 against 60.8/35.9/36.6, where the first
PR run pays pipeline specialisation. Same predicted class as the CPU path,
output finite.

Validation

cargo test -p tract-metal --release: 86 passed on the M1 Pro and on the M4,
with one failure on both — test_mfa_attention_causal_const_is_noop, which is
pre-existing on main and unrelated (#2546). New tests cover 1×1, 3×3 valid and
same, strided, dilated, f16 and unaligned channels against the CPU op, plus
end-to-end cases that run a conv through MetalTransform from each of the three
kernel layouts — OHWI, OIHW and HWIO — and check the result against CPU.
That last one matters: the first version of this assumed OIHW, which is right
for ONNX and wrong for TF, and the reorder produced a kernel whose spatial dims
were channel counts. It did not crash or fail a unit test — it just ran the
model 470× slower, which only the end-to-end run surfaced. fmt and clippy clean.

Not covered, left on the direct kernel: NCHW, grouped and depthwise convs, 3D,
and the specialised non-general implicit-GEMM and Winograd variants mlx also
has. Those are follow-ups if this shape of thing is welcome.

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

🍍

[czoli1976]

Updated with end-to-end numbers and a fix. The first version of this reordered the conv kernel assuming OIHW, which is right for ONNX and wrong for TF — the resulting kernel had channel counts where its spatial dims should be. No test failed; the model just ran 470× slower. Inception v3 now goes 477.7 → 44.8 ms on an M1 Pro and 158.4 → 35.9 ms on an M4, and there are end-to-end tests for all three kernel layouts.

[czoli1976]

Depthwise is stacked on this as #2550 — it covers the case this PR declines (MobileNet v2 is 17/18 depthwise), 2–14× at the kernel and 1.34× end to end on MobileNet.