diff --git a/backends/aoti/passes/replace_view_copy_with_view.py b/backends/aoti/passes/replace_view_copy_with_view.py index c881e650e7a..6363910e43f 100644 --- a/backends/aoti/passes/replace_view_copy_with_view.py +++ b/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, } diff --git a/backends/cuda/runtime/cuda_backend.cpp b/backends/cuda/runtime/cuda_backend.cpp index ce429a9a65a..2dcda0004f0 100644 --- a/backends/cuda/runtime/cuda_backend.cpp +++ b/backends/cuda/runtime/cuda_backend.cpp @@ -934,6 +934,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 @@ -944,7 +956,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); } } }); @@ -1168,7 +1180,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. @@ -1188,7 +1200,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; } diff --git a/backends/cuda/tests/test_cuda_export.py b/backends/cuda/tests/test_cuda_export.py index 637ff4036b1..3ed4878ec5a 100644 --- a/backends/cuda/tests/test_cuda_export.py +++ b/backends/cuda/tests/test_cuda_export.py @@ -766,6 +766,41 @@ def forward(self, x, mask, column, idx, counts): for output, expected in zip(outputs, model(*inputs), strict=True): torch.testing.assert_close(output.cpu(), expected) + @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