operation_state.hpp

Go to the documentation of this file.

#ifndef KOKKOS_EXECUTION_EXECUTION_SPACE_OPERATION_STATE_HPP
#define KOKKOS_EXECUTION_EXECUTION_SPACE_OPERATION_STATE_HPP
 
#include "kokkos-execution/stdexec.hpp"
 
#include "kokkos-execution/execution_space/domain.hpp"
#include "kokkos-execution/execution_space/get_exec.hpp"
#include "kokkos-execution/impl/dispatch_label.hpp"
#include "kokkos-execution/impl/env.hpp"
#include "kokkos-execution/impl/event.hpp"
#include "kokkos-execution/impl/sender_concepts.hpp"
 
namespace Kokkos::Execution::ExecutionSpaceImpl {
 
template <typename Clsr>
concept Closure = requires(const Clsr& clsr) {
    typename Clsr::execution_space;
 
    { clsr.execute() } -> std::same_as<void>;
 
    { clsr.get_policy() };
    requires Kokkos::ExecutionPolicy<std::remove_cvref_t<decltype(clsr.get_policy())>>;
 
    requires std::same_as<std::remove_cvref_t<decltype(clsr.get_policy().space())>, typename Clsr::execution_space>;
};

template <stdexec::receiver Rcvr, typename OpState>
struct RequiresSynchronization {
    static constexpr bool successor_handles_sync = stdexec::__is_instance_of<Rcvr, ScheduleFromReceiver>
                                                || stdexec::__is_instance_of<Rcvr, SyncWaitReceiver>;

    constexpr bool operator()(const OpState&) const noexcept requires(successor_handles_sync)
    {
        return false;
    }

    bool operator()(const OpState& opstate) const noexcept
        requires(!successor_handles_sync && stdexec::__queryable_with<stdexec::env_of_t<Rcvr>, get_exec_t>)
    {
        if constexpr (
            std::same_as<
                std::remove_cvref_t<stdexec::__query_result_t<stdexec::env_of_t<Rcvr>, get_exec_t>>,
                stdexec::__query_result_t<OpState, get_exec_t>
            >) {
            return opstate.query(get_exec).get() != get_exec(stdexec::get_env(opstate.rcvr)).get();
        }
        return true;
    }

    constexpr bool operator()(const OpState&) const noexcept
        requires(!successor_handles_sync && !stdexec::__queryable_with<stdexec::env_of_t<Rcvr>, get_exec_t>)
    {
        return true;
    }
};

template <typename Rcvr, typename Exec>
concept delegate_completion_with_event =
    Kokkos::Execution::Impl::support_events<Exec>
    && stdexec::__queryable_with<stdexec::env_of_t<Rcvr>, stdexec::get_delegation_scheduler_t>;
 
template <
    stdexec::receiver Rcvr,
    Kokkos::ExecutionSpace Exec,
    bool Delegate = delegate_completion_with_event<Rcvr, Exec>
>
struct MayDelegateCompletionWithEvent;
 
template <stdexec::receiver Rcvr, Kokkos::ExecutionSpace Exec>
struct WaitEventReceiver {
    using receiver_concept = stdexec::receiver_t;
    using event_t = Impl::Event<Exec>;
 
    MayDelegateCompletionWithEvent<Rcvr, Exec>* opstate;
    event_t event;
 
    void set_value() && noexcept {
        try {
            event.wait();
            opstate->storage.__destroy();
            stdexec::set_value(std::move(opstate->rcvr));
        } catch (...) {
            opstate->storage.__destroy();
            stdexec::set_error(std::move(opstate->rcvr), std::current_exception());
        }
    }
};
 
template <stdexec::receiver Rcvr, Kokkos::ExecutionSpace Exec>
struct MayDelegateCompletionWithEvent<Rcvr, Exec, false> {
    static constexpr auto label = Impl::dispatch_label<Exec, ": after dispatch">();
 
    Rcvr rcvr;
 
