test_parallel_for.cpp

Go to the documentation of this file.

#include <bit>
 
#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/category.hpp"
#include "tests/utils/check_rcvr_env_queryable_with.hpp"
#include "tests/utils/execution_space_context.hpp"
#include "tests/utils/functors/counter.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/just_stopped.hpp"
#include "tests/utils/sink_receiver.hpp"
#include "tests/utils/sync_wait.hpp"
#include "tests/utils/tracking_allocator.hpp"

 
namespace Tests::ExecutionSpaceImpl {
 
using namespace Kokkos::utils::callbacks;
 
class TEST_CATEGORY(ParallelForTest)
    : 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,
        BeginSingleEvent,
        Kokkos::Execution::Impl::RecordEvent,
        Kokkos::Execution::Impl::WaitEvent
    >;
};

template <
    template <typename...> typename SndrAdptr,
    bool ExptIsDispatchingSender,
    template <typename> typename ExptTransformSenderType
>
consteval bool test_sndr_traits() {
    using schd_sndr_t = typename TEST_CATEGORY(ParallelForTest)::schedule_sender_t;

    using label_t = std::string;
    using functor_t = Tests::Utils::Functors::SumIndices<typename TEST_CATEGORY(ParallelForTest)::view_s_t>;
    using policy_t = Kokkos::RangePolicy<TEST_EXECUTION_SPACE>;
    using pfor_data_t =
        Kokkos::Execution::Impl::ParallelForData<functor_t, policy_t, Kokkos::Execution::Impl::Label<label_t>>;
    using pfor_sndr_t = SndrAdptr<schd_sndr_t, TEST_EXECUTION_SPACE, label_t, functor_t, policy_t>;

    static_assert(Kokkos::Execution::ExecutionSpaceImpl::execution_space_completing_sender<pfor_sndr_t>);
    static_assert(std::same_as<Kokkos::Execution::Impl::exec_of_t<pfor_sndr_t>, TEST_EXECUTION_SPACE>);

    static_assert(Kokkos::Execution::Impl::dispatching_sender<pfor_sndr_t> == ExptIsDispatchingSender);

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

    static_assert(std::same_as<
                  stdexec::__completion_domain_of_t<stdexec::set_value_t, pfor_sndr_t, stdexec::env<>>,
                  Kokkos::Execution::ExecutionSpaceImpl::Domain
    >);

    static_assert(std::same_as<
                  Kokkos::Execution::Impl::completion_scheduler_of_t<stdexec::set_value_t, pfor_sndr_t>,
                  Kokkos::Execution::ExecutionSpaceImpl::Scheduler<TEST_EXECUTION_SPACE>
    >);

    static_assert(stdexec::sender_to<pfor_sndr_t, Tests::Utils::SinkReceiver>);
 
    static_assert(std::same_as<
                  stdexec::transform_sender_result_t<pfor_sndr_t, stdexec::env_of_t<Tests::Utils::SinkReceiver>>,
                  typename ExptTransformSenderType<Kokkos::Execution::ExecutionSpaceImpl::ParallelForSender< // NOLINT(modernize-type-traits)
                      Kokkos::Execution::parallel_for_t,
                      schd_sndr_t,
                      TEST_EXECUTION_SPACE,
                      pfor_data_t
                  >>::type
    >);

    static_assert(stdexec::__nothrow_connectable<pfor_sndr_t, Tests::Utils::SinkReceiver>);
 
    return true;
}
 
template <typename Sndr, typename Exec, typename Label, typename Functor, typename Policy>
using pfor_generic_t = Kokkos::Execution::Impl::ParallelForSender<
    Sndr,
    Kokkos::Execution::Impl::ParallelForData<Functor, Policy, Kokkos::Execution::Impl::Label<Label>>
>;
 
template <typename Sndr, typename Exec, typename Label, typename Functor, typename Policy>
using pfor_impl_t = Kokkos::Execution::ExecutionSpaceImpl::ParallelForSender<
    Kokkos::Execution::parallel_for_t,
    Sndr,
    Exec,
    Kokkos::Execution::Impl::ParallelForData<Functor, Policy, Kokkos::Execution::Impl::Label<Label>>
