Optimize FuseViewCopyTransform pass (#20591)

This is an optimizing and refactoring patch targeting the
FuseViewCopyTransform pass.

Prior to this patch, the pass worked in two full iterations and
recompilations of the graph: The first iteration updated view nodes,
recompiled the graph and the second iteration removed any redundant view
nodes.

Optimize this pass to instead do the same work in just one graph
iteration and compilation. The pass now works by finding fusible chains,
removing redundant view nodes, and form updated node edges within a
single iteration.

The table below shows the improvement when lowering models on a computer
using the TOSA FP profile.

| Model        | Before | After | Fuse delta      | Total delta |
|--------------|-------:|------:|----------------:|------------:|
| mobilenet_v2 |  3.21s | 1.61s | -1.61s / -50.0% | -2.7%       |
| deit_tiny    | 10.51s | 5.27s | -5.24s / -49.9% | -1.6%       |
| swin2sr      |  2.55s | 1.26s | -1.29s / -50.7% | -1.2%       |

cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils
@Sebastian-Larsson @robell @rascani

Signed-off-by: Martin Lindström <Martin.Lindstroem@arm.com>
Authored-by: Martin Lindström <Martin.Lindstroem@arm.com>

Author: Erik-Lundell

backends/transforms/fuse_view_copy.py

@@ -11,15 +11,35 @@
 
 import torch
 from executorch.exir.dialects._ops import ops as exir_ops
+from executorch.exir.dialects.edge._ops import EdgeOpOverload
 from executorch.exir.pass_base import ExportPass, PassResult
 
 
 class FuseViewCopyTransform(ExportPass):
+    """Fuse redundant ``view_copy`` nodes through supported unary elementwise
+    chains.
+
+    Example:
+        let view(sN) be a view node with output shape sN
+
+        Input:
+        view(s0) -> view(s1) -> abs -> view(s2) -> sqrt -> view(s3) -> user0
+                                                               `-----> user1
+        Output:
+        view(s3) -> abs -> sqrt -> user0
+                             `---> user1
+
+        Note that all view nodes except the first one have been elimininated.
+        That first node now outputs the final view shape s3 instead. user0 and
+        user1 have been reconnected to the last retained node sqrt because the
+        former connection has been broken with the removal of the view nodes.
+    """
+
     _passes_required_after: Set[Type[ExportPass]] = set()
 
-    VIEW_OP = exir_ops.edge.aten.view_copy.default
+    VIEW_OP: EdgeOpOverload = exir_ops.edge.aten.view_copy.default
 
-    UNARY_ELEMENTWISE_OPS = [
+    UNARY_ELEMENTWISE_OPS: list[EdgeOpOverload] = [
         exir_ops.edge.aten.alias_copy.default,
         exir_ops.edge.aten.clone.default,
         exir_ops.edge.dim_order_ops._clone_dim_order.default,
@@ -46,71 +66,72 @@ class FuseViewCopyTransform(ExportPass):
         exir_ops.edge.aten.log.default,
     ]
 
-    def merge_view_copy_chains(
-        self, graph: torch.fx.Graph
-    ) -> tuple[torch.fx.Graph, bool]:
-        """
-        Find chains of view_copy nodes and unary elementwise ops and set all
-        view_copy nodes to have the final shape. The views will then be removed
-        by the remove_noop_view_copy call.
+    def _find_view_copy_chain(
+        self, start_node: torch.fx.Node, ops: list[EdgeOpOverload]
+    ) -> tuple[torch.fx.Node, list[torch.fx.Node], list[torch.fx.Node]]:
+        """Collect a fusible chain starting at ``start_node``.
 
-        Only merges view_copy nodes that are not used by any other nodes.
+        Returns:
+            list[torch.fx.Node]: View nodes following ``node`` that can be removed.
         """
-        view_op = self.VIEW_OP
-        modified = False
-        ops = self.UNARY_ELEMENTWISE_OPS + [view_op]
-        for node in graph.nodes:
-            if node.op == "call_function" and node.target == view_op:
-                # Find a chain of unary elementwise ops and save all view_copy nodes
-                end_node = node
-                view_ops = [node]
-                while (
-                    end_node.op == "call_function"
-                    and end_node.target in ops
-                    and len(end_node.users) == 1
-                    and list(end_node.users)[0].target in ops
-                ):
-                    end_node = list(end_node.users)[0]
-                    if end_node.target == view_op:
-                        view_ops.append(end_node)
-
-                # Set all view_copy nodes to have the final shape
-                if len(view_ops) > 1:
-                    final_shape = view_ops[-1].args[1]
-                    for node in view_ops:
-                        new_args = (node.args[0], final_shape)
-                        node.args = new_args
-                    modified = True
-        if modified:
-            graph.eliminate_dead_code()
-        return graph, modified
 
-    def remove_noop_view_copy(
+        end_node = start_node
+        view_nodes: list[torch.fx.Node] = []
+        while (
+            end_node.op == "call_function"
+            and end_node.target in ops
+            and len(end_node.users) == 1
+            and (end_node_user := next(iter(end_node.users))).target in ops
+        ):
+            if end_node_user.target == self.VIEW_OP:
+                view_nodes.append(end_node_user)
+
+            end_node = end_node_user
+
+        return view_nodes
+
+    def _merge_view_copy_chains(
         self, graph: torch.fx.Graph
     ) -> tuple[torch.fx.Graph, bool]:
+        """Merge redundant view nodes in linear chains.
+
+        For each view node, search forward through a single-user chain of view
+        and supported unary elementwise ops. If the chain contains later view
+        nodes, update the first view to the final view shape and bypass each
+        later view by replacing its uses with its input. The view nodes are
+        then eliminated.
+
+        Args:
+            graph (torch.fx.Graph): Graph to rewrite.
+
+        Returns:
+            tuple[torch.fx.Graph, bool]: The rewritten graph and whether it was
+                modified.
         """
-        Remove view_copy nodes that are no-ops.
-        """
-        view_op = self.VIEW_OP
         modified = False
-        for node in graph.nodes:
-            if node.op == "call_function" and node.target == view_op:
-                input_shape = list(node.args[0].meta["val"].shape)
-                target_shape = list(node.meta["val"].shape)
-                if input_shape == target_shape:
-                    node.replace_all_uses_with(node.args[0])
-                    modified = True
+        ops: list[EdgeOpOverload] = self.UNARY_ELEMENTWISE_OPS + [self.VIEW_OP]
+        for node in graph.find_nodes(op="call_function", target=self.VIEW_OP):
+            view_nodes_to_remove = self._find_view_copy_chain(node, ops)
+
+            if len(view_nodes_to_remove) > 0:
+                modified = True
+
+                # Set the first view node in the chain to have the final shape
+                final_shape = view_nodes_to_remove[-1].args[1]
+                new_args = (node.args[0], final_shape)
+                node.args = new_args
+
+                # Redirect output edges from removed view nodes to bypass them
+                for view_node in view_nodes_to_remove:
+                    view_node.replace_all_uses_with(view_node.args[0])
+
         if modified:
             graph.eliminate_dead_code()
+
         return graph, modified
 
     def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
-        graph_module.graph, modified = self.merge_view_copy_chains(graph_module.graph)
-        if modified:
-            graph_module.recompile()
-            graph_module = super().call(graph_module).graph_module
-
-        graph_module.graph, modified = self.remove_noop_view_copy(graph_module.graph)
+        graph_module.graph, modified = self._merge_view_copy_chains(graph_module.graph)
         if modified:
             graph_module.recompile()
             graph_module = super().call(graph_module).graph_module