Skip to content

GitLab

Commit 6a4ac837 authored by peastman's avatar peastman
Browse files

Parallelized CpuCustomGBForce

parent bb70341a
Hide whitespace changes
Inline Side-by-side
Showing with 380 additions and 231 deletions
platforms/cpu/include/CpuCustomGBForce.h
View file @ 6a4ac837
...@@ -29,6 +29,7 @@ ...@@ -29,6 +29,7 @@
#include "CpuNeighborList.h" #include "CpuNeighborList.h"
#include "lepton/CompiledExpression.h" #include "lepton/CompiledExpression.h"
#include "openmm/CustomGBForce.h" #include "openmm/CustomGBForce.h"
#include "openmm/internal/ThreadPool.h"
#include "openmm/internal/vectorize.h" #include "openmm/internal/vectorize.h"
#include <map> #include <map>
#include <set> #include <set>
...@@ -38,59 +39,63 @@ namespace OpenMM { ...@@ -38,59 +39,63 @@ namespace OpenMM {
class CpuCustomGBForce { class CpuCustomGBForce {
private: private:
class ComputeForceTask;
class ThreadData;
bool cutoff; bool cutoff;
bool periodic; bool periodic;
const CpuNeighborList* neighborList; const CpuNeighborList* neighborList;
float periodicBoxSize[3]; float periodicBoxSize[3];
float cutoffDistance; float cutoffDistance;
CompiledExpressionSet expressionSet; const std::vector<std::set<int> > exclusions;
std::vector<Lepton::CompiledExpression> valueExpressions;
std::vector<std::vector<Lepton::CompiledExpression> > valueDerivExpressions;
std::vector<std::vector<Lepton::CompiledExpression> > valueGradientExpressions;
std::vector<std::string> valueNames; std::vector<std::string> valueNames;
std::vector<int> valueIndex;
std::vector<CustomGBForce::ComputationType> valueTypes; 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<std::string> paramNames;
std::vector<int> paramIndex;
std::vector<CustomGBForce::ComputationType> energyTypes; std::vector<CustomGBForce::ComputationType> energyTypes;
std::vector<int> particleParamIndex; ThreadPool& threads;
std::vector<int> particleValueIndex; std::vector<ThreadData*> threadData;
int xindex, yindex, zindex, rindex; std::vector<double> threadEnergy;
// Workspace vectors // Workspace vectors
std::vector<std::vector<float> > values, dEdV; 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 * Calculate a computed value of type SingleParticle
* *
* @param index the index of the value to compute * @param index the index of the value to compute
* @param data workspace for the current thread
* @param numAtoms number of atoms * @param numAtoms number of atoms
* @param posq atom coordinates * @param posq atom coordinates
* @param values the vector to store computed values into
* @param atomParameters atomParameters[atomIndex][paramterIndex] * @param atomParameters atomParameters[atomIndex][paramterIndex]
*/ */
void calculateSingleParticleValue(int index, int numAtoms, float* posq, std::vector<std::vector<float> >& values, void calculateSingleParticleValue(int index, ThreadData& data, int numAtoms, float* posq, RealOpenMM** atomParameters);
RealOpenMM** atomParameters);
/** /**
* Calculate a computed value that is based on particle pairs * Calculate a computed value that is based on particle pairs
* *
* @param index the index of the value to compute * @param index the index of the value to compute
* @param data workspace for the current thread
* @param numAtoms number of atoms * @param numAtoms number of atoms
* @param posq atom coordinates * @param posq atom coordinates
* @param atomParameters atomParameters[atomIndex][paramterIndex] * @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 * @param useExclusions specifies whether to use exclusions
*/ */
void calculateParticlePairValue(int index, int numAtoms, float* posq, RealOpenMM** atomParameters, void calculateParticlePairValue(int index, ThreadData& data, int numAtoms, float* posq, RealOpenMM** atomParameters,
std::vector<std::vector<float> >& values, bool useExclusions, const fvec4& boxSize, const fvec4& invBoxSize);
const std::vector<std::set<int> >& exclusions, bool useExclusions, const fvec4& boxSize, const fvec4& invBoxSize);
/** /**
* Evaluate a single atom pair as part of calculating a computed value * Evaluate a single atom pair as part of calculating a computed value
...@@ -98,51 +103,43 @@ private: ...@@ -98,51 +103,43 @@ private:
* @param index the index of the value to compute * @param index the index of the value to compute
* @param atom1 the index of the first atom in the pair * @param atom1 the index of the first atom in the pair
* @param atom2 the index of the second 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 posq atom coordinates
* @param atomParameters atomParameters[atomIndex][paramterIndex] * @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, void calculateOnePairValue(int index, int atom1, int atom2, ThreadData& data, float* posq, RealOpenMM** atomParameters,
std::vector<std::vector<float> >& values, const fvec4& boxSize, const fvec4& invBoxSize); std::vector<float>& valueArray, const fvec4& boxSize, const fvec4& invBoxSize);
/** /**
* Calculate an energy term of type SingleParticle * Calculate an energy term of type SingleParticle
* *
* @param index the index of the value to compute * @param index the index of the value to compute
* @param data workspace for the current thread
* @param numAtoms number of atoms * @param numAtoms number of atoms
* @param posq atom coordinates * @param posq atom coordinates
* @param values the vector containing computed values
* @param atomParameters atomParameters[atomIndex][paramterIndex] * @param atomParameters atomParameters[atomIndex][paramterIndex]
* @param forces forces on atoms are added to this * @param forces forces on atoms are added to this
* @param totalEnergy the energy contribution is 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, void calculateSingleParticleEnergyTerm(int index, ThreadData& data, int numAtoms, float* posq, RealOpenMM** atomParameters, float* forces, double& totalEnergy);
RealOpenMM** atomParameters, float* forces,
double* totalEnergy, std::vector<std::vector<float> >& dEdV);
/** /**
* Calculate an energy term that is based on particle pairs * Calculate an energy term that is based on particle pairs
* *
* @param index the index of the term to compute * @param index the index of the term to compute
* @param data workspace for the current thread
* @param numAtoms number of atoms * @param numAtoms number of atoms
* @param posq atom coordinates * @param posq atom coordinates
* @param atomParameters atomParameters[atomIndex][paramterIndex] * @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 useExclusions specifies whether to use exclusions
* @param forces forces on atoms are added to this * @param forces forces on atoms are added to this
* @param totalEnergy the energy contribution is 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, void calculateParticlePairEnergyTerm(int index, ThreadData& data, int numAtoms, float* posq, RealOpenMM** atomParameters,
const std::vector<std::vector<float> >& values, bool useExclusions, float* forces, double& totalEnergy, const fvec4& boxSize, const fvec4& invBoxSize);
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);
/** /**
* Evaluate a single atom pair as part of calculating an energy term * Evaluate a single atom pair as part of calculating an energy term
...@@ -150,54 +147,43 @@ private: ...@@ -150,54 +147,43 @@ private:
* @param index the index of the term to compute * @param index the index of the term to compute
* @param atom1 the index of the first atom in the pair * @param atom1 the index of the first atom in the pair
* @param atom2 the index of the second 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 posq atom coordinates
* @param atomParameters atomParameters[atomIndex][paramterIndex] * @param atomParameters atomParameters[atomIndex][paramterIndex]
* @param values the vector containing computed values
* @param forces forces on atoms are added to this * @param forces forces on atoms are added to this
* @param totalEnergy the energy contribution is 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, void calculateOnePairEnergyTerm(int index, int atom1, int atom2, ThreadData& data, float* posq, RealOpenMM** atomParameters,
const std::vector<std::vector<float> >& values, float* forces, double& totalEnergy, const fvec4& boxSize, const fvec4& invBoxSize);
float* forces, double* totalEnergy, std::vector<std::vector<float> >& dEdV,
const fvec4& boxSize, const fvec4& invBoxSize);
/** /**
* Apply the chain rule to compute forces on atoms * Apply the chain rule to compute forces on atoms
* *
* @param data workspace for the current thread
* @param numAtoms number of atoms * @param numAtoms number of atoms
* @param posq atom coordinates * @param posq atom coordinates
* @param atomParameters atomParameters[atomIndex][paramterIndex] * @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 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, void calculateChainRuleForces(ThreadData& data, int numAtoms, float* posq, RealOpenMM** atomParameters,
const std::vector<std::vector<float> >& values, float* forces, const fvec4& boxSize, const fvec4& invBoxSize);
const std::vector<std::set<int> >& exclusions,
float* forces, std::vector<std::vector<float> >& dEdV,
const fvec4& boxSize, const fvec4& invBoxSize);
/** /**
* Evaluate a single atom pair as part of applying the chain rule * Evaluate a single atom pair as part of applying the chain rule
* *
* @param atom1 the index of the first atom in the pair * @param atom1 the index of the first atom in the pair
* @param atom2 the index of the second 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 posq atom coordinates
* @param atomParameters atomParameters[atomIndex][paramterIndex] * @param atomParameters atomParameters[atomIndex][paramterIndex]
* @param values the vector containing computed values
* @param forces forces on atoms are added to this * @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 * @param isExcluded specifies whether this is an excluded pair
*/ */
void calculateOnePairChainRule(int atom1, int atom2, float* posq, RealOpenMM** atomParameters, void calculateOnePairChainRule(int atom1, int atom2, ThreadData& data, float* posq, RealOpenMM** atomParameters,
const std::vector<std::vector<float> >& values, float* forces, bool isExcluded, const fvec4& boxSize, const fvec4& invBoxSize);
float* forces, std::vector<std::vector<float> >& dEdV,
bool isExcluded, const fvec4& boxSize, const fvec4& invBoxSize);
/** /**
* Compute the displacement and squared distance between two points, optionally using * Compute the displacement and squared distance between two points, optionally using
...@@ -211,16 +197,17 @@ public: ...@@ -211,16 +197,17 @@ public:
* Construct a new CpuCustomGBForce. * 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<std::vector<Lepton::CompiledExpression> >& valueDerivExpressions, const std::vector<Lepton::CompiledExpression>& valueExpressions,
const std::vector<std::vector<Lepton::CompiledExpression> >& valueGradientExpressions, const std::vector<std::vector<Lepton::CompiledExpression> >& valueDerivExpressions,
const std::vector<std::string>& valueNames, const std::vector<std::vector<Lepton::CompiledExpression> >& valueGradientExpressions,
const std::vector<CustomGBForce::ComputationType>& valueTypes, const std::vector<std::string>& valueNames,
const std::vector<Lepton::CompiledExpression>& energyExpressions, const std::vector<CustomGBForce::ComputationType>& valueTypes,
const std::vector<std::vector<Lepton::CompiledExpression> >& energyDerivExpressions, const std::vector<Lepton::CompiledExpression>& energyExpressions,
const std::vector<std::vector<Lepton::CompiledExpression> >& energyGradientExpressions, const std::vector<std::vector<Lepton::CompiledExpression> >& energyDerivExpressions,
const std::vector<CustomGBForce::ComputationType>& energyTypes, const std::vector<std::vector<Lepton::CompiledExpression> >& energyGradientExpressions,
const std::vector<std::string>& parameterNames); const std::vector<CustomGBForce::ComputationType>& energyTypes,
const std::vector<std::string>& parameterNames, ThreadPool& threads);
~CpuCustomGBForce(); ~CpuCustomGBForce();
...@@ -249,14 +236,42 @@ public: ...@@ -249,14 +236,42 @@ public:
* @param numberOfAtoms number of atoms * @param numberOfAtoms number of atoms
* @param posq atom coordinates * @param posq atom coordinates
* @param atomParameters atomParameters[atomIndex][paramterIndex] * @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 globalParameters the values of global parameters
* @param forces force array (forces added) * @param forces force array (forces added)
* @param totalEnergy total energy * @param totalEnergy total energy
*/ */
void calculateIxn(int numberOfAtoms, float* posq, RealOpenMM** atomParameters, const std::vector<std::set<int> >& exclusions, void calculateIxn(int numberOfAtoms, float* posq, RealOpenMM** atomParameters,
std::map<std::string, double>& globalParameters, float* forces, double* totalEnergy); 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 } // namespace OpenMM
......
platforms/cpu/src/CpuCustomGBForce.cpp
View file @ 6a4ac837
...@@ -30,23 +30,34 @@ ...@@ -30,23 +30,34 @@
#include "SimTKOpenMMUtilities.h" #include "SimTKOpenMMUtilities.h"
#include "ReferenceForce.h" #include "ReferenceForce.h"
#include "CpuCustomGBForce.h" #include "CpuCustomGBForce.h"
#include "gmx_atomic.h"
using namespace OpenMM; using namespace OpenMM;
using namespace std; using namespace std;
CpuCustomGBForce::CpuCustomGBForce(int numAtoms, const vector<Lepton::CompiledExpression>& valueExpressions, class CpuCustomGBForce::ComputeForceTask : public ThreadPool::Task {
const vector<vector<Lepton::CompiledExpression> >& valueDerivExpressions, public:
const vector<vector<Lepton::CompiledExpression> >& valueGradientExpressions, ComputeForceTask(CpuCustomGBForce& owner) : owner(owner) {
const vector<string>& valueNames, }
const vector<CustomGBForce::ComputationType>& valueTypes, void execute(ThreadPool& threads, int threadIndex) {
const vector<Lepton::CompiledExpression>& energyExpressions, owner.threadComputeForce(threads, threadIndex);
const vector<vector<Lepton::CompiledExpression> >& energyDerivExpressions, }
const vector<vector<Lepton::CompiledExpression> >& energyGradientExpressions, CpuCustomGBForce& owner;
const vector<CustomGBForce::ComputationType>& energyTypes, };
const vector<string>& parameterNames) :
cutoff(false), periodic(false), valueExpressions(valueExpressions), valueDerivExpressions(valueDerivExpressions), valueGradientExpressions(valueGradientExpressions), CpuCustomGBForce::ThreadData::ThreadData(int numAtoms, int numThreads, int threadIndex,
valueNames(valueNames), valueTypes(valueTypes), energyExpressions(energyExpressions), energyDerivExpressions(energyDerivExpressions), energyGradientExpressions(energyGradientExpressions), const vector<Lepton::CompiledExpression>& valueExpressions,
energyTypes(energyTypes), paramNames(parameterNames) { const vector<vector<Lepton::CompiledExpression> >& valueDerivExpressions,
const vector<vector<Lepton::CompiledExpression> >& valueGradientExpressions,
const vector<string>& valueNames,
const vector<Lepton::CompiledExpression>& energyExpressions,
const vector<vector<Lepton::CompiledExpression> >& energyDerivExpressions,
const vector<vector<Lepton::CompiledExpression> >& energyGradientExpressions,
const vector<string>& parameterNames) :
valueExpressions(valueExpressions), valueDerivExpressions(valueDerivExpressions), valueGradientExpressions(valueGradientExpressions),
energyExpressions(energyExpressions), energyDerivExpressions(energyDerivExpressions), energyGradientExpressions(energyGradientExpressions) {
firstAtom = (threadIndex*(long long) numAtoms)/numThreads;
lastAtom = ((threadIndex+1)*(long long) numAtoms)/numThreads;
for (int i = 0; i < (int) valueExpressions.size(); i++) for (int i = 0; i < (int) valueExpressions.size(); i++)
expressionSet.registerExpression(this->valueExpressions[i]); expressionSet.registerExpression(this->valueExpressions[i]);
for (int i = 0; i < (int) valueDerivExpressions.size(); i++) for (int i = 0; i < (int) valueDerivExpressions.size(); i++)
...@@ -67,11 +78,11 @@ CpuCustomGBForce::CpuCustomGBForce(int numAtoms, const vector<Lepton::CompiledEx ...@@ -67,11 +78,11 @@ CpuCustomGBForce::CpuCustomGBForce(int numAtoms, const vector<Lepton::CompiledEx
yindex = expressionSet.getVariableIndex("y"); yindex = expressionSet.getVariableIndex("y");
zindex = expressionSet.getVariableIndex("z"); zindex = expressionSet.getVariableIndex("z");
rindex = expressionSet.getVariableIndex("r"); rindex = expressionSet.getVariableIndex("r");
for (int i = 0; i < (int) paramNames.size(); i++) { for (int i = 0; i < (int) parameterNames.size(); i++) {
paramIndex.push_back(expressionSet.getVariableIndex(paramNames[i])); paramIndex.push_back(expressionSet.getVariableIndex(parameterNames[i]));
for (int j = 1; j < 3; j++) { for (int j = 1; j < 3; j++) {
stringstream name; stringstream name;
name << paramNames[i] << j; name << parameterNames[i] << j;
particleParamIndex.push_back(expressionSet.getVariableIndex(name.str())); particleParamIndex.push_back(expressionSet.getVariableIndex(name.str()));
} }
} }
...@@ -83,12 +94,10 @@ CpuCustomGBForce::CpuCustomGBForce(int numAtoms, const vector<Lepton::CompiledEx ...@@ -83,12 +94,10 @@ CpuCustomGBForce::CpuCustomGBForce(int numAtoms, const vector<Lepton::CompiledEx
particleValueIndex.push_back(expressionSet.getVariableIndex(name.str())); particleValueIndex.push_back(expressionSet.getVariableIndex(name.str()));
} }
} }
values.resize(valueTypes.size()); value0.resize(numAtoms);
dEdV.resize(valueTypes.size()); dEdV.resize(valueNames.size());
for (int i = 0; i < (int) values.size(); i++) { for (int i = 0; i < (int) dEdV.size(); i++)
values[i].resize(numAtoms);
dEdV[i].resize(numAtoms); dEdV[i].resize(numAtoms);
}
dVdX.resize(valueDerivExpressions.size()); dVdX.resize(valueDerivExpressions.size());
dVdY.resize(valueDerivExpressions.size()); dVdY.resize(valueDerivExpressions.size());
dVdZ.resize(valueDerivExpressions.size()); dVdZ.resize(valueDerivExpressions.size());
...@@ -96,7 +105,33 @@ CpuCustomGBForce::CpuCustomGBForce(int numAtoms, const vector<Lepton::CompiledEx ...@@ -96,7 +105,33 @@ CpuCustomGBForce::CpuCustomGBForce(int numAtoms, const vector<Lepton::CompiledEx
dVdR2.resize(valueDerivExpressions.size()); dVdR2.resize(valueDerivExpressions.size());
} }
CpuCustomGBForce::CpuCustomGBForce(int numAtoms, const std::vector<std::set<int> >& exclusions,
const vector<Lepton::CompiledExpression>& valueExpressions,
const vector<vector<Lepton::CompiledExpression> >& valueDerivExpressions,
const vector<vector<Lepton::CompiledExpression> >& valueGradientExpressions,
const vector<string>& valueNames,
const vector<CustomGBForce::ComputationType>& valueTypes,
const vector<Lepton::CompiledExpression>& energyExpressions,
const vector<vector<Lepton::CompiledExpression> >& energyDerivExpressions,
const vector<vector<Lepton::CompiledExpression> >& energyGradientExpressions,
const vector<CustomGBForce::ComputationType>& energyTypes,
const vector<string>& parameterNames, ThreadPool& threads) :
exclusions(exclusions), cutoff(false), periodic(false), valueNames(valueNames), valueTypes(valueTypes),
energyTypes(energyTypes), paramNames(parameterNames), threads(threads) {
for (int i = 0; i < threads.getNumThreads(); i++)
threadData.push_back(new ThreadData(numAtoms, threads.getNumThreads(), i, valueExpressions, valueDerivExpressions, valueGradientExpressions, valueNames,
energyExpressions, energyDerivExpressions, energyGradientExpressions, parameterNames));
values.resize(valueNames.size());
dEdV.resize(valueNames.size());
for (int i = 0; i < (int) values.size(); i++) {
values[i].resize(numAtoms);
dEdV[i].resize(numAtoms);
}
}
CpuCustomGBForce::~CpuCustomGBForce() { CpuCustomGBForce::~CpuCustomGBForce() {
for (int i = 0; i < (int) threadData.size(); i++)
delete threadData[i];
} }
void CpuCustomGBForce::setUseCutoff(float distance, const CpuNeighborList& neighbors) { void CpuCustomGBForce::setUseCutoff(float distance, const CpuNeighborList& neighbors) {
...@@ -106,7 +141,6 @@ void CpuCustomGBForce::setUseCutoff(float distance, const CpuNeighborList& neigh ...@@ -106,7 +141,6 @@ void CpuCustomGBForce::setUseCutoff(float distance, const CpuNeighborList& neigh
} }
void CpuCustomGBForce::setPeriodic(RealVec& boxSize) { void CpuCustomGBForce::setPeriodic(RealVec& boxSize) {
if (cutoff) { if (cutoff) {
assert(boxSize[0] >= 2.0*cutoffDistance); assert(boxSize[0] >= 2.0*cutoffDistance);
assert(boxSize[1] >= 2.0*cutoffDistance); assert(boxSize[1] >= 2.0*cutoffDistance);
...@@ -119,68 +153,157 @@ void CpuCustomGBForce::setPeriodic(RealVec& boxSize) { ...@@ -119,68 +153,157 @@ void CpuCustomGBForce::setPeriodic(RealVec& boxSize) {
} }
void CpuCustomGBForce::calculateIxn(int numberOfAtoms, float* posq, RealOpenMM** atomParameters, void CpuCustomGBForce::calculateIxn(int numberOfAtoms, float* posq, RealOpenMM** atomParameters,
const vector<set<int> >& exclusions, map<string, double>& globalParameters, float* forces, map<string, double>& globalParameters, vector<AlignedArray<float> >& threadForce,
double* totalEnergy) { bool includeForce, bool includeEnergy, double& totalEnergy) {
// Record the parameters for the threads.
this->numberOfAtoms = numberOfAtoms;
this->posq = posq;
this->atomParameters = atomParameters;
this->globalParameters = &globalParameters;
this->threadForce = &threadForce;
this->includeForce = includeForce;
this->includeEnergy = includeEnergy;
threadEnergy.resize(threads.getNumThreads());
gmx_atomic_t counter;
this->atomicCounter = &counter;
// Calculate the first computed value.
ComputeForceTask task(*this);
gmx_atomic_set(&counter, 0);
threads.execute(task);
threads.waitForThreads();
// Calculate the remaining computed values.
threads.resumeThreads();
threads.waitForThreads();
// Calculate the energy terms.
for (int i = 0; i < (int) threadData[0]->energyExpressions.size(); i++) {
gmx_atomic_set(&counter, 0);
threads.execute(task);
threads.waitForThreads();
}
// Sum the energy derivatives.
threads.resumeThreads();
threads.waitForThreads();
// Apply the chain rule to evaluate forces.
gmx_atomic_set(&counter, 0);
threads.resumeThreads();
threads.waitForThreads();
// Combine the energies from all the threads.
if (includeEnergy) {
int numThreads = threads.getNumThreads();
for (int i = 0; i < numThreads; i++)
totalEnergy += threadEnergy[i];
}
}
void CpuCustomGBForce::threadComputeForce(ThreadPool& threads, int threadIndex) {
// Compute this thread's subset of interactions.
int numThreads = threads.getNumThreads();
threadEnergy[threadIndex] = 0;
double& energy = threadEnergy[threadIndex];
float* forces = &(*threadForce)[threadIndex][0];
ThreadData& data = *threadData[threadIndex];
fvec4 boxSize(periodicBoxSize[0], periodicBoxSize[1], periodicBoxSize[2], 0); fvec4 boxSize(periodicBoxSize[0], periodicBoxSize[1], periodicBoxSize[2], 0);
fvec4 invBoxSize((1/periodicBoxSize[0]), (1/periodicBoxSize[1]), (1/periodicBoxSize[2]), 0); fvec4 invBoxSize((1/periodicBoxSize[0]), (1/periodicBoxSize[1]), (1/periodicBoxSize[2]), 0);
for (map<string, double>::const_iterator iter = globalParameters.begin(); iter != globalParameters.end(); ++iter) for (map<string, double>::const_iterator iter = globalParameters->begin(); iter != globalParameters->end(); ++iter)
expressionSet.setVariable(expressionSet.getVariableIndex(iter->first), iter->second); data.expressionSet.setVariable(data.expressionSet.getVariableIndex(iter->first), iter->second);
// Calculate the first computed value.
for (int i = 0; i < (int) data.value0.size(); i++)
data.value0[i] = 0.0f;
if (valueTypes[0] == CustomGBForce::ParticlePair)
calculateParticlePairValue(0, data, numberOfAtoms, posq, atomParameters, true, boxSize, invBoxSize);
else
calculateParticlePairValue(0, data, numberOfAtoms, posq, atomParameters, false, boxSize, invBoxSize);
threads.syncThreads();
// Sum the first computed value.
for (int atom = data.firstAtom; atom < data.lastAtom; atom++) {
float sum = 0.0f;
for (int j = 0; j < (int) threadData.size(); j++)
sum += threadData[j]->value0[atom];
values[0][atom] = sum;
}
// First calculate the computed values. // Calculate the remaining computed values.
int numValues = valueTypes.size(); int numValues = valueTypes.size();
for (int valueIndex = 0; valueIndex < numValues; valueIndex++) { for (int i = 1; i < numValues; i++)
if (valueTypes[valueIndex] == CustomGBForce::SingleParticle) calculateSingleParticleValue(i, data, numberOfAtoms, posq, atomParameters);
calculateSingleParticleValue(valueIndex, numberOfAtoms, posq, values, atomParameters); threads.syncThreads();
else if (valueTypes[valueIndex] == CustomGBForce::ParticlePair)
calculateParticlePairValue(valueIndex, numberOfAtoms, posq, atomParameters, values, exclusions, true, boxSize, invBoxSize);
else
calculateParticlePairValue(valueIndex, numberOfAtoms, posq, atomParameters, values, exclusions, false, boxSize, invBoxSize);
}
// Now calculate the energy and its derivatives. // Now calculate the energy and its derivatives.
for (int i = 0; i < (int) dEdV.size(); i++) for (int i = 0; i < (int) data.dEdV.size(); i++)
for (int j = 0; j < (int) dEdV[i].size(); j++) for (int j = 0; j < (int) data.dEdV[i].size(); j++)
dEdV[i][j] = 0.0; data.dEdV[i][j] = 0.0;
for (int termIndex = 0; termIndex < (int) energyExpressions.size(); termIndex++) { for (int termIndex = 0; termIndex < (int) data.energyExpressions.size(); termIndex++) {
if (energyTypes[termIndex] == CustomGBForce::SingleParticle) if (energyTypes[termIndex] == CustomGBForce::SingleParticle)
calculateSingleParticleEnergyTerm(termIndex, numberOfAtoms, posq, values, atomParameters, forces, totalEnergy, dEdV); calculateSingleParticleEnergyTerm(termIndex, data, numberOfAtoms, posq, atomParameters, forces, energy);
else if (energyTypes[termIndex] == CustomGBForce::ParticlePair) else if (energyTypes[termIndex] == CustomGBForce::ParticlePair)
calculateParticlePairEnergyTerm(termIndex, numberOfAtoms, posq, atomParameters, values, exclusions, true, forces, totalEnergy, dEdV, boxSize, invBoxSize); calculateParticlePairEnergyTerm(termIndex, data, numberOfAtoms, posq, atomParameters, true, forces, energy, boxSize, invBoxSize);
else else
calculateParticlePairEnergyTerm(termIndex, numberOfAtoms, posq, atomParameters, values, exclusions, false, forces, totalEnergy, dEdV, boxSize, invBoxSize); calculateParticlePairEnergyTerm(termIndex, data, numberOfAtoms, posq, atomParameters, false, forces, energy, boxSize, invBoxSize);
threads.syncThreads();
}
// Sum the energy derivatives.
for (int atom = data.firstAtom; atom < data.lastAtom; atom++) {
for (int i = 0; i < (int) dEdV.size(); i++) {
float sum = 0.0f;
for (int j = 0; j < (int) threadData.size(); j++)
sum += threadData[j]->dEdV[i][atom];
dEdV[i][atom] = sum;
}
} }
threads.syncThreads();
// Apply the chain rule to evaluate forces. // Apply the chain rule to evaluate forces.
calculateChainRuleForces(numberOfAtoms, posq, atomParameters, values, exclusions, forces, dEdV, boxSize, invBoxSize); calculateChainRuleForces(data, numberOfAtoms, posq, atomParameters, forces, boxSize, invBoxSize);
} }
void CpuCustomGBForce::calculateSingleParticleValue(int index, int numAtoms, float* posq, vector<vector<float> >& values, void CpuCustomGBForce::calculateSingleParticleValue(int index, ThreadData& data, int numAtoms, float* posq, RealOpenMM** atomParameters) {
RealOpenMM** atomParameters) { for (int i = data.firstAtom; i < data.lastAtom; i++) {
values[index].resize(numAtoms); data.expressionSet.setVariable(data.xindex, posq[4*i]);
for (int i = 0; i < numAtoms; i++) { data.expressionSet.setVariable(data.yindex, posq[4*i+1]);
expressionSet.setVariable(xindex, posq[4*i]); data.expressionSet.setVariable(data.zindex, posq[4*i+2]);
expressionSet.setVariable(yindex, posq[4*i+1]);
expressionSet.setVariable(zindex, posq[4*i+2]);
for (int j = 0; j < (int) paramNames.size(); j++) for (int j = 0; j < (int) paramNames.size(); j++)
expressionSet.setVariable(paramIndex[j], atomParameters[i][j]); data.expressionSet.setVariable(data.paramIndex[j], atomParameters[i][j]);
for (int j = 0; j < index; j++) for (int j = 0; j < index; j++)
expressionSet.setVariable(valueIndex[j], values[j][i]); data.expressionSet.setVariable(data.valueIndex[j], values[j][i]);
values[index][i] = (float) valueExpressions[index].evaluate(); values[index][i] = (float) data.valueExpressions[index].evaluate();
} }
} }
void CpuCustomGBForce::calculateParticlePairValue(int index, int numAtoms, float* posq, RealOpenMM** atomParameters, void CpuCustomGBForce::calculateParticlePairValue(int index, ThreadData& data, int numAtoms, float* posq, RealOpenMM** atomParameters,
vector<vector<float> >& values, const vector<set<int> >& exclusions, bool useExclusions, const fvec4& boxSize, const fvec4& invBoxSize) { bool useExclusions, const fvec4& boxSize, const fvec4& invBoxSize) {
values[index].resize(numAtoms);
for (int i = 0; i < numAtoms; i++) for (int i = 0; i < numAtoms; i++)
values[index][i] = 0.0f; values[index][i] = 0.0f;
vector<float>& valueArray = (index == 0 ? data.value0 : values[index]);
if (cutoff) { if (cutoff) {
// Loop over all pairs in the neighbor list. // Loop over all pairs in the neighbor list.
for (int blockIndex = 0; blockIndex < neighborList->getNumBlocks(); blockIndex++) { while (true) {
int blockIndex = gmx_atomic_fetch_add(reinterpret_cast<gmx_atomic_t*>(atomicCounter), 1);
if (blockIndex >= neighborList->getNumBlocks())
break;
const int* blockAtom = &neighborList->getSortedAtoms()[4*blockIndex]; const int* blockAtom = &neighborList->getSortedAtoms()[4*blockIndex];
const vector<int>& neighbors = neighborList->getBlockNeighbors(blockIndex); const vector<int>& neighbors = neighborList->getBlockNeighbors(blockIndex);
const vector<char>& blockExclusions = neighborList->getBlockExclusions(blockIndex); const vector<char>& blockExclusions = neighborList->getBlockExclusions(blockIndex);
...@@ -191,8 +314,8 @@ void CpuCustomGBForce::calculateParticlePairValue(int index, int numAtoms, float ...@@ -191,8 +314,8 @@ void CpuCustomGBForce::calculateParticlePairValue(int index, int numAtoms, float
int second = blockAtom[k]; int second = blockAtom[k];
if (useExclusions && exclusions[first].find(second) != exclusions[first].end()) if (useExclusions && exclusions[first].find(second) != exclusions[first].end())
continue; continue;
calculateOnePairValue(index, first, second, posq, atomParameters, values, boxSize, invBoxSize); calculateOnePairValue(index, first, second, data, posq, atomParameters, valueArray, boxSize, invBoxSize);
calculateOnePairValue(index, second, first, posq, atomParameters, values, boxSize, invBoxSize); calculateOnePairValue(index, second, first, data, posq, atomParameters, valueArray, boxSize, invBoxSize);
} }
} }
} }
...@@ -201,19 +324,22 @@ void CpuCustomGBForce::calculateParticlePairValue(int index, int numAtoms, float ...@@ -201,19 +324,22 @@ void CpuCustomGBForce::calculateParticlePairValue(int index, int numAtoms, float
else { else {
// Perform an O(N^2) loop over all atom pairs. // Perform an O(N^2) loop over all atom pairs.
for (int i = 0; i < numAtoms; i++) { while (true) {
int i = gmx_atomic_fetch_add(reinterpret_cast<gmx_atomic_t*>(atomicCounter), 1);
if (i >= numAtoms)
break;
for (int j = i+1; j < numAtoms; j++) { for (int j = i+1; j < numAtoms; j++) {
if (useExclusions && exclusions[i].find(j) != exclusions[i].end()) if (useExclusions && exclusions[i].find(j) != exclusions[i].end())
continue; continue;
calculateOnePairValue(index, i, j, posq, atomParameters, values, boxSize, invBoxSize); calculateOnePairValue(index, i, j, data, posq, atomParameters, valueArray, boxSize, invBoxSize);
calculateOnePairValue(index, j, i, posq, atomParameters, values, boxSize, invBoxSize); calculateOnePairValue(index, j, i, data, posq, atomParameters, valueArray, boxSize, invBoxSize);
} }
} }
} }
} }
void CpuCustomGBForce::calculateOnePairValue(int index, int atom1, int atom2, float* posq, RealOpenMM** atomParameters, void CpuCustomGBForce::calculateOnePairValue(int index, int atom1, int atom2, ThreadData& data, float* posq, RealOpenMM** atomParameters,
vector<vector<float> >& values, const fvec4& boxSize, const fvec4& invBoxSize) { vector<float>& valueArray, const fvec4& boxSize, const fvec4& invBoxSize) {
fvec4 deltaR; fvec4 deltaR;
fvec4 pos1(posq+4*atom1); fvec4 pos1(posq+4*atom1);
fvec4 pos2(posq+4*atom2); fvec4 pos2(posq+4*atom2);
...@@ -223,45 +349,46 @@ void CpuCustomGBForce::calculateOnePairValue(int index, int atom1, int atom2, fl ...@@ -223,45 +349,46 @@ void CpuCustomGBForce::calculateOnePairValue(int index, int atom1, int atom2, fl
if (cutoff && r >= cutoffDistance) if (cutoff && r >= cutoffDistance)
return; return;
for (int i = 0; i < (int) paramNames.size(); i++) { for (int i = 0; i < (int) paramNames.size(); i++) {
expressionSet.setVariable(particleParamIndex[i*2], atomParameters[atom1][i]); data.expressionSet.setVariable(data.particleParamIndex[i*2], atomParameters[atom1][i]);
expressionSet.setVariable(particleParamIndex[i*2+1], atomParameters[atom2][i]); data.expressionSet.setVariable(data.particleParamIndex[i*2+1], atomParameters[atom2][i]);
} }
expressionSet.setVariable(rindex, r); data.expressionSet.setVariable(data.rindex, r);
for (int i = 0; i < index; i++) { for (int i = 0; i < index; i++) {
expressionSet.setVariable(particleValueIndex[i*2], values[i][atom1]); data.expressionSet.setVariable(data.particleValueIndex[i*2], values[i][atom1]);
expressionSet.setVariable(particleValueIndex[i*2+1], values[i][atom2]); data.expressionSet.setVariable(data.particleValueIndex[i*2+1], values[i][atom2]);
} }
values[index][atom1] += (float) valueExpressions[index].evaluate(); valueArray[atom1] += (float) data.valueExpressions[index].evaluate();
} }
void CpuCustomGBForce::calculateSingleParticleEnergyTerm(int index, int numAtoms, float* posq, const vector<vector<float> >& values, void CpuCustomGBForce::calculateSingleParticleEnergyTerm(int index, ThreadData& data, int numAtoms, float* posq,
RealOpenMM** atomParameters, float* forces, double* totalEnergy, RealOpenMM** atomParameters, float* forces, double& totalEnergy) {
vector<vector<float> >& dEdV) { for (int i = data.firstAtom; i < data.lastAtom; i++) {
for (int i = 0; i < numAtoms; i++) { data.expressionSet.setVariable(data.xindex, posq[4*i]);
expressionSet.setVariable(xindex, posq[4*i]); data.expressionSet.setVariable(data.yindex, posq[4*i+1]);
expressionSet.setVariable(yindex, posq[4*i+1]); data.expressionSet.setVariable(data.zindex, posq[4*i+2]);
expressionSet.setVariable(zindex, posq[4*i+2]);
for (int j = 0; j < (int) paramNames.size(); j++) for (int j = 0; j < (int) paramNames.size(); j++)
expressionSet.setVariable(paramIndex[j], atomParameters[i][j]); data.expressionSet.setVariable(data.paramIndex[j], atomParameters[i][j]);
for (int j = 0; j < (int) valueNames.size(); j++) for (int j = 0; j < (int) valueNames.size(); j++)
expressionSet.setVariable(valueIndex[j], values[j][i]); data.expressionSet.setVariable(data.valueIndex[j], values[j][i]);
if (totalEnergy != NULL) if (includeEnergy)
*totalEnergy += (float) energyExpressions[index].evaluate(); totalEnergy += (float) data.energyExpressions[index].evaluate();
for (int j = 0; j < (int) valueNames.size(); j++) for (int j = 0; j < (int) valueNames.size(); j++)
dEdV[j][i] += (float) energyDerivExpressions[index][j].evaluate(); data.dEdV[j][i] += (float) data.energyDerivExpressions[index][j].evaluate();
forces[4*i+0] -= (float) energyGradientExpressions[index][0].evaluate(); forces[4*i+0] -= (float) data.energyGradientExpressions[index][0].evaluate();
forces[4*i+1] -= (float) energyGradientExpressions[index][1].evaluate(); forces[4*i+1] -= (float) data.energyGradientExpressions[index][1].evaluate();
forces[4*i+2] -= (float) energyGradientExpressions[index][2].evaluate(); forces[4*i+2] -= (float) data.energyGradientExpressions[index][2].evaluate();
} }
} }
void CpuCustomGBForce::calculateParticlePairEnergyTerm(int index, int numAtoms, float* posq, RealOpenMM** atomParameters, void CpuCustomGBForce::calculateParticlePairEnergyTerm(int index, ThreadData& data, int numAtoms, float* posq, RealOpenMM** atomParameters,
const vector<vector<float> >& values, const vector<set<int> >& exclusions, bool useExclusions, bool useExclusions, float* forces, double& totalEnergy, const fvec4& boxSize, const fvec4& invBoxSize) {
float* forces, double* totalEnergy, vector<vector<float> >& dEdV, const fvec4& boxSize, const fvec4& invBoxSize) {
if (cutoff) { if (cutoff) {
// Loop over all pairs in the neighbor list. // Loop over all pairs in the neighbor list.
for (int blockIndex = 0; blockIndex < neighborList->getNumBlocks(); blockIndex++) { while (true) {
int blockIndex = gmx_atomic_fetch_add(reinterpret_cast<gmx_atomic_t*>(atomicCounter), 1);
if (blockIndex >= neighborList->getNumBlocks())
break;
const int* blockAtom = &neighborList->getSortedAtoms()[4*blockIndex]; const int* blockAtom = &neighborList->getSortedAtoms()[4*blockIndex];
const vector<int>& neighbors = neighborList->getBlockNeighbors(blockIndex); const vector<int>& neighbors = neighborList->getBlockNeighbors(blockIndex);
const vector<char>& blockExclusions = neighborList->getBlockExclusions(blockIndex); const vector<char>& blockExclusions = neighborList->getBlockExclusions(blockIndex);
...@@ -272,7 +399,7 @@ void CpuCustomGBForce::calculateParticlePairEnergyTerm(int index, int numAtoms, ...@@ -272,7 +399,7 @@ void CpuCustomGBForce::calculateParticlePairEnergyTerm(int index, int numAtoms,
int second = blockAtom[k]; int second = blockAtom[k];
if (useExclusions && exclusions[first].find(second) != exclusions[first].end()) if (useExclusions && exclusions[first].find(second) != exclusions[first].end())
continue; continue;
calculateOnePairEnergyTerm(index, first, second, posq, atomParameters, values, forces, totalEnergy, dEdV, boxSize, invBoxSize); calculateOnePairEnergyTerm(index, first, second, data, posq, atomParameters, forces, totalEnergy, boxSize, invBoxSize);
} }
} }
} }
...@@ -281,19 +408,21 @@ void CpuCustomGBForce::calculateParticlePairEnergyTerm(int index, int numAtoms, ...@@ -281,19 +408,21 @@ void CpuCustomGBForce::calculateParticlePairEnergyTerm(int index, int numAtoms,
else { else {
// Perform an O(N^2) loop over all atom pairs. // Perform an O(N^2) loop over all atom pairs.
for (int i = 0; i < numAtoms; i++) { while (true) {
int i = gmx_atomic_fetch_add(reinterpret_cast<gmx_atomic_t*>(atomicCounter), 1);
if (i >= numAtoms)
break;
for (int j = i+1; j < numAtoms; j++) { for (int j = i+1; j < numAtoms; j++) {
if (useExclusions && exclusions[i].find(j) != exclusions[i].end()) if (useExclusions && exclusions[i].find(j) != exclusions[i].end())
continue; continue;
calculateOnePairEnergyTerm(index, i, j, posq, atomParameters, values, forces, totalEnergy, dEdV, boxSize, invBoxSize); calculateOnePairEnergyTerm(index, i, j, data, posq, atomParameters, forces, totalEnergy, boxSize, invBoxSize);
} }
} }
} }
} }
void CpuCustomGBForce::calculateOnePairEnergyTerm(int index, int atom1, int atom2, float* posq, RealOpenMM** atomParameters, void CpuCustomGBForce::calculateOnePairEnergyTerm(int index, int atom1, int atom2, ThreadData& data, float* posq, RealOpenMM** atomParameters,
const vector<vector<float> >& values, float* forces, double* totalEnergy, float* forces, double& totalEnergy, const fvec4& boxSize, const fvec4& invBoxSize) {
vector<vector<float> >& dEdV, const fvec4& boxSize, const fvec4& invBoxSize) {
// Compute the displacement. // Compute the displacement.
fvec4 deltaR; fvec4 deltaR;
...@@ -308,37 +437,39 @@ void CpuCustomGBForce::calculateOnePairEnergyTerm(int index, int atom1, int atom ...@@ -308,37 +437,39 @@ void CpuCustomGBForce::calculateOnePairEnergyTerm(int index, int atom1, int atom
// Record variables for evaluating expressions. // Record variables for evaluating expressions.
for (int i = 0; i < (int) paramNames.size(); i++) { for (int i = 0; i < (int) paramNames.size(); i++) {
expressionSet.setVariable(particleParamIndex[i*2], atomParameters[atom1][i]); data.expressionSet.setVariable(data.particleParamIndex[i*2], atomParameters[atom1][i]);
expressionSet.setVariable(particleParamIndex[i*2+1], atomParameters[atom2][i]); data.expressionSet.setVariable(data.particleParamIndex[i*2+1], atomParameters[atom2][i]);
} }
expressionSet.setVariable(rindex, r); data.expressionSet.setVariable(data.rindex, r);
for (int i = 0; i < (int) valueNames.size(); i++) { for (int i = 0; i < (int) valueNames.size(); i++) {
expressionSet.setVariable(particleValueIndex[i*2], values[i][atom1]); data.expressionSet.setVariable(data.particleValueIndex[i*2], values[i][atom1]);
expressionSet.setVariable(particleValueIndex[i*2+1], values[i][atom2]); data.expressionSet.setVariable(data.particleValueIndex[i*2+1], values[i][atom2]);
} }
// Evaluate the energy and its derivatives. // Evaluate the energy and its derivatives.
if (totalEnergy != NULL) if (includeEnergy)
*totalEnergy += (float) energyExpressions[index].evaluate(); totalEnergy += (float) data.energyExpressions[index].evaluate();
float dEdR = (float) energyDerivExpressions[index][0].evaluate(); float dEdR = (float) data.energyDerivExpressions[index][0].evaluate();
dEdR *= 1/r; dEdR *= 1/r;
fvec4 result = deltaR*dEdR; fvec4 result = deltaR*dEdR;
(fvec4(forces+4*atom1)-result).store(forces+4*atom1); (fvec4(forces+4*atom1)-result).store(forces+4*atom1);
(fvec4(forces+4*atom2)+result).store(forces+4*atom2); (fvec4(forces+4*atom2)+result).store(forces+4*atom2);
for (int i = 0; i < (int) valueNames.size(); i++) { for (int i = 0; i < (int) valueNames.size(); i++) {
dEdV[i][atom1] += (float) energyDerivExpressions[index][2*i+1].evaluate(); data.dEdV[i][atom1] += (float) data.energyDerivExpressions[index][2*i+1].evaluate();
dEdV[i][atom2] += (float) energyDerivExpressions[index][2*i+2].evaluate(); data.dEdV[i][atom2] += (float) data.energyDerivExpressions[index][2*i+2].evaluate();
} }
} }
void CpuCustomGBForce::calculateChainRuleForces(int numAtoms, float* posq, RealOpenMM** atomParameters, void CpuCustomGBForce::calculateChainRuleForces(ThreadData& data, int numAtoms, float* posq, RealOpenMM** atomParameters,
const vector<vector<float> >& values, const vector<set<int> >& exclusions, float* forces, vector<vector<float> >& dEdV, float* forces, const fvec4& boxSize, const fvec4& invBoxSize) {
const fvec4& boxSize, const fvec4& invBoxSize) {
if (cutoff) { if (cutoff) {
// Loop over all pairs in the neighbor list. // Loop over all pairs in the neighbor list.
for (int blockIndex = 0; blockIndex < neighborList->getNumBlocks(); blockIndex++) { while (true) {
int blockIndex = gmx_atomic_fetch_add(reinterpret_cast<gmx_atomic_t*>(atomicCounter), 1);
if (blockIndex >= neighborList->getNumBlocks())
break;
const int* blockAtom = &neighborList->getSortedAtoms()[4*blockIndex]; const int* blockAtom = &neighborList->getSortedAtoms()[4*blockIndex];
const vector<int>& neighbors = neighborList->getBlockNeighbors(blockIndex); const vector<int>& neighbors = neighborList->getBlockNeighbors(blockIndex);
const vector<char>& blockExclusions = neighborList->getBlockExclusions(blockIndex); const vector<char>& blockExclusions = neighborList->getBlockExclusions(blockIndex);
...@@ -348,8 +479,8 @@ void CpuCustomGBForce::calculateChainRuleForces(int numAtoms, float* posq, RealO ...@@ -348,8 +479,8 @@ void CpuCustomGBForce::calculateChainRuleForces(int numAtoms, float* posq, RealO
if ((blockExclusions[i] & (1<<k)) == 0) { if ((blockExclusions[i] & (1<<k)) == 0) {
int second = blockAtom[k]; int second = blockAtom[k];
bool isExcluded = (exclusions[first].find(second) != exclusions[first].end()); bool isExcluded = (exclusions[first].find(second) != exclusions[first].end());
calculateOnePairChainRule(first, second, posq, atomParameters, values, forces, dEdV, isExcluded, boxSize, invBoxSize); calculateOnePairChainRule(first, second, data, posq, atomParameters, forces, isExcluded, boxSize, invBoxSize);
calculateOnePairChainRule(second, first, posq, atomParameters, values, forces, dEdV, isExcluded, boxSize, invBoxSize); calculateOnePairChainRule(second, first, data, posq, atomParameters, forces, isExcluded, boxSize, invBoxSize);
} }
} }
} }
...@@ -358,46 +489,49 @@ void CpuCustomGBForce::calculateChainRuleForces(int numAtoms, float* posq, RealO ...@@ -358,46 +489,49 @@ void CpuCustomGBForce::calculateChainRuleForces(int numAtoms, float* posq, RealO
else { else {
// Perform an O(N^2) loop over all atom pairs. // Perform an O(N^2) loop over all atom pairs.
for (int i = 0; i < numAtoms; i++) { while (true) {
int i = gmx_atomic_fetch_add(reinterpret_cast<gmx_atomic_t*>(atomicCounter), 1);
if (i >= numAtoms)
break;
for (int j = i+1; j < numAtoms; j++) { for (int j = i+1; j < numAtoms; j++) {
bool isExcluded = (exclusions[i].find(j) != exclusions[i].end()); bool isExcluded = (exclusions[i].find(j) != exclusions[i].end());
calculateOnePairChainRule(i, j, posq, atomParameters, values, forces, dEdV, isExcluded, boxSize, invBoxSize); calculateOnePairChainRule(i, j, data, posq, atomParameters, forces, isExcluded, boxSize, invBoxSize);
calculateOnePairChainRule(j, i, posq, atomParameters, values, forces, dEdV, isExcluded, boxSize, invBoxSize); calculateOnePairChainRule(j, i, data, posq, atomParameters, forces, isExcluded, boxSize, invBoxSize);
} }
} }
} }
// Compute chain rule terms for computed values that depend explicitly on particle coordinates. // Compute chain rule terms for computed values that depend explicitly on particle coordinates.
for (int i = 0; i < numAtoms; i++) { for (int i = data.firstAtom; i < data.lastAtom; i++) {
expressionSet.setVariable(xindex, posq[4*i]); data.expressionSet.setVariable(data.xindex, posq[4*i]);
expressionSet.setVariable(yindex, posq[4*i+1]); data.expressionSet.setVariable(data.yindex, posq[4*i+1]);
expressionSet.setVariable(zindex, posq[4*i+2]); data.expressionSet.setVariable(data.zindex, posq[4*i+2]);
for (int j = 0; j < (int) paramNames.size(); j++) for (int j = 0; j < (int) paramNames.size(); j++)
expressionSet.setVariable(paramIndex[j], atomParameters[i][j]); data.expressionSet.setVariable(data.paramIndex[j], atomParameters[i][j]);
for (int j = 1; j < (int) valueNames.size(); j++) { for (int j = 1; j < (int) valueNames.size(); j++) {
expressionSet.setVariable(valueIndex[j-1], values[j-1][i]); data.expressionSet.setVariable(data.valueIndex[j-1], values[j-1][i]);
dVdX[j] = 0.0; data.dVdX[j] = 0.0;
dVdY[j] = 0.0; data.dVdY[j] = 0.0;
dVdZ[j] = 0.0; data.dVdZ[j] = 0.0;
for (int k = 1; k < j; k++) { for (int k = 1; k < j; k++) {
float dVdV = (float) valueDerivExpressions[j][k].evaluate(); float dVdV = (float) data.valueDerivExpressions[j][k].evaluate();
dVdX[j] += dVdV*dVdX[k]; data.dVdX[j] += dVdV*data.dVdX[k];
dVdY[j] += dVdV*dVdY[k]; data.dVdY[j] += dVdV*data.dVdY[k];
dVdZ[j] += dVdV*dVdZ[k]; data.dVdZ[j] += dVdV*data.dVdZ[k];
} }
dVdX[j] += (float) valueGradientExpressions[j][0].evaluate(); data.dVdX[j] += (float) data.valueGradientExpressions[j][0].evaluate();
dVdY[j] += (float) valueGradientExpressions[j][1].evaluate(); data.dVdY[j] += (float) data.valueGradientExpressions[j][1].evaluate();
dVdZ[j] += (float) valueGradientExpressions[j][2].evaluate(); data.dVdZ[j] += (float) data.valueGradientExpressions[j][2].evaluate();
forces[4*i+0] -= dEdV[j][i]*dVdX[j]; forces[4*i+0] -= dEdV[j][i]*data.dVdX[j];
forces[4*i+1] -= dEdV[j][i]*dVdY[j]; forces[4*i+1] -= dEdV[j][i]*data.dVdY[j];
forces[4*i+2] -= dEdV[j][i]*dVdZ[j]; forces[4*i+2] -= dEdV[j][i]*data.dVdZ[j];
} }
} }
} }
void CpuCustomGBForce::calculateOnePairChainRule(int atom1, int atom2, float* posq, RealOpenMM** atomParameters, void CpuCustomGBForce::calculateOnePairChainRule(int atom1, int atom2, ThreadData& data, float* posq, RealOpenMM** atomParameters,
const vector<vector<float> >& values, float* forces, vector<vector<float> >& dEdV, bool isExcluded, const fvec4& boxSize, const fvec4& invBoxSize) { float* forces, bool isExcluded, const fvec4& boxSize, const fvec4& invBoxSize) {
// Compute the displacement. // Compute the displacement.
fvec4 deltaR; fvec4 deltaR;
...@@ -412,40 +546,40 @@ void CpuCustomGBForce::calculateOnePairChainRule(int atom1, int atom2, float* po ...@@ -412,40 +546,40 @@ void CpuCustomGBForce::calculateOnePairChainRule(int atom1, int atom2, float* po
// Record variables for evaluating expressions. // Record variables for evaluating expressions.
for (int i = 0; i < (int) paramNames.size(); i++) { for (int i = 0; i < (int) paramNames.size(); i++) {
expressionSet.setVariable(particleParamIndex[i*2], atomParameters[atom1][i]); data.expressionSet.setVariable(data.particleParamIndex[i*2], atomParameters[atom1][i]);
expressionSet.setVariable(particleParamIndex[i*2+1], atomParameters[atom2][i]); data.expressionSet.setVariable(data.particleParamIndex[i*2+1], atomParameters[atom2][i]);
} }
expressionSet.setVariable(rindex, r); data.expressionSet.setVariable(data.rindex, r);
expressionSet.setVariable(particleValueIndex[0], values[0][atom1]); data.expressionSet.setVariable(data.particleValueIndex[0], values[0][atom1]);
expressionSet.setVariable(particleValueIndex[1], values[0][atom2]); data.expressionSet.setVariable(data.particleValueIndex[1], values[0][atom2]);
// Evaluate the derivative of each parameter with respect to position and apply forces. // Evaluate the derivative of each parameter with respect to position and apply forces.
float rinv = 1/r; float rinv = 1/r;
deltaR *= rinv; deltaR *= rinv;
if (!isExcluded || valueTypes[0] != CustomGBForce::ParticlePair) { if (!isExcluded || valueTypes[0] != CustomGBForce::ParticlePair) {
dVdR1[0] = (float) valueDerivExpressions[0][0].evaluate(); data.dVdR1[0] = (float) data.valueDerivExpressions[0][0].evaluate();
dVdR2[0] = -dVdR1[0]; data.dVdR2[0] = -data.dVdR1[0];
(fvec4(forces+4*atom1)-deltaR*(dEdV[0][atom1]*dVdR1[0])).store(forces+4*atom1); (fvec4(forces+4*atom1)-deltaR*(dEdV[0][atom1]*data.dVdR1[0])).store(forces+4*atom1);
(fvec4(forces+4*atom2)-deltaR*(dEdV[0][atom1]*dVdR2[0])).store(forces+4*atom2); (fvec4(forces+4*atom2)-deltaR*(dEdV[0][atom1]*data.dVdR2[0])).store(forces+4*atom2);
} }
for (int i = 0; i < (int) paramNames.size(); i++) for (int i = 0; i < (int) paramNames.size(); i++)
expressionSet.setVariable(paramIndex[i], atomParameters[atom1][i]); data.expressionSet.setVariable(data.paramIndex[i], atomParameters[atom1][i]);
expressionSet.setVariable(valueIndex[0], values[0][atom1]); data.expressionSet.setVariable(data.valueIndex[0], values[0][atom1]);
for (int i = 1; i < (int) valueNames.size(); i++) { for (int i = 1; i < (int) valueNames.size(); i++) {
expressionSet.setVariable(valueIndex[i], values[i][atom1]); data.expressionSet.setVariable(data.valueIndex[i], values[i][atom1]);
expressionSet.setVariable(xindex, posq[4*atom1]); data.expressionSet.setVariable(data.xindex, posq[4*atom1]);
expressionSet.setVariable(yindex, posq[4*atom1+1]); data.expressionSet.setVariable(data.yindex, posq[4*atom1+1]);
expressionSet.setVariable(zindex, posq[4*atom1+2]); data.expressionSet.setVariable(data.zindex, posq[4*atom1+2]);
dVdR1[i] = 0.0; data.dVdR1[i] = 0.0;
dVdR2[i] = 0.0; data.dVdR2[i] = 0.0;
for (int j = 0; j < i; j++) { for (int j = 0; j < i; j++) {
float dVdV = (float) valueDerivExpressions[i][j].evaluate(); float dVdV = (float) data.valueDerivExpressions[i][j].evaluate();
dVdR1[i] += dVdV*dVdR1[j]; data.dVdR1[i] += dVdV*data.dVdR1[j];
dVdR2[i] += dVdV*dVdR2[j]; data.dVdR2[i] += dVdV*data.dVdR2[j];
} }
(fvec4(forces+4*atom1)-deltaR*(dEdV[i][atom1]*dVdR1[i])).store(forces+4*atom1); (fvec4(forces+4*atom1)-deltaR*(dEdV[i][atom1]*data.dVdR1[i])).store(forces+4*atom1);
(fvec4(forces+4*atom2)-deltaR*(dEdV[i][atom1]*dVdR2[i])).store(forces+4*atom2); (fvec4(forces+4*atom2)-deltaR*(dEdV[i][atom1]*data.dVdR2[i])).store(forces+4*atom2);
} }
} }
......
platforms/cpu/src/CpuKernels.cpp
View file @ 6a4ac837
...@@ -957,8 +957,8 @@ void CpuCalcCustomGBForceKernel::initialize(const System& system, const CustomGB ...@@ -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++) for (map<string, Lepton::CustomFunction*>::iterator iter = functions.begin(); iter != functions.end(); iter++)
delete iter->second; delete iter->second;
ixn = new CpuCustomGBForce(numParticles, valueExpressions, valueDerivExpressions, valueGradientExpressions, valueNames, valueTypes, energyExpressions, ixn = new CpuCustomGBForce(numParticles, exclusions, valueExpressions, valueDerivExpressions, valueGradientExpressions, valueNames, valueTypes, energyExpressions,
energyDerivExpressions, energyGradientExpressions, energyTypes, particleParameterNames); energyDerivExpressions, energyGradientExpressions, energyTypes, particleParameterNames, data.threads);
data.isPeriodic = (force.getNonbondedMethod() == CustomGBForce::CutoffPeriodic); data.isPeriodic = (force.getNonbondedMethod() == CustomGBForce::CutoffPeriodic);
} }
...@@ -977,7 +977,7 @@ double CpuCalcCustomGBForceKernel::execute(ContextImpl& context, bool includeFor ...@@ -977,7 +977,7 @@ double CpuCalcCustomGBForceKernel::execute(ContextImpl& context, bool includeFor
map<string, double> globalParameters; map<string, double> globalParameters;
for (int i = 0; i < (int) globalParameterNames.size(); i++) for (int i = 0; i < (int) globalParameterNames.size(); i++)
globalParameters[globalParameterNames[i]] = context.getParameter(globalParameterNames[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; return energy;
} }
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment