Cortex-M: avoid redundant convolution and pooling work

backends/cortex_m/passes/aten_to_cortex_m_pass.py

@@ -526,6 +526,17 @@ def _get_convolution_replacement(
     in_channels = param_weight_tensor.shape[1] * groups
     out_channels = param_weight_tensor.shape[0]
     is_depthwise = (in_channels == groups) and (out_channels % in_channels == 0)
+    # CMSIS-NN MVE already repacks these weights and runs regular convolution on
+    # every inference. Emit that layout at export to avoid repeated repacking
+    # and its scratch storage. The >8 limit covers both compiler thresholds.
+    assert isinstance(dialect_pass, AtenToCortexMPass)
+    if (
+        is_depthwise
+        and dialect_pass.target_config.backend == cmsis_nn.Backend.MVE
+        and in_channels == 1
+        and out_channels > 8
+    ):
+        is_depthwise = False
 
     # Only use DW path if batch_size==1, as CMSIS-NN DW falls back to
     # unoptimized implementation otherwise.
@@ -880,7 +891,7 @@ def _get_avg_pool2d_replacement(
     output_mult, output_shift = quantize_multiplier_aot(input_scale)
 
     avg_padding = padding
-    if count_include_pad:
+    if count_include_pad and any(padding):
         pad_h, pad_w = padding
         if explicit_nhwc:
             pre_pad = post_pad = [0, pad_h, pad_w, 0]

backends/cortex_m/test/models/test_ds_cnn.py

@@ -23,11 +23,11 @@
 ops_after_transforms: dict[str, int] = {
     "executorch_exir_dialects_edge__ops_aten_view_copy_default": 1,
     "executorch_exir_dialects_edge__ops_cortex_m_dequantize_per_tensor_default": 3,
-    "executorch_exir_dialects_edge__ops_cortex_m_pad_default": 1,
+    "executorch_exir_dialects_edge__ops_cortex_m_pad_default": 0,
     "executorch_exir_dialects_edge__ops_cortex_m_quantize_per_tensor_default": 3,
     "executorch_exir_dialects_edge__ops_cortex_m_quantized_avg_pool2d_default": 1,
-    "executorch_exir_dialects_edge__ops_cortex_m_quantized_conv2d_default": 4,
-    "executorch_exir_dialects_edge__ops_cortex_m_quantized_depthwise_conv2d_default": 5,
+    "executorch_exir_dialects_edge__ops_cortex_m_quantized_conv2d_default": 5,
+    "executorch_exir_dialects_edge__ops_cortex_m_quantized_depthwise_conv2d_default": 4,
     "executorch_exir_dialects_edge__ops_cortex_m_quantized_linear_default": 1,
 }
 

backends/cortex_m/test/ops/test_avg_pool2d.py

@@ -52,6 +52,10 @@ def forward(self, x):  # noqa: D102
 
 # Prepare test cases: simple 2x2 pool on 4x4, and 3x3 stride 1 on 3x3
 test_cases = {
+    "avgpool_no_padding_included": McuTestCase(
+        CortexMAvgPool2d(kernel_size=2, stride=2, count_include_pad=True),
+        (ramp_tensor(-5, 5, (1, 2, 4, 4)).to(memory_format=torch.channels_last),),
+    ),
     "avgpool_2x2": McuTestCase(
         CortexMAvgPool2d(kernel_size=2, stride=2), (ramp_tensor(0, 15, (1, 1, 4, 4)),)
     ),
@@ -115,7 +119,7 @@ def test_dialect_avg_pool2d(test_case, cortex_m_target):
         test_case.model, test_case.example_inputs, target_config=cortex_m_target
     )
     ops_after = dict(test_case.model.ops_after_transforms)
-    if test_case.model.pool.count_include_pad:
+    if test_case.model.pool.count_include_pad and test_case.model.pool.padding != 0:
         ops_after["executorch_exir_dialects_edge__ops_cortex_m_pad_default"] = 1
     tester.test_dialect(
         test_case.model.ops_before_transforms,

backends/cortex_m/test/ops/test_conv.py

@@ -4,8 +4,11 @@
 # LICENSE file in the root directory of this source tree.
 
 
+import pytest
 import torch
 from executorch.backends.arm.test.common import parametrize, xfail_type
+from executorch.backends.cortex_m.library import cmsis_nn
+from executorch.backends.cortex_m.target_config import CortexM, CortexMTargetConfig
 from executorch.backends.cortex_m.test.tester import (
     CortexMTester,
     McuTestCase,
@@ -389,3 +392,33 @@ def test_grouped_conv2d_bias_is_populated(test_case, cortex_m_target):
         if n.op == "call_function" and n.target in grouped_convs
     ]
     assert conv_node.args[2] is not None
+
+
+@pytest.mark.parametrize(
+    "backend", [cmsis_nn.Backend.SCALAR, cmsis_nn.Backend.DSP, cmsis_nn.Backend.MVE]
+)
+@pytest.mark.parametrize("out_channels", [1, 8, 9, 64])
+def test_dialect_single_channel_dispatch(backend, out_channels):
+    model = torch.nn.Conv2d(1, out_channels, 3, bias=False).eval()
+    inputs = (torch.randn(1, 1, 8, 8).to(memory_format=torch.channels_last),)
+    config = CortexMTargetConfig(cpu=CortexM.M55, isa=backend)
+    tester = CortexMTester(model, inputs, target_config=config)
+    tester.quantize().export().to_edge().run_passes()
+    tester.run_method_and_compare_outputs(inputs=inputs, qtol=1)
+    graph = tester.get_artifact(StageType.RUN_PASSES).exported_program().graph
+    regular = backend == cmsis_nn.Backend.MVE and out_channels > 8
+    expected = (
+        exir_ops.edge.cortex_m.quantized_conv2d.default
+        if regular
+        else exir_ops.edge.cortex_m.quantized_depthwise_conv2d.default
+    )
+    assert sum(node.target == expected for node in graph.nodes) == 1
+
+
+@pytest.mark.parametrize("out_channels", [1, 8, 9, 64])
+def test_implementation_single_channel_dispatch(out_channels, cortex_m_target):
+    model = torch.nn.Conv2d(1, out_channels, 3, bias=False).eval()
+    inputs = (torch.randn(1, 1, 8, 8).to(memory_format=torch.channels_last),)
+    CortexMTester(model, inputs, target_config=cortex_m_target).test_implementation(
+        qtol=1
+    )