Cortex-M backend: thread target CPU/ISA through the AOT pass manager

rascani · rascani · commit a2b556e92d9b · 2026-05-11T14:37:25.000-07:00
Introduce a CortexMCompileConfig dataclass (cpu + isa) that carries Cortex-M target information from the --target=cortex-m<variant>+int8 CLI string into CortexMPassManager. The full standard Cortex-M lineup is registered (M0, M0+, M3, M4, M7, M23, M33, M35P, M52, M55, M85), each with a sensible default ISA; the optional-DSP M33/M35P and optional-MVE M52/M55/M85 cases can be expressed via the isa= kwarg. No pass reads the config yet, so this change is purely plumbing — but it positions both the upcoming AOT scratch-buffer sizing work (#16580) and the M0+ (#17646) / M33 (#17644) backend support to plug in without re-plumbing the call site. Actually building for the new variants still requires Phase 2's MPS2 platform glue. CortexMTester gains an optional config kwarg, and the Pico 2 MLP example now constructs CortexMPassManager with cpu='cortex-m33' to match the RP2350 hardware it targets. Authored with Claude.
diff --git a/backends/arm/scripts/aot_arm_compiler.py b/backends/arm/scripts/aot_arm_compiler.py
@@ -33,6 +33,7 @@
 from executorch.backends.arm.util._factory import create_partitioner, create_quantizer
 
 from executorch.backends.arm.vgf import VgfCompileSpec
+from executorch.backends.cortex_m.compile_config import CortexMCompileConfig
 from executorch.backends.cortex_m.passes.cortex_m_pass_manager import CortexMPassManager
 
 from executorch.backends.cortex_m.passes.replace_quant_nodes_pass import (
@@ -465,7 +466,17 @@ def forward(self, x):
     "TOSA-1.0+INT",
     "TOSA-1.0+FP",
     "TOSA-1.0+INT+int16",
+    "cortex-m0+int8",
+    "cortex-m0plus+int8",
+    "cortex-m3+int8",
+    "cortex-m4+int8",
+    "cortex-m7+int8",
+    "cortex-m23+int8",
+    "cortex-m33+int8",
+    "cortex-m35p+int8",
+    "cortex-m52+int8",
     "cortex-m55+int8",
+    "cortex-m85+int8",
 ]
 
 
@@ -566,7 +577,7 @@ def _get_args():
         required=False,
         default="ethos-u55-128",
         choices=TARGETS,
-        help=f"Target backend. For delegated models: Ethos-U/VGF/TOSA variants. For non-delegated: cortex-m55+int8 (CMSIS-NN portable kernels). Valid targets: {TARGETS}",
+        help=f"Target backend. For delegated models: Ethos-U/VGF/TOSA variants. For non-delegated: cortex-m<variant>+int8 (CMSIS-NN portable kernels). Valid targets: {TARGETS}",
     )
     # TODO: Remove --evaluate and --evaluate_config completely after a suitable time.
     # They are deprecated and no longer functional in this script.
@@ -860,9 +871,14 @@ def _to_edge_cortex_m(
     model: GraphModule,
     example_inputs: Tuple[torch.Tensor],
     calibration_samples: Optional[List[Tuple[torch.Tensor, ...]]],
+    config: CortexMCompileConfig,
 ):
     """Cortex-M/CMSIS-NN compilation path with no delegation."""
-    logging.info("Using Cortex-M/CMSIS-NN compilation path (no delegation)")
+    logging.info(
+        "Using Cortex-M/CMSIS-NN compilation path for cpu=%s isa=%s",
+        config.cpu,
+        config.isa,
+    )
 
     def _to_channels_last(x):
         if isinstance(x, torch.Tensor):
@@ -915,7 +931,7 @@ def _to_channels_last(x):
         ),
     )
 
-    pass_manager = CortexMPassManager(edge.exported_program())
+    pass_manager = CortexMPassManager(edge.exported_program(), config=config)
     edge._edge_programs["forward"] = pass_manager.transform()
 
     return model_quant, edge
@@ -1007,12 +1023,14 @@ def main() -> None:  # noqa: C901
     else:
         quant_mode = None
 
-    if args.target == "cortex-m55+int8":
+    if args.target.startswith("cortex-m"):
         # Cortex-M path: CMSIS-NN portable kernels, no delegation
