Finished parallelizing CPU implementation of CustomNonbondedForce

platforms/cpu/include/CpuCustomNonbondedForce.h

@@ -38,54 +38,6 @@
 namespace OpenMM {
     
 class CpuCustomNonbondedForce {
-
-   private:
-      class ComputeForceTask;
-      class ThreadData;
-
-      bool cutoff;
-      bool useSwitch;
-      bool periodic;
-      const CpuNeighborList* neighborList;
-      RealOpenMM periodicBoxSize[3];
-      RealOpenMM cutoffDistance, switchingDistance;
-      ThreadPool& threads;
-      std::vector<ThreadData*> threadData;
-      std::vector<std::string> paramNames;
-      std::vector<std::pair<std::set<int>, std::set<int> > > interactionGroups;
-      std::vector<double> threadEnergy;
-    // The following variables are used to make information accessible to the individual threads.
-    int numberOfAtoms;
-    float* posq;
-    RealVec const* atomCoordinates;
-    RealOpenMM** atomParameters;        
-    const std::map<std::string, double>* globalParameters;
-    std::set<int> const* exclusions;
-    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 custom pair ixn between two atoms
-
-         @param atom1            the index of the first atom
-         @param atom2            the index of the second atom
-         @param atomCoordinates  atom coordinates
-         @param atomParameters   atomParameters[atomIndex][parameterIndex]
-         @param forces           force array (forces added)
-         @param totalEnergy      total energy
-
-         --------------------------------------------------------------------------------------- */
-
-      void calculateOneIxn(int atom1, int atom2, ThreadData& data, float* forces, double& totalEnergy, const fvec4& boxSize, const fvec4& invBoxSize);
-
-
    public:
 
       /**---------------------------------------------------------------------------------------
@@ -95,7 +47,7 @@ class CpuCustomNonbondedForce {
          --------------------------------------------------------------------------------------- */
 
        CpuCustomNonbondedForce(const Lepton::CompiledExpression& energyExpression, const Lepton::CompiledExpression& forceExpression,
-                                   const std::vector<std::string>& parameterNames, ThreadPool& threads);
+                                   const std::vector<std::string>& parameterNames, const std::vector<std::set<int> >& exclusions, ThreadPool& threads);
 
       /**---------------------------------------------------------------------------------------
 
@@ -103,7 +55,7 @@ class CpuCustomNonbondedForce {
 
          --------------------------------------------------------------------------------------- */
 
-       ~CpuCustomNonbondedForce( );
+       ~CpuCustomNonbondedForce();
 
       /**---------------------------------------------------------------------------------------
 
@@ -114,7 +66,7 @@ class CpuCustomNonbondedForce {
 
          --------------------------------------------------------------------------------------- */
 
-      void setUseCutoff( RealOpenMM distance, const CpuNeighborList& neighbors );
+      void setUseCutoff(RealOpenMM distance, const CpuNeighborList& neighbors);
 
       /**---------------------------------------------------------------------------------------
 
@@ -135,7 +87,7 @@ class CpuCustomNonbondedForce {
       
          --------------------------------------------------------------------------------------- */
       
-      void setUseSwitchingFunction( RealOpenMM distance );
+      void setUseSwitchingFunction(RealOpenMM distance);
 
       /**---------------------------------------------------------------------------------------
 
@@ -147,7 +99,7 @@ class CpuCustomNonbondedForce {
 
          --------------------------------------------------------------------------------------- */
 
-      void setPeriodic( OpenMM::RealVec& boxSize );
+      void setPeriodic(OpenMM::RealVec& boxSize);
 
       /**---------------------------------------------------------------------------------------
 
@@ -157,8 +109,6 @@ class CpuCustomNonbondedForce {
          @param posq             atom coordinates in float format
          @param atomCoordinates  atom coordinates
          @param atomParameters   atom parameters (charges, c6, c12, ...)     atomParameters[atomIndex][paramterIndex]
-         @param exclusions       atom exclusion indices
-                                 exclusions[atomIndex] contains the list of exclusions for that atom
          @param fixedParameters  non atom parameters (not currently used)
          @param globalParameters the values of global parameters
          @param forces           force array (forces added)
@@ -167,10 +117,52 @@ class CpuCustomNonbondedForce {
 
          --------------------------------------------------------------------------------------- */
 
-      void calculatePairIxn(int numberOfAtoms, float* posq, std::vector<OpenMM::RealVec>& atomCoordinates,
-                            RealOpenMM** atomParameters, std::vector<std::set<int> >& exclusions,
+    void calculatePairIxn(int numberOfAtoms, float* posq, std::vector<OpenMM::RealVec>& atomCoordinates, RealOpenMM** atomParameters,
                           RealOpenMM* fixedParameters, const std::map<std::string, double>& globalParameters,
                           std::vector<AlignedArray<float> >& threadForce, bool includeForce, bool includeEnergy, double& totalEnergy);
+private:
+    class ComputeForceTask;
+    class ThreadData;
+
+    bool cutoff;
+    bool useSwitch;
+    bool periodic;
+    const CpuNeighborList* neighborList;
+    RealOpenMM periodicBoxSize[3];
+    RealOpenMM cutoffDistance, switchingDistance;
+    ThreadPool& threads;
+    const std::vector<std::set<int> > exclusions;
+    std::vector<ThreadData*> threadData;
+    std::vector<std::string> paramNames;
+    std::vector<std::pair<int, int> > groupInteractions;
+    std::vector<double> threadEnergy;
+    // The following variables are used to make information accessible to the individual threads.
+    int numberOfAtoms;
+    float* posq;
+    RealVec const* atomCoordinates;
+    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 the interaction between two atoms.
+     * 
+     * @param atom1            the index of the first atom
+     * @param atom2            the index of the second atom
+     * @param data             workspace for the current thread
+     * @param forces           force array (forces added)
+     * @param totalEnergy      total energy
+     * @param boxSize          the size of the periodic box
+     * @param boxSize          the inverse size of the periodic box
+     */
+    void calculateOneIxn(int atom1, int atom2, ThreadData& data, float* forces, double& totalEnergy, const fvec4& boxSize, const fvec4& invBoxSize);
 
     /**
      * Compute the displacement and squared distance between two points, optionally using

platforms/cpu/src/CpuCustomNonbondedForce.cpp

@@ -60,8 +60,8 @@ CpuCustomNonbondedForce::ThreadData::ThreadData(const Lepton::CompiledExpression
 }
 
 CpuCustomNonbondedForce::CpuCustomNonbondedForce(const Lepton::CompiledExpression& energyExpression,
-            const Lepton::CompiledExpression& forceExpression, const vector<string>& parameterNames, ThreadPool& threads) :
-            cutoff(false), useSwitch(false), periodic(false), paramNames(parameterNames), threads(threads) {
+            const Lepton::CompiledExpression& forceExpression, const vector<string>& parameterNames, const vector<set<int> >& exclusions,ThreadPool& threads) :
+            cutoff(false), useSwitch(false), periodic(false), paramNames(parameterNames), exclusions(exclusions), threads(threads) {
     for (int i = 0; i < threads.getNumThreads(); i++)
         threadData.push_back(new ThreadData(energyExpression, forceExpression, parameterNames));
 }
@@ -78,7 +78,19 @@ void CpuCustomNonbondedForce::setUseCutoff(RealOpenMM distance, const CpuNeighbo
   }
 
 void CpuCustomNonbondedForce::setInteractionGroups(const vector<pair<set<int>, set<int> > >& groups) {
-    interactionGroups = groups;
+    for (int group = 0; group < (int) groups.size(); group++) {
+        const set<int>& set1 = groups[group].first;
+        const set<int>& set2 = groups[group].second;
+        for (set<int>::const_iterator atom1 = set1.begin(); atom1 != set1.end(); ++atom1) {
+            for (set<int>::const_iterator atom2 = set2.begin(); atom2 != set2.end(); ++atom2) {
+                if (*atom1 == *atom2 || exclusions[*atom1].find(*atom2) != exclusions[*atom1].end())
+                    continue; // This is an excluded interaction.
+                if (*atom1 > *atom2 && set1.find(*atom2) != set1.end() && set2.find(*atom1) != set2.end())
+                    continue; // Both atoms are in both sets, so skip duplicate interactions.
+                groupInteractions.push_back(make_pair(*atom1, *atom2));
+            }
+        }
+    }
 }
 
 void CpuCustomNonbondedForce::setUseSwitchingFunction(RealOpenMM distance) {
@@ -98,8 +110,7 @@ void CpuCustomNonbondedForce::setPeriodic(RealVec& boxSize) {
 }
 
 
-void CpuCustomNonbondedForce::calculatePairIxn(int numberOfAtoms, float* posq, vector<RealVec>& atomCoordinates,
-                                             RealOpenMM** atomParameters, vector<set<int> >& exclusions,
+void CpuCustomNonbondedForce::calculatePairIxn(int numberOfAtoms, float* posq, vector<RealVec>& atomCoordinates, RealOpenMM** atomParameters,
                                              RealOpenMM* fixedParameters, const map<string, double>& globalParameters,
                                              vector<AlignedArray<float> >& threadForce, bool includeForce, bool includeEnergy, double& totalEnergy) {
     // Record the parameters for the threads.
@@ -109,7 +120,6 @@ void CpuCustomNonbondedForce::calculatePairIxn(int numberOfAtoms, float* posq, v
     this->atomCoordinates = &atomCoordinates[0];
     this->atomParameters = atomParameters;
     this->globalParameters = &globalParameters;
-    this->exclusions = &exclusions[0];
     this->threadForce = &threadForce;
     this->includeForce = includeForce;
     this->includeEnergy = includeEnergy;
@@ -147,29 +157,22 @@ void CpuCustomNonbondedForce::threadComputeForce(ThreadPool& threads, int thread
     }
     fvec4 boxSize(periodicBoxSize[0], periodicBoxSize[1], periodicBoxSize[2], 0);
     fvec4 invBoxSize((1/periodicBoxSize[0]), (1/periodicBoxSize[1]), (1/periodicBoxSize[2]), 0);
-    if (interactionGroups.size() > 0) {
-        if (threadIndex > 0)
-            return;
+    if (groupInteractions.size() > 0) {
         // The user has specified interaction groups, so compute only the requested interactions.
         
-        for (int group = 0; group < (int) interactionGroups.size(); group++) {
-            const set<int>& set1 = interactionGroups[group].first;
-            const set<int>& set2 = interactionGroups[group].second;
-            for (set<int>::const_iterator atom1 = set1.begin(); atom1 != set1.end(); ++atom1) {
-                for (set<int>::const_iterator atom2 = set2.begin(); atom2 != set2.end(); ++atom2) {
-                    if (*atom1 == *atom2 || exclusions[*atom1].find(*atom2) != exclusions[*atom1].end())
-                        continue; // This is an excluded interaction.
-                    if (*atom1 > *atom2 && set1.find(*atom2) != set1.end() && set2.find(*atom1) != set2.end())
-                        continue; // Both atoms are in both sets, so skip duplicate interactions.
+        while (true) {
+            int i = gmx_atomic_fetch_add(reinterpret_cast<gmx_atomic_t*>(atomicCounter), 1);
+            if (i >= groupInteractions.size())
+                break;
+            int atom1 = groupInteractions[i].first;
+            int atom2 = groupInteractions[i].second;
             for (int j = 0; j < (int) paramNames.size(); j++) {
-                        ReferenceForce::setVariable(data.energyParticleParams[j*2], atomParameters[*atom1][j]);
-                        ReferenceForce::setVariable(data.energyParticleParams[j*2+1], atomParameters[*atom2][j]);
-                        ReferenceForce::setVariable(data.forceParticleParams[j*2], atomParameters[*atom1][j]);
-                        ReferenceForce::setVariable(data.forceParticleParams[j*2+1], atomParameters[*atom2][j]);
-                    }
-                    calculateOneIxn(*atom1, *atom2, data, forces, energy, boxSize, invBoxSize);
-                }
+                ReferenceForce::setVariable(data.energyParticleParams[j*2], atomParameters[atom1][j]);
+                ReferenceForce::setVariable(data.energyParticleParams[j*2+1], atomParameters[atom2][j]);
+                ReferenceForce::setVariable(data.forceParticleParams[j*2], atomParameters[atom1][j]);
+                ReferenceForce::setVariable(data.forceParticleParams[j*2+1], atomParameters[atom2][j]);
             }
+            calculateOneIxn(atom1, atom2, data, forces, energy, boxSize, invBoxSize);
         }
     }
     else if (cutoff) {
@@ -206,9 +209,10 @@ void CpuCustomNonbondedForce::threadComputeForce(ThreadPool& threads, int thread
     else {
         // Every particle interacts with every other one.
         
-        if (threadIndex > 0)
-            return;
-        for (int ii = 0; ii < numberOfAtoms; ii++) {
+        while (true) {
+            int ii = gmx_atomic_fetch_add(reinterpret_cast<gmx_atomic_t*>(atomicCounter), 1);
+            if (ii >= numberOfAtoms)
+                break;
             for (int jj = ii+1; jj < numberOfAtoms; jj++) {
                 if (exclusions[jj].find(ii) == exclusions[jj].end()) {
                     for (int j = 0; j < (int) paramNames.size(); j++) {

platforms/cpu/src/CpuKernels.cpp

@@ -719,7 +719,7 @@ void CpuCalcCustomNonbondedForceKernel::initialize(const System& system, const C
         interactionGroups.push_back(make_pair(set1, set2));
     }
     data.isPeriodic = (nonbondedMethod == CutoffPeriodic);
-    nonbonded = new CpuCustomNonbondedForce(energyExpression, forceExpression, parameterNames, data.threads);
+    nonbonded = new CpuCustomNonbondedForce(energyExpression, forceExpression, parameterNames, exclusions, data.threads);
 }
 
 double CpuCalcCustomNonbondedForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) {
@@ -750,7 +750,7 @@ double CpuCalcCustomNonbondedForceKernel::execute(ContextImpl& context, bool inc
     }
     if (useSwitchingFunction)
         nonbonded->setUseSwitchingFunction(switchingDistance);
-    nonbonded->calculatePairIxn(numParticles, &data.posq[0], posData, particleParamArray, exclusions, 0, globalParamValues, data.threadForce, includeForces, includeEnergy, energy);
+    nonbonded->calculatePairIxn(numParticles, &data.posq[0], posData, particleParamArray, 0, globalParamValues, data.threadForce, includeForces, includeEnergy, energy);
     
     // Add in the long range correction.