test_on.cpp

Go to the documentation of this file.

#include "kokkos-utils/callbacks/RecorderListener.hpp"
#include "kokkos-utils/tests/scoped/callbacks/Manager.hpp"
 
#include "kokkos-execution/execution_space.hpp"
 
#include "tests/utils/callback_matchers.hpp"
#include "tests/utils/check_scheduler_type.hpp"
#include "tests/utils/execution_space_context.hpp"
#include "tests/utils/functors/increment.hpp"
#include "tests/utils/kokkos.hpp"
#include "tests/utils/sync_wait.hpp"

 
using host_execution_space = Kokkos::DefaultHostExecutionSpace;
 
namespace Tests::ExecutionSpaceImpl {
 
using namespace Kokkos::utils::callbacks;
 
class OnTest
    : public Tests::Utils::ExecutionSpaceContextTest<TEST_EXECUTION_SPACE>
    , public Kokkos::utils::tests::scoped::callbacks::Manager {
   public:
    using recorder_listener_t = RecorderListener<
        EventDiscardMatcher<TEST_EXECUTION_SPACE>,
        BeginFenceEvent,
        BeginParallelForEvent,
        Kokkos::Execution::Impl::RecordEvent,
        Kokkos::Execution::Impl::WaitEvent
    >;
};

TEST_F(OnTest, on_same_execution_space_instance) {
    const view_s_t data(Kokkos::view_alloc("data", exec));
 
    const context_t esc{exec};
 
    auto chain = stdexec::schedule(esc.get_scheduler())
               | Tests::Utils::check_scheduler_type<stdexec::set_value_t, scheduler_t>() | THEN_INCREMENT(data)
               | stdexec::on(
                     esc.get_scheduler(),
                     Tests::Utils::check_scheduler_type<stdexec::set_value_t, scheduler_t>() | THEN_INCREMENT(data))
               | Tests::Utils::check_scheduler_type<stdexec::set_value_t, scheduler_t>() | THEN_INCREMENT(data);
 
    ASSERT_EQ(data(), 0) << "Eager execution is not allowed.";
 
    ASSERT_THAT(
        Tests::Utils::record_sync_wait<recorder_listener_t>(std::move(chain)),
        testing::ElementsAre(
            MATCHER_FOR_BEGIN_PFOR(exec, dispatch_label(exec, "then")),
            MATCHER_FOR_BEGIN_PFOR(exec, dispatch_label(exec, "then")),
            MATCHER_FOR_BEGIN_PFOR(exec, dispatch_label(exec, "then")),
            MATCHER_FOR_BEGIN_FENCE(exec, dispatch_label(exec, "sync_wait"))));
 
    ASSERT_EQ(data(), 3);
}

TEST_F(OnTest, on_another_execution_space_instance_same_type) {
    const view_s_t data(Kokkos::view_alloc("data", exec));
 
    const auto [exec_A, exec_B] = Kokkos::Experimental::partition_space(exec, 1, 1);
 
    const context_t esc_A{exec_A};
    Tests::Utils::show_exec_space_id(exec_A, "exec_A");
    const context_t esc_B{exec_B};
    Tests::Utils::show_exec_space_id(exec_B, "exec_B");
 
    auto chain = stdexec::schedule(esc_A.get_scheduler()) | THEN_INCREMENT(data)
               | stdexec::on(esc_B.get_scheduler(), THEN_INCREMENT(data)) | THEN_INCREMENT(data);
 
    static_assert(!Tests::Utils::has_completion_scheduler_for<decltype(chain), stdexec::set_value_t>);
    static_assert(stdexec::dependent_sender<decltype(chain)>);
    static_assert(Tests::Utils::has_completion_scheduler_for<decltype(chain), stdexec::set_value_t, stdexec::env<>>);
 
    ASSERT_EQ(data(), 0) << "Eager execution is not allowed.";
 
    const auto recorded_events = Tests::Utils::record_sync_wait<recorder_listener_t>(std::move(chain));
 
    if constexpr (Kokkos::Execution::Impl::has_exec_wait_event<TEST_EXECUTION_SPACE>) {
        ASSERT_THAT(
            recorded_events,
            testing::ElementsAre(
                MATCHER_FOR_BEGIN_PFOR(exec_A, dispatch_label(exec, "then")),
                MATCHER_FOR_RECORD_EVENT(exec_A),
                MATCHER_FOR_WAIT_EXEC_EVENT(exec_B, recorded_events.at(1)),
                MATCHER_FOR_BEGIN_PFOR(exec_B, dispatch_label(exec, "then")),
                MATCHER_FOR_RECORD_EVENT(exec_B),
                MATCHER_FOR_WAIT_EXEC_EVENT(exec_A, recorded_events.at(4)),
                MATCHER_FOR_BEGIN_PFOR(exec_A, dispatch_label(exec, "then")),
                MATCHER_FOR_BEGIN_FENCE(exec_A, dispatch_label(exec, "sync_wait"))));
    } else {
        if (Tests::Utils::are_same_instances(exec_A, exec_B)) {
            ASSERT_THAT(
                recorded_events,
                testing::ElementsAre(
                    MATCHER_FOR_BEGIN_PFOR(exec_A, dispatch_label(exec, "then")),
                    MATCHER_FOR_BEGIN_PFOR(exec_B, dispatch_label(exec, "then")),
                    MATCHER_FOR_BEGIN_PFOR(exec_A, dispatch_label(exec, "then")),
                    MATCHER_FOR_BEGIN_FENCE(exec_A, dispatch_label(exec, "sync_wait"))));
        } else {
            ASSERT_THAT(
                recorded_events,
                testing::ElementsAre(
                    MATCHER_FOR_BEGIN_PFOR(exec_A, dispatch_label(exec, "then")),
                    MATCHER_FOR_BEGIN_FENCE(exec_A, dispatch_label(exec, "dependency")),
                    MATCHER_FOR_BEGIN_PFOR(exec_B, dispatch_label(exec, "then")),
                    MATCHER_FOR_BEGIN_FENCE(exec_B, dispatch_label(exec, "dependency")),
                    MATCHER_FOR_BEGIN_PFOR(exec_A, dispatch_label(exec, "then")),
                    MATCHER_FOR_BEGIN_FENCE(exec_A, dispatch_label(exec, "sync_wait"))));
        }
    }
 
    ASSERT_EQ(data(), 3) << "A synchronization is missing.";
}

TEST_F(OnTest, many_execution_space_instances_of_different_type) {
    const view_s_t data(Kokkos::view_alloc(exec, "data - shared space"));
 
    const host_execution_space exec_h{};
 
    const Kokkos::Execution::ExecutionSpaceContext esc_h{exec_h};
    const context_t esc{exec};
 
    Tests::Utils::show_exec_space_id(exec, "exec");
    Tests::Utils::show_exec_space_id(exec_h, "exec_h");
 
    auto chain = stdexec::schedule(esc.get_scheduler()) | THEN_INCREMENT(data)
               | stdexec::on(esc_h.get_scheduler(), THEN_INCREMENT(data)) | THEN_INCREMENT(data);
 
    ASSERT_EQ(data(), 0) << "Eager execution is not allowed.";
 
    const auto recorded_events = Tests::Utils::record_sync_wait<recorder_listener_t>(std::move(chain));
 
    if (Tests::Utils::are_same_instances(exec, exec_h)) {
        ASSERT_THAT(
            recorded_events,
            testing::ElementsAre(
                MATCHER_FOR_BEGIN_PFOR(exec, dispatch_label(exec, "then")),
                MATCHER_FOR_BEGIN_PFOR(exec_h, dispatch_label(exec_h, "then")),
                MATCHER_FOR_BEGIN_PFOR(exec, dispatch_label(exec, "then")),
                MATCHER_FOR_BEGIN_FENCE(exec, dispatch_label(exec, "sync_wait"))));
    } else if constexpr (
        Kokkos::Execution::Impl::has_exec_wait_event<TEST_EXECUTION_SPACE>
        && std::same_as<host_execution_space, TEST_EXECUTION_SPACE>) {
        ASSERT_THAT(
            recorded_events,
            testing::ElementsAre(
                MATCHER_FOR_BEGIN_PFOR(exec, dispatch_label(exec, "then")),
                MATCHER_FOR_RECORD_EVENT(exec),
                MATCHER_FOR_WAIT_EXEC_EVENT(exec_h, recorded_events.at(1)),
                MATCHER_FOR_BEGIN_PFOR(exec_h, dispatch_label(exec_h, "then")),
                MATCHER_FOR_RECORD_EVENT(exec_h),
                MATCHER_FOR_WAIT_EXEC_EVENT(exec, recorded_events.at(4)),
                MATCHER_FOR_BEGIN_PFOR(exec, dispatch_label(exec, "then")),
                MATCHER_FOR_BEGIN_FENCE(exec, dispatch_label(exec, "sync_wait"))));
    } else {
        ASSERT_THAT(
            recorded_events,
            testing::ElementsAre(
                MATCHER_FOR_BEGIN_PFOR(exec, dispatch_label(exec, "then")),
                MATCHER_FOR_BEGIN_FENCE(exec, dispatch_label(exec, "dependency")),
                MATCHER_FOR_BEGIN_PFOR(exec_h, dispatch_label(exec_h, "then")),
                MATCHER_FOR_BEGIN_FENCE(exec_h, dispatch_label(exec_h, "dependency")),
                MATCHER_FOR_BEGIN_PFOR(exec, dispatch_label(exec, "then")),
                MATCHER_FOR_BEGIN_FENCE(exec, dispatch_label(exec, "sync_wait"))));
    }
 
    ASSERT_EQ(data(), 3) << "A synchronization is missing.";
}
 
} // namespace Tests::ExecutionSpaceImpl