PreconditionerConvertFieldTypeAdapter.hpp

/*
  Copyright 2022-2023 SINTEF AS

  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_PRECONDITIONERCONVERTOFLOATADAPTER_HPP
#define OPM_PRECONDITIONERCONVERTOFLOATADAPTER_HPP
#include <cusparse.h>
#include <dune/istl/bcrsmatrix.hh>
#include <dune/istl/preconditioner.hh>
#include <opm/simulators/linalg/PreconditionerWithUpdate.hpp>
#include <opm/simulators/linalg/gpuistl/GpuSparseMatrixWrapper.hpp>
#include <opm/simulators/linalg/gpuistl/detail/CuMatrixDescription.hpp>
#include <opm/simulators/linalg/gpuistl/detail/CuSparseHandle.hpp>
#include <opm/simulators/linalg/gpuistl/detail/CuSparseResource.hpp>
#include <opm/simulators/linalg/gpuistl/detail/gpu_constants.hpp>
#include <opm/simulators/linalg/gpuistl/detail/cusparse_safe_call.hpp>
#include <opm/simulators/linalg/gpuistl/detail/preconditioner_should_call_post_pre.hpp>


namespace Opm::gpuistl
{
template <class CudaPreconditionerType, class M, class X, class Y, int l = 1>
class PreconditionerConvertFieldTypeAdapter : public Dune::PreconditionerWithUpdate<X, Y>
{
public:
    using matrix_type = typename std::remove_const<M>::type;

    using domain_type = X;
    using range_type = Y;
    using field_type = typename X::field_type;


    using domain_type_to = typename CudaPreconditionerType::domain_type;
    using range_type_to = typename CudaPreconditionerType::range_type;
    using field_type_to = typename domain_type_to::field_type;


    using block_type = typename domain_type::block_type;

    using XTo = Dune::BlockVector<Dune::FieldVector<field_type_to, block_type::dimension>>;
    using YTo = Dune::BlockVector<Dune::FieldVector<field_type_to, block_type::dimension>>;
    using matrix_type_to =
        typename Dune::BCRSMatrix<Dune::FieldMatrix<field_type_to, block_type::dimension, block_type::dimension>>;

    explicit PreconditionerConvertFieldTypeAdapter(const M& matrix)
        : m_matrix(matrix)
        , m_convertedMatrix(createConvertedMatrix())
    {
    }

    virtual void pre([[maybe_unused]] X& x, [[maybe_unused]] Y& b) override
    {
        static_assert(!detail::shouldCallPreconditionerPre<CudaPreconditionerType>(),
                      "We currently do not support Preconditioner::pre().");
    }

    virtual void apply(X& v, const Y& d) override
    {
        OPM_ERROR_IF(!m_underlyingPreconditioner,
                     "You need to set the underlying preconditioner with setUnderlyingPreconditioner.");
        XTo convertedV(v.N());
        for (size_t i = 0; i < v.N(); ++i) {
            for (size_t j = 0; j < block_type::dimension; ++j) {
                // This is probably unnecessary, but doing it anyway:
                convertedV[i][j] = field_type_to(v[i][j]);
            }
        }
        YTo convertedD(d.N());
        for (size_t i = 0; i < d.N(); ++i) {
            for (size_t j = 0; j < block_type::dimension; ++j) {
                convertedD[i][j] = field_type_to(d[i][j]);
            }
        }

        m_underlyingPreconditioner->apply(convertedV, convertedD);

        for (size_t i = 0; i < v.N(); ++i) {
            for (size_t j = 0; j < block_type::dimension; ++j) {
                v[i][j] = field_type(convertedV[i][j]);
            }
        }
    }

    virtual void post([[maybe_unused]] X& x) override
    {
        static_assert(!detail::shouldCallPreconditionerPost<CudaPreconditionerType>(),
                      "We currently do not support Preconditioner::post().");
    }

    virtual Dune::SolverCategory::Category category() const override
    {
        return m_underlyingPreconditioner->category();
    }

    virtual void update() override
    {
        OPM_ERROR_IF(!m_underlyingPreconditioner,
                     "You need to set the underlying preconditioner with setUnderlyingPreconditioner.");
        updateMatrix();
        m_underlyingPreconditioner->update();
    }

    const matrix_type_to& getConvertedMatrix() const
    {
        return m_convertedMatrix;
    }

    void setUnderlyingPreconditioner(const std::shared_ptr<CudaPreconditionerType>& conditioner)
    {
        m_underlyingPreconditioner = conditioner;
    }

    virtual bool hasPerfectUpdate() const override {
        return m_underlyingPreconditioner->hasPerfectUpdate();
    }


private:
    void updateMatrix()
    {
        const auto nnz = m_matrix.nonzeroes() * m_matrix[0][0].N() * m_matrix[0][0].N();
        const auto dataPointerIn = static_cast<const field_type*>(&((m_matrix[0][0][0][0])));
        auto dataPointerOut = static_cast<field_type_to*>(&((m_convertedMatrix[0][0][0][0])));

        std::vector<field_type_to> buffer(nnz, 0);
        for (size_t i = 0; i < nnz; ++i) {
            dataPointerOut[i] = field_type_to(dataPointerIn[i]);
        }
    }
    matrix_type_to createConvertedMatrix()
    {
        // TODO: Check if this whole conversion can be done more efficiently.
        const auto N = m_matrix.N();
        matrix_type_to matrixBuilder(N, N, m_matrix.nonzeroes(), matrix_type_to::row_wise);
        {
            auto rowIn = m_matrix.begin();
            for (auto rowOut = matrixBuilder.createbegin(); rowOut != matrixBuilder.createend(); ++rowOut) {
                for (auto column = rowIn->begin(); column != rowIn->end(); ++column) {
                    rowOut.insert(column.index());
                }
                ++rowIn;
            }
        }

        for (auto row = m_matrix.begin(); row != m_matrix.end(); ++row) {
            for (auto column = row->begin(); column != row->end(); ++column) {
                for (size_t i = 0; i < block_type::dimension; ++i) {
                    for (size_t j = 0; j < block_type::dimension; ++j) {
                        matrixBuilder[row.index()][column.index()][i][j]
                            = field_type_to(m_matrix[row.index()][column.index()][i][j]);
                    }
                }
            }
        }

        return matrixBuilder;
    }
    const M& m_matrix;
    matrix_type_to m_convertedMatrix;
    std::shared_ptr<CudaPreconditionerType> m_underlyingPreconditioner;
};
} // end namespace Opm::gpuistl

#endif