Parallelized CpuCustomGBForce

platforms/cpu/include/CpuCustomGBForce.h

@@ -29,6 +29,7 @@
 #include "CpuNeighborList.h"
 #include "lepton/CompiledExpression.h"
 #include "openmm/CustomGBForce.h"
+#include "openmm/internal/ThreadPool.h"
 #include "openmm/internal/vectorize.h"
 #include <map>
 #include <set>
@@ -38,59 +39,63 @@ namespace OpenMM {
 
 class CpuCustomGBForce {
 private:
+    class ComputeForceTask;
+    class ThreadData;
+
     bool cutoff;
     bool periodic;
     const CpuNeighborList* neighborList;
     float periodicBoxSize[3];
     float cutoffDistance;
-    CompiledExpressionSet expressionSet;
-    std::vector<Lepton::CompiledExpression> valueExpressions;
-    std::vector<std::vector<Lepton::CompiledExpression> > valueDerivExpressions;
-    std::vector<std::vector<Lepton::CompiledExpression> > valueGradientExpressions;
+    const std::vector<std::set<int> > exclusions;
     std::vector<std::string> valueNames;
-    std::vector<int> valueIndex;
     std::vector<CustomGBForce::ComputationType> valueTypes;
-    std::vector<Lepton::CompiledExpression> energyExpressions;
-    std::vector<std::vector<Lepton::CompiledExpression> > energyDerivExpressions;
-    std::vector<std::vector<Lepton::CompiledExpression> > energyGradientExpressions;
     std::vector<std::string> paramNames;
-    std::vector<int> paramIndex;
     std::vector<CustomGBForce::ComputationType> energyTypes;
-    std::vector<int> particleParamIndex;
-    std::vector<int> particleValueIndex;
-    int xindex, yindex, zindex, rindex;
+    ThreadPool& threads;
+    std::vector<ThreadData*> threadData;
+    std::vector<double> threadEnergy;
     // Workspace vectors
     std::vector<std::vector<float> > values, dEdV;
-    std::vector<float> dVdR1, dVdR2, dVdX, dVdY, dVdZ;
+    // The following variables are used to make information accessible to the individual threads.
+    int numberOfAtoms;
+    float* posq;
+    RealOpenMM** atomParameters;
+    const std::map<std::string, double>* globalParameters;
+    std::vector<AlignedArray<float> >* threadForce;
+    bool includeForce, includeEnergy;
+    void* atomicCounter;
+    
+    /**
+     * This routine contains the code executed by each thread.
+     */
+    void threadComputeForce(ThreadPool& threads, int threadIndex);
     
     /**
      * Calculate a computed value of type SingleParticle
      * 
      * @param index            the index of the value to compute
+     * @param data             workspace for the current thread
      * @param numAtoms         number of atoms
      * @param posq             atom coordinates
-     * @param values           the vector to store computed values into
      * @param atomParameters   atomParameters[atomIndex][paramterIndex]
      */
 
-    void calculateSingleParticleValue(int index, int numAtoms, float* posq, std::vector<std::vector<float> >& values,
-                                      RealOpenMM** atomParameters);
+    void calculateSingleParticleValue(int index, ThreadData& data, int numAtoms, float* posq, RealOpenMM** atomParameters);
 
     /**
      * Calculate a computed value that is based on particle pairs
      * 
      * @param index            the index of the value to compute
+     * @param data             workspace for the current thread
      * @param numAtoms         number of atoms
      * @param posq             atom coordinates
      * @param atomParameters   atomParameters[atomIndex][paramterIndex]
-     * @param values           the vector to store computed values into
-     * @param exclusions       exclusions[i] is the set of excluded indices for atom i
      * @param useExclusions    specifies whether to use exclusions
      */
 
-    void calculateParticlePairValue(int index, int numAtoms, float* posq, RealOpenMM** atomParameters,
-                                    std::vector<std::vector<float> >& values,
-                                    const std::vector<std::set<int> >& exclusions, bool useExclusions, const fvec4& boxSize, const fvec4& invBoxSize);
+    void calculateParticlePairValue(int index, ThreadData& data, int numAtoms, float* posq, RealOpenMM** atomParameters,
+                                    bool useExclusions, const fvec4& boxSize, const fvec4& invBoxSize);
 
     /**
      * Evaluate a single atom pair as part of calculating a computed value
@@ -98,51 +103,43 @@ private:
      * @param index            the index of the value to compute
      * @param atom1            the index of the first atom in the pair
      * @param atom2            the index of the second atom in the pair
+     * @param data             workspace for the current thread
      * @param posq             atom coordinates
      * @param atomParameters   atomParameters[atomIndex][paramterIndex]
-     * @param values           the vector to store computed values into
      */
 
-    void calculateOnePairValue(int index, int atom1, int atom2, float* posq, RealOpenMM** atomParameters,
-                               std::vector<std::vector<float> >& values, const fvec4& boxSize, const fvec4& invBoxSize);
+    void calculateOnePairValue(int index, int atom1, int atom2, ThreadData& data, float* posq, RealOpenMM** atomParameters,
+                               std::vector<float>& valueArray, const fvec4& boxSize, const fvec4& invBoxSize);
 
     /**
      * Calculate an energy term of type SingleParticle
      * 
      * @param index            the index of the value to compute
+     * @param data             workspace for the current thread
      * @param numAtoms         number of atoms
      * @param posq             atom coordinates
-     * @param values           the vector containing computed values
      * @param atomParameters   atomParameters[atomIndex][paramterIndex]
      * @param forces           forces on atoms are added to this
      * @param totalEnergy      the energy contribution is added to this
-     * @param dEdV             the derivative of energy with respect to computed values is stored in this
      */
 
-    void calculateSingleParticleEnergyTerm(int index, int numAtoms, float* posq, const std::vector<std::vector<float> >& values,
-                                      RealOpenMM** atomParameters, float* forces,
-                                      double* totalEnergy, std::vector<std::vector<float> >& dEdV);
+    void calculateSingleParticleEnergyTerm(int index, ThreadData& data, int numAtoms, float* posq, RealOpenMM** atomParameters, float* forces, double& totalEnergy);
 
     /**
      * Calculate an energy term that is based on particle pairs
      * 
      * @param index            the index of the term to compute
+     * @param data             workspace for the current thread
      * @param numAtoms         number of atoms
      * @param posq             atom coordinates
      * @param atomParameters   atomParameters[atomIndex][paramterIndex]
-     * @param values           the vector containing computed values
-     * @param exclusions       exclusions[i] is the set of excluded indices for atom i
      * @param useExclusions    specifies whether to use exclusions
      * @param forces           forces on atoms are added to this
      * @param totalEnergy      the energy contribution is added to this
-     * @param dEdV             the derivative of energy with respect to computed values is stored in this
      */
 
-    void calculateParticlePairEnergyTerm(int index, int numAtoms, float* posq, RealOpenMM** atomParameters,
-                                    const std::vector<std::vector<float> >& values,
-                                    const std::vector<std::set<int> >& exclusions, bool useExclusions,
-                                    float* forces, double* totalEnergy, std::vector<std::vector<float> >& dEdV,
-                                    const fvec4& boxSize, const fvec4& invBoxSize);
+    void calculateParticlePairEnergyTerm(int index, ThreadData& data, int numAtoms, float* posq, RealOpenMM** atomParameters,
+                                    bool useExclusions, float* forces, double& totalEnergy, const fvec4& boxSize, const fvec4& invBoxSize);
 
     /**
      * Evaluate a single atom pair as part of calculating an energy term
@@ -150,54 +147,43 @@ private:
      * @param index            the index of the term to compute
      * @param atom1            the index of the first atom in the pair
      * @param atom2            the index of the second atom in the pair
+     * @param data             workspace for the current thread
      * @param posq             atom coordinates
      * @param atomParameters   atomParameters[atomIndex][paramterIndex]
-     * @param values           the vector containing computed values
      * @param forces           forces on atoms are added to this
      * @param totalEnergy      the energy contribution is added to this
-     * @param dEdV             the derivative of energy with respect to computed values is stored in this
      */
 
-    void calculateOnePairEnergyTerm(int index, int atom1, int atom2, float* posq, RealOpenMM** atomParameters,
-                               const std::vector<std::vector<float> >& values,
-                               float* forces, double* totalEnergy, std::vector<std::vector<float> >& dEdV,
-                               const fvec4& boxSize, const fvec4& invBoxSize);
+    void calculateOnePairEnergyTerm(int index, int atom1, int atom2, ThreadData& data, float* posq, RealOpenMM** atomParameters,
+                               float* forces, double& totalEnergy, const fvec4& boxSize, const fvec4& invBoxSize);
 
     /**
      * Apply the chain rule to compute forces on atoms
      * 
+     * @param data             workspace for the current thread
      * @param numAtoms         number of atoms
      * @param posq             atom coordinates
      * @param atomParameters   atomParameters[atomIndex][paramterIndex]
-     * @param values           the vector containing computed values
-     * @param exclusions       exclusions[i] is the set of excluded indices for atom i
      * @param forces           forces on atoms are added to this
-     * @param dEdV             the derivative of energy with respect to computed values is stored in this
      */
 
-    void calculateChainRuleForces(int numAtoms, float* posq, RealOpenMM** atomParameters,
-                                    const std::vector<std::vector<float> >& values,
-                                    const std::vector<std::set<int> >& exclusions,
-                                    float* forces, std::vector<std::vector<float> >& dEdV,
-                                    const fvec4& boxSize, const fvec4& invBoxSize);
+    void calculateChainRuleForces(ThreadData& data, int numAtoms, float* posq, RealOpenMM** atomParameters,
+                                    float* forces, const fvec4& boxSize, const fvec4& invBoxSize);
 
     /**
      * Evaluate a single atom pair as part of applying the chain rule
      * 
      * @param atom1            the index of the first atom in the pair
      * @param atom2            the index of the second atom in the pair
+     * @param data             workspace for the current thread
      * @param posq             atom coordinates
      * @param atomParameters   atomParameters[atomIndex][paramterIndex]
-     * @param values           the vector containing computed values
      * @param forces           forces on atoms are added to this
-     * @param dEdV             the derivative of energy with respect to computed values is stored in this
      * @param isExcluded       specifies whether this is an excluded pair
      */
 
-    void calculateOnePairChainRule(int atom1, int atom2, float* posq, RealOpenMM** atomParameters,
-                               const std::vector<std::vector<float> >& values,
-                               float* forces, std::vector<std::vector<float> >& dEdV,
-                               bool isExcluded, const fvec4& boxSize, const fvec4& invBoxSize);
+    void calculateOnePairChainRule(int atom1, int atom2, ThreadData& data, float* posq, RealOpenMM** atomParameters,
+                               float* forces, bool isExcluded, const fvec4& boxSize, const fvec4& invBoxSize);
 
     /**
      * Compute the displacement and squared distance between two points, optionally using
@@ -211,7 +197,8 @@ public:
      * Construct a new CpuCustomGBForce.
      */
 
-     CpuCustomGBForce(int numAtoms, const std::vector<Lepton::CompiledExpression>& valueExpressions,
+     CpuCustomGBForce(int numAtoms, const std::vector<std::set<int> >& exclusions,
+                        const std::vector<Lepton::CompiledExpression>& valueExpressions,
                         const std::vector<std::vector<Lepton::CompiledExpression> >& valueDerivExpressions,
                         const std::vector<std::vector<Lepton::CompiledExpression> >& valueGradientExpressions,
                         const std::vector<std::string>& valueNames,
@@ -220,7 +207,7 @@ public:
                         const std::vector<std::vector<Lepton::CompiledExpression> >& energyDerivExpressions,
                         const std::vector<std::vector<Lepton::CompiledExpression> >& energyGradientExpressions,
                         const std::vector<CustomGBForce::ComputationType>& energyTypes,
-                          const std::vector<std::string>& parameterNames);
+                        const std::vector<std::string>& parameterNames, ThreadPool& threads);
 
      ~CpuCustomGBForce();
 
@@ -249,14 +236,42 @@ public:
      * @param numberOfAtoms    number of atoms
      * @param posq             atom coordinates
      * @param atomParameters   atomParameters[atomIndex][paramterIndex]
-     * @param exclusions       exclusions[i] is the set of excluded indices for atom i
      * @param globalParameters the values of global parameters
      * @param forces           force array (forces added)
      * @param totalEnergy      total energy
      */
 
-    void calculateIxn(int numberOfAtoms, float* posq, RealOpenMM** atomParameters, const std::vector<std::set<int> >& exclusions,
-                     std::map<std::string, double>& globalParameters, float* forces, double* totalEnergy);
+    void calculateIxn(int numberOfAtoms, float* posq, RealOpenMM** atomParameters,
+                     std::map<std::string, double>& globalParameters, std::vector<AlignedArray<float> >& threadForce, bool includeForce, bool includeEnergy, double& totalEnergy);
+};
+
+class CpuCustomGBForce::ThreadData {
+public:
+    ThreadData(int numAtoms, int numThreads, int threadIndex,
+               const std::vector<Lepton::CompiledExpression>& valueExpressions,
+               const std::vector<std::vector<Lepton::CompiledExpression> >& valueDerivExpressions,
+               const std::vector<std::vector<Lepton::CompiledExpression> >& valueGradientExpressions,
+               const std::vector<std::string>& valueNames,
+               const std::vector<Lepton::CompiledExpression>& energyExpressions,
+               const std::vector<std::vector<Lepton::CompiledExpression> >& energyDerivExpressions,
+               const std::vector<std::vector<Lepton::CompiledExpression> >& energyGradientExpressions,
+               const std::vector<std::string>& parameterNames);
+    CompiledExpressionSet expressionSet;
+    std::vector<Lepton::CompiledExpression> valueExpressions;
+    std::vector<std::vector<Lepton::CompiledExpression> > valueDerivExpressions;
+    std::vector<std::vector<Lepton::CompiledExpression> > valueGradientExpressions;
+    std::vector<int> valueIndex;
+    std::vector<Lepton::CompiledExpression> energyExpressions;
+    std::vector<std::vector<Lepton::CompiledExpression> > energyDerivExpressions;
+    std::vector<std::vector<Lepton::CompiledExpression> > energyGradientExpressions;
+    std::vector<int> paramIndex;
+    std::vector<int> particleParamIndex;
+    std::vector<int> particleValueIndex;
+    int xindex, yindex, zindex, rindex;
+    int firstAtom, lastAtom;
+    // Workspace vectors
+    std::vector<float> value0, dVdR1, dVdR2, dVdX, dVdY, dVdZ;
+    std::vector<std::vector<float> > dEdV;
 };
 
 } // namespace OpenMM

