Make SETTLE setup code platform-agnostic
[alexxy/gromacs.git] / src / gromacs / mdlib / settle_gpu.cpp
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 /*! \internal \file
36  *
37  * \brief Implements SETTLE using GPU
38  *
39  * This file contains implementation for the data management of GPU version of SETTLE constraints algorithm.
40  *
41  * \author Artem Zhmurov <zhmurov@gmail.com>
42  *
43  * \ingroup module_mdlib
44  */
45 #include "gmxpre.h"
46
47 #include "settle_gpu.h"
48
49 #include <assert.h>
50 #include <stdio.h>
51
52 #include <cmath>
53
54 #include <algorithm>
55
56 #include "gromacs/gpu_utils/devicebuffer.h"
57 #include "gromacs/gpu_utils/gputraits.h"
58 #include "gromacs/math/functions.h"
59 #include "gromacs/math/vec.h"
60 #include "gromacs/mdlib/settle_gpu_internal.h"
61 #include "gromacs/pbcutil/pbc.h"
62
63 namespace gmx
64 {
65
66 void SettleGpu::apply(const DeviceBuffer<Float3> d_x,
67                       DeviceBuffer<Float3>       d_xp,
68                       const bool                 updateVelocities,
69                       DeviceBuffer<Float3>       d_v,
70                       const real                 invdt,
71                       const bool                 computeVirial,
72                       tensor                     virialScaled,
73                       const PbcAiuc              pbcAiuc)
74 {
75
76     // Early exit if no settles
77     if (numSettles_ == 0)
78     {
79         return;
80     }
81
82     if (computeVirial)
83     {
84         // Fill with zeros so the values can be reduced to it
85         // Only 6 values are needed because virial is symmetrical
86         clearDeviceBufferAsync(&d_virialScaled_, 0, 6, deviceStream_);
87     }
88
89     launchSettleGpuKernel(numSettles_,
90                           d_atomIds_,
91                           settleParameters_,
92                           d_x,
93                           d_xp,
94                           updateVelocities,
95                           d_v,
96                           invdt,
97                           computeVirial,
98                           d_virialScaled_,
99                           pbcAiuc,
100                           deviceStream_);
101
102
103     if (computeVirial)
104     {
105         copyFromDeviceBuffer(
106                 h_virialScaled_.data(), &d_virialScaled_, 0, 6, deviceStream_, GpuApiCallBehavior::Sync, nullptr);
107
108         // Mapping [XX, XY, XZ, YY, YZ, ZZ] internal format to a tensor object
109         virialScaled[XX][XX] += h_virialScaled_[0];
110         virialScaled[XX][YY] += h_virialScaled_[1];
111         virialScaled[XX][ZZ] += h_virialScaled_[2];
112
113         virialScaled[YY][XX] += h_virialScaled_[1];
114         virialScaled[YY][YY] += h_virialScaled_[3];
115         virialScaled[YY][ZZ] += h_virialScaled_[4];
116
117         virialScaled[ZZ][XX] += h_virialScaled_[2];
118         virialScaled[ZZ][YY] += h_virialScaled_[4];
119         virialScaled[ZZ][ZZ] += h_virialScaled_[5];
120     }
121
122     return;
123 }
124
125 SettleGpu::SettleGpu(const gmx_mtop_t& mtop, const DeviceContext& deviceContext, const DeviceStream& deviceStream) :
126     deviceContext_(deviceContext),
127     deviceStream_(deviceStream)
128 {
129     static_assert(sizeof(real) == sizeof(float),
130                   "Real numbers should be in single precision in GPU code.");
131
132     // This is to prevent the assertion failure for the systems without water
133     int totalSettles = 0;
134     for (unsigned mt = 0; mt < mtop.moltype.size(); mt++)
135     {
136         const int        nral1           = 1 + NRAL(F_SETTLE);
137         InteractionList  interactionList = mtop.moltype[mt].ilist[F_SETTLE];
138         std::vector<int> iatoms          = interactionList.iatoms;
139         totalSettles += iatoms.size() / nral1;
140     }
141     if (totalSettles == 0)
142     {
143         return;
144     }
145
146     // TODO This should be lifted to a separate subroutine that gets the values of Oxygen and
147     // Hydrogen masses, checks if they are consistent across the topology and if there is no more
148     // than two values for each mass if the free energy perturbation is enabled. In later case,
149     // masses may need to be updated on a regular basis (i.e. in set(...) method).
150     // TODO Do the checks for FEP
151     real mO = -1.0;
152     real mH = -1.0;
153
154     for (unsigned mt = 0; mt < mtop.moltype.size(); mt++)
155     {
156         const int        nral1           = 1 + NRAL(F_SETTLE);
157         InteractionList  interactionList = mtop.moltype[mt].ilist[F_SETTLE];
158         std::vector<int> iatoms          = interactionList.iatoms;
159         for (unsigned i = 0; i < iatoms.size() / nral1; i++)
160         {
161             WaterMolecule settler;
162             settler.ow1 = iatoms[i * nral1 + 1]; // Oxygen index
163             settler.hw2 = iatoms[i * nral1 + 2]; // First hydrogen index
164             settler.hw3 = iatoms[i * nral1 + 3]; // Second hydrogen index
165             t_atom ow1  = mtop.moltype[mt].atoms.atom[settler.ow1];
166             t_atom hw2  = mtop.moltype[mt].atoms.atom[settler.hw2];
167             t_atom hw3  = mtop.moltype[mt].atoms.atom[settler.hw3];
168
169             if (mO < 0)
170             {
171                 mO = ow1.m;
172             }
173             if (mH < 0)
174             {
175                 mH = hw2.m;
176             }
177             GMX_RELEASE_ASSERT(mO == ow1.m,
178                                "Topology has different values for oxygen mass. Should be identical "
179                                "in order to use SETTLE.");
180             GMX_RELEASE_ASSERT(hw2.m == hw3.m && hw2.m == mH,
181                                "Topology has different values for hydrogen mass. Should be "
182                                "identical in order to use SETTLE.");
183         }
184     }
185
186     GMX_RELEASE_ASSERT(mO > 0, "Could not find oxygen mass in the topology. Needed in SETTLE.");
187     GMX_RELEASE_ASSERT(mH > 0, "Could not find hydrogen mass in the topology. Needed in SETTLE.");
188
189     // TODO Very similar to SETTLE initialization on CPU. Should be lifted to a separate method
190     // (one that gets dOH and dHH values and checks them for consistency)
191     int settle_type = -1;
192     for (unsigned mt = 0; mt < mtop.moltype.size(); mt++)
193     {
194         const int       nral1           = 1 + NRAL(F_SETTLE);
195         InteractionList interactionList = mtop.moltype[mt].ilist[F_SETTLE];
196         for (int i = 0; i < interactionList.size(); i += nral1)
197         {
198             if (settle_type == -1)
199             {
200                 settle_type = interactionList.iatoms[i];
201             }
202             else if (interactionList.iatoms[i] != settle_type)
203             {
204                 gmx_fatal(FARGS,
205                           "The [molecules] section of your topology specifies more than one block "
206                           "of\n"
207                           "a [moleculetype] with a [settles] block. Only one such is allowed.\n"
208                           "If you are trying to partition your solvent into different *groups*\n"
209                           "(e.g. for freezing, T-coupling, etc.), you are using the wrong "
210                           "approach. Index\n"
211                           "files specify groups. Otherwise, you may wish to change the least-used\n"
212                           "block of molecules with SETTLE constraints into 3 normal constraints.");
213             }
214         }
215     }
216
217     GMX_RELEASE_ASSERT(settle_type >= 0, "settle_init called without settles");
218
219     real dOH = mtop.ffparams.iparams[settle_type].settle.doh;
220     real dHH = mtop.ffparams.iparams[settle_type].settle.dhh;
221
222     settleParameters_ = settleParameters(mO, mH, 1.0 / mO, 1.0 / mH, dOH, dHH);
223
224     allocateDeviceBuffer(&d_virialScaled_, 6, deviceContext_);
225     h_virialScaled_.resize(6);
226 }
227
228 SettleGpu::~SettleGpu()
229 {
230     // Early exit if there is no settles
231     if (numSettles_ == 0)
232     {
233         return;
234     }
235     freeDeviceBuffer(&d_virialScaled_);
236     if (numAtomIdsAlloc_ > 0)
237     {
238         freeDeviceBuffer(&d_atomIds_);
239     }
240 }
241
242 void SettleGpu::set(const InteractionDefinitions& idef)
243 {
244     const int              nral1     = 1 + NRAL(F_SETTLE);
245     const InteractionList& il_settle = idef.il[F_SETTLE];
246     ArrayRef<const int>    iatoms    = il_settle.iatoms;
247     numSettles_                      = il_settle.size() / nral1;
248
249     reallocateDeviceBuffer(&d_atomIds_, numSettles_, &numAtomIds_, &numAtomIdsAlloc_, deviceContext_);
250     h_atomIds_.resize(numSettles_);
251     for (int i = 0; i < numSettles_; i++)
252     {
253         WaterMolecule settler;
254         settler.ow1   = iatoms[i * nral1 + 1]; // Oxygen index
255         settler.hw2   = iatoms[i * nral1 + 2]; // First hydrogen index
256         settler.hw3   = iatoms[i * nral1 + 3]; // Second hydrogen index
257         h_atomIds_[i] = settler;
258     }
259     copyToDeviceBuffer(
260             &d_atomIds_, h_atomIds_.data(), 0, numSettles_, deviceStream_, GpuApiCallBehavior::Sync, nullptr);
261 }
262
263 } // namespace gmx