when_all.hpp

Go to the documentation of this file.

#ifndef KOKKOS_EXECUTION_EXECUTION_SPACE_WHEN_ALL_HPP
#define KOKKOS_EXECUTION_EXECUTION_SPACE_WHEN_ALL_HPP
 
#include "kokkos-execution/execution_space/execution_space_fwd.hpp"
 
#include "kokkos-execution/execution_space/env.hpp"
#include "kokkos-execution/impl/get_exec.hpp"
#include "kokkos-execution/stdexec.hpp"
 
namespace Kokkos::Execution::ExecutionSpaceImpl {
 
struct when_all_t { };
 
} // namespace Kokkos::Execution::ExecutionSpaceImpl
 
namespace stdexec {
template <>
struct __sexpr_impl<Kokkos::Execution::ExecutionSpaceImpl::when_all_t> : stdexec::__when_all::__when_all_impl {
    static constexpr auto __get_env = // NOLINT(bugprone-reserved-identifier)
        []<typename State>(stdexec::__ignore, const State& state) noexcept {
            auto env = stdexec::__when_all::__when_all_impl::__get_env(stdexec::__ignore{}, state);
 
            using exec_env_policy_t =
                Kokkos::Execution::ExecutionSpaceImpl::exec_env_policy_t<stdexec::env_of_t<decltype(state.__rcvr_)>>;
 
            if constexpr (std::same_as<exec_env_policy_t, Kokkos::Execution::ExecutionSpaceImpl::WithExecEnvPolicy>) {
                return Kokkos::Execution::ExecutionSpaceImpl::join_env_with_exec(
                    std::move(env), Kokkos::Execution::Impl::get_exec(stdexec::get_env(state.__rcvr_)).get());
            } else {
                return std::move(env);
            }
        };

    template <typename Self, typename... Envs>
    static consteval auto __get_completion_signatures() { // NOLINT(bugprone-reserved-identifier)
        if constexpr (stdexec::__minvocable_q<stdexec::__when_all::__when_all_impl::__completions_t, Self, Envs...>) {
            return stdexec::__when_all::__when_all_impl::__completions_t<Self, Envs...>{};
        } else if constexpr (sizeof...(Envs) == 0) {
            return stdexec::__throw_dependent_sender_error<Self>();
        }
    }
};
 
} // namespace stdexec
 
namespace Kokkos::Execution::ExecutionSpaceImpl {
 
template <>
struct TransformSenderFor<stdexec::when_all_t> {
    template <typename Env, typename... Sndrs>
    auto operator()(const Env&, stdexec::when_all_t, stdexec::__, Sndrs&&... sndrs) const {
        return stdexec::__make_sexpr<Kokkos::Execution::ExecutionSpaceImpl::when_all_t>(
            stdexec::__{}, std::forward<Sndrs>(sndrs)...);
    }
};
 
} // namespace Kokkos::Execution::ExecutionSpaceImpl
 
#endif // KOKKOS_EXECUTION_EXECUTION_SPACE_WHEN_ALL_HPP