ic.hpp

// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors
//
// SPDX-License-Identifier: BSD-3-Clause
 
#ifndef GKO_PUBLIC_CORE_PRECONDITIONER_IC_HPP_
#define GKO_PUBLIC_CORE_PRECONDITIONER_IC_HPP_
 
 
#include <memory>
#include <type_traits>
 
 
#include <ginkgo/core/base/abstract_factory.hpp>
#include <ginkgo/core/base/composition.hpp>
#include <ginkgo/core/base/exception.hpp>
#include <ginkgo/core/base/exception_helpers.hpp>
#include <ginkgo/core/base/lin_op.hpp>
#include <ginkgo/core/base/precision_dispatch.hpp>
#include <ginkgo/core/base/std_extensions.hpp>
#include <ginkgo/core/config/config.hpp>
#include <ginkgo/core/config/registry.hpp>
#include <ginkgo/core/factorization/par_ic.hpp>
#include <ginkgo/core/matrix/dense.hpp>
#include <ginkgo/core/preconditioner/isai.hpp>
#include <ginkgo/core/preconditioner/utils.hpp>
#include <ginkgo/core/solver/gmres.hpp>
#include <ginkgo/core/solver/ir.hpp>
#include <ginkgo/core/solver/solver_traits.hpp>
#include <ginkgo/core/solver/triangular.hpp>
#include <ginkgo/core/stop/combined.hpp>
#include <ginkgo/core/stop/iteration.hpp>
#include <ginkgo/core/stop/residual_norm.hpp>
 
 
namespace gko {
namespace preconditioner {
namespace detail {
 
 
template <typename Type>
constexpr bool support_ic_parse =
    is_instantiation_of<Type, solver::LowerTrs>::value ||
    is_instantiation_of<Type, solver::Ir>::value ||
    is_instantiation_of<Type, solver::Gmres>::value ||
    is_instantiation_of<Type, preconditioner::LowerIsai>::value;
 
 
template <
    typename Ic,
    std::enable_if_t<!support_ic_parse<typename Ic::l_solver_type>>* = nullptr>
typename Ic::parameters_type ic_parse(
    const config::pnode& config, const config::registry& context,
    const config::type_descriptor& td_for_child)
{
    GKO_INVALID_STATE(
        "preconditioner::Ic only supports limited type for parse.");
}
 
template <
    typename Ic,
    std::enable_if_t<support_ic_parse<typename Ic::l_solver_type>>* = nullptr>
typename Ic::parameters_type ic_parse(
    const config::pnode& config, const config::registry& context,
    const config::type_descriptor& td_for_child);
 
 
}  // namespace detail
 
template <typename LSolverType = solver::LowerTrs<>, typename IndexType = int32>
class Ic : public EnableLinOp<Ic<LSolverType, IndexType>>, public Transposable {
    friend class EnableLinOp<Ic>;
    friend class EnablePolymorphicObject<Ic, LinOp>;
 
public:
    static_assert(
        std::is_same<typename LSolverType::transposed_type::transposed_type,
                     LSolverType>::value,
        "LSolverType::transposed_type must be symmetric");
    using value_type = typename LSolverType::value_type;
    using l_solver_type = LSolverType;
    using lh_solver_type = typename LSolverType::transposed_type;
    using index_type = IndexType;
    using transposed_type = Ic<LSolverType, IndexType>;
 
    class Factory;
 
    struct parameters_type
        : public enable_parameters_type<parameters_type, Factory> {
        std::shared_ptr<const typename l_solver_type::Factory>
            l_solver_factory{};
 
        std::shared_ptr<const LinOpFactory> factorization_factory{};
 
        GKO_DEPRECATED("use with_l_solver instead")
        parameters_type& with_l_solver_factory(
            deferred_factory_parameter<const typename l_solver_type::Factory>
                solver)
        {
            return with_l_solver(std::move(solver));
        }
 
        parameters_type& with_l_solver(
            deferred_factory_parameter<const typename l_solver_type::Factory>
                solver)
        {
            this->l_solver_generator = std::move(solver);
            this->deferred_factories["l_solver"] = [](const auto& exec,
                                                      auto& params) {
                if (!params.l_solver_generator.is_empty()) {
                    params.l_solver_factory =
                        params.l_solver_generator.on(exec);
                }
            };
            return *this;
        }
 
        GKO_DEPRECATED("use with_factorization instead")
        parameters_type& with_factorization_factory(
            deferred_factory_parameter<const LinOpFactory> factorization)
        {
            return with_factorization(std::move(factorization));
        }
 
        parameters_type& with_factorization(
            deferred_factory_parameter<const LinOpFactory> factorization)
        {
            this->factorization_generator = std::move(factorization);
            this->deferred_factories["factorization"] = [](const auto& exec,
                                                           auto& params) {
                if (!params.factorization_generator.is_empty()) {
                    params.factorization_factory =
                        params.factorization_generator.on(exec);
                }
            };
            return *this;
        }
 
    private:
        deferred_factory_parameter<const typename l_solver_type::Factory>
            l_solver_generator;
 
        deferred_factory_parameter<const LinOpFactory> factorization_generator;
    };
 
