GpuPressureTransferPolicy.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_GPU_PRESSURE_TRANSFER_POLICY_HEADER_INCLUDED
#define OPM_GPU_PRESSURE_TRANSFER_POLICY_HEADER_INCLUDED
 
#include <opm/simulators/linalg/PropertyTree.hpp>
#include <opm/simulators/linalg/WellOperators.hpp>
#include <opm/simulators/linalg/gpuistl/GpuSparseMatrixWrapper.hpp>
#include <opm/simulators/linalg/gpuistl/GpuVector.hpp>
#include <opm/simulators/linalg/gpuistl/detail/cpr_amg_operations.hpp>
#include <opm/simulators/linalg/twolevelmethodcpr.hh>
 
#include <cstddef>
 
namespace Opm
{
namespace Details
{
    template <class Scalar>
    using GpuPressureMatrixType = gpuistl::GpuSparseMatrixWrapper<Scalar>;
    template <class Scalar>
    using GpuPressureVectorType = gpuistl::GpuVector<Scalar>;
    template <class Scalar>
    using GpuSeqCoarseOperatorType = Dune::
        MatrixAdapter<GpuPressureMatrixType<Scalar>, GpuPressureVectorType<Scalar>, GpuPressureVectorType<Scalar>>;
 
    template <class Scalar, class Comm>
    using GpuCoarseOperatorType = GpuSeqCoarseOperatorType<Scalar>;
} // namespace Details
} // namespace Opm
 
namespace Opm::gpuistl
{
template <class FineOperator, class Communication, class Scalar, bool transpose = false>
class GpuPressureTransferPolicy
    : public Dune::Amg::LevelTransferPolicyCpr<FineOperator, Details::GpuCoarseOperatorType<Scalar, Communication>>
{
public:
    using CoarseOperator = typename Details::GpuCoarseOperatorType<Scalar, Communication>;
    using ParentType = Dune::Amg::LevelTransferPolicyCpr<FineOperator, CoarseOperator>;
    using ParallelInformation = Communication;
    using FineVectorType = typename FineOperator::domain_type;
 
public:
    GpuPressureTransferPolicy(const Communication& comm,
                              const FineVectorType& weights,
                              [[maybe_unused]] const PropertyTree& prm,
                              int pressure_var_index)
        : communication_(&const_cast<Communication&>(comm))
        , weights_(weights)
        , pressure_var_index_(pressure_var_index)
    {
    }
 
 
    void createCoarseLevelSystem(const FineOperator& fineOperator) override
    {
        using CoarseMatrix = typename CoarseOperator::matrix_type;
        const auto& fineLevelMatrix = fineOperator.getmat();
        // Create coarse level matrix on GPU with block size 1 but same sparsity pattern
        coarseLevelMatrix_ = std::make_shared<CoarseMatrix>(fineLevelMatrix.getRowIndices(),
                                                            fineLevelMatrix.getColumnIndices(),
                                                            1 // block size 1
        );
 
        // Calculate entries for coarse matrix
        calculateCoarseEntries(fineOperator);
        coarseLevelCommunication_.reset(communication_, [](Communication*) {});
 
        this->lhs_.resize(coarseLevelMatrix_->N());
        this->rhs_.resize(coarseLevelMatrix_->N());
 
        // Create a MatrixAdapter that wraps the matrix without copying it
        this->operator_ = std::make_shared<CoarseOperator>(*coarseLevelMatrix_);
    }
 
    void calculateCoarseEntries(const FineOperator& fineOperator) override
    {
        const auto& fineLevelMatrix = fineOperator.getmat();
 
        // Set coarse matrix to zero
        coarseLevelMatrix_->getNonZeroValues() = 0.0;
 
        // Calculate coarse entries on GPU
        detail::calculateCoarseEntries<Scalar, transpose>(fineLevelMatrix, *coarseLevelMatrix_, weights_, pressure_var_index_);
    }
 
    void moveToCoarseLevel(const typename ParentType::FineRangeType& fine) override
    {
        // Set coarse vector to zero
        this->rhs_ = 0;
 
        // Restrict vector on GPU
        detail::restrictVector<Scalar, transpose>(fine, this->rhs_, weights_, pressure_var_index_);
 
        this->lhs_ = 0;
    }
 
    void moveToFineLevel(typename ParentType::FineDomainType& fine) override
    {
        // Prolongate vector on GPU
        detail::prolongateVector<Scalar, transpose>(this->lhs_, fine, weights_, pressure_var_index_);
    }
 
    GpuPressureTransferPolicy* clone() const override
    {
        return new GpuPressureTransferPolicy(*this);
    }
 
    const Communication& getCoarseLevelCommunication() const
    {
        return *coarseLevelCommunication_;
    }
 
    std::size_t getPressureIndex() const
    {
        return pressure_var_index_;
    }
 
private:
    Communication* communication_;
    const FineVectorType& weights_;
    const std::size_t pressure_var_index_;
    std::shared_ptr<Communication> coarseLevelCommunication_;
    std::shared_ptr<typename CoarseOperator::matrix_type> coarseLevelMatrix_;
};
 
} // namespace Opm::gpuistl
 
#endif // OPM_GPU_PRESSURE_TRANSFER_POLICY_HEADER_INCLUDED