src/gromacs/nbnxm/cuda/nbnxm_cuda_kernel_utils.cuh

   1 /*
   2  * This file is part of the GROMACS molecular simulation package.
   3  *
   4  * Copyright (c) 2012,2013,2014,2016,2017 by the GROMACS development team.
   5  * Copyright (c) 2018,2019,2020,2021, by the GROMACS development team, led by
   6  * Mark Abraham, David van der Spoel, Berk Hess, and Erik Lindahl,
   7  * and including many others, as listed in the AUTHORS file in the
   8  * top-level source directory and at http://www.gromacs.org.
   9  *
  10  * GROMACS is free software; you can redistribute it and/or
  11  * modify it under the terms of the GNU Lesser General Public License
  12  * as published by the Free Software Foundation; either version 2.1
  13  * of the License, or (at your option) any later version.
  14  *
  15  * GROMACS is distributed in the hope that it will be useful,
  16  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  17  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  18  * Lesser General Public License for more details.
  19  *
  20  * You should have received a copy of the GNU Lesser General Public
  21  * License along with GROMACS; if not, see
  22  * http://www.gnu.org/licenses, or write to the Free Software Foundation,
  23  * Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA.
  24  *
  25  * If you want to redistribute modifications to GROMACS, please
  26  * consider that scientific software is very special. Version
  27  * control is crucial - bugs must be traceable. We will be happy to
  28  * consider code for inclusion in the official distribution, but
  29  * derived work must not be called official GROMACS. Details are found
  30  * in the README & COPYING files - if they are missing, get the
  31  * official version at http://www.gromacs.org.
  32  *
  33  * To help us fund GROMACS development, we humbly ask that you cite
  34  * the research papers on the package. Check out http://www.gromacs.org.
  35  */
  36
  37 /*! \internal \file
  38  *  \brief
  39  *  Utility constant and function declaration for the CUDA non-bonded kernels.
  40  *  This header should be included once at the top level, just before the
  41  *  kernels are included (has to be preceded by nbnxn_cuda_types.h).
  42  *
  43  *  \author Szilárd Páll <pall.szilard@gmail.com>
  44  *  \ingroup module_nbnxm
  45  */
  46 #include <assert.h>
  47
  48 /* Note that floating-point constants in CUDA code should be suffixed
  49  * with f (e.g. 0.5f), to stop the compiler producing intermediate
  50  * code that is in double precision.
  51  */
  52
  53 #include "gromacs/gpu_utils/cuda_arch_utils.cuh"
  54 #include "gromacs/gpu_utils/cuda_kernel_utils.cuh"
  55 #include "gromacs/gpu_utils/vectype_ops.cuh"
  56
  57 #include "nbnxm_cuda_types.h"
  58
  59 #ifndef NBNXM_CUDA_KERNEL_UTILS_CUH
  60 #    define NBNXM_CUDA_KERNEL_UTILS_CUH
  61
  62 /*! \brief Log of the i and j cluster size.
  63  *  change this together with c_clSize !*/
  64 static const int __device__ c_clSizeLog2 = 3;
  65 /*! \brief Square of cluster size. */
  66 static const int __device__ c_clSizeSq = c_clSize * c_clSize;
  67 /*! \brief j-cluster size after split (4 in the current implementation). */
  68 static const int __device__ c_splitClSize = c_clSize / c_nbnxnGpuClusterpairSplit;
  69 /*! \brief Stride in the force accumualation buffer */
  70 static const int __device__ c_fbufStride = c_clSizeSq;
  71 /*! \brief i-cluster interaction mask for a super-cluster with all c_nbnxnGpuNumClusterPerSupercluster=8 bits set */
  72 static const unsigned __device__ superClInteractionMask =
  73         ((1U << c_nbnxnGpuNumClusterPerSupercluster) - 1U);
  74
  75 static const float __device__ c_oneSixth    = 0.16666667F;
  76 static const float __device__ c_oneTwelveth = 0.08333333F;
  77
  78
  79 /*! Convert LJ sigma,epsilon parameters to C6,C12. */
  80 static __forceinline__ __device__ void
  81 convert_sigma_epsilon_to_c6_c12(const float sigma, const float epsilon, float* c6, float* c12)
  82 {
  83     float sigma2, sigma6;
  84
  85     sigma2 = sigma * sigma;
  86     sigma6 = sigma2 * sigma2 * sigma2;
  87     *c6    = epsilon * sigma6;
  88     *c12   = *c6 * sigma6;
  89 }
  90
  91 /*! Apply force switch,  force + energy version. */
  92 static __forceinline__ __device__ void
  93 calculate_force_switch_F(const NBParamGpu nbparam, float c6, float c12, float inv_r, float r2, float* F_invr)
  94 {
  95     float r, r_switch;
  96
  97     /* force switch constants */
  98     float disp_shift_V2 = nbparam.dispersion_shift.c2;
  99     float disp_shift_V3 = nbparam.dispersion_shift.c3;
 100     float repu_shift_V2 = nbparam.repulsion_shift.c2;
 101     float repu_shift_V3 = nbparam.repulsion_shift.c3;
 102
 103     r        = r2 * inv_r;
 104     r_switch = r - nbparam.rvdw_switch;
 105     r_switch = r_switch >= 0.0F ? r_switch : 0.0F;
 106
 107     *F_invr += -c6 * (disp_shift_V2 + disp_shift_V3 * r_switch) * r_switch * r_switch * inv_r
 108                + c12 * (repu_shift_V2 + repu_shift_V3 * r_switch) * r_switch * r_switch * inv_r;
 109 }
 110
 111 /*! Apply force switch, force-only version. */
 112 static __forceinline__ __device__ void calculate_force_switch_F_E(const NBParamGpu nbparam,
 113                                                                   float            c6,
 114                                                                   float            c12,
 115                                                                   float            inv_r,
 116                                                                   float            r2,
 117                                                                   float*           F_invr,
 118                                                                   float*           E_lj)
 119 {
 120     float r, r_switch;
 121
 122     /* force switch constants */
 123     float disp_shift_V2 = nbparam.dispersion_shift.c2;
 124     float disp_shift_V3 = nbparam.dispersion_shift.c3;
 125     float repu_shift_V2 = nbparam.repulsion_shift.c2;
 126     float repu_shift_V3 = nbparam.repulsion_shift.c3;
 127
 128     float disp_shift_F2 = nbparam.dispersion_shift.c2 / 3;
 129     float disp_shift_F3 = nbparam.dispersion_shift.c3 / 4;
 130     float repu_shift_F2 = nbparam.repulsion_shift.c2 / 3;
 131     float repu_shift_F3 = nbparam.repulsion_shift.c3 / 4;
 132
 133     r        = r2 * inv_r;
 134     r_switch = r - nbparam.rvdw_switch;
 135     r_switch = r_switch >= 0.0F ? r_switch : 0.0F;
 136
 137     *F_invr += -c6 * (disp_shift_V2 + disp_shift_V3 * r_switch) * r_switch * r_switch * inv_r
 138                + c12 * (repu_shift_V2 + repu_shift_V3 * r_switch) * r_switch * r_switch * inv_r;
 139     *E_lj += c6 * (disp_shift_F2 + disp_shift_F3 * r_switch) * r_switch * r_switch * r_switch
 140              - c12 * (repu_shift_F2 + repu_shift_F3 * r_switch) * r_switch * r_switch * r_switch;
 141 }
 142
 143 /*! Apply potential switch, force-only version. */
 144 static __forceinline__ __device__ void calculate_potential_switch_F(const NBParamGpu& nbparam,
 145                                                                     float             inv_r,
 146                                                                     float             r2,
 147                                                                     float*            F_invr,
 148                                                                     const float*      E_lj)
 149 {
 150     float r, r_switch;
 151     float sw, dsw;
 152
 153     /* potential switch constants */
 154     float switch_V3 = nbparam.vdw_switch.c3;
 155     float switch_V4 = nbparam.vdw_switch.c4;
 156     float switch_V5 = nbparam.vdw_switch.c5;
 157     float switch_F2 = 3 * nbparam.vdw_switch.c3;
 158     float switch_F3 = 4 * nbparam.vdw_switch.c4;
 159     float switch_F4 = 5 * nbparam.vdw_switch.c5;
 160
 161     r        = r2 * inv_r;
 162     r_switch = r - nbparam.rvdw_switch;
 163
 164     /* Unlike in the F+E kernel, conditional is faster here */
 165     if (r_switch > 0.0F)
 166     {
 167         sw  = 1.0F + (switch_V3 + (switch_V4 + switch_V5 * r_switch) * r_switch) * r_switch * r_switch * r_switch;
 168         dsw = (switch_F2 + (switch_F3 + switch_F4 * r_switch) * r_switch) * r_switch * r_switch;
 169
 170         *F_invr = (*F_invr) * sw - inv_r * (*E_lj) * dsw;
 171     }
 172 }
 173
 174 /*! Apply potential switch, force + energy version. */
 175 static __forceinline__ __device__ void
 176 calculate_potential_switch_F_E(const NBParamGpu nbparam, float inv_r, float r2, float* F_invr, float* E_lj)
 177 {
 178     float r, r_switch;
 179     float sw, dsw;
 180
 181     /* potential switch constants */
 182     float switch_V3 = nbparam.vdw_switch.c3;
 183     float switch_V4 = nbparam.vdw_switch.c4;
 184     float switch_V5 = nbparam.vdw_switch.c5;
 185     float switch_F2 = 3 * nbparam.vdw_switch.c3;
 186     float switch_F3 = 4 * nbparam.vdw_switch.c4;
 187     float switch_F4 = 5 * nbparam.vdw_switch.c5;
 188
 189     r        = r2 * inv_r;
 190     r_switch = r - nbparam.rvdw_switch;
 191     r_switch = r_switch >= 0.0F ? r_switch : 0.0F;
 192
 193     /* Unlike in the F-only kernel, masking is faster here */
 194     sw  = 1.0F + (switch_V3 + (switch_V4 + switch_V5 * r_switch) * r_switch) * r_switch * r_switch * r_switch;
 195     dsw = (switch_F2 + (switch_F3 + switch_F4 * r_switch) * r_switch) * r_switch * r_switch;
 196
 197     *F_invr = (*F_invr) * sw - inv_r * (*E_lj) * dsw;
 198     *E_lj *= sw;
 199 }
 200
 201
 202 /*! \brief Fetch C6 grid contribution coefficients and return the product of these.
 203  *
 204  *  Depending on what is supported, it fetches parameters either
 205  *  using direct load, texture objects, or texrefs.
 206  */
 207 static __forceinline__ __device__ float calculate_lj_ewald_c6grid(const NBParamGpu nbparam, int typei, int typej)
 208 {
 209 #    if DISABLE_CUDA_TEXTURES
 210     float c6_i = LDG(&nbparam.nbfp_comb[typei]).x;
 211     float c6_j = LDG(&nbparam.nbfp_comb[typej]).x;
 212 #    else
 213     float c6_i = tex1Dfetch<float2>(nbparam.nbfp_comb_texobj, typei).x;
 214     float c6_j = tex1Dfetch<float2>(nbparam.nbfp_comb_texobj, typej).x;
 215 #    endif /* DISABLE_CUDA_TEXTURES */
 216     return c6_i * c6_j;
 217 }
 218
 219
 220 /*! Calculate LJ-PME grid force contribution with
 221  *  geometric combination rule.
 222  */
 223 static __forceinline__ __device__ void calculate_lj_ewald_comb_geom_F(const NBParamGpu nbparam,
 224                                                                       int              typei,
 225                                                                       int              typej,
 226                                                                       float            r2,
 227                                                                       float            inv_r2,
 228                                                                       float            lje_coeff2,
 229                                                                       float            lje_coeff6_6,
 230                                                                       float*           F_invr)
 231 {
 232     float c6grid, inv_r6_nm, cr2, expmcr2, poly;
 233
 234     c6grid = calculate_lj_ewald_c6grid(nbparam, typei, typej);
 235
 236     /* Recalculate inv_r6 without exclusion mask */
 237     inv_r6_nm = inv_r2 * inv_r2 * inv_r2;
 238     cr2       = lje_coeff2 * r2;
 239     expmcr2   = expf(-cr2);
 240     poly      = 1.0F + cr2 + 0.5F * cr2 * cr2;
 241
 242     /* Subtract the grid force from the total LJ force */
 243     *F_invr += c6grid * (inv_r6_nm - expmcr2 * (inv_r6_nm * poly + lje_coeff6_6)) * inv_r2;
 244 }
 245
 246
 247 /*! Calculate LJ-PME grid force + energy contribution with
 248  *  geometric combination rule.
 249  */
 250 static __forceinline__ __device__ void calculate_lj_ewald_comb_geom_F_E(const NBParamGpu nbparam,
 251                                                                         int              typei,
 252                                                                         int              typej,
 253                                                                         float            r2,
 254                                                                         float            inv_r2,
 255                                                                         float            lje_coeff2,
 256                                                                         float  lje_coeff6_6,
 257                                                                         float  int_bit,
 258                                                                         float* F_invr,
 259                                                                         float* E_lj)
 260 {
 261     float c6grid, inv_r6_nm, cr2, expmcr2, poly, sh_mask;
 262
 263     c6grid = calculate_lj_ewald_c6grid(nbparam, typei, typej);
 264
 265     /* Recalculate inv_r6 without exclusion mask */
 266     inv_r6_nm = inv_r2 * inv_r2 * inv_r2;
 267     cr2       = lje_coeff2 * r2;
 268     expmcr2   = expf(-cr2);
 269     poly      = 1.0F + cr2 + 0.5F * cr2 * cr2;
 270
 271     /* Subtract the grid force from the total LJ force */
 272     *F_invr += c6grid * (inv_r6_nm - expmcr2 * (inv_r6_nm * poly + lje_coeff6_6)) * inv_r2;
 273
 274     /* Shift should be applied only to real LJ pairs */
 275     sh_mask = nbparam.sh_lj_ewald * int_bit;
 276     *E_lj += c_oneSixth * c6grid * (inv_r6_nm * (1.0F - expmcr2 * poly) + sh_mask);
 277 }
 278
 279 /*! Fetch per-type LJ parameters.
 280  *
 281  *  Depending on what is supported, it fetches parameters either
 282  *  using direct load, texture objects, or texrefs.
 283  */
 284 static __forceinline__ __device__ float2 fetch_nbfp_comb_c6_c12(const NBParamGpu nbparam, int type)
 285 {
 286 #    if DISABLE_CUDA_TEXTURES
 287     return LDG(&nbparam.nbfp_comb[type]);
 288 #    else
 289     return tex1Dfetch<float2>(nbparam.nbfp_comb_texobj, type);
 290 #    endif /* DISABLE_CUDA_TEXTURES */
 291 }
 292
 293
 294 /*! Calculate LJ-PME grid force + energy contribution (if E_lj != nullptr) with
 295  *  Lorentz-Berthelot combination rule.
 296  *  We use a single F+E kernel with conditional because the performance impact
 297  *  of this is pretty small and LB on the CPU is anyway very slow.
 298  */
 299 static __forceinline__ __device__ void calculate_lj_ewald_comb_LB_F_E(const NBParamGpu nbparam,
 300                                                                       int              typei,
 301                                                                       int              typej,
 302                                                                       float            r2,
 303                                                                       float            inv_r2,
 304                                                                       float            lje_coeff2,
 305                                                                       float            lje_coeff6_6,
 306                                                                       float            int_bit,
 307                                                                       float*           F_invr,
 308                                                                       float*           E_lj)
 309 {
 310     float c6grid, inv_r6_nm, cr2, expmcr2, poly;
 311     float sigma, sigma2, epsilon;
 312
 313     /* sigma and epsilon are scaled to give 6*C6 */
 314     float2 c6c12_i = fetch_nbfp_comb_c6_c12(nbparam, typei);
 315     float2 c6c12_j = fetch_nbfp_comb_c6_c12(nbparam, typej);
 316
 317     sigma   = c6c12_i.x + c6c12_j.x;
 318     epsilon = c6c12_i.y * c6c12_j.y;
 319
 320     sigma2 = sigma * sigma;
 321     c6grid = epsilon * sigma2 * sigma2 * sigma2;
 322
 323     /* Recalculate inv_r6 without exclusion mask */
 324     inv_r6_nm = inv_r2 * inv_r2 * inv_r2;
 325     cr2       = lje_coeff2 * r2;
 326     expmcr2   = expf(-cr2);
 327     poly      = 1.0F + cr2 + 0.5F * cr2 * cr2;
 328
 329     /* Subtract the grid force from the total LJ force */
 330     *F_invr += c6grid * (inv_r6_nm - expmcr2 * (inv_r6_nm * poly + lje_coeff6_6)) * inv_r2;
 331
 332     if (E_lj != nullptr)
 333     {
 334         float sh_mask;
 335
 336         /* Shift should be applied only to real LJ pairs */
 337         sh_mask = nbparam.sh_lj_ewald * int_bit;
 338         *E_lj += c_oneSixth * c6grid * (inv_r6_nm * (1.0F - expmcr2 * poly) + sh_mask);
 339     }
 340 }
 341
 342
 343 /*! Fetch two consecutive values from the Ewald correction F*r table.
 344  *
 345  *  Depending on what is supported, it fetches parameters either
 346  *  using direct load, texture objects, or texrefs.
 347  */
 348 static __forceinline__ __device__ float2 fetch_coulomb_force_r(const NBParamGpu nbparam, int index)
 349 {
 350     float2 d;
 351
 352 #    if DISABLE_CUDA_TEXTURES
 353     /* Can't do 8-byte fetch because some of the addresses will be misaligned. */
 354     d.x = LDG(&nbparam.coulomb_tab[index]);
 355     d.y = LDG(&nbparam.coulomb_tab[index + 1]);
 356 #    else
 357     d.x   = tex1Dfetch<float>(nbparam.coulomb_tab_texobj, index);
 358     d.y   = tex1Dfetch<float>(nbparam.coulomb_tab_texobj, index + 1);
 359 #    endif // DISABLE_CUDA_TEXTURES
 360
 361     return d;
 362 }
 363
 364 /*! Linear interpolation using exactly two FMA operations.
 365  *
 366  *  Implements numeric equivalent of: (1-t)*d0 + t*d1
 367  *  Note that CUDA does not have fnms, otherwise we'd use
 368  *  fma(t, d1, fnms(t, d0, d0)
 369  *  but input modifiers are designed for this and are fast.
 370  */
 371 template<typename T>
 372 __forceinline__ __host__ __device__ T lerp(T d0, T d1, T t)
 373 {
 374     return fma(t, d1, fma(-t, d0, d0));
 375 }
 376
 377 /*! Interpolate Ewald coulomb force correction using the F*r table.
 378  */
 379 static __forceinline__ __device__ float interpolate_coulomb_force_r(const NBParamGpu nbparam, float r)
 380 {
 381     float normalized = nbparam.coulomb_tab_scale * r;
 382     int   index      = static_cast<int>(normalized);
 383     float fraction   = normalized - index;
 384
 385     float2 d01 = fetch_coulomb_force_r(nbparam, index);
 386
 387     return lerp(d01.x, d01.y, fraction);
 388 }
 389
 390 /*! Fetch C6 and C12 from the parameter table.
 391  *
 392  *  Depending on what is supported, it fetches parameters either
 393  *  using direct load, texture objects, or texrefs.
 394  */
 395 // NOLINTNEXTLINE(google-runtime-references)
 396 static __forceinline__ __device__ void fetch_nbfp_c6_c12(float& c6, float& c12, const NBParamGpu nbparam, int baseIndex)
 397 {
 398     float2 c6c12;
 399 #    if DISABLE_CUDA_TEXTURES
 400     c6c12 = LDG(&nbparam.nbfp[baseIndex]);
 401 #    else
 402     c6c12 = tex1Dfetch<float2>(nbparam.nbfp_texobj, baseIndex);
 403 #    endif // DISABLE_CUDA_TEXTURES
 404     c6  = c6c12.x;
 405     c12 = c6c12.y;
 406 }
 407
 408
 409 /*! Calculate analytical Ewald correction term. */
 410 static __forceinline__ __device__ float pmecorrF(float z2)
 411 {
 412     const float FN6 = -1.7357322914161492954e-8F;
 413     const float FN5 = 1.4703624142580877519e-6F;
 414     const float FN4 = -0.000053401640219807709149F;
 415     const float FN3 = 0.0010054721316683106153F;
 416     const float FN2 = -0.019278317264888380590F;
 417     const float FN1 = 0.069670166153766424023F;
 418     const float FN0 = -0.75225204789749321333F;
 419
 420     const float FD4 = 0.0011193462567257629232F;
 421     const float FD3 = 0.014866955030185295499F;
 422     const float FD2 = 0.11583842382862377919F;
 423     const float FD1 = 0.50736591960530292870F;
 424     const float FD0 = 1.0F;
 425
 426     float z4;
 427     float polyFN0, polyFN1, polyFD0, polyFD1;
 428
 429     z4 = z2 * z2;
 430
 431     polyFD0 = FD4 * z4 + FD2;
 432     polyFD1 = FD3 * z4 + FD1;
 433     polyFD0 = polyFD0 * z4 + FD0;
 434     polyFD0 = polyFD1 * z2 + polyFD0;
 435
 436     polyFD0 = 1.0F / polyFD0;
 437
 438     polyFN0 = FN6 * z4 + FN4;
 439     polyFN1 = FN5 * z4 + FN3;
 440     polyFN0 = polyFN0 * z4 + FN2;
 441     polyFN1 = polyFN1 * z4 + FN1;
 442     polyFN0 = polyFN0 * z4 + FN0;
 443     polyFN0 = polyFN1 * z2 + polyFN0;
 444
 445     return polyFN0 * polyFD0;
 446 }
 447
 448 /*! Final j-force reduction; this generic implementation works with
 449  *  arbitrary array sizes.
 450  */
 451 static __forceinline__ __device__ void
 452 reduce_force_j_generic(const float* f_buf, float3* fout, int tidxi, int tidxj, int aidx)
 453 {
 454     if (tidxi < 3)
 455     {
 456         float f = 0.0F;
 457         for (int j = tidxj * c_clSize; j < (tidxj + 1) * c_clSize; j++)
 458         {
 459             f += f_buf[c_fbufStride * tidxi + j];
 460         }
 461
 462         atomicAdd((&fout[aidx].x) + tidxi, f);
 463     }
 464 }
 465
 466 /*! Final j-force reduction; this implementation only with power of two
 467  *  array sizes.
 468  */
 469 static __forceinline__ __device__ void
 470 reduce_force_j_warp_shfl(float3 f, float3* fout, int tidxi, int aidx, const unsigned int activemask)
 471 {
 472     f.x += __shfl_down_sync(activemask, f.x, 1);
 473     f.y += __shfl_up_sync(activemask, f.y, 1);
 474     f.z += __shfl_down_sync(activemask, f.z, 1);
 475
 476     if (tidxi & 1)
 477     {
 478         f.x = f.y;
 479     }
 480
 481     f.x += __shfl_down_sync(activemask, f.x, 2);
 482     f.z += __shfl_up_sync(activemask, f.z, 2);
 483
 484     if (tidxi & 2)
 485     {
 486         f.x = f.z;
 487     }
 488
 489     f.x += __shfl_down_sync(activemask, f.x, 4);
 490
 491     if (tidxi < 3)
 492     {
 493         atomicAdd((&fout[aidx].x) + tidxi, f.x);
 494     }
 495 }
 496
 497 /*! Final i-force reduction; this generic implementation works with
 498  *  arbitrary array sizes.
 499  * TODO: add the tidxi < 3 trick
 500  */
 501 static __forceinline__ __device__ void reduce_force_i_generic(const float* f_buf,
 502                                                               float3*      fout,
 503                                                               float*       fshift_buf,
 504                                                               bool         bCalcFshift,
 505                                                               int          tidxi,
 506                                                               int          tidxj,
 507                                                               int          aidx)
 508 {
 509     if (tidxj < 3)
 510     {
 511         float f = 0.0F;
 512         for (int j = tidxi; j < c_clSizeSq; j += c_clSize)
 513         {
 514             f += f_buf[tidxj * c_fbufStride + j];
 515         }
 516
 517         atomicAdd(&fout[aidx].x + tidxj, f);
 518
 519         if (bCalcFshift)
 520         {
 521             *fshift_buf += f;
 522         }
 523     }
 524 }
 525
 526 /*! Final i-force reduction; this implementation works only with power of two
 527  *  array sizes.
 528  */
 529 static __forceinline__ __device__ void reduce_force_i_pow2(volatile float* f_buf,
 530                                                            float3*         fout,
 531                                                            float*          fshift_buf,
 532                                                            bool            bCalcFshift,
 533                                                            int             tidxi,
 534                                                            int             tidxj,
 535                                                            int             aidx)
 536 {
 537     int   i, j;
 538     float f;
 539
 540     static_assert(c_clSize == 1 << c_clSizeLog2);
 541
 542     /* Reduce the initial c_clSize values for each i atom to half
 543      * every step by using c_clSize * i threads.
 544      * Can't just use i as loop variable because than nvcc refuses to unroll.
 545      */
 546     i = c_clSize / 2;
 547 #    pragma unroll 5
 548     for (j = c_clSizeLog2 - 1; j > 0; j--)
 549     {
 550         if (tidxj < i)
 551         {
 552
 553             f_buf[tidxj * c_clSize + tidxi] += f_buf[(tidxj + i) * c_clSize + tidxi];
 554             f_buf[c_fbufStride + tidxj * c_clSize + tidxi] +=
 555                     f_buf[c_fbufStride + (tidxj + i) * c_clSize + tidxi];
 556             f_buf[2 * c_fbufStride + tidxj * c_clSize + tidxi] +=
 557                     f_buf[2 * c_fbufStride + (tidxj + i) * c_clSize + tidxi];
 558         }
 559         i >>= 1;
 560     }
 561
 562     /* i == 1, last reduction step, writing to global mem */
 563     if (tidxj < 3)
 564     {
 565         /* tidxj*c_fbufStride selects x, y or z */
 566         f = f_buf[tidxj * c_fbufStride + tidxi] + f_buf[tidxj * c_fbufStride + i * c_clSize + tidxi];
 567
 568         atomicAdd(&(fout[aidx].x) + tidxj, f);
 569
 570         if (bCalcFshift)
 571         {
 572             *fshift_buf += f;
 573         }
 574     }
 575 }
 576
 577 /*! Final i-force reduction wrapper; calls the generic or pow2 reduction depending
 578  *  on whether the size of the array to be reduced is power of two or not.
 579  */
 580 static __forceinline__ __device__ void
 581 reduce_force_i(float* f_buf, float3* f, float* fshift_buf, bool bCalcFshift, int tidxi, int tidxj, int ai)
 582 {
 583     if ((c_clSize & (c_clSize - 1)))
 584     {
 585         reduce_force_i_generic(f_buf, f, fshift_buf, bCalcFshift, tidxi, tidxj, ai);
 586     }
 587     else
 588     {
 589         reduce_force_i_pow2(f_buf, f, fshift_buf, bCalcFshift, tidxi, tidxj, ai);
 590     }
 591 }
 592
 593 /*! Final i-force reduction; this implementation works only with power of two
 594  *  array sizes.
 595  */
 596 static __forceinline__ __device__ void reduce_force_i_warp_shfl(float3             fin,
 597                                                                 float3*            fout,
 598                                                                 float*             fshift_buf,
 599                                                                 bool               bCalcFshift,
 600                                                                 int                tidxj,
 601                                                                 int                aidx,
 602                                                                 const unsigned int activemask)
 603 {
 604     fin.x += __shfl_down_sync(activemask, fin.x, c_clSize);
 605     fin.y += __shfl_up_sync(activemask, fin.y, c_clSize);
 606     fin.z += __shfl_down_sync(activemask, fin.z, c_clSize);
 607
 608     if (tidxj & 1)
 609     {
 610         fin.x = fin.y;
 611     }
 612
 613     fin.x += __shfl_down_sync(activemask, fin.x, 2 * c_clSize);
 614     fin.z += __shfl_up_sync(activemask, fin.z, 2 * c_clSize);
 615
 616     if (tidxj & 2)
 617     {
 618         fin.x = fin.z;
 619     }
 620
 621     /* Threads 0,1,2 and 4,5,6 increment x,y,z for their warp */
 622     if ((tidxj & 3) < 3)
 623     {
 624         atomicAdd(&fout[aidx].x + (tidxj & 3), fin.x);
 625
 626         if (bCalcFshift)
 627         {
 628             *fshift_buf += fin.x;
 629         }
 630     }
 631 }
 632
 633 /*! Energy reduction; this implementation works only with power of two
 634  *  array sizes.
 635  */
 636 static __forceinline__ __device__ void
 637 reduce_energy_pow2(volatile float* buf, float* e_lj, float* e_el, unsigned int tidx)
 638 {
 639     float e1, e2;
 640
 641     unsigned int i = warp_size / 2;
 642
 643     /* Can't just use i as loop variable because than nvcc refuses to unroll. */
 644 #    pragma unroll 10
 645     for (int j = warp_size_log2 - 1; j > 0; j--)
 646     {
 647         if (tidx < i)
 648         {
 649             buf[tidx] += buf[tidx + i];
 650             buf[c_fbufStride + tidx] += buf[c_fbufStride + tidx + i];
 651         }
 652         i >>= 1;
 653     }
 654
 655     // TODO do this on two threads - requires e_lj and e_el to be stored on adjascent
 656     // memory locations to make sense
 657     /* last reduction step, writing to global mem */
 658     if (tidx == 0)
 659     {
 660         e1 = buf[tidx] + buf[tidx + i];
 661         e2 = buf[c_fbufStride + tidx] + buf[c_fbufStride + tidx + i];
 662
 663         atomicAdd(e_lj, e1);
 664         atomicAdd(e_el, e2);
 665     }
 666 }
 667
 668 /*! Energy reduction; this implementation works only with power of two
 669  *  array sizes.
 670  */
 671 static __forceinline__ __device__ void
 672 reduce_energy_warp_shfl(float E_lj, float E_el, float* e_lj, float* e_el, int tidx, const unsigned int activemask)
 673 {
 674     int i, sh;
 675
 676     sh = 1;
 677 #    pragma unroll 5
 678     for (i = 0; i < 5; i++)
 679     {
 680         E_lj += __shfl_down_sync(activemask, E_lj, sh);
 681         E_el += __shfl_down_sync(activemask, E_el, sh);
 682         sh += sh;
 683     }
 684
 685     /* The first thread in the warp writes the reduced energies */
 686     if (tidx == 0 || tidx == warp_size)
 687     {
 688         atomicAdd(e_lj, E_lj);
 689         atomicAdd(e_el, E_el);
 690     }
 691 }
 692
 693 #endif /* NBNXN_CUDA_KERNEL_UTILS_CUH */