Fix dim_order derivation for channels-last tensors with size-1 dims (#22745)

Fixes #22520. Stride sorting can give a channels-last tensor with a
size-1 channel an unsupported dimension order, such as `(0, 2, 1, 3)`
for the KWS input `(1, 1, 49, 10)`. Pass sizes through tensor specs,
weight emission, and serialization so these orders can be corrected when
strides exactly match channels-last. Preserve existing default orders,
which keeps physically equivalent singleton layouts compatible with
portable kernels. Reuse PyTorch's stride calculation; memory-format pass
behavior is unchanged.

Add coverage for the convolution export, mixed-format addition, 5D ReLU,
and symbolic sizes. Re-enable the NXP channels-last PTQ/QAT cases with
the portable C++ reference and the comparator's standard finite MSE
tolerance (`3.5e-4`).

Validation: 287 tests passed across tensor, dim-order, passes, native
serialization, EXIR serialization, and layout transforms; 49 skipped and
1 expected failure. The unrelated `test_to_out_variant_none_output`
failure was confirmed with pre-PR stride inference. Lintrunner is clean
on all touched files. Portable runtime probes match Python for addition,
5D ReLU, and PTQ/QAT convolution and DSCNN KWS models (XNNPACK PT2E
quantization, no delegation). The full NXP simulator suite requires CI
and was not run locally.

This PR was authored with AI assistance (Muse Code and Codex).

Author: JakeStevens

backends/native/serialization/graph_serialize.py

@@ -205,7 +205,7 @@ def _dim_order(t: torch.Tensor) -> list[int]:
     strides = tuple(t.stride())
     sizes = list(t.shape)
     # dim_order_from_stride handles symbolic strides and rejects stride-0 layouts.
-    dim_order = [int(d) for d in dim_order_from_stride(strides)]
+    dim_order = [int(d) for d in dim_order_from_stride(strides, tuple(sizes))]
     expected = stride_from_dim_order(sizes, dim_order)
     for i in range(ndim):
         # A size-1 dim only ever indexes 0, so its stride is arbitrary and need not

backends/native/test/test_serialize.py

@@ -571,10 +571,9 @@ def test_permuted_contiguous(self):
         self.assertEqual(_dim_order(t), [0, 2, 1])
 
     def test_channels_last_with_size_one_channel(self):
-        # channels-last leaves the size-1 channel with an arbitrary stride, which
-        # must not be treated as a non-expressible layout.
+        # A size-1 channel must not make channels-last non-expressible.
         t = torch.randn(2, 1, 3, 4).to(memory_format=torch.channels_last)
-        self.assertEqual(_dim_order(t), [0, 2, 1, 3])
+        self.assertEqual(_dim_order(t), [0, 2, 3, 1])
 
     def test_sliced_layout_raises(self):
         t = torch.randn(4, 8)[:, :4]

backends/nxp/tests/models/test_mlperf_tiny_keyword_spotting.py

@@ -31,11 +31,11 @@
 
 BOUNDS_MSE = {
     "PTQ": {
-        "channels-last": np.inf,
+        "channels-last": 3.5e-4,
         "channels-first": 2.0e-6,
     },
     "QAT": {
-        "channels-last": np.inf,
+        "channels-last": 3.5e-4,
         "channels-first": 5.0e-5,
     },
 }
@@ -47,19 +47,7 @@ def reseed_model_per_test_run():
     np.random.seed(23)
 
 
-@pytest.mark.parametrize(
-    "channels_last",
-    [
-        False,
-        pytest.param(
-            True,
-            marks=pytest.mark.xfail(
-                reason="EIEX-1082, don't forget to readjust bounds when it start working",
-                strict=True,
-            ),
-        ),
-    ],
-)
+@pytest.mark.parametrize("channels_last", [False, True])
 def test_mlperf_tiny_kws_mse_cpu_vs_npu(mocker, request, channels_last, use_qat):
     # approx. 5 samples per class
     num_samples = 60
@@ -89,16 +77,6 @@ def test_mlperf_tiny_kws_mse_cpu_vs_npu(mocker, request, channels_last, use_qat)
         partial(kws.train_model_fn, channels_last=channels_last) if use_qat else None
     )
 
