row_gatherer.hpp

// SPDX-FileCopyrightText: 2017 - 2026 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/event.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>
class RowGatherer;
 
template <typename LocalIndexType>
class RowScatterer;
 
 
namespace detail {
 
 
// give access to test function on protected function
template <typename LocalIndexType>
std::shared_ptr<const gko::detail::Event> apply_prepare(
    const RowGatherer<LocalIndexType>* rg, ptr_param<const LinOp> b);
 
// give access to test function on protected function
template <typename LocalIndexType>
std::shared_ptr<const gko::detail::Event> apply_prepare(
    const RowGatherer<LocalIndexType>* rg, ptr_param<const LinOp> b,
    gko::detail::GenericDenseCache& workspace);
 
// give access to test function on protected function
template <typename LocalIndexType>
mpi::request apply_finalize(const RowGatherer<LocalIndexType>* rg,
                            ptr_param<const LinOp> b, ptr_param<LinOp> x,
                            std::shared_ptr<const gko::detail::Event>);
 
// give access to test function on protected function
template <typename LocalIndexType>
mpi::request apply_finalize(const RowGatherer<LocalIndexType>* rg,
                            ptr_param<const LinOp> b, ptr_param<LinOp> x,
                            std::shared_ptr<const gko::detail::Event>,
                            gko::detail::GenericDenseCache& workspace);
 
 
}  // namespace detail
 
 
template <typename LocalIndexType = int32>
class RowGatherer final : public PolymorphicObject,
                          public EnableCloneable<RowGatherer<LocalIndexType>>,
                          public DistributedBase {
    friend class EnableCloneable<RowGatherer>;
    template <typename ValueT, typename LocalIndexT, typename GlobalIndexT>
    friend class Matrix;
    // for test purpose
    friend std::shared_ptr<const gko::detail::Event>
    detail::apply_prepare<LocalIndexType>(const RowGatherer* rg,
                                          ptr_param<const LinOp> b);
    friend std::shared_ptr<const gko::detail::Event> detail::apply_prepare<
        LocalIndexType>(const RowGatherer* rg, ptr_param<const LinOp> b,
                        gko::detail::GenericDenseCache& workspace);
    friend mpi::request detail::apply_finalize<LocalIndexType>(
        const RowGatherer* rg, ptr_param<const LinOp> b, ptr_param<LinOp> x,
        std::shared_ptr<const gko::detail::Event>);
    friend mpi::request detail::apply_finalize<LocalIndexType>(
        const RowGatherer* rg, ptr_param<const LinOp> b, ptr_param<LinOp> x,
        std::shared_ptr<const gko::detail::Event>,
        gko::detail::GenericDenseCache& workspace);
 
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,
        gko::detail::GenericDenseCache& 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));
    }
 
    static std::unique_ptr<RowGatherer> create_from_scatterer(
        std::shared_ptr<const Executor> exec,
        ptr_param<const RowScatterer<LocalIndexType>> scatterer);
 
    static std::unique_ptr<RowGatherer> create(
        std::shared_ptr<const Executor> exec, mpi::communicator comm);
 
    static std::unique_ptr<RowGatherer> create(
        std::shared_ptr<const Executor> exec,
        std::shared_ptr<const mpi::CollectiveCommunicator> coll_comm_template);
 
    RowGatherer(const RowGatherer& o);
 
    RowGatherer(RowGatherer&& o);
 
    RowGatherer& operator=(const RowGatherer& o);
 
    RowGatherer& operator=(RowGatherer&& o);
 
protected:
    std::shared_ptr<const gko::detail::Event> apply_prepare(
        ptr_param<const LinOp> b) const;
 
    std::shared_ptr<const gko::detail::Event> apply_prepare(
        ptr_param<const LinOp> b,
        gko::detail::GenericDenseCache& workspace) const;
 
    mpi::request apply_finalize(
        ptr_param<const LinOp> b, ptr_param<LinOp> x,
        std::shared_ptr<const gko::detail::Event>) const;
 
    mpi::request apply_finalize(
        ptr_param<const LinOp> b, ptr_param<LinOp> x,
        std::shared_ptr<const gko::detail::Event>,
        gko::detail::GenericDenseCache& workspace) const;
 
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);
 
    RowGatherer(
        std::shared_ptr<const Executor> exec,
        std::shared_ptr<const mpi::CollectiveCommunicator> coll_comm_template);
 
    RowGatherer(std::shared_ptr<const Executor> exec,
                std::shared_ptr<const mpi::CollectiveCommunicator> coll_comm,
                array<LocalIndexType> send_idxs, dim<2> size);
 
    dim<2> size_;
    std::shared_ptr<const mpi::CollectiveCommunicator> coll_comm_;
    array<LocalIndexType> send_idxs_;
    mutable gko::detail::GenericDenseCache send_cache_;
};
 
 
}  // namespace distributed
}  // namespace experimental
}  // namespace gko
 
#endif
#endif  // GKO_PUBLIC_CORE_DISTRIBUTED_ROW_GATHERER_HPP_