src/gromacs/ewald/tests/pmetestcommon.cpp

   1 /*
   2  * This file is part of the GROMACS molecular simulation package.
   3  *
   4  * Copyright (c) 2016,2017,2018,2019, by the GROMACS development team, led by
   5  * Mark Abraham, David van der Spoel, Berk Hess, and Erik Lindahl,
   6  * and including many others, as listed in the AUTHORS file in the
   7  * top-level source directory and at http://www.gromacs.org.
   8  *
   9  * GROMACS is free software; you can redistribute it and/or
  10  * modify it under the terms of the GNU Lesser General Public License
  11  * as published by the Free Software Foundation; either version 2.1
  12  * of the License, or (at your option) any later version.
  13  *
  14  * GROMACS is distributed in the hope that it will be useful,
  15  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  16  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  17  * Lesser General Public License for more details.
  18  *
  19  * You should have received a copy of the GNU Lesser General Public
  20  * License along with GROMACS; if not, see
  21  * http://www.gnu.org/licenses, or write to the Free Software Foundation,
  22  * Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA.
  23  *
  24  * If you want to redistribute modifications to GROMACS, please
  25  * consider that scientific software is very special. Version
  26  * control is crucial - bugs must be traceable. We will be happy to
  27  * consider code for inclusion in the official distribution, but
  28  * derived work must not be called official GROMACS. Details are found
  29  * in the README & COPYING files - if they are missing, get the
  30  * official version at http://www.gromacs.org.
  31  *
  32  * To help us fund GROMACS development, we humbly ask that you cite
  33  * the research papers on the package. Check out http://www.gromacs.org.
  34  */
  35 /*! \internal \file
  36  * \brief
  37  * Implements common routines for PME tests.
  38  *
  39  * \author Aleksei Iupinov <a.yupinov@gmail.com>
  40  * \ingroup module_ewald
  41  */
  42 #include "gmxpre.h"
  43
  44 #include "pmetestcommon.h"
  45
  46 #include <cstring>
  47
  48 #include <algorithm>
  49
  50 #include "gromacs/domdec/domdec.h"
  51 #include "gromacs/ewald/pme_gather.h"
  52 #include "gromacs/ewald/pme_gpu_internal.h"
  53 #include "gromacs/ewald/pme_grid.h"
  54 #include "gromacs/ewald/pme_internal.h"
  55 #include "gromacs/ewald/pme_redistribute.h"
  56 #include "gromacs/ewald/pme_solve.h"
  57 #include "gromacs/ewald/pme_spread.h"
  58 #include "gromacs/fft/parallel_3dfft.h"
  59 #include "gromacs/gpu_utils/gpu_utils.h"
  60 #include "gromacs/math/invertmatrix.h"
  61 #include "gromacs/mdtypes/commrec.h"
  62 #include "gromacs/pbcutil/pbc.h"
  63 #include "gromacs/topology/topology.h"
  64 #include "gromacs/utility/exceptions.h"
  65 #include "gromacs/utility/gmxassert.h"
  66 #include "gromacs/utility/logger.h"
  67 #include "gromacs/utility/stringutil.h"
  68
  69 #include "testutils/testasserts.h"
  70
  71 namespace gmx
  72 {
  73 namespace test
  74 {
  75
  76 bool pmeSupportsInputForMode(const gmx_hw_info_t &hwinfo,
  77                              const t_inputrec    *inputRec,
  78                              CodePath             mode)
  79 {
  80     bool       implemented;
  81     gmx_mtop_t mtop;
  82     switch (mode)
  83     {
  84         case CodePath::CPU:
  85             implemented = true;
  86             break;
  87
  88         case CodePath::GPU:
  89             implemented = (pme_gpu_supports_build(nullptr) &&
  90                            pme_gpu_supports_hardware(hwinfo, nullptr) &&
  91                            pme_gpu_supports_input(*inputRec, mtop, nullptr));
  92             break;
  93
  94         default:
  95             GMX_THROW(InternalError("Test not implemented for this mode"));
  96     }
  97     return implemented;
  98 }
  99
 100 uint64_t getSplineModuliDoublePrecisionUlps(int splineOrder)
 101 {
 102     /* Arbitrary ulp tolerance for sine/cosine implementation. It's
 103      * hard to know what to pick without testing lots of
 104      * implementations. */
 105     const uint64_t sineUlps = 10;
 106     return 4 * (splineOrder - 2) + 2 * sineUlps * splineOrder;
 107 }
 108
 109 //! PME initialization
 110 PmeSafePointer pmeInitWrapper(const t_inputrec         *inputRec,
 111                               const CodePath            mode,
 112                               const gmx_device_info_t  *gpuInfo,
 113                               PmeGpuProgramHandle       pmeGpuProgram,
 114                               const Matrix3x3          &box,
 115                               const real                ewaldCoeff_q,
 116                               const real                ewaldCoeff_lj)
 117 {
 118     const MDLogger dummyLogger;
 119     const auto     runMode       = (mode == CodePath::CPU) ? PmeRunMode::CPU : PmeRunMode::Mixed;
 120     t_commrec      dummyCommrec  = {0};
 121     NumPmeDomains  numPmeDomains = { 1, 1 };
 122     gmx_pme_t     *pmeDataRaw    = gmx_pme_init(&dummyCommrec, numPmeDomains, inputRec, false, false, true,
 123                                                 ewaldCoeff_q, ewaldCoeff_lj, 1, runMode, nullptr, gpuInfo, pmeGpuProgram, dummyLogger);
 124     PmeSafePointer pme(pmeDataRaw); // taking ownership
 125
 126     // TODO get rid of this with proper matrix type
 127     matrix boxTemp;
 128     for (int i = 0; i < DIM; i++)
 129     {
 130         for (int j = 0; j < DIM; j++)
 131         {
 132             boxTemp[i][j] = box[i * DIM + j];
 133         }
 134     }
 135     const char *boxError = check_box(-1, boxTemp);
 136     GMX_RELEASE_ASSERT(boxError == nullptr, boxError);
 137
 138     switch (mode)
 139     {
 140         case CodePath::CPU:
 141             invertBoxMatrix(boxTemp, pme->recipbox);
 142             break;
 143
 144         case CodePath::GPU:
 145             pme_gpu_set_testing(pme->gpu, true);
 146             pme_gpu_update_input_box(pme->gpu, boxTemp);
 147             break;
 148
 149         default:
 150             GMX_THROW(InternalError("Test not implemented for this mode"));
 151     }
 152
 153     return pme;
 154 }
 155
 156 //! Simple PME initialization based on input, no atom data
 157 PmeSafePointer pmeInitEmpty(const t_inputrec         *inputRec,
 158                             const CodePath            mode,
 159                             const gmx_device_info_t  *gpuInfo,
 160                             PmeGpuProgramHandle       pmeGpuProgram,
 161                             const Matrix3x3          &box,
 162                             real                      ewaldCoeff_q,
 163                             real                      ewaldCoeff_lj
 164                             )
 165 {
 166     return pmeInitWrapper(inputRec, mode, gpuInfo, pmeGpuProgram, box, ewaldCoeff_q, ewaldCoeff_lj);
 167     // hiding the fact that PME actually needs to know the number of atoms in advance
 168 }
 169
 170 //! Make a GPU state-propagator manager
 171 std::unique_ptr<StatePropagatorDataGpu>
 172 makeStatePropagatorDataGpu(const gmx_pme_t &pme)
 173 {
 174     // TODO: Pin the host buffer and use async memory copies
 175     // TODO: Special constructor for PME-only rank / PME-tests is used here. There should be a mechanism to
 176     //       restrict one from using other constructor here.
 177     return std::make_unique<StatePropagatorDataGpu>(pme_gpu_get_device_stream(&pme),
 178                                                     pme_gpu_get_device_context(&pme),
 179                                                     GpuApiCallBehavior::Sync,
 180                                                     pme_gpu_get_padding_size(&pme));
 181 }
 182
 183 //! PME initialization with atom data
 184 void pmeInitAtoms(gmx_pme_t               *pme,
 185                   StatePropagatorDataGpu  *stateGpu,
 186                   const CodePath           mode,
 187                   const CoordinatesVector &coordinates,
 188                   const ChargesVector     &charges)
 189 {
 190     const index  atomCount = coordinates.size();
 191     GMX_RELEASE_ASSERT(atomCount == charges.ssize(), "Mismatch in atom data");
 192     PmeAtomComm *atc = nullptr;
 193
 194     switch (mode)
 195     {
 196         case CodePath::CPU:
 197             atc              = &(pme->atc[0]);
 198             atc->x           = coordinates;
 199             atc->coefficient = charges;
 200             gmx_pme_reinit_atoms(pme, atomCount, charges.data());
 201             /* With decomposition there would be more boilerplate atc code here, e.g. do_redist_pos_coeffs */
 202             break;
 203
 204         case CodePath::GPU:
 205             // TODO: Avoid use of atc in the GPU code path
 206             atc              = &(pme->atc[0]);
 207             // We need to set atc->n for passing the size in the tests
 208             atc->setNumAtoms(atomCount);
 209             gmx_pme_reinit_atoms(pme, atomCount, charges.data());
 210
 211             stateGpu->reinit(atomCount, atomCount);
 212             stateGpu->copyCoordinatesToGpu(arrayRefFromArray(coordinates.data(), coordinates.size()), gmx::AtomLocality::All);
 213             pme_gpu_set_kernelparam_coordinates(pme->gpu, stateGpu->getCoordinates());
 214
 215             break;
 216
 217         default:
 218             GMX_THROW(InternalError("Test not implemented for this mode"));
 219     }
 220 }
 221
 222 //! Getting local PME real grid pointer for test I/O
 223 static real *pmeGetRealGridInternal(const gmx_pme_t *pme)
 224 {
 225     const size_t gridIndex = 0;
 226     return pme->fftgrid[gridIndex];
 227 }
 228
 229 //! Getting local PME real grid dimensions
 230 static void pmeGetRealGridSizesInternal(const gmx_pme_t      *pme,
 231                                         CodePath              mode,
 232                                         IVec                 &gridSize,       //NOLINT(google-runtime-references)
 233                                         IVec                 &paddedGridSize) //NOLINT(google-runtime-references)
 234 {
 235     const size_t gridIndex = 0;
 236     IVec         gridOffsetUnused;
 237     switch (mode)
 238     {
 239         case CodePath::CPU:
 240             gmx_parallel_3dfft_real_limits(pme->pfft_setup[gridIndex], gridSize, gridOffsetUnused, paddedGridSize);
 241             break;
 242
 243         case CodePath::GPU:
 244             pme_gpu_get_real_grid_sizes(pme->gpu, &gridSize, &paddedGridSize);
 245             break;
 246
 247         default:
 248             GMX_THROW(InternalError("Test not implemented for this mode"));
 249     }
 250 }
 251
 252 //! Getting local PME complex grid pointer for test I/O
 253 static t_complex *pmeGetComplexGridInternal(const gmx_pme_t *pme)
 254 {
 255     const size_t gridIndex = 0;
 256     return pme->cfftgrid[gridIndex];
 257 }
 258
 259 //! Getting local PME complex grid dimensions
 260 static void pmeGetComplexGridSizesInternal(const gmx_pme_t      *pme,
 261                                            IVec                 &gridSize,       //NOLINT(google-runtime-references)
 262                                            IVec                 &paddedGridSize) //NOLINT(google-runtime-references)
 263 {
 264     const size_t gridIndex = 0;
 265     IVec         gridOffsetUnused, complexOrderUnused;
 266     gmx_parallel_3dfft_complex_limits(pme->pfft_setup[gridIndex], complexOrderUnused, gridSize, gridOffsetUnused, paddedGridSize); //TODO: what about YZX ordering?
 267 }
 268
 269 //! Getting the PME grid memory buffer and its sizes - template definition
 270 template<typename ValueType> static void pmeGetGridAndSizesInternal(const gmx_pme_t * /*unused*/, CodePath /*unused*/, ValueType * & /*unused*/, IVec & /*unused*/, IVec & /*unused*/) //NOLINT(google-runtime-references)
 271 {
 272     GMX_THROW(InternalError("Deleted function call"));
 273     // explicitly deleting general template does not compile in clang/icc, see https://llvm.org/bugs/show_bug.cgi?id=17537
 274 }
 275
 276 //! Getting the PME real grid memory buffer and its sizes
 277 template<> void pmeGetGridAndSizesInternal<real>(const gmx_pme_t *pme, CodePath mode, real * &grid, IVec &gridSize, IVec &paddedGridSize)
 278 {
 279     grid = pmeGetRealGridInternal(pme);
 280     pmeGetRealGridSizesInternal(pme, mode, gridSize, paddedGridSize);
 281 }
 282
 283 //! Getting the PME complex grid memory buffer and its sizes
 284 template<> void pmeGetGridAndSizesInternal<t_complex>(const gmx_pme_t *pme, CodePath /*unused*/, t_complex * &grid, IVec &gridSize, IVec &paddedGridSize)
 285 {
 286     grid = pmeGetComplexGridInternal(pme);
 287     pmeGetComplexGridSizesInternal(pme, gridSize, paddedGridSize);
 288 }
 289
 290 //! PME spline calculation and charge spreading
 291 void pmePerformSplineAndSpread(gmx_pme_t *pme, CodePath mode, // TODO const qualifiers elsewhere
 292                                bool computeSplines, bool spreadCharges)
 293 {
 294     GMX_RELEASE_ASSERT(pme != nullptr, "PME data is not initialized");
 295     PmeAtomComm    *atc                          = &(pme->atc[0]);
 296     const size_t    gridIndex                    = 0;
 297     const bool      computeSplinesForZeroCharges = true;
 298     real           *fftgrid                      = spreadCharges ? pme->fftgrid[gridIndex] : nullptr;
 299     real           *pmegrid                      = pme->pmegrid[gridIndex].grid.grid;
 300
 301     switch (mode)
 302     {
 303         case CodePath::CPU:
 304             spread_on_grid(pme, atc, &pme->pmegrid[gridIndex], computeSplines, spreadCharges,
 305                            fftgrid, computeSplinesForZeroCharges, gridIndex);
 306             if (spreadCharges && !pme->bUseThreads)
 307             {
 308                 wrap_periodic_pmegrid(pme, pmegrid);
 309                 copy_pmegrid_to_fftgrid(pme, pmegrid, fftgrid, gridIndex);
 310             }
 311             break;
 312
 313         case CodePath::GPU:
 314         {
 315             // no synchronization needed as x is transferred in the PME stream
 316             GpuEventSynchronizer *xReadyOnDevice = nullptr;
 317             pme_gpu_spread(pme->gpu, xReadyOnDevice, gridIndex, fftgrid, computeSplines, spreadCharges);
 318         }
 319         break;
 320
 321         default:
 322             GMX_THROW(InternalError("Test not implemented for this mode"));
 323     }
 324 }
 325
 326 //! Getting the internal spline data buffer pointer
 327 static real *pmeGetSplineDataInternal(const gmx_pme_t *pme, PmeSplineDataType type, int dimIndex)
 328 {
 329     GMX_ASSERT((0 <= dimIndex) && (dimIndex < DIM), "Invalid dimension index");
 330     const PmeAtomComm    *atc          = &(pme->atc[0]);
 331     const size_t          threadIndex  = 0;
 332     real                 *splineBuffer = nullptr;
 333     switch (type)
 334     {
 335         case PmeSplineDataType::Values:
 336             splineBuffer = atc->spline[threadIndex].theta.coefficients[dimIndex];
 337             break;
 338
 339         case PmeSplineDataType::Derivatives:
 340             splineBuffer = atc->spline[threadIndex].dtheta.coefficients[dimIndex];
 341             break;
 342
 343         default:
 344             GMX_THROW(InternalError("Unknown spline data type"));
 345     }
 346     return splineBuffer;
 347 }
 348
 349 //! PME solving
 350 void pmePerformSolve(const gmx_pme_t *pme, CodePath mode,
 351                      PmeSolveAlgorithm method, real cellVolume,
 352                      GridOrdering gridOrdering, bool computeEnergyAndVirial)
 353 {
 354     t_complex      *h_grid                 = pmeGetComplexGridInternal(pme);
 355     const bool      useLorentzBerthelot    = false;
 356     const size_t    threadIndex            = 0;
 357     switch (mode)
 358     {
 359         case CodePath::CPU:
 360             if (gridOrdering != GridOrdering::YZX)
 361             {
 362                 GMX_THROW(InternalError("Test not implemented for this mode"));
 363             }
 364             switch (method)
 365             {
 366                 case PmeSolveAlgorithm::Coulomb:
 367                     solve_pme_yzx(pme, h_grid, cellVolume,
 368                                   computeEnergyAndVirial, pme->nthread, threadIndex);
 369                     break;
 370
 371                 case PmeSolveAlgorithm::LennardJones:
 372                     solve_pme_lj_yzx(pme, &h_grid, useLorentzBerthelot,
 373                                      cellVolume, computeEnergyAndVirial, pme->nthread, threadIndex);
 374                     break;
 375
 376                 default:
 377                     GMX_THROW(InternalError("Test not implemented for this mode"));
 378             }
 379             break;
 380
 381         case CodePath::GPU:
 382             switch (method)
 383             {
 384                 case PmeSolveAlgorithm::Coulomb:
 385                     pme_gpu_solve(pme->gpu, h_grid, gridOrdering, computeEnergyAndVirial);
 386                     break;
 387
 388                 default:
 389                     GMX_THROW(InternalError("Test not implemented for this mode"));
 390             }
 391             break;
 392
 393         default:
 394             GMX_THROW(InternalError("Test not implemented for this mode"));
 395     }
 396 }
 397
 398 //! PME force gathering
 399 void pmePerformGather(gmx_pme_t *pme, CodePath mode,
 400                       PmeForceOutputHandling inputTreatment, ForcesVector &forces)
 401 {
 402     PmeAtomComm    *atc                     = &(pme->atc[0]);
 403     const index     atomCount               = atc->numAtoms();
 404     GMX_RELEASE_ASSERT(forces.ssize() == atomCount, "Invalid force buffer size");
 405     const bool      forceReductionWithInput = (inputTreatment == PmeForceOutputHandling::ReduceWithInput);
 406     const real      scale                   = 1.0;
 407     const size_t    threadIndex             = 0;
 408     const size_t    gridIndex               = 0;
 409     real           *pmegrid                 = pme->pmegrid[gridIndex].grid.grid;
 410     real           *fftgrid                 = pme->fftgrid[gridIndex];
 411
 412     switch (mode)
 413     {
 414         case CodePath::CPU:
 415             atc->f = forces;
 416             if (atc->nthread == 1)
 417             {
 418                 // something which is normally done in serial spline computation (make_thread_local_ind())
 419                 atc->spline[threadIndex].n = atomCount;
 420             }
 421             copy_fftgrid_to_pmegrid(pme, fftgrid, pmegrid, gridIndex, pme->nthread, threadIndex);
 422             unwrap_periodic_pmegrid(pme, pmegrid);
 423             gather_f_bsplines(pme, pmegrid, !forceReductionWithInput, atc, &atc->spline[threadIndex], scale);
 424             break;
 425
 426         case CodePath::GPU:
 427         {
 428             // Variable initialization needs a non-switch scope
 429             PmeOutput output = pme_gpu_getOutput(*pme, GMX_PME_CALC_F);
 430             GMX_ASSERT(forces.size() == output.forces_.size(), "Size of force buffers did not match");
 431             if (forceReductionWithInput)
 432             {
 433                 std::copy(std::begin(forces), std::end(forces), std::begin(output.forces_));
 434             }
 435             pme_gpu_gather(pme->gpu, inputTreatment, reinterpret_cast<float *>(fftgrid));
 436             std::copy(std::begin(output.forces_), std::end(output.forces_), std::begin(forces));
 437         }
 438         break;
 439
 440         default:
 441             GMX_THROW(InternalError("Test not implemented for this mode"));
 442     }
 443 }
 444
 445 //! PME test finalization before fetching the outputs
 446 void pmeFinalizeTest(const gmx_pme_t *pme, CodePath mode)
 447 {
 448     switch (mode)
 449     {
 450         case CodePath::CPU:
 451             break;
 452
 453         case CodePath::GPU:
 454             pme_gpu_synchronize(pme->gpu);
 455             break;
 456
 457         default:
 458             GMX_THROW(InternalError("Test not implemented for this mode"));
 459     }
 460 }
 461
 462 //! Setting atom spline values/derivatives to be used in spread/gather
 463 void pmeSetSplineData(const gmx_pme_t *pme, CodePath mode,
 464                       const SplineParamsDimVector &splineValues, PmeSplineDataType type, int dimIndex)
 465 {
 466     const PmeAtomComm    *atc         = &(pme->atc[0]);
 467     const index           atomCount   = atc->numAtoms();
 468     const index           pmeOrder    = pme->pme_order;
 469     const index           dimSize     = pmeOrder * atomCount;
 470     GMX_RELEASE_ASSERT(dimSize == splineValues.ssize(), "Mismatch in spline data");
 471     real                 *splineBuffer = pmeGetSplineDataInternal(pme, type, dimIndex);
 472
 473     switch (mode)
 474     {
 475         case CodePath::CPU:
 476             std::copy(splineValues.begin(), splineValues.end(), splineBuffer);
 477             break;
 478
 479         case CodePath::GPU:
 480             std::copy(splineValues.begin(), splineValues.end(), splineBuffer);
 481             pme_gpu_transform_spline_atom_data(pme->gpu, atc, type, dimIndex, PmeLayoutTransform::HostToGpu);
 482             break;
 483
 484         default:
 485             GMX_THROW(InternalError("Test not implemented for this mode"));
 486     }
 487 }
 488
 489 //! Setting gridline indices to be used in spread/gather
 490 void pmeSetGridLineIndices(gmx_pme_t *pme, CodePath mode,
 491                            const GridLineIndicesVector &gridLineIndices)
 492 {
 493     PmeAtomComm                *atc         = &(pme->atc[0]);
 494     const index                 atomCount   = atc->numAtoms();
 495     GMX_RELEASE_ASSERT(atomCount == gridLineIndices.ssize(), "Mismatch in gridline indices size");
 496
 497     IVec paddedGridSizeUnused, gridSize(0, 0, 0);
 498     pmeGetRealGridSizesInternal(pme, mode, gridSize, paddedGridSizeUnused);
 499
 500     for (const auto &index : gridLineIndices)
 501     {
 502         for (int i = 0; i < DIM; i++)
 503         {
 504             GMX_RELEASE_ASSERT((0 <= index[i]) && (index[i] < gridSize[i]), "Invalid gridline index");
 505         }
 506     }
 507
 508     switch (mode)
 509     {
 510         case CodePath::GPU:
 511             memcpy(pme->gpu->staging.h_gridlineIndices, gridLineIndices.data(), atomCount * sizeof(gridLineIndices[0]));
 512             break;
 513
 514         case CodePath::CPU:
 515             atc->idx.resize(gridLineIndices.size());
 516             std::copy(gridLineIndices.begin(), gridLineIndices.end(), atc->idx.begin());
 517             break;
 518         default:
 519             GMX_THROW(InternalError("Test not implemented for this mode"));
 520     }
 521 }
 522
 523 //! Getting plain index into the complex 3d grid
 524 inline size_t pmeGetGridPlainIndexInternal(const IVec &index, const IVec &paddedGridSize, GridOrdering gridOrdering)
 525 {
 526     size_t result;
 527     switch (gridOrdering)
 528     {
 529         case GridOrdering::YZX:
 530             result = (index[YY] * paddedGridSize[ZZ] + index[ZZ]) * paddedGridSize[XX] + index[XX];
 531             break;
 532
 533         case GridOrdering::XYZ:
 534             result = (index[XX] * paddedGridSize[YY] + index[YY]) * paddedGridSize[ZZ] + index[ZZ];
 535             break;
 536
 537         default:
 538             GMX_THROW(InternalError("Test not implemented for this mode"));
 539     }
 540     return result;
 541 }
 542
 543 //! Setting real or complex grid
 544 template<typename ValueType>
 545 static void pmeSetGridInternal(const gmx_pme_t *pme, CodePath mode,
 546                                GridOrdering gridOrdering,
 547                                const SparseGridValuesInput<ValueType> &gridValues)
 548 {
 549     IVec       gridSize(0, 0, 0), paddedGridSize(0, 0, 0);
 550     ValueType *grid;
 551     pmeGetGridAndSizesInternal<ValueType>(pme, mode, grid, gridSize, paddedGridSize);
 552
 553     switch (mode)
 554     {
 555         case CodePath::GPU: // intentional absence of break, the grid will be copied from the host buffer in testing mode
 556         case CodePath::CPU:
 557             std::memset(grid, 0, paddedGridSize[XX] * paddedGridSize[YY] * paddedGridSize[ZZ] * sizeof(ValueType));
 558             for (const auto &gridValue : gridValues)
 559             {
 560                 for (int i = 0; i < DIM; i++)
 561                 {
 562                     GMX_RELEASE_ASSERT((0 <= gridValue.first[i]) && (gridValue.first[i] < gridSize[i]), "Invalid grid value index");
 563                 }
 564                 const size_t gridValueIndex = pmeGetGridPlainIndexInternal(gridValue.first, paddedGridSize, gridOrdering);
 565                 grid[gridValueIndex] = gridValue.second;
 566             }
 567             break;
 568
 569         default:
 570             GMX_THROW(InternalError("Test not implemented for this mode"));
 571     }
 572 }
 573
 574 //! Setting real grid to be used in gather
 575 void pmeSetRealGrid(const gmx_pme_t *pme, CodePath mode,
 576                     const SparseRealGridValuesInput &gridValues)
 577 {
 578     pmeSetGridInternal<real>(pme, mode, GridOrdering::XYZ, gridValues);
 579 }
 580
 581 //! Setting complex grid to be used in solve
 582 void pmeSetComplexGrid(const gmx_pme_t *pme, CodePath mode,
 583                        GridOrdering gridOrdering,
 584                        const SparseComplexGridValuesInput &gridValues)
 585 {
 586     pmeSetGridInternal<t_complex>(pme, mode, gridOrdering, gridValues);
 587 }
 588
 589 //! Getting the single dimension's spline values or derivatives
 590 SplineParamsDimVector pmeGetSplineData(const gmx_pme_t *pme, CodePath mode,
 591                                        PmeSplineDataType type, int dimIndex)
 592 {
 593     GMX_RELEASE_ASSERT(pme != nullptr, "PME data is not initialized");
 594     const PmeAtomComm       *atc         = &(pme->atc[0]);
 595     const size_t             atomCount   = atc->numAtoms();
 596     const size_t             pmeOrder    = pme->pme_order;
 597     const size_t             dimSize     = pmeOrder * atomCount;
 598
 599     real                    *sourceBuffer = pmeGetSplineDataInternal(pme, type, dimIndex);
 600     SplineParamsDimVector    result;
 601     switch (mode)
 602     {
 603         case CodePath::GPU:
 604             pme_gpu_transform_spline_atom_data(pme->gpu, atc, type, dimIndex, PmeLayoutTransform::GpuToHost);
 605             result = arrayRefFromArray(sourceBuffer, dimSize);
 606             break;
 607
 608         case CodePath::CPU:
 609             result = arrayRefFromArray(sourceBuffer, dimSize);
 610             break;
 611
 612         default:
 613             GMX_THROW(InternalError("Test not implemented for this mode"));
 614     }
 615     return result;
 616 }
 617
 618 //! Getting the gridline indices
 619 GridLineIndicesVector pmeGetGridlineIndices(const gmx_pme_t *pme, CodePath mode)
 620 {
 621     GMX_RELEASE_ASSERT(pme != nullptr, "PME data is not initialized");
 622     const PmeAtomComm    *atc         = &(pme->atc[0]);
 623     const size_t          atomCount   = atc->numAtoms();
 624
 625     GridLineIndicesVector gridLineIndices;
 626     switch (mode)
 627     {
 628         case CodePath::GPU:
 629             gridLineIndices = arrayRefFromArray(reinterpret_cast<IVec *>(pme->gpu->staging.h_gridlineIndices), atomCount);
 630             break;
 631
 632         case CodePath::CPU:
 633             gridLineIndices = atc->idx;
 634             break;
 635
 636         default:
 637             GMX_THROW(InternalError("Test not implemented for this mode"));
 638     }
 639     return gridLineIndices;
 640 }
 641
 642 //! Getting real or complex grid - only non zero values
 643 template<typename ValueType>
 644 static SparseGridValuesOutput<ValueType> pmeGetGridInternal(const gmx_pme_t *pme, CodePath mode, GridOrdering gridOrdering)
 645 {
 646     IVec       gridSize(0, 0, 0), paddedGridSize(0, 0, 0);
 647     ValueType *grid;
 648     pmeGetGridAndSizesInternal<ValueType>(pme, mode, grid, gridSize, paddedGridSize);
 649     SparseGridValuesOutput<ValueType> gridValues;
 650     switch (mode)
 651     {
 652         case CodePath::GPU: // intentional absence of break
 653         case CodePath::CPU:
 654             gridValues.clear();
 655             for (int ix = 0; ix < gridSize[XX]; ix++)
 656             {
 657                 for (int iy = 0; iy < gridSize[YY]; iy++)
 658                 {
 659                     for (int iz = 0; iz < gridSize[ZZ]; iz++)
 660                     {
 661                         IVec            temp(ix, iy, iz);
 662                         const size_t    gridValueIndex = pmeGetGridPlainIndexInternal(temp, paddedGridSize, gridOrdering);
 663                         const ValueType value          = grid[gridValueIndex];
 664                         if (value != ValueType {})
 665                         {
 666                             auto key = formatString("Cell %d %d %d", ix, iy, iz);
 667                             gridValues[key] = value;
 668                         }
 669                     }
 670                 }
 671             }
 672             break;
 673
 674         default:
 675             GMX_THROW(InternalError("Test not implemented for this mode"));
 676     }
 677     return gridValues;
 678 }
 679
 680 //! Getting the real grid (spreading output of pmePerformSplineAndSpread())
 681 SparseRealGridValuesOutput pmeGetRealGrid(const gmx_pme_t *pme, CodePath mode)
 682 {
 683     return pmeGetGridInternal<real>(pme, mode, GridOrdering::XYZ);
 684 }
 685
 686 //! Getting the complex grid output of pmePerformSolve()
 687 SparseComplexGridValuesOutput pmeGetComplexGrid(const gmx_pme_t *pme, CodePath mode,
 688                                                 GridOrdering gridOrdering)
 689 {
 690     return pmeGetGridInternal<t_complex>(pme, mode, gridOrdering);
 691 }
 692
 693 //! Getting the reciprocal energy and virial
 694 PmeOutput pmeGetReciprocalEnergyAndVirial(const gmx_pme_t *pme, CodePath mode,
 695                                           PmeSolveAlgorithm method)
 696 {
 697     PmeOutput output;
 698     switch (mode)
 699     {
 700         case CodePath::CPU:
 701             switch (method)
 702             {
 703                 case PmeSolveAlgorithm::Coulomb:
 704                     get_pme_ener_vir_q(pme->solve_work, pme->nthread, &output);
 705                     break;
 706
 707                 case PmeSolveAlgorithm::LennardJones:
 708                     get_pme_ener_vir_lj(pme->solve_work, pme->nthread, &output);
 709                     break;
 710
 711                 default:
 712                     GMX_THROW(InternalError("Test not implemented for this mode"));
 713             }
 714             break;
 715         case CodePath::GPU:
 716             switch (method)
 717             {
 718                 case PmeSolveAlgorithm::Coulomb:
 719                     pme_gpu_getEnergyAndVirial(*pme, &output);
 720                     break;
 721
 722                 default:
 723                     GMX_THROW(InternalError("Test not implemented for this mode"));
 724             }
 725             break;
 726
 727         default:
 728             GMX_THROW(InternalError("Test not implemented for this mode"));
 729     }
 730     return output;
 731 }
 732
 733 }  // namespace test
 734 }  // namespace gmx