    template <typename OpState>
    void delegate(OpState* const opstate) noexcept {
        if (RequiresSynchronization<Rcvr, OpState>{}(*opstate)) {
            opstate->query(get_exec).get().fence(std::string(label));
        }
        stdexec::set_value(std::move(this->rcvr));
    }
};
 
template <stdexec::receiver Rcvr, Kokkos::ExecutionSpace Exec>
struct MayDelegateCompletionWithEvent<Rcvr, Exec, true> {
    using receiver_t = WaitEventReceiver<Rcvr, Exec>;
    using opstate_t = stdexec::connect_result_t<
        stdexec::schedule_result_t<
            stdexec::__query_result_t<stdexec::env_of_t<Rcvr>, stdexec::get_delegation_scheduler_t>
        >,
        receiver_t
    >;
 
    Rcvr rcvr;
    stdexec::__manual_lifetime<opstate_t> storage{};
 
    template <typename OpState>
    void delegate(OpState* const opstate) noexcept {
        if (RequiresSynchronization<Rcvr, OpState>{}(*opstate)) {
            storage.__construct_from(
                stdexec::connect,
                stdexec::schedule(stdexec::get_delegation_scheduler(stdexec::get_env(this->rcvr))),
                receiver_t{.opstate = opstate, .event = typename receiver_t::event_t{opstate->query(get_exec).get()}});
            stdexec::start(storage.__get());
        } else {
            stdexec::set_value(std::move(this->rcvr));
        }
    }
};
 
template <stdexec::receiver Rcvr, Closure Clsr, Closure... Clsrs>
requires(std::same_as<typename Clsr::execution_space, typename Clsrs::execution_space> && ...)
struct OpStateBase
    : public stdexec::__immovable
    , public MayDelegateCompletionWithEvent<Rcvr, typename Clsr::execution_space> {
    using execution_space = typename Clsr::execution_space;
    using receiver_t = Rcvr;
    using closures_t = stdexec::__tuple<Clsr, Clsrs...>;
 
    using may_delegate_completion_with_event_t = MayDelegateCompletionWithEvent<Rcvr, execution_space>;
 
    closures_t clsrs;
 
    constexpr explicit OpStateBase(Rcvr rcvr_, Clsr clsr_, Clsrs... clsrs_) noexcept(
        std::is_nothrow_constructible_v<may_delegate_completion_with_event_t, Rcvr>
        && std::is_nothrow_move_constructible_v<Clsr> && (std::is_nothrow_move_constructible_v<Clsrs> && ...))
        : may_delegate_completion_with_event_t{std::move(rcvr_)}
        , clsrs(std::move(clsr_), std::move(clsrs_)...) {
    }
 
    void propagate_completion_signal(stdexec::set_value_t) noexcept {
        try {
            stdexec::__apply([](auto&... clsr) { (clsr.execute(), ...); }, clsrs);
        } catch (...) {
            this->propagate_completion_signal(stdexec::set_error, std::current_exception());
            return;
        }
 
        may_delegate_completion_with_event_t::delegate(this);
    }
 
    template <typename Error>
    void propagate_completion_signal(stdexec::set_error_t, Error&& error) noexcept {
        stdexec::set_error(std::move(this->rcvr), std::forward<Error>(error));
    }
 
    void propagate_completion_signal(stdexec::set_stopped_t) noexcept {
        stdexec::set_stopped(std::move(this->rcvr));
    }

    [[nodiscard]]
    constexpr auto query(get_exec_t) const noexcept -> ExecutionSpaceRef<execution_space> {
        return ExecutionSpaceRef<execution_space>{stdexec::__get<0>(clsrs).get_policy().space()};
    }
 
    KOKKOS_EXECUTION_FORWARDING_GET_ENV(Rcvr, this->rcvr)
};
 
template <typename ParentOp>
struct OpStateReceiver {
    using receiver_concept = stdexec::receiver_t;
 
    ParentOp* parent_op;
 
    void set_value() && noexcept {
        parent_op->propagate_completion_signal(stdexec::set_value);
    }
 
    template <typename Error>
    void set_error(Error&& error) && noexcept {
        parent_op->propagate_completion_signal(stdexec::set_error, std::forward<Error>(error));
    }
 
    void set_stopped() && noexcept {
        parent_op->propagate_completion_signal(stdexec::set_stopped);
    }
 
    KOKKOS_EXECUTION_UPSERT_EXEC(
        typename ParentOp::execution_space,
        parent_op->query(get_exec).get(),
        typename ParentOp::receiver_t,
        parent_op->rcvr)
};
 
template <stdexec::sender Sndr, stdexec::receiver Rcvr, Closure... Clsrs>
requires(!Impl::dispatching_sender<Sndr>)
struct OpState : public OpStateBase<Rcvr, Clsrs...> {
    using operation_state_concept = stdexec::operation_state_t;
 
