Skip to content

Commit aa6f591

Browse files
authored
[Offload] Implement hasPendingWork on CUDA (#152728)
Following on from #152304, implement the new query in the CUDA plugin
1 parent 187f296 commit aa6f591

File tree

1 file changed

+12
-2
lines changed
  • offload/plugins-nextgen/cuda/src

1 file changed

+12
-2
lines changed

offload/plugins-nextgen/cuda/src/rtl.cpp

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -914,9 +914,19 @@ struct CUDADeviceTy : public GenericDeviceTy {
914914
return Plugin::check(Res, "error in cuStreamWaitEvent: %s");
915915
}
916916

917-
// TODO: This should be implementable on CUDA
918917
Expected<bool> hasPendingWorkImpl(AsyncInfoWrapperTy &AsyncInfo) override {
919-
return true;
918+
CUstream Stream;
919+
if (auto Err = getStream(AsyncInfo, Stream))
920+
return Err;
921+
922+
CUresult Ret = cuStreamQuery(Stream);
923+
if (Ret == CUDA_SUCCESS)
924+
return false;
925+
926+
if (Ret == CUDA_ERROR_NOT_READY)
927+
return true;
928+
929+
return Plugin::check(Ret, "error in cuStreamQuery: %s");
920930
}
921931

922932
/// Synchronize the current thread with the event.

0 commit comments

Comments
 (0)