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