+        cortex_m_config = CortexMCompileConfig.from_target_string(args.target)
         if args.delegate:
             logging.warning(
-                "--delegate is ignored for target 'cortex-m55+int8' "
-                "(this target does not use delegated ops)."
+                "--delegate is ignored for target %r "
+                "(this target does not use delegated ops).",
+                args.target,
             )
             args.delegate = False
         model_quant, edge = _to_edge_cortex_m(
@@ -1021,6 +1039,7 @@ def main() -> None:  # noqa: C901
             model,
             example_inputs,
             calibration_samples,
+            cortex_m_config,
         )
     elif args.delegate:
         # As we can target multiple output encodings, one must
diff --git a/backends/cortex_m/compile_config.py b/backends/cortex_m/compile_config.py
@@ -0,0 +1,98 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Literal
+
+Cpu = Literal[
+    "cortex-m0",
+    "cortex-m0plus",
+    "cortex-m3",
+    "cortex-m4",
+    "cortex-m7",
+    "cortex-m23",
+    "cortex-m33",
+    "cortex-m35p",
+    "cortex-m52",
+    "cortex-m55",
+    "cortex-m85",
+]
+Isa = Literal["scalar", "dsp", "mve"]
+
+# Default ISA per CPU follows the most common configuration each core is
+# shipped with. M33/M35P optionally lack DSP, and M52/M55/M85 optionally
+# lack MVE; callers can pass `isa=` explicitly to override.
+_CPU_DEFAULT_ISA: dict[str, str] = {
+    "cortex-m0": "scalar",
+    "cortex-m0plus": "scalar",
+    "cortex-m3": "scalar",
+    "cortex-m4": "dsp",
+    "cortex-m7": "dsp",
+    "cortex-m23": "scalar",
+    "cortex-m33": "dsp",
+    "cortex-m35p": "dsp",
+    "cortex-m52": "mve",
+    "cortex-m55": "mve",
+    "cortex-m85": "mve",
+}
+
+_SUPPORTED_FEATURES: frozenset[str] = frozenset({"int8"})
+
+
+@dataclass(frozen=True)
+class CortexMCompileConfig:
+    """AOT compile configuration for the Cortex-M backend.
+
+    `cpu` and `isa` are consumed by passes that need to differ by target — most
+    notably any future AOT scratch-buffer sizing — and threaded through the
+    build system as the `-mcpu=` value.
+
+    The current default matches pre-config behavior (M55 + MVE) so callers that
+    don't opt in see no change.
+    """
+
+    cpu: Cpu = "cortex-m55"
+    isa: Isa | None = None
+
+    def __post_init__(self) -> None:
+        if self.cpu not in _CPU_DEFAULT_ISA:
+            raise ValueError(
+                f"Unsupported Cortex-M CPU: {self.cpu!r}. "
+                f"Supported: {sorted(_CPU_DEFAULT_ISA)}"
+            )
+        if self.isa is None:
+            # frozen dataclass: use object.__setattr__ to fill default ISA.
+            object.__setattr__(self, "isa", _CPU_DEFAULT_ISA[self.cpu])
+
+    @classmethod
+    def from_target_string(cls, target: str) -> CortexMCompileConfig:
+        """Parse `cortex-m<variant>+int8` strings used by `aot_arm_compiler.py`.
+
+        Today only `+int8` is supported. The suffix is required so the target
+        string remains explicit about the data type contract.
+        """
+        cpu, sep, features = target.partition("+")
+        if not sep:
+            raise ValueError(
+                f"Cortex-M target string must include a feature suffix "
+                f"(e.g. '+int8'), got: {target!r}"
+            )
+        feature_set = set(features.split("+"))
+        unknown = feature_set - _SUPPORTED_FEATURES
+        if unknown or "int8" not in feature_set:
+            raise ValueError(
+                f"Cortex-M target string must be '<cpu>+int8' "
+                f"(supported features: {sorted(_SUPPORTED_FEATURES)}), "
+                f"got: {target!r}"
+            )
+        if cpu not in _CPU_DEFAULT_ISA:
+            raise ValueError(
+                f"Unsupported Cortex-M CPU in target string: {cpu!r}. "
+                f"Supported: {sorted(_CPU_DEFAULT_ISA)}"
+            )
+        return cls(cpu=cpu)  # type: ignore[arg-type]
diff --git a/backends/cortex_m/passes/cortex_m_pass_manager.py b/backends/cortex_m/passes/cortex_m_pass_manager.py
@@ -11,6 +11,7 @@
     FoldAndAnnotateQParamsPass,
     ScalarsToAttributePass,
 )
