src/gromacs/ewald/pme_gpu_internal.cpp

   1 /*
   2  * This file is part of the GROMACS molecular simulation package.
   3  *
   4  * Copyright (c) 2016,2017,2018,2019,2020, by the GROMACS development team.
   5  * Copyright (c) 2021, by the GROMACS development team, led by
   6  * Mark Abraham, David van der Spoel, Berk Hess, and Erik Lindahl,
   7  * and including many others, as listed in the AUTHORS file in the
   8  * top-level source directory and at http://www.gromacs.org.
   9  *
  10  * GROMACS is free software; you can redistribute it and/or
  11  * modify it under the terms of the GNU Lesser General Public License
  12  * as published by the Free Software Foundation; either version 2.1
  13  * of the License, or (at your option) any later version.
  14  *
  15  * GROMACS is distributed in the hope that it will be useful,
  16  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  17  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  18  * Lesser General Public License for more details.
  19  *
  20  * You should have received a copy of the GNU Lesser General Public
  21  * License along with GROMACS; if not, see
  22  * http://www.gnu.org/licenses, or write to the Free Software Foundation,
  23  * Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA.
  24  *
  25  * If you want to redistribute modifications to GROMACS, please
  26  * consider that scientific software is very special. Version
  27  * control is crucial - bugs must be traceable. We will be happy to
  28  * consider code for inclusion in the official distribution, but
  29  * derived work must not be called official GROMACS. Details are found
  30  * in the README & COPYING files - if they are missing, get the
  31  * official version at http://www.gromacs.org.
  32  *
  33  * To help us fund GROMACS development, we humbly ask that you cite
  34  * the research papers on the package. Check out http://www.gromacs.org.
  35  */
  36
  37 /*! \internal \file
  38  *
  39  * \brief This file contains internal function implementations
  40  * for performing the PME calculations on GPU.
  41  *
  42  * Note that this file is compiled as regular C++ source in OpenCL builds, but
  43  * it is treated as CUDA source in CUDA-enabled GPU builds.
  44  *
  45  * \author Aleksei Iupinov <a.yupinov@gmail.com>
  46  * \ingroup module_ewald
  47  */
  48
  49 #include "gmxpre.h"
  50
  51 #include "pme_gpu_internal.h"
  52
  53 #include "config.h"
  54
  55 #include <list>
  56 #include <memory>
  57 #include <string>
  58
  59 #include "gromacs/ewald/ewald_utils.h"
  60 #include "gromacs/gpu_utils/device_context.h"
  61 #include "gromacs/gpu_utils/device_stream.h"
  62 #include "gromacs/gpu_utils/gpu_utils.h"
  63 #include "gromacs/gpu_utils/pmalloc.h"
  64 #if GMX_GPU_SYCL
  65 #    include "gromacs/gpu_utils/syclutils.h"
  66 #endif
  67 #include "gromacs/hardware/device_information.h"
  68 #include "gromacs/math/invertmatrix.h"
  69 #include "gromacs/math/units.h"
  70 #include "gromacs/timing/gpu_timing.h"
  71 #include "gromacs/utility/exceptions.h"
  72 #include "gromacs/utility/fatalerror.h"
  73 #include "gromacs/utility/gmxassert.h"
  74 #include "gromacs/utility/logger.h"
  75 #include "gromacs/utility/stringutil.h"
  76 #include "gromacs/ewald/pme.h"
  77
  78 #if GMX_GPU_CUDA
  79 #    include "pme.cuh"
  80 #endif
  81
  82 #include "pme_gpu_3dfft.h"
  83 #include "pme_gpu_calculate_splines.h"
  84 #include "pme_gpu_constants.h"
  85 #include "pme_gpu_program_impl.h"
  86 #include "pme_gpu_timings.h"
  87 #include "pme_gpu_types.h"
  88 #include "pme_gpu_types_host.h"
  89 #include "pme_gpu_types_host_impl.h"
  90 #include "pme_grid.h"
  91 #include "pme_internal.h"
  92 #include "pme_solve.h"
  93
  94 /*! \brief
  95  * CUDA only
  96  * Atom limit above which it is advantageous to turn on the
  97  * recalculating of the splines in the gather and using less threads per atom in the spline and spread
  98  */
  99 constexpr int c_pmeGpuPerformanceAtomLimit = 23000;
 100
 101 /*! \internal \brief
 102  * Wrapper for getting a pointer to the plain C++ part of the GPU kernel parameters structure.
 103  *
 104  * \param[in] pmeGpu  The PME GPU structure.
 105  * \returns The pointer to the kernel parameters.
 106  */
 107 static PmeGpuKernelParamsBase* pme_gpu_get_kernel_params_base_ptr(const PmeGpu* pmeGpu)
 108 {
 109     // reinterpret_cast is needed because the derived CUDA structure is not known in this file
 110     auto* kernelParamsPtr = reinterpret_cast<PmeGpuKernelParamsBase*>(pmeGpu->kernelParams.get());
 111     return kernelParamsPtr;
 112 }
 113
 114 /*! \brief
 115  * Atom data block size (in terms of number of atoms).
 116  * This is the least common multiple of number of atoms processed by
 117  * a single block/workgroup of the spread and gather kernels.
 118  * The GPU atom data buffers must be padded, which means that
 119  * the numbers of atoms used for determining the size of the memory
 120  * allocation must be divisible by this.
 121  */
 122 constexpr int c_pmeAtomDataBlockSize = 64;
 123
 124 int pme_gpu_get_atom_data_block_size()
 125 {
 126     return c_pmeAtomDataBlockSize;
 127 }
 128
 129 void pme_gpu_synchronize(const PmeGpu* pmeGpu)
 130 {
 131     pmeGpu->archSpecific->pmeStream_.synchronize();
 132 }
 133
 134 void pme_gpu_alloc_energy_virial(PmeGpu* pmeGpu)
 135 {
 136     const size_t energyAndVirialSize = c_virialAndEnergyCount * sizeof(float);
 137
 138     GMX_ASSERT(
 139             pmeGpu->common->ngrids == 1 || pmeGpu->common->ngrids == 2,
 140             "Only one (normal Coulomb PME) or two (FEP coulomb PME) PME grids can be used on GPU");
 141
 142     for (int gridIndex = 0; gridIndex < pmeGpu->common->ngrids; gridIndex++)
 143     {
 144         allocateDeviceBuffer(&pmeGpu->kernelParams->constants.d_virialAndEnergy[gridIndex],
 145                              c_virialAndEnergyCount,
 146                              pmeGpu->archSpecific->deviceContext_);
 147         pmalloc(reinterpret_cast<void**>(&pmeGpu->staging.h_virialAndEnergy[gridIndex]), energyAndVirialSize);
 148     }
 149 }
 150
 151 void pme_gpu_free_energy_virial(PmeGpu* pmeGpu)
 152 {
 153     for (int gridIndex = 0; gridIndex < pmeGpu->common->ngrids; gridIndex++)
 154     {
 155         freeDeviceBuffer(&pmeGpu->kernelParams->constants.d_virialAndEnergy[gridIndex]);
 156         pfree(pmeGpu->staging.h_virialAndEnergy[gridIndex]);
 157         pmeGpu->staging.h_virialAndEnergy[gridIndex] = nullptr;
 158     }
 159 }
 160
 161 void pme_gpu_clear_energy_virial(const PmeGpu* pmeGpu)
 162 {
 163     for (int gridIndex = 0; gridIndex < pmeGpu->common->ngrids; gridIndex++)
 164     {
 165         clearDeviceBufferAsync(&pmeGpu->kernelParams->constants.d_virialAndEnergy[gridIndex],
 166                                0,
 167                                c_virialAndEnergyCount,
 168                                pmeGpu->archSpecific->pmeStream_);
 169     }
 170 }
 171
 172 void pme_gpu_realloc_and_copy_bspline_values(PmeGpu* pmeGpu, const int gridIndex)
 173 {
 174     GMX_ASSERT(
 175             pmeGpu->common->ngrids == 1 || pmeGpu->common->ngrids == 2,
 176             "Only one (normal Coulomb PME) or two (FEP coulomb PME) PME grids can be used on GPU");
 177     GMX_ASSERT(gridIndex < pmeGpu->common->ngrids,
 178                "Invalid combination of gridIndex and number of grids");
 179
 180     const int splineValuesOffset[DIM] = { 0,
 181                                           pmeGpu->kernelParams->grid.realGridSize[XX],
 182                                           pmeGpu->kernelParams->grid.realGridSize[XX]
 183                                                   + pmeGpu->kernelParams->grid.realGridSize[YY] };
 184     memcpy(&pmeGpu->kernelParams->grid.splineValuesOffset, &splineValuesOffset, sizeof(splineValuesOffset));
 185
 186     const int newSplineValuesSize = pmeGpu->kernelParams->grid.realGridSize[XX]
 187                                     + pmeGpu->kernelParams->grid.realGridSize[YY]
 188                                     + pmeGpu->kernelParams->grid.realGridSize[ZZ];
 189     const bool shouldRealloc = (newSplineValuesSize > pmeGpu->archSpecific->splineValuesSize[gridIndex]);
 190     reallocateDeviceBuffer(&pmeGpu->kernelParams->grid.d_splineModuli[gridIndex],
 191                            newSplineValuesSize,
 192                            &pmeGpu->archSpecific->splineValuesSize[gridIndex],
 193                            &pmeGpu->archSpecific->splineValuesCapacity[gridIndex],
 194                            pmeGpu->archSpecific->deviceContext_);
 195     if (shouldRealloc)
 196     {
 197         /* Reallocate the host buffer */
 198         pfree(pmeGpu->staging.h_splineModuli[gridIndex]);
 199         pmalloc(reinterpret_cast<void**>(&pmeGpu->staging.h_splineModuli[gridIndex]),
 200                 newSplineValuesSize * sizeof(float));
 201     }
 202     for (int i = 0; i < DIM; i++)
 203     {
 204         memcpy(pmeGpu->staging.h_splineModuli[gridIndex] + splineValuesOffset[i],
 205                pmeGpu->common->bsp_mod[i].data(),
 206                pmeGpu->common->bsp_mod[i].size() * sizeof(float));
 207     }
 208     /* TODO: pin original buffer instead! */
 209     copyToDeviceBuffer(&pmeGpu->kernelParams->grid.d_splineModuli[gridIndex],
 210                        pmeGpu->staging.h_splineModuli[gridIndex],
 211                        0,
 212                        newSplineValuesSize,
 213                        pmeGpu->archSpecific->pmeStream_,
 214                        pmeGpu->settings.transferKind,
 215                        nullptr);
 216 }
 217
 218 void pme_gpu_free_bspline_values(const PmeGpu* pmeGpu)
 219 {
 220     for (int gridIndex = 0; gridIndex < pmeGpu->common->ngrids; gridIndex++)
 221     {
 222         pfree(pmeGpu->staging.h_splineModuli[gridIndex]);
 223         freeDeviceBuffer(&pmeGpu->kernelParams->grid.d_splineModuli[gridIndex]);
 224     }
 225 }
 226
 227 void pme_gpu_realloc_forces(PmeGpu* pmeGpu)
 228 {
 229     const size_t newForcesSize = pmeGpu->nAtomsAlloc;
 230     GMX_ASSERT(newForcesSize > 0, "Bad number of atoms in PME GPU");
 231     reallocateDeviceBuffer(&pmeGpu->kernelParams->atoms.d_forces,
 232                            newForcesSize,
 233                            &pmeGpu->archSpecific->forcesSize,
 234                            &pmeGpu->archSpecific->forcesSizeAlloc,
 235                            pmeGpu->archSpecific->deviceContext_);
 236     pmeGpu->staging.h_forces.reserveWithPadding(pmeGpu->nAtomsAlloc);
 237     pmeGpu->staging.h_forces.resizeWithPadding(pmeGpu->kernelParams->atoms.nAtoms);
 238 }
 239
 240 void pme_gpu_free_forces(const PmeGpu* pmeGpu)
 241 {
 242     freeDeviceBuffer(&pmeGpu->kernelParams->atoms.d_forces);
 243 }
 244
 245 void pme_gpu_copy_input_forces(PmeGpu* pmeGpu)
 246 {
 247     GMX_ASSERT(pmeGpu->kernelParams->atoms.nAtoms > 0, "Bad number of atoms in PME GPU");
 248     copyToDeviceBuffer(&pmeGpu->kernelParams->atoms.d_forces,
 249                        pmeGpu->staging.h_forces.data(),
 250                        0,
 251                        pmeGpu->kernelParams->atoms.nAtoms,
 252                        pmeGpu->archSpecific->pmeStream_,
 253                        pmeGpu->settings.transferKind,
 254                        nullptr);
 255 }
 256
 257 void pme_gpu_copy_output_forces(PmeGpu* pmeGpu)
 258 {
 259     GMX_ASSERT(pmeGpu->kernelParams->atoms.nAtoms > 0, "Bad number of atoms in PME GPU");
 260     copyFromDeviceBuffer(pmeGpu->staging.h_forces.data(),
 261                          &pmeGpu->kernelParams->atoms.d_forces,
 262                          0,
 263                          pmeGpu->kernelParams->atoms.nAtoms,
 264                          pmeGpu->archSpecific->pmeStream_,
 265                          pmeGpu->settings.transferKind,
 266                          nullptr);
 267 }
 268
 269 void pme_gpu_realloc_and_copy_input_coefficients(const PmeGpu* pmeGpu,
 270                                                  const float*  h_coefficients,
 271                                                  const int     gridIndex)
 272 {
 273     GMX_ASSERT(h_coefficients, "Bad host-side charge buffer in PME GPU");
 274     const size_t newCoefficientsSize = pmeGpu->nAtomsAlloc;
 275     GMX_ASSERT(newCoefficientsSize > 0, "Bad number of atoms in PME GPU");
 276     reallocateDeviceBuffer(&pmeGpu->kernelParams->atoms.d_coefficients[gridIndex],
 277                            newCoefficientsSize,
 278                            &pmeGpu->archSpecific->coefficientsSize[gridIndex],
 279                            &pmeGpu->archSpecific->coefficientsCapacity[gridIndex],
 280                            pmeGpu->archSpecific->deviceContext_);
 281     copyToDeviceBuffer(&pmeGpu->kernelParams->atoms.d_coefficients[gridIndex],
 282                        const_cast<float*>(h_coefficients),
 283                        0,
 284                        pmeGpu->kernelParams->atoms.nAtoms,
 285                        pmeGpu->archSpecific->pmeStream_,
 286                        pmeGpu->settings.transferKind,
 287                        nullptr);
 288
 289     const size_t paddingIndex = pmeGpu->kernelParams->atoms.nAtoms;
 290     const size_t paddingCount = pmeGpu->nAtomsAlloc - paddingIndex;
 291     if (paddingCount > 0)
 292     {
 293         clearDeviceBufferAsync(&pmeGpu->kernelParams->atoms.d_coefficients[gridIndex],
 294                                paddingIndex,
 295                                paddingCount,
 296                                pmeGpu->archSpecific->pmeStream_);
 297     }
 298 }
 299
 300 void pme_gpu_free_coefficients(const PmeGpu* pmeGpu)
 301 {
 302     for (int gridIndex = 0; gridIndex < pmeGpu->common->ngrids; gridIndex++)
 303     {
 304         freeDeviceBuffer(&pmeGpu->kernelParams->atoms.d_coefficients[gridIndex]);
 305     }
 306 }
 307
 308 void pme_gpu_realloc_spline_data(PmeGpu* pmeGpu)
 309 {
 310     const int order             = pmeGpu->common->pme_order;
 311     const int newSplineDataSize = DIM * order * pmeGpu->nAtomsAlloc;
 312     GMX_ASSERT(newSplineDataSize > 0, "Bad number of atoms in PME GPU");
 313     /* Two arrays of the same size */
 314     const bool shouldRealloc        = (newSplineDataSize > pmeGpu->archSpecific->splineDataSize);
 315     int        currentSizeTemp      = pmeGpu->archSpecific->splineDataSize;
 316     int        currentSizeTempAlloc = pmeGpu->archSpecific->splineDataSizeAlloc;
 317     reallocateDeviceBuffer(&pmeGpu->kernelParams->atoms.d_theta,
 318                            newSplineDataSize,
 319                            &currentSizeTemp,
 320                            &currentSizeTempAlloc,
 321                            pmeGpu->archSpecific->deviceContext_);
 322     reallocateDeviceBuffer(&pmeGpu->kernelParams->atoms.d_dtheta,
 323                            newSplineDataSize,
 324                            &pmeGpu->archSpecific->splineDataSize,
 325                            &pmeGpu->archSpecific->splineDataSizeAlloc,
 326                            pmeGpu->archSpecific->deviceContext_);
 327     // the host side reallocation
 328     if (shouldRealloc)
 329     {
 330         pfree(pmeGpu->staging.h_theta);
 331         pmalloc(reinterpret_cast<void**>(&pmeGpu->staging.h_theta), newSplineDataSize * sizeof(float));
 332         pfree(pmeGpu->staging.h_dtheta);
 333         pmalloc(reinterpret_cast<void**>(&pmeGpu->staging.h_dtheta), newSplineDataSize * sizeof(float));
 334     }
 335 }
 336
 337 void pme_gpu_free_spline_data(const PmeGpu* pmeGpu)
 338 {
 339     /* Two arrays of the same size */
 340     freeDeviceBuffer(&pmeGpu->kernelParams->atoms.d_theta);
 341     freeDeviceBuffer(&pmeGpu->kernelParams->atoms.d_dtheta);
 342     pfree(pmeGpu->staging.h_theta);
 343     pfree(pmeGpu->staging.h_dtheta);
 344 }
 345
 346 void pme_gpu_realloc_grid_indices(PmeGpu* pmeGpu)
 347 {
 348     const size_t newIndicesSize = DIM * pmeGpu->nAtomsAlloc;
 349     GMX_ASSERT(newIndicesSize > 0, "Bad number of atoms in PME GPU");
 350     reallocateDeviceBuffer(&pmeGpu->kernelParams->atoms.d_gridlineIndices,
 351                            newIndicesSize,
 352                            &pmeGpu->archSpecific->gridlineIndicesSize,
 353                            &pmeGpu->archSpecific->gridlineIndicesSizeAlloc,
 354                            pmeGpu->archSpecific->deviceContext_);
 355     pfree(pmeGpu->staging.h_gridlineIndices);
 356     pmalloc(reinterpret_cast<void**>(&pmeGpu->staging.h_gridlineIndices), newIndicesSize * sizeof(int));
 357 }
 358
 359 void pme_gpu_free_grid_indices(const PmeGpu* pmeGpu)
 360 {
 361     freeDeviceBuffer(&pmeGpu->kernelParams->atoms.d_gridlineIndices);
 362     pfree(pmeGpu->staging.h_gridlineIndices);
 363 }
 364
 365 void pme_gpu_realloc_grids(PmeGpu* pmeGpu)
 366 {
 367     auto* kernelParamsPtr = pmeGpu->kernelParams.get();
 368
 369     const int newRealGridSize = kernelParamsPtr->grid.realGridSizePadded[XX]
 370                                 * kernelParamsPtr->grid.realGridSizePadded[YY]
 371                                 * kernelParamsPtr->grid.realGridSizePadded[ZZ];
 372     const int newComplexGridSize = kernelParamsPtr->grid.complexGridSizePadded[XX]
 373                                    * kernelParamsPtr->grid.complexGridSizePadded[YY]
 374                                    * kernelParamsPtr->grid.complexGridSizePadded[ZZ] * 2;
 375     for (int gridIndex = 0; gridIndex < pmeGpu->common->ngrids; gridIndex++)
 376     {
 377         // Multiplied by 2 because we count complex grid size for complex numbers, but all allocations/pointers are float
 378         if (pmeGpu->archSpecific->performOutOfPlaceFFT)
 379         {
 380             /* 2 separate grids */
 381             reallocateDeviceBuffer(&kernelParamsPtr->grid.d_fourierGrid[gridIndex],
 382                                    newComplexGridSize,
 383                                    &pmeGpu->archSpecific->complexGridSize[gridIndex],
 384                                    &pmeGpu->archSpecific->complexGridCapacity[gridIndex],
 385                                    pmeGpu->archSpecific->deviceContext_);
 386             reallocateDeviceBuffer(&kernelParamsPtr->grid.d_realGrid[gridIndex],
 387                                    newRealGridSize,
 388                                    &pmeGpu->archSpecific->realGridSize[gridIndex],
 389                                    &pmeGpu->archSpecific->realGridCapacity[gridIndex],
 390                                    pmeGpu->archSpecific->deviceContext_);
 391         }
 392         else
 393         {
 394             /* A single buffer so that any grid will fit */
 395             const int newGridsSize = std::max(newRealGridSize, newComplexGridSize);
 396             reallocateDeviceBuffer(&kernelParamsPtr->grid.d_realGrid[gridIndex],
 397                                    newGridsSize,
 398                                    &pmeGpu->archSpecific->realGridSize[gridIndex],
 399                                    &pmeGpu->archSpecific->realGridCapacity[gridIndex],
 400                                    pmeGpu->archSpecific->deviceContext_);
 401             kernelParamsPtr->grid.d_fourierGrid[gridIndex] = kernelParamsPtr->grid.d_realGrid[gridIndex];
 402             pmeGpu->archSpecific->complexGridSize[gridIndex] =
 403                     pmeGpu->archSpecific->realGridSize[gridIndex];
 404             // the size might get used later for copying the grid
 405         }
 406     }
 407 }
 408
 409 void pme_gpu_free_grids(const PmeGpu* pmeGpu)
 410 {
 411     for (int gridIndex = 0; gridIndex < pmeGpu->common->ngrids; gridIndex++)
 412     {
 413         if (pmeGpu->archSpecific->performOutOfPlaceFFT)
 414         {
 415             freeDeviceBuffer(&pmeGpu->kernelParams->grid.d_fourierGrid[gridIndex]);
 416         }
 417         freeDeviceBuffer(&pmeGpu->kernelParams->grid.d_realGrid[gridIndex]);
 418     }
 419 }
 420
 421 void pme_gpu_clear_grids(const PmeGpu* pmeGpu)
 422 {
 423     for (int gridIndex = 0; gridIndex < pmeGpu->common->ngrids; gridIndex++)
 424     {
 425         clearDeviceBufferAsync(&pmeGpu->kernelParams->grid.d_realGrid[gridIndex],
 426                                0,
 427                                pmeGpu->archSpecific->realGridSize[gridIndex],
 428                                pmeGpu->archSpecific->pmeStream_);
 429     }
 430 }
 431
 432 void pme_gpu_realloc_and_copy_fract_shifts(PmeGpu* pmeGpu)
 433 {
 434     pme_gpu_free_fract_shifts(pmeGpu);
 435
 436     auto* kernelParamsPtr = pmeGpu->kernelParams.get();
 437
 438     const int nx                  = kernelParamsPtr->grid.realGridSize[XX];
 439     const int ny                  = kernelParamsPtr->grid.realGridSize[YY];
 440     const int nz                  = kernelParamsPtr->grid.realGridSize[ZZ];
 441     const int cellCount           = c_pmeNeighborUnitcellCount;
 442     const int gridDataOffset[DIM] = { 0, cellCount * nx, cellCount * (nx + ny) };
 443
 444     memcpy(kernelParamsPtr->grid.tablesOffsets, &gridDataOffset, sizeof(gridDataOffset));
 445
 446     const int newFractShiftsSize = cellCount * (nx + ny + nz);
 447
 448     initParamLookupTable(&kernelParamsPtr->grid.d_fractShiftsTable,
 449                          &kernelParamsPtr->fractShiftsTableTexture,
 450                          pmeGpu->common->fsh.data(),
 451                          newFractShiftsSize,
 452                          pmeGpu->archSpecific->deviceContext_);
 453
 454     initParamLookupTable(&kernelParamsPtr->grid.d_gridlineIndicesTable,
 455                          &kernelParamsPtr->gridlineIndicesTableTexture,
 456                          pmeGpu->common->nn.data(),
 457                          newFractShiftsSize,
 458                          pmeGpu->archSpecific->deviceContext_);
 459 }
 460
 461 void pme_gpu_free_fract_shifts(const PmeGpu* pmeGpu)
 462 {
 463     auto* kernelParamsPtr = pmeGpu->kernelParams.get();
 464 #if GMX_GPU_CUDA
 465     destroyParamLookupTable(&kernelParamsPtr->grid.d_fractShiftsTable,
 466                             &kernelParamsPtr->fractShiftsTableTexture);
 467     destroyParamLookupTable(&kernelParamsPtr->grid.d_gridlineIndicesTable,
 468                             &kernelParamsPtr->gridlineIndicesTableTexture);
 469 #elif GMX_GPU_OPENCL || GMX_GPU_SYCL
 470     freeDeviceBuffer(&kernelParamsPtr->grid.d_fractShiftsTable);
 471     freeDeviceBuffer(&kernelParamsPtr->grid.d_gridlineIndicesTable);
 472 #endif
 473 }
 474
 475 bool pme_gpu_stream_query(const PmeGpu* pmeGpu)
 476 {
 477     return haveStreamTasksCompleted(pmeGpu->archSpecific->pmeStream_);
 478 }
 479
 480 void pme_gpu_copy_input_gather_grid(const PmeGpu* pmeGpu, const float* h_grid, const int gridIndex)
 481 {
 482     copyToDeviceBuffer(&pmeGpu->kernelParams->grid.d_realGrid[gridIndex],
 483                        h_grid,
 484                        0,
 485                        pmeGpu->archSpecific->realGridSize[gridIndex],
 486                        pmeGpu->archSpecific->pmeStream_,
 487                        pmeGpu->settings.transferKind,
 488                        nullptr);
 489 }
 490
 491 void pme_gpu_copy_output_spread_grid(const PmeGpu* pmeGpu, float* h_grid, const int gridIndex)
 492 {
 493     copyFromDeviceBuffer(h_grid,
 494                          &pmeGpu->kernelParams->grid.d_realGrid[gridIndex],
 495                          0,
 496                          pmeGpu->archSpecific->realGridSize[gridIndex],
 497                          pmeGpu->archSpecific->pmeStream_,
 498                          pmeGpu->settings.transferKind,
 499                          nullptr);
 500     pmeGpu->archSpecific->syncSpreadGridD2H.markEvent(pmeGpu->archSpecific->pmeStream_);
 501 }
 502
 503 void pme_gpu_copy_output_spread_atom_data(const PmeGpu* pmeGpu)
 504 {
 505     const size_t splinesCount    = DIM * pmeGpu->nAtomsAlloc * pmeGpu->common->pme_order;
 506     auto*        kernelParamsPtr = pmeGpu->kernelParams.get();
 507     copyFromDeviceBuffer(pmeGpu->staging.h_dtheta,
 508                          &kernelParamsPtr->atoms.d_dtheta,
 509                          0,
 510                          splinesCount,
 511                          pmeGpu->archSpecific->pmeStream_,
 512                          pmeGpu->settings.transferKind,
 513                          nullptr);
 514     copyFromDeviceBuffer(pmeGpu->staging.h_theta,
 515                          &kernelParamsPtr->atoms.d_theta,
 516                          0,
 517                          splinesCount,
 518                          pmeGpu->archSpecific->pmeStream_,
 519                          pmeGpu->settings.transferKind,
 520                          nullptr);
 521     copyFromDeviceBuffer(pmeGpu->staging.h_gridlineIndices,
 522                          &kernelParamsPtr->atoms.d_gridlineIndices,
 523                          0,
 524                          kernelParamsPtr->atoms.nAtoms * DIM,
 525                          pmeGpu->archSpecific->pmeStream_,
 526                          pmeGpu->settings.transferKind,
 527                          nullptr);
 528 }
 529
 530 void pme_gpu_copy_input_gather_atom_data(const PmeGpu* pmeGpu)
 531 {
 532     const size_t splinesCount    = DIM * pmeGpu->nAtomsAlloc * pmeGpu->common->pme_order;
 533     auto*        kernelParamsPtr = pmeGpu->kernelParams.get();
 534
 535     // TODO: could clear only the padding and not the whole thing, but this is a test-exclusive code anyway
 536     clearDeviceBufferAsync(&kernelParamsPtr->atoms.d_gridlineIndices,
 537                            0,
 538                            pmeGpu->nAtomsAlloc * DIM,
 539                            pmeGpu->archSpecific->pmeStream_);
 540     clearDeviceBufferAsync(&kernelParamsPtr->atoms.d_dtheta,
 541                            0,
 542                            pmeGpu->nAtomsAlloc * pmeGpu->common->pme_order * DIM,
 543                            pmeGpu->archSpecific->pmeStream_);
 544     clearDeviceBufferAsync(&kernelParamsPtr->atoms.d_theta,
 545                            0,
 546                            pmeGpu->nAtomsAlloc * pmeGpu->common->pme_order * DIM,
 547                            pmeGpu->archSpecific->pmeStream_);
 548
 549     copyToDeviceBuffer(&kernelParamsPtr->atoms.d_dtheta,
 550                        pmeGpu->staging.h_dtheta,
 551                        0,
 552                        splinesCount,
 553                        pmeGpu->archSpecific->pmeStream_,
 554                        pmeGpu->settings.transferKind,
 555                        nullptr);
 556     copyToDeviceBuffer(&kernelParamsPtr->atoms.d_theta,
 557                        pmeGpu->staging.h_theta,
 558                        0,
 559                        splinesCount,
 560                        pmeGpu->archSpecific->pmeStream_,
 561                        pmeGpu->settings.transferKind,
 562                        nullptr);
 563     copyToDeviceBuffer(&kernelParamsPtr->atoms.d_gridlineIndices,
 564                        pmeGpu->staging.h_gridlineIndices,
 565                        0,
 566                        kernelParamsPtr->atoms.nAtoms * DIM,
 567                        pmeGpu->archSpecific->pmeStream_,
 568                        pmeGpu->settings.transferKind,
 569                        nullptr);
 570 }
 571
 572 void pme_gpu_sync_spread_grid(const PmeGpu* pmeGpu)
 573 {
 574     pmeGpu->archSpecific->syncSpreadGridD2H.waitForEvent();
 575 }
 576
 577 /*! \brief Internal GPU initialization for PME.
 578  *
 579  * \param[in]  pmeGpu         GPU PME data.
 580  * \param[in]  deviceContext  GPU context.
 581  * \param[in]  deviceStream   GPU stream.
 582  */
 583 static void pme_gpu_init_internal(PmeGpu* pmeGpu, const DeviceContext& deviceContext, const DeviceStream& deviceStream)
 584 {
 585     /* Allocate the target-specific structures */
 586     pmeGpu->archSpecific.reset(new PmeGpuSpecific(deviceContext, deviceStream));
 587     pmeGpu->kernelParams.reset(new PmeGpuKernelParams());
 588
 589     pmeGpu->archSpecific->performOutOfPlaceFFT = true;
 590     /* This should give better performance, according to the cuFFT documentation.
 591      * The performance seems to be the same though.
 592      * TODO: PME could also try to pick up nice grid sizes (with factors of 2, 3, 5, 7).
 593      */
 594
 595 #if GMX_GPU_CUDA
 596     pmeGpu->maxGridWidthX = deviceContext.deviceInfo().prop.maxGridSize[0];
 597 #else
 598     // Use this path for any non-CUDA GPU acceleration
 599     // TODO: is there no really global work size limit in OpenCL?
 600     pmeGpu->maxGridWidthX = INT32_MAX / 2;
 601 #endif
 602 }
 603
 604 void pme_gpu_reinit_3dfft(const PmeGpu* pmeGpu)
 605 {
 606     if (pme_gpu_settings(pmeGpu).performGPUFFT)
 607     {
 608         pmeGpu->archSpecific->fftSetup.resize(0);
 609         const bool        useDecomposition     = pme_gpu_settings(pmeGpu).useDecomposition;
 610         const bool        performOutOfPlaceFFT = pmeGpu->archSpecific->performOutOfPlaceFFT;
 611         PmeGpuGridParams& grid                 = pme_gpu_get_kernel_params_base_ptr(pmeGpu)->grid;
 612         for (int gridIndex = 0; gridIndex < pmeGpu->common->ngrids; gridIndex++)
 613         {
 614             pmeGpu->archSpecific->fftSetup.push_back(
 615                     std::make_unique<GpuParallel3dFft>(grid.realGridSize,
 616                                                        grid.realGridSizePadded,
 617                                                        grid.complexGridSizePadded,
 618                                                        useDecomposition,
 619                                                        performOutOfPlaceFFT,
 620                                                        pmeGpu->archSpecific->deviceContext_,
 621                                                        pmeGpu->archSpecific->pmeStream_,
 622                                                        grid.d_realGrid[gridIndex],
 623                                                        grid.d_fourierGrid[gridIndex]));
 624         }
 625     }
 626 }
 627
 628 void pme_gpu_destroy_3dfft(const PmeGpu* pmeGpu)
 629 {
 630     pmeGpu->archSpecific->fftSetup.resize(0);
 631 }
 632
 633 void pme_gpu_getEnergyAndVirial(const gmx_pme_t& pme, const float lambda, PmeOutput* output)
 634 {
 635     const PmeGpu* pmeGpu = pme.gpu;
 636
 637     GMX_ASSERT(lambda == 1.0 || pmeGpu->common->ngrids == 2,
 638                "Invalid combination of lambda and number of grids");
 639
 640     for (int gridIndex = 0; gridIndex < pmeGpu->common->ngrids; gridIndex++)
 641     {
 642         for (int j = 0; j < c_virialAndEnergyCount; j++)
 643         {
 644             GMX_ASSERT(std::isfinite(pmeGpu->staging.h_virialAndEnergy[gridIndex][j]),
 645                        "PME GPU produces incorrect energy/virial.");
 646         }
 647     }
 648     for (int dim1 = 0; dim1 < DIM; dim1++)
 649     {
 650         for (int dim2 = 0; dim2 < DIM; dim2++)
 651         {
 652             output->coulombVirial_[dim1][dim2] = 0;
 653         }
 654     }
 655     output->coulombEnergy_ = 0;
 656     float scale            = 1.0;
 657     for (int gridIndex = 0; gridIndex < pmeGpu->common->ngrids; gridIndex++)
 658     {
 659         if (pmeGpu->common->ngrids == 2)
 660         {
 661             scale = gridIndex == 0 ? (1.0 - lambda) : lambda;
 662         }
 663         output->coulombVirial_[XX][XX] +=
 664                 scale * 0.25F * pmeGpu->staging.h_virialAndEnergy[gridIndex][0];
 665         output->coulombVirial_[YY][YY] +=
 666                 scale * 0.25F * pmeGpu->staging.h_virialAndEnergy[gridIndex][1];
 667         output->coulombVirial_[ZZ][ZZ] +=
 668                 scale * 0.25F * pmeGpu->staging.h_virialAndEnergy[gridIndex][2];
 669         output->coulombVirial_[XX][YY] +=
 670                 scale * 0.25F * pmeGpu->staging.h_virialAndEnergy[gridIndex][3];
 671         output->coulombVirial_[YY][XX] +=
 672                 scale * 0.25F * pmeGpu->staging.h_virialAndEnergy[gridIndex][3];
 673         output->coulombVirial_[XX][ZZ] +=
 674                 scale * 0.25F * pmeGpu->staging.h_virialAndEnergy[gridIndex][4];
 675         output->coulombVirial_[ZZ][XX] +=
 676                 scale * 0.25F * pmeGpu->staging.h_virialAndEnergy[gridIndex][4];
 677         output->coulombVirial_[YY][ZZ] +=
 678                 scale * 0.25F * pmeGpu->staging.h_virialAndEnergy[gridIndex][5];
 679         output->coulombVirial_[ZZ][YY] +=
 680                 scale * 0.25F * pmeGpu->staging.h_virialAndEnergy[gridIndex][5];
 681         output->coulombEnergy_ += scale * 0.5F * pmeGpu->staging.h_virialAndEnergy[gridIndex][6];
 682     }
 683     if (pmeGpu->common->ngrids > 1)
 684     {
 685         output->coulombDvdl_ = 0.5F
 686                                * (pmeGpu->staging.h_virialAndEnergy[FEP_STATE_B][6]
 687                                   - pmeGpu->staging.h_virialAndEnergy[FEP_STATE_A][6]);
 688     }
 689 }
 690
 691 /*! \brief Sets the force-related members in \p output
 692  *
 693  * \param[in]   pmeGpu      PME GPU data structure
 694  * \param[out]  output      Pointer to PME output data structure
 695  */
 696 static void pme_gpu_getForceOutput(PmeGpu* pmeGpu, PmeOutput* output)
 697 {
 698     output->haveForceOutput_ = !pmeGpu->settings.useGpuForceReduction;
 699     if (output->haveForceOutput_)
 700     {
 701         output->forces_ = pmeGpu->staging.h_forces;
 702     }
 703 }
 704
 705 PmeOutput pme_gpu_getOutput(const gmx_pme_t& pme, const bool computeEnergyAndVirial, const real lambdaQ)
 706 {
 707     PmeGpu* pmeGpu = pme.gpu;
 708
 709     PmeOutput output;
 710
 711     pme_gpu_getForceOutput(pmeGpu, &output);
 712
 713     if (computeEnergyAndVirial)
 714     {
 715         if (pme_gpu_settings(pmeGpu).performGPUSolve)
 716         {
 717             pme_gpu_getEnergyAndVirial(pme, lambdaQ, &output);
 718         }
 719         else
 720         {
 721             get_pme_ener_vir_q(pme.solve_work, pme.nthread, &output);
 722         }
 723     }
 724     return output;
 725 }
 726
 727 void pme_gpu_update_input_box(PmeGpu gmx_unused* pmeGpu, const matrix gmx_unused box)
 728 {
 729 #if GMX_DOUBLE
 730     GMX_THROW(gmx::NotImplementedError("PME is implemented for single-precision only on GPU"));
 731 #else
 732     matrix scaledBox;
 733     pmeGpu->common->boxScaler->scaleBox(box, scaledBox);
 734     auto* kernelParamsPtr              = pme_gpu_get_kernel_params_base_ptr(pmeGpu);
 735     kernelParamsPtr->current.boxVolume = scaledBox[XX][XX] * scaledBox[YY][YY] * scaledBox[ZZ][ZZ];
 736     GMX_ASSERT(kernelParamsPtr->current.boxVolume != 0.0F, "Zero volume of the unit cell");
 737     matrix recipBox;
 738     gmx::invertBoxMatrix(scaledBox, recipBox);
 739
 740     /* The GPU recipBox is transposed as compared to the CPU recipBox.
 741      * Spread uses matrix columns (while solve and gather use rows).
 742      * There is no particular reason for this; it might be further rethought/optimized for better access patterns.
 743      */
 744     const real newRecipBox[DIM][DIM] = { { recipBox[XX][XX], recipBox[YY][XX], recipBox[ZZ][XX] },
 745                                          { 0.0, recipBox[YY][YY], recipBox[ZZ][YY] },
 746                                          { 0.0, 0.0, recipBox[ZZ][ZZ] } };
 747     memcpy(kernelParamsPtr->current.recipBox, newRecipBox, sizeof(matrix));
 748 #endif
 749 }
 750
 751 /*! \brief \libinternal
 752  * (Re-)initializes all the PME GPU data related to the grid size and cut-off.
 753  *
 754  * \param[in] pmeGpu            The PME GPU structure.
 755  */
 756 static void pme_gpu_reinit_grids(PmeGpu* pmeGpu)
 757 {
 758     auto* kernelParamsPtr = pme_gpu_get_kernel_params_base_ptr(pmeGpu);
 759
 760     GMX_ASSERT(
 761             pmeGpu->common->ngrids == 1 || pmeGpu->common->ngrids == 2,
 762             "Only one (normal Coulomb PME) or two (FEP coulomb PME) PME grids can be used on GPU");
 763
 764     kernelParamsPtr->grid.ewaldFactor =
 765             (M_PI * M_PI) / (pmeGpu->common->ewaldcoeff_q * pmeGpu->common->ewaldcoeff_q);
 766     /* The grid size variants */
 767     for (int i = 0; i < DIM; i++)
 768     {
 769         kernelParamsPtr->grid.realGridSize[i] = pmeGpu->common->nk[i];
 770         kernelParamsPtr->grid.realGridSizeFP[i] =
 771                 static_cast<float>(kernelParamsPtr->grid.realGridSize[i]);
 772         kernelParamsPtr->grid.realGridSizePadded[i] = kernelParamsPtr->grid.realGridSize[i];
 773
 774         // The complex grid currently uses no padding;
 775         // if it starts to do so, then another test should be added for that
 776         kernelParamsPtr->grid.complexGridSize[i]       = kernelParamsPtr->grid.realGridSize[i];
 777         kernelParamsPtr->grid.complexGridSizePadded[i] = kernelParamsPtr->grid.realGridSize[i];
 778     }
 779     /* FFT: n real elements correspond to (n / 2 + 1) complex elements in minor dimension */
 780     if (!pme_gpu_settings(pmeGpu).performGPUFFT)
 781     {
 782         // This allows for GPU spreading grid and CPU fftgrid to have the same layout, so that we can copy the data directly
 783         kernelParamsPtr->grid.realGridSizePadded[ZZ] =
 784                 (kernelParamsPtr->grid.realGridSize[ZZ] / 2 + 1) * 2;
 785     }
 786     /* GPU FFT: n real elements correspond to (n / 2 + 1) complex elements in minor dimension */
 787     kernelParamsPtr->grid.complexGridSize[ZZ] /= 2;
 788     kernelParamsPtr->grid.complexGridSize[ZZ]++;
 789     kernelParamsPtr->grid.complexGridSizePadded[ZZ] = kernelParamsPtr->grid.complexGridSize[ZZ];
 790
 791     pme_gpu_realloc_and_copy_fract_shifts(pmeGpu);
 792     for (int gridIndex = 0; gridIndex < pmeGpu->common->ngrids; gridIndex++)
 793     {
 794         pme_gpu_realloc_and_copy_bspline_values(pmeGpu, gridIndex);
 795     }
 796
 797     pme_gpu_realloc_grids(pmeGpu);
 798     pme_gpu_reinit_3dfft(pmeGpu);
 799 }
 800
 801 /* Several GPU functions that refer to the CPU PME data live here.
 802  * We would like to keep these away from the GPU-framework specific code for clarity,
 803  * as well as compilation issues with MPI.
 804  */
 805
 806 /*! \brief \libinternal
 807  * Copies everything useful from the PME CPU to the PME GPU structure.
 808  * The goal is to minimize interaction with the PME CPU structure in the GPU code.
 809  *
 810  * \param[in] pme         The PME structure.
 811  */
 812 static void pme_gpu_copy_common_data_from(const gmx_pme_t* pme)
 813 {
 814     /* TODO: Consider refactoring the CPU PME code to use the same structure,
 815      * so that this function becomes 2 lines */
 816     PmeGpu* pmeGpu               = pme->gpu;
 817     pmeGpu->common->ngrids       = pme->bFEP_q ? 2 : 1;
 818     pmeGpu->common->epsilon_r    = pme->epsilon_r;
 819     pmeGpu->common->ewaldcoeff_q = pme->ewaldcoeff_q;
 820     pmeGpu->common->nk[XX]       = pme->nkx;
 821     pmeGpu->common->nk[YY]       = pme->nky;
 822     pmeGpu->common->nk[ZZ]       = pme->nkz;
 823     pmeGpu->common->pme_order    = pme->pme_order;
 824     if (pmeGpu->common->pme_order != c_pmeGpuOrder)
 825     {
 826         GMX_THROW(gmx::NotImplementedError("pme_order != 4 is not implemented!"));
 827     }
 828     for (int i = 0; i < DIM; i++)
 829     {
 830         pmeGpu->common->bsp_mod[i].assign(pme->bsp_mod[i], pme->bsp_mod[i] + pmeGpu->common->nk[i]);
 831     }
 832     const int cellCount = c_pmeNeighborUnitcellCount;
 833     pmeGpu->common->fsh.resize(0);
 834     pmeGpu->common->fsh.insert(pmeGpu->common->fsh.end(), pme->fshx, pme->fshx + cellCount * pme->nkx);
 835     pmeGpu->common->fsh.insert(pmeGpu->common->fsh.end(), pme->fshy, pme->fshy + cellCount * pme->nky);
 836     pmeGpu->common->fsh.insert(pmeGpu->common->fsh.end(), pme->fshz, pme->fshz + cellCount * pme->nkz);
 837     pmeGpu->common->nn.resize(0);
 838     pmeGpu->common->nn.insert(pmeGpu->common->nn.end(), pme->nnx, pme->nnx + cellCount * pme->nkx);
 839     pmeGpu->common->nn.insert(pmeGpu->common->nn.end(), pme->nny, pme->nny + cellCount * pme->nky);
 840     pmeGpu->common->nn.insert(pmeGpu->common->nn.end(), pme->nnz, pme->nnz + cellCount * pme->nkz);
 841     pmeGpu->common->runMode       = pme->runMode;
 842     pmeGpu->common->isRankPmeOnly = !pme->bPPnode;
 843     pmeGpu->common->boxScaler     = pme->boxScaler.get();
 844 }
 845
 846 /*! \libinternal \brief
 847  * uses heuristics to select the best performing PME gather and scatter kernels
 848  *
 849  * \param[in,out] pmeGpu         The PME GPU structure.
 850  */
 851 static void pme_gpu_select_best_performing_pme_spreadgather_kernels(PmeGpu* pmeGpu)
 852 {
 853     if (GMX_GPU_CUDA && pmeGpu->kernelParams->atoms.nAtoms > c_pmeGpuPerformanceAtomLimit)
 854     {
 855         pmeGpu->settings.threadsPerAtom     = ThreadsPerAtom::Order;
 856         pmeGpu->settings.recalculateSplines = true;
 857     }
 858     else
 859     {
 860         pmeGpu->settings.threadsPerAtom     = ThreadsPerAtom::OrderSquared;
 861         pmeGpu->settings.recalculateSplines = false;
 862     }
 863 }
 864
 865
 866 /*! \libinternal \brief
 867  * Initializes the PME GPU data at the beginning of the run.
 868  * TODO: this should become PmeGpu::PmeGpu()
 869  *
 870  * \param[in,out] pme            The PME structure.
 871  * \param[in]     deviceContext  The GPU context.
 872  * \param[in]     deviceStream   The GPU stream.
 873  * \param[in,out] pmeGpuProgram  The handle to the program/kernel data created outside (e.g. in unit tests/runner)
 874  */
 875 static void pme_gpu_init(gmx_pme_t*           pme,
 876                          const DeviceContext& deviceContext,
 877                          const DeviceStream&  deviceStream,
 878                          const PmeGpuProgram* pmeGpuProgram)
 879 {
 880     pme->gpu       = new PmeGpu();
 881     PmeGpu* pmeGpu = pme->gpu;
 882     changePinningPolicy(&pmeGpu->staging.h_forces, pme_get_pinning_policy());
 883     pmeGpu->common = std::make_shared<PmeShared>();
 884
 885     /* These settings are set here for the whole run; dynamic ones are set in pme_gpu_reinit() */
 886     /* A convenience variable. */
 887     pmeGpu->settings.useDecomposition = (pme->nnodes != 1);
 888     /* TODO: CPU gather with GPU spread is broken due to different theta/dtheta layout. */
 889     pmeGpu->settings.performGPUGather = true;
 890     // By default GPU-side reduction is off (explicitly set here for tests, otherwise reset per-step)
 891     pmeGpu->settings.useGpuForceReduction = false;
 892
 893     pme_gpu_set_testing(pmeGpu, false);
 894
 895     GMX_ASSERT(pmeGpuProgram != nullptr, "GPU kernels must be already compiled");
 896     pmeGpu->programHandle_ = pmeGpuProgram;
 897
 898     pmeGpu->initializedClfftLibrary_ = std::make_unique<gmx::ClfftInitializer>();
 899
 900     pme_gpu_init_internal(pmeGpu, deviceContext, deviceStream);
 901
 902     pme_gpu_copy_common_data_from(pme);
 903     pme_gpu_alloc_energy_virial(pmeGpu);
 904
 905     GMX_ASSERT(pmeGpu->common->epsilon_r != 0.0F, "PME GPU: bad electrostatic coefficient");
 906
 907     auto* kernelParamsPtr               = pme_gpu_get_kernel_params_base_ptr(pmeGpu);
 908     kernelParamsPtr->constants.elFactor = gmx::c_one4PiEps0 / pmeGpu->common->epsilon_r;
 909 }
 910
 911 void pme_gpu_get_real_grid_sizes(const PmeGpu* pmeGpu, gmx::IVec* gridSize, gmx::IVec* paddedGridSize)
 912 {
 913     GMX_ASSERT(gridSize != nullptr, "");
 914     GMX_ASSERT(paddedGridSize != nullptr, "");
 915     GMX_ASSERT(pmeGpu != nullptr, "");
 916     auto* kernelParamsPtr = pme_gpu_get_kernel_params_base_ptr(pmeGpu);
 917     for (int i = 0; i < DIM; i++)
 918     {
 919         (*gridSize)[i]       = kernelParamsPtr->grid.realGridSize[i];
 920         (*paddedGridSize)[i] = kernelParamsPtr->grid.realGridSizePadded[i];
 921     }
 922 }
 923
 924 void pme_gpu_reinit(gmx_pme_t*           pme,
 925                     const DeviceContext* deviceContext,
 926                     const DeviceStream*  deviceStream,
 927                     const PmeGpuProgram* pmeGpuProgram)
 928 {
 929     GMX_ASSERT(pme != nullptr, "Need valid PME object");
 930
 931     if (!pme->gpu)
 932     {
 933         GMX_RELEASE_ASSERT(deviceContext != nullptr,
 934                            "Device context can not be nullptr when setting up PME on GPU.");
 935         GMX_RELEASE_ASSERT(deviceStream != nullptr,
 936                            "Device stream can not be nullptr when setting up PME on GPU.");
 937         /* First-time initialization */
 938         pme_gpu_init(pme, *deviceContext, *deviceStream, pmeGpuProgram);
 939     }
 940     else
 941     {
 942         /* After this call nothing in the GPU code should refer to the gmx_pme_t *pme itself - until the next pme_gpu_reinit */
 943         pme_gpu_copy_common_data_from(pme);
 944     }
 945     /* GPU FFT will only get used for a single rank.*/
 946     pme->gpu->settings.performGPUFFT =
 947             (pme->gpu->common->runMode == PmeRunMode::GPU) && !pme->gpu->settings.useDecomposition;
 948     pme->gpu->settings.performGPUSolve = (pme->gpu->common->runMode == PmeRunMode::GPU);
 949
 950     /* Reinit active timers */
 951     pme_gpu_reinit_timings(pme->gpu);
 952
 953     pme_gpu_reinit_grids(pme->gpu);
 954     // Note: if timing the reinit launch overhead becomes more relevant
 955     // (e.g. with regulat PP-PME re-balancing), we should pass wcycle here.
 956     pme_gpu_reinit_computation(pme, nullptr);
 957     /* Clear the previous box - doesn't hurt, and forces the PME CPU recipbox
 958      * update for mixed mode on grid switch. TODO: use shared recipbox field.
 959      */
 960     std::memset(pme->gpu->common->previousBox, 0, sizeof(pme->gpu->common->previousBox));
 961 }
 962
 963 void pme_gpu_destroy(PmeGpu* pmeGpu)
 964 {
 965     /* Free lots of data */
 966     pme_gpu_free_energy_virial(pmeGpu);
 967     pme_gpu_free_bspline_values(pmeGpu);
 968     pme_gpu_free_forces(pmeGpu);
 969     pme_gpu_free_coefficients(pmeGpu);
 970     pme_gpu_free_spline_data(pmeGpu);
 971     pme_gpu_free_grid_indices(pmeGpu);
 972     pme_gpu_free_fract_shifts(pmeGpu);
 973     pme_gpu_free_grids(pmeGpu);
 974
 975     pme_gpu_destroy_3dfft(pmeGpu);
 976
 977     delete pmeGpu;
 978 }
 979
 980 void pme_gpu_reinit_atoms(PmeGpu* pmeGpu, const int nAtoms, const real* chargesA, const real* chargesB)
 981 {
 982     auto* kernelParamsPtr         = pme_gpu_get_kernel_params_base_ptr(pmeGpu);
 983     kernelParamsPtr->atoms.nAtoms = nAtoms;
 984     const int  block_size         = pme_gpu_get_atom_data_block_size();
 985     const int  nAtomsNewPadded    = ((nAtoms + block_size - 1) / block_size) * block_size;
 986     const bool haveToRealloc      = (pmeGpu->nAtomsAlloc < nAtomsNewPadded);
 987     pmeGpu->nAtomsAlloc           = nAtomsNewPadded;
 988
 989 #if GMX_DOUBLE
 990     GMX_RELEASE_ASSERT(false, "Only single precision supported");
 991     GMX_UNUSED_VALUE(charges);
 992 #else
 993     int gridIndex = 0;
 994     /* Could also be checked for haveToRealloc, but the copy always needs to be performed */
 995     pme_gpu_realloc_and_copy_input_coefficients(pmeGpu, reinterpret_cast<const float*>(chargesA), gridIndex);
 996     gridIndex++;
 997     if (chargesB != nullptr)
 998     {
 999         pme_gpu_realloc_and_copy_input_coefficients(
1000                 pmeGpu, reinterpret_cast<const float*>(chargesB), gridIndex);
1001     }
1002     else
1003     {
1004         /* Fill the second set of coefficients with chargesA as well to be able to avoid
1005          * conditionals in the GPU kernels */
1006         /* FIXME: This should be avoided by making a separate templated version of the
1007          * relevant kernel(s) (probably only pme_gather_kernel). That would require a
1008          * reduction of the current number of templated parameters of that kernel. */
1009         pme_gpu_realloc_and_copy_input_coefficients(
1010                 pmeGpu, reinterpret_cast<const float*>(chargesA), gridIndex);
1011     }
1012 #endif
1013
1014     if (haveToRealloc)
1015     {
1016         pme_gpu_realloc_forces(pmeGpu);
1017         pme_gpu_realloc_spline_data(pmeGpu);
1018         pme_gpu_realloc_grid_indices(pmeGpu);
1019     }
1020     pme_gpu_select_best_performing_pme_spreadgather_kernels(pmeGpu);
1021 }
1022
1023 /*! \internal \brief
1024  * Returns raw timing event from the corresponding GpuRegionTimer (if timings are enabled).
1025  * In CUDA result can be nullptr stub, per GpuRegionTimer implementation.
1026  *
1027  * \param[in] pmeGpu         The PME GPU data structure.
1028  * \param[in] pmeStageId     The PME GPU stage gtPME_ index from the enum in src/gromacs/timing/gpu_timing.h
1029  */
1030 static CommandEvent* pme_gpu_fetch_timing_event(const PmeGpu* pmeGpu, PmeStage pmeStageId)
1031 {
1032     CommandEvent* timingEvent = nullptr;
1033     if (pme_gpu_timings_enabled(pmeGpu))
1034     {
1035         GMX_ASSERT(pmeStageId < PmeStage::Count, "Wrong PME GPU timing event index");
1036         timingEvent = pmeGpu->archSpecific->timingEvents[pmeStageId].fetchNextEvent();
1037     }
1038     return timingEvent;
1039 }
1040
1041 void pme_gpu_3dfft(const PmeGpu* pmeGpu, gmx_fft_direction dir, const int grid_index)
1042 {
1043     PmeStage timerId = (dir == GMX_FFT_REAL_TO_COMPLEX) ? PmeStage::FftTransformR2C
1044                                                         : PmeStage::FftTransformC2R;
1045
1046     pme_gpu_start_timing(pmeGpu, timerId);
1047     pmeGpu->archSpecific->fftSetup[grid_index]->perform3dFft(
1048             dir, pme_gpu_fetch_timing_event(pmeGpu, timerId));
1049     pme_gpu_stop_timing(pmeGpu, timerId);
1050 }
1051
1052 /*! \brief
1053  * Given possibly large \p blockCount, returns a compact 1D or 2D grid for kernel scheduling,
1054  * to minimize number of unused blocks.
1055  */
1056 std::pair<int, int> inline pmeGpuCreateGrid(const PmeGpu* pmeGpu, int blockCount)
1057 {
1058     // How many maximum widths in X do we need (hopefully just one)
1059     const int minRowCount = (blockCount + pmeGpu->maxGridWidthX - 1) / pmeGpu->maxGridWidthX;
1060     // Trying to make things even
1061     const int colCount = (blockCount + minRowCount - 1) / minRowCount;
1062     GMX_ASSERT((colCount * minRowCount - blockCount) >= 0, "pmeGpuCreateGrid: totally wrong");
1063     GMX_ASSERT((colCount * minRowCount - blockCount) < minRowCount,
1064                "pmeGpuCreateGrid: excessive blocks");
1065     return std::pair<int, int>(colCount, minRowCount);
1066 }
1067
1068 /*! \brief
1069  * Returns a pointer to appropriate spline and spread kernel based on the input bool values
1070  *
1071  * \param[in]  pmeGpu                   The PME GPU structure.
1072  * \param[in]  threadsPerAtom           Controls whether we should use order or order*order threads per atom
1073  * \param[in]  writeSplinesToGlobal     bool controlling if we should write spline data to global memory
1074  * \param[in]  numGrids                 Number of grids to use. numGrids == 2 if Coulomb is perturbed.
1075  *
1076  * \return Pointer to CUDA kernel
1077  */
1078 static auto selectSplineAndSpreadKernelPtr(const PmeGpu*  pmeGpu,
1079                                            ThreadsPerAtom threadsPerAtom,
1080                                            bool           writeSplinesToGlobal,
1081                                            const int      numGrids)
1082 {
1083     PmeGpuProgramImpl::PmeKernelHandle kernelPtr = nullptr;
1084     if (writeSplinesToGlobal)
1085     {
1086         if (threadsPerAtom == ThreadsPerAtom::Order)
1087         {
1088             if (numGrids == 2)
1089             {
1090                 kernelPtr = pmeGpu->programHandle_->impl_->splineAndSpreadKernelWriteSplinesThPerAtom4Dual;
1091             }
1092             else
1093             {
1094                 kernelPtr = pmeGpu->programHandle_->impl_->splineAndSpreadKernelWriteSplinesThPerAtom4Single;
1095             }
1096         }
1097         else
1098         {
1099             if (numGrids == 2)
1100             {
1101                 kernelPtr = pmeGpu->programHandle_->impl_->splineAndSpreadKernelWriteSplinesDual;
1102             }
1103             else
1104             {
1105                 kernelPtr = pmeGpu->programHandle_->impl_->splineAndSpreadKernelWriteSplinesSingle;
1106             }
1107         }
1108     }
1109     else
1110     {
1111         if (threadsPerAtom == ThreadsPerAtom::Order)
1112         {
1113             if (numGrids == 2)
1114             {
1115                 kernelPtr = pmeGpu->programHandle_->impl_->splineAndSpreadKernelThPerAtom4Dual;
1116             }
1117             else
1118             {
1119                 kernelPtr = pmeGpu->programHandle_->impl_->splineAndSpreadKernelThPerAtom4Single;
1120             }
1121         }
1122         else
1123         {
1124             if (numGrids == 2)
1125             {
1126                 kernelPtr = pmeGpu->programHandle_->impl_->splineAndSpreadKernelDual;
1127             }
1128             else
1129             {
1130                 kernelPtr = pmeGpu->programHandle_->impl_->splineAndSpreadKernelSingle;
1131             }
1132         }
1133     }
1134
1135     return kernelPtr;
1136 }
1137
1138 /*! \brief
1139  * Returns a pointer to appropriate spline kernel based on the input bool values
1140  *
1141  * \param[in]  pmeGpu                   The PME GPU structure.
1142  * \param[in]  threadsPerAtom           Controls whether we should use order or order*order threads per atom
1143  * \param[in]  writeSplinesToGlobal     bool controlling if we should write spline data to global memory
1144  * \param[in]  numGrids                 Number of grids to use. numGrids == 2 if Coulomb is perturbed.
1145  *
1146  * \return Pointer to CUDA kernel
1147  */
1148 static auto selectSplineKernelPtr(const PmeGpu*   pmeGpu,
1149                                   ThreadsPerAtom  threadsPerAtom,
1150                                   bool gmx_unused writeSplinesToGlobal,
1151                                   const int       numGrids)
1152 {
1153     PmeGpuProgramImpl::PmeKernelHandle kernelPtr = nullptr;
1154     GMX_ASSERT(
1155             writeSplinesToGlobal,
1156             "Spline data should always be written to global memory when just calculating splines");
1157
1158     if (threadsPerAtom == ThreadsPerAtom::Order)
1159     {
1160         if (numGrids == 2)
1161         {
1162             kernelPtr = pmeGpu->programHandle_->impl_->splineKernelThPerAtom4Dual;
1163         }
1164         else
1165         {
1166             kernelPtr = pmeGpu->programHandle_->impl_->splineKernelThPerAtom4Single;
1167         }
1168     }
1169     else
1170     {
1171         if (numGrids == 2)
1172         {
1173             kernelPtr = pmeGpu->programHandle_->impl_->splineKernelDual;
1174         }
1175         else
1176         {
1177             kernelPtr = pmeGpu->programHandle_->impl_->splineKernelSingle;
1178         }
1179     }
1180     return kernelPtr;
1181 }
1182
1183 /*! \brief
1184  * Returns a pointer to appropriate spread kernel based on the input bool values
1185  *
1186  * \param[in]  pmeGpu                   The PME GPU structure.
1187  * \param[in]  threadsPerAtom           Controls whether we should use order or order*order threads per atom
1188  * \param[in]  writeSplinesToGlobal     bool controlling if we should write spline data to global memory
1189  * \param[in]  numGrids                 Number of grids to use. numGrids == 2 if Coulomb is perturbed.
1190  *
1191  * \return Pointer to CUDA kernel
1192  */
1193 static auto selectSpreadKernelPtr(const PmeGpu*  pmeGpu,
1194                                   ThreadsPerAtom threadsPerAtom,
1195                                   bool           writeSplinesToGlobal,
1196                                   const int      numGrids)
1197 {
1198     PmeGpuProgramImpl::PmeKernelHandle kernelPtr = nullptr;
1199     if (writeSplinesToGlobal)
1200     {
1201         if (threadsPerAtom == ThreadsPerAtom::Order)
1202         {
1203             if (numGrids == 2)
1204             {
1205                 kernelPtr = pmeGpu->programHandle_->impl_->spreadKernelThPerAtom4Dual;
1206             }
1207             {
1208                 kernelPtr = pmeGpu->programHandle_->impl_->spreadKernelThPerAtom4Single;
1209             }
1210         }
1211         else
1212         {
1213             if (numGrids == 2)
1214             {
1215                 kernelPtr = pmeGpu->programHandle_->impl_->spreadKernelDual;
1216             }
1217             else
1218             {
1219                 kernelPtr = pmeGpu->programHandle_->impl_->spreadKernelSingle;
1220             }
1221         }
1222     }
1223     else
1224     {
1225         /* if we are not saving the spline data we need to recalculate it
1226            using the spline and spread Kernel */
1227         if (threadsPerAtom == ThreadsPerAtom::Order)
1228         {
1229             if (numGrids == 2)
1230             {
1231                 kernelPtr = pmeGpu->programHandle_->impl_->splineAndSpreadKernelThPerAtom4Dual;
1232             }
1233             else
1234             {
1235                 kernelPtr = pmeGpu->programHandle_->impl_->splineAndSpreadKernelThPerAtom4Single;
1236             }
1237         }
1238         else
1239         {
1240             if (numGrids == 2)
1241             {
1242                 kernelPtr = pmeGpu->programHandle_->impl_->splineAndSpreadKernelDual;
1243             }
1244             else
1245             {
1246                 kernelPtr = pmeGpu->programHandle_->impl_->splineAndSpreadKernelSingle;
1247             }
1248         }
1249     }
1250     return kernelPtr;
1251 }
1252
1253 void pme_gpu_spread(const PmeGpu*         pmeGpu,
1254                     GpuEventSynchronizer* xReadyOnDevice,
1255                     real**                h_grids,
1256                     bool                  computeSplines,
1257                     bool                  spreadCharges,
1258                     const real            lambda)
1259 {
1260     GMX_ASSERT(
1261             pmeGpu->common->ngrids == 1 || pmeGpu->common->ngrids == 2,
1262             "Only one (normal Coulomb PME) or two (FEP coulomb PME) PME grids can be used on GPU");
1263
1264     GMX_ASSERT(computeSplines || spreadCharges,
1265                "PME spline/spread kernel has invalid input (nothing to do)");
1266     auto* kernelParamsPtr = pmeGpu->kernelParams.get();
1267     GMX_ASSERT(kernelParamsPtr->atoms.nAtoms > 0, "No atom data in PME GPU spread");
1268
1269     const size_t blockSize = pmeGpu->programHandle_->impl_->spreadWorkGroupSize;
1270
1271     const int order = pmeGpu->common->pme_order;
1272     GMX_ASSERT(order == c_pmeGpuOrder, "Only PME order 4 is implemented");
1273     const bool writeGlobal = pmeGpu->settings.copyAllOutputs;
1274     const int  threadsPerAtom =
1275             (pmeGpu->settings.threadsPerAtom == ThreadsPerAtom::Order ? order : order * order);
1276     const bool recalculateSplines = pmeGpu->settings.recalculateSplines;
1277
1278     GMX_ASSERT(!GMX_GPU_OPENCL || pmeGpu->settings.threadsPerAtom == ThreadsPerAtom::OrderSquared,
1279                "Only 16 threads per atom supported in OpenCL");
1280     GMX_ASSERT(!GMX_GPU_OPENCL || !recalculateSplines,
1281                "Recalculating splines not supported in OpenCL");
1282
1283     const int atomsPerBlock = blockSize / threadsPerAtom;
1284
1285     // TODO: pick smaller block size in runtime if needed
1286     // (e.g. on 660 Ti where 50% occupancy is ~25% faster than 100% occupancy with RNAse (~17.8k atoms))
1287     // If doing so, change atomsPerBlock in the kernels as well.
1288     // TODO: test varying block sizes on modern arch-s as well
1289     // TODO: also consider using cudaFuncSetCacheConfig() for preferring shared memory on older architectures
1290     //(for spline data mostly)
1291     GMX_ASSERT(!(c_pmeAtomDataBlockSize % atomsPerBlock),
1292                "inconsistent atom data padding vs. spreading block size");
1293
1294     // Ensure that coordinates are ready on the device before launching spread;
1295     // only needed with CUDA on PP+PME ranks, not on separate PME ranks, in unit tests
1296     // nor in OpenCL as these cases use a single stream (hence xReadyOnDevice == nullptr).
1297     GMX_ASSERT(!GMX_GPU_CUDA || xReadyOnDevice != nullptr || pmeGpu->common->isRankPmeOnly
1298                        || pme_gpu_settings(pmeGpu).copyAllOutputs,
1299                "Need a valid coordinate synchronizer on PP+PME ranks with CUDA.");
1300
1301     if (xReadyOnDevice)
1302     {
1303         xReadyOnDevice->enqueueWaitEvent(pmeGpu->archSpecific->pmeStream_);
1304     }
1305
1306     const int blockCount = pmeGpu->nAtomsAlloc / atomsPerBlock;
1307     auto      dimGrid    = pmeGpuCreateGrid(pmeGpu, blockCount);
1308
1309     if (pmeGpu->common->ngrids == 1)
1310     {
1311         kernelParamsPtr->current.scale = 1.0;
1312     }
1313     else
1314     {
1315         kernelParamsPtr->current.scale = 1.0 - lambda;
1316     }
1317
1318     KernelLaunchConfig config;
1319     config.blockSize[0] = order;
1320     config.blockSize[1] = (pmeGpu->settings.threadsPerAtom == ThreadsPerAtom::Order ? 1 : order);
1321     config.blockSize[2] = atomsPerBlock;
1322     config.gridSize[0]  = dimGrid.first;
1323     config.gridSize[1]  = dimGrid.second;
1324
1325     PmeStage                           timingId;
1326     PmeGpuProgramImpl::PmeKernelHandle kernelPtr = nullptr;
1327     if (computeSplines)
1328     {
1329         if (spreadCharges)
1330         {
1331             timingId  = PmeStage::SplineAndSpread;
1332             kernelPtr = selectSplineAndSpreadKernelPtr(pmeGpu,
1333                                                        pmeGpu->settings.threadsPerAtom,
1334                                                        writeGlobal || (!recalculateSplines),
1335                                                        pmeGpu->common->ngrids);
1336         }
1337         else
1338         {
1339             timingId  = PmeStage::Spline;
1340             kernelPtr = selectSplineKernelPtr(pmeGpu,
1341                                               pmeGpu->settings.threadsPerAtom,
1342                                               writeGlobal || (!recalculateSplines),
1343                                               pmeGpu->common->ngrids);
1344         }
1345     }
1346     else
1347     {
1348         timingId  = PmeStage::Spread;
1349         kernelPtr = selectSpreadKernelPtr(pmeGpu,
1350                                           pmeGpu->settings.threadsPerAtom,
1351                                           writeGlobal || (!recalculateSplines),
1352                                           pmeGpu->common->ngrids);
1353     }
1354
1355
1356     pme_gpu_start_timing(pmeGpu, timingId);
1357     auto* timingEvent = pme_gpu_fetch_timing_event(pmeGpu, timingId);
1358 #if c_canEmbedBuffers
1359     const auto kernelArgs = prepareGpuKernelArguments(kernelPtr, config, kernelParamsPtr);
1360 #else
1361     const auto kernelArgs =
1362             prepareGpuKernelArguments(kernelPtr,
1363                                       config,
1364                                       kernelParamsPtr,
1365                                       &kernelParamsPtr->atoms.d_theta,
1366                                       &kernelParamsPtr->atoms.d_dtheta,
1367                                       &kernelParamsPtr->atoms.d_gridlineIndices,
1368                                       &kernelParamsPtr->grid.d_realGrid[FEP_STATE_A],
1369                                       &kernelParamsPtr->grid.d_realGrid[FEP_STATE_B],
1370                                       &kernelParamsPtr->grid.d_fractShiftsTable,
1371                                       &kernelParamsPtr->grid.d_gridlineIndicesTable,
1372                                       &kernelParamsPtr->atoms.d_coefficients[FEP_STATE_A],
1373                                       &kernelParamsPtr->atoms.d_coefficients[FEP_STATE_B],
1374                                       &kernelParamsPtr->atoms.d_coordinates);
1375 #endif
1376
1377     launchGpuKernel(
1378             kernelPtr, config, pmeGpu->archSpecific->pmeStream_, timingEvent, "PME spline/spread", kernelArgs);
1379     pme_gpu_stop_timing(pmeGpu, timingId);
1380
1381     const auto& settings    = pmeGpu->settings;
1382     const bool copyBackGrid = spreadCharges && (!settings.performGPUFFT || settings.copyAllOutputs);
1383     if (copyBackGrid)
1384     {
1385         for (int gridIndex = 0; gridIndex < pmeGpu->common->ngrids; gridIndex++)
1386         {
1387             float* h_grid = h_grids[gridIndex];
1388             pme_gpu_copy_output_spread_grid(pmeGpu, h_grid, gridIndex);
1389         }
1390     }
1391     const bool copyBackAtomData =
1392             computeSplines && (!settings.performGPUGather || settings.copyAllOutputs);
1393     if (copyBackAtomData)
1394     {
1395         pme_gpu_copy_output_spread_atom_data(pmeGpu);
1396     }
1397 }
1398
1399 void pme_gpu_solve(const PmeGpu* pmeGpu,
1400                    const int     gridIndex,
1401                    t_complex*    h_grid,
1402                    GridOrdering  gridOrdering,
1403                    bool          computeEnergyAndVirial)
1404 {
1405     GMX_ASSERT(
1406             pmeGpu->common->ngrids == 1 || pmeGpu->common->ngrids == 2,
1407             "Only one (normal Coulomb PME) or two (FEP coulomb PME) PME grids can be used on GPU");
1408     GMX_ASSERT(gridIndex < pmeGpu->common->ngrids,
1409                "Invalid combination of gridIndex and number of grids");
1410
1411     const auto& settings               = pmeGpu->settings;
1412     const bool  copyInputAndOutputGrid = !settings.performGPUFFT || settings.copyAllOutputs;
1413
1414     auto* kernelParamsPtr = pmeGpu->kernelParams.get();
1415
1416     float* h_gridFloat = reinterpret_cast<float*>(h_grid);
1417     if (copyInputAndOutputGrid)
1418     {
1419         copyToDeviceBuffer(&kernelParamsPtr->grid.d_fourierGrid[gridIndex],
1420                            h_gridFloat,
1421                            0,
1422                            pmeGpu->archSpecific->complexGridSize[gridIndex],
1423                            pmeGpu->archSpecific->pmeStream_,
1424                            pmeGpu->settings.transferKind,
1425                            nullptr);
1426     }
1427
1428     int majorDim = -1, middleDim = -1, minorDim = -1;
1429     switch (gridOrdering)
1430     {
1431         case GridOrdering::YZX:
1432             majorDim  = YY;
1433             middleDim = ZZ;
1434             minorDim  = XX;
1435             break;
1436
1437         case GridOrdering::XYZ:
1438             majorDim  = XX;
1439             middleDim = YY;
1440             minorDim  = ZZ;
1441             break;
1442
1443         default: GMX_ASSERT(false, "Implement grid ordering here and below for the kernel launch");
1444     }
1445
1446     const int maxBlockSize = pmeGpu->programHandle_->impl_->solveMaxWorkGroupSize;
1447
1448     const int gridLineSize      = pmeGpu->kernelParams->grid.complexGridSize[minorDim];
1449     const int gridLinesPerBlock = std::max(maxBlockSize / gridLineSize, 1);
1450     const int blocksPerGridLine = (gridLineSize + maxBlockSize - 1) / maxBlockSize;
1451     int       cellsPerBlock;
1452     if (blocksPerGridLine == 1)
1453     {
1454         cellsPerBlock = gridLineSize * gridLinesPerBlock;
1455     }
1456     else
1457     {
1458         cellsPerBlock = (gridLineSize + blocksPerGridLine - 1) / blocksPerGridLine;
1459     }
1460     const int warpSize  = pmeGpu->programHandle_->warpSize();
1461     const int blockSize = (cellsPerBlock + warpSize - 1) / warpSize * warpSize;
1462
1463     static_assert(!GMX_GPU_CUDA || c_solveMaxWarpsPerBlock / 2 >= 4,
1464                   "The CUDA solve energy kernels needs at least 4 warps. "
1465                   "Here we launch at least half of the max warps.");
1466
1467     KernelLaunchConfig config;
1468     config.blockSize[0] = blockSize;
1469     config.gridSize[0]  = blocksPerGridLine;
1470     // rounding up to full warps so that shuffle operations produce defined results
1471     config.gridSize[1] = (pmeGpu->kernelParams->grid.complexGridSize[middleDim] + gridLinesPerBlock - 1)
1472                          / gridLinesPerBlock;
1473     config.gridSize[2] = pmeGpu->kernelParams->grid.complexGridSize[majorDim];
1474
1475     PmeStage                           timingId  = PmeStage::Solve;
1476     PmeGpuProgramImpl::PmeKernelHandle kernelPtr = nullptr;
1477     if (gridOrdering == GridOrdering::YZX)
1478     {
1479         if (gridIndex == 0)
1480         {
1481             kernelPtr = computeEnergyAndVirial ? pmeGpu->programHandle_->impl_->solveYZXEnergyKernelA
1482                                                : pmeGpu->programHandle_->impl_->solveYZXKernelA;
1483         }
1484         else
1485         {
1486             kernelPtr = computeEnergyAndVirial ? pmeGpu->programHandle_->impl_->solveYZXEnergyKernelB
1487                                                : pmeGpu->programHandle_->impl_->solveYZXKernelB;
1488         }
1489     }
1490     else if (gridOrdering == GridOrdering::XYZ)
1491     {
1492         if (gridIndex == 0)
1493         {
1494             kernelPtr = computeEnergyAndVirial ? pmeGpu->programHandle_->impl_->solveXYZEnergyKernelA
1495                                                : pmeGpu->programHandle_->impl_->solveXYZKernelA;
1496         }
1497         else
1498         {
1499             kernelPtr = computeEnergyAndVirial ? pmeGpu->programHandle_->impl_->solveXYZEnergyKernelB
1500                                                : pmeGpu->programHandle_->impl_->solveXYZKernelB;
1501         }
1502     }
1503
1504     pme_gpu_start_timing(pmeGpu, timingId);
1505     auto* timingEvent = pme_gpu_fetch_timing_event(pmeGpu, timingId);
1506 #if c_canEmbedBuffers
1507     const auto kernelArgs = prepareGpuKernelArguments(kernelPtr, config, kernelParamsPtr);
1508 #else
1509     const auto kernelArgs =
1510             prepareGpuKernelArguments(kernelPtr,
1511                                       config,
1512                                       kernelParamsPtr,
1513                                       &kernelParamsPtr->grid.d_splineModuli[gridIndex],
1514                                       &kernelParamsPtr->constants.d_virialAndEnergy[gridIndex],
1515                                       &kernelParamsPtr->grid.d_fourierGrid[gridIndex]);
1516 #endif
1517     launchGpuKernel(kernelPtr, config, pmeGpu->archSpecific->pmeStream_, timingEvent, "PME solve", kernelArgs);
1518     pme_gpu_stop_timing(pmeGpu, timingId);
1519
1520     if (computeEnergyAndVirial)
1521     {
1522         copyFromDeviceBuffer(pmeGpu->staging.h_virialAndEnergy[gridIndex],
1523                              &kernelParamsPtr->constants.d_virialAndEnergy[gridIndex],
1524                              0,
1525                              c_virialAndEnergyCount,
1526                              pmeGpu->archSpecific->pmeStream_,
1527                              pmeGpu->settings.transferKind,
1528                              nullptr);
1529     }
1530
1531     if (copyInputAndOutputGrid)
1532     {
1533         copyFromDeviceBuffer(h_gridFloat,
1534                              &kernelParamsPtr->grid.d_fourierGrid[gridIndex],
1535                              0,
1536                              pmeGpu->archSpecific->complexGridSize[gridIndex],
1537                              pmeGpu->archSpecific->pmeStream_,
1538                              pmeGpu->settings.transferKind,
1539                              nullptr);
1540     }
1541 }
1542
1543 /*! \brief
1544  * Returns a pointer to appropriate gather kernel based on the inputvalues
1545  *
1546  * \param[in]  pmeGpu                   The PME GPU structure.
1547  * \param[in]  threadsPerAtom           Controls whether we should use order or order*order threads per atom
1548  * \param[in]  readSplinesFromGlobal    bool controlling if we should write spline data to global memory
1549  * \param[in]  numGrids                 Number of grids to use. numGrids == 2 if Coulomb is perturbed.
1550  *
1551  * \return Pointer to CUDA kernel
1552  */
1553 inline auto selectGatherKernelPtr(const PmeGpu*  pmeGpu,
1554                                   ThreadsPerAtom threadsPerAtom,
1555                                   bool           readSplinesFromGlobal,
1556                                   const int      numGrids)
1557
1558 {
1559     PmeGpuProgramImpl::PmeKernelHandle kernelPtr = nullptr;
1560
1561     if (readSplinesFromGlobal)
1562     {
1563         if (threadsPerAtom == ThreadsPerAtom::Order)
1564         {
1565             if (numGrids == 2)
1566             {
1567                 kernelPtr = pmeGpu->programHandle_->impl_->gatherKernelReadSplinesThPerAtom4Dual;
1568             }
1569             else
1570             {
1571                 kernelPtr = pmeGpu->programHandle_->impl_->gatherKernelReadSplinesThPerAtom4Single;
1572             }
1573         }
1574         else
1575         {
1576             if (numGrids == 2)
1577             {
1578                 kernelPtr = pmeGpu->programHandle_->impl_->gatherKernelReadSplinesDual;
1579             }
1580             else
1581             {
1582                 kernelPtr = pmeGpu->programHandle_->impl_->gatherKernelReadSplinesSingle;
1583             }
1584         }
1585     }
1586     else
1587     {
1588         if (threadsPerAtom == ThreadsPerAtom::Order)
1589         {
1590             if (numGrids == 2)
1591             {
1592                 kernelPtr = pmeGpu->programHandle_->impl_->gatherKernelThPerAtom4Dual;
1593             }
1594             else
1595             {
1596                 kernelPtr = pmeGpu->programHandle_->impl_->gatherKernelThPerAtom4Single;
1597             }
1598         }
1599         else
1600         {
1601             if (numGrids == 2)
1602             {
1603                 kernelPtr = pmeGpu->programHandle_->impl_->gatherKernelDual;
1604             }
1605             else
1606             {
1607                 kernelPtr = pmeGpu->programHandle_->impl_->gatherKernelSingle;
1608             }
1609         }
1610     }
1611     return kernelPtr;
1612 }
1613
1614 void pme_gpu_gather(PmeGpu* pmeGpu, real** h_grids, const float lambda)
1615 {
1616     GMX_ASSERT(
1617             pmeGpu->common->ngrids == 1 || pmeGpu->common->ngrids == 2,
1618             "Only one (normal Coulomb PME) or two (FEP coulomb PME) PME grids can be used on GPU");
1619
1620     const auto& settings = pmeGpu->settings;
1621
1622     if (!settings.performGPUFFT || settings.copyAllOutputs)
1623     {
1624         for (int gridIndex = 0; gridIndex < pmeGpu->common->ngrids; gridIndex++)
1625         {
1626             float* h_grid = const_cast<float*>(h_grids[gridIndex]);
1627             pme_gpu_copy_input_gather_grid(pmeGpu, h_grid, gridIndex);
1628         }
1629     }
1630
1631     if (settings.copyAllOutputs)
1632     {
1633         pme_gpu_copy_input_gather_atom_data(pmeGpu);
1634     }
1635
1636     /* Set if we have unit tests */
1637     const bool   readGlobal = pmeGpu->settings.copyAllOutputs;
1638     const size_t blockSize  = pmeGpu->programHandle_->impl_->gatherWorkGroupSize;
1639     const int    order      = pmeGpu->common->pme_order;
1640     GMX_ASSERT(order == c_pmeGpuOrder, "Only PME order 4 is implemented");
1641     const int threadsPerAtom =
1642             (pmeGpu->settings.threadsPerAtom == ThreadsPerAtom::Order ? order : order * order);
1643     const bool recalculateSplines = pmeGpu->settings.recalculateSplines;
1644
1645     GMX_ASSERT(!GMX_GPU_OPENCL || pmeGpu->settings.threadsPerAtom == ThreadsPerAtom::OrderSquared,
1646                "Only 16 threads per atom supported in OpenCL");
1647     GMX_ASSERT(!GMX_GPU_OPENCL || !recalculateSplines,
1648                "Recalculating splines not supported in OpenCL");
1649
1650     const int atomsPerBlock = blockSize / threadsPerAtom;
1651
1652     GMX_ASSERT(!(c_pmeAtomDataBlockSize % atomsPerBlock),
1653                "inconsistent atom data padding vs. gathering block size");
1654
1655     const int blockCount = pmeGpu->nAtomsAlloc / atomsPerBlock;
1656     auto      dimGrid    = pmeGpuCreateGrid(pmeGpu, blockCount);
1657
1658     KernelLaunchConfig config;
1659     config.blockSize[0] = order;
1660     config.blockSize[1] = (pmeGpu->settings.threadsPerAtom == ThreadsPerAtom::Order ? 1 : order);
1661     config.blockSize[2] = atomsPerBlock;
1662     config.gridSize[0]  = dimGrid.first;
1663     config.gridSize[1]  = dimGrid.second;
1664
1665     // TODO test different cache configs
1666
1667     PmeStage                           timingId = PmeStage::Gather;
1668     PmeGpuProgramImpl::PmeKernelHandle kernelPtr =
1669             selectGatherKernelPtr(pmeGpu,
1670                                   pmeGpu->settings.threadsPerAtom,
1671                                   readGlobal || (!recalculateSplines),
1672                                   pmeGpu->common->ngrids);
1673     // TODO design kernel selection getters and make PmeGpu a friend of PmeGpuProgramImpl
1674
1675     pme_gpu_start_timing(pmeGpu, timingId);
1676     auto* timingEvent     = pme_gpu_fetch_timing_event(pmeGpu, timingId);
1677     auto* kernelParamsPtr = pmeGpu->kernelParams.get();
1678     if (pmeGpu->common->ngrids == 1)
1679     {
1680         kernelParamsPtr->current.scale = 1.0;
1681     }
1682     else
1683     {
1684         kernelParamsPtr->current.scale = 1.0 - lambda;
1685     }
1686
1687 #if c_canEmbedBuffers
1688     const auto kernelArgs = prepareGpuKernelArguments(kernelPtr, config, kernelParamsPtr);
1689 #else
1690     const auto kernelArgs =
1691             prepareGpuKernelArguments(kernelPtr,
1692                                       config,
1693                                       kernelParamsPtr,
1694                                       &kernelParamsPtr->atoms.d_coefficients[FEP_STATE_A],
1695                                       &kernelParamsPtr->atoms.d_coefficients[FEP_STATE_B],
1696                                       &kernelParamsPtr->grid.d_realGrid[FEP_STATE_A],
1697                                       &kernelParamsPtr->grid.d_realGrid[FEP_STATE_B],
1698                                       &kernelParamsPtr->atoms.d_theta,
1699                                       &kernelParamsPtr->atoms.d_dtheta,
1700                                       &kernelParamsPtr->atoms.d_gridlineIndices,
1701                                       &kernelParamsPtr->atoms.d_forces);
1702 #endif
1703     launchGpuKernel(kernelPtr, config, pmeGpu->archSpecific->pmeStream_, timingEvent, "PME gather", kernelArgs);
1704     pme_gpu_stop_timing(pmeGpu, timingId);
1705
1706     if (pmeGpu->settings.useGpuForceReduction)
1707     {
1708         pmeGpu->archSpecific->pmeForcesReady.markEvent(pmeGpu->archSpecific->pmeStream_);
1709     }
1710     else
1711     {
1712         pme_gpu_copy_output_forces(pmeGpu);
1713     }
1714 }
1715
1716 DeviceBuffer<gmx::RVec> pme_gpu_get_kernelparam_forces(const PmeGpu* pmeGpu)
1717 {
1718     if (pmeGpu && pmeGpu->kernelParams)
1719     {
1720         return pmeGpu->kernelParams->atoms.d_forces;
1721     }
1722     else
1723     {
1724         return DeviceBuffer<gmx::RVec>{};
1725     }
1726 }
1727
1728 void pme_gpu_set_kernelparam_coordinates(const PmeGpu* pmeGpu, DeviceBuffer<gmx::RVec> d_x)
1729 {
1730     GMX_ASSERT(pmeGpu && pmeGpu->kernelParams,
1731                "PME GPU device buffer can not be set in non-GPU builds or before the GPU PME was "
1732                "initialized.");
1733
1734     GMX_ASSERT(checkDeviceBuffer(d_x, pmeGpu->kernelParams->atoms.nAtoms),
1735                "The device-side buffer can not be set.");
1736
1737     pmeGpu->kernelParams->atoms.d_coordinates = d_x;
1738 }
1739
1740 GpuEventSynchronizer* pme_gpu_get_forces_ready_synchronizer(const PmeGpu* pmeGpu)
1741 {
1742     if (pmeGpu && pmeGpu->kernelParams)
1743     {
1744         return &pmeGpu->archSpecific->pmeForcesReady;
1745     }
1746     else
1747     {
1748         return nullptr;
1749     }
1750 }