Fold non-matching end permutes instead of discarding the region (pytorch#22293)

Differential Revision: D117939968

Pull Request resolved: pytorch#22293

Author: abeakkas

backends/cadence/aot/tests/test_remove_ops_passes.py

@@ -477,11 +477,21 @@ def test_keep_permutes_around_elemwise_ops_add(self) -> None:
         )
         builder.output([permute])
         original = builder.get_graph_module()
+        gm_before = copy.deepcopy(original)
         p = RemovePermutesAroundElementwiseOps()
         graph_after_passes = cast(PassResult, p(original)).graph_module
-        # Ensure no permutes were removed, since the dimensions don't fit the expected pattern
+        # The end permute is not the inverse of the start one, so it cannot be
+        # dropped. It absorbs the start permute instead, leaving one behind.
         self.assertEqual(
-            count_node(graph_after_passes, exir_ops.edge.aten.permute_copy.default), 2
+            count_node(graph_after_passes, exir_ops.edge.aten.permute_copy.default), 1
+        )
+
+        sample_inputs = [torch.randn(1, 8, 4, 4, dtype=torch.float32)]
+        validate(
+            gm_before,
+            graph_after_passes,
+            sample_inputs,
+            "RemovePermutesAroundElementwiseOps",
         )
 
     def test_remove_permutes_around_elemwise_ops_add_mean(self) -> None:

backends/transforms/remove_permutes_around_elementwise_ops.py

@@ -41,6 +41,12 @@ class Subgraph:
         edges_in: set[tuple[torch.fx.Node, torch.fx.Node]] = field(default_factory=set)
         # Outgoing edges of the subgraph to permute nodes.
         edges_out: set[tuple[torch.fx.Node, torch.fx.Node]] = field(default_factory=set)
