From 2d0fde292e3c74c7abd959c62d7314cd8c63dc50 Mon Sep 17 00:00:00 2001 From: RJ Ascani Date: Wed, 19 Aug 2026 21:51:07 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- .../test/all_ops/generate_test_models.py | 36 +- backends/arm/scripts/aot_arm_compiler.py | 36 +- backends/cortex_m/edge_compile_config.py | 5 +- backends/cortex_m/passes/BUCK | 5 + .../cortex_m/passes/aten_to_cortex_m_pass.py | 86 ++- .../cortex_m/passes/cortex_m_pass_manager.py | 39 +- .../cortex_m/passes/explicit_layout_pass.py | 142 +++++ .../cortex_m/quantizer/pattern_checkers.py | 50 +- backends/cortex_m/quantizer/quantizer.py | 31 +- .../cortex_m/quantizer/quantizer_support.py | 24 + .../cortex_m/test/misc/test_target_config.py | 13 + backends/cortex_m/test/targets.bzl | 32 + .../cortex_m/test/test_explicit_layout.py | 559 ++++++++++++++++++ backends/cortex_m/test/tester.py | 67 ++- 14 files changed, 1039 insertions(+), 86 deletions(-) create mode 100644 backends/cortex_m/passes/explicit_layout_pass.py create mode 100644 backends/cortex_m/test/test_explicit_layout.py 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..513b58808b9 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 @@ -14,7 +14,6 @@ Coverage is reconciled up front: every discovered component must have either a recipe or an explicit skip reason, otherwise the run fails (no silent gaps). - """ from __future__ import annotations @@ -44,7 +43,8 @@ def _assert_executorch_from_source(source_dir: Path) -> None: - """Fail fast if the imported ``executorch`` is not the one under source_dir. + """Fail fast if the imported ``executorch`` is not the one under + source_dir. The pack ships C++ kernels copied straight from the repo tree, but the test models are exported by importing ``executorch`` as a Python package. If that @@ -54,7 +54,6 @@ def _assert_executorch_from_source(source_dir: Path) -> None: tensor) while the pack registers the CURRENT arity -- so the op links but fails at runtime with a KernelCall arity mismatch. Catch that skew here instead of on the FVP. - """ import executorch # noqa: PLC0415 @@ -169,7 +168,6 @@ def _compute_test_threshold(actual, expected): the other multi-output recipes return tuples, and the comparison and logical recipes return bool tensors, which cannot be subtracted. Both are compared leaf-wise in float space, so an exact match yields 0.0/0.0. - """ a_leaves = actual if isinstance(actual, (tuple, list)) else (actual,) e_leaves = expected if isinstance(expected, (tuple, list)) else (expected,) @@ -197,12 +195,15 @@ def _export_cortex_m( display_metadata: bool = False, target_core: str = "m55", ) -> ExecutorchProgramManager: + 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 EdgeCompileConfig, to_edge_transform_and_lower + from executorch.exir import to_edge_transform_and_lower from torchao.quantization.pt2e import move_exported_model_to_eval from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e @@ -235,17 +236,7 @@ def _export_cortex_m( edge = to_edge_transform_and_lower( exported, - compile_config=EdgeCompileConfig( - preserve_ops=[ - torch.ops.aten.linear.default, - torch.ops.aten.hardsigmoid.default, - torch.ops.aten.hardsigmoid_.default, - torch.ops.aten.hardswish.default, - torch.ops.aten.hardswish_.default, - ], - _check_ir_validity=False, - _core_aten_ops_exception_list=[torch.ops.aten.max_pool2d.default], - ), + compile_config=cortex_m_edge_compile_config(), constant_methods=metadata, ) edge._edge_programs["forward"] = CortexMPassManager( @@ -272,7 +263,6 @@ def _export_ethos_u( all, so they cannot be exported for any target -- as are the few in op_recipes.ETHOS_U_SKIPS, which this flow cannot lower for the documented backend reasons but which still run on the CPU variants. - """ from executorch.backends.arm.ethosu import EthosUCompileSpec, EthosUPartitioner from executorch.backends.arm.quantizer import ( @@ -341,7 +331,6 @@ def _strip_guards_fn(gm) -> None: The Arm annotation/decomposition passes iterate over the graph and reject unexpected ``call_module`` nodes (``DecomposeSelectScatterPass: call_module is not supported``), so remove it before lowering to the Ethos-U delegate. - """ changed = False for node in list(gm.graph.nodes): @@ -355,9 +344,8 @@ def _strip_guards_fn(gm) -> None: def _is_delegated(program: ExecutorchProgramManager) -> bool: """True if the forward graph contains an Ethos-U delegate call, i.e. Vela - took at least part of the graph onto the NPU (the rest, if any, stays on the - host core). - """ + took at least part of the graph onto the NPU (the rest, if any, stays on + the host core).""" return any( node.op == "call_function" and "executorch_call_delegate" in str(node.target) for node in program.exported_program("forward").graph_module.graph.nodes @@ -374,7 +362,6 @@ def _assert_kernel_present( quantize/dequantize (this happened for the conv recipes when their inputs were channel-first). Require the op's own cortex_m:: kernel in the final forward graph so such fallbacks fail loudly at export time. - """ if category != "Cortex-M": return @@ -423,9 +410,8 @@ def _assert_arity_matches_schema( shipped schema -- the signature-skew symptom of exporting with a stale executorch (see _assert_executorch_from_source). - Such a .pte links against the pack but fails at runtime with a KernelCall - arity mismatch. - + Such a .pte links against the pack but fails at runtime with a + KernelCall arity mismatch. """ if category != "Cortex-M": return diff --git a/backends/arm/scripts/aot_arm_compiler.py b/backends/arm/scripts/aot_arm_compiler.py index 250606f9a4b..a9c8241cac3 100644 --- a/backends/arm/scripts/aot_arm_compiler.py +++ b/backends/arm/scripts/aot_arm_compiler.py @@ -156,7 +156,6 @@ def _load_python_module_model( """Load a model and inputs from a Python source file. The file must define `ModelUnderTest` and `ModelInputs` attributes. - """ if not model_name.endswith(".py"): return None @@ -223,7 +222,6 @@ def get_model_and_inputs_from_name( Raises: RuntimeError: If the model cannot be resolved or required inputs are missing. - """ example_inputs = _load_example_inputs(model_input) @@ -346,8 +344,7 @@ def quantize( calibration_samples: Optional[List[Tuple[torch.Tensor, ...]]] = None, ) -> GraphModule: """This is the official recommended flow for quantization in pytorch 2.0 - export. - """ + export.""" logging.info("Quantizing Model...") logging.debug(f"Original model: {model}") @@ -626,6 +623,14 @@ def _get_args(): choices=TARGETS, help=f"Target backend. For delegated models: Ethos-U/VGF/TOSA variants. For non-delegated: cortex-m (CMSIS-NN portable kernels). Valid targets: {TARGETS}", ) + parser.add_argument( + "--cortex_m_explicit_layout", + action="store_true", + help=( + "Use explicit NCHW/NHWC permutes for Cortex-M instead of dim-order " + "operators. This is an experimental Cortex-M-only option." + ), + ) # TODO: Remove --evaluate and --evaluate_config completely after a suitable time. # They are deprecated and no longer functional in this script. parser.add_argument( @@ -921,9 +926,11 @@ def _to_edge_cortex_m( target_config: CortexMTargetConfig, ): """Cortex-M/CMSIS-NN compilation path with no delegation.""" + use_explicit_layout = args.cortex_m_explicit_layout logging.info( f"Using Cortex-M/CMSIS-NN compilation path for cpu={target_config.cpu.name} " - f"backend={target_config.backend.name}" + f"backend={target_config.backend.name} " + f"layout={'explicit' if use_explicit_layout else 'dim-order'}" ) def _to_channels_last(x): @@ -949,17 +956,20 @@ def _to_channels_last(x): ) model_quant = None else: - model = model.to(memory_format=torch.channels_last) # type: ignore[call-overload] - example_inputs = tuple(_to_channels_last(x) for x in example_inputs) + if not use_explicit_layout: + model = model.to(memory_format=torch.channels_last) # type: ignore[call-overload] + example_inputs = tuple(_to_channels_last(x) for x in example_inputs) - quantizer = CortexMQuantizer() + quantizer = CortexMQuantizer(use_explicit_layout=use_explicit_layout) prepared = prepare_pt2e(model, quantizer) if calibration_samples is None: calibration_samples = [example_inputs] for sample in calibration_samples: - prepared(*tuple(_to_channels_last(x) for x in sample)) + if not use_explicit_layout: + sample = tuple(_to_channels_last(x) for x in sample) + prepared(*sample) model_quant = convert_pt2e(prepared) @@ -969,11 +979,15 @@ def _to_channels_last(x): edge = to_edge_transform_and_lower( exported_program, - compile_config=cortex_m_edge_compile_config(), + compile_config=cortex_m_edge_compile_config( + use_explicit_layout=use_explicit_layout + ), ) pass_manager = CortexMPassManager( - edge.exported_program(), target_config=target_config + edge.exported_program(), + target_config=target_config, + use_explicit_layout=use_explicit_layout, ) edge._edge_programs["forward"] = pass_manager.transform() diff --git a/backends/cortex_m/edge_compile_config.py b/backends/cortex_m/edge_compile_config.py index c691e3cf97f..948b2942ada 100644 --- a/backends/cortex_m/edge_compile_config.py +++ b/backends/cortex_m/edge_compile_config.py @@ -21,7 +21,9 @@ ) -def cortex_m_edge_compile_config() -> EdgeCompileConfig: +def cortex_m_edge_compile_config( + use_explicit_layout: bool = False, +) -> EdgeCompileConfig: """The to_edge configuration the Cortex-M backend requires. Shared by the AOT compiler and the test harness so the two cannot drift: an @@ -37,4 +39,5 @@ def cortex_m_edge_compile_config() -> EdgeCompileConfig: return EdgeCompileConfig( preserve_ops=list(_PRESERVE_OPS), _check_ir_validity=False, + _skip_dim_order=use_explicit_layout, ) diff --git a/backends/cortex_m/passes/BUCK b/backends/cortex_m/passes/BUCK index 09dfb5942f4..dadf8979ed4 100644 --- a/backends/cortex_m/passes/BUCK +++ b/backends/cortex_m/passes/BUCK @@ -35,6 +35,7 @@ fbcode_target(_kind = runtime.python_library, "cortex_m_pass_manager.py", "decompose_hardswish_pass.py", "decompose_mean_pass.py", + "explicit_layout_pass.py", "matmul_to_bmm_pass.py", "quantized_clamp_activation_pass.py", ], @@ -48,8 +49,12 @@ fbcode_target(_kind = runtime.python_library, "//executorch/backends/cortex_m/passes:replace_quant_nodes_pass", "//executorch/backends/cortex_m/passes:scratch_buffer_sizes", "//executorch/backends/transforms:aten_to_dialect_pass", + "//executorch/backends/transforms:channels_last_ops", + "//executorch/backends/transforms:convert_conv1d_to_conv2d_pass", "//executorch/backends/transforms:remove_getitem_op", "//executorch/backends/transforms:replace_scalar_with_tensor", + "//executorch/backends/transforms:replace_ops_with_channels_last_variants", + "//executorch/backends/transforms:to_contiguous_channels_last_pass", "//executorch/backends/transforms:utils", "//executorch/exir:lib", "//executorch/exir:pass_base", 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 08484ce3b38..333e76e18b2 100644 --- a/backends/cortex_m/passes/aten_to_cortex_m_pass.py +++ b/backends/cortex_m/passes/aten_to_cortex_m_pass.py @@ -10,6 +10,7 @@ from typing import cast, Optional 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 @@ -441,6 +442,9 @@ def _get_linear_replacement( return DialectNodeSpec(exir_ops.edge.cortex_m.quantized_linear.default, args) +@AtenToCortexMPass.register_dialect_substitution( + exir_ops.edge.channels_last.convolution.default +) @AtenToCortexMPass.register_dialect_substitution(exir_ops.edge.aten.convolution.default) def _get_convolution_replacement( node: Node, dialect_pass: AtenToDialectPass @@ -448,6 +452,8 @@ def _get_convolution_replacement( if not _has_qparams(node): return None + explicit_nhwc = node.target == exir_ops.edge.channels_last.convolution.default + exported_program = dialect_pass.exported_program conv_args = node.args ( @@ -605,7 +611,11 @@ def _get_convolution_replacement( scratch, ) return DialectNodeSpec( - exir_ops.edge.cortex_m.quantized_depthwise_conv2d.default, + ( + exir_ops.edge.cortex_m.quantized_depthwise_conv2d_nhwc.default + if explicit_nhwc + else exir_ops.edge.cortex_m.quantized_depthwise_conv2d.default + ), depthwise_args, ) @@ -627,7 +637,14 @@ def _get_convolution_replacement( output_qmax, scratch, ) - return DialectNodeSpec(exir_ops.edge.cortex_m.quantized_conv2d.default, conv2d_args) + return DialectNodeSpec( + ( + exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default + if explicit_nhwc + else exir_ops.edge.cortex_m.quantized_conv2d.default + ), + conv2d_args, + ) def _get_transpose_conv2d_replacement( @@ -639,6 +656,7 @@ def _get_transpose_conv2d_replacement( if not _has_qparams(node): return None + explicit_nhwc = node.target == exir_ops.edge.channels_last.convolution.default exported_program = dialect_pass.exported_program conv_t_args = node.args ( @@ -751,7 +769,12 @@ def _get_transpose_conv2d_replacement( output_scratch, ) return DialectNodeSpec( - exir_ops.edge.cortex_m.quantized_transpose_conv2d.default, new_args + ( + exir_ops.edge.cortex_m.quantized_transpose_conv2d_nhwc.default + if explicit_nhwc + else exir_ops.edge.cortex_m.quantized_transpose_conv2d.default + ), + new_args, ) @@ -818,12 +841,16 @@ def _get_bmm_replacement( @AtenToCortexMPass.register_dialect_substitution(exir_ops.edge.aten.avg_pool2d.default) +@AtenToCortexMPass.register_dialect_substitution( + exir_ops.edge.channels_last.avg_pool2d.default +) def _get_avg_pool2d_replacement( node: Node, dialect_pass: AtenToDialectPass ) -> DialectNodeSpec | None: if not _has_qparams(node): return None + explicit_nhwc = node.target == exir_ops.edge.channels_last.avg_pool2d.default exported_program = dialect_pass.exported_program pool_args = node.args kernel_size = cast(list[int], pool_args[1]) @@ -844,12 +871,19 @@ def _get_avg_pool2d_replacement( avg_padding = padding if count_include_pad: pad_h, pad_w = padding - input_tensor = get_first_fake_tensor(input_node) - pre_pad = post_pad = to_physical_order([0, 0, pad_h, pad_w], input_tensor) + if explicit_nhwc: + pre_pad = post_pad = [0, pad_h, pad_w, 0] + else: + input_tensor = get_first_fake_tensor(input_node) + pre_pad = post_pad = to_physical_order([0, 0, pad_h, pad_w], input_tensor) with node.graph.inserting_before(node): input_node = node.graph.create_node( "call_function", - target=exir_ops.edge.cortex_m.pad.default, + target=( + exir_ops.edge.cortex_m.pad_nhwc.default + if explicit_nhwc + else exir_ops.edge.cortex_m.pad.default + ), args=(input_node, pre_pad, post_pad, int(input_zp)), ) avg_padding = [0, 0] @@ -868,7 +902,12 @@ def _get_avg_pool2d_replacement( scratch, ) return DialectNodeSpec( - exir_ops.edge.cortex_m.quantized_avg_pool2d.default, new_args + ( + exir_ops.edge.cortex_m.quantized_avg_pool2d_nhwc.default + if explicit_nhwc + else exir_ops.edge.cortex_m.quantized_avg_pool2d.default + ), + new_args, ) @@ -1054,10 +1093,14 @@ def _get_softmax_replacement( @AtenToCortexMPass.register_dialect_substitution(exir_ops.edge.aten.max_pool2d.default) +@AtenToCortexMPass.register_dialect_substitution( + exir_ops.edge.channels_last.max_pool2d.default +) def _get_max_pool2d_replacement( node: Node, dialect_pass: AtenToDialectPass ) -> DialectNodeSpec | None: del dialect_pass + explicit_nhwc = node.target == exir_ops.edge.channels_last.max_pool2d.default input_qparams = node.meta.get("input_qparams", {}).get(0) cortex_m_meta = node.meta.get("custom", {}).get("cortex_m", {}) if input_qparams is None or cortex_m_meta.get("skip_quantized_max_pool2d", False): @@ -1115,6 +1158,12 @@ def _get_max_pool2d_replacement( activation_min, activation_max, ) + if explicit_nhwc: + quantized_op = getattr( + exir_ops.edge.cortex_m, "quantized_max_pool2d_nhwc", None + ) + if quantized_op is None: + return None return DialectNodeSpec(quantized_op.default, args) @@ -1143,6 +1192,9 @@ def _get_maximum_replacement( @AtenToCortexMPass.register_dialect_substitution( exir_ops.edge.aten.permute_copy.default ) +@AtenToCortexMPass.register_dialect_substitution( + exir_ops.edge.channels_last.permute_copy.default +) def _get_permute_replacement( node: Node, dialect_pass: AtenToDialectPass ) -> DialectNodeSpec | None: @@ -1161,10 +1213,14 @@ def _get_permute_replacement( @AtenToCortexMPass.register_dialect_substitution( exir_ops.edge.aten.constant_pad_nd.default ) +@AtenToCortexMPass.register_dialect_substitution( + exir_ops.edge.channels_last.constant_pad_nd.default +) def _get_pad_replacement( node: Node, dialect_pass: AtenToDialectPass ) -> DialectNodeSpec | None: del dialect_pass + explicit_nhwc = node.target == exir_ops.edge.channels_last.constant_pad_nd.default input_qparams = node.meta.get("input_qparams", {}) if not input_qparams: return None @@ -1181,6 +1237,8 @@ def _get_pad_replacement( input_tensor = _get_input_tensor_data(node) rank = len(input_tensor.shape) + if explicit_nhwc and rank != 4: + return None assert 1 <= rank <= 4, f"cortex_m pad: expected rank in [1, 4], got {rank}" n_pairs = len(padding) // 2 assert ( @@ -1194,8 +1252,16 @@ def _get_pad_replacement( pre_pad[dim_4d] = int(padding[2 * i]) post_pad[dim_4d] = int(padding[2 * i + 1]) - pre_pad = to_physical_order(pre_pad, input_tensor) - post_pad = to_physical_order(post_pad, input_tensor) + if not explicit_nhwc: + pre_pad = to_physical_order(pre_pad, input_tensor) + post_pad = to_physical_order(post_pad, input_tensor) args = (node.args[0], pre_pad, post_pad, int(quantized_pad_value)) - return DialectNodeSpec(exir_ops.edge.cortex_m.pad.default, args) + return DialectNodeSpec( + ( + exir_ops.edge.cortex_m.pad_nhwc.default + if explicit_nhwc + else exir_ops.edge.cortex_m.pad.default + ), + args, + ) diff --git a/backends/cortex_m/passes/cortex_m_pass_manager.py b/backends/cortex_m/passes/cortex_m_pass_manager.py index 892baf136ed..d6d8b566079 100644 --- a/backends/cortex_m/passes/cortex_m_pass_manager.py +++ b/backends/cortex_m/passes/cortex_m_pass_manager.py @@ -13,6 +13,9 @@ ScalarsToAttributePass, ) from executorch.backends.cortex_m.target_config import CortexM, CortexMTargetConfig +from executorch.backends.transforms.convert_conv1d_to_conv2d_pass import ( + ConvertConv1dToConv2dPass, +) from executorch.backends.transforms.remove_getitem_op import RemoveGetItemPass from executorch.backends.transforms.replace_scalar_with_tensor import ( ReplaceScalarWithTensorArgPass, @@ -27,6 +30,7 @@ from .clamp_hardswish_pass import ClampHardswishPass from .decompose_hardswish_pass import DecomposeHardswishPass from .decompose_mean_pass import DecomposeMeanPass +from .explicit_layout_pass import CortexMExplicitLayoutPass from .matmul_to_bmm_pass import MatmulToBmmPass from .quantized_clamp_activation_pass import QuantizedClampActivationPass from .replace_quant_nodes_pass import ReplaceQuantNodesPass @@ -35,7 +39,7 @@ class CortexMPassManager(PassManager): - pass_list: list[PassClass] = [ + legacy_pass_list: list[PassClass] = [ # Run before folding so qparams attach to max_pool2d values, not tuple + getitem. RemoveGetItemPass, FoldAndAnnotateQParamsPass, @@ -47,6 +51,22 @@ class CortexMPassManager(PassManager): AtenToCortexMPass, ] + explicit_layout_pass_list: list[PassClass] = [ + # Run before folding so qparams attach to max_pool2d values, not tuple + getitem. + RemoveGetItemPass, + FoldAndAnnotateQParamsPass, + ReplaceScalarWithTensorArgPass, + ActivationFusionPass, + QuantizedClampActivationPass, + DecomposeHardswishPass, + ConvertConv1dToConv2dPass, + CortexMExplicitLayoutPass, + ReplaceQuantNodesPass, + AtenToCortexMPass, + ] + + pass_list = legacy_pass_list + pass_list_transform_for_annotation: list[PassClass] = [ ScalarsToAttributePass, ReplaceScalarWithTensorArgPass, @@ -61,6 +81,7 @@ def __init__( exported_program: ExportedProgram | None, passes: Optional[list[PassClass]] = None, target_config: Optional[CortexMTargetConfig] = None, + use_explicit_layout: bool = False, ) -> None: """Initialize the Cortex-M pass manager. @@ -68,22 +89,30 @@ def __init__( exported_program: The exported program to transform. Required before calling ``transform()``; may be ``None`` for callers that only use ``transform_for_annotation()``. - passes: Optional override of the pass list. Defaults to - ``CortexMPassManager.pass_list``. + passes: Optional override of the pass list. Defaults to the legacy + or explicit-layout pass list selected by ``use_explicit_layout``. target_config: Compilation target for passes that need it. Defaults to ``CortexMTargetConfig(cpu=CortexM.M55)``, which resolves through cmsis_nn to the MVE backend — matching the pre-config historical behaviour. + use_explicit_layout: Run channels-last dialect region formation. + Legacy dim-order 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 self.pass_list # type: ignore[assignment] + passes if passes is not None else default_passes # type: ignore[assignment] ) self.target_config: CortexMTargetConfig = target_config or CortexMTargetConfig( cpu=CortexM.M55 ) + self.use_explicit_layout = use_explicit_layout def transform_for_annotation(self, model): passes = self.pass_list_transform_for_annotation @@ -112,6 +141,8 @@ def transform(self) -> ExportedProgram: kwargs["exported_program"] = exported_program if "target_config" in signature.parameters: kwargs["target_config"] = self.target_config + if "use_explicit_layout" in signature.parameters: + kwargs["use_explicit_layout"] = self.use_explicit_layout transform_pass = pass_cls(**kwargs) exported_program = _transform(exported_program, transform_pass) diff --git a/backends/cortex_m/passes/explicit_layout_pass.py b/backends/cortex_m/passes/explicit_layout_pass.py new file mode 100644 index 00000000000..4b908ebb3f7 --- /dev/null +++ b/backends/cortex_m/passes/explicit_layout_pass.py @@ -0,0 +1,142 @@ +# 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 executorch.backends.transforms.channels_last_ops # noqa: F401 + +import torch + +from executorch.backends.transforms.replace_ops_with_channels_last_variants import ( + ChannelsLastOpSpec, +) +from executorch.backends.transforms.to_contiguous_channels_last_pass import ( + ToContiguousChannelsLastPass, +) +from executorch.exir import ExportedProgram +from executorch.exir.dialects._ops import ops as exir_ops + + +def _is_rank4(node: torch.fx.Node) -> bool: + return len(node.meta["val"].shape) == 4 + + +def _has_input_qparams(node: torch.fx.Node) -> bool: + return bool(node.meta.get("input_qparams")) + + +def _has_per_tensor_qparam(node: torch.fx.Node, key: str, index: int) -> bool: + qparam = node.meta.get(key, {}).get(index) + return qparam is not None and not getattr(qparam, "per_channel", False) + + +def _has_input_and_output_qparams(node: torch.fx.Node) -> bool: + return _has_per_tensor_qparam(node, "input_qparams", 0) and _has_per_tensor_qparam( + node, "output_qparams", 0 + ) + + +def _supports_avg_pool2d(node: torch.fx.Node) -> bool: + divisor_override = node.args[6] if len(node.args) > 6 else None + return ( + _is_rank4(node) + and _has_input_and_output_qparams(node) + and divisor_override is None + ) + + +def _to_pair(value, default: tuple[int, int]) -> tuple[int, int]: + if value is None or value == []: + return default + if isinstance(value, int): + return (value, value) + if isinstance(value, (list, tuple)) and len(value) == 1: + return (int(value[0]), int(value[0])) + if isinstance(value, (list, tuple)) and len(value) == 2: + return (int(value[0]), int(value[1])) + return default + + +def _supports_max_pool2d(node: torch.fx.Node) -> bool: + if not _is_rank4(node) or not _has_per_tensor_qparam(node, "input_qparams", 0): + return False + if ( + node.meta.get("custom", {}) + .get("cortex_m", {}) + .get("skip_quantized_max_pool2d", False) + ): + return False + + dilation = _to_pair(node.args[4] if len(node.args) > 4 else None, (1, 1)) + ceil_mode = bool(node.args[5]) if len(node.args) > 5 else False + if dilation != (1, 1) or ceil_mode: + return False + + input_qparams = node.meta["input_qparams"].get(0) + output_qparams = node.meta.get("output_qparams", {}).get(0) + if input_qparams is None or output_qparams is None: + return input_qparams is not None + return ( + not getattr(output_qparams, "per_channel", False) + and abs(float(input_qparams.scale) - float(output_qparams.scale)) <= 1e-6 + and int(input_qparams.zp) == int(output_qparams.zp) + ) + + +def _can_propagate(node: torch.fx.Node) -> bool: + if node.target == exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default: + return False + if node.target not in { + exir_ops.edge.aten.add.Tensor, + exir_ops.edge.aten.mul.Tensor, + }: + return True + if len(node.args) < 2: + return False + input1, input2 = node.args[:2] + if not isinstance(input1, torch.fx.Node) or not isinstance(input2, torch.fx.Node): + return False + tensor1 = input1.meta.get("val") + tensor2 = input2.meta.get("val") + if tensor1 is None or tensor2 is None: + return False + return tensor1.shape == tensor2.shape or _has_input_and_output_qparams(node) + + +class CortexMExplicitLayoutPass(ToContiguousChannelsLastPass): + """Configure the common explicit-layout pipeline for Cortex-M kernels.""" + + def __init__( + self, + exported_program: ExportedProgram, + strict: bool = False, + ) -> None: + super().__init__( + exported_program, + op_map={ + exir_ops.edge.aten.convolution.default: ChannelsLastOpSpec( + target=exir_ops.edge.channels_last.convolution.default, + input_indices=[0], + output_indices=[0], + filter_fn=lambda node: ( + _is_rank4(node) and _has_input_and_output_qparams(node) + ), + ), + exir_ops.edge.aten.avg_pool2d.default: ChannelsLastOpSpec( + target=exir_ops.edge.channels_last.avg_pool2d.default, + input_indices=[0], + output_indices=[0], + filter_fn=_supports_avg_pool2d, + ), + exir_ops.edge.aten.max_pool2d.default: ChannelsLastOpSpec( + target=exir_ops.edge.channels_last.max_pool2d.default, + input_indices=[0], + output_indices=[0], + filter_fn=_supports_max_pool2d, + ), + }, + can_propagate=_can_propagate, + layout_pad_target=exir_ops.edge.channels_last.constant_pad_nd.default, + strict=strict, + ) diff --git a/backends/cortex_m/quantizer/pattern_checkers.py b/backends/cortex_m/quantizer/pattern_checkers.py index cc89715b537..34d1c0dcd29 100644 --- a/backends/cortex_m/quantizer/pattern_checkers.py +++ b/backends/cortex_m/quantizer/pattern_checkers.py @@ -101,15 +101,11 @@ def check_quantization_config( class CortexMConv2DCheck(PatternCheck): @classmethod def check_pattern(cls, pattern): - """ - Checks that all nodes of the pattern use channels_last memory format. - """ - for node in pattern: - tensor = get_first_fake_tensor(node) - if not is_channels_last(tensor): - return False - - return True + return all( + get_first_fake_tensor(node).dim() == 4 + and is_channels_last(get_first_fake_tensor(node)) + for node in pattern + ) @classmethod def check_quantization_config( @@ -127,6 +123,18 @@ def check_quantization_config( return is_int8 and is_ch_axis_0 +class CortexMExplicitConv2DCheck(CortexMConv2DCheck): + @classmethod + def check_pattern(cls, pattern): + return all(get_first_fake_tensor(node).dim() == 4 for node in pattern) + + +class CortexMExplicitConv1DCheck(CortexMConv2DCheck): + @classmethod + def check_pattern(cls, pattern): + return all(get_first_fake_tensor(node).dim() == 3 for node in pattern) + + class CortexMLinearCheck(PatternCheck): @classmethod def check_quantization_config( @@ -215,8 +223,10 @@ def check_quantization_config( class CortexMConvTranspose2DCheck(PatternCheck): + require_channels_last = True + @classmethod - def _check_node(cls, node: Node) -> bool: + def _check_node(cls, node: Node) -> bool: # noqa: C901 if node is None: return False # Reject if node is None @@ -224,9 +234,10 @@ def _check_node(cls, node: Node) -> bool: if tensor is None: return False # Reject if no tensor found - # REJECT if using NCHW format (we need channels_last/NHWC) - if not is_channels_last(tensor): - return False # Reject NCHW + if tensor.dim() != 4: + return False + if cls.require_channels_last and not is_channels_last(tensor): + return False # For aten.conv_transpose2d.input: # (input, weight, bias, stride, padding, output_padding, groups, dilation) @@ -253,10 +264,9 @@ def _check_node(cls, node: Node) -> bool: def check_pattern(cls, pattern): """ Positive filter function for transpose conv to REJECT: - 1. NCHW memory format (we only support channels_last/NHWC) - 2. Grouped convolutions (groups > 1) - not supported by CMSIS-NN - 3. Non-zero output_padding - not supported by CMSIS-NN - 4. Dilation != 1 - produces incorrect results with CMSIS-NN + 1. Grouped convolutions (groups > 1) - not supported by CMSIS-NN + 2. Non-zero output_padding - not supported by CMSIS-NN + 3. Dilation != 1 - produces incorrect results with CMSIS-NN Returns True to ACCEPT the node, False to REJECT. """ @@ -264,7 +274,7 @@ def check_pattern(cls, pattern): if not cls._check_node(node): return False # REJECT invalid transpose conv - return True # ACCEPT channels_last transpose conv + return True @classmethod def check_quantization_config( @@ -284,6 +294,10 @@ def check_quantization_config( return is_int8 and is_ch_axis_1 +class CortexMExplicitConvTranspose2DCheck(CortexMConvTranspose2DCheck): + require_channels_last = False + + class CortexMAvgPool2DCheck(PatternCheck): @classmethod def check_pattern(cls, pattern): diff --git a/backends/cortex_m/quantizer/quantizer.py b/backends/cortex_m/quantizer/quantizer.py index d3f49114144..e1386511502 100644 --- a/backends/cortex_m/quantizer/quantizer.py +++ b/backends/cortex_m/quantizer/quantizer.py @@ -18,6 +18,11 @@ GlobalNodeFinder, NodeTargetNodeFinder, ) +from executorch.backends.cortex_m.quantizer.pattern_checkers import ( + CortexMExplicitConv1DCheck, + CortexMExplicitConv2DCheck, + CortexMExplicitConvTranspose2DCheck, +) from executorch.backends.cortex_m.quantizer.pattern_matcher import PatternMatcher from executorch.backends.cortex_m.quantizer.quantization_configs import ( INT8_PER_CHANNEL_CONFIG, @@ -25,6 +30,7 @@ ) from executorch.backends.cortex_m.quantizer.quantizer_support import ( __name__ as cortex_m_quantizer_support_module, + CONV1D_OP_PATTERNS, CONV_OP_PATTERNS, CONV_TRANSPOSE_OP_PATTERNS, CORTEX_M_QUANTIZER_SUPPORT_DICT, @@ -46,7 +52,11 @@ def mark_node_as_annotated( class CortexMQuantizer(ComposableQuantizer): - def __init__(self, per_tensor_config: Optional[QuantizationConfig] = None) -> None: + def __init__( + self, + per_tensor_config: Optional[QuantizationConfig] = None, + use_explicit_layout: bool = False, + ) -> None: """Cortex-M PT2E quantizer. Args: @@ -57,20 +67,35 @@ def __init__(self, per_tensor_config: Optional[QuantizationConfig] = None) -> No ``INT8_PER_TENSOR_CONFIG``; pass ``INT16_PER_TENSOR_CONFIG`` to quantize the ops that support it (e.g. ``quantized_div``) with int16 activations. + use_explicit_layout: Allow contiguous NCHW Conv2d and + ConvTranspose2d patterns. Legacy mode still requires + channels-last tensors during quantization. """ per_tensor_config = per_tensor_config or INT8_PER_TENSOR_CONFIG conv_targets: set[OpOverload] = set() - for key in CONV_OP_PATTERNS.keys() | CONV_TRANSPOSE_OP_PATTERNS.keys(): + conv_patterns = CONV_OP_PATTERNS.keys() | CONV_TRANSPOSE_OP_PATTERNS.keys() + if use_explicit_layout: + conv_patterns |= CONV1D_OP_PATTERNS.keys() + for key in conv_patterns: conv_targets.update(key) support_dict_name = ( cortex_m_quantizer_support_module + ".CORTEX_M_QUANTIZER_SUPPORT_DICT" ) + support_dict = dict(CORTEX_M_QUANTIZER_SUPPORT_DICT) + if use_explicit_layout: + for pattern in CONV1D_OP_PATTERNS: + support_dict[pattern] = CortexMExplicitConv1DCheck + for pattern in CONV_OP_PATTERNS: + support_dict[pattern] = CortexMExplicitConv2DCheck + for pattern in CONV_TRANSPOSE_OP_PATTERNS: + support_dict[pattern] = CortexMExplicitConvTranspose2DCheck + pattern_matcher = PatternMatcher( cast( dict[tuple[OpOverload, ...], Optional[type[PatternCheck]]], - CORTEX_M_QUANTIZER_SUPPORT_DICT, + support_dict, ), support_dict_name=support_dict_name, ) diff --git a/backends/cortex_m/quantizer/quantizer_support.py b/backends/cortex_m/quantizer/quantizer_support.py index aaaf6414d06..835a53af7e9 100644 --- a/backends/cortex_m/quantizer/quantizer_support.py +++ b/backends/cortex_m/quantizer/quantizer_support.py @@ -82,6 +82,30 @@ (torch.ops.aten.conv2d.default, torch.ops.aten.clamp_.default): CortexMConv2DCheck, } +CONV1D_OP_PATTERNS = { + (torch.ops.aten.conv1d.default,): CortexMConv2DCheck, + (torch.ops.aten.conv1d.default, torch.ops.aten.relu.default): CortexMConv2DCheck, + (torch.ops.aten.conv1d.default, torch.ops.aten.relu_.default): CortexMConv2DCheck, + ( + torch.ops.aten.conv1d.default, + torch.ops.aten.hardtanh.default, + ): CortexMConv2DCheck, + ( + torch.ops.aten.conv1d.default, + torch.ops.aten.hardtanh_.default, + ): CortexMConv2DCheck, + ( + torch.ops.aten.conv1d.default, + torch.ops.aten.hardsigmoid.default, + ): CortexMConv2DCheck, + ( + torch.ops.aten.conv1d.default, + torch.ops.aten.hardsigmoid_.default, + ): CortexMConv2DCheck, + (torch.ops.aten.conv1d.default, torch.ops.aten.clamp.default): CortexMConv2DCheck, + (torch.ops.aten.conv1d.default, torch.ops.aten.clamp_.default): CortexMConv2DCheck, +} + CONV_TRANSPOSE_OP_PATTERNS = { (torch.ops.aten.conv_transpose2d.input,): CortexMConvTranspose2DCheck, ( diff --git a/backends/cortex_m/test/misc/test_target_config.py b/backends/cortex_m/test/misc/test_target_config.py index 472d1927886..24675cdb89e 100644 --- a/backends/cortex_m/test/misc/test_target_config.py +++ b/backends/cortex_m/test/misc/test_target_config.py @@ -6,6 +6,9 @@ # LICENSE file in the root directory of this source tree. import pytest +from executorch.backends.cortex_m.edge_compile_config import ( + cortex_m_edge_compile_config, +) from executorch.backends.cortex_m.library import cmsis_nn from executorch.backends.cortex_m.target_config import CortexM, CortexMTargetConfig @@ -115,3 +118,13 @@ def test_explicit_target_config_threaded(self): pm = CortexMPassManager(exported_program=None, target_config=target_config) assert pm.target_config.cpu == CortexM.M33 assert pm.target_config.backend == cmsis_nn.Backend.DSP + + def test_explicit_layout_is_opt_in(self): + from executorch.backends.cortex_m.passes.cortex_m_pass_manager import ( + CortexMPassManager, + ) + + assert not cortex_m_edge_compile_config()._skip_dim_order + assert cortex_m_edge_compile_config(use_explicit_layout=True)._skip_dim_order + assert not CortexMPassManager(None).use_explicit_layout + assert CortexMPassManager(None, use_explicit_layout=True).use_explicit_layout diff --git a/backends/cortex_m/test/targets.bzl b/backends/cortex_m/test/targets.bzl index b80f8d3f31c..592abe01d91 100644 --- a/backends/cortex_m/test/targets.bzl +++ b/backends/cortex_m/test/targets.bzl @@ -5,6 +5,7 @@ # LICENSE file in the root directory of this source tree. load("@fbcode_macros//build_defs:python_unittest.bzl", "python_unittest") +load("@fbcode_macros//build_defs:python_library.bzl", "python_library") load("@fbcode_macros//build_defs:python_pytest.bzl", "python_pytest") load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") load("@fbsource//tools/build_defs:platform_defs.bzl", "CXX") @@ -37,6 +38,22 @@ def define_common_targets(is_fbcode = False): define_operator_test_target(op) if is_fbcode: + python_library( + name = "tester", + srcs = ["tester.py"], + deps = [ + "//caffe2:torch", + "//executorch/backends/arm/test:arm_tester", + "//executorch/backends/arm/test:common", + "//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/backends/test/harness:tester", + "//executorch/backends/transforms:duplicate_dynamic_quant_chain", + ], + ) + python_unittest( name = "test_replace_quant_nodes", srcs = [ @@ -67,3 +84,18 @@ def define_common_targets(is_fbcode = False): "fbsource//third-party/pypi/pytest:pytest", ], ) + + python_pytest( + name = "test_explicit_layout_host", + srcs = ["test_explicit_layout.py"], + compile = "with-source", + pytest_cmd_args = ["-k", "not runs_on_fvp"], + typing = False, + deps = [ + "//caffe2:torch", + ":tester", + "//executorch/backends/cortex_m:target_config", + "//executorch/exir/dialects:lib", + "fbsource//third-party/pypi/pytest:pytest", + ], + ) diff --git a/backends/cortex_m/test/test_explicit_layout.py b/backends/cortex_m/test/test_explicit_layout.py new file mode 100644 index 00000000000..50cade0f97f --- /dev/null +++ b/backends/cortex_m/test/test_explicit_layout.py @@ -0,0 +1,559 @@ +# 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 copy + +import torch +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 StageType +from executorch.exir.dialects._ops import ops as exir_ops + + +class Conv2d(torch.nn.Module): + def __init__( + self, + in_channels=3, + out_channels=4, + groups=1, + kernel_size=3, + stride=1, + padding=1, + ): + super().__init__() + self.conv = torch.nn.Conv2d( + in_channels, + out_channels, + kernel_size, + stride=stride, + padding=padding, + groups=groups, + ) + + def forward(self, x): + return self.conv(x) + + +class Conv1d(torch.nn.Module): + def __init__(self): + super().__init__() + self.conv = torch.nn.Conv1d(2, 4, 3, padding=1) + + def forward(self, x): + return self.conv(x) + + +class TwoConv2d(torch.nn.Module): + def __init__(self): + super().__init__() + self.conv1 = torch.nn.Conv2d(3, 4, 3, padding=1) + self.conv2 = torch.nn.Conv2d(4, 5, 3, padding=1) + + def forward(self, x): + return self.conv2(self.conv1(x)) + + +class ConvPoolConv(torch.nn.Module): + def __init__(self, pool): + super().__init__() + self.conv1 = torch.nn.Conv2d(3, 4, 3, padding=1) + self.pool = pool + self.conv2 = torch.nn.Conv2d(4, 5, 3, padding=1) + + def forward(self, x): + return self.conv2(self.pool(self.conv1(x))) + + +class ConvTranspose2d(torch.nn.Module): + def __init__( + self, + in_channels=3, + out_channels=4, + kernel_size=3, + stride=2, + padding=1, + bias=True, + ): + super().__init__() + self.conv = torch.nn.ConvTranspose2d( + in_channels, + out_channels, + kernel_size, + stride=stride, + padding=padding, + bias=bias, + ) + + def forward(self, x): + return self.conv(x) + + +class ConvPadConv(torch.nn.Module): + def __init__(self): + super().__init__() + self.conv1 = torch.nn.Conv2d(3, 4, 3, padding=1) + self.conv2 = torch.nn.Conv2d(4, 5, 3, padding=1) + + def forward(self, x): + return self.conv2(torch.nn.functional.pad(self.conv1(x), (1, 1, 1, 1))) + + +class ConvBiasConv(torch.nn.Module): + def __init__(self): + super().__init__() + self.conv1 = torch.nn.Conv2d(3, 4, 3, padding=1) + self.conv2 = torch.nn.Conv2d(4, 5, 3, padding=1) + self.bias = torch.nn.Parameter(torch.randn(1, 4, 1, 1)) + + def forward(self, x): + return self.conv2(torch.relu(self.conv1(x)) + self.bias) + + +class ConvForkAdd(torch.nn.Module): + def __init__(self): + super().__init__() + self.stem = torch.nn.Conv2d(3, 8, 3, padding=1) + self.branch1 = torch.nn.Conv2d(8, 8, 3, padding=1) + self.branch2 = torch.nn.Conv2d(8, 8, 3, padding=1) + + def forward(self, x): + stem = self.stem(x) + return self.branch1(stem) + self.branch2(stem) + + +class ConvSoftmax(torch.nn.Module): + def __init__(self): + super().__init__() + self.conv = torch.nn.Conv2d(3, 4, 3, padding=1) + + def forward(self, x): + return torch.softmax(self.conv(x), dim=-1) + + +def _lower(module, inputs): + tester = CortexMTester(module, inputs, use_explicit_layout=True) + tester.quantize().export().to_edge() + + edge_program = tester.get_artifact(StageType.TO_EDGE).exported_program() + assert all( + getattr(node.target, "namespace", None) != "dim_order_ops" + for node in edge_program.graph.nodes + if node.op == "call_function" + ) + + tester.run_passes() + tester.run_method_and_compare_outputs(inputs=inputs, qtol=2) + return tester.get_artifact(StageType.RUN_PASSES).exported_program() + + +def _lower_legacy(module, inputs): + tester = CortexMTester(module, inputs) + tester.quantize().export().to_edge().run_passes() + return tester.get_artifact(StageType.RUN_PASSES).exported_program() + + +def _run_explicit_layout_on_fvp(module, inputs, target, qtol=2): + tester = CortexMTester(module, inputs, use_explicit_layout=True) + tester.quantize().export().to_edge().run_passes() + program = tester.get_artifact(StageType.RUN_PASSES).exported_program() + assert _count(program, target) == 1 + + tester.to_executorch().serialize() + tester.run_method_and_compare_outputs(inputs=inputs, qtol=qtol) + + +def _count(exported_program, target): + return sum( + node.op == "call_function" and node.target == target + for node in exported_program.graph.nodes + ) + + +def _planned_buffer_sizes( + module, + inputs, + use_explicit_layout, + target_config, + expected_ops, +): + module = copy.deepcopy(module).eval() + inputs = tuple(value.clone() for value in inputs) + if not use_explicit_layout: + # Legacy kernels rely on channels-last capture for zero-copy CMSIS input. + module.to(memory_format=torch.channels_last) + inputs = tuple( + value.to(memory_format=torch.channels_last) if value.dim() == 4 else value + for value in inputs + ) + + tester = CortexMTester( + module, + inputs, + target_config=target_config, + use_explicit_layout=use_explicit_layout, + ) + tester.quantize().export().to_edge().run_passes() + exported_program = tester.get_artifact(StageType.RUN_PASSES).exported_program() + for target, count in expected_ops.items(): + assert _count(exported_program, target) == count + + tester.to_executorch() + program = tester.get_artifact(StageType.TO_EXECUTORCH).executorch_program + return tuple(program.execution_plan[0].non_const_buffer_sizes) + + +def test_conv2d_uses_explicit_nhwc_operator(): + x = torch.randn(1, 3, 8, 8) + program = _lower(Conv2d(), (x,)) + + assert _count(program, exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default) == 1 + assert _count(program, exir_ops.edge.cortex_m.quantized_conv2d.default) == 0 + assert _count(program, exir_ops.edge.cortex_m.transpose.default) == 2 + assert program.module()(x).shape == torch.Size([1, 4, 8, 8]) + + +def test_conv1d_reuses_explicit_conv2d_region(): + x = torch.randn(1, 2, 8) + tester = CortexMTester(Conv1d(), (x,), use_explicit_layout=True) + tester.quantize().export().to_edge().run_passes() + tester.run_method_and_compare_outputs(inputs=(x,), qtol=1) + program = tester.get_artifact(StageType.RUN_PASSES).exported_program() + + assert _count(program, exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default) == 1 + assert _count(program, exir_ops.edge.cortex_m.transpose.default) == 2 + assert program.module()(x).shape == torch.Size([1, 4, 8]) + + +def test_legacy_mode_does_not_select_explicit_layout(): + x = torch.randn(1, 3, 8, 8) + program = _lower_legacy(Conv2d(), (x,)) + + assert _count(program, exir_ops.edge.aten.convolution.default) == 1 + assert _count(program, exir_ops.edge.cortex_m.quantized_conv2d.default) == 0 + assert _count(program, exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default) == 0 + + +def test_depthwise_conv2d_uses_explicit_nhwc_operator(): + x = torch.randn(1, 4, 8, 8) + program = _lower(Conv2d(4, 4, groups=4), (x,)) + + assert ( + _count( + program, + exir_ops.edge.cortex_m.quantized_depthwise_conv2d_nhwc.default, + ) + == 1 + ) + assert ( + _count(program, exir_ops.edge.cortex_m.quantized_depthwise_conv2d.default) == 0 + ) + assert _count(program, exir_ops.edge.cortex_m.transpose.default) == 2 + assert program.module()(x).shape == torch.Size([1, 4, 8, 8]) + + +def test_adjacent_convolutions_eliminate_internal_copies(): + x = torch.randn(1, 3, 8, 8) + program = _lower(TwoConv2d(), (x,)) + + assert _count(program, exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default) == 2 + assert _count(program, exir_ops.edge.cortex_m.transpose.default) == 2 + + +def test_avg_pool2d_joins_explicit_layout_region(): + x = torch.randn(1, 3, 8, 8) + program = _lower(ConvPoolConv(torch.nn.AvgPool2d(2, 2)), (x,)) + + assert _count(program, exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default) == 2 + assert ( + _count(program, exir_ops.edge.cortex_m.quantized_avg_pool2d_nhwc.default) == 1 + ) + assert _count(program, exir_ops.edge.cortex_m.transpose.default) == 2 + + +def test_max_pool2d_joins_explicit_layout_region(): + x = torch.randn(1, 3, 8, 8) + program = _lower(ConvPoolConv(torch.nn.MaxPool2d(2, 2)), (x,)) + + assert _count(program, exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default) == 2 + assert ( + _count(program, exir_ops.edge.cortex_m.quantized_max_pool2d_nhwc.default) == 1 + ) + assert _count(program, exir_ops.edge.cortex_m.transpose.default) == 2 + + +def test_transpose_conv2d_uses_explicit_nhwc_operator(): + x = torch.randn(1, 3, 5, 5) + program = _lower(ConvTranspose2d(), (x,)) + + assert ( + _count( + program, + exir_ops.edge.cortex_m.quantized_transpose_conv2d_nhwc.default, + ) + == 1 + ) + assert ( + _count(program, exir_ops.edge.cortex_m.quantized_transpose_conv2d.default) == 0 + ) + assert _count(program, exir_ops.edge.cortex_m.transpose.default) == 2 + assert program.module()(x).shape == torch.Size([1, 4, 9, 9]) + + +def test_pad_is_remapped_inside_explicit_layout_region(): + x = torch.randn(1, 3, 8, 8) + program = _lower(ConvPadConv(), (x,)) + [pad] = [ + node + for node in program.graph.nodes + if node.target == exir_ops.edge.cortex_m.pad_nhwc.default + ] + + assert pad.args[1:3] == ([0, 1, 1, 0], [0, 1, 1, 0]) + assert _count(program, exir_ops.edge.cortex_m.transpose.default) == 2 + + +def test_softmax_stays_outside_explicit_layout_region(): + x = torch.randn(1, 3, 8, 8) + program = _lower(ConvSoftmax(), (x,)) + [softmax] = [ + node + for node in program.graph.nodes + if node.target == exir_ops.edge.cortex_m.softmax.default + ] + + assert softmax.args[1] in (-1, 3) + assert _count(program, exir_ops.edge.cortex_m.transpose.default) == 2 + + +def test_unquantized_channel_bias_stays_outside_explicit_layout_region(): + x = torch.randn(1, 3, 8, 8) + program = _lower(ConvBiasConv(), (x,)) + [add] = [ + node + for node in program.graph.nodes + if node.target == exir_ops.edge.aten.add.Tensor + ] + + assert add.args[1].meta["val"].shape == torch.Size([1, 4, 1, 1]) + assert _count(program, exir_ops.edge.cortex_m.quantized_add.default) == 0 + assert _count(program, exir_ops.edge.cortex_m.transpose.default) == 4 + + +def test_explicit_layout_does_not_increase_planned_memory_for_float_qdq(): + torch.manual_seed(0) + m33 = CortexMTargetConfig(cpu=CortexM.M33) + m55 = CortexMTargetConfig(cpu=CortexM.M55) + cases = ( + ( + Conv2d(), + (torch.randn(1, 3, 8, 8),), + m55, + {exir_ops.edge.cortex_m.quantized_conv2d.default: 1}, + {exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default: 1}, + ), + ( + TwoConv2d(), + (torch.randn(1, 3, 8, 8),), + m55, + {exir_ops.edge.cortex_m.quantized_conv2d.default: 2}, + {exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default: 2}, + ), + ( + Conv2d(4, 4, groups=4), + (torch.randn(1, 4, 8, 8),), + m55, + {exir_ops.edge.cortex_m.quantized_depthwise_conv2d.default: 1}, + {exir_ops.edge.cortex_m.quantized_depthwise_conv2d_nhwc.default: 1}, + ), + ( + ConvTranspose2d(), + (torch.randn(1, 3, 5, 5),), + m55, + {exir_ops.edge.cortex_m.quantized_transpose_conv2d.default: 1}, + {exir_ops.edge.cortex_m.quantized_transpose_conv2d_nhwc.default: 1}, + ), + ( + ConvPoolConv(torch.nn.AvgPool2d(2, 2)), + (torch.randn(1, 3, 8, 8),), + m55, + { + exir_ops.edge.cortex_m.quantized_conv2d.default: 2, + exir_ops.edge.cortex_m.quantized_avg_pool2d.default: 1, + }, + { + exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default: 2, + exir_ops.edge.cortex_m.quantized_avg_pool2d_nhwc.default: 1, + }, + ), + ( + ConvPoolConv(torch.nn.MaxPool2d(2, 2)), + (torch.randn(1, 3, 8, 8),), + m55, + { + exir_ops.edge.cortex_m.quantized_conv2d.default: 2, + exir_ops.edge.cortex_m.quantized_max_pool2d.default: 1, + }, + { + exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default: 2, + exir_ops.edge.cortex_m.quantized_max_pool2d_nhwc.default: 1, + }, + ), + ( + ConvForkAdd(), + (torch.randn(1, 3, 8, 8),), + m55, + {exir_ops.edge.cortex_m.quantized_conv2d.default: 3}, + {exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default: 3}, + ), + ( + ConvPadConv(), + (torch.randn(1, 3, 8, 8),), + m55, + { + exir_ops.edge.cortex_m.quantized_conv2d.default: 2, + exir_ops.edge.cortex_m.pad.default: 1, + }, + { + exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default: 2, + exir_ops.edge.cortex_m.pad_nhwc.default: 1, + }, + ), + ( + Conv2d(), + (torch.randn(1, 3, 8, 8),), + m33, + {exir_ops.edge.cortex_m.quantized_conv2d.default: 1}, + {exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default: 1}, + ), + ) + + for module, inputs, target_config, legacy_ops, explicit_ops in cases: + legacy = _planned_buffer_sizes( + module, + inputs, + use_explicit_layout=False, + target_config=target_config, + expected_ops=legacy_ops, + ) + explicit = _planned_buffer_sizes( + module, + inputs, + use_explicit_layout=True, + target_config=target_config, + expected_ops=explicit_ops, + ) + assert len(explicit) == len(legacy), type(module).__name__ + assert all( + explicit_size <= legacy_size + for explicit_size, legacy_size in zip(explicit, legacy) + ), type(module).__name__ + + +def test_explicit_nhwc_conv2d_runs_on_fvp(): + x = torch.linspace(-5, 5, steps=1 * 3 * 7 * 10).reshape(1, 3, 7, 10) + _run_explicit_layout_on_fvp( + Conv2d(kernel_size=(3, 2), stride=(2, 1), padding=(1, 0)), + (x,), + exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default, + ) + + +def test_explicit_nhwc_depthwise_conv2d_runs_on_fvp(): + x = torch.linspace(-5, 5, steps=1 * 4 * 7 * 10).reshape(1, 4, 7, 10) + _run_explicit_layout_on_fvp( + Conv2d( + 4, + 4, + groups=4, + kernel_size=(3, 2), + stride=(2, 1), + padding=(1, 0), + ), + (x,), + exir_ops.edge.cortex_m.quantized_depthwise_conv2d_nhwc.default, + ) + + +def test_explicit_nhwc_transpose_conv2d_runs_on_fvp(): + module = ConvTranspose2d(2, 4, kernel_size=(2, 4), stride=1, padding=0, bias=False) + module.conv.weight.data.fill_(1.0) + x = torch.linspace(-8, 8, steps=1 * 2 * 5 * 5).reshape(1, 2, 5, 5) + _run_explicit_layout_on_fvp( + module, + (x,), + exir_ops.edge.cortex_m.quantized_transpose_conv2d_nhwc.default, + ) + + +def test_explicit_nhwc_avg_pool2d_runs_on_fvp(): + x = torch.linspace(-5, 5, steps=1 * 3 * 7 * 9).reshape(1, 3, 7, 9) + _run_explicit_layout_on_fvp( + ConvPoolConv( + torch.nn.AvgPool2d(kernel_size=(2, 3), stride=(2, 1), padding=(0, 1)) + ), + (x,), + exir_ops.edge.cortex_m.quantized_avg_pool2d_nhwc.default, + ) + + +def test_explicit_nhwc_max_pool2d_runs_on_fvp(): + x = torch.linspace(-5, 5, steps=1 * 3 * 7 * 9).reshape(1, 3, 7, 9) + _run_explicit_layout_on_fvp( + ConvPoolConv( + torch.nn.MaxPool2d(kernel_size=(2, 3), stride=(2, 1), padding=(0, 1)) + ), + (x,), + exir_ops.edge.cortex_m.quantized_max_pool2d_nhwc.default, + ) + + +def test_explicit_nhwc_pad_runs_on_fvp_with_singleton_height(): + x = torch.linspace(-5, 5, steps=1 * 3 * 1 * 7).reshape(1, 3, 1, 7) + _run_explicit_layout_on_fvp( + ConvPadConv(), + (x,), + exir_ops.edge.cortex_m.pad_nhwc.default, + ) + + +def test_aot_explicit_layout_conv1d_runs_on_fvp(): + from types import SimpleNamespace + + from executorch.backends.arm.scripts.aot_arm_compiler import _to_edge_cortex_m + from executorch.backends.cortex_m.test.tester import CortexMSerialize + + module = Conv1d().eval() + inputs = (torch.linspace(-5, 5, steps=1 * 2 * 8).reshape(1, 2, 8),) + exported_program = torch.export.export(module, inputs, strict=True) + target_config = CortexMTargetConfig(cpu=CortexM.M55) + model_quant, edge, runtime_inputs = _to_edge_cortex_m( + exported_program, + SimpleNamespace( + cortex_m_explicit_layout=True, + quantize=True, + strict_export=True, + ), + exported_program.module(), + inputs, + None, + target_config, + ) + program = edge.exported_program() + + assert model_quant is not None + assert _count(program, exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default) == 1 + assert all( + getattr(node.target, "namespace", None) != "dim_order_ops" + for node in program.graph.nodes + if node.op == "call_function" + ) + + serialized = CortexMSerialize(target_config) + serialized.run(edge.to_executorch()) + [actual] = serialized.run_artifact(runtime_inputs) + expected = model_quant(*runtime_inputs) + torch.testing.assert_close(actual, expected, atol=0.05, rtol=1e-3) diff --git a/backends/cortex_m/test/tester.py b/backends/cortex_m/test/tester.py index a1b5245b80b..7290a463568 100644 --- a/backends/cortex_m/test/tester.py +++ b/backends/cortex_m/test/tester.py @@ -31,24 +31,39 @@ class CortexMQuantize(Quantize): - def __init__(self, calibration_samples=None): - quantizer = CortexMQuantizer() + def __init__(self, calibration_samples=None, use_explicit_layout: bool = False): + quantizer = CortexMQuantizer(use_explicit_layout=use_explicit_layout) super().__init__(quantizer, calibration_samples=calibration_samples) class CortexMToEdge(ToEdge): - def __init__(self): - super().__init__(cortex_m_edge_compile_config()) + def __init__(self, use_explicit_layout: bool = False): + super().__init__( + cortex_m_edge_compile_config(use_explicit_layout=use_explicit_layout) + ) class CortexMRunPasses(RunPasses): - def __init__(self, target_config: Optional[CortexMTargetConfig] = None): + def __init__( + self, + target_config: Optional[CortexMTargetConfig] = None, + use_explicit_layout: bool = False, + ): 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. + pass_list = ( + CortexMPassManager.explicit_layout_pass_list + if use_explicit_layout + else CortexMPassManager.legacy_pass_list + ) super().__init__( - partial(CortexMPassManager, target_config=target_config), # type: ignore[arg-type] - CortexMPassManager.pass_list, # type: ignore[arg-type] + partial( + CortexMPassManager, + target_config=target_config, + use_explicit_layout=use_explicit_layout, + ), # type: ignore[arg-type] + pass_list, # type: ignore[arg-type] ) @@ -60,9 +75,19 @@ class CortexMToEdgeTransformAndLower(ToEdgeTransformAndLower): lowering entry point the shared backend test suite drives. """ - def __init__(self, target_config: Optional[CortexMTargetConfig] = None): - super().__init__(edge_compile_config=cortex_m_edge_compile_config()) - self._run_passes = CortexMRunPasses(target_config) + def __init__( + self, + target_config: Optional[CortexMTargetConfig] = None, + use_explicit_layout: bool = False, + ): + super().__init__( + edge_compile_config=cortex_m_edge_compile_config( + use_explicit_layout=use_explicit_layout + ) + ) + self._run_passes = CortexMRunPasses( + target_config, use_explicit_layout=use_explicit_layout + ) def run(self, artifact, inputs=None, generate_etrecord: bool = False) -> None: super().run(artifact, inputs, generate_etrecord=generate_etrecord) @@ -105,20 +130,32 @@ def __init__( example_inputs, target_config: Optional[CortexMTargetConfig] = None, timeout: int = 120, + use_explicit_layout: bool = False, ): if callable(example_inputs): resolved_example_inputs = example_inputs() else: resolved_example_inputs = example_inputs target_config = target_config or CortexMTargetConfig(cpu=CortexM.M55) + self.use_explicit_layout = use_explicit_layout stage_classes: dict[StageType, Callable[..., Any]] = dict( cortex_m_stage_classes ) + stage_classes[StageType.QUANTIZE] = lambda: CortexMQuantize( + use_explicit_layout=use_explicit_layout + ) + stage_classes[StageType.TO_EDGE] = lambda: CortexMToEdge( + use_explicit_layout=use_explicit_layout + ) stage_classes[StageType.RUN_PASSES] = lambda: CortexMRunPasses( - target_config=target_config + target_config=target_config, + use_explicit_layout=use_explicit_layout, ) stage_classes[StageType.TO_EDGE_TRANSFORM_AND_LOWER] = ( - lambda: CortexMToEdgeTransformAndLower(target_config=target_config) + lambda: CortexMToEdgeTransformAndLower( + target_config=target_config, + use_explicit_layout=use_explicit_layout, + ) ) stage_classes[StageType.SERIALIZE] = lambda: CortexMSerialize( target_config=target_config, timeout=timeout @@ -138,7 +175,8 @@ def test_dialect( """ if calibration_samples is not None: quantization_stage = CortexMQuantize( - calibration_samples=calibration_samples + calibration_samples=calibration_samples, + use_explicit_layout=self.use_explicit_layout, ) else: quantization_stage = None @@ -160,7 +198,8 @@ def test_implementation(self, qtol=0, atol=1e-03, calibration_samples=None): if calibration_samples is not None: quantization_stage = CortexMQuantize( - calibration_samples=calibration_samples + calibration_samples=calibration_samples, + use_explicit_layout=self.use_explicit_layout, ) else: quantization_stage = None