get_exec.hpp

Go to the documentation of this file.

#ifndef KOKKOS_EXECUTION_EXECUTION_SPACE_GET_EXEC_HPP
#define KOKKOS_EXECUTION_EXECUTION_SPACE_GET_EXEC_HPP
 
#include "Kokkos_Core.hpp"
 
#include "kokkos-execution/impl/env.hpp"
 
namespace Kokkos::Execution::ExecutionSpaceImpl {

struct get_exec_t
    : public stdexec::__query<get_exec_t>
    , stdexec::forwarding_query_t {
    using stdexec::__query<get_exec_t>::operator();
};
 
inline constexpr get_exec_t get_exec{};

template <Kokkos::ExecutionSpace Exec>
struct ExecutionSpaceRef {
    Exec const * m_exec_ptr;
 
    explicit constexpr ExecutionSpaceRef(const Exec& exec) noexcept
        : m_exec_ptr(&exec) {
    }
 
    ExecutionSpaceRef() = delete;
    ExecutionSpaceRef(const ExecutionSpaceRef&) noexcept = default;
    ExecutionSpaceRef& operator=(const ExecutionSpaceRef&) noexcept = default;
    ExecutionSpaceRef(ExecutionSpaceRef&&) noexcept = default;
    ExecutionSpaceRef& operator=(ExecutionSpaceRef&&) noexcept = default;
    ~ExecutionSpaceRef() noexcept = default;
 
    const Exec& get() const noexcept {
        return *m_exec_ptr;
    }
 
    friend constexpr auto operator<=>(const ExecutionSpaceRef&, const ExecutionSpaceRef&) noexcept = default;
 
    [[nodiscard]]
    constexpr const ExecutionSpaceRef& query(const get_exec_t&) const noexcept {
        return *this;
    }
};
 
// NOLINTNEXTLINE(cppcoreguidelines-macro-usage)
#define KOKKOS_EXECUTION_UPSERT_EXEC_TYPE(_exec_type_, _rcvr_type_)                                                    \
    Impl::upsert_in_env_or_join_t<                                                                                     \
        Kokkos::Execution::ExecutionSpaceImpl::get_exec_t,                                                             \
        stdexec::__fwd_env_t<stdexec::env_of_t<_rcvr_type_>>,                                                          \
        Kokkos::Execution::ExecutionSpaceImpl::ExecutionSpaceRef<_exec_type_>                                          \
    >
 
// NOLINTNEXTLINE(cppcoreguidelines-macro-usage)
#define KOKKOS_EXECUTION_UPSERT_EXEC(_exec_type_, _exec_, _rcvr_type_, _rcvr_)                                         \
    [[nodiscard]]                                                                                                      \
    constexpr auto get_env() const noexcept -> KOKKOS_EXECUTION_UPSERT_EXEC_TYPE(_exec_type_, _rcvr_type_) {           \
        return Impl::upsert_in_env_or_join(                                                                            \
            Kokkos::Execution::ExecutionSpaceImpl::get_exec,                                                           \
            stdexec::__fwd_env(stdexec::get_env(_rcvr_)),                                                              \
            Kokkos::Execution::ExecutionSpaceImpl::ExecutionSpaceRef{_exec_});                                         \
    }
 
} // namespace Kokkos::Execution::ExecutionSpaceImpl
 
#endif // KOKKOS_EXECUTION_EXECUTION_SPACE_GET_EXEC_HPP