XGCa
 All Classes Namespaces Files Functions Variables Typedefs Enumerations Enumerator Macros Pages
transpose_particles.hpp
Go to the documentation of this file.
1 #ifndef TRANSPOSE_PARTICLES_HPP
2 #define TRANSPOSE_PARTICLES_HPP
3 
4 #include "space_settings.hpp"
5 
6 template<class DataType>
7 void transpose_particles_from_AoS_to_AoSoA(Cabana::AoSoA<DataType,DeviceType,VEC_LEN>& local_particles, int ioffset, int n_vec){
8  const std::size_t PTL_N_DBL = sizeof( Cabana::Tuple<DataType>)/8;
9 
10  // Pointer to local particles
11  VecParticlesSimple<PTL_N_DBL>* vptl = (VecParticlesSimple<PTL_N_DBL>*)local_particles.data() + ioffset;
12 
13 #ifdef USE_GPU
14  int team_size = VEC_LEN;
15 #else
16  int team_size = 1;
17 #endif
18  int league_size = n_vec;
19 
20  typedef Kokkos::View<double[VEC_LEN*PTL_N_DBL],ExSpace::scratch_memory_space, Kokkos::MemoryTraits<Kokkos::Unmanaged>> PtlVec;
21  size_t shmem_size = PtlVec::shmem_size(); // Could halve this, but may worsen memory access
22  Kokkos::parallel_for("transpose_particles_from_AoS_to_AoSoA",
23  Kokkos::TeamPolicy<ExSpace> (league_size, team_size).set_scratch_size(KOKKOS_TEAM_SCRATCH_OPT,Kokkos::PerTeam(shmem_size)),
24  KOKKOS_LAMBDA (Kokkos::TeamPolicy<ExSpace>::member_type team_member){
25  // Locally shared: global index in population
26  PtlVec shmem_ptl_vec(team_member.team_scratch(KOKKOS_TEAM_SCRATCH_OPT));
27 
28  int ivec = team_member.league_rank(); // vector assigned to this team
29  int ith = team_member.team_rank(); // This thread's rank in the team
30 
31  // Copy vector to shared memory
32  for(int idbl = ith; idbl<VEC_LEN*PTL_N_DBL; idbl+=team_size){ // Loop through doubles
33  shmem_ptl_vec(idbl) = vptl[ivec].data[idbl];
34  }
35 
36  // Write transposed particles back to particle array
37  for(int iptl = ith; iptl<VEC_LEN; iptl+=team_size){ // Loop through particles
38  for(int iprop = 0; iprop<PTL_N_DBL; iprop++){ // Loop through properties
39  vptl[ivec].data[iprop*VEC_LEN + iptl] = shmem_ptl_vec(iptl*PTL_N_DBL + iprop);
40  }
41  }
42  });
43 
44  Kokkos::fence();
45 }
46 
47 template<class DataType>
48 void transpose_particles_from_AoSoA_to_AoS(Cabana::AoSoA<DataType,DeviceType,VEC_LEN>& local_particles, int ioffset, int n_vec){
49  const std::size_t PTL_N_DBL = sizeof( Cabana::Tuple<DataType>)/8;
50 
51  // Pointer to local particles
52  VecParticlesSimple<PTL_N_DBL>* vptl = (VecParticlesSimple<PTL_N_DBL>*)local_particles.data() + ioffset;
53 
54 #ifdef USE_GPU
55  int team_size = VEC_LEN;
56 #else
57  int team_size = 1;
58 #endif
59  int league_size = n_vec;
60 
61  typedef Kokkos::View<double[VEC_LEN*PTL_N_DBL],ExSpace::scratch_memory_space, Kokkos::MemoryTraits<Kokkos::Unmanaged>> PtlVec;
62  size_t shmem_size = PtlVec::shmem_size(); // Could halve this, but may worsen memory access
63  Kokkos::parallel_for("transpose_particles_from_AoS_to_AoSoA",
64  Kokkos::TeamPolicy<ExSpace> (league_size, team_size).set_scratch_size(KOKKOS_TEAM_SCRATCH_OPT,Kokkos::PerTeam(shmem_size)),
65  KOKKOS_LAMBDA (Kokkos::TeamPolicy<ExSpace>::member_type team_member){
66  // Locally shared: global index in population
67  PtlVec shmem_ptl_vec(team_member.team_scratch(KOKKOS_TEAM_SCRATCH_OPT));
68 
69  int ivec = team_member.league_rank(); // vector assigned to this team
70  int ith = team_member.team_rank(); // This thread's rank in the team
71 
72  // Transpose into shared memory
73  for(int iptl = ith; iptl<VEC_LEN; iptl+=team_size){ // Loop through particles
74  for(int iprop = 0; iprop<PTL_N_DBL; iprop++){ // Loop through properties
75  shmem_ptl_vec(iptl*PTL_N_DBL + iprop) = vptl[ivec].data[iprop*VEC_LEN + iptl];
76  }
77  }
78 
79  // Copy transposed particles back to particle array
80  for(int idbl = ith; idbl<VEC_LEN*PTL_N_DBL; idbl+=team_size){ // Loop through doubles
81  vptl[ivec].data[idbl] = shmem_ptl_vec(idbl);
82  }
83  });
84 
85  Kokkos::fence();
86 }
87 
88 #endif
Definition: particles.hpp:96
double data[T *VEC_LEN]
Definition: particles.hpp:97
void transpose_particles_from_AoSoA_to_AoS(Cabana::AoSoA< DataType, DeviceType, VEC_LEN > &local_particles, int ioffset, int n_vec)
Definition: transpose_particles.hpp:48
void transpose_particles_from_AoS_to_AoSoA(Cabana::AoSoA< DataType, DeviceType, VEC_LEN > &local_particles, int ioffset, int n_vec)
Definition: transpose_particles.hpp:7
#define KOKKOS_TEAM_SCRATCH_OPT
Definition: space_settings.hpp:68
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