Emit Conv1d output layout copies after their squeeze (#24256)

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:

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