particles.hpp

Go to the documentation of this file.

#ifndef PARTICLES_HPP
#define PARTICLES_HPP
#include <Cabana_AoSoA.hpp>
#include <Kokkos_Core.hpp>
#include "globals.hpp"

// Set the type for the particle AoSoA.
#ifdef ESC_PTL
using ParticleDataTypes = Cabana::MemberTypes<double[6],double[3],long long int, long long int>;
#else
using ParticleDataTypes = Cabana::MemberTypes<double[6],double[3],long long int>;
#endif
// Just the phase
using PhaseDataTypes = Cabana::MemberTypes<double[6]>;

// 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;
};

// 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];
};

// 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

inline int add_vec_buffer(int n_ptl) { return divide_and_round_up(n_ptl, VEC_LEN)*VEC_LEN;}

#include "particles.tpp"
#endif