Events.hpp

Go to the documentation of this file.

#ifndef KOKKOS_UTILS_CALLBACKS_EVENTS_HPP
#define KOKKOS_UTILS_CALLBACKS_EVENTS_HPP
 
#include <iomanip>
 
#include "impl/Kokkos_Profiling.hpp"
 
#include "kokkos-utils/impl/type_list.hpp"
 
namespace Kokkos::utils::callbacks
{

struct BeginParallelForEvent
{
    std::string name {};
    uint32_t dev_id = 0;
    uint64_t event_id = 0;
};

struct EndParallelForEvent
{
    uint64_t event_id = 0;
};

struct BeginParallelReduceEvent
{
    std::string name {};
    uint32_t dev_id = 0;
    uint64_t event_id = 0;
};

struct EndParallelReduceEvent
{
    uint64_t event_id = 0;
};

struct BeginParallelScanEvent
{
    std::string name {};
    uint32_t dev_id = 0;
    uint64_t event_id = 0;
};

struct EndParallelScanEvent
{
    uint64_t event_id = 0;
};

struct BeginFenceEvent
{
    std::string name {};
    uint32_t dev_id = 0;
    uint64_t event_id = 0;
};

struct EndFenceEvent
{
    uint64_t event_id = 0;
};

struct AllocDescriptor
{
    Kokkos_Profiling_SpaceHandle kpsh {};
    std::string name {};
    const void* ptr = nullptr;
    uint64_t size = 0;
};

struct AllocateDataEvent
{
    AllocDescriptor alloc {};
};

struct DeallocateDataEvent
{
    AllocDescriptor alloc {};
};

struct BeginDeepCopyEvent
{
    AllocDescriptor dst {};
    AllocDescriptor src {};
};

struct EndDeepCopyEvent {};

struct CreateProfileSectionEvent
{
    std::string name {};
    uint32_t section_id = 0;
};

struct StartProfileSectionEvent
{
    uint32_t section_id = 0;
};

struct StopProfileSectionEvent
{
    uint32_t section_id = 0;
};

struct DestroyProfileSectionEvent
{
    uint32_t section_id = 0;
};

struct PushRegionEvent
{
    std::string name {};
};

struct PopRegionEvent {};

struct ProfileEvent
{
    std::string name {};
};


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

template <typename T>
concept Event = impl::type_list_contains_v<T, EventTypeList>;

template <typename EventType>
concept BeginEvent =
    Event<EventType> &&
    std::constructible_from<EventType, const char*, uint32_t, uint64_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<uint64_t&>;
    };

template <typename EventType>
concept EndEvent =
    Event<EventType> && (! BeginEvent<EventType>) &&
    std::constructible_from<EventType, uint64_t> &&
    requires (EventType event) {
        { event.event_id } -> std::same_as<uint64_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 T, typename EventTypeListType>
concept EventOneOf = impl::type_list_contains_v<T, EventTypeListType>;

template <Event EventType>
constexpr auto get_name() {
    return Kokkos::Impl::TypeInfo<EventType>::name().substr(26); // erase "Kokkos::utils::callbacks::"
}
 
template <Event EventType>
std::ostream& operator<<(std::ostream &out, const EventType& /* event */) {
    return out << get_name<EventType>() << ": "
               << "{}";
}
 
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 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 ProfileEvent& event) {
    return out << get_name<ProfileEvent>() << ": "
               << "{name = " << std::quoted(event.name) << "}";
}


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(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(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