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