diff --git a/paddle/fluid/platform/stream_callback_manager.h b/paddle/fluid/platform/stream_callback_manager.h index 6c984065aa5fa1a8875aebe84051ab396bc417ec..0e88a439cf6ca83e3d98725f58875adeeea86be0 100644 --- a/paddle/fluid/platform/stream_callback_manager.h +++ b/paddle/fluid/platform/stream_callback_manager.h @@ -24,8 +24,6 @@ namespace paddle { namespace platform { -using StreamCallback = std::function; - class StreamCallbackManager; struct StreamCallbackContext { @@ -35,7 +33,7 @@ struct StreamCallbackContext { : manager_(manager), callback_(callback) {} const StreamCallbackManager *manager_; // do not own - StreamCallback callback_; + std::function callback_; }; class StreamCallbackManager { @@ -45,16 +43,18 @@ class StreamCallbackManager { template inline void AddCallback(Callback &&callback) const { - AddCallbackWithStreamAndErrorInfo( - [=](cudaStream_t, cudaError_t) { callback(); }); - } - - template - inline void AddCallbackWithStreamAndErrorInfo(Callback &&callback) const { - auto *stream_callback_context = new StreamCallbackContext(this, callback); - PADDLE_ENFORCE(cudaStreamAddCallback( - stream_, StreamCallbackManager::StreamCallbackFunc, - stream_callback_context, 0)); + auto *stream_callback_context = + new StreamCallbackContext(this, std::forward(callback)); + PADDLE_ENFORCE( +#if CUDA_VERSION >= 10000 + cudaLaunchHostFunc(stream_, StreamCallbackManager::StreamCallbackFunc, + stream_callback_context) +#else + cudaStreamAddCallback(stream_, + StreamCallbackManager::StreamCallbackFunc, + stream_callback_context, 0) +#endif + ); // NOLINT } void Wait() const { thread_pool_.reset(new ThreadPool(1)); } @@ -63,17 +63,21 @@ class StreamCallbackManager { const cudaStream_t stream_; mutable std::unique_ptr thread_pool_; - // cudaStreamCallback cannot call CUDA API inside, so we have to use - // thread_pool here +// cudaStreamCallback cannot call CUDA API inside, so we have to use +// thread_pool here +#if CUDA_VERSION >= 10000 + static void CUDART_CB StreamCallbackFunc(void *user_data) +#else static void CUDART_CB StreamCallbackFunc(cudaStream_t stream, - cudaError_t status, - void *user_data) { + cudaError_t status, void *user_data) +#endif + { auto *callback_context_ptr = reinterpret_cast(user_data); callback_context_ptr->manager_->thread_pool_->enqueue([=]() { std::unique_ptr callback_context( callback_context_ptr); - callback_context->callback_(stream, status); + callback_context->callback_(); }); } };