Merge pull request #1558 from peastman/variables

Optimizations to custom forces/integrators on CPU platform

libraries/lepton/include/lepton/CompiledExpression.h

@@ -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

@@ -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

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

@@ -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

@@ -59,48 +59,69 @@ 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());
     for (int i = 0; i < (int) dEdV.size(); i++)
@@ -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

@@ -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

@@ -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

@@ -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

 #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.
  */
@@ -101,6 +101,18 @@ void verifyEvaluation(const string& expression, double x, double y, double expec
         compiled.getVariableReference("y") = y;
     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.