platforms/cpu/src/CpuKernels.cpp

@@ -957,8 +957,8 @@ void CpuCalcCustomGBForceKernel::initialize(const System& system, const CustomGB
 
     for (map<string, Lepton::CustomFunction*>::iterator iter = functions.begin(); iter != functions.end(); iter++)
         delete iter->second;
-    ixn = new CpuCustomGBForce(numParticles, valueExpressions, valueDerivExpressions, valueGradientExpressions, valueNames, valueTypes, energyExpressions,
-        energyDerivExpressions, energyGradientExpressions, energyTypes, particleParameterNames);
+    ixn = new CpuCustomGBForce(numParticles, exclusions, valueExpressions, valueDerivExpressions, valueGradientExpressions, valueNames, valueTypes, energyExpressions,
+        energyDerivExpressions, energyGradientExpressions, energyTypes, particleParameterNames, data.threads);
     data.isPeriodic = (force.getNonbondedMethod() == CustomGBForce::CutoffPeriodic);
 }
 
@@ -977,7 +977,7 @@ double CpuCalcCustomGBForceKernel::execute(ContextImpl& context, bool includeFor
     map<string, double> globalParameters;
     for (int i = 0; i < (int) globalParameterNames.size(); i++)
         globalParameters[globalParameterNames[i]] = context.getParameter(globalParameterNames[i]);
-    ixn->calculateIxn(numParticles, &data.posq[0], particleParamArray, exclusions, globalParameters, &data.threadForce[0][0], includeEnergy ? &energy : NULL);
+    ixn->calculateIxn(numParticles, &data.posq[0], particleParamArray, globalParameters, data.threadForce, includeForces, includeEnergy, energy);
     return energy;
 }