-    # This model does not work in channels-last format when running with portable kernels.
-    # See more information below.
-    # Github issue: https://github.com/pytorch/executorch/issues/22520
-    # NXP internal issue ID: EIEX-1074
-    ref_model = (
-        ReferenceModel.QUANTIZED_EDGE_PYTHON
-        if channels_last
-        else ReferenceModel.QUANTIZED_EXECUTORCH_CPP
-    )
-
     lower_run_compare(
         model,
         [input_spec],
@@ -107,7 +85,7 @@ def test_mlperf_tiny_kws_mse_cpu_vs_npu(mocker, request, channels_last, use_qat)
         dataset_creator=dataset_creator,
         output_comparator=comparator,
         mocker=mocker,
-        reference_model=ref_model,
+        reference_model=ReferenceModel.QUANTIZED_EXECUTORCH_CPP,
         use_qat=use_qat,
         train_fn=train_fn,
     )

exir/emit/_emitter.py

@@ -2077,7 +2077,7 @@ def _is_buffer(node: Node, graph_signature: ExportGraphSignature) -> bool:
 
                 spec.storage = real_tensor.untyped_storage()
                 spec.stride = real_tensor.stride()
-                spec.dim_order = dim_order_from_stride(spec.stride)
+                spec.dim_order = dim_order_from_stride(spec.stride, tuple(spec.shape))
             # User inputs and mutable buffers are not constants, other buffers or parameters are.
             if initialize_buffer and is_mutable_buffer:
                 spec.const = True

exir/tensor.py

@@ -49,7 +49,9 @@ def contiguous_stride_from_shape(shape: torch.Size) -> Tuple[int]:
     return tuple(reversed(strides))
 
 
-def dim_order_from_stride(stride: Tuple[int]) -> Tuple[bytes]:
+def dim_order_from_stride(
+    stride: Tuple[int], sizes: Optional[Tuple[int]] = None
+) -> Tuple[bytes]:
     """
     Dimension order represents how dimensions are laid out in memory,
     starting from the outer-most to the inner-most dimension.
@@ -65,8 +67,16 @@ def dim_order_from_stride(stride: Tuple[int]) -> Tuple[bytes]:
     in original order. Thus when strides = (4, 3, 1, 1) returned value is (0, 1, 2, 3)
     Another example is: sizes = (1, 3, 1, 1) with strides = (3, 1, 3, 3), returned
     value is (0, 2, 3, 1)
+
+    With sizes, ambiguous non-canonical orders are corrected when strides
+    exactly match channels-last. Preserve existing default orders: portable
+    kernels require matching dim orders even for physically equivalent layouts.
     """
-    from torch.fx.experimental.symbolic_shapes import guard_or_false, guard_or_true
+    from torch.fx.experimental.symbolic_shapes import (
+        guard_or_false,
+        guard_or_true,
+        statically_known_true,
+    )
 
     for s in stride:
         torch._check(s != 0, lambda: "0 in strides is not supported for ExecuTorch.")
@@ -93,6 +103,18 @@ def __lt__(self, other):
     sorted_dims = [
         i[0] for i in sorted(enumerate(stride), key=lambda x: K(x[1]), reverse=True)
     ]
+    ndim = len(stride)
+    if (
+        sizes is not None
+        and len(sizes) == ndim
+        and ndim in (4, 5)
+        and sorted_dims != list(range(ndim))
+    ):
+        from torch._prims_common import make_channels_last_strides_for
+
+        expected = make_channels_last_strides_for(sizes)
+        if all(statically_known_true(s == e) for s, e in zip(stride, expected)):
+            sorted_dims = [0, *range(2, ndim), 1]
     return tuple(typing.cast(Tuple[bytes], sorted_dims))
 
 
