XGCa
 All Classes Namespaces Files Functions Variables Typedefs Enumerations Enumerator Friends Macros Pages
plasma.hpp
Go to the documentation of this file.
1 #ifndef PLASMA_HPP
2 #define PLASMA_HPP
3 #include "my_subview.hpp"
4 #include "species.hpp"
5 #include "vgrid_distribution.hpp"
6 
7 extern "C" void f0_init_decomposed_arrays(double* f0_T_ev_cpp, double* f0_inv_grid_vol_cpp, double* f0_grid_vol_cpp,
8  double* f0_grid_vol_vonly_cpp, double* f0_n_Ta_cpp, double* f0_den_cpp,
9  double* f0_flow_cpp, double* f0_B_B0_cpp);
10 extern "C" void f0_init_global_arrays(double* f0_den_global, double* f0_temp_global);
11 extern "C" void f0_set_ptrs(int nnode, double* f0_delta_n_cpp, double* f0_delta_u_cpp, double* f0_delta_T_cpp);
12 extern "C" void set_f0_f0g_ptr(int f0_inode1, int f0_inode2, double* f0_f0g_loc);
13 extern "C" void f0_initialize_wrap();
14 
15 class Plasma{
16 
17  // Variables for managing the device particle memory allocation
21 
22  public:
23 
24  std::vector<Species<DeviceType>> all_species;
25 
26  // Poloidally decomposed f0 values
27  // These will probably become a member of species later, but the fortran code expects
28  // the species to be an array index
30 
31  private:
32  /* Contains f0 values that are poloidally decomposed but don't need to be transferred between ranks during load rebalancing
33  * */
35  Kokkos::View<double**,Kokkos::LayoutRight, HostType> f0_T_ev;
36  Kokkos::View<double**,Kokkos::LayoutRight, HostType> f0_inv_grid_vol;
37  Kokkos::View<double**,Kokkos::LayoutRight, HostType> f0_grid_vol;
38  Kokkos::View<double**,Kokkos::LayoutRight, HostType> f0_grid_vol_vonly;
39  Kokkos::View<double**,Kokkos::LayoutRight, HostType> f0_n_Ta;
40  Kokkos::View<double**,Kokkos::LayoutRight, HostType> f0_den;
41  Kokkos::View<double**,Kokkos::LayoutRight, HostType> f0_flow;
42 
43  // f0_B_B0 is not plasma-related, should be moved elsewhere (e.g. Grid) - ALS
44  Kokkos::View<double*,Kokkos::LayoutRight, HostType> f0_B_B0;
45 
47 
49  : f0_T_ev("f0_T_ev", nsp, pol_decomp.nnodes),
50  f0_inv_grid_vol("f0_inv_grid_vol", nsp, pol_decomp.nnodes),
51  f0_grid_vol("f0_grid_vol", nsp, pol_decomp.nnodes),
52  f0_grid_vol_vonly("f0_grid_vol_vonly", nsp, pol_decomp.nnodes),
53  f0_n_Ta("f0_n_Ta", nsp, pol_decomp.nnodes),
54  f0_den("f0_den", nsp, pol_decomp.nnodes),
55  f0_flow("f0_flow", nsp, pol_decomp.nnodes),
56  f0_B_B0("f0_B_B0", pol_decomp.nnodes)
57  {
58  // Fortran arrays are set to point to these Views
59  // The calculations are currently done in Fortran too
61  f0_inv_grid_vol.data(),
62  f0_grid_vol.data(),
63  f0_grid_vol_vonly.data(),
64  f0_n_Ta.data(),
65  f0_den.data(),
66  f0_flow.data(),
67  f0_B_B0.data());
68  }
69  };
70 
71  public:
72 
74 
75  // Not poloidally decomposed
76  Kokkos::View<double**,Kokkos::LayoutRight, HostType> f0_delta_n;
77  Kokkos::View<double**,Kokkos::LayoutRight, HostType> f0_delta_u;
78  Kokkos::View<double**,Kokkos::LayoutRight, HostType> f0_delta_T;
79 
80  View<double**,CLayout, HostType> f0_den_global;
81  View<double**,CLayout, HostType> f0_temp_global;
82 
83  int nspecies;
85 
86  // Constructors
88  : particles_d_has_owner(false),
90  nspecies(0), // Initialize with 0 species
91  n_nonadiabatic_species(0) // Initialize with 0 species
92  {}
93 
94  template<class Device> // Host grid can be used as input
95  Plasma(NLReader::NamelistReader& nlr, bool use_f0_grid, const Grid<Device>& grid, const MagneticField<DeviceType>& magnetic_field, const DomainDecomposition<DeviceType>& pol_decomp, const VelocityGrid& vgrid)
96  : particles_d_has_owner(false)
97  {
98  if(nlr.namelist_present("performance_param")){
99  nlr.use_namelist("performance_param");
100  species_share_particles_d_ownership = !(nlr.get<bool>("particles_resident_on_device", default_residence_option()));
101  }else{
103  }
104 #ifndef USE_GPU
106 #endif
107 
108  nlr.use_namelist("ptl_param");
109  nspecies = nlr.get<int>("ptl_nsp",1);
110  nspecies += 1; // Add electrons
111 
112  // Add each species
113  int i_nonadiabatic_species = 0;
114  for(int isp = 0; isp<nspecies; isp++){
115  all_species.push_back(Species<DeviceType>(nlr, grid, magnetic_field, pol_decomp, isp, i_nonadiabatic_species));
116  if(!all_species[isp].is_adiabatic){
117  i_nonadiabatic_species++;
118  }
119  }
120 
121  // Save count of non-adiabatic species
122  n_nonadiabatic_species = i_nonadiabatic_species;
123 
124  /*** f0 arrays ***/
125  int grid_nnode = pol_decomp.gvid0_pid_h(pol_decomp.gvid0_pid_h.size()-1) - 1;
126  init_global_f0_arrays(grid_nnode);
127 
128  if(use_f0_grid){
129  resize_f0_f0g(pol_decomp, vgrid);
131 
132  // Reallocate and recalculate arrays that are decomposed but recalculated rather than
133  // transferred by the load balance (f0_T_ev etc)
135  }
136  }
137 
141  };
142 
146  };
147 
148  // Loop over all species
149  template<typename F>
150  inline void for_all_species(F func, DevicePtlOpt device_ptl_opt = UseDevicePtl){
151  for(int isp = 0; isp<all_species.size(); isp++){
152  manage_particle_ownership(isp, device_ptl_opt);
153  func(all_species[isp]);
154  }
155  }
156 
157  // Loop over all non-adiabatic species
158  template<typename F>
159  inline void for_all_nonadiabatic_species(F func, DevicePtlOpt device_ptl_opt = UseDevicePtl){
160  for(int isp = 0; isp<all_species.size(); isp++){
161  if(isp>10) { // ISP ERROR
162  printf("ISP ERROR in for_all_nonadiabatic_sepcies: isp=%d",isp);
163  fflush(stdout);
164  }
165  if(!all_species[isp].is_adiabatic){
166  manage_particle_ownership(isp, device_ptl_opt);
167  func(all_species[isp]);
168  }
169  }
170  }
171 
172  // Loop over electrons
173  template<typename F>
174  inline void for_electrons(F func, DevicePtlOpt device_ptl_opt = UseDevicePtl){
175  for(int isp = 0; isp<all_species.size(); isp++){
176  if(all_species[isp].is_electron){
177  manage_particle_ownership(isp, device_ptl_opt);
178  func(all_species[isp]);
179  }
180  }
181  }
182 
183  // Loop over ions
184  template<typename F>
185  inline void for_all_ions(F func, DevicePtlOpt device_ptl_opt = UseDevicePtl){
186  for(int isp = 0; isp<all_species.size(); isp++){
187  if(!all_species[isp].is_electron){
188  manage_particle_ownership(isp, device_ptl_opt);
189  func(all_species[isp]);
190  }
191  }
192  }
193 
194  // Loop over electrons or ions as specified
195  template<typename F>
196  inline void for_all(ParticleType particle_type, F func, DevicePtlOpt device_ptl_opt = UseDevicePtl){
197  for(int isp = 0; isp<all_species.size(); isp++){
198  if((particle_type==Electrons && all_species[isp].is_electron) ||
199  (particle_type==Ions && !all_species[isp].is_electron)){
200  manage_particle_ownership(isp, device_ptl_opt);
201  func(all_species[isp]);
202  }
203  }
204  }
205 
206  // Operate on one species
207  template<typename F>
208  inline void for_one_species(int isp, F func, DevicePtlOpt device_ptl_opt = UseDevicePtl){
209  manage_particle_ownership(isp, device_ptl_opt);
210  func(all_species[isp]);
211  }
212 
213  int largest_n_ptl(bool check_backup){
214  int max_n_ptl = 0;
215 
216  // Loop over all species
218  if(species.is_electron && check_backup){
219  // Check on backup particles instead, if that's what's getting used
220  max_n_ptl = std::max(max_n_ptl,species.n_backup_particles);
221  } else {
222  // For now, need to access fortran object for n_ptl
223  max_n_ptl = std::max(max_n_ptl,species.n_ptl);
224  }
225  }, NoDevicePtl);
226 
227  return max_n_ptl;
228  }
229 
230  /* Deallocates the device particles and resets ownership tracking variables
231  * */
233  for_all_nonadiabatic_species([&](Species<DeviceType>& species){
234  if(species.owns_particles_d){
235  species.particles_d = Cabana::AoSoA<ParticleDataTypes,DeviceType,VEC_LEN>();
236  species.owns_particles_d = false;
237  }
238  });
239  particles_d_has_owner = false;
240  }
241 
242  /* Reallocate and recalculate arrays that are decomposed but recalculated rather than
243  * transferred by the load balance (f0_T_ev etc) */
245  decomposed_recalculable_f0_arrays = DecomposedRecalculableF0Arrays(n_nonadiabatic_species, pol_decomp);
246 
247  // Set all_species unmanaged views
248  int f0_species_cnt=0;
249  for_all_nonadiabatic_species([&](Species<DeviceType>& species){
250  set_unmanaged_f0_species_view(decomposed_recalculable_f0_arrays.f0_T_ev, f0_species_cnt, species.f0.temp_ev_h);
251  set_unmanaged_f0_species_view(decomposed_recalculable_f0_arrays.f0_flow, f0_species_cnt, species.f0.flow_h);
252  set_unmanaged_f0_species_view(decomposed_recalculable_f0_arrays.f0_grid_vol, f0_species_cnt, species.f0.grid_vol_h);
253  set_unmanaged_f0_species_view(decomposed_recalculable_f0_arrays.f0_grid_vol_vonly, f0_species_cnt, species.f0.grid_vol_vonly_h);
254  set_unmanaged_f0_species_view(decomposed_recalculable_f0_arrays.f0_inv_grid_vol, f0_species_cnt, species.f0.inv_grid_vol_h);
255  set_unmanaged_f0_species_view(decomposed_recalculable_f0_arrays.f0_n_Ta, f0_species_cnt, species.f0.n_Ta_h);
256  set_unmanaged_f0_species_view(decomposed_recalculable_f0_arrays.f0_den, f0_species_cnt, species.f0.den_h);
257  f0_species_cnt++;
258  }, NoDevicePtl);
259  }
260 
261  private:
262 
263  /* Allocate and calculate global f0 arrays. Should be part of a constructor */
264  void init_global_f0_arrays(int grid_nnode){
265  // Allocate for all species (even though most simulations don't need it for adiabatic species)
266  f0_delta_n = View<double**,CLayout, HostType>("f0_delta_n", nspecies, grid_nnode);
267  f0_delta_u = View<double**,CLayout, HostType>("f0_delta_u", nspecies, grid_nnode);
268  f0_delta_T = View<double**,CLayout, HostType>("f0_delta_T", nspecies, grid_nnode);
269  // Set all_species unmanaged views
270  for_all_species([&](Species<DeviceType>& species){
271  set_unmanaged_f0_species_view(f0_delta_n, species.idx, species.f0.delta_n_h);
272  set_unmanaged_f0_species_view(f0_delta_u, species.idx, species.f0.delta_u_h);
273  set_unmanaged_f0_species_view(f0_delta_T, species.idx, species.f0.delta_T_h);
274  }, NoDevicePtl);
275 
276  // Set fortran pointers
277  f0_set_ptrs(grid_nnode, f0_delta_n.data(), f0_delta_u.data(), f0_delta_T.data());
278 
279  // These are constant for the simulation:
280  f0_den_global = View<double**,CLayout, HostType>("f0_den_global", n_nonadiabatic_species, grid_nnode);
281  f0_temp_global = View<double**,CLayout, HostType>("f0_temp_global", n_nonadiabatic_species, grid_nnode);
282 
283  // Fortran arrays are set to point to these Views
284  // The calculations are currently done in Fortran too
285  f0_init_global_arrays(f0_den_global.data(), f0_temp_global.data());
286 
287  for_all_nonadiabatic_species([&](Species<DeviceType>& species){
288  set_unmanaged_f0_species_view(f0_den_global, species.nonadiabatic_idx, species.f0.den_global_h);
289  set_unmanaged_f0_species_view(f0_temp_global, species.nonadiabatic_idx, species.f0.temp_global_h);
290  }, NoDevicePtl);
291  }
292 
293  // Creates an unmanaged view pointing to the subview specified by isp
294  template<typename T_in, typename T_out>
295  void set_unmanaged_f0_species_view(const T_in& view_in, int isp, T_out& view_out){
296  auto view_in_subview = my_subview(view_in, isp);
297  view_out = T_out(view_in_subview.data(), view_in_subview.layout());
298  }
299 
300  public:
301 
302  void resize_f0_f0g(const DomainDecomposition<DeviceType>& pol_decomp, const VelocityGrid& vgrid){
303  // Create new object rather than resizing since we don't need to preserve data
304  f0_f0g = VGridDistribution<HostType>(n_nonadiabatic_species, vgrid, pol_decomp);
305 
306  int f0_inode1 = pol_decomp.first_node; // 1-indexed
307  int f0_inode2 = f0_inode1 + pol_decomp.nnodes - 1; // 1-indexed
308  set_f0_f0g_ptr(f0_inode1, f0_inode2, f0_f0g.data());
309 
310  // Set all_species unmanaged views
311  int f0_species_cnt=0;
312  for_all_nonadiabatic_species([&](Species<DeviceType>& species){
313  set_unmanaged_f0_species_view(f0_f0g.f, f0_species_cnt, species.f0.f0g_h);
314  f0_species_cnt++;
315  }, NoDevicePtl);
316  }
317 
318  private:
319 
321  if(!particles_d_has_owner){
322  // Set new owner and initialize with 0 particles
323  particles_d_owner = isp;
324 
325  all_species[particles_d_owner].particles_d = Cabana::AoSoA<ParticleDataTypes,DeviceType,VEC_LEN>("particles_d", 0);
326  all_species[particles_d_owner].owns_particles_d = true;
327  particles_d_has_owner = true;
328  }else{
329  // No-op if species is handing off to itself
330  if(particles_d_owner == isp) return;
331 
332  // Shallow copy to new owner
333  all_species[isp].particles_d = all_species[particles_d_owner].particles_d;
334 
335  // Delete original
336  all_species[particles_d_owner].particles_d = Cabana::AoSoA<ParticleDataTypes,DeviceType,VEC_LEN>();
337  all_species[particles_d_owner].owns_particles_d = false;
338 
339  // Set new owner
340  particles_d_owner = isp;
341  all_species[particles_d_owner].owns_particles_d = true;
342  }
343  }
344 
345  void manage_particle_ownership(int isp, DevicePtlOpt device_ptl_opt){
346  if((!all_species[isp].is_adiabatic) && device_ptl_opt==UseDevicePtl){
347  if(species_share_particles_d_ownership){
348  transfer_particles_d_ownership(isp);
349  }else{
350  if(!all_species[isp].owns_particles_d){
351  all_species[isp].particles_d = Cabana::AoSoA<ParticleDataTypes,DeviceType,VEC_LEN>("particles_d", 0);
352  all_species[isp].owns_particles_d = true;
353  }
354  }
355 
356  // Resize device particles if they are used
357  all_species[isp].resize_device_particles();
358  }
359  }
360 
361 };
362 
363 #endif
Kokkos::View< double **, Kokkos::LayoutRight, HostType > f0_n_Ta
Equilibrium n_Ta at nodes.
Definition: plasma.hpp:39
void for_all_ions(F func, DevicePtlOpt device_ptl_opt=UseDevicePtl)
Definition: plasma.hpp:185
Kokkos::View< double *, Kokkos::LayoutRight, HostType > f0_B_B0
Bfield normalized to B0 at nodes.
Definition: plasma.hpp:44
bool owns_particles_d
Whether the species owns the device particle allocation right now.
Definition: species.hpp:97
View< double **, CLayout, HostType > f0_den_global
Equilibrium density at vertices.
Definition: plasma.hpp:80
Distribution< Device > f0
Species distribution in velocity space on local mesh nodes.
Definition: species.hpp:116
void manage_particle_ownership(int isp, DevicePtlOpt device_ptl_opt)
Definition: plasma.hpp:345
Plasma(NLReader::NamelistReader &nlr, bool use_f0_grid, const Grid< Device > &grid, const MagneticField< DeviceType > &magnetic_field, const DomainDecomposition< DeviceType > &pol_decomp, const VelocityGrid &vgrid)
Definition: plasma.hpp:95
bool is_electron
Whether this species is the electrons.
Definition: species.hpp:75
void f0_set_ptrs(int nnode, double *f0_delta_n_cpp, double *f0_delta_u_cpp, double *f0_delta_T_cpp)
ParticleType
Definition: plasma.hpp:143
void for_one_species(int isp, F func, DevicePtlOpt device_ptl_opt=UseDevicePtl)
Definition: plasma.hpp:208
void update_decomposed_f0_calculations(const DomainDecomposition< DeviceType > &pol_decomp)
Definition: plasma.hpp:244
T get(const string &param, const T default_val, int val_ind=0)
Definition: NamelistReader.hpp:361
View< double **, CLayout, HostType > f0_temp_global
Equilibrium temperature at vertices.
Definition: plasma.hpp:81
Definition: velocity_grid.hpp:7
bool particles_d_has_owner
Whether a species owns the device particles allocation.
Definition: plasma.hpp:19
void deallocate_device_ptl()
Definition: plasma.hpp:232
Definition: plasma.hpp:140
Definition: NamelistReader.hpp:163
Definition: magnetic_field.hpp:12
DecomposedRecalculableF0Arrays decomposed_recalculable_f0_arrays
Contains f0 values that are poloidally decomposed but don&#39;t need to be transferred between ranks duri...
Definition: plasma.hpp:73
int idx
Index in all_species.
Definition: species.hpp:74
VGridDistribution< HostType > f0_f0g
Definition: plasma.hpp:29
void for_all_nonadiabatic_species(F func, DevicePtlOpt device_ptl_opt=UseDevicePtl)
Definition: plasma.hpp:159
Definition: grid.hpp:50
void f0_initialize_wrap()
int nonadiabatic_idx
Index of species skipping adiabatic species (for compatibility with fortran arrays) ...
Definition: species.hpp:77
bool default_residence_option()
Definition: species.hpp:31
int n_ptl
Number of particles.
Definition: species.hpp:91
Kokkos::View< double **, Kokkos::LayoutRight, HostType > f0_delta_n
Flux-surface averaged change of density.
Definition: plasma.hpp:76
std::vector< Species< DeviceType > > all_species
Every particle species in the simulation.
Definition: plasma.hpp:24
void resize_f0_f0g(const DomainDecomposition< DeviceType > &pol_decomp, const VelocityGrid &vgrid)
Definition: plasma.hpp:302
int nnodes
Number of nodes belonging to this MPI rank.
Definition: domain_decomposition.hpp:40
Kokkos::View< double **, Kokkos::LayoutRight, HostType > f0_flow
Equilibrium flow at nodes.
Definition: plasma.hpp:41
Definition: plasma.hpp:144
Cabana::AoSoA< ParticleDataTypes, Device, VEC_LEN > particles_d
Particles on device.
Definition: species.hpp:95
Kokkos::View< double **, Kokkos::LayoutRight, HostType > f0_inv_grid_vol
Inverse grid volume at nodes.
Definition: plasma.hpp:36
void set_unmanaged_f0_species_view(const T_in &view_in, int isp, T_out &view_out)
Definition: plasma.hpp:295
int particles_d_owner
Which species, if any, owns the device particles allocation.
Definition: plasma.hpp:20
void use_namelist(const string &namelist)
Definition: NamelistReader.hpp:330
Kokkos::View< double **, Kokkos::LayoutRight, HostType > f0_grid_vol_vonly
Grid volume (v only) at nodes.
Definition: plasma.hpp:38
void f0_init_global_arrays(double *f0_den_global, double *f0_temp_global)
int largest_n_ptl(bool check_backup)
Definition: plasma.hpp:213
DevicePtlOpt
Definition: plasma.hpp:138
int nspecies
Number of species including electrons.
Definition: plasma.hpp:83
void for_all_species(F func, DevicePtlOpt device_ptl_opt=UseDevicePtl)
Definition: plasma.hpp:150
int n_nonadiabatic_species
Number of nonadiabatic species.
Definition: plasma.hpp:84
Kokkos::View< double **, Kokkos::LayoutRight, HostType > f0_T_ev
Equilibrium temperature at nodes.
Definition: plasma.hpp:35
Kokkos::View< double **, Kokkos::LayoutRight, HostType > f0_grid_vol
Grid volume at nodes.
Definition: plasma.hpp:37
void set_f0_f0g_ptr(int f0_inode1, int f0_inode2, double *f0_f0g_loc)
Kokkos::View< double **, Kokkos::LayoutRight, HostType > f0_delta_u
Flux-surface averaged change of parallel flow.
Definition: plasma.hpp:77
bool namelist_present(const string &namelist)
Definition: NamelistReader.hpp:320
Definition: plasma.hpp:139
DecomposedRecalculableF0Arrays()
Definition: plasma.hpp:46
void f0_init_decomposed_arrays(double *f0_T_ev_cpp, double *f0_inv_grid_vol_cpp, double *f0_grid_vol_cpp, double *f0_grid_vol_vonly_cpp, double *f0_n_Ta_cpp, double *f0_den_cpp, double *f0_flow_cpp, double *f0_B_B0_cpp)
Kokkos::View< T *, Kokkos::LayoutRight, Device > my_subview(const Kokkos::View< T ****, Kokkos::LayoutRight, Device > &view, int i, int j, int k)
Definition: my_subview.hpp:8
Definition: magnetic_field.F90:1
int n_backup_particles
Definition: species.hpp:113
void for_electrons(F func, DevicePtlOpt device_ptl_opt=UseDevicePtl)
Definition: plasma.hpp:174
double * data()
Definition: vgrid_distribution.hpp:68
Definition: plasma.hpp:15
Kokkos::View< double **, Kokkos::LayoutRight, HostType > f0_den
Equilibrium density at nodes.
Definition: plasma.hpp:40
Definition: plasma.hpp:145
void for_all(ParticleType particle_type, F func, DevicePtlOpt device_ptl_opt=UseDevicePtl)
Definition: plasma.hpp:196
bool species_share_particles_d_ownership
Whether to use the device particles sharing scheme.
Definition: plasma.hpp:18
Plasma()
Definition: plasma.hpp:87
Definition: species.hpp:71
void transfer_particles_d_ownership(int isp)
Definition: plasma.hpp:320
Kokkos::View< double **, Kokkos::LayoutRight, HostType > f0_delta_T
Flux-surface averaged change of temperature.
Definition: plasma.hpp:78
void init_global_f0_arrays(int grid_nnode)
Definition: plasma.hpp:264
int first_node
First mesh node belonging to this MPI rank.
Definition: domain_decomposition.hpp:39
DecomposedRecalculableF0Arrays(int nsp, const DomainDecomposition< DeviceType > &pol_decomp)
Definition: plasma.hpp:48
Kokkos::View< int *, Kokkos::LayoutRight, HostType > gvid0_pid_h
Which processors get which vertices (host)
Definition: domain_decomposition.hpp:43