Cortex-M: add an opt-in NHWC method I/O contract

backends/cortex_m/passes/BUCK

@@ -48,6 +48,7 @@ fbcode_target(_kind = runtime.python_library,
         "//executorch/backends/cortex_m/passes:passes_utils",
         "//executorch/backends/cortex_m/passes:replace_quant_nodes_pass",
         "//executorch/backends/cortex_m/passes:scratch_buffer_sizes",
+        "//executorch/backends/transforms:absorb_boundary_layout_copies",
         "//executorch/backends/transforms:aten_to_dialect_pass",
         "//executorch/backends/transforms:channels_last_ops",
         "//executorch/backends/transforms:convert_conv1d_to_conv2d_pass",

backends/cortex_m/passes/cortex_m_pass_manager.py

@@ -13,6 +13,10 @@
     ScalarsToAttributePass,
 )
 from executorch.backends.cortex_m.target_config import CortexM, CortexMTargetConfig
+from executorch.backends.transforms.absorb_boundary_layout_copies import (
+    AbsorbBoundaryLayoutCopies,
+    BoundaryLayoutContract,
+)
 from executorch.backends.transforms.convert_conv1d_to_conv2d_pass import (
     ConvertConv1dToConv2dPass,
 )
@@ -69,6 +73,14 @@ class CortexMPassManager(PassManager):
         AtenToCortexMPass,
     ]
 
+    # Absorption has to see the layout copies while they are still dialect
+    # nodes, so it goes directly after region formation.
+    nhwc_io_pass_list: list[PassClass] = list(explicit_layout_pass_list)
+    nhwc_io_pass_list.insert(
+        nhwc_io_pass_list.index(CortexMExplicitLayoutPass) + 1,
+        AbsorbBoundaryLayoutCopies,
+    )
+
     pass_list = legacy_pass_list
 
     pass_list_transform_for_annotation: list[PassClass] = [
@@ -86,6 +98,7 @@ def __init__(
         passes: Optional[list[PassClass]] = None,
         target_config: Optional[CortexMTargetConfig] = None,
         use_explicit_layout: bool = False,
+        use_nhwc_io: bool = False,
     ) -> None:
         """Initialize the Cortex-M pass manager.
 
@@ -101,22 +114,36 @@ def __init__(
                 pre-config historical behaviour.
             use_explicit_layout: Run channels-last dialect region formation.
                 Legacy dim-order lowering remains the default.
+            use_nhwc_io: Give the method an NHWC input and output contract
+                instead of transposing at the boundary. Requires
+                ``use_explicit_layout``; see ``boundary_layout_contract``.
         """
+        if use_nhwc_io and not use_explicit_layout:
+            raise ValueError(
+                "use_nhwc_io absorbs channels-last dialect copies and so needs "
+                "use_explicit_layout; the legacy path carries layout in dim "
+                "order, where there are no copies to absorb."
+            )
         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
-        )
+        if use_nhwc_io:
+            default_passes = self.nhwc_io_pass_list
+        elif use_explicit_layout:
+            default_passes = self.explicit_layout_pass_list
+        else:
+            default_passes = self.legacy_pass_list
         self.passes: list[PassClass] = (  # 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
+        self.use_nhwc_io = use_nhwc_io
+        # Populated by transform(); tells callers which method inputs and
+        # outputs changed layout, since that is not inferable from the graph.
+        self.boundary_layout_contract = BoundaryLayoutContract()
 
     def transform_for_annotation(self, model):
         passes = self.pass_list_transform_for_annotation
@@ -150,6 +177,8 @@ def transform(self) -> ExportedProgram:
 
             transform_pass = pass_cls(**kwargs)
             exported_program = _transform(exported_program, transform_pass)
+            if isinstance(transform_pass, AbsorbBoundaryLayoutCopies):
+                self.boundary_layout_contract = transform_pass.contract
 
         # All constant tensors should be lifted to buffers at this point, re-run
         # lift_constant_tensor_pass in case new ones have been introduced.

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/backends/transforms:absorb_boundary_layout_copies",
                 "//executorch/backends/transforms:duplicate_dynamic_quant_chain",
             ],
         )
