adds neighborhood implementation of collective communicator

core/CMakeLists.txt

@@ -107,6 +107,7 @@ if(GINKGO_BUILD_MPI)
  mpi/exception.cpp
  distributed/matrix.cpp
  distributed/partition_helpers.cpp
+ distributed/neighborhood_communicator.cpp
  distributed/vector.cpp
  distributed/preconditioner/schwarz.cpp)
 endif()

core/distributed/neighborhood_communicator.cpp

@@ -0,0 +1,223 @@
+// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors
+//
+// SPDX-License-Identifier: BSD-3-Clause
+
+#include <ginkgo/core/distributed/neighborhood_communicator.hpp>
+
+
+#include <ginkgo/core/base/precision_dispatch.hpp>
+#include <ginkgo/core/matrix/dense.hpp>
+
+
+namespace gko {
+namespace experimental {
+namespace mpi {
+
+
+/**
+ * \brief Computes the inverse envelope (target-ids, sizes) for a given
+ * one-sided communication pattern.
+ *
+ * \param exec the executor, this will always use the host executor
+ * \param comm communicator
+ * \param ids target ids of the one-sided operation
+ * \param sizes number of elements send to each id
+ *
+ * \return the inverse envelope consisting of the target-ids and the sizes
+ */
+std::tuple<std::vector<comm_index_type>, std::vector<comm_index_type>>
+communicate_inverse_envelope(std::shared_ptr<const Executor> exec,
+ mpi::communicator comm,
+ const std::vector<comm_index_type>& ids,
+ const std::vector<comm_index_type>& sizes)
+{
+ auto host_exec = exec->get_master();
+ std::vector<comm_index_type> inverse_sizes_full(comm.size());
+ mpi::window<comm_index_type> window(host_exec, inverse_sizes_full.data(),
+ inverse_sizes_full.size(), comm,
+ sizeof(comm_index_type), MPI_INFO_ENV);
+ window.fence();
+ for (int i = 0; i < ids.size(); ++i) {
+ window.put(host_exec, sizes.data() + i, 1, ids[i], comm.rank(), 1);
+ }
+ window.fence();
+
+ std::vector<comm_index_type> inverse_sizes;
+ std::vector<comm_index_type> inverse_ids;
+ for (int i = 0; i < inverse_sizes_full.size(); ++i) {
+ if (inverse_sizes_full[i] > 0) {
+ inverse_ids.push_back(i);
+ inverse_sizes.push_back(inverse_sizes_full[i]);
+ }
+ }
+
+ return std::make_tuple(std::move(inverse_ids), std::move(inverse_sizes));
+}
+
+
+/**
+ * Creates a distributed graph communicator based on the input sources and
+ * destinations.
+ *
+ * The graph is unweighted and has the same rank ordering as the input
+ * communicator.
+ */
+mpi::communicator create_neighborhood_comm(
+ mpi::communicator base, const std::vector<comm_index_type>& sources,
+ const std::vector<comm_index_type>& destinations)
+{
+ auto in_degree = static_cast<comm_index_type>(sources.size());
+ auto out_degree = static_cast<comm_index_type>(destinations.size());
+
+ // adjacent constructor guarantees that querying sources/destinations
+ // will result in the array having the same order as defined here
+ MPI_Comm graph_comm;
+ MPI_Info info;
+ GKO_ASSERT_NO_MPI_ERRORS(MPI_Info_dup(MPI_INFO_ENV, &info));
+ GKO_ASSERT_NO_MPI_ERRORS(MPI_Dist_graph_create_adjacent(
+ base.get(), in_degree, sources.data(),
+ in_degree ? MPI_UNWEIGHTED : MPI_WEIGHTS_EMPTY, out_degree,
+ destinations.data(), out_degree ? MPI_UNWEIGHTED : MPI_WEIGHTS_EMPTY,
+ info, false, &graph_comm));
+ GKO_ASSERT_NO_MPI_ERRORS(MPI_Info_free(&info));
+
+ return mpi::communicator::create_owning(graph_comm,
+ base.force_host_buffer());
+}
+
+std::unique_ptr<collective_communicator>
+neighborhood_communicator::create_inverse() const
+{
+ auto base_comm = this->get_base_communicator();
+ distributed::comm_index_type num_sources;
+ distributed::comm_index_type num_destinations;
+ distributed::comm_index_type weighted;
+ GKO_ASSERT_NO_MPI_ERRORS(MPI_Dist_graph_neighbors_count(
+ comm_.get(), &num_sources, &num_destinations, &weighted));
+
+ std::vector<distributed::comm_index_type> sources(num_sources);
+ std::vector<distributed::comm_index_type> destinations(num_destinations);
+ GKO_ASSERT_NO_MPI_ERRORS(MPI_Dist_graph_neighbors(
+ comm_.get(), num_sources, sources.data(), MPI_UNWEIGHTED,
+ num_destinations, destinations.data(), MPI_UNWEIGHTED));
+
+ return std::unique_ptr<collective_communicator>{
+ new neighborhood_communicator(base_comm, destinations, recv_sizes_,
+ recv_offsets_, sources, send_sizes_,
+ send_offsets_)};
+}
+
+
+comm_index_type neighborhood_communicator::get_recv_size() const
+{
+ return recv_offsets_.back();
+}
+
+
+comm_index_type neighborhood_communicator::get_send_size() const
+{
+ return send_offsets_.back();
+}
+
+
+neighborhood_communicator::neighborhood_communicator(
+ communicator base, const std::vector<distributed::comm_index_type>& sources,
+ const std::vector<comm_index_type>& send_sizes,
+ const std::vector<comm_index_type>& send_offsets,
+ const std::vector<distributed::comm_index_type>& destinations,
+ const std::vector<comm_index_type>& recv_sizes,
+ const std::vector<comm_index_type>& recv_offsets)
+ : collective_communicator(base), comm_(MPI_COMM_NULL)
+{
+ comm_ = create_neighborhood_comm(base, sources, destinations);
+ send_sizes_ = send_sizes;
+ send_offsets_ = send_offsets;
+ recv_sizes_ = recv_sizes;
+ recv_offsets_ = recv_offsets;
+}
+
+
+neighborhood_communicator::neighborhood_communicator(communicator base)
+ : collective_communicator(std::move(base)),
+ comm_(MPI_COMM_SELF),
+ send_sizes_(),
+ send_offsets_(1),
+ recv_sizes_(),
+ recv_offsets_(1)
+{
+ // ensure that comm_ always has the correct topology
+ std::vector<comm_index_type> non_nullptr(1);
+ non_nullptr.resize(0);
+ comm_ = create_neighborhood_comm(this->get_base_communicator(), non_nullptr,
+ non_nullptr);
+}
+
+
+request neighborhood_communicator::i_all_to_all_v(
+ std::shared_ptr<const Executor> exec, const void* send_buffer,
+ MPI_Datatype send_type, void* recv_buffer, MPI_Datatype recv_type) const
+{
+ auto guard = exec->get_scoped_device_id_guard();
+ request req;
+ GKO_ASSERT_NO_MPI_ERRORS(MPI_Ineighbor_alltoallv(
+ send_buffer, send_sizes_.data(), send_offsets_.data(), send_type,
+ recv_buffer, recv_sizes_.data(), recv_offsets_.data(), recv_type,
+ comm_.get(), req.get()));
+ return req;
+}
+
+
+template <typename LocalIndexType, typename GlobalIndexType>
+neighborhood_communicator::neighborhood_communicator(
+ communicator base,
+ const distributed::index_map<LocalIndexType, GlobalIndexType>& imap)
+ : collective_communicator(base),
+ comm_(MPI_COMM_SELF),
+ recv_sizes_(imap.get_remote_target_ids().get_size()),
+ recv_offsets_(recv_sizes_.size() + 1),
+ send_offsets_(1)
+{
+ 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());
+ std::vector<comm_index_type> recv_target_ids(
+ recv_target_ids_arr->get_size());
+ std::copy_n(recv_target_ids_arr->get_const_data(),
+ recv_target_ids_arr->get_size(), recv_target_ids.begin());
+ std::transform(imap.get_remote_global_idxs().begin(),
+ imap.get_remote_global_idxs().end(), recv_sizes_.begin(),
+ [](const auto& a) { return a.get_size(); });
+ auto send_envelope =
+ communicate_inverse_envelope(exec, base, recv_target_ids, recv_sizes_);
+ const auto& send_target_ids = std::get<0>(send_envelope);
+ send_sizes_ = std::move(std::get<1>(send_envelope));
+
+ send_offsets_.resize(send_sizes_.size() + 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);
+
+ comm_ = create_neighborhood_comm(base, recv_target_ids, send_target_ids);
+}
+
+
+#define GKO_DECLARE_NEIGHBORHOOD_CONSTRUCTOR(LocalIndexType, GlobalIndexType) \
+ neighborhood_communicator::neighborhood_communicator( \
+ communicator base, \
+ const distributed::index_map<LocalIndexType, GlobalIndexType>& imap)
+
+GKO_INSTANTIATE_FOR_EACH_LOCAL_GLOBAL_INDEX_TYPE(
+ GKO_DECLARE_NEIGHBORHOOD_CONSTRUCTOR);
+
+#undef GKO_DECLARE_NEIGHBORHOOD_CONSTRUCTOR
+
+
+} // namespace mpi
+} // namespace experimental
+} // namespace gko

