Implemented PME reciprocal convolution kernel

platforms/cuda/src/kernels/gpu.cpp

@@ -780,6 +780,64 @@ void gpuSetPMEParameters(gpuContext gpu, float alpha)
     gpu->sim.pPmeParticleIndex = gpu->psPmeParticleIndex->_pDevData;
     gpu->psPmeParticleFraction = new CUDAStream<float4>(gpu->natoms, 1, "PmeParticleFraction");
     gpu->sim.pPmeParticleFraction = gpu->psPmeParticleFraction->_pDevData;
+
+    // Initialize the b-spline moduli.
+
+    int maxSize = max(max(gridSize.x, gridSize.y), gridSize.z);
+    vector<double> data(PME_ORDER);
+    vector<double> ddata(PME_ORDER);
+    vector<double> bsplines_data(maxSize);
+    data[PME_ORDER-1] = 0.0;
+    data[1] = 0.0;
+    data[0] = 1.0;
+    for (int i = 3; i < PME_ORDER; i++)
+    {
+        double div = 1.0/(i-1.0);
+        data[i-1] = 0.0;
+        for (int j = 1; j < (i-1); j++)
+            data[i-j-1] = div*(j*data[i-j-2]+(i-j)*data[i-j-1]);
+        data[0] = div*data[0];
+    }
+
+    // Differentiate.
+
+    ddata[0] = -data[0];
+    for (int i = 1; i < PME_ORDER; i++)
+        ddata[i] = data[i-1]-data[i];
+    double div = 1.0/(PME_ORDER-1);
+    data[PME_ORDER-1] = 0.0;
+    for (int i = 1; i < (PME_ORDER-1); i++)
+        data[PME_ORDER-i-1] = div*(i*data[PME_ORDER-i-2]+(PME_ORDER-i)*data[PME_ORDER-i-1]);
+    data[0] = div*data[0];
+    for (int i = 0; i < maxSize; i++)
+        bsplines_data[i] = 0.0;
+    for (int i = 1; i <= PME_ORDER; i++)
+        bsplines_data[i] = data[i-1];
+
+    // Evaluate the actual bspline moduli for X/Y/Z.
+
+    for(int dim = 0; dim < 3; dim++)
+    {
+        int ndata = (dim == 0 ? gridSize.x : dim == 1 ? gridSize.y : gridSize.z);
+        for (int i = 0; i < ndata; i++)
+        {
+            double sc = 0.0;
+            double ss = 0.0;
+            for (int j = 0; j < ndata; j++)
+            {
+                double arg = (2.0*M_PI*i*j)/ndata;
+                sc += bsplines_data[j]*cos(arg);
+                ss += bsplines_data[j]*sin(arg);
+            }
+            (*gpu->psPmeBsplineModuli[dim])[i] = sc*sc+ss*ss;
+        }
+        for (int i = 0; i < ndata; i++)
+        {
+            if ((*gpu->psPmeBsplineModuli[dim])[i] < 1.0e-7)
+                (*gpu->psPmeBsplineModuli[dim])[i] = ((*gpu->psPmeBsplineModuli[dim])[i-1]+(*gpu->psPmeBsplineModuli[dim])[i+1])*0.5;
+        }
+        gpu->psPmeBsplineModuli[dim]->Upload();
+    }
 }
 
 extern "C"

platforms/cuda/src/kernels/kCalculatePME.cu

@@ -262,6 +262,39 @@ __global__ void kGridSpreadCharge_kernel()
     }
 }
 
+__global__ void kReciprocalConvolution_kernel()
+{
+    const unsigned int gridSize = cSim.pmeGridSize.x*cSim.pmeGridSize.y*cSim.pmeGridSize.z;
+    float expFactor = PI*PI/(cSim.alphaEwald*cSim.alphaEwald);
+    float scaleFactor = 1.0/(PI*cSim.periodicBoxSizeX*cSim.periodicBoxSizeY*cSim.periodicBoxSizeZ);
+    float energy = 0.0f;
+    for (int index = blockIdx.x*blockDim.x+threadIdx.x; index < gridSize; index += blockDim.x*gridDim.x)
+    {
+        int kx = index/(cSim.pmeGridSize.y*cSim.pmeGridSize.z);
+        int remainder = index-kx*cSim.pmeGridSize.y*cSim.pmeGridSize.z;
+        int ky = remainder/cSim.pmeGridSize.z;
+        int kz = remainder-ky*cSim.pmeGridSize.z;
+        if (kx == 0 && ky == 0 && kz == 0)
+            continue;
+        int mx = (kx < (cSim.pmeGridSize.x+1)/2) ? kx : (kx-cSim.pmeGridSize.x);
+        int my = (ky < (cSim.pmeGridSize.y+1)/2) ? ky : (ky-cSim.pmeGridSize.y);
+        int mz = (kz < (cSim.pmeGridSize.z+1)/2) ? kz : (kz-cSim.pmeGridSize.z);
+        float mhx = mx/cSim.periodicBoxSizeX;
+        float mhy = my/cSim.periodicBoxSizeY;
+        float mhz = mz/cSim.periodicBoxSizeZ;
+        float bx = cSim.pPmeBsplineModuli[0][kx];
+        float by = cSim.pPmeBsplineModuli[1][ky];
+        float bz = cSim.pPmeBsplineModuli[2][kz];
+        cuComplex grid = cSim.pPmeGrid[index];
+        float m2 = mhx*mhx+mhy*mhy+mhz*mhz;
+        float denom = m2*bx*by*bz;
+        float eterm = scaleFactor*exp(-expFactor*m2)/denom;
+        cSim.pPmeGrid[index] = make_cuComplex(grid.x*eterm, grid.y*eterm);
+        energy += eterm*(grid.x*grid.x + grid.y*grid.y);
+    }
+    cSim.pEnergy[blockIdx.x*blockDim.x+threadIdx.x] += energy;
+}
+
 void kCalculatePME(gpuContext gpu)
 {
 //    printf("kCalculatePME\n");
@@ -272,4 +305,6 @@ void kCalculatePME(gpuContext gpu)
     kGridSpreadCharge_kernel<<<gpu->sim.blocks, 64, 64*(sizeof(float4)+sizeof(int4))>>>();
     LAUNCHERROR("kUpdateBsplines");
     cufftExecC2C(gpu->fftplan, gpu->psPmeGrid->_pDevData, gpu->psPmeGrid->_pDevData, CUFFT_FORWARD);
+    kReciprocalConvolution_kernel<<<gpu->sim.blocks, gpu->sim.nonbond_threads_per_block>>>();
+    LAUNCHERROR("kUpdateGridIndexAndFraction");
 }