test_bulk.cpp

Go to the documentation of this file.

#include "kokkos-utils/callbacks/ConjunctionMatcher.hpp"
#include "kokkos-utils/callbacks/RecorderListener.hpp"
#include "kokkos-utils/tests/scoped/callbacks/Manager.hpp"
 
#include "kokkos-execution/graph.hpp"
#include "kokkos-execution/impl/event.hpp"
 
#include "tests/graph/events.hpp"
#include "tests/utils/callback_matchers.hpp"
#include "tests/utils/functors/increment.hpp"
#include "tests/utils/functors/load_check_add.hpp"
#include "tests/utils/functors/no_op.hpp"
#include "tests/utils/functors/sum_indices.hpp"
#include "tests/utils/functors/throws_when_copied.hpp"
#include "tests/utils/graph_context.hpp"
#include "tests/utils/just_stopped.hpp"
#include "tests/utils/sink_receiver.hpp"
#include "tests/utils/stdexec.hpp"
#include "tests/utils/sync_wait.hpp"

 
namespace Tests::GraphImpl {
 
using namespace Kokkos::utils::callbacks;
 
class BulkTest
    : public Tests::Utils::GraphContextTest<TEST_EXECUTION_SPACE>
    , public Kokkos::utils::tests::scoped::callbacks::Manager {
   public:
    using recorder_listener_t = RecorderListener<
        ConjunctionMatcher<EventDiscardMatcher<TEST_EXECUTION_SPACE>, GraphEventDiscardMatcher<TEST_EXECUTION_SPACE>>,
        BeginFenceEvent,
        BeginParallelForEvent,
        AllocateDataEvent,
        DeallocateDataEvent,
        Kokkos::Execution::Impl::RecordEvent,
        Kokkos::Execution::Impl::WaitEvent,
        Kokkos::Execution::GraphImpl::GraphAddNodeEvent,
        Kokkos::Execution::GraphImpl::GraphCreateEvent,
        Kokkos::Execution::GraphImpl::GraphInstantiateEvent,
        Kokkos::Execution::GraphImpl::GraphSubmitEvent
    >;
};

consteval bool test_sndr_traits() {
    using schd_sndr_t = typename BulkTest::schedule_sender_t;

    using functor_t = Tests::Utils::Functors::SumIndices<BulkTest::view_s_t>;
    using bulk_sndr_t = stdexec::transform_sender_result_t<
        decltype(stdexec::bulk(std::declval<schd_sndr_t>(), stdexec::par, 42, std::declval<functor_t>())),
        stdexec::env<>
    >;

    static_assert(Kokkos::Execution::GraphImpl::graph_completing_sender<bulk_sndr_t>);
    static_assert(std::same_as<Kokkos::Execution::Impl::exec_of_t<bulk_sndr_t>, TEST_EXECUTION_SPACE>);

    static_assert(!Kokkos::Execution::Impl::dispatching_sender<bulk_sndr_t>);
 
    return true;
}
 
static_assert(test_sndr_traits());

static_assert(Tests::Utils::check_continues_on_after_just_stopped<
              typename BulkTest::scheduler_t,
              stdexec::bulk_t,
              stdexec::parallel_policy,
              int,
              Tests::Utils::Functors::NoOp<false, false, false>
>());

consteval bool test_sndr_nothrow_transformable() {
    using sndr_bulk_t =
        decltype(stdexec::schedule(std::declval<typename BulkTest::scheduler_t>()) | stdexec::bulk(stdexec::par, 42, Tests::Utils::Functors::NoOp<false, false, false>{}));
 
    static_assert(
        std::same_as<
            sndr_bulk_t,
            stdexec::__sexpr<stdexec::bulk_t, stdexec::__data_of<sndr_bulk_t>, typename BulkTest::schedule_sender_t>
        >);
 
    static_assert(stdexec::__detail::__has_nothrow_transform_sender<
                  Kokkos::Execution::GraphImpl::Domain,
                  stdexec::set_value_t,
                  sndr_bulk_t&&,
                  stdexec::env<>
    >);
 
    using sndr_bulk_maythrow_on_move_t =
        decltype(stdexec::schedule(std::declval<typename BulkTest::scheduler_t>()) | stdexec::bulk(stdexec::par, 42, Tests::Utils::Functors::NoOp<false, false, true>{}));
 
    static_assert(!stdexec::__detail::__has_nothrow_transform_sender<
                  Kokkos::Execution::GraphImpl::Domain,
                  stdexec::set_value_t,
                  sndr_bulk_maythrow_on_move_t&&,
                  stdexec::env<>
    >);
 
    return true;
}
 
static_assert(test_sndr_nothrow_transformable());

consteval bool test_sndr_nothrow_connectable() {
    static_assert(!std::is_nothrow_constructible_v<Kokkos::Experimental::Graph<TEST_EXECUTION_SPACE>>);
 
    using sndr_bulk_t =
        decltype(stdexec::schedule(std::declval<typename BulkTest::scheduler_t>()) | stdexec::bulk(stdexec::par, 42, Tests::Utils::Functors::NoOp<false, false, false>{}));
 
    static_assert(!stdexec::__nothrow_connectable<sndr_bulk_t, Tests::Utils::SinkReceiver>);
 
    return true;
}
 
static_assert(test_sndr_nothrow_connectable());

consteval bool test_bulk_op_state_traits() {
    using functor_t = Tests::Utils::Functors::NoOp<false, false, false>;
 
    using sndr_t =
        decltype(stdexec::schedule(std::declval<typename BulkTest::scheduler_t>()) | stdexec::bulk(stdexec::par, 42, functor_t{}) | stdexec::bulk(stdexec::par, 42, functor_t{}));
    using connect_result_t = stdexec::connect_result_t<sndr_t, Tests::Utils::SinkReceiver>;
 
    static_assert(stdexec::__is_instance_of<connect_result_t, Kokkos::Execution::GraphImpl::OpState>);
 
    using root_t = typename BulkTest::graph_t::root_t;
    using range_policy_t =
        Kokkos::RangePolicy<TEST_EXECUTION_SPACE, Kokkos::IndexType<int>, Kokkos::Impl::IsGraphKernelTag>;
    using bulk_A_t = Kokkos::Experimental::GraphNodeRef<
        TEST_EXECUTION_SPACE,
        Kokkos::Impl::GraphNodeKernelImpl<TEST_EXECUTION_SPACE, range_policy_t, functor_t, Kokkos::ParallelForTag>,
        root_t
    >;
    using bulk_B_t = Kokkos::Experimental::GraphNodeRef<
        TEST_EXECUTION_SPACE,
        Kokkos::Impl::GraphNodeKernelImpl<TEST_EXECUTION_SPACE, range_policy_t, functor_t, Kokkos::ParallelForTag>,
        bulk_A_t
    >;
 
    static_assert(std::same_as<typename connect_result_t::predecessor_t, root_t>);
    static_assert(std::same_as<typename connect_result_t::node_t, bulk_B_t>);
 
    return true;
}
 
static_assert(test_bulk_op_state_traits());

TEST_F(BulkTest, bulk_schedule) {
    const view_s_t data(Kokkos::view_alloc(exec, "data - shared space"));
 
    const context_t gctx{exec};
 
    using load_check_add_t = Tests::Utils::Functors::LoadCheckAdd<value_t, on_device>;
    using sum_indices_t = Tests::Utils::Functors::SumIndices<view_s_t>;
 
    auto sndr = stdexec::schedule(gctx.get_scheduler()) | BULK_SUM_INDICES(4, data)
              | stdexec::then(load_check_add_t{.prev = 6, .value = 3, .data = data.data()}) | BULK_SUM_INDICES(4, data);
 
    using sndr_t = decltype(sndr);
 
    using bulk_data_t = stdexec::__data_of<sndr_t>;
 
    static_assert(std::same_as<
                  stdexec::transform_sender_result_t<sndr_t, stdexec::env<>>,
                  Kokkos::Execution::GraphImpl::ParallelForSender<
                      stdexec::then_t::__sender<
                          load_check_add_t,
                          stdexec::__sexpr<stdexec::bulk_t, bulk_data_t, typename BulkTest::schedule_sender_t>
                      >,
                      Kokkos::Execution::GraphImpl::node_props_t<TEST_EXECUTION_SPACE, std::string>,
                      sum_indices_t,
                      Kokkos::RangePolicy<TEST_EXECUTION_SPACE, Kokkos::IndexType<int>>
                  >
    >);

    static_assert(std::same_as<stdexec::__domain_of_t<stdexec::env_of_t<sndr_t>>, stdexec::default_domain>);
    static_assert(std::same_as<
                  stdexec::__detail::__completing_domain_t<stdexec::set_value_t, sndr_t>,
                  Kokkos::Execution::GraphImpl::Domain
    >);

    static_assert(std::same_as<
                  Kokkos::Execution::Impl::completion_scheduler_of_t<stdexec::set_value_t, sndr_t>,
                  typename BulkTest::scheduler_t
    >);
 
    ASSERT_EQ(data(), 0) << "Eager execution is not allowed.";
 
    const auto recorded_events = Tests::Utils::record_sync_wait<recorder_listener_t>(
        std::move(sndr)); // NOLINT(performance-move-const-arg)
 
    ASSERT_THAT(
        recorded_events,
        testing::ElementsAre(
            MATCHER_FOR_GRAPH_CREATE(device_handle),
            MATCHER_FOR_GRAPH_ADDNODE(
                recorded_events.at(0),
                device_handle,
                MATCHER_FOR_GRAPH_ROOT_NODE_OF(recorded_events.at(0)),
                dispatch_label(exec, "bulk")),
            MATCHER_FOR_GRAPH_ADDNODE(
                recorded_events.at(0),
                device_handle,
                MATCHER_FOR_GRAPH_NODE_OF(recorded_events.at(1)),
                dispatch_label(exec, "then")),
            MATCHER_FOR_GRAPH_ADDNODE(
                recorded_events.at(0),
                device_handle,
                MATCHER_FOR_GRAPH_NODE_OF(recorded_events.at(2)),
                dispatch_label(exec, "bulk")),
            MATCHER_FOR_GRAPH_SUBMIT(exec, recorded_events.at(0)),
            MATCHER_FOR_BEGIN_FENCE(exec, dispatch_label(exec, "sync_wait"))));
 
    ASSERT_EQ(data(), 15);
}

TEST_F(BulkTest, bulk_starts_on) {
    const view_s_t data(Kokkos::view_alloc(exec, "data - shared space"));
 
    const context_t gctx{exec};

    auto chain = stdexec::just() | BULK_SUM_INDICES(5, data);

    using chain_t = decltype(chain);
 
    static_assert(
        stdexec::get_completion_signatures<chain_t>()
        == stdexec::completion_signatures<stdexec::set_value_t(), stdexec::set_error_t(std::exception_ptr)>{});
    static_assert(!Tests::Utils::has_completion_scheduler_for<chain_t, stdexec::set_value_t>);
    static_assert(
        std::same_as<stdexec::__completion_domain_of_t<stdexec::set_value_t, chain_t>, stdexec::indeterminate_domain<>>);

    auto starts_on = stdexec::starts_on(gctx.get_scheduler(), std::move(chain));
 
    using starts_on_t = decltype(starts_on);

    static_assert(stdexec::__has_eptr_completion<chain_t>);
    static_assert(!stdexec::dependent_sender<starts_on_t>);
    static_assert(
        stdexec::get_completion_signatures<starts_on_t>()
        == stdexec::completion_signatures<stdexec::set_value_t(), stdexec::set_error_t(std::exception_ptr)>{});
 

    static_assert(stdexec::__completes_where_it_starts<stdexec::set_value_t, stdexec::env_of_t<chain_t>>);
    static_assert(Tests::Utils::has_completion_scheduler_for<starts_on_t, stdexec::set_value_t, stdexec::env<>>);
    static_assert(std::same_as<
                  Kokkos::Execution::Impl::completion_scheduler_of_t<stdexec::set_value_t, starts_on_t, stdexec::env<>>,
                  scheduler_t
    >);
    static_assert(std::same_as<
                  stdexec::__completion_domain_of_t<stdexec::set_value_t, starts_on_t, stdexec::env<>>,
                  Kokkos::Execution::GraphImpl::Domain
    >);
 
    ASSERT_EQ(data(), 0) << "Eager execution is not allowed.";
 
    const auto recorded_events = Tests::Utils::record_sync_wait<recorder_listener_t>(std::move(starts_on));
 
    ASSERT_THAT(
        recorded_events,
        testing::ElementsAre(
            MATCHER_FOR_GRAPH_CREATE(device_handle),
            MATCHER_FOR_GRAPH_ADDNODE(
                recorded_events.at(0), device_handle, MATCHER_FOR_GRAPH_ROOT_NODE_OF(recorded_events.at(0))),
            MATCHER_FOR_GRAPH_SUBMIT(exec, recorded_events.at(0)),
            MATCHER_FOR_BEGIN_FENCE(exec, dispatch_label(exec, "after dispatch"))));
 
    ASSERT_EQ(data(), 10);
}
 
} // namespace Tests::GraphImpl