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