include/ginkgo/core/distributed/neighborhood_communicator.hpp

@@ -0,0 +1,131 @@
+// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors
+//
+// SPDX-License-Identifier: BSD-3-Clause
+
+#ifndef GKO_PUBLIC_CORE_DISTRIBUTED_NEIGHBORHOOD_COMMUNICATOR_HPP_
+#define GKO_PUBLIC_CORE_DISTRIBUTED_NEIGHBORHOOD_COMMUNICATOR_HPP_
+
+
+#include <ginkgo/config.hpp>
+
+
+#if GINKGO_BUILD_MPI
+
+#include <ginkgo/core/base/mpi.hpp>
+#include <ginkgo/core/distributed/collective_communicator.hpp>
+#include <ginkgo/core/distributed/index_map.hpp>
+
+
+namespace gko {
+namespace experimental {
+namespace mpi {
+
+
+/**
+ * A collective_communicator that uses a neighborhood topology.
+ *
+ * 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 neighborhood_communicator final : public collective_communicator {
+public:
+ using collective_communicator::i_all_to_all_v;
+
+ /**
+ * Default constructor with empty communication pattern
+ * @param base the base communicator
+ */
+ explicit neighborhood_communicator(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 <typename LocalIndexType, typename GlobalIndexType>
+ neighborhood_communicator(
+ communicator base,
+ const distributed::index_map<LocalIndexType, GlobalIndexType>& 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
+ */
+ neighborhood_communicator(
+ communicator base,
+ const std::vector<distributed::comm_index_type>& sources,
+ const std::vector<comm_index_type>& recv_sizes,
+ const std::vector<comm_index_type>& recv_offsets,
+ const std::vector<distributed::comm_index_type>& destinations,
+ const std::vector<comm_index_type>& send_sizes,
+ const std::vector<comm_index_type>& send_offsets);
+
+ /**
+ * Communicate data from all ranks to all other ranks using the
+ * neighboorhood communication MPI_Ineighbor_alltoallv. See MPI
+ * documentation for more details
+ *
+ * @param exec The executor, on which the message buffers are located.
+ * @param send_buffer the buffer to send
+ * @param send_type the MPI_Datatype for the send buffer
+ * @param recv_buffer the buffer to gather into
+ * @param recv_type the MPI_Datatype for the recv buffer
+ *
+ * @return the request handle for the call
+ */
+ request i_all_to_all_v(std::shared_ptr<const Executor> exec,
+ const void* send_buffer, MPI_Datatype send_type,
+ void* recv_buffer,
+ MPI_Datatype recv_type) const override;
+
+ /**
+ * Creates the inverse neighborhood_communicator by switching sources
+ * and destinations.
+ *
+ * @return collective_communicator with the inverse communication pattern
+ */
+ std::unique_ptr<collective_communicator> create_inverse() const override;
+
+ /**
+ * @copydoc collective_communicator::get_recv_size
+ */
+ comm_index_type get_recv_size() const override;
+
+ /**
+ * @copydoc collective_communicator::get_recv_size
+ */
+ comm_index_type get_send_size() const override;
+
+private:
+ communicator comm_;
+
+ std::vector<distributed::comm_index_type> send_sizes_;
+ std::vector<distributed::comm_index_type> send_offsets_;
+ std::vector<distributed::comm_index_type> recv_sizes_;
+ std::vector<distributed::comm_index_type> recv_offsets_;
+};
+
+
+} // namespace mpi
+} // namespace experimental
+} // namespace gko
+
+#endif
+#endif // GKO_PUBLIC_CORE_DISTRIBUTED_NEIGHBORHOOD_COMMUNICATOR_HPP_

include/ginkgo/ginkgo.hpp

@@ -60,6 +60,7 @@
 #include <ginkgo/core/distributed/index_map.hpp>
 #include <ginkgo/core/distributed/lin_op.hpp>
 #include <ginkgo/core/distributed/matrix.hpp>
+#include <ginkgo/core/distributed/neighborhood_communicator.hpp>
 #include <ginkgo/core/distributed/partition.hpp>
 #include <ginkgo/core/distributed/partition_helpers.hpp>
 #include <ginkgo/core/distributed/polymorphic_object.hpp>