[aotinductor] Solves a problem where a tensor is returned more than once (#112177)

desertfire · pytorchmergebot · commit 67256d5c1c22 · 2023-11-06T20:12:25.000Z
Pull Request resolved: #112177 Approved by: https://github.com/zhxchen17
diff --git a/test/inductor/test_aot_inductor.py b/test/inductor/test_aot_inductor.py
@@ -1011,6 +1011,18 @@ def forward(self, x):
         x = torch.randn(5, device=self.device)
         self.check_model(Model(self.device), (x,))
 
+    def test_repeat_output(self):
+        class Model(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+
+            def forward(self, x):
+                y = torch.sin(x)
+                return y, y
+
+        example_inputs = (torch.randn(3, 10, device=self.device),)
+        self.check_model(Model(), example_inputs)
+
 
 class AOTInductorTestABICompatibleCpu(TestCase):
     device = "cpu"
@@ -1036,6 +1048,8 @@ class AOTInductorTestABICompatibleCpu(TestCase):
         "test_freezing": TestFailure(("abi_compatible_cpu",), is_skip=True),
         "test_normal_functional": TestFailure(("abi_compatible_cpu",)),
         "test_poi_multiple_dynamic": TestFailure(("abi_compatible_cpu",)),
+        # There is a double-free issue which will be fixed in another PR
+        "test_repeat_output": TestFailure(("abi_compatible_cpu",), is_skip=True),
         "test_sdpa": TestFailure(("abi_compatible_cpu",)),
         "test_sdpa_2": TestFailure(("abi_compatible_cpu",)),
         "test_simple_dynamic": TestFailure(("abi_compatible_cpu",)),
@@ -1058,6 +1072,8 @@ class AOTInductorTestABICompatibleCuda(TestCase):
     {
         "test_dup_unbacked_sym_decl": TestFailure(("abi_compatible_cuda",)),
         "test_normal_functional": TestFailure(("abi_compatible_cuda",)),
+        # There is a double-free issue which will be fixed in another PR
+        "test_repeat_output": TestFailure(("abi_compatible_cuda",), is_skip=True),
     },
 )
 
diff --git a/torch/_export/exported_program.py b/torch/_export/exported_program.py
@@ -71,7 +71,9 @@ def _unlift(gm, inp_pos_to_param_buffer_name, in_spec, out_spec, state_dict, buf
         # Step 2: Find the all the buffers that were mutated and update them
         if node.op == "output":
             user_output_nodes = []
-            for return_node in node.all_input_nodes:
+            # In the case that the same node is returned multiple times,
+            # node.all_input_nodes will only iterate that node once
+            for return_node in pytree.tree_flatten(node.args)[0]:
                 return_node_name = return_node.name
                 # we found a param/buffer mutation
                 if return_node_name in buffers_to_mutate:
