diff --git a/src/arch/cuda/hybridAPI/hapi_impl.cpp b/src/arch/cuda/hybridAPI/hapi_impl.cpp index beeea38d24..309bd6e23b 100644 --- a/src/arch/cuda/hybridAPI/hapi_impl.cpp +++ b/src/arch/cuda/hybridAPI/hapi_impl.cpp @@ -1367,6 +1367,7 @@ void hapiPollEvents(void* param) { std::queue& queue = CpvAccess(hapi_event_queue); while (!queue.empty()) { hapiEvent hev = queue.front(); + queue.pop(); // TODO: investigate possible race condition with charm4py futures - temporarily resolved by popping here if (cudaEventQuery(hev.event) == cudaSuccess) { // invoke Charm++ callback if one was given hev.cb.send(hev.cb_msg); @@ -1376,7 +1377,6 @@ void hapiPollEvents(void* param) { hapiWorkRequestCleanup(hev.wr); } cudaEventDestroy(hev.event); - queue.pop(); CpvAccess(n_hapi_events)--; // inform QD that an event was processed diff --git a/src/ck-core/charm.h b/src/ck-core/charm.h index cee65d539a..b645027626 100644 --- a/src/ck-core/charm.h +++ b/src/ck-core/charm.h @@ -510,6 +510,11 @@ extern int CkDisableTracing(int epIdx); extern void CkEnableTracing(int epIdx); extern void CkCallWhenIdle(int epIdx, void* obj); + +#if CMK_CHARM4PY && CMK_CUDA +extern void CkHapiAddCallback(long stream, void (*cb)(void*, void*), int fid); +#endif + #ifdef __cplusplus } #endif diff --git a/src/ck-core/ck.C b/src/ck-core/ck.C index 33bfc82bcf..8e6aa5e470 100644 --- a/src/ck-core/ck.C +++ b/src/ck-core/ck.C @@ -2637,6 +2637,18 @@ void CkArrayExtSend_multi(int aid, int *idx, int ndims, int epIdx, int num_bufs, } } + +#if CMK_CUDA +#include "hapi.h" +void CkHapiAddCallback(long stream, void (*cb)(void*, void*), int fid) +{ + cudaStream_t stream_ptr = (cudaStream_t)stream; + CkCallback callback(cb, (void *) fid); + + hapiAddCallback(stream_ptr, callback, NULL); +} +#endif // CMK_CUDA + #endif //------------------- Message Watcher (record/replay) ----------------