diff --git a/backends/cortex_m/passes/BUCK b/backends/cortex_m/passes/BUCK index c0405f667f5..85f92c6ac37 100644 --- a/backends/cortex_m/passes/BUCK +++ b/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", diff --git a/backends/cortex_m/passes/cortex_m_pass_manager.py b/backends/cortex_m/passes/cortex_m_pass_manager.py index 5d4185b57bc..1c2102568d7 100644 --- a/backends/cortex_m/passes/cortex_m_pass_manager.py +++ b/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,15 +114,25 @@ 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] ) @@ -117,6 +140,10 @@ def __init__( 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. diff --git a/backends/cortex_m/test/targets.bzl b/backends/cortex_m/test/targets.bzl index 592abe01d91..15ba6022dfc 100644 --- a/backends/cortex_m/test/targets.bzl +++ b/backends/cortex_m/test/targets.bzl @@ -50,6 +50,7 @@ def define_common_targets(is_fbcode = False): "//executorch/backends/cortex_m/passes:cortex_passes", "//executorch/backends/cortex_m/quantizer:quantizer", "//executorch/backends/test/harness:tester", + "//executorch/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", ], diff --git a/backends/cortex_m/test/test_explicit_layout.py b/backends/cortex_m/test/test_explicit_layout.py index 2eb724d0992..86865de9a14 100644 --- a/backends/cortex_m/test/test_explicit_layout.py +++ b/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) diff --git a/backends/cortex_m/test/tester.py b/backends/cortex_m/test/tester.py index 7290a463568..93ab7c40f8c 100644 --- a/backends/cortex_m/test/tester.py +++ b/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,6 +157,7 @@ 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() @@ -138,6 +165,7 @@ def __init__( 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, diff --git a/backends/transforms/absorb_boundary_layout_copies.py b/backends/transforms/absorb_boundary_layout_copies.py index d97ef62e900..c60e10b7bf4 100644 --- a/backends/transforms/absorb_boundary_layout_copies.py +++ b/backends/transforms/absorb_boundary_layout_copies.py @@ -166,8 +166,11 @@ def _absorb_inputs(self, graph_module) -> dict[int, tuple[int, ...]]: continue interior, copies, dims = found + # .contiguous() is load-bearing: a bare permute leaves the old + # strides, which serialize as a permuted dim order, and the kernels + # reading these tensors require plain contiguous NHWC. for member in (node, *interior): - member.meta["val"] = member.meta["val"].permute(dims) + member.meta["val"] = member.meta["val"].permute(dims).contiguous() for copy in copies: copy.replace_all_uses_with(copy.args[0]) graph_module.graph.erase_node(copy) @@ -190,7 +193,9 @@ def _absorb_outputs(self, graph_module) -> dict[int, tuple[int, ...]]: interior, copy, source, dims = found for member in interior: - member.meta["val"] = member.meta["val"].permute(_inverse(dims)) + member.meta["val"] = ( + member.meta["val"].permute(_inverse(dims)).contiguous() + ) if interior: interior[-1].replace_input_with(copy, source) else: