src/gromacs/nbnxm/cuda/nbnxm_cuda_kernel_pruneonly.cuh

   1 /*
   2  * This file is part of the GROMACS molecular simulation package.
   3  *
   4  * Copyright (c) 2016,2017,2018,2019,2020,2021, 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  *  CUDA non-bonded prune-only kernel.
  39  *
  40  *  Unlike the non-bonded interaction kernels, this is not preprocessor-generated,
  41  *  the two flavors achieved by templating.
  42  *
  43  *  \author Szilárd Páll <pall.szilard@gmail.com>
  44  *  \author Berk Hess <hess@kth.se>
  45  *  \ingroup module_nbnxm
  46  */
  47 #include "gmxpre.h"
  48
  49 #include "gromacs/gpu_utils/cuda_arch_utils.cuh"
  50 #include "gromacs/gpu_utils/typecasts.cuh"
  51 #include "gromacs/math/utilities.h"
  52 #include "gromacs/pbcutil/ishift.h"
  53
  54 #include "nbnxm_cuda_kernel_utils.cuh"
  55 #include "nbnxm_cuda_types.h"
  56
  57 /* Note that floating-point constants in CUDA code should be suffixed
  58  * with f (e.g. 0.5f), to stop the compiler producing intermediate
  59  * code that is in double precision.
  60  */
  61
  62 /**@{*/
  63 /*! \brief Compute capability dependent definition of kernel launch configuration parameters.
  64  *
  65  * Kernel launch bounds for different compute capabilities. The value of NTHREAD_Z
  66  * represents the j-concurrency, hence it determines the number of threads per block.
  67  * It is chosen such that 100% occupancy is maintained (on Maxwell and later for any NTHREAD_Z,
  68  * requires >=4 warp/block, NTHREAD_Z>=2 on Kepler).
  69  *
  70  * Hence, values NTHREAD_Z >= 2 trade inter- for intra-block parallelism
  71  * which has the advantage of lowering the overhead of starting up a block, filling shmem
  72  * and registers, etc. Ideally we'd want to expose as much intra-block work as possible
  73  * As we also split lists to cater for the block-parallelization needed by the register-
  74  * limited non-bonded kernels, for very short j-loops large NTHREAD_Z will cause slowdown
  75  * as it leads to intra-block warp imbalance. Ideally, we'd want to auto-tune the choice
  76  * of NTHREAD_Z, but for now we instead pick a reasonable tradeoff-value.
  77  *
  78  * Note that given the above input size tradeoffs and that performance depends on
  79  * additional factors including GPU arch, #SM's, we'll accept performance tradeoffs
  80  * of using a fixed NTHREAD_Z=4. The following outliers have been observed:
  81  *   - up to 25% faster (rolling) prune kernels with NTHREAD_Z=8 in the regime where lists
  82  *     are not split (much), but the rolling chunks are small;
  83  *   - with large inputs NTHREAD_Z=1 is 2-3% faster (on CC>=5.0)
  84  */
  85 #define NTHREAD_Z (GMX_NBNXN_PRUNE_KERNEL_J4_CONCURRENCY)
  86 #define THREADS_PER_BLOCK (c_clSize * c_clSize * NTHREAD_Z)
  87 // we want 100% occupancy, so max threads/block
  88 #define MIN_BLOCKS_PER_MP (GMX_CUDA_MAX_THREADS_PER_MP / THREADS_PER_BLOCK)
  89 /**@}*/
  90
  91 /*! \brief Nonbonded list pruning kernel.
  92  *
  93  *  The \p haveFreshList template parameter defines the two flavors of the kernel; when
  94  *  true a new list from immediately after pair-list generation is pruned using rlistOuter,
  95  *  the pruned masks are stored in a separate buffer and the outer-list is pruned
  96  *  using the rlistInner distance; when false only the pruning with rlistInner is performed.
  97  *
  98  *  Kernel launch parameters:
  99  *   - #blocks   = #pair lists, blockId = pair list Id
 100  *   - #threads  = NTHREAD_Z * c_clSize^2
 101  *   - shmem     = see nbnxn_cuda.cu:calc_shmem_required_prune()
 102  *
 103  *   Each thread calculates an i-j atom distance..
 104  */
 105 template<bool haveFreshList>
 106 __launch_bounds__(THREADS_PER_BLOCK, MIN_BLOCKS_PER_MP) __global__
 107         void nbnxn_kernel_prune_cuda(NBAtomDataGpu    atdat,
 108                                      NBParamGpu       nbparam,
 109                                      Nbnxm::gpu_plist plist,
 110                                      int              numParts,
 111                                      int              part)
 112 #ifdef FUNCTION_DECLARATION_ONLY
 113                 ; /* Only do function declaration, omit the function body. */
 114
 115 // Add extern declarations so each translation unit understands that
 116 // there will be a definition provided.
 117 extern template __global__ void
 118 nbnxn_kernel_prune_cuda<true>(const NBAtomDataGpu, const NBParamGpu, const Nbnxm::gpu_plist, int, int);
 119 extern template __global__ void
 120 nbnxn_kernel_prune_cuda<false>(const NBAtomDataGpu, const NBParamGpu, const Nbnxm::gpu_plist, int, int);
 121 #else
 122 {
 123
 124     /* convenience variables */
 125     const nbnxn_sci_t* pl_sci    = plist.sci;
 126     nbnxn_cj4_t*       pl_cj4    = plist.cj4;
 127     const float4*      xq        = atdat.xq;
 128     const float3*      shift_vec = asFloat3(atdat.shiftVec);
 129
 130     float rlistOuter_sq = nbparam.rlistOuter_sq;
 131     float rlistInner_sq = nbparam.rlistInner_sq;
 132
 133     /* thread/block/warp id-s */
 134     unsigned int tidxi = threadIdx.x;
 135     unsigned int tidxj = threadIdx.y;
 136 #    if NTHREAD_Z == 1
 137     unsigned int tidxz = 0;
 138 #    else
 139     unsigned int tidxz = threadIdx.z;
 140 #    endif
 141     unsigned int bidx  = blockIdx.x;
 142     unsigned int widx  = (threadIdx.y * c_clSize) / warp_size; /* warp index */
 143
 144     // cj preload is off in the following cases:
 145     // - sm_70 (V100), sm_8x (A100, GA100), sm_75 (TU102)
 146     // - for future arch (> 8.6 at the time of writing) we assume it is better to keep it off
 147     constexpr bool c_preloadCj = (GMX_PTX_ARCH < 700);
 148
 149     /*********************************************************************
 150      * Set up shared memory pointers.
 151      * sm_nextSlotPtr should always be updated to point to the "next slot",
 152      * that is past the last point where data has been stored.
 153      */
 154     extern __shared__ char sm_dynamicShmem[];
 155     char*                  sm_nextSlotPtr = sm_dynamicShmem;
 156     static_assert(sizeof(char) == 1,
 157                   "The shared memory offset calculation assumes that char is 1 byte");
 158
 159     /* shmem buffer for i x+q pre-loading */
 160     float4* xib = reinterpret_cast<float4*>(sm_nextSlotPtr);
 161     sm_nextSlotPtr += (c_nbnxnGpuNumClusterPerSupercluster * c_clSize * sizeof(*xib));
 162
 163     /* shmem buffer for cj, for each warp separately */
 164     int* cjs = reinterpret_cast<int*>(sm_nextSlotPtr);
 165     if (c_preloadCj)
 166     {
 167         /* the cjs buffer's use expects a base pointer offset for pairs of warps in the j-concurrent execution */
 168         cjs += tidxz * c_nbnxnGpuClusterpairSplit * c_nbnxnGpuJgroupSize;
 169         sm_nextSlotPtr += (NTHREAD_Z * c_nbnxnGpuClusterpairSplit * c_nbnxnGpuJgroupSize * sizeof(*cjs));
 170     }
 171     /*********************************************************************/
 172
 173
 174     nbnxn_sci_t nb_sci =
 175             pl_sci[bidx * numParts + part]; /* my i super-cluster's index = sciOffset + current bidx * numParts + part */
 176     int sci        = nb_sci.sci;           /* super-cluster */
 177     int cij4_start = nb_sci.cj4_ind_start; /* first ...*/
 178     int cij4_end   = nb_sci.cj4_ind_end;   /* and last index of j clusters */
 179
 180     if (tidxz == 0)
 181     {
 182         /* Pre-load i-atom x and q into shared memory */
 183         int ci = sci * c_nbnxnGpuNumClusterPerSupercluster + tidxj;
 184         int ai = ci * c_clSize + tidxi;
 185
 186         /* We don't need q, but using float4 in shmem avoids bank conflicts.
 187            (but it also wastes L2 bandwidth). */
 188         float4 tmp                    = xq[ai];
 189         float4 xi                     = tmp + shift_vec[nb_sci.shift];
 190         xib[tidxj * c_clSize + tidxi] = xi;
 191     }
 192     __syncthreads();
 193
 194     /* loop over the j clusters = seen by any of the atoms in the current super-cluster;
 195      * The loop stride NTHREAD_Z ensures that consecutive warps-pairs are assigned
 196      * consecutive j4's entries.
 197      */
 198     for (int j4 = cij4_start + tidxz; j4 < cij4_end; j4 += NTHREAD_Z)
 199     {
 200         unsigned int imaskFull, imaskCheck, imaskNew;
 201
 202         if (haveFreshList)
 203         {
 204             /* Read the mask from the list transferred from the CPU */
 205             imaskFull = pl_cj4[j4].imei[widx].imask;
 206             /* We attempt to prune all pairs present in the original list */
 207             imaskCheck = imaskFull;
 208             imaskNew   = 0;
 209         }
 210         else
 211         {
 212             /* Read the mask from the "warp-pruned" by rlistOuter mask array */
 213             imaskFull = plist.imask[j4 * c_nbnxnGpuClusterpairSplit + widx];
 214             /* Read the old rolling pruned mask, use as a base for new */
 215             imaskNew = pl_cj4[j4].imei[widx].imask;
 216             /* We only need to check pairs with different mask */
 217             imaskCheck = (imaskNew ^ imaskFull);
 218         }
 219
 220         if (imaskCheck)
 221         {
 222             if (c_preloadCj)
 223             {
 224                 /* Pre-load cj into shared memory on both warps separately */
 225                 if ((tidxj == 0 || tidxj == 4) && tidxi < c_nbnxnGpuJgroupSize)
 226                 {
 227                     cjs[tidxi + tidxj * c_nbnxnGpuJgroupSize / c_splitClSize] = pl_cj4[j4].cj[tidxi];
 228                 }
 229                 __syncwarp(c_fullWarpMask);
 230             }
 231
 232 #    pragma unroll 4
 233             for (int jm = 0; jm < c_nbnxnGpuJgroupSize; jm++)
 234             {
 235                 if (imaskCheck & (superClInteractionMask << (jm * c_nbnxnGpuNumClusterPerSupercluster)))
 236                 {
 237                     unsigned int mask_ji = (1U << (jm * c_nbnxnGpuNumClusterPerSupercluster));
 238                     int cj = c_preloadCj ? cjs[jm + (tidxj & 4) * c_nbnxnGpuJgroupSize / c_splitClSize]
 239                                          : pl_cj4[j4].cj[jm];
 240                     int aj = cj * c_clSize + tidxj;
 241
 242                     /* load j atom data */
 243                     float4 tmp = xq[aj];
 244                     float3 xj  = make_float3(tmp.x, tmp.y, tmp.z);
 245
 246 #    pragma unroll 8
 247                     for (int i = 0; i < c_nbnxnGpuNumClusterPerSupercluster; i++)
 248                     {
 249                         if (imaskCheck & mask_ji)
 250                         {
 251                             /* load i-cluster coordinates from shmem */
 252                             float4 xi = xib[i * c_clSize + tidxi];
 253
 254
 255                             /* distance between i and j atoms */
 256                             float3 rv = make_float3(xi.x, xi.y, xi.z) - xj;
 257                             float  r2 = norm2(rv);
 258
 259                             /* If _none_ of the atoms pairs are in rlistOuter
 260                                range, the bit corresponding to the current
 261                                cluster-pair in imask gets set to 0. */
 262                             if (haveFreshList && !__any_sync(c_fullWarpMask, r2 < rlistOuter_sq))
 263                             {
 264                                 imaskFull &= ~mask_ji;
 265                             }
 266                             /* If any atom pair is within range, set the bit
 267                                corresponding to the current cluster-pair. */
 268                             if (__any_sync(c_fullWarpMask, r2 < rlistInner_sq))
 269                             {
 270                                 imaskNew |= mask_ji;
 271                             }
 272                         }
 273
 274                         /* shift the mask bit by 1 */
 275                         mask_ji += mask_ji;
 276                     }
 277                 }
 278             }
 279
 280             if (haveFreshList)
 281             {
 282                 /* copy the list pruned to rlistOuter to a separate buffer */
 283                 plist.imask[j4 * c_nbnxnGpuClusterpairSplit + widx] = imaskFull;
 284             }
 285             /* update the imask with only the pairs up to rlistInner */
 286             plist.cj4[j4].imei[widx].imask = imaskNew;
 287         }
 288         if (c_preloadCj)
 289         {
 290             // avoid shared memory WAR hazards on sm_cjs between loop iterations
 291             __syncwarp(c_fullWarpMask);
 292         }
 293     }
 294 }
 295 #endif /* FUNCTION_DECLARATION_ONLY */
 296
 297 #undef NTHREAD_Z
 298 #undef MIN_BLOCKS_PER_MP
 299 #undef THREADS_PER_BLOCK