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