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