particles.hpp

Go to the documentation of this file.

#ifndef PARTICLES_HPP
#define PARTICLES_HPP
#include "globals.hpp"
 
// Phase
struct SimdPhase{
    Simd<double> r;   
    Simd<double> z;   
    Simd<double> phi; 
    Simd<double> rho; 
    Simd<double> w1;  
    Simd<double> w2;  
 
    /* Returns a SimdVector2D alias to r and z
     * */
    KOKKOS_INLINE_FUNCTION SimdVector2D& x() {
        return *(reinterpret_cast<SimdVector2D*>(this));
    }
 
    /* Returns a SimdVector2D alias to r and z
     * */
    KOKKOS_INLINE_FUNCTION const SimdVector2D& x() const {
        return *(reinterpret_cast<const SimdVector2D*>(this));
    }
 
    /* Returns a SimdVector alias to r, z, and phi
     * */
    KOKKOS_INLINE_FUNCTION SimdVector& v() {
        return *(reinterpret_cast<SimdVector*>(this));
    }
 
    /* Returns a SimdVector alias to r, z, and phi
     * */
    KOKKOS_INLINE_FUNCTION const SimdVector& v() const {
        return *(reinterpret_cast<const SimdVector*>(this));
    }
};
 
// Constants
struct SimdConstants{
    Simd<double> mu; 
    Simd<double> w0; 
    Simd<double> f0; 
#ifdef PTL_G2
    Simd<double> g2;
#endif
};
 
// SoA particle structure, with inner array of length SIMD_SIZE (generally, VEC_LEN on CPU, 1 on GPU)
struct SimdParticles{
    SimdPhase ph;
    SimdConstants ct;
    Simd<long long int> gid;
#ifdef ESC_PTL
    Simd<long long int> flag; 
#endif    
 
    KOKKOS_INLINE_FUNCTION bool is_active(int i_simd) const{
        return (gid[i_simd]>0);
    }
 
    KOKKOS_INLINE_FUNCTION bool is_inactive(int i_simd) const{
        return (gid[i_simd]<=0);
    }
 
    KOKKOS_INLINE_FUNCTION void deactivate(const Simd<bool>& mask){
        for (int i_simd = 0; i_simd<SIMD_SIZE; i_simd++){
            // Don't reactivate particle
            gid[i_simd] = (mask[i_simd] && is_active(i_simd)) ? -gid[i_simd] : gid[i_simd];
        }
    }
 
    KOKKOS_INLINE_FUNCTION void deactivate(int i_simd){
        // Don't reactivate particle
        gid[i_simd] = is_active(i_simd) ? -gid[i_simd] : gid[i_simd];
    }
};
 
// Phase
struct VecPhase{
    double r[VEC_LEN];
    double z[VEC_LEN];
    double phi[VEC_LEN];
    double rho[VEC_LEN];
    double w1[VEC_LEN];
    double w2[VEC_LEN];
};
 
// Constants
struct VecConstants{
    double mu[VEC_LEN];
    double w0[VEC_LEN];
    double f0[VEC_LEN];
#ifdef PTL_G2
    double g2[VEC_LEN];
#endif
};
 
// SoA particle structure, with inner array of length VEC_LEN (just like the AoSoA)
struct VecParticles{
    VecPhase ph;
    VecConstants ct;
    long long int gid[VEC_LEN];
#ifdef ESC_PTL
    long long int flag[VEC_LEN];
#endif
 
    KOKKOS_INLINE_FUNCTION bool is_active(int i_vec) const{
        return (gid[i_vec]>0);
    }
 
    KOKKOS_INLINE_FUNCTION bool is_inactive(int i_vec) const{
        return (gid[i_vec]<=0);
    }
};
 
// Other useful ways to access particles
template<int T>
struct VecParticlesSimple{
    double data[T*VEC_LEN];
};
 
template<int T>
struct OneParticle{
    double data[T];
};
 
KOKKOS_INLINE_FUNCTION void update_phases(SimdPhase &ph_new, const SimdPhase &ph, double local_dt, const SimdPhase &dph);
 
KOKKOS_INLINE_FUNCTION void simd_2_AoSoA(VecParticles *part, const SimdParticles &part_one, int a_vec, int s_vec);
 
KOKKOS_INLINE_FUNCTION void AoSoA_2_simd(SimdParticles &part_one, const VecParticles *part, int a_vec, int s_vec);
 
KOKKOS_INLINE_FUNCTION void update_dphm(SimdPhase &dphm,const SimdPhase &dpht);
 
KOKKOS_INLINE_FUNCTION void update_dpht(SimdPhase &dpht,const SimdPhase &dph,const SimdPhase &dphm);
 
template<class Device>
struct AoSoAIndices{
    int s; 
    int a; 
 
    KOKKOS_INLINE_FUNCTION AoSoAIndices(int idx);
};
 
