Date: 2025-03-29T05:58:26+00:00

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 recomputed 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:   PetscFunctionBegin;
     17:   /* get work vectors needed by PIPEPRCG */
     18:   PetscCall(KSPSetWorkVecs(ksp, 9));
     19:   PetscFunctionReturn(PETSC_SUCCESS);
     20: }

     22: static PetscErrorCode KSPSetFromOptions_PIPEPRCG(KSP ksp, PetscOptionItems PetscOptionsObject)
     23: {
     24:   KSP_CG_PIPE_PR *prcg = (KSP_CG_PIPE_PR *)ksp->data;
     25:   PetscBool       flag = PETSC_FALSE;

     27:   PetscFunctionBegin;
     28:   PetscOptionsHeadBegin(PetscOptionsObject, "KSP PIPEPRCG options");
     29:   PetscCall(PetscOptionsBool("-recompute_w", "-recompute w_k with Ar_k? (default = True)", "", prcg->rc_w_q, &prcg->rc_w_q, &flag));
     30:   if (!flag) prcg->rc_w_q = PETSC_TRUE;
     31:   PetscOptionsHeadEnd();
     32:   PetscFunctionReturn(PETSC_SUCCESS);
     33: }

     35: /*
     36:  KSPSolve_PIPEPRCG - This routine actually applies the pipelined predict and recompute conjugate gradient method
     37: */
     38: static PetscErrorCode KSPSolve_PIPEPRCG(KSP ksp)
     39: {
     40:   PetscInt        i;
     41:   KSP_CG_PIPE_PR *prcg  = (KSP_CG_PIPE_PR *)ksp->data;
     42:   PetscScalar     alpha = 0.0, beta = 0.0, nu = 0.0, nu_old = 0.0, mudelgam[3], *mu_p, *delta_p, *gamma_p;
     43:   PetscReal       dp = 0.0;
     44:   Vec             X, B, R, RT, W, WT, P, S, ST, U, UT, PRTST[3];
     45:   Mat             Amat, Pmat;
     46:   PetscBool       diagonalscale, rc_w_q = prcg->rc_w_q;

     48:   /* note that these are pointers to entries of muldelgam, different than nu */
     49:   mu_p    = &mudelgam[0];
     50:   delta_p = &mudelgam[1];
     51:   gamma_p = &mudelgam[2];

     53:   PetscFunctionBegin;
     54:   PetscCall(PCGetDiagonalScale(ksp->pc, &diagonalscale));
     55:   PetscCheck(!diagonalscale, PetscObjectComm((PetscObject)ksp), PETSC_ERR_SUP, "Krylov method %s does not support diagonal scaling", ((PetscObject)ksp)->type_name);

     57:   X  = ksp->vec_sol;
     58:   B  = ksp->vec_rhs;
     59:   R  = ksp->work[0];
     60:   RT = ksp->work[1];
     61:   W  = ksp->work[2];
     62:   WT = ksp->work[3];
     63:   P  = ksp->work[4];
     64:   S  = ksp->work[5];
     65:   ST = ksp->work[6];
     66:   U  = ksp->work[7];
     67:   UT = ksp->work[8];

     69:   PetscCall(PCGetOperators(ksp->pc, &Amat, &Pmat));

     71:   /* initialize */
     72:   ksp->its = 0;
     73:   if (!ksp->guess_zero) {
     74:     PetscCall(KSP_MatMult(ksp, Amat, X, R)); /*   r <- b - Ax  */
     75:     PetscCall(VecAYPX(R, -1.0, B));
     76:   } else {
     77:     PetscCall(VecCopy(B, R)); /*   r <- b       */
     78:   }

     80:   PetscCall(KSP_PCApply(ksp, R, RT));       /*   rt <- Br     */
     81:   PetscCall(KSP_MatMult(ksp, Amat, RT, W)); /*   w <- A rt    */
     82:   PetscCall(KSP_PCApply(ksp, W, WT));       /*   wt <- B w    */

     84:   PetscCall(VecCopy(RT, P));  /*   p <- rt      */
     85:   PetscCall(VecCopy(W, S));   /*   p <- rt      */
     86:   PetscCall(VecCopy(WT, ST)); /*   p <- rt      */

     88:   PetscCall(KSP_MatMult(ksp, Amat, ST, U)); /*   u <- Ast     */
     89:   PetscCall(KSP_PCApply(ksp, U, UT));       /*   ut <- Bu     */

     91:   PetscCall(VecDotBegin(RT, R, &nu));
     92:   PetscCall(VecDotBegin(P, S, mu_p));
     93:   PetscCall(VecDotBegin(ST, S, gamma_p));

     95:   PetscCall(VecDotEnd(RT, R, &nu));     /*   nu    <- (rt,r)  */
     96:   PetscCall(VecDotEnd(P, S, mu_p));     /*   mu    <- (p,s)   */
     97:   PetscCall(VecDotEnd(ST, S, gamma_p)); /*   gamma <- (st,s)  */
     98:   *delta_p = *mu_p;

    100:   i = 0;
    101:   do {
    102:     /* Compute appropriate norm */
    103:     switch (ksp->normtype) {
    104:     case KSP_NORM_PRECONDITIONED:
    105:       PetscCall(VecNormBegin(RT, NORM_2, &dp));
    106:       PetscCall(PetscCommSplitReductionBegin(PetscObjectComm((PetscObject)RT)));
    107:       PetscCall(VecNormEnd(RT, NORM_2, &dp));
    108:       break;
    109:     case KSP_NORM_UNPRECONDITIONED:
    110:       PetscCall(VecNormBegin(R, NORM_2, &dp));
    111:       PetscCall(PetscCommSplitReductionBegin(PetscObjectComm((PetscObject)R)));
    112:       PetscCall(VecNormEnd(R, NORM_2, &dp));
    113:       break;
    114:     case KSP_NORM_NATURAL:
    115:       dp = PetscSqrtReal(PetscAbsScalar(nu));
    116:       break;
    117:     case KSP_NORM_NONE:
    118:       dp = 0.0;
    119:       break;
    120:     default:
    121:       SETERRQ(PetscObjectComm((PetscObject)ksp), PETSC_ERR_SUP, "%s", KSPNormTypes[ksp->normtype]);
    122:     }

    124:     ksp->rnorm = dp;
    125:     PetscCall(KSPLogResidualHistory(ksp, dp));
    126:     PetscCall(KSPMonitor(ksp, i, dp));
    127:     PetscCall((*ksp->converged)(ksp, i, dp, &ksp->reason, ksp->cnvP));
    128:     if (ksp->reason) PetscFunctionReturn(PETSC_SUCCESS);

    130:     /* update scalars */
    131:     alpha  = nu / *mu_p;
    132:     nu_old = nu;
    133:     nu     = nu_old - 2. * alpha * (*delta_p) + (alpha * alpha) * (*gamma_p);
    134:     beta   = nu / nu_old;

    136:     /* update vectors */
    137:     PetscCall(VecAXPY(X, alpha, P));    /*   x  <- x  + alpha * p   */
    138:     PetscCall(VecAXPY(R, -alpha, S));   /*   r  <- r  - alpha * s   */
    139:     PetscCall(VecAXPY(RT, -alpha, ST)); /*   rt <- rt - alpha * st  */
    140:     PetscCall(VecAXPY(W, -alpha, U));   /*   w  <- w  - alpha * u   */
    141:     PetscCall(VecAXPY(WT, -alpha, UT)); /*   wt <- wt - alpha * ut  */
    142:     PetscCall(VecAYPX(P, beta, RT));    /*   p  <- rt + beta  * p   */
    143:     PetscCall(VecAYPX(S, beta, W));     /*   s  <- w  + beta  * s   */
    144:     PetscCall(VecAYPX(ST, beta, WT));   /*   st <- wt + beta  * st  */

    146:     PetscCall(VecDotBegin(RT, R, &nu));

    148:     PRTST[0] = P;
    149:     PRTST[1] = RT;
    150:     PRTST[2] = ST;

    152:     PetscCall(VecMDotBegin(S, 3, PRTST, mudelgam));

    154:     PetscCall(PetscCommSplitReductionBegin(PetscObjectComm((PetscObject)R)));

    156:     PetscCall(KSP_MatMult(ksp, Amat, ST, U)); /*   u  <- A st             */
    157:     PetscCall(KSP_PCApply(ksp, U, UT));       /*   ut <- B u              */

    159:     /* predict-and-recompute */
    160:     /* ideally this is combined with the previous matvec; i.e. equivalent of MDot */
    161:     if (rc_w_q) {
    162:       PetscCall(KSP_MatMult(ksp, Amat, RT, W)); /*   w  <- A rt             */
    163:       PetscCall(KSP_PCApply(ksp, W, WT));       /*   wt <- B w              */
    164:     }

    166:     PetscCall(VecDotEnd(RT, R, &nu));
    167:     PetscCall(VecMDotEnd(S, 3, PRTST, mudelgam));

    169:     i++;
    170:     ksp->its = i;

    172:   } while (i <= ksp->max_it);
    173:   if (!ksp->reason) ksp->reason = KSP_DIVERGED_ITS;
    174:   PetscFunctionReturn(PETSC_SUCCESS);
    175: }

    177: /*MC
    178:    KSPPIPEPRCG - Pipelined predict-and-recompute conjugate gradient Krylov method {cite}`chen2020predict`. [](sec_pipelineksp)

    180:    Options Database Key:
    181: .  -ksp_pipeprcg_recompute_w - recompute the $w_k$ with $Ar_k$, default is true

    183:    Level: intermediate

    185:    Notes:
    186:    This method has only a single non-blocking reduction per iteration, compared to 2 blocking for standard `KSPCG`.
    187:    The non-blocking reduction is overlapped by the matrix-vector product and preconditioner application.

    189:    MPI configuration may be necessary for reductions to make asynchronous progress, which is important for performance of pipelined methods.
    190:    See [](doc_faq_pipelined)

    192:    Contributed by:
    193:    Tyler Chen, University of Washington, Applied Mathematics Department

    195:    Acknowledgments:
    196:    This material is based upon work supported by the National Science Foundation Graduate Research Fellowship Program under Grant No. DGE-1762114.
    197:    Any opinions, findings, and conclusions or recommendations expressed in this material are those of the author and do not necessarily
    198:    reflect the views of the National Science Foundation.

    200: .seealso: [](ch_ksp), [](doc_faq_pipelined), [](sec_pipelineksp), `KSPCreate()`, `KSPSetType()`, `KSPCG`, `KSPPIPECG`, `KSPPIPECR`, `KSPGROPPCG`, `KSPPGMRES`, `KSPCG`, `KSPCGUseSingleReduction()`
    201: M*/
    202: PETSC_EXTERN PetscErrorCode KSPCreate_PIPEPRCG(KSP ksp)
    203: {
    204:   KSP_CG_PIPE_PR *prcg = NULL;
    205:   PetscBool       cite = PETSC_FALSE;

    207:   PetscFunctionBegin;
    208:   PetscCall(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  "
    209:                                    "archivePrefix = {arXiv},\n  primaryClass = {cs.NA}\n}",
    210:                                    &cite));

    212:   PetscCall(PetscNew(&prcg));
    213:   ksp->data = (void *)prcg;

    215:   PetscCall(KSPSetSupportedNorm(ksp, KSP_NORM_UNPRECONDITIONED, PC_LEFT, 2));
    216:   PetscCall(KSPSetSupportedNorm(ksp, KSP_NORM_PRECONDITIONED, PC_LEFT, 2));
    217:   PetscCall(KSPSetSupportedNorm(ksp, KSP_NORM_NATURAL, PC_LEFT, 2));
    218:   PetscCall(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:   PetscFunctionReturn(PETSC_SUCCESS);
    228: }