diff --git a/src/utils/device.cpp b/src/utils/device.cpp index 6e447ba6..f217b213 100644 --- a/src/utils/device.cpp +++ b/src/utils/device.cpp @@ -37,18 +37,60 @@ WithDevice::~WithDevice() { cudaCheckErrorNoThrow(cudaSetDevice(original_device_id_)); } +namespace { + +//! Returns true if `stream` belongs to the current device. +//! cudaStreamQuery accepts streams from any device in the process, so it cannot answer this on its +//! own. +bool streamIsOnCurrentDevice(cudaStream_t stream) { +#if CUDART_VERSION >= 12080 + int currentDevice = -1; + if (cudaGetDevice(¤tDevice) != cudaSuccess) { + cudaGetLastError(); + return false; + } + int streamDevice = -1; + if (cudaStreamGetDevice(stream, &streamDevice) != cudaSuccess) { + cudaGetLastError(); + return false; + } + return streamDevice == currentDevice; +#else + // cudaStreamGetDevice needs CUDA 12.8. Below that, probe with an event instead: an event and the + // stream it is recorded into must share a CUDA context, so recording an event created on the + // current device fails with cudaErrorInvalidResourceHandle if the stream belongs to another one. + cudaEvent_t probe = nullptr; + if (cudaEventCreateWithFlags(&probe, cudaEventDisableTiming) != cudaSuccess) { + cudaGetLastError(); + return false; + } + const cudaError_t recordErr = cudaEventRecord(probe, stream); + cudaCheckErrorNoThrow(cudaEventDestroy(probe)); + if (recordErr != cudaSuccess) { + cudaGetLastError(); + return false; + } + return true; +#endif +} + +} // namespace + std::optional acquireExternalStream(std::uintptr_t streamPtr) { auto stream = reinterpret_cast(streamPtr); if (streamPtr == 0) { return stream; } cudaError_t err = cudaStreamQuery(stream); - if (err == cudaSuccess || err == cudaErrorNotReady) { - return stream; + if (err != cudaSuccess && err != cudaErrorNotReady) { + // Clear the sticky error state + cudaGetLastError(); + return std::nullopt; + } + if (!streamIsOnCurrentDevice(stream)) { + return std::nullopt; } - // Clear the sticky error state - cudaGetLastError(); - return std::nullopt; + return stream; } size_t getDeviceFreeMemory() { diff --git a/tests/test_fmcs_primitives.cu b/tests/test_fmcs_primitives.cu index 324371c8..bb4338f3 100644 --- a/tests/test_fmcs_primitives.cu +++ b/tests/test_fmcs_primitives.cu @@ -13,9 +13,9 @@ #include #include -#include "fmcs_cuda/fmcs_seed.cuh" -#include "fmcs_cuda/fmcs_seed_queue.cuh" -#include "mcs_common/mcs_cooperative_copy.cuh" +#include "src/mcs/fmcs_cuda/fmcs_seed.cuh" +#include "src/mcs/fmcs_cuda/fmcs_seed_queue.cuh" +#include "src/mcs/mcs_common/mcs_cooperative_copy.cuh" #include "src/utils/device_vector.h" namespace {