AmgxPreconditioner.hpp

Go to the documentation of this file.

/*
  Copyright 2024 SINTEF AS
  Copyright 2024 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_PRECONDITIONER_HEADER_INCLUDED
#define OPM_AMGX_PRECONDITIONER_HEADER_INCLUDED
 
#include <opm/common/ErrorMacros.hpp>
#include <opm/common/TimingMacros.hpp>
#include <opm/simulators/linalg/PreconditionerWithUpdate.hpp>
#include <opm/simulators/linalg/PropertyTree.hpp>
 
#include <dune/common/fmatrix.hh>
#include <dune/istl/bcrsmatrix.hh>
 
#include <amgx_c.h>
 
#include <vector>
 
namespace Amgx {
 
struct AmgxConfig {
    int determinism_flag = 0;
    int print_grid_stats = 0;
    int print_solve_stats = 0;
    std::string solver = "AMG";
    std::string algorithm = "CLASSICAL";
    std::string interpolator = "D2";
    std::string selector = "PMIS";
    std::string smoother = "BLOCK_JACOBI";
    int presweeps = 3;
    int postsweeps = 3;
    double strength_threshold = 0.5;
    int max_iters = 1;
 
    explicit AmgxConfig(const Opm::PropertyTree& prm) {
        determinism_flag = prm.get<int>("determinism_flag", determinism_flag);
        print_grid_stats = prm.get<int>("print_grid_stats", print_grid_stats);
        print_solve_stats = prm.get<int>("print_solve_stats", print_solve_stats);
        solver = prm.get<std::string>("solver", solver);
        algorithm = prm.get<std::string>("algorithm", algorithm);
        interpolator = prm.get<std::string>("interpolator", interpolator);
        selector = prm.get<std::string>("selector", selector);
        smoother = prm.get<std::string>("smoother", smoother);
        presweeps = prm.get<int>("presweeps", presweeps);
        postsweeps = prm.get<int>("postsweeps", postsweeps);
        strength_threshold = prm.get<double>("strength_threshold", strength_threshold);
        max_iters = prm.get<int>("max_iters", max_iters);
    }
 
    std::string toString() const {
        return "config_version=2, "
               "determinism_flag=" + std::to_string(determinism_flag) + ", "
               "print_grid_stats=" + std::to_string(print_grid_stats) + ", "
               "print_solve_stats=" + std::to_string(print_solve_stats) + ", "
               "solver=" + solver + ", "
               "algorithm=" + algorithm + ", "
               "interpolator=" + interpolator + ", "
               "selector=" + selector + ", "
               "smoother=" + smoother + ", "
               "presweeps=" + std::to_string(presweeps) + ", "
               "postsweeps=" + std::to_string(postsweeps) + ", "
               "strength_threshold=" + std::to_string(strength_threshold) + ", "
               "max_iters=" + std::to_string(max_iters);
    }
};
 
template<class M, class X, class Y>
class AmgxPreconditioner : public Dune::PreconditionerWithUpdate<X,Y>
{
public:
    using matrix_type = M;
    using matrix_field_type = typename M::field_type;
    using domain_type = X;
    using range_type = Y;
    using vector_field_type = typename X::field_type;
 
    static constexpr int block_size = 1;
 
    AmgxPreconditioner(const M& A, const Opm::PropertyTree prm)
        : A_(A)
        , N_(A.N())
        , nnz_(A.nonzeroes())
    {
        OPM_TIMEBLOCK(prec_construct);
 
        // Create configuration
        AmgxConfig config(prm);
        AMGX_SAFE_CALL(AMGX_config_create(&cfg_, config.toString().c_str()));
        AMGX_SAFE_CALL(AMGX_resources_create_simple(&rsrc_, cfg_));
 
        // Setup frequency is set in the property tree
        setup_frequency_ = prm.get<int>("setup_frequency", 30);
 
        // Select appropriate AMGX mode based on matrix and vector scalar types
        AMGX_Mode amgx_mode;
        if constexpr (std::is_same_v<matrix_field_type, double> && std::is_same_v<vector_field_type, double>) {
            amgx_mode = AMGX_mode_dDDI;
        } else if constexpr (std::is_same_v<matrix_field_type, float> && std::is_same_v<vector_field_type, double>) {
            amgx_mode = AMGX_mode_dDFI;
        } else if constexpr (std::is_same_v<matrix_field_type, float> && std::is_same_v<vector_field_type, float>) {
            amgx_mode = AMGX_mode_dFFI;
        } else {
            OPM_THROW(std::runtime_error, "Unsupported combination of matrix and vector types in AmgxPreconditioner");
        }
 
        // Create solver and matrix/vector handles with selected mode
        AMGX_SAFE_CALL(AMGX_solver_create(&solver_, rsrc_, amgx_mode, cfg_));
        AMGX_SAFE_CALL(AMGX_matrix_create(&A_amgx_, rsrc_, amgx_mode));
        AMGX_SAFE_CALL(AMGX_vector_create(&x_amgx_, rsrc_, amgx_mode));
        AMGX_SAFE_CALL(AMGX_vector_create(&b_amgx_, rsrc_, amgx_mode));
 
        // Setup matrix structure
        std::vector<int> row_ptrs(N_ + 1);
        std::vector<int> col_indices(nnz_);
        setupSparsityPattern(row_ptrs, col_indices);
 
        // initialize matrix with values
        const matrix_field_type* values = &(A_[0][0][0][0]);
        AMGX_SAFE_CALL(AMGX_pin_memory(const_cast<matrix_field_type*>(values), sizeof(matrix_field_type) * nnz_ * block_size * block_size));
        AMGX_SAFE_CALL(AMGX_matrix_upload_all(A_amgx_, N_, nnz_, block_size, block_size,
                                             row_ptrs.data(), col_indices.data(),
                                             values, nullptr));
        update();
    }
 
    ~AmgxPreconditioner()
    {
        const matrix_field_type* values = &(A_[0][0][0][0]);
        AMGX_SAFE_CALL(AMGX_unpin_memory(const_cast<matrix_field_type*>(values)));
        if (solver_) {
            AMGX_SAFE_CALL(AMGX_solver_destroy(solver_));
        }
        if (x_amgx_) {
            AMGX_SAFE_CALL(AMGX_vector_destroy(x_amgx_));
        }
        if (b_amgx_) {
            AMGX_SAFE_CALL(AMGX_vector_destroy(b_amgx_));
        }
        if (A_amgx_) {
            AMGX_SAFE_CALL(AMGX_matrix_destroy(A_amgx_));
        }
        // Destroying resources and config crashes when reinitializing
        //if (rsrc_) {
        //    AMGX_SAFE_CALL(AMGX_resources_destroy(rsrc_));
        //}
        //if (cfg_) {
        //    AMGX_SAFE_CALL(AMGX_config_destroy(cfg_));
        //}
    }
 
    void pre(X& /*v*/, Y& /*d*/) override {
    }
 
    void apply(X& v, const Y& d) override
    {
        OPM_TIMEBLOCK(prec_apply);
 
        // Upload vectors to AMGX
        AMGX_SAFE_CALL(AMGX_vector_upload(x_amgx_, N_, block_size, &v[0][0]));
        AMGX_SAFE_CALL(AMGX_vector_upload(b_amgx_, N_, block_size, &d[0][0]));
 
        // Apply preconditioner
        AMGX_SAFE_CALL(AMGX_solver_solve(solver_, b_amgx_, x_amgx_));
 
        // Download result
        AMGX_SAFE_CALL(AMGX_vector_download(x_amgx_, &v[0][0]));
    }
 
    void post(X& /*v*/) override {
    }
 
    void update() override
    {
        OPM_TIMEBLOCK(prec_update);
        copyMatrixToAmgx();
        if (update_counter_ == 0) {
            AMGX_SAFE_CALL(AMGX_solver_setup(solver_, A_amgx_));
        } else {
            AMGX_SAFE_CALL(AMGX_solver_resetup(solver_, A_amgx_));
        }
 
        ++update_counter_;
        if (update_counter_ >= setup_frequency_) {
            update_counter_ = 0;
        }
    }
 
    Dune::SolverCategory::Category category() const override
    {
        return Dune::SolverCategory::sequential;
    }
 
    bool hasPerfectUpdate() const override
    {
        // The AMG hierarchy of the Amgx preconditioner can depend on the values of the matrix, so it must be recreated
        // when the matrix values change, at given frequency. Since this is handled internally, we return true.
        return true;
    }
 
