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