block_operator.hpp

// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors
//
// SPDX-License-Identifier: BSD-3-Clause
 
#ifndef GINKGO_BLOCK_MATRIX_HPP
#define GINKGO_BLOCK_MATRIX_HPP
 
#include <ginkgo/config.hpp>
#include <ginkgo/core/base/dense_cache.hpp>
#include <ginkgo/core/base/lin_op.hpp>
#include <ginkgo/core/base/polymorphic_object.hpp>
 
 
namespace gko {
namespace detail {
 
 
struct value_span {
    constexpr value_span(size_type point) noexcept
        : value_span{point, point + 1}
    {}
 
    constexpr value_span(size_type begin, size_type end) noexcept
        : begin{begin}, end{end}
    {}
 
    constexpr operator span() const { return {begin, end}; }
 
    constexpr value_span(const span& s) noexcept : value_span(s.begin, s.end) {}
 
    constexpr bool is_valid() const { return begin <= end; }
 
    constexpr size_type length() const { return end - begin; }
 
    size_type begin;
    size_type end;
};
 
 
}  // namespace detail
 
 
class BlockOperator final : public EnableLinOp<BlockOperator> {
    friend class EnablePolymorphicObject<BlockOperator, LinOp>;
 
public:
    dim<2> get_block_size() const { return block_size_; }
 
    const LinOp* block_at(size_type i, size_type j) const
    {
        GKO_ENSURE_IN_DIMENSION_BOUNDS(i, j, block_size_);
        return blocks_[i * block_size_[1] + j].get();
    }
 
    static std::unique_ptr<BlockOperator> create(
        std::shared_ptr<const Executor> exec);
 
    static std::unique_ptr<BlockOperator> create(
        std::shared_ptr<const Executor> exec,
        std::vector<std::vector<std::shared_ptr<const LinOp>>> blocks);
 
    BlockOperator(const BlockOperator& other);
 
    BlockOperator(BlockOperator&& other) noexcept;
 
    BlockOperator& operator=(const BlockOperator& other);
 
    BlockOperator& operator=(BlockOperator&& other);
 
private:
    explicit BlockOperator(std::shared_ptr<const Executor> exec);
 
    BlockOperator(
        std::shared_ptr<const Executor> exec,
        std::vector<std::vector<std::shared_ptr<const LinOp>>> blocks);
 
    void apply_impl(const LinOp* b, LinOp* x) const override;
 
    void apply_impl(const LinOp* alpha, const LinOp* b, const LinOp* beta,
                    LinOp* x) const override;
 
    dim<2> block_size_;
    std::vector<detail::value_span> row_spans_;
    std::vector<detail::value_span> col_spans_;
    std::vector<std::shared_ptr<const LinOp>> blocks_;
    detail::DenseCache<default_precision> one_;
};
 
 
}  // namespace gko
 
#endif  // GINKGO_BLOCK_MATRIX_HPP