lin_op.hpp

// SPDX-FileCopyrightText: 2017 - 2026 The Ginkgo authors
//
// SPDX-License-Identifier: BSD-3-Clause
 
#ifndef GKO_PUBLIC_CORE_BASE_LIN_OP_HPP_
#define GKO_PUBLIC_CORE_BASE_LIN_OP_HPP_
 
 
#include <memory>
#include <type_traits>
#include <utility>
 
#include <ginkgo/core/base/abstract_factory.hpp>
#include <ginkgo/core/base/device_matrix_data.hpp>
#include <ginkgo/core/base/dim.hpp>
#include <ginkgo/core/base/exception_helpers.hpp>
#include <ginkgo/core/base/math.hpp>
#include <ginkgo/core/base/matrix_assembly_data.hpp>
#include <ginkgo/core/base/matrix_data.hpp>
#include <ginkgo/core/base/polymorphic_object.hpp>
#include <ginkgo/core/base/types.hpp>
#include <ginkgo/core/base/utils.hpp>
#include <ginkgo/core/log/logger.hpp>
 
 
namespace gko {
namespace matrix {
 
 
template <typename ValueType>
class Diagonal;
 
 
}  // namespace matrix
 
 
class LinOp : public PolymorphicObject {
public:
    void apply(ptr_param<const LinOp> b, ptr_param<LinOp> x) const
    {
        this->template log<log::Logger::linop_apply_started>(this, b.get(),
                                                             x.get());
        this->validate_application_parameters(b.get(), x.get());
        auto exec = this->get_executor();
        this->apply_impl(make_temporary_clone(exec, b).get(),
                         make_temporary_clone(exec, x).get());
        this->template log<log::Logger::linop_apply_completed>(this, b.get(),
                                                               x.get());
    }
 
 
    void apply(ptr_param<const LinOp> alpha, ptr_param<const LinOp> b,
               ptr_param<const LinOp> beta, ptr_param<LinOp> x) const
    {
        this->template log<log::Logger::linop_advanced_apply_started>(
            this, alpha.get(), b.get(), beta.get(), x.get());
        this->validate_application_parameters(alpha.get(), b.get(), beta.get(),
                                              x.get());
        auto exec = this->get_executor();
        this->apply_impl(make_temporary_clone(exec, alpha).get(),
                         make_temporary_clone(exec, b).get(),
                         make_temporary_clone(exec, beta).get(),
                         make_temporary_clone(exec, x).get());
        this->template log<log::Logger::linop_advanced_apply_completed>(
            this, alpha.get(), b.get(), beta.get(), x.get());
    }
 
    const dim<2>& get_size() const noexcept { return size_; }
 
    virtual bool apply_uses_initial_guess() const { return false; }
 
    LinOp& operator=(const LinOp&) = default;
 
    LinOp& operator=(LinOp&& other)
    {
        if (this != &other) {
            PolymorphicObject::operator=(std::move(other));
            this->set_size(other.get_size());
            other.set_size({});
        }
        return *this;
    }
 
    LinOp(const LinOp&) = default;
 
    LinOp(LinOp&& other)
        : PolymorphicObject(std::move(other)),
          size_{std::exchange(other.size_, dim<2>{})}
    {}
 
protected:
    explicit LinOp(std::shared_ptr<const Executor> exec,
                   const dim<2>& size = dim<2>{})
        : PolymorphicObject(exec), size_{size}
    {}
 
    void set_size(const dim<2>& value) noexcept { size_ = value; }
 
    virtual void apply_impl(const LinOp* b, LinOp* x) const = 0;
 
    virtual void apply_impl(const LinOp* alpha, const LinOp* b,
                            const LinOp* beta, LinOp* x) const = 0;
 
    void validate_application_parameters(const LinOp* b, const LinOp* x) const
    {
        GKO_ASSERT_CONFORMANT(this, b);
        GKO_ASSERT_EQUAL_ROWS(this, x);
        GKO_ASSERT_EQUAL_COLS(b, x);
    }
 
