test_let_value.cpp

Go to the documentation of this file.

#include "kokkos-execution/utils/ignore_warnings.hpp"
PRAGMA_DIAGNOSTIC_PUSH
KOKKOS_EXECUTION_STDEXEC_PRAGMA_DIAGNOSTIC_IGNORED
#include "exec/single_thread_context.hpp"
PRAGMA_DIAGNOSTIC_POP
 
#include "kokkos-utils/callbacks/RecorderListener.hpp"
#include "kokkos-utils/concepts/View.hpp"
#include "kokkos-utils/tests/scoped/callbacks/Manager.hpp"
 
#include "kokkos-execution/execution_space.hpp"
 
#include "tests/utils/callback_matchers.hpp"
#include "tests/utils/execution_space_context.hpp"
#include "tests/utils/stdexec.hpp"
#include "tests/utils/sync_wait.hpp"

 
namespace Tests::ExecutionSpaceImpl {
 
using namespace Kokkos::utils::callbacks;
 
class LetValueTest
    : 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,
        AllocateDataEvent,
        DeallocateDataEvent,
        Kokkos::Execution::Impl::RecordEvent,
        Kokkos::Execution::Impl::WaitEvent
    >;
    using variant_t = typename recorder_listener_t::event_variant_t;
};

TEST_F(LetValueTest, scoped_allocation) {
    const view_s_t data(Kokkos::view_alloc(exec, "data - shared space"));
 
    const context_t esc{exec};
 
    experimental::execution::single_thread_context stc{};
 
    using view_of_5_t = Kokkos::View<value_t[5], typename TEST_EXECUTION_SPACE::memory_space>;

    auto allocate = stdexec::schedule(esc.get_scheduler())
                  | stdexec::let_value([this]() noexcept -> stdexec::sender auto {
                        return stdexec::just(
                            view_of_5_t{Kokkos::view_alloc(exec, "scratch", Kokkos::WithoutInitializing)});
                    });
 
    static_assert(stdexec::dependent_sender<decltype(allocate)>);

    static_assert(Tests::Utils::has_completion_signatures<
                  decltype(allocate),
                  stdexec::__mset<stdexec::set_value_t(view_of_5_t)>,
                  stdexec::env<>
    >);

    auto run = std::move(allocate) // NOLINT(performance-move-const-arg)
             | stdexec::let_value([&esc, &data, &stc](auto&& scratch) noexcept -> stdexec::sender auto {
                   static_assert(std::same_as<decltype(scratch), view_of_5_t&>);
                   EXPECT_EQ(scratch.use_count(), 1);
                   return stdexec::schedule(esc.get_scheduler())
                        | stdexec::bulk(
                              stdexec::par,
                              5,
                              KOKKOS_LAMBDA<std::integral T>(const T index) {
                                  KOKKOS_IF_ON_HOST(EXPECT_EQ(scratch.use_count(), 2);)
                                  scratch(index) = index;
                                  Kokkos::atomic_add(data.data(), scratch(index));
                              })
                        | stdexec::continues_on(stc.get_scheduler())
                        | stdexec::then([&scratch]() mutable -> Kokkos::utils::concepts::ViewOfRank<1> auto {
                              EXPECT_EQ(scratch.use_count(), 2);
                              return std::move(scratch);
                          });
               });
 
    static_assert(Tests::Utils::has_completion_signatures<
                  decltype(run),
                  stdexec::__mset<stdexec::set_value_t(view_of_5_t), stdexec::set_error_t(std::exception_ptr)>,
                  stdexec::env<>
    >);

    auto check = std::move(run) // NOLINT(performance-move-const-arg)
               | stdexec::let_value([&esc, &data](auto&& scratch) noexcept -> stdexec::sender auto {
                     static_assert(std::same_as<decltype(scratch), view_of_5_t&>);
                     EXPECT_EQ(scratch.use_count(), 2);
                     return stdexec::schedule(esc.get_scheduler())
                          | stdexec::bulk(
                                stdexec::par, 5, KOKKOS_LAMBDA<std::integral T>(const T index) {
                                    Kokkos::atomic_add(data.data(), scratch(index));
                                    KOKKOS_IF_ON_HOST(EXPECT_EQ(scratch.use_count(), 2);)
                                });
                 });
 
    static_assert(Tests::Utils::has_completion_signatures<
                  decltype(check),
                  stdexec::__mset<stdexec::set_value_t(), stdexec::set_error_t(std::exception_ptr)>,
                  stdexec::env<>
    >);
 
    ASSERT_EQ(data(), 0) << "Eager execution is not allowed.";
 
    KOKKOS_EXECUTION_THREADS_THROWS_ON_SYNC_WAIT_ASSERT_AND_SKIP(check)
 
    const auto recorded_events = Tests::Utils::record_sync_wait<recorder_listener_t>(
        std::move(check)); // NOLINT(performance-move-const-arg)
 
    ASSERT_THAT(recorded_events, [&]() {
        if constexpr (Kokkos::Execution::Impl::has_non_blocking_dispatch<TEST_EXECUTION_SPACE>) {
            return ContainsInOrder<variant_t>(
                Kokkos::utils::callbacks::AAllocateDataEvent(
                    testing::Field(
                        &Kokkos::utils::callbacks::AllocateDataEvent::alloc,
                        testing::Field(&Kokkos::utils::callbacks::AllocDescriptor::name, testing::StrEq("scratch")))),
                MATCHER_FOR_BEGIN_PFOR(exec, dispatch_label(exec, "bulk")),
                MATCHER_FOR_BEGIN_FENCE(exec, dispatch_label(exec, "schedule_from")),
                MATCHER_FOR_BEGIN_PFOR(exec, dispatch_label(exec, "bulk")),
                MATCHER_FOR_RECORD_EVENT(exec),
                MATCHER_FOR_WAIT_EVENT(recorded_events.at(4)),
                Kokkos::utils::callbacks::ADeallocateDataEvent(
                    testing::Field(
                        &Kokkos::utils::callbacks::DeallocateDataEvent::alloc,
                        testing::Field(&Kokkos::utils::callbacks::AllocDescriptor::name, testing::StrEq("scratch")))));
        } else {
            return ContainsInOrder<variant_t>(
                Kokkos::utils::callbacks::AAllocateDataEvent(
                    testing::Field(
                        &Kokkos::utils::callbacks::AllocateDataEvent::alloc,
                        testing::Field(&Kokkos::utils::callbacks::AllocDescriptor::name, testing::StrEq("scratch")))),
                MATCHER_FOR_BEGIN_PFOR(exec, dispatch_label(exec, "bulk")),
                MATCHER_FOR_BEGIN_FENCE(exec, dispatch_label(exec, "schedule_from")),
                MATCHER_FOR_BEGIN_PFOR(exec, dispatch_label(exec, "bulk")),
                MATCHER_FOR_BEGIN_FENCE(exec, dispatch_label(exec, "after dispatch")),
                Kokkos::utils::callbacks::ADeallocateDataEvent(
                    testing::Field(
                        &Kokkos::utils::callbacks::DeallocateDataEvent::alloc,
                        testing::Field(&Kokkos::utils::callbacks::AllocDescriptor::name, testing::StrEq("scratch")))));
        }
    }());
 
    ASSERT_EQ(data(), 2 * (0 + 1 + 2 + 3 + 4));
}
 
} // namespace Tests::ExecutionSpaceImpl