dense_communicator.hpp

// SPDX-FileCopyrightText: 2017 - 2025 The Ginkgo authors
//
// SPDX-License-Identifier: BSD-3-Clause
 
#ifndef GKO_PUBLIC_CORE_DISTRIBUTED_DENSE_COMMUNICATOR_HPP_
#define GKO_PUBLIC_CORE_DISTRIBUTED_DENSE_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 {
 
 
class DenseCommunicator final : public CollectiveCommunicator {
public:
    using CollectiveCommunicator::i_all_to_all_v;
 
    DenseCommunicator(const DenseCommunicator& other) = default;
 
    DenseCommunicator(DenseCommunicator&& other);
 
    DenseCommunicator& operator=(const DenseCommunicator& other) = default;
 
    DenseCommunicator& operator=(DenseCommunicator&& other);
 
    explicit DenseCommunicator(communicator base);
 
    template <typename LocalIndexType, typename GlobalIndexType>
    DenseCommunicator(
        communicator base,
        const distributed::index_map<LocalIndexType, GlobalIndexType>& imap);
 
    [[nodiscard]] std::unique_ptr<CollectiveCommunicator> create_with_same_type(
        communicator base, index_map_ptr imap) const override;
 
    [[nodiscard]] std::unique_ptr<CollectiveCommunicator> create_inverse()
        const override;
 
    [[nodiscard]] comm_index_type get_recv_size() const override;
 
    [[nodiscard]] comm_index_type get_send_size() const override;
 
    friend bool operator==(const DenseCommunicator& a,
                           const DenseCommunicator& b);
 
    friend bool operator!=(const DenseCommunicator& a,
                           const DenseCommunicator& b);
 
protected:
    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 override;
 
private:
    communicator comm_;
 
    std::vector<comm_index_type> send_sizes_;
    std::vector<comm_index_type> send_offsets_;
    std::vector<comm_index_type> recv_sizes_;
    std::vector<comm_index_type> recv_offsets_;
};
 
 
}  // namespace mpi
}  // namespace experimental
}  // namespace gko
 
 
#endif
#endif  // GKO_PUBLIC_CORE_DISTRIBUTED_DENSE_COMMUNICATOR_HPP_