Optimized CPU implementation of CustomIntegrator

platforms/reference/include/ReferenceCustomDynamics.h

 
-/* 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

 
-/* 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)