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: PetscErrorCode amgx_output_messages(PC_AMGX *amgx)
175: {

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:     PetscPrintf(amgx->comm, "AMGX: %s", amgx_output.c_str());

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

186:   return 0;
187: }

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

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

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

211:    Application Interface Routine: PCSetUp()

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

223:   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:       MatMPIAIJGetLocalMat(Pmat, MAT_REUSE_MATRIX, &amgx->localA);
231:     } else {
232:       MatMPIAIJGetLocalMat(Pmat, MAT_INITIAL_MATRIX, &amgx->localA);
233:     }

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

240:   if (is_dev_ptrs) {
241:     MatSeqAIJCUSPARSEGetArrayRead(amgx->localA, &amgx->values);
242:   } else {
243:     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:       AMGX_config_create(&amgx->cfg, amgx->cfg_contents.c_str());
251:       AMGX_resources_create(&amgx->rsrc, amgx->cfg, &amgx->comm, 1, &amgx->devID);
252:       amgx->rsrc_init = true;
253:     }

256:     AMGX_matrix_create(&amgx->A, amgx->rsrc, AMGX_mode_dDDI);
257:     AMGX_vector_create(&amgx->sol, amgx->rsrc, AMGX_mode_dDDI);
258:     AMGX_vector_create(&amgx->rhs, amgx->rsrc, AMGX_mode_dDDI);
259:     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:     MatGetRowIJ(amgx->localA, 0, PETSC_FALSE, PETSC_FALSE, &amgx->nLocalRows, &rowOffsets, &colIndices, &done);

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


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:     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:     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:     AMGX_distribution_create(&dist, amgx->cfg);
295:     AMGX_distribution_set_32bit_colindices(dist, true);
296:     AMGX_distribution_set_partition_data(dist, AMGX_DIST_PARTITION_OFFSETS, partitionOffsets.data());
297:     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:     AMGX_solver_setup(amgx->solver, amgx->A);
299:     AMGX_vector_bind(amgx->sol, amgx->A);
300:     AMGX_vector_bind(amgx->rhs, amgx->A);

302:     PetscInt nlr = 0;
303:     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:     AMGX_matrix_replace_coefficients(amgx->A, amgx->nLocalRows, amgx->nnz, amgx->values, NULL);
307:     AMGX_solver_resetup(amgx->solver, amgx->A);
308:   }

310:   if (is_dev_ptrs) {
311:     MatSeqAIJCUSPARSERestoreArrayRead(amgx->localA, &amgx->values);
312:   } else {
313:     MatSeqAIJRestoreArrayRead(amgx->localA, &amgx->values);
314:   }
315:   amgx_output_messages(amgx);
316:   return 0;
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:   PetscObjectTypeCompareAny((PetscObject)x, &is_dev_ptrs, VECCUDA, VECMPICUDA, VECSEQCUDA, "");

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

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

352:   AMGX_SOLVE_STATUS status;
353:   AMGX_solver_get_status(amgx->solver, &status);
354:   PCSetErrorIfFailure(pc, static_cast<PetscBool>(status == AMGX_SOLVE_FAILED));
356:   AMGX_vector_download(amgx->sol, x_);

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

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

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

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

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

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

398:   /* decrease the number of instances, only the last instance need to destroy resource and finalizing AmgX */
399:   if (s_count == 1) {
400:     /* can put this in a PCAMGXInitializePackage method */
402:     AMGX_resources_destroy(amgx->rsrc);
403:     /* destroy config (need to use AMGX_SAFE_CALL after this point) */
404:     AMGX_config_destroy(amgx->cfg);
405:     AMGX_finalize_plugins();
406:     AMGX_finalize();
407:     MPI_Comm_free(&amgx->comm);
408:   } else {
409:     AMGX_config_destroy(amgx->cfg);
410:   }
411:   s_count -= 1;
412:   PetscFree(amgx);
413:   return 0;
414: }

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

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

431:   PetscOptionsHeadBegin(PetscOptionsObject, "AmgX options");
432:   amgx->cfg_contents = "config_version=2,";
433:   amgx->cfg_contents += "determinism_flag=1,";

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

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

441:   // Set method
442:   std::string def_amg_method = map_reverse_lookup(AmgXControlMap::AMGMethods, amgx->amg_method);
443:   PetscStrcpy(option, def_amg_method.c_str());
444:   PetscOptionsString("-pc_amgx_amg_method", "AmgX AMG Method", "", option, option, MAX_PARAM_LEN, NULL);
446:   amgx->amg_method = AmgXControlMap::AMGMethods.at(option);
447:   amgx->cfg_contents += "amg:algorithm=" + std::string(option) + ",";

449:   // Set cycle
450:   std::string def_amg_cycle = map_reverse_lookup(AmgXControlMap::AMGCycles, amgx->amg_cycle);
451:   PetscStrcpy(option, def_amg_cycle.c_str());
452:   PetscOptionsString("-pc_amgx_amg_cycle", "AmgX AMG Cycle", "", option, option, MAX_PARAM_LEN, NULL);
454:   amgx->amg_cycle = AmgXControlMap::AMGCycles.at(option);
455:   amgx->cfg_contents += "amg:cycle=" + std::string(option) + ",";

