ginkgo.hpp

Go to the documentation of this file.

// SPDX-FileCopyrightText: 2024 - 2026 NeoN authors
//
// SPDX-License-Identifier: MIT
 
#pragma once
 
#if NF_WITH_GINKGO
 
#include <chrono>
#include <optional>
#include <string>
 
#include <ginkgo/ginkgo.hpp>
#include <ginkgo/extensions/kokkos.hpp>
#include <ginkgo/extensions/config/json_config.hpp>
 
#include "NeoN/fields/field.hpp"
#include "NeoN/core/dictionary.hpp"
#include "NeoN/core/executor/executor.hpp"
#include "NeoN/linearAlgebra/solver.hpp"
#include "NeoN/linearAlgebra/linearSystem.hpp"
#include "NeoN/linearAlgebra/utilities.hpp"
 
 
namespace NeoN::la::ginkgo
{
 
std::shared_ptr<gko::Executor> getGkoExecutor(Executor exec);
 
gko::config::pnode parse(const Dictionary& dict);
 
template<typename NeoNMatrixType>
std::shared_ptr<const gko::LinOp> createGkoMtx(const NeoNMatrixType& mtx);
 
#ifdef NF_WITH_MPI_SUPPORT
#include "NeoN/distributed/communicationPattern.hpp"
 
template<typename IndexType>
std::shared_ptr<const gko::LinOp> createGkoMtxDist(
    std::shared_ptr<const gko::Executor> exec,
    const gko::experimental::mpi::communicator& comm,
    const CSRMatrix<scalar, IndexType>& mtx,
    const COOMatrix<scalar, IndexType>& bmtx,
    const CommunicationPattern& commPattern,
    std::shared_ptr<gko::experimental::distributed::index_map<label, gko::int64>>& imapCache,
    std::shared_ptr<gko::matrix::Coo<scalar, IndexType>>& nonLocalMtxCache
);
#endif // NF_WITH_MPI_SUPPORT
 
template<typename T>
gko::array<T> gkoArrayView(std::shared_ptr<const gko::Executor> exec, std::span<T> values)
{
    return gko::make_array_view(exec, values.size(), values.data());
}
 
template<typename T>
std::shared_ptr<gko::matrix::Dense<scalar>>
gkoVecView(std::shared_ptr<const gko::Executor> exec, T* ptr, localIdx s)
{
    constexpr std::size_t cols = std::is_same_v<T, Vec3> ? 3 : 1;
    auto size = static_cast<std::size_t>(s);
    return gko::share(gko::matrix::Dense<scalar>::create(
        exec,
        gko::dim<2> {size, cols},
        gkoArrayView<scalar>(exec, std::span<scalar> {reinterpret_cast<scalar*>(ptr), cols * size}),
        cols
    ));
}
 
template<typename T>
std::shared_ptr<const gko::matrix::Dense<scalar>>
gkoVecView(std::shared_ptr<const gko::Executor> exec, const T* ptr, localIdx s)
{
    constexpr std::size_t cols = std::is_same_v<T, Vec3> ? 3 : 1;
    auto size = static_cast<std::size_t>(s);
    return gko::share(gko::matrix::Dense<scalar>::create_const(
        exec,
        gko::dim<2> {size, cols},
        gko::array<scalar>::const_view(exec, cols * size, reinterpret_cast<const scalar*>(ptr)),
        cols
    ));
}
 
template<typename InType>
scalar retrieve(const InType& in)
{
    using vec = gko::matrix::Dense<scalar>;
    auto host = vec::create(in->get_executor()->get_master(), gko::dim<2> {1});
    return host->copy_from(in)->at(0);
}
 
struct L1ResidualControl
{
    scalar tolerance;        
    scalar relTol;           
    gko::size_type maxIter;  
    gko::size_type minIter;  
    localIdx checkFrequency; 
};
 
struct L1ResidualResult
{
    gko::size_type numIter = 0; 
    scalar initResNorm = 0.0;   
    scalar finalResNorm = 0.0;  
    bool fired = false;         
    // Per-column scaled residuals; populated only for multi-RHS (Vec3) solves (size == ncols).
    std::vector<scalar> perColInitNorms;
    std::vector<scalar> perColFinalNorms;
};
 
inline constexpr const char* l1CriterionKey = "neon::l1ScaledResidual";
 
std::shared_ptr<gko::stop::CriterionFactory> makeL1CriterionFactory(
    std::shared_ptr<const gko::Executor> exec,
    const L1ResidualControl& control,
    L1ResidualResult* report
);
 
L1ResidualResult solveWithL1Stop(
    std::shared_ptr<const gko::Executor> exec,
    std::shared_ptr<const gko::LinOp> mtx,
    std::shared_ptr<const gko::matrix::Dense<scalar>> b,
    std::shared_ptr<gko::matrix::Dense<scalar>> x,
    gko::LinOp* solver,
    const L1ResidualControl& control
);
 
#ifdef NF_WITH_MPI_SUPPORT
L1ResidualResult solveWithL1StopDist(
    std::shared_ptr<const gko::Executor> exec,
    std::shared_ptr<const gko::LinOp> mtx,
    std::shared_ptr<const gko::experimental::distributed::Vector<scalar>> b,
    std::shared_ptr<gko::experimental::distributed::Vector<scalar>> x,
    gko::LinOp* solver,
    const L1ResidualControl& control
);
#endif
 
inline std::optional<L1ResidualControl> readL1ResidualControl(const Dictionary& cfg)
{
    const std::string flag = "l1ScaledResidual";
    if (!cfg.contains(flag))
    {
        return std::nullopt;
    }
    // A boolean read from a dictionary file is stored as a word/string, not a bool;
    // accept the common representations rather than assuming a single type.
    bool enabled = false;
    if (cfg.isType<bool>(flag))
    {
        enabled = cfg.get<bool>(flag);
    }
    else if (cfg.isType<int>(flag))
    {
        enabled = cfg.get<int>(flag) != 0;
    }
    else if (cfg.isType<std::string>(flag))
    {
        const std::string v = cfg.get<std::string>(flag);
        enabled = (v == "true" || v == "yes" || v == "on" || v == "1");
    }
    if (!enabled)
    {
        return std::nullopt;
    }
 
    L1ResidualControl control {0.0, 0.0, 1000, 0, 1};
 
    // criteria entries may be stored as int, label or scalar depending on source
    auto readScalar = [](const Dictionary& d, const std::string& key, scalar fallback)
    {
        if (!d.contains(key)) return fallback;
        if (d.isType<int>(key)) return scalar(d.get<int>(key));
        if (d.isType<label>(key)) return scalar(d.get<label>(key));
        return d.get<scalar>(key);
    };
    auto readInt = [](const Dictionary& d, const std::string& key, localIdx fallback)
    {
        if (!d.contains(key)) return fallback;
        if (d.isType<int>(key)) return localIdx(d.get<int>(key));
        if (d.isType<scalar>(key)) return localIdx(d.get<scalar>(key));
        return d.get<localIdx>(key);
    };
 
    // OpenFOAM-style top-level keys (the configFile solver path keeps them: there is no
    // mapped "criteria" subdict there). Read first so they seed the control...
    control.tolerance = readScalar(cfg, "tolerance", control.tolerance);
    control.relTol = readScalar(cfg, "relTol", control.relTol);
    control.maxIter = readInt(cfg, "maxIter", control.maxIter);
 
    // ...then let a mapped "criteria" subdict (Ginkgo-style keys) override when present,
    // preserving the existing behaviour for fvSolution-mapped solvers.
    if (cfg.contains("criteria"))
    {
        const Dictionary& criteria = cfg.subDict("criteria");
        control.tolerance = readScalar(criteria, "absolute_residual_norm", control.tolerance);
        control.relTol = readScalar(criteria, "initial_residual_norm", control.relTol);
        control.maxIter = readInt(criteria, "iteration", control.maxIter);
    }
    control.minIter = readInt(cfg, "minIter", control.minIter);
    control.checkFrequency = readInt(cfg, "checkFrequency", control.checkFrequency);
 
    return control;
}
 
inline bool pnodeReferencesString(const gko::config::pnode& node, const std::string& name)
{
    using tag = gko::config::pnode::tag_t;
    switch (node.get_tag())
    {
    case tag::string:
        return node.get_string() == name;
    case tag::array:
        for (const auto& e : node.get_array())
            if (pnodeReferencesString(e, name)) return true;
        return false;
    case tag::map:
        for (const auto& kv : node.get_map())
            if (pnodeReferencesString(kv.second, name)) return true;
        return false;
    default:
        return false;
    }
}
 
class GinkgoSolver : public SolverFactory::template Register<GinkgoSolver>
{
 
