From 3a14e9e769c4c3ce95e8b413dc71e74a4b17a327 Mon Sep 17 00:00:00 2001 From: RJ Ascani Date: Wed, 19 Aug 2026 21:50:46 -0700 Subject: [PATCH 1/3] Update [ghstack-poisoned] --- backends/transforms/BUCK | 16 + .../fuse_cascaded_transpose_or_permute_ops.py | 13 + .../postpone_permute_below_squeeze_view.py | 13 +- ...replace_ops_with_channels_last_variants.py | 10 +- backends/transforms/targets.bzl | 38 ++ .../test_to_contiguous_channels_last_pass.py | 297 +++++-------- ...st_to_contiguous_channels_last_pipeline.py | 417 ++++++++++++++++++ .../to_contiguous_channels_last_pass.py | 307 +++++++++++++ 8 files changed, 916 insertions(+), 195 deletions(-) create mode 100644 backends/transforms/test/test_to_contiguous_channels_last_pipeline.py create mode 100644 backends/transforms/to_contiguous_channels_last_pass.py diff --git a/backends/transforms/BUCK b/backends/transforms/BUCK index f5029903c21..8bcff0e0563 100644 --- a/backends/transforms/BUCK +++ b/backends/transforms/BUCK @@ -1,6 +1,22 @@ load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target", "non_fbcode_target") +load("@fbcode_macros//build_defs:python_pytest.bzl", "python_pytest") load(":targets.bzl", "define_common_targets") oncall("executorch") fbcode_target(_kind = define_common_targets,) + +fbcode_target( + _kind = python_pytest, + name = "test_to_contiguous_channels_last_pipeline", + srcs = ["test/test_to_contiguous_channels_last_pipeline.py"], + compile = "with-source", + typing = False, + deps = [ + "//caffe2:torch", + ":to_contiguous_channels_last_pass", + "//executorch/exir:lib", + "//executorch/exir/dialects:lib", + "fbsource//third-party/pypi/pytest:pytest", + ], +) diff --git a/backends/transforms/fuse_cascaded_transpose_or_permute_ops.py b/backends/transforms/fuse_cascaded_transpose_or_permute_ops.py index ab75a6e0b5f..5c76b078036 100644 --- a/backends/transforms/fuse_cascaded_transpose_or_permute_ops.py +++ b/backends/transforms/fuse_cascaded_transpose_or_permute_ops.py @@ -6,6 +6,10 @@ # pyre-unsafe +from collections.abc import Callable + +import torch + from executorch.backends.transforms.channels_last_layout import ( composed_permute_target, is_channels_last_input_normalization_pair, @@ -39,6 +43,13 @@ class FuseCascadedTransposeOrPermuteOps(RemoveOrReplacePassInterface): exir_ops.edge.aten.view.default, } + def __init__( + self, + can_propagate: Callable[[torch.fx.Node], bool] | None = None, + ) -> None: + super().__init__() + self.can_propagate = can_propagate + @property def targets(self) -> list[EdgeOpOverload]: return list(self.transpose_or_permute_target) @@ -111,6 +122,8 @@ def _apply_view_to_dims( def _fuse_across_view(self, node: Node, view_node: Node) -> bool: # noqa: C901 """Fuse permute -> view(squeeze/unsqueeze) -> permute into a view_copy.""" + if self.can_propagate is not None and not self.can_propagate(view_node): + return False # view_node must have exactly one user (this permute node) if len(view_node.users) != 1: return False diff --git a/backends/transforms/postpone_permute_below_squeeze_view.py b/backends/transforms/postpone_permute_below_squeeze_view.py index 226924cbea7..fe6ea618d91 100644 --- a/backends/transforms/postpone_permute_below_squeeze_view.py +++ b/backends/transforms/postpone_permute_below_squeeze_view.py @@ -7,6 +7,7 @@ # pyre-unsafe +from collections.abc import Callable from typing import cast, List import torch @@ -35,6 +36,13 @@ class PostponePermuteOpBelowSqueezeOrUnsqueezeLikeView(RemoveOrReplacePassInterf mean the view_copy is normalized from squeeze or unsqueeze. """ + def __init__( + self, + can_propagate: Callable[[torch.fx.Node], bool] | None = None, + ) -> None: + super().__init__() + self.can_propagate = can_propagate + @property def targets(self) -> list[EdgeOpOverload]: return list(PERMUTE_COPY_TARGETS) @@ -69,6 +77,9 @@ def maybe_remove_or_replace(self, node: torch.fx.Node) -> bool: exir_ops.edge.aten.view.default, ): return False + view_node = users[0] + if self.can_propagate is not None and not self.can_propagate(view_node): + return False # If the permute_node/view_node was newly added to the # graph, it may not have the meta["val"] FakeTensor. @@ -79,8 +90,6 @@ def maybe_remove_or_replace(self, node: torch.fx.Node) -> bool: permute_node_shape = [*cast(list, get_shape(node.graph.owning_module, node))] permute_dims = cast(list, node.args[1]) - view_node = users[0] - if view_node.meta.get("val") is None: return False diff --git a/backends/transforms/replace_ops_with_channels_last_variants.py b/backends/transforms/replace_ops_with_channels_last_variants.py index 641c3925b65..40f38df574b 100644 --- a/backends/transforms/replace_ops_with_channels_last_variants.py +++ b/backends/transforms/replace_ops_with_channels_last_variants.py @@ -151,6 +151,8 @@ def __init__( self.op_map: dict[Target, ChannelsLastOpSpec] = ( op_map if op_map is not None else dict(_DEFAULT_OP_MAP) ) + self.candidate_count = 0 + self.replacement_count = 0 @staticmethod def _permute_node_input( @@ -200,6 +202,8 @@ def _permute_node_output( original_node_output.replace_all_uses_with(output) def call(self, graph_module: torch.fx.GraphModule) -> PassResult: + self.candidate_count = 0 + self.replacement_count = 0 modified = False graph = graph_module.graph @@ -208,13 +212,14 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: continue if (spec := self.op_map.get(node.target)) is None: continue + if spec.filter_fn is not None and not spec.filter_fn(node): + continue + self.candidate_count += 1 val = node.meta["val"] val = val[0] if isinstance(val, (list, tuple)) else val contiguous_dim_order = tuple(range(val.dim())) if val.dim_order() != contiguous_dim_order: continue - if spec.filter_fn is not None and not spec.filter_fn(node): - continue # In case of implicit batch size, insert also `unsqueeze_copy.default` and `squeeze_copy.dims` operators. # With `convolution`, this already happens during lowering to edge. But it doesn't happen for example with @@ -274,6 +279,7 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: graph.erase_node(node) modified = True + self.replacement_count += 1 if modified: graph.eliminate_dead_code() diff --git a/backends/transforms/targets.bzl b/backends/transforms/targets.bzl index 4283c949e13..426c96174f6 100644 --- a/backends/transforms/targets.bzl +++ b/backends/transforms/targets.bzl @@ -537,6 +537,31 @@ def define_common_targets(): ], ) + runtime.python_library( + name = "to_contiguous_channels_last_pass", + srcs = [ + "to_contiguous_channels_last_pass.py", + ], + visibility = [ + "//executorch/backends/...", + ], + deps = [ + "//caffe2:torch", + ":channels_last_layout", + ":fuse_cascaded_transpose_or_permute_ops", + ":fuse_cascaded_view_ops", + ":fuse_transpose_or_permute_op_pairs_pass", + ":postpone_permute_below_squeeze_view", + ":remove_permutes_around_elementwise_ops", + ":replace_nop_transpose_or_permute_with_view", + ":replace_ops_with_channels_last_variants", + ":replace_squeeze_unsqueeze_with_view", + "//executorch/exir:lib", + "//executorch/exir:pass_base", + "//executorch/exir/dialects:lib", + ], + ) + runtime.python_test( name = "test_replace_ops_with_channels_last_variants", srcs = [ @@ -552,6 +577,19 @@ def define_common_targets(): ], ) + runtime.python_test( + name = "test_to_contiguous_channels_last_pass", + srcs = [ + "test/test_to_contiguous_channels_last_pass.py", + ], + deps = [ + "//caffe2:torch", + ":to_contiguous_channels_last_pass", + "//executorch/exir:lib", + "//executorch/exir/dialects:lib", + ], + ) + runtime.python_test( name = "test_convert_conv1d_to_conv2d_pass", srcs = [ diff --git a/backends/transforms/test/test_to_contiguous_channels_last_pass.py b/backends/transforms/test/test_to_contiguous_channels_last_pass.py index 65dd39f09b1..54df4f51f6c 100644 --- a/backends/transforms/test/test_to_contiguous_channels_last_pass.py +++ b/backends/transforms/test/test_to_contiguous_channels_last_pass.py @@ -3,17 +3,16 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +import unittest from dataclasses import dataclass from typing import Any, Tuple -import pytest import torch -from executorch.backends.transforms.test import common -from executorch.exir import to_edge_transform_and_lower +from executorch.backends.transforms.to_contiguous_channels_last_pass import ( + ToContiguousChannelsLastPass, +) +from executorch.exir import EdgeCompileConfig, to_edge from executorch.exir.dialects._ops import ops as exir_ops -from executorch.exir.pass_base import ExportPass -from torch.fx import GraphModule -from torch.fx.passes.infra.pass_base import PassResult InputT = Tuple[Any, ...] @@ -353,9 +352,11 @@ def forward(self, x: torch.Tensor): ), "conv1d_rank3": PermuteCountTestCase(Conv1dModule(), (torch.randn(1, 2, 8),), 0), "conv2d_rank3": PermuteCountTestCase( - Conv2dModule(), (torch.randn(2, 8, 8),), 0, 2, 0, 2 + Conv2dModule(), (torch.randn(2, 8, 8),), 0, 2, 2, 2 + ), + "conv2d_rank4": PermuteCountTestCase( + Conv2dModule(), (torch.randn(1, 2, 8, 8),), 0, 0, 2, 0 ), - "conv2d_rank4": PermuteCountTestCase(Conv2dModule(), (torch.randn(1, 2, 8, 8),), 0), "conv3d_rank4": PermuteCountTestCase( Conv3dModule(), (torch.randn(2, 6, 6, 6),), 0, 2, 0, 2 ), @@ -414,25 +415,34 @@ def forward(self, x: torch.Tensor): GroupedConvModule(), (torch.randn(1, 4, 8, 8),), 0, + 0, + 2, + 0, ), "transpose_conv": PermuteCountTestCase( TransposeConvModule(), (torch.randn(1, 2, 8, 8),), 0, + 0, + 2, + 0, ), - "views": PermuteCountTestCase(ViewsModule(), (torch.rand(1, 2, 2, 4),), 0, 2, 0, 2), + "views": PermuteCountTestCase(ViewsModule(), (torch.rand(1, 2, 2, 4),), 0, 2, 4, 2), "transposes": PermuteCountTestCase( TransposesModule(), (torch.randn(1, 2, 3, 4),), 2, 0, - 2, + 1, 0, ), "maxpool2d_dilation": PermuteCountTestCase( MaxPool2dDilatedModule(), (torch.randn(1, 2, 8, 8),), 0, + 0, + 2, + 0, ), "lstm": PermuteCountTestCase( LstmModule(), @@ -440,7 +450,7 @@ def forward(self, x: torch.Tensor): 7, 19, 7, - 19, + 16, ), "groupnorm": PermuteCountTestCase( GroupNormModule(), @@ -452,16 +462,16 @@ def forward(self, x: torch.Tensor): (torch.randn(4, 8),), 11, 24, - 11, - 24, + 8, + 14, ), "multihead_attention_rank3": PermuteCountTestCase( MultiheadAttentionModule(), (torch.randn(2, 4, 8),), 12, 20, - 12, - 20, + 10, + 18, ), "cumsum_rank3_dim0": PermuteCountTestCase( CumsumModule(), @@ -474,227 +484,132 @@ def forward(self, x: torch.Tensor): 0, ), "model_1_conv_maxpool_residual_linear": PermuteCountTestCase( - Model1ConvMaxPoolResidualLinear(), (torch.randn(2, 8, 64),), 2, 7, 2, 7 + Model1ConvMaxPoolResidualLinear(), (torch.randn(2, 8, 64),), 2, 7, 6, 7 ), "model_2_conv_mha_linear_layernorm": PermuteCountTestCase( - Model2ConvMhaLinearLayerNorm(), (torch.randn(2, 8, 32),), 14, 23, 14, 23 + Model2ConvMhaLinearLayerNorm(), (torch.randn(2, 8, 32),), 14, 23, 11, 21 ), "model_3_lstm_linear": PermuteCountTestCase( - Model3LstmLinear(), (torch.randn(2, 16, 8),), 20, 58, 20, 58 + Model3LstmLinear(), (torch.randn(2, 16, 8),), 20, 58, 20, 55 ), "model_4_conv_lstm_linear_layernorm": PermuteCountTestCase( - Model4ConvLstmLinearLayerNorm(), (torch.randn(2, 8, 32),), 37, 106, 37, 106 + Model4ConvLstmLinearLayerNorm(), (torch.randn(2, 8, 32),), 37, 106, 36, 103 ), "model_5_dwconv_gelu_layernorm_avgpool": PermuteCountTestCase( - Model5DwConvGeluLayerNormAvgPool(), (torch.randn(1, 8, 16, 16),), 2, 0, 2, 0 + Model5DwConvGeluLayerNormAvgPool(), (torch.randn(1, 8, 16, 16),), 2, 0, 4, 0 ), "model_6_gru_linear": PermuteCountTestCase( - Model6GruLinear(), (torch.randn(2, 16, 8),), 20, 56, 20, 56 + Model6GruLinear(), (torch.randn(2, 16, 8),), 20, 56, 20, 55 ), "model_7_dwconv_batchnorm_linear": PermuteCountTestCase( Model7DwConvBatchNormLinear(), (torch.randn(2, 8, 64),), 2, 3, 2, 3 ), "model_8_conv_batchnorm_maxpool_residual": PermuteCountTestCase( - Model8ConvBatchNormMaxPoolResidual(), (torch.randn(1, 8, 16, 16),), 0 - ), - "model_9_dilated_conv_batchnorm_avgpool_residual": PermuteCountTestCase( - Model9DilatedConvBatchNormAvgPoolResidual(), (torch.randn(1, 8, 16, 16),), 0 - ), - "model_10_dwconv_batchnorm_linear_cat": PermuteCountTestCase( - Model10DwConvBatchNormLinearCat(), (torch.randn(2, 8, 64),), 3, 6, 3, 6 - ), - "permute_silu_permute": PermuteCountTestCase( - PermuteSiluPermute(), - (torch.randn(1, 2, 3, 4),), - 2, + Model8ConvBatchNormMaxPoolResidual(), + (torch.randn(1, 8, 16, 16),), 0, - 2, 0, - ), -} - - -cases_channels_last = { - "conv2d_rank4_channels_last": PermuteCountTestCase( - Conv2dModule(), - (torch.randn(1, 2, 8, 8).to(memory_format=torch.channels_last),), + 5, 0, ), - "conv3d_rank4_channels_last": PermuteCountTestCase( - Conv3dModule(), - (torch.randn(2, 6, 6, 6).to(memory_format=torch.channels_last),), - 0, - 2, - 0, - 2, - ), - "conv3d_rank5_channels_last": PermuteCountTestCase( - Conv3dModule(), - (torch.randn(1, 2, 6, 6, 6).to(memory_format=torch.channels_last_3d),), - 0, - ), - "linear_rank4_channels_last": PermuteCountTestCase( - LinearModule(), - (torch.randn(1, 2, 2, 8).to(memory_format=torch.channels_last),), - 1, - 3, - 1, - 3, - ), - "matmul_rank4_channels_last": PermuteCountTestCase( - MatmulModule(), - ( - torch.randn(2, 2, 2, 3).to(memory_format=torch.channels_last), - torch.randn(2, 2, 3, 4).to(memory_format=torch.channels_last), - ), + "model_9_dilated_conv_batchnorm_avgpool_residual": PermuteCountTestCase( + Model9DilatedConvBatchNormAvgPoolResidual(), + (torch.randn(1, 8, 16, 16),), 0, - 3, 0, - 3, - ), - "pixel_shuffle_channels_last": PermuteCountTestCase( - PixelShuffleModule(), - (torch.randn(1, 8, 2, 2).to(memory_format=torch.channels_last),), - 1, - 2, - 1, - 2, - ), - "grouped_conv_channels_last": PermuteCountTestCase( - GroupedConvModule(), - (torch.randn(1, 4, 8, 8).to(memory_format=torch.channels_last),), - 0, - ), - "transpose_conv_channels_last": PermuteCountTestCase( - TransposeConvModule(), - (torch.randn(1, 2, 8, 8).to(memory_format=torch.channels_last),), + 5, 0, ), - "views_channels_last": PermuteCountTestCase( - ViewsModule(), - (torch.rand(1, 2, 2, 4).to(memory_format=torch.channels_last),), - -1, # The test crashes before reaching the transpose count + "model_10_dwconv_batchnorm_linear_cat": PermuteCountTestCase( + Model10DwConvBatchNormLinearCat(), (torch.randn(2, 8, 64),), 3, 6, 3, 6 ), - "transposes_channels_last": PermuteCountTestCase( - TransposesModule(), - (torch.randn(1, 2, 3, 4).to(memory_format=torch.channels_last),), - 2, - 0, + "permute_silu_permute": PermuteCountTestCase( + PermuteSiluPermute(), + (torch.randn(1, 2, 3, 4),), 2, 0, - ), - "maxpool2d_dilation_channels_last": PermuteCountTestCase( - MaxPool2dDilatedModule(), - (torch.randn(1, 2, 8, 8).to(memory_format=torch.channels_last),), 0, - ), - "groupnorm_channels_last": PermuteCountTestCase( - GroupNormModule(), - (torch.randn(1, 4, 4, 4).to(memory_format=torch.channels_last),), - 0, - ), - "cumsum_rank4_dim3_channels_last": PermuteCountTestCase( - CumsumModule(), - (torch.randn(1, 2, 3, 4).to(memory_format=torch.channels_last), 3), 0, ), } -class ToContiguousChannelsLastPassTestPass(ExportPass): - """ - A test pass which runs the pass pipeline intended to and verifies that permutes and - views are fused as expected. - - TODO: Currently no permute-view passes are implemented, proof of concept only. - """ - - _PERMUTE_TARGETS = { - exir_ops.edge.aten.permute.default, - exir_ops.edge.aten.permute_copy.default, - exir_ops.edge.aten.transpose.int, - exir_ops.edge.aten.transpose_copy.int, - } - _VIEW_TARGETS = { - exir_ops.edge.aten._unsafe_view.default, - exir_ops.edge.aten.reshape.default, - exir_ops.edge.aten.squeeze.default, - exir_ops.edge.aten.squeeze.dim, - exir_ops.edge.aten.squeeze.dims, - exir_ops.edge.aten.squeeze_copy.default, - exir_ops.edge.aten.squeeze_copy.dim, - exir_ops.edge.aten.squeeze_copy.dims, - exir_ops.edge.aten.unsqueeze.default, - exir_ops.edge.aten.unsqueeze_copy.default, - exir_ops.edge.aten.view.default, - exir_ops.edge.aten.view_copy.default, - } +_PERMUTE_TARGETS = { + exir_ops.edge.aten.permute.default, + exir_ops.edge.aten.permute_copy.default, + exir_ops.edge.aten.transpose.int, + exir_ops.edge.aten.transpose_copy.int, + exir_ops.edge.channels_last.permute_copy.default, +} +_VIEW_TARGETS = { + exir_ops.edge.aten._unsafe_view.default, + exir_ops.edge.aten.reshape.default, + exir_ops.edge.aten.squeeze.default, + exir_ops.edge.aten.squeeze.dim, + exir_ops.edge.aten.squeeze.dims, + exir_ops.edge.aten.squeeze_copy.default, + exir_ops.edge.aten.squeeze_copy.dim, + exir_ops.edge.aten.squeeze_copy.dims, + exir_ops.edge.aten.unsqueeze.default, + exir_ops.edge.aten.unsqueeze_copy.default, + exir_ops.edge.aten.view.default, + exir_ops.edge.aten.view_copy.default, +} - def __init__(self): - super().__init__() - self.initial_permutes = 0 - self.initial_views = 0 - self.final_permutes = 0 - self.final_views = 0 - - def count_ops(self, graph_module: GraphModule, targets: set) -> int: - return sum( - 1 - for node in graph_module.graph.nodes - if node.op == "call_function" and node.target in targets - ) - def call(self, graph_module: GraphModule) -> PassResult: - self.initial_permutes = self.count_ops(graph_module, self._PERMUTE_TARGETS) - self.initial_views = self.count_ops(graph_module, self._VIEW_TARGETS) - result = super().call(graph_module) - self.final_permutes = self.count_ops(result.graph_module, self._PERMUTE_TARGETS) - self.final_views = self.count_ops(result.graph_module, self._VIEW_TARGETS) - return result +def _count_ops(graph_module: torch.fx.GraphModule, targets: set) -> int: + return sum( + node.op == "call_function" and node.target in targets + for node in graph_module.graph.nodes + ) def run_test(case: PermuteCountTestCase) -> None: case.module.eval() with torch.no_grad(): exported_program = torch.export.export(case.module, case.inputs) - test_pass = ToContiguousChannelsLastPassTestPass() - edge_program = to_edge_transform_and_lower( - exported_program, transform_passes=[test_pass] + edge_program = to_edge( + exported_program, + compile_config=EdgeCompileConfig( + _check_ir_validity=False, + _skip_dim_order=True, + ), ) - - if not ( - (test_pass.initial_permutes == case.expected_initial_permutes) - and (test_pass.initial_views == case.expected_initial_views) - and (test_pass.final_permutes == case.expected_final_permutes) - and (test_pass.final_views == case.expected_final_views) - ): - raise AssertionError( - f"Operator counts do not match for case {case.module.__class__.__name__}\n" - f"Expected initial permutes: {case.expected_initial_permutes}, got: {test_pass.initial_permutes}\n" - f"Expected initial views: {case.expected_initial_views}, got: {test_pass.initial_views}\n" - f"Expected final permutes: {case.expected_final_permutes}, got: {test_pass.final_permutes}\n" - f"Expected final views: {case.expected_final_views}, got: {test_pass.final_views}\n" - ) - + initial_graph = edge_program.exported_program().graph_module + initial_permutes = _count_ops(initial_graph, _PERMUTE_TARGETS) + initial_views = _count_ops(initial_graph, _VIEW_TARGETS) + + layout_pass = ToContiguousChannelsLastPass(edge_program.exported_program()) + transformed = edge_program.transform([layout_pass]) + final_graph = transformed.exported_program().graph_module + final_permutes = _count_ops(final_graph, _PERMUTE_TARGETS) + final_views = _count_ops(final_graph, _VIEW_TARGETS) + + assert initial_permutes == case.expected_initial_permutes + assert initial_views == case.expected_initial_views + assert final_permutes == case.expected_final_permutes + assert final_views == case.expected_final_views ref_result = exported_program.module()(*case.inputs) - edge_result = edge_program.exported_program().module()(*case.inputs) + edge_result = transformed.exported_program().module()(*case.inputs) assert torch.allclose(ref_result, edge_result, atol=1e-6) -@pytest.mark.skip( - reason="Proof of concept - currently no permute-view passes implemented." -) -@common.parametrize("case", cases) -def test_permute_view_counts(case: PermuteCountTestCase) -> None: - run_test(case) +_EXPECTED_AGGREGATE_COUNTS = (138, 339, 156, 315) -xfails = {"views_channels_last": "Views are not supported for channels last tensors"} - +class TestToContiguousChannelsLastPass(unittest.TestCase): + def test_aggregate_permute_view_counts(self) -> None: + self.assertEqual( + ( + sum(case.expected_initial_permutes for case in cases.values()), + sum(case.expected_initial_views for case in cases.values()), + sum(case.expected_final_permutes for case in cases.values()), + sum(case.expected_final_views for case in cases.values()), + ), + _EXPECTED_AGGREGATE_COUNTS, + ) -@pytest.mark.skip( - reason="Proof of concept - currently no permute-view passes implemented." -) -@common.parametrize("case", cases_channels_last, xfails=xfails) -def test_permute_view_counts_channels_last(case: PermuteCountTestCase) -> None: - run_test(case) + def test_permute_view_counts(self) -> None: + for name, case in cases.items(): + with self.subTest(name=name): + run_test(case) diff --git a/backends/transforms/test/test_to_contiguous_channels_last_pipeline.py b/backends/transforms/test/test_to_contiguous_channels_last_pipeline.py new file mode 100644 index 00000000000..e9dcece4293 --- /dev/null +++ b/backends/transforms/test/test_to_contiguous_channels_last_pipeline.py @@ -0,0 +1,417 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import pytest +import torch + +from executorch.backends.transforms.to_contiguous_channels_last_pass import ( + ToContiguousChannelsLastPass, +) +from executorch.exir import EdgeCompileConfig, to_edge +from executorch.exir.dialects._ops import ops as exir_ops + + +class ConvChain(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv1 = torch.nn.Conv2d(4, 4, 3, padding=1) + self.conv2 = torch.nn.Conv2d(4, 4, 3, padding=1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.conv2(torch.relu(self.conv1(x))) + + +class ConvOnly(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv = torch.nn.Conv2d(4, 4, 3, padding=1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.conv(x) + + +class DynamicViewConv(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv = torch.nn.Conv2d(4, 4, 3, padding=1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.conv(x.view(x.shape[0], x.shape[1], x.shape[2], x.shape[3])) + + +class ConvThenLinear(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv = torch.nn.Conv2d(4, 4, 3, padding=1) + self.linear = torch.nn.Linear(4 * 8 * 8, 3) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.linear(self.conv(x).flatten(1)) + + +class UserPermute(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x.permute(0, 2, 3, 1) + + +class ConvChannelBias(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv = torch.nn.Conv2d(4, 4, 3, padding=1) + self.register_buffer("channel_bias", torch.randn(1, 4, 1, 1)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.conv(x) + self.channel_bias + + +class PadConv(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv = torch.nn.Conv2d(4, 4, 3) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.conv(torch.nn.functional.pad(x, (1, 1, 1, 1))) + + +class ConvSoftmax(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv = torch.nn.Conv2d(4, 4, 3, padding=1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.softmax(self.conv(x), dim=-1) + + +class ConvRuntimeBiasConv(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv1 = torch.nn.Conv2d(4, 4, 3, padding=1) + self.conv2 = torch.nn.Conv2d(4, 4, 3, padding=1) + + def forward(self, x: torch.Tensor, bias: torch.Tensor) -> torch.Tensor: + return self.conv2(self.conv1(x) + bias) + + +class ConvSpatialBufferConv(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv1 = torch.nn.Conv2d(4, 4, 3, padding=1) + self.conv2 = torch.nn.Conv2d(4, 4, 3, padding=1) + self.register_buffer("bias", torch.randn(4, 8, 8)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.conv2(self.conv1(x) + self.bias) + + +def _edge(module: torch.nn.Module, inputs: tuple[torch.Tensor, ...]): + exported = torch.export.export(module.eval(), inputs) + return to_edge( + exported, + compile_config=EdgeCompileConfig( + _check_ir_validity=False, + _skip_dim_order=True, + ), + ) + + +def _count(graph_module: torch.fx.GraphModule, target: object) -> int: + return sum( + node.op == "call_function" and node.target == target + for node in graph_module.graph.nodes + ) + + +def test_conv_chain_folds_to_boundary_copies() -> None: + torch.manual_seed(0) + module = ConvChain().eval() + inputs = (torch.randn(1, 4, 8, 8),) + expected = module(*inputs) + edge = _edge(module, inputs) + layout_pass = ToContiguousChannelsLastPass(edge.exported_program()) + + transformed = edge.transform([layout_pass]) + graph_module = transformed.exported_program().graph_module + + assert _count(graph_module, exir_ops.edge.channels_last.convolution.default) == 2 + assert _count(graph_module, exir_ops.edge.channels_last.permute_copy.default) == 2 + assert layout_pass.report.candidate_anchor_count == 2 + assert layout_pass.report.converted_anchor_count == 2 + assert layout_pass.report.inserted_copy_count == 4 + assert layout_pass.report.eliminated_copy_count == 2 + assert layout_pass.report.boundary_copy_count == 2 + assert layout_pass.report.internal_copy_count == 0 + assert layout_pass.report.unknown_copy_count == 0 + assert layout_pass.report.boundary_copy_bytes == 2048 + assert layout_pass.report.internal_copy_bytes == 0 + assert layout_pass.report.unknown_copy_bytes == 0 + assert layout_pass.report.copies_with_unknown_size == 0 + actual = transformed.exported_program().module()(*inputs) + assert torch.allclose(actual, expected, atol=1e-6) + + +def test_strict_rejects_internal_layout_copy() -> None: + inputs = (torch.randn(1, 4, 8, 8),) + edge = _edge(ConvThenLinear().eval(), inputs) + + with pytest.raises(RuntimeError, match="left .* internal"): + ToContiguousChannelsLastPass(edge.exported_program(), strict=True).call( + edge.exported_program().graph_module + ) + + +def test_strict_rejects_supported_anchor_with_noncontiguous_dim_order() -> None: + module = ConvOnly().eval().to(memory_format=torch.channels_last) + inputs = (torch.randn(1, 4, 8, 8).to(memory_format=torch.channels_last),) + exported = torch.export.export(module, inputs) + edge = to_edge( + exported, + compile_config=EdgeCompileConfig(_check_ir_validity=False), + ) + layout_pass = ToContiguousChannelsLastPass( + edge.exported_program(), + strict=True, + ) + + with pytest.raises(RuntimeError, match="0 converted of 1 candidate anchors"): + layout_pass.call(edge.exported_program().graph_module) + + assert layout_pass.report.candidate_anchor_count == 1 + assert layout_pass.report.converted_anchor_count == 0 + + +def test_strict_rejects_unknown_boundary_copy_size() -> None: + inputs = (torch.randn(1, 4, 8, 8),) + exported = torch.export.export( + ConvOnly().eval(), + inputs, + dynamic_shapes={"x": {2: torch.export.Dim("height", min=4, max=16)}}, + ) + edge = to_edge( + exported, + compile_config=EdgeCompileConfig( + _check_ir_validity=False, + _skip_dim_order=True, + ), + ) + layout_pass = ToContiguousChannelsLastPass(edge.exported_program(), strict=True) + + with pytest.raises(RuntimeError, match="unknown sizes"): + layout_pass.call(edge.exported_program().graph_module) + + assert layout_pass.report.boundary_copy_count == 2 + assert layout_pass.report.boundary_copy_bytes == 0 + assert layout_pass.report.copies_with_unknown_size == 2 + + +def test_dynamic_conv_chain_eliminates_internal_layout_copies() -> None: + inputs = (torch.randn(1, 4, 8, 8),) + exported = torch.export.export( + ConvChain().eval(), + inputs, + dynamic_shapes={"x": {2: torch.export.Dim("height", min=4, max=16)}}, + ) + edge = to_edge( + exported, + compile_config=EdgeCompileConfig( + _check_ir_validity=False, + _skip_dim_order=True, + ), + ) + layout_pass = ToContiguousChannelsLastPass(edge.exported_program()) + + edge.transform([layout_pass]) + + assert layout_pass.report.boundary_copy_count == 2 + assert layout_pass.report.internal_copy_count == 0 + assert layout_pass.report.copies_with_unknown_size == 2 + + +def test_dynamic_view_does_not_hide_input_boundary_copy() -> None: + inputs = (torch.randn(1, 4, 8, 8),) + exported = torch.export.export( + DynamicViewConv().eval(), + inputs, + dynamic_shapes={"x": {2: torch.export.Dim("height", min=4, max=16)}}, + ) + edge = to_edge( + exported, + compile_config=EdgeCompileConfig( + _check_ir_validity=False, + _skip_dim_order=True, + ), + ) + layout_pass = ToContiguousChannelsLastPass(edge.exported_program()) + + edge.transform([layout_pass]) + + assert layout_pass.report.boundary_copy_count == 2 + assert layout_pass.report.internal_copy_count == 0 + assert layout_pass.report.copies_with_unknown_size == 2 + + +def test_backend_can_block_layout_propagation_at_a_node() -> None: + torch.manual_seed(0) + module = ConvChain().eval() + inputs = (torch.randn(1, 4, 8, 8),) + expected = module(*inputs) + edge = _edge(module, inputs) + layout_pass = ToContiguousChannelsLastPass( + edge.exported_program(), + can_propagate=lambda node: node.target != exir_ops.edge.aten.relu.default, + ) + + transformed = edge.transform([layout_pass]) + + assert layout_pass.report.boundary_copy_count == 2 + assert layout_pass.report.internal_copy_count == 2 + actual = transformed.exported_program().module()(*inputs) + torch.testing.assert_close(actual, expected) + + +def test_backend_barrier_blocks_view_reordering() -> None: + module = DynamicViewConv().eval() + inputs = (torch.randn(1, 4, 8, 8),) + expected = module(*inputs) + edge = _edge(module, inputs) + layout_pass = ToContiguousChannelsLastPass( + edge.exported_program(), + can_propagate=lambda node: node.target + not in ( + exir_ops.edge.aten.view.default, + exir_ops.edge.aten.view_copy.default, + ), + ) + + transformed = edge.transform([layout_pass]) + + assert layout_pass.report.boundary_copy_count == 1 + assert layout_pass.report.internal_copy_count == 1 + actual = transformed.exported_program().module()(*inputs) + torch.testing.assert_close(actual, expected) + + +@pytest.mark.parametrize("module", [ConvChannelBias(), PadConv()]) +def test_one_sided_propagation_reaches_graph_boundary( + module: torch.nn.Module, +) -> None: + torch.manual_seed(0) + module.eval() + inputs = (torch.randn(1, 4, 8, 8),) + expected = module(*inputs) + edge = _edge(module, inputs) + layout_pass = ToContiguousChannelsLastPass( + edge.exported_program(), + layout_pad_target=( + exir_ops.edge.channels_last.constant_pad_nd.default + if isinstance(module, PadConv) + else None + ), + ) + + transformed = edge.transform([layout_pass]) + + assert layout_pass.report.boundary_copy_count == 2 + assert layout_pass.report.internal_copy_count == 0 + assert layout_pass.report.unknown_copy_count == 0 + actual = transformed.exported_program().module()(*inputs) + assert torch.allclose(actual, expected, atol=1e-6) + + if isinstance(module, PadConv): + assert ( + _count( + transformed.exported_program().graph_module, + exir_ops.edge.channels_last.constant_pad_nd.default, + ) + == 1 + ) + + +def test_softmax_blocks_layout_propagation() -> None: + module = ConvSoftmax().eval() + inputs = (torch.randn(1, 4, 8, 8),) + expected = module(*inputs) + edge = _edge(module, inputs) + layout_pass = ToContiguousChannelsLastPass(edge.exported_program()) + + transformed = edge.transform([layout_pass]) + + assert layout_pass.report.boundary_copy_count == 1 + assert layout_pass.report.internal_copy_count == 1 + softmax = next( + node + for node in transformed.exported_program().graph.nodes + if node.target + in (exir_ops.edge.aten._softmax.default, exir_ops.edge.aten.softmax.int) + ) + assert softmax.args[1] in (-1, 3) + actual = transformed.exported_program().module()(*inputs) + torch.testing.assert_close(actual, expected) + + +@pytest.mark.parametrize( + "module, inputs", + [ + ( + ConvRuntimeBiasConv(), + (torch.randn(1, 4, 8, 8), torch.randn(4, 1, 1)), + ), + (ConvSpatialBufferConv(), (torch.randn(1, 4, 8, 8),)), + ], +) +def test_boundary_propagation_rejects_unsafe_broadcast_rewrites( + module: torch.nn.Module, inputs: tuple[torch.Tensor, ...] +) -> None: + module.eval() + expected = module(*inputs) + edge = _edge(module, inputs) + layout_pass = ToContiguousChannelsLastPass(edge.exported_program()) + + transformed = edge.transform([layout_pass]) + + assert layout_pass.report.boundary_copy_count == 2 + assert layout_pass.report.internal_copy_count == 2 + actual = transformed.exported_program().module()(*inputs) + torch.testing.assert_close(actual, expected) + + +def test_layout_copy_report_is_idempotent() -> None: + inputs = (torch.randn(1, 4, 8, 8),) + edge = _edge(ConvChain().eval(), inputs) + first_pass = ToContiguousChannelsLastPass(edge.exported_program()) + transformed = edge.transform([first_pass]) + second_pass = ToContiguousChannelsLastPass(transformed.exported_program()) + + transformed.transform([second_pass]) + + assert second_pass.report.inserted_copy_count == 0 + assert second_pass.report.eliminated_copy_count == 0 + assert second_pass.report.candidate_anchor_count == 0 + assert second_pass.report.converted_anchor_count == 0 + + +def test_user_permute_is_not_reported_as_layout_copy() -> None: + inputs = (torch.randn(1, 4, 8, 8),) + edge = _edge(UserPermute(), inputs) + layout_pass = ToContiguousChannelsLastPass(edge.exported_program(), op_map={}) + + transformed = edge.transform([layout_pass]) + permutes = [ + node + for node in transformed.exported_program().graph.nodes + if node.target == exir_ops.edge.aten.permute_copy.default + ] + + assert len(permutes) == 1 + assert ( + _count( + transformed.exported_program().graph_module, + exir_ops.edge.channels_last.permute_copy.default, + ) + == 0 + ) + assert layout_pass.report.inserted_copy_count == 0 + assert layout_pass.report.boundary_copy_count == 0 + assert layout_pass.report.internal_copy_count == 0 diff --git a/backends/transforms/to_contiguous_channels_last_pass.py b/backends/transforms/to_contiguous_channels_last_pass.py new file mode 100644 index 00000000000..1df7dcd49d1 --- /dev/null +++ b/backends/transforms/to_contiguous_channels_last_pass.py @@ -0,0 +1,307 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-unsafe + +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any + +import torch + +from executorch.backends.transforms.channels_last_layout import is_layout_copy +from executorch.backends.transforms.fuse_cascaded_transpose_or_permute_ops import ( + FuseCascadedTransposeOrPermuteOps, +) +from executorch.backends.transforms.fuse_cascaded_view_ops import FuseCascadedViewOps +from executorch.backends.transforms.fuse_transpose_or_permute_op_pairs_pass import ( + FuseTransposeOrPermuteOpPairsPass, +) +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_nop_transpose_or_permute_with_view import ( + ReplaceNopTransposeOrPermuteWithViewPass, +) +from executorch.backends.transforms.replace_ops_with_channels_last_variants import ( + ChannelsLastOpSpec, + ReplaceOpsWithChannelsLastVariants, +) +from executorch.backends.transforms.replace_squeeze_unsqueeze_with_view import ( + ReplaceSqueezeAndUnsqueezeWithViewPass, +) +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.node import Target + +_BOUNDARY_TRANSPARENT_TARGETS = { + exir_ops.edge.aten.squeeze_copy.default, + exir_ops.edge.aten.squeeze_copy.dim, + exir_ops.edge.aten.squeeze_copy.dims, + exir_ops.edge.aten.unsqueeze_copy.default, + exir_ops.edge.aten.view.default, + exir_ops.edge.aten.view_copy.default, +} +try: + _BOUNDARY_TRANSPARENT_TARGETS.update( + { + exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + } + ) +except AttributeError: + pass + + +@dataclass(frozen=True) +class ChannelsLastLayoutReport: + candidate_anchor_count: int = 0 + converted_anchor_count: int = 0 + inserted_copy_count: int = 0 + eliminated_copy_count: int = 0 + boundary_copy_count: int = 0 + internal_copy_count: int = 0 + unknown_copy_count: int = 0 + boundary_copy_bytes: int = 0 + internal_copy_bytes: int = 0 + unknown_copy_bytes: int = 0 + copies_with_unknown_size: int = 0 + internal_copy_nodes: tuple[str, ...] = () + unknown_copy_nodes: tuple[str, ...] = () + + +class ToContiguousChannelsLastPass(ExportPass): + """Build and optimize explicit contiguous-NHWC regions. + + The pass replaces selected NCHW operators with channels-last dialect + anchors surrounded by ``channels_last.permute_copy`` nodes. It then runs + the common data-movement optimizers to a fixed point and reports only the + surviving layout copies. Strict mode rejects structurally unsafe copies and + supported source anchors that were not converted. It does not reject + user-authored permutes or estimate peak arena usage after memory planning. + Additional permutable ops must be layout-equivariant without argument + remapping. ``can_propagate`` is consulted by every transform that moves a + layout copy across a graph node. + """ + + _MAX_OPTIMIZATION_ITERATIONS = 8 + + def __init__( + self, + exported_program: ExportedProgram, + op_map: dict[Target, ChannelsLastOpSpec] | None = None, + extra_permutable_ops: set[Target] | None = None, + can_propagate: Callable[[torch.fx.Node], bool] | None = None, + layout_pad_target: Target | None = None, + strict: bool = False, + ) -> None: + super().__init__() + self.exported_program = exported_program + self.op_map = op_map + self.extra_permutable_ops: set[Target] = set() + if extra_permutable_ops: + self.extra_permutable_ops |= extra_permutable_ops + self.can_propagate = can_propagate + self.layout_pad_target = layout_pad_target + self.strict = strict + self.report = ChannelsLastLayoutReport() + + def call(self, graph_module: torch.fx.GraphModule) -> PassResult: + existing_copy_count = len(self._layout_copy_nodes(graph_module)) + replacement_pass = ReplaceOpsWithChannelsLastVariants( + self.exported_program, + op_map=self.op_map, + ) + replacement = replacement_pass.call(graph_module) + graph_module = replacement.graph_module + modified = replacement.modified + preoptimization_copy_count = len(self._layout_copy_nodes(graph_module)) + inserted_copy_count = max(0, preoptimization_copy_count - existing_copy_count) + + for iteration in range(self._MAX_OPTIMIZATION_ITERATIONS): + iteration_modified = False + for transform in self._optimization_passes(): + result = transform.call(graph_module) + graph_module = result.graph_module + iteration_modified |= result.modified + + modified |= iteration_modified + if not iteration_modified: + break + if iteration == self._MAX_OPTIMIZATION_ITERATIONS - 1: + raise RuntimeError( + "Channels-last layout optimization did not converge after " + f"{self._MAX_OPTIMIZATION_ITERATIONS} iterations." + ) + + self.report = self._build_report( + graph_module, + replacement_pass.candidate_count, + replacement_pass.replacement_count, + inserted_copy_count, + preoptimization_copy_count, + ) + if self.strict and ( + self.report.candidate_anchor_count != self.report.converted_anchor_count + or self.report.internal_copy_count + or self.report.unknown_copy_count + or self.report.copies_with_unknown_size + ): + raise RuntimeError( + "Channels-last layout optimization left " + f"{self.report.converted_anchor_count} converted of " + f"{self.report.candidate_anchor_count} candidate anchors, " + f"{self.report.internal_copy_count} internal and " + f"{self.report.unknown_copy_count} unknown copies, with " + f"{self.report.copies_with_unknown_size} unknown sizes. " + f"Internal nodes: {self.report.internal_copy_nodes}; " + f"unknown nodes: {self.report.unknown_copy_nodes}." + ) + + return PassResult(graph_module, modified) + + def _optimization_passes(self) -> tuple[ExportPass, ...]: + return ( + ReplaceSqueezeAndUnsqueezeWithViewPass(), + ReplaceNopTransposeOrPermuteWithViewPass(), + PostponePermuteOpBelowSqueezeOrUnsqueezeLikeView( + can_propagate=self.can_propagate + ), + FuseCascadedViewOps(), + FuseCascadedTransposeOrPermuteOps(can_propagate=self.can_propagate), + RemovePermutesAroundElementwiseOps( + self.extra_permutable_ops, + exported_program=self.exported_program, + allow_layout_boundary_propagation=True, + layout_pad_target=self.layout_pad_target, + can_propagate=self.can_propagate, + ), + FuseTransposeOrPermuteOpPairsPass(can_propagate=self.can_propagate), + FuseCascadedViewOps(), + FuseCascadedTransposeOrPermuteOps(can_propagate=self.can_propagate), + ) + + @staticmethod + def _layout_copy_nodes( + graph_module: torch.fx.GraphModule, + ) -> list[torch.fx.Node]: + return [node for node in graph_module.graph.nodes if is_layout_copy(node)] + + def _build_report( + self, + graph_module: torch.fx.GraphModule, + candidate_anchor_count: int, + converted_anchor_count: int, + inserted_copy_count: int, + preoptimization_copy_count: int, + ) -> ChannelsLastLayoutReport: + boundary_nodes: list[torch.fx.Node] = [] + internal_nodes: list[torch.fx.Node] = [] + unknown_nodes: list[torch.fx.Node] = [] + + for node in self._layout_copy_nodes(graph_module): + dims = self._normalized_dims(node) + if dims is None: + unknown_nodes.append(node) + elif self._reaches_user_input(node) or self._reaches_graph_output(node): + boundary_nodes.append(node) + else: + internal_nodes.append(node) + + boundary_bytes, boundary_unknown = self._copy_bytes(boundary_nodes) + internal_bytes, internal_unknown = self._copy_bytes(internal_nodes) + unknown_bytes, unknown_unknown = self._copy_bytes(unknown_nodes) + surviving_count = len(boundary_nodes) + len(internal_nodes) + len(unknown_nodes) + return ChannelsLastLayoutReport( + candidate_anchor_count=candidate_anchor_count, + converted_anchor_count=converted_anchor_count, + inserted_copy_count=inserted_copy_count, + eliminated_copy_count=max(0, preoptimization_copy_count - surviving_count), + boundary_copy_count=len(boundary_nodes), + internal_copy_count=len(internal_nodes), + unknown_copy_count=len(unknown_nodes), + boundary_copy_bytes=boundary_bytes, + internal_copy_bytes=internal_bytes, + unknown_copy_bytes=unknown_bytes, + copies_with_unknown_size=( + boundary_unknown + internal_unknown + unknown_unknown + ), + internal_copy_nodes=tuple(node.name for node in internal_nodes), + unknown_copy_nodes=tuple(node.name for node in unknown_nodes), + ) + + def _reaches_user_input(self, node: torch.fx.Node) -> bool: + current = node.args[0] if node.args else None + visited: set[torch.fx.Node] = set() + while isinstance(current, torch.fx.Node) and current not in visited: + visited.add(current) + if current.op == "placeholder": + return current.name in self.exported_program.graph_signature.user_inputs + if self.can_propagate is not None and not self.can_propagate(current): + return False + if ( + current.op != "call_function" + or current.target not in _BOUNDARY_TRANSPARENT_TARGETS + or not current.args + or not isinstance(current.args[0], torch.fx.Node) + ): + return False + current = current.args[0] + return False + + def _reaches_graph_output(self, node: torch.fx.Node) -> bool: + pending = list(node.users) + visited: set[torch.fx.Node] = set() + reached_output = False + while pending: + current = pending.pop() + if current in visited: + continue + visited.add(current) + if current.op == "output": + reached_output = True + continue + if self.can_propagate is not None and not self.can_propagate(current): + return False + if ( + current.op != "call_function" + or current.target not in _BOUNDARY_TRANSPARENT_TARGETS + or not current.users + ): + return False + pending.extend(current.users) + return reached_output + + @staticmethod + def _normalized_dims(node: torch.fx.Node) -> list[int] | None: + if len(node.args) < 2 or not isinstance(node.args[1], (list, tuple)): + return None + dims = list(node.args[1]) + if not all(isinstance(dim, int) for dim in dims): + return None + rank = len(dims) + normalized = [dim + rank if dim < 0 else dim for dim in dims] + if sorted(normalized) != list(range(rank)): + return None + return normalized + + @staticmethod + def _copy_bytes(nodes: list[torch.fx.Node]) -> tuple[int, int]: + known_bytes = 0 + unknown_count = 0 + for node in nodes: + val: Any = node.meta.get("val") + if not isinstance(val, torch.Tensor) or not all( + isinstance(dim, int) for dim in val.shape + ): + unknown_count += 1 + continue + known_bytes += val.numel() * val.element_size() + return known_bytes, unknown_count From 577732aaaeccee729a2029a8b420c2b75fa2b108 Mon Sep 17 00:00:00 2001 From: RJ Ascani Date: Fri, 21 Aug 2026 11:54:57 -0700 Subject: [PATCH 2/3] Update [ghstack-poisoned] --- .../transforms/test/test_to_contiguous_channels_last_pass.py | 1 - 1 file changed, 1 deletion(-) diff --git a/backends/transforms/test/test_to_contiguous_channels_last_pass.py b/backends/transforms/test/test_to_contiguous_channels_last_pass.py index 2f5a55ac1b2..c02f068dcc1 100644 --- a/backends/transforms/test/test_to_contiguous_channels_last_pass.py +++ b/backends/transforms/test/test_to_contiguous_channels_last_pass.py @@ -1001,7 +1001,6 @@ def test_one_sided_propagation_reaches_graph_boundary( assert torch.allclose(actual, expected, atol=1e-6) if isinstance(module, PadConv): - # The pad stays on the aten operator; only its argument is remapped. assert ( _count( transformed.exported_program().graph_module, From 07f2fd373113fcda514319cdd07078c141d38d11 Mon Sep 17 00:00:00 2001 From: RJ Ascani Date: Fri, 21 Aug 2026 12:12:54 -0700 Subject: [PATCH 3/3] Update [ghstack-poisoned] --- .../absorb_boundary_layout_copies.py | 208 ++++++++++++++++ backends/transforms/targets.bzl | 34 +++ .../test_absorb_boundary_layout_copies.py | 228 ++++++++++++++++++ 3 files changed, 470 insertions(+) create mode 100644 backends/transforms/absorb_boundary_layout_copies.py create mode 100644 backends/transforms/test/test_absorb_boundary_layout_copies.py diff --git a/backends/transforms/absorb_boundary_layout_copies.py b/backends/transforms/absorb_boundary_layout_copies.py new file mode 100644 index 00000000000..d97ef62e900 --- /dev/null +++ b/backends/transforms/absorb_boundary_layout_copies.py @@ -0,0 +1,208 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-unsafe + +from dataclasses import dataclass, field + +import torch + +from executorch.backends.transforms.channels_last_layout import is_layout_copy +from executorch.exir import ExportedProgram +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import ExportPass, PassResult + +# Elementwise per-tensor quantization is layout-agnostic, so a layout copy on +# the far side of one is still a boundary copy. Per-channel quantization is not: +# its axis is dimension-dependent. +_LAYOUT_AGNOSTIC_TARGETS = frozenset( + { + exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + } +) + + +def _inverse(dims: tuple[int, ...]) -> list[int]: + inverse = [0] * len(dims) + for position, dim in enumerate(dims): + inverse[dim] = position + return inverse + + +@dataclass(frozen=True) +class BoundaryLayoutContract: + """Which method inputs and outputs changed layout, and to what. + + An entry ``{0: (0, 2, 3, 1)}`` in ``inputs`` means argument 0 must now be + passed as ``argument.permute(0, 2, 3, 1)``. An entry in ``outputs`` means + the returned tensor needs the same permutation applied to recover what the + method used to return. + """ + + inputs: dict[int, tuple[int, ...]] = field(default_factory=dict) + outputs: dict[int, tuple[int, ...]] = field(default_factory=dict) + + def __bool__(self) -> bool: + return bool(self.inputs or self.outputs) + + +class AbsorbBoundaryLayoutCopies(ExportPass): + """Move layout copies that sit on the method boundary into the signature. + + A layout region formed by ``ToContiguousChannelsLastPass`` is bracketed by + ``channels_last.permute_copy``. Those in the interior cancel against each + other; the ones on the boundary have nothing to cancel against and survive. + Deleting them and declaring the corresponding method input or output to be + channels-last moves the transpose to the caller, which is free whenever the + caller already has the data in that layout. + + Run this *after* region formation. Permuting the boundary first and hoping + the copies cancel is measurably worse: it inserts copies into graphs that + have no anchors at all, where nothing can cancel them. + + A copy need not touch the boundary directly. Quantized graphs put a + per-tensor ``quantize``/``dequantize`` in between, which reorders nothing, + so the search walks through those and relabels them on the way. + + Changing a method's layout is caller-visible, so the applied changes are + reported in ``contract`` rather than assumed. + """ + + def __init__(self, exported_program: ExportedProgram) -> None: + super().__init__() + self.exported_program = exported_program + self.contract = BoundaryLayoutContract() + + def call(self, graph_module: torch.fx.GraphModule) -> PassResult: + if graph_module is not self.exported_program.graph_module: + raise RuntimeError( + "AbsorbBoundaryLayoutCopies rewrites the ExportedProgram's graph " + "and signature together; run it as its own transform rather than " + "after a pass that replaces the graph module." + ) + + inputs = self._absorb_inputs(graph_module) + outputs = self._absorb_outputs(graph_module) + self.contract = BoundaryLayoutContract(inputs=inputs, outputs=outputs) + + modified = bool(self.contract) + if modified: + graph_module.graph.eliminate_dead_code() + graph_module.recompile() + return PassResult(graph_module, modified) + + def _forward_to_copies(self, node: torch.fx.Node): + """Follow every path out of ``node`` until it reaches a layout copy. + + Returns the layout-agnostic nodes crossed, the copies terminating the + paths, and the permutation they agree on; ``None`` if any path ends + somewhere else or the copies disagree. + """ + interior: list[torch.fx.Node] = [] + copies: list[torch.fx.Node] = [] + dims: tuple[int, ...] | None = None + frontier = [node] + + while frontier: + users = list(frontier.pop().users) + if not users: + return None + for user in users: + if is_layout_copy(user): + user_dims = tuple(user.args[1]) + if dims is not None and user_dims != dims: + # A residual block feeds one value to several branches, + # each with its own copy. They collapse to a single + # contract entry only if they agree. + return None + dims = user_dims + copies.append(user) + elif user.target in _LAYOUT_AGNOSTIC_TARGETS: + if user in interior: + continue + interior.append(user) + frontier.append(user) + else: + return None + + return None if dims is None else (interior, copies, dims) + + def _backward_to_copy(self, result: torch.fx.Node): + """Walk back from a returned value to the layout copy that produced it.""" + interior: list[torch.fx.Node] = [] + current = result + + while True: + if is_layout_copy(current): + if len(current.users) != 1: + return None + source = current.args[0] + if not isinstance(source, torch.fx.Node): + return None + return interior, current, source, tuple(current.args[1]) + if ( + current.target not in _LAYOUT_AGNOSTIC_TARGETS + or len(current.users) != 1 + ): + return None + interior.append(current) + current = current.args[0] + if not isinstance(current, torch.fx.Node): + return None + + def _absorb_inputs(self, graph_module) -> dict[int, tuple[int, ...]]: + user_inputs = list(self.exported_program.graph_signature.user_inputs) + absorbed: dict[int, tuple[int, ...]] = {} + + for node in list(graph_module.graph.nodes): + if node.op != "placeholder" or node.name not in user_inputs: + continue + found = self._forward_to_copies(node) + if found is None: + continue + interior, copies, dims = found + + for member in (node, *interior): + member.meta["val"] = member.meta["val"].permute(dims) + for copy in copies: + copy.replace_all_uses_with(copy.args[0]) + graph_module.graph.erase_node(copy) + absorbed[user_inputs.index(node.name)] = dims + + return absorbed + + def _absorb_outputs(self, graph_module) -> dict[int, tuple[int, ...]]: + output_node = graph_module.graph.output_node() + results = list(output_node.args[0]) + specs = self.exported_program.graph_signature.output_specs + absorbed: dict[int, tuple[int, ...]] = {} + + for index, result in enumerate(results): + if not isinstance(result, torch.fx.Node): + continue + found = self._backward_to_copy(result) + if found is None: + continue + interior, copy, source, dims = found + + for member in interior: + member.meta["val"] = member.meta["val"].permute(_inverse(dims)) + if interior: + interior[-1].replace_input_with(copy, source) + else: + results[index] = source + # The manager re-derives the signature, but the direct + # exported_program= path does not. + if index < len(specs) and getattr(specs[index].arg, "name", None) == ( + result.name + ): + specs[index].arg.name = source.name + absorbed[index] = dims + + if absorbed: + output_node.args = (results,) + return absorbed diff --git a/backends/transforms/targets.bzl b/backends/transforms/targets.bzl index e15e089b051..fb7b3f54838 100644 --- a/backends/transforms/targets.bzl +++ b/backends/transforms/targets.bzl @@ -562,6 +562,40 @@ def define_common_targets(): ], ) + runtime.python_library( + name = "absorb_boundary_layout_copies", + srcs = [ + "absorb_boundary_layout_copies.py", + ], + visibility = [ + "//executorch/backends/...", + ], + deps = [ + "//caffe2:torch", + ":channels_last_layout", + "//executorch/exir:lib", + "//executorch/exir:pass_base", + "//executorch/exir/dialects:lib", + ], + ) + + runtime.python_test( + name = "test_absorb_boundary_layout_copies", + srcs = [ + "test/test_absorb_boundary_layout_copies.py", + # The permute-count matrix the absorption totals are measured on. + "test/test_to_contiguous_channels_last_pass.py", + ], + deps = [ + "//caffe2:torch", + ":absorb_boundary_layout_copies", + ":to_contiguous_channels_last_pass", + "//executorch/exir:lib", + "//executorch/exir/dialects:lib", + "fbsource//third-party/pypi/pytest:pytest", + ], + ) + runtime.python_test( name = "test_replace_ops_with_channels_last_variants", srcs = [ diff --git a/backends/transforms/test/test_absorb_boundary_layout_copies.py b/backends/transforms/test/test_absorb_boundary_layout_copies.py new file mode 100644 index 00000000000..e3675d51d05 --- /dev/null +++ b/backends/transforms/test/test_absorb_boundary_layout_copies.py @@ -0,0 +1,228 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import pytest +import torch + +from executorch.backends.transforms.absorb_boundary_layout_copies import ( + AbsorbBoundaryLayoutCopies, +) +from executorch.backends.transforms.to_contiguous_channels_last_pass import ( + ToContiguousChannelsLastPass, +) +from executorch.exir import EdgeCompileConfig, to_edge +from executorch.exir.dialects._ops import ops as exir_ops + +_LAYOUT_COPY = exir_ops.edge.channels_last.permute_copy.default + + +class Conv(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv = torch.nn.Conv2d(4, 4, 3, padding=1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.conv(x) + + +class ResidualConvPool(torch.nn.Module): + """One input feeding two branches, so the region brackets it twice.""" + + def __init__(self) -> None: + super().__init__() + self.conv = torch.nn.Conv2d(4, 4, 3, padding=1) + self.pool = torch.nn.MaxPool2d(3, stride=1, padding=1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.conv(x) + self.pool(x) + + +def _count(graph_module, target) -> int: + return sum( + node.op == "call_function" and node.target == target + for node in graph_module.graph.nodes + ) + + +def _lower(module, inputs): + module.eval() + with torch.no_grad(): + exported = torch.export.export(module, inputs) + edge = to_edge( + exported, + compile_config=EdgeCompileConfig( + _check_ir_validity=False, _skip_dim_order=True + ), + ) + edge = edge.transform([ToContiguousChannelsLastPass(edge.exported_program())]) + return edge + + +def _absorb(edge): + layout_pass = AbsorbBoundaryLayoutCopies(edge.exported_program()) + return edge.transform([layout_pass]), layout_pass.contract + + +def _run(edge, contract, inputs): + args = list(inputs) + for index, dims in contract.inputs.items(): + args[index] = args[index].permute(list(dims)).contiguous() + result = edge.exported_program().module()(*args) + results = list(result) if isinstance(result, (tuple, list)) else [result] + for index, dims in contract.outputs.items(): + results[index] = results[index].permute(list(dims)) + return results[0] if len(results) == 1 else results + + +@pytest.mark.parametrize("module", [Conv(), ResidualConvPool()]) +def test_boundary_copies_are_absorbed_and_numerics_hold(module) -> None: + inputs = (torch.randn(1, 4, 8, 8),) + expected = module.eval()(*inputs) + edge = _lower(module, inputs) + assert _count(edge.exported_program().graph_module, _LAYOUT_COPY) > 0 + + edge, contract = _absorb(edge) + + assert contract.inputs and contract.outputs + assert _count(edge.exported_program().graph_module, _LAYOUT_COPY) == 0 + assert torch.allclose(_run(edge, contract, inputs), expected, atol=1e-6) + + +def test_fan_out_collapses_to_one_contract_entry() -> None: + """Both branches of a residual share the input, so one entry covers them.""" + module = ResidualConvPool() + inputs = (torch.randn(1, 4, 8, 8),) + edge = _lower(module, inputs) + copies_on_input = [ + node + for node in edge.exported_program().graph_module.graph.nodes + if node.op == "placeholder" + and node.name in edge.exported_program().graph_signature.user_inputs + for _ in node.users + ] + assert len(copies_on_input) > 1 + + _, contract = _absorb(edge) + + assert list(contract.inputs) == [0] + + +def test_mixed_users_are_left_alone() -> None: + """An input consumed both by a layout region and directly is not a boundary.""" + + class MixedUse(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv = torch.nn.Conv2d(4, 4, 3, padding=1) + + def forward(self, x): + return self.conv(x) + x + + inputs = (torch.randn(1, 4, 8, 8),) + module = MixedUse() + expected = module.eval()(*inputs) + edge = _lower(module, inputs) + + edge, contract = _absorb(edge) + + assert 0 not in contract.inputs + assert torch.allclose(_run(edge, contract, inputs), expected, atol=1e-6) + + +def test_absorbing_is_idempotent() -> None: + inputs = (torch.randn(1, 4, 8, 8),) + edge = _lower(Conv(), inputs) + edge, first = _absorb(edge) + edge, second = _absorb(edge) + + assert first + assert not second + + +def test_signature_stays_valid() -> None: + inputs = (torch.randn(1, 4, 8, 8),) + edge = _lower(Conv(), inputs) + edge, contract = _absorb(edge) + + assert contract + edge.exported_program()._validate() + + +# The layout pass is measured against these 36 models in +# test_to_contiguous_channels_last_pass.py, which pins its own per-case counts +# but stops before absorption. These are the totals across that matrix. +_MATRIX_BASELINE_PERMUTES = 138 +_MATRIX_LAYOUT_ONLY_PERMUTES = 156 +_MATRIX_ABSORBED_PERMUTES = 137 + +# Absorption does not destroy 19 permutes, it moves them across the method +# boundary: 17 become obligations on the caller. Pinning both numbers keeps the +# in-graph count honest, since it could otherwise be driven to zero by handing +# the caller unlimited work. The gap is the genuine saving, and it comes from +# fan-out — one placeholder feeding several branches needs several copies but +# only one contract entry. +_MATRIX_CONTRACT_ENTRIES = 17 + +# Cases still above baseline after absorbing. Every one of these is an internal +# copy left by a region that a non-permutable op (batch_norm, linear) cut in +# two, which is region-merging work rather than boundary work. +_MATRIX_RESIDUAL_REGRESSIONS = { + "conv2d_rank3", + "model_1_conv_maxpool_residual_linear", + "model_8_conv_batchnorm_maxpool_residual", + "model_9_dilated_conv_batchnorm_avgpool_residual", + "views", +} + +_PERMUTE_TARGETS = { + exir_ops.edge.aten.permute_copy.default, + exir_ops.edge.channels_last.permute_copy.default, +} + + +def _permutes(edge) -> int: + return sum( + node.op == "call_function" and node.target in _PERMUTE_TARGETS + for node in edge.exported_program().graph.nodes + ) + + +def test_absorption_pays_for_the_layout_pass_across_the_model_matrix() -> None: + from executorch.backends.transforms.test.test_to_contiguous_channels_last_pass import ( + cases, + ) + + baseline = layout_only = absorbed = contract_entries = 0 + regressions = set() + for name, case in cases.items(): + case.module.eval() + with torch.no_grad(): + exported = torch.export.export(case.module, case.inputs) + config = EdgeCompileConfig(_check_ir_validity=False, _skip_dim_order=True) + case_baseline = _permutes(to_edge(exported, compile_config=config)) + + edge = to_edge(exported, compile_config=config) + edge = edge.transform( + [ToContiguousChannelsLastPass(edge.exported_program())] + ) + case_layout = _permutes(edge) + + absorb = AbsorbBoundaryLayoutCopies(edge.exported_program()) + edge = edge.transform([absorb]) + case_absorbed = _permutes(edge) + + baseline += case_baseline + layout_only += case_layout + absorbed += case_absorbed + contract_entries += len(absorb.contract.inputs) + len(absorb.contract.outputs) + if case_absorbed > case_baseline: + regressions.add(name) + + assert baseline == _MATRIX_BASELINE_PERMUTES + assert layout_only == _MATRIX_LAYOUT_ONLY_PERMUTES + assert absorbed == _MATRIX_ABSORBED_PERMUTES + assert contract_entries == _MATRIX_CONTRACT_ENTRIES + assert regressions == _MATRIX_RESIDUAL_REGRESSIONS