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