Skip to content

GitLab

Commit 25d21454 authored by peastman's avatar peastman
Browse files

Merge pull request #1471 from peastman/compiledintegrator 

Optimized CPU implementation of CustomIntegrator
parents 7e877240 803e826b
Hide whitespace changes
Inline Side-by-side
Showing with 128 additions and 95 deletions
libraries/lepton/src/CompiledExpression.cpp
View file @ 25d21454
......@@ -78,7 +78,8 @@ CompiledExpression& CompiledExpression::operator=(const CompiledExpression& expr
for (int i = 0; i < (int) operation.size(); i++)
operation[i] = expression.operation[i]->clone();
#ifdef LEPTON_USE_JIT
generateJitCode();
if (workspace.size() > 0)
generateJitCode();
#endif
return *this;
}
......
platforms/reference/include/ReferenceCustomDynamics.h
View file @ 25d21454
/* Portions copyright (c) 2011-2015 Stanford University and Simbios.
/* Portions copyright (c) 2011-2016 Stanford University and Simbios.
* Contributors: Peter Eastman
*
* Permission is hereby granted, free of charge, to any person obtaining
......@@ -29,7 +29,8 @@
#include "openmm/CustomIntegrator.h"
#include "openmm/internal/ContextImpl.h"
#include "openmm/internal/CustomIntegratorUtilities.h"
#include "lepton/ExpressionProgram.h"
#include "openmm/internal/CompiledExpressionSet.h"
#include "lepton/CompiledExpression.h"
#include <map>
#include <string>
......@@ -44,23 +45,27 @@ private:
std::vector<RealOpenMM> inverseMasses;
std::vector<OpenMM::RealVec> sumBuffer, oldPos;
std::vector<OpenMM::CustomIntegrator::ComputationType> stepType;
std::vector<std::string> stepVariable, forceName, energyName;
std::vector<std::vector<Lepton::ExpressionProgram> > stepExpressions;
std::vector<std::string> stepVariable;
std::vector<std::vector<Lepton::CompiledExpression> > stepExpressions;
std::vector<CustomIntegratorUtilities::Comparison> comparisons;
std::vector<bool> invalidatesForces, needsForces, needsEnergy, computeBothForceAndEnergy;
std::vector<int> forceGroupFlags, blockEnd;
RealOpenMM energy;
Lepton::ExpressionProgram kineticEnergyExpression;
Lepton::CompiledExpression kineticEnergyExpression;
bool kineticEnergyNeedsForce;
CompiledExpressionSet expressionSet;
int xIndex, vIndex, mIndex, fIndex, energyIndex, gaussianIndex, uniformIndex;
std::vector<int> forceVariableIndex, energyVariableIndex, perDofVariableIndex, stepVariableIndex;
void initialize(OpenMM::ContextImpl& context, std::vector<RealOpenMM>& masses, std::map<std::string, RealOpenMM>& globals);
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::map<std::string, RealOpenMM>& globals, const std::vector<std::vector<OpenMM::RealVec> >& perDof,
const Lepton::ExpressionProgram& expression, const std::string& forceName);
const std::vector<std::vector<OpenMM::RealVec> >& perDof, const Lepton::CompiledExpression& expression, int forceIndex);
void recordChangedParameters(OpenMM::ContextImpl& context, std::map<std::string, RealOpenMM>& globals);
bool evaluateCondition(int step, std::map<std::string, RealOpenMM>& globals);
bool evaluateCondition(int step);
public:
......
platforms/reference/src/SimTKReference/ReferenceCustomDynamics.cpp
View file @ 25d21454
/* Portions copyright (c) 2011-2015 Stanford University and Simbios.
/* Portions copyright (c) 2011-2016 Stanford University and Simbios.
* Contributors: Peter Eastman
*
* Permission is hereby granted, free of charge, to any person obtaining
......@@ -56,13 +56,11 @@ ReferenceCustomDynamics::ReferenceCustomDynamics(int numberOfAtoms, const Custom
string expression;
integrator.getComputationStep(i, stepType[i], stepVariable[i], expression);
}
kineticEnergyExpression = Lepton::Parser::parse(integrator.getKineticEnergyExpression()).optimize().createProgram();
kineticEnergyExpression = Lepton::Parser::parse(integrator.getKineticEnergyExpression()).optimize().createCompiledExpression();
expressionSet.registerExpression(kineticEnergyExpression);
kineticEnergyNeedsForce = false;
for (int i = 0; i < kineticEnergyExpression.getNumOperations(); i++) {
const Lepton::Operation& op = kineticEnergyExpression.getOperation(i);
if (op.getId() == Lepton::Operation::VARIABLE && op.getName() == "f")
kineticEnergyNeedsForce = true;
}
if (kineticEnergyExpression.getVariables().find("f") != kineticEnergyExpression.getVariables().end())
kineticEnergyNeedsForce = true;
}
/**---------------------------------------------------------------------------------------
......@@ -74,6 +72,75 @@ ReferenceCustomDynamics::ReferenceCustomDynamics(int numberOfAtoms, const Custom
ReferenceCustomDynamics::~ReferenceCustomDynamics() {
}
void ReferenceCustomDynamics::initialize(ContextImpl& context, vector<RealOpenMM>& masses, map<string, RealOpenMM>& globals) {
// 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.
int numSteps = stepType.size();
vector<int> forceGroup;
vector<vector<Lepton::ParsedExpression> > expressions;
CustomIntegratorUtilities::analyzeComputations(context, integrator, expressions, comparisons, blockEnd, invalidatesForces, needsForces, needsEnergy, computeBothForceAndEnergy, forceGroup);
stepExpressions.resize(expressions.size());
for (int i = 0; i < numSteps; i++) {
stepExpressions[i].resize(expressions[i].size());
for (int j = 0; j < (int) expressions[i].size(); j++) {
stepExpressions[i][j] = expressions[i][j].createCompiledExpression();
expressionSet.registerExpression(stepExpressions[i][j]);
}
if (stepType[i] == CustomIntegrator::WhileBlockStart)
blockEnd[blockEnd[i]] = i; // Record where to branch back to.
}
// Record the variable names and flags for the force and energy in 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]]);
if (forceGroup[i] > -1)
forceGroupFlags[i] = 1<<forceGroup[i];
}
// Build the list of inverse masses.
int numberOfAtoms = masses.size();
inverseMasses.resize(numberOfAtoms);
for (int i = 0; i < numberOfAtoms; i++) {
if (masses[i] == 0.0)
inverseMasses[i] = 0.0;
else
inverseMasses[i] = 1.0/masses[i];
}
// Record indices of other 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++)
stepVariableIndex.push_back(expressionSet.getVariableIndex(stepVariable[i]));
}
/**---------------------------------------------------------------------------------------
Update -- driver routine for performing Custom dynamics update of coordinates
......@@ -93,58 +160,13 @@ ReferenceCustomDynamics::~ReferenceCustomDynamics() {
void ReferenceCustomDynamics::update(ContextImpl& context, int numberOfAtoms, vector<RealVec>& atomCoordinates,
vector<RealVec>& velocities, vector<RealVec>& forces, vector<RealOpenMM>& masses,
map<string, RealOpenMM>& globals, vector<vector<RealVec> >& perDof, bool& forcesAreValid, RealOpenMM tolerance) {
if (invalidatesForces.size() == 0)
initialize(context, masses, globals);
int numSteps = stepType.size();
globals.insert(context.getParameters().begin(), context.getParameters().end());
for (map<string, RealOpenMM>::const_iterator iter = globals.begin(); iter != globals.end(); ++iter)
expressionSet.setVariable(expressionSet.getVariableIndex(iter->first), iter->second);
oldPos = atomCoordinates;
if (invalidatesForces.size() == 0) {
// 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 this method is called.
vector<int> forceGroup;
vector<vector<Lepton::ParsedExpression> > expressions;
CustomIntegratorUtilities::analyzeComputations(context, integrator, expressions, comparisons, blockEnd, invalidatesForces, needsForces, needsEnergy, computeBothForceAndEnergy, forceGroup);
stepExpressions.resize(expressions.size());
for (int i = 0; i < numSteps; i++) {
for (int j = 0; j < (int) expressions[i].size(); j++)
stepExpressions[i].push_back(expressions[i][j].createProgram());
if (stepType[i] == CustomIntegrator::WhileBlockStart)
blockEnd[blockEnd[i]] = i; // Record where to branch back to.
}
// Record the variable names and flags for the force and energy in each step.
forceGroupFlags.resize(numSteps, -1);
forceName.resize(numSteps, "f");
energyName.resize(numSteps, "energy");
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)
forceName[i] = forceGroupName[forceGroup[i]];
if (needsEnergy[i] && forceGroup[i] > -1)
energyName[i] = energyGroupName[forceGroup[i]];
if (forceGroup[i] > -1)
forceGroupFlags[i] = 1<<forceGroup[i];
}
// Build the list of inverse masses.
inverseMasses.resize(numberOfAtoms);
for (int i = 0; i < numberOfAtoms; i++) {
if (masses[i] == 0.0)
inverseMasses[i] = 0.0;
else
inverseMasses[i] = 1.0/masses[i];
}
}
// Loop over steps and execute them.
......@@ -160,42 +182,44 @@ void ReferenceCustomDynamics::update(ContextImpl& context, int numberOfAtoms, ve
energy = e;
forcesAreValid = true;
}
globals[energyName[step]] = energy;
expressionSet.setVariable(energyVariableIndex[step], energy);
// Execute the step.
int nextStep = step+1;
switch (stepType[step]) {
case CustomIntegrator::ComputeGlobal: {
map<string, RealOpenMM> variables = globals;
variables["uniform"] = SimTKOpenMMUtilities::getUniformlyDistributedRandomNumber();
variables["gaussian"] = SimTKOpenMMUtilities::getNormallyDistributedRandomNumber();
globals[stepVariable[step]] = stepExpressions[step][0].evaluate(variables);
expressionSet.setVariable(uniformIndex, SimTKOpenMMUtilities::getUniformlyDistributedRandomNumber());
expressionSet.setVariable(gaussianIndex, SimTKOpenMMUtilities::getNormallyDistributedRandomNumber());
RealOpenMM result = stepExpressions[step][0].evaluate();
globals[stepVariable[step]] = result;
expressionSet.setVariable(stepVariableIndex[step], result);
break;
}
case CustomIntegrator::ComputePerDof: {
vector<RealVec>* results = NULL;
if (stepVariable[step] == "x")
if (stepVariableIndex[step] == xIndex)
results = &atomCoordinates;
else if (stepVariable[step] == "v")
else if (stepVariableIndex[step] == vIndex)
results = &velocities;
else {
for (int j = 0; j < integrator.getNumPerDofVariables(); j++)
if (stepVariable[step] == integrator.getPerDofVariableName(j))
if (stepVariableIndex[step] == perDofVariableIndex[j])
results = &perDof[j];
}
if (results == NULL)
throw OpenMMException("Illegal per-DOF output variable: "+stepVariable[step]);
computePerDof(numberOfAtoms, *results, atomCoordinates, velocities, forces, masses, globals, perDof, stepExpressions[step][0], forceName[step]);
computePerDof(numberOfAtoms, *results, atomCoordinates, velocities, forces, masses, perDof, stepExpressions[step][0], forceVariableIndex[step]);
break;
}
case CustomIntegrator::ComputeSum: {
computePerDof(numberOfAtoms, sumBuffer, atomCoordinates, velocities, forces, masses, globals, perDof, stepExpressions[step][0], forceName[step]);
computePerDof(numberOfAtoms, sumBuffer, atomCoordinates, velocities, forces, masses, perDof, stepExpressions[step][0], forceVariableIndex[step]);
RealOpenMM sum = 0.0;
for (int j = 0; j < numberOfAtoms; j++)
if (masses[j] != 0.0)
sum += sumBuffer[j][0]+sumBuffer[j][1]+sumBuffer[j][2];
globals[stepVariable[step]] = sum;
expressionSet.setVariable(stepVariableIndex[step], sum);
break;
}
case CustomIntegrator::ConstrainPositions: {
......@@ -211,15 +235,17 @@ void ReferenceCustomDynamics::update(ContextImpl& context, int numberOfAtoms, ve
recordChangedParameters(context, globals);
context.updateContextState();
globals.insert(context.getParameters().begin(), context.getParameters().end());
for (map<string, RealOpenMM>::const_iterator iter = globals.begin(); iter != globals.end(); ++iter)
expressionSet.setVariable(expressionSet.getVariableIndex(iter->first), iter->second);
break;
}
case CustomIntegrator::IfBlockStart: {
if (!evaluateCondition(step, globals))
if (!evaluateCondition(step))
nextStep = blockEnd[step]+1;
break;
}
case CustomIntegrator::WhileBlockStart: {
if (!evaluateCondition(step, globals))
if (!evaluateCondition(step))
nextStep = blockEnd[step]+1;
break;
}
......@@ -240,36 +266,33 @@ 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 map<string, RealOpenMM>& globals, const vector<vector<RealVec> >& perDof,
const Lepton::ExpressionProgram& expression, const std::string& forceName) {
const vector<vector<RealVec> >& perDof, const Lepton::CompiledExpression& expression, int forceIndex) {
// Loop over all degrees of freedom.
map<string, RealOpenMM> variables = globals;
for (int i = 0; i < numberOfAtoms; i++) {
if (masses[i] != 0.0) {
variables["m"] = masses[i];
expressionSet.setVariable(mIndex, masses[i]);
for (int j = 0; j < 3; j++) {
// Compute the expression.
variables["x"] = atomCoordinates[i][j];
variables["v"] = velocities[i][j];
variables[forceName] = forces[i][j];
variables["uniform"] = SimTKOpenMMUtilities::getUniformlyDistributedRandomNumber();
variables["gaussian"] = SimTKOpenMMUtilities::getNormallyDistributedRandomNumber();
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());
for (int k = 0; k < (int) perDof.size(); k++)
variables[integrator.getPerDofVariableName(k)] = perDof[k][i][j];
results[i][j] = expression.evaluate(variables);
expressionSet.setVariable(perDofVariableIndex[k], perDof[k][i][j]);
results[i][j] = expression.evaluate();
}
}
}
}
bool ReferenceCustomDynamics::evaluateCondition(int step, map<string, RealOpenMM>& globals) {
map<string, RealOpenMM> variables = globals;
variables["uniform"] = SimTKOpenMMUtilities::getUniformlyDistributedRandomNumber();
variables["gaussian"] = SimTKOpenMMUtilities::getNormallyDistributedRandomNumber();
double lhs = stepExpressions[step][0].evaluate(variables);
double rhs = stepExpressions[step][1].evaluate(variables);
bool ReferenceCustomDynamics::evaluateCondition(int step) {
expressionSet.setVariable(uniformIndex, SimTKOpenMMUtilities::getUniformlyDistributedRandomNumber());
expressionSet.setVariable(gaussianIndex, SimTKOpenMMUtilities::getNormallyDistributedRandomNumber());
double lhs = stepExpressions[step][0].evaluate();
double rhs = stepExpressions[step][1].evaluate();
switch (comparisons[step]) {
case CustomIntegratorUtilities::EQUAL:
return (lhs == rhs);
......@@ -318,12 +341,16 @@ void ReferenceCustomDynamics::recordChangedParameters(OpenMM::ContextImpl& conte
double ReferenceCustomDynamics::computeKineticEnergy(OpenMM::ContextImpl& context, int numberOfAtoms, std::vector<OpenMM::RealVec>& atomCoordinates,
std::vector<OpenMM::RealVec>& velocities, std::vector<OpenMM::RealVec>& forces, std::vector<RealOpenMM>& masses,
std::map<std::string, RealOpenMM>& globals, std::vector<std::vector<OpenMM::RealVec> >& perDof, bool& forcesAreValid) {
if (invalidatesForces.size() == 0)
initialize(context, masses, globals);
globals.insert(context.getParameters().begin(), context.getParameters().end());
for (map<string, RealOpenMM>::const_iterator iter = globals.begin(); iter != globals.end(); ++iter)
expressionSet.setVariable(expressionSet.getVariableIndex(iter->first), iter->second);
if (kineticEnergyNeedsForce) {
energy = context.calcForcesAndEnergy(true, true, -1);
forcesAreValid = true;
}
computePerDof(numberOfAtoms, sumBuffer, atomCoordinates, velocities, forces, masses, globals, perDof, kineticEnergyExpression, "f");
computePerDof(numberOfAtoms, sumBuffer, atomCoordinates, velocities, forces, masses, perDof, kineticEnergyExpression, fIndex);
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