src/gromacs/ewald/pme-solve.cu

   1 /*
   2  * This file is part of the GROMACS molecular simulation package.
   3  *
   4  * Copyright (c) 2016,2017,2018, 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 Implements PME GPU Fourier grid solving in CUDA.
  38  *
  39  *  \author Aleksei Iupinov <a.yupinov@gmail.com>
  40  */
  41
  42 #include "gmxpre.h"
  43
  44 #include "gromacs/gpu_utils/cuda_arch_utils.cuh"
  45 #include "gromacs/gpu_utils/cudautils.cuh"
  46 #include "gromacs/gpu_utils/devicebuffer.h"
  47 #include "gromacs/utility/exceptions.h"
  48 #include "gromacs/utility/gmxassert.h"
  49
  50 #include "pme.cuh"
  51 #include "pme-timings.cuh"
  52
  53 //! Solving kernel max block width in warps picked among powers of 2 (2, 4, 8, 16) for max. occupancy and min. runtime
  54 //! (560Ti (CC2.1), 660Ti (CC3.0) and 750 (CC5.0)))
  55 constexpr int c_solveMaxWarpsPerBlock = 8;
  56 //! Solving kernel max block size in threads
  57 constexpr int c_solveMaxThreadsPerBlock = (c_solveMaxWarpsPerBlock * warp_size);
  58
  59 /*! \brief
  60  * PME complex grid solver kernel function.
  61  *
  62  * \tparam[in] gridOrdering             Specifies the dimension ordering of the complex grid.
  63  * \tparam[in] computeEnergyAndVirial   Tells if the reciprocal energy and virial should be computed.
  64  * \param[in]  kernelParams             Input PME CUDA data in constant memory.
  65  */
  66 template<
  67     GridOrdering gridOrdering,
  68     bool computeEnergyAndVirial
  69     >
  70 __launch_bounds__(c_solveMaxThreadsPerBlock)
  71 __global__ void pme_solve_kernel(const struct PmeGpuCudaKernelParams kernelParams)
  72 {
  73     /* This kernel supports 2 different grid dimension orderings: YZX and XYZ */
  74     int majorDim, middleDim, minorDim;
  75     switch (gridOrdering)
  76     {
  77         case GridOrdering::YZX:
  78             majorDim  = YY;
  79             middleDim = ZZ;
  80             minorDim  = XX;
  81             break;
  82
  83         case GridOrdering::XYZ:
  84             majorDim  = XX;
  85             middleDim = YY;
  86             minorDim  = ZZ;
  87             break;
  88
  89         default:
  90             assert(false);
  91     }
  92
  93     /* Global memory pointers */
  94     const float * __restrict__ gm_splineValueMajor    = kernelParams.grid.d_splineModuli + kernelParams.grid.splineValuesOffset[majorDim];
  95     const float * __restrict__ gm_splineValueMiddle   = kernelParams.grid.d_splineModuli + kernelParams.grid.splineValuesOffset[middleDim];
  96     const float * __restrict__ gm_splineValueMinor    = kernelParams.grid.d_splineModuli + kernelParams.grid.splineValuesOffset[minorDim];
  97     float * __restrict__       gm_virialAndEnergy     = kernelParams.constants.d_virialAndEnergy;
  98     float2 * __restrict__      gm_grid                = (float2 *)kernelParams.grid.d_fourierGrid;
  99
 100     /* Various grid sizes and indices */
 101     const int localOffsetMinor = 0, localOffsetMajor = 0, localOffsetMiddle = 0; //unused
 102     const int localSizeMinor   = kernelParams.grid.complexGridSizePadded[minorDim];
 103     const int localSizeMiddle  = kernelParams.grid.complexGridSizePadded[middleDim];
 104     const int localCountMiddle = kernelParams.grid.complexGridSize[middleDim];
 105     const int localCountMinor  = kernelParams.grid.complexGridSize[minorDim];
 106     const int nMajor           = kernelParams.grid.realGridSize[majorDim];
 107     const int nMiddle          = kernelParams.grid.realGridSize[middleDim];
 108     const int nMinor           = kernelParams.grid.realGridSize[minorDim];
 109     const int maxkMajor        = (nMajor + 1) / 2;  // X or Y
 110     const int maxkMiddle       = (nMiddle + 1) / 2; // Y OR Z => only check for !YZX
 111     const int maxkMinor        = (nMinor + 1) / 2;  // Z or X => only check for YZX
 112
 113     /* Each thread works on one cell of the Fourier space complex 3D grid (gm_grid).
 114      * Each block handles up to c_solveMaxThreadsPerBlock cells -
 115      * depending on the grid contiguous dimension size,
 116      * that can range from a part of a single gridline to several complete gridlines.
 117      */
 118     const int threadLocalId     = threadIdx.x;
 119     const int gridLineSize      = localCountMinor;
 120     const int gridLineIndex     = threadLocalId / gridLineSize;
 121     const int gridLineCellIndex = threadLocalId - gridLineSize * gridLineIndex;
 122     const int gridLinesPerBlock = blockDim.x / gridLineSize;
 123     const int activeWarps       = (blockDim.x / warp_size);
 124     const int indexMinor        = blockIdx.x * blockDim.x + gridLineCellIndex;
 125     const int indexMiddle       = blockIdx.y * gridLinesPerBlock + gridLineIndex;
 126     const int indexMajor        = blockIdx.z;
 127
 128     /* Optional outputs */
 129     float energy = 0.0f;
 130     float virxx  = 0.0f;
 131     float virxy  = 0.0f;
 132     float virxz  = 0.0f;
 133     float viryy  = 0.0f;
 134     float viryz  = 0.0f;
 135     float virzz  = 0.0f;
 136
 137     assert(indexMajor < kernelParams.grid.complexGridSize[majorDim]);
 138     if ((indexMiddle < localCountMiddle) & (indexMinor < localCountMinor) & (gridLineIndex < gridLinesPerBlock))
 139     {
 140         /* The offset should be equal to the global thread index for coalesced access */
 141         const int             gridIndex     = (indexMajor * localSizeMiddle + indexMiddle) * localSizeMinor + indexMinor;
 142         float2 * __restrict__ gm_gridCell   = gm_grid + gridIndex;
 143
 144         const int             kMajor  = indexMajor + localOffsetMajor;
 145         /* Checking either X in XYZ, or Y in YZX cases */
 146         const float           mMajor  = (kMajor < maxkMajor) ? kMajor : (kMajor - nMajor);
 147
 148         const int             kMiddle = indexMiddle + localOffsetMiddle;
 149         float                 mMiddle = kMiddle;
 150         /* Checking Y in XYZ case */
 151         if (gridOrdering == GridOrdering::XYZ)
 152         {
 153             mMiddle = (kMiddle < maxkMiddle) ? kMiddle : (kMiddle - nMiddle);
 154         }
 155         const int             kMinor  = localOffsetMinor + indexMinor;
 156         float                 mMinor  = kMinor;
 157         /* Checking X in YZX case */
 158         if (gridOrdering == GridOrdering::YZX)
 159         {
 160             mMinor = (kMinor < maxkMinor) ? kMinor : (kMinor - nMinor);
 161         }
 162         /* We should skip the k-space point (0,0,0) */
 163         const bool notZeroPoint  = (kMinor > 0) | (kMajor > 0) | (kMiddle > 0);
 164
 165         float      mX, mY, mZ;
 166         switch (gridOrdering)
 167         {
 168             case GridOrdering::YZX:
 169                 mX = mMinor;
 170                 mY = mMajor;
 171                 mZ = mMiddle;
 172                 break;
 173
 174             case GridOrdering::XYZ:
 175                 mX = mMajor;
 176                 mY = mMiddle;
 177                 mZ = mMinor;
 178                 break;
 179
 180             default:
 181                 assert(false);
 182         }
 183
 184         /* 0.5 correction factor for the first and last components of a Z dimension */
 185         float corner_fac = 1.0f;
 186         switch (gridOrdering)
 187         {
 188             case GridOrdering::YZX:
 189                 if ((kMiddle == 0) | (kMiddle == maxkMiddle))
 190                 {
 191                     corner_fac = 0.5f;
 192                 }
 193                 break;
 194
 195             case GridOrdering::XYZ:
 196                 if ((kMinor == 0) | (kMinor == maxkMinor))
 197                 {
 198                     corner_fac = 0.5f;
 199                 }
 200                 break;
 201
 202             default:
 203                 assert(false);
 204         }
 205
 206         if (notZeroPoint)
 207         {
 208             const float mhxk = mX * kernelParams.current.recipBox[XX][XX];
 209             const float mhyk = mX * kernelParams.current.recipBox[XX][YY] + mY * kernelParams.current.recipBox[YY][YY];
 210             const float mhzk = mX * kernelParams.current.recipBox[XX][ZZ] + mY * kernelParams.current.recipBox[YY][ZZ] + mZ * kernelParams.current.recipBox[ZZ][ZZ];
 211
 212             const float m2k        = mhxk * mhxk + mhyk * mhyk + mhzk * mhzk;
 213             assert(m2k != 0.0f);
 214             //TODO: use LDG/textures for gm_splineValue
 215             float       denom = m2k * float(M_PI) * kernelParams.current.boxVolume * gm_splineValueMajor[kMajor] * gm_splineValueMiddle[kMiddle] * gm_splineValueMinor[kMinor];
 216             assert(isfinite(denom));
 217             assert(denom != 0.0f);
 218             const float   tmp1   = expf(-kernelParams.grid.ewaldFactor * m2k);
 219             const float   etermk = kernelParams.constants.elFactor * tmp1 / denom;
 220
 221             float2        gridValue    = *gm_gridCell;
 222             const float2  oldGridValue = gridValue;
 223             gridValue.x   *= etermk;
 224             gridValue.y   *= etermk;
 225             *gm_gridCell   = gridValue;
 226
 227             if (computeEnergyAndVirial)
 228             {
 229                 const float tmp1k = 2.0f * (gridValue.x * oldGridValue.x + gridValue.y * oldGridValue.y);
 230
 231                 float       vfactor = (kernelParams.grid.ewaldFactor + 1.0f / m2k) * 2.0f;
 232                 float       ets2    = corner_fac * tmp1k;
 233                 energy = ets2;
 234
 235                 float ets2vf  = ets2 * vfactor;
 236
 237                 virxx   = ets2vf * mhxk * mhxk - ets2;
 238                 virxy   = ets2vf * mhxk * mhyk;
 239                 virxz   = ets2vf * mhxk * mhzk;
 240                 viryy   = ets2vf * mhyk * mhyk - ets2;
 241                 viryz   = ets2vf * mhyk * mhzk;
 242                 virzz   = ets2vf * mhzk * mhzk - ets2;
 243             }
 244         }
 245     }
 246
 247     /* Optional energy/virial reduction */
 248     if (computeEnergyAndVirial)
 249     {
 250 #if (GMX_PTX_ARCH >= 300)
 251         /* A tricky shuffle reduction inspired by reduce_force_j_warp_shfl.
 252          * The idea is to reduce 7 energy/virial components into a single variable (aligned by 8).
 253          * We will reduce everything into virxx.
 254          */
 255
 256         /* We can only reduce warp-wise */
 257         const int          width      = warp_size;
 258         const unsigned int activeMask = c_fullWarpMask;
 259
 260         /* Making pair sums */
 261         virxx  += gmx_shfl_down_sync(activeMask, virxx, 1, width);
 262         viryy  += gmx_shfl_up_sync  (activeMask, viryy, 1, width);
 263         virzz  += gmx_shfl_down_sync(activeMask, virzz, 1, width);
 264         virxy  += gmx_shfl_up_sync  (activeMask, virxy, 1, width);
 265         virxz  += gmx_shfl_down_sync(activeMask, virxz, 1, width);
 266         viryz  += gmx_shfl_up_sync  (activeMask, viryz, 1, width);
 267         energy += gmx_shfl_down_sync(activeMask, energy, 1, width);
 268         if (threadLocalId & 1)
 269         {
 270             virxx = viryy; // virxx now holds virxx and viryy pair sums
 271             virzz = virxy; // virzz now holds virzz and virxy pair sums
 272             virxz = viryz; // virxz now holds virxz and viryz pair sums
 273         }
 274
 275         /* Making quad sums */
 276         virxx  += gmx_shfl_down_sync(activeMask, virxx, 2, width);
 277         virzz  += gmx_shfl_up_sync  (activeMask, virzz, 2, width);
 278         virxz  += gmx_shfl_down_sync(activeMask, virxz, 2, width);
 279         energy += gmx_shfl_up_sync  (activeMask, energy, 2, width);
 280         if (threadLocalId & 2)
 281         {
 282             virxx = virzz;  // virxx now holds quad sums of virxx, virxy, virzz and virxy
 283             virxz = energy; // virxz now holds quad sums of virxz, viryz, energy and unused paddings
 284         }
 285
 286         /* Making octet sums */
 287         virxx += gmx_shfl_down_sync(activeMask, virxx, 4, width);
 288         virxz += gmx_shfl_up_sync  (activeMask, virxz, 4, width);
 289         if (threadLocalId & 4)
 290         {
 291             virxx = virxz; // virxx now holds all 7 components' octet sums + unused paddings
 292         }
 293
 294         /* We only need to reduce virxx now */
 295 #pragma unroll
 296         for (int delta = 8; delta < width; delta <<= 1)
 297         {
 298             virxx += gmx_shfl_down_sync(activeMask, virxx, delta, width);
 299         }
 300         /* Now first 7 threads of each warp have the full output contributions in virxx */
 301
 302         const int        componentIndex      = threadLocalId & (warp_size - 1);
 303         const bool       validComponentIndex = (componentIndex < c_virialAndEnergyCount);
 304         /* Reduce 7 outputs per warp in the shared memory */
 305         const int        stride              = 8; // this is c_virialAndEnergyCount==7 rounded up to power of 2 for convenience, hence the assert
 306         assert(c_virialAndEnergyCount == 7);
 307         const int        reductionBufferSize = (c_solveMaxThreadsPerBlock / warp_size) * stride;
 308         __shared__ float sm_virialAndEnergy[reductionBufferSize];
 309
 310         if (validComponentIndex)
 311         {
 312             const int warpIndex = threadLocalId / warp_size;
 313             sm_virialAndEnergy[warpIndex * stride + componentIndex] = virxx;
 314         }
 315         __syncthreads();
 316
 317         /* Reduce to the single warp size */
 318         const int targetIndex = threadLocalId;
 319 #pragma unroll
 320         for (int reductionStride = reductionBufferSize >> 1; reductionStride >= warp_size; reductionStride >>= 1)
 321         {
 322             const int sourceIndex = targetIndex + reductionStride;
 323             if ((targetIndex < reductionStride) & (sourceIndex < activeWarps * stride))
 324             {
 325                 // TODO: the second conditional is only needed on first iteration, actually - see if compiler eliminates it!
 326                 sm_virialAndEnergy[targetIndex] += sm_virialAndEnergy[sourceIndex];
 327             }
 328             __syncthreads();
 329         }
 330
 331         /* Now use shuffle again */
 332         if (threadLocalId < warp_size)
 333         {
 334             float output = sm_virialAndEnergy[threadLocalId];
 335 #pragma unroll
 336             for (int delta = stride; delta < warp_size; delta <<= 1)
 337             {
 338                 output += gmx_shfl_down_sync(activeMask, output, delta, warp_size);
 339             }
 340             /* Final output */
 341             if (validComponentIndex)
 342             {
 343                 assert(isfinite(output));
 344                 atomicAdd(gm_virialAndEnergy + componentIndex, output);
 345             }
 346         }
 347 #else
 348         /* Shared memory reduction with atomics for compute capability < 3.0.
 349          * Each component is first reduced into warp_size positions in the shared memory;
 350          * Then first c_virialAndEnergyCount warps reduce everything further and add to the global memory.
 351          * This can likely be improved, but is anyway faster than the previous straightforward reduction,
 352          * which was using too much shared memory (for storing all 7 floats on each thread).
 353          * [48KB (shared mem limit per SM on CC2.x) / sizeof(float) (4) / c_solveMaxThreadsPerBlock (256) / c_virialAndEnergyCount (7) ==
 354          * 6 blocks per SM instead of 16 which is maximum on CC2.x].
 355          */
 356
 357         const int        lane      = threadLocalId & (warp_size - 1);
 358         const int        warpIndex = threadLocalId / warp_size;
 359         const bool       firstWarp = (warpIndex == 0);
 360         __shared__ float sm_virialAndEnergy[c_virialAndEnergyCount * warp_size];
 361         if (firstWarp)
 362         {
 363             sm_virialAndEnergy[0 * warp_size + lane] = virxx;
 364             sm_virialAndEnergy[1 * warp_size + lane] = viryy;
 365             sm_virialAndEnergy[2 * warp_size + lane] = virzz;
 366             sm_virialAndEnergy[3 * warp_size + lane] = virxy;
 367             sm_virialAndEnergy[4 * warp_size + lane] = virxz;
 368             sm_virialAndEnergy[5 * warp_size + lane] = viryz;
 369             sm_virialAndEnergy[6 * warp_size + lane] = energy;
 370         }
 371         __syncthreads();
 372         if (!firstWarp)
 373         {
 374             atomicAdd(sm_virialAndEnergy + 0 * warp_size + lane, virxx);
 375             atomicAdd(sm_virialAndEnergy + 1 * warp_size + lane, viryy);
 376             atomicAdd(sm_virialAndEnergy + 2 * warp_size + lane, virzz);
 377             atomicAdd(sm_virialAndEnergy + 3 * warp_size + lane, virxy);
 378             atomicAdd(sm_virialAndEnergy + 4 * warp_size + lane, virxz);
 379             atomicAdd(sm_virialAndEnergy + 5 * warp_size + lane, viryz);
 380             atomicAdd(sm_virialAndEnergy + 6 * warp_size + lane, energy);
 381         }
 382         __syncthreads();
 383
 384         GMX_UNUSED_VALUE(activeWarps);
 385         assert(activeWarps >= c_virialAndEnergyCount); // we need to cover all components, or have multiple iterations otherwise
 386         const int componentIndex = warpIndex;
 387         if (componentIndex < c_virialAndEnergyCount)
 388         {
 389             const int targetIndex = threadLocalId;
 390 #pragma unroll
 391             for (int reductionStride = warp_size >> 1; reductionStride >= 1; reductionStride >>= 1)
 392             {
 393                 if (lane < reductionStride)
 394                 {
 395                     sm_virialAndEnergy[targetIndex] += sm_virialAndEnergy[targetIndex + reductionStride];
 396                 }
 397             }
 398             if (lane == 0)
 399             {
 400                 atomicAdd(gm_virialAndEnergy + componentIndex, sm_virialAndEnergy[targetIndex]);
 401             }
 402         }
 403 #endif
 404     }
 405 }
 406
 407 void pme_gpu_solve(const PmeGpu *pmeGpu, t_complex *h_grid,
 408                    GridOrdering gridOrdering, bool computeEnergyAndVirial)
 409 {
 410     const bool   copyInputAndOutputGrid = pme_gpu_is_testing(pmeGpu) || !pme_gpu_performs_FFT(pmeGpu);
 411
 412     cudaStream_t stream          = pmeGpu->archSpecific->pmeStream;
 413     auto        *kernelParamsPtr = pmeGpu->kernelParams.get();
 414
 415     float       *h_gridFloat = reinterpret_cast<float *>(h_grid);
 416     if (copyInputAndOutputGrid)
 417     {
 418         copyToDeviceBuffer(&kernelParamsPtr->grid.d_fourierGrid, h_gridFloat,
 419                            0, pmeGpu->archSpecific->complexGridSize,
 420                            stream, pmeGpu->settings.transferKind, nullptr);
 421     }
 422
 423     int majorDim = -1, middleDim = -1, minorDim = -1;
 424     switch (gridOrdering)
 425     {
 426         case GridOrdering::YZX:
 427             majorDim  = YY;
 428             middleDim = ZZ;
 429             minorDim  = XX;
 430             break;
 431
 432         case GridOrdering::XYZ:
 433             majorDim  = XX;
 434             middleDim = YY;
 435             minorDim  = ZZ;
 436             break;
 437
 438         default:
 439             GMX_ASSERT(false, "Implement grid ordering here and below for the kernel launch");
 440     }
 441
 442     const int maxBlockSize      = c_solveMaxThreadsPerBlock;
 443     const int gridLineSize      = pmeGpu->kernelParams->grid.complexGridSize[minorDim];
 444     const int gridLinesPerBlock = std::max(maxBlockSize / gridLineSize, 1);
 445     const int blocksPerGridLine = (gridLineSize + maxBlockSize - 1) / maxBlockSize;
 446     const int cellsPerBlock     = gridLineSize * gridLinesPerBlock;
 447     const int blockSize         = (cellsPerBlock + warp_size - 1) / warp_size * warp_size;
 448     // rounding up to full warps so that shuffle operations produce defined results
 449     dim3 threads(blockSize);
 450     dim3 blocks(blocksPerGridLine,
 451                 (pmeGpu->kernelParams->grid.complexGridSize[middleDim] + gridLinesPerBlock - 1) / gridLinesPerBlock,
 452                 pmeGpu->kernelParams->grid.complexGridSize[majorDim]);
 453
 454     pme_gpu_start_timing(pmeGpu, gtPME_SOLVE);
 455     if (gridOrdering == GridOrdering::YZX)
 456     {
 457         if (computeEnergyAndVirial)
 458         {
 459             pme_solve_kernel<GridOrdering::YZX, true> <<< blocks, threads, 0, stream>>> (*kernelParamsPtr);
 460         }
 461         else
 462         {
 463             pme_solve_kernel<GridOrdering::YZX, false> <<< blocks, threads, 0, stream>>> (*kernelParamsPtr);
 464         }
 465     }
 466     else if (gridOrdering == GridOrdering::XYZ)
 467     {
 468         if (computeEnergyAndVirial)
 469         {
 470             pme_solve_kernel<GridOrdering::XYZ, true> <<< blocks, threads, 0, stream>>> (*kernelParamsPtr);
 471         }
 472         else
 473         {
 474             pme_solve_kernel<GridOrdering::XYZ, false> <<< blocks, threads, 0, stream>>> (*kernelParamsPtr);
 475         }
 476     }
 477     CU_LAUNCH_ERR("pme_solve_kernel");
 478     pme_gpu_stop_timing(pmeGpu, gtPME_SOLVE);
 479
 480     if (computeEnergyAndVirial)
 481     {
 482         copyFromDeviceBuffer(pmeGpu->staging.h_virialAndEnergy, &kernelParamsPtr->constants.d_virialAndEnergy,
 483                              0, c_virialAndEnergyCount,
 484                              stream, pmeGpu->settings.transferKind, nullptr);
 485     }
 486
 487     if (copyInputAndOutputGrid)
 488     {
 489         copyFromDeviceBuffer(h_gridFloat, &kernelParamsPtr->grid.d_fourierGrid,
 490                              0, pmeGpu->archSpecific->complexGridSize,
 491                              stream, pmeGpu->settings.transferKind, nullptr);
 492     }
 493 }