Began creating OpenCL implementation of CustomCVForce

openmmapi/include/openmm/internal/ContextImpl.h

@@ -266,6 +266,7 @@ public:
     static std::vector<std::vector<int> > findMolecules(int numParticles, std::vector<std::vector<int> >& particleBonds);
 private:
     friend class Context;
+    void initialize();
     Context& owner;
     const System& system;
     Integrator& integrator;

openmmapi/src/Context.cpp

@@ -42,14 +42,17 @@ using namespace std;
 
 Context::Context(const System& system, Integrator& integrator) : properties(map<string, string>()) {
     impl = new ContextImpl(*this, system, integrator, 0, properties);
+    impl->initialize();
 }
 
 Context::Context(const System& system, Integrator& integrator, Platform& platform) : properties(map<string, string>()) {
     impl = new ContextImpl(*this, system, integrator, &platform, properties);
+    impl->initialize();
 }
 
 Context::Context(const System& system, Integrator& integrator, Platform& platform, const map<string, string>& properties) : properties(properties) {
     impl = new ContextImpl(*this, system, integrator, &platform, properties);
+    impl->initialize();
 }
 
 Context::~Context() {
@@ -240,6 +243,7 @@ void Context::reinitialize() {
     integrator.cleanup();
     delete impl;
     impl = new ContextImpl(*this, system, integrator, &platform, properties);
+    impl->initialize();
 }
 
 void Context::createCheckpoint(ostream& stream) {

openmmapi/src/ContextImpl.cpp

@@ -161,7 +161,9 @@ ContextImpl::ContextImpl(Context& owner, const System& system, Integrator& integ
             throw;
         }
     }
+}
 
