Push the streams context onto the driver stack when calling cumemcpy

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<void>(::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",

libcudacxx/test/libcudacxx/cuda/ccclrt/utility/driver_api.c2h.cu

@@ -9,6 +9,8 @@
 //===----------------------------------------------------------------------===//
 
 #include <cuda/__driver/driver_api.h>
+#include <cuda/devices>
+#include <cuda/memory_pool>
 
 #include <testing.cuh>
 
@@ -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<int>(stream, stream_device, 1, cuda::no_init);
+    auto dst          = cuda::make_device_buffer<int>(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);
+}