Fix reinplace liveness for returned tensors (pytorch#23488)

The reinplace pass did not include graph outputs in reverse liveness
analysis. When a model returned both an original tensor and the result
of index_put, the pass could rewrite index_put to index_put_ and change
the original output. Treating output nodes as readers keeps returned
tensors live and prevents this incorrect rewrite while preserving valid
in-place rewrites.

Regression tests cover returned originals, nested outputs, later
readers, immutable inputs, view-copy behavior, and repeated
mutable-buffer executions. Validation: 20 tests passed in
exir/tests/test_reinplace_pass.py; repository commit lint hook passed.

This PR was authored with assistance from Codex (OpenAI AI assistant).

Author: JakeStevens

exir/passes/reinplace.py

@@ -337,7 +337,7 @@ def reinplace_pass(  # noqa: C901
 
     Safety rules:
       * The mutated arg must not be used by any later node in the
-        graph.
+        graph or be returned as an unmodified output.
       * If the mutated arg is a placeholder (program input), it must
         be a *mutable* input — i.e., declared in
         `graph_signature.user_inputs_to_mutate` or
@@ -425,7 +425,7 @@ def reinplace_pass(  # noqa: C901
 
     for node in reversed(ep.graph.nodes):
         if node.op != "call_function" or node.target not in resolved:
-            if node.op == "call_function":
+            if node.op in ("call_function", "output"):
                 seen_nodes.update(node.all_input_nodes)
             continue
 

exir/tests/test_reinplace_pass.py

@@ -90,6 +90,116 @@ def forward(
             )
         )
 
+    def test_returned_original_is_not_reinplaced(self) -> None:
+        class ReturnsOldAndNew(torch.nn.Module):
+            def forward(self, x, indices, values):
+                original = x + 1
+                updated = original.index_put((indices,), values)
+                return original, updated
+
+        model = ReturnsOldAndNew()
+        first = (torch.zeros(5), torch.tensor([0]), torch.tensor([9.0]))
+        edge = to_edge(export(model, first, strict=True))
+        et = edge.to_executorch(ExecutorchBackendConfig(run_reinplace_pass=True))
+        self.assertEqual(len(_find_nodes(et.exported_program(), "index_put_")), 0)
+
+        loaded = _load_for_executorch_from_buffer(et.buffer)
+        for inputs in (
+            first,
+            (torch.full((5,), 2.0), torch.tensor([3]), torch.tensor([-2.0])),
+        ):
+            with self.subTest(inputs=inputs):
+                for actual, expected in zip(loaded.forward(inputs), model(*inputs)):
+                    torch.testing.assert_close(actual, expected)
+
+    def test_nested_output_keeps_original_live(self) -> None:
+        class ReturnsNested(torch.nn.Module):
+            def forward(self, x, indices, values):
+                original = x + 1
+                updated = original.index_put((indices,), values)
+                return original, (original, updated)
+
+        inputs = (torch.zeros(5), torch.tensor([0]), torch.tensor([9.0]))
+        edge = to_edge(export(ReturnsNested(), inputs, strict=True))
+        ep = reinplace_pass(edge.exported_program())
+        self.assertEqual(len(_find_nodes(ep, "index_put_")), 0)
+        actual = ep.module()(*inputs)
+        expected = ReturnsNested()(*inputs)
+        for result, reference in zip(
+            (actual[0], *actual[1]), (expected[0], *expected[1])
+        ):
+            torch.testing.assert_close(result, reference)
+
+    def test_output_view_copy_preserves_original(self) -> None:
+        class ReturnsViewAndUpdated(torch.nn.Module):
+            def forward(self, x, indices, values):
+                original = x + 1
+                view = original.view(-1)
+                updated = original.index_put((indices,), values)
+                return view, updated
+
+        inputs = (torch.zeros(5), torch.tensor([0]), torch.tensor([9.0]))
+        edge = to_edge(export(ReturnsViewAndUpdated(), inputs, strict=True))
+        et = edge.to_executorch(ExecutorchBackendConfig(run_reinplace_pass=True))
+        # The view is materialized before the update, so it keeps the old value.
+        self.assertEqual(len(_find_nodes(et.exported_program(), "index_put_")), 1)
+        loaded = _load_for_executorch_from_buffer(et.buffer)
+        for actual, expected in zip(
+            loaded.forward(inputs), ReturnsViewAndUpdated()(*inputs)
+        ):
+            torch.testing.assert_close(actual, expected)
+
+    def test_later_reader_keeps_original_live(self) -> None:
+        class ReadsOriginalLater(torch.nn.Module):
+            def forward(self, x, indices, values):
+                original = x + 1
+                updated = original.index_put((indices,), values)
+                return updated, original * 2
+
+        inputs = (torch.zeros(5), torch.tensor([0]), torch.tensor([9.0]))
+        edge = to_edge(export(ReadsOriginalLater(), inputs, strict=True))
+        ep = reinplace_pass(edge.exported_program())
+        self.assertEqual(len(_find_nodes(ep, "index_put_")), 0)
+        for actual, expected in zip(
+            ep.module()(*inputs), ReadsOriginalLater()(*inputs)
+        ):
+            torch.testing.assert_close(actual, expected)
+
+    def test_immutable_input_is_not_reinplaced(self) -> None:
+        class UpdatesCopy(torch.nn.Module):
+            def forward(self, x, indices, values):
+                return x.index_put((indices,), values)
+
+        inputs = (torch.zeros(5), torch.tensor([0]), torch.tensor([9.0]))
+        edge = to_edge(export(UpdatesCopy(), inputs, strict=True))
+        ep = reinplace_pass(edge.exported_program())
+        self.assertEqual(len(_find_nodes(ep, "index_put_")), 0)
+        torch.testing.assert_close(ep.module()(*inputs), UpdatesCopy()(*inputs))
+        torch.testing.assert_close(inputs[0], torch.zeros(5))
+
+    def test_mutable_buffer_reinplace_across_executions(self) -> None:
+        class UpdatesBuffer(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.register_buffer("state", torch.zeros(5))
+
+            def forward(self, indices, values):
+                self.state.index_put_((indices,), values)
+                return self.state
+
+        first = (torch.tensor([0]), torch.tensor([1.0]))
+        edge = to_edge(export(UpdatesBuffer(), first, strict=True))
+        et = edge.to_executorch(ExecutorchBackendConfig(run_reinplace_pass=True))
+        self.assertEqual(len(_find_nodes(et.exported_program(), "index_put_")), 1)
+
+        loaded = _load_for_executorch_from_buffer(et.buffer)
+        reference = UpdatesBuffer()
+        for inputs in (first, (torch.tensor([2]), torch.tensor([3.0]))):
+            with self.subTest(inputs=inputs):
+                torch.testing.assert_close(
+                    loaded.forward(inputs)[0], reference(*inputs)
+                )
+
     def test_cant_reinplace(self) -> None:
         """Test that index_put on a mutable buffer that is viewed later is not safe."""