Implemented per-DOF uniform random numbers for OpenCL. Also fixed some bugs...

Implemented per-DOF uniform random numbers for OpenCL.  Also fixed some bugs in getting and setting per-DOF variables.

openmmapi/src/CustomIntegrator.cpp

@@ -55,8 +55,11 @@ void CustomIntegrator::initialize(ContextImpl& contextRef) {
     kernel = context->getPlatform().createKernel(IntegrateCustomStepKernel::Name(), contextRef);
     dynamic_cast<IntegrateCustomStepKernel&>(kernel.getImpl()).initialize(contextRef.getSystem(), *this);
     dynamic_cast<IntegrateCustomStepKernel&>(kernel.getImpl()).setGlobalVariables(contextRef, globalValues);
-    for (int i = 0; i < (int) perDofValues.size(); i++)
+    for (int i = 0; i < (int) perDofValues.size(); i++) {
+        if (perDofValues[i].size() == 1)
+            perDofValues[i].resize(context->getSystem().getNumParticles(), perDofValues[i][0]);
         dynamic_cast<IntegrateCustomStepKernel&>(kernel.getImpl()).setPerDofVariable(contextRef, i, perDofValues[i]);
+    }
 }
 
 void CustomIntegrator::stateChanged(State::DataType changed) {
@@ -146,6 +149,8 @@ void CustomIntegrator::getPerDofVariable(int index, vector<Vec3>& values) const
 void CustomIntegrator::setPerDofVariable(int index, const vector<Vec3>& values) {
     if (index < 0 || index >= perDofNames.size())
         throw OpenMMException("Index out of range");
+    if (values.size() != context->getSystem().getNumParticles())
+        throw OpenMMException("setPerDofVariable() called with wrong number of values");
     if (owner == NULL)
         perDofValues[index] = values;
     else

platforms/opencl/src/OpenCLKernels.cpp

@@ -3509,6 +3509,10 @@ OpenCLIntegrateCustomStepKernel::~OpenCLIntegrateCustomStepKernel() {
         delete sumBuffer;
     if (energy != NULL)
         delete energy;
+    if (uniformRandoms != NULL)
+        delete uniformRandoms;
+    if (randomSeed != NULL)
+        delete randomSeed;
     if (perDofValues != NULL)
         delete perDofValues;
 }
@@ -3576,6 +3580,7 @@ string OpenCLIntegrateCustomStepKernel::createPerDofComputation(const string& va
     variables["v"] = "velocity"+suffix;
     variables["f"] = "f"+suffix;
     variables["gaussian"] = "gaussian"+suffix;
+    variables["uniform"] = "uniform"+suffix;
     variables["m"] = "mass";
     variables["dt"] = "stepSize";
     variables["energy"] = "energy[0]";
@@ -3607,12 +3612,34 @@ void OpenCLIntegrateCustomStepKernel::execute(ContextImpl& context, CustomIntegr
         }
         contextParameterValues->upload();
         kernels.resize(integrator.getNumComputations());
-        requiredRandoms.resize(integrator.getNumComputations());
+        requiredGaussian.resize(integrator.getNumComputations(), 0);
+        requiredUniform.resize(integrator.getNumComputations(), 0);
         needsForces.resize(numSteps, false);
         needsEnergy.resize(numSteps, false);
         invalidatesForces.resize(numSteps, false);
         merged.resize(numSteps, false);
         modifiesParameters = false;
+        map<string, string> defines;
+        defines["NUM_ATOMS"] = intToString(cl.getNumAtoms());
+        defines["WORK_GROUP_SIZE"] = intToString(OpenCLContext::ThreadBlockSize);
+        
+        // Initialize the random number generator.
+        
+        uniformRandoms = new OpenCLArray<mm_float4>(cl, cl.getNumAtoms(), "uniformRandoms");
+        randomSeed = new OpenCLArray<mm_int4>(cl, cl.getNumThreadBlocks()*OpenCLContext::ThreadBlockSize, "randomSeed");
+        vector<mm_int4> seed(randomSeed->getSize());
+        unsigned int r = integrator.getRandomNumberSeed()+1;
+        for (int i = 0; i < randomSeed->getSize(); i++) {
+            seed[i].x = r = (1664525*r + 1013904223) & 0xFFFFFFFF;
+            seed[i].y = r = (1664525*r + 1013904223) & 0xFFFFFFFF;
+            seed[i].z = r = (1664525*r + 1013904223) & 0xFFFFFFFF;
+            seed[i].w = r = (1664525*r + 1013904223) & 0xFFFFFFFF;
+        }
+        randomSeed->upload(seed);
+        cl::Program randomProgram = cl.createProgram(OpenCLKernelSources::customIntegrator, defines);
+        randomKernel = cl::Kernel(randomProgram, "generateRandomNumbers");
+        randomKernel.setArg<cl::Buffer>(0, uniformRandoms->getDeviceBuffer());
+        randomKernel.setArg<cl::Buffer>(1, randomSeed->getDeviceBuffer());
         
         // Build a list of all variables that affect the forces, so we can tell which
         // steps invalidate them.
@@ -3624,9 +3651,6 @@ void OpenCLIntegrateCustomStepKernel::execute(ContextImpl& context, CustomIntegr
             for (map<string, double>::const_iterator param = params.begin(); param != params.end(); ++param)
                 affectsForce.insert(param->first);
         }
-        map<string, string> defines;
-        defines["NUM_ATOMS"] = intToString(cl.getNumAtoms());
-        defines["WORK_GROUP_SIZE"] = intToString(OpenCLContext::ThreadBlockSize);
         
         // Record information about all the computation steps.
         
@@ -3672,7 +3696,8 @@ void OpenCLIntegrateCustomStepKernel::execute(ContextImpl& context, CustomIntegr
             if (stepType[step-1] == CustomIntegrator::ComputeGlobal && stepType[step] == CustomIntegrator::ComputeGlobal)
                 merged[step] = true;
             if (stepType[step-1] == CustomIntegrator::ComputePerDof && stepType[step] == CustomIntegrator::ComputePerDof &&
-                    storePosAsDelta[step-1] == storePosAsDelta[step] && loadPosAsDelta[step-1] == loadPosAsDelta[step])
+                    storePosAsDelta[step-1] == storePosAsDelta[step] && loadPosAsDelta[step-1] == loadPosAsDelta[step] &&
+                    !usesVariable(expression[step], "uniform"))
                 merged[step] = true;
         }
         
@@ -3690,7 +3715,7 @@ void OpenCLIntegrateCustomStepKernel::execute(ContextImpl& context, CustomIntegr
                     compute << buffer.getType()<<" perDofz"<<intToString(i+1)<<" = perDofValues"<<intToString(i+1)<<"[3*index+2];\n";
                 }
                 string convert = (cl.getSupportsDoublePrecision() ? "convert_float4(" : "(");
-                int randoms = 0;
+                int numGaussian = 0, numUniform = 0;
                 for (int j = step; j < numSteps && (j == step || merged[j]); j++) {
                     compute << "{\n";
                     for (int i = 0; i < 3; i++)
@@ -3714,7 +3739,8 @@ void OpenCLIntegrateCustomStepKernel::execute(ContextImpl& context, CustomIntegr
                         compute << "perDofValues"<<intToString(i+1)<<"[3*index+2] = perDofz"<<intToString(i+1)<<";\n";
                     }
                     compute << "}\n";
-                    randoms += numAtoms*usesVariable(expression[j], "gaussian");
+                    numGaussian += numAtoms*usesVariable(expression[j], "gaussian");
+                    numUniform += numAtoms*usesVariable(expression[j], "uniform");
                 }
                 map<string, string> replacements;
                 replacements["COMPUTE_STEP"] = compute.str();
@@ -3732,9 +3758,8 @@ void OpenCLIntegrateCustomStepKernel::execute(ContextImpl& context, CustomIntegr
                 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(expression[step], "uniform"))
-                    throw OpenMMException("OpenCL platform does not currently support per-DOF uniform random numbers");
-                requiredRandoms[step].push_back(randoms);
+                requiredGaussian[step] = numGaussian;
+                requiredUniform[step] = numUniform;
                 int index = 0;
                 kernel.setArg<cl::Buffer>(index++, cl.getPosq().getDeviceBuffer());
                 kernel.setArg<cl::Buffer>(index++, integration.getPosDelta().getDeviceBuffer());
@@ -3746,6 +3771,7 @@ void OpenCLIntegrateCustomStepKernel::execute(ContextImpl& context, CustomIntegr
                 kernel.setArg<cl::Buffer>(index++, sumBuffer->getDeviceBuffer());
                 kernel.setArg<cl::Buffer>(index++, integration.getRandom().getDeviceBuffer());
                 index++;
+                kernel.setArg<cl::Buffer>(index++, uniformRandoms->getDeviceBuffer());
                 kernel.setArg<cl::Buffer>(index++, energy->getDeviceBuffer());
                 for (int i = 0; i < (int) perDofValues->getBuffers().size(); i++)
                     kernel.setArg<cl::Memory>(index++, perDofValues->getBuffers()[i].getMemory());
@@ -3867,7 +3893,9 @@ void OpenCLIntegrateCustomStepKernel::execute(ContextImpl& context, CustomIntegr
             forcesAreValid = true;
         }
         if (stepType[i] == CustomIntegrator::ComputePerDof && !merged[i]) {
-            kernels[i][0].setArg<cl_uint>(9, integration.prepareRandomNumbers(requiredRandoms[i][0]));
+            kernels[i][0].setArg<cl_uint>(9, integration.prepareRandomNumbers(requiredGaussian[i]));
+            if (requiredUniform[i] > 0)
+                cl.executeKernel(randomKernel, numAtoms);
             cl.executeKernel(kernels[i][0], numAtoms);
         }
         else if (stepType[i] == CustomIntegrator::ComputeGlobal && !merged[i]) {
@@ -3876,7 +3904,9 @@ void OpenCLIntegrateCustomStepKernel::execute(ContextImpl& context, CustomIntegr
             cl.executeKernel(kernels[i][0], 1);
         }
         else if (stepType[i] == CustomIntegrator::ComputeSum) {
-            kernels[i][0].setArg<cl_uint>(9, integration.prepareRandomNumbers(requiredRandoms[i][0]));
+            kernels[i][0].setArg<cl_uint>(9, integration.prepareRandomNumbers(requiredGaussian[i]));
+            if (requiredUniform[i] > 0)
+                cl.executeKernel(randomKernel, numAtoms);
             cl.executeKernel(kernels[i][0], numAtoms);
             cl.executeKernel(kernels[i][1], OpenCLContext::ThreadBlockSize, OpenCLContext::ThreadBlockSize);
         }
@@ -3934,7 +3964,7 @@ void OpenCLIntegrateCustomStepKernel::getPerDofVariable(ContextImpl& context, in
     values.resize(perDofValues->getNumObjects()/3);
     for (int i = 0; i < (int) values.size(); i++)
         for (int j = 0; j < 3; j++)
-            values[i][j] = localPerDofValues[variable][3*i+j];
+            values[i][j] = localPerDofValues[3*i+j][variable];
 }
 
 void OpenCLIntegrateCustomStepKernel::setPerDofVariable(ContextImpl& context, int variable, const vector<Vec3>& values) {
@@ -3944,7 +3974,7 @@ void OpenCLIntegrateCustomStepKernel::setPerDofVariable(ContextImpl& context, in
     }
     for (int i = 0; i < (int) values.size(); i++)
         for (int j = 0; j < 3; j++)
-            localPerDofValues[variable][3*i+j] = (float) values[i][j];
+            localPerDofValues[3*i+j][variable] = (float) values[i][j];
     deviceValuesAreCurrent = false;
 }
 

platforms/opencl/src/OpenCLKernels.h

@@ -883,7 +883,8 @@ private:
 class OpenCLIntegrateCustomStepKernel : public IntegrateCustomStepKernel {
 public:
     OpenCLIntegrateCustomStepKernel(std::string name, const Platform& platform, OpenCLContext& cl) : IntegrateCustomStepKernel(name, platform), cl(cl),
-            hasInitializedKernels(false), localValuesAreCurrent(false), globalValues(NULL), contextParameterValues(NULL), sumBuffer(NULL), energy(NULL), perDofValues(NULL) {
+            hasInitializedKernels(false), localValuesAreCurrent(false), globalValues(NULL), contextParameterValues(NULL), sumBuffer(NULL), energy(NULL),
+            uniformRandoms(NULL), randomSeed(NULL), perDofValues(NULL) {
     }
     ~OpenCLIntegrateCustomStepKernel();
     /**
@@ -947,16 +948,19 @@ private:
     OpenCLArray<cl_float>* contextParameterValues;
     OpenCLArray<cl_float>* sumBuffer;
     OpenCLArray<cl_float>* energy;
+    OpenCLArray<mm_float4>* uniformRandoms;
+    OpenCLArray<mm_int4>* randomSeed;
     OpenCLParameterSet* perDofValues;
     mutable std::vector<std::vector<cl_float> > localPerDofValues;
     std::vector<std::vector<cl::Kernel> > kernels;
-    cl::Kernel sumEnergyKernel;
+    cl::Kernel sumEnergyKernel, randomKernel;
     std::vector<CustomIntegrator::ComputationType> stepType;
     std::vector<bool> needsForces;
     std::vector<bool> needsEnergy;
     std::vector<bool> invalidatesForces;
     std::vector<bool> merged;
-    std::vector<std::vector<int> > requiredRandoms;
+    std::vector<int> requiredGaussian;
+    std::vector<int> requiredUniform;
     std::vector<std::string> parameterNames;
 };
 

platforms/opencl/src/kernels/customIntegrator.cl

@@ -22,3 +22,47 @@ __kernel void applyPositionDeltas(__global float4* restrict posq, __global float
         posDelta[index] = (float4) 0.0f;
     }
 }
+
+__kernel void generateRandomNumbers(__global float4* restrict random, __global uint4* restrict seed) {
+    uint4 state = seed[get_global_id(0)];
+    unsigned int carry = 0;
+    for (int index = get_global_id(0); index < NUM_ATOMS; index += get_global_size(0)) {
+        // Generate three uniform random numbers.
+
+        state.x = state.x * 69069 + 1;
+        state.y ^= state.y << 13;
+        state.y ^= state.y >> 17;
+        state.y ^= state.y << 5;
+        unsigned int k = (state.z >> 2) + (state.w >> 3) + (carry >> 2);
+        unsigned int m = state.w + state.w + state.z + carry;
+        state.z = state.w;
+        state.w = m;
+        carry = k >> 30;
+        float x1 = (float)max(state.x + state.y + state.w, 0x00000001u) / (float)0xffffffff;
+        state.x = state.x * 69069 + 1;
+        state.y ^= state.y << 13;
+        state.y ^= state.y >> 17;
+        state.y ^= state.y << 5;
+        k = (state.z >> 2) + (state.w >> 3) + (carry >> 2);
+        m = state.w + state.w + state.z + carry;
+        state.z = state.w;
+        state.w = m;
+        carry = k >> 30;
+        float x2 = (float)max(state.x + state.y + state.w, 0x00000001u) / (float)0xffffffff;
+        state.x = state.x * 69069 + 1;
+        state.y ^= state.y << 13;
+        state.y ^= state.y >> 17;
+        state.y ^= state.y << 5;
+        k = (state.z >> 2) + (state.w >> 3) + (carry >> 2);
+        m = state.w + state.w + state.z + carry;
+        state.z = state.w;
+        state.w = m;
+        carry = k >> 30;
+        float x3 = (float)max(state.x + state.y + state.w, 0x00000001u) / (float)0xffffffff;
+
+        // Record the values.
+
+        random[index] = (float4) (x1, x2, x3, 0.0f);
+    }
+    seed[get_global_id(0)] = state;
+}

platforms/opencl/src/kernels/customIntegratorPerDof.cl

@@ -4,8 +4,8 @@
 
 __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,
-        unsigned int randomIndex, __global const float* restrict energy
+        __global const float* restrict params, __global float* restrict sum, __global const float4* restrict gaussianValues,
+        unsigned int randomIndex, __global const float4* restrict uniformValues, __global const float* restrict energy
         PARAMETER_ARGUMENTS) {
     float stepSize = dt[0].y;
     int index = get_global_id(0);
@@ -30,7 +30,8 @@ __kernel void computePerDof(__global float4* restrict posq, __global float4* res
         float4 f = force[index];
         float mass = 1.0f/velocity.w;
 #endif
-        float4 gaussian = random[randomIndex];
+        float4 gaussian = gaussianValues[randomIndex];
+        float4 uniform = uniformValues[index];
         COMPUTE_STEP
         randomIndex += get_global_size(0);
         index += get_global_size(0);

platforms/opencl/tests/TestOpenCLCustomIntegrator.cpp

@@ -398,6 +398,75 @@ void testParameter() {
     }
 }
 
+/**
+ * Test random number distributions.
+ */
+void testRandomDistributions() {
+    const int numParticles = 100;
+    const int numBins = 20;
+    const int numSteps = 100;
+    OpenCLPlatform platform;
+    System system;
+    for (int i = 0; i < numParticles; i++)
+        system.addParticle(1.0);
+    CustomIntegrator integrator(0.1);
+    integrator.addPerDofVariable("a", 0);
+    integrator.addPerDofVariable("b", 0);
+    integrator.addComputePerDof("a", "uniform");
+    integrator.addComputePerDof("b", "gaussian");
+    Context context(system, integrator, platform);
+    
+    // See if the random numbers are distributed correctly.
+    
+    vector<int> bins(numBins);
+    double mean = 0.0;
+    double var = 0.0;
+    double skew = 0.0;
+    double kurtosis = 0.0;
+    vector<Vec3> values;
+    for (int i = 0; i < numSteps; i++) {
+        integrator.step(1);
+        integrator.getPerDofVariable(0, values);
+        for (int i = 0; i < numParticles; i++)
+            for (int j = 0; j < 3; j++) {
+                double v = values[i][j];
+                ASSERT(v >= 0 && v < 1);
+                bins[(int) (v*numBins)]++;
+            }
+        integrator.getPerDofVariable(1, values);
+        for (int i = 0; i < numParticles; i++)
+            for (int j = 0; j < 3; j++) {
+                double v = values[i][j];
+                mean += v;
+                var += v*v;
+                skew += v*v*v;
+                kurtosis += v*v*v*v;
+            }
+    }
+    
+    // Check the distribution of uniform randoms.
+    
+    int numValues = numParticles*numSteps*3;
+    double expected = numValues/(double) numBins;
+    double tol = 4*sqrt(expected);
+    for (int i = 0; i < numBins; i++)
+        ASSERT(bins[i] >= expected-tol && bins[i] <= expected+tol);
+    
+    // Check the distribution of gaussian randoms.
+    
+    mean /= numValues;
+    var /= numValues;
+    skew /= numValues;
+    kurtosis /= numValues;
+    double c2 = var-mean*mean;
+    double c3 = skew-3*var*mean+2*mean*mean*mean;
+    double c4 = kurtosis-4*skew*mean-3*var*var+12*var*mean*mean-6*mean*mean*mean*mean;
+    ASSERT_EQUAL_TOL(0.0, mean, 3.0/sqrt((double) numValues));
+    ASSERT_EQUAL_TOL(1.0, c2, 3.0/pow(numValues, 1.0/3.0));
+    ASSERT_EQUAL_TOL(0.0, c3, 3.0/pow(numValues, 1.0/4.0));
+    ASSERT_EQUAL_TOL(0.0, c4, 3.0/pow(numValues, 1.0/4.0));
+}
+
 int main() {
     try {
         testSingleBond();
@@ -407,6 +476,7 @@ int main() {
         testMonteCarlo();
         testSum();
         testParameter();
+        testRandomDistributions();
     }
     catch(const exception& e) {
         cout << "exception: " << e.what() << endl;

platforms/reference/tests/TestReferenceCustomIntegrator.cpp

@@ -384,6 +384,75 @@ void testParameter() {
     }
 }
 
+/**
+ * Test random number distributions.
+ */
+void testRandomDistributions() {
+    const int numParticles = 100;
+    const int numBins = 20;
+    const int numSteps = 100;
+    ReferencePlatform platform;
+    System system;
+    for (int i = 0; i < numParticles; i++)
+        system.addParticle(1.0);
+    CustomIntegrator integrator(0.1);
+    integrator.addPerDofVariable("a", 0);
+    integrator.addPerDofVariable("b", 0);
+    integrator.addComputePerDof("a", "uniform");
+    integrator.addComputePerDof("b", "gaussian");
+    Context context(system, integrator, platform);
+    
+    // See if the random numbers are distributed correctly.
+    
+    vector<int> bins(numBins);
+    double mean = 0.0;
+    double var = 0.0;
+    double skew = 0.0;
+    double kurtosis = 0.0;
+    vector<Vec3> values;
+    for (int i = 0; i < numSteps; i++) {
+        integrator.step(1);
+        integrator.getPerDofVariable(0, values);
+        for (int i = 0; i < numParticles; i++)
+            for (int j = 0; j < 3; j++) {
+                double v = values[i][j];
+                ASSERT(v >= 0 && v < 1);
+                bins[(int) (v*numBins)]++;
+            }
+        integrator.getPerDofVariable(1, values);
+        for (int i = 0; i < numParticles; i++)
+            for (int j = 0; j < 3; j++) {
+                double v = values[i][j];
+                mean += v;
+                var += v*v;
+                skew += v*v*v;
+                kurtosis += v*v*v*v;
+            }
+    }
+    
+    // Check the distribution of uniform randoms.
+    
+    int numValues = numParticles*numSteps*3;
+    double expected = numValues/(double) numBins;
+    double tol = 4*sqrt(expected);
+    for (int i = 0; i < numBins; i++)
+        ASSERT(bins[i] >= expected-tol && bins[i] <= expected+tol);
+    
+    // Check the distribution of gaussian randoms.
+    
+    mean /= numValues;
+    var /= numValues;
+    skew /= numValues;
+    kurtosis /= numValues;
+    double c2 = var-mean*mean;
+    double c3 = skew-3*var*mean+2*mean*mean*mean;
+    double c4 = kurtosis-4*skew*mean-3*var*var+12*var*mean*mean-6*mean*mean*mean*mean;
+    ASSERT_EQUAL_TOL(0.0, mean, 3.0/sqrt((double) numValues));
+    ASSERT_EQUAL_TOL(1.0, c2, 3.0/pow(numValues, 1.0/3.0));
+    ASSERT_EQUAL_TOL(0.0, c3, 3.0/pow(numValues, 1.0/4.0));
+    ASSERT_EQUAL_TOL(0.0, c4, 3.0/pow(numValues, 1.0/4.0));
+}
+
 int main() {
     try {
         testSingleBond();
@@ -393,6 +462,7 @@ int main() {
         testMonteCarlo();
         testSum();
         testParameter();
+        testRandomDistributions();
     }
     catch(const exception& e) {
         cout << "exception: " << e.what() << endl;