+        # Outgoing edges to permutes that do not match end_permute. Those are
+        # kept and their permutation rewritten to absorb the removed start
+        # permute, as (producer, permute node, new permutation).
+        edges_out_to_update: set[
+            tuple[torch.fx.Node, torch.fx.Node, tuple[int, ...]]
+        ] = field(default_factory=set)
         # Incoming edges from constant nodes that need a compensating permute.
         constant_edges_in: set[tuple[torch.fx.Node, torch.fx.Node]] = field(
             default_factory=set
@@ -372,7 +378,6 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:  # noqa: C901
             # Expected end permutation for the subgraph.
             end_permute = [start_permute.index(i) for i in range(len(start_permute))]
 
-            # Try direct users first (same-rank matching)
             for user in node.users:
                 if (
                     not self.is_node_permutable(user)
@@ -385,50 +390,6 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:  # noqa: C901
                     for n in subgraph.nodes:
                         processed_nodes.add(n)
 
-            # Also try: permute → view(squeeze/unsqueeze) → chain → ...
-            # If the permute's sole user is a squeeze/unsqueeze view,
-            # adapt the permutation across the view and search for a
-            # matching end permute at the new rank.
-            users = list(node.users.keys())
-            if (
-                len(users) == 1
-                and self._is_squeeze_unsqueeze_view(users[0])
-                and node not in processed_nodes
-            ):
-                view_node = users[0]
-                adapted_start = self._adapt_permute_across_view(
-                    start_permute, view_node
-                )
-                if adapted_start is not None:
-                    adapted_end = [
-                        adapted_start.index(i) for i in range(len(adapted_start))
-                    ]
-                    for view_user in view_node.users:
-                        if (
-                            not self.is_node_permutable(view_user)
-                            and self._interleave_triple(view_user) is None
-                        ):
-                            continue
-                        subgraph = self.Subgraph(adapted_start, adapted_end)
-                        # Include the view in the subgraph
-                        subgraph.nodes.add(view_node)
-                        subgraph.node_end_permute[view_node] = adapted_end
-                        # Use the ORIGINAL start_permute for the view node
-                        # so update_view_copy can remap its shape correctly
-                        subgraph.node_start_permute[view_node] = start_permute
-                        # The start permute feeds into the view
-                        subgraph.edges_in.add((node, view_node))
-                        if self.visit(
-                            view_user,
-                            subgraph,
-                            processed_nodes,
-                            adapted_end,
-                            adapted_start,
-                        ):
-                            subgraphs_found.append(subgraph)
-                            for n in subgraph.nodes:
-                                processed_nodes.add(n)
-
         modified = False
         for subgraph in subgraphs_found:
             if self.permute_subgraph(subgraph):
@@ -526,29 +487,17 @@ def visit(  # noqa: C901
                 if user_perm == downstream_end:
                     subgraph.edges_out.add((users_source, user))
                 else:
-                    # Check if permute → view(squeeze/unsqueeze) forms an
-                    # end boundary at a different rank.
-                    user_users = list(user.users.keys())
-                    if len(user_users) == 1 and self._is_squeeze_unsqueeze_view(
-                        user_users[0]
-                    ):
-                        view_after: torch.fx.Node = user_users[0]
-                        # Adapt the start permute across the view and derive
-                        # the expected end permute as its inverse.
-                        adapted_start_after = self._adapt_permute_across_view(
-                            downstream_start, view_after
+                    # Non-matching permute: keep it and fold the start permute into it
+                    # rather than discarding the region.
+                    if user_perm is None or len(user_perm) != len(downstream_start):
+                        return False
+                    subgraph.edges_out_to_update.add(
+                        (
+                            users_source,
+                            user,
+                            tuple(downstream_start[d] for d in user_perm),
                         )
-                        if adapted_start_after is not None:
-                            adapted = [
-                                adapted_start_after.index(i)
-                                for i in range(len(adapted_start_after))
-                            ]
-                            if user_perm == adapted:
-                                # Include both the permute and the view as end edges
-                                subgraph.edges_out.add((users_source, user))
-                                # Mark the view for inclusion so it gets preserved
-                                continue
-                    return False
+                    )
             elif user.op == "output":
                 return False
             elif self._is_permutation_sink_view(user):
@@ -682,11 +631,27 @@ def _is_constant_pad(self, node: torch.fx.Node) -> bool:
 
         return True
 
+    def _removes_a_permute(self, subgraph: Subgraph) -> bool:
+        """Whether rewriting this region reduces the number of permutes."""
+        if subgraph.edges_out:
+            return True
+        rewired = set(subgraph.edges_in)
+        for permute, _ in subgraph.edges_in:
+            if all((permute, user) in rewired for user in permute.users):
+                return True
+        return False
+
     def permute_subgraph(self, subgraph: Subgraph) -> bool:  # noqa: C901
         # Ensure that the subgraph's edges have not been modified by an earlier rewrite before applying changes.
         if not self._subgraph_edges_are_current(subgraph):
             return False
 
+        # Folding an end permute only pays for itself if some permute goes away.
+        # Otherwise the region is rewritten for nothing, and the composed
+        # permutation is a worse fusion candidate for the passes downstream.
+        if subgraph.edges_out_to_update and not self._removes_a_permute(subgraph):
+            return False
+
         # Nodes belonging to a repeat_interleave triple are rewritten as a unit
         # below, so they must skip the per-node dim handling and the view rank
         # check (the triple's interior ranks intentionally differ from the
@@ -794,6 +759,11 @@ def permute_subgraph(self, subgraph: Subgraph) -> bool:  # noqa: C901
             assert out.target in PERMUTE_COPY_TARGETS
             out.replace_all_uses_with(inp)
 
+        # Update outgoing permutes that can't be eliminated.
+        for _, out, new_permutation in subgraph.edges_out_to_update:
+            assert out.target in PERMUTE_COPY_TARGETS
+            set_arg(out, "dims", list(new_permutation))
+
         return True
 
     def _subgraph_edges_are_current(self, subgraph: Subgraph) -> bool:
@@ -802,10 +772,20 @@ def _subgraph_edges_are_current(self, subgraph: Subgraph) -> bool:
             if inp.target not in PERMUTE_COPY_TARGETS or inp not in out.all_input_nodes:
                 return False
 
+            # edges_out_to_update can rewrite a permute in place, leaving it wired.
+            if self.get_permutation(inp) != subgraph.node_start_permute.get(
+                out, subgraph.start_permute
+            ):
+                return False
+
         for inp, out in subgraph.edges_out:
             if out.target not in PERMUTE_COPY_TARGETS or out not in inp.users:
                 return False
 
+        for inp, out, _ in subgraph.edges_out_to_update:
+            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:
             if const_node not in user_node.all_input_nodes:
                 return False