test_continues_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_continues_on.hpp"
#include "tests/utils/check_rcvr_env.hpp"
#include "tests/utils/execution_space_context.hpp"
#include "tests/utils/functors/increment.hpp"
#include "tests/utils/functors/labeled.hpp"
#include "tests/utils/kokkos.hpp"
#include "tests/utils/sink_receiver.hpp"
#include "tests/utils/stdexec.hpp"
#include "tests/utils/sync_wait.hpp"

 
using host_execution_space = Kokkos::DefaultHostExecutionSpace;
 
namespace Tests::ExecutionSpaceImpl {
 
using namespace Kokkos::utils::callbacks;
 
class ContinuesOnTest
    : 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,
        BeginSingleEvent,
        Kokkos::Execution::Impl::RecordEvent,
        Kokkos::Execution::Impl::WaitEvent
    >;
};

consteval bool test_schedule_from_sndr_traits() {
    using schd_sndr_t = typename ContinuesOnTest::schedule_sender_t;
 
    using schedule_from_sndr_t = Kokkos::Execution::ExecutionSpaceImpl::ScheduleFromSender<schd_sndr_t>;
 
    static_assert(stdexec::__nothrow_connectable<schedule_from_sndr_t, Tests::Utils::SinkReceiver>);
 
    return true;
}
 
static_assert(test_schedule_from_sndr_traits());

consteval bool test_continues_on_sndr_traits() {
    static_assert(Tests::Utils::check_continues_on<typename ContinuesOnTest::scheduler_t>());
 
    return true;
}
 
static_assert(test_continues_on_sndr_traits());

