Replace legacy MovePermuteAfterConcat with the ExecuTorch pass (#21593)

Differential Revision: D114631054

Pull Request resolved: #21593

Author: DrJessop

backends/cadence/aot/reorder_ops.py

@@ -738,6 +738,9 @@ def _match_inputs(
             # pyre-ignore[6]
             list[torch.fx.Node] | tuple[torch.fx.Node, ...],
         )
+        # Moving one reused permute after cat makes it process the larger output.
+        if inputs and all(inp is inputs[0] for inp in inputs):
+            return None
 
         unpermuted_inputs: List[torch.fx.Node] = []
         common_permutation: Optional[Tuple[int, ...]] = None

backends/cadence/aot/tests/test_reorder_ops_passes.py

@@ -682,6 +682,32 @@ def test_matching_single_user_permutes_moved_after_cat(self) -> None:
         placeholder_nodes = converted.graph.find_nodes(op="placeholder")
         self.assertEqual(cat_nodes[0].args[0], placeholder_nodes)
 
+    def test_repeated_permute_input_not_moved(self) -> None:
+        x_data = torch.randn(1, 16, 32)
+        builder = GraphBuilder()
+        x = builder.placeholder("x", x_data)
+        permuted_x = builder.call_operator(
+            exir_ops.edge.aten.permute_copy.default,
+            args=(x, [1, 2, 0]),
+        )
+        cat = builder.call_operator(
+            exir_ops.edge.aten.cat.default,
+            args=([permuted_x, permuted_x], 2),
+        )
+        builder.output([cat])
+
+        result = transform_and_check_numerics(
+            builder.get_graph_module(),
+            (x_data,),
+            MovePermuteAfterConcat(),
+        )
+
+        self.assertFalse(result.modified)
+        self.assertEqual(
+            get_compute_nodes_in_gm(result.graph_module),
+            [exir_ops.edge.aten.permute_copy, exir_ops.edge.aten.cat],
+        )
+
     def test_different_permutations_not_moved(self) -> None:
         builder = GraphBuilder()
         x = builder.placeholder("x", torch.randn(2, 3, 4))