Add all2all functionn and tests

Author: Yuuichi Asahi

distributed/src/KokkosFFT_Distributed_All2All.hpp

@@ -0,0 +1,60 @@
+// SPDX-FileCopyrightText: (C) The kokkos-fft development team, see COPYRIGHT.md file
+//
+// SPDX-License-Identifier: MIT OR Apache-2.0 WITH LLVM-exception
+
+#ifndef KOKKOSFFT_DISTRIBUTED_ALL2ALL_HPP
+#define KOKKOSFFT_DISTRIBUTED_ALL2ALL_HPP
+
+#include <mpi.h>
+#include <Kokkos_Core.hpp>
+#include <Kokkos_Profiling_ScopedRegion.hpp>
+#include "KokkosFFT_Distributed_MPI_Types.hpp"
+
+namespace KokkosFFT {
+namespace Distributed {
+namespace Impl {
+
+template <typename ExecutionSpace, typename ViewType>
+struct All2All {
+  static_assert(ViewType::rank() >= 2,
+                "All2All: View rank must be larger than or equal to 2");
+  using value_type = typename ViewType::non_const_value_type;
+  using LayoutType = typename ViewType::array_layout;
+
+  ExecutionSpace m_exec_space;
+  MPI_Comm m_comm;
+  MPI_Datatype m_mpi_data_type;
+
+  All2All(const ViewType& send, const ViewType& recv,
+          const MPI_Comm& comm            = MPI_COMM_WORLD,
+          const ExecutionSpace exec_space = ExecutionSpace())
+      : m_exec_space(exec_space),
+        m_comm(comm),
+        m_mpi_data_type(MPIDataType<value_type>::type()) {
+    // Compute the outermost dimension size
+    int send_count = 0;
+    if (std::is_same_v<LayoutType, Kokkos::LayoutLeft>) {
+      send_count = send.size() / send.extent(ViewType::rank() - 1);
+    } else {
+      send_count = send.size() / send.extent(0);
+    }
+
+    ::MPI_Alltoall(send.data(), send_count, m_mpi_data_type, recv.data(),
+                   send_count, m_mpi_data_type, m_comm);
+  }
+};
+
+template <typename ExecutionSpace, typename ViewType>
+void all2all(const ExecutionSpace& exec_space, const ViewType& send,
+             const ViewType& recv, const MPI_Comm& comm) {
+  static_assert(ViewType::rank() >= 2,
+                "all2all: View rank must be larger than or equal to 2");
+  Kokkos::Profiling::ScopedRegion region("KokkosFFT::Distributed::all2all");
+  All2All(send, recv, comm, exec_space);
+}
+
+}  // namespace Impl
+}  // namespace Distributed
+}  // namespace KokkosFFT
+
+#endif

distributed/unit_test/CMakeLists.txt

@@ -4,6 +4,7 @@
 
 set(SOURCES
         Test_Main.cpp
+        Test_All2All.cpp
         Test_MPI_Types.cpp
     )
 

distributed/unit_test/Test_All2All.cpp

