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_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
561: .    -pc_amgx_jacobi_relaxation_factor - set the relaxation factor for Jacobi smoothing
562: .    -pc_amgx_gs_symmetric - enforce symmetric Gauss-Seidel smoothing (only applies if GS smoothing is selected)
563: .    -pc_amgx_selector <SIZE_2,SIZE_4,SIZE_8,MULTI_PAIRWISE,PMIS,HMIS> - set the AMG coarse selector
564: .    -pc_amgx_presweeps - set the number of AMG pre-sweeps
565: .    -pc_amgx_postsweeps - set the number of AMG post-sweeps
566: .    -pc_amgx_max_levels - set the maximum number of levels in the AMG level hierarchy
567: .    -pc_amgx_strength_threshold - set the strength threshold for the AMG coarsening
568: .    -pc_amgx_aggressive_levels - set the number of levels (from the finest) that should apply aggressive coarsening
569: .    -pc_amgx_coarse_solver <DENSE_LU_SOLVER,NOSOLVER> - set the coarse solve
570: .    -pc_amgx_print_grid_stats - output the AMG grid hierarchy to stdout
571: -    -pc_amgx_verbose - enable AmgX output

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` (for list of available types), `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: }