Actual source code: pch2opus.c

      1: #include <petsc/private/pcimpl.h>
      2: #include <petsc/private/matimpl.h>
      3: #include <petscdm.h>
      4: #include <h2opusconf.h>

      6: /* Use GPU only if H2OPUS is configured for GPU */
      7: #if defined(PETSC_HAVE_CUDA) && defined(H2OPUS_USE_GPU)
      8:   #define PETSC_H2OPUS_USE_GPU
      9: #endif

     11: typedef struct {
     12:   Mat         A;
     13:   Mat         M;
     14:   PetscScalar s0;

     16:   /* sampler for Newton-Schultz */
     17:   Mat      S;
     18:   PetscInt hyperorder;
     19:   Vec      wns[4];
     20:   Mat      wnsmat[4];

     22:   /* convergence testing */
     23:   Mat T;
     24:   Vec w;

     26:   /* Support for PCSetCoordinates */
     27:   PetscInt   sdim;
     28:   PetscInt   nlocc;
     29:   PetscReal *coords;

     31:   /* Newton-Schultz customization */
     32:   PetscInt  maxits;
     33:   PetscReal rtol, atol;
     34:   PetscBool monitor;
     35:   PetscBool useapproximatenorms;
     36:   NormType  normtype;

     38:   /* Used when pmat != MATH2OPUS */
     39:   PetscReal eta;
     40:   PetscInt  leafsize;
     41:   PetscInt  max_rank;
     42:   PetscInt  bs;
     43:   PetscReal mrtol;

     45:   /* CPU/GPU */
     46:   PetscBool forcecpu;
     47:   PetscBool boundtocpu;
     48: } PC_H2OPUS;

     50: PETSC_EXTERN PetscErrorCode MatNorm_H2OPUS(Mat, NormType, PetscReal *);

     52: static PetscErrorCode PCH2OpusInferCoordinates_Private(PC pc)
     53: {
     54:   PC_H2OPUS *pch2opus = (PC_H2OPUS *)pc->data;
     55:   DM         dm;
     56:   PetscBool  isdmda;

     58:   PetscFunctionBegin;
     59:   if (pch2opus->sdim) PetscFunctionReturn(PETSC_SUCCESS);
     60:   PetscCall(PCGetDM(pc, &dm));
     61:   if (!dm) PetscCall(MatGetDM(pc->useAmat ? pc->mat : pc->pmat, &dm));
     62:   PetscCall(PetscObjectTypeCompare((PetscObject)dm, DMDA, &isdmda));
     63:   if (isdmda) {
     64:     Vec                c;
     65:     const PetscScalar *coords;
     66:     PetscInt           n, sdim;

     68:     PetscCall(DMGetCoordinates(dm, &c));
     69:     PetscCall(DMGetDimension(dm, &sdim));
     70:     PetscCall(VecGetLocalSize(c, &n));
     71:     PetscCall(VecGetArrayRead(c, &coords));
     72:     PetscCall(PCSetCoordinates(pc, sdim, n / sdim, (PetscScalar *)coords));
     73:     PetscCall(VecRestoreArrayRead(c, &coords));
     74:   }
     75:   PetscFunctionReturn(PETSC_SUCCESS);
     76: }

     78: static PetscErrorCode PCReset_H2OPUS(PC pc)
     79: {
     80:   PC_H2OPUS *pch2opus = (PC_H2OPUS *)pc->data;

     82:   PetscFunctionBegin;
     83:   pch2opus->sdim  = 0;
     84:   pch2opus->nlocc = 0;
     85:   PetscCall(PetscFree(pch2opus->coords));
     86:   PetscFunctionReturn(PETSC_SUCCESS);
     87: }

     89: static PetscErrorCode PCSetCoordinates_H2OPUS(PC pc, PetscInt sdim, PetscInt nlocc, PetscReal *coords)
     90: {
     91:   PC_H2OPUS *pch2opus = (PC_H2OPUS *)pc->data;
     92:   PetscBool  reset    = PETSC_TRUE;

     94:   PetscFunctionBegin;
     95:   if (pch2opus->sdim && sdim == pch2opus->sdim && nlocc == pch2opus->nlocc) {
     96:     PetscCall(PetscArraycmp(pch2opus->coords, coords, sdim * nlocc, &reset));
     97:     reset = (PetscBool)!reset;
     98:   }
     99:   PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, &reset, 1, MPIU_BOOL, MPI_LOR, PetscObjectComm((PetscObject)pc)));
    100:   if (reset) {
    101:     PetscCall(PCReset_H2OPUS(pc));
    102:     PetscCall(PetscMalloc1(sdim * nlocc, &pch2opus->coords));
    103:     PetscCall(PetscArraycpy(pch2opus->coords, coords, sdim * nlocc));
    104:     pch2opus->sdim  = sdim;
    105:     pch2opus->nlocc = nlocc;
    106:   }
    107:   PetscFunctionReturn(PETSC_SUCCESS);
    108: }

    110: static PetscErrorCode PCDestroy_H2OPUS(PC pc)
    111: {
    112:   PC_H2OPUS *pch2opus = (PC_H2OPUS *)pc->data;

    114:   PetscFunctionBegin;
    115:   PetscCall(PCReset_H2OPUS(pc));
    116:   PetscCall(MatDestroy(&pch2opus->A));
    117:   PetscCall(MatDestroy(&pch2opus->M));
    118:   PetscCall(MatDestroy(&pch2opus->T));
    119:   PetscCall(VecDestroy(&pch2opus->w));
    120:   PetscCall(MatDestroy(&pch2opus->S));
    121:   PetscCall(VecDestroy(&pch2opus->wns[0]));
    122:   PetscCall(VecDestroy(&pch2opus->wns[1]));
    123:   PetscCall(VecDestroy(&pch2opus->wns[2]));
    124:   PetscCall(VecDestroy(&pch2opus->wns[3]));
    125:   PetscCall(MatDestroy(&pch2opus->wnsmat[0]));
    126:   PetscCall(MatDestroy(&pch2opus->wnsmat[1]));
    127:   PetscCall(MatDestroy(&pch2opus->wnsmat[2]));
    128:   PetscCall(MatDestroy(&pch2opus->wnsmat[3]));
    129:   PetscCall(PetscObjectComposeFunction((PetscObject)pc, "PCSetCoordinates_C", NULL));
    130:   PetscCall(PetscFree(pc->data));
    131:   PetscFunctionReturn(PETSC_SUCCESS);
    132: }

    134: static PetscErrorCode PCSetFromOptions_H2OPUS(PC pc, PetscOptionItems *PetscOptionsObject)
    135: {
    136:   PC_H2OPUS *pch2opus = (PC_H2OPUS *)pc->data;

    138:   PetscFunctionBegin;
    139:   PetscOptionsHeadBegin(PetscOptionsObject, "H2OPUS options");
    140:   PetscCall(PetscOptionsInt("-pc_h2opus_maxits", "Maximum number of iterations for Newton-Schultz", NULL, pch2opus->maxits, &pch2opus->maxits, NULL));
    141:   PetscCall(PetscOptionsBool("-pc_h2opus_monitor", "Monitor Newton-Schultz convergence", NULL, pch2opus->monitor, &pch2opus->monitor, NULL));
    142:   PetscCall(PetscOptionsReal("-pc_h2opus_atol", "Absolute tolerance", NULL, pch2opus->atol, &pch2opus->atol, NULL));
    143:   PetscCall(PetscOptionsReal("-pc_h2opus_rtol", "Relative tolerance", NULL, pch2opus->rtol, &pch2opus->rtol, NULL));
    144:   PetscCall(PetscOptionsEnum("-pc_h2opus_norm_type", "Norm type for convergence monitoring", NULL, NormTypes, (PetscEnum)pch2opus->normtype, (PetscEnum *)&pch2opus->normtype, NULL));
    145:   PetscCall(PetscOptionsInt("-pc_h2opus_hyperorder", "Hyper power order of sampling", NULL, pch2opus->hyperorder, &pch2opus->hyperorder, NULL));
    146:   PetscCall(PetscOptionsInt("-pc_h2opus_leafsize", "Leaf size when constructed from kernel", NULL, pch2opus->leafsize, &pch2opus->leafsize, NULL));
    147:   PetscCall(PetscOptionsReal("-pc_h2opus_eta", "Admissibility condition tolerance", NULL, pch2opus->eta, &pch2opus->eta, NULL));
    148:   PetscCall(PetscOptionsInt("-pc_h2opus_maxrank", "Maximum rank when constructed from matvecs", NULL, pch2opus->max_rank, &pch2opus->max_rank, NULL));
    149:   PetscCall(PetscOptionsInt("-pc_h2opus_samples", "Number of samples to be taken concurrently when constructing from matvecs", NULL, pch2opus->bs, &pch2opus->bs, NULL));
    150:   PetscCall(PetscOptionsReal("-pc_h2opus_mrtol", "Relative tolerance for construction from sampling", NULL, pch2opus->mrtol, &pch2opus->mrtol, NULL));
    151:   PetscCall(PetscOptionsBool("-pc_h2opus_forcecpu", "Force construction on CPU", NULL, pch2opus->forcecpu, &pch2opus->forcecpu, NULL));
    152:   PetscOptionsHeadEnd();
    153:   PetscFunctionReturn(PETSC_SUCCESS);
    154: }

    156: typedef struct {
    157:   Mat A;
    158:   Mat M;
    159:   Vec w;
    160: } AAtCtx;

    162: static PetscErrorCode MatMult_AAt(Mat A, Vec x, Vec y)
    163: {
    164:   AAtCtx *aat;

    166:   PetscFunctionBegin;
    167:   PetscCall(MatShellGetContext(A, (void *)&aat));
    168:   /* PetscCall(MatMultTranspose(aat->M,x,aat->w)); */
    169:   PetscCall(MatMultTranspose(aat->A, x, aat->w));
    170:   PetscCall(MatMult(aat->A, aat->w, y));
    171:   PetscFunctionReturn(PETSC_SUCCESS);
    172: }

    174: static PetscErrorCode PCH2OpusSetUpInit(PC pc)
    175: {
    176:   PC_H2OPUS *pch2opus = (PC_H2OPUS *)pc->data;
    177:   Mat        A        = pc->useAmat ? pc->mat : pc->pmat, AAt;
    178:   PetscInt   M, m;
    179:   VecType    vtype;
    180:   PetscReal  n;
    181:   AAtCtx     aat;

    183:   PetscFunctionBegin;
    184:   aat.A = A;
    185:   aat.M = pch2opus->M; /* unused so far */
    186:   PetscCall(MatCreateVecs(A, NULL, &aat.w));
    187:   PetscCall(MatGetSize(A, &M, NULL));
    188:   PetscCall(MatGetLocalSize(A, &m, NULL));
    189:   PetscCall(MatCreateShell(PetscObjectComm((PetscObject)A), m, m, M, M, &aat, &AAt));
    190:   PetscCall(MatBindToCPU(AAt, pch2opus->boundtocpu));
    191:   PetscCall(MatShellSetOperation(AAt, MATOP_MULT, (void (*)(void))MatMult_AAt));
    192:   PetscCall(MatShellSetOperation(AAt, MATOP_MULT_TRANSPOSE, (void (*)(void))MatMult_AAt));
    193:   PetscCall(MatShellSetOperation(AAt, MATOP_NORM, (void (*)(void))MatNorm_H2OPUS));
    194:   PetscCall(MatGetVecType(A, &vtype));
    195:   PetscCall(MatShellSetVecType(AAt, vtype));
    196:   PetscCall(MatNorm(AAt, NORM_1, &n));
    197:   pch2opus->s0 = 1. / n;
    198:   PetscCall(VecDestroy(&aat.w));
    199:   PetscCall(MatDestroy(&AAt));
    200:   PetscFunctionReturn(PETSC_SUCCESS);
    201: }

    203: static PetscErrorCode PCApplyKernel_H2OPUS(PC pc, Vec x, Vec y, PetscBool t)
    204: {
    205:   PC_H2OPUS *pch2opus = (PC_H2OPUS *)pc->data;

    207:   PetscFunctionBegin;
    208:   if (t) PetscCall(MatMultTranspose(pch2opus->M, x, y));
    209:   else PetscCall(MatMult(pch2opus->M, x, y));
    210:   PetscFunctionReturn(PETSC_SUCCESS);
    211: }

    213: static PetscErrorCode PCApplyMatKernel_H2OPUS(PC pc, Mat X, Mat Y, PetscBool t)
    214: {
    215:   PC_H2OPUS *pch2opus = (PC_H2OPUS *)pc->data;

    217:   PetscFunctionBegin;
    218:   if (t) PetscCall(MatTransposeMatMult(pch2opus->M, X, MAT_REUSE_MATRIX, PETSC_DETERMINE, &Y));
    219:   else PetscCall(MatMatMult(pch2opus->M, X, MAT_REUSE_MATRIX, PETSC_CURRENT, &Y));
    220:   PetscFunctionReturn(PETSC_SUCCESS);
    221: }

    223: static PetscErrorCode PCApplyMat_H2OPUS(PC pc, Mat X, Mat Y)
    224: {
    225:   PetscFunctionBegin;
    226:   PetscCall(PCApplyMatKernel_H2OPUS(pc, X, Y, PETSC_FALSE));
    227:   PetscFunctionReturn(PETSC_SUCCESS);
    228: }

    230: static PetscErrorCode PCApplyTransposeMat_H2OPUS(PC pc, Mat X, Mat Y)
    231: {
    232:   PetscFunctionBegin;
    233:   PetscCall(PCApplyMatKernel_H2OPUS(pc, X, Y, PETSC_TRUE));
    234:   PetscFunctionReturn(PETSC_SUCCESS);
    235: }

    237: static PetscErrorCode PCApply_H2OPUS(PC pc, Vec x, Vec y)
    238: {
    239:   PetscFunctionBegin;
    240:   PetscCall(PCApplyKernel_H2OPUS(pc, x, y, PETSC_FALSE));
    241:   PetscFunctionReturn(PETSC_SUCCESS);
    242: }

    244: static PetscErrorCode PCApplyTranspose_H2OPUS(PC pc, Vec x, Vec y)
    245: {
    246:   PetscFunctionBegin;
    247:   PetscCall(PCApplyKernel_H2OPUS(pc, x, y, PETSC_TRUE));
    248:   PetscFunctionReturn(PETSC_SUCCESS);
    249: }

    251: /* used to test the norm of (M^-1 A - I) */
    252: static PetscErrorCode MatMultKernel_MAmI(Mat M, Vec x, Vec y, PetscBool t)
    253: {
    254:   PC         pc;
    255:   Mat        A;
    256:   PC_H2OPUS *pch2opus;
    257:   PetscBool  sideleft = PETSC_TRUE;

    259:   PetscFunctionBegin;
    260:   PetscCall(MatShellGetContext(M, (void *)&pc));
    261:   pch2opus = (PC_H2OPUS *)pc->data;
    262:   if (!pch2opus->w) PetscCall(MatCreateVecs(pch2opus->M, &pch2opus->w, NULL));
    263:   A = pch2opus->A;
    264:   PetscCall(VecBindToCPU(pch2opus->w, pch2opus->boundtocpu));
    265:   if (t) {
    266:     if (sideleft) {
    267:       PetscCall(PCApplyTranspose_H2OPUS(pc, x, pch2opus->w));
    268:       PetscCall(MatMultTranspose(A, pch2opus->w, y));
    269:     } else {
    270:       PetscCall(MatMultTranspose(A, x, pch2opus->w));
    271:       PetscCall(PCApplyTranspose_H2OPUS(pc, pch2opus->w, y));
    272:     }
    273:   } else {
    274:     if (sideleft) {
    275:       PetscCall(MatMult(A, x, pch2opus->w));
    276:       PetscCall(PCApply_H2OPUS(pc, pch2opus->w, y));
    277:     } else {
    278:       PetscCall(PCApply_H2OPUS(pc, x, pch2opus->w));
    279:       PetscCall(MatMult(A, pch2opus->w, y));
    280:     }
    281:   }
    282:   PetscCall(VecAXPY(y, -1.0, x));
    283:   PetscFunctionReturn(PETSC_SUCCESS);
    284: }

    286: static PetscErrorCode MatMult_MAmI(Mat A, Vec x, Vec y)
    287: {
    288:   PetscFunctionBegin;
    289:   PetscCall(MatMultKernel_MAmI(A, x, y, PETSC_FALSE));
    290:   PetscFunctionReturn(PETSC_SUCCESS);
    291: }

    293: static PetscErrorCode MatMultTranspose_MAmI(Mat A, Vec x, Vec y)
    294: {
    295:   PetscFunctionBegin;
    296:   PetscCall(MatMultKernel_MAmI(A, x, y, PETSC_TRUE));
    297:   PetscFunctionReturn(PETSC_SUCCESS);
    298: }

    300: /* HyperPower kernel:
    301: Y = R = x
    302: for i = 1 . . . l - 1 do
    303:   R = (I - A * Xk) * R
    304:   Y = Y + R
    305: Y = Xk * Y
    306: */
    307: static PetscErrorCode MatMultKernel_Hyper(Mat M, Vec x, Vec y, PetscBool t)
    308: {
    309:   PC         pc;
    310:   Mat        A;
    311:   PC_H2OPUS *pch2opus;

    313:   PetscFunctionBegin;
    314:   PetscCall(MatShellGetContext(M, (void *)&pc));
    315:   pch2opus = (PC_H2OPUS *)pc->data;
    316:   A        = pch2opus->A;
    317:   PetscCall(MatCreateVecs(pch2opus->M, pch2opus->wns[0] ? NULL : &pch2opus->wns[0], pch2opus->wns[1] ? NULL : &pch2opus->wns[1]));
    318:   PetscCall(MatCreateVecs(pch2opus->M, pch2opus->wns[2] ? NULL : &pch2opus->wns[2], pch2opus->wns[3] ? NULL : &pch2opus->wns[3]));
    319:   PetscCall(VecBindToCPU(pch2opus->wns[0], pch2opus->boundtocpu));
    320:   PetscCall(VecBindToCPU(pch2opus->wns[1], pch2opus->boundtocpu));
    321:   PetscCall(VecBindToCPU(pch2opus->wns[2], pch2opus->boundtocpu));
    322:   PetscCall(VecBindToCPU(pch2opus->wns[3], pch2opus->boundtocpu));
    323:   PetscCall(VecCopy(x, pch2opus->wns[0]));
    324:   PetscCall(VecCopy(x, pch2opus->wns[3]));
    325:   if (t) {
    326:     for (PetscInt i = 0; i < pch2opus->hyperorder - 1; i++) {
    327:       PetscCall(MatMultTranspose(A, pch2opus->wns[0], pch2opus->wns[1]));
    328:       PetscCall(PCApplyTranspose_H2OPUS(pc, pch2opus->wns[1], pch2opus->wns[2]));
    329:       PetscCall(VecAXPY(pch2opus->wns[0], -1., pch2opus->wns[2]));
    330:       PetscCall(VecAXPY(pch2opus->wns[3], 1., pch2opus->wns[0]));
    331:     }
    332:     PetscCall(PCApplyTranspose_H2OPUS(pc, pch2opus->wns[3], y));
    333:   } else {
    334:     for (PetscInt i = 0; i < pch2opus->hyperorder - 1; i++) {
    335:       PetscCall(PCApply_H2OPUS(pc, pch2opus->wns[0], pch2opus->wns[1]));
    336:       PetscCall(MatMult(A, pch2opus->wns[1], pch2opus->wns[2]));
    337:       PetscCall(VecAXPY(pch2opus->wns[0], -1., pch2opus->wns[2]));
    338:       PetscCall(VecAXPY(pch2opus->wns[3], 1., pch2opus->wns[0]));
    339:     }
    340:     PetscCall(PCApply_H2OPUS(pc, pch2opus->wns[3], y));
    341:   }
    342:   PetscFunctionReturn(PETSC_SUCCESS);
    343: }

    345: static PetscErrorCode MatMult_Hyper(Mat M, Vec x, Vec y)
    346: {
    347:   PetscFunctionBegin;
    348:   PetscCall(MatMultKernel_Hyper(M, x, y, PETSC_FALSE));
    349:   PetscFunctionReturn(PETSC_SUCCESS);
    350: }

    352: static PetscErrorCode MatMultTranspose_Hyper(Mat M, Vec x, Vec y)
    353: {
    354:   PetscFunctionBegin;
    355:   PetscCall(MatMultKernel_Hyper(M, x, y, PETSC_TRUE));
    356:   PetscFunctionReturn(PETSC_SUCCESS);
    357: }

    359: /* Hyper power kernel, MatMat version */
    360: static PetscErrorCode MatMatMultKernel_Hyper(Mat M, Mat X, Mat Y, PetscBool t)
    361: {
    362:   PC         pc;
    363:   Mat        A;
    364:   PC_H2OPUS *pch2opus;
    365:   PetscInt   i;

    367:   PetscFunctionBegin;
    368:   PetscCall(MatShellGetContext(M, (void *)&pc));
    369:   pch2opus = (PC_H2OPUS *)pc->data;
    370:   A        = pch2opus->A;
    371:   if (pch2opus->wnsmat[0] && pch2opus->wnsmat[0]->cmap->N != X->cmap->N) {
    372:     PetscCall(MatDestroy(&pch2opus->wnsmat[0]));
    373:     PetscCall(MatDestroy(&pch2opus->wnsmat[1]));
    374:   }
    375:   if (!pch2opus->wnsmat[0]) {
    376:     PetscCall(MatDuplicate(X, MAT_SHARE_NONZERO_PATTERN, &pch2opus->wnsmat[0]));
    377:     PetscCall(MatDuplicate(Y, MAT_SHARE_NONZERO_PATTERN, &pch2opus->wnsmat[1]));
    378:   }
    379:   if (pch2opus->wnsmat[2] && pch2opus->wnsmat[2]->cmap->N != X->cmap->N) {
    380:     PetscCall(MatDestroy(&pch2opus->wnsmat[2]));
    381:     PetscCall(MatDestroy(&pch2opus->wnsmat[3]));
    382:   }
    383:   if (!pch2opus->wnsmat[2]) {
    384:     PetscCall(MatDuplicate(X, MAT_SHARE_NONZERO_PATTERN, &pch2opus->wnsmat[2]));
    385:     PetscCall(MatDuplicate(Y, MAT_SHARE_NONZERO_PATTERN, &pch2opus->wnsmat[3]));
    386:   }
    387:   PetscCall(MatCopy(X, pch2opus->wnsmat[0], SAME_NONZERO_PATTERN));
    388:   PetscCall(MatCopy(X, pch2opus->wnsmat[3], SAME_NONZERO_PATTERN));
    389:   if (t) {
    390:     for (i = 0; i < pch2opus->hyperorder - 1; i++) {
    391:       PetscCall(MatTransposeMatMult(A, pch2opus->wnsmat[0], MAT_REUSE_MATRIX, PETSC_CURRENT, &pch2opus->wnsmat[1]));
    392:       PetscCall(PCApplyTransposeMat_H2OPUS(pc, pch2opus->wnsmat[1], pch2opus->wnsmat[2]));
    393:       PetscCall(MatAXPY(pch2opus->wnsmat[0], -1., pch2opus->wnsmat[2], SAME_NONZERO_PATTERN));
    394:       PetscCall(MatAXPY(pch2opus->wnsmat[3], 1., pch2opus->wnsmat[0], SAME_NONZERO_PATTERN));
    395:     }
    396:     PetscCall(PCApplyTransposeMat_H2OPUS(pc, pch2opus->wnsmat[3], Y));
    397:   } else {
    398:     for (i = 0; i < pch2opus->hyperorder - 1; i++) {
    399:       PetscCall(PCApplyMat_H2OPUS(pc, pch2opus->wnsmat[0], pch2opus->wnsmat[1]));
    400:       PetscCall(MatMatMult(A, pch2opus->wnsmat[1], MAT_REUSE_MATRIX, PETSC_CURRENT, &pch2opus->wnsmat[2]));
    401:       PetscCall(MatAXPY(pch2opus->wnsmat[0], -1., pch2opus->wnsmat[2], SAME_NONZERO_PATTERN));
    402:       PetscCall(MatAXPY(pch2opus->wnsmat[3], 1., pch2opus->wnsmat[0], SAME_NONZERO_PATTERN));
    403:     }
    404:     PetscCall(PCApplyMat_H2OPUS(pc, pch2opus->wnsmat[3], Y));
    405:   }
    406:   PetscFunctionReturn(PETSC_SUCCESS);
    407: }

    409: static PetscErrorCode MatMatMultNumeric_Hyper(Mat M, Mat X, Mat Y, void *ctx)
    410: {
    411:   PetscFunctionBegin;
    412:   PetscCall(MatMatMultKernel_Hyper(M, X, Y, PETSC_FALSE));
    413:   PetscFunctionReturn(PETSC_SUCCESS);
    414: }

    416: /* Basic Newton-Schultz sampler: (2 * I - M * A)*M */
    417: static PetscErrorCode MatMultKernel_NS(Mat M, Vec x, Vec y, PetscBool t)
    418: {
    419:   PC         pc;
    420:   Mat        A;
    421:   PC_H2OPUS *pch2opus;

    423:   PetscFunctionBegin;
    424:   PetscCall(MatShellGetContext(M, (void *)&pc));
    425:   pch2opus = (PC_H2OPUS *)pc->data;
    426:   A        = pch2opus->A;
    427:   PetscCall(MatCreateVecs(pch2opus->M, pch2opus->wns[0] ? NULL : &pch2opus->wns[0], pch2opus->wns[1] ? NULL : &pch2opus->wns[1]));
    428:   PetscCall(VecBindToCPU(pch2opus->wns[0], pch2opus->boundtocpu));
    429:   PetscCall(VecBindToCPU(pch2opus->wns[1], pch2opus->boundtocpu));
    430:   if (t) {
    431:     PetscCall(PCApplyTranspose_H2OPUS(pc, x, y));
    432:     PetscCall(MatMultTranspose(A, y, pch2opus->wns[1]));
    433:     PetscCall(PCApplyTranspose_H2OPUS(pc, pch2opus->wns[1], pch2opus->wns[0]));
    434:     PetscCall(VecAXPBY(y, -1., 2., pch2opus->wns[0]));
    435:   } else {
    436:     PetscCall(PCApply_H2OPUS(pc, x, y));
    437:     PetscCall(MatMult(A, y, pch2opus->wns[0]));
    438:     PetscCall(PCApply_H2OPUS(pc, pch2opus->wns[0], pch2opus->wns[1]));
    439:     PetscCall(VecAXPBY(y, -1., 2., pch2opus->wns[1]));
    440:   }
    441:   PetscFunctionReturn(PETSC_SUCCESS);
    442: }

    444: static PetscErrorCode MatMult_NS(Mat M, Vec x, Vec y)
    445: {
    446:   PetscFunctionBegin;
    447:   PetscCall(MatMultKernel_NS(M, x, y, PETSC_FALSE));
    448:   PetscFunctionReturn(PETSC_SUCCESS);
    449: }

    451: static PetscErrorCode MatMultTranspose_NS(Mat M, Vec x, Vec y)
    452: {
    453:   PetscFunctionBegin;
    454:   PetscCall(MatMultKernel_NS(M, x, y, PETSC_TRUE));
    455:   PetscFunctionReturn(PETSC_SUCCESS);
    456: }

    458: /* Basic Newton-Schultz sampler: (2 * I - M * A)*M, MatMat version */
    459: static PetscErrorCode MatMatMultKernel_NS(Mat M, Mat X, Mat Y, PetscBool t)
    460: {
    461:   PC         pc;
    462:   Mat        A;
    463:   PC_H2OPUS *pch2opus;

    465:   PetscFunctionBegin;
    466:   PetscCall(MatShellGetContext(M, (void *)&pc));
    467:   pch2opus = (PC_H2OPUS *)pc->data;
    468:   A        = pch2opus->A;
    469:   if (pch2opus->wnsmat[0] && pch2opus->wnsmat[0]->cmap->N != X->cmap->N) {
    470:     PetscCall(MatDestroy(&pch2opus->wnsmat[0]));
    471:     PetscCall(MatDestroy(&pch2opus->wnsmat[1]));
    472:   }
    473:   if (!pch2opus->wnsmat[0]) {
    474:     PetscCall(MatDuplicate(X, MAT_SHARE_NONZERO_PATTERN, &pch2opus->wnsmat[0]));
    475:     PetscCall(MatDuplicate(Y, MAT_SHARE_NONZERO_PATTERN, &pch2opus->wnsmat[1]));
    476:   }
    477:   if (t) {
    478:     PetscCall(PCApplyTransposeMat_H2OPUS(pc, X, Y));
    479:     PetscCall(MatTransposeMatMult(A, Y, MAT_REUSE_MATRIX, PETSC_CURRENT, &pch2opus->wnsmat[1]));
    480:     PetscCall(PCApplyTransposeMat_H2OPUS(pc, pch2opus->wnsmat[1], pch2opus->wnsmat[0]));
    481:     PetscCall(MatScale(Y, 2.));
    482:     PetscCall(MatAXPY(Y, -1., pch2opus->wnsmat[0], SAME_NONZERO_PATTERN));
    483:   } else {
    484:     PetscCall(PCApplyMat_H2OPUS(pc, X, Y));
    485:     PetscCall(MatMatMult(A, Y, MAT_REUSE_MATRIX, PETSC_CURRENT, &pch2opus->wnsmat[0]));
    486:     PetscCall(PCApplyMat_H2OPUS(pc, pch2opus->wnsmat[0], pch2opus->wnsmat[1]));
    487:     PetscCall(MatScale(Y, 2.));
    488:     PetscCall(MatAXPY(Y, -1., pch2opus->wnsmat[1], SAME_NONZERO_PATTERN));
    489:   }
    490:   PetscFunctionReturn(PETSC_SUCCESS);
    491: }

    493: static PetscErrorCode MatMatMultNumeric_NS(Mat M, Mat X, Mat Y, void *ctx)
    494: {
    495:   PetscFunctionBegin;
    496:   PetscCall(MatMatMultKernel_NS(M, X, Y, PETSC_FALSE));
    497:   PetscFunctionReturn(PETSC_SUCCESS);
    498: }

    500: static PetscErrorCode PCH2OpusSetUpSampler_Private(PC pc)
    501: {
    502:   PC_H2OPUS *pch2opus = (PC_H2OPUS *)pc->data;
    503:   Mat        A        = pc->useAmat ? pc->mat : pc->pmat;

    505:   PetscFunctionBegin;
    506:   if (!pch2opus->S) {
    507:     PetscInt M, N, m, n;

    509:     PetscCall(MatGetSize(A, &M, &N));
    510:     PetscCall(MatGetLocalSize(A, &m, &n));
    511:     PetscCall(MatCreateShell(PetscObjectComm((PetscObject)A), m, n, M, N, pc, &pch2opus->S));
    512:     PetscCall(MatSetBlockSizesFromMats(pch2opus->S, A, A));
    513: #if defined(PETSC_H2OPUS_USE_GPU)
    514:     PetscCall(MatShellSetVecType(pch2opus->S, VECCUDA));
    515: #endif
    516:   }
    517:   if (pch2opus->hyperorder >= 2) {
    518:     PetscCall(MatShellSetOperation(pch2opus->S, MATOP_MULT, (void (*)(void))MatMult_Hyper));
    519:     PetscCall(MatShellSetOperation(pch2opus->S, MATOP_MULT_TRANSPOSE, (void (*)(void))MatMultTranspose_Hyper));
    520:     PetscCall(MatShellSetMatProductOperation(pch2opus->S, MATPRODUCT_AB, NULL, MatMatMultNumeric_Hyper, NULL, MATDENSE, MATDENSE));
    521:     PetscCall(MatShellSetMatProductOperation(pch2opus->S, MATPRODUCT_AB, NULL, MatMatMultNumeric_Hyper, NULL, MATDENSECUDA, MATDENSECUDA));
    522:   } else {
    523:     PetscCall(MatShellSetOperation(pch2opus->S, MATOP_MULT, (void (*)(void))MatMult_NS));
    524:     PetscCall(MatShellSetOperation(pch2opus->S, MATOP_MULT_TRANSPOSE, (void (*)(void))MatMultTranspose_NS));
    525:     PetscCall(MatShellSetMatProductOperation(pch2opus->S, MATPRODUCT_AB, NULL, MatMatMultNumeric_NS, NULL, MATDENSE, MATDENSE));
    526:     PetscCall(MatShellSetMatProductOperation(pch2opus->S, MATPRODUCT_AB, NULL, MatMatMultNumeric_NS, NULL, MATDENSECUDA, MATDENSECUDA));
    527:   }
    528:   PetscCall(MatPropagateSymmetryOptions(A, pch2opus->S));
    529:   PetscCall(MatBindToCPU(pch2opus->S, pch2opus->boundtocpu));
    530:   /* XXX */
    531:   PetscCall(MatSetOption(pch2opus->S, MAT_SYMMETRIC, PETSC_TRUE));
    532:   PetscFunctionReturn(PETSC_SUCCESS);
    533: }

    535: static PetscErrorCode PCSetUp_H2OPUS(PC pc)
    536: {
    537:   PC_H2OPUS *pch2opus = (PC_H2OPUS *)pc->data;
    538:   Mat        A        = pc->useAmat ? pc->mat : pc->pmat;
    539:   NormType   norm     = pch2opus->normtype;
    540:   PetscReal  initerr  = 0.0, err;
    541:   PetscBool  ish2opus;

    543:   PetscFunctionBegin;
    544:   if (!pch2opus->T) {
    545:     PetscInt    M, N, m, n;
    546:     const char *prefix;

    548:     PetscCall(PCGetOptionsPrefix(pc, &prefix));
    549:     PetscCall(MatGetSize(A, &M, &N));
    550:     PetscCall(MatGetLocalSize(A, &m, &n));
    551:     PetscCall(MatCreateShell(PetscObjectComm((PetscObject)pc->pmat), m, n, M, N, pc, &pch2opus->T));
    552:     PetscCall(MatSetBlockSizesFromMats(pch2opus->T, A, A));
    553:     PetscCall(MatShellSetOperation(pch2opus->T, MATOP_MULT, (void (*)(void))MatMult_MAmI));
    554:     PetscCall(MatShellSetOperation(pch2opus->T, MATOP_MULT_TRANSPOSE, (void (*)(void))MatMultTranspose_MAmI));
    555:     PetscCall(MatShellSetOperation(pch2opus->T, MATOP_NORM, (void (*)(void))MatNorm_H2OPUS));
    556: #if defined(PETSC_H2OPUS_USE_GPU)
    557:     PetscCall(MatShellSetVecType(pch2opus->T, VECCUDA));
    558: #endif
    559:     PetscCall(MatSetOptionsPrefix(pch2opus->T, prefix));
    560:     PetscCall(MatAppendOptionsPrefix(pch2opus->T, "pc_h2opus_est_"));
    561:   }
    562:   PetscCall(MatDestroy(&pch2opus->A));
    563:   PetscCall(PetscObjectTypeCompare((PetscObject)A, MATH2OPUS, &ish2opus));
    564:   if (ish2opus) {
    565:     PetscCall(PetscObjectReference((PetscObject)A));
    566:     pch2opus->A = A;
    567:   } else {
    568:     const char *prefix;

    570:     /* See if we can infer coordinates from the DM */
    571:     if (!pch2opus->sdim) PetscCall(PCH2OpusInferCoordinates_Private(pc));
    572:     PetscCall(MatCreateH2OpusFromMat(A, pch2opus->sdim, pch2opus->coords, PETSC_FALSE, pch2opus->eta, pch2opus->leafsize, pch2opus->max_rank, pch2opus->bs, pch2opus->mrtol, &pch2opus->A));
    573:     /* XXX */
    574:     PetscCall(MatSetOption(pch2opus->A, MAT_SYMMETRIC, PETSC_TRUE));
    575:     PetscCall(PCGetOptionsPrefix(pc, &prefix));
    576:     PetscCall(MatSetOptionsPrefix(pch2opus->A, prefix));
    577:     PetscCall(MatAppendOptionsPrefix(pch2opus->A, "pc_h2opus_init_"));
    578:     PetscCall(MatSetFromOptions(pch2opus->A));
    579:     PetscCall(MatAssemblyBegin(pch2opus->A, MAT_FINAL_ASSEMBLY));
    580:     PetscCall(MatAssemblyEnd(pch2opus->A, MAT_FINAL_ASSEMBLY));
    581:     /* XXX */
    582:     PetscCall(MatSetOption(pch2opus->A, MAT_SYMMETRIC, PETSC_TRUE));

    584:     /* always perform construction on the GPU unless forcecpu is true */
    585:     PetscCall(MatBindToCPU(pch2opus->A, pch2opus->forcecpu));
    586:   }
    587: #if defined(PETSC_H2OPUS_USE_GPU)
    588:   pch2opus->boundtocpu = pch2opus->forcecpu ? PETSC_TRUE : pch2opus->A->boundtocpu;
    589: #endif
    590:   PetscCall(MatBindToCPU(pch2opus->T, pch2opus->boundtocpu));
    591:   if (pch2opus->M) { /* see if we can reuse M as initial guess */
    592:     PetscReal reuse;

    594:     PetscCall(MatBindToCPU(pch2opus->M, pch2opus->boundtocpu));
    595:     PetscCall(MatNorm(pch2opus->T, norm, &reuse));
    596:     if (reuse >= 1.0) PetscCall(MatDestroy(&pch2opus->M));
    597:   }
    598:   if (!pch2opus->M) {
    599:     const char *prefix;
    600:     PetscCall(MatDuplicate(pch2opus->A, MAT_COPY_VALUES, &pch2opus->M));
    601:     PetscCall(PCGetOptionsPrefix(pc, &prefix));
    602:     PetscCall(MatSetOptionsPrefix(pch2opus->M, prefix));
    603:     PetscCall(MatAppendOptionsPrefix(pch2opus->M, "pc_h2opus_inv_"));
    604:     PetscCall(MatSetFromOptions(pch2opus->M));
    605:     PetscCall(PCH2OpusSetUpInit(pc));
    606:     PetscCall(MatScale(pch2opus->M, pch2opus->s0));
    607:   }
    608:   /* A and M have the same h2 matrix nonzero structure, save on reordering routines */
    609:   PetscCall(MatH2OpusSetNativeMult(pch2opus->A, PETSC_TRUE));
    610:   PetscCall(MatH2OpusSetNativeMult(pch2opus->M, PETSC_TRUE));
    611:   if (norm == NORM_1 || norm == NORM_2 || norm == NORM_INFINITY) PetscCall(MatNorm(pch2opus->T, norm, &initerr));
    612:   if (PetscIsInfOrNanReal(initerr)) pc->failedreason = PC_SETUP_ERROR;
    613:   err = initerr;
    614:   if (pch2opus->monitor) PetscCall(PetscPrintf(PetscObjectComm((PetscObject)pc), "%" PetscInt_FMT ": ||M*A - I|| NORM%s abs %g rel %g\n", 0, NormTypes[norm], (double)err, (double)(err / initerr)));
    615:   if (initerr > pch2opus->atol && !pc->failedreason) {
    616:     PetscInt i;

    618:     PetscCall(PCH2OpusSetUpSampler_Private(pc));
    619:     for (i = 0; i < pch2opus->maxits; i++) {
    620:       Mat         M;
    621:       const char *prefix;

    623:       PetscCall(MatDuplicate(pch2opus->M, MAT_SHARE_NONZERO_PATTERN, &M));
    624:       PetscCall(MatGetOptionsPrefix(pch2opus->M, &prefix));
    625:       PetscCall(MatSetOptionsPrefix(M, prefix));
    626:       PetscCall(MatH2OpusSetSamplingMat(M, pch2opus->S, PETSC_DECIDE, PETSC_DECIDE));
    627:       PetscCall(MatSetFromOptions(M));
    628:       PetscCall(MatH2OpusSetNativeMult(M, PETSC_TRUE));
    629:       PetscCall(MatAssemblyBegin(M, MAT_FINAL_ASSEMBLY));
    630:       PetscCall(MatAssemblyEnd(M, MAT_FINAL_ASSEMBLY));
    631:       PetscCall(MatDestroy(&pch2opus->M));
    632:       pch2opus->M = M;
    633:       if (norm == NORM_1 || norm == NORM_2 || norm == NORM_INFINITY) PetscCall(MatNorm(pch2opus->T, norm, &err));
    634:       if (pch2opus->monitor) PetscCall(PetscPrintf(PetscObjectComm((PetscObject)pc), "%" PetscInt_FMT ": ||M*A - I|| NORM%s abs %g rel %g\n", i + 1, NormTypes[norm], (double)err, (double)(err / initerr)));
    635:       if (PetscIsInfOrNanReal(err)) pc->failedreason = PC_SETUP_ERROR;
    636:       if (err < pch2opus->atol || err < pch2opus->rtol * initerr || pc->failedreason) break;
    637:     }
    638:   }
    639:   /* cleanup setup workspace */
    640:   PetscCall(MatH2OpusSetNativeMult(pch2opus->A, PETSC_FALSE));
    641:   PetscCall(MatH2OpusSetNativeMult(pch2opus->M, PETSC_FALSE));
    642:   PetscCall(VecDestroy(&pch2opus->wns[0]));
    643:   PetscCall(VecDestroy(&pch2opus->wns[1]));
    644:   PetscCall(VecDestroy(&pch2opus->wns[2]));
    645:   PetscCall(VecDestroy(&pch2opus->wns[3]));
    646:   PetscCall(MatDestroy(&pch2opus->wnsmat[0]));
    647:   PetscCall(MatDestroy(&pch2opus->wnsmat[1]));
    648:   PetscCall(MatDestroy(&pch2opus->wnsmat[2]));
    649:   PetscCall(MatDestroy(&pch2opus->wnsmat[3]));
    650:   PetscFunctionReturn(PETSC_SUCCESS);
    651: }

    653: static PetscErrorCode PCView_H2OPUS(PC pc, PetscViewer viewer)
    654: {
    655:   PC_H2OPUS *pch2opus = (PC_H2OPUS *)pc->data;
    656:   PetscBool  isascii;

    658:   PetscFunctionBegin;
    659:   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &isascii));
    660:   if (isascii) {
    661:     if (pch2opus->A && pch2opus->A != pc->mat && pch2opus->A != pc->pmat) {
    662:       PetscCall(PetscViewerASCIIPrintf(viewer, "Initial approximation matrix\n"));
    663:       PetscCall(PetscViewerASCIIPushTab(viewer));
    664:       PetscCall(PetscViewerPushFormat(viewer, PETSC_VIEWER_ASCII_INFO_DETAIL));
    665:       PetscCall(MatView(pch2opus->A, viewer));
    666:       PetscCall(PetscViewerPopFormat(viewer));
    667:       PetscCall(PetscViewerASCIIPopTab(viewer));
    668:     }
    669:     if (pch2opus->M) {
    670:       PetscCall(PetscViewerASCIIPrintf(viewer, "Inner matrix constructed\n"));
    671:       PetscCall(PetscViewerASCIIPushTab(viewer));
    672:       PetscCall(PetscViewerPushFormat(viewer, PETSC_VIEWER_ASCII_INFO_DETAIL));
    673:       PetscCall(MatView(pch2opus->M, viewer));
    674:       PetscCall(PetscViewerPopFormat(viewer));
    675:       PetscCall(PetscViewerASCIIPopTab(viewer));
    676:     }
    677:     PetscCall(PetscViewerASCIIPrintf(viewer, "Initial scaling: %g\n", pch2opus->s0));
    678:   }
    679:   PetscFunctionReturn(PETSC_SUCCESS);
    680: }

    682: /*MC
    683:      PCH2OPUS = "h2opus" - A preconditioner type for, `MATH2OPUS`, hierarchical matrices using the H2Opus package.

    685:    Options Database Keys:
    686: +   -pc_type h2opus - pc type to "h2opus" during a call to `PCSetFromOptions()`
    687: .   -pc_h2opus_maxits - maximum number of iterations for Newton-Schultz
    688: .   -pc_h2opus_monitor - monitor Newton-Schultz convergence
    689: .   -pc_h2opus_atol - absolute tolerance
    690: .   -pc_h2opus_rtol - relative tolerance
    691: .   -pc_h2opus_norm_type - normtype
    692: .   -pc_h2opus_hyperorder - Hyper power order of sampling
    693: .   -pc_h2opus_leafsize - leaf size when constructed from kernel
    694: .   -pc_h2opus_eta - admissibility condition tolerance
    695: .   -pc_h2opus_maxrank - maximum rank when constructed from matvecs
    696: .   -pc_h2opus_samples - number of samples to be taken concurrently when constructing from matvecs
    697: .   -pc_h2opus_mrtol - relative tolerance for construction from sampling
    698: -   -pc_h2opus_forcecpu - force construction of preconditioner on CPU

    700:    Level: intermediate

    702: .seealso: [](ch_ksp), `MATH2OPUS`, `MATHTOOL`, `MATDENSE`, `MatCreateH2OpusFromKernel()`, `MatCreateH2OpusFromMat()`
    703: M*/

    705: PETSC_EXTERN PetscErrorCode PCCreate_H2OPUS(PC pc)
    706: {
    707:   PC_H2OPUS *pch2opus;

    709:   PetscFunctionBegin;
    710:   PetscCall(PetscNew(&pch2opus));
    711:   pc->data = (void *)pch2opus;

    713:   pch2opus->atol       = 1.e-2;
    714:   pch2opus->rtol       = 1.e-6;
    715:   pch2opus->maxits     = 50;
    716:   pch2opus->hyperorder = 1; /* defaults to basic Newton-Schultz */
    717:   pch2opus->normtype   = NORM_2;

    719:   /* these are needed when we are sampling the pmat */
    720:   pch2opus->eta      = PETSC_DECIDE;
    721:   pch2opus->leafsize = PETSC_DECIDE;
    722:   pch2opus->max_rank = PETSC_DECIDE;
    723:   pch2opus->bs       = PETSC_DECIDE;
    724:   pch2opus->mrtol    = PETSC_DECIDE;
    725: #if defined(PETSC_H2OPUS_USE_GPU)
    726:   pch2opus->boundtocpu = PETSC_FALSE;
    727: #else
    728:   pch2opus->boundtocpu = PETSC_TRUE;
    729: #endif
    730:   pc->ops->destroy        = PCDestroy_H2OPUS;
    731:   pc->ops->setup          = PCSetUp_H2OPUS;
    732:   pc->ops->apply          = PCApply_H2OPUS;
    733:   pc->ops->matapply       = PCApplyMat_H2OPUS;
    734:   pc->ops->applytranspose = PCApplyTranspose_H2OPUS;
    735:   pc->ops->reset          = PCReset_H2OPUS;
    736:   pc->ops->setfromoptions = PCSetFromOptions_H2OPUS;
    737:   pc->ops->view           = PCView_H2OPUS;

    739:   PetscCall(PetscObjectComposeFunction((PetscObject)pc, "PCSetCoordinates_C", PCSetCoordinates_H2OPUS));
    740:   PetscFunctionReturn(PETSC_SUCCESS);
    741: }