when_all.hpp

Go to the documentation of this file.

#ifndef KOKKOS_EXECUTION_GRAPH_WHEN_ALL_HPP
#define KOKKOS_EXECUTION_GRAPH_WHEN_ALL_HPP
 
#include "kokkos-execution/stdexec.hpp"
 
#if defined(KOKKOS_EXECUTION_ENABLE_DEBUG_LOGGING)
#    include "plog/Log.h"
#endif
 
#include "Kokkos_Macros.hpp"
#ifdef KOKKOS_ENABLE_EXPERIMENTAL_CXX20_MODULES
import kokkos.core;
import kokkos.graph;
#else
#    include "Kokkos_Core.hpp"
#    include "Kokkos_Graph.hpp"
#endif
 
#include "kokkos-execution/graph/graph_fwd.hpp"
 
#include "kokkos-execution/graph/events.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/compile_time_optional.hpp"
#include "kokkos-execution/impl/completion_signatures.hpp"
#include "kokkos-execution/impl/diagnostics.hpp"
#include "kokkos-execution/impl/dispatch_label.hpp"
#include "kokkos-execution/impl/env.hpp"
#include "kokkos-execution/impl/macros.hpp"
#include "kokkos-execution/impl/queryable.hpp"
#include "kokkos-execution/impl/sender_concepts.hpp"
#include "kokkos-execution/impl/sender_introspection.hpp"
 
