operation_state.hpp

Go to the documentation of this file.

#ifndef KOKKOS_EXECUTION_GRAPH_OPERATION_STATE_HPP
#define KOKKOS_EXECUTION_GRAPH_OPERATION_STATE_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;
import kokkos.graph_impl;
#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/get_graph.hpp"
#include "kokkos-execution/graph/get_node.hpp"
#include "kokkos-execution/graph/sender_concepts.hpp"
#include "kokkos-execution/impl/completion_signal.hpp"
#include "kokkos-execution/impl/diagnostics.hpp"
#include "kokkos-execution/impl/dispatch_label.hpp"
#include "kokkos-execution/impl/immovable.hpp"
#include "kokkos-execution/impl/make_op_state.hpp"
#include "kokkos-execution/impl/receiver.hpp"
#include "kokkos-execution/impl/sender_concepts.hpp"
#include "kokkos-execution/impl/submitted.hpp"
#include "kokkos-execution/impl/type_info.hpp"
 
namespace Kokkos::Execution::GraphImpl {
 
template <typename Clsr>
concept Closure = requires {
    typename Clsr::execution_space;
    requires std::same_as<typename Clsr::device_handle_t, Kokkos::Impl::DeviceHandle<typename Clsr::execution_space>>;
    typename Clsr::node_props_t;
};

template <stdexec::operation_state OpState, Kokkos::ExecutionSpace Exec>
requires (
    requires { typename OpState::inner_op_state_t; } &&
    requires(const OpState& op_state) {
        { op_state.query(Kokkos::Execution::GraphImpl::get_node) };
        { op_state.query(Kokkos::Execution::GraphImpl::get_graph) };
    } &&
    requires { typename OpState::execution_space; } &&
    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(stdexec::__is_instance_of<OpState, Kokkos::Execution::GraphImpl::Scheduler<Exec>::template OpState>)
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::Scheduler<Exec>::template OpState>)
struct RemainsOnGraphFor<OpState, Exec> : public std::true_type { };

template <stdexec::operation_state OpState, Kokkos::ExecutionSpace Exec>
requires(graph_operation_state_for<OpState, Exec> && requires { typename OpState::inner_op_state_t; })
struct RemainsOnGraphFor<OpState, Exec> : public RemainsOnGraphFor<typename OpState::inner_op_state_t, Exec> {
 
    using base_t = RemainsOnGraphFor<typename OpState::inner_op_state_t, Exec>;
 
    static constexpr auto diagnose() noexcept {
        if constexpr (base_t::value) {
            return stdexec::__msuccess{};
        } else {
            return stdexec::__mexception<
                stdexec::_WHAT_(OPERATION_STATE_DOES_NOT_REMAIN_ON_GRAPH),
                stdexec::_WHY_(BECAUSE_THE_INNER_OPERATION_STATE_DOES_NOT_REMAIN_ON_GRAPH),
                Impl::diagnose_t<base_t>,
                Impl::WITH_OPSTATE<OpState>,
                Impl::WITH_EXECUTION_SPACE<Exec>
            >{};
        }
    }
};
 
template <typename GraphCompositionPolicy, Kokkos::ExecutionSpace Exec>
struct State;

template <Kokkos::ExecutionSpace Exec>
struct State<GraphComposition::Attach, Exec> {
    using graph_composition_policy_t = GraphComposition::Attach;
 
    using graph_t = Kokkos::Experimental::Graph<Exec>;
 
    explicit State(const Kokkos::Impl::DeviceHandle<Exec>&) {
    }
};

template <Kokkos::ExecutionSpace Exec>
struct State<GraphComposition::Create, Exec> {
    using graph_composition_policy_t = GraphComposition::Create;
 
    using device_handle_t = Kokkos::Impl::DeviceHandle<Exec>;
 
    using graph_t = Kokkos::Experimental::Graph<Exec>;
    using root_t = typename graph_t::root_t;
 
    graph_t graph;
    mutable root_t m_root{};
 
    explicit State(const device_handle_t& device_handle)
        : graph(Kokkos::Execution::GraphImpl::create_graph(device_handle)) {
    }
 
    const device_handle_t& get_device_handle() const {
        return graph.get_device_handle();
    }
 
    const root_t& get_root_node() const {
        if (Kokkos::Impl::GraphAccess::get_node_ptr(m_root) == nullptr) {
            m_root = graph.root_node();
        }
        return m_root;
    }
};

template <Kokkos::ExecutionSpace Exec, stdexec::receiver Rcvr>
struct OpStateBase {
    static consteval auto select_completion_signal_policy() noexcept {
        if constexpr (Impl::supports_submitted_order_on<Rcvr>) {
            return Impl::SubmittedPolicy::OrderOnExec{};
        } else {
            return Impl::SyncPolicy::InlineFenceExec{};
        }
    }
 
