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