Transforms: make layout-copy rewrites backend-configurable

backends/arm/test/passes/test_canonicalize_view_copy_permute_pass.py

@@ -571,3 +571,61 @@ def test_canonicalize_unsupported_end_view_does_not_block_prefix() -> None:
     )
     assert _count_node(result.graph_module, exir_ops.edge.aten.view_copy.default) == 1
     _validate_numerics(gm_before, result.graph_module, (x_data,))
+
+
+def test_canonicalize_mixed_permute_dialects_is_idempotent() -> None:
+    builder = GraphBuilder()
+    x_data = torch.randn(1, 2, 3, 4)
+    x = builder.placeholder("x", x_data)
+    layout_permute = builder.call_operator(
+        op=exir_ops.edge.channels_last.permute_copy.default,
+        args=(x, [0, 2, 3, 1]),
+    )
+    view = builder.call_operator(
+        op=exir_ops.edge.aten.view_copy.default,
+        args=(layout_permute, [1, 12, 2]),
+    )
+    permute = builder.call_operator(
+        op=exir_ops.edge.aten.permute_copy.default,
+        args=(view, [0, 2, 1]),
+    )
+    builder.output([permute])
+    graph_module = builder.get_graph_module()
+
+    pass_instance = CanonicalizeViewCopyPermutePass(
+        permute_targets={
+            exir_ops.edge.channels_last.permute_copy.default,
+            exir_ops.edge.aten.permute_copy.default,
+        }
+    )
+    first = cast(PassResult, pass_instance.call(graph_module))
+    second = cast(PassResult, pass_instance.call(first.graph_module))
+
+    assert not first.modified
+    assert not second.modified
+    assert _compute_nodes(second.graph_module) == [
+        exir_ops.edge.channels_last.permute_copy.default,
+        exir_ops.edge.aten.view_copy.default,
+        exir_ops.edge.aten.permute_copy.default,
+    ]
+
+
+def test_canonicalize_refreshes_backend_permute_metadata() -> None:
+    builder = GraphBuilder()
+    x = builder.placeholder("x", torch.randn(1, 2, 3, 4))
+    permute = builder.call_operator(
+        op=exir_ops.edge.channels_last.permute_copy.default,
+        args=(x, [0, 2, 3, 1]),
+    )
+
+    pass_instance = CanonicalizeViewCopyPermutePass(
+        permute_targets={exir_ops.edge.channels_last.permute_copy.default}
+    )
+    pass_instance._set_node_op(
+        permute.node,
+        exir_ops.edge.channels_last.permute_copy.default,
+        x.node,
+        [0, 3, 1, 2],
+    )
+
+    assert permute.node.meta["val"].shape == torch.Size([1, 4, 2, 3])

backends/transforms/canonicalize_view_copy_permute_pass.py

@@ -43,8 +43,7 @@ class CanonicalizeViewCopyPermutePass(ExportPass):
     def __init__(self, permute_targets: Iterable[Any] | None = None) -> None:
         super().__init__()
         # Which targets count as a permute. A backend carrying its own layout
-        # dialect passes them here; the pass still emits _PERMUTE_TARGET when it
-        # has to create one.
+        # dialect passes them here. Mixed-dialect chains are left unchanged.
         self._permute_targets = frozenset(permute_targets or (self._PERMUTE_TARGET,))
         self._targets = {self._VIEW_TARGET} | self._permute_targets
 
@@ -55,6 +54,12 @@ def call(self, graph_module: GraphModule) -> PassResult:
         modified = False
 
         for chain in self._collect_chains(graph_module):
+            permute_targets = {
+                node.target for node in chain if node.target in self._permute_targets
+            }
+            if len(permute_targets) > 1:
+                continue
+
             updated_chain = chain
 
             while True:
@@ -169,7 +174,7 @@ def _fuse_sequential_ops(
                             _normalize_dim(dim, len(self._shape(input_node)))
                             for dim in dims
                         ]
-                        self._set_node_op(node, self._PERMUTE_TARGET, input_node, dims)
+                        self._set_node_op(node, node.target, input_node, dims)
                         changed = True
                         any_changed = True
 
