src/gromacs/ewald/pme-gpu.cpp

   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 high-level PME GPU functions which do not require GPU framework-specific code.
  38  *
  39  * \author Aleksei Iupinov <a.yupinov@gmail.com>
  40  * \ingroup module_ewald
  41  */
  42
  43 #include "gmxpre.h"
  44
  45 #include <list>
  46
  47 #include "gromacs/ewald/ewald-utils.h"
  48 #include "gromacs/ewald/pme.h"
  49 #include "gromacs/fft/parallel_3dfft.h"
  50 #include "gromacs/math/invertmatrix.h"
  51 #include "gromacs/mdtypes/inputrec.h"
  52 #include "gromacs/utility/exceptions.h"
  53 #include "gromacs/utility/fatalerror.h"
  54 #include "gromacs/utility/gmxassert.h"
  55 #include "gromacs/utility/stringutil.h"
  56
  57 #include "pme-gpu-internal.h"
  58 #include "pme-grid.h"
  59 #include "pme-internal.h"
  60 #include "pme-solve.h"
  61
  62 void pme_gpu_reset_timings(const gmx_pme_t *pme)
  63 {
  64     if (pme_gpu_active(pme))
  65     {
  66         pme_gpu_reset_timings(pme->gpu);
  67     }
  68 }
  69
  70 void pme_gpu_get_timings(const gmx_pme_t *pme, gmx_wallclock_gpu_pme_t *timings)
  71 {
  72     if (pme_gpu_active(pme))
  73     {
  74         pme_gpu_get_timings(pme->gpu, timings);
  75     }
  76 }
  77
  78 /*! \brief
  79  * A convenience wrapper for launching either the GPU or CPU FFT.
  80  *
  81  * \param[in] pme            The PME structure.
  82  * \param[in] gridIndex      The grid index - should currently always be 0.
  83  * \param[in] dir            The FFT direction enum.
  84  * \param[in] wcycle         The wallclock counter.
  85  */
  86 void inline parallel_3dfft_execute_gpu_wrapper(gmx_pme_t              *pme,
  87                                                const int               gridIndex,
  88                                                enum gmx_fft_direction  dir,
  89                                                gmx_wallcycle_t         wcycle)
  90 {
  91     GMX_ASSERT(gridIndex == 0, "Only single grid supported");
  92     if (pme_gpu_performs_FFT(pme->gpu))
  93     {
  94         wallcycle_start_nocount(wcycle, ewcLAUNCH_GPU);
  95         wallcycle_sub_start_nocount(wcycle, ewcsLAUNCH_GPU_PME);
  96         pme_gpu_3dfft(pme->gpu, dir, gridIndex);
  97         wallcycle_sub_stop(wcycle, ewcsLAUNCH_GPU_PME);
  98         wallcycle_stop(wcycle, ewcLAUNCH_GPU);
  99     }
 100     else
 101     {
 102         wallcycle_start(wcycle, ewcPME_FFT_MIXED_MODE);
 103 #pragma omp parallel for num_threads(pme->nthread) schedule(static)
 104         for (int thread = 0; thread < pme->nthread; thread++)
 105         {
 106             gmx_parallel_3dfft_execute(pme->pfft_setup[gridIndex], dir, thread, wcycle);
 107         }
 108         wallcycle_stop(wcycle, ewcPME_FFT_MIXED_MODE);
 109     }
 110 }
 111
 112 /* The PME computation code split into a few separate functions. */
 113
 114 void pme_gpu_prepare_computation(gmx_pme_t            *pme,
 115                                  bool                  needToUpdateBox,
 116                                  const matrix          box,
 117                                  gmx_wallcycle        *wcycle,
 118                                  int                   flags)
 119 {
 120     GMX_ASSERT(pme_gpu_active(pme), "This should be a GPU run of PME but it is not enabled.");
 121     GMX_ASSERT(pme->nnodes > 0, "");
 122     GMX_ASSERT(pme->nnodes == 1 || pme->ndecompdim > 0, "");
 123
 124     PmeGpu *pmeGpu = pme->gpu;
 125     pmeGpu->settings.currentFlags = flags;
 126     // TODO these flags are only here to honor the CPU PME code, and probably should be removed
 127
 128     bool shouldUpdateBox = false;
 129     for (int i = 0; i < DIM; ++i)
 130     {
 131         for (int j = 0; j <= i; ++j)
 132         {
 133             shouldUpdateBox                  |= (pmeGpu->common->previousBox[i][j] != box[i][j]);
 134             pmeGpu->common->previousBox[i][j] = box[i][j];
 135         }
 136     }
 137
 138     if (needToUpdateBox || shouldUpdateBox) // || is to make the first computation always update
 139     {
 140         wallcycle_start_nocount(wcycle, ewcLAUNCH_GPU);
 141         wallcycle_sub_start_nocount(wcycle, ewcsLAUNCH_GPU_PME);
 142         pme_gpu_update_input_box(pmeGpu, box);
 143         wallcycle_sub_stop(wcycle, ewcsLAUNCH_GPU_PME);
 144         wallcycle_stop(wcycle, ewcLAUNCH_GPU);
 145
 146         if (!pme_gpu_performs_solve(pmeGpu))
 147         {
 148             // TODO remove code duplication and add test coverage
 149             matrix scaledBox;
 150             pmeGpu->common->boxScaler->scaleBox(box, scaledBox);
 151             gmx::invertBoxMatrix(scaledBox, pme->recipbox);
 152             pme->boxVolume = scaledBox[XX][XX] * scaledBox[YY][YY] * scaledBox[ZZ][ZZ];
 153         }
 154     }
 155 }
 156
 157
 158 void pme_gpu_launch_spread(gmx_pme_t            *pme,
 159                            const rvec           *x,
 160                            gmx_wallcycle        *wcycle)
 161 {
 162     GMX_ASSERT(pme_gpu_active(pme), "This should be a GPU run of PME but it is not enabled.");
 163
 164     PmeGpu *pmeGpu = pme->gpu;
 165
 166     // The only spot of PME GPU where LAUNCH_GPU counter increases call-count
 167     wallcycle_start(wcycle, ewcLAUNCH_GPU);
 168     // The only spot of PME GPU where ewcsLAUNCH_GPU_PME subcounter increases call-count
 169     wallcycle_sub_start(wcycle, ewcsLAUNCH_GPU_PME);
 170     pme_gpu_copy_input_coordinates(pmeGpu, x);
 171     wallcycle_sub_stop(wcycle, ewcsLAUNCH_GPU_PME);
 172     wallcycle_stop(wcycle, ewcLAUNCH_GPU);
 173
 174     const unsigned int gridIndex  = 0;
 175     real              *fftgrid    = pme->fftgrid[gridIndex];
 176     if (pmeGpu->settings.currentFlags & GMX_PME_SPREAD)
 177     {
 178         /* Spread the coefficients on a grid */
 179         const bool computeSplines = true;
 180         const bool spreadCharges  = true;
 181         wallcycle_start_nocount(wcycle, ewcLAUNCH_GPU);
 182         wallcycle_sub_start_nocount(wcycle, ewcsLAUNCH_GPU_PME);
 183         pme_gpu_spread(pmeGpu, gridIndex, fftgrid, computeSplines, spreadCharges);
 184         wallcycle_sub_stop(wcycle, ewcsLAUNCH_GPU_PME);
 185         wallcycle_stop(wcycle, ewcLAUNCH_GPU);
 186     }
 187 }
 188
 189 void pme_gpu_launch_complex_transforms(gmx_pme_t      *pme,
 190                                        gmx_wallcycle  *wcycle)
 191 {
 192     PmeGpu            *pmeGpu                 = pme->gpu;
 193     const bool         computeEnergyAndVirial = pmeGpu->settings.currentFlags & GMX_PME_CALC_ENER_VIR;
 194     const bool         performBackFFT         = pmeGpu->settings.currentFlags & (GMX_PME_CALC_F | GMX_PME_CALC_POT);
 195     const unsigned int gridIndex              = 0;
 196     t_complex         *cfftgrid               = pme->cfftgrid[gridIndex];
 197
 198     if (pmeGpu->settings.currentFlags & GMX_PME_SPREAD)
 199     {
 200         if (!pme_gpu_performs_FFT(pmeGpu))
 201         {
 202             wallcycle_start(wcycle, ewcWAIT_GPU_PME_SPREAD);
 203             pme_gpu_sync_spread_grid(pme->gpu);
 204             wallcycle_stop(wcycle, ewcWAIT_GPU_PME_SPREAD);
 205         }
 206     }
 207
 208     try
 209     {
 210         if (pmeGpu->settings.currentFlags & GMX_PME_SOLVE)
 211         {
 212             /* do R2C 3D-FFT */
 213             parallel_3dfft_execute_gpu_wrapper(pme, gridIndex, GMX_FFT_REAL_TO_COMPLEX, wcycle);
 214
 215             /* solve in k-space for our local cells */
 216             if (pme_gpu_performs_solve(pmeGpu))
 217             {
 218                 const auto gridOrdering = pme_gpu_uses_dd(pmeGpu) ? GridOrdering::YZX : GridOrdering::XYZ;
 219                 wallcycle_start_nocount(wcycle, ewcLAUNCH_GPU);
 220                 wallcycle_sub_start_nocount(wcycle, ewcsLAUNCH_GPU_PME);
 221                 pme_gpu_solve(pmeGpu, cfftgrid, gridOrdering, computeEnergyAndVirial);
 222                 wallcycle_sub_stop(wcycle, ewcsLAUNCH_GPU_PME);
 223                 wallcycle_stop(wcycle, ewcLAUNCH_GPU);
 224             }
 225             else
 226             {
 227                 wallcycle_start(wcycle, ewcPME_SOLVE_MIXED_MODE);
 228 #pragma omp parallel for num_threads(pme->nthread) schedule(static)
 229                 for (int thread = 0; thread < pme->nthread; thread++)
 230                 {
 231                     solve_pme_yzx(pme, cfftgrid, pme->boxVolume,
 232                                   computeEnergyAndVirial, pme->nthread, thread);
 233                 }
 234                 wallcycle_stop(wcycle, ewcPME_SOLVE_MIXED_MODE);
 235             }
 236         }
 237
 238         if (performBackFFT)
 239         {
 240             parallel_3dfft_execute_gpu_wrapper(pme, gridIndex, GMX_FFT_COMPLEX_TO_REAL, wcycle);
 241         }
 242     } GMX_CATCH_ALL_AND_EXIT_WITH_FATAL_ERROR;
 243 }
 244
 245 void pme_gpu_launch_gather(const gmx_pme_t                 *pme,
 246                            gmx_wallcycle gmx_unused        *wcycle,
 247                            PmeForceOutputHandling           forceTreatment)
 248 {
 249     GMX_ASSERT(pme_gpu_active(pme), "This should be a GPU run of PME but it is not enabled.");
 250
 251     if (!pme_gpu_performs_gather(pme->gpu))
 252     {
 253         return;
 254     }
 255
 256     wallcycle_start_nocount(wcycle, ewcLAUNCH_GPU);
 257     wallcycle_sub_start_nocount(wcycle, ewcsLAUNCH_GPU_PME);
 258     const unsigned int gridIndex  = 0;
 259     real              *fftgrid    = pme->fftgrid[gridIndex];
 260     pme_gpu_gather(pme->gpu, forceTreatment, reinterpret_cast<float *>(fftgrid));
 261     wallcycle_sub_stop(wcycle, ewcsLAUNCH_GPU_PME);
 262     wallcycle_stop(wcycle, ewcLAUNCH_GPU);
 263 }
 264
 265 /*! \brief Reduce staged virial and energy outputs.
 266  *
 267  * \param[in]  pme            The PME data structure.
 268  * \param[out] forces         Output forces pointer, the internal ArrayRef pointers gets assigned to it.
 269  * \param[out] virial         The output virial matrix.
 270  * \param[out] energy         The output energy.
 271  */
 272 static void pme_gpu_get_staged_results(const gmx_pme_t                *pme,
 273                                        gmx::ArrayRef<const gmx::RVec> *forces,
 274                                        matrix                          virial,
 275                                        real                           *energy)
 276 {
 277     const bool haveComputedEnergyAndVirial = pme->gpu->settings.currentFlags & GMX_PME_CALC_ENER_VIR;
 278     *forces = pme_gpu_get_forces(pme->gpu);
 279
 280     if (haveComputedEnergyAndVirial)
 281     {
 282         if (pme_gpu_performs_solve(pme->gpu))
 283         {
 284             pme_gpu_get_energy_virial(pme->gpu, energy, virial);
 285         }
 286         else
 287         {
 288             get_pme_ener_vir_q(pme->solve_work, pme->nthread, energy, virial);
 289         }
 290     }
 291 }
 292
 293 bool pme_gpu_try_finish_task(const gmx_pme_t                *pme,
 294                              gmx_wallcycle                  *wcycle,
 295                              gmx::ArrayRef<const gmx::RVec> *forces,
 296                              matrix                          virial,
 297                              real                           *energy,
 298                              GpuTaskCompletion               completionKind)
 299 {
 300     GMX_ASSERT(pme_gpu_active(pme), "This should be a GPU run of PME but it is not enabled.");
 301
 302     wallcycle_start_nocount(wcycle, ewcWAIT_GPU_PME_GATHER);
 303
 304     if (completionKind == GpuTaskCompletion::Check)
 305     {
 306         // Query the PME stream for completion of all tasks enqueued and
 307         // if we're not done, stop the timer before early return.
 308         if (!pme_gpu_stream_query(pme->gpu))
 309         {
 310             wallcycle_stop(wcycle, ewcWAIT_GPU_PME_GATHER);
 311             return false;
 312         }
 313     }
 314     else
 315     {
 316         // Synchronize the whole PME stream at once, including D2H result transfers.
 317         pme_gpu_synchronize(pme->gpu);
 318     }
 319     wallcycle_stop(wcycle, ewcWAIT_GPU_PME_GATHER);
 320
 321     // Time the final staged data handling separately with a counting call to get
 322     // the call count right.
 323     wallcycle_start(wcycle, ewcWAIT_GPU_PME_GATHER);
 324     pme_gpu_update_timings(pme->gpu);
 325     pme_gpu_get_staged_results(pme, forces, virial, energy);
 326     wallcycle_stop(wcycle, ewcWAIT_GPU_PME_GATHER);
 327
 328     return true;
 329 }
 330
 331 void pme_gpu_wait_finish_task(const gmx_pme_t                *pme,
 332                               gmx_wallcycle                  *wcycle,
 333                               gmx::ArrayRef<const gmx::RVec> *forces,
 334                               matrix                          virial,
 335                               real                           *energy)
 336 {
 337     pme_gpu_try_finish_task(pme, wcycle, forces, virial, energy, GpuTaskCompletion::Wait);
 338 }
 339
 340 void pme_gpu_reinit_computation(const gmx_pme_t *pme,
 341                                 gmx_wallcycle   *wcycle)
 342 {
 343     GMX_ASSERT(pme_gpu_active(pme), "This should be a GPU run of PME but it is not enabled.");
 344
 345     wallcycle_start_nocount(wcycle, ewcLAUNCH_GPU);
 346     wallcycle_sub_start_nocount(wcycle, ewcsLAUNCH_GPU_PME);
 347
 348     pme_gpu_clear_grids(pme->gpu);
 349     pme_gpu_clear_energy_virial(pme->gpu);
 350
 351     wallcycle_sub_stop(wcycle, ewcsLAUNCH_GPU_PME);
 352     wallcycle_stop(wcycle, ewcLAUNCH_GPU);
 353 }