test: refactor + add new MD view helpers

Refactored implementation of view builder helpers and added new ones,
all now MD-capable.
- Refactored `ViewBuilder` struct into a free function that allows
building views of arbitrary type and dimension, contiguous or not.
- Added view initialization and error counting helpers. Basically, these
were the problematic pieces of code that forced us to have dedicated
test functions for each view dimension we wanted to test. These now
abstract this away and are MD-capable.

Theses changes make it possible to write only the high-level test
function, and pass different kind of views to it (which better mirrors
user code that Kokkos Comm is designed to handle correctly). It's now
trivial to test views of higher dimensions (>2), contiguous or not,
provided that we support them (which is not the case when they aren't
contiguous, due to limitations caused by a similar implementation
strategy as what was refactored here).

Author: dssgabriel

unit_tests/CMakeLists.txt

@@ -46,7 +46,7 @@ list(APPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/../cmake")
 include(kc-test)
 
 # --- Core API unit tests --- #
-kc_add_unit_test(test.core.p2p CORE NUM_PES 2 FILES test_main.cpp test_sendrecv.cpp)
+kc_add_unit_test(test.core.send-recv CORE NUM_PES 2 FILES test_main.cpp test_send_recv.cpp)
 kc_add_unit_test(test.core.broadcast CORE NUM_PES 2 FILES test_main.cpp test_broadcast.cpp)
 kc_add_unit_test(test.core.all-gather CORE NUM_PES 2 FILES test_main.cpp test_allgather.cpp)
 kc_add_unit_test(test.core.reduce CORE NUM_PES 2 FILES test_main.cpp test_reduce.cpp)

unit_tests/test_send_recv.cpp

@@ -0,0 +1,76 @@
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+// SPDX-FileCopyrightText: Copyright Contributors to the Kokkos project
+
+#include <gtest/gtest.h>
+#include <KokkosComm/KokkosComm.hpp>
+
+#include "view_utils.hpp"
+#if defined(KOKKOSCOMM_ENABLE_NCCL)
+#include "nccl/utils.hpp"
+#endif
+
+namespace {
+
+template <typename T>
+class SendRecv : public testing::Test {
+ public:
+  using Scalar = T;
+};
+
+#if defined(KOKKOSCOMM_ENABLE_NCCL)
+using ScalarTypes = ::testing::Types<float, double, int, int64_t>;
+#else
+using ScalarTypes =
+    ::testing::Types<float, double, Kokkos::complex<float>, Kokkos::complex<double>, int, unsigned, int64_t, size_t>;
+#endif
+TYPED_TEST_SUITE(SendRecv, ScalarTypes);
+
+template <KokkosComm::KokkosView View>
+void test_core_send_recv(const View& v) {
+  using Exec = Kokkos::DefaultExecutionSpace;
+  using Comm = KokkosComm::DefaultCommunicationSpace;
+#if defined(KOKKOSCOMM_ENABLE_NCCL)
+  auto nccl_ctx = test_utils::nccl::Ctx::init();
+  auto raw_comm = nccl_ctx.comm();
+#else
+  auto raw_comm = MPI_COMM_WORLD;
+#endif
+  auto exec      = Exec();
+  auto comm      = KokkosComm::Communicator<Comm, Exec>::duplicate(raw_comm, exec).value();
+  const int size = comm.size();
+  const int rank = comm.rank();
+  if (size < 2) {
+    GTEST_SKIP() << "Requires >= 2 ranks (" << size << " provided)";
+  }
+  const int src = 0;
+  const int dst = 1;
+
+  if (rank == src) {
+    test_utils::init_view(exec, v);
+    exec.fence();
+    KokkosComm::send(comm, v, dst).wait();
+  } else if (rank == dst) {
+    KokkosComm::recv(comm, v, src).wait();
+    int errs = test_utils::count_errors(v);
+    ASSERT_EQ(errs, 0);
+  }
+}
+
+TYPED_TEST(SendRecv, Contig1D) {
+  auto v = test_utils::build_view<typename TestFixture::Scalar, 1>(test_utils::Contig{}, "v", 1013);
+  test_core_send_recv(v);
+}
+TYPED_TEST(SendRecv, NonContig1D) {
+  auto v = test_utils::build_view<typename TestFixture::Scalar, 1>(test_utils::NonContig{}, "v", 1013);
+  test_core_send_recv(v);
+}
+TYPED_TEST(SendRecv, Contig2D) {
+  auto v = test_utils::build_view<typename TestFixture::Scalar, 2>(test_utils::Contig{}, "v", 137, 17);
+  test_core_send_recv(v);
+}
+TYPED_TEST(SendRecv, NonContig2D) {
+  auto v = test_utils::build_view<typename TestFixture::Scalar, 2>(test_utils::NonContig{}, "v", 137, 17);
+  test_core_send_recv(v);
+}
+
+}  // namespace

unit_tests/test_sendrecv.cpp

@@ -0,0 +1,123 @@
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+// SPDX-FileCopyrightText: Copyright Contributors to the Kokkos project
+
+#pragma once
+
+#include <cstddef>
+#include <string>
+#include <utility>
+
+#include <Kokkos_Core.hpp>
+
+namespace test_utils {
+
+struct Contig {};
+struct NonContig {};
+
+namespace Impl {
+
+template <typename T, size_t N>
+struct AddPtrs {
+  using type = typename AddPtrs<T*, N - 1>::type;
+};
+template <typename T>
+struct AddPtrs<T, 0> {
+  using type = T;
+};
+template <typename T, size_t N>
+using Stars = typename AddPtrs<T, N>::type;
+
+template <size_t Rank>
+auto all_then_one() {
+  return [&]<size_t... Is>(std::index_sequence<Is...>) {
+    return std::make_tuple((void(Is), Kokkos::ALL)..., 1);
+  }(std::make_index_sequence<Rank>{});
+}
+template <typename ViewType, typename Tuple, size_t... Is>
+auto apply_subview(ViewType& v, Tuple&& t, std::index_sequence<Is...>) {
+  return Kokkos::subview(v, std::get<Is>(std::forward<Tuple>(t))...);
+}
+
+template <typename View, size_t... Is>
+struct InitFunctor {
+  View v;
+  std::array<int, sizeof...(Is)> exts;
+
+  KOKKOS_FUNCTION void operator()(decltype(static_cast<int>(std::declval<View>().extent(Is)))... idxs) const {
+    int val   = 0;
+    int ids[] = {int(idxs)...};
+    ((val = ids[Is] + exts[Is] * val), ...);
+    [&]<size_t... Js>(std::index_sequence<Js...>) { v(ids[Js]...) = val; }(std::index_sequence<Is...>{});
+  }
+};
+
+template <typename View, size_t... Is>
+struct CountFunctor {
+  using Scalar = typename View::value_type;
+  View v;
+  std::array<int, sizeof...(Is)> exts;
+
+  KOKKOS_FUNCTION void operator()(
+      decltype(static_cast<int>(std::declval<View>().extent(Is)))... idxs, int& lsum
+  ) const {
+    int val   = 0;
+    int ids[] = {int(idxs)...};
+    ((val = ids[Is] + exts[Is] * val), ...);
+    [&]<size_t... Js>(std::index_sequence<Js...>) {
+      lsum += v(ids[Js]...) != Scalar(val);
+    }(std::index_sequence<Is...>{});
+  }
+};
+
+}  // namespace Impl
+
+template <typename T, size_t R, typename... Extents>
+auto build_view(Contig, const std::string& name, Extents... exts) {
+  static_assert(sizeof...(exts) == R, "Number of extents must match Rank");
+  return Kokkos::View<Impl::Stars<T, R>>(name, exts...);
+}
+template <typename T, size_t R, typename... Extents>
+auto build_view(NonContig, const std::string& name, Extents... exts) {
+  static_assert(sizeof...(exts) == R, "Number of extents must match Rank");
+  Kokkos::View<Impl::Stars<T, R + 1>, Kokkos::LayoutRight> v(name, exts..., 2);
+  auto args = Impl::all_then_one<R>();
+  return Impl::apply_subview(v, args, std::make_index_sequence<R + 1>{});
+}
+
+template <typename Exec, typename View>
+auto init_view(const Exec& exec, const View& v) -> void {
+  constexpr size_t R = v.rank();
+  if constexpr (R == 1) {
+    Kokkos::parallel_for(Kokkos::RangePolicy<Exec>(exec, 0, v.extent(0)), KOKKOS_LAMBDA(const int i) { v(i) = i; });
+  } else {
+    [&]<size_t... Is>(std::index_sequence<Is...>) {
+      std::array<int, R> exts = {static_cast<int>(v.extent(Is))...};
+      Kokkos::parallel_for(
+          Kokkos::MDRangePolicy<Exec, Kokkos::Rank<R>>(exec, {(void(Is), 0)...}, {static_cast<int>(v.extent(Is))...}),
+          Impl::InitFunctor<View, Is...>{v, exts}
+      );
+    }(std::make_index_sequence<R>{});
+  }
+  exec.fence();
+}
+
+template <typename View>
+auto count_errors(const View& v) -> int {
+  constexpr size_t R = v.rank();
+  int errs           = 0;
+  using Scalar       = typename View::value_type;
+  if constexpr (R == 1) {
+    Kokkos::parallel_reduce(v.extent(0), KOKKOS_LAMBDA(const int i, int& lsum) { lsum += v(i) != Scalar(i); }, errs);
+  } else {
+    [&]<size_t... Is>(std::index_sequence<Is...>) {
+      std::array<int, R> exts = {static_cast<int>(v.extent(Is))...};
+      Kokkos::parallel_reduce(
+          Kokkos::MDRangePolicy<Kokkos::Rank<R>>({(void(Is), 0)...}, {static_cast<int>(v.extent(Is))...}),
+          Impl::CountFunctor<View, Is...>{v, exts}, errs
+      );
+    }(std::make_index_sequence<R>{});
+  }
+  return errs;
+}
+
+}  // namespace test_utils