>;
 
static_assert(test_sndr_traits<pfor_generic_t, true, std::type_identity>());
static_assert(test_sndr_traits<pfor_impl_t, false, std::add_rvalue_reference>());

consteval bool test_sndr_is_transformable() {
    using schd_sndr_t = typename TEST_CATEGORY(ParallelForTest)::schedule_sender_t;

    using label_t = Kokkos::Execution::Impl::Label<std::string>;
    using functor_t = Tests::Utils::Functors::SumIndices<typename TEST_CATEGORY(ParallelForTest)::view_s_t>;
    using policy_t = Kokkos::RangePolicy<TEST_EXECUTION_SPACE>;
    using pfor_data_t = Kokkos::Execution::Impl::ParallelForData<functor_t, policy_t, label_t>;
    using pfor_sndr_t = Kokkos::Execution::Impl::ParallelForSender<schd_sndr_t, pfor_data_t>;

    static_assert(stdexec::__applicable<
                  Kokkos::Execution::ExecutionSpaceImpl::TransformSenderFor<stdexec::tag_of_t<pfor_sndr_t>>,
                  pfor_sndr_t,
                  const stdexec::env<>&
    >);
 
    return true;
}
 
static_assert(test_sndr_is_transformable());

template <typename IterSpace>
consteval bool test_pfor_sndr_with_index_space_transforms_into_pfor_sndr_with_range_policy() {
    using schd_sndr_t = typename TEST_CATEGORY(ParallelForTest)::schedule_sender_t;

    using label_t = Kokkos::Execution::Impl::Label<std::string>;
    using functor_t = Tests::Utils::Functors::SumIndices<typename TEST_CATEGORY(ParallelForTest)::view_s_t>;
    using pfor_data_t = Kokkos::Execution::Impl::ParallelForData<functor_t, IterSpace, label_t>;
    using pfor_sndr_t = Kokkos::Execution::Impl::ParallelForSender<schd_sndr_t, pfor_data_t>;
 
    static_assert(stdexec::__detail::__has_nothrow_transform_sender<
                  Kokkos::Execution::ExecutionSpaceImpl::Domain,
                  stdexec::set_value_t,
                  pfor_sndr_t&&,
                  stdexec::env<>
    >);

    static_assert(
        std::same_as<
            decltype(std::declval<typename stdexec::transform_sender_result_t<pfor_sndr_t, stdexec::env<>>::closure_t>()
                         .get_policy()),
            Kokkos::RangePolicy<TEST_EXECUTION_SPACE>
        >);
 
    return true;
}
 
static_assert(
    test_pfor_sndr_with_index_space_transforms_into_pfor_sndr_with_range_policy<Kokkos::range_slice<size_t, size_t>>());
static_assert(test_pfor_sndr_with_index_space_transforms_into_pfor_sndr_with_range_policy<Kokkos::dims<1>>());

template <typename ViewType>
consteval bool test_closure_traits() {
    using functor_t = Tests::Utils::Functors::SumIndices<ViewType>;
    using policy_t = Kokkos::RangePolicy<TEST_EXECUTION_SPACE>;
    using pfor_data_t =
        Kokkos::Execution::Impl::ParallelForData<functor_t, policy_t, Kokkos::Execution::Impl::Label<std::string_view>>;
    using closure_t = Kokkos::Execution::ExecutionSpaceImpl::ParallelForClosure<TEST_EXECUTION_SPACE, pfor_data_t>;

    static_assert(Kokkos::Execution::ExecutionSpaceImpl::Closure<closure_t>);
 
    static_assert(std::is_nothrow_move_constructible_v<closure_t>);
 
    return true;
}
 
static_assert(test_closure_traits<typename TEST_CATEGORY(ParallelForTest)::view_s_t>());
static_assert(test_closure_traits<std::span<int>>());

static_assert(Tests::Utils::check_continues_on_after_just_stopped<
              typename TEST_CATEGORY(ParallelForTest)::scheduler_t,
              Kokkos::Execution::parallel_for_t,
              Kokkos::RangePolicy<TEST_EXECUTION_SPACE>,
              Tests::Utils::Functors::NoOp<false, false, false>
>());

