combination.hpp

// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors
//
// SPDX-License-Identifier: BSD-3-Clause
 
#ifndef GKO_PUBLIC_CORE_BASE_COMBINATION_HPP_
#define GKO_PUBLIC_CORE_BASE_COMBINATION_HPP_
 
 
#include <type_traits>
#include <vector>
 
#include <ginkgo/core/base/lin_op.hpp>
 
 
namespace gko {
 
 
template <typename ValueType = default_precision>
class Combination : public EnableLinOp<Combination<ValueType>>,
                    public EnableCreateMethod<Combination<ValueType>>,
                    public Transposable {
    friend class EnablePolymorphicObject<Combination, LinOp>;
    friend class EnableCreateMethod<Combination>;
 
public:
    using value_type = ValueType;
    using transposed_type = Combination<ValueType>;
 
    const std::vector<std::shared_ptr<const LinOp>>& get_coefficients()
        const noexcept
    {
        return coefficients_;
    }
 
    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;
 
    Combination& operator=(const Combination&);
 
    Combination& operator=(Combination&&);
 
    Combination(const Combination&);
 
    Combination(Combination&&);
 
protected:
    void add_operators() {}
 
    template <typename... Rest>
    void add_operators(std::shared_ptr<const LinOp> coef,
                       std::shared_ptr<const LinOp> oper, Rest&&... rest)
    {
        GKO_ASSERT_EQUAL_DIMENSIONS(coef, dim<2>(1, 1));
        GKO_ASSERT_EQUAL_DIMENSIONS(oper, this->get_size());
        auto exec = this->get_executor();
        coefficients_.push_back(std::move(coef));
        operators_.push_back(std::move(oper));
        if (coefficients_.back()->get_executor() != exec) {
            coefficients_.back() = gko::clone(exec, coefficients_.back());
        }
        if (operators_.back()->get_executor() != exec) {
            operators_.back() = gko::clone(exec, operators_.back());
        }
        add_operators(std::forward<Rest>(rest)...);
    }
 
    explicit Combination(std::shared_ptr<const Executor> exec)
        : EnableLinOp<Combination>(exec)
    {}
 
    template <
        typename CoefficientIterator, typename OperatorIterator,
        typename = std::void_t<
            typename std::iterator_traits<
                CoefficientIterator>::iterator_category,
            typename std::iterator_traits<OperatorIterator>::iterator_category>>
    explicit Combination(CoefficientIterator coefficient_begin,
                         CoefficientIterator coefficient_end,
                         OperatorIterator operator_begin,
                         OperatorIterator operator_end)
        : EnableLinOp<Combination>([&] {
              if (operator_begin == operator_end) {
                  throw OutOfBoundsError(__FILE__, __LINE__, 1, 0);
              }
              return (*operator_begin)->get_executor();
          }())
    {
        GKO_ASSERT_EQ(std::distance(coefficient_begin, coefficient_end),
                      std::distance(operator_begin, operator_end));
        this->set_size((*operator_begin)->get_size());
        auto coefficient_it = coefficient_begin;
        for (auto operator_it = operator_begin; operator_it != operator_end;
             ++operator_it) {
            add_operators(*coefficient_it, *operator_it);
            ++coefficient_it;
        }
    }
 
    template <typename... Rest>
    explicit Combination(std::shared_ptr<const LinOp> coef,
                         std::shared_ptr<const LinOp> oper, Rest&&... rest)
        : Combination(oper->get_executor())
    {
        this->set_size(oper->get_size());
        add_operators(std::move(coef), 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>> coefficients_;
    std::vector<std::shared_ptr<const LinOp>> operators_;
 
    // TODO: solve race conditions when multithreading
    mutable struct cache_struct {
        cache_struct() = default;
        ~cache_struct() = default;
        cache_struct(const cache_struct& other) {}
        cache_struct& operator=(const cache_struct& other) { return *this; }
 
        std::unique_ptr<LinOp> zero;
        std::unique_ptr<LinOp> one;
        std::unique_ptr<LinOp> intermediate_x;
    } cache_;
};
 
 
}  // namespace gko
 
 
#endif  // GKO_PUBLIC_CORE_BASE_COMBINATION_HPP_