【Allocator】Add compact func (PaddlePaddle#76223)

liuruyan · web-flow · commit 42b814627a8c · 2025-11-14T20:52:32.000+08:00
* add compact func

* fix ut

* fix ci

* add ut

* fix ci

* fix ci

* fix ci

* fix ci

* fix ci

* fix ci
diff --git a/paddle/fluid/pybind/pybind.cc b/paddle/fluid/pybind/pybind.cc
@@ -3605,6 +3605,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 & {
diff --git a/paddle/phi/core/memory/allocation/allocator.h b/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) {
diff --git a/paddle/phi/core/memory/allocation/allocator_facade.cc b/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));
diff --git a/paddle/phi/core/memory/allocation/allocator_facade.h b/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,
diff --git a/paddle/phi/core/memory/allocation/retry_allocator.h b/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_;
diff --git a/paddle/phi/core/memory/allocation/stat_allocator.h b/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_;
 };
diff --git a/paddle/phi/core/memory/allocation/stream_safe_cuda_allocator.cc b/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.
diff --git a/paddle/phi/core/memory/allocation/stream_safe_cuda_allocator.h b/paddle/phi/core/memory/allocation/stream_safe_cuda_allocator.h
@@ -84,6 +84,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();
diff --git a/paddle/phi/core/memory/malloc.cc b/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) {
diff --git a/paddle/phi/core/memory/malloc.h b/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);
 
diff --git a/paddle/phi/core/platform/device/gpu/gpu_info.cc b/paddle/phi/core/platform/device/gpu/gpu_info.cc
@@ -615,6 +615,8 @@ void EmptyCache() {
   }
 }
 
+void VmmCompact() { memory::Compact(phi::GPUPlace(GetCurrentDeviceId())); }
+
 bool IsGPUManagedMemorySupported(int dev_id) {
   return phi::backends::gpu::IsGPUManagedMemorySupported(dev_id);
 }
diff --git a/paddle/phi/core/platform/device/gpu/gpu_info.h b/paddle/phi/core/platform/device/gpu/gpu_info.h
@@ -179,6 +179,9 @@ PADDLE_API bool IsGpuMallocRecorded(int dev_id);
 //! Empty idle cached memory held by the allocator.
 PADDLE_API void EmptyCache(void);
 
+//! Compact memory of free blocks held by the VmmAllocator.
+PADDLE_API void VmmCompact(void);
+
 PADDLE_API bool IsGPUManagedMemorySupported(int dev_id);
 
 PADDLE_API bool IsGPUManagedMemoryOversubscriptionSupported(int dev_id);
diff --git a/test/cpp/phi/CMakeLists.txt b/test/cpp/phi/CMakeLists.txt
@@ -4,3 +4,4 @@ add_subdirectory(common)
 add_subdirectory(core)
 add_subdirectory(kernels)
 add_subdirectory(ops)
