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