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/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;
 
 
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,
    array<char>& 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>,
                            array<char>& workspace);
 
 
}  // namespace detail
 

template <typename LocalIndexType = int32>
class RowGatherer final
    : public EnablePolymorphicObject<RowGatherer<LocalIndexType>>,
      public EnablePolymorphicAssignment<RowGatherer<LocalIndexType>>,
      public DistributedBase {
    friend class EnablePolymorphicObject<RowGatherer, PolymorphicObject>;
    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,
                                          array<char>& 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>, array<char>& 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,
                                           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);
 
    /*
     * Create method for an empty RowGatherer with an template for the
     * collective communicator.
     *
     * This is mainly used for creating a new RowGatherer with the same runtime
     * type for the collective communicator, e.g.:
     * ```c++
     * auto rg = RowGatherer<>::create(
     *   exec, std::make_shared<mpi::NeighborhoodCommunicator>(comm));
     * ...
     * rg = RowGatherer<>::create(
     *   exec,
     *   rg->get_collective_communicator()->create_with_same_type(comm, &imap),
     *   imap);
     * ```
     */
    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) noexcept;
 
    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, array<char>& 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>,
                                array<char>& 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);
 
    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_