Postpone conv1d output permutes below their squeeze (#24256)

Summary:

On the explicit-layout pass list, a Conv1d is lowered as a Conv2d, so each convolution's NHWC output goes through a rank-4 layout copy followed by a squeeze back to rank 3. `RemovePermutesAroundElementwiseOps` cancels permutes around a region of elementwise ops, but it can only cross a squeeze or unsqueeze while walking downstream. When a region walks upstream into a neighbouring convolution's output, for example through the second operand of a residual add, it reaches that squeeze from below, carrying a rank-3 permutation with a rank-4 permute behind it. It discards the whole region, and the layout copies between consecutive convolutions survive as transposes.

Add `CortexMPostponeLayoutCopyBelowSqueezePass`, a restricted `PostponePermuteOpBelowSqueezeOrUnsqueezeLikeView`, and run it just before `RemovePermutesAroundElementwiseOps` in the explicit-layout pass list. It only moves the layout copies this pipeline inserts (its `targets` is `LAYOUT_PERMUTE_COPY`), leaving user permutes in place, and never moves one below an unsqueeze, which would raise its rank, up to a rank-5 layout copy that Cortex-M cannot run. After sinking, a convolution output is a squeeze followed by a rank-3 layout copy, so regions meet copies of their own rank and cancel as they do for Conv2d. As its own pass-list entry, the pass gets the pass manager's dtype-preserving retrace, so the folded int8 dtypes survive the sinking.

Authored with an AI coding assistant (Claude).

Differential Revision: D128001279

Author: pssrawat

backends/cortex_m/passes/BUCK

@@ -57,6 +57,7 @@ fbcode_target(_kind = runtime.python_library,
         "//executorch/backends/transforms:channels_last_layout",
         "//executorch/backends/transforms:channels_last_ops",
         "//executorch/backends/transforms:convert_conv1d_to_conv2d_pass",
+        "//executorch/backends/transforms:postpone_permute_below_squeeze_view",
         "//executorch/backends/transforms:remove_getitem_op",
         "//executorch/backends/transforms:remove_permutes_around_elementwise_ops",
         "//executorch/backends/transforms:remove_unused_constants_pass",

backends/cortex_m/passes/cortex_m_pass_manager.py

@@ -55,6 +55,7 @@
 from .decompose_sdpa_pass import DecomposeSDPAPass
 from .explicit_layout_pass import (
     CortexMCanonicalizeViewCopyPermutePass,
+    CortexMPostponeLayoutCopyBelowSqueezePass,
     CortexMReplaceOpsWithChannelsLastVariants,
     ValidateCortexMExplicitLayoutPass,
 )
@@ -175,6 +176,8 @@ class CortexMPassManager(ExportedProgramPassManager):
         ConvertConv1dToConv2dPass,
         CortexMReplaceOpsWithChannelsLastVariants,
         ReplaceSqueezeAndUnsqueezeWithViewPass,
+        # Let elementwise regions meet layout copies of their own rank.
+        CortexMPostponeLayoutCopyBelowSqueezePass,
         # Move layout copies across pads before singleton permutations become views.
         RemovePermutesAroundElementwiseOps,
         CortexMCanonicalizeViewCopyPermutePass,

backends/cortex_m/passes/explicit_layout_pass.py

@@ -19,12 +19,16 @@
     LAYOUT_PERMUTE_COPY,
     PERMUTE_COPY_TARGETS,
 )
+from executorch.backends.transforms.postpone_permute_below_squeeze_view import (
+    PostponePermuteOpBelowSqueezeOrUnsqueezeLikeView,
+)
 from executorch.backends.transforms.replace_ops_with_channels_last_variants import (
     ChannelsLastOpSpec,
     ReplaceOpsWithChannelsLastVariants,
 )
 from executorch.exir import ExportedProgram
 from executorch.exir.dialects._ops import ops as exir_ops
+from executorch.exir.dialects.edge._ops import EdgeOpOverload
 from executorch.exir.pass_base import ExportPass, PassResult
 from torch.fx import GraphModule, Node
 from torch.fx.node import Target
@@ -117,6 +121,33 @@ def _set_node_op(self, node, target, input_node, arg) -> None:
             )
 
 
