Events.hpp

Go to the documentation of this file.

#ifndef KOKKOS_UTILS_CALLBACKS_EVENTS_HPP
#define KOKKOS_UTILS_CALLBACKS_EVENTS_HPP
 
#include <cstring>
#include <iomanip>
#include <memory>
 
#include "Kokkos_Macros.hpp"
#ifdef KOKKOS_ENABLE_EXPERIMENTAL_CXX20_MODULES
import kokkos.core;
#else
#include "Kokkos_Core.hpp"
#endif
#include "Kokkos_TypeInfo.hpp"
 
#include "kokkos-utils/impl/type_list.hpp"
 
namespace Kokkos::utils::callbacks
{

struct EventTraits {
    using event_id_t = uint64_t;
 
    static constexpr event_id_t invalid_event_id = Kokkos::Experimental::finite_max_v<event_id_t>;
};

bool operator==(const Kokkos::Tools::SpaceHandle& fst, const Kokkos::Tools::SpaceHandle& snd) {
    return strcmp(fst.name, snd.name) == 0;
};

struct BeginParallelForEvent
{
    std::string name {};
    uint32_t dev_id = 0;
    EventTraits::event_id_t event_id = EventTraits::invalid_event_id;
 
    bool operator==(const BeginParallelForEvent&) const = default;
};

struct EndParallelForEvent
{
    EventTraits::event_id_t event_id = EventTraits::invalid_event_id;
 
    bool operator==(const EndParallelForEvent&) const = default;
};

struct BeginParallelReduceEvent
{
    std::string name {};
    uint32_t dev_id = 0;
    EventTraits::event_id_t event_id = EventTraits::invalid_event_id;
 
    bool operator==(const BeginParallelReduceEvent&) const = default;
};

struct EndParallelReduceEvent
{
    EventTraits::event_id_t event_id = EventTraits::invalid_event_id;
 
    bool operator==(const EndParallelReduceEvent&) const = default;
};

struct BeginParallelScanEvent
{
    std::string name {};
    uint32_t dev_id = 0;
    EventTraits::event_id_t event_id = EventTraits::invalid_event_id;
 
    bool operator==(const BeginParallelScanEvent&) const = default;
};

struct EndParallelScanEvent
{
    EventTraits::event_id_t event_id = EventTraits::invalid_event_id;
 
    bool operator==(const EndParallelScanEvent&) const = default;
};

struct BeginSingleEvent
{
    std::string name {};
    uint32_t dev_id = 0;
    EventTraits::event_id_t event_id = EventTraits::invalid_event_id;
 
    bool operator==(const BeginSingleEvent&) const = default;
};

struct EndSingleEvent
{
    EventTraits::event_id_t event_id = EventTraits::invalid_event_id;
 
    bool operator==(const EndSingleEvent&) const = default;
};

struct BeginFenceEvent
{
    std::string name {};
    uint32_t dev_id = 0;
    EventTraits::event_id_t event_id = EventTraits::invalid_event_id;
 
    bool operator==(const BeginFenceEvent&) const = default;
};

struct EndFenceEvent
{
    EventTraits::event_id_t event_id = EventTraits::invalid_event_id;
 
    bool operator==(const EndFenceEvent&) const = default;
};

struct AllocDescriptor
{
    Kokkos::Tools::SpaceHandle kpsh {};
    std::string name {};
    const void* ptr = nullptr;
    uint64_t size = 0;
 
    bool operator==(const AllocDescriptor&) const = default;
};

struct AllocateDataEvent
{
    AllocDescriptor alloc {};
 
    bool operator==(const AllocateDataEvent&) const = default;
};

struct DeallocateDataEvent
{
    AllocDescriptor alloc {};
 
    bool operator==(const DeallocateDataEvent&) const = default;
};

struct BeginDeepCopyEvent
{
    AllocDescriptor dst {};
    AllocDescriptor src {};
 
    bool operator==(const BeginDeepCopyEvent&) const = default;
};

struct EndDeepCopyEvent
{
    bool operator==(const EndDeepCopyEvent&) const = default;
};

struct CreateProfileSectionEvent
{
    std::string name {};
    uint32_t section_id = 0;
 
    bool operator==(const CreateProfileSectionEvent&) const = default;
};

struct StartProfileSectionEvent
{
    uint32_t section_id = 0;
 
