3 #include <Cabana_AoSoA.hpp>
4 #include <Cabana_DeepCopy.hpp>
5 #include <Kokkos_Core.hpp>
70 template<
class Device>
92 Cabana::AoSoA<ParticleDataTypes,HostType,VEC_LEN>
particles;
108 Cabana::AoSoA<PhaseDataTypes,HostType,VEC_LEN>
phase0;
109 Cabana::AoSoA<PhaseDataTypes,Device,VEC_LEN>
phase0_d;
124 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,
127 template<
class Gr
idDevice>
174 auto ph = Cabana::slice<PtlSlice::Ph>(
particles);
175 auto ct = Cabana::slice<PtlSlice::Ct>(
particles);
176 auto gid = Cabana::slice<PtlSlice::Gid>(
particles);
178 auto flag = Cabana::slice<PtlSlice::Flag>(
particles);
185 long long int gid_offset = 0;
189 for (
int i=0;i<
n_ptl;i++){
191 gid(i) = gid_offset + i+1;
196 for (
int j=0;j<6;j++) ph(i, j) = gid(i) + (j)*0.1;
197 for (
int j=0;j<3;j++) ct(i, j) = gid(i) + (j+6)*0.1;
204 for (
int j=0;j<6;j++) ph(i, j) = gid(i) + (j)*0.1;
205 for (
int j=0;j<3;j++) ct(i, j) = gid(i) + (j+6)*0.1;
216 particles_d = Cabana::AoSoA<ParticleDataTypes,Device,VEC_LEN>();
248 particles = Cabana::AoSoA<ParticleDataTypes,HostType,VEC_LEN>();
285 #ifdef PRINT_COPA_ARTIFACTS
286 if(
is_rank_zero()) printf(
"\nCopying Cabana AoSoA particle data to GPU with Cabana::deep_copy\n");
305 #ifdef PRINT_COPA_ARTIFACTS
306 if(
is_rank_zero()) printf(
"\nCopying Cabana AoSoA particle data from GPU with Cabana::deep_copy\n");
350 int last_ptl_index =
n_ptl - 1;
351 auto ph = Cabana::slice<PtlSlice::Ph>(
particles_d);
352 auto ct = Cabana::slice<PtlSlice::Ct>(
particles_d);
353 auto gid = Cabana::slice<PtlSlice::Gid>(
particles_d);
355 auto flag = Cabana::slice<PtlSlice::Flag>(
particles_d);
360 for (
int j=0;j<6;j++) ph(i, j) = ph(last_ptl_index, j);
361 for (
int j=0;j<3;j++) ct(i, j) = ct(last_ptl_index, j);
364 flag(i) = flag(last_ptl_index);
379 int last_ptl_index =
n_ptl - 1;
380 auto ph = Cabana::slice<PtlSlice::Ph>(
phase0_d);
384 for (
int j=0;j<6;j++) ph(i, j) = ph(last_ptl_index, j);
402 #ifdef PRINT_COPA_ARTIFACTS
403 if(
is_rank_zero()) printf(
"\nLaunching GPU kernel '%s' on Cabana AoSoA of particles\n", label.c_str());
419 if(!
owns_particles_d)
exit_XGC(
"\nSpecies tried to loop over particles on device, but doesn't own the device array.");
421 #ifdef PRINT_COPA_ARTIFACTS
422 if(
is_rank_zero()) printf(
"\nLaunching GPU kernel '%s' on Cabana AoSoA of particles\n", label.c_str());
427 use_streaming =
false;
436 if((!send_ptl) && (!return_ptl)) use_streaming =
false;
451 #ifdef USE_EPUSH_LAUNCH_BOUNDS
452 # if !defined(PUSH_MAX_THREADS_PER_BLOCK) || !defined(PUSH_MIN_WARPS_PER_EU)
453 # error "USE_EPUSH_LAUNCH_BOUNDS requires PUSH_MAX_THREADS_PER_BLOCK and PUSH_MIN_WARPS_PER_EU to be defined"
455 Kokkos::RangePolicy<ExSpace, Kokkos::LaunchBounds<PUSH_MAX_THREADS_PER_BLOCK, PUSH_MIN_WARPS_PER_EU>>
459 exit_XGC(
"\nERROR: LaunchBounds::Custom specified, but USE_EPUSH_LAUNCH_BOUNDS is not defined\n");
462 Kokkos::RangePolicy<ExSpace>
484 for (
int i_simd = 0; i_simd<SIMD_SIZE; i_simd++){
485 int p_vec = inds.
a + i_simd;
486 ph0_loc[inds.
s].
r[p_vec] = ptl_loc[inds.
s].
ph.
r[p_vec];
487 ph0_loc[inds.
s].
z[p_vec] = ptl_loc[inds.
s].
ph.
z[p_vec];
488 ph0_loc[inds.
s].
phi[p_vec] = ptl_loc[inds.
s].
ph.
phi[p_vec];
489 ph0_loc[inds.
s].
rho[p_vec] = ptl_loc[inds.
s].
ph.
rho[p_vec];
490 ph0_loc[inds.
s].
w1[p_vec] = ptl_loc[inds.
s].
ph.
w1[p_vec];
491 ph0_loc[inds.
s].
w2[p_vec] = ptl_loc[inds.
s].
ph.
w2[p_vec];
539 backup_particles = Cabana::AoSoA<ParticleDataTypes,HostType,VEC_LEN>(
"backup_particles", 0);
557 for (
int i_simd = 0; i_simd<SIMD_SIZE; i_simd++){
558 int p_vec = inds.
a + i_simd;
559 part_one.
ph.
r[i_simd] = ph0_loc[inds.
s].
r[p_vec];
560 part_one.
ph.
z[i_simd] = ph0_loc[inds.
s].
z[p_vec];
561 part_one.
ph.
phi[i_simd] = ph0_loc[inds.
s].
phi[p_vec];
562 part_one.
ph.
rho[i_simd] = ph0_loc[inds.
s].
rho[p_vec];
563 part_one.
ph.
w1[i_simd] = ph0_loc[inds.
s].
w1[p_vec];
564 part_one.
ph.
w2[i_simd] = ph0_loc[inds.
s].
w2[p_vec];
570 long long int tmp_n_ptl =
n_ptl;
571 long long int out_n_ptl = 0;
572 MPI_Allreduce(&tmp_n_ptl, &out_n_ptl, 1, MPI_LONG_LONG_INT, MPI_SUM,
SML_COMM_WORLD);
575 return (
long long int)(
n_ptl);
581 int tmp_n_ptl =
n_ptl;
583 MPI_Allreduce(&tmp_n_ptl, &out_n_ptl, 1, MPI_INT, MPI_MAX,
SML_COMM_WORLD);
Cabana::AoSoA< PhaseDataTypes, HostType, VEC_LEN > phase0
Definition: species.hpp:108
bool stream_particles
Whether to stream particles between host and device if possible.
Definition: species.hpp:100
Definition: globals.hpp:77
KOKKOS_INLINE_FUNCTION VecPhase * ph0() const
Definition: species.hpp:475
KOKKOS_INLINE_FUNCTION int divide_and_round_up(int a, int b)
Definition: globals.hpp:126
bool owns_particles_d
Whether the species owns the device particle allocation right now.
Definition: species.hpp:97
subroutine adjust_n_ptl_for_core_ptl(n_ptl)
Definition: load.F90:445
void set_spall_num_and_ptr(int idx, int n_ptl, int n_vecs, VecParticles *ptl)
bool is_rank_zero()
Definition: globals.hpp:26
Distribution< Device > f0
Species distribution in velocity space on local mesh nodes.
Definition: species.hpp:116
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:112
bool is_electron
Whether this species is the electrons.
Definition: species.hpp:75
void for_all_particles(const std::string label, F lambda_func, const PtlMvmt mvmt, LaunchBounds launch_bounds=LaunchBounds::Default)
Definition: species.hpp:417
void save_backup_particles()
Definition: species.hpp:496
double c2_2m
c2/2m
Definition: species.hpp:83
double rho[VEC_LEN]
Definition: particles.hpp:72
void copy_to_phase0(Species< Device > &species)
Definition: species.hpp:479
Definition: species.hpp:54
Simd< double > w1
Definition: particles.hpp:22
double c_m
c/m
Definition: species.hpp:82
Definition: species.hpp:53
bool default_streaming_option()
Definition: species.hpp:24
Eq::Profile< Device > eq_den
Definition: species.hpp:121
Definition: globals.hpp:82
KOKKOS_INLINE_FUNCTION VecParticles * ptl() const
Definition: species.hpp:471
Definition: NamelistReader.hpp:163
KinType kintype
Whether the species is gyrokinetic or drift kinetic.
Definition: species.hpp:78
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:74
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:105
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
Definition: streamed_parallel_for.hpp:16
Definition: streamed_parallel_for.hpp:14
void set_buffer_phase0_d()
Definition: species.hpp:377
long long int get_total_n_ptl()
Definition: species.hpp:568
void set_buffer_particles_d()
Definition: species.hpp:348
Simd< double > rho
Definition: particles.hpp:21
Definition: species.hpp:55
int p_range< DeviceType >(int num_particle)
Definition: particles.hpp:157
double charge_eu
Particle charge in eu.
Definition: species.hpp:81
Definition: species.hpp:47
void resize_particles(int new_n_ptl)
Definition: species.hpp:210
double mass
Particle mass.
Definition: species.hpp:79
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:401
Cabana::AoSoA< ParticleDataTypes, Device, VEC_LEN > particles_d
Particles on device.
Definition: species.hpp:95
Definition: species.hpp:44
double w2[VEC_LEN]
Definition: particles.hpp:74
#define TIMER(N, F)
Definition: timer_macro.hpp:29
RKRestorationMethod
Definition: species.hpp:64
idx
Definition: diag_f0_df_port1.hpp:32
void copy_particles_to_device_if_not_resident()
Definition: species.hpp:331
RKRestorationMethod RK_restoration_method
Currently, electrons must use first method and ions must use second.
Definition: species.hpp:103
Simd< double > r
Definition: particles.hpp:18
void resize_host_particles_to_match_device()
Definition: species.hpp:230
Definition: species.hpp:65
ReturnOpt return_opt
Definition: species.hpp:59
ReturnOpt
Definition: species.hpp:52
Option
Definition: streamed_parallel_for.hpp:13
void restore_particles_from_backup()
Definition: species.hpp:517
Definition: globals.hpp:83
SendOpt send_opt
Definition: species.hpp:58
double charge
Particle charge.
Definition: species.hpp:80
SimdPhase ph
Definition: particles.hpp:59
void copy_particles_from_device()
Definition: species.hpp:302
void copy_particles_from_device_if_not_resident()
Definition: species.hpp:337
void unassign_host_particles()
Definition: species.hpp:247
int ncycles_between_sorts
Number of subcycles between sorts.
Definition: species.hpp:88
Definition: particles.hpp:58
Cabana::AoSoA< PhaseDataTypes, Device, VEC_LEN > phase0_d
Definition: species.hpp:109
int SML_COMM_RANK
Definition: my_mpi.cpp:5
KinType
Definition: globals.hpp:81
Species(SpeciesType sp_type, int n_ptl)
Definition: species.hpp:131
bool is_deltaf
Whether this species is deltaf.
Definition: species.hpp:85
VecPhase ph
Definition: particles.hpp:86
Definition: species.hpp:40
Definition: species.hpp:66
constexpr double PROTON_MASS
Definition: globals.hpp:121
void set_min_max_num(int isp, int n_ptl)
int minimum_ptl_reservation
The minimum reservation size for particles.
Definition: species.hpp:90
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:319
Definition: species.hpp:48
void resize_device_particles(int new_n_ptl)
Definition: species.hpp:268
Definition: species.hpp:49
void exit_XGC(std::string msg)
Definition: globals.hpp:36
void copy_particles_from_device_if_resident()
Definition: species.hpp:325
bool is_adiabatic
Whether this species is adiabatic.
Definition: species.hpp:76
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:74
int n_backup_particles
Definition: species.hpp:113
Eq::Profile< Device > eq_flow
Definition: species.hpp:122
Definition: streamed_parallel_for.hpp:15
SendOpt
Definition: species.hpp:46
void copy_particles_to_device()
Definition: species.hpp:282
KOKKOS_INLINE_FUNCTION void restore_phase_from_phase0(const AoSoAIndices< Device > &inds, SimdParticles &part_one) const
Definition: species.hpp:555
Species(int n_ptl_in)
Definition: species.hpp:160
double phi[VEC_LEN]
Definition: particles.hpp:71
Simd< double > w2
Definition: particles.hpp:23
double r[VEC_LEN]
Definition: particles.hpp:69
Definition: species.hpp:71
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:40
Eq::Profile< Device > eq_temp
Definition: species.hpp:120
bool particles_resident_on_device
Whether the particles can reside on device.
Definition: species.hpp:99
PtlMvmt(SendOpt send_opt, ReturnOpt return_opt)
Definition: species.hpp:61
int ncycles
Number of subcycles.
Definition: species.hpp:87
Definition: profile.hpp:65
int collision_grid_index
Which collision grid to use.
Definition: species.hpp:118
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:118
Definition: distribution.hpp:10
void resize_device_particles()
Definition: species.hpp:253
int get_max_n_ptl()
Definition: species.hpp:579
LaunchBounds
Definition: species.hpp:390
Cabana::AoSoA< ParticleDataTypes, HostType, VEC_LEN > particles
Particles.
Definition: species.hpp:92
double w1[VEC_LEN]
Definition: particles.hpp:73
Definition: species.hpp:40