    using Base = SolverFactory::template Register<GinkgoSolver>;
 
public:
 
    GinkgoSolver(Executor exec, const Dictionary& solverConfig)
        : Base(exec), gkoExec_(getGkoExecutor(exec)), coupled_(solverConfig.get("coupled", false)),
          l1Control_(readL1ResidualControl(solverConfig)), config_(parse(solverConfig))
    {
        // Register NeoN's L1-scaled residual criterion in the Ginkgo config registry so a
        // configFile can name it (l1CriterionKey) in its "criteria" array. Only needed when
        // the L1 stop is requested (l1ScaledResidual); the factory carries the tolerances and
        // a report sink. If the parsed config actually references it, the criterion lives
        // INSIDE the built solver (l1InConfig_) and reports via l1Report_; otherwise the
        // existing post-hoc attach in solve() handles the flag-only case unchanged.
        gko::config::registry reg;
        if (l1Control_)
        {
            l1CritFactory_ = makeL1CriterionFactory(gkoExec_, *l1Control_, &l1Report_);
            reg.emplace(std::string(l1CriterionKey), l1CritFactory_);
            l1InConfig_ = pnodeReferencesString(config_, l1CriterionKey);
        }
        factory_ = gko::config::parse(config_, reg, gko::config::make_type_descriptor<scalar>())
                       .on(gkoExec_);
    }
 
