diff --git a/backends/cortex_m/passes/explicit_layout_pass.py b/backends/cortex_m/passes/explicit_layout_pass.py index 659b67d729e..971f7d83526 100644 --- a/backends/cortex_m/passes/explicit_layout_pass.py +++ b/backends/cortex_m/passes/explicit_layout_pass.py @@ -5,6 +5,8 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +from typing import cast + import executorch.backends.transforms.channels_last_ops # noqa: F401 import torch @@ -26,7 +28,7 @@ from executorch.exir import ExportedProgram from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.pass_base import ExportPass, PassResult -from torch.fx import GraphModule, Node +from torch.fx import Graph, GraphModule, Node from torch.fx.node import Target @@ -100,6 +102,42 @@ class CortexMReplaceOpsWithChannelsLastVariants(ReplaceOpsWithChannelsLastVarian def __init__(self, exported_program: ExportedProgram) -> None: super().__init__(exported_program, op_map=dict(_EXPLICIT_LAYOUT_OP_MAP)) + @staticmethod + def _is_nchw_to_ncw_squeeze(node: Node) -> bool: + # The squeeze ConvertConv1dToConv2dPass emits: squeeze_copy.dim(conv, 2). + return ( + node.target == exir_ops.edge.aten.squeeze_copy.dim + # Squeezes NCHW H; % 4 normalizes a negative dim of the rank-4 input. + and cast(int, node.args[1]) % 4 == 2 + # Squeezing a non-unit dim is a no-op that keeps rank 4. + and node.meta["val"].dim() == 3 + ) + + # Emit a Conv1d-as-Conv2d output as squeeze -> rank-3 layout copy rather + # than rank-4 copy -> squeeze, so RemovePermutesAroundElementwiseOps can + # cancel the copy against the elementwise ops that consume it. + @staticmethod + def _permute_node_output( + graph: Graph, + node_output: Node, + original_node_output: Node, + implicit_batch: bool, + ) -> None: + users = list(original_node_output.users) + is_squeeze = CortexMReplaceOpsWithChannelsLastVariants._is_nchw_to_ncw_squeeze + if len(users) != 1 or not is_squeeze(users[0]): + ReplaceOpsWithChannelsLastVariants._permute_node_output( + graph, node_output, original_node_output, implicit_batch + ) + return + squeeze = users[0] + squeeze.args = (node_output, 1) + with graph.inserting_after(squeeze): + permute = graph.call_function(LAYOUT_PERMUTE_COPY, (squeeze, [0, 2, 1])) + squeeze.replace_all_uses_with( + permute, delete_user_cb=lambda n: n is not permute + ) + class CortexMCanonicalizeViewCopyPermutePass(CanonicalizeViewCopyPermutePass): def __init__(self) -> None: diff --git a/backends/cortex_m/test/test_explicit_layout_pipeline.py b/backends/cortex_m/test/test_explicit_layout_pipeline.py index d29248e5647..101c8eadcb7 100644 --- a/backends/cortex_m/test/test_explicit_layout_pipeline.py +++ b/backends/cortex_m/test/test_explicit_layout_pipeline.py @@ -13,6 +13,9 @@ CortexMPassManager, LiftConstantTensorsPass, ) +from executorch.backends.cortex_m.passes.explicit_layout_pass import ( + CortexMReplaceOpsWithChannelsLastVariants, +) from executorch.backends.cortex_m.quantizer.quantizer import CortexMQuantizer from executorch.backends.cortex_m.target_config import CortexM, CortexMTargetConfig from executorch.backends.cortex_m.test.tester import ( @@ -88,6 +91,18 @@ def forward(self, x): return self.conv2(self.conv1(torch.tanh(x)).view(1, 4, 32)) +class Conv1dChain(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv1 = torch.nn.Conv1d(8, 8, 3, padding=1) + self.conv2 = torch.nn.Conv1d(8, 8, 3, padding=1) + self.conv3 = torch.nn.Conv1d(8, 8, 1) + + def forward(self, x): + y = torch.tanh(self.conv1(x)) + return self.conv3(torch.relu(self.conv2(y) + y)) + + class Apply(torch.nn.Module): def __init__(self, fn) -> None: super().__init__() @@ -669,3 +684,56 @@ def test_unquantized_constant_is_not_precomputed(use_explicit_layout): assert _count(program, exir_ops.edge.aten.full.default) == 1 tester.run_method_and_compare_outputs(inputs=inputs) + + +def test_conv1d_chain_keeps_only_the_boundary_layout_copies(): + torch.manual_seed(0) + inputs = (torch.randn(1, 8, 16),) + tester = _run_explicit_layout_passes(CortexMTester(Conv1dChain().eval(), inputs)) + program = tester.get_artifact(StageType.RUN_PASSES).exported_program() + + assert _count(program, exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default) == 3 + assert [ + node.args[1] + for node in program.graph.find_nodes( + op="call_function", target=exir_ops.edge.cortex_m.transpose.default + ) + ] == [[0, 2, 1], [0, 2, 1]] + assert _count(program, exir_ops.edge.aten.permute_copy.default) == 0 + tester.run_method_and_compare_outputs(inputs=inputs, qtol=1) + + +def test_conv1d_output_is_squeezed_before_its_layout_copy(): + torch.manual_seed(0) + inputs = (torch.randn(1, 8, 16),) + tester = _run_explicit_layout_passes( + CortexMTester(torch.nn.Conv1d(8, 4, 3).eval(), inputs) + ) + program = tester.get_artifact(StageType.RUN_PASSES).exported_program() + + (conv,) = program.graph.find_nodes( + op="call_function", target=exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default + ) + (squeeze,) = conv.users + (layout_copy,) = squeeze.users + assert squeeze.target == exir_ops.edge.aten.view_copy.default + assert squeeze.meta["val"].shape == (1, 14, 4) + assert layout_copy.target == exir_ops.edge.cortex_m.transpose.default + assert layout_copy.args[1] == [0, 2, 1] + tester.run_method_and_compare_outputs(inputs=inputs, qtol=1) + + +@pytest.mark.parametrize( + "dim,expected", [(2, True), (-2, True), (1, False), (-1, False)] +) +def test_only_a_height_squeeze_is_moved_above_the_layout_copy(dim, expected): + graph = torch.fx.Graph() + squeeze = graph.call_function( + exir_ops.edge.aten.squeeze_copy.dim, (graph.placeholder("x"), dim) + ) + squeeze.meta["val"] = torch.empty(1, 8, 16) + + assert ( + CortexMReplaceOpsWithChannelsLastVariants._is_nchw_to_ncw_squeeze(squeeze) + is expected + )