+void ContextImpl::initialize() {
     // Create and initialize kernels and other objects.
     
     initializeForcesKernel = platform->createKernel(CalcForcesAndEnergyKernel::Name(), *this);
@@ -181,7 +183,7 @@ ContextImpl::ContextImpl(Context& owner, const System& system, Integrator& integ
         parameters.insert(forceParameters.begin(), forceParameters.end());
     }
     integrator.initialize(*this);
-    updateStateDataKernel.getAs<UpdateStateDataKernel>().setVelocities(*this, vector<Vec3>(numParticles));
+    updateStateDataKernel.getAs<UpdateStateDataKernel>().setVelocities(*this, vector<Vec3>(system.getNumParticles()));
 }
 
 ContextImpl::~ContextImpl() {

platforms/opencl/include/OpenCLKernels.h

@@ -37,6 +37,7 @@
 #include "openmm/internal/CompiledExpressionSet.h"
 #include "openmm/internal/CustomIntegratorUtilities.h"
 #include "lepton/CompiledExpression.h"
+#include "lepton/ExpressionProgram.h"
 #include "openmm/System.h"
 
 namespace OpenMM {
@@ -1207,6 +1208,53 @@ private:
     cl::Kernel framesKernel, blockBoundsKernel, neighborsKernel, forceKernel, torqueKernel;
 };
 
+/**
+ * This kernel is invoked by CustomCVForce to calculate the forces acting on the system and the energy of the system.
+ */
+class OpenCLCalcCustomCVForceKernel : public CalcCustomCVForceKernel {
+public:
+    OpenCLCalcCustomCVForceKernel(std::string name, const Platform& platform, OpenCLContext& cl) : CalcCustomCVForceKernel(name, platform),
+            cl(cl), hasInitializedKernels(false), invAtomOrder(NULL), innerInvAtomOrder(NULL) {
+    }
+    ~OpenCLCalcCustomCVForceKernel();
+    /**
+     * Initialize the kernel.
+     *
+     * @param system     the System this kernel will be applied to
+     * @param force      the CustomCVForce this kernel will be used for
+     */
+    void initialize(const System& system, const CustomCVForce& force);
+    /**
+     * Execute the kernel to calculate the forces and/or energy.
+     *
+     * @param context        the context in which to execute this kernel
+     * @param innerContext   the context created by the CustomCVForce for computing collective variables
+     * @param includeForces  true if forces should be calculated
+     * @param includeEnergy  true if the energy should be calculated
+     * @return the potential energy due to the force
+     */
+    double execute(ContextImpl& context, ContextImpl& innerContext, bool includeForces, bool includeEnergy);
+    /**
+     * Copy state information to the inner context.
+     *
+     * @param context        the context in which to execute this kernel
+     * @param innerContext   the context created by the CustomCVForce for computing collective variables
+     */
+    void copyState(ContextImpl& context, ContextImpl& innerContext);
+private:
+    class ReorderListener;
+    OpenCLContext& cl;
+    bool hasInitializedKernels;
+    Lepton::ExpressionProgram energyExpression;
+    std::vector<std::string> variableNames, paramDerivNames, globalParameterNames;
+    std::vector<Lepton::ExpressionProgram> variableDerivExpressions;
+    std::vector<Lepton::ExpressionProgram> paramDerivExpressions;
+    std::vector<OpenCLArray*> cvForces;
+    OpenCLArray* invAtomOrder;
+    OpenCLArray* innerInvAtomOrder;
+    cl::Kernel copyStateKernel, copyForcesKernel, addForcesKernel;
+};
+
 /**
  * This kernel is invoked by VerletIntegrator to take one time step.
  */

platforms/opencl/src/OpenCLKernelFactory.cpp

@@ -106,6 +106,8 @@ KernelImpl* OpenCLKernelFactory::createKernelImpl(std::string name, const Platfo
         return new OpenCLCalcCustomCentroidBondForceKernel(name, platform, cl, context.getSystem());
     if (name == CalcCustomCompoundBondForceKernel::Name())
         return new OpenCLCalcCustomCompoundBondForceKernel(name, platform, cl, context.getSystem());
+    if (name == CalcCustomCVForceKernel::Name())
+        return new OpenCLCalcCustomCVForceKernel(name, platform, cl);
     if (name == CalcCustomManyParticleForceKernel::Name())
         return new OpenCLCalcCustomManyParticleForceKernel(name, platform, cl, context.getSystem());
     if (name == CalcGayBerneForceKernel::Name())

platforms/opencl/src/OpenCLKernels.cpp

@@ -6874,6 +6874,185 @@ void OpenCLCalcGayBerneForceKernel::sortAtoms() {
     exclusionStartIndex->upload(startIndexVec);
 }
 
+class OpenCLCalcCustomCVForceKernel::ReorderListener : public OpenCLContext::ReorderListener {
+public:
+    ReorderListener(OpenCLContext& cl, OpenCLArray& invAtomOrder) : cl(cl), invAtomOrder(invAtomOrder) {
+    }
+    void execute() {
+        vector<cl_int> invOrder(cl.getNumAtoms());
+        const vector<int>& order = cl.getAtomIndex();
+        for (int i = 0; i < order.size(); i++)
+            invOrder[order[i]] = i;
+        invAtomOrder.upload(invOrder);
+        cl.getQueue().finish();
+    }
+private:
+    OpenCLContext& cl;
+    OpenCLArray& invAtomOrder;
+};
+
+OpenCLCalcCustomCVForceKernel::~OpenCLCalcCustomCVForceKernel() {
+    for (auto force : cvForces)
+        delete force;
+    if (invAtomOrder != NULL)
+        delete invAtomOrder;
+    if (innerInvAtomOrder != NULL)
+        delete innerInvAtomOrder;
+}
+
+void OpenCLCalcCustomCVForceKernel::initialize(const System& system, const CustomCVForce& force) {
+    int numCVs = force.getNumCollectiveVariables();
+    cl.addForce(new OpenCLForceInfo(1));
+    for (int i = 0; i < force.getNumGlobalParameters(); i++)
+        globalParameterNames.push_back(force.getGlobalParameterName(i));
+    
+    // Create custom functions for the tabulated functions.
+
+    map<string, Lepton::CustomFunction*> functions;
+    for (int i = 0; i < (int) force.getNumTabulatedFunctions(); i++)
+        functions[force.getTabulatedFunctionName(i)] = createReferenceTabulatedFunction(force.getTabulatedFunction(i));
+
+    // Create the expressions.
+
+    Lepton::ParsedExpression energyExpr = Lepton::Parser::parse(force.getEnergyFunction(), functions);
+    energyExpression = energyExpr.createProgram();
+    for (int i = 0; i < numCVs; i++) {
+        string name = force.getCollectiveVariableName(i);
+        variableNames.push_back(name);
+        variableDerivExpressions.push_back(energyExpr.differentiate(name).optimize().createProgram());
+    }
+    for (int i = 0; i < force.getNumEnergyParameterDerivatives(); i++) {
+        string name = force.getEnergyParameterDerivativeName(i);
+        paramDerivNames.push_back(name);
+        paramDerivExpressions.push_back(energyExpr.differentiate(name).optimize().createProgram());
+    }
+
+    // Delete the custom functions.
+
+    for (auto& function : functions)
+        delete function.second;
+    
+    // Create arrays for storing information.
+    
+    int elementSize = (cl.getUseDoublePrecision() || cl.getUseMixedPrecision() ? sizeof(double) : sizeof(float));
+    for (int i = 0; i < numCVs; i++)
+        cvForces.push_back(new OpenCLArray(cl, cl.getNumAtoms(), elementSize, "cvForce"));
+    invAtomOrder = OpenCLArray::create<cl_int>(cl, cl.getNumAtoms(), "invAtomOrder");
+    innerInvAtomOrder = OpenCLArray::create<cl_int>(cl, cl.getNumAtoms(), "innerInvAtomOrder");
+    
+    // Create the kernels.
+    
+    stringstream args, add;
+    for (int i = 0; i < numCVs; i++) {
+        args << ", __global real4* restrict force" << i << ", real dEdV" << i;
+        add << "f += force" << i << "[i]*dEdV" << i << ";\n";
+    }
+    map<string, string> replacements;
+    replacements["PARAMETER_ARGUMENTS"] = args.str();
+    replacements["ADD_FORCES"] = add.str();
+    cl::Program program = cl.createProgram(cl.replaceStrings(OpenCLKernelSources::customCVForce, replacements));
+    copyStateKernel = cl::Kernel(program, "copyState");
+    copyForcesKernel = cl::Kernel(program, "copyForces");
+    addForcesKernel = cl::Kernel(program, "addForces");
+}
+
+double OpenCLCalcCustomCVForceKernel::execute(ContextImpl& context, ContextImpl& innerContext, bool includeForces, bool includeEnergy) {
+    copyState(context, innerContext);
+    int numCVs = variableNames.size();
+    int numAtoms = cl.getNumAtoms();
+    OpenCLContext& cl2 = *reinterpret_cast<OpenCLPlatform::PlatformData*>(innerContext.getPlatformData())->contexts[0];
+    vector<double> cvValues;
+    vector<map<string, double> > cvDerivs(numCVs);
+    for (int i = 0; i < numCVs; i++) {
+        cvValues.push_back(innerContext.calcForcesAndEnergy(true, true, 1<<i));
+        copyForcesKernel.setArg<cl::Buffer>(0, cvForces[i]->getDeviceBuffer());
+        cl.executeKernel(copyForcesKernel, numAtoms);
+        innerContext.getEnergyParameterDerivatives(cvDerivs[i]);
+    }
+    
+    // Compute the energy and forces.
+    
+    map<string, double> variables;
+    for (auto& name : globalParameterNames)
+        variables[name] = context.getParameter(name);
+    for (int i = 0; i < numCVs; i++)
+        variables[variableNames[i]] = cvValues[i];
+    double energy = energyExpression.evaluate(variables);
+    for (int i = 0; i < numCVs; i++) {
+        double dEdV = variableDerivExpressions[i].evaluate(variables);
+        if (cl.getUseDoublePrecision())
+            addForcesKernel.setArg<cl_double>(2*i+3, dEdV);
+        else
+            addForcesKernel.setArg<cl_float>(2*i+3, dEdV);
+        cl.executeKernel(addForcesKernel, numAtoms);
+    }
+    
+    // Compute the energy parameter derivatives.
+    
+    map<string, double>& energyParamDerivs = cl.getEnergyParamDerivWorkspace();
+    for (int i = 0; i < paramDerivExpressions.size(); i++)
+        energyParamDerivs[paramDerivNames[i]] += paramDerivExpressions[i].evaluate(variables);
+    for (int i = 0; i < numCVs; i++) {
+        double dEdV = variableDerivExpressions[i].evaluate(variables);
+        for (auto& deriv : cvDerivs[i])
+            energyParamDerivs[deriv.first] += dEdV*deriv.second;
+    }
+    return energy;
+}
+
+void OpenCLCalcCustomCVForceKernel::copyState(ContextImpl& context, ContextImpl& innerContext) {
+    int numAtoms = cl.getNumAtoms();
+    if (!hasInitializedKernels) {
+        hasInitializedKernels = true;
+        
+        // Initialize the listeners.
+        
+        OpenCLContext& cl2 = *reinterpret_cast<OpenCLPlatform::PlatformData*>(innerContext.getPlatformData())->contexts[0];
+        ReorderListener* listener1 = new ReorderListener(cl, *invAtomOrder);
+        ReorderListener* listener2 = new ReorderListener(cl2, *innerInvAtomOrder);
+        cl.addReorderListener(listener1);
+        cl2.addReorderListener(listener2);
+        listener1->execute();
+        listener2->execute();
+        
+        // Initialize the kernels.
+        
+        copyStateKernel.setArg<cl::Buffer>(0, cl.getPosq().getDeviceBuffer());
+        copyStateKernel.setArg<cl::Buffer>(2, cl.getVelm().getDeviceBuffer());
+        copyStateKernel.setArg<cl::Buffer>(3, cl.getAtomIndexArray().getDeviceBuffer());
+        copyStateKernel.setArg<cl::Buffer>(4, cl2.getPosq().getDeviceBuffer());
+        copyStateKernel.setArg<cl::Buffer>(6, cl2.getVelm().getDeviceBuffer());
+        copyStateKernel.setArg<cl::Buffer>(7, innerInvAtomOrder->getDeviceBuffer());
+        copyStateKernel.setArg<cl_int>(8, numAtoms);
+        if (cl.getUseMixedPrecision()) {
+            copyStateKernel.setArg<cl::Buffer>(1, cl.getPosqCorrection().getDeviceBuffer());
+            copyStateKernel.setArg<cl::Buffer>(5, cl2.getPosqCorrection().getDeviceBuffer());
+        }
+        else {
+            copyStateKernel.setArg<void*>(1, NULL);
+            copyStateKernel.setArg<void*>(5, NULL);
+        }
+
+        copyForcesKernel.setArg<cl::Buffer>(1, invAtomOrder->getDeviceBuffer());
+        copyForcesKernel.setArg<cl::Buffer>(2, cl2.getForce().getDeviceBuffer());
+        copyForcesKernel.setArg<cl::Buffer>(3, cl2.getAtomIndexArray().getDeviceBuffer());
+        copyForcesKernel.setArg<cl_int>(4, numAtoms);
+
+        addForcesKernel.setArg<cl::Buffer>(0, cl.getForce().getDeviceBuffer());
+        addForcesKernel.setArg<cl_int>(1, numAtoms);
+        for (int i = 0; i < cvForces.size(); i++)
+            addForcesKernel.setArg<cl::Buffer>(2*i+2, cvForces[i]->getDeviceBuffer());
+    }
+    cl.executeKernel(copyStateKernel, numAtoms);
+    Vec3 a, b, c;
+    context.getPeriodicBoxVectors(a, b, c);
+    innerContext.setPeriodicBoxVectors(a, b, c);
+    innerContext.setTime(context.getTime());
+    map<string, double> innerParameters = innerContext.getParameters();
+    for (auto& param : innerParameters)
+        innerContext.setParameter(param.first, context.getParameter(param.first));
+}
+
 OpenCLIntegrateVerletStepKernel::~OpenCLIntegrateVerletStepKernel() {
 }
 

platforms/opencl/src/OpenCLPlatform.cpp

@@ -82,6 +82,7 @@ OpenCLPlatform::OpenCLPlatform() {
     registerKernelFactory(CalcCustomHbondForceKernel::Name(), factory);
     registerKernelFactory(CalcCustomCentroidBondForceKernel::Name(), factory);
     registerKernelFactory(CalcCustomCompoundBondForceKernel::Name(), factory);
+    registerKernelFactory(CalcCustomCVForceKernel::Name(), factory);
     registerKernelFactory(CalcCustomManyParticleForceKernel::Name(), factory);
     registerKernelFactory(CalcGayBerneForceKernel::Name(), factory);
     registerKernelFactory(IntegrateVerletStepKernel::Name(), factory);

platforms/opencl/src/kernels/customCVForce.cl

+/**
+ * Copy the positions and velocities to the inner context.
+ */
+__kernel void copyState(__global real4* posq, __global real4* posqCorrection, __global mixed4* velm, __global int* restrict atomOrder,
+        __global real4* innerPosq, __global real4* innerPosqCorrection, __global mixed4* innerVelm, __global int* restrict innerInvAtomOrder,
+        int numAtoms) {
+    for (int i = get_global_id(0); i < numAtoms; i += get_global_size(0)) {
+        int index = innerInvAtomOrder[atomOrder[i]];
+        innerPosq[index] = posq[i];
+        innerVelm[index] = velm[i];
+#ifdef USE_MIXED_PRECISION
+        innerPosqCorrection[index] = posqCorrection[i];
+#endif
+    }
+}
+
+/**
+ * Copy the forces back to the main context.
+ */
+__kernel void copyForces(__global real4* forces, __global int* restrict invAtomOrder, __global real4* innerForces,
+        __global int* restrict innerAtomOrder, int numAtoms) {
+    for (int i = get_global_id(0); i < numAtoms; i += get_global_size(0)) {
+        int index = invAtomOrder[innerAtomOrder[i]];
+        forces[index] = innerForces[i];
+    }
+}
+
+/**
+ * Add all the forces from the CVs.
+ */
+__kernel void addForces(__global real4* forces, int numAtoms, int numCVs
+    PARAMETER_ARGUMENTS) {
+    for (int i = get_global_id(0); i < numAtoms; i += get_global_size(0)) {
+        real4 f = forces[i];
+        ADD_FORCES
+        forces[i] = f;
+    }
+}
\ No newline at end of file

platforms/opencl/tests/TestOpenCLCustomCVForce.cpp

+/* -------------------------------------------------------------------------- *
+ *                                   OpenMM                                   *
+ * -------------------------------------------------------------------------- *
+ * This is part of the OpenMM molecular simulation toolkit originating from   *
+ * Simbios, the NIH National Center for Physics-Based Simulation of           *
+ * Biological Structures at Stanford, funded under the NIH Roadmap for        *
+ * Medical Research, grant U54 GM072970. See https://simtk.org.               *
+ *                                                                            *
+ * Portions copyright (c) 2017 Stanford University and the Authors.           *
+ * Authors: Peter Eastman                                                     *
+ * Contributors:                                                              *
+ *                                                                            *
+ * Permission is hereby granted, free of charge, to any person obtaining a    *
+ * copy of this software and associated documentation files (the "Software"), *
+ * to deal in the Software without restriction, including without limitation  *
+ * the rights to use, copy, modify, merge, publish, distribute, sublicense,   *
+ * and/or sell copies of the Software, and to permit persons to whom the      *
+ * Software is furnished to do so, subject to the following conditions:       *
+ *                                                                            *
+ * The above copyright notice and this permission notice shall be included in *
+ * all copies or substantial portions of the Software.                        *
+ *                                                                            *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR *
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,   *
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL    *
+ * THE AUTHORS, CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,    *
+ * DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR      *
+ * OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE  *
+ * USE OR OTHER DEALINGS IN THE SOFTWARE.                                     *
+ * -------------------------------------------------------------------------- */
+
+#include "OpenCLTests.h"
+#include "TestCustomCVForce.h"
+
+void runPlatformTests() {
+}

platforms/reference/src/ReferenceKernels.cpp

@@ -2045,10 +2045,10 @@ double ReferenceCalcCustomCVForceKernel::execute(ContextImpl& context, ContextIm
 void ReferenceCalcCustomCVForceKernel::copyState(ContextImpl& context, ContextImpl& innerContext) {
     extractPositions(innerContext) = extractPositions(context);
     extractVelocities(innerContext) = extractVelocities(context);
-    extractBoxSize(innerContext) = extractBoxSize(context);
-    for (int i = 0; i < 3; i++)
-        extractBoxVectors(innerContext)[i] = extractBoxVectors(context)[i];
-    reinterpret_cast<ReferencePlatform::PlatformData*>(innerContext.getPlatformData())->time = reinterpret_cast<ReferencePlatform::PlatformData*>(context.getPlatformData())->time;
+    Vec3 a, b, c;
+    context.getPeriodicBoxVectors(a, b, c);
+    innerContext.setPeriodicBoxVectors(a, b, c);
+    innerContext.setTime(context.getTime());
     map<string, double> innerParameters = innerContext.getParameters();
     for (auto& param : innerParameters)
         innerContext.setParameter(param.first, context.getParameter(param.first));