feat(platform): lazy initialization of devicecontext in pool (#14067)

* feat(platform): lazy initialization of devicecontext in pool

Use std::async(deferer, []{...}) to lazy initialize DeviceContext in Pool

test=develop

* Add future includes

test=develop

Author: reyoung

paddle/fluid/framework/parallel_executor.cc

@@ -303,10 +303,8 @@ void ParallelExecutor::FeedAndSplitTensorIntoLocalScopes(
 }
 
 ParallelExecutor::~ParallelExecutor() {
-  const auto dev_ctxs =
-      platform::DeviceContextPool::Instance().GetAllDeviceContexts();
-  for (auto &dev_ctx : dev_ctxs) {
-    dev_ctx->Wait();
+  for (auto &p : member_->places_) {
+    platform::DeviceContextPool::Instance().Get(p)->Wait();
   }
 
   if (member_->own_local_scope_) {

paddle/fluid/platform/device_context.cc

@@ -32,23 +32,25 @@ platform::DeviceContext* DeviceContextPool::Get(const platform::Place& place) {
         "'Place' is not supported, Please re-compile with WITH_GPU "
         "option");
   }
-  return it->second.get();
+  return it->second.get().get();
 }
 
-const std::vector<const DeviceContext*>
-DeviceContextPool::GetAllDeviceContexts() const {
-  std::vector<const DeviceContext*> all_device_ctx;
-  all_device_ctx.reserve(device_contexts_.size());
-  for (auto& dev_ctx : device_contexts_) {
-    all_device_ctx.emplace_back(dev_ctx.second.get());
-  }
-  return all_device_ctx;
+template <typename DevCtx, typename PlaceType>
+inline void EmplaceDeviceContext(
+    std::map<Place, std::shared_future<std::unique_ptr<DeviceContext>>>*
+        map_ptr,
+    platform::Place p) {
+  using PtrType = std::unique_ptr<DeviceContext>;
+  map_ptr->emplace(p, std::async(std::launch::deferred, [=] {
+                     // lazy evaluation. i.e., only create device context at
+                     // first `Get`
+                     return PtrType(new DevCtx(boost::get<PlaceType>(p)));
+                   }));
 }
 
 DeviceContextPool::DeviceContextPool(
     const std::vector<platform::Place>& places) {
   PADDLE_ENFORCE_GT(places.size(), 0);
-  using PtrType = std::unique_ptr<DeviceContext>;
   std::set<Place> set;
   for (auto& p : places) {
     set.insert(p);
@@ -57,26 +59,22 @@ DeviceContextPool::DeviceContextPool(
   for (auto& p : set) {
     if (platform::is_cpu_place(p)) {
 #ifdef PADDLE_WITH_MKLDNN
-      device_contexts_.emplace(
-          p, PtrType(new MKLDNNDeviceContext(boost::get<CPUPlace>(p))));
+      EmplaceDeviceContext<MKLDNNDeviceContext, CPUPlace>(&device_contexts_, p);
 #else
-      device_contexts_.emplace(
-          p, PtrType(new CPUDeviceContext(boost::get<CPUPlace>(p))));
+      EmplaceDeviceContext<CPUDeviceContext, CPUPlace>(&device_contexts_, p);
 #endif
     } else if (platform::is_gpu_place(p)) {
 #ifdef PADDLE_WITH_CUDA
-      device_contexts_.emplace(
-          p, PtrType(new CUDADeviceContext(boost::get<CUDAPlace>(p))));
+      EmplaceDeviceContext<CUDADeviceContext, CUDAPlace>(&device_contexts_, p);
 #else
       PADDLE_THROW(
           "'CUDAPlace' is not supported, Please re-compile with WITH_GPU "
           "option");
 #endif
     } else if (platform::is_cuda_pinned_place(p)) {
 #ifdef PADDLE_WITH_CUDA
-      device_contexts_.emplace(
-          p,
-          PtrType(new CUDAPinnedDeviceContext(boost::get<CUDAPinnedPlace>(p))));
+      EmplaceDeviceContext<CUDAPinnedDeviceContext, CUDAPinnedPlace>(
+          &device_contexts_, p);
 #else
       PADDLE_THROW(
           "'CUDAPlace' is not supported, Please re-compile with WITH_GPU "

paddle/fluid/platform/device_context.h

@@ -10,6 +10,7 @@ See the License for the specific language governing permissions and
 limitations under the License. */
 #pragma once
 
+#include <future>  // NOLINT
 #include <memory>
 #include <mutex>  // NOLINT
 #include <string>
@@ -223,9 +224,6 @@ class DeviceContextPool {
   /*! \brief  Return handle of single device context. */
   platform::DeviceContext* Get(const platform::Place& place);
 
-  /*! \brief  Return all the device contexts. */
-  const std::vector<const DeviceContext*> GetAllDeviceContexts() const;
-
   template <typename Place>
   const typename DefaultDeviceContextType<Place>::TYPE* GetByPlace(
       const Place& place) {
@@ -237,7 +235,8 @@ class DeviceContextPool {
 
  private:
   static DeviceContextPool* pool;
-  std::map<Place, std::unique_ptr<DeviceContext>> device_contexts_;
+  std::map<Place, std::shared_future<std::unique_ptr<DeviceContext>>>
+      device_contexts_;
   DISABLE_COPY_AND_ASSIGN(DeviceContextPool);
 };