GpuSparseMatrix.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_GPUSPARSEMATRIX_HPP
#define OPM_GPUSPARSEMATRIX_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/GpuSparseMatrixGeneric.hpp>

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

namespace Opm::gpuistl
{

template <typename T>
class GpuSparseMatrix

{
public:
    using field_type = T;

    static constexpr int max_block_size = 6;

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

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

    GpuSparseMatrix(const GpuSparseMatrix&);

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

    virtual ~GpuSparseMatrix();

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

    void setUpperTriangular();

    void setLowerTriangular();

    void setUnitDiagonal();


    void setNonUnitDiagonal();

    std::size_t N() 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_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()
    {
        if (m_genericMatrixForBlockSize1) {
            return m_genericMatrixForBlockSize1->getNonZeroValues();
        }
        return m_nonZeroElements;
    }

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

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

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

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

    const GpuVector<int>& getColumnIndices() const
    {
        if (m_genericMatrixForBlockSize1) {
            return m_genericMatrixForBlockSize1->getColumnIndices();
        }
        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);
    }

    detail::GpuSparseMatrixDescription& getDescription()
    {
        return *m_matrixDescription;
    }

    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 GpuSparseMatrix<T>& matrix);

     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;

    detail::GpuSparseMatrixDescriptionPtr m_matrixDescription;
    detail::CuSparseHandle& m_cusparseHandle;

    // For blockSize == 1, we use the generic API
    std::unique_ptr<GpuSparseMatrixGeneric<T>> m_genericMatrixForBlockSize1;

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

    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