HyprePreconditioner.hpp

/*
  Copyright 2024 SINTEF AS
  Copyright 2024-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_HYPRE_PRECONDITIONER_HEADER_INCLUDED
#define OPM_HYPRE_PRECONDITIONER_HEADER_INCLUDED
 
#include <opm/common/ErrorMacros.hpp>
#include <opm/common/TimingMacros.hpp>
#include <opm/simulators/linalg/PreconditionerWithUpdate.hpp>
#include <opm/simulators/linalg/PropertyTree.hpp>
#include <opm/simulators/linalg/gpuistl/HypreInterface.hpp>
#include <opm/simulators/linalg/gpuistl/detail/gpu_type_detection.hpp>
 
#include <dune/common/fmatrix.hh>
#include <dune/istl/bcrsmatrix.hh>
 
#include <HYPRE.h>
#include <HYPRE_krylov.h>
#include <HYPRE_parcsr_ls.h>
#include <_hypre_utilities.h>
 
#include <numeric>
#include <vector>
 
namespace Hypre
{
 
namespace HypreInterface = Opm::gpuistl::HypreInterface;

template <class M, class X, class Y, class Comm>
class HyprePreconditioner : public Dune::PreconditionerWithUpdate<X, Y>
{
public:
    using matrix_type = M;
    using matrix_field_type = typename M::field_type;
    using domain_type = X;
    using range_type = Y;
    using vector_field_type = typename X::field_type;

    HyprePreconditioner(const M& A, const Opm::PropertyTree prm, const Comm& comm)
        : A_(A)
        , comm_(comm)
    {
        OPM_TIMEBLOCK(prec_construct);
        int size;
        int rank;
        MPI_Comm mpi_comm;
        if constexpr (std::is_same_v<Comm, Dune::Amg::SequentialInformation>) {
            mpi_comm = MPI_COMM_SELF;
        } else {
            mpi_comm = comm.communicator();
        }
        MPI_Comm_size(mpi_comm, &size);
        MPI_Comm_rank(mpi_comm, &rank);
        if (size > 1) {
            assert(size == comm.communicator().size());
            assert(rank == comm.communicator().rank());
        }
        // Set use_gpu_backend_ to user value if specified, otherwise match input type
        use_gpu_backend_ = prm.get<bool>("use_gpu", Opm::gpuistl::is_gpu_type<M>::value);
 
        // Initialize Hypre library with backend configuration
        HypreInterface::initialize(use_gpu_backend_);
 
        // Create solver
        solver_ = HypreInterface::createAMGSolver();
        HypreInterface::setSolverParameters(solver_, prm, use_gpu_backend_);
 
        // Setup parallel info and mappings
        par_info_ = HypreInterface::setupHypreParallelInfo(comm_, A_);
 
        // Setup sparsity pattern
        sparsity_pattern_ = HypreInterface::setupSparsityPattern(A_, par_info_, par_info_.owner_first);
 
        // Setup host arrays
        host_arrays_.row_indexes = HypreInterface::computeRowIndexes(
            A_, sparsity_pattern_.ncols, par_info_.local_dune_to_local_hypre, par_info_.owner_first);
 
        // Create indices for vector operations - simple sequential indices for owned DOFs
        host_arrays_.indices.resize(par_info_.N_owned);
        std::iota(host_arrays_.indices.begin(), host_arrays_.indices.end(), par_info_.dof_offset);
 
        // Setup continuous vector values buffer - only needed for non-owner-first
        if (!par_info_.owner_first) {
            host_arrays_.continuous_vector_values.resize(par_info_.N_owned);
        }
 
        // Allocate device arrays if using GPU backend
        if (use_gpu_backend_) {
#if HYPRE_USING_CUDA || HYPRE_USING_HIP
            device_arrays_.ncols_device = hypre_CTAlloc(HYPRE_Int, par_info_.N_owned, HYPRE_MEMORY_DEVICE);
            device_arrays_.rows_device = hypre_CTAlloc(HYPRE_BigInt, par_info_.N_owned, HYPRE_MEMORY_DEVICE);
            device_arrays_.cols_device = hypre_CTAlloc(HYPRE_BigInt, sparsity_pattern_.nnz, HYPRE_MEMORY_DEVICE);
            device_arrays_.row_indexes_device = hypre_CTAlloc(HYPRE_Int, par_info_.N_owned, HYPRE_MEMORY_DEVICE);
            device_arrays_.indices_device = hypre_CTAlloc(HYPRE_BigInt, par_info_.N_owned, HYPRE_MEMORY_DEVICE);
            device_arrays_.vector_buffer_device = hypre_CTAlloc(HYPRE_Real, par_info_.N_owned, HYPRE_MEMORY_DEVICE);
            if constexpr (!Opm::gpuistl::is_gpu_type<M>::value) {
                // For CPU input and GPU backend we need to allocate space for transfering the matrix values
                // Note that the buffer must be allocated with the number of nonzeroes in the matrix, not the
                // sparsity_pattern.nnz because we need to copy the entire matrix values from the host to the device.
                device_arrays_.matrix_buffer_device = hypre_CTAlloc(HYPRE_Real, A_.nonzeroes(), HYPRE_MEMORY_DEVICE);
            }
            // Copy data to device
            hypre_TMemcpy(device_arrays_.ncols_device, sparsity_pattern_.ncols.data(), HYPRE_Int, par_info_.N_owned, HYPRE_MEMORY_DEVICE, HYPRE_MEMORY_HOST);
            hypre_TMemcpy(device_arrays_.rows_device, sparsity_pattern_.rows.data(), HYPRE_BigInt, par_info_.N_owned, HYPRE_MEMORY_DEVICE, HYPRE_MEMORY_HOST);
            hypre_TMemcpy(device_arrays_.cols_device, sparsity_pattern_.cols.data(), HYPRE_BigInt, sparsity_pattern_.nnz, HYPRE_MEMORY_DEVICE, HYPRE_MEMORY_HOST);
            hypre_TMemcpy(device_arrays_.row_indexes_device, host_arrays_.row_indexes.data(), HYPRE_Int, par_info_.N_owned, HYPRE_MEMORY_DEVICE, HYPRE_MEMORY_HOST);
            hypre_TMemcpy(device_arrays_.indices_device, host_arrays_.indices.data(), HYPRE_BigInt, par_info_.N_owned, HYPRE_MEMORY_DEVICE, HYPRE_MEMORY_HOST);
#endif
        }
 
        // Create Hypre matrix and vectors
        A_hypre_ = HypreInterface::createMatrix(par_info_.N_owned, par_info_.dof_offset, comm_);
        x_hypre_ = HypreInterface::createVector(par_info_.N_owned, par_info_.dof_offset, comm_);
        b_hypre_ = HypreInterface::createVector(par_info_.N_owned, par_info_.dof_offset, comm_);
 
        // Perform initial update
        update();
    }

    ~HyprePreconditioner()
    {
        // Clean up device arrays if allocated
        if (use_gpu_backend_) {
#if HYPRE_USING_CUDA || HYPRE_USING_HIP
            if (device_arrays_.ncols_device) {
                hypre_TFree(device_arrays_.ncols_device, HYPRE_MEMORY_DEVICE);
            }
            if (device_arrays_.rows_device) {
                hypre_TFree(device_arrays_.rows_device, HYPRE_MEMORY_DEVICE);
            }
            if (device_arrays_.cols_device) {
                hypre_T