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


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

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

    return sum;
}

// Finds the max of all elements of a view of any shape
template<typename T>
inline typename T::value_type max_view(const T& view){
    // Cast to 1D with an unmanaged view for a generic element-wise operation
    View<typename T::value_type*,CLayout, typename T::device_type, Kokkos::MemoryTraits<Kokkos::Unmanaged>> view_1D(view.data(), view.size());

    typename T::value_type maxval = 0;
    Kokkos::parallel_reduce("maxval", Kokkos::RangePolicy<typename T::execution_space>( 0,view.size() ), KOKKOS_LAMBDA( const int i, typename T::value_type& maxval_l){
        maxval_l = max(maxval_l, view_1D(i));
    }, Kokkos::Max<typename T::value_type,HostType>(maxval));
    Kokkos::fence();

    return maxval;
}

#endif