    using base_t = OpStateBase<Rcvr, Clsrs...>;
 
    using inner_opstate_t = stdexec::connect_result_t<Sndr, OpStateReceiver<base_t>>;
 
    static constexpr bool opstate_base_is_nothrow_constructible =
        std::is_nothrow_constructible_v<base_t, Rcvr&&, Clsrs&&...>;
 
    static constexpr bool inner_opstate_is_nothrow_constructible =
        stdexec::__nothrow_connectable<Sndr&&, OpStateReceiver<base_t>>;
 
    inner_opstate_t inner_opstate;
 
    constexpr explicit OpState(
        Sndr&& sndr, // NOLINT(cppcoreguidelines-rvalue-reference-param-not-moved)
        Rcvr rcvr_,
        Clsrs... clsrs_) noexcept(opstate_base_is_nothrow_constructible && inner_opstate_is_nothrow_constructible)
        : base_t(std::move(rcvr_), std::move(clsrs_)...)
        , inner_opstate(stdexec::connect(std::forward<Sndr>(sndr), OpStateReceiver<base_t>{this})) {
    }
 
    void start() & noexcept {
        stdexec::start(inner_opstate);
    }
};
 
template <typename Sndr, typename Rcvr, typename... Clsrs>
struct MakeOpStateFn {
    using type = OpState<Sndr, Rcvr, Clsrs...>;
 
    // NOLINTNEXTLINE(cppcoreguidelines-rvalue-reference-param-not-moved)
    constexpr auto operator()(Sndr&& sndr, Rcvr rcvr, Clsrs... clsrs) const
        noexcept(std::is_nothrow_constructible_v<type, Sndr&&, Rcvr&&, Clsrs&&...>) -> type {
        return type(std::forward<Sndr>(sndr), std::move(rcvr), std::move(clsrs)...);
    }
};
 
template <Impl::dispatching_sender Sndr, typename Rcvr, typename... Clsrs>
struct MakeOpStateFn<Sndr, Rcvr, Clsrs...> {
    using child_of_sndr_t = stdexec::__child_of<Sndr>;
    using clsr_of_sndr_t = typename stdexec::transform_sender_result_t<Sndr, stdexec::env_of_t<Rcvr>>::closure_t;
 
    using make_opstate_fn_t = MakeOpStateFn<child_of_sndr_t, Rcvr, clsr_of_sndr_t, Clsrs...>;
    using type = typename make_opstate_fn_t::type;
 
    static constexpr bool sndr_has_nothrow_transform_sender = stdexec::__detail::__has_nothrow_transform_sender<
        Kokkos::Execution::ExecutionSpaceImpl::Domain,
        stdexec::set_value_t,
        Sndr&&,
        stdexec::env_of_t<Rcvr>
    >;
 
    static constexpr bool is_nothrow_make_opstate =
        std::is_nothrow_invocable_v<make_opstate_fn_t, child_of_sndr_t&&, Rcvr&&, clsr_of_sndr_t&&, Clsrs&&...>;

    // NOLINTNEXTLINE(cppcoreguidelines-rvalue-reference-param-not-moved)
    constexpr auto operator()(Sndr&& sndr, Rcvr&& rcvr, Clsrs... clsrs) const
        noexcept(sndr_has_nothrow_transform_sender && is_nothrow_make_opstate) -> type {
        auto trnsfrmd_sndr = stdexec::transform_sender(std::forward<Sndr>(sndr), stdexec::get_env(rcvr));
        return make_opstate_fn_t{}(
            stdexec::__forward_like<Sndr>(trnsfrmd_sndr.sndr),
            std::forward<Rcvr>(rcvr),
            std::move(trnsfrmd_sndr.clsr),
            std::move(clsrs)...);
    }
};
 
template <typename Sndr, typename Rcvr, typename... Clsrs>
using opstate_t = typename MakeOpStateFn<Sndr, Rcvr, Clsrs...>::type;
 
} // namespace Kokkos::Execution::ExecutionSpaceImpl
 
#endif // KOKKOS_EXECUTION_EXECUTION_SPACE_OPERATION_STATE_HPP