From c62e2b54538e985915af76a1ededde851deee032 Mon Sep 17 00:00:00 2001 From: Erik Lundell Date: Tue, 14 Apr 2026 10:31:16 +0200 Subject: [PATCH] Cortex-M backend: Add CortexMConfiguration The configuration mirrors cmsis_nn's CortexM enum, with an added ANY value, and a .backend convenience function. Add the configuration to CortexM's pass manager as an optional parameter that defaults to ANY. If a pass wants it, pass it (analogous to the current exported_program handling). Signed-off-by: Erik Lundell Change-Id: I7c4fac9fca8c58672ca344ecfbc240ac17fdf063 --- backends/cortex_m/passes/__init__.py | 1 + .../cortex_m/passes/cortex_m_configuration.py | 37 ++++++++++++++ backends/cortex_m/passes/cortex_m_pass.py | 32 ++++++++++++ .../cortex_m/passes/cortex_m_pass_manager.py | 51 ++++++++++++++----- 4 files changed, 107 insertions(+), 14 deletions(-) create mode 100644 backends/cortex_m/passes/cortex_m_configuration.py create mode 100644 backends/cortex_m/passes/cortex_m_pass.py diff --git a/backends/cortex_m/passes/__init__.py b/backends/cortex_m/passes/__init__.py index b1b67add413..c5b4fa249ca 100644 --- a/backends/cortex_m/passes/__init__.py +++ b/backends/cortex_m/passes/__init__.py @@ -36,6 +36,7 @@ def _ensure_cortex_m_dependencies() -> None: from .activation_fusion_pass import ActivationFusionPass # noqa from .clamp_hardswish_pass import ClampHardswishPass # noqa from .convert_to_cortex_m_pass import ConvertToCortexMPass # noqa +from .cortex_m_configuration import CortexMConfiguration # noqa from .decompose_hardswish_pass import DecomposeHardswishPass # noqa from .decompose_mean_pass import DecomposeMeanPass # noqa from .quantized_clamp_activation_pass import QuantizedClampActivationPass # noqa diff --git a/backends/cortex_m/passes/cortex_m_configuration.py b/backends/cortex_m/passes/cortex_m_configuration.py new file mode 100644 index 00000000000..4aa98809de0 --- /dev/null +++ b/backends/cortex_m/passes/cortex_m_configuration.py @@ -0,0 +1,37 @@ +# Copyright 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. + +from enum import auto, Enum +from typing import ClassVar, Mapping + +import cmsis_nn # type: ignore[import-not-found, import-untyped] + + +class CortexMConfiguration(Enum): + M0 = auto() + M0PLUS = auto() + M3 = auto() + M4 = auto() + M7 = auto() + M23 = auto() + M33 = auto() + M35P = auto() + M55 = auto() + M85 = auto() + ANY = auto() # Guaranteed to work on any Cortex-M. + __members__: ClassVar[Mapping[str, "CortexMConfiguration"]] + + @property + def backend(self) -> cmsis_nn.Backend: + if self == CortexMConfiguration.ANY: + # Currently, MVE is all we support. We can just return the MVE backend. + return cmsis_nn.Backend.MVE + + cmsis_nn_cortex_m = cmsis_nn.CortexM.__members__.get(self.name, None) + if cmsis_nn_cortex_m is None: + raise ValueError( + f"CortexM configuration {self.name} is not supported by cmsis_nn." + ) + return cmsis_nn.resolve_backend(cmsis_nn_cortex_m) diff --git a/backends/cortex_m/passes/cortex_m_pass.py b/backends/cortex_m/passes/cortex_m_pass.py new file mode 100644 index 00000000000..6fad2f2b431 --- /dev/null +++ b/backends/cortex_m/passes/cortex_m_pass.py @@ -0,0 +1,32 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# Copyright 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. + +from executorch.exir.pass_base import ExportPass +from torch.export import ExportedProgram + +from .cortex_m_configuration import CortexMConfiguration + + +class CortexMPass(ExportPass): + """ + An abstract interface for CortexM backend passes. + """ + + def __init__( + self, exported_program: ExportedProgram, cortex_m_config: CortexMConfiguration + ) -> None: + super().__init__() + self._exported_program = exported_program + self._cortex_m_config = cortex_m_config + + @property + def exported_program(self) -> ExportedProgram: + return self._exported_program + + @property + def cortex_m_config(self) -> CortexMConfiguration: + return self._cortex_m_config diff --git a/backends/cortex_m/passes/cortex_m_pass_manager.py b/backends/cortex_m/passes/cortex_m_pass_manager.py index 074eb6118d0..987ae4a84e0 100644 --- a/backends/cortex_m/passes/cortex_m_pass_manager.py +++ b/backends/cortex_m/passes/cortex_m_pass_manager.py @@ -5,7 +5,9 @@ import inspect -from typing import Callable, cast, Optional, Type +from typing import Any, Optional, Type + +import cmsis_nn # type: ignore[import-not-found, import-untyped] from executorch.backends.arm._passes import ( FoldAndAnnotateQParamsPass, @@ -19,13 +21,11 @@ from executorch.exir.pass_manager import PassManager from executorch.exir.program._program import _transform, lift_constant_tensor_pass from torch.export import ExportedProgram -from torch.fx.passes.infra.pass_base import PassResult - -from torch.nn import Module from .activation_fusion_pass import ActivationFusionPass from .clamp_hardswish_pass import ClampHardswishPass from .convert_to_cortex_m_pass import ConvertToCortexMPass +from .cortex_m_configuration import CortexMConfiguration from .decompose_hardswish_pass import DecomposeHardswishPass from .decompose_mean_pass import DecomposeMeanPass from .quantized_clamp_activation_pass import QuantizedClampActivationPass @@ -57,34 +57,57 @@ class CortexMPassManager(PassManager): ] def __init__( - self, exported_program, passes: Optional[list[PassClass]] = None + self, + exported_program: ExportedProgram | None, + passes: Optional[list[PassClass]] = None, + cortex_m: CortexMConfiguration = CortexMConfiguration.ANY, ) -> None: super().__init__(passes=[]) self.exported_program = exported_program + self.cortex_m_config = cortex_m + if self.cortex_m_config.backend != cmsis_nn.Backend.MVE: + raise NotImplementedError( + "Currently, the Cortex-M pass manager only supports MVE." + f"Got {self.cortex_m_config.name} with {self.cortex_m_config.backend.name}" + ) + # 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] ) def transform_for_annotation(self, model): + passes = self.pass_list_transform_for_annotation for p in passes: model = p().call(model).graph_module return model def transform(self) -> ExportedProgram: - ep = self.exported_program + exported_program = self.exported_program + if not isinstance(exported_program, ExportedProgram): + raise ValueError( + f"{self.__class__.__name__} needs an exported_program to run transform, got {exported_program=}" + ) + for pass_cls in self.passes: + if not isinstance(pass_cls, type): + raise ValueError( + f"{self.__class__.__name__} can't have instansiated passes in pass list, got {pass_cls}." + ) + signature = inspect.signature(pass_cls) + kwargs: dict[str, Any] = {} if "exported_program" in signature.parameters: - ep_pass_ctor = cast(Callable[[ExportedProgram], ExportPass], pass_cls) - transform_pass = ep_pass_ctor(ep) - else: - transform_pass = pass_cls() - pass_callable = cast(Callable[[Module], PassResult], transform_pass) - ep = _transform(ep, pass_callable) + kwargs["exported_program"] = exported_program + if "cortex_m_config" in signature.parameters: + kwargs["cortex_m_config"] = self.cortex_m_config + + transform_pass = pass_cls(**kwargs) + exported_program = _transform(exported_program, transform_pass) # All constant tensors should be lifted to buffers at this point, re-run # lift_constant_tensor_pass in case new ones have been introduced by the passes above. - ep = lift_constant_tensor_pass(ep) - return ep + exported_program = lift_constant_tensor_pass(exported_program) + + return exported_program