view_arithmetic.hpp

Go to the documentation of this file.

#ifndef VIEW_ARITHMETIC_HPP
#define VIEW_ARITHMETIC_HPP
#include "space_settings.hpp"
#include "globals.hpp"

// Functions to facilitate element-wise addition and subtraction of Views. Could try overloading views, but this works and is simple enough

// Struct to hold raw device pointers (At some point there was an OMP bug that disallowed raw pointers sent in lambdas
template<typename T>
struct ViewArithmeticPointers{
    T* x;
    T* y;

    ViewArithmeticPointers(T* x, T* y) : x(x), y(y){}
};

// Add added_view*coeff to dest_view
template<typename V1, typename V2, typename T>
inline void host_add_ay_to_x(V1& dest_view, const V2& added_view, T coeff){
    if (dest_view.size() != added_view.size()) exit_XGC("\nWARNING: dest_view and added_view are not the same size");

    // Access raw pointer of View so that we can loop over the whole array without worrying about dimensions
    const ViewArithmeticPointers<T> ptrs(dest_view.data(), added_view.data());
    Kokkos::parallel_for("host_add_ay_to_x", Kokkos::RangePolicy<HostExSpace>( 0, dest_view.size()), KOKKOS_LAMBDA(const int idx){
        ptrs.x[idx] += coeff*ptrs.y[idx];
    });
}

// Divide x by y
template<typename V1, typename V2>
inline void host_divide_x_by_y(V1& dest_view, const V2& divisor_view){
    if (dest_view.size() != divisor_view.size()) exit_XGC("\nWARNING: dest_view and divisor_view are not the same size");

    // Access raw pointer of View so that we can loop over the whole array without worrying about dimensions
    const ViewArithmeticPointers<double> ptrs(dest_view.data(), divisor_view.data());
    Kokkos::parallel_for("host_divide_x_by_y", Kokkos::RangePolicy<HostExSpace>( 0, dest_view.size()), KOKKOS_LAMBDA(const int idx){
        ptrs.x[idx] /= ptrs.y[idx];
    });
}

// Divide x by a
template<typename V1, typename T>
inline void host_divide_x_by_a(V1& dest_view, T coeff){
    double coeff_inv = 1.0/coeff;

    // Access raw pointer of View so that we can loop over the whole array without worrying about dimensions
    const ViewArithmeticPointers<double> ptrs(dest_view.data(), dest_view.data());
    Kokkos::parallel_for("host_divide_x_by_a", Kokkos::RangePolicy<HostExSpace>( 0, dest_view.size()), KOKKOS_LAMBDA(const int idx){
        ptrs.x[idx] *= coeff_inv;
    });
}

template<typename T>
inline double sum_view(const T& view){
    // Cast to 1D with an unmanaged view for a generic element-wise operation
    View<double*,CLayout, typename T::device_type, Kokkos::MemoryTraits<Kokkos::Unmanaged>> view_1D(view.data(), view.size());

    double sum = 0.0;
    Kokkos::parallel_reduce("sum", Kokkos::RangePolicy<typename T::execution_space>( 0,view.size() ), KOKKOS_LAMBDA( const int i, double& sum_l){
        sum_l += view_1D(i);
    }, sum);
    Kokkos::fence();

    return sum;
}

#endif