[OpenMP] Reset async stream properly upon failure
Summary:
If the call to `synchronize` fails, it will currently block the stream indefinitely if execution is continued from this point. Additionally, if the program exits it will trigger an assertion on the non-null value of the async queue and prevent the runtime from printing debugging information.
Reviewers: jdoerfert
Differential Revision: https://reviews.llvm.org/D99443
GitOrigin-RevId: 16064e71e934547ca68b3d8f4bd71e9a2b7f4248
diff --git a/libomptarget/plugins/cuda/src/rtl.cpp b/libomptarget/plugins/cuda/src/rtl.cpp
index 3d0424f..a2f3cf4 100644
--- a/libomptarget/plugins/cuda/src/rtl.cpp
+++ b/libomptarget/plugins/cuda/src/rtl.cpp
@@ -1035,13 +1035,6 @@
int synchronize(const int DeviceId, __tgt_async_info *AsyncInfo) const {
CUstream Stream = reinterpret_cast<CUstream>(AsyncInfo->Queue);
CUresult Err = cuStreamSynchronize(Stream);
- if (Err != CUDA_SUCCESS) {
- REPORT("Error when synchronizing stream. stream = " DPxMOD
- ", async info ptr = " DPxMOD "\n",
- DPxPTR(Stream), DPxPTR(AsyncInfo));
- CUDA_ERR_STRING(Err);
- return OFFLOAD_FAIL;
- }
// Once the stream is synchronized, return it to stream pool and reset
// AsyncInfo. This is to make sure the synchronization only works for its
@@ -1050,7 +1043,13 @@
reinterpret_cast<CUstream>(AsyncInfo->Queue));
AsyncInfo->Queue = nullptr;
- return OFFLOAD_SUCCESS;
+ if (Err != CUDA_SUCCESS) {
+ REPORT("Error when synchronizing stream. stream = " DPxMOD
+ ", async info ptr = " DPxMOD "\n",
+ DPxPTR(Stream), DPxPTR(AsyncInfo));
+ CUDA_ERR_STRING(Err);
+ }
+ return (Err == CUDA_SUCCESS) ? OFFLOAD_SUCCESS : OFFLOAD_FAIL;
}
};