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;
178: // If AmgX output is enabled and we have a message, output it
179: if (amgx->verbose && !amgx_output.empty()) {
180: // Only a single rank to output the AmgX messages
181: PetscCall(PetscPrintf(amgx->comm, "AMGX: %s", amgx_output.c_str()));
183: // Note that all ranks clear their received output
184: amgx_output.clear();
185: }
187: PetscFunctionReturn(PETSC_SUCCESS);
188: }
190: // XXX Need to add call in AmgX API that gracefully destroys everything
191: // without abort etc.
192: #define PetscCallAmgX(rc) \
193: do { \
194: AMGX_RC err = (rc); \
195: char msg[4096]; \
196: switch (err) { \
197: case AMGX_RC_OK: \
198: break; \
199: default: \
200: AMGX_get_error_string(err, msg, 4096); \
201: SETERRQ(amgx->comm, PETSC_ERR_LIB, "%s", msg); \
202: } \
203: } while (0)
205: /*
206: PCSetUp_AMGX - Prepares for the use of the AmgX preconditioner
207: by setting data structures and options.
209: Input Parameter:
210: . pc - the preconditioner context
212: Application Interface Routine: PCSetUp()
214: Note:
215: The interface routine PCSetUp() is not usually called directly by
216: the user, but instead is called by PCApply() if necessary.
217: */
218: static PetscErrorCode PCSetUp_AMGX(PC pc)
219: {
220: PC_AMGX *amgx = (PC_AMGX *)pc->data;
221: Mat Pmat = pc->pmat;
222: PetscBool is_dev_ptrs;
224: PetscFunctionBegin;
225: PetscCall(PetscObjectTypeCompareAny((PetscObject)Pmat, &is_dev_ptrs, MATAIJCUSPARSE, MATSEQAIJCUSPARSE, MATMPIAIJCUSPARSE, ""));
227: // At the present time, an AmgX matrix is a sequential matrix
228: // Non-sequential/MPI matrices must be adapted to extract the local matrix
229: bool partial_setup_allowed = (pc->setupcalled && pc->flag != DIFFERENT_NONZERO_PATTERN);
230: if (amgx->nranks > 1) {
231: if (partial_setup_allowed) {
232: PetscCall(MatMPIAIJGetLocalMat(Pmat, MAT_REUSE_MATRIX, &amgx->localA));
233: } else {
234: PetscCall(MatMPIAIJGetLocalMat(Pmat, MAT_INITIAL_MATRIX, &amgx->localA));
235: }
237: if (is_dev_ptrs) PetscCall(MatConvert(amgx->localA, MATSEQAIJCUSPARSE, MAT_INPLACE_MATRIX, &amgx->localA));
238: } else {
239: amgx->localA = Pmat;
240: }
242: if (is_dev_ptrs) {
243: PetscCall(MatSeqAIJCUSPARSEGetArrayRead(amgx->localA, &amgx->values));
244: } else {
245: PetscCall(MatSeqAIJGetArrayRead(amgx->localA, &amgx->values));
246: }
248: if (!partial_setup_allowed) {
249: // Initialise resources and matrices
250: if (!amgx->rsrc_init) {
251: // Read configuration file
252: PetscCallAmgX(AMGX_config_create(&amgx->cfg, amgx->cfg_contents.c_str()));
253: PetscCallAmgX(AMGX_resources_create(&amgx->rsrc, amgx->cfg, &amgx->comm, 1, &amgx->devID));
254: amgx->rsrc_init = true;
255: }
257: PetscCheck(!amgx->solve_state_init, amgx->comm, PETSC_ERR_PLIB, "AmgX solve state initialisation already called.");
258: PetscCallAmgX(AMGX_matrix_create(&amgx->A, amgx->rsrc, AMGX_mode_dDDI));
259: PetscCallAmgX(AMGX_vector_create(&amgx->sol, amgx->rsrc, AMGX_mode_dDDI));
260: PetscCallAmgX(AMGX_vector_create(&amgx->rhs, amgx->rsrc, AMGX_mode_dDDI));
261: PetscCallAmgX(AMGX_solver_create(&amgx->solver, amgx->rsrc, AMGX_mode_dDDI, amgx->cfg));
262: amgx->solve_state_init = true;
264: // Extract the CSR data
265: PetscBool done;
266: const PetscInt *colIndices;
267: const PetscInt *rowOffsets;
268: PetscCall(MatGetRowIJ(amgx->localA, 0, PETSC_FALSE, PETSC_FALSE, &amgx->nLocalRows, &rowOffsets, &colIndices, &done));
269: PetscCheck(done, amgx->comm, PETSC_ERR_PLIB, "MatGetRowIJ was not successful");
270: 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);
272: if (is_dev_ptrs) {
273: PetscCallCUDA(cudaMemcpy(&amgx->nnz, &rowOffsets[amgx->nLocalRows], sizeof(int), cudaMemcpyDefault));
274: } else {
275: amgx->nnz = rowOffsets[amgx->nLocalRows];
276: }
278: 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);
280: // Allocate space for some partition offsets
281: std::vector<PetscInt> partitionOffsets(amgx->nranks + 1);
283: // Fetch the number of local rows per rank
284: partitionOffsets[0] = 0; /* could use PetscLayoutGetRanges */
285: PetscCallMPI(MPI_Allgather(&amgx->nLocalRows, 1, MPIU_INT, partitionOffsets.data() + 1, 1, MPIU_INT, amgx->comm));
286: std::partial_sum(partitionOffsets.begin(), partitionOffsets.end(), partitionOffsets.begin());
288: // Fetch the number of global rows
289: amgx->nGlobalRows = partitionOffsets[amgx->nranks];
291: PetscCall(MatGetBlockSize(Pmat, &amgx->bSize));
293: // XXX Currently constrained to 32-bit indices, to be changed in the future
294: // Create the distribution and upload the matrix data
295: AMGX_distribution_handle dist;
296: PetscCallAmgX(AMGX_distribution_create(&dist, amgx->cfg));
297: PetscCallAmgX(AMGX_distribution_set_32bit_colindices(dist, true));
298: PetscCallAmgX(AMGX_distribution_set_partition_data(dist, AMGX_DIST_PARTITION_OFFSETS, partitionOffsets.data()));
299: 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));
300: PetscCallAmgX(AMGX_solver_setup(amgx->solver, amgx->A));
301: PetscCallAmgX(AMGX_vector_bind(amgx->sol, amgx->A));
302: PetscCallAmgX(AMGX_vector_bind(amgx->rhs, amgx->A));
304: PetscInt nlr = 0;
305: PetscCall(MatRestoreRowIJ(amgx->localA, 0, PETSC_FALSE, PETSC_FALSE, &nlr, &rowOffsets, &colIndices, &done));
306: } else {
307: // The fast path for if the sparsity pattern persists
308: PetscCallAmgX(AMGX_matrix_replace_coefficients(amgx->A, amgx->nLocalRows, amgx->nnz, amgx->values, NULL));
309: PetscCallAmgX(AMGX_solver_resetup(amgx->solver, amgx->A));
310: }
312: if (is_dev_ptrs) {
313: PetscCall(MatSeqAIJCUSPARSERestoreArrayRead(amgx->localA, &amgx->values));
314: } else {
315: PetscCall(MatSeqAIJRestoreArrayRead(amgx->localA, &amgx->values));
316: }
317: PetscCall(amgx_output_messages(amgx));
318: PetscFunctionReturn(PETSC_SUCCESS);
319: }
321: /*
322: PCApply_AMGX - Applies the AmgX preconditioner to a vector.
324: Input Parameters:
325: . pc - the preconditioner context
326: . b - rhs vector
328: Output Parameter:
329: . x - solution vector
331: Application Interface Routine: PCApply()
332: */
333: static PetscErrorCode PCApply_AMGX(PC pc, Vec b, Vec x)
334: {
335: PC_AMGX *amgx = (PC_AMGX *)pc->data;
336: PetscScalar *x_;
337: const PetscScalar *b_;
338: PetscBool is_dev_ptrs;
340: PetscFunctionBegin;
341: PetscCall(PetscObjectTypeCompareAny((PetscObject)x, &is_dev_ptrs, VECCUDA, VECMPICUDA, VECSEQCUDA, ""));
343: if (is_dev_ptrs) {
344: PetscCall(VecCUDAGetArrayWrite(x, &x_));
345: PetscCall(VecCUDAGetArrayRead(b, &b_));
346: } else {
347: PetscCall(VecGetArrayWrite(x, &x_));
348: PetscCall(VecGetArrayRead(b, &b_));
349: }
351: PetscCallAmgX(AMGX_vector_upload(amgx->sol, amgx->nLocalRows, 1, x_));
352: PetscCallAmgX(AMGX_vector_upload(amgx->rhs, amgx->nLocalRows, 1, b_));
353: PetscCallAmgX(AMGX_solver_solve_with_0_initial_guess(amgx->solver, amgx->rhs, amgx->sol));
355: AMGX_SOLVE_STATUS status;
356: PetscCallAmgX(AMGX_solver_get_status(amgx->solver, &status));
357: PetscCall(PCSetErrorIfFailure(pc, static_cast<PetscBool>(status == AMGX_SOLVE_FAILED)));
358: PetscCheck(status != AMGX_SOLVE_FAILED, amgx->comm, PETSC_ERR_CONV_FAILED, "AmgX solver failed to solve the system! The error code is %d.", status);
359: PetscCallAmgX(AMGX_vector_download(amgx->sol, x_));
361: if (is_dev_ptrs) {
362: PetscCall(VecCUDARestoreArrayWrite(x, &x_));
363: PetscCall(VecCUDARestoreArrayRead(b, &b_));
364: } else {
365: PetscCall(VecRestoreArrayWrite(x, &x_));
366: PetscCall(VecRestoreArrayRead(b, &b_));
367: }
368: PetscCall(amgx_output_messages(amgx));
369: PetscFunctionReturn(PETSC_SUCCESS);
370: }
372: static PetscErrorCode PCReset_AMGX(PC pc)
373: {
374: PC_AMGX *amgx = (PC_AMGX *)pc->data;
376: PetscFunctionBegin;
377: if (amgx->solve_state_init) {
378: PetscCallAmgX(AMGX_solver_destroy(amgx->solver));
379: PetscCallAmgX(AMGX_matrix_destroy(amgx->A));
380: PetscCallAmgX(AMGX_vector_destroy(amgx->sol));
381: PetscCallAmgX(AMGX_vector_destroy(amgx->rhs));
382: if (amgx->nranks > 1) PetscCall(MatDestroy(&amgx->localA));
383: PetscCall(amgx_output_messages(amgx));
384: amgx->solve_state_init = false;
385: }
386: PetscFunctionReturn(PETSC_SUCCESS);
387: }
389: /*
390: PCDestroy_AMGX - Destroys the private context for the AmgX preconditioner
391: that was created with PCCreate_AMGX().
393: Input Parameter:
394: . pc - the preconditioner context
396: Application Interface Routine: PCDestroy()
397: */
398: static PetscErrorCode PCDestroy_AMGX(PC pc)
399: {
400: PC_AMGX *amgx = (PC_AMGX *)pc->data;
402: PetscFunctionBegin;
403: /* decrease the number of instances, only the last instance need to destroy resource and finalizing AmgX */
404: if (s_count == 1) {
405: /* can put this in a PCAMGXInitializePackage method */
406: PetscCheck(amgx->rsrc != nullptr, PETSC_COMM_SELF, PETSC_ERR_PLIB, "s_rsrc == NULL");
407: PetscCallAmgX(AMGX_resources_destroy(amgx->rsrc));
408: /* destroy config (need to use AMGX_SAFE_CALL after this point) */
409: PetscCallAmgX(AMGX_config_destroy(amgx->cfg));
410: PetscCallAmgX(AMGX_finalize_plugins());
411: PetscCallAmgX(AMGX_finalize());
412: PetscCallMPI(MPI_Comm_free(&amgx->comm));
413: } else {
414: PetscCallAmgX(AMGX_config_destroy(amgx->cfg));
415: }
416: s_count -= 1;
417: PetscCall(PetscFree(amgx));
418: PetscFunctionReturn(PETSC_SUCCESS);
419: }
421: template <class T>
422: std::string map_reverse_lookup(const std::map<std::string, T> &map, const T &key)
423: {
424: for (auto const &m : map) {
425: if (m.second == key) return m.first;
426: }
427: return "";
428: }
430: static PetscErrorCode PCSetFromOptions_AMGX(PC pc, PetscOptionItems *PetscOptionsObject)
431: {
432: PC_AMGX *amgx = (PC_AMGX *)pc->data;
433: constexpr int MAX_PARAM_LEN = 128;
434: char option[MAX_PARAM_LEN];
436: PetscFunctionBegin;
437: PetscOptionsHeadBegin(PetscOptionsObject, "AmgX options");
438: amgx->cfg_contents = "config_version=2,";
439: amgx->cfg_contents += "determinism_flag=1,";
441: // Set exact coarse solve
442: PetscCall(PetscOptionsBool("-pc_amgx_exact_coarse_solve", "AmgX AMG Exact Coarse Solve", "", amgx->exact_coarse_solve, &amgx->exact_coarse_solve, NULL));
443: if (amgx->exact_coarse_solve) amgx->cfg_contents += "exact_coarse_solve=1,";
445: amgx->cfg_contents += "solver(amg)=AMG,";
447: // Set method
448: std::string def_amg_method = map_reverse_lookup(AmgXControlMap::AMGMethods, amgx->amg_method);
449: PetscCall(PetscStrncpy(option, def_amg_method.c_str(), sizeof(option)));
450: PetscCall(PetscOptionsString("-pc_amgx_amg_method", "AmgX AMG Method", "", option, option, MAX_PARAM_LEN, NULL));
451: PetscCheck(AmgXControlMap::AMGMethods.count(option) == 1, PETSC_COMM_SELF, PETSC_ERR_PLIB, "AMG Method %s not registered for AmgX.", option);
452: amgx->amg_method = AmgXControlMap::AMGMethods.at(option);
453: amgx->cfg_contents += "amg:algorithm=" + std::string(option) + ",";
455: // Set cycle
456: std::string def_amg_cycle = map_reverse_lookup(AmgXControlMap::AMGCycles, amgx->amg_cycle);
457: PetscCall(PetscStrncpy(option, def_amg_cycle.c_str(), sizeof(option)));
458: PetscCall(PetscOptionsString("-pc_amgx_amg_cycle", "AmgX AMG Cycle", "", option, option, MAX_PARAM_LEN, NULL));
459: PetscCheck(AmgXControlMap::AMGCycles.count(option) == 1, PETSC_COMM_SELF, PETSC_ERR_PLIB, "AMG Cycle %s not registered for AmgX.", option);
460: amgx->amg_cycle = AmgXControlMap::AMGCycles.at(option);
461: amgx->cfg_contents += "amg:cycle=" + std::string(option) + ",";
463: // Set smoother
464: std::string def_smoother = map_reverse_lookup(AmgXControlMap::Smoothers, amgx->smoother);
465: PetscCall(PetscStrncpy(option, def_smoother.c_str(), sizeof(option)));
466: PetscCall(PetscOptionsString("-pc_amgx_smoother", "AmgX Smoother", "", option, option, MAX_PARAM_LEN, NULL));
467: PetscCheck(AmgXControlMap::Smoothers.count(option) == 1, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Smoother %s not registered for AmgX.", option);
468: amgx->smoother = AmgXControlMap::Smoothers.at(option);
469: amgx->cfg_contents += "amg:smoother(smooth)=" + std::string(option) + ",";
471: if (amgx->smoother == AmgXSmoother::JacobiL1 || amgx->smoother == AmgXSmoother::BlockJacobi) {
472: PetscCall(PetscOptionsScalar("-pc_amgx_jacobi_relaxation_factor", "AmgX AMG Jacobi Relaxation Factor", "", amgx->jacobi_relaxation_factor, &amgx->jacobi_relaxation_factor, NULL));
473: amgx->cfg_contents += "smooth:relaxation_factor=" + std::to_string(amgx->jacobi_relaxation_factor) + ",";
474: } else if (amgx->smoother == AmgXSmoother::GS || amgx->smoother == AmgXSmoother::MulticolorGS) {
475: PetscCall(PetscOptionsScalar("-pc_amgx_gs_symmetric", "AmgX AMG Gauss Seidel Symmetric", "", amgx->gs_symmetric, &amgx->gs_symmetric, NULL));
476: amgx->cfg_contents += "smooth:symmetric_GS=" + std::to_string(amgx->gs_symmetric) + ",";
477: }
479: // Set selector
480: std::string def_selector = map_reverse_lookup(AmgXControlMap::Selectors, amgx->selector);
481: PetscCall(PetscStrncpy(option, def_selector.c_str(), sizeof(option)));
482: PetscCall(PetscOptionsString("-pc_amgx_selector", "AmgX Selector", "", option, option, MAX_PARAM_LEN, NULL));
483: PetscCheck(AmgXControlMap::Selectors.count(option) == 1, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Selector %s not registered for AmgX.", option);
485: // Double check that the user has selected an appropriate selector for the AMG method
486: if (amgx->amg_method == AmgXAMGMethod::Classical) {
487: 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);
488: amgx->cfg_contents += "amg:interpolator=D2,";
489: } else if (amgx->amg_method == AmgXAMGMethod::Aggregation) {
490: 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");
491: }
492: amgx->selector = AmgXControlMap::Selectors.at(option);
493: amgx->cfg_contents += "amg:selector=" + std::string(option) + ",";
495: // Set presweeps
496: PetscCall(PetscOptionsInt("-pc_amgx_presweeps", "AmgX AMG Presweep Count", "", amgx->presweeps, &amgx->presweeps, NULL));
497: amgx->cfg_contents += "amg:presweeps=" + std::to_string(amgx->presweeps) + ",";
499: // Set postsweeps
500: PetscCall(PetscOptionsInt("-pc_amgx_postsweeps", "AmgX AMG Postsweep Count", "", amgx->postsweeps, &amgx->postsweeps, NULL));
501: amgx->cfg_contents += "amg:postsweeps=" + std::to_string(amgx->postsweeps) + ",";
503: // Set max levels
504: PetscCall(PetscOptionsInt("-pc_amgx_max_levels", "AmgX AMG Max Level Count", "", amgx->max_levels, &amgx->max_levels, NULL));
505: amgx->cfg_contents += "amg:max_levels=" + std::to_string(amgx->max_levels) + ",";
507: // Set dense LU num rows
508: 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));
509: amgx->cfg_contents += "amg:dense_lu_num_rows=" + std::to_string(amgx->dense_lu_num_rows) + ",";
511: // Set strength threshold
512: PetscCall(PetscOptionsScalar("-pc_amgx_strength_threshold", "AmgX AMG Strength Threshold", "", amgx->strength_threshold, &amgx->strength_threshold, NULL));
513: amgx->cfg_contents += "amg:strength_threshold=" + std::to_string(amgx->strength_threshold) + ",";
515: // Set aggressive_levels
516: PetscCall(PetscOptionsInt("-pc_amgx_aggressive_levels", "AmgX AMG Presweep Count", "", amgx->aggressive_levels, &amgx->aggressive_levels, NULL));
517: if (amgx->aggressive_levels > 0) amgx->cfg_contents += "amg:aggressive_levels=" + std::to_string(amgx->aggressive_levels) + ",";
519: // Set coarse solver
520: std::string def_coarse_solver = map_reverse_lookup(AmgXControlMap::CoarseSolvers, amgx->coarse_solver);
521: PetscCall(PetscStrncpy(option, def_coarse_solver.c_str(), sizeof(option)));
522: PetscCall(PetscOptionsString("-pc_amgx_coarse_solver", "AmgX CoarseSolver", "", option, option, MAX_PARAM_LEN, NULL));
523: PetscCheck(AmgXControlMap::CoarseSolvers.count(option) == 1, PETSC_COMM_SELF, PETSC_ERR_PLIB, "CoarseSolver %s not registered for AmgX.", option);
524: amgx->coarse_solver = AmgXControlMap::CoarseSolvers.at(option);
525: amgx->cfg_contents += "amg:coarse_solver=" + std::string(option) + ",";
527: // Set max iterations
528: amgx->cfg_contents += "amg:max_iters=1,";
530: // Set output control parameters
531: PetscCall(PetscOptionsBool("-pc_amgx_print_grid_stats", "AmgX Print Grid Stats", "", amgx->print_grid_stats, &amgx->print_grid_stats, NULL));
533: if (amgx->print_grid_stats) amgx->cfg_contents += "amg:print_grid_stats=1,";
534: amgx->cfg_contents += "amg:monitor_residual=0";
536: // Set whether AmgX output will be seen
537: PetscCall(PetscOptionsBool("-pc_amgx_verbose", "Enable output from AmgX", "", amgx->verbose, &amgx->verbose, NULL));
538: PetscOptionsHeadEnd();
539: PetscFunctionReturn(PETSC_SUCCESS);
540: }
542: static PetscErrorCode PCView_AMGX(PC pc, PetscViewer viewer)
543: {
544: PC_AMGX *amgx = (PC_AMGX *)pc->data;
545: PetscBool iascii;
547: PetscFunctionBegin;
548: PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &iascii));
549: if (iascii) {
550: std::string output_cfg(amgx->cfg_contents);
551: std::replace(output_cfg.begin(), output_cfg.end(), ',', '\n');
552: PetscCall(PetscViewerASCIIPrintf(viewer, "\n%s\n", output_cfg.c_str()));
553: }
554: PetscFunctionReturn(PETSC_SUCCESS);
555: }
557: /*MC
558: PCAMGX - Interface to NVIDIA's AmgX algebraic multigrid
560: Options Database Keys:
561: + -pc_amgx_amg_method <CLASSICAL,AGGREGATION> - set the AMG algorithm to use
562: . -pc_amgx_amg_cycle <V,W,F,CG> - set the AMG cycle type
563: . -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
564: . -pc_amgx_jacobi_relaxation_factor - set the relaxation factor for Jacobi smoothing
565: . -pc_amgx_gs_symmetric - enforce symmetric Gauss-Seidel smoothing (only applies if GS smoothing is selected)
566: . -pc_amgx_selector <SIZE_2,SIZE_4,SIZE_8,MULTI_PAIRWISE,PMIS,HMIS> - set the AMG coarse selector
567: . -pc_amgx_presweeps - set the number of AMG pre-sweeps
568: . -pc_amgx_postsweeps - set the number of AMG post-sweeps
569: . -pc_amgx_max_levels - set the maximum number of levels in the AMG level hierarchy
570: . -pc_amgx_strength_threshold - set the strength threshold for the AMG coarsening
571: . -pc_amgx_aggressive_levels - set the number of levels (from the finest) that should apply aggressive coarsening
572: . -pc_amgx_coarse_solver <DENSE_LU_SOLVER,NOSOLVER> - set the coarse solve
573: . -pc_amgx_print_grid_stats - output the AMG grid hierarchy to stdout
574: - -pc_amgx_verbose - enable AmgX output
576: Level: intermediate
578: Note:
579: 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.
581: .seealso: [](ch_ksp), `PCGAMG`, `PCHYPRE`, `PCMG`, `PCAmgXGetResources()`, `PCCreate()`, `PCSetType()`, `PCType` (for list of available types), `PC`
582: M*/
584: PETSC_EXTERN PetscErrorCode PCCreate_AMGX(PC pc)
585: {
586: PC_AMGX *amgx;
588: PetscFunctionBegin;
589: PetscCall(PetscNew(&amgx));
590: pc->ops->apply = PCApply_AMGX;
591: pc->ops->setfromoptions = PCSetFromOptions_AMGX;
592: pc->ops->setup = PCSetUp_AMGX;
593: pc->ops->view = PCView_AMGX;
594: pc->ops->destroy = PCDestroy_AMGX;
595: pc->ops->reset = PCReset_AMGX;
596: pc->data = (void *)amgx;
598: // Set the defaults
599: amgx->selector = AmgXSelector::PMIS;
600: amgx->smoother = AmgXSmoother::BlockJacobi;
601: amgx->amg_method = AmgXAMGMethod::Classical;
602: amgx->coarse_solver = AmgXCoarseSolver::DenseLU;
603: amgx->amg_cycle = AmgXAMGCycle::V;
604: amgx->exact_coarse_solve = PETSC_TRUE;
605: amgx->presweeps = 1;
606: amgx->postsweeps = 1;
607: amgx->max_levels = 100;
608: amgx->strength_threshold = 0.5;
609: amgx->aggressive_levels = 0;
610: amgx->dense_lu_num_rows = 1;
611: amgx->jacobi_relaxation_factor = 0.9;
612: amgx->gs_symmetric = PETSC_FALSE;
613: amgx->print_grid_stats = PETSC_FALSE;
614: amgx->verbose = PETSC_FALSE;
615: amgx->rsrc_init = false;
616: amgx->solve_state_init = false;
618: s_count++;
620: PetscCallCUDA(cudaGetDevice(&amgx->devID));
621: if (s_count == 1) {
622: PetscCallAmgX(AMGX_initialize());
623: PetscCallAmgX(AMGX_initialize_plugins());
624: PetscCallAmgX(AMGX_register_print_callback(&print_callback));
625: PetscCallAmgX(AMGX_install_signal_handler());
626: }
627: /* This communicator is not yet known to this system, so we duplicate it and make an internal communicator */
628: PetscCallMPI(MPI_Comm_dup(PetscObjectComm((PetscObject)pc), &amgx->comm));
629: PetscCallMPI(MPI_Comm_size(amgx->comm, &amgx->nranks));
630: PetscCallMPI(MPI_Comm_rank(amgx->comm, &amgx->rank));
632: PetscCall(amgx_output_messages(amgx));
633: PetscFunctionReturn(PETSC_SUCCESS);
634: }
636: /*@C
637: PCAmgXGetResources - get AMGx's internal resource object
639: Not Collective
641: Input Parameter:
642: . pc - the PC
644: Output Parameter:
645: . rsrc_out - pointer to the AMGx resource object
647: Level: advanced
649: .seealso: [](ch_ksp), `PCAMGX`, `PC`, `PCGAMG`
650: @*/
651: PETSC_EXTERN PetscErrorCode PCAmgXGetResources(PC pc, void *rsrc_out)
652: {
653: PC_AMGX *amgx = (PC_AMGX *)pc->data;
655: PetscFunctionBegin;
656: if (!amgx->rsrc_init) {
657: // Read configuration file
658: PetscCallAmgX(AMGX_config_create(&amgx->cfg, amgx->cfg_contents.c_str()));
659: PetscCallAmgX(AMGX_resources_create(&amgx->rsrc, amgx->cfg, &amgx->comm, 1, &amgx->devID));
660: amgx->rsrc_init = true;
661: }
662: *static_cast<AMGX_resources_handle *>(rsrc_out) = amgx->rsrc;
663: PetscFunctionReturn(PETSC_SUCCESS);
664: }