AmgxInterface.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_AMGX_INTERFACE_HPP
#define OPM_AMGX_INTERFACE_HPP

#include <opm/common/ErrorMacros.hpp>
#include <opm/simulators/linalg/gpuistl/GpuSparseMatrixWrapper.hpp>
#include <opm/simulators/linalg/gpuistl/GpuVector.hpp>
#include <opm/simulators/linalg/gpuistl/detail/gpu_safe_call.hpp>
#include <opm/simulators/linalg/gpuistl/detail/gpu_type_detection.hpp>

#include <amgx_c.h>
#include <cuda_runtime.h>

#include <fmt/core.h>

#include <string>
#include <vector>

namespace Opm::gpuistl
{

class AmgxError : public std::runtime_error
{
public:
    explicit AmgxError(const std::string& msg)
        : std::runtime_error(msg)
    {
    }
};

inline std::string
getAmgxErrorMessage(
    AMGX_RC err, const std::string& expression, const std::string& file, const std::string& function, int line)
{
    char amgx_err_msg[4096];
    AMGX_get_error_string(err, amgx_err_msg, sizeof(amgx_err_msg));

    return fmt::format(
        "AMGX error in expression: {}\nError code: {}\nError message: {}\nLocation: {}:{} in function {}",
        expression,
        static_cast<int>(err),
        amgx_err_msg,
        file,
        line,
        function);
}

inline void
amgxSafeCall(AMGX_RC rc, const std::string& expression, const std::string& file, const std::string& function, int line)
{
    if (rc != AMGX_RC_OK) {
        throw AmgxError(getAmgxErrorMessage(rc, expression, file, function, line));
    }
}

#define OPM_AMGX_SAFE_CALL(expr) ::Opm::gpuistl::amgxSafeCall((expr), #expr, __FILE__, __func__, __LINE__)


class AmgxInterface
{
public:
    static void initialize()
    {
        OPM_AMGX_SAFE_CALL(AMGX_initialize());
    }

    static void finalize()
    {
        OPM_AMGX_SAFE_CALL(AMGX_finalize());
    }

    static AMGX_config_handle createConfig(const std::string& config_string)
    {
        AMGX_config_handle config;
        OPM_AMGX_SAFE_CALL(AMGX_config_create(&config, config_string.c_str()));
        return config;
    }

    static AMGX_resources_handle createResources(AMGX_config_handle config)
    {
        AMGX_resources_handle resources;
        OPM_AMGX_SAFE_CALL(AMGX_resources_create_simple(&resources, config));
        return resources;
    }

    static AMGX_solver_handle createSolver(AMGX_resources_handle resources, AMGX_Mode mode, AMGX_config_handle config)
    {
        AMGX_solver_handle solver;
        OPM_AMGX_SAFE_CALL(AMGX_solver_create(&solver, resources, mode, config));
        return solver;
    }

    static AMGX_matrix_handle createMatrix(AMGX_resources_handle resources, AMGX_Mode mode)
    {
        AMGX_matrix_handle matrix;
        OPM_AMGX_SAFE_CALL(AMGX_matrix_create(&matrix, resources, mode));
        return matrix;
    }

    static AMGX_vector_handle createVector(AMGX_resources_handle resources, AMGX_Mode mode)
    {
        AMGX_vector_handle vector;
        OPM_AMGX_SAFE_CALL(AMGX_vector_create(&vector, resources, mode));
        return vector;
    }

    static void destroyConfig(AMGX_config_handle config)
    {
        if (config) {
            OPM_AMGX_SAFE_CALL(AMGX_config_destroy(config));
        }
    }

    static void destroyResources(AMGX_resources_handle resources)
    {
        if (resources) {
            OPM_AMGX_SAFE_CALL(AMGX_resources_destroy(resources));
        }
    }

    static void destroySolver(AMGX_solver_handle solver)
    {
        if (solver) {
            OPM_AMGX_SAFE_CALL(AMGX_solver_destroy(solver));
        }
    }

    template <typename MatrixType>
    static void destroyMatrix(AMGX_matrix_handle amgx_matrix, const MatrixType& matrix)
    {
        if constexpr (!is_gpu_type<MatrixType>::value) {
            // Unpin memory for CPU matrices
            using T = typename MatrixType::field_type;
            const T* values = &(matrix[0][0][0][0]);
            OPM_AMGX_SAFE_CALL(AMGX_unpin_memory(const_cast<T*>(values)));
        }

        if (amgx_matrix) {
            OPM_AMGX_SAFE_CALL(AMGX_matrix_destroy(amgx_matrix));
        }
    }

    static void destroyVector(AMGX_vector_handle vector)
    {
        if (vector) {
            OPM_AMGX_SAFE_CALL(AMGX_vector_destroy(vector));
        }
    }

