Pad 2D conv input channels to reach the specialized Metal kernel

[eyupcanakman]

Proposed changes

A conv whose input channels are not a multiple of 16 misses the specialized implicit gemm kernel and falls to the general one that #2258 added for these shapes. Making that path faster was left open there.

Zero padding the input channels up to a multiple of 16 and running the specialized kernel beats the general kernel on the original shape, even after paying for the padded work.

The gate is narrow. It needs a stride of 1, an output large enough for implicit gemm, a kernel area of at least 9, and output channels that are already aligned. conv_transpose2d maps its stride onto an input dilation, so only the stride-1 transpose qualifies.

Leaving the output channels alone lets the kernel write straight into out, so nothing is sliced or copied back and the result stays row contiguous.

Weights get the same zero padding, so the added channels contribute nothing to the sum. Output is bit identical to before in fp32, fp16 and bf16.

Peak memory during the conv grows by the padded copies of the input and the weights, both released with the command buffer. Over 220 shapes that hit the gate it ran a median of 1.8x and at most 2.4x, worst where the channel count sits just above a multiple of 16 and the kernel is large.

Added two shapes to test_conv2d_unaligned_channels.

Benchmarks

M5 Pro, macOS 27.0, fp32. One shape per process, best of 150 iterations, the two builds measured turn by turn, min of 7 rounds.

input x weight	before	after
(4, 32, 32, 21) x (128, 3, 3, 21)	0.194 ms	0.175 ms	1.11x
(4, 32, 32, 370) x (128, 7, 7, 370)	3.532 ms	2.281 ms	1.55x
(16, 56, 56, 40) x (64, 3, 3, 40)	0.854 ms	0.616 ms	1.39x
(8, 56, 56, 24) x (32, 5, 5, 24)	0.511 ms	0.344 ms	1.49x
(16, 28, 28, 200) x (256, 3, 3, 200)	2.237 ms	1.925 ms	1.16x
(16, 224, 224, 40) x (128, 3, 3, 40)	22.570 ms	13.945 ms	1.62x
(8, 128, 128, 40) x (128, 7, 7, 40)	19.232 ms	10.857 ms	1.77x

The first two rows are shapes 1 and 4 from benchmarks/python/conv_unaligned_bench.py, and the last two are larger, where the padded work amortizes and the gain grows. Shapes that miss the gate run the same code as before and measured 0.96x to 1.05x, the noise floor of the measurement.

Checklist

Put an x in the boxes that apply.

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (if needed)

[jagrit06]

It needs a stride of 1

Just to clarify, does that mean performance improvements are only seen when the convolution is run with stride 1 ?
Did you have a chance to try it with other small strides ?

[eyupcanakman]

I tried stride 2 in an earlier prototype, where 3x3 regressed to about 0.89x while 5x5 and 7x7 improved to 1.06x and 1.17x. I haven't re-benchmarked that on the final input-only design, so this PR keeps the stride-1 gate and leaves a size-aware stride-2 gate as a follow-up.