    static std::string name() { return "Ginkgo"; }
 
    static std::string doc() { return "TBD"; }
 
    static std::string schema() { return "none"; }
 
    virtual SolverStats solve(
        const LinearSystem<scalar, scalar, CSRMatrix<scalar, localIdx>>& sys, Vector<scalar>& x
    ) const final;
 
    virtual SolverStats solve(
        const LinearSystem<Vec3, Vec3, CSRMatrix<Vec3, localIdx>>& sys, Vector<Vec3>& x
    ) const final;
 
    virtual SolverStats solve(
        const LinearSystem<scalar, Vec3, CSRMatrix<scalar, localIdx>, COOMatrix<scalar, localIdx>>&
            sys,
        Vector<Vec3>& x
    ) const final;
 
#ifdef NF_WITH_MPI_SUPPORT
    virtual SolverStats solveDist(
        const LinearSystem<scalar, scalar, CSRMatrix<scalar, localIdx>>& sys, Vector<scalar>& x
    ) const final;
 
    virtual SolverStats solveDist(
        const LinearSystem<Vec3, Vec3, CSRMatrix<Vec3, localIdx>>& sys, Vector<Vec3>& x
    ) const final;
 
    virtual SolverStats solveDist(
        const LinearSystem<scalar, Vec3, CSRMatrix<scalar, localIdx>, COOMatrix<scalar, localIdx>>&
            sys,
        Vector<Vec3>& x
    ) const final;
#endif
 
    // TODO why use a smart pointer here?
    virtual std::unique_ptr<SolverFactory> clone() const final
    {
        NF_ERROR_EXIT("Not implemented");
        return {};
    }
 
private:
 
    std::shared_ptr<const gko::Executor> gkoExec_;
    bool coupled_;
    std::optional<L1ResidualControl> l1Control_;
    gko::config::pnode config_;
    std::shared_ptr<const gko::LinOpFactory> factory_;
    // L1-scaled residual criterion registered into the config registry (l1Control_ set).
    std::shared_ptr<gko::stop::CriterionFactory> l1CritFactory_;
    // True when config_ names the L1 criterion, so it is built into factory_'s solver and
    // reports through l1Report_ (the in-config path); false keeps the post-hoc attach.
    bool l1InConfig_ = false;
    // Report sink the in-config criterion writes its scaled L1 residual / iters into.
    mutable L1ResidualResult l1Report_;
#ifdef NF_WITH_MPI_SUPPORT
    // Both caches are null until the first solve; after that topology is fixed.
    mutable std::shared_ptr<gko::experimental::distributed::index_map<label, gko::int64>>
        cachedImap_;
    mutable std::shared_ptr<gko::matrix::Coo<scalar, localIdx>> cachedNonLocalMtx_;
#endif
};
 
 
}
 
#endif