Reducing sorting overhead for small systems with PME (#5027)

platforms/common/include/openmm/common/CommonCalcNonbondedForce.h

@@ -47,7 +47,7 @@ namespace OpenMM {
 class CommonCalcNonbondedForceKernel : public CalcNonbondedForceKernel {
 public:
     CommonCalcNonbondedForceKernel(std::string name, const Platform& platform, ComputeContext& cc, const System& system) : CalcNonbondedForceKernel(name, platform),
-            hasInitializedKernel(false), cc(cc), pmeio(NULL) {
+            hasInitializedKernel(false), cc(cc), pmeio(NULL), stepsToSort(0) {
     }
     ~CommonCalcNonbondedForceKernel();
     /**
@@ -168,6 +168,7 @@ private:
     double ewaldSelfEnergy, dispersionCoefficient, alpha, dispersionAlpha, totalCharge;
     int gridSizeX, gridSizeY, gridSizeZ;
     int dispersionGridSizeX, dispersionGridSizeY, dispersionGridSizeZ;
+    int stepsToSort;
     bool usePmeQueue, deviceIsCpu, useFixedPointChargeSpreading, useCpuPme;
     bool hasCoulomb, hasLJ, doLJPME, usePosqCharges, recomputeParams, hasOffsets;
     NonbondedMethod nonbondedMethod;

platforms/common/src/CommonCalcNonbondedForce.cpp

@@ -960,6 +960,7 @@ double CommonCalcNonbondedForceKernel::execute(ContextImpl& context, bool includ
         // Execute the reciprocal space kernels.
 
         if (hasCoulomb) {
+            if (stepsToSort <= 0 || doLJPME || cc.getNumAtoms() > 15000) {
                 setPeriodicBoxArgs(cc, pmeGridIndexKernel, 2);
                 if (cc.getUseDoublePrecision()) {
                     pmeGridIndexKernel->setArg(7, recipBoxVectors[0]);
@@ -973,6 +974,10 @@ double CommonCalcNonbondedForceKernel::execute(ContextImpl& context, bool includ
                 }
                 pmeGridIndexKernel->execute(cc.getNumAtoms());
                 sort->sort(pmeAtomGridIndex);
+                stepsToSort = 3;
+            }
+            else
+                stepsToSort--;
             setPeriodicBoxArgs(cc, pmeSpreadChargeKernel, 2);
             if (cc.getUseDoublePrecision()) {
                 pmeSpreadChargeKernel->setArg(7, recipBoxVectors[0]);
@@ -1038,7 +1043,6 @@ double CommonCalcNonbondedForceKernel::execute(ContextImpl& context, bool includ
                 pmeDispersionGridIndexKernel->setArg(9, recipBoxVectorsFloat[2]);
             }
             pmeDispersionGridIndexKernel->execute(cc.getNumAtoms());
-            if (!hasCoulomb)
             sort->sort(pmeAtomGridIndex);
             if (useFixedPointChargeSpreading)
                 cc.clearBuffer(pmeGrid2);