TEST_F(TEST_CATEGORY(ParallelForTest), team_policy) {
    constexpr int size = 32;
 
    const auto [num_teams, team_size] = [&]() {
#if defined(KOKKOS_ENABLE_HPX)
        if constexpr (std::same_as<TEST_EXECUTION_SPACE, Kokkos::Experimental::HPX>) {
            return std::make_tuple(size, 1);
        }
#endif
        const int team_size_ = std::bit_floor(static_cast<unsigned short>(std::min(exec.concurrency(), size / 2)));
        return std::make_tuple(size / team_size_, team_size_);
    }();
 
    ASSERT_EQ(team_size * num_teams, size);
 
    const view_s_t witness(Kokkos::view_alloc(exec, "data - shared space"));
 
    const context_t esc{exec};
 
    auto chain = stdexec::schedule(esc.get_scheduler())
               | Kokkos::Execution::parallel_for(
                     "hello from pfor",
                     Kokkos::TeamPolicy<TEST_EXECUTION_SPACE>(num_teams, team_size),
                     Tests::Utils::Functors::SumIndices{.data = witness});
 
    stdexec::sync_wait(std::move(chain));
 
    ASSERT_EQ(witness(), size / 2 * (size - 1));
}
 
template <typename ViewType, Kokkos::ExecutionSpace Exec>
auto closure_object_creation_overloads(
    const size_t size,
    const ViewType& witness,
    const Kokkos::Execution::ExecutionSpaceContext<Exec>& esc) -> stdexec::sender auto {
    return stdexec::schedule(esc.get_scheduler())
         | Kokkos::Execution::parallel_for(
               "passing label, execution policy and functor",
               Kokkos::RangePolicy<Exec>(0, size),
               Tests::Utils::Functors::SumIndices{.data = witness})
         | Kokkos::Execution::parallel_for(
               Kokkos::RangePolicy<Exec>(0, size), Tests::Utils::Functors::SumIndices{.data = witness})
         | Kokkos::Execution::parallel_for(
               "passing label, work count and functor", size, Tests::Utils::Functors::SumIndices{.data = witness})
         | Kokkos::Execution::parallel_for(size, Tests::Utils::Functors::SumIndices{.data = witness})
         | Kokkos::Execution::parallel_for(
               Kokkos::range_slice<size_t, size_t>{4, 6}, Tests::Utils::Functors::SumIndices{.data = witness});
}

TEST_F(TEST_CATEGORY(ParallelForTest), closure_object_creation_overloads) {
    constexpr size_t size = 10;
 
    const view_s_t witness(Kokkos::view_alloc(exec, "data - shared space"));
 
    const auto recorded_events = Tests::Utils::record_sync_wait<recorder_listener_t>(
        closure_object_creation_overloads(size, witness, context_t{exec}));
 
    using functor_t = Tests::Utils::Functors::SumIndices<view_s_t>;
 
    ASSERT_THAT(
        recorded_events,
        testing::ElementsAre(
            MATCHER_FOR_BEGIN_PFOR(exec, "passing label, execution policy and functor"),
            MATCHER_FOR_BEGIN_PFOR(exec, Kokkos::Execution::Impl::display_string_of<functor_t>()),
            MATCHER_FOR_BEGIN_PFOR(exec, "passing label, work count and functor"),
            MATCHER_FOR_BEGIN_PFOR(exec, Kokkos::Execution::Impl::display_string_of<functor_t>()),
            MATCHER_FOR_BEGIN_PFOR(exec, Kokkos::Execution::Impl::display_string_of<functor_t>()),
            MATCHER_FOR_BEGIN_FENCE(exec, dispatch_label(exec, "sync_wait"))));
 
    ASSERT_EQ(witness(), 4 * size / 2 * (size - 1) + (4 + 5));
}

