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

// Multiply x by y
template<typename V1, typename V2>
inline void host_multiply_x_by_y(V1& dest_view, const V2& multiplier_view){
    if (dest_view.size() != multiplier_view.size()) exit_XGC("\nWARNING: dest_view and multiplier_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(), multiplier_view.data());
    Kokkos::parallel_for("host_multiply_x_by_y", Kokkos::RangePolicy<HostExSpace>( 0, dest_view.size()), KOKKOS_LAMBDA(const int idx){
        ptrs.x[idx] *= 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. a is converted to a double to take the inverse, so it must be convertible. The view value type must be multipliable by a double.
template<typename V1, typename T>
inline void divide_x_by_a(const V1& dest_view, T coeff){
    double coeff_inv = 1.0/coeff;

    // Cast to 1D with an unmanaged view for a generic element-wise operation
    View<typename V1::value_type*,CLayout, typename V1::device_type, Kokkos::MemoryTraits<Kokkos::Unmanaged>> view_1D(dest_view.data(), dest_view.size());

    // Multiply each element by inverse of coeff
    Kokkos::parallel_for("divide_x_by_a", Kokkos::RangePolicy<typename V1::execution_space>( 0, view_1D.size()), KOKKOS_LAMBDA(const int idx){
        view_1D(idx) *= coeff_inv;
    });
    Kokkos::fence();
}

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

    // Cast to 1D with an unmanaged view for a generic element-wise operation
    View<typename V1::value_type*,CLayout, typename V1::device_type, Kokkos::MemoryTraits<Kokkos::Unmanaged>> dest_view_1D(dest_view.data(), dest_view.size());
    View<typename V2::value_type*,CLayout, typename V2::device_type, Kokkos::MemoryTraits<Kokkos::Unmanaged>> added_view_1D(added_view.data(), added_view.size());

    // Add each element of added_view to dest_view
    Kokkos::parallel_for("add_y_to_x", Kokkos::RangePolicy<typename V1::execution_space>( 0, dest_view_1D.size()), KOKKOS_LAMBDA(const int idx){
        dest_view_1D(idx) += added_view_1D(idx);
    });
    Kokkos::fence();
}

// Add device view to host view. Copies device view to a temporary host view, then adds them
template<typename V1, typename V2>
inline void add_device_view_to_host_view(V1& dest_view, const V2& device_view){
    if (dest_view.size() != device_view.size()) exit_XGC("\nWARNING: dest_view and device_view are not the same size");

    // Create 1D host view and copy device data
    View<double*,CLayout,HostType> view_tmp_h(NoInit("view_tmp_h"), device_view.size());
    View<double*,CLayout,typename V2::device_type, Kokkos::MemoryTraits<Kokkos::Unmanaged>> view_tmp_d(device_view.data(),device_view.size());
    Kokkos::deep_copy(view_tmp_h, view_tmp_d);

    // 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(), view_tmp_h.data());
    Kokkos::parallel_for("host_view_sum", Kokkos::RangePolicy<HostExSpace>( 0, dest_view.size()), KOKKOS_LAMBDA(const int idx){
        ptrs.x[idx] += ptrs.y[idx];
    });
}

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