simd.hpp

Go to the documentation of this file.

#ifndef SIMD_HPP
#define SIMD_HPP
#include <Kokkos_Core.hpp>

/*** SIMD structures ****/

// Use simd pragma on CPU
#ifdef USE_GPU
#define FORCE_SIMD
#else
#define FORCE_SIMD _Pragma("omp simd")
#endif



// A 1-dimension array of size SIMD_SIZE and type T
template<typename T>
struct Simd{
    T x[SIMD_SIZE];

    // Constructors
    KOKKOS_INLINE_FUNCTION Simd<T>(){}
    KOKKOS_INLINE_FUNCTION Simd<T>(T initval){
        for (int i_simd = 0; i_simd<SIMD_SIZE; i_simd++) x[i_simd] = initval;
    }

    // Allows array-like access pattern
    KOKKOS_INLINE_FUNCTION T operator [](int i) const {return x[i];}
    KOKKOS_INLINE_FUNCTION T& operator [](int i) {return x[i];}

    // Basic operations
    // Assignment from scalar
    template<typename Trhs>
    KOKKOS_INLINE_FUNCTION void operator = (const Trhs rhs) {
        for (int i_simd = 0; i_simd<SIMD_SIZE; i_simd++) x[i_simd] = rhs;
        // Void return, chaining not enabled
    }

    KOKKOS_INLINE_FUNCTION Simd<T> operator - () const {
        Simd<T> out;
        for (int i_simd = 0; i_simd<SIMD_SIZE; i_simd++) out.x[i_simd] = -x[i_simd];
        return out;
    }

    template<typename Trhs>
    KOKKOS_INLINE_FUNCTION Simd<T> operator - (const Simd<Trhs>& rhs) const {
        Simd<T> out;
        for (int i_simd = 0; i_simd<SIMD_SIZE; i_simd++) out.x[i_simd] = x[i_simd] - rhs[i_simd];
        return out;
    }

    template<typename Trhs>
    KOKKOS_INLINE_FUNCTION Simd<T> operator + (const Simd<Trhs>& rhs) const {
        Simd<T> out;
        for (int i_simd = 0; i_simd<SIMD_SIZE; i_simd++) out.x[i_simd] = x[i_simd] + rhs[i_simd];
        return out;
    }

    template<typename Trhs>
    KOKKOS_INLINE_FUNCTION Simd<T> operator * (const Simd<Trhs>& rhs) const {
        Simd<T> out;
        for (int i_simd = 0; i_simd<SIMD_SIZE; i_simd++) out.x[i_simd] = x[i_simd] * rhs[i_simd];
        return out;
    }

    template<typename Trhs>
    KOKKOS_INLINE_FUNCTION Simd<T> operator / (const Simd<Trhs>& rhs) const {
        Simd<T> out;
        for (int i_simd = 0; i_simd<SIMD_SIZE; i_simd++) out.x[i_simd] = x[i_simd] / rhs[i_simd];
        return out;
    }

    // Other operations

    // Sets values to 0.0
    KOKKOS_INLINE_FUNCTION void zero(){
        for (int i_simd = 0; i_simd<SIMD_SIZE; i_simd++) x[i_simd] = 0.0;
    }
};

// Operators with non-Simd argument
template<typename T>
KOKKOS_INLINE_FUNCTION Simd<T> operator - (const Simd<T>& lhs, const T rhs) {
    Simd<T> out;
    for (int i_simd = 0; i_simd<SIMD_SIZE; i_simd++) out.x[i_simd] = lhs.x[i_simd] - rhs;
    return out;
}

template<typename T>
KOKKOS_INLINE_FUNCTION Simd<T> operator - (const T lhs, const Simd<T>& rhs) {
    Simd<T> out;
    for (int i_simd = 0; i_simd<SIMD_SIZE; i_simd++) out.x[i_simd] = lhs - rhs.x[i_simd];
    return out;
}

template<typename T>
KOKKOS_INLINE_FUNCTION Simd<T> operator + (const Simd<T>& lhs, const T rhs) {
    Simd<T> out;
    for (int i_simd = 0; i_simd<SIMD_SIZE; i_simd++) out.x[i_simd] = lhs.x[i_simd] + rhs;
    return out;
}

template<typename T>
KOKKOS_INLINE_FUNCTION Simd<T> operator + (const T lhs, const Simd<T>& rhs) {
    Simd<T> out;
    for (int i_simd = 0; i_simd<SIMD_SIZE; i_simd++) out.x[i_simd] = lhs + rhs.x[i_simd];
    return out;
}

template<typename T>
KOKKOS_INLINE_FUNCTION Simd<T> operator * (const Simd<T>& lhs, const T rhs) {
    Simd<T> out;
    for (int i_simd = 0; i_simd<SIMD_SIZE; i_simd++) out.x[i_simd] = lhs.x[i_simd] * rhs;
    return out;
}

template<typename T>
KOKKOS_INLINE_FUNCTION Simd<T> operator * (const T lhs, const Simd<T>& rhs) {
    Simd<T> out;
    for (int i_simd = 0; i_simd<SIMD_SIZE; i_simd++) out.x[i_simd] = lhs * rhs.x[i_simd];
    return out;
}

template<typename T>
KOKKOS_INLINE_FUNCTION Simd<T> operator / (const Simd<T>& lhs, const T rhs) {
    Simd<T> out;
    for (int i_simd = 0; i_simd<SIMD_SIZE; i_simd++) out.x[i_simd] = lhs.x[i_simd] / rhs;
    return out;
}

template<typename T>
KOKKOS_INLINE_FUNCTION Simd<T> operator / (const T lhs, const Simd<T>& rhs) {
    Simd<T> out;
    for (int i_simd = 0; i_simd<SIMD_SIZE; i_simd++) out.x[i_simd] = lhs / rhs.x[i_simd];
    return out;
}


struct SimdVector2D{
    Simd<double> r;
    Simd<double> z;

    // Get magnitude of a single 2D vector
    KOKKOS_INLINE_FUNCTION double magnitude(const int i_simd) const{
        return sqrt(r[i_simd]*r[i_simd] + z[i_simd]*z[i_simd]);
    }
};

struct SimdVector{
    Simd<double> r;
    Simd<double> z;
    Simd<double> phi;

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

    // Vector operation, sets SimdVector to 0.0
    KOKKOS_INLINE_FUNCTION void zero(){
        for (int i_simd = 0; i_simd<SIMD_SIZE; i_simd++){
            r[i_simd] = 0.0;
            z[i_simd] = 0.0;
            phi[i_simd] = 0.0;
        }
    }

    // Get magnitude of a single vector
    KOKKOS_INLINE_FUNCTION double magnitude(const int i_simd) const{
        return sqrt(r[i_simd]*r[i_simd] + z[i_simd]*z[i_simd] + phi[i_simd]*phi[i_simd]);
    }

    // Get magnitude of a single vector projected onto a poloidal plane
    KOKKOS_INLINE_FUNCTION double poloidal_magnitude(const int i_simd) const{
        return x().magnitude(i_simd);
    }
};

#endif