Unify NB atoms and staging data structures in OpenCL, CUDA and SYCL
[alexxy/gromacs.git] / src / gromacs / nbnxm / opencl / nbnxm_ocl.cpp
1 /*
2  * This file is part of the GROMACS molecular simulation package.
3  *
4  * Copyright (c) 2012,2013,2014,2015,2016 by the GROMACS development team.
5  * Copyright (c) 2017,2018,2019,2020,2021, by the GROMACS development team, led by
6  * Mark Abraham, David van der Spoel, Berk Hess, and Erik Lindahl,
7  * and including many others, as listed in the AUTHORS file in the
8  * top-level source directory and at http://www.gromacs.org.
9  *
10  * GROMACS is free software; you can redistribute it and/or
11  * modify it under the terms of the GNU Lesser General Public License
12  * as published by the Free Software Foundation; either version 2.1
13  * of the License, or (at your option) any later version.
14  *
15  * GROMACS is distributed in the hope that it will be useful,
16  * but WITHOUT ANY WARRANTY; without even the implied warranty of
17  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
18  * Lesser General Public License for more details.
19  *
20  * You should have received a copy of the GNU Lesser General Public
21  * License along with GROMACS; if not, see
22  * http://www.gnu.org/licenses, or write to the Free Software Foundation,
23  * Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA.
24  *
25  * If you want to redistribute modifications to GROMACS, please
26  * consider that scientific software is very special. Version
27  * control is crucial - bugs must be traceable. We will be happy to
28  * consider code for inclusion in the official distribution, but
29  * derived work must not be called official GROMACS. Details are found
30  * in the README & COPYING files - if they are missing, get the
31  * official version at http://www.gromacs.org.
32  *
33  * To help us fund GROMACS development, we humbly ask that you cite
34  * the research papers on the package. Check out http://www.gromacs.org.
35  */
36 /*! \internal \file
37  *  \brief Define OpenCL implementation of nbnxm_gpu.h
38  *
39  *  \author Anca Hamuraru <anca@streamcomputing.eu>
40  *  \author Teemu Virolainen <teemu@streamcomputing.eu>
41  *  \author Dimitrios Karkoulis <dimitris.karkoulis@gmail.com>
42  *  \author Szilárd Páll <pall.szilard@gmail.com>
43  *  \ingroup module_nbnxm
44  *
45  *  TODO (psz):
46  *  - Add a static const cl_uint c_pruneKernelWorkDim / c_nbnxnKernelWorkDim = 3;
47  *  - Rework the copying of OCL data structures done before every invocation of both
48  *    nb and prune kernels (using fillin_ocl_structures); also consider at the same
49  *    time calling clSetKernelArg only on the updated parameters (if tracking changed
50  *    parameters is feasible);
51  *  - Consider using the event_wait_list argument to clEnqueueNDRangeKernel to mark
52  *    dependencies on the kernel launched: e.g. the non-local nb kernel's dependency
53  *    on the misc_ops_and_local_H2D_done event could be better expressed this way.
54  *
55  *  - Consider extracting common sections of the OpenCL and CUDA nbnxn logic, e.g:
56  *    - in nbnxn_gpu_launch_kernel_pruneonly() the pre- and post-kernel launch logic
57  *      is identical in the two implementations, so a 3-way split might allow sharing
58  *      code;
59  *    -
60  *
61  */
62 #include "gmxpre.h"
63
64 #include <assert.h>
65 #include <stdlib.h>
66
67 #if defined(_MSVC)
68 #    include <limits>
69 #endif
70
71 #include "gromacs/gpu_utils/device_context.h"
72 #include "gromacs/gpu_utils/gputraits_ocl.h"
73 #include "gromacs/gpu_utils/oclutils.h"
74 #include "gromacs/hardware/device_information.h"
75 #include "gromacs/hardware/hw_info.h"
76 #include "gromacs/mdtypes/simulation_workload.h"
77 #include "gromacs/nbnxm/atomdata.h"
78 #include "gromacs/nbnxm/gpu_common.h"
79 #include "gromacs/nbnxm/gpu_common_utils.h"
80 #include "gromacs/nbnxm/gpu_data_mgmt.h"
81 #include "gromacs/nbnxm/nbnxm.h"
82 #include "gromacs/nbnxm/nbnxm_gpu.h"
83 #include "gromacs/nbnxm/pairlist.h"
84 #include "gromacs/pbcutil/ishift.h"
85 #include "gromacs/timing/gpu_timing.h"
86 #include "gromacs/utility/cstringutil.h"
87 #include "gromacs/utility/fatalerror.h"
88 #include "gromacs/utility/gmxassert.h"
89
90 #include "nbnxm_ocl_types.h"
91
92 namespace Nbnxm
93 {
94
95 /*! \brief Convenience constants */
96 //@{
97 static constexpr int c_clSize = c_nbnxnGpuClusterSize;
98 //@}
99
100
101 /*! \brief Validates the input global work size parameter.
102  */
103 static inline void validate_global_work_size(const KernelLaunchConfig& config,
104                                              int                       work_dim,
105                                              const DeviceInformation*  dinfo)
106 {
107     cl_uint device_size_t_size_bits;
108     cl_uint host_size_t_size_bits;
109
110     GMX_ASSERT(dinfo, "Need a valid device info object");
111
112     size_t global_work_size[3];
113     GMX_ASSERT(work_dim <= 3, "Not supporting hyper-grids just yet");
114     for (int i = 0; i < work_dim; i++)
115     {
116         global_work_size[i] = config.blockSize[i] * config.gridSize[i];
117     }
118
119     /* Each component of a global_work_size must not exceed the range given by the
120        sizeof(device size_t) for the device on which the kernel execution will
121        be enqueued. See:
122        https://www.khronos.org/registry/cl/sdk/1.0/docs/man/xhtml/clEnqueueNDRangeKernel.html
123      */
124     device_size_t_size_bits = dinfo->adress_bits;
125     host_size_t_size_bits   = static_cast<cl_uint>(sizeof(size_t) * 8);
126
127     /* If sizeof(host size_t) <= sizeof(device size_t)
128             => global_work_size components will always be valid
129        else
130             => get device limit for global work size and
131             compare it against each component of global_work_size.
132      */
133     if (host_size_t_size_bits > device_size_t_size_bits)
134     {
135         size_t device_limit;
136
137         device_limit = (1ULL << device_size_t_size_bits) - 1;
138
139         for (int i = 0; i < work_dim; i++)
140         {
141             if (global_work_size[i] > device_limit)
142             {
143                 gmx_fatal(
144                         FARGS,
145                         "Watch out, the input system is too large to simulate!\n"
146                         "The number of nonbonded work units (=number of super-clusters) exceeds the"
147                         "device capabilities. Global work size limit exceeded (%zu > %zu)!",
148                         global_work_size[i],
149                         device_limit);
150             }
151         }
152     }
153 }
154
155 /* Constant arrays listing non-bonded kernel function names. The arrays are
156  * organized in 2-dim arrays by: electrostatics and VDW type.
157  *
158  *  Note that the row- and column-order of function pointers has to match the
159  *  order of corresponding enumerated electrostatics and vdw types, resp.,
160  *  defined in nbnxm_ocl_types.h.
161  */
162
163 /*! \brief Force-only kernel function names. */
164 static const char* nb_kfunc_noener_noprune_ptr[c_numElecTypes][c_numVdwTypes] = {
165     { "nbnxn_kernel_ElecCut_VdwLJ_F_opencl",
166       "nbnxn_kernel_ElecCut_VdwLJCombGeom_F_opencl",
167       "nbnxn_kernel_ElecCut_VdwLJCombLB_F_opencl",
168       "nbnxn_kernel_ElecCut_VdwLJFsw_F_opencl",
169       "nbnxn_kernel_ElecCut_VdwLJPsw_F_opencl",
170       "nbnxn_kernel_ElecCut_VdwLJEwCombGeom_F_opencl",
171       "nbnxn_kernel_ElecCut_VdwLJEwCombLB_F_opencl" },
172     { "nbnxn_kernel_ElecRF_VdwLJ_F_opencl",
173       "nbnxn_kernel_ElecRF_VdwLJCombGeom_F_opencl",
174       "nbnxn_kernel_ElecRF_VdwLJCombLB_F_opencl",
175       "nbnxn_kernel_ElecRF_VdwLJFsw_F_opencl",
176       "nbnxn_kernel_ElecRF_VdwLJPsw_F_opencl",
177       "nbnxn_kernel_ElecRF_VdwLJEwCombGeom_F_opencl",
178       "nbnxn_kernel_ElecRF_VdwLJEwCombLB_F_opencl" },
179     { "nbnxn_kernel_ElecEwQSTab_VdwLJ_F_opencl",
180       "nbnxn_kernel_ElecEwQSTab_VdwLJCombGeom_F_opencl",
181       "nbnxn_kernel_ElecEwQSTab_VdwLJCombLB_F_opencl",
182       "nbnxn_kernel_ElecEwQSTab_VdwLJFsw_F_opencl",
183       "nbnxn_kernel_ElecEwQSTab_VdwLJPsw_F_opencl",
184       "nbnxn_kernel_ElecEwQSTab_VdwLJEwCombGeom_F_opencl",
185       "nbnxn_kernel_ElecEwQSTab_VdwLJEwCombLB_F_opencl" },
186     { "nbnxn_kernel_ElecEwQSTabTwinCut_VdwLJ_F_opencl",
187       "nbnxn_kernel_ElecEwQSTabTwinCut_VdwLJCombGeom_F_opencl",
188       "nbnxn_kernel_ElecEwQSTabTwinCut_VdwLJCombLB_F_opencl",
189       "nbnxn_kernel_ElecEwQSTabTwinCut_VdwLJFsw_F_opencl",
190       "nbnxn_kernel_ElecEwQSTabTwinCut_VdwLJPsw_F_opencl",
191       "nbnxn_kernel_ElecEwQSTabTwinCut_VdwLJEwCombGeom_F_opencl",
192       "nbnxn_kernel_ElecEwQSTabTwinCut_VdwLJEwCombLB_F_opencl" },
193     { "nbnxn_kernel_ElecEw_VdwLJ_F_opencl",
194       "nbnxn_kernel_ElecEw_VdwLJCombGeom_F_opencl",
195       "nbnxn_kernel_ElecEw_VdwLJCombLB_F_opencl",
196       "nbnxn_kernel_ElecEw_VdwLJFsw_F_opencl",
197       "nbnxn_kernel_ElecEw_VdwLJPsw_F_opencl",
198       "nbnxn_kernel_ElecEw_VdwLJEwCombGeom_F_opencl",
199       "nbnxn_kernel_ElecEw_VdwLJEwCombLB_F_opencl" },
200     { "nbnxn_kernel_ElecEwTwinCut_VdwLJ_F_opencl",
201       "nbnxn_kernel_ElecEwTwinCut_VdwLJCombGeom_F_opencl",
202       "nbnxn_kernel_ElecEwTwinCut_VdwLJCombLB_F_opencl",
203       "nbnxn_kernel_ElecEwTwinCut_VdwLJFsw_F_opencl",
204       "nbnxn_kernel_ElecEwTwinCut_VdwLJPsw_F_opencl",
205       "nbnxn_kernel_ElecEwTwinCut_VdwLJEwCombGeom_F_opencl",
206       "nbnxn_kernel_ElecEwTwinCut_VdwLJEwCombLB_F_opencl" }
207 };
208
209 /*! \brief Force + energy kernel function pointers. */
210 static const char* nb_kfunc_ener_noprune_ptr[c_numElecTypes][c_numVdwTypes] = {
211     { "nbnxn_kernel_ElecCut_VdwLJ_VF_opencl",
212       "nbnxn_kernel_ElecCut_VdwLJCombGeom_VF_opencl",
213       "nbnxn_kernel_ElecCut_VdwLJCombLB_VF_opencl",
214       "nbnxn_kernel_ElecCut_VdwLJFsw_VF_opencl",
215       "nbnxn_kernel_ElecCut_VdwLJPsw_VF_opencl",
216       "nbnxn_kernel_ElecCut_VdwLJEwCombGeom_VF_opencl",
217       "nbnxn_kernel_ElecCut_VdwLJEwCombLB_VF_opencl" },
218     { "nbnxn_kernel_ElecRF_VdwLJ_VF_opencl",
219       "nbnxn_kernel_ElecRF_VdwLJCombGeom_VF_opencl",
220       "nbnxn_kernel_ElecRF_VdwLJCombLB_VF_opencl",
221       "nbnxn_kernel_ElecRF_VdwLJFsw_VF_opencl",
222       "nbnxn_kernel_ElecRF_VdwLJPsw_VF_opencl",
223       "nbnxn_kernel_ElecRF_VdwLJEwCombGeom_VF_opencl",
224       "nbnxn_kernel_ElecRF_VdwLJEwCombLB_VF_opencl" },
225     { "nbnxn_kernel_ElecEwQSTab_VdwLJ_VF_opencl",
226       "nbnxn_kernel_ElecEwQSTab_VdwLJCombGeom_VF_opencl",
227       "nbnxn_kernel_ElecEwQSTab_VdwLJCombLB_VF_opencl",
228       "nbnxn_kernel_ElecEwQSTab_VdwLJFsw_VF_opencl",
229       "nbnxn_kernel_ElecEwQSTab_VdwLJPsw_VF_opencl",
230       "nbnxn_kernel_ElecEwQSTab_VdwLJEwCombGeom_VF_opencl",
231       "nbnxn_kernel_ElecEwQSTab_VdwLJEwCombLB_VF_opencl" },
232     { "nbnxn_kernel_ElecEwQSTabTwinCut_VdwLJ_VF_opencl",
233       "nbnxn_kernel_ElecEwQSTabTwinCut_VdwLJCombGeom_VF_opencl",
234       "nbnxn_kernel_ElecEwQSTabTwinCut_VdwLJCombLB_VF_opencl",
235       "nbnxn_kernel_ElecEwQSTabTwinCut_VdwLJFsw_VF_opencl",
236       "nbnxn_kernel_ElecEwQSTabTwinCut_VdwLJPsw_VF_opencl",
237       "nbnxn_kernel_ElecEwQSTabTwinCut_VdwLJEwCombGeom_VF_opencl",
238       "nbnxn_kernel_ElecEwQSTabTwinCut_VdwLJEwCombLB_VF_opencl" },
239     { "nbnxn_kernel_ElecEw_VdwLJ_VF_opencl",
240       "nbnxn_kernel_ElecEw_VdwLJCombGeom_VF_opencl",
241       "nbnxn_kernel_ElecEw_VdwLJCombLB_VF_opencl",
242       "nbnxn_kernel_ElecEw_VdwLJFsw_VF_opencl",
243       "nbnxn_kernel_ElecEw_VdwLJPsw_VF_opencl",
244       "nbnxn_kernel_ElecEw_VdwLJEwCombGeom_VF_opencl",
245       "nbnxn_kernel_ElecEw_VdwLJEwCombLB_VF_opencl" },
246     { "nbnxn_kernel_ElecEwTwinCut_VdwLJ_VF_opencl",
247       "nbnxn_kernel_ElecEwTwinCut_VdwLJCombGeom_VF_opencl",
248       "nbnxn_kernel_ElecEwTwinCut_VdwLJCombLB_VF_opencl",
249       "nbnxn_kernel_ElecEwTwinCut_VdwLJFsw_VF_opencl",
250       "nbnxn_kernel_ElecEwTwinCut_VdwLJPsw_VF_opencl",
251       "nbnxn_kernel_ElecEwTwinCut_VdwLJEwCombGeom_VF_opencl",
252       "nbnxn_kernel_ElecEwTwinCut_VdwLJEwCombLB_VF_opencl" }
253 };
254
255 /*! \brief Force + pruning kernel function pointers. */
256 static const char* nb_kfunc_noener_prune_ptr[c_numElecTypes][c_numVdwTypes] = {
257     { "nbnxn_kernel_ElecCut_VdwLJ_F_prune_opencl",
258       "nbnxn_kernel_ElecCut_VdwLJCombGeom_F_prune_opencl",
259       "nbnxn_kernel_ElecCut_VdwLJCombLB_F_prune_opencl",
260       "nbnxn_kernel_ElecCut_VdwLJFsw_F_prune_opencl",
261       "nbnxn_kernel_ElecCut_VdwLJPsw_F_prune_opencl",
262       "nbnxn_kernel_ElecCut_VdwLJEwCombGeom_F_prune_opencl",
263       "nbnxn_kernel_ElecCut_VdwLJEwCombLB_F_prune_opencl" },
264     { "nbnxn_kernel_ElecRF_VdwLJ_F_prune_opencl",
265       "nbnxn_kernel_ElecRF_VdwLJCombGeom_F_prune_opencl",
266       "nbnxn_kernel_ElecRF_VdwLJCombLB_F_prune_opencl",
267       "nbnxn_kernel_ElecRF_VdwLJFsw_F_prune_opencl",
268       "nbnxn_kernel_ElecRF_VdwLJPsw_F_prune_opencl",
269       "nbnxn_kernel_ElecRF_VdwLJEwCombGeom_F_prune_opencl",
270       "nbnxn_kernel_ElecRF_VdwLJEwCombLB_F_prune_opencl" },
271     { "nbnxn_kernel_ElecEwQSTab_VdwLJ_F_prune_opencl",
272       "nbnxn_kernel_ElecEwQSTab_VdwLJCombGeom_F_prune_opencl",
273       "nbnxn_kernel_ElecEwQSTab_VdwLJCombLB_F_prune_opencl",
274       "nbnxn_kernel_ElecEwQSTab_VdwLJFsw_F_prune_opencl",
275       "nbnxn_kernel_ElecEwQSTab_VdwLJPsw_F_prune_opencl",
276       "nbnxn_kernel_ElecEwQSTab_VdwLJEwCombGeom_F_prune_opencl",
277       "nbnxn_kernel_ElecEwQSTab_VdwLJEwCombLB_F_prune_opencl" },
278     { "nbnxn_kernel_ElecEwQSTabTwinCut_VdwLJ_F_prune_opencl",
279       "nbnxn_kernel_ElecEwQSTabTwinCut_VdwLJCombGeom_F_prune_opencl",
280       "nbnxn_kernel_ElecEwQSTabTwinCut_VdwLJCombLB_F_prune_opencl",
281       "nbnxn_kernel_ElecEwQSTabTwinCut_VdwLJFsw_F_prune_opencl",
282       "nbnxn_kernel_ElecEwQSTabTwinCut_VdwLJPsw_F_prune_opencl",
283       "nbnxn_kernel_ElecEwQSTabTwinCut_VdwLJEwCombGeom_F_prune_opencl",
284       "nbnxn_kernel_ElecEwQSTabTwinCut_VdwLJEwCombLB_F_prune_opencl" },
285     { "nbnxn_kernel_ElecEw_VdwLJ_F_prune_opencl",
286       "nbnxn_kernel_ElecEw_VdwLJCombGeom_F_prune_opencl",
287       "nbnxn_kernel_ElecEw_VdwLJCombLB_F_prune_opencl",
288       "nbnxn_kernel_ElecEw_VdwLJFsw_F_prune_opencl",
289       "nbnxn_kernel_ElecEw_VdwLJPsw_F_prune_opencl",
290       "nbnxn_kernel_ElecEw_VdwLJEwCombGeom_F_prune_opencl",
291       "nbnxn_kernel_ElecEw_VdwLJEwCombLB_F_prune_opencl" },
292     { "nbnxn_kernel_ElecEwTwinCut_VdwLJ_F_prune_opencl",
293       "nbnxn_kernel_ElecEwTwinCut_VdwLJCombGeom_F_prune_opencl",
294       "nbnxn_kernel_ElecEwTwinCut_VdwLJCombLB_F_prune_opencl",
295       "nbnxn_kernel_ElecEwTwinCut_VdwLJFsw_F_prune_opencl",
296       "nbnxn_kernel_ElecEwTwinCut_VdwLJPsw_F_prune_opencl",
297       "nbnxn_kernel_ElecEwTwinCut_VdwLJEwCombGeom_F_prune_opencl",
298       "nbnxn_kernel_ElecEwTwinCut_VdwLJEwCombLB_F_prune_opencl" }
299 };
300
301 /*! \brief Force + energy + pruning kernel function pointers. */
302 static const char* nb_kfunc_ener_prune_ptr[c_numElecTypes][c_numVdwTypes] = {
303     { "nbnxn_kernel_ElecCut_VdwLJ_VF_prune_opencl",
304       "nbnxn_kernel_ElecCut_VdwLJCombGeom_VF_prune_opencl",
305       "nbnxn_kernel_ElecCut_VdwLJCombLB_VF_prune_opencl",
306       "nbnxn_kernel_ElecCut_VdwLJFsw_VF_prune_opencl",
307       "nbnxn_kernel_ElecCut_VdwLJPsw_VF_prune_opencl",
308       "nbnxn_kernel_ElecCut_VdwLJEwCombGeom_VF_prune_opencl",
309       "nbnxn_kernel_ElecCut_VdwLJEwCombLB_VF_prune_opencl" },
310     { "nbnxn_kernel_ElecRF_VdwLJ_VF_prune_opencl",
311       "nbnxn_kernel_ElecRF_VdwLJCombGeom_VF_prune_opencl",
312       "nbnxn_kernel_ElecRF_VdwLJCombLB_VF_prune_opencl",
313       "nbnxn_kernel_ElecRF_VdwLJFsw_VF_prune_opencl",
314       "nbnxn_kernel_ElecRF_VdwLJPsw_VF_prune_opencl",
315       "nbnxn_kernel_ElecRF_VdwLJEwCombGeom_VF_prune_opencl",
316       "nbnxn_kernel_ElecRF_VdwLJEwCombLB_VF_prune_opencl" },
317     { "nbnxn_kernel_ElecEwQSTab_VdwLJ_VF_prune_opencl",
318       "nbnxn_kernel_ElecEwQSTab_VdwLJCombGeom_VF_prune_opencl",
319       "nbnxn_kernel_ElecEwQSTab_VdwLJCombLB_VF_prune_opencl",
320       "nbnxn_kernel_ElecEwQSTab_VdwLJFsw_VF_prune_opencl",
321       "nbnxn_kernel_ElecEwQSTab_VdwLJPsw_VF_prune_opencl",
322       "nbnxn_kernel_ElecEwQSTab_VdwLJEwCombGeom_VF_prune_opencl",
323       "nbnxn_kernel_ElecEwQSTab_VdwLJEwCombLB_VF_prune_opencl" },
324     { "nbnxn_kernel_ElecEwQSTabTwinCut_VdwLJ_VF_prune_opencl",
325       "nbnxn_kernel_ElecEwQSTabTwinCut_VdwLJCombGeom_VF_prune_opencl",
326       "nbnxn_kernel_ElecEwQSTabTwinCut_VdwLJCombLB_VF_prune_opencl",
327       "nbnxn_kernel_ElecEwQSTabTwinCut_VdwLJFsw_VF_prune_opencl",
328       "nbnxn_kernel_ElecEwQSTabTwinCut_VdwLJPsw_VF_prune_opencl",
329       "nbnxn_kernel_ElecEwQSTabTwinCut_VdwLJEwCombGeom_VF_prune_opencl",
330       "nbnxn_kernel_ElecEwQSTabTwinCut_VdwLJEwCombLB_VF_prune_opencl" },
331     { "nbnxn_kernel_ElecEw_VdwLJ_VF_prune_opencl",
332       "nbnxn_kernel_ElecEw_VdwLJCombGeom_VF_prune_opencl",
333       "nbnxn_kernel_ElecEw_VdwLJCombLB_VF_prune_opencl",
334       "nbnxn_kernel_ElecEw_VdwLJFsw_VF_prune_opencl",
335       "nbnxn_kernel_ElecEw_VdwLJPsw_VF_prune_opencl",
336       "nbnxn_kernel_ElecEw_VdwLJEwCombGeom_VF_prune_opencl",
337       "nbnxn_kernel_ElecEw_VdwLJEwCombLB_VF_prune_opencl" },
338     { "nbnxn_kernel_ElecEwTwinCut_VdwLJ_VF_prune_opencl",
339       "nbnxn_kernel_ElecEwTwinCut_VdwLJCombGeom_VF_prune_opencl",
340       "nbnxn_kernel_ElecEwTwinCut_VdwLJCombLB_VF_prune_opencl",
341       "nbnxn_kernel_ElecEwTwinCut_VdwLJFsw_VF_prune_opencl",
342       "nbnxn_kernel_ElecEwTwinCut_VdwLJPsw_VF_prune_opencl",
343       "nbnxn_kernel_ElecEwTwinCut_VdwLJEwCombGeom_VF_prune_opencl",
344       "nbnxn_kernel_ElecEwTwinCut_VdwLJEwCombLB_VF_prune_opencl" }
345 };
346
347 /*! \brief Return a pointer to the prune kernel version to be executed at the current invocation.
348  *
349  * \param[in] kernel_pruneonly  array of prune kernel objects
350  * \param[in] firstPrunePass    true if the first pruning pass is being executed
351  */
352 static inline cl_kernel selectPruneKernel(cl_kernel kernel_pruneonly[], bool firstPrunePass)
353 {
354     cl_kernel* kernelPtr;
355
356     if (firstPrunePass)
357     {
358         kernelPtr = &(kernel_pruneonly[epruneFirst]);
359     }
360     else
361     {
362         kernelPtr = &(kernel_pruneonly[epruneRolling]);
363     }
364     // TODO: consider creating the prune kernel object here to avoid a
365     // clCreateKernel for the rolling prune kernel if this is not needed.
366     return *kernelPtr;
367 }
368
369 /*! \brief Return a pointer to the kernel version to be executed at the current step.
370  *  OpenCL kernel objects are cached in nb. If the requested kernel is not
371  *  found in the cache, it will be created and the cache will be updated.
372  */
373 static inline cl_kernel
374 select_nbnxn_kernel(NbnxmGpu* nb, enum ElecType elecType, enum VdwType vdwType, bool bDoEne, bool bDoPrune)
375 {
376     const char* kernel_name_to_run;
377     cl_kernel*  kernel_ptr;
378     cl_int      cl_error;
379
380     const int elecTypeIdx = static_cast<int>(elecType);
381     const int vdwTypeIdx  = static_cast<int>(vdwType);
382
383     GMX_ASSERT(elecTypeIdx < c_numElecTypes,
384                "The electrostatics type requested is not implemented in the OpenCL kernels.");
385     GMX_ASSERT(vdwTypeIdx < c_numVdwTypes,
386                "The VdW type requested is not implemented in the OpenCL kernels.");
387
388     if (bDoEne)
389     {
390         if (bDoPrune)
391         {
392             kernel_name_to_run = nb_kfunc_ener_prune_ptr[elecTypeIdx][vdwTypeIdx];
393             kernel_ptr         = &(nb->kernel_ener_prune_ptr[elecTypeIdx][vdwTypeIdx]);
394         }
395         else
396         {
397             kernel_name_to_run = nb_kfunc_ener_noprune_ptr[elecTypeIdx][vdwTypeIdx];
398             kernel_ptr         = &(nb->kernel_ener_noprune_ptr[elecTypeIdx][vdwTypeIdx]);
399         }
400     }
401     else
402     {
403         if (bDoPrune)
404         {
405             kernel_name_to_run = nb_kfunc_noener_prune_ptr[elecTypeIdx][vdwTypeIdx];
406             kernel_ptr         = &(nb->kernel_noener_prune_ptr[elecTypeIdx][vdwTypeIdx]);
407         }
408         else
409         {
410             kernel_name_to_run = nb_kfunc_noener_noprune_ptr[elecTypeIdx][vdwTypeIdx];
411             kernel_ptr         = &(nb->kernel_noener_noprune_ptr[elecTypeIdx][vdwTypeIdx]);
412         }
413     }
414
415     if (nullptr == kernel_ptr[0])
416     {
417         *kernel_ptr = clCreateKernel(nb->dev_rundata->program, kernel_name_to_run, &cl_error);
418         GMX_ASSERT(cl_error == CL_SUCCESS,
419                    ("clCreateKernel failed: " + ocl_get_error_string(cl_error)
420                     + " for kernel named " + kernel_name_to_run)
421                            .c_str());
422     }
423
424     return *kernel_ptr;
425 }
426
427 /*! \brief Calculates the amount of shared memory required by the nonbonded kernel in use.
428  */
429 static inline int calc_shmem_required_nonbonded(enum VdwType vdwType, bool bPrefetchLjParam)
430 {
431     int shmem;
432
433     /* size of shmem (force-buffers/xq/atom type preloading) */
434     /* NOTE: with the default kernel on sm3.0 we need shmem only for pre-loading */
435     /* i-atom x+q in shared memory */
436     shmem = c_nbnxnGpuNumClusterPerSupercluster * c_clSize * sizeof(float) * 4; /* xqib */
437     /* cj in shared memory, for both warps separately
438      * TODO: in the "nowarp kernels we load cj only once  so the factor 2 is not needed.
439      */
440     shmem += 2 * c_nbnxnGpuJgroupSize * sizeof(int); /* cjs  */
441     if (bPrefetchLjParam)
442     {
443         if (useLjCombRule(vdwType))
444         {
445             /* i-atom LJ combination parameters in shared memory */
446             shmem += c_nbnxnGpuNumClusterPerSupercluster * c_clSize * 2
447                      * sizeof(float); /* atib abused for ljcp, float2 */
448         }
449         else
450         {
451             /* i-atom types in shared memory */
452             shmem += c_nbnxnGpuNumClusterPerSupercluster * c_clSize * sizeof(int); /* atib */
453         }
454     }
455     /* force reduction buffers in shared memory */
456     shmem += c_clSize * c_clSize * 3 * sizeof(float); /* f_buf */
457     /* Warp vote. In fact it must be * number of warps in block.. */
458     shmem += sizeof(cl_uint) * 2; /* warp_any */
459     return shmem;
460 }
461
462 /*! \brief Initializes data structures that are going to be sent to the OpenCL device.
463  *
464  *  The device can't use the same data structures as the host for two main reasons:
465  *  - OpenCL restrictions (pointers are not accepted inside data structures)
466  *  - some host side fields are not needed for the OpenCL kernels.
467  *
468  *  This function is called before the launch of both nbnxn and prune kernels.
469  */
470 static void fillin_ocl_structures(NBParamGpu* nbp, cl_nbparam_params_t* nbparams_params)
471 {
472     nbparams_params->coulomb_tab_scale = nbp->coulomb_tab_scale;
473     nbparams_params->c_rf              = nbp->c_rf;
474     nbparams_params->dispersion_shift  = nbp->dispersion_shift;
475     nbparams_params->elecType          = nbp->elecType;
476     nbparams_params->epsfac            = nbp->epsfac;
477     nbparams_params->ewaldcoeff_lj     = nbp->ewaldcoeff_lj;
478     nbparams_params->ewald_beta        = nbp->ewald_beta;
479     nbparams_params->rcoulomb_sq       = nbp->rcoulomb_sq;
480     nbparams_params->repulsion_shift   = nbp->repulsion_shift;
481     nbparams_params->rlistOuter_sq     = nbp->rlistOuter_sq;
482     nbparams_params->rvdw_sq           = nbp->rvdw_sq;
483     nbparams_params->rlistInner_sq     = nbp->rlistInner_sq;
484     nbparams_params->rvdw_switch       = nbp->rvdw_switch;
485     nbparams_params->sh_ewald          = nbp->sh_ewald;
486     nbparams_params->sh_lj_ewald       = nbp->sh_lj_ewald;
487     nbparams_params->two_k_rf          = nbp->two_k_rf;
488     nbparams_params->vdwType           = nbp->vdwType;
489     nbparams_params->vdw_switch        = nbp->vdw_switch;
490 }
491
492 void nbnxnInsertNonlocalGpuDependency(NbnxmGpu* nb, const InteractionLocality interactionLocality)
493 {
494     const DeviceStream& deviceStream = *nb->deviceStreams[interactionLocality];
495
496     /* When we get here all misc operations issued in the local stream as well as
497        the local xq H2D are done,
498        so we record that in the local stream and wait for it in the nonlocal one.
499        This wait needs to precede any PP tasks, bonded or nonbonded, that may
500        compute on interactions between local and nonlocal atoms.
501      */
502     if (nb->bUseTwoStreams)
503     {
504         if (interactionLocality == InteractionLocality::Local)
505         {
506             nb->misc_ops_and_local_H2D_done.markEvent(deviceStream);
507
508             /* Based on the v1.2 section 5.13 of the OpenCL spec, a flush is needed
509              * in the local stream in order to be able to sync with the above event
510              * from the non-local stream.
511              */
512             cl_int gmx_used_in_debug cl_error = clFlush(deviceStream.stream());
513             GMX_ASSERT(cl_error == CL_SUCCESS,
514                        ("clFlush failed: " + ocl_get_error_string(cl_error)).c_str());
515         }
516         else
517         {
518             nb->misc_ops_and_local_H2D_done.enqueueWaitEvent(deviceStream);
519         }
520     }
521 }
522
523 /*! \brief Launch asynchronously the xq buffer host to device copy. */
524 void gpu_copy_xq_to_gpu(NbnxmGpu* nb, const nbnxn_atomdata_t* nbatom, const AtomLocality atomLocality)
525 {
526     GMX_ASSERT(nb, "Need a valid nbnxn_gpu object");
527
528     const InteractionLocality iloc = gpuAtomToInteractionLocality(atomLocality);
529
530     /* local/nonlocal offset and length used for xq and f */
531     int adat_begin, adat_len;
532
533     NBAtomData*         adat         = nb->atdat;
534     gpu_plist*          plist        = nb->plist[iloc];
535     cl_timers_t*        t            = nb->timers;
536     const DeviceStream& deviceStream = *nb->deviceStreams[iloc];
537
538     bool bDoTime = nb->bDoTime;
539
540     /* Don't launch the non-local H2D copy if there is no dependent
541        work to do: neither non-local nor other (e.g. bonded) work
542        to do that has as input the nbnxn coordinates.
543        Doing the same for the local kernel is more complicated, since the
544        local part of the force array also depends on the non-local kernel.
545        So to avoid complicating the code and to reduce the risk of bugs,
546        we always call the local local x+q copy (and the rest of the local
547        work in nbnxn_gpu_launch_kernel().
548      */
549     if ((iloc == InteractionLocality::NonLocal) && !haveGpuShortRangeWork(*nb, iloc))
550     {
551         plist->haveFreshList = false;
552
553         // The event is marked for Local interactions unconditionally,
554         // so it has to be released here because of the early return
555         // for NonLocal interactions.
556         nb->misc_ops_and_local_H2D_done.reset();
557
558         return;
559     }
560
561     /* calculate the atom data index range based on locality */
562     if (atomLocality == AtomLocality::Local)
563     {
564         adat_begin = 0;
565         adat_len   = adat->numAtomsLocal;
566     }
567     else
568     {
569         adat_begin = adat->numAtomsLocal;
570         adat_len   = adat->numAtoms - adat->numAtomsLocal;
571     }
572
573     /* beginning of timed HtoD section */
574     if (bDoTime)
575     {
576         t->xf[atomLocality].nb_h2d.openTimingRegion(deviceStream);
577     }
578
579     /* HtoD x, q */
580     static_assert(sizeof(float) == sizeof(*nbatom->x().data()),
581                   "The size of the xyzq buffer element should be equal to the size of float4.");
582     copyToDeviceBuffer(&adat->xq,
583                        reinterpret_cast<const Float4*>(nbatom->x().data()) + adat_begin,
584                        adat_begin,
585                        adat_len,
586                        deviceStream,
587                        GpuApiCallBehavior::Async,
588                        bDoTime ? t->xf[atomLocality].nb_h2d.fetchNextEvent() : nullptr);
589
590     if (bDoTime)
591     {
592         t->xf[atomLocality].nb_h2d.closeTimingRegion(deviceStream);
593     }
594
595     /* When we get here all misc operations issued in the local stream as well as
596        the local xq H2D are done,
597        so we record that in the local stream and wait for it in the nonlocal one.
598        This wait needs to precede any PP tasks, bonded or nonbonded, that may
599        compute on interactions between local and nonlocal atoms.
600      */
601     nbnxnInsertNonlocalGpuDependency(nb, iloc);
602 }
603
604
605 /*! \brief Launch GPU kernel
606
607    As we execute nonbonded workload in separate queues, before launching
608    the kernel we need to make sure that he following operations have completed:
609    - atomdata allocation and related H2D transfers (every nstlist step);
610    - pair list H2D transfer (every nstlist step);
611    - shift vector H2D transfer (every nstlist step);
612    - force (+shift force and energy) output clearing (every step).
613
614    These operations are issued in the local queue at the beginning of the step
615    and therefore always complete before the local kernel launch. The non-local
616    kernel is launched after the local on the same device/context, so this is
617    inherently scheduled after the operations in the local stream (including the
618    above "misc_ops").
619    However, for the sake of having a future-proof implementation, we use the
620    misc_ops_done event to record the point in time when the above  operations
621    are finished and synchronize with this event in the non-local stream.
622  */
623 void gpu_launch_kernel(NbnxmGpu* nb, const gmx::StepWorkload& stepWork, const Nbnxm::InteractionLocality iloc)
624 {
625     NBAtomData*         adat         = nb->atdat;
626     NBParamGpu*         nbp          = nb->nbparam;
627     gpu_plist*          plist        = nb->plist[iloc];
628     cl_timers_t*        t            = nb->timers;
629     const DeviceStream& deviceStream = *nb->deviceStreams[iloc];
630
631     bool bDoTime = nb->bDoTime;
632
633     cl_nbparam_params_t nbparams_params;
634
635     /* Don't launch the non-local kernel if there is no work to do.
636        Doing the same for the local kernel is more complicated, since the
637        local part of the force array also depends on the non-local kernel.
638        So to avoid complicating the code and to reduce the risk of bugs,
639        we always call the local kernel and later (not in
640        this function) the stream wait, local f copyback and the f buffer
641        clearing. All these operations, except for the local interaction kernel,
642        are needed for the non-local interactions. The skip of the local kernel
643        call is taken care of later in this function. */
644     if (canSkipNonbondedWork(*nb, iloc))
645     {
646         plist->haveFreshList = false;
647
648         return;
649     }
650
651     if (nbp->useDynamicPruning && plist->haveFreshList)
652     {
653         /* Prunes for rlistOuter and rlistInner, sets plist->haveFreshList=false
654            (that's the way the timing accounting can distinguish between
655            separate prune kernel and combined force+prune).
656          */
657         Nbnxm::gpu_launch_kernel_pruneonly(nb, iloc, 1);
658     }
659
660     if (plist->nsci == 0)
661     {
662         /* Don't launch an empty local kernel (is not allowed with OpenCL).
663          */
664         return;
665     }
666
667     /* beginning of timed nonbonded calculation section */
668     if (bDoTime)
669     {
670         t->interaction[iloc].nb_k.openTimingRegion(deviceStream);
671     }
672
673     /* kernel launch config */
674
675     KernelLaunchConfig config;
676     config.sharedMemorySize = calc_shmem_required_nonbonded(nbp->vdwType, nb->bPrefetchLjParam);
677     config.blockSize[0]     = c_clSize;
678     config.blockSize[1]     = c_clSize;
679     config.gridSize[0]      = plist->nsci;
680
681     validate_global_work_size(config, 3, &nb->deviceContext_->deviceInfo());
682
683     if (debug)
684     {
685         fprintf(debug,
686                 "Non-bonded GPU launch configuration:\n\tLocal work size: %zux%zux%zu\n\t"
687                 "Global work size : %zux%zu\n\t#Super-clusters/clusters: %d/%d (%d)\n",
688                 config.blockSize[0],
689                 config.blockSize[1],
690                 config.blockSize[2],
691                 config.blockSize[0] * config.gridSize[0],
692                 config.blockSize[1] * config.gridSize[1],
693                 plist->nsci * c_nbnxnGpuNumClusterPerSupercluster,
694                 c_nbnxnGpuNumClusterPerSupercluster,
695                 plist->na_c);
696     }
697
698     fillin_ocl_structures(nbp, &nbparams_params);
699
700     auto*          timingEvent  = bDoTime ? t->interaction[iloc].nb_k.fetchNextEvent() : nullptr;
701     constexpr char kernelName[] = "k_calc_nb";
702     const auto     kernel =
703             select_nbnxn_kernel(nb,
704                                 nbp->elecType,
705                                 nbp->vdwType,
706                                 stepWork.computeEnergy,
707                                 (plist->haveFreshList && !nb->timers->interaction[iloc].didPrune));
708
709
710     // The OpenCL kernel takes int as second to last argument because bool is
711     // not supported as a kernel argument type (sizeof(bool) is implementation defined).
712     const int computeFshift = static_cast<int>(stepWork.computeVirial);
713     if (useLjCombRule(nb->nbparam->vdwType))
714     {
715         const auto kernelArgs = prepareGpuKernelArguments(kernel,
716                                                           config,
717                                                           &nbparams_params,
718                                                           &adat->xq,
719                                                           &adat->f,
720                                                           &adat->eLJ,
721                                                           &adat->eElec,
722                                                           &adat->fShift,
723                                                           &adat->ljComb,
724                                                           &adat->shiftVec,
725                                                           &nbp->nbfp,
726                                                           &nbp->nbfp_comb,
727                                                           &nbp->coulomb_tab,
728                                                           &plist->sci,
729                                                           &plist->cj4,
730                                                           &plist->excl,
731                                                           &computeFshift);
732
733         launchGpuKernel(kernel, config, deviceStream, timingEvent, kernelName, kernelArgs);
734     }
735     else
736     {
737         const auto kernelArgs = prepareGpuKernelArguments(kernel,
738                                                           config,
739                                                           &adat->numTypes,
740                                                           &nbparams_params,
741                                                           &adat->xq,
742                                                           &adat->f,
743                                                           &adat->eLJ,
744                                                           &adat->eElec,
745                                                           &adat->fShift,
746                                                           &adat->atomTypes,
747                                                           &adat->shiftVec,
748                                                           &nbp->nbfp,
749                                                           &nbp->nbfp_comb,
750                                                           &nbp->coulomb_tab,
751                                                           &plist->sci,
752                                                           &plist->cj4,
753                                                           &plist->excl,
754                                                           &computeFshift);
755         launchGpuKernel(kernel, config, deviceStream, timingEvent, kernelName, kernelArgs);
756     }
757
758     if (bDoTime)
759     {
760         t->interaction[iloc].nb_k.closeTimingRegion(deviceStream);
761     }
762 }
763
764
765 /*! \brief Calculates the amount of shared memory required by the prune kernel.
766  *
767  *  Note that for the sake of simplicity we use the CUDA terminology "shared memory"
768  *  for OpenCL local memory.
769  *
770  * \param[in] num_threads_z cj4 concurrency equal to the number of threads/work items in the 3-rd
771  * dimension. \returns   the amount of local memory in bytes required by the pruning kernel
772  */
773 static inline int calc_shmem_required_prune(const int num_threads_z)
774 {
775     int shmem;
776
777     /* i-atom x in shared memory (for convenience we load all 4 components including q) */
778     shmem = c_nbnxnGpuNumClusterPerSupercluster * c_clSize * sizeof(float) * 4;
779     /* cj in shared memory, for each warp separately
780      * Note: only need to load once per wavefront, but to keep the code simple,
781      * for now we load twice on AMD.
782      */
783     shmem += num_threads_z * c_nbnxnGpuClusterpairSplit * c_nbnxnGpuJgroupSize * sizeof(int);
784     /* Warp vote, requires one uint per warp/32 threads per block. */
785     shmem += sizeof(cl_uint) * 2 * num_threads_z;
786
787     return shmem;
788 }
789
790 /*! \brief
791  * Launch the pairlist prune only kernel for the given locality.
792  * \p numParts tells in how many parts, i.e. calls the list will be pruned.
793  */
794 void gpu_launch_kernel_pruneonly(NbnxmGpu* nb, const InteractionLocality iloc, const int numParts)
795 {
796     NBAtomData*         adat         = nb->atdat;
797     NBParamGpu*         nbp          = nb->nbparam;
798     gpu_plist*          plist        = nb->plist[iloc];
799     cl_timers_t*        t            = nb->timers;
800     const DeviceStream& deviceStream = *nb->deviceStreams[iloc];
801     bool                bDoTime      = nb->bDoTime;
802
803     if (plist->haveFreshList)
804     {
805         GMX_ASSERT(numParts == 1, "With first pruning we expect 1 part");
806
807         /* Set rollingPruningNumParts to signal that it is not set */
808         plist->rollingPruningNumParts = 0;
809         plist->rollingPruningPart     = 0;
810     }
811     else
812     {
813         if (plist->rollingPruningNumParts == 0)
814         {
815             plist->rollingPruningNumParts = numParts;
816         }
817         else
818         {
819             GMX_ASSERT(numParts == plist->rollingPruningNumParts,
820                        "It is not allowed to change numParts in between list generation steps");
821         }
822     }
823
824     /* Use a local variable for part and update in plist, so we can return here
825      * without duplicating the part increment code.
826      */
827     int part = plist->rollingPruningPart;
828
829     plist->rollingPruningPart++;
830     if (plist->rollingPruningPart >= plist->rollingPruningNumParts)
831     {
832         plist->rollingPruningPart = 0;
833     }
834
835     /* Compute the number of list entries to prune in this pass */
836     int numSciInPart = (plist->nsci - part) / numParts;
837
838     /* Don't launch the kernel if there is no work to do. */
839     if (numSciInPart <= 0)
840     {
841         plist->haveFreshList = false;
842
843         return;
844     }
845
846     GpuRegionTimer* timer = nullptr;
847     if (bDoTime)
848     {
849         timer = &(plist->haveFreshList ? t->interaction[iloc].prune_k : t->interaction[iloc].rollingPrune_k);
850     }
851
852     /* beginning of timed prune calculation section */
853     if (bDoTime)
854     {
855         timer->openTimingRegion(deviceStream);
856     }
857
858     /* Kernel launch config:
859      * - The thread block dimensions match the size of i-clusters, j-clusters,
860      *   and j-cluster concurrency, in x, y, and z, respectively.
861      * - The 1D block-grid contains as many blocks as super-clusters.
862      */
863     int num_threads_z = c_pruneKernelJ4Concurrency;
864     /* kernel launch config */
865     KernelLaunchConfig config;
866     config.sharedMemorySize = calc_shmem_required_prune(num_threads_z);
867     config.blockSize[0]     = c_clSize;
868     config.blockSize[1]     = c_clSize;
869     config.blockSize[2]     = num_threads_z;
870     config.gridSize[0]      = numSciInPart;
871
872     validate_global_work_size(config, 3, &nb->deviceContext_->deviceInfo());
873
874     if (debug)
875     {
876         fprintf(debug,
877                 "Pruning GPU kernel launch configuration:\n\tLocal work size: %zux%zux%zu\n\t"
878                 "\tGlobal work size: %zux%zu\n\t#Super-clusters/clusters: %d/%d (%d)\n"
879                 "\tShMem: %zu\n",
880                 config.blockSize[0],
881                 config.blockSize[1],
882                 config.blockSize[2],
883                 config.blockSize[0] * config.gridSize[0],
884                 config.blockSize[1] * config.gridSize[1],
885                 plist->nsci * c_nbnxnGpuNumClusterPerSupercluster,
886                 c_nbnxnGpuNumClusterPerSupercluster,
887                 plist->na_c,
888                 config.sharedMemorySize);
889     }
890
891     cl_nbparam_params_t nbparams_params;
892     fillin_ocl_structures(nbp, &nbparams_params);
893
894     auto*          timingEvent  = bDoTime ? timer->fetchNextEvent() : nullptr;
895     constexpr char kernelName[] = "k_pruneonly";
896     const auto     pruneKernel  = selectPruneKernel(nb->kernel_pruneonly, plist->haveFreshList);
897     const auto     kernelArgs   = prepareGpuKernelArguments(pruneKernel,
898                                                       config,
899                                                       &nbparams_params,
900                                                       &adat->xq,
901                                                       &adat->shiftVec,
902                                                       &plist->sci,
903                                                       &plist->cj4,
904                                                       &plist->imask,
905                                                       &numParts,
906                                                       &part);
907     launchGpuKernel(pruneKernel, config, deviceStream, timingEvent, kernelName, kernelArgs);
908
909     if (plist->haveFreshList)
910     {
911         plist->haveFreshList = false;
912         /* Mark that pruning has been done */
913         nb->timers->interaction[iloc].didPrune = true;
914     }
915     else
916     {
917         /* Mark that rolling pruning has been done */
918         nb->timers->interaction[iloc].didRollingPrune = true;
919     }
920
921     if (bDoTime)
922     {
923         timer->closeTimingRegion(deviceStream);
924     }
925 }
926
927 /*! \brief
928  * Launch asynchronously the download of nonbonded forces from the GPU
929  * (and energies/shift forces if required).
930  */
931 void gpu_launch_cpyback(NbnxmGpu*                nb,
932                         struct nbnxn_atomdata_t* nbatom,
933                         const gmx::StepWorkload& stepWork,
934                         const AtomLocality       aloc)
935 {
936     GMX_ASSERT(nb, "Need a valid nbnxn_gpu object");
937
938     cl_int gmx_unused cl_error;
939     int               adat_begin, adat_len; /* local/nonlocal offset and length used for xq and f */
940
941     /* determine interaction locality from atom locality */
942     const InteractionLocality iloc = gpuAtomToInteractionLocality(aloc);
943     GMX_ASSERT(iloc == InteractionLocality::Local
944                        || (iloc == InteractionLocality::NonLocal && nb->bNonLocalStreamDoneMarked == false),
945                "Non-local stream is indicating that the copy back event is enqueued at the "
946                "beginning of the copy back function.");
947
948     NBAtomData*         adat         = nb->atdat;
949     cl_timers_t*        t            = nb->timers;
950     bool                bDoTime      = nb->bDoTime;
951     const DeviceStream& deviceStream = *nb->deviceStreams[iloc];
952
953     /* don't launch non-local copy-back if there was no non-local work to do */
954     if ((iloc == InteractionLocality::NonLocal) && !haveGpuShortRangeWork(*nb, iloc))
955     {
956         /* TODO An alternative way to signal that non-local work is
957            complete is to use a clEnqueueMarker+clEnqueueBarrier
958            pair. However, the use of bNonLocalStreamDoneMarked has the
959            advantage of being local to the host, so probably minimizes
960            overhead. Curiously, for NVIDIA OpenCL with an empty-domain
961            test case, overall simulation performance was higher with
962            the API calls, but this has not been tested on AMD OpenCL,
963            so could be worth considering in future. */
964         nb->bNonLocalStreamDoneMarked = false;
965         return;
966     }
967
968     getGpuAtomRange(adat, aloc, &adat_begin, &adat_len);
969
970     /* beginning of timed D2H section */
971     if (bDoTime)
972     {
973         t->xf[aloc].nb_d2h.openTimingRegion(deviceStream);
974     }
975
976     /* With DD the local D2H transfer can only start after the non-local
977        has been launched. */
978     if (iloc == InteractionLocality::Local && nb->bNonLocalStreamDoneMarked)
979     {
980         nb->nonlocal_done.enqueueWaitEvent(deviceStream);
981         nb->bNonLocalStreamDoneMarked = false;
982     }
983
984     /* DtoH f */
985     GMX_ASSERT(sizeof(*nbatom->out[0].f.data()) == sizeof(float),
986                "The host force buffer should be in single precision to match device data size.");
987     copyFromDeviceBuffer(reinterpret_cast<Float3*>(nbatom->out[0].f.data()) + adat_begin,
988                          &adat->f,
989                          adat_begin,
990                          adat_len,
991                          deviceStream,
992                          GpuApiCallBehavior::Async,
993                          bDoTime ? t->xf[aloc].nb_d2h.fetchNextEvent() : nullptr);
994
995     /* kick off work */
996     cl_error = clFlush(deviceStream.stream());
997     GMX_ASSERT(cl_error == CL_SUCCESS, ("clFlush failed: " + ocl_get_error_string(cl_error)).c_str());
998
999     /* After the non-local D2H is launched the nonlocal_done event can be
1000        recorded which signals that the local D2H can proceed. This event is not
1001        placed after the non-local kernel because we first need the non-local
1002        data back first. */
1003     if (iloc == InteractionLocality::NonLocal)
1004     {
1005         nb->nonlocal_done.markEvent(deviceStream);
1006         nb->bNonLocalStreamDoneMarked = true;
1007     }
1008
1009     /* only transfer energies in the local stream */
1010     if (iloc == InteractionLocality::Local)
1011     {
1012         /* DtoH fshift when virial is needed */
1013         if (stepWork.computeVirial)
1014         {
1015             static_assert(
1016                     sizeof(*nb->nbst.fShift) == sizeof(Float3),
1017                     "Sizes of host- and device-side shift vector elements should be the same.");
1018             copyFromDeviceBuffer(nb->nbst.fShift,
1019                                  &adat->fShift,
1020                                  0,
1021                                  SHIFTS,
1022                                  deviceStream,
1023                                  GpuApiCallBehavior::Async,
1024                                  bDoTime ? t->xf[aloc].nb_d2h.fetchNextEvent() : nullptr);
1025         }
1026
1027         /* DtoH energies */
1028         if (stepWork.computeEnergy)
1029         {
1030             static_assert(sizeof(*nb->nbst.eLJ) == sizeof(float),
1031                           "Sizes of host- and device-side LJ energy terms should be the same.");
1032             copyFromDeviceBuffer(nb->nbst.eLJ,
1033                                  &adat->eLJ,
1034                                  0,
1035                                  1,
1036                                  deviceStream,
1037                                  GpuApiCallBehavior::Async,
1038                                  bDoTime ? t->xf[aloc].nb_d2h.fetchNextEvent() : nullptr);
1039             static_assert(sizeof(*nb->nbst.eElec) == sizeof(float),
1040                           "Sizes of host- and device-side electrostatic energy terms should be the "
1041                           "same.");
1042             copyFromDeviceBuffer(nb->nbst.eElec,
1043                                  &adat->eElec,
1044                                  0,
1045                                  1,
1046                                  deviceStream,
1047                                  GpuApiCallBehavior::Async,
1048                                  bDoTime ? t->xf[aloc].nb_d2h.fetchNextEvent() : nullptr);
1049         }
1050     }
1051
1052     if (bDoTime)
1053     {
1054         t->xf[aloc].nb_d2h.closeTimingRegion(deviceStream);
1055     }
1056 }
1057
1058 } // namespace Nbnxm