Further cleanup to handling of random numbers by CustomIntegrator

platforms/cuda/src/CudaKernels.cpp

@@ -4844,8 +4844,7 @@ void CudaIntegrateCustomStepKernel::prepareForComputation(ContextImpl& context,
                 continue;
             if (stepType[step-1] == CustomIntegrator::ComputeGlobal && stepType[step] == CustomIntegrator::ComputeGlobal)
                 merged[step] = true;
-            if (stepType[step-1] == CustomIntegrator::ComputePerDof && stepType[step] == CustomIntegrator::ComputePerDof &&
-                    !usesVariable(expression[step], "uniform"))
+            if (stepType[step-1] == CustomIntegrator::ComputePerDof && stepType[step] == CustomIntegrator::ComputePerDof)
                 merged[step] = true;
         }
         
@@ -5123,7 +5122,8 @@ void CudaIntegrateCustomStepKernel::execute(ContextImpl& context, CustomIntegrat
 
     // Loop over computation steps in the integrator and execute them.
 
-    void* randomArgs[] = {&uniformRandoms->getDevicePointer(), &randomSeed->getDevicePointer()};
+    int maxUniformRandoms = uniformRandoms->getSize();
+    void* randomArgs[] = {&maxUniformRandoms, &uniformRandoms->getDevicePointer(), &randomSeed->getDevicePointer()};
     CUdeviceptr posCorrection = (cu.getUseMixedPrecision() ? cu.getPosqCorrection().getDevicePointer() : 0);
     for (int i = 0; i < numSteps; i++) {
         int lastForceGroups = context.getLastForceGroups();

platforms/cuda/src/kernels/customIntegrator.cu

@@ -52,10 +52,10 @@ extern "C" __global__ void applyPositionDeltas(real4* __restrict__ posq, real4*
     }
 }
 
-extern "C" __global__ void generateRandomNumbers(float4* __restrict__ random, uint4* __restrict__ seed) {
+extern "C" __global__ void generateRandomNumbers(int numValues, float4* __restrict__ random, uint4* __restrict__ seed) {
     uint4 state = seed[blockIdx.x*blockDim.x+threadIdx.x];
     unsigned int carry = 0;
-    for (int index = blockIdx.x*blockDim.x+threadIdx.x; index < NUM_ATOMS; index += blockDim.x*gridDim.x) {
+    for (int index = blockIdx.x*blockDim.x+threadIdx.x; index < numValues; index += blockDim.x*gridDim.x) {
         // Generate three uniform random numbers.
 
         state.x = state.x * 69069 + 1;

platforms/opencl/src/OpenCLKernels.cpp

@@ -5072,8 +5072,7 @@ void OpenCLIntegrateCustomStepKernel::prepareForComputation(ContextImpl& context
                 continue;
             if (stepType[step-1] == CustomIntegrator::ComputeGlobal && stepType[step] == CustomIntegrator::ComputeGlobal)
                 merged[step] = true;
-            if (stepType[step-1] == CustomIntegrator::ComputePerDof && stepType[step] == CustomIntegrator::ComputePerDof &&
-                    !usesVariable(expression[step], "uniform"))
+            if (stepType[step-1] == CustomIntegrator::ComputePerDof && stepType[step] == CustomIntegrator::ComputePerDof)
                 merged[step] = true;
         }
         
@@ -5235,8 +5234,9 @@ void OpenCLIntegrateCustomStepKernel::prepareForComputation(ContextImpl& context
         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());
+        randomKernel.setArg<cl_int>(0, maxUniformRandoms);
+        randomKernel.setArg<cl::Buffer>(1, uniformRandoms->getDeviceBuffer());
+        randomKernel.setArg<cl::Buffer>(2, randomSeed->getDeviceBuffer());
         
         // Create the kernel for summing the potential energy.
 

platforms/opencl/src/kernels/customIntegrator.cl

@@ -52,10 +52,10 @@ __kernel void applyPositionDeltas(__global real4* restrict posq, __global real4*
     }
 }
 
-__kernel void generateRandomNumbers(__global float4* restrict random, __global uint4* restrict seed) {
+__kernel void generateRandomNumbers(int numValues, __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)) {
+    for (int index = get_global_id(0); index < numValues; index += get_global_size(0)) {
         // Generate three uniform random numbers.
 
         state.x = state.x * 69069 + 1;