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", ], )