From 2a6aeab18fa707167f01a86bb6918c2ee3e0b849 Mon Sep 17 00:00:00 2001 From: Prashant Rawat Date: Tue, 3 Nov 2026 13:59:55 -0800 Subject: [PATCH] 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 --- .../cortex_m/passes/explicit_layout_pass.py | 40 ++++++++++- .../test/test_explicit_layout_pipeline.py | 68 +++++++++++++++++++ 2 files changed, 107 insertions(+), 1 deletion(-) 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 + )