diff --git a/core/CMakeLists.txt b/core/CMakeLists.txt index 5725813bbac..c4bea01dbd0 100644 --- a/core/CMakeLists.txt +++ b/core/CMakeLists.txt @@ -131,6 +131,7 @@ if(GINKGO_BUILD_MPI) PRIVATE mpi/exception.cpp distributed/collective_communicator.cpp + distributed/dense_communicator.cpp distributed/matrix.cpp distributed/partition_helpers.cpp distributed/vector.cpp diff --git a/core/distributed/dense_communicator.cpp b/core/distributed/dense_communicator.cpp new file mode 100644 index 00000000000..d3c7241f8aa --- /dev/null +++ b/core/distributed/dense_communicator.cpp @@ -0,0 +1,130 @@ +// SPDX-FileCopyrightText: 2024 The Ginkgo authors +// +// SPDX-License-Identifier: BSD-3-Clause + +#include "ginkgo/core/distributed/dense_communicator.hpp" + +namespace gko { +namespace experimental { +namespace mpi { + + +DenseCommunicator::DenseCommunicator(communicator base) + : CollectiveCommunicator(base), + comm_(base), + recv_sizes_(comm_.size()), + recv_offsets_(comm_.size() + 1), + send_sizes_(comm_.size()), + send_offsets_(comm_.size() + 1) +{} + + +template +DenseCommunicator::DenseCommunicator( + communicator base, + const distributed::index_map& imap) + : DenseCommunicator(base) +{ + auto exec = imap.get_executor(); + if (!exec) { + return; + } + auto host_exec = exec->get_master(); + + auto recv_target_ids_arr = + make_temporary_clone(host_exec, &imap.get_remote_target_ids()); + auto remote_idx_offsets_arr = make_temporary_clone( + host_exec, &imap.get_remote_global_idxs().get_offsets()); + for (size_type seg_id = 0; + seg_id < imap.get_remote_global_idxs().get_segment_count(); ++seg_id) { + recv_sizes_[recv_target_ids_arr->get_const_data()[seg_id]] = + remote_idx_offsets_arr->get_const_data()[seg_id + 1] - + remote_idx_offsets_arr->get_const_data()[seg_id]; + } + + comm_.all_to_all(host_exec, recv_sizes_.data(), 1, send_sizes_.data(), 1); + + std::partial_sum(send_sizes_.begin(), send_sizes_.end(), + send_offsets_.begin() + 1); + std::partial_sum(recv_sizes_.begin(), recv_sizes_.end(), + recv_offsets_.begin() + 1); +} + +#define GKO_DECLARE_DENSE_CONSTRUCTOR(LocalIndexType, GlobalIndexType) \ + DenseCommunicator::DenseCommunicator( \ + communicator base, \ + const distributed::index_map& imap) + +GKO_INSTANTIATE_FOR_EACH_LOCAL_GLOBAL_INDEX_TYPE(GKO_DECLARE_DENSE_CONSTRUCTOR); + +#undef GKO_DECLARE_DENSE_CONSTRUCTOR + + +DenseCommunicator::DenseCommunicator( + communicator base, const std::vector& recv_sizes, + const std::vector& recv_offsets, + const std::vector& send_sizes, + const std::vector& send_offsets) + : CollectiveCommunicator(base), + comm_(base), + recv_sizes_(recv_sizes), + recv_offsets_(recv_offsets), + send_sizes_(send_sizes), + send_offsets_(send_offsets) +{} + + +request DenseCommunicator::i_all_to_all_v(std::shared_ptr exec, + const void* send_buffer, + MPI_Datatype send_type, + void* recv_buffer, + MPI_Datatype recv_type) const +{ +#ifdef GINKGO_FORCE_SPMV_BLOCKING_COMM + comm_.all_to_all_v(exec, send_buffer, send_sizes_.data(), + send_offsets_.data(), send_type, recv_buffer, + recv_sizes_.data(), recv_offsets_.data(), recv_type); + return {}; +#else + return comm_.i_all_to_all_v( + exec, send_buffer, send_sizes_.data(), send_offsets_.data(), send_type, + recv_buffer, recv_sizes_.data(), recv_offsets_.data(), recv_type); +#endif +} + + +std::unique_ptr +DenseCommunicator::create_with_same_type( + communicator base, const distributed::index_map_variant& imap) const +{ + return std::visit( + [base](const auto& imap) { + return std::make_unique(base, imap); + }, + imap); +} + + +std::unique_ptr DenseCommunicator::create_inverse() + const +{ + return std::make_unique( + comm_, send_sizes_, send_offsets_, recv_sizes_, recv_offsets_); +} + + +comm_index_type DenseCommunicator::get_recv_size() const +{ + return recv_offsets_.back(); +} + + +comm_index_type DenseCommunicator::get_send_size() const +{ + return send_offsets_.back(); +} + + +} // namespace mpi +} // namespace experimental +} // namespace gko diff --git a/core/test/mpi/distributed/CMakeLists.txt b/core/test/mpi/distributed/CMakeLists.txt index 8cff893408a..433b3dd1fc6 100644 --- a/core/test/mpi/distributed/CMakeLists.txt +++ b/core/test/mpi/distributed/CMakeLists.txt @@ -1,5 +1,6 @@ ginkgo_create_test(helpers MPI_SIZE 1) ginkgo_create_test(matrix MPI_SIZE 1) +ginkgo_create_test(dense_communicator MPI_SIZE 6) add_subdirectory(preconditioner) add_subdirectory(solver) diff --git a/core/test/mpi/distributed/dense_communicator.cpp b/core/test/mpi/distributed/dense_communicator.cpp new file mode 100644 index 00000000000..1a868a047b9 --- /dev/null +++ b/core/test/mpi/distributed/dense_communicator.cpp @@ -0,0 +1,232 @@ +// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors +// +// SPDX-License-Identifier: BSD-3-Clause + +#include + +#include + +#include "core/test/utils/assertions.hpp" + +using gko::experimental::mpi::comm_index_type; + +class DenseCommunicator : public ::testing::Test { +protected: + using part_type = gko::experimental::distributed::Partition; + using map_type = gko::experimental::distributed::index_map; + + void SetUp() override { ASSERT_EQ(comm.size(), 6); } + + std::shared_ptr ref = gko::ReferenceExecutor::create(); + gko::experimental::mpi::communicator comm = MPI_COMM_WORLD; + int rank = comm.rank(); +}; + + +TEST_F(DenseCommunicator, CanDefaultConstruct) +{ + gko::experimental::mpi::DenseCommunicator nhcomm{comm}; + + ASSERT_EQ(nhcomm.get_base_communicator(), comm); + ASSERT_EQ(nhcomm.get_send_size(), 0); + ASSERT_EQ(nhcomm.get_recv_size(), 0); +} + + +TEST_F(DenseCommunicator, CanConstructFromIndexMap) +{ + auto part = gko::share(part_type::build_from_global_size_uniform( + ref, comm.size(), comm.size() * 3)); + gko::array recv_connections[] = {{ref, {3, 5, 10, 11}}, + {ref, {0, 1, 7, 12, 13}}, + {ref, {3, 4, 17}}, + {ref, {1, 2, 12, 14}}, + {ref, {4, 5, 9, 10, 16, 15}}, + {ref, {8, 12, 13, 14}}}; + auto imap = map_type{ref, part, comm.rank(), recv_connections[rank]}; + + gko::experimental::mpi::DenseCommunicator spcomm{comm, imap}; + + std::array send_sizes = {4, 6, 2, 4, 7, 3}; + ASSERT_EQ(spcomm.get_recv_size(), recv_connections[rank].get_size()); + ASSERT_EQ(spcomm.get_send_size(), send_sizes[rank]); +} + + +TEST_F(DenseCommunicator, CanConstructFromEnvelopData) +{ + // clang-format off + std::vector recv_sizes[] = { + {0, 2, 2, + 0, 0, 0}, + {2, 0, 1, + 2, 0, 0}, + {0, 2, 0, + 0, 0, 1}, + {2, 0, 0, + 0, 2, 0}, + {0, 2, 0, + 2, 0, 2}, + {0, 0, 1, + 0, 3, 0}}; + std::vector send_sizes[] = { + {0, 2, 0, + 2, 0, 0}, + {2, 0, 2, + 0, 2, 0}, + {0, 1, 0, + 0, 0, 1}, + {2, 0, 0, + 0, 2, 0}, + {0, 2, 0, + 2, 0, 3}, + {0, 0, 1, + 0, 2, 0}}; + // clang-format on + std::vector recv_offsets(recv_sizes[rank].size() + 1); + std::vector send_offsets(send_sizes[rank].size() + 1); + std::partial_sum(recv_sizes[rank].begin(), recv_sizes[rank].end(), + recv_offsets.begin() + 1); + std::partial_sum(send_sizes[rank].begin(), send_sizes[rank].end(), + send_offsets.begin() + 1); + + gko::experimental::mpi::DenseCommunicator spcomm{ + comm, recv_sizes[rank], recv_offsets, send_sizes[rank], send_offsets, + }; + + ASSERT_EQ(spcomm.get_recv_size(), recv_offsets.back()); + ASSERT_EQ(spcomm.get_send_size(), send_offsets.back()); +} + + +TEST_F(DenseCommunicator, CanConstructFromEmptyIndexMap) +{ + auto imap = map_type{ref}; + + gko::experimental::mpi::DenseCommunicator spcomm{comm, imap}; + + ASSERT_EQ(spcomm.get_recv_size(), 0); + ASSERT_EQ(spcomm.get_send_size(), 0); +} + + +TEST_F(DenseCommunicator, CanConstructFromIndexMapWithoutConnection) +{ + auto part = gko::share(part_type::build_from_global_size_uniform( + ref, comm.size(), comm.size() * 3)); + auto imap = map_type{ref, part, comm.rank(), {ref, 0}}; + + gko::experimental::mpi::DenseCommunicator spcomm{comm, imap}; + + ASSERT_EQ(spcomm.get_recv_size(), 0); + ASSERT_EQ(spcomm.get_send_size(), 0); +} + + +TEST_F(DenseCommunicator, CanConstructFromEmptyEnvelopData) +{ + std::vector recv_sizes; + std::vector send_sizes; + std::vector recv_offsets{0}; + std::vector send_offsets{0}; + + gko::experimental::mpi::DenseCommunicator spcomm{ + comm, recv_sizes, recv_offsets, send_sizes, send_offsets, + }; + + ASSERT_EQ(spcomm.get_recv_size(), 0); + ASSERT_EQ(spcomm.get_send_size(), 0); +} + + +TEST_F(DenseCommunicator, CanCommunicateIalltoall) +{ + auto part = gko::share(part_type::build_from_global_size_uniform( + ref, comm.size(), comm.size() * 3)); + gko::array recv_connections[] = {{ref, {3, 5, 10, 11}}, + {ref, {0, 1, 7, 12, 13}}, + {ref, {3, 4, 17}}, + {ref, {1, 2, 12, 14}}, + {ref, {4, 5, 9, 10, 16, 15}}, + {ref, {8, 12, 13, 14}}}; + auto imap = map_type{ref, part, comm.rank(), recv_connections[rank]}; + gko::experimental::mpi::DenseCommunicator spcomm{comm, imap}; + gko::array recv_buffer{ref, recv_connections[rank].get_size()}; + gko::array send_buffers[] = {{ref, {0, 1, 1, 2}}, + {ref, {3, 5, 3, 4, 4, 5}}, + {ref, {7, 8}}, + {ref, {10, 11, 9, 10}}, + {ref, {12, 13, 12, 14, 12, 13, 14}}, + {ref, {17, 16, 15}}}; + + auto req = spcomm.i_all_to_all_v(ref, send_buffers[rank].get_const_data(), + recv_buffer.get_data()); + req.wait(); + + GKO_ASSERT_ARRAY_EQ(recv_buffer, recv_connections[rank]); +} + + +TEST_F(DenseCommunicator, CanCommunicateIalltoallWhenEmpty) +{ + gko::experimental::mpi::DenseCommunicator spcomm{comm}; + + auto req = spcomm.i_all_to_all_v(ref, static_cast(nullptr), + static_cast(nullptr)); + req.wait(); +} + + +TEST_F(DenseCommunicator, CanCreateInverse) +{ + auto part = gko::share(part_type::build_from_global_size_uniform( + ref, comm.size(), comm.size() * 3)); + gko::array recv_connections[] = {{ref, {3, 5, 10, 11}}, + {ref, {0, 1, 7, 12, 13}}, + {ref, {3, 4, 17}}, + {ref, {1, 2, 12, 14}}, + {ref, {4, 5, 9, 10, 16, 15}}, + {ref, {8, 12, 13, 14}}}; + auto imap = map_type{ref, part, comm.rank(), recv_connections[rank]}; + gko::experimental::mpi::DenseCommunicator spcomm{comm, imap}; + + auto inverse = spcomm.create_inverse(); + + ASSERT_EQ(inverse->get_recv_size(), spcomm.get_send_size()); + ASSERT_EQ(inverse->get_send_size(), spcomm.get_recv_size()); +} + + +TEST_F(DenseCommunicator, CanCommunicateRoundTrip) +{ + auto part = gko::share(part_type::build_from_global_size_uniform( + ref, comm.size(), comm.size() * 3)); + gko::array recv_connections[] = {{ref, {3, 5, 10, 11}}, + {ref, {0, 1, 7, 12, 13}}, + {ref, {3, 4, 17}}, + {ref, {1, 2, 12, 14}}, + {ref, {4, 5, 9, 10, 16, 15}}, + {ref, {8, 12, 13, 14}}}; + auto imap = map_type{ref, part, comm.rank(), recv_connections[rank]}; + gko::experimental::mpi::DenseCommunicator spcomm{comm, imap}; + auto inverse = spcomm.create_inverse(); + gko::array send_buffers[] = {{ref, {1, 2, 3, 4}}, + {ref, {5, 6, 7, 8, 9, 10}}, + {ref, {11, 12}}, + {ref, {13, 14, 15, 16}}, + {ref, {17, 18, 19, 20, 21, 22, 23}}, + {ref, {24, 25, 26}}}; + gko::array recv_buffer{ref, recv_connections[rank].get_size()}; + gko::array round_trip{ref, send_buffers[rank].get_size()}; + + spcomm + .i_all_to_all_v(ref, send_buffers[rank].get_const_data(), + recv_buffer.get_data()) + .wait(); + inverse + ->i_all_to_all_v(ref, recv_buffer.get_const_data(), + round_trip.get_data()) + .wait(); + + GKO_ASSERT_ARRAY_EQ(send_buffers[rank], round_trip); +} diff --git a/include/ginkgo/core/distributed/dense_communicator.hpp b/include/ginkgo/core/distributed/dense_communicator.hpp new file mode 100644 index 00000000000..20a8ad99114 --- /dev/null +++ b/include/ginkgo/core/distributed/dense_communicator.hpp @@ -0,0 +1,126 @@ +// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors +// +// SPDX-License-Identifier: BSD-3-Clause + +#ifndef GKO_PUBLIC_CORE_DISTRIBUTED_DENSE_COMMUNICATOR_HPP_ +#define GKO_PUBLIC_CORE_DISTRIBUTED_DENSE_COMMUNICATOR_HPP_ + + +#include + + +#if GINKGO_BUILD_MPI + +#include +#include +#include + + +namespace gko { +namespace experimental { +namespace mpi { + + +/** + * A collective_communicator that uses a dense communication. + * + * The neighborhood communicator is defined by a list of neighbors this + * rank sends data to and a list of neighbors this rank receives data from. + * No communication with any ranks that is not in one of those lists will + * take place. + */ +class DenseCommunicator final : public CollectiveCommunicator { +public: + using CollectiveCommunicator::i_all_to_all_v; + + /** + * Default constructor with empty communication pattern + * @param base the base communicator + */ + explicit DenseCommunicator(communicator base); + + /** + * Create a neighborhood_communicator from an index map. + * + * The receive neighbors are defined by the remote indices and their + * owning ranks of the index map. The send neighbors are deduced + * from that through collective communication. + * + * @tparam LocalIndexType the local index type of the map + * @tparam GlobalIndexType the global index type of the map + * @param base the base communicator + * @param imap the index map that defines the communication pattern + */ + template + DenseCommunicator( + communicator base, + const distributed::index_map& imap); + + /** + * Create a neighborhood_communicator by explicitly defining the + * neighborhood lists and sizes/offsets. + * + * @param base the base communicator + * @param sources the ranks to receive from + * @param recv_sizes the number of elements to recv for each source + * @param recv_offsets the offset for each source + * @param destinations the ranks to send to + * @param send_sizes the number of elements to send for each destination + * @param send_offsets the offset for each destination + */ + DenseCommunicator(communicator base, + const std::vector& recv_sizes, + const std::vector& recv_offsets, + const std::vector& send_sizes, + const std::vector& send_offsets); + + /** + * @copydoc CollectiveCommunicator::i_all_to_all_v + * + * This implementation uses the neighborhood communication + * MPI_Ineighbor_alltoallv. See MPI documentation for more details. + */ + request i_all_to_all_v(std::shared_ptr exec, + const void* send_buffer, MPI_Datatype send_type, + void* recv_buffer, + MPI_Datatype recv_type) const override; + + [[nodiscard]] std::unique_ptr create_with_same_type( + communicator base, + const distributed::index_map_variant& imap) const override; + + /** + * Creates the inverse neighborhood_communicator by switching sources + * and destinations. + * + * @return collective_communicator with the inverse communication pattern + */ + [[nodiscard]] std::unique_ptr create_inverse() + const override; + + /** + * @copydoc collective_communicator::get_recv_size + */ + [[nodiscard]] comm_index_type get_recv_size() const override; + + /** + * @copydoc collective_communicator::get_recv_size + */ + [[nodiscard]] comm_index_type get_send_size() const override; + +private: + communicator comm_; + + std::vector send_sizes_; + std::vector send_offsets_; + std::vector recv_sizes_; + std::vector recv_offsets_; +}; + + +} // namespace mpi +} // namespace experimental +} // namespace gko + +#endif +#endif // GKO_PUBLIC_CORE_DISTRIBUTED_DENSE_COMMUNICATOR_HPP_ diff --git a/include/ginkgo/ginkgo.hpp b/include/ginkgo/ginkgo.hpp index 2a480496c96..7bf6a8c477d 100644 --- a/include/ginkgo/ginkgo.hpp +++ b/include/ginkgo/ginkgo.hpp @@ -60,6 +60,7 @@ #include #include +#include #include #include #include