    template <typename T>
    static void updateAmgxFromGpuVector(const GpuVector<T>& gpu_vec, AMGX_vector_handle amgx_vec)
    {
        // Get vector size from AMGX to verify compatibility
        int n, block_dim;
        OPM_AMGX_SAFE_CALL(AMGX_vector_get_size(amgx_vec, &n, &block_dim));

        if (n > 0 && static_cast<size_t>(n * block_dim) != gpu_vec.dim()) {
            throw AmgxError(fmt::format("Vector size mismatch in updateAmgxFromGpuVector: "
                                        "AMGX vector size {} vs. GpuVector size {}",
                                        n * block_dim,
                                        gpu_vec.dim()));
        }

        // Get raw device pointer directly from GpuVector
        const T* device_ptr = gpu_vec.data();

        // Update AMGX vector directly with the device pointer
        OPM_AMGX_SAFE_CALL(AMGX_vector_upload(amgx_vec, n, block_dim, const_cast<T*>(device_ptr)));
    }

    template <typename T>
    static void updateGpuVectorFromAmgx(AMGX_vector_handle amgx_vec, GpuVector<T>& gpu_vec)
    {
        // Get vector size from AMGX to verify compatibility
        int n, block_dim;
        OPM_AMGX_SAFE_CALL(AMGX_vector_get_size(amgx_vec, &n, &block_dim));

        if (static_cast<size_t>(n * block_dim) != gpu_vec.dim()) {
            throw AmgxError(fmt::format("Vector size mismatch in updateGpuVectorFromAmgx: "
                                        "AMGX vector size {} vs. GpuVector size {}",
                                        n * block_dim,
                                        gpu_vec.dim()));
        }

        // Get destination device pointer from GpuVector
        T* dst_device_ptr = gpu_vec.data();

        // Download data directly from AMGX vector to GpuVector's device memory
        OPM_AMGX_SAFE_CALL(AMGX_vector_download(amgx_vec, dst_device_ptr));
    }

    template <typename VectorType>
    static void transferVectorToAmgx(const VectorType& vec, AMGX_vector_handle amgx_vec)
    {
        if constexpr (is_gpu_type<VectorType>::value) {
            // GPU implementation - device-to-device transfer
            updateAmgxFromGpuVector(vec, amgx_vec);
        } else {
            // CPU implementation - host-to-device transfer
            // For BlockVector, get dimensions
            const int N = vec.size();
            const int block_size = 1; // Assumes scalar vectors for CPU case

            // Upload vector to AMGX
            OPM_AMGX_SAFE_CALL(AMGX_vector_upload(amgx_vec, N, block_size, &vec[0][0]));
        }
    }

    template <typename VectorType>
    static void transferVectorFromAmgx(AMGX_vector_handle amgx_vec, VectorType& vec)
    {
        if constexpr (is_gpu_type<VectorType>::value) {
            // GPU implementation - device-to-device transfer
            updateGpuVectorFromAmgx(amgx_vec, vec);
        } else {
            // CPU implementation - device-to-host transfer
            OPM_AMGX_SAFE_CALL(AMGX_vector_download(amgx_vec, &vec[0][0]));
        }
    }

    template <typename T>
    static void updateAmgxMatrixFromGpuSparseMatrix(const GpuSparseMatrixWrapper<T>& gpuSparseMatrix,
                                                    AMGX_matrix_handle amgxMatrix)
    {
        // Get matrix dimensions and sparsity information
        auto n = detail::to_int(gpuSparseMatrix.N());
        auto nnz = detail::to_int(gpuSparseMatrix.nonzeroes());
        auto block_size = detail::to_int(gpuSparseMatrix.blockSize());

        // Get device pointers directly from GpuSparseMatrix
        const T* values = gpuSparseMatrix.getNonZeroValues().data();
        const int* row_ptrs = gpuSparseMatrix.getRowIndices().data();
        const int* col_indices = gpuSparseMatrix.getColumnIndices().data();

        // Update AMGX matrix with the device pointers
        OPM_AMGX_SAFE_CALL(
            AMGX_matrix_upload_all(amgxMatrix, n, nnz, block_size, block_size, row_ptrs, col_indices, values, nullptr));
    }

    template <typename T>
    static void updateAmgxMatrixCoefficientsFromGpuSparseMatrix(const GpuSparseMatrixWrapper<T>& gpuSparseMatrix,
                                                                AMGX_matrix_handle amgxMatrix)
    {
        // Get matrix dimensions and sparsity information
        auto n = detail::to_int(gpuSparseMatrix.N());
        auto nnz = detail::to_int(gpuSparseMatrix.nonzeroes());

        // Get device pointer to values directly from GpuSparseMatrix
        const T* values = gpuSparseMatrix.getNonZeroValues().data();

        // Update coefficients in AMGX matrix with the device pointer
        OPM_AMGX_SAFE_CALL(AMGX_matrix_replace_coefficients(amgxMatrix, n, nnz, values, nullptr));
    }

