From ca5fa2f6c4d6791a9d16556868c4ccec01d0e421 Mon Sep 17 00:00:00 2001 From: PyTorch Bot Date: Mon, 21 Sep 2026 23:45:37 -0700 Subject: [PATCH] Borrow the device through the shared CUDA guard The ExecuTorch delegate selects a CUDA device in three places: when a program is loaded, when it runs, and when its engine is freed. Two of those three selected the device and never put the caller's back, so loading or freeing a program moved the calling thread to another card and left it there. The next unrelated allocation on that thread then landed on the wrong card. Only the third one, the call that does the work, restored anything, and it did so through a small struct written for that one function. This matters for a program split between this delegate and the CUDA one, with engines on different cards, because the thread that loads one delegate is not doing that delegate's work. ExecuTorch now ships a scope guard for this in its CUDA extension, beside the caller-stream helper this file already uses. It selects a device, restores the previous one when it goes out of scope, refuses an index it cannot select, and keys the restore on what the caller asked for rather than on whether the selection reported success, because a failed selection can still move the calling thread. All three places use it now, and the local struct is gone. The destructor deliberately ignores the result. It can run while an arena is torn down, on a thread that was working on another card, and a destructor has no way to report a failure. The guard logs one instead. The vendored build of the extension carried only the caller-stream source, so the guard's definitions would have been absent from the shared library it produces. It builds both sources now, and exports both headers. A test covers the fix. It exports a coalesced program with its engine on the second card, puts the caller on the first, and reads the current device after the load, after the method is prepared, after the run, and after the free. It skips with a reason on a machine with one card, because with one card there is nothing to switch to and so nothing to restore. Test plan: built an ExecuTorch wheel carrying the guard, installed it into a fresh environment, and built this file against that wheel alone. It compiles with no warnings, and the two symbols it needs for the guard are both exported by that wheel's own library, so a build with only the wheel can link it. --- .../executorch/TensorRTBackend.cpp | 47 ++++-------- .../test_cuda_partitioner_composition.py | 71 +++++++++++++++++++ third_party/executorch/BUILD | 32 +++++++-- 3 files changed, 111 insertions(+), 39 deletions(-) diff --git a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp index bbcd5e196c..e58edd2389 100644 --- a/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp +++ b/cpp/src/torch_tensorrt/executorch/TensorRTBackend.cpp @@ -28,6 +28,7 @@ #include #include +#include #include #include #include @@ -78,7 +79,11 @@ void TRTLogger::log(Severity severity, const char* msg) noexcept { } EngineHandle::~EngineHandle() { - cudaSetDevice(device_id); + // Freeing runs on the engine's device. Borrowed rather than selected outright, + // because this can run from arena teardown on a thread that was working + // elsewhere, and a destructor has no way to report a failure. + const auto guard = ::executorch::extension::cuda::CUDAGuard::create(device_id); + (void)guard; // No wait here: execute already waited, and a device-wide one blocks unrelated work. for (void* p : cached_input_ptrs) { if (p != nullptr) { @@ -328,16 +333,15 @@ Result TensorRTBackend::init( handle->output_binding_names = std::move(header.output_binding_names); handle->device_id = header.device_id; - cudaError_t cuda_err = cudaSetDevice(handle->device_id); - if (cuda_err != cudaSuccess) { - ET_LOG( - Error, "TensorRTBackend::init: cudaSetDevice(%d) failed: %s", handle->device_id, cudaGetErrorString(cuda_err)); + auto device_guard = ::executorch::extension::cuda::CUDAGuard::create(handle->device_id); + if (!device_guard.ok()) { + ET_LOG(Error, "TensorRTBackend::init: cannot select device %d", handle->device_id); return Error::InvalidProgram; } // Whether this device can reach pageable host memory at all. Speed is the query below. int pageable_access = 0; - cuda_err = cudaDeviceGetAttribute(&pageable_access, cudaDevAttrPageableMemoryAccess, handle->device_id); + cudaError_t cuda_err = cudaDeviceGetAttribute(&pageable_access, cudaDevAttrPageableMemoryAccess, handle->device_id); if (cuda_err != cudaSuccess) { ET_LOG( Info, @@ -649,36 +653,15 @@ Error TensorRTBackend::execute(BackendExecutionContext& context, DelegateHandle* return Error::InvalidArgument; } - int entry_device = -1; - cudaError_t cuda_err = cudaGetDevice(&entry_device); - if (cuda_err != cudaSuccess) { - ET_LOG(Error, "TensorRTBackend::execute: cudaGetDevice failed: %s", cudaGetErrorString(cuda_err)); - return Error::InvalidProgram; - } // Put the engine on its own device for multi-GPU correctness, restoring the // caller's device on exit; green-context confinement rides the selected stream, // independent of the current device/context. - const bool switch_device = (entry_device != engine->device_id); - if (switch_device) { - cuda_err = cudaSetDevice(engine->device_id); - if (cuda_err != cudaSuccess) { - ET_LOG( - Error, - "TensorRTBackend::execute: cudaSetDevice(%d) failed: %s", - engine->device_id, - cudaGetErrorString(cuda_err)); - return Error::InvalidProgram; - } + auto device_guard = ::executorch::extension::cuda::CUDAGuard::create(engine->device_id); + if (!device_guard.ok()) { + ET_LOG(Error, "TensorRTBackend::execute: cannot select device %d", engine->device_id); + return Error::InvalidProgram; } - struct DeviceRestore { - int device; - bool active; - ~DeviceRestore() { - if (active) { - cudaSetDevice(device); - } - } - } device_restore{entry_device, switch_device}; + cudaError_t cuda_err = cudaSuccess; std::unique_lock lock(engine->mu); diff --git a/tests/py/dynamo/executorch/test_cuda_partitioner_composition.py b/tests/py/dynamo/executorch/test_cuda_partitioner_composition.py index 93585144a2..264eff3f39 100644 --- a/tests/py/dynamo/executorch/test_cuda_partitioner_composition.py +++ b/tests/py/dynamo/executorch/test_cuda_partitioner_composition.py @@ -1,6 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: BSD-3-Clause +import gc import importlib.util import os import shutil @@ -390,3 +391,73 @@ def test_the_partitioner_refuses_a_device_it_cannot_run_on(requested, accepted): else: with pytest.raises(ValueError, match="not a device this delegate runs on"): TensorRTPartitioner(compile_specs=specs) + + +def test_the_caller_keeps_its_device_across_load_run_and_free(tmp_path): + """A program whose engine sits on another card must not move the caller. + + Loading, running and freeing all select the engine's device. Each one has to put + the caller's device back, because a coalesced program holds several delegates on + several cards and the thread that loads one is not doing that delegate's work. + Needs two cards: with one there is nothing to switch to, so nothing to restore. + """ + if torch.cuda.device_count() < 2: + pytest.skip( + "needs two CUDA devices to tell a restored device from an unchanged one" + ) + + import torch_tensorrt + + class Model(torch.nn.Module): + def forward(self, x): + return torch.cos(torch.erfinv(torch.tanh(x))) + + engine_device = 1 + caller_device = 0 + + with torch.cuda.device(engine_device): + model = Model().eval().to(f"cuda:{engine_device}") + inputs = (torch.randn(64, 64, device=f"cuda:{engine_device}"),) + exported = torch.export.export(model, inputs) + trt_gm = torch_tensorrt.dynamo.compile( + exported, inputs=list(inputs), min_block_size=1, truncate_double=True + ) + out = tmp_path / "other_card.pte" + torch_tensorrt.save( + trt_gm, + str(out), + output_format="executorch", + retrace=False, + arg_inputs=list(inputs), + partitioners=[_cuda_partitioner()], + ) + + delegate_ids = _delegate_ids(out) + assert ( + "TensorRTBackend" in delegate_ids + ), f"nothing went to TensorRT; {delegate_ids}" + + import torch_tensorrt_executorch_runtime # noqa: F401 + from executorch.runtime import Runtime + + torch.cuda.set_device(caller_device) + + program = Runtime.get().load_program(out) + assert ( + torch.cuda.current_device() == caller_device + ), "loading moved the caller's device" + + method = program.load_method("forward") + assert torch.cuda.current_device() == caller_device, "preparing the method moved it" + + method.execute((torch.randn(64, 64, device=f"cuda:{engine_device}"),)) + assert ( + torch.cuda.current_device() == caller_device + ), "running moved the caller's device" + + del method + del program + gc.collect() + assert ( + torch.cuda.current_device() == caller_device + ), "freeing moved the caller's device" diff --git a/third_party/executorch/BUILD b/third_party/executorch/BUILD index 98d66cd02b..97b77d0d4f 100644 --- a/third_party/executorch/BUILD +++ b/third_party/executorch/BUILD @@ -60,19 +60,23 @@ cmake( # build the implementation once as a cc_binary(linkshared=True), re-import the # .so via cc_import, and expose it through a srcs-less cc_library that also # carries the headers (with include remap) and the cudart dep. Consumers depend -# on ":extension_cuda"; caller_stream.cpp is only ever linked into the one .so. +# on ":extension_cuda"; the sources are only ever linked into the one .so. # Private implementation. It carries hdrs + include_prefix/strip_include_prefix -# so caller_stream.cpp can find its own header +# so each source can find its own header # (the physical tree has an extra executorch/ dir; a cc_binary's srcs headers -# get no such remap). alwayslink=True guarantees the object -- and thus -# getCallerStream/CallerStreamGuard -- is pulled into the shared object below, -# even though that source-less cc_binary references none of its symbols. +# get no such remap). alwayslink=True guarantees the objects -- and thus the +# caller-stream and device-guard symbols -- are pulled into the shared object +# below, even though that source-less cc_binary references none of them. cc_library( name = "extension_cuda_impl", - srcs = ["executorch/extension/cuda/caller_stream.cpp"], + srcs = [ + "executorch/extension/cuda/caller_stream.cpp", + "executorch/extension/cuda/device_guard.cpp", + ], hdrs = [ "executorch/extension/cuda/caller_stream.h", + "executorch/extension/cuda/device_guard.h", "executorch/extension/cuda/export.h", ], local_defines = [ @@ -84,7 +88,15 @@ cc_library( alwayslink = True, linkstatic = True, visibility = ["//visibility:private"], - deps = ["@cuda//:cudart"], + # device_guard.cpp returns an Error and logs, so it needs the runtime's headers, + # which caller_stream.cpp beside it does not. Only the headers: the logging symbols + # themselves resolve at load time from the runtime the consumer already loads. + # Linking the static core here instead would give the process a second copy of + # every registry it carries. + deps = [ + ":executorch_headers", + "@cuda//:cudart", + ], ) # The single shared object. The explicit -soname pins DT_SONAME to the bare @@ -115,12 +127,18 @@ cc_library( name = "extension_cuda", hdrs = [ "executorch/extension/cuda/caller_stream.h", + "executorch/extension/cuda/device_guard.h", "executorch/extension/cuda/export.h", ], include_prefix = "executorch", strip_include_prefix = "executorch", deps = [ ":extension_cuda_shared", + # The device guard logs, so the shared object above leaves the runtime's + # logging symbols undefined. Listed here rather than by each consumer so the + # runtime always lands after that shared object on a link line, which is the + # only order in which a static archive can satisfy it. + ":executorch_core", "@cuda//:cudart", ], )