Memory Pool Improvement For Variadic Sized Inputs

src/allocator.cpp

@@ -33,6 +33,7 @@ class PoolAllocatorPrivate
     Mutex budgets_lock;
     Mutex payouts_lock;
     unsigned int size_compare_ratio; // 0~256
+    static const size_t size_threshold = 10;
     std::list<std::pair<size_t, void*> > budgets;
     std::list<std::pair<size_t, void*> > payouts;
 };
@@ -104,11 +105,20 @@ void* PoolAllocator::fastMalloc(size_t size)
     d->budgets_lock.lock();
 
     // find free budget
-    std::list<std::pair<size_t, void*> >::iterator it = d->budgets.begin();
+    std::list<std::pair<size_t, void*> >::iterator it = d->budgets.begin(), it_max = d->budgets.end(), it_min = d->budgets.end();
     for (; it != d->budgets.end(); ++it)
     {
         size_t bs = it->first;
 
+        if (it_min == d->budgets.end() || it->first < it_min->first)
+        {
+            it_min = it;
+        }
+        if (it_max == d->budgets.end() || it->first > it_max->first)
+        {
+            it_max = it;
+        }
+
         // size_compare_ratio ~ 100%
         if (bs >= size && ((bs * d->size_compare_ratio) >> 8) <= size)
         {
@@ -128,6 +138,26 @@ void* PoolAllocator::fastMalloc(size_t size)
         }
     }
 
+    if (d->budgets.size() >= d->size_threshold)
+    {
+        // All chunks in pool are not chosen. Then try to drop some outdated
+        // chunks and return them to OS.
+        if (it_max->first < size)
+        {
+            // Current query is asking for a chunk larger than any cached chunks.
+            // Then remove the smallest one.
+            ncnn::fastFree(it_min->second);
+            d->budgets.erase(it_min);
+        }
+        else if (it_min->first > size)
+        {
+            // Current query is asking for a chunk smaller than any cached chunks.
+            // Then remove the largest one.
+            ncnn::fastFree(it_max->second);
+            d->budgets.erase(it_max);
+        }
+    }
+
     d->budgets_lock.unlock();
 
     // new
@@ -178,6 +208,7 @@ class UnlockedPoolAllocatorPrivate
 {
 public:
     unsigned int size_compare_ratio; // 0~256
+    static const size_t size_threshold = 10;
     std::list<std::pair<size_t, void*> > budgets;
     std::list<std::pair<size_t, void*> > payouts;
 };
@@ -243,11 +274,20 @@ void UnlockedPoolAllocator::set_size_compare_ratio(float scr)
 void* UnlockedPoolAllocator::fastMalloc(size_t size)
 {
     // find free budget
-    std::list<std::pair<size_t, void*> >::iterator it = d->budgets.begin();
+    std::list<std::pair<size_t, void*> >::iterator it = d->budgets.begin(), it_max = d->budgets.begin(), it_min = d->budgets.begin();
     for (; it != d->budgets.end(); ++it)
     {
         size_t bs = it->first;
 
+        if (it->first > it_max->first)
+        {
+            it_max = it;
+        }
+        if (it->first < it_min->first)
+        {
+            it_min = it;
+        }
+
         // size_compare_ratio ~ 100%
         if (bs >= size && ((bs * d->size_compare_ratio) >> 8) <= size)
         {
@@ -261,6 +301,20 @@ void* UnlockedPoolAllocator::fastMalloc(size_t size)
         }
     }
 
+    if (d->budgets.size() >= d->size_threshold)
+    {
+        if (it_max->first < size)
+        {
+            ncnn::fastFree(it_min->second);
+            d->budgets.erase(it_min);
+        }
+        else if (it_min->first > size)
+        {
+            ncnn::fastFree(it_max->second);
+            d->budgets.erase(it_max);
+        }
+    }
+
     // new
     void* ptr = ncnn::fastMalloc(size);