TEST_F(ContinuesOnTest, queryable_get_exec) {
    const auto [exec_A, exec_B] = Kokkos::Experimental::partition_space(exec, 1, 1);
 
    const host_execution_space exec_h{};
 
    const Kokkos::Execution::ExecutionSpaceContext esc_h{exec_h};
    const context_t esc_A{exec_A}, esc_B{exec_B};
 
    auto schs_A = stdexec::schedule(esc_A.get_scheduler());

    static_assert(!stdexec::__queryable_with<decltype(stdexec::get_env(schs_A)), Kokkos::Execution::Impl::get_exec_t>);
 
    auto schs_A_then = std::move(schs_A) | THEN_LABELED('A'); // NOLINT(performance-move-const-arg)
 
    const auto sch_B = esc_B.get_scheduler();
    auto schs_A_then_con_B = std::move(schs_A_then) // NOLINT(performance-move-const-arg)
                           | stdexec::continues_on(sch_B);

    ASSERT_EQ(stdexec::get_completion_scheduler<stdexec::set_value_t>(stdexec::get_env(schs_A_then_con_B)), sch_B);
 
    auto schs_A_then_con_B_then = std::move(schs_A_then_con_B) // NOLINT(performance-move-const-arg)
                                | THEN_LABELED('B');
 
    static_assert(std::same_as<
                  decltype(schs_A_then_con_B_then),
                  stdexec::then_t::__sender<
                      Tests::Utils::Functors::Labeled<'B'>,
                      stdexec::__sexpr<
                          stdexec::continues_on_t,
                          scheduler_t,
                          stdexec::__sexpr<
                              stdexec::schedule_from_t,
                              stdexec::__,
                              stdexec::then_t::__sender<Tests::Utils::Functors::Labeled<'A'>, schedule_sender_t>
                          >
                      >
                  >
    >);

    const auto sch_h = esc_h.get_scheduler();
    auto schs_A_then_con_B_then_con_h_then = std::move(schs_A_then_con_B_then) // NOLINT(performance-move-const-arg)
                                           | stdexec::continues_on(sch_h) | THEN_LABELED('h');
    ASSERT_EQ(
        stdexec::get_completion_scheduler<stdexec::set_value_t>(stdexec::get_env(schs_A_then_con_B_then_con_h_then)),
        sch_h);
 
    const auto op_state = stdexec::connect(
        std::move(schs_A_then_con_B_then_con_h_then), // NOLINT(performance-move-const-arg)
        Kokkos::Execution::Impl::SyncWait::Receiver<host_execution_space, std::true_type>{
            .state = std::addressof(esc_h.m_state), .runloop_state = nullptr, .result = nullptr});

    const auto& then_op_state = op_state;
    using then_rcvr_t = Kokkos::Execution::Impl::Receiver<Kokkos::Execution::ExecutionSpaceImpl::OpStateBase<
        Kokkos::Execution::Impl::SyncWait::Receiver<host_execution_space, std::true_type>,
        Kokkos::Execution::ExecutionSpaceImpl::ThenClosure<
            host_execution_space,
            std::string_view,
            Tests::Utils::Functors::Labeled<'h'>
        >
    >>;
    static_assert(std::same_as<decltype(then_op_state.inner_op_state.rcvr), then_rcvr_t>);
    static_assert(!stdexec::__queryable_with<stdexec::env_of_t<then_rcvr_t>, Kokkos::Execution::Impl::get_exec_t>);

    const auto& con_h_then_op_state = then_op_state.inner_op_state;
    static_assert(stdexec::__is_instance_of<
                  std::remove_cvref_t<decltype(con_h_then_op_state)>,
                  Kokkos::Execution::ExecutionSpaceImpl::ContinuesOnOpState
    >);
    using con_h_then_rcvr_t = decltype(con_h_then_op_state.inner_op_state.rcvr);
    static_assert(stdexec::__queryable_with<stdexec::env_of_t<con_h_then_rcvr_t>, Kokkos::Execution::Impl::get_exec_t>);
    static_assert(std::same_as<
                  stdexec::env_of_t<con_h_then_rcvr_t>,
                  stdexec::env<
                      stdexec::prop<
                          Kokkos::Execution::Impl::get_exec_t,
                          Kokkos::Execution::Impl::ExecutionSpaceRef<host_execution_space>
                      >,
                      stdexec::__env::__fwd<Kokkos::Execution::Impl::SyncWait::env>
                  >
    >);
    ASSERT_EQ(
        Kokkos::Execution::Impl::get_exec(stdexec::get_env(con_h_then_op_state.inner_op_state.rcvr)).get(), exec_h);
 
    const auto& sfrom_con_h_then_op_state = con_h_then_op_state.inner_op_state;
    static_assert(stdexec::__is_instance_of<
                  std::remove_cvref_t<decltype(sfrom_con_h_then_op_state)>,
                  Kokkos::Execution::ExecutionSpaceImpl::ScheduleFromOpState
    >);
    using sfrom_con_h_then_rcvr_t = decltype(sfrom_con_h_then_op_state.inner_op_state.completion_signal.rcvr);
    static_assert(
        !stdexec::__queryable_with<stdexec::env_of_t<sfrom_con_h_then_rcvr_t>, Kokkos::Execution::Impl::get_exec_t>);
 
    const auto& then_sfrom_con_h_then_op_state = sfrom_con_h_then_op_state.inner_op_state;
    static_assert(stdexec::__is_instance_of<
                  std::remove_cvref_t<decltype(then_sfrom_con_h_then_op_state)>,
                  Kokkos::Execution::ExecutionSpaceImpl::OpState
    >);
    using then_sfrom_con_h_then_rcvr_t = decltype(then_sfrom_con_h_then_op_state.inner_op_state.rcvr);
    static_assert(!stdexec::__queryable_with<
                  stdexec::env_of_t<then_sfrom_con_h_then_rcvr_t>,
                  Kokkos::Execution::Impl::get_exec_t
    >);
 
    const auto& con_B_then_sfrom_con_h_then_op_state = then_sfrom_con_h_then_op_state.inner_op_state;
    static_assert(stdexec::__is_instance_of<
                  std::remove_cvref_t<decltype(con_B_then_sfrom_con_h_then_op_state)>,
                  Kokkos::Execution::ExecutionSpaceImpl::ContinuesOnOpState
    >);
    using con_B_then_sfrom_con_h_then_rcvr_t = decltype(con_B_then_sfrom_con_h_then_op_state.inner_op_state.rcvr);
    static_assert(stdexec::__queryable_with<
                  stdexec::env_of_t<con_B_then_sfrom_con_h_then_rcvr_t>,
                  Kokkos::Execution::Impl::get_exec_t
    >);
    static_assert(std::same_as<
                  stdexec::env_of_t<con_B_then_sfrom_con_h_then_rcvr_t>,
                  stdexec::env<
                      stdexec::prop<
                          Kokkos::Execution::Impl::get_exec_t,
                          Kokkos::Execution::Impl::ExecutionSpaceRef<TEST_EXECUTION_SPACE>
                      >,
                      stdexec::__env::__fwd<stdexec::env<
                          stdexec::prop<
                              Kokkos::Execution::Impl::get_exec_t,
                              Kokkos::Execution::Impl::ExecutionSpaceRef<host_execution_space>
                          >,
                          stdexec::__env::__fwd<Kokkos::Execution::Impl::SyncWait::env>
                      >>
                  >
    >);
    ASSERT_EQ(
        Kokkos::Execution::Impl::get_exec(stdexec::get_env(con_B_then_sfrom_con_h_then_op_state.inner_op_state.rcvr))
            .get(),
        exec_B);
 
    const auto& sfrom_con_B_then_sfrom_con_h_then_op_state = con_B_then_sfrom_con_h_then_op_state.inner_op_state;
    static_assert(stdexec::__is_instance_of<
                  std::remove_cvref_t<decltype(sfrom_con_B_then_sfrom_con_h_then_op_state)>,
                  Kokkos::Execution::ExecutionSpaceImpl::ScheduleFromOpState
    >);
    using sfrom_con_B_then_sfrom_con_h_then_rcvr_t = decltype(sfrom_con_B_then_sfrom_con_h_then_op_state.inner_op_state
                                                                  .completion_signal.rcvr);
    static_assert(!stdexec::__queryable_with<
                  stdexec::env_of_t<sfrom_con_B_then_sfrom_con_h_then_rcvr_t>,
                  Kokkos::Execution::Impl::get_exec_t
    >);
 
    const auto& then_sfrom_B_then_sfrom_con_h_then_op_state = sfrom_con_B_then_sfrom_con_h_then_op_state.inner_op_state;
    static_assert(stdexec::__is_instance_of<
                  std::remove_cvref_t<decltype(then_sfrom_B_then_sfrom_con_h_then_op_state)>,
                  Kokkos::Execution::ExecutionSpaceImpl::OpState
    >);
    using then_sfrom_con_B_then_sfrom_con_h_then_rcvr_t = decltype(then_sfrom_B_then_sfrom_con_h_then_op_state
                                                                       .inner_op_state.rcvr);
    static_assert(!stdexec::__queryable_with<
                  stdexec::env_of_t<then_sfrom_con_B_then_sfrom_con_h_then_rcvr_t>,
                  Kokkos::Execution::Impl::get_exec_t
    >);
}

