diff --git a/backends/cortex_m/passes/BUCK b/backends/cortex_m/passes/BUCK index 4d758703cb4..d2bda57a491 100644 --- a/backends/cortex_m/passes/BUCK +++ b/backends/cortex_m/passes/BUCK @@ -36,6 +36,7 @@ fbcode_target(_kind = runtime.python_library, "decompose_hardswish_pass.py", "decompose_mean_pass.py", "explicit_layout_pass.py", + "initialize_scratch_buffers_pass.py", "matmul_to_bmm_pass.py", "quantized_clamp_activation_pass.py", "scratch_buffer_sizes.py", 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..435d7bc2d8f 100644 --- a/backends/cortex_m/passes/aten_to_cortex_m_pass.py +++ b/backends/cortex_m/passes/aten_to_cortex_m_pass.py @@ -11,7 +11,6 @@ import executorch.backends.cortex_m.ops.operators # noqa import executorch.backends.transforms.channels_last_ops # noqa: F401 -import executorch.exir as exir import torch import torch.fx from executorch.backends.arm._passes.arm_pass_utils import get_first_fake_tensor @@ -25,9 +24,6 @@ SHIFT_INT8, to_physical_order, ) -from executorch.backends.cortex_m.passes.scratch_buffer_sizes import ( - required_cmsis_nn_buffer_sizes, -) from executorch.backends.cortex_m.quantizer.quantization_configs import ( CMSIS_SOFTMAX_SCALE, CMSIS_SOFTMAX_ZERO_POINT, @@ -80,33 +76,9 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: raise RuntimeError( f"Cortex-M lowering left {node.target} in the graph." ) - self._initialize_alloc_node_size(node) return PassResult(result.graph_module, result.modified or max_pool_modified) - def _initialize_alloc_node_size(self, node: torch.fx.Node) -> None: - """Initialize trailing scratch alloc nodes for CMSIS-NN kernels.""" - scratch_buffer_sizes = required_cmsis_nn_buffer_sizes( - node, self.target_config.backend - ) - if scratch_buffer_sizes is None: - return - - for i, scratch_buffer_size in enumerate(reversed(scratch_buffer_sizes)): - scratch_arg = node.args[-(i + 1)] - if ( - not isinstance(scratch_arg, torch.fx.Node) - or scratch_arg.target != exir.memory.alloc - ): - raise RuntimeError( - f"Expected scratch alloc node as final argument(s) for {node.target}, got {scratch_arg}." - ) - - scratch_arg.args = (((scratch_buffer_size,), torch.uint8),) - scratch_arg.meta["val"] = torch.empty( - (scratch_buffer_size,), dtype=torch.uint8, device="meta" - ) - def _create_uninitialized_alloc_node( node: Node, exported_program: ExportedProgram @@ -526,6 +498,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 +863,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/passes/cortex_m_pass_manager.py b/backends/cortex_m/passes/cortex_m_pass_manager.py index a15a7262660..6efd551c0a9 100644 --- a/backends/cortex_m/passes/cortex_m_pass_manager.py +++ b/backends/cortex_m/passes/cortex_m_pass_manager.py @@ -41,6 +41,7 @@ CortexMReplaceOpsWithChannelsLastVariants, ValidateCortexMExplicitLayoutPass, ) +from .initialize_scratch_buffers_pass import InitializeScratchBuffersPass from .matmul_to_bmm_pass import MatmulToBmmPass from .quantized_clamp_activation_pass import QuantizedClampActivationPass from .replace_quant_nodes_pass import ReplaceQuantNodesPass @@ -59,6 +60,7 @@ class CortexMPassManager(PassManager): QuantizedClampActivationPass, DecomposeHardswishPass, AtenToCortexMPass, + InitializeScratchBuffersPass, ] explicit_layout_pass_list: list[PassClass] = [ @@ -77,6 +79,7 @@ class CortexMPassManager(PassManager): ValidateCortexMExplicitLayoutPass, ReplaceQuantNodesPass, AtenToCortexMPass, + InitializeScratchBuffersPass, ] pass_list = legacy_pass_list diff --git a/backends/cortex_m/passes/initialize_scratch_buffers_pass.py b/backends/cortex_m/passes/initialize_scratch_buffers_pass.py new file mode 100644 index 00000000000..f3d78ec1b45 --- /dev/null +++ b/backends/cortex_m/passes/initialize_scratch_buffers_pass.py @@ -0,0 +1,37 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# Copyright 2025-2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import torch +from executorch.backends.cortex_m.passes.cortex_m_pass import CortexMPass +from executorch.backends.cortex_m.passes.scratch_buffer_sizes import ( + required_cmsis_nn_buffer_sizes, +) +from executorch.exir.memory import alloc +from executorch.exir.pass_base import PassResult + + +class InitializeScratchBuffersPass(CortexMPass): + def call(self, graph_module: torch.fx.GraphModule) -> PassResult: + modified = False + for node in graph_module.graph.nodes: + sizes = required_cmsis_nn_buffer_sizes(node, self.target_config.backend) + if sizes is None: + continue + for index, size in enumerate(reversed(sizes)): + scratch = node.args[-(index + 1)] + if not isinstance(scratch, torch.fx.Node) or scratch.target != alloc: + raise RuntimeError( + f"Expected scratch alloc node as final argument(s) for {node.target}, got {scratch}." + ) + scratch.args = (((size,), torch.uint8),) + scratch.meta["val"] = torch.empty( + (size,), dtype=torch.uint8, device="meta" + ) + modified = True + if modified: + graph_module.recompile() + return PassResult(graph_module, modified) 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 + )