Continuing to implement OpenCL CustomIntegrator: added position constraints,...

Continuing to implement OpenCL CustomIntegrator: added position constraints, and use double precision where possible

platforms/opencl/src/OpenCLExpressionUtilities.cpp

@@ -47,38 +47,38 @@ string OpenCLExpressionUtilities::intToString(int value) {
 }
 
 string OpenCLExpressionUtilities::createExpressions(const map<string, ParsedExpression>& expressions, const map<string, string>& variables,
-        const vector<pair<string, string> >& functions, const string& prefix, const string& functionParams) {
+        const vector<pair<string, string> >& functions, const string& prefix, const string& functionParams, const string& tempType) {
     vector<pair<ExpressionTreeNode, string> > variableNodes;
     for (map<string, string>::const_iterator iter = variables.begin(); iter != variables.end(); ++iter)
         variableNodes.push_back(make_pair(ExpressionTreeNode(new Operation::Variable(iter->first)), iter->second));
-    return createExpressions(expressions, variableNodes, functions, prefix, functionParams);
+    return createExpressions(expressions, variableNodes, functions, prefix, functionParams, tempType);
 }
 
 string OpenCLExpressionUtilities::createExpressions(const map<string, ParsedExpression>& expressions, const vector<pair<ExpressionTreeNode, string> >& variables,
-        const vector<pair<string, string> >& functions, const string& prefix, const string& functionParams) {
+        const vector<pair<string, string> >& functions, const string& prefix, const string& functionParams, const string& tempType) {
     stringstream out;
     vector<ParsedExpression> allExpressions;
     for (map<string, ParsedExpression>::const_iterator iter = expressions.begin(); iter != expressions.end(); ++iter)
         allExpressions.push_back(iter->second);
     vector<pair<ExpressionTreeNode, string> > temps = variables;
     for (map<string, ParsedExpression>::const_iterator iter = expressions.begin(); iter != expressions.end(); ++iter) {
-        processExpression(out, iter->second.getRootNode(), temps, functions, prefix, functionParams, allExpressions);
+        processExpression(out, iter->second.getRootNode(), temps, functions, prefix, functionParams, allExpressions, tempType);
         out << iter->first << getTempName(iter->second.getRootNode(), temps) << ";\n";
     }
     return out.str();
 }
 
 void OpenCLExpressionUtilities::processExpression(stringstream& out, const ExpressionTreeNode& node, vector<pair<ExpressionTreeNode, string> >& temps,
-        const vector<pair<string, string> >& functions, const string& prefix, const string& functionParams, const vector<ParsedExpression>& allExpressions) {
+        const vector<pair<string, string> >& functions, const string& prefix, const string& functionParams, const vector<ParsedExpression>& allExpressions, const string& tempType) {
     for (int i = 0; i < (int) temps.size(); i++)
         if (temps[i].first == node)
             return;
     for (int i = 0; i < (int) node.getChildren().size(); i++)
-        processExpression(out, node.getChildren()[i], temps, functions, prefix, functionParams, allExpressions);
+        processExpression(out, node.getChildren()[i], temps, functions, prefix, functionParams, allExpressions, tempType);
     string name = prefix+intToString(temps.size());
     bool hasRecordedNode = false;
     
-    out << "float " << name << " = ";
+    out << tempType << " " << name << " = ";
     switch (node.getOperation().getId()) {
         case Operation::CONSTANT:
             out << doubleToString(dynamic_cast<const Operation::Constant*>(&node.getOperation())->getValue());
@@ -108,7 +108,7 @@ void OpenCLExpressionUtilities::processExpression(stringstream& out, const Expre
             string derivName = name;
             if (valueNode != NULL && derivNode != NULL) {
                 string name2 = prefix+intToString(temps.size());
-                out << "float " << name2 << " = 0.0f;\n";
+                out << tempType << " " << name2 << " = 0.0f;\n";
                 if (isDeriv) {
                     valueName = name2;
                     temps.push_back(make_pair(*valueNode, name2));
@@ -266,7 +266,7 @@ void OpenCLExpressionUtilities::processExpression(stringstream& out, const Expre
                         string name2 = prefix+intToString(temps.size());
                         names.push_back(name2);
                         temps.push_back(make_pair(*iter->second, name2));
-                        out << "float " << name2 << " = 0.0f;\n";
+                        out << tempType << " " << name2 << " = 0.0f;\n";
                     }
                 }
                 out << "{\n";

platforms/opencl/src/OpenCLExpressionUtilities.h

@@ -54,9 +54,10 @@ public:
      * @param functions      defines the variable name for each tabulated function that may appear in the expressions
      * @param prefix         a prefix to put in front of temporary variables
      * @param functionParams the variable name containing the parameters for each tabulated function
+     * @param tempType       the type of value to use for temporary variables (defaults to "float")
      */
     static std::string createExpressions(const std::map<std::string, Lepton::ParsedExpression>& expressions, const std::map<std::string, std::string>& variables,
-            const std::vector<std::pair<std::string, std::string> >& functions, const std::string& prefix, const std::string& functionParams);
+            const std::vector<std::pair<std::string, std::string> >& functions, const std::string& prefix, const std::string& functionParams, const std::string& tempType="float");
     /**
      * Generate the source code for calculating a set of expressions.
      *
@@ -66,9 +67,10 @@ public:
      * @param functions      defines the variable name for each tabulated function that may appear in the expressions
      * @param prefix         a prefix to put in front of temporary variables
      * @param functionParams the variable name containing the parameters for each tabulated function
+     * @param tempType       the type of value to use for temporary variables (defaults to "float")
      */
     static std::string createExpressions(const std::map<std::string, Lepton::ParsedExpression>& expressions, const std::vector<std::pair<Lepton::ExpressionTreeNode, std::string> >& variables,
-            const std::vector<std::pair<std::string, std::string> >& functions, const std::string& prefix, const std::string& functionParams);
+            const std::vector<std::pair<std::string, std::string> >& functions, const std::string& prefix, const std::string& functionParams, const std::string& tempType="float");
     /**
      * Calculate the spline coefficients for a tabulated function that appears in expressions.
      *
@@ -91,7 +93,7 @@ private:
     static void processExpression(std::stringstream& out, const Lepton::ExpressionTreeNode& node,
             std::vector<std::pair<Lepton::ExpressionTreeNode, std::string> >& temps,
             const std::vector<std::pair<std::string, std::string> >& functions, const std::string& prefix, const std::string& functionParams,
-            const std::vector<Lepton::ParsedExpression>& allExpressions);
+            const std::vector<Lepton::ParsedExpression>& allExpressions, const std::string& tempType);
     static std::string getTempName(const Lepton::ExpressionTreeNode& node, const std::vector<std::pair<Lepton::ExpressionTreeNode, std::string> >& temps);
     static void findRelatedTabulatedFunctions(const Lepton::ExpressionTreeNode& node, const Lepton::ExpressionTreeNode& searchNode,
             const Lepton::ExpressionTreeNode*& valueNode, const Lepton::ExpressionTreeNode*& derivNode);

platforms/opencl/src/OpenCLKernels.cpp

@@ -3584,7 +3584,8 @@ string OpenCLIntegrateCustomStepKernel::createPerDofComputation(const string& va
     for (int i = 0; i < (int) parameterNames.size(); i++)
         variables[parameterNames[i]] = "params["+intToString(i)+"]";
     vector<pair<string, string> > functions;
-    return OpenCLExpressionUtilities::createExpressions(expressions, variables, functions, "temp"+intToString(component)+"_", "");
+    string tempType = (cl.getSupportsDoublePrecision() ? "double" : "float");
+    return OpenCLExpressionUtilities::createExpressions(expressions, variables, functions, "temp"+intToString(component)+"_", "", tempType);
 }
 
 void OpenCLIntegrateCustomStepKernel::execute(ContextImpl& context, CustomIntegrator& integrator, bool& forcesAreValid) {
@@ -3624,15 +3625,43 @@ void OpenCLIntegrateCustomStepKernel::execute(ContextImpl& context, CustomIntegr
         defines["NUM_ATOMS"] = intToString(cl.getNumAtoms());
         defines["WORK_GROUP_SIZE"] = intToString(OpenCLContext::ThreadBlockSize);
         
+        // Record information about all the computation steps.
+        
+        stepType.resize(numSteps);
+        vector<string> variable(numSteps);
+        vector<Lepton::ParsedExpression> expression(numSteps);
+        for (int step = 0; step < numSteps; step++) {
+            string expr;
+            integrator.getComputationStep(step, stepType[step], variable[step], expr);
+            if (expr.size() > 0)
+                expression[step] = Lepton::Parser::parse(expr).optimize();
+        }
+        
+        // Determine how each step will represent the position (as just a value, or a value plus a delta).
+        
+        vector<bool> storePosAsDelta(numSteps, false);
+        vector<bool> loadPosAsDelta(numSteps, false);
+        bool beforeConstrain = false;
+        for (int step = numSteps-1; step >= 0; step--) {
+            if (stepType[step] == CustomIntegrator::ConstrainPositions)
+                beforeConstrain = true;
+            else if (stepType[step] == CustomIntegrator::ComputePerDof && variable[step] == "x" && beforeConstrain)
+                storePosAsDelta[step] = true;
+        }
+        bool storedAsDelta = false;
+        for (int step = 0; step < numSteps; step++) {
+            loadPosAsDelta[step] = storedAsDelta;
+            if (storePosAsDelta[step] == true)
+                storedAsDelta = true;
+            if (stepType[step] == CustomIntegrator::ConstrainPositions)
+                storedAsDelta = false;
+        }
+        
         // Loop over all steps and create the kernels for them.
         
         for (int step = 0; step < numSteps; step++) {
-            CustomIntegrator::ComputationType type;
-            string variable, expression;
-            integrator.getComputationStep(step, type, variable, expression);
-            stepType.push_back(type);
-            invalidatesForces[step] = (type == CustomIntegrator::ConstrainPositions || affectsForce.find(variable) != affectsForce.end());
-            if (type == CustomIntegrator::ComputePerDof || type == CustomIntegrator::ComputeSum) {
+            invalidatesForces[step] = (stepType[step] == CustomIntegrator::ConstrainPositions || affectsForce.find(variable[step]) != affectsForce.end());
+            if (stepType[step] == CustomIntegrator::ComputePerDof || stepType[step] == CustomIntegrator::ComputeSum) {
                 // Compute a per-DOF value.
                 
                 stringstream compute;
@@ -3642,15 +3671,19 @@ void OpenCLIntegrateCustomStepKernel::execute(ContextImpl& context, CustomIntegr
                     compute << buffer.getType()<<" perDofy"<<intToString(i+1)<<" = perDofValues"<<intToString(i+1)<<"[3*index+1];\n";
                     compute << buffer.getType()<<" perDofz"<<intToString(i+1)<<" = perDofValues"<<intToString(i+1)<<"[3*index+2];\n";
                 }
-                Lepton::ParsedExpression expr = Lepton::Parser::parse(expression).optimize();
-                needsForces[step] = usesVariable(expr, "f");
-                needsEnergy[step] = usesVariable(expr, "energy");
+                needsForces[step] = usesVariable(expression[step], "f");
+                needsEnergy[step] = usesVariable(expression[step], "energy");
                 for (int i = 0; i < 3; i++)
-                    compute << createPerDofComputation(type == CustomIntegrator::ComputePerDof ? variable : "", expr, i, integrator);
-                if (variable == "x")
-                    compute << "posq[index] = position;\n";
-                else if (variable == "v")
-                    compute << "velm[index] = velocity;\n";
+                    compute << createPerDofComputation(stepType[step] == CustomIntegrator::ComputePerDof ? variable[step] : "", expression[step], i, integrator);
+                string convert = (cl.getSupportsDoublePrecision() ? "convert_float4(" : "(");
+                if (variable[step] == "x") {
+                    if (storePosAsDelta[step])
+                        compute << "posDelta[index] = " << convert << "position-posq[index]);\n";
+                    else
+                        compute << "posq[index] = " << convert << "position);\n";
+                }
+                else if (variable[step] == "v")
+                    compute << "velm[index] = " << convert << "velocity);\n";
                 for (int i = 0; i < (int) perDofValues->getBuffers().size(); i++) {
                     const OpenCLNonbondedUtilities::ParameterInfo& buffer = perDofValues->getBuffers()[i];
                     compute << "perDofValues"<<intToString(i+1)<<"[3*index] = perDofx"<<intToString(i+1)<<";\n";
@@ -3666,12 +3699,16 @@ void OpenCLIntegrateCustomStepKernel::execute(ContextImpl& context, CustomIntegr
                     args << ", __global " << buffer.getType() << "* restrict " << valueName;
                 }
                 replacements["PARAMETER_ARGUMENTS"] = args.str();
+                if (loadPosAsDelta[step])
+                    defines["LOAD_POS_AS_DELTA"] = "1";
+                else if (defines.find("LOAD_POS_AS_DELTA") != defines.end())
+                    defines.erase("LOAD_POS_AS_DELTA");
                 cl::Program program = cl.createProgram(cl.replaceStrings(OpenCLKernelSources::customIntegratorPerDof, replacements), defines);
                 cl::Kernel kernel = cl::Kernel(program, "computePerDof");
                 kernels[step].push_back(kernel);
-                if (usesVariable(expr, "uniform"))
+                if (usesVariable(expression[step], "uniform"))
                     throw OpenMMException("OpenCL platform does not currently support per-DOF uniform random numbers");
-                requiredRandoms[step].push_back(numAtoms*usesVariable(expr, "gaussian"));
+                requiredRandoms[step].push_back(numAtoms*usesVariable(expression[step], "gaussian"));
                 int index = 0;
                 kernel.setArg<cl::Buffer>(index++, cl.getPosq().getDeviceBuffer());
                 kernel.setArg<cl::Buffer>(index++, integration.getPosDelta().getDeviceBuffer());
@@ -3685,39 +3722,38 @@ void OpenCLIntegrateCustomStepKernel::execute(ContextImpl& context, CustomIntegr
                 index += 2;
                 for (int i = 0; i < (int) perDofValues->getBuffers().size(); i++)
                     kernel.setArg<cl::Memory>(index++, perDofValues->getBuffers()[i].getMemory());
-                if (type == CustomIntegrator::ComputeSum) {
+                if (stepType[step] == CustomIntegrator::ComputeSum) {
                     // Create a second kernel for this step that sums the values.
 
-                    program = cl.createProgram(OpenCLKernelSources::customIntegratorSum, defines);
+                    program = cl.createProgram(OpenCLKernelSources::customIntegrator, defines);
                     kernel = cl::Kernel(program, "computeSum");
                     kernels[step].push_back(kernel);
                     index = 0;
                     kernel.setArg<cl::Buffer>(index++, sumBuffer->getDeviceBuffer());
                     bool found = false;
                     for (int j = 0; j < integrator.getNumGlobalVariables() && !found; j++)
-                        if (variable == integrator.getGlobalVariableName(j)) {
+                        if (variable[step] == integrator.getGlobalVariableName(j)) {
                             kernel.setArg<cl::Buffer>(index++, globalValues->getDeviceBuffer());
                             kernel.setArg<cl_uint>(index++, j);
                             found = true;
                         }
                     for (int j = 0; j < (int) parameterNames.size() && !found; j++)
-                        if (variable == parameterNames[j]) {
+                        if (variable[step] == parameterNames[j]) {
                             kernel.setArg<cl::Buffer>(index++, contextParameterValues->getDeviceBuffer());
                             kernel.setArg<cl_uint>(index++, j);
                             found = true;
                             modifiesParameters = true;
                         }
                     if (!found)
-                        throw OpenMMException("Unknown global variable: "+variable);
+                        throw OpenMMException("Unknown global variable: "+variable[step]);
                 }
             }
-            else if (type == CustomIntegrator::ComputeGlobal) {
+            else if (stepType[step] == CustomIntegrator::ComputeGlobal) {
                 // Compute a global value.
 
                 stringstream compute;
-                Lepton::ParsedExpression expr = Lepton::Parser::parse(expression).optimize();
-                needsEnergy[step] = usesVariable(expr, "energy");
-                compute << createGlobalComputation(variable, expr, integrator);
+                needsEnergy[step] = usesVariable(expression[step], "energy");
+                compute << createGlobalComputation(variable[step], expression[step], integrator);
                 map<string, string> replacements;
                 replacements["COMPUTE_STEP"] = compute.str();
                 stringstream args;
@@ -3729,6 +3765,16 @@ void OpenCLIntegrateCustomStepKernel::execute(ContextImpl& context, CustomIntegr
                 kernel.setArg<cl::Buffer>(index++, globalValues->getDeviceBuffer());
                 kernel.setArg<cl::Buffer>(index++, contextParameterValues->getDeviceBuffer());
             }
+            else if (stepType[step] == CustomIntegrator::ConstrainPositions) {
+                // Apply position constraints.
+
+                cl::Program program = cl.createProgram(OpenCLKernelSources::customIntegrator, defines);
+                cl::Kernel kernel = cl::Kernel(program, "applyPositionDeltas");
+                kernels[step].push_back(kernel);
+                int index = 0;
+                kernel.setArg<cl::Buffer>(index++, cl.getPosq().getDeviceBuffer());
+                kernel.setArg<cl::Buffer>(index++, integration.getPosDelta().getDeviceBuffer());
+            }
         }
     }
     
@@ -3803,6 +3849,10 @@ void OpenCLIntegrateCustomStepKernel::execute(ContextImpl& context, CustomIntegr
             recordChangedParameters(context);
             context.updateContextState();
         }
+        else if (stepType[i] == CustomIntegrator::ConstrainPositions) {
+            cl.getIntegrationUtilities().applyConstraints(integrator.getConstraintTolerance());
+            cl.executeKernel(kernels[i][0], numAtoms);
+        }
         if (invalidatesForces[i])
             forcesAreValid = false;
     }

platforms/opencl/src/kernels/customIntegratorSum.cl → platforms/opencl/src/kernels/customIntegrator.cl

@@ -13,3 +13,11 @@ __kernel void computeSum(__global const float* restrict sumBuffer, __global floa
     if (thread == 0)
         result[outputIndex] = tempBuffer[0];
 }
+
+__kernel void applyPositionDeltas(__global float4* restrict posq, __global float4* restrict posDelta) {
+    for (unsigned int index = get_local_id(0); index < NUM_ATOMS; index += get_global_size(0)) {
+        float4 position = posq[index];
+        position.xyz += posDelta[index].xyz;
+        posq[index] = position;
+    }
+}

platforms/opencl/src/kernels/customIntegratorPerDof.cl

+#ifdef SUPPORTS_DOUBLE_PRECISION
+#pragma OPENCL EXTENSION cl_khr_fp64 : enable
+#endif
+
 __kernel void computePerDof(__global float4* restrict posq, __global float4* restrict posDelta, __global float4* restrict velm,
         __global const float4* restrict force, __global const float2* restrict dt, __global const float* restrict globals,
         __global const float* restrict params, __global float* restrict sum, __global const float4* restrict random,
@@ -7,11 +11,26 @@ __kernel void computePerDof(__global float4* restrict posq, __global float4* res
     int index = get_global_id(0);
     randomIndex += index;
     while (index < NUM_ATOMS) {
+#ifdef SUPPORTS_DOUBLE_PRECISION
+#ifdef LOAD_POS_AS_DELTA
+        double4 position = convert_double4(posq[index]+posDelta[index]);
+#else
+        double4 position = convert_double4(posq[index]);
+#endif
+        double4 velocity = convert_double4(velm[index]);
+        double4 f = convert_double4(force[index]);
+        double mass = 1.0/velocity.w;
+#else
+#ifdef LOAD_POS_AS_DELTA
+        float4 position = posq[index]+posDelta[index];
+#else
         float4 position = posq[index];
+#endif
         float4 velocity = velm[index];
         float4 f = force[index];
-        float4 gaussian = random[randomIndex];
         float mass = 1.0f/velocity.w;
+#endif
+        float4 gaussian = random[randomIndex];
         COMPUTE_STEP
         randomIndex += get_global_size(0);
         index += get_global_size(0);

platforms/opencl/tests/TestOpenCLCustomIntegrator.cpp

@@ -382,7 +382,7 @@ void testParameter() {
 int main() {
     try {
        testSingleBond();
-//        testConstraints();
+        testConstraints();
 //        testVelocityConstraints();
         testWithThermostat();
         testMonteCarlo();