diff --git a/hipnv/include/hip/nvidia_detail/nvidia_hip_runtime_api.h b/hipnv/include/hip/nvidia_detail/nvidia_hip_runtime_api.h index d770182632..4c8be9af3d 100644 --- a/hipnv/include/hip/nvidia_detail/nvidia_hip_runtime_api.h +++ b/hipnv/include/hip/nvidia_detail/nvidia_hip_runtime_api.h @@ -2507,6 +2507,20 @@ inline static hipError_t hipStreamAddCallback(hipStream_t stream, hipStreamCallb cudaStreamAddCallback(stream, (cudaStreamCallback_t)callback, userData, flags)); } +inline static hipError_t hipStreamGetDevice(hipStream_t stream, hipDevice_t* device) { + hipCtx_t context; + auto err = hipCUResultTohipError(cuStreamGetCtx(stream, &context)); + if (err != hipSuccess) return err; + + err = hipCUResultTohipError(cuCtxPushCurrent(context)); + if (err != hipSuccess) return err; + + err = hipCUResultTohipError(cuCtxGetDevice(device)); + if (err != hipSuccess) return err; + + return hipCUResultTohipError(cuCtxPopCurrent(&context)); +} + inline static hipError_t hipDriverGetVersion(int* driverVersion) { return hipCUDAErrorTohipError(cudaDriverGetVersion(driverVersion)); }