1 // clang-format off
2 /* ----------------------------------------------------------------------
3    LAMMPS - Large-scale Atomic/Molecular Massively Parallel Simulator
4    https://www.lammps.org/, Sandia National Laboratories
5    Steve Plimpton, sjplimp@sandia.gov
6 
7    Copyright (2003) Sandia Corporation.  Under the terms of Contract
8    DE-AC04-94AL85000 with Sandia Corporation, the U.S. Government retains
9    certain rights in this software.  This software is distributed under
10    the GNU General Public License.
11 
12    See the README file in the top-level LAMMPS directory.
13 ------------------------------------------------------------------------- */
14 
15 /* ----------------------------------------------------------------------
16    Contributing author: Stan Moore (SNL)
17 ------------------------------------------------------------------------- */
18 
19 #include "angle_harmonic_kokkos.h"
20 
21 #include "atom_kokkos.h"
22 #include "atom_masks.h"
23 #include "comm.h"
24 #include "force.h"
25 #include "math_const.h"
26 #include "memory_kokkos.h"
27 #include "neighbor_kokkos.h"
28 
29 #include <cmath>
30 
31 using namespace LAMMPS_NS;
32 using namespace MathConst;
33 
34 #define SMALL 0.001
35 
36 /* ---------------------------------------------------------------------- */
37 
38 template<class DeviceType>
AngleHarmonicKokkos(LAMMPS * lmp)39 AngleHarmonicKokkos<DeviceType>::AngleHarmonicKokkos(LAMMPS *lmp) : AngleHarmonic(lmp)
40 {
41   atomKK = (AtomKokkos *) atom;
42   neighborKK = (NeighborKokkos *) neighbor;
43   execution_space = ExecutionSpaceFromDevice<DeviceType>::space;
44   datamask_read = X_MASK | F_MASK | ENERGY_MASK | VIRIAL_MASK;
45   datamask_modify = F_MASK | ENERGY_MASK | VIRIAL_MASK;
46 
47   centroidstressflag = CENTROID_NOTAVAIL;
48 }
49 
50 /* ---------------------------------------------------------------------- */
51 
52 template<class DeviceType>
~AngleHarmonicKokkos()53 AngleHarmonicKokkos<DeviceType>::~AngleHarmonicKokkos()
54 {
55   if (!copymode) {
56     memoryKK->destroy_kokkos(k_eatom,eatom);
57     memoryKK->destroy_kokkos(k_vatom,vatom);
58   }
59 }
60 
61 /* ---------------------------------------------------------------------- */
62 
63 template<class DeviceType>
compute(int eflag_in,int vflag_in)64 void AngleHarmonicKokkos<DeviceType>::compute(int eflag_in, int vflag_in)
65 {
66   eflag = eflag_in;
67   vflag = vflag_in;
68 
69   ev_init(eflag,vflag,0);
70 
71   // reallocate per-atom arrays if necessary
72 
73   if (eflag_atom) {
74     memoryKK->destroy_kokkos(k_eatom,eatom);
75     memoryKK->create_kokkos(k_eatom,eatom,maxeatom,"angle:eatom");
76     d_eatom = k_eatom.template view<DeviceType>();
77   }
78   if (vflag_atom) {
79     memoryKK->destroy_kokkos(k_vatom,vatom);
80     memoryKK->create_kokkos(k_vatom,vatom,maxvatom,"angle:vatom");
81     d_vatom = k_vatom.template view<DeviceType>();
82   }
83 
84   //atomKK->sync(execution_space,datamask_read);
85   k_k.template sync<DeviceType>();
86   k_theta0.template sync<DeviceType>();
87   //  if (eflag || vflag) atomKK->modified(execution_space,datamask_modify);
88   //  else atomKK->modified(execution_space,F_MASK);
89 
90   x = atomKK->k_x.template view<DeviceType>();
91   f = atomKK->k_f.template view<DeviceType>();
92   neighborKK->k_anglelist.template sync<DeviceType>();
93   anglelist = neighborKK->k_anglelist.template view<DeviceType>();
94   int nanglelist = neighborKK->nanglelist;
95   nlocal = atom->nlocal;
96   newton_bond = force->newton_bond;
97 
98   copymode = 1;
99 
100   // loop over neighbors of my atoms
101 
102   EV_FLOAT ev;
103 
104   if (evflag) {
105     if (newton_bond) {
106       Kokkos::parallel_reduce(Kokkos::RangePolicy<DeviceType, TagAngleHarmonicCompute<1,1> >(0,nanglelist),*this,ev);
107     } else {
108       Kokkos::parallel_reduce(Kokkos::RangePolicy<DeviceType, TagAngleHarmonicCompute<0,1> >(0,nanglelist),*this,ev);
109     }
110   } else {
111     if (newton_bond) {
112       Kokkos::parallel_for(Kokkos::RangePolicy<DeviceType, TagAngleHarmonicCompute<1,0> >(0,nanglelist),*this);
113     } else {
114       Kokkos::parallel_for(Kokkos::RangePolicy<DeviceType, TagAngleHarmonicCompute<0,0> >(0,nanglelist),*this);
115     }
116   }
117 
118   if (eflag_global) energy += ev.evdwl;
119   if (vflag_global) {
120     virial[0] += ev.v[0];
121     virial[1] += ev.v[1];
122     virial[2] += ev.v[2];
123     virial[3] += ev.v[3];
124     virial[4] += ev.v[4];
125     virial[5] += ev.v[5];
126   }
127 
128   if (eflag_atom) {
129     k_eatom.template modify<DeviceType>();
130     k_eatom.template sync<LMPHostType>();
131   }
132 
133   if (vflag_atom) {
134     k_vatom.template modify<DeviceType>();
135     k_vatom.template sync<LMPHostType>();
136   }
137 
138   copymode = 0;
139 }
140 
141 template<class DeviceType>
142 template<int NEWTON_BOND, int EVFLAG>
143 KOKKOS_INLINE_FUNCTION
operator ()(TagAngleHarmonicCompute<NEWTON_BOND,EVFLAG>,const int & n,EV_FLOAT & ev) const144 void AngleHarmonicKokkos<DeviceType>::operator()(TagAngleHarmonicCompute<NEWTON_BOND,EVFLAG>, const int &n, EV_FLOAT& ev) const {
145 
146   // The f array is atomic
147   Kokkos::View<F_FLOAT*[3], typename DAT::t_f_array::array_layout,typename KKDevice<DeviceType>::value,Kokkos::MemoryTraits<Kokkos::Atomic|Kokkos::Unmanaged> > a_f = f;
148 
149   const int i1 = anglelist(n,0);
150   const int i2 = anglelist(n,1);
151   const int i3 = anglelist(n,2);
152   const int type = anglelist(n,3);
153 
154   // 1st bond
155 
156   const F_FLOAT delx1 = x(i1,0) - x(i2,0);
157   const F_FLOAT dely1 = x(i1,1) - x(i2,1);
158   const F_FLOAT delz1 = x(i1,2) - x(i2,2);
159 
160   const F_FLOAT rsq1 = delx1*delx1 + dely1*dely1 + delz1*delz1;
161   const F_FLOAT r1 = sqrt(rsq1);
162 
163   // 2nd bond
164 
165   const F_FLOAT delx2 = x(i3,0) - x(i2,0);
166   const F_FLOAT dely2 = x(i3,1) - x(i2,1);
167   const F_FLOAT delz2 = x(i3,2) - x(i2,2);
168 
169   const F_FLOAT rsq2 = delx2*delx2 + dely2*dely2 + delz2*delz2;
170   const F_FLOAT r2 = sqrt(rsq2);
171 
172   // angle (cos and sin)
173 
174   F_FLOAT c = delx1*delx2 + dely1*dely2 + delz1*delz2;
175   c /= r1*r2;
176 
177   if (c > 1.0) c = 1.0;
178   if (c < -1.0) c = -1.0;
179 
180   F_FLOAT s = sqrt(1.0 - c*c);
181   if (s < SMALL) s = SMALL;
182   s = 1.0/s;
183 
184   // force & energy
185 
186   const F_FLOAT dtheta = acos(c) - d_theta0[type];
187   const F_FLOAT tk = d_k[type] * dtheta;
188 
189   F_FLOAT eangle = 0.0;
190   if (eflag) eangle = tk*dtheta;
191 
192   const F_FLOAT a = -2.0 * tk * s;
193   const F_FLOAT a11 = a*c / rsq1;
194   const F_FLOAT a12 = -a / (r1*r2);
195   const F_FLOAT a22 = a*c / rsq2;
196 
197   F_FLOAT f1[3],f3[3];
198   f1[0] = a11*delx1 + a12*delx2;
199   f1[1] = a11*dely1 + a12*dely2;
200   f1[2] = a11*delz1 + a12*delz2;
201   f3[0] = a22*delx2 + a12*delx1;
202   f3[1] = a22*dely2 + a12*dely1;
203   f3[2] = a22*delz2 + a12*delz1;
204 
205   // apply force to each of 3 atoms
206 
207   if (NEWTON_BOND || i1 < nlocal) {
208     a_f(i1,0) += f1[0];
209     a_f(i1,1) += f1[1];
210     a_f(i1,2) += f1[2];
211   }
212 
213   if (NEWTON_BOND || i2 < nlocal) {
214     a_f(i2,0) -= f1[0] + f3[0];
215     a_f(i2,1) -= f1[1] + f3[1];
216     a_f(i2,2) -= f1[2] + f3[2];
217   }
218 
219   if (NEWTON_BOND || i3 < nlocal) {
220     a_f(i3,0) += f3[0];
221     a_f(i3,1) += f3[1];
222     a_f(i3,2) += f3[2];
223   }
224 
225   if (EVFLAG) ev_tally(ev,i1,i2,i3,eangle,f1,f3,
226                        delx1,dely1,delz1,delx2,dely2,delz2);
227 }
228 
229 template<class DeviceType>
230 template<int NEWTON_BOND, int EVFLAG>
231 KOKKOS_INLINE_FUNCTION
operator ()(TagAngleHarmonicCompute<NEWTON_BOND,EVFLAG>,const int & n) const232 void AngleHarmonicKokkos<DeviceType>::operator()(TagAngleHarmonicCompute<NEWTON_BOND,EVFLAG>, const int &n) const {
233   EV_FLOAT ev;
234   this->template operator()<NEWTON_BOND,EVFLAG>(TagAngleHarmonicCompute<NEWTON_BOND,EVFLAG>(), n, ev);
235 }
236 
237 /* ---------------------------------------------------------------------- */
238 
239 template<class DeviceType>
allocate()240 void AngleHarmonicKokkos<DeviceType>::allocate()
241 {
242   AngleHarmonic::allocate();
243 
244   int n = atom->nangletypes;
245   k_k = typename ArrayTypes<DeviceType>::tdual_ffloat_1d("AngleHarmonic::k",n+1);
246   k_theta0 = typename ArrayTypes<DeviceType>::tdual_ffloat_1d("AngleHarmonic::theta0",n+1);
247 
248   d_k = k_k.template view<DeviceType>();
249   d_theta0 = k_theta0.template view<DeviceType>();
250 }
251 
252 /* ----------------------------------------------------------------------
253    set coeffs for one or more types
254 ------------------------------------------------------------------------- */
255 
256 template<class DeviceType>
coeff(int narg,char ** arg)257 void AngleHarmonicKokkos<DeviceType>::coeff(int narg, char **arg)
258 {
259   AngleHarmonic::coeff(narg, arg);
260 
261   int n = atom->nangletypes;
262   for (int i = 1; i <= n; i++) {
263     k_k.h_view[i] = k[i];
264     k_theta0.h_view[i] = theta0[i];
265   }
266 
267   k_k.template modify<LMPHostType>();
268   k_theta0.template modify<LMPHostType>();
269 }
270 
271 /* ----------------------------------------------------------------------
272    proc 0 reads coeffs from restart file, bcasts them
273 ------------------------------------------------------------------------- */
274 
275 template<class DeviceType>
read_restart(FILE * fp)276 void AngleHarmonicKokkos<DeviceType>::read_restart(FILE *fp)
277 {
278   AngleHarmonic::read_restart(fp);
279 
280   int n = atom->nangletypes;
281   for (int i = 1; i <= n; i++) {
282     k_k.h_view[i] = k[i];
283     k_theta0.h_view[i] = theta0[i];
284   }
285 
286   k_k.template modify<LMPHostType>();
287   k_theta0.template modify<LMPHostType>();
288 }
289 
290 /* ----------------------------------------------------------------------
291    tally energy and virial into global and per-atom accumulators
292    virial = r1F1 + r2F2 + r3F3 = (r1-r2) F1 + (r3-r2) F3 = del1*f1 + del2*f3
293 ------------------------------------------------------------------------- */
294 
295 template<class DeviceType>
296 //template<int NEWTON_BOND>
297 KOKKOS_INLINE_FUNCTION
ev_tally(EV_FLOAT & ev,const int i,const int j,const int k,F_FLOAT & eangle,F_FLOAT * f1,F_FLOAT * f3,const F_FLOAT & delx1,const F_FLOAT & dely1,const F_FLOAT & delz1,const F_FLOAT & delx2,const F_FLOAT & dely2,const F_FLOAT & delz2) const298 void AngleHarmonicKokkos<DeviceType>::ev_tally(EV_FLOAT &ev, const int i, const int j, const int k,
299                      F_FLOAT &eangle, F_FLOAT *f1, F_FLOAT *f3,
300                      const F_FLOAT &delx1, const F_FLOAT &dely1, const F_FLOAT &delz1,
301                      const F_FLOAT &delx2, const F_FLOAT &dely2, const F_FLOAT &delz2) const
302 {
303   E_FLOAT eanglethird;
304   F_FLOAT v[6];
305 
306   // The eatom and vatom arrays are atomic
307   Kokkos::View<E_FLOAT*, typename DAT::t_efloat_1d::array_layout,typename KKDevice<DeviceType>::value,Kokkos::MemoryTraits<Kokkos::Atomic|Kokkos::Unmanaged> > v_eatom = k_eatom.template view<DeviceType>();
308   Kokkos::View<F_FLOAT*[6], typename DAT::t_virial_array::array_layout,typename KKDevice<DeviceType>::value,Kokkos::MemoryTraits<Kokkos::Atomic|Kokkos::Unmanaged> > v_vatom = k_vatom.template view<DeviceType>();
309 
310   if (eflag_either) {
311     if (eflag_global) {
312       if (newton_bond) ev.evdwl += eangle;
313       else {
314         eanglethird = THIRD*eangle;
315 
316         if (i < nlocal) ev.evdwl += eanglethird;
317         if (j < nlocal) ev.evdwl += eanglethird;
318         if (k < nlocal) ev.evdwl += eanglethird;
319       }
320     }
321     if (eflag_atom) {
322       eanglethird = THIRD*eangle;
323 
324       if (newton_bond || i < nlocal) v_eatom[i] += eanglethird;
325       if (newton_bond || j < nlocal) v_eatom[j] += eanglethird;
326       if (newton_bond || k < nlocal) v_eatom[k] += eanglethird;
327     }
328   }
329 
330   if (vflag_either) {
331     v[0] = delx1*f1[0] + delx2*f3[0];
332     v[1] = dely1*f1[1] + dely2*f3[1];
333     v[2] = delz1*f1[2] + delz2*f3[2];
334     v[3] = delx1*f1[1] + delx2*f3[1];
335     v[4] = delx1*f1[2] + delx2*f3[2];
336     v[5] = dely1*f1[2] + dely2*f3[2];
337 
338     if (vflag_global) {
339       if (newton_bond) {
340         ev.v[0] += v[0];
341         ev.v[1] += v[1];
342         ev.v[2] += v[2];
343         ev.v[3] += v[3];
344         ev.v[4] += v[4];
345         ev.v[5] += v[5];
346       } else {
347         if (i < nlocal) {
348           ev.v[0] += THIRD*v[0];
349           ev.v[1] += THIRD*v[1];
350           ev.v[2] += THIRD*v[2];
351           ev.v[3] += THIRD*v[3];
352           ev.v[4] += THIRD*v[4];
353           ev.v[5] += THIRD*v[5];
354         }
355         if (j < nlocal) {
356           ev.v[0] += THIRD*v[0];
357           ev.v[1] += THIRD*v[1];
358           ev.v[2] += THIRD*v[2];
359           ev.v[3] += THIRD*v[3];
360           ev.v[4] += THIRD*v[4];
361           ev.v[5] += THIRD*v[5];
362         }
363         if (k < nlocal) {
364           ev.v[0] += THIRD*v[0];
365 
366           ev.v[1] += THIRD*v[1];
367           ev.v[2] += THIRD*v[2];
368           ev.v[3] += THIRD*v[3];
369           ev.v[4] += THIRD*v[4];
370           ev.v[5] += THIRD*v[5];
371         }
372       }
373     }
374 
375     if (vflag_atom) {
376       if (newton_bond || i < nlocal) {
377         v_vatom(i,0) += THIRD*v[0];
378         v_vatom(i,1) += THIRD*v[1];
379         v_vatom(i,2) += THIRD*v[2];
380         v_vatom(i,3) += THIRD*v[3];
381         v_vatom(i,4) += THIRD*v[4];
382         v_vatom(i,5) += THIRD*v[5];
383       }
384       if (newton_bond || j < nlocal) {
385         v_vatom(j,0) += THIRD*v[0];
386         v_vatom(j,1) += THIRD*v[1];
387         v_vatom(j,2) += THIRD*v[2];
388         v_vatom(j,3) += THIRD*v[3];
389         v_vatom(j,4) += THIRD*v[4];
390         v_vatom(j,5) += THIRD*v[5];
391       }
392       if (newton_bond || k < nlocal) {
393         v_vatom(k,0) += THIRD*v[0];
394         v_vatom(k,1) += THIRD*v[1];
395         v_vatom(k,2) += THIRD*v[2];
396         v_vatom(k,3) += THIRD*v[3];
397         v_vatom(k,4) += THIRD*v[4];
398         v_vatom(k,5) += THIRD*v[5];
399 
400       }
401     }
402   }
403 }
404 
405 /* ---------------------------------------------------------------------- */
406 
407 namespace LAMMPS_NS {
408 template class AngleHarmonicKokkos<LMPDeviceType>;
409 #ifdef LMP_KOKKOS_GPU
410 template class AngleHarmonicKokkos<LMPHostType>;
411 #endif
412 }
413 
414