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 "tests/utils/callback_matchers.hpp"
#include "tests/utils/execution_space_context.hpp"
#include "tests/utils/functors/load_check_add.hpp"
#include "tests/utils/functors/sum_indices.hpp"
#include "tests/utils/kokkos.hpp"
#include "tests/utils/sink_receiver.hpp"

 
namespace Tests::ExecutionSpaceImpl {
 
using namespace Kokkos::utils::callbacks;
 
class 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>;
    using variant_t = typename recorder_listener_t::event_variant_t;
 
    static constexpr bool on_device = Tests::Utils::on_device<TEST_EXECUTION_SPACE>();
};

template <template <typename, typename, typename, typename> class SndrAdptr>
consteval bool test_sndr_traits() {
    using schd_sndr_t = typename ParallelForTest::schedule_sender_t;

    using label_t = std::string;
    using functor_t = Tests::Utils::Functors::SumIndices<typename ParallelForTest::view_s_t>;
    using policy_t = Kokkos::RangePolicy<TEST_EXECUTION_SPACE>;
    using pfor_sndr_t = SndrAdptr<schd_sndr_t, 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::ExecutionSpaceImpl::exec_of_t<pfor_sndr_t>, TEST_EXECUTION_SPACE>);

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

    using completion_signatures_t = stdexec::__completion_signatures_of_t<pfor_sndr_t, stdexec::env<>>;
 
    static_assert(stdexec::__mset_eq<
                  stdexec::__mset<stdexec::set_value_t(), stdexec::set_error_t(std::exception_ptr)>,
                  completion_signatures_t
    >);

    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<
                  stdexec::__completion_scheduler_of_t<stdexec::set_value_t, pfor_sndr_t, stdexec::env<>>,
                  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>>,
                  Kokkos::Execution::ExecutionSpaceImpl::ParallelForSender<schd_sndr_t, label_t, functor_t, policy_t>
    >);

#if defined(KOKKOS_COMPILER_CLANG) && defined(KOKKOS_ENABLE_CUDA)
    static_assert(!std::is_nothrow_move_constructible_v<typename ParallelForTest::view_s_t>);
    static_assert(!std::is_nothrow_move_constructible_v<functor_t>);
    using closure_t = Kokkos::Execution::ExecutionSpaceImpl::ParallelForClosure<label_t, functor_t, policy_t>;
    static_assert(!std::is_nothrow_move_constructible_v<closure_t>);
    static_assert(!stdexec::__nothrow_connectable<pfor_sndr_t, Tests::Utils::SinkReceiver>);
#else
    static_assert(stdexec::__nothrow_connectable<pfor_sndr_t, Tests::Utils::SinkReceiver>);
#endif
 
    return true;
}
static_assert(test_sndr_traits<Kokkos::Execution::Impl::ParallelForSender>());
static_assert(test_sndr_traits<Kokkos::Execution::ExecutionSpaceImpl::ParallelForSender>());

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

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

    static_assert(std::same_as<stdexec::tag_of_t<pfor_sndr_t>, Kokkos::Execution::parallel_for_t>);
 
    static_assert(std::same_as<
                  stdexec::__data_of<pfor_sndr_t>,
                  Kokkos::Execution::Impl::ParallelForData<label_t, functor_t, policy_t>
    >);
 
    static_assert(stdexec::__nbr_children_of<pfor_sndr_t> == 1);
    static_assert(std::same_as<stdexec::__child_of<pfor_sndr_t>, schd_sndr_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_decomposition());

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

    static_assert(Kokkos::Execution::ExecutionSpaceImpl::Closure<closure_t>);
 
    static_assert(std::is_nothrow_move_constructible_v<closure_t> == ExpectNoThrowMoveConstructible);
 
    return true;
}

#if defined(KOKKOS_COMPILER_CLANG) && defined(KOKKOS_ENABLE_CUDA)
static_assert(test_closure_traits<typename ParallelForTest::view_s_t, false>());
#else
static_assert(test_closure_traits<typename ParallelForTest::view_s_t, true>());
#endif
static_assert(test_closure_traits<std::span<int>, true>());

