ilu.hpp

// SPDX-FileCopyrightText: 2017 - 2025 The Ginkgo authors
//
// SPDX-License-Identifier: BSD-3-Clause
 
#ifndef GKO_PUBLIC_CORE_PRECONDITIONER_ILU_HPP_
#define GKO_PUBLIC_CORE_PRECONDITIONER_ILU_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/type_traits.hpp>
#include <ginkgo/core/config/config.hpp>
#include <ginkgo/core/config/registry.hpp>
#include <ginkgo/core/factorization/par_ilu.hpp>
#include <ginkgo/core/matrix/dense.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_ilu_parse =
    std::is_same_v<typename Type::l_solver_type, LinOp>&&
        std::is_same_v<typename Type::u_solver_type, LinOp>;
 
 
template <typename Ilu, std::enable_if_t<!support_ilu_parse<Ilu>>* = nullptr>
typename Ilu::parameters_type ilu_parse(
    const config::pnode& config, const config::registry& context,
    const config::type_descriptor& td_for_child)
{
    GKO_INVALID_STATE(
        "preconditioner::Ilu only supports limited type for parse.");
}
 
template <typename Ilu, std::enable_if_t<support_ilu_parse<Ilu>>* = nullptr>
typename Ilu::parameters_type ilu_parse(
    const config::pnode& config, const config::registry& context,
    const config::type_descriptor& td_for_child);
 
}  // namespace detail
 
 
template <typename LSolverTypeOrValueType = solver::LowerTrs<>,
          typename USolverTypeOrValueType =
              gko::detail::transposed_type<LSolverTypeOrValueType>,
          bool ReverseApply = false, typename IndexType = int32>
