From 1bdb5cee24039e86b784d8499e9513a36eb6e047 Mon Sep 17 00:00:00 2001 From: Jacob Faibussowitsch Date: Fri, 17 Jul 2026 10:52:00 -0400 Subject: [PATCH 1/2] Push the streams context onto the driver stack when calling cumemcpy --- libcudacxx/include/cuda/__driver/driver_api.h | 41 +++++++++++++++++++ 1 file changed, 41 insertions(+) diff --git a/libcudacxx/include/cuda/__driver/driver_api.h b/libcudacxx/include/cuda/__driver/driver_api.h index ed726a46b8a..22f6be865ad 100644 --- a/libcudacxx/include/cuda/__driver/driver_api.h +++ b/libcudacxx/include/cuda/__driver/driver_api.h @@ -339,10 +339,51 @@ _CCCL_HOST_API inline ::CUcontext __ctxPop() // Memory management +[[nodiscard]] _CCCL_HOST_API inline ::CUcontext __streamGetCtx(::CUstream __stream); + +// The driver memcpy entry points validate their device pointer arguments against the *current* +// context, NOT against the context the supplied stream belongs to. In a multi-device program +// the current context is frequently some other device's context (e.g. nothing was pushed, or a +// previous operation left a different device current), in which case a perfectly valid device +// pointer is rejected with `CUDA_ERROR_INVALID_VALUE` ("invalid argument"). Pushing the +// stream's own context around the call guarantees the device pointer is validated against the +// context that actually owns it. +// +// The null/default stream has no retrievable context; for it we leave the current context +// untouched, matching the default-stream's documented "current context" semantics. +struct [[maybe_unused]] __stream_ctx_guard +{ + bool __pushed_ = false; + + _CCCL_HOST_API explicit __stream_ctx_guard(::CUstream __stream) + { + if (__stream != nullptr) + { + ::cuda::__driver::__ctxPush(::cuda::__driver::__streamGetCtx(__stream)); + __pushed_ = true; + } + } + + __stream_ctx_guard(__stream_ctx_guard&&) = delete; + __stream_ctx_guard(const __stream_ctx_guard&) = delete; + __stream_ctx_guard& operator=(__stream_ctx_guard&&) = delete; + __stream_ctx_guard& operator=(const __stream_ctx_guard&) = delete; + + _CCCL_HOST_API ~__stream_ctx_guard() // NOLINT(bugprone-exception-escape) + { + if (__pushed_) + { + static_cast(::cuda::__driver::__ctxPop()); + } + } +}; + _CCCL_HOST_API inline void __memcpyAsync(void* __dst, const void* __src, ::cuda::std::size_t __count, ::CUstream __stream) { static auto __driver_fn = _CCCLRT_GET_DRIVER_FUNCTION(cuMemcpyAsync); + const auto __guard = ::cuda::__driver::__stream_ctx_guard{__stream}; + ::cuda::__driver::__call_driver_fn( __driver_fn, "Failed to perform a memcpy", From 155a9ed35c765dd1f8bef0fb97428789def7e6d6 Mon Sep 17 00:00:00 2001 From: Jacob Faibussowitsch Date: Fri, 17 Jul 2026 11:14:28 -0400 Subject: [PATCH 2/2] fixup! Push the streams context onto the driver stack when calling cumemcpy --- .../cuda/ccclrt/utility/driver_api.c2h.cu | 38 +++++++++++++++++++ 1 file changed, 38 insertions(+) diff --git a/libcudacxx/test/libcudacxx/cuda/ccclrt/utility/driver_api.c2h.cu b/libcudacxx/test/libcudacxx/cuda/ccclrt/utility/driver_api.c2h.cu index 520befaa13c..b2d70666175 100644 --- a/libcudacxx/test/libcudacxx/cuda/ccclrt/utility/driver_api.c2h.cu +++ b/libcudacxx/test/libcudacxx/cuda/ccclrt/utility/driver_api.c2h.cu @@ -9,6 +9,8 @@ //===----------------------------------------------------------------------===// #include +#include +#include #include @@ -74,3 +76,39 @@ C2H_TEST("Call each driver api", "[utility]") CUDART(driver::__streamDestroyNoThrow(stream)); } + +C2H_CCCLRT_TEST("memcpy async uses the stream context", "[utility][multi_gpu]") +{ + namespace driver = ::cuda::__driver; + + if (cuda::devices.size() < 2) + { + SKIP("Need at least 2 devices"); + } + + cuda::device_ref current_device{0}; + cuda::device_ref stream_device{1}; + const auto initial_stack_depth = test::count_driver_stack(); + + { + const auto stream = cuda::stream{stream_device}; + auto src = cuda::make_device_buffer(stream, stream_device, 1, cuda::no_init); + auto dst = cuda::make_device_buffer(stream, stream_device, 1, cuda::no_init); + + { + const auto guard = cuda::__ensure_current_context{current_device}; + const auto current_context = driver::__ctxGetCurrent(); + const auto stack_depth = test::count_driver_stack(); + + REQUIRE(current_context != driver::__streamGetCtx(stream.get())); + + driver::__memcpyAsync(dst.data(), src.data(), sizeof(*src.data()), stream.get()); + + REQUIRE(driver::__ctxGetCurrent() == current_context); + REQUIRE(test::count_driver_stack() == stack_depth); + } + + stream.sync(); + } + REQUIRE(test::count_driver_stack() == initial_stack_depth); +}