Channels last: add the common layout pipeline

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",
+    ],
+)

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

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
 

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

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 = [