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

Summary:

On the explicit-layout pass list, a Conv1d is lowered as a Conv2d with a unit height. `CortexMReplaceOpsWithChannelsLastVariants` wrapped each convolution output in a rank-4 NHWC-to-NCHW layout copy, followed by the squeeze that `ConvertConv1dToConv2dPass` adds to get back to rank 3. `RemovePermutesAroundElementwiseOps` cannot match a rank-3 region against a rank-4 copy behind a squeeze, so in residual or chained Conv1d blocks the layout copies between convolutions survived as transposes.

`CortexMReplaceOpsWithChannelsLastVariants` now emits such an output as a squeeze followed by a rank-3 layout copy (`[0, 2, 1]`), reusing the existing squeeze node. It does this only when the convolution's sole user is a `squeeze_copy.dim(conv, 2)` with a rank-3 result; every other output keeps the rank-4 copy.

One side effect: when user code transposes a Conv1d output around elementwise ops, the layout copy and the user's `transpose(1, 2)` are now directly adjacent. `RemovePermutesAroundElementwiseOps` cannot merge adjacent permutes and cancels the user's two transposes against each other instead, so both layout copies survive: Conv1d -> transpose -> relu -> transpose -> Conv1d gets four transposes instead of two. The next change fixes this by running `FuseCascadedTransposeOrPermuteOps` before `RemovePermutesAroundElementwiseOps`.

A Conv1d -> tanh -> Conv1d -> add -> relu -> Conv1d chain now lowers with only the two boundary transposes instead of six.

This makes more regions reachable for `RemovePermutesAroundElementwiseOps`, including regions with lower-rank constants, such as a `[C, L]` or `[L]` buffer added to two Conv1d outputs. The pass's compensation for those constants is fixed separately in the shared pass, and this change depends on that fix.

Authored with an AI coding assistant (Claude).

Differential Revision: D128001279

Pull Request resolved: #24256

Author: pssrawat

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
 
 
@@ -74,6 +76,17 @@ def _supports_max_pool2d(node: Node) -> bool:
     )
 
 
+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
+    )
+
+
 _EXPLICIT_LAYOUT_OP_MAP: dict[Target, ChannelsLastOpSpec] = {
     exir_ops.edge.aten.convolution.default: ChannelsLastOpSpec(
         target=exir_ops.edge.channels_last.convolution.default,
@@ -100,6 +113,30 @@ class CortexMReplaceOpsWithChannelsLastVariants(ReplaceOpsWithChannelsLastVarian
     def __init__(self, exported_program: ExportedProgram) -> None:
         super().__init__(exported_program, op_map=dict(_EXPLICIT_LAYOUT_OP_MAP))
 
+    # 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)
+        if len(users) != 1 or not _is_nchw_to_ncw_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 (
+    _is_nchw_to_ncw_squeeze,
+)
 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,53 @@ 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 _is_nchw_to_ncw_squeeze(squeeze) is expected