src/gromacs/ewald/tests/pmetestcommon.cpp

   1 /*
   2  * This file is part of the GROMACS molecular simulation package.
   3  *
   4  * Copyright (c) 2016,2017,2018,2019, by the GROMACS development team, led by
   5  * Mark Abraham, David van der Spoel, Berk Hess, and Erik Lindahl,
   6  * and including many others, as listed in the AUTHORS file in the
   7  * top-level source directory and at http://www.gromacs.org.
   8  *
   9  * GROMACS is free software; you can redistribute it and/or
  10  * modify it under the terms of the GNU Lesser General Public License
  11  * as published by the Free Software Foundation; either version 2.1
  12  * of the License, or (at your option) any later version.
  13  *
  14  * GROMACS is distributed in the hope that it will be useful,
  15  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  16  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  17  * Lesser General Public License for more details.
  18  *
  19  * You should have received a copy of the GNU Lesser General Public
  20  * License along with GROMACS; if not, see
  21  * http://www.gnu.org/licenses, or write to the Free Software Foundation,
  22  * Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA.
  23  *
  24  * If you want to redistribute modifications to GROMACS, please
  25  * consider that scientific software is very special. Version
  26  * control is crucial - bugs must be traceable. We will be happy to
  27  * consider code for inclusion in the official distribution, but
  28  * derived work must not be called official GROMACS. Details are found
  29  * in the README & COPYING files - if they are missing, get the
  30  * official version at http://www.gromacs.org.
  31  *
  32  * To help us fund GROMACS development, we humbly ask that you cite
  33  * the research papers on the package. Check out http://www.gromacs.org.
  34  */
  35 /*! \internal \file
  36  * \brief
  37  * Implements common routines for PME tests.
  38  *
  39  * \author Aleksei Iupinov <a.yupinov@gmail.com>
  40  * \ingroup module_ewald
  41  */
  42 #include "gmxpre.h"
  43
  44 #include "pmetestcommon.h"
  45
  46 #include <cstring>
  47
  48 #include <algorithm>
  49
  50 #include "gromacs/domdec/domdec.h"
  51 #include "gromacs/ewald/pme_gather.h"
  52 #include "gromacs/ewald/pme_gpu_internal.h"
  53 #include "gromacs/ewald/pme_grid.h"
  54 #include "gromacs/ewald/pme_internal.h"
  55 #include "gromacs/ewald/pme_redistribute.h"
  56 #include "gromacs/ewald/pme_solve.h"
  57 #include "gromacs/ewald/pme_spread.h"
  58 #include "gromacs/fft/parallel_3dfft.h"
  59 #include "gromacs/gpu_utils/gpu_utils.h"
  60 #include "gromacs/math/invertmatrix.h"
  61 #include "gromacs/mdtypes/commrec.h"
  62 #include "gromacs/pbcutil/pbc.h"
  63 #include "gromacs/topology/topology.h"
  64 #include "gromacs/utility/exceptions.h"
  65 #include "gromacs/utility/gmxassert.h"
  66 #include "gromacs/utility/logger.h"
  67 #include "gromacs/utility/stringutil.h"
  68
  69 #include "testutils/testasserts.h"
  70
  71 namespace gmx
  72 {
  73 namespace test
  74 {
  75
  76 bool pmeSupportsInputForMode(const gmx_hw_info_t &hwinfo,
  77                              const t_inputrec    *inputRec,
  78                              CodePath             mode)
  79 {
  80     bool       implemented;
  81     gmx_mtop_t mtop;
  82     switch (mode)
  83     {
  84         case CodePath::CPU:
  85             implemented = true;
  86             break;
  87
  88         case CodePath::GPU:
  89             implemented = (pme_gpu_supports_build(nullptr) &&
  90                            pme_gpu_supports_hardware(hwinfo, nullptr) &&
  91                            pme_gpu_supports_input(*inputRec, mtop, nullptr));
  92             break;
  93
  94         default:
  95             GMX_THROW(InternalError("Test not implemented for this mode"));
  96     }
  97     return implemented;
  98 }
  99
 100 uint64_t getSplineModuliDoublePrecisionUlps(int splineOrder)
 101 {
 102     /* Arbitrary ulp tolerance for sine/cosine implementation. It's
 103      * hard to know what to pick without testing lots of
 104      * implementations. */
 105     const uint64_t sineUlps = 10;
 106     return 4 * (splineOrder - 2) + 2 * sineUlps * splineOrder;
 107 }
 108
 109 //! PME initialization - internal
 110 static PmeSafePointer pmeInitInternal(const t_inputrec         *inputRec,
 111                                       CodePath                  mode,
 112                                       const gmx_device_info_t  *gpuInfo,
 113                                       PmeGpuProgramHandle       pmeGpuProgram,
 114                                       const Matrix3x3          &box,
 115                                       real                      ewaldCoeff_q = 1.0f,
 116                                       real                      ewaldCoeff_lj = 1.0f
 117                                       )
 118 {
 119     const MDLogger dummyLogger;
 120     const auto     runMode       = (mode == CodePath::CPU) ? PmeRunMode::CPU : PmeRunMode::Mixed;
 121     t_commrec      dummyCommrec  = {0};
 122     NumPmeDomains  numPmeDomains = { 1, 1 };
 123     gmx_pme_t     *pmeDataRaw    = gmx_pme_init(&dummyCommrec, numPmeDomains, inputRec, false, false, true,
 124                                                 ewaldCoeff_q, ewaldCoeff_lj, 1, runMode, nullptr, gpuInfo, pmeGpuProgram, dummyLogger);
 125     PmeSafePointer pme(pmeDataRaw); // taking ownership
 126
 127     // TODO get rid of this with proper matrix type
 128     matrix boxTemp;
 129     for (int i = 0; i < DIM; i++)
 130     {
 131         for (int j = 0; j < DIM; j++)
 132         {
 133             boxTemp[i][j] = box[i * DIM + j];
 134         }
 135     }
 136     const char *boxError = check_box(-1, boxTemp);
 137     GMX_RELEASE_ASSERT(boxError == nullptr, boxError);
 138
 139     switch (mode)
 140     {
 141         case CodePath::CPU:
 142             invertBoxMatrix(boxTemp, pme->recipbox);
 143             break;
 144
 145         case CodePath::GPU:
 146             pme_gpu_set_testing(pme->gpu, true);
 147             pme_gpu_update_input_box(pme->gpu, boxTemp);
 148             break;
 149
 150         default:
 151             GMX_THROW(InternalError("Test not implemented for this mode"));
 152     }
 153
 154     return pme;
 155 }
 156
 157 //! Simple PME initialization based on input, no atom data
 158 PmeSafePointer pmeInitEmpty(const t_inputrec         *inputRec,
 159                             CodePath                  mode,
 160                             const gmx_device_info_t  *gpuInfo,
 161                             PmeGpuProgramHandle       pmeGpuProgram,
 162                             const Matrix3x3          &box,
 163                             real                      ewaldCoeff_q,
 164                             real                      ewaldCoeff_lj
 165                             )
 166 {
 167     return pmeInitInternal(inputRec, mode, gpuInfo, pmeGpuProgram, box, ewaldCoeff_q, ewaldCoeff_lj);
 168     // hiding the fact that PME actually needs to know the number of atoms in advance
 169 }
 170
 171 //! PME initialization with atom data
 172 PmeSafePointer pmeInitAtoms(const t_inputrec         *inputRec,
 173                             CodePath                  mode,
 174                             const gmx_device_info_t  *gpuInfo,
 175                             PmeGpuProgramHandle       pmeGpuProgram,
 176                             const CoordinatesVector  &coordinates,
 177                             const ChargesVector      &charges,
 178                             const Matrix3x3          &box
 179                             )
 180 {
 181     const index     atomCount = coordinates.size();
 182     GMX_RELEASE_ASSERT(atomCount == charges.ssize(), "Mismatch in atom data");
 183     PmeSafePointer  pmeSafe = pmeInitInternal(inputRec, mode, gpuInfo, pmeGpuProgram, box);
 184     PmeAtomComm    *atc     = nullptr;
 185
 186     switch (mode)
 187     {
 188         case CodePath::CPU:
 189             atc              = &(pmeSafe->atc[0]);
 190             atc->x           = coordinates;
 191             atc->coefficient = charges;
 192             gmx_pme_reinit_atoms(pmeSafe.get(), atomCount, charges.data());
 193             /* With decomposition there would be more boilerplate atc code here, e.g. do_redist_pos_coeffs */
 194             break;
 195
 196         case CodePath::GPU:
 197             // TODO: Avoid use of atc in the GPU code path
 198             atc              = &(pmeSafe->atc[0]);
 199             // We need to set atc->n for passing the size in the tests
 200             atc->setNumAtoms(atomCount);
 201             gmx_pme_reinit_atoms(pmeSafe.get(), atomCount, charges.data());
 202             pme_gpu_copy_input_coordinates(pmeSafe->gpu, as_rvec_array(coordinates.data()));
 203             break;
 204
 205         default:
 206             GMX_THROW(InternalError("Test not implemented for this mode"));
 207     }
 208
 209     return pmeSafe;
 210 }
 211
 212 //! Getting local PME real grid pointer for test I/O
 213 static real *pmeGetRealGridInternal(const gmx_pme_t *pme)
 214 {
 215     const size_t gridIndex = 0;
 216     return pme->fftgrid[gridIndex];
 217 }
 218
 219 //! Getting local PME real grid dimensions
 220 static void pmeGetRealGridSizesInternal(const gmx_pme_t      *pme,
 221                                         CodePath              mode,
 222                                         IVec                 &gridSize,       //NOLINT(google-runtime-references)
 223                                         IVec                 &paddedGridSize) //NOLINT(google-runtime-references)
 224 {
 225     const size_t gridIndex = 0;
 226     IVec         gridOffsetUnused;
 227     switch (mode)
 228     {
 229         case CodePath::CPU:
 230             gmx_parallel_3dfft_real_limits(pme->pfft_setup[gridIndex], gridSize, gridOffsetUnused, paddedGridSize);
 231             break;
 232
 233         case CodePath::GPU:
 234             pme_gpu_get_real_grid_sizes(pme->gpu, &gridSize, &paddedGridSize);
 235             break;
 236
 237         default:
 238             GMX_THROW(InternalError("Test not implemented for this mode"));
 239     }
 240 }
 241
 242 //! Getting local PME complex grid pointer for test I/O
 243 static t_complex *pmeGetComplexGridInternal(const gmx_pme_t *pme)
 244 {
 245     const size_t gridIndex = 0;
 246     return pme->cfftgrid[gridIndex];
 247 }
 248
 249 //! Getting local PME complex grid dimensions
 250 static void pmeGetComplexGridSizesInternal(const gmx_pme_t      *pme,
 251                                            IVec                 &gridSize,       //NOLINT(google-runtime-references)
 252                                            IVec                 &paddedGridSize) //NOLINT(google-runtime-references)
 253 {
 254     const size_t gridIndex = 0;
 255     IVec         gridOffsetUnused, complexOrderUnused;
 256     gmx_parallel_3dfft_complex_limits(pme->pfft_setup[gridIndex], complexOrderUnused, gridSize, 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> static void pmeGetGridAndSizesInternal(const gmx_pme_t * /*unused*/, CodePath /*unused*/, ValueType * & /*unused*/, IVec & /*unused*/, IVec & /*unused*/) //NOLINT(google-runtime-references)
 261 {
 262     GMX_THROW(InternalError("Deleted function call"));
 263     // explicitly deleting general template does not compile in clang/icc, see https://llvm.org/bugs/show_bug.cgi?id=17537
 264 }
 265
 266 //! Getting the PME real grid memory buffer and its sizes
 267 template<> void pmeGetGridAndSizesInternal<real>(const gmx_pme_t *pme, CodePath mode, real * &grid, IVec &gridSize, IVec &paddedGridSize)
 268 {
 269     grid = pmeGetRealGridInternal(pme);
 270     pmeGetRealGridSizesInternal(pme, mode, gridSize, paddedGridSize);
 271 }
 272
 273 //! Getting the PME complex grid memory buffer and its sizes
 274 template<> void pmeGetGridAndSizesInternal<t_complex>(const gmx_pme_t *pme, CodePath /*unused*/, t_complex * &grid, IVec &gridSize, IVec &paddedGridSize)
 275 {
 276     grid = pmeGetComplexGridInternal(pme);
 277     pmeGetComplexGridSizesInternal(pme, gridSize, paddedGridSize);
 278 }
 279
 280 //! PME spline calculation and charge spreading
 281 void pmePerformSplineAndSpread(gmx_pme_t *pme, CodePath mode, // TODO const qualifiers elsewhere
 282                                bool computeSplines, bool spreadCharges)
 283 {
 284     GMX_RELEASE_ASSERT(pme != nullptr, "PME data is not initialized");
 285     PmeAtomComm    *atc                          = &(pme->atc[0]);
 286     const size_t    gridIndex                    = 0;
 287     const bool      computeSplinesForZeroCharges = true;
 288     real           *fftgrid                      = spreadCharges ? pme->fftgrid[gridIndex] : nullptr;
 289     real           *pmegrid                      = pme->pmegrid[gridIndex].grid.grid;
 290
 291     switch (mode)
 292     {
 293         case CodePath::CPU:
 294             spread_on_grid(pme, atc, &pme->pmegrid[gridIndex], computeSplines, spreadCharges,
 295                            fftgrid, computeSplinesForZeroCharges, gridIndex);
 296             if (spreadCharges && !pme->bUseThreads)
 297             {
 298                 wrap_periodic_pmegrid(pme, pmegrid);
 299                 copy_pmegrid_to_fftgrid(pme, pmegrid, fftgrid, gridIndex);
 300             }
 301             break;
 302
 303         case CodePath::GPU:
 304             pme_gpu_spread(pme->gpu, gridIndex, fftgrid, computeSplines, spreadCharges);
 305             break;
 306
 307         default:
 308             GMX_THROW(InternalError("Test not implemented for this mode"));
 309     }
 310 }
 311
 312 //! Getting the internal spline data buffer pointer
 313 static real *pmeGetSplineDataInternal(const gmx_pme_t *pme, PmeSplineDataType type, int dimIndex)
 314 {
 315     GMX_ASSERT((0 <= dimIndex) && (dimIndex < DIM), "Invalid dimension index");
 316     const PmeAtomComm    *atc          = &(pme->atc[0]);
 317     const size_t          threadIndex  = 0;
 318     real                 *splineBuffer = nullptr;
 319     switch (type)
 320     {
 321         case PmeSplineDataType::Values:
 322             splineBuffer = atc->spline[threadIndex].theta.coefficients[dimIndex];
 323             break;
 324
 325         case PmeSplineDataType::Derivatives:
 326             splineBuffer = atc->spline[threadIndex].dtheta.coefficients[dimIndex];
 327             break;
 328
 329         default:
 330             GMX_THROW(InternalError("Unknown spline data type"));
 331     }
 332     return splineBuffer;
 333 }
 334
 335 //! PME solving
 336 void pmePerformSolve(const gmx_pme_t *pme, CodePath mode,
 337                      PmeSolveAlgorithm method, real cellVolume,
 338                      GridOrdering gridOrdering, bool computeEnergyAndVirial)
 339 {
 340     t_complex      *h_grid                 = pmeGetComplexGridInternal(pme);
 341     const bool      useLorentzBerthelot    = false;
 342     const size_t    threadIndex            = 0;
 343     switch (mode)
 344     {
 345         case CodePath::CPU:
 346             if (gridOrdering != GridOrdering::YZX)
 347             {
 348                 GMX_THROW(InternalError("Test not implemented for this mode"));
 349             }
 350             switch (method)
 351             {
 352                 case PmeSolveAlgorithm::Coulomb:
 353                     solve_pme_yzx(pme, h_grid, cellVolume,
 354                                   computeEnergyAndVirial, pme->nthread, threadIndex);
 355                     break;
 356
 357                 case PmeSolveAlgorithm::LennardJones:
 358                     solve_pme_lj_yzx(pme, &h_grid, useLorentzBerthelot,
 359                                      cellVolume, computeEnergyAndVirial, pme->nthread, threadIndex);
 360                     break;
 361
 362                 default:
 363                     GMX_THROW(InternalError("Test not implemented for this mode"));
 364             }
 365             break;
 366
 367         case CodePath::GPU:
 368             switch (method)
 369             {
 370                 case PmeSolveAlgorithm::Coulomb:
 371                     pme_gpu_solve(pme->gpu, h_grid, gridOrdering, computeEnergyAndVirial);
 372                     break;
 373
 374                 default:
 375                     GMX_THROW(InternalError("Test not implemented for this mode"));
 376             }
 377             break;
 378
 379         default:
 380             GMX_THROW(InternalError("Test not implemented for this mode"));
 381     }
 382 }
 383
 384 //! PME force gathering
 385 void pmePerformGather(gmx_pme_t *pme, CodePath mode,
 386                       PmeForceOutputHandling inputTreatment, ForcesVector &forces)
 387 {
 388     PmeAtomComm    *atc                     = &(pme->atc[0]);
 389     const index     atomCount               = atc->numAtoms();
 390     GMX_RELEASE_ASSERT(forces.ssize() == atomCount, "Invalid force buffer size");
 391     const bool      forceReductionWithInput = (inputTreatment == PmeForceOutputHandling::ReduceWithInput);
 392     const real      scale                   = 1.0;
 393     const size_t    threadIndex             = 0;
 394     const size_t    gridIndex               = 0;
 395     real           *pmegrid                 = pme->pmegrid[gridIndex].grid.grid;
 396     real           *fftgrid                 = pme->fftgrid[gridIndex];
 397
 398     switch (mode)
 399     {
 400         case CodePath::CPU:
 401             atc->f = forces;
 402             if (atc->nthread == 1)
 403             {
 404                 // something which is normally done in serial spline computation (make_thread_local_ind())
 405                 atc->spline[threadIndex].n = atomCount;
 406             }
 407             copy_fftgrid_to_pmegrid(pme, fftgrid, pmegrid, gridIndex, pme->nthread, threadIndex);
 408             unwrap_periodic_pmegrid(pme, pmegrid);
 409             gather_f_bsplines(pme, pmegrid, !forceReductionWithInput, atc, &atc->spline[threadIndex], scale);
 410             break;
 411
 412         case CodePath::GPU:
 413         {
 414             // Variable initialization needs a non-switch scope
 415             PmeOutput output = pme_gpu_getOutput(*pme, GMX_PME_CALC_F);
 416             GMX_ASSERT(forces.size() == output.forces_.size(), "Size of force buffers did not match");
 417             if (forceReductionWithInput)
 418             {
 419                 std::copy(std::begin(forces), std::end(forces), std::begin(output.forces_));
 420             }
 421             pme_gpu_gather(pme->gpu, inputTreatment, reinterpret_cast<float *>(fftgrid));
 422             std::copy(std::begin(output.forces_), std::end(output.forces_), std::begin(forces));
 423         }
 424         break;
 425
 426         default:
 427             GMX_THROW(InternalError("Test not implemented for this mode"));
 428     }
 429 }
 430
 431 //! PME test finalization before fetching the outputs
 432 void pmeFinalizeTest(const gmx_pme_t *pme, CodePath mode)
 433 {
 434     switch (mode)
 435     {
 436         case CodePath::CPU:
 437             break;
 438
 439         case CodePath::GPU:
 440             pme_gpu_synchronize(pme->gpu);
 441             break;
 442
 443         default:
 444             GMX_THROW(InternalError("Test not implemented for this mode"));
 445     }
 446 }
 447
 448 //! Setting atom spline values/derivatives to be used in spread/gather
 449 void pmeSetSplineData(const gmx_pme_t *pme, CodePath mode,
 450                       const SplineParamsDimVector &splineValues, PmeSplineDataType type, int dimIndex)
 451 {
 452     const PmeAtomComm    *atc         = &(pme->atc[0]);
 453     const index           atomCount   = atc->numAtoms();
 454     const index           pmeOrder    = pme->pme_order;
 455     const index           dimSize     = pmeOrder * atomCount;
 456     GMX_RELEASE_ASSERT(dimSize == splineValues.ssize(), "Mismatch in spline data");
 457     real                 *splineBuffer = pmeGetSplineDataInternal(pme, type, dimIndex);
 458
 459     switch (mode)
 460     {
 461         case CodePath::CPU:
 462             std::copy(splineValues.begin(), splineValues.end(), splineBuffer);
 463             break;
 464
 465         case CodePath::GPU:
 466             std::copy(splineValues.begin(), splineValues.end(), splineBuffer);
 467             pme_gpu_transform_spline_atom_data(pme->gpu, atc, type, dimIndex, PmeLayoutTransform::HostToGpu);
 468             break;
 469
 470         default:
 471             GMX_THROW(InternalError("Test not implemented for this mode"));
 472     }
 473 }
 474
 475 //! Setting gridline indices to be used in spread/gather
 476 void pmeSetGridLineIndices(gmx_pme_t *pme, CodePath mode,
 477                            const GridLineIndicesVector &gridLineIndices)
 478 {
 479     PmeAtomComm                *atc         = &(pme->atc[0]);
 480     const index                 atomCount   = atc->numAtoms();
 481     GMX_RELEASE_ASSERT(atomCount == gridLineIndices.ssize(), "Mismatch in gridline indices size");
 482
 483     IVec paddedGridSizeUnused, gridSize(0, 0, 0);
 484     pmeGetRealGridSizesInternal(pme, mode, gridSize, paddedGridSizeUnused);
 485
 486     for (const auto &index : gridLineIndices)
 487     {
 488         for (int i = 0; i < DIM; i++)
 489         {
 490             GMX_RELEASE_ASSERT((0 <= index[i]) && (index[i] < gridSize[i]), "Invalid gridline index");
 491         }
 492     }
 493
 494     switch (mode)
 495     {
 496         case CodePath::GPU:
 497             memcpy(pme->gpu->staging.h_gridlineIndices, gridLineIndices.data(), atomCount * sizeof(gridLineIndices[0]));
 498             break;
 499
 500         case CodePath::CPU:
 501             atc->idx.resize(gridLineIndices.size());
 502             std::copy(gridLineIndices.begin(), gridLineIndices.end(), atc->idx.begin());
 503             break;
 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 PmeAtomComm       *atc         = &(pme->atc[0]);
 581     const size_t             atomCount   = atc->numAtoms();
 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 PmeAtomComm    *atc         = &(pme->atc[0]);
 609     const size_t          atomCount   = atc->numAtoms();
 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 = atc->idx;
 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 PmeOutput pmeGetReciprocalEnergyAndVirial(const gmx_pme_t *pme, CodePath mode,
 681                                           PmeSolveAlgorithm method)
 682 {
 683     PmeOutput output;
 684     switch (mode)
 685     {
 686         case CodePath::CPU:
 687             switch (method)
 688             {
 689                 case PmeSolveAlgorithm::Coulomb:
 690                     get_pme_ener_vir_q(pme->solve_work, pme->nthread, &output);
 691                     break;
 692
 693                 case PmeSolveAlgorithm::LennardJones:
 694                     get_pme_ener_vir_lj(pme->solve_work, pme->nthread, &output);
 695                     break;
 696
 697                 default:
 698                     GMX_THROW(InternalError("Test not implemented for this mode"));
 699             }
 700             break;
 701         case CodePath::GPU:
 702             switch (method)
 703             {
 704                 case PmeSolveAlgorithm::Coulomb:
 705                     output = pme_gpu_getEnergyAndVirial(*pme);
 706                     break;
 707
 708                 default:
 709                     GMX_THROW(InternalError("Test not implemented for this mode"));
 710             }
 711             break;
 712
 713         default:
 714             GMX_THROW(InternalError("Test not implemented for this mode"));
 715     }
 716     return output;
 717 }
 718
 719 }  // namespace test
 720 }  // namespace gmx