457:   // Set smoother
458:   std::string def_smoother = map_reverse_lookup(AmgXControlMap::Smoothers, amgx->smoother);
459:   PetscStrcpy(option, def_smoother.c_str());
460:   PetscOptionsString("-pc_amgx_smoother", "AmgX Smoother", "", option, option, MAX_PARAM_LEN, NULL);
462:   amgx->smoother = AmgXControlMap::Smoothers.at(option);
463:   amgx->cfg_contents += "amg:smoother(smooth)=" + std::string(option) + ",";

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

473:   // Set selector
474:   std::string def_selector = map_reverse_lookup(AmgXControlMap::Selectors, amgx->selector);
475:   PetscStrcpy(option, def_selector.c_str());
476:   PetscOptionsString("-pc_amgx_selector", "AmgX Selector", "", option, option, MAX_PARAM_LEN, NULL);

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

489:   // Set presweeps
490:   PetscOptionsInt("-pc_amgx_presweeps", "AmgX AMG Presweep Count", "", amgx->presweeps, &amgx->presweeps, NULL);
491:   amgx->cfg_contents += "amg:presweeps=" + std::to_string(amgx->presweeps) + ",";

493:   // Set postsweeps
494:   PetscOptionsInt("-pc_amgx_postsweeps", "AmgX AMG Postsweep Count", "", amgx->postsweeps, &amgx->postsweeps, NULL);
495:   amgx->cfg_contents += "amg:postsweeps=" + std::to_string(amgx->postsweeps) + ",";

497:   // Set max levels
498:   PetscOptionsInt("-pc_amgx_max_levels", "AmgX AMG Max Level Count", "", amgx->max_levels, &amgx->max_levels, NULL);
499:   amgx->cfg_contents += "amg:max_levels=100,";

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

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

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

513:   // Set coarse solver
514:   std::string def_coarse_solver = map_reverse_lookup(AmgXControlMap::CoarseSolvers, amgx->coarse_solver);
515:   PetscStrcpy(option, def_coarse_solver.c_str());
516:   PetscOptionsString("-pc_amgx_coarse_solver", "AmgX CoarseSolver", "", option, option, MAX_PARAM_LEN, NULL);
518:   amgx->coarse_solver = AmgXControlMap::CoarseSolvers.at(option);
519:   amgx->cfg_contents += "amg:coarse_solver=" + std::string(option) + ",";

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

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

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

530:   // Set whether AmgX output will be seen
531:   PetscOptionsBool("-pc_amgx_verbose", "Enable output from AmgX", "", amgx->verbose, &amgx->verbose, NULL);
532:   PetscOptionsHeadEnd();
533:   return 0;
534: }

536: static PetscErrorCode PCView_AMGX(PC pc, PetscViewer viewer)
537: {
538:   PC_AMGX  *amgx = (PC_AMGX *)pc->data;
539:   PetscBool iascii;

541:   PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &iascii);
542:   if (iascii) {
543:     std::string output_cfg(amgx->cfg_contents);
544:     std::replace(output_cfg.begin(), output_cfg.end(), ',', '\n');
545:     PetscViewerASCIIPrintf(viewer, "\n%s\n", output_cfg.c_str());
546:   }
547:   return 0;
548: }

550: /*MC
551:      PCAMGX - Interface to NVIDIA's AmgX algebraic multigrid

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

569:    Level: intermediate

571:    Note:
572:      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.

574: .seealso: `PCGAMG`, `PCHYPRE`, `PCMG`, `PCAmgXGetResources()`, `PCCreate()`, `PCSetType()`, `PCType` (for list of available types), `PC`
575: M*/

577: PETSC_EXTERN PetscErrorCode PCCreate_AMGX(PC pc)
578: {
579:   PC_AMGX *amgx;

581:   PetscNew(&amgx);
582:   pc->ops->apply          = PCApply_AMGX;
583:   pc->ops->setfromoptions = PCSetFromOptions_AMGX;
584:   pc->ops->setup          = PCSetUp_AMGX;
585:   pc->ops->view           = PCView_AMGX;
586:   pc->ops->destroy        = PCDestroy_AMGX;
587:   pc->ops->reset          = PCReset_AMGX;
588:   pc->data                = (void *)amgx;

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

610:   s_count++;

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

624:   amgx_output_messages(amgx);
625:   return 0;
626: }

628: /*@C
629:    PCAmgXGetResources - get AMGx's internal resource object

631:     Not Collective

633:    Input Parameter:
634: .  pc - the PC

636:    Output Parameter:
637: .  rsrc_out - pointer to the AMGx resource object

639:    Level: advanced

641: .seealso: `PCAMGX`, `PC`, `PCGAMG`
642: @*/
643: PETSC_EXTERN PetscErrorCode PCAmgXGetResources(PC pc, void *rsrc_out)
644: {
645:   PC_AMGX *amgx = (PC_AMGX *)pc->data;

647:   if (!amgx->rsrc_init) {
648:     // Read configuration file
649:     AMGX_config_create(&amgx->cfg, amgx->cfg_contents.c_str());
650:     AMGX_resources_create(&amgx->rsrc, amgx->cfg, &amgx->comm, 1, &amgx->devID);
651:     amgx->rsrc_init = true;
652:   }
653:   *static_cast<AMGX_resources_handle *>(rsrc_out) = amgx->rsrc;
654:   return 0;
655: }