Merge branch 'optimize-priority-queue'

regtest/basic/rt-make-9/main.cpp

@@ -1,4 +1,4 @@
-#include "plumed/tools/Tools.h"
+#include "plumed/tools/MergeVectorTools.h"
 #include <fstream>
 #include <iostream>
 #include <vector>
@@ -20,7 +20,7 @@ void run(bool queue) {
     vecs.push_back(&v2);
     vecs.push_back(&v3);
     vecs.push_back(&v4);
-    Tools::mergeSortedVectors(vecs,result);
+    mergeVectorTools::mergeSortedVectors(vecs,result);
     for(const auto &v : result) ofs<<" "<<v; ofs<<"\n";
   }
 
@@ -39,7 +39,7 @@ void run(bool queue) {
     vecs.push_back(&v2);
     vecs.push_back(&v3);
     vecs.push_back(&v4);
-    Tools::mergeSortedVectors(vecs,result);
+    mergeVectorTools::mergeSortedVectors(vecs,result);
     for(const auto &v : result) ofs<<" "<<v; ofs<<"\n";
   }
 }

src/core/DomainDecomposition.cpp

@@ -26,6 +26,9 @@
 #include "PlumedMain.h"
 #include "ActionSet.h"
 
+#include "small_vector/small_vector.h"
+#include "tools/MergeVectorTools.h"
+
 //+PLUMEDOC ANALYSIS DOMAIN_DECOMPOSITION
 /*
 Pass domain decomposed properties of atoms to PLUMED
@@ -37,13 +40,16 @@ Pass domain decomposed properties of atoms to PLUMED
 
 namespace PLMD {
 
-/// Use a priority_queue to merge unique vectors.
-/// export PLUMED_MERGE_VECTORS_PRIORITY_QUEUE=yes to use a priority_queue.
-/// Might be faster with some settings, but appears to not be in practice.
-/// This option is for testing and might be removed.
-static bool getenvMergeVectorsPriorityQueue() noexcept {
-  static const auto* res=std::getenv("PLUMED_MERGE_VECTORS_PRIORITY_QUEUE");
-  return res;
+namespace {
+
+enum class Option {automatic, no, yes };
+
+Option interpretEnvString(const char* env,const char* str) {
+  if(!str) return Option::automatic;
+  if(!std::strcmp(str,"yes"))return Option::yes;
+  if(!std::strcmp(str,"no"))return Option::no;
+  if(!std::strcmp(str,"auto"))return Option::automatic;
+  plumed_error()<<"Cannot understand env var "<<env<<"\nPossible values: yes/no/auto\nActual value: "<<str;
 }
 
 /// Use unique list of atoms to manipulate forces and positions.
@@ -55,9 +61,12 @@ static bool getenvMergeVectorsPriorityQueue() noexcept {
 /// export PLUMED_FORCE_UNIQUE=no   # enforce not using the unique list in serial runs
 /// export PLUMED_FORCE_UNIQUE=auto # choose heuristically
 /// default: auto
-static const char* getenvForceUnique() noexcept {
-  static const auto* res=std::getenv("PLUMED_FORCE_UNIQUE");
-  return res;
+Option getenvForceUnique() {
+  static const char* name="PLUMED_FORCE_UNIQUE";
+  static const auto opt = interpretEnvString(name,std::getenv(name));
+  return opt;
+}
+
 }
 
 PLUMED_REGISTER_ACTION(DomainDecomposition,"DOMAIN_DECOMPOSITION")
@@ -238,7 +247,7 @@ void DomainDecomposition::share() {
     shareAll(); return;
   }
 
-  if(!getenvForceUnique() || !std::strcmp(getenvForceUnique(),"auto")) {
+  if(getenvForceUnique()==Option::automatic) {
     unsigned largest=0;
     for(unsigned i=0; i<actions.size(); i++) {
       if(actions[i]->isActive()) {
@@ -248,20 +257,20 @@ void DomainDecomposition::share() {
     }
     if(largest*2<getNumberOfAtoms()) unique_serial=true;
     else unique_serial=false;
-  } else if(!std::strcmp(getenvForceUnique(),"yes")) {
+  } else if(getenvForceUnique()==Option::yes) {
     unique_serial=true;
-  } else if(!std::strcmp(getenvForceUnique(),"no")) {
+  } else if(getenvForceUnique()==Option::no) {
     unique_serial=false;
   } else {
-    plumed_error()<<"PLUMED_FORCE_UNIQUE set to unknown value "<<getenvForceUnique();
+    plumed_error();
   }
 
   if(unique_serial || !(int(gatindex.size())==getNumberOfAtoms() && shuffledAtoms==0)) {
     for(unsigned i=0; i<actions.size(); i++) {
       if( actions[i]->unique_local_needs_update ) actions[i]->updateUniqueLocal( !(dd && shuffledAtoms>0), g2l );
     }
     // Now reset unique for the new step
-    std::vector<const std::vector<AtomNumber>*> vectors;
+    gch::small_vector<const std::vector<AtomNumber>*,32> vectors;
     vectors.reserve(actions.size());
     for(unsigned i=0; i<actions.size(); i++) {
       if(actions[i]->isActive()) {
@@ -273,7 +282,7 @@ void DomainDecomposition::share() {
     }
     if(!vectors.empty()) atomsNeeded=true;
     unique.clear();
-    Tools::mergeSortedVectors(vectors,unique,getenvMergeVectorsPriorityQueue());
+    mergeVectorTools::mergeSortedVectors(vectors.data(),vectors.size(),unique);
   } else {
     for(unsigned i=0; i<actions.size(); i++) {
       if(actions[i]->isActive()) {
@@ -434,7 +443,8 @@ const std::vector<int>& DomainDecomposition::getGatindex() const {
 }
 
 void DomainDecomposition::getAllActiveAtoms( std::vector<AtomNumber>& u ) {
-  std::vector<const std::vector<AtomNumber>*> vectors;
+  gch::small_vector<const std::vector<AtomNumber>*,32> vectors;
+  vectors.reserve(actions.size());
   for(unsigned i=0; i<actions.size(); i++) {
     if(actions[i]->isActive()) {
       if(!actions[i]->getUnique().empty()) {
@@ -444,7 +454,7 @@ void DomainDecomposition::getAllActiveAtoms( std::vector<AtomNumber>& u ) {
     }
   }
   u.clear();
-  Tools::mergeSortedVectors(vectors,u,getenvMergeVectorsPriorityQueue());
+  mergeVectorTools::mergeSortedVectors(vectors.data(),vectors.size(),u);
 }
 
 void DomainDecomposition::createFullList(const TypesafePtr & n) {

src/tools/MergeVectorTools.cpp

@@ -0,0 +1,23 @@
+/* +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
+   Copyright (c) 2011-2023 The plumed team
+   (see the PEOPLE file at the root of the distribution for a list of names)
+
+   See http://www.plumed.org for more information.
+
+   This file is part of plumed, version 2.
+
+   plumed is free software: you can redistribute it and/or modify
+   it under the terms of the GNU Lesser General Public License as published by
+   the Free Software Foundation, either version 3 of the License, or
+   (at your option) any later version.
+
+   plumed is distributed in the hope that it will be useful,
+   but WITHOUT ANY WARRANTY; without even the implied warranty of
+   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+   GNU Lesser General Public License for more details.
+
+   You should have received a copy of the GNU Lesser General Public License
+   along with plumed.  If not, see <http://www.gnu.org/licenses/>.
++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ */
+
+#include "MergeVectorTools.h"

