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 isascii;
545: PetscFunctionBegin;
546: PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &isascii));
547: if (isascii) {
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, No Fortran Support
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: }