event.hpp

Go to the documentation of this file.

#ifndef KOKKOS_EXECUTION_IMPL_EVENT_HPP
#define KOKKOS_EXECUTION_IMPL_EVENT_HPP
 
#include <concepts>
#include <optional>
 
#include "Kokkos_Macros.hpp"
#ifdef KOKKOS_ENABLE_EXPERIMENTAL_CXX20_MODULES
import kokkos.core;
#else
#    include "Kokkos_Core.hpp"
#endif
 
#if defined(KOKKOS_EXECUTION_ENABLE_EVENT_DISPATCH)
#    include "kokkos-utils/callbacks/Manager.hpp"
#endif
 
#include "kokkos-execution/impl/optional_ref.hpp"

 
namespace Kokkos::Execution::Impl {

template <typename EventType, typename Exec>
concept event = Kokkos::ExecutionSpace<Exec> && std::default_initializable<EventType> && !std::copyable<EventType>
             && !std::movable<EventType> && requires(EventType& obj, const Exec& exec) {
                    { obj.record(exec) } -> std::same_as<void>;
                } && requires(const EventType& obj) {
                    { obj.wait() } -> std::same_as<void>;
                    { obj.m_event_id } -> std::same_as<const uint64_t&>;
                };

template <Kokkos::ExecutionSpace Exec>
struct Event;

template <Kokkos::ExecutionSpace Exec>
struct HasNonBlockingDispatch : std::false_type { };
 
template <typename Exec>
concept has_non_blocking_dispatch = HasNonBlockingDispatch<Exec>::value;
 
static constexpr auto invalid_dev_id = Kokkos::Experimental::finite_max_v<uint32_t>;
static constexpr auto invalid_event_id = Kokkos::Experimental::finite_max_v<uint64_t>;

struct RecordEvent {
    uint32_t dev_id = 0;
    uint64_t event_id = 0;
 
    constexpr auto operator<=>(const RecordEvent&) const = default;
 
    friend std::ostream& operator<<(std::ostream& out, const RecordEvent& event) {
        return out << "RecordEvent: {dev_id = " << event.dev_id << ", event_id = " << event.event_id << '}';
    }
};
 
template <Kokkos::ExecutionSpace Exec>
void record_event(const Exec& exec, uint64_t& event_id) {
#if defined(KOKKOS_EXECUTION_ENABLE_EVENT_DISPATCH)
    event_id = Kokkos::utils::callbacks::get_next_event_id();
    Kokkos::utils::callbacks::dispatch(
        RecordEvent{.dev_id = Kokkos::Tools::Experimental::device_id(exec), .event_id = event_id});
#endif
}

struct WaitEvent {
    uint32_t dev_id = 0;
    uint64_t event_id = 0;
 
    constexpr auto operator<=>(const WaitEvent&) const = default;
 
    friend std::ostream& operator<<(std::ostream& out, const WaitEvent& event) {
        return out << "WaitEvent: {dev_id = " << event.dev_id << ", event_id = " << event.event_id << '}';
    }
};
 
template <Kokkos::ExecutionSpace Exec>
void wait_event(const Event<Exec>& event) {
#if defined(KOKKOS_EXECUTION_ENABLE_EVENT_DISPATCH)
    Kokkos::utils::callbacks::dispatch(WaitEvent{.dev_id = invalid_dev_id, .event_id = event.m_event_id});
#endif
}
 
template <Kokkos::ExecutionSpace ExecTo, Kokkos::ExecutionSpace ExecFrom>
void wait_event(const ExecTo& exec, const Event<ExecFrom>& event) {
#if defined(KOKKOS_EXECUTION_ENABLE_EVENT_DISPATCH)
    Kokkos::utils::callbacks::dispatch(
        WaitEvent{.dev_id = Kokkos::Tools::Experimental::device_id(exec), .event_id = event.m_event_id});
#endif
}

template <Kokkos::ExecutionSpace Exec>
struct Event {
    static_assert(!has_non_blocking_dispatch<Exec>);
 
    std::atomic<bool> m_recorded{false};
 
    uint64_t m_event_id = invalid_event_id;
 
    Event() = default;
    Event(const Event&) = delete;
    Event& operator=(const Event&) = delete;
    Event(Event&&) noexcept = delete;
    Event& operator=(Event&&) noexcept = delete;
    ~Event() = default;
 
    void record(const Exec&) {
        m_recorded.store(true, std::memory_order_release);
    }
 
    void wait() const {
        while (!m_recorded.load(std::memory_order_acquire)) {
        }
    }
};

template <Kokkos::ExecutionSpace Exec>
void record(Event<Exec>& event, const Exec& exec) {
    record_event(exec, event.m_event_id);
    event.record(exec);
}
 
template <Kokkos::ExecutionSpace... Exec>
void impl_wait(const Event<Exec>&... events) {
    (events.wait(), ...);
}

template <Kokkos::ExecutionSpace... Exec>
void wait(const Event<Exec>&... events) {
    (wait_event(events), ...);
    impl_wait(events...);
}
 
template <Kokkos::ExecutionSpace ExecTo, Kokkos::ExecutionSpace... ExecFrom>
void impl_wait(const ExecTo&, const Event<ExecFrom>&... events) {
    (events.wait(), ...);
}

template <Kokkos::ExecutionSpace ExecTo, Kokkos::ExecutionSpace... ExecFrom>
void wait(const ExecTo& exec, const Event<ExecFrom>&... events) {
    (wait_event(exec, events), ...);
    impl_wait(exec, events...);
}
 
template <Kokkos::ExecutionSpace Exec>
using event_storage_t = std::optional<Event<Exec>>;

template <Kokkos::ExecutionSpace Exec>
using OptionalConstEventRef = OptionalRef<const Event<Exec>>;
 
} // namespace Kokkos::Execution::Impl
 
#if defined(KOKKOS_ENABLE_CUDA)
#    include "kokkos-execution/impl/Cuda/event.hpp"
#endif
#if defined(KOKKOS_ENABLE_HIP)
#    include "kokkos-execution/impl/HIP/event.hpp"
#endif
#if defined(KOKKOS_ENABLE_HPX)
#    include "kokkos-execution/impl/HPX/event.hpp"
#endif
#if defined(KOKKOS_ENABLE_SYCL)
#    include "kokkos-execution/impl/SYCL/event.hpp"
#endif
 
#endif // KOKKOS_EXECUTION_IMPL_EVENT_HPP