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: }