@@ -214,13 +219,17 @@ def _fuse_sequential_ops(
                         any_changed = True
                         continue
 
-                    if self._is_permute(node) and self._is_permute(next_node):
+                    if (
+                        self._is_permute(node)
+                        and self._is_permute(next_node)
+                        and node.target == next_node.target
+                    ):
                         # Fuse consecutive permutes
                         dims = self._permute_dims(node)
                         next_dims = self._permute_dims(next_node)
                         self._set_node_op(
                             node,
-                            self._PERMUTE_TARGET,
+                            node.target,
                             input_node,
                             [dims[dim] for dim in next_dims],
                         )
@@ -372,7 +381,7 @@ def _set_node_op(
     ) -> None:
         node.target = target
         node.args = (input_node, list(arg))
-        refresh_permute_view_meta(node)
+        refresh_permute_view_meta(node, self._permute_targets)
 
     def _permute_dims(self, node: Node) -> list[int]:
         assert self._is_permute(node), "Expected permute node"

backends/transforms/fuse_identical_input_transforms_pass.py

@@ -259,7 +259,7 @@ def _sink_identical_input_transforms(self, node: Node) -> bool:
                 kwargs=dict(transform.kwargs),
             )
         new_node.meta = self._new_transform_meta(node, transform)
-        refresh_permute_view_meta(new_node)
+        refresh_permute_view_meta(new_node, self._permute_targets)
 
         for user in list(node.users):
             if user is not new_node:

backends/transforms/permute_view_meta.py

@@ -6,8 +6,8 @@
 
 # pyre-unsafe
 
-from collections.abc import Sequence
-from typing import cast
+from collections.abc import Iterable, Sequence
+from typing import Any, cast
 
 import torch
 import torch.fx
@@ -19,16 +19,26 @@
 from executorch.exir.dialects._ops import ops as exir_ops
 
 
-def refresh_permute_view_meta(node: torch.fx.Node) -> None:
+def refresh_permute_view_meta(
+    node: torch.fx.Node, permute_targets: Iterable[Any] | None = None
+) -> None:
     """Compute new meta-vals, specifically preserving SymInts for view/permute
     nodes.
     """
     input_node = node.all_input_nodes[0]
     input_val = input_node.meta.get("val")
-    if input_val is None or node.target not in {
-        exir_ops.edge.aten.view_copy.default,
-        exir_ops.edge.aten.permute_copy.default,
-    }:
+    if input_val is None:
+        return
+
+    permute_targets = frozenset(
+        (exir_ops.edge.aten.permute_copy.default,)
+        if permute_targets is None
+        else permute_targets
+    )
+    if (
+        node.target != exir_ops.edge.aten.view_copy.default
+        and node.target not in permute_targets
+    ):
         return
 
     if not isinstance(input_val, torch.Tensor):
@@ -45,7 +55,7 @@ def refresh_permute_view_meta(node: torch.fx.Node) -> None:
                     )
                 )
             )
-        case exir_ops.edge.aten.permute_copy.default:
+        case target if target in permute_targets:
             dims = _normalize_dims(
                 cast(Sequence[int], node.args[1]), len(input_val.shape)
             )

backends/transforms/propagate_view_copy_permute_pass.py

@@ -40,7 +40,9 @@ class PropagateViewCopyPermutePass(ExportPass, ABC):
     """Abstract implementation of a permute/view_copy propagation pass.
 
     To be used for upwards/downwards propagation by implementing the abstract
-    methods for the direction of propagation.
+    methods for the direction of propagation. Backends may supply a closed set
+    of equivalent permute targets; synthesized permutes retain their source
+    target.
 
     """
 
