Actual source code: pipeprcg.c
petsc-3.13.6 2020-09-29
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
42: Input Parameter:
43: . ksp - the Krylov space object that was set to use conjugate gradient, by, for
44: example, KSPCreate(MPI_Comm,KSP *ksp); KSPSetType(ksp,KSPCG);
45: */
46: static PetscErrorCode KSPSolve_PIPEPRCG(KSP ksp)
47: {
49: PetscInt i;
50: KSP_CG_PIPE_PR *prcg=(KSP_CG_PIPE_PR*)ksp->data;
51: PetscScalar alpha = 0.0, beta = 0.0, nu = 0.0, nu_old = 0.0, mudelgam[3], *mu_p, *delta_p, *gamma_p;
52: PetscReal dp = 0.0;
53: Vec X,B,R,RT,W,WT,P,S,ST,U,UT,PRTST[3];
54: Mat Amat,Pmat;
55: PetscBool diagonalscale,rc_w_q=prcg->rc_w_q;
57: /* note that these are pointers to entries of muldelgam, different than nu */
58: mu_p=&mudelgam[0];delta_p=&mudelgam[1];gamma_p=&mudelgam[2];
62: PCGetDiagonalScale(ksp->pc,&diagonalscale);
63: if (diagonalscale) SETERRQ1(PetscObjectComm((PetscObject)ksp),PETSC_ERR_SUP,"Krylov method %s does not support diagonal scaling",((PetscObject)ksp)->type_name);
65: X = ksp->vec_sol;
66: B = ksp->vec_rhs;
67: R = ksp->work[0];
68: RT = ksp->work[1];
69: W = ksp->work[2];
70: WT = ksp->work[3];
71: P = ksp->work[4];
72: S = ksp->work[5];
73: ST = ksp->work[6];
74: U = ksp->work[7];
75: UT = ksp->work[8];
77: PCGetOperators(ksp->pc,&Amat,&Pmat);
79: /* initialize */
80: ksp->its = 0;
81: if (!ksp->guess_zero) {
82: KSP_MatMult(ksp,Amat,X,R); /* r <- b - Ax */
83: VecAYPX(R,-1.0,B);
84: } else {
85: VecCopy(B,R); /* r <- b */
86: }
88: KSP_PCApply(ksp,R,RT); /* rt <- Br */
89: KSP_MatMult(ksp,Amat,RT,W); /* w <- A rt */
90: KSP_PCApply(ksp,W,WT); /* wt <- B w */
92: VecCopy(RT,P); /* p <- rt */
93: VecCopy(W,S); /* p <- rt */
94: VecCopy(WT,ST); /* p <- rt */
96: KSP_MatMult(ksp,Amat,ST,U); /* u <- Ast */
97: KSP_PCApply(ksp,U,UT); /* ut <- Bu */
99: VecDotBegin(RT,R,&nu);
100: VecDotBegin(P,S,mu_p);
101: VecDotBegin(ST,S,gamma_p);
103: VecDotEnd(RT,R,&nu); /* nu <- (rt,r) */
104: VecDotEnd(P,S,mu_p); /* mu <- (p,s) */
105: VecDotEnd(ST,S,gamma_p); /* gamma <- (st,s) */
106: *delta_p = *mu_p;
108: i = 0;
109: do {
111: /* Compute appropriate norm */
112: switch (ksp->normtype) {
113: case KSP_NORM_PRECONDITIONED:
114: VecNormBegin(RT,NORM_2,&dp);
115: PetscCommSplitReductionBegin(PetscObjectComm((PetscObject)RT));
116: VecNormEnd(RT,NORM_2,&dp);
117: break;
118: case KSP_NORM_UNPRECONDITIONED:
119: VecNormBegin(R,NORM_2,&dp);
120: PetscCommSplitReductionBegin(PetscObjectComm((PetscObject)R));
121: VecNormEnd(R,NORM_2,&dp);
122: break;
123: case KSP_NORM_NATURAL:
124: dp = PetscSqrtReal(PetscAbsScalar(nu));
125: break;
126: case KSP_NORM_NONE:
127: dp = 0.0;
128: break;
129: default: SETERRQ1(PetscObjectComm((PetscObject)ksp),PETSC_ERR_SUP,"%s",KSPNormTypes[ksp->normtype]);
130: }
132: ksp->rnorm = dp;
133: KSPLogResidualHistory(ksp,dp);
134: KSPMonitor(ksp,i,dp);
135: (*ksp->converged)(ksp,i,dp,&ksp->reason,ksp->cnvP);
136: if (ksp->reason) break;
138: /* update scalars */
139: alpha = nu / *mu_p;
140: nu_old = nu;
141: nu = nu_old - 2*alpha*(*delta_p) + (alpha*alpha)*(*gamma_p);
142: beta = nu/nu_old;
144: /* update vectors */
145: VecAXPY(X, alpha,P); /* x <- x + alpha * p */
146: VecAXPY(R,-alpha,S); /* r <- r - alpha * s */
147: VecAXPY(RT,-alpha,ST); /* rt <- rt - alpha * st */
148: VecAXPY(W,-alpha,U); /* w <- w - alpha * u */
149: VecAXPY(WT,-alpha,UT); /* wt <- wt - alpha * ut */
150: VecAYPX(P,beta,RT); /* p <- rt + beta * p */
151: VecAYPX(S,beta,W); /* s <- w + beta * s */
152: VecAYPX(ST,beta,WT); /* st <- wt + beta * st */
154: VecDotBegin(RT,R,&nu);
156: PRTST[0] = P; PRTST[1] = RT; PRTST[2] = ST;
158: VecMDotBegin(S,3,PRTST,mudelgam);
160: PetscCommSplitReductionBegin(PetscObjectComm((PetscObject)R));
162: KSP_MatMult(ksp,Amat,ST,U); /* u <- A st */
163: KSP_PCApply(ksp,U,UT); /* ut <- B u */
165: /* predict-and-recompute */
166: /* ideally this is combined with the previous matvec; i.e. equivalent of MDot */
167: if ( rc_w_q ) {
168: KSP_MatMult(ksp,Amat,RT,W); /* w <- A rt */
169: KSP_PCApply(ksp,W,WT); /* wt <- B w */
170: }
172: VecDotEnd(RT,R,&nu);
173: VecMDotEnd(S,3,PRTST,mudelgam);
175: i++;
176: ksp->its = i;
178: } while (i<ksp->max_it);
179: if (i >= ksp->max_it) ksp->reason = KSP_DIVERGED_ITS;
180: return(0);
181: }
184: /*MC
185: KSPPIPEPRCG - Pipelined predict-and-recompute conjugate gradient method.
187: This method has only a single non-blocking reduction per iteration, compared to 2 blocking for standard CG.
188: The non-blocking reduction is overlapped by the matrix-vector product and preconditioner Section 1.5 Writing Application Codes with PETSc.
190: Level: intermediate
192: Notes:
193: MPI configuration may be necessary for reductions to make asynchronous progress, which is important for performance of pipelined methods.
194: See the FAQ on the PETSc website for details.
196: Contributed by:
197: Tyler Chen, University of Washington, Applied Mathematics Department
199: Reference:
200: "Predict-and-recompute conjugate gradient variants". Tyler Chen and Erin C. Carson. In preparation.
202: Acknowledgments:
203: 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.
205: .seealso: KSPCreate(), KSPSetType(), KSPPIPECG, KSPPIPECR, KSPGROPPCG, KSPPGMRES, KSPCG, KSPCGUseSingleReduction()
206: M*/
207: PETSC_EXTERN PetscErrorCode KSPCreate_PIPEPRCG(KSP ksp)
208: {
210: KSP_CG_PIPE_PR *prcg=NULL;
211: PetscBool cite=PETSC_FALSE;
215: 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);
217: PetscNewLog(ksp,&prcg);
218: ksp->data = (void*)prcg;
220: KSPSetSupportedNorm(ksp,KSP_NORM_UNPRECONDITIONED,PC_LEFT,2);
221: KSPSetSupportedNorm(ksp,KSP_NORM_PRECONDITIONED,PC_LEFT,2);
222: KSPSetSupportedNorm(ksp,KSP_NORM_NATURAL,PC_LEFT,2);
223: KSPSetSupportedNorm(ksp,KSP_NORM_NONE,PC_LEFT,1);
225: ksp->ops->setup = KSPSetUp_PIPEPRCG;
226: ksp->ops->solve = KSPSolve_PIPEPRCG;
227: ksp->ops->destroy = KSPDestroyDefault;
228: ksp->ops->view = 0;
229: ksp->ops->setfromoptions = KSPSetFromOptions_PIPEPRCG;
230: ksp->ops->buildsolution = KSPBuildSolutionDefault;
231: ksp->ops->buildresidual = KSPBuildResidualDefault;
232: return(0);
233: }