+add_subdirectory(memory)
diff --git a/test/cpp/phi/memory/CMakeLists.txt b/test/cpp/phi/memory/CMakeLists.txt
@@ -0,0 +1,10 @@
+if(WITH_GPU)
+  if(WIN32)
+    message(STATUS "Skip compact_allocator_test on Windows")
+  else()
+    nv_test(
+      compact_allocator_test
+      SRCS compact_allocator_test.cc
+      DEPS phi common)
+  endif()
+endif()
diff --git a/test/cpp/phi/memory/compact_allocator_test.cc b/test/cpp/phi/memory/compact_allocator_test.cc
@@ -0,0 +1,43 @@
+// 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
+//
+//     http://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/phi/core/memory/allocation/allocator.h"
+#include "paddle/phi/core/memory/allocation/cuda_virtual_mem_allocator.h"
+#include "paddle/phi/core/memory/allocation/retry_allocator.h"
+#include "paddle/phi/core/memory/allocation/virtual_memory_auto_growth_best_fit_allocator.h"
+#include "paddle/phi/core/platform/device/gpu/gpu_info.h"
+#ifdef PADDLE_WITH_CUDA
+#include <cuda.h>
+#include <cuda_runtime.h>
+#endif
+#include "gtest/gtest.h"
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+TEST(VirtualMemoryAutoGrowthBestFitAllocator, TestCompact) {
+  auto vmm_cuda_allocator =
+      std::make_shared<CUDAVirtualMemAllocator>(phi::GPUPlace());
+  auto vma_allocator =
+      std::make_shared<VirtualMemoryAutoGrowthBestFitAllocator>(
+          vmm_cuda_allocator, platform::GpuMinChunkSize(), phi::GPUPlace());
+  size_t mb = (1 << 20);
+  vma_allocator->Allocate(1 * mb);
+  vma_allocator->Allocate(2 * mb);
+  vma_allocator->Compact(phi::GPUPlace());
+}
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle
diff --git a/test/legacy_test/test_vmm_compact.py b/test/legacy_test/test_vmm_compact.py
@@ -0,0 +1,81 @@
+# 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
+#
+#     http://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.
+
+import unittest
+
+import numpy as np
+
+import paddle
+
+
+@unittest.skipIf(
+    (not paddle.is_compiled_with_cuda()) or paddle.is_compiled_with_rocm(),
+    'should compile with cuda.',
+)
+class TestVmmCompact(unittest.TestCase):
+    def setUp(self):
+        self.GB = 1000**3
+        self.MB = 1000**2
+        paddle.set_flags({'FLAGS_use_virtual_memory_auto_growth': True})
+        self.cmds = [
+            ["Alloc", 1 * self.GB, "0x100000000"],
+            ["Alloc", 2 * self.GB, "0x100000001"],
+            ["Alloc", 2 * self.GB, "0x100000002"],
+            ["Alloc", 1 * self.GB, "0x100000003"],
+            ["Free", 1 * self.GB, "0x100000000"],
+            ["Free", 1 * self.GB, "0x100000003"],
+        ]
+
+    def test_compact_api(self):
+        params = {}
+        (
+            old_tensor1,
+            new_tensor1,
+        ) = 0, 0
+
+        for op, size, ptr in self.cmds:
+            paddle.device.synchronize()
+            paddle_reserved1 = paddle.device.cuda.memory_reserved() // self.MB
+
+            if op == "Alloc":
+                params[ptr] = paddle.randn(
+                    [int(int(size) / 4)], dtype='float32'
+                )
+            if op == "Free" and ptr in params:
+                del params[ptr]
+
+            if ptr == '0x100000001':
+                old_tensor1 = params['0x100000001'].numpy()[0:100]
+
+            paddle.device.synchronize()
+            paddle_reserved2 = paddle.device.cuda.memory_reserved() // self.MB
+            paddle_allocated2 = paddle.device.cuda.memory_allocated() // self.MB
+            paddle_max_reserved = (
+                paddle.device.cuda.max_memory_reserved() // self.MB
+            )
+            paddle_max_allocated = (
+                paddle.device.cuda.max_memory_allocated() // self.MB
+            )
+
+            print(
+                f"reserved = {paddle_reserved2} allocated = {paddle_allocated2} auto growth = {paddle_reserved2 - paddle_reserved1} max_allocated = {paddle_max_allocated} max_reserved = {paddle_max_reserved}"
+            )
+
+        paddle.core.vmm_compact()
+        new_tensor1 = params['0x100000001'].numpy()[0:100]
+        np.testing.assert_array_equal(old_tensor1, new_tensor1)
+
+
+if __name__ == '__main__':
+    unittest.main()

Original file line number	Diff line number	Diff line change
`@@ -3605,6 +3605,7 @@ All parameter, weight, gradient are variables in Paddle.`
`3605`	`3605`	`}`
`3606`	`3606`	`platform::EmptyCache();`
`3607`	`3607`	`});`
	`3608`	`+ m.def("vmm_compact", [] { platform::VmmCompact(); });`
`3608`	`3609`	`m.def(`
`3609`	`3610`	`"get_device_properties",`
`3610`	`3611`	`[](int id) -> const gpuDeviceProp & {`
Original file line number	Diff line number	Diff line change
`@@ -58,6 +58,9 @@ class PADDLE_API RetryAllocator : public Allocator {`
`58`	`58`	`uint64_t ReleaseImpl(const phi::Place& place) override {`
`59`	`59`	`return underlying_allocator_->Release(place);`
`60`	`60`	`}`
	`61`	`+ size_t CompactImpl(const phi::Place& place) override {`
	`62`	`+ return underlying_allocator_->Compact(place);`
	`63`	`+ }`
`61`	`64`
`62`	`65`	`private:`
`63`	`66`	`std::shared_ptr<Allocator> underlying_allocator_;`