Skip to content

GitLab

Commit 27550631 authored by peastman's avatar peastman Committed by GitHub
Browse files

Merge pull request #1558 from peastman/variables 

Optimizations to custom forces/integrators on CPU platform
parents 5a8a8aa9 bf0b43c5
Show whitespace changes
Inline Side-by-side
Showing with 221 additions and 150 deletions
libraries/lepton/include/lepton/CompiledExpression.h
View file @ 27550631
......@@ -9,7 +9,7 @@
* Biological Structures at Stanford, funded under the NIH Roadmap for *
* Medical Research, grant U54 GM072970. See https://simtk.org. *
* *
* Portions copyright (c) 2013 Stanford University and the Authors. *
* Portions copyright (c) 2013-2016 Stanford University and the Authors. *
* Authors: Peter Eastman *
* Contributors: *
* *
......@@ -37,6 +37,7 @@
#include <map>
#include <set>
#include <string>
#include <utility>
#include <vector>
#ifdef LEPTON_USE_JIT
#include "asmjit.h"
......@@ -73,6 +74,12 @@ public:
* to set the value of the variable before calling evaluate().
*/
double& getVariableReference(const std::string& name);
/**
* You can optionally specify the memory locations from which the values of variables should be read.
* This is useful, for example, when several expressions all use the same variable. You can then set
* the value of that variable in one place, and it will be seen by all of them.
*/
void setVariableLocations(std::map<std::string, double*>& variableLocations);
/**
* Evaluate the expression. The values of all variables should have been set before calling this.
*/
......@@ -82,6 +89,8 @@ private:
CompiledExpression(const ParsedExpression& expression);
void compileExpression(const ExpressionTreeNode& node, std::vector<std::pair<ExpressionTreeNode, int> >& temps);
int findTempIndex(const ExpressionTreeNode& node, std::vector<std::pair<ExpressionTreeNode, int> >& temps);
std::map<std::string, double*> variablePointers;
std::vector<std::pair<double*, double*> > variablesToCopy;
std::vector<std::vector<int> > arguments;
std::vector<int> target;
std::vector<Operation*> operation;
......
libraries/lepton/src/CompiledExpression.cpp
View file @ 27550631
......@@ -6,7 +6,7 @@
* Biological Structures at Stanford, funded under the NIH Roadmap for *
* Medical Research, grant U54 GM072970. See https://simtk.org. *
* *
* Portions copyright (c) 2013 Stanford University and the Authors. *
* Portions copyright (c) 2013-2016 Stanford University and the Authors. *
* Authors: Peter Eastman *
* Contributors: *
* *
......@@ -77,10 +77,7 @@ CompiledExpression& CompiledExpression::operator=(const CompiledExpression& expr
operation.resize(expression.operation.size());
for (int i = 0; i < (int) operation.size(); i++)
operation[i] = expression.operation[i]->clone();
#ifdef LEPTON_USE_JIT
if (workspace.size() > 0)
generateJitCode();
#endif
setVariableLocations(variablePointers);
return *this;
}
......@@ -138,16 +135,41 @@ const set<string>& CompiledExpression::getVariables() const {
}
double& CompiledExpression::getVariableReference(const string& name) {
map<string, double*>::iterator pointer = variablePointers.find(name);
if (pointer != variablePointers.end())
return *pointer->second;
map<string, int>::iterator index = variableIndices.find(name);
if (index == variableIndices.end())
throw Exception("getVariableReference: Unknown variable '"+name+"'");
return workspace[index->second];
}
void CompiledExpression::setVariableLocations(map<string, double*>& variableLocations) {
variablePointers = variableLocations;
#ifdef LEPTON_USE_JIT
// Rebuild the JIT code.
if (workspace.size() > 0)
generateJitCode();
#else
// Make a list of all variables we will need to copy before evaluating the expression.
variablesToCopy.clear();
for (map<string, int>::const_iterator iter = variableIndices.begin(); iter != variableIndices.end(); ++iter) {
map<string, double*>::iterator pointer = variablePointers.find(iter->first);
if (pointer != variablePointers.end())
variablesToCopy.push_back(make_pair(&workspace[iter->second], pointer->second));
}
#endif
}
double CompiledExpression::evaluate() const {
#ifdef LEPTON_USE_JIT
return ((double (*)()) jitCode)();
#else
for (int i = 0; i < variablesToCopy.size(); i++)
*variablesToCopy[i].first = *variablesToCopy[i].second;
// Loop over the operations and evaluate each one.
for (int step = 0; step < operation.size(); step++) {
......@@ -176,16 +198,16 @@ void CompiledExpression::generateJitCode() {
vector<X86XmmVar> workspaceVar(workspace.size());
for (int i = 0; i < (int) workspaceVar.size(); i++)
workspaceVar[i] = c.newXmmVar(kX86VarTypeXmmSd);
X86GpVar workspacePointer(c);
X86GpVar argsPointer(c);
c.mov(workspacePointer, imm_ptr(&workspace[0]));
c.mov(argsPointer, imm_ptr(&argValues[0]));
// Load the arguments into variables.
for (set<string>::const_iterator iter = variableNames.begin(); iter != variableNames.end(); ++iter) {
map<string, int>::iterator index = variableIndices.find(*iter);
c.movsd(workspaceVar[index->second], x86::ptr(workspacePointer, 8*index->second, 0));
X86GpVar variablePointer(c);
c.mov(variablePointer, imm_ptr(&getVariableReference(index->first)));
c.movsd(workspaceVar[index->second], x86::ptr(variablePointer, 0, 0));
}
// Make a list of all constants that will be needed for evaluation.
......
platforms/cpu/include/CpuCustomGBForce.h
View file @ 27550631
/* Portions copyright (c) 2009-2014 Stanford University and Simbios.
/* Portions copyright (c) 2009-2016 Stanford University and Simbios.
* Contributors: Peter Eastman
*
* Permission is hereby granted, free of charge, to any person obtaining
......@@ -254,15 +254,15 @@ public:
std::vector<std::vector<Lepton::CompiledExpression> > valueDerivExpressions;
std::vector<std::vector<Lepton::CompiledExpression> > valueGradientExpressions;
std::vector<std::vector<Lepton::CompiledExpression> > valueParamDerivExpressions;
std::vector<int> valueIndex;
std::vector<double> value;
std::vector<Lepton::CompiledExpression> energyExpressions;
std::vector<std::vector<Lepton::CompiledExpression> > energyDerivExpressions;
std::vector<std::vector<Lepton::CompiledExpression> > energyGradientExpressions;
std::vector<std::vector<Lepton::CompiledExpression> > energyParamDerivExpressions;
std::vector<int> paramIndex;
std::vector<int> particleParamIndex;
std::vector<int> particleValueIndex;
int xindex, yindex, zindex, rindex;
std::vector<double> param;
std::vector<double> particleParam;
std::vector<double> particleValue;
double x, y, z, r;
int firstAtom, lastAtom;
// Workspace vectors
std::vector<float> value0, dVdR1, dVdR2, dVdX, dVdY, dVdZ;
......
platforms/cpu/include/CpuCustomNonbondedForce.h
View file @ 27550631
......@@ -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 @ 27550631
......@@ -59,47 +59,68 @@ CpuCustomGBForce::ThreadData::ThreadData(int numAtoms, int numThreads, int threa
energyGradientExpressions(energyGradientExpressions), energyParamDerivExpressions(energyParamDerivExpressions) {
firstAtom = (threadIndex*(long long) numAtoms)/numThreads;
lastAtom = ((threadIndex+1)*(long long) numAtoms)/numThreads;
for (int i = 0; i < (int) valueExpressions.size(); i++)
map<string, double*> variableLocations;
variableLocations["x"] = &x;
variableLocations["y"] = &y;
variableLocations["z"] = &z;
variableLocations["r"] = &r;
param.resize(parameterNames.size());
particleParam.resize(parameterNames.size()*2);
for (int i = 0; i < (int) parameterNames.size(); i++) {
variableLocations[parameterNames[i]] = &param[i];
for (int j = 0; j < 2; j++) {
stringstream name;
name << parameterNames[i] << (j+1);
variableLocations[name.str()] = &particleParam[2*i+j];
}
}
value.resize(valueNames.size());
particleValue.resize(valueNames.size()*2);
for (int i = 0; i < (int) valueNames.size(); i++) {
variableLocations[valueNames[i]] = &value[i];
for (int j = 0; j < 2; j++) {
stringstream name;
name << valueNames[i] << (j+1);
variableLocations[name.str()] = &particleValue[2*i+j];
}
}
for (int i = 0; i < (int) valueExpressions.size(); i++) {
this->valueExpressions[i].setVariableLocations(variableLocations);
expressionSet.registerExpression(this->valueExpressions[i]);
}
for (int i = 0; i < (int) valueDerivExpressions.size(); i++)
for (int j = 0; j < (int) valueDerivExpressions[i].size(); j++)
for (int j = 0; j < (int) valueDerivExpressions[i].size(); j++) {
this->valueDerivExpressions[i][j].setVariableLocations(variableLocations);
expressionSet.registerExpression(this->valueDerivExpressions[i][j]);
}
for (int i = 0; i < (int) valueGradientExpressions.size(); i++)
for (int j = 0; j < (int) valueGradientExpressions[i].size(); j++)
for (int j = 0; j < (int) valueGradientExpressions[i].size(); j++) {
this->valueGradientExpressions[i][j].setVariableLocations(variableLocations);
expressionSet.registerExpression(this->valueGradientExpressions[i][j]);
}
for (int i = 0; i < (int) valueParamDerivExpressions.size(); i++)
for (int j = 0; j < (int) valueParamDerivExpressions[i].size(); j++)
for (int j = 0; j < (int) valueParamDerivExpressions[i].size(); j++) {
this->valueParamDerivExpressions[i][j].setVariableLocations(variableLocations);
expressionSet.registerExpression(this->valueParamDerivExpressions[i][j]);
for (int i = 0; i < (int) energyExpressions.size(); i++)
}
for (int i = 0; i < (int) energyExpressions.size(); i++) {
this->energyExpressions[i].setVariableLocations(variableLocations);
expressionSet.registerExpression(this->energyExpressions[i]);
}
for (int i = 0; i < (int) energyDerivExpressions.size(); i++)
for (int j = 0; j < (int) energyDerivExpressions[i].size(); j++)
for (int j = 0; j < (int) energyDerivExpressions[i].size(); j++) {
this->energyDerivExpressions[i][j].setVariableLocations(variableLocations);
expressionSet.registerExpression(this->energyDerivExpressions[i][j]);
}
for (int i = 0; i < (int) energyGradientExpressions.size(); i++)
for (int j = 0; j < (int) energyGradientExpressions[i].size(); j++)
for (int j = 0; j < (int) energyGradientExpressions[i].size(); j++) {
this->energyGradientExpressions[i][j].setVariableLocations(variableLocations);
expressionSet.registerExpression(this->energyGradientExpressions[i][j]);
}
for (int i = 0; i < (int) energyParamDerivExpressions.size(); i++)
for (int j = 0; j < (int) energyParamDerivExpressions[i].size(); j++)
for (int j = 0; j < (int) energyParamDerivExpressions[i].size(); j++) {
this->energyParamDerivExpressions[i][j].setVariableLocations(variableLocations);
expressionSet.registerExpression(this->energyParamDerivExpressions[i][j]);
xindex = expressionSet.getVariableIndex("x");
yindex = expressionSet.getVariableIndex("y");
zindex = expressionSet.getVariableIndex("z");
rindex = expressionSet.getVariableIndex("r");
for (int i = 0; i < (int) parameterNames.size(); i++) {
paramIndex.push_back(expressionSet.getVariableIndex(parameterNames[i]));
for (int j = 1; j < 3; j++) {
stringstream name;
name << parameterNames[i] << j;
particleParamIndex.push_back(expressionSet.getVariableIndex(name.str()));
}
}
for (int i = 0; i < (int) valueNames.size(); i++) {
valueIndex.push_back(expressionSet.getVariableIndex(valueNames[i]));
for (int j = 1; j < 3; j++) {
stringstream name;
name << valueNames[i] << j;
particleValueIndex.push_back(expressionSet.getVariableIndex(name.str()));
}
}
value0.resize(numAtoms);
dEdV.resize(valueNames.size());
......@@ -283,13 +304,13 @@ void CpuCustomGBForce::threadComputeForce(ThreadPool& threads, int threadIndex)
for (int j = 0; j < (int) threadData.size(); j++)
sum += threadData[j]->value0[atom];
values[0][atom] = sum;
data.expressionSet.setVariable(data.xindex, posq[4*atom]);
data.expressionSet.setVariable(data.yindex, posq[4*atom+1]);
data.expressionSet.setVariable(data.zindex, posq[4*atom+2]);
data.x = posq[4*atom];
data.y = posq[4*atom+1];
data.z = posq[4*atom+2];
for (int j = 0; j < numParams; j++)
data.expressionSet.setVariable(data.paramIndex[j], atomParameters[atom][j]);
data.param[j] = atomParameters[atom][j];
for (int i = 1; i < numValues; i++) {
data.expressionSet.setVariable(data.valueIndex[i-1], values[i-1][atom]);
data.value[i-1] = values[i-1][atom];
values[i][atom] = (float) data.valueExpressions[i].evaluate();
// Calculate derivatives with respect to parameters.
......@@ -397,15 +418,14 @@ void CpuCustomGBForce::calculateOnePairValue(int index, int atom1, int atom2, Th
getDeltaR(pos2, pos1, deltaR, r2, periodic, boxSize, invBoxSize);
if (cutoff && r2 >= cutoffDistance2)
return;
float r = sqrtf(r2);
data.r = sqrtf(r2);
for (int i = 0; i < numParams; i++) {
data.expressionSet.setVariable(data.particleParamIndex[i*2], atomParameters[atom1][i]);
data.expressionSet.setVariable(data.particleParamIndex[i*2+1], atomParameters[atom2][i]);
data.particleParam[i*2] = atomParameters[atom1][i];
data.particleParam[i*2+1] = atomParameters[atom2][i];
}
data.expressionSet.setVariable(data.rindex, r);
for (int i = 0; i < index; i++) {
data.expressionSet.setVariable(data.particleValueIndex[i*2], values[i][atom1]);
data.expressionSet.setVariable(data.particleValueIndex[i*2+1], values[i][atom2]);
data.particleValue[i*2] = values[i][atom1];
data.particleValue[i*2+1] = values[i][atom2];
}
valueArray[atom1] += (float) data.valueExpressions[index].evaluate();
......@@ -418,13 +438,13 @@ void CpuCustomGBForce::calculateOnePairValue(int index, int atom1, int atom2, Th
void CpuCustomGBForce::calculateSingleParticleEnergyTerm(int index, ThreadData& data, int numAtoms, float* posq,
RealOpenMM** atomParameters, float* forces, double& totalEnergy) {
for (int i = data.firstAtom; i < data.lastAtom; i++) {
data.expressionSet.setVariable(data.xindex, posq[4*i]);
data.expressionSet.setVariable(data.yindex, posq[4*i+1]);
data.expressionSet.setVariable(data.zindex, posq[4*i+2]);
data.x = posq[4*i];
data.y = posq[4*i+1];
data.z = posq[4*i+2];
for (int j = 0; j < numParams; j++)
data.expressionSet.setVariable(data.paramIndex[j], atomParameters[i][j]);
data.param[j] = atomParameters[i][j];
for (int j = 0; j < (int) values.size(); j++)
data.expressionSet.setVariable(data.valueIndex[j], values[j][i]);
data.value[j] = values[j][i];
if (includeEnergy)
totalEnergy += (float) data.energyExpressions[index].evaluate();
for (int j = 0; j < (int) values.size(); j++)
......@@ -494,17 +514,17 @@ void CpuCustomGBForce::calculateOnePairEnergyTerm(int index, int atom1, int atom
if (cutoff && r2 >= cutoffDistance2)
return;
float r = sqrtf(r2);
data.r = r;
// Record variables for evaluating expressions.
for (int i = 0; i < numParams; i++) {
data.expressionSet.setVariable(data.particleParamIndex[i*2], atomParameters[atom1][i]);
data.expressionSet.setVariable(data.particleParamIndex[i*2+1], atomParameters[atom2][i]);
data.particleParam[i*2] = atomParameters[atom1][i];
data.particleParam[i*2+1] = atomParameters[atom2][i];
}
data.expressionSet.setVariable(data.rindex, r);
for (int i = 0; i < (int) values.size(); i++) {
data.expressionSet.setVariable(data.particleValueIndex[i*2], values[i][atom1]);
data.expressionSet.setVariable(data.particleValueIndex[i*2+1], values[i][atom2]);
data.particleValue[i*2] = values[i][atom1];
data.particleValue[i*2+1] = values[i][atom2];
}
// Evaluate the energy and its derivatives.
......@@ -571,13 +591,13 @@ void CpuCustomGBForce::calculateChainRuleForces(ThreadData& data, int numAtoms,
// Compute chain rule terms for computed values that depend explicitly on particle coordinates.
for (int i = data.firstAtom; i < data.lastAtom; i++) {
data.expressionSet.setVariable(data.xindex, posq[4*i]);
data.expressionSet.setVariable(data.yindex, posq[4*i+1]);
data.expressionSet.setVariable(data.zindex, posq[4*i+2]);
data.x = posq[4*i];
data.y = posq[4*i+1];
data.z = posq[4*i+2];
for (int j = 0; j < numParams; j++)
data.expressionSet.setVariable(data.paramIndex[j], atomParameters[i][j]);
data.param[j] = atomParameters[i][j];
for (int j = 1; j < (int) values.size(); j++) {
data.expressionSet.setVariable(data.valueIndex[j-1], values[j-1][i]);
data.value[j-1] = values[j-1][i];
data.dVdX[j] = 0.0;
data.dVdY[j] = 0.0;
data.dVdZ[j] = 0.0;
......@@ -599,7 +619,7 @@ void CpuCustomGBForce::calculateChainRuleForces(ThreadData& data, int numAtoms,
// Compute chain rule terms for derivatives with respect to parameters.
for (int i = data.firstAtom; i < data.lastAtom; i++)
for (int j = 0; j < data.valueIndex.size(); j++)
for (int j = 0; j < data.value.size(); j++)
for (int k = 0; k < dValuedParam[j].size(); k++)
data.energyParamDerivs[k] += dEdV[j][i]*dValuedParam[j][k][i];
}
......@@ -616,21 +636,21 @@ void CpuCustomGBForce::calculateOnePairChainRule(int atom1, int atom2, ThreadDat
if (cutoff && r2 >= cutoffDistance2)
return;
float r = sqrtf(r2);
data.r = r;
// Record variables for evaluating expressions.
for (int i = 0; i < numParams; i++) {
data.expressionSet.setVariable(data.particleParamIndex[i*2], atomParameters[atom1][i]);
data.expressionSet.setVariable(data.particleParamIndex[i*2+1], atomParameters[atom2][i]);
data.expressionSet.setVariable(data.paramIndex[i], atomParameters[atom1][i]);
}
data.expressionSet.setVariable(data.valueIndex[0], values[0][atom1]);
data.expressionSet.setVariable(data.xindex, posq[4*atom1]);
data.expressionSet.setVariable(data.yindex, posq[4*atom1+1]);
data.expressionSet.setVariable(data.zindex, posq[4*atom1+2]);
data.expressionSet.setVariable(data.rindex, r);
data.expressionSet.setVariable(data.particleValueIndex[0], values[0][atom1]);
data.expressionSet.setVariable(data.particleValueIndex[1], values[0][atom2]);
data.particleParam[i*2] = atomParameters[atom1][i];
data.particleParam[i*2+1] = atomParameters[atom2][i];
data.param[i] = atomParameters[atom1][i];
}
data.value[0] = values[0][atom1];
data.x = posq[4*atom1];
data.y = posq[4*atom1+1];
data.z = posq[4*atom1+2];
data.particleValue[0] = values[0][atom1];
data.particleValue[1] = values[0][atom2];
// Evaluate the derivative of each parameter with respect to position and apply forces.
......@@ -644,7 +664,7 @@ void CpuCustomGBForce::calculateOnePairChainRule(int atom1, int atom2, ThreadDat
f2 -= deltaR*(dEdV[0][atom1]*data.dVdR2[0]);
}
for (int i = 1; i < (int) values.size(); i++) {
data.expressionSet.setVariable(data.valueIndex[i], values[i][atom1]);
data.value[i] = values[i][atom1];
data.dVdR1[i] = 0.0;
data.dVdR2[i] = 0.0;
for (int j = 0; j < i; j++) {
......
platforms/cpu/src/CpuCustomNonbondedForce.cpp
View file @ 27550631
......@@ -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 @ 27550631
......@@ -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 @ 27550631
......@@ -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)
......
tests/TestParser.cpp
View file @ 27550631
#include "../libraries/lepton/include/Lepton.h"
#include "openmm/internal/AssertionUtilities.h"
#include <iostream>
#include <limits>
#include <map>
using namespace Lepton;
using namespace OpenMM;
using namespace std;
#define ASSERT_EQUAL_TOL(expected, found, tol) {double _scale_ = std::fabs(expected) > 1.0 ? std::fabs(expected) : 1.0; if (!(std::fabs((expected)-(found))/_scale_ <= (tol))) throw exception();};
/**
* This is a custom function equal to f(x,y) = 2*x*y.
*/
......@@ -102,6 +102,18 @@ void verifyEvaluation(const string& expression, double x, double y, double expec
value = compiled.evaluate();
ASSERT_EQUAL_TOL(expectedValue, value, 1e-10);
// Try specifying memory locations for the compiled expression.
map<string, double*> variablePointers;
variablePointers["x"] = &x;
variablePointers["y"] = &y;
CompiledExpression compiled2 = parsed.createCompiledExpression();
compiled2.setVariableLocations(variablePointers);
value = compiled2.evaluate();
ASSERT_EQUAL_TOL(expectedValue, value, 1e-10);
ASSERT_EQUAL(&x, &compiled2.getVariableReference("x"));
ASSERT_EQUAL(&y, &compiled2.getVariableReference("y"));
// Make sure that variable renaming works.
variables.clear();
......
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