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