Insert write-back copy_ nodes at the earliest safe point

exir/passes/insert_write_back_for_buffers_pass.py

@@ -4,7 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-from typing import Dict, List, Optional, Tuple
+import operator
+from typing import Dict, List, Optional, Set, Tuple
 
 import torch
 from executorch.exir.operator.convert import is_inplace_variant
@@ -21,22 +22,139 @@
 from torchgen.model import SchemaKind
 
 
+def _may_alias_input(node: torch.fx.Node) -> bool:
+    """
+    Whether the value produced by this node may alias one of its inputs. When
+    we cannot tell (no schema, getitem, submodule calls, etc.) we
+    conservatively answer True.
+    """
+    if node.op != "call_function":
+        return True
+    if node.target is operator.getitem:
+        return True
+    schema = getattr(node.target, "_schema", None)
+    if schema is None:
+        return True
+    return any(ret.alias_info is not None for ret in schema.returns)
+
+
+def _mutates_input(node: torch.fx.Node, input_node: torch.fx.Node) -> bool:
+    """
+    Whether this node may mutate the value passed to it as input_node. When we
+    cannot tell we conservatively answer True.
+    """
+    if node.op == "output":
+        return False
+    if node.op != "call_function":
+        return True
+    schema = getattr(node.target, "_schema", None)
+    if schema is None:
+        return True
+    for i, arg in enumerate(node.args):
+        if arg is input_node and i < len(schema.arguments):
+            alias_info = schema.arguments[i].alias_info
+            if alias_info is not None and alias_info.is_write:
+                return True
+    schema_kwargs = {a.name: a for a in schema.arguments}
+    for name, arg in node.kwargs.items():
+        if arg is input_node and name in schema_kwargs:
+            alias_info = schema_kwargs[name].alias_info
+            if alias_info is not None and alias_info.is_write:
+                return True
+    return False
+
+
+def _collect_aliases(
+    seed: torch.fx.Node, node_order: Dict[torch.fx.Node, int]
+) -> Set[torch.fx.Node]:
+    """
+    The set of nodes whose values may alias the value of seed, found by
+    walking forward through the graph.
+    """
+    aliases = {seed}
+    for node in node_order:
+        if node in aliases:
+            continue
+        if any(arg in aliases for arg in node.all_input_nodes) and _may_alias_input(
+            node
+        ):
+            aliases.add(node)
+    return aliases
+
+
+def _insertion_point(
+    mutated_node: torch.fx.Node,
+    return_node: torch.fx.Node,
+    node_order: Dict[torch.fx.Node, int],
+    last_placeholder: torch.fx.Node,
+) -> torch.fx.Node:
+    """
+    The earliest node after which it is safe to insert
+    copy_(mutated_node, return_node), preserving the semantics of inserting it
+    at the end of the graph. The copy_ must come after:
+
+     * return_node itself, and any node that may mutate it (or an alias of
+       it), so that we write back the final value;
+     * every reader of mutated_node or an alias of it, since they must observe
+       the old value of the buffer (this also orders us after anything that
+       may mutate the buffer);
+     * all placeholders.
+    """
+    latest = last_placeholder
+    if node_order[return_node] > node_order[latest]:
+        latest = return_node
+
+    for alias in _collect_aliases(mutated_node, node_order):
+        for user in alias.users:
+            # Users not in node_order are copy_ nodes inserted by us for other
+            # buffers; ordering with respect to them is handled by the
+            # independence check in _insert_copy.
+            if (
+                user.op != "output"
+                and user in node_order
+                and node_order[user] > node_order[latest]
+            ):
+                latest = user
+
+    for alias in _collect_aliases(return_node, node_order):
+        for user in alias.users:
+            if (
+                user in node_order
+                and _mutates_input(user, alias)
+                and node_order[user] > node_order[latest]
+            ):
+                latest = user
+
+    return latest
+
+
 def _insert_copy(
     gm: torch.fx.GraphModule,
     mutated_outputs: List[Optional[str]],
     input_name_to_node: Dict[str, torch.fx.Node],
 ):
     """
     Find the all the buffers and inputs that were mutated and insert copy_
-    operators to reflect mutations.
+    operators to reflect mutations. Each copy_ is inserted at the earliest
+    point at which it is safe, rather than at the end of the graph, so that
+    the memory planner does not have to arbitrarily extend the lifetime of the
+    value written back.
     """
     output_node = gm.graph.output_node()
     assert output_node is not None
     outputs = pytree.tree_flatten(output_node.args)[0]
     assert len(outputs) == len(mutated_outputs)
 
+    node_order: Dict[torch.fx.Node, int] = {
+        node: i for i, node in enumerate(gm.graph.nodes)
+    }
+    last_placeholder = [node for node in gm.graph.nodes if node.op == "placeholder"][
+        -1
+    ]
+
+    # Pair up the returns with the nodes they mutate.
+    copies: List[Tuple[torch.fx.Node, torch.fx.Node]] = []
     user_output_nodes = []
-    buffer_output_nodes = []
     for return_node, mutated_node_name in zip(outputs, mutated_outputs):
         # User output, leave alone
         if mutated_node_name is None:
