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