[release/1.6] Delete duplicate CUDA output handles once

backends/aoti/passes/replace_view_copy_with_view.py

@@ -33,8 +33,8 @@
     ops.edge.aten.split_copy.Tensor: ops.edge.aten.split.Tensor,
     # Edge turns ``detach`` into ``detach_copy``, which has no c-shim either.
     # Without autograd it is ``alias_copy``, which Inductor can decompose,
-    # unlike ``detach_copy``. A view would let two outputs share one handle,
-    # and ``clone`` fails Edge verification in dim order mode.
+    # unlike ``detach_copy``. ``clone`` fails Edge verification in dim order
+    # mode.
     torch.ops.aten.detach_copy.default: torch.ops.aten.alias_copy.default,
     ops.edge.aten.detach_copy.default: ops.edge.aten.alias_copy.default,
 }

backends/cuda/runtime/cuda_backend.cpp

@@ -956,6 +956,18 @@ class ET_EXPERIMENTAL CudaBackend final
 
     bool run_called = false;
 
+    // AOTI puts the same handle in every slot of an output that it returns
+    // more than once, for example y.detach() returned twice. So only the last
+    // slot that holds a handle deletes it, as PyTorch's AOTI runner does.
+    const auto delete_output = [&](size_t i) noexcept {
+      for (size_t j = i + 1; j < slim_outputs.size(); j++) {
+        if (slim_outputs[j] == slim_outputs[i]) {
+          return;
+        }
+      }
+      delete slim_outputs[i];
+    };
+
     // Scope guard: deletes any non-null slim_outputs on exit. Normal paths
     // null entries as they take ownership, so the guard only fires on
     // early-return error paths. Also cleans up inputs if run() was never
@@ -966,7 +978,7 @@ class ET_EXPERIMENTAL CudaBackend final
       }
       for (size_t i = 0; i < slim_outputs.size(); i++) {
         if (slim_outputs[i]) {
-          delete slim_outputs[i];
+          delete_output(i);
         }
       }
     });
@@ -1175,7 +1187,7 @@ class ET_EXPERIMENTAL CudaBackend final
       if (pre_run_outputs[i] == slim_outputs[i]) {
         // Case (a): AOTI wrote directly into our planned ET buffer. The result
         // is already in place — nothing to copy, just drop the view wrapper.
-        delete slim_outputs[i];
+        delete_output(i);
       } else {
         // Case (b): AOTI returned its own buffer. D2D copy the result into the
         // planned GPU ET buffer, then free AOTI's buffer.
@@ -1195,7 +1207,7 @@ class ET_EXPERIMENTAL CudaBackend final
                 slim_outputs[i], et_output, cuda_stream),
             "Failed to D2D copy GPU output %zu into ETensor",
             i);
-        delete slim_outputs[i];
+        delete_output(i);
       }
       slim_outputs[i] = nullptr;
     }

backends/cuda/tests/test_cuda_export.py

@@ -6,6 +6,7 @@
 
 import contextlib
 import os
+import tempfile
 import threading
 import unittest
 from typing import Tuple
@@ -708,6 +709,41 @@ def forward(self, counts: torch.Tensor, idx: torch.Tensor):
         self.assertIn("aoti_torch_cuda_scatter_reduce_two_out", message)
         self.assertIn("aoti_torch_index_put_out", message)
 
+    @unittest.skipIf(
+        torch.version.hip is not None,
+        "loads through the Python runtime, which the ROCm build does not include",
+    )
+    def test_same_tensor_returned_twice(self):
+        """AOTI gives two outputs one handle when they are the same tensor. With
+        CUDA graphs on, three calls run eagerly, the fourth captures a graph and
+        the fifth replays it."""
+        from executorch.runtime import Runtime
+
+        class Twice(torch.nn.Module):
+            def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, ...]:
+                y = x * 2
+                return y.detach(), x + 1, y.detach()
+
+        model = Twice()
+        edge = self._export_to_cuda_with_lower(model, (torch.randn(4, 8),))
+        runtime = Runtime.get()
+        option = "enable_cuda_graph_for_method"
+        with tempfile.TemporaryDirectory() as directory:
+            path = os.path.join(directory, "model.pte")
+            with open(path, "wb") as f:
+                edge.to_executorch().write_to_file(f)
+            runtime.backend_registry.set_option("CudaBackend", {option: "forward"})
+            try:
+                method = runtime.load_program(path).load_method("forward")
+            finally:
+                # The option is process wide.
+                runtime.backend_registry.set_option("CudaBackend", {option: ""})
+            for _ in range(5):
+                x = torch.randn(4, 8)
+                outputs = method.execute([x])
+                for output, expected in zip(outputs, model(x), strict=True):
+                    torch.testing.assert_close(output.cpu(), expected)
+
     def test_device_info_propagated_to_cuda_delegate_outputs(self):
         """
         Verify that, for a CUDA-delegated graph, every memory-planned tensor's