Optimized computing long range correction for CustomNonbondedForce (#3236)

openmmapi/include/openmm/internal/CustomNonbondedForceImpl.h

@@ -9,7 +9,7 @@
  * Biological Structures at Stanford, funded under the NIH Roadmap for        *
  * Medical Research, grant U54 GM072970. See https://simtk.org.               *
  *                                                                            *
- * Portions copyright (c) 2008-2016 Stanford University and the Authors.      *
+ * Portions copyright (c) 2008-2021 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -35,6 +35,7 @@
 #include "ForceImpl.h"
 #include "openmm/CustomNonbondedForce.h"
 #include "openmm/Kernel.h"
+#include "openmm/internal/ThreadPool.h"
 #include "lepton/CompiledExpression.h"
 #include <utility>
 #include <map>
@@ -48,6 +49,7 @@ namespace OpenMM {
 
 class OPENMM_EXPORT CustomNonbondedForceImpl : public ForceImpl {
 public:
+    class LongRangeCorrectionData;
     CustomNonbondedForceImpl(const CustomNonbondedForce& owner);
     ~CustomNonbondedForceImpl();
     void initialize(ContextImpl& context);
@@ -61,12 +63,19 @@ public:
     std::map<std::string, double> getDefaultParameters();
     std::vector<std::string> getKernelNames();
     void updateParametersInContext(ContextImpl& context);
+    /**
+     * Prepare for computing the long range correction.  This function pre-computes anything
+     * that depends only on the Force (such as particle parameters) but not on information in
+     * the Context (such as global parameters).  This allows the coefficient to be updated
+     * more quickly when global parameters change.
+     */
+    static LongRangeCorrectionData prepareLongRangeCorrection(const CustomNonbondedForce& force);
     /**
      * Compute the coefficient which, when divided by the periodic box volume, gives the
      * long range correction to the energy.  If the Force computes parameter derivatives,
      * also compute the corresponding derivatives of the correction.
      */
-    static void calcLongRangeCorrection(const CustomNonbondedForce& force, const Context& context, double& coefficient, std::vector<double>& derivatives);
+    static void calcLongRangeCorrection(const CustomNonbondedForce& force, LongRangeCorrectionData& data, const Context& context, double& coefficient, std::vector<double>& derivatives, ThreadPool& threads);
 private:
     static double integrateInteraction(Lepton::CompiledExpression& expression, const std::vector<double>& params1, const std::vector<double>& params2,
             const CustomNonbondedForce& force, const Context& context, const std::vector<std::string>& paramNames);
@@ -74,6 +83,16 @@ private:
     Kernel kernel;
 };
 
+class CustomNonbondedForceImpl::LongRangeCorrectionData {
+public:
+    CustomNonbondedForce::NonbondedMethod method;
+    std::vector<std::vector<double> > classes;
+    std::vector<std::string> paramNames;
+    std::map<std::pair<int, int>, long long int> interactionCount;
+    Lepton::CompiledExpression energyExpression;
+    std::vector<Lepton::CompiledExpression> derivExpressions;
+};
+
 } // namespace OpenMM
 
 #endif /*OPENMM_CUSTOMNONBONDEDFORCEIMPL_H_*/

openmmapi/src/CustomNonbondedForceImpl.cpp

@@ -6,7 +6,7 @@
  * Biological Structures at Stanford, funded under the NIH Roadmap for        *
  * Medical Research, grant U54 GM072970. See https://simtk.org.               *
  *                                                                            *
- * Portions copyright (c) 2008-2016 Stanford University and the Authors.      *
+ * Portions copyright (c) 2008-2021 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -40,6 +40,7 @@
 #include "ReferenceTabulatedFunction.h"
 #include "lepton/ParsedExpression.h"
 #include "lepton/Parser.h"
+#include <atomic>
 #include <cmath>
 #include <sstream>
 #include <utility>
@@ -158,16 +159,15 @@ void CustomNonbondedForceImpl::updateParametersInContext(ContextImpl& context) {
     context.systemChanged();
 }
 
-void CustomNonbondedForceImpl::calcLongRangeCorrection(const CustomNonbondedForce& force, const Context& context, double& coefficient, vector<double>& derivatives) {
-    if (force.getNonbondedMethod() == CustomNonbondedForce::NoCutoff || force.getNonbondedMethod() == CustomNonbondedForce::CutoffNonPeriodic) {
-        coefficient = 0.0;
-        return;
-    }
+CustomNonbondedForceImpl::LongRangeCorrectionData CustomNonbondedForceImpl::prepareLongRangeCorrection(const CustomNonbondedForce& force) {
+    LongRangeCorrectionData data;
+    data.method = force.getNonbondedMethod();
+    if (data.method == CustomNonbondedForce::NoCutoff || data.method == CustomNonbondedForce::CutoffNonPeriodic)
+        return data;
     
     // Identify all particle classes (defined by parameters), and record the class of each particle.
     
     int numParticles = force.getNumParticles();
-    vector<vector<double> > classes;
     map<vector<double>, int> classIndex;
     vector<int> atomClass(numParticles);
     vector<double> parameters;
@@ -175,18 +175,17 @@ void CustomNonbondedForceImpl::calcLongRangeCorrection(const CustomNonbondedForc
         force.getParticleParameters(i, parameters);
         map<vector<double>, int>::iterator entry = classIndex.find(parameters);
         if (entry == classIndex.end()) {
-            classIndex[parameters] = classes.size();
-            atomClass[i] = classes.size();
-            classes.push_back(parameters);
+            classIndex[parameters] = data.classes.size();
+            atomClass[i] = data.classes.size();
+            data.classes.push_back(parameters);
         }
         else
             atomClass[i] = entry->second;
     }
-    int numClasses = classes.size();
+    int numClasses = data.classes.size();
     
     // Count the total number of particle pairs for each pair of classes.
     
-    map<pair<int, int>, long long int> interactionCount;
     if (force.getNumInteractionGroups() == 0) {
         // Count the particles of each class.
         
@@ -194,9 +193,9 @@ void CustomNonbondedForceImpl::calcLongRangeCorrection(const CustomNonbondedForc
         for (int i = 0; i < numParticles; i++)
             classCounts[atomClass[i]]++;
         for (int i = 0; i < numClasses; i++) {
-            interactionCount[make_pair(i, i)] = (classCounts[i]*(classCounts[i]+1))/2;
+            data.interactionCount[make_pair(i, i)] = (classCounts[i]*(classCounts[i]+1))/2;
             for (int j = i+1; j < numClasses; j++)
-                interactionCount[make_pair(i, j)] = classCounts[i]*classCounts[j];
+                data.interactionCount[make_pair(i, j)] = classCounts[i]*classCounts[j];
         }
     }
     else {
@@ -204,7 +203,7 @@ void CustomNonbondedForceImpl::calcLongRangeCorrection(const CustomNonbondedForc
         
         for (int i = 0; i < numClasses; i++) {
             for (int j = i; j < numClasses; j++)
-                interactionCount[make_pair(i, j)] = 0;
+                data.interactionCount[make_pair(i, j)] = 0;
         }
         
         // Loop over interaction groups and count the interactions in each one.
@@ -218,44 +217,80 @@ void CustomNonbondedForceImpl::calcLongRangeCorrection(const CustomNonbondedForc
                         continue;
                     int class1 = atomClass[*a1];
                     int class2 = atomClass[*a2];
-                    interactionCount[make_pair(min(class1, class2), max(class1, class2))]++;
+                    data.interactionCount[make_pair(min(class1, class2), max(class1, class2))]++;
                 }
         }
     }
     
-    // Compute the coefficient.
+    // Prepare for evaluating the expressions.
     
     map<string, Lepton::CustomFunction*> functions;
     for (int i = 0; i < force.getNumFunctions(); i++)
         functions[force.getTabulatedFunctionName(i)] = createReferenceTabulatedFunction(force.getTabulatedFunction(i));
-    double nPart = (double) numParticles;
-    double numInteractions = (nPart*(nPart+1))/2;
-    Lepton::CompiledExpression expression = Lepton::Parser::parse(force.getEnergyFunction(), functions).createCompiledExpression();
-    vector<string> paramNames;
+    data.energyExpression = Lepton::Parser::parse(force.getEnergyFunction(), functions).createCompiledExpression();
+    for (int k = 0; k < force.getNumEnergyParameterDerivatives(); k++)
+        data.derivExpressions.push_back(Lepton::Parser::parse(force.getEnergyFunction(), functions).differentiate(force.getEnergyParameterDerivativeName(k)).createCompiledExpression());
     for (int i = 0; i < force.getNumPerParticleParameters(); i++) {
         stringstream name1, name2;
         name1 << force.getPerParticleParameterName(i) << 1;
         name2 << force.getPerParticleParameterName(i) << 2;
-        paramNames.push_back(name1.str());
-        paramNames.push_back(name2.str());
+        data.paramNames.push_back(name1.str());
+        data.paramNames.push_back(name2.str());
+    }
+    return data;
+}
+
+void CustomNonbondedForceImpl::calcLongRangeCorrection(const CustomNonbondedForce& force, LongRangeCorrectionData& data, const Context& context, double& coefficient, vector<double>& derivatives, ThreadPool& threads) {
+    if (data.method == CustomNonbondedForce::NoCutoff || data.method == CustomNonbondedForce::CutoffNonPeriodic) {
+        coefficient = 0.0;
+        return;
     }
+
+    // Compute the coefficient.  Use multiple threads to compute the integrals in parallel.
+
+    int numClasses = data.classes.size();
+    double nPart = (double) context.getSystem().getNumParticles();
+    double numInteractions = (nPart*(nPart+1))/2;
+    vector<double> threadSum(threads.getNumThreads(), 0.0);
+    atomic<int> atomicCounter(0);
+    threads.execute([&] (ThreadPool& threads, int threadIndex) {
+        Lepton::CompiledExpression expression = data.energyExpression;
+        while (true) {
+            int i = atomicCounter++;
+            if (i >= numClasses)
+                break;
+            for (int j = i; j < numClasses; j++)
+                threadSum[threadIndex] += data.interactionCount.at(make_pair(i, j))*integrateInteraction(expression, data.classes[i], data.classes[j], force, context, data.paramNames);
+        }
+    });
+    threads.waitForThreads();
     double sum = 0;
-    for (int i = 0; i < numClasses; i++)
-        for (int j = i; j < numClasses; j++)
-            sum += interactionCount[make_pair(i, j)]*integrateInteraction(expression, classes[i], classes[j], force, context, paramNames);
+    for (int i = 0; i < threadSum.size(); i++)
+        sum += threadSum[i];
     sum /= numInteractions;
     coefficient = 2*M_PI*nPart*nPart*sum;
     
     // Now do the same for parameter derivatives.
     
-    int numDerivs = force.getNumEnergyParameterDerivatives();
+    int numDerivs = data.derivExpressions.size();
     derivatives.resize(numDerivs);
     for (int k = 0; k < numDerivs; k++) {
-        expression = Lepton::Parser::parse(force.getEnergyFunction(), functions).differentiate(force.getEnergyParameterDerivativeName(k)).createCompiledExpression();
+        atomicCounter = 0;
+        threads.execute([&] (ThreadPool& threads, int threadIndex) {
+            threadSum[threadIndex] = 0;
+            Lepton::CompiledExpression expression = data.derivExpressions[k];
+            while (true) {
+                int i = atomicCounter++;
+                if (i >= numClasses)
+                    break;
+                for (int j = i; j < numClasses; j++)
+                    threadSum[threadIndex] += data.interactionCount.at(make_pair(i, j))*integrateInteraction(expression, data.classes[i], data.classes[j], force, context, data.paramNames);
+            }
+        });
+        threads.waitForThreads();
         sum = 0;
-        for (int i = 0; i < numClasses; i++)
-            for (int j = i; j < numClasses; j++)
-                sum += interactionCount[make_pair(i, j)]*integrateInteraction(expression, classes[i], classes[j], force, context, paramNames);
+        for (int i = 0; i < threadSum.size(); i++)
+            sum += threadSum[i];
         sum /= numInteractions;
         derivatives[k] = 2*M_PI*nPart*nPart*sum;
     }

platforms/common/include/openmm/common/CommonKernels.h

@@ -9,7 +9,7 @@
  * Biological Structures at Stanford, funded under the NIH Roadmap for        *
  * Medical Research, grant U54 GM072970. See https://simtk.org.               *
  *                                                                            *
- * Portions copyright (c) 2008-2020 Stanford University and the Authors.      *
+ * Portions copyright (c) 2008-2021 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -34,6 +34,7 @@
 #include "openmm/kernels.h"
 #include "openmm/internal/CompiledExpressionSet.h"
 #include "openmm/internal/CustomIntegratorUtilities.h"
+#include "openmm/internal/CustomNonbondedForceImpl.h"
 #include "lepton/CompiledExpression.h"
 #include "lepton/ExpressionProgram.h"
 
@@ -616,6 +617,8 @@ public:
     void copyParametersToContext(ContextImpl& context, const CustomNonbondedForce& force);
 private:
     class ForceInfo;
+    class LongRangePostComputation;
+    class LongRangeTask;
     void initInteractionGroups(const CustomNonbondedForce& force, const std::string& interactionSource, const std::vector<std::string>& tableTypes);
     ComputeContext& cc;
     ForceInfo* info;
@@ -631,6 +634,7 @@ private:
     bool hasInitializedLongRangeCorrection, hasInitializedKernel, hasParamDerivs, useNeighborList;
     int numGroupThreadBlocks;
     CustomNonbondedForce* forceCopy;
+    CustomNonbondedForceImpl::LongRangeCorrectionData longRangeCorrectionData;
     const System& system;
 };
 

platforms/common/src/CommonKernels.cpp

@@ -34,7 +34,6 @@
 #include "openmm/internal/CustomCompoundBondForceImpl.h"
 #include "openmm/internal/CustomHbondForceImpl.h"
 #include "openmm/internal/CustomManyParticleForceImpl.h"
-#include "openmm/internal/CustomNonbondedForceImpl.h"
 #include "CommonKernelSources.h"
 #include "lepton/CustomFunction.h"
 #include "lepton/ExpressionTreeNode.h"
@@ -1771,6 +1770,47 @@ private:
     vector<set<int> > groupsForParticle;
 };
 
+class CommonCalcCustomNonbondedForceKernel::LongRangePostComputation : public ComputeContext::ForcePostComputation {
+public:
+    LongRangePostComputation(ComputeContext& cc, double& longRangeCoefficient, vector<double>& longRangeCoefficientDerivs, CustomNonbondedForce* force) :
+            cc(cc), longRangeCoefficient(longRangeCoefficient), longRangeCoefficientDerivs(longRangeCoefficientDerivs), force(force) {
+    }
+    double computeForceAndEnergy(bool includeForces, bool includeEnergy, int groups) {
+        cc.getWorkThread().flush();
+        Vec3 a, b, c;
+        cc.getPeriodicBoxVectors(a, b, c);
+        double volume = a[0]*b[1]*c[2];
+        map<string, double>& derivs = cc.getEnergyParamDerivWorkspace();
+        for (int i = 0; i < longRangeCoefficientDerivs.size(); i++)
+            derivs[force->getEnergyParameterDerivativeName(i)] += longRangeCoefficientDerivs[i]/volume;
+        return longRangeCoefficient/volume;
+    }
+private:
+    ComputeContext& cc;
+    double& longRangeCoefficient;
+    vector<double>& longRangeCoefficientDerivs;
+    CustomNonbondedForce* force;
+};
+
+class CommonCalcCustomNonbondedForceKernel::LongRangeTask : public ComputeContext::WorkTask {
+public:
+    LongRangeTask(ComputeContext& cc, Context& context, CustomNonbondedForceImpl::LongRangeCorrectionData& data,
+                  double& longRangeCoefficient, vector<double>& longRangeCoefficientDerivs, CustomNonbondedForce* force) :
+                        cc(cc), context(context), data(data), longRangeCoefficient(longRangeCoefficient),
+                        longRangeCoefficientDerivs(longRangeCoefficientDerivs), force(force) {
+    }
+    void execute() {
+        CustomNonbondedForceImpl::calcLongRangeCorrection(*force, data, context, longRangeCoefficient, longRangeCoefficientDerivs, cc.getThreadPool());
+    }
+private:
+    ComputeContext& cc;
+    Context& context;
+    CustomNonbondedForceImpl::LongRangeCorrectionData& data;
+    double& longRangeCoefficient;
+    vector<double>& longRangeCoefficientDerivs;
+    CustomNonbondedForce* force;
+};
+
 CommonCalcCustomNonbondedForceKernel::~CommonCalcCustomNonbondedForceKernel() {
     cc.setAsCurrent();
     if (params != NULL)
@@ -1912,6 +1952,8 @@ void CommonCalcCustomNonbondedForceKernel::initialize(const System& system, cons
     
     if (force.getNonbondedMethod() == CustomNonbondedForce::CutoffPeriodic && force.getUseLongRangeCorrection() && cc.getContextIndex() == 0) {
         forceCopy = new CustomNonbondedForce(force);
+        longRangeCorrectionData = CustomNonbondedForceImpl::prepareLongRangeCorrection(force);
+        cc.addPostComputation(new LongRangePostComputation(cc, longRangeCoefficient, longRangeCoefficientDerivs, forceCopy));
         hasInitializedLongRangeCorrection = false;
     }
     else {
@@ -2199,6 +2241,7 @@ double CommonCalcCustomNonbondedForceKernel::execute(ContextImpl& context, bool
         // When using a neighbor list, run the whole calculation on a single device.
         return 0.0;
     }
+    bool recomputeLongRangeCorrection = !hasInitializedLongRangeCorrection;
     if (globals.isInitialized()) {
         bool changed = false;
         for (int i = 0; i < (int) globalParamNames.size(); i++) {
@@ -2209,14 +2252,12 @@ double CommonCalcCustomNonbondedForceKernel::execute(ContextImpl& context, bool
         }
         if (changed) {
             globals.upload(globalParamValues);
-            if (forceCopy != NULL) {
-                CustomNonbondedForceImpl::calcLongRangeCorrection(*forceCopy, context.getOwner(), longRangeCoefficient, longRangeCoefficientDerivs);
-                hasInitializedLongRangeCorrection = true;
-            }
+            if (forceCopy != NULL)
+                recomputeLongRangeCorrection = true;
         }
     }
-    if (!hasInitializedLongRangeCorrection) {
-        CustomNonbondedForceImpl::calcLongRangeCorrection(*forceCopy, context.getOwner(), longRangeCoefficient, longRangeCoefficientDerivs);
+    if (recomputeLongRangeCorrection) {
+        cc.getWorkThread().addTask(new LongRangeTask(cc, context.getOwner(), longRangeCorrectionData, longRangeCoefficient, longRangeCoefficientDerivs, forceCopy));
         hasInitializedLongRangeCorrection = true;
     }
     if (interactionGroupData.isInitialized()) {
@@ -2264,13 +2305,7 @@ double CommonCalcCustomNonbondedForceKernel::execute(ContextImpl& context, bool
         setPeriodicBoxArgs(cc, interactionGroupKernel, 6);
         interactionGroupKernel->execute(numGroupThreadBlocks*forceThreadBlockSize, forceThreadBlockSize);
     }
-    Vec3 a, b, c;
-    cc.getPeriodicBoxVectors(a, b, c);
-    double volume = a[0]*b[1]*c[2];
-    map<string, double>& derivs = cc.getEnergyParamDerivWorkspace();
-    for (int i = 0; i < longRangeCoefficientDerivs.size(); i++)
-        derivs[forceCopy->getEnergyParameterDerivativeName(i)] += longRangeCoefficientDerivs[i]/volume;
-    return longRangeCoefficient/volume;
+    return 0;
 }
 
 void CommonCalcCustomNonbondedForceKernel::copyParametersToContext(ContextImpl& context, const CustomNonbondedForce& force) {
@@ -2296,8 +2331,9 @@ void CommonCalcCustomNonbondedForceKernel::copyParametersToContext(ContextImpl&
     // If necessary, recompute the long range correction.
     
     if (forceCopy != NULL) {
-        CustomNonbondedForceImpl::calcLongRangeCorrection(force, context.getOwner(), longRangeCoefficient, longRangeCoefficientDerivs);
-        hasInitializedLongRangeCorrection = true;
+        longRangeCorrectionData = CustomNonbondedForceImpl::prepareLongRangeCorrection(force);
+        CustomNonbondedForceImpl::calcLongRangeCorrection(force, longRangeCorrectionData, context.getOwner(), longRangeCoefficient, longRangeCoefficientDerivs, cc.getThreadPool());
+        hasInitializedLongRangeCorrection = false;
         *forceCopy = force;
     }
     

platforms/cpu/include/CpuKernels.h

@@ -9,7 +9,7 @@
  * Biological Structures at Stanford, funded under the NIH Roadmap for        *
  * Medical Research, grant U54 GM072970. See https://simtk.org.               *
  *                                                                            *
- * Portions copyright (c) 2013-2020 Stanford University and the Authors.      *
+ * Portions copyright (c) 2013-2021 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -45,6 +45,7 @@
 #include "CpuPlatform.h"
 #include "openmm/kernels.h"
 #include "openmm/System.h"
+#include "openmm/internal/CustomNonbondedForceImpl.h"
 #include <array>
 #include <tuple>
 
@@ -326,6 +327,7 @@ private:
     double nonbondedCutoff, switchingDistance, periodicBoxSize[3], longRangeCoefficient;
     bool useSwitchingFunction, hasInitializedLongRangeCorrection;
     CustomNonbondedForce* forceCopy;
+    CustomNonbondedForceImpl::LongRangeCorrectionData longRangeCorrectionData;
     std::map<std::string, double> globalParamValues;
     std::vector<std::set<int> > exclusions;
     std::vector<std::string> parameterNames, globalParameterNames, energyParamDerivNames;

platforms/cpu/src/CpuKernels.cpp

@@ -6,7 +6,7 @@
  * Biological Structures at Stanford, funded under the NIH Roadmap for        *
  * Medical Research, grant U54 GM072970. See https://simtk.org.               *
  *                                                                            *
- * Portions copyright (c) 2013-2020 Stanford University and the Authors.      *
+ * Portions copyright (c) 2013-2021 Stanford University and the Authors.      *
  * Authors: Peter Eastman                                                     *
  * Contributors:                                                              *
  *                                                                            *
@@ -43,7 +43,6 @@
 #include "openmm/OpenMMException.h"
 #include "openmm/Vec3.h"
 #include "openmm/internal/ContextImpl.h"
-#include "openmm/internal/CustomNonbondedForceImpl.h"
 #include "openmm/internal/NonbondedForceImpl.h"
 #include "openmm/internal/vectorize.h"
 #include "lepton/CompiledExpression.h"
@@ -975,10 +974,13 @@ double CpuCalcCustomNonbondedForceKernel::execute(ContextImpl& context, bool inc
     
     // Add in the long range correction.
     
-    if (!hasInitializedLongRangeCorrection || (globalParamsChanged && forceCopy != NULL)) {
-        CustomNonbondedForceImpl::calcLongRangeCorrection(*forceCopy, context.getOwner(), longRangeCoefficient, longRangeCoefficientDerivs);
+    if (!hasInitializedLongRangeCorrection) {
+        longRangeCorrectionData = CustomNonbondedForceImpl::prepareLongRangeCorrection(*forceCopy);
+        CustomNonbondedForceImpl::calcLongRangeCorrection(*forceCopy, longRangeCorrectionData, context.getOwner(), longRangeCoefficient, longRangeCoefficientDerivs, data.threads);
         hasInitializedLongRangeCorrection = true;
     }
+    else if (globalParamsChanged && forceCopy != NULL)
+        CustomNonbondedForceImpl::calcLongRangeCorrection(*forceCopy, longRangeCorrectionData, context.getOwner(), longRangeCoefficient, longRangeCoefficientDerivs, data.threads);
     double volume = boxVectors[0][0]*boxVectors[1][1]*boxVectors[2][2];
     energy += longRangeCoefficient/volume;
     for (int i = 0; i < longRangeCoefficientDerivs.size(); i++)
@@ -1004,7 +1006,8 @@ void CpuCalcCustomNonbondedForceKernel::copyParametersToContext(ContextImpl& con
     // If necessary, recompute the long range correction.
     
     if (forceCopy != NULL) {
-        CustomNonbondedForceImpl::calcLongRangeCorrection(force, context.getOwner(), longRangeCoefficient, longRangeCoefficientDerivs);
+        longRangeCorrectionData = CustomNonbondedForceImpl::prepareLongRangeCorrection(force);
+        CustomNonbondedForceImpl::calcLongRangeCorrection(force, longRangeCorrectionData, context.getOwner(), longRangeCoefficient, longRangeCoefficientDerivs, data.threads);
         hasInitializedLongRangeCorrection = true;
         *forceCopy = force;
     }

platforms/reference/include/ReferenceKernels.h

@@ -34,6 +34,7 @@
 
 #include "ReferencePlatform.h"
 #include "openmm/kernels.h"
+#include "openmm/internal/CustomNonbondedForceImpl.h"
 #include "SimTKOpenMMRealType.h"
 #include "ReferenceNeighborList.h"
 #include "lepton/CompiledExpression.h"
@@ -686,6 +687,7 @@ private:
     double nonbondedCutoff, switchingDistance, periodicBoxSize[3], longRangeCoefficient;
     bool useSwitchingFunction, hasInitializedLongRangeCorrection;
     CustomNonbondedForce* forceCopy;
+    CustomNonbondedForceImpl::LongRangeCorrectionData longRangeCorrectionData;
     std::map<std::string, double> globalParamValues;
     std::vector<std::set<int> > exclusions;
     Lepton::CompiledExpression energyExpression, forceExpression;

platforms/reference/src/ReferenceKernels.cpp

@@ -76,7 +76,6 @@
 #include "openmm/internal/CustomCentroidBondForceImpl.h"
 #include "openmm/internal/CustomCompoundBondForceImpl.h"
 #include "openmm/internal/CustomHbondForceImpl.h"
-#include "openmm/internal/CustomNonbondedForceImpl.h"
 #include "openmm/internal/CMAPTorsionForceImpl.h"
 #include "openmm/internal/NonbondedForceImpl.h"
 #include "openmm/Integrator.h"
@@ -1260,10 +1259,16 @@ double ReferenceCalcCustomNonbondedForceKernel::execute(ContextImpl& context, bo
     
     // Add in the long range correction.
     
-    if (!hasInitializedLongRangeCorrection || (globalParamsChanged && forceCopy != NULL)) {
-        CustomNonbondedForceImpl::calcLongRangeCorrection(*forceCopy, context.getOwner(), longRangeCoefficient, longRangeCoefficientDerivs);
+    if (!hasInitializedLongRangeCorrection) {
+        longRangeCorrectionData = CustomNonbondedForceImpl::prepareLongRangeCorrection(*forceCopy);
+        ThreadPool threads;
+        CustomNonbondedForceImpl::calcLongRangeCorrection(*forceCopy, longRangeCorrectionData, context.getOwner(), longRangeCoefficient, longRangeCoefficientDerivs, threads);
         hasInitializedLongRangeCorrection = true;
     }
+    else if (globalParamsChanged && forceCopy != NULL) {
+        ThreadPool threads;
+        CustomNonbondedForceImpl::calcLongRangeCorrection(*forceCopy, longRangeCorrectionData, context.getOwner(), longRangeCoefficient, longRangeCoefficientDerivs, threads);
+    }
     double volume = boxVectors[0][0]*boxVectors[1][1]*boxVectors[2][2];
     energy += longRangeCoefficient/volume;
     for (int i = 0; i < longRangeCoefficientDerivs.size(); i++)
@@ -1289,7 +1294,9 @@ void ReferenceCalcCustomNonbondedForceKernel::copyParametersToContext(ContextImp
     // If necessary, recompute the long range correction.
     
     if (forceCopy != NULL) {
-        CustomNonbondedForceImpl::calcLongRangeCorrection(force, context.getOwner(), longRangeCoefficient, longRangeCoefficientDerivs);
+        longRangeCorrectionData = CustomNonbondedForceImpl::prepareLongRangeCorrection(force);
+        ThreadPool threads;
+        CustomNonbondedForceImpl::calcLongRangeCorrection(force, longRangeCorrectionData, context.getOwner(), longRangeCoefficient, longRangeCoefficientDerivs, threads);
         hasInitializedLongRangeCorrection = true;
         *forceCopy = force;
     }