template<class Device>
KOKKOS_INLINE_FUNCTION AoSoAIndices<Device>::AoSoAIndices(int idx){
    // CPU: access an array of particles
    s = idx;
    a = 0;
}
 
#ifdef USE_GPU
template<>
KOKKOS_INLINE_FUNCTION AoSoAIndices<DeviceType>::AoSoAIndices(int idx){
    // GPU: access an individual particle
    s = idx / VEC_LEN;
    a = idx % VEC_LEN;
}
#endif
 
template<class Device>
inline int p_range(int num_particle)
{
    return divide_and_round_up(num_particle, VEC_LEN);
}
 
#ifdef USE_GPU
template<>
inline int p_range<DeviceType>(int num_particle)
{
    return num_particle;
}
#endif
 
template <typename T>
inline long long int add_vec_buffer(T n_ptl) {
    long long int n = static_cast<long long int>(n_ptl);
    long long int vec = static_cast<long long int>(VEC_LEN);
    return divide_and_round_up(n, vec)*vec;
}
 
#ifdef USE_CABANA
#include <Cabana_AoSoA.hpp>
#include <Cabana_DeepCopy.hpp>
 
// Set the type for the particle AoSoA.
# ifdef ESC_PTL
using ParticleDataTypes = Cabana::MemberTypes<double[PTL_NPHASE],double[PTL_NCONST],long long int, long long int>;
# else
using ParticleDataTypes = Cabana::MemberTypes<double[PTL_NPHASE],double[PTL_NCONST],long long int>;
# endif
// Just the phase
using PhaseDataTypes = Cabana::MemberTypes<double[PTL_NPHASE]>;
#else
// Consolidate these when Cabana is removed
using ParticleDataTypes = VecParticles;
using PhaseDataTypes = VecPhase;
#endif
 
// Used for Cabana slices (getting one particle property at a time)
namespace PtlSlice{
#ifdef ESC_PTL
enum{Ph=0,Ct,Gid,Flag};
#else
enum{Ph=0,Ct,Gid};
#endif
}
 
 
namespace Particles{
 
#ifdef USE_CABANA
// Cabana Wrappers
template<class DataType, class Device> using Array = typename Cabana::AoSoA<DataType,Device,VEC_LEN>;
 
template<class DataType, class Device, class Device2>
void deep_copy(const Array<DataType,Device>& dest, const Array<DataType,Device2>& src){
    Cabana::deep_copy(dest, src);
}
 
template<class DataType>
constexpr size_t sizeof_ptl(){
    return sizeof( Cabana::Tuple<DataType> );
}
 
template<int SliceInd, typename T>
typename T::template member_slice_type<SliceInd> slice(const T& ptl){
    return Cabana::slice<SliceInd>(ptl);
}
#else
 
// A non-Cabana AoSoA implementation
template<class DataType, class Device>
class Array{
    public:
 
    using vec_data_type = std::conditional_t<std::is_same_v<DataType, ParticleDataTypes>, VecParticles, VecPhase>;
 
    // Same (used in streamed_parallel_for.hpp (should be removed)
    using soa_type = vec_data_type;
 
    size_t n;
 
    View<vec_data_type*,CLayout,Device> view;
 
    Array(){}
 
    Array(const std::string& label, size_t n_in)
      : n(n_in),
        view(NoInit(label), divide_and_round_up(n, VEC_LEN)) {}
 
    size_t size() const{
        return n;
    }
 
    vec_data_type* data() const{
        return view.data();
    }
 
    KOKKOS_INLINE_FUNCTION vec_data_type& access(size_t ind) const{
        return view(ind);
    }
 
    void reserve(size_t n_reserve){
        // Number of vectors requested to be reserved
        size_t n_vec_reserve = divide_and_round_up(n_reserve, VEC_LEN);
 
        // Number of vectors currently reserved
        size_t n_vec = view.size();
 
        // Reserve only increases, never decreases
        if(n_vec_reserve>n_vec){
            // Make new allocation
            View<vec_data_type*,CLayout,Device> resized_view(NoInit(view.label()), n_vec_reserve);
 
            // Copy data to new allocation
            size_t vecs_in_use = divide_and_round_up(n, VEC_LEN);
            Kokkos::deep_copy(
                Kokkos::subview(resized_view, Kokkos::pair<size_t, size_t>( 0, vecs_in_use)),
                Kokkos::subview(view,         Kokkos::pair<size_t, size_t>( 0, vecs_in_use)) );
 
            // Reassign view to new allocation
            view = resized_view;
        }
    }
 
