AMReX_CTOParallelForImpl.H

Go to the documentation of this file.

#ifndef AMREX_CTO_PARALLEL_FOR_H_
#define AMREX_CTO_PARALLEL_FOR_H_
 
#include <AMReX_BLassert.H>
#include <AMReX_Box.H>
#include <AMReX_TypeList.H>
 
#include <array>
#include <concepts>
#include <type_traits>
#include <utility>
 
/* This header is not for the users to include directly.  It's meant to be
 * included in AMReX_GpuLaunch.H, which has included the headers needed
 * here. */
 
/* Thank Maikel Nadolski and Alex Sinn for the techniques used here! */
 
namespace amrex {
 
template <int... ctr>
struct CompileTimeOptions {
    // TypeList is defined in AMReX_TypeList.H
    using list_type = TypeList<std::integral_constant<int, ctr>...>;
};
 
namespace detail
{
    template<class F, int... ctr>
    struct CTOWrapper {
        F f;
 
        template<class... Args>
        AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
        auto operator() (Args... args) const noexcept
            -> decltype(f(args..., std::integral_constant<int, ctr>{}...)) {
            return f(args..., std::integral_constant<int, ctr>{}...);
        }
 
        AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
        static constexpr
        std::array<int, sizeof...(ctr)> GetOptions () noexcept {
            return {ctr...};
        }
    };
 
    template <class L, typename... As, class... Fs>
    bool
    AnyCTO_helper2 (const L& l, TypeList<As...>,
                    std::array<int,sizeof...(As)> const& runtime_options, const Fs&...cto_functs)
    {
        if (runtime_options == std::array<int,sizeof...(As)>{As::value...}) {
            if constexpr (sizeof...(cto_functs) != 0) {
                // Apply the CTOWrapper to each function that was given in cto_functs
                // and call the CPU function l with all of them
                l(CTOWrapper<Fs, As::value...>{cto_functs}...);
            } else {
                // No functions in cto_functs so we call l directly with the compile time arguments
                l(As{}...);
            }
            return true;
        } else {
            return false;
        }
    }
 
    template <class L, typename... PPs, typename RO, class...Fs>
    void
    AnyCTO_helper1 (const L& l, TypeList<PPs...>,
                    RO const& runtime_options, const Fs&...cto_functs)
    {
        bool found_option = (false || ... ||
                                AnyCTO_helper2(l, PPs{}, runtime_options, cto_functs...));
        amrex::ignore_unused(found_option);
        AMREX_ASSERT(found_option);
    }
 
    template <int Begin, int End, int... I>
    constexpr auto MakeCTOImpl(std::integer_sequence<int, I...>)
        -> CompileTimeOptions<Begin + I...>;
}
 
template <int Begin, int End>
using CTOSeq = decltype(detail::MakeCTOImpl<Begin, End>(
                            std::make_integer_sequence<int, End - Begin>{}));
 
template <class L, class... Fs, typename... CTOs>
void AnyCTO ([[maybe_unused]] TypeList<CTOs...> list_of_compile_time_options,
             std::array<int,sizeof...(CTOs)> const& runtime_options,
             L&& l, Fs&&...cto_functs)
{
    detail::AnyCTO_helper1(std::forward<L>(l),
                           CartesianProduct(typename CTOs::list_type{}...),
                           runtime_options,
                           std::forward<Fs>(cto_functs)...);
}
 
template <int MT, std::integral T, class F, typename... CTOs>
void ParallelFor (TypeList<CTOs...> ctos,
             std::array<int,sizeof...(CTOs)> const& runtime_options,
             T N, F&& f)
{
    AnyCTO(ctos, runtime_options,
        [&](auto cto_func){
            ParallelFor<MT>(N, cto_func);
        },
        std::forward<F>(f)
    );
}
 
template <int MT, class F, int dim, typename... CTOs>
void ParallelFor (TypeList<CTOs...> ctos,
                  std::array<int,sizeof...(CTOs)> const& runtime_options,
                  BoxND<dim> const& box, F&& f)
{
    AnyCTO(ctos, runtime_options,
        [&](auto cto_func){
            ParallelFor<MT>(box, cto_func);
        },
        std::forward<F>(f)
    );
}
 
template <int MT, std::integral T, class F, int dim, typename... CTOs>
void ParallelFor (TypeList<CTOs...> ctos,
             std::array<int,sizeof...(CTOs)> const& runtime_options,
             BoxND<dim> const& box, T ncomp, F&& f)
{
    AnyCTO(ctos, runtime_options,
        [&](auto cto_func){
            ParallelFor<MT>(box, ncomp, cto_func);
        },
        std::forward<F>(f)
    );
}
 
template <std::integral T, class F, typename... CTOs>
void ParallelFor (TypeList<CTOs...> ctos,
             std::array<int,sizeof...(CTOs)> const& option,
             T N, F&& f)
{
    ParallelFor<AMREX_GPU_MAX_THREADS>(ctos, option, N, std::forward<F>(f));
}
 
template <class F, int dim, typename... CTOs>
void ParallelFor (TypeList<CTOs...> ctos,
                  std::array<int,sizeof...(CTOs)> const& option,
                  BoxND<dim> const& box, F&& f)
{
    ParallelFor<AMREX_GPU_MAX_THREADS>(ctos, option, box, std::forward<F>(f));
}
 
template <std::integral T, class F, int dim, typename... CTOs>
void ParallelFor (TypeList<CTOs...> ctos,
             std::array<int,sizeof...(CTOs)> const& option,
             BoxND<dim> const& box, T ncomp, F&& f)
{
    ParallelFor<AMREX_GPU_MAX_THREADS>(ctos, option, box, ncomp, std::forward<F>(f));
}
 
}
 
#endif