Actual source code: bjkokkos.kokkos.cxx

      1: #define PETSC_SKIP_CXX_COMPLEX_FIX // Kokkos::complex does not need the PetscComplex fix

      3: #include <petsc/private/pcbjkokkosimpl.h>

      5: #include <petsc/private/kspimpl.h>
      6: #include <petscksp.h>
      7: #include <../src/mat/impls/aij/mpi/mpiaij.h>
      8: #include <../src/mat/impls/aij/seq/kokkos/aijkok.hpp>
      9: #include <petscsection.h>
     10: #include <petscdmcomposite.h>

     12: #include <../src/mat/impls/aij/seq/aij.h>
     13: #include <../src/mat/impls/aij/seq/kokkos/aijkok.hpp>

     15: #include <petscdevice_cupm.h>

     17: static PetscErrorCode PCBJKOKKOSCreateKSP_BJKOKKOS(PC pc)
     18: {
     19:   const char    *prefix;
     20:   PC_PCBJKOKKOS *jac = (PC_PCBJKOKKOS *)pc->data;
     21:   DM             dm;

     23:   PetscFunctionBegin;
     24:   PetscCall(KSPCreate(PetscObjectComm((PetscObject)pc), &jac->ksp));
     25:   PetscCall(KSPSetNestLevel(jac->ksp, pc->kspnestlevel));
     26:   PetscCall(KSPSetErrorIfNotConverged(jac->ksp, pc->erroriffailure));
     27:   PetscCall(PetscObjectIncrementTabLevel((PetscObject)jac->ksp, (PetscObject)pc, 1));
     28:   PetscCall(PCGetOptionsPrefix(pc, &prefix));
     29:   PetscCall(KSPSetOptionsPrefix(jac->ksp, prefix));
     30:   PetscCall(KSPAppendOptionsPrefix(jac->ksp, "pc_bjkokkos_"));
     31:   PetscCall(PCGetDM(pc, &dm));
     32:   if (dm) {
     33:     PetscCall(KSPSetDM(jac->ksp, dm));
     34:     PetscCall(KSPSetDMActive(jac->ksp, KSP_DMACTIVE_ALL, PETSC_FALSE));
     35:   }
     36:   jac->reason       = PETSC_FALSE;
     37:   jac->monitor      = PETSC_FALSE;
     38:   jac->batch_target = 0;
     39:   jac->rank_target  = 0;
     40:   jac->nsolves_team = 1;
     41:   jac->ksp->max_it  = 50; // this is really for GMRES w/o restarts
     42:   PetscFunctionReturn(PETSC_SUCCESS);
     43: }

     45: // y <-- Ax
     46: KOKKOS_INLINE_FUNCTION PetscErrorCode MatMult(const team_member team, const PetscInt *glb_Aai, const PetscInt *glb_Aaj, const PetscScalar *glb_Aaa, const PetscInt *r, const PetscInt *ic, const PetscInt start, const PetscInt end, const PetscScalar *x_loc, PetscScalar *y_loc)
     47: {
     48:   Kokkos::parallel_for(Kokkos::TeamThreadRange(team, start, end), [=](const int rowb) {
     49:     int                rowa = ic[rowb];
     50:     int                n    = glb_Aai[rowa + 1] - glb_Aai[rowa];
     51:     const PetscInt    *aj   = glb_Aaj + glb_Aai[rowa]; // global
     52:     const PetscScalar *aa   = glb_Aaa + glb_Aai[rowa];
     53:     PetscScalar        sum;
     54:     Kokkos::parallel_reduce(Kokkos::ThreadVectorRange(team, n), [=](const int i, PetscScalar &lsum) { lsum += aa[i] * x_loc[r[aj[i]] - start]; }, sum);
     55:     Kokkos::single(Kokkos::PerThread(team), [=]() { y_loc[rowb - start] = sum; });
     56:   });
     57:   team.team_barrier();
     58:   return PETSC_SUCCESS;
     59: }

     61: // temp buffer per thread with reduction at end?
     62: KOKKOS_INLINE_FUNCTION PetscErrorCode MatMultTranspose(const team_member team, const PetscInt *glb_Aai, const PetscInt *glb_Aaj, const PetscScalar *glb_Aaa, const PetscInt *r, const PetscInt *ic, const PetscInt start, const PetscInt end, const PetscScalar *x_loc, PetscScalar *y_loc)
     63: {
     64:   Kokkos::parallel_for(Kokkos::TeamVectorRange(team, end - start), [=](int i) { y_loc[i] = 0; });
     65:   team.team_barrier();
     66:   Kokkos::parallel_for(Kokkos::TeamThreadRange(team, start, end), [=](const int rowb) {
     67:     int                rowa = ic[rowb];
     68:     int                n    = glb_Aai[rowa + 1] - glb_Aai[rowa];
     69:     const PetscInt    *aj   = glb_Aaj + glb_Aai[rowa]; // global
     70:     const PetscScalar *aa   = glb_Aaa + glb_Aai[rowa];
     71:     const PetscScalar  xx   = x_loc[rowb - start]; // rowb = ic[rowa] = ic[r[rowb]]
     72:     Kokkos::parallel_for(Kokkos::ThreadVectorRange(team, n), [=](const int &i) {
     73:       PetscScalar val = aa[i] * xx;
     74:       Kokkos::atomic_fetch_add(&y_loc[r[aj[i]] - start], val);
     75:     });
     76:   });
     77:   team.team_barrier();
     78:   return PETSC_SUCCESS;
     79: }

     81: typedef struct Batch_MetaData_TAG {
     82:   PetscInt           flops;
     83:   PetscInt           its;
     84:   KSPConvergedReason reason;
     85: } Batch_MetaData;

     87: // Solve A(BB^-1)x = y with TFQMR. Right preconditioned to get un-preconditioned residual
     88: static KOKKOS_INLINE_FUNCTION PetscErrorCode BJSolve_TFQMR(const team_member team, const PetscInt *glb_Aai, const PetscInt *glb_Aaj, const PetscScalar *glb_Aaa, const PetscInt *r, const PetscInt *ic, PetscScalar *work_space_global, const int stride_global, const int nShareVec, PetscScalar *work_space_shared, const int stride_shared, PetscReal rtol, PetscReal atol, PetscReal dtol, PetscInt maxit, Batch_MetaData *metad, const PetscInt start, const PetscInt end, const PetscScalar glb_idiag[], const PetscScalar *glb_b, PetscScalar *glb_x, bool monitor)
     89: {
     90:   using Kokkos::parallel_for;
     91:   using Kokkos::parallel_reduce;
     92:   int                Nblk = end - start, it, m, stride = stride_shared, idx = 0;
     93:   PetscReal          dp, dpold, w, dpest, tau, psi, cm, r0;
     94:   const PetscScalar *Diag = &glb_idiag[start];
     95:   PetscScalar       *ptr  = work_space_shared, rho, rhoold, a, s, b, eta, etaold, psiold, cf, dpi;

     97:   if (idx++ == nShareVec) {
     98:     ptr    = work_space_global;
     99:     stride = stride_global;
    100:   }
    101:   PetscScalar *XX = ptr;
    102:   ptr += stride;
    103:   if (idx++ == nShareVec) {
    104:     ptr    = work_space_global;
    105:     stride = stride_global;
    106:   }
    107:   PetscScalar *R = ptr;
    108:   ptr += stride;
    109:   if (idx++ == nShareVec) {
    110:     ptr    = work_space_global;
    111:     stride = stride_global;
    112:   }
    113:   PetscScalar *RP = ptr;
    114:   ptr += stride;
    115:   if (idx++ == nShareVec) {
    116:     ptr    = work_space_global;
    117:     stride = stride_global;
    118:   }
    119:   PetscScalar *V = ptr;
    120:   ptr += stride;
    121:   if (idx++ == nShareVec) {
    122:     ptr    = work_space_global;
    123:     stride = stride_global;
    124:   }
    125:   PetscScalar *T = ptr;
    126:   ptr += stride;
    127:   if (idx++ == nShareVec) {
    128:     ptr    = work_space_global;
    129:     stride = stride_global;
    130:   }
    131:   PetscScalar *Q = ptr;
    132:   ptr += stride;
    133:   if (idx++ == nShareVec) {
    134:     ptr    = work_space_global;
    135:     stride = stride_global;
    136:   }
    137:   PetscScalar *P = ptr;
    138:   ptr += stride;
    139:   if (idx++ == nShareVec) {
    140:     ptr    = work_space_global;
    141:     stride = stride_global;
    142:   }
    143:   PetscScalar *U = ptr;
    144:   ptr += stride;
    145:   if (idx++ == nShareVec) {
    146:     ptr    = work_space_global;
    147:     stride = stride_global;
    148:   }
    149:   PetscScalar *D = ptr;
    150:   ptr += stride;
    151:   if (idx++ == nShareVec) {
    152:     ptr    = work_space_global;
    153:     stride = stride_global;
    154:   }
    155:   PetscScalar *AUQ = V;

    157:   // init: get b, zero x
    158:   parallel_for(Kokkos::TeamVectorRange(team, start, end), [=](int rowb) {
    159:     int rowa         = ic[rowb];
    160:     R[rowb - start]  = glb_b[rowa];
    161:     XX[rowb - start] = 0;
    162:   });
    163:   team.team_barrier();
    164:   parallel_reduce(Kokkos::TeamVectorRange(team, Nblk), [=](const int idx, PetscScalar &lsum) { lsum += R[idx] * PetscConj(R[idx]); }, dpi);
    165:   team.team_barrier();
    166:   r0 = dp = PetscSqrtReal(PetscRealPart(dpi));
    167:   // diagnostics
    168: #if defined(PETSC_USE_DEBUG) && !defined(PETSC_HAVE_SYCL)
    169:   if (monitor) Kokkos::single(Kokkos::PerTeam(team), [=]() { printf("%3d KSP Residual norm %14.12e\n", 0, (double)dp); });
    170: #endif
    171:   if (dp < atol) {
    172:     metad->reason = KSP_CONVERGED_ATOL_NORMAL_EQUATIONS;
    173:     it            = 0;
    174:     goto done;
    175:   }
    176:   if (0 == maxit) {
    177:     metad->reason = KSP_CONVERGED_ITS;
    178:     it            = 0;
    179:     goto done;
    180:   }

    182:   /* Make the initial Rp = R */
    183:   parallel_for(Kokkos::TeamVectorRange(team, Nblk), [=](int idx) { RP[idx] = R[idx]; });
    184:   team.team_barrier();
    185:   /* Set the initial conditions */
    186:   etaold = 0.0;
    187:   psiold = 0.0;
    188:   tau    = dp;
    189:   dpold  = dp;

    191:   /* rhoold = (r,rp)     */
    192:   parallel_reduce(Kokkos::TeamVectorRange(team, Nblk), [=](const int idx, PetscScalar &dot) { dot += R[idx] * PetscConj(RP[idx]); }, rhoold);
    193:   team.team_barrier();
    194:   parallel_for(Kokkos::TeamVectorRange(team, Nblk), [=](int idx) {
    195:     U[idx] = R[idx];
    196:     P[idx] = R[idx];
    197:     T[idx] = Diag[idx] * P[idx];
    198:     D[idx] = 0;
    199:   });
    200:   team.team_barrier();
    201:   static_cast<void>(MatMult(team, glb_Aai, glb_Aaj, glb_Aaa, r, ic, start, end, T, V));

    203:   it = 0;
    204:   do {
    205:     /* s <- (v,rp)          */
    206:     parallel_reduce(Kokkos::TeamVectorRange(team, Nblk), [=](const int idx, PetscScalar &dot) { dot += V[idx] * PetscConj(RP[idx]); }, s);
    207:     team.team_barrier();
    208:     if (s == 0) {
    209:       metad->reason = KSP_CONVERGED_HAPPY_BREAKDOWN;
    210:       goto done;
    211:     }
    212:     a = rhoold / s; /* a <- rho / s         */
    213:     /* q <- u - a v    VecWAXPY(w,alpha,x,y): w = alpha x + y.     */
    214:     /* t <- u + q           */
    215:     parallel_for(Kokkos::TeamVectorRange(team, Nblk), [=](int idx) {
    216:       Q[idx] = U[idx] - a * V[idx];
    217:       T[idx] = U[idx] + Q[idx];
    218:     });
    219:     team.team_barrier();
    220:     // KSP_PCApplyBAorAB
    221:     parallel_for(Kokkos::TeamVectorRange(team, Nblk), [=](int idx) { T[idx] = Diag[idx] * T[idx]; });
    222:     team.team_barrier();
    223:     static_cast<void>(MatMult(team, glb_Aai, glb_Aaj, glb_Aaa, r, ic, start, end, T, AUQ));
    224:     /* r <- r - a K (u + q) */
    225:     parallel_for(Kokkos::TeamVectorRange(team, Nblk), [=](int idx) { R[idx] = R[idx] - a * AUQ[idx]; });
    226:     team.team_barrier();
    227:     parallel_reduce(Kokkos::TeamVectorRange(team, Nblk), [=](const int idx, PetscScalar &lsum) { lsum += R[idx] * PetscConj(R[idx]); }, dpi);
    228:     team.team_barrier();
    229:     dp = PetscSqrtReal(PetscRealPart(dpi));
    230:     for (m = 0; m < 2; m++) {
    231:       if (!m) w = PetscSqrtReal(dp * dpold);
    232:       else w = dp;
    233:       psi = w / tau;
    234:       cm  = 1.0 / PetscSqrtReal(1.0 + psi * psi);
    235:       tau = tau * psi * cm;
    236:       eta = cm * cm * a;
    237:       cf  = psiold * psiold * etaold / a;
    238:       if (!m) {
    239:         /* D = U + cf D */
    240:         parallel_for(Kokkos::TeamVectorRange(team, Nblk), [=](int idx) { D[idx] = U[idx] + cf * D[idx]; });
    241:       } else {
    242:         /* D = Q + cf D */
    243:         parallel_for(Kokkos::TeamVectorRange(team, Nblk), [=](int idx) { D[idx] = Q[idx] + cf * D[idx]; });
    244:       }
    245:       team.team_barrier();
    246:       parallel_for(Kokkos::TeamVectorRange(team, Nblk), [=](int idx) { XX[idx] = XX[idx] + eta * D[idx]; });
    247:       team.team_barrier();
    248:       dpest = PetscSqrtReal(2 * it + m + 2.0) * tau;
    249: #if defined(PETSC_USE_DEBUG) && !defined(PETSC_HAVE_SYCL)
    250:       if (monitor && m == 1) Kokkos::single(Kokkos::PerTeam(team), [=]() { printf("%3d KSP Residual norm %14.12e\n", it + 1, (double)dpest); });
    251: #endif
    252:       if (dpest < atol) {
    253:         metad->reason = KSP_CONVERGED_ATOL_NORMAL_EQUATIONS;
    254:         goto done;
    255:       }
    256:       if (dpest / r0 < rtol) {
    257:         metad->reason = KSP_CONVERGED_RTOL_NORMAL_EQUATIONS;
    258:         goto done;
    259:       }
    260: #if defined(PETSC_USE_DEBUG) && !defined(PETSC_HAVE_SYCL)
    261:       if (dpest / r0 > dtol) {
    262:         metad->reason = KSP_DIVERGED_DTOL;
    263:         Kokkos::single(Kokkos::PerTeam(team), [=]() { printf("ERROR block %d diverged: %d it, res=%e, r_0=%e\n", team.league_rank(), it, dpest, r0); });
    264:         goto done;
    265:       }
    266: #else
    267:       if (dpest / r0 > dtol) {
    268:         metad->reason = KSP_DIVERGED_DTOL;
    269:         goto done;
    270:       }
    271: #endif
    272:       if (it + 1 == maxit) {
    273:         metad->reason = KSP_CONVERGED_ITS;
    274: #if defined(PETSC_USE_DEBUG) && !defined(PETSC_HAVE_SYCL)
    275:         Kokkos::single(Kokkos::PerTeam(team), [=]() { printf("ERROR block %d diverged: TFQMR %d:%d it, res=%e, r_0=%e r_res=%e\n", team.league_rank(), it, m, dpest, r0, dpest / r0); });
    276: #endif
    277:         goto done;
    278:       }
    279:       etaold = eta;
    280:       psiold = psi;
    281:     }

    283:     /* rho <- (r,rp)       */
    284:     parallel_reduce(Kokkos::TeamVectorRange(team, Nblk), [=](const int idx, PetscScalar &dot) { dot += R[idx] * PetscConj(RP[idx]); }, rho);
    285:     team.team_barrier();
    286:     if (rho == 0) {
    287:       metad->reason = KSP_CONVERGED_HAPPY_BREAKDOWN;
    288:       goto done;
    289:     }
    290:     b = rho / rhoold; /* b <- rho / rhoold   */
    291:     /* u <- r + b q        */
    292:     /* p <- u + b(q + b p) */
    293:     parallel_for(Kokkos::TeamVectorRange(team, Nblk), [=](int idx) {
    294:       U[idx] = R[idx] + b * Q[idx];
    295:       Q[idx] = Q[idx] + b * P[idx];
    296:       P[idx] = U[idx] + b * Q[idx];
    297:     });
    298:     /* v <- K p  */
    299:     team.team_barrier();
    300:     parallel_for(Kokkos::TeamVectorRange(team, Nblk), [=](int idx) { T[idx] = Diag[idx] * P[idx]; });
    301:     team.team_barrier();
    302:     static_cast<void>(MatMult(team, glb_Aai, glb_Aaj, glb_Aaa, r, ic, start, end, T, V));

    304:     rhoold = rho;
    305:     dpold  = dp;

    307:     it++;
    308:   } while (it < maxit);
    309: done:
    310:   // KSPUnwindPreconditioner
    311:   parallel_for(Kokkos::TeamVectorRange(team, Nblk), [=](int idx) { XX[idx] = Diag[idx] * XX[idx]; });
    312:   team.team_barrier();
    313:   // put x into Plex order
    314:   parallel_for(Kokkos::TeamVectorRange(team, start, end), [=](int rowb) {
    315:     int rowa    = ic[rowb];
    316:     glb_x[rowa] = XX[rowb - start];
    317:   });
    318:   metad->its = it;
    319:   if (1) {
    320:     int nnz;
    321:     parallel_reduce(Kokkos::TeamVectorRange(team, start, end), [=](const int idx, int &lsum) { lsum += (glb_Aai[idx + 1] - glb_Aai[idx]); }, nnz);
    322:     metad->flops = 2 * (metad->its * (10 * Nblk + 2 * nnz) + 5 * Nblk);
    323:   } else {
    324:     metad->flops = 2 * (metad->its * (10 * Nblk + 2 * 50 * Nblk) + 5 * Nblk); // guess
    325:   }
    326:   return PETSC_SUCCESS;
    327: }

    329: // Solve Ax = y with biCG
    330: static KOKKOS_INLINE_FUNCTION PetscErrorCode BJSolve_BICG(const team_member team, const PetscInt *glb_Aai, const PetscInt *glb_Aaj, const PetscScalar *glb_Aaa, const PetscInt *r, const PetscInt *ic, PetscScalar *work_space_global, const int stride_global, const int nShareVec, PetscScalar *work_space_shared, const int stride_shared, PetscReal rtol, PetscReal atol, PetscReal dtol, PetscInt maxit, Batch_MetaData *metad, const PetscInt start, const PetscInt end, const PetscScalar glb_idiag[], const PetscScalar *glb_b, PetscScalar *glb_x, bool monitor)
    331: {
    332:   using Kokkos::parallel_for;
    333:   using Kokkos::parallel_reduce;
    334:   int                Nblk = end - start, it, stride = stride_shared, idx = 0; // start in shared mem
    335:   PetscReal          dp, r0;
    336:   const PetscScalar *Di  = &glb_idiag[start];
    337:   PetscScalar       *ptr = work_space_shared, dpi, a = 1.0, beta, betaold = 1.0, t1, t2;

    339:   if (idx++ == nShareVec) {
    340:     ptr    = work_space_global;
    341:     stride = stride_global;
    342:   }
    343:   PetscScalar *XX = ptr;
    344:   ptr += stride;
    345:   if (idx++ == nShareVec) {
    346:     ptr    = work_space_global;
    347:     stride = stride_global;
    348:   }
    349:   PetscScalar *Rl = ptr;
    350:   ptr += stride;
    351:   if (idx++ == nShareVec) {
    352:     ptr    = work_space_global;
    353:     stride = stride_global;
    354:   }
    355:   PetscScalar *Zl = ptr;
    356:   ptr += stride;
    357:   if (idx++ == nShareVec) {
    358:     ptr    = work_space_global;
    359:     stride = stride_global;
    360:   }
    361:   PetscScalar *Pl = ptr;
    362:   ptr += stride;
    363:   if (idx++ == nShareVec) {
    364:     ptr    = work_space_global;
    365:     stride = stride_global;
    366:   }
    367:   PetscScalar *Rr = ptr;
    368:   ptr += stride;
    369:   if (idx++ == nShareVec) {
    370:     ptr    = work_space_global;
    371:     stride = stride_global;
    372:   }
    373:   PetscScalar *Zr = ptr;
    374:   ptr += stride;
    375:   if (idx++ == nShareVec) {
    376:     ptr    = work_space_global;
    377:     stride = stride_global;
    378:   }
    379:   PetscScalar *Pr = ptr;
    380:   ptr += stride;

    382:   /*     r <- b (x is 0) */
    383:   parallel_for(Kokkos::TeamVectorRange(team, start, end), [=](int rowb) {
    384:     int rowa         = ic[rowb];
    385:     Rl[rowb - start] = Rr[rowb - start] = glb_b[rowa];
    386:     XX[rowb - start]                    = 0;
    387:   });
    388:   team.team_barrier();
    389:   /*     z <- Br         */
    390:   parallel_for(Kokkos::TeamVectorRange(team, Nblk), [=](int idx) {
    391:     Zr[idx] = Di[idx] * Rr[idx];
    392:     Zl[idx] = Di[idx] * Rl[idx];
    393:   });
    394:   team.team_barrier();
    395:   /*    dp <- r'*r       */
    396:   parallel_reduce(Kokkos::TeamVectorRange(team, Nblk), [=](const int idx, PetscScalar &lsum) { lsum += Rr[idx] * PetscConj(Rr[idx]); }, dpi);
    397:   team.team_barrier();
    398:   r0 = dp = PetscSqrtReal(PetscRealPart(dpi));
    399: #if defined(PETSC_USE_DEBUG) && !defined(PETSC_HAVE_SYCL)
    400:   if (monitor) Kokkos::single(Kokkos::PerTeam(team), [=]() { printf("%3d KSP Residual norm %14.12e\n", 0, (double)dp); });
    401: #endif
    402:   if (dp < atol) {
    403:     metad->reason = KSP_CONVERGED_ATOL_NORMAL_EQUATIONS;
    404:     it            = 0;
    405:     goto done;
    406:   }
    407:   if (0 == maxit) {
    408:     metad->reason = KSP_CONVERGED_ITS;
    409:     it            = 0;
    410:     goto done;
    411:   }

    413:   it = 0;
    414:   do {
    415:     /*     beta <- r'z     */
    416:     parallel_reduce(Kokkos::TeamVectorRange(team, Nblk), [=](const int idx, PetscScalar &dot) { dot += Zr[idx] * PetscConj(Rl[idx]); }, beta);
    417:     team.team_barrier();
    418: #if PCBJKOKKOS_VERBOSE_LEVEL >= 6
    419:   #if defined(PETSC_USE_DEBUG) && !defined(PETSC_HAVE_SYCL)
    420:     Kokkos::single(Kokkos::PerTeam(team), [=]() { printf("%7d beta = Z.R = %22.14e \n", i, (double)beta); });
    421:   #endif
    422: #endif
    423:     if (beta == 0.0) {
    424:       metad->reason = KSP_CONVERGED_HAPPY_BREAKDOWN;
    425:       goto done;
    426:     }
    427:     if (it == 0) {
    428:       /*     p <- z          */
    429:       parallel_for(Kokkos::TeamVectorRange(team, Nblk), [=](int idx) {
    430:         Pr[idx] = Zr[idx];
    431:         Pl[idx] = Zl[idx];
    432:       });
    433:     } else {
    434:       t1 = beta / betaold;
    435:       /*     p <- z + b* p   */
    436:       t2 = PetscConj(t1);
    437:       parallel_for(Kokkos::TeamVectorRange(team, Nblk), [=](int idx) {
    438:         Pr[idx] = t1 * Pr[idx] + Zr[idx];
    439:         Pl[idx] = t2 * Pl[idx] + Zl[idx];
    440:       });
    441:     }
    442:     team.team_barrier();
    443:     betaold = beta;
    444:     /*     z <- Kp         */
    445:     static_cast<void>(MatMult(team, glb_Aai, glb_Aaj, glb_Aaa, r, ic, start, end, Pr, Zr));
    446:     static_cast<void>(MatMultTranspose(team, glb_Aai, glb_Aaj, glb_Aaa, r, ic, start, end, Pl, Zl));
    447:     /*     dpi <- z'p      */
    448:     parallel_reduce(Kokkos::TeamVectorRange(team, Nblk), [=](const int idx, PetscScalar &lsum) { lsum += Zr[idx] * PetscConj(Pl[idx]); }, dpi);
    449:     team.team_barrier();
    450:     if (dpi == 0) {
    451:       metad->reason = KSP_CONVERGED_HAPPY_BREAKDOWN;
    452:       goto done;
    453:     }
    454:     //
    455:     a  = beta / dpi; /*     a = beta/p'z    */
    456:     t1 = -a;
    457:     t2 = PetscConj(t1);
    458:     /*     x <- x + ap     */
    459:     parallel_for(Kokkos::TeamVectorRange(team, Nblk), [=](int idx) {
    460:       XX[idx] = XX[idx] + a * Pr[idx];
    461:       Rr[idx] = Rr[idx] + t1 * Zr[idx];
    462:       Rl[idx] = Rl[idx] + t2 * Zl[idx];
    463:     });
    464:     team.team_barrier();
    465:     team.team_barrier();
    466:     /*    dp <- r'*r       */
    467:     parallel_reduce(Kokkos::TeamVectorRange(team, Nblk), [=](const int idx, PetscScalar &lsum) { lsum += Rr[idx] * PetscConj(Rr[idx]); }, dpi);
    468:     team.team_barrier();
    469:     dp = PetscSqrtReal(PetscRealPart(dpi));
    470: #if defined(PETSC_USE_DEBUG) && !defined(PETSC_HAVE_SYCL)
    471:     if (monitor) Kokkos::single(Kokkos::PerTeam(team), [=]() { printf("%3d KSP Residual norm %14.12e\n", it + 1, (double)dp); });
    472: #endif
    473:     if (dp < atol) {
    474:       metad->reason = KSP_CONVERGED_ATOL_NORMAL_EQUATIONS;
    475:       goto done;
    476:     }
    477:     if (dp / r0 < rtol) {
    478:       metad->reason = KSP_CONVERGED_RTOL_NORMAL_EQUATIONS;
    479:       goto done;
    480:     }
    481: #if defined(PETSC_USE_DEBUG) && !defined(PETSC_HAVE_SYCL)
    482:     if (dp / r0 > dtol) {
    483:       metad->reason = KSP_DIVERGED_DTOL;
    484:       Kokkos::single(Kokkos::PerTeam(team), [=]() { printf("ERROR block %d diverged: %d it, res=%e, r_0=%e (BICG does this)\n", team.league_rank(), it, dp, r0); });
    485:       goto done;
    486:     }
    487: #else
    488:     if (dp / r0 > dtol) {
    489:       metad->reason = KSP_DIVERGED_DTOL;
    490:       goto done;
    491:     }
    492: #endif
    493:     if (it + 1 == maxit) {
    494:       metad->reason = KSP_CONVERGED_ITS; // don't worry about hitting max iterations
    495: #if defined(PETSC_USE_DEBUG) && !defined(PETSC_HAVE_SYCL)
    496:       Kokkos::single(Kokkos::PerTeam(team), [=]() { printf("ERROR block %d diverged: BICG %d it, res=%e, r_0=%e r_res=%e\n", team.league_rank(), it, dp, r0, dp / r0); });
    497: #endif
    498:       goto done;
    499:     }
    500:     /* z <- Br  */
    501:     parallel_for(Kokkos::TeamVectorRange(team, Nblk), [=](int idx) {
    502:       Zr[idx] = Di[idx] * Rr[idx];
    503:       Zl[idx] = Di[idx] * Rl[idx];
    504:     });

    506:     it++;
    507:   } while (it < maxit);
    508: done:
    509:   // put x back into Plex order
    510:   parallel_for(Kokkos::TeamVectorRange(team, start, end), [=](int rowb) {
    511:     int rowa    = ic[rowb];
    512:     glb_x[rowa] = XX[rowb - start];
    513:   });
    514:   metad->its = it;
    515:   if (1) {
    516:     int nnz;
    517:     parallel_reduce(Kokkos::TeamVectorRange(team, start, end), [=](const int idx, int &lsum) { lsum += (glb_Aai[idx + 1] - glb_Aai[idx]); }, nnz);
    518:     metad->flops = 2 * (metad->its * (10 * Nblk + 2 * nnz) + 5 * Nblk);
    519:   } else {
    520:     metad->flops = 2 * (metad->its * (10 * Nblk + 2 * 50 * Nblk) + 5 * Nblk); // guess
    521:   }
    522:   return PETSC_SUCCESS;
    523: }

    525: // KSP solver solve Ax = b; xout is output, bin is input
    526: static PetscErrorCode PCApply_BJKOKKOS(PC pc, Vec bin, Vec xout)
    527: {
    528:   PC_PCBJKOKKOS *jac = (PC_PCBJKOKKOS *)pc->data;
    529:   Mat            A = pc->pmat, Aseq = A;
    530:   PetscMPIInt    rank;

    532:   PetscFunctionBegin;
    533:   PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)A), &rank));
    534:   if (!A->spptr) Aseq = ((Mat_MPIAIJ *)A->data)->A; // MPI
    535:   PetscCall(MatSeqAIJKokkosSyncDevice(Aseq));
    536:   {
    537:     PetscInt           maxit = jac->ksp->max_it;
    538:     const PetscInt     conc = Kokkos::DefaultExecutionSpace().concurrency(), openmp = !!(conc < 1000), team_size = (openmp == 0 && PCBJKOKKOS_VEC_SIZE != 1) ? PCBJKOKKOS_TEAM_SIZE : 1;
    539:     const PetscInt     nwork = jac->nwork, nBlk = jac->nBlocks;
    540:     PetscScalar       *glb_xdata = NULL, *dummy;
    541:     PetscReal          rtol = jac->ksp->rtol, atol = jac->ksp->abstol, dtol = jac->ksp->divtol;
    542:     const PetscScalar *glb_idiag = jac->d_idiag_k->data(), *glb_bdata = NULL;
    543:     const PetscInt    *glb_Aai, *glb_Aaj, *d_bid_eqOffset = jac->d_bid_eqOffset_k->data();
    544:     const PetscScalar *glb_Aaa;
    545:     const PetscInt    *d_isicol = jac->d_isicol_k->data(), *d_isrow = jac->d_isrow_k->data();
    546:     PCFailedReason     pcreason;
    547:     KSPIndex           ksp_type_idx = jac->ksp_type_idx;
    548:     PetscMemType       mtype;
    549:     PetscContainer     container;
    550:     PetscInt           batch_sz;                // the number of repeated DMs, [DM_e_1, DM_e_2, DM_e_batch_sz, DM_i_1, ...]
    551:     VecScatter         plex_batch = NULL;       // not used
    552:     Vec                bvec;                    // a copy of b for scatter (just alias to bin now)
    553:     PetscBool          monitor  = jac->monitor; // captured
    554:     PetscInt           view_bid = jac->batch_target;
    555:     MatInfo            info;

    557:     PetscCall(MatSeqAIJGetCSRAndMemType(Aseq, &glb_Aai, &glb_Aaj, &dummy, &mtype));
    558:     jac->max_nits = 0;
    559:     glb_Aaa       = dummy;
    560:     if (jac->rank_target != rank) view_bid = -1; // turn off all but one process
    561:     PetscCall(MatGetInfo(A, MAT_LOCAL, &info));
    562:     // get field major is to map plex IO to/from block/field major
    563:     PetscCall(PetscObjectQuery((PetscObject)A, "plex_batch_is", (PetscObject *)&container));
    564:     if (container) {
    565:       PetscCall(VecDuplicate(bin, &bvec));
    566:       PetscCall(PetscContainerGetPointer(container, &plex_batch));
    567:       PetscCall(VecScatterBegin(plex_batch, bin, bvec, INSERT_VALUES, SCATTER_FORWARD));
    568:       PetscCall(VecScatterEnd(plex_batch, bin, bvec, INSERT_VALUES, SCATTER_FORWARD));
    569:       SETERRQ(PetscObjectComm((PetscObject)A), PETSC_ERR_USER, "No plex_batch_is -- require NO field major ordering for now");
    570:     } else {
    571:       bvec = bin;
    572:     }
    573:     // get x
    574:     PetscCall(VecGetArrayAndMemType(xout, &glb_xdata, &mtype));
    575: #if defined(PETSC_HAVE_CUDA)
    576:     PetscCheck(PetscMemTypeDevice(mtype), PetscObjectComm((PetscObject)pc), PETSC_ERR_ARG_WRONG, "No GPU data for x %d != %d", static_cast<int>(mtype), static_cast<int>(PETSC_MEMTYPE_DEVICE));
    577: #endif
    578:     PetscCall(VecGetArrayReadAndMemType(bvec, &glb_bdata, &mtype));
    579: #if defined(PETSC_HAVE_CUDA)
    580:     PetscCheck(PetscMemTypeDevice(mtype), PetscObjectComm((PetscObject)pc), PETSC_ERR_ARG_WRONG, "No GPU data for b");
    581: #endif
    582:     // get batch size
    583:     PetscCall(PetscObjectQuery((PetscObject)A, "batch size", (PetscObject *)&container));
    584:     if (container) {
    585:       PetscInt *pNf = NULL;
    586:       PetscCall(PetscContainerGetPointer(container, &pNf));
    587:       batch_sz = *pNf; // number of times to repeat the DMs
    588:     } else batch_sz = 1;
    589:     PetscCheck(nBlk % batch_sz == 0, PetscObjectComm((PetscObject)pc), PETSC_ERR_ARG_WRONG, "batch_sz = %" PetscInt_FMT ", nBlk = %" PetscInt_FMT, batch_sz, nBlk);
    590:     if (ksp_type_idx == BATCH_KSP_GMRESKK_IDX) {
    591:       // KK solver - move PETSc data into Kokkos Views, setup solver, solve, move data out of Kokkos, process metadata (convergence tests, etc.)
    592: #if defined(PETSC_HAVE_KOKKOS_KERNELS_BATCH)
    593:       PetscCall(PCApply_BJKOKKOSKERNELS(pc, glb_bdata, glb_xdata, glb_Aai, glb_Aaj, glb_Aaa, team_size, info, batch_sz, &pcreason));
    594: #else
    595:       PetscCheck(ksp_type_idx != BATCH_KSP_GMRESKK_IDX, PetscObjectComm((PetscObject)pc), PETSC_ERR_ARG_WRONG, "Type: BATCH_KSP_GMRES not supported for complex");
    596: #endif
    597:     } else { // Kokkos Krylov
    598:       using scr_mem_t    = Kokkos::DefaultExecutionSpace::scratch_memory_space;
    599:       using vect2D_scr_t = Kokkos::View<PetscScalar **, Kokkos::LayoutLeft, scr_mem_t>;
    600:       Kokkos::View<Batch_MetaData *, Kokkos::DefaultExecutionSpace> d_metadata("solver meta data", nBlk);
    601:       int                                                           stride_shared, stride_global, global_buff_words;
    602:       d_bid_eqOffset = jac->d_bid_eqOffset_k->data();
    603:       // solve each block independently
    604:       int scr_bytes_team_shared = 0, nShareVec = 0, nGlobBVec = 0;
    605:       if (jac->const_block_size) { // use shared memory for work vectors only if constant block size - TODO: test efficiency loss
    606:         size_t      maximum_shared_mem_size = 64000;
    607:         PetscDevice device;
    608:         PetscCall(PetscDeviceGetDefault_Internal(&device));
    609:         PetscCall(PetscDeviceGetAttribute(device, PETSC_DEVICE_ATTR_SIZE_T_SHARED_MEM_PER_BLOCK, &maximum_shared_mem_size));
    610:         stride_shared = jac->const_block_size;                                                   // captured
    611:         nShareVec     = maximum_shared_mem_size / (jac->const_block_size * sizeof(PetscScalar)); // integer floor, number of vectors that fit in shared
    612:         if (nShareVec > nwork) nShareVec = nwork;
    613:         else nGlobBVec = nwork - nShareVec;
    614:         global_buff_words     = jac->n * nGlobBVec;
    615:         scr_bytes_team_shared = jac->const_block_size * nShareVec * sizeof(PetscScalar);
    616:       } else {
    617:         scr_bytes_team_shared = 0;
    618:         stride_shared         = 0;
    619:         global_buff_words     = jac->n * nwork;
    620:         nGlobBVec             = nwork; // not needed == fix
    621:       }
    622:       stride_global = jac->n; // captured
    623: #if defined(PETSC_HAVE_CUDA)
    624:       nvtxRangePushA("batch-kokkos-solve");
    625: #endif
    626:       Kokkos::View<PetscScalar *, Kokkos::DefaultExecutionSpace> d_work_vecs_k("workvectors", global_buff_words); // global work vectors
    627: #if PCBJKOKKOS_VERBOSE_LEVEL > 1
    628:       PetscCall(PetscInfo(pc, "\tn = %d. %d shared bytes/team, %d global mem bytes, rtol=%e, num blocks %d, team_size=%d, %d vector threads, %d shared vectors, %d global vectors\n", (int)jac->n, scr_bytes_team_shared, global_buff_words, rtol, (int)nBlk, (int)team_size, PCBJKOKKOS_VEC_SIZE, nShareVec, nGlobBVec));
    629: #endif
    630:       PetscScalar *d_work_vecs = d_work_vecs_k.data();
    631:       Kokkos::parallel_for(
    632:         "Solve", Kokkos::TeamPolicy<Kokkos::LaunchBounds<256, 4>>(nBlk, team_size, PCBJKOKKOS_VEC_SIZE).set_scratch_size(PCBJKOKKOS_SHARED_LEVEL, Kokkos::PerTeam(scr_bytes_team_shared)), KOKKOS_LAMBDA(const team_member team) {
    633:           const int    blkID = team.league_rank(), start = d_bid_eqOffset[blkID], end = d_bid_eqOffset[blkID + 1];
    634:           vect2D_scr_t work_vecs_shared(team.team_scratch(PCBJKOKKOS_SHARED_LEVEL), end - start, nShareVec);
    635:           PetscScalar *work_buff_shared = work_vecs_shared.data();
    636:           PetscScalar *work_buff_global = &d_work_vecs[start]; // start inc'ed in
    637:           bool         print            = monitor && (blkID == view_bid);
    638:           switch (ksp_type_idx) {
    639:           case BATCH_KSP_BICG_IDX:
    640:             static_cast<void>(BJSolve_BICG(team, glb_Aai, glb_Aaj, glb_Aaa, d_isrow, d_isicol, work_buff_global, stride_global, nShareVec, work_buff_shared, stride_shared, rtol, atol, dtol, maxit, &d_metadata[blkID], start, end, glb_idiag, glb_bdata, glb_xdata, print));
    641:             break;
    642:           case BATCH_KSP_TFQMR_IDX:
    643:             static_cast<void>(BJSolve_TFQMR(team, glb_Aai, glb_Aaj, glb_Aaa, d_isrow, d_isicol, work_buff_global, stride_global, nShareVec, work_buff_shared, stride_shared, rtol, atol, dtol, maxit, &d_metadata[blkID], start, end, glb_idiag, glb_bdata, glb_xdata, print));
    644:             break;
    645:           default:
    646: #if defined(PETSC_USE_DEBUG) && !defined(PETSC_HAVE_SYCL)
    647:             printf("Unknown KSP type %d\n", ksp_type_idx);
    648: #else
    649:             /* void */;
    650: #endif
    651:           }
    652:         });
    653:       Kokkos::fence();
    654: #if defined(PETSC_HAVE_CUDA)
    655:       nvtxRangePop();
    656:       nvtxRangePushA("Post-solve-metadata");
    657: #endif
    658:       auto h_metadata = Kokkos::create_mirror(Kokkos::HostSpace::memory_space(), d_metadata);
    659:       Kokkos::deep_copy(h_metadata, d_metadata);
    660:       PetscInt max_nnit = -1;
    661: #if PCBJKOKKOS_VERBOSE_LEVEL > 1
    662:       PetscInt mbid = 0;
    663: #endif
    664:       int in[2], out[2];
    665:       if (jac->reason) { // -pc_bjkokkos_ksp_converged_reason
    666: #if PCBJKOKKOS_VERBOSE_LEVEL >= 3
    667:   #if PCBJKOKKOS_VERBOSE_LEVEL >= 4
    668:         PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Iterations\n"));
    669:   #endif
    670:         // assume species major
    671:   #if PCBJKOKKOS_VERBOSE_LEVEL == 3
    672:         if (batch_sz != 1) PetscCall(PetscPrintf(PetscObjectComm((PetscObject)A), "%s: max iterations per species:", ksp_type_idx == BATCH_KSP_BICG_IDX ? "bicg" : "tfqmr"));
    673:         else PetscCall(PetscPrintf(PetscObjectComm((PetscObject)A), "    Linear solve converged due to %s iterations ", ksp_type_idx == BATCH_KSP_BICG_IDX ? "bicg" : "tfqmr"));
    674:   #endif
    675:         for (PetscInt dmIdx = 0, head = 0, s = 0; dmIdx < jac->num_dms; dmIdx += batch_sz) {
    676:           for (PetscInt f = 0, idx = head; f < jac->dm_Nf[dmIdx]; f++, idx++, s++) {
    677:             for (int bid = 0; bid < batch_sz; bid++) {
    678:   #if PCBJKOKKOS_VERBOSE_LEVEL >= 4
    679:               jac->max_nits += h_metadata[idx + bid * jac->dm_Nf[dmIdx]].its; // report total number of iterations with high verbose
    680:               if (h_metadata[idx + bid * jac->dm_Nf[dmIdx]].its > max_nnit) {
    681:                 max_nnit = h_metadata[idx + bid * jac->dm_Nf[dmIdx]].its;
    682:                 mbid     = bid;
    683:               }
    684:   #else
    685:               if (h_metadata[idx + bid * jac->dm_Nf[dmIdx]].its > max_nnit) {
    686:                 jac->max_nits = max_nnit = h_metadata[idx + bid * jac->dm_Nf[dmIdx]].its;
    687:                 mbid                     = bid;
    688:               }
    689:   #endif
    690:             }
    691:   #if PCBJKOKKOS_VERBOSE_LEVEL >= 4
    692:             PetscCall(PetscPrintf(PetscObjectComm((PetscObject)A), "%2" PetscInt_FMT ":", s));
    693:             for (int bid = 0; bid < batch_sz; bid++) PetscCall(PetscPrintf(PetscObjectComm((PetscObject)A), "%3" PetscInt_FMT " ", h_metadata[idx + bid * jac->dm_Nf[dmIdx]].its));
    694:             PetscCall(PetscPrintf(PetscObjectComm((PetscObject)A), "\n"));
    695:   #else // == 3
    696:             PetscCall(PetscPrintf(PetscObjectComm((PetscObject)A), "%3" PetscInt_FMT " ", max_nnit));
    697:   #endif
    698:           }
    699:           head += batch_sz * jac->dm_Nf[dmIdx];
    700:         }
    701:   #if PCBJKOKKOS_VERBOSE_LEVEL == 3
    702:         PetscCall(PetscPrintf(PetscObjectComm((PetscObject)A), "\n"));
    703:   #endif
    704: #endif
    705:         if (max_nnit == -1) { // < 3
    706:           for (int blkID = 0; blkID < nBlk; blkID++) {
    707:             if (h_metadata[blkID].its > max_nnit) {
    708:               jac->max_nits = max_nnit = h_metadata[blkID].its;
    709: #if PCBJKOKKOS_VERBOSE_LEVEL > 1
    710:               mbid = blkID;
    711: #endif
    712:             }
    713:           }
    714:         }
    715:         in[0] = max_nnit;
    716:         in[1] = rank;
    717:         PetscCallMPI(MPIU_Allreduce(in, out, 1, MPI_2INT, MPI_MAXLOC, PetscObjectComm((PetscObject)A)));
    718: #if PCBJKOKKOS_VERBOSE_LEVEL > 1
    719:         if (0 == rank) {
    720:           if (batch_sz != 1)
    721:             PetscCall(PetscPrintf(PETSC_COMM_SELF, "[%d] Linear solve converged due to %s iterations %d (max), on block %" PetscInt_FMT ", species %" PetscInt_FMT " (max)\n", out[1], KSPConvergedReasons[h_metadata[mbid].reason], out[0], mbid % batch_sz, mbid / batch_sz));
    722:           else PetscCall(PetscPrintf(PETSC_COMM_SELF, "[%d] Linear solve converged due to %s iterations %d (max), on block %" PetscInt_FMT "\n", out[1], KSPConvergedReasons[h_metadata[mbid].reason], out[0], mbid));
    723:         }
    724: #endif
    725:       }
    726:       for (int blkID = 0; blkID < nBlk; blkID++) {
    727:         PetscCall(PetscLogGpuFlops((PetscLogDouble)h_metadata[blkID].flops));
    728:         PetscCheck(h_metadata[blkID].reason >= 0 || !jac->ksp->errorifnotconverged, PetscObjectComm((PetscObject)pc), PETSC_ERR_CONV_FAILED, "ERROR reason=%s, its=%" PetscInt_FMT ". species %" PetscInt_FMT ", batch %" PetscInt_FMT,
    729:                    KSPConvergedReasons[h_metadata[blkID].reason], h_metadata[blkID].its, blkID / batch_sz, blkID % batch_sz);
    730:       }
    731:       {
    732:         int errsum;
    733:         Kokkos::parallel_reduce(
    734:           nBlk,
    735:           KOKKOS_LAMBDA(const int idx, int &lsum) {
    736:             if (d_metadata[idx].reason < 0) ++lsum;
    737:           },
    738:           errsum);
    739:         pcreason = errsum ? PC_SUBPC_ERROR : PC_NOERROR;
    740:         if (!errsum && !jac->max_nits) { // set max its to give back to top KSP
    741:           for (int blkID = 0; blkID < nBlk; blkID++) {
    742:             if (h_metadata[blkID].its > jac->max_nits) jac->max_nits = h_metadata[blkID].its;
    743:           }
    744:         } else if (errsum) {
    745:           PetscCall(PetscPrintf(PETSC_COMM_SELF, "[%d] ERROR Kokkos batch solver did not converge in all solves\n", (int)rank));
    746:         }
    747:       }
    748: #if defined(PETSC_HAVE_CUDA)
    749:       nvtxRangePop();
    750: #endif
    751:     } // end of Kokkos (not Kernels) solvers block
    752:     PetscCall(VecRestoreArrayAndMemType(xout, &glb_xdata));
    753:     PetscCall(VecRestoreArrayReadAndMemType(bvec, &glb_bdata));
    754:     PetscCall(PCSetFailedReason(pc, pcreason));
    755:     // map back to Plex space - not used
    756:     if (plex_batch) {
    757:       PetscCall(VecCopy(xout, bvec));
    758:       PetscCall(VecScatterBegin(plex_batch, bvec, xout, INSERT_VALUES, SCATTER_REVERSE));
    759:       PetscCall(VecScatterEnd(plex_batch, bvec, xout, INSERT_VALUES, SCATTER_REVERSE));
    760:       PetscCall(VecDestroy(&bvec));
    761:     }
    762:   }
    763:   PetscFunctionReturn(PETSC_SUCCESS);
    764: }

    766: static PetscErrorCode PCSetUp_BJKOKKOS(PC pc)
    767: {
    768:   PC_PCBJKOKKOS *jac = (PC_PCBJKOKKOS *)pc->data;
    769:   Mat            A = pc->pmat, Aseq = A; // use filtered block matrix, really "P"
    770:   PetscBool      flg;

    772:   PetscFunctionBegin;
    773:   PetscCheck(A, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONG, "No matrix - A is used above");
    774:   PetscCall(PetscObjectTypeCompareAny((PetscObject)A, &flg, MATSEQAIJKOKKOS, MATMPIAIJKOKKOS, MATAIJKOKKOS, ""));
    775:   PetscCheck(flg, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONG, "must use '-[dm_]mat_type aijkokkos -[dm_]vec_type kokkos' for -pc_type bjkokkos");
    776:   if (!A->spptr) Aseq = ((Mat_MPIAIJ *)A->data)->A; // MPI
    777:   PetscCall(MatSeqAIJKokkosSyncDevice(Aseq));
    778:   {
    779:     PetscInt    Istart, Iend;
    780:     PetscMPIInt rank;
    781:     PetscCallMPI(MPI_Comm_rank(PetscObjectComm((PetscObject)A), &rank));
    782:     PetscCall(MatGetOwnershipRange(A, &Istart, &Iend));
    783:     if (!jac->vec_diag) {
    784:       Vec     *subX = NULL;
    785:       DM       pack, *subDM = NULL;
    786:       PetscInt nDMs, n, *block_sizes = NULL;
    787:       IS       isrow, isicol;
    788:       { // Permute the matrix to get a block diagonal system: d_isrow_k, d_isicol_k
    789:         MatOrderingType rtype;
    790:         const PetscInt *rowindices, *icolindices;
    791:         rtype = MATORDERINGRCM;
    792:         // get permutation. And invert. should we convert to local indices?
    793:         PetscCall(MatGetOrdering(Aseq, rtype, &isrow, &isicol)); // only seems to work for seq matrix
    794:         PetscCall(ISDestroy(&isrow));
    795:         PetscCall(ISInvertPermutation(isicol, PETSC_DECIDE, &isrow)); // THIS IS BACKWARD -- isrow is inverse
    796:         // if (rank==1) PetscCall(ISView(isicol, PETSC_VIEWER_STDOUT_SELF));
    797:         if (0) {
    798:           Mat mat_block_order; // debug
    799:           PetscCall(ISShift(isicol, Istart, isicol));
    800:           PetscCall(MatCreateSubMatrix(A, isicol, isicol, MAT_INITIAL_MATRIX, &mat_block_order));
    801:           PetscCall(ISShift(isicol, -Istart, isicol));
    802:           PetscCall(MatViewFromOptions(mat_block_order, NULL, "-ksp_batch_reorder_view"));
    803:           PetscCall(MatDestroy(&mat_block_order));
    804:         }
    805:         PetscCall(ISGetIndices(isrow, &rowindices)); // local idx
    806:         PetscCall(ISGetIndices(isicol, &icolindices));
    807:         const Kokkos::View<PetscInt *, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_isrow_k((PetscInt *)rowindices, A->rmap->n);
    808:         const Kokkos::View<PetscInt *, Kokkos::HostSpace, Kokkos::MemoryTraits<Kokkos::Unmanaged>> h_isicol_k((PetscInt *)icolindices, A->rmap->n);
    809:         jac->d_isrow_k  = new Kokkos::View<PetscInt *>(Kokkos::create_mirror(DefaultMemorySpace(), h_isrow_k));
    810:         jac->d_isicol_k = new Kokkos::View<PetscInt *>(Kokkos::create_mirror(DefaultMemorySpace(), h_isicol_k));
    811:         Kokkos::deep_copy(*jac->d_isrow_k, h_isrow_k);
    812:         Kokkos::deep_copy(*jac->d_isicol_k, h_isicol_k);
    813:         PetscCall(ISRestoreIndices(isrow, &rowindices));
    814:         PetscCall(ISRestoreIndices(isicol, &icolindices));
    815:         // if (rank==1) PetscCall(ISView(isicol, PETSC_VIEWER_STDOUT_SELF));
    816:       }
    817:       // get block sizes & allocate vec_diag
    818:       PetscCall(PCGetDM(pc, &pack));
    819:       if (pack) {
    820:         PetscCall(PetscObjectTypeCompare((PetscObject)pack, DMCOMPOSITE, &flg));
    821:         if (flg) {
    822:           PetscCall(DMCompositeGetNumberDM(pack, &nDMs));
    823:           PetscCall(DMCreateGlobalVector(pack, &jac->vec_diag));
    824:         } else pack = NULL; // flag for no DM
    825:       }
    826:       if (!jac->vec_diag) { // get 'nDMs' and sizes 'block_sizes' w/o DMComposite. TODO: User could provide ISs
    827:         PetscInt        bsrt, bend, ncols, ntot = 0;
    828:         const PetscInt *colsA, nloc = Iend - Istart;
    829:         const PetscInt *rowindices, *icolindices;
    830:         PetscCall(PetscMalloc1(nloc, &block_sizes)); // very inefficient, to big
    831:         PetscCall(ISGetIndices(isrow, &rowindices));
    832:         PetscCall(ISGetIndices(isicol, &icolindices));
    833:         nDMs = 0;
    834:         bsrt = 0;
    835:         bend = 1;
    836:         for (PetscInt row_B = 0; row_B < nloc; row_B++) { // for all rows in block diagonal space
    837:           PetscInt rowA = icolindices[row_B], minj = PETSC_INT_MAX, maxj = 0;
    838:           //PetscCall(PetscPrintf(PETSC_COMM_SELF, "\t[%d] rowA = %d\n",rank,rowA));
    839:           PetscCall(MatGetRow(Aseq, rowA, &ncols, &colsA, NULL)); // not sorted in permutation
    840:           PetscCheck(ncols, PetscObjectComm((PetscObject)pc), PETSC_ERR_ARG_WRONG, "Empty row not supported: %" PetscInt_FMT, row_B);
    841:           for (PetscInt colj = 0; colj < ncols; colj++) {
    842:             PetscInt colB = rowindices[colsA[colj]]; // use local idx
    843:             //PetscCall(PetscPrintf(PETSC_COMM_SELF, "\t\t[%d] colB = %d\n",rank,colB));
    844:             PetscCheck(colB >= 0 && colB < nloc, PetscObjectComm((PetscObject)pc), PETSC_ERR_ARG_WRONG, "colB < 0: %" PetscInt_FMT, colB);
    845:             if (colB > maxj) maxj = colB;
    846:             if (colB < minj) minj = colB;
    847:           }
    848:           PetscCall(MatRestoreRow(Aseq, rowA, &ncols, &colsA, NULL));
    849:           if (minj >= bend) { // first column is > max of last block -- new block or last block
    850:             //PetscCall(PetscPrintf(PetscObjectComm((PetscObject)A), "\t\t finish block %d, N loc = %d (%d,%d)\n", nDMs+1, bend - bsrt,bsrt,bend));
    851:             block_sizes[nDMs] = bend - bsrt;
    852:             ntot += block_sizes[nDMs];
    853:             PetscCheck(minj == bend, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "minj != bend: %" PetscInt_FMT " != %" PetscInt_FMT, minj, bend);
    854:             bsrt = bend;
    855:             bend++; // start with size 1 in new block
    856:             nDMs++;
    857:           }
    858:           if (maxj + 1 > bend) bend = maxj + 1;
    859:           PetscCheck(minj >= bsrt || row_B == Iend - 1, PetscObjectComm((PetscObject)pc), PETSC_ERR_ARG_WRONG, "%" PetscInt_FMT ") minj < bsrt: %" PetscInt_FMT " != %" PetscInt_FMT, rowA, minj, bsrt);
    860:           //PetscCall(PetscPrintf(PETSC_COMM_SELF, "[%d] %d) row %d.%d) cols %d : %d ; bsrt = %d, bend = %d\n",rank,row_B,nDMs,rowA,minj,maxj,bsrt,bend));
    861:         }
    862:         // do last block
    863:         //PetscCall(PetscPrintf(PETSC_COMM_SELF, "\t\t\t [%d] finish block %d, N loc = %d (%d,%d)\n", rank, nDMs+1, bend - bsrt,bsrt,bend));
    864:         block_sizes[nDMs] = bend - bsrt;
    865:         ntot += block_sizes[nDMs];
    866:         nDMs++;
    867:         // cleanup
    868:         PetscCheck(ntot == nloc, PetscObjectComm((PetscObject)pc), PETSC_ERR_ARG_WRONG, "n total != n local: %" PetscInt_FMT " != %" PetscInt_FMT, ntot, nloc);
    869:         PetscCall(ISRestoreIndices(isrow, &rowindices));
    870:         PetscCall(ISRestoreIndices(isicol, &icolindices));
    871:         PetscCall(PetscRealloc(sizeof(PetscInt) * nDMs, &block_sizes));
    872:         PetscCall(MatCreateVecs(A, &jac->vec_diag, NULL));
    873:         PetscCall(PetscInfo(pc, "Setup Matrix based meta data (not DMComposite not attached to PC) %" PetscInt_FMT " sub domains\n", nDMs));
    874:       }
    875:       PetscCall(ISDestroy(&isrow));
    876:       PetscCall(ISDestroy(&isicol));
    877:       jac->num_dms = nDMs;
    878:       PetscCall(VecGetLocalSize(jac->vec_diag, &n));
    879:       jac->n         = n;
    880:       jac->d_idiag_k = new Kokkos::View<PetscScalar *, Kokkos::LayoutRight>("idiag", n);
    881:       // options
    882:       PetscCall(PCBJKOKKOSCreateKSP_BJKOKKOS(pc));
    883:       PetscCall(KSPSetFromOptions(jac->ksp));
    884:       PetscCall(PetscObjectTypeCompareAny((PetscObject)jac->ksp, &flg, KSPBICG, ""));
    885:       if (flg) {
    886:         jac->ksp_type_idx = BATCH_KSP_BICG_IDX;
    887:         jac->nwork        = 7;
    888:       } else {
    889:         PetscCall(PetscObjectTypeCompareAny((PetscObject)jac->ksp, &flg, KSPTFQMR, ""));
    890:         if (flg) {
    891:           jac->ksp_type_idx = BATCH_KSP_TFQMR_IDX;
    892:           jac->nwork        = 10;
    893:         } else {
    894: #if defined(PETSC_HAVE_KOKKOS_KERNELS_BATCH)
    895:           PetscCall(PetscObjectTypeCompareAny((PetscObject)jac->ksp, &flg, KSPGMRES, ""));
    896:           PetscCheck(flg, PetscObjectComm((PetscObject)A), PETSC_ERR_ARG_WRONG, "Unsupported batch ksp type");
    897:           jac->ksp_type_idx = BATCH_KSP_GMRESKK_IDX;
    898:           jac->nwork        = 0;
    899: #else
    900:           KSPType ksptype;
    901:           PetscCall(KSPGetType(jac->ksp, &ksptype));
    902:           PetscCheck(flg, PetscObjectComm((PetscObject)pc), PETSC_ERR_ARG_WRONG, "Type: %s not supported in complex", ksptype);
    903: #endif
    904:         }
    905:       }
    906:       PetscOptionsBegin(PetscObjectComm((PetscObject)jac->ksp), ((PetscObject)jac->ksp)->prefix, "Options for Kokkos batch solver", "none");
    907:       PetscCall(PetscOptionsBool("-ksp_converged_reason", "", "bjkokkos.kokkos.cxx.c", jac->reason, &jac->reason, NULL));
    908:       PetscCall(PetscOptionsBool("-ksp_monitor", "", "bjkokkos.kokkos.cxx.c", jac->monitor, &jac->monitor, NULL));
    909:       PetscCall(PetscOptionsInt("-ksp_batch_target", "", "bjkokkos.kokkos.cxx.c", jac->batch_target, &jac->batch_target, NULL));
    910:       PetscCall(PetscOptionsInt("-ksp_rank_target", "", "bjkokkos.kokkos.cxx.c", jac->rank_target, &jac->rank_target, NULL));
    911:       PetscCall(PetscOptionsInt("-ksp_batch_nsolves_team", "", "bjkokkos.kokkos.cxx.c", jac->nsolves_team, &jac->nsolves_team, NULL));
    912:       PetscCheck(jac->batch_target < jac->num_dms, PETSC_COMM_WORLD, PETSC_ERR_ARG_WRONG, "-ksp_batch_target (%" PetscInt_FMT ") >= number of DMs (%" PetscInt_FMT ")", jac->batch_target, jac->num_dms);
    913:       PetscOptionsEnd();
    914:       // get blocks - jac->d_bid_eqOffset_k
    915:       if (pack) {
    916:         PetscCall(PetscMalloc(sizeof(*subX) * nDMs, &subX));
    917:         PetscCall(PetscMalloc(sizeof(*subDM) * nDMs, &subDM));
    918:       }
    919:       PetscCall(PetscMalloc(sizeof(*jac->dm_Nf) * nDMs, &jac->dm_Nf));
    920:       PetscCall(PetscInfo(pc, "Have %" PetscInt_FMT " blocks, n=%" PetscInt_FMT " rtol=%g type = %s\n", nDMs, n, (double)jac->ksp->rtol, ((PetscObject)jac->ksp)->type_name));
    921:       if (pack) PetscCall(DMCompositeGetEntriesArray(pack, subDM));
    922:       jac->nBlocks = 0;
    923:       for (PetscInt ii = 0; ii < nDMs; ii++) {
    924:         PetscInt Nf;
    925:         if (subDM) {
    926:           DM           dm = subDM[ii];
    927:           PetscSection section;
    928:           PetscCall(DMGetLocalSection(dm, &section));
    929:           PetscCall(PetscSectionGetNumFields(section, &Nf));
    930:         } else Nf = 1;
    931:         jac->nBlocks += Nf;
    932: #if PCBJKOKKOS_VERBOSE_LEVEL <= 2
    933:         if (ii == 0) PetscCall(PetscInfo(pc, "%" PetscInt_FMT ") %" PetscInt_FMT " blocks (%" PetscInt_FMT " total)\n", ii, Nf, jac->nBlocks));
    934: #else
    935:         PetscCall(PetscInfo(pc, "%" PetscInt_FMT ") %" PetscInt_FMT " blocks (%" PetscInt_FMT " total)\n", ii, Nf, jac->nBlocks));
    936: #endif
    937:         jac->dm_Nf[ii] = Nf;
    938:       }
    939:       { // d_bid_eqOffset_k
    940:         Kokkos::View<PetscInt *, Kokkos::LayoutRight, Kokkos::HostSpace> h_block_offsets("block_offsets", jac->nBlocks + 1);
    941:         if (pack) PetscCall(DMCompositeGetAccessArray(pack, jac->vec_diag, nDMs, NULL, subX));
    942:         h_block_offsets[0]    = 0;
    943:         jac->const_block_size = -1;
    944:         for (PetscInt ii = 0, idx = 0; ii < nDMs; ii++) {
    945:           PetscInt nloc, nblk;
    946:           if (pack) PetscCall(VecGetSize(subX[ii], &nloc));
    947:           else nloc = block_sizes[ii];
    948:           nblk = nloc / jac->dm_Nf[ii];
    949:           PetscCheck(nloc % jac->dm_Nf[ii] == 0, PetscObjectComm((PetscObject)pc), PETSC_ERR_USER, "nloc%%jac->dm_Nf[ii] (%" PetscInt_FMT ") != 0 DMs", nloc % jac->dm_Nf[ii]);
    950:           for (PetscInt jj = 0; jj < jac->dm_Nf[ii]; jj++, idx++) {
    951:             h_block_offsets[idx + 1] = h_block_offsets[idx] + nblk;
    952: #if PCBJKOKKOS_VERBOSE_LEVEL <= 2
    953:             if (idx == 0) PetscCall(PetscInfo(pc, "Add first of %" PetscInt_FMT " blocks with %" PetscInt_FMT " equations\n", jac->nBlocks, nblk));
    954: #else
    955:             PetscCall(PetscInfo(pc, "\t%" PetscInt_FMT ") Add block with %" PetscInt_FMT " equations of %" PetscInt_FMT "\n", idx + 1, nblk, jac->nBlocks));
    956: #endif
    957:             if (jac->const_block_size == -1) jac->const_block_size = nblk;
    958:             else if (jac->const_block_size > 0 && jac->const_block_size != nblk) jac->const_block_size = 0;
    959:           }
    960:         }
    961:         if (pack) {
    962:           PetscCall(DMCompositeRestoreAccessArray(pack, jac->vec_diag, jac->nBlocks, NULL, subX));
    963:           PetscCall(PetscFree(subX));
    964:           PetscCall(PetscFree(subDM));
    965:         }
    966:         jac->d_bid_eqOffset_k = new Kokkos::View<PetscInt *, Kokkos::LayoutRight>(Kokkos::create_mirror(Kokkos::DefaultExecutionSpace::memory_space(), h_block_offsets));
    967:         Kokkos::deep_copy(*jac->d_bid_eqOffset_k, h_block_offsets);
    968:       }
    969:       if (!pack) PetscCall(PetscFree(block_sizes));
    970:     }
    971:     { // get jac->d_idiag_k (PC setup),
    972:       const PetscInt    *d_ai, *d_aj;
    973:       const PetscScalar *d_aa;
    974:       const PetscInt     conc = Kokkos::DefaultExecutionSpace().concurrency(), openmp = !!(conc < 1000), team_size = (openmp == 0 && PCBJKOKKOS_VEC_SIZE != 1) ? PCBJKOKKOS_TEAM_SIZE : 1;
    975:       const PetscInt    *d_bid_eqOffset = jac->d_bid_eqOffset_k->data(), *r = jac->d_isrow_k->data(), *ic = jac->d_isicol_k->data();
    976:       PetscScalar       *d_idiag = jac->d_idiag_k->data(), *dummy;
    977:       PetscMemType       mtype;
    978:       PetscCall(MatSeqAIJGetCSRAndMemType(Aseq, &d_ai, &d_aj, &dummy, &mtype));
    979:       d_aa = dummy;
    980:       Kokkos::parallel_for(
    981:         "Diag", Kokkos::TeamPolicy<>(jac->nBlocks, team_size, PCBJKOKKOS_VEC_SIZE), KOKKOS_LAMBDA(const team_member team) {
    982:           const PetscInt blkID = team.league_rank();
    983:           Kokkos::parallel_for(Kokkos::TeamThreadRange(team, d_bid_eqOffset[blkID], d_bid_eqOffset[blkID + 1]), [=](const int rowb) {
    984:             const PetscInt     rowa = ic[rowb], ai = d_ai[rowa], *aj = d_aj + ai; // grab original data
    985:             const PetscScalar *aa   = d_aa + ai;
    986:             const PetscInt     nrow = d_ai[rowa + 1] - ai;
    987:             int                found;
    988:             Kokkos::parallel_reduce(
    989:               Kokkos::ThreadVectorRange(team, nrow),
    990:               [=](const int &j, int &count) {
    991:                 const PetscInt colb = r[aj[j]];
    992:                 if (colb == rowb) {
    993:                   d_idiag[rowb] = 1. / aa[j];
    994:                   count++;
    995:                 }
    996:               },
    997:               found);
    998: #if defined(PETSC_USE_DEBUG) && !defined(PETSC_HAVE_SYCL)
    999:             if (found != 1) Kokkos::single(Kokkos::PerThread(team), [=]() { printf("ERRORrow %d) found = %d\n", rowb, found); });
    1000: #endif
    1001:           });
    1002:         });
    1003:     }
    1004:   }
    1005:   PetscFunctionReturn(PETSC_SUCCESS);
    1006: }

    1008: /* Default destroy, if it has never been setup */
    1009: static PetscErrorCode PCReset_BJKOKKOS(PC pc)
    1010: {
    1011:   PC_PCBJKOKKOS *jac = (PC_PCBJKOKKOS *)pc->data;

    1013:   PetscFunctionBegin;
    1014:   PetscCall(KSPDestroy(&jac->ksp));
    1015:   PetscCall(VecDestroy(&jac->vec_diag));
    1016:   if (jac->d_bid_eqOffset_k) delete jac->d_bid_eqOffset_k;
    1017:   if (jac->d_idiag_k) delete jac->d_idiag_k;
    1018:   if (jac->d_isrow_k) delete jac->d_isrow_k;
    1019:   if (jac->d_isicol_k) delete jac->d_isicol_k;
    1020:   jac->d_bid_eqOffset_k = NULL;
    1021:   jac->d_idiag_k        = NULL;
    1022:   jac->d_isrow_k        = NULL;
    1023:   jac->d_isicol_k       = NULL;
    1024:   PetscCall(PetscObjectComposeFunction((PetscObject)pc, "PCBJKOKKOSGetKSP_C", NULL)); // not published now (causes configure errors)
    1025:   PetscCall(PetscObjectComposeFunction((PetscObject)pc, "PCBJKOKKOSSetKSP_C", NULL));
    1026:   PetscCall(PetscFree(jac->dm_Nf));
    1027:   jac->dm_Nf = NULL;
    1028:   if (jac->rowOffsets) delete jac->rowOffsets;
    1029:   if (jac->colIndices) delete jac->colIndices;
    1030:   if (jac->batch_b) delete jac->batch_b;
    1031:   if (jac->batch_x) delete jac->batch_x;
    1032:   if (jac->batch_values) delete jac->batch_values;
    1033:   jac->rowOffsets   = NULL;
    1034:   jac->colIndices   = NULL;
    1035:   jac->batch_b      = NULL;
    1036:   jac->batch_x      = NULL;
    1037:   jac->batch_values = NULL;
    1038:   PetscFunctionReturn(PETSC_SUCCESS);
    1039: }

    1041: static PetscErrorCode PCDestroy_BJKOKKOS(PC pc)
    1042: {
    1043:   PetscFunctionBegin;
    1044:   PetscCall(PCReset_BJKOKKOS(pc));
    1045:   PetscCall(PetscFree(pc->data));
    1046:   PetscFunctionReturn(PETSC_SUCCESS);
    1047: }

    1049: static PetscErrorCode PCView_BJKOKKOS(PC pc, PetscViewer viewer)
    1050: {
    1051:   PC_PCBJKOKKOS *jac = (PC_PCBJKOKKOS *)pc->data;
    1052:   PetscBool      isascii;

    1054:   PetscFunctionBegin;
    1055:   if (!jac->ksp) PetscCall(PCBJKOKKOSCreateKSP_BJKOKKOS(pc));
    1056:   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &isascii));
    1057:   if (isascii) {
    1058:     PetscCall(PetscViewerASCIIPrintf(viewer, "  Batched device linear solver: Krylov (KSP) method with Jacobi preconditioning\n"));
    1059:     PetscCall(PetscViewerASCIIPrintf(viewer, "\t\tnwork = %" PetscInt_FMT ", rel tol = %e, abs tol = %e, div tol = %e, max it =%" PetscInt_FMT ", type = %s\n", jac->nwork, jac->ksp->rtol, jac->ksp->abstol, jac->ksp->divtol, jac->ksp->max_it,
    1060:                                      ((PetscObject)jac->ksp)->type_name));
    1061:   }
    1062:   PetscFunctionReturn(PETSC_SUCCESS);
    1063: }

    1065: static PetscErrorCode PCSetFromOptions_BJKOKKOS(PC pc, PetscOptionItems PetscOptionsObject)
    1066: {
    1067:   PetscFunctionBegin;
    1068:   PetscOptionsHeadBegin(PetscOptionsObject, "PC BJKOKKOS options");
    1069:   PetscOptionsHeadEnd();
    1070:   PetscFunctionReturn(PETSC_SUCCESS);
    1071: }

    1073: static PetscErrorCode PCBJKOKKOSSetKSP_BJKOKKOS(PC pc, KSP ksp)
    1074: {
    1075:   PC_PCBJKOKKOS *jac = (PC_PCBJKOKKOS *)pc->data;

    1077:   PetscFunctionBegin;
    1078:   PetscCall(PetscObjectReference((PetscObject)ksp));
    1079:   PetscCall(KSPDestroy(&jac->ksp));
    1080:   jac->ksp = ksp;
    1081:   PetscFunctionReturn(PETSC_SUCCESS);
    1082: }

    1084: /*@
    1085:   PCBJKOKKOSSetKSP - Sets the `KSP` context for `PCBJKOKKOS`

    1087:   Collective

    1089:   Input Parameters:
    1090: + pc  - the `PCBJKOKKOS` preconditioner context
    1091: - ksp - the `KSP` solver

    1093:   Level: advanced

    1095:   Notes:
    1096:   The `PC` and the `KSP` must have the same communicator

    1098:   If the `PC` is not `PCBJKOKKOS` this function returns without doing anything

    1100: .seealso: [](ch_ksp), `PCBJKOKKOSGetKSP()`, `PCBJKOKKOS`
    1101: @*/
    1102: PetscErrorCode PCBJKOKKOSSetKSP(PC pc, KSP ksp)
    1103: {
    1104:   PetscFunctionBegin;
    1107:   PetscCheckSameComm(pc, 1, ksp, 2);
    1108:   PetscTryMethod(pc, "PCBJKOKKOSSetKSP_C", (PC, KSP), (pc, ksp));
    1109:   PetscFunctionReturn(PETSC_SUCCESS);
    1110: }

    1112: static PetscErrorCode PCBJKOKKOSGetKSP_BJKOKKOS(PC pc, KSP *ksp)
    1113: {
    1114:   PC_PCBJKOKKOS *jac = (PC_PCBJKOKKOS *)pc->data;

    1116:   PetscFunctionBegin;
    1117:   if (!jac->ksp) PetscCall(PCBJKOKKOSCreateKSP_BJKOKKOS(pc));
    1118:   *ksp = jac->ksp;
    1119:   PetscFunctionReturn(PETSC_SUCCESS);
    1120: }

    1122: /*@
    1123:   PCBJKOKKOSGetKSP - Gets the `KSP` context for the `PCBJKOKKOS` preconditioner

    1125:   Not Collective but `KSP` returned is parallel if `PC` was parallel

    1127:   Input Parameter:
    1128: . pc - the preconditioner context

    1130:   Output Parameter:
    1131: . ksp - the `KSP` solver

    1133:   Level: advanced

    1135:   Notes:
    1136:   You must call `KSPSetUp()` before calling `PCBJKOKKOSGetKSP()`.

    1138:   If the `PC` is not a `PCBJKOKKOS` object it raises an error

    1140: .seealso: [](ch_ksp), `PCBJKOKKOS`, `PCBJKOKKOSSetKSP()`
    1141: @*/
    1142: PetscErrorCode PCBJKOKKOSGetKSP(PC pc, KSP *ksp)
    1143: {
    1144:   PetscFunctionBegin;
    1146:   PetscAssertPointer(ksp, 2);
    1147:   PetscUseMethod(pc, "PCBJKOKKOSGetKSP_C", (PC, KSP *), (pc, ksp));
    1148:   PetscFunctionReturn(PETSC_SUCCESS);
    1149: }

    1151: static PetscErrorCode PCPostSolve_BJKOKKOS(PC pc, KSP ksp, Vec b, Vec x)
    1152: {
    1153:   PC_PCBJKOKKOS *jac = (PC_PCBJKOKKOS *)pc->data;

    1155:   PetscFunctionBegin;
    1157:   ksp->its = jac->max_nits;
    1158:   PetscFunctionReturn(PETSC_SUCCESS);
    1159: }

    1161: static PetscErrorCode PCPreSolve_BJKOKKOS(PC pc, KSP ksp, Vec b, Vec x)
    1162: {
    1163:   PC_PCBJKOKKOS *jac = (PC_PCBJKOKKOS *)pc->data;

    1165:   PetscFunctionBegin;
    1167:   jac->ksp->errorifnotconverged = ksp->errorifnotconverged;
    1168:   PetscFunctionReturn(PETSC_SUCCESS);
    1169: }

    1171: /*MC
    1172:      PCBJKOKKOS - A batched Krylov/block Jacobi solver that runs a solve of each diagaonl block of a block diagonal `MATSEQAIJ` in a Kokkos thread group

    1174:    Options Database Key:
    1175: .  -pc_bjkokkos_ - options prefix for its `KSP` options

    1177:    Level: intermediate

    1179:    Note:
    1180:    For use with `-ksp_type preonly` to bypass any computation on the CPU

    1182:    Developer Notes:
    1183:    The entire Krylov (TFQMR or BICG) with diagonal preconditioning for each block of a block diagnaol matrix runs in a Kokkos thread group (eg, one block per SM on NVIDIA). It supports taking a non-block diagonal matrix but this is not tested. One should create an explicit block diagonal matrix and use that as the matrix for constructing the preconditioner in the outer `KSP` solver. Variable block size are supported and tested in src/ts/utils/dmplexlandau/tutorials/ex[1|2].c

    1185: .seealso: [](ch_ksp), `PCCreate()`, `PCSetType()`, `PCType`, `PC`, `PCBJACOBI`,
    1186:           `PCSHELL`, `PCCOMPOSITE`, `PCSetUseAmat()`, `PCBJKOKKOSGetKSP()`
    1187: M*/

    1189: PETSC_EXTERN PetscErrorCode PCCreate_BJKOKKOS(PC pc)
    1190: {
    1191:   PC_PCBJKOKKOS *jac;

    1193:   PetscFunctionBegin;
    1194:   PetscCall(PetscNew(&jac));
    1195:   pc->data = (void *)jac;

    1197:   jac->ksp              = NULL;
    1198:   jac->vec_diag         = NULL;
    1199:   jac->d_bid_eqOffset_k = NULL;
    1200:   jac->d_idiag_k        = NULL;
    1201:   jac->d_isrow_k        = NULL;
    1202:   jac->d_isicol_k       = NULL;
    1203:   jac->nBlocks          = 1;
    1204:   jac->max_nits         = 0;

    1206:   PetscCall(PetscMemzero(pc->ops, sizeof(struct _PCOps)));
    1207:   pc->ops->apply          = PCApply_BJKOKKOS;
    1208:   pc->ops->applytranspose = NULL;
    1209:   pc->ops->setup          = PCSetUp_BJKOKKOS;
    1210:   pc->ops->reset          = PCReset_BJKOKKOS;
    1211:   pc->ops->destroy        = PCDestroy_BJKOKKOS;
    1212:   pc->ops->setfromoptions = PCSetFromOptions_BJKOKKOS;
    1213:   pc->ops->view           = PCView_BJKOKKOS;
    1214:   pc->ops->postsolve      = PCPostSolve_BJKOKKOS;
    1215:   pc->ops->presolve       = PCPreSolve_BJKOKKOS;

    1217:   jac->rowOffsets   = NULL;
    1218:   jac->colIndices   = NULL;
    1219:   jac->batch_b      = NULL;
    1220:   jac->batch_x      = NULL;
    1221:   jac->batch_values = NULL;

    1223:   PetscCall(PetscObjectComposeFunction((PetscObject)pc, "PCBJKOKKOSGetKSP_C", PCBJKOKKOSGetKSP_BJKOKKOS));
    1224:   PetscCall(PetscObjectComposeFunction((PetscObject)pc, "PCBJKOKKOSSetKSP_C", PCBJKOKKOSSetKSP_BJKOKKOS));
    1225:   PetscFunctionReturn(PETSC_SUCCESS);
    1226: }