Cortex-M: size scratch buffers in a dedicated pass

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",

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]

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

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)

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
+    )