    using completion_signal_policy_t = decltype(select_completion_signal_policy());
    using completion_signal_t = Impl::CompletionSignal<completion_signal_policy_t, Exec, Rcvr>;
 
    completion_signal_t completion_signal;
 
    constexpr explicit OpStateBase(Rcvr rcvr) noexcept(std::is_nothrow_constructible_v<completion_signal_t, Rcvr&&>)
        : completion_signal(std::move(rcvr)) {
    }
 
    template <typename Error>
    void complete(stdexec::set_error_t, Error&& error) noexcept {
        stdexec::set_error(std::move(completion_signal.rcvr), std::forward<Error>(error));
    }
 
    void complete(stdexec::set_stopped_t) noexcept {
        stdexec::set_stopped(std::move(completion_signal.rcvr));
    }
};

template <typename Predecessor, Closure FirstClosure>
requires NodeRef<std::remove_cvref_t<Predecessor>>
static auto add_nodes(Predecessor&& predecessor, FirstClosure&& clsr) {
    return std::forward<FirstClosure>(clsr).add_node(std::forward<Predecessor>(predecessor));
}
 
template <typename Predecessor, Closure FirstClosure, Closure... RestOfClosures>
requires NodeRef<std::remove_cvref_t<Predecessor>> && (sizeof...(RestOfClosures) > 0)
static auto add_nodes(Predecessor&& predecessor, FirstClosure&& clsr, RestOfClosures&&... clsrs) {
    return add_nodes(
        std::forward<FirstClosure>(clsr).add_node(std::forward<Predecessor>(predecessor)),
        std::forward<RestOfClosures>(clsrs)...);
}

template <stdexec::sender Sndr, stdexec::receiver Rcvr, Closure FirstClosure, Closure... RestOfClosures>
struct OpState
    : public Impl::Immovable
    , public OpStateBase<typename FirstClosure::execution_space, Rcvr> {
    using operation_state_concept = Impl::SubmittedOperationStateTag;
 
    using execution_space = typename FirstClosure::execution_space;
    using device_handle_t = typename FirstClosure::device_handle_t;

    static_assert((std::same_as<typename RestOfClosures::execution_space, execution_space> && ...));
 
    using rcvr_t = Impl::Receiver<OpState, stdexec::env_of_t<Rcvr>>;
    using inner_op_state_t = stdexec::connect_result_t<Sndr, rcvr_t>;
    using graph_composition_policy_t = GraphComposition::policy_t<inner_op_state_t, Rcvr>;
    using state_t = State<graph_composition_policy_t, execution_space>;
    using predecessor_t = GraphComposition::node_t<execution_space, inner_op_state_t, Rcvr>;
 
    static constexpr bool is_graph_create = std::same_as<graph_composition_policy_t, GraphComposition::Create>;
 
    using node_t = decltype(add_nodes(
        std::declval<predecessor_t>(),
        std::declval<FirstClosure>(),
        std::declval<RestOfClosures>()...));
 
    inner_op_state_t inner_op_state;
    [[no_unique_address]]
    state_t state;
    node_t node;

    constexpr OpState(
        Sndr&& sndr, // NOLINT(cppcoreguidelines-rvalue-reference-param-not-moved)
        Rcvr rcvr,
        FirstClosure clsr,
        RestOfClosures... clsrs) noexcept(false)
        : OpStateBase<execution_space, Rcvr>(std::move(rcvr))
        , inner_op_state(stdexec::connect(std::forward<Sndr>(sndr), rcvr_t{this}))
        , state{Kokkos::Impl::get_property<device_handle_t>(clsr.node_props)}
        , node{add_nodes(this->get_predecessor(), std::move(clsr), std::move(clsrs)...)} {
#if defined(KOKKOS_EXECUTION_ENABLE_DEBUG_LOGGING)
        PLOG_INFO << "Operation state graph composition policy is "
                  << Impl::display_string_of<graph_composition_policy_t>() << ", the receiver is of type "
                  << Impl::display_string_of<Rcvr>() << " and the inner operation state is of type "
                  << Impl::display_string_of<inner_op_state_t>() << '.';
#endif
    }

    [[nodiscard]]
    auto get_predecessor() const noexcept -> const predecessor_t& {
        if constexpr (is_graph_create) {
#if defined(KOKKOS_EXECUTION_ENABLE_DEBUG_LOGGING)
            PLOG_INFO << "The predecessor is the root node of graph " << get_graph_impl_ptr(state.get_root_node())
                      << '.';
#endif
            return state.get_root_node();
        } else if constexpr (stdexec::__queryable_with<inner_op_state_t, get_node_t>) {
            return inner_op_state.query(get_node);
        } else {
            return this->completion_signal.rcvr.query(get_node);
        }
    }
 
    void complete(stdexec::set_value_t) noexcept {
        this->submit();
    }
 
    template <typename Error>
    void complete(stdexec::set_error_t, Error&& error) noexcept {
        stdexec::set_error(std::move(this->completion_signal.rcvr), std::forward<Error>(error));
    }
 
    void complete(stdexec::set_stopped_t) noexcept {
        stdexec::set_stopped(std::move(this->completion_signal.rcvr));
    }
 
    void submit() noexcept {
        if constexpr (is_graph_create) {
#if defined(KOKKOS_EXECUTION_ENABLE_DEBUG_LOGGING)
            PLOG_INFO << "Submitting graph " << get_graph_impl_ptr(state.get_root_node()) << " on "
                      << Kokkos::Tools::Experimental::device_id(state.get_device_handle().m_exec) << '.';
#endif
            try {
                submit_graph(state.graph, state.get_device_handle().m_exec);
            } catch (...) {
                this->complete(stdexec::set_error, std::current_exception());
                return;
            }
        }
        this->completion_signal.propagate(Impl::get_exec(*this).get());
    }
 
    [[nodiscard]]
    auto query(get_node_t) const & noexcept -> const node_t& {
        return node;
    }
 
    [[nodiscard]]
    auto query(get_graph_t) const & noexcept -> const typename state_t::graph_t& {
        if constexpr (is_graph_create) {
            return state.graph;
        } else {
            return inner_op_state.query(get_graph);
        }
    }
 
    [[nodiscard]]
    constexpr auto
        query(Impl::get_exec_t) const noexcept -> Impl::ExecutionSpaceRef<execution_space> requires is_graph_create
    {
        return Impl::ExecutionSpaceRef{state.get_device_handle().m_exec};
    }
 
    [[nodiscard]]
    constexpr auto query(Impl::get_exec_t) const noexcept -> decltype(auto) requires(!is_graph_create)
    {
        return Impl::get_exec(inner_op_state);
    }
 
    void start() & noexcept {
        stdexec::start(inner_op_state);
    }
 
    KOKKOS_EXECUTION_GET_ENV(Rcvr, this->completion_signal.rcvr)
};
 
template <typename Sndr, typename Rcvr, typename... Clsrs>
using make_op_state_t = Impl::MakeOpState<Domain, OpState>::Huddle<Sndr, Rcvr, Clsrs...>;
 
template <typename Sndr, typename Rcvr, typename... Clsrs>
using op_state_t = typename make_op_state_t<Sndr, Rcvr, Clsrs...>::type;
 
#define KOKKOS_EXECUTION_GRAPH_OPERATION_STATE_CONNECT                                                                 \
    template <stdexec::receiver Rcvr>                                                                                  \
    constexpr auto connect(Rcvr rcvr) && noexcept(noexcept(make_op_state_t<Sndr, Rcvr, closure_t>{}(                   \
        std::declval<Sndr>(), std::declval<Rcvr>(), std::declval<closure_t>())))                                       \
        -> op_state_t<Sndr, Rcvr, closure_t> {                                                                         \
        return make_op_state_t<Sndr, Rcvr, closure_t>{}(std::forward<Sndr>(sndr), std::move(rcvr), std::move(clsr));   \
    }
 
#if defined(KOKKOS_EXECUTION_ENABLE_DEBUG_LOGGING) // NOLINTNEXTLINE(cppcoreguidelines-macro-usage)
#    define KOKKOS_EXECUTION_IMPL_GRAPH_ADD_NODE_DEBUG_LOGGING(_type_, _node_, _predecessor_)                          \
        PLOG_INFO << "Adding '" << _type_ << "' node " << get_node_ptr(_node_) << " to graph "                         \
                  << get_graph_impl_ptr(_node_) << " after " << get_node_ptr(_predecessor_) << " on device "           \
                  << Kokkos::Tools::Experimental::device_id(get_node_ptr(_node_)->get_device_handle().m_exec) << '.';
#else
#    define KOKKOS_EXECUTION_IMPL_GRAPH_ADD_NODE_DEBUG_LOGGING(_type_, _node_, _predecessor_)
#endif
 
} // namespace Kokkos::Execution::GraphImpl
 
#endif // KOKKOS_EXECUTION_GRAPH_OPERATION_STATE_HPP