Actual source code: pipeprcg.c
1: #include <petsc/private/kspimpl.h>
3: typedef struct KSP_CG_PIPE_PR_s KSP_CG_PIPE_PR;
4: struct KSP_CG_PIPE_PR_s {
5: PetscBool rc_w_q; /* flag to determine whether w_k should be recomputed with A r_k */
6: };
8: /*
9: KSPSetUp_PIPEPRCG - Sets up the workspace needed by the PIPEPRCG method.
11: This is called once, usually automatically by KSPSolve() or KSPSetUp()
12: but can be called directly by KSPSetUp()
13: */
14: static PetscErrorCode KSPSetUp_PIPEPRCG(KSP ksp)
15: {
16: PetscFunctionBegin;
17: /* get work vectors needed by PIPEPRCG */
18: PetscCall(KSPSetWorkVecs(ksp, 9));
19: PetscFunctionReturn(PETSC_SUCCESS);
20: }
22: static PetscErrorCode KSPSetFromOptions_PIPEPRCG(KSP ksp, PetscOptionItems PetscOptionsObject)
23: {
24: KSP_CG_PIPE_PR *prcg = (KSP_CG_PIPE_PR *)ksp->data;
25: PetscBool flag = PETSC_FALSE;
27: PetscFunctionBegin;
28: PetscOptionsHeadBegin(PetscOptionsObject, "KSP PIPEPRCG options");
29: PetscCall(PetscOptionsBool("-recompute_w", "-recompute w_k with Ar_k? (default = True)", "", prcg->rc_w_q, &prcg->rc_w_q, &flag));
30: if (!flag) prcg->rc_w_q = PETSC_TRUE;
31: PetscOptionsHeadEnd();
32: PetscFunctionReturn(PETSC_SUCCESS);
33: }
35: /*
36: KSPSolve_PIPEPRCG - This routine actually applies the pipelined predict and recompute conjugate gradient method
37: */
38: static PetscErrorCode KSPSolve_PIPEPRCG(KSP ksp)
39: {
40: PetscInt i;
41: KSP_CG_PIPE_PR *prcg = (KSP_CG_PIPE_PR *)ksp->data;
42: PetscScalar alpha = 0.0, beta = 0.0, nu = 0.0, nu_old = 0.0, mudelgam[3], *mu_p, *delta_p, *gamma_p;
43: PetscReal dp = 0.0;
44: Vec X, B, R, RT, W, WT, P, S, ST, U, UT, PRTST[3];
45: Mat Amat, Pmat;
46: PetscBool rc_w_q = prcg->rc_w_q;
48: /* note that these are pointers to entries of muldelgam, different than nu */
49: mu_p = &mudelgam[0];
50: delta_p = &mudelgam[1];
51: gamma_p = &mudelgam[2];
53: PetscFunctionBegin;
54: X = ksp->vec_sol;
55: B = ksp->vec_rhs;
56: R = ksp->work[0];
57: RT = ksp->work[1];
58: W = ksp->work[2];
59: WT = ksp->work[3];
60: P = ksp->work[4];
61: S = ksp->work[5];
62: ST = ksp->work[6];
63: U = ksp->work[7];
64: UT = ksp->work[8];
66: PetscCall(PCGetOperators(ksp->pc, &Amat, &Pmat));
68: /* initialize */
69: ksp->its = 0;
70: if (!ksp->guess_zero) {
71: PetscCall(KSP_MatMult(ksp, Amat, X, R)); /* r <- b - Ax */
72: PetscCall(VecAYPX(R, -1.0, B));
73: } else {
74: PetscCall(VecCopy(B, R)); /* r <- b */
75: }
77: PetscCall(KSP_PCApply(ksp, R, RT)); /* rt <- Br */
78: PetscCall(KSP_MatMult(ksp, Amat, RT, W)); /* w <- A rt */
79: PetscCall(KSP_PCApply(ksp, W, WT)); /* wt <- B w */
81: PetscCall(VecCopy(RT, P)); /* p <- rt */
82: PetscCall(VecCopy(W, S)); /* p <- rt */
83: PetscCall(VecCopy(WT, ST)); /* p <- rt */
85: PetscCall(KSP_MatMult(ksp, Amat, ST, U)); /* u <- Ast */
86: PetscCall(KSP_PCApply(ksp, U, UT)); /* ut <- Bu */
88: PetscCall(VecDotBegin(RT, R, &nu));
89: PetscCall(VecDotBegin(P, S, mu_p));
90: PetscCall(VecDotBegin(ST, S, gamma_p));
92: PetscCall(VecDotEnd(RT, R, &nu)); /* nu <- (rt,r) */
93: PetscCall(VecDotEnd(P, S, mu_p)); /* mu <- (p,s) */
94: PetscCall(VecDotEnd(ST, S, gamma_p)); /* gamma <- (st,s) */
95: *delta_p = *mu_p;
97: i = 0;
98: do {
99: /* Compute appropriate norm */
100: switch (ksp->normtype) {
101: case KSP_NORM_PRECONDITIONED:
102: PetscCall(VecNormBegin(RT, NORM_2, &dp));
103: PetscCall(PetscCommSplitReductionBegin(PetscObjectComm((PetscObject)RT)));
104: PetscCall(VecNormEnd(RT, NORM_2, &dp));
105: break;
106: case KSP_NORM_UNPRECONDITIONED:
107: PetscCall(VecNormBegin(R, NORM_2, &dp));
108: PetscCall(PetscCommSplitReductionBegin(PetscObjectComm((PetscObject)R)));
109: PetscCall(VecNormEnd(R, NORM_2, &dp));
110: break;
111: case KSP_NORM_NATURAL:
112: dp = PetscSqrtReal(PetscAbsScalar(nu));
113: break;
114: case KSP_NORM_NONE:
115: dp = 0.0;
116: break;
117: default:
118: SETERRQ(PetscObjectComm((PetscObject)ksp), PETSC_ERR_SUP, "%s", KSPNormTypes[ksp->normtype]);
119: }
121: ksp->rnorm = dp;
122: PetscCall(KSPLogResidualHistory(ksp, dp));
123: PetscCall(KSPMonitor(ksp, i, dp));
124: PetscCall((*ksp->converged)(ksp, i, dp, &ksp->reason, ksp->cnvP));
125: if (ksp->reason) PetscFunctionReturn(PETSC_SUCCESS);
127: /* update scalars */
128: alpha = nu / *mu_p;
129: nu_old = nu;
130: nu = nu_old - 2. * alpha * (*delta_p) + (alpha * alpha) * (*gamma_p);
131: beta = nu / nu_old;
133: /* update vectors */
134: PetscCall(VecAXPY(X, alpha, P)); /* x <- x + alpha * p */
135: PetscCall(VecAXPY(R, -alpha, S)); /* r <- r - alpha * s */
136: PetscCall(VecAXPY(RT, -alpha, ST)); /* rt <- rt - alpha * st */
137: PetscCall(VecAXPY(W, -alpha, U)); /* w <- w - alpha * u */
138: PetscCall(VecAXPY(WT, -alpha, UT)); /* wt <- wt - alpha * ut */
139: PetscCall(VecAYPX(P, beta, RT)); /* p <- rt + beta * p */
140: PetscCall(VecAYPX(S, beta, W)); /* s <- w + beta * s */
141: PetscCall(VecAYPX(ST, beta, WT)); /* st <- wt + beta * st */
143: PetscCall(VecDotBegin(RT, R, &nu));
145: PRTST[0] = P;
146: PRTST[1] = RT;
147: PRTST[2] = ST;
149: PetscCall(VecMDotBegin(S, 3, PRTST, mudelgam));
151: PetscCall(PetscCommSplitReductionBegin(PetscObjectComm((PetscObject)R)));
153: PetscCall(KSP_MatMult(ksp, Amat, ST, U)); /* u <- A st */
154: PetscCall(KSP_PCApply(ksp, U, UT)); /* ut <- B u */
156: /* predict-and-recompute */
157: /* ideally this is combined with the previous matvec; i.e. equivalent of MDot */
158: if (rc_w_q) {
159: PetscCall(KSP_MatMult(ksp, Amat, RT, W)); /* w <- A rt */
160: PetscCall(KSP_PCApply(ksp, W, WT)); /* wt <- B w */
161: }
163: PetscCall(VecDotEnd(RT, R, &nu));
164: PetscCall(VecMDotEnd(S, 3, PRTST, mudelgam));
166: i++;
167: ksp->its = i;
169: } while (i <= ksp->max_it);
170: if (!ksp->reason) ksp->reason = KSP_DIVERGED_ITS;
171: PetscFunctionReturn(PETSC_SUCCESS);
172: }
174: /*MC
175: KSPPIPEPRCG - Pipelined predict-and-recompute conjugate gradient Krylov method {cite}`chen2020predict`. [](sec_pipelineksp)
177: Options Database Key:
178: . -ksp_pipeprcg_recompute_w - recompute the $w_k$ with $Ar_k$, default is true
180: Level: intermediate
182: Notes:
183: This method has only a single non-blocking reduction per iteration, compared to 2 blocking for standard `KSPCG`.
184: The non-blocking reduction is overlapped by the matrix-vector product and preconditioner application.
186: MPI configuration may be necessary for reductions to make asynchronous progress, which is important for performance of pipelined methods.
187: See [](doc_faq_pipelined)
189: Contributed by:
190: Tyler Chen, University of Washington, Applied Mathematics Department
192: Acknowledgments:
193: This material is based upon work supported by the National Science Foundation Graduate Research Fellowship Program under Grant No. DGE-1762114.
194: Any opinions, findings, and conclusions or recommendations expressed in this material are those of the author and do not necessarily
195: reflect the views of the National Science Foundation.
197: .seealso: [](ch_ksp), [](doc_faq_pipelined), [](sec_pipelineksp), `KSPCreate()`, `KSPSetType()`, `KSPCG`, `KSPPIPECG`, `KSPPIPECR`, `KSPGROPPCG`, `KSPPGMRES`, `KSPCGUseSingleReduction()`
198: M*/
199: PETSC_EXTERN PetscErrorCode KSPCreate_PIPEPRCG(KSP ksp)
200: {
201: KSP_CG_PIPE_PR *prcg = NULL;
202: PetscBool cite = PETSC_FALSE;
204: PetscFunctionBegin;
205: PetscCall(PetscCitationsRegister("@article{predict_and_recompute_cg,\n author = {Tyler Chen and Erin C. Carson},\n title = {Predict-and-recompute conjugate gradient variants},\n journal = {},\n year = {2020},\n eprint = {1905.01549},\n "
206: "archivePrefix = {arXiv},\n primaryClass = {cs.NA}\n}",
207: &cite));
209: PetscCall(PetscNew(&prcg));
210: ksp->data = (void *)prcg;
212: PetscCall(KSPSetSupportedNorm(ksp, KSP_NORM_UNPRECONDITIONED, PC_LEFT, 2));
213: PetscCall(KSPSetSupportedNorm(ksp, KSP_NORM_PRECONDITIONED, PC_LEFT, 2));
214: PetscCall(KSPSetSupportedNorm(ksp, KSP_NORM_NATURAL, PC_LEFT, 2));
215: PetscCall(KSPSetSupportedNorm(ksp, KSP_NORM_NONE, PC_LEFT, 1));
217: ksp->ops->setup = KSPSetUp_PIPEPRCG;
218: ksp->ops->solve = KSPSolve_PIPEPRCG;
219: ksp->ops->destroy = KSPDestroyDefault;
220: ksp->ops->view = NULL;
221: ksp->ops->setfromoptions = KSPSetFromOptions_PIPEPRCG;
222: ksp->ops->buildsolution = KSPBuildSolutionDefault;
223: ksp->ops->buildresidual = KSPBuildResidualDefault;
224: PetscFunctionReturn(PETSC_SUCCESS);
225: }