@@ -95,6 +96,7 @@ def define_common_targets(is_fbcode = False):
                 "//caffe2:torch",
                 ":tester",
                 "//executorch/backends/cortex_m:target_config",
+                "//executorch/backends/cortex_m/passes:cortex_passes",
                 "//executorch/exir/dialects:lib",
                 "fbsource//third-party/pypi/pytest:pytest",
             ],

backends/cortex_m/test/test_explicit_layout.py

@@ -6,7 +6,10 @@
 
 import copy
 
+import pytest
+
 import torch
+from executorch.backends.cortex_m.passes.cortex_m_pass_manager import CortexMPassManager
 from executorch.backends.cortex_m.target_config import CortexM, CortexMTargetConfig
 
 from executorch.backends.cortex_m.test.tester import CortexMTester
@@ -743,3 +746,121 @@ def test_aot_explicit_layout_conv1d_runs_on_fvp():
     [actual] = serialized.run_artifact(runtime_inputs)
     expected = model_quant(*runtime_inputs)
     torch.testing.assert_close(actual, expected, atol=0.05, rtol=1e-3)
+
+
+def _lower_nhwc_io(module, inputs, target_config=None):
+    tester = CortexMTester(
+        module,
+        inputs,
+        target_config=target_config,
+        use_explicit_layout=True,
+        use_nhwc_io=True,
+    )
+    tester.quantize().export().to_edge().run_passes()
+    return tester
+
+
+def test_nhwc_io_removes_the_boundary_transposes():
+    x = torch.randn(1, 3, 8, 8)
+    tester = _lower_nhwc_io(Conv2d(), (x,))
+    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) == 0
+    assert program.module()(x.permute(0, 2, 3, 1).contiguous()).shape == torch.Size(
+        [1, 8, 8, 4]
+    )
+
+
+def test_nhwc_io_reports_the_contract():
+    tester = _lower_nhwc_io(Conv2d(), (torch.randn(1, 3, 8, 8),))
+    contract = tester.boundary_layout_contract
+
+    assert contract.inputs == {0: (0, 2, 3, 1)}
+    assert contract.outputs == {0: (0, 3, 1, 2)}
+
+
+def test_nhwc_io_leaves_layout_free_models_alone():
+    class Linear(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.linear = torch.nn.Linear(8, 4)
+
+        def forward(self, x):
+            return self.linear(x)
+
+    tester = _lower_nhwc_io(Linear(), (torch.randn(2, 8),))
+
+    assert not tester.boundary_layout_contract
+
+
+def test_nhwc_io_keeps_a_shared_input_on_one_contract_entry():
+    x = torch.randn(1, 3, 8, 8)
+    tester = _lower_nhwc_io(ConvForkAdd(), (x,))
+    program = tester.get_artifact(StageType.RUN_PASSES).exported_program()
+
+    assert tester.boundary_layout_contract.inputs == {0: (0, 2, 3, 1)}
+    assert _count(program, exir_ops.edge.cortex_m.transpose.default) == 0
+
+
+def test_nhwc_io_requires_explicit_layout():
+    with pytest.raises(ValueError, match="use_explicit_layout"):
+        CortexMPassManager(None, use_nhwc_io=True)
+
+
+def test_nhwc_io_does_not_increase_planned_memory():
+    # Measured against legacy rather than plain explicit layout. Absorption
+    # removes two tensors, which reshuffles what the greedy planner packs
+    # where; on Conv2d that happens to cost 128 bytes against explicit even
+    # though there is strictly less to place. Legacy is the bar that matters.
+    torch.manual_seed(0)
+    m55 = CortexMTargetConfig(cpu=CortexM.M55)
+    for module, inputs in (
+        (Conv2d(), (torch.randn(1, 3, 8, 8),)),
+        (TwoConv2d(), (torch.randn(1, 3, 8, 8),)),
+        (ConvPoolConv(torch.nn.AvgPool2d(2, 2)), (torch.randn(1, 3, 8, 8),)),
+        (ConvPoolConv(torch.nn.MaxPool2d(2, 2)), (torch.randn(1, 3, 8, 8),)),
+        (ConvForkAdd(), (torch.randn(1, 3, 8, 8),)),
+    ):
+        legacy = _planned_buffer_sizes(
+            module,
+            inputs,
+            use_explicit_layout=False,
+            target_config=m55,
+            expected_ops={},
+        )
+        tester = _lower_nhwc_io(
+            copy.deepcopy(module).eval(),
+            tuple(value.clone() for value in inputs),
+            target_config=m55,
+        )
+        tester.to_executorch()
+        program = tester.get_artifact(StageType.TO_EXECUTORCH).executorch_program
+        nhwc_io = tuple(program.execution_plan[0].non_const_buffer_sizes)
+
+        assert len(nhwc_io) == len(legacy), type(module).__name__
+        assert all(
+            nhwc_size <= legacy_size for nhwc_size, legacy_size in zip(nhwc_io, legacy)
+        ), type(module).__name__
+
+
+def test_nhwc_io_conv2d_runs_on_fvp():
+    x = torch.linspace(-5, 5, steps=1 * 3 * 7 * 10).reshape(1, 3, 7, 10)
+    tester = _lower_nhwc_io(
+        Conv2d(kernel_size=(3, 2), stride=(2, 1), padding=(1, 0)), (x,)
+    )
+    program = tester.get_artifact(StageType.RUN_PASSES).exported_program()
+    assert _count(program, exir_ops.edge.cortex_m.transpose.default) == 0
+
+    tester.to_executorch().serialize()
+    tester.run_method_and_compare_outputs(inputs=(x,), qtol=2)
+
+
+def test_nhwc_io_conv_pool_conv_runs_on_fvp():
+    x = torch.linspace(-5, 5, steps=1 * 3 * 8 * 8).reshape(1, 3, 8, 8)
+    tester = _lower_nhwc_io(ConvPoolConv(torch.nn.MaxPool2d(2, 2)), (x,))
+    program = tester.get_artifact(StageType.RUN_PASSES).exported_program()
+    assert _count(program, exir_ops.edge.cortex_m.transpose.default) == 0
+
+    tester.to_executorch().serialize()
+    tester.run_method_and_compare_outputs(inputs=(x,), qtol=2)

backends/cortex_m/test/tester.py

@@ -28,6 +28,16 @@
     ToEdgeTransformAndLower,
     ToExecutorch,
 )
