Improved logic for deciding when to reorder atoms (#3721)

platforms/common/include/openmm/common/ComputeContext.h

@@ -248,8 +248,18 @@ public:
     /**
      * Reorder the internal arrays of atoms to try to keep spatially contiguous atoms close
      * together in the arrays.
+     * 
+     * Calling this method might or might not actually change the atom order.  It uses
+     * internal heuristics to decide when and how often to update the order.  If you
+     * want to guarantee that reordering will definitely be done, call forceReorder() before
+     * calling this.
      */
     void reorderAtoms();
+    /**
+     * Calling this method guarantees that the next call to reorderAtoms() will actually
+     * perform reordering.
+     */
+    void forceReorder();
     /**
      * Add a listener that should be called whenever atoms get reordered.  The OpenCLContext
      * assumes ownership of the object, and deletes it when the context itself is deleted.
@@ -507,7 +517,7 @@ protected:
     double time;
     int numAtoms, paddedNumAtoms, computeForceCount, stepsSinceReorder;
     long long stepCount;
-    bool atomsWereReordered, forcesValid;
+    bool forceNextReorder, atomsWereReordered, forcesValid;
     std::vector<ComputeForceInfo*> forces;
     std::vector<Molecule> molecules;
     std::vector<MoleculeGroup> moleculeGroups;

platforms/common/src/ComputeContext.cpp

@@ -44,7 +44,7 @@ const int ComputeContext::ThreadBlockSize = 64;
 const int ComputeContext::TileSize = 32;
 
 ComputeContext::ComputeContext(const System& system) : system(system), time(0.0), stepCount(0), computeForceCount(0), stepsSinceReorder(99999),
-        atomsWereReordered(false), forcesValid(false), thread(NULL) {
+        forceNextReorder(false), atomsWereReordered(false), forcesValid(false), thread(NULL) {
     thread = new WorkThread();
 }
 
@@ -439,16 +439,22 @@ bool ComputeContext::invalidateMolecules(ComputeForceInfo* force) {
     findMoleculeGroups();
     for (auto listener : reorderListeners)
         listener->execute();
+    forceNextReorder = true;
     reorderAtoms();
     return true;
 }
 
+void ComputeContext::forceReorder() {
+    forceNextReorder = true;
+}
+
 void ComputeContext::reorderAtoms() {
     atomsWereReordered = false;
-    if (numAtoms == 0 || !getNonbondedUtilities().getUseCutoff() || stepsSinceReorder < 250) {
+    if (numAtoms == 0 || !getNonbondedUtilities().getUseCutoff() || (stepsSinceReorder < 250 && !forceNextReorder)) {
         stepsSinceReorder++;
         return;
     }
+    forceNextReorder = false;
     atomsWereReordered = true;
     stepsSinceReorder = 0;
     if (getUseDoublePrecision())

platforms/cuda/include/CudaNonbondedUtilities.h

@@ -349,7 +349,7 @@ private:
     double lastCutoff;
     bool useCutoff, usePeriodic, anyExclusions, usePadding, forceRebuildNeighborList, canUsePairList;
     int startTileIndex, startBlockIndex, numBlocks, maxExclusions, numForceThreadBlocks, forceThreadBlockSize, numAtoms, groupFlags;
-    unsigned int maxTiles, maxSinglePairs;
+    unsigned int maxTiles, maxSinglePairs, tilesAfterReorder;
     long long numTiles;
     std::string kernelSource;
 };

platforms/cuda/src/CudaNonbondedUtilities.cpp

@@ -423,6 +423,10 @@ void CudaNonbondedUtilities::computeInteractions(int forceGroups, bool includeFo
 bool CudaNonbondedUtilities::updateNeighborListSize() {
     if (!useCutoff)
         return false;
+    if (context.getStepsSinceReorder() == 0)
+        tilesAfterReorder = pinnedCountBuffer[0];
+    else if (context.getStepsSinceReorder() > 25 && pinnedCountBuffer[0] > 1.1*tilesAfterReorder)
+        context.forceReorder();
     if (pinnedCountBuffer[0] <= maxTiles && pinnedCountBuffer[1] <= maxSinglePairs)
         return false;
 

platforms/opencl/include/OpenCLNonbondedUtilities.h

@@ -333,6 +333,7 @@ private:
     bool useCutoff, usePeriodic, deviceIsCpu, anyExclusions, usePadding, forceRebuildNeighborList;
     int numForceBuffers, startTileIndex, startBlockIndex, numBlocks, maxExclusions, numForceThreadBlocks;
     int forceThreadBlockSize, interactingBlocksThreadBlockSize, groupFlags;
+    unsigned int tilesAfterReorder;
     long long numTiles;
     std::string kernelSource;
 };

platforms/opencl/src/OpenCLNonbondedUtilities.cpp

@@ -395,6 +395,10 @@ void OpenCLNonbondedUtilities::computeInteractions(int forceGroups, bool include
 bool OpenCLNonbondedUtilities::updateNeighborListSize() {
     if (!useCutoff)
         return false;
+    if (context.getStepsSinceReorder() == 0)
+        tilesAfterReorder = pinnedCountMemory[0];
+    else if (context.getStepsSinceReorder() > 25 && pinnedCountMemory[0] > 1.1*tilesAfterReorder)
+        context.forceReorder();
     if (pinnedCountMemory[0] <= interactingTiles.getSize())
         return false;