GpuSparseMatrixWrapper.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_GPUSPARSEMATRIXWRAPPER_HPP
#define OPM_GPUSPARSEMATRIXWRAPPER_HPP

#include <opm/common/ErrorMacros.hpp>
#include <opm/simulators/linalg/gpuistl/detail/CuMatrixDescription.hpp>
#include <opm/simulators/linalg/gpuistl/detail/CuSparseHandle.hpp>
#include <opm/simulators/linalg/gpuistl/detail/safe_conversion.hpp>
#include <opm/simulators/linalg/gpuistl/GpuVector.hpp>
#include <opm/simulators/linalg/gpuistl/GpuSparseMatrix.hpp>
#include <opm/simulators/linalg/gpuistl/GpuSparseMatrixGeneric.hpp>

#include <cstddef>
#include <cuda.h>
#include <cusparse.h>
#include <memory>
#include <stdexcept>
#include <type_traits>

namespace Opm::gpuistl
{

template <typename T, bool ForceLegacy = false>
class GpuSparseMatrixWrapper
{
public:
    using field_type = T;

    /*
        Here is the primary function of this class.
        Since the generic API for CUDA/HIP is primaryly supported on CUDA 13 (and not yet HIP) for blocked
        matrices, places wanting to use blocked matrices can invoke this class which handles which API to use.
        Basically we just check if HIP is present, or if we are using cuda and a version prior to 13.
        If ForceLegacy is true, always use the legacy API (GpuSparseMatrix) regardless of CUDA version.
    */
#if USE_HIP || (!USE_HIP && CUDA_VERSION < 13000)
     using matrix_type = GpuSparseMatrix<T>;
#else
     using matrix_type = std::conditional_t<ForceLegacy,
                                            GpuSparseMatrix<T>,
                                            GpuSparseMatrixGeneric<T>>;
#endif

    // Arrow operator overloads for direct access to the underlying matrix
    matrix_type* operator->() {
        if (!m_matrix) {
            throw std::runtime_error("GpuSparseMatrixWrapper: underlying matrix is nullptr.");
        }
        return m_matrix.get();
    }
    const matrix_type* operator->() const {
        if (!m_matrix) {
            throw std::runtime_error("GpuSparseMatrixWrapper: underlying matrix is nullptr.");
        }
        return m_matrix.get();
    }

    static constexpr int max_block_size = 6;

    GpuSparseMatrixWrapper(const T* nonZeroElements,
                   const int* rowIndices,
                   const int* columnIndices,
                   std::size_t numberOfNonzeroBlocks,
                   std::size_t blockSize,
                   std::size_t numberOfRows)
    {
        m_matrix = std::make_unique<matrix_type>(nonZeroElements,
                                                 rowIndices,
                                                 columnIndices,
                                                 numberOfNonzeroBlocks,
                                                 blockSize,
                                                 numberOfRows);
    }

    GpuSparseMatrixWrapper(const GpuVector<int>& rowIndices,
                   const GpuVector<int>& columnIndices,
                   std::size_t blockSize)
    {
        m_matrix = std::make_unique<matrix_type>(rowIndices, columnIndices, blockSize);
    }

    GpuSparseMatrixWrapper(const GpuSparseMatrixWrapper& other)
    {
        if (!other.m_matrix) {
            throw std::runtime_error("Internal error, other.m_matrix is a nullptr.");
        }
        m_matrix = std::make_unique<matrix_type>(*other.m_matrix);
    }

    // We want to have this as non-mutable as possible, that is we do not want
    // to deal with changing matrix sizes and sparsity patterns.
    GpuSparseMatrixWrapper& operator=(const GpuSparseMatrixWrapper&) = delete;

    virtual ~GpuSparseMatrixWrapper() = default;

    GpuSparseMatrixWrapper() = default;

    const matrix_type& get() const {
        if (!m_matrix) {
            throw std::runtime_error("GpuSparseMatrixWrapper: underlying matrix is nullptr.");
        }
        return *m_matrix;
    }

