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 "tests/utils/callback_matchers.hpp"
#include "tests/utils/execution_space_context.hpp"
#include "tests/utils/stdexec.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<BeginFenceEvent, BeginParallelForEvent, AllocateDataEvent, DeallocateDataEvent>;
    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(
        Tests::Utils::has_completion_signatures<decltype(allocate), stdexec::__mset<stdexec::set_value_t(view_of_5_t)>>);

#if defined(KOKKOS_COMPILER_CLANG) && defined(KOKKOS_ENABLE_CUDA)
    static_assert(Tests::Utils::has_completion_signatures<
                  decltype(allocate),
                  stdexec::__mset<stdexec::set_value_t(view_of_5_t), stdexec::set_error_t(std::exception_ptr)>,
                  stdexec::env<>
    >);
#else
    static_assert(Tests::Utils::has_completion_signatures<
                  decltype(allocate),
                  stdexec::__mset<stdexec::set_value_t(view_of_5_t)>,
                  stdexec::env<>
    >);
#endif

    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.";
 
    ASSERT_THAT(
        recorder_listener_t::record([check = std::move(check)]() mutable { // NOLINT(performance-move-const-arg)
            stdexec::sync_wait(std::move(check));                          // NOLINT(performance-move-const-arg)
        }),
        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, "sync_wait")),
            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