src/gromacs/listed_forces/gpubondedkernels.cu

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