Extend GPU traits class
[alexxy/gromacs.git] / src / gromacs / mdtypes / state_propagator_data_gpu_impl_gpu.cpp
1 /*
2  * This file is part of the GROMACS molecular simulation package.
3  *
4  * Copyright (c) 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  *
37  * \brief Definitions of interfaces for GPU state data propagator object.
38  *
39  * \author Artem Zhmurov <zhmurov@gmail.com>
40  *
41  * \ingroup module_mdtypes
42  */
43 #include "gmxpre.h"
44
45 #include "config.h"
46
47 #if GMX_GPU != GMX_GPU_NONE
48
49 #if GMX_GPU == GMX_GPU_CUDA
50 #include "gromacs/gpu_utils/cudautils.cuh"
51 #endif
52 #include "gromacs/gpu_utils/devicebuffer.h"
53 #if GMX_GPU == GMX_GPU_OPENCL
54 #include "gromacs/gpu_utils/oclutils.h"
55 #endif
56 #include "gromacs/math/vectypes.h"
57 #include "gromacs/mdtypes/state_propagator_data_gpu.h"
58 #include "gromacs/utility/classhelpers.h"
59
60 #include "state_propagator_data_gpu_impl.h"
61
62 namespace gmx
63 {
64
65 StatePropagatorDataGpu::Impl::Impl(gmx_unused const void *commandStream,
66                                    gmx_unused const void *deviceContext,
67                                    GpuApiCallBehavior     transferKind,
68                                    int                    paddingSize) :
69     transferKind_(transferKind),
70     paddingSize_(paddingSize)
71 {
72
73     GMX_RELEASE_ASSERT(getenv("GMX_USE_GPU_BUFFER_OPS") == nullptr, "GPU buffer ops are not supported in this build.");
74
75     // Set the stream-context pair for the OpenCL builds,
76     // use the nullptr stream for CUDA builds
77 #if GMX_GPU == GMX_GPU_OPENCL
78     if (commandStream != nullptr)
79     {
80         commandStream_ = *static_cast<const CommandStream*>(commandStream);
81     }
82     if (deviceContext != nullptr)
83     {
84         deviceContext_ = *static_cast<const DeviceContext*>(deviceContext);
85     }
86 #endif
87
88 }
89
90 StatePropagatorDataGpu::Impl::~Impl()
91 {
92 }
93
94 void StatePropagatorDataGpu::Impl::reinit(int numAtomsLocal, int numAtomsAll)
95 {
96 #if GMX_GPU == GMX_GPU_OPENCL
97     GMX_ASSERT(deviceContext_ != nullptr, "GPU context should be set in OpenCL builds.");
98 #endif
99     numAtomsLocal_ = numAtomsLocal;
100     numAtomsAll_   = numAtomsAll;
101
102     int numAtomsPadded;
103     if (paddingSize_ > 0)
104     {
105         numAtomsPadded = ((numAtomsAll_ + paddingSize_ - 1 ) / paddingSize_ )*paddingSize_;
106     }
107     else
108     {
109         numAtomsPadded = numAtomsAll_;
110     }
111
112     reallocateDeviceBuffer(&d_x_, DIM*numAtomsPadded, &d_xSize_, &d_xCapacity_, deviceContext_);
113
114     const size_t paddingAllocationSize = numAtomsPadded - numAtomsAll_;
115     if (paddingAllocationSize > 0)
116     {
117         clearDeviceBufferAsync(&d_x_, DIM*numAtomsAll_, DIM*paddingAllocationSize, commandStream_);
118     }
119
120     reallocateDeviceBuffer(&d_v_, DIM*numAtomsAll_, &d_vSize_, &d_vCapacity_, deviceContext_);
121     reallocateDeviceBuffer(&d_f_, DIM*numAtomsAll_, &d_fSize_, &d_fCapacity_, deviceContext_);
122
123 }
124
125 std::tuple<int, int> StatePropagatorDataGpu::Impl::getAtomRangesFromAtomLocality(AtomLocality  atomLocality)
126 {
127     int atomsStartAt   = 0;
128     int numAtomsToCopy = 0;
129     switch (atomLocality)
130     {
131         case AtomLocality::All:
132             atomsStartAt    = 0;
133             numAtomsToCopy  = numAtomsAll_;
134             break;
135         case AtomLocality::Local:
136             atomsStartAt    = 0;
137             numAtomsToCopy  = numAtomsLocal_;
138             break;
139         case AtomLocality::NonLocal:
140             atomsStartAt    = numAtomsLocal_;
141             numAtomsToCopy  = numAtomsAll_ - numAtomsLocal_;
142             break;
143         default:
144             GMX_RELEASE_ASSERT(false, "Wrong range of atoms requested in GPU state data manager. Should be All, Local or NonLocal.");
145     }
146     GMX_ASSERT(atomsStartAt   >= 0, "The first elemtnt to copy has negative index. Probably, the GPU propagator state was not initialized.");
147     GMX_ASSERT(numAtomsToCopy >= 0, "Number of atoms to copy is negative. Probably, the GPU propagator state was not initialized.");
148     return std::make_tuple(atomsStartAt, numAtomsToCopy);
149 }
150
151 void StatePropagatorDataGpu::Impl::copyToDevice(DeviceBuffer<float>                   d_data,
152                                                 const gmx::ArrayRef<const gmx::RVec>  h_data,
153                                                 int                                   dataSize,
154                                                 AtomLocality                          atomLocality)
155 {
156
157 #if GMX_GPU == GMX_GPU_OPENCL
158     GMX_ASSERT(deviceContext_ != nullptr, "GPU context should be set in OpenCL builds.");
159 #endif
160
161     GMX_UNUSED_VALUE(dataSize);
162
163     GMX_ASSERT(dataSize >= 0, "Trying to copy to device buffer before it was allocated.");
164
165     int atomsStartAt, numAtomsToCopy;
166     std::tie(atomsStartAt, numAtomsToCopy) = getAtomRangesFromAtomLocality(atomLocality);
167
168     int elementsStartAt   = atomsStartAt*DIM;
169     int numElementsToCopy = numAtomsToCopy*DIM;
170
171     if (numAtomsToCopy != 0)
172     {
173         GMX_ASSERT(elementsStartAt + numElementsToCopy <= dataSize, "The device allocation is smaller than requested copy range.");
174         GMX_ASSERT(atomsStartAt + numAtomsToCopy <= h_data.ssize(), "The host buffer is smaller than the requested copy range.");
175
176         // TODO: Use the proper stream
177         copyToDeviceBuffer(&d_data, reinterpret_cast<const float *>(&h_data.data()[atomsStartAt]),
178                            elementsStartAt, numElementsToCopy,
179                            commandStream_, transferKind_, nullptr);
180     }
181 }
182
183 void StatePropagatorDataGpu::Impl::copyFromDevice(gmx::ArrayRef<gmx::RVec>  h_data,
184                                                   DeviceBuffer<float>       d_data,
185                                                   int                       dataSize,
186                                                   AtomLocality              atomLocality)
187 {
188
189 #if GMX_GPU == GMX_GPU_OPENCL
190     GMX_ASSERT(deviceContext_ != nullptr, "GPU context should be set in OpenCL builds.");
191 #endif
192
193     GMX_UNUSED_VALUE(dataSize);
194
195     GMX_ASSERT(dataSize >= 0, "Trying to copy from device buffer before it was allocated.");
196
197     int atomsStartAt, numAtomsToCopy;
198     std::tie(atomsStartAt, numAtomsToCopy) = getAtomRangesFromAtomLocality(atomLocality);
199
200     int elementsStartAt   = atomsStartAt*DIM;
201     int numElementsToCopy = numAtomsToCopy*DIM;
202
203     if (numAtomsToCopy != 0)
204     {
205         GMX_ASSERT(elementsStartAt + numElementsToCopy <= dataSize, "The device allocation is smaller than requested copy range.");
206         GMX_ASSERT(atomsStartAt + numAtomsToCopy <= h_data.ssize(), "The host buffer is smaller than the requested copy range.");
207
208         // TODO: Use the proper stream
209         copyFromDeviceBuffer(reinterpret_cast<float*>(&h_data.data()[atomsStartAt]), &d_data,
210                              elementsStartAt, numElementsToCopy,
211                              commandStream_, transferKind_, nullptr);
212
213     }
214 }
215
216 DeviceBuffer<float> StatePropagatorDataGpu::Impl::getCoordinates()
217 {
218     return d_x_;
219 }
220
221 void StatePropagatorDataGpu::Impl::copyCoordinatesToGpu(const gmx::ArrayRef<const gmx::RVec>  h_x,
222                                                         AtomLocality                          atomLocality)
223 {
224     copyToDevice(d_x_, h_x, d_xSize_, atomLocality);
225 }
226
227 void StatePropagatorDataGpu::Impl::copyCoordinatesFromGpu(gmx::ArrayRef<gmx::RVec>  h_x,
228                                                           AtomLocality              atomLocality)
229 {
230     copyFromDevice(h_x, d_x_, d_xSize_, atomLocality);
231 }
232
233
234 DeviceBuffer<float> StatePropagatorDataGpu::Impl::getVelocities()
235 {
236     return d_v_;
237 }
238
239 void StatePropagatorDataGpu::Impl::copyVelocitiesToGpu(const gmx::ArrayRef<const gmx::RVec>  h_v,
240                                                        AtomLocality                          atomLocality)
241 {
242     copyToDevice(d_v_, h_v, d_vSize_, atomLocality);
243 }
244
245 void StatePropagatorDataGpu::Impl::copyVelocitiesFromGpu(gmx::ArrayRef<gmx::RVec>  h_v,
246                                                          AtomLocality              atomLocality)
247 {
248     copyFromDevice(h_v, d_v_, d_vSize_, atomLocality);
249 }
250
251
252 DeviceBuffer<float> StatePropagatorDataGpu::Impl::getForces()
253 {
254     return d_f_;
255 }
256
257 void StatePropagatorDataGpu::Impl::copyForcesToGpu(const gmx::ArrayRef<const gmx::RVec>  h_f,
258                                                    AtomLocality                          atomLocality)
259 {
260     copyToDevice(d_f_, h_f, d_fSize_, atomLocality);
261 }
262
263 void StatePropagatorDataGpu::Impl::copyForcesFromGpu(gmx::ArrayRef<gmx::RVec>  h_f,
264                                                      AtomLocality              atomLocality)
265 {
266     copyFromDevice(h_f, d_f_, d_fSize_, atomLocality);
267 }
268
269 void StatePropagatorDataGpu::Impl::synchronizeStream()
270 {
271     gpuStreamSynchronize(commandStream_);
272 }
273
274 int StatePropagatorDataGpu::Impl::numAtomsLocal()
275 {
276     return numAtomsLocal_;
277 }
278
279 int StatePropagatorDataGpu::Impl::numAtomsAll()
280 {
281     return numAtomsAll_;
282 }
283
284
285
286 StatePropagatorDataGpu::StatePropagatorDataGpu(const void        *commandStream,
287                                                const void        *deviceContext,
288                                                GpuApiCallBehavior transferKind,
289                                                int                paddingSize)
290     : impl_(new Impl(commandStream,
291                      deviceContext,
292                      transferKind,
293                      paddingSize))
294 {
295 }
296
297 StatePropagatorDataGpu::StatePropagatorDataGpu(StatePropagatorDataGpu && /* other */) noexcept = default;
298
299 StatePropagatorDataGpu &StatePropagatorDataGpu::operator=(StatePropagatorDataGpu && /* other */) noexcept = default;
300
301 StatePropagatorDataGpu::~StatePropagatorDataGpu() = default;
302
303
304 void StatePropagatorDataGpu::reinit(int numAtomsLocal, int numAtomsAll)
305 {
306     return impl_->reinit(numAtomsLocal, numAtomsAll);
307 }
308
309 std::tuple<int, int> StatePropagatorDataGpu::getAtomRangesFromAtomLocality(AtomLocality  atomLocality)
310 {
311     return impl_->getAtomRangesFromAtomLocality(atomLocality);
312 }
313
314
315 DeviceBuffer<float> StatePropagatorDataGpu::getCoordinates()
316 {
317     return impl_->getCoordinates();
318 }
319
320 void StatePropagatorDataGpu::copyCoordinatesToGpu(const gmx::ArrayRef<const gmx::RVec>  h_x,
321                                                   AtomLocality                          atomLocality)
322 {
323     return impl_->copyCoordinatesToGpu(h_x, atomLocality);
324 }
325
326 void StatePropagatorDataGpu::copyCoordinatesFromGpu(gmx::ArrayRef<RVec>  h_x,
327                                                     AtomLocality         atomLocality)
328 {
329     return impl_->copyCoordinatesFromGpu(h_x, atomLocality);
330 }
331
332
333 DeviceBuffer<float> StatePropagatorDataGpu::getVelocities()
334 {
335     return impl_->getVelocities();
336 }
337
338 void StatePropagatorDataGpu::copyVelocitiesToGpu(const gmx::ArrayRef<const gmx::RVec>  h_v,
339                                                  AtomLocality                          atomLocality)
340 {
341     return impl_->copyVelocitiesToGpu(h_v, atomLocality);
342 }
343
344 void StatePropagatorDataGpu::copyVelocitiesFromGpu(gmx::ArrayRef<RVec>  h_v,
345                                                    AtomLocality         atomLocality)
346 {
347     return impl_->copyVelocitiesFromGpu(h_v, atomLocality);
348 }
349
350
351 DeviceBuffer<float> StatePropagatorDataGpu::getForces()
352 {
353     return impl_->getForces();
354 }
355
356 void StatePropagatorDataGpu::copyForcesToGpu(const gmx::ArrayRef<const gmx::RVec>  h_f,
357                                              AtomLocality                          atomLocality)
358 {
359     return impl_->copyForcesToGpu(h_f, atomLocality);
360 }
361
362 void StatePropagatorDataGpu::copyForcesFromGpu(gmx::ArrayRef<RVec>  h_f,
363                                                AtomLocality         atomLocality)
364 {
365     return impl_->copyForcesFromGpu(h_f, atomLocality);
366 }
367
368 void StatePropagatorDataGpu::synchronizeStream()
369 {
370     return impl_->synchronizeStream();
371 }
372
373 int StatePropagatorDataGpu::numAtomsLocal()
374 {
375     return impl_->numAtomsLocal();
376 }
377
378 int StatePropagatorDataGpu::numAtomsAll()
379 {
380     return impl_->numAtomsAll();
381 }
382
383 }      // namespace gmx
384
385 #endif // GMX_GPU == GMX_GPU_NONE