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 recomputer 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: /* get work vectors needed by PIPEPRCG */
17: KSPSetWorkVecs(ksp,9);
19: return 0;
20: }
22: static PetscErrorCode KSPSetFromOptions_PIPEPRCG(PetscOptionItems *PetscOptionsObject,KSP ksp)
23: {
24: KSP_CG_PIPE_PR *prcg=(KSP_CG_PIPE_PR*)ksp->data;
25: PetscBool flag=PETSC_FALSE;
27: PetscOptionsHead(PetscOptionsObject,"KSP PIPEPRCG options");
28: PetscOptionsBool("-recompute_w","-recompute w_k with Ar_k? (default = True)","",prcg->rc_w_q,&prcg->rc_w_q,&flag);
29: if (!flag) prcg->rc_w_q = PETSC_TRUE;
30: PetscOptionsTail();
31: return 0;
32: }
34: /*
35: KSPSolve_PIPEPRCG - This routine actually applies the pipelined predict and recompute conjugate gradient method
36: */
37: static PetscErrorCode KSPSolve_PIPEPRCG(KSP ksp)
38: {
39: PetscInt i;
40: KSP_CG_PIPE_PR *prcg=(KSP_CG_PIPE_PR*)ksp->data;
41: PetscScalar alpha = 0.0, beta = 0.0, nu = 0.0, nu_old = 0.0, mudelgam[3], *mu_p, *delta_p, *gamma_p;
42: PetscReal dp = 0.0;
43: Vec X,B,R,RT,W,WT,P,S,ST,U,UT,PRTST[3];
44: Mat Amat,Pmat;
45: PetscBool diagonalscale,rc_w_q=prcg->rc_w_q;
47: /* note that these are pointers to entries of muldelgam, different than nu */
48: mu_p=&mudelgam[0];delta_p=&mudelgam[1];gamma_p=&mudelgam[2];
51: PCGetDiagonalScale(ksp->pc,&diagonalscale);
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: PCGetOperators(ksp->pc,&Amat,&Pmat);
68: /* initialize */
69: ksp->its = 0;
70: if (!ksp->guess_zero) {
71: KSP_MatMult(ksp,Amat,X,R); /* r <- b - Ax */
72: VecAYPX(R,-1.0,B);
73: } else {
74: VecCopy(B,R); /* r <- b */
75: }
77: KSP_PCApply(ksp,R,RT); /* rt <- Br */
78: KSP_MatMult(ksp,Amat,RT,W); /* w <- A rt */
79: KSP_PCApply(ksp,W,WT); /* wt <- B w */
81: VecCopy(RT,P); /* p <- rt */
82: VecCopy(W,S); /* p <- rt */
83: VecCopy(WT,ST); /* p <- rt */
85: KSP_MatMult(ksp,Amat,ST,U); /* u <- Ast */
86: KSP_PCApply(ksp,U,UT); /* ut <- Bu */
88: VecDotBegin(RT,R,&nu);
89: VecDotBegin(P,S,mu_p);
90: VecDotBegin(ST,S,gamma_p);
92: VecDotEnd(RT,R,&nu); /* nu <- (rt,r) */
93: VecDotEnd(P,S,mu_p); /* mu <- (p,s) */
94: 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: VecNormBegin(RT,NORM_2,&dp);
103: PetscCommSplitReductionBegin(PetscObjectComm((PetscObject)RT));
104: VecNormEnd(RT,NORM_2,&dp);
105: break;
106: case KSP_NORM_UNPRECONDITIONED:
107: VecNormBegin(R,NORM_2,&dp);
108: PetscCommSplitReductionBegin(PetscObjectComm((PetscObject)R));
109: 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: SETERRQ(PetscObjectComm((PetscObject)ksp),PETSC_ERR_SUP,"%s",KSPNormTypes[ksp->normtype]);
118: }
120: ksp->rnorm = dp;
121: KSPLogResidualHistory(ksp,dp);
122: KSPMonitor(ksp,i,dp);
123: (*ksp->converged)(ksp,i,dp,&ksp->reason,ksp->cnvP);
124: if (ksp->reason) return 0;
126: /* update scalars */
127: alpha = nu / *mu_p;
128: nu_old = nu;
129: nu = nu_old - 2.*alpha*(*delta_p) + (alpha*alpha)*(*gamma_p);
130: beta = nu/nu_old;
132: /* update vectors */
133: VecAXPY(X, alpha,P); /* x <- x + alpha * p */
134: VecAXPY(R,-alpha,S); /* r <- r - alpha * s */
135: VecAXPY(RT,-alpha,ST); /* rt <- rt - alpha * st */
136: VecAXPY(W,-alpha,U); /* w <- w - alpha * u */
137: VecAXPY(WT,-alpha,UT); /* wt <- wt - alpha * ut */
138: VecAYPX(P,beta,RT); /* p <- rt + beta * p */
139: VecAYPX(S,beta,W); /* s <- w + beta * s */
140: VecAYPX(ST,beta,WT); /* st <- wt + beta * st */
142: VecDotBegin(RT,R,&nu);
144: PRTST[0] = P; PRTST[1] = RT; PRTST[2] = ST;
146: VecMDotBegin(S,3,PRTST,mudelgam);
148: PetscCommSplitReductionBegin(PetscObjectComm((PetscObject)R));
150: KSP_MatMult(ksp,Amat,ST,U); /* u <- A st */
151: KSP_PCApply(ksp,U,UT); /* ut <- B u */
153: /* predict-and-recompute */
154: /* ideally this is combined with the previous matvec; i.e. equivalent of MDot */
155: if (rc_w_q) {
156: KSP_MatMult(ksp,Amat,RT,W); /* w <- A rt */
157: KSP_PCApply(ksp,W,WT); /* wt <- B w */
158: }
160: VecDotEnd(RT,R,&nu);
161: VecMDotEnd(S,3,PRTST,mudelgam);
163: i++;
164: ksp->its = i;
166: } while (i<=ksp->max_it);
167: if (!ksp->reason) ksp->reason = KSP_DIVERGED_ITS;
168: return 0;
169: }
171: /*MC
172: KSPPIPEPRCG - Pipelined predict-and-recompute conjugate gradient method.
174: This method has only a single non-blocking reduction per iteration, compared to 2 blocking for standard CG.
175: The non-blocking reduction is overlapped by the matrix-vector product and preconditioner application.
177: Level: intermediate
179: Notes:
180: MPI configuration may be necessary for reductions to make asynchronous progress, which is important for performance of pipelined methods.
181: See the FAQ on the PETSc website for details.
183: Contributed by:
184: Tyler Chen, University of Washington, Applied Mathematics Department
186: Reference:
187: Tyler Chen and Erin Carson. "Predict-and-recompute conjugate gradient variants." SIAM Journal on Scientific Computing 42.5 (2020): A3084-A3108.
189: Acknowledgments:
190: This material is based upon work supported by the National Science Foundation Graduate Research Fellowship Program under Grant No. DGE-1762114. Any opinions, findings, and conclusions or recommendations expressed in this material are those of the author and do not necessarily reflect the views of the National Science Foundation.
192: .seealso: KSPCreate(), KSPSetType(), KSPPIPECG, KSPPIPECR, KSPGROPPCG, KSPPGMRES, KSPCG, KSPCGUseSingleReduction()
193: M*/
194: PETSC_EXTERN PetscErrorCode KSPCreate_PIPEPRCG(KSP ksp)
195: {
196: KSP_CG_PIPE_PR *prcg=NULL;
197: PetscBool cite=PETSC_FALSE;
200: 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 archivePrefix = {arXiv},\n primaryClass = {cs.NA}\n}",&cite);
202: PetscNewLog(ksp,&prcg);
203: ksp->data = (void*)prcg;
205: KSPSetSupportedNorm(ksp,KSP_NORM_UNPRECONDITIONED,PC_LEFT,2);
206: KSPSetSupportedNorm(ksp,KSP_NORM_PRECONDITIONED,PC_LEFT,2);
207: KSPSetSupportedNorm(ksp,KSP_NORM_NATURAL,PC_LEFT,2);
208: KSPSetSupportedNorm(ksp,KSP_NORM_NONE,PC_LEFT,1);
210: ksp->ops->setup = KSPSetUp_PIPEPRCG;
211: ksp->ops->solve = KSPSolve_PIPEPRCG;
212: ksp->ops->destroy = KSPDestroyDefault;
213: ksp->ops->view = NULL;
214: ksp->ops->setfromoptions = KSPSetFromOptions_PIPEPRCG;
215: ksp->ops->buildsolution = KSPBuildSolutionDefault;
216: ksp->ops->buildresidual = KSPBuildResidualDefault;
217: return 0;
218: }