collective_communicator.hpp

// SPDX-FileCopyrightText: 2017 - 2026 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 <variant>
 
#include <ginkgo/core/base/mpi.hpp>
#include <ginkgo/core/distributed/index_map.hpp>
 
 
namespace gko {
namespace experimental {
namespace mpi {
 
 
class CollectiveCommunicator {
public:
    using index_map_ptr =
        std::variant<const distributed::index_map<int32, int32>*,
                     const distributed::index_map<int32, int64>*,
                     const distributed::index_map<int64, int64>*>;
 
    virtual ~CollectiveCommunicator() = default;
 
    explicit CollectiveCommunicator(communicator base = MPI_COMM_NULL);
 
    [[nodiscard]] const communicator& get_base_communicator() const;
 
    template <typename SendType, typename RecvType>
    [[nodiscard]] request i_all_to_all_v(std::shared_ptr<const Executor> exec,
                                         const SendType* send_buffer,
                                         RecvType* recv_buffer) const;
 
    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;
 
    [[nodiscard]] virtual std::unique_ptr<CollectiveCommunicator>
    create_with_same_type(communicator base, index_map_ptr imap) const = 0;
 
    [[nodiscard]] virtual std::unique_ptr<CollectiveCommunicator>
    create_inverse() const = 0;
 
 
    [[nodiscard]] virtual std::pair<std::unique_ptr<CollectiveCommunicator>,
                                    std::vector<comm_index_type>>
    resize(std::shared_ptr<const Executor> exec,
           const std::vector<comm_index_type>& send_factors) const = 0;
 
    [[nodiscard]] virtual comm_index_type get_recv_size() const = 0;
 
    [[nodiscard]] virtual comm_index_type get_send_size() const = 0;
 
protected:
    virtual request 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 = 0;
 
private:
    communicator base_;
};
 
 
template <typename SendType, typename RecvType>
request CollectiveCommunicator::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());
}
 
 
namespace detail {
 
 
std::shared_ptr<CollectiveCommunicator> create_default_collective_communicator(
    communicator base);
 
 
}  // namespace detail
}  // namespace mpi
}  // namespace experimental
}  // namespace gko
 
 
#endif
#endif  // GKO_PUBLIC_CORE_DISTRIBUTED_COLLECTIVE_COMMUNICATOR_HPP_