From c472672f3fc8d16dd19ef2852003a6ce1fa303fc Mon Sep 17 00:00:00 2001 From: RJ Ascani Date: Mon, 21 Sep 2026 14:55:44 -0700 Subject: [PATCH] [Cortex-M] Support public edge transforms Migrate CortexMPassManager to ExportedProgramPassManager so callers can use edge.transform() and to_edge_transform_and_lower() without replacing private program entries. Preserve the existing lowering sequence and legacy entry point, update callers, and cover program state and serialization in tests. Authored with OpenAI Codex. --- .claude/skills/cortex-m/SKILL.md | 4 +- .../test/all_ops/generate_test_models.py | 4 +- backends/arm/scripts/aot_arm_compiler.py | 11 +- backends/cortex_m/passes/cortex_m_pass.py | 2 +- .../cortex_m/passes/cortex_m_pass_manager.py | 85 +++++---- backends/cortex_m/quantizer/quantizer.py | 2 +- .../cortex_m/test/misc/test_target_config.py | 4 +- backends/cortex_m/test/targets.bzl | 21 +++ .../test/test_explicit_layout_pipeline.py | 23 +-- backends/cortex_m/test/test_pass_manager.py | 167 ++++++++++++++++++ backends/cortex_m/test/tester.py | 42 +++-- backends/test/harness/stages/run_passes.py | 4 +- .../arm-cortex-m/arm-cortex-m-overview.md | 8 +- examples/arduino/export_model.py | 4 +- examples/arm/cortex_m_mv2_example.ipynb | 2 +- .../ethos_u_cmsis_nn_fallback_example.ipynb | 4 +- .../export/export_model.py | 3 +- .../pico2/export_mlp_mnist_cmsis.py | 5 +- 18 files changed, 296 insertions(+), 99 deletions(-) create mode 100644 backends/cortex_m/test/test_pass_manager.py diff --git a/.claude/skills/cortex-m/SKILL.md b/.claude/skills/cortex-m/SKILL.md index d8dcd2a9aff..4eea5db5655 100644 --- a/.claude/skills/cortex-m/SKILL.md +++ b/.claude/skills/cortex-m/SKILL.md @@ -33,10 +33,8 @@ exported = export(quantized, example_inputs) edge = to_edge_transform_and_lower( exported, compile_config=cortex_m_edge_compile_config(), + transform_passes=CortexMPassManager(), ) -edge._edge_programs["forward"] = CortexMPassManager( - edge.exported_program(), CortexMPassManager.pass_list -).transform() et_program = edge.to_executorch() ``` diff --git a/backends/arm/cmsis_pack/test/all_ops/generate_test_models.py b/backends/arm/cmsis_pack/test/all_ops/generate_test_models.py index b9e34bf61da..97ba1f34d1b 100755 --- a/backends/arm/cmsis_pack/test/all_ops/generate_test_models.py +++ b/backends/arm/cmsis_pack/test/all_ops/generate_test_models.py @@ -247,10 +247,8 @@ def _export_cortex_m( _core_aten_ops_exception_list=[torch.ops.aten.max_pool2d.default], ), constant_methods=metadata, + transform_passes=CortexMPassManager(target_config=target_config), ) - edge._edge_programs["forward"] = CortexMPassManager( - edge.exported_program(), target_config=target_config - ).transform() return edge.to_executorch() diff --git a/backends/arm/scripts/aot_arm_compiler.py b/backends/arm/scripts/aot_arm_compiler.py index 9dc09276acf..2298435d137 100644 --- a/backends/arm/scripts/aot_arm_compiler.py +++ b/backends/arm/scripts/aot_arm_compiler.py @@ -983,15 +983,12 @@ def _to_channels_last(x): edge = to_edge_transform_and_lower( exported_program, compile_config=cortex_m_edge_compile_config(), + transform_passes=CortexMPassManager( + target_config=target_config, + use_explicit_layout=args.cortex_m_explicit_layout, + ), ) - pass_manager = CortexMPassManager( - edge.exported_program(), - target_config=target_config, - use_explicit_layout=args.cortex_m_explicit_layout, - ) - edge._edge_programs["forward"] = pass_manager.transform() - return model_quant, edge, example_inputs diff --git a/backends/cortex_m/passes/cortex_m_pass.py b/backends/cortex_m/passes/cortex_m_pass.py index 5b3e98e9a96..b88f11a92c2 100644 --- a/backends/cortex_m/passes/cortex_m_pass.py +++ b/backends/cortex_m/passes/cortex_m_pass.py @@ -13,7 +13,7 @@ class CortexMPass(ExportPass): """Base class for passes that need the Cortex-M target config. Passes that subclass this declare `exported_program` and `target_config` - in their `__init__`; `CortexMPassManager.transform()` injects both + in their `__init__`; `CortexMPassManager` injects both automatically when running the pass list. """ diff --git a/backends/cortex_m/passes/cortex_m_pass_manager.py b/backends/cortex_m/passes/cortex_m_pass_manager.py index 1e007c32cc9..8d44cd9b298 100644 --- a/backends/cortex_m/passes/cortex_m_pass_manager.py +++ b/backends/cortex_m/passes/cortex_m_pass_manager.py @@ -26,8 +26,12 @@ from executorch.backends.transforms.replace_squeeze_unsqueeze_with_view import ( ReplaceSqueezeAndUnsqueezeWithViewPass, ) -from executorch.exir.pass_base import ExportPass -from executorch.exir.pass_manager import PassManager +from executorch.exir.pass_base import ( + ExportedProgramPassBase, + ExportedProgramPassResult, + ExportPass, +) +from executorch.exir.pass_manager import ExportedProgramPassManager, PassType from executorch.exir.program._program import _transform, lift_constant_tensor_pass from torch.export import ExportedProgram @@ -50,7 +54,36 @@ PassClass = Type[ExportPass] -class CortexMPassManager(PassManager): +class _CortexMLoweringPass(ExportedProgramPassBase): + def __init__( + self, pass_classes: list[PassClass], target_config: CortexMTargetConfig + ) -> None: + self.pass_classes = pass_classes + self.target_config = target_config + + def call(self, exported_program: ExportedProgram) -> ExportedProgramPassResult: + modified = False + for pass_cls in self.pass_classes: + signature = inspect.signature(pass_cls) + kwargs: dict[str, Any] = {} + if "exported_program" in signature.parameters: + kwargs["exported_program"] = exported_program + if "target_config" in signature.parameters: + kwargs["target_config"] = self.target_config + + transform_pass = pass_cls(**kwargs) + transformed = _transform(exported_program, transform_pass) + modified |= transformed is not exported_program + exported_program = transformed + + # Passes can introduce tensor attributes that must become program inputs. + buffer_count = len(exported_program.graph_signature.buffers) + exported_program = lift_constant_tensor_pass(exported_program) + modified |= len(exported_program.graph_signature.buffers) != buffer_count + return ExportedProgramPassResult(exported_program, modified) + + +class CortexMPassManager(ExportedProgramPassManager): legacy_pass_list: list[PassClass] = [ # Run before folding so qparams attach to max_pool2d values, not tuple + getitem. RemoveGetItemPass, @@ -98,7 +131,7 @@ class CortexMPassManager(PassManager): def __init__( self, - exported_program: ExportedProgram | None, + exported_program: ExportedProgram | None = None, passes: Optional[list[PassClass]] = None, target_config: Optional[CortexMTargetConfig] = None, use_explicit_layout: bool = False, @@ -106,9 +139,8 @@ def __init__( """Initialize the Cortex-M pass manager. Args: - exported_program: The exported program to transform. Required - before calling ``transform()``; may be ``None`` for callers - that only use ``transform_for_annotation()``. + exported_program: Optional program for the legacy ``transform()`` + entry point. Omit when using ``edge.transform(pass_manager)``. passes: Optional override of the pass list. Defaults to the legacy or explicit-layout pass list selected by ``use_explicit_layout``. @@ -119,20 +151,26 @@ def __init__( use_explicit_layout: Select the experimental explicit-layout pass sequence. Legacy lowering remains the default. """ - 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. default_passes = ( self.explicit_layout_pass_list if use_explicit_layout else self.legacy_pass_list ) - self.passes: list[PassClass] = ( # type: ignore[assignment] - passes if passes is not None else default_passes # type: ignore[assignment] - ) + pass_classes = passes if passes is not None else default_passes + for pass_cls in pass_classes: + if not isinstance(pass_cls, type): + raise ValueError( + f"{type(self).__name__} expects pass classes, not instances; " + f"got {pass_cls!r}" + ) self.target_config: CortexMTargetConfig = target_config or CortexMTargetConfig( cpu=CortexM.M55 ) + lowering_passes: list[PassType] = [ + _CortexMLoweringPass(pass_classes, self.target_config) + ] + super().__init__(lowering_passes) def transform_for_annotation(self, model): passes = self.pass_list_transform_for_annotation @@ -148,24 +186,5 @@ def transform(self) -> ExportedProgram: f"got {exported_program!r}" ) - for pass_cls in self.passes: - if not isinstance(pass_cls, type): - raise ValueError( - f"{type(self).__name__} expects pass classes, not instances; " - f"got {pass_cls!r}" - ) - - signature = inspect.signature(pass_cls) - kwargs: dict[str, Any] = {} - if "exported_program" in signature.parameters: - kwargs["exported_program"] = exported_program - if "target_config" in signature.parameters: - kwargs["target_config"] = self.target_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. - exported_program = lift_constant_tensor_pass(exported_program) - return exported_program + result = self(exported_program) + return result.exported_program if result.modified else exported_program diff --git a/backends/cortex_m/quantizer/quantizer.py b/backends/cortex_m/quantizer/quantizer.py index cd3c5e47d6e..1fa972dd63d 100644 --- a/backends/cortex_m/quantizer/quantizer.py +++ b/backends/cortex_m/quantizer/quantizer.py @@ -117,5 +117,5 @@ def validate(self, model: GraphModule) -> None: return None def transform_for_annotation(self, model: GraphModule) -> GraphModule: - pass_manager = CortexMPassManager(None) + pass_manager = CortexMPassManager() return pass_manager.transform_for_annotation(model) diff --git a/backends/cortex_m/test/misc/test_target_config.py b/backends/cortex_m/test/misc/test_target_config.py index 472d1927886..1d285ebb581 100644 --- a/backends/cortex_m/test/misc/test_target_config.py +++ b/backends/cortex_m/test/misc/test_target_config.py @@ -102,7 +102,7 @@ def test_default_target_config_is_m55(self): CortexMPassManager, ) - pm = CortexMPassManager(exported_program=None) + pm = CortexMPassManager() assert pm.target_config.cpu == CortexM.M55 assert pm.target_config.backend == cmsis_nn.Backend.MVE @@ -112,6 +112,6 @@ def test_explicit_target_config_threaded(self): ) target_config = CortexMTargetConfig(cpu=CortexM.M33) - pm = CortexMPassManager(exported_program=None, target_config=target_config) + pm = CortexMPassManager(target_config=target_config) assert pm.target_config.cpu == CortexM.M33 assert pm.target_config.backend == cmsis_nn.Backend.DSP diff --git a/backends/cortex_m/test/targets.bzl b/backends/cortex_m/test/targets.bzl index e6a8a8ce252..2713566b5e0 100644 --- a/backends/cortex_m/test/targets.bzl +++ b/backends/cortex_m/test/targets.bzl @@ -50,6 +50,7 @@ def define_common_targets(is_fbcode = False): "//executorch/backends/cortex_m/passes:cortex_passes", "//executorch/backends/cortex_m/quantizer:quantizer", "//executorch/backends/test/harness:tester", + "//executorch/exir:lib", ], ) @@ -101,6 +102,26 @@ def define_common_targets(is_fbcode = False): ], ) + python_pytest( + name = "test_pass_manager", + srcs = ["test_pass_manager.py"], + compile = "with-source", + typing = False, + deps = [ + "//caffe2:torch", + "//pytorch/ao:torchao", # @manual + "//executorch/backends/cortex_m:edge_compile_config", + "//executorch/backends/cortex_m:target_config", + "//executorch/backends/cortex_m/passes:cortex_passes", + "//executorch/backends/cortex_m/quantizer:quantizer", + "//executorch/exir:lib", + "//executorch/exir:pass_base", + "//executorch/exir/_serialize:lib", + "//executorch/exir/dialects:lib", + "fbsource//third-party/pypi/pytest:pytest", + ], + ) + python_pytest( diff --git a/backends/cortex_m/test/test_explicit_layout_pipeline.py b/backends/cortex_m/test/test_explicit_layout_pipeline.py index ce01ceba8b9..b02dead803c 100644 --- a/backends/cortex_m/test/test_explicit_layout_pipeline.py +++ b/backends/cortex_m/test/test_explicit_layout_pipeline.py @@ -4,15 +4,12 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -from functools import partial - import pytest import torch -from executorch.backends.cortex_m.passes.cortex_m_pass_manager import CortexMPassManager from executorch.backends.cortex_m.quantizer.quantizer import CortexMQuantizer from executorch.backends.cortex_m.target_config import CortexM, CortexMTargetConfig -from executorch.backends.cortex_m.test.tester import CortexMTester -from executorch.backends.test.harness.stages import Quantize, RunPasses, StageType +from executorch.backends.cortex_m.test.tester import CortexMRunPasses, CortexMTester +from executorch.backends.test.harness.stages import Quantize, StageType from executorch.exir.dialects._ops import ops as exir_ops from torch.fx import Node @@ -75,13 +72,9 @@ def _count(exported_program, target) -> int: def _run_explicit_layout_pass_manager(tester: CortexMTester) -> CortexMTester: target_config = CortexMTargetConfig(cpu=CortexM.M55) tester.run_passes( - RunPasses( - partial( - CortexMPassManager, - target_config=target_config, - use_explicit_layout=True, - ), # type: ignore[arg-type] - CortexMPassManager.explicit_layout_pass_list, # type: ignore[arg-type] + CortexMRunPasses( + target_config=target_config, + use_explicit_layout=True, ) ) return tester @@ -216,5 +209,7 @@ def test_explicit_layout_rejects_unsupported_spatial_operator(): with pytest.raises(Exception) as caught: _run_explicit_layout_passes(tester) - assert caught.value.__cause__ is not None - assert "NHWC-eligible" in str(caught.value.__cause__) + error = caught.value + while error.__cause__ is not None: + error = error.__cause__ + assert "NHWC-eligible" in str(error) diff --git a/backends/cortex_m/test/test_pass_manager.py b/backends/cortex_m/test/test_pass_manager.py new file mode 100644 index 00000000000..b60ecb59e30 --- /dev/null +++ b/backends/cortex_m/test/test_pass_manager.py @@ -0,0 +1,167 @@ +# 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. + +import pytest +import torch +from executorch.backends.cortex_m.edge_compile_config import ( + cortex_m_edge_compile_config, +) +from executorch.backends.cortex_m.passes.cortex_m_pass_manager import CortexMPassManager +from executorch.backends.cortex_m.quantizer.quantizer import CortexMQuantizer +from executorch.backends.cortex_m.target_config import CortexM, CortexMTargetConfig +from executorch.exir import to_edge, to_edge_transform_and_lower +from executorch.exir._serialize._program import deserialize_pte_binary +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import ExportPass +from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e + + +def _quantize(model, inputs, use_explicit_layout=False): + prepared = prepare_pt2e( + torch.export.export(model.eval(), inputs).module(), + CortexMQuantizer(use_explicit_layout=use_explicit_layout), + ) + with torch.no_grad(): + prepared(*inputs) + quantized = convert_pt2e(prepared) + expected = quantized(*inputs) + output_scale = [ + node.args[1] + for node in quantized.graph.nodes + if node.target == torch.ops.quantized_decomposed.dequantize_per_tensor.default + ][-1] + return torch.export.export(quantized, inputs), expected, output_scale + + +@pytest.mark.parametrize("use_explicit_layout", [False, True]) +@pytest.mark.parametrize("entry_point", ["transform", "combined", "legacy"]) +def test_lowers_and_serializes(entry_point, use_explicit_layout): + model = torch.nn.Conv2d(3, 4, 3, padding=1) + inputs = (torch.randn(1, 3, 8, 8),) + if not use_explicit_layout: + model = model.to(memory_format=torch.channels_last) + inputs = (inputs[0].to(memory_format=torch.channels_last),) + exported, expected, output_scale = _quantize(model, inputs, use_explicit_layout) + config = cortex_m_edge_compile_config() + manager = CortexMPassManager(use_explicit_layout=use_explicit_layout) + + if entry_point == "combined": + lowered = to_edge_transform_and_lower( + exported, compile_config=config, transform_passes=manager + ) + else: + lowered = to_edge(exported, compile_config=config) + if entry_point == "transform": + lowered = lowered.transform(manager) + else: + manager = CortexMPassManager( + lowered.exported_program(), use_explicit_layout=use_explicit_layout + ) + lowered._edge_programs["forward"] = manager.transform() + + program = lowered.exported_program() + program.validate() + conv = ( + exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default + if use_explicit_layout + else exir_ops.edge.cortex_m.quantized_conv2d.default + ) + assert sum(node.target == conv for node in program.graph.nodes) == 1 + assert all(node.op != "get_attr" for node in program.graph.nodes) + torch.testing.assert_close( + program.module()(*inputs), expected, rtol=0, atol=2 * output_scale + ) + serialized = deserialize_pte_binary(lowered.to_executorch().buffer).program + assert serialized.execution_plan[0].name == "forward" + assert not serialized.execution_plan[0].delegates + + +def test_transform_uses_each_method_program(): + inputs = { + "forward": (torch.randn(2, 4),), + "other": (torch.randn(1, 8),), + } + models = { + "forward": torch.nn.Linear(4, 3), + "other": torch.nn.Linear(8, 2), + } + programs = {} + expected = {} + output_scales = {} + for name, model in models.items(): + programs[name], expected[name], output_scales[name] = _quantize( + model, inputs[name] + ) + edge = to_edge( + programs, + compile_config=cortex_m_edge_compile_config(), + constant_methods={"version": 17}, + ) + # A legacy bound program must not override the method supplied by transform(). + manager = CortexMPassManager( + edge.exported_program(), target_config=CortexMTargetConfig(cpu=CortexM.M33) + ) + lowered = edge.transform(manager) + assert lowered.methods == {"forward", "other"} + assert lowered.config_methods == {"version"} + for name in models: + program = lowered.exported_program(name) + program.validate() + [linear] = [ + node + for node in program.graph.nodes + if node.target == exir_ops.edge.cortex_m.quantized_linear.default + ] + assert linear.args[3] is None # The M33 does not use the MVE kernel sum. + torch.testing.assert_close( + program.module()(*inputs[name]), + expected[name], + rtol=0, + atol=2 * output_scales[name], + ) + serialized = deserialize_pte_binary(lowered.to_executorch().buffer).program + assert {plan.name for plan in serialized.execution_plan} == { + "forward", + "other", + "version", + } + + +def test_empty_pass_list_reports_no_change(): + inputs = (torch.tensor([-1.0, 2.0]),) + edge = to_edge(torch.export.export(torch.nn.ReLU(), inputs)) + program = edge.exported_program() + manager = CortexMPassManager(passes=[]) + assert not manager(program).modified + assert edge.transform(manager).exported_program() is program + assert CortexMPassManager(program, passes=[]).transform() is program + with pytest.raises(ValueError, match="needs a real ExportedProgram"): + manager.transform() + + +def test_empty_pass_list_still_lifts_constants(): + class AddConstant(ExportPass): + def call_operator(self, op, args, kwargs, meta): + result = super().call_operator(op, args, kwargs, meta) + return super().call_operator( + exir_ops.edge.aten.add.Tensor, + (result, torch.tensor([2.0])), + {}, + meta, + ) + + inputs = (torch.tensor([-1.0, 2.0]),) + edge = to_edge(torch.export.export(torch.nn.ReLU(), inputs)) + program = edge.transform([AddConstant()]).exported_program() + assert any(node.op == "get_attr" for node in program.graph.nodes) + + result = CortexMPassManager(passes=[])(program) + assert result.modified + lowered = result.exported_program + lowered.validate() + assert all(node.op != "get_attr" for node in lowered.graph.nodes) + assert len(lowered.graph_signature.buffers) == 1 + torch.testing.assert_close(lowered.module()(*inputs), torch.tensor([2.0, 4.0])) diff --git a/backends/cortex_m/test/tester.py b/backends/cortex_m/test/tester.py index b644db4e6c1..c0dbeab1e98 100644 --- a/backends/cortex_m/test/tester.py +++ b/backends/cortex_m/test/tester.py @@ -6,7 +6,6 @@ from collections.abc import Callable from dataclasses import dataclass -from functools import partial from typing import Any, Optional import torch @@ -28,6 +27,8 @@ ToEdgeTransformAndLower, ToExecutorch, ) +from executorch.exir import EdgeProgramManager, to_edge_transform_and_lower +from torch.export import ExportedProgram class CortexMQuantize(Quantize): @@ -42,31 +43,36 @@ def __init__(self): class CortexMRunPasses(RunPasses): - def __init__(self, target_config: Optional[CortexMTargetConfig] = None): - target_config = target_config or CortexMTargetConfig(cpu=CortexM.M55) - # The base RunPasses constructs the pass manager as `cls(ep, pass_list)`. - # Pre-bind the target_config so it flows through that 2-arg call. - super().__init__( - partial(CortexMPassManager, target_config=target_config), # type: ignore[arg-type] - CortexMPassManager.pass_list, # type: ignore[arg-type] + def __init__( + self, + target_config: Optional[CortexMTargetConfig] = None, + use_explicit_layout: bool = False, + ): + super().__init__(CortexMPassManager) + self.pass_manager = CortexMPassManager( + target_config=target_config, + use_explicit_layout=use_explicit_layout, ) + def run(self, artifact: EdgeProgramManager | ExportedProgram, inputs=None) -> None: + if isinstance(artifact, EdgeProgramManager): + self.edge_or_aten_program = artifact.transform(self.pass_manager) + else: + self.edge_or_aten_program = self.pass_manager(artifact).exported_program -class CortexMToEdgeTransformAndLower(ToEdgeTransformAndLower): - """to_edge with no partitioner, then CortexMPassManager. - - Cortex-M rewrites edge operators in place rather than delegating a subgraph, - so this is its equivalent of to_edge_transform_and_lower, which is the only - lowering entry point the shared backend test suite drives. - """ +class CortexMToEdgeTransformAndLower(ToEdgeTransformAndLower): def __init__(self, target_config: Optional[CortexMTargetConfig] = None): super().__init__(edge_compile_config=cortex_m_edge_compile_config()) - self._run_passes = CortexMRunPasses(target_config) + self.pass_manager = CortexMPassManager(target_config=target_config) def run(self, artifact, inputs=None, generate_etrecord: bool = False) -> None: - super().run(artifact, inputs, generate_etrecord=generate_etrecord) - self._run_passes.run(self.edge_dialect_program, inputs) # type: ignore[arg-type] + self.edge_dialect_program = to_edge_transform_and_lower( + artifact, + compile_config=self.edge_compile_conf, + transform_passes=self.pass_manager, + generate_etrecord=generate_etrecord, + ) class CortexMSerialize(Serialize): diff --git a/backends/test/harness/stages/run_passes.py b/backends/test/harness/stages/run_passes.py index b72c40d5337..c59b083acd1 100644 --- a/backends/test/harness/stages/run_passes.py +++ b/backends/test/harness/stages/run_passes.py @@ -17,7 +17,9 @@ def __init__( self.pass_manager_cls = pass_manager_cls self.pass_list = pass_list self.pass_functions = pass_functions - self.edge_or_aten_program = None + self.edge_or_aten_program: Optional[ + Union[EdgeProgramManager, ExportedProgram] + ] = None def stage_type(self) -> StageType: return StageType.RUN_PASSES diff --git a/docs/source/backends/arm-cortex-m/arm-cortex-m-overview.md b/docs/source/backends/arm-cortex-m/arm-cortex-m-overview.md index 47eb2fceeb7..ab58461f254 100644 --- a/docs/source/backends/arm-cortex-m/arm-cortex-m-overview.md +++ b/docs/source/backends/arm-cortex-m/arm-cortex-m-overview.md @@ -105,7 +105,7 @@ to use the same memory format as the export example. ### 2. Lower to edge and apply Cortex-M passes -Lower to the edge dialect with the backend's `EdgeCompileConfig`, then run the `CortexMPassManager` to replace quantized subgraphs with CMSIS-NN operator implementations: +Lower to the edge dialect with the backend's `EdgeCompileConfig`, then apply `CortexMPassManager` through `EdgeProgramManager.transform()` to replace quantized subgraphs with CMSIS-NN operator implementations: ```python from executorch.exir import ExecutorchBackendConfig, to_edge @@ -122,10 +122,12 @@ config = cortex_m_edge_compile_config() edge_program_manager = to_edge(quantized_exported_program, compile_config=config) -pass_manager = CortexMPassManager(edge_program_manager.exported_program()) -edge_program_manager._edge_programs["forward"] = pass_manager.transform() +edge_program_manager = edge_program_manager.transform(CortexMPassManager()) ``` +Alternatively, pass `transform_passes=CortexMPassManager()` to +`to_edge_transform_and_lower()` with the same compile configuration. + ### 3. Serialize to .pte ```python diff --git a/examples/arduino/export_model.py b/examples/arduino/export_model.py index 688f58dfecc..8c5bb750843 100644 --- a/examples/arduino/export_model.py +++ b/examples/arduino/export_model.py @@ -161,11 +161,9 @@ def export_model( compile_config=cortex_m_edge_compile_config(), ) pm = CortexMPassManager( - edge.exported_program(), - CortexMPassManager.pass_list, target_config=CortexMTargetConfig(cpu=cpu), ) - edge._edge_programs["forward"] = pm.transform() + edge = edge.transform(pm) et = edge.to_executorch() return et.buffer diff --git a/examples/arm/cortex_m_mv2_example.ipynb b/examples/arm/cortex_m_mv2_example.ipynb index 36844b4e5fd..fb4bb6de270 100644 --- a/examples/arm/cortex_m_mv2_example.ipynb +++ b/examples/arm/cortex_m_mv2_example.ipynb @@ -91,7 +91,7 @@ "execution_count": null, "metadata": {}, "outputs": [], - "source": "import os\nfrom executorch.exir import (\n EdgeCompileConfig,\n ExecutorchBackendConfig,\n to_edge,\n)\nfrom executorch.backends.cortex_m.passes.cortex_m_pass_manager import CortexMPassManager\n\n# Create compile config for Cortex-M lowering\nconfig = EdgeCompileConfig(\n preserve_ops=[\n torch.ops.aten.linear.default,\n torch.ops.aten.hardsigmoid.default,\n torch.ops.aten.hardsigmoid_.default,\n torch.ops.aten.hardswish.default,\n torch.ops.aten.hardswish_.default,\n ],\n _check_ir_validity=False,\n _core_aten_ops_exception_list=[torch.ops.aten.max_pool2d.default],\n )\n\n# Lower the exported program for the Cortex-M backend - note to_edge usage rather than to_edge_transform_and_lower, currently required to use preserve_ops w/o partitioner.\nedge_program_manager = to_edge(\n quantized_exported_program,\n compile_config=config,\n )\n\n# Run pass manager on the forward graph_module - use of pass_manager.transform() over edge_program_mangager.transform() is currently required to ensure that the passes can modify the exported_program and not only the graph_module.\npass_manager = CortexMPassManager(edge_program_manager.exported_program())\nedge_program_manager._edge_programs[\"forward\"] = pass_manager.transform()\n\n# Test converted edge program running python implementations of cortex-m dialect.\noutput_data = edge_program_manager.exported_program().module()(transformed_img)\nprint_result(output_data, label)\n\n" + "source": "import os\nfrom executorch.exir import (\n EdgeCompileConfig,\n ExecutorchBackendConfig,\n to_edge,\n)\nfrom executorch.backends.cortex_m.passes.cortex_m_pass_manager import CortexMPassManager\n\n# Create compile config for Cortex-M lowering\nconfig = EdgeCompileConfig(\n preserve_ops=[\n torch.ops.aten.linear.default,\n torch.ops.aten.hardsigmoid.default,\n torch.ops.aten.hardsigmoid_.default,\n torch.ops.aten.hardswish.default,\n torch.ops.aten.hardswish_.default,\n ],\n _check_ir_validity=False,\n _core_aten_ops_exception_list=[torch.ops.aten.max_pool2d.default],\n )\n\n# Lower the exported program to the edge dialect.\nedge_program_manager = to_edge(\n quantized_exported_program,\n compile_config=config,\n )\n\nedge_program_manager = edge_program_manager.transform(CortexMPassManager())\n\n# Test converted edge program running python implementations of cortex-m dialect.\noutput_data = edge_program_manager.exported_program().module()(transformed_img)\nprint_result(output_data, label)\n\n" }, { "cell_type": "code", diff --git a/examples/arm/ethos_u_cmsis_nn_fallback_example.ipynb b/examples/arm/ethos_u_cmsis_nn_fallback_example.ipynb index 0dd8f7045fb..ec113b5d2d1 100644 --- a/examples/arm/ethos_u_cmsis_nn_fallback_example.ipynb +++ b/examples/arm/ethos_u_cmsis_nn_fallback_example.ipynb @@ -160,9 +160,7 @@ "outputs": [], "source": [ "\n", - "edge_program_manager._edge_programs[\"forward\"] = CortexMPassManager(\n", - " edge_program_manager.exported_program()\n", - ").transform()\n", + "edge_program_manager = edge_program_manager.transform(CortexMPassManager())\n", "\n", "executorch_program = edge_program_manager.to_executorch(\n", " config=ExecutorchBackendConfig(extract_delegate_segments=False)\n", diff --git a/examples/arm/person_detection_example_cortex_m/export/export_model.py b/examples/arm/person_detection_example_cortex_m/export/export_model.py index 774034569a9..615791af89a 100644 --- a/examples/arm/person_detection_example_cortex_m/export/export_model.py +++ b/examples/arm/person_detection_example_cortex_m/export/export_model.py @@ -103,11 +103,10 @@ def export_cortex_m(model: MicroYolo): passes=[QuantizeInputs(edge, [0]), QuantizeOutputs(edge, [0])] ) pass_manager = CortexMPassManager( - edge.exported_program(), target_config=CortexMTargetConfig(cpu=CortexM.M55), use_explicit_layout=True, ) - edge._edge_programs["forward"] = pass_manager.transform() + edge = edge.transform(pass_manager) return edge diff --git a/examples/raspberry_pi/pico2/export_mlp_mnist_cmsis.py b/examples/raspberry_pi/pico2/export_mlp_mnist_cmsis.py index 9d24048fed6..7f917391738 100644 --- a/examples/raspberry_pi/pico2/export_mlp_mnist_cmsis.py +++ b/examples/raspberry_pi/pico2/export_mlp_mnist_cmsis.py @@ -96,12 +96,9 @@ def export_to_pte(quantized_model, example_input, output_path: str): logger.info("Applying Cortex-M optimization passes...") pass_manager = CortexMPassManager( - edge_program.exported_program(), target_config=CortexMTargetConfig(cpu=CortexM.M33), ) - transformed_ep = pass_manager.transform() - - edge_program = to_edge(transformed_ep, compile_config=edge_config) + edge_program = edge_program.transform(pass_manager) logger.info("Converting to ExecuTorch format...") exec_program = edge_program.to_executorch(