Replace view_copy with view (2/3) (#2462)

metascroy · facebook-github-bot · commit 1d6105d860f6 · 2024-03-15T12:50:50.000-07:00
Summary: Design: https://docs.google.com/document/d/1l9x925EOrE8mHFJdRCC59nBJXyqBdnoeK-EgNQScXD0/edit#heading=h.kocb2mvchnib This stack replaces view_copy nodes with memory.view nodes. In the first diff (D54816555), I write a pass to normalize view_copy nodes by making their base point to the upstream non-view node. This means if we have something like op -> view_copy1 -> view_copy2, then after normalization, both view copies will point to op in their base (assuming op is not a view node). Note that this pass combined with dead-code elimination removes redundant view copies. This is because a redundant view copy will have no users have this pass. In the second diff (D54827305), I write a pass to convert view_copy nodes to memory.view nodes. A memory.view is similar to torch.ops.aten.view.default, but it is its own function so that we can handle it specially during memory planning and emission. A memory.view node has a special TensorSpec of type _MemoryViewSpec. This spec is immutable and dynamically looks up non-size related fields from its base's TensorSpec. Because it is immutable, fields on a _MemoryViewSpec cannot be set, but if a field is updated on the base spec, this update is reflected in the memory.view node's _MemoryViewSpec. Not all view_copy nodes are converted to memory.view nodes. Only static nodes that are memory planned are converted. Not all static nodes are memory planned in ExecuTorch. For example, there is an option to turn off memory planning for input nodes, and outputs from some higher order ops like cond are not memory planned. Which nodes are memory planned is not easily available, and I did not try to cover all cases of nodes that can be converted. We can expand this list over time. In the third diff (D54827438), I implement the actual view_copy elimination. In the ExecutorchBackendConfig, there is a new option remove_static_view_copy. If remove_static_view_copy = True, the memory planning passes are [NormalizeViewCopyBasePass(), ReplaceViewCopyWithMemoryViewPass(), config.to_out_var_pass, config.memory_planning_pass]; if remove_static_view_copy = False, the memory planning passes are [config.to_out_var_pass, config.memory_planning_pass] (state today). Let's look at the flow when remove_static_view_copy = True: NormalizeViewCopyBasePass(), ReplaceViewCopyWithMemoryViewPass(), config.to_out_var_pass, config.memory_planning_pass. The first two steps are the just the first and second diff described above. In config.to_out_var_pass, the memory.view nodes are skipped. In config.memory_planning_pass, when a spec is requested for a memory.view node (e.g., to update the lifetime), we return the spec of its base. Returning the spec for the base means that whenever we see a memory.view node, we actually update the lifetime of the base to cover it. Moreover, the memory.view node's special _MemoryViewSpec sees this update reflected. (Note that an exception would be thrown if we kept the usual flow and returned the spec for the memory.view node. This is because the special _MemoryViewSpec is immutable and would not allow the memory_planning_pass to update its lifetime.) Finally, during emission the memory.view is emitted as an evalue. There are two more diffs on the stack D54866523 and D54866539. The first of these replaces the old RemoveRedundantViewCopy pass with a NormalizeViewCopyBasePass + dead code elimination. The second converts view-like ops (squeeze, unsqueeze, slice) to view ops when safe to do so to take advantage of the view_copy elimination. Differential Revision: D54827305
diff --git a/exir/memory.py b/exir/memory.py
@@ -39,3 +39,12 @@ def free(spec: TensorSpec) -> None:
     E.g., it can be the target of call_function node.
     """
     pass
+
+
+def view(base: torch.Tensor, size: List[int]) -> torch.Tensor:
+    """
+    This function mimics torch.ops.aten.view.default.
+
+    It is used to elide memory-planned, static view_copy nodes.
+    """
+    return base.view(size)
diff --git a/exir/memory_planning.py b/exir/memory_planning.py
@@ -35,6 +35,16 @@
 
 REGISTERED_ALGOS: Dict[str, Callable[..., List[int]]] = {}
 
+SPECIAL_TARGETS = [  # pyre-ignore
+    memory.alloc,
+    memory.view,
+    operator.getitem,
+    torch.ops.higher_order.cond,
+    exir_while,
+    torch.ops.higher_order.map_impl,
+    executorch_call_delegate,
+]
+
 
 class Verifier:
     """
@@ -393,16 +403,7 @@ def collect_specs_from_nodes(  # noqa: C901
 
         if do_assertion:
             internal_assert(
-                node.op in ("placeholder", "output")
-                or node.target
-                in [
-                    memory.alloc,
-                    operator.getitem,
-                    torch.ops.higher_order.cond,
-                    exir_while,
-                    torch.ops.higher_order.map_impl,
-                    executorch_call_delegate,
-                ],
+                node.op in ("placeholder", "output") or node.target in SPECIAL_TARGETS,
                 f"Unexpected op {node.op}, target {node.target}",
             )
         for spec in specs:
@@ -689,6 +690,24 @@ def get_input_specs(graph_module: fx.GraphModule) -> Set[TensorSpec]:
     return input_specs
 
 
+def is_op_memory_planned(node: torch.fx.Node) -> bool:
+    """
+    Return true if this call function node is memory planned.
+    We are cautious in this function and are OK returning False even if a node
+    may be memory planned.
+    """
+    assert node.op == "call_function"
+    if node.target in SPECIAL_TARGETS:
+        # Assume nodes in SPECIAL_TARGETS are not memory planned,
+        # unless they are in the exceptions below
+        exceptions = [
+            executorch_call_delegate,
+        ]
+        return node.target in exceptions
+
+    return True
+
+
 def insert_calls_to_free(
     graph_module: fx.GraphModule, allspecs: Set[TensorSpec]
 ) -> None:
diff --git a/exir/passes/TARGETS b/exir/passes/TARGETS
@@ -323,3 +323,17 @@ python_library(
         "//executorch/exir/dialects:lib",
     ],
 )
+
+python_library(
+    name = "replace_view_copy_with_memory_view_pass",
+    srcs = [
+        "replace_view_copy_with_memory_view_pass.py",
+    ],
+    deps = [
+        "//caffe2:torch",
+        "//executorch/exir:memory",
+        "//executorch/exir:memory_planning",
+        "//executorch/exir:tensor",
+        "//executorch/exir/dialects:lib",
+    ],
+)
diff --git a/exir/passes/replace_view_copy_with_memory_view_pass.py b/exir/passes/replace_view_copy_with_memory_view_pass.py
@@ -0,0 +1,218 @@
+# 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 math
+from typing import Any, Dict, List, Optional, Tuple
+
+import torch
+
+from executorch.exir import memory
+from executorch.exir.dialects._ops import ops
+from executorch.exir.memory_planning import is_op_memory_planned
+from executorch.exir.tensor import (
+    contiguous_stride_from_shape,
+    determine_tensor_dynanism,
+    dim_order_from_stride,
+    TensorShapeDynamism,
+    TensorSpec,
+)
+from torch.export.exported_program import ExportedProgram
+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 _MemoryViewSpec(TensorSpec):
+    def __init__(self, base: TensorSpec, shape: List[int]) -> None:
+        """
+        A MemoryViewSpec is an immutable TensorSpec that mirrors its base for non-size
+        related information.
+        """
+
+        if math.prod(base.shape) != math.prod(shape):
+            raise Exception(
+                f"Cannot create a MemoryViewSpec because the provided shape {shape} is not consistent with the number of elements in the provided base ({math.prod(base.shape)})."
+            )
+
+        if not base.is_static_shape_tensor:
+            raise Exception(
+                "Cannot create a MemoryViewSpec because the provided base is not a static shape tensor."
+            )
+
+        self._init_setters = [
+            "_frozen",
+            "_base",
+            "_guards",
+            "shape",
+            "stride",
+            "dim_order",
+            "shape_dynamism",
+        ]
+        self._frozen = False
+        self._base = base
+        self.shape: List[int] = shape
+        self.stride: Tuple[int] = contiguous_stride_from_shape(torch.Size(self.shape))
+        self.dim_order: Tuple[bytes] = dim_order_from_stride(self.stride)
+        self.shape_dynamism: TensorShapeDynamism = determine_tensor_dynanism(
+            torch.Size(self.shape)
+        )
+
+        # This spec gives a view into its base.
+        # The base can be modified (e.g., mem_id) and this spec will
+        # update accordingly, but certain fields we do not expect to change
+        # We create guards for these
+        self._guards: Dict[str, Any] = {
+            "shape_dynamism": base.shape_dynamism,
+            "scalar_type": base.scalar_type,
+            "layout": base.layout,
+            "is_sparse": base.is_sparse,
+        }
+        self._frozen = True
+
+    def _check_guards(self) -> None:
+        for name in self._guards:
+            if getattr(self._base, name) != self._guards[name]:
+                raise Exception(
+                    f"The guarded attribute '{name}' has changed value.  At creation of the MemoryViewSpec, it was {self._guards[name]}, but it is now {getattr(self._base, name)}."
+                )
+
+    def __getattribute__(self, name):  # pyre-ignore
+        if name in [
+            "_init_setters",
+            "_frozen",
+            "_base",
+            "_guards",
+            "_check_guards",
+            # Adding debug is needed so that view_spec.debug() shows the right id in
+            # its string (if debug is excluded, it shows the id(view_spec._base) instead
+            # of id(view_spec))
+            "debug",
+        ]:
+            return object.__getattribute__(self, name)
+
+        # Guard check after freeze
+        if self._frozen:
+            self._check_guards()
+
+        # self._init_setters attributes come from self, others come from base
+        if name in self._init_setters:
+            return object.__getattribute__(self, name)
+        return getattr(self._base, name)
+
+    def __setattr__(self, name: str, val) -> None:  # pyre-ignore
+        if name in ["_init_setters", "_frozen"]:
+            object.__setattr__(self, name, val)
+            return
+
+        # Allow setting during initialization
+        if name in self._init_setters and not self._frozen:
+            object.__setattr__(self, name, val)
+            return
+
+        if name in self._init_setters:
+            raise Exception(
+                f"MemoryViewSpec is immutable.  Cannot set the attribute '{name}' after creation."
+            )
+
+        raise Exception(
+            f"MemoryViewSpec is immutable.  To update the non-size related attribute '{name}', update the base."
+        )
+
+
+class ReplaceViewCopyWithMemoryViewPass(PassBase):
+    def __init__(self) -> None:
+        super().__init__()
+        self._program: Optional[ExportedProgram] = None
+
+    def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
+        """
+        This pass replaces all static, memory-planned view_copy nodes with special
+        memory.view nodes.
+
+        This should be run after the NormalizeViewCopyBasePass.
+
+        During memory planning, memory.view nodes share the same storage as their base.
+
+        During emission, memory.view nodes are not emitted as operators, but are instead
+        directly emitted as evalues.  They share the same allocation_info/storage as their
+        base.
+        """
+
+        n_replaced = 0
+        for module in graph_module.modules():
+            if not isinstance(module, torch.fx.GraphModule):
+                continue
+            for node in module.graph.nodes:
+                if self._is_replaceable_view_copy(node):
+                    base, shape = node.args
+                    node.target = memory.view
+
+                    # Create spec for the node.
+                    # _MemoryViewSpec is an immutable TensorSpec gives a view into
+                    # its base spec for non-size related information.
+                    node.meta["spec"] = _MemoryViewSpec(base.meta["spec"], shape)
+
+                    n_replaced += 1
+
+            module.recompile()
+
+        logging.debug(f"Replaced {n_replaced} view_copy nodes with memory.view nodes.")
+        return PassResult(graph_module, n_replaced > 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:
+                assert not self._is_replaceable_view_copy(node)
+                if node.op == "call_function" and node.target == memory.view:
+                    assert isinstance(node.meta["spec"], _MemoryViewSpec)
+
+    def requires(self, graph_module: torch.fx.GraphModule) -> None:
+        """
+        This pass should be called after NormalizeViewCopyBasePass.
+        We check that all view_copy nodes have been normalized.
+        """
+        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)
+
+    def set_program(self, program: ExportedProgram) -> None:
+        self._program = program
+
+    def _is_replaceable_view_copy(self, node: torch.fx.Node) -> bool:
+        if not _is_view_copy(node):
+            return False
+
+        base = node.args[0]
+        assert isinstance(base, torch.fx.Node)
+
+        is_base_memory_planned = False  # until proven otherwise
+        if base.op == "call_function":
+            is_base_memory_planned = is_op_memory_planned(base)
+
+        if base.op == "placeholder":
+            # For now, we only assume placeholder parameters are memory planned.
+            # We cautiously assume that general user inputs and buffers are not memory planned.
+            # Memory planning for user inputs can be turned off, see
+            # https://github.com/pytorch/executorch/blob/main/exir/passes/memory_planning_pass.py#L32
+            assert isinstance(self._program, ExportedProgram)
+            if base.name in self._program.graph_signature.inputs_to_parameters:
+                is_base_memory_planned = True
+
+        return is_base_memory_planned and node.meta["spec"].is_static_shape_tensor
diff --git a/exir/tests/TARGETS b/exir/tests/TARGETS
@@ -208,6 +208,7 @@ python_unittest(
         "//executorch/exir:memory",
         "//executorch/exir:memory_planning",
         "//executorch/exir:pass_base",
+        "//executorch/exir:schema",
         "//executorch/exir:tensor",
         "//executorch/exir/dialects:lib",
         "//executorch/exir/dialects/edge:lib",
@@ -220,6 +221,7 @@ python_unittest(
         "//executorch/exir/passes:remove_graph_asserts_pass",
         "//executorch/exir/passes:remove_mixed_type_operators",
         "//executorch/exir/passes:replace_edge_with_backend_pass",
+        "//executorch/exir/passes:replace_view_copy_with_memory_view_pass",
         "//executorch/exir/passes:scalar_to_tensor_pass",
         "//executorch/exir/passes:spec_prop_pass",
         "//executorch/exir/passes:sym_to_tensor_pass",
diff --git a/exir/tests/test_passes.py b/exir/tests/test_passes.py
