[Offload] Implement hasPendingWork on CUDA

[callumfare]

Following on from #152304, implement the new query in the CUDA plugin

[llvmbot]

@llvm/pr-subscribers-offload

Author: Callum Fare (callumfare)

Changes

Following on from #152304, implement the new query in the CUDA plugin

---

Full diff: https://github.com/llvm/llvm-project/pull/152728.diff

1 Files Affected:

• (modified) offload/plugins-nextgen/cuda/src/rtl.cpp (+12-2)

diff --git a/offload/plugins-nextgen/cuda/src/rtl.cpp b/offload/plugins-nextgen/cuda/src/rtl.cpp
index f3f3783b3ce7c..6ae17c4265894 100644
--- a/offload/plugins-nextgen/cuda/src/rtl.cpp
+++ b/offload/plugins-nextgen/cuda/src/rtl.cpp
@@ -914,9 +914,19 @@ struct CUDADeviceTy : public GenericDeviceTy {
     return Plugin::check(Res, "error in cuStreamWaitEvent: %s");
   }
 
-  // TODO: This should be implementable on CUDA
   Expected<bool> hasPendingWorkImpl(AsyncInfoWrapperTy &AsyncInfo) override {
-    return true;
+    CUstream Stream;
+    if (auto Err = getStream(AsyncInfo, Stream))
+      return Err;
+
+    CUresult Ret = cuStreamQuery(Stream);
+    if (Ret == CUDA_SUCCESS)
+      return false;
+
+    if (Ret == CUDA_ERROR_NOT_READY)
+      return true;
+
+    return Plugin::check(Ret, "error in cuStreamQuery: %s");
   }
 
   /// Synchronize the current thread with the event.