From 5ea636e0dd4177b162cb5ad1a40bb6fe302d4201 Mon Sep 17 00:00:00 2001
From: hangangqiang <hangangqiang2@huawei.com>
Date: Tue, 22 Dec 2020 19:35:07 +0800
Subject: [PATCH] fix mempool bug

---
 mindspore/lite/src/runtime/allocator.cc       | 20 ++-----------------
 mindspore/lite/src/runtime/allocator.h        |  5 +++--
 .../src/runtime/opencl/opencl_allocator.cc    |  2 +-
 .../src/runtime/opencl/opencl_allocator.h     |  2 +-
 4 files changed, 7 insertions(+), 22 deletions(-)

diff --git a/mindspore/lite/src/runtime/allocator.cc b/mindspore/lite/src/runtime/allocator.cc
index 75a033de9e..48021eaf3e 100644
--- a/mindspore/lite/src/runtime/allocator.cc
+++ b/mindspore/lite/src/runtime/allocator.cc
@@ -49,7 +49,7 @@ void *DefaultAllocator::Malloc(size_t size) {
     MS_LOG(ERROR) << "MallocData out of max_size, size: " << size;
     return nullptr;
   }
-  if (this->GetTotalSize() >= MAX_THREAD_POOL_SIZE) {
+  if (this->total_size_ >= MAX_THREAD_POOL_SIZE) {
     MS_LOG(ERROR) << "Memory pool is exhausted";
     return nullptr;
   }
@@ -69,6 +69,7 @@ void *DefaultAllocator::Malloc(size_t size) {
     UnLock();
     return nullptr;
   }
+  this->total_size_ += size;
   membuf->size = size;
   membuf->buf = reinterpret_cast<char *>(membuf.get()) + sizeof(MemBuf);
   auto bufPtr = membuf->buf;
@@ -95,23 +96,6 @@ void DefaultAllocator::Free(void *buf) {
   buf = nullptr;
 }
 
-size_t DefaultAllocator::GetTotalSize() {
-  Lock();
-  size_t totalSize = 0;
-
-  for (auto &it : allocatedList_) {
-    auto membuf = it.second;
-    totalSize += membuf->size;
-  }
-
-  for (auto &it : freeList_) {
-    auto membuf = it.second;
-    totalSize += membuf->size;
-  }
-  UnLock();
-  return totalSize;
-}
-
 void DefaultAllocator::Clear() {
   Lock();
 
diff --git a/mindspore/lite/src/runtime/allocator.h b/mindspore/lite/src/runtime/allocator.h
index 9d324747bd..90fb76b144 100644
--- a/mindspore/lite/src/runtime/allocator.h
+++ b/mindspore/lite/src/runtime/allocator.h
@@ -38,7 +38,7 @@ class Allocator {
   virtual void *Malloc(size_t size) = 0;
   virtual void Free(void *ptr) = 0;
   virtual void SetContext(const AllocatorContext &ctx) {}
-  virtual size_t GetTotalSize() { return 0; }
+  virtual size_t total_size() = 0;
   static std::shared_ptr<Allocator> Create();
   virtual void *Prepare(void *ptr) { return ptr; }
   std::string name;
@@ -51,7 +51,7 @@ class DefaultAllocator : public Allocator {
   void SetContext(const AllocatorContext &ctx) override;
   void *Malloc(size_t size) override;
   void Free(void *ptr) override;
-  size_t GetTotalSize() override;
+  size_t total_size() override { return this->total_size_; }
   void Clear();
 
  private:
@@ -63,6 +63,7 @@ class DefaultAllocator : public Allocator {
   };
 
   std::mutex lock_;
+  size_t total_size_ = 0;
   // <membuf->buf, membuf>
   std::unordered_map<void *, MemBuf *> allocatedList_;
   std::multimap<size_t, MemBuf *> freeList_;
diff --git a/mindspore/lite/src/runtime/opencl/opencl_allocator.cc b/mindspore/lite/src/runtime/opencl/opencl_allocator.cc
index 84ffd87c0d..dab5258cdf 100644
--- a/mindspore/lite/src/runtime/opencl/opencl_allocator.cc
+++ b/mindspore/lite/src/runtime/opencl/opencl_allocator.cc
@@ -225,7 +225,7 @@ void OpenCLAllocator::Free(void *buf) {
   MS_LOG(WARNING) << "Host ptr " << buf << " has freed";
 }
 
-size_t OpenCLAllocator::GetTotalSize() {
+size_t OpenCLAllocator::total_size() {
   Lock();
   size_t totalSize = 0;
 
diff --git a/mindspore/lite/src/runtime/opencl/opencl_allocator.h b/mindspore/lite/src/runtime/opencl/opencl_allocator.h
index f339f15b35..84246fe410 100644
--- a/mindspore/lite/src/runtime/opencl/opencl_allocator.h
+++ b/mindspore/lite/src/runtime/opencl/opencl_allocator.h
@@ -40,7 +40,7 @@ class OpenCLAllocator : public Allocator {
   void *Malloc(size_t size) override;
   void *Malloc(size_t size, const std::vector<size_t> &img_size, void *data = nullptr);
   void Free(void *ptr) override;
-  size_t GetTotalSize() override;
+  size_t total_size() override;
 
   void Clear();
   void *GetImage(void *host_ptr);