GpuSparseMatrixGeneric.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_GPUSPARSEMATRIXGENERIC_HPP
#define OPM_GPUSPARSEMATRIXGENERIC_HPP

#include <opm/common/ErrorMacros.hpp>
#include <opm/simulators/linalg/gpuistl/detail/CuSparseHandle.hpp>
#include <opm/simulators/linalg/gpuistl/detail/cusparse_safe_call.hpp>
#include <opm/simulators/linalg/gpuistl/detail/gpusparse_matrix_utilities.hpp>
#include <opm/simulators/linalg/gpuistl/detail/safe_conversion.hpp>
#include <opm/simulators/linalg/gpuistl/GpuBuffer.hpp>
#include <opm/simulators/linalg/gpuistl/GpuVector.hpp>

#include <cstddef>
#include <cusparse.h>

namespace Opm::gpuistl
{

template <typename T>
class GpuSparseMatrixGeneric
{
public:
    using field_type = T;

    static constexpr int max_block_size = 6;

    GpuSparseMatrixGeneric(const T* nonZeroElements,
                           const int* rowIndices,
                           const int* columnIndices,
                           std::size_t numberOfNonzeroBlocks,
                           std::size_t blockSize,
                           std::size_t numberOfRows);

    GpuSparseMatrixGeneric(const GpuVector<int>& rowIndices, const GpuVector<int>& columnIndices, std::size_t blockSize);

    GpuSparseMatrixGeneric(const GpuSparseMatrixGeneric&);

    void preprocessSpMV();

    // 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.
    GpuSparseMatrixGeneric& operator=(const GpuSparseMatrixGeneric&) = delete;

    virtual ~GpuSparseMatrixGeneric();

    template <class MatrixType>
    static GpuSparseMatrixGeneric<T> fromMatrix(const MatrixType& matrix, bool copyNonZeroElementsDirectly = false);

    std::size_t N() const
    {
        return detail::to_size_t(m_numberOfRows);
    }

    std::size_t nonzeroes() const
    {
        // Technically this safe conversion is not needed since we enforce these to be
        // non-negative in the constructor, but keeping them for added sanity for now.
        //
        // We don't believe this will yield any performance penality (it's used too far away from the inner loop),
        // but should that be false, they can be removed.
        return detail::to_size_t(m_numberOfNonzeroBlocks);
    }

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

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

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

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

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

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

    std::size_t dim() const
    {
        // Technically this safe conversion is not needed since we enforce these to be
        // non-negative in the constructor, but keeping them for added sanity for now.
        //
        // We don't believe this will yield any performance penality (it's used too far away from the inner loop),
        // but should that be false, they can be removed.
        return detail::to_size_t(m_blockSize) * detail::to_size_t(m_numberOfRows);
    }

    std::size_t blockSize() const
    {
        // Technically this safe conversion is not needed since we enforce these to be
        // non-negative in the constructor, but keeping them for added sanity for now.
        //
        // We don't believe this will yield any performance penality (it's used too far away from the inner loop),
        // but should that be false, they can be removed.
        return detail::to_size_t(m_blockSize);
    }

    virtual void mv(const GpuVector<T>& x, GpuVector<T>& y) const;

    virtual void umv(const GpuVector<T>& x, GpuVector<T>& y) const;


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

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

    void updateNonzeroValues(const GpuSparseMatrixGeneric<T>& matrix);

    void setToZero();

     template<class FunctionType>
     auto dispatchOnBlocksize(FunctionType function) const
     {
        return dispatchOnBlocksizeImpl<max_block_size>(function);
     }

private:
    GpuVector<T> m_nonZeroElements;
    GpuVector<int> m_columnIndices;
    GpuVector<int> m_rowIndices;

    // Notice that we store these three as int to make sure we are cusparse compatible.
    //
    // This gives the added benefit of checking the size constraints at construction of the matrix
    // rather than in some call to cusparse.
    const int m_numberOfNonzeroBlocks;
    const int m_numberOfRows;
    const int m_blockSize;

    // Generic API descriptors
    decltype(detail::makeSafeMatrixDescriptor()) m_matrixDescriptor;
    detail::CuSparseHandle& m_cusparseHandle;

    // Cached buffer for operations
    mutable GpuBuffer<std::byte> m_buffer;

    // Helper methods for SpMV operations
    void spMV(T alpha, const GpuVector<T>& x, T beta, GpuVector<T>& y) const;

    // Initialize matrix descriptor based on block size
    void initializeMatrixDescriptor();

    template <class VectorType>
    void assertSameSize(const VectorType& vector) const;

    // Helper to get cuSPARSE data type from C++ type
    constexpr cudaDataType getDataType() const
    {
        if constexpr (std::is_same_v<T, float>) {
            return CUDA_R_32F;
        } else if constexpr (std::is_same_v<T, double>) {
            return CUDA_R_64F;
        } else {
            static_assert(std::is_same_v<T, float> || std::is_same_v<T, double>, "Only float and double are supported");
            return CUDA_R_32F; // Unreachable, but needed to compile
        }
    }

    template<int blockSizeCompileTime, class FunctionType>
    auto dispatchOnBlocksizeImpl(FunctionType function) const
    {
        if (blockSizeCompileTime == m_blockSize) {
            return function(std::integral_constant<int, blockSizeCompileTime>());
        }

        if constexpr (blockSizeCompileTime > 1) {
            return dispatchOnBlocksizeImpl<blockSizeCompileTime - 1>(function);
        } else {
            OPM_THROW(std::runtime_error, fmt::format(fmt::runtime("Unsupported block size: {}"), m_blockSize));
        }
    }
};
} // namespace Opm::gpuistl
#endif