@@ -77,8 +79,7 @@ def __init__(
         self.exported_program = exported_program
         self.compile_spec = compile_spec
         # Which targets count as a permute. A backend carrying its own layout
-        # dialect passes them here; the pass still emits self._PERMUTE_TARGET
-        # when it has to create one.
+        # dialect passes them here.
         self._permute_targets = frozenset(permute_targets or (self._PERMUTE_TARGET,))
         self._targets = {
             self._VIEW_TARGET,
@@ -764,7 +765,7 @@ def _maybe_split_upwards_cat_fanout(
             output_shape = [input_val.shape[dim] for dim in permute_args[0]]
             with next_node.graph.inserting_before(next_node):
                 permute = next_node.graph.call_function(
-                    self._PERMUTE_TARGET,
+                    cast(Any, node.target),
                     args=(input_node, permute_args[0]),
                 )
             permute.meta = dict(input_node.meta)

backends/transforms/remove_permutes_around_elementwise_ops.py

@@ -7,15 +7,13 @@
 
 # pyre-unsafe
 
+from collections.abc import Callable
 from dataclasses import dataclass, field
-from typing import cast
+from typing import Any, cast
 
 import torch
 import torch.fx
-from executorch.backends.transforms.channels_last_layout import (
-    is_permute_copy,
-    PERMUTE_COPY_TARGETS,
-)
+from executorch.backends.transforms.channels_last_layout import 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
@@ -29,12 +27,21 @@ class RemovePermutesAroundElementwiseOps(ExportPass):
     based on the permute's parameter such as mean, cat, and slice.
     The repeat_interleave idiom (unsqueeze -> expand_copy -> merging view_copy) is
     recognised as a single rank-preserving unit; see _interleave_triple.
+
+    ``extra_permutable_ops`` must be layout-equivariant without argument
+    remapping. ``can_propagate`` lets a backend reject nodes that the shared
+    pass would otherwise treat as layout-equivariant.
+
+    ``permute_targets`` is the closed family of layout-copy operators the
+    backend considers equivalent. A rewritten region must use one member of
+    that family consistently, and any synthesized copies retain that member.
     """
 
     @dataclass()
     class Subgraph:
         start_permute: list[int]
         end_permute: list[int]
+        permute_target: Any
         # Nodes in the subgraph, does not include permutes.
         nodes: set[torch.fx.Node] = field(default_factory=set)
         # Incoming edges to the subgraph from permute nodes.
@@ -62,8 +69,16 @@ class Subgraph:
             torch.fx.Node, tuple[int, int, torch.fx.Node, torch.fx.Node]
         ] = field(default_factory=dict)
 
-    def __init__(self, extra_permutable_ops: set | None = None) -> None:
+    def __init__(
+        self,
+        extra_permutable_ops: set | None = None,
+        *,
+        can_propagate: Callable[[torch.fx.Node], bool] | None = None,
+        permute_targets: set | frozenset | None = None,
+    ) -> None:
         super().__init__()
+        self.can_propagate = can_propagate
+        self._permute_targets = frozenset(permute_targets or PERMUTE_COPY_TARGETS)
         self._permutable_ops = {
             exir_ops.edge.aten.add.Tensor,
             exir_ops.edge.aten.mul.Tensor,
@@ -198,7 +213,7 @@ def _sink_users_are_layout_invariant(self, sink: torch.fx.Node) -> bool:
 
             if node.op == "output":
                 continue
-            if node.target == exir_ops.edge.aten.permute_copy.default:
+            if node.target in self._permute_targets:
                 # This explicit transform re-establishes the downstream layout,
                 # so consumers beyond it do not depend on the sink's layout.
                 continue
@@ -370,7 +385,7 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:  # noqa: C901
         subgraphs_found: list[RemovePermutesAroundElementwiseOps.Subgraph] = []
         processed_nodes: set[torch.fx.Node] = set()
         for node in graph_module.graph.nodes:
-            if not is_permute_copy(node):
+            if node.target not in self._permute_targets:
                 continue
             start_permute = self.get_permutation(node)
             if start_permute is None:
@@ -384,7 +399,7 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:  # noqa: C901
                     and self._interleave_triple(user) is None
                 ):
                     continue
-                subgraph = self.Subgraph(start_permute, end_permute)
+                subgraph = self.Subgraph(start_permute, end_permute, node.target)
                 if self.visit(user, subgraph, processed_nodes):
                     subgraphs_found.append(subgraph)
                     for n in subgraph.nodes:
@@ -482,7 +497,9 @@ def visit(  # noqa: C901
 
         # Traverse downstream:
         for user in users_source.users:
-            if user.target in PERMUTE_COPY_TARGETS:
+            if user.target in self._permute_targets:
+                if user.target != subgraph.permute_target:
+                    return False
                 user_perm = self.get_permutation(user)
                 if user_perm == downstream_end:
                     subgraph.edges_out.add((users_source, user))
@@ -500,6 +517,8 @@ def visit(  # noqa: C901
                     )
             elif user.op == "output":
                 return False
+            elif self.can_propagate is not None and not self.can_propagate(user):
+                return False
             elif self._is_permutation_sink_view(user):
                 # The tensor's element order is invariant at this reshape, but
                 # its output shape can still carry broadcast-axis meaning.
@@ -513,8 +532,11 @@ def visit(  # noqa: C901
 
         # Traverse upstream:
         for inp in node.all_input_nodes:
-            if inp.target in PERMUTE_COPY_TARGETS:
-                if self.get_permutation(inp) != current_start_permute:
+            if inp.target in self._permute_targets:
+                if (
+                    inp.target != subgraph.permute_target
+                    or self.get_permutation(inp) != current_start_permute
+                ):
                     return False
                 subgraph.edges_in.add((inp, node))
             elif (inp_val := inp.meta.get("val")) is not None and inp_val.numel() == 1:
@@ -599,6 +621,8 @@ def _is_pointwise(target) -> bool:
         return False
 
     def is_node_permutable(self, node: torch.fx.Node) -> bool:
+        if self.can_propagate is not None and not self.can_propagate(node):
+            return False
         if node.target in self._PAD_OPS and not self._is_constant_pad(node):
             return False
         if node.target in self._permutable_ops:
@@ -713,7 +737,7 @@ def permute_subgraph(self, subgraph: Subgraph) -> bool:  # noqa: C901
 
         # Skip incoming permutes.
         for inp, out in subgraph.edges_in:
-            assert inp.target in PERMUTE_COPY_TARGETS
+            assert inp.target in self._permute_targets
             if len(inp.args) >= 1:
                 out.replace_input_with(inp, cast(torch.fx.Node, inp.args[0]))
             else:
@@ -733,7 +757,7 @@ def permute_subgraph(self, subgraph: Subgraph) -> bool:  # noqa: C901
                 if const_rank is not None and const_rank == permute_rank:
                     new_node = graph.create_node(
                         "call_function",
-                        exir_ops.edge.aten.permute_copy.default,
+                        subgraph.permute_target,
                         args=(const_node, node_end_perm),
                     )
                 elif (
@@ -756,7 +780,7 @@ def permute_subgraph(self, subgraph: Subgraph) -> bool:  # noqa: C901
 
         # Skip outgoing permutes.
         for inp, out in subgraph.edges_out:
-            assert out.target in PERMUTE_COPY_TARGETS
+            assert out.target in self._permute_targets
             out.replace_all_uses_with(inp)
 
         # Update outgoing permutes that can't be eliminated.
@@ -769,7 +793,10 @@ 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 not in PERMUTE_COPY_TARGETS or inp not in out.all_input_nodes:
+            if (
+                inp.target not in self._permute_targets
+                or inp not in out.all_input_nodes
+            ):
                 return False
 
             # edges_out_to_update can rewrite a permute in place, leaving it wired.
@@ -779,7 +806,7 @@ def _subgraph_edges_are_current(self, subgraph: Subgraph) -> bool:
                 return False
 
         for inp, out in subgraph.edges_out:
-            if out.target not in PERMUTE_COPY_TARGETS or out not in inp.users:
+            if out.target not in self._permute_targets or out not in inp.users:
                 return False
 
         for inp, out, _ in subgraph.edges_out_to_update:
@@ -902,7 +929,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 in PERMUTE_COPY_TARGETS
+        assert permute_node.target in self._permute_targets
         raw_permute: list[int]
         if len(permute_node.args) >= 2:
             raw_permute = list(cast(list[int], permute_node.args[1]))