[CustomDevice] fix stream_safe_custom_device_allocator (#63369)

ronny1996 · web-flow · commit 1b2ee65394c5 · 2024-04-10T15:37:25.000+08:00
diff --git a/paddle/fluid/memory/allocation/stream_safe_custom_device_allocator.cc b/paddle/fluid/memory/allocation/stream_safe_custom_device_allocator.cc
@@ -50,35 +50,11 @@ void StreamSafeCustomDeviceAllocation::RecordStream(
     outstanding_event_map_[stream]->Init(place());
     VLOG(9) << "Create a new event "
             << outstanding_event_map_[stream]->raw_event();
-    auto stream_wrapper = phi::stream::Stream(place(), stream);
-    VLOG(8) << "Record event " << outstanding_event_map_[stream]->raw_event()
-            << " to stream " << stream;
-    outstanding_event_map_[stream]->Record(&stream_wrapper);
-  }
-}
-
-void StreamSafeCustomDeviceAllocation::MarkAsWillBeFreed() {
-  std::lock_guard<SpinLock> lock_guard(outstanding_event_map_lock_);
-  if (!will_be_freed_) {
-    will_be_freed_ = false;
-    VLOG(8) << "ptr: " << ptr() << " will be freed";
-    if (phi::DeviceManager::HasDeviceType(place_.GetDeviceType()) &&
-        outstanding_event_map_.find(owning_stream_) ==
-            outstanding_event_map_.end()) {
-      std::call_once(once_flag_,
-                     [this] { phi::DeviceManager::SetDevice(place_); });
-      outstanding_event_map_.insert(
-          {owning_stream_, std::make_shared<phi::event::Event>()});
-      outstanding_event_map_[owning_stream_]->Init(place_);
-      VLOG(9) << "Create a new event "
-              << outstanding_event_map_[owning_stream_]->raw_event();
-      auto stream_wrapper = phi::stream::Stream(place_, owning_stream_);
-      VLOG(8) << "Record event "
-              << outstanding_event_map_[owning_stream_]->raw_event()
-              << " to stream " << owning_stream_;
-      outstanding_event_map_[owning_stream_]->Record(&stream_wrapper);
-    }
   }
+  auto stream_wrapper = phi::stream::Stream(place(), stream);
+  VLOG(8) << "Record event " << outstanding_event_map_[stream]->raw_event()
+          << " to stream " << stream;
+  outstanding_event_map_[stream]->Record(&stream_wrapper);
 }
 
 bool StreamSafeCustomDeviceAllocation::CanBeFreed() {
@@ -190,7 +166,6 @@ void StreamSafeCustomDeviceAllocator::FreeImpl(phi::Allocation* allocation) {
                               phi::DeviceContextPool::Instance().Get(place_))
                               ->stream());
   }
-  stream_safe_cuda_allocation->MarkAsWillBeFreed();
   if (stream_safe_cuda_allocation->CanBeFreed()) {
     VLOG(9) << "Directly delete allocation";
     delete stream_safe_cuda_allocation;
diff --git a/paddle/fluid/memory/allocation/stream_safe_custom_device_allocator.h b/paddle/fluid/memory/allocation/stream_safe_custom_device_allocator.h
@@ -37,7 +37,6 @@ class StreamSafeCustomDeviceAllocation : public Allocation {
 
   void RecordStream(phi::stream::stream_t stream);
   bool CanBeFreed();
-  void MarkAsWillBeFreed();
   phi::stream::stream_t GetOwningStream() const;
   void SetOwningStream(phi::stream::stream_t s);
 
