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

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: }