TEST_F(TEST_CATEGORY(ParallelForTest), two_parallel_regions) {
    constexpr size_t size = 10;
 
    const view_s_t witness(Kokkos::view_alloc(exec, "data - shared space"));
 
    const context_t esc{exec};
 
    auto chain =
        stdexec::schedule(esc.get_scheduler())
        | Kokkos::Execution::parallel_for(
            std::format("{}: hello from pfor", Kokkos::Execution::Impl::display_string_of<TEST_EXECUTION_SPACE>()),
            Kokkos::RangePolicy<TEST_EXECUTION_SPACE>(0, size),
            Tests::Utils::Functors::SumIndices{.data = witness})
        | stdexec::then(
            Tests::Utils::Functors::LoadCheckAdd<value_t, on_device>{
                .prev = size / 2 * (size - 1), .value = 4, .data = witness.data()})
        | Kokkos::Execution::parallel_for(
            std::format(
                "{}: hello again from pfor", Kokkos::Execution::Impl::display_string_of<TEST_EXECUTION_SPACE>()),
            Kokkos::RangePolicy<TEST_EXECUTION_SPACE>(0, 2 * size),
            Tests::Utils::Functors::SumIndices{.data = witness});
 
    ASSERT_THAT(
        Tests::Utils::record_sync_wait<recorder_listener_t>(std::move(chain)),
        testing::ElementsAre(
            MATCHER_FOR_BEGIN_PFOR(exec, dispatch_label(exec, "hello from pfor")),
            MATCHER_FOR_BEGIN_SINGLE(exec, dispatch_label(exec, "then")),
            MATCHER_FOR_BEGIN_PFOR(exec, dispatch_label(exec, "hello again from pfor")),
            MATCHER_FOR_BEGIN_FENCE(exec, dispatch_label(exec, "sync_wait"))));
 
    ASSERT_EQ(witness(), size / 2 * (size - 1) + 4 + 2 * size * (2 * size - 1) / 2);
}

TEST_F(TEST_CATEGORY(ParallelForTest), starts_on_parallel_region) {
    constexpr size_t size = 10;
 
    const view_s_t witness(Kokkos::view_alloc(exec, "data - shared space"));
 
    auto chain =
        stdexec::just()
        | Kokkos::Execution::parallel_for(
            std::format("{}: hello from pfor", Kokkos::Execution::Impl::display_string_of<TEST_EXECUTION_SPACE>()),
            Kokkos::RangePolicy<TEST_EXECUTION_SPACE>(0, size),
            Tests::Utils::Functors::SumIndices{.data = witness});
 
    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<>>);
 
    const context_t esc{exec};
    auto starts_on = stdexec::starts_on(esc.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::__is_instance_of<
                  stdexec::transform_sender_result_t<
                      decltype(starts_on),
                      stdexec::env_of_t<Kokkos::Execution::Impl::SyncWait::Receiver<TEST_EXECUTION_SPACE>>
                  >,
                  stdexec::__seq::__sndr
    >);
 
    const auto recorded_events = Tests::Utils::record_sync_wait<recorder_listener_t>(std::move(starts_on));
 
    ASSERT_THAT(recorded_events, [&]() {
        if constexpr (Kokkos::Execution::Impl::has_non_blocking_dispatch<TEST_EXECUTION_SPACE>) {
            return testing::ElementsAre(
                MATCHER_FOR_BEGIN_PFOR(exec, dispatch_label(exec, "hello from pfor")),
                MATCHER_FOR_RECORD_EVENT(exec),
                MATCHER_FOR_WAIT_EVENT(recorded_events.at(1)));
        } else {
            return testing::ElementsAre(
                MATCHER_FOR_BEGIN_PFOR(exec, dispatch_label(exec, "hello from pfor")),
                MATCHER_FOR_BEGIN_FENCE(exec, dispatch_label(exec, "after dispatch")));
        }
    }());
 
    ASSERT_EQ(witness(), size / 2 * (size - 1));
}

