Add specific communicator for neighborhood communication

CMakeLists.txt

@@ -239,7 +239,7 @@ if(GINKGO_BUILD_HWLOC AND (MSVC OR WIN32 OR CYGWIN OR APPLE))
 endif()
 
 set(GINKGO_HAVE_GPU_AWARE_MPI OFF)
-set(GINKGO_FORCE_SPMV_BLOCKING_COMM OFF)
+set(GINKGO_HAVE_OPENMPI_PRE_4_1_X OFF)
 if(GINKGO_BUILD_MPI)
     find_package(MPI 3.1 COMPONENTS CXX REQUIRED)
     if(GINKGO_FORCE_GPU_AWARE_MPI)
@@ -264,12 +264,14 @@ if(GINKGO_BUILD_MPI)
         )
         if(NOT valid_openmpi_version)
             message(WARNING
-                "OpenMPI v4.0.x has a bug that forces us to use blocking communication in our distributed "
+                "OpenMPI v4.0.x has several bugs that forces us to use non-optimal communication in our distributed "
                 "matrix class. To enable faster, non-blocking communication, consider updating your OpenMPI version or "
                 "switch to a different vendor.")
-            set(GINKGO_FORCE_SPMV_BLOCKING_COMM ON)
+            set(GINKGO_HAVE_OPENMPI_PRE_4_1_X ON)
         endif()
+        unset(valid_openmpi_version)
     endif()
+    unset(uses_openmpi)
 endif()
 
 # Try to find the third party packages before using our subdirectories

core/CMakeLists.txt

@@ -1,6 +1,6 @@
 add_subdirectory(device_hooks) # placeholders for disabled modules
 