private:
    void setupSparsityPattern(std::vector<int>& row_ptrs, std::vector<int>& col_indices)
    {
        int pos = 0;
        row_ptrs[0] = 0;
        for (auto row = A_.begin(); row != A_.end(); ++row) {
            for (auto col = row->begin(); col != row->end(); ++col) {
                col_indices[pos++] = col.index();
            }
            row_ptrs[row.index() + 1] = pos;
        }
    }
 
    void copyMatrixToAmgx()
    {
        // Get direct pointer to matrix values
        const matrix_field_type* values = &(A_[0][0][0][0]);
        // Indexing explanation:
        // A_[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
        // update matrix with new values, assuming the sparsity structure is the same
        AMGX_SAFE_CALL(AMGX_matrix_replace_coefficients(A_amgx_, N_, nnz_, values, nullptr));
    }
 
    const M& A_; 
    const int N_; 
    const int nnz_; 
    // Internal variables to control AMGX setup and reuse frequency
    int setup_frequency_ = -1; 
    int update_counter_ = 0; 
 
    AMGX_config_handle cfg_ = nullptr; 
    AMGX_resources_handle rsrc_ = nullptr; 
    AMGX_solver_handle solver_ = nullptr; 
    AMGX_matrix_handle A_amgx_ = nullptr; 
    AMGX_vector_handle x_amgx_ = nullptr; 
    AMGX_vector_handle b_amgx_ = nullptr; 
};
 
} // namespace Amgx
 
#endif // OPM_AMGX_PRECONDITIONER_HEADER_INCLUDED