row_gatherer.hpp

// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors
//
// SPDX-License-Identifier: BSD-3-Clause
 
#ifndef GKO_PUBLIC_CORE_MATRIX_ROW_GATHERER_HPP_
#define GKO_PUBLIC_CORE_MATRIX_ROW_GATHERER_HPP_
 
 
#include <algorithm>
#include <memory>
#include <numeric>
#include <vector>
 
 
#include <ginkgo/core/base/array.hpp>
#include <ginkgo/core/base/exception.hpp>
#include <ginkgo/core/base/exception_helpers.hpp>
#include <ginkgo/core/base/executor.hpp>
#include <ginkgo/core/base/lin_op.hpp>
#include <ginkgo/core/base/types.hpp>
#include <ginkgo/core/base/utils.hpp>
 
 
namespace gko {
namespace matrix {
 
 
template <typename IndexType = int32>
class RowGatherer : public EnableLinOp<RowGatherer<IndexType>> {
    friend class EnablePolymorphicObject<RowGatherer, LinOp>;
 
public:
    using index_type = IndexType;
 
    index_type* get_row_idxs() noexcept { return row_idxs_.get_data(); }
 
    const index_type* get_const_row_idxs() const noexcept
    {
        return row_idxs_.get_const_data();
    }
 
    static std::unique_ptr<RowGatherer> create(
        std::shared_ptr<const Executor> exec, const dim<2>& size = {});
 
    static std::unique_ptr<RowGatherer> create(
        std::shared_ptr<const Executor> exec, const dim<2>& size,
        array<index_type> row_idxs);
 
    static std::unique_ptr<const RowGatherer> create_const(
        std::shared_ptr<const Executor> exec, const dim<2>& size,
        gko::detail::const_array_view<IndexType>&& row_idxs);
 
protected:
    RowGatherer(std::shared_ptr<const Executor> exec, const dim<2>& size = {});
 
    RowGatherer(std::shared_ptr<const Executor> exec, const dim<2>& size,
                array<index_type> row_idxs);
 
    void apply_impl(const LinOp* in, LinOp* out) const override;
 
    void apply_impl(const LinOp* alpha, const LinOp* in, const LinOp* beta,
                    LinOp* out) const override;
 
private:
    gko::array<index_type> row_idxs_;
};
 
 
}  // namespace matrix
}  // namespace gko
 
 
#endif  // GKO_PUBLIC_CORE_MATRIX_ROW_GATHERER_HPP_