class Ilu
    : public EnableLinOp<Ilu<LSolverTypeOrValueType, USolverTypeOrValueType,
                             ReverseApply, IndexType>>,
      public Transposable {
    friend class EnableLinOp<Ilu>;
    friend class EnablePolymorphicObject<Ilu, LinOp>;
 
public:
    static_assert(
        std::is_same_v<gko::detail::get_value_type<LSolverTypeOrValueType>,
                       gko::detail::get_value_type<USolverTypeOrValueType>>,
        "Both the L- and the U-solver must use the same `value_type`!");
    using value_type = gko::detail::get_value_type<LSolverTypeOrValueType>;
    using l_solver_type =
        std::conditional_t<gko::detail::is_ginkgo_linop<LSolverTypeOrValueType>,
                           LSolverTypeOrValueType, LinOp>;
    using u_solver_type =
        std::conditional_t<gko::detail::is_ginkgo_linop<USolverTypeOrValueType>,
                           USolverTypeOrValueType, LinOp>;
    static constexpr bool performs_reverse_apply = ReverseApply;
    using index_type = IndexType;
    using transposed_type =
        Ilu<gko::detail::transposed_type<USolverTypeOrValueType>,
            gko::detail::transposed_type<LSolverTypeOrValueType>, ReverseApply,
            IndexType>;
 
    class Factory;
 
    struct parameters_type
        : public enable_parameters_type<parameters_type, Factory> {
        std::shared_ptr<const gko::detail::factory_type<l_solver_type>>
            l_solver_factory{};
 
        std::shared_ptr<const gko::detail::factory_type<u_solver_type>>
            u_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 gko::detail::factory_type<l_solver_type>>
                solver)
        {
            return with_l_solver(std::move(solver));
        }
 
        parameters_type& with_l_solver(
            deferred_factory_parameter<
                const gko::detail::factory_type<l_solver_type>>
                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_u_solver instead")
        parameters_type& with_u_solver_factory(
            deferred_factory_parameter<
                const gko::detail::factory_type<u_solver_type>>
                solver)
        {
            return with_u_solver(std::move(solver));
        }
 
        parameters_type& with_u_solver(
            deferred_factory_parameter<
                const gko::detail::factory_type<u_solver_type>>
                solver)
        {
            this->u_solver_generator = std::move(solver);
            this->deferred_factories["u_solver"] = [](const auto& exec,
                                                      auto& params) {
                if (!params.u_solver_generator.is_empty()) {
                    params.u_solver_factory =
                        params.u_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 gko::detail::factory_type<l_solver_type>>
            l_solver_generator;
 
        deferred_factory_parameter<
            const gko::detail::factory_type<u_solver_type>>
            u_solver_generator;
 
        deferred_factory_parameter<const LinOpFactory> factorization_generator;
    };
 
    GKO_ENABLE_LIN_OP_FACTORY(Ilu, 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>())
    {
        // parse is not templated, so we can only use SFINAE later
        return detail::ilu_parse<Ilu>(config, context, td_for_child);
    }
 
    std::shared_ptr<const l_solver_type> get_l_solver() const
    {
        return l_solver_;
    }
 
    std::shared_ptr<const u_solver_type> get_u_solver() const
    {
        return u_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<gko::detail::transposed_type<u_solver_type>>(
                as<Transposable>(this->get_u_solver())->transpose()));
        transposed->u_solver_ =
            share(as<gko::detail::transposed_type<l_solver_type>>(
                as<Transposable>(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<gko::detail::transposed_type<u_solver_type>>(
                as<Transposable>(this->get_u_solver())->conj_transpose()));
        transposed->u_solver_ =
            share(as<gko::detail::transposed_type<l_solver_type>>(
                as<Transposable>(this->get_l_solver())->conj_transpose()));
 
        return std::move(transposed);
    }
 
    Ilu& operator=(const Ilu& other)
    {
        if (&other != this) {
            EnableLinOp<Ilu>::operator=(other);
            auto exec = this->get_executor();
            l_solver_ = other.l_solver_;
            u_solver_ = other.u_solver_;
            parameters_ = other.parameters_;
            if (other.get_executor() != exec) {
                l_solver_ = gko::clone(exec, l_solver_);
                u_solver_ = gko::clone(exec, u_solver_);
            }
        }
        return *this;
    }
 
    Ilu& operator=(Ilu&& other)
    {
        if (&other != this) {
            EnableLinOp<Ilu>::operator=(other);
            auto exec = this->get_executor();
            l_solver_ = std::move(other.l_solver_);
            u_solver_ = std::move(other.u_solver_);
            parameters_ = std::exchange(other.parameters_, parameters_type{});
            if (other.get_executor() != exec) {
                l_solver_ = gko::clone(exec, l_solver_);
                u_solver_ = gko::clone(exec, u_solver_);
            }
        }
        return *this;
    }
 
    Ilu(const Ilu& other) : Ilu{other.get_executor()} { *this = other; }
 
    Ilu(Ilu&& other) : Ilu{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);
                if (!ReverseApply) {
                    l_solver_->apply(dense_b, cache_.intermediate);
                    if (u_solver_->apply_uses_initial_guess()) {
                        dense_x->copy_from(cache_.intermediate);
                    }
                    u_solver_->apply(cache_.intermediate, dense_x);
                } else {
                    u_solver_->apply(dense_b, cache_.intermediate);
                    if (l_solver_->apply_uses_initial_guess()) {
                        dense_x->copy_from(cache_.intermediate);
                    }
                    l_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);
                if (!ReverseApply) {
                    l_solver_->apply(dense_b, cache_.intermediate);
                    u_solver_->apply(dense_alpha, cache_.intermediate,
                                     dense_beta, dense_x);
                } else {
                    u_solver_->apply(dense_b, cache_.intermediate);
                    l_solver_->apply(dense_alpha, cache_.intermediate,
                                     dense_beta, dense_x);
                }
            },
            alpha, b, beta, x);
    }
 
    explicit Ilu(std::shared_ptr<const Executor> exec)
        : EnableLinOp<Ilu>(std::move(exec))
    {}
 
    explicit Ilu(const Factory* factory, std::shared_ptr<const LinOp> lin_op)
        : EnableLinOp<Ilu>(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;
        std::shared_ptr<const LinOp> u_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::ParIlu<value_type, index_type>::build().on(
                        exec);
            }
            auto fact = std::shared_ptr<const LinOp>(
                parameters_.factorization_factory->generate(lin_op));
            // ensure that the result is a composition
            comp = as<const Composition<value_type>>(fact);
        }
        if (comp->get_operators().size() == 2) {
            l_factor = comp->get_operators()[0];
            u_factor = comp->get_operators()[1];
        } else {
            GKO_NOT_SUPPORTED(comp);
        }
        GKO_ASSERT_EQUAL_DIMENSIONS(l_factor, u_factor);
 
        auto exec = this->get_executor();
 
        // If no factories are provided, generate default ones
        if (!parameters_.l_solver_factory) {
            // when l_solver_type is LinOp, use LowerTrs as the default one
            l_solver_ = generate_default_solver<std::conditional_t<
                std::is_same_v<l_solver_type, LinOp>,
                solver::LowerTrs<value_type, index_type>, l_solver_type>>(
                exec, l_factor);
        } else {
            l_solver_ = parameters_.l_solver_factory->generate(l_factor);
        }
        if (!parameters_.u_solver_factory) {
            // when u_solver_type is LinOp, use UpperTrs as the default one
            u_solver_ = generate_default_solver<std::conditional_t<
                std::is_same_v<u_solver_type, LinOp>,
                solver::UpperTrs<value_type, index_type>, u_solver_type>>(
                exec, u_factor);
        } else {
            u_solver_ = parameters_.u_solver_factory->generate(u_factor);
        }
    }
 
    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)
    {
        // half can not use constexpr constructor
        const 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 u_solver_type> u_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_ILU_HPP_