Remove static view_copy (1/4): normalize_view_copy_base_pass

Summary:
Design: https://docs.google.com/document/d/1l9x925EOrE8mHFJdRCC59nBJXyqBdnoeK-EgNQScXD0/edit#heading=h.kocb2mvchnib

When remove_static_view_copy is turned off (state today), the pass flow in to_executorch is:
1. config.to_out_var_pass
2. config.memory_planning_pass

When remove_static_view_copy is turned on, the pass flow in to_executorch becomes:
1. NormalizeViewCopyBasePass()
2. ReplaceStaticViewCopyWithMemoryViewPass() (introduces executorch.exir.memory.view)
3. config.to_out_var_pass (skips executorch.exir.memory.view)
4. config.memory_planning_pass
5. ReplaceMemoryViewWithAllocPass() (removes executorch.exir.memory.view)

The basic idea is to replace view_copy with a new operator executorch.exir.memory.view before memory planning (ReplaceStaticViewCopyWithMemoryViewPass).  These nodes share the same spec as their base so that lifetimes are updated appropriately during memory planning.

After memory planning, these nodes are converted to executorch.exir.memory.alloc nodes before emission.  They are not converted to alloc nodes before memory planning.  Before memory planning, memory.view nodes share the same spec as their base, but after memory planning, they get new specs when they're converted to memory.alloc (pointing to the same storage as base).

Differential Revision: https://internalfb.com/D54816555

Author: Scott Roy

exir/passes/TARGETS

@@ -312,3 +312,14 @@ python_library(
         "//executorch/exir/dialects:lib",
     ],
 )
+
+python_library(
+    name = "normalize_view_copy_base_pass",
+    srcs = [
+        "normalize_view_copy_base_pass.py",
+    ],
+    deps = [
+        "//caffe2:torch",
+        "//executorch/exir/dialects:lib",
+    ],
+)

exir/passes/normalize_view_copy_base_pass.py

@@ -0,0 +1,64 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-strict
+
+import logging
+
+import torch
+
+from executorch.exir.dialects._ops import ops
+from torch.fx.passes.infra.pass_base import PassBase, PassResult
+
+
+def _is_view_copy(node: torch.fx.Node) -> bool:
+    return node.op == "call_function" and node.target in (
+        torch.ops.aten.view_copy.default,
+        ops.edge.aten.view_copy.default,
+    )
+
+
+class NormalizeViewCopyBasePass(PassBase):
+    """
+    Point each view_copy to the first upstream non-view.
+
+    After this pass, the base of each view_copy is not a view_copy.
+
+    When combined with dead-code elimination, this pass removes redundant
+    view_copy nodes.
+
+    TODO: replace RemoveRedundantViewCopyPass with NormalizeViewCopyBasePass + dead code elimination.
+    """
+
+    def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
+        n_updated = 0
+        for module in graph_module.modules():
+            if not isinstance(module, torch.fx.GraphModule):
+                continue
+            for node in module.graph.nodes:
+                if _is_view_copy(node):
+                    base, size = node.args
+                    if _is_view_copy(base):
+                        # Point base to bases's base and update node's args
+                        # Base's base will not be a view_copy because we iterate
+                        # through the graph in topological order, replacing as we go.
+                        base = base.args[0]
+                        node.args = (base, size)
+                        n_updated += 1
+
+            module.recompile()
+
+        logging.debug(f"Updated the base on {n_updated} view_copy nodes.")
+        return PassResult(graph_module, n_updated > 0)
+
+    def ensures(self, graph_module: torch.fx.GraphModule) -> None:
+        for module in graph_module.modules():
+            if not isinstance(module, torch.fx.GraphModule):
+                continue
+            for node in module.graph.nodes:
+                if _is_view_copy(node):
+                    base, size = node.args
+                    assert not _is_view_copy(base)

exir/tests/TARGETS

@@ -216,6 +216,7 @@ python_unittest(
         "//executorch/exir/passes:debug_handle_generator_pass",
         "//executorch/exir/passes:insert_write_back_for_buffers_pass",
         "//executorch/exir/passes:lib",
+        "//executorch/exir/passes:normalize_view_copy_base_pass",
         "//executorch/exir/passes:remove_graph_asserts_pass",
         "//executorch/exir/passes:remove_mixed_type_operators",
         "//executorch/exir/passes:replace_edge_with_backend_pass",

exir/tests/test_passes.py

@@ -36,6 +36,9 @@
 from executorch.exir.passes.insert_write_back_for_buffers_pass import (
     insert_write_back_for_buffers_pass,
 )
+from executorch.exir.passes.normalize_view_copy_base_pass import (
+    NormalizeViewCopyBasePass,
+)
 from executorch.exir.passes.remove_graph_asserts_pass import RemoveGraphAssertsPass
 from executorch.exir.passes.remove_mixed_type_operators import RemoveMixedTypeOperators
 from executorch.exir.passes.replace_edge_with_backend_pass import EdgeToBackendOpsPass
@@ -1310,3 +1313,53 @@ def forward(self, x):
         # The 3 views on z are collapsed to 1 view
         # In total, 2 view remain
         self.assertEqual(n_views_after, 2)
+
+    def test_normalize_view_copy_base_pass(self) -> None:
+
+        class ViewChain(torch.nn.Module):
+            def forward(self, x):
+                x = torch.ops.aten.view_copy.default(x, [30, 1])
+                x = torch.ops.aten.view_copy.default(x, [5, 6])
+                x = torch.ops.aten.view_copy.default(x, [2, 15])
+                x = torch.ops.aten.view_copy.default(x, [3, -1])
+                return x
+
+        def is_view_copy(node: torch.fx.Node) -> bool:
+            return (
+                node.op == "call_function"
+                and node.target == torch.ops.aten.view_copy.default
+            )
+
+        gm = export(ViewChain(), (torch.ones(30),)).graph_module
+
+        # Check before transformation
+        n_view_copy_before = 0
+        n_view_copy_bases_before = 0
+        for node in gm.graph.nodes:
+            if is_view_copy(node):
+                n_view_copy_before += 1
+                base = node.args[0]
+                if is_view_copy(base):
+                    n_view_copy_bases_before += 1
+
+        self.assertEqual(n_view_copy_before, 4)
+        self.assertEqual(n_view_copy_bases_before, 3)
+
+        # Do transformation
+        p = NormalizeViewCopyBasePass()
+        gm_res = p(gm)
+        assert gm_res is not None
+        gm = gm_res.graph_module
+
+        # Check after transformation
+        n_view_copy_after = 0
+        n_view_copy_bases_after = 0
+        for node in gm.graph.nodes:
+            if is_view_copy(node):
+                n_view_copy_after += 1
+                base = node.args[0]
+                if is_view_copy(base):
+                    n_view_copy_bases_after += 1
+
+        self.assertEqual(n_view_copy_after, 4)
+        self.assertEqual(n_view_copy_bases_after, 0)