Don't drop a permute that an earlier region already bypassed (#24067)

"Arm backend: Map linear rows to conv2d width" (9a81281) makes
DecomposeLinearPass emit permutes around the 1x1 conv. In a gated MLP
(`u * v`, where `v` also goes through a transpose), this lets
RemovePermutesAroundElementwiseOps plan two regions that share one
permute:

- Region A treats the v-branch transpose as an incoming permute and
bypasses it, so `mul` reads the transpose's producer directly.
- Region B was planned before A was applied. It treats the same
transpose as its outgoing permute, removes its own start permute and
changes that producer's layout.

`_subgraph_edges_are_current` only checked that the outgoing permute was
still wired to its producer. A bypassed permute stays wired even though
nothing uses it, so B was applied too. `mul` was left with mismatched
operands (`[1,144,14] * [1,14,144]`), and
PropagateViewCopyPermuteDownPass failed.

The fix records the permutes each applied region bypasses. A later
region whose outgoing (or rewritten) permute is in that set is skipped.

Test plan:
- New unit test `test_permute_bypassed_by_one_region_is_kept_by_another`
passes. With the check disabled it fails with the same broadcast
RuntimeError.
- A gated-MLP model lowered through TosaPipelineINT failed with the same
error before the fix and passes after.

Internal mirror: D126717094.

Author: Gasoonjia

backends/transforms/remove_permutes_around_elementwise_ops.py

@@ -120,6 +120,7 @@ def __init__(
             tuple[int, int, torch.fx.Node, torch.fx.Node] | None,
         ] = {}
         self._pending_visits = None
+        self._bypassed_permutes: set[torch.fx.Node] = set()
 
     _VIEW_OPS = (
         exir_ops.edge.aten.view_copy.default,
@@ -417,6 +418,7 @@ def _adapt_permute_across_view(
     def call(self, graph_module: torch.fx.GraphModule) -> PassResult:  # noqa: C901
         self._sq_unsq_cache.clear()
         self._interleave_cache.clear()
+        self._bypassed_permutes.clear()
         subgraphs_found: list[RemovePermutesAroundElementwiseOps.Subgraph] = []
         processed_nodes: set[torch.fx.Node] = set()
         for node in graph_module.graph.nodes:
@@ -829,6 +831,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
+            self._bypassed_permutes.add(inp)
             if len(inp.args) >= 1:
                 out.replace_input_with(inp, cast(torch.fx.Node, inp.args[0]))
             else:
@@ -897,7 +900,13 @@ def _subgraph_edges_are_current(self, subgraph: Subgraph) -> bool:
             (inp, out) for inp, out, _ in subgraph.edges_out_to_update
         ]
         for inp, out in outgoing:
-            if out.target not in PERMUTE_COPY_TARGETS or out not in inp.users:
+            # A bypassed permute stays wired to its producer, but an earlier
+            # region now reads that producer directly and relies on its layout.
+            if (
+                out.target not in PERMUTE_COPY_TARGETS
+                or out not in inp.users
+                or out in self._bypassed_permutes
+            ):
                 return False
 
         for const_node, user_node in subgraph.constant_edges_in:

backends/transforms/test/test_permute_optimization_passes.py

@@ -2500,6 +2500,43 @@ def test_chained_regions_absorb_into_last_permute(self) -> None:
             "chained_regions_absorb_into_last_permute",
         )
 
+    def test_permute_bypassed_by_one_region_is_kept_by_another(self) -> None:
+        """Region A bypasses pb as an incoming permute, so mul reads tanh
+        directly. Region B would remove pb as its outgoing permute and change
+        tanh's layout, so it must be skipped."""
+        builder = GraphBuilder()
+        x_data = torch.randn(2, 3, 4)
+        y_data = torch.randn(2, 4, 3)
+        x = builder.placeholder("x", x_data)
+        y = builder.placeholder("y", y_data)
+        pa = builder.call_operator(
+            op=exir_ops.edge.aten.permute_copy.default, args=(x, [0, 2, 1])
+        )
+        pc = builder.call_operator(
+            op=exir_ops.edge.aten.permute_copy.default, args=(y, [0, 2, 1])
+        )
+        tanh = builder.call_operator(op=exir_ops.edge.aten.tanh.default, args=(pc,))
+        pb = builder.call_operator(
+            op=exir_ops.edge.aten.permute_copy.default, args=(tanh, [0, 2, 1])
+        )
+        mul = builder.call_operator(op=exir_ops.edge.aten.mul.Tensor, args=(pa, pb))
+        out = builder.call_operator(
+            op=exir_ops.edge.aten.permute_copy.default, args=(mul, [0, 2, 1])
+        )
+        builder.output([out])
+        original = builder.get_graph_module()
+        gm_before = copy.deepcopy(original)
+
+        result = cast(PassResult, RemovePermutesAroundElementwiseOps()(original))
+
+        self.assertTrue(result.modified)
+        validate_numerics(
+            gm_before,
+            result.graph_module,
+            [x_data, y_data],
+            "permute_bypassed_by_one_region_is_kept_by_another",
+        )
+
 
 class LayoutPermuteVisibilityTest(unittest.TestCase):
     """The data-movement passes must see both permute dialects.