    GKO_ENABLE_LIN_OP_FACTORY(Ic, parameters, Factory);
    GKO_ENABLE_BUILD_METHOD(Factory);
 
    static parameters_type parse(
        const config::pnode& config, const config::registry& context,
        const config::type_descriptor& td_for_child =
            config::make_type_descriptor<value_type, index_type>())
    {
        return detail::ic_parse<Ic>(config, context, td_for_child);
    }
 
    std::shared_ptr<const l_solver_type> get_l_solver() const
    {
        return l_solver_;
    }
 
    std::shared_ptr<const lh_solver_type> get_lh_solver() const
    {
        return lh_solver_;
    }
 
    std::unique_ptr<LinOp> transpose() const override
    {
        std::unique_ptr<transposed_type> transposed{
            new transposed_type{this->get_executor()}};
        transposed->set_size(gko::transpose(this->get_size()));
        transposed->l_solver_ =
            share(as<typename lh_solver_type::transposed_type>(
                this->get_lh_solver()->transpose()));
        transposed->lh_solver_ =
            share(as<typename l_solver_type::transposed_type>(
                this->get_l_solver()->transpose()));
 
        return std::move(transposed);
    }
 
    std::unique_ptr<LinOp> conj_transpose() const override
    {
        std::unique_ptr<transposed_type> transposed{
            new transposed_type{this->get_executor()}};
        transposed->set_size(gko::transpose(this->get_size()));
        transposed->l_solver_ =
            share(as<typename lh_solver_type::transposed_type>(
                this->get_lh_solver()->conj_transpose()));
        transposed->lh_solver_ =
            share(as<typename l_solver_type::transposed_type>(
                this->get_l_solver()->conj_transpose()));
 
        return std::move(transposed);
    }
 
    Ic& operator=(const Ic& other)
    {
        if (&other != this) {
            EnableLinOp<Ic>::operator=(other);
            auto exec = this->get_executor();
            l_solver_ = other.l_solver_;
            lh_solver_ = other.lh_solver_;
            parameters_ = other.parameters_;
            if (other.get_executor() != exec) {
                l_solver_ = gko::clone(exec, l_solver_);
                lh_solver_ = gko::clone(exec, lh_solver_);
            }
        }
        return *this;
    }
 
    Ic& operator=(Ic&& other)
    {
        if (&other != this) {
            EnableLinOp<Ic>::operator=(other);
            auto exec = this->get_executor();
            l_solver_ = std::move(other.l_solver_);
            lh_solver_ = std::move(other.lh_solver_);
            parameters_ = std::exchange(other.parameters_, parameters_type{});
            if (other.get_executor() != exec) {
                l_solver_ = gko::clone(exec, l_solver_);
                lh_solver_ = gko::clone(exec, lh_solver_);
            }
        }
        return *this;
    }
 
    Ic(const Ic& other) : Ic{other.get_executor()} { *this = other; }
 
    Ic(Ic&& other) : Ic{other.get_executor()} { *this = std::move(other); }
 
protected:
    void apply_impl(const LinOp* b, LinOp* x) const override
    {
        // take care of real-to-complex apply
        precision_dispatch_real_complex<value_type>(
            [&](auto dense_b, auto dense_x) {
                this->set_cache_to(dense_b);
                l_solver_->apply(dense_b, cache_.intermediate);
                if (lh_solver_->apply_uses_initial_guess()) {
                    dense_x->copy_from(cache_.intermediate);
                }
                lh_solver_->apply(cache_.intermediate, dense_x);
            },
            b, x);
    }
 
    void apply_impl(const LinOp* alpha, const LinOp* b, const LinOp* beta,
                    LinOp* x) const override
    {
        precision_dispatch_real_complex<value_type>(
            [&](auto dense_alpha, auto dense_b, auto dense_beta, auto dense_x) {
                this->set_cache_to(dense_b);
                l_solver_->apply(dense_b, cache_.intermediate);
                lh_solver_->apply(dense_alpha, cache_.intermediate, dense_beta,
                                  dense_x);
            },
            alpha, b, beta, x);
    }
 
