From 93ce7f6952fa0614afe32bfb14e3721749d6fc26 Mon Sep 17 00:00:00 2001
From: Leo Chen <chenqiuliang@baidu.com>
Date: Mon, 11 Jan 2021 19:40:05 +0800
Subject: [PATCH]  [cherry-pick] Async drop scope in executor (#29714) #30285

[cherry-pick] Async drop scope in executor (#29714)
---
 paddle/fluid/framework/executor.cc         | 40 +++++++++++++++-------
 paddle/fluid/framework/garbage_collector.h |  4 +++
 2 files changed, 31 insertions(+), 13 deletions(-)

diff --git a/paddle/fluid/framework/executor.cc b/paddle/fluid/framework/executor.cc
index bce4cb018d..755b3bff76 100644
--- a/paddle/fluid/framework/executor.cc
+++ b/paddle/fluid/framework/executor.cc
@@ -478,21 +478,35 @@ void Executor::RunPartialPreparedContext(ExecutorPrepareContext* ctx,
     }
   }
 
-  platform::DeviceContextPool::Instance().Get(place_)->Wait();
+  auto callback = [scope, local_scope, keep_kids]() {
+    if (local_scope != scope) {
+      VLOG(4) << "Delete scope: " << local_scope;
+      scope->DeleteScope(local_scope);
+    } else {
+      if (!keep_kids) {
+        VLOG(4) << "Drop kids: " << scope;
+        // By default, we should delete all kid scopes after run executor
+        // because
+        // some operators may create local scope when running, such as while_op.
+        // But when while_op also create a local executor to run it's sub block,
+        // the sub scopes it created should not be dropped immediately, because
+        // while_grad_op will use some variables created during while_op run, so
+        // we need to keep the kids and wait for the outer executor to drop
+        // them.
+
+        scope->DropKids();
+      }
+      VLOG(4) << "Keep kids: " << scope;
+    }
+  };
 
-  if (local_scope != scope) {
-    scope->DeleteScope(local_scope);
+  if (gc) {
+    VLOG(4) << "Async deleting scope";
+    gc->DirectClearCallback(callback);
   } else {
-    if (!keep_kids) {
-      // By default, we should delete all kid scopes after run executor because
-      // some operators may create local scope when running, such as while_op.
-      // But when while_op also create a local executor to run it's sub block,
-      // the sub scopes it created should not be dropped immediately, because
-      // while_grad_op will use some variables created during while_op run, so
-      // we need to keep the kids and wait for the outer executor to drop them.
-
-      scope->DropKids();
-    }
+    VLOG(4) << "Sync deleting scope";
+    platform::DeviceContextPool::Instance().Get(place_)->Wait();
+    callback();
   }
 }
 
diff --git a/paddle/fluid/framework/garbage_collector.h b/paddle/fluid/framework/garbage_collector.h
index 884d230816..0b5fdc4745 100644
--- a/paddle/fluid/framework/garbage_collector.h
+++ b/paddle/fluid/framework/garbage_collector.h
@@ -48,6 +48,10 @@ class GarbageCollector {
   template <typename Container, typename Callback>
   void Add(Container &&objs, Callback &&callback);
 
+  void DirectClearCallback(const std::function<void()> &callback) {
+    ClearCallback(callback);
+  }
+
  protected:
   virtual void ClearCallback(const std::function<void()> &callback) = 0;
 
-- 
GitLab