sync_wait.hpp

Go to the documentation of this file.

#ifndef KOKKOS_EXECUTION_IMPL_SYNC_WAIT_HPP
#define KOKKOS_EXECUTION_IMPL_SYNC_WAIT_HPP
 
#include "kokkos-execution/stdexec.hpp"
 
#include "kokkos-execution/impl/completion_signal.hpp"
#include "kokkos-execution/impl/dispatch_label.hpp"
#include "kokkos-execution/impl/event.hpp"
#include "kokkos-execution/impl/sender_introspection.hpp"
#include "kokkos-execution/impl/state.hpp"
 
namespace Kokkos::Execution::Impl::SyncWait {

struct env {
    stdexec::run_loop::scheduler schd;
 
    [[nodiscard]]
    auto query(stdexec::get_scheduler_t) const noexcept -> stdexec::run_loop::scheduler {
        return schd;
    }
 
    [[nodiscard]]
    auto query(stdexec::get_delegation_scheduler_t) const noexcept -> stdexec::run_loop::scheduler {
        return schd;
    }
};

template <typename HasErrorPtr>
struct State;
 
template <>
struct State<std::true_type> {
    std::exception_ptr error;
    stdexec::run_loop loop;
};
 
template <>
struct State<std::false_type> {
    stdexec::run_loop loop;
};

template <Kokkos::ExecutionSpace Exec, typename HasErrorPtr = std::false_type, typename ResultType = std::tuple<>>
struct Receiver {
    using receiver_concept = SubmittedReceiverTag;
 
    static constexpr auto label = Impl::dispatch_label<Exec, ": sync_wait">();
 
    Impl::State<Exec> const * state;
    State<HasErrorPtr>* runloop_state;
    std::optional<ResultType>* result;
 
    template <typename... Args>
    void set_value(Args&&... args) && noexcept {
        result->emplace(std::forward<Args>(args)...);
        runloop_state->loop.finish();
    }
 
    template <typename Error>
    void set_error(Error&& err) && noexcept requires HasErrorPtr::value
    {
        runloop_state->error = std::forward<Error>(err);
        state->exec.fence(std::string(label));
        runloop_state->loop.finish();
    }
 
    void set_stopped() && noexcept {
        state->exec.fence(std::string(label));
        runloop_state->loop.finish();
    }
 
    void submitted() && noexcept {
        state->exec.fence(std::string(label));
        result->emplace();
        runloop_state->loop.finish();
    }
 
    void submitted(OptionalConstEventRef<Exec> dep) && noexcept {
        if (dep.has_value()) {
            Impl::wait(dep.get());
        }
        result->emplace();
        runloop_state->loop.finish();
    }
 
    [[nodiscard]]
    constexpr auto get_env() const noexcept -> env {
        return {.schd = runloop_state->loop.get_scheduler()};
    }
};
 
struct SyncWait {
    template <typename Sndr>
    using sends_error = std::bool_constant<stdexec::__sends<stdexec::set_error_t, Sndr, env>>;
 
    template <typename Sndr>
    using result_t = stdexec::__sync_wait::__value_tuple_for_t<Sndr>;
 
    template <typename Sndr>
    using receiver_t = Receiver<Impl::exec_of_t<Sndr, env>, sends_error<Sndr>, result_t<Sndr>>;
 
    template <typename Sndr>
    static constexpr bool is_nothrow_connectable = stdexec::__nothrow_connectable<Sndr, receiver_t<Sndr>>;

    template <stdexec::sender Sndr>
    auto operator()(Sndr&& sndr) const noexcept(!sends_error<Sndr&&>::value && is_nothrow_connectable<Sndr&&>)
        -> std::optional<result_t<Sndr&&>> {
        static_assert(noexcept(std::declval<State<sends_error<Sndr&&>>>().loop.run()));
 
        State<sends_error<Sndr&&>> runloop_state;
 
        std::optional<result_t<Sndr&&>> result{};
 
        Receiver rcvr{
            .state = stdexec::get_completion_scheduler<stdexec::set_value_t>(
                         stdexec::get_env(sndr), env{.schd = runloop_state.loop.get_scheduler()})
                         .state,
            .runloop_state = std::addressof(runloop_state),
            .result = std::addressof(result)};
 
        static_assert(std::same_as<decltype(rcvr), receiver_t<Sndr&&>>);
 
        auto op_state = stdexec::connect(std::forward<Sndr>(sndr), std::move(rcvr));
 
        stdexec::start(op_state);
 
        runloop_state.loop.run();
 
        if constexpr (sends_error<Sndr&&>::value)
            if (runloop_state.error)
                std::rethrow_exception(std::move(runloop_state.error));
 
        return result;
    }
};

template <template <typename...> typename SndrTrait>
struct ApplySenderFor {
    template <typename Sndr>
    requires SndrTrait<Sndr, env>::value
    auto operator()(Sndr&& sndr) const noexcept(std::is_nothrow_invocable_v<SyncWait, Sndr&&>) {
        return SyncWait{}(std::forward<Sndr>(sndr));
    }
};
 
} // namespace Kokkos::Execution::Impl::SyncWait
 
#endif // KOKKOS_EXECUTION_IMPL_SYNC_WAIT_HPP