namespace Kokkos::Execution::GraphImpl {

template <Kokkos::ExecutionSpace Exec, stdexec::receiver Rcvr, stdexec::sender... Sndrs>
struct WhenAllOpState
    : public Impl::Immovable
    , public OpStateBase<Exec, Rcvr> {
    using operation_state_concept = stdexec::operation_state_tag;
 
    using base_t = OpStateBase<Exec, Rcvr>;
    using execution_space = Exec;
 
    using predecessor_t = GraphComposition::node_t<Exec, Rcvr>;

    struct WhenAllChildReceiver : public Impl::Receiver<WhenAllOpState, stdexec::env_of_t<Rcvr>> {
        [[nodiscard]]
        auto query(get_node_t) const & noexcept -> const predecessor_t& {
            return this->parent_op->get_predecessor();
        }
    };
 
    using children_op_states_t = stdexec::__tuple<stdexec::connect_result_t<Sndrs, WhenAllChildReceiver>...>;

    static constexpr bool as_one =
        (connect_result_remains_on_graph_for<execution_space, Sndrs, WhenAllChildReceiver> && ...);

    static constexpr bool at_least_one_child_adds_a_node = stdexec::__mapply<
        stdexec::__many_of<stdexec::__q<Impl::queryable_for<get_node_t>::type>>,
        children_op_states_t
    >::value;

    static constexpr std::size_t num_branches = sizeof...(Sndrs);

    static constexpr std::size_t num_branches_remaining_on_graph = stdexec::__mapply<
        stdexec::__mcount_if<stdexec::__mbind_back_q<RemainsOnGraphFor, Exec>>,
        children_op_states_t
    >::value;
    static constexpr std::size_t num_branches_not_remaining_on_graph = num_branches - num_branches_remaining_on_graph;
 
    static_assert(
        (as_one && num_branches_not_remaining_on_graph == 0) || (!as_one && num_branches_not_remaining_on_graph > 0));

    struct JoinOpStateNodes {
        template <typename OpState>
        static constexpr auto node_of(const OpState& opstate) noexcept {
            if constexpr (stdexec::__queryable_with<OpState, get_node_t>) {
                using query_t = stdexec::__query_result_t<OpState, get_node_t>;
                static_assert(std::is_lvalue_reference_v<query_t>);
                return std::tuple<query_t>{opstate.query(get_node)};
            } else {
                return std::tuple<>{};
            }
        }
 
        static constexpr auto join() noexcept {
            return Impl::Empty{};
        }
 
        template <typename Node>
        static constexpr const Node& join(const Node& node) noexcept {
            return node;
        }
 
        template <typename... Nodes>
        requires(sizeof...(Nodes) >= 2)
        static constexpr auto join(const Nodes&... nodes) {
            auto agg = Kokkos::Experimental::when_all(nodes...);
            graph_add_aggregate_node_event(agg, nodes...);
            return agg;
        }
 
        template <typename ChildrenOpStates>
        constexpr decltype(auto) operator()(const ChildrenOpStates& children) const {
            return stdexec::__apply(
                [](const auto&... opstates) -> decltype(auto) {
                    return std::apply(
                        [](const auto&... nodes) -> decltype(auto) { return join(nodes...); },
                        std::tuple_cat(node_of(opstates)...));
                },
                children);
        }
    };
 
    using node_t = std::invoke_result_t<JoinOpStateNodes, const children_op_states_t&>;

    using graph_composition_policy_t =
        std::conditional_t<at_least_one_child_adds_a_node, GraphComposition::policy_t<Rcvr>, GraphComposition::Attach>;
    using state_t = State<graph_composition_policy_t, execution_space>;
 
    [[no_unique_address]]
    state_t state;
    children_op_states_t children_op_states;
    PRAGMA_DIAGNOSTIC_PUSH
    KOKKOS_EXECUTION_PRAGMA_DIAGNOSTIC_IGNORED_DANGLING_REFERENCE
    [[no_unique_address]]
    node_t node;
    PRAGMA_DIAGNOSTIC_POP
    [[no_unique_address]]
    Impl::CompileTimeOptional<!as_one, std::atomic<size_t>> count{num_branches_not_remaining_on_graph};

    WhenAllOpState(stdexec::__tuple<Sndrs...>&& sndrs_, Rcvr&& rcvr_)
        : base_t(std::move(rcvr_))
        , state{Kokkos::Experimental::get_device_handle(execution_space{})}
        , children_op_states(
              stdexec::__apply(
                  [this](auto&&... children) -> children_op_states_t {
                      static_assert((std::same_as<decltype(children), Sndrs&&> && ...));
                      PRAGMA_DIAGNOSTIC_PUSH
                      PRAGMA_DIAGNOSTIC_IGNORED("-Wmissing-braces")
                      return children_op_states_t{
                          stdexec::connect(std::forward<Sndrs>(children), WhenAllChildReceiver{this})...};
                      PRAGMA_DIAGNOSTIC_POP
                  },
                  std::move(sndrs_)))
        , node(JoinOpStateNodes{}(children_op_states)) {
    }
 
    const predecessor_t& get_predecessor() const noexcept {
        if constexpr (std::same_as<graph_composition_policy_t, GraphComposition::Create>) {
            return state.get_root_node();
        } else {
            return this->completion_signal.rcvr.query(get_node);
        }
    }
 
    [[nodiscard]]
    auto query(get_node_t) const & noexcept -> const node_t& requires(!Impl::empty<node_t>)
    {
        return node;
    }
 
    [[nodiscard]]
    auto query(get_graph_t) const & noexcept -> const
        typename state_t::graph_t& requires std::same_as<graph_composition_policy_t, GraphComposition::Create>
    {
        return state.graph;
    }

    void submit() & noexcept requires(!as_one)
    {
        if (count.value.fetch_sub(1) == 1) {
            this->submit_graph();
        }
    }
 
    void submit() & noexcept requires(as_one)
    {
        this->submit_graph();
    }
 
    void complete(stdexec::set_value_t) & noexcept {
        this->submit();
    }
 
    template <typename Tag, typename... Args>
    requires(!std::same_as<Tag, stdexec::set_value_t>)
    void complete(Tag, Args&&... args) & noexcept {
        base_t::complete(Tag{}, std::forward<Args>(args)...);
    }
 
    void start() & noexcept requires(!as_one)
    {
#if defined(KOKKOS_EXECUTION_ENABLE_DEBUG_LOGGING)
        PLOG_INFO << "Starting all branches before submission.";
#endif
#if defined(__cpp_expansion_statements) && __cpp_expansion_statements >= 202506L
        template for (auto& op: children_op_states) {
            if constexpr (!remains_on_graph_for<std::remove_cvref_t<decltype(op)>, Exec>)
                stdexec::start(op);
        }
#else
        stdexec::__apply(
            []<typename... ChildOpStates>([[maybe_unused]] ChildOpStates&... ops) noexcept {
                (
                    [&] {
                        if constexpr (!remains_on_graph_for<ChildOpStates, Exec>)
                            stdexec::start(ops);
                    }(),
                    ...);
            },
            children_op_states);
#endif
    }

    void start() & noexcept requires(as_one)
    {
#if defined(KOKKOS_EXECUTION_ENABLE_DEBUG_LOGGING)
        PLOG_INFO << "Submit the graph directly without starting the branches.";
#endif
        this->submit();
    }
 
    void submit_graph() & noexcept {
        stdexec::set_value(std::move(this->completion_signal.rcvr));
    }
 
    void submit_graph() & noexcept
        requires(std::is_same_v<graph_composition_policy_t, GraphComposition::Create> && at_least_one_child_adds_a_node)
    {
#if defined(KOKKOS_EXECUTION_ENABLE_DEBUG_LOGGING)
        PLOG_INFO << "Submitting graph " << get_graph_impl_ptr(state.graph.root_node()) << " on "
                  << Kokkos::Tools::Experimental::device_id(state.get_device_handle().m_exec) << '.';
#endif
        try {
            Kokkos::Execution::GraphImpl::submit_graph(state.graph, state.get_device_handle().m_exec);
        } catch (...) {
            stdexec::set_error(std::move(this->completion_signal.rcvr), std::current_exception());
            return;
        }
        this->completion_signal.propagate(state.get_device_handle().m_exec);
    }
 
    KOKKOS_EXECUTION_GET_ENV(Rcvr, this->completion_signal.rcvr)
};

template <stdexec::operation_state OpState, Kokkos::ExecutionSpace Exec>
requires(
    stdexec::__is_instance_of<OpState, Kokkos::Execution::GraphImpl::WhenAllOpState>
    && std::same_as<typename OpState::execution_space, Exec>)
struct GraphOperationStateFor<OpState, Exec> : public std::true_type { };

template <stdexec::operation_state OpState, Kokkos::ExecutionSpace Exec>
requires(
    graph_operation_state_for<OpState, Exec>
    && stdexec::__is_instance_of<OpState, Kokkos::Execution::GraphImpl::WhenAllOpState>)
struct RemainsOnGraphFor<OpState, Exec> {
    template <stdexec::operation_state ChildOpState>
    using RemainsOnGraphForChild = RemainsOnGraphFor<ChildOpState, Exec>;
 
