ISTLSolverGPUISTL.hpp

/*
  Copyright 2025 Equinor ASA

  This file is part of the Open Porous Media project (OPM).
  OPM is free software: you can redistribute it and/or modify
  it under the terms of the GNU General Public License as published by
  the Free Software Foundation, either version 3 of the License, or
  (at your option) any later version.
  OPM is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  GNU General Public License for more details.
  You should have received a copy of the GNU General Public License
  along with OPM.  If not, see <http://www.gnu.org/licenses/>.
*/

#ifndef OPM_ISTLSOLVERGPUISTL_HEADER_INCLUDED
#define OPM_ISTLSOLVERGPUISTL_HEADER_INCLUDED

#include <dune/istl/operators.hh>
#include <memory>
#include <optional>
#include <opm/grid/utility/ElementChunks.hpp>
#include <opm/simulators/linalg/AbstractISTLSolver.hpp>
#include <opm/simulators/linalg/getQuasiImpesWeights.hpp>
#include <opm/simulators/linalg/ISTLSolver.hpp>

#if USE_HIP
#include <opm/simulators/linalg/gpuistl_hip/GpuSparseMatrixWrapper.hpp>
#include <opm/simulators/linalg/gpuistl_hip/GpuVector.hpp>
#include <opm/simulators/linalg/gpuistl_hip/PinnedMemoryHolder.hpp>
#else
#include <opm/simulators/linalg/gpuistl/GpuSparseMatrixWrapper.hpp>
#include <opm/simulators/linalg/gpuistl/GpuVector.hpp>
#include <opm/simulators/linalg/gpuistl/PinnedMemoryHolder.hpp>
#endif

#include <opm/simulators/linalg/ExtractParallelGridInformationToISTL.hpp>
#include <opm/simulators/linalg/ParallelIstlInformation.hpp>
#include <opm/simulators/linalg/findOverlapRowsAndColumns.hpp>
#include <opm/simulators/linalg/gpuistl/detail/FlexibleSolverWrapper.hpp>
#include <opm/simulators/linalg/printlinearsolverparameter.hpp>

