From acd31c58b64e2b34792e008673cb719470783819 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Mon, 2 Nov 2026 08:41:50 -0800 Subject: [PATCH] CUDA backend: delete each output handle once when a program returns one tensor twice (#24420) ### What is wrong today When a CUDA program returns the same tensor more than once, for example `y.detach()` twice, AOTI (AOTInductor, PyTorch's ahead-of-time compiler that generates the CUDA code) puts one output 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 that slot's handle. So it deleted the shared handle twice, and the process crashed, for example with: ``` 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 ``` ### What this change does A handle is now deleted only at the last output slot that holds it, as PyTorch's own AOTI runner does. The cleanup on error paths follows the same rule. With CUDA graphs on, a successful capture keeps every output handle for the graph and deletes none of them, so nothing changes there. If ending the capture or instantiating the graph fails, the cleanup deletes the handles by the same rule. A comment in `ReplaceViewCopyWithViewPass` said that a view would make two outputs share one handle. That was misleading, because two `alias_copy` outputs of one tensor share one handle too, so the comment no longer says it. ### What was tested `test_same_tensor_returned_twice` returns one tensor in two output slots that are not next to each other, with CUDA graphs on: three eager calls, a graph capture and a replay, each compared with eager. With torch 2.14 on an A100 it passes, and it crashes without this change. The other `backends/cuda` tests that CI runs in the same job still pass. (cherry picked from commit 5c05ad471d159858a814f8e72126280cd586682a) --- .../passes/replace_view_copy_with_view.py | 6 ++ backends/cuda/runtime/cuda_backend.cpp | 18 +++- backends/cuda/tests/test_cuda_export.py | 93 +++++++++++++++++++ 3 files changed, 114 insertions(+), 3 deletions(-) diff --git a/backends/aoti/passes/replace_view_copy_with_view.py b/backends/aoti/passes/replace_view_copy_with_view.py index 460d1f09981..cee89943f9e 100644 --- a/backends/aoti/passes/replace_view_copy_with_view.py +++ b/backends/aoti/passes/replace_view_copy_with_view.py @@ -31,6 +31,12 @@ # natively by inductor, so map it here like the other view-copy ops. torch.ops.aten.split_copy.Tensor: torch.ops.aten.split.Tensor, 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``. ``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 921c15ae7af..1ef86060a5f 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); } } }); @@ -1136,7 +1148,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. @@ -1145,7 +1157,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 8329f42e522..ffb455c4664 100644 --- a/backends/cuda/tests/test_cuda_export.py +++ b/backends/cuda/tests/test_cuda_export.py @@ -562,6 +562,99 @@ def test_triton_kernel_mode_off(self): "SDPA kernel export with triton_kernel_mode=OFF failed", ) + def test_missing_scatter_kernel_fails_at_export(self): + """Inductor leaves int64 scatter_add to an ATen kernel the CUDA shim lacks.""" + + class Int64ScatterAdd(torch.nn.Module): + def forward(self, counts: torch.Tensor, idx: torch.Tensor): + return counts.scatter_add(0, idx, torch.ones_like(idx)) + + inputs = (torch.zeros(8, dtype=torch.long), torch.randint(0, 8, (32,))) + with self.assertRaises(RuntimeError) as context: + self._export_to_cuda_with_lower(Int64ScatterAdd(), inputs) + self.assertIn("aoti_torch_cuda_scatter_reduce_two_out", str(context.exception)) + + @unittest.skipIf( + torch.version.hip is not None, + "loads through the Python runtime, which the ROCm build does not include", + ) + def test_mask_index_and_index_put(self): + """A mask index makes the result size depend on the data, and Inductor + has no atomic add for int64, so these run in the index shims.""" + from executorch.runtime import Runtime + + class MaskIndex(torch.nn.Module): + def forward(self, x, mask, column, idx, counts): + out = x.clone() + out[mask] = x[mask] * 2 + # The mask and column index dims 0 and 2, which are not adjacent. + picked = x[mask, :, column].sum(0) + # A transposed input gives a result that is not row-major. + rows = x.transpose(1, 2)[mask] * 3 + ones = torch.ones_like(idx) + counts = counts.index_put((idx,), ones, accumulate=True) + return out, picked, rows, counts + + model = MaskIndex() + x = torch.randn(4, 5, 6) + column = torch.tensor([2]) + idx = torch.tensor([0, 3, 3, 1, 3]) + # Nonzero counts check that repeated writes add to the old values. + counts = torch.arange(8) * 5 + masks = ( + torch.tensor([True, False, True, True]), + torch.zeros(4, dtype=torch.bool), + torch.ones(4, dtype=torch.bool), + ) + edge = self._export_to_cuda_with_lower( + model, (x, masks[0], column, idx, counts) + ) + 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) + method = Runtime.get().load_program(path).load_method("forward") + for mask in masks: + inputs = (x, mask, column, idx, counts) + outputs = method.execute(list(inputs)) + 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