diff --git a/backends/cortex_m/passes/aten_to_cortex_m_pass.py b/backends/cortex_m/passes/aten_to_cortex_m_pass.py index 5799896e17c..122468f04d2 100644 --- a/backends/cortex_m/passes/aten_to_cortex_m_pass.py +++ b/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] diff --git a/backends/cortex_m/test/models/test_ds_cnn.py b/backends/cortex_m/test/models/test_ds_cnn.py index ba5b2da6b78..c179fbaa1a8 100644 --- a/backends/cortex_m/test/models/test_ds_cnn.py +++ b/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, } diff --git a/backends/cortex_m/test/ops/test_avg_pool2d.py b/backends/cortex_m/test/ops/test_avg_pool2d.py index 20a5fc92222..580a5b028df 100644 --- a/backends/cortex_m/test/ops/test_avg_pool2d.py +++ b/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, diff --git a/backends/cortex_m/test/ops/test_conv.py b/backends/cortex_m/test/ops/test_conv.py index f11734faa6c..f2ac0df0cc2 100644 --- a/backends/cortex_m/test/ops/test_conv.py +++ b/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 + )