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

     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:    Tyler Chen and Erin Carson. "Predict-and-recompute conjugate gradient variants." SIAM Journal on Scientific Computing 42.5 (2020): A3084-A3108.

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