    bool operator==(const StartProfileSectionEvent&) const = default;
};

struct StopProfileSectionEvent
{
    uint32_t section_id = 0;
 
    bool operator==(const StopProfileSectionEvent&) const = default;
};

struct DestroyProfileSectionEvent
{
    uint32_t section_id = 0;
 
    bool operator==(const DestroyProfileSectionEvent&) const = default;
};

struct PushRegionEvent
{
    std::string name {};
 
    bool operator==(const PushRegionEvent&) const = default;
};

struct PopRegionEvent
{
    bool operator==(const PopRegionEvent&) const = default;
};

struct ProfileEvent
{
    std::string name {};
 
    bool operator==(const ProfileEvent&) const = default;
};


using EventTypeList = Kokkos::Impl::type_list<
    BeginParallelForEvent,     EndParallelForEvent,
    BeginParallelReduceEvent,  EndParallelReduceEvent,
    BeginParallelScanEvent,    EndParallelScanEvent,
    BeginSingleEvent,          EndSingleEvent,
    BeginFenceEvent,           EndFenceEvent,
    AllocateDataEvent,         DeallocateDataEvent,
    BeginDeepCopyEvent,        EndDeepCopyEvent,
    CreateProfileSectionEvent, DestroyProfileSectionEvent,
    StartProfileSectionEvent,  StopProfileSectionEvent,
    PushRegionEvent,           PopRegionEvent,
    ProfileEvent
>;

template <typename T>
concept Event = std::default_initializable<T>
    && std::copyable<T>
    && std::movable<T>
    && std::equality_comparable<T>;

template <typename EventType>
concept BeginEvent =
    Event<EventType> &&
    std::constructible_from<EventType, const char*, uint32_t, EventTraits::event_id_t> &&
    requires (EventType event) {
        { event.name }     -> std::same_as<std::string&>;
        { event.dev_id }   -> std::same_as<uint32_t&>;
        { event.event_id } -> std::same_as<EventTraits::event_id_t&>;
    };

template <typename EventType>
concept EndEvent =
    Event<EventType> && (! BeginEvent<EventType>) &&
    std::constructible_from<EventType, EventTraits::event_id_t> &&
    requires (EventType event) {
        { event.event_id } -> std::same_as<EventTraits::event_id_t&>;
    };

template <typename EventType>
concept DataEvent =
    Event<EventType> &&
    std::constructible_from<EventType, AllocDescriptor> &&
    requires (EventType event) {
        { event.alloc } -> std::same_as<AllocDescriptor&>;
    };

template <typename EventType>
concept ProfileSectionManipulationEvent =
    Event<EventType> &&
    std::constructible_from<EventType, uint32_t> &&
    requires (EventType event) {
        { event.section_id } -> std::same_as<uint32_t&>;
    };

template <typename EventType>
concept NamedEvent =
    Event<EventType> &&
    requires (EventType event) {
        { event.name } -> std::same_as<std::string&>;
    };

template <typename EventType>
concept IndexedEvent = Event<EventType> && requires (EventType event) {
    { event.event_id } -> std::same_as<EventTraits::event_id_t&>;
};

template <typename EventType>
concept EnqueuedEvent = Event<EventType> && requires (EventType event) {
    { event.dev_id } -> std::same_as<uint32_t&>;
};

template <typename T, typename... EventTypes>
concept EventOneOf = impl::type_list_contains_v<T, EventTypes...>;

template <Event EventType>
struct PairedEventType;
 
#define PAIRED_EVENT_TYPE(__begin_eventtype__, __end_eventtype__) \
    template <> \
    struct PairedEventType<__begin_eventtype__> { using type = __end_eventtype__; };
 
PAIRED_EVENT_TYPE(BeginParallelForEvent,     EndParallelForEvent)
PAIRED_EVENT_TYPE(BeginParallelReduceEvent,  EndParallelReduceEvent)
PAIRED_EVENT_TYPE(BeginParallelScanEvent,    EndParallelScanEvent)
PAIRED_EVENT_TYPE(BeginSingleEvent,          EndSingleEvent)
PAIRED_EVENT_TYPE(BeginFenceEvent,           EndFenceEvent)
PAIRED_EVENT_TYPE(AllocateDataEvent,         DeallocateDataEvent)
PAIRED_EVENT_TYPE(BeginDeepCopyEvent,        EndDeepCopyEvent)
PAIRED_EVENT_TYPE(CreateProfileSectionEvent, DestroyProfileSectionEvent)
PAIRED_EVENT_TYPE(StartProfileSectionEvent,  StopProfileSectionEvent)
PAIRED_EVENT_TYPE(PushRegionEvent,           PopRegionEvent)
 
template <Event EventType>
using paired_event_t = typename PairedEventType<EventType>::type;


template <typename T>
constexpr auto get_name() {
    return Kokkos::Impl::TypeInfo<T>::name().substr(26);
}
 
template <BeginEvent EventType>
std::ostream& operator<<(std::ostream &out, const EventType& event) {
    return out << get_name<EventType>() << ": "
               << "{name = " << std::quoted(event.name) << ", dev_id = " << event.dev_id << ", event_id = " << event.event_id  << "}";
}
 
template <EndEvent EventType>
std::ostream& operator<<(std::ostream &out, const EventType& event) {
    return out << get_name<EventType>() << ": "
               << "{event_id = " << event.event_id  << "}";
}
 
template <DataEvent EventType>
std::ostream& operator<<(std::ostream &out, const EventType& event) {
    const auto& alloc = event.alloc;
    return out << get_name<EventType>() << ": "
               << "{name = " << std::quoted(alloc.name) << ", space = " << std::quoted(alloc.kpsh.name) << ", ptr = " << alloc.ptr << ", size = " << alloc.size << "}";
}
 
inline std::ostream& operator<<(std::ostream &out, const BeginDeepCopyEvent& event) {
    const auto& dst = event.dst;
    const auto& src = event.src;
    return out << get_name<BeginDeepCopyEvent>() << ": "
               << "{src = " << std::quoted(src.name) << " (" << src.kpsh.name << ", " << src.ptr << ") -> "
               << "dst = "  << std::quoted(dst.name) << " (" << dst.kpsh.name << ", " << dst.ptr << ") of size " << src.size << "}";
}
 
inline std::ostream& operator<<(std::ostream &out, const EndDeepCopyEvent& /* event */) {
    return out << get_name<EndDeepCopyEvent>() << ": "
               << "{}";
}
 
inline std::ostream& operator<<(std::ostream &out, const CreateProfileSectionEvent& event) {
    return out << get_name<CreateProfileSectionEvent>() << ": "
               << "{name = " << std::quoted(event.name) << ", section_id = " << event.section_id << "}";
}
 
template <ProfileSectionManipulationEvent EventType>
std::ostream& operator<<(std::ostream &out, const EventType& event) {
    return out << get_name<EventType>() << ": "
               << "{section_id = " << event.section_id << "}";
}
 
inline std::ostream& operator<<(std::ostream &out, const PushRegionEvent& event) {
    return out << get_name<PushRegionEvent>() << ": "
               << "{name = " << std::quoted(event.name) << "}";
}
 
inline std::ostream& operator<<(std::ostream &out, const PopRegionEvent& /* event */) {
    return out << get_name<PopRegionEvent>() << ": "
               << "{}";
}
 
inline std::ostream& operator<<(std::ostream &out, const ProfileEvent& event) {
    return out << get_name<ProfileEvent>() << ": "
               << "{name = " << std::quoted(event.name) << "}";
}
 
inline std::ostream& operator<<(std::ostream& out, const AllocDescriptor& descr) {
    return out << get_name<AllocDescriptor>() << ": "
               << "{name = " << std::quoted(descr.name) << " (" << descr.kpsh.name << ", " << descr.ptr << ") of size " << descr.size << '}';
}


template <Event EventType>
auto get_callback_from_eventset(const Kokkos::Tools::Experimental::EventSet& event_set);

#define GET_CALLBACK_FROM_EVENTSET(__eventtype__, __callback__) \
    template <> \
    inline auto get_callback_from_eventset<__eventtype__>(const Kokkos::Tools::Experimental::EventSet& event_set) { return event_set.__callback__; }
 
GET_CALLBACK_FROM_EVENTSET(BeginParallelForEvent,      begin_parallel_for)
GET_CALLBACK_FROM_EVENTSET(EndParallelForEvent,        end_parallel_for)
GET_CALLBACK_FROM_EVENTSET(BeginParallelReduceEvent,   begin_parallel_reduce)
GET_CALLBACK_FROM_EVENTSET(EndParallelReduceEvent,     end_parallel_reduce)
GET_CALLBACK_FROM_EVENTSET(BeginParallelScanEvent,     begin_parallel_scan)
GET_CALLBACK_FROM_EVENTSET(EndParallelScanEvent,       end_parallel_scan)
GET_CALLBACK_FROM_EVENTSET(BeginSingleEvent,           begin_single)
GET_CALLBACK_FROM_EVENTSET(EndSingleEvent,             end_single)
GET_CALLBACK_FROM_EVENTSET(BeginFenceEvent,            begin_fence)
GET_CALLBACK_FROM_EVENTSET(EndFenceEvent,              end_fence)
GET_CALLBACK_FROM_EVENTSET(AllocateDataEvent,          allocate_data)
GET_CALLBACK_FROM_EVENTSET(DeallocateDataEvent,        deallocate_data)
GET_CALLBACK_FROM_EVENTSET(BeginDeepCopyEvent,         begin_deep_copy)
GET_CALLBACK_FROM_EVENTSET(EndDeepCopyEvent,           end_deep_copy)
GET_CALLBACK_FROM_EVENTSET(CreateProfileSectionEvent,  create_profile_section)
GET_CALLBACK_FROM_EVENTSET(StartProfileSectionEvent,   start_profile_section)
GET_CALLBACK_FROM_EVENTSET(StopProfileSectionEvent,    stop_profile_section)
GET_CALLBACK_FROM_EVENTSET(DestroyProfileSectionEvent, destroy_profile_section)
GET_CALLBACK_FROM_EVENTSET(PushRegionEvent,            push_region)
GET_CALLBACK_FROM_EVENTSET(PopRegionEvent,             pop_region)
GET_CALLBACK_FROM_EVENTSET(ProfileEvent,               profile_event)

template <Event EventType>
auto get_callback_setter();

#define GET_CALLBACK_SETTER(__eventtype__, __callback_setter__) \
    template <> \
    inline auto get_callback_setter<__eventtype__>() { return &Kokkos::Tools::Experimental::__callback_setter__; }
 
GET_CALLBACK_SETTER(BeginParallelForEvent,      set_begin_parallel_for_callback)
GET_CALLBACK_SETTER(EndParallelForEvent,        set_end_parallel_for_callback)
GET_CALLBACK_SETTER(BeginParallelReduceEvent,   set_begin_parallel_reduce_callback)
GET_CALLBACK_SETTER(EndParallelReduceEvent,     set_end_parallel_reduce_callback)
GET_CALLBACK_SETTER(BeginParallelScanEvent,     set_begin_parallel_scan_callback)
GET_CALLBACK_SETTER(EndParallelScanEvent,       set_end_parallel_scan_callback)
GET_CALLBACK_SETTER(BeginSingleEvent,           set_begin_single_callback)
GET_CALLBACK_SETTER(EndSingleEvent,             set_end_single_callback)
GET_CALLBACK_SETTER(BeginFenceEvent,            set_begin_fence_callback)
GET_CALLBACK_SETTER(EndFenceEvent,              set_end_fence_callback)
GET_CALLBACK_SETTER(AllocateDataEvent,          set_allocate_data_callback)
GET_CALLBACK_SETTER(DeallocateDataEvent,        set_deallocate_data_callback)
GET_CALLBACK_SETTER(BeginDeepCopyEvent,         set_begin_deep_copy_callback)
GET_CALLBACK_SETTER(EndDeepCopyEvent,           set_end_deep_copy_callback)
GET_CALLBACK_SETTER(CreateProfileSectionEvent,  set_create_profile_section_callback)
GET_CALLBACK_SETTER(StartProfileSectionEvent,   set_start_profile_section_callback)
GET_CALLBACK_SETTER(StopProfileSectionEvent,    set_stop_profile_section_callback)
GET_CALLBACK_SETTER(DestroyProfileSectionEvent, set_destroy_profile_section_callback)
GET_CALLBACK_SETTER(PushRegionEvent,            set_push_region_callback)
GET_CALLBACK_SETTER(PopRegionEvent,             set_pop_region_callback)
GET_CALLBACK_SETTER(ProfileEvent,               set_profile_event_callback)
 
} // namespace Kokkos::utils::callbacks
 
#endif // KOKKOS_UTILS_CALLBACKS_EVENTS_HPP