Cortex-M: add opt-in explicit layout lowering

backends/arm/cmsis_pack/test/all_ops/generate_test_models.py

@@ -1,3 +1,3 @@
 #!/usr/bin/env python3
 # Copyright 2026 Arm Limited and/or its affiliates.
 #
@@ -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 @@
     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 @@
     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 @@
     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 @@
 
     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 @@
     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 @@
     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 _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 @@
     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 @@
     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

backends/arm/scripts/aot_arm_compiler.py

@@ -1,3 +1,3 @@
 #!/usr/bin/env python3
 
 # Copyright (c) Meta Platforms, Inc. and affiliates.
@@ -156,7 +156,6 @@
     """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 @@
     Raises:
         RuntimeError: If the model cannot be resolved or required inputs are
             missing.
-
     """
     example_inputs = _load_example_inputs(model_input)
 
@@ -346,8 +344,7 @@
     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 @@
         choices=TARGETS,
         help=f"Target backend. For delegated models: Ethos-U/VGF/TOSA variants. For non-delegated: cortex-m<variant> (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 @@
     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 @@
         )
         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 @@
 
     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()
 

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

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