Improves FLANN's heap allocations by a memory pool

modules/flann/include/opencv2/flann/heap.h

@@ -36,9 +36,21 @@
 #include <algorithm>
 #include <vector>
 
+#include <unordered_map>
+
 namespace cvflann
 {
 
+// TODO: Define x > y operator and use std::greater<T> instead
+template <typename T>
+struct greater
+{
+    bool operator()(const T& x, const T& y) const
+    {
+        return y < x;
+    }
+};
+
 /**
  * Priority Queue Implementation
  *
@@ -49,117 +61,180 @@ namespace cvflann
 template <typename T>
 class Heap
 {
-
     /**
      * Storage array for the heap.
      * Type T must be comparable.
      */
     std::vector<T> heap;
-    int length;
-
-    /**
-     * Number of element in the heap
-     */
-    int count;
-
-
-
 public:
     /**
-     * Constructor.
+     * \brief Constructs a heap with a pre-allocated capacity
      *
-     * Params:
+     * \param capacity heap maximum capacity
-     *     sz = heap size
      */
-
+    Heap(const int capacity)
-    Heap(int sz)
     {
-        length = sz;
+        reserve(capacity);
-        heap.reserve(length);
+    }
-        count = 0;
+
+    /**
+     * \brief Move-constructs a heap from an external vector
+     *
+     * \param vec external vector
+     */
+    Heap(std::vector<T>&& vec)
+        : heap(std::move(vec))
+    {
+        std::make_heap(heap.begin(), heap.end(), greater<T>());
     }
 
     /**
      *
-     * Returns: heap size
+     * \returns heap size
      */
-    int size()
+    int size() const
     {
-        return count;
+        return (int)heap.size();
     }
 
     /**
-     * Tests if the heap is empty
      *
-     * Returns: true is heap empty, false otherwise
+     * \returns heap capacity
+     */
+    int capacity() const
+    {
+        return (int)heap.capacity();
+    }
+
+    /**
+     * \brief Tests if the heap is empty
+     *
+     * \returns true is heap empty, false otherwise
      */
     bool empty()
     {
-        return size()==0;
+        return heap.empty();
     }
 
     /**
-     * Clears the heap.
+     * \brief Clears the heap.
      */
     void clear()
     {
         heap.clear();
-        count = 0;
     }
 
-    struct CompareT
+    /**
+     * \brief Sets the heap maximum capacity.
+     *
+     * \param capacity heap maximum capacity
+     */
+    void reserve(const int capacity)
     {
-        bool operator()(const T& t_1, const T& t_2) const
+        heap.reserve(capacity);
-        {
+    }
-            return t_2 < t_1;
-        }
-    };
 
     /**
-     * Insert a new element in the heap.
+     * \brief Inserts a new element in the heap.
      *
      * We select the next empty leaf node, and then keep moving any larger
      * parents down until the right location is found to store this element.
      *
-     * Params:
+     * \param value the new element to be inserted in the heap
-     *     value = the new element to be inserted in the heap
      */
     void insert(T value)
     {
         /* If heap is full, then return without adding this element. */
-        if (count == length) {
+        if (size() == capacity()) {
             return;
         }
 
         heap.push_back(value);
-        static CompareT compareT;
+        std::push_heap(heap.begin(), heap.end(), greater<T>());
-        std::push_heap(heap.begin(), heap.end(), compareT);
-        ++count;
     }
 
-
-
     /**
-     * Returns the node of minimum value from the heap (top of the heap).
+     * \brief Returns the node of minimum value from the heap (top of the heap).
      *
-     * Params:
+     * \param[out] value parameter used to return the min element
-     *     value = out parameter used to return the min element
+     * \returns false if heap empty
-     * Returns: false if heap empty
      */
     bool popMin(T& value)
     {
-        if (count == 0) {
+        if (empty()) {
             return false;
         }
 
         value = heap[0];
-        static CompareT compareT;
+        std::pop_heap(heap.begin(), heap.end(), greater<T>());
-        std::pop_heap(heap.begin(), heap.end(), compareT);
         heap.pop_back();
-        --count;
 
         return true;  /* Return old last node. */
     }
+
+    /**
+     * \brief Returns a shared heap for the given memory pool ID.
+     *
+     * It constructs the heap if it does not already exists.
+     *
+     * \param poolId a user-chosen hashable ID for identifying the heap.
+     *     For thread-safe operations, using current thread ID is a good choice.
+     * \param capacity heap maximum capacity
+     * \param iterThreshold remove heaps that were not reused for more than specified iterations count
+     *        if iterThreshold value is less 2, it will be internally adjusted to twice the number of CPU threads
+     * \returns pointer to the heap
+     */
+    template <typename HashableT>
+    static cv::Ptr<Heap<T>> getPooledInstance(
+        const HashableT& poolId, const int capacity, int iterThreshold = 0)
+    {
+        static cv::Mutex mutex;
+        const cv::AutoLock lock(mutex);
+
+        struct HeapMapValueType {
+            cv::Ptr<Heap<T>> heapPtr;
+            int iterCounter;
+        };
+        typedef std::unordered_map<HashableT, HeapMapValueType> HeapMapType;
+
+        static HeapMapType heapsPool;
+        typename HeapMapType::iterator heapIt = heapsPool.find(poolId);
+
+        if (heapIt == heapsPool.end())
+        {
+            // Construct the heap as it does not already exists
+            HeapMapValueType heapAndTimePair = {cv::makePtr<Heap<T>>(capacity), 0};
+            const std::pair<typename HeapMapType::iterator, bool>& emplaceResult = heapsPool.emplace(poolId, std::move(heapAndTimePair));
+            CV_CheckEQ(static_cast<int>(emplaceResult.second), 1, "Failed to insert the heap into its memory pool");
+            heapIt = emplaceResult.first;
+        }
+        else
+        {
+            CV_CheckEQ(heapIt->second.heapPtr.use_count(), 1, "Cannot modify a heap that is currently accessed by another caller");
+            heapIt->second.heapPtr->clear();
+            heapIt->second.heapPtr->reserve(capacity);
+            heapIt->second.iterCounter = 0;
+        }
+
+        if (iterThreshold <= 1) {
+            iterThreshold = 2 * cv::getNumThreads();
+        }
+
+        // Remove heaps that were not reused for more than given iterThreshold
+        typename HeapMapType::iterator cleanupIt = heapsPool.begin();
+        while (cleanupIt != heapsPool.end())
+        {
+            if (cleanupIt->second.iterCounter++ > iterThreshold)
+            {
+                CV_Assert(cleanupIt != heapIt);
+                cleanupIt = heapsPool.erase(cleanupIt);
+                continue;
+            }
+            ++cleanupIt;
+        }
+
+        return heapIt->second.heapPtr;
+    }
 };
 
 }

modules/flann/include/opencv2/flann/hierarchical_clustering_index.h

@@ -532,7 +532,7 @@ public:
         const bool explore_all_trees = get_param(searchParams,"explore_all_trees",false);
 
         // Priority queue storing intermediate branches in the best-bin-first search
-        Heap<BranchSt>* heap = new Heap<BranchSt>((int)size_);
+        const cv::Ptr<Heap<BranchSt>>& heap = Heap<BranchSt>::getPooledInstance(cv::utils::getThreadID(), (int)size_);
 
         std::vector<bool> checked(size_,false);
         int checks = 0;
@@ -548,8 +548,6 @@ public:
             findNN(node, result, vec, checks, maxChecks, heap, checked, false);
         }
 
-        delete heap;
-
         CV_Assert(result.full());
     }
 
@@ -742,7 +740,7 @@ private:
 
 
     void findNN(NodePtr node, ResultSet<DistanceType>& result, const ElementType* vec, int& checks, int maxChecks,
-                Heap<BranchSt>* heap, std::vector<bool>& checked, bool explore_all_trees = false)
+                const cv::Ptr<Heap<BranchSt>>& heap, std::vector<bool>& checked, bool explore_all_trees = false)
     {
         if (node->childs==NULL) {
             if (!explore_all_trees && (checks>=maxChecks) && result.full()) {

modules/flann/include/opencv2/flann/kdtree_index.h

@@ -445,11 +445,12 @@ private:
     {
         int i;
         BranchSt branch;
-
         int checkCount = 0;
-        Heap<BranchSt>* heap = new Heap<BranchSt>((int)size_);
         DynamicBitset checked(size_);
 
+        // Priority queue storing intermediate branches in the best-bin-first search
+        const cv::Ptr<Heap<BranchSt>>& heap = Heap<BranchSt>::getPooledInstance(cv::utils::getThreadID(), (int)size_);
+
         /* Search once through each tree down to root. */
         for (i = 0; i < trees_; ++i) {
             searchLevel(result, vec, tree_roots_[i], 0, checkCount, maxCheck,
@@ -464,8 +465,6 @@ private:
                         epsError, heap, checked, false);
         }
 
-        delete heap;
-
         CV_Assert(result.full());
     }
 
@@ -476,7 +475,7 @@ private:
      *  at least "mindistsq".
      */
     void searchLevel(ResultSet<DistanceType>& result_set, const ElementType* vec, NodePtr node, DistanceType mindist, int& checkCount, int maxCheck,
-                     float epsError, Heap<BranchSt>* heap, DynamicBitset& checked, bool explore_all_trees = false)
+                     float epsError, const cv::Ptr<Heap<BranchSt>>& heap, DynamicBitset& checked, bool explore_all_trees = false)
     {
         if (result_set.worstDist()<mindist) {
             //			printf("Ignoring branch, too far\n");

modules/flann/include/opencv2/flann/kmeans_index.h

@@ -528,7 +528,7 @@ public:
         }
         else {
             // Priority queue storing intermediate branches in the best-bin-first search
-            Heap<BranchSt>* heap = new Heap<BranchSt>((int)size_);
+            const cv::Ptr<Heap<BranchSt>>& heap = Heap<BranchSt>::getPooledInstance(cv::utils::getThreadID(), (int)size_);
 
             int checks = 0;
             for (int i=0; i<trees_; ++i) {
@@ -542,8 +542,6 @@ public:
                 KMeansNodePtr node = branch.node;
                 findNN(node, result, vec, checks, maxChecks, heap);
             }
-            delete heap;
-
             CV_Assert(result.full());
         }
     }
@@ -1529,7 +1527,7 @@ private:
 
 
     void findNN(KMeansNodePtr node, ResultSet<DistanceType>& result, const ElementType* vec, int& checks, int maxChecks,
-                Heap<BranchSt>* heap)
+                const cv::Ptr<Heap<BranchSt>>& heap)
     {
         // Ignore those clusters that are too far away
         {
@@ -1577,7 +1575,7 @@ private:
      *     distances = array with the distances to each child node.
      * Returns:
      */
-    int exploreNodeBranches(KMeansNodePtr node, const ElementType* q, DistanceType* domain_distances, Heap<BranchSt>* heap)
+    int exploreNodeBranches(KMeansNodePtr node, const ElementType* q, DistanceType* domain_distances, const cv::Ptr<Heap<BranchSt>>& heap)
     {
 
         int best_index = 0;