+from executorch.backends.transforms.absorb_boundary_layout_copies import (
+    BoundaryLayoutContract,
+)
+
+
+def _inverse_permutation(dims: tuple[int, ...]) -> list[int]:
+    inverse = [0] * len(dims)
+    for position, dim in enumerate(dims):
+        inverse[dim] = position
+    return inverse
 
 
 class CortexMQuantize(Quantize):
@@ -48,24 +58,40 @@ def __init__(
         self,
         target_config: Optional[CortexMTargetConfig] = None,
         use_explicit_layout: bool = False,
+        use_nhwc_io: 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
+        if use_nhwc_io:
+            pass_list = CortexMPassManager.nhwc_io_pass_list
+        elif use_explicit_layout:
+            pass_list = CortexMPassManager.explicit_layout_pass_list
+        else:
+            pass_list = CortexMPassManager.legacy_pass_list
+        # The base RunPasses constructs the pass manager as `cls(ep, pass_list)`
+        # and then discards it, so keep a handle: the NHWC I/O contract is only
+        # readable from the manager.
+        self.pass_manager: Optional[CortexMPassManager] = None
+        self._new_pass_manager = partial(
+            CortexMPassManager,
+            target_config=target_config,
+            use_explicit_layout=use_explicit_layout,
+            use_nhwc_io=use_nhwc_io,
         )
         super().__init__(
-            partial(
-                CortexMPassManager,
-                target_config=target_config,
-                use_explicit_layout=use_explicit_layout,
-            ),  # type: ignore[arg-type]
+            self._build_pass_manager,  # type: ignore[arg-type]
             pass_list,  # type: ignore[arg-type]
         )
 
+    def _build_pass_manager(self, exported_program, pass_list) -> CortexMPassManager:
+        self.pass_manager = self._new_pass_manager(exported_program, pass_list)
+        return self.pass_manager
+
+    @property
+    def boundary_layout_contract(self) -> BoundaryLayoutContract:
+        if self.pass_manager is None:
+            return BoundaryLayoutContract()
+        return self.pass_manager.boundary_layout_contract
+
 
 class CortexMToEdgeTransformAndLower(ToEdgeTransformAndLower):
     """to_edge with no partitioner, then CortexMPassManager.
@@ -131,13 +157,15 @@ def __init__(
         target_config: Optional[CortexMTargetConfig] = None,
         timeout: int = 120,
         use_explicit_layout: bool = False,
+        use_nhwc_io: 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
+        self.use_nhwc_io = use_nhwc_io
         stage_classes: dict[StageType, Callable[..., Any]] = dict(
             cortex_m_stage_classes
         )
@@ -150,6 +178,7 @@ def __init__(
         stage_classes[StageType.RUN_PASSES] = lambda: CortexMRunPasses(
             target_config=target_config,
             use_explicit_layout=use_explicit_layout,
+            use_nhwc_io=use_nhwc_io,
         )
         stage_classes[StageType.TO_EDGE_TRANSFORM_AND_LOWER] = (
             lambda: CortexMToEdgeTransformAndLower(
@@ -162,6 +191,59 @@ def __init__(
         )
         super().__init__(module, resolved_example_inputs, stage_classes)
 
+    @property
+    def boundary_layout_contract(self) -> BoundaryLayoutContract:
+        stage = self.stages[StageType.RUN_PASSES]
+        if not isinstance(stage, CortexMRunPasses):
+            return BoundaryLayoutContract()
+        return stage.boundary_layout_contract
+
+    def run_method_and_compare_outputs(self, *args, inputs=None, **kwargs):
+        """Keep tests in NCHW terms even when the method contract is NHWC."""
+        contract = self.boundary_layout_contract
+        if contract:
+            if inputs is None:
+                raise ValueError(
+                    "An NHWC I/O contract needs explicit inputs to permute; "
+                    "randomly generated ones would be fed to the method in the "
+                    "wrong layout."
+                )
+            inputs = tuple(
+                (
+                    value.permute(contract.inputs[index]).contiguous()
+                    if index in contract.inputs
+                    else value
+                )
+                for index, value in enumerate(inputs)
+            )
+        return super().run_method_and_compare_outputs(*args, inputs=inputs, **kwargs)
+
+    def _calculate_reference_output(self, program, inputs):
+        """Restate the NCHW reference in the lowered method's NHWC terms.
+
+        ``run_method_and_compare_outputs`` has already flipped the inputs for
+        the method under test; the reference program is still NCHW, so undo the
+        flip going in and apply the output flip coming out.
+        """
+        contract = self.boundary_layout_contract
+        if not contract:
+            return TesterBase._calculate_reference_output(program, inputs)
+
+        inputs = tuple(
+            (
+                value.permute(_inverse_permutation(contract.inputs[index])).contiguous()
+                if index in contract.inputs
+                else value
+            )
+            for index, value in enumerate(inputs)
+        )
+        output, scale = TesterBase._calculate_reference_output(program, inputs)
+        was_tensor = isinstance(output, torch.Tensor)
+        outputs = [output] if was_tensor else list(output)
+        for index, dims in contract.outputs.items():
+            outputs[index] = outputs[index].permute(_inverse_permutation(dims))
+        return (outputs[0] if was_tensor else tuple(outputs)), scale
+
     def test_dialect(
         self,
         ops_before_transforms,