row_gatherer.hpp

// SPDX-FileCopyrightText: 2017 - 2025 The Ginkgo authors
//
// SPDX-License-Identifier: BSD-3-Clause
 
#ifndef GKO_PUBLIC_CORE_DISTRIBUTED_ROW_GATHERER_HPP_
#define GKO_PUBLIC_CORE_DISTRIBUTED_ROW_GATHERER_HPP_
 
 
#include <ginkgo/config.hpp>
 
 
#if GINKGO_BUILD_MPI
 
 
#include <ginkgo/core/base/dense_cache.hpp>
#include <ginkgo/core/base/lin_op.hpp>
#include <ginkgo/core/base/mpi.hpp>
#include <ginkgo/core/distributed/base.hpp>
#include <ginkgo/core/distributed/collective_communicator.hpp>
#include <ginkgo/core/distributed/index_map.hpp>
 
 
namespace gko {
namespace experimental {
namespace distributed {
 
 
template <typename LocalIndexType = int32>
class RowGatherer final
    : public EnablePolymorphicObject<RowGatherer<LocalIndexType>>,
      public EnablePolymorphicAssignment<RowGatherer<LocalIndexType>>,
      public DistributedBase {
    friend class EnablePolymorphicObject<RowGatherer, PolymorphicObject>;
 
public:
    [[nodiscard]] mpi::request apply_async(ptr_param<const LinOp> b,
                                           ptr_param<LinOp> x) const;
 
    [[nodiscard]] mpi::request apply_async(ptr_param<const LinOp> b,
                                           ptr_param<LinOp> x,
                                           array<char>& workspace) const;
 
    dim<2> get_size() const;
 
    std::shared_ptr<const mpi::CollectiveCommunicator>
    get_collective_communicator() const;
 
    const LocalIndexType* get_const_send_idxs() const;
 
    size_type get_num_send_idxs() const;
 
    template <typename GlobalIndexType = int64,
              typename = std::enable_if_t<sizeof(GlobalIndexType) >=
                                          sizeof(LocalIndexType)>>
    static std::unique_ptr<RowGatherer> create(
        std::shared_ptr<const Executor> exec,
        std::shared_ptr<const mpi::CollectiveCommunicator> coll_comm,
        const index_map<LocalIndexType, GlobalIndexType>& imap)
    {
        return std::unique_ptr<RowGatherer>(
            new RowGatherer(std::move(exec), std::move(coll_comm), imap));
    }
 
    /*
     * Create method for an empty RowGatherer.
     */
    static std::unique_ptr<RowGatherer> create(
        std::shared_ptr<const Executor> exec, mpi::communicator comm);
 
    RowGatherer(const RowGatherer& o);
 
    RowGatherer(RowGatherer&& o) noexcept;
 
    RowGatherer& operator=(const RowGatherer& o);
 
    RowGatherer& operator=(RowGatherer&& o);
 
private:
    template <typename GlobalIndexType>
    RowGatherer(std::shared_ptr<const Executor> exec,
                std::shared_ptr<const mpi::CollectiveCommunicator> coll_comm,
                const index_map<LocalIndexType, GlobalIndexType>& imap);
 
    RowGatherer(std::shared_ptr<const Executor> exec, mpi::communicator comm);
 
    dim<2> size_;
    std::shared_ptr<const mpi::CollectiveCommunicator> coll_comm_;
    array<LocalIndexType> send_idxs_;
    mutable array<char> send_workspace_;
};
 
 
}  // namespace distributed
}  // namespace experimental
}  // namespace gko
 
#endif
#endif  // GKO_PUBLIC_CORE_DISTRIBUTED_ROW_GATHERER_HPP_