    static constexpr bool value = stdexec::__mapply<
        stdexec::__mall_of<stdexec::__q<RemainsOnGraphForChild>>,
        typename OpState::children_op_states_t
    >::value;
 
    static constexpr auto diagnose() noexcept {
        using children_t = stdexec::__mapply<
            stdexec::__mtransform<stdexec::__q<RemainsOnGraphForChild>, stdexec::__q<Impl::diagnose_t>>,
            typename OpState::children_op_states_t
        >;
 
        if constexpr (stdexec::__ok<Impl::first_error_t<children_t>>) {
            return stdexec::__msuccess{};
        } else {
            return stdexec::__mexception<
                stdexec::_WHAT_(OPERATION_STATE_DOES_NOT_REMAIN_ON_GRAPH),
                stdexec::_WHY_(BECAUSE_AT_LEAST_ONE_CHILD_DOES_NOT_REMAIN_ON_GRAPH),
                Impl::rebind_t<Impl::WITH_CHILDREN_ERRORS, Impl::keep_errors_t<children_t>>,
                Impl::WITH_OPSTATE<OpState>
            >{};
        }
    }
};

template <Kokkos::ExecutionSpace Exec, stdexec::sender... Sndrs>
struct WhenAllSender {
    using sender_concept = stdexec::sender_tag;
 
    using sndrs_t = stdexec::__tuple<Sndrs...>;
 