TEST_F(TEST_CATEGORY(ParallelForTest), forwarding_env) {
    constexpr size_t size = 10;
 
    const view_s_t data(Kokkos::view_alloc(exec, "data - shared space"));
 
    std::atomic<size_t> count = 0;
 
    int value;
 
    const context_t esc{exec};
 
    stdexec::sender auto sndr =
        stdexec::read_env(stdexec::get_allocator)
        | stdexec::then([&value](auto allocator) { value = Tests::Utils::round_trip_allocate(allocator, 42); })
        | stdexec::continues_on(esc.get_scheduler())
        | Tests::Utils::check_rcvr_env_queryable_with<stdexec::get_allocator_t>()
        | Kokkos::Execution::parallel_for(
            "my pfor",
            Kokkos::RangePolicy<TEST_EXECUTION_SPACE>(0, size),
            Tests::Utils::Functors::SumIndices{.data = data})
        | stdexec::write_env(stdexec::prop{stdexec::get_allocator, Tests::Utils::TrackingAllocator<int>{&count}});
 
    ASSERT_EQ(data(), 0) << "Eager execution is not allowed.";
 
    const auto recorded_events = Tests::Utils::record_sync_wait<recorder_listener_t>(std::move(sndr));
 
    ASSERT_THAT(recorded_events, [&]() {
        if constexpr (Kokkos::Execution::Impl::has_non_blocking_dispatch<TEST_EXECUTION_SPACE>) {
            return testing::ElementsAre(
                MATCHER_FOR_BEGIN_PFOR(exec, "my pfor"),
                MATCHER_FOR_RECORD_EVENT(exec),
                MATCHER_FOR_WAIT_EVENT(recorded_events.at(1)));
        } else {
            return testing::ElementsAre(
                MATCHER_FOR_BEGIN_PFOR(exec, "my pfor"),
                MATCHER_FOR_BEGIN_FENCE(exec, dispatch_label(exec, "after dispatch")));
        }
    }());
 
    ASSERT_EQ(data(), size / 2 * (size - 1));
 
    ASSERT_EQ(value, 42);
    ASSERT_EQ(count, 1);
}

TEST_F(TEST_CATEGORY(ParallelForTest), perfect_forwarding_connect) {
    using counter_t = Tests::Utils::Functors::Counter;
 
    const context_t esc{exec};
 
    const Kokkos::RangePolicy<TEST_EXECUTION_SPACE> policy(0, 1);
    counter_t counter;
 
    counter_t::reset();
 
    [[maybe_unused]]
    const stdexec::operation_state auto opstate_from_lvalue_functor = stdexec::connect(
        stdexec::schedule(esc.get_scheduler()) | Kokkos::Execution::parallel_for("label", policy, counter),
        Tests::Utils::SinkReceiver{});
 
    ASSERT_EQ(counter_t::copy_assignments, 0);
    ASSERT_EQ(counter_t::copy_constructions, 1);
    ASSERT_EQ(counter_t::move_assignments, 0);
    ASSERT_EQ(counter_t::move_constructions, 6);
 
    counter_t::reset();
 
    [[maybe_unused]]
    const stdexec::operation_state auto closure_from_xvalue_functor = stdexec::connect(
        stdexec::schedule(esc.get_scheduler()) | Kokkos::Execution::parallel_for("label", policy, std::move(counter)),
        Tests::Utils::SinkReceiver{});
 
    ASSERT_EQ(counter_t::copy_assignments, 0);
    ASSERT_EQ(counter_t::copy_constructions, 0);
    ASSERT_EQ(counter_t::move_assignments, 0);
    ASSERT_EQ(counter_t::move_constructions, 7);
 
    counter_t::reset();
 
    [[maybe_unused]]
    const stdexec::operation_state auto closure_from_prvalue_functor = stdexec::connect(
        stdexec::schedule(esc.get_scheduler()) | Kokkos::Execution::parallel_for("label", policy, counter_t{}),
        Tests::Utils::SinkReceiver{});
 
    ASSERT_EQ(counter_t::copy_assignments, 0);
    ASSERT_EQ(counter_t::copy_constructions, 0);
    ASSERT_EQ(counter_t::move_assignments, 0);
    ASSERT_EQ(counter_t::move_constructions, 7);
}
 
} // namespace Tests::ExecutionSpaceImpl