src/gromacs/ewald/pme.cu

   1 /*
   2  * This file is part of the GROMACS molecular simulation package.
   3  *
   4  * Copyright (c) 2016,2017, 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 This file contains internal CUDA function implementations
  38  * for performing the PME calculations on GPU.
  39  *
  40  * \author Aleksei Iupinov <a.yupinov@gmail.com>
  41  */
  42
  43 #include "gmxpre.h"
  44
  45 #include <cmath>
  46
  47 /* The rest */
  48 #include "pme.h"
  49
  50 #include "gromacs/gpu_utils/cudautils.cuh"
  51 #include "gromacs/gpu_utils/pmalloc_cuda.h"
  52 #include "gromacs/utility/gmxassert.h"
  53 #include "gromacs/utility/smalloc.h"
  54
  55 #include "pme.cuh"
  56 #include "pme-3dfft.cuh"
  57 #include "pme-grid.h"
  58
  59 int pme_gpu_get_atom_data_alignment(const PmeGpu *pmeGPU)
  60 {
  61     const int order = pmeGPU->common->pme_order;
  62     GMX_ASSERT(order > 0, "Invalid PME order");
  63     return PME_ATOM_DATA_ALIGNMENT;
  64 }
  65
  66 int pme_gpu_get_atoms_per_warp(const PmeGpu *pmeGPU)
  67 {
  68     const int order = pmeGPU->common->pme_order;
  69     GMX_ASSERT(order > 0, "Invalid PME order");
  70     return PME_SPREADGATHER_ATOMS_PER_WARP;
  71 }
  72
  73 void pme_gpu_synchronize(const PmeGpu *pmeGPU)
  74 {
  75     cudaError_t stat = cudaStreamSynchronize(pmeGPU->archSpecific->pmeStream);
  76     CU_RET_ERR(stat, "Failed to synchronize the PME GPU stream!");
  77 }
  78
  79 void pme_gpu_alloc_energy_virial(const PmeGpu *pmeGPU)
  80 {
  81     const size_t energyAndVirialSize = c_virialAndEnergyCount * sizeof(float);
  82     cudaError_t  stat                = cudaMalloc((void **)&pmeGPU->kernelParams->constants.d_virialAndEnergy, energyAndVirialSize);
  83     CU_RET_ERR(stat, "cudaMalloc failed on PME energy and virial");
  84     pmalloc((void **)&pmeGPU->staging.h_virialAndEnergy, energyAndVirialSize);
  85 }
  86
  87 void pme_gpu_free_energy_virial(PmeGpu *pmeGPU)
  88 {
  89     cudaError_t stat = cudaFree(pmeGPU->kernelParams->constants.d_virialAndEnergy);
  90     CU_RET_ERR(stat, "cudaFree failed on PME energy and virial");
  91     pmeGPU->kernelParams->constants.d_virialAndEnergy = nullptr;
  92     pfree(pmeGPU->staging.h_virialAndEnergy);
  93     pmeGPU->staging.h_virialAndEnergy = nullptr;
  94 }
  95
  96 void pme_gpu_clear_energy_virial(const PmeGpu *pmeGPU)
  97 {
  98     cudaError_t stat = cudaMemsetAsync(pmeGPU->kernelParams->constants.d_virialAndEnergy, 0,
  99                                        c_virialAndEnergyCount * sizeof(float), pmeGPU->archSpecific->pmeStream);
 100     CU_RET_ERR(stat, "PME energy/virial cudaMemsetAsync error");
 101 }
 102
 103 void pme_gpu_realloc_and_copy_bspline_values(const PmeGpu *pmeGPU)
 104 {
 105     const int splineValuesOffset[DIM] = {
 106         0,
 107         pmeGPU->kernelParams->grid.realGridSize[XX],
 108         pmeGPU->kernelParams->grid.realGridSize[XX] + pmeGPU->kernelParams->grid.realGridSize[YY]
 109     };
 110     memcpy((void *)&pmeGPU->kernelParams->grid.splineValuesOffset, &splineValuesOffset, sizeof(splineValuesOffset));
 111
 112     const int newSplineValuesSize = pmeGPU->kernelParams->grid.realGridSize[XX] +
 113         pmeGPU->kernelParams->grid.realGridSize[YY] +
 114         pmeGPU->kernelParams->grid.realGridSize[ZZ];
 115     const bool shouldRealloc = (newSplineValuesSize > pmeGPU->archSpecific->splineValuesSize);
 116     cu_realloc_buffered((void **)&pmeGPU->kernelParams->grid.d_splineModuli, nullptr, sizeof(float),
 117                         &pmeGPU->archSpecific->splineValuesSize, &pmeGPU->archSpecific->splineValuesSizeAlloc, newSplineValuesSize, pmeGPU->archSpecific->pmeStream, true);
 118     if (shouldRealloc)
 119     {
 120         /* Reallocate the host buffer */
 121         pfree(pmeGPU->staging.h_splineModuli);
 122         pmalloc((void **)&pmeGPU->staging.h_splineModuli, newSplineValuesSize * sizeof(float));
 123     }
 124     for (int i = 0; i < DIM; i++)
 125     {
 126         memcpy(pmeGPU->staging.h_splineModuli + splineValuesOffset[i], pmeGPU->common->bsp_mod[i].data(), pmeGPU->common->bsp_mod[i].size() * sizeof(float));
 127     }
 128     /* TODO: pin original buffer instead! */
 129     cu_copy_H2D(pmeGPU->kernelParams->grid.d_splineModuli, pmeGPU->staging.h_splineModuli,
 130                 newSplineValuesSize * sizeof(float), pmeGPU->settings.transferKind, pmeGPU->archSpecific->pmeStream);
 131 }
 132
 133 void pme_gpu_free_bspline_values(const PmeGpu *pmeGPU)
 134 {
 135     pfree(pmeGPU->staging.h_splineModuli);
 136     cu_free_buffered(pmeGPU->kernelParams->grid.d_splineModuli, &pmeGPU->archSpecific->splineValuesSize,
 137                      &pmeGPU->archSpecific->splineValuesSizeAlloc);
 138 }
 139
 140 void pme_gpu_realloc_forces(const PmeGpu *pmeGPU)
 141 {
 142     const size_t newForcesSize = pmeGPU->nAtomsAlloc * DIM;
 143     GMX_ASSERT(newForcesSize > 0, "Bad number of atoms in PME GPU");
 144     cu_realloc_buffered((void **)&pmeGPU->kernelParams->atoms.d_forces, nullptr, sizeof(float),
 145                         &pmeGPU->archSpecific->forcesSize, &pmeGPU->archSpecific->forcesSizeAlloc, newForcesSize, pmeGPU->archSpecific->pmeStream, true);
 146 }
 147
 148 void pme_gpu_free_forces(const PmeGpu *pmeGPU)
 149 {
 150     cu_free_buffered(pmeGPU->kernelParams->atoms.d_forces, &pmeGPU->archSpecific->forcesSize, &pmeGPU->archSpecific->forcesSizeAlloc);
 151 }
 152
 153 void pme_gpu_copy_input_forces(const PmeGpu *pmeGPU, const float *h_forces)
 154 {
 155     GMX_ASSERT(h_forces, "nullptr host forces pointer in PME GPU");
 156     const size_t forcesSize = DIM * pmeGPU->kernelParams->atoms.nAtoms * sizeof(float);
 157     GMX_ASSERT(forcesSize > 0, "Bad number of atoms in PME GPU");
 158     cu_copy_H2D(pmeGPU->kernelParams->atoms.d_forces, const_cast<float *>(h_forces), forcesSize, pmeGPU->settings.transferKind, pmeGPU->archSpecific->pmeStream);
 159 }
 160
 161 void pme_gpu_copy_output_forces(const PmeGpu *pmeGPU, float *h_forces)
 162 {
 163     GMX_ASSERT(h_forces, "nullptr host forces pointer in PME GPU");
 164     const size_t forcesSize   = DIM * pmeGPU->kernelParams->atoms.nAtoms * sizeof(float);
 165     GMX_ASSERT(forcesSize > 0, "Bad number of atoms in PME GPU");
 166     cu_copy_D2H(h_forces, pmeGPU->kernelParams->atoms.d_forces, forcesSize, pmeGPU->settings.transferKind, pmeGPU->archSpecific->pmeStream);
 167 }
 168
 169 void pme_gpu_realloc_coordinates(const PmeGpu *pmeGPU)
 170 {
 171     const size_t newCoordinatesSize = pmeGPU->nAtomsAlloc * DIM;
 172     GMX_ASSERT(newCoordinatesSize > 0, "Bad number of atoms in PME GPU");
 173     cu_realloc_buffered((void **)&pmeGPU->kernelParams->atoms.d_coordinates, nullptr, sizeof(float),
 174                         &pmeGPU->archSpecific->coordinatesSize, &pmeGPU->archSpecific->coordinatesSizeAlloc, newCoordinatesSize, pmeGPU->archSpecific->pmeStream, true);
 175     if (c_usePadding)
 176     {
 177         const size_t paddingIndex = DIM * pmeGPU->kernelParams->atoms.nAtoms;
 178         const size_t paddingCount = DIM * pmeGPU->nAtomsAlloc - paddingIndex;
 179         if (paddingCount > 0)
 180         {
 181             cudaError_t stat = cudaMemsetAsync(pmeGPU->kernelParams->atoms.d_coordinates + paddingIndex, 0, paddingCount * sizeof(float), pmeGPU->archSpecific->pmeStream);
 182             CU_RET_ERR(stat, "PME failed to clear the padded coordinates");
 183         }
 184     }
 185 }
 186
 187 void pme_gpu_copy_input_coordinates(const PmeGpu *pmeGPU, const rvec *h_coordinates)
 188 {
 189     GMX_ASSERT(h_coordinates, "Bad host-side coordinate buffer in PME GPU");
 190 #if GMX_DOUBLE
 191     GMX_RELEASE_ASSERT(false, "Only single precision is supported");
 192     GMX_UNUSED_VALUE(h_coordinates);
 193 #else
 194     cu_copy_H2D(pmeGPU->kernelParams->atoms.d_coordinates, const_cast<rvec *>(h_coordinates),
 195                 pmeGPU->kernelParams->atoms.nAtoms * sizeof(rvec), pmeGPU->settings.transferKind, pmeGPU->archSpecific->pmeStream);
 196 #endif
 197 }
 198
 199 void pme_gpu_free_coordinates(const PmeGpu *pmeGPU)
 200 {
 201     cu_free_buffered(pmeGPU->kernelParams->atoms.d_coordinates, &pmeGPU->archSpecific->coordinatesSize, &pmeGPU->archSpecific->coordinatesSizeAlloc);
 202 }
 203
 204 void pme_gpu_realloc_and_copy_input_coefficients(const PmeGpu *pmeGPU, const float *h_coefficients)
 205 {
 206     GMX_ASSERT(h_coefficients, "Bad host-side charge buffer in PME GPU");
 207     const size_t newCoefficientsSize = pmeGPU->nAtomsAlloc;
 208     GMX_ASSERT(newCoefficientsSize > 0, "Bad number of atoms in PME GPU");
 209     cu_realloc_buffered((void **)&pmeGPU->kernelParams->atoms.d_coefficients, nullptr, sizeof(float),
 210                         &pmeGPU->archSpecific->coefficientsSize, &pmeGPU->archSpecific->coefficientsSizeAlloc,
 211                         newCoefficientsSize, pmeGPU->archSpecific->pmeStream, true);
 212     cu_copy_H2D(pmeGPU->kernelParams->atoms.d_coefficients, const_cast<float *>(h_coefficients),
 213                 pmeGPU->kernelParams->atoms.nAtoms * sizeof(float), pmeGPU->settings.transferKind, pmeGPU->archSpecific->pmeStream);
 214     if (c_usePadding)
 215     {
 216         const size_t paddingIndex = pmeGPU->kernelParams->atoms.nAtoms;
 217         const size_t paddingCount = pmeGPU->nAtomsAlloc - paddingIndex;
 218         if (paddingCount > 0)
 219         {
 220             cudaError_t stat = cudaMemsetAsync(pmeGPU->kernelParams->atoms.d_coefficients + paddingIndex, 0, paddingCount * sizeof(float), pmeGPU->archSpecific->pmeStream);
 221             CU_RET_ERR(stat, "PME failed to clear the padded charges");
 222         }
 223     }
 224 }
 225
 226 void pme_gpu_free_coefficients(const PmeGpu *pmeGPU)
 227 {
 228     cu_free_buffered(pmeGPU->kernelParams->atoms.d_coefficients, &pmeGPU->archSpecific->coefficientsSize, &pmeGPU->archSpecific->coefficientsSizeAlloc);
 229 }
 230
 231 void pme_gpu_realloc_spline_data(const PmeGpu *pmeGPU)
 232 {
 233     const int    order             = pmeGPU->common->pme_order;
 234     const int    alignment         = pme_gpu_get_atoms_per_warp(pmeGPU);
 235     const size_t nAtomsPadded      = ((pmeGPU->nAtomsAlloc + alignment - 1) / alignment) * alignment;
 236     const int    newSplineDataSize = DIM * order * nAtomsPadded;
 237     GMX_ASSERT(newSplineDataSize > 0, "Bad number of atoms in PME GPU");
 238     /* Two arrays of the same size */
 239     const bool shouldRealloc        = (newSplineDataSize > pmeGPU->archSpecific->splineDataSize);
 240     int        currentSizeTemp      = pmeGPU->archSpecific->splineDataSize;
 241     int        currentSizeTempAlloc = pmeGPU->archSpecific->splineDataSizeAlloc;
 242     cu_realloc_buffered((void **)&pmeGPU->kernelParams->atoms.d_theta, nullptr, sizeof(float),
 243                         &currentSizeTemp, &currentSizeTempAlloc, newSplineDataSize, pmeGPU->archSpecific->pmeStream, true);
 244     cu_realloc_buffered((void **)&pmeGPU->kernelParams->atoms.d_dtheta, nullptr, sizeof(float),
 245                         &pmeGPU->archSpecific->splineDataSize, &pmeGPU->archSpecific->splineDataSizeAlloc, newSplineDataSize, pmeGPU->archSpecific->pmeStream, true);
 246     // the host side reallocation
 247     if (shouldRealloc)
 248     {
 249         pfree(pmeGPU->staging.h_theta);
 250         pmalloc((void **)&pmeGPU->staging.h_theta, newSplineDataSize * sizeof(float));
 251         pfree(pmeGPU->staging.h_dtheta);
 252         pmalloc((void **)&pmeGPU->staging.h_dtheta, newSplineDataSize * sizeof(float));
 253     }
 254 }
 255
 256 void pme_gpu_free_spline_data(const PmeGpu *pmeGPU)
 257 {
 258     /* Two arrays of the same size */
 259     cu_free_buffered(pmeGPU->kernelParams->atoms.d_theta);
 260     cu_free_buffered(pmeGPU->kernelParams->atoms.d_dtheta, &pmeGPU->archSpecific->splineDataSize, &pmeGPU->archSpecific->splineDataSizeAlloc);
 261     pfree(pmeGPU->staging.h_theta);
 262     pfree(pmeGPU->staging.h_dtheta);
 263 }
 264
 265 void pme_gpu_realloc_grid_indices(const PmeGpu *pmeGPU)
 266 {
 267     const size_t newIndicesSize = DIM * pmeGPU->nAtomsAlloc;
 268     GMX_ASSERT(newIndicesSize > 0, "Bad number of atoms in PME GPU");
 269     cu_realloc_buffered((void **)&pmeGPU->kernelParams->atoms.d_gridlineIndices, nullptr, sizeof(int),
 270                         &pmeGPU->archSpecific->gridlineIndicesSize, &pmeGPU->archSpecific->gridlineIndicesSizeAlloc, newIndicesSize, pmeGPU->archSpecific->pmeStream, true);
 271     pfree(pmeGPU->staging.h_gridlineIndices);
 272     pmalloc((void **)&pmeGPU->staging.h_gridlineIndices, newIndicesSize * sizeof(int));
 273 }
 274
 275 void pme_gpu_free_grid_indices(const PmeGpu *pmeGPU)
 276 {
 277     cu_free_buffered(pmeGPU->kernelParams->atoms.d_gridlineIndices, &pmeGPU->archSpecific->gridlineIndicesSize, &pmeGPU->archSpecific->gridlineIndicesSizeAlloc);
 278     pfree(pmeGPU->staging.h_gridlineIndices);
 279 }
 280
 281 void pme_gpu_realloc_grids(PmeGpu *pmeGPU)
 282 {
 283     auto     *kernelParamsPtr = pmeGPU->kernelParams.get();
 284     const int newRealGridSize = kernelParamsPtr->grid.realGridSizePadded[XX] *
 285         kernelParamsPtr->grid.realGridSizePadded[YY] *
 286         kernelParamsPtr->grid.realGridSizePadded[ZZ];
 287     const int newComplexGridSize = kernelParamsPtr->grid.complexGridSizePadded[XX] *
 288         kernelParamsPtr->grid.complexGridSizePadded[YY] *
 289         kernelParamsPtr->grid.complexGridSizePadded[ZZ] * 2;
 290     // Multiplied by 2 because we count complex grid size for complex numbers, but all allocations/pointers are float
 291     if (pmeGPU->archSpecific->performOutOfPlaceFFT)
 292     {
 293         /* 2 separate grids */
 294         cu_realloc_buffered((void **)&kernelParamsPtr->grid.d_fourierGrid, nullptr, sizeof(float),
 295                             &pmeGPU->archSpecific->complexGridSize, &pmeGPU->archSpecific->complexGridSizeAlloc,
 296                             newComplexGridSize, pmeGPU->archSpecific->pmeStream, true);
 297         cu_realloc_buffered((void **)&kernelParamsPtr->grid.d_realGrid, nullptr, sizeof(float),
 298                             &pmeGPU->archSpecific->realGridSize, &pmeGPU->archSpecific->realGridSizeAlloc,
 299                             newRealGridSize, pmeGPU->archSpecific->pmeStream, true);
 300     }
 301     else
 302     {
 303         /* A single buffer so that any grid will fit */
 304         const int newGridsSize = std::max(newRealGridSize, newComplexGridSize);
 305         cu_realloc_buffered((void **)&kernelParamsPtr->grid.d_realGrid, nullptr, sizeof(float),
 306                             &pmeGPU->archSpecific->realGridSize, &pmeGPU->archSpecific->realGridSizeAlloc,
 307                             newGridsSize, pmeGPU->archSpecific->pmeStream, true);
 308         kernelParamsPtr->grid.d_fourierGrid   = kernelParamsPtr->grid.d_realGrid;
 309         pmeGPU->archSpecific->complexGridSize = pmeGPU->archSpecific->realGridSize;
 310         // the size might get used later for copying the grid
 311     }
 312 }
 313
 314 void pme_gpu_free_grids(const PmeGpu *pmeGPU)
 315 {
 316     if (pmeGPU->archSpecific->performOutOfPlaceFFT)
 317     {
 318         cu_free_buffered(pmeGPU->kernelParams->grid.d_fourierGrid);
 319     }
 320     cu_free_buffered(pmeGPU->kernelParams->grid.d_realGrid,
 321                      &pmeGPU->archSpecific->realGridSize, &pmeGPU->archSpecific->realGridSizeAlloc);
 322 }
 323
 324 void pme_gpu_clear_grids(const PmeGpu *pmeGPU)
 325 {
 326     cudaError_t stat = cudaMemsetAsync(pmeGPU->kernelParams->grid.d_realGrid, 0,
 327                                        pmeGPU->archSpecific->realGridSize * sizeof(float), pmeGPU->archSpecific->pmeStream);
 328     /* Should the complex grid be cleared in some weird case? */
 329     CU_RET_ERR(stat, "cudaMemsetAsync on the PME grid error");
 330 }
 331
 332 void pme_gpu_realloc_and_copy_fract_shifts(PmeGpu *pmeGPU)
 333 {
 334     pme_gpu_free_fract_shifts(pmeGPU);
 335
 336     auto        *kernelParamsPtr = pmeGPU->kernelParams.get();
 337
 338     const int    nx                  = kernelParamsPtr->grid.realGridSize[XX];
 339     const int    ny                  = kernelParamsPtr->grid.realGridSize[YY];
 340     const int    nz                  = kernelParamsPtr->grid.realGridSize[ZZ];
 341     const int    cellCount           = c_pmeNeighborUnitcellCount;
 342     const int    gridDataOffset[DIM] = {0, cellCount * nx, cellCount * (nx + ny)};
 343
 344     memcpy(kernelParamsPtr->grid.tablesOffsets, &gridDataOffset, sizeof(gridDataOffset));
 345
 346     const int    newFractShiftsSize  = cellCount * (nx + ny + nz);
 347
 348     initParamLookupTable(kernelParamsPtr->grid.d_fractShiftsTable,
 349                          kernelParamsPtr->fractShiftsTableTexture,
 350                          &pme_gpu_get_fract_shifts_texref(),
 351                          pmeGPU->common->fsh.data(),
 352                          newFractShiftsSize,
 353                          pmeGPU->deviceInfo);
 354
 355     initParamLookupTable(kernelParamsPtr->grid.d_gridlineIndicesTable,
 356                          kernelParamsPtr->gridlineIndicesTableTexture,
 357                          &pme_gpu_get_gridline_texref(),
 358                          pmeGPU->common->nn.data(),
 359                          newFractShiftsSize,
 360                          pmeGPU->deviceInfo);
 361 }
 362
 363 void pme_gpu_free_fract_shifts(const PmeGpu *pmeGPU)
 364 {
 365     auto *kernelParamsPtr = pmeGPU->kernelParams.get();
 366     destroyParamLookupTable(kernelParamsPtr->grid.d_fractShiftsTable,
 367                             kernelParamsPtr->fractShiftsTableTexture,
 368                             &pme_gpu_get_fract_shifts_texref(),
 369                             pmeGPU->deviceInfo);
 370     destroyParamLookupTable(kernelParamsPtr->grid.d_gridlineIndicesTable,
 371                             kernelParamsPtr->gridlineIndicesTableTexture,
 372                             &pme_gpu_get_gridline_texref(),
 373                             pmeGPU->deviceInfo);
 374 }
 375
 376 void pme_gpu_copy_input_gather_grid(const PmeGpu *pmeGpu, float *h_grid)
 377 {
 378     const size_t gridSize = pmeGpu->archSpecific->realGridSize * sizeof(float);
 379     cu_copy_H2D(pmeGpu->kernelParams->grid.d_realGrid, h_grid, gridSize, pmeGpu->settings.transferKind, pmeGpu->archSpecific->pmeStream);
 380 }
 381
 382 void pme_gpu_copy_output_spread_grid(const PmeGpu *pmeGpu, float *h_grid)
 383 {
 384     const size_t gridSize = pmeGpu->archSpecific->realGridSize * sizeof(float);
 385     cu_copy_D2H(h_grid, pmeGpu->kernelParams->grid.d_realGrid, gridSize, pmeGpu->settings.transferKind, pmeGpu->archSpecific->pmeStream);
 386     cudaError_t  stat = cudaEventRecord(pmeGpu->archSpecific->syncSpreadGridD2H, pmeGpu->archSpecific->pmeStream);
 387     CU_RET_ERR(stat, "PME spread grid sync event record failure");
 388 }
 389
 390 void pme_gpu_copy_output_spread_atom_data(const PmeGpu *pmeGpu)
 391 {
 392     const int    alignment       = pme_gpu_get_atoms_per_warp(pmeGpu);
 393     const size_t nAtomsPadded    = ((pmeGpu->nAtomsAlloc + alignment - 1) / alignment) * alignment;
 394     const size_t splinesSize     = DIM * nAtomsPadded * pmeGpu->common->pme_order * sizeof(float);
 395     auto        *kernelParamsPtr = pmeGpu->kernelParams.get();
 396     cu_copy_D2H(pmeGpu->staging.h_dtheta, kernelParamsPtr->atoms.d_dtheta, splinesSize, pmeGpu->settings.transferKind, pmeGpu->archSpecific->pmeStream);
 397     cu_copy_D2H(pmeGpu->staging.h_theta, kernelParamsPtr->atoms.d_theta, splinesSize, pmeGpu->settings.transferKind, pmeGpu->archSpecific->pmeStream);
 398     cu_copy_D2H(pmeGpu->staging.h_gridlineIndices, kernelParamsPtr->atoms.d_gridlineIndices,
 399                 kernelParamsPtr->atoms.nAtoms * DIM * sizeof(int), pmeGpu->settings.transferKind, pmeGpu->archSpecific->pmeStream);
 400 }
 401
 402 void pme_gpu_copy_input_gather_atom_data(const PmeGpu *pmeGpu)
 403 {
 404     const int    alignment       = pme_gpu_get_atoms_per_warp(pmeGpu);
 405     const size_t nAtomsPadded    = ((pmeGpu->nAtomsAlloc + alignment - 1) / alignment) * alignment;
 406     const size_t splinesSize     = DIM * nAtomsPadded * pmeGpu->common->pme_order * sizeof(float);
 407     auto        *kernelParamsPtr = pmeGpu->kernelParams.get();
 408     if (c_usePadding)
 409     {
 410         const size_t gridlineIndicesSizePerAtom = DIM * sizeof(int);
 411         const size_t splineDataSizePerAtom      = pmeGpu->common->pme_order * DIM * sizeof(float);
 412         // TODO: could clear only the padding and not the whole thing, but this is a test-exclusive code anyway
 413         CU_RET_ERR(cudaMemsetAsync(kernelParamsPtr->atoms.d_gridlineIndices, 0, pmeGpu->nAtomsAlloc * gridlineIndicesSizePerAtom, pmeGpu->archSpecific->pmeStream),
 414                    "PME failed to clear the gridline indices");
 415         CU_RET_ERR(cudaMemsetAsync(kernelParamsPtr->atoms.d_dtheta, 0, pmeGpu->nAtomsAlloc * splineDataSizePerAtom, pmeGpu->archSpecific->pmeStream),
 416                    "PME failed to clear the spline derivatives");
 417         CU_RET_ERR(cudaMemsetAsync(kernelParamsPtr->atoms.d_theta, 0, pmeGpu->nAtomsAlloc * splineDataSizePerAtom, pmeGpu->archSpecific->pmeStream),
 418                    "PME failed to clear the spline values");
 419     }
 420     cu_copy_H2D(kernelParamsPtr->atoms.d_dtheta, pmeGpu->staging.h_dtheta, splinesSize, pmeGpu->settings.transferKind, pmeGpu->archSpecific->pmeStream);
 421     cu_copy_H2D(kernelParamsPtr->atoms.d_theta, pmeGpu->staging.h_theta, splinesSize, pmeGpu->settings.transferKind, pmeGpu->archSpecific->pmeStream);
 422     cu_copy_H2D(kernelParamsPtr->atoms.d_gridlineIndices, pmeGpu->staging.h_gridlineIndices,
 423                 kernelParamsPtr->atoms.nAtoms * DIM * sizeof(int), pmeGpu->settings.transferKind, pmeGpu->archSpecific->pmeStream);
 424 }
 425
 426 void pme_gpu_sync_spread_grid(const PmeGpu *pmeGPU)
 427 {
 428     cudaError_t stat = cudaEventSynchronize(pmeGPU->archSpecific->syncSpreadGridD2H);
 429     CU_RET_ERR(stat, "Error while waiting for the PME GPU spread grid to be copied to the host");
 430 }
 431
 432 void pme_gpu_init_internal(PmeGpu *pmeGPU)
 433 {
 434     /* Allocate the target-specific structures */
 435     pmeGPU->archSpecific.reset(new PmeGpuSpecific());
 436     pmeGPU->kernelParams.reset(new PmeGpuKernelParams());
 437
 438     pmeGPU->archSpecific->performOutOfPlaceFFT = true;
 439     /* This should give better performance, according to the cuFFT documentation.
 440      * The performance seems to be the same though.
 441      * TODO: PME could also try to pick up nice grid sizes (with factors of 2, 3, 5, 7).
 442      */
 443
 444     pmeGPU->archSpecific->useTiming = (getenv("GMX_DISABLE_CUDA_TIMING") == nullptr) &&
 445         (getenv("GMX_DISABLE_GPU_TIMING") == nullptr);
 446     /* TODO: multiple CUDA streams on same GPU cause nonsense cudaEvent_t timings.
 447      * This should probably also check for gpuId exclusivity?
 448      */
 449
 450     /* Creating a PME CUDA stream */
 451     cudaError_t stat;
 452     int         highest_priority, lowest_priority;
 453     stat = cudaDeviceGetStreamPriorityRange(&lowest_priority, &highest_priority);
 454     CU_RET_ERR(stat, "PME cudaDeviceGetStreamPriorityRange failed");
 455     stat = cudaStreamCreateWithPriority(&pmeGPU->archSpecific->pmeStream,
 456                                         cudaStreamDefault, //cudaStreamNonBlocking,
 457                                         highest_priority);
 458     CU_RET_ERR(stat, "cudaStreamCreateWithPriority on the PME stream failed");
 459 }
 460
 461 void pme_gpu_destroy_specific(const PmeGpu *pmeGPU)
 462 {
 463     /* Destroy the CUDA stream */
 464     cudaError_t stat = cudaStreamDestroy(pmeGPU->archSpecific->pmeStream);
 465     CU_RET_ERR(stat, "PME cudaStreamDestroy error");
 466 }
 467
 468 void pme_gpu_init_sync_events(const PmeGpu *pmeGPU)
 469 {
 470     const auto  eventFlags = cudaEventDisableTiming;
 471     CU_RET_ERR(cudaEventCreateWithFlags(&pmeGPU->archSpecific->syncSpreadGridD2H, eventFlags), "cudaEventCreate on syncSpreadGridD2H failed");
 472 }
 473
 474 void pme_gpu_destroy_sync_events(const PmeGpu *pmeGPU)
 475 {
 476     CU_RET_ERR(cudaEventDestroy(pmeGPU->archSpecific->syncSpreadGridD2H), "cudaEventDestroy failed on syncSpreadGridD2H");
 477 }
 478
 479 void pme_gpu_reinit_3dfft(const PmeGpu *pmeGPU)
 480 {
 481     if (pme_gpu_performs_FFT(pmeGPU))
 482     {
 483         pmeGPU->archSpecific->fftSetup.resize(0);
 484         for (int i = 0; i < pmeGPU->common->ngrids; i++)
 485         {
 486             pmeGPU->archSpecific->fftSetup.push_back(std::unique_ptr<GpuParallel3dFft>(new GpuParallel3dFft(pmeGPU)));
 487         }
 488     }
 489 }
 490
 491 void pme_gpu_destroy_3dfft(const PmeGpu *pmeGPU)
 492 {
 493     pmeGPU->archSpecific->fftSetup.resize(0);
 494 }