Remove channels_last requirement for pool ops via explicit permutes (#21796)

backends/cortex_m/ops/cortex_m_ops_common.h

@@ -143,6 +143,45 @@ inline bool is_channels_last_tensor(const Tensor& tensor) {
   return tensor.dim_order() == channels_last_order;
 }
 
+// Strict channels_last check: true only when the tensor explicitly carries the
+// channels_last dim_order [0, 2, 3, 1]. Unlike is_channels_last_tensor(), it
+// does NOT treat unit-dim tensors as ambiguously channels_last, so a
+// dim-order-free contiguous NHWC tensor (dim_order [0, 1, 2, 3]) is never
+// misclassified.
+inline bool has_channels_last_dim_order(const Tensor& tensor) {
+  if (tensor.dim() != 4) {
+    return false;
+  }
+  constexpr executorch::aten::DimOrderType kChannelsLastDimOrder[] = {
+      0, 2, 3, 1};
+  executorch::aten::ArrayRef<executorch::aten::DimOrderType>
+      channels_last_order(kChannelsLastDimOrder, 4);
+  return tensor.dim_order() == channels_last_order;
+}
+
+// Logical (batch, height, width, channels) of a 4-D activation tensor under
+// either supported contract:
+//   - legacy channels_last: logical [N, C, H, W] (dim_order [0, 2, 3, 1]);
+//   - dim-order-free:       logical [N, H, W, C] (plain contiguous).
+// The CMSIS-NN kernels always operate on physically-NHWC data, so reading dims
+// this way lets them serve both previously-generated channels_last .pte files
+// and new contiguous ones through one code path.
+struct NhwcDims {
+  int64_t n;
+  int64_t h;
+  int64_t w;
+  int64_t c;
+};
+
+inline NhwcDims read_nhwc_dims(const Tensor& tensor) {
+  if (has_channels_last_dim_order(tensor)) {
+    return NhwcDims{
+        tensor.size(0), tensor.size(2), tensor.size(3), tensor.size(1)};
+  }
+  return NhwcDims{
+      tensor.size(0), tensor.size(1), tensor.size(2), tensor.size(3)};
+}
+
 inline bool is_channel_broadcast(const Tensor& tensor1, const Tensor& tensor2) {
   if (tensor1.dim() != tensor2.dim()) {
     return false;
@@ -203,7 +242,6 @@ inline bool prepare_cmsis_pool2d_config(
     int64_t activation_min,
     int64_t activation_max,
     CmsisPool2DConfig& config,
-    bool require_channels_last = true,
     bool allow_ceil_mode = false) {
   if (input.dim() != 4 || output.dim() != 4) {
     ET_LOG(Error, "%s: tensors must be 4-D", op_name);
@@ -218,7 +256,11 @@ inline bool prepare_cmsis_pool2d_config(
     return false;
   }
 
-  if (input.size(0) != output.size(0) || input.size(1) != output.size(1)) {
+  // Read dims under either the legacy channels_last or the dim-order-free
+  // contiguous contract (both physically NHWC for CMSIS-NN).
+  const NhwcDims in_dims = read_nhwc_dims(input);
+  const NhwcDims out_dims = read_nhwc_dims(output);
+  if (in_dims.n != out_dims.n || in_dims.c != out_dims.c) {
     ET_LOG(
         Error,
         "%s: batch and channel dimensions must match between input and output",
@@ -227,15 +269,6 @@ inline bool prepare_cmsis_pool2d_config(
     return false;
   }
 
-  if (require_channels_last) {
-    if (!is_channels_last_tensor(input) || !is_channels_last_tensor(output)) {
-      ET_LOG(
-          Error, "%s: tensors must use channels_last dimension order", op_name);
-      context.fail(Error::InvalidArgument);
-      return false;
-    }
-  }
-
   auto check_tuple_len = [&](const Int64ArrayRef& arr,
                              const char* name) -> bool {
     if (arr.size() != 2) {
@@ -314,17 +347,17 @@ inline bool prepare_cmsis_pool2d_config(
 
   int32_t batch, channels, input_h, input_w, output_h, output_w;
   if (!check_int32_within_range(
-          context, op_name, input.size(0), "input batch", batch) ||
+          context, op_name, in_dims.n, "input batch", batch) ||
       !check_int32_within_range(
-          context, op_name, input.size(1), "input channels", channels) ||
+          context, op_name, in_dims.h, "input height", input_h) ||
       !check_int32_within_range(
-          context, op_name, input.size(2), "input height", input_h) ||
+          context, op_name, in_dims.w, "input width", input_w) ||
       !check_int32_within_range(
-          context, op_name, input.size(3), "input width", input_w) ||
+          context, op_name, in_dims.c, "input channels", channels) ||
       !check_int32_within_range(
-          context, op_name, output.size(2), "output height", output_h) ||
+          context, op_name, out_dims.h, "output height", output_h) ||
       !check_int32_within_range(
-          context, op_name, output.size(3), "output width", output_w)) {
+          context, op_name, out_dims.w, "output width", output_w)) {
     return false;
   }
 

backends/cortex_m/ops/op_quantized_avg_pool2d.cpp

@@ -97,8 +97,7 @@ Tensor& quantized_avg_pool2d_out(
           activation_min,
           activation_max,
           pool_config,
-          true,
-          true)) {
+          /*allow_ceil_mode=*/true)) {
     return out;
   }
 

backends/cortex_m/ops/operators.py

@@ -1317,19 +1317,23 @@ def quantized_avg_pool2d_meta(
     kernel = _ensure_tuple2(kernel_size)
     stride_vals = _ensure_tuple2(stride)
     padding_vals = _ensure_tuple2(padding)
+    # Input is NHWC-shaped [N, H, W, C] and plain contiguous. Derive the output
+    # spatial dims via the NCHW pooling helper, then emit an NHWC-shaped output.
+    n, h, w, c = input.shape
+    nchw = torch.empty((n, c, h, w), dtype=torch.float, device=input.device)
     output = F.avg_pool2d(
-        input.to(torch.float),
+        nchw,
         kernel,
         stride=stride_vals,
         padding=padding_vals,
         ceil_mode=ceil_mode,
         count_include_pad=False,
     )
+    out_n, out_c, out_h, out_w = output.shape
     return torch.empty(
-        output.shape,
+        (out_n, out_h, out_w, out_c),
         dtype=torch.int8,
         device=input.device,
-        memory_format=torch.channels_last,
     )
 
 
@@ -1346,6 +1350,8 @@ def quantized_avg_pool2d_impl(
     scratch: torch.Tensor,
 ) -> torch.Tensor:
     dequant_input = dequantize_per_tensor_cmsis(input, zero_point, multiplier, shift)
+    # Input is NHWC-shaped [N, H, W, C]; convert to NCHW for the reference pool.
+    dequant_input = dequant_input.permute(0, 3, 1, 2).contiguous()
 
     kernel = _ensure_tuple2(kernel_size)
     stride_vals = _ensure_tuple2(stride)
@@ -1362,7 +1368,8 @@ def quantized_avg_pool2d_impl(
     )
     result = quantize_per_tensor_cmsis(result, zero_point, multiplier, shift)
     output = torch.clamp(result, -128, 127)
-    return output.to(torch.int8)
+    # result is NCHW; return NHWC-shaped output to match the kernel contract.
+    return output.to(torch.int8).permute(0, 2, 3, 1).contiguous()
 
 
 # ===================================================================
@@ -1416,10 +1423,11 @@ def _compute_max_pool2d_output_shape(
     padding: Sequence[int],
     dilation: Sequence[int],
 ) -> torch.Size:
+    # Input is NHWC-shaped [N, H, W, C].
     batch = input_shape[0]
-    channels = input_shape[1]
-    in_height = input_shape[2]
-    in_width = input_shape[3]
+    in_height = input_shape[1]
+    in_width = input_shape[2]
+    channels = input_shape[3]
 
     kernel_height, kernel_width = kernel_size
     stride_h, stride_w = stride
@@ -1432,7 +1440,7 @@ def _compute_max_pool2d_output_shape(
     out_width = (
         in_width + 2 * pad_w - dilation_w * (kernel_width - 1) - 1
     ) // stride_w + 1
-    return torch.Size([batch, channels, out_height, out_width])
+    return torch.Size([batch, out_height, out_width, channels])
 
 
 @register_fake("cortex_m::quantized_max_pool2d")  # type: ignore[misc]
@@ -1453,14 +1461,14 @@ def quantized_max_pool2d_meta(
     padding_vals = _ensure_tuple2(padding)
     dilation_vals = _ensure_tuple2(dilation)
 
+    # Input/output are NHWC-shaped [N, H, W, C] and plain contiguous.
     output_shape = _compute_max_pool2d_output_shape(
         input.shape, kernel, stride_vals, padding_vals, dilation_vals
     )
     return torch.empty(
         output_shape,
         dtype=torch.int8,
         device=input.device,
-        memory_format=torch.channels_last,
     )
 
 
@@ -1495,13 +1503,16 @@ def quantized_max_pool2d_impl(
     if ceil_mode:
         raise RuntimeError("quantized_max_pool2d does not support ceil_mode=True")
 
+    # Input is NHWC-shaped [N, H, W, C]; convert to NCHW for the reference pool.
+    input_nchw = input.permute(0, 3, 1, 2).contiguous()
     result = F.max_pool2d(
-        input,
+        input_nchw,
         kernel,
         stride=stride_vals,
         padding=padding_vals,
         dilation=dilation_vals,
         ceil_mode=ceil_mode,
     )
     result = torch.clamp(result, activation_min, activation_max)
-    return result.to(torch.int8).contiguous(memory_format=torch.channels_last)
+    # result is NCHW; return NHWC-shaped output to match the kernel contract.
+    return result.to(torch.int8).permute(0, 2, 3, 1).contiguous()

backends/cortex_m/passes/aten_to_cortex_m_pass.py

@@ -828,23 +828,32 @@ def _get_avg_pool2d_replacement(
     input_scale = node.meta["input_qparams"][0].scale
     output_mult, output_shift = quantize_multiplier_aot(input_scale)
 
+    # The cortex_m pool op consumes/produces NHWC-shaped contiguous tensors, so
+    # permute the NCHW graph input to NHWC before the op.
+    with node.graph.inserting_before(node):
+        pool_input = node.graph.create_node(
+            "call_function",
+            target=exir_ops.edge.aten.permute_copy.default,
+            args=(input_node, [0, 2, 3, 1]),
+        )
+
     avg_padding = padding
     if count_include_pad:
         pad_h, pad_w = padding
-        input_tensor = get_first_fake_tensor(input_node)
-        pre_pad = post_pad = to_physical_order([0, 0, pad_h, pad_w], input_tensor)
+        # Input is now NHWC [N, H, W, C]; pad H and W in NHWC physical order.
+        pre_pad = post_pad = [0, int(pad_h), int(pad_w), 0]
         with node.graph.inserting_before(node):
-            input_node = node.graph.create_node(
+            pool_input = node.graph.create_node(
                 "call_function",
                 target=exir_ops.edge.cortex_m.pad.default,
-                args=(input_node, pre_pad, post_pad, int(input_zp)),
+                args=(pool_input, pre_pad, post_pad, int(input_zp)),
             )
         avg_padding = [0, 0]
 
     scratch = _create_uninitialized_alloc_node(node, exported_program)
 
     new_args = (
-        input_node,
+        pool_input,
         kernel_size,
         stride,
         avg_padding,
@@ -854,8 +863,19 @@ def _get_avg_pool2d_replacement(
         int(output_shift),
         scratch,
     )
+    with node.graph.inserting_before(node):
+        pool_node = node.graph.create_node(
+            "call_function",
+            target=exir_ops.edge.cortex_m.quantized_avg_pool2d.default,
+            args=new_args,
+        )
+    for key in ("input_qparams", "output_qparams"):
+        if key in node.meta:
+            pool_node.meta[key] = node.meta[key]
+
+    # Permute the NHWC-shaped op output back to NCHW so downstream stays NCHW.
     return DialectNodeSpec(
-        exir_ops.edge.cortex_m.quantized_avg_pool2d.default, new_args
+        exir_ops.edge.aten.permute_copy.default, (pool_node, [0, 3, 1, 2])
     )
 
 
@@ -1090,19 +1110,39 @@ def _get_max_pool2d_replacement(
     if quantized_op is None:
         return None
 
-    args = (
-        node.args[0],
-        kernel_size,
-        stride,
-        padding,
-        dilation,
-        ceil_mode,
-        input_zero_point,
-        output_zero_point,
-        activation_min,
-        activation_max,
+    input_node = cast(Node, node.args[0])
+    # The cortex_m pool op consumes/produces NHWC-shaped contiguous tensors, so
+    # permute the NCHW graph input to NHWC before the op.
+    with node.graph.inserting_before(node):
+        pool_input = node.graph.create_node(
+            "call_function",
+            target=exir_ops.edge.aten.permute_copy.default,
+            args=(input_node, [0, 2, 3, 1]),
+        )
+        pool_node = node.graph.create_node(
+            "call_function",
+            target=quantized_op.default,
+            args=(
+                pool_input,
+                kernel_size,
+                stride,
+                padding,
+                dilation,
+                ceil_mode,
+                input_zero_point,
+                output_zero_point,
+                activation_min,
+                activation_max,
+            ),
+        )
+    for key in ("input_qparams", "output_qparams"):
+        if key in node.meta:
+            pool_node.meta[key] = node.meta[key]
+
+    # Permute the NHWC-shaped op output back to NCHW so downstream stays NCHW.
+    return DialectNodeSpec(
+        exir_ops.edge.aten.permute_copy.default, (pool_node, [0, 3, 1, 2])
     )
-    return DialectNodeSpec(quantized_op.default, args)
 
 
 @AtenToCortexMPass.register_dialect_substitution(exir_ops.edge.aten.minimum.default)

backends/cortex_m/passes/scratch_buffer_sizes.py

@@ -251,10 +251,10 @@ def cmsis_nn_avgpool_buffer_size(
 ) -> list[int]:
     x = cast(torch.fx.Node, pool_node.args[0])
 
-    # Input is NCHW (PyTorch); CMSIS-NN's avgpool buffer sizer only needs the
-    # input channel count and output width.
-    _, c_in, _, _ = _shape_from_node(x)
-    _, _, _, out_w = _shape_from_node(pool_node)
+    # The cortex_m pool op consumes/produces NHWC-shaped [N, H, W, C] tensors.
+    # The avgpool buffer sizer only needs the input channel count and output width.
+    _, _, _, c_in = _shape_from_node(x)
+    _, _, out_w, _ = _shape_from_node(pool_node)
 
     return [
         int(

backends/cortex_m/test/build_test_runner.sh

@@ -43,10 +43,14 @@ join_by_comma() {
 ops_list=(
     aten::add.out
     aten::clamp.out
+    aten::clone.out
     aten::mul.out
     aten::convolution.out
     aten::max_pool2d_with_indices.out
+    # Keep legacy PTEs generated before _skip_dim_order=True loadable.
     dim_order_ops::_clone_dim_order.out
+    dim_order_ops::_empty_dim_order.out
+    dim_order_ops::_to_dim_order_copy.out
     aten::cat.out
     aten::full.out
     aten::ge.Tensor_out