    template <typename T>
    static void updateGpuSparseMatrixFromAmgxMatrix(AMGX_matrix_handle amgxMatrix, GpuSparseMatrixWrapper<T>& gpuSparseMatrix)
    {
        // Get matrix dimensions from AMGX
        int n, nnz, block_sizex, block_sizey;
        OPM_AMGX_SAFE_CALL(AMGX_matrix_get_size(amgxMatrix, &n, &block_sizex, &block_sizey));
        OPM_AMGX_SAFE_CALL(AMGX_matrix_get_nnz(amgxMatrix, &nnz));

        // Allocate temporary device memory for row pointers and column indices (required by AMGX API)
        int* temp_row_ptrs;
        int* temp_col_indices;
        OPM_GPU_SAFE_CALL(cudaMalloc(&temp_row_ptrs, (n + 1) * sizeof(int)));
        OPM_GPU_SAFE_CALL(cudaMalloc(&temp_col_indices, nnz * sizeof(int)));

        // Get device pointer to values in the GpuSparseMatrix
        T* gpu_values = gpuSparseMatrix.getNonZeroValues().data();

        // AMGX requires a valid pointer for diagonal data (even if unused)
        void* diag_data_ptr = nullptr;

        // Download matrix values directly from AMGX to GpuSparseMatrix
        OPM_AMGX_SAFE_CALL(AMGX_matrix_download_all(amgxMatrix,
                                                    temp_row_ptrs,
                                                    temp_col_indices,
                                                    gpu_values,
                                                    &diag_data_ptr));

        // Clean up temporary device memory
        OPM_GPU_SAFE_CALL(cudaFree(temp_row_ptrs));
        OPM_GPU_SAFE_CALL(cudaFree(temp_col_indices));
    }

    template <typename MatrixType>
    static void initializeMatrix(const MatrixType& matrix, AMGX_matrix_handle amgx_matrix)
    {
        if constexpr (is_gpu_type<MatrixType>::value) {
            // GPU implementation - device-to-device transfer
            updateAmgxMatrixFromGpuSparseMatrix(matrix, amgx_matrix);
        } else {
            // CPU implementation - host-to-device transfer
            auto N = detail::to_int(matrix.N());
            auto nnz = detail::to_int(matrix.nonzeroes());
            const int block_size = 1; // Assumes scalar matrices for CPU case

            // Setup sparsity pattern
            std::vector<int> row_ptrs(N + 1);
            std::vector<int> col_indices(nnz);

            int pos = 0;
            row_ptrs[0] = 0;
            for (auto row = matrix.begin(); row != matrix.end(); ++row) {
                for (auto col = row->begin(); col != row->end(); ++col) {
                    col_indices[pos++] = col.index();
                }
                row_ptrs[row.index() + 1] = pos;
            }

            // Get values pointer
            using T = typename MatrixType::field_type;
            const T* values = &(matrix[0][0][0][0]);
            // Indexing explanation:
            // matrix[0]             - First row of the matrix
            //          [0]          - First block in that row
            //            [0]        - First row within the 1x1 block
            //              [0]      - First column within the 1x1 block

            // Pin memory for CPU-to-GPU transfers
            OPM_AMGX_SAFE_CALL(AMGX_pin_memory(const_cast<T*>(values), sizeof(T) * nnz * block_size * block_size));

            // Upload to AMGX
            OPM_AMGX_SAFE_CALL(AMGX_matrix_upload_all(amgx_matrix,
                                                      N,
                                                      nnz,
                                                      block_size,
                                                      block_size,
                                                      row_ptrs.data(),
                                                      col_indices.data(),
                                                      const_cast<T*>(values),
                                                      nullptr));
        }
    }

    static void initializeVector(int N, int block_size, AMGX_vector_handle amgx_vector)
    {
        OPM_AMGX_SAFE_CALL(AMGX_vector_set_zero(amgx_vector, N, block_size));
    }

    template <typename MatrixType>
    static void updateMatrixValues(const MatrixType& matrix, AMGX_matrix_handle amgx_matrix)
    {
        if constexpr (is_gpu_type<MatrixType>::value) {
            // GPU implementation - device-to-device transfer
            updateAmgxMatrixCoefficientsFromGpuSparseMatrix(matrix, amgx_matrix);
        } else {
            // CPU implementation - host-to-device transfer
            using T = typename MatrixType::field_type;
            const T* values = &(matrix[0][0][0][0]);
            OPM_AMGX_SAFE_CALL(AMGX_matrix_replace_coefficients(
                amgx_matrix, matrix.N(), matrix.nonzeroes(), const_cast<T*>(values), nullptr));
        }
    }

    template <typename MatrixFieldType, typename VectorFieldType>
    static AMGX_Mode determineAmgxMode()
    {
        if constexpr (std::is_same_v<MatrixFieldType, double> && std::is_same_v<VectorFieldType, double>) {
            return AMGX_mode_dDDI;
        } else if constexpr (std::is_same_v<MatrixFieldType, float> && std::is_same_v<VectorFieldType, double>) {
            return AMGX_mode_dDFI;
        } else if constexpr (std::is_same_v<MatrixFieldType, float> && std::is_same_v<VectorFieldType, float>) {
            return AMGX_mode_dFFI;
        } else {
            OPM_THROW(std::runtime_error,
                      "Unsupported combination of matrix and vector types for AMGX: "
                          + std::string(typeid(MatrixFieldType).name()) + " and "
                          + std::string(typeid(VectorFieldType).name()));
        }
    }
};

} // namespace Opm::gpuistl

#endif // OPM_AMGX_INTERFACE_HPP