TEST_F(ContinuesOnTest, then_sync_wait) {
    const view_s_t data(Kokkos::view_alloc(exec, "data - shared space"));
 
    const context_t esc{exec};
 
    stdexec::sender auto chain = stdexec::just() | stdexec::continues_on(esc.get_scheduler()) | 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_SINGLE(exec, dispatch_label(exec, "then")),
            MATCHER_FOR_BEGIN_FENCE(exec, dispatch_label(exec, "sync_wait"))));
 
    ASSERT_EQ(data(), 1);
}

TEST_F(ContinuesOnTest, transition_to_same_execution_space_instance) {
    const view_s_t data(Kokkos::view_alloc(exec, "data - shared space"));
 
    const context_t esc{exec};
 
    auto chain = stdexec::just() | stdexec::continues_on(esc.get_scheduler()) | THEN_INCREMENT(data)
               | stdexec::continues_on(esc.get_scheduler()) | THEN_INCREMENT(data)
               | stdexec::continues_on(esc.get_scheduler()) | 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_SINGLE(exec, dispatch_label(exec, "then")),
            MATCHER_FOR_BEGIN_SINGLE(exec, dispatch_label(exec, "then")),
            MATCHER_FOR_BEGIN_SINGLE(exec, dispatch_label(exec, "then")),
            MATCHER_FOR_BEGIN_FENCE(exec, dispatch_label(exec, "sync_wait"))));
 
    ASSERT_EQ(data(), 3) << "A synchronization is missing.";
}