+from executorch.backends.cortex_m.compile_config import CortexMCompileConfig
 from executorch.backends.transforms.remove_getitem_op import RemoveGetItemPass
 from executorch.backends.transforms.replace_scalar_with_tensor import (
     ReplaceScalarWithTensorArgPass,
@@ -57,14 +58,18 @@ class CortexMPassManager(PassManager):
     ]
 
     def __init__(
-        self, exported_program, passes: Optional[list[PassClass]] = None
+        self,
+        exported_program,
+        passes: Optional[list[PassClass]] = None,
+        config: Optional[CortexMCompileConfig] = None,
     ) -> None:
         super().__init__(passes=[])
         self.exported_program = exported_program
         # PassManager.passes is typed as callables; this manager stores pass classes which are initialized at transform time with the exported_program.
         self.passes: list[PassClass] = (  # type: ignore[assignment]
             passes if passes is not None else self.pass_list  # type: ignore[assignment]
         )
+        self.config: CortexMCompileConfig = config or CortexMCompileConfig()
 
     def transform_for_annotation(self, model):
         passes = self.pass_list_transform_for_annotation
diff --git a/backends/cortex_m/test/misc/test_compile_config.py b/backends/cortex_m/test/misc/test_compile_config.py
@@ -0,0 +1,103 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from importlib.util import find_spec
+
+import pytest
+
+from executorch.backends.cortex_m.compile_config import CortexMCompileConfig
+
+_HAS_CMSIS_NN = find_spec("cmsis_nn") is not None
+
+
+class TestCortexMCompileConfig:
+    def test_default_is_m55_mve(self):
+        config = CortexMCompileConfig()
+        assert config.cpu == "cortex-m55"
+        assert config.isa == "mve"
+
+    @pytest.mark.parametrize(
+        "target_string,expected_cpu,expected_isa",
+        [
+            ("cortex-m0+int8", "cortex-m0", "scalar"),
+            ("cortex-m0plus+int8", "cortex-m0plus", "scalar"),
+            ("cortex-m3+int8", "cortex-m3", "scalar"),
+            ("cortex-m4+int8", "cortex-m4", "dsp"),
+            ("cortex-m7+int8", "cortex-m7", "dsp"),
+            ("cortex-m23+int8", "cortex-m23", "scalar"),
+            ("cortex-m33+int8", "cortex-m33", "dsp"),
+            ("cortex-m35p+int8", "cortex-m35p", "dsp"),
+            ("cortex-m52+int8", "cortex-m52", "mve"),
+            ("cortex-m55+int8", "cortex-m55", "mve"),
+            ("cortex-m85+int8", "cortex-m85", "mve"),
+        ],
+    )
+    def test_from_target_string(self, target_string, expected_cpu, expected_isa):
+        config = CortexMCompileConfig.from_target_string(target_string)
+        assert config.cpu == expected_cpu
+        assert config.isa == expected_isa
+
+    def test_from_target_string_rejects_unknown_cpu(self):
+        with pytest.raises(ValueError, match="cortex-m999"):
+            CortexMCompileConfig.from_target_string("cortex-m999+int8")
+
+    @pytest.mark.parametrize(
+        "target_string",
+        [
+            "cortex-m55",  # missing feature suffix
+            "cortex-m55+int8+int16",  # unsupported extra feature
+            "cortex-m55+",  # trailing plus
+            "cortex-m55+fp16",  # unknown feature
+        ],
+    )
+    def test_from_target_string_rejects_invalid_features(self, target_string):
+        with pytest.raises(ValueError):
+            CortexMCompileConfig.from_target_string(target_string)
+
+    def test_default_matches_m55_target_string(self):
+        # Regression guard: pre-Phase-1 behavior was M55+MVE; the default
+        # constructor must remain equivalent to parsing the existing target.
+        assert CortexMCompileConfig() == CortexMCompileConfig.from_target_string(
+            "cortex-m55+int8"
+        )
+
+    def test_is_hashable_and_frozen(self):
+        from dataclasses import FrozenInstanceError
+
+        config = CortexMCompileConfig(cpu="cortex-m33")
+        assert hash(config) == hash(CortexMCompileConfig(cpu="cortex-m33"))
+        assert {config, CortexMCompileConfig(cpu="cortex-m33")} == {config}
+        with pytest.raises(FrozenInstanceError):
+            config.cpu = "cortex-m55"  # type: ignore[misc]
+
+    def test_explicit_isa_override(self):
+        config = CortexMCompileConfig(cpu="cortex-m33", isa="scalar")
+        assert config.cpu == "cortex-m33"
+        assert config.isa == "scalar"
+
+
+@pytest.mark.skipif(
+    not _HAS_CMSIS_NN, reason="cortex_m passes require cmsis_nn"
+)
+class TestPassManagerConfigWiring:
+    def test_default_config_is_m55(self):
+        from executorch.backends.cortex_m.passes.cortex_m_pass_manager import (
+            CortexMPassManager,
+        )
+
+        pm = CortexMPassManager(exported_program=None)
+        assert pm.config.cpu == "cortex-m55"
+        assert pm.config.isa == "mve"
+
+    def test_explicit_config_threaded(self):
+        from executorch.backends.cortex_m.passes.cortex_m_pass_manager import (
+            CortexMPassManager,
+        )
+
+        config = CortexMCompileConfig(cpu="cortex-m33")
+        pm = CortexMPassManager(exported_program=None, config=config)
+        assert pm.config.cpu == "cortex-m33"
+        assert pm.config.isa == "dsp"
diff --git a/backends/cortex_m/test/tester.py b/backends/cortex_m/test/tester.py
@@ -6,11 +6,13 @@
 
 from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Any
