fix is_inplace_node check

Summary: Oops it could be an edge op which failed the opOverload check. Switch to just hasattr

Differential Revision: D77462717

Author: JacobSzwejbka

exir/passes/insert_write_back_for_buffers_pass.py

@@ -76,30 +76,31 @@ def _is_inplace_node(node: torch.fx.Node) -> bool:
     """Check if a node is an inplace node."""
     return (
         node.op == "call_function"
-        and isinstance(node.target, torch._ops.OpOverload)
+        and hasattr(node.target, "_schema")
         and is_inplace_variant(
-            node.target._schema.name, node.target._schema.overload_name
+            node.target._schema.name, node.target._schema.overload_name #pyre-ignore
         )
     )
 
 
 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:
+        print("Wrong kind")
         return False
 
     while output_arg.op != "placeholder":
         if _is_inplace_node(output_arg):
             # From looking at native_functions.yaml, inplace ops always have self as the first arg
             output_arg = output_arg.args[0]  # pyre-ignore
         else:
+            print("Not inplace " , output_arg)
             return False
 
     # If the output arg was a buffer then it needs to reach a buffer placeholder
@@ -160,7 +161,6 @@ def insert_write_back_for_buffers_pass(
             # if the arg and target are not the same, we add a copy_ node.
             not _inplace_lineage(
                 output_node.args[0][i],
-                gm,
                 ep.graph_signature,
                 ep.graph_signature.output_specs[i].kind,
             )

exir/tests/test_reinplace_pass.py

@@ -61,6 +61,15 @@ def forward(
 
         self.assertIsNotNone(index_put_node, "Should find an index_put_ node")
 
+        # Find the copy_ node
+        copy_node = None
+        for node in et.exported_program().graph.nodes:
+            if node.op == "call_function" and "copy_" in str(node.target):
+                copy_node = node
+                break
+
+        self.assertIsNone(copy_node, "Shoulnt find an copy_ node")
+
         e = _load_for_executorch_from_buffer(et.buffer)
         self.assertTrue(
             torch.allclose(