From 1af26c75335d6989ea74e59f3fbab135e7c4c916 Mon Sep 17 00:00:00 2001 From: Jacky Li <86073892+JPL11@users.noreply.github.com> Date: Mon, 6 Jul 2026 10:13:47 -0700 Subject: [PATCH] Insert write-back copy_ nodes at the earliest safe point The insert_write_back_for_buffers pass placed every write-back copy_ at the end of the graph, arbitrarily extending the lifetime of the value being written back and wasting space in the memory plan. Now each copy_(buffer, value) is inserted at the earliest point that preserves the end-of-graph semantics: after the value is computed, after every reader of the buffer or any alias of it (they must observe the old contents), and after any mutation of the value or any alias of it (so we snapshot the final value). Aliases are found with a forward walk using schema alias_info, treating getitem, submodule calls, and schema-less targets conservatively. If the value written back by one copy may alias the buffer mutated by another, all copies fall back to the old end-of-graph placement in their original order. Fixes #7345 --- .../insert_write_back_for_buffers_pass.py | 156 +++++++++++++++++- exir/tests/test_passes.py | 73 +++++++- 2 files changed, 223 insertions(+), 6 deletions(-) diff --git a/exir/passes/insert_write_back_for_buffers_pass.py b/exir/passes/insert_write_back_for_buffers_pass.py index 4dce40ae57c..39e6d2dbfab 100644 --- a/exir/passes/insert_write_back_for_buffers_pass.py +++ b/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,6 +22,112 @@ 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]], @@ -28,15 +135,26 @@ def _insert_copy( ): """ 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) ) diff --git a/exir/tests/test_passes.py b/exir/tests/test_passes.py index 20906fe92e9..a93c8c28636 100644 --- a/exir/tests/test_passes.py +++ b/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):