namespace Opm::gpuistl
{

template <class TypeTag>
class ISTLSolverGPUISTL : public AbstractISTLSolver<GetPropType<TypeTag, Properties::SparseMatrixAdapter>,
                                                    GetPropType<TypeTag, Properties::GlobalEqVector>>
{
public:
    using SparseMatrixAdapter = GetPropType<TypeTag, Properties::SparseMatrixAdapter>;
    using Vector = GetPropType<TypeTag, Properties::GlobalEqVector>;
    using Simulator = GetPropType<TypeTag, Properties::Simulator>;
    using ElementMapper = GetPropType<TypeTag, Properties::ElementMapper>;
    using Matrix = typename SparseMatrixAdapter::IstlMatrix;
    using ThreadManager = GetPropType<TypeTag, Properties::ThreadManager>;
    using GridView = GetPropType<TypeTag, Properties::GridView>;
    using ElementContext = GetPropType<TypeTag, Properties::ElementContext>;
    using ElementChunksType = Opm::ElementChunks<GridView, Dune::Partitions::All>;

    using real_type = typename Vector::field_type;

    using GPUMatrix = Opm::gpuistl::GpuSparseMatrixWrapper<real_type>;
    using GPUVector = Opm::gpuistl::GpuVector<real_type>;
    using GPUVectorInt = Opm::gpuistl::GpuVector<int>;

    constexpr static std::size_t pressureIndex = GetPropType<TypeTag, Properties::Indices>::pressureSwitchIdx;


#if HAVE_MPI
    using CommunicationType = Dune::OwnerOverlapCopyCommunication<int, int>;
#else
    using CommunicationType = Dune::Communication<int>;
#endif

    using SolverType = Opm::gpuistl::detail::FlexibleSolverWrapper<GPUMatrix, GPUVector, CommunicationType>;

    ISTLSolverGPUISTL(const Simulator& simulator,
                      const FlowLinearSolverParameters& parameters,
                      bool forceSerial = false)
        : m_parameters(parameters)
        , m_simulator(simulator)
        , m_forceSerial(forceSerial)
        , m_element_chunks(simulator.vanguard().gridView(), Dune::Partitions::all, ThreadManager::maxThreads())
    {
#if HAVE_MPI
        m_comm = std::make_shared<CommunicationType>(simulator.vanguard().grid().comm());
        // Extract parallel grid information to populate index sets
        extractParallelGridInformationToISTL(simulator.vanguard().grid(), m_parallelInformation);
        // Set up element mapper manually
        ElementMapper elemMapper(simulator.vanguard().gridView(), Dune::mcmgElementLayout());
        // Overlap rows are needed for making overlap rows invalid in parallel mode
        Opm::detail::findOverlapAndInterior(simulator.vanguard().grid(), elemMapper, m_overlapRows, m_interiorRows);
        if (isParallel()) {
            const std::size_t size = simulator.vanguard().grid().leafGridView().size(0);
            // Copy parallel information to communication object (index sets and remote indices)
            Opm::detail::copyParValues(m_parallelInformation, size, *m_comm);
        }

#else
        m_comm = std::make_shared<CommunicationType>(simulator.gridView().comm());
#endif
        m_parameters.init(simulator.vanguard().eclState().getSimulationConfig().useCPR());
        m_propertyTree = setupPropertyTree(m_parameters,
                                           Parameters::IsSet<Parameters::LinearSolverMaxIter>(),
                                           Parameters::IsSet<Parameters::LinearSolverReduction>());
        if (!Parameters::Get<Parameters::MatrixAddWellContributions>()) {
            OPM_THROW(std::logic_error, "Well operators are currently not supported for the GPU backend. "
            "Use --matrix-add-well-contributions=true to add well contributions to the matrix instead.");
        }

        Opm::detail::printLinearSolverParameters(m_parameters, m_propertyTree, simulator.gridView().comm());
    }

    explicit ISTLSolverGPUISTL(const Simulator& simulator)
        : ISTLSolverGPUISTL(simulator, FlowLinearSolverParameters(), false)
    {
    }

    void eraseMatrix() override
    {
        // Nothing, this is the same as the ISTLSolver
    }

    void setActiveSolver(int num) override
    {
        if (num != 0) {
            OPM_THROW(std::logic_error, "Only one solver available for the GPU backend.");
        }
    }

    int numAvailableSolvers() const override
    {
        return 1;
    }

    void prepare(const SparseMatrixAdapter& M, Vector& b) override
    {
        prepare(M.istlMatrix(), b);
    }

    void prepare(const Matrix& M, Vector& b) override
    {
        try {
            if (isParallel() && !m_overlapRows.empty()) {
                Opm::detail::makeOverlapRowsInvalid(const_cast<Matrix&>(M), m_overlapRows);
            }
            updateMatrix(M);
            updateRhs(b);
        }
        OPM_CATCH_AND_RETHROW_AS_CRITICAL_ERROR("This is likely due to a faulty linear solver JSON specification. "
                                                "Check for errors related to missing nodes.");
    }

    void setResidual(Vector&) override
    {
        // Should be handled in prepare() instead.
    }

    void getResidual(Vector& b) const override
    {
        if (!m_rhs) {
            OPM_THROW(std::runtime_error, "m_rhs not initialized, prepare(matrix, rhs); needs to be called");
        }
        m_rhs->copyToHost(b);
    }

    void setMatrix(const SparseMatrixAdapter&) override
    {
        // Should be handled in prepare() instead.
    }

    bool solve(Vector& x) override
    {
        // TODO: Write matrix to disk if needed
        Dune::InverseOperatorResult result;
        if (!m_matrix) {
            OPM_THROW(std::runtime_error, "m_matrix not initialized, prepare(matrix, rhs); needs to be called");
        }
        if (!m_rhs) {
            OPM_THROW(std::runtime_error, "m_rhs not initialized, prepare(matrix, rhs); needs to be called");
        }
        if (!m_gpuSolver) {
            OPM_THROW(std::runtime_error,
                      "m_gpuFlexibleSolver not initialized, prepare(matrix, rhs); needs to be called");
        }

        if (!m_x) {
            m_x = std::make_unique<GPUVector>(x);
            m_pinnedXMemory = std::make_unique<PinnedMemoryHolder<real_type>>(
                const_cast<real_type*>(&x[0][0]),
                x.dim()
            );
        } else {
            // copy from host to device using main stream and asynchronous transfer
            m_x->copyFromHostAsync(x);
        }
        m_gpuSolver->apply(*m_x, *m_rhs, result);

        m_x->copyToHost(x);

        ++m_solveCount;

        m_lastSeenIterations = result.iterations;
        return checkConvergence(result);
    }

    int iterations() const override
    {
        return m_lastSeenIterations;
    }

     const CommunicationType* comm() const override
     {
         return m_comm.get();
     }

     bool isParallel() const
     {
 #if HAVE_MPI
         return !m_forceSerial && m_comm->communicator().size() > 1;
 #else
         return false;
 #endif
     }

    int getSolveCount() const override
    {
        return m_solveCount;
    }

private:
    bool checkConvergence(const Dune::InverseOperatorResult& result) const
    {
        return AbstractISTLSolver<SparseMatrixAdapter, Vector>::checkConvergence(result, m_parameters);
    }

    // Weights to make approximate pressure equations.
    std::function<GPUVector&()> getWeightsCalculator()
    {
        std::function<GPUVector&()> weightsCalculator;

        using namespace std::string_literals;

        auto preconditionerType = m_propertyTree.get("preconditioner.type"s, "cpr"s);
        // Make the preconditioner type lowercase for internal canonical representation
        std::transform(preconditionerType.begin(), preconditionerType.end(), preconditionerType.begin(), ::tolower);
        if (preconditionerType == "cpr" || preconditionerType == "cprt"
            || preconditionerType == "cprw" || preconditionerType == "cprwt") {
            const bool transpose = preconditionerType == "cprt" || preconditionerType == "cprwt";
            const auto weightsType = m_propertyTree.get("preconditioner.weight_type"s, "quasiimpes"s);
            if (weightsType == "quasiimpes") {
                m_weights.emplace(m_matrix->N() * m_matrix->blockSize());
                // Pre-compute diagonal indices once when setting up the calculator
                auto diagonalIndices = Amg::precomputeDiagonalIndices(*m_matrix);
                m_diagonalIndices.emplace(diagonalIndices);

                if (transpose) {
                    weightsCalculator = [this]() -> GPUVector& {
                        Amg::getQuasiImpesWeights<real_type, true>(
                            *m_matrix, pressureIndex, *m_weights, *m_diagonalIndices);
                        return *m_weights;
                    };
                } else {
                    weightsCalculator = [this]() -> GPUVector& {
                        Amg::getQuasiImpesWeights<real_type, false>(
                            *m_matrix, pressureIndex, *m_weights, *m_diagonalIndices);
                        return *m_weights;
                    };
                }
            } else if (weightsType == "trueimpes") {
                // Create CPU vector for the weights and initialize GPU vector
                m_cpuWeights.resize(m_matrix->N());
                m_pinnedWeightsMemory = std::make_unique<PinnedMemoryHolder<real_type>>(
                    const_cast<real_type*>(&m_cpuWeights[0][0]), m_cpuWeights.dim());
                m_weights.emplace(m_cpuWeights);

                const bool enableThreadParallel = m_parameters.cpr_weights_thread_parallel_;
                // CPU implementation wrapped for GPU
                weightsCalculator = [this, enableThreadParallel]() -> GPUVector& {
                    // Use the CPU implementation to calculate the weights
                    ElementContext elemCtx(m_simulator);
                    Amg::getTrueImpesWeights(pressureIndex,
                                             m_cpuWeights,
                                             elemCtx, m_simulator.model(),
                                             m_element_chunks,
                                             enableThreadParallel);

                    // Copy CPU vector to GPU vector using main stream and asynchronous transfer
                    m_weights->copyFromHostAsync(m_cpuWeights);
                    return *m_weights;
                };
            } else if (weightsType == "trueimpesanalytic") {
                // Create CPU vector for the weights and initialize GPU vector
                m_cpuWeights.resize(m_matrix->N());
                m_pinnedWeightsMemory = std::make_unique<PinnedMemoryHolder<real_type>>(
                    const_cast<real_type*>(&m_cpuWeights[0][0]), m_cpuWeights.dim());
                m_weights.emplace(m_cpuWeights);
                const bool enableThreadParallel = m_parameters.cpr_weights_thread_parallel_;
                // CPU implementation wrapped for GPU
                weightsCalculator = [this, enableThreadParallel]() -> GPUVector& {
                    // Use the CPU implementation to calculate the weights
                    ElementContext elemCtx(m_simulator);
                    Amg::getTrueImpesWeightsAnalytic(pressureIndex,
                                                     m_cpuWeights,
                                                     elemCtx, m_simulator.model(),
                                                     m_element_chunks,
                                                     enableThreadParallel);

                    // Copy CPU vector to GPU vector using main stream and asynchronous transfer
                    m_weights->copyFromHostAsync(m_cpuWeights);
                    return *m_weights;
                };
            } else {
                OPM_THROW(std::invalid_argument,
                          "Weights type " + weightsType
                              + " not implemented for cpr."
                                " Please use quasiimpes, trueimpes or trueimpesanalytic.");
            }
        }
        return weightsCalculator;
    }

    void updateMatrix(const Matrix& M)
    {
        if (!m_matrix) {
            m_matrix.reset(new auto(GPUMatrix::fromMatrix(M)));
            m_pinnedMatrixMemory = std::make_unique<PinnedMemoryHolder<real_type>>(
                const_cast<real_type*>(&M[0][0][0][0]),
                M.nonzeroes() * M[0][0].N() * M[0][0].M()
            );
            std::function<GPUVector&()> weightsCalculator = getWeightsCalculator();
            m_gpuSolver = std::make_unique<SolverType>(
                *m_matrix, isParallel(), m_propertyTree, pressureIndex, weightsCalculator, m_forceSerial, m_comm.get());
        } else {
            m_matrix->updateNonzeroValues(M, true);
            m_gpuSolver->update();
        }
    }

    void updateRhs(const Vector& b)
    {
        if (!m_rhs) {
            m_rhs = std::make_unique<GPUVector>(b);
            m_pinnedRhsMemory = std::make_unique<PinnedMemoryHolder<real_type>>(
                const_cast<real_type*>(&b[0][0]),
                b.dim()
            );
        } else {
            // copy from host to device using main stream and asynchronous transfer
            m_rhs->copyFromHostAsync(b);
        }
    }

    FlowLinearSolverParameters m_parameters;
    const Simulator& m_simulator;
    const bool m_forceSerial;
    PropertyTree m_propertyTree;
    ElementChunksType m_element_chunks;

    int m_lastSeenIterations = 0;
    int m_solveCount = 0;

    std::unique_ptr<GPUMatrix> m_matrix;

    std::unique_ptr<SolverType> m_gpuSolver;

    std::unique_ptr<GPUVector> m_rhs;
    std::unique_ptr<GPUVector> m_x;

    std::unique_ptr<PinnedMemoryHolder<real_type>> m_pinnedMatrixMemory;
    std::unique_ptr<PinnedMemoryHolder<real_type>> m_pinnedRhsMemory;
    std::unique_ptr<PinnedMemoryHolder<real_type>> m_pinnedXMemory;
    std::unique_ptr<PinnedMemoryHolder<real_type>> m_pinnedWeightsMemory;

    Vector m_cpuWeights;
    std::optional<GPUVector> m_weights;
    std::optional<GPUVectorInt> m_diagonalIndices;

    std::shared_ptr<CommunicationType> m_comm;
    std::vector<int> m_interiorRows;
    std::vector<int> m_overlapRows;
    std::any m_parallelInformation;

};
} // namespace Opm::gpuistl

#endif