Metal backend: honor channels-last strides in conv2d - #22952
abdelaziz-mahdy wants to merge 6 commits into
Conversation
aoti_torch_mps_convolution read its operands as contiguous NCHW / OIHW regardless of their strides. Inductor's layout optimization gives conv2d channels-last input and weight and expects a channels-last output, so every conv2d reaching the kernel returned wrong values without raising an error. 1x1 convs were unaffected only because inductor lowers them to mm. Classify the input and the weight from their strides, declare channels-last buffers to MPSGraph in their physical NHWC / OHWI shape and reorder them inside the graph, and write the result channels-last when either operand is, as ATen does. Other stride patterns now return InvalidArgument instead of computing garbage. Adds conv2d modules (plain, bias, strided, depthwise, stacked) to the Metal module tests, which covered conv1d only.
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/22952
Note: Links to docs will display an error until the docs builds have been completed. ❗ 1 Active SEVsThere are 1 currently active SEVs. If your PR is affected, please view them below: ❌ 1 Awaiting Approval, 1 New Failure, 1 PendingAs of commit 4bf45ef with merge base 11120c8 ( NEW FAILURE - The following job has failed:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
This PR needs a
|
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
Unresolved critical output-boundary and bias-cache issues block approval.
Get a fresh assessment by requesting another Copilot review.
Review effort: Lite
Findings: 2
Open (2)
What changed in this PR
This PR updates Metal AOTI conv2d to honor channels-last layouts and adds regression coverage.
Changes:
- Detects supported tensor layouts and reorders channels-last operands for MPSGraph.
- Preserves channels-last outputs and extends graph cache keys.
- Adds float32/bfloat16 tests for bias, stride, depthwise, and stacked convolutions.
| File | Summary |
|---|---|
backends/apple/metal/tests/test_modules.py |
Adds five Conv2d regression modules and consistency tests. |
backends/apple/metal/runtime/ops/op_convolution.mm |
Implements layout-aware convolution handling. Findings remain: two critical issues involving output copying and bias cache-key handling, one moderate ambiguous-layout issue, and two nit findings about truncated stride diagnostics. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
A graph built without a bias placeholder could be reused for a same-shaped convolution that has one, which would silently drop the bias (and the reverse would run a graph whose bias placeholder is never fed). Inductor currently adds the bias outside the kernel for MPS, so this was latent.
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
The convolution graph cache key omits output_padding, risking reuse of incompatible transposed-convolution graphs.
Get a fresh assessment by requesting another Copilot review.
Review effort: Lite
Findings: 1
Open (1)
Resolved since last review (2)
Output padding changes a transposed convolution's descriptor, so two transposed convolutions that differ only in output padding must not share a cached graph.
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
Singleton-channel inputs can still produce incorrect output layout metadata and need a regression test.
Get a fresh assessment by requesting another Copilot review.
Review effort: Lite
Findings: 1
Resolved since last review (1)
With one channel, contiguous and channels-last strides describe the same memory, so the kernel reports such a tensor as contiguous. Inductor does the same: it leaves a single-channel conv input contiguous and expects a contiguous result, and it expects a channels-last single-channel result when the input was channels-last. Pin both down.
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
The bias path can produce incorrect results because the rank-1 bias broadcasts along the width dimension instead of channels.
Get a fresh assessment by requesting another Copilot review.
Review effort: Lite
Findings: 1
Resolved since last review (1)
The kernel added the rank-1 [C_out] bias straight to the NCHW result. MPSGraph broadcasts from the trailing dimension, so the bias lined up with W: when C_out happened to equal W every channel got the wrong offsets, and otherwise MPSGraph aborted on the shape mismatch. Reshape the bias to [1, C_out, 1, 1] first. Latent so far, because inductor adds the bias itself for MPS and always calls the kernel with bias=nullptr.
|
Found an additional conv2d case this fix still misses at torch.manual_seed(2026)
dtype = torch.float32 # also reproduces with torch.bfloat16
model = torch.nn.Sequential(
torch.nn.Conv2d(3, 1, 1),
torch.nn.Conv2d(1, 4, 3, padding=1),
).eval().to(dtype)
inputs = (torch.ones(2, 3, 7, 9, dtype=dtype),)Exporting with the existing
The disabled-layout controls pass the module test tolerances. The baseline has the same failure, so this is an incomplete fix rather than a new regression. All 12 added module/dtype cases pass on this PR. The 1×1 convolution lowers to addmm. For the second convolution, the generated wrapper expects channels-last output strides
Validation performed with Codex on Apple M1 Pro, macOS 26.6.2, PyTorch 2.14.0, using the PR's native convolution implementation. |
A 1x1 conv producing one channel lowers to addmm, and its result reaches the
next conv as a {2, 1, 7, 9} tensor with channels-last strides {63, 1, 9, 1}.
Those strides describe the same memory as contiguous ones, so the kernel
treated the input as contiguous and returned a contiguous output, while the
generated wrapper expects channels-last: max abs error 1.3 on the chain
Conv2d(3, 1, 1) -> Conv2d(1, 4, 3, padding=1).
Two things were needed:
- The kernel now resolves the ambiguous case the way PyTorch's MPS rule does
(suggest_memory_format of input or weight is channels-last), with a port of
c10's is_channels_last_strides_2d, which looks at the actual stride values.
- Those values did not survive tensor creation. make_tensor_ptr derives the
strides again from a dim order it gets by sorting them, and a size-1
dimension ties with its neighbour, so {63, 1, 9, 1} came back as
{63, 9, 9, 1}. The Metal shims now pass an explicit dim order that puts the
size-1 dimension last among equal strides, which gives back the strides
inductor asked for.
Adds the chain as conv2d_pointwise_to_single_channel.
|
Reproduced, including the
Your chain is added as One thing unrelated to this PR: on my machine (macOS 27.0) the 10 int4 module tests fail with and without this change, because the Metal compiler rejects the 4-bit shader source: |
|
Re-reviewed The constant below has shape Repro using the existing module-test export helper: from pathlib import Path
import torch
from executorch.backends.apple.metal.tests.test_modules import export_model_to_pte
class Repro(torch.nn.Module):
def __init__(self):
super().__init__()
self.register_buffer(
"weight",
torch.arange(8).reshape(2, 4, 1).float().transpose(1, 2),
)
def forward(self, x):
return torch.bmm(self.weight, x)
out = Path("singleton-bmm-repro")
out.mkdir(exist_ok=True)
for dtype in (torch.float32, torch.bfloat16):
model = Repro().eval().to(dtype)
inputs = (torch.ones(2, 4, 3, dtype=dtype),)
name = f"repro_{str(dtype).split('.')[-1]}"
pte, expected = export_model_to_pte(model, inputs, out, name)
print(pte, expected)Run the same exported PTE with the baseline and PR runtimes, e.g.: executor_runner --model_path singleton-bmm-repro/repro_float32.pte --output_file singleton-bmm-repro/output
The PR reports: There is a second instance of the same problem in A local diagnostic build keeps the convolution stride-preservation fix and makes BMM's contiguity checks ignore size-one axes, while the copy check ignores a stride difference only when both tensors have size one on that axis. Both native repros and both exported BMM cases then pass, and all 18 previously tested convolution cases still pass, including the original single-channel chain. Please add regression coverage and update these consumers alongside the shared tensor-construction change, or land the consumer updates as a prerequisite. The strict checks predate this PR, but changing their inputs makes this a new regression for an existing model. Validation performed with Codex on Apple M1 Pro, macOS 26.6.2, PyTorch 2.14.0, using the current PR's rebuilt native runtime. |
Summary
Fixes #22951
aoti_torch_mps_convolutionread its operands as contiguous NCHW / OIHW regardless of their strides. Inductor's layout optimization gives conv2d channels-last input and weight and expects a channels-last output, so every conv2d reaching the kernel returned wrong values without raising an error (MobileNetV3: top-1 logit 3094 instead of 8.146). 1x1 convs were unaffected only because inductor lowers them tomm.The kernel now classifies the input and the weight from their strides (contiguous or channels-last, ignoring size-1 dimensions, whose stride carries no information). Channels-last buffers are declared to MPSGraph in their physical NHWC / OHWI shape and reordered inside the graph, the convolution itself is still described as NCHW / OIHW, and the result is written in channels-last order with matching strides when either operand is channels-last, as ATen does. The two layout flags are part of the graph cache key. Any other stride pattern now returns
InvalidArgumentinstead of computing garbage, matching howaoti_torch_mps_bmm_outtreats layouts it does not support.Contiguous operands, conv1d and unbatched 3D input take the same path as before.
Test plan
Added five conv2d modules to
MODULE_REGISTRYinbackends/apple/metal/tests/test_modules.py(plain, bias with batch 2, stride 2 on a non-square input, depthwise, and two stacked convolutions so a channels-last output feeds the next conv). Each runs in float32 and bfloat16.conv2d_*_output_consistencytests fail withmax_atolbetween 1.6 and 2.7; everything else passes.Ran 134 tests ... OK..pte(exported by the 1.5.0 release, layout optimization on) now matches eager throughexecutor_runneron an all-ones input: first logits-0.0298, -0.1163, 0.2346, -0.1151, 0.2841, 1.3327, -1.2021, -0.4182on both sides.lintrunnerclean on the touched files.