src/gromacs/nbnxm/kernels_simd_2xmm/kernel_inner.h

   1 /*
   2  * This file is part of the GROMACS molecular simulation package.
   3  *
   4  * Copyright (c) 2012,2013,2014,2015,2016,2017,2018,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
  36 /* This is the innermost loop contents for the 4 x N atom simd kernel.
  37  * This flavor of the kernel duplicates the data for N j-particles in
  38  * 2xN wide simd registers to do operate on 2 i-particles at once.
  39  * This leads to 4/2=2 sets of most instructions. Therefore we call
  40  * this kernel 2x(N+N) = 2xnn
  41  *
  42  * This 2xnn kernel is basically the 4xn equivalent with half the registers
  43  * and instructions removed.
  44  *
  45  * An alternative would be to load to different cluster of N j-particles
  46  * into simd registers, giving a 4x(N+N) kernel. This doubles the amount
  47  * of instructions, which could lead to better scheduling. But we actually
  48  * observed worse scheduling for the AVX-256 4x8 normal analytical PME
  49  * kernel, which has a lower pair throughput than 2x(4+4) with gcc 4.7.
  50  * It could be worth trying this option, but it takes some more effort.
  51  * This 2xnn kernel is basically the 4xn equivalent with
  52  */
  53
  54
  55 /* When calculating RF or Ewald interactions we calculate the electrostatic/LJ
  56  * forces on excluded atom pairs here in the non-bonded loops.
  57  * But when energies and/or virial is required we calculate them
  58  * separately to as then it is easier to separate the energy and virial
  59  * contributions.
  60  */
  61 #if defined CHECK_EXCLS && (defined CALC_COULOMB || defined LJ_EWALD_GEOM)
  62 #    define EXCL_FORCES
  63 #endif
  64
  65 {
  66     int cj, aj, ajx, ajy, ajz;
  67
  68 #ifdef ENERGY_GROUPS
  69     /* Energy group indices for two atoms packed into one int */
  70     int egp_jj[UNROLLJ / 2];
  71 #endif
  72
  73 #ifdef CHECK_EXCLS
  74     /* Interaction (non-exclusion) mask of all 1's or 0's */
  75     SimdBool interact_S0;
  76     SimdBool interact_S2;
  77 #endif
  78
  79     SimdReal jx_S, jy_S, jz_S;
  80     SimdReal dx_S0, dy_S0, dz_S0;
  81     SimdReal dx_S2, dy_S2, dz_S2;
  82     SimdReal tx_S0, ty_S0, tz_S0;
  83     SimdReal tx_S2, ty_S2, tz_S2;
  84     SimdReal rsq_S0, rinv_S0, rinvsq_S0;
  85     SimdReal rsq_S2, rinv_S2, rinvsq_S2;
  86     /* wco: within cut-off, mask of all 1's or 0's */
  87     SimdBool wco_S0;
  88     SimdBool wco_S2;
  89 #ifdef VDW_CUTOFF_CHECK
  90     SimdBool wco_vdw_S0;
  91 #    ifndef HALF_LJ
  92     SimdBool wco_vdw_S2;
  93 #    endif
  94 #endif
  95
  96 #if (defined CALC_COULOMB && defined CALC_COUL_TAB) || defined LJ_FORCE_SWITCH || defined LJ_POT_SWITCH
  97     SimdReal                                                        r_S0;
  98 #    if (defined CALC_COULOMB && defined CALC_COUL_TAB) || !defined HALF_LJ
  99     SimdReal                                                        r_S2;
 100 #    endif
 101 #endif
 102
 103 #if defined LJ_FORCE_SWITCH || defined LJ_POT_SWITCH
 104     SimdReal rsw_S0, rsw2_S0;
 105 #    ifndef HALF_LJ
 106     SimdReal rsw_S2, rsw2_S2;
 107 #    endif
 108 #endif
 109
 110 #ifdef CALC_COULOMB
 111 #    ifdef CHECK_EXCLS
 112     /* 1/r masked with the interaction mask */
 113     SimdReal rinv_ex_S0;
 114     SimdReal rinv_ex_S2;
 115 #    endif
 116     SimdReal jq_S;
 117     SimdReal qq_S0;
 118     SimdReal qq_S2;
 119 #    ifdef CALC_COUL_TAB
 120     /* The force (PME mesh force) we need to subtract from 1/r^2 */
 121     SimdReal fsub_S0;
 122     SimdReal fsub_S2;
 123 #    endif
 124 #    ifdef CALC_COUL_EWALD
 125     SimdReal brsq_S0, brsq_S2;
 126     SimdReal ewcorr_S0, ewcorr_S2;
 127 #    endif
 128
 129     /* frcoul = (1/r - fsub)*r */
 130     SimdReal frcoul_S0;
 131     SimdReal frcoul_S2;
 132 #    ifdef CALC_COUL_TAB
 133     /* For tables: r, rs=r/sp, rf=floor(rs), frac=rs-rf */
 134     SimdReal rs_S0, rf_S0, frac_S0;
 135     SimdReal rs_S2, rf_S2, frac_S2;
 136     /* Table index: rs truncated to an int */
 137     SimdInt32 ti_S0, ti_S2;
 138     /* Linear force table values */
 139     SimdReal ctab0_S0, ctab1_S0;
 140     SimdReal ctab0_S2, ctab1_S2;
 141 #        ifdef CALC_ENERGIES
 142     /* Quadratic energy table value */
 143     SimdReal ctabv_S0, dum_S0;
 144     SimdReal ctabv_S2, dum_S2;
 145 #        endif
 146 #    endif
 147 #    if defined CALC_ENERGIES && (defined CALC_COUL_EWALD || defined CALC_COUL_TAB)
 148     /* The potential (PME mesh) we need to subtract from 1/r */
 149     SimdReal vc_sub_S0;
 150     SimdReal vc_sub_S2;
 151 #    endif
 152 #    ifdef CALC_ENERGIES
 153     /* Electrostatic potential */
 154     SimdReal vcoul_S0;
 155     SimdReal vcoul_S2;
 156 #    endif
 157 #endif
 158     /* The force times 1/r */
 159     SimdReal fscal_S0;
 160     SimdReal fscal_S2;
 161
 162 #ifdef CALC_LJ
 163 #    ifdef LJ_COMB_LB
 164     /* LJ sigma_j/2 and sqrt(epsilon_j) */
 165     SimdReal hsig_j_S, seps_j_S;
 166     /* LJ sigma_ij and epsilon_ij */
 167     SimdReal sig_S0, eps_S0;
 168 #        ifndef HALF_LJ
 169     SimdReal sig_S2, eps_S2;
 170 #        endif
 171 #        ifdef CALC_ENERGIES
 172     SimdReal sig2_S0, sig6_S0;
 173 #            ifndef HALF_LJ
 174     SimdReal sig2_S2, sig6_S2;
 175 #            endif
 176 #        endif /* LJ_COMB_LB */
 177 #    endif     /* CALC_LJ */
 178
 179 #    ifdef LJ_COMB_GEOM
 180     SimdReal c6s_j_S, c12s_j_S;
 181 #    endif
 182
 183 #    if defined LJ_COMB_GEOM || defined LJ_COMB_LB || defined LJ_EWALD_GEOM
 184     /* Index for loading LJ parameters, complicated when interleaving */
 185     int aj2;
 186 #    endif
 187
 188     /* Intermediate variables for LJ calculation */
 189 #    ifndef LJ_COMB_LB
 190     SimdReal rinvsix_S0;
 191 #        ifndef HALF_LJ
 192     SimdReal rinvsix_S2;
 193 #        endif
 194 #    endif
 195 #    ifdef LJ_COMB_LB
 196     SimdReal sir_S0, sir2_S0, sir6_S0;
 197 #        ifndef HALF_LJ
 198     SimdReal sir_S2, sir2_S2, sir6_S2;
 199 #        endif
 200 #    endif
 201
 202     SimdReal FrLJ6_S0, FrLJ12_S0, frLJ_S0;
 203 #    ifndef HALF_LJ
 204     SimdReal FrLJ6_S2, FrLJ12_S2, frLJ_S2;
 205 #    endif
 206 #endif /* CALC_LJ */
 207
 208     /* j-cluster index */
 209     cj = l_cj[cjind].cj;
 210
 211     /* Atom indices (of the first atom in the cluster) */
 212     aj = cj * UNROLLJ;
 213 #if defined CALC_LJ && (defined LJ_COMB_GEOM || defined LJ_COMB_LB || defined LJ_EWALD_GEOM)
 214     aj2 = aj * 2;
 215 #endif
 216     ajx = aj * DIM;
 217     ajy = ajx + STRIDE;
 218     ajz = ajy + STRIDE;
 219
 220 #ifdef CHECK_EXCLS
 221     gmx_load_simd_2xnn_interactions(static_cast<int>(l_cj[cjind].excl), filter_S0, filter_S2,
 222                                     &interact_S0, &interact_S2);
 223 #endif /* CHECK_EXCLS */
 224
 225     /* load j atom coordinates */
 226     jx_S = loadDuplicateHsimd(x + ajx);
 227     jy_S = loadDuplicateHsimd(x + ajy);
 228     jz_S = loadDuplicateHsimd(x + ajz);
 229
 230     /* Calculate distance */
 231     dx_S0 = ix_S0 - jx_S;
 232     dy_S0 = iy_S0 - jy_S;
 233     dz_S0 = iz_S0 - jz_S;
 234     dx_S2 = ix_S2 - jx_S;
 235     dy_S2 = iy_S2 - jy_S;
 236     dz_S2 = iz_S2 - jz_S;
 237
 238     /* rsq = dx*dx+dy*dy+dz*dz */
 239     rsq_S0 = norm2(dx_S0, dy_S0, dz_S0);
 240     rsq_S2 = norm2(dx_S2, dy_S2, dz_S2);
 241
 242     /* Do the cut-off check */
 243     wco_S0 = (rsq_S0 < rc2_S);
 244     wco_S2 = (rsq_S2 < rc2_S);
 245
 246 #ifdef CHECK_EXCLS
 247 #    ifdef EXCL_FORCES
 248     /* Only remove the (sub-)diagonal to avoid double counting */
 249 #        if UNROLLJ == UNROLLI
 250     if (cj == ci_sh)
 251     {
 252         wco_S0 = wco_S0 && diagonal_mask_S0;
 253         wco_S2 = wco_S2 && diagonal_mask_S2;
 254     }
 255 #        else
 256 #            if UNROLLJ == 2 * UNROLLI
 257     if (cj * 2 == ci_sh)
 258     {
 259         wco_S0 = wco_S0 && diagonal_mask0_S0;
 260         wco_S2 = wco_S2 && diagonal_mask0_S2;
 261     }
 262     else if (cj * 2 + 1 == ci_sh)
 263     {
 264         wco_S0 = wco_S0 && diagonal_mask1_S0;
 265         wco_S2 = wco_S2 && diagonal_mask1_S2;
 266     }
 267 #            else
 268 #                error "only UNROLLJ == UNROLLI*(1 or 2) currently supported in 2xnn kernels"
 269 #            endif
 270 #        endif
 271 #    else /* EXCL_FORCES */
 272     /* No exclusion forces: remove all excluded atom pairs from the list */
 273     wco_S0 = wco_S0 && interact_S0;
 274     wco_S2 = wco_S2 && interact_S2;
 275 #    endif
 276 #endif
 277
 278 #ifdef COUNT_PAIRS
 279     {
 280         int                              i, j;
 281         alignas(GMX_SIMD_ALIGNMENT) real tmp[GMX_SIMD_REAL_WIDTH];
 282
 283         for (i = 0; i < UNROLLI; i += 2)
 284         {
 285             store(tmp, rc2_S - (i == 0 ? rsq_S0 : rsq_S2));
 286             for (j = 0; j < 2 * UNROLLJ; j++)
 287             {
 288                 if (tmp[j] >= 0)
 289                 {
 290                     npair++;
 291                 }
 292             }
 293         }
 294     }
 295 #endif
 296
 297     // Ensure the distances do not fall below the limit where r^-12 overflows.
 298     // This should never happen for normal interactions.
 299     rsq_S0 = max(rsq_S0, minRsq_S);
 300     rsq_S2 = max(rsq_S2, minRsq_S);
 301
 302     /* Calculate 1/r */
 303     rinv_S0 = invsqrt(rsq_S0);
 304     rinv_S2 = invsqrt(rsq_S2);
 305
 306 #ifdef CALC_COULOMB
 307     /* Load parameters for j atom */
 308     jq_S  = loadDuplicateHsimd(q + aj);
 309     qq_S0 = iq_S0 * jq_S;
 310     qq_S2 = iq_S2 * jq_S;
 311 #endif
 312
 313 #ifdef CALC_LJ
 314 #    if !defined LJ_COMB_GEOM && !defined LJ_COMB_LB && !defined FIX_LJ_C
 315     SimdReal                                                     c6_S0, c12_S0;
 316     gatherLoadTransposeHsimd<c_simdBestPairAlignment>(nbfp0, nbfp1, type + aj, &c6_S0, &c12_S0);
 317 #        ifndef HALF_LJ
 318     SimdReal c6_S2, c12_S2;
 319     gatherLoadTransposeHsimd<c_simdBestPairAlignment>(nbfp2, nbfp3, type + aj, &c6_S2, &c12_S2);
 320 #        endif
 321 #    endif /* not defined any LJ rule */
 322
 323 #    ifdef LJ_COMB_GEOM
 324     c6s_j_S        = loadDuplicateHsimd(ljc + aj2);
 325     c12s_j_S       = loadDuplicateHsimd(ljc + aj2 + STRIDE);
 326     SimdReal c6_S0 = c6s_S0 * c6s_j_S;
 327 #        ifndef HALF_LJ
 328     SimdReal c6_S2 = c6s_S2 * c6s_j_S;
 329 #        endif
 330     SimdReal c12_S0 = c12s_S0 * c12s_j_S;
 331 #        ifndef HALF_LJ
 332     SimdReal c12_S2 = c12s_S2 * c12s_j_S;
 333 #        endif
 334 #    endif /* LJ_COMB_GEOM */
 335
 336 #    ifdef LJ_COMB_LB
 337     hsig_j_S = loadDuplicateHsimd(ljc + aj2);
 338     seps_j_S = loadDuplicateHsimd(ljc + aj2 + STRIDE);
 339
 340     sig_S0 = hsig_i_S0 + hsig_j_S;
 341     eps_S0 = seps_i_S0 * seps_j_S;
 342 #        ifndef HALF_LJ
 343     sig_S2 = hsig_i_S2 + hsig_j_S;
 344     eps_S2 = seps_i_S2 * seps_j_S;
 345 #        endif
 346 #    endif /* LJ_COMB_LB */
 347
 348 #endif /* CALC_LJ */
 349
 350     /* Set rinv to zero for r beyond the cut-off */
 351     rinv_S0 = selectByMask(rinv_S0, wco_S0);
 352     rinv_S2 = selectByMask(rinv_S2, wco_S2);
 353
 354     rinvsq_S0 = rinv_S0 * rinv_S0;
 355     rinvsq_S2 = rinv_S2 * rinv_S2;
 356
 357 #ifdef CALC_COULOMB
 358     /* Note that here we calculate force*r, not the usual force/r.
 359      * This allows avoiding masking the reaction-field contribution,
 360      * as frcoul is later multiplied by rinvsq which has been
 361      * masked with the cut-off check.
 362      */
 363
 364 #    ifdef EXCL_FORCES
 365     /* Only add 1/r for non-excluded atom pairs */
 366     rinv_ex_S0 = selectByMask(rinv_S0, interact_S0);
 367     rinv_ex_S2 = selectByMask(rinv_S2, interact_S2);
 368 #    else
 369     /* No exclusion forces, we always need 1/r */
 370 #        define rinv_ex_S0 rinv_S0
 371 #        define rinv_ex_S2 rinv_S2
 372 #    endif
 373
 374 #    ifdef CALC_COUL_RF
 375     /* Electrostatic interactions */
 376     frcoul_S0 = qq_S0 * fma(rsq_S0, mrc_3_S, rinv_ex_S0);
 377     frcoul_S2 = qq_S2 * fma(rsq_S2, mrc_3_S, rinv_ex_S2);
 378
 379 #        ifdef CALC_ENERGIES
 380     vcoul_S0 = qq_S0 * (rinv_ex_S0 + fma(rsq_S0, hrc_3_S, moh_rc_S));
 381     vcoul_S2 = qq_S2 * (rinv_ex_S2 + fma(rsq_S2, hrc_3_S, moh_rc_S));
 382 #        endif
 383 #    endif
 384
 385 #    ifdef CALC_COUL_EWALD
 386     /* We need to mask (or limit) rsq for the cut-off,
 387      * as large distances can cause an overflow in gmx_pmecorrF/V.
 388      */
 389     brsq_S0   = beta2_S * selectByMask(rsq_S0, wco_S0);
 390     brsq_S2   = beta2_S * selectByMask(rsq_S2, wco_S2);
 391     ewcorr_S0 = beta_S * pmeForceCorrection(brsq_S0);
 392     ewcorr_S2 = beta_S * pmeForceCorrection(brsq_S2);
 393     frcoul_S0 = qq_S0 * fma(ewcorr_S0, brsq_S0, rinv_ex_S0);
 394     frcoul_S2 = qq_S2 * fma(ewcorr_S2, brsq_S2, rinv_ex_S2);
 395
 396 #        ifdef CALC_ENERGIES
 397     vc_sub_S0 = beta_S * pmePotentialCorrection(brsq_S0);
 398     vc_sub_S2 = beta_S * pmePotentialCorrection(brsq_S2);
 399 #        endif
 400
 401 #    endif /* CALC_COUL_EWALD */
 402
 403 #    ifdef CALC_COUL_TAB
 404     /* Electrostatic interactions */
 405     r_S0 = rsq_S0 * rinv_S0;
 406     r_S2 = rsq_S2 * rinv_S2;
 407     /* Convert r to scaled table units */
 408     rs_S0 = r_S0 * invtsp_S;
 409     rs_S2 = r_S2 * invtsp_S;
 410     /* Truncate scaled r to an int */
 411     ti_S0 = cvttR2I(rs_S0);
 412     ti_S2 = cvttR2I(rs_S2);
 413
 414     rf_S0 = trunc(rs_S0);
 415     rf_S2 = trunc(rs_S2);
 416
 417     frac_S0 = rs_S0 - rf_S0;
 418     frac_S2 = rs_S2 - rf_S2;
 419
 420     /* Load and interpolate table forces and possibly energies.
 421      * Force and energy can be combined in one table, stride 4: FDV0
 422      * or in two separate tables with stride 1: F and V
 423      * Currently single precision uses FDV0, double F and V.
 424      */
 425 #        ifndef CALC_ENERGIES
 426 #            ifdef TAB_FDV0
 427     gatherLoadBySimdIntTranspose<4>(tab_coul_F, ti_S0, &ctab0_S0, &ctab1_S0);
 428     gatherLoadBySimdIntTranspose<4>(tab_coul_F, ti_S2, &ctab0_S2, &ctab1_S2);
 429 #            else
 430     gatherLoadUBySimdIntTranspose<1>(tab_coul_F, ti_S0, &ctab0_S0, &ctab1_S0);
 431     gatherLoadUBySimdIntTranspose<1>(tab_coul_F, ti_S2, &ctab0_S2, &ctab1_S2);
 432     ctab1_S0  = ctab1_S0 - ctab0_S0;
 433     ctab1_S2  = ctab1_S2 - ctab0_S2;
 434 #            endif
 435 #        else
 436 #            ifdef TAB_FDV0
 437     gatherLoadBySimdIntTranspose<4>(tab_coul_F, ti_S0, &ctab0_S0, &ctab1_S0, &ctabv_S0, &dum_S0);
 438     gatherLoadBySimdIntTranspose<4>(tab_coul_F, ti_S2, &ctab0_S2, &ctab1_S2, &ctabv_S2, &dum_S2);
 439 #            else
 440     gatherLoadUBySimdIntTranspose<1>(tab_coul_F, ti_S0, &ctab0_S0, &ctab1_S0);
 441     gatherLoadUBySimdIntTranspose<1>(tab_coul_V, ti_S0, &ctabv_S0, &dum_S0);
 442     gatherLoadUBySimdIntTranspose<1>(tab_coul_F, ti_S2, &ctab0_S2, &ctab1_S2);
 443     gatherLoadUBySimdIntTranspose<1>(tab_coul_V, ti_S2, &ctabv_S2, &dum_S2);
 444     ctab1_S0 = ctab1_S0 - ctab0_S0;
 445     ctab1_S2 = ctab1_S2 - ctab0_S2;
 446 #            endif
 447 #        endif
 448     fsub_S0   = fma(frac_S0, ctab1_S0, ctab0_S0);
 449     fsub_S2   = fma(frac_S2, ctab1_S2, ctab0_S2);
 450     frcoul_S0 = qq_S0 * fnma(fsub_S0, r_S0, rinv_ex_S0);
 451     frcoul_S2 = qq_S2 * fnma(fsub_S2, r_S2, rinv_ex_S2);
 452
 453 #        ifdef CALC_ENERGIES
 454     vc_sub_S0 = fma((mhalfsp_S * frac_S0), (ctab0_S0 + fsub_S0), ctabv_S0);
 455     vc_sub_S2 = fma((mhalfsp_S * frac_S2), (ctab0_S2 + fsub_S2), ctabv_S2);
 456 #        endif
 457 #    endif /* CALC_COUL_TAB */
 458
 459 #    if defined CALC_ENERGIES && (defined CALC_COUL_EWALD || defined CALC_COUL_TAB)
 460 #        ifndef NO_SHIFT_EWALD
 461     /* Add Ewald potential shift to vc_sub for convenience */
 462 #            ifdef CHECK_EXCLS
 463     vc_sub_S0 = vc_sub_S0 + selectByMask(sh_ewald_S, interact_S0);
 464     vc_sub_S2 = vc_sub_S2 + selectByMask(sh_ewald_S, interact_S2);
 465 #            else
 466     vc_sub_S0 = vc_sub_S0 + sh_ewald_S;
 467     vc_sub_S2 = vc_sub_S2 + sh_ewald_S;
 468 #            endif
 469 #        endif
 470
 471     vcoul_S0 = qq_S0 * (rinv_ex_S0 - vc_sub_S0);
 472     vcoul_S2 = qq_S2 * (rinv_ex_S2 - vc_sub_S2);
 473
 474 #    endif
 475
 476 #    ifdef CALC_ENERGIES
 477     /* Mask energy for cut-off and diagonal */
 478     vcoul_S0 = selectByMask(vcoul_S0, wco_S0);
 479     vcoul_S2 = selectByMask(vcoul_S2, wco_S2);
 480 #    endif
 481
 482 #endif /* CALC_COULOMB */
 483
 484 #ifdef CALC_LJ
 485     /* Lennard-Jones interaction */
 486
 487 #    ifdef VDW_CUTOFF_CHECK
 488     wco_vdw_S0 = (rsq_S0 < rcvdw2_S);
 489 #        ifndef HALF_LJ
 490     wco_vdw_S2 = (rsq_S2 < rcvdw2_S);
 491 #        endif
 492 #    else
 493     /* Same cut-off for Coulomb and VdW, reuse the registers */
 494 #        define wco_vdw_S0 wco_S0
 495 #        define wco_vdw_S2 wco_S2
 496 #    endif
 497
 498 #    ifndef LJ_COMB_LB
 499     rinvsix_S0 = rinvsq_S0 * rinvsq_S0 * rinvsq_S0;
 500 #        ifdef EXCL_FORCES
 501     rinvsix_S0 = selectByMask(rinvsix_S0, interact_S0);
 502 #        endif
 503 #        ifndef HALF_LJ
 504     rinvsix_S2 = rinvsq_S2 * rinvsq_S2 * rinvsq_S2;
 505 #            ifdef EXCL_FORCES
 506     rinvsix_S2 = selectByMask(rinvsix_S2, interact_S2);
 507 #            endif
 508 #        endif
 509
 510 #        if defined LJ_CUT || defined LJ_POT_SWITCH
 511     /* We have plain LJ or LJ-PME with simple C6/6 C12/12 coefficients */
 512     FrLJ6_S0 = c6_S0 * rinvsix_S0;
 513 #            ifndef HALF_LJ
 514     FrLJ6_S2 = c6_S2 * rinvsix_S2;
 515 #            endif
 516     FrLJ12_S0 = c12_S0 * rinvsix_S0 * rinvsix_S0;
 517 #            ifndef HALF_LJ
 518     FrLJ12_S2 = c12_S2 * rinvsix_S2 * rinvsix_S2;
 519 #            endif
 520 #        endif
 521
 522 #        if defined LJ_FORCE_SWITCH || defined LJ_POT_SWITCH
 523     /* We switch the LJ force */
 524     r_S0    = rsq_S0 * rinv_S0;
 525     rsw_S0  = max(r_S0 - rswitch_S, zero_S);
 526     rsw2_S0 = rsw_S0 * rsw_S0;
 527 #            ifndef HALF_LJ
 528     r_S2    = rsq_S2 * rinv_S2;
 529     rsw_S2  = max(r_S2 - rswitch_S, zero_S);
 530     rsw2_S2 = rsw_S2 * rsw_S2;
 531 #            endif
 532 #        endif
 533
 534 #        ifdef LJ_FORCE_SWITCH
 535
 536 #            define add_fr_switch(fr, rsw, rsw2_r, c2, c3) fma(fma(c3, rsw, c2), rsw2_r, fr)
 537     SimdReal rsw2_r_S0 = rsw2_S0 * r_S0;
 538     FrLJ6_S0           = c6_S0 * add_fr_switch(rinvsix_S0, rsw_S0, rsw2_r_S0, p6_fc2_S, p6_fc3_S);
 539 #            ifndef HALF_LJ
 540     SimdReal rsw2_r_S2 = rsw2_S2 * r_S2;
 541     FrLJ6_S2           = c6_S2 * add_fr_switch(rinvsix_S2, rsw_S2, rsw2_r_S2, p6_fc2_S, p6_fc3_S);
 542 #            endif
 543     FrLJ12_S0 = c12_S0 * add_fr_switch(rinvsix_S0 * rinvsix_S0, rsw_S0, rsw2_r_S0, p12_fc2_S, p12_fc3_S);
 544 #            ifndef HALF_LJ
 545     FrLJ12_S2 = c12_S2 * add_fr_switch(rinvsix_S2 * rinvsix_S2, rsw_S2, rsw2_r_S2, p12_fc2_S, p12_fc3_S);
 546 #            endif
 547 #            undef add_fr_switch
 548 #        endif /* LJ_FORCE_SWITCH */
 549
 550 #    endif /* not LJ_COMB_LB */
 551
 552 #    ifdef LJ_COMB_LB
 553     sir_S0 = sig_S0 * rinv_S0;
 554 #        ifndef HALF_LJ
 555     sir_S2 = sig_S2 * rinv_S2;
 556 #        endif
 557     sir2_S0 = sir_S0 * sir_S0;
 558 #        ifndef HALF_LJ
 559     sir2_S2 = sir_S2 * sir_S2;
 560 #        endif
 561     sir6_S0 = sir2_S0 * sir2_S0 * sir2_S0;
 562 #        ifdef EXCL_FORCES
 563     sir6_S0 = selectByMask(sir6_S0, interact_S0);
 564 #        endif
 565 #        ifndef HALF_LJ
 566     sir6_S2 = sir2_S2 * sir2_S2 * sir2_S2;
 567 #            ifdef EXCL_FORCES
 568     sir6_S2 = selectByMask(sir6_S2, interact_S2);
 569 #            endif
 570 #        endif
 571 #        ifdef VDW_CUTOFF_CHECK
 572     sir6_S0 = selectByMask(sir6_S0, wco_vdw_S0);
 573 #            ifndef HALF_LJ
 574     sir6_S2 = selectByMask(sir6_S2, wco_vdw_S2);
 575 #            endif
 576 #        endif
 577     FrLJ6_S0 = eps_S0 * sir6_S0;
 578 #        ifndef HALF_LJ
 579     FrLJ6_S2 = eps_S2 * sir6_S2;
 580 #        endif
 581     FrLJ12_S0 = FrLJ6_S0 * sir6_S0;
 582 #        ifndef HALF_LJ
 583     FrLJ12_S2 = FrLJ6_S2 * sir6_S2;
 584 #        endif
 585 #        if defined CALC_ENERGIES
 586     /* We need C6 and C12 to calculate the LJ potential shift */
 587     sig2_S0 = sig_S0 * sig_S0;
 588 #            ifndef HALF_LJ
 589     sig2_S2 = sig_S2 * sig_S2;
 590 #            endif
 591     sig6_S0 = sig2_S0 * sig2_S0 * sig2_S0;
 592 #            ifndef HALF_LJ
 593     sig6_S2 = sig2_S2 * sig2_S2 * sig2_S2;
 594 #            endif
 595     SimdReal c6_S0 = eps_S0 * sig6_S0;
 596 #            ifndef HALF_LJ
 597     SimdReal c6_S2 = eps_S2 * sig6_S2;
 598 #            endif
 599     SimdReal c12_S0 = c6_S0 * sig6_S0;
 600 #            ifndef HALF_LJ
 601     SimdReal c12_S2 = c6_S2 * sig6_S2;
 602 #            endif
 603 #        endif
 604 #    endif /* LJ_COMB_LB */
 605
 606     /* Determine the total scalar LJ force*r */
 607     frLJ_S0 = FrLJ12_S0 - FrLJ6_S0;
 608 #    ifndef HALF_LJ
 609     frLJ_S2 = FrLJ12_S2 - FrLJ6_S2;
 610 #    endif
 611
 612 #    if (defined LJ_CUT || defined LJ_FORCE_SWITCH) && defined CALC_ENERGIES
 613
 614 #        ifdef LJ_CUT
 615     /* Calculate the LJ energies, with constant potential shift */
 616     SimdReal VLJ6_S0 = sixth_S * fma(c6_S0, p6_cpot_S, FrLJ6_S0);
 617 #            ifndef HALF_LJ
 618     SimdReal VLJ6_S2 = sixth_S * fma(c6_S2, p6_cpot_S, FrLJ6_S2);
 619 #            endif
 620     SimdReal VLJ12_S0 = twelveth_S * fma(c12_S0, p12_cpot_S, FrLJ12_S0);
 621 #            ifndef HALF_LJ
 622     SimdReal VLJ12_S2 = twelveth_S * fma(c12_S2, p12_cpot_S, FrLJ12_S2);
 623 #            endif
 624 #        endif /* LJ_CUT */
 625
 626 #        ifdef LJ_FORCE_SWITCH
 627 #            define v_fswitch_pr(rsw, rsw2, c0, c3, c4) fma(fma(c4, rsw, c3), (rsw2) * (rsw), c0)
 628
 629     SimdReal VLJ6_S0 =
 630             c6_S0 * fma(sixth_S, rinvsix_S0, v_fswitch_pr(rsw_S0, rsw2_S0, p6_6cpot_S, p6_vc3_S, p6_vc4_S));
 631 #            ifndef HALF_LJ
 632     SimdReal VLJ6_S2 =
 633             c6_S2 * fma(sixth_S, rinvsix_S2, v_fswitch_pr(rsw_S2, rsw2_S2, p6_6cpot_S, p6_vc3_S, p6_vc4_S));
 634 #            endif
 635     SimdReal VLJ12_S0 = c12_S0
 636                         * fma(twelveth_S, rinvsix_S0 * rinvsix_S0,
 637                               v_fswitch_pr(rsw_S0, rsw2_S0, p12_12cpot_S, p12_vc3_S, p12_vc4_S));
 638 #            ifndef HALF_LJ
 639     SimdReal VLJ12_S2 = c12_S2
 640                         * fma(twelveth_S, rinvsix_S2 * rinvsix_S2,
 641                               v_fswitch_pr(rsw_S2, rsw2_S2, p12_12cpot_S, p12_vc3_S, p12_vc4_S));
 642 #            endif
 643 #            undef v_fswitch_pr
 644 #        endif /* LJ_FORCE_SWITCH */
 645
 646     /* Add up the repulsion and dispersion */
 647     SimdReal VLJ_S0 = VLJ12_S0 - VLJ6_S0;
 648 #        ifndef HALF_LJ
 649     SimdReal VLJ_S2 = VLJ12_S2 - VLJ6_S2;
 650 #        endif
 651
 652 #    endif /* (LJ_CUT || LJ_FORCE_SWITCH) && CALC_ENERGIES */
 653
 654 #    ifdef LJ_POT_SWITCH
 655     /* We always need the potential, since it is needed for the force */
 656     SimdReal VLJ_S0 = fnma(sixth_S, FrLJ6_S0, twelveth_S * FrLJ12_S0);
 657 #        ifndef HALF_LJ
 658     SimdReal VLJ_S2 = fnma(sixth_S, FrLJ6_S2, twelveth_S * FrLJ12_S2);
 659 #        endif
 660
 661     {
 662         SimdReal sw_S0, dsw_S0;
 663 #        ifndef HALF_LJ
 664         SimdReal sw_S2, dsw_S2;
 665 #        endif
 666
 667 #        define switch_pr(rsw, rsw2, c3, c4, c5) \
 668             fma(fma(fma(c5, rsw, c4), rsw, c3), (rsw2) * (rsw), one_S)
 669 #        define dswitch_pr(rsw, rsw2, c2, c3, c4) fma(fma(c4, rsw, c3), rsw, c2) * (rsw2)
 670
 671         sw_S0  = switch_pr(rsw_S0, rsw2_S0, swV3_S, swV4_S, swV5_S);
 672         dsw_S0 = dswitch_pr(rsw_S0, rsw2_S0, swF2_S, swF3_S, swF4_S);
 673 #        ifndef HALF_LJ
 674         sw_S2  = switch_pr(rsw_S2, rsw2_S2, swV3_S, swV4_S, swV5_S);
 675         dsw_S2 = dswitch_pr(rsw_S2, rsw2_S2, swF2_S, swF3_S, swF4_S);
 676 #        endif
 677         frLJ_S0 = fnma(dsw_S0 * VLJ_S0, r_S0, sw_S0 * frLJ_S0);
 678 #        ifndef HALF_LJ
 679         frLJ_S2 = fnma(dsw_S2 * VLJ_S2, r_S2, sw_S2 * frLJ_S2);
 680 #        endif
 681 #        ifdef CALC_ENERGIES
 682         VLJ_S0 = sw_S0 * VLJ_S0;
 683 #            ifndef HALF_LJ
 684         VLJ_S2 = sw_S2 * VLJ_S2;
 685 #            endif
 686 #        endif
 687
 688 #        undef switch_pr
 689 #        undef dswitch_pr
 690     }
 691 #    endif /* LJ_POT_SWITCH */
 692
 693 #    if defined CALC_ENERGIES && defined CHECK_EXCLS
 694     /* The potential shift should be removed for excluded pairs */
 695     VLJ_S0 = selectByMask(VLJ_S0, interact_S0);
 696 #        ifndef HALF_LJ
 697     VLJ_S2 = selectByMask(VLJ_S2, interact_S2);
 698 #        endif
 699 #    endif
 700
 701 #    ifdef LJ_EWALD_GEOM
 702     {
 703         SimdReal c6s_j_S;
 704         SimdReal c6grid_S0, rinvsix_nm_S0, cr2_S0, expmcr2_S0, poly_S0;
 705 #        ifndef HALF_LJ
 706         SimdReal c6grid_S2, rinvsix_nm_S2, cr2_S2, expmcr2_S2, poly_S2;
 707 #        endif
 708 #        ifdef CALC_ENERGIES
 709         SimdReal sh_mask_S0;
 710 #            ifndef HALF_LJ
 711         SimdReal sh_mask_S2;
 712 #            endif
 713 #        endif
 714
 715         /* Determine C6 for the grid using the geometric combination rule */
 716         c6s_j_S   = loadDuplicateHsimd(ljc + aj2);
 717         c6grid_S0 = c6s_S0 * c6s_j_S;
 718 #        ifndef HALF_LJ
 719         c6grid_S2 = c6s_S2 * c6s_j_S;
 720 #        endif
 721
 722 #        ifdef CHECK_EXCLS
 723         /* Recalculate rinvsix without exclusion mask (compiler might optimize) */
 724         rinvsix_nm_S0 = rinvsq_S0 * rinvsq_S0 * rinvsq_S0;
 725 #            ifndef HALF_LJ
 726         rinvsix_nm_S2 = rinvsq_S2 * rinvsq_S2 * rinvsq_S2;
 727 #            endif
 728 #        else
 729         /* We didn't use a mask, so we can copy */
 730         rinvsix_nm_S0 = rinvsix_S0;
 731 #            ifndef HALF_LJ
 732         rinvsix_nm_S2 = rinvsix_S2;
 733 #            endif
 734 #        endif
 735
 736         /* Mask for the cut-off to avoid overflow of cr2^2 */
 737         cr2_S0 = lje_c2_S * selectByMask(rsq_S0, wco_vdw_S0);
 738 #        ifndef HALF_LJ
 739         cr2_S2 = lje_c2_S * selectByMask(rsq_S2, wco_vdw_S2);
 740 #        endif
 741         // Unsafe version of our exp() should be fine, since these arguments should never
 742         // be smaller than -127 for any reasonable choice of cutoff or ewald coefficients.
 743         expmcr2_S0 = exp<MathOptimization::Unsafe>(-cr2_S0);
 744 #        ifndef HALF_LJ
 745         expmcr2_S2 = exp<MathOptimization::Unsafe>(-cr2_S2);
 746 #        endif
 747
 748         /* 1 + cr2 + 1/2*cr2^2 */
 749         poly_S0 = fma(fma(half_S, cr2_S0, one_S), cr2_S0, one_S);
 750 #        ifndef HALF_LJ
 751         poly_S2 = fma(fma(half_S, cr2_S2, one_S), cr2_S2, one_S);
 752 #        endif
 753
 754         /* We calculate LJ F*r = (6*C6)*(r^-6 - F_mesh/6), we use:
 755          * r^-6*cexp*(1 + cr2 + cr2^2/2 + cr2^3/6) = cexp*(r^-6*poly + c^6/6)
 756          */
 757         frLJ_S0 = fma(c6grid_S0,
 758                       fnma(expmcr2_S0, fma(rinvsix_nm_S0, poly_S0, lje_c6_6_S), rinvsix_nm_S0), frLJ_S0);
 759 #        ifndef HALF_LJ
 760         frLJ_S2 = fma(c6grid_S2,
 761                       fnma(expmcr2_S2, fma(rinvsix_nm_S2, poly_S2, lje_c6_6_S), rinvsix_nm_S2), frLJ_S2);
 762 #        endif
 763
 764 #        ifdef CALC_ENERGIES
 765 #            ifdef CHECK_EXCLS
 766         sh_mask_S0 = selectByMask(lje_vc_S, interact_S0);
 767 #                ifndef HALF_LJ
 768         sh_mask_S2 = selectByMask(lje_vc_S, interact_S2);
 769 #                endif
 770 #            else
 771         sh_mask_S0 = lje_vc_S;
 772 #                ifndef HALF_LJ
 773         sh_mask_S2 = lje_vc_S;
 774 #                endif
 775 #            endif
 776
 777         VLJ_S0 = fma(sixth_S * c6grid_S0,
 778                      fma(rinvsix_nm_S0, fnma(expmcr2_S0, poly_S0, one_S), sh_mask_S0), VLJ_S0);
 779 #            ifndef HALF_LJ
 780         VLJ_S2 = fma(sixth_S * c6grid_S2,
 781                      fma(rinvsix_nm_S2, fnma(expmcr2_S2, poly_S2, one_S), sh_mask_S2), VLJ_S2);
 782 #            endif
 783 #        endif /* CALC_ENERGIES */
 784     }
 785 #    endif /* LJ_EWALD_GEOM */
 786
 787 #    if defined VDW_CUTOFF_CHECK
 788     /* frLJ is multiplied later by rinvsq, which is masked for the Coulomb
 789      * cut-off, but if the VdW cut-off is shorter, we need to mask with that.
 790      */
 791     frLJ_S0 = selectByMask(frLJ_S0, wco_vdw_S0);
 792 #        ifndef HALF_LJ
 793     frLJ_S2 = selectByMask(frLJ_S2, wco_vdw_S2);
 794 #        endif
 795 #    endif
 796
 797 #    ifdef CALC_ENERGIES
 798     /* The potential shift should be removed for pairs beyond cut-off */
 799     VLJ_S0 = selectByMask(VLJ_S0, wco_vdw_S0);
 800 #        ifndef HALF_LJ
 801     VLJ_S2 = selectByMask(VLJ_S2, wco_vdw_S2);
 802 #        endif
 803 #    endif
 804
 805 #endif /* CALC_LJ */
 806
 807 #ifdef CALC_ENERGIES
 808 #    ifdef ENERGY_GROUPS
 809     /* Extract the group pair index per j pair.
 810      * Energy groups are stored per i-cluster, so things get
 811      * complicated when the i- and j-cluster size don't match.
 812      */
 813     {
 814         int egps_j;
 815 #        if UNROLLJ == 2
 816         egps_j    = nbatParams.energrp[cj >> 1];
 817         egp_jj[0] = ((egps_j >> ((cj & 1) * egps_jshift)) & egps_jmask) * egps_jstride;
 818 #        else
 819         /* We assume UNROLLI <= UNROLLJ */
 820         int jdi;
 821         for (jdi = 0; jdi < UNROLLJ / UNROLLI; jdi++)
 822         {
 823             int jj;
 824             egps_j = nbatParams.energrp[cj * (UNROLLJ / UNROLLI) + jdi];
 825             for (jj = 0; jj < (UNROLLI / 2); jj++)
 826             {
 827                 egp_jj[jdi * (UNROLLI / 2) + jj] =
 828                         ((egps_j >> (jj * egps_jshift)) & egps_jmask) * egps_jstride;
 829             }
 830         }
 831 #        endif
 832     }
 833 #    endif
 834
 835 #    ifdef CALC_COULOMB
 836 #        ifndef ENERGY_GROUPS
 837     vctot_S = vctot_S + vcoul_S0 + vcoul_S2;
 838 #        else
 839     add_ener_grp_halves(vcoul_S0, vctp[0], vctp[1], egp_jj);
 840     add_ener_grp_halves(vcoul_S2, vctp[2], vctp[3], egp_jj);
 841 #        endif
 842 #    endif
 843
 844 #    ifdef CALC_LJ
 845 #        ifndef ENERGY_GROUPS
 846     Vvdwtot_S = Vvdwtot_S + VLJ_S0
 847 #            ifndef HALF_LJ
 848                 + VLJ_S2
 849 #            endif
 850             ;
 851 #        else
 852     add_ener_grp_halves(VLJ_S0, vvdwtp[0], vvdwtp[1], egp_jj);
 853 #            ifndef HALF_LJ
 854     add_ener_grp_halves(VLJ_S2, vvdwtp[2], vvdwtp[3], egp_jj);
 855 #            endif
 856 #        endif
 857 #    endif /* CALC_LJ */
 858 #endif     /* CALC_ENERGIES */
 859
 860 #ifdef CALC_LJ
 861 #    ifdef CALC_COULOMB
 862     fscal_S0 = rinvsq_S0 * (frcoul_S0 + frLJ_S0);
 863 #    else
 864     fscal_S0 = rinvsq_S0 * frLJ_S0;
 865 #    endif
 866 #else
 867     fscal_S0 = rinvsq_S0 * frcoul_S0;
 868 #endif /* CALC_LJ */
 869 #if defined CALC_LJ && !defined HALF_LJ
 870 #    ifdef CALC_COULOMB
 871     fscal_S2 = rinvsq_S2 * (frcoul_S2 + frLJ_S2);
 872 #    else
 873     fscal_S2 = rinvsq_S2 * frLJ_S2;
 874 #    endif
 875 #else
 876     /* Atom 2 and 3 don't have LJ, so only add Coulomb forces */
 877     fscal_S2 = rinvsq_S2 * frcoul_S2;
 878 #endif
 879
 880     /* Calculate temporary vectorial force */
 881     tx_S0 = fscal_S0 * dx_S0;
 882     tx_S2 = fscal_S2 * dx_S2;
 883     ty_S0 = fscal_S0 * dy_S0;
 884     ty_S2 = fscal_S2 * dy_S2;
 885     tz_S0 = fscal_S0 * dz_S0;
 886     tz_S2 = fscal_S2 * dz_S2;
 887
 888     /* Increment i atom force */
 889     fix_S0 = fix_S0 + tx_S0;
 890     fix_S2 = fix_S2 + tx_S2;
 891     fiy_S0 = fiy_S0 + ty_S0;
 892     fiy_S2 = fiy_S2 + ty_S2;
 893     fiz_S0 = fiz_S0 + tz_S0;
 894     fiz_S2 = fiz_S2 + tz_S2;
 895
 896     /* Decrement j atom force */
 897     decrHsimd(f + ajx, tx_S0 + tx_S2);
 898     decrHsimd(f + ajy, ty_S0 + ty_S2);
 899     decrHsimd(f + ajz, tz_S0 + tz_S2);
 900 }
 901
 902 #undef rinv_ex_S0
 903 #undef rinv_ex_S2
 904
 905 #undef wco_vdw_S0
 906 #undef wco_vdw_S2
 907
 908 #undef EXCL_FORCES