[Metal][Performance] Avoid zero work in stride-2 ConvTranspose3d

[ternaus]

Proposed changes

Fixes #4341.

conv_transpose3d(..., kernel_size=2, stride=2, padding=0, output_padding=0) is represented internally as a flipped conv_general with input_dilation=2. The existing Metal 3D dispatch falls through to explicit unfold + GEMM; for this exact configuration, each output coordinate has one valid kernel phase and seven zero phases.

This PR adds a narrowly gated Metal path that:

• splits the output into its eight spatial parity phases;
• reuses the input as an [M, C] matrix for each phase;
• exposes the corresponding [C, O] weight slice as a zero-copy view;
• runs the existing Steel GEMM and copies each phase into its strided output view;
• leaves every other transposed-convolution configuration on the existing fallback.

No public API or Python behavior changes. The path is restricted to groups=1, flip=true, input_dilation=(2,2,2), kernel_dilation=(1,1,1), stride=(1,1,1), kernel=(2,2,2), transformed low padding (1,1,1), and output shape exactly twice the input shape.

Correctness

• Added the exact stride-2/kernel-2 case to the existing PyTorch comparison test.
• Built MLX with the Metal backend and ran the repository test against the Apple M4 Max GPU.
• float32 and float16 non-square 3D cases with unequal input/output channels both matched PyTorch with max_abs_err=0.0.
• The general explicit path remains unchanged for output padding, non-zero user padding, other kernels/strides, dilation, and groups.

Benchmark

Same locally built MLX revision, same M4 Max, float16, five warmup evaluations followed by twenty timed mx.eval evaluations:

input (N,D,H,W,C)	output	upstream main	this PR	speedup
(1,32,64,64,60)	(1,64,128,128,60)	23.04 ms	1.79 ms	12.9x
(1,64,64,64,60)	(1,128,128,128,60)	45.75 ms	3.54 ms	12.9x

The improvement comes from avoiding the dense unfold matrix: for this exact configuration, seven of eight phase blocks are zeros.

Validation

• pre-commit run --files mlx/backend/metal/conv.cpp python/tests/test_conv_transpose.py
• Full Metal C++ build and metallib: passed.
• Metal C++ suite: 277 test cases / 3714 assertions passed.
• Repository Python test: test_torch_conv_transpose_3D passed.

Checklist

• I have read the CONTRIBUTING document
• I have run the configured pre-commit checks on changed files
• I have added a test that covers the new dispatch shape
• I have validated the Metal build and GPU behavior locally
• I have updated documentation (not needed; no public API change)

[zcbenz]

What is this configuration used for?

[ternaus]

What is this configuration used for?

At the public API level, this fast path covers the non-overlapping learned 2× upsampling operation:

conv_transpose3d(
    ...,
    kernel_size=(2, 2, 2),
    stride=(2, 2, 2),
    padding=(0, 0, 0),
    output_padding=(0, 0, 0),
    dilation=(1, 1, 1),
    groups=1,
)

This is a canonical decoder operation in 3D U-Net-style segmentation networks. The original 3D U-Net paper specifies a 2×2×2 up-convolution with stride 2 in every spatial dimension at each level of the synthesis path.

nnU-Net uses the same pattern. Its public PyTorch implementation selects nn.ConvTranspose3d for 3D networks, defaults the spatial upsampling kernels to (2, 2, 2), and passes the same tuple as both the transposed-convolution kernel size and stride:

• ConvTranspose3d and the default (2, 2, 2) kernels
• kernel size and stride taken from the same pool_op_kernel_sizes entry

My workload that exposed the MLX bottleneck was a FracSegNet-derived 3D pelvic-fracture segmentation model used for PENGWIN ML Competition. The FracSegNet paper integrates its distance-weighted objective into a 3D U-Net, and its public training code instantiates nnU-Net’s Generic_UNet with convolutional upsampling enabled.

This architecture family is used across a broad range of volumetric biomedical segmentation tasks. The nnU-Net study evaluated it on 23 public competition datasets: the ten Medical Segmentation Decathlon tasks, BCV, PROMISE12, ACDC, LiTS, the MS lesion challenge, CHAOS, KiTS19, SegTHOR, CREMI, and four Cell Tracking Challenge datasets.

The broader Medical Segmentation Decathlon analysis found that U-Net was the base architecture for 64% of the participating methods that provided complete algorithmic information. nnU-Net won the Decathlon and subsequently won 33 of 53 additional segmentation tasks observed over the following two years. The reported applications include BraTS 2020, KiTS19, and COVID-19-20.

Some 3D segmentation networks use interpolation-based upsampling, and anisotropic nnU-Net stages can use strides such as (1, 2, 2). That is why this optimization remains narrowly gated. The exact (2, 2, 2) case nevertheless recurs throughout isotropic 3D U-Net and nnU-Net decoders and represents a mainstream volumetric-segmentation workload rather than a model-specific special case.

[zcbenz]

Thanks for the clarification, and nice improvement!