@@ -201,7 +223,7 @@ def from_tensor(cls, tensor: torch.Tensor, const: bool = False) -> TensorSpec:
             is_sparse=tensor.is_sparse,
         )
         spec.stride = tensor.stride()
-        spec.dim_order = dim_order_from_stride(spec.stride)
+        spec.dim_order = dim_order_from_stride(spec.stride, tuple(spec.shape))
         spec.requires_grad = tensor.requires_grad
         spec.storage = tensor.untyped_storage() if const else None
 

exir/tensor_layout.py

@@ -35,5 +35,5 @@ def from_tensor(cls, tensor: torch.Tensor) -> "TensorLayout":
         return TensorLayout(
             scalar_type=scalar_type_enum(tensor.dtype),
             sizes=list(tensor.shape),
-            dim_order=list(dim_order_from_stride(tensor.stride())),
+            dim_order=list(dim_order_from_stride(tensor.stride(), tuple(tensor.shape))),
         )

exir/tests/test_dim_order_utils.py

@@ -7,9 +7,12 @@
 # pyre-strict
 import unittest
 
+from typing import List
+
 import torch
-from executorch.exir import to_edge_transform_and_lower
+from executorch.exir import to_edge, to_edge_transform_and_lower
 from executorch.exir.dim_order_utils import get_dim_order, get_memory_format
+from executorch.exir.schema import KernelCall, Tensor
 
 
 class TestDimOrderUtils(unittest.TestCase):
@@ -47,3 +50,77 @@ def forward(self, t1, t2):
         expo_prog = torch.export.export(M, (x, y))
         edge_prog = to_edge_transform_and_lower(expo_prog)
         edge_prog.to_executorch()
+
+    def test_channels_last_single_channel_conv_dim_order(self) -> None:
+        class Conv(torch.nn.Module):
+            def __init__(self) -> None:
+                super().__init__()
+                self.conv = torch.nn.Conv2d(1, 8, kernel_size=3, padding=1)
+
+            def forward(self, x: torch.Tensor) -> torch.Tensor:
+                return self.conv(x)
+
+        model = Conv().to(memory_format=torch.channels_last).eval()
+        inputs = (torch.randn(1, 1, 12, 12).to(memory_format=torch.channels_last),)
+        edge_prog = to_edge(torch.export.export(model, inputs, strict=True))
+        et_program = edge_prog.to_executorch().executorch_program
+
+        conv_found = False
+        for inst in et_program.execution_plan[0].chains[0].instructions:
+            kernel = inst.instr_args
+            if not isinstance(kernel, KernelCall):
+                continue
+            op = et_program.execution_plan[0].operators[kernel.op_index]
+            if op.name != "aten::convolution":
+                continue
+            conv_found = True
+            tensors: List[Tensor] = []
+            for arg in dict.fromkeys(kernel.args):
+                val = et_program.execution_plan[0].values[arg].val
+                if isinstance(val, Tensor) and len(val.sizes) == 4:
+                    tensors.append(val)
+            # Input, weight, and output.
+            self.assertEqual(len(tensors), 3)
+            for tensor in tensors:
+                self.assertIn(
+                    list(tensor.dim_order),
+                    ([0, 1, 2, 3], [0, 2, 3, 1]),
+                )
+            self.assertEqual(list(tensors[0].dim_order), list(tensors[-1].dim_order))
+        self.assertTrue(conv_found)
+
+    def test_singleton_dims_preserve_default_order(self) -> None:
+        class Add(torch.nn.Module):
+            def forward(self, x, y):
+                return x + y
+
+        cases = (
+            (
+                Add(),
+                (
+                    torch.randn(2, 1, 3, 1).to(memory_format=torch.channels_last),
+                    torch.randn(2, 1, 3, 1),
+                ),
+            ),
+            (
+                torch.nn.ReLU(),
+                (torch.randn(2, 1, 3, 1, 1).to(memory_format=torch.channels_last_3d),),
+            ),
+        )
+        for model, inputs in cases:
+            with self.subTest(model=type(model).__name__):
+                program = (
+                    to_edge_transform_and_lower(torch.export.export(model, inputs))
+                    .to_executorch()
+                    .executorch_program
+                )
+                tensors = [
+                    value.val
+                    for value in program.execution_plan[0].values
+                    if isinstance(value.val, Tensor)
+                ]
+                self.assertEqual(len(tensors), len(inputs) + 1)
+                for tensor in tensors:
+                    self.assertEqual(
+                        list(tensor.dim_order), list(range(len(tensor.sizes)))
+                    )

