transpose_particles.hpp

Go to the documentation of this file.

#ifndef TRANSPOSE_PARTICLES_HPP
#define TRANSPOSE_PARTICLES_HPP

#include "space_settings.hpp"

template<class DataType>
void transpose_particles_from_AoS_to_AoSoA(Cabana::AoSoA<DataType,DeviceType,VEC_LEN>& local_particles, int ioffset, int n_vec){
    // Early return if VEC_LEN=1 since no transposition needed
    if (VEC_LEN == 1) return;

    const std::size_t PTL_N_DBL = sizeof( Cabana::Tuple<DataType>)/8;

    // Pointer to local particles
    VecParticlesSimple<PTL_N_DBL>* vptl = (VecParticlesSimple<PTL_N_DBL>*)local_particles.data() + ioffset;

#ifdef USE_GPU
    int team_size = VEC_LEN;
#else
    int team_size = 1;
#endif
    int league_size = n_vec;

    typedef Kokkos::View<double[VEC_LEN*PTL_N_DBL],ExSpace::scratch_memory_space, Kokkos::MemoryTraits<Kokkos::Unmanaged>> PtlVec;
    size_t shmem_size = PtlVec::shmem_size(); // Could halve this, but may worsen memory access
    Kokkos::parallel_for("transpose_particles_from_AoS_to_AoSoA",
                         Kokkos::TeamPolicy<ExSpace> (league_size, team_size).set_scratch_size(KOKKOS_TEAM_SCRATCH_OPT,Kokkos::PerTeam(shmem_size)),
                         KOKKOS_LAMBDA (Kokkos::TeamPolicy<ExSpace>::member_type team_member){
        // Locally shared: global index in population
        PtlVec shmem_ptl_vec(team_member.team_scratch(KOKKOS_TEAM_SCRATCH_OPT));

        int ivec = team_member.league_rank(); // vector assigned to this team
        int ith = team_member.team_rank(); // This thread's rank in the team

        // Copy vector to shared memory
        for(int idbl = ith; idbl<VEC_LEN*PTL_N_DBL; idbl+=team_size){ // Loop through doubles
            shmem_ptl_vec(idbl) = vptl[ivec].data[idbl];
        }

        // Write transposed particles back to particle array
        for(int iptl = ith; iptl<VEC_LEN; iptl+=team_size){ // Loop through particles 
            for(int iprop = 0; iprop<PTL_N_DBL; iprop++){ // Loop through properties
                vptl[ivec].data[iprop*VEC_LEN + iptl] = shmem_ptl_vec(iptl*PTL_N_DBL + iprop);
            }
        }
    });

    Kokkos::fence();
}

template<class DataType>
void transpose_particles_from_AoSoA_to_AoS(Cabana::AoSoA<DataType,DeviceType,VEC_LEN>& local_particles, int ioffset, int n_vec){
    // Early return if VEC_LEN=1 since no transposition needed
    if (VEC_LEN == 1) return;

    const std::size_t PTL_N_DBL = sizeof( Cabana::Tuple<DataType>)/8;

    // Pointer to local particles
    VecParticlesSimple<PTL_N_DBL>* vptl = (VecParticlesSimple<PTL_N_DBL>*)local_particles.data() + ioffset;

#ifdef USE_GPU
    int team_size = VEC_LEN;
#else
    int team_size = 1;
#endif
    int league_size = n_vec;

    typedef Kokkos::View<double[VEC_LEN*PTL_N_DBL],ExSpace::scratch_memory_space, Kokkos::MemoryTraits<Kokkos::Unmanaged>> PtlVec;
    size_t shmem_size = PtlVec::shmem_size(); // Could halve this, but may worsen memory access
    Kokkos::parallel_for("transpose_particles_from_AoS_to_AoSoA",
                         Kokkos::TeamPolicy<ExSpace> (league_size, team_size).set_scratch_size(KOKKOS_TEAM_SCRATCH_OPT,Kokkos::PerTeam(shmem_size)),
                         KOKKOS_LAMBDA (Kokkos::TeamPolicy<ExSpace>::member_type team_member){
        // Locally shared: global index in population
        PtlVec shmem_ptl_vec(team_member.team_scratch(KOKKOS_TEAM_SCRATCH_OPT));

        int ivec = team_member.league_rank(); // vector assigned to this team
        int ith = team_member.team_rank(); // This thread's rank in the team

        // Transpose into shared memory
        for(int iptl = ith; iptl<VEC_LEN; iptl+=team_size){ // Loop through particles
            for(int iprop = 0; iprop<PTL_N_DBL; iprop++){ // Loop through properties
                shmem_ptl_vec(iptl*PTL_N_DBL + iprop) = vptl[ivec].data[iprop*VEC_LEN + iptl];
            }
        }

        // Copy transposed particles back to particle array
        for(int idbl = ith; idbl<VEC_LEN*PTL_N_DBL; idbl+=team_size){ // Loop through doubles
            vptl[ivec].data[idbl] = shmem_ptl_vec(idbl);
        }
    });

    Kokkos::fence();
}

#endif