Set temperature scaling groups upon construction of integrator
[alexxy/gromacs.git] / src / gromacs / mdlib / leapfrog_gpu.h
1 /*
2  * This file is part of the GROMACS molecular simulation package.
3  *
4  * Copyright (c) 2019,2020,2021, 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 /*! \libinternal \file
36  *
37  * \brief Declarations for GPU implementation of Leap-Frog.
38  *
39  * \author Artem Zhmurov <zhmurov@gmail.com>
40  *
41  * \ingroup module_mdlib
42  * \inlibraryapi
43  */
44 #ifndef GMX_MDLIB_LEAPFROG_GPU_H
45 #define GMX_MDLIB_LEAPFROG_GPU_H
46
47 #include "config.h"
48
49 #if GMX_GPU_CUDA
50 #    include "gromacs/gpu_utils/devicebuffer.cuh"
51 #    include "gromacs/gpu_utils/gputraits.cuh"
52 #endif
53 #if GMX_GPU_SYCL
54 #    include "gromacs/gpu_utils/devicebuffer_sycl.h"
55 #    include "gromacs/gpu_utils/gputraits_sycl.h"
56 #endif
57
58 #include <memory>
59
60 #include "gromacs/gpu_utils/hostallocator.h"
61 #include "gromacs/pbcutil/pbc.h"
62 #include "gromacs/pbcutil/pbc_aiuc.h"
63 #include "gromacs/utility/arrayref.h"
64
65 class DeviceContext;
66 class DeviceStream;
67 struct t_grp_tcstat;
68
69 namespace gmx
70 {
71
72
73 /*! \brief Sets the number of different temperature coupling values
74  *
75  *  This is needed to template the kernel
76  *  \todo Unify with similar enum in CPU update module
77  */
78 enum class NumTempScaleValues
79 {
80     None     = 0, //!< No temperature coupling
81     Single   = 1, //!< Single T-scaling value (one group)
82     Multiple = 2, //!< Multiple T-scaling values, need to use T-group indices
83     Count    = 3  //!< Number of valid values
84 };
85
86 /*! \brief Different variants of the Parrinello-Rahman velocity scaling
87  *
88  *  This is needed to template the kernel
89  *  \todo Unify with similar enum in CPU update module
90  */
91 enum class VelocityScalingType
92 {
93     None     = 0, //!< Do not apply velocity scaling (not a PR-coupling run or step)
94     Diagonal = 1, //!< Apply velocity scaling using a diagonal matrix
95     Count    = 2  //!< Number of valid values
96 };
97
98 class LeapFrogGpu
99 {
100
101 public:
102     /*! \brief Constructor.
103      *
104      * \param[in] deviceContext       Device context (dummy in CUDA).
105      * \param[in] deviceStream        Device stream to use.
106      * \param[in] numTempScaleValues  Number of temperature scale groups.
107      */
108     LeapFrogGpu(const DeviceContext& deviceContext, const DeviceStream& deviceStream, int numTempScaleValues);
109     ~LeapFrogGpu();
110
111     /*! \brief Integrate
112      *
113      * Integrates the equation of motion using Leap-Frog algorithm.
114      * Updates coordinates and velocities on the GPU. The current coordinates are saved for constraints.
115      *
116      * \param[in,out] d_x                      Coordinates to update
117      * \param[out]    d_xp                     Place to save the values of initial coordinates coordinates to.
118      * \param[in,out] d_v                      Velocities (will be updated).
119      * \param[in]     d_f                      Forces.
120      * \param[in]     dt                       Timestep.
121      * \param[in]     doTemperatureScaling     If velocities should be scaled for temperature coupling.
122      * \param[in]     tcstat                   Temperature coupling data.
123      * \param[in]     doParrinelloRahman       If current step is a Parrinello-Rahman pressure coupling step.
124      * \param[in]     dtPressureCouple         Period between pressure coupling steps
125      * \param[in]     prVelocityScalingMatrix  Parrinello-Rahman velocity scaling matrix
126      */
127     void integrate(const DeviceBuffer<float3>        d_x,
128                    DeviceBuffer<float3>              d_xp,
129                    DeviceBuffer<float3>              d_v,
130                    const DeviceBuffer<float3>        d_f,
131                    const real                        dt,
132                    const bool                        doTemperatureScaling,
133                    gmx::ArrayRef<const t_grp_tcstat> tcstat,
134                    const bool                        doParrinelloRahman,
135                    const float                       dtPressureCouple,
136                    const matrix                      prVelocityScalingMatrix);
137
138     /*! \brief Set the integrator
139      *
140      * Allocates memory for inverse masses, and, if needed for temperature scaling factor(s)
141      * and temperature coupling groups. Copies inverse masses and temperature coupling groups
142      * to the GPU.
143      *
144      * \param[in] numAtoms        Number of atoms in the system.
145      * \param[in] inverseMasses   Inverse masses of atoms.
146      * \param[in] tempScaleGroups Maps the atom index to temperature scale value.
147      */
148     void set(const int numAtoms, const real* inverseMasses, const unsigned short* tempScaleGroups);
149
150     /*! \brief Class with hardware-specific interfaces and implementations.*/
151     class Impl;
152
153 private:
154     //! GPU context object
155     const DeviceContext& deviceContext_;
156     //! GPU stream
157     const DeviceStream& deviceStream_;
158     //! GPU kernel launch config
159     KernelLaunchConfig kernelLaunchConfig_;
160     //! Number of atoms
161     int numAtoms_;
162
163     //! 1/mass for all atoms (GPU)
164     DeviceBuffer<float> d_inverseMasses_;
165     //! Current size of the reciprocal masses array
166     int numInverseMasses_ = -1;
167     //! Maximum size of the reciprocal masses array
168     int numInverseMassesAlloc_ = -1;
169
170     //! Number of temperature coupling groups (zero = no coupling)
171     int numTempScaleValues_ = 0;
172     /*! \brief Array with temperature scaling factors.
173      * This is temporary solution to remap data from t_grp_tcstat into plain array.
174      * Not used in SYCL.
175      * \todo Replace with better solution.
176      */
177     gmx::HostVector<float> h_lambdas_;
178     //! Device-side temperature scaling factors
179     DeviceBuffer<float> d_lambdas_;
180     //! Current size of the array with temperature scaling factors (lambdas)
181     int numLambdas_ = -1;
182     //! Maximum size of the array with temperature scaling factors (lambdas)
183     int numLambdasAlloc_ = -1;
184
185
186     //! Array that maps atom index onto the temperature scaling group to get scaling parameter
187     DeviceBuffer<unsigned short> d_tempScaleGroups_;
188     //! Current size of the temperature coupling groups array
189     int numTempScaleGroups_ = -1;
190     //! Maximum size of the temperature coupling groups array
191     int numTempScaleGroupsAlloc_ = -1;
192
193     //! Vector with diagonal elements of the Parrinello-Rahman pressure coupling velocity rescale factors
194     float3 prVelocityScalingMatrixDiagonal_;
195 };
196
197 } // namespace gmx
198
199 #endif