then.hpp

Go to the documentation of this file.

#ifndef KOKKOS_EXECUTION_EXECUTION_SPACE_THEN_HPP
#define KOKKOS_EXECUTION_EXECUTION_SPACE_THEN_HPP
 
#include "kokkos-execution/execution_space/execution_space_fwd.hpp"
 
#include "kokkos-execution/execution_space/parallel_for.hpp"
#include "kokkos-execution/execution_space/sender_concepts.hpp"
#include "kokkos-execution/impl/dispatch_label.hpp"
#include "kokkos-execution/impl/sender_introspection.hpp"
 
namespace Kokkos::Execution::ExecutionSpaceImpl {

template <typename Functor>
requires(std::same_as<void, std::invoke_result_t<Functor>>)
struct ThenWrapper {
    Functor functor; // NOLINT(cppcoreguidelines-avoid-const-or-ref-data-members)
 
    template <std::integral T>
    KOKKOS_FUNCTION void operator()(const T) const {
        functor();
    }
};
 
template <>
struct TransformSenderFor<stdexec::then_t> {
    template <typename Sndr, typename Env>
    using policy_t = Kokkos::RangePolicy<Impl::exec_of_t<Sndr, Env>, Kokkos::LaunchBounds<1>>;
 
    template <typename Env, typename Functor, typename Sndr>
    using trnsfrmd_sndr_t =
        ParallelForSender<stdexec::then_t, Sndr, std::string_view, ThenWrapper<Functor>, policy_t<Sndr, Env>>;
 
    template <typename Env, typename Functor, typename Sndr>
    requires stdexec::__sends<stdexec::set_value_t, Sndr, Env>
    auto operator()(const Env& env, stdexec::then_t, Functor&& functor, Sndr&& sndr) const
        noexcept(std::is_nothrow_constructible_v<
                 trnsfrmd_sndr_t<Env, Functor, Sndr>,
                 typename trnsfrmd_sndr_t<Env, Functor, Sndr>::closure_t&&,
                 Sndr&&
        >) {
        if constexpr (execution_space_completing_sender<Sndr, Env>) {
            auto schd = stdexec::get_completion_scheduler<stdexec::set_value_t>(stdexec::get_env(sndr), env);
 
            return trnsfrmd_sndr_t<Env, Functor, Sndr>{
                {{Impl::dispatch_label<Impl::exec_of_t<Sndr, Env>, ": then">(),
                  ThenWrapper<Functor>{std::forward<Functor>(functor)},
                  policy_t<Sndr, Env>(schd.state->exec, 0, 1)}},
                std::forward<Sndr>(sndr)};
        } else {
            return no_execution_space_scheduler_in_env<stdexec::then_t, Sndr, Env>();
        }
    }
};
 
} // namespace Kokkos::Execution::ExecutionSpaceImpl
 
#endif // KOKKOS_EXECUTION_EXECUTION_SPACE_THEN_HPP