exir/tests/test_tensor.py

@@ -321,6 +321,69 @@ def test_dim_order_from_stride_unbacked(self) -> None:
         with self.assertRaises(RuntimeError):
             dim_order_from_stride((u0, 0, 1))
 
+    def test_dim_order_from_stride_with_sizes(self) -> None:
+        dim_order = dim_order_from_stride((490, 1, 10, 1), (1, 1, 49, 10))
+        self.assertEqual((0, 2, 3, 1), dim_order)
+        self.assertEqual(
+            [490, 1, 10, 1],
+            stride_from_dim_order([1, 1, 49, 10], [0, 2, 3, 1]),
+        )
+
+        self.assertEqual(
+            (0, 1, 2, 3),
+            dim_order_from_stride((490, 490, 10, 1), (1, 1, 49, 10)),
+        )
+
+        self.assertEqual(
+            (0, 1, 2, 3), dim_order_from_stride((1, 1, 1, 1), (2, 1, 1, 1))
+        )
+
+        self.assertEqual(
+            (0, 2, 3, 4, 1),
+            dim_order_from_stride((120, 1, 30, 6, 1), (1, 1, 4, 5, 6)),
+        )
+
+        self.assertEqual(
+            (3, 1, 2, 0),
+            dim_order_from_stride((1, 20, 5, 60), (2, 3, 4, 5)),
+        )
+
+        self.assertEqual((0, 2, 1, 3), dim_order_from_stride((490, 1, 10, 1)))
+
+        t = torch.empty(1, 1, 49, 10).to(memory_format=torch.channels_last)
+        spec = TensorSpec.from_tensor(t)
+        self.assertEqual((0, 2, 3, 1), spec.dim_order)
+
+    def test_dim_order_from_stride_preserves_supported_orders(self) -> None:
+        for sizes, strides in (
+            ((2, 1, 3, 1), (3, 1, 1, 1)),
+            ((2, 1, 3, 1, 1), (3, 1, 1, 1, 1)),
+            ((2, 3, 4, 5), (60, 1, 15, 3)),
+            ((2, 3, 4, 5, 6), (360, 1, 90, 18, 3)),
+        ):
+            with self.subTest(sizes=sizes):
+                self.assertEqual(
+                    dim_order_from_stride(strides),
+                    dim_order_from_stride(strides, sizes),
+                )
+
+    def test_dim_order_from_stride_with_symbolic_sizes(self) -> None:
+        from torch.fx.experimental.symbolic_shapes import ShapeEnv
+
+        shape_env = ShapeEnv()
+        height = shape_env.create_unbacked_symint()
+        torch._check_is_size(height)
+        torch._check(height >= 2)
+        self.assertEqual(
+            (0, 2, 3, 1),
+            dim_order_from_stride((10 * height, 1, 10, 1), (2, 1, height, 10)),
+        )
+        self.assertEqual(
+            (0, 2, 1, 3),
+            dim_order_from_stride((490, 1, 10, 1), (2, 1, height, 10)),
+        )
+        self.assertEqual(shape_env.guards, [])
+
     def test_strides_from_dim_order(self) -> None:
         sizes = []
         dim_order = []