-set(config_source 
+set(config_source
     config/factorization_config.cpp
     config/multigrid_config.cpp
     config/preconditioner_config.cpp
@@ -142,7 +142,10 @@ if(GINKGO_BUILD_MPI)
         distributed/vector_cache.cpp
         mpi/exception.cpp
         distributed/assembly.cpp
+        distributed/collective_communicator.cpp
+        distributed/dense_communicator.cpp
         distributed/matrix.cpp
+        distributed/neighborhood_communicator.cpp
         distributed/partition_helpers.cpp
         distributed/vector.cpp
         distributed/preconditioner/schwarz.cpp)
@@ -158,7 +161,7 @@ if((MSVC OR MINGW) AND BUILD_SHARED_LIBS)
     if(GINKGO_CHECK_CIRCULAR_DEPS)
         ginkgo_check_headers(ginkgo "")
     endif()
-else() 
+else()
     target_sources(${ginkgo_core} PRIVATE ${config_source})
 endif()
 

core/distributed/collective_communicator.cpp

@@ -0,0 +1,37 @@
+// SPDX-FileCopyrightText: 2024 The Ginkgo authors
+//
+// SPDX-License-Identifier: BSD-3-Clause
+
+#include "ginkgo/core/distributed/collective_communicator.hpp"
+
+#include <mpi.h>
+
+
+namespace gko {
+namespace experimental {
+namespace mpi {
+
+
+CollectiveCommunicator::CollectiveCommunicator(communicator base)
+    : base_(std::move(base))
+{}
+
+
+const communicator& CollectiveCommunicator::get_base_communicator() const
+{
+    return base_;
+}
+
+
+request CollectiveCommunicator::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
+{
+    return this->i_all_to_all_v_impl(std::move(exec), send_buffer, send_type,
+                                     recv_buffer, recv_type);
+}
+
+
+}  // namespace mpi
+}  // namespace experimental
+}  // namespace gko

core/distributed/dense_communicator.cpp

@@ -0,0 +1,166 @@
+// 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 {
+
+
+size_type get_comm_size_safe(const communicator& comm)
+{
+    if (comm.get() == MPI_COMM_NULL) {
+        return 0;
+    }
+    return comm.size();
+}
+
+
+DenseCommunicator::DenseCommunicator(communicator base)
+    : CollectiveCommunicator(base),
+      comm_(base),
+      recv_sizes_(get_comm_size_safe(comm_)),
+      recv_offsets_(get_comm_size_safe(comm_) + 1),
+      send_sizes_(get_comm_size_safe(comm_)),
+      send_offsets_(get_comm_size_safe(comm_) + 1)
+{}
+
+
+template <typename LocalIndexType, typename GlobalIndexType>
+DenseCommunicator::DenseCommunicator(
+    communicator base,
+    const distributed::index_map<LocalIndexType, GlobalIndexType>& 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<LocalIndexType, GlobalIndexType>& imap)
+
+GKO_INSTANTIATE_FOR_EACH_LOCAL_GLOBAL_INDEX_TYPE(GKO_DECLARE_DENSE_CONSTRUCTOR);
+
+#undef GKO_DECLARE_DENSE_CONSTRUCTOR
+
+
+DenseCommunicator::DenseCommunicator(DenseCommunicator&& other)
+    : DenseCommunicator(other.get_base_communicator())
+{
+    *this = std::move(other);
+}
+
+
+DenseCommunicator& DenseCommunicator::operator=(DenseCommunicator&& other)
+{
+    if (this != &other) {
+        comm_ = std::exchange(other.comm_, MPI_COMM_NULL);
+        send_sizes_ =
+            std::exchange(other.send_sizes_, std::vector<comm_index_type>{});
+        send_offsets_ =
+            std::exchange(other.send_offsets_, std::vector<comm_index_type>{0});
+        recv_sizes_ =
+            std::exchange(other.recv_sizes_, std::vector<comm_index_type>{});
+        recv_offsets_ =
+            std::exchange(other.recv_offsets_, std::vector<comm_index_type>{0});
+    }
+    return *this;
+}
+
+
+request DenseCommunicator::i_all_to_all_v_impl(
+    std::shared_ptr<const Executor> exec, const void* send_buffer,
+    MPI_Datatype send_type, void* recv_buffer, MPI_Datatype recv_type) const
+{
+#ifdef GINKGO_HAVE_OPENMPI_PRE_4_1_X
+    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<CollectiveCommunicator>
+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<DenseCommunicator>(base, imap);
+        },
+        imap);
+}
+
+
+std::unique_ptr<CollectiveCommunicator> DenseCommunicator::create_inverse()
+    const
+{
+    auto inv = std::make_unique<DenseCommunicator>(comm_);
+    inv->send_sizes_ = recv_sizes_;
+    inv->send_offsets_ = recv_offsets_;
+    inv->recv_sizes_ = send_sizes_;
+    inv->recv_offsets_ = send_offsets_;
+    return inv;
+}
+
+
+comm_index_type DenseCommunicator::get_recv_size() const
+{
+    return recv_offsets_.back();
+}
+
+
+comm_index_type DenseCommunicator::get_send_size() const
+{
+    return send_offsets_.back();
+}
+
+
+bool operator==(const DenseCommunicator& a, const DenseCommunicator& b)
+{
+    return (a.comm_.is_identical(b.comm_) || a.comm_.is_congruent(b.comm_)) &&
+           a.send_sizes_ == b.send_sizes_ && a.recv_sizes_ == b.recv_sizes_ &&
+           a.send_offsets_ == b.send_offsets_ &&
+           a.recv_offsets_ == b.recv_offsets_;
+}
+
+
+bool operator!=(const DenseCommunicator& a, const DenseCommunicator& b)
+{
+    return !(a == b);
+}
+
+
+}  // namespace mpi
+}  // namespace experimental
+}  // namespace gko

core/distributed/device_partition.hpp

@@ -2,8 +2,8 @@
 //
 // SPDX-License-Identifier: BSD-3-Clause
 
-#ifndef GINKGO_PARTITION_HPP
-#define GINKGO_PARTITION_HPP
+#ifndef GKO_CORE_DISTRIBUTED_PARTITION_HPP
+#define GKO_CORE_DISTRIBUTED_PARTITION_HPP
 
 #include <ginkgo/core/distributed/partition.hpp>
 
@@ -89,4 +89,4 @@ to_device_const(
 }  // namespace gko
 
 
-#endif  // GINKGO_PARTITION_HPP
+#endif  // GKO_CORE_DISTRIBUTED_PARTITION_HPP

core/distributed/matrix.cpp

@@ -472,7 +472,7 @@ mpi::request Matrix<ValueType, LocalIndexType, GlobalIndexType>::communicate(
     auto recv_ptr = use_host_buffer ? host_recv_buffer_->get_values()
                                     : recv_buffer_->get_values();
     exec->synchronize();
-#ifdef GINKGO_FORCE_SPMV_BLOCKING_COMM
+#ifdef GINKGO_HAVE_OPENMPI_PRE_4_1_X
     comm.all_to_all_v(use_host_buffer ? exec->get_master() : exec, send_ptr,
                       send_sizes_.data(), send_offsets_.data(), type.get(),
                       recv_ptr, recv_sizes_.data(), recv_offsets_.data(),