@@ -50,9 +168,37 @@ def _insert_copy(
             raise RuntimeError(
                 f"Could not find {mutated_node_name} in either buffer or input nodes"
             )
+        copies.append((mutated_node, return_node))
+
+    # The copies themselves mutate the buffers. If the value written back by
+    # one copy may alias the buffer mutated by another, then the order of the
+    # copies (and their position relative to everything else) matters in ways
+    # the insertion points below do not track, so fall back to inserting all
+    # of them at the end of the graph, in their original order, as before.
+    independent = True
+    if len(copies) > 1:
+        mutated_aliases: Set[torch.fx.Node] = set()
+        return_aliases: Set[torch.fx.Node] = set()
+        for i, (mutated_node, return_node) in enumerate(copies):
+            mutated_alias = _collect_aliases(mutated_node, node_order)
+            return_alias = _collect_aliases(return_node, node_order)
+            if mutated_alias & return_aliases or return_alias & mutated_aliases:
+                independent = False
+                break
+            mutated_aliases |= mutated_alias
+            return_aliases |= return_alias
 
-        # insert copy
-        with gm.graph.inserting_before(output_node):
+    # insert the copies
+    buffer_output_nodes = []
+    for mutated_node, return_node in copies:
+        if independent:
+            insert_after = _insertion_point(
+                mutated_node, return_node, node_order, last_placeholder
+            )
+            insertion = gm.graph.inserting_after(insert_after)
+        else:
+            insertion = gm.graph.inserting_before(output_node)
+        with insertion:
             buffer_output = gm.graph.call_function(
                 torch.ops.aten.copy_.default, (mutated_node, return_node)
             )

exir/tests/test_passes.py

@@ -1935,14 +1935,85 @@ def forward(self, x):
         #     %b_direct_copy_from_input : [num_users=1] = placeholder[target=b_direct_copy_from_input]
         #     %_lifted_tensor_constant2 : [num_users=1] = placeholder[target=_lifted_tensor_constant2]
         #     %x : [num_users=2] = placeholder[target=x]
+        #     %copy__default_1 : [num_users=1] = call_function[target=torch.ops.aten.copy_.default](args = (%b_direct_copy_from_input, %x), kwargs = {})
         #     %aten_add_tensor : [num_users=1] = call_function[target=executorch.exir.dialects.edge._ops.aten.add.Tensor](args = (%x, %b_state), kwargs = {})
         #     %dim_order_ops__to_dim_order_copy_default : [num_users=1] = call_function[target=executorch.exir.dialects.edge._ops.dim_order_ops._to_dim_order_copy.default](args = (%_lifted_tensor_constant2,), kwargs = {dtype: torch.float32, dim_order: []})
         #     %aten_add_tensor_1 : [num_users=1] = call_function[target=executorch.exir.dialects.edge._ops.aten.add.Tensor](args = (%b_state, %dim_order_ops__to_dim_order_copy_default), kwargs = {})
         #     %copy__default : [num_users=1] = call_function[target=torch.ops.aten.copy_.default](args = (%b_state, %aten_add_tensor_1), kwargs = {})
-        #     %copy__default_1 : [num_users=1] = call_function[target=torch.ops.aten.copy_.default](args = (%b_direct_copy_from_input, %x), kwargs = {})
         #     return (copy__default, copy__default_1, aten_add_tensor)
         self.assertEqual(count_copies(gm), 2)
 
+    def test_mutable_buffers_write_back_is_inserted_early(self) -> None:
+        class EarlyMutationModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.register_buffer("state", torch.zeros(1))
+
+            def forward(self, x):
+                # The buffer's new value is computed at the very start, so its
+                # write-back can happen immediately, letting the memory
+                # planner reuse the space during the rest of the graph.
+                self.state.add_(1)
+                y = x + 1
+                y = y + 1
+                y = y + 1
+                return y
+
+        model = to_edge(
+            export(EarlyMutationModule(), (torch.zeros(1),), strict=True)
+        )
+        gm, _ = insert_write_back_for_buffers_pass(model.exported_program())
+
+        node_order = {node: i for i, node in enumerate(gm.graph.nodes)}
+        copies = [
+            node
+            for node in gm.graph.nodes
+            if node.target == torch.ops.aten.copy_.default
+        ]
+        self.assertEqual(len(copies), 1)
+        copy = copies[0]
+        # The copy_ comes right after the value it writes back, not at the end
+        # of the graph: every user computation (the adds on x) is after it.
+        self.assertEqual(node_order[copy], node_order[copy.args[1]] + 1)
+        output_node = gm.graph.output_node()
+        user_return = output_node.args[0][1]
+        self.assertLess(node_order[copy], node_order[user_return])
+
+    def test_mutable_buffers_write_back_after_old_value_reads(self) -> None:
+        class ReadOldValueModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.register_buffer("state", torch.zeros(1))
+
+            def forward(self, x):
+                # The buffer's new value is computed before the old value is
+                # read, so the write-back must not simply follow the value it
+                # writes: it must wait for the read of the old value.
+                new_state = x * 2
+                old_plus = self.state + x
+                self.state.copy_(new_state)
+                return old_plus
+
+        model = to_edge(
+            export(ReadOldValueModule(), (torch.zeros(1),), strict=True)
+        )
+        gm, _ = insert_write_back_for_buffers_pass(model.exported_program())
+
+        node_order = {node: i for i, node in enumerate(gm.graph.nodes)}
+        copies = [
+            node
+            for node in gm.graph.nodes
+            if node.target == torch.ops.aten.copy_.default
+        ]
+        self.assertEqual(len(copies), 1)
+        copy = copies[0]
+        buffer_placeholder = copy.args[0]
+        self.assertEqual(buffer_placeholder.op, "placeholder")
+        # Every read of the buffer's old value stays before the write-back.
+        for user in buffer_placeholder.users:
+            if user is not copy and user.op != "output":
+                self.assertLess(node_order[user], node_order[copy])
+
     def test_remove_quantized_op_noop_pass(self) -> None:
         class TestAddSliceNoop(torch.nn.Module):
             def __init__(self):