composition.hpp

// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors
//
// SPDX-License-Identifier: BSD-3-Clause
 
#ifndef GKO_PUBLIC_CORE_BASE_COMPOSITION_HPP_
#define GKO_PUBLIC_CORE_BASE_COMPOSITION_HPP_
 
 
#include <vector>
 
 
#include <ginkgo/core/base/executor.hpp>
#include <ginkgo/core/base/lin_op.hpp>
 
 
namespace gko {
 
 
template <typename ValueType = default_precision>
class Composition : public EnableLinOp<Composition<ValueType>>,
                    public EnableCreateMethod<Composition<ValueType>>,
                    public Transposable {
    friend class EnablePolymorphicObject<Composition, LinOp>;
    friend class EnableCreateMethod<Composition>;
 
public:
    using value_type = ValueType;
    using transposed_type = Composition<ValueType>;
 
    const std::vector<std::shared_ptr<const LinOp>>& get_operators()
        const noexcept
    {
        return operators_;
    }
 
    std::unique_ptr<LinOp> transpose() const override;
 
    std::unique_ptr<LinOp> conj_transpose() const override;
 
    Composition& operator=(const Composition&);
 
    Composition& operator=(Composition&&);
 
    Composition(const Composition&);
 
    Composition(Composition&&);
 
protected:
    void add_operators() {}
 
    template <typename... Rest>
    void add_operators(std::shared_ptr<const LinOp> oper, Rest&&... rest)
    {
        if (!operators_.empty()) {
            GKO_ASSERT_CONFORMANT(this, oper);
        }
        auto exec = this->get_executor();
        operators_.push_back(std::move(oper));
        if (operators_.back()->get_executor() != exec) {
            operators_.back() = gko::clone(exec, operators_.back());
        }
        this->set_size(dim<2>{operators_.front()->get_size()[0],
                              operators_.back()->get_size()[1]});
        add_operators(std::forward<Rest>(rest)...);
    }
 
    explicit Composition(std::shared_ptr<const Executor> exec)
        : EnableLinOp<Composition>(exec), storage_{exec}
    {}
 
    template <typename Iterator,
              typename = xstd::void_t<
                  typename std::iterator_traits<Iterator>::iterator_category>>
    explicit Composition(Iterator begin, Iterator end)
        : EnableLinOp<Composition>([&] {
              if (begin == end) {
                  throw OutOfBoundsError(__FILE__, __LINE__, 1, 0);
              }
              return (*begin)->get_executor();
          }()),
          storage_{this->get_executor()}
    {
        for (auto it = begin; it != end; ++it) {
            add_operators(*it);
        }
    }
 
    template <typename... Rest>
    explicit Composition(std::shared_ptr<const LinOp> oper, Rest&&... rest)
        : Composition(oper->get_executor())
    {
        add_operators(std::move(oper), std::forward<Rest>(rest)...);
    }
 
    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;
 
private:
    std::vector<std::shared_ptr<const LinOp>> operators_;
    mutable array<ValueType> storage_;
};
 
 
template <typename ValueType = default_precision>
class UseComposition {
public:
    using value_type = ValueType;
    std::shared_ptr<Composition<ValueType>> get_composition() const
    {
        return composition_;
    }
 
    std::shared_ptr<const LinOp> get_operator_at(size_type index) const
    {
        if (composition_ == nullptr) {
            return nullptr;
        } else {
            return composition_->get_operators().at(index);
        }
    }
 
protected:
    template <typename... LinOp>
    void set_composition(LinOp&&... linop)
    {
        composition_ =
            Composition<ValueType>::create(std::forward<LinOp>(linop)...);
    }
 
private:
    std::shared_ptr<Composition<ValueType>> composition_;
};
 
 
}  // namespace gko
 
 
#endif  // GKO_PUBLIC_CORE_BASE_COMPOSITION_HPP_