3 #include <Cabana_AoSoA.hpp>
4 #include <Cabana_DeepCopy.hpp>
5 #include <Kokkos_Core.hpp>
71 template<
class Device>
93 Cabana::AoSoA<ParticleDataTypes,HostType,VEC_LEN>
particles;
109 Cabana::AoSoA<PhaseDataTypes,HostType,VEC_LEN>
phase0;
110 Cabana::AoSoA<PhaseDataTypes,Device,VEC_LEN>
phase0_d;
130 Species(
int idx_in,
int nonadiabatic_idx_in,
bool is_electron_in,
bool is_adiabatic_in,
KinType kintype_in,
double mass_in,
double charge_in,
double charge_eu_in,
bool is_deltaf_in,
133 template<
class Gr
idDevice>
181 auto ph = Cabana::slice<PtlSlice::Ph>(
particles);
182 auto ct = Cabana::slice<PtlSlice::Ct>(
particles);
183 auto gid = Cabana::slice<PtlSlice::Gid>(
particles);
185 auto flag = Cabana::slice<PtlSlice::Flag>(
particles);
192 long long int gid_offset = 0;
196 for (
int i=0;i<
n_ptl;i++){
198 gid(i) = gid_offset + i+1;
203 for (
int j=0;j<6;j++) ph(i, j) = gid(i) + (j)*0.1;
204 for (
int j=0;j<3;j++) ct(i, j) = gid(i) + (j+6)*0.1;
211 for (
int j=0;j<6;j++) ph(i, j) = gid(i) + (j)*0.1;
212 for (
int j=0;j<3;j++) ct(i, j) = gid(i) + (j+6)*0.1;
223 particles_d = Cabana::AoSoA<ParticleDataTypes,Device,VEC_LEN>();
255 particles = Cabana::AoSoA<ParticleDataTypes,HostType,VEC_LEN>();
292 #ifdef PRINT_COPA_ARTIFACTS
293 if(
is_rank_zero()) printf(
"\nCopying Cabana AoSoA particle data to GPU with Cabana::deep_copy\n");
312 #ifdef PRINT_COPA_ARTIFACTS
313 if(
is_rank_zero()) printf(
"\nCopying Cabana AoSoA particle data from GPU with Cabana::deep_copy\n");
357 int last_ptl_index =
n_ptl - 1;
358 auto ph = Cabana::slice<PtlSlice::Ph>(
particles_d);
359 auto ct = Cabana::slice<PtlSlice::Ct>(
particles_d);
360 auto gid = Cabana::slice<PtlSlice::Gid>(
particles_d);
362 auto flag = Cabana::slice<PtlSlice::Flag>(
particles_d);
367 for (
int j=0;j<6;j++) ph(i, j) = ph(last_ptl_index, j);
368 for (
int j=0;j<3;j++) ct(i, j) = ct(last_ptl_index, j);
371 flag(i) = flag(last_ptl_index);
386 int last_ptl_index =
n_ptl - 1;
387 auto ph = Cabana::slice<PtlSlice::Ph>(
phase0_d);
391 for (
int j=0;j<6;j++) ph(i, j) = ph(last_ptl_index, j);
409 #ifdef PRINT_COPA_ARTIFACTS
410 if(
is_rank_zero()) printf(
"\nLaunching GPU kernel '%s' on Cabana AoSoA of particles\n", label.c_str());
416 inline void back_up_SoA(Cabana::AoSoA<ParticleDataTypes,Device,VEC_LEN>& backup_SoA,
int offset,
int n)
const{
417 auto ph_b = Cabana::slice<PtlSlice::Ph>(backup_SoA);
418 auto ct_b = Cabana::slice<PtlSlice::Ct>(backup_SoA);
419 auto gid_b = Cabana::slice<PtlSlice::Gid>(backup_SoA);
421 auto flag_b = Cabana::slice<PtlSlice::Flag>(backup_SoA);
424 auto ph = Cabana::slice<PtlSlice::Ph>(
particles_d);
425 auto ct = Cabana::slice<PtlSlice::Ct>(
particles_d);
426 auto gid = Cabana::slice<PtlSlice::Gid>(
particles_d);
428 auto flag = Cabana::slice<PtlSlice::Flag>(
particles_d);
431 Kokkos::parallel_for(
"backup_first_soa", Kokkos::RangePolicy<ExSpace>( 0, n ), KOKKOS_LAMBDA(
const int i ){
432 int i_offset = i + offset;
434 for (
int j=0;j<6;j++) ph_b(i, j) = ph(i_offset, j);
435 for (
int j=0;j<3;j++) ct_b(i, j) = ct(i_offset, j);
436 gid_b(i) = gid(i_offset);
438 flag_b(i) = flag(i_offset);
445 inline void restore_backup_SoA(Cabana::AoSoA<ParticleDataTypes,Device,VEC_LEN>& backup_SoA,
int offset,
int n)
const{
446 auto ph_b = Cabana::slice<PtlSlice::Ph>(backup_SoA);
447 auto ct_b = Cabana::slice<PtlSlice::Ct>(backup_SoA);
448 auto gid_b = Cabana::slice<PtlSlice::Gid>(backup_SoA);
450 auto flag_b = Cabana::slice<PtlSlice::Flag>(backup_SoA);
453 auto ph = Cabana::slice<PtlSlice::Ph>(
particles_d);
454 auto ct = Cabana::slice<PtlSlice::Ct>(
particles_d);
455 auto gid = Cabana::slice<PtlSlice::Gid>(
particles_d);
457 auto flag = Cabana::slice<PtlSlice::Flag>(
particles_d);
460 Kokkos::parallel_for(
"backup_first_soa", Kokkos::RangePolicy<ExSpace>( 0, n ), KOKKOS_LAMBDA(
const int i ){
461 int i_offset = i + offset;
463 for (
int j=0;j<6;j++) ph(i_offset, j) = ph_b(i, j);
464 for (
int j=0;j<3;j++) ct(i_offset, j) = ct_b(i, j);
465 gid(i_offset) = gid_b(i);
467 flag(i_offset) = flag_b(i);
478 inline void for_particle_range(
int begin_idx,
int end_idx,
const std::string label, F lambda_func)
const {
479 if(end_idx <= begin_idx)
return;
482 int first_soa = begin_idx/VEC_LEN;
483 int n_other_ptl_in_first_soa = begin_idx - first_soa*VEC_LEN;
484 bool first_soa_is_partial = (n_other_ptl_in_first_soa>0);
485 int last_soa = (end_idx-1)/VEC_LEN;
486 int n_other_ptl_in_last_soa = (last_soa+1)*VEC_LEN - end_idx;
487 bool last_soa_is_partial = (n_other_ptl_in_last_soa>0);
490 int first_item_in_shifted_range = first_soa*VEC_LEN;
492 int first_item_in_shifted_range = first_soa;
495 Cabana::AoSoA<ParticleDataTypes,Device,VEC_LEN> ptl_first_soa;
496 Cabana::AoSoA<ParticleDataTypes,Device,VEC_LEN> ptl_last_soa;
497 if(first_soa_is_partial){
499 ptl_first_soa = Cabana::AoSoA<ParticleDataTypes,Device,VEC_LEN>(
"ptl_first_soa", n_other_ptl_in_first_soa);
500 back_up_SoA(ptl_first_soa, first_soa*VEC_LEN, n_other_ptl_in_first_soa);
502 if(last_soa_is_partial){
504 ptl_last_soa = Cabana::AoSoA<ParticleDataTypes,Device,VEC_LEN>(
"ptl_last_soa", n_other_ptl_in_last_soa);
505 back_up_SoA(ptl_last_soa, end_idx, n_other_ptl_in_last_soa);
509 Kokkos::RangePolicy<ExSpace> particle_range_policy( first_item_in_shifted_range,
p_range<DeviceType>(end_idx) );
513 if(first_soa_is_partial){
516 if(last_soa_is_partial){
531 if(!
owns_particles_d)
exit_XGC(
"\nSpecies tried to loop over particles on device, but doesn't own the device array.");
533 #ifdef PRINT_COPA_ARTIFACTS
534 if(
is_rank_zero()) printf(
"\nLaunching GPU kernel '%s' on Cabana AoSoA of particles\n", label.c_str());
539 use_streaming =
false;
548 if((!send_ptl) && (!return_ptl)) use_streaming =
false;
563 #ifdef USE_EPUSH_LAUNCH_BOUNDS
564 # if !defined(PUSH_MAX_THREADS_PER_BLOCK) || !defined(PUSH_MIN_WARPS_PER_EU)
565 # error "USE_EPUSH_LAUNCH_BOUNDS requires PUSH_MAX_THREADS_PER_BLOCK and PUSH_MIN_WARPS_PER_EU to be defined"
567 Kokkos::RangePolicy<ExSpace, Kokkos::LaunchBounds<PUSH_MAX_THREADS_PER_BLOCK, PUSH_MIN_WARPS_PER_EU>>
571 exit_XGC(
"\nERROR: LaunchBounds::Custom specified, but USE_EPUSH_LAUNCH_BOUNDS is not defined\n");
574 Kokkos::RangePolicy<ExSpace>
596 for (
int i_simd = 0; i_simd<SIMD_SIZE; i_simd++){
597 int p_vec = inds.
a + i_simd;
598 ph0_loc[inds.
s].
r[p_vec] = ptl_loc[inds.
s].
ph.
r[p_vec];
599 ph0_loc[inds.
s].
z[p_vec] = ptl_loc[inds.
s].
ph.
z[p_vec];
600 ph0_loc[inds.
s].
phi[p_vec] = ptl_loc[inds.
s].
ph.
phi[p_vec];
601 ph0_loc[inds.
s].
rho[p_vec] = ptl_loc[inds.
s].
ph.
rho[p_vec];
602 ph0_loc[inds.
s].
w1[p_vec] = ptl_loc[inds.
s].
ph.
w1[p_vec];
603 ph0_loc[inds.
s].
w2[p_vec] = ptl_loc[inds.
s].
ph.
w2[p_vec];
651 backup_particles = Cabana::AoSoA<ParticleDataTypes,HostType,VEC_LEN>(
"backup_particles", 0);
669 for (
int i_simd = 0; i_simd<SIMD_SIZE; i_simd++){
670 int p_vec = inds.
a + i_simd;
671 part_one.
ph.
r[i_simd] = ph0_loc[inds.
s].
r[p_vec];
672 part_one.
ph.
z[i_simd] = ph0_loc[inds.
s].
z[p_vec];
673 part_one.
ph.
phi[i_simd] = ph0_loc[inds.
s].
phi[p_vec];
674 part_one.
ph.
rho[i_simd] = ph0_loc[inds.
s].
rho[p_vec];
675 part_one.
ph.
w1[i_simd] = ph0_loc[inds.
s].
w1[p_vec];
676 part_one.
ph.
w2[i_simd] = ph0_loc[inds.
s].
w2[p_vec];
682 long long int tmp_n_ptl =
n_ptl;
683 long long int out_n_ptl = 0;
684 MPI_Allreduce(&tmp_n_ptl, &out_n_ptl, 1, MPI_LONG_LONG_INT, MPI_SUM,
SML_COMM_WORLD);
687 return (
long long int)(
n_ptl);
693 int tmp_n_ptl =
n_ptl;
695 MPI_Allreduce(&tmp_n_ptl, &out_n_ptl, 1, MPI_INT, MPI_MAX,
SML_COMM_WORLD);
constexpr double EV_2_J
Conversion rate ev to J.
Definition: globals.hpp:172
Cabana::AoSoA< PhaseDataTypes, HostType, VEC_LEN > phase0
Definition: species.hpp:109
bool stream_particles
Whether to stream particles between host and device if possible.
Definition: species.hpp:101
Definition: globals.hpp:77
KOKKOS_INLINE_FUNCTION VecPhase * ph0() const
Definition: species.hpp:587
KOKKOS_INLINE_FUNCTION int divide_and_round_up(int a, int b)
Definition: globals.hpp:179
bool owns_particles_d
Whether the species owns the device particle allocation right now.
Definition: species.hpp:98
void back_up_SoA(Cabana::AoSoA< ParticleDataTypes, Device, VEC_LEN > &backup_SoA, int offset, int n) const
Definition: species.hpp:416
KOKKOS_INLINE_FUNCTION double get_f0_eq_gyro_radius(int inode, double smu_n, double bfield) const
Definition: species.hpp:706
subroutine adjust_n_ptl_for_core_ptl(n_ptl)
Definition: load.F90:442
void set_spall_num_and_ptr(int idx, int n_ptl, int n_vecs, VecParticles *ptl)
bool is_rank_zero()
Definition: globals.hpp:26
void for_particle_range(int begin_idx, int end_idx, const std::string label, F lambda_func) const
Definition: species.hpp:478
Distribution< Device > f0
Species distribution in velocity space on local mesh nodes.
Definition: species.hpp:117
MPI_Comm SML_COMM_WORLD
Definition: my_mpi.cpp:4
Cabana::AoSoA< ParticleDataTypes, HostType, VEC_LEN > backup_particles
Copy of particles to be restored for RK2.
Definition: species.hpp:113
bool is_electron
Whether this species is the electrons.
Definition: species.hpp:76
void for_all_particles(const std::string label, F lambda_func, const PtlMvmt mvmt, LaunchBounds launch_bounds=LaunchBounds::Default)
Definition: species.hpp:529
void save_backup_particles()
Definition: species.hpp:608
double c2_2m
c2/2m
Definition: species.hpp:84
double rho[VEC_LEN]
Definition: particles.hpp:72
void copy_to_phase0(Species< Device > &species)
Definition: species.hpp:591
Definition: species.hpp:55
Simd< double > w1
Definition: particles.hpp:22
double c_m
c/m
Definition: species.hpp:83
Definition: species.hpp:54
bool default_streaming_option()
Definition: species.hpp:25
Eq::Profile< Device > eq_den
Definition: species.hpp:122
Definition: globals.hpp:82
KOKKOS_INLINE_FUNCTION VecParticles * ptl() const
Definition: species.hpp:583
Definition: NamelistReader.hpp:193
KinType kintype
Whether the species is gyrokinetic or drift kinetic.
Definition: species.hpp:79
Definition: magnetic_field.hpp:12
int add_vec_buffer(int n_ptl)
Definition: particles.hpp:164
int idx
Index in all_species.
Definition: species.hpp:75
Definition: particles.hpp:68
int a
The index in the inner array of the AoSoA.
Definition: particles.hpp:120
Definition: particles.hpp:85
bool particles_are_backed_up
Whether particles are currently backed up.
Definition: species.hpp:106
int nonadiabatic_idx
Index of species skipping adiabatic species (for compatibility with fortran arrays) ...
Definition: species.hpp:78
bool default_residence_option()
Definition: species.hpp:32
int n_ptl
Number of particles.
Definition: species.hpp:92
Definition: streamed_parallel_for.hpp:16
Definition: streamed_parallel_for.hpp:14
void set_buffer_phase0_d()
Definition: species.hpp:384
long long int get_total_n_ptl()
Definition: species.hpp:680
void set_buffer_particles_d()
Definition: species.hpp:355
Simd< double > rho
Definition: particles.hpp:21
Definition: species.hpp:56
int p_range< DeviceType >(int num_particle)
Definition: particles.hpp:157
int eq_flow_type
Definition: species.hpp:124
double charge_eu
Particle charge in eu.
Definition: species.hpp:82
Definition: species.hpp:48
void resize_particles(int new_n_ptl)
Definition: species.hpp:217
double mass
Particle mass.
Definition: species.hpp:80
KOKKOS_INLINE_FUNCTION double get_f0_eq_thermal_velocity(int inode) const
Definition: species.hpp:713
Species(int idx_in, int nonadiabatic_idx_in, bool is_electron_in, bool is_adiabatic_in, KinType kintype_in, double mass_in, double charge_in, double charge_eu_in, bool is_deltaf_in, int ncycles_in)
Definition: species.tpp:22
void for_all_particles(const std::string label, F lambda_func) const
Definition: species.hpp:408
Cabana::AoSoA< ParticleDataTypes, Device, VEC_LEN > particles_d
Particles on device.
Definition: species.hpp:96
Definition: species.hpp:45
double w2[VEC_LEN]
Definition: particles.hpp:74
#define TIMER(N, F)
Definition: timer_macro.hpp:24
RKRestorationMethod
Definition: species.hpp:65
idx
Definition: diag_f0_df_port1.hpp:32
void copy_particles_to_device_if_not_resident()
Definition: species.hpp:338
RKRestorationMethod RK_restoration_method
Currently, electrons must use first method and ions must use second.
Definition: species.hpp:104
Simd< double > r
Definition: particles.hpp:18
void resize_host_particles_to_match_device()
Definition: species.hpp:237
Definition: species.hpp:66
ReturnOpt return_opt
Definition: species.hpp:60
ReturnOpt
Definition: species.hpp:53
Option
Definition: streamed_parallel_for.hpp:13
void restore_particles_from_backup()
Definition: species.hpp:629
Definition: globals.hpp:83
SendOpt send_opt
Definition: species.hpp:59
double charge
Particle charge.
Definition: species.hpp:81
SimdPhase ph
Definition: particles.hpp:59
void copy_particles_from_device()
Definition: species.hpp:309
void copy_particles_from_device_if_not_resident()
Definition: species.hpp:344
void unassign_host_particles()
Definition: species.hpp:254
int ncycles_between_sorts
Number of subcycles between sorts.
Definition: species.hpp:89
Definition: particles.hpp:58
Cabana::AoSoA< PhaseDataTypes, Device, VEC_LEN > phase0_d
Definition: species.hpp:110
int SML_COMM_RANK
Definition: my_mpi.cpp:5
KinType
Definition: globals.hpp:81
Species(SpeciesType sp_type, int n_ptl)
Definition: species.hpp:137
bool is_deltaf
Whether this species is deltaf.
Definition: species.hpp:86
VecPhase ph
Definition: particles.hpp:86
Definition: species.hpp:41
Definition: species.hpp:67
constexpr double PROTON_MASS
Definition: globals.hpp:174
void set_min_max_num(int isp, int n_ptl)
int minimum_ptl_reservation
The minimum reservation size for particles.
Definition: species.hpp:91
int s
The index in the outer array of the AoSoA.
Definition: particles.hpp:119
Simd< double > z
Definition: particles.hpp:19
void copy_particles_to_device_if_resident()
Definition: species.hpp:326
Definition: species.hpp:49
void resize_device_particles(int new_n_ptl)
Definition: species.hpp:275
Definition: species.hpp:50
KOKKOS_INLINE_FUNCTION double get_f0_eq_thermal_velocity_lnode_h(int inode) const
Definition: species.hpp:719
void exit_XGC(std::string msg)
Definition: globals.hpp:36
void copy_particles_from_device_if_resident()
Definition: species.hpp:332
bool is_adiabatic
Whether this species is adiabatic.
Definition: species.hpp:77
Simd< double > phi
Definition: particles.hpp:20
Definition: magnetic_field.F90:1
static constexpr const Kokkos::Experimental::WorkItemProperty::HintLightWeight_t Async
Definition: space_settings.hpp:81
int n_backup_particles
Definition: species.hpp:114
Eq::Profile< Device > eq_flow
Definition: species.hpp:123
Definition: streamed_parallel_for.hpp:15
SendOpt
Definition: species.hpp:47
void copy_particles_to_device()
Definition: species.hpp:289
KOKKOS_INLINE_FUNCTION void restore_phase_from_phase0(const AoSoAIndices< Device > &inds, SimdParticles &part_one) const
Definition: species.hpp:667
Species(int n_ptl_in)
Definition: species.hpp:167
double phi[VEC_LEN]
Definition: particles.hpp:71
Simd< double > w2
Definition: particles.hpp:23
double r[VEC_LEN]
Definition: particles.hpp:69
GyroAverageMatrices< HostType > gyro_avg_matrices
Definition: species.hpp:126
Definition: species.hpp:72
void parallel_for(const std::string name, int n_ptl, Function func, Option option, HostAoSoA aosoa_h, DeviceAoSoA aosoa_d)
Definition: streamed_parallel_for.hpp:252
Definition: species.hpp:41
Eq::Profile< Device > eq_temp
Definition: species.hpp:121
bool particles_resident_on_device
Whether the particles can reside on device.
Definition: species.hpp:100
PtlMvmt(SendOpt send_opt, ReturnOpt return_opt)
Definition: species.hpp:62
int ncycles
Number of subcycles.
Definition: species.hpp:88
Definition: profile.hpp:65
int collision_grid_index
Which collision grid to use.
Definition: species.hpp:119
Definition: particles.hpp:118
SpeciesType
Definition: globals.hpp:76
double z[VEC_LEN]
Definition: particles.hpp:70
constexpr double UNIT_CHARGE
Charge of an electron (C)
Definition: globals.hpp:171
void restore_backup_SoA(Cabana::AoSoA< ParticleDataTypes, Device, VEC_LEN > &backup_SoA, int offset, int n) const
Definition: species.hpp:445
Definition: distribution.hpp:10
void resize_device_particles()
Definition: species.hpp:260
int get_max_n_ptl()
Definition: species.hpp:691
LaunchBounds
Definition: species.hpp:397
Cabana::AoSoA< ParticleDataTypes, HostType, VEC_LEN > particles
Particles.
Definition: species.hpp:93
double w1[VEC_LEN]
Definition: particles.hpp:73
Definition: species.hpp:41