reduce_indices.hpp

Go to the documentation of this file.

#ifndef KOKKOS_EXECUTION_TESTS_UTILS_FUNCTORS_REDUCE_INDICES_HPP
#define KOKKOS_EXECUTION_TESTS_UTILS_FUNCTORS_REDUCE_INDICES_HPP
 
#include <concepts>
#include <type_traits>
 
#include "kokkos-execution/stdexec.hpp"
 
#include "Kokkos_Macros.hpp"
#ifdef KOKKOS_ENABLE_EXPERIMENTAL_CXX20_MODULES
import kokkos.core;
#else
#    include "Kokkos_Core.hpp"
#endif
 
namespace Tests::Utils::Functors {

struct ReduceIndices {
    template <std::integral T, typename U>
    KOKKOS_FUNCTION void operator()(const T index, U& value) const noexcept {
        value += index;
    }
};

template <Kokkos::ExecutionSpace Exec, typename ValueType>
struct ReduceIndicesWithTeam {
    using team_handle_t = typename Kokkos::TeamPolicy<Exec>::member_type;
 
    KOKKOS_FUNCTION void operator()(const team_handle_t& team_handle, ValueType& value) const {
        const auto start_index = team_handle.league_rank() * team_handle.team_size();
        Kokkos::parallel_for(
            Kokkos::TeamThreadRange(team_handle, team_handle.team_size()),
            [&]<std::integral T>(const T index) { value += start_index + index; });
    }
};
 
} // namespace Tests::Utils::Functors
 
#endif // KOKKOS_EXECUTION_TESTS_UTILS_FUNCTORS_REDUCE_INDICES_HPP