From af9d68a6682997926cf18d4e98de945df5cb4f28 Mon Sep 17 00:00:00 2001 From: Carl Pearson Date: Mon, 22 Apr 2024 15:06:21 -0600 Subject: [PATCH 1/4] mpi: 1D contiguous MPI_Allgather --- docs/api/core.rst | 16 ++++++ src/KokkosComm_collective.hpp | 6 +++ src/impl/KokkosComm_allgather.hpp | 51 +++++++++++++++++++ unit_tests/CMakeLists.txt | 1 + unit_tests/test_allgather.cpp | 82 +++++++++++++++++++++++++++++++ 5 files changed, 156 insertions(+) create mode 100644 src/impl/KokkosComm_allgather.hpp create mode 100644 unit_tests/test_allgather.cpp diff --git a/docs/api/core.rst b/docs/api/core.rst index fb3a895a..1b9e2b34 100644 --- a/docs/api/core.rst +++ b/docs/api/core.rst @@ -97,6 +97,22 @@ Collective :tparam RecvView: A Kokkos::View to recv :tparam ExecSpace: A Kokkos execution space to operate in + +.. cpp:function:: template \ + void KokkosComm::allgather(const ExecSpace &space, const SendView &sv, const RecvView &rv, MPI_Comm comm) + + MPI_Allgather wrapper + + :param space: The execution space to operate in + :param sv: The data to send + :param rv: The view to receive into + :param comm: the MPI communicator + :tparam SendView: A Kokkos::View to send. Contiguous and rank less than 2. + :tparam RecvView: A Kokkos::View to recv. Contiguous and rank 1. + :tparam ExecSpace: A Kokkos execution space to operate in + + If sv_ is a rank-0 view, the value from the jth rank will be placed in index j of rv_. + Related Types ------------- diff --git a/src/KokkosComm_collective.hpp b/src/KokkosComm_collective.hpp index 92ccf50b..e9af5ee3 100644 --- a/src/KokkosComm_collective.hpp +++ b/src/KokkosComm_collective.hpp @@ -21,6 +21,7 @@ #include "KokkosComm_concepts.hpp" #include "KokkosComm_alltoall.hpp" #include "KokkosComm_reduce.hpp" +#include "KokkosComm_allgather.hpp" namespace KokkosComm { @@ -29,4 +30,9 @@ void reduce(const ExecSpace &space, const SendView &sv, const RecvView &rv, MPI_ return Impl::reduce(space, sv, rv, op, root, comm); } +template +void allgather(const ExecSpace &space, const SendView &sv, const RecvView &rv, MPI_Comm comm) { + return Impl::allgather(space, sv, rv, comm); +} + } // namespace KokkosComm diff --git a/src/impl/KokkosComm_allgather.hpp b/src/impl/KokkosComm_allgather.hpp new file mode 100644 index 00000000..f09b6251 --- /dev/null +++ b/src/impl/KokkosComm_allgather.hpp @@ -0,0 +1,51 @@ +//@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 "KokkosComm_pack_traits.hpp" +#include "KokkosComm_traits.hpp" + +// impl +#include "KokkosComm_include_mpi.hpp" +#include "KokkosComm_types.hpp" + +namespace KokkosComm::Impl { +template +void allgather(const ExecSpace &space, const SendView &sv, const RecvView &rv, MPI_Comm comm) { + Kokkos::Tools::pushRegion("KokkosComm::Impl::allgather"); + + using ST = KokkosComm::Traits; + using RT = KokkosComm::Traits; + using SendScalar = typename SendView::value_type; + using RecvScalar = typename RecvView::value_type; + + static_assert(ST::rank() <= 1, "allgather for SendView::rank > 1 not supported"); + static_assert(RT::rank() <= 1, "allgather for RecvView::rank > 1 not supported"); + + if (KokkosComm::PackTraits::needs_pack(sv) || KokkosComm::PackTraits::needs_pack(rv)) { + throw std::runtime_error("allgather for non-contiguous views not implemented"); + } else { + const int count = ST::span(sv); // all ranks send/recv same count + MPI_Allgather(ST::data_handle(sv), count, mpi_type_v, RT::data_handle(rv), count, + mpi_type_v, comm); + } + + Kokkos::Tools::popRegion(); +} +} // namespace KokkosComm::Impl diff --git a/unit_tests/CMakeLists.txt b/unit_tests/CMakeLists.txt index de8c2241..8cb9d1d2 100644 --- a/unit_tests/CMakeLists.txt +++ b/unit_tests/CMakeLists.txt @@ -50,6 +50,7 @@ add_executable(test-main test_main.cpp test_barrier.cpp test_alltoall.cpp test_reduce.cpp + test_allgather.cpp ) target_link_libraries(test-main KokkosComm::KokkosComm gtest) if(KOKKOSCOMM_ENABLE_TESTS) diff --git a/unit_tests/test_allgather.cpp b/unit_tests/test_allgather.cpp new file mode 100644 index 00000000..926134af --- /dev/null +++ b/unit_tests/test_allgather.cpp @@ -0,0 +1,82 @@ +//@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" + +template +class Allgather : public testing::Test { + public: + using Scalar = T; +}; + +using ScalarTypes = ::testing::Types, Kokkos::complex>; +TYPED_TEST_SUITE(Allgather, ScalarTypes); + +TYPED_TEST(Allgather, 0D) { + using TestScalar = typename TestFixture::Scalar; + + int rank, size; + MPI_Comm_rank(MPI_COMM_WORLD, &rank); + MPI_Comm_size(MPI_COMM_WORLD, &size); + + const int nContrib = 10; + + Kokkos::View sv("sv"); + Kokkos::View rv("rv", size); + + // fill send buffer + Kokkos::parallel_for( + sv.extent(0), KOKKOS_LAMBDA(const int i) { sv() = rank; }); + + KokkosComm::allgather(Kokkos::DefaultExecutionSpace(), sv, rv, MPI_COMM_WORLD); + + int errs; + Kokkos::parallel_reduce( + rv.extent(0), KOKKOS_LAMBDA(const int &src, int &lsum) { lsum += rv(src) != src; }, errs); + EXPECT_EQ(errs, 0); +} + +TYPED_TEST(Allgather, 1D_contig) { + using TestScalar = typename TestFixture::Scalar; + + int rank, size; + MPI_Comm_rank(MPI_COMM_WORLD, &rank); + MPI_Comm_size(MPI_COMM_WORLD, &size); + + const int nContrib = 10; + + Kokkos::View sv("sv", nContrib); + Kokkos::View rv("rv", size * nContrib); + + // fill send buffer + Kokkos::parallel_for( + sv.extent(0), KOKKOS_LAMBDA(const int i) { sv(i) = rank + i; }); + + KokkosComm::allgather(Kokkos::DefaultExecutionSpace(), sv, rv, MPI_COMM_WORLD); + + int errs; + Kokkos::parallel_reduce( + rv.extent(0), + KOKKOS_LAMBDA(const int &i, int &lsum) { + const int src = i / nContrib; + const int j = i % nContrib; + lsum += rv(i) != src + j; + }, + errs); + EXPECT_EQ(errs, 0); +} From bd4f7ddd31ac6c3068f729bccfc8edd61a37cc64 Mon Sep 17 00:00:00 2001 From: Carl Pearson Date: Wed, 15 May 2024 12:55:10 -0600 Subject: [PATCH 2/4] mpi allgather docs: fix param reference --- docs/api/core.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/api/core.rst b/docs/api/core.rst index 1b9e2b34..d229dc5d 100644 --- a/docs/api/core.rst +++ b/docs/api/core.rst @@ -111,7 +111,7 @@ Collective :tparam RecvView: A Kokkos::View to recv. Contiguous and rank 1. :tparam ExecSpace: A Kokkos execution space to operate in - If sv_ is a rank-0 view, the value from the jth rank will be placed in index j of rv_. + If ``sv`` is a rank-0 view, the value from the jth rank will be placed in index j of ``rv``. Related Types ------------- From 200ebaea119aac51e22f1f66467fd35c0e3b372c Mon Sep 17 00:00:00 2001 From: Carl Pearson Date: Fri, 17 May 2024 09:32:31 -0600 Subject: [PATCH 3/4] mpi: allgather: low-level API, space handling --- src/impl/KokkosComm_allgather.hpp | 31 +++++++++++++++++++++++++------ 1 file changed, 25 insertions(+), 6 deletions(-) diff --git a/src/impl/KokkosComm_allgather.hpp b/src/impl/KokkosComm_allgather.hpp index f09b6251..75e41f61 100644 --- a/src/impl/KokkosComm_allgather.hpp +++ b/src/impl/KokkosComm_allgather.hpp @@ -26,8 +26,9 @@ #include "KokkosComm_types.hpp" namespace KokkosComm::Impl { -template -void allgather(const ExecSpace &space, const SendView &sv, const RecvView &rv, MPI_Comm comm) { + +template +void allgather(const SendView &sv, const RecvView &rv, MPI_Comm comm) { Kokkos::Tools::pushRegion("KokkosComm::Impl::allgather"); using ST = KokkosComm::Traits; @@ -38,12 +39,30 @@ void allgather(const ExecSpace &space, const SendView &sv, const RecvView &rv, M static_assert(ST::rank() <= 1, "allgather for SendView::rank > 1 not supported"); static_assert(RT::rank() <= 1, "allgather for RecvView::rank > 1 not supported"); - if (KokkosComm::PackTraits::needs_pack(sv) || KokkosComm::PackTraits::needs_pack(rv)) { + if (!ST::is_contiguous(sv)) { + throw std::runtime_error("low-level allgather requires contiguous send view"); + } + if (!RT::is_contiguous(rv)) { + throw std::runtime_error("low-level allgather requires contiguous recv view"); + } + const int count = ST::span(sv); // all ranks send/recv same count + MPI_Allgather(ST::data_handle(sv), count, mpi_type_v, RT::data_handle(rv), count, mpi_type_v, + comm); + + Kokkos::Tools::popRegion(); +} + +template +void allgather(const ExecSpace &space, const SendView &sv, const RecvView &rv, MPI_Comm comm) { + Kokkos::Tools::pushRegion("KokkosComm::Impl::allgather"); + using ST = KokkosComm::Traits; + using RT = KokkosComm::Traits; + + if (ST::needs_pack(sv) || RT::needs_pack(rv)) { throw std::runtime_error("allgather for non-contiguous views not implemented"); } else { - const int count = ST::span(sv); // all ranks send/recv same count - MPI_Allgather(ST::data_handle(sv), count, mpi_type_v, RT::data_handle(rv), count, - mpi_type_v, comm); + space.fence(); // work in space may have been used to produce send view data + allgather(sv, rv, comm); } Kokkos::Tools::popRegion(); From cd8f4ff217e25c3395baa73683f4f2817b3ab377 Mon Sep 17 00:00:00 2001 From: Carl Pearson Date: Fri, 17 May 2024 09:36:56 -0600 Subject: [PATCH 4/4] mpi: allgather: fix trait alias --- src/impl/KokkosComm_allgather.hpp | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/impl/KokkosComm_allgather.hpp b/src/impl/KokkosComm_allgather.hpp index 75e41f61..ae9c4265 100644 --- a/src/impl/KokkosComm_allgather.hpp +++ b/src/impl/KokkosComm_allgather.hpp @@ -55,10 +55,10 @@ void allgather(const SendView &sv, const RecvView &rv, MPI_Comm comm) { template void allgather(const ExecSpace &space, const SendView &sv, const RecvView &rv, MPI_Comm comm) { Kokkos::Tools::pushRegion("KokkosComm::Impl::allgather"); - using ST = KokkosComm::Traits; - using RT = KokkosComm::Traits; + using SPT = KokkosComm::PackTraits; + using RPT = KokkosComm::PackTraits; - if (ST::needs_pack(sv) || RT::needs_pack(rv)) { + if (SPT::needs_pack(sv) || RPT::needs_pack(rv)) { throw std::runtime_error("allgather for non-contiguous views not implemented"); } else { space.fence(); // work in space may have been used to produce send view data