+class CortexMPostponeLayoutCopyBelowSqueezePass(
+    PostponePermuteOpBelowSqueezeOrUnsqueezeLikeView
+):
+    """Sink a layout copy below a following squeeze, never below an unsqueeze.
+
+    Only the layout copies this pipeline inserts are moved; user permutes keep
+    their place. Sinking below an unsqueeze would raise the copy's rank, up to
+    a rank-5 layout copy that Cortex-M cannot run.
+    """
+
+    @property
+    def targets(self) -> list[EdgeOpOverload]:
+        return [LAYOUT_PERMUTE_COPY]
+
+    def maybe_remove_or_replace(self, node: Node) -> bool:
+        users = list(node.users)
+        user_value = users[0].meta.get("val") if len(users) == 1 else None
+        value = node.meta.get("val")
+        if (
+            isinstance(user_value, torch.Tensor)
+            and isinstance(value, torch.Tensor)
+            and user_value.dim() > value.dim()
+        ):
+            return False
+        return super().maybe_remove_or_replace(node)
+
+
 class ValidateCortexMExplicitLayoutPass(ExportPass):
     def call(self, graph_module: GraphModule) -> PassResult:
         for node in graph_module.graph.nodes:

backends/cortex_m/test/targets.bzl

@@ -83,6 +83,8 @@ 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/test:graph_builder",
+                "//executorch/backends/transforms:channels_last_layout",
                 "//executorch/backends/transforms:remove_unused_constants_pass",
                 "//executorch/exir:lib",
                 "//executorch/exir/dialects:lib",

backends/cortex_m/test/test_explicit_layout_pipeline.py

@@ -13,14 +13,19 @@
     CortexMPassManager,
     LiftConstantTensorsPass,
 )
+from executorch.backends.cortex_m.passes.explicit_layout_pass import (
+    CortexMPostponeLayoutCopyBelowSqueezePass,
+)
 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 (
     CortexMQuantize,
     CortexMRunPasses,
     CortexMTester,
 )
+from executorch.backends.test.graph_builder import GraphBuilder
 from executorch.backends.test.harness.stages import Quantize, StageType
+from executorch.backends.transforms.channels_last_layout import LAYOUT_PERMUTE_COPY
 from executorch.backends.transforms.remove_unused_constants_pass import (
     RemoveUnusedConstantsPass,
 )
@@ -88,6 +93,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 +686,62 @@ 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 _count(program, exir_ops.edge.cortex_m.transpose.default) == 2
+    assert _count(program, exir_ops.edge.aten.permute_copy.default) == 0
+    tester.run_method_and_compare_outputs(inputs=inputs, qtol=1)
+
+
+def _permute_then_view(permute, shape, view_shape) -> torch.fx.GraphModule:
+    builder = GraphBuilder()
+    x = builder.placeholder("x", torch.randn(shape))
+    permuted = builder.call_operator(op=permute, args=(x, [0, 3, 1, 2]))
+    view = builder.call_operator(
+        op=exir_ops.edge.aten.view_copy.default, args=(permuted, view_shape)
+    )
+    builder.output([view])
+    return builder.get_graph_module()
+
+
+def test_layout_copy_is_sunk_below_a_squeeze():
+    x = torch.randn(1, 1, 4, 6)
+    graph_module = _permute_then_view(LAYOUT_PERMUTE_COPY, x.shape, [1, 6, 4])
+
+    result = CortexMPostponeLayoutCopyBelowSqueezePass()(graph_module)
+
+    assert result.modified
+    assert [
+        (node.target, node.args[1])
+        for node in result.graph_module.graph.nodes
+        if node.op == "call_function"
+    ] == [
+        (exir_ops.edge.aten.view_copy.default, [1, 4, 6]),
+        (LAYOUT_PERMUTE_COPY, [0, 2, 1]),
+    ]
+    torch.testing.assert_close(
+        result.graph_module(x)[0], x.permute(0, 3, 1, 2).reshape(1, 6, 4)
+    )
+
+
+@pytest.mark.parametrize(
+    "permute,shape,view_shape",
+    [
+        (LAYOUT_PERMUTE_COPY, (1, 5, 4, 6), [1, 1, 6, 5, 4]),
+        (exir_ops.edge.aten.permute_copy.default, (1, 1, 4, 6), [1, 6, 4]),
+    ],
+    ids=["layout_copy_unsqueeze", "user_permute_squeeze"],
+)
+def test_other_permutes_are_not_sunk(permute, shape, view_shape):
+    graph_module = _permute_then_view(permute, shape, view_shape)
+
+    result = CortexMPostponeLayoutCopyBelowSqueezePass()(graph_module)
+
+    assert not result.modified