!10345 fix mempool bug

From: @hangangqiang
Reviewed-by: @zhanghaibo5,@zhang_xue_tong
Signed-off-by: @zhanghaibo5

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();
 

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_;

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;
 

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);