src/tools/MergeVectorTools.h

@@ -0,0 +1,127 @@
+/* +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
+   Copyright (c) 2011-2023 The plumed team
+   (see the PEOPLE file at the root of the distribution for a list of names)
+
+   See http://www.plumed.org for more information.
+
+   This file is part of plumed, version 2.
+
+   plumed is free software: you can redistribute it and/or modify
+   it under the terms of the GNU Lesser General Public License as published by
+   the Free Software Foundation, either version 3 of the License, or
+   (at your option) any later version.
+
+   plumed is distributed in the hope that it will be useful,
+   but WITHOUT ANY WARRANTY; without even the implied warranty of
+   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+   GNU Lesser General Public License for more details.
+
+   You should have received a copy of the GNU Lesser General Public License
+   along with plumed.  If not, see <http://www.gnu.org/licenses/>.
++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ */
+#ifndef __PLUMED_tools_MergeVectorTools_h
+#define __PLUMED_tools_MergeVectorTools_h
+
+#include "small_vector/small_vector.h"
+#include <vector>
+#include <algorithm>
+
+namespace PLMD {
+
+namespace mergeVectorTools {
+
+/// Merge sorted vectors.
+/// Takes a vector of pointers to containers and merge them.
+/// Containers should be already sorted.
+/// The content is appended to the result vector.
+/// Optionally, uses a priority_queue implementation.
+template<class C>
+static void mergeSortedVectors(const C* const* vecs, std::size_t size, std::vector<typename C::value_type> & result) {
+
+  /// local class storing the range of remaining objects to be pushed
+  class Entry
+  {
+    typename C::const_iterator fwdIt,endIt;
+
+  public:
+    explicit Entry(C const& v) : fwdIt(v.begin()), endIt(v.end()) {}
+    /// check if this vector still contains something to be pushed
+    bool empty() const { return fwdIt == endIt; }
+    /// to allow using a priority_queu, which selects the highest element.
+    /// we here (counterintuitively) define < as >
+    bool operator< (Entry const& rhs) const { return top() > rhs.top(); }
+    const auto & top() const { return *fwdIt; }
+    void next() { ++fwdIt;};
+  };
+
+  // preprocessing
+  {
+    std::size_t maxsize=0;
+    for(unsigned i=0; i<size; i++) {
+      // find the largest
+      maxsize=std::max(maxsize,vecs[i]->size());
+    }
+    // this is just to decrease the number of reallocations on push_back
+    result.reserve(maxsize);
+    // if vectors are empty we are done
+    if(maxsize==0) return;
+  }
+
+  // start
+  // heap containing the ranges to be pushed
+  // avoid allocations when it's small
+  gch::small_vector<Entry,32> heap;
+
+  {
+    for(unsigned i=0; i<size; i++) {
+      if(!vecs[i]->empty()) {
+        // add entry at the end of the array
+        heap.emplace_back(Entry(*vecs[i]));
+      }
+    }
+    // make a sorted heap
+    std::make_heap(std::begin(heap),std::end(heap));
+  }
+
+  // first iteration, to avoid an extra "if" in main iteration
+  // explanations are below
+  {
+    std::pop_heap(std::begin(heap), std::end(heap));
+    auto & tmp=heap.back();
+    const auto val=tmp.top();
+    // here, we append inconditionally
+    result.push_back(val);
+    tmp.next();
+    if(tmp.empty()) heap.pop_back();
+    else std::push_heap(std::begin(heap), std::end(heap));
+  }
+
+  while(!heap.empty()) {
+    // move entry with the smallest element to the back of the array
+    std::pop_heap(std::begin(heap), std::end(heap));
+    // entry
+    auto & tmp=heap.back();
+    // element
+    const auto val=tmp.top();
+    // if the element is larger than the current largest element,
+    // push it to result
+    if(val > result.back()) result.push_back(val);
+    // move forward the used entry
+    tmp.next();
+    // if this entry is exhausted, remove it from the array
+    if(tmp.empty()) heap.pop_back();
+    // otherwise, sort again the array
+    else std::push_heap(std::begin(heap), std::end(heap));
+  }
+}
+
+template<class C>
+static void mergeSortedVectors(const std::vector<C*> & vecs, std::vector<typename C::value_type> & result) {
+  mergeSortedVectors(vecs.data(),vecs.size(),result);
+}
+}
+
+}
+
+#endif
+

