From d3a43477fbb391e93299f2f4c11c05b705635ae4 Mon Sep 17 00:00:00 2001
From: ronnywang <524019753@qq.com>
Date: Thu, 24 Mar 2022 13:08:39 +0800
Subject: [PATCH] [custom runtime] clear headers (#40845)

---
 paddle/phi/backends/callback_manager.cc | 6 ++++--
 paddle/phi/backends/callback_manager.h  | 6 +++---
 paddle/phi/backends/device_base.h       | 2 ++
 3 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/paddle/phi/backends/callback_manager.cc b/paddle/phi/backends/callback_manager.cc
index e21e8502d8..4a958ef73b 100644
--- a/paddle/phi/backends/callback_manager.cc
+++ b/paddle/phi/backends/callback_manager.cc
@@ -16,16 +16,18 @@
 #include "paddle/fluid/platform/device/device_wrapper.h"
 #include "paddle/fluid/platform/enforce.h"
 
+#include <ThreadPool.h>
+
 namespace phi {
 
 CallbackManager::CallbackManager(stream::Stream *stream)
-    : stream_(stream), thread_pool_(1) {}
+    : stream_(stream), thread_pool_(new ::ThreadPool(1)) {}
 
 void CallbackManager::AddCallback(std::function<void()> callback) const {
   auto *callback_func = new std::function<void()>(std::move(callback));
   auto *func = new std::function<void()>([this, callback_func] {
     std::lock_guard<std::mutex> lock(mtx_);
-    last_future_ = thread_pool_.enqueue([callback_func] {
+    last_future_ = thread_pool_->enqueue([callback_func] {
       std::unique_ptr<std::function<void()>> releaser(callback_func);
       (*callback_func)();
     });
diff --git a/paddle/phi/backends/callback_manager.h b/paddle/phi/backends/callback_manager.h
index 359958b7c9..2bb2674528 100644
--- a/paddle/phi/backends/callback_manager.h
+++ b/paddle/phi/backends/callback_manager.h
@@ -14,8 +14,6 @@
 
 #pragma once
 
-#include <ThreadPool.h>
-
 #ifdef PADDLE_WITH_CUDA
 #include <cuda.h>
 #include <cuda_runtime.h>
@@ -30,6 +28,8 @@
 #include <memory>
 #include <mutex>  // NOLINT
 
+class ThreadPool;
+
 namespace phi {
 
 namespace stream {
@@ -50,7 +50,7 @@ class CallbackManager {
 
  private:
   stream::Stream* stream_;
-  mutable ::ThreadPool thread_pool_;
+  mutable std::shared_ptr<::ThreadPool> thread_pool_;
   mutable std::mutex mtx_;
   mutable std::future<void> last_future_;
 };
diff --git a/paddle/phi/backends/device_base.h b/paddle/phi/backends/device_base.h
index b4964708df..8cc6e49806 100644
--- a/paddle/phi/backends/device_base.h
+++ b/paddle/phi/backends/device_base.h
@@ -14,6 +14,8 @@
 
 #pragma once
 #ifdef PADDLE_WITH_CUSTOM_DEVICE
+#include <vector>
+
 #include "paddle/phi/backends/event.h"
 #include "paddle/phi/backends/stream.h"
 
-- 
GitLab