1 /*
2 This file implements an AmgX preconditioner in PETSc as part of PC.
3 */
4
5 /*
6 Include files needed for the AmgX preconditioner:
7 pcimpl.h - private include file intended for use by all preconditioners
8 */
9
10 #include <petsc/private/pcimpl.h> /*I "petscpc.h" I*/
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
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 };
57
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 };
65
66 const std::map<std::string, AmgXAMGMethod> AmgXControlMap::AMGMethods = {
67 {"CLASSICAL", AmgXAMGMethod::Classical },
68 {"AGGREGATION", AmgXAMGMethod::Aggregation}
69 };
70
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 };
86
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 };
95
96 const std::map<std::string, AmgXCoarseSolver> AmgXControlMap::CoarseSolvers = {
97 {"DENSE_LU_SOLVER", AmgXCoarseSolver::DenseLU },
98 {"NOSOLVER", AmgXCoarseSolver::NoSolver}
99 };
100
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 };
108
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;
119
120 AMGX_matrix_handle A;
121 AMGX_vector_handle sol;
122 AMGX_vector_handle rhs;
123
124 MPI_Comm comm;
125 PetscMPIInt rank = 0;
126 PetscMPIInt nranks = 0;
127 int devID = 0;
128
129 void *lib_handle = 0;
130 std::string cfg_contents;
131
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;
139
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;
154
155 // Smoother control parameters
156 PetscScalar jacobi_relaxation_factor;
157 PetscScalar gs_symmetric;
158 };
159
160 static PetscInt s_count = 0;
161
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{};
165
166 // A print callback that allows AmgX to return status messages
print_callback(const char * msg,int length)167 static void print_callback(const char *msg, int length)
168 {
169 amgx_output.append(msg);
170 }
171
172 // Outputs messages from the AmgX message buffer and clears it
amgx_output_messages(PC_AMGX * amgx)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()));
180
181 // Note that all ranks clear their received output
182 amgx_output.clear();
183 }
184 PetscFunctionReturn(PETSC_SUCCESS);
185 }
186
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)
201
202 /*
203 PCSetUp_AMGX - Prepares for the use of the AmgX preconditioner
204 by setting data structures and options.
205
206 Input Parameter:
207 . pc - the preconditioner context
208
209 Application Interface Routine: PCSetUp()
210
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 */
PCSetUp_AMGX(PC pc)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;
220
221 PetscFunctionBegin;
222 PetscCall(PetscObjectTypeCompareAny((PetscObject)Pmat, &is_dev_ptrs, MATAIJCUSPARSE, MATSEQAIJCUSPARSE, MATMPIAIJCUSPARSE, ""));
223
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 }
233
234 if (is_dev_ptrs) PetscCall(MatConvert(amgx->localA, MATSEQAIJCUSPARSE, MAT_INPLACE_MATRIX, &amgx->localA));
235 } else {
236 amgx->localA = Pmat;
237 }
238
239 if (is_dev_ptrs) {
240 PetscCall(MatSeqAIJCUSPARSEGetArrayRead(amgx->localA, &amgx->values));
241 } else {
242 PetscCall(MatSeqAIJGetArrayRead(amgx->localA, &amgx->values));
243 }
244
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 }
253
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;
260
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);
268
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 }
274
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);
276
277 // Allocate space for some partition offsets
278 std::vector<PetscInt> partitionOffsets(amgx->nranks + 1);
279
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());
284
285 // Fetch the number of global rows
286 amgx->nGlobalRows = partitionOffsets[amgx->nranks];
287
288 PetscCall(MatGetBlockSize(Pmat, &amgx->bSize));
289
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));
300
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 }
308
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 }
317
318 /*
319 PCApply_AMGX - Applies the AmgX preconditioner to a vector.
320
321 Input Parameters:
322 . pc - the preconditioner context
323 . b - rhs vector
324
325 Output Parameter:
326 . x - solution vector
327
328 Application Interface Routine: PCApply()
329 */
PCApply_AMGX(PC pc,Vec b,Vec x)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;
336
337 PetscFunctionBegin;
338 PetscCall(PetscObjectTypeCompareAny((PetscObject)x, &is_dev_ptrs, VECCUDA, VECMPICUDA, VECSEQCUDA, ""));
339
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 }
347
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));
351
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_));
357
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 }
368
PCReset_AMGX(PC pc)369 static PetscErrorCode PCReset_AMGX(PC pc)
370 {
371 PC_AMGX *amgx = (PC_AMGX *)pc->data;
372
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 }
385
386 /*
387 PCDestroy_AMGX - Destroys the private context for the AmgX preconditioner
388 that was created with PCCreate_AMGX().
389
390 Input Parameter:
391 . pc - the preconditioner context
392
393 Application Interface Routine: PCDestroy()
394 */
PCDestroy_AMGX(PC pc)395 static PetscErrorCode PCDestroy_AMGX(PC pc)
396 {
397 PC_AMGX *amgx = (PC_AMGX *)pc->data;
398
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 }
417
418 template <class T>
map_reverse_lookup(const std::map<std::string,T> & map,const T & key)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 }
426
PCSetFromOptions_AMGX(PC pc,PetscOptionItems PetscOptionsObject)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];
432
433 PetscFunctionBegin;
434 PetscOptionsHeadBegin(PetscOptionsObject, "AmgX options");
435 amgx->cfg_contents = "config_version=2,";
436 amgx->cfg_contents += "determinism_flag=1,";
437
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,";
441
442 amgx->cfg_contents += "solver(amg)=AMG,";
443
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) + ",";
451
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) + ",";
459
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) + ",";
467
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 }
475
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);
481
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) + ",";
491
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) + ",";
495
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) + ",";
499
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) + ",";
503
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) + ",";
507
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) + ",";
511
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) + ",";
515
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) + ",";
523
524 // Set max iterations
525 amgx->cfg_contents += "amg:max_iters=1,";
526
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));
529
530 if (amgx->print_grid_stats) amgx->cfg_contents += "amg:print_grid_stats=1,";
531 amgx->cfg_contents += "amg:monitor_residual=0";
532
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 }
538
PCView_AMGX(PC pc,PetscViewer viewer)539 static PetscErrorCode PCView_AMGX(PC pc, PetscViewer viewer)
540 {
541 PC_AMGX *amgx = (PC_AMGX *)pc->data;
542 PetscBool isascii;
543
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 }
553
554 /*MC
555 PCAMGX - Interface to NVIDIA's AmgX algebraic multigrid
556
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
572
573 Level: intermediate
574
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.
577
578 .seealso: [](ch_ksp), `PCGAMG`, `PCHYPRE`, `PCMG`, `PCAmgXGetResources()`, `PCCreate()`, `PCSetType()`, `PCType` (for list of available types), `PC`
579 M*/
580
PCCreate_AMGX(PC pc)581 PETSC_EXTERN PetscErrorCode PCCreate_AMGX(PC pc)
582 {
583 PC_AMGX *amgx;
584
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;
594
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;
614
615 s_count++;
616
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));
628
629 PetscCall(amgx_output_messages(amgx));
630 PetscFunctionReturn(PETSC_SUCCESS);
631 }
632
633 /*@C
634 PCAmgXGetResources - get AMGx's internal resource object
635
636 Not Collective, No Fortran Support
637
638 Input Parameter:
639 . pc - the PC
640
641 Output Parameter:
642 . rsrc_out - pointer to the AMGx resource object
643
644 Level: advanced
645
646 .seealso: [](ch_ksp), `PCAMGX`, `PC`, `PCGAMG`
647 @*/
PCAmgXGetResources(PC pc,void * rsrc_out)648 PETSC_EXTERN PetscErrorCode PCAmgXGetResources(PC pc, void *rsrc_out)
649 {
650 PC_AMGX *amgx = (PC_AMGX *)pc->data;
651
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 }
662