1 // clang-format off
2 /* -*- c++ -*- ----------------------------------------------------------
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 authors: Stefan Paquay & Matthew Peterson (Brandeis University)
17 ------------------------------------------------------------------------- */
18
19 #include "fix_enforce2d_kokkos.h"
20 #include "atom_masks.h"
21 #include "atom_kokkos.h"
22 #include "comm.h"
23 #include "error.h"
24
25
26 using namespace LAMMPS_NS;
27
28
29 template <class DeviceType>
FixEnforce2DKokkos(LAMMPS * lmp,int narg,char ** arg)30 FixEnforce2DKokkos<DeviceType>::FixEnforce2DKokkos(LAMMPS *lmp, int narg, char **arg) :
31 FixEnforce2D(lmp, narg, arg)
32 {
33 kokkosable = 1;
34 atomKK = (AtomKokkos *) atom;
35 execution_space = ExecutionSpaceFromDevice<DeviceType>::space;
36
37 datamask_read = V_MASK | F_MASK | OMEGA_MASK | MASK_MASK
38 | TORQUE_MASK | ANGMOM_MASK;
39
40 datamask_modify = V_MASK | F_MASK | OMEGA_MASK
41 | TORQUE_MASK | ANGMOM_MASK;
42 }
43
44
45 template <class DeviceType>
setup(int vflag)46 void FixEnforce2DKokkos<DeviceType>::setup(int vflag)
47 {
48 post_force(vflag);
49 }
50
51
52 template <class DeviceType>
post_force(int)53 void FixEnforce2DKokkos<DeviceType>::post_force(int /*vflag*/)
54 {
55 atomKK->sync(execution_space,datamask_read);
56
57 v = atomKK->k_v.view<DeviceType>();
58 f = atomKK->k_f.view<DeviceType>();
59
60 if (atomKK->omega_flag)
61 omega = atomKK->k_omega.view<DeviceType>();
62
63 if (atomKK->angmom_flag)
64 angmom = atomKK->k_angmom.view<DeviceType>();
65
66 if (atomKK->torque_flag)
67 torque = atomKK->k_torque.view<DeviceType>();
68
69
70 mask = atomKK->k_mask.view<DeviceType>();
71
72 int nlocal = atomKK->nlocal;
73 if (igroup == atomKK->firstgroup) nlocal = atomKK->nfirst;
74
75 int flag_mask = 0;
76 if (atomKK->omega_flag) flag_mask |= 1;
77 if (atomKK->angmom_flag) flag_mask |= 2;
78 if (atomKK->torque_flag) flag_mask |= 4;
79
80 copymode = 1;
81 switch (flag_mask) {
82 case 0:{
83 FixEnforce2DKokkosPostForceFunctor<DeviceType,0,0,0> functor(this);
84 Kokkos::parallel_for(nlocal,functor);
85 break;
86 }
87 case 1:{
88 FixEnforce2DKokkosPostForceFunctor<DeviceType,1,0,0> functor(this);
89 Kokkos::parallel_for(nlocal,functor);
90 break;
91 }
92 case 2:{
93 FixEnforce2DKokkosPostForceFunctor<DeviceType,0,1,0> functor(this);
94 Kokkos::parallel_for(nlocal,functor);
95 break;
96 }
97 case 3:{
98 FixEnforce2DKokkosPostForceFunctor<DeviceType,1,1,0> functor(this);
99 Kokkos::parallel_for(nlocal,functor);
100 break;
101 }
102 case 4:{
103 FixEnforce2DKokkosPostForceFunctor<DeviceType,0,0,1> functor(this);
104 Kokkos::parallel_for(nlocal,functor);
105 break;
106 }
107 case 5:{
108 FixEnforce2DKokkosPostForceFunctor<DeviceType,1,0,1> functor(this);
109 Kokkos::parallel_for(nlocal,functor);
110 break;
111 }
112 case 6:{
113 FixEnforce2DKokkosPostForceFunctor<DeviceType,0,1,1> functor(this);
114 Kokkos::parallel_for(nlocal,functor);
115 break;
116 }
117 case 7:{
118 FixEnforce2DKokkosPostForceFunctor<DeviceType,1,1,1> functor(this);
119 Kokkos::parallel_for(nlocal,functor);
120 break;
121 }
122 default:
123 error->all(FLERR, "Flag in fix_enforce2d_kokkos outside of what it should be");
124 }
125 copymode = 0;
126
127 atomKK->modified(execution_space,datamask_modify);
128
129 for (int m = 0; m < nfixlist; m++) {
130 atomKK->sync(flist[m]->execution_space,flist[m]->datamask_read);
131 flist[m]->enforce2d();
132 atomKK->modified(flist[m]->execution_space,flist[m]->datamask_modify);
133 }
134
135 }
136
137
138 template <class DeviceType>
139 template <int omega_flag, int angmom_flag, int torque_flag>
140 KOKKOS_INLINE_FUNCTION
post_force_item(int i) const141 void FixEnforce2DKokkos<DeviceType>::post_force_item( int i ) const
142 {
143 if (mask[i] & groupbit) {
144 v(i,2) = 0.0;
145 f(i,2) = 0.0;
146
147 if (omega_flag) {
148 omega(i,0) = 0.0;
149 omega(i,1) = 0.0;
150 }
151
152 if (angmom_flag) {
153 angmom(i,0) = 0.0;
154 angmom(i,1) = 0.0;
155 }
156
157 if (torque_flag) {
158 torque(i,0) = 0.0;
159 torque(i,1) = 0.0;
160 }
161 }
162 }
163
164
165 namespace LAMMPS_NS {
166 template class FixEnforce2DKokkos<LMPDeviceType>;
167 #ifdef LMP_KOKKOS_GPU
168 template class FixEnforce2DKokkos<LMPHostType>;
169 #endif
170 }
171