Cortex-M: add opt-in explicit layout lowering

backends/arm/scripts/aot_arm_compiler.py

@@ -626,6 +626,14 @@ def _get_args():
         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 +929,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 +959,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 +982,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()
 

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,13 @@ 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_clone_ops",
         "//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",

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
@@ -60,9 +61,11 @@ def __init__(
         self,
         exported_program: ExportedProgram,
         target_config: CortexMTargetConfig,
+        use_explicit_layout: bool = False,
     ) -> None:
         super().__init__(exported_program=exported_program)
         self.target_config = target_config
+        self.use_explicit_layout = use_explicit_layout
 
     def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
         result = super().call(graph_module)
@@ -441,13 +444,18 @@ 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
 ) -> DialectNodeSpec | None:
     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 +613,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 +639,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 +658,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 +771,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 +843,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 +873,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_contiguous.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 +904,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 +1095,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 +1160,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 +1194,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:
@@ -1164,7 +1218,7 @@ def _get_permute_replacement(
 def _get_pad_replacement(
     node: Node, dialect_pass: AtenToDialectPass
 ) -> DialectNodeSpec | None:
-    del dialect_pass
+    contiguous = cast(AtenToCortexMPass, dialect_pass).use_explicit_layout
     input_qparams = node.meta.get("input_qparams", {})
     if not input_qparams:
         return None
@@ -1194,8 +1248,17 @@ 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 contiguous:
+        # The legacy entry point infers the layout, so hand it physical order.
+        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_contiguous.default
+            if contiguous
+            else exir_ops.edge.cortex_m.pad.default
+        ),
+        args,
+    )