Skip to content

GitLab

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

Parallelized CpuCustomGBForce

parent bb70341a
Expand all 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
This diff is collapsed.
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