Skip to content

GitLab

Commit bf0b43c5 authored by peastman's avatar peastman
Browse files

Optimization to CustomNonbondedForce and CustomIntegrator on CPU

parent 1c416f54
Show whitespace changes
Inline Side-by-side
Showing with 81 additions and 71 deletions
platforms/cpu/include/CpuCustomNonbondedForce.h
View file @ bf0b43c5
......@@ -183,8 +183,8 @@ public:
Lepton::CompiledExpression forceExpression;
std::vector<Lepton::CompiledExpression> energyParamDerivExpressions;
CompiledExpressionSet expressionSet;
std::vector<int> particleParamIndex;
int rIndex;
std::vector<double> particleParam;
double r;
std::vector<RealOpenMM> energyParamDerivs;
};
......
platforms/cpu/src/CpuCustomGBForce.cpp
View file @ bf0b43c5
......@@ -513,7 +513,8 @@ void CpuCustomGBForce::calculateOnePairEnergyTerm(int index, int atom1, int atom
getDeltaR(pos2, pos1, deltaR, r2, periodic, boxSize, invBoxSize);
if (cutoff && r2 >= cutoffDistance2)
return;
data.r = sqrtf(r2);
float r = sqrtf(r2);
data.r = r;
// Record variables for evaluating expressions.
......@@ -531,7 +532,7 @@ void CpuCustomGBForce::calculateOnePairEnergyTerm(int index, int atom1, int atom
if (includeEnergy)
totalEnergy += (float) data.energyExpressions[index].evaluate();
float dEdR = (float) data.energyDerivExpressions[index][0].evaluate();
dEdR *= 1/data.r;
dEdR *= 1/r;
fvec4 result = deltaR*dEdR;
(fvec4(forces+4*atom1)-result).store(forces+4*atom1);
(fvec4(forces+4*atom2)+result).store(forces+4*atom2);
......@@ -634,7 +635,8 @@ void CpuCustomGBForce::calculateOnePairChainRule(int atom1, int atom2, ThreadDat
getDeltaR(pos2, pos1, deltaR, r2, periodic, boxSize, invBoxSize);
if (cutoff && r2 >= cutoffDistance2)
return;
data.r = sqrtf(r2);
float r = sqrtf(r2);
data.r = r;
// Record variables for evaluating expressions.
......@@ -652,7 +654,7 @@ void CpuCustomGBForce::calculateOnePairChainRule(int atom1, int atom2, ThreadDat
// Evaluate the derivative of each parameter with respect to position and apply forces.
float rinv = 1/data.r;
float rinv = 1/r;
deltaR *= rinv;
fvec4 f1(0.0f), f2(0.0f);
if (!isExcluded || valueTypes[0] != CustomGBForce::ParticlePair) {
......
platforms/cpu/src/CpuCustomNonbondedForce.cpp
View file @ bf0b43c5
......@@ -46,19 +46,25 @@ public:
CpuCustomNonbondedForce::ThreadData::ThreadData(const Lepton::CompiledExpression& energyExpression, const Lepton::CompiledExpression& forceExpression,
const vector<string>& parameterNames, const std::vector<Lepton::CompiledExpression> energyParamDerivExpressions) :
energyExpression(energyExpression), forceExpression(forceExpression), energyParamDerivExpressions(energyParamDerivExpressions) {
expressionSet.registerExpression(this->energyExpression);
expressionSet.registerExpression(this->forceExpression);
for (int i = 0; i < this->energyParamDerivExpressions.size(); i++)
expressionSet.registerExpression(this->energyParamDerivExpressions[i]);
rIndex = expressionSet.getVariableIndex("r");
map<string, double*> variableLocations;
variableLocations["r"] = &r;
particleParam.resize(2*parameterNames.size());
for (int i = 0; i < (int) parameterNames.size(); i++) {
for (int j = 1; j < 3; j++) {
for (int j = 0; j < 2; j++) {
stringstream name;
name << parameterNames[i] << j;
particleParamIndex.push_back(expressionSet.getVariableIndex(name.str()));
name << parameterNames[i] << (j+1);
variableLocations[name.str()] = &particleParam[i*2+j];
}
}
energyParamDerivs.resize(energyParamDerivExpressions.size());
this->energyExpression.setVariableLocations(variableLocations);
this->forceExpression.setVariableLocations(variableLocations);
expressionSet.registerExpression(this->energyExpression);
expressionSet.registerExpression(this->forceExpression);
for (int i = 0; i < this->energyParamDerivExpressions.size(); i++) {
this->energyParamDerivExpressions[i].setVariableLocations(variableLocations);
expressionSet.registerExpression(this->energyParamDerivExpressions[i]);
}
}
CpuCustomNonbondedForce::CpuCustomNonbondedForce(const Lepton::CompiledExpression& energyExpression,
......@@ -187,8 +193,8 @@ void CpuCustomNonbondedForce::threadComputeForce(ThreadPool& threads, int thread
int atom1 = groupInteractions[i].first;
int atom2 = groupInteractions[i].second;
for (int j = 0; j < (int) paramNames.size(); j++) {
data.expressionSet.setVariable(data.particleParamIndex[j*2], atomParameters[atom1][j]);
data.expressionSet.setVariable(data.particleParamIndex[j*2+1], atomParameters[atom2][j]);
data.particleParam[j*2] = atomParameters[atom1][j];
data.particleParam[j*2+1] = atomParameters[atom2][j];
}
calculateOneIxn(atom1, atom2, data, forces, energy, boxSize, invBoxSize);
}
......@@ -207,12 +213,12 @@ void CpuCustomNonbondedForce::threadComputeForce(ThreadPool& threads, int thread
for (int i = 0; i < (int) neighbors.size(); i++) {
int first = neighbors[i];
for (int j = 0; j < (int) paramNames.size(); j++)
data.expressionSet.setVariable(data.particleParamIndex[j*2], atomParameters[first][j]);
data.particleParam[j*2] = atomParameters[first][j];
for (int k = 0; k < blockSize; k++) {
if ((exclusions[i] & (1<<k)) == 0) {
int second = blockAtom[k];
for (int j = 0; j < (int) paramNames.size(); j++)
data.expressionSet.setVariable(data.particleParamIndex[j*2+1], atomParameters[second][j]);
data.particleParam[j*2+1] = atomParameters[second][j];
calculateOneIxn(first, second, data, forces, energy, boxSize, invBoxSize);
}
}
......@@ -229,8 +235,8 @@ void CpuCustomNonbondedForce::threadComputeForce(ThreadPool& threads, int thread
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++) {
data.expressionSet.setVariable(data.particleParamIndex[j*2], atomParameters[ii][j]);
data.expressionSet.setVariable(data.particleParamIndex[j*2+1], atomParameters[jj][j]);
data.particleParam[j*2] = atomParameters[ii][j];
data.particleParam[j*2+1] = atomParameters[jj][j];
}
calculateOneIxn(ii, jj, data, forces, energy, boxSize, invBoxSize);
}
......@@ -251,10 +257,10 @@ void CpuCustomNonbondedForce::calculateOneIxn(int ii, int jj, ThreadData& data,
if (cutoff && r2 >= cutoffDistance*cutoffDistance)
return;
float r = sqrtf(r2);
data.r = r;
// accumulate forces
data.expressionSet.setVariable(data.rIndex, r);
double dEdR = (includeForce ? data.forceExpression.evaluate()/r : 0.0);
double energy = (includeEnergy ? data.energyExpression.evaluate() : 0.0);
double switchValue = 1.0;
......
platforms/reference/include/ReferenceCustomDynamics.h
View file @ bf0b43c5
......@@ -51,13 +51,14 @@ private:
std::vector<CustomIntegratorUtilities::Comparison> comparisons;
std::vector<bool> invalidatesForces, needsForces, needsEnergy, computeBothForceAndEnergy;
std::vector<int> forceGroupFlags, blockEnd;
RealOpenMM energy;
std::map<std::string, double> energyParamDerivs;
Lepton::CompiledExpression kineticEnergyExpression;
bool kineticEnergyNeedsForce;
CompiledExpressionSet expressionSet;
int xIndex, vIndex, mIndex, fIndex, energyIndex, gaussianIndex, uniformIndex;
std::vector<int> forceVariableIndex, energyVariableIndex, perDofVariableIndex, stepVariableIndex;
double x, v, m, f, energy, gaussian, uniform;
int xIndex, vIndex;
std::vector<int> perDofVariableIndex, stepVariableIndex;
std::vector<double> perDofVariable;
void initialize(OpenMM::ContextImpl& context, std::vector<RealOpenMM>& masses, std::map<std::string, RealOpenMM>& globals);
......@@ -65,7 +66,7 @@ private:
void computePerDof(int numberOfAtoms, std::vector<OpenMM::RealVec>& results, const std::vector<OpenMM::RealVec>& atomCoordinates,
const std::vector<OpenMM::RealVec>& velocities, const std::vector<OpenMM::RealVec>& forces, const std::vector<RealOpenMM>& masses,
const std::vector<std::vector<OpenMM::RealVec> >& perDof, const Lepton::CompiledExpression& expression, int forceIndex);
const std::vector<std::vector<OpenMM::RealVec> >& perDof, const Lepton::CompiledExpression& expression);
void recordChangedParameters(OpenMM::ContextImpl& context, std::map<std::string, RealOpenMM>& globals);
......
platforms/reference/src/SimTKReference/ReferenceCustomDynamics.cpp
View file @ bf0b43c5
......@@ -78,11 +78,6 @@ ReferenceCustomDynamics::ReferenceCustomDynamics(int numberOfAtoms, const Custom
string expression;
integrator.getComputationStep(i, stepType[i], stepVariable[i], expression);
}
kineticEnergyExpression = Parser::parse(integrator.getKineticEnergyExpression()).optimize().createCompiledExpression();
expressionSet.registerExpression(kineticEnergyExpression);
kineticEnergyNeedsForce = false;
if (kineticEnergyExpression.getVariables().find("f") != kineticEnergyExpression.getVariables().end())
kineticEnergyNeedsForce = true;
}
/**---------------------------------------------------------------------------------------
......@@ -98,6 +93,28 @@ void ReferenceCustomDynamics::initialize(ContextImpl& context, vector<RealOpenMM
// Some initialization can't be done in the constructor, since we need a ContextImpl from which to get the list of
// Context parameters. Instead, we do it the first time update() or computeKineticEnergy() is called.
std::map<std::string, double*> variableLocations;
variableLocations["x"] = &x;
variableLocations["v"] = &v;
variableLocations["m"] = &m;
variableLocations["f"] = &f;
variableLocations["energy"] = &energy;
variableLocations["gaussian"] = &gaussian;
variableLocations["uniform"] = &uniform;
perDofVariable.resize(integrator.getNumPerDofVariables());
for (int i = 0; i < integrator.getNumPerDofVariables(); i++)
variableLocations[integrator.getPerDofVariableName(i)] = &perDofVariable[i];
for (int i = 0; i < 32; i++) {
stringstream fname;
fname << "f" << i;
variableLocations[fname.str()] = &f;
stringstream ename;
ename << "energy" << i;
variableLocations[ename.str()] = &energy;
}
// Parse the expressions.
int numSteps = stepType.size();
vector<int> forceGroup;
vector<vector<ParsedExpression> > expressions;
......@@ -107,37 +124,25 @@ void ReferenceCustomDynamics::initialize(ContextImpl& context, vector<RealOpenMM
stepExpressions[i].resize(expressions[i].size());
for (int j = 0; j < (int) expressions[i].size(); j++) {
stepExpressions[i][j] = ParsedExpression(replaceDerivFunctions(expressions[i][j].getRootNode(), context)).createCompiledExpression();
stepExpressions[i][j].setVariableLocations(variableLocations);
expressionSet.registerExpression(stepExpressions[i][j]);
}
if (stepType[i] == CustomIntegrator::WhileBlockStart)
blockEnd[blockEnd[i]] = i; // Record where to branch back to.
}
kineticEnergyExpression = Parser::parse(integrator.getKineticEnergyExpression()).optimize().createCompiledExpression();
kineticEnergyExpression.setVariableLocations(variableLocations);
expressionSet.registerExpression(kineticEnergyExpression);
kineticEnergyNeedsForce = false;
if (kineticEnergyExpression.getVariables().find("f") != kineticEnergyExpression.getVariables().end())
kineticEnergyNeedsForce = true;
// Record the variable names and flags for the force and energy in each step.
// Record the force group flags for each step.
forceGroupFlags.resize(numSteps, -1);
fIndex = expressionSet.getVariableIndex("f");
energyIndex = expressionSet.getVariableIndex("energy");
forceVariableIndex.resize(numSteps, fIndex);
energyVariableIndex.resize(numSteps, energyIndex);
vector<string> forceGroupName;
vector<string> energyGroupName;
for (int i = 0; i < 32; i++) {
stringstream fname;
fname << "f" << i;
forceGroupName.push_back(fname.str());
stringstream ename;
ename << "energy" << i;
energyGroupName.push_back(ename.str());
}
for (int i = 0; i < numSteps; i++) {
if (needsForces[i] && forceGroup[i] > -1)
forceVariableIndex[i] = expressionSet.getVariableIndex(forceGroupName[forceGroup[i]]);
if (needsEnergy[i] && forceGroup[i] > -1)
energyVariableIndex[i] = expressionSet.getVariableIndex(energyGroupName[forceGroup[i]]);
for (int i = 0; i < numSteps; i++)
if (forceGroup[i] > -1)
forceGroupFlags[i] = 1<<forceGroup[i];
}
// Build the list of inverse masses.
......@@ -150,13 +155,10 @@ void ReferenceCustomDynamics::initialize(ContextImpl& context, vector<RealOpenMM
inverseMasses[i] = 1.0/masses[i];
}
// Record indices of other variables.
// Record indices of variables.
xIndex = expressionSet.getVariableIndex("x");
vIndex = expressionSet.getVariableIndex("v");
mIndex = expressionSet.getVariableIndex("m");
gaussianIndex = expressionSet.getVariableIndex("gaussian");
uniformIndex = expressionSet.getVariableIndex("uniform");
for (int i = 0; i < integrator.getNumPerDofVariables(); i++)
perDofVariableIndex.push_back(expressionSet.getVariableIndex(integrator.getPerDofVariableName(i)));
for (int i = 0; i < stepVariable.size(); i++)
......@@ -222,15 +224,14 @@ void ReferenceCustomDynamics::update(ContextImpl& context, int numberOfAtoms, ve
}
forcesAreValid = true;
}
expressionSet.setVariable(energyVariableIndex[step], energy);
// Execute the step.
int nextStep = step+1;
switch (stepType[step]) {
case CustomIntegrator::ComputeGlobal: {
expressionSet.setVariable(uniformIndex, SimTKOpenMMUtilities::getUniformlyDistributedRandomNumber());
expressionSet.setVariable(gaussianIndex, SimTKOpenMMUtilities::getNormallyDistributedRandomNumber());
uniform = SimTKOpenMMUtilities::getUniformlyDistributedRandomNumber();
gaussian = SimTKOpenMMUtilities::getNormallyDistributedRandomNumber();
RealOpenMM result = stepExpressions[step][0].evaluate();
globals[stepVariable[step]] = result;
expressionSet.setVariable(stepVariableIndex[step], result);
......@@ -249,11 +250,11 @@ void ReferenceCustomDynamics::update(ContextImpl& context, int numberOfAtoms, ve
}
if (results == NULL)
throw OpenMMException("Illegal per-DOF output variable: "+stepVariable[step]);
computePerDof(numberOfAtoms, *results, atomCoordinates, velocities, forces, masses, perDof, stepExpressions[step][0], forceVariableIndex[step]);
computePerDof(numberOfAtoms, *results, atomCoordinates, velocities, forces, masses, perDof, stepExpressions[step][0]);
break;
}
case CustomIntegrator::ComputeSum: {
computePerDof(numberOfAtoms, sumBuffer, atomCoordinates, velocities, forces, masses, perDof, stepExpressions[step][0], forceVariableIndex[step]);
computePerDof(numberOfAtoms, sumBuffer, atomCoordinates, velocities, forces, masses, perDof, stepExpressions[step][0]);
RealOpenMM sum = 0.0;
for (int j = 0; j < numberOfAtoms; j++)
if (masses[j] != 0.0)
......@@ -306,22 +307,22 @@ void ReferenceCustomDynamics::update(ContextImpl& context, int numberOfAtoms, ve
void ReferenceCustomDynamics::computePerDof(int numberOfAtoms, vector<RealVec>& results, const vector<RealVec>& atomCoordinates,
const vector<RealVec>& velocities, const vector<RealVec>& forces, const vector<RealOpenMM>& masses,
const vector<vector<RealVec> >& perDof, const CompiledExpression& expression, int forceIndex) {
const vector<vector<RealVec> >& perDof, const CompiledExpression& expression) {
// Loop over all degrees of freedom.
for (int i = 0; i < numberOfAtoms; i++) {
if (masses[i] != 0.0) {
expressionSet.setVariable(mIndex, masses[i]);
m = masses[i];
for (int j = 0; j < 3; j++) {
// Compute the expression.
expressionSet.setVariable(xIndex, atomCoordinates[i][j]);
expressionSet.setVariable(vIndex, velocities[i][j]);
expressionSet.setVariable(forceIndex, forces[i][j]);
expressionSet.setVariable(uniformIndex, SimTKOpenMMUtilities::getUniformlyDistributedRandomNumber());
expressionSet.setVariable(gaussianIndex, SimTKOpenMMUtilities::getNormallyDistributedRandomNumber());
x = atomCoordinates[i][j];
v = velocities[i][j];
f = forces[i][j];
uniform = SimTKOpenMMUtilities::getUniformlyDistributedRandomNumber();
gaussian = SimTKOpenMMUtilities::getNormallyDistributedRandomNumber();
for (int k = 0; k < (int) perDof.size(); k++)
expressionSet.setVariable(perDofVariableIndex[k], perDof[k][i][j]);
perDofVariable[k] = perDof[k][i][j];
results[i][j] = expression.evaluate();
}
}
......@@ -329,8 +330,8 @@ void ReferenceCustomDynamics::computePerDof(int numberOfAtoms, vector<RealVec>&
}
bool ReferenceCustomDynamics::evaluateCondition(int step) {
expressionSet.setVariable(uniformIndex, SimTKOpenMMUtilities::getUniformlyDistributedRandomNumber());
expressionSet.setVariable(gaussianIndex, SimTKOpenMMUtilities::getNormallyDistributedRandomNumber());
uniform = SimTKOpenMMUtilities::getUniformlyDistributedRandomNumber();
gaussian = SimTKOpenMMUtilities::getNormallyDistributedRandomNumber();
double lhs = stepExpressions[step][0].evaluate();
double rhs = stepExpressions[step][1].evaluate();
switch (comparisons[step]) {
......@@ -390,7 +391,7 @@ double ReferenceCustomDynamics::computeKineticEnergy(OpenMM::ContextImpl& contex
energy = context.calcForcesAndEnergy(true, true, -1);
forcesAreValid = true;
}
computePerDof(numberOfAtoms, sumBuffer, atomCoordinates, velocities, forces, masses, perDof, kineticEnergyExpression, fIndex);
computePerDof(numberOfAtoms, sumBuffer, atomCoordinates, velocities, forces, masses, perDof, kineticEnergyExpression);
RealOpenMM sum = 0.0;
for (int j = 0; j < numberOfAtoms; j++)
if (masses[j] != 0.0)
......
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