    explicit Ic(std::shared_ptr<const Executor> exec)
        : EnableLinOp<Ic>(std::move(exec))
    {}
 
    explicit Ic(const Factory* factory, std::shared_ptr<const LinOp> lin_op)
        : EnableLinOp<Ic>(factory->get_executor(), lin_op->get_size()),
          parameters_{factory->get_parameters()}
    {
        auto comp =
            std::dynamic_pointer_cast<const Composition<value_type>>(lin_op);
        std::shared_ptr<const LinOp> l_factor;
 
        // build factorization if we weren't passed a composition
        if (!comp) {
            auto exec = lin_op->get_executor();
            if (!parameters_.factorization_factory) {
                parameters_.factorization_factory =
                    factorization::ParIc<value_type, index_type>::build()
                        .with_both_factors(false)
                        .on(exec);
            }
            auto fact = std::shared_ptr<const LinOp>(
                parameters_.factorization_factory->generate(lin_op));
            // ensure that the result is a composition
            comp =
                std::dynamic_pointer_cast<const Composition<value_type>>(fact);
            if (!comp) {
                GKO_NOT_SUPPORTED(comp);
            }
        }
        // comp must contain one or two factors
        if (comp->get_operators().size() > 2 || comp->get_operators().empty()) {
            GKO_NOT_SUPPORTED(comp);
        }
        l_factor = comp->get_operators()[0];
        GKO_ASSERT_IS_SQUARE_MATRIX(l_factor);
 
        auto exec = this->get_executor();
 
        // If no factories are provided, generate default ones
        if (!parameters_.l_solver_factory) {
            l_solver_ = generate_default_solver<l_solver_type>(exec, l_factor);
            // If comp contains both factors: use the transposed factor to avoid
            // transposing twice
            if (comp->get_operators().size() == 2) {
                auto lh_factor = comp->get_operators()[1];
                GKO_ASSERT_EQUAL_DIMENSIONS(l_factor, lh_factor);
                lh_solver_ = as<lh_solver_type>(l_solver_->conj_transpose());
            } else {
                lh_solver_ = as<lh_solver_type>(l_solver_->conj_transpose());
            }
        } else {
            l_solver_ = parameters_.l_solver_factory->generate(l_factor);
            lh_solver_ = as<lh_solver_type>(l_solver_->conj_transpose());
        }
    }
 
    void set_cache_to(const LinOp* b) const
    {
        if (cache_.intermediate == nullptr) {
            cache_.intermediate =
                matrix::Dense<value_type>::create(this->get_executor());
        }
        // Use b as the initial guess for the first triangular solve
        cache_.intermediate->copy_from(b);
    }
 
 
    template <typename SolverType>
    static std::enable_if_t<solver::has_with_criteria<SolverType>::value,
                            std::unique_ptr<SolverType>>
    generate_default_solver(const std::shared_ptr<const Executor>& exec,
                            const std::shared_ptr<const LinOp>& mtx)
    {
        constexpr gko::remove_complex<value_type> default_reduce_residual{1e-4};
        const unsigned int default_max_iters{
            static_cast<unsigned int>(mtx->get_size()[0])};
 
        return SolverType::build()
            .with_criteria(
                gko::stop::Iteration::build().with_max_iters(default_max_iters),
                gko::stop::ResidualNorm<value_type>::build()
                    .with_reduction_factor(default_reduce_residual))
            .on(exec)
            ->generate(mtx);
    }
 
    template <typename SolverType>
    static std::enable_if_t<!solver::has_with_criteria<SolverType>::value,
                            std::unique_ptr<SolverType>>
    generate_default_solver(const std::shared_ptr<const Executor>& exec,
                            const std::shared_ptr<const LinOp>& mtx)
    {
        return SolverType::build().on(exec)->generate(mtx);
    }
 
private:
    std::shared_ptr<const l_solver_type> l_solver_{};
    std::shared_ptr<const lh_solver_type> lh_solver_{};
    mutable struct cache_struct {
        cache_struct() = default;
        ~cache_struct() = default;
        cache_struct(const cache_struct&) {}
        cache_struct(cache_struct&&) {}
        cache_struct& operator=(const cache_struct&) { return *this; }
        cache_struct& operator=(cache_struct&&) { return *this; }
        std::unique_ptr<LinOp> intermediate{};
    } cache_;
};
 
 
}  // namespace preconditioner
}  // namespace gko
 
 
#endif  // GKO_PUBLIC_CORE_PRECONDITIONER_IC_HPP_