    void validate_application_parameters(const LinOp* alpha, const LinOp* b,
                                         const LinOp* beta,
                                         const LinOp* x) const
    {
        this->validate_application_parameters(b, x);
        GKO_ASSERT_EQUAL_DIMENSIONS(alpha, dim<2>(1, 1));
        GKO_ASSERT_EQUAL_DIMENSIONS(beta, dim<2>(1, 1));
    }
 
private:
    dim<2> size_{};
};
 
 
class LinOpFactory
    : public AbstractFactory<LinOp, std::shared_ptr<const LinOp>> {
public:
    using AbstractFactory::AbstractFactory;
 
    std::unique_ptr<LinOp> generate(std::shared_ptr<const LinOp> input) const
    {
        this->template log<log::Logger::linop_factory_generate_started>(
            this, input.get());
        const auto exec = this->get_executor();
        std::unique_ptr<LinOp> generated;
        if (input->get_executor() == exec) {
            generated = this->AbstractFactory::generate(input);
        } else {
            generated =
                this->AbstractFactory::generate(gko::clone(exec, input));
        }
        this->template log<log::Logger::linop_factory_generate_completed>(
            this, input.get(), generated.get());
        return generated;
    }
 
    class ReuseData {
    public:
        ReuseData();
 
        virtual ~ReuseData();
 
        ReuseData(const ReuseData&) = delete;
 
        ReuseData(ReuseData&&) = delete;
 
        ReuseData& operator=(const ReuseData&) = delete;
 
        ReuseData& operator=(ReuseData&&) = delete;
    };
 
    virtual std::unique_ptr<ReuseData> create_empty_reuse_data() const;
 
    std::unique_ptr<LinOp> generate_reuse(std::shared_ptr<const LinOp> input,
                                          ReuseData& reuse_data) const;
 
    virtual void check_reuse_consistent(const LinOp* input,
                                        const ReuseData& reuse_data) const;
 
protected:
    virtual std::unique_ptr<LinOp> generate_reuse_impl(
        std::shared_ptr<const LinOp> input, ReuseData& reuse_data) const;
};
 
 
class Transposable {
public:
    virtual ~Transposable() = default;
 
    virtual std::unique_ptr<LinOp> transpose() const = 0;
 
    virtual std::unique_ptr<LinOp> conj_transpose() const = 0;
};
 
 
template <typename IndexType>
class Permutable {
public:
    virtual ~Permutable() = default;
 
    virtual std::unique_ptr<LinOp> permute(
        const array<IndexType>* permutation_indices) const
    {
        return as<Permutable>(this->row_permute(permutation_indices))
            ->column_permute(permutation_indices);
    }
 
    virtual std::unique_ptr<LinOp> inverse_permute(
        const array<IndexType>* permutation_indices) const
    {
        return as<Permutable>(this->inverse_row_permute(permutation_indices))
            ->inverse_column_permute(permutation_indices);
    }
 
    virtual std::unique_ptr<LinOp> row_permute(
        const array<IndexType>* permutation_indices) const = 0;
 
    virtual std::unique_ptr<LinOp> column_permute(
        const array<IndexType>* permutation_indices) const = 0;
 
    virtual std::unique_ptr<LinOp> inverse_row_permute(
        const array<IndexType>* permutation_indices) const = 0;
 
    virtual std::unique_ptr<LinOp> inverse_column_permute(
        const array<IndexType>* permutation_indices) const = 0;
};
 
 
template <typename ValueType, typename IndexType>
class ReadableFromMatrixData {
public:
    using value_type = ValueType;
    using index_type = IndexType;
 
    virtual ~ReadableFromMatrixData() = default;
 
    virtual void read(const matrix_data<ValueType, IndexType>& data) = 0;
 
    void read(const matrix_assembly_data<ValueType, IndexType>& data)
    {
        this->read(data.get_ordered_data());
    }
 
    virtual void read(const device_matrix_data<ValueType, IndexType>& data)
    {
        this->read(data.copy_to_host());
    }
 
    virtual void read(device_matrix_data<ValueType, IndexType>&& data)
    {
        this->read(data.copy_to_host());
        data.empty_out();
    }
};
 
 
template <typename ValueType, typename IndexType>
class WritableToMatrixData {
public:
    using value_type = ValueType;
    using index_type = IndexType;
 
    virtual ~WritableToMatrixData() = default;
 
    virtual void write(matrix_data<ValueType, IndexType>& data) const = 0;
 
    virtual void write(device_matrix_data<ValueType, IndexType>& data) const
    {
        matrix_data<ValueType, IndexType> host_data;
        this->write(host_data);
        data = device_matrix_data<ValueType, IndexType>::create_from_host(
            data.get_executor(), host_data);
    }
};
 
 
class Preconditionable {
public:
    virtual ~Preconditionable() = default;
 