+from functools import partial
+from typing import Any, Optional
 
 import torch
 from executorch.backends.arm.test.common import get_u55_compile_spec
 from executorch.backends.arm.test.tester.arm_tester import Serialize
+from executorch.backends.cortex_m.compile_config import CortexMCompileConfig
 from executorch.backends.cortex_m.passes.cortex_m_pass_manager import CortexMPassManager
 from executorch.backends.cortex_m.quantizer.quantizer import CortexMQuantizer
 from executorch.backends.test.harness import Tester as TesterBase
@@ -48,9 +50,12 @@ def __init__(self):
 
 
 class CortexMRunPasses(RunPasses):
-    def __init__(self):
+    def __init__(self, config: Optional[CortexMCompileConfig] = None):
+        config = config or CortexMCompileConfig()
+        # The base RunPasses constructs the pass manager as `cls(ep, pass_list)`.
+        # Pre-bind the config so it flows through that 2-arg call.
         super().__init__(
-            CortexMPassManager,
+            partial(CortexMPassManager, config=config),  # type: ignore[arg-type]
             CortexMPassManager.pass_list,
         )
 
@@ -73,12 +78,20 @@ def __init__(self):
 
 
 class CortexMTester(TesterBase):
-    def __init__(self, module, example_inputs):
+    def __init__(
+        self,
+        module,
+        example_inputs,
+        config: Optional[CortexMCompileConfig] = None,
+    ):
         if callable(example_inputs):
             resolved_example_inputs = example_inputs()
         else:
             resolved_example_inputs = example_inputs
-        super().__init__(module, resolved_example_inputs, cortex_m_stage_classes)
+        config = config or CortexMCompileConfig()
+        stage_classes = dict(cortex_m_stage_classes)
+        stage_classes[StageType.RUN_PASSES] = lambda: CortexMRunPasses(config=config)
+        super().__init__(module, resolved_example_inputs, stage_classes)
 
     def test_dialect(
         self,
diff --git a/examples/raspberry_pi/pico2/export_mlp_mnist_cmsis.py b/examples/raspberry_pi/pico2/export_mlp_mnist_cmsis.py
@@ -24,6 +24,7 @@
 
 import torch
 
+from executorch.backends.cortex_m.compile_config import CortexMCompileConfig
 from executorch.backends.cortex_m.passes.cortex_m_pass_manager import CortexMPassManager
 from executorch.backends.cortex_m.quantizer.quantizer import CortexMQuantizer
 from executorch.exir import EdgeCompileConfig, ExecutorchBackendConfig, to_edge
@@ -94,7 +95,10 @@ def export_to_pte(quantized_model, example_input, output_path: str):
     logger.info("Edge program created")
 
     logger.info("Applying Cortex-M optimization passes...")
-    pass_manager = CortexMPassManager(edge_program.exported_program())
+    pass_manager = CortexMPassManager(
+        edge_program.exported_program(),
+        config=CortexMCompileConfig(cpu="cortex-m33"),
+    )
     transformed_ep = pass_manager.transform()
 
     edge_program = to_edge(transformed_ep, compile_config=edge_config)
