test_RecorderListener.cpp

Go to the documentation of this file.

#include "gmock/gmock.h"
#include "gtest/gtest.h"
 
#include "Kokkos_Macros.hpp"
#ifdef KOKKOS_ENABLE_EXPERIMENTAL_CXX20_MODULES
import kokkos.core;
#else
#include "Kokkos_Core.hpp"
#endif
 
#include "kokkos-utils/callbacks/EventInProfileSectionMatcher.hpp"
#include "kokkos-utils/callbacks/EventRegexMatcher.hpp"
#include "kokkos-utils/callbacks/Helpers.hpp"
#include "kokkos-utils/callbacks/RecorderListener.hpp"
#include "kokkos-utils/impl/type_traits.hpp"
#include "kokkos-utils/tests/scoped/ExecutionSpace.hpp"
#include "kokkos-utils/tests/scoped/callbacks/Manager.hpp"
 
#include "tests/callbacks/Helpers.hpp"
#include "tests/callbacks/TestWorkload.hpp"

 
using execution_space = Kokkos::DefaultExecutionSpace;
 
namespace Kokkos::utils::tests::callbacks
{
 
using namespace Kokkos::utils::callbacks;

using event_in_profile_section_recorder_t = RecorderListener<EventInProfileSectionMatcher<EventRegexMatcher>, EventTypeList>;
 
class RecorderListenerTest : public ::testing::Test,
                             public scoped::callbacks::Manager,
                             public scoped::ExecutionSpace<execution_space>
{};

TEST(RecorderListener, traits)
{
    static_assert(std::same_as<
        listener_event_type_list_t<event_in_profile_section_recorder_t>,
        EventTypeList
    >);
 
    static_assert(Listener   <event_in_profile_section_recorder_t>);
    static_assert(ListenerFor<event_in_profile_section_recorder_t, EventTypeList>);
    static_assert(ListenerFor<event_in_profile_section_recorder_t, PushRegionEvent, PopRegionEvent>);
}
 
template <Event EvenType>
class RecorderListenerSingleEventTypeTest : public ::testing::Test,
                                            public scoped::callbacks::Manager {};
 
TYPED_TEST_SUITE(RecorderListenerSingleEventTypeTest, EventTestTypes);

TYPED_TEST(RecorderListenerSingleEventTypeTest, record)
{
    using recorder_listener_t = RecorderListener<TypeParam>;
    static_assert(std::same_as<typename recorder_listener_t::event_type_list_t, Kokkos::Impl::type_list<TypeParam>>);
    static_assert(std::same_as<typename recorder_listener_t::matcher_t, AnyEventMatcher>);
 
    const auto recorder = std::make_shared<recorder_listener_t>();
 
    Manager::  register_listener(recorder);
    Manager::unregister_listener(recorder.get());
 
    ASSERT_EQ(recorder->recorded_events.size(), 0);
}

TEST_F(RecorderListenerTest, recorded_events)
{
    EventInProfileSectionMatcher matcher{.matcher = EventRegexMatcher{.regex = std::regex("profile section")}};
    const auto recorder = std::make_shared<event_in_profile_section_recorder_t>(std::move(matcher));
 
    using recorder_t = RecorderListener<EventTypeList>;
 
    const auto any_event_recorder = std::make_shared<recorder_t>();
 
    Kokkos::utils::callbacks::Manager::register_listener(recorder);
    Kokkos::utils::callbacks::Manager::register_listener(any_event_recorder);
 
    MyWorkload<execution_space>{}.execute(exec);

    ASSERT_THAT(
        recorder->recorded_events,
        ::testing::Contains(ABeginParallelForEventWithName(::testing::StrEq("computation - level 0 - pfor")))
    );
 
    ASSERT_THAT(
        recorder->recorded_events,
        ContainsInOrder<typename recorder_t::event_variant_t>(
            APushRegionEventWithName(::testing::StrEq("computation - level 0")),
            ABeginParallelForEventWithName(::testing::StrEq("computation - level 0 - pfor")),
            AEndParallelForEvent(),
            ABeginFenceEventWithName(::testing::StrEq("computation - level 0 - fence after pfor")),
            AEndFenceEvent(),
            AProfileEventWithName(::testing::StrEq("buried marker")),
            APushRegionEventWithName(::testing::StrEq("computation - level 1")),
            ABeginParallelReduceEventWithName(::testing::StrEq("computation - level 1 - preduce on default exec")),
            AEndParallelReduceEvent(),
            ABeginFenceEventWithName(::testing::StrEq("Kokkos::parallel_reduce: fence due to result being value, not view")),
            AEndFenceEvent(),
            APopRegionEvent(),
            APopRegionEvent()
        )
    );
 
    ASSERT_GT(any_event_recorder->recorded_events.size(), recorder->recorded_events.size());
 
    const auto fence_event_outside_profile_section = ABeginFenceEventWithName(::testing::StrEq("other fence after stopping the profile section"));
 
    ASSERT_THAT(
        recorder->recorded_events,
        ::testing::Not(::testing::Contains(fence_event_outside_profile_section))
    );
 
    ASSERT_THAT(
        any_event_recorder->recorded_events,
        ::testing::Contains(fence_event_outside_profile_section)
    );
 
    const auto output = [&recorder] () -> std::string {
        std::ostringstream oss;
        recorder->report(oss);
        return oss.str();
    }();
 
    ASSERT_THAT(output, ::testing::HasSubstr("computation - level 0 - pfor"));
    ASSERT_THAT(output, ::testing::HasSubstr("computation - level 0 - fence after pfor"));
 
    Kokkos::utils::callbacks::Manager::unregister_listener(recorder.get());
    Kokkos::utils::callbacks::Manager::unregister_listener(any_event_recorder.get());
}

TEST_F(RecorderListenerTest, record)
{
    ASSERT_THAT(
        RecorderListener<EventTypeList>::record([&exec = this->exec]{ MyWorkload<execution_space>{}.execute(exec); }),
        ::testing::Contains(
            ABeginParallelForEvent(
                ::testing::Field(&BeginParallelForEvent::name,   ::testing::StrEq("computation - level 0 - pfor")),
                ::testing::Field(&BeginParallelForEvent::dev_id, ::testing::Eq(Kokkos::Tools::Experimental::device_id(this->exec)))
            )
        )
    );

    ASSERT_THAT(
        RecorderListener<EventTypeList>::record([&exec = this->exec]{ MyWorkload<execution_space>{}.execute(exec); }),
        ::testing::Contains(
            ABeginParallelForEvent(
                ::testing::FieldsAre(::testing::StrEq("computation - level 0 - pfor"), ::testing::Eq(Kokkos::Tools::Experimental::device_id(this->exec)), ::testing::_)
            )
        )
    );
}

using fence_finder_t = RecorderListener<EventRegexMatcher, BeginFenceEvent>;

TEST(FenceFinder, traits)
{
    static_assert(std::same_as<
        listener_event_type_list_t<fence_finder_t>,
        Kokkos::Impl::type_list<BeginFenceEvent>
    >);
 
    static_assert(Listener<fence_finder_t>);
}

TEST_F_WITH_CB_MGR(FenceFinderTest, recorded_events)
{
    const auto exec = Kokkos::Experimental::partition_space(execution_space{}, 1)[0];
 
    const auto fence_finder = std::make_shared<fence_finder_t>(EventRegexMatcher(std::regex("computation - level 0 - fence after pfor")));
 
    Kokkos::utils::callbacks::Manager::register_listener(fence_finder);
 
    MyWorkload<execution_space>{}.execute(exec);

    ASSERT_THAT(
        fence_finder->recorded_events,
        ::testing::Contains(
            ABeginFenceEvent(
                ::testing::Field(&BeginFenceEvent::name,   ::testing::StrEq("computation - level 0 - fence after pfor")),
                ::testing::Field(&BeginFenceEvent::dev_id, ::testing::Eq(Kokkos::Tools::Experimental::device_id(exec)))
            )
        )
    );
 
    Kokkos::utils::callbacks::Manager::unregister_listener(fence_finder.get());
}
 
struct CustomEvent {
    EventTraits::event_id_t event_id = EventTraits::invalid_event_id;
 