    virtual std::shared_ptr<const LinOp> get_preconditioner() const
    {
        return preconditioner_;
    }
 
    virtual void set_preconditioner(std::shared_ptr<const LinOp> new_precond)
    {
        preconditioner_ = new_precond;
    }
 
private:
    std::shared_ptr<const LinOp> preconditioner_{};
};
 
 
class DiagonalLinOpExtractable {
public:
    virtual ~DiagonalLinOpExtractable() = default;
 
    virtual std::unique_ptr<LinOp> extract_diagonal_linop() const = 0;
};
 
 
template <typename ValueType>
class DiagonalExtractable : public DiagonalLinOpExtractable {
public:
    using value_type = ValueType;
 
    virtual ~DiagonalExtractable() = default;
 
    std::unique_ptr<LinOp> extract_diagonal_linop() const override;
 
    virtual std::unique_ptr<matrix::Diagonal<ValueType>> extract_diagonal()
        const = 0;
};
 
 
class AbsoluteComputable {
public:
    virtual ~AbsoluteComputable() = default;
 
    virtual std::unique_ptr<LinOp> compute_absolute_linop() const = 0;
 
    virtual void compute_absolute_inplace() = 0;
};
 
 
template <typename AbsoluteLinOp>
class EnableAbsoluteComputation : public AbsoluteComputable {
public:
    using absolute_type = AbsoluteLinOp;
 
    virtual ~EnableAbsoluteComputation() = default;
 
    std::unique_ptr<LinOp> compute_absolute_linop() const override
    {
        return this->compute_absolute();
    }
 
    virtual std::unique_ptr<absolute_type> compute_absolute() const = 0;
};
 
 
class ScaledIdentityAddable {
public:
    virtual ~ScaledIdentityAddable() = default;
 
    void add_scaled_identity(ptr_param<const LinOp> const a,
                             ptr_param<const LinOp> const b)
    {
        GKO_ASSERT_IS_SCALAR(a);
        GKO_ASSERT_IS_SCALAR(b);
        auto ae = make_temporary_clone(as<LinOp>(this)->get_executor(), a);
        auto be = make_temporary_clone(as<LinOp>(this)->get_executor(), b);
        add_scaled_identity_impl(ae.get(), be.get());
    }
 
private:
    virtual void add_scaled_identity_impl(const LinOp* a, const LinOp* b) = 0;
};
 
 
template <typename ConcreteFactory, typename ConcreteLinOp,
          typename ParametersType, typename PolymorphicBase = LinOpFactory>
using EnableDefaultLinOpFactory =
    EnableDefaultFactory<ConcreteFactory, ConcreteLinOp, ParametersType,
                         PolymorphicBase>;
 
 
template <typename ConcreteFactory, typename ConcreteLinOp,
          typename ParametersType, typename ReuseDataType,
          typename PolymorphicBase = LinOpFactory>
class EnableDefaultReusableLinOpFactory
    : public EnableDefaultLinOpFactory<ConcreteFactory, ConcreteLinOp,
                                       ParametersType, PolymorphicBase> {
    using base_type =
        EnableDefaultLinOpFactory<ConcreteFactory, ConcreteLinOp,
                                  ParametersType, PolymorphicBase>;
 
public:
    std::unique_ptr<ConcreteLinOp> generate_reuse(
        std::shared_ptr<const LinOp> input,
        LinOpFactory::ReuseData& reuse_data) const
    {
        return std::unique_ptr<ConcreteLinOp>(static_cast<ConcreteLinOp*>(
            this->LinOpFactory::generate_reuse(std::move(input), reuse_data)
                .release()));
    }
 
    std::unique_ptr<LinOpFactory::ReuseData> create_empty_reuse_data()
        const override
    {
        return std::make_unique<ReuseDataType>();
    }
 
    void check_reuse_consistent(
        const LinOp* input,
        const LinOpFactory::ReuseData& reuse_data) const override
    {
        ConcreteLinOp::check_reuse_consistent(
            static_cast<const ConcreteFactory*>(this), input,
            *as<ReuseDataType>(&reuse_data));
    }
 
protected:
    explicit EnableDefaultReusableLinOpFactory(
        std::shared_ptr<const Executor> exec,
        const ParametersType& parameters = {})
        : base_type(std::move(exec), parameters)
    {}
 
    std::unique_ptr<LinOp> generate_reuse_impl(
        std::shared_ptr<const LinOp> input,
        LinOpFactory::ReuseData& reuse_data) const override
    {
        return std::unique_ptr<LinOp>(new ConcreteLinOp(
            static_cast<const ConcreteFactory*>(this), std::move(input),
            *as<ReuseDataType>(&reuse_data)));
    }
};
 
 
#define GKO_ENABLE_LIN_OP_FACTORY(_lin_op, _parameters_name, _factory_name)  \
public:                                                                      \
    const _parameters_name##_type& get_##_parameters_name() const            \
    {                                                                        \
        return _parameters_name##_;                                          \
    }                                                                        \
                                                                             \
    class _factory_name                                                      \
        : public ::gko::EnableDefaultLinOpFactory<_factory_name, _lin_op,    \
                                                  _parameters_name##_type> { \
        friend class ::gko::enable_parameters_type<_parameters_name##_type,  \
                                                   _factory_name>;           \
        explicit _factory_name(std::shared_ptr<const ::gko::Executor> exec)  \
            : ::gko::EnableDefaultLinOpFactory<_factory_name, _lin_op,       \
                                               _parameters_name##_type>(     \
                  std::move(exec))                                           \
        {}                                                                   \
        explicit _factory_name(std::shared_ptr<const ::gko::Executor> exec,  \
                               const _parameters_name##_type& parameters)    \
            : ::gko::EnableDefaultLinOpFactory<_factory_name, _lin_op,       \
                                               _parameters_name##_type>(     \
                  std::move(exec), parameters)                               \
        {}                                                                   \
    };                                                                       \
    friend ::gko::EnableDefaultLinOpFactory<_factory_name, _lin_op,          \
                                            _parameters_name##_type>;        \
                                                                             \
                                                                             \
