From 8f80f5bc794d8900f9d57b51eea167f4dde2903c Mon Sep 17 00:00:00 2001
From: liaogang <liaogang@baidu.com>
Date: Wed, 16 Aug 2017 19:46:12 +0800
Subject: [PATCH] FIX: Release CPU/GPU memory via deleter

---
 paddle/memory/memory.cc | 59 ++++++++++++++++++++++++++++-------------
 1 file changed, 41 insertions(+), 18 deletions(-)

diff --git a/paddle/memory/memory.cc b/paddle/memory/memory.cc
index 207025f9b1..5946c3ea4a 100644
--- a/paddle/memory/memory.cc
+++ b/paddle/memory/memory.cc
@@ -16,19 +16,31 @@ limitations under the License. */
 #include "paddle/memory/detail/buddy_allocator.h"
 #include "paddle/memory/detail/system_allocator.h"
 
-#include <cstring>  // for memcpy
+#include <algorithm>  // for transfrom
+#include <cstring>    // for memcpy
+#include <mutex>      // for call_once
+
+#include "glog/logging.h"
 
 namespace paddle {
 namespace memory {
 
-detail::BuddyAllocator* GetCPUBuddyAllocator() {
-  static detail::BuddyAllocator* a = nullptr;
-  if (a == nullptr) {
-    a = new detail::BuddyAllocator(new detail::CPUAllocator,
-                                   platform::CpuMinChunkSize(),
-                                   platform::CpuMaxChunkSize());
-  }
-  return a;
+using BuddyAllocator = detail::BuddyAllocator;
+
+std::once_flag cpu_alloctor_flag;
+std::once_flag gpu_alloctor_flag;
+
+BuddyAllocator* GetCPUBuddyAllocator() {
+  static std::unique_ptr<BuddyAllocator, void (*)(BuddyAllocator*)> a{
+      nullptr, [](BuddyAllocator* p) { delete p; }};
+
+  std::call_once(cpu_alloctor_flag, [&]() {
+    a.reset(new BuddyAllocator(new detail::CPUAllocator,
+                               platform::CpuMinChunkSize(),
+                               platform::CpuMaxChunkSize()));
+  });
+
+  return a.get();
 }
 
 template <>
@@ -48,20 +60,31 @@ size_t Used<platform::CPUPlace>(platform::CPUPlace place) {
 
 #ifndef PADDLE_ONLY_CPU
 
-detail::BuddyAllocator* GetGPUBuddyAllocator(int gpu_id) {
-  static detail::BuddyAllocator** as = NULL;
-  if (as == NULL) {
+BuddyAllocator* GetGPUBuddyAllocator(int gpu_id) {
+  using BuddyAllocVec = std::vector<BuddyAllocator*>;
+  static std::unique_ptr<BuddyAllocVec, void (*)(BuddyAllocVec * p)> as{
+      new std::vector<BuddyAllocator*>, [](BuddyAllocVec* p) {
+        std::for_each(p->begin(), p->end(),
+                      [](BuddyAllocator* p) { delete p; });
+      }};
+
+  // GPU buddy alloctors
+  auto& alloctors = *as.get();
+
+  // GPU buddy allocator initialization
+  std::call_once(gpu_alloctor_flag, [&]() {
     int gpu_num = platform::GetDeviceCount();
-    as = new detail::BuddyAllocator*[gpu_num];
+    alloctors.reserve(gpu_num);
     for (int gpu = 0; gpu < gpu_num; gpu++) {
       platform::SetDeviceId(gpu);
-      as[gpu] = new detail::BuddyAllocator(new detail::GPUAllocator,
-                                           platform::GpuMinChunkSize(),
-                                           platform::GpuMaxChunkSize());
+      alloctors.emplace_back(new BuddyAllocator(new detail::GPUAllocator,
+                                                platform::GpuMinChunkSize(),
+                                                platform::GpuMaxChunkSize()));
     }
-  }
+  });
+
   platform::SetDeviceId(gpu_id);
-  return as[gpu_id];
+  return alloctors[gpu_id];
 }
 
 template <>
-- 
GitLab