From 2214ce3e2cdf9d5b18919d1d875f5de014d46ca1 Mon Sep 17 00:00:00 2001
From: Yutian Li <hotpxless@gmail.com>
Date: Mon, 10 Aug 2015 23:51:44 +0800
Subject: [PATCH] [storage] storage managers

---
 src/storage/naive_storage_manager.h  |  2 +-
 src/storage/pooled_storage_manager.h | 42 ++++++++++++++++++++++++++--
 src/storage/storage_manager.h        |  3 +-
 3 files changed, 43 insertions(+), 4 deletions(-)

diff --git a/src/storage/naive_storage_manager.h b/src/storage/naive_storage_manager.h
index c0cad9355262..d145cada3f43 100644
--- a/src/storage/naive_storage_manager.h
+++ b/src/storage/naive_storage_manager.h
@@ -27,7 +27,7 @@ class NaiveStorageManager final : public StorageManager {
    */
   ~NaiveStorageManager() = default;
   void* Alloc(size_t size) override;
-  void Free(void* ptr) override;
+  void Free(void* ptr, size_t) override;
 
  private:
   DISALLOW_COPY_AND_ASSIGN(NaiveStorageManager);
diff --git a/src/storage/pooled_storage_manager.h b/src/storage/pooled_storage_manager.h
index 6420aa86125b..fb168391dc5e 100644
--- a/src/storage/pooled_storage_manager.h
+++ b/src/storage/pooled_storage_manager.h
@@ -6,6 +6,8 @@
 #ifndef MXNET_STORAGE_POOLED_STORAGE_MANAGER_H_
 #define MXNET_STORAGE_POOLED_STORAGE_MANAGER_H_
 
+#include <unordered_map>
+#include <vector>
 #include "./storage_manager.h"
 #include "mxnet/base.h"
 
@@ -15,7 +17,7 @@ namespace storage {
 /*!
  * \brief Storage manager with a memory pool.
  */
-template <class DeviceStorage>
+template <class DeviceStorage, size_t kThreshold>
 class PooledStorageManager final : public StorageManager {
  public:
   /*!
@@ -27,12 +29,48 @@ class PooledStorageManager final : public StorageManager {
    */
   ~PooledStorageManager() = default;
   void* Alloc(size_t size) override;
-  void Free(void* ptr) override;
+  void Free(void* ptr, size_t size) override;
 
  private:
+  void ReleaseAll();
+  size_t used_memory_ = 0;
+  std::unordered_map<size_t, std::vector<void*>> memory_pool_;
   DISALLOW_COPY_AND_ASSIGN(PooledStorageManager);
 };
 
+templace <class DeviceStorage, size_t kThreshold>
+void* PooledStorageManager<DeviceStorage, kThreshold>::Alloc(size_t size) {
+  auto&& reuse_it = memory_pool_.find(size);
+  if (reuse_it == memory_pool_.end() || reuse_it->second.size() == 0) {
+    if (kThreshold <= used_memory_) {
+        ReleaseAll();
+    }
+    used_memory_ += size;
+    return DeviceStorage::Alloc(size);
+  } else {
+    auto&& reuse_pool = reuse_it->second;
+    auto ret = reuse_pool.back();
+    reuse_pool.pop_back();
+    return ret;
+  }
+}
+
+templace <class DeviceStorage, size_t kThreshold>
+void PooledStorageManager<DeviceStorage, kThreshold>::Free(void* ptr, size_t size) {
+  auto&& reuse_pool = memory_pool_[size];
+  reuse_pool.push_back(ptr);
+}
+
+template <class DeviceStorage, size_t kThreshold>
+void PooledStorageManager<DeviceStorage, kThreshold>::ReleaseAll() {
+  for (auto&& i : memory_pool_) {
+    for (auto&& j : i.second) {
+      DeviceStorage::Free(i.second);
+      used_memory_ -= i.first;
+    }
+  }
+}
+
 }  // namespace storage
 }  // namespace mxnet
 
diff --git a/src/storage/storage_manager.h b/src/storage/storage_manager.h
index 98129675ddf5..94967a210e80 100644
--- a/src/storage/storage_manager.h
+++ b/src/storage/storage_manager.h
@@ -25,8 +25,9 @@ class StorageManager {
   /*!
    * \brief Deallocation.
    * \param ptr Pointer to deallocate.
+   * \param size Size of the storage.
    */
-  virtual void* Free(void* ptr) = 0;
+  virtual void* Free(void* ptr, size_t size) = 0;
   /*!
    * \brief Destructor.
    */