Actual source code: pipeprcg.c
petsc-3.14.3 2021-01-09
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: {
19: /* get work vectors needed by PIPEPRCG */
20: KSPSetWorkVecs(ksp,9);
22: return(0);
23: }
25: static PetscErrorCode KSPSetFromOptions_PIPEPRCG(PetscOptionItems *PetscOptionsObject,KSP ksp)
26: {
27: PetscInt ierr=0;
28: KSP_CG_PIPE_PR *prcg=(KSP_CG_PIPE_PR*)ksp->data;
29: PetscBool flag=PETSC_FALSE;
32: PetscOptionsHead(PetscOptionsObject,"KSP PIPEPRCG options");
33: PetscOptionsBool("-recompute_w","-recompute w_k with Ar_k? (default = True)","",prcg->rc_w_q,&prcg->rc_w_q,&flag);
34: if (!flag) prcg->rc_w_q = PETSC_TRUE;
35: PetscOptionsTail();
36: return(0);
37: }
39: /*
40: KSPSolve_PIPEPRCG - This routine actually applies the pipelined predict and recompute conjugate gradient method
41: */
42: static PetscErrorCode KSPSolve_PIPEPRCG(KSP ksp)
43: {
45: PetscInt i;
46: KSP_CG_PIPE_PR *prcg=(KSP_CG_PIPE_PR*)ksp->data;
47: PetscScalar alpha = 0.0, beta = 0.0, nu = 0.0, nu_old = 0.0, mudelgam[3], *mu_p, *delta_p, *gamma_p;
48: PetscReal dp = 0.0;
49: Vec X,B,R,RT,W,WT,P,S,ST,U,UT,PRTST[3];
50: Mat Amat,Pmat;
51: PetscBool diagonalscale,rc_w_q=prcg->rc_w_q;
53: /* note that these are pointers to entries of muldelgam, different than nu */
54: mu_p=&mudelgam[0];delta_p=&mudelgam[1];gamma_p=&mudelgam[2];
58: PCGetDiagonalScale(ksp->pc,&diagonalscale);
59: if (diagonalscale) SETERRQ1(PetscObjectComm((PetscObject)ksp),PETSC_ERR_SUP,"Krylov method %s does not support diagonal scaling",((PetscObject)ksp)->type_name);
61: X = ksp->vec_sol;
62: B = ksp->vec_rhs;
63: R = ksp->work[0];
64: RT = ksp->work[1];
65: W = ksp->work[2];
66: WT = ksp->work[3];
67: P = ksp->work[4];
68: S = ksp->work[5];
69: ST = ksp->work[6];
70: U = ksp->work[7];
71: UT = ksp->work[8];
73: PCGetOperators(ksp->pc,&Amat,&Pmat);
75: /* initialize */
76: ksp->its = 0;
77: if (!ksp->guess_zero) {
78: KSP_MatMult(ksp,Amat,X,R); /* r <- b - Ax */
79: VecAYPX(R,-1.0,B);
80: } else {
81: VecCopy(B,R); /* r <- b */
82: }
84: KSP_PCApply(ksp,R,RT); /* rt <- Br */
85: KSP_MatMult(ksp,Amat,RT,W); /* w <- A rt */
86: KSP_PCApply(ksp,W,WT); /* wt <- B w */
88: VecCopy(RT,P); /* p <- rt */
89: VecCopy(W,S); /* p <- rt */
90: VecCopy(WT,ST); /* p <- rt */
92: KSP_MatMult(ksp,Amat,ST,U); /* u <- Ast */
93: KSP_PCApply(ksp,U,UT); /* ut <- Bu */
95: VecDotBegin(RT,R,&nu);
96: VecDotBegin(P,S,mu_p);
97: VecDotBegin(ST,S,gamma_p);
99: VecDotEnd(RT,R,&nu); /* nu <- (rt,r) */
100: VecDotEnd(P,S,mu_p); /* mu <- (p,s) */
101: VecDotEnd(ST,S,gamma_p); /* gamma <- (st,s) */
102: *delta_p = *mu_p;
104: i = 0;
105: do {
106: /* Compute appropriate norm */
107: switch (ksp->normtype) {
108: case KSP_NORM_PRECONDITIONED:
109: VecNormBegin(RT,NORM_2,&dp);
110: PetscCommSplitReductionBegin(PetscObjectComm((PetscObject)RT));
111: VecNormEnd(RT,NORM_2,&dp);
112: break;
113: case KSP_NORM_UNPRECONDITIONED:
114: VecNormBegin(R,NORM_2,&dp);
115: PetscCommSplitReductionBegin(PetscObjectComm((PetscObject)R));
116: VecNormEnd(R,NORM_2,&dp);
117: break;
118: case KSP_NORM_NATURAL:
119: dp = PetscSqrtReal(PetscAbsScalar(nu));
120: break;
121: case KSP_NORM_NONE:
122: dp = 0.0;
123: break;
124: default: SETERRQ1(PetscObjectComm((PetscObject)ksp),PETSC_ERR_SUP,"%s",KSPNormTypes[ksp->normtype]);
125: }
127: ksp->rnorm = dp;
128: KSPLogResidualHistory(ksp,dp);
129: KSPMonitor(ksp,i,dp);
130: (*ksp->converged)(ksp,i,dp,&ksp->reason,ksp->cnvP);
131: if (ksp->reason) return(0);
133: /* update scalars */
134: alpha = nu / *mu_p;
135: nu_old = nu;
136: nu = nu_old - 2*alpha*(*delta_p) + (alpha*alpha)*(*gamma_p);
137: beta = nu/nu_old;
139: /* update vectors */
140: VecAXPY(X, alpha,P); /* x <- x + alpha * p */
141: VecAXPY(R,-alpha,S); /* r <- r - alpha * s */
142: VecAXPY(RT,-alpha,ST); /* rt <- rt - alpha * st */
143: VecAXPY(W,-alpha,U); /* w <- w - alpha * u */
144: VecAXPY(WT,-alpha,UT); /* wt <- wt - alpha * ut */
145: VecAYPX(P,beta,RT); /* p <- rt + beta * p */
146: VecAYPX(S,beta,W); /* s <- w + beta * s */
147: VecAYPX(ST,beta,WT); /* st <- wt + beta * st */
149: VecDotBegin(RT,R,&nu);
151: PRTST[0] = P; PRTST[1] = RT; PRTST[2] = ST;
153: VecMDotBegin(S,3,PRTST,mudelgam);
155: PetscCommSplitReductionBegin(PetscObjectComm((PetscObject)R));
157: KSP_MatMult(ksp,Amat,ST,U); /* u <- A st */
158: KSP_PCApply(ksp,U,UT); /* ut <- B u */
160: /* predict-and-recompute */
161: /* ideally this is combined with the previous matvec; i.e. equivalent of MDot */
162: if (rc_w_q) {
163: KSP_MatMult(ksp,Amat,RT,W); /* w <- A rt */
164: KSP_PCApply(ksp,W,WT); /* wt <- B w */
165: }
167: VecDotEnd(RT,R,&nu);
168: VecMDotEnd(S,3,PRTST,mudelgam);
170: i++;
171: ksp->its = i;
173: } while (i<=ksp->max_it);
174: if (!ksp->reason) ksp->reason = KSP_DIVERGED_ITS;
175: return(0);
176: }
179: /*MC
180: KSPPIPEPRCG - Pipelined predict-and-recompute conjugate gradient method.
182: This method has only a single non-blocking reduction per iteration, compared to 2 blocking for standard CG.
183: The non-blocking reduction is overlapped by the matrix-vector product and preconditioner application.
185: Level: intermediate
187: Notes:
188: MPI configuration may be necessary for reductions to make asynchronous progress, which is important for performance of pipelined methods.
189: See the FAQ on the PETSc website for details.
191: Contributed by:
192: Tyler Chen, University of Washington, Applied Mathematics Department
194: Reference:
195: "Predict-and-recompute conjugate gradient variants". Tyler Chen and Erin C. Carson. In preparation.
197: Acknowledgments:
198: 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.
200: .seealso: KSPCreate(), KSPSetType(), KSPPIPECG, KSPPIPECR, KSPGROPPCG, KSPPGMRES, KSPCG, KSPCGUseSingleReduction()
201: M*/
202: PETSC_EXTERN PetscErrorCode KSPCreate_PIPEPRCG(KSP ksp)
203: {
205: KSP_CG_PIPE_PR *prcg=NULL;
206: PetscBool cite=PETSC_FALSE;
210: 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);
212: PetscNewLog(ksp,&prcg);
213: ksp->data = (void*)prcg;
215: KSPSetSupportedNorm(ksp,KSP_NORM_UNPRECONDITIONED,PC_LEFT,2);
216: KSPSetSupportedNorm(ksp,KSP_NORM_PRECONDITIONED,PC_LEFT,2);
217: KSPSetSupportedNorm(ksp,KSP_NORM_NATURAL,PC_LEFT,2);
218: KSPSetSupportedNorm(ksp,KSP_NORM_NONE,PC_LEFT,1);
220: ksp->ops->setup = KSPSetUp_PIPEPRCG;
221: ksp->ops->solve = KSPSolve_PIPEPRCG;
222: ksp->ops->destroy = KSPDestroyDefault;
223: ksp->ops->view = NULL;
224: ksp->ops->setfromoptions = KSPSetFromOptions_PIPEPRCG;
225: ksp->ops->buildsolution = KSPBuildSolutionDefault;
226: ksp->ops->buildresidual = KSPBuildResidualDefault;
227: return(0);
228: }