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