parallel_for.hpp

Go to the documentation of this file.

#ifndef KOKKOS_EXECUTION_GRAPH_PARALLEL_FOR_HPP
#define KOKKOS_EXECUTION_GRAPH_PARALLEL_FOR_HPP
 
#include "kokkos-execution/stdexec.hpp"
 
#include "kokkos-execution/graph/diagnostics.hpp"
#include "kokkos-execution/graph/node_props.hpp"
#include "kokkos-execution/graph/operation_state.hpp"
#include "kokkos-execution/graph/sender_concepts.hpp"
#include "kokkos-execution/impl/attributes.hpp"
#include "kokkos-execution/impl/make_policy.hpp"
#include "kokkos-execution/impl/sender_introspection.hpp"
#include "kokkos-execution/parallel_for.hpp"
 
namespace Kokkos::Execution::GraphImpl {
 
template <typename NodeProps, typename Functor, Impl::execution_policy_or_iteration_space ExecPolicyOrIterSpace>
struct ParallelForClosure {
    using node_props_t = NodeProps;
    using device_handle_t = device_handle_of_t<NodeProps>;
    using execution_space = exec_of_t<device_handle_t>;
    using policy_t = Impl::make_policy_result_t<Impl::make_policy_t<execution_space>, ExecPolicyOrIterSpace&&>;
 
    template <NodeRef Predecessor>
    using node_t = decltype(std::declval<const Predecessor&>().then_parallel_for(
        std::declval<node_props_t&&>(),
        std::declval<policy_t&&>(),
        std::declval<Functor&&>()));

    template <NodeRef Predecessor>
    auto add_node(const Predecessor& predecessor) && noexcept(false) -> node_t<Predecessor> {
#if defined(KOKKOS_EXECUTION_ENABLE_EVENT_DISPATCH) || defined(KOKKOS_EXECUTION_ENABLE_DEBUG_LOGGING)
        const auto label = get_label_or_display_string_of<Functor>(node_props);
#endif
        auto node = predecessor.then_parallel_for(
            std::move(node_props),
            Impl::make_policy<execution_space>(std::move(policy_or_iter_space)),
            std::move(functor));
#if defined(KOKKOS_EXECUTION_ENABLE_EVENT_DISPATCH) || defined(KOKKOS_EXECUTION_ENABLE_DEBUG_LOGGING)
        KOKKOS_EXECUTION_IMPL_GRAPH_ADD_NODE_DEBUG_LOGGING(label, node, predecessor)
        graph_add_node_event(std::move(label), predecessor, node);
#endif
        return node;
    }
 
    node_props_t node_props;
    Functor functor;
    ExecPolicyOrIterSpace policy_or_iter_space;
};
 
template <
    stdexec::sender Sndr,
    typename NodeProps,
    typename Functor,
    Impl::execution_policy_or_iteration_space ExecPolicyOrIterSpace
>
struct ParallelForSender {
    using sender_concept = stdexec::sender_tag;
 
    using closure_t = ParallelForClosure<NodeProps, Functor, ExecPolicyOrIterSpace>;
 
    KOKKOS_EXECUTION_COMPL_SIGS_ADD(ParallelForSender, Sndr, stdexec::set_error_t(std::exception_ptr))
 
    KOKKOS_EXECUTION_GRAPH_OPERATION_STATE_CONNECT
 
    KOKKOS_EXECUTION_IMPL_FORWARDING_ATTRIBUTES_GET_ENV(Sndr, sndr)
 
    closure_t clsr;
    Sndr sndr; // NOLINT(cppcoreguidelines-avoid-const-or-ref-data-members)
};
 
template <>
struct TransformSenderFor<Kokkos::Execution::parallel_for_t> {
    template <typename Env, typename Data, typename Sndr>
    using trnsfrmd_sndr_t = ParallelForSender<
        Sndr,
        make_node_props_t<Impl::exec_of_t<Sndr, Env>, Data>,
        typename std::remove_cvref_t<Data>::functor_t,
        typename std::remove_cvref_t<Data>::policy_or_iter_space_t
    >;
 
    template <typename Env, typename Data, typename Sndr>
    requires(stdexec::__sends<stdexec::set_value_t, Sndr, Env> && graph_completing_sender<Sndr, Env>)
    auto operator()(
        const Env& env,
        Kokkos::Execution::parallel_for_t,
        Data&& data, // NOLINT(cppcoreguidelines-missing-std-forward)
        Sndr&& sndr) const
        noexcept(std::is_nothrow_constructible_v<
                 trnsfrmd_sndr_t<Env, Data, Sndr>,
                 typename trnsfrmd_sndr_t<Env, Data, Sndr>::closure_t&&,
                 Sndr&&
        >) {
        const auto schd = stdexec::get_completion_scheduler<stdexec::set_value_t>(stdexec::get_env(sndr), env);
 
        return trnsfrmd_sndr_t<Env, Data, Sndr>{
            .clsr =
                {.node_props = make_node_props(schd.state->exec, std::forward<Data>(data)),
                       .functor = stdexec::__forward_like<Data>(data.functor), // NOLINT(bugprone-use-after-move)
                 .policy_or_iter_space = stdexec::__forward_like<Data>(data.policy_or_iter_space)},
            .sndr = std::forward<Sndr>(sndr)
        };
    }
 
    template <typename Env, typename Data, typename Sndr>
    requires stdexec::__sends<stdexec::set_value_t, Sndr, Env>
    auto operator()(const Env&, Kokkos::Execution::parallel_for_t, Data&&, Sndr&&) const noexcept {
        return no_graph_scheduler_in_env<parallel_for_t, Sndr, Env>();
    }
};
 
} // namespace Kokkos::Execution::GraphImpl
 
#endif // KOKKOS_EXECUTION_GRAPH_PARALLEL_FOR_HPP