[coll-comm] adds interface for collective communicator

Co-authored-by: Pratik Nayak <pratik.nayak@kit.edu>

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
@@ -130,6 +130,7 @@ if(GINKGO_BUILD_MPI)
     target_sources(${ginkgo_core}
         PRIVATE
         mpi/exception.cpp
+        distributed/collective_communicator.cpp
         distributed/matrix.cpp
         distributed/partition_helpers.cpp
         distributed/vector.cpp
@@ -146,7 +147,7 @@ if(MSVC 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,26 @@
+// SPDX-FileCopyrightText: 2024 The Ginkgo authors
+//
+// SPDX-License-Identifier: BSD-3-Clause
+
+#include "ginkgo/core/distributed/collective_communicator.hpp"
+
+
+namespace gko {
+namespace experimental {
+namespace mpi {
+
+
+CollectiveCommunicator::CollectiveCommunicator(communicator base)
+    : base_(std::move(base))
+{}
+
+
+const communicator& CollectiveCommunicator::get_base_communicator() const
+{
+    return base_;
+}
+
+
+}  // namespace mpi
+}  // namespace experimental
+}  // namespace gko

include/ginkgo/core/distributed/collective_communicator.hpp

@@ -0,0 +1,116 @@
+// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors
+//
+// SPDX-License-Identifier: BSD-3-Clause
+
+#ifndef GKO_PUBLIC_CORE_DISTRIBUTED_COLLECTIVE_COMMUNICATOR_HPP_
+#define GKO_PUBLIC_CORE_DISTRIBUTED_COLLECTIVE_COMMUNICATOR_HPP_
+
+
+#include <ginkgo/config.hpp>
+
+
+#if GINKGO_BUILD_MPI
+
+#include <ginkgo/core/base/mpi.hpp>
+#include <ginkgo/core/distributed/index_map_fwd.hpp>
+
+
+namespace gko {
+namespace experimental {
+namespace mpi {
+
+/**
+ * Interface for an collective communicator.
+ *
+ * A collective communicator only provides routines for collective
+ * communications. At the moment this is restricted to the variable all-to-all.
+ */
+class CollectiveCommunicator {
+public:
+    virtual ~CollectiveCommunicator() = default;
+
+    explicit CollectiveCommunicator(communicator base = MPI_COMM_NULL);
+
+    [[nodiscard]] const communicator& get_base_communicator() const;
+
+    /**
+     * Non-blocking all-to-all communication.
+     *
+     * The send_buffer must have size get_send_size, and the recv_buffer
+     * must have size get_recv_size.
+     *
+     * @tparam SendType  the type of the elements to send
+     * @tparam RecvType  the type of the elements to receive
+     * @param exec  the executor for the communication
+     * @param send_buffer  the send buffer
+     * @param recv_buffer  the receive buffer
+     * @return  a request handle
+     */
+    template <typename SendType, typename RecvType>
+    request i_all_to_all_v(std::shared_ptr<const Executor> exec,
+                           const SendType* send_buffer,
+                           RecvType* recv_buffer) const
+    {
+        return this->i_all_to_all_v(
+            std::move(exec), send_buffer, type_impl<SendType>::get_type(),
+            recv_buffer, type_impl<RecvType>::get_type());
+    }
+
+    /**
+     * @copydoc i_all_to_all_v(std::shared_ptr<const Executor>, const SendType*
+     *          send_buffer, RecvType* recv_buffer)
+     */
+    virtual 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 = 0;
+
+    /**
+     * Creates a new collective_communicator with the same dynamic type.
+     *
+     * @param base  The base communicator
+     * @param imap  The index_map that defines the communication pattern
+     *
+     * @return  a collective_communicator with the same dynamic type
+     */
+    [[nodiscard]] virtual std::unique_ptr<CollectiveCommunicator>
+    create_with_same_type(communicator base,
+                          const distributed::index_map_variant& imap) const = 0;
+
+    /**
+     * Creates a collective_communicator with the inverse communication pattern
+     * than this object.
+     *
+     * @return  a collective_communicator with the inverse communication
+     * pattern.
+     */
+    [[nodiscard]] virtual std::unique_ptr<CollectiveCommunicator>
+    create_inverse() const = 0;
+
+    /**
+     * Get the total number of received elements this communication patterns
+     * expects.
+     *
+     * @return  number of received elements.
+     */
+    [[nodiscard]] virtual comm_index_type get_recv_size() const = 0;
+
+    /**
+     * Get the total number of sent elements this communication patterns
+     * expects.
+     *
+     * @return  number of sent elements.
+     */
+    [[nodiscard]] virtual comm_index_type get_send_size() const = 0;
+
+private:
+    communicator base_;
+};
+
+
+}  // namespace mpi
+}  // namespace experimental
+}  // namespace gko
+
+#endif
+#endif  // GKO_PUBLIC_CORE_DISTRIBUTED_COLLECTIVE_COMMUNICATOR_HPP_

include/ginkgo/ginkgo.hpp

@@ -59,6 +59,7 @@
 #include <ginkgo/core/config/type_descriptor.hpp>
 
 #include <ginkgo/core/distributed/base.hpp>
+#include <ginkgo/core/distributed/collective_communicator.hpp>
 #include <ginkgo/core/distributed/index_map.hpp>
 #include <ginkgo/core/distributed/index_map_fwd.hpp>
 #include <ginkgo/core/distributed/lin_op.hpp>