src/gromacs/listed_forces/gpubondedkernels.cu

   1 /*
   2  * This file is part of the GROMACS molecular simulation package.
   3  *
   4  * Copyright (c) 2018,2019, by the GROMACS development team, led by
   5  * Mark Abraham, David van der Spoel, Berk Hess, and Erik Lindahl,
   6  * and including many others, as listed in the AUTHORS file in the
   7  * top-level source directory and at http://www.gromacs.org.
   8  *
   9  * GROMACS is free software; you can redistribute it and/or
  10  * modify it under the terms of the GNU Lesser General Public License
  11  * as published by the Free Software Foundation; either version 2.1
  12  * of the License, or (at your option) any later version.
  13  *
  14  * GROMACS is distributed in the hope that it will be useful,
  15  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  16  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  17  * Lesser General Public License for more details.
  18  *
  19  * You should have received a copy of the GNU Lesser General Public
  20  * License along with GROMACS; if not, see
  21  * http://www.gnu.org/licenses, or write to the Free Software Foundation,
  22  * Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA.
  23  *
  24  * If you want to redistribute modifications to GROMACS, please
  25  * consider that scientific software is very special. Version
  26  * control is crucial - bugs must be traceable. We will be happy to
  27  * consider code for inclusion in the official distribution, but
  28  * derived work must not be called official GROMACS. Details are found
  29  * in the README & COPYING files - if they are missing, get the
  30  * official version at http://www.gromacs.org.
  31  *
  32  * To help us fund GROMACS development, we humbly ask that you cite
  33  * the research papers on the package. Check out http://www.gromacs.org.
  34  */
  35 /*! \internal \file
  36  *
  37  * \brief Implements CUDA bonded functionality
  38  *
  39  * \author Jon Vincent <jvincent@nvidia.com>
  40  * \author Magnus Lundborg <lundborg.magnus@gmail.com>
  41  * \author Berk Hess <hess@kth.se>
  42  * \author Szilárd Páll <pall.szilard@gmail.com>
  43  * \author Alan Gray <alang@nvidia.com>
  44  * \author Mark Abraham <mark.j.abraham@gmail.com>
  45  *
  46  * \ingroup module_listed_forces
  47  */
  48
  49 #include "gmxpre.h"
  50
  51 #include <cassert>
  52
  53 #include <math_constants.h>
  54
  55 #include "gromacs/gpu_utils/cudautils.cuh"
  56 #include "gromacs/gpu_utils/gpu_vec.cuh"
  57 #include "gromacs/listed_forces/gpubonded.h"
  58 #include "gromacs/math/units.h"
  59 #include "gromacs/mdlib/force_flags.h"
  60 #include "gromacs/mdtypes/forcerec.h"
  61 #include "gromacs/mdtypes/simulation_workload.h"
  62 #include "gromacs/pbcutil/pbc.h"
  63 #include "gromacs/pbcutil/pbc_aiuc_cuda.cuh"
  64 #include "gromacs/utility/gmxassert.h"
  65
  66 #include "gpubonded_impl.h"
  67
  68 #if defined(_MSVC)
  69 #include <limits>
  70 #endif
  71
  72 //CUDA threads per block
  73 #define TPB_BONDED 256
  74
  75 /*-------------------------------- CUDA kernels-------------------------------- */
  76 /*------------------------------------------------------------------------------*/
  77
  78 #define CUDA_DEG2RAD_F          (CUDART_PI_F/180.0f)
  79
  80 /*---------------- BONDED CUDA kernels--------------*/
  81
  82 /* Harmonic */
  83 __device__ __forceinline__
  84 static void harmonic_gpu(const float kA, const float xA, const float x, float *V, float *F)
  85 {
  86     constexpr float half = 0.5f;
  87     float           dx, dx2;
  88
  89     dx    = x-xA;
  90     dx2   = dx*dx;
  91
  92     *F = -kA*dx;
  93     *V = half*kA*dx2;
  94 }
  95
  96 template <bool calcVir, bool calcEner>
  97 __device__
  98 void bonds_gpu(const int i, float *vtot_loc, const int numBonds,
  99                const t_iatom d_forceatoms[], const t_iparams d_forceparams[],
 100                const float4 gm_xq[], fvec gm_f[], fvec sm_fShiftLoc[],
 101                const PbcAiuc pbcAiuc)
 102 {
 103     if (i < numBonds)
 104     {
 105         int3  bondData = *(int3 *)(d_forceatoms + 3 * i);
 106         int   type     = bondData.x;
 107         int   ai       = bondData.y;
 108         int   aj       = bondData.z;
 109
 110         /* dx = xi - xj, corrected for periodic boundary conditions. */
 111         fvec  dx;
 112         int   ki = pbcDxAiuc<calcVir>(pbcAiuc, gm_xq[ai], gm_xq[aj], dx);
 113
 114         float dr2 = iprod_gpu(dx, dx);
 115         float dr  = sqrt(dr2);
 116
 117         float vbond;
 118         float fbond;
 119         harmonic_gpu(d_forceparams[type].harmonic.krA,
 120                      d_forceparams[type].harmonic.rA,
 121                      dr, &vbond, &fbond);
 122
 123         if (calcEner)
 124         {
 125             *vtot_loc += vbond;
 126         }
 127
 128         if (dr2 != 0.0f)
 129         {
 130             fbond *= rsqrtf(dr2);
 131
 132 #pragma unroll
 133             for (int m = 0; m < DIM; m++)
 134             {
 135                 float fij = fbond*dx[m];
 136                 atomicAdd(&gm_f[ai][m], fij);
 137                 atomicAdd(&gm_f[aj][m], -fij);
 138                 if (calcVir && ki != CENTRAL)
 139                 {
 140                     atomicAdd(&sm_fShiftLoc[ki][m], fij);
 141                     atomicAdd(&sm_fShiftLoc[CENTRAL][m], -fij);
 142                 }
 143             }
 144         }
 145     }
 146 }
 147
 148 template <bool returnShift>
 149 __device__ __forceinline__
 150 static float bond_angle_gpu(const float4 xi, const float4 xj, const float4 xk,
 151                             const PbcAiuc &pbcAiuc,
 152                             fvec r_ij, fvec r_kj, float *costh,
 153                             int *t1, int *t2)
 154 /* Return value is the angle between the bonds i-j and j-k */
 155 {
 156     *t1      = pbcDxAiuc<returnShift>(pbcAiuc, xi, xj, r_ij);
 157     *t2      = pbcDxAiuc<returnShift>(pbcAiuc, xk, xj, r_kj);
 158
 159     *costh   = cos_angle_gpu(r_ij, r_kj);
 160     float th = acosf(*costh);
 161
 162     return th;
 163 }
 164
 165 template <bool calcVir, bool calcEner>
 166 __device__
 167 void angles_gpu(const int i, float *vtot_loc, const int numBonds,
 168                 const t_iatom d_forceatoms[], const t_iparams d_forceparams[],
 169                 const float4 gm_xq[], fvec gm_f[], fvec sm_fShiftLoc[],
 170                 const PbcAiuc pbcAiuc)
 171 {
 172     if (i < numBonds)
 173     {
 174         int4  angleData = *(int4 *)(d_forceatoms + 4 * i);
 175         int   type      = angleData.x;
 176         int   ai        = angleData.y;
 177         int   aj        = angleData.z;
 178         int   ak        = angleData.w;
 179
 180         fvec  r_ij;
 181         fvec  r_kj;
 182         float cos_theta;
 183         int   t1;
 184         int   t2;
 185         float theta =
 186             bond_angle_gpu<calcVir>(gm_xq[ai], gm_xq[aj], gm_xq[ak], pbcAiuc,
 187                                     r_ij, r_kj, &cos_theta, &t1, &t2);
 188
 189         float va;
 190         float dVdt;
 191         harmonic_gpu(d_forceparams[type].harmonic.krA,
 192                      d_forceparams[type].harmonic.rA*CUDA_DEG2RAD_F,
 193                      theta, &va, &dVdt);
 194
 195         if (calcEner)
 196         {
 197             *vtot_loc += va;
 198         }
 199
 200         float cos_theta2 = cos_theta*cos_theta;
 201         if (cos_theta2 < 1.0f)
 202         {
 203             float st    = dVdt*rsqrtf(1.0f - cos_theta2);
 204             float sth   = st*cos_theta;
 205             float nrij2 = iprod_gpu(r_ij, r_ij);
 206             float nrkj2 = iprod_gpu(r_kj, r_kj);
 207
 208             float nrij_1 = rsqrtf(nrij2);
 209             float nrkj_1 = rsqrtf(nrkj2);
 210
 211             float cik = st*nrij_1*nrkj_1;
 212             float cii = sth*nrij_1*nrij_1;
 213             float ckk = sth*nrkj_1*nrkj_1;
 214
 215             fvec  f_i;
 216             fvec  f_k;
 217             fvec  f_j;
 218 #pragma unroll
 219             for (int m = 0; m < DIM; m++)
 220             {
 221                 f_i[m]    = -(cik*r_kj[m] - cii*r_ij[m]);
 222                 f_k[m]    = -(cik*r_ij[m] - ckk*r_kj[m]);
 223                 f_j[m]    = -f_i[m] - f_k[m];
 224                 atomicAdd(&gm_f[ai][m], f_i[m]);
 225                 atomicAdd(&gm_f[aj][m], f_j[m]);
 226                 atomicAdd(&gm_f[ak][m], f_k[m]);
 227                 if (calcVir)
 228                 {
 229                     atomicAdd(&sm_fShiftLoc[t1][m], f_i[m]);
 230                     atomicAdd(&sm_fShiftLoc[CENTRAL][m], f_j[m]);
 231                     atomicAdd(&sm_fShiftLoc[t2][m], f_k[m]);
 232                 }
 233             }
 234         }
 235
 236     }
 237 }
 238
 239 template <bool calcVir, bool calcEner>
 240 __device__
 241 void urey_bradley_gpu(const int i, float *vtot_loc, const int numBonds,
 242                       const t_iatom d_forceatoms[], const t_iparams d_forceparams[],
 243                       const float4 gm_xq[], fvec gm_f[], fvec sm_fShiftLoc[],
 244                       const PbcAiuc pbcAiuc)
 245 {
 246     if (i < numBonds)
 247     {
 248         int4  ubData = *(int4 *)(d_forceatoms + 4 * i);
 249         int   type   = ubData.x;
 250         int   ai     = ubData.y;
 251         int   aj     = ubData.z;
 252         int   ak     = ubData.w;
 253
 254         float th0A  = d_forceparams[type].u_b.thetaA*CUDA_DEG2RAD_F;
 255         float kthA  = d_forceparams[type].u_b.kthetaA;
 256         float r13A  = d_forceparams[type].u_b.r13A;
 257         float kUBA  = d_forceparams[type].u_b.kUBA;
 258
 259         fvec  r_ij;
 260         fvec  r_kj;
 261         float cos_theta;
 262         int   t1;
 263         int   t2;
 264         float theta = bond_angle_gpu<calcVir>(gm_xq[ai], gm_xq[aj], gm_xq[ak], pbcAiuc,
 265                                               r_ij, r_kj, &cos_theta, &t1, &t2);
 266
 267         float va;
 268         float dVdt;
 269         harmonic_gpu(kthA, th0A, theta, &va, &dVdt);
 270
 271         if (calcEner)
 272         {
 273             *vtot_loc += va;
 274         }
 275
 276         fvec  r_ik;
 277         int   ki = pbcDxAiuc<calcVir>(pbcAiuc, gm_xq[ai], gm_xq[ak], r_ik);
 278
 279         float dr2  = iprod_gpu(r_ik, r_ik);
 280         float dr   = dr2*rsqrtf(dr2);
 281
 282         float vbond;
 283         float fbond;
 284         harmonic_gpu(kUBA, r13A, dr, &vbond, &fbond);
 285
 286         float cos_theta2 = cos_theta*cos_theta;
 287         if (cos_theta2 < 1.0f)
 288         {
 289             float st    = dVdt*rsqrtf(1.0f - cos_theta2);
 290             float sth   = st*cos_theta;
 291
 292             float nrkj2 = iprod_gpu(r_kj, r_kj);
 293             float nrij2 = iprod_gpu(r_ij, r_ij);
 294
 295             float cik   = st*rsqrtf(nrkj2*nrij2);
 296             float cii   = sth/nrij2;
 297             float ckk   = sth/nrkj2;
 298
 299             fvec  f_i;
 300             fvec  f_j;
 301             fvec  f_k;
 302 #pragma unroll
 303             for (int m = 0; m < DIM; m++)
 304             {
 305                 f_i[m]    = -(cik*r_kj[m]-cii*r_ij[m]);
 306                 f_k[m]    = -(cik*r_ij[m]-ckk*r_kj[m]);
 307                 f_j[m]    = -f_i[m]-f_k[m];
 308                 atomicAdd(&gm_f[ai][m], f_i[m]);
 309                 atomicAdd(&gm_f[aj][m], f_j[m]);
 310                 atomicAdd(&gm_f[ak][m], f_k[m]);
 311                 if (calcVir)
 312                 {
 313                     atomicAdd(&sm_fShiftLoc[t1][m], f_i[m]);
 314                     atomicAdd(&sm_fShiftLoc[CENTRAL][m], f_j[m]);
 315                     atomicAdd(&sm_fShiftLoc[t2][m], f_k[m]);
 316                 }
 317             }
 318         }
 319
 320         /* Time for the bond calculations */
 321         if (dr2 != 0.0f)
 322         {
 323             if (calcEner)
 324             {
 325                 *vtot_loc += vbond;
 326             }
 327
 328             fbond *= rsqrtf(dr2);
 329
 330 #pragma unroll
 331             for (int m = 0; m < DIM; m++)
 332             {
 333                 float fik = fbond*r_ik[m];
 334                 atomicAdd(&gm_f[ai][m], fik);
 335                 atomicAdd(&gm_f[ak][m], -fik);
 336
 337                 if (calcVir && ki != CENTRAL)
 338                 {
 339                     atomicAdd(&sm_fShiftLoc[ki][m], fik);
 340                     atomicAdd(&sm_fShiftLoc[CENTRAL][m], -fik);
 341                 }
 342             }
 343         }
 344     }
 345 }
 346
 347 template <bool returnShift, typename T>
 348 __device__ __forceinline__
 349 static float dih_angle_gpu(const T xi, const T xj, const T xk, const T xl,
 350                            const PbcAiuc &pbcAiuc,
 351                            fvec r_ij, fvec r_kj, fvec r_kl, fvec m, fvec n,
 352                            int *t1, int *t2, int *t3)
 353 {
 354     *t1 = pbcDxAiuc<returnShift>(pbcAiuc, xi, xj, r_ij);
 355     *t2 = pbcDxAiuc<returnShift>(pbcAiuc, xk, xj, r_kj);
 356     *t3 = pbcDxAiuc<returnShift>(pbcAiuc, xk, xl, r_kl);
 357
 358     cprod_gpu(r_ij, r_kj, m);
 359     cprod_gpu(r_kj, r_kl, n);
 360     float phi  = gmx_angle_gpu(m, n);
 361     float ipr  = iprod_gpu(r_ij, n);
 362     float sign = (ipr < 0.0f) ? -1.0f : 1.0f;
 363     phi        = sign*phi;
 364
 365     return phi;
 366 }
 367
 368
 369 __device__ __forceinline__
 370 static void dopdihs_gpu(const float cpA, const float phiA, const int mult,
 371                         const float phi, float *v, float *f)
 372 {
 373     float mdphi, sdphi;
 374
 375     mdphi = mult*phi - phiA*CUDA_DEG2RAD_F;
 376     sdphi = sinf(mdphi);
 377     *v    = cpA * (1.0f + cosf(mdphi));
 378     *f    = -cpA*mult*sdphi;
 379 }
 380
 381 template <bool calcVir>
 382 __device__
 383 static void do_dih_fup_gpu(const int i, const int j, const int k, const int l,
 384                            const float ddphi, const fvec r_ij, const fvec r_kj, const fvec r_kl,
 385                            const fvec m, const fvec n, fvec gm_f[], fvec sm_fShiftLoc[],
 386                            const PbcAiuc &pbcAiuc,
 387                            const float4 gm_xq[], const int t1, const int t2, const int gmx_unused t3)
 388 {
 389     float iprm  = iprod_gpu(m, m);
 390     float iprn  = iprod_gpu(n, n);
 391     float nrkj2 = iprod_gpu(r_kj, r_kj);
 392     float toler = nrkj2*GMX_REAL_EPS;
 393     if ((iprm > toler) && (iprn > toler))
 394     {
 395         float nrkj_1 = rsqrtf(nrkj2); // replacing std::invsqrt call
 396         float nrkj_2 = nrkj_1*nrkj_1;
 397         float nrkj   = nrkj2*nrkj_1;
 398         float a      = -ddphi*nrkj/iprm;
 399         fvec  f_i;
 400         svmul_gpu(a, m, f_i);
 401         float b      = ddphi*nrkj/iprn;
 402         fvec  f_l;
 403         svmul_gpu(b, n, f_l);
 404         float p      = iprod_gpu(r_ij, r_kj);
 405         p           *= nrkj_2;
 406         float q      = iprod_gpu(r_kl, r_kj);
 407         q           *= nrkj_2;
 408         fvec  uvec;
 409         svmul_gpu(p, f_i, uvec);
 410         fvec  vvec;
 411         svmul_gpu(q, f_l, vvec);
 412         fvec  svec;
 413         fvec_sub_gpu(uvec, vvec, svec);
 414         fvec  f_j;
 415         fvec_sub_gpu(f_i, svec, f_j);
 416         fvec  f_k;
 417         fvec_add_gpu(f_l, svec, f_k);
 418 #pragma unroll
 419         for (int m = 0; (m < DIM); m++)
 420         {
 421             atomicAdd(&gm_f[i][m], f_i[m]);
 422             atomicAdd(&gm_f[j][m], -f_j[m]);
 423             atomicAdd(&gm_f[k][m], -f_k[m]);
 424             atomicAdd(&gm_f[l][m], f_l[m]);
 425         }
 426
 427         if (calcVir)
 428         {
 429             fvec dx_jl;
 430             int  t3 = pbcDxAiuc<calcVir>(pbcAiuc, gm_xq[l], gm_xq[j], dx_jl);
 431
 432 #pragma unroll
 433             for (int m = 0; (m < DIM); m++)
 434             {
 435                 atomicAdd(&sm_fShiftLoc[t1][m], f_i[m]);
 436                 atomicAdd(&sm_fShiftLoc[CENTRAL][m], -f_j[m]);
 437                 atomicAdd(&sm_fShiftLoc[t2][m], -f_k[m]);
 438                 atomicAdd(&sm_fShiftLoc[t3][m], f_l[m]);
 439             }
 440         }
 441     }
 442 }
 443
 444 template <bool calcVir, bool calcEner>
 445 __device__
 446 void  pdihs_gpu(const int i, float *vtot_loc, const int numBonds,
 447                 const t_iatom d_forceatoms[], const t_iparams d_forceparams[],
 448                 const float4 gm_xq[], fvec gm_f[], fvec sm_fShiftLoc[],
 449                 const PbcAiuc pbcAiuc)
 450 {
 451     if (i < numBonds)
 452     {
 453         int   type = d_forceatoms[5*i];
 454         int   ai   = d_forceatoms[5*i + 1];
 455         int   aj   = d_forceatoms[5*i + 2];
 456         int   ak   = d_forceatoms[5*i + 3];
 457         int   al   = d_forceatoms[5*i + 4];
 458
 459         fvec  r_ij;
 460         fvec  r_kj;
 461         fvec  r_kl;
 462         fvec  m;
 463         fvec  n;
 464         int   t1;
 465         int   t2;
 466         int   t3;
 467         float phi  =
 468             dih_angle_gpu<calcVir>(gm_xq[ai], gm_xq[aj], gm_xq[ak], gm_xq[al], pbcAiuc,
 469                                    r_ij, r_kj, r_kl, m, n, &t1, &t2, &t3);
 470
 471         float vpd;
 472         float ddphi;
 473         dopdihs_gpu(d_forceparams[type].pdihs.cpA,
 474                     d_forceparams[type].pdihs.phiA,
 475                     d_forceparams[type].pdihs.mult,
 476                     phi, &vpd, &ddphi);
 477
 478         if (calcEner)
 479         {
 480             *vtot_loc += vpd;
 481         }
 482
 483         do_dih_fup_gpu<calcVir>(ai, aj, ak, al,
 484                                 ddphi, r_ij, r_kj, r_kl, m, n,
 485                                 gm_f, sm_fShiftLoc, pbcAiuc,
 486                                 gm_xq, t1, t2, t3);
 487
 488     }
 489 }
 490
 491 template <bool calcVir, bool calcEner>
 492 __device__
 493 void rbdihs_gpu(const int i, float *vtot_loc, const int numBonds,
 494                 const t_iatom d_forceatoms[], const t_iparams d_forceparams[],
 495                 const float4 gm_xq[], fvec gm_f[], fvec sm_fShiftLoc[],
 496                 const PbcAiuc pbcAiuc)
 497 {
 498     constexpr float  c0 = 0.0f, c1 = 1.0f, c2 = 2.0f, c3 = 3.0f, c4 = 4.0f, c5 = 5.0f;
 499
 500     if (i < numBonds)
 501     {
 502         int   type = d_forceatoms[5*i];
 503         int   ai   = d_forceatoms[5*i+1];
 504         int   aj   = d_forceatoms[5*i+2];
 505         int   ak   = d_forceatoms[5*i+3];
 506         int   al   = d_forceatoms[5*i+4];
 507
 508         fvec  r_ij;
 509         fvec  r_kj;
 510         fvec  r_kl;
 511         fvec  m;
 512         fvec  n;
 513         int   t1;
 514         int   t2;
 515         int   t3;
 516         float phi  =
 517             dih_angle_gpu<calcVir>(gm_xq[ai], gm_xq[aj], gm_xq[ak], gm_xq[al], pbcAiuc,
 518                                    r_ij, r_kj, r_kl, m, n, &t1, &t2, &t3);
 519
 520         /* Change to polymer convention */
 521         if (phi < c0)
 522         {
 523             phi += CUDART_PI_F;
 524         }
 525         else
 526         {
 527             phi -= CUDART_PI_F;
 528
 529         }
 530         float cos_phi = cosf(phi);
 531         /* Beware of accuracy loss, cannot use 1-sqrt(cos^2) ! */
 532         float sin_phi = sinf(phi);
 533
 534         float parm[NR_RBDIHS];
 535         for (int j = 0; j < NR_RBDIHS; j++)
 536         {
 537             parm[j]  = d_forceparams[type].rbdihs.rbcA[j];
 538         }
 539         /* Calculate cosine powers */
 540         /* Calculate the energy */
 541         /* Calculate the derivative */
 542         float v      = parm[0];
 543         float ddphi  = c0;
 544         float cosfac = c1;
 545
 546         float rbp    = parm[1];
 547         ddphi       += rbp*cosfac;
 548         cosfac      *= cos_phi;
 549         if (calcEner)
 550         {
 551             v       += cosfac*rbp;
 552         }
 553         rbp          = parm[2];
 554         ddphi       += c2*rbp*cosfac;
 555         cosfac      *= cos_phi;
 556         if (calcEner)
 557         {
 558             v       += cosfac*rbp;
 559         }
 560         rbp          = parm[3];
 561         ddphi       += c3*rbp*cosfac;
 562         cosfac      *= cos_phi;
 563         if (calcEner)
 564         {
 565             v       += cosfac*rbp;
 566         }
 567         rbp          = parm[4];
 568         ddphi       += c4*rbp*cosfac;
 569         cosfac      *= cos_phi;
 570         if (calcEner)
 571         {
 572             v       += cosfac*rbp;
 573         }
 574         rbp          = parm[5];
 575         ddphi       += c5*rbp*cosfac;
 576         cosfac      *= cos_phi;
 577         if (calcEner)
 578         {
 579             v       += cosfac*rbp;
 580         }
 581
 582         ddphi = -ddphi*sin_phi;
 583
 584         do_dih_fup_gpu<calcVir>(ai, aj, ak, al,
 585                                 ddphi, r_ij, r_kj, r_kl, m, n,
 586                                 gm_f, sm_fShiftLoc, pbcAiuc,
 587                                 gm_xq, t1, t2, t3);
 588         if (calcEner)
 589         {
 590             *vtot_loc += v;
 591         }
 592     }
 593 }
 594
 595 __device__ __forceinline__
 596 static void make_dp_periodic_gpu(float *dp)
 597 {
 598     /* dp cannot be outside (-pi,pi) */
 599     if (*dp >= CUDART_PI_F)
 600     {
 601         *dp -= 2.0f*CUDART_PI_F;
 602     }
 603     else if (*dp < -CUDART_PI_F)
 604     {
 605         *dp += 2.0f*CUDART_PI_F;
 606     }
 607 }
 608
 609 template <bool calcVir, bool calcEner>
 610 __device__
 611 void  idihs_gpu(const int i, float *vtot_loc, const int numBonds,
 612                 const t_iatom d_forceatoms[], const t_iparams d_forceparams[],
 613                 const float4 gm_xq[], fvec gm_f[], fvec sm_fShiftLoc[],
 614                 const PbcAiuc pbcAiuc)
 615 {
 616     if (i < numBonds)
 617     {
 618         int   type = d_forceatoms[5*i];
 619         int   ai   = d_forceatoms[5*i + 1];
 620         int   aj   = d_forceatoms[5*i + 2];
 621         int   ak   = d_forceatoms[5*i + 3];
 622         int   al   = d_forceatoms[5*i + 4];
 623
 624         fvec  r_ij;
 625         fvec  r_kj;
 626         fvec  r_kl;
 627         fvec  m;
 628         fvec  n;
 629         int   t1;
 630         int   t2;
 631         int   t3;
 632         float phi  =
 633             dih_angle_gpu<calcVir>(gm_xq[ai], gm_xq[aj], gm_xq[ak], gm_xq[al], pbcAiuc,
 634                                    r_ij, r_kj, r_kl, m, n, &t1, &t2, &t3);
 635
 636         /* phi can jump if phi0 is close to Pi/-Pi, which will cause huge
 637          * force changes if we just apply a normal harmonic.
 638          * Instead, we first calculate phi-phi0 and take it modulo (-Pi,Pi).
 639          * This means we will never have the periodicity problem, unless
 640          * the dihedral is Pi away from phiO, which is very unlikely due to
 641          * the potential.
 642          */
 643         float kA   = d_forceparams[type].harmonic.krA;
 644         float pA   = d_forceparams[type].harmonic.rA;
 645
 646         float phi0 = pA*CUDA_DEG2RAD_F;
 647
 648         float dp   = phi - phi0;
 649
 650         make_dp_periodic_gpu(&dp);
 651
 652         float ddphi = -kA*dp;
 653
 654         do_dih_fup_gpu<calcVir>(ai, aj, ak, al,
 655                                 -ddphi, r_ij, r_kj, r_kl, m, n,
 656                                 gm_f, sm_fShiftLoc, pbcAiuc,
 657                                 gm_xq, t1, t2, t3);
 658
 659         if (calcEner)
 660         {
 661             *vtot_loc += -0.5f*ddphi*dp;
 662         }
 663     }
 664 }
 665
 666 template <bool calcVir, bool calcEner>
 667 __device__
 668 void pairs_gpu(const int i, const int numBonds,
 669                const t_iatom d_forceatoms[], const t_iparams iparams[],
 670                const float4 gm_xq[], fvec gm_f[], fvec sm_fShiftLoc[],
 671                const PbcAiuc pbcAiuc,
 672                const float scale_factor,
 673                float *vtotVdw_loc, float *vtotElec_loc)
 674 {
 675     if (i <  numBonds)
 676     {
 677         int3  pairData = *(int3 *)(d_forceatoms + 3 * i);
 678         int   type     = pairData.x;
 679         int   ai       = pairData.y;
 680         int   aj       = pairData.z;
 681
 682         float qq    = gm_xq[ai].w*gm_xq[aj].w;
 683         float c6    = iparams[type].lj14.c6A;
 684         float c12   = iparams[type].lj14.c12A;
 685
 686         /* Do we need to apply full periodic boundary conditions? */
 687         fvec  dr;
 688         int   fshift_index = pbcDxAiuc<calcVir>(pbcAiuc, gm_xq[ai], gm_xq[aj], dr);
 689
 690         float r2    = norm2_gpu(dr);
 691         float rinv  = rsqrtf(r2);
 692         float rinv2 = rinv*rinv;
 693         float rinv6 = rinv2*rinv2*rinv2;
 694
 695         /* Calculate the Coulomb force * r */
 696         float velec = scale_factor*qq*rinv;
 697
 698         /* Calculate the LJ force * r and add it to the Coulomb part */
 699         float fr    = (12.0f*c12*rinv6 - 6.0f*c6)*rinv6 + velec;
 700
 701         float finvr = fr * rinv2;
 702         fvec  f;
 703         svmul_gpu(finvr, dr, f);
 704
 705         /* Add the forces */
 706 #pragma unroll
 707         for (int m = 0; m < DIM; m++)
 708         {
 709             atomicAdd(&gm_f[ai][m], f[m]);
 710             atomicAdd(&gm_f[aj][m], -f[m]);
 711             if (calcVir && fshift_index != CENTRAL)
 712             {
 713                 atomicAdd(&sm_fShiftLoc[fshift_index][m], f[m]);
 714                 atomicAdd(&sm_fShiftLoc[CENTRAL][m], -f[m]);
 715             }
 716         }
 717
 718         if (calcEner)
 719         {
 720             *vtotVdw_loc  += (c12*rinv6 - c6)*rinv6;
 721             *vtotElec_loc += velec;
 722         }
 723     }
 724 }
 725
 726 namespace gmx
 727 {
 728
 729 template <bool calcVir, bool calcEner>
 730 __global__
 731 void exec_kernel_gpu(BondedCudaKernelParameters kernelParams)
 732 {
 733     assert(blockDim.y == 1 && blockDim.z == 1);
 734     const int       tid          = blockIdx.x*blockDim.x+threadIdx.x;
 735     float           vtot_loc     = 0;
 736     float           vtotVdw_loc  = 0;
 737     float           vtotElec_loc = 0;
 738     __shared__ fvec sm_fShiftLoc[SHIFTS];
 739
 740     if (calcVir)
 741     {
 742         if (threadIdx.x < SHIFTS)
 743         {
 744             sm_fShiftLoc[threadIdx.x][XX] = 0.0f;
 745             sm_fShiftLoc[threadIdx.x][YY] = 0.0f;
 746             sm_fShiftLoc[threadIdx.x][ZZ] = 0.0f;
 747         }
 748         __syncthreads();
 749     }
 750
 751     int  fType;
 752     bool threadComputedPotential = false;
 753 #pragma unroll
 754     for (int j = 0; j < numFTypesOnGpu; j++)
 755     {
 756         if (tid >= kernelParams.fTypeRangeStart[j] && tid <= kernelParams.fTypeRangeEnd[j])
 757         {
 758             const int      numBonds = kernelParams.numFTypeBonds[j];
 759             int            fTypeTid = tid - kernelParams.fTypeRangeStart[j];
 760             const t_iatom *iatoms   = kernelParams.d_iatoms[j];
 761             fType                   = kernelParams.fTypesOnGpu[j];
 762             if (calcEner)
 763             {
 764                 threadComputedPotential         = true;
 765             }
 766
 767             switch (fType)
 768             {
 769                 case F_BONDS:
 770                     bonds_gpu<calcVir, calcEner>(fTypeTid, &vtot_loc, numBonds, iatoms, kernelParams.d_forceParams,
 771                                                  kernelParams.d_xq, kernelParams.d_f, sm_fShiftLoc, kernelParams.pbcAiuc);
 772                     break;
 773                 case F_ANGLES:
 774                     angles_gpu<calcVir, calcEner>(fTypeTid, &vtot_loc, numBonds, iatoms, kernelParams.d_forceParams,
 775                                                   kernelParams.d_xq, kernelParams.d_f, sm_fShiftLoc, kernelParams.pbcAiuc);
 776                     break;
 777                 case F_UREY_BRADLEY:
 778                     urey_bradley_gpu<calcVir, calcEner>(fTypeTid, &vtot_loc, numBonds, iatoms, kernelParams.d_forceParams,
 779                                                         kernelParams.d_xq, kernelParams.d_f, sm_fShiftLoc, kernelParams.pbcAiuc);
 780                     break;
 781                 case F_PDIHS:
 782                 case F_PIDIHS:
 783                     pdihs_gpu<calcVir, calcEner>(fTypeTid, &vtot_loc, numBonds, iatoms, kernelParams.d_forceParams,
 784                                                  kernelParams.d_xq, kernelParams.d_f, sm_fShiftLoc, kernelParams.pbcAiuc);
 785                     break;
 786                 case F_RBDIHS:
 787                     rbdihs_gpu<calcVir, calcEner>(fTypeTid, &vtot_loc, numBonds, iatoms, kernelParams.d_forceParams,
 788                                                   kernelParams.d_xq, kernelParams.d_f, sm_fShiftLoc, kernelParams.pbcAiuc);
 789                     break;
 790                 case F_IDIHS:
 791                     idihs_gpu<calcVir, calcEner>(fTypeTid, &vtot_loc, numBonds, iatoms, kernelParams.d_forceParams,
 792                                                  kernelParams.d_xq, kernelParams.d_f, sm_fShiftLoc, kernelParams.pbcAiuc);
 793                     break;
 794                 case F_LJ14:
 795                     pairs_gpu<calcVir, calcEner>(fTypeTid, numBonds, iatoms, kernelParams.d_forceParams,
 796                                                  kernelParams.d_xq, kernelParams.d_f, sm_fShiftLoc, kernelParams.pbcAiuc,
 797                                                  kernelParams.scaleFactor, &vtotVdw_loc, &vtotElec_loc);
 798                     break;
 799             }
 800             break;
 801         }
 802     }
 803
 804     if (threadComputedPotential)
 805     {
 806         float *vtotVdw  = kernelParams.d_vTot + F_LJ14;
 807         float *vtotElec = kernelParams.d_vTot + F_COUL14;
 808         atomicAdd(kernelParams.d_vTot + fType, vtot_loc);
 809         atomicAdd(vtotVdw, vtotVdw_loc);
 810         atomicAdd(vtotElec, vtotElec_loc);
 811     }
 812     /* Accumulate shift vectors from shared memory to global memory on the first SHIFTS threads of the block. */
 813     if (calcVir)
 814     {
 815         __syncthreads();
 816         if (threadIdx.x < SHIFTS)
 817         {
 818             fvec_inc_atomic(kernelParams.d_fShift[threadIdx.x], sm_fShiftLoc[threadIdx.x]);
 819         }
 820     }
 821 }
 822
 823
 824 /*-------------------------------- End CUDA kernels-----------------------------*/
 825
 826
 827 template <bool calcVir, bool calcEner>
 828 void
 829 GpuBonded::Impl::launchKernel(const t_forcerec *fr,
 830                               const matrix      box)
 831 {
 832     GMX_ASSERT(haveInteractions_,
 833                "Cannot launch bonded GPU kernels unless bonded GPU work was scheduled");
 834     static_assert(TPB_BONDED >= SHIFTS, "TPB_BONDED must be >= SHIFTS for the virial kernel (calcVir=true)");
 835
 836     PbcAiuc       pbcAiuc;
 837     setPbcAiuc(fr->bMolPBC ? ePBC2npbcdim(fr->ePBC) : 0, box, &pbcAiuc);
 838
 839     int                fTypeRangeEnd = kernelParams_.fTypeRangeEnd[numFTypesOnGpu - 1];
 840
 841     if (fTypeRangeEnd < 0)
 842     {
 843         return;
 844     }
 845
 846     KernelLaunchConfig config;
 847     config.blockSize[0] = TPB_BONDED;
 848     config.blockSize[1] = 1;
 849     config.blockSize[2] = 1;
 850     config.gridSize[0]  = (fTypeRangeEnd + TPB_BONDED)/TPB_BONDED;
 851     config.gridSize[1]  = 1;
 852     config.gridSize[2]  = 1;
 853     config.stream       = stream_;
 854
 855     auto kernelPtr            = exec_kernel_gpu<calcVir, calcEner>;
 856     kernelParams_.scaleFactor = fr->ic->epsfac*fr->fudgeQQ;
 857     kernelParams_.pbcAiuc     = pbcAiuc;
 858
 859     const auto kernelArgs     = prepareGpuKernelArguments(kernelPtr, config, &kernelParams_);
 860
 861     launchGpuKernel(kernelPtr, config, nullptr, "exec_kernel_gpu<calcVir, calcEner>", kernelArgs);
 862 }
 863
 864 void
 865 GpuBonded::launchKernel(const t_forcerec        *fr,
 866                         const gmx::StepWorkload &stepWork,
 867                         const matrix             box)
 868 {
 869     if (stepWork.computeEnergy)
 870     {
 871         // When we need the energy, we also need the virial
 872         impl_->launchKernel<true, true>
 873             (fr, box);
 874     }
 875     else if (stepWork.computeVirial)
 876     {
 877         impl_->launchKernel<true, false>
 878             (fr, box);
 879     }
 880     else
 881     {
 882         impl_->launchKernel<false, false>
 883             (fr, box);
 884     }
 885 }
 886
 887 } // namespace gmx