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 `CortexMRemovePermutesAroundElementwiseOps`, which sinks each layout copy below a following squeeze and then runs `RemovePermutesAroundElementwiseOps`. The sinking is done by `CortexMPostponeLayoutCopyBelowSqueezePass`, a restricted `PostponePermuteOpBelowSqueezeOrUnsqueezeLikeView`: it only moves the layout copies this pipeline inserts, 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. The sinking runs inside the pass, outside the pass manager's dtype-preserving retrace, so when only the sinking changed the graph the pass retraces once more to restore the folded int8 dtypes.

Authored with an AI coding assistant (Claude).

Differential Revision: D128001279

Author: pssrawat

backends/cortex_m/passes/BUCK

@@ -58,6 +58,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

@@ -21,9 +21,6 @@
     ConvertConv1dToConv2dPass,
 )
 from executorch.backends.transforms.remove_getitem_op import RemoveGetItemPass
-from executorch.backends.transforms.remove_permutes_around_elementwise_ops import (
-    RemovePermutesAroundElementwiseOps,
-)
 from executorch.backends.transforms.remove_unused_constants_pass import (
     RemoveUnusedConstantsPass,
 )
@@ -55,6 +52,7 @@
 from .expand_as_to_expand_pass import ExpandAsToExpandPass
 from .explicit_layout_pass import (
     CortexMCanonicalizeViewCopyPermutePass,
+    CortexMRemovePermutesAroundElementwiseOps,
     CortexMReplaceOpsWithChannelsLastVariants,
     ValidateCortexMExplicitLayoutPass,
 )
@@ -176,7 +174,7 @@ class CortexMPassManager(ExportedProgramPassManager):
         CortexMReplaceOpsWithChannelsLastVariants,
         ReplaceSqueezeAndUnsqueezeWithViewPass,
         # Move layout copies across pads before singleton permutations become views.
-        RemovePermutesAroundElementwiseOps,
+        CortexMRemovePermutesAroundElementwiseOps,
         CortexMCanonicalizeViewCopyPermutePass,
         ValidateCortexMExplicitLayoutPass,
         ReplaceQuantNodesPass,

backends/cortex_m/passes/explicit_layout_pass.py

@@ -19,6 +19,12 @@
     LAYOUT_PERMUTE_COPY,
     PERMUTE_COPY_TARGETS,
 )
+from executorch.backends.transforms.postpone_permute_below_squeeze_view import (
+    PostponePermuteOpBelowSqueezeOrUnsqueezeLikeView,
+)
+from executorch.backends.transforms.remove_permutes_around_elementwise_ops import (
+    RemovePermutesAroundElementwiseOps,
+)
 from executorch.backends.transforms.replace_ops_with_channels_last_variants import (
     ChannelsLastOpSpec,
     ReplaceOpsWithChannelsLastVariants,
@@ -117,6 +123,46 @@ 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.
+    """
+
+    def maybe_remove_or_replace(self, node: Node) -> bool:
+        if node.target != LAYOUT_PERMUTE_COPY:
+            return False
+        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 CortexMRemovePermutesAroundElementwiseOps(RemovePermutesAroundElementwiseOps):
+    """Remove permutes around elementwise regions, after first sinking each layout
+    copy below a following squeeze so regions meet copies of their own rank.
+    """
+
+    def call(self, graph_module: GraphModule) -> PassResult:
+        sunk = CortexMPostponeLayoutCopyBelowSqueezePass()(graph_module)
+        result = super().call(sunk.graph_module)
+        if sunk.modified and not result.modified:
+            # The sinking retraced outside this pass's call_operator; retrace
+            # through it so the folded int8 dtypes are restored.
+            result = ExportPass.call(self, result.graph_module)
+        return PassResult(result.graph_module, sunk.modified or result.modified)
+
+
 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

@@ -14,6 +14,9 @@
     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 (
@@ -25,6 +28,8 @@
 from executorch.backends.transforms.remove_unused_constants_pass import (
     RemoveUnusedConstantsPass,
 )
+from executorch.backends.test.graph_builder import GraphBuilder
+from executorch.backends.transforms.channels_last_layout import LAYOUT_PERMUTE_COPY
 from executorch.exir import to_edge
 from executorch.exir.dialects._ops import ops as exir_ops
 from torch.export import ExportedProgram
@@ -89,6 +94,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__()
@@ -650,3 +667,40 @@ 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)
+
+
+@pytest.mark.parametrize(
+    "permute,shape,view_shape,sunk",
+    [
+        (LAYOUT_PERMUTE_COPY, (1, 1, 4, 6), [1, 6, 4], True),
+        (LAYOUT_PERMUTE_COPY, (1, 5, 4, 6), [1, 1, 6, 5, 4], False),
+        (exir_ops.edge.aten.permute_copy.default, (1, 1, 4, 6), [1, 6, 4], False),
+    ],
+    ids=["layout_copy_squeeze", "layout_copy_unsqueeze", "user_permute_squeeze"],
+)
+def test_only_layout_copies_are_sunk_below_a_squeeze(
+    permute, shape, view_shape, sunk
+):
+    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])
+
+    result = CortexMPostponeLayoutCopyBelowSqueezePass()(builder.get_graph_module())
+
+    assert result.modified == sunk