1 /* ----------------------------------------------------------------------
2    SPARTA - Stochastic PArallel Rarefied-gas Time-accurate Analyzer
3    http://sparta.sandia.gov
4    Steve Plimpton, sjplimp@sandia.gov, Michael Gallis, magalli@sandia.gov
5    Sandia National Laboratories
6 
7    Copyright (2014) 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 SPARTA directory.
13 ------------------------------------------------------------------------- */
14 
15 /* ----------------------------------------------------------------------
16    Contributing author: Alan Stagg (Sandia)
17 ------------------------------------------------------------------------- */
18 
19 #include "string.h"
20 #include "compute_lambda_grid_kokkos.h"
21 #include "update.h"
22 #include "grid_kokkos.h"
23 #include "domain.h"
24 #include "collide.h"
25 #include "modify.h"
26 #include "fix.h"
27 #include "compute.h"
28 #include "math_const.h"
29 #include "memory.h"
30 #include "memory_kokkos.h"
31 #include "error.h"
32 #include "kokkos.h"
33 
34 using namespace SPARTA_NS;
35 using namespace MathConst;
36 
37 enum{NONE,COMPUTE,FIX};
38 enum{KNONE,KALL,KX,KY,KZ};
39 
40 #define INVOKED_PER_GRID 16
41 #define BIG 1.0e20
42 
43 /* ---------------------------------------------------------------------- */
44 
ComputeLambdaGridKokkos(SPARTA * sparta,int narg,char ** arg)45 ComputeLambdaGridKokkos::ComputeLambdaGridKokkos(SPARTA *sparta, int narg, char **arg) :
46   ComputeLambdaGrid(sparta, narg, arg)
47 {
48   kokkos_flag = 1;
49 }
50 
51 /* ---------------------------------------------------------------------- */
52 
~ComputeLambdaGridKokkos()53 ComputeLambdaGridKokkos::~ComputeLambdaGridKokkos()
54 {
55   if (copymode) return;
56 
57   if (kflag == KNONE)
58     memoryKK->destroy_kokkos(k_vector_grid,vector_grid);
59   else
60     memoryKK->destroy_kokkos(k_array_grid,array_grid);
61   vector_grid = NULL;
62   array_grid = NULL;
63 }
64 
65 /* ---------------------------------------------------------------------- */
66 
compute_per_grid()67 void ComputeLambdaGridKokkos::compute_per_grid()
68 {
69   if (sparta->kokkos->prewrap)
70     ComputeLambdaGrid::compute_per_grid();
71   else
72     compute_per_grid_kokkos();
73 }
74 
75 /* ---------------------------------------------------------------------- */
76 
compute_per_grid_kokkos()77 void ComputeLambdaGridKokkos::compute_per_grid_kokkos()
78 {
79   invoked_per_grid = update->ntimestep;
80 
81   if (nrhowhich == FIX && update->ntimestep % fnrho->per_grid_freq)
82     error->all(FLERR,"Compute lambda/grid fix not computed at compatible time");
83   if (tempwhich == FIX && update->ntimestep % ftemp->per_grid_freq)
84     error->all(FLERR,"Compute lambda/grid fix not computed at compatible time");
85 
86   // grab nrho and temp values from compute or fix
87   // invoke nrho and temp computes as needed
88   if (nrhowhich == COMPUTE) {
89     if (!cnrho->kokkos_flag)
90       error->all(FLERR,"Cannot (yet) use non-Kokkos computes with compute lambda/grid/kk");
91     KokkosBase* computeKKBase = dynamic_cast<KokkosBase*>(cnrho);
92     if (!(cnrho->invoked_flag & INVOKED_PER_GRID)) {
93       computeKKBase->compute_per_grid_kokkos();
94       cnrho->invoked_flag |= INVOKED_PER_GRID;
95     }
96 
97     if (cnrho->post_process_grid_flag)
98       computeKKBase->post_process_grid_kokkos(nrhoindex,1,DAT::t_float_2d_lr(),NULL,DAT::t_float_1d_strided());
99 
100     if (nrhoindex == 0 || cnrho->post_process_grid_flag)
101       Kokkos::deep_copy(d_nrho_vector, computeKKBase->d_vector);
102     else {
103       d_array = computeKKBase->d_array_grid;
104       copymode = 1;
105       Kokkos::parallel_for(Kokkos::RangePolicy<DeviceType, TagComputeLambdaGrid_LoadNrhoVecFromArray>(0,nglocal),*this);
106       DeviceType().fence();
107       copymode = 0;
108     }
109   } else if (nrhowhich == FIX){
110     if (!fnrho->kokkos_flag)
111       error->all(FLERR,"Cannot (yet) use non-Kokkos fixes with compute lambda/grid/kk");
112     KokkosBase* computeKKBase = dynamic_cast<KokkosBase*>(fnrho);
113     if (nrhoindex == 0)
114       d_nrho_vector = computeKKBase->d_vector;
115     else {
116       d_array = computeKKBase->d_array_grid;
117       copymode = 1;
118       Kokkos::parallel_for(Kokkos::RangePolicy<DeviceType, TagComputeLambdaGrid_LoadNrhoVecFromArray>(0,nglocal),*this);
119       DeviceType().fence();
120       copymode = 0;
121     }
122   }
123 
124   if (tempwhich == COMPUTE) {
125     if (!ctemp->kokkos_flag)
126       error->all(FLERR,"Cannot (yet) use non-Kokkos computes with compute lambda/grid/kk");
127     KokkosBase* computeKKBase = dynamic_cast<KokkosBase*>(ctemp);
128     if (!(ctemp->invoked_flag & INVOKED_PER_GRID)) {
129       computeKKBase->compute_per_grid_kokkos();
130       ctemp->invoked_flag |= INVOKED_PER_GRID;
131     }
132 
133     if (ctemp->post_process_grid_flag)
134       computeKKBase->post_process_grid_kokkos(tempindex,1,DAT::t_float_2d_lr(),NULL,DAT::t_float_1d_strided());
135 
136     if (tempindex == 0 || ctemp->post_process_grid_flag)
137       Kokkos::deep_copy(d_temp_vector, computeKKBase->d_vector);
138     else{
139       d_array = computeKKBase->d_array_grid;
140       copymode = 1;
141       Kokkos::parallel_for(Kokkos::RangePolicy<DeviceType, TagComputeLambdaGrid_LoadTempVecFromArray>(0,nglocal),*this);
142       DeviceType().fence();
143       copymode = 0;
144     }
145   } else if (tempwhich == FIX){
146     if (!ftemp->kokkos_flag)
147       error->all(FLERR,"Cannot (yet) use non-Kokkos fixes with compute lambda/grid/kk");
148     KokkosBase* computeKKBase = dynamic_cast<KokkosBase*>(ftemp);
149     if (tempindex == 0)
150       d_temp_vector = computeKKBase->d_vector;
151     else {
152       d_array = computeKKBase->d_array_grid;
153       copymode = 1;
154       Kokkos::parallel_for(Kokkos::RangePolicy<DeviceType, TagComputeLambdaGrid_LoadTempVecFromArray>(0,nglocal),*this);
155       DeviceType().fence();
156       copymode = 0;
157     }
158   }
159 
160   GridKokkos* grid_kk = ((GridKokkos*)grid);
161   d_cells = grid_kk->k_cells.d_view;
162   dimension = domain->dimension;
163   copymode = 1;
164   Kokkos::parallel_for(Kokkos::RangePolicy<DeviceType, TagComputeLambdaGrid_ComputePerGrid>(0,nglocal),*this);
165   DeviceType().fence();
166   copymode = 0;
167 
168   if (kflag == KNONE) {
169     k_vector_grid.modify_device();
170     k_vector_grid.sync_host();
171   } else {
172     k_array_grid.modify_device();
173     k_array_grid.sync_host();
174   }
175 }
176 
177 /* ---------------------------------------------------------------------- */
178 
179 KOKKOS_INLINE_FUNCTION
operator ()(TagComputeLambdaGrid_LoadNrhoVecFromArray,const int & i) const180 void ComputeLambdaGridKokkos::operator()(TagComputeLambdaGrid_LoadNrhoVecFromArray, const int &i) const {
181   d_nrho_vector(i) = d_array(i,nrhoindex-1);
182 }
183 
184 /* ---------------------------------------------------------------------- */
185 
186 KOKKOS_INLINE_FUNCTION
operator ()(TagComputeLambdaGrid_LoadTempVecFromArray,const int & i) const187 void ComputeLambdaGridKokkos::operator()(TagComputeLambdaGrid_LoadTempVecFromArray, const int &i) const {
188   d_temp_vector(i) = d_array(i,tempindex-1);
189 }
190 
191 /* ---------------------------------------------------------------------- */
192 
193 KOKKOS_INLINE_FUNCTION
operator ()(TagComputeLambdaGrid_ComputePerGrid,const int & i) const194 void ComputeLambdaGridKokkos::operator()(TagComputeLambdaGrid_ComputePerGrid, const int &i) const {
195   double lambda;
196 
197   if (d_nrho_vector(i) == 0.0) lambda = BIG;
198   else if (tempwhich == NONE || d_temp_vector(i) == 0.0)
199     lambda = 1.0 / (prefactor * d_nrho_vector(i));
200   else
201     lambda = 1.0 / (prefactor * d_nrho_vector(i) * pow(tref/d_temp_vector(i),omega-0.5));
202 
203   if (kflag == KNONE) d_vector(i) = lambda;
204   else d_array_grid(i,0) = lambda;
205 
206   // calculate per-cell Knudsen number
207   if (kflag == KNONE) return;
208 
209   if (kflag == KALL) {
210     double size;
211     size =  (d_cells(i).hi[0] - d_cells(i).lo[0]);
212     size += (d_cells(i).hi[1] - d_cells(i).lo[1]);
213     if (dimension == 2) size *= 0.5;
214     else {
215       size += (d_cells(i).hi[2] - d_cells(i).lo[2]);
216       size /= 3.0;
217     }
218     d_array_grid(i,1) = d_array_grid(i,0) / size;
219   } else if (kflag == KX) {
220     d_array_grid(i,1) = d_array_grid(i,0) / (d_cells(i).hi[0] - d_cells(i).lo[0]);
221   } else if (kflag == KY) {
222     d_array_grid(i,1) = d_array_grid(i,0) / (d_cells(i).hi[1] - d_cells(i).lo[1]);
223   } else if (kflag == KZ) {
224     d_array_grid(i,1) = d_array_grid(i,0) / (d_cells(i).hi[2] - d_cells(i).lo[2]);
225   }
226 }
227 
228 /* ---------------------------------------------------------------------- */
229 
reallocate()230 void ComputeLambdaGridKokkos::reallocate()
231 {
232   if (grid->nlocal == nglocal) return;
233   nglocal = grid->nlocal;
234   if (kflag == KNONE) {
235     memoryKK->destroy_kokkos(k_vector_grid,vector_grid);
236     memoryKK->create_kokkos(k_vector_grid,vector_grid,nglocal,"lambda/grid:vector_grid");
237     d_vector = k_vector_grid.d_view;
238   } else {
239     memoryKK->destroy_kokkos(k_array_grid,array_grid);
240     memoryKK->create_kokkos(k_array_grid,array_grid,nglocal,2,"lambda/grid:array_grid");
241     d_array_grid = k_array_grid.d_view;
242   }
243   d_nrho_vector = DAT::t_float_1d ("d_nrho_vector", nglocal);
244   d_temp_vector = DAT::t_float_1d ("d_temp_vector", nglocal);
245 }
246