diff --git a/mindspore/lite/src/runtime/allocator.cc b/mindspore/lite/src/runtime/allocator.cc index 75a033de9ed..48021eaf3e8 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(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 9d324747bde..90fb76b1441 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 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; // buf, membuf> std::unordered_map allocatedList_; std::multimap freeList_; diff --git a/mindspore/lite/src/runtime/opencl/opencl_allocator.cc b/mindspore/lite/src/runtime/opencl/opencl_allocator.cc index 84ffd87c0de..dab5258cdf9 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 f339f15b350..84246fe410b 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 &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);