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