access_add.hpp

Go to the documentation of this file.

#ifndef ACCESS_ADD_HPP
#define ACCESS_ADD_HPP
#include "space_settings.hpp"

/* ScatterType template specifies whether the class uses atomics or array replication
 * */
enum class ScatterType{
    Atomic,
    Replication
};

#ifdef USE_GPU
# ifdef USE_ARRAY_REPLICATION
#  error Cannot use ARRAY REPLICATION on GPU
# endif
#else
# ifndef USE_ARRAY_REPLICATION
#  error Must use ARRAY REPLICATION if executing on CPU for now
# endif
#endif

#ifdef USE_ARRAY_REPLICATION
constexpr ScatterType SCATTER_TYPE_GLOBAL = ScatterType::Replication; 
#else
constexpr ScatterType SCATTER_TYPE_GLOBAL = ScatterType::Atomic;
#endif

// returns omp thread if on CPU (using array replication strategy for scatter) or
// returns 0 on GPU (uses atomics, no replication needed)
KOKKOS_INLINE_FUNCTION int get_thread(){
#if defined(USE_ARRAY_REPLICATION) && defined(USE_OMP)
    return omp_get_thread_num();
#else
    return 0;
#endif
}

// access_add (could use scatter view instead): Uses atomic if array replication is off
template<typename T>
KOKKOS_INLINE_FUNCTION void access_add(T* addr, T val){
#ifdef USE_ARRAY_REPLICATION
    *addr += val;
#else
    Kokkos::atomic_add(addr, val);
#endif
}

/* Sums the contents of a View into the 0th index of its first dimension if USE_ARRAY_REPLICATION is on
 */
template<typename T>
void reduce_replicated_array(T& view){
#ifdef USE_ARRAY_REPLICATION
    int n_threads = view.extent(0);
    int size_per_thread = view.size()/n_threads;

    auto thread_0_ptr = view.data();
    auto thread_i_ptr = view.data();

    for(int i = 1; i<n_threads; i++){
        thread_i_ptr += size_per_thread;
        Kokkos::parallel_for("reduce_replicated_array", Kokkos::RangePolicy<HostExSpace>( 0, size_per_thread), KOKKOS_LAMBDA(const int idx){
            thread_0_ptr[idx] += thread_i_ptr[idx];
        });
    }
#endif
}

#endif