Cortex-M: add a contiguous-layout pad operator

backends/cortex_m/ops/op_pad.cpp

@@ -17,21 +17,19 @@ namespace {
 
 constexpr size_t kMaxSupportedDims = 4;
 
-} // namespace
-
-// cppcheck-suppress unusedFunction
-Tensor& pad_out(
+Tensor& pad_out_impl(
     KernelRuntimeContext& context,
     const Tensor& input,
     const Int64ArrayRef pre_pad,
     const Int64ArrayRef post_pad,
     int64_t pad_value,
+    bool require_contiguous,
     Tensor& out) {
   if (input.scalar_type() != ScalarType::Char ||
       out.scalar_type() != ScalarType::Char) {
     ET_LOG(
         Error,
-        "pad_out: only int8 tensors are supported (input=%d, out=%d)",
+        "cortex_m::pad: only int8 tensors are supported (input=%d, out=%d)",
         static_cast<int>(input.scalar_type()),
         static_cast<int>(out.scalar_type()));
     context.fail(Error::InvalidArgument);
@@ -42,22 +40,48 @@ Tensor& pad_out(
   if (rank == 0 || rank > kMaxSupportedDims) {
     ET_LOG(
         Error,
-        "pad_out: expected tensor rank in [1, %zu], got %zu",
+        "cortex_m::pad: expected tensor rank in [1, %zu], got %zu",
         kMaxSupportedDims,
         rank);
     context.fail(Error::InvalidArgument);
     return out;
   }
+  if (pre_pad.size() != kMaxSupportedDims ||
+      post_pad.size() != kMaxSupportedDims) {
+    ET_LOG(Error, "cortex_m::pad: pre_pad and post_pad must have length 4");
+    context.fail(Error::InvalidArgument);
+    return out;
+  }
+
+  if (require_contiguous) {
+    // This entry point infers nothing: it requires the dim order to say the
+    // tensor is contiguous, and then indexes the padding by logical axis.
+    if (!executorch::runtime::is_contiguous_dim_order(
+            input.dim_order().data(), input.dim_order().size()) ||
+        !executorch::runtime::is_contiguous_dim_order(
+            out.dim_order().data(), out.dim_order().size())) {
+      ET_LOG(
+          Error,
+          "cortex_m::pad_contiguous: input and output must use contiguous dim order");
+      context.fail(Error::InvalidArgument);
+      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);
+  // Only the legacy entry point infers the layout. Its predicate is tolerant on
+  // purpose: the tolerance short-circuits before the dim order is consulted,
+  // which is what keeps it agreeing with the AOT pass for shapes whose
+  // serialized dim order cannot name the channel axis.
+  const bool legacy_channels_last =
+      !require_contiguous && is_channels_last_tensor(input);
 
   int32_t dims[kMaxSupportedDims] = {1, 1, 1, 1};
   for (size_t i = 0; i < rank; ++i) {
-    const size_t src = nhwc ? kNhwcDimOrder[offset + i] : i;
+    const size_t src = legacy_channels_last ? kNhwcDimOrder[offset + i] : i;
     dims[offset + i] = static_cast<int32_t>(input.size(src));
   }
 
@@ -87,7 +111,7 @@ Tensor& pad_out(
   if (status != ARM_CMSIS_NN_SUCCESS) {
     ET_LOG(
         Error,
-        "pad_out: arm_pad_s8 failed with status [%d]",
+        "cortex_m::pad: arm_pad_s8 failed with status [%d]",
         static_cast<int>(status));
     context.fail(Error::Internal);
     return out;
@@ -96,5 +120,43 @@ Tensor& pad_out(
   return out;
 }
 
+} // namespace
+
+// cppcheck-suppress unusedFunction
+Tensor& pad_out(
+    KernelRuntimeContext& context,
+    const Tensor& input,
+    const Int64ArrayRef pre_pad,
+    const Int64ArrayRef post_pad,
+    int64_t pad_value,
+    Tensor& out) {
+  return pad_out_impl(
+      context,
+      input,
+      pre_pad,
+      post_pad,
+      pad_value,
+      /*require_contiguous=*/false,
+      out);
+}
+
+// cppcheck-suppress unusedFunction
+Tensor& pad_contiguous_out(
+    KernelRuntimeContext& context,
+    const Tensor& input,
+    const Int64ArrayRef pre_pad,
+    const Int64ArrayRef post_pad,
+    int64_t pad_value,
+    Tensor& out) {
+  return pad_out_impl(
+      context,
+      input,
+      pre_pad,
+      post_pad,
+      pad_value,
+      /*require_contiguous=*/true,
+      out);
+}
+
 } // namespace native
 } // namespace cortex_m

backends/cortex_m/ops/operators.py

@@ -1,3 +1,3 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
 # Copyright 2025-2026 Arm Limited and/or its affiliates.
@@ -663,6 +663,18 @@
     "pad.out(Tensor input, int[] pre_pad, int[] post_pad, int pad_value, "
     "*, Tensor(a!) out) -> Tensor(a!)"
 )
+lib.define(
+    "pad_contiguous(Tensor input, int[] pre_pad, int[] post_pad, int pad_value) -> Tensor"
+)
+lib.define(
+    "pad_contiguous.out(Tensor input, int[] pre_pad, int[] post_pad, int pad_value, "
+    "*, Tensor(a!) out) -> Tensor(a!)"
+)
+
+
+_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):
@@ -717,6 +729,58 @@
     return F.pad(input, padding, mode="constant", value=pad_value)
 
 
+@register_fake("cortex_m::pad_contiguous")  # type: ignore[misc]
+def pad_contiguous_meta(
+    input: torch.Tensor,
+    pre_pad: list[int],
+    post_pad: list[int],
+    pad_value: int,
+) -> torch.Tensor:
+    del pad_value
+    rank = input.dim()
+    if rank == 0 or rank > 4:
+        raise RuntimeError(
+            f"cortex_m.pad_contiguous expects a rank in [1, 4], got {rank}"
+        )
+    if len(pre_pad) != 4 or len(post_pad) != 4:
+        raise RuntimeError(
+            "cortex_m.pad_contiguous expects four padding values per side"
+        )
+    offset = 4 - rank
+    output_shape = [
+        input.shape[dim] + pre_pad[offset + dim] + post_pad[offset + dim]
+        for dim in range(rank)
+    ]
+    return torch.empty(output_shape, dtype=input.dtype, device=input.device)
+
+
+@impl(lib, "pad_contiguous", "CompositeExplicitAutograd")  # type: ignore[misc]
+def pad_contiguous_impl(
+    input: torch.Tensor,
+    pre_pad: list[int],
+    post_pad: list[int],
+    pad_value: int,
+) -> torch.Tensor:
+    rank = input.dim()
+    if rank == 0 or rank > 4:
+        raise RuntimeError(
+            f"cortex_m.pad_contiguous expects a rank in [1, 4], got {rank}"
+        )
+    if len(pre_pad) != 4 or len(post_pad) != 4:
+        raise RuntimeError(
+            "cortex_m.pad_contiguous expects four padding values per side"
+        )
+    offset = 4 - rank
+    padding = []
+    for dim in reversed(range(rank)):
+        padding.extend([pre_pad[offset + dim], post_pad[offset + dim]])
+    return F.pad(input, padding, mode="constant", value=pad_value)
+
+
+# ===================================================================
+# QUANTIZED CONV2D OPERATION DEFINITION
+# ===================================================================
+
 lib.define(
     "quantized_conv2d("
     "Tensor input, "

backends/cortex_m/ops/operators.yaml

@@ -77,6 +77,12 @@
     - arg_meta: null
       kernel_name: cortex_m::pad_out
 
+- func: cortex_m::pad_contiguous.out(Tensor input, int[] pre_pad, int[] post_pad, int pad_value, *, Tensor(a!) out) -> Tensor(a!)
+  variants: function
+  kernels:
+    - arg_meta: null
+      kernel_name: cortex_m::pad_contiguous_out
+
 - func: cortex_m::quantized_conv2d.out(Tensor input, Tensor weight, Tensor? bias, int[] stride, int[] padding, int[] dilation, int input_offset, int output_offset, Tensor requantize_multipliers, Tensor requantize_shifts, int activation_min, int activation_max, Tensor scratch, *, Tensor(a!) out) -> Tensor(a!)
   variants: function
   kernels:

backends/cortex_m/test/build_test_runner.sh

@@ -65,6 +65,7 @@ ops_list=(
     cortex_m::softmax.out
     cortex_m::transpose.out
     cortex_m::pad.out
+    cortex_m::pad_contiguous.out
     cortex_m::quantized_conv2d.out
     cortex_m::quantized_conv2d_nhwc.out
     cortex_m::quantized_depthwise_conv2d.out

backends/cortex_m/test/ops/test_explicit_nhwc_runtime.py

@@ -217,6 +217,16 @@ def forward(self, x):
         )
 
 
+class PadNhwc(torch.nn.Module):
+    def forward(self, x):
+        return torch.ops.cortex_m.pad_contiguous.default(
+            x,
+            [0, 1, 2, 0],
+            [0, 2, 1, 0],
+            -7,
+        )
+
+
 def test_conv2d_nhwc_runs_on_fvp(cortex_m_target):
     _run_on_fvp(
         Conv2dNhwc(),
@@ -276,3 +286,11 @@ def test_max_pool2d_nhwc_runs_on_fvp(cortex_m_target):
         cortex_m_target,
     )
 
+
+def test_pad_contiguous_runs_on_fvp_with_singleton_height(cortex_m_target):
+    _run_on_fvp(
+        PadNhwc(),
+        _int8_values((1, 1, 7, 3)),
+        exir_ops.edge.cortex_m.pad_contiguous.default,
+        cortex_m_target,
+    )

backends/cortex_m/test/ops/test_pad.py

@@ -77,6 +77,19 @@ def forward(self, x):
         CortexMPad((1, 2, 3, 4)),
         (ramp_tensor(-1.0, 1.0, (1, 3, 4, 5)).to(memory_format=torch.channels_last),),
     ),
+    # A channels-last tensor with one channel serializes its dim order as
+    # (0, 2, 1, 3), which names no channel axis. Deriving the physical sizes
+    # from it instead of from the shape sizes the pad write wrongly.
+    "pad_rank4_single_channel_channels_last": McuTestCase(
+        CortexMPad((1, 1, 2, 2)),
+        (ramp_tensor(-0.5, 0.5, (1, 1, 3, 4)).to(memory_format=torch.channels_last),),
+    ),
+    # With one channel and unit width the dim order collapses all the way to
+    # (0, 1, 2, 3), making the tensor indistinguishable from a contiguous one.
+    "pad_rank4_single_channel_unit_width_channels_last": McuTestCase(
+        CortexMPad((0, 0, 2, 2)),
+        (ramp_tensor(-0.5, 0.5, (1, 1, 8, 1)).to(memory_format=torch.channels_last),),
+    ),
 }