Channels last: teach the data-movement passes both permute dialects

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

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:

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,
         )

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)

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

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,25 +759,19 @@ 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
 
     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]))

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.

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 = {}

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",