insert write back now tries to avoid inserting if all operations were inplace (#12008)

Summary:

Before only considered if the placeholder was the output. Now allow if all operations are inplace.

Differential Revision: D77204116

Author: JacobSzwejbka

exir/passes/insert_write_back_for_buffers_pass.py

@@ -15,8 +15,12 @@
     OutputKind,
     OutputSpec,
 )
+from executorch.exir.operator.convert import (
+    _get_overload_schema,
+)
 from torch.export.graph_signature import TensorArgument
 from torch.utils import _pytree as pytree
+from torchgen.model import SchemaKind
 
 
 def _insert_copy(
@@ -69,6 +73,28 @@ def _insert_copy(
         gm.graph.erase_node(output_node)
     return buffer_output_nodes
 
+def _inplace_lineage(output_arg: torch.fx.Node, gm: torch.fx.GraphModule, gs: ExportGraphSignature, kind: SchemaKind) -> bool:
+    """
+    Walk the graph backwards to see if output_arg is ultimately the same as an input.
+    """
+    if kind != OutputKind.BUFFER_MUTATION and kind != OutputKind.USER_INPUT_MUTATION:
+        return False
+
+    while output_arg.op != "placeholder":
+        if output_arg.op == "call_function" and _get_overload_schema(output_arg.target).kind() == SchemaKind.inplace: #pyre-ignore
+            # From looking at native_functions.yaml, inplace ops always have self as the first arg
+            output_arg = output_arg.args[0] #pyre-ignore
+        else:
+            return False
+
+    # If the output arg was a buffer then it needs to reach a buffer placeholder
+    if kind == OutputKind.BUFFER_MUTATION:
+        assert output_arg.target in gs.inputs_to_buffers
+        return True
+    # If the output arg was a user input then it needs to reach a user input placeholder
+    assert output_arg.target in gs.user_inputs
+    return True
+
 
 def insert_write_back_for_buffers_pass(
     ep: ExportedProgram,
@@ -99,9 +125,16 @@ def insert_write_back_for_buffers_pass(
         if lifted_node is not None:
             input_name_to_node[lifted_node] = input_node
 
+
+    output_node = None
+    for node in gm.graph.nodes:
+        if node.op == "output":
+            output_node = node
+            break
+
     # Grab the mutable buffer nodes in the outputs,
     mutated_outputs: List[Optional[str]] = []
-    for out_spec in ep.graph_signature.output_specs:
+    for i, out_spec in enumerate(ep.graph_signature.output_specs):
         # if the output arg is the input value then all operations on it are in-place
         # so there's no need to add a copy_ node
         if (
@@ -112,7 +145,7 @@ def insert_write_back_for_buffers_pass(
             out_spec.target in input_name_to_node
             and
             # if the arg and target are not the same, we add a copy_ node.
-            out_spec.arg.name != input_name_to_node[out_spec.target].name
+            not _inplace_lineage(output_node.args[0][i], gm, ep.graph_signature, ep.graph_signature.output_specs[i].kind)
         ):
             mutated_outputs.append(out_spec.target)
         else:

exir/tests/TARGETS

@@ -145,6 +145,7 @@ python_unittest(
         "//caffe2:torch",
         "//executorch/exir:lib",
         "//executorch/exir/passes:lib",
+        "//executorch/extension/pybindings:portable_lib",
     ],
 )