    bool operator==(const CustomEvent&) const = default;
};
 
inline std::ostream& operator<<(std::ostream& out, const CustomEvent& event) {
    return out << "CustomEvent: " << "{event_id = " << event.event_id << "}";
}
 
struct CustomEventWithPayload {
    std::string payload {};
    EventTraits::event_id_t event_id = EventTraits::invalid_event_id;
 
    bool operator==(const CustomEventWithPayload&) const = default;
};
 
inline std::ostream& operator<<(std::ostream& out, const CustomEventWithPayload& event) {
    return out << "CustomEventWithPayload: " << "{payload = " << event.payload << ", event_id = " << event.event_id << "}";
}
 
static_assert(Kokkos::utils::callbacks::Event<CustomEvent>);
static_assert(Kokkos::utils::callbacks::Event<CustomEventWithPayload>);
 
DEFINE_EVENT_MATCHER_IN(Kokkos::utils::tests::callbacks, CustomEvent)
DEFINE_EVENT_MATCHER_IN(Kokkos::utils::tests::callbacks, CustomEventWithPayload)
 

TEST_F(RecorderListenerTest, record_mix_kokkos_and_custom_events) {
    using recorder_listener_t = RecorderListener<BeginParallelForEvent, CustomEvent, CustomEventWithPayload>;
 
    static_assert(std::same_as<
        Kokkos::utils::callbacks::listener_event_type_list_t<recorder_listener_t>,
        Kokkos::Impl::type_list<BeginParallelForEvent, CustomEvent, CustomEventWithPayload>
    >);
 
    static_assert(Kokkos::utils::callbacks::impl::ImplicitListener<recorder_listener_t>);
    static_assert(Kokkos::utils::callbacks::impl::ExplicitListener<recorder_listener_t>);
    static_assert(Kokkos::utils::callbacks::Listener<recorder_listener_t>);
    static_assert(Kokkos::utils::callbacks::ListenerFor<recorder_listener_t, BeginParallelForEvent, CustomEvent, CustomEventWithPayload>);
 
    ASSERT_THAT(
        (recorder_listener_t::record([&exec = this->exec]{
            MyWorkload<execution_space>{}.execute(exec);
            Kokkos::utils::callbacks::dispatch(CustomEventWithPayload{
                .payload = "My taylor is rich.",
                .event_id = Kokkos::utils::callbacks::get_next_event_id()
            });
            Kokkos::utils::callbacks::dispatch(CustomEvent{
                .event_id = Kokkos::utils::callbacks::get_next_event_id()
            });
        })),
        ContainsInOrder<typename recorder_listener_t::event_variant_t>(
            ABeginParallelForEvent(
                ::testing::Field(&BeginParallelForEvent::name, ::testing::StrEq("computation - level 0 - pfor")),
                ::testing::Field(&BeginParallelForEvent::dev_id, ::testing::Eq(Kokkos::Tools::Experimental::device_id(this->exec))),
                ::testing::Field(&BeginParallelForEvent::event_id, ::testing::Eq(0))
            ),
            ACustomEventWithPayload(
                ::testing::Field(&CustomEventWithPayload::payload, ::testing::StrEq("My taylor is rich.")),
                ::testing::Field(&CustomEventWithPayload::event_id, ::testing::Eq(1))
            ),
            ACustomEvent(
                ::testing::Field(&CustomEvent::event_id, ::testing::Eq(2))
            )
        )
    );
}

TEST_F(RecorderListenerTest, record_custom_event) {
    using recorder_listener_t = RecorderListener<CustomEventWithPayload>;
 
    static_assert(std::same_as<
        Kokkos::utils::callbacks::listener_event_type_list_t<recorder_listener_t>,
        Kokkos::Impl::type_list<CustomEventWithPayload>
    >);
 
    static_assert(!Kokkos::utils::callbacks::impl::ImplicitListener<recorder_listener_t>);
    static_assert(Kokkos::utils::callbacks::impl::ExplicitListener<recorder_listener_t>);
    static_assert(Kokkos::utils::callbacks::Listener<recorder_listener_t>);
    static_assert(Kokkos::utils::callbacks::ListenerFor<recorder_listener_t, CustomEventWithPayload>);
    static_assert(!Kokkos::utils::callbacks::ListenerFor<recorder_listener_t, CustomEvent>);
 
    ASSERT_THAT(
        (recorder_listener_t::record([]{
            Kokkos::utils::callbacks::dispatch(CustomEventWithPayload{
                .payload = "My taylor is rich.",
                .event_id = Kokkos::utils::callbacks::get_next_event_id()
            });
        })),
        ::testing::ElementsAre(
            ACustomEventWithPayload(
                ::testing::Field(&CustomEventWithPayload::payload, ::testing::StrEq("My taylor is rich.")),
                ::testing::Field(&CustomEventWithPayload::event_id, ::testing::Eq(0))
            )
        )
    );
}
 
} // namespace Kokkos::utils::tests::callbacks