Actual source code: pipeprcg.c

petsc-master 2020-02-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

 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 application.

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