TEST_F(ContinuesOnTest, transition_to_another_execution_space_instance_and_back_same_type) {
    const view_s_t data(Kokkos::view_alloc(exec, "data - shared space"));
 
    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::just() | stdexec::continues_on(esc_A.get_scheduler()) | THEN_INCREMENT(data)
               | stdexec::continues_on(esc_B.get_scheduler()) | THEN_INCREMENT(data)
               | stdexec::continues_on(esc_A.get_scheduler()) | 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 constexpr (Kokkos::Execution::Impl::has_exec_wait_event<TEST_EXECUTION_SPACE>) {
        ASSERT_THAT(
            recorded_events,
            ::testing::ElementsAre(
                MATCHER_FOR_BEGIN_SINGLE(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_SINGLE(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_SINGLE(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_SINGLE(exec_A, dispatch_label(exec, "then")),
                    MATCHER_FOR_BEGIN_SINGLE(exec_B, dispatch_label(exec, "then")),
                    MATCHER_FOR_BEGIN_SINGLE(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_SINGLE(exec_A, dispatch_label(exec, "then")),
                    MATCHER_FOR_BEGIN_FENCE(exec_A, dispatch_label(exec, "dependency")),
                    MATCHER_FOR_BEGIN_SINGLE(exec_B, dispatch_label(exec, "then")),
                    MATCHER_FOR_BEGIN_FENCE(exec_B, dispatch_label(exec, "dependency")),
                    MATCHER_FOR_BEGIN_SINGLE(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(ContinuesOnTest, transition_to_another_execution_space_instance_and_back_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");
 
    using level_C_env_t = Kokkos::Execution::Impl::SyncWait::env;
    using level_B_env_t = stdexec::__env::__fwd<stdexec::env<
        stdexec::prop<
            Kokkos::Execution::Impl::get_exec_t,
            Kokkos::Execution::Impl::ExecutionSpaceRef<TEST_EXECUTION_SPACE>
        >,
        stdexec::__env::__fwd<level_C_env_t>
    >>;
    using level_A_env_t = stdexec::__env::__fwd<stdexec::env<
        stdexec::prop<
            Kokkos::Execution::Impl::get_exec_t,
            Kokkos::Execution::Impl::ExecutionSpaceRef<host_execution_space>
        >,
        level_B_env_t
    >>;
    auto chain = stdexec::just() | stdexec::continues_on(esc.get_scheduler()) | THEN_INCREMENT(data)
               | Tests::Utils::check_rcvr_env<level_A_env_t>() | stdexec::continues_on(esc_h.get_scheduler())
               | THEN_INCREMENT(data) | Tests::Utils::check_rcvr_env<level_B_env_t>()
               | stdexec::continues_on(esc.get_scheduler()) | THEN_INCREMENT(data)
               | Tests::Utils::check_rcvr_env<level_C_env_t>();
 
    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_SINGLE(exec, dispatch_label(exec, "then")),
                MATCHER_FOR_BEGIN_SINGLE(exec_h, dispatch_label(exec_h, "then")),
                MATCHER_FOR_BEGIN_SINGLE(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_SINGLE(exec, dispatch_label(exec, "then")),
                MATCHER_FOR_RECORD_EVENT(exec),
                MATCHER_FOR_WAIT_EXEC_EVENT(exec_h, recorded_events.at(1)),
                MATCHER_FOR_BEGIN_SINGLE(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_SINGLE(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_SINGLE(exec, dispatch_label(exec, "then")),
                MATCHER_FOR_BEGIN_FENCE(exec, dispatch_label(exec, "dependency")),
                MATCHER_FOR_BEGIN_SINGLE(exec_h, dispatch_label(exec_h, "then")),
                MATCHER_FOR_BEGIN_FENCE(exec_h, dispatch_label(exec_h, "dependency")),
                MATCHER_FOR_BEGIN_SINGLE(exec, dispatch_label(exec, "then")),
                MATCHER_FOR_BEGIN_FENCE(exec, dispatch_label(exec, "sync_wait"))));
    }
 
    ASSERT_EQ(data(), 3) << "A synchronization is missing.";
}

consteval bool test_sndr_nothrow_transformable() {
    using continues_on_sndr_t =
        decltype(stdexec::just() | stdexec::continues_on(std::declval<typename ContinuesOnTest::scheduler_t>()));
 
    static_assert(std::same_as<
                  continues_on_sndr_t,
                  stdexec::__sexpr<
                      stdexec::continues_on_t,
                      typename ContinuesOnTest::scheduler_t,
                      stdexec::__sexpr<stdexec::schedule_from_t, stdexec::__, stdexec::just_t::__sender<>>
                  >
    >);
 
    static_assert(stdexec::__detail::__has_nothrow_transform_sender<
                  Kokkos::Execution::ExecutionSpaceImpl::Domain,
                  stdexec::set_value_t,
                  continues_on_sndr_t&&,
                  stdexec::env<>
    >);
 
    using schedule_from_sndr_t = decltype(stdexec::schedule_from(
        stdexec::schedule(std::declval<typename ContinuesOnTest::scheduler_t>())));
 
    static_assert(std::same_as<
                  schedule_from_sndr_t,
                  stdexec::__sexpr<stdexec::schedule_from_t, stdexec::__, typename ContinuesOnTest::schedule_sender_t>
    >);
 
    static_assert(stdexec::__detail::__has_nothrow_transform_sender<
                  Kokkos::Execution::ExecutionSpaceImpl::Domain,
                  stdexec::set_value_t,
                  schedule_from_sndr_t&&,
                  stdexec::env<>
    >);
 
    return true;
}
 
static_assert(test_sndr_nothrow_transformable());
 
} // namespace Tests::ExecutionSpaceImpl