@@ -0,0 +1,109 @@
+// SPDX-FileCopyrightText: (C) The kokkos-fft development team, see COPYRIGHT.md file
+//
+// SPDX-License-Identifier: MIT OR Apache-2.0 WITH LLVM-exception
+
+#include <gtest/gtest.h>
+#include <Kokkos_Core.hpp>
+#include "KokkosFFT_Distributed_All2All.hpp"
+
+namespace {
+using execution_space = Kokkos::DefaultExecutionSpace;
+using test_types      = ::testing::Types<std::pair<float, Kokkos::LayoutLeft>,
+                                    std::pair<float, Kokkos::LayoutRight>,
+                                    std::pair<double, Kokkos::LayoutLeft>,
+                                    std::pair<double, Kokkos::LayoutRight>>;
+
+// Basically the same fixtures, used for labeling tests
+template <typename T>
+struct TestAll2All : public ::testing::Test {
+  using float_type  = typename T::first_type;
+  using layout_type = typename T::second_type;
+
+  int m_rank   = 0;
+  int m_nprocs = 1;
+
+  virtual void SetUp() {
+    ::MPI_Comm_rank(MPI_COMM_WORLD, &m_rank);
+    ::MPI_Comm_size(MPI_COMM_WORLD, &m_nprocs);
+  }
+};
+
+template <typename T, typename LayoutType>
+void test_all2all_view2D(int rank, int nprocs) {
+  using View3DType = Kokkos::View<T***, LayoutType, execution_space>;
+
+  const std::size_t n0 = 16, n1 = 15;
+  const std::size_t n0_local = ((n0 - 1) / nprocs) + 1;
+  const std::size_t n1_local = ((n1 - 1) / nprocs) + 1;
+
+  int n0_buffer = 0, n1_buffer = 0, n2_buffer = 0;
+  if constexpr (std::is_same_v<LayoutType, Kokkos::LayoutLeft>) {
+    n0_buffer = n0_local;
+    n1_buffer = n1_local;
+    n2_buffer = nprocs;
+  } else {
+    n0_buffer = nprocs;
+    n1_buffer = n0_local;
+    n2_buffer = n1_local;
+  }
+
+  View3DType send("send", n0_buffer, n1_buffer, n2_buffer),
+      recv("recv", n0_buffer, n1_buffer, n2_buffer),
+      ref("ref", n0_buffer, n1_buffer, n2_buffer);
+
+  auto h_send = Kokkos::create_mirror_view(send);
+  auto h_ref  = Kokkos::create_mirror_view(ref);
+
+  for (std::size_t i2 = 0; i2 < send.extent(2); i2++) {
+    for (std::size_t i1 = 0; i1 < send.extent(1); i1++) {
+      for (std::size_t i0 = 0; i0 < send.extent(0); i0++) {
+        if constexpr (std::is_same_v<LayoutType, Kokkos::LayoutLeft>) {
+          T value =
+              static_cast<T>(rank * send.size() + i0 + i1 * send.extent(0) +
+                             i2 * send.extent(0) * send.extent(1));
+          T value_T =
+              static_cast<T>(i2 * send.size() + i0 + i1 * send.extent(0) +
+                             rank * send.extent(0) * send.extent(1));
+          h_send(i0, i1, i2) = value;
+          h_ref(i0, i1, i2)  = value_T;
+        } else {
+          T value =
+              static_cast<T>(rank * send.size() + i2 + i1 * send.extent(2) +
+                             i0 * send.extent(2) * send.extent(1));
+          T value_T =
+              static_cast<T>(i0 * send.size() + i2 + i1 * send.extent(2) +
+                             rank * send.extent(2) * send.extent(1));
+          h_send(i0, i1, i2) = value;
+          h_ref(i0, i1, i2)  = value_T;
+        }
+      }
+    }
+  }
+
+  Kokkos::deep_copy(send, h_send);
+  Kokkos::deep_copy(ref, h_ref);
+
+  execution_space exec;
+  KokkosFFT::Distributed::Impl::all2all(exec, send, recv, MPI_COMM_WORLD);
+  auto h_recv = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), recv);
+
+  T epsilon = std::numeric_limits<T>::epsilon() * 100;
+  for (std::size_t i2 = 0; i2 < send.extent(2); i2++) {
+    for (std::size_t i1 = 0; i1 < send.extent(1); i1++) {
+      for (std::size_t i0 = 0; i0 < send.extent(0); i0++) {
+        auto diff = Kokkos::abs(h_recv(i0, i1, i2) - h_ref(i0, i1, i2));
+        EXPECT_LE(diff, epsilon);
+      }
+    }
+  }
+}
+
+}  // namespace
+
+TYPED_TEST_SUITE(TestAll2All, test_types);
+
+TYPED_TEST(TestAll2All, View2D) {
+  using float_type  = typename TestFixture::float_type;
+  using layout_type = typename TestFixture::layout_type;
+  test_all2all_view2D<float_type, layout_type>(this->m_rank, this->m_nprocs);
+}