add compact func

Author: liuruyan

paddle/fluid/pybind/pybind.cc

@@ -3587,6 +3587,7 @@ All parameter, weight, gradient are variables in Paddle.
     }
     platform::EmptyCache();
   });
+  m.def("vmm_compact", [] { platform::VmmCompact(); });
   m.def(
       "get_device_properties",
       [](int id) -> const gpuDeviceProp & {

paddle/phi/core/memory/allocation/allocator.h

@@ -200,11 +200,16 @@ class PADDLE_API Allocator : public phi::Allocator {
   }
 
   uint64_t Release(const phi::Place& place) { return ReleaseImpl(place); }
+  size_t Compact(const phi::Place& place) { return CompactImpl(place); }
 
  protected:
   virtual phi::Allocation* AllocateImpl(size_t size) = 0;
   virtual void FreeImpl(phi::Allocation* allocation);
   virtual uint64_t ReleaseImpl(const phi::Place& place UNUSED) { return 0; }
+  virtual size_t CompactImpl(const phi::Place& place UNUSED) {
+    LOG(INFO) << "Compact is not supported";
+    return 0;
+  }
 };
 
 inline size_t AlignedSize(size_t size, size_t alignment) {

paddle/phi/core/memory/allocation/allocator_facade.cc

@@ -1663,6 +1663,12 @@ uint64_t AllocatorFacade::Release(const phi::Place& place) {
       ->Release(place);
 }
 
+size_t AllocatorFacade::Compact(const phi::Place& place) {
+  return GetPrivate()
+      ->GetAllocator(place, /* A non-zero num to choose allocator_ */ 1)
+      ->Compact(place);
+}
+
 std::shared_ptr<phi::Allocation> AllocatorFacade::AllocShared(
     const phi::Place& place, size_t size, const phi::Stream& stream) {
   return std::shared_ptr<phi::Allocation>(Alloc(place, size, stream));

paddle/phi/core/memory/allocation/allocator_facade.h

@@ -70,6 +70,8 @@ class AllocatorFacade {
   PADDLE_API AllocationPtr Alloc(const phi::Place& place, size_t size);
   // Release unused memory pool.
   uint64_t Release(const phi::Place& place);
+  // Compact memory of free blocks held by the VmmAllocator.
+  size_t Compact(const phi::Place& place);
 
   std::shared_ptr<Allocation> AllocShared(const phi::Place& place,
                                           size_t size,

paddle/phi/core/memory/allocation/retry_allocator.h

@@ -58,6 +58,9 @@ class PADDLE_API RetryAllocator : public Allocator {
   uint64_t ReleaseImpl(const phi::Place& place) override {
     return underlying_allocator_->Release(place);
   }
+  size_t CompactImpl(const phi::Place& place) override {
+    return underlying_allocator_->Compact(place);
+  }
 
  private:
   std::shared_ptr<Allocator> underlying_allocator_;

paddle/phi/core/memory/allocation/stat_allocator.h

@@ -69,6 +69,10 @@ class StatAllocator : public Allocator {
     return underlying_allocator_->Release(place);
   }
 
+  size_t CompactImpl(const phi::Place& place) override {
+    return underlying_allocator_->Compact(place);
+  }
+
  private:
   std::shared_ptr<Allocator> underlying_allocator_;
 };

paddle/phi/core/memory/allocation/stream_safe_cuda_allocator.cc

@@ -259,6 +259,17 @@ uint64_t StreamSafeCUDAAllocator::ReleaseImpl(const phi::Place& place) {
   return released_size;
 }
 
+size_t StreamSafeCUDAAllocator::CompactImpl(const phi::Place& place) {
+  std::lock_guard<SpinLock> lock_guard(allocator_map_lock_);
+  VLOG(4) << "enter StreamSafeCUDAAllocator compact!!";
+  std::vector<StreamSafeCUDAAllocator*>& allocators = allocator_map_[place];
+  size_t compact_free_size = 0;
+  for (StreamSafeCUDAAllocator* allocator : allocators) {
+    compact_free_size += allocator->underlying_allocator_->Compact(place_);
+  }
+  return compact_free_size;
+}
+
 void StreamSafeCUDAAllocator::ProcessUnfreedAllocations() {
   // NOTE(Ruibiao): This condition is to reduce lock completion. It does not
   // need to be thread-safe since here occasional misjudgments are permissible.

paddle/phi/core/memory/allocation/stream_safe_cuda_allocator.h

@@ -77,6 +77,7 @@ class StreamSafeCUDAAllocator
   phi::Allocation *AllocateImpl(size_t size) override;
   void FreeImpl(phi::Allocation *allocation) override;
   uint64_t ReleaseImpl(const phi::Place &place) override;
+  size_t CompactImpl(const phi::Place &place) override;
 
  private:
   void ProcessUnfreedAllocations();

paddle/phi/core/memory/malloc.cc

@@ -32,6 +32,10 @@ uint64_t Release(const phi::Place& place) {
   return allocation::AllocatorFacade::Instance().Release(place);
 }
 
+size_t Compact(const phi::GPUPlace& place) {
+  return allocation::AllocatorFacade::Instance().Compact(place);
+}
+
 std::shared_ptr<Allocation> AllocShared(const phi::Place& place,
                                         size_t size,
                                         const phi::Stream& stream) {

paddle/phi/core/memory/malloc.h

@@ -41,6 +41,9 @@ PADDLE_API extern AllocationPtr Alloc(const phi::Place& place, size_t size);
 
 PADDLE_API extern uint64_t Release(const phi::Place& place);
 
+// Compact memory of free blocks held by the VmmAllocator.
+PADDLE_API extern size_t Compact(const phi::GPUPlace& place);
+
 PADDLE_API extern std::shared_ptr<Allocation> AllocShared(
     const phi::Place& place, size_t size, const phi::Stream& stream);