src/gromacs/mdlib/gpuforcereduction_impl.cu

   1 /*
   2  * This file is part of the GROMACS molecular simulation package.
   3  *
   4  * Copyright (c) 2020, 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 /*! \internal \file
  36  *
  37  * \brief Implements GPU Force Reduction using CUDA
  38  *
  39  * \author Alan Gray <alang@nvidia.com>
  40  *
  41  * \ingroup module_mdlib
  42  */
  43
  44 #include "gmxpre.h"
  45
  46 #include "gpuforcereduction_impl.cuh"
  47
  48 #include <stdio.h>
  49
  50 #include "gromacs/gpu_utils/cudautils.cuh"
  51 #include "gromacs/gpu_utils/device_context.h"
  52 #include "gromacs/gpu_utils/devicebuffer.h"
  53 #include "gromacs/gpu_utils/gpu_utils.h"
  54 #include "gromacs/gpu_utils/gpueventsynchronizer.cuh"
  55 #include "gromacs/gpu_utils/typecasts.cuh"
  56 #include "gromacs/gpu_utils/vectype_ops.cuh"
  57 #include "gromacs/utility/gmxassert.h"
  58
  59 #include "gpuforcereduction.h"
  60
  61 namespace gmx
  62 {
  63
  64 constexpr static int c_threadsPerBlock = 128;
  65
  66 typedef struct rvecDeviceForceData rvecDeviceForceData_t;
  67
  68
  69 template<bool addRvecForce, bool accumulateForce>
  70 static __global__ void reduceKernel(const float3* __restrict__ gm_nbnxmForce,
  71                                     const float3* __restrict__ rvecForceToAdd,
  72                                     float3*    gm_fTotal,
  73                                     const int* gm_cell,
  74                                     const int  numAtoms)
  75 {
  76
  77     // map particle-level parallelism to 1D CUDA thread and block index
  78     const int threadIndex = blockIdx.x * blockDim.x + threadIdx.x;
  79
  80     // perform addition for each particle
  81     if (threadIndex < numAtoms)
  82     {
  83
  84         float3* gm_fDest = &gm_fTotal[threadIndex];
  85         float3  temp;
  86
  87         // Accumulate or set nbnxm force
  88         if (accumulateForce)
  89         {
  90             temp = *gm_fDest;
  91             temp += gm_nbnxmForce[gm_cell[threadIndex]];
  92         }
  93         else
  94         {
  95             temp = gm_nbnxmForce[gm_cell[threadIndex]];
  96         }
  97
  98         if (addRvecForce)
  99         {
 100             temp += rvecForceToAdd[threadIndex];
 101         }
 102
 103         *gm_fDest = temp;
 104     }
 105     return;
 106 }
 107
 108 GpuForceReduction::Impl::Impl(const DeviceContext& deviceContext, const DeviceStream& deviceStream) :
 109     deviceContext_(deviceContext),
 110     deviceStream_(deviceStream){};
 111
 112 void GpuForceReduction::Impl::reinit(float3*               baseForcePtr,
 113                                      const int             numAtoms,
 114                                      ArrayRef<const int>   cell,
 115                                      const int             atomStart,
 116                                      const bool            accumulate,
 117                                      GpuEventSynchronizer* completionMarker)
 118 {
 119     GMX_ASSERT((baseForcePtr != nullptr), "Input base force for reduction has no data");
 120     baseForce_        = &(baseForcePtr[atomStart]);
 121     numAtoms_         = numAtoms;
 122     atomStart_        = atomStart;
 123     accumulate_       = accumulate;
 124     completionMarker_ = completionMarker;
 125     cellInfo_.cell    = cell.data();
 126     reallocateDeviceBuffer(&cellInfo_.d_cell, numAtoms_, &cellInfo_.cellSize,
 127                            &cellInfo_.cellSizeAlloc, deviceContext_);
 128     copyToDeviceBuffer(&cellInfo_.d_cell, &(cellInfo_.cell[atomStart]), 0, numAtoms_, deviceStream_,
 129                        GpuApiCallBehavior::Async, nullptr);
 130
 131     dependencyList_.clear();
 132 };
 133
 134 void GpuForceReduction::Impl::registerNbnxmForce(DeviceBuffer<RVec> forcePtr)
 135 {
 136     GMX_ASSERT((forcePtr != nullptr), "Input force for reduction has no data");
 137     nbnxmForceToAdd_ = forcePtr;
 138 };
 139
 140 void GpuForceReduction::Impl::registerRvecForce(DeviceBuffer<RVec> forcePtr)
 141 {
 142     GMX_ASSERT((forcePtr != nullptr), "Input force for reduction has no data");
 143     rvecForceToAdd_ = forcePtr;
 144 };
 145
 146 void GpuForceReduction::Impl::addDependency(GpuEventSynchronizer* const dependency)
 147 {
 148     dependencyList_.push_back(dependency);
 149 }
 150
 151 void GpuForceReduction::Impl::execute()
 152 {
 153
 154     if (numAtoms_ == 0)
 155     {
 156         return;
 157     }
 158
 159     GMX_ASSERT((nbnxmForceToAdd_ != nullptr), "Nbnxm force for reduction has no data");
 160
 161     // Enqueue wait on all dependencies passed
 162     for (auto const synchronizer : dependencyList_)
 163     {
 164         synchronizer->enqueueWaitEvent(deviceStream_);
 165     }
 166
 167     float3* d_nbnxmForce     = asFloat3(nbnxmForceToAdd_);
 168     float3* d_rvecForceToAdd = &(asFloat3(rvecForceToAdd_)[atomStart_]);
 169
 170     // Configure and launch kernel
 171     KernelLaunchConfig config;
 172     config.blockSize[0]     = c_threadsPerBlock;
 173     config.blockSize[1]     = 1;
 174     config.blockSize[2]     = 1;
 175     config.gridSize[0]      = ((numAtoms_ + 1) + c_threadsPerBlock - 1) / c_threadsPerBlock;
 176     config.gridSize[1]      = 1;
 177     config.gridSize[2]      = 1;
 178     config.sharedMemorySize = 0;
 179
 180     auto kernelFn = (rvecForceToAdd_ != nullptr)
 181                             ? (accumulate_ ? reduceKernel<true, true> : reduceKernel<true, false>)
 182                             : (accumulate_ ? reduceKernel<false, true> : reduceKernel<false, false>);
 183
 184     const auto kernelArgs = prepareGpuKernelArguments(kernelFn, config, &d_nbnxmForce, &d_rvecForceToAdd,
 185                                                       &baseForce_, &cellInfo_.d_cell, &numAtoms_);
 186
 187     launchGpuKernel(kernelFn, config, deviceStream_, nullptr, "Force Reduction", kernelArgs);
 188
 189     // Mark that kernel has been launched
 190     if (completionMarker_ != nullptr)
 191     {
 192         completionMarker_->markEvent(deviceStream_);
 193     }
 194 }
 195
 196 GpuForceReduction::Impl::~Impl(){};
 197
 198 GpuForceReduction::GpuForceReduction(const DeviceContext& deviceContext, const DeviceStream& deviceStream) :
 199     impl_(new Impl(deviceContext, deviceStream))
 200 {
 201 }
 202
 203 void GpuForceReduction::registerNbnxmForce(void* forcePtr)
 204 {
 205     impl_->registerNbnxmForce(reinterpret_cast<DeviceBuffer<RVec>>(forcePtr));
 206 }
 207
 208 void GpuForceReduction::registerRvecForce(void* forcePtr)
 209 {
 210     impl_->registerRvecForce(reinterpret_cast<DeviceBuffer<RVec>>(forcePtr));
 211 }
 212
 213 void GpuForceReduction::addDependency(GpuEventSynchronizer* const dependency)
 214 {
 215     impl_->addDependency(dependency);
 216 }
 217
 218 void GpuForceReduction::reinit(DeviceBuffer<RVec>    baseForcePtr,
 219                                const int             numAtoms,
 220                                ArrayRef<const int>   cell,
 221                                const int             atomStart,
 222                                const bool            accumulate,
 223                                GpuEventSynchronizer* completionMarker)
 224 {
 225     impl_->reinit(asFloat3(baseForcePtr), numAtoms, cell, atomStart, accumulate, completionMarker);
 226 }
 227 void GpuForceReduction::execute()
 228 {
 229     impl_->execute();
 230 }
 231
 232 GpuForceReduction::~GpuForceReduction() = default;
 233
 234 } // namespace gmx