From 3b5e286d12d02feab7d8df10f1a164fff6308267 Mon Sep 17 00:00:00 2001 From: Carl Pearson Date: Mon, 13 May 2024 15:23:47 -0600 Subject: [PATCH] mpi: basic contiguous Irecv --- src/KokkosComm_traits.hpp | 1 + src/impl/KokkosComm_irecv.hpp | 47 ++++++++++++++ unit_tests/CMakeLists.txt | 1 + unit_tests/test_isendirecv.cpp | 109 +++++++++++++++++++++++++++++++++ unit_tests/test_main.cpp | 1 - unit_tests/view_builder.hpp | 46 ++++++++++++++ 6 files changed, 204 insertions(+), 1 deletion(-) create mode 100644 src/impl/KokkosComm_irecv.hpp create mode 100644 unit_tests/test_isendirecv.cpp create mode 100644 unit_tests/view_builder.hpp diff --git a/src/KokkosComm_traits.hpp b/src/KokkosComm_traits.hpp index f34aa5b4..2968ff93 100644 --- a/src/KokkosComm_traits.hpp +++ b/src/KokkosComm_traits.hpp @@ -32,6 +32,7 @@ struct Traits { /*! \brief This can be specialized to do custom behavior for a particular view*/ template struct Traits { + // product of extents is span static bool is_contiguous(const View &v) { return v.span_is_contiguous(); } static auto data_handle(const View &v) { return v.data(); } diff --git a/src/impl/KokkosComm_irecv.hpp b/src/impl/KokkosComm_irecv.hpp new file mode 100644 index 00000000..cdffc34f --- /dev/null +++ b/src/impl/KokkosComm_irecv.hpp @@ -0,0 +1,47 @@ +//@HEADER +// ************************************************************************ +// +// Kokkos v. 4.0 +// Copyright (2022) National Technology & Engineering +// Solutions of Sandia, LLC (NTESS). +// +// Under the terms of Contract DE-NA0003525 with NTESS, +// the U.S. Government retains certain rights in this software. +// +// Part of Kokkos, under the Apache License v2.0 with LLVM Exceptions. +// See https://kokkos.org/LICENSE for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//@HEADER + +#pragma once + +#include + +#include + +#include "KokkosComm_pack_traits.hpp" +#include "KokkosComm_traits.hpp" + +// impl +#include "KokkosComm_include_mpi.hpp" + +namespace KokkosComm::Impl { + +// low-level API +template +void irecv(RecvView &rv, int src, int tag, MPI_Comm comm, MPI_Request &req) { + Kokkos::Tools::pushRegion("KokkosComm::Impl::irecv"); + + using KCT = KokkosComm::Traits; + + if (KCT::is_contiguous(rv)) { + using RecvScalar = typename RecvView::value_type; + MPI_Irecv(KCT::data_handle(rv), KCT::span(rv), mpi_type_v, src, tag, comm, &req); + } else { + throw std::runtime_error("Only contiguous irecv viewsupported"); + } + + Kokkos::Tools::popRegion(); +} +} // namespace KokkosComm::Impl diff --git a/unit_tests/CMakeLists.txt b/unit_tests/CMakeLists.txt index 029ff87d..09212db1 100644 --- a/unit_tests/CMakeLists.txt +++ b/unit_tests/CMakeLists.txt @@ -44,6 +44,7 @@ target_link_libraries(test-mpi MPI::MPI_CXX) # Kokkos Comm tests add_executable(test-main test_main.cpp test_gtest_mpi.cpp + test_isendirecv.cpp test_isendrecv.cpp test_reduce.cpp test_sendrecv.cpp diff --git a/unit_tests/test_isendirecv.cpp b/unit_tests/test_isendirecv.cpp new file mode 100644 index 00000000..1a351e85 --- /dev/null +++ b/unit_tests/test_isendirecv.cpp @@ -0,0 +1,109 @@ +//@HEADER +// ************************************************************************ +// +// Kokkos v. 4.0 +// Copyright (2022) National Technology & Engineering +// Solutions of Sandia, LLC (NTESS). +// +// Under the terms of Contract DE-NA0003525 with NTESS, +// the U.S. Government retains certain rights in this software. +// +// Part of Kokkos, under the Apache License v2.0 with LLVM Exceptions. +// See https://kokkos.org/LICENSE for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//@HEADER + +#include + +#include "KokkosComm.hpp" +#include "KokkosComm_irecv.hpp" + +#include "view_builder.hpp" + +namespace { + +template +class IsendIrecv : public testing::Test { + public: + using Scalar = T; +}; + +using ScalarTypes = + ::testing::Types, Kokkos::complex, int, unsigned, int64_t, size_t>; +TYPED_TEST_SUITE(IsendIrecv, ScalarTypes); + +template +void test_1d(const View1D &a) { + static_assert(View1D::rank == 1, ""); + using Scalar = typename View1D::non_const_value_type; + + int rank, size; + MPI_Comm_rank(MPI_COMM_WORLD, &rank); + MPI_Comm_size(MPI_COMM_WORLD, &size); + if (size < 2) { + GTEST_SKIP() << "Requires >= 2 ranks (" << size << " provided)"; + } + + if (0 == rank) { + int dst = 1; + Kokkos::parallel_for( + a.extent(0), KOKKOS_LAMBDA(const int i) { a(i) = i; }); + KokkosComm::Req req = KokkosComm::isend(Kokkos::DefaultExecutionSpace(), a, dst, 0, MPI_COMM_WORLD); + req.wait(); + } else if (1 == rank) { + int src = 0; + MPI_Request req; + KokkosComm::Impl::irecv(a, src, 0, MPI_COMM_WORLD, req); + MPI_Wait(&req, MPI_STATUS_IGNORE); + int errs; + Kokkos::parallel_reduce( + a.extent(0), KOKKOS_LAMBDA(const int &i, int &lsum) { lsum += a(i) != Scalar(i); }, errs); + ASSERT_EQ(errs, 0); + } +} + +template +void test_2d(const View2D &a) { + static_assert(View2D::rank == 2, ""); + using Scalar = typename View2D::non_const_value_type; + + int rank, size; + MPI_Comm_rank(MPI_COMM_WORLD, &rank); + MPI_Comm_size(MPI_COMM_WORLD, &size); + if (size < 2) { + GTEST_SKIP() << "Requires >= 2 ranks (" << size << " provided)"; + } + + using Policy = Kokkos::MDRangePolicy>; + Policy policy({0, 0}, {a.extent(0), a.extent(1)}); + + if (0 == rank) { + int dst = 1; + Kokkos::parallel_for( + policy, KOKKOS_LAMBDA(int i, int j) { a(i, j) = i * a.extent(0) + j; }); + KokkosComm::Req req = KokkosComm::isend(Kokkos::DefaultExecutionSpace(), a, dst, 0, MPI_COMM_WORLD); + req.wait(); + } else if (1 == rank) { + int src = 0; + MPI_Request req; + KokkosComm::Impl::irecv(a, src, 0, MPI_COMM_WORLD, req); + MPI_Wait(&req, MPI_STATUS_IGNORE); + int errs; + Kokkos::parallel_reduce( + policy, KOKKOS_LAMBDA(int i, int j, int &lsum) { lsum += a(i, j) != Scalar(i * a.extent(0) + j); }, errs); + ASSERT_EQ(errs, 0); + } +} + +TYPED_TEST(IsendIrecv, 1D_contig) { + auto a = ViewBuilder::view(contig{}, "a", 1013); + test_1d(a); +} + +TYPED_TEST(IsendIrecv, 2D_contig) { + auto a = ViewBuilder::view(contig{}, "a", 137, 17); + test_2d(a); +} + +} // namespace \ No newline at end of file diff --git a/unit_tests/test_main.cpp b/unit_tests/test_main.cpp index 27c03e5b..d5dea47a 100644 --- a/unit_tests/test_main.cpp +++ b/unit_tests/test_main.cpp @@ -88,7 +88,6 @@ int main(int argc, char *argv[]) { if (0 == rank) { std::cerr << argv[0] << " (KokkosComm " << KOKKOSCOMM_VERSION_MAJOR << "." << KOKKOSCOMM_VERSION_MINOR << "." << KOKKOSCOMM_VERSION_PATCH << ")\n"; - std::cerr << "size=" << size << "\n"; } Kokkos::initialize(); diff --git a/unit_tests/view_builder.hpp b/unit_tests/view_builder.hpp new file mode 100644 index 00000000..f4ba4fbf --- /dev/null +++ b/unit_tests/view_builder.hpp @@ -0,0 +1,46 @@ +//@HEADER +// ************************************************************************ +// +// Kokkos v. 4.0 +// Copyright (2022) National Technology & Engineering +// Solutions of Sandia, LLC (NTESS). +// +// Under the terms of Contract DE-NA0003525 with NTESS, +// the U.S. Government retains certain rights in this software. +// +// Part of Kokkos, under the Apache License v2.0 with LLVM Exceptions. +// See https://kokkos.org/LICENSE for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//@HEADER + +#pragma once + +#include + +struct contig {}; +struct noncontig {}; + +template +struct ViewBuilder; + +template +struct ViewBuilder { + static auto view(noncontig, const std::string &name, int e0) { + // this is C-style layout, i.e. v(0,0) is next to v(0,1) + Kokkos::View v(name, e0, 2); + return Kokkos::subview(v, Kokkos::ALL, 1); // take column 1 + } + + static auto view(contig, const std::string &name, int e0) { return Kokkos::View(name, e0); } +}; + +template +struct ViewBuilder { + static auto view(noncontig, const std::string &name, int e0, int e1) { + Kokkos::View v(name, e0, e1, 2); + return Kokkos::subview(v, Kokkos::ALL, Kokkos::ALL, 1); + } + + static auto view(contig, const std::string &name, int e0, int e1) { return Kokkos::View(name, e0, e1); } +};