Cleanup to CUDA dispersion PME code

openmmapi/include/openmm/internal/NonbondedForceImpl.h

@@ -75,7 +75,7 @@ public:
      * This is a utility routine that calculates the values to use for alpha and grid size when using
      * Particle Mesh Ewald.
      */
-    static void calcPMEParameters(const System& system, const NonbondedForce& force, double& alpha, int& xsize, int& ysize, int& zsize, bool LJ);
+    static void calcPMEParameters(const System& system, const NonbondedForce& force, double& alpha, int& xsize, int& ysize, int& zsize, bool lj);
     /**
      * Compute the coefficient which, when divided by the periodic box volume, gives the
      * long range dispersion correction to the energy.

openmmapi/src/NonbondedForceImpl.cpp

@@ -151,12 +151,11 @@ void NonbondedForceImpl::calcEwaldParameters(const System& system, const Nonbond
         kmaxz++;
 }
 
-void NonbondedForceImpl::calcPMEParameters(const System& system, const NonbondedForce& force, double& alpha, int& xsize, int& ysize, int& zsize, bool LJ) {
-    if(LJ) {
+void NonbondedForceImpl::calcPMEParameters(const System& system, const NonbondedForce& force, double& alpha, int& xsize, int& ysize, int& zsize, bool lj) {
+    if (lj)
         force.getLJPMEParameters(alpha, xsize, ysize, zsize);
-    } else {
+    else
         force.getPMEParameters(alpha, xsize, ysize, zsize);
-    }
     if (alpha == 0.0) {
         Vec3 boxVectors[3];
         system.getDefaultPeriodicBoxVectors(boxVectors[0], boxVectors[1], boxVectors[2]);

platforms/cuda/include/CudaKernels.h

@@ -598,10 +598,9 @@ private:
 class CudaCalcNonbondedForceKernel : public CalcNonbondedForceKernel {
 public:
     CudaCalcNonbondedForceKernel(std::string name, const Platform& platform, CudaContext& cu, const System& system) : CalcNonbondedForceKernel(name, platform),
-            cu(cu), hasInitializedFFT(false), sigmaEpsilon(NULL), C6s(NULL), exceptionParams(NULL), cosSinSums(NULL), directPmeGrid(NULL), reciprocalPmeGrid(NULL),
-            directDispersionPmeGrid(NULL), reciprocalDispersionPmeGrid(NULL),
-            pmeBsplineModuliX(NULL), pmeBsplineModuliY(NULL), pmeBsplineModuliZ(NULL),  pmeAtomRange(NULL), pmeAtomGridIndex(NULL), pmeAtomDispersionGridIndex(NULL),
-            pmeEnergyBuffer(NULL), dispersionPmeEnergyBuffer(NULL), sort(NULL), dispersionFft(NULL), fft(NULL), pmeio(NULL), dispersionPmeio(NULL) {
+            cu(cu), hasInitializedFFT(false), sigmaEpsilon(NULL), exceptionParams(NULL), cosSinSums(NULL), directPmeGrid(NULL), reciprocalPmeGrid(NULL),
+            pmeBsplineModuliX(NULL), pmeBsplineModuliY(NULL), pmeBsplineModuliZ(NULL),  pmeAtomRange(NULL), pmeAtomGridIndex(NULL),
+            pmeEnergyBuffer(NULL), sort(NULL), dispersionFft(NULL), fft(NULL), pmeio(NULL) {
     }
     ~CudaCalcNonbondedForceKernel();
     /**
@@ -666,28 +665,22 @@ private:
     CudaContext& cu;
     bool hasInitializedFFT;
     CudaArray* sigmaEpsilon;
-    CudaArray* C6s;
     CudaArray* exceptionParams;
     CudaArray* cosSinSums;
     CudaArray* directPmeGrid;
     CudaArray* reciprocalPmeGrid;
-    CudaArray* directDispersionPmeGrid;
-    CudaArray* reciprocalDispersionPmeGrid;
     CudaArray* pmeBsplineModuliX;
     CudaArray* pmeBsplineModuliY;
     CudaArray* pmeBsplineModuliZ;
     CudaArray* pmeAtomRange;
     CudaArray* pmeAtomGridIndex;
-    CudaArray* pmeAtomDispersionGridIndex;
     CudaArray* pmeEnergyBuffer;
-    CudaArray* dispersionPmeEnergyBuffer;
     CudaSort* sort;
     Kernel cpuPme;
     Kernel cpuDispersionPme;
     PmeIO* pmeio;
-    PmeIO* dispersionPmeio;
-    CUstream pmeStream, dispersionPmeStream;
-    CUevent pmeSyncEvent, dispersionPmeSyncEvent;
+    CUstream pmeStream;
+    CUevent pmeSyncEvent;
     CudaFFT3D* fft;
     cufftHandle fftForward;
     cufftHandle fftBackward;
@@ -710,7 +703,7 @@ private:
     CUfunction pmeInterpolateDispersionForceKernel;
     std::map<std::string, std::string> pmeDefines;
     std::vector<std::pair<int, int> > exceptionAtoms;
-    double ewaldSelfEnergy, dispersionSelfEnergy, dispersionCoefficient, alpha, dispersionAlpha;
+    double ewaldSelfEnergy, dispersionCoefficient, alpha, dispersionAlpha;
     int interpolateForceThreads;
     int gridSizeX, gridSizeY, gridSizeZ;
     int dispersionGridSizeX, dispersionGridSizeY, dispersionGridSizeZ;

platforms/cuda/src/kernels/coulombLennardJones.cu

@@ -30,7 +30,8 @@
             const real dar6 = dar4*dar2;
             const real invR2 = invR*invR;
             const real expDar2 = EXP(-dar2);
-            const real c6 = C6s1*C6s2;
+            const real2 sigExpProd = sigmaEpsilon1*sigmaEpsilon2;
+            const real c6 = 64*sigExpProd.x*sigExpProd.x*sigExpProd.x*sigExpProd.y;
             const real coef = invR2*invR2*invR2*c6;
             const real eprefac = 1.0f + dar2 + 0.5f*dar4;
             const real dprefac = eprefac + dar6/6.0f;

platforms/cuda/src/kernels/ljpme.cu

@@ -22,7 +22,7 @@ extern "C" __global__ void findAtomDispersionGridIndex(const real4* __restrict__
 extern "C" __global__ void gridSpreadC6(const real4* __restrict__ posq, real* __restrict__ originalPmeGrid,
         real4 periodicBoxSize, real4 invPeriodicBoxSize, real4 periodicBoxVecX, real4 periodicBoxVecY, real4 periodicBoxVecZ,
         real3 recipBoxVecX, real3 recipBoxVecY, real3 recipBoxVecZ, const int2* __restrict__ pmeAtomGridIndex,
-        const real* __restrict__ C6s) {
+        const real2* __restrict__ sigmaEpsilon) {
     real3 data[PME_ORDER];
     const real scale = RECIP(PME_ORDER-1);
     
@@ -63,7 +63,9 @@ extern "C" __global__ void gridSpreadC6(const real4* __restrict__ posq, real* __
         data[0] = scale*(make_real3(1)-dr)*data[0];
         
         // Spread the charge from this atom onto each grid point.
-         
+        
+        const real2 sigEps = sigmaEpsilon[atom];
+        const real c6 = 8*sigEps.x*sigEps.x*sigEps.x*sigEps.y;
         for (int ix = 0; ix < PME_ORDER; ix++) {
             int xbase = gridIndex.x+ix;
             xbase -= (xbase >= DISPERSION_GRID_SIZE_X ? DISPERSION_GRID_SIZE_X : 0);
@@ -81,8 +83,7 @@ extern "C" __global__ void gridSpreadC6(const real4* __restrict__ posq, real* __
                     zindex -= (zindex >= DISPERSION_GRID_SIZE_Z ? DISPERSION_GRID_SIZE_Z : 0);
                     int index = ybase + zindex;
 
-                    // We need to grab the C6 coefficient from the array
-                    real add = C6s[atom]*dx*dy*data[iz].z;
+                    real add = c6*dx*dy*data[iz].z;
 #ifdef USE_DOUBLE_PRECISION
                     unsigned long long * ulonglong_p = (unsigned long long *) originalPmeGrid;
                     atomicAdd(&ulonglong_p[index],  static_cast<unsigned long long>((long long) (add*0x100000000)));
@@ -241,7 +242,7 @@ gridEvaluateDispersionEnergy(real2* __restrict__ halfcomplex_pmeGrid, mixed* __r
 extern "C" __global__
 void gridInterpolateDispersionForce(const real4* __restrict__ posq, unsigned long long* __restrict__ forceBuffers, const real* __restrict__ originalPmeGrid,
         real4 periodicBoxSize, real4 invPeriodicBoxSize, real4 periodicBoxVecX, real4 periodicBoxVecY, real4 periodicBoxVecZ,
-        real3 recipBoxVecX, real3 recipBoxVecY, real3 recipBoxVecZ, const int2* __restrict__ pmeAtomGridIndex, const real* __restrict__ C6s) {
+        real3 recipBoxVecX, real3 recipBoxVecY, real3 recipBoxVecZ, const int2* __restrict__ pmeAtomGridIndex, const real2* __restrict__ sigmaEpsilon) {
     real3 data[PME_ORDER];
     real3 ddata[PME_ORDER];
     const real scale = RECIP(PME_ORDER-1);
@@ -287,7 +288,7 @@ void gridInterpolateDispersionForce(const real4* __restrict__ posq, unsigned lon
 
         
         // Compute the force on this atom.
-         
+        
         for (int ix = 0; ix < PME_ORDER; ix++) {
             int xbase = gridIndex.x+ix;
             xbase -= (xbase >= DISPERSION_GRID_SIZE_X ? DISPERSION_GRID_SIZE_X : 0);
@@ -313,7 +314,8 @@ void gridInterpolateDispersionForce(const real4* __restrict__ posq, unsigned lon
                 }
             }
         }
-        real q = C6s[atom];
+        const real2 sigEps = sigmaEpsilon[atom];
+        real q = 8*sigEps.x*sigEps.x*sigEps.x*sigEps.y;
         real forceX = -q*(force.x*DISPERSION_GRID_SIZE_X*recipBoxVecX.x);
         real forceY = -q*(force.x*DISPERSION_GRID_SIZE_X*recipBoxVecY.x+force.y*DISPERSION_GRID_SIZE_Y*recipBoxVecY.y);
         real forceZ = -q*(force.x*DISPERSION_GRID_SIZE_X*recipBoxVecZ.x+force.y*DISPERSION_GRID_SIZE_Y*recipBoxVecZ.y+force.z*DISPERSION_GRID_SIZE_Z*recipBoxVecZ.z);