TEST_F(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 {
    auto chain = 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});
 
    if constexpr (std::same_as<Exec, Kokkos::DefaultExecutionSpace>) {
        return std::move(chain)
             | 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});
    } else {
        return chain;
    }
}

TEST_F(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 = recorder_listener_t::record(
        [chain = closure_object_creation_overloads(size, witness, context_t{exec})]() mutable {
            stdexec::sync_wait(std::move(chain));
        });
 
    unsigned short int ievent = 0;
 
    ASSERT_GE(recorded_events.size(), 3);
 
    using functor_t = Tests::Utils::Functors::SumIndices<view_s_t>;
 
    ASSERT_THAT(
        recorded_events,
        ElementAt<variant_t>(ievent++, MATCHER_FOR_BEGIN_PFOR(exec, "passing label, execution policy and functor")));
    ASSERT_THAT(
        recorded_events,
        ElementAt<variant_t>(ievent++, MATCHER_FOR_BEGIN_PFOR(exec, Kokkos::Impl::TypeInfo<functor_t>::name())));
 
    if constexpr (std::same_as<TEST_EXECUTION_SPACE, Kokkos::DefaultExecutionSpace>) {
        ASSERT_THAT(
            recorded_events.at(ievent++), MATCHER_FOR_BEGIN_PFOR(exec, "passing label, work count and functor"));
        ASSERT_THAT(
            recorded_events.at(ievent++), MATCHER_FOR_BEGIN_PFOR(exec, Kokkos::Impl::TypeInfo<functor_t>::name()));
    }
 
    ASSERT_THAT(recorded_events.at(ievent), MATCHER_FOR_BEGIN_FENCE(exec, dispatch_label(exec, "sync_wait")));
 
    ASSERT_EQ(witness(), ievent * size / 2 * (size - 1));
}

TEST_F(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::Impl::TypeInfo<TEST_EXECUTION_SPACE>::name()),
                     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::Impl::TypeInfo<TEST_EXECUTION_SPACE>::name()),
                     Kokkos::RangePolicy<TEST_EXECUTION_SPACE>(0, 2 * size),
                     Tests::Utils::Functors::SumIndices{.data = witness});
 
    ASSERT_THAT(
        recorder_listener_t::record([chain = std::move(chain)]() mutable { stdexec::sync_wait(std::move(chain)); }),
        testing::ElementsAre(
            MATCHER_FOR_BEGIN_PFOR(exec, dispatch_label(exec, "hello from pfor")),
            MATCHER_FOR_BEGIN_PFOR(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(ParallelForTest, starts_on_parallel_region) {
    constexpr size_t size = 10;
 
    const view_s_t witness(Kokkos::view_alloc(exec, "data - shared space"));
 
    auto sndr = stdexec::just()
              | Kokkos::Execution::parallel_for(
                    std::format("{}: hello from pfor", Kokkos::Impl::TypeInfo<TEST_EXECUTION_SPACE>::name()),
                    Kokkos::RangePolicy<TEST_EXECUTION_SPACE>(0, size),
                    Tests::Utils::Functors::SumIndices{.data = witness});
 
    const context_t esc{exec};
    auto starts_on = stdexec::starts_on(esc.get_scheduler(), std::move(sndr));
 
    ASSERT_THAT(
        recorder_listener_t::record(
            [starts_on = std::move(starts_on)]() mutable { stdexec::sync_wait(std::move(starts_on)); }),
        testing::ElementsAre(
            MATCHER_FOR_BEGIN_PFOR(exec, dispatch_label(exec, "hello from pfor")),
            MATCHER_FOR_BEGIN_FENCE(exec, dispatch_label(exec, "sync_wait"))));
 
    ASSERT_EQ(witness(), size / 2 * (size - 1));
}
 
} // namespace Tests::ExecutionSpaceImpl