Cortex-M: Fix pad op to support channels_last memory format (pytorch#18429)

### Summary
Fix pad_meta to propagate channels_last from input to output tensor.
Fix pad_out (C++) to use dim_order() to permute logical dims and padding
into physical memory order for arm_pad_s8.
Add channels_last test cases to test_pad.

### Test Plan
```
pytest backends/cortex_m/test/ops/test_pad.py
```

---------

Co-authored-by: Claude <noreply@anthropic.com>

Author: rascani

backends/cortex_m/ops/op_pad.cpp

@@ -48,14 +48,19 @@ Tensor& pad_out(
     return out;
   }
 
+  // Permute logical sizes to physical memory order.
+  // Padding is already in physical order from the AOT pass.
+  constexpr size_t kNhwcDimOrder[] = {0, 2, 3, 1};
   const size_t offset = kMaxSupportedDims - rank;
+  const bool nhwc = is_channels_last_tensor(input);
 
-  cmsis_nn_dims input_dims = {1, 1, 1, 1};
-  int32_t* d = &input_dims.n;
+  int32_t dims[kMaxSupportedDims] = {1, 1, 1, 1};
   for (size_t i = 0; i < rank; ++i) {
-    d[offset + i] = static_cast<int32_t>(input.size(i));
+    const size_t src = nhwc ? kNhwcDimOrder[offset + i] : i;
+    dims[offset + i] = static_cast<int32_t>(input.size(src));
   }
 
+  cmsis_nn_dims input_dims = {dims[0], dims[1], dims[2], dims[3]};
   cmsis_nn_dims cmsis_pre_pad = {
       static_cast<int32_t>(pre_pad[0]),
       static_cast<int32_t>(pre_pad[1]),

backends/cortex_m/ops/operators.py

@@ -14,6 +14,7 @@
 from executorch.backends.cortex_m.passes.passes_utils import (
     dequantize_per_tensor_cmsis,
     is_channel_broadcast,
+    is_channels_last,
     quantize_per_tensor_cmsis,
     requantize_cmsis,
     SHIFT_INT8,
@@ -564,6 +565,16 @@ def transpose_impl(input: torch.Tensor, perm: Sequence[int]) -> torch.Tensor:
 )
 
 
+_NHWC_INV_ORDER = [0, 3, 1, 2]
+
+
+def _pad_to_logical_order(physical_pad: list[int], input: torch.Tensor) -> list[int]:
+    """Inverse of _to_physical_order: map physical-order padding back to logical."""
+    if not is_channels_last(input):
+        return list(physical_pad)
+    return [physical_pad[_NHWC_INV_ORDER[i]] for i in range(4)]
+
+
 @register_fake("cortex_m::pad")  # type: ignore[misc]
 def pad_meta(
     input: torch.Tensor,
@@ -573,10 +584,16 @@ def pad_meta(
 ) -> torch.Tensor:
     rank = input.dim()
     offset = 4 - rank
+    logical_pre = _pad_to_logical_order(pre_pad, input)
+    logical_post = _pad_to_logical_order(post_pad, input)
+
     output_shape = list(input.shape)
     for i in range(rank):
-        output_shape[i] += pre_pad[offset + i] + post_pad[offset + i]
-    return torch.empty(output_shape, dtype=input.dtype, device=input.device)
+        output_shape[i] += logical_pre[offset + i] + logical_post[offset + i]
+    result = torch.empty(output_shape, dtype=input.dtype, device=input.device)
+    if is_channels_last(input):
+        result = result.to(memory_format=torch.channels_last)
+    return result
 
 
 @impl(lib, "pad", "CompositeExplicitAutograd")  # type: ignore[misc]
@@ -588,9 +605,12 @@ def pad_impl(
 ) -> torch.Tensor:
     rank = input.dim()
     offset = 4 - rank
+    logical_pre = _pad_to_logical_order(pre_pad, input)
+    logical_post = _pad_to_logical_order(post_pad, input)
+
     padding = []
     for i in reversed(range(rank)):
-        padding.extend([pre_pad[offset + i], post_pad[offset + i]])
+        padding.extend([logical_pre[offset + i], logical_post[offset + i]])
     return F.pad(input, padding, mode="constant", value=pad_value)
 
 

backends/cortex_m/passes/quantized_op_fusion_pass.py

@@ -10,6 +10,7 @@
 
 import torch
 from executorch.backends.cortex_m.passes.passes_utils import (
+    is_channels_last,
     quantize_multiplier_aot,
     quantize_val,
     SHIFT_INT8,
@@ -37,6 +38,14 @@ class QuantizedOpFusionPass(ExportPass):
 
     _SOFTMAX_INPUT_INTEGER_BITS = 5
 
+    _NHWC_DIM_ORDER = [0, 2, 3, 1]
+
+    def _to_physical_order(self, logical_pad: list[int], tensor_data) -> list[int]:
+        """Permute a 4-element logical-dim-order list to physical memory order."""
+        if not is_channels_last(tensor_data):
+            return logical_pad
+        return [logical_pad[self._NHWC_DIM_ORDER[i]] for i in range(4)]
+
     def _get_add_replacement(self, args, meta):
         if (
             meta.data.get("input_qparams", {}) == {}
@@ -329,6 +338,8 @@ def _get_avg_pool2d_replacement(self, args, meta):
             pad_h, pad_w = padding
             pre_pad = [0, 0, pad_h, pad_w]
             post_pad = [0, 0, pad_h, pad_w]
+            pre_pad = self._to_physical_order(pre_pad, args[0].data)
+            post_pad = self._to_physical_order(post_pad, args[0].data)
             input_arg = super().call_operator(
                 exir_ops.edge.cortex_m.pad.default,
                 (input_arg, pre_pad, post_pad, int(zero_point)),
@@ -379,6 +390,9 @@ def _get_pad_replacement(self, args, meta):
             pre_pad[dim_4d] = int(padding[2 * i])
             post_pad[dim_4d] = int(padding[2 * i + 1])
 
+        pre_pad = self._to_physical_order(pre_pad, args[0].data)
+        post_pad = self._to_physical_order(post_pad, args[0].data)
+
         new_args = (args[0], pre_pad, post_pad, int(quantized_pad_value))
         return exir_ops.edge.cortex_m.pad.default, new_args
 

backends/cortex_m/test/ops/test_pad.py

@@ -69,6 +69,14 @@ def forward(self, x):
         CortexMPad((0, 0, 0, 0)),
         (ramp_tensor(-0.5, 0.5, (2, 3, 4, 5)),),
     ),
+    "pad_rank4_all_dims_channels_last": McuTestCase(
+        CortexMPad((1, 1, 2, 2, 1, 0, 0, 1)),
+        (ramp_tensor(-0.5, 0.5, (1, 2, 3, 4)).to(memory_format=torch.channels_last),),
+    ),
+    "pad_rank4_last_two_dims_channels_last": McuTestCase(
+        CortexMPad((1, 2, 3, 4)),
+        (ramp_tensor(-1.0, 1.0, (1, 3, 4, 5)).to(memory_format=torch.channels_last),),
+    ),
 }