Fix dim_order derivation for channels-last tensors with size-1 dims

backends/native/serialization/graph_serialize.py

@@ -203,7 +203,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

@@ -554,10 +554,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

@@ -29,11 +29,11 @@
 
 BOUNDS_MSE = {
     "PTQ": {
-        "channels-last": np.inf,
+        "channels-last": 3.5e-4,
         "channels-first": 5.5e-7,
     },
     "QAT": {
-        "channels-last": np.inf,
+        "channels-last": 3.5e-4,
         "channels-first": 3.3e-5,
     },
 }
@@ -45,19 +45,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
@@ -87,16 +75,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],
@@ -105,7 +83,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

@@ -2034,7 +2034,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 = []