row_scatterer.hpp

// SPDX-FileCopyrightText: 2017 - 2026 The Ginkgo authors
//
// SPDX-License-Identifier: BSD-3-Clause
 
#ifndef GKO_PUBLIC_CORE_DISTRIBUTED_ROW_SCATTERER_HPP_
#define GKO_PUBLIC_CORE_DISTRIBUTED_ROW_SCATTERER_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>
#include <ginkgo/core/distributed/row_gatherer.hpp>
 
 
namespace gko {
namespace experimental {
namespace distributed {
 
 
template <typename LocalIndexType = int32>
class RowScatterer final : public PolymorphicObject,
                           public EnableCloneable<RowScatterer<LocalIndexType>>,
                           public DistributedBase {
    friend class EnableCloneable<RowScatterer>;
 
public:
    void apply_async(ptr_param<const LinOp> local_values) const;
 
    void wait_and_accumulate(ptr_param<LinOp> distributed_target) const;
 
    dim<2> get_size() const;
 
    std::shared_ptr<const mpi::CollectiveCommunicator>
    get_collective_communicator() const;
 
    const LocalIndexType* get_const_recv_idxs() const;
 
    size_type get_num_recv_idxs() const;
 
    template <typename GlobalIndexType = int64,
              typename = std::enable_if_t<sizeof(GlobalIndexType) >=
                                          sizeof(LocalIndexType)>>
    static std::unique_ptr<RowScatterer> 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<RowScatterer>(
            new RowScatterer(std::move(exec), std::move(coll_comm), imap));
    }
 
    static std::unique_ptr<RowScatterer> create_from_gatherer(
        std::shared_ptr<const Executor> exec,
        ptr_param<const RowGatherer<LocalIndexType>> gatherer);
 
    static std::unique_ptr<RowScatterer> create(
        std::shared_ptr<const Executor> exec, mpi::communicator comm);
 
    RowScatterer(const RowScatterer& o);
 
    RowScatterer(RowScatterer&& o);
 
    RowScatterer& operator=(const RowScatterer& o);
 
    RowScatterer& operator=(RowScatterer&& o);
 
private:
    template <typename GlobalIndexType>
    RowScatterer(std::shared_ptr<const Executor> exec,
                 std::shared_ptr<const mpi::CollectiveCommunicator> coll_comm,
                 const index_map<LocalIndexType, GlobalIndexType>& imap);
 
    RowScatterer(std::shared_ptr<const Executor> exec, mpi::communicator comm);
 
    RowScatterer(std::shared_ptr<const Executor> exec,
                 std::shared_ptr<const mpi::CollectiveCommunicator> coll_comm,
                 array<LocalIndexType> recv_idxs, dim<2> size);
 
    dim<2> size_;
    std::shared_ptr<const mpi::CollectiveCommunicator> coll_comm_;
    array<LocalIndexType> recv_idxs_;  // local indices to accumulate into
    mutable gko::detail::GenericDenseCache send_cache_;
    // request and receive buffer of the active apply_async, consumed by
    // wait_and_accumulate
    mutable mpi::request req_;
    mutable std::shared_ptr<LinOp> recv_buffer_;
    mutable gko::detail::GenericDenseCache recv_cache_;
};
 
 
}  // namespace distributed
}  // namespace experimental
}  // namespace gko
 
#endif
#endif  // GKO_PUBLIC_CORE_DISTRIBUTED_ROW_SCATTERER_HPP_