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