From 898cf0090a654c4f9c65537ef10d9dbf88b467b2 Mon Sep 17 00:00:00 2001
From: Marcel Koch <marcel.koch@kit.edu>
Date: Thu, 4 Apr 2024 10:05:03 +0000
Subject: [PATCH] adds interface for collective communicator

---
 .../distributed/collective_communicator.hpp   | 104 ++++++++++++++++++
 include/ginkgo/ginkgo.hpp                     |   1 +
 2 files changed, 105 insertions(+)
 create mode 100644 include/ginkgo/core/distributed/collective_communicator.hpp

diff --git a/include/ginkgo/core/distributed/collective_communicator.hpp b/include/ginkgo/core/distributed/collective_communicator.hpp
new file mode 100644
index 00000000000..2183cefb5a3
--- /dev/null
+++ b/include/ginkgo/core/distributed/collective_communicator.hpp
@@ -0,0 +1,104 @@
+// 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>
+
+
+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 collective_communicator {
+public:
+    virtual ~collective_communicator() = default;
+
+    explicit collective_communicator(communicator base) : base_(std::move(base))
+    {}
+
+    const communicator& get_base_communicator() const { return base_; }
+
+    /**
+     * 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 collective_communicator with the inverse communication pattern
+     * than this object.
+     *
+     * @return  a collective_communicator with the inverse communication
+     * pattern.
+     */
+    virtual std::unique_ptr<collective_communicator> create_inverse() const = 0;
+
+    /**
+     * Get the total number of received elements this communication patterns
+     * expects.
+     *
+     * @return  number of received elements.
+     */
+    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.
+     */
+    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_
diff --git a/include/ginkgo/ginkgo.hpp b/include/ginkgo/ginkgo.hpp
index 4335ee6d3c2..7a9d3927261 100644
--- a/include/ginkgo/ginkgo.hpp
+++ b/include/ginkgo/ginkgo.hpp
@@ -56,6 +56,7 @@
 #include <ginkgo/core/config/property_tree.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/lin_op.hpp>
 #include <ginkgo/core/distributed/matrix.hpp>