lin_op.hpp

// SPDX-FileCopyrightText: 2017 - 2025 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 EnableAbstractPolymorphicObject<LinOp> {
public:
    LinOp* apply(ptr_param<const LinOp> b, ptr_param<LinOp> x)
    {
        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());
        return this;
    }
 
    const LinOp* 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());
        return this;
    }
 
    LinOp* apply(ptr_param<const LinOp> alpha, ptr_param<const LinOp> b,
                 ptr_param<const LinOp> beta, ptr_param<LinOp> x)
    {
        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());
        return this;
    }
 
    const LinOp* 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());
        return this;
    }
 
    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) {
            EnableAbstractPolymorphicObject<LinOp>::operator=(std::move(other));
            this->set_size(other.get_size());
            other.set_size({});
        }
        return *this;
    }
 
    LinOp(const LinOp&) = default;
 
    LinOp(LinOp&& other)
        : EnableAbstractPolymorphicObject<LinOp>(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>{})
        : EnableAbstractPolymorphicObject<LinOp>(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<LinOp, std::shared_ptr<const LinOp>>::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 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;
};
 
 
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 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:
    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 ConcreteLinOp, typename PolymorphicBase = LinOp>
class EnableLinOp
    : public EnablePolymorphicObject<ConcreteLinOp, PolymorphicBase>,
      public EnablePolymorphicAssignment<ConcreteLinOp> {
public:
    using EnablePolymorphicObject<ConcreteLinOp,
                                  PolymorphicBase>::EnablePolymorphicObject;
 
    const ConcreteLinOp* apply(ptr_param<const LinOp> b,
                               ptr_param<LinOp> x) const
    {
        PolymorphicBase::apply(b, x);
        return self();
    }
 
    ConcreteLinOp* apply(ptr_param<const LinOp> b, ptr_param<LinOp> x)
    {
        PolymorphicBase::apply(b, x);
        return self();
    }
 
    const ConcreteLinOp* apply(ptr_param<const LinOp> alpha,
                               ptr_param<const LinOp> b,
                               ptr_param<const LinOp> beta,
                               ptr_param<LinOp> x) const
    {
        PolymorphicBase::apply(alpha, b, beta, x);
        return self();
    }
 
    ConcreteLinOp* apply(ptr_param<const LinOp> alpha, ptr_param<const LinOp> b,
                         ptr_param<const LinOp> beta, ptr_param<LinOp> x)
    {
        PolymorphicBase::apply(alpha, b, beta, x);
        return self();
    }
 
protected:
    GKO_ENABLE_SELF(ConcreteLinOp);
};
 
 
template <typename ConcreteFactory, typename ConcreteLinOp,
          typename ParametersType, typename PolymorphicBase = LinOpFactory>
using EnableDefaultLinOpFactory =
    EnableDefaultFactory<ConcreteFactory, ConcreteLinOp, ParametersType,
                         PolymorphicBase>;
 
 
#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::EnablePolymorphicObject<_factory_name,           \
                                                    ::gko::LinOpFactory>;    \
        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")
 
 
}  // namespace gko
 
 
#endif  // GKO_PUBLIC_CORE_BASE_LIN_OP_HPP_