    void resize(size_t new_n_ptl){
        // Copy existing particles to a larger allocation if needed
        reserve(new_n_ptl);
 
        n = new_n_ptl;
    }
};
 
template<class DataType, class Device, class Device2>
void deep_copy(const Array<DataType,Device>& dest, const Array<DataType,Device2>& src){
    if(dest.n != src.n){
        printf("Error in Particles::deep_copy: %s.n=(%d) must equal %s.n=(%d).\n", dest.view.label().c_str(), dest.n, src.view.label().c_str(), src.n);
        fflush(stdout);
        exit_XGC("Error: Size mismatch in Particles::deep_copy\n");
    }
    size_t vecs_in_use = divide_and_round_up(dest.n, VEC_LEN);
    Kokkos::deep_copy(
        Kokkos::subview(dest.view, Kokkos::pair<size_t, size_t>( 0, vecs_in_use)),
        Kokkos::subview(src.view,  Kokkos::pair<size_t, size_t>( 0, vecs_in_use)) );
}
 
// Hard-coded sizes of ParticleDataTypes and PhaseDataTypes
template<class DataType> constexpr size_t sizeof_ptl();
#ifdef ESC_PTL
template<> constexpr size_t sizeof_ptl<ParticleDataTypes>(){return 11*8;};
#else
template<> constexpr size_t sizeof_ptl<ParticleDataTypes>(){return 10*8;};
#endif
template<> constexpr size_t sizeof_ptl<PhaseDataTypes>(){return 6*8;};
 
 
template<int SliceInd, class DataType, class Device> class Slice;
 
// Hardcode these slices. This is temporary - remove slices completely.
//enum{Ph=0,Ct,Gid,Flag};
 
template<class Device>
class Slice<PtlSlice::Ph, PhaseDataTypes, Device>{
    static constexpr int ind = PtlSlice::Ph;
    using ArrayType = Array<PhaseDataTypes, Device>;
    const ArrayType array;
    public:
    Slice(const ArrayType& array_in) : array(array_in) {}
    KOKKOS_INLINE_FUNCTION double& operator ()(int ip, int j) const{
        int ivec = ip/VEC_LEN;
        double* ptr = (double*)(&array.view(ivec));
        int istr = ip%VEC_LEN;
        return ptr[j*VEC_LEN + istr];
    }
};
 
template<class Device>
class Slice<PtlSlice::Ph, ParticleDataTypes, Device>{
    static constexpr int ind = PtlSlice::Ph;
    using ArrayType = Array<ParticleDataTypes, Device>;
    const ArrayType array;
    public:
    Slice(const ArrayType& array_in) : array(array_in) {}
    KOKKOS_INLINE_FUNCTION double& operator ()(int ip, int j) const{
        int ivec = ip/VEC_LEN;
        double* ptr = (double*)(&array.view(ivec).ph);
        int istr = ip%VEC_LEN;
        return ptr[j*VEC_LEN + istr];
    }
};
 
template<class Device>
class Slice<PtlSlice::Ct, ParticleDataTypes, Device>{
    static constexpr int ind = PtlSlice::Ct;
    using ArrayType = Array<ParticleDataTypes, Device>;
    const ArrayType array;
    public:
    Slice(const ArrayType& array_in) : array(array_in) {}
    KOKKOS_INLINE_FUNCTION double& operator ()(int ip, int j) const{
        int ivec = ip/VEC_LEN;
        double* ptr = (double*)(&array.view(ivec).ct);
        int istr = ip%VEC_LEN; 
        return ptr[j*VEC_LEN + istr];
    }
};
 
template<class Device>
class Slice<PtlSlice::Gid, ParticleDataTypes, Device>{
    static constexpr int ind = PtlSlice::Gid;
    using ArrayType = Array<ParticleDataTypes, Device>;
    const ArrayType array;
    public:
    Slice(const ArrayType& array_in) : array(array_in) {}
    KOKKOS_INLINE_FUNCTION long long int& operator ()(int ip) const{
        int ivec = ip/VEC_LEN;
        long long int* ptr = (long long int*)(&array.view(ivec).gid);
        int istr = ip%VEC_LEN; 
        return ptr[istr];
    }
};
 
#ifdef ESC_PTL
template<class Device>
class Slice<PtlSlice::Flag, ParticleDataTypes, Device>{
    static constexpr int ind = PtlSlice::Flag;
    using ArrayType = Array<ParticleDataTypes, Device>;
    const ArrayType array;
    public:
    Slice(const ArrayType& array_in) : array(array_in) {}
    KOKKOS_INLINE_FUNCTION long long int& operator ()(int ip) const{
        int ivec = ip/VEC_LEN;
        long long int* ptr = (long long int*)(&array.view(ivec).flag);
        int istr = ip%VEC_LEN; 
        return ptr[istr];
    }
};
#endif
 
template<int SliceInd, class DataType, class Device>
Slice<SliceInd, DataType, Device> slice(const Array<DataType, Device>& array){
    return Slice<SliceInd, DataType, Device>(array);
}
 
#endif
 
};
 
 
#include "particles.tpp"
#endif