diff --git a/backends/transforms/channels_last_layout.py b/backends/transforms/channels_last_layout.py new file mode 100644 index 00000000000..eab72d9eef2 --- /dev/null +++ b/backends/transforms/channels_last_layout.py @@ -0,0 +1,32 @@ +# 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 executorch.backends.transforms.channels_last_ops # noqa: F401 + +import torch + +from executorch.exir.dialects._ops import ops as exir_ops +from torch.fx.node import Target + +ATEN_PERMUTE_COPY = exir_ops.edge.aten.permute_copy.default +LAYOUT_PERMUTE_COPY = exir_ops.edge.channels_last.permute_copy.default +PERMUTE_COPY_TARGETS: frozenset[Target] = frozenset( + (ATEN_PERMUTE_COPY, LAYOUT_PERMUTE_COPY) +) + + +def is_permute_copy(node: torch.fx.Node) -> bool: + return node.op == "call_function" and node.target in PERMUTE_COPY_TARGETS + + +def is_layout_copy(node: torch.fx.Node) -> bool: + return node.op == "call_function" and node.target == LAYOUT_PERMUTE_COPY + + +def composed_permute_target(first: torch.fx.Node, second: torch.fx.Node) -> Target: + if is_layout_copy(first) and is_layout_copy(second): + return LAYOUT_PERMUTE_COPY + return ATEN_PERMUTE_COPY diff --git a/backends/transforms/fuse_cascaded_transpose_or_permute_ops.py b/backends/transforms/fuse_cascaded_transpose_or_permute_ops.py index f350120e7eb..6b08f1ef415 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 executorch.backends.transforms.channels_last_layout import ( + composed_permute_target, + PERMUTE_COPY_TARGETS, +) from executorch.backends.transforms.permute_pass_utils import ( get_arg, get_permuted_dims, @@ -26,7 +30,7 @@ class FuseCascadedTransposeOrPermuteOps(RemoveOrReplacePassInterface): transpose_or_permute_target = { exir_ops.edge.aten.transpose_copy.int, - exir_ops.edge.aten.permute_copy.default, + *PERMUTE_COPY_TARGETS, } _VIEW_OPS = { @@ -71,10 +75,10 @@ def _fuse_direct(self, node: Node, parent_node: Node) -> bool: else: with node.graph.inserting_before(node): new_permute = node.graph.call_function( - exir_ops.edge.aten.permute_copy.default, + composed_permute_target(parent_node, node), args=(input_of_parent, dims), ) - new_permute.meta = node.meta + new_permute.meta = dict(node.meta) node.replace_all_uses_with(new_permute) return True @@ -141,7 +145,7 @@ def _fuse_across_view(self, node: Node, view_node: Node) -> bool: # noqa: C901 node_dims = list(range(len(dims))) node_dims = get_transposed_dims(node, node_dims) dims = [dims[d] for d in node_dims] - elif node.target == exir_ops.edge.aten.permute_copy.default: + elif node.target in PERMUTE_COPY_TARGETS: perm = get_arg(node, "dims") dims = [dims[d] for d in perm] else: diff --git a/backends/transforms/fuse_transpose_or_permute_op_pairs_pass.py b/backends/transforms/fuse_transpose_or_permute_op_pairs_pass.py index 7270736f5bd..83b33533995 100644 --- a/backends/transforms/fuse_transpose_or_permute_op_pairs_pass.py +++ b/backends/transforms/fuse_transpose_or_permute_op_pairs_pass.py @@ -10,6 +10,10 @@ import torch import torch.fx +from executorch.backends.transforms.channels_last_layout import ( + ATEN_PERMUTE_COPY, + LAYOUT_PERMUTE_COPY, +) from executorch.backends.transforms.permute_pass_utils import ( FuseOpPairsAcrossBranchesPass, get_permuted_dims, @@ -55,7 +59,8 @@ def can_fuse_for_chain( # this mapping helps to handle both transpose and permutations f: dict[Any, Callable] = { exir_ops.edge.aten.transpose_copy.int: get_transposed_dims, - exir_ops.edge.aten.permute_copy.default: get_permuted_dims, + ATEN_PERMUTE_COPY: get_permuted_dims, + LAYOUT_PERMUTE_COPY: get_permuted_dims, } in_dims = f[producer.target](producer, ident_dims) out_dims = f[consumer.target](consumer, in_dims) @@ -80,6 +85,7 @@ def get_fused_node( (consumer.args[0], output_shape), {}, ) + view.meta = dict(consumer.meta) return view def call(self, graph_module: torch.fx.GraphModule) -> PassResult: @@ -89,10 +95,12 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: producer_op_packets={ exir_ops.edge.aten.transpose_copy, exir_ops.edge.aten.permute_copy, + exir_ops.edge.channels_last.permute_copy, }, consumer_op_packets={ exir_ops.edge.aten.transpose_copy, exir_ops.edge.aten.permute_copy, + exir_ops.edge.channels_last.permute_copy, }, bypass_ops=self.bypass_ops, ) diff --git a/backends/transforms/permute_pass_utils.py b/backends/transforms/permute_pass_utils.py index fca8946165e..97588beb423 100644 --- a/backends/transforms/permute_pass_utils.py +++ b/backends/transforms/permute_pass_utils.py @@ -18,6 +18,7 @@ import torch import torch.fx +from executorch.backends.transforms.channels_last_layout import PERMUTE_COPY_TARGETS from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.dialects.edge._ops import EdgeOpOverload, EdgeOpOverloadPacket from executorch.exir.pass_base import ExportPass, PassResult @@ -76,7 +77,7 @@ def get_transposed_dims( def get_permuted_dims(node: torch.fx.Node, dims: List[int]) -> List[int]: """Applies the permutation as given by node onto the dimensions given in input.""" - assert node.target == exir_ops.edge.aten.permute_copy.default + assert node.target in PERMUTE_COPY_TARGETS # pyre-fixme[6]: This combined typecheck isn't supported yet. permute_dims: List[int] = list(node.args[1]) assert all(isinstance(x, int) for x in permute_dims) diff --git a/backends/transforms/postpone_permute_below_squeeze_view.py b/backends/transforms/postpone_permute_below_squeeze_view.py index e0e9a3ec198..226924cbea7 100644 --- a/backends/transforms/postpone_permute_below_squeeze_view.py +++ b/backends/transforms/postpone_permute_below_squeeze_view.py @@ -11,6 +11,7 @@ import torch import torch.fx +from executorch.backends.transforms.channels_last_layout import PERMUTE_COPY_TARGETS from executorch.backends.transforms.permute_pass_utils import ( get_shape, RemoveOrReplacePassInterface, @@ -36,7 +37,7 @@ class PostponePermuteOpBelowSqueezeOrUnsqueezeLikeView(RemoveOrReplacePassInterf @property def targets(self) -> list[EdgeOpOverload]: - return [exir_ops.edge.aten.permute_copy.default] + return list(PERMUTE_COPY_TARGETS) # If list1 and list2 are same (same values and in same order) except # list1 has one more element with value of 1. Return index of the extra 1. @@ -182,7 +183,7 @@ def _insert_nodes( permute_target, args=(new_view_node, new_permute_dims), ) - new_permute_node.meta = view_node.meta + new_permute_node.meta = dict(view_node.meta) view_node.replace_all_uses_with(new_permute_node) # view_node is user of permute_node, so must erase view_node first diff --git a/backends/transforms/remove_permutes_around_elementwise_ops.py b/backends/transforms/remove_permutes_around_elementwise_ops.py index 6e916dfe50a..34068e97ecd 100644 --- a/backends/transforms/remove_permutes_around_elementwise_ops.py +++ b/backends/transforms/remove_permutes_around_elementwise_ops.py @@ -12,6 +12,10 @@ import torch import torch.fx +from executorch.backends.transforms.channels_last_layout import ( + is_permute_copy, + PERMUTE_COPY_TARGETS, +) from executorch.backends.transforms.permute_pass_utils import get_arg, set_arg from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.pass_base import ExportPass, PassResult @@ -325,9 +329,9 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901 self._interleave_cache.clear() subgraphs_found: list[RemovePermutesAroundElementwiseOps.Subgraph] = [] processed_nodes: set[torch.fx.Node] = set() - for node in graph_module.graph.find_nodes( - op="call_function", target=exir_ops.edge.aten.permute_copy.default - ): + for node in graph_module.graph.nodes: + if not is_permute_copy(node): + continue start_permute = self.get_permutation(node) if start_permute is None: continue @@ -483,7 +487,7 @@ def visit( # noqa: C901 # Traverse downstream: for user in users_source.users: - if user.target == exir_ops.edge.aten.permute_copy.default: + if user.target in PERMUTE_COPY_TARGETS: user_perm = self.get_permutation(user) if user_perm == downstream_end: subgraph.edges_out.add((users_source, user)) @@ -528,7 +532,7 @@ def visit( # noqa: C901 # Traverse upstream: for inp in node.all_input_nodes: - if inp.target == exir_ops.edge.aten.permute_copy.default: + if inp.target in PERMUTE_COPY_TARGETS: if self.get_permutation(inp) != current_start_permute: return False subgraph.edges_in.add((inp, node)) @@ -712,7 +716,7 @@ def permute_subgraph(self, subgraph: Subgraph) -> bool: # noqa: C901 # Skip incoming permutes. for inp, out in subgraph.edges_in: - assert inp.target == exir_ops.edge.aten.permute_copy.default + assert inp.target in PERMUTE_COPY_TARGETS if len(inp.args) >= 1: out.replace_input_with(inp, cast(torch.fx.Node, inp.args[0])) else: @@ -755,7 +759,7 @@ def permute_subgraph(self, subgraph: Subgraph) -> bool: # noqa: C901 # Skip outgoing permutes. for inp, out in subgraph.edges_out: - assert out.target == exir_ops.edge.aten.permute_copy.default + assert out.target in PERMUTE_COPY_TARGETS out.replace_all_uses_with(inp) return True @@ -763,17 +767,11 @@ def permute_subgraph(self, subgraph: Subgraph) -> bool: # noqa: C901 def _subgraph_edges_are_current(self, subgraph: Subgraph) -> bool: """Return false if an earlier rewrite invalidated this candidate.""" for inp, out in subgraph.edges_in: - if ( - inp.target != exir_ops.edge.aten.permute_copy.default - or inp not in out.all_input_nodes - ): + if inp.target not in PERMUTE_COPY_TARGETS or inp not in out.all_input_nodes: return False for inp, out in subgraph.edges_out: - if ( - out.target != exir_ops.edge.aten.permute_copy.default - or out not in inp.users - ): + if out.target not in PERMUTE_COPY_TARGETS or out not in inp.users: return False for const_node, user_node in subgraph.constant_edges_in: @@ -892,7 +890,7 @@ def update_view_copy(self, node: torch.fx.Node, start_permute: list[int]) -> Non node.update_arg(1, new_shape) def get_permutation(self, permute_node: torch.fx.Node) -> list[int] | None: - assert permute_node.target == exir_ops.edge.aten.permute_copy.default + assert permute_node.target in PERMUTE_COPY_TARGETS raw_permute: list[int] if len(permute_node.args) >= 2: raw_permute = list(cast(list[int], permute_node.args[1])) diff --git a/backends/transforms/replace_nop_transpose_or_permute_with_view.py b/backends/transforms/replace_nop_transpose_or_permute_with_view.py index ccfb4ebe8b9..aeff987954b 100644 --- a/backends/transforms/replace_nop_transpose_or_permute_with_view.py +++ b/backends/transforms/replace_nop_transpose_or_permute_with_view.py @@ -10,6 +10,7 @@ import torch import torch.fx +from executorch.backends.transforms.channels_last_layout import PERMUTE_COPY_TARGETS from executorch.backends.transforms.permute_pass_utils import ( RemoveOrReplacePassInterface, ) @@ -28,7 +29,7 @@ class ReplaceNopTransposeOrPermuteWithViewPass(RemoveOrReplacePassInterface): def targets(self) -> list[EdgeOpOverload]: return [ exir_ops.edge.aten.transpose_copy.int, - exir_ops.edge.aten.permute_copy.default, + *PERMUTE_COPY_TARGETS, ] def maybe_remove_or_replace(self, node: torch.fx.Node) -> bool: @@ -61,7 +62,7 @@ def maybe_remove_or_replace(self, node: torch.fx.Node) -> bool: node.replace_all_uses_with(new_node) return True - elif node.target == exir_ops.edge.aten.permute_copy.default: + elif node.target in PERMUTE_COPY_TARGETS: old_dims = list(range(len(in_shape))) new_dims = cast(Sequence[int], node.args[1]) # If the permute does not change anything, return the input as output. diff --git a/backends/transforms/replace_ops_with_channels_last_variants.py b/backends/transforms/replace_ops_with_channels_last_variants.py index 5e8a4612163..641c3925b65 100644 --- a/backends/transforms/replace_ops_with_channels_last_variants.py +++ b/backends/transforms/replace_ops_with_channels_last_variants.py @@ -11,6 +11,7 @@ import torch +from executorch.backends.transforms.channels_last_layout import LAYOUT_PERMUTE_COPY from executorch.exir import ExportedProgram from executorch.exir.dialects._ops import ops as exir_ops @@ -166,7 +167,7 @@ def _permute_node_input( res = graph.create_node( "call_function", - target=exir_ops.edge.channels_last.permute_copy.default, + target=LAYOUT_PERMUTE_COPY, args=(node_input, _NCHW_TO_NHWC_PERM), ) res.meta = {} @@ -182,7 +183,7 @@ def _permute_node_output( ): output = graph.create_node( "call_function", - target=exir_ops.edge.channels_last.permute_copy.default, + target=LAYOUT_PERMUTE_COPY, args=(node_output, _NHWC_TO_NCHW_PERM), ) output.meta = {} diff --git a/backends/transforms/targets.bzl b/backends/transforms/targets.bzl index f72d847a8ef..4283c949e13 100644 --- a/backends/transforms/targets.bzl +++ b/backends/transforms/targets.bzl @@ -221,6 +221,21 @@ def define_common_targets(): ], ) + runtime.python_library( + name = "channels_last_layout", + srcs = [ + "channels_last_layout.py", + ], + visibility = [ + "//executorch/backends/...", + ], + deps = [ + "//caffe2:torch", + ":channels_last_ops", + "//executorch/exir/dialects:lib", + ], + ) + runtime.python_library( name = "decompose_channels_last_pass", srcs = [ @@ -375,6 +390,7 @@ def define_common_targets(): "//executorch/backends/...", ], deps = [ + ":channels_last_layout", "//caffe2:torch", "//executorch/exir:pass_base", "//executorch/exir/dialects:lib", @@ -388,6 +404,7 @@ def define_common_targets(): "//executorch/backends/...", ], deps = [ + ":channels_last_layout", "//caffe2:torch", "//executorch/exir/dialects:lib", ":permute_pass_utils", @@ -414,6 +431,7 @@ def define_common_targets(): "//executorch/backends/...", ], deps = [ + ":channels_last_layout", "//caffe2:torch", "//executorch/exir:pass_base", "//executorch/exir/dialects:lib", @@ -429,8 +447,10 @@ def define_common_targets(): "@EXECUTORCH_CLIENTS", ], deps = [ + ":channels_last_layout", ":permute_pass_utils", "//caffe2:torch", + "//executorch/exir:lib", "//executorch/exir:pass_base", "//executorch/exir/dialects:lib", ], @@ -443,6 +463,7 @@ def define_common_targets(): "//executorch/backends/...", ], deps = [ + ":channels_last_layout", "//caffe2:torch", "//executorch/exir:pass_base", "//executorch/exir/dialects:lib", @@ -457,6 +478,7 @@ def define_common_targets(): "//executorch/backends/...", ], deps = [ + ":channels_last_layout", "//caffe2:torch", "//executorch/exir/dialects:lib", ":permute_pass_utils", @@ -508,6 +530,7 @@ def define_common_targets(): ], deps = [ "//caffe2:torch", + ":channels_last_layout", ":channels_last_ops", "//executorch/exir:pass_base", "//executorch/exir:lib", diff --git a/backends/transforms/test/test_permute_optimization_passes.py b/backends/transforms/test/test_permute_optimization_passes.py index cf4583357d7..c774852cf63 100644 --- a/backends/transforms/test/test_permute_optimization_passes.py +++ b/backends/transforms/test/test_permute_optimization_passes.py @@ -11,6 +11,8 @@ import unittest from typing import cast +import executorch.backends.transforms.channels_last_ops # noqa: F401 + import torch from executorch.backends.test.graph_builder import GraphBuilder, single_op_builder from executorch.backends.transforms.fuse_cascaded_transpose_or_permute_ops import ( @@ -2024,3 +2026,43 @@ def test_no_permutes_is_noop(self) -> None: self.assertEqual( count_node(result.graph_module, exir_ops.edge.aten.permute_copy.default), 0 ) + + +class LayoutPermuteVisibilityTest(unittest.TestCase): + """The data-movement passes must see both permute dialects. + + Before the shared target set these passes matched only + ``aten.permute_copy``, so a ``channels_last.permute_copy`` pair inserted by + the layout replacement was invisible and survived untouched. + """ + + def test_inverse_layout_copy_pair_around_elementwise_is_removed(self) -> None: + builder = GraphBuilder() + x = builder.placeholder("x", torch.randn(1, 2, 3, 4)) + to_nhwc = builder.call_operator( + op=exir_ops.edge.channels_last.permute_copy.default, + args=(x, [0, 2, 3, 1]), + ) + relu = builder.call_operator( + op=exir_ops.edge.aten.hardtanh.default, args=(to_nhwc,) + ) + to_nchw = builder.call_operator( + op=exir_ops.edge.channels_last.permute_copy.default, + args=(relu, [0, 3, 1, 2]), + ) + builder.output([to_nchw]) + original = builder.get_graph_module() + gm_before = copy.deepcopy(original) + + result = cast(PassResult, RemovePermutesAroundElementwiseOps()(original)) + self.assertTrue(result.modified) + gm = result.graph_module + self.assertEqual( + count_node(gm, exir_ops.edge.channels_last.permute_copy.default), 0 + ) + validate_numerics( + gm_before, + gm, + (torch.randn(1, 2, 3, 4),), + "RemovePermutesAroundElementwiseOps", + )