Actual source code: pipeprcg.c

petsc-master 2020-09-19
Report Typos and Errors
  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: }