[CustomDevice] Fix blocking cpu in ~XCCLTask (#58038)

ronny1996 · web-flow · commit ccdf25df0361 · 2023-10-13T10:17:17.000+08:00
diff --git a/paddle/fluid/distributed/collective/process_group_custom.cc b/paddle/fluid/distributed/collective/process_group_custom.cc
@@ -32,23 +32,57 @@ PD_DECLARE_bool(use_stream_safe_cuda_allocator);
 namespace paddle {
 namespace distributed {
 
+static std::mutex g_unfinished_xccl_task_events_mutex;
+static std::list<std::unique_ptr<phi::event::Event>>
+    g_unfinished_xccl_task_events;
+
 ProcessGroupCustom::XCCLTask::XCCLTask(const Place& place,
                                        int rank,
                                        CommType comm_type,
                                        bool sync_op,
                                        bool use_calc_stream)
     : TaskStream(rank, comm_type, sync_op, use_calc_stream),
-      task_place_(place) {
-  comm_event_.Init(place);
+      task_place_(place),
+      comm_event_(std::make_unique<phi::event::Event>()) {
+  comm_event_->Init(task_place_);
+}
+
+ProcessGroupCustom::XCCLTask::XCCLTask(
+    const std::vector<Place>& places,
+    int rank,
+    CommType CommType,
+    const std::vector<phi::DenseTensor>& inputs)
+    : TaskStream(rank, inputs, CommType),
+      task_place_(places[0]),
+      comm_event_(std::make_unique<phi::event::Event>()) {
+  comm_event_->Init(task_place_);
 }
 
-ProcessGroupCustom::XCCLTask::~XCCLTask() = default;
+ProcessGroupCustom::XCCLTask::~XCCLTask() {
+  if (!IsCompleted()) {
+    std::lock_guard<std::mutex> lock(g_unfinished_xccl_task_events_mutex);
+    g_unfinished_xccl_task_events.push_back(std::move(comm_event_));
+  }
+}
 
-bool ProcessGroupCustom::XCCLTask::IsCompleted() { return comm_event_.Query(); }
+bool ProcessGroupCustom::XCCLTask::IsCompleted() {
+  return comm_event_->Query();
+}
 
 void ProcessGroupCustom::XCCLTask::UpdateWaitChain(
     const phi::DeviceContext& ctx) {
-  comm_event_.Record(
+  {
+    std::lock_guard<std::mutex> lock(g_unfinished_xccl_task_events_mutex);
+    for (auto iter = g_unfinished_xccl_task_events.begin();
+         iter != g_unfinished_xccl_task_events.end();) {
+      if ((*iter)->Query()) {
+        iter = g_unfinished_xccl_task_events.erase(iter);
+      } else {
+        iter++;
+      }
+    }
+  }
+  comm_event_->Record(
       reinterpret_cast<const phi::CustomContext&>(ctx).GetStream().get());
 }
 
@@ -62,7 +96,7 @@ bool ProcessGroupCustom::XCCLTask::Wait(std::chrono::milliseconds timeout) {
 
   const auto* calc_ctx = reinterpret_cast<phi::CustomContext*>(
       platform::DeviceContextPool::Instance().Get(task_place_));
-  calc_ctx->GetStream()->WaitEvent(&comm_event_);
+  calc_ctx->GetStream()->WaitEvent(comm_event_.get());
 
   if (IsBlockCPUInWait()) {
     // If we use the work to do barrier, we should block cpu
@@ -590,15 +624,6 @@ std::shared_ptr<ProcessGroupCustom::XCCLTask> ProcessGroupCustom::CreateTask(
       places, rank, comm_type, inputs);
 }
 
-ProcessGroupCustom::XCCLTask::XCCLTask(
-    const std::vector<Place>& places,
-    int rank,
-    CommType CommType,
-    const std::vector<phi::DenseTensor>& inputs)
-    : TaskStream(rank, inputs, CommType), task_place_(places[0]) {
-  comm_event_.Init(places[0]);
-}
-
 // create XCCLManager cache for places_key
 void ProcessGroupCustom::CreateXCCLManagerCache(
     const std::string& places_key, const std::vector<Place>& places) {
@@ -676,7 +701,7 @@ std::shared_ptr<ProcessGroup::Task> ProcessGroupCustom::Collective(
   {
     std::lock_guard<std::mutex> lock(mutex_);
     if (place_to_comm_ctx_.find(key) == place_to_comm_ctx_.end()) {
-      CreateXCCLManagerCache(key, places);
+      CreateXCCLEnvCache(places[0], key);
     }
   }
 
diff --git a/paddle/fluid/distributed/collective/process_group_custom.h b/paddle/fluid/distributed/collective/process_group_custom.h
@@ -61,8 +61,8 @@ class ProcessGroupCustom final : public ProcessGroupWithStream {
 
    private:
     bool block_cpu_in_wait_{false};
-    phi::event::Event comm_event_;  // event on comm stream
     Place task_place_;
+    std::unique_ptr<phi::event::Event> comm_event_;  // event on comm stream
   };
 
  public:
diff --git a/paddle/phi/backends/event.cc b/paddle/phi/backends/event.cc
@@ -46,6 +46,7 @@ Event::Event(const Place& place, event_t event)
       own_data_(false) {}
 
 Event::~Event() {
+  Synchronize();
   Destroy();
   std::unique_lock lock(g_events_mutex);
   g_events.remove(this);
@@ -77,14 +78,35 @@ void Event::Destroy() {
     own_data_ = false;
     event_ = nullptr;
     device_ = nullptr;
+    is_recorded_ = false;
   }
 }
 
-void Event::Record(const stream::Stream* stream) { stream->RecordEvent(this); }
+void Event::Record(const stream::Stream* stream) {
+  if (device_) {
+    is_recorded_ = true;  // synchronize the event during detroy
+    stream->RecordEvent(this);
+  }
+}
 
-bool Event::Query() const { return device_->QueryEvent(this); }
+bool Event::Query() const {
+  if (device_ && is_recorded_) {
+    bool ret = device_->QueryEvent(this);
+    if (ret) {
+      is_recorded_ =
+          false;  // event completed, do not need to synchronize the event.
+    }
+    return ret;
+  } else {
+    return true;
+  }
+}
 
-void Event::Synchronize() const { device_->SynchronizeEvent(this); }
+void Event::Synchronize() const {
+  if (device_ && is_recorded_) {
+    device_->SynchronizeEvent(this);
+  }
+}
 
 const Place& Event::GetPlace() const { return place_; }
 
diff --git a/paddle/phi/backends/event.h b/paddle/phi/backends/event.h
@@ -57,6 +57,7 @@ class Event {
   Device* device_;
   event_t event_;
   bool own_data_ = true;
+  mutable bool is_recorded_ = false;
 };
 }  // namespace event
 
