[NewExecutor] implement async check gc (#73448)

huangjiyi · web-flow · commit ee8d360cf66f · 2025-06-20T15:09:41.000+08:00
diff --git a/paddle/common/flags.cc b/paddle/common/flags.cc
@@ -432,6 +432,24 @@ PHI_DEFINE_EXPORTED_bool(
     "Fast eager deletion mode. If enabled, memory would release "
     "immediately without waiting GPU kernel ends.");
 
+/**
+ * Memory related FLAG
+ * Name: FLAGS_enable_async_fast_gc
+ * Since Version: 3.1.0
+ * Value Range: bool, default=false
+ * Example:
+ * Note: Enable async fast garbage collection mode. If enabled, allocation will
+ *       be released asynchronously, which makes the garbage collection process
+ *       faster. This flag is valid when fast_eager_deletion_mode is enabled.
+ */
+PHI_DEFINE_EXPORTED_bool(
+    enable_async_fast_gc,
+    false,
+    "Enable async fast garbage collection mode. If enabled, allocation will "
+    "be released asynchronously, which make the garbage collection process "
+    "non-blocking. This flag is only valid when FLAGS_fast_eager_deletion_mode "
+    "is true.");
+
 /**
  * Memory related FLAG
  * Name: FLAGS_memory_fraction_of_eager_deletion
diff --git a/paddle/fluid/framework/new_executor/garbage_collector/async_fast_garbage_collector.cc b/paddle/fluid/framework/new_executor/garbage_collector/async_fast_garbage_collector.cc
@@ -0,0 +1,123 @@
+// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/fluid/framework/new_executor/garbage_collector/async_fast_garbage_collector.h"
+#include "paddle/phi/api/profiler/event_tracing.h"
+#include "paddle/phi/core/threadpool.h"
+
+namespace paddle {
+namespace framework {
+
+SingleThreadLockFreeWorker::SingleThreadLockFreeWorker(int capacity)
+    : capacity_(capacity), head_(0), tail_(0), running_(true) {
+  tasks_queue_.resize(capacity);
+  worker_ = std::thread([this]() { this->WorkerLoop(); });
+}
+
+void SingleThreadLockFreeWorker::AddTask(Task task) {
+  tasks_queue_[tail_] = task;
+  tail_++;
+  if (tail_ >= tasks_queue_.size()) {
+    tasks_queue_.resize(tasks_queue_.size() + capacity_);
+  }
+}
+
+void SingleThreadLockFreeWorker::Wait() {
+  running_ = false;
+  if (worker_.joinable()) worker_.join();
+}
+
+void SingleThreadLockFreeWorker::WorkerLoop() {
+  while (true) {
+    if (head_ < tail_) {
+      Task task = tasks_queue_[head_];
+      task();
+      head_++;
+    } else if (head_ == tail_ && running_) {
+      std::this_thread::yield();
+    } else {
+      break;
+    }
+  }
+}
+
+InterpreterCoreAsyncFastGarbageCollector::
+    InterpreterCoreAsyncFastGarbageCollector(int num_instructions) {
+  async_worker_ =
+      std::make_unique<SingleThreadLockFreeWorker>(num_instructions);
+}
+
+void FreeVariable(Variable* var) {
+  if (var == nullptr) {
+    return;
+  }
+
+  if (var->IsType<phi::DenseTensor>()) {
+    Garbage garbage = var->GetMutable<phi::DenseTensor>()->MoveMemoryHolder();
+  } else if (
+      var->IsType<
+          operators::reader::
+              OrderedMultiDeviceDenseTensorBlockingQueueHolder>()) {  // NOLINT
+    // TODO(xiongkun03) in old executor, this type of variable is not support
+    // eager deletion. so we just leave it here ?
+  } else if (var->IsType<phi::SelectedRows>()) {
+    Garbage garbage = var->GetMutable<phi::SelectedRows>()
+                          ->mutable_value()
+                          ->MoveMemoryHolder();
+    var->GetMutable<phi::SelectedRows>()->mutable_rows()->clear();
+  } else if (var->IsType<phi::TensorArray>()) {
+    auto* tensor_arr = var->GetMutable<phi::TensorArray>();
+    for (auto& t : *tensor_arr) {
+      Garbage garbage = t.MoveMemoryHolder();
+    }
+    tensor_arr->clear();
+  } else if (var->IsType<phi::SparseCooTensor>()) {
+    Garbage indices = var->GetMutable<phi::SparseCooTensor>()
+                          ->mutable_indices()
+                          ->MoveMemoryHolder();
+    Garbage values = var->GetMutable<phi::SparseCooTensor>()
+                         ->mutable_values()
+                         ->MoveMemoryHolder();
+  } else if (var->IsType<phi::SparseCsrTensor>()) {
+    Garbage cols = var->GetMutable<phi::SparseCsrTensor>()
+                       ->mutable_cols()
+                       ->MoveMemoryHolder();
+    Garbage crows = var->GetMutable<phi::SparseCsrTensor>()
+                        ->mutable_crows()
+                        ->MoveMemoryHolder();
+    Garbage values = var->GetMutable<phi::SparseCsrTensor>()
+                         ->mutable_values()
+                         ->MoveMemoryHolder();
+  } else if (var->IsType<std::vector<Scope*>>()) {
+    // NOTE(@xiongkun03) conditional_op / while_op will create a STEP_SCOPE
+    // refer to executor.cc to see what old garbage collector does.
+    // do nothing, because the sub scope will be deleted by sub-executor.
+  } else {
+    PADDLE_THROW(common::errors::Unimplemented(
+        "The variable(%s) is not supported in eager deletion.",
+        framework::ToTypeName(var->Type())));
+  }
+}
+
+void InterpreterCoreAsyncFastGarbageCollector::Add(
+    const std::vector<Variable*>& vars) {
+  async_worker_->AddTask([vars]() {
+    for (const auto& var : vars) {
+      FreeVariable(var);
+    }
+  });
+}
+
+}  // namespace framework
+}  // namespace paddle
diff --git a/paddle/fluid/framework/new_executor/garbage_collector/async_fast_garbage_collector.h b/paddle/fluid/framework/new_executor/garbage_collector/async_fast_garbage_collector.h
@@ -0,0 +1,54 @@
+// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+#pragma once
+
+#include "paddle/fluid/framework/new_executor/garbage_collector/garbage_collector.h"
+
+namespace paddle {
+namespace framework {
+
+class SingleThreadLockFreeWorker {
+ public:
+  using Task = std::function<void()>;
+
+  explicit SingleThreadLockFreeWorker(int capacity);
+
+  ~SingleThreadLockFreeWorker() { Wait(); }
+
+  void AddTask(Task task);
+
+  void Wait();
+
+ private:
+  void WorkerLoop();
+
+  const int capacity_;
+  std::thread worker_;
+  std::vector<Task> tasks_queue_;
+  std::atomic<int> head_;
+  std::atomic<int> tail_;
+  std::atomic<bool> running_;
+};
+
+class InterpreterCoreAsyncFastGarbageCollector {
+ public:
+  explicit InterpreterCoreAsyncFastGarbageCollector(int num_instructions);
+
+  void Add(const std::vector<Variable*>& vars);
+
+ private:
+  std::unique_ptr<SingleThreadLockFreeWorker> async_worker_;
+};
+}  // namespace framework
+}  // namespace paddle
diff --git a/paddle/fluid/framework/new_executor/pir_interpreter.cc b/paddle/fluid/framework/new_executor/pir_interpreter.cc
@@ -48,6 +48,7 @@
 #ifdef PADDLE_WITH_CUSTOM_DEVICE
 #include "paddle/fluid/framework/new_executor/instruction/custom_engine_instruction.h"
 #endif
+#include "paddle/fluid/framework/new_executor/garbage_collector/async_fast_garbage_collector.h"
 #include "paddle/fluid/framework/new_executor/instruction/builtin_combine_instruction.h"
 #include "paddle/fluid/framework/new_executor/instruction/control_flow/assert_instruction.h"
 #include "paddle/fluid/framework/new_executor/instruction/control_flow/has_elements_instruction.h"
@@ -90,6 +91,7 @@ COMMON_DECLARE_bool(enable_pir_in_executor_trace_run);
 COMMON_DECLARE_bool(enable_collect_shape);
 COMMON_DECLARE_int32(low_precision_op_list);
 COMMON_DECLARE_bool(pir_interpreter_record_stream_for_gc_cache);
+COMMON_DECLARE_bool(enable_async_fast_gc);
 
 #define CREATE_INSTR(instr_name)                                   \
   vec_instruction_base_.emplace_back(std::make_unique<instr_name>( \
@@ -142,6 +144,7 @@ PirInterpreter::PirInterpreter(const phi::Place& place,
       exception_notifier_(nullptr),
       completion_notifier_(nullptr),
       gc_(nullptr),
+      async_gc_{nullptr},
       last_live_ops_(),
       dependency_count_(nullptr),
       deps_(),
@@ -233,6 +236,7 @@ PirInterpreter::PirInterpreter(
       exception_notifier_(nullptr),
       completion_notifier_(nullptr),
       gc_(nullptr),
+      async_gc_{nullptr},
       last_live_ops_(),
       dependency_count_(nullptr),
       deps_(),
@@ -300,6 +304,7 @@ PirInterpreter::PirInterpreter(
 PirInterpreter::~PirInterpreter() {
   // cancel gc's thread
   gc_.reset(nullptr);
+  async_gc_.reset(nullptr);
   async_work_queue_.reset();
   VLOG(4) << "~PirInterpreter(): " << this << " on " << place_;
 
@@ -1283,6 +1288,7 @@ void PirInterpreter::CheckGC(InstructionBase* instr) {
   RecordStreamForGC(instr);
 #endif
 
+  std::vector<Variable*> gc_vars;
   for (auto var_id : instr->GCCheckVars()) {
     VLOG(4) << "GC:" << value_exe_info_->GetNameById(static_cast<int>(var_id))
             << ", id:" << var_id << ", ref:" << refs_[var_id]->DynamicRef();
@@ -1298,10 +1304,18 @@ void PirInterpreter::CheckGC(InstructionBase* instr) {
     if (is_ready) {
       VLOG(6) << "Async delete variable with name : "
               << value_exe_info_->GetNameById(static_cast<int>(var_id));
-      gc_->Add(refs_[var_id]->Var(), instr);
+      if (use_trace_run_ && FLAGS_enable_async_fast_gc) {
+        gc_vars.push_back(refs_[var_id]->Var());
+      } else {
+        gc_->Add(refs_[var_id]->Var(), instr);
+      }
     }
   }
 
+  if (use_trace_run_ && FLAGS_enable_async_fast_gc) {
+    async_gc_->Add(gc_vars);
+  }
+
   for (auto var : instr->EagerGCVars()) {
     gc_->Add(var, instr);
   }
@@ -1501,7 +1515,7 @@ paddle::framework::FetchList PirInterpreter::Run(
     PreAnalysis();
     VLOG(4) << "Done PreAnalysis";
 
-    if (UseTraceRun(execution_config_, onednn_op_num_, sync_op_num_)) {
+    if (use_trace_run_) {
       LOG_FIRST_N(INFO, 1) << "pir interpreter is running by trace mode ...";
       TraceRunImpl();
     } else {
@@ -1513,7 +1527,7 @@ paddle::framework::FetchList PirInterpreter::Run(
     is_build_ = true;
     is_shared_results_build_ = true;
   } else {
-    if (UseTraceRun(execution_config_, onednn_op_num_, sync_op_num_)) {
+    if (use_trace_run_) {
       TraceRunImpl();
     } else {
       MultiThreadRunImpl();
@@ -1579,7 +1593,7 @@ FetchList PirInterpreter::Run(const std::vector<std::string>& feed_names,
     VLOG(4) << "Done PreAnalysis";
 
     // Run
-    if (UseTraceRun(execution_config_, onednn_op_num_, sync_op_num_)) {
+    if (use_trace_run_) {
       LOG_FIRST_N(INFO, 1) << "pir interpreter is running by trace mode ...";
       TraceRunImpl();
     } else {
@@ -1591,7 +1605,7 @@ FetchList PirInterpreter::Run(const std::vector<std::string>& feed_names,
     is_build_ = true;
     is_shared_results_build_ = true;
   } else {
-    if (UseTraceRun(execution_config_, onednn_op_num_, sync_op_num_)) {
+    if (use_trace_run_) {
       TraceRunImpl();
     } else {
       MultiThreadRunImpl();
@@ -1624,6 +1638,11 @@ void PirInterpreter::TraceRunImpl() {
     gc_ = CreateInterpreterCoreGarbageCollector(place_, vec_instruction_base_);
   }
 
+  if (FLAGS_enable_async_fast_gc) {
+    async_gc_ = std::make_unique<InterpreterCoreAsyncFastGarbageCollector>(
+        vec_instruction_base_.size());
+  }
+
   interpreter::ResetAtomicGuard guard(&deps_, &refs_);
   VLOG(4) << "Tracing Instruction List";
 
@@ -2036,6 +2055,8 @@ void PirInterpreter::PreAnalysis() {
 
   UpdateOneDNNOpNum();
   VLOG(4) << "Done UpdateOneDNNOpNum";
+
+  use_trace_run_ = UseTraceRun(execution_config_, onednn_op_num_, sync_op_num_);
 }
 
 ::pir::Value PirInterpreter::GetValueByName(const std::string& var_name) {
diff --git a/paddle/fluid/framework/new_executor/pir_interpreter.h b/paddle/fluid/framework/new_executor/pir_interpreter.h
@@ -29,6 +29,7 @@ class Block;
 namespace paddle {
 namespace framework {
 class ValueExecutionInfo;
+class InterpreterCoreAsyncFastGarbageCollector;
 class PirInterpreter : public InterpreterBaseImpl {
   using ExecutionConfig = interpreter::ExecutionConfig;
   using InstructionSchedulingPriorityLess = std::function<bool(size_t, size_t)>;
@@ -188,6 +189,7 @@ class PirInterpreter : public InterpreterBaseImpl {
   std::shared_ptr<EventsWaiter::EventNotifier> completion_notifier_{nullptr};
 
   std::unique_ptr<InterpreterCoreGarbageCollector> gc_;
+  std::unique_ptr<InterpreterCoreAsyncFastGarbageCollector> async_gc_;
 
   // last_live_ops_[i] contains the id of operators that last access the i-th
   // var
@@ -201,6 +203,7 @@ class PirInterpreter : public InterpreterBaseImpl {
   std::vector<std::shared_ptr<interpreter::VarRefInfo>> refs_;
 
   // used for Trace
+  bool use_trace_run_{false};
   int64_t sync_op_num_{-1};
   int64_t nccl_op_num_{-1};
   int64_t onednn_op_num_{-1};