    template <class MatrixType>
    static GpuSparseMatrixWrapper<T, ForceLegacy> fromMatrix(const MatrixType& matrix, bool copyNonZeroElementsDirectly = false)
    {
        GpuSparseMatrixWrapper<T, ForceLegacy> gpuSparseMatrixWrapper;
        gpuSparseMatrixWrapper.m_matrix = std::make_unique<matrix_type>(
            matrix_type::fromMatrix(matrix, copyNonZeroElementsDirectly));
        return gpuSparseMatrixWrapper;
    }

    // Only participates in overload resolution when matrix_type == GpuSparseMatrix<T>
    template <class M = matrix_type,
              typename = std::enable_if_t<std::is_same_v<M, GpuSparseMatrix<T>>>>
    void setUpperTriangular()
    {
        m_matrix->setUpperTriangular();
    }

    template <class M = matrix_type,
              typename = std::enable_if_t<std::is_same_v<M, GpuSparseMatrix<T>>>>
    void setLowerTriangular()
    {
        m_matrix->setLowerTriangular();
    }

    template <class M = matrix_type,
              typename = std::enable_if_t<std::is_same_v<M, GpuSparseMatrix<T>>>>
    void setUnitDiagonal()
    {
        m_matrix->setUnitDiagonal();
    }

    template <class M = matrix_type,
              typename = std::enable_if_t<std::is_same_v<M, GpuSparseMatrix<T>>>>
    void setNonUnitDiagonal()
    {
        m_matrix->setNonUnitDiagonal();
    }

    std::size_t N() const
    {
        return m_matrix->N();
    }

    std::size_t nonzeroes() const
    {
        return m_matrix->nonzeroes();
    }

    GpuVector<T>& getNonZeroValues()
    {
        return m_matrix->getNonZeroValues();
    }

    const GpuVector<T>& getNonZeroValues() const
    {
        return m_matrix->getNonZeroValues();
    }

    GpuVector<int>& getRowIndices()
    {
        return m_matrix->getRowIndices();
    }

    const GpuVector<int>& getRowIndices() const
    {
        return m_matrix->getRowIndices();
    }

    GpuVector<int>& getColumnIndices()
    {
        return m_matrix->getColumnIndices();
    }

    const GpuVector<int>& getColumnIndices() const
    {
        return m_matrix->getColumnIndices();
    }

    std::size_t dim() const
    {
        return m_matrix->dim();
    }

    std::size_t blockSize() const
    {
        return m_matrix->blockSize();
    }

    detail::GpuSparseMatrixDescription& getDescription()
    {
        return m_matrix->getDescription();
    }

    virtual void mv(const GpuVector<T>& x, GpuVector<T>& y) const
    {
        m_matrix->mv(x, y);
    }

    virtual void umv(const GpuVector<T>& x, GpuVector<T>& y) const
    {
        m_matrix->umv(x, y);
    }


    virtual void usmv(T alpha, const GpuVector<T>& x, GpuVector<T>& y) const
    {
        m_matrix->usmv(alpha, x, y);
    }

    template <class MatrixType>
    void updateNonzeroValues(const MatrixType& matrix, bool copyNonZeroElementsDirectly = false)
    {
        m_matrix->updateNonzeroValues(matrix, copyNonZeroElementsDirectly);
    }

    template <bool OtherForceLegacy>
    void updateNonzeroValues(const GpuSparseMatrixWrapper<T, OtherForceLegacy>& matrix)
    {
        m_matrix->updateNonzeroValues(matrix.get());
    }

    void setToZero()
    {
        m_matrix->setToZero();
    }

    template<class FunctionType>
    auto dispatchOnBlocksize(FunctionType function) const
    {
        return m_matrix->dispatchOnBlocksize(function);
    }

private:
    std::unique_ptr<matrix_type> m_matrix{};
};

} // namespace Opm::gpuistl

#endif