get_exec.hpp

Go to the documentation of this file.

#ifndef KOKKOS_EXECUTION_IMPL_GET_EXEC_HPP
#define KOKKOS_EXECUTION_IMPL_GET_EXEC_HPP
 
#include "Kokkos_Core.hpp"
 
#include "kokkos-execution/impl/sender_introspection.hpp"
 
namespace Kokkos::Execution::Impl {

struct get_exec_t : public stdexec::__query<get_exec_t> {
    using stdexec::__query<get_exec_t>::operator();
};
 
inline constexpr get_exec_t get_exec{};
 
template <typename... Args>
struct ExecOf;

template <stdexec::sender Sndr, typename... Env>
struct ExecOf<Sndr, Env...> {
    using type = typename completion_scheduler_of_t<stdexec::set_value_t, Sndr, Env...>::execution_space;
};

template <stdexec::__queryable_with<get_exec_t> T>
struct ExecOf<T> {
    using type = typename std::remove_cvref_t<stdexec::__query_result_t<T, get_exec_t>>::execution_space;
};
 
template <typename... Args>
using exec_of_t = typename ExecOf<Args...>::type;

template <Kokkos::ExecutionSpace Exec>
struct ExecutionSpaceRef {
    using execution_space = Exec;
 
    execution_space const * m_exec_ptr;
 
    explicit constexpr ExecutionSpaceRef(const execution_space& 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 execution_space& 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;
    }
};
 
} // namespace Kokkos::Execution::Impl
 
#endif // KOKKOS_EXECUTION_IMPL_GET_EXEC_HPP