src/gromacs/mdlib/gpuforcereduction_impl.cu

   1 /*
   2  * This file is part of the GROMACS molecular simulation package.
   3  *
   4  * Copyright (c) 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 /*! \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,
 109                               const DeviceStream&  deviceStream,
 110                               gmx_wallcycle*       wcycle) :
 111     deviceContext_(deviceContext),
 112     deviceStream_(deviceStream),
 113     wcycle_(wcycle){};
 114
 115 void GpuForceReduction::Impl::reinit(float3*               baseForcePtr,
 116                                      const int             numAtoms,
 117                                      ArrayRef<const int>   cell,
 118                                      const int             atomStart,
 119                                      const bool            accumulate,
 120                                      GpuEventSynchronizer* completionMarker)
 121 {
 122     GMX_ASSERT((baseForcePtr != nullptr), "Input base force for reduction has no data");
 123     baseForce_        = &(baseForcePtr[atomStart]);
 124     numAtoms_         = numAtoms;
 125     atomStart_        = atomStart;
 126     accumulate_       = accumulate;
 127     completionMarker_ = completionMarker;
 128     cellInfo_.cell    = cell.data();
 129
 130     wallcycle_start_nocount(wcycle_, ewcLAUNCH_GPU);
 131     reallocateDeviceBuffer(
 132             &cellInfo_.d_cell, numAtoms_, &cellInfo_.cellSize, &cellInfo_.cellSizeAlloc, deviceContext_);
 133     copyToDeviceBuffer(&cellInfo_.d_cell,
 134                        &(cellInfo_.cell[atomStart]),
 135                        0,
 136                        numAtoms_,
 137                        deviceStream_,
 138                        GpuApiCallBehavior::Async,
 139                        nullptr);
 140     wallcycle_stop(wcycle_, ewcLAUNCH_GPU);
 141
 142     dependencyList_.clear();
 143 };
 144
 145 void GpuForceReduction::Impl::registerNbnxmForce(DeviceBuffer<RVec> forcePtr)
 146 {
 147     GMX_ASSERT((forcePtr != nullptr), "Input force for reduction has no data");
 148     nbnxmForceToAdd_ = forcePtr;
 149 };
 150
 151 void GpuForceReduction::Impl::registerRvecForce(DeviceBuffer<RVec> forcePtr)
 152 {
 153     GMX_ASSERT((forcePtr != nullptr), "Input force for reduction has no data");
 154     rvecForceToAdd_ = forcePtr;
 155 };
 156
 157 void GpuForceReduction::Impl::addDependency(GpuEventSynchronizer* const dependency)
 158 {
 159     dependencyList_.push_back(dependency);
 160 }
 161
 162 void GpuForceReduction::Impl::execute()
 163 {
 164     wallcycle_start_nocount(wcycle_, ewcLAUNCH_GPU);
 165     wallcycle_sub_start(wcycle_, ewcsLAUNCH_GPU_NB_F_BUF_OPS);
 166
 167     if (numAtoms_ == 0)
 168     {
 169         return;
 170     }
 171
 172     GMX_ASSERT((nbnxmForceToAdd_ != nullptr), "Nbnxm force for reduction has no data");
 173
 174     // Enqueue wait on all dependencies passed
 175     for (auto const synchronizer : dependencyList_)
 176     {
 177         synchronizer->enqueueWaitEvent(deviceStream_);
 178     }
 179
 180     float3* d_nbnxmForce     = asFloat3(nbnxmForceToAdd_);
 181     float3* d_rvecForceToAdd = &(asFloat3(rvecForceToAdd_)[atomStart_]);
 182
 183     // Configure and launch kernel
 184     KernelLaunchConfig config;
 185     config.blockSize[0]     = c_threadsPerBlock;
 186     config.blockSize[1]     = 1;
 187     config.blockSize[2]     = 1;
 188     config.gridSize[0]      = ((numAtoms_ + 1) + c_threadsPerBlock - 1) / c_threadsPerBlock;
 189     config.gridSize[1]      = 1;
 190     config.gridSize[2]      = 1;
 191     config.sharedMemorySize = 0;
 192
 193     auto kernelFn = (rvecForceToAdd_ != nullptr)
 194                             ? (accumulate_ ? reduceKernel<true, true> : reduceKernel<true, false>)
 195                             : (accumulate_ ? reduceKernel<false, true> : reduceKernel<false, false>);
 196
 197     const auto kernelArgs = prepareGpuKernelArguments(
 198             kernelFn, config, &d_nbnxmForce, &d_rvecForceToAdd, &baseForce_, &cellInfo_.d_cell, &numAtoms_);
 199
 200     launchGpuKernel(kernelFn, config, deviceStream_, nullptr, "Force Reduction", kernelArgs);
 201
 202     // Mark that kernel has been launched
 203     if (completionMarker_ != nullptr)
 204     {
 205         completionMarker_->markEvent(deviceStream_);
 206     }
 207
 208     wallcycle_sub_stop(wcycle_, ewcsLAUNCH_GPU_NB_F_BUF_OPS);
 209     wallcycle_stop(wcycle_, ewcLAUNCH_GPU);
 210 }
 211
 212 GpuForceReduction::Impl::~Impl(){};
 213
 214 GpuForceReduction::GpuForceReduction(const DeviceContext& deviceContext,
 215                                      const DeviceStream&  deviceStream,
 216                                      gmx_wallcycle*       wcycle) :
 217     impl_(new Impl(deviceContext, deviceStream, wcycle))
 218 {
 219 }
 220
 221 void GpuForceReduction::registerNbnxmForce(DeviceBuffer<Float3> forcePtr)
 222 {
 223     impl_->registerNbnxmForce(forcePtr);
 224 }
 225
 226 void GpuForceReduction::registerRvecForce(void* forcePtr)
 227 {
 228     impl_->registerRvecForce(reinterpret_cast<DeviceBuffer<RVec>>(forcePtr));
 229 }
 230
 231 void GpuForceReduction::addDependency(GpuEventSynchronizer* const dependency)
 232 {
 233     impl_->addDependency(dependency);
 234 }
 235
 236 void GpuForceReduction::reinit(DeviceBuffer<RVec>    baseForcePtr,
 237                                const int             numAtoms,
 238                                ArrayRef<const int>   cell,
 239                                const int             atomStart,
 240                                const bool            accumulate,
 241                                GpuEventSynchronizer* completionMarker)
 242 {
 243     impl_->reinit(asFloat3(baseForcePtr), numAtoms, cell, atomStart, accumulate, completionMarker);
 244 }
 245 void GpuForceReduction::execute()
 246 {
 247     impl_->execute();
 248 }
 249
 250 GpuForceReduction::~GpuForceReduction() = default;
 251
 252 } // namespace gmx