2 * This file is part of the GROMACS molecular simulation package.
4 * Copyright (c) 1991-2000, University of Groningen, The Netherlands.
5 * Copyright (c) 2001-2004, The GROMACS development team.
6 * Copyright (c) 2013,2014,2015,2017,2018 by the GROMACS development team.
7 * Copyright (c) 2019,2020,2021, by the GROMACS development team, led by
8 * Mark Abraham, David van der Spoel, Berk Hess, and Erik Lindahl,
9 * and including many others, as listed in the AUTHORS file in the
10 * top-level source directory and at http://www.gromacs.org.
12 * GROMACS is free software; you can redistribute it and/or
13 * modify it under the terms of the GNU Lesser General Public License
14 * as published by the Free Software Foundation; either version 2.1
15 * of the License, or (at your option) any later version.
17 * GROMACS is distributed in the hope that it will be useful,
18 * but WITHOUT ANY WARRANTY; without even the implied warranty of
19 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
20 * Lesser General Public License for more details.
22 * You should have received a copy of the GNU Lesser General Public
23 * License along with GROMACS; if not, see
24 * http://www.gnu.org/licenses, or write to the Free Software Foundation,
25 * Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
27 * If you want to redistribute modifications to GROMACS, please
28 * consider that scientific software is very special. Version
29 * control is crucial - bugs must be traceable. We will be happy to
30 * consider code for inclusion in the official distribution, but
31 * derived work must not be called official GROMACS. Details are found
32 * in the README & COPYING files - if they are missing, get the
33 * official version at http://www.gromacs.org.
35 * To help us fund GROMACS development, we humbly ask that you cite
36 * the research papers on the package. Check out http://www.gromacs.org.
41 #include "pme_gather.h"
43 #include "gromacs/math/vec.h"
44 #include "gromacs/simd/simd.h"
45 #include "gromacs/utility/basedefinitions.h"
46 #include "gromacs/utility/gmxassert.h"
47 #include "gromacs/utility/smalloc.h"
48 #include "gromacs/utility/typetraits.h"
50 #include "pme_internal.h"
52 #include "pme_spline_work.h"
54 using namespace gmx; // TODO: Remove when this file is moved into gmx namespace
56 /* Spline function. Goals: 1) Force compiler to instantiate function separately
57 for each compile-time value of order and once more for any possible (runtime)
58 value. 2) Allow overloading for specific compile-time values.
59 The Int template argument can be either int (runtime) or an object of type
60 integral_constant<int, N> (compile-time). Common runtime values can be
61 converted to compile-time values with a switch statement. For the compile
62 value the compiler is required to instantiate the function separately for
63 each value. The function can be overloaded for specific compile-time values
64 using integral_constant<int, N> where N is either a specific value or an
65 enable_if constrained non-type template parameter. The most specific overload
66 (specific value > int template parameter > general function) is called. Inside
67 the function the order argument can be used as regular int because
68 integral_constant has a proper conversion.
70 SIMD do_fspline() template funtions will be used for PME order 4 and 5
71 when the SIMD module has support for SIMD4 for the architecture used.
72 For SIMD4 without unaligned load/store support:
73 order 4 and 5 use the order 4+5 aligned SIMD template
74 For SIMD4 with unaligned load/store support:
75 order 4 uses the order 4 unaligned SIMD template
76 order 5 uses the order 4+5 aligned SIMD template
80 do_fspline(const gmx_pme_t* pme,
81 const real* gmx_restrict grid,
82 const PmeAtomComm* gmx_restrict atc,
83 const splinedata_t* gmx_restrict spline,
85 pme(pme), grid(grid), atc(atc), spline(spline), nn(nn)
89 template<typename Int>
90 RVec operator()(Int order) const
92 static_assert(isIntegralConstant<Int, int>::value || std::is_same_v<Int, int>,
93 "'order' needs to be either of type integral_constant<int,N> or int.");
95 const int norder = nn * order;
97 /* Pointer arithmetic alert, next six statements */
98 const real* const gmx_restrict thx = spline->theta.coefficients[XX] + norder;
99 const real* const gmx_restrict thy = spline->theta.coefficients[YY] + norder;
100 const real* const gmx_restrict thz = spline->theta.coefficients[ZZ] + norder;
101 const real* const gmx_restrict dthx = spline->dtheta.coefficients[XX] + norder;
102 const real* const gmx_restrict dthy = spline->dtheta.coefficients[YY] + norder;
103 const real* const gmx_restrict dthz = spline->dtheta.coefficients[ZZ] + norder;
107 for (int ithx = 0; (ithx < order); ithx++)
109 const int index_x = (idxX + ithx) * gridNY * gridNZ;
110 const real tx = thx[ithx];
111 const real dx = dthx[ithx];
113 for (int ithy = 0; (ithy < order); ithy++)
115 const int index_xy = index_x + (idxY + ithy) * gridNZ;
116 const real ty = thy[ithy];
117 const real dy = dthy[ithy];
118 real fxy1 = 0, fz1 = 0;
120 for (int ithz = 0; (ithz < order); ithz++)
122 const real gval = grid[index_xy + (idxZ + ithz)];
123 fxy1 += thz[ithz] * gval;
124 fz1 += dthz[ithz] * gval;
126 f[XX] += dx * ty * fxy1;
127 f[YY] += tx * dy * fxy1;
128 f[ZZ] += tx * ty * fz1;
135 // TODO: Consider always have at least a dummy implementation of Simd (enough for first phase of two-phase lookup) and then use enable_if instead of #ifdef
136 #if PME_4NSIMD_GATHER
137 /* Gather for one charge with pme_order=4 with unaligned SIMD4 load+store.
138 * Uses 4N SIMD where N is SIMD_WIDTH/4 to operate on all of z and N of y.
139 * This code does not assume any memory alignment for the grid.
141 RVec operator()(std::integral_constant<int, 4> /*unused*/) const
143 const int norder = nn * 4;
144 /* Pointer arithmetic alert, next six statements */
145 const real* const gmx_restrict thx = spline->theta.coefficients[XX] + norder;
146 const real* const gmx_restrict thy = spline->theta.coefficients[YY] + norder;
147 const real* const gmx_restrict thz = spline->theta.coefficients[ZZ] + norder;
148 const real* const gmx_restrict dthx = spline->dtheta.coefficients[XX] + norder;
149 const real* const gmx_restrict dthy = spline->dtheta.coefficients[YY] + norder;
150 const real* const gmx_restrict dthz = spline->dtheta.coefficients[ZZ] + norder;
152 Simd4NReal fx_S = setZero();
153 Simd4NReal fy_S = setZero();
154 Simd4NReal fz_S = setZero();
156 /* With order 4 the z-spline is actually aligned */
157 const Simd4NReal tz_S = load4DuplicateN(thz);
158 const Simd4NReal dz_S = load4DuplicateN(dthz);
160 for (int ithx = 0; ithx < 4; ithx++)
162 const int index_x = (idxX + ithx) * gridNY * gridNZ;
163 const Simd4NReal tx_S = Simd4NReal(thx[ithx]);
164 const Simd4NReal dx_S = Simd4NReal(dthx[ithx]);
166 for (int ithy = 0; ithy < 4; ithy += GMX_SIMD4N_REAL_WIDTH / 4)
168 const int index_xy = index_x + (idxY + ithy) * gridNZ;
170 const Simd4NReal ty_S = loadUNDuplicate4(thy + ithy);
171 const Simd4NReal dy_S = loadUNDuplicate4(dthy + ithy);
173 const Simd4NReal gval_S = loadU4NOffset(grid + index_xy + idxZ, gridNZ);
176 const Simd4NReal fxy1_S = tz_S * gval_S;
177 const Simd4NReal fz1_S = dz_S * gval_S;
179 fx_S = fma(dx_S * ty_S, fxy1_S, fx_S);
180 fy_S = fma(tx_S * dy_S, fxy1_S, fy_S);
181 fz_S = fma(tx_S * ty_S, fz1_S, fz_S);
185 return { reduce(fx_S), reduce(fy_S), reduce(fz_S) };
189 #ifdef PME_SIMD4_SPREAD_GATHER
190 /* Load order elements from unaligned memory into two 4-wide SIMD */
192 static inline void loadOrderU(const real* data,
193 std::integral_constant<int, order> /*unused*/,
198 # ifdef PME_SIMD4_UNALIGNED
199 *S0 = load4U(data - offset);
200 *S1 = load4U(data - offset + 4);
202 alignas(GMX_SIMD_ALIGNMENT) real buf_aligned[GMX_SIMD4_WIDTH * 2];
203 /* Copy data to an aligned buffer */
204 for (int i = 0; i < order; i++)
206 buf_aligned[offset + i] = data[i];
208 *S0 = load4(buf_aligned);
209 *S1 = load4(buf_aligned + 4);
214 #ifdef PME_SIMD4_SPREAD_GATHER
215 /* This code assumes that the grid is allocated 4-real aligned
216 * and that pme->pmegrid_nz is a multiple of 4.
217 * This code supports pme_order <= 5.
220 std::enable_if_t<Order == 4 || Order == 5, RVec> operator()(std::integral_constant<int, Order> order) const
222 const int norder = nn * order;
223 GMX_ASSERT(gridNZ % 4 == 0,
224 "For aligned SIMD4 operations the grid size has to be padded up to a multiple "
226 /* Pointer arithmetic alert, next six statements */
227 const real* const gmx_restrict thx = spline->theta.coefficients[XX] + norder;
228 const real* const gmx_restrict thy = spline->theta.coefficients[YY] + norder;
229 const real* const gmx_restrict thz = spline->theta.coefficients[ZZ] + norder;
230 const real* const gmx_restrict dthx = spline->dtheta.coefficients[XX] + norder;
231 const real* const gmx_restrict dthy = spline->dtheta.coefficients[YY] + norder;
232 const real* const gmx_restrict dthz = spline->dtheta.coefficients[ZZ] + norder;
234 struct pme_spline_work* const work = pme->spline_work;
236 const int offset = idxZ & 3;
238 Simd4Real fx_S = setZero();
239 Simd4Real fy_S = setZero();
240 Simd4Real fz_S = setZero();
242 Simd4Real tz_S0, tz_S1, dz_S0, dz_S1;
243 loadOrderU(thz, order, offset, &tz_S0, &tz_S1);
244 loadOrderU(dthz, order, offset, &dz_S0, &dz_S1);
246 tz_S0 = selectByMask(tz_S0, work->mask_S0[offset]);
247 dz_S0 = selectByMask(dz_S0, work->mask_S0[offset]);
248 tz_S1 = selectByMask(tz_S1, work->mask_S1[offset]);
249 dz_S1 = selectByMask(dz_S1, work->mask_S1[offset]);
251 for (int ithx = 0; (ithx < order); ithx++)
253 const int index_x = (idxX + ithx) * gridNY * gridNZ;
254 const Simd4Real tx_S = Simd4Real(thx[ithx]);
255 const Simd4Real dx_S = Simd4Real(dthx[ithx]);
257 for (int ithy = 0; (ithy < order); ithy++)
259 const int index_xy = index_x + (idxY + ithy) * gridNZ;
260 const Simd4Real ty_S = Simd4Real(thy[ithy]);
261 const Simd4Real dy_S = Simd4Real(dthy[ithy]);
263 const Simd4Real gval_S0 = load4(grid + index_xy + idxZ - offset);
264 const Simd4Real gval_S1 = load4(grid + index_xy + idxZ - offset + 4);
266 const Simd4Real fxy1_S0 = tz_S0 * gval_S0;
267 const Simd4Real fz1_S0 = dz_S0 * gval_S0;
268 const Simd4Real fxy1_S1 = tz_S1 * gval_S1;
269 const Simd4Real fz1_S1 = dz_S1 * gval_S1;
271 const Simd4Real fxy1_S = fxy1_S0 + fxy1_S1;
272 const Simd4Real fz1_S = fz1_S0 + fz1_S1;
274 fx_S = fma(dx_S * ty_S, fxy1_S, fx_S);
275 fy_S = fma(tx_S * dy_S, fxy1_S, fy_S);
276 fz_S = fma(tx_S * ty_S, fz1_S, fz_S);
280 return { reduce(fx_S), reduce(fy_S), reduce(fz_S) };
284 const gmx_pme_t* const pme;
285 const real* const gmx_restrict grid;
286 const PmeAtomComm* const gmx_restrict atc;
287 const splinedata_t* const gmx_restrict spline;
290 const int gridNY = pme->pmegrid_ny;
291 const int gridNZ = pme->pmegrid_nz;
293 const int* const idxptr = atc->idx[spline->ind[nn]];
294 const int idxX = idxptr[XX];
295 const int idxY = idxptr[YY];
296 const int idxZ = idxptr[ZZ];
300 void gather_f_bsplines(const gmx_pme_t* pme,
303 const PmeAtomComm* atc,
304 const splinedata_t* spline,
307 /* sum forces for local particles */
309 const int order = pme->pme_order;
310 const int nx = pme->nkx;
311 const int ny = pme->nky;
312 const int nz = pme->nkz;
314 const real rxx = pme->recipbox[XX][XX];
315 const real ryx = pme->recipbox[YY][XX];
316 const real ryy = pme->recipbox[YY][YY];
317 const real rzx = pme->recipbox[ZZ][XX];
318 const real rzy = pme->recipbox[ZZ][YY];
319 const real rzz = pme->recipbox[ZZ][ZZ];
321 /* Extract the buffer for force output */
322 rvec* gmx_restrict force = as_rvec_array(atc->f.data());
324 /* Note that unrolling this loop by templating this function on order
325 * deteriorates performance significantly with gcc5/6/7.
327 for (int nn = 0; nn < spline->n; nn++)
329 const int n = spline->ind[nn];
330 const real coefficient = scale * atc->coefficient[n];
338 if (coefficient != 0)
341 const auto spline_func = do_fspline(pme, grid, atc, spline, nn);
345 case 4: f = spline_func(std::integral_constant<int, 4>()); break;
346 case 5: f = spline_func(std::integral_constant<int, 5>()); break;
347 default: f = spline_func(order); break;
350 force[n][XX] += -coefficient * (f[XX] * nx * rxx);
351 force[n][YY] += -coefficient * (f[XX] * nx * ryx + f[YY] * ny * ryy);
352 force[n][ZZ] += -coefficient * (f[XX] * nx * rzx + f[YY] * ny * rzy + f[ZZ] * nz * rzz);
355 /* Since the energy and not forces are interpolated
356 * the net force might not be exactly zero.
357 * This can be solved by also interpolating F, but
358 * that comes at a cost.
359 * A better hack is to remove the net force every
360 * step, but that must be done at a higher level
361 * since this routine doesn't see all atoms if running
362 * in parallel. Don't know how important it is? EL 990726
367 real gather_energy_bsplines(gmx_pme_t* pme, const real* grid, PmeAtomComm* atc)
369 splinedata_t* spline;
370 int ithx, ithy, ithz, i0, j0, k0;
371 int index_x, index_xy;
373 real energy, pot, tx, ty, coefficient, gval;
374 real * thx, *thy, *thz;
378 spline = &atc->spline[0];
380 order = pme->pme_order;
383 for (int n = 0; n < atc->numAtoms(); n++)
385 coefficient = atc->coefficient[n];
387 if (coefficient != 0)
389 idxptr = atc->idx[n];
396 /* Pointer arithmetic alert, next three statements */
397 thx = spline->theta.coefficients[XX] + norder;
398 thy = spline->theta.coefficients[YY] + norder;
399 thz = spline->theta.coefficients[ZZ] + norder;
402 for (ithx = 0; (ithx < order); ithx++)
404 index_x = (i0 + ithx) * pme->pmegrid_ny * pme->pmegrid_nz;
407 for (ithy = 0; (ithy < order); ithy++)
409 index_xy = index_x + (j0 + ithy) * pme->pmegrid_nz;
412 for (ithz = 0; (ithz < order); ithz++)
414 gval = grid[index_xy + (k0 + ithz)];
415 pot += tx * ty * thz[ithz] * gval;
420 energy += pot * coefficient;