Metal backend: conv2d returns wrong results (`aoti_torch_mps_convolution` ignores tensor strides)

[abdelaziz-mahdy]

🐛 Describe the bug

Every conv2d that AOTInductor routes through aoti_torch_mps_convolution returns wrong values on the Metal backend. Export succeeds and the model runs without any error, so the failure is silent. torchvision mobilenet_v3_small exported for Metal produces a top-1 logit of 3094 where eager PyTorch (CPU and MPS) and the XNNPACK export of the same model produce 8.146.

Cause. For graphs with 4D convolutions inductor applies its layout optimization: it converts the input and the weight to channels-last and expects a channels-last output. From the generated wrapper for a single Conv2d(8, 8, 3, padding=1, groups=8):

static constexpr int64_t int_array_0[] = {1LL, 8LL, 16LL, 16LL};
static constexpr int64_t int_array_1[] = {2048LL, 1LL, 128LL, 8LL};   // channels-last strides
aoti_torch_empty_strided(4, int_array_0, int_array_1, ..., &buf0_handle);
...
aoti_torch_mps_convolution(buf0, weight, nullptr, ...);
assert_size_stride(buf1, {1LL, 8LL, 16LL, 16LL}, {2048LL, 1LL, 128LL, 8LL}, ...);  // output expected channels-last too

backends/apple/metal/runtime/ops/op_convolution.mm never reads strides. It takes sizes() and builds the MPSGraph op as NCHW / OIHW, so NHWC-ordered memory is interpreted as NCHW on the way in, and an NCHW result is handed to a wrapper that goes on to index it as NHWC.

This is easy to miss: conv1d (3D tensors) does not trigger the layout optimization, and 1x1 conv2d is lowered to mm and never reaches the kernel. The module tests in backends/apple/metal/tests/test_modules.py cover conv1d only.

Repro. Add a 3x3 nn.Conv2d to MODULE_REGISTRY in backends/apple/metal/tests/test_modules.py, for example:

class Conv2dNoBias(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(3, 8, kernel_size=3, stride=1, padding=1, bias=False)

    def forward(self, x):
        return self.conv(x)


MODULE_REGISTRY["conv2d_nobias"] = {
    "model_class": Conv2dNoBias,
    "input_shapes": [(1, 3, 16, 16)],
    "description": "Conv2d layer with 3 input channels, 8 output channels, 3x3 kernel",
}

backends/apple/metal/tests/run_metal_test.sh --build
python -m unittest backends.apple.metal.tests.test_modules.TestMetalBackendModules -k conv2d

AssertionError: False is not true : conv2d_nobias (float32): Output mismatch - max_atol=1.7575674057006836, max_rtol=1.9980312585830688

Plain, biased, strided, depthwise and stacked 3x3 convolutions all fail the same way in float32 and bfloat16 (10 of 10 tests), with a max absolute error around 2 on outputs of magnitude around 1. The existing conv1d tests pass.

Exporting under torch._inductor.config.patch({"layout_optimization": False}) keeps the tensors contiguous and makes all of them, and MobileNetV3, match eager. That confirms the cause and works as a stopgap, but the kernel should handle the layout inductor gives it, the way aoti_torch_mps_mm_out and aoti_torch_mps_addmm_out already detect transposed inputs.

I have a fix with tests ready and will open a PR.

Versions

ExecuTorch: main @ 11120c8dff (also reproduced with the 1.5.0 release: PyPI wheel for export, runtime built from the v1.5.0 tag)
PyTorch version: 2.14.0
torchao: built from source with TORCHAO_BUILD_EXPERIMENTAL_MPS=1
OS: macOS 27.0 (arm64), Apple M2 Pro
Clang version: 21.0.0 (clang-2100.1.1.101)
CMake version: 4.4.3
Python version: 3.10.11

[nil-is-all]

Hi @abdelaziz-mahdy, are you still working on #22952 as a fix for this issue you created?