From dc628cc316198698a76c10489662f673a4b17581 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sun, 1 Nov 2026 22:20:15 -0800 Subject: [PATCH] Delete each CUDA output handle once when AOTI repeats it When a program returns the same tensor more than once, for example y.detach() twice, AOTI puts one handle in each of those output slots: output_handles[0] = buf0.release(); output_handles[1] = output_handles[0]; The CUDA backend copied each output into its ExecuTorch tensor and then deleted the slot's handle, so the shared handle was deleted twice. The second copy read a freed handle, and the process crashed: SlimTensor storage_offset must not be negative, got -9007650542493611503 Failed to D2D copy GPU output 1 into ETensor free(): double free detected in tcache 2 Now a handle is deleted only at the last slot that holds it, as PyTorch's own AOTI runner does. The cleanup on error paths follows the same rule. With CUDA graphs on, the capture step keeps every output handle for the graph and deletes none of them, so it needs no change. The comment in ReplaceViewCopyWithViewPass no longer says that a view would make two outputs share one handle: two detached copies of one tensor become two alias_copy outputs, and those share one handle too. test_same_tensor_returned_twice returns one tensor in two output slots that are not next to each other, with CUDA graphs on. So it runs the method eagerly, then captures a graph and replays it, and compares every output with eager. It crashes without this change. --- .../passes/replace_view_copy_with_view.py | 4 +-- backends/cuda/runtime/cuda_backend.cpp | 18 ++++++++-- backends/cuda/tests/test_cuda_export.py | 35 +++++++++++++++++++ 3 files changed, 52 insertions(+), 5 deletions(-) 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