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