    struct attrs {
        template <typename... Env>
        [[nodiscard]]
        constexpr auto
            query(stdexec::get_completion_domain_t<stdexec::set_value_t>, const Env&...) const noexcept -> Domain {
            return {};
        }
    };

    template <typename Self, typename... Env>
    static consteval auto get_completion_signatures() {
        return stdexec::completion_signatures<stdexec::set_value_t(), stdexec::set_error_t(std::exception_ptr)>{};
    }
 
    template <stdexec::receiver Rcvr>
        auto connect(Rcvr rcvr)
        && KOKKOS_EXECUTION_RETURNS_BRACED((WhenAllOpState<Exec, Rcvr, Sndrs...>), std::move(sndrs), std::move(rcvr));
 
    constexpr auto get_env() const noexcept -> attrs {
        return {};
    }
 
    sndrs_t sndrs;
};
 
struct BECAUSE_THE_EXECUTION_SPACE_TYPE_IS_NOT_HOMOGENEOUS;
 
template <size_t Index, typename Sndr>
struct WITH_SENDER_AT_INDEX { };
 
template <>
struct TransformSenderFor<stdexec::when_all_t> {
    template <typename Env, typename... Sndrs>
    using trnsfrmd_sndr_t = WhenAllSender<Impl::exec_of_t<stdexec::__m_at_c<0, Sndrs...>, Env>, Sndrs...>;
 
    template <typename Env, typename... Sndrs>
    auto operator()(const Env&, stdexec::when_all_t, stdexec::__ignore, Sndrs&&... sndrs) const
        noexcept(std::is_nothrow_constructible_v<typename trnsfrmd_sndr_t<Env, Sndrs...>::sndrs_t, Sndrs&&...>) {
        if constexpr ((graph_completing_sender<Sndrs, Env> && ...)) {
            using execution_space = Impl::exec_of_t<stdexec::__m_at_c<0, Sndrs...>, Env>;

            if constexpr ((std::same_as<Impl::exec_of_t<Sndrs, Env>, execution_space> && ...)) {
                PRAGMA_DIAGNOSTIC_PUSH
                PRAGMA_DIAGNOSTIC_IGNORED("-Wmissing-braces")
                return trnsfrmd_sndr_t<Env, Sndrs...>{.sndrs = {std::forward<Sndrs>(sndrs)...}};
                PRAGMA_DIAGNOSTIC_POP
            } else {
                STDEXEC_CONSTEXPR_LOCAL bool map[] = {
                    !std::same_as<Impl::exec_of_t<stdexec::__m_at_c<0, Sndrs>, Env>, execution_space>...};
                STDEXEC_CONSTEXPR_LOCAL std::size_t index = stdexec::__pos_of(map, map + sizeof...(Sndrs));
                using invalid_sndr_t = stdexec::__m_at_c<index, Sndrs...>;
                return stdexec::__not_a_sender<
                    stdexec::_WHAT_(CANNOT_DISPATCH_THIS_ALGORITHM_TO_THE_GRAPH_SCHEDULER),
                    stdexec::_WHY_(BECAUSE_THE_EXECUTION_SPACE_TYPE_IS_NOT_HOMOGENEOUS),
                    stdexec::_WHERE_(stdexec::_IN_ALGORITHM_, stdexec::when_all_t),
                    WITH_SENDER_AT_INDEX<index, invalid_sndr_t>,
                    stdexec::_WITH_SENDERS_<Sndrs...>,
                    stdexec::_WITH_ENVIRONMENT_(Env)
                >{};
            }
        } else {
            STDEXEC_CONSTEXPR_LOCAL bool map[] = {!graph_completing_sender<Sndrs, Env>...};
            STDEXEC_CONSTEXPR_LOCAL std::size_t index = stdexec::__pos_of(map, map + sizeof...(Sndrs));
            using invalid_sndr_t = stdexec::__m_at_c<index, Sndrs...>;
            return no_graph_scheduler_in_env<stdexec::when_all_t, invalid_sndr_t, Env>();
        }
    }
};
 
} // namespace Kokkos::Execution::GraphImpl
 
#endif // KOKKOS_EXECUTION_GRAPH_WHEN_ALL_HPP