private:                                                                     \
    _parameters_name##_type _parameters_name##_;                             \
                                                                             \
public:                                                                      \
    static_assert(true,                                                      \
                  "This assert is used to counter the false positive extra " \
                  "semi-colon warnings")
 
 
#define GKO_ENABLE_LIN_OP_FACTORY_WITH_REUSE(_lin_op, _parameters_name,        \
                                             _factory_name, _reuse_data_type)  \
public:                                                                        \
    const _parameters_name##_type& get_##_parameters_name() const              \
    {                                                                          \
        return _parameters_name##_;                                            \
    }                                                                          \
                                                                               \
    class _factory_name : public ::gko::EnableDefaultReusableLinOpFactory<     \
                              _factory_name, _lin_op, _parameters_name##_type, \
                              _reuse_data_type> {                              \
        friend class ::gko::enable_parameters_type<_parameters_name##_type,    \
                                                   _factory_name>;             \
        explicit _factory_name(std::shared_ptr<const ::gko::Executor> exec)    \
            : ::gko::EnableDefaultReusableLinOpFactory<                        \
                  _factory_name, _lin_op, _parameters_name##_type,             \
                  _reuse_data_type>(std::move(exec))                           \
        {}                                                                     \
        explicit _factory_name(std::shared_ptr<const ::gko::Executor> exec,    \
                               const _parameters_name##_type& parameters)      \
            : ::gko::EnableDefaultReusableLinOpFactory<                        \
                  _factory_name, _lin_op, _parameters_name##_type,             \
                  _reuse_data_type>(std::move(exec), parameters)               \
        {}                                                                     \
    };                                                                         \
    friend ::gko::EnableDefaultLinOpFactory<_factory_name, _lin_op,            \
                                            _parameters_name##_type>;          \
    friend ::gko::EnableDefaultReusableLinOpFactory<                           \
        _factory_name, _lin_op, _parameters_name##_type, _reuse_data_type>;    \
                                                                               \
                                                                               \
private:                                                                       \
    _parameters_name##_type _parameters_name##_;                               \
                                                                               \
public:                                                                        \
    static_assert(true,                                                        \
                  "This assert is used to counter the false positive extra "   \
                  "semi-colon warnings")
 
 
}  // namespace gko
 
 
#endif  // GKO_PUBLIC_CORE_BASE_LIN_OP_HPP_