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