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