From 9b917d55c1c5458f2e08a79889e9c078f5c93deb Mon Sep 17 00:00:00 2001 From: zhangtongshe Date: Sat, 21 May 2022 19:52:24 +0800 Subject: [PATCH] adding memory usage limits to allocators, and export the c_api. --- python/src/pybind11_allocator.h | 8 +++ src/allocator.cpp | 97 ++++++++++++++++++++++++++------- src/allocator.h | 12 +++- src/c_api.cpp | 76 +++++++++++++++++++++++++- src/c_api.h | 5 ++ tools/modelwriter.h | 4 ++ 6 files changed, 178 insertions(+), 24 deletions(-) diff --git a/python/src/pybind11_allocator.h b/python/src/pybind11_allocator.h index 7c568209cfe..7b7205dee32 100644 --- a/python/src/pybind11_allocator.h +++ b/python/src/pybind11_allocator.h @@ -31,6 +31,10 @@ class PyAllocator : public Base { PYBIND11_OVERLOAD_PURE(void, Base, fastFree, ptr); } + void emptyCache() override + { + PYBIND11_OVERLOAD_PURE(void, Base, emptyCache); + } }; template @@ -46,6 +50,10 @@ class PyAllocatorOther : public PyAllocator { PYBIND11_OVERLOAD(void, Other, fastFree, ptr); } + void emptyCache() override + { + PYBIND11_OVERLOAD_PURE(void, Other, emptyCache); + } }; #if NCNN_VULKAN diff --git a/src/allocator.cpp b/src/allocator.cpp index d89bdf2596d..0b1eedc9d63 100644 --- a/src/allocator.cpp +++ b/src/allocator.cpp @@ -41,6 +41,9 @@ PoolAllocator::PoolAllocator() : Allocator(), d(new PoolAllocatorPrivate) { d->size_compare_ratio = 192; // 0.75f * 256 + mem_usage = 0; + low_thresh = 0; + high_thresh = 0; } PoolAllocator::~PoolAllocator() @@ -64,7 +67,7 @@ PoolAllocator::~PoolAllocator() } PoolAllocator::PoolAllocator(const PoolAllocator&) - : d(0) + : d(0), mem_usage(0), low_thresh(0), high_thresh(0) { } @@ -75,17 +78,7 @@ PoolAllocator& PoolAllocator::operator=(const PoolAllocator&) void PoolAllocator::clear() { - d->budgets_lock.lock(); - - std::list >::iterator it = d->budgets.begin(); - for (; it != d->budgets.end(); ++it) - { - void* ptr = it->second; - ncnn::fastFree(ptr); - } - d->budgets.clear(); - - d->budgets_lock.unlock(); + emptyCache(); } void PoolAllocator::set_size_compare_ratio(float scr) @@ -99,6 +92,12 @@ void PoolAllocator::set_size_compare_ratio(float scr) d->size_compare_ratio = (unsigned int)(scr * 256); } +void PoolAllocator::set_memory_limit(size_t low, size_t high) +{ + low_thresh = low; + high_thresh = high; +} + void* PoolAllocator::fastMalloc(size_t size) { d->budgets_lock.lock(); @@ -136,9 +135,22 @@ void* PoolAllocator::fastMalloc(size_t size) d->payouts_lock.lock(); d->payouts.push_back(std::make_pair(size, ptr)); - + mem_usage += size; d->payouts_lock.unlock(); + if (high_thresh > 0 && mem_usage > high_thresh) + { + d->budgets_lock.lock(); + while (low_thresh > 0 && mem_usage > low_thresh && !d->budgets.empty()) + { + std::list >::iterator it = d->budgets.begin(); + mem_usage -= it->first; + ncnn::fastFree(it->second); + d->budgets.erase(it); + } + d->budgets_lock.unlock(); + } + return ptr; } @@ -174,6 +186,21 @@ void PoolAllocator::fastFree(void* ptr) ncnn::fastFree(ptr); } +void PoolAllocator::emptyCache() +{ + d->budgets_lock.lock(); + + std::list >::iterator it = d->budgets.begin(); + for (; it != d->budgets.end(); ++it) + { + mem_usage -= it->first; + ncnn::fastFree(it->second); + } + + d->budgets.clear(); + d->budgets_lock.unlock(); +} + class UnlockedPoolAllocatorPrivate { public: @@ -186,6 +213,9 @@ UnlockedPoolAllocator::UnlockedPoolAllocator() : Allocator(), d(new UnlockedPoolAllocatorPrivate) { d->size_compare_ratio = 192; // 0.75f * 256 + mem_usage = 0; + low_thresh = 0; + high_thresh = 0; } UnlockedPoolAllocator::~UnlockedPoolAllocator() @@ -209,7 +239,7 @@ UnlockedPoolAllocator::~UnlockedPoolAllocator() } UnlockedPoolAllocator::UnlockedPoolAllocator(const UnlockedPoolAllocator&) - : d(0) + : d(0), mem_usage(0), low_thresh(0), high_thresh(0) { } @@ -220,13 +250,7 @@ UnlockedPoolAllocator& UnlockedPoolAllocator::operator=(const UnlockedPoolAlloca void UnlockedPoolAllocator::clear() { - std::list >::iterator it = d->budgets.begin(); - for (; it != d->budgets.end(); ++it) - { - void* ptr = it->second; - ncnn::fastFree(ptr); - } - d->budgets.clear(); + emptyCache(); } void UnlockedPoolAllocator::set_size_compare_ratio(float scr) @@ -240,6 +264,12 @@ void UnlockedPoolAllocator::set_size_compare_ratio(float scr) d->size_compare_ratio = (unsigned int)(scr * 256); } +void UnlockedPoolAllocator::set_memory_limit(size_t low, size_t high) +{ + low_thresh = low; + high_thresh = high; +} + void* UnlockedPoolAllocator::fastMalloc(size_t size) { // find free budget @@ -265,6 +295,18 @@ void* UnlockedPoolAllocator::fastMalloc(size_t size) void* ptr = ncnn::fastMalloc(size); d->payouts.push_back(std::make_pair(size, ptr)); + mem_usage += size; + + if (high_thresh > 0 && mem_usage > high_thresh) + { + while (low_thresh > 0 && mem_usage > low_thresh && !d->budgets.empty()) + { + std::list >::iterator it = d->budgets.begin(); + mem_usage -= it->first; + ncnn::fastFree(it->second); + d->budgets.erase(it); + } + } return ptr; } @@ -291,6 +333,19 @@ void UnlockedPoolAllocator::fastFree(void* ptr) ncnn::fastFree(ptr); } +void UnlockedPoolAllocator::emptyCache() +{ + std::list >::iterator it = d->budgets.begin(); + for (; it != d->budgets.end(); ++it) + { + mem_usage -= it->first; + void* ptr = it->second; + ncnn::fastFree(ptr); + } + + d->budgets.clear(); +} + #if NCNN_VULKAN VkAllocator::VkAllocator(const VulkanDevice* _vkdev) : vkdev(_vkdev) diff --git a/src/allocator.h b/src/allocator.h index c9fcf90d181..8de3f86bd46 100644 --- a/src/allocator.h +++ b/src/allocator.h @@ -160,6 +160,7 @@ class NCNN_EXPORT Allocator virtual ~Allocator(); virtual void* fastMalloc(size_t size) = 0; virtual void fastFree(void* ptr) = 0; + virtual void emptyCache() = 0; }; class PoolAllocatorPrivate; @@ -172,12 +173,13 @@ class NCNN_EXPORT PoolAllocator : public Allocator // ratio range 0 ~ 1 // default cr = 0.75 void set_size_compare_ratio(float scr); - + void set_memory_limit(size_t low, size_t high); // release all budgets immediately void clear(); virtual void* fastMalloc(size_t size); virtual void fastFree(void* ptr); + virtual void emptyCache(); private: PoolAllocator(const PoolAllocator&); @@ -185,6 +187,9 @@ class NCNN_EXPORT PoolAllocator : public Allocator private: PoolAllocatorPrivate* const d; + size_t mem_usage; + size_t low_thresh; + size_t high_thresh; }; class UnlockedPoolAllocatorPrivate; @@ -197,12 +202,14 @@ class NCNN_EXPORT UnlockedPoolAllocator : public Allocator // ratio range 0 ~ 1 // default cr = 0.75 void set_size_compare_ratio(float scr); + void set_memory_limit(size_t low, size_t high); // release all budgets immediately void clear(); virtual void* fastMalloc(size_t size); virtual void fastFree(void* ptr); + virtual void emptyCache(); private: UnlockedPoolAllocator(const UnlockedPoolAllocator&); @@ -210,6 +217,9 @@ class NCNN_EXPORT UnlockedPoolAllocator : public Allocator private: UnlockedPoolAllocatorPrivate* const d; + size_t mem_usage; + size_t low_thresh; + size_t high_thresh; }; #if NCNN_VULKAN diff --git a/src/c_api.cpp b/src/c_api.cpp index ff5aa91c2f2..cf160fb6c3b 100644 --- a/src/c_api.cpp +++ b/src/c_api.cpp @@ -70,6 +70,16 @@ class PoolAllocator_c_api : public ncnn::PoolAllocator { return allocator->fast_free(allocator, ptr); } + + virtual void emptyCache() + { + return allocator->empty_cache(allocator); + } + + void set_memory_limit(size_t low, size_t high) + { + ((ncnn::PoolAllocator*)(allocator->pthis))->set_memory_limit(low, high); + } public: ncnn_allocator_t allocator; @@ -85,6 +95,11 @@ static void __ncnn_PoolAllocator_fast_free(ncnn_allocator_t allocator, void* ptr ((ncnn::PoolAllocator*)allocator->pthis)->ncnn::PoolAllocator::fastFree(ptr); } +static void __ncnn_PoolAllocator_empty_cache(ncnn_allocator_t allocator) +{ + ((ncnn::PoolAllocator*)allocator->pthis)->ncnn::PoolAllocator::emptyCache(); +} + class UnlockedPoolAllocator_c_api : public ncnn::UnlockedPoolAllocator { public: @@ -104,6 +119,16 @@ class UnlockedPoolAllocator_c_api : public ncnn::UnlockedPoolAllocator return allocator->fast_free(allocator, ptr); } + virtual void emptyCache() + { + return allocator->empty_cache(allocator); + } + + void set_memory_limit(size_t low, size_t high) + { + ((ncnn::UnlockedPoolAllocator*)(allocator->pthis))->set_memory_limit(low, high); + } + public: ncnn_allocator_t allocator; }; @@ -118,24 +143,71 @@ static void __ncnn_UnlockedPoolAllocator_fast_free(ncnn_allocator_t allocator, v ((ncnn::UnlockedPoolAllocator*)allocator->pthis)->ncnn::UnlockedPoolAllocator::fastFree(ptr); } +static void __ncnn_UnlockedPoolAllocator_empty_cache(ncnn_allocator_t allocator) +{ + ((ncnn::UnlockedPoolAllocator*)allocator->pthis)->ncnn::UnlockedPoolAllocator::emptyCache(); +} + ncnn_allocator_t ncnn_allocator_create_pool_allocator() { ncnn_allocator_t allocator = (ncnn_allocator_t)malloc(sizeof(struct __ncnn_allocator_t)); - allocator->pthis = (void*)(new PoolAllocator_c_api(allocator)); + PoolAllocator_c_api* pthis = new PoolAllocator_c_api(allocator); + allocator->pthis = (void*)(pthis); + allocator->fast_malloc = __ncnn_PoolAllocator_fast_malloc; allocator->fast_free = __ncnn_PoolAllocator_fast_free; + allocator->empty_cache = __ncnn_PoolAllocator_empty_cache; + return allocator; +} + +ncnn_allocator_t ncnn_allocator_create_pool_allocator_with_memory_limit(size_t low, size_t high) +{ + ncnn_allocator_t allocator = (ncnn_allocator_t)malloc(sizeof(struct __ncnn_allocator_t)); + PoolAllocator_c_api* pthis = new PoolAllocator_c_api(allocator); + allocator->pthis = (void*)(pthis); + + pthis->set_memory_limit(low, high); + + allocator->fast_malloc = __ncnn_PoolAllocator_fast_malloc; + allocator->fast_free = __ncnn_PoolAllocator_fast_free; + allocator->empty_cache = __ncnn_PoolAllocator_empty_cache; return allocator; } ncnn_allocator_t ncnn_allocator_create_unlocked_pool_allocator() { ncnn_allocator_t allocator = (ncnn_allocator_t)malloc(sizeof(struct __ncnn_allocator_t)); - allocator->pthis = (void*)(new UnlockedPoolAllocator_c_api(allocator)); + UnlockedPoolAllocator_c_api* pthis = new UnlockedPoolAllocator_c_api(allocator); + allocator->pthis = (void*)(pthis); + allocator->fast_malloc = __ncnn_UnlockedPoolAllocator_fast_malloc; allocator->fast_free = __ncnn_UnlockedPoolAllocator_fast_free; + allocator->empty_cache = __ncnn_UnlockedPoolAllocator_empty_cache; return allocator; } +ncnn_allocator_t ncnn_allocator_create_unlocked_pool_allocator_with_memory_limit(size_t low, size_t high) +{ + ncnn_allocator_t allocator = (ncnn_allocator_t)malloc(sizeof(struct __ncnn_allocator_t)); + UnlockedPoolAllocator_c_api* pthis = new UnlockedPoolAllocator_c_api(allocator); + allocator->pthis = (void*)(pthis); + + pthis->set_memory_limit(low, high); + + allocator->fast_malloc = __ncnn_UnlockedPoolAllocator_fast_malloc; + allocator->fast_free = __ncnn_UnlockedPoolAllocator_fast_free; + allocator->empty_cache = __ncnn_UnlockedPoolAllocator_empty_cache; + return allocator; +} + +void ncnn_allocator_empty_cache(ncnn_allocator_t allocator) +{ + if (allocator) + { + ((Allocator*)(allocator->pthis))->emptyCache(); + } +} + void ncnn_allocator_destroy(ncnn_allocator_t allocator) { if (allocator) diff --git a/src/c_api.h b/src/c_api.h index a276adf5763..5e34d133121 100644 --- a/src/c_api.h +++ b/src/c_api.h @@ -36,10 +36,15 @@ struct NCNN_EXPORT __ncnn_allocator_t void* (*fast_malloc)(ncnn_allocator_t allocator, size_t size); void (*fast_free)(ncnn_allocator_t allocator, void* ptr); + void (*empty_cache)(ncnn_allocator_t allocator); }; NCNN_EXPORT ncnn_allocator_t ncnn_allocator_create_pool_allocator(); +NCNN_EXPORT ncnn_allocator_t ncnn_allocator_create_pool_allocator_with_memory_limit(size_t low, size_t high); NCNN_EXPORT ncnn_allocator_t ncnn_allocator_create_unlocked_pool_allocator(); +NCNN_EXPORT ncnn_allocator_t ncnn_allocator_create_unlocked_pool_allocator_with_memory_limit(size_t low, size_t high); + +NCNN_EXPORT void ncnn_allocator_empty_cache(ncnn_allocator_t allocator); NCNN_EXPORT void ncnn_allocator_destroy(ncnn_allocator_t allocator); /* option api */ diff --git a/tools/modelwriter.h b/tools/modelwriter.h index a2d1d5153e2..2946ec3dd28 100644 --- a/tools/modelwriter.h +++ b/tools/modelwriter.h @@ -139,6 +139,10 @@ class MemoryFootprintAllocator : public ncnn::Allocator bookkeeper.erase(bookkeeper.find(ptr)); ncnn::fastFree(ptr); } + + virtual void emptyCache() + { + } public: int current_memory_usage;