adding memory usage limits to allocators, and export the c_api.

Tencent · May 21, 2022 · 9b917d5 · 9b917d5
1 parent ab4b012
commit 9b917d5
Show file tree

Hide file tree

Showing 6 changed files with 178 additions and 24 deletions.
diff --git 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<class Other>
@@ -46,6 +50,10 @@ class PyAllocatorOther : public PyAllocator<Other>
     {
         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
@@ -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<std::pair<size_t, void*> >::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<std::pair<size_t, void*> >::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<std::pair<size_t, void*> >::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<std::pair<size_t, void*> >::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<std::pair<size_t, void*> >::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<std::pair<size_t, void*> >::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
@@ -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,19 +173,23 @@ 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&);
     PoolAllocator& operator=(const PoolAllocator&);
 
 private:
     PoolAllocatorPrivate* const d;
+    size_t mem_usage;
+    size_t low_thresh;
+    size_t high_thresh;
 };
 
 class UnlockedPoolAllocatorPrivate;
@@ -197,19 +202,24 @@ 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&);
     UnlockedPoolAllocator& operator=(const UnlockedPoolAllocator&);
 
 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
@@ -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
@@ -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 */