src/gromacs/ewald/tests/pmetestcommon.cpp

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