2 * This file is part of the GROMACS molecular simulation package.
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.
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.
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.
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.
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.
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.
38 * \brief This file contains internal function definitions for performing the PME calculations on GPU.
39 * These are not meant to be exposed outside of the PME GPU code.
40 * As of now, their bodies are still in the common pme_gpu.cpp files.
42 * \author Aleksei Iupinov <a.yupinov@gmail.com>
43 * \ingroup module_ewald
46 #ifndef GMX_EWALD_PME_GPU_INTERNAL_H
47 #define GMX_EWALD_PME_GPU_INTERNAL_H
49 #include "gromacs/fft/fft.h" // for the gmx_fft_direction enum
50 #include "gromacs/gpu_utils/gpu_macros.h" // for the GPU_FUNC_ macros
51 #include "gromacs/utility/arrayref.h"
53 #include "pme_gpu_types_host.h" // for the inline functions accessing PmeGpu members
57 struct gmx_pme_t; // only used in pme_gpu_reinit
58 struct gmx_wallclock_gpu_pme_t;
67 //! Type of spline data
68 enum class PmeSplineDataType
71 Derivatives, // dtheta
72 }; // TODO move this into new and shiny pme.h (pme-types.h?)
74 //! PME grid dimension ordering (from major to minor)
75 enum class GridOrdering
81 /*! \libinternal \brief
82 * Returns the number of atoms per chunk in the atom charges/coordinates data layout.
83 * Depends on CUDA-specific block sizes, needed for the atom data padding.
85 * \param[in] pmeGpu The PME GPU structure.
86 * \returns Number of atoms in a single GPU atom data chunk.
88 int pme_gpu_get_atom_data_alignment(const PmeGpu* pmeGpu);
90 /*! \libinternal \brief
91 * Returns the number of atoms per chunk in the atom spline theta/dtheta data layout.
93 * \param[in] pmeGpu The PME GPU structure.
94 * \returns Number of atoms in a single GPU atom spline data chunk.
96 int pme_gpu_get_atoms_per_warp(const PmeGpu* pmeGpu);
98 /*! \libinternal \brief
99 * Synchronizes the current computation, waiting for the GPU kernels/transfers to finish.
101 * \param[in] pmeGpu The PME GPU structure.
103 GPU_FUNC_QUALIFIER void pme_gpu_synchronize(const PmeGpu* GPU_FUNC_ARGUMENT(pmeGpu)) GPU_FUNC_TERM;
105 /*! \libinternal \brief
106 * Allocates the fixed size energy and virial buffer both on GPU and CPU.
108 * \param[in,out] pmeGpu The PME GPU structure.
110 void pme_gpu_alloc_energy_virial(PmeGpu* pmeGpu);
112 /*! \libinternal \brief
113 * Frees the energy and virial memory both on GPU and CPU.
115 * \param[in] pmeGpu The PME GPU structure.
117 void pme_gpu_free_energy_virial(PmeGpu* pmeGpu);
119 /*! \libinternal \brief
120 * Clears the energy and virial memory on GPU with 0.
121 * Should be called at the end of PME computation which returned energy/virial.
123 * \param[in] pmeGpu The PME GPU structure.
125 void pme_gpu_clear_energy_virial(const PmeGpu* pmeGpu);
127 /*! \libinternal \brief
128 * Reallocates and copies the pre-computed B-spline values to the GPU.
130 * \param[in,out] pmeGpu The PME GPU structure.
132 void pme_gpu_realloc_and_copy_bspline_values(PmeGpu* pmeGpu);
134 /*! \libinternal \brief
135 * Frees the pre-computed B-spline values on the GPU (and the transfer CPU buffers).
137 * \param[in] pmeGpu The PME GPU structure.
139 void pme_gpu_free_bspline_values(const PmeGpu* pmeGpu);
141 /*! \libinternal \brief
142 * Reallocates the GPU buffer for the PME forces.
144 * \param[in] pmeGpu The PME GPU structure.
146 void pme_gpu_realloc_forces(PmeGpu* pmeGpu);
148 /*! \libinternal \brief
149 * Frees the GPU buffer for the PME forces.
151 * \param[in] pmeGpu The PME GPU structure.
153 void pme_gpu_free_forces(const PmeGpu* pmeGpu);
155 /*! \libinternal \brief
156 * Copies the forces from the CPU buffer to the GPU (to reduce them with the PME GPU gathered
157 * forces). To be called e.g. after the bonded calculations.
159 * \param[in] pmeGpu The PME GPU structure.
161 void pme_gpu_copy_input_forces(PmeGpu* pmeGpu);
163 /*! \libinternal \brief
164 * Copies the forces from the GPU to the CPU buffer. To be called after the gathering stage.
166 * \param[in] pmeGpu The PME GPU structure.
168 void pme_gpu_copy_output_forces(PmeGpu* pmeGpu);
170 /*! \libinternal \brief
171 * Checks whether work in the PME GPU stream has completed.
173 * \param[in] pmeGpu The PME GPU structure.
175 * \returns True if work in the PME stream has completed.
177 bool pme_gpu_stream_query(const PmeGpu* pmeGpu);
179 /*! \libinternal \brief
180 * Reallocates the input coordinates buffer on the GPU (and clears the padded part if needed).
182 * \param[in] pmeGpu The PME GPU structure.
184 * Needs to be called on every DD step/in the beginning.
186 void pme_gpu_realloc_coordinates(const PmeGpu* pmeGpu);
188 /*! \libinternal \brief
189 * Frees the coordinates on the GPU.
191 * \param[in] pmeGpu The PME GPU structure.
193 void pme_gpu_free_coordinates(const PmeGpu* pmeGpu);
195 /*! \libinternal \brief
196 * Reallocates the buffer on the GPU and copies the charges/coefficients from the CPU buffer.
197 * Clears the padded part if needed.
199 * \param[in] pmeGpu The PME GPU structure.
200 * \param[in] h_coefficients The input atom charges/coefficients.
202 * Does not need to be done for every PME computation, only whenever the local charges change.
203 * (So, in the beginning of the run, or on DD step).
205 void pme_gpu_realloc_and_copy_input_coefficients(const PmeGpu* pmeGpu, const float* h_coefficients);
207 /*! \libinternal \brief
208 * Frees the charges/coefficients on the GPU.
210 * \param[in] pmeGpu The PME GPU structure.
212 void pme_gpu_free_coefficients(const PmeGpu* pmeGpu);
214 /*! \libinternal \brief
215 * Reallocates the buffers on the GPU and the host for the atoms spline data.
217 * \param[in,out] pmeGpu The PME GPU structure.
219 void pme_gpu_realloc_spline_data(PmeGpu* pmeGpu);
221 /*! \libinternal \brief
222 * Frees the buffers on the GPU for the atoms spline data.
224 * \param[in] pmeGpu The PME GPU structure.
226 void pme_gpu_free_spline_data(const PmeGpu* pmeGpu);
228 /*! \libinternal \brief
229 * Reallocates the buffers on the GPU and the host for the particle gridline indices.
231 * \param[in,out] pmeGpu The PME GPU structure.
233 void pme_gpu_realloc_grid_indices(PmeGpu* pmeGpu);
235 /*! \libinternal \brief
236 * Frees the buffer on the GPU for the particle gridline indices.
238 * \param[in] pmeGpu The PME GPU structure.
240 void pme_gpu_free_grid_indices(const PmeGpu* pmeGpu);
242 /*! \libinternal \brief
243 * Reallocates the real space grid and the complex reciprocal grid (if needed) on the GPU.
245 * \param[in] pmeGpu The PME GPU structure.
247 void pme_gpu_realloc_grids(PmeGpu* pmeGpu);
249 /*! \libinternal \brief
250 * Frees the real space grid and the complex reciprocal grid (if needed) on the GPU.
252 * \param[in] pmeGpu The PME GPU structure.
254 void pme_gpu_free_grids(const PmeGpu* pmeGpu);
256 /*! \libinternal \brief
257 * Clears the real space grid on the GPU.
258 * Should be called at the end of each computation.
260 * \param[in] pmeGpu The PME GPU structure.
262 void pme_gpu_clear_grids(const PmeGpu* pmeGpu);
264 /*! \libinternal \brief
265 * Reallocates and copies the pre-computed fractional coordinates' shifts to the GPU.
267 * \param[in] pmeGpu The PME GPU structure.
269 void pme_gpu_realloc_and_copy_fract_shifts(PmeGpu* pmeGpu);
271 /*! \libinternal \brief
272 * Frees the pre-computed fractional coordinates' shifts on the GPU.
274 * \param[in] pmeGpu The PME GPU structure.
276 void pme_gpu_free_fract_shifts(const PmeGpu* pmeGpu);
278 /*! \libinternal \brief
279 * Copies the input real-space grid from the host to the GPU.
281 * \param[in] pmeGpu The PME GPU structure.
282 * \param[in] h_grid The host-side grid buffer.
284 void pme_gpu_copy_input_gather_grid(const PmeGpu* pmeGpu, float* h_grid);
286 /*! \libinternal \brief
287 * Copies the output real-space grid from the GPU to the host.
289 * \param[in] pmeGpu The PME GPU structure.
290 * \param[out] h_grid The host-side grid buffer.
292 void pme_gpu_copy_output_spread_grid(const PmeGpu* pmeGpu, float* h_grid);
294 /*! \libinternal \brief
295 * Copies the spread output spline data and gridline indices from the GPU to the host.
297 * \param[in] pmeGpu The PME GPU structure.
299 void pme_gpu_copy_output_spread_atom_data(const PmeGpu* pmeGpu);
301 /*! \libinternal \brief
302 * Copies the gather input spline data and gridline indices from the host to the GPU.
304 * \param[in] pmeGpu The PME GPU structure.
306 void pme_gpu_copy_input_gather_atom_data(const PmeGpu* pmeGpu);
308 /*! \libinternal \brief
309 * Waits for the grid copying to the host-side buffer after spreading to finish.
311 * \param[in] pmeGpu The PME GPU structure.
313 void pme_gpu_sync_spread_grid(const PmeGpu* pmeGpu);
315 /*! \libinternal \brief
316 * Does the one-time GPU-framework specific PME initialization.
317 * For CUDA, the PME stream is created with the highest priority.
319 * \param[in] pmeGpu The PME GPU structure.
321 void pme_gpu_init_internal(PmeGpu* pmeGpu);
323 /*! \libinternal \brief
324 * Destroys the PME GPU-framework specific data.
325 * Should be called last in the PME GPU destructor.
327 * \param[in] pmeGpu The PME GPU structure.
329 void pme_gpu_destroy_specific(const PmeGpu* pmeGpu);
331 /*! \libinternal \brief
332 * Initializes the CUDA FFT structures.
334 * \param[in] pmeGpu The PME GPU structure.
336 void pme_gpu_reinit_3dfft(const PmeGpu* pmeGpu);
338 /*! \libinternal \brief
339 * Destroys the CUDA FFT structures.
341 * \param[in] pmeGpu The PME GPU structure.
343 void pme_gpu_destroy_3dfft(const PmeGpu* pmeGpu);
345 /* Several GPU event-based timing functions that live in pme_gpu_timings.cpp */
347 /*! \libinternal \brief
348 * Finalizes all the active PME GPU stage timings for the current computation. Should be called at the end of every computation.
350 * \param[in] pmeGpu The PME GPU structure.
352 void pme_gpu_update_timings(const PmeGpu* pmeGpu);
354 /*! \libinternal \brief
355 * Updates the internal list of active PME GPU stages (if timings are enabled).
357 * \param[in] pmeGpu The PME GPU data structure.
359 void pme_gpu_reinit_timings(const PmeGpu* pmeGpu);
362 * Resets the PME GPU timings. To be called at the reset MD step.
364 * \param[in] pmeGpu The PME GPU structure.
366 void pme_gpu_reset_timings(const PmeGpu* pmeGpu);
368 /*! \libinternal \brief
369 * Copies the PME GPU timings to the gmx_wallclock_gpu_t structure (for log output). To be called at the run end.
371 * \param[in] pmeGpu The PME GPU structure.
372 * \param[in] timings The gmx_wallclock_gpu_pme_t structure.
374 void pme_gpu_get_timings(const PmeGpu* pmeGpu, gmx_wallclock_gpu_pme_t* timings);
376 /* The PME stages themselves */
378 /*! \libinternal \brief
379 * A GPU spline computation and charge spreading function.
381 * \param[in] pmeGpu The PME GPU structure.
382 * \param[in] xReadyOnDevice Event synchronizer indicating that the coordinates are ready in the device memory;
383 * can be nullptr when invoked on a separate PME rank or from PME tests.
384 * \param[in] gridIndex Index of the PME grid - unused, assumed to be 0.
385 * \param[out] h_grid The host-side grid buffer (used only if the result of the spread is expected on the host,
386 * e.g. testing or host-side FFT)
387 * \param[in] computeSplines Should the computation of spline parameters and gridline indices be performed.
388 * \param[in] spreadCharges Should the charges/coefficients be spread on the grid.
390 GPU_FUNC_QUALIFIER void pme_gpu_spread(const PmeGpu* GPU_FUNC_ARGUMENT(pmeGpu),
391 GpuEventSynchronizer* GPU_FUNC_ARGUMENT(xReadyOnDevice),
392 int GPU_FUNC_ARGUMENT(gridIndex),
393 real* GPU_FUNC_ARGUMENT(h_grid),
394 bool GPU_FUNC_ARGUMENT(computeSplines),
395 bool GPU_FUNC_ARGUMENT(spreadCharges)) GPU_FUNC_TERM;
397 /*! \libinternal \brief
398 * 3D FFT R2C/C2R routine.
400 * \param[in] pmeGpu The PME GPU structure.
401 * \param[in] direction Transform direction (real-to-complex or complex-to-real)
402 * \param[in] gridIndex Index of the PME grid - unused, assumed to be 0.
404 void pme_gpu_3dfft(const PmeGpu* pmeGpu, enum gmx_fft_direction direction, int gridIndex);
406 /*! \libinternal \brief
407 * A GPU Fourier space solving function.
409 * \param[in] pmeGpu The PME GPU structure.
410 * \param[in,out] h_grid The host-side input and output Fourier grid buffer (used only with testing or host-side FFT)
411 * \param[in] gridOrdering Specifies the dimenion ordering of the complex grid. TODO: store this information?
412 * \param[in] computeEnergyAndVirial Tells if the energy and virial computation should also be performed.
414 GPU_FUNC_QUALIFIER void pme_gpu_solve(const PmeGpu* GPU_FUNC_ARGUMENT(pmeGpu),
415 t_complex* GPU_FUNC_ARGUMENT(h_grid),
416 GridOrdering GPU_FUNC_ARGUMENT(gridOrdering),
417 bool GPU_FUNC_ARGUMENT(computeEnergyAndVirial)) GPU_FUNC_TERM;
419 /*! \libinternal \brief
420 * A GPU force gathering function.
422 * \param[in] pmeGpu The PME GPU structure.
423 * \param[in] forceTreatment Tells how data in h_forces should be treated.
424 * TODO: determine efficiency/balance of host/device-side
425 * reductions. \param[in] h_grid The host-side grid buffer (used only in testing mode)
427 GPU_FUNC_QUALIFIER void pme_gpu_gather(PmeGpu* GPU_FUNC_ARGUMENT(pmeGpu),
428 PmeForceOutputHandling GPU_FUNC_ARGUMENT(forceTreatment),
429 const float* GPU_FUNC_ARGUMENT(h_grid)) GPU_FUNC_TERM;
431 /*! \brief Return pointer to device copy of coordinate data.
432 * \param[in] pmeGpu The PME GPU structure.
433 * \returns Pointer to coordinate data
435 GPU_FUNC_QUALIFIER DeviceBuffer<float> pme_gpu_get_kernelparam_coordinates(const PmeGpu* GPU_FUNC_ARGUMENT(pmeGpu))
436 GPU_FUNC_TERM_WITH_RETURN(DeviceBuffer<float>{});
438 /*! \brief Sets the device pointer to coordinate data
439 * \param[in] pmeGpu The PME GPU structure.
440 * \param[in] d_x Pointer to coordinate data
442 GPU_FUNC_QUALIFIER void pme_gpu_set_kernelparam_coordinates(const PmeGpu* GPU_FUNC_ARGUMENT(pmeGpu),
443 DeviceBuffer<float> GPU_FUNC_ARGUMENT(d_x)) GPU_FUNC_TERM;
445 /*! \brief Return pointer to device copy of force data.
446 * \param[in] pmeGpu The PME GPU structure.
447 * \returns Pointer to force data
449 GPU_FUNC_QUALIFIER void* pme_gpu_get_kernelparam_forces(const PmeGpu* GPU_FUNC_ARGUMENT(pmeGpu))
450 GPU_FUNC_TERM_WITH_RETURN(nullptr);
452 /*! \brief Return pointer to GPU stream.
453 * \param[in] pmeGpu The PME GPU structure.
454 * \returns Pointer to stream object.
456 GPU_FUNC_QUALIFIER void* pme_gpu_get_stream(const PmeGpu* GPU_FUNC_ARGUMENT(pmeGpu))
457 GPU_FUNC_TERM_WITH_RETURN(nullptr);
459 /*! \brief Return pointer to GPU context (for OpenCL builds).
460 * \param[in] pmeGpu The PME GPU structure.
461 * \returns Pointer to context object.
463 GPU_FUNC_QUALIFIER void* pme_gpu_get_context(const PmeGpu* GPU_FUNC_ARGUMENT(pmeGpu))
464 GPU_FUNC_TERM_WITH_RETURN(nullptr);
466 /*! \brief Return pointer to the sync object triggered after the PME force calculation completion
467 * \param[in] pmeGpu The PME GPU structure.
468 * \returns Pointer to sync object
470 GPU_FUNC_QUALIFIER GpuEventSynchronizer* pme_gpu_get_forces_ready_synchronizer(
471 const PmeGpu* GPU_FUNC_ARGUMENT(pmeGpu)) GPU_FUNC_TERM_WITH_RETURN(nullptr);
473 /* The inlined convenience PME GPU status getters */
475 /*! \libinternal \brief
476 * Tells if PME runs on multiple GPUs with the decomposition.
478 * \param[in] pmeGpu The PME GPU structure.
479 * \returns True if PME runs on multiple GPUs, false otherwise.
481 inline bool pme_gpu_uses_dd(const PmeGpu* pmeGpu)
483 return !pmeGpu->settings.useDecomposition;
486 /*! \libinternal \brief
487 * Tells if PME performs the gathering stage on GPU.
489 * \param[in] pmeGpu The PME GPU structure.
490 * \returns True if the gathering is performed on GPU, false otherwise.
492 inline bool pme_gpu_performs_gather(const PmeGpu* pmeGpu)
494 return pmeGpu->settings.performGPUGather;
497 /*! \libinternal \brief
498 * Tells if PME performs the FFT stages on GPU.
500 * \param[in] pmeGpu The PME GPU structure.
501 * \returns True if FFT is performed on GPU, false otherwise.
503 inline bool pme_gpu_performs_FFT(const PmeGpu* pmeGpu)
505 return pmeGpu->settings.performGPUFFT;
508 /*! \libinternal \brief
509 * Tells if PME performs the grid (un-)wrapping on GPU.
511 * \param[in] pmeGpu The PME GPU structure.
512 * \returns True if (un-)wrapping is performed on GPU, false otherwise.
514 inline bool pme_gpu_performs_wrapping(const PmeGpu* pmeGpu)
516 return pmeGpu->settings.useDecomposition;
519 /*! \libinternal \brief
520 * Tells if PME performs the grid solving on GPU.
522 * \param[in] pmeGpu The PME GPU structure.
523 * \returns True if solving is performed on GPU, false otherwise.
525 inline bool pme_gpu_performs_solve(const PmeGpu* pmeGpu)
527 return pmeGpu->settings.performGPUSolve;
530 /*! \libinternal \brief
531 * Enables or disables the testing mode.
532 * Testing mode only implies copying all the outputs, even the intermediate ones, to the host,
533 * and also makes the copies synchronous.
535 * \param[in] pmeGpu The PME GPU structure.
536 * \param[in] testing Should the testing mode be enabled, or disabled.
538 inline void pme_gpu_set_testing(PmeGpu* pmeGpu, bool testing)
542 pmeGpu->settings.copyAllOutputs = testing;
543 pmeGpu->settings.transferKind = testing ? GpuApiCallBehavior::Sync : GpuApiCallBehavior::Async;
547 /*! \libinternal \brief
548 * Tells if PME is in the testing mode.
550 * \param[in] pmeGpu The PME GPU structure.
551 * \returns true if testing mode is enabled, false otherwise.
553 inline bool pme_gpu_is_testing(const PmeGpu* pmeGpu)
555 return pmeGpu->settings.copyAllOutputs;
558 /* A block of C++ functions that live in pme_gpu_internal.cpp */
560 /*! \libinternal \brief
561 * Returns the energy and virial GPU outputs, useful for testing.
563 * It is the caller's responsibility to be aware of whether the GPU
564 * handled the solve stage.
566 * \param[in] pme The PME structure.
567 * \param[out] output Pointer to output where energy and virial should be stored.
569 GPU_FUNC_QUALIFIER void pme_gpu_getEnergyAndVirial(const gmx_pme_t& GPU_FUNC_ARGUMENT(pme),
570 PmeOutput* GPU_FUNC_ARGUMENT(output)) GPU_FUNC_TERM;
572 /*! \libinternal \brief
573 * Returns the GPU outputs (forces, energy and virial)
575 * \param[in] pme The PME structure.
576 * \param[in] flags The combination of flags that affected this PME computation.
577 * The flags are the GMX_PME_ flags from pme.h.
578 * \returns The output object.
580 GPU_FUNC_QUALIFIER PmeOutput pme_gpu_getOutput(const gmx_pme_t& GPU_FUNC_ARGUMENT(pme),
581 int GPU_FUNC_ARGUMENT(flags))
582 GPU_FUNC_TERM_WITH_RETURN(PmeOutput{});
584 /*! \libinternal \brief
585 * Updates the unit cell parameters. Does not check if update is necessary - that is done in pme_gpu_prepare_computation().
587 * \param[in] pmeGpu The PME GPU structure.
588 * \param[in] box The unit cell box.
590 GPU_FUNC_QUALIFIER void pme_gpu_update_input_box(PmeGpu* GPU_FUNC_ARGUMENT(pmeGpu),
591 const matrix GPU_FUNC_ARGUMENT(box)) GPU_FUNC_TERM;
593 /*! \libinternal \brief
594 * Finishes the PME GPU computation, waiting for the output forces and/or energy/virial to be copied to the host.
595 * If forces were computed, they will have arrived at the external host buffer provided to gather.
596 * If virial/energy were computed, they will have arrived into the internal staging buffer
597 * (even though that should have already happened before even launching the gather).
598 * Finally, cudaEvent_t based GPU timers get updated if enabled. They also need stream synchronization for correctness.
599 * Additionally, device-side buffers are cleared asynchronously for the next computation.
601 * \param[in] pmeGpu The PME GPU structure.
603 void pme_gpu_finish_computation(const PmeGpu* pmeGpu);
605 //! A binary enum for spline data layout transformation
606 enum class PmeLayoutTransform
612 /*! \libinternal \brief
613 * Rearranges the atom spline data between the GPU and host layouts.
614 * Only used for test purposes so far, likely to be horribly slow.
616 * \param[in] pmeGpu The PME GPU structure.
617 * \param[out] atc The PME CPU atom data structure (with a single-threaded layout).
618 * \param[in] type The spline data type (values or derivatives).
619 * \param[in] dimIndex Dimension index.
620 * \param[in] transform Layout transform type
622 GPU_FUNC_QUALIFIER void pme_gpu_transform_spline_atom_data(const PmeGpu* GPU_FUNC_ARGUMENT(pmeGpu),
623 const PmeAtomComm* GPU_FUNC_ARGUMENT(atc),
624 PmeSplineDataType GPU_FUNC_ARGUMENT(type),
625 int GPU_FUNC_ARGUMENT(dimIndex),
626 PmeLayoutTransform GPU_FUNC_ARGUMENT(transform)) GPU_FUNC_TERM;
628 /*! \libinternal \brief
629 * Gets a unique index to an element in a spline parameter buffer (theta/dtheta),
630 * which is laid out for GPU spread/gather kernels. The index is wrt the execution block,
631 * in range(0, atomsPerBlock * order * DIM).
632 * This is a wrapper, only used in unit tests.
633 * \param[in] order PME order
634 * \param[in] splineIndex Spline contribution index (from 0 to \p order - 1)
635 * \param[in] dimIndex Dimension index (from 0 to 2)
636 * \param[in] atomIndex Atom index wrt the block.
637 * \param[in] atomsPerWarp Number of atoms processed by a warp.
639 * \returns Index into theta or dtheta array using GPU layout.
641 int getSplineParamFullIndex(int order, int splineIndex, int dimIndex, int atomIndex, int atomsPerWarp);
643 /*! \libinternal \brief
644 * Get the normal/padded grid dimensions of the real-space PME grid on GPU. Only used in tests.
646 * \param[in] pmeGpu The PME GPU structure.
647 * \param[out] gridSize Pointer to the grid dimensions to fill in.
648 * \param[out] paddedGridSize Pointer to the padded grid dimensions to fill in.
650 GPU_FUNC_QUALIFIER void pme_gpu_get_real_grid_sizes(const PmeGpu* GPU_FUNC_ARGUMENT(pmeGpu),
651 gmx::IVec* GPU_FUNC_ARGUMENT(gridSize),
652 gmx::IVec* GPU_FUNC_ARGUMENT(paddedGridSize)) GPU_FUNC_TERM;
654 /*! \libinternal \brief
655 * (Re-)initializes the PME GPU data at the beginning of the run or on DLB.
657 * \param[in,out] pme The PME structure.
658 * \param[in] gpuInfo The GPU information structure.
659 * \param[in] pmeGpuProgram The PME GPU program data
660 * \throws gmx::NotImplementedError if this generally valid PME structure is not valid for GPU runs.
662 GPU_FUNC_QUALIFIER void pme_gpu_reinit(gmx_pme_t* GPU_FUNC_ARGUMENT(pme),
663 const gmx_device_info_t* GPU_FUNC_ARGUMENT(gpuInfo),
664 PmeGpuProgramHandle GPU_FUNC_ARGUMENT(pmeGpuProgram)) GPU_FUNC_TERM;
666 /*! \libinternal \brief
667 * Destroys the PME GPU data at the end of the run.
669 * \param[in] pmeGpu The PME GPU structure.
671 GPU_FUNC_QUALIFIER void pme_gpu_destroy(PmeGpu* GPU_FUNC_ARGUMENT(pmeGpu)) GPU_FUNC_TERM;
673 /*! \libinternal \brief
674 * Reallocates the local atoms data (charges, coordinates, etc.). Copies the charges to the GPU.
676 * \param[in] pmeGpu The PME GPU structure.
677 * \param[in] nAtoms The number of particles.
678 * \param[in] charges The pointer to the host-side array of particle charges.
680 * This is a function that should only be called in the beginning of the run and on domain
681 * decomposition. Should be called before the pme_gpu_set_io_ranges.
683 GPU_FUNC_QUALIFIER void pme_gpu_reinit_atoms(PmeGpu* GPU_FUNC_ARGUMENT(pmeGpu),
684 int GPU_FUNC_ARGUMENT(nAtoms),
685 const real* GPU_FUNC_ARGUMENT(charges)) GPU_FUNC_TERM;
687 /*! \brief \libinternal
688 * The PME GPU reinitialization function that is called both at the end of any PME computation and on any load balancing.
690 * This clears the device-side working buffers in preparation for new computation.
692 * \param[in] pmeGpu The PME GPU structure.
694 void pme_gpu_reinit_computation(const PmeGpu* pmeGpu);
697 * Blocks until PME GPU tasks are completed, and gets the output forces and virial/energy
698 * (if they were to be computed).
700 * \param[in] pme The PME data structure.
701 * \param[in] flags The combination of flags to affect this PME computation.
702 * The flags are the GMX_PME_ flags from pme.h.
703 * \param[out] wcycle The wallclock counter.
704 * \return The output forces, energy and virial
706 GPU_FUNC_QUALIFIER PmeOutput pme_gpu_wait_finish_task(gmx_pme_t* GPU_FUNC_ARGUMENT(pme),
707 int GPU_FUNC_ARGUMENT(flags),
708 gmx_wallcycle* GPU_FUNC_ARGUMENT(wcycle))
709 GPU_FUNC_TERM_WITH_RETURN(PmeOutput{});