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>

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    160: static PetscInt s_count = 0;

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

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

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

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

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

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

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

    209:    Application Interface Routine: PCSetUp()

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

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

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

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

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

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

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

    261:     // Extract the CSR data
    262:     PetscBool       done;
    263:     const PetscInt *colIndices;
    264:     const PetscInt *rowOffsets;
    265:     PetscCall(MatGetRowIJ(amgx->localA, 0, PETSC_FALSE, PETSC_FALSE, &amgx->nLocalRows, &rowOffsets, &colIndices, &done));
    266:     PetscCheck(done, amgx->comm, PETSC_ERR_PLIB, "MatGetRowIJ was not successful");
    267:     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);

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

    275:     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);

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

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

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

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

    290:     // XXX Currently constrained to 32-bit indices, to be changed in the future
    291:     // Create the distribution and upload the matrix data
    292:     AMGX_distribution_handle dist;
    293:     PetscCallAmgX(AMGX_distribution_create(&dist, amgx->cfg));
    294:     PetscCallAmgX(AMGX_distribution_set_32bit_colindices(dist, true));
    295:     PetscCallAmgX(AMGX_distribution_set_partition_data(dist, AMGX_DIST_PARTITION_OFFSETS, partitionOffsets.data()));
    296:     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));
    297:     PetscCallAmgX(AMGX_solver_setup(amgx->solver, amgx->A));
    298:     PetscCallAmgX(AMGX_vector_bind(amgx->sol, amgx->A));
    299:     PetscCallAmgX(AMGX_vector_bind(amgx->rhs, amgx->A));

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    482:   // Double check that the user has selected an appropriate selector for the AMG method
    483:   if (amgx->amg_method == AmgXAMGMethod::Classical) {
    484:     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);
    485:     amgx->cfg_contents += "amg:interpolator=D2,";
    486:   } else if (amgx->amg_method == AmgXAMGMethod::Aggregation) {
    487:     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");
    488:   }
    489:   amgx->selector = AmgXControlMap::Selectors.at(option);
    490:   amgx->cfg_contents += "amg:selector=" + std::string(option) + ",";

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

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

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

    504:   // Set dense LU num rows
    505:   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));
    506:   amgx->cfg_contents += "amg:dense_lu_num_rows=" + std::to_string(amgx->dense_lu_num_rows) + ",";

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

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

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

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

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

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

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

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

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

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

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

    573:    Level: intermediate

    575:    Note:
    576:    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.

    578: .seealso: [](ch_ksp), `PCGAMG`, `PCHYPRE`, `PCMG`, `PCAmgXGetResources()`, `PCCreate()`, `PCSetType()`, `PCType`, `PC`
    579: M*/

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

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

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

    615:   s_count++;

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

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

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

    636:   Not Collective, No Fortran Support

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

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

    644:   Level: advanced

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

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