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