Date: 2025-03-28T21:10:55+00:00

Actual source code: amgx.cxx

      1: /*
      2:      This file implements an AmgX preconditioner in PETSc as part of PC.
      3:  */

      5: /*
      6:    Include files needed for the AmgX preconditioner:
      7:      pcimpl.h - private include file intended for use by all preconditioners
      8: */

     10: #include <petsc/private/pcimpl.h>
     11: #include <petscdevice_cuda.h>
     12: #include <amgx_c.h>
     13: #include <limits>
     14: #include <vector>
     15: #include <algorithm>
     16: #include <map>
     17: #include <numeric>
     18: #include "cuda_runtime.h"

     20: enum class AmgXSmoother {
     21:   PCG,
     22:   PCGF,
     23:   PBiCGStab,
     24:   GMRES,
     25:   FGMRES,
     26:   JacobiL1,
     27:   BlockJacobi,
     28:   GS,
     29:   MulticolorGS,
     30:   MulticolorILU,
     31:   MulticolorDILU,
     32:   ChebyshevPoly,
     33:   NoSolver
     34: };
     35: enum class AmgXAMGMethod {
     36:   Classical,
     37:   Aggregation
     38: };
     39: enum class AmgXSelector {
     40:   Size2,
     41:   Size4,
     42:   Size8,
     43:   MultiPairwise,
     44:   PMIS,
     45:   HMIS
     46: };
     47: enum class AmgXCoarseSolver {
     48:   DenseLU,
     49:   NoSolver
     50: };
     51: enum class AmgXAMGCycle {
     52:   V,
     53:   W,
     54:   F,
     55:   CG,
     56:   CGF
     57: };

     59: struct AmgXControlMap {
     60:   static const std::map<std::string, AmgXAMGMethod>    AMGMethods;
     61:   static const std::map<std::string, AmgXSmoother>     Smoothers;
     62:   static const std::map<std::string, AmgXSelector>     Selectors;
     63:   static const std::map<std::string, AmgXCoarseSolver> CoarseSolvers;
     64:   static const std::map<std::string, AmgXAMGCycle>     AMGCycles;
     65: };

     67: const std::map<std::string, AmgXAMGMethod> AmgXControlMap::AMGMethods = {
     68:   {"CLASSICAL",   AmgXAMGMethod::Classical  },
     69:   {"AGGREGATION", AmgXAMGMethod::Aggregation}
     70: };

     72: const std::map<std::string, AmgXSmoother> AmgXControlMap::Smoothers = {
     73:   {"PCG",             AmgXSmoother::PCG           },
     74:   {"PCGF",            AmgXSmoother::PCGF          },
     75:   {"PBICGSTAB",       AmgXSmoother::PBiCGStab     },
     76:   {"GMRES",           AmgXSmoother::GMRES         },
     77:   {"FGMRES",          AmgXSmoother::FGMRES        },
     78:   {"JACOBI_L1",       AmgXSmoother::JacobiL1      },
     79:   {"BLOCK_JACOBI",    AmgXSmoother::BlockJacobi   },
     80:   {"GS",              AmgXSmoother::GS            },
     81:   {"MULTICOLOR_GS",   AmgXSmoother::MulticolorGS  },
     82:   {"MULTICOLOR_ILU",  AmgXSmoother::MulticolorILU },
     83:   {"MULTICOLOR_DILU", AmgXSmoother::MulticolorDILU},
     84:   {"CHEBYSHEV_POLY",  AmgXSmoother::ChebyshevPoly },
     85:   {"NOSOLVER",        AmgXSmoother::NoSolver      }
     86: };

     88: const std::map<std::string, AmgXSelector> AmgXControlMap::Selectors = {
     89:   {"SIZE_2",         AmgXSelector::Size2        },
     90:   {"SIZE_4",         AmgXSelector::Size4        },
     91:   {"SIZE_8",         AmgXSelector::Size8        },
     92:   {"MULTI_PAIRWISE", AmgXSelector::MultiPairwise},
     93:   {"PMIS",           AmgXSelector::PMIS         },
     94:   {"HMIS",           AmgXSelector::HMIS         }
     95: };

     97: const std::map<std::string, AmgXCoarseSolver> AmgXControlMap::CoarseSolvers = {
     98:   {"DENSE_LU_SOLVER", AmgXCoarseSolver::DenseLU },
     99:   {"NOSOLVER",        AmgXCoarseSolver::NoSolver}
    100: };

    102: const std::map<std::string, AmgXAMGCycle> AmgXControlMap::AMGCycles = {
    103:   {"V",   AmgXAMGCycle::V  },
    104:   {"W",   AmgXAMGCycle::W  },
    105:   {"F",   AmgXAMGCycle::F  },
    106:   {"CG",  AmgXAMGCycle::CG },
    107:   {"CGF", AmgXAMGCycle::CGF}
    108: };

    110: /*
    111:    Private context (data structure) for the AMGX preconditioner.
    112: */
    113: struct PC_AMGX {
    114:   AMGX_solver_handle    solver;
    115:   AMGX_config_handle    cfg;
    116:   AMGX_resources_handle rsrc;
    117:   bool                  solve_state_init;
    118:   bool                  rsrc_init;
    119:   PetscBool             verbose;

    121:   AMGX_matrix_handle A;
    122:   AMGX_vector_handle sol;
    123:   AMGX_vector_handle rhs;

    125:   MPI_Comm    comm;
    126:   PetscMPIInt rank   = 0;
    127:   PetscMPIInt nranks = 0;
    128:   int         devID  = 0;

    130:   void       *lib_handle = 0;
    131:   std::string cfg_contents;

    133:   // Cached state for re-setup
    134:   PetscInt           nnz;
    135:   PetscInt           nLocalRows;
    136:   PetscInt           nGlobalRows;
    137:   PetscInt           bSize;
    138:   Mat                localA;
    139:   const PetscScalar *values;

    141:   // AMG Control parameters
    142:   AmgXSmoother     smoother;
    143:   AmgXAMGMethod    amg_method;
    144:   AmgXSelector     selector;
    145:   AmgXCoarseSolver coarse_solver;
    146:   AmgXAMGCycle     amg_cycle;
    147:   PetscInt         presweeps;
    148:   PetscInt         postsweeps;
    149:   PetscInt         max_levels;
    150:   PetscInt         aggressive_levels;
    151:   PetscInt         dense_lu_num_rows;
    152:   PetscScalar      strength_threshold;
    153:   PetscBool        print_grid_stats;
    154:   PetscBool        exact_coarse_solve;

    156:   // Smoother control parameters
    157:   PetscScalar jacobi_relaxation_factor;
    158:   PetscScalar gs_symmetric;
    159: };

    161: static PetscInt s_count = 0;

    163: // Buffer of messages from AmgX
    164: // Currently necessary hack before we adapt AmgX to print from single rank only
    165: static std::string amgx_output{};

    167: // A print callback that allows AmgX to return status messages
    168: static void print_callback(const char *msg, int length)
    169: {
    170:   amgx_output.append(msg);
    171: }

    173: // Outputs messages from the AmgX message buffer and clears it
    174: static PetscErrorCode amgx_output_messages(PC_AMGX *amgx)
    175: {
    176:   PetscFunctionBegin;
    177:   // If AmgX output is enabled and we have a message, output it
    178:   if (amgx->verbose && !amgx_output.empty()) {
    179:     // Only a single rank to output the AmgX messages
    180:     PetscCall(PetscPrintf(amgx->comm, "AMGX: %s", amgx_output.c_str()));

    182:     // Note that all ranks clear their received output
    183:     amgx_output.clear();
    184:   }
    185:   PetscFunctionReturn(PETSC_SUCCESS);
    186: }

    188: // XXX Need to add call in AmgX API that gracefully destroys everything
    189: // without abort etc.
    190: #define PetscCallAmgX(rc) \
    191:   do { \
    192:     AMGX_RC err = (rc); \
    193:     char    msg[4096]; \
    194:     switch (err) { \
    195:     case AMGX_RC_OK: \
    196:       break; \
    197:     default: \
    198:       AMGX_get_error_string(err, msg, 4096); \
    199:       SETERRQ(amgx->comm, PETSC_ERR_LIB, "%s", msg); \
    200:     } \
    201:   } while (0)

    203: /*
    204:    PCSetUp_AMGX - Prepares for the use of the AmgX preconditioner
    205:                     by setting data structures and options.

    207:    Input Parameter:
    208: .  pc - the preconditioner context

    210:    Application Interface Routine: PCSetUp()

    212:    Note:
    213:    The interface routine PCSetUp() is not usually called directly by
    214:    the user, but instead is called by PCApply() if necessary.
    215: */
    216: static PetscErrorCode PCSetUp_AMGX(PC pc)
    217: {
    218:   PC_AMGX  *amgx = (PC_AMGX *)pc->data;
    219:   Mat       Pmat = pc->pmat;
    220:   PetscBool is_dev_ptrs;

    222:   PetscFunctionBegin;
    223:   PetscCall(PetscObjectTypeCompareAny((PetscObject)Pmat, &is_dev_ptrs, MATAIJCUSPARSE, MATSEQAIJCUSPARSE, MATMPIAIJCUSPARSE, ""));

    225:   // At the present time, an AmgX matrix is a sequential matrix
    226:   // Non-sequential/MPI matrices must be adapted to extract the local matrix
    227:   bool partial_setup_allowed = (pc->setupcalled && pc->flag != DIFFERENT_NONZERO_PATTERN);
    228:   if (amgx->nranks > 1) {
    229:     if (partial_setup_allowed) {
    230:       PetscCall(MatMPIAIJGetLocalMat(Pmat, MAT_REUSE_MATRIX, &amgx->localA));
    231:     } else {
    232:       PetscCall(MatMPIAIJGetLocalMat(Pmat, MAT_INITIAL_MATRIX, &amgx->localA));
    233:     }

    235:     if (is_dev_ptrs) PetscCall(MatConvert(amgx->localA, MATSEQAIJCUSPARSE, MAT_INPLACE_MATRIX, &amgx->localA));
    236:   } else {
    237:     amgx->localA = Pmat;
    238:   }

    240:   if (is_dev_ptrs) {
    241:     PetscCall(MatSeqAIJCUSPARSEGetArrayRead(amgx->localA, &amgx->values));
    242:   } else {
    243:     PetscCall(MatSeqAIJGetArrayRead(amgx->localA, &amgx->values));
    244:   }

    246:   if (!partial_setup_allowed) {
    247:     // Initialise resources and matrices
    248:     if (!amgx->rsrc_init) {
    249:       // Read configuration file
    250:       PetscCallAmgX(AMGX_config_create(&amgx->cfg, amgx->cfg_contents.c_str()));
    251:       PetscCallAmgX(AMGX_resources_create(&amgx->rsrc, amgx->cfg, &amgx->comm, 1, &amgx->devID));
    252:       amgx->rsrc_init = true;
    253:     }

    255:     PetscCheck(!amgx->solve_state_init, amgx->comm, PETSC_ERR_PLIB, "AmgX solve state initialisation already called.");
    256:     PetscCallAmgX(AMGX_matrix_create(&amgx->A, amgx->rsrc, AMGX_mode_dDDI));
    257:     PetscCallAmgX(AMGX_vector_create(&amgx->sol, amgx->rsrc, AMGX_mode_dDDI));
    258:     PetscCallAmgX(AMGX_vector_create(&amgx->rhs, amgx->rsrc, AMGX_mode_dDDI));
    259:     PetscCallAmgX(AMGX_solver_create(&amgx->solver, amgx->rsrc, AMGX_mode_dDDI, amgx->cfg));
    260:     amgx->solve_state_init = true;

    262:     // Extract the CSR data
    263:     PetscBool       done;
    264:     const PetscInt *colIndices;
    265:     const PetscInt *rowOffsets;
    266:     PetscCall(MatGetRowIJ(amgx->localA, 0, PETSC_FALSE, PETSC_FALSE, &amgx->nLocalRows, &rowOffsets, &colIndices, &done));
    267:     PetscCheck(done, amgx->comm, PETSC_ERR_PLIB, "MatGetRowIJ was not successful");
    268:     PetscCheck(amgx->nLocalRows < std::numeric_limits<int>::max(), PETSC_COMM_SELF, PETSC_ERR_PLIB, "AmgX restricted to int local rows but nLocalRows = %" PetscInt_FMT " > max<int>", amgx->nLocalRows);

    270:     if (is_dev_ptrs) {
    271:       PetscCallCUDA(cudaMemcpy(&amgx->nnz, &rowOffsets[amgx->nLocalRows], sizeof(int), cudaMemcpyDefault));
    272:     } else {
    273:       amgx->nnz = rowOffsets[amgx->nLocalRows];
    274:     }

    276:     PetscCheck(amgx->nnz < std::numeric_limits<int>::max(), PETSC_COMM_SELF, PETSC_ERR_PLIB, "Support for 64-bit integer nnz not yet implemented, nnz = %" PetscInt_FMT ".", amgx->nnz);

    278:     // Allocate space for some partition offsets
    279:     std::vector<PetscInt> partitionOffsets(amgx->nranks + 1);

    281:     // Fetch the number of local rows per rank
    282:     partitionOffsets[0] = 0; /* could use PetscLayoutGetRanges */
    283:     PetscCallMPI(MPI_Allgather(&amgx->nLocalRows, 1, MPIU_INT, partitionOffsets.data() + 1, 1, MPIU_INT, amgx->comm));
    284:     std::partial_sum(partitionOffsets.begin(), partitionOffsets.end(), partitionOffsets.begin());

    286:     // Fetch the number of global rows
    287:     amgx->nGlobalRows = partitionOffsets[amgx->nranks];

    289:     PetscCall(MatGetBlockSize(Pmat, &amgx->bSize));

    291:     // XXX Currently constrained to 32-bit indices, to be changed in the future
    292:     // Create the distribution and upload the matrix data
    293:     AMGX_distribution_handle dist;
    294:     PetscCallAmgX(AMGX_distribution_create(&dist, amgx->cfg));
    295:     PetscCallAmgX(AMGX_distribution_set_32bit_colindices(dist, true));
    296:     PetscCallAmgX(AMGX_distribution_set_partition_data(dist, AMGX_DIST_PARTITION_OFFSETS, partitionOffsets.data()));
    297:     PetscCallAmgX(AMGX_matrix_upload_distributed(amgx->A, amgx->nGlobalRows, (int)amgx->nLocalRows, (int)amgx->nnz, amgx->bSize, amgx->bSize, rowOffsets, colIndices, amgx->values, NULL, dist));
    298:     PetscCallAmgX(AMGX_solver_setup(amgx->solver, amgx->A));
    299:     PetscCallAmgX(AMGX_vector_bind(amgx->sol, amgx->A));
    300:     PetscCallAmgX(AMGX_vector_bind(amgx->rhs, amgx->A));

    302:     PetscInt nlr = 0;
    303:     PetscCall(MatRestoreRowIJ(amgx->localA, 0, PETSC_FALSE, PETSC_FALSE, &nlr, &rowOffsets, &colIndices, &done));
    304:   } else {
    305:     // The fast path for if the sparsity pattern persists
    306:     PetscCallAmgX(AMGX_matrix_replace_coefficients(amgx->A, amgx->nLocalRows, amgx->nnz, amgx->values, NULL));
    307:     PetscCallAmgX(AMGX_solver_resetup(amgx->solver, amgx->A));
    308:   }

    310:   if (is_dev_ptrs) {
    311:     PetscCall(MatSeqAIJCUSPARSERestoreArrayRead(amgx->localA, &amgx->values));
    312:   } else {
    313:     PetscCall(MatSeqAIJRestoreArrayRead(amgx->localA, &amgx->values));
    314:   }
    315:   PetscCall(amgx_output_messages(amgx));
    316:   PetscFunctionReturn(PETSC_SUCCESS);
    317: }

    319: /*
    320:    PCApply_AMGX - Applies the AmgX preconditioner to a vector.

    322:    Input Parameters:
    323: .  pc - the preconditioner context
    324: .  b - rhs vector

    326:    Output Parameter:
    327: .  x - solution vector

    329:    Application Interface Routine: PCApply()
    330:  */
    331: static PetscErrorCode PCApply_AMGX(PC pc, Vec b, Vec x)
    332: {
    333:   PC_AMGX           *amgx = (PC_AMGX *)pc->data;
    334:   PetscScalar       *x_;
    335:   const PetscScalar *b_;
    336:   PetscBool          is_dev_ptrs;

    338:   PetscFunctionBegin;
    339:   PetscCall(PetscObjectTypeCompareAny((PetscObject)x, &is_dev_ptrs, VECCUDA, VECMPICUDA, VECSEQCUDA, ""));

    341:   if (is_dev_ptrs) {
    342:     PetscCall(VecCUDAGetArrayWrite(x, &x_));
    343:     PetscCall(VecCUDAGetArrayRead(b, &b_));
    344:   } else {
    345:     PetscCall(VecGetArrayWrite(x, &x_));
    346:     PetscCall(VecGetArrayRead(b, &b_));
    347:   }

    349:   PetscCallAmgX(AMGX_vector_upload(amgx->sol, amgx->nLocalRows, 1, x_));
    350:   PetscCallAmgX(AMGX_vector_upload(amgx->rhs, amgx->nLocalRows, 1, b_));
    351:   PetscCallAmgX(AMGX_solver_solve_with_0_initial_guess(amgx->solver, amgx->rhs, amgx->sol));

    353:   AMGX_SOLVE_STATUS status;
    354:   PetscCallAmgX(AMGX_solver_get_status(amgx->solver, &status));
    355:   PetscCall(PCSetErrorIfFailure(pc, static_cast<PetscBool>(status == AMGX_SOLVE_FAILED)));
    356:   PetscCheck(status != AMGX_SOLVE_FAILED, amgx->comm, PETSC_ERR_CONV_FAILED, "AmgX solver failed to solve the system! The error code is %d.", status);
    357:   PetscCallAmgX(AMGX_vector_download(amgx->sol, x_));

    359:   if (is_dev_ptrs) {
    360:     PetscCall(VecCUDARestoreArrayWrite(x, &x_));
    361:     PetscCall(VecCUDARestoreArrayRead(b, &b_));
    362:   } else {
    363:     PetscCall(VecRestoreArrayWrite(x, &x_));
    364:     PetscCall(VecRestoreArrayRead(b, &b_));
    365:   }
    366:   PetscCall(amgx_output_messages(amgx));
    367:   PetscFunctionReturn(PETSC_SUCCESS);
    368: }

    370: static PetscErrorCode PCReset_AMGX(PC pc)
    371: {
    372:   PC_AMGX *amgx = (PC_AMGX *)pc->data;

    374:   PetscFunctionBegin;
    375:   if (amgx->solve_state_init) {
    376:     PetscCallAmgX(AMGX_solver_destroy(amgx->solver));
    377:     PetscCallAmgX(AMGX_matrix_destroy(amgx->A));
    378:     PetscCallAmgX(AMGX_vector_destroy(amgx->sol));
    379:     PetscCallAmgX(AMGX_vector_destroy(amgx->rhs));
    380:     if (amgx->nranks > 1) PetscCall(MatDestroy(&amgx->localA));
    381:     PetscCall(amgx_output_messages(amgx));
    382:     amgx->solve_state_init = false;
    383:   }
    384:   PetscFunctionReturn(PETSC_SUCCESS);
    385: }

    387: /*
    388:    PCDestroy_AMGX - Destroys the private context for the AmgX preconditioner
    389:    that was created with PCCreate_AMGX().

    391:    Input Parameter:
    392: .  pc - the preconditioner context

    394:    Application Interface Routine: PCDestroy()
    395: */
    396: static PetscErrorCode PCDestroy_AMGX(PC pc)
    397: {
    398:   PC_AMGX *amgx = (PC_AMGX *)pc->data;

    400:   PetscFunctionBegin;
    401:   /* decrease the number of instances, only the last instance need to destroy resource and finalizing AmgX */
    402:   if (s_count == 1) {
    403:     /* can put this in a PCAMGXInitializePackage method */
    404:     PetscCheck(amgx->rsrc != nullptr, PETSC_COMM_SELF, PETSC_ERR_PLIB, "s_rsrc == NULL");
    405:     PetscCallAmgX(AMGX_resources_destroy(amgx->rsrc));
    406:     /* destroy config (need to use AMGX_SAFE_CALL after this point) */
    407:     PetscCallAmgX(AMGX_config_destroy(amgx->cfg));
    408:     PetscCallAmgX(AMGX_finalize_plugins());
    409:     PetscCallAmgX(AMGX_finalize());
    410:     PetscCallMPI(MPI_Comm_free(&amgx->comm));
    411:   } else {
    412:     PetscCallAmgX(AMGX_config_destroy(amgx->cfg));
    413:   }
    414:   s_count -= 1;
    415:   PetscCall(PetscFree(amgx));
    416:   PetscFunctionReturn(PETSC_SUCCESS);
    417: }

    419: template <class T>
    420: std::string map_reverse_lookup(const std::map<std::string, T> &map, const T &key)
    421: {
    422:   for (auto const &m : map) {
    423:     if (m.second == key) return m.first;
    424:   }
    425:   return "";
    426: }

    428: static PetscErrorCode PCSetFromOptions_AMGX(PC pc, PetscOptionItems *PetscOptionsObject)
    429: {
    430:   PC_AMGX      *amgx          = (PC_AMGX *)pc->data;
    431:   constexpr int MAX_PARAM_LEN = 128;
    432:   char          option[MAX_PARAM_LEN];

    434:   PetscFunctionBegin;
    435:   PetscOptionsHeadBegin(PetscOptionsObject, "AmgX options");
    436:   amgx->cfg_contents = "config_version=2,";
    437:   amgx->cfg_contents += "determinism_flag=1,";

    439:   // Set exact coarse solve
    440:   PetscCall(PetscOptionsBool("-pc_amgx_exact_coarse_solve", "AmgX AMG Exact Coarse Solve", "", amgx->exact_coarse_solve, &amgx->exact_coarse_solve, NULL));
    441:   if (amgx->exact_coarse_solve) amgx->cfg_contents += "exact_coarse_solve=1,";

    443:   amgx->cfg_contents += "solver(amg)=AMG,";

    445:   // Set method
    446:   std::string def_amg_method = map_reverse_lookup(AmgXControlMap::AMGMethods, amgx->amg_method);
    447:   PetscCall(PetscStrncpy(option, def_amg_method.c_str(), sizeof(option)));
    448:   PetscCall(PetscOptionsString("-pc_amgx_amg_method", "AmgX AMG Method", "", option, option, MAX_PARAM_LEN, NULL));
    449:   PetscCheck(AmgXControlMap::AMGMethods.count(option) == 1, PETSC_COMM_SELF, PETSC_ERR_PLIB, "AMG Method %s not registered for AmgX.", option);
    450:   amgx->amg_method = AmgXControlMap::AMGMethods.at(option);
    451:   amgx->cfg_contents += "amg:algorithm=" + std::string(option) + ",";

    453:   // Set cycle
    454:   std::string def_amg_cycle = map_reverse_lookup(AmgXControlMap::AMGCycles, amgx->amg_cycle);
    455:   PetscCall(PetscStrncpy(option, def_amg_cycle.c_str(), sizeof(option)));
    456:   PetscCall(PetscOptionsString("-pc_amgx_amg_cycle", "AmgX AMG Cycle", "", option, option, MAX_PARAM_LEN, NULL));
    457:   PetscCheck(AmgXControlMap::AMGCycles.count(option) == 1, PETSC_COMM_SELF, PETSC_ERR_PLIB, "AMG Cycle %s not registered for AmgX.", option);
    458:   amgx->amg_cycle = AmgXControlMap::AMGCycles.at(option);
    459:   amgx->cfg_contents += "amg:cycle=" + std::string(option) + ",";

    461:   // Set smoother
    462:   std::string def_smoother = map_reverse_lookup(AmgXControlMap::Smoothers, amgx->smoother);
    463:   PetscCall(PetscStrncpy(option, def_smoother.c_str(), sizeof(option)));
    464:   PetscCall(PetscOptionsString("-pc_amgx_smoother", "AmgX Smoother", "", option, option, MAX_PARAM_LEN, NULL));
    465:   PetscCheck(AmgXControlMap::Smoothers.count(option) == 1, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Smoother %s not registered for AmgX.", option);
    466:   amgx->smoother = AmgXControlMap::Smoothers.at(option);
    467:   amgx->cfg_contents += "amg:smoother(smooth)=" + std::string(option) + ",";

    469:   if (amgx->smoother == AmgXSmoother::JacobiL1 || amgx->smoother == AmgXSmoother::BlockJacobi) {
    470:     PetscCall(PetscOptionsScalar("-pc_amgx_jacobi_relaxation_factor", "AmgX AMG Jacobi Relaxation Factor", "", amgx->jacobi_relaxation_factor, &amgx->jacobi_relaxation_factor, NULL));
    471:     amgx->cfg_contents += "smooth:relaxation_factor=" + std::to_string(amgx->jacobi_relaxation_factor) + ",";
    472:   } else if (amgx->smoother == AmgXSmoother::GS || amgx->smoother == AmgXSmoother::MulticolorGS) {
    473:     PetscCall(PetscOptionsScalar("-pc_amgx_gs_symmetric", "AmgX AMG Gauss Seidel Symmetric", "", amgx->gs_symmetric, &amgx->gs_symmetric, NULL));
    474:     amgx->cfg_contents += "smooth:symmetric_GS=" + std::to_string(amgx->gs_symmetric) + ",";
    475:   }

    477:   // Set selector
    478:   std::string def_selector = map_reverse_lookup(AmgXControlMap::Selectors, amgx->selector);
    479:   PetscCall(PetscStrncpy(option, def_selector.c_str(), sizeof(option)));
    480:   PetscCall(PetscOptionsString("-pc_amgx_selector", "AmgX Selector", "", option, option, MAX_PARAM_LEN, NULL));
    481:   PetscCheck(AmgXControlMap::Selectors.count(option) == 1, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Selector %s not registered for AmgX.", option);

    483:   // Double check that the user has selected an appropriate selector for the AMG method
    484:   if (amgx->amg_method == AmgXAMGMethod::Classical) {
    485:     PetscCheck(amgx->selector == AmgXSelector::PMIS || amgx->selector == AmgXSelector::HMIS, amgx->comm, PETSC_ERR_PLIB, "Chosen selector is not used for AmgX Classical AMG: selector=%s", option);
    486:     amgx->cfg_contents += "amg:interpolator=D2,";
    487:   } else if (amgx->amg_method == AmgXAMGMethod::Aggregation) {
    488:     PetscCheck(amgx->selector == AmgXSelector::Size2 || amgx->selector == AmgXSelector::Size4 || amgx->selector == AmgXSelector::Size8 || amgx->selector == AmgXSelector::MultiPairwise, amgx->comm, PETSC_ERR_PLIB, "Chosen selector is not used for AmgX Aggregation AMG");
    489:   }
    490:   amgx->selector = AmgXControlMap::Selectors.at(option);
    491:   amgx->cfg_contents += "amg:selector=" + std::string(option) + ",";

    493:   // Set presweeps
    494:   PetscCall(PetscOptionsInt("-pc_amgx_presweeps", "AmgX AMG Presweep Count", "", amgx->presweeps, &amgx->presweeps, NULL));
    495:   amgx->cfg_contents += "amg:presweeps=" + std::to_string(amgx->presweeps) + ",";

    497:   // Set postsweeps
    498:   PetscCall(PetscOptionsInt("-pc_amgx_postsweeps", "AmgX AMG Postsweep Count", "", amgx->postsweeps, &amgx->postsweeps, NULL));
    499:   amgx->cfg_contents += "amg:postsweeps=" + std::to_string(amgx->postsweeps) + ",";

    501:   // Set max levels
    502:   PetscCall(PetscOptionsInt("-pc_amgx_max_levels", "AmgX AMG Max Level Count", "", amgx->max_levels, &amgx->max_levels, NULL));
    503:   amgx->cfg_contents += "amg:max_levels=" + std::to_string(amgx->max_levels) + ",";

    505:   // Set dense LU num rows
    506:   PetscCall(PetscOptionsInt("-pc_amgx_dense_lu_num_rows", "AmgX Dense LU Number of Rows", "", amgx->dense_lu_num_rows, &amgx->dense_lu_num_rows, NULL));
    507:   amgx->cfg_contents += "amg:dense_lu_num_rows=" + std::to_string(amgx->dense_lu_num_rows) + ",";

    509:   // Set strength threshold
    510:   PetscCall(PetscOptionsScalar("-pc_amgx_strength_threshold", "AmgX AMG Strength Threshold", "", amgx->strength_threshold, &amgx->strength_threshold, NULL));
    511:   amgx->cfg_contents += "amg:strength_threshold=" + std::to_string(amgx->strength_threshold) + ",";

    513:   // Set aggressive_levels
    514:   PetscCall(PetscOptionsInt("-pc_amgx_aggressive_levels", "AmgX AMG Presweep Count", "", amgx->aggressive_levels, &amgx->aggressive_levels, NULL));
    515:   if (amgx->aggressive_levels > 0) amgx->cfg_contents += "amg:aggressive_levels=" + std::to_string(amgx->aggressive_levels) + ",";

    517:   // Set coarse solver
    518:   std::string def_coarse_solver = map_reverse_lookup(AmgXControlMap::CoarseSolvers, amgx->coarse_solver);
    519:   PetscCall(PetscStrncpy(option, def_coarse_solver.c_str(), sizeof(option)));
    520:   PetscCall(PetscOptionsString("-pc_amgx_coarse_solver", "AmgX CoarseSolver", "", option, option, MAX_PARAM_LEN, NULL));
    521:   PetscCheck(AmgXControlMap::CoarseSolvers.count(option) == 1, PETSC_COMM_SELF, PETSC_ERR_PLIB, "CoarseSolver %s not registered for AmgX.", option);
    522:   amgx->coarse_solver = AmgXControlMap::CoarseSolvers.at(option);
    523:   amgx->cfg_contents += "amg:coarse_solver=" + std::string(option) + ",";

    525:   // Set max iterations
    526:   amgx->cfg_contents += "amg:max_iters=1,";

    528:   // Set output control parameters
    529:   PetscCall(PetscOptionsBool("-pc_amgx_print_grid_stats", "AmgX Print Grid Stats", "", amgx->print_grid_stats, &amgx->print_grid_stats, NULL));

    531:   if (amgx->print_grid_stats) amgx->cfg_contents += "amg:print_grid_stats=1,";
    532:   amgx->cfg_contents += "amg:monitor_residual=0";

    534:   // Set whether AmgX output will be seen
    535:   PetscCall(PetscOptionsBool("-pc_amgx_verbose", "Enable output from AmgX", "", amgx->verbose, &amgx->verbose, NULL));
    536:   PetscOptionsHeadEnd();
    537:   PetscFunctionReturn(PETSC_SUCCESS);
    538: }

    540: static PetscErrorCode PCView_AMGX(PC pc, PetscViewer viewer)
    541: {
    542:   PC_AMGX  *amgx = (PC_AMGX *)pc->data;
    543:   PetscBool iascii;

    545:   PetscFunctionBegin;
    546:   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &iascii));
    547:   if (iascii) {
    548:     std::string output_cfg(amgx->cfg_contents);
    549:     std::replace(output_cfg.begin(), output_cfg.end(), ',', '\n');
    550:     PetscCall(PetscViewerASCIIPrintf(viewer, "\n%s\n", output_cfg.c_str()));
    551:   }
    552:   PetscFunctionReturn(PETSC_SUCCESS);
    553: }

    555: /*MC
    556:      PCAMGX - Interface to NVIDIA's AmgX algebraic multigrid

    558:    Options Database Keys:
    559: +    -pc_amgx_amg_method <CLASSICAL,AGGREGATION> - set the AMG algorithm to use
    560: .    -pc_amgx_amg_cycle <V,W,F,CG> - set the AMG cycle type
    561: .    -pc_amgx_smoother <PCG,PCGF,PBICGSTAB,GMRES,FGMRES,JACOBI_L1,BLOCK_JACOBI,GS,MULTICOLOR_GS,MULTICOLOR_ILU,MULTICOLOR_DILU,CHEBYSHEV_POLY,NOSOLVER> - set the AMG pre/post smoother
    562: .    -pc_amgx_jacobi_relaxation_factor - set the relaxation factor for Jacobi smoothing
    563: .    -pc_amgx_gs_symmetric - enforce symmetric Gauss-Seidel smoothing (only applies if GS smoothing is selected)
    564: .    -pc_amgx_selector <SIZE_2,SIZE_4,SIZE_8,MULTI_PAIRWISE,PMIS,HMIS> - set the AMG coarse selector
    565: .    -pc_amgx_presweeps - set the number of AMG pre-sweeps
    566: .    -pc_amgx_postsweeps - set the number of AMG post-sweeps
    567: .    -pc_amgx_max_levels - set the maximum number of levels in the AMG level hierarchy
    568: .    -pc_amgx_strength_threshold - set the strength threshold for the AMG coarsening
    569: .    -pc_amgx_aggressive_levels - set the number of levels (from the finest) that should apply aggressive coarsening
    570: .    -pc_amgx_coarse_solver <DENSE_LU_SOLVER,NOSOLVER> - set the coarse solve
    571: .    -pc_amgx_print_grid_stats - output the AMG grid hierarchy to stdout
    572: -    -pc_amgx_verbose - enable AmgX output

    574:    Level: intermediate

    576:    Note:
    577:      Implementation will accept host or device pointers, but good performance will require that the `KSP` is also GPU accelerated so that data is not frequently transferred between host and device.

    579: .seealso: [](ch_ksp), `PCGAMG`, `PCHYPRE`, `PCMG`, `PCAmgXGetResources()`, `PCCreate()`, `PCSetType()`, `PCType` (for list of available types), `PC`
    580: M*/

    582: PETSC_EXTERN PetscErrorCode PCCreate_AMGX(PC pc)
    583: {
    584:   PC_AMGX *amgx;

    586:   PetscFunctionBegin;
    587:   PetscCall(PetscNew(&amgx));
    588:   pc->ops->apply          = PCApply_AMGX;
    589:   pc->ops->setfromoptions = PCSetFromOptions_AMGX;
    590:   pc->ops->setup          = PCSetUp_AMGX;
    591:   pc->ops->view           = PCView_AMGX;
    592:   pc->ops->destroy        = PCDestroy_AMGX;
    593:   pc->ops->reset          = PCReset_AMGX;
    594:   pc->data                = (void *)amgx;

    596:   // Set the defaults
    597:   amgx->selector                 = AmgXSelector::PMIS;
    598:   amgx->smoother                 = AmgXSmoother::BlockJacobi;
    599:   amgx->amg_method               = AmgXAMGMethod::Classical;
    600:   amgx->coarse_solver            = AmgXCoarseSolver::DenseLU;
    601:   amgx->amg_cycle                = AmgXAMGCycle::V;
    602:   amgx->exact_coarse_solve       = PETSC_TRUE;
    603:   amgx->presweeps                = 1;
    604:   amgx->postsweeps               = 1;
    605:   amgx->max_levels               = 100;
    606:   amgx->strength_threshold       = 0.5;
    607:   amgx->aggressive_levels        = 0;
    608:   amgx->dense_lu_num_rows        = 1;
    609:   amgx->jacobi_relaxation_factor = 0.9;
    610:   amgx->gs_symmetric             = PETSC_FALSE;
    611:   amgx->print_grid_stats         = PETSC_FALSE;
    612:   amgx->verbose                  = PETSC_FALSE;
    613:   amgx->rsrc_init                = false;
    614:   amgx->solve_state_init         = false;

    616:   s_count++;

    618:   PetscCallCUDA(cudaGetDevice(&amgx->devID));
    619:   if (s_count == 1) {
    620:     PetscCallAmgX(AMGX_initialize());
    621:     PetscCallAmgX(AMGX_initialize_plugins());
    622:     PetscCallAmgX(AMGX_register_print_callback(&print_callback));
    623:     PetscCallAmgX(AMGX_install_signal_handler());
    624:   }
    625:   /* This communicator is not yet known to this system, so we duplicate it and make an internal communicator */
    626:   PetscCallMPI(MPI_Comm_dup(PetscObjectComm((PetscObject)pc), &amgx->comm));
    627:   PetscCallMPI(MPI_Comm_size(amgx->comm, &amgx->nranks));
    628:   PetscCallMPI(MPI_Comm_rank(amgx->comm, &amgx->rank));

    630:   PetscCall(amgx_output_messages(amgx));
    631:   PetscFunctionReturn(PETSC_SUCCESS);
    632: }

    634: /*@C
    635:   PCAmgXGetResources - get AMGx's internal resource object

    637:   Not Collective, No Fortran Support

    639:   Input Parameter:
    640: . pc - the PC

    642:   Output Parameter:
    643: . rsrc_out - pointer to the AMGx resource object

    645:   Level: advanced

    647: .seealso: [](ch_ksp), `PCAMGX`, `PC`, `PCGAMG`
    648: @*/
    649: PETSC_EXTERN PetscErrorCode PCAmgXGetResources(PC pc, void *rsrc_out)
    650: {
    651:   PC_AMGX *amgx = (PC_AMGX *)pc->data;

    653:   PetscFunctionBegin;
    654:   if (!amgx->rsrc_init) {
    655:     // Read configuration file
    656:     PetscCallAmgX(AMGX_config_create(&amgx->cfg, amgx->cfg_contents.c_str()));
    657:     PetscCallAmgX(AMGX_resources_create(&amgx->rsrc, amgx->cfg, &amgx->comm, 1, &amgx->devID));
    658:     amgx->rsrc_init = true;
    659:   }
    660:   *static_cast<AMGX_resources_handle *>(rsrc_out) = amgx->rsrc;
    661:   PetscFunctionReturn(PETSC_SUCCESS);
    662: }