src/tools/Tools.h

@@ -246,94 +246,6 @@ class Tools {
     set_to_zero(&vec[0](0,0),s*n*m);
   }
 
-
-
-
-  /// Merge sorted vectors.
-  /// Takes a vector of pointers to containers and merge them.
-  /// Containers should be already sorted.
-  /// The content is appended to the result vector.
-  /// Optionally, uses a priority_queue implementation.
-  template<class C>
-  static void mergeSortedVectors(const std::vector<C*> & vecs, std::vector<typename C::value_type> & result,bool priority_queue=false) {
-
-    /// local class storing the range of remaining objects to be pushed
-    struct Entry
-    {
-      typename C::const_iterator fwdIt,endIt;
-
-      explicit Entry(C const& v) : fwdIt(v.begin()), endIt(v.end()) {}
-      /// check if this vector still contains something to be pushed
-      explicit operator bool () const { return fwdIt != endIt; }
-      /// to allow using a priority_queu, which selects the highest element.
-      /// we here (counterintuitively) define < as >
-      bool operator< (Entry const& rhs) const { return *fwdIt > *rhs.fwdIt; }
-    };
-
-    if(priority_queue) {
-      std::priority_queue<Entry> queue;
-      // note: queue does not have reserve() method
-
-      // add vectors to the queue
-      {
-        std::size_t maxsize=0;
-        for(unsigned i=0; i<vecs.size(); i++) {
-          if(vecs[i]->size()>maxsize) maxsize=vecs[i]->size();
-          if(!vecs[i]->empty())queue.push(Entry(*vecs[i]));
-        }
-        // this is just to save multiple reallocations on push_back
-        result.reserve(maxsize);
-      }
-
-      // first iteration (to avoid a if in the main loop)
-      if(queue.empty()) return;
-      auto tmp=queue.top();
-      queue.pop();
-      result.push_back(*tmp.fwdIt);
-      tmp.fwdIt++;
-      if(tmp) queue.push(tmp);
-
-      // main loop
-      while(!queue.empty()) {
-        auto tmp=queue.top();
-        queue.pop();
-        if(result.back() < *tmp.fwdIt) result.push_back(*tmp.fwdIt);
-        tmp.fwdIt++;
-        if(tmp) queue.push(tmp);
-      }
-    } else {
-
-      std::vector<Entry> entries;
-      entries.reserve(vecs.size());
-
-      {
-        std::size_t maxsize=0;
-        for(int i=0; i<vecs.size(); i++) {
-          if(vecs[i]->size()>maxsize) maxsize=vecs[i]->size();
-          if(!vecs[i]->empty())entries.push_back(Entry(*vecs[i]));
-        }
-        // this is just to save multiple reallocations on push_back
-        result.reserve(maxsize);
-      }
-
-      while(!entries.empty()) {
-        // find smallest pending element
-        // we use max_element instead of min_element because we are defining < as > (see above)
-        const auto minval=*std::max_element(entries.begin(),entries.end())->fwdIt;
-
-        // push it
-        result.push_back(minval);
-
-        // fast forward vectors with elements equal to minval (to avoid duplicates)
-        for(auto & e : entries) while(e && *e.fwdIt==minval) ++e.fwdIt;
-
-        // remove from the entries vector all exhausted vectors
-        auto erase=std::remove_if(entries.begin(),entries.end(),[](const Entry & e) {return !e;});
-        entries.erase(erase,entries.end());
-      }
-    }
-
-  }
   static std::unique_ptr<std::lock_guard